From 79ff3d55bd2ce69a96a21e26630cfa3afa9ec1b2 Mon Sep 17 00:00:00 2001 From: Avery Huang Date: Thu, 5 Nov 2026 12:44:16 -0800 Subject: [PATCH] perf(cake_dense_projection_det_gemm): drop the 3-split router weight-gradient rule at H 6144 The 3-split count of the FP32 router weight gradient at H 6144 (T 16129-16384, N 256, sm_100a and sm_103a at 148 / 152 SMs) was chosen on an idle GPU, where it is 1-3 % faster than the default 6 splits. With a kernel on another stream holding 8, 16 or 32 SMs, as a communication kernel does in training, a single measurement of it runs up to 1.38x slower than 6 splits, and 6 splits is 1.06-1.17x faster with 16 or 32 SMs held (geometric mean per token count on GB200, GB300, B200 and B300). - H 6144: back to the default 6 splits on sm_100a and sm_103a (the four table keys are removed). - H 5120 keeps 9 splits: faster than 6 splits alone and, as a geometric mean, with 8, 16 or 32 SMs held. - The CPU rule test pins every table key and its count. - Two router tests that run on CPU stand-ins now also stand in the device facts: they asked device 0 whether it supports programmatic dependent launch and failed on a host without a GPU. The H 6144 rows change bits back to the default count's partition; they stay bit-exact run to run and keep the router accuracy budget. --- .../dense_projection_det_gemm/README.md | 4 +-- .../dense_projection_det_gemm/cake_backend.py | 9 +++---- .../test_cake_dense_projection_det_gemm.py | 25 ++++++++++++++++--- 3 files changed, 25 insertions(+), 13 deletions(-) diff --git a/flashinfer/experimental/dense_projection_det_gemm/README.md b/flashinfer/experimental/dense_projection_det_gemm/README.md index 66e6e21a368..0ed1ff3aac8 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/README.md +++ b/flashinfer/experimental/dense_projection_det_gemm/README.md @@ -232,9 +232,7 @@ split take `enable_pdl` too and wait before their first global access. device: by architecture, SM count, layout, hidden size H and token-count window T 16129-16384: `sm_100a` at 148 / 152 SMs, weight gradient H 5120: 9 splits; - `sm_100a` at 148 / 152 SMs, weight gradient H 6144: 3 splits; - `sm_103a` at 148 / 152 SMs, weight gradient H 5120: 9 splits; - `sm_103a` at 148 / 152 SMs, weight gradient H 6144: 3 splits. + `sm_103a` at 148 / 152 SMs, weight gradient H 5120: 9 splits. * Every `launch()` re-splits the second operand into the retained `[3, outer, inner]` BF16 stack in the operand's own layout (exact three-way split, no allocation: one launch of the fused split program diff --git a/flashinfer/experimental/dense_projection_det_gemm/cake_backend.py b/flashinfer/experimental/dense_projection_det_gemm/cake_backend.py index 2e97bf28e23..a44cd358210 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/cake_backend.py +++ b/flashinfer/experimental/dense_projection_det_gemm/cake_backend.py @@ -231,19 +231,16 @@ # the inclusive window [T_lo, T_hi], for N = ROUTER_RULE_N. The window T 16129-16384 holds every token count with the # forward's 128 M tiles (505-512 K blocks of the weight gradient). Each count was measured faster than the default one # on GB200 / GB300 (152 SMs) and on B200 / B300 (148 SMs) for the whole router launch (the BF16x3 split, the program and the sum of -# the partials); other SM counts keep the default. A different split count is a different K partition: its outputs are -# bit-exact run to run but not bitwise equal to the default count's. +# the partials), alone and while a kernel on another stream holds 8, 16 or 32 SMs (a count that wins alone but loses +# with fewer free SMs is not listed); other SM counts keep the default. A different split count is a different K +# partition: its outputs are bit-exact run to run but not bitwise equal to the default count's. ROUTER_SPLIT_RULES: dict[ tuple[str, int, str, int], tuple[tuple[int, int, int], ...] ] = { ("sm_100a", 148, "nn_t", 5120): ((16129, 16384, 9),), - ("sm_100a", 148, "nn_t", 6144): ((16129, 16384, 3),), ("sm_100a", 152, "nn_t", 5120): ((16129, 16384, 9),), - ("sm_100a", 152, "nn_t", 6144): ((16129, 16384, 3),), ("sm_103a", 148, "nn_t", 5120): ((16129, 16384, 9),), - ("sm_103a", 148, "nn_t", 6144): ((16129, 16384, 3),), ("sm_103a", 152, "nn_t", 5120): ((16129, 16384, 9),), - ("sm_103a", 152, "nn_t", 6144): ((16129, 16384, 3),), } diff --git a/tests/experimental/test_cake_dense_projection_det_gemm.py b/tests/experimental/test_cake_dense_projection_det_gemm.py index b72da9cd57e..24758a8d358 100644 --- a/tests/experimental/test_cake_dense_projection_det_gemm.py +++ b/tests/experimental/test_cake_dense_projection_det_gemm.py @@ -5566,6 +5566,14 @@ def test_router_split_rules_are_measured_windows_and_resolve_by_device(): SMs) and a ``kk`` / ``nn_t`` row in the T 16129-16384 window; ``router_splits`` returns it inside the window and the layout's default outside it, for another SM count or N, and without a device.""" assert ROUTER_SPLIT_RULES + # the measured rows and their counts, pinned here rather than read from the table under test + assert {key: windows[0][2] for key, windows in ROUTER_SPLIT_RULES.items()} == { + ("sm_100a", 148, "nn_t", 5120): 9, + ("sm_100a", 152, "nn_t", 5120): 9, + ("sm_103a", 148, "nn_t", 5120): 9, + ("sm_103a", 152, "nn_t", 5120): 9, + } + assert all(len(windows) == 1 for windows in ROUTER_SPLIT_RULES.values()) for (arch, sm_count, layout, H), windows in ROUTER_SPLIT_RULES.items(): assert arch in ("sm_100a", "sm_103a") and sm_count in (148, 152), arch assert layout in ("kk", "nn_t") and H in (4096, 5120, 6144), layout @@ -6157,7 +6165,7 @@ def grows_then_fails(arena): def test_convenience_router_launch_holds_the_arena_lock(monkeypatch): """The K2 hit path holds its arena's lock from the carve of the workspace views through the launch and the - split-K reduction (CPU stand-ins for the device, the stream and the FFI entry).""" + split-K reduction (CPU stand-ins for the device and its facts, the stream and the FFI entry).""" monkeypatch.setattr(cake_backend, "_K1_CACHE", OrderedDict()) monkeypatch.setattr(cake_backend, "_K2_CACHE", OrderedDict()) monkeypatch.setattr(cake_backend, "_ARENAS", {}) @@ -6173,6 +6181,10 @@ def entry_call(*arguments): cached.entry = entry_call key = convenience_key(0, 7, (4,), v["A"], v["B"], v["out"]) cake_backend._K2_CACHE[key] = cached + facts = cake_backend.DeviceFacts( + index=0, arch="sm_100a", sm_count=148, l2_bytes=L2_BYTES + ) + monkeypatch.setattr(cake_backend, "device_facts", lambda device=None: facts) monkeypatch.setattr(cake_backend, "_same_cuda_device", lambda A, B, out: out.device) monkeypatch.setattr(cake_backend, "_cuda_index", lambda d: 0) monkeypatch.setattr(cake_backend, "_current_raw_stream", lambda index: 7) @@ -6219,9 +6231,10 @@ def locked_split(src, parts, resid): def _cpu_router_registry(monkeypatch, *, fused): - """CPU stand-ins for a K2 plan on ``sm_100a``: the device check and the architecture, the stream and FFI - context, and a module loader whose entries record ``(module, arguments, an arena lock was held)``; with - ``fused`` the registry holds a stand-in fused split program, without it none. Returns the recorded calls.""" + """CPU stand-ins for a K2 plan on ``sm_100a``: the device check, the architecture and the device facts, the + stream and FFI context, and a module loader whose entries record ``(module, arguments, an arena lock was + held)``; with ``fused`` the registry holds a stand-in fused split program, without it none. Returns the + recorded calls.""" calls = [] kernels = dict(cake_jit.KERNELS["sm_100a"]) kernels.pop(cake_backend.SPLIT_TEMPLATE, None) @@ -6229,6 +6242,10 @@ def _cpu_router_registry(monkeypatch, *, fused): kernels[cake_backend.SPLIT_TEMPLATE] = _FAKE_SPLIT_MODULE monkeypatch.setitem(cake_jit.MODULES, _FAKE_SPLIT_MODULE, _FAKE_SPLIT_RECORD) monkeypatch.setitem(cake_jit.KERNELS, "sm_100a", kernels) + facts = cake_backend.DeviceFacts( + index=0, arch="sm_100a", sm_count=148, l2_bytes=L2_BYTES + ) + monkeypatch.setattr(cake_backend, "device_facts", lambda device=None: facts) monkeypatch.setattr(cake_backend, "_same_cuda_device", lambda A, B, out: out.device) monkeypatch.setattr(cake_backend, "require_arch", lambda device: "sm_100a") monkeypatch.setattr(cake_backend, "_cuda_index", lambda d: 0)