minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Fix] Support 128-aligned hidden sizes in the W4AFP8 DeepEP low-latency requant kernel (#35593)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alex Nails
2026-08-19 22:18:17 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent f744607567
commit b6dcd393d6
2 changed files with 105 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/moe/ep_moe_kernels.py
+23 -5
View File
@@ -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,
)
test/registered/kernels/ops/moe/test_fp8_per_token_to_per_tensor_quant.py
+82
View File
@@ -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"]))