diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 0cebb241b..bc8aaba8d 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -1944,6 +1944,7 @@ def _fp8_per_token_quant_to_per_tensor_quant_kernel( k, K_SCALE_BLOCK_SIZE: tl.constexpr, K_BLOCK_SIZE: tl.constexpr, + HAS_K_TAIL: tl.constexpr, ): pid_k, pid_m, pid_e = ( tl.program_id(axis=0), @@ -1959,6 +1960,11 @@ def _fp8_per_token_quant_to_per_tensor_quant_kernel( return output_scale_val_inv = 1.0 / tl.load(output_scale_ptr).to(tl.float32) k_offsets = pid_k * K_BLOCK_SIZE + tl.arange(0, K_BLOCK_SIZE) + # k only has to be a multiple of the 128-wide scale group (e.g. 3584), so the + # last k block can be partial. Specialize on it: hidden sizes that fill + # every block keep the unmasked loads, and their codegen is unchanged. + if HAS_K_TAIL: + k_mask = k_offsets < k scale_offsets = (k_offsets // K_SCALE_BLOCK_SIZE) * x_scale_stride2 x_ptrs = x_ptr + pid_e * m * k + k_offsets @@ -1966,10 +1972,22 @@ def _fp8_per_token_quant_to_per_tensor_quant_kernel( x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets for tok_idx in tl.range(token_id, last_effective_id, pid_m_dim): - hidden = tl.load(x_ptrs + tok_idx * k).to(tl.float32) - scale_fp32 = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1).to(tl.float32) + if HAS_K_TAIL: + hidden = tl.load(x_ptrs + tok_idx * k, mask=k_mask, other=0.0) + x_scale = tl.load( + x_scale_ptrs + tok_idx * x_scale_stride1, mask=k_mask, other=0.0 + ) + else: + hidden = tl.load(x_ptrs + tok_idx * k) + x_scale = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1) + hidden = hidden.to(tl.float32) + scale_fp32 = x_scale.to(tl.float32) hidden = hidden * scale_fp32 * output_scale_val_inv - tl.store(output_ptrs + tok_idx * k, hidden.to(output_ptr.dtype.element_ty)) + quantized = hidden.to(output_ptr.dtype.element_ty) + if HAS_K_TAIL: + tl.store(output_ptrs + tok_idx * k, quantized, mask=k_mask) + else: + tl.store(output_ptrs + tok_idx * k, quantized) def fp8_per_token_to_per_tensor_quant_triton( @@ -1986,8 +2004,7 @@ def fp8_per_token_to_per_tensor_quant_triton( assert output_scale.numel() == 1 K_BLOCK_SIZE = 1024 - assert x.size(2) % K_BLOCK_SIZE == 0 - grid = (x.size(2) // K_BLOCK_SIZE, 32, x.size(0)) + grid = (triton.cdiv(x.size(2), K_BLOCK_SIZE), 32, x.size(0)) _fp8_per_token_quant_to_per_tensor_quant_kernel[grid]( x, x_scale, @@ -1999,6 +2016,7 @@ def fp8_per_token_to_per_tensor_quant_triton( x.size(2), K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE, K_BLOCK_SIZE=K_BLOCK_SIZE, + HAS_K_TAIL=x.size(2) % K_BLOCK_SIZE != 0, num_warps=8, ) diff --git a/test/registered/kernels/ops/moe/test_fp8_per_token_to_per_tensor_quant.py b/test/registered/kernels/ops/moe/test_fp8_per_token_to_per_tensor_quant.py new file mode 100644 index 000000000..17215dfa8 --- /dev/null +++ b/test/registered/kernels/ops/moe/test_fp8_per_token_to_per_tensor_quant.py @@ -0,0 +1,82 @@ +"""Unit test for ``fp8_per_token_to_per_tensor_quant_triton`` across hidden sizes. + +W4AFP8 DeepEP low-latency requantizes the fp8 dispatch payload with this kernel +before the first CUTLASS grouped GEMM. The payload's hidden size is only +guaranteed to be a multiple of the fp8 scale-group size (128) -- e.g. 3584 for +Kimi-K3 -- so the kernel must handle a ``k`` tail that does not fill a whole +``K_BLOCK_SIZE`` (1024) block, and must still leave the rows past ``masked_m`` +untouched. +""" + +import pytest +import torch + +from sglang.kernels.ops.moe.ep_moe_kernels import ( + fp8_per_token_to_per_tensor_quant_triton, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +dev = "cuda" +FP8 = torch.float8_e4m3fn +K_SCALE_BLOCK_SIZE = 128 +# Every value the kernel can produce below is a multiple of 0.25, so this +# sentinel cannot be matched by a kernel that wrongly writes a padding row. +SENTINEL = 0.375 +OUTPUT_SCALE = 2.0 + + +def _build(num_experts, m, k, seed): + g = torch.Generator(device="cpu").manual_seed(seed) + # Integers in [-8, 8] with power-of-two per-token-group scales keep every + # intermediate exactly representable in e4m3, so the reference below matches + # bit-for-bit regardless of the rounding mode of the final cast. + x = torch.randint(-8, 9, (num_experts, m, k), generator=g).float() + exps = torch.randint(-1, 2, (num_experts, m, k // K_SCALE_BLOCK_SIZE), generator=g) + x_scale = torch.pow(2.0, exps.float()) + return x.to(dev).to(FP8), x_scale.to(dev) + + +def _ref(x, x_scale): + dequant = x.float() * x_scale.repeat_interleave(K_SCALE_BLOCK_SIZE, dim=2) + return (dequant * (1.0 / OUTPUT_SCALE)).to(FP8) + + +# 7168: exact multiple of K_BLOCK_SIZE (the DeepSeek-V3 hidden size). +# 3584 / 1152: only 128-aligned, so the last k block is partially masked. +@pytest.mark.parametrize("k", [7168, 3584, 1152]) +def test_masked_rows_and_k_tail(k): + num_experts, m = 4, 48 + masked = [0, 1, 17, m] + + x, x_scale = _build(num_experts, m, k, seed=k) + masked_m = torch.tensor(masked, dtype=torch.int32, device=dev) + output_scale = torch.tensor([OUTPUT_SCALE], dtype=torch.float32, device=dev) + output = torch.full((num_experts, m, k), SENTINEL, device=dev).to(FP8) + + fp8_per_token_to_per_tensor_quant_triton( + x=x, + x_scale=x_scale, + masked_m=masked_m, + output_scale=output_scale, + output=output, + ) + + ref = _ref(x, x_scale) + for e, valid in enumerate(masked): + torch.testing.assert_close( + output[e, :valid].float(), ref[e, :valid].float(), rtol=0, atol=0 + ) + # Padding rows are not part of any expert's GEMM problem size and must + # stay as the caller left them. + padding = output[e, valid:].float() + torch.testing.assert_close( + padding, torch.full_like(padding, SENTINEL), rtol=0, atol=0 + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-s"]))