perf(cake_dense_projection_det_gemm): drop the 3-split router weight-gradient rule at H 6144

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

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),),
 }
 
 

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,16 +6231,21 @@ 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)
     if 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)