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,
)