From fad225396dbf7f5f515bb6d818606e8c1a4cefc5 Mon Sep 17 00:00:00 2001 From: Xiangyi Zhang Date: Sun, 12 Jul 2026 22:45:30 -0700 Subject: [PATCH 1/3] perf(deepseek-v4): select two-stage fused mHC launch by token count The vendored fused_hc op exposes backend/tile_n/num_k_splits launch parameters that the wrapper previously hardcoded: allinone-fma for M<=32 and allinone-mma above. A 768-combination sweep on B200 (graph-replay timed, parity-checked against the composed two-stage path) found: * backend=1 (fma_ksplit + mhcBigFuseKernel), tile_n=2, num_k_splits=2 runs 15.2us/call at M=32 vs 22.2us for the current default (-31%). * the allinone-mma backend selected for M>32 runs 206-241us at every M; the routing cliff fixed by the previous commit is a bad backend choice, not a large-M property. backend=1 with tile_n=4 stays flat (18-24us). * the k-split backend accumulates num_k_splits x M partial rows into the y_acc/r_acc workspaces. The fused workspace sized them at M rows, so the spill corrupted small-M outputs and crashed a production p8 run with cudaErrorIllegalAddress; accumulators now reserve FUSED_HC_MAX_K_SPLITS x max_bs rows (a few KB per buffer set). Certified on DeepSeek-V4-Pro TP8 SWE-smith golden p8 (b200-79, order- reversed pairs vs a 7-run baseline plateau): TPOT -4~6%, steady throughput +6~8%, acceptance rate unchanged, 189/189 completions in all runs. An exact decode trace shows the mHC kernel category alone moving (3.43 -> 2.87 ms/step) with every other category flat. Signed-off-by: Xiangyi Zhang --- .../runtime/layers/deepseek_v4_mhc.py | 11 +- test/runtime/test_deepseek_v4_mhc_routing.py | 19 ++ .../tokenspeed_kernel/ops/mhc/__init__.py | 2 + .../tokenspeed_kernel/ops/mhc/trtllm.py | 48 ++++- .../test/test_deepseek_v4_mhc_fused_launch.py | 175 ++++++++++++++++++ 5 files changed, 247 insertions(+), 8 deletions(-) create mode 100644 tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py diff --git a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py index 7236d5532..e9934c1e5 100644 --- a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py +++ b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py @@ -11,6 +11,7 @@ import triton import triton.language as tl from tokenspeed_kernel.ops.mhc import ( + FUSED_HC_MAX_K_SPLITS, deep_gemm_mhc_prenorm_gemm, has_deep_gemm_mhc, supports_trtllm_mhc, @@ -399,6 +400,10 @@ def __init__( ): n2 = hc_mult * hc_mult shape_n = hc_mult * (2 + hc_mult) + # The k-split fused_hc backend accumulates num_k_splits x num_tokens + # partial rows into y_acc/r_acc, so the accumulator workspaces need + # FUSED_HC_MAX_K_SPLITS x max_bs rows, not max_bs like allinone. + acc_rows = FUSED_HC_MAX_K_SPLITS * max_bs self.bufs = tuple( ( torch.empty( @@ -407,9 +412,9 @@ def __init__( torch.empty(max_bs, hc_mult, dtype=torch.float32, device=device), torch.empty(max_bs, n2, dtype=torch.float32, device=device), torch.empty(max_bs, hidden_size, dtype=torch.bfloat16, device=device), - torch.empty(max_bs, shape_n, dtype=torch.float32, device=device), - torch.empty(max_bs, dtype=torch.float32, device=device), - torch.empty(max_bs, dtype=torch.int32, device=device), + torch.empty(acc_rows, shape_n, dtype=torch.float32, device=device), + torch.empty(acc_rows, dtype=torch.float32, device=device), + torch.empty(acc_rows, dtype=torch.int32, device=device), ) for _ in range(2) ) diff --git a/test/runtime/test_deepseek_v4_mhc_routing.py b/test/runtime/test_deepseek_v4_mhc_routing.py index a9c1ef461..e54b20f8a 100644 --- a/test/runtime/test_deepseek_v4_mhc_routing.py +++ b/test/runtime/test_deepseek_v4_mhc_routing.py @@ -110,3 +110,22 @@ def test_fused_hc_routing_logs_once_per_shape(caplog): _call(m) routing_logs = [r for r in caplog.records if "fused_hc routing" in r.getMessage()] assert len(routing_logs) == 1 + + +def test_fused_workspace_accumulators_sized_for_k_splits(): + # The k-split fused_hc backend writes FUSED_HC_MAX_K_SPLITS x num_tokens + # partial rows into y_acc/r_acc. Accumulators sized to max_bs rows + # overflow into neighboring allocations (observed as a production + # cudaErrorIllegalAddress before the sizing was fixed). + from tokenspeed_kernel.ops.mhc import FUSED_HC_MAX_K_SPLITS + + max_bs, hidden = 32, 64 + pp = mhc._FusedHcPingPong(max_bs, HC, hidden, torch.device("cpu")) + for bufs in pp.bufs: + residual_cur, post_cur, comb_cur, layer_input, y_acc, r_acc, dc = bufs + assert residual_cur.shape[0] == max_bs + assert layer_input.shape[0] == max_bs + expected = FUSED_HC_MAX_K_SPLITS * max_bs + assert y_acc.shape[0] == expected + assert r_acc.shape[0] == expected + assert dc.shape[0] == expected diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py index c44c7b6d3..c0964fba8 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/__init__.py @@ -25,6 +25,7 @@ has_deep_gemm_mhc, ) from tokenspeed_kernel.ops.mhc.trtllm import ( + FUSED_HC_MAX_K_SPLITS, has_trtllm_mhc, supports_trtllm_mhc, trtllm_mhc_big_fuse, @@ -33,6 +34,7 @@ ) __all__ = [ + "FUSED_HC_MAX_K_SPLITS", "deep_gemm_mhc_prenorm_gemm", "has_deep_gemm_mhc", "has_trtllm_mhc", diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py index 31995de27..77a25c600 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py @@ -88,6 +88,47 @@ def supports_trtllm_mhc( return _is_blackwell_device(device_index) +# fused_hc launch table, selected by a B200 sweep over the vendored kernel's +# backend/tile/k-split parameters (all 768 combinations, graph-replay timed, +# parity-checked against the composed post_mapping + prenorm-GEMM + big_fuse +# path). backend=1 is the two-stage fma_ksplit + big_fuse organization; it +# beats the allinone default by ~31% at M=32 (15.2us vs 22.2us) and stays +# flat past the M>32 allinone-mma cliff. Two constraints shape the table: +# * backend=1 accumulates num_k_splits x M partial rows into the y_acc and +# r_acc workspaces; callers must size them accordingly (see +# FUSED_HC_MAX_K_SPLITS). Undersized accumulators corrupt small-M +# outputs and crash at scale. +# * small token counts stay on the allinone-fma default, which is at +# least as fast there. +_FUSED_HC_SMALL_M_MAX = 12 +_FUSED_HC_MEDIUM_M_MAX = 32 + +FUSED_HC_MAX_K_SPLITS = 2 +"""Largest ``num_k_splits`` the launch table may select. + +The two-stage k-split backend writes ``num_k_splits * num_tokens`` partial +rows into the ``y_acc``/``r_acc`` accumulator workspaces, so callers must +allocate ``FUSED_HC_MAX_K_SPLITS * max_tokens`` rows for them. +""" + + +def _select_fused_hc_launch(num_tokens: int) -> tuple[int, int, int]: + """Pick the fused_hc backend and tile configuration for a token count. + + Args: + num_tokens: Number of tokens (rows) in the fused mHC call. + + Returns: + Tuple of ``(backend, tile_n, num_k_splits)`` kernel launch parameters. + """ + + if num_tokens <= _FUSED_HC_SMALL_M_MAX: + return 3, 1, 1 # fused_all_fma (allinone) default + if num_tokens <= _FUSED_HC_MEDIUM_M_MAX: + return 1, 2, 2 # fma_ksplit + big_fuse + return 1, 4, 2 + + trtllm_mhc_big_fuse = error_fn trtllm_mhc_fused_hc = error_fn trtllm_mhc_post_mapping = error_fn @@ -268,10 +309,7 @@ def trtllm_mhc_fused_hc( """ num_tokens, hc_mult, hidden_size = residual_prev.shape - if num_tokens <= 32: - backend, tile_n = 3, 1 # fused_all_fma - else: - backend, tile_n = 2, 0 # fused_all_mma + backend, tile_n, num_k_splits = _select_fused_hc_launch(num_tokens) _mhc_fused_hc( x_prev, residual_prev, @@ -297,7 +335,7 @@ def trtllm_mhc_fused_hc( sinkhorn_iters, backend, tile_n, - 1, + num_k_splits, 0, 1, None, diff --git a/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py new file mode 100644 index 000000000..41cd27b76 --- /dev/null +++ b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026 LightSeek Foundation + +"""Launch-table and GPU parity tests for the fused mHC backend selection. + +The launch table was measured on B200 (768-combination sweep, graph-replay +timed, parity-checked against the composed two-stage path). These tests pin +its central constraint: backend=1 (fma_ksplit + big_fuse) writes +``num_k_splits * M`` partial rows into the y_acc/r_acc accumulator +workspaces, so they must be sized ``FUSED_HC_MAX_K_SPLITS x max_tokens``. +With M-row accumulators the spill corrupts small-M outputs and crashed a +production p8 run with ``cudaErrorIllegalAddress``. +""" + +import pytest +import torch +from tokenspeed_kernel.ops.mhc.trtllm import ( + FUSED_HC_MAX_K_SPLITS, + _select_fused_hc_launch, +) + +HC, HID = 4, 7168 +MIX = (2 + HC) * HC +RMS_EPS, HC_EPS, SINK = 1e-6, 1e-2, 4 +REL_TOL = 0.03 + + +def _is_sm100() -> bool: + if not torch.cuda.is_available() or torch.version.hip is not None: + return False + return torch.cuda.get_device_capability(0) == (10, 0) + + +class TestSelectFusedHcLaunch: + def test_small_m_stays_on_allinone_default(self): + # The allinone default is at least as fast up to M=12, so the table + # only switches organizations where the sweep showed a real win. + for m in range(1, 13): + assert _select_fused_hc_launch(m) == (3, 1, 1) + + def test_medium_m_uses_ksplit_tile2(self): + for m in (13, 16, 24, 32): + assert _select_fused_hc_launch(m) == (1, 2, 2) + + def test_large_m_uses_ksplit_tile4(self): + for m in (33, 36, 48, 64, 128): + assert _select_fused_hc_launch(m) == (1, 4, 2) + + def test_cliff_backend_never_selected(self): + # backend=2 (allinone tf32 mma) measured 206-241us at every M on + # B200 -- the launch table must never dispatch to it. + for m in range(1, 129): + assert _select_fused_hc_launch(m)[0] != 2 + + def test_k_splits_bounded_by_workspace_contract(self): + for m in range(1, 129): + assert _select_fused_hc_launch(m)[2] <= FUSED_HC_MAX_K_SPLITS + + +def _make_inputs(m, dev): + torch.manual_seed(0) + x = torch.randn(m, HID, device=dev, dtype=torch.bfloat16) + res = torch.randn(m, HC, HID, device=dev, dtype=torch.bfloat16) * 0.5 + post = torch.rand(m, HC, device=dev, dtype=torch.float32) + comb = torch.rand(m, HC, HC, device=dev, dtype=torch.float32) / HC + fn = torch.randn(MIX, HC * HID, device=dev, dtype=torch.float32) * 0.02 + scale = torch.ones(3, device=dev, dtype=torch.float32) + base = torch.zeros(MIX, device=dev, dtype=torch.float32) + return x, res, post, comb, fn, scale, base + + +def _raw_fused_hc(inputs, m, backend, tile_n, ksp, acc_rows=None): + """Call the raw op with explicit launch params; return outputs and acc.""" + + x, res, post, comb, fn, scale, base = inputs + dev = x.device + rows = acc_rows if acc_rows is not None else FUSED_HC_MAX_K_SPLITS * m + shape_n = HC * (2 + HC) + residual_cur = torch.empty(m, HC, HID, dtype=torch.bfloat16, device=dev) + post_cur = torch.empty(m, HC, dtype=torch.float32, device=dev) + comb_cur = torch.empty(m, HC * HC, dtype=torch.float32, device=dev) + layer_input = torch.empty(m, HID, dtype=torch.bfloat16, device=dev) + y_acc = torch.full((rows, shape_n), 12345.678, dtype=torch.float32, device=dev) + r_acc = torch.full((rows,), 12345.678, dtype=torch.float32, device=dev) + dc = torch.zeros(rows, dtype=torch.int32, device=dev) + torch.ops.trtllm.mhc_fused_hc( + x, + res, + post, + comb, + fn, + scale, + base, + residual_cur, + post_cur, + comb_cur, + layer_input, + y_acc, + r_acc, + dc, + m, + HID, + HC, + RMS_EPS, + HC_EPS, + HC_EPS, + 2.0, + SINK, + backend, + tile_n, + ksp, + 0, + 1, + None, + 0.0, + ) + torch.cuda.synchronize() + return (residual_cur, layer_input, post_cur, comb_cur), (y_acc, r_acc) + + +def _assert_close(got, ref): + for g, r in zip(got, ref): + g = g.float() + r = r.float() + denom = r.abs().max().clamp_min(1e-6) + rel = ((g - r).abs().max() / denom).item() + assert rel <= REL_TOL, f"rel err {rel} > {REL_TOL}" + + +@pytest.mark.skipif(not _is_sm100(), reason="TRT-LLM mHC kernels need sm100") +class TestFusedHcLaunchTableParityGpu: + @pytest.fixture(autouse=True) + def _load_kernels(self): + pytest.importorskip("trtllm_kernel") + from tokenspeed_kernel.ops.mhc.trtllm import has_trtllm_mhc + + if not has_trtllm_mhc(): + pytest.skip("trtllm mHC kernels unavailable") + + @pytest.mark.parametrize("m", [13, 16, 24, 32, 36, 48, 64]) + def test_selected_launch_matches_allinone_reference(self, m): + # The allinone-fma backend (3) is the numerically trusted incumbent + # at every M; the sweep-selected launch must agree with it. + inputs = _make_inputs(m, "cuda:0") + ref, _ = _raw_fused_hc(inputs, m, 3, 1, 1) + backend, tile_n, ksp = _select_fused_hc_launch(m) + got, _ = _raw_fused_hc(inputs, m, backend, tile_n, ksp) + _assert_close(got, ref) + + @pytest.mark.parametrize("m", [16, 32]) + def test_ksplit_writes_beyond_m_rows_but_within_contract(self, m): + # Documents the workspace contract: the k-split backend spills + # num_k_splits x M partial rows -- more than M (so M-row buffers + # overflow) but never more than FUSED_HC_MAX_K_SPLITS x M. + inputs = _make_inputs(m, "cuda:0") + backend, tile_n, ksp = _select_fused_hc_launch(m) + assert backend == 1 + _, (y_acc, r_acc) = _raw_fused_hc( + inputs, m, backend, tile_n, ksp, acc_rows=FUSED_HC_MAX_K_SPLITS * m + ) + y_rows = int((y_acc != 12345.678).any(dim=1).sum().item()) + assert y_rows > m + assert y_rows <= FUSED_HC_MAX_K_SPLITS * m + + @pytest.mark.parametrize("m", list(range(1, 9))) + def test_small_m_ksplit_parity_with_contract_sized_workspace(self, m): + # With accumulators sized to the FUSED_HC_MAX_K_SPLITS x M contract, + # backend=1 matches the default at small M too. (With M-row + # accumulators its spill corrupts these outputs -- the same + # undersizing that crashed production at scale.) The launch table + # still keeps small M on the default backend because it is at least + # as fast there. + inputs = _make_inputs(m, "cuda:0") + ref, _ = _raw_fused_hc(inputs, m, 3, 1, 1) + got, _ = _raw_fused_hc(inputs, m, 1, 2, 2) + _assert_close(got, ref) From a3b42c02e0fa8bd3cc2263b0bdb6648e0b074f7e Mon Sep 17 00:00:00 2001 From: Xiangyi Zhang Date: Sun, 12 Jul 2026 23:20:42 -0700 Subject: [PATCH 2/3] perf(deepseek-v4): tune big-fuse block size in fused mHC launch table Sweeping bigfuse_block_size (fixed at the kernel default in the first sweep) over the selected launch configs on B200: 512 improves the two-stage path at every production decode shape (graph-replay microbench): 15.1 -> 14.7us at M=32, 11.9 -> 9.6us at M=16, 17.9 -> 17.2us at M=48. Endpoint-neutral by magnitude (~0.1 ms/step at p8); claims are kernel-level only. The same sweep ran the full backend=0 (atomic-routing GEMM) tile grid at M in {32, 64}: nothing beats the backend=1 table at production shapes, so the atomic path stays unselected. Signed-off-by: Xiangyi Zhang --- .../tokenspeed_kernel/ops/mhc/trtllm.py | 19 +++++++++------ .../test/test_deepseek_v4_mhc_fused_launch.py | 24 ++++++++++++------- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py index 77a25c600..59eb42abc 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py @@ -112,21 +112,24 @@ def supports_trtllm_mhc( """ -def _select_fused_hc_launch(num_tokens: int) -> tuple[int, int, int]: +def _select_fused_hc_launch(num_tokens: int) -> tuple[int, int, int, int]: """Pick the fused_hc backend and tile configuration for a token count. Args: num_tokens: Number of tokens (rows) in the fused mHC call. Returns: - Tuple of ``(backend, tile_n, num_k_splits)`` kernel launch parameters. + Tuple of ``(backend, tile_n, num_k_splits, bigfuse_block_size)`` + kernel launch parameters. """ if num_tokens <= _FUSED_HC_SMALL_M_MAX: - return 3, 1, 1 # fused_all_fma (allinone) default + return 3, 1, 1, 0 # fused_all_fma (allinone) default + # bigfuse_block_size=512 over the kernel default: 15.1 -> 14.7us at M=32, + # 11.9 -> 9.6us at M=16 (graph-replay microbench; endpoint-neutral). if num_tokens <= _FUSED_HC_MEDIUM_M_MAX: - return 1, 2, 2 # fma_ksplit + big_fuse - return 1, 4, 2 + return 1, 2, 2, 512 # fma_ksplit + big_fuse + return 1, 4, 2, 512 trtllm_mhc_big_fuse = error_fn @@ -309,7 +312,9 @@ def trtllm_mhc_fused_hc( """ num_tokens, hc_mult, hidden_size = residual_prev.shape - backend, tile_n, num_k_splits = _select_fused_hc_launch(num_tokens) + backend, tile_n, num_k_splits, bigfuse_block_size = _select_fused_hc_launch( + num_tokens + ) _mhc_fused_hc( x_prev, residual_prev, @@ -336,7 +341,7 @@ def trtllm_mhc_fused_hc( backend, tile_n, num_k_splits, - 0, + bigfuse_block_size, 1, None, 0.0, diff --git a/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py index 41cd27b76..4ca00c0c2 100644 --- a/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py +++ b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py @@ -35,15 +35,15 @@ def test_small_m_stays_on_allinone_default(self): # The allinone default is at least as fast up to M=12, so the table # only switches organizations where the sweep showed a real win. for m in range(1, 13): - assert _select_fused_hc_launch(m) == (3, 1, 1) + assert _select_fused_hc_launch(m) == (3, 1, 1, 0) def test_medium_m_uses_ksplit_tile2(self): for m in (13, 16, 24, 32): - assert _select_fused_hc_launch(m) == (1, 2, 2) + assert _select_fused_hc_launch(m) == (1, 2, 2, 512) def test_large_m_uses_ksplit_tile4(self): for m in (33, 36, 48, 64, 128): - assert _select_fused_hc_launch(m) == (1, 4, 2) + assert _select_fused_hc_launch(m) == (1, 4, 2, 512) def test_cliff_backend_never_selected(self): # backend=2 (allinone tf32 mma) measured 206-241us at every M on @@ -68,7 +68,7 @@ def _make_inputs(m, dev): return x, res, post, comb, fn, scale, base -def _raw_fused_hc(inputs, m, backend, tile_n, ksp, acc_rows=None): +def _raw_fused_hc(inputs, m, backend, tile_n, ksp, bfbs=0, acc_rows=None): """Call the raw op with explicit launch params; return outputs and acc.""" x, res, post, comb, fn, scale, base = inputs @@ -108,7 +108,7 @@ def _raw_fused_hc(inputs, m, backend, tile_n, ksp, acc_rows=None): backend, tile_n, ksp, - 0, + bfbs, 1, None, 0.0, @@ -142,8 +142,8 @@ def test_selected_launch_matches_allinone_reference(self, m): # at every M; the sweep-selected launch must agree with it. inputs = _make_inputs(m, "cuda:0") ref, _ = _raw_fused_hc(inputs, m, 3, 1, 1) - backend, tile_n, ksp = _select_fused_hc_launch(m) - got, _ = _raw_fused_hc(inputs, m, backend, tile_n, ksp) + backend, tile_n, ksp, bfbs = _select_fused_hc_launch(m) + got, _ = _raw_fused_hc(inputs, m, backend, tile_n, ksp, bfbs=bfbs) _assert_close(got, ref) @pytest.mark.parametrize("m", [16, 32]) @@ -152,10 +152,16 @@ def test_ksplit_writes_beyond_m_rows_but_within_contract(self, m): # num_k_splits x M partial rows -- more than M (so M-row buffers # overflow) but never more than FUSED_HC_MAX_K_SPLITS x M. inputs = _make_inputs(m, "cuda:0") - backend, tile_n, ksp = _select_fused_hc_launch(m) + backend, tile_n, ksp, bfbs = _select_fused_hc_launch(m) assert backend == 1 _, (y_acc, r_acc) = _raw_fused_hc( - inputs, m, backend, tile_n, ksp, acc_rows=FUSED_HC_MAX_K_SPLITS * m + inputs, + m, + backend, + tile_n, + ksp, + bfbs=bfbs, + acc_rows=FUSED_HC_MAX_K_SPLITS * m, ) y_rows = int((y_acc != 12345.678).any(dim=1).sum().item()) assert y_rows > m From d5bc88f3b90c28650913ff1efa242b6164013c8a Mon Sep 17 00:00:00 2001 From: Xiangyi Zhang Date: Mon, 13 Jul 2026 03:26:48 -0700 Subject: [PATCH 3/3] perf(deepseek-v4): extend two-stage fused mHC launch to small token counts With contract-sized accumulator workspaces the two-stage k-split path wins at every measured M, not just M>=13: 7.6-9.2us vs 9.7-11.4us for the allinone default at M in [1, 12] (graph-replay microbench, parity within 0.5% of the composed reference at every point). The earlier small-M gate guarded against wrong values that were actually the workspace undersizing corrupting the k-split spill, so it can go. Low-concurrency decode (p1/p2, M=4/8) sits on this path and also drops the per-call fhcZeroWorkspacesKernel prologue the allinone backend launches; estimated ~0.3-0.45 ms/step at p1/p2 (kernel-level claim, to be verified by decode-trace kernel counts). Signed-off-by: Xiangyi Zhang --- .../python/tokenspeed_kernel/ops/mhc/trtllm.py | 9 ++++----- .../test/test_deepseek_v4_mhc_fused_launch.py | 15 ++++++--------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py index 59eb42abc..2a301321a 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py @@ -98,9 +98,10 @@ def supports_trtllm_mhc( # r_acc workspaces; callers must size them accordingly (see # FUSED_HC_MAX_K_SPLITS). Undersized accumulators corrupt small-M # outputs and crash at scale. -# * small token counts stay on the allinone-fma default, which is at -# least as fast there. -_FUSED_HC_SMALL_M_MAX = 12 +# * the two-stage path wins at every measured M (1..64): 19-24% below +# M=13 (e.g. 8.6 vs 11.4us at M=8), 26-35% at M in [16, 32]. The +# small-M gate that previously kept M<=12 on the allinone default was +# an artifact of the undersized workspace corrupting small-M outputs. _FUSED_HC_MEDIUM_M_MAX = 32 FUSED_HC_MAX_K_SPLITS = 2 @@ -123,8 +124,6 @@ def _select_fused_hc_launch(num_tokens: int) -> tuple[int, int, int, int]: kernel launch parameters. """ - if num_tokens <= _FUSED_HC_SMALL_M_MAX: - return 3, 1, 1, 0 # fused_all_fma (allinone) default # bigfuse_block_size=512 over the kernel default: 15.1 -> 14.7us at M=32, # 11.9 -> 9.6us at M=16 (graph-replay microbench; endpoint-neutral). if num_tokens <= _FUSED_HC_MEDIUM_M_MAX: diff --git a/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py index 4ca00c0c2..a5c6fdf6a 100644 --- a/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py +++ b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py @@ -31,14 +31,11 @@ def _is_sm100() -> bool: class TestSelectFusedHcLaunch: - def test_small_m_stays_on_allinone_default(self): - # The allinone default is at least as fast up to M=12, so the table - # only switches organizations where the sweep showed a real win. - for m in range(1, 13): - assert _select_fused_hc_launch(m) == (3, 1, 1, 0) - - def test_medium_m_uses_ksplit_tile2(self): - for m in (13, 16, 24, 32): + def test_small_and_medium_m_use_ksplit_tile2(self): + # The two-stage path wins at every measured M >= 1 with + # contract-sized workspaces (the earlier small-M gate was an + # artifact of the workspace undersizing). + for m in (1, 2, 4, 8, 12, 13, 16, 24, 32): assert _select_fused_hc_launch(m) == (1, 2, 2, 512) def test_large_m_uses_ksplit_tile4(self): @@ -136,7 +133,7 @@ def _load_kernels(self): if not has_trtllm_mhc(): pytest.skip("trtllm mHC kernels unavailable") - @pytest.mark.parametrize("m", [13, 16, 24, 32, 36, 48, 64]) + @pytest.mark.parametrize("m", [1, 4, 8, 13, 16, 24, 32, 36, 48, 64]) def test_selected_launch_matches_allinone_reference(self, m): # The allinone-fma backend (3) is the numerically trusted incumbent # at every M; the sweep-selected launch must agree with it.