diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json index 3676c0136..2042b0584 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json @@ -36,11 +36,11 @@ "audio_channels": 2, "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, - "dit_quant_scheme": "fp8-sgl", - "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors", + "dit_quant_scheme": "fp8-f16-accum", + "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/minimax_h3_dit_fp8_f16_accum.safetensors", "video_vae_quantized": true, - "video_vae_quant_scheme": "fp8-sgl", - "video_vae_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors", + "video_vae_quant_scheme": "fp8-f16-accum", + "video_vae_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/minimax_h3_video_vae_fp8_f16_accum.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { diff --git a/docs/EN/source/method_tutorials/quantization.md b/docs/EN/source/method_tutorials/quantization.md index da355a92a..142d3b62d 100644 --- a/docs/EN/source/method_tutorials/quantization.md +++ b/docs/EN/source/method_tutorials/quantization.md @@ -13,6 +13,7 @@ LightX2V supports quantized inference for DIT, T5, and CLIP models, reducing mem | `fp8-vllm` | FP8 channel symmetric | FP8 channel dynamic symmetric | [VLLM](https://github.com/vllm-project/vllm) | H100/H200/H800, RTX 40 series, etc. | | `int8-vllm` | INT8 channel symmetric | INT8 channel dynamic symmetric | [VLLM](https://github.com/vllm-project/vllm) | A100/A800, RTX 30/40 series, etc. | | `fp8-sgl` | FP8 channel symmetric | FP8 channel dynamic symmetric | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | H100/H200/H800, RTX 40 series, etc. | +| `fp8-f16-accum` | FP8 channel symmetric | FP8 row-wise dynamic symmetric | CUTLASS FP16 accumulation | RTX 5090 (SM120) | | `int8-sgl` | INT8 channel symmetric | INT8 channel dynamic symmetric | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | A100/A800, RTX 30/40 series, etc. | | `fp8-q8f` | FP8 channel symmetric | FP8 channel dynamic symmetric | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40 series, L40S, etc. | | `int8-q8f` | INT8 channel symmetric | INT8 channel dynamic symmetric | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40 series, L40S, etc. | @@ -67,7 +68,7 @@ For detailed quantization tool usage, refer to: [Model Conversion Documentation] #### Supported Quantization Modes -DIT quantization modes (`dit_quant_scheme`) support: `fp8-vllm`, `int8-vllm`, `fp8-sgl`, `int8-sgl`, `fp8-q8f`, `int8-q8f`, `int8-torchao`, `int4-g128-marlin`, `fp8-b128-deepgemm` +DIT quantization modes (`dit_quant_scheme`) support: `fp8-vllm`, `int8-vllm`, `fp8-sgl`, `fp8-f16-accum`, `int8-sgl`, `fp8-q8f`, `int8-q8f`, `int8-torchao`, `int4-g128-marlin`, `fp8-b128-deepgemm` #### Configuration Example @@ -81,6 +82,72 @@ DIT quantization modes (`dit_quant_scheme`) support: `fp8-vllm`, `int8-vllm`, `f > 💡 **Tip**: When there's only one DIT model in the script's `model_path`, `dit_quantized_ckpt` doesn't need to be specified separately. +#### MiniMax-H3 FP8 with FP16 Accumulation + +On RTX 5090, MiniMax-H3 can use FP8 inputs with FP16 accumulation through `fp8-f16-accum`. Convert +the weights with the `h3-fp8-f16-accum` profile; regular `fp8-sgl` checkpoints are not compatible. +DiT and Video VAE decoder are converted separately. The profile selects the qmax-14 projections, +keeps standard FP8 quantization for the remaining layers, and records the policy in safetensors +metadata. + +```bash +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/transformer \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_dit_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3 \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file + +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/vae \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_video_vae_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3_video_vae_decoder \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file +``` + +```json +{ + "dit_quantized": true, + "dit_quant_scheme": "fp8-f16-accum", + "dit_quantized_ckpt": "/path/to/minimax_h3_dit_fp8_f16_accum.safetensors", + "video_vae_quantized": true, + "video_vae_quant_scheme": "fp8-f16-accum", + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_f16_accum.safetensors" +} +``` + +Activations use dynamic row-wise quantization with `scale = max(abs(x)) / qmax`. Reducing qmax +increases the scale and lowers the raw values accumulated in FP16, at the cost of fewer effective FP8 +levels. In the validated MiniMax-H3 workload, DiT produced non-finite FFN-out values with qmax 14 and +12, while qmax 7 completed every denoising step. Video VAE decoder remained finite and had the lowest +error with qmax 14. The current H3 policy therefore fixes activation qmax to 7 for DiT and 14 for +Video VAE, avoiding mismatches between runtime configuration and checkpoint conversion. + +The kernel is enabled only for DiT Q/K/V, attention output, and FFN projections, and for Video VAE +packed QKV, attention output, and FFN projections. It falls back to `fp8-sgl` when the extension is +unavailable or the device is not SM120. DiT tensor parallel also currently uses the `fp8-sgl` fallback. +Initialization logs report the effective scope or fallback reason. All other pipeline settings are +independent of this quantization mode. + +The kernel automatically tunes its CUTLASS tile and swizzle for each exact GEMM shape. The first use +of an unseen shape benchmarks the built-in candidates in C++ and keeps the winner in a process-local +C++ cache; later calls in the same process perform only a cache lookup. Enabling `warmup` moves the +tuning cost out of the first request when warmup covers the production shapes. A restarted process +tunes its shapes again and does not write to the user's cache directory. + ### T5 Model Quantization #### Supported Quantization Modes diff --git a/docs/ZH_CN/source/method_tutorials/quantization.md b/docs/ZH_CN/source/method_tutorials/quantization.md index 311367cc6..63a8b1660 100644 --- a/docs/ZH_CN/source/method_tutorials/quantization.md +++ b/docs/ZH_CN/source/method_tutorials/quantization.md @@ -13,6 +13,7 @@ LightX2V 支持对 DIT、T5 和 CLIP 模型进行量化推理，通过降低模 | `fp8-vllm` | FP8 通道对称 | FP8 通道动态对称 | [VLLM](https://github.com/vllm-project/vllm) | H100/H200/H800, RTX 40系等 | | `int8-vllm` | INT8 通道对称 | INT8 通道动态对称 | [VLLM](https://github.com/vllm-project/vllm) | A100/A800, RTX 30/40系等 | | `fp8-sgl` | FP8 通道对称 | FP8 通道动态对称 | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | H100/H200/H800, RTX 40系等 | +| `fp8-f16-accum` | FP8 通道对称 | FP8 行动态对称 | CUTLASS FP16 累加 | RTX 5090（SM120） | | `int8-sgl` | INT8 通道对称 | INT8 通道动态对称 | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | A100/A800, RTX 30/40系等 | | `fp8-q8f` | FP8 通道对称 | FP8 通道动态对称 | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40系, L40S等 | | `int8-q8f` | INT8 通道对称 | INT8 通道动态对称 | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40系, L40S等 | @@ -67,7 +68,7 @@ huggingface-cli download lightx2v/Encoders-Lightx2v \ #### 支持的量化模式 -DIT 量化模式（`dit_quant_scheme`）支持：`fp8-vllm`、`int8-vllm`、`fp8-sgl`、`int8-sgl`、`fp8-q8f`、`int8-q8f`、`int8-torchao`、`int4-g128-marlin`、`fp8-b128-deepgemm` +DIT 量化模式（`dit_quant_scheme`）支持：`fp8-vllm`、`int8-vllm`、`fp8-sgl`、`fp8-f16-accum`、`int8-sgl`、`fp8-q8f`、`int8-q8f`、`int8-torchao`、`int4-g128-marlin`、`fp8-b128-deepgemm` #### 配置示例 @@ -81,6 +82,67 @@ DIT 量化模式（`dit_quant_scheme`）支持：`fp8-vllm`、`int8-vllm`、`fp8 > 💡 **提示**：当运行脚本的 `model_path` 中只有一个 DIT 模型时，`dit_quantized_ckpt` 可以不用单独指定。 +#### MiniMax-H3 FP8 FP16 累加 + +RTX 5090 上的 MiniMax-H3 可以通过 `fp8-f16-accum` 使用 FP8 输入和 FP16 累加。权重需要用 +`h3-fp8-f16-accum` profile 转换；普通 `fp8-sgl` checkpoint 不兼容。DiT 和 Video VAE decoder +分别转换，profile 会独立选择使用 qmax 14 的投影层、保留其他层的标准 FP8 量化，并把策略写入 +safetensors metadata。 + +```bash +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/transformer \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_dit_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3 \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file + +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/vae \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_video_vae_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3_video_vae_decoder \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file +``` + +```json +{ + "dit_quantized": true, + "dit_quant_scheme": "fp8-f16-accum", + "dit_quantized_ckpt": "/path/to/minimax_h3_dit_fp8_f16_accum.safetensors", + "video_vae_quantized": true, + "video_vae_quant_scheme": "fp8-f16-accum", + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_f16_accum.safetensors" +} +``` + +激活按行动态量化，`scale = max(abs(x)) / qmax`。减小 qmax 会扩大 scale，从而降低 FP16 +累加器中的原始数值范围，但也会减少 FP8 有效量化级数。实测中 DiT 的 qmax 14 和 12 会在 FFN-out +产生非有限值，qmax 7 可完成全部去噪步骤；Video VAE decoder 在 qmax 14 下保持有限且误差最小。因此 +当前 H3 策略固定使用 DiT activation qmax 7 和 Video VAE activation qmax 14，避免运行配置与 +checkpoint 的转换策略错配。 + +DiT 仅对 Q/K/V、attention output 和 FFN projection 启用该内核，Video VAE 仅对 packed QKV、 +attention output 和 FFN projection 启用。扩展不可用或设备不是 SM120 时会回退到 `fp8-sgl`； +DiT tensor parallel 当前也回退到 `fp8-sgl`。初始化日志会打印实际启用范围或回退原因。 +其他 pipeline 配置与该量化模式相互独立。 + +该内核会按精确 GEMM shape 自动调优 CUTLASS tile 和 swizzle。首次遇到新 shape 时在 C++ 内遍历 +内置候选，winner 保存在当前进程的 C++ cache 中；同一进程的后续调用只执行 cache 查询。若 warmup +覆盖正式请求的 shape，首次调优开销会在请求前完成。进程重启后会重新调优一次，不写用户目录。 + ### T5 模型量化 #### 支持的量化模式 diff --git a/lightx2v/common/ops/mm/fp8_f16_accum.py b/lightx2v/common/ops/mm/fp8_f16_accum.py new file mode 100644 index 000000000..f04034f3b --- /dev/null +++ b/lightx2v/common/ops/mm/fp8_f16_accum.py @@ -0,0 +1,49 @@ +import math + +import torch + +from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton + +try: + from lightx2v_kernel.gemm import FP8_F16_ACCUM_MM_AVAILABLE, cutlass_scaled_fp8_mm_f16_accum +except ImportError: + FP8_F16_ACCUM_MM_AVAILABLE = False + cutlass_scaled_fp8_mm_f16_accum = None + + +def fp8_f16_accum_mm_unavailable_reason(): + if not FP8_F16_ACCUM_MM_AVAILABLE: + return "the lightx2v-kernel extension does not provide the FP8-F16 accumulation op" + if not torch.cuda.is_available(): + return "CUDA is unavailable" + capability = torch.cuda.get_device_capability() + if capability != (12, 0): + return f"SM120 is required, but the current CUDA capability is SM{capability[0]}{capability[1]}" + return None + + +def fp8_f16_accum_mm_available(): + return fp8_f16_accum_mm_unavailable_reason() is None + + +def validate_fp8_f16_accum_qmax(activation_qmax): + activation_qmax = float(activation_qmax) + fp8_max = torch.finfo(torch.float8_e4m3fn).max + if not math.isfinite(activation_qmax) or not 0 < activation_qmax <= fp8_max: + raise ValueError(f"FP8 activation qmax must be finite and in (0, {fp8_max}], got {activation_qmax}") + return activation_qmax + + +def fp8_f16_accum_linear(input_tensor, weight, weight_scale, bias, activation_qmax): + input_shape = input_tensor.shape + input_matrix = input_tensor.reshape(-1, input_shape[-1]) + quantized, activation_scale = fp8_quantize_range_triton(input_matrix, activation_qmax) + output = cutlass_scaled_fp8_mm_f16_accum( + quantized, + weight, + activation_scale, + weight_scale.float(), + input_tensor.dtype, + bias, + ) + return output.view(*input_shape[:-1], weight.shape[1]) diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 61613dd5a..d2a84facb 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -11,6 +11,11 @@ except ImportError: magi_register_custom_op = None +from lightx2v.common.ops.mm.fp8_f16_accum import ( + fp8_f16_accum_linear, + fp8_f16_accum_mm_available, + validate_fp8_f16_accum_qmax, +) from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_scaled_mm, sgl_fp8_scaled_mm_meta from lightx2v.common.ops.mm.triton_kernels import ( fp8_gemm_bias_triton, @@ -1977,6 +1982,33 @@ def apply(self, input_tensor): return output_tensor +@MM_WEIGHT_REGISTER("fp8-f16-accum") +class MMWeightWfp8channelAfp8channelF16Accum(MMWeightWfp8channelAfp8channeldynamicSgl): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.fp8_activation_qmax = None + + def enable_fp8_f16_accum(self, activation_qmax): + activation_qmax = validate_fp8_f16_accum_qmax(activation_qmax) + if fp8_f16_accum_mm_available(): + self.fp8_activation_qmax = activation_qmax + + def apply(self, input_tensor): + if self.fp8_activation_qmax is None: + return super().apply(input_tensor) + + output_tensor = fp8_f16_accum_linear( + input_tensor, + self.weight, + self.weight_scale, + self._get_actual_bias(), + self.fp8_activation_qmax, + ) + if self.has_lora_branch: + return output_tensor + self.apply_lora(input_tensor) + return output_tensor + + @MM_WEIGHT_REGISTER("int8-sgl") class MMWeightWint8channelAint8channeldynamicSglActVllm(MMWeightQuantTemplate): """ diff --git a/lightx2v/common/ops/mm/triton_kernels.py b/lightx2v/common/ops/mm/triton_kernels.py index 7e3fa5f7c..a846e9167 100755 --- a/lightx2v/common/ops/mm/triton_kernels.py +++ b/lightx2v/common/ops/mm/triton_kernels.py @@ -69,6 +69,24 @@ def fp8_quantize_triton(x): return quantized.view(x_shape_orig), scales.view(x_shape_orig[:-1]) +def fp8_quantize_range_triton(x, qmax): + x_shape = x.shape + x = x.reshape(-1, x_shape[-1]).contiguous() + quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn) + scales = torch.empty(x.shape[0], dtype=torch.float32, device=x.device) + block_size = next_power_of_2(x_shape[-1]) + fp8_quantize_kernel[(x.shape[0],)]( + x, + quantized, + scales, + x_shape[-1], + block_size, + FP8_MAX_VAL=qmax, + num_warps=8, + ) + return quantized.view(x_shape), scales.view(x_shape[:-1]) + + def upcast_if_fp8(a): if "fp8" in str(a): return torch.float16 diff --git a/lightx2v/models/input_encoders/hf/q_linear.py b/lightx2v/models/input_encoders/hf/q_linear.py index 2142b392a..ed972c17c 100755 --- a/lightx2v/models/input_encoders/hf/q_linear.py +++ b/lightx2v/models/input_encoders/hf/q_linear.py @@ -31,6 +31,11 @@ except ImportError: fp8_linear = None +from lightx2v.common.ops.mm.fp8_f16_accum import ( + fp8_f16_accum_linear, + fp8_f16_accum_mm_available, + validate_fp8_f16_accum_qmax, +) from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_scaled_mm from lightx2v.common.ops.mm.triton_kernels import fp8_gemm_bias_triton, fp8_gemm_triton, fp8_quantize_triton, int8_gemm_bias_triton, int8_gemm_triton, int8_quantize_triton from lightx2v_platform.ops.mm.mthreads_musa.fp8_scaled_mm import fp8_linear as musa_fp8_linear @@ -310,6 +315,28 @@ def maybe_cast(t): return self +class F16AccumQuantLinearFp8(SglQuantLinearFp8): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.fp8_activation_qmax = None + + def enable_fp8_f16_accum(self, activation_qmax): + activation_qmax = validate_fp8_f16_accum_qmax(activation_qmax) + if fp8_f16_accum_mm_available(): + self.fp8_activation_qmax = activation_qmax + + def forward(self, input_tensor): + if self.fp8_activation_qmax is None: + return super().forward(input_tensor) + return fp8_f16_accum_linear( + input_tensor, + self.weight.t(), + self.weight_scale, + self.bias, + self.fp8_activation_qmax, + ) + + class MusaQuantLinearFp8(nn.Module): """MUSA W8A8 FP8 linear with per-channel weights and per-token inputs.""" diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 35a8ceaf7..c666842f2 100644 --- a/lightx2v/models/networks/base_model.py +++ b/lightx2v/models/networks/base_model.py @@ -122,6 +122,7 @@ def _check_dit_quantized(self): "int8-q8f", "int8-convrot", "fp8-b128-deepgemm", + "fp8-f16-accum", "fp8-sgl", "int8-sgl", "int8-torchao", diff --git a/lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py b/lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py new file mode 100644 index 000000000..9757e601e --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py @@ -0,0 +1,36 @@ +from pathlib import Path + +from safetensors import safe_open + +FP8_F16_ACCUM_WEIGHT_QMAX = 14.0 +DIT_FP8_F16_ACCUM_ACTIVATION_QMAX = 7.0 +VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX = 14.0 +FP8_F16_ACCUM_QUANTIZATION_PROFILE = "h3-fp8-f16-accum" +FP8_F16_ACCUM_PROJECTION_SUFFIXES = ( + ".attn.to_q", + ".attn.to_k", + ".attn.to_v", + ".attn.to_out.0", + ".ff.net.0.proj", + ".ff.net.2", +) + + +def validate_fp8_f16_accum_checkpoint(checkpoint_path): + checkpoint_path = Path(checkpoint_path) + files = (checkpoint_path,) if checkpoint_path.is_file() else tuple(sorted(checkpoint_path.glob("*.safetensors"))) + if not files: + raise FileNotFoundError(f"No safetensors weights found in FP8 checkpoint: {checkpoint_path}") + + for filename in files: + with safe_open(filename, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + profile = metadata.get("quantization_profile") + if profile != FP8_F16_ACCUM_QUANTIZATION_PROFILE: + raise ValueError(f"{filename} requires quantization profile {FP8_F16_ACCUM_QUANTIZATION_PROFILE!r}, got {profile!r}") + try: + weight_qmax = float(metadata.get("weight_qmax")) + except (TypeError, ValueError): + weight_qmax = None + if weight_qmax != FP8_F16_ACCUM_WEIGHT_QMAX: + raise ValueError(f"{filename} requires weight_qmax={FP8_F16_ACCUM_WEIGHT_QMAX}, got {metadata.get('weight_qmax')!r}") diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 19ca82476..250044b14 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -7,7 +7,13 @@ from loguru import logger from safetensors import safe_open +from lightx2v.common.ops.mm.fp8_f16_accum import fp8_f16_accum_mm_unavailable_reason from lightx2v.models.networks.base_model import BaseTransformerModel +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + DIT_FP8_F16_ACCUM_ACTIVATION_QMAX, + FP8_F16_ACCUM_WEIGHT_QMAX, + validate_fp8_f16_accum_checkpoint, +) from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3SequenceParallelState from lightx2v.models.networks.minimax_h3.infer.offload import MiniMaxH3OffloadTransformerInfer from lightx2v.models.networks.minimax_h3.infer.post_infer import MiniMaxH3PostInfer @@ -23,6 +29,7 @@ H3_CHANNEL_QUANT_SCHEMES = { "fp8-q8f", + "fp8-f16-accum", "fp8-musa", "fp8-sgl", "fp8-torchao", @@ -64,6 +71,19 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 raise NotImplementedError(f"MiniMax-H3 quantized inference requires a per-output-channel FP8/INT8 scheme; got {quant_scheme!r}. Supported schemes: {sorted(H3_CHANNEL_QUANT_SCHEMES)}") if not config.get("dit_quantized_ckpt"): raise ValueError("MiniMax-H3 quantized inference requires dit_quantized_ckpt") + if quant_scheme == "fp8-f16-accum": + validate_fp8_f16_accum_checkpoint(config["dit_quantized_ckpt"]) + fallback_reason = fp8_f16_accum_mm_unavailable_reason() + if config.get("tensor_parallel", False): + logger.info("MiniMax-H3 DiT FP8-F16 accumulation falls back to FP8-SGL under tensor parallel") + elif fallback_reason is not None: + logger.warning("MiniMax-H3 DiT FP8-F16 accumulation requested but {}; falling back to FP8-SGL", fallback_reason) + else: + logger.info( + "MiniMax-H3 DiT FP8-F16 accumulation enabled for Q/K/V, attention output, and FFN projections (weight qmax={}, activation qmax={})", + FP8_F16_ACCUM_WEIGHT_QMAX, + DIT_FP8_F16_ACCUM_ACTIVATION_QMAX, + ) elif config.get("dit_quant_scheme", "Default") != "Default": raise ValueError("MiniMax-H3 dit_quant_scheme requires a dit_quantized_ckpt") if config.get("cpu_offload", False) and config.get("offload_granularity", "model") not in {"model", "block"}: diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 9634eb595..e654cfa6b 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -2,19 +2,24 @@ import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + DIT_FP8_F16_ACCUM_ACTIVATION_QMAX, + FP8_F16_ACCUM_PROJECTION_SUFFIXES, +) from lightx2v.models.networks.minimax_h3.infer.triton_ops import MiniMaxH3TritonRope # noqa: F401 from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): lora_prefix = "transformer_blocks" + quant_scheme = config.get("dit_quant_scheme", "Default") if config.get("tensor_parallel", False) and tp_split is not None: tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") tp_mm_type = config.get("tp_mm_type", "TensorParallel") return MM_WEIGHT_REGISTER[tp_mm_type]( weight_name=f"{name}.weight", bias_name=f"{name}.bias" if bias else None, - mm_type=config.get("dit_quant_scheme", "Default"), + mm_type=quant_scheme, tp_group=tp_group, tp_rank=dist.get_rank(tp_group), tp_size=dist.get_world_size(tp_group), @@ -23,12 +28,16 @@ def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): create_cuda_buffer=create_cuda_buffer, lora_prefix=lora_prefix, ) - return MM_WEIGHT_REGISTER[config.get("dit_quant_scheme", "Default")]( + + linear = MM_WEIGHT_REGISTER[quant_scheme]( f"{name}.weight", f"{name}.bias" if bias else None, create_cuda_buffer=create_cuda_buffer, lora_prefix=lora_prefix, ) + if quant_scheme == "fp8-f16-accum" and name.endswith(FP8_F16_ACCUM_PROJECTION_SUFFIXES): + linear.enable_fp8_f16_accum(DIT_FP8_F16_ACCUM_ACTIVATION_QMAX) + return linear def _rms(config, name, eps, create_cuda_buffer=False): diff --git a/lightx2v/models/runners/default_runner.py b/lightx2v/models/runners/default_runner.py index 2d5ff4c0f..d13b73d40 100755 --- a/lightx2v/models/runners/default_runner.py +++ b/lightx2v/models/runners/default_runner.py @@ -252,6 +252,7 @@ def init_modules(self): self.run_input_encoder = self._run_input_encoder_local_i2av elif self.config["task"] == "sr": self.run_input_encoder = self._run_input_encoder_local_sr + self.config.lock() # lock config to avoid modification def set_init_device(self): diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 318348b2d..7658e5ca4 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -43,6 +43,12 @@ import torch.nn.functional as F from loguru import logger +from lightx2v.common.ops.mm.fp8_f16_accum import fp8_f16_accum_mm_unavailable_reason +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + FP8_F16_ACCUM_WEIGHT_QMAX, + VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX, + validate_fp8_f16_accum_checkpoint, +) from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( SafetensorsSubsetReport, load_safetensors_subset, @@ -552,7 +558,7 @@ def __init__( attn_type: str = "torch_sdpa", ) -> None: super().__init__() - if quant_scheme not in {None, "fp8-musa", "fp8-sgl"}: + if quant_scheme not in {None, "fp8-f16-accum", "fp8-musa", "fp8-sgl"}: raise NotImplementedError(f"Unsupported MiniMax-H3 video VAE quantization scheme: {quant_scheme!r}") if attn_type not in {"torch_sdpa", "sage_attn2"}: raise ValueError(f"Unsupported MiniMax-H3 video VAE attention type: {attn_type!r}; expected torch_sdpa or sage_attn2") @@ -649,8 +655,18 @@ def _pack_decoder_fp8_qkv(self) -> None: for block in self.decoder.transformer_blocks: block.attn._pack_fp8_qkv() + def _configure_fp8_f16_accum_linears(self) -> None: + # Packed QKV, attention output, and FFN use the validated H3 shapes. + for block in self.decoder.transformer_blocks: + block.attn.to_qkv.enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + block.attn.to_out[0].enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + block.ff.net[0].proj.enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + block.ff.net[2].enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + def _make_fp8_linear(self, linear: nn.Linear) -> nn.Module: - if self.quant_scheme == "fp8-musa": + if self.quant_scheme == "fp8-f16-accum": + from lightx2v.models.input_encoders.hf.q_linear import F16AccumQuantLinearFp8 as linear_cls + elif self.quant_scheme == "fp8-musa": from lightx2v.models.input_encoders.hf.q_linear import MusaQuantLinearFp8 as linear_cls elif self.quant_scheme == "fp8-sgl": from lightx2v.models.input_encoders.hf.q_linear import SglQuantLinearFp8 as linear_cls @@ -703,6 +719,17 @@ def from_pretrained( if (checkpoint_path is None) != (quant_scheme is None): raise ValueError("MiniMax-H3 video VAE checkpoint_path and quant_scheme must be configured together") weight_path = checkpoint_path if checkpoint_path is not None else vae_dir + if quant_scheme == "fp8-f16-accum": + validate_fp8_f16_accum_checkpoint(weight_path) + fallback_reason = fp8_f16_accum_mm_unavailable_reason() + if fallback_reason is None: + logger.info( + "MiniMax-H3 Video VAE FP8-F16 accumulation enabled for packed QKV, attention output, and FFN projections (weight qmax={}, activation qmax={})", + FP8_F16_ACCUM_WEIGHT_QMAX, + VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX, + ) + else: + logger.warning("MiniMax-H3 Video VAE FP8-F16 accumulation requested but {}; falling back to FP8-SGL", fallback_reason) with (vae_dir / "config.json").open("r", encoding="utf-8") as handle: config = json.load(handle) @@ -723,6 +750,8 @@ def from_pretrained( if quant_scheme is not None: # Pack only after loading the checkpoint's original Q/K/V keys. model._pack_decoder_fp8_qkv() + if quant_scheme == "fp8-f16-accum": + model._configure_fp8_f16_accum_linears() model._prepare_inference_dtypes() model.eval().requires_grad_(False) if not cpu_offload: diff --git a/lightx2v_kernel/CMakeLists.txt b/lightx2v_kernel/CMakeLists.txt index 369f4da9c..990a1b53a 100644 --- a/lightx2v_kernel/CMakeLists.txt +++ b/lightx2v_kernel/CMakeLists.txt @@ -82,6 +82,7 @@ list(APPEND LIGHTX2V_KERNEL_CUDA_FLAGS set(SOURCES + "csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu" "csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu" "csrc/gemm/nvfp4_quant_kernels_sm120.cu" "csrc/gemm/mxfp4_quant_kernels_sm120.cu" diff --git a/lightx2v_kernel/csrc/common_extension.cc b/lightx2v_kernel/csrc/common_extension.cc index 3bfdc746f..cdded046a 100644 --- a/lightx2v_kernel/csrc/common_extension.cc +++ b/lightx2v_kernel/csrc/common_extension.cc @@ -6,6 +6,23 @@ TORCH_LIBRARY_FRAGMENT(lightx2v_kernel, m) { + m.def( + "cutlass_scaled_fp8_mm_f16_accum_sm120(Tensor activation, Tensor weight, Tensor activation_scale, " + "Tensor weight_scale, ScalarType out_dtype, Tensor? bias=None) -> Tensor"); + m.impl( + "cutlass_scaled_fp8_mm_f16_accum_sm120", + torch::kCUDA, + &cutlass_scaled_fp8_mm_f16_accum_sm120); + + m.def( + "cutlass_scaled_fp8_mm_f16_accum_with_config_sm120(Tensor activation, Tensor weight, " + "Tensor activation_scale, Tensor weight_scale, ScalarType out_dtype, Tensor? bias, int config_id) -> Tensor"); + m.impl( + "cutlass_scaled_fp8_mm_f16_accum_with_config_sm120", + torch::kCUDA, + &cutlass_scaled_fp8_mm_f16_accum_with_config_sm120); + + m.def( "cutlass_scaled_nvfp4_mm_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, Tensor " "alpha, Tensor? bias) -> ()"); diff --git a/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu b/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu new file mode 100644 index 000000000..790f09b74 --- /dev/null +++ b/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu @@ -0,0 +1,673 @@ +// SM120 FP8 GEMM with FP16 accumulation. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cutlass/epilogue/fusion/operations.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" + +using namespace cute; + +namespace { + +template < + typename TileShape_, + typename ElementD_ = cutlass::bfloat16_t, + bool FuseBias_ = false> +struct GemmDefinition { + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementD = ElementD_; + using ElementC = void; + using ElementAccumulator = cutlass::half_t; + using TileShape = TileShape_; + static constexpr bool FuseBias = FuseBias_; + using ClusterShape = Shape<_1, _1, _1>; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + static constexpr int AlignmentAB = 16; + static constexpr int AlignmentD = 16 / sizeof(ElementD); + + using Accum = cutlass::epilogue::fusion::Sm90AccFetch; + using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< + 0, + TileShape, + float, + float, + Stride, Int<0>, Int<0>>>; + using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast< + 0, + TileShape, + float, + float, + Stride, Int<1>, Int<0>>>; + using Multiply = cutlass::epilogue::fusion::Sm90Compute< + cutlass::multiplies, + float, + float, + cutlass::FloatRoundStyle::round_to_nearest>; + using MultiplyOutput = cutlass::epilogue::fusion::Sm90Compute< + cutlass::multiplies, + ElementD, + float, + cutlass::FloatRoundStyle::round_to_nearest>; + using AddBias = cutlass::epilogue::fusion::Sm90Compute< + cutlass::plus, + ElementD, + float, + cutlass::FloatRoundStyle::round_to_nearest>; + using Bias = cutlass::epilogue::fusion::Sm90RowBroadcast< + 0, + TileShape, + ElementD, + float, + Stride, Int<1>, Int<0>>, + AlignmentD>; + using ScaleBAccum = + cutlass::epilogue::fusion::Sm90EVT; + using ScaledEVT = cutlass::epilogue::fusion::Sm90EVT< + MultiplyOutput, + ScaleA, + ScaleBAccum>; + using OutputEVT = cutlass::epilogue::fusion::Sm90EVT< + AddBias, + ScaledEVT, + Bias>; + using EpilogueEVT = std::conditional_t; + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, + cutlass::arch::OpClassTensorOp, + TileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + float, + ElementC, + LayoutC, + AlignmentD, + ElementD, + LayoutD, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + EpilogueEVT>::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, + cutlass::arch::OpClassTensorOp, + ElementA, + LayoutA, + AlignmentAB, + ElementB, + LayoutB, + AlignmentAB, + ElementAccumulator, + TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static typename EpilogueEVT::Arguments prepare_epilogue( + float* activation_scale, + float* weight_scale, + ElementD const* bias) { + typename ScaleA::Arguments activation_arguments{activation_scale}; + typename ScaleB::Arguments weight_arguments{weight_scale}; + typename ScaleBAccum::Arguments scaled_accumulator{ + weight_arguments, + {}, + {}, + }; + typename ScaledEVT::Arguments scaled_output{ + activation_arguments, + scaled_accumulator, + {}, + }; + if constexpr (FuseBias) { + typename Bias::Arguments bias_arguments{bias}; + return typename OutputEVT::Arguments{ + scaled_output, + bias_arguments, + {}, + }; + } else { + return scaled_output; + } + } +}; +using NarrowGemm = GemmDefinition>; +using WideGemm = GemmDefinition>; +using NarrowGemmFp16 = GemmDefinition, cutlass::half_t>; +using WideGemmFp16 = GemmDefinition, cutlass::half_t>; + +using NarrowGemmWithBias = + GemmDefinition, cutlass::bfloat16_t, true>; +using WideGemmWithBias = + GemmDefinition, cutlass::bfloat16_t, true>; +using NarrowGemmFp16WithBias = + GemmDefinition, cutlass::half_t, true>; +using WideGemmFp16WithBias = + GemmDefinition, cutlass::half_t, true>; + +struct KernelConfig { + bool wide_tile; + int swizzle; +}; + +constexpr std::array kKernelConfigs = {{ + {false, 1}, + {false, 2}, + {false, 4}, + {false, 8}, + {true, 1}, + {true, 2}, + {true, 4}, + {true, 8}, +}}; + +KernelConfig const& kernel_config(int64_t config_id) { + TORCH_CHECK( + config_id >= 0 && + config_id < static_cast(kKernelConfigs.size()), + "FP8-F16 GEMM config_id must be in [0, ", + kKernelConfigs.size(), + "), got ", + config_id); + return kKernelConfigs[config_id]; +} + +struct AutotuneKey { + int device_index; + int32_t m; + int32_t n; + int32_t k; + torch::ScalarType output_dtype; + bool has_bias; + + bool operator==(AutotuneKey const& other) const { + return device_index == other.device_index && m == other.m && + n == other.n && k == other.k && + output_dtype == other.output_dtype && has_bias == other.has_bias; + } +}; + +struct AutotuneKeyHash { + size_t operator()(AutotuneKey const& key) const { + size_t value = std::hash{}(key.device_index); + value = value * 31 + std::hash{}(key.m); + value = value * 31 + std::hash{}(key.n); + value = value * 31 + std::hash{}(key.k); + value = value * 31 + std::hash{}(static_cast(key.output_dtype)); + return value * 31 + std::hash{}(key.has_bias); + } +}; + +using AutotuneCache = + std::unordered_map; + +AutotuneCache& autotune_cache() { + static AutotuneCache cache; + return cache; +} + +std::shared_mutex& autotune_cache_mutex() { + static std::shared_mutex mutex; + return mutex; +} + +std::mutex& autotune_measurement_mutex() { + static std::mutex mutex; + return mutex; +} + +std::optional cached_config_id(AutotuneKey const& key) { + std::shared_lock lock(autotune_cache_mutex()); + auto entry = autotune_cache().find(key); + if (entry == autotune_cache().end()) { + return std::nullopt; + } + return entry->second; +} + +void cache_config(AutotuneKey const& key, int64_t config_id) { + std::unique_lock lock(autotune_cache_mutex()); + autotune_cache()[key] = config_id; +} + +template +void launch( + torch::Tensor output, + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + int swizzle) { + using Gemm = typename Definition::Gemm; + using GemmKernel = typename Definition::GemmKernel; + using ElementA = typename Definition::ElementA; + using ElementB = typename Definition::ElementB; + using ElementD = typename Definition::ElementD; + using ElementC = typename Definition::ElementC; + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideC; + using StrideD = typename GemmKernel::StrideD; + + int32_t m = activation.size(0); + int32_t k = activation.size(1); + int32_t n = weight.size(1); + StrideA stride_a = make_stride( + int64_t(activation.stride(0)), Int<1>{}, int64_t(0)); + StrideB stride_b = make_stride( + int64_t(weight.stride(1)), Int<1>{}, int64_t(0)); + auto stride_c = cutlass::make_cute_packed_stride( + StrideC{}, make_shape(m, n, 1)); + auto stride_d = cutlass::make_cute_packed_stride( + StrideD{}, make_shape(m, n, 1)); + + typename GemmKernel::MainloopArguments mainloop{ + reinterpret_cast(activation.data_ptr()), + stride_a, + reinterpret_cast(weight.data_ptr()), + stride_b, + }; + typename GemmKernel::EpilogueArguments epilogue{ + Definition::prepare_epilogue( + activation_scale.data_ptr(), + weight_scale.data_ptr(), + bias ? reinterpret_cast(bias->data_ptr()) : nullptr), + nullptr, + stride_c, + reinterpret_cast(output.data_ptr()), + stride_d, + }; + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {m, n, k, 1}, + mainloop, + epilogue, + }; + arguments.scheduler.max_swizzle_size = swizzle; + + Gemm gemm; + auto status = gemm.can_implement(arguments); + TORCH_CHECK( + status == cutlass::Status::kSuccess, + "CUTLASS cannot implement this shape"); + TORCH_CHECK( + Gemm::get_workspace_size(arguments) == 0, + "Unexpected CUTLASS scheduler workspace"); + auto stream = at::cuda::getCurrentCUDAStream(activation.device().index()); + status = gemm.run(arguments, nullptr, stream); + TORCH_CHECK( + status == cutlass::Status::kSuccess, + "CUTLASS kernel launch failed"); +} + +void validate( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + torch::ScalarType output_dtype) { + TORCH_CHECK( + activation.is_cuda() && weight.is_cuda() && + activation_scale.is_cuda() && weight_scale.is_cuda(), + "inputs and scales must be CUDA tensors"); + TORCH_CHECK( + activation.device() == weight.device() && + activation.device() == activation_scale.device() && + activation.device() == weight_scale.device(), + "inputs and scales must be on the same CUDA device"); + TORCH_CHECK( + activation.scalar_type() == torch::kFloat8_e4m3fn && + weight.scalar_type() == torch::kFloat8_e4m3fn, + "activation and weight must be float8_e4m3fn"); + TORCH_CHECK( + activation.dim() == 2 && weight.dim() == 2, + "activation and weight must be matrices"); + TORCH_CHECK( + activation.stride(1) == 1 && + activation.stride(0) == activation.size(1), + "activation must be contiguous"); + TORCH_CHECK( + weight.stride(0) == 1 && weight.stride(1) == weight.size(0), + "weight must be a transposed contiguous matrix"); + TORCH_CHECK( + activation.size(1) == weight.size(0), + "K dimensions must match"); + TORCH_CHECK( + activation.size(0) > 0 && + activation.size(0) <= std::numeric_limits::max() && + weight.size(1) > 0 && + weight.size(1) <= std::numeric_limits::max() && + activation.size(1) > 0 && + activation.size(1) <= std::numeric_limits::max(), + "M, N and K must be positive int32 values"); + TORCH_CHECK( + activation_scale.scalar_type() == torch::kFloat32 && + weight_scale.scalar_type() == torch::kFloat32, + "scales must be float32"); + TORCH_CHECK( + activation_scale.is_contiguous() && weight_scale.is_contiguous(), + "scales must be contiguous"); + TORCH_CHECK( + activation_scale.numel() == activation.size(0) && + weight_scale.numel() == weight.size(1), + "scale sizes must match the activation rows and weight columns"); + if (bias) { + TORCH_CHECK( + bias->is_cuda() && bias->device() == activation.device(), + "bias must be on the same CUDA device as the inputs"); + TORCH_CHECK( + bias->scalar_type() == output_dtype, + "bias dtype must match the output dtype"); + TORCH_CHECK( + bias->is_contiguous() && bias->dim() == 1 && + bias->size(0) == weight.size(1), + "bias must be a contiguous vector matching the output columns"); + } +} + +template +void launch_config( + torch::Tensor output, + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + int64_t config_id) { + KernelConfig const& config = kernel_config(config_id); + if (config.wide_tile) { + launch( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config.swizzle); + } else { + launch( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config.swizzle); + } +} + +template < + typename NarrowDefinition, + typename NarrowDefinitionWithBias, + typename WideDefinition, + typename WideDefinitionWithBias> +int64_t tune_config( + AutotuneKey const& key, + torch::Tensor output, + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias) { + std::lock_guard measurement_lock(autotune_measurement_mutex()); + if (auto cached = cached_config_id(key)) { + return *cached; + } + + auto launch_candidate = [&](int64_t config_id) { + if (bias) { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config_id); + } else { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config_id); + } + }; + + constexpr int kWarmups = 2; + constexpr int kTrials = 3; + constexpr int kRepeats = 5; + constexpr int kConfigCount = static_cast(kKernelConfigs.size()); + std::array, kConfigCount> timings{}; + + for (int warmup = 0; warmup < kWarmups; ++warmup) { + for (int index = 0; index < kConfigCount; ++index) { + launch_candidate((index + warmup) % kConfigCount); + } + } + + auto stream = at::cuda::getCurrentCUDAStream(activation.device().index()); + int initial_offset = + static_cast((key.m + key.n + key.k) % kConfigCount); + for (int trial = 0; trial < kTrials; ++trial) { + int offset = (initial_offset + trial * 3) % kConfigCount; + for (int index = 0; index < kConfigCount; ++index) { + int config_id = (index + offset) % kConfigCount; + c10::cuda::CUDAEvent start(cudaEventDefault); + c10::cuda::CUDAEvent end(cudaEventDefault); + start.record(stream); + for (int repeat = 0; repeat < kRepeats; ++repeat) { + launch_candidate(config_id); + } + end.record(stream); + end.synchronize(); + timings[config_id][trial] = + start.elapsed_time(end) / static_cast(kRepeats); + } + } + + int64_t best_config_id = 0; + float best_time = std::numeric_limits::max(); + for (int config_id = 0; config_id < kConfigCount; ++config_id) { + auto samples = timings[config_id]; + std::sort(samples.begin(), samples.end()); + if (samples[kTrials / 2] < best_time) { + best_time = samples[kTrials / 2]; + best_config_id = config_id; + } + } + cache_config(key, best_config_id); + return best_config_id; +} + +template < + typename NarrowDefinition, + typename NarrowDefinitionWithBias, + typename WideDefinition, + typename WideDefinitionWithBias> +torch::Tensor run( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + torch::ScalarType output_dtype, + c10::optional config_id = c10::nullopt) { + validate( + activation, + weight, + activation_scale, + weight_scale, + bias, + output_dtype); + c10::cuda::CUDAGuard guard(activation.device()); + int32_t m = activation.size(0); + int32_t k = activation.size(1); + int32_t n = weight.size(1); + auto output = torch::empty( + {m, n}, + activation.options().dtype(output_dtype)); + + if (!config_id) { + AutotuneKey key{ + activation.device().index(), + m, + n, + k, + output_dtype, + bias.has_value(), + }; + if (auto cached = cached_config_id(key)) { + config_id = *cached; + } else { + config_id = tune_config< + NarrowDefinition, + NarrowDefinitionWithBias, + WideDefinition, + WideDefinitionWithBias>( + key, + output, + activation, + weight, + activation_scale, + weight_scale, + bias); + } + } + + if (bias) { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + *config_id); + } else { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + *config_id); + } + return output; +} + +torch::Tensor run_with_dtype( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias, + c10::optional config_id = c10::nullopt) { + if (out_dtype == torch::kBFloat16) { + return run< + NarrowGemm, + NarrowGemmWithBias, + WideGemm, + WideGemmWithBias>( + activation, + weight, + activation_scale, + weight_scale, + bias, + out_dtype, + config_id); + } + TORCH_CHECK( + out_dtype == torch::kFloat16, + "output dtype must be bfloat16 or float16"); + return run< + NarrowGemmFp16, + NarrowGemmFp16WithBias, + WideGemmFp16, + WideGemmFp16WithBias>( + activation, + weight, + activation_scale, + weight_scale, + bias, + out_dtype, + config_id); +} + +} // namespace + +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias) { + return run_with_dtype( + activation, + weight, + activation_scale, + weight_scale, + out_dtype, + bias); +} + +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_with_config_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias, + int64_t config_id) { + return run_with_dtype( + activation, + weight, + activation_scale, + weight_scale, + out_dtype, + bias, + config_id); +} diff --git a/lightx2v_kernel/include/lightx2v_kernel_ops.h b/lightx2v_kernel/include/lightx2v_kernel_ops.h index 04b380596..a8eaaafdf 100644 --- a/lightx2v_kernel/include/lightx2v_kernel_ops.h +++ b/lightx2v_kernel/include/lightx2v_kernel_ops.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include @@ -42,6 +43,24 @@ limitations under the License. /* * From csrc/gemm */ +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias = c10::nullopt); + +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_with_config_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias, + int64_t config_id); + + void scaled_nvfp4_quant_sm120( torch::Tensor& output, torch::Tensor const& input, torch::Tensor& output_sf, torch::Tensor const& input_sf); diff --git a/lightx2v_kernel/python/lightx2v_kernel/gemm.py b/lightx2v_kernel/python/lightx2v_kernel/gemm.py index 8ae4b956e..9ea4f1deb 100644 --- a/lightx2v_kernel/python/lightx2v_kernel/gemm.py +++ b/lightx2v_kernel/python/lightx2v_kernel/gemm.py @@ -1,6 +1,24 @@ import torch +def _fp8_f16_accum_meta(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): + del scales_a, scales_b, bias + return torch.empty((mat_a.shape[0], mat_b.shape[1]), dtype=out_dtype, device=mat_a.device) + + +FP8_F16_ACCUM_MM_AVAILABLE = hasattr(torch.ops.lightx2v_kernel, "cutlass_scaled_fp8_mm_f16_accum_sm120") +if FP8_F16_ACCUM_MM_AVAILABLE: + _fp8_f16_accum_op = torch.ops.lightx2v_kernel.cutlass_scaled_fp8_mm_f16_accum_sm120.default + if not _fp8_f16_accum_op.has_kernel_for_dispatch_key("Meta"): + torch.library.register_fake(_fp8_f16_accum_op, _fp8_f16_accum_meta) + + +def cutlass_scaled_fp8_mm_f16_accum(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): + if not FP8_F16_ACCUM_MM_AVAILABLE: + raise ImportError("lightx2v-kernel was built without the SM120 FP8 GEMM with FP16 accumulation") + return _fp8_f16_accum_op(mat_a, mat_b, scales_a, scales_b, out_dtype, bias) + + def cutlass_scaled_nvfp4_mm(mat_a, mat_b, scales_a, scales_b, alpha, bias=None): m, n = mat_a.shape[0], mat_b.shape[0] out = torch.empty((m, n), dtype=torch.bfloat16, device=mat_a.device) diff --git a/tools/convert/converter.py b/tools/convert/converter.py index 6fae91829..d728060e1 100755 --- a/tools/convert/converter.py +++ b/tools/convert/converter.py @@ -24,6 +24,10 @@ if quant_path not in sys.path: sys.path.insert(0, quant_path) +from h3_fp8_f16_accum import ( # noqa: E402 + FP8_F16_ACCUM_QUANTIZATION_PROFILE, + create_h3_fp8_f16_accum_quantization, +) from quant import * # noqa: E402 from lightx2v.utils.lora_loader import LoRALoader # noqa: E402 @@ -310,9 +314,9 @@ def get_key_mapping_rules(direction, model_type): return [rule["backward"] for rule in unified_rules] else: raise ValueError(f"Invalid direction: {direction}") - elif model_type == "h3": - # MiniMax-H3 checkpoints under the Diffusers ``transformer`` or - # ``transformer_ref`` directory already use LightX2V's runtime keys. + elif model_type in {"h3", "h3_video_vae_decoder"}: + # MiniMax-H3 transformer and Video VAE decoder checkpoints already use + # LightX2V's runtime keys. return [] else: raise ValueError(f"Unsupported model type: {model_type}") @@ -331,6 +335,7 @@ def quantize_model( preserve_non_quant_dtype=False, comfyui_mode=False, comfyui_keys=[], + quantization_policy=None, ): """ Quantize model weights in-place @@ -407,7 +412,10 @@ def quantize_model( # Quantize tensor and store results quantizer = CONVERT_WEIGHT_REGISTER[linear_type](tensor) - w_q, scales, extra = quantizer.weight_quant_func(tensor, comfyui_mode) + if quantization_policy is None: + w_q, scales, extra = quantizer.weight_quant_func(tensor, comfyui_mode) + else: + w_q, scales, extra = quantization_policy.quantize_weight(key, tensor, quantizer.weight_quant_func) weight_global_scale = extra.get("weight_global_scale", None) # For nvfp4 convrot_groupsize = extra.get("convrot_groupsize", None) @@ -448,6 +456,8 @@ def quantize_model( logger.info(f"Total final model size: {total_final_size_mb:.2f} MB") logger.info(f"Size reduction in quantized tensors: {size_reduction_mb:.2f} MB ({size_reduction_mb / original_size_mb * 100:.1f}%)") + if quantization_policy is not None: + quantization_policy.validate() if comfyui_mode: weights["scaled_fp8"] = torch.zeros(2, dtype=torch.float8_e4m3fn) @@ -827,6 +837,7 @@ def convert_key(key): preserve_non_quant_dtype=getattr(args, "preserve_non_quant_dtype", False), comfyui_mode=args.comfyui_mode, comfyui_keys=args.comfyui_keys, + quantization_policy=args.quantization_policy, ) os.makedirs(args.output, exist_ok=True) @@ -852,7 +863,8 @@ def convert_key(key): logger.warning("Consider using --save_by_block or default chunked saving for better memory efficiency.") # Save the entire model as a single file - st.save_file(converted_weights, output_path) + metadata = args.quantization_policy.metadata if args.quantization_policy is not None else None + st.save_file(converted_weights, output_path, metadata=metadata) logger.info(f"Model saved successfully to: {output_path} ({total_size_gb:.2f}GB)") except MemoryError: @@ -975,7 +987,7 @@ def main(): parser.add_argument( "-t", "--model_type", - choices=["wan_dit", "h3", "h3_text_encoder", "hunyuan_dit", "wan_t5", "wan_clip", "wan_animate_dit", "qwen_image_dit", "qwen25vl_llm", "z_image_dit", "self_forcing"], + choices=["wan_dit", "h3", "h3_video_vae_decoder", "h3_text_encoder", "hunyuan_dit", "wan_t5", "wan_clip", "wan_animate_dit", "qwen_image_dit", "qwen25vl_llm", "z_image_dit", "self_forcing"], default="wan_dit", help="Model type", ) @@ -1011,6 +1023,7 @@ def main(): parser.add_argument("--comfyui_mode", action="store_true") parser.add_argument("--full_quantized", action="store_true") parser.add_argument("--quantized", action="store_true") + parser.add_argument("--quantization_profile", choices=[FP8_F16_ACCUM_QUANTIZATION_PROFILE]) parser.add_argument("--bits", type=int, default=8, choices=[8], help="Quantization bit width") parser.add_argument( "--device", @@ -1113,6 +1126,12 @@ def _parse_csv_override(v: str | None) -> list[str] | None: # every tensor outside the quantized block linears. "preserve_non_quant_dtype": True, }, + "h3_video_vae_decoder": { + "key_idx": 1, + "target_keys": ["transformer_blocks", "proj_out"], + "ignore_key": None, + "preserve_non_quant_dtype": True, + }, "self_forcing": { "key_idx": 3, "target_keys": ["self_attn", "cross_attn", "ffn"], @@ -1189,6 +1208,15 @@ def _parse_csv_override(v: str | None) -> list[str] | None: else: args.ignore_quant_keys = None + args.quantization_policy = None + if args.quantization_profile is not None: + if not args.quantized or args.linear_type != "fp8" or not args.single_file or args.output_ext != ".safetensors" or args.comfyui_mode: + parser.error("H3 FP8-F16 accumulation conversion requires --quantized --linear_type fp8 --output_ext .safetensors --single_file without --comfyui_mode") + try: + args.quantization_policy = create_h3_fp8_f16_accum_quantization(args.quantization_profile, args.model_type) + except ValueError as profile_error: + parser.error(str(profile_error)) + if os.path.isfile(args.output): raise ValueError("Output path must be a directory, not a file") diff --git a/tools/convert/quant/h3_fp8_f16_accum.py b/tools/convert/quant/h3_fp8_f16_accum.py new file mode 100644 index 000000000..70e009d05 --- /dev/null +++ b/tools/convert/quant/h3_fp8_f16_accum.py @@ -0,0 +1,66 @@ +"""MiniMax-H3 checkpoint policy for FP8 GEMM with FP16 accumulation.""" + +import torch + +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + FP8_F16_ACCUM_PROJECTION_SUFFIXES, + FP8_F16_ACCUM_QUANTIZATION_PROFILE, + FP8_F16_ACCUM_WEIGHT_QMAX, +) + +# DiT has 50 x (six main projections + AdaLN); VAE has 36 x six projections + proj_out. +_EXPECTED_QUANTIZED_COUNTS = { + "h3": 350, + "h3_video_vae_decoder": 217, +} +_EXPECTED_QMAX14_COUNTS = { + "h3": 300, + "h3_video_vae_decoder": 216, +} + + +class H3FP8F16AccumQuantization: + """Assign qmax14 only to H3 projections using FP16 accumulation.""" + + def __init__(self, model_type): + if model_type not in _EXPECTED_QUANTIZED_COUNTS: + raise ValueError(f"{FP8_F16_ACCUM_QUANTIZATION_PROFILE} does not support model_type={model_type!r}") + self.model_type = model_type + self.quantized_count = 0 + self.qmax14_count = 0 + + @property + def metadata(self): + return { + "format": "pt", + "quantization_profile": FP8_F16_ACCUM_QUANTIZATION_PROFILE, + "weight_qmax": str(FP8_F16_ACCUM_WEIGHT_QMAX), + } + + def quantize_weight(self, name, weight, default_quantize): + projection_name = name.removesuffix(".weight") + uses_reduced_range = projection_name.endswith(FP8_F16_ACCUM_PROJECTION_SUFFIXES) + self.quantized_count += 1 + if not uses_reduced_range: + return default_quantize(weight) + + values = weight.float() + scales = values.abs().amax(dim=1, keepdim=True).clamp_min_(1e-8).div_(FP8_F16_ACCUM_WEIGHT_QMAX) + values.div_(scales).clamp_(-FP8_F16_ACCUM_WEIGHT_QMAX, FP8_F16_ACCUM_WEIGHT_QMAX) + + self.qmax14_count += 1 + return values.to(torch.float8_e4m3fn), scales, {} + + def validate(self): + expected_quantized = _EXPECTED_QUANTIZED_COUNTS[self.model_type] + expected_qmax14 = _EXPECTED_QMAX14_COUNTS[self.model_type] + if self.quantized_count != expected_quantized or self.qmax14_count != expected_qmax14: + raise ValueError(f"Unexpected {self.model_type} FP8 conversion coverage: quantized={self.quantized_count}/{expected_quantized}, qmax14={self.qmax14_count}/{expected_qmax14}") + + +def create_h3_fp8_f16_accum_quantization(profile, model_type): + if profile is None: + return None + if profile != FP8_F16_ACCUM_QUANTIZATION_PROFILE: + raise ValueError(f"Unsupported quantization profile: {profile}") + return H3FP8F16AccumQuantization(model_type) diff --git a/tools/convert/readme.md b/tools/convert/readme.md index 16f03c259..bafb617d7 100755 --- a/tools/convert/readme.md +++ b/tools/convert/readme.md @@ -49,6 +49,8 @@ A powerful model weight conversion tool that supports format conversion, quantiz - `mxfp4`: MXFP4 quantization - `mxfp6`: MXFP6 quantization - `mxfp8`: MXFP8 quantization +- `--quantization_profile`: Optional model-specific policy. `h3-fp8-f16-accum` supports `h3` and + `h3_video_vae_decoder` with `--quantized --linear_type fp8 --single_file`. - `--non_linear_dtype`: Non-linear layer data type - `torch.bfloat16`: BF16 - `torch.float16`: FP16 diff --git a/tools/convert/readme_zh.md b/tools/convert/readme_zh.md index 41f9f0be7..56bdca36f 100755 --- a/tools/convert/readme_zh.md +++ b/tools/convert/readme_zh.md @@ -41,6 +41,8 @@ - `int8`（torch.int8） - `fp8`（torch.float8_e4m3fn） - `nvfp4` / `mxfp4` / `mxfp6` / `mxfp8` +- `--quantization_profile`：可选的模型专用策略。`h3-fp8-f16-accum` 支持 `h3` 和 + `h3_video_vae_decoder`，需配合 `--quantized --linear_type fp8 --single_file`。 - `--non_linear_dtype`：非线性层数据类型（`torch.bfloat16` / `torch.float16` / `torch.float32` 默认） - `--device`：量化设备 `cpu` 或 `cuda`（默认） - `--comfyui_mode`：ComfyUI 兼容模式（仅 int8、fp8）