feat: MNNVL Allreduce quant fusion and performance optimization

csrc/trtllm_mnnvl_allreduce.cu

@@ -3,6 +3,7 @@
 
 using namespace flashinfer::trtllm_mnnvl_allreduce;
 
+using flashinfer::QuantizationSFLayout;
 using tvm::ffi::Optional;
 
 #define DISPATCH_FLOATING_TYPES_FOR_MNNVL_ALLREDUCE(dtype, c_type, ...)             \
@@ -30,23 +31,36 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
                                    int64_t buffer_ptrs_dev, int64_t buffer_ptr_local,
                                    TensorView buffer_flags_mnnvl, int64_t nranks, int64_t rank,
                                    bool rmsnorm_fusion, bool launch_with_pdl, bool use_oneshot,
-                                   TensorView output, Optional<TensorView> residual_out,
+                                   Optional<TensorView> output, Optional<TensorView> residual_out,
                                    Optional<TensorView> residual_in, Optional<TensorView> gamma,
-                                   Optional<double> epsilon, Optional<double> weight_bias) {
+                                   Optional<double> epsilon, Optional<double> weight_bias,
+                                   Optional<int64_t> quant_type, Optional<TensorView> quant_out,
+                                   Optional<TensorView> sf_out, Optional<TensorView> output_scale,
+                                   Optional<int64_t> layout_code) {
   ffi::CUDADeviceGuard device_guard(input.device().device_id);
   auto stream = get_stream(input.device());
 
   DISPATCH_FLOATING_TYPES_FOR_MNNVL_ALLREDUCE(input.dtype(), c_type, [&] {
     // Extract parameters from tensors
     int64_t num_tokens = input.size(0);
     int64_t token_dim = input.size(1);
+    auto quant_type_enum =
+        quant_type.has_value() ? static_cast<QuantType>(quant_type.value()) : QuantType::kNone;
+    auto sf_layout = layout_code.has_value()
+                         ? static_cast<QuantizationSFLayout>(layout_code.value())
+                         : QuantizationSFLayout::SWIZZLED_128x4;
 
     // Validate input parameters
     TVM_FFI_ICHECK_EQ(token_dim % (sizeof(float4) / sizeof(c_type)), 0)
         << "token_dim must be divisible by " << sizeof(float4) / sizeof(c_type);
-    TVM_FFI_ICHECK(output.size(0) == input.size(0) && output.size(1) == input.size(1))
-        << "output shape mismatch: expected (" << input.size(0) << ", " << input.size(1)
-        << ") but got (" << output.size(0) << ", " << output.size(1) << ")";
+    TVM_FFI_ICHECK(output.has_value() || quant_type_enum != QuantType::kNone)
+        << "output must be provided unless quantization fusion is enabled";
+    if (output.has_value()) {
+      TVM_FFI_ICHECK(output.value().size(0) == input.size(0) &&
+                     output.value().size(1) == input.size(1))
+          << "output shape mismatch: expected (" << input.size(0) << ", " << input.size(1)
+          << ") but got (" << output.value().size(0) << ", " << output.value().size(1) << ")";
+    }
     TVM_FFI_ICHECK(nranks >= 2 && nranks <= 64)
         << "nranks must be between 2 and 64, got " << nranks;
     TVM_FFI_ICHECK(rank >= 0 && rank < nranks)
@@ -56,6 +70,17 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
                    !rmsnorm_fusion)
         << "residual_in, residual_out, gamma, and epsilon must be provided if rmsnorm_fusion is "
            "true";
+    TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || rmsnorm_fusion)
+        << "MNNVL quantization fusion requires rmsnorm_fusion=true";
+    TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone ||
+                   (output_scale.has_value() &&
+                    encode_dlpack_dtype(output_scale.value().dtype()) == float32_code))
+        << "output_scale must be provided for MNNVL quantization fusion and must be float32";
+    TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone ||
+                   sf_layout != QuantizationSFLayout::SWIZZLED_8x4)
+        << "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts";
+    TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || quant_out.has_value())
+        << "quant_out must be provided when quantization fusion is enabled";
 
     if (rmsnorm_fusion) {
       TVM_FFI_ICHECK(residual_in.value().size(0) == num_tokens &&
@@ -72,6 +97,44 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
           << "gamma must have the same shape as token dimension (" << token_dim << ") but got ("
           << gamma.value().size(0) << ")";
     }
+    switch (quant_type_enum) {
+      case QuantType::kNone:
+        break;
+      case QuantType::kFP8:
+        TVM_FFI_ICHECK(quant_out.value().size(0) == num_tokens &&
+                       quant_out.value().size(1) == token_dim)
+            << "quant_out shape mismatch for FP8: expected (" << num_tokens << ", " << token_dim
+            << ") but got (" << quant_out.value().size(0) << ", " << quant_out.value().size(1)
+            << ")";
+        TVM_FFI_ICHECK(encode_dlpack_dtype(quant_out.value().dtype()) == float8_e4m3fn_code)
+            << "quant_out for FP8 must have dtype float8_e4m3fn";
+        break;
+      case QuantType::kFP4:
+        TVM_FFI_ICHECK(sizeof(c_type) == 2)
+            << "NVFP4 MNNVL quantization fusion is only supported for FP16/BF16 inputs";
+        TVM_FFI_ICHECK(token_dim % 16 == 0)
+            << "NVFP4 MNNVL quantization fusion requires token_dim divisible by 16";
+        TVM_FFI_ICHECK(quant_out.value().size(0) == num_tokens &&
+                       quant_out.value().size(1) == token_dim / 2)
+            << "quant_out shape mismatch for FP4: expected (" << num_tokens << ", " << token_dim / 2
+            << ") but got (" << quant_out.value().size(0) << ", " << quant_out.value().size(1)
+            << ")";
+        TVM_FFI_ICHECK(encode_dlpack_dtype(quant_out.value().dtype()) == uint8_code ||
+                       encode_dlpack_dtype(quant_out.value().dtype()) ==
+                           encode_dlpack_dtype(dl_float4_e2m1fn_x2))
+            << "quant_out for FP4 must have dtype uint8 or float4_e2m1fn_x2";
+        TVM_FFI_ICHECK(sf_out.has_value())
+            << "sf_out must be provided for NVFP4 MNNVL quantization fusion";
+        TVM_FFI_ICHECK(encode_dlpack_dtype(sf_out.value().dtype()) == float8_e4m3fn_code)
+            << "sf_out for FP4 must have dtype float8_e4m3fn";
+        TVM_FFI_ICHECK(sf_out.value().numel() >= num_tokens * token_dim / 16)
+            << "sf_out is too small for FP4: expected at least " << num_tokens * token_dim / 16
+            << " elements but got " << sf_out.value().numel();
+        break;
+      default:
+        TVM_FFI_LOG_AND_THROW(NotImplementedError)
+            << "Unsupported MNNVL quantization type " << static_cast<int>(quant_type_enum);
+    }
 
     // Create the parameters struct
     AllReduceFusionParams params;
@@ -87,6 +150,8 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
     params.bufferFlags = reinterpret_cast<uint32_t*>(buffer_flags_mnnvl.data_ptr());
     params.rmsNormFusion = rmsnorm_fusion;
     params.launchWithPdl = launch_with_pdl;
+    params.sfLayout = sf_layout;
+    params.quantType = quant_type_enum;
 
     // input data
     params.input = const_cast<void const*>(input.data_ptr());
@@ -95,11 +160,18 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
     params.gamma = gamma.has_value() ? const_cast<void const*>(gamma.value().data_ptr()) : nullptr;
     params.epsilon = epsilon.has_value() ? epsilon.value() : 1e-5;
     params.weightBias = weight_bias.has_value() ? static_cast<float>(weight_bias.value()) : 0.0f;
+    params.outputScale = output_scale.has_value()
+                             ? reinterpret_cast<float*>(output_scale.value().data_ptr())
+                             : nullptr;
 
     // output data
-    params.output = const_cast<void*>(output.data_ptr());
+    params.output = output.has_value() ? const_cast<void*>(output.value().data_ptr()) : nullptr;
     params.residualOut =
         residual_out.has_value() ? const_cast<void*>(residual_out.value().data_ptr()) : nullptr;
+    params.quantOut =
+        quant_out.has_value() ? reinterpret_cast<void*>(quant_out.value().data_ptr()) : nullptr;
+    params.scalingFactorOut =
+        sf_out.has_value() ? reinterpret_cast<void*>(sf_out.value().data_ptr()) : nullptr;
     params.stream = stream;
 
     cudaError_t status;

docs/api/comm.rst

@@ -140,6 +140,7 @@ TensorRT-LLM MNNVL AllReduce
     trtllm_mnnvl_all_reduce
     trtllm_mnnvl_allreduce
     trtllm_mnnvl_fused_allreduce_add_rmsnorm
+    trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant
     trtllm_mnnvl_fused_allreduce_rmsnorm
     mpi_barrier
 

flashinfer/comm/allreduce.py

@@ -74,10 +74,13 @@
 # Note: AllReduceFusionPattern and QuantizationSFLayout are pseudo-types (classes with int constants)
 # Import them for runtime use but type hint as int for mypy compatibility
 from .trtllm_ar import AllReduceFusionPattern
+from .trtllm_ar import QuantizationSFLayout
 from .trtllm_mnnvl_ar import MNNVLAllReduceFusionWorkspace
 from .trtllm_mnnvl_ar import MNNVLAllreduceFusionStrategy
+from .trtllm_mnnvl_ar import MNNVLQuantType
 from .trtllm_mnnvl_ar import trtllm_mnnvl_allreduce
 from .trtllm_mnnvl_ar import trtllm_mnnvl_fused_allreduce_add_rmsnorm
+from .trtllm_mnnvl_ar import trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant
 
 # ============================================================================
 # WORKSPACE IMPLEMENTATIONS
@@ -270,7 +273,8 @@ def _workspace_creation_heuristic(
 
     # Single-node scenarios
     # From benchmarking data, we can see that MNNVL is either on par (smaller problem sizes) or significantly faster than TRTLLM (larger problem sizes such as hidden_dim=8192, token_num=64 for TP=4), for single-node scenarios.
-    # However, trtllm has a larger support surface (more fusion patterns, more quantization support, etc.)
+    # TRTLLM still has the larger specialized-fusion surface, such as MoE
+    # patterns and packed group FP8 quantization.
     if "mnnvl" in suitable_backends:
         return ["mnnvl"]
     else:
@@ -493,14 +497,15 @@ def allreduce_fusion(
     weight_bias: float = 0.0,
 ) -> torch.Tensor:
     """
-    AllReduce + RMSNorm fusion operation.
+    AllReduce + RMSNorm fusion operation, with optional FP8/NVFP4
+    quantization for supported backends.
 
     Backend is automatically determined from workspace type. If you need another backend, create the workspace for the desired backend.
 
     Supports multiple fusion patterns:
     - AllReduce only
     - AllReduce + Residual + RMSNorm
-    - AllReduce + Residual + RMSNorm + Quantization (FP8/FP4)
+    - AllReduce + Residual + RMSNorm + Quantization (FP8/NVFP4)
 
     **Note on Workspace Reusability:**
     You can reuse the same workspace with different (token_num, hidden_dim) combinations
@@ -509,7 +514,7 @@ def allreduce_fusion(
     Args:
         input: Input tensor [token_num, hidden_dim]
         workspace: Workspace object (type determines backend, see create_allreduce_fusion_workspace)
-        pattern: Fusion pattern (AllReduceFusionPattern constant, 0-7)
+        pattern: Fusion pattern (AllReduceFusionPattern constant, 0-9)
                  - kAllReduce = 0
                  - kARResidualRMSNorm = 1
                  - kARResidualRMSNormFP8Quant = 2
@@ -520,8 +525,8 @@ def allreduce_fusion(
                  - kMoEFinalizeARResidualRMSNorm = 7 (trtllm only)
                  - kARResidualRMSNormPerTokenGroupFP8PackedQuant = 8 (trtllm only)
                  - kARResidualRMSNormOutPerTokenGroupFP8PackedQuant = 9 (trtllm only)
-                 Note: MNNVL only supports patterns 0 and 1
-                 Note: MOE patterns (6-7) only support trtllm backend
+                 MNNVL supports the standard FP8/NVFP4 quant patterns (2-5).
+                 MoE and packed group quant patterns remain TRTLLM-only.
         launch_with_pdl: Use Programmatic Dependent Launch
         trigger_completion_at_end: [trtllm only] Controls when PDL completion is signaled.
                      True (default): signal completion after the kernel finishes (safe, no overlap).
@@ -533,8 +538,10 @@ def allreduce_fusion(
         output: AllReduce output [token_num, hidden_dim]
         residual_out: Prenorm output (after residual add, before norm) [token_num, hidden_dim]
         norm_out: Normalized output [token_num, hidden_dim]
-        quant_out: Quantized output [token_num, hidden_dim] [trtllm only]
-        scale_out: Quantization scale factors [trtllm only]
+        quant_out: Quantized output. FP8 uses shape [token_num, hidden_dim]
+                   and NVFP4 uses shape [token_num, hidden_dim / 2].
+        scale_out: NVFP4 quantization scale factors. Not used by per-tensor
+                   FP8 quantization.
 
         # ===== INPUT parameters =====
         residual_in: Residual tensor to ADD [token_num, hidden_dim]
@@ -543,11 +550,14 @@ def allreduce_fusion(
         weight_bias: Bias added to rms_gamma before scaling.
                      0.0 (default) -> standard RMSNorm (out = gamma * x * rsqrt(...)).
                      1.0           -> Gemma / Qwen3.5 RMSNorm (out = (1 + gamma) * x * rsqrt(...)).
-                     Supported by both TRTLLM and MNNVL backends for kARResidualRMSNorm
-                     plus all TRTLLM RMSNorm variants (quant + MoE Reduction/Finalize).
+                     Supported by both TRTLLM and MNNVL backends for standard
+                     RMSNorm and quant patterns (1-5), and by TRTLLM for MoE
+                     RMSNorm variants.
                      Ignored for kAllReduce (no normalization).
-        scale_factor: Input scale factor for quantization [trtllm only]
-        layout_code: Scale factor layout (QuantizationSFLayout) [trtllm only]
+        scale_factor: Output scale used by FP8/NVFP4 quantization.
+        layout_code: NVFP4 scale factor layout (QuantizationSFLayout).
+                     MNNVL supports SWIZZLED_128x4 and LINEAR; SWIZZLED_8x4
+                     remains TRTLLM-only.
 
         # ===== Control parameters =====
         use_oneshot: Use oneshot strategy vs twoshot
@@ -569,7 +579,9 @@ def allreduce_fusion(
         shared_expert_output: Optional shared expert output to add [token_num, hidden_dim]
 
     Returns:
-        Output tensor (typically norm_out for fusion cases, output otherwise)
+        Output tensor for the selected pattern. Quant patterns return
+        quant_out, RMSNorm patterns return norm_out, and kAllReduce returns
+        output.
 
     Examples:
         >>> # Basic AllReduce + Residual + RMSNorm
@@ -599,17 +611,15 @@ def allreduce_fusion(
         ... )
         >>> # output == normed (final result)
 
-        >>> # With FP8 quantization
+        >>> # With FP8 quantization using either TRTLLM or MNNVL standard quant
+        >>> # patterns, depending on the workspace backend.
         >>> quant = torch.empty_like(hidden_states, dtype=torch.float8_e4m3fn)
-        >>> scales = torch.empty(token_num * hidden_dim // 16, dtype=torch.float16)
         >>>
         >>> output = allreduce_fusion(
         ...     input=hidden_states,
         ...     workspace=workspace,
         ...     pattern=AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
-        ...     norm_out=normed,
         ...     quant_out=quant,
-        ...     scale_out=scales,
         ...     residual_in=residual,
         ...     rms_gamma=norm_weight,
         ...     scale_factor=scale_tensor
@@ -683,9 +693,9 @@ def allreduce_fusion(
                 residual_in=residual_in,
                 rms_gamma=rms_gamma,
                 rms_eps=rms_eps,
-                scale_factor=scale_factor
-                if isinstance(scale_factor, (int, float))
-                else 1.0,
+                scale_factor=(
+                    scale_factor if isinstance(scale_factor, (int, float)) else 1.0
+                ),
                 moe_reduction_device_num_experts=moe_reduction_device_num_experts,
                 moe_reduction_scale_input=moe_reduction_scale_input,
                 moe_reduction_active_experts_token_input=moe_reduction_active_experts_token_input,
@@ -865,17 +875,51 @@ def _flatten_checked(t, name):
             return output
 
     elif isinstance(workspace, MNNVLAllReduceFusionWorkspace):
-        if (
-            pattern != AllReduceFusionPattern.kARResidualRMSNorm
-            and pattern != AllReduceFusionPattern.kAllReduce
+        strategy = (
+            MNNVLAllreduceFusionStrategy.AUTO
+            if use_oneshot is None
+            else (
+                MNNVLAllreduceFusionStrategy.ONESHOT
+                if use_oneshot
+                else MNNVLAllreduceFusionStrategy.TWOSHOT
+            )
+        )
+        mnnvl_quant_patterns = {
+            AllReduceFusionPattern.kARResidualRMSNormFP8Quant: (
+                MNNVLQuantType.FP8,
+                False,
+            ),
+            AllReduceFusionPattern.kARResidualRMSNormFP4Quant: (
+                MNNVLQuantType.NVFP4,
+                False,
+            ),
+            AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant: (
+                MNNVLQuantType.FP8,
+                True,
+            ),
+            AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant: (
+                MNNVLQuantType.NVFP4,
+                True,
+            ),
+        }
+        if pattern not in (
+            AllReduceFusionPattern.kAllReduce,
+            AllReduceFusionPattern.kARResidualRMSNorm,
+            *mnnvl_quant_patterns.keys(),
         ):
             raise ValueError(
                 f"MNNVL AllReduce+RMS fusion does not support pattern {pattern}. Please try the TRTLLM backend instead."
             )
 
-        if layout_code is not None:
+        mnnvl_layout_code = (
+            QuantizationSFLayout.SWIZZLED_128x4 if layout_code is None else layout_code
+        )
+        if (
+            pattern in mnnvl_quant_patterns
+            and mnnvl_layout_code == QuantizationSFLayout.SWIZZLED_8x4
+        ):
             raise ValueError(
-                "MNNVL AllReduce does not support quantization fusion and thus no layout_code"
+                "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts, not SWIZZLED_8x4"
             )
 
         # MNNVL backend implementation
@@ -888,6 +932,7 @@ def _flatten_checked(t, name):
                 workspace=workspace,
                 launch_with_pdl=launch_with_pdl,
                 output=output,
+                strategy=strategy,
             )
             return output
 
@@ -915,10 +960,46 @@ def _flatten_checked(t, name):
                 output=norm_out,
                 residual_out=residual_out,
                 launch_with_pdl=launch_with_pdl,
+                strategy=strategy,
                 weight_bias=weight_bias,
             )
             return norm_result
 
+        elif pattern in mnnvl_quant_patterns:
+            if residual_in is None:
+                raise ValueError(
+                    "MNNVL quantized AllReduce+RMS fusion requires residual_in"
+                )
+            if rms_gamma is None:
+                raise ValueError(
+                    "MNNVL quantized AllReduce+RMS fusion requires rms_gamma"
+                )
+
+            quant_type, has_norm_out = mnnvl_quant_patterns[pattern]
+            if has_norm_out and norm_out is None:
+                norm_out = torch.empty_like(input)
+            if residual_out is None:
+                residual_out = torch.empty_like(input)
+
+            quant_result, _, _, _ = trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant(
+                input=input,
+                residual_in=residual_in,
+                gamma=rms_gamma,
+                workspace=workspace,
+                epsilon=rms_eps,
+                output=norm_out if has_norm_out else None,
+                residual_out=residual_out,
+                quant_out=quant_out,
+                scale_out=scale_out,
+                output_scale=scale_factor,
+                layout_code=mnnvl_layout_code,
+                quant_type=quant_type,
+                launch_with_pdl=launch_with_pdl,
+                strategy=strategy,
+                weight_bias=weight_bias,
+            )
+            return quant_result
+
         else:
             raise ValueError(f"Unsupported pattern for MNNVL backend: {pattern}")