diff --git a/csrc/trtllm_mnnvl_allreduce.cu b/csrc/trtllm_mnnvl_allreduce.cu index 03badc47b72..0bdecba8883 100644 --- a/csrc/trtllm_mnnvl_allreduce.cu +++ b/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,9 +31,12 @@ 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 residual_out, + Optional output, Optional residual_out, Optional residual_in, Optional gamma, - Optional epsilon, Optional weight_bias) { + Optional epsilon, Optional weight_bias, + Optional quant_type, Optional quant_out, + Optional sf_out, Optional output_scale, + Optional layout_code) { ffi::CUDADeviceGuard device_guard(input.device().device_id); auto stream = get_stream(input.device()); @@ -40,13 +44,28 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt // 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(quant_type.value()) : QuantType::kNone; + auto sf_layout = QuantizationSFLayout::SWIZZLED_128x4; + if (layout_code.has_value()) { + auto const sf_layout_code = layout_code.value(); + TVM_FFI_ICHECK(sf_layout_code == static_cast(QuantizationSFLayout::LINEAR) || + sf_layout_code == static_cast(QuantizationSFLayout::SWIZZLED_128x4)) + << "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts"; + sf_layout = static_cast(sf_layout_code); + } // 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 +75,14 @@ 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 || 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 +99,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(quant_type_enum); + } // Create the parameters struct AllReduceFusionParams params; @@ -87,6 +152,8 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt params.bufferFlags = reinterpret_cast(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(input.data_ptr()); @@ -95,11 +162,18 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt params.gamma = gamma.has_value() ? const_cast(gamma.value().data_ptr()) : nullptr; params.epsilon = epsilon.has_value() ? epsilon.value() : 1e-5; params.weightBias = weight_bias.has_value() ? static_cast(weight_bias.value()) : 0.0f; + params.outputScale = output_scale.has_value() + ? reinterpret_cast(output_scale.value().data_ptr()) + : nullptr; // output data - params.output = const_cast(output.data_ptr()); + params.output = output.has_value() ? const_cast(output.value().data_ptr()) : nullptr; params.residualOut = residual_out.has_value() ? const_cast(residual_out.value().data_ptr()) : nullptr; + params.quantOut = + quant_out.has_value() ? reinterpret_cast(quant_out.value().data_ptr()) : nullptr; + params.scalingFactorOut = + sf_out.has_value() ? reinterpret_cast(sf_out.value().data_ptr()) : nullptr; params.stream = stream; cudaError_t status; diff --git a/docs/api/comm.rst b/docs/api/comm.rst index 1ac4b1669cb..54edefad483 100644 --- a/docs/api/comm.rst +++ b/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 diff --git a/flashinfer/comm/allreduce.py b/flashinfer/comm/allreduce.py index c0bb663bd58..aec144126a8 100644 --- a/flashinfer/comm/allreduce.py +++ b/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}") diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 32464ca98f1..5b2d2783517 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -7,7 +7,7 @@ import math import logging from types import SimpleNamespace -from typing import Optional, Tuple +from typing import Optional, Tuple, Union from enum import Enum import torch @@ -18,7 +18,9 @@ from ..jit import gen_trtllm_mnnvl_comm_module from ..utils import register_custom_op +from ..fp4_quantization import _compute_swizzled_layout_sf_size from .mnnvl import CommBackend, MPIBackend +from .trtllm_ar import QuantizationSFLayout from .workspace_base import AllReduceFusionWorkspace from .torch_symmetric_memory import _alloc_symm_buffer_bytes @@ -50,6 +52,12 @@ def select_strategy( MNNVL_ONE_SHOT_THRESHOLD = 64 * 1024 * 8 * 2 +class MNNVLQuantType: + NONE = 0 + FP8 = 1 + NVFP4 = 2 + + class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): NUM_LAMPORT_BUFFERS = 3 @@ -280,6 +288,11 @@ def get_trtllm_mnnvl_comm_module(): "residual_in", "gamma", "epsilon", + "quant_type", + "quant_out", + "sf_out", + "output_scale", + "layout_code", ], ) def trtllm_mnnvl_allreduce_fusion( @@ -293,12 +306,17 @@ def trtllm_mnnvl_allreduce_fusion( rmsnorm_fusion: bool, launch_with_pdl: bool, use_oneshot: bool, - output: torch.Tensor, + output: Optional[torch.Tensor], residual_out: Optional[torch.Tensor], residual_in: Optional[torch.Tensor], gamma: Optional[torch.Tensor], epsilon: Optional[float], weight_bias: Optional[float] = None, + quant_type: int = MNNVLQuantType.NONE, + quant_out: Optional[torch.Tensor] = None, + sf_out: Optional[torch.Tensor] = None, + output_scale: Optional[torch.Tensor] = None, + layout_code: int = QuantizationSFLayout.SWIZZLED_128x4, ) -> None: """ Perform a multi-node NVLink all-reduce operation with fusion. @@ -335,6 +353,11 @@ def trtllm_mnnvl_allreduce_fusion( gamma, epsilon, weight_bias, + quant_type, + quant_out, + sf_out, + output_scale, + layout_code, ) return SimpleNamespace( @@ -349,18 +372,42 @@ def trtllm_mnnvl_allreduce( output: Optional[torch.Tensor] = None, strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO, ) -> torch.Tensor: - """Perform a multi-node NVLink all-reduce operation across multiple GPUs. - - This function performs an all-reduce (sum) operation using NVIDIA's multi-node NVLink (MNNVL) - technology to efficiently combine tensors across multiple GPUs and nodes. - - There are 2 variants: One-shot and Two-shot: - - One-shot: Each rank stores local shard to all other ranks. Each ranks will receive all shards at the end of the communication round and perfom local reduction. Suitable for small data size and is optimized for low latency. - - Two-shot: There will be 3 steps: - 1. Scatter each GPU's input shard to other ranks. Each rank will received all shards of a slice of tokens. - 2. Each rank perform reduction on the local tokens. - 3. Each rank broadcast the result to all ranks. - Suitable for large data size and is optimized for balancing throughput and latency. + """Perform an MNNVL all-reduce sum across tensor-parallel ranks. + + ``input`` must be a 2-D local shard with shape ``[num_tokens, hidden_dim]``. + The result has the same shape and is written to ``output`` when provided, or + to a newly allocated tensor otherwise. ``workspace`` must be an + :class:`MNNVLAllReduceFusionWorkspace` created for the same tensor-parallel + group and with enough buffer capacity for the selected strategy. + + MNNVL supports two execution strategies: + + * ``ONESHOT`` stores each rank's local shard to the peer-visible workspace + and each rank performs the reduction locally. This is the low-latency path + used for smaller payloads. + * ``TWOSHOT`` scatters token slices, reduces each slice on its destination + rank, then broadcasts the reduced result. This is the throughput-oriented + path for larger payloads. + + With ``AUTO``, FlashInfer selects the strategy from the payload size. Both + ``ONESHOT`` and ``TWOSHOT`` are fully deterministic across ranks. + + This allreduce-only helper does not perform quantization. Use + :func:`trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant`, or the unified + :func:`flashinfer.comm.allreduce_fusion` API with FP8/NVFP4 + ``AllReduceFusionPattern`` values, when the post-RMSNorm quantized output is + needed. + + Determinism: + + * ``ONESHOT`` and ``TWOSHOT`` use the exact same reduction order on each + rank. + * ``ONESHOT`` keeps the local rank's value in registers; only remote ranks + are volatile-loaded from the Lamport buffer. + * ``ONESHOT`` uses a rank-specialized fast path for ``world_size <= 8``. + Larger world sizes use a compact deterministic fallback because the + runtime benefit is thin and specializing every rank significantly + increases JIT compile time. Args: input: Local Input Shard [num_tokens, hidden_dim] @@ -519,6 +566,225 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm( return output, residual_out +def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( + input: torch.Tensor, + residual_in: torch.Tensor, + gamma: torch.Tensor, + workspace: MNNVLAllReduceFusionWorkspace, + epsilon: Optional[float] = None, + output: Optional[torch.Tensor] = None, + residual_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, + scale_out: Optional[torch.Tensor] = None, + output_scale: Union[torch.Tensor, float, None] = None, + layout_code: int = QuantizationSFLayout.SWIZZLED_128x4, + quant_type: int = MNNVLQuantType.NONE, + launch_with_pdl: bool = False, + strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO, + weight_bias: float = 0.0, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """Perform MNNVL AllReduce + Residual + RMSNorm + FP8/NVFP4 quantization. + + Quantization is applied after RMSNorm. ``output`` is optional; pass it only + when the normalized non-quantized tensor is also needed. The quantized result + is always returned as ``quant_out``. + + Args: + input: Input tensor with shape ``[num_tokens, hidden_dim]``. + residual_in: Residual tensor with the same shape as ``input``. + gamma: RMSNorm gamma tensor with shape ``[hidden_dim]``. + workspace: MNNVL workspace for the tensor-parallel group. + epsilon: RMSNorm epsilon. Defaults to ``torch.finfo(input.dtype).eps``. + output: Optional normalized output tensor with shape + ``[num_tokens, hidden_dim]``. + residual_out: Optional residual output tensor with shape + ``[num_tokens, hidden_dim]``. + quant_out: Optional quantized output. For FP8, shape must be + ``[num_tokens, hidden_dim]`` and dtype ``torch.float8_e4m3fn``. For + NVFP4, shape must be ``[num_tokens, hidden_dim // 2]`` and dtype + ``torch.uint8`` or ``torch.float4_e2m1fn_x2``. + scale_out: Optional NVFP4 scale output. For ``LINEAR`` layout, shape is + ``[num_tokens, hidden_dim // 16]``. For ``SWIZZLED_128x4``, provide a + 1-D tensor large enough for the padded swizzled scale layout. FP8 + ignores this argument. + output_scale: Scalar float or float32 tensor used as the quantization + output scale. Defaults to ``1.0``. + layout_code: NVFP4 scale layout. MNNVL supports ``SWIZZLED_128x4`` and + ``LINEAR``; ``SWIZZLED_8x4`` is not supported. + quant_type: ``MNNVLQuantType.FP8`` or ``MNNVLQuantType.NVFP4``. + launch_with_pdl: Whether to launch with PDL. + strategy: MNNVL execution strategy. ``AUTO`` uses internal heuristics. + weight_bias: Bias added to gamma before scaling. ``0.0`` for standard + RMSNorm; ``1.0`` for Gemma / Qwen3.5 RMSNorm. + + Returns: + A tuple ``(quant_out, scale_out, residual_out, output)``. ``scale_out`` + is ``None`` for FP8, and ``output`` is ``None`` unless requested. + """ + + if epsilon is None: + epsilon = torch.finfo(input.dtype).eps + if output_scale is None: + output_scale = 1.0 + + if len(input.shape) != 2: + raise ValueError( + f"The input tensor must be 2D, got {len(input.shape)}D. The shape is {input.shape}." + ) + if len(residual_in.shape) != 2: + raise ValueError( + f"The residual input tensor must be 2D, got {len(residual_in.shape)}D. The shape is {residual_in.shape}." + ) + if residual_in.shape != input.shape: + raise ValueError( + f"residual_in shape must match input shape, got {residual_in.shape} and {input.shape}." + ) + if gamma.numel() != input.shape[1]: + raise ValueError( + f"The gamma tensor must have the same number of elements as the hidden dimension, got {gamma.numel()} elements but expected {input.shape[1]} elements." + ) + + if layout_code == QuantizationSFLayout.SWIZZLED_8x4: + raise ValueError( + "MNNVL NVFP4 quantization supports SWIZZLED_128x4 or LINEAR scale layouts, not SWIZZLED_8x4." + ) + + if residual_out is None: + residual_out = torch.empty_like(input) + elif residual_out.shape != input.shape: + raise ValueError( + f"residual_out shape must match input shape, got {residual_out.shape} and {input.shape}." + ) + + if output is not None and output.shape != input.shape: + raise ValueError( + f"output shape must match input shape, got {output.shape} and {input.shape}." + ) + + if isinstance(output_scale, torch.Tensor): + if output_scale.numel() < 1: + raise ValueError("output_scale must contain at least one element") + output_scale_tensor = output_scale.to(device=input.device, dtype=torch.float32) + else: + output_scale_tensor = torch.tensor( + [float(output_scale)], dtype=torch.float32, device=input.device + ) + + token_num, hidden_dim = input.shape + if quant_type == MNNVLQuantType.FP8: + if quant_out is None: + quant_out = torch.empty_like(input, dtype=torch.float8_e4m3fn) + elif quant_out.shape != input.shape: + raise ValueError( + f"quant_out shape must be {tuple(input.shape)} for FP8, got {tuple(quant_out.shape)}." + ) + if quant_out.dtype != torch.float8_e4m3fn: + raise ValueError( + f"quant_out dtype for FP8 must be float8_e4m3fn, got {quant_out.dtype}." + ) + scale_out = None + elif quant_type == MNNVLQuantType.NVFP4: + if input.dtype == torch.float32: + raise ValueError("MNNVL NVFP4 quantization requires fp16 or bf16 input.") + if hidden_dim % 16 != 0: + raise ValueError( + f"MNNVL NVFP4 quantization requires hidden_dim divisible by 16, got {hidden_dim}." + ) + expected_quant_shape = (token_num, hidden_dim // 2) + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if quant_out is None: + quant_out = torch.empty( + expected_quant_shape, dtype=torch.uint8, device=input.device + ) + elif tuple(quant_out.shape) != expected_quant_shape: + raise ValueError( + f"quant_out shape must be {expected_quant_shape} for NVFP4, got {tuple(quant_out.shape)}." + ) + if quant_out.dtype != torch.uint8 and not ( + fp4_dtype is not None and quant_out.dtype == fp4_dtype + ): + raise ValueError( + f"quant_out dtype for NVFP4 must be uint8 or float4_e2m1fn_x2, got {quant_out.dtype}." + ) + + expected_scale_out_numel = ( + token_num * hidden_dim // 16 + if layout_code == QuantizationSFLayout.LINEAR + else _compute_swizzled_layout_sf_size(token_num, hidden_dim // 16) + ) + if scale_out is None: + if layout_code == QuantizationSFLayout.LINEAR: + scale_out = torch.empty( + (token_num, hidden_dim // 16), + dtype=torch.float8_e4m3fn, + device=input.device, + ) + else: + scale_out = torch.empty( + _compute_swizzled_layout_sf_size(token_num, hidden_dim // 16), + dtype=torch.float8_e4m3fn, + device=input.device, + ) + elif scale_out.numel() < expected_scale_out_numel: + raise ValueError( + f"scale_out is too small for NVFP4: got {scale_out.numel()} elements, need at least {expected_scale_out_numel}." + ) + if scale_out.dtype != torch.float8_e4m3fn: + raise ValueError( + f"scale_out dtype for NVFP4 must be float8_e4m3fn, got {scale_out.dtype}." + ) + else: + raise ValueError(f"Unsupported MNNVL quant_type: {quant_type}") + + if strategy == MNNVLAllreduceFusionStrategy.AUTO: + strategy = MNNVLAllreduceFusionStrategy.select_strategy( + workspace.tp_size, token_num, hidden_dim, input.dtype + ) + if not workspace.is_buffer_size_sufficient( + workspace.tp_size, token_num, hidden_dim, input.dtype, strategy + ): + raise ValueError( + f"The buffer size in the given workspace is insufficient for the given problem size. Buffer: {workspace.buffer_size_bytes} bytes, Required: {workspace.get_required_buffer_size_bytes(workspace.tp_size, token_num, hidden_dim, input.dtype, strategy)} bytes." + ) + + for tensor, name in ( + (residual_in, "residual_in"), + (gamma, "gamma"), + (output, "output"), + (residual_out, "residual_out"), + (quant_out, "quant_out"), + (scale_out, "scale_out"), + ): + if tensor is not None and tensor.device != input.device: + raise ValueError(f"{name} must be on {input.device}, got {tensor.device}.") + + module = get_trtllm_mnnvl_comm_module() + module.trtllm_mnnvl_allreduce_fusion( + input, + workspace.mc_ptr, + workspace.uc_ptrs_dev, + workspace.uc_ptr_local, + workspace.buffer_flags, + workspace.tp_size, + workspace.rank, + True, + launch_with_pdl, + strategy == MNNVLAllreduceFusionStrategy.ONESHOT, + output, + residual_out, + residual_in, + gamma, + epsilon, + weight_bias, + quant_type, + quant_out, + scale_out, + output_scale_tensor, + layout_code, + ) + return quant_out, scale_out, residual_out, output + + # Legacy API that has been deprecated; Left for backward compatibility @deprecated( "get_allreduce_mnnvl_workspace is deprecated, use MNNVLAllReduceFusionWorkspace class to manage the workspace instead" @@ -586,7 +852,14 @@ def trtllm_mnnvl_all_reduce( launch_with_pdl: bool, out: Optional[torch.Tensor] = None, ) -> None: - """Perform a multi-node NVLink all-reduce operation across multiple GPUs. + """Deprecated pointer-based MNNVL all-reduce API. + + Deprecated: + Use :func:`trtllm_mnnvl_allreduce` with + :class:`MNNVLAllReduceFusionWorkspace` instead. This legacy function is + kept for compatibility and may be removed in a future release. + + Perform a multi-node NVLink all-reduce operation across multiple GPUs. This function performs an all-reduce (sum) operation using NVIDIA's multi-node NVLink (MNNVL) technology to efficiently combine tensors across multiple GPUs and nodes. diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index 834f77f921b..e43ec428e65 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -35,11 +35,22 @@ def gen_comm_alltoall_module() -> JitSpec: def gen_trtllm_mnnvl_comm_module() -> JitSpec: + # This MNNVL path is supported on Hopper/Blackwell datacenter targets only. + # Thor (SM11x) and RTX/consumer (SM12x) are intentionally excluded here. + nvcc_flags = current_compilation_context.get_nvcc_flags_list( + supported_major_versions=[9, 10] + ) + # MNNVL allreduce is fully deterministic across ranks. Oneshot specializes + # the common TP<=8 cases to keep the local value in registers and + # volatile-load only peers from the Lamport buffer. Larger world sizes use a + # compact deterministic fallback: the runtime benefit is thin there, while + # rank-specializing every case significantly increases JIT compile time. return gen_jit_spec( "trtllm_mnnvl_comm", [ jit_env.FLASHINFER_CSRC_DIR / "trtllm_mnnvl_allreduce.cu", ], + extra_cuda_cflags=nvcc_flags, ) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 75c157147a1..429555a05bc 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -14,21 +14,38 @@ * limitations under the License. */ +#pragma once + #include #include #include +#include #include #include +#include #include +#include #include +#include #include "../exception.h" +#include "../fp4_layout.cuh" #include "../logging.h" #include "../utils.cuh" +#include "trtllm_allreduce_fusion.cuh" + namespace flashinfer { namespace trtllm_mnnvl_allreduce { +using flashinfer::QuantizationSFLayout; + +enum class QuantType : int { + kNone = 0, + kFP8 = 1, + kFP4 = 2, +}; + struct AllReduceFusionParams { int nRanks; int rank; @@ -48,12 +65,40 @@ struct AllReduceFusionParams { // 0 for standard RMSNorm (out = gamma * x * rsqrt(...)), // 1 for Gemma / Qwen3.5 (out = (1 + gamma) * x * rsqrt(...)). float weightBias = 0.f; + float* outputScale = nullptr; + QuantizationSFLayout sfLayout = QuantizationSFLayout::SWIZZLED_128x4; + QuantType quantType = QuantType::kNone; void* residualOut; void* output; + void* quantOut = nullptr; + void* scalingFactorOut = nullptr; cudaStream_t stream = nullptr; }; +template +struct AllReduceKernelParams { + T* outputPtr; + T* prenormedPtr; + T const* shardPtr; + T const* residualInPtr; + T const* gammaPtr; + T** inputPtrs; + T* mcastPtr; + T* localBufferPtr; + int nRanks; + int rank; + int numTokens; + int tokenDim; + float epsilon; + float weightBias; + float* outputScalePtr; + void* quantOutPtr; + void* scalingFactorOutPtr; + QuantizationSFLayout sfLayout; + uint32_t* bufferFlags; +}; + namespace utils { constexpr uint16_t kNEGZERO_FP16 = 0x8000U; @@ -183,7 +228,7 @@ struct LamportBufferLayout { namespace cg = cooperative_groups; // PackedType is the one used in kernel for Lamport buffer (LDG.128 or LDG.64) -template +template __device__ struct __attribute__((aligned(32))) LamportFlags { public: __device__ explicit LamportFlags(uint32_t* bufferFlags, uint32_t numStages = 1) @@ -228,42 +273,38 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { } __device__ void ctaArrive() { - int tid{0}; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - - cg::cluster_group cluster = cg::this_cluster(); - // We update the atomic counter per cluster - tid = cluster.thread_rank(); - cluster.sync(); -#else - tid = threadIdx.x; - __syncthreads(); -#endif - if (tid == 0) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) - asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) - : "memory"); -#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) - asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) - : "memory"); -#else - atomicAdd(mFlagAccessPtr, 1); -#endif + if constexpr (UseCGA) { + cg::cluster_group cluster = cg::this_cluster(); + __cluster_barrier_arrive(); + if (cluster.block_rank() == 0 && threadIdx.x < 32) { + __cluster_barrier_wait(); + arriveCounter(threadIdx.x); + } + return; + } + uint32_t const barrierThreads = round_up(static_cast(blockDim.x), 32u); + // Named CTA barrier avoids a full __syncthreads() while still ordering payload stores + // before the per-CTA Lamport arrival counter update. + if (threadIdx.x < 32) { + asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); + arriveCounter(threadIdx.x); + } else { + asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); } } __device__ void waitAndUpdate(uint4 bytesToClearPerStage) { bool isLastCtaT0{false}; int targetCount{0}; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cg::grid_group grid = cg::this_grid(); // Use the first thread instead of the last thread as the last thread may exit early isLastCtaT0 = grid.thread_rank() == 0; - targetCount = grid.num_clusters(); -#else - isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0; - targetCount = gridDim.x * gridDim.y; -#endif + if constexpr (UseCGA) { + cg::cluster_group cluster = cg::this_cluster(); + targetCount = gridDim.x * gridDim.y * gridDim.z / cluster.num_blocks(); + } else { + targetCount = gridDim.x * gridDim.y * gridDim.z; + } if (isLastCtaT0) { uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); while (*reinterpret_cast(mFlagAccessPtr) < targetCount) { @@ -287,6 +328,20 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { alignas(16) std::array mBytesToClear; LamportBufferLayout mCurBufferLayout, mDirtyBufferLayout; + __device__ void arriveCounter(int tid) { + if (tid == 0) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) + : "memory"); +#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) + asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) + : "memory"); +#else + atomicAdd(mFlagAccessPtr, 1); +#endif + } + } + inline __device__ void clearPackedBuf(void* bufferBasePtr, uint32_t globalTid, uint32_t numThreads, uint32_t bytesToClear, uint8_t dirtyIndex, uint8_t stageIdx) { @@ -323,6 +378,16 @@ union PackedVec { } }; +template +inline __device__ void sanitizeLamportPayload(PackedVec& value) { +#pragma unroll + for (int i = 0; i < sizeof(PackedType) / sizeof(T); i++) { + if (isNegZero(value.elements[i])) { + value.elements[i] = fromFloat(0.f); + } + } +} + template inline __device__ PackedType loadPacked(T* ptr) { return *reinterpret_cast(ptr); @@ -334,29 +399,198 @@ inline __device__ const PackedType loadPacked(T const* ptr) { } template -inline __device__ PackedType loadPackedVolatile(void const* ptr) { +union VolatilePackedLoad { + PackedType packed; + uint32_t words[sizeof(PackedType) / sizeof(uint32_t)]; +}; + +template +inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { static_assert(sizeof(PackedType) == 0, "Not implemented"); - return PackedType{}; + return {}; } template <> -inline __device__ float4 loadPackedVolatile(void const* ptr) { - float4 returnValue; - asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" - : "=f"(returnValue.x), "=f"(returnValue.y), "=f"(returnValue.z), "=f"(returnValue.w) - : "l"(ptr)); +inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { + VolatilePackedLoad returnValue; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]), "=r"(returnValue.words[2]), + "=r"(returnValue.words[3]) + : "l"(ptr) + : "memory"); return returnValue; } template <> -inline __device__ float2 loadPackedVolatile(void const* ptr) { - float2 returnValue; - asm volatile("ld.volatile.global.v2.f32 {%0, %1}, [%2];\n" - : "=f"(returnValue.x), "=f"(returnValue.y) - : "l"(ptr)); +inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { + VolatilePackedLoad returnValue; + asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];\n" + : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]) + : "l"(ptr) + : "memory"); return returnValue; } +template +inline __device__ bool isLamportDirty(VolatilePackedLoad const& value) { + // The dirty sentinel is a raw word; typed fp compares can flush nearby bit patterns. + // ptx memeory model only gaurantee atomicity for 64bits granularity so we check every element + // here to make sure the packed payload is consumed. + bool dirty = false; +#pragma unroll + for (int i = 0; i < sizeof(PackedType) / sizeof(uint32_t); i++) { + dirty |= value.words[i] == kNEGZERO_FP32; + } + return dirty; +} + +template +inline __device__ bool pollOneshotRemoteRank(PackedVec* remoteValues, + T* stagePtrLocal, int token, int tokenDim, + int packedIdx) { + if constexpr (Rank == LocalRank) { + return true; + } else { + auto loaded = loadPackedVolatile( + &stagePtrLocal[token * tokenDim * WorldSize + Rank * tokenDim + + packedIdx * kELTS_PER_THREAD]); + remoteValues[Rank].packed = loaded.packed; + return !isLamportDirty(loaded); + } +} + +template +inline __device__ bool pollOneshotRemoteRanks(PackedVec* remoteValues, + T* stagePtrLocal, int token, int tokenDim, + int packedIdx, std::integer_sequence) { + bool valid = true; + ((valid &= pollOneshotRemoteRank( + remoteValues, stagePtrLocal, token, tokenDim, packedIdx)), + ...); + return valid; +} + +template +inline __device__ void accumulatePacked(float (&accum)[kELTS_PER_THREAD], + PackedVec const& value) { +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] += toFloat(value.elements[i]); + } +} + +template +inline __device__ void accumulateOneshotRank(float (&accum)[kELTS_PER_THREAD], + PackedVec const* remoteValues, + PackedVec const& localValue) { + if constexpr (Rank == LocalRank) { + accumulatePacked(accum, localValue); + } else { + accumulatePacked(accum, remoteValues[Rank]); + } +} + +template +inline __device__ void accumulateOneshotRanks(float (&accum)[kELTS_PER_THREAD], + PackedVec const* remoteValues, + PackedVec const& localValue, + std::integer_sequence) { + (accumulateOneshotRank(accum, remoteValues, + localValue), + ...); +} + +template +inline __device__ void waitOneshotRemoteRanks(PackedVec* remoteValues, + T* stagePtrLocal, int token, int tokenDim, + int packedIdx) { + static_assert(LocalRank < WorldSize); + while (1) { + bool const valid = + pollOneshotRemoteRanks( + remoteValues, stagePtrLocal, token, tokenDim, packedIdx, + std::make_integer_sequence{}); + if (valid) { + break; + } + } +} + +template +inline __device__ PackedVec reduceOneshotDeterministic( + PackedVec const* remoteValues, PackedVec const& localValue) { + static_assert(LocalRank < WorldSize); + float accum[kELTS_PER_THREAD]; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] = 0.f; + } + accumulateOneshotRanks( + accum, remoteValues, localValue, std::make_integer_sequence{}); + + PackedVec packedAccum; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + packedAccum.elements[i] = fromFloat(accum[i]); + } + return packedAccum; +} + +template +inline __device__ PackedVec reduceLamportRanksChunked(T* buffer, int token, + int tokenDim, + int tokenOffset) { + static_assert(kRankChunk > 0, "kRankChunk must be positive"); + static_assert(WorldSize % kRankChunk == 0, "WorldSize must be divisible by kRankChunk"); + // Chunk ranks for large world sizes to avoid register spills from keeping every rank payload + // live at once. + float accum[kELTS_PER_THREAD]; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] = 0.f; + } + +#pragma unroll 1 + for (int rankBase = 0; rankBase < WorldSize; rankBase += kRankChunk) { + float chunkAccum[kELTS_PER_THREAD]; + while (1) { + bool valid = true; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + chunkAccum[i] = 0.f; + } +#pragma unroll + for (int rr = 0; rr < kRankChunk; rr++) { + int const r = rankBase + rr; + auto loaded = loadPackedVolatile( + &buffer[token * tokenDim * WorldSize + r * tokenDim + tokenOffset]); + PackedVec value; + value.packed = loaded.packed; + valid &= !isLamportDirty(loaded); + accumulatePacked(chunkAccum, value); + } + if (valid) { + break; + } + } +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] += chunkAccum[i]; + } + } + + PackedVec packedAccum; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + packedAccum.elements[i] = fromFloat(accum[i]); + } + return packedAccum; +} + template inline __device__ void copyF4(T_IN* dst, T_IN const* src) { float4* dst4 = reinterpret_cast(dst); @@ -457,47 +691,60 @@ inline __device__ T blockReduceSum(T val) { return blockReduceSumFull(val); } } -// A helper function to tune the grid configuration for fused oneshot and rmsnorm kernels -// Return (block_size, cluster_size, loads_per_thread) +// Tune the grid configuration for fused oneshot and RMSNorm kernels. +// Return (block_size, cluster_size, loads_per_thread). std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerThread, - int smVersionMajor) { - // Start with preferred block_size and cluster_size - int clusterSize = smVersionMajor >= 9 ? 8 : 1; + bool useCluster) { + // Step 1: start from the widest cluster we are willing to launch. MNNVL JIT only + // targets SM90/SM100, but RMSNorm may request a no-CGA fallback for multi-load rows. + int clusterSize = useCluster ? 8 : 1; int blockSize = 128; - // ========================== Adjust the grid configuration ========================== int threadsNeeded = ceil_div(dim, eltsPerThread); int loadsPerThread = 1; blockSize = ceil_div(threadsNeeded, clusterSize); - if (smVersionMajor >= 9) { + if (useCluster) { + // Step 2: shrink the cluster until the hidden dimension partitions cleanly across CTAs. + // This keeps each CTA responsible for an integral chunk of packed elements. while (threadsNeeded % clusterSize != 0 && clusterSize > 1) { clusterSize /= 2; } + int const maxDivisibleClusterSize = clusterSize; blockSize = ceil_div(threadsNeeded, clusterSize); + // Step 3: if divisibility leaves each CTA too small, trade cluster width for at least + // a 128-thread CTA. This improves occupancy and avoids tiny CTAs when the row is narrow. while (blockSize < 128 && clusterSize >= 2) { blockSize *= 2; clusterSize /= 2; } int smCount = GetCudaMultiProcessorCount(); + // Step 4: if the token grid already has enough CTAs to cover the GPU, reduce cluster + // width and make CTAs larger. This avoids over-partitioning one token across too many CTAs. while (numTokens * clusterSize > smCount && clusterSize > 1 && blockSize <= 512) { blockSize *= 2; clusterSize /= 2; } - } - // Trying to scale up use multiple loads or CGA - while (blockSize > 1024) { - if (smVersionMajor >= 9) { - if (clusterSize < 8) { - clusterSize = clusterSize << 1; - } else { + // Step 5: if the token grid still underfills the GPU, restore cluster width up to the + // divisibility limit. We accept 64-thread CTAs here to expose more CTAs per token. + while (clusterSize < maxDivisibleClusterSize) { + int const candidateClusterSize = clusterSize * 2; + int const candidateBlockSize = ceil_div(threadsNeeded, candidateClusterSize); + if (candidateBlockSize < 64 || numTokens * candidateClusterSize > smCount) { break; } + clusterSize = candidateClusterSize; + blockSize = candidateBlockSize; + } + } + // Step 6: for very wide rows, first increase cluster width on SM90+ and then increase + // per-thread loads. The goal is to keep block_size within CUDA's 1024-thread limit. + while (blockSize > 1024) { + if (useCluster && clusterSize < 8) { + clusterSize = clusterSize << 1; + } else if (loadsPerThread < 8) { + loadsPerThread += 1; } else { - if (loadsPerThread < 8) { - loadsPerThread += 1; - } else { - break; - } + break; } blockSize = ceil_div(threadsNeeded, clusterSize * loadsPerThread); } @@ -507,24 +754,76 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh using utils::blockReduceSum; using utils::fromFloat; +using utils::isLamportDirty; using utils::isNegZero; using utils::LamportFlags; using utils::loadPacked; using utils::loadPackedVolatile; using utils::PackedVec; +using utils::sanitizeLamportPayload; using utils::toFloat; -template +namespace quant { + +template +inline __device__ void quant_fp8(PackedVec packedAccum, void* quantOutPtr, + float invOutputScale, uint32_t threadOffset) { + static_assert(ELTS_PER_THREAD == 8 || ELTS_PER_THREAD == 4, "ELTS_PER_THREAD must be 8 or 4"); + using QuantizedPackedType = std::conditional_t; + + auto quantOut = reinterpret_cast<__nv_fp8_e4m3*>(quantOutPtr); + PackedVec quantizedAccum; +#pragma unroll + for (int i = 0; i < ELTS_PER_THREAD; i++) { + quantizedAccum.elements[i] = + __nv_fp8_e4m3(toFloat(packedAccum.elements[i]) * invOutputScale); + } + reinterpret_cast(&quantOut[threadOffset])[0] = quantizedAccum.packed; +} + +template +inline __device__ void quant_nvfp4(PackedVec packedAccum, void* quantOutPtr, + void* sfOutPtr, float* outputScale, uint32_t tokenIdx, + uint32_t tokenDim, uint32_t packedIdx, + QuantizationSFLayout sfLayout) { +#if CUDA_VERSION >= 12080 + static_assert( + ELTS_PER_THREAD == 8 && (std::is_same_v || std::is_same_v), + "NVFP4 quantization fusion is only supported for FP16/BF16!"); + + auto packedAccumVec = *reinterpret_cast*>(&packedAccum); + auto sfOut = trtllm_allreduce_fusion::utils::cvt_quant_to_fp4_get_sf_out_offset< + uint32_t, trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE / ELTS_PER_THREAD>( + cuda::std::nullopt, tokenIdx, packedIdx, cuda::std::nullopt, tokenDim, + reinterpret_cast(sfOutPtr), sfLayout); + + uint32_t quantOutOffset = tokenIdx * tokenDim / ELTS_PER_THREAD + packedIdx; + reinterpret_cast(quantOutPtr)[quantOutOffset] = + trtllm_allreduce_fusion::utils::cvt_warp_fp16_to_fp4(packedAccumVec, + *outputScale, sfOut); +#else + (void)packedAccum; + (void)quantOutPtr; + (void)sfOutPtr; + (void)outputScale; + (void)tokenIdx; + (void)tokenDim; + (void)packedIdx; + (void)sfLayout; +#endif +} + +} // namespace quant + +template __global__ void __launch_bounds__(1024) - oneshotAllreduceFusionKernel(T* outputPtr, T* prenormedPtr, T const* shardPtr, - T const* residualInPtr, T const* gammaPtr, T** inputPtrs, - T* mcastPtr, int const numTokens, int const tokenDim, - float epsilon, float weightBias, int const rank, - uint32_t* bufferFlags) { + oneshotAllreduceFusionKernel(AllReduceKernelParams params) { + static_assert(QType == QuantType::kNone || RMSNormFusion, + "Quant-only pattern without RMSNorm is not supported!"); constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); - constexpr int kLAMPORT_ELTS_PER_PACKED = sizeof(PackedType) / sizeof(float); constexpr uint32_t kELT_SIZE = sizeof(T); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + int const tokenDim = params.tokenDim; namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); int packedIdx = cluster.thread_rank(); @@ -532,93 +831,82 @@ __global__ void __launch_bounds__(1024) int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; cudaGridDependencySynchronize(); -#else - int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; - int token = blockIdx.x; - // Offset w.r.t. the input shard - int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; -#endif - // We only use 1 stage for the oneshot allreduce - LamportFlags flag(bufferFlags, 1); - T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(mcastPtr, 0)); - T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], 0)); + constexpr bool kUseLamportCGA = true; + LamportFlags flag(params.bufferFlags, 1); + T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); + T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); if (packedIdx * kELTS_PER_THREAD >= tokenDim) { flag.ctaArrive(); - flag.clearDirtyLamportBuf(inputPtrs[rank], -1); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); return; } // ==================== Broadcast tokens to each rank ============================= PackedVec val; - val.packed = loadPacked(&shardPtr[threadOffset]); -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) val.elements[i] = fromFloat(0.f); - } + val.packed = loadPacked(¶ms.shardPtr[threadOffset]); + sanitizeLamportPayload(val); reinterpret_cast( - &stagePtrMcast[token * tokenDim * WorldSize + rank * tokenDim])[packedIdx] = val.packed; + &stagePtrMcast[token * tokenDim * WorldSize + params.rank * tokenDim])[packedIdx] = + val.packed; flag.ctaArrive(); // ======================= Lamport Sync and clear the output buffer from previous iteration // ============================= - flag.clearDirtyLamportBuf(inputPtrs[rank], -1); - - PackedVec valuesLamport[WorldSize]; - while (1) { - bool valid = true; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - valuesLamport[r].packed = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + - packedIdx * kELTS_PER_THREAD]); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); -#pragma unroll - for (int i = 0; i < kLAMPORT_ELTS_PER_PACKED; i++) { - valid &= !isNegZero(valuesLamport[r].elements[i]); - } - } - if (valid) { - break; - } - } - - auto values = reinterpret_cast*>(valuesLamport); // ======================= Reduction ============================= - float accum[kELTS_PER_THREAD]; + // Fully deterministic: every rank uses the exact same reduction order. + // For WorldSize <= 8, specialize the local slot so the fast path reuses `val` + // from registers without a dynamic `remoteValues[params.rank]` store. Larger + // world sizes use the compact fallback because the benefit is thin but + // specializing every rank significantly increases JIT compile time. PackedVec packedAccum; - -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] = toFloat(values[0].elements[i]); - } - -#pragma unroll - for (int r = 1; r < WorldSize; r++) { -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] += toFloat(values[r].elements[i]); + if constexpr (WorldSize <= 8) { + packedAccum = val; +#define RUN_ONESHOT_LOCAL_RANK(LOCAL_RANK) \ + case LOCAL_RANK: \ + if constexpr (WorldSize > LOCAL_RANK) { \ + PackedVec remoteValues[WorldSize]; \ + utils::waitOneshotRemoteRanks( \ + remoteValues, stagePtrLocal, token, tokenDim, packedIdx); \ + packedAccum = utils::reduceOneshotDeterministic(remoteValues, val); \ + } \ + break + + switch (params.rank) { + RUN_ONESHOT_LOCAL_RANK(0); + RUN_ONESHOT_LOCAL_RANK(1); + RUN_ONESHOT_LOCAL_RANK(2); + RUN_ONESHOT_LOCAL_RANK(3); + RUN_ONESHOT_LOCAL_RANK(4); + RUN_ONESHOT_LOCAL_RANK(5); + RUN_ONESHOT_LOCAL_RANK(6); + RUN_ONESHOT_LOCAL_RANK(7); } +#undef RUN_ONESHOT_LOCAL_RANK + } else { + // Chunk large-world reductions to avoid register spills from a live values[WorldSize] array. + packedAccum = utils::reduceLamportRanksChunked( + stagePtrLocal, token, tokenDim, packedIdx * kELTS_PER_THREAD); } -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum.elements[i] = fromFloat(accum[i]); - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); -#endif if constexpr (RMSNormFusion) { // =============================== Residual =============================== PackedVec residualIn; - residualIn.packed = *reinterpret_cast(&residualInPtr[threadOffset]); + residualIn.packed = *reinterpret_cast(¶ms.residualInPtr[threadOffset]); packedAccum += residualIn; - *reinterpret_cast(&prenormedPtr[threadOffset]) = packedAccum.packed; + if (params.prenormedPtr != nullptr) { + *reinterpret_cast(¶ms.prenormedPtr[threadOffset]) = packedAccum.packed; + } // =============================== Rmsnorm ================================ PackedVec gamma; - gamma.packed = *reinterpret_cast(&gammaPtr[packedIdx * kELTS_PER_THREAD]); + gamma.packed = + *reinterpret_cast(¶ms.gammaPtr[packedIdx * kELTS_PER_THREAD]); float threadSum = 0.F; #pragma unroll @@ -630,9 +918,6 @@ __global__ void __launch_bounds__(1024) __shared__ float sharedVal[8]; // Temporary variable to share the sum within block float fullSum = blockSum; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); int const numBlocks = cluster.num_blocks(); if (numBlocks > 1) { fullSum = 0.F; @@ -646,17 +931,30 @@ __global__ void __launch_bounds__(1024) fullSum += sharedVal[i]; } } -#endif - float rcpRms = rsqrtf(fullSum / tokenDim + epsilon); + float rcpRms = rsqrtf(fullSum / tokenDim + params.epsilon); #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { packedAccum.elements[i] = fromFloat(toFloat(packedAccum.elements[i]) * rcpRms * - (weightBias + toFloat(gamma.elements[i]))); + (params.weightBias + toFloat(gamma.elements[i]))); } } - reinterpret_cast(&outputPtr[threadOffset])[0] = packedAccum.packed; + if (params.outputPtr != nullptr) { + reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = packedAccum.packed; + } + if constexpr (QType == QuantType::kFP8) { + float invOutputScale = 1.0f / (*params.outputScalePtr); + quant::quant_fp8(packedAccum, params.quantOutPtr, + invOutputScale, threadOffset); + } +#if CUDA_VERSION >= 12080 + else if constexpr (QType == QuantType::kFP4) { + quant::quant_nvfp4( + packedAccum, params.quantOutPtr, params.scalingFactorOutPtr, params.outputScalePtr, token, + tokenDim, packedIdx, params.sfLayout); + } +#endif flag.waitAndUpdate( - {static_cast(numTokens * tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); + {static_cast(params.numTokens * tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); } using utils::adjustGridConfig; @@ -667,10 +965,8 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) int const tokenDim = params.tokenDim; int const eltsPerThread = sizeof(float4) / sizeof(T); - static const int kSMVersionMajor = GetCudaComputeCapability().first; - auto [blockSize, clusterSize, loadsPerThread] = - adjustGridConfig(numTokens, tokenDim, eltsPerThread, kSMVersionMajor); + adjustGridConfig(numTokens, tokenDim, eltsPerThread, true); dim3 grid(numTokens, clusterSize, 1); FLASHINFER_LOG_DEBUG( @@ -683,7 +979,37 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_CHECK(blockSize <= 1024 && loadsPerThread == 1, "Hidden Dimension %d exceeds the maximum supported hidden dimension (%d)", - tokenDim, 1024 * (kSMVersionMajor >= 9 ? 8 : 1) * eltsPerThread); + tokenDim, 1024 * 8 * eltsPerThread); + + if (!params.rmsNormFusion && params.quantType != QuantType::kNone) { + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] Quantization requires RMSNorm fusion"); + return cudaErrorInvalidValue; + } + if (params.quantType != QuantType::kNone) { + if (params.quantOut == nullptr || params.outputScale == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceOneShot] quantOut and outputScale must be non-null when " + "quantization is enabled"); + return cudaErrorInvalidValue; + } + } + if (params.quantType == QuantType::kFP4) { +#if CUDA_VERSION < 12080 + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] FP4 quantization requires CUDA 12.8+"); + return cudaErrorInvalidValue; +#else + if (params.scalingFactorOut == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceOneShot] scalingFactorOut is required for FP4 quantization"); + return cudaErrorInvalidValue; + } + if (tokenDim % trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE != 0) { + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] FP4 quantization requires tokenDim divisible by " + + std::to_string(trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE)); + return cudaErrorInvalidValue; + } +#endif + } cudaLaunchAttribute attrs[2]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; @@ -699,31 +1025,63 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .dynamicSmemBytes = 0, .stream = params.stream, .attrs = attrs, - .numAttrs = kSMVersionMajor >= 9 ? 2 : 1, + .numAttrs = 2, }; -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &config, &oneshotAllreduceFusionKernel, output, residualOut, input, \ - residualIn, gamma, ucPtrs, mcPtr, numTokens, tokenDim, static_cast(params.epsilon), \ - params.weightBias, params.rank, params.bufferFlags)); -#define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ - if (params.rmsNormFusion) { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true); \ - } else { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, false); \ - } - - T** ucPtrs = reinterpret_cast(params.bufferPtrsDev); - T* mcPtr = reinterpret_cast(params.multicastPtr); - T* output = reinterpret_cast(params.output); - T* residualOut = reinterpret_cast(params.residualOut); - T const* input = reinterpret_cast(params.input); - T const* residualIn = reinterpret_cast(params.residualIn); - T const* gamma = reinterpret_cast(params.gamma); +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM, QTYPE) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &config, &oneshotAllreduceFusionKernel, kernelParams)); +#define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ + if (params.rmsNormFusion) { \ + switch (params.quantType) { \ + case QuantType::kFP8: \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true, QuantType::kFP8); \ + break; \ + case QuantType::kFP4: \ + if constexpr (std::is_same_v || std::is_same_v) { \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true, QuantType::kFP4); \ + } else { \ + FLASHINFER_ERROR( \ + "[MNNVL AllReduceOneShot] FP4 quantization is only " \ + "supported for FP16/BF16"); \ + return cudaErrorInvalidValue; \ + } \ + break; \ + case QuantType::kNone: \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true, QuantType::kNone); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] Unsupported quant type " + \ + std::to_string(static_cast(params.quantType))); \ + return cudaErrorInvalidValue; \ + } \ + } else { \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, false, QuantType::kNone); \ + } + + AllReduceKernelParams kernelParams{ + .outputPtr = reinterpret_cast(params.output), + .prenormedPtr = reinterpret_cast(params.residualOut), + .shardPtr = reinterpret_cast(params.input), + .residualInPtr = reinterpret_cast(params.residualIn), + .gammaPtr = reinterpret_cast(params.gamma), + .inputPtrs = reinterpret_cast(params.bufferPtrsDev), + .mcastPtr = reinterpret_cast(params.multicastPtr), + .localBufferPtr = reinterpret_cast(params.bufferPtrLocal), + .nRanks = params.nRanks, + .rank = params.rank, + .numTokens = params.numTokens, + .tokenDim = params.tokenDim, + .epsilon = static_cast(params.epsilon), + .weightBias = params.weightBias, + .outputScalePtr = params.outputScale, + .quantOutPtr = params.quantOut, + .scalingFactorOutPtr = params.scalingFactorOut, + .sfLayout = params.sfLayout, + .bufferFlags = params.bufferFlags, + }; switch (params.nRanks) { - // FIXME: Do we need other world sizes? case 2: DISPATCH_ALLREDUCE_KERNEL(2); break; @@ -747,6 +1105,7 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) ". Supported sizes: {2, 4, 8, 16, 32, 64}"); return cudaErrorInvalidValue; } +#undef DISPATCH_ALLREDUCE_KERNEL #undef LAUNCH_ALLREDUCE_KERNEL return cudaSuccess; } @@ -757,152 +1116,85 @@ enum MNNVLTwoShotStage : uint8_t { NUM_STAGES = 2, }; -template -__global__ __launch_bounds__(128) void twoshotAllreduceKernel( - T* outputPtr, T const* shardPtr, T** inputPtrs, T* mcastPtr, uint32_t const numTokens, - uint32_t const tokenDim, uint32_t const rank, uint32_t* bufferFlags, - bool const wait_for_results) { +using utils::copyF4; + +template +__global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelParams params) { constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); - constexpr int kLAMPORT_ELTS_PER_PACKED = sizeof(PackedType) / sizeof(float); constexpr uint32_t kELT_SIZE = sizeof(T); - int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; - int token = blockIdx.x; - // Offset w.r.t. the input shard - int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; + int const packedIdx = blockIdx.y * blockDim.x + threadIdx.x; + int const token = blockIdx.x; + int const tokenOffset = packedIdx * kELTS_PER_THREAD; + int const threadOffset = token * params.tokenDim + tokenOffset; + bool const inBounds = tokenOffset < params.tokenDim; + int const destRank = token % WorldSize; + int const destTokenOffset = token / WorldSize; - int destRank = token % WorldSize; - int destTokenOffset = token / WorldSize; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); -#endif - LamportFlags flag(bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - T* scatterBufLocal = - reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::SCATTER)); - T* scatterBufDest = - reinterpret_cast(flag.getCurLamportBuf(inputPtrs[destRank], MNNVLTwoShotStage::SCATTER)); + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + T* scatterBufLocal = reinterpret_cast( + flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER)); + T* scatterBufDest = reinterpret_cast( + flag.getCurLamportBuf(params.inputPtrs[destRank], MNNVLTwoShotStage::SCATTER)); T* broadcastBufW = - reinterpret_cast(flag.getCurLamportBuf(mcastPtr, MNNVLTwoShotStage::BROADCAST)); - T* broadcastBufR = - reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::BROADCAST)); - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - // Make sure the clear function is called before OOB thread exits - if (packedIdx * kELTS_PER_THREAD >= tokenDim) { - flag.clearDirtyLamportBuf(inputPtrs[rank], -1); - return; - } - - // =============================== Scatter =============================== + reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST)); + T* broadcastBufR = reinterpret_cast( + flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST)); - // Load vectorized data PackedVec val; - val.packed = loadPacked(&shardPtr[threadOffset]); -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) { - val.elements[i] = fromFloat(0.F); - } - } + if (inBounds) { + val.packed = loadPacked(¶ms.shardPtr[threadOffset]); + sanitizeLamportPayload(val); - // Store vectorized data - reinterpret_cast( - &scatterBufDest[destTokenOffset * tokenDim * WorldSize + rank * tokenDim])[packedIdx] = - val.packed; - - flag.clearDirtyLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::SCATTER); - - // =============================== Reduction and Broadcast =============================== - - if ((token % WorldSize) == rank) { - int localToken = token / WorldSize; - float accum[kELTS_PER_THREAD] = {0.F}; - - // Use float as we only check each float value for validity - PackedVec valuesLamport[WorldSize]; - while (1) { - bool valid = true; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - valuesLamport[r].packed = loadPackedVolatile( - &scatterBufLocal[localToken * tokenDim * WorldSize + r * tokenDim + - packedIdx * kELTS_PER_THREAD]); - - // Check validity across all elements -#pragma unroll - for (int i = 0; i < kLAMPORT_ELTS_PER_PACKED; i++) { - valid &= !isNegZero(valuesLamport[r].elements[i]); - } - } - if (valid) { - break; - } - } - - // Now we view it as the value for reduction - auto values = reinterpret_cast*>(valuesLamport); -#pragma unroll - for (int r = 0; r < WorldSize; r++) { -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] += toFloat(values[r].elements[i]); - } - } + reinterpret_cast(&scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + + params.rank * params.tokenDim])[packedIdx] = + val.packed; + } - // Store vectorized result - PackedVec packedAccum; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum.elements[i] = fromFloat(accum[i]); - } - reinterpret_cast(&broadcastBufW[token * tokenDim])[packedIdx] = packedAccum.packed; + cudaTriggerProgrammaticLaunchCompletion(); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); + + if (inBounds && destRank == params.rank) { + int const localToken = token / WorldSize; + // Fully deterministic: every rank uses the exact same reduction order. Chunking avoids + // register spills for large world sizes. + constexpr int kRankChunk = WorldSize < 16 ? WorldSize : 16; + PackedVec packedAccum = + utils::reduceLamportRanksChunked( + scatterBufLocal, localToken, params.tokenDim, tokenOffset); + // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. + sanitizeLamportPayload(packedAccum); + reinterpret_cast(&broadcastBufW[token * params.tokenDim])[packedIdx] = + packedAccum.packed; } - flag.clearDirtyLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::BROADCAST); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); - // Optionally wait for results if the next layer isn't doing the Lamport check - if (wait_for_results) { - // Update the atomic counter to indicate the block has read the offsets + if constexpr (WaitForResults) { + // OOB threads still arrive so waitAndUpdate sees every CTA. flag.ctaArrive(); - PackedVec valLamport; - valLamport.packed = loadPackedVolatile(&broadcastBufR[threadOffset]); - while (isNegZero(valLamport.elements[0])) { - valLamport.packed = loadPackedVolatile(&broadcastBufR[threadOffset]); - } - if (outputPtr) { - reinterpret_cast(&outputPtr[threadOffset])[0] = valLamport.packed; + if (inBounds) { + auto loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); + while (isLamportDirty(loaded)) { + loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); + } + reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = loaded.packed; } - // Update the buffer flags flag.waitAndUpdate( - {static_cast(round_up(numTokens, WorldSize) * tokenDim * - kELT_SIZE), // Clear Size for scatter stage - static_cast(numTokens * tokenDim * kELT_SIZE), // Clear Size for broadcast stage - 0, 0}); - // If not wait for results, we will rely on the following kernel to update the buffer + {static_cast(round_up(params.numTokens, WorldSize) * params.tokenDim * kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } } -using utils::copyF4; -// This kernel works performant when loads_per_thread is 1. -// For this mode, we are able to support up to 1024 (threads) x 8 (elements) = 8192 hidden -// dimension. There are two options for further scaling up: -// 1. Use CGA if supported. It expands the hidden dimension to 8k x 8 = 64k. -// 2. Set loads_per_thread >1. Which can be used if CGA is not supported. Note that this will -// be limited by the shared memory size and register count. -template -__global__ __launch_bounds__(1024) void rmsNormLamport(T_IN* outputPreNorm, T_OUT* outputNorm, - T_IN* bufferInput, T_IN const* gamma, - float epsilon, float weightBias, - T_IN const* residual, uint32_t numTokens, - uint32_t dim, uint32_t worldSize, - uint32_t* bufferFlags) { - static_assert(std::is_same_v, "T_IN and T_OUT must be the same type"); - static int const kELTS_PER_LOAD = sizeof(float4) / sizeof(T_IN); +template +__global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams params) { + constexpr int kELTS_PER_LOAD = sizeof(PackedType) / sizeof(T); + constexpr uint32_t kELT_SIZE = sizeof(T); uint32_t const token = blockIdx.x; uint32_t const blockSize = blockDim.x; @@ -910,157 +1202,156 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(T_IN* outputPreNorm, T_OU uint32_t numThreads = blockSize; uint32_t clusterSize = 1; - uint32_t blockOffset = 0; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - numThreads = cluster.num_threads(); - clusterSize = cluster.num_blocks(); - blockOffset = cluster.block_rank(); -#endif - uint32_t const dimPadded = round_up(dim, kELTS_PER_LOAD * numThreads); + uint32_t clusterBlockRank = 0; + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + numThreads = cluster.num_threads(); + clusterSize = cluster.num_blocks(); + clusterBlockRank = cluster.block_rank(); + } + + uint32_t const dimPadded = round_up(params.tokenDim, kELTS_PER_LOAD * numThreads); uint32_t const elemsPerThread = dimPadded / numThreads; uint32_t const loadStride = blockSize; + uint32_t const blockChunkSize = + ceil_div(params.tokenDim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; + uint32_t const baseTokenOffset = clusterBlockRank * blockChunkSize; extern __shared__ uint8_t smem[]; - float rInput[LoadsPerThread * kELTS_PER_LOAD]; - uint32_t offsets[LoadsPerThread * kELTS_PER_LOAD]; - - uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T_IN); - T_IN* smemInput = (T_IN*)&smem[0]; - T_IN* smemResidual = (T_IN*)&smem[smemBufferSize]; - T_IN* smemGamma = (T_IN*)&smem[2 * smemBufferSize]; + uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T); + T* smemInput = reinterpret_cast(&smem[0]); + T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); + T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); - LamportFlags flag(bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - T_IN* input = reinterpret_cast( - flag.getCurLamportBuf(reinterpret_cast(bufferInput), MNNVLTwoShotStage::BROADCAST)); - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); -#endif - // The offset that current thread should load from. Note that the hidden dimension is split by CGA - // size and each block loads a contiguous chunk; The size of chunk that each block processes - uint32_t const blockChunkSize = ceil_div(dim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; - uint32_t const blockLoadOffset = token * dim + blockOffset * blockChunkSize; -#pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - // Each block load a contiguous chunk of tokens - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - offsets[i] = blockLoadOffset + threadLoadOffset; - } + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + T* input = reinterpret_cast( + flag.getCurLamportBuf(params.localBufferPtr, MNNVLTwoShotStage::BROADCAST)); #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - copyF4(&smemResidual[threadLoadOffset], &residual[blockLoadOffset + threadLoadOffset]); + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemResidual[chunkOffset], + ¶ms.residualInPtr[token * params.tokenDim + tokenOffset]); } } __pipeline_commit(); #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - copyF4(&smemGamma[threadLoadOffset], &gamma[blockOffset * blockChunkSize + threadLoadOffset]); + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemGamma[chunkOffset], ¶ms.gammaPtr[tokenOffset]); } } __pipeline_commit(); flag.ctaArrive(); - bool valid = false; - // ACQBLK if not lamport - while (!valid) { - valid = true; -#pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - float4* dst4 = reinterpret_cast(&smemInput[threadLoadOffset]); - float4 const* src4 = reinterpret_cast(&input[offsets[i]]); - float4 value = loadPackedVolatile(src4); - // Assume that the 16B were written atomically, so we only need to check one value - valid &= !isNegZero(value.x); - *dst4 = value; +#pragma unroll + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + auto loaded = loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); + while (isLamportDirty(loaded)) { + loaded = loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); } + reinterpret_cast(&smemInput[chunkOffset])[0] = loaded.packed; } } + float rInput[LoadsPerThread * kELTS_PER_LOAD]; + float threadSum = 0.f; + + // Residual and gamma staging use normal cp.async groups; residual is consumed first. __pipeline_wait_prior(1); - __syncthreads(); - float threadSum = 0.f; #pragma unroll - for (int i = 0; i < LoadsPerThread; i++) { - int threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - PackedVec inp{.packed = loadPacked(&smemInput[threadLoadOffset])}; - PackedVec res{.packed = loadPacked(&smemResidual[threadLoadOffset])}; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; + PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; + PackedVec inpPlusRes = inp + res; + if (params.prenormedPtr != nullptr) { + reinterpret_cast( + ¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = inpPlusRes.packed; + } - PackedVec inp_plus_res = inp + res; #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { - rInput[i * kELTS_PER_LOAD + j] = toFloat(inp_plus_res.elements[j]); - threadSum += toFloat(inp_plus_res.elements[j] * inp_plus_res.elements[j]); + float const value = toFloat(inpPlusRes.elements[j]); + rInput[i * kELTS_PER_LOAD + j] = value; + threadSum += value * value; } - - *reinterpret_cast(&outputPreNorm[blockLoadOffset + threadLoadOffset]) = - inp_plus_res.packed; } } __pipeline_wait_prior(0); float blockSum = blockReduceSum(threadSum); - float fullSum = blockSum; __shared__ float sharedVal[8]; - // Use CGA Reduction if supported -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - int const numBlocks = cluster.num_blocks(); - if (numBlocks > 1) { - fullSum = 0.F; - // Need to reduce over the entire cluster - int const blockRank = cluster.block_rank(); - if (threadIdx.x < numBlocks) { - cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; - } - cluster.barrier_wait(cluster.barrier_arrive()); - for (int i = 0; i < numBlocks; ++i) { - fullSum += sharedVal[i]; + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const numBlocks = cluster.num_blocks(); + if (numBlocks > 1) { + fullSum = 0.F; + int const blockRank = cluster.block_rank(); + if (threadIdx.x < numBlocks) { + cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; + } + cluster.barrier_wait(cluster.barrier_arrive()); + for (int i = 0; i < numBlocks; ++i) { + fullSum += sharedVal[i]; + } } } -#endif - - float rcpRms = rsqrtf(fullSum / dim + epsilon); + float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); #pragma unroll - for (int i = 0; i < LoadsPerThread; i++) { - PackedVec r_out; - uint32_t threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - PackedVec gamma = {.packed = loadPacked(&smemGamma[threadLoadOffset])}; - + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec gamma{.packed = loadPacked(&smemGamma[chunkOffset])}; + PackedVec out; #pragma unroll - for (uint32_t j = 0; j < kELTS_PER_LOAD; j++) { - r_out.elements[j] = fromFloat((weightBias + toFloat(gamma.elements[j])) * - rInput[i * kELTS_PER_LOAD + j] * rcpRms); + for (int j = 0; j < kELTS_PER_LOAD; j++) { + out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * + rInput[i * kELTS_PER_LOAD + j] * rcpRms); } - - *reinterpret_cast(&outputNorm[blockLoadOffset + threadLoadOffset]) = r_out.packed; + if (params.outputPtr != nullptr) { + reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = + out.packed; + } + if constexpr (QType == QuantType::kFP8) { + float invOutputScale = 1.0f / (*params.outputScalePtr); + quant::quant_fp8(out, params.quantOutPtr, invOutputScale, + token * params.tokenDim + tokenOffset); + } +#if CUDA_VERSION >= 12080 + else if constexpr (QType == QuantType::kFP4) { + quant::quant_nvfp4( + out, params.quantOutPtr, params.scalingFactorOutPtr, params.outputScalePtr, token, + params.tokenDim, tokenOffset / kELTS_PER_LOAD, params.sfLayout); + } +#endif } } - constexpr int kELTS_SIZE = sizeof(T_IN); - // Assume the previous kernel does not modify the buffer_flags. -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); -#endif - // Update the buffer pointers - flag.waitAndUpdate({static_cast(round_up(numTokens, worldSize) * dim * kELTS_SIZE), - static_cast(numTokens * dim * kELTS_SIZE), 0, 0}); + + flag.waitAndUpdate({static_cast(round_up(params.numTokens, params.nRanks) * + params.tokenDim * kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } template @@ -1071,9 +1362,61 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_CHECK(tokenDim % numEltsPerThread == 0, "[MNNVL AllReduceTwoShot] token_dim must be divisible by %d", numEltsPerThread); + if (!params.rmsNormFusion && params.quantType != QuantType::kNone) { + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Quantization requires RMSNorm fusion"); + return cudaErrorInvalidValue; + } + if (params.quantType != QuantType::kNone) { + if (params.quantOut == nullptr || params.outputScale == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceTwoShot] quantOut and outputScale must be non-null when " + "quantization is enabled"); + return cudaErrorInvalidValue; + } + } + if (params.quantType == QuantType::kFP4) { +#if CUDA_VERSION < 12080 + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] FP4 quantization requires CUDA 12.8+"); + return cudaErrorInvalidValue; +#else + if (params.scalingFactorOut == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceTwoShot] scalingFactorOut is required for FP4 quantization"); + return cudaErrorInvalidValue; + } + if (tokenDim % trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE != 0) { + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] FP4 quantization requires tokenDim divisible by " + + std::to_string(trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE)); + return cudaErrorInvalidValue; + } +#endif + } + int const arNumThreads = ceil_div(tokenDim, numEltsPerThread); int const arNumBlocksPerToken = ceil_div(arNumThreads, 128); + AllReduceKernelParams kernelParams{ + .outputPtr = reinterpret_cast(params.output), + .prenormedPtr = reinterpret_cast(params.residualOut), + .shardPtr = reinterpret_cast(params.input), + .residualInPtr = reinterpret_cast(params.residualIn), + .gammaPtr = reinterpret_cast(params.gamma), + .inputPtrs = reinterpret_cast(params.bufferPtrsDev), + .mcastPtr = reinterpret_cast(params.multicastPtr), + .localBufferPtr = reinterpret_cast(params.bufferPtrLocal), + .nRanks = params.nRanks, + .rank = params.rank, + .numTokens = params.numTokens, + .tokenDim = params.tokenDim, + .epsilon = static_cast(params.epsilon), + .weightBias = params.weightBias, + .outputScalePtr = params.outputScale, + .quantOutPtr = params.quantOut, + .scalingFactorOutPtr = params.scalingFactorOut, + .sfLayout = params.sfLayout, + .bufferFlags = params.bufferFlags, + }; + dim3 arGrid(numTokens, arNumBlocksPerToken); cudaLaunchAttribute arAttrs[1]; @@ -1092,131 +1435,168 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_LOG_DEBUG("[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", numTokens, arNumBlocksPerToken); -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &arConfig, &twoshotAllreduceKernel, output, input, ucPtrs, mcastPtr, \ - numTokens, tokenDim, params.rank, params.bufferFlags, (!params.rmsNormFusion))); - T** ucPtrs = reinterpret_cast(params.bufferPtrsDev); - T* mcastPtr = reinterpret_cast(params.multicastPtr); - T* output = reinterpret_cast(params.output); - T const* input = reinterpret_cast(params.input); - switch (params.nRanks) { - case 2: - LAUNCH_ALLREDUCE_KERNEL(2); - break; - case 4: - LAUNCH_ALLREDUCE_KERNEL(4); - break; - case 8: - LAUNCH_ALLREDUCE_KERNEL(8); - break; - case 16: - LAUNCH_ALLREDUCE_KERNEL(16); - break; - case 32: - LAUNCH_ALLREDUCE_KERNEL(32); - break; - case 64: - LAUNCH_ALLREDUCE_KERNEL(64); - break; - default: - FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + - std::to_string(params.nRanks) + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); - return cudaErrorInvalidValue; +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, WAIT_FOR_RESULTS) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &arConfig, &twoshotAllreduceKernel, kernelParams)); +#define DISPATCH_ALLREDUCE_KERNEL(WAIT_FOR_RESULTS) \ + switch (params.nRanks) { \ + case 2: \ + LAUNCH_ALLREDUCE_KERNEL(2, WAIT_FOR_RESULTS); \ + break; \ + case 4: \ + LAUNCH_ALLREDUCE_KERNEL(4, WAIT_FOR_RESULTS); \ + break; \ + case 8: \ + LAUNCH_ALLREDUCE_KERNEL(8, WAIT_FOR_RESULTS); \ + break; \ + case 16: \ + LAUNCH_ALLREDUCE_KERNEL(16, WAIT_FOR_RESULTS); \ + break; \ + case 32: \ + LAUNCH_ALLREDUCE_KERNEL(32, WAIT_FOR_RESULTS); \ + break; \ + case 64: \ + LAUNCH_ALLREDUCE_KERNEL(64, WAIT_FOR_RESULTS); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + \ + std::to_string(params.nRanks) + \ + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ + return cudaErrorInvalidValue; \ } -#undef LAUNCH_ALLREDUCE_KERNEL - // Launch the rmsnorm lamport kernel if fusion is enabled if (params.rmsNormFusion) { - static const int kSMVersionMajor = GetCudaComputeCapability().first; - auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, kSMVersionMajor); - int rnBlockSize = std::get<0>(gridConfig); - int rnClusterSize = std::get<1>(gridConfig); - int rnLoadsPerThread = std::get<2>(gridConfig); - - int rnNumThreads = rnClusterSize * rnBlockSize; - dim3 rnGrid(numTokens, rnClusterSize, 1); - cudaLaunchConfig_t rnConfig; - cudaLaunchAttribute rnAttrs[2]; - rnConfig.stream = params.stream; - rnConfig.gridDim = rnGrid; - rnConfig.blockDim = rnBlockSize; - rnConfig.attrs = rnAttrs; - rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - rnAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; + DISPATCH_ALLREDUCE_KERNEL(false); + } else { + DISPATCH_ALLREDUCE_KERNEL(true); + return cudaSuccess; + } +#undef DISPATCH_ALLREDUCE_KERNEL +#undef LAUNCH_ALLREDUCE_KERNEL + + auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, true); + int rnBlockSize = std::get<0>(gridConfig); + int rnClusterSize = std::get<1>(gridConfig); + int rnLoadsPerThread = std::get<2>(gridConfig); + + bool rnUseCGA = rnClusterSize > 1 && rnLoadsPerThread == 1; + if (rnClusterSize > 1 && !rnUseCGA) { + gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, false); + rnBlockSize = std::get<0>(gridConfig); + rnClusterSize = std::get<1>(gridConfig); + rnLoadsPerThread = std::get<2>(gridConfig); + rnUseCGA = false; + } + + if (rnBlockSize > 1024 || rnLoadsPerThread > 8) { + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported hidden dimension " + + std::to_string(tokenDim)); + return cudaErrorInvalidValue; + } + + int const rnNumThreads = rnClusterSize * rnBlockSize; + int const dimPadded = round_up(tokenDim, numEltsPerThread * rnNumThreads); + int const iters = dimPadded / rnNumThreads; + size_t const smemSize = 3 * rnBlockSize * iters * sizeof(T); + + dim3 rnGrid(numTokens, rnClusterSize, 1); + cudaLaunchConfig_t rnConfig; + cudaLaunchAttribute rnAttrs[2]; + rnConfig.stream = params.stream; + rnConfig.gridDim = rnGrid; + rnConfig.blockDim = rnBlockSize; + rnConfig.dynamicSmemBytes = smemSize; + rnConfig.attrs = rnAttrs; + rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + rnAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; + rnConfig.numAttrs = 1; + if (rnUseCGA) { rnAttrs[1].id = cudaLaunchAttributeClusterDimension; rnAttrs[1].val.clusterDim.x = 1; rnAttrs[1].val.clusterDim.y = rnClusterSize; rnAttrs[1].val.clusterDim.z = 1; - rnConfig.numAttrs = kSMVersionMajor >= 9 ? 2 : 1; - - bool const rnUseCGA = kSMVersionMajor >= 9 && rnClusterSize > 1; - int const dimPadded = round_up(tokenDim, numEltsPerThread * rnNumThreads); - int const iters = dimPadded / rnNumThreads; - - size_t const smemSize = 3 * rnBlockSize * iters * sizeof(T); - - FLASHINFER_LOG_DEBUG( - "[MNNVL AllReduceTwoShotRMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, " - "cluster_size: %d, " - "loads_per_thread: %d, " - "threads_needed: %d", - numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, - ceil_div(tokenDim, numEltsPerThread)); - -#define RUN_RMSNORM_KERNEL(LOADS_PER_THREAD) \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - smemSize)); \ - rnConfig.dynamicSmemBytes = smemSize; \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &rnConfig, &rmsNormLamport, residualOut, output, bufferInput, gamma, \ - static_cast(params.epsilon), params.weightBias, residualIn, numTokens, tokenDim, \ - params.nRanks, params.bufferFlags)); - - T* residualOut = reinterpret_cast(params.residualOut); - T* output = reinterpret_cast(params.output); - T* bufferInput = reinterpret_cast(params.bufferPtrLocal); - T const* gamma = reinterpret_cast(params.gamma); - T const* residualIn = reinterpret_cast(params.residualIn); - if (rnUseCGA) { - RUN_RMSNORM_KERNEL(1); - } else { - switch (rnLoadsPerThread) { - case 1: - RUN_RMSNORM_KERNEL(1); - break; - case 2: - RUN_RMSNORM_KERNEL(2); - break; - case 3: - RUN_RMSNORM_KERNEL(3); - break; - case 4: - RUN_RMSNORM_KERNEL(4); - break; - case 5: - RUN_RMSNORM_KERNEL(5); - break; - case 6: - RUN_RMSNORM_KERNEL(6); - break; - case 7: - RUN_RMSNORM_KERNEL(7); - break; - case 8: - RUN_RMSNORM_KERNEL(8); - break; - default: - FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread" + - std::to_string(rnLoadsPerThread) + - ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); - return cudaErrorInvalidValue; - } // switch (rnLoadsPerThread) - } // if (rnUseCGA) -#undef RUN_RMSNORM_KERNEL - - } // if (params.rmsNormFusion) + rnConfig.numAttrs = 2; + } + + FLASHINFER_LOG_DEBUG( + "[MNNVL AllReduceTwoShotRMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, " + "cluster_size: %d, loads_per_thread: %d, threads_needed: %d", + numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, + ceil_div(tokenDim, numEltsPerThread)); + +#define LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QTYPE) \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + smemSize)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &rnConfig, &rmsNormLamport, kernelParams)); + +#define DISPATCH_QUANT(USE_CGA, LOADS_PER_THREAD) \ + switch (params.quantType) { \ + case QuantType::kFP8: \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QuantType::kFP8); \ + break; \ + case QuantType::kFP4: \ + if constexpr (std::is_same_v || std::is_same_v) { \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QuantType::kFP4); \ + } else { \ + FLASHINFER_ERROR( \ + "[MNNVL AllReduceTwoShotRMSNorm] FP4 quantization is only " \ + "supported for FP16/BF16"); \ + return cudaErrorInvalidValue; \ + } \ + break; \ + case QuantType::kNone: \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QuantType::kNone); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported quant type " + \ + std::to_string(static_cast(params.quantType))); \ + return cudaErrorInvalidValue; \ + } + +#define DISPATCH_LOADS(USE_CGA) \ + switch (rnLoadsPerThread) { \ + case 1: \ + DISPATCH_QUANT(USE_CGA, 1); \ + break; \ + case 2: \ + DISPATCH_QUANT(false, 2); \ + break; \ + case 3: \ + DISPATCH_QUANT(false, 3); \ + break; \ + case 4: \ + DISPATCH_QUANT(false, 4); \ + break; \ + case 5: \ + DISPATCH_QUANT(false, 5); \ + break; \ + case 6: \ + DISPATCH_QUANT(false, 6); \ + break; \ + case 7: \ + DISPATCH_QUANT(false, 7); \ + break; \ + case 8: \ + DISPATCH_QUANT(false, 8); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread " + \ + std::to_string(rnLoadsPerThread) + \ + ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ + return cudaErrorInvalidValue; \ + } + + if (rnUseCGA) { + DISPATCH_LOADS(true); + } else { + DISPATCH_LOADS(false); + } + +#undef DISPATCH_LOADS +#undef DISPATCH_QUANT +#undef LAUNCH_RMSNORM_KERNEL return cudaSuccess; } } // namespace trtllm_mnnvl_allreduce diff --git a/tests/comm/test_trtllm_mnnvl_allreduce.py b/tests/comm/test_trtllm_mnnvl_allreduce.py index e0f78775026..1e10e36df5f 100644 --- a/tests/comm/test_trtllm_mnnvl_allreduce.py +++ b/tests/comm/test_trtllm_mnnvl_allreduce.py @@ -6,9 +6,12 @@ import torch import torch.distributed as dist +import flashinfer.comm as comm import flashinfer.comm.trtllm_mnnvl_ar as trtllm_mnnvl_ar from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import TorchDistBackend +from flashinfer.fp4_quantization import _compute_swizzled_layout_sf_size +from flashinfer.utils import get_compute_capability # Use flashinfer.norm.rmsnorm as reference implementation. from flashinfer.norm import rmsnorm @@ -17,10 +20,64 @@ from tests.test_helpers.comm import ( init_torch_distributed_from_mpi, ) +from tests.test_helpers.utils_fp4 import ( + cast_from_fp4, + recover_swizzled_scales, + ref_fp4_quant, +) # Note: torch.distributed cleanup is handled by tests/comm/conftest.py +def fp8_quant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + finfo = torch.finfo(torch.float8_e4m3fn) + qinput = (x.float() * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max) + return qinput.to(torch.float8_e4m3fn) + + +def dequant(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return (x.to(torch.float32) * scale).to(dtype) + + +def _mnnvl_nvfp4_supported_on_all_ranks(device: torch.device) -> bool: + local_supported = get_compute_capability(device)[0] >= 10 + if not dist.is_initialized(): + return local_supported + + support_flags = [False] * dist.get_world_size() + dist.all_gather_object(support_flags, local_supported) + return all(support_flags) + + +def test_mnnvl_nvfp4_default_swizzled_scale_out_requires_padded_size(): + token_num = 1 + hidden_dim = 16 + linear_scale_size = token_num * hidden_dim // 16 + swizzled_scale_size = _compute_swizzled_layout_sf_size(token_num, hidden_dim // 16) + assert linear_scale_size < swizzled_scale_size + + input_tensor = torch.randn((token_num, hidden_dim), dtype=torch.float16) + residual = torch.randn_like(input_tensor) + gamma = torch.randn((hidden_dim,), dtype=torch.float16) + scale_out = torch.empty(linear_scale_size, dtype=torch.float8_e4m3fn) + + with pytest.raises( + ValueError, + match=( + f"scale_out is too small for NVFP4: got {linear_scale_size} " + f"elements, need at least {swizzled_scale_size}" + ), + ): + trtllm_mnnvl_ar.trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( + input=input_tensor, + residual_in=residual, + gamma=gamma, + workspace=object(), # type: ignore[arg-type] + scale_out=scale_out, + quant_type=trtllm_mnnvl_ar.MNNVLQuantType.NVFP4, + ) + + @torch.inference_mode() def row_linear_residual_norm_fusion_forward( x: torch.Tensor, @@ -496,6 +553,422 @@ def run_mnnvl_ar_full( dist.barrier() +def _prepare_quant_test_data( + seq_len: int, hidden_size: int, dtype: torch.dtype, eps: float +): + rank = dist.get_rank() + world_size = dist.get_world_size() + + if rank == 0: + x_full = torch.randn((world_size, seq_len, hidden_size), dtype=dtype) + residual = torch.randn((seq_len, hidden_size), dtype=dtype) + norm_weight = torch.randn((hidden_size,), dtype=dtype) + else: + x_full = None + residual = None + norm_weight = None + + data_list = [x_full, residual, norm_weight] + dist.broadcast_object_list(data_list, src=0) + x_full, residual, norm_weight = data_list + + x_full = x_full.cuda() + residual = residual.cuda() + norm_weight = norm_weight.cuda() + x_local = x_full[rank, :, :] + + allreduce_result = torch.sum(x_full, dim=0) + residual_out = allreduce_result + residual + norm_out = rmsnorm(residual_out, norm_weight, eps, enable_pdl=False) + return (x_local, residual, norm_weight), (norm_out, residual_out) + + +def _assert_quant_close( + quant_out: torch.Tensor, + scale_out: Optional[torch.Tensor], + ref_norm_out: torch.Tensor, + quant_type: int, + global_scale: torch.Tensor, + layout_code: int, +): + if quant_type == trtllm_mnnvl_ar.MNNVLQuantType.FP8: + ref_dequant = dequant( + fp8_quant(ref_norm_out, global_scale), global_scale, torch.float32 + ) + out_dequant = dequant(quant_out, global_scale, torch.float32) + mismatch_tol = 0.002 + else: + assert scale_out is not None + ref_fp4, ref_sf = ref_fp4_quant(ref_norm_out, global_scale, block_size=16) + if layout_code == comm.QuantizationSFLayout.SWIZZLED_128x4: + padded_rows = ((ref_norm_out.shape[0] + 127) // 128) * 128 + padded_cols = ((ref_norm_out.shape[1] // 16 + 3) // 4) * 4 + scale_out = recover_swizzled_scales( + scale_out.reshape(padded_rows, padded_cols), + ref_norm_out.shape[0], + ref_norm_out.shape[1], + 16, + ) + ref_dequant = ( + ref_fp4.to(torch.float32) + * ref_sf.repeat_interleave(16, dim=-1).to(torch.float32) + * global_scale.to(torch.float32) + ) + out_fp4 = cast_from_fp4(quant_out.view(torch.uint8)).to(ref_norm_out.device) + out_dequant = ( + out_fp4.to(torch.float32) + * scale_out.to(torch.float32).repeat_interleave(16, dim=-1) + * global_scale.to(torch.float32) + ) + mismatch_tol = 0.01 + + rtol, atol = 0.05, 0.15 + diff = torch.abs(out_dequant - ref_dequant) + max_diff = torch.maximum( + torch.abs(ref_dequant) * rtol, torch.tensor(atol, device=ref_dequant.device) + ) + mismatch_ratio = (diff > max_diff).sum().item() / out_dequant.numel() + assert mismatch_ratio <= mismatch_tol, ( + f"Mismatch ratio {mismatch_ratio:.4%} exceeds {mismatch_tol:.4%} threshold" + ) + + +# Quant fusion reuses the same MNNVL allreduce and RMSNorm paths that the +# exhaustive non-quant tests cover below. Keep this matrix pruned to the axes +# that are quant-specific or easy to regress: FP8/NVFP4, oneshot/twoshot, +# fp16/bf16, norm_out optionality, FP4 scale layout, and representative +# sequence/hidden sizes for small, mixed, wide, and large-token offsets. +MNNVL_QUANT_TEST_CASES = [ + pytest.param( + [1], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 2880, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-oneshot-fp16-h2880-s1", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 5120, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-twoshot-bf16-h5120-s4", + ), + pytest.param( + [15], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-twoshot-fp16-h7168-s15", + ), + pytest.param( + [27, 11, 24, 256], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-oneshot-bf16-h8192-mixed", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 16384, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-oneshot-fp16-h16384-s127", + ), + pytest.param( + [998, 2048], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-twoshot-bf16-h8192-large", + ), + pytest.param( + [1], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 2880, + [comm.QuantizationSFLayout.LINEAR, comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-oneshot-fp16-h2880-s1", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 5120, + [comm.QuantizationSFLayout.LINEAR, comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-twoshot-bf16-h5120-s4", + ), + pytest.param( + [15], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.LINEAR], + id="fp4-twoshot-fp16-h7168-s15", + ), + pytest.param( + [27, 11, 24, 256], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-oneshot-bf16-h8192-mixed", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 16384, + [comm.QuantizationSFLayout.LINEAR], + id="fp4-oneshot-fp16-h16384-s127", + ), + pytest.param( + [998, 2048], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-twoshot-bf16-h8192-large", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.bfloat16, + 2880, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="out-fp8-oneshot-bf16-h2880-s4", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="out-fp8-twoshot-fp16-h7168-s127", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 2880, + [comm.QuantizationSFLayout.LINEAR], + id="out-fp4-twoshot-bf16-h2880-s4", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="out-fp4-oneshot-fp16-h7168-s127", + ), +] + + +@pytest.mark.parametrize( + "seq_lens,pattern,strategy,dtype,hidden_size,layout_codes", + MNNVL_QUANT_TEST_CASES, +) +def test_mnnvl_allreduce_quant_unified( + seq_lens: list[int], + pattern: int, + strategy: trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy, + dtype: torch.dtype, + hidden_size: int, + layout_codes: list[int], +): + if torch.cuda.device_count() == 0: + pytest.skip("MNNVL quant test requires CUDA") + + init_torch_distributed_from_mpi() + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + pytest.skip(f"This test requires at least 2 ranks, got {world_size}") + + mapping = Mapping( + world_size=world_size, + rank=rank, + gpus_per_node=torch.cuda.device_count(), + tp_size=world_size, + ) + torch.cuda.set_device(mapping.local_rank) + is_fp4 = pattern in ( + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + ) + if is_fp4 and not _mnnvl_nvfp4_supported_on_all_ranks( + torch.device("cuda", torch.cuda.current_device()) + ): + pytest.skip("MNNVL NVFP4 quantization requires SM100+ on all ranks") + + comm_backend = TorchDistBackend() + eps = torch.finfo(dtype).eps + + workspace = trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace( + mapping, + max_num_tokens=max(seq_lens), + hidden_dim=hidden_size, + dtype=dtype, + comm_backend=comm_backend, + buffer_size_in_bytes=trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace.get_required_buffer_size_bytes( + world_size, max(seq_lens), hidden_size, dtype, strategy + ), + ) + try: + for seq_len in seq_lens: + (x_local, residual, norm_weight), (ref_norm, ref_residual) = ( + _prepare_quant_test_data(seq_len, hidden_size, dtype, eps) + ) + global_scale = torch.ones(1, dtype=torch.float32, device=x_local.device) + residual_out = torch.empty_like(x_local) + has_norm_out = pattern in ( + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant, + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + ) + norm_out = torch.empty_like(x_local) if has_norm_out else None + + for layout_code in layout_codes: + if is_fp4: + quant_type = trtllm_mnnvl_ar.MNNVLQuantType.NVFP4 + quant_out = torch.empty( + (seq_len, hidden_size // 2), + dtype=torch.uint8, + device=x_local.device, + ) + if layout_code == comm.QuantizationSFLayout.LINEAR: + scale_out = torch.empty( + (seq_len, hidden_size // 16), + dtype=torch.float8_e4m3fn, + device=x_local.device, + ) + else: + scale_out = torch.empty( + _compute_swizzled_layout_sf_size( + seq_len, hidden_size // 16 + ), + dtype=torch.float8_e4m3fn, + device=x_local.device, + ) + else: + quant_type = trtllm_mnnvl_ar.MNNVLQuantType.FP8 + quant_out = torch.empty_like(x_local, dtype=torch.float8_e4m3fn) + scale_out = None + + result = comm.allreduce_fusion( + input=x_local, + workspace=workspace, + pattern=pattern, + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + quant_out=quant_out, + scale_out=scale_out, + rms_gamma=norm_weight, + rms_eps=eps, + scale_factor=global_scale, + layout_code=layout_code, + use_oneshot=strategy + == trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + ) + assert result is quant_out + torch.testing.assert_close( + residual_out, ref_residual, rtol=0.05, atol=0.15 + ) + if has_norm_out: + assert norm_out is not None + torch.testing.assert_close(norm_out, ref_norm, rtol=0.05, atol=0.15) + _assert_quant_close( + quant_out, + scale_out, + ref_norm, + quant_type, + global_scale, + layout_code, + ) + dist.barrier() + finally: + workspace.destroy() + dist.barrier() + + +@pytest.mark.parametrize( + "strategy", + [ + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + ], +) +def test_mnnvl_nvfp4_rejects_fp32( + strategy: trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy, +): + if torch.cuda.device_count() == 0: + pytest.skip("MNNVL quant test requires CUDA") + + init_torch_distributed_from_mpi() + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + pytest.skip(f"This test requires at least 2 ranks, got {world_size}") + + mapping = Mapping( + world_size=world_size, + rank=rank, + gpus_per_node=torch.cuda.device_count(), + tp_size=world_size, + ) + torch.cuda.set_device(mapping.local_rank) + workspace = trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace( + mapping, + max_num_tokens=4, + hidden_dim=2048, + dtype=torch.float32, + comm_backend=TorchDistBackend(), + buffer_size_in_bytes=trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace.get_required_buffer_size_bytes( + world_size, 4, 2048, torch.float32, strategy + ), + ) + try: + x = torch.randn((4, 2048), dtype=torch.float32, device="cuda") + residual = torch.randn_like(x) + gamma = torch.randn((2048,), dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="NVFP4"): + comm.allreduce_fusion( + input=x, + workspace=workspace, + pattern=comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + residual_in=residual, + rms_gamma=gamma, + scale_factor=torch.ones(1, dtype=torch.float32, device="cuda"), + use_oneshot=strategy + == trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + ) + finally: + workspace.destroy() + dist.barrier() + + """Test with default workspace size""" # Multi-gpu test: mpirun -np 4 pytest tests/comm/test_trtllm_mnnvl_allreduce.py -vv -s