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..2a301321a 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/mhc/trtllm.py @@ -88,6 +88,49 @@ 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. +# * 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 +"""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, 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, bigfuse_block_size)`` + kernel launch parameters. + """ + + # 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, 512 # fma_ksplit + big_fuse + return 1, 4, 2, 512 + + trtllm_mhc_big_fuse = error_fn trtllm_mhc_fused_hc = error_fn trtllm_mhc_post_mapping = error_fn @@ -268,10 +311,9 @@ 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, bigfuse_block_size = _select_fused_hc_launch( + num_tokens + ) _mhc_fused_hc( x_prev, residual_prev, @@ -297,8 +339,8 @@ def trtllm_mhc_fused_hc( sinkhorn_iters, backend, tile_n, - 1, - 0, + num_k_splits, + 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 new file mode 100644 index 000000000..a5c6fdf6a --- /dev/null +++ b/tokenspeed-kernel/test/test_deepseek_v4_mhc_fused_launch.py @@ -0,0 +1,178 @@ +# 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_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): + for m in (33, 36, 48, 64, 128): + 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 + # 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, 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 + 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, + bfbs, + 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", [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. + inputs = _make_inputs(m, "cuda:0") + ref, _ = _raw_fused_hc(inputs, m, 3, 1, 1) + 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]) + 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, bfbs = _select_fused_hc_launch(m) + assert backend == 1 + _, (y_acc, r_acc) = _raw_fused_hc( + 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 + 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)