diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 7cc615b06..20263ae3a 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/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, - HAS_K_TAIL: tl.constexpr, + G_BLOCK_SIZE: 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 - output_ptrs = output_ptr + pid_e * m * k + k_offsets - x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets + # Tile whole scale groups: one scalar scale load per group. DeepEP scales + # are column-major in the last two dims, so element-axis loads would gather. + 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): - 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 - 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) + # Phase 1: this expert's rows below row_cap, strided over the m-grid. + last_effective_id = tl.load(masked_m_ptr + pid_e) + for row in tl.range(pid_m, min(last_effective_id, row_cap), m_grid): + _requant_row( + x_ptr, + x_scale_ptr, + x_scale_stride0, + x_scale_stride1, + output_ptr, + m, + k, + pid_e, + row, + output_scale_val_inv, + k_offsets, + scale_g_offsets, + 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 - grid = (triton.cdiv(x.size(2), K_BLOCK_SIZE), 32, x.size(0)) + num_experts = 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, - HAS_K_TAIL=x.size(2) % K_BLOCK_SIZE != 0, - num_warps=8, + G_BLOCK_SIZE=g_block, + HAS_G_TAIL=(num_groups % g_block != 0), + EXPERT_BLOCK=triton.next_power_of_2(num_experts), + num_warps=_REQUANT_NUM_WARPS, ) diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 3cc372057..0efbd7878 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/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) diff --git a/python/sglang/srt/layers/quantization/w4afp8.py b/python/sglang/srt/layers/quantization/w4afp8.py index cdf5b1e52..05ac19a51 100644 --- a/python/sglang/srt/layers/quantization/w4afp8.py +++ b/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 diff --git a/test/registered/kernels/benchmark/moe/bench_fp8_per_token_to_per_tensor_quant.py b/test/registered/kernels/benchmark/moe/bench_fp8_per_token_to_per_tensor_quant.py new file mode 100644 index 000000000..750006712 --- /dev/null +++ b/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() 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 index 17215dfa8..dcbe1eeb6 100644 --- 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 @@ -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 -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. +The hidden size is only guaranteed to be a multiple of the scale group (128), +rows past ``masked_m`` must stay untouched, and every launch geometry must +produce the same bytes. """ 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 - # intermediate exactly representable in e4m3, so the reference below matches - # bit-for-bit regardless of the rounding mode of the final cast. + # Integers in [-8, 8] with power-of-two group scales keep every intermediate + # exactly representable in e4m3, so the reference matches bit-for-bit. 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) + x_scale = torch.pow(2.0, exps.float()).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). -# 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] +def _assert_output(output, x, x_scale, masked): + 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 + ) + # 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) - 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 - ) + _assert_output(output, x, x_scale, masked) + + +# 7168: fills every tile of scale groups (the DeepSeek-V3 hidden size). +# 3584 / 1152: only 128-aligned, so the last tile is partially masked. +@pytest.mark.parametrize("k", [7168, 3584, 1152]) +# None: shape-independent grid; 4 and 64: both ends of the m-grid heuristic. +@pytest.mark.parametrize("expected_rows", [None, 4, 64]) +@pytest.mark.parametrize("column_major_scales", [False, True]) +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__": diff --git a/test/registered/unit/layers/moe/test_w4afp8_requant_geometry.py b/test/registered/unit/layers/moe/test_w4afp8_requant_geometry.py new file mode 100644 index 000000000..f0872ffa1 --- /dev/null +++ b/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()