[Perf] Tune the W4AFP8 DeepEP low-latency requant launch geometry (#35760)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/ops/moe/ep_moe_kernels.py

@@ -1,4 +1,5 @@
 import logging
+from functools import lru_cache
 from typing import Optional, Tuple
 
 import torch
@@ -2085,6 +2086,45 @@ def silu_and_mul_masked_post_per_tensor_quant_fwd(
     return output
 
 
+@triton.jit
+def _requant_row(
+    x_ptr,
+    x_scale_ptr,
+    x_scale_stride0,
+    x_scale_stride1,
+    output_ptr,
+    m,
+    k,
+    expert,
+    row,
+    output_scale_val_inv,
+    k_offsets,
+    scale_g_offsets,
+    g_mask,
+    HAS_G_TAIL: tl.constexpr,
+):
+    """Requantize one row; shared by both phases so they write rows identically."""
+    row_base = expert.to(tl.int64) * m + row
+    x_ptrs = x_ptr + row_base * k + k_offsets
+    output_ptrs = output_ptr + row_base * k + k_offsets
+    x_scale_ptrs = (
+        x_scale_ptr + expert * x_scale_stride0 + row * x_scale_stride1 + scale_g_offsets
+    )
+    if HAS_G_TAIL:
+        hidden = tl.load(x_ptrs, mask=g_mask[:, None], other=0.0)
+        group_scale = tl.load(x_scale_ptrs, mask=g_mask, other=0.0)
+    else:
+        hidden = tl.load(x_ptrs)
+        group_scale = tl.load(x_scale_ptrs)
+    scaled = hidden.to(tl.float32) * group_scale.to(tl.float32)[:, None]
+    scaled = scaled * output_scale_val_inv
+    quantized = scaled.to(output_ptr.dtype.element_ty)
+    if HAS_G_TAIL:
+        tl.store(output_ptrs, quantized, mask=g_mask[:, None])
+    else:
+        tl.store(output_ptrs, quantized)
+
+
 @triton.jit
 def _fp8_per_token_quant_to_per_tensor_quant_kernel(
     x_ptr,
@@ -2097,52 +2137,154 @@ def _fp8_per_token_quant_to_per_tensor_quant_kernel(
     output_ptr,
     m,
     k,
+    num_experts,
+    row_cap,
     K_SCALE_BLOCK_SIZE: tl.constexpr,
-    K_BLOCK_SIZE: tl.constexpr,
+    G_BLOCK_SIZE: tl.constexpr,
-    HAS_K_TAIL: tl.constexpr,
+    HAS_G_TAIL: tl.constexpr,
+    EXPERT_BLOCK: tl.constexpr,
 ):
-    pid_k, pid_m, pid_e = (
+    pid_g, pid_m, pid_e = (
         tl.program_id(axis=0),
         tl.program_id(axis=1),
         tl.program_id(axis=2),
     )
-    pid_m_dim = tl.num_programs(1)
+    m_grid = tl.num_programs(1)
 
-    token_id = pid_m
-    last_effective_id = tl.load(masked_m_ptr + pid_e)
-
-    if token_id >= last_effective_id:
-        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
+    # Tile whole scale groups: one scalar scale load per group.  DeepEP scales
-    output_ptrs = output_ptr + pid_e * m * k + k_offsets
+    # are column-major in the last two dims, so element-axis loads would gather.
-    x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets
+    g_offsets = pid_g * G_BLOCK_SIZE + tl.arange(0, G_BLOCK_SIZE)
+    k_offsets = (
+        g_offsets[:, None] * K_SCALE_BLOCK_SIZE
+        + tl.arange(0, K_SCALE_BLOCK_SIZE)[None, :]
+    )
+    g_mask = g_offsets < k // K_SCALE_BLOCK_SIZE
+    scale_g_offsets = g_offsets * x_scale_stride2
 
-    for tok_idx in tl.range(token_id, last_effective_id, pid_m_dim):
+    # Phase 1: this expert's rows below row_cap, strided over the m-grid.
-        if HAS_K_TAIL:
+    last_effective_id = tl.load(masked_m_ptr + pid_e)
-            hidden = tl.load(x_ptrs + tok_idx * k, mask=k_mask, other=0.0)
+    for row in tl.range(pid_m, min(last_effective_id, row_cap), m_grid):
-            x_scale = tl.load(
+        _requant_row(
-                x_scale_ptrs + tok_idx * x_scale_stride1, mask=k_mask, other=0.0
+            x_ptr,
-            )
+            x_scale_ptr,
-        else:
+            x_scale_stride0,
-            hidden = tl.load(x_ptrs + tok_idx * k)
+            x_scale_stride1,
-            x_scale = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1)
+            output_ptr,
-        hidden = hidden.to(tl.float32)
+            m,
-        scale_fp32 = x_scale.to(tl.float32)
+            k,
-        hidden = hidden * scale_fp32 * output_scale_val_inv
+            pid_e,
-        quantized = hidden.to(output_ptr.dtype.element_ty)
+            row,
-        if HAS_K_TAIL:
+            output_scale_val_inv,
-            tl.store(output_ptrs + tok_idx * k, quantized, mask=k_mask)
+            k_offsets,
-        else:
+            scale_g_offsets,
-            tl.store(output_ptrs + tok_idx * k, quantized)
+            g_mask,
+            HAS_G_TAIL,
+        )
+
+    # Phase 2: rows above row_cap are shared across the whole launch, so a hot
+    # expert cannot serialize; a batch with no overflow pays one reduction here.
+    expert_ids = tl.arange(0, EXPERT_BLOCK)
+    counts = tl.load(masked_m_ptr + expert_ids, mask=expert_ids < num_experts, other=0)
+    overflow = tl.maximum(counts - row_cap, 0)
+    total_overflow = tl.sum(overflow)
+    if total_overflow == 0:
+        return
+
+    # The inclusive prefix sum maps flat index i to (expert, row): the owner is
+    # however many experts finish at or before i; zero-overflow experts drop out.
+    overflow_before = tl.cumsum(overflow)
+    flat_id = pid_e * m_grid + pid_m
+    num_programs = m_grid * num_experts
+    for i in tl.range(flat_id, total_overflow, num_programs):
+        expert = tl.sum((overflow_before <= i).to(tl.int32))
+        started = tl.max(tl.where(overflow_before <= i, overflow_before, 0))
+        _requant_row(
+            x_ptr,
+            x_scale_ptr,
+            x_scale_stride0,
+            x_scale_stride1,
+            output_ptr,
+            m,
+            k,
+            expert,
+            row_cap + (i - started),
+            output_scale_val_inv,
+            k_offsets,
+            scale_g_offsets,
+            g_mask,
+            HAS_G_TAIL,
+        )
+
+
+# Tuned in bytes per lane, not elements: warp width differs by vendor, and
+# 16 B/lane measured best on both H200 (2048 elems) and MI350X (4096).
+# Below _REQUANT_MANY_EXPERTS the grid underfills NVIDIA parts and a half
+# tile buys k-block parallelism; that costs MI350X up to 5% there.
+_REQUANT_BYTES_PER_LANE = 16
+_REQUANT_BYTES_PER_LANE_FEW_EXPERTS = 8
+_REQUANT_MANY_EXPERTS = 32
+_REQUANT_NUM_WARPS = 4
+_REQUANT_DEFAULT_WARP_SIZE = 32
+_REQUANT_M_GRID_MAX = 32
+_REQUANT_M_GRID_MIN = 4
+# Program target on the (m-grid x expert) plane while rows are scarce; measured,
+# and deliberately not scaled to core count (8 per core was worse on MI350X).
+_REQUANT_TARGET_PROGRAMS = 1024
+# Past this many rows per expert the capped-away programs would carry real work.
+_REQUANT_ROWS_SATURATED = 64
+# Rows past slack * expected_rows go to the shared phase.  2x keeps ordinary
+# variation per-expert; measured 4% at even load and removes the skew regression.
+_REQUANT_ROW_CAP_SLACK = 2
+
+
+def _floor_pow2(value: int) -> int:
+    return 1 << (max(1, value).bit_length() - 1)
+
+
+@lru_cache(maxsize=None)
+def requant_warp_size(device: torch.device) -> int:
+    """Lanes per warp, which sets the tile width the requant launches."""
+    return torch.cuda.get_device_properties(device).warp_size
+
+
+def requant_launch_geometry(
+    num_groups: int,
+    num_experts: int,
+    group_size: int = 128,
+    expected_rows: Optional[int] = None,
+    warp_size: int = _REQUANT_DEFAULT_WARP_SIZE,
+    max_rows: int = 1 << 30,
+) -> Tuple[int, int, int]:
+    """Pick (groups per program, m-grid, row cap) for the requant.
+
+    All three are launch hints: any values produce the same bytes.  The row
+    estimate rounds down to a power of two because ``dispatch_a`` reports
+    ``(rows + num_experts) // num_experts``, one high at exact averages.
+    ``warp_size`` scales the tile to keep bytes per lane constant.
+    """
+    # The payload is fp8, so a byte per lane is an element per lane.
+    bytes_per_lane = (
+        _REQUANT_BYTES_PER_LANE
+        if num_experts >= _REQUANT_MANY_EXPERTS
+        else _REQUANT_BYTES_PER_LANE_FEW_EXPERTS
+    )
+    tile_elems = bytes_per_lane * _REQUANT_NUM_WARPS * warp_size
+    # Clamp to the payload.  Non-pow2 group counts (40, 48) leave the last tile
+    # partly masked, up to 8% behind a narrower tile on MI350X; accepted, since
+    # per-width constants only moved the loss.
+    g_block = min(_floor_pow2(tile_elems // group_size), _floor_pow2(num_groups))
+    if expected_rows is None:
+        # Nothing to place a cap against, so leave every row with its own expert.
+        return g_block, _REQUANT_M_GRID_MAX, max_rows
+    m_grid = min(_REQUANT_M_GRID_MAX, _floor_pow2(expected_rows))
+    if expected_rows < _REQUANT_ROWS_SATURATED:
+        m_grid = min(
+            m_grid, _floor_pow2(_REQUANT_TARGET_PROGRAMS // max(1, num_experts))
+        )
+    row_cap = min(max_rows, max(1, expected_rows) * _REQUANT_ROW_CAP_SLACK)
+    return g_block, max(_REQUANT_M_GRID_MIN, m_grid), row_cap
 
 
 def fp8_per_token_to_per_tensor_quant_triton(
@@ -2151,15 +2293,35 @@ def fp8_per_token_to_per_tensor_quant_triton(
     masked_m: torch.Tensor,
     output_scale: torch.Tensor,
     output: torch.Tensor,
+    expected_rows: Optional[int] = None,
 ):
+    # The 2-D tile indexes within a group via tl.arange, so the group width
+    # must be a power of two.
     K_SCALE_BLOCK_SIZE = 128
     assert len(x.shape) == 3 and x.size(2) % K_SCALE_BLOCK_SIZE == 0
     assert x.is_contiguous()
+    assert output.shape == x.shape and output.is_contiguous()
+    # Addressing flattens (expert, row) by raw strides; a shape mismatch reads
+    # out of bounds rather than failing.
+    assert masked_m.shape[0] == x.size(0)
+    assert x_scale.size(0) == x.size(0) and x_scale.size(1) == x.size(1)
     assert x_scale.size(2) == x.size(2) // K_SCALE_BLOCK_SIZE
+    # Under `use_ue8m0` DeepEP returns int32-packed UE8M0 scales; reinterpreting
+    # those as fp32 would quantize against garbage.
+    assert x_scale.dtype == torch.float32
     assert output_scale.numel() == 1
 
-    K_BLOCK_SIZE = 1024
+    num_experts = x.size(0)
-    grid = (triton.cdiv(x.size(2), K_BLOCK_SIZE), 32, x.size(0))
+    num_groups = x.size(2) // K_SCALE_BLOCK_SIZE
+    g_block, m_grid, row_cap = requant_launch_geometry(
+        num_groups=num_groups,
+        num_experts=num_experts,
+        group_size=K_SCALE_BLOCK_SIZE,
+        expected_rows=expected_rows,
+        warp_size=requant_warp_size(x.device),
+        max_rows=x.size(1),
+    )
+    grid = (triton.cdiv(num_groups, g_block), m_grid, num_experts)
     _fp8_per_token_quant_to_per_tensor_quant_kernel[grid](
         x,
         x_scale,
@@ -2169,10 +2331,13 @@ def fp8_per_token_to_per_tensor_quant_triton(
         output,
         x.size(1),
         x.size(2),
+        num_experts,
+        row_cap,
         K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE,
-        K_BLOCK_SIZE=K_BLOCK_SIZE,
+        G_BLOCK_SIZE=g_block,
-        HAS_K_TAIL=x.size(2) % K_BLOCK_SIZE != 0,
+        HAS_G_TAIL=(num_groups % g_block != 0),
-        num_warps=8,
+        EXPERT_BLOCK=triton.next_power_of_2(num_experts),
+        num_warps=_REQUANT_NUM_WARPS,
     )
 
 

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -456,6 +456,7 @@ def cutlass_w4a8_moe_deepep_ll(
     problem_sizes2: torch.Tensor,
     a1_scale: Optional[torch.Tensor] = None,
     a2_scale: Optional[torch.Tensor] = None,
+    expected_m: Optional[int] = None,
 ) -> torch.Tensor:
     """
     This function computes a w4a8-quantized Mixture of Experts (MoE) layer
@@ -492,6 +493,8 @@ def cutlass_w4a8_moe_deepep_ll(
         Shape: scalar or [1, N]
     - apply_router_weight_on_input (bool): When true, the topk weights are
         applied directly on the inputs. This is only applicable when topk is 1.
+    - expected_m (Optional[int]): Dispatcher's expected rows per expert; a
+        requant launch hint only, any value is correct.
 
     Returns:
     - torch.Tensor: The fp8 output tensor after applying the MoE layer.
@@ -532,6 +535,7 @@ def cutlass_w4a8_moe_deepep_ll(
         masked_m=masked_m,
         output_scale=a1_scale,
         output=gateup_input,
+        expected_rows=expected_m,
     )
     c1 = torch.empty((num_experts, m, n * 2), device=device, dtype=torch.bfloat16)
     c2 = torch.empty((num_experts, m, k), device=device, dtype=torch.bfloat16)

python/sglang/srt/layers/quantization/w4afp8.py

@@ -342,7 +342,9 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         layer: DeepEPMoE,
         dispatch_output: DeepEPLLDispatchOutput,
     ) -> torch.Tensor:
-        hidden_states, hidden_scales, topk_ids, _, masked_m, _ = dispatch_output
+        hidden_states, hidden_scales, topk_ids, _, masked_m, expected_m = (
+            dispatch_output
+        )
 
         if hidden_scales is None:
             raise RuntimeError(
@@ -376,6 +378,7 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
             layer.quant_method.problem_sizes2,
             layer.w13_input_scale,
             layer.w2_input_scale,
+            expected_m=expected_m,
         )
 
         return output

test/registered/kernels/benchmark/moe/bench_fp8_per_token_to_per_tensor_quant.py

@@ -0,0 +1,135 @@
+"""Benchmark the W4AFP8 DeepEP low-latency requant against its previous geometry.
+
+``legacy-geometry`` launches the current kernel with its previous launch
+parameters (1024-element tile, 8 warps, 32 programs per expert); ``tuned`` goes
+through the wrapper.  ``skew`` concentrates rows on one hot expert, which the
+m-grid cannot see because ``expected_m`` is a dispatch-wide average.
+"""
+
+import torch
+import triton
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.ops.moe.ep_moe_kernels import (
+    _fp8_per_token_quant_to_per_tensor_quant_kernel,
+    fp8_per_token_to_per_tensor_quant_triton,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=45, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+FP8 = torch.float8_e4m3fn
+K_SCALE_BLOCK_SIZE = 128
+LEGACY_G_BLOCK = 8  # 1024 hidden elements
+LEGACY_WARPS = 8
+LEGACY_M_GRID = 32
+
+
+def _expected_m(num_experts, dispatched_rows):
+    """``dispatch_a`` reports ``(rows + num_experts) // num_experts``, one high
+    at exact averages; benchmark with what production would pass."""
+    return (dispatched_rows + num_experts) // num_experts
+
+
+def _row_counts(num_experts, rows, skew):
+    """One hot expert at ``skew * rows``, the rest share the fixed remainder; a
+    dispatch redistributes rows, so skew cannot exceed ``num_experts``."""
+    if skew == 1:
+        return [rows] * num_experts
+    total = num_experts * rows
+    counts = [(total - rows * skew) // (num_experts - 1)] * num_experts
+    counts[0] = rows * skew
+    return counts
+
+
+def _build(num_experts, m, k, rows, skew=1):
+    x = (torch.randn(num_experts, m, k, device="cuda") * 4).to(FP8)
+    # DeepEP returns the last two scale dims column-major (for TMA).
+    x_scale = (
+        torch.rand(num_experts, m, k // K_SCALE_BLOCK_SIZE, device="cuda")
+        .add_(0.5)
+        .permute(0, 2, 1)
+        .contiguous()
+        .permute(0, 2, 1)
+    )
+    counts = _row_counts(num_experts, rows, skew)
+    masked_m = torch.tensor(counts, dtype=torch.int32, device="cuda")
+    output_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
+    output = torch.empty((num_experts, m, k), dtype=FP8, device="cuda")
+    return (
+        x,
+        x_scale,
+        masked_m,
+        output_scale,
+        output,
+        _expected_m(num_experts, sum(counts)),
+    )
+
+
+def _tuned(x, x_scale, masked_m, output_scale, output, expected_m):
+    fp8_per_token_to_per_tensor_quant_triton(
+        x=x,
+        x_scale=x_scale,
+        masked_m=masked_m,
+        output_scale=output_scale,
+        output=output,
+        expected_rows=expected_m,
+    )
+    return output
+
+
+def _legacy_geometry(x, x_scale, masked_m, output_scale, output, expected_m):
+    num_groups = x.size(2) // K_SCALE_BLOCK_SIZE
+    grid = (triton.cdiv(num_groups, LEGACY_G_BLOCK), LEGACY_M_GRID, x.size(0))
+    _fp8_per_token_quant_to_per_tensor_quant_kernel[grid](
+        x,
+        x_scale,
+        *x_scale.stride(),
+        masked_m,
+        output_scale,
+        output,
+        x.size(1),
+        x.size(2),
+        x.size(0),
+        # row_cap = m keeps every row on its own expert, as the old launch did.
+        x.size(1),
+        K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE,
+        G_BLOCK_SIZE=LEGACY_G_BLOCK,
+        HAS_G_TAIL=(num_groups % LEGACY_G_BLOCK != 0),
+        EXPERT_BLOCK=triton.next_power_of_2(x.size(0)),
+        num_warps=LEGACY_WARPS,
+    )
+    return output
+
+
+FN_MAP = {"tuned": _tuned, "legacy-geometry": _legacy_geometry}
+
+# (hidden, local experts, padded rows): DeepSeek-V3 at EP8, then a 3584 hidden
+# size at a low and a high local-expert count.
+SHAPES = [(7168, 8, 1024), (3584, 8, 1024), (3584, 56, 256)]
+
+
+@marker.parametrize("hidden,num_experts,m", SHAPES, [(7168, 8, 1024), (3584, 56, 256)])
+@marker.parametrize("rows", [8, 32, 128, 256], [8, 32, 256])
+@marker.parametrize("skew", [1, 4, 16], [1, 16])
+@marker.benchmark("impl", ["tuned", "legacy-geometry"])
+def benchmark(hidden: int, num_experts: int, m: int, rows: int, skew: int, impl: str):
+    if skew > num_experts:
+        marker.skip("one expert cannot hold more than the whole dispatch")
+    if rows * skew > m:
+        marker.skip("more live rows than the payload holds")
+    args = _build(num_experts, m, hidden, rows, skew)
+    return marker.do_bench(
+        FN_MAP[impl],
+        input_args=args,
+        graph_clone_args=(0, 1),
+        memory_args=None,
+        # Tensor-size bandwidth would be off by the padding factor.
+        disable_log_bandwidth=True,
+    )
+
+
+if __name__ == "__main__":
+    benchmark.run()

test/registered/kernels/ops/moe/test_fp8_per_token_to_per_tensor_quant.py

@@ -1,22 +1,21 @@
-"""Unit test for ``fp8_per_token_to_per_tensor_quant_triton`` across hidden sizes.
+"""Unit test for ``fp8_per_token_to_per_tensor_quant_triton``.
 
-W4AFP8 DeepEP low-latency requantizes the fp8 dispatch payload with this kernel
+The hidden size is only guaranteed to be a multiple of the scale group (128),
-before the first CUTLASS grouped GEMM.  The payload's hidden size is only
+rows past ``masked_m`` must stay untouched, and every launch geometry must
-guaranteed to be a multiple of the fp8 scale-group size (128) -- e.g. 3584 for
+produce the same bytes.
-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
+import triton
 
 from sglang.kernels.ops.moe.ep_moe_kernels import (
+    _fp8_per_token_quant_to_per_tensor_quant_kernel,
     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")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 dev = "cuda"
 FP8 = torch.float8_e4m3fn
@@ -27,15 +26,17 @@ SENTINEL = 0.375
 OUTPUT_SCALE = 2.0
 
 
-def _build(num_experts, m, k, seed):
+def _build(num_experts, m, k, seed, column_major_scales=False):
     g = torch.Generator(device="cpu").manual_seed(seed)
-    # Integers in [-8, 8] with power-of-two per-token-group scales keep every
+    # Integers in [-8, 8] with power-of-two group scales keep every intermediate
-    # intermediate exactly representable in e4m3, so the reference below matches
+    # exactly representable in e4m3, so the reference matches bit-for-bit.
-    # 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())
+    x_scale = torch.pow(2.0, exps.float()).to(dev)
-    return x.to(dev).to(FP8), x_scale.to(dev)
+    if column_major_scales:
+        # DeepEP returns the last two scale dims column-major (for TMA).
+        x_scale = x_scale.permute(0, 2, 1).contiguous().permute(0, 2, 1)
+    return x.to(dev).to(FP8), x_scale
 
 
 def _ref(x, x_scale):
@@ -43,14 +44,23 @@ def _ref(x, x_scale):
     return (dequant * (1.0 / OUTPUT_SCALE)).to(FP8)
 
 
-# 7168: exact multiple of K_BLOCK_SIZE (the DeepSeek-V3 hidden size).
+def _assert_output(output, x, x_scale, masked):
-# 3584 / 1152: only 128-aligned, so the last k block is partially masked.
+    ref = _ref(x, x_scale)
-@pytest.mark.parametrize("k", [7168, 3584, 1152])
+    for e, valid in enumerate(masked):
-def test_masked_rows_and_k_tail(k):
+        torch.testing.assert_close(
-    num_experts, m = 4, 48
+            output[e, :valid].float(), ref[e, :valid].float(), rtol=0, atol=0
-    masked = [0, 1, 17, m]
+        )
+        # Rows past masked_m 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
+        )
 
-    x, x_scale = _build(num_experts, m, k, seed=k)
+
+def _run_and_check(num_experts, m, k, masked, expected_rows, column_major_scales):
+    x, x_scale = _build(
+        num_experts, m, k, seed=k + num_experts, column_major_scales=column_major_scales
+    )
     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)
@@ -61,19 +71,100 @@ def test_masked_rows_and_k_tail(k):
         masked_m=masked_m,
         output_scale=output_scale,
         output=output,
+        expected_rows=expected_rows,
     )
 
-    ref = _ref(x, x_scale)
+    _assert_output(output, x, x_scale, masked)
-    for e, valid in enumerate(masked):
+
-        torch.testing.assert_close(
+
-            output[e, :valid].float(), ref[e, :valid].float(), rtol=0, atol=0
+# 7168: fills every tile of scale groups (the DeepSeek-V3 hidden size).
-        )
+# 3584 / 1152: only 128-aligned, so the last tile is partially masked.
-        # Padding rows are not part of any expert's GEMM problem size and must
+@pytest.mark.parametrize("k", [7168, 3584, 1152])
-        # stay as the caller left them.
+# None: shape-independent grid; 4 and 64: both ends of the m-grid heuristic.
-        padding = output[e, valid:].float()
+@pytest.mark.parametrize("expected_rows", [None, 4, 64])
-        torch.testing.assert_close(
+@pytest.mark.parametrize("column_major_scales", [False, True])
-            padding, torch.full_like(padding, SENTINEL), rtol=0, atol=0
+def test_masked_rows_and_group_tail(k, expected_rows, column_major_scales):
-        )
+    _run_and_check(
+        num_experts=4,
+        m=48,
+        k=k,
+        masked=[0, 1, 17, 48],
+        expected_rows=expected_rows,
+        column_major_scales=column_major_scales,
+    )
+
+
+# Row estimates chosen so the expert-count cap binds: 40 experts cap at 16
+# programs, 128 at 8; uncapped these would be 32 and 16.
+@pytest.mark.parametrize("num_experts,expected_rows", [(40, 32), (128, 16)])
+def test_many_experts(num_experts, expected_rows):
+    m = 48
+    # Every expert gets a different row count, as a real dispatch would.
+    masked = [(e * 7) % (m + 1) for e in range(num_experts)]
+    _run_and_check(
+        num_experts=num_experts,
+        m=m,
+        k=3584,
+        masked=masked,
+        expected_rows=expected_rows,
+        column_major_scales=True,
+    )
+
+
+# The wrapper only launches the running vendor's tile, so drive the kernel
+# directly across every width either vendor can pick, plus the degenerate
+# single-group tile.
+@pytest.mark.parametrize("g_block", [1, 8, 16, 32])
+@pytest.mark.parametrize("k", [7168, 3584])
+@pytest.mark.parametrize("m_grid", [1, 4, 32])
+# 0 sends every row to the shared overflow path, 48 keeps every row on its own
+# expert, and 4 splits the batch across both.
+@pytest.mark.parametrize("row_cap", [0, 4, 48])
+def test_every_launch_geometry_agrees(g_block, k, m_grid, row_cap):
+    num_experts, m = 4, 48
+    masked = [0, 1, 17, 48]
+    x, x_scale = _build(num_experts, m, k, seed=k + g_block, column_major_scales=True)
+    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)
+
+    num_groups = k // K_SCALE_BLOCK_SIZE
+    grid = (triton.cdiv(num_groups, g_block), m_grid, num_experts)
+    _fp8_per_token_quant_to_per_tensor_quant_kernel[grid](
+        x,
+        x_scale,
+        *x_scale.stride(),
+        masked_m,
+        output_scale,
+        output,
+        m,
+        k,
+        num_experts,
+        row_cap,
+        K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE,
+        G_BLOCK_SIZE=g_block,
+        HAS_G_TAIL=(num_groups % g_block != 0),
+        EXPERT_BLOCK=triton.next_power_of_2(num_experts),
+        num_warps=4,
+    )
+
+    _assert_output(output, x, x_scale, masked)
+
+
+# Experts with no live rows must be stepped over by the prefix-sum mapping,
+# the case most likely to be off by one.
+@pytest.mark.parametrize(
+    "masked", [[0, 0, 0, 0], [0, 5, 0, 7], [9, 0, 0, 0], [0, 0, 0, 9]]
+)
+def test_experts_with_no_rows_are_skipped(masked):
+    _run_and_check(
+        num_experts=4,
+        m=48,
+        k=3584,
+        masked=masked,
+        expected_rows=2,
+        column_major_scales=True,
+    )
 
 
 if __name__ == "__main__":

test/registered/unit/layers/moe/test_w4afp8_requant_geometry.py

@@ -0,0 +1,145 @@
+"""CPU tests for the W4AFP8 low-latency requant launch geometry."""
+
+import unittest
+
+from sglang.kernels.ops.moe.ep_moe_kernels import requant_launch_geometry
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+DSV3_GROUPS = 7168 // 128  # 56
+K3_GROUPS = 3584 // 128  # 28
+PREVIOUS_FIXED_M_GRID = 32
+ROW_CAP_SLACK = 2
+
+
+class TestRequantLaunchGeometry(CustomTestCase):
+    def test_cap_leaves_ordinary_variation_to_the_owning_expert(self):
+        """Rows below the cap stay on their expert; the shared path costs a
+        lookup per row and only pays under real imbalance."""
+        for expected_rows in (1, 4, 16, 64, 256):
+            _, _, row_cap = requant_launch_geometry(
+                DSV3_GROUPS, 64, expected_rows=expected_rows
+            )
+            self.assertGreaterEqual(row_cap, expected_rows * ROW_CAP_SLACK)
+
+    def test_cap_never_exceeds_the_payload(self):
+        """A cap past the padded rows would leave the shared path unreachable."""
+        for max_rows in (1, 8, 128):
+            for expected_rows in (1, 64, 4096):
+                _, _, row_cap = requant_launch_geometry(
+                    DSV3_GROUPS, 64, expected_rows=expected_rows, max_rows=max_rows
+                )
+                self.assertLessEqual(row_cap, max_rows)
+
+    def test_unknown_row_count_keeps_every_row_with_its_expert(self):
+        """With no estimate there is nothing to place a cap against."""
+        for num_experts in (8, 56):
+            _, m_grid, row_cap = requant_launch_geometry(
+                K3_GROUPS, num_experts, max_rows=128
+            )
+            self.assertEqual(m_grid, PREVIOUS_FIXED_M_GRID)
+            self.assertEqual(row_cap, 128)
+
+    def test_m_grid_never_exceeds_the_previous_fixed_grid(self):
+        """The estimate only ever shrinks the grid, so no batch can regress."""
+        for expected_rows in (1, 8, 32, 33, 1024):
+            for num_experts in (8, 56, 256):
+                _, m_grid, _ = requant_launch_geometry(
+                    DSV3_GROUPS, num_experts, expected_rows=expected_rows
+                )
+                self.assertLessEqual(m_grid, PREVIOUS_FIXED_M_GRID)
+
+    def test_m_grid_shrinks_as_the_expert_axis_fills_the_grid(self):
+        """Hundreds of experts already saturate the grid without 32 programs each."""
+        scarce = [
+            requant_launch_geometry(DSV3_GROUPS, num_experts, expected_rows=32)[1]
+            for num_experts in (8, 64, 128, 256)
+        ]
+        self.assertEqual(scarce, sorted(scarce, reverse=True))
+        self.assertEqual(scarce[0], PREVIOUS_FIXED_M_GRID)
+        self.assertLess(scarce[-1], scarce[0])
+
+    def test_expert_cap_lifts_once_rows_carry_the_work(self):
+        """Past the row threshold the extra programs are not just early exits."""
+        for num_experts in (8, 128, 512):
+            _, m_grid, _ = requant_launch_geometry(
+                DSV3_GROUPS, num_experts, expected_rows=1024
+            )
+            self.assertEqual(m_grid, PREVIOUS_FIXED_M_GRID)
+
+    def test_dispatcher_round_up_does_not_bump_the_grid(self):
+        """dispatch_a reports (rows + num_experts) // num_experts, one high;
+        rounding up as well would double the launch at every power of two."""
+        for rows in (4, 8, 16, 32):
+            exact = requant_launch_geometry(DSV3_GROUPS, 8, expected_rows=rows)[1]
+            reported = requant_launch_geometry(DSV3_GROUPS, 8, expected_rows=rows + 1)[
+                1
+            ]
+            self.assertEqual(reported, exact, f"rows={rows}")
+
+    def test_m_grid_is_bounded_and_monotonic(self):
+        previous = 0
+        for expected_rows in range(1, 512):
+            _, m_grid, _ = requant_launch_geometry(
+                K3_GROUPS, 8, expected_rows=expected_rows
+            )
+            # The floor keeps a one-row batch from serializing an expert into
+            # one program (measured several times slower).
+            self.assertGreaterEqual(m_grid, 4)
+            self.assertLessEqual(m_grid, PREVIOUS_FIXED_M_GRID)
+            self.assertGreaterEqual(m_grid, previous)
+            previous = m_grid
+
+    def test_tile_never_exceeds_the_payload(self):
+        """A 512-wide hidden size is 4 groups; a wider tile would be mostly masked."""
+        for num_groups in (1, 4, 12):
+            for num_experts in (8, 56):
+                g_block, _, _ = requant_launch_geometry(
+                    num_groups, num_experts, expected_rows=64
+                )
+                self.assertLessEqual(g_block, num_groups)
+
+    def test_tile_holds_bytes_per_lane_across_warp_widths(self):
+        """The tuned unit is bytes per lane: 2048 elements at warp 32 must become
+        4096 at warp 64, or a wave64 part gets half the bytes per lane."""
+        for warp_size, want_elems in ((32, 2048), (64, 4096)):
+            for group_size in (64, 128, 256, 512):
+                g_block, _, _ = requant_launch_geometry(
+                    num_groups=7168 // group_size,
+                    num_experts=56,
+                    group_size=group_size,
+                    expected_rows=16,
+                    warp_size=warp_size,
+                )
+                self.assertEqual(
+                    g_block * group_size, want_elems, f"warp_size={warp_size}"
+                )
+
+    def test_few_experts_halve_the_tile_on_either_warp_width(self):
+        """A grid too small to fill the part buys k-blocks by halving the tile."""
+        for warp_size, want_elems in ((32, 1024), (64, 2048)):
+            g_block, _, _ = requant_launch_geometry(
+                DSV3_GROUPS, 8, expected_rows=16, warp_size=warp_size
+            )
+            self.assertEqual(g_block * 128, want_elems, f"warp_size={warp_size}")
+
+    def test_warp_width_does_not_move_the_m_grid(self):
+        """The two knobs are independent: the m-grid answers to rows and experts."""
+        for num_experts in (8, 56, 256):
+            for expected_rows in (1, 8, 32, 1024):
+                grids = {
+                    requant_launch_geometry(
+                        DSV3_GROUPS,
+                        num_experts,
+                        expected_rows=expected_rows,
+                        warp_size=warp_size,
+                    )[1]
+                    for warp_size in (32, 64)
+                }
+                self.assertEqual(len(grids), 1, f"E={num_experts} rows={expected_rows}")
+
+
+if __name__ == "__main__":
+    unittest.main()