diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh index f76936782..6378825da 100644 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh +++ b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh @@ -118,7 +118,8 @@ cvt_fp16_to_fp4( uint32_t* output_scale_offset_by_experts, int32_t* mask, int n_experts, - bool low_latency) { + bool low_latency, + bool use_silu_and_mul) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using PackedVec = PackedVec; static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); @@ -127,11 +128,9 @@ cvt_fp16_to_fp4( // Input tensor row/col loops. int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; - // TODO(kaixih@nvidia): For now, we assume mask is used together with - // silu_and_mal. Maybe we want a more general behavior of mask later. In the - // silu case, the input last dim doubles. bool use_mask = mask != nullptr; - int actualColsPerRow = use_mask ? colsPerRow * 2 : colsPerRow; + // When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated). + int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { @@ -188,7 +187,7 @@ cvt_fp16_to_fp4( int64_t inOffset = rowIdx * actualColsPerRow + colIdx; PackedVec in_vec = reinterpret_cast(in)[inOffset]; - if (use_mask) { + if (use_mask || use_silu_and_mul) { PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; silu_and_mul(in_vec, in_vec_mul); } @@ -335,7 +334,8 @@ cvt_fp16_to_fp4( uint32_t* input_offset_by_experts, uint32_t* output_scale_offset_by_experts, int32_t* mask, - int n_experts) { + int n_experts, + bool use_silu_and_mul) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using PackedVec = PackedVec; static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); @@ -363,7 +363,8 @@ cvt_fp16_to_fp4( int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; bool use_mask = mask != nullptr; - int actualColsPerRow = use_mask ? colsPerRow * 2 : colsPerRow; + // When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated). + int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { @@ -402,7 +403,7 @@ cvt_fp16_to_fp4( int64_t inOffset = rowIdx * actualColsPerRow + colIdx; PackedVec in_vec = reinterpret_cast(in)[inOffset]; - if (use_mask) { + if (use_mask || use_silu_and_mul) { PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; silu_and_mul(in_vec, in_vec_mul); } @@ -488,7 +489,8 @@ void quant_impl( reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), - n_experts); + n_experts, + use_silu_and_mul); } else { cvt_fp16_to_fp4<<>>( m_topk, @@ -500,7 +502,8 @@ void quant_impl( reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), - n_experts); + n_experts, + use_silu_and_mul); } } else { if (n_experts >= 16) { @@ -515,7 +518,8 @@ void quant_impl( reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), n_experts, - /* bool low_latency */ true); + /* bool low_latency */ true, + use_silu_and_mul); } else { cvt_fp16_to_fp4<<>>( m_topk, @@ -528,7 +532,8 @@ void quant_impl( reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), n_experts, - /* bool low_latency */ true); + /* bool low_latency */ true, + use_silu_and_mul); } } } @@ -710,3 +715,92 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a( stream); } } + +void silu_and_mul_scaled_fp4_experts_quant_packed_sm100a( + tvm::ffi::TensorView output, + tvm::ffi::TensorView output_scale, + tvm::ffi::TensorView input, + tvm::ffi::TensorView input_global_scale, + tvm::ffi::TensorView input_offset_by_experts, + tvm::ffi::TensorView output_scale_offset_by_experts) { + auto MTopK = SymbolicSize{"m_topk"}; + auto KBy2 = SymbolicSize{"k_by_2"}; + auto OutputCols = SymbolicSize{"output_cols"}; + auto OutputScaleRows = SymbolicSize{"output_scale_rows"}; + auto OutputScaleCols = SymbolicSize{"output_scale_cols"}; + auto NExperts = SymbolicSize{"n_experts"}; + auto OffsetSize = SymbolicSize{"offset_size"}; + auto device = SymbolicDevice{}; + + TensorMatcher({MTopK, KBy2}) // + .with_dtype() + .template with_device(device) + .verify(input); + TensorMatcher({MTopK, OutputCols}) // + .with_dtype() + .with_device(device) + .verify(output); + TensorMatcher({OutputScaleRows, OutputScaleCols}) // + .with_dtype() + .with_device(device) + .verify(output_scale); + TensorMatcher({NExperts}) // + .with_dtype() + .with_device(device) + .verify(input_global_scale); + TensorMatcher({OffsetSize}) // + .with_dtype() + .with_device(device) + .verify(input_offset_by_experts) + .verify(output_scale_offset_by_experts); + + const int device_id = input.device().device_id; + RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); + + const int BLOCK_SIZE = 16; + const auto m_topk = static_cast(MTopK.unwrap()); + const auto k_by_2 = static_cast(KBy2.unwrap()); + // Input last dim is 2*k (gate+up concatenated). The kernel does SiLU(gate)*up + // then FP4-quantizes the k-dim result. + RuntimeCheck(k_by_2 % 2 == 0, "input last dim must be even (2*k)"); + const int k = k_by_2 / 2; + RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); + const auto n_experts = static_cast(NExperts.unwrap()); + const auto offset_size = static_cast(OffsetSize.unwrap()); + RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch"); + RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); + const int scales_k = k / BLOCK_SIZE; + const int padded_k = (scales_k + 3) / 4 * 4; + RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); + + const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); + if (host::is_type(input.dtype())) { + quant_impl( + output.data_ptr(), + output_scale.data_ptr(), + input.data_ptr(), + input_global_scale.data_ptr(), + input_offset_by_experts.data_ptr(), + output_scale_offset_by_experts.data_ptr(), + nullptr, // mask + true, // use_silu_and_mul + m_topk, + k, + n_experts, + stream); + } else { + quant_impl<__nv_bfloat16>( + output.data_ptr(), + output_scale.data_ptr(), + input.data_ptr(), + input_global_scale.data_ptr(), + input_offset_by_experts.data_ptr(), + output_scale_offset_by_experts.data_ptr(), + nullptr, // mask + true, // use_silu_and_mul + m_topk, + k, + n_experts, + stream); + } +} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh index 29b06dfc0..f2bd76787 100644 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh +++ b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh @@ -38,6 +38,14 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a( tvm::ffi::TensorView mask, bool use_silu_and_mul); +void silu_and_mul_scaled_fp4_experts_quant_packed_sm100a( + tvm::ffi::TensorView output, + tvm::ffi::TensorView output_scale, + tvm::ffi::TensorView input, + tvm::ffi::TensorView input_global_scale, + tvm::ffi::TensorView input_offset_by_experts, + tvm::ffi::TensorView output_scale_offset_by_experts); + void scaled_fp4_quant( tvm::ffi::TensorView output, tvm::ffi::TensorView input, @@ -66,3 +74,14 @@ void silu_and_mul_scaled_fp4_experts_quant( bool use_silu_and_mul) { silu_and_mul_scaled_fp4_experts_quant_sm100a(output, output_scale, input, input_global_scale, mask, use_silu_and_mul); } + +void silu_and_mul_scaled_fp4_experts_quant_packed( + tvm::ffi::TensorView output, + tvm::ffi::TensorView output_scale, + tvm::ffi::TensorView input, + tvm::ffi::TensorView input_global_scale, + tvm::ffi::TensorView input_offset_by_experts, + tvm::ffi::TensorView output_scale_offset_by_experts) { + silu_and_mul_scaled_fp4_experts_quant_packed_sm100a( + output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts); +} diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py index ff3c20072..a5d4e9d20 100644 --- a/python/sglang/jit_kernel/nvfp4.py +++ b/python/sglang/jit_kernel/nvfp4.py @@ -90,6 +90,10 @@ def _jit_nvfp4_expert_quant_module() -> Module: "silu_and_mul_scaled_fp4_experts_quant", "silu_and_mul_scaled_fp4_experts_quant_sm100a", ), + ( + "silu_and_mul_scaled_fp4_experts_quant_packed", + "silu_and_mul_scaled_fp4_experts_quant_packed_sm100a", + ), ], extra_dependencies=["cutlass"], extra_cuda_cflags=_nvfp4_cuda_flags(), @@ -355,6 +359,95 @@ def scaled_fp4_experts_quant( return output, output_scales +@register_custom_op( + op_name="silu_and_mul_scaled_fp4_experts_quant_packed", + mutates_args=["output", "output_scales"], +) +def _silu_and_mul_scaled_fp4_experts_quant_packed_custom_op( + output: torch.Tensor, + output_scales: torch.Tensor, + input_tensor: torch.Tensor, + input_global_scale: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, +) -> None: + module = _jit_nvfp4_expert_quant_module() + module.silu_and_mul_scaled_fp4_experts_quant_packed( + output, + output_scales, + input_tensor, + input_global_scale, + expert_offsets, + blockscale_offsets, + ) + + +@debug_kernel_api +def silu_and_mul_scaled_fp4_experts_quant_packed( + input_tensor: torch.Tensor, + input_global_scale: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + topk: int, + expert_map: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused SiLU+mul then FP4 quant for packed MoE inputs (expert_offsets aware). + + Input shape is (m, 2*k) — gate+up concatenated. The kernel does SiLU(gate)*up + then FP4-quantizes the k-dim result. + """ + assert ( + input_tensor.ndim == 2 + ), f"input.ndim needs to be == 2, but got {input_tensor.ndim}." + if expert_map is not None: + m, k = input_tensor.shape + output_tensor_shape = (m * topk, k) + input_tensor = _shuffle_rows_torch( + input_tensor, expert_map, output_tensor_shape + ) + + m_numtopk, k_input_doubled = input_tensor.shape + k = k_input_doubled // 2 + + max_tokens_per_expert = int(os.environ.get("MODELOPT_MAX_TOKENS_PER_EXPERT", 65536)) + assert m_numtopk <= max_tokens_per_expert * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT({max_tokens_per_expert})" + f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" + " MODELOPT_MAX_TOKENS_PER_EXPERT to set this value." + ) + scales_k = k // 16 + padded_k_in_int32 = (scales_k + 3) // 4 + + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + if padded_k_in_int32 * 4 > scales_k: + output_scales = torch.zeros( + max_tokens_per_expert * topk, + padded_k_in_int32, + dtype=torch.int32, + device=input_tensor.device, + ) + else: + output_scales = torch.empty( + max_tokens_per_expert * topk, + padded_k_in_int32, + dtype=torch.int32, + device=input_tensor.device, + ) + + _silu_and_mul_scaled_fp4_experts_quant_packed_custom_op( + output, + output_scales, + input_tensor, + input_global_scale, + expert_offsets, + blockscale_offsets, + ) + output_scales = output_scales.view(torch.float8_e4m3fn) + return output, output_scales + + @register_custom_op( op_name="scaled_fp4_grouped_quant", mutates_args=["output", "output_scales"], diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index 05cfe00fc..0a0136316 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -23,6 +23,7 @@ if _is_cuda: from sglang.jit_kernel.nvfp4 import ( cutlass_fp4_group_mm, scaled_fp4_experts_quant, + silu_and_mul_scaled_fp4_experts_quant_packed, ) @@ -468,19 +469,15 @@ def cutlass_moe_fp4( ) del rep_a_fp4, rep_a_blockscale - # hidden size dimension is split to one half sized tensor. - intermediate = torch.empty( - (m_a * num_topk, w1_fp4.shape[1] // 2), device=device, dtype=out_dtype - ) - silu_and_mul(c1, intermediate) - - int_fp4, int_blockscale = scaled_fp4_experts_quant( - intermediate, + # fused: SiLU + mul then FP4 quant (expert-packed) + int_fp4, int_blockscale = silu_and_mul_scaled_fp4_experts_quant_packed( + c1, a2_gscale, params.expert_offsets, params.blockscale_offsets, num_topk, ) + c2 = cutlass_fp4_group_mm( int_fp4, w2_fp4, diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 59394c17a..69d979f0c 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -65,6 +65,8 @@ class MoeRunner: self.runner_core = None # FlashInfer CUTLASS only supports fused path elif runner_backend.is_flashinfer_mxfp4(): self.runner_core = None # FlashInfer MXFP4 only supports fused path + elif runner_backend.is_cutlass(): + self.runner_core = None # CUTLASS uses the direct cutlass_moe_fp4 path else: raise NotImplementedError(f"Unsupported runner backend: {runner_backend}") diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 5a9b62441..cb52e7672 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2268,7 +2268,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): if moe_runner_backend.is_flashinfer_cutlass(): import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401 - self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + # The plain CUTLASS backend uses the direct cutlass_moe_fp4 fused path + # (see apply()), not a registered MoeRunner fused func, so skip creating + # a MoeRunner for it -- constructing one would fail the fused-func check. + if not moe_runner_backend.is_cutlass(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) def apply( self, diff --git a/test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py b/test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py new file mode 100644 index 000000000..9be9ba92e --- /dev/null +++ b/test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit test for the fused JIT op ``silu_and_mul_scaled_fp4_experts_quant_packed`` +(introduced in PR #18612). + +On the CUTLASS NVFP4 MoE intermediate, the op fuses the previous two-step path + + intermediate = silu_and_mul(c1) # SiLU(gate) * up + fp4, sf = scaled_fp4_experts_quant(intermediate) # NVFP4 expert quant + +into a single kernel + + fp4, sf = silu_and_mul_scaled_fp4_experts_quant_packed(c1, ...) + +This test compares the fused op against that exact unfused +``silu_and_mul`` + ``scaled_fp4_experts_quant`` path **with uneven expert offsets**: +experts deliberately receive very different token counts, including tiny experts and +experts whose row count is not a multiple of the 128-row block-scale padding. That is +precisely the regime that stresses the per-expert ``expert_offsets`` / +``blockscale_offsets`` indexing the fusion has to get right. + +It follows the two existing siblings: + * ``test_silu_and_mul_quantize_to_fp4_grouped`` (the grouped/masked variant) -- the + unfused path is the reference, and the fused output must match it bit-exactly + (packed FP4 nibbles + recovered block scales), and + * ``test_nvfp4_blockwise_moe`` (the expert-offset variant) -- offsets are built from + an explicit, non-uniform per-expert token list. + +A high-precision ``F.silu(gate) * up`` check additionally grounds the unfused path so a +bug shared by both kernels cannot produce a false (vacuous) pass. + + pytest python/sglang/jit_kernel/tests/test_silu_and_mul_scaled_fp4_experts_quant_packed.py -v +""" + +import sys + +import pytest +import torch +import triton +from torch.nn import functional as F + +from sglang.jit_kernel.activation import silu_and_mul +from sglang.jit_kernel.nvfp4 import ( + scaled_fp4_experts_quant, + silu_and_mul_scaled_fp4_experts_quant_packed, +) +from sglang.test.ci.ci_register import register_cuda_ci + +# The NVFP4 expert-quant kernels are Blackwell-only (sm100a), so this runs on +# the B200 unit suite. +register_cuda_ci(est_time=20, suite="base-b-kernel-unit-1-gpu-b200") + +FLOAT8_E4M3_MAX = 448.0 +FLOAT4_E2M1_MAX = 6.0 +BLOCK_SIZE = 16 +kE2M1ToFloat = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 +) + + +def _nvfp4_supported() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) + + +def _round_up(x: int, y: int) -> int: + return ((x + y - 1) // y) * y + + +# --------------------------------------------------------------------------- # +# Offset builders (mirror test_nvfp4_blockwise_moe.py). +# expert_offsets: cumulative *actual* per-expert rows ([E+1] int32) +# blockscale_offsets: cumulative rows padded up to 128 per expert ([E+1] int32) +# A non-uniform ``m_per_expert`` makes both offset tensors uneven. +# --------------------------------------------------------------------------- # +def _build_expert_offsets(m_per_expert, device) -> torch.Tensor: + offsets = [0] + for m in m_per_expert: + offsets.append(offsets[-1] + m) + return torch.tensor(offsets, dtype=torch.int32, device=device) + + +def _build_blockscale_offsets(m_per_expert, device) -> torch.Tensor: + offsets = [0] + for m in m_per_expert: + offsets.append(offsets[-1] + _round_up(m, 128)) + return torch.tensor(offsets, dtype=torch.int32, device=device) + + +# --------------------------------------------------------------------------- # +# FP4 dequant / scale-recovery helpers (mirror test/registered/kernels/test_fp4_moe.py) +# --------------------------------------------------------------------------- # +def break_fp4_bytes(a: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + assert a.dtype == torch.uint8 + m, n = a.shape + a_flat = a.flatten() + high = (a_flat & 0xF0) >> 4 + low = a_flat & 0x0F + combined = torch.stack((low, high), dim=1).flatten() + signs = (combined & 0x08).to(torch.bool) + abs_vals = (combined & 0x07).to(torch.long) + kE2M1 = kE2M1ToFloat.to(device=a.device) + values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) + return values.reshape(m, n * 2).to(dtype=dtype) + + +def convert_swizzled_to_linear( + a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int +) -> torch.Tensor: + """De-swizzle one expert's block-scale region and drop the 128-row padding tail.""" + m_tiles = (m + 128 - 1) // 128 + f = block_size * 4 + k_tiles = (k + f - 1) // f + tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4)) + tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) + out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size) + return out[0:m, 0:k] + + +def dequantize_nvfp4_to_dtype( + tensor_fp4: torch.Tensor, + tensor_sf: torch.Tensor, + global_scale: torch.Tensor, + dtype: torch.dtype, + device: torch.device, + block_size: int = 16, +) -> torch.Tensor: + """Dequantize one expert's packed FP4 (m, k//2) + swizzled block scales.""" + assert tensor_fp4.dtype == torch.uint8 + m, packed_k = tensor_fp4.shape + k = packed_k * 2 + tensor_f32 = break_fp4_bytes(tensor_fp4, dtype) + tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) + tensor_sf = tensor_sf.view(torch.float8_e4m3fn) + tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) + tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale + out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) + return out.to(dtype=dtype) + + +def _recover_block_scales( + sf: torch.Tensor, s0: int, s1: int, m_e: int, n: int +) -> torch.Tensor: + """De-swizzled, un-padded block scales (float32) for one expert's region.""" + block = sf[s0:s1].contiguous().view(torch.float8_e4m3fn) + return convert_swizzled_to_linear(block, m_e, n, BLOCK_SIZE).to(torch.float32) + + +def _rel_l2(a: torch.Tensor, b: torch.Tensor) -> float: + return (a.float() - b.float()).norm().item() / b.float().norm().clamp_min( + 1e-9 + ).item() + + +# --------------------------------------------------------------------------- # +# Uneven per-expert token counts. Each list is deliberately NON-uniform so that +# expert_offsets / blockscale_offsets are uneven, exercising: +# * tiny experts (1, 5, 7 tokens), +# * an exactly-128 expert (no padding), +# * experts straddling the 128-row block-scale padding (130, 200, 384). +# --------------------------------------------------------------------------- # +UNEVEN_M_PER_EXPERT = [ + [33, 17, 48, 29], # all < 128 (matches test_nvfp4_blockwise_moe) + [1, 128, 200, 5, 64], # tiny + exactly-128 + cross-128 + [130, 1, 384, 17, 96, 7], # heavy skew, large dynamic range +] +NS = [256, 768] # 768 == Qwen3-30B-A3B moe_intermediate_size +DTYPES = [torch.bfloat16, torch.float16] + + +@pytest.mark.skipif( + not _nvfp4_supported(), + reason="NVFP4 fused expert-quant kernel requires compute capability >= 10.0 (B200/SM100).", +) +@pytest.mark.parametrize("m_per_expert", UNEVEN_M_PER_EXPERT) +@pytest.mark.parametrize("n", NS) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_fused_matches_unfused_uneven_offsets(m_per_expert, n, dtype): + torch.manual_seed(0) + device = torch.device("cuda") + num_experts = len(m_per_expert) + + # --- uneven expert offsets --- + expert_offsets = _build_expert_offsets(m_per_expert, device) + blockscale_offsets = _build_blockscale_offsets(m_per_expert, device) + total_m = int(expert_offsets[-1].item()) + counts = torch.tensor(m_per_expert) + assert ( + counts.max() >= 2 * counts.min() + ), "expert offsets must be uneven for this test" + + # gate+up concatenated input (m, 2n); /5 keeps values in a sane FP4 range. + c1 = torch.randn((total_m, 2 * n), dtype=dtype, device=device) / 5.0 + gate, up = c1[:, :n].float(), c1[:, n:].float() + ref = F.silu(gate) * up # high-precision SiLU(gate) * up, (total_m, n) fp32 + + # Per-expert global scale, exactly like cutlass_moe builds a2_gscale. + gscale = torch.empty(num_experts, dtype=torch.float32, device=device) + for e in range(num_experts): + r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1]) + amax = ref[r0:r1].abs().max().clamp_min(1e-6) + gscale[e] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax + + # topk only gates the buffer-size assertion in the wrapper; the per-expert + # layout is driven entirely by the offsets. Both paths use the same value. + topk = 1 + + # ---- fused (new op) ---- + fused_fp4, fused_sf = silu_and_mul_scaled_fp4_experts_quant_packed( + c1, gscale, expert_offsets, blockscale_offsets, topk + ) + + # ---- unfused (the exact path the op replaced) ---- + intermediate = torch.empty((total_m, n), dtype=dtype, device=device) + silu_and_mul(c1, intermediate) + unf_fp4, unf_sf = scaled_fp4_experts_quant( + intermediate, gscale, expert_offsets, blockscale_offsets, topk + ) + + assert fused_fp4.shape == unf_fp4.shape == (total_m, n // 2) + + # Per-expert, bit-exact comparison honoring the uneven offsets. + for e in range(num_experts): + r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1]) + s0, s1 = int(blockscale_offsets[e]), int(blockscale_offsets[e + 1]) + m_e = r1 - r0 + + # (1) Packed FP4 nibbles are identical: same NVFP4 quantizer, same fused + # SiLU(gate)*up rounded to the storage dtype before quantization. + torch.testing.assert_close( + fused_fp4[r0:r1], unf_fp4[r0:r1], msg=f"FP4 bytes differ for expert {e}" + ) + + # (2) Recovered (de-swizzled, un-padded) block scales are identical. + torch.testing.assert_close( + _recover_block_scales(fused_sf, s0, s1, m_e, n), + _recover_block_scales(unf_sf, s0, s1, m_e, n), + msg=f"block scales differ for expert {e}", + ) + + # (3) Grounding: the unfused path really reproduces SiLU(gate)*up within FP4 + # error, so (1)/(2) cannot pass vacuously on a bug shared by both kernels. + deq = dequantize_nvfp4_to_dtype( + unf_fp4[r0:r1].contiguous(), + unf_sf[s0:s1].contiguous(), + gscale[e], + dtype, + device, + BLOCK_SIZE, + ) + assert ( + _rel_l2(deq, ref[r0:r1]) < 0.2 + ), f"expert {e}: unfused dequant does not match SiLU(gate)*up reference" + + +# --------------------------------------------------------------------------- # +# Performance. The fusion removes, on the MoE down-projection input, one +# intermediate buffer allocation, one extra kernel launch, and a full HBM +# round-trip of the SiLU(gate)*up result. The speedup is measured under CUDA +# graphs -- the steady-state GPU memory-traffic saving, matching how SGLang +# executes graphed decode (the credible, low-noise number; eager wall-clock is +# dominated by launch/dispatch overhead and is too noisy to assert on). The +# assert is only a conservative regression floor; the printed speedup is the real +# result. Mirrors test_cutedsl_gdn_performance, in the same kernel unit suite. +# +# Tokens are spread evenly across experts here (the representative throughput +# case) at realistic Qwen3-30B-A3B MoE dims (n=768, 128 experts), swept from a +# decode batch up to a prefill chunk; the uneven-offset corner cases are covered +# by the correctness test above. +# --------------------------------------------------------------------------- # +PERF_SHAPES = [ + (1024, 768, 128), + (4096, 768, 128), + (16384, 768, 128), +] + + +def _even_offsets(total_tokens, num_experts, device): + base, rem = divmod(total_tokens, num_experts) + m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)] + return ( + _build_expert_offsets(m_per_expert, device), + _build_blockscale_offsets(m_per_expert, device), + ) + + +@pytest.mark.skipif( + not _nvfp4_supported(), + reason="NVFP4 fused expert-quant kernel requires compute capability >= 10.0 (B200/SM100).", +) +@pytest.mark.parametrize("total_tokens,n,num_experts", PERF_SHAPES) +@torch.inference_mode() +def test_fused_perf_not_regressed(total_tokens, n, num_experts): + device = torch.device("cuda") + dtype = torch.bfloat16 + expert_offsets, blockscale_offsets = _even_offsets( + total_tokens, num_experts, device + ) + c1 = torch.randn((total_tokens, 2 * n), dtype=dtype, device=device) / 5.0 + gscale = torch.empty(num_experts, dtype=torch.float32, device=device) + for e in range(num_experts): + r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1]) + amax = c1[r0:r1].abs().max().to(torch.float32).clamp_min(1e-6) + gscale[e] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax + topk = 1 + + def fused(): + silu_and_mul_scaled_fp4_experts_quant_packed( + c1, gscale, expert_offsets, blockscale_offsets, topk + ) + + def unfused(): + # The exact path the op replaced: alloc the intermediate, SiLU*mul into + # it, then quantize it -- one extra buffer + kernel + HBM round-trip. + intermediate = torch.empty((total_tokens, n), dtype=dtype, device=device) + silu_and_mul(c1, intermediate) + scaled_fp4_experts_quant( + intermediate, gscale, expert_offsets, blockscale_offsets, topk + ) + + g_f = triton.testing.do_bench_cudagraph(fused) # ms, median + g_u = triton.testing.do_bench_cudagraph(unfused) + cuda_graph_speedup = g_u / g_f + print( + f"\n [PERF] tokens={total_tokens:>6} n={n} E={num_experts}: " + f"unfused {g_u * 1e3:6.1f}us fused {g_f * 1e3:6.1f}us " + f"cuda-graph speedup = {cuda_graph_speedup:.2f}x" + ) + # Regression guard only: the fusion must not make this op slower. The actual + # win carries a wide margin over this floor, so shared-runner noise cannot + # flake it. + assert ( + cuda_graph_speedup >= 1.05 + ), f"fused regressed under cuda-graph: {cuda_graph_speedup:.2f}x" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))