From 98009e5e7120696de0104b82a6e3bd093db72b12 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 5 Nov 2026 21:20:42 -0800 Subject: [PATCH 01/13] test(dense_projection_gemm): call deterministic entries by their det names The non-deterministic dense projection GEMM tests and benchmark used this package's default (deterministic omitted = the deterministic package) as the reference for the deterministic result. Call the deterministic package's canonical entries (projection_det_*, router_det_*, router_fp32_det_gemm, dense_projection_det_gemm) directly instead, and drop the test that asserted the deterministic package's former names without "det". The tests and the benchmark no longer depend on those former names or on the default mode, so the deterministic package can drop its aliases. No library code changes; every changed call runs the same deterministic program with the same arguments as before. --- .../bench_cake_dense_projection_gemm.py | 27 +++-- .../test_cake_dense_projection_gemm.py | 98 ++++++++++--------- 2 files changed, 72 insertions(+), 53 deletions(-) diff --git a/benchmarks/bench_cake_dense_projection_gemm.py b/benchmarks/bench_cake_dense_projection_gemm.py index f9f5ceee0da..050bc43a898 100644 --- a/benchmarks/bench_cake_dense_projection_gemm.py +++ b/benchmarks/bench_cake_dense_projection_gemm.py @@ -315,18 +315,28 @@ def prepare_atomic(spec, inp, out, **knobs): ) -def convenience(spec, inp, out, *, deterministic): - """The convenience entry point of the row in one mode (cached plan, arena workspace).""" +def convenience_det(spec, inp, out): + """The deterministic package's convenience entry point of the row (cached plan, arena workspace).""" + if spec["family"] == "router": + return lambda: det_backend.router_fp32_det_gemm(inp["A"], inp["B"], out) + if spec["family"] == "proj" and spec["op"] == "wgrad": + return lambda: det_backend.projection_det_wgrad(inp["G"], inp["X"], out) + return lambda: det_backend.dense_projection_det_gemm(inp["A"], inp["B"], out) + + +def convenience_atomic(spec, inp, out): + """``deterministic=False`` through this package's convenience entry point of the row (cached plan, arena + workspace).""" if spec["family"] == "router": return lambda: cake_backend.router_fp32_gemm( - inp["A"], inp["B"], out, deterministic=deterministic + inp["A"], inp["B"], out, deterministic=False ) if spec["family"] == "proj" and spec["op"] == "wgrad": return lambda: cake_backend.projection_wgrad( - inp["G"], inp["X"], out, deterministic=deterministic + inp["G"], inp["X"], out, deterministic=False ) return lambda: cake_backend.dense_projection_gemm( - inp["A"], inp["B"], out, deterministic=deterministic + inp["A"], inp["B"], out, deterministic=False ) @@ -1346,9 +1356,12 @@ def _measure_row(entry, spec, args, libs, resident, uuid): mem["atomic_prepared_retained"], mem["atomic_prepared_peak_temp"] = memory_of( lambda: prepare_atomic(spec, inp, outs["atomic"]).launch ) - for mode, deterministic in (("det", True), ("atomic", False)): + for mode, convenience in ( + ("det", convenience_det), + ("atomic", convenience_atomic), + ): cake_backend.clear_convenience_caches() - fn = convenience(spec, inp, outs[mode], deterministic=deterministic) + fn = convenience(spec, inp, outs[mode]) retained, peak = memory_of(lambda fn=fn: first_call(fn)) mem[f"{mode}_convenience_retained"] = retained mem[f"{mode}_convenience_peak_temp"] = peak diff --git a/tests/experimental/test_cake_dense_projection_gemm.py b/tests/experimental/test_cake_dense_projection_gemm.py index 4248bdc6a2f..31ab28ba5b6 100644 --- a/tests/experimental/test_cake_dense_projection_gemm.py +++ b/tests/experimental/test_cake_dense_projection_gemm.py @@ -81,7 +81,7 @@ "router_wgrad", "split_fp32_to_bf16x3_", ) -# former name -> the deterministic package's canonical name (its alias block, kept by #8046) +# former name -> the deterministic package's canonical name FORMER_TO_DET = { "GemmPlan": "DetGemmPlan", "PreparedGemm": "PreparedDetGemm", @@ -217,16 +217,6 @@ def test_entry_points_take_keyword_only_deterministic_true(name): assert parameter.default is True -def test_the_deterministic_names_map(): - assert { - k: v for k, v in cake_backend.DETERMINISTIC_NAMES.items() if k in FORMER_TO_DET - } == FORMER_TO_DET - for former, canonical in FORMER_TO_DET.items(): - assert getattr(det_backend, former) is getattr(det_backend, canonical) - for name in ENTRY_POINTS: - assert getattr(cake_backend, name).__module__ == cake_backend.__name__ - - @pytest.mark.parametrize("name", FORMER_MODULE_NAMES) def test_every_former_module_name_resolves(name): obj = getattr(cake_backend, name) @@ -7690,7 +7680,7 @@ def test_gpu_k1_budgets_over_reruns(row, out_dtype, splits, k1_rule): device, arch = _atomic_device_or_skip() G, X, out = _k1_wgrad_inputs(row, 16231, out_dtype, device, SEED) ref = G.double().t() @ X.double() - det_out = cake_backend.projection_wgrad(G, X, torch.empty_like(out)) + det_out = det_backend.projection_det_wgrad(G, X, torch.empty_like(out)) prepared = _k1_prepare(G, X, out, arch, splits) _assert_ran(prepared.launch, prepared.module_name) k1_rule(G, X, out, splits) @@ -7937,7 +7927,7 @@ def test_gpu_k1_ragged_token_counts_reuse_poisoned_scratch(k1_rule): _k1_check(out, G.double().t() @ X.double(), f"T={T}") assert _arena_tickets_zero(), T if not cake_backend.k1_splits_fit(3, T): - assert torch.equal(out, cake_backend.projection_wgrad(G, X)) + assert torch.equal(out, det_backend.projection_det_wgrad(G, X)) assert len(cake_backend._ATOMIC_ARENAS) == 1 (arena,) = cake_backend._ATOMIC_ARENAS.values() assert arena.version >= 2 # created, then grown or reused; never shrunk @@ -7958,7 +7948,7 @@ def test_gpu_torch_deterministic_mode_rule(torch_determinism): launch; with ``warn_only=True`` it warns and returns the deterministic bits.""" device = _det_device_or_skip() G, X, out = _k1_wgrad_inputs("indexer_hw", 4097, torch.bfloat16, device, SEED + 40) - expected = cake_backend.projection_wgrad(G, X, torch.empty_like(out)) + expected = det_backend.projection_det_wgrad(G, X, torch.empty_like(out)) torch.use_deterministic_algorithms(True) with pytest.raises(RuntimeError, match="deterministic=True"): cake_backend.projection_wgrad(G, X, out, deterministic=False) @@ -7976,7 +7966,7 @@ def test_gpu_det_then_atomic_share_a_stream_not_a_workspace(k1_rule): k1_rule(G, X, out, 2) ref = G.double().t() @ X.double() for _ in range(3): - det = cake_backend.projection_wgrad(G, X, torch.empty_like(out)) + det = det_backend.projection_det_wgrad(G, X, torch.empty_like(out)) atomic = cake_backend.projection_wgrad(G, X, out, deterministic=False) torch.cuda.synchronize() _k1_check(det, ref) @@ -8011,19 +8001,26 @@ def check(T, seed, pad_out=False): _bf16((N, K), device, g, K**-0.5), _bf16((T, N), device, g), ) + # op -> (out shape, out dtype, the deterministic package's call, this package's deterministic=False call) # fmt: off calls = { - "fwd": ((T, N), torch.bfloat16, lambda o, d: cake_backend.projection_forward(X, W, o, deterministic=d)), - "dgrad": ((T, K), torch.bfloat16, lambda o, d: cake_backend.projection_dgrad(G, W, o, deterministic=d)), - "dgrad_f32": ((T, K), torch.float32, lambda o, d: cake_backend.projection_dgrad(G, W, o, deterministic=d)), - "wgrad": ((N, K), torch.bfloat16, lambda o, d: cake_backend.projection_wgrad(G, X, o, deterministic=d)), - "wgrad_f32": ((N, K), torch.float32, lambda o, d: cake_backend.projection_wgrad(G, X, o, deterministic=d)), - "wgrad_sk": ((N, K), torch.float32, lambda o, d: cake_backend.projection_wgrad(G, X, o, sk=False, deterministic=d)), + "fwd": ((T, N), torch.bfloat16, lambda o: det_backend.projection_det_forward(X, W, o), + lambda o: cake_backend.projection_forward(X, W, o, deterministic=False)), + "dgrad": ((T, K), torch.bfloat16, lambda o: det_backend.projection_det_dgrad(G, W, o), + lambda o: cake_backend.projection_dgrad(G, W, o, deterministic=False)), + "dgrad_f32": ((T, K), torch.float32, lambda o: det_backend.projection_det_dgrad(G, W, o), + lambda o: cake_backend.projection_dgrad(G, W, o, deterministic=False)), + "wgrad": ((N, K), torch.bfloat16, lambda o: det_backend.projection_det_wgrad(G, X, o), + lambda o: cake_backend.projection_wgrad(G, X, o, deterministic=False)), + "wgrad_f32": ((N, K), torch.float32, lambda o: det_backend.projection_det_wgrad(G, X, o), + lambda o: cake_backend.projection_wgrad(G, X, o, deterministic=False)), + "wgrad_sk": ((N, K), torch.float32, lambda o: det_backend.projection_det_wgrad(G, X, o, sk=False), + lambda o: cake_backend.projection_wgrad(G, X, o, sk=False, deterministic=False)), } # fmt: on arch = det_backend.arch_for(device) sm_count = det_backend.device_facts(device).sm_count - for op, ((rows_, cols), dtype, call) in calls.items(): + for op, ((rows_, cols), dtype, det_call, call) in calls.items(): f32 = dtype == torch.float32 if op.startswith("wgrad") and cake_backend.k1_plan( arch=arch, L=1, M=K, N=N, K=T, out_f32=f32, sm_count=sm_count @@ -8036,10 +8033,10 @@ def check(T, seed, pad_out=False): out = storage[:, :cols] else: out = torch.empty(rows_, cols, device=device, dtype=dtype) - want = call(out, True).clone() + want = det_call(out).clone() for _ in range(2): out.fill_(float("nan")) - assert call(out, False) is out + assert call(out) is out assert torch.equal(out, want), (name, op, T, pad_out) if pad_out: assert bool((storage[:, cols:] == 7.0).all()), (name, op) @@ -8058,7 +8055,7 @@ def check(T, seed, pad_out=False): ) if det_backend._swap_ab_launch(A, B, out, False) is not None: swapped.add((T, shape)) - want = cake_backend.dense_projection_gemm(A, B, out).clone() + want = det_backend.dense_projection_det_gemm(A, B, out).clone() out.fill_(float("nan")) assert ( cake_backend.dense_projection_gemm(A, B, out, deterministic=False) @@ -8103,7 +8100,7 @@ def check(T, seed, pad_out=False): Xm = storage[1 : 1 + 1001 * 4096].view(1001, 4096) Xm.copy_(X) with pytest.raises(ValueError, match="16-byte aligned") as refused: - cake_backend.projection_forward(Xm, W) + det_backend.projection_det_forward(Xm, W) with pytest.raises(ValueError, match=re.escape(str(refused.value))): cake_backend.projection_forward(Xm, W, deterministic=False) @@ -8163,7 +8160,7 @@ def test_gpu_k1_pdl_waits_for_an_early_producer(splits, k1_rule): for _ in range(3): G.fill_(float("nan")) out.fill_(float("nan")) - cake_backend.projection_forward(H, Wg, G, enable_pdl=True) + det_backend.projection_det_forward(H, Wg, G, enable_pdl=True) cake_backend.projection_wgrad(G, X, out, deterministic=False, enable_pdl=True) torch.cuda.synchronize() _k1_check(out, G.double().t() @ X.double()) @@ -8184,12 +8181,12 @@ def test_gpu_k1_bound_deterministic_program_waits_for_an_early_producer(): W = _bf16((N, K), device, g, K**-0.5) G = torch.empty(T, N, device=device, dtype=torch.bfloat16) dX = torch.empty(T, K, device=device, dtype=torch.bfloat16) - cake_backend.projection_forward(H, Wg, G, enable_pdl=False) - want = cake_backend.projection_dgrad(G, W, enable_pdl=False).clone() + det_backend.projection_det_forward(H, Wg, G, enable_pdl=False) + want = det_backend.projection_det_dgrad(G, W, enable_pdl=False).clone() for _ in range(3): G.fill_(float("nan")) dX.fill_(float("nan")) - cake_backend.projection_forward(H, Wg, G, enable_pdl=True) + det_backend.projection_det_forward(H, Wg, G, enable_pdl=True) assert ( cake_backend.projection_dgrad( G, W, dX, deterministic=False, enable_pdl=True @@ -8405,7 +8402,7 @@ def _router_gate(op, splits, spec, router_rule, T, hidden, seed): assert worst <= ROUTER_REL_FRO_MAX, worst if splits == 1: det_out = torch.empty_like(out) - cake_backend.router_fp32_gemm(A, B, det_out, splits=1) + det_backend.router_fp32_det_gemm(A, B, det_out, splits=1) assert torch.equal(out, det_out) return plan @@ -8522,7 +8519,7 @@ def test_gpu_router_tmem_head_tail_program_in_its_windows(hidden, T): launch.launch() torch.cuda.synchronize() det_out = torch.empty_like(out) - cake_backend.router_fp32_gemm(A, B, det_out, splits=1) + det_backend.router_fp32_det_gemm(A, B, det_out, splits=1) assert torch.equal(stores[ts], stores[c3]) and torch.equal(stores[ts], det_out) @@ -8887,7 +8884,7 @@ def test_gpu_router_input_gradient_windows_run_at_their_edges(): prepared = cake_backend.prepare_router_fp32_gemm( G, W, torch.empty(T, hidden, device=device), deterministic=False ) - det_out = cake_backend.router_dgrad(G, W) + det_out = det_backend.router_det_dgrad(G, W) first = None for _ in range(2): out = torch.full((T, hidden), float("nan"), device=device) @@ -9014,7 +9011,7 @@ def test_gpu_router_forward_simt_windows_keep_the_fp64_rule(monkeypatch): assert torch.equal(out, got), (hidden, T) del graph ref = X.double() @ W.double().t() - base = cake_backend.router_forward(X, W) + base = det_backend.router_det_forward(X, W) err, base_err = _rel_fro(got, ref), _rel_fro(base, ref) assert err <= ROUTER_REL_FRO_MAX, (hidden, T, seed, err) gamma = hidden * 2.0**-24 / (1 - hidden * 2.0**-24) @@ -9139,7 +9136,7 @@ def test_gpu_router_entry_points_reach_the_stream_k_program(router_rule): router_rule("fwd", X, W.t(), "forward") router_rule("wgrad", G.t(), X, "pairs") first = cake_backend.router_forward(X, W, deterministic=False) - det_fwd = cake_backend.router_forward(X, W) + det_fwd = det_backend.router_det_forward(X, W) ref_fwd, ref_wgrad = X.double() @ W.double().t(), G.double().t() @ X.double() for name, out in (("atomic", first), ("det", det_fwd)): differ, ties = _top8_rows_differ(out, ref_fwd) @@ -9320,7 +9317,7 @@ def test_gpu_k1_forward_rows_run_their_plan(): torch.cuda.synchronize() assert torch.equal(out, again), (K, N, T) det_out = torch.full_like(out, float("nan")) - cake_backend.projection_forward(X, W, det_out) + det_backend.projection_det_forward(X, W, det_out) torch.cuda.synchronize() _k1_check(det_out, ref, f"deterministic forward {K} {N} {T}") ran += 1 @@ -9767,8 +9764,8 @@ def test_gpu_router_rows_without_a_verdict_run_the_deterministic_program_unplann wgrad = cake_backend.router_wgrad(G, X, deterministic=False) assert not cake_backend._K2_ATOMIC_CACHE assert len(cake_backend._K2_SPLITK_CACHE) == 2 - assert torch.equal(fwd, cake_backend.router_forward(X, W)) - assert torch.equal(wgrad, cake_backend.router_wgrad(G, X)) + assert torch.equal(fwd, det_backend.router_det_forward(X, W)) + assert torch.equal(wgrad, det_backend.router_det_wgrad(G, X)) A, B, out = _router_case("fwd", X, W, G) assert cake_backend.router_fp32_gemm(A, B, out, deterministic=False) is out assert torch.equal(out, fwd) and not cake_backend._K2_ATOMIC_CACHE @@ -9901,24 +9898,28 @@ def check(T, seed, pad_out=False, pad_g=False): if pad_g: G = torch.cat([G, torch.zeros(T, 16, device=device)], 1)[:, :256] shapes = {"fwd": (T, 256), "dgrad": (T, hidden), "wgrad": (256, hidden)} + # op -> (the deterministic package's call, this package's deterministic=False call) # fmt: off calls = { - "fwd": lambda out, d: cake_backend.router_forward(X, W, out, deterministic=d), - "dgrad": lambda out, d: cake_backend.router_dgrad(G, W, out, deterministic=d), - "wgrad": lambda out, d: cake_backend.router_wgrad(G, X, out, deterministic=d), + "fwd": (lambda out: det_backend.router_det_forward(X, W, out), + lambda out: cake_backend.router_forward(X, W, out, deterministic=False)), + "dgrad": (lambda out: det_backend.router_det_dgrad(G, W, out), + lambda out: cake_backend.router_dgrad(G, W, out, deterministic=False)), + "wgrad": (lambda out: det_backend.router_det_wgrad(G, X, out), + lambda out: cake_backend.router_wgrad(G, X, out, deterministic=False)), } # fmt: on - for op, call in calls.items(): + for op, (det_call, call) in calls.items(): rows, cols = shapes[op] if pad_out: storage = torch.full((rows, cols + 64), 7.0, device=device) out = storage[:, :cols] else: out = torch.empty(rows, cols, device=device) - want = call(out, True).clone() + want = det_call(out).clone() for _ in range(2): out.fill_(float("nan")) - assert call(out, False) is out + assert call(out) is out assert torch.equal(out, want), (op, T, pad_out, pad_g) if pad_out: assert bool((storage[:, cols:] == 7.0).all()), op @@ -9955,7 +9956,7 @@ def check(T, seed, pad_out=False, pad_g=False): storage = torch.empty(1001 * hidden + 4, device=device) out = storage[1 : 1 + 1001 * hidden].view(1001, hidden) with pytest.raises(ValueError, match="16-byte aligned") as refused: - cake_backend.router_dgrad(G, W, out) + det_backend.router_det_dgrad(G, W, out) with pytest.raises(ValueError, match=re.escape(str(refused.value))): cake_backend.router_dgrad(G, W, out, deterministic=False) @@ -10333,6 +10334,11 @@ def test_gpu_router_entry_points_give_their_paths_results(monkeypatch, hidden, T "dgrad": cake_backend.router_dgrad, "wgrad": cake_backend.router_wgrad, } + det_calls = { + "fwd": det_backend.router_det_forward, + "dgrad": det_backend.router_det_dgrad, + "wgrad": det_backend.router_det_wgrad, + } shapes = {"fwd": (T, 256), "dgrad": (T, hidden), "wgrad": (256, hidden)} side = torch.cuda.Stream(device) ran = {"bits": 0, "gate": 0} @@ -10372,7 +10378,7 @@ def test_gpu_router_entry_points_give_their_paths_results(monkeypatch, hidden, T assert torch.equal(out, want), (op, splits, pad) if not entry.atomic: det_kw = {} if op == "dgrad" or splits is None else {"splits": splits} - det_out = calls[op](P, Q, torch.empty_like(out), **det_kw) + det_out = det_calls[op](P, Q, torch.empty_like(out), **det_kw) stream.synchronize() assert torch.equal(out, det_out), (op, splits, pad) ran["bits" if _router_entry_bits(entry) else "gate"] += 1 From 372208ba3b5fe9b65cf730be96967b521a8d2af3 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 5 Nov 2026 21:42:44 -0800 Subject: [PATCH 02/13] refactor(dense_projection_gemm): explicit deterministic=False, no det aliases or routing The non-deterministic dense projection GEMM package no longer routes to the deterministic package and no longer exposes deterministic names: - Every compute entry point keeps a keyword-only `deterministic: bool = True` but accepts only `deterministic=False`. Any other value (the default, True, None, 0, 1, "False") raises ValueError naming the deterministic entry point in flashinfer.experimental.dense_projection_det_gemm. - Under torch.use_deterministic_algorithms(True), PyTorch's alert is raised right before an order-varying launch only (K1 split-K at three or more splits, K2 stream-K with more than two partials per element): RuntimeError, or a UserWarning with warn_only=True after which the launch runs. Fixed-order launches do not alert, and the setting never changes the program that runs. - Removed the module __getattr__ read-through, DETERMINISTIC_NAMES, the instance_symbol and split_fp32_to_bf16x3_ wrappers, and the forwarding of deterministic templates and programs in select_module, the loader and bind_launch. router_symbol takes no mode keyword. - README, package and module docs describe the mode and dispatch rules. --- .../dense_projection_gemm/README.md | 130 ++-- .../dense_projection_gemm/__init__.py | 20 +- .../dense_projection_gemm/cake_backend.py | 611 +++++++++--------- .../dense_projection_gemm/cake_jit.py | 46 +- 4 files changed, 413 insertions(+), 394 deletions(-) diff --git a/flashinfer/experimental/dense_projection_gemm/README.md b/flashinfer/experimental/dense_projection_gemm/README.md index 8385fb0d0b0..3b79eed9096 100644 --- a/flashinfer/experimental/dense_projection_gemm/README.md +++ b/flashinfer/experimental/dense_projection_gemm/README.md @@ -1,42 +1,73 @@ -# Cake dense projection GEMM backend with an explicit numerical mode (SM100 / SM103 / SM107) +# Cake dense projection GEMM backend, non-deterministic (SM100 / SM103 / SM107) Experimental backend for the GLM-5.2 training GEMMs -- the BF16 dense projection GEMM (K1: BF16 operands, FP32 accumulation, BF16 or FP32 output, strided views, ragged token counts, batched head projections) and the FP32 -router-gate GEMM through split-BF16x3 tensor-core emulation (K2) -- with a selectable numerical mode. Family issue: -flashinfer-ai/flashinfer#5677. +router-gate GEMM through split-BF16x3 tensor-core emulation (K2) -- with programs whose split-K reductions add the +partials in the order they finish. Family issue: flashinfer-ai/flashinfer#5677. The deterministic implementation is the package `flashinfer.experimental.dense_projection_det_gemm` (flashinfer-ai/flashinfer#5677, #7043; see its README for the view contract, the host paths and the per-architecture -programs). This package keeps its operations, dtypes, accuracy budgets and host paths, and adds programs whose -split-K reductions add the partials in the order they finish. Only run-to-run bitwise reproducibility is relaxed, -and only where a row adds more than two partials. +programs). This package keeps its operations, dtypes, accuracy budgets and host paths. Only run-to-run bitwise +reproducibility is relaxed. Both the API and the backend are experimental: no compatibility guarantee. -## Numerical mode +## Mode and dispatch rules -Every entry point of `cake_backend` takes the keyword-only `deterministic: bool = True`. +Every compute entry point of `cake_backend` takes the keyword-only `deterministic: bool = True` and runs only with +`deterministic=False`. The entry points: `dense_projection_gemm`, `plan_dense_projection_gemm`, +`prepare_dense_projection_gemm`, `projection_forward`, `projection_dgrad`, `prepare_projection_wgrad`, +`projection_wgrad`, `plan_router_fp32_gemm`, `prepare_router_fp32_gemm`, `router_fp32_gemm`, `router_forward`, +`router_dgrad`, `router_wgrad`. -| call | runs | -|---|---| -| `deterministic=True` (default) | the deterministic package's canonical entry point with the same arguments: the same programs, the same bits, the same errors | -| `deterministic=False` | this package's programs on the rows that carry a measured verdict for the device (below); the deterministic programs for every other row and operation | -| `deterministic=False` under `torch.use_deterministic_algorithms(True)` | `RuntimeError`; with `warn_only=True` a `UserWarning` and the deterministic path | - -*Table: what a call runs, by its `deterministic` argument. The default does not read PyTorch's setting: it is -deterministic already.* +| call | `deterministic=False` | any other value (the default, `True`, `None`, `0`, `1`, `"False"`, a NumPy bool) | under `torch.use_deterministic_algorithms(True)` (strict or `warn_only=True`) | +|---|---|---|---| +| this package's compute entry points (above) | runs this package's routes (table below) | `ValueError` naming the deterministic entry, before anything is planned, allocated or launched | the same routes and programs. Right before a launch whose result can vary run to run (the K1 split-K program at three or more splits; a K2 stream-K partition that gives an element more than two partials), PyTorch's alert: `RuntimeError` (nothing launched, `out` untouched), or a `UserWarning` and then the launch under `warn_only=True`. Fixed-order launches run without an alert | +| this package's helpers (`router_symbol`, `bind_launch`, the planners `plan_atomic_gemm` / `plan_atomic_router`, the rule tables, `clear_convenience_caches`, `use_deterministic_backend`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged | +| `flashinfer.experimental.dense_projection_det_gemm` (`projection_det_wgrad`, `router_fp32_det_gemm`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged: every launch is fixed-order | + +*Table: the mode rule per call and value.* + +The `ValueError` message (one line, wrapped here), for example: + +```text +projection_wgrad: flashinfer.experimental.dense_projection_gemm runs the non-deterministic package only and needs +deterministic=False (got True): pass deterministic=False for the non-deterministic path; the deterministic entry is +flashinfer.experimental.dense_projection_det_gemm.cake_backend.projection_det_wgrad +``` + +Rules: + +- The `deterministic` keyword never selects the deterministic package, and no call falls back to it. For a + deterministic result, call the deterministic package's entry point (`projection_det_wgrad`, ...) yourself. +- Every call states the mode: pass `deterministic=False` on every compute entry point call, the `plan_*` and + `prepare_*` entry points included. +- A `deterministic=False` call reuses the deterministic package's fixed-order programs on the rows the table below + marks "deterministic program". It launches them through this package's binding, or, on such a row, through the + deterministic package's entry point of the same operation where the binding cannot rebuild the call's launch views + (two operands that alias) and for a batched forward on that package's swapped small-token program: the same + program, the same bits. On those rows `plan_*` / `prepare_*` return the deterministic package's plan or prepared + launch (`DetGemmPlan`, `PreparedDetGemm`, `RouterDetPlan`, `PreparedRouterDetGemm`). A row that runs a program + of this package never goes to the deterministic package. +- A result of this package carries no run-to-run bitwise guarantee on any row, also on a row that runs a + fixed-order program today. The "run-to-run bits" rows below describe the current programs, not a contract. +- The PyTorch setting never changes which program a call runs. +- The tensors passed between the forward and the gradients carry no mode: recompute a checkpointed forward with the + package that made the saved one (same package only). Neither package has an autograd Function or saved tape. +- This package imports the deterministic package at run time, as its sibling (a vendored copy binds it with + `cake_backend.use_deterministic_backend`). This package re-exports no name of the deterministic package. | operation | `deterministic=False` | |---|---| -| K1 forward (`projection_forward`, `kk` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order; a recomputed forward reproduces its result). The rows of `K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM `sm_100a` / `sm_103a` parts, the K1 forward table below): the deterministic package's single-CTA program, every tile's whole K in one CTA -- other bits than `deterministic=True`, the same on every run | +| K1 forward (`projection_forward`, `kk` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order; a recomputed forward reproduces its result). The rows of `K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM `sm_100a` / `sm_103a` parts, the K1 forward table below): the deterministic package's single-CTA program, every tile's whole K in one CTA -- other bits than the deterministic package's plan, the same on every run | | K1 input gradient (`projection_dgrad`, `kn` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order) | | K1 weight gradient (`projection_wgrad`, `nn` layouts) | the rows of `K1_SPLIT_RULES`: split-K over the ragged token count -- every split writes its FP32 partial tile to a workspace slot; the last split of a tile to arrive (a per-tile ticket counter) adds the other slots in arrival order and rounds once into `out`, then returns the counter to zero. Every other row: deterministic program | | K2 forward (`router_forward`, `kk`) | the token windows of its `ROUTER_RULES` rows (152-SM `sm_100a` / `sm_103a` parts, 256 experts, hidden 4096, 5120 and 6144, 1 to 192 tokens; `sm_100a` hidden 4096: 1 to 160): the exact-FP32 SIMT program `router_fp32_fwd_simt_c8_*` -- `K` split over the eight CTAs of a thread-block cluster, FFMA chains of the FP32 operands read in place (up to 16 tokens in a compensated form), the eight partials added in cluster-rank order through distributed shared memory, every element of `out` stored once. Every other row, token count and part: deterministic program (its four splits); an explicit `splits` of 1 or 2 runs this package's stream-K program (at most two partials per element) | | K2 input gradient (`router_dgrad`, `kn`) | the token windows of its `ROUTER_RULES` rows (on 212-SM `sm_107a` and 148-SM `sm_100a` / `sm_103a` parts `ROUTER_RULES_BY_SM`): the exact-FP32 SIMT program `router_fp32_dgrad_simt_*` -- one FFMA per product of the FP32 operands read in place, every element of `out` stored once. Every other row, token count and part: deterministic program (one split) | | K2 weight gradient (`router_wgrad`, `nn_t`) | the SIMT token windows of its rows in the router tables (`ROUTER_RULES`; on 212-SM `sm_107a` and 148-SM `sm_100a` / `sm_103a` parts `ROUTER_RULES_BY_SM`): the exact-FP32 SIMT program `router_fp32_gemm_simt_*` on the CUDA cores, every element of `out` stored once; on 148-SM parts at one token at hidden 4096 its compensated form. The rows of `ROUTER_RULES` with a partition: stream-K over the K blocks, in spans alike or in long head spans and short tail spans (the head/tail program) -- whole tiles store the one-split result, tiles split between units TMA reduce-add into `out` (zeroed before the program, on the same stream). Every other row: deterministic program (its six splits) | -| BF16x3 split (`split_fp32_to_bf16x3_`) | the deterministic stage (exact) | +| BF16x3 split of the K2 second operand (inside the K2 launches) | the deterministic stage (exact) | -*Table: the routing of `deterministic=False` (`cake_backend.ROUTES`).* +*Table: the routing of `deterministic=False` (`cake_backend.ROUTES`), per operation and row.* ## Numerics, architectures and parallelism @@ -46,11 +77,12 @@ deterministic already.* | accuracy | the deterministic package's budgets against FP64 of the exact operands, worst over repeated runs: BF16 out `atol = rtol = 1e-2`; FP32 out relative Frobenius error `<= 5e-5` and max error `<= 1e-3 * max abs(ref)`; router `<= 2e-6` relative Frobenius error and no row whose top-8 experts differ from FP64 (outside near ties of the 8th and 9th logits). At the largest token counts some thin K1 rows miss the BF16 (FP32) budget with the deterministic program and with this package's program alike; there the row keeps its route. The deterministic program: BF16 output of `sm_100a` `4096 -> 32`, `4096 -> 48`, `4096 -> 64`, `5120 -> 32`, `5120 -> 48` and of `sm_103a` `4096 -> 32`, `4096 -> 48`, `5120 -> 32`, `5120 -> 48`, `5120 -> 64`, `5120 -> 128` (from 45058, 49153, 65536 or 65537 tokens up by row: the K1 token tables below), and FP32 output of the rows of the odd-block K1 table at an odd number of K blocks from 98305 tokens up. This package's program: BF16 output of `sm_100a` `4096 -> 64` from about 43009 to 65535 tokens and `5120 -> 32` at 47105 to 65536 tokens, of `sm_103a` `4096 -> 32` from about 49153 to 65535 tokens, and on 148-SM parts of `sm_103a` `5120 -> 48` from 49152 tokens, `5120 -> 64` at 47104 and from 49153 tokens and `5120 -> 128` from 65537 tokens | | rounding | each output element is rounded once, from its complete FP32 sum | | subnormals | kept: K1 adds the partials with IEEE round-to-nearest adds, K2 reduces with the TMA reduce-add; a K2 SIMT window (forward, input and weight gradient) accumulates with round-to-nearest FMAs and adds, merges the sums of a CTA's slices (and in a compensated form their compensations) in a fixed order through shared memory (the forward: the sums of a cluster's eight CTAs, through distributed shared memory) and stores every element of `out` once, with no reduce-add; nothing flushes to zero | +| `torch.use_deterministic_algorithms(True)` | PyTorch's alert right before an order-varying launch: K1 at three or more splits (in the shipped tables only `sm_107a` rows, see "run-to-run bits, K1"; or a caller's `splits >= 3`), K2 stream-K with more than two partials per element (the `ROUTER_RULES` partitions, the `"waves"` windows, a caller's `splits >= 3` on `nn_t`). `RuntimeError`, or a `UserWarning` with `warn_only=True` and then the launch. The alert is checked on every launch (prepared and convenience), never at plan or prepare time. The two-split K1 rows, the K2 forward, the K2 SIMT windows and every deterministic program run without an alert | | run-to-run bits, K1 | two splits (every `sm_100a` and `sm_103a` row, `sm_107a` `8192 -> 128`): repeat bit for bit (one add of two values); one split (`sm_107a` rows marked `p = 1` below): repeat bit for bit (whole tiles); three or four splits (the other `sm_107a` rows): vary in the last bits | | run-to-run bits, K2 | forward: in a SIMT window repeats bit for bit (every element stored once, the lanes and the cluster's CTAs combined in a fixed order), else the deterministic program's bits; weight gradient on a `ROUTER_RULES` partition: varies in the last bits (more than two partials per element); a SIMT window (weight or input gradient): repeats bit for bit (every element stored once, its slices merged in a fixed order); with an explicit forward `splits` of 1 or 2: repeats bit for bit (`0 + a + b == 0 + b + a`), so recomputed router logits and top-k match | | architectures | `sm_100a`, `sm_103a`, `sm_107a`: one `-gencode=arch=compute_,code=sm_` build per program, minimum CUDA 12.8 / 12.9 / 13.4 (`build.json`) | | streams and threads | planning takes a module lock; the convenience path keeps one workspace arena per (device, stream); prepared launches own their workspace | -| CUDA graphs | prepared launches allocate nothing on the device and capture into a CUDA graph (K2: the split, the memset of `out`, the program); no host synchronization. A `deterministic=False` convenience-path call that this package launches from its arenas (its own programs, and the deterministic programs it launches through its binding, see [Workspace](#workspace)), captured into a CUDA graph after a warm-up call on the capturing stream (so the capture plans and allocates nothing), stays valid until `clear_convenience_caches()` or `use_deterministic_backend`, which invalidate it: the arenas keep every buffer a launch was bound to, through growths and cache evictions. A call this package hands to the deterministic package's convenience path -- operands whose launch views the binding cannot rebuild from the call's tensors (for example two operands that alias), a batched forward on that package's swapped small-token program, a call under `torch.use_deterministic_algorithms(True, warn_only=True)` -- follows that package's rule (capture its prepared launch, not its convenience path). The graphs captured on one stream share its arena (`torch.cuda.graph` without `stream=` captures every graph on one shared stream): replay them one at a time, and not concurrently with calls on that stream | +| CUDA graphs | prepared launches allocate nothing on the device and capture into a CUDA graph (K2: the split, the memset of `out`, the program); no host synchronization. A `deterministic=False` convenience-path call that this package launches from its arenas (its own programs, and the deterministic programs it launches through its binding, see [Workspace](#workspace)), captured into a CUDA graph after a warm-up call on the capturing stream (so the capture plans and allocates nothing), stays valid until `clear_convenience_caches()` or `use_deterministic_backend`, which invalidate it: the arenas keep every buffer a launch was bound to, through growths and cache evictions. A call on a row that runs the deterministic program, which this package hands to the deterministic package's entry point of the same operation -- operands whose launch views the binding cannot rebuild from the call's tensors (for example two operands that alias), a batched forward on that package's swapped small-token program -- follows that package's rule (capture its prepared launch, not its convenience path). The graphs captured on one stream share its arena (`torch.cuda.graph` without `stream=` captures every graph on one shared stream): replay them one at a time, and not concurrently with calls on that stream | | other kernels on the device | the programs walk their work persistently over at most `sm_count // 2` CTA pairs and wait on no CTA that is not resident (the forward's SIMT program waits only on the CTAs of its own cluster, which the hardware schedules together); every routed row was measured next to a kernel holding 8, 16 or 32 SMs on another stream (no more than 3 % slower than the deterministic plan there) | | context parallelism (KCP) | not applicable: a local operator on the caller's CP-local shards; no collectives; the reduction of parameter gradients across ranks stays with the caller | | autograd, checkpointing, gradient accumulation | none in either package: callers compose the forward and both gradients; `out` is overwritten, not accumulated into | @@ -815,7 +847,7 @@ The K1 forward rows of `K1_FORWARD_RULES_BY_SM` run, with `deterministic=False`, single-CTA program (`K1_FORWARD_SINGLE_CTA_KNOBS`: one 128-row tile per CTA, 128-wide column tiles, no TMA cache hints; every tile's whole K in one CTA) in place of that package's plan of the row (CTA pairs of 256-wide tiles with the K range split two ways, the partials added in a fixed order). The forward stays deterministic: it repeats bit for bit -and a recomputed forward reproduces its result, with other bits than `deterministic=True`; both meet the BF16 budget +and a recomputed forward reproduces its result, with other bits than the deterministic plan; both meet the BF16 budget against FP64 (10 seeds, 20 runs each, worst run). A row is in the table only at a token count and SM count where it was measured and met this rule in every run: the forward alone faster than the deterministic plan (a one-sided paired sign test, p < 0.05, median ratio below 0.99, eager launches and CUDA graph replays) and no more than 3 % slower than @@ -900,46 +932,42 @@ re-times `det` and `atomic` next to a kernel holding `--resident-sms` SMs on ano ## Names -This import path held the dense projection GEMM package until #8046 renamed it to `dense_projection_det_gemm`; #8046 -left it as an alias of the renamed modules. The names keep their spelling and, with the default, their behavior: +This package has no alias, re-export or read-through of a deterministic name. Each name below is this package's +own object. The deterministic package's names (`*_det_*`, `DetGemmPlan`, `ROW_RULES`, `stream_k_plan`, +`wgrad_views`, ...) are reached only in `flashinfer.experimental.dense_projection_det_gemm`; reading one of them +here raises `AttributeError`. -| name before #8046 (`dense_projection_gemm.cake_backend`) | deterministic package (`dense_projection_det_gemm.cake_backend`) | this path | -|---|---|---| -| `dense_projection_gemm` | `dense_projection_det_gemm` | `dense_projection_gemm(..., deterministic=True)` | -| `projection_forward`, `projection_dgrad`, `projection_wgrad` | `projection_det_forward`, `projection_det_dgrad`, `projection_det_wgrad` | the same names with `deterministic=` | -| `prepare_dense_projection_gemm`, `prepare_projection_wgrad` | `prepare_dense_projection_det_gemm`, `prepare_projection_det_wgrad` | the same names with `deterministic=` | -| `plan_dense_projection_gemm` | `plan_dense_projection_det_gemm` | the same name with `deterministic=` | -| `router_fp32_gemm`, `router_forward`, `router_dgrad`, `router_wgrad` | `router_fp32_det_gemm`, `router_det_forward`, `router_det_dgrad`, `router_det_wgrad` | the same names with `deterministic=` | -| `prepare_router_fp32_gemm`, `plan_router_fp32_gemm` | `prepare_router_fp32_det_gemm`, `plan_router_fp32_det_gemm` | the same names with `deterministic=` | -| `split_fp32_to_bf16x3_` | `split_fp32_to_bf16x3_det_` | the same name with `deterministic=` (the deterministic stage in both modes) | -| `instance_symbol`, `router_symbol` | `instance_symbol`, `router_symbol` | the same names with `deterministic=` | -| `GemmPlan`, `PreparedGemm` | `DetGemmPlan`, `PreparedDetGemm` | read through; `deterministic=False` plans `AtomicGemmPlan`, `PreparedAtomicGemm` | -| `RouterPlan`, `PreparedRouterGemm` | `RouterDetPlan`, `PreparedRouterDetGemm` | read through; `deterministic=False` plans `AtomicRouterPlan`, `PreparedAtomicRouterGemm` | -| every other module name (`ROW_RULES`, `stream_k_plan`, `wgrad_views`, ...) | the same name | read through to the deterministic package | -| program templates `dense_proj_gemm_*`, `router_fp32_gemm_`, `split_fp32_bf16x3` | `dense_proj_det_gemm_*`, `router_fp32_det_gemm_`, `split_fp32_det_bf16x3` | this package's templates are new: `dense_proj_gemm_atomic_nn_n<64\|128\|256>[_f32]_t`, `dense_proj_gemm_atomic_nn_n192_f32` (row-major store, `sm_107a`), `router_fp32_gemm_atomic_{kk,nn_t,nn_t_ht_c3_nreg}`, `router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nreg` (`sm_103a`) | - -*Table: the names before #8046, their deterministic names, and what this path offers -(`cake_backend.DETERMINISTIC_NAMES`).* - -| name (`flashinfer.experimental.dense_projection_gemm`) | default | `deterministic=False` | +| operation | this package (`deterministic=False` only) | deterministic package (no `deterministic` keyword) | |---|---|---| -| the entry points of the table above | the deterministic package's `*_det_*` object | this package's schedule (routing table above) | -| `cake_backend.clear_convenience_caches` | clears the caches and arenas of both packages (invalidates the CUDA graphs that captured a convenience-path call) | -- | -| `cake_jit.{MODULES, KERNELS, ARG_PLANS}` | this package's programs only | -- | -| `cake_jit.select_module`, `cake_jit.load_cake_dense_projection_gemm_module` | a deterministic template or program resolves through the deterministic package | -- | -| `cake_build` | this package's loader (byte-identical to every Cake package's) | -- | +| K1 GEMM | `dense_projection_gemm`, `plan_dense_projection_gemm`, `prepare_dense_projection_gemm` | `dense_projection_det_gemm`, `plan_dense_projection_det_gemm`, `prepare_dense_projection_det_gemm` | +| K1 forward, input gradient, weight gradient | `projection_forward`, `projection_dgrad`, `projection_wgrad`, `prepare_projection_wgrad` | `projection_det_forward`, `projection_det_dgrad`, `projection_det_wgrad`, `prepare_projection_det_wgrad` | +| K2 GEMM | `router_fp32_gemm`, `plan_router_fp32_gemm`, `prepare_router_fp32_gemm` | `router_fp32_det_gemm`, `plan_router_fp32_det_gemm`, `prepare_router_fp32_det_gemm` | +| K2 forward, input gradient, weight gradient | `router_forward`, `router_dgrad`, `router_wgrad` | `router_det_forward`, `router_det_dgrad`, `router_det_wgrad` | +| BF16x3 split | -- (inside the K2 launches) | `split_fp32_to_bf16x3_det_` | +| plan and prepared types | `AtomicGemmPlan`, `PreparedAtomicGemm`, `AtomicRouterPlan`, `PreparedAtomicRouterGemm` (a row on the deterministic program returns that package's types, see [Mode and dispatch rules](#mode-and-dispatch-rules)) | `DetGemmPlan`, `PreparedDetGemm`, `RouterDetPlan`, `PreparedRouterDetGemm` | +| program templates | `dense_proj_gemm_atomic_nn_n<64\|128\|256>[_f32]_t`, `dense_proj_gemm_atomic_nn_n<128\|256>_f32_t_c3`, `dense_proj_gemm_atomic_nn_n192_f32[_c3]` (row-major store, `sm_107a`), `router_fp32_gemm_atomic_{kk,nn_t,nn_t_ht_c3_nreg}`, `router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nreg`, `router_fp32_gemm_simt_*`, `router_fp32_dgrad_simt_*`, `router_fp32_fwd_simt_*`, `router_partials_sum`; `router_symbol(a_mn, b_mn, out_t)` names the stream-K template of `kk` / `nn_t` | `dense_proj_det_gemm_*`, `router_fp32_det_gemm_`, `split_fp32_det_bf16x3`, ... | + +*Table: the entry points of the two packages, per operation.* + +| name (`flashinfer.experimental.dense_projection_gemm`) | reaches | +|---|---| +| `cake_backend.clear_convenience_caches` | clears the caches and arenas of both packages (invalidates the CUDA graphs that captured a convenience-path call) | +| `cake_backend.bind_launch` | binds a program of this package (a deterministic program binds through that package's `bind_launch`) | +| `cake_jit.{MODULES, KERNELS, ARG_PLANS}`, `cake_jit.select_module`, `cake_jit.load_cake_dense_projection_gemm_module` | this package's programs only (a deterministic template or program name raises; use the deterministic package's `cake_jit`) | +| `cake_backend.use_deterministic_backend`, `cake_backend.deterministic_backend` | bind and return the deterministic package's `cake_backend` (the package whose fixed-order programs this package reuses) | +| `cake_build` | this package's loader (byte-identical to every Cake package's) | -*Table: the names of this path and what they reach.* +*Table: the helpers of this package and what they reach.* -A module-level assignment through this path (for example `CONVENIENCE_CACHE_SIZE`) does not reach the deterministic -package; set it there. +A module-level assignment in this package (for example a rule table) does not reach the deterministic package; set +the deterministic package's names there. ## Layout of this package | file | content | |---|---| | `__init__.py` | imports nothing; `cake_backend`, `cake_jit`, `cake_build` load on first attribute access | -| `cake_backend.py` | entry points, mode resolution, routing, the planners, rule tables and the convenience state of this package's programs | +| `cake_backend.py` | entry points, mode rules, routing, the planners, rule tables and the convenience state of this package's programs | | `cake_jit.py` | the registry of this package's programs (`ARG_PLANS`, `MODULES`, `KERNELS`) and the loader | | `cake_build.py`, `build.json`, `csrc/cake_dense_projection_gemm/` | the `cake-build/1` loader, manifest and generated sources of this package's programs (written by the Cake exporter) | diff --git a/flashinfer/experimental/dense_projection_gemm/__init__.py b/flashinfer/experimental/dense_projection_gemm/__init__.py index 49a43e611c9..7bbaa3feae9 100644 --- a/flashinfer/experimental/dense_projection_gemm/__init__.py +++ b/flashinfer/experimental/dense_projection_gemm/__init__.py @@ -14,16 +14,16 @@ limitations under the License. """ -# Experimental Cake backend: the GLM-5.2 dense projection GEMMs for training -# and the FP32 router-gate GEMM on SM100 / SM103 / SM107 with an explicit -# numerical mode (flashinfer-ai/flashinfer#5677). Every entry point of -# ``cake_backend`` takes the keyword-only ``deterministic: bool = True``: the -# default calls the deterministic package -# (``flashinfer.experimental.dense_projection_det_gemm``: the same programs and -# the same bits as before this package existed); ``deterministic=False`` selects -# the programs of this package, which add split-K partials in the order they -# finish. Importing this package imports nothing: ``cake_backend``, -# ``cake_jit`` and ``cake_build`` load on first attribute access. +# Experimental Cake backend, non-deterministic: the GLM-5.2 dense projection +# GEMMs for training and the FP32 router-gate GEMM on SM100 / SM103 / SM107 +# (flashinfer-ai/flashinfer#5677), with programs that add split-K partials in +# the order they finish. Every compute entry point of ``cake_backend`` takes +# the keyword-only ``deterministic: bool = True`` and accepts only +# ``deterministic=False``: any other value, the default included, raises +# ``ValueError`` naming the entry point of the deterministic package +# ``flashinfer.experimental.dense_projection_det_gemm``. Importing this package +# imports nothing: ``cake_backend``, ``cake_jit`` and ``cake_build`` load on +# first attribute access. __all__ = ["cake_backend", "cake_build", "cake_jit"] diff --git a/flashinfer/experimental/dense_projection_gemm/cake_backend.py b/flashinfer/experimental/dense_projection_gemm/cake_backend.py index b654e12ccac..d55b279363d 100644 --- a/flashinfer/experimental/dense_projection_gemm/cake_backend.py +++ b/flashinfer/experimental/dense_projection_gemm/cake_backend.py @@ -13,25 +13,28 @@ See the License for the specific language governing permissions and limitations under the License. -Cake backend with an explicit numerical mode: the GLM-5.2 dense projection GEMMs for training (K1) and the FP32 -router-gate GEMM (K2) on SM100 / SM103 / SM107. +Cake backend, non-deterministic: the GLM-5.2 dense projection GEMMs for training (K1) and the FP32 router-gate GEMM +(K2) on SM100 / SM103 / SM107. The deterministic implementation is the package ``dense_projection_det_gemm`` (flashinfer-ai/flashinfer#5677 and -#7043; named so by #8046): its stream-K and split-K reductions add the partials in a fixed order, so every result is -bitwise reproducible run to run. This package keeps that package's operations, view contract, dtypes, accuracy -budgets and host paths (prepared and convenience) and adds a second numerical mode whose split-K reductions add the -partials in the order they finish. Only run-to-run bitwise reproducibility is relaxed: the products, the FP32 -accumulation, the rounding points (an output is rounded once, from the complete FP32 sum), the masks, the layouts -and the dynamic shapes are the deterministic package's, and no reduction flushes a subnormal. - -Numerical mode --------------- -Every entry point takes the keyword-only argument ``deterministic: bool = True``: - -* ``deterministic=True`` (the default) calls the deterministic package's canonical object with the same arguments - (:data:`DETERMINISTIC_NAMES`: ``projection_forward`` -> ``projection_det_forward``, ...): the same planner, the - same programs, the same bits and the same errors. The names this import path had before #8046 keep their behavior. -* ``deterministic=False`` runs the schedule of this package for an operation :data:`ROUTES` marks ``"atomic"``: +#7043): its stream-K and split-K reductions add the partials in a fixed order, so every result is bitwise +reproducible run to run. This package keeps that package's operations, view contract, dtypes, accuracy budgets and +host paths (prepared and convenience) and adds programs whose split-K reductions add the partials in the order they +finish. The products, the FP32 accumulation, the rounding points (an output is rounded once, from the complete FP32 +sum), the masks, the layouts and the dynamic shapes are the deterministic package's, and no reduction flushes a +subnormal. A result of this package carries no run-to-run bitwise guarantee on any row, also on a row that runs a +fixed-order program today. + +Mode and dispatch rules +----------------------- +Every compute entry point takes the keyword-only argument ``deterministic: bool = True`` and accepts only +``deterministic=False``: + +* Any other value -- the default, ``True``, ``None``, ``0``, ``1``, ``"False"`` -- raises ``ValueError`` before + anything is planned, allocated or launched. The message names the deterministic package's entry point of the + operation (``projection_wgrad`` -> ``projection_det_wgrad``, ...). The keyword never selects the deterministic + package, and nothing falls back to it. +* ``deterministic=False`` runs the routes of :data:`ROUTES`: - the K1 weight gradient (both operands MN-major, ``K = T`` ragged), launched as the transposed-store GEMM ``out[l, n, m]``: each of ``L x m_tiles x n_tiles`` tiles is cut into ``num_splits`` K ranges; every split stores @@ -58,33 +61,41 @@ stored once (other bits than the deterministic program's, within the deterministic package's router FP64 budget); every other forward row and token count runs the deterministic program. - Every other operation -- the K1 forward and input gradient, the BF16x3 split -- runs the deterministic programs in - both modes: a recomputed forward must reproduce its result, and where the deterministic plan of a K1 forward or input - gradient splits K across CTAs (its automatic two-way split), it adds the partials in a fixed order. The K1 forward - rows of :data:`K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM ``sm_100a`` / ``sm_103a`` parts) run another - plan of those programs with ``deterministic=False``: the single-CTA program, every tile's whole K in one CTA. On the - convenience path a call of ``deterministic=False`` that - runs a deterministic program (the K1 forward and input gradient, a K2 row without a verdict) launches it through a - binding of this package (the deterministic package's plan and programs, and the reduction of its split-K partials in - the deterministic package's order by this package's program: the same bits, fewer host operations per call). -* With ``torch.use_deterministic_algorithms(True)`` an explicit ``deterministic=False`` raises ``RuntimeError``; with - ``warn_only=True`` it warns and runs the deterministic path (:func:`_resolve_deterministic`). The default does not - consult the PyTorch setting: it is deterministic already. + Every other operation -- the K1 forward and input gradient, the BF16x3 split -- runs the deterministic programs: a + recomputed forward must reproduce its result, and where the deterministic plan of a K1 forward or input gradient + splits K across CTAs (its automatic two-way split), it adds the partials in a fixed order. The K1 forward rows of + :data:`K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM ``sm_100a`` / ``sm_103a`` parts) run another plan of + those programs: the single-CTA program, every tile's whole K in one CTA. +* Under ``torch.use_deterministic_algorithms(True)`` a launch whose result can vary run to run -- the K1 split-K + program at three or more splits, a K2 stream-K partition that gives an element more than two partials -- raises + PyTorch's alert right before it launches (:func:`_alert_order_varying`): ``RuntimeError``, or with + ``warn_only=True`` a ``UserWarning``, after which the launch runs. The fixed-order launches (every other route) + do not alert. The PyTorch setting never changes which program a call runs. + +The helpers of this package (the planners, the rule tables, :func:`router_symbol`, :func:`bind_launch`, ...) take no +mode keyword. The two packages ---------------- This package registers only the programs of its own schedule (``cake_jit``: programs -``cake_dense_projection_gemm_``, templates ``dense_proj_gemm_atomic_*`` / ``router_fp32_gemm_atomic_*``); every -deterministic program, and every planner rule of the deterministic package, is that package's (resolved as this -package's sibling on first use, never through an absolute ``flashinfer`` import; a vendored copy binds it with -:func:`use_deterministic_backend`). The deterministic package never imports this one. The module-level names this -path had before #8046 that are not entry points (``ROW_RULES``, ``stream_k_plan``, ``wgrad_views``, the plan and -prepared types ``GemmPlan`` / ``PreparedGemm`` / ``RouterPlan`` / ``PreparedRouterGemm``, ...) read through to the -deterministic package (module ``__getattr__``); assigning one here does not change the deterministic package's. -:func:`clear_convenience_caches` clears the convenience caches and workspace arenas of both packages; a convenience-path -call that this package launches from its arenas, captured into a CUDA graph, stays valid until then (the arenas keep -every buffer a launch was bound to, through growths and cache evictions); a call it hands to the deterministic package's -convenience path follows that package's rule (:func:`clear_convenience_caches`). +``cake_dense_projection_gemm_``, templates ``dense_proj_gemm_atomic_*`` / ``router_fp32_gemm_*``); every +deterministic program, and every planner rule of the deterministic package, is that package's. A +``deterministic=False`` call reuses them on the rows that run the deterministic programs (above): it launches the +deterministic plan and program through a binding of this package (the same bits, fewer host operations per call), +or, where that binding cannot rebuild the call's launch views from its tensors (two operands that alias) and for a +batched forward that the deterministic package runs on its swapped small-token program, it calls the deterministic +package's entry point of the same operation (the same program, the same bits). On such a row ``plan_*`` / +``prepare_*`` return the deterministic package's plan or prepared launch. A row that runs a program of this package +never goes to the deterministic package. The deterministic package is imported at run time as this package's +sibling, never through an absolute ``flashinfer`` import; a vendored copy binds it with +:func:`use_deterministic_backend`. The deterministic package never imports this one, and this module re-exports no +name of it. :func:`clear_convenience_caches` clears the convenience caches and workspace arenas of both packages; a +convenience-path call that this package launches from its arenas, captured into a CUDA graph, stays valid until then +(the arenas keep every buffer a launch was bound to, through growths and cache evictions); a call it hands to the +deterministic package's entry point follows that package's rule (:func:`clear_convenience_caches`). + +The tensors passed between the forward and the gradients carry no mode: recompute a checkpointed forward with the +package that made the saved one (same package only). """ from __future__ import annotations @@ -97,7 +108,7 @@ from collections import OrderedDict from dataclasses import dataclass, field from functools import cached_property -from types import ModuleType +from types import FrameType, ModuleType from typing import Any, Callable, Optional, TypeVar, Union import torch @@ -105,7 +116,7 @@ from . import cake_jit -# The registry names this module had before #8046, re-exported from this package's own registry. +# This package's own registry and loader, re-exported (its programs only). from .cake_jit import KERNELS as KERNELS from .cake_jit import MODULES as MODULES from .cake_jit import TRACKING_ISSUE as TRACKING_ISSUE @@ -116,37 +127,25 @@ from .cake_jit import select_module as select_module # --------------------------------------------------------------------------- -# The deterministic package and the numerical mode +# The deterministic package and the mode rules # --------------------------------------------------------------------------- -#: Each name of this module's entry points and types -> its canonical object in the deterministic package's -#: ``cake_backend`` (the former names #8046 kept there as aliases). ``deterministic=True`` calls that object. -DETERMINISTIC_NAMES: dict[str, str] = { - "GemmPlan": "DetGemmPlan", - "PreparedGemm": "PreparedDetGemm", - "PreparedRouterGemm": "PreparedRouterDetGemm", - "RouterPlan": "RouterDetPlan", - "dense_projection_gemm": "dense_projection_det_gemm", - "plan_dense_projection_gemm": "plan_dense_projection_det_gemm", - "plan_router_fp32_gemm": "plan_router_fp32_det_gemm", - "prepare_dense_projection_gemm": "prepare_dense_projection_det_gemm", - "prepare_projection_wgrad": "prepare_projection_det_wgrad", - "prepare_router_fp32_gemm": "prepare_router_fp32_det_gemm", - "projection_dgrad": "projection_det_dgrad", - "projection_forward": "projection_det_forward", - "projection_wgrad": "projection_det_wgrad", - "router_dgrad": "router_det_dgrad", - "router_forward": "router_det_forward", - "router_fp32_gemm": "router_fp32_det_gemm", - "router_wgrad": "router_det_wgrad", - "split_fp32_to_bf16x3_": "split_fp32_to_bf16x3_det_", - "instance_symbol": "instance_symbol", - "router_symbol": "router_symbol", -} -#: The types among :data:`DETERMINISTIC_NAMES`: read through to the deterministic package's classes. -_DET_TYPES = ("GemmPlan", "PreparedGemm", "PreparedRouterGemm", "RouterPlan") -#: Names spelling the ``det`` token live in the deterministic package only; this module does not read them through. -_DET_SPELLED = re.compile(r"(?:^|_)det(?:_|$)|Det") +#: The deterministic package's entry points this package calls on the rows that run the deterministic programs +#: (see "The two packages" above); :func:`use_deterministic_backend` checks that a bound module has them. +_DET_ENTRIES = ( + "dense_projection_det_gemm", + "plan_dense_projection_det_gemm", + "plan_router_fp32_det_gemm", + "prepare_dense_projection_det_gemm", + "prepare_projection_det_wgrad", + "prepare_router_fp32_det_gemm", + "projection_det_dgrad", + "projection_det_forward", + "projection_det_wgrad", + "router_fp32_det_gemm", + "split_fp32_to_bf16x3_det_", + "clear_convenience_caches", +) def use_deterministic_backend(module: ModuleType) -> None: @@ -157,9 +156,7 @@ def use_deterministic_backend(module: ModuleType) -> None: caches and arenas (:func:`_drop_convenience_entries`): their entries hold the plans and programs of the package bound before. Like :func:`clear_convenience_caches`, it invalidates the CUDA graphs that captured a convenience-path call of this package.""" - missing = [ - name for key, name in DETERMINISTIC_NAMES.items() if not hasattr(module, name) - ] + missing = [name for name in _DET_ENTRIES if not hasattr(module, name)] if missing: raise TypeError( f"{getattr(module, '__name__', module)!r} is not the deterministic dense projection GEMM backend: " @@ -188,33 +185,85 @@ def deterministic_backend() -> ModuleType: return module -def _det_entry(name: str) -> Callable[..., Any]: - return getattr(deterministic_backend(), DETERMINISTIC_NAMES[name]) +def _det_package_name() -> str: + """The dotted name of the deterministic package (the bound one, else this package's sibling), as text: the mode + messages name it without importing it.""" + package = cake_jit._det_package + if package is not None: + return package + parent = __name__.rpartition(".")[0].rpartition(".")[0] + return ( + f"{parent}.{cake_jit.DETERMINISTIC_PACKAGE}" + if parent + else cake_jit.DETERMINISTIC_PACKAGE + ) -def _resolve_deterministic(deterministic: bool, api: str) -> bool: - """Whether a call of ``api`` runs the deterministic path. +def _require_false(deterministic: Any, entry: str, det_entry: str) -> None: + """The mode rule of every compute entry point, its first statement: this package computes non-deterministically, + so ``entry`` takes only ``deterministic=False``. Any other value (the default ``True``, ``None``, ``0``, ``1``, + ``"False"``, a NumPy bool) raises ``ValueError`` naming the deterministic package's entry point ``det_entry``; + the keyword never selects the deterministic package.""" + if deterministic is not False: + package = __name__.rpartition(".")[0] or __name__ + raise ValueError( + f"{entry}: {package} runs the non-deterministic package only and needs deterministic=False (got " + f"{deterministic!r}): pass deterministic=False for the non-deterministic path; the deterministic entry " + f"is {_det_package_name()}.cake_backend.{det_entry}" + ) - ``deterministic=True`` does. ``deterministic=False`` runs this package's schedule unless - ``torch.use_deterministic_algorithms(True)`` is in effect: then the call raises ``RuntimeError``, or, with - ``warn_only=True``, warns (``UserWarning``) and runs the deterministic path. A call that asked for determinism - never runs a program of this package. - """ - if deterministic: - return True + +def _warning_stacklevel() -> int: + """The ``stacklevel`` of a warning raised in this module that points at the first frame outside it (the caller's + line): the depth of this module's frames between the warning and the caller.""" + level = 1 + frame: Optional[FrameType] = sys._getframe(1) + while frame is not None and frame.f_globals.get("__name__") == __name__: + level += 1 + frame = frame.f_back + return level + + +def _alert_order_varying(launch: str) -> None: + """PyTorch's alert for an operation without a deterministic implementation, raised right before this package + launches ``launch``, a program whose result can vary run to run (the K1 split-K program at three or more splits, + a K2 stream-K partition that gives an element more than two partials), while PyTorch's deterministic algorithms + are enabled: a ``RuntimeError``, or a ``UserWarning`` under ``torch.use_deterministic_algorithms(True, + warn_only=True)``, after which the launch runs; nothing otherwise. The alert of PyTorch's own operators (PyTorch + exposes no public Python entry for it). Fixed-order launches do not call it, and it never changes the program + that runs.""" if not torch.are_deterministic_algorithms_enabled(): - return False - message = ( - f"{__name__}.{api}(deterministic=False) adds split-K partials in the order they finish, but " - "torch.use_deterministic_algorithms(True) is in effect; call it with deterministic=True or use " - f"{cake_jit.DETERMINISTIC_PACKAGE}.cake_backend.{DETERMINISTIC_NAMES.get(api, api)}" - ) + return if torch.is_deterministic_algorithms_warn_only_enabled(): warnings.warn( - f"{message}; running the deterministic path", UserWarning, stacklevel=3 + f"{__name__}: {launch} does not have a deterministic implementation, but you set " + "'torch.use_deterministic_algorithms(True, warn_only=True)'; " + f"{_det_package_name()} is the deterministic package", + UserWarning, + stacklevel=_warning_stacklevel(), ) - return True - raise RuntimeError(message) + return + raise RuntimeError( + f"{__name__}: {launch} does not have a deterministic implementation, but you set " + f"'torch.use_deterministic_algorithms(True)'; {_det_package_name()} is the deterministic package, or use " + "the 'warn_only=True' option if that is acceptable for your application" + ) + + +def _k1_launch_text(plan: Any) -> str: + """The launch an order-varying K1 plan (:class:`AtomicGemmPlan`) names in :func:`_alert_order_varying`.""" + return ( + f"the K1 split-K program {plan.template} at {plan.splits} splits (the partials added in the order the " + "splits finish)" + ) + + +def _k2_launch_text(plan: Any) -> str: + """The launch an order-varying K2 plan (:class:`AtomicRouterPlan`) names in :func:`_alert_order_varying`.""" + return ( + f"the router stream-K program {plan.template} with {plan.units} units over {plan.tiles} tiles (more than " + "two partials per element, added in the order they finish)" + ) # --------------------------------------------------------------------------- @@ -223,7 +272,8 @@ def _resolve_deterministic(deterministic: bool, api: str) -> bool: #: The program family each operation runs with ``deterministic=False``: ``"atomic"`` -- this package's planner and #: programs (the planner keeps every row without a measured verdict on the deterministic program, see -#: :func:`k1_plan` / :func:`router_partition`) -- or ``"det"`` -- the deterministic package's programs in both modes. K1 +#: :func:`k1_plan` / :func:`router_partition`) -- or ``"det"`` -- the deterministic package's programs (fixed order). +#: The mode keyword never reads this table: it only takes ``deterministic=False`` (:func:`_require_false`). K1 #: operations are keyed by the operand layouts of the launched GEMM (``kk`` forward, ``kn`` input gradient, ``nn`` #: weight gradient; ``nk`` has no training use), K2 by the layout class of ``router_layout_class``. ROUTES: dict[str, str] = { @@ -1059,7 +1109,7 @@ def k1_block_n(N: int) -> int: #: ``sm_103a`` parts the deterministic plan of these rows splits K two ways over CTA pairs of 256-wide tiles (stream-K, #: the partials added in a fixed order); the single-CTA program (:data:`K1_FORWARD_SINGLE_CTA_KNOBS`) keeps every tile's #: whole K in one CTA, so a forward still repeats bit for bit and a recomputed forward reproduces its result, with other -#: bits than ``deterministic=True`` (within the deterministic package's BF16 budget against FP64). Measured at exactly +#: bits than the deterministic package's plan (within the deterministic package's BF16 budget against FP64). Measured at exactly #: these token counts (no other token count, SM count or row): the single op faster than the deterministic plan in every #: run (a one-sided paired sign test, p < 0.05, median ratio below 0.99, eager and CUDA graph), no more than 3 % slower #: than it next to a kernel holding 8, 16 or 32 SMs on another stream, and the training steps that contain the forward @@ -1166,6 +1216,12 @@ def tiles(self) -> int: def items(self) -> int: return self.tiles * self.splits + @property + def reproducible(self) -> bool: + """Whether the result repeats bit for bit: one split stores whole tiles (its chains added in order), two add + ``own + other`` (one add of two values); three or more add the partials in the order the splits finish.""" + return self.splits <= 2 + @property def grid(self) -> tuple[int, int, int]: return (2 * min(self.items, self.sm_pairs), 1, 1) @@ -2067,8 +2123,8 @@ def router_partition( if layout == "kk" and s > ROUTER_FORWARD_MAX_SPLITS: raise ValueError( f"router_fp32_gemm(deterministic=False): the forward takes at most {ROUTER_FORWARD_MAX_SPLITS} " - f"splits (more partials per element do not repeat bit for bit under recompute), got {s}; use " - "deterministic=True for more" + f"splits (more partials per element do not repeat bit for bit under recompute), got {s}; the " + f"deterministic entry {_det_package_name()}.cake_backend.router_fp32_det_gemm takes more" ) spec = f"S{s}" else: @@ -2873,11 +2929,9 @@ def _atomic_bind( def bind_launch( module_name: str, values: dict[str, Any], grid: tuple[int, int, int], arch: str ) -> Any: - """Order ``values`` by the generated argument plan of ``module_name`` and load its entry: a program of the - deterministic package binds through that package (its ``bind_launch``), one of this package here. Fails - closed on a keyword the kernel expects that the host does not provide.""" - if cake_jit.is_det_program(module_name): - return deterministic_backend().bind_launch(module_name, values, grid, arch) + """Order ``values`` by the generated argument plan of ``module_name``, a program of this package, and load its + entry (the deterministic package's programs bind through that package's ``bind_launch``). Fails closed on a + keyword the kernel expects that the host does not provide. Takes no mode keyword.""" entry, arguments = _atomic_bind(module_name, values, grid, arch) return _AtomicLaunch(module_name, entry, arguments, grid) @@ -2931,6 +2985,8 @@ def reset_counters(self) -> None: def launch(self) -> torch.Tensor: assert self._launch is not None + if not self.plan.reproducible: + _alert_order_varying(_k1_launch_text(self.plan)) with deterministic_backend()._ffi_stream_context(self.device_index): self._launch() return self.out @@ -2976,6 +3032,9 @@ def fused_split(self) -> bool: def launch(self) -> torch.Tensor: assert self._launch is not None + if not self.plan.reproducible: + # before the split, the fill of ``out`` and the program: a raise leaves ``out`` untouched + _alert_order_varying(_k2_launch_text(self.plan)) det = deterministic_backend() index = self.device_index stream = det._current_raw_stream(index) @@ -3202,9 +3261,10 @@ def _atomic_k2( _, arch, sm_count = _device_and_arch(A, B, out) plan, _ = plan_atomic_router(A, B, out, arch=arch, sm_count=sm_count, splits=splits) if plan is None: + det = deterministic_backend() if op == "plan": - return _det_entry("plan_router_fp32_gemm")(A, B, out, splits=splits) - return _det_entry("prepare_router_fp32_gemm")( + return det.plan_router_fp32_det_gemm(A, B, out, splits=splits) + return det.prepare_router_fp32_det_gemm( A, B, out, splits=splits, **_pdl_kw(enable_pdl) ) if op == "plan": @@ -3216,8 +3276,7 @@ def _atomic_k2( # # The deterministic package's design (one least-recently-used cache per kind, keyed by device, stream and view # geometry; one growable workspace arena per (device, raw stream) shared by every cached entry), with caches and -# arenas of their own: a buffer of this package is never a buffer of the deterministic package, and the names below -# are not names this path had before #8046 (those read through to the deterministic package's state). The cache +# arenas of their own: a buffer of this package is never a buffer of the deterministic package. The cache # bound is the deterministic package's ``CONVENIENCE_CACHE_SIZE``, read on every insertion, so one setting bounds # both packages. A cached entry holds its plan and the resolved program entry, never a caller's tensor: the views # are rebuilt from the caller's tensors on every call. The launches of one stream share its arena; launches on @@ -3590,6 +3649,8 @@ class _CachedAtomic: k1: Optional[_K1Binding] = field(default=None, repr=False) det_plan: Any = field(default=None, repr=False) memo: Any = field(default=None, repr=False) + # the plan's launch can vary run to run (``not plan.reproducible``): PyTorch's alert before every launch + order_varying: bool = False #: An entry of the convenience caches (:data:`_K1_ATOMIC_CACHE`, :data:`_K1_DET_CACHE`, :data:`_K2_ATOMIC_CACHE`, @@ -3763,6 +3824,8 @@ def _k1_cached_launch( entry = det._cache_get(cache, _k1_key(index, stream, (*head, pdl, *tail), *sources)) if entry is None or entry.k1 is None or entry.entry is None: return False + if entry.order_varying: + _alert_order_varying(_k1_launch_text(entry.plan)) _launch_k1_binding(det, entry.entry, entry.k1, entry.arena, sources, index, stream) return True @@ -3837,6 +3900,8 @@ class _CachedDetK1: entry: Optional[Callable[..., Any]] = field(default=None, repr=False) arena: Optional[_AtomicArena] = field(default=None, repr=False) k1: Optional[_K1Binding] = field(default=None, repr=False) + # a deterministic program is fixed-order: never alerts (read by :func:`_k1_cached_launch`) + order_varying: bool = False def _plan_det_k1( @@ -4036,13 +4101,20 @@ def _launch_cached_atomic_k1( k1=_bind_atomic_k1( plan, module_name, arg_plan, A, B, out, arena, pdl ), + order_varying=not plan.reproducible, ) entry = _cache_insert(_K1_ATOMIC_CACHE, key, entry) bound, fn = entry.k1, entry.entry if bound is None or fn is None: if entry.plan is None: + # a row on the deterministic program whose views the binding cannot rebuild: the same program through + # the deterministic package's entry point return det_call() + if entry.order_varying: + _alert_order_varying(_k1_launch_text(entry.plan)) return _launch_atomic_k1_unbound(entry, A, B, out, index, stream, pdl) + if entry.order_varying: + _alert_order_varying(_k1_launch_text(entry.plan)) _launch_k1_binding(det, fn, bound, entry.arena, (A, B, out), index, stream) return out @@ -4261,6 +4333,7 @@ def _plan_atomic_k2( split_arg_plan=split_arg_plan, k2=binding, memo=memo, + order_varying=not plan.reproducible, ) @@ -4293,6 +4366,9 @@ def _launch_atomic_k2_binding( weight gradient's output gradient, the input gradient's router weight) and ``out``.""" arena, fn, bound = entry.arena, entry.entry, entry.k2 assert arena is not None and fn is not None and bound is not None + if entry.order_varying: + # before the split, the fill of ``out`` and the program: a raise leaves ``out`` untouched + _alert_order_varying(_k2_launch_text(entry.plan)) split_fn = entry.split_entry # the split, the memset and the kernel share the arena's BF16x3 stack with arena.lock: @@ -4508,9 +4584,7 @@ def _launch_cached_splitk_k2( _plan_splitk_k2(A, B, out, splits, index, stream, pdl), ) if entry.k2 is None or entry.entry is None or entry.arena is None: - return getattr(det, DETERMINISTIC_NAMES["router_fp32_gemm"])( - A, B, out, splits=splits, enable_pdl=bool(pdl) - ) + return det.router_fp32_det_gemm(A, B, out, splits=splits, enable_pdl=bool(pdl)) a_view, split_src = _launch_views(A, B, entry.plan, views) _launch_splitk_k2_binding(det, entry, out, a_view, split_src, index, stream) return out @@ -4769,10 +4843,10 @@ def clear_convenience_caches() -> None: workspace, so capture such graphs again after it. A captured ``deterministic=False`` convenience-path call that this package launches from its arenas (its programs, and the deterministic programs it launches through its binding) stays valid until then: its arena keeps every buffer a launch was bound to, through growths and cache - evictions. A call this package hands to the deterministic package's convenience path (launch views its binding - cannot rebuild from the call's tensors, such as two operands that alias; a batched forward on that package's - swapped small-token program; ``torch.use_deterministic_algorithms(True, warn_only=True)``) follows that package's - rule: a CUDA graph captures its prepared launch, not its convenience path. The graphs captured on one stream share + evictions. A call on a row that runs the deterministic program, which this package hands to the deterministic + package's entry point of the same operation (launch views its binding cannot rebuild from the call's tensors, such + as two operands that alias; a batched forward on that package's swapped small-token program), follows that + package's rule: a CUDA graph captures its prepared launch, not its convenience path. The graphs captured on one stream share its arena: replay them one at a time, and not concurrently with calls on that stream.""" _drop_convenience_entries() deterministic_backend().clear_convenience_caches() @@ -4800,26 +4874,17 @@ def _splits_kw(splits: Optional[int]) -> dict[str, Any]: return {} if splits is None else {"splits": splits} -def instance_symbol(key: tuple, *, deterministic: bool = True) -> str: - """The instance template of a K1 instance key (``deterministic_backend().instance_key``): the deterministic - package's symbol in both modes. This package's K1 programs are not keyed by the deterministic instance keys; - their templates are :func:`atomic_k1_template` (``dense_proj_gemm_atomic_nn_n[_f32]_t``).""" - _resolve_deterministic(deterministic, "instance_symbol") - return _det_entry("instance_symbol")(key) - - -def router_symbol( - a_mn: bool, b_mn: bool, out_t: bool, *, deterministic: bool = True -) -> str: - """Kernel symbol / registry template of a router instance: ``router_fp32_det_gemm_[_t]`` of the - deterministic package, or with ``deterministic=False`` ``router_fp32_gemm_atomic_[_t]`` for the layout - classes this package reduces (``kk``, ``nn_t``) and the deterministic symbol for ``kn`` (no stream-K program; its - SIMT programs are named per tile, :func:`router_simt_template`).""" - det = _resolve_deterministic(deterministic, "router_symbol") - det_symbol = _det_entry("router_symbol")(a_mn, b_mn, out_t) +def router_symbol(a_mn: bool, b_mn: bool, out_t: bool) -> str: + """Kernel symbol / registry template of this package's router stream-K program for the layout class of + ``(a_mn, b_mn)``: ``router_fp32_gemm_atomic_[_t]`` for the classes this package reduces with it (``kk``, + ``nn_t``). Takes no mode keyword. ``kn`` has no stream-K program (its SIMT programs are named per tile, + :func:`router_simt_template`) and the deterministic planner refuses ``nk``: ``ValueError``.""" layout = _k2_class(a_mn, b_mn) - if det or layout == "k2_kn" or _route(layout) == "det": - return det_symbol + if layout not in ("k2_kk", "k2_nn_t"): + raise ValueError( + f"router_symbol: the layout class {layout[3:]!r} has no stream-K program of this package (the input " + "gradient's SIMT programs are router_simt_template(tile, 'kn'))" + ) return ( cake_jit.TEMPLATE_PREFIXES[1] + ("n" if a_mn else "k") @@ -4836,16 +4901,19 @@ def plan_dense_projection_gemm( deterministic: bool = True, **kwargs: Any, ) -> Any: - """Validate the views and plan the K1 launch: the deterministic package's ``plan_dense_projection_det_gemm`` - with the same keywords (``sm_count``, ``l2_bytes``, ``transposed_out``, the planner knobs, ``arch``); with - ``deterministic=False`` a weight-gradient layout (both operands MN-major) is planned for this package's split-K - program (:class:`AtomicGemmPlan`; the knobs ``splits`` / ``block_n`` / ``chains`` are this package's, the others - apply when the row runs the deterministic program).""" - if ( - _resolve_deterministic(deterministic, "plan_dense_projection_gemm") - or _route(_k1_layout(A, B)) == "det" - ): - return _det_entry("plan_dense_projection_gemm")(A, B, out, **kwargs) + """Validate the views and plan the K1 launch: with ``deterministic=False`` (the only value accepted) a + weight-gradient layout (both operands MN-major) is planned for this package's split-K program + (:class:`AtomicGemmPlan`; the knobs ``splits`` / ``block_n`` / ``chains`` are this package's, the others -- + ``sm_count``, ``l2_bytes``, ``transposed_out``, the planner knobs, ``arch`` -- apply when the row runs the + deterministic program). A row the planner keeps on the deterministic program, and every other layout, returns the + deterministic package's plan (``plan_dense_projection_det_gemm``).""" + _require_false( + deterministic, "plan_dense_projection_gemm", "plan_dense_projection_det_gemm" + ) + if _route(_k1_layout(A, B)) == "det": + return deterministic_backend().plan_dense_projection_det_gemm( + A, B, out, **kwargs + ) own, det_kwargs = _split_knobs(kwargs) return _atomic_k1( "plan", @@ -4856,7 +4924,7 @@ def plan_dense_projection_gemm( own=own, arch=det_kwargs.get("arch"), sm_count=det_kwargs.get("sm_count"), - det_call=lambda: _det_entry("plan_dense_projection_gemm")( + det_call=lambda: deterministic_backend().plan_dense_projection_det_gemm( A, B, out, **det_kwargs ), ) @@ -4870,17 +4938,22 @@ def prepare_dense_projection_gemm( deterministic: bool = True, **kwargs: Any, ) -> Any: - """Validate one K1 binding, plan it for the device and prepare its allocation-free launch: the deterministic - package's ``prepare_dense_projection_det_gemm`` with the same keywords (``transposed_out`` and the planner - knobs); with ``deterministic=False`` a weight-gradient layout prepares this package's split-K program (its - workspace is owned by the prepared object; ``splits`` / ``block_n`` / ``chains`` as for - :func:`plan_dense_projection_gemm`). - ``enable_pdl`` (default: the device supports it) makes either launch a programmatic dependent launch.""" - if ( - _resolve_deterministic(deterministic, "prepare_dense_projection_gemm") - or _route(_k1_layout(A, B)) == "det" - ): - return _det_entry("prepare_dense_projection_gemm")(A, B, out, **kwargs) + """Validate one K1 binding, plan it for the device and prepare its allocation-free launch: with + ``deterministic=False`` (the only value accepted) a weight-gradient layout prepares this package's split-K program + (:class:`PreparedAtomicGemm`; its workspace is owned by the prepared object; ``splits`` / ``block_n`` / ``chains`` + as for :func:`plan_dense_projection_gemm`). A row the planner keeps on the deterministic program, and every other + layout, prepares the deterministic package's ``prepare_dense_projection_det_gemm`` (its keywords: + ``transposed_out`` and the planner knobs). ``enable_pdl`` (default: the device supports it) makes either launch + a programmatic dependent launch.""" + _require_false( + deterministic, + "prepare_dense_projection_gemm", + "prepare_dense_projection_det_gemm", + ) + if _route(_k1_layout(A, B)) == "det": + return deterministic_backend().prepare_dense_projection_det_gemm( + A, B, out, **kwargs + ) own, det_kwargs = _split_knobs(kwargs) return _atomic_k1( "prepare", @@ -4889,7 +4962,7 @@ def prepare_dense_projection_gemm( out, transposed_out=bool(det_kwargs.get("transposed_out", False)), own=own, - det_call=lambda: _det_entry("prepare_dense_projection_gemm")( + det_call=lambda: deterministic_backend().prepare_dense_projection_det_gemm( A, B, out, **det_kwargs ), enable_pdl=det_kwargs.get("enable_pdl"), @@ -4907,24 +4980,16 @@ def dense_projection_gemm( weights_static: bool = False, ) -> torch.Tensor: """``out[l, m, n] = sum_k A[l, m, k] B[l, k, n]`` for strided views (the cached convenience path; see the - deterministic package for the view contract). ``deterministic=False`` runs a weight-gradient layout (both - operands MN-major) on this package's split-K program where the planner takes the row; every other layout and - row runs the deterministic program. ``enable_pdl`` (default: the device supports it) launches either program - as a programmatic dependent launch; each waits for the preceding kernel before its first global access, and - ``weights_static=True`` lets the deterministic swapped small-token forward start its loads of ``B`` before the - wait (see the deterministic package's ``dense_projection_det_gemm``).""" - if _resolve_deterministic(deterministic, "dense_projection_gemm"): - return _det_entry("dense_projection_gemm")( - A, - B, - out, - transposed_out=transposed_out, - enable_pdl=enable_pdl, - **_weights_kw(weights_static), - ) + deterministic package for the view contract). ``deterministic=False`` (the only value accepted) runs a + weight-gradient layout (both operands MN-major) on this package's split-K program where the planner takes the + row; every other layout and row runs the deterministic program. ``enable_pdl`` (default: the device supports it) + launches either program as a programmatic dependent launch; each waits for the preceding kernel before its first + global access, and ``weights_static=True`` lets the deterministic swapped small-token forward start its loads of + ``B`` before the wait (see the deterministic package's ``dense_projection_det_gemm``).""" + _require_false(deterministic, "dense_projection_gemm", "dense_projection_det_gemm") def det_call() -> Any: - return _det_entry("dense_projection_gemm")( + return deterministic_backend().dense_projection_det_gemm( A, B, out, @@ -4968,15 +5033,13 @@ def projection_forward( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, N] = X[T, K] @ W[N, K].T`` (bf16 out by default). The deterministic package's programs in both modes - (a recomputed forward must reproduce its result; where the deterministic plan splits K across CTAs, it adds the - partials in a fixed order); with ``deterministic=False`` through this package's binding of them (:func:`_det_k1`): - the deterministic plan, program and bits, except on a row of :data:`K1_FORWARD_RULES_BY_SM` for the device, which - runs the program the deterministic planner gives it with the table's knobs (:func:`k1_forward_knobs`; other bits, - still the same on every run). - ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - if _resolve_deterministic(deterministic, "projection_forward"): - return _det_entry("projection_forward")(X, W, out, **_pdl_kw(enable_pdl)) + """``out[T, N] = X[T, K] @ W[N, K].T`` (bf16 out by default), ``deterministic=False`` only. The deterministic + package's programs (a recomputed forward must reproduce its result; where the deterministic plan splits K across + CTAs, it adds the partials in a fixed order) through this package's binding of them (:func:`_det_k1`): the + deterministic plan, program and bits, except on a row of :data:`K1_FORWARD_RULES_BY_SM` for the device, which runs + the program the deterministic planner gives it with the table's knobs (:func:`k1_forward_knobs`; other bits, still + the same on every run). ``enable_pdl`` as for :func:`dense_projection_gemm`.""" + _require_false(deterministic, "projection_forward", "projection_det_forward") if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.bfloat16) o = out @@ -4988,7 +5051,9 @@ def projection_forward( (X, W, o), ("fwd",), lambda: (X, W.t(), o, False, _k1_forward_knobs_of(X, W, o)), - lambda: _det_entry("projection_forward")(X, W, o, **_pdl_kw(enable_pdl)), + lambda: deterministic_backend().projection_det_forward( + X, W, o, **_pdl_kw(enable_pdl) + ), enable_pdl, ) @@ -5002,13 +5067,10 @@ def projection_dgrad( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, K] = G[T, N] @ W[N, K]``. The deterministic program in both modes (whole output tiles); with - ``deterministic=False`` through this package's binding of it (:func:`_det_k1`: the same plan, program and bits). - ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - if _resolve_deterministic(deterministic, "projection_dgrad"): - return _det_entry("projection_dgrad")( - G, W, out, out_dtype=out_dtype, **_pdl_kw(enable_pdl) - ) + """``out[T, K] = G[T, N] @ W[N, K]``, ``deterministic=False`` only. The deterministic program (whole output tiles) + through this package's binding of it (:func:`_det_k1`: the same plan, program and bits). ``enable_pdl`` as for + :func:`dense_projection_gemm`.""" + _require_false(deterministic, "projection_dgrad", "projection_det_dgrad") if out is None: out = torch.empty( G.shape[0], @@ -5025,7 +5087,7 @@ def projection_dgrad( (G, W, o), ("dgrad",), lambda: (G, W, o, False, ()), - lambda: _det_entry("projection_dgrad")( + lambda: deterministic_backend().projection_det_dgrad( G, W, o, out_dtype=out_dtype, **_pdl_kw(enable_pdl) ), enable_pdl, @@ -5042,14 +5104,13 @@ def prepare_projection_wgrad( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> Any: - """Prepared ``out[N, K] = G[T, N].T @ X[T, K]``. ``deterministic=False`` prepares this package's split-K - program as the transposed-store GEMM ``X.T @ G`` (``M = K``, ``N = N``) where the planner takes the row; a row it - does not take prepares the deterministic package's ``prepare_projection_det_wgrad`` (its orientation, ``sk``, + """Prepared ``out[N, K] = G[T, N].T @ X[T, K]``, ``deterministic=False`` only: this package's split-K program as + the transposed-store GEMM ``X.T @ G`` (``M = K``, ``N = N``) where the planner takes the row; a row it does not + take prepares the deterministic package's ``prepare_projection_det_wgrad`` (its orientation, ``sk``, ``cta_rows``). ``enable_pdl`` as for :func:`prepare_dense_projection_gemm`.""" - if _resolve_deterministic(deterministic, "prepare_projection_wgrad"): - return _det_entry("prepare_projection_wgrad")( - G, X, out, sk=sk, cta_rows=cta_rows, **_pdl_kw(enable_pdl) - ) + _require_false( + deterministic, "prepare_projection_wgrad", "prepare_projection_det_wgrad" + ) return _atomic_k1( "prepare", X.t(), @@ -5057,7 +5118,7 @@ def prepare_projection_wgrad( out, transposed_out=True, own={}, - det_call=lambda: _det_entry("prepare_projection_wgrad")( + det_call=lambda: deterministic_backend().prepare_projection_det_wgrad( G, X, out, sk=sk, cta_rows=cta_rows, **_pdl_kw(enable_pdl) ), enable_pdl=enable_pdl, @@ -5075,21 +5136,12 @@ def projection_wgrad( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[N, K] = G[T, N].T @ X[T, K]`` (ragged reduction length ``T``). ``deterministic=False`` runs this + """``out[N, K] = G[T, N].T @ X[T, K]`` (ragged reduction length ``T``), ``deterministic=False`` only: this package's split-K program where the planner takes the row (the split partials summed in the order they finish, the output rounded once), else the deterministic package's ``projection_det_wgrad`` program (its orientation, ``sk`` and ``cta_rows``) through this package's binding of it (:func:`_launch_cached_atomic_k1`); ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - if _resolve_deterministic(deterministic, "projection_wgrad"): - return _det_entry("projection_wgrad")( - G, - X, - out, - out_dtype=out_dtype, - sk=sk, - cta_rows=cta_rows, - **_pdl_kw(enable_pdl), - ) + _require_false(deterministic, "projection_wgrad", "projection_det_wgrad") if out is None: out = torch.empty( int(G.shape[1]), @@ -5115,7 +5167,7 @@ def projection_wgrad( out, transposed_out=True, own={}, - det_call=lambda: _det_entry("projection_wgrad")( + det_call=lambda: deterministic_backend().projection_det_wgrad( G, X, out, @@ -5154,14 +5206,15 @@ def plan_router_fp32_gemm( splits: Optional[int] = None, deterministic: bool = True, ) -> Any: - """Classify the fp32 views and plan the K2 launch (the deterministic package's ``plan_router_fp32_det_gemm``); - ``deterministic=False`` plans this package's stream-K program for the ``kk`` / ``nn_t`` classes and the SIMT - programs in their windows (:class:`AtomicRouterPlan`; ``splits`` sets ``splits`` units per tile, see - :func:`router_units`).""" - if _resolve_deterministic(deterministic, "plan_router_fp32_gemm") or ( - _route(_k2_layout(A, B)) == "det" - ): - return _det_entry("plan_router_fp32_gemm")(A, B, out, splits=splits) + """Classify the fp32 views and plan the K2 launch, ``deterministic=False`` only: this package's stream-K program + for the ``kk`` / ``nn_t`` classes and the SIMT programs in their windows (:class:`AtomicRouterPlan`; ``splits`` + sets ``splits`` units per tile, see :func:`router_units`); a row the planner keeps on the deterministic program, + and a layout class routed ``"det"``, returns the deterministic package's plan (``plan_router_fp32_det_gemm``).""" + _require_false(deterministic, "plan_router_fp32_gemm", "plan_router_fp32_det_gemm") + if _route(_k2_layout(A, B)) == "det": + return deterministic_backend().plan_router_fp32_det_gemm( + A, B, out, splits=splits + ) return _atomic_k2("plan", A, B, out, splits=splits) @@ -5174,16 +5227,18 @@ def prepare_router_fp32_gemm( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> Any: - """Validate one fp32 binding, plan it and prepare its allocation-free launch (the deterministic package's - ``prepare_router_fp32_det_gemm``); ``deterministic=False`` prepares this package's stream-K program for the - ``kk`` / ``nn_t`` classes (its launch splits the second operand, zeroes ``out`` on the stream when a segment adds - into it, and runs the program: no partial buffers), and a SIMT program in its windows (the ``nn_t`` and ``kn`` - classes: one kernel over the fp32 operands). ``enable_pdl`` (default: the device supports it) launches the split - and the stream-K program as programmatic dependent launches.""" - if _resolve_deterministic(deterministic, "prepare_router_fp32_gemm") or ( - _route(_k2_layout(A, B)) == "det" - ): - return _det_entry("prepare_router_fp32_gemm")( + """Validate one fp32 binding, plan it and prepare its allocation-free launch, ``deterministic=False`` only: this + package's stream-K program for the ``kk`` / ``nn_t`` classes (its launch splits the second operand, zeroes + ``out`` on the stream when a segment adds into it, and runs the program: no partial buffers), and a SIMT program + in its windows (the ``nn_t`` and ``kn`` classes: one kernel over the fp32 operands); a row the planner keeps on the + deterministic program prepares the deterministic package's ``prepare_router_fp32_det_gemm``. ``enable_pdl`` + (default: the device supports it) launches the split and the stream-K program as programmatic dependent + launches.""" + _require_false( + deterministic, "prepare_router_fp32_gemm", "prepare_router_fp32_det_gemm" + ) + if _route(_k2_layout(A, B)) == "det": + return deterministic_backend().prepare_router_fp32_det_gemm( A, B, out, splits=splits, **_pdl_kw(enable_pdl) ) return _atomic_k2("prepare", A, B, out, splits=splits, enable_pdl=enable_pdl) @@ -5194,18 +5249,12 @@ def _router_gemm( B: torch.Tensor, out: torch.Tensor, splits: Optional[int], - det: bool, enable_pdl: Optional[bool] = None, views: Optional[tuple] = None, ) -> torch.Tensor: - """K2 on the convenience path: ``det`` calls the deterministic package's; with ``deterministic=False`` a layout - class routed ``"det"`` runs the deterministic program through this package's binding (:func:`_splitk_k2`), the - others this package's K2 path (``views``: an entry point's own launch views, :func:`_launch_views`). - ``enable_pdl`` reaches every route.""" - if det: - return _det_entry("router_fp32_gemm")( - A, B, out, splits=splits, **_pdl_kw(enable_pdl) - ) + """K2 on the convenience path: a layout class routed ``"det"`` runs the deterministic program through this + package's binding (:func:`_splitk_k2`), the others this package's K2 path (``views``: an entry point's own launch + views, :func:`_launch_views`). ``enable_pdl`` reaches every route.""" if _route(_k2_layout(A, B)) == "det": return _splitk_k2(A, B, out, splits, views, enable_pdl) return _atomic_k2( @@ -5223,14 +5272,14 @@ def router_fp32_gemm( enable_pdl: Optional[bool] = None, ) -> torch.Tensor: """``out[m, n] = sum_k A[m, k] B[k, n]`` for fp32 views through the BF16x3 emulation (the cached convenience - path). ``deterministic=False`` runs the ``kk`` / ``nn_t`` rows the planner takes (:func:`router_units`) on this + path), ``deterministic=False`` only: the ``kk`` / ``nn_t`` rows the planner takes (:func:`router_units`) on this package's stream-K program, whose part tiles add into ``out`` in the order they finish, instead of the fixed-order split-K sum, the ``nn_t`` / ``kn`` token counts of a SIMT window on its exact-FP32 program (:func:`router_simt`) and the ``kk`` token counts of a forward SIMT window on the forward's cluster program (:func:`router_forward_simt`); every other row runs the deterministic program. ``enable_pdl`` as for :func:`prepare_router_fp32_gemm`.""" - det = _resolve_deterministic(deterministic, "router_fp32_gemm") - return _router_gemm(A, B, out, splits, det, enable_pdl) + _require_false(deterministic, "router_fp32_gemm", "router_fp32_det_gemm") + return _router_gemm(A, B, out, splits, enable_pdl) def router_forward( @@ -5242,17 +5291,13 @@ def router_forward( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, N] = X[T, K] @ W[N, K].T`` in FP32 through the emulation; ``splits`` defaults to the deterministic - package's production count (4). With ``deterministic=False`` a row runs this package's program only with an - explicit ``splits`` (at most two partials per element: the logits repeat bit for bit) or a verdict of - :data:`ROUTER_RULES` (the token windows of the forward's SIMT program, :func:`router_forward_simt`), else the - deterministic program (its plan for ``router_fp32_gemm(X, W.T, out)``, through this package's binding, - :func:`_launch_cached_splitk_k2`): without a split count, while no router table gives the forward a verdict, - without the per-call check of the row.""" - if _resolve_deterministic(deterministic, "router_forward"): - return _det_entry("router_forward")( - X, W, out, **_splits_kw(splits), **_pdl_kw(enable_pdl) - ) + """``out[T, N] = X[T, K] @ W[N, K].T`` in FP32 through the emulation, ``deterministic=False`` only. A row runs + this package's program only with an explicit ``splits`` (at most two partials per element: the logits repeat bit + for bit) or a verdict of :data:`ROUTER_RULES` (the token windows of the forward's SIMT program, + :func:`router_forward_simt`), else the deterministic program (its plan for ``router_fp32_gemm(X, W.T, out)``, + through this package's binding, :func:`_launch_cached_splitk_k2`): without a split count, while no router table + gives the forward a verdict, without the per-call check of the row.""" + _require_false(deterministic, "router_forward", "router_det_forward") if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.float32) if _k2_entry("fwd", X, W, out, splits, enable_pdl) is not None: @@ -5262,7 +5307,7 @@ def router_forward( # no router table gives the forward a verdict: the deterministic program (the plan the planner would keep), # without the per-call check of the row return _splitk_k2(X, W.t(), out, None, views, enable_pdl) - return _router_gemm(X, W.t(), out, splits, False, enable_pdl, views) + return _router_gemm(X, W.t(), out, splits, enable_pdl, views) def router_dgrad( @@ -5273,22 +5318,18 @@ def router_dgrad( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, K] = G[T, N] @ W[N, K]`` in FP32: the deterministic program (one split: nothing to reduce) through the - emulation; with ``deterministic=False`` in the token windows of the row's input-gradient verdict in - :data:`ROUTER_RULES` (or, on a part with a table of its own, :data:`ROUTER_RULES_BY_SM`) the exact-FP32 SIMT - program (:func:`router_simt`: compensated sums of the products of ``G`` and ``W`` read in place, every element of - ``out`` stored once), and at every other row and token count the deterministic program through this package's - binding of it (:func:`_launch_cached_splitk_k2`: the same plan, program and bits). ``enable_pdl`` as for - :func:`prepare_router_fp32_gemm`.""" - if _resolve_deterministic(deterministic, "router_dgrad"): - return _det_entry("router_dgrad")(G, W, out, **_pdl_kw(enable_pdl)) + """``out[T, K] = G[T, N] @ W[N, K]`` in FP32, ``deterministic=False`` only: in the token windows of the row's + input-gradient verdict in :data:`ROUTER_RULES` (or, on a part with a table of its own, :data:`ROUTER_RULES_BY_SM`) + the exact-FP32 SIMT program (:func:`router_simt`: compensated sums of the products of ``G`` and ``W`` read in + place, every element of ``out`` stored once), and at every other row and token count the deterministic program + (one split: nothing to reduce) through this package's binding of it (:func:`_launch_cached_splitk_k2`: the same + plan, program and bits). ``enable_pdl`` as for :func:`prepare_router_fp32_gemm`.""" + _require_false(deterministic, "router_dgrad", "router_det_dgrad") if out is None: out = torch.empty(G.shape[0], W.shape[1], device=G.device, dtype=torch.float32) if _k2_entry("dgrad", G, W, out, None, enable_pdl) is not None: return out - return _router_gemm( - G, W, out, None, False, enable_pdl, (G.unsqueeze(0), W, _DGRAD_VIEWS) - ) + return _router_gemm(G, W, out, None, enable_pdl, (G.unsqueeze(0), W, _DGRAD_VIEWS)) def router_wgrad( @@ -5300,44 +5341,16 @@ def router_wgrad( deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[N, K] = G[T, N].T @ X[T, K]`` in FP32: the swapped GEMM ``X.T @ G`` with a transposed store; - ``splits`` defaults to the deterministic package's production count (6), or with ``deterministic=False`` to the - partition of the row's verdict in :data:`ROUTER_RULES` or, on a part with a table of its own, in - :data:`ROUTER_RULES_BY_SM` (:func:`router_partition`; a row or token count without one runs the deterministic - program, through this package's binding of it without planning the row here, :func:`_launch_cached_splitk_k2`).""" - if _resolve_deterministic(deterministic, "router_wgrad"): - return _det_entry("router_wgrad")( - G, X, out, **_splits_kw(splits), **_pdl_kw(enable_pdl) - ) + """``out[N, K] = G[T, N].T @ X[T, K]`` in FP32, ``deterministic=False`` only: the swapped GEMM ``X.T @ G`` with a + transposed store; ``splits`` defaults to the partition of the row's verdict in :data:`ROUTER_RULES` or, on a part + with a table of its own, in :data:`ROUTER_RULES_BY_SM` (:func:`router_partition`; a row or token count without one + runs the deterministic program, through this package's binding of it without planning the row here, + :func:`_launch_cached_splitk_k2`).""" + _require_false(deterministic, "router_wgrad", "router_det_wgrad") if out is None: out = torch.empty(G.shape[1], X.shape[1], device=G.device, dtype=torch.float32) if _k2_entry("wgrad", G, X, out, splits, enable_pdl) is not None: return out # a row without a verdict (kept per device and geometry, :func:`_router_memo`) runs the deterministic program views = (X.unsqueeze(0), G, _WGRAD_VIEWS) - return _router_gemm(G.t(), X, out, splits, False, enable_pdl, views) - - -def split_fp32_to_bf16x3_( - src: torch.Tensor, - parts: torch.Tensor, - resid: torch.Tensor, - *, - deterministic: bool = True, -) -> torch.Tensor: - """Exact three-way BF16 split ``src = p1 + p2 + p3`` into the retained ``parts [3, *src.shape]`` stack with the - fp32 scratch ``resid``. The same deterministic stage in both modes.""" - _resolve_deterministic(deterministic, "split_fp32_to_bf16x3_") - return _det_entry("split_fp32_to_bf16x3_")(src, parts, resid) - - -def __getattr__(name: str) -> Any: - """Read the other module-level names of this path through to the deterministic package (``ROW_RULES``, - ``stream_k_plan``, ``wgrad_views``, the plan and prepared types, ...). Names spelling the ``det`` token are not - read through: they live in the deterministic package only.""" - if (name.startswith("__") and name.endswith("__")) or _DET_SPELLED.search(name): - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - try: - return getattr(deterministic_backend(), DETERMINISTIC_NAMES.get(name, name)) - except AttributeError: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + return _router_gemm(G.t(), X, out, splits, enable_pdl, views) diff --git a/flashinfer/experimental/dense_projection_gemm/cake_jit.py b/flashinfer/experimental/dense_projection_gemm/cake_jit.py index 869866ba266..7464506f800 100644 --- a/flashinfer/experimental/dense_projection_gemm/cake_jit.py +++ b/flashinfer/experimental/dense_projection_gemm/cake_jit.py @@ -43,9 +43,8 @@ # ``router_fp32_fwd_simt_*``) and the sum of the # deterministic router program's split-K partials in a fixed order (template ``router_partials_sum``), which this # package's binding of that program launches after it. Every program the deterministic package registers -# (``cake_dense_projection_det_gemm_``, templates ``dense_proj_det_gemm_*``, ``router_fp32_det_gemm_*``, -# ``split_fp32_det_bf16x3``) is served by that package: :func:`select_module` and -# :func:`load_cake_dense_projection_gemm_module` forward its templates and names to it. +# (``cake_dense_projection_det_gemm_``) is served by that package's registry and loader only: :func:`select_module` +# and :func:`load_cake_dense_projection_gemm_module` know this package's programs, nothing else. # fmt: off ARG_PLANS: dict[str, list[list[str]]] = { @@ -98,12 +97,12 @@ # The architectures this package registers programs for (``KERNELS`` keys and every record's ``arches`` are drawn # from it). ARCHES = ("sm_100a", "sm_103a", "sm_107a") -TRACKING_ISSUE = "flashinfer-ai/flashinfer#5677" # the dense projection GEMM family (both numerical modes) +# The tracking issue of the dense projection GEMM family (both packages). +TRACKING_ISSUE = "flashinfer-ai/flashinfer#5677" -#: The program-name prefix of this package and of the deterministic package. +#: The program-name prefix of this package. PROGRAM_PREFIX = "cake_dense_projection_gemm_" -DET_PROGRAM_PREFIX = "cake_dense_projection_det_gemm_" -#: The instance templates of this package's programs, and of the deterministic package's. The token ``atomic`` +#: The instance templates of this package's programs. The token ``atomic`` #: keeps every template of this package different from every template the deterministic package registers and #: from every template it registered before its rename (``dense_proj_gemm_...``, ``router_fp32_gemm_``); #: ``router_fp32_gemm_ts_`` names the router weight gradient's head/tail programs with the ``A`` parts in tensor memory, @@ -120,11 +119,6 @@ "router_partials_sum", "router_fp32_fwd_simt_", ) -DET_TEMPLATE_PREFIXES = ( - "dense_proj_det_gemm_", - "router_fp32_det_gemm_", - "split_fp32_det_bf16x3", -) #: The deterministic package: a sibling of this package (``.dense_projection_det_gemm``), resolved on first #: use and never through an absolute ``flashinfer`` import; a vendored copy binds it explicitly @@ -147,7 +141,7 @@ def deterministic_module(name: str) -> ModuleType: parent = __name__.rpartition(".")[0].rpartition(".")[0] if not parent: raise RuntimeError( - f"{__name__.rpartition('.')[0]} delegates its deterministic path to the package " + f"{__name__.rpartition('.')[0]} reuses the programs of the deterministic package " f"{DETERMINISTIC_PACKAGE!r}, which a vendored copy has no sibling of; vendor that package too and " "bind it with cake_backend.use_deterministic_backend()" ) @@ -155,16 +149,6 @@ def deterministic_module(name: str) -> ModuleType: return importlib.import_module(f"{package}.{name}") -def is_det_template(template: str) -> bool: - """Whether ``template`` names an instance of the deterministic package.""" - return template.startswith(DET_TEMPLATE_PREFIXES) - - -def is_det_program(name: str) -> bool: - """Whether ``name`` names a program of the deterministic package.""" - return name.startswith(DET_PROGRAM_PREFIX) - - def arg_plan_key(plan) -> str: """Registry key of an argument plan: the first 12 hex digits of the SHA-256 of its compact JSON (``[[kind, name], ...]`` with ``,`` / ``:`` separators) -- the export computes the same key.""" @@ -196,10 +180,8 @@ def registered_templates(arch: str) -> tuple[str, ...]: def select_module(arch: str, template: str) -> str: - """Return the registered module name serving ``template`` on ``arch``: a template of the deterministic package - resolves through that package's registry, any other through this package's.""" - if is_det_template(template): - return deterministic_module("cake_jit").select_module(arch, template) + """Return the registered module name of this package serving ``template`` on ``arch`` (a template of the + deterministic package is not registered here: that package's ``cake_jit.select_module`` resolves it).""" name = KERNELS.get(arch, {}).get(template) if name is None: registered = registered_templates(arch) @@ -237,13 +219,9 @@ def _load_module(name: str, arch: str): def load_cake_dense_projection_gemm_module(name: str, arch: str): - """Load (once per program and architecture) the generated program ``name`` for ``arch``: this package's - prebuilt library under ``aot/`` when it ships one, else FlashInfer's JIT build of its ``build.json`` entry; a - program of the deterministic package is loaded by that package's loader.""" - if is_det_program(name): - return deterministic_module( - "cake_jit" - ).load_cake_dense_projection_det_gemm_module(name, arch) + """Load (once per program and architecture) this package's generated program ``name`` for ``arch``: its + prebuilt library under ``aot/`` when it ships one, else FlashInfer's JIT build of its ``build.json`` entry (a + program of the deterministic package is loaded by that package's ``load_cake_dense_projection_det_gemm_module``).""" return _load_module(name, arch) From e0d562f55d4502a1a89f8559a0db2be06d3367e6 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 5 Nov 2026 21:42:44 -0800 Subject: [PATCH 03/13] test(dense_projection_gemm): explicit deterministic=False, alerts before order-varying launches - Absence tests for every removed name (former-name map, module __getattr__, mode resolver, deterministic wrappers, template/program forwarding). - ValueError matrix over every compute entry point: omitted, True, None, 0, 1, "False", 0.0, "" and NumPy bools, with PyTorch's deterministic mode off, strict and warn_only; nothing is planned or launched and nothing warns. - torch.use_deterministic_algorithms(True) alerts (CPU, monkeypatched launches): RuntimeError before an order-varying K1 (three splits) or K2 (more than two partials per element) launch, UserWarning then the launch under warn_only, no alert for fixed-order launches; prepared and convenience paths. - GPU: two splits run silently and repeat bit for bit under strict mode; three splits raise with out untouched, or warn and pass the accuracy rule. - Drop the default-path and deterministic-mode rerouting tests. --- .../test_cake_dense_projection_gemm.py | 666 ++++++++++++------ 1 file changed, 451 insertions(+), 215 deletions(-) diff --git a/tests/experimental/test_cake_dense_projection_gemm.py b/tests/experimental/test_cake_dense_projection_gemm.py index 31ab28ba5b6..5b1958cbc35 100644 --- a/tests/experimental/test_cake_dense_projection_gemm.py +++ b/tests/experimental/test_cake_dense_projection_gemm.py @@ -14,9 +14,10 @@ limitations under the License. """ -# The dense projection GEMM package with an explicit numerical mode: the dispatch of every entry point -# (``deterministic=True`` -> the deterministic package's canonical object, ``deterministic=False`` -> the routes of -# ``cake_backend.ROUTES``), PyTorch's deterministic-algorithms rule, the names of the former import path, the K1 and +# The non-deterministic dense projection GEMM package: the mode rule of every compute entry point (only +# ``deterministic=False``; any other value raises ``ValueError`` naming the deterministic entry point), the routes of +# ``cake_backend.ROUTES``, PyTorch's deterministic-algorithms alert before an order-varying launch, the absence of +# every deterministic name (no alias, map or read-through), the K1 and # K2 planners (the per-architecture planners, the rule tables, the token-count guard, the stream-K partition), the # registry and file names of this package's programs and its workspace kinds. The CPU tests launch nothing. The GPU # tests (``test_gpu_*``) need a device the deterministic package registers programs for; the ones of this package's @@ -62,10 +63,10 @@ RERUNS = 20 SEED = 20261031 -# The names of ``cake_backend`` before #8046 that are entry points: each takes ``deterministic`` here. +# The compute entry points of ``cake_backend``: each takes the keyword-only ``deterministic: bool = True`` and accepts +# only ``deterministic=False``. ENTRY_POINTS = ( "dense_projection_gemm", - "instance_symbol", "plan_dense_projection_gemm", "plan_router_fp32_gemm", "prepare_dense_projection_gemm", @@ -77,11 +78,9 @@ "router_dgrad", "router_forward", "router_fp32_gemm", - "router_symbol", "router_wgrad", - "split_fp32_to_bf16x3_", ) -# former name -> the deterministic package's canonical name +# name of this package (or a former name) -> the deterministic package's canonical name FORMER_TO_DET = { "GemmPlan": "DetGemmPlan", "PreparedGemm": "PreparedDetGemm", @@ -126,6 +125,40 @@ "wgrad_swapped", "wgrad_views", ) # fmt: on +# The names among FORMER_MODULE_NAMES that this package defines itself (its entry points, helpers and constants); every +# other one is absent here (no read-through to the deterministic package). +OWN_FORMER_NAMES = frozenset( + ( + *ENTRY_POINTS, + "router_symbol", + "bind_launch", + "clear_convenience_caches", + "ROUTER_BLOCK_K", + "ROUTER_BLOCK_M", + "ROUTER_BLOCK_N", + ) +) +# The names this package no longer has: the former-name map, its read-through and type tables, the mode resolver, the +# wrappers that only returned or ran a deterministic object, and the forwarding of deterministic templates and programs. +REMOVED_BACKEND_NAMES = ( + "DETERMINISTIC_NAMES", + "_DET_TYPES", + "_DET_SPELLED", + "_det_entry", + "_resolve_deterministic", + "instance_symbol", + "split_fp32_to_bf16x3_", + "__getattr__", +) +REMOVED_JIT_NAMES = ( + "is_det_template", + "is_det_program", + "DET_TEMPLATE_PREFIXES", + "DET_PROGRAM_PREFIX", +) +# The values of ``deterministic`` other than ``False``: each raises ``ValueError`` (``_OMITTED``: no keyword). +_OMITTED = object() +NOT_FALSE = (_OMITTED, True, None, 0, 1, "False", 0.0, "") DET_TOKEN = re.compile(r"(?:^|[^a-z])det(?:[^a-z]|$)", re.IGNORECASE) # the mode is named by deterministic= NEGATED_TOKEN = re.compile(r"non_?" + "det", re.IGNORECASE) @@ -218,21 +251,35 @@ def test_entry_points_take_keyword_only_deterministic_true(name): @pytest.mark.parametrize("name", FORMER_MODULE_NAMES) -def test_every_former_module_name_resolves(name): - obj = getattr(cake_backend, name) - if name in ENTRY_POINTS or name in ("bind_launch", "clear_convenience_caches"): - assert obj.__module__ == cake_backend.__name__ +def test_no_former_module_name_reads_through(name): + """A name ``cake_backend`` had before #8046 is either this package's own object (defined in this module) or + absent: no alias and no read-through to the deterministic package.""" + if name in OWN_FORMER_NAMES: + assert name in vars(cake_backend) + obj = vars(cake_backend)[name] + if callable(obj): + assert obj.__module__ == cake_backend.__name__ else: - assert obj is getattr( - det_backend, cake_backend.DETERMINISTIC_NAMES.get(name, name) - ) + assert name not in vars(cake_backend) and name not in dir(cake_backend) + with pytest.raises(AttributeError): + getattr(cake_backend, name) + + +def test_the_former_types_are_absent(): + for name in ("GemmPlan", "PreparedGemm", "RouterPlan", "PreparedRouterGemm"): + assert not hasattr(cake_backend, name), name -def test_former_types_are_the_deterministic_classes(): - assert cake_backend.GemmPlan is det_backend.DetGemmPlan - assert cake_backend.PreparedGemm is det_backend.PreparedDetGemm - assert cake_backend.RouterPlan is det_backend.RouterDetPlan - assert cake_backend.PreparedRouterGemm is det_backend.PreparedRouterDetGemm +def test_the_removed_names_are_absent(): + """No former-name map, read-through (module ``__getattr__``), mode resolver, deterministic wrapper or forwarding + of deterministic templates and programs is left in this package.""" + for name in REMOVED_BACKEND_NAMES: + assert name not in vars(cake_backend), name + for name in REMOVED_BACKEND_NAMES[:-1]: + with pytest.raises(AttributeError): + getattr(cake_backend, name) + for name in REMOVED_JIT_NAMES: + assert not hasattr(cake_jit, name), name @pytest.mark.parametrize( @@ -241,15 +288,33 @@ def test_former_types_are_the_deterministic_classes(): "projection_det_forward", "DetGemmPlan", "router_fp32_det_gemm", + "ROW_RULES", + "stream_k_plan", + "wgrad_views", "__wrapped__", "no_such_name", ], ) -def test_deterministic_spellings_are_not_read_through(name): +def test_deterministic_names_are_not_reachable(name): with pytest.raises(AttributeError): getattr(cake_backend, name) +def test_no_deterministic_object_is_reachable_through_this_package(): + """No function or class the deterministic package defines is an attribute of this package's modules (the + deterministic package is reached only through ``deterministic_backend()``).""" + det_objects = { + id(obj) + for module in (det_backend, det_jit) + for obj in vars(module).values() + if (inspect.isfunction(obj) or inspect.isclass(obj)) + and obj.__module__ == module.__name__ + } + for module in (cake_backend, cake_jit, package): + leaked = sorted(n for n, obj in vars(module).items() if id(obj) in det_objects) + assert not leaked, (module.__name__, leaked) + + def test_the_deterministic_package_resolves_as_a_sibling(monkeypatch): """In a FlashInfer tree and in a copy vendored with the deterministic package beside it (the vendored gate's layout) the deterministic package resolves as this package's sibling.""" @@ -266,16 +331,25 @@ def test_use_deterministic_backend_binds_a_vendored_copy(monkeypatch): ModuleType("vendored_det_pkg"), ModuleType("vendored_det_pkg.cake_backend"), ) - for name in cake_backend.DETERMINISTIC_NAMES.values(): + for name in cake_backend._DET_ENTRIES: setattr(fake, name, Recorder()) monkeypatch.setitem(sys.modules, "vendored_det_pkg", fake_pkg) monkeypatch.setitem(sys.modules, "vendored_det_pkg.cake_backend", fake) cake_backend.use_deterministic_backend(fake) assert cake_backend.deterministic_backend() is fake + # a forward layout (routed "det") plans with the bound package's deterministic planner + A, B, out = _k1_views("kk") assert ( - cake_backend.projection_forward("X", "W") is fake.projection_det_forward.result + cake_backend.plan_dense_projection_gemm(A, B, out, deterministic=False) + is fake.plan_dense_projection_det_gemm.result ) - assert fake.projection_det_forward.calls == [(("X", "W", None), {})] + assert fake.plan_dense_projection_det_gemm.calls == [((A, B, out), {})] + # the mode rule names the bound package + with pytest.raises( + ValueError, match=r"vendored_det_pkg\.cake_backend\.projection_det_forward" + ): + cake_backend.projection_forward("X", "W") + assert not fake.projection_det_forward.calls with pytest.raises(TypeError, match="lacks"): cake_backend.use_deterministic_backend(ModuleType("not_a_backend.cake_backend")) @@ -284,99 +358,130 @@ def test_use_deterministic_backend_binds_a_vendored_copy(monkeypatch): # Dispatch (CPU) # --------------------------------------------------------------------------- -# entry point -> (positional arguments, keyword arguments, the call the deterministic object receives) +# entry point -> its positional arguments (sentinels: the mode rule raises before anything reads them) and the +# keyword arguments of the call # fmt: off -DEFAULT_CALLS = { - "dense_projection_gemm": ( - ("A", "B", "out"), - {"transposed_out": True, "enable_pdl": True}, - (("A", "B", "out"), {"transposed_out": True, "enable_pdl": True}), - ), - "plan_dense_projection_gemm": ( - ("A", "B", "out"), - {"sm_count": 148, "l2_bytes": 1 << 20, "transposed_out": True, "block_n": 128}, - (("A", "B", "out"), {"sm_count": 148, "l2_bytes": 1 << 20, "transposed_out": True, "block_n": 128}), - ), - "prepare_dense_projection_gemm": (("A", "B", "out"), {"sk": False}, (("A", "B", "out"), {"sk": False})), - "projection_forward": (("X", "W"), {}, (("X", "W", None), {})), - "projection_dgrad": (("G", "W", "out"), {"out_dtype": "f32"}, (("G", "W", "out"), {"out_dtype": "f32"})), - "projection_wgrad": ( - ("G", "X"), - {"out_dtype": "f32", "sk": True, "cta_rows": 64}, - (("G", "X", None), {"out_dtype": "f32", "sk": True, "cta_rows": 64}), - ), - "prepare_projection_wgrad": ( - ("G", "X", "out"), - {"cta_rows": 128}, - (("G", "X", "out"), {"sk": "auto", "cta_rows": 128}), - ), - "plan_router_fp32_gemm": (("A", "B", "out"), {"splits": 3}, (("A", "B", "out"), {"splits": 3})), - "prepare_router_fp32_gemm": (("A", "B", "out"), {}, (("A", "B", "out"), {"splits": None})), - "router_fp32_gemm": (("A", "B", "out"), {"splits": 2}, (("A", "B", "out"), {"splits": 2})), - "router_forward": (("X", "W"), {}, (("X", "W", None), {})), - "router_dgrad": (("G", "W", "out"), {}, (("G", "W", "out"), {})), - "router_wgrad": (("G", "X"), {"splits": 9}, (("G", "X", None), {"splits": 9})), - "split_fp32_to_bf16x3_": (("src", "parts", "resid"), {}, (("src", "parts", "resid"), {})), +MODE_CALLS = { + "dense_projection_gemm": (("A", "B", "out"), {"transposed_out": True, "enable_pdl": True}), + "plan_dense_projection_gemm": (("A", "B", "out"), {"sm_count": 148, "transposed_out": True, "block_n": 128}), + "prepare_dense_projection_gemm": (("A", "B", "out"), {"sk": False}), + "projection_forward": (("X", "W"), {}), + "projection_dgrad": (("G", "W", "out"), {"out_dtype": "f32"}), + "projection_wgrad": (("G", "X"), {"out_dtype": "f32", "sk": True, "cta_rows": 64}), + "prepare_projection_wgrad": (("G", "X", "out"), {"cta_rows": 128}), + "plan_router_fp32_gemm": (("A", "B", "out"), {"splits": 3}), + "prepare_router_fp32_gemm": (("A", "B", "out"), {}), + "router_fp32_gemm": (("A", "B", "out"), {"splits": 2}), + "router_forward": (("X", "W"), {}), + "router_dgrad": (("G", "W", "out"), {}), + "router_wgrad": (("G", "X"), {"splits": 9}), } # fmt: on -@pytest.mark.parametrize("name", sorted(DEFAULT_CALLS)) -def test_the_default_calls_the_deterministic_object( - monkeypatch, torch_determinism, name -): - """The default path is the deterministic package's object with the same arguments (so the same bits), with and - without PyTorch's deterministic mode, and never warns.""" - args, kwargs, expected = DEFAULT_CALLS[name] - recorder = _patch_det(monkeypatch, name) - for enabled in (False, True): - torch.use_deterministic_algorithms(enabled) - with _no_warnings(): - assert getattr(cake_backend, name)(*args, **kwargs) is recorder.result - assert ( - getattr(cake_backend, name)(*args, deterministic=True, **kwargs) - is recorder.result - ) - assert recorder.calls == [expected] * 4 - - -def test_the_default_router_forward_and_wgrad_keep_the_production_splits(monkeypatch): - forward, wgrad = ( - _patch_det(monkeypatch, "router_forward"), - _patch_det(monkeypatch, "router_wgrad"), - ) - cake_backend.router_forward("X", "W", "out") - cake_backend.router_wgrad("G", "X", "out", splits=6) - assert forward.calls == [ - (("X", "W", "out"), {}) - ] # the deterministic default (4 splits) applies - assert wgrad.calls == [(("G", "X", "out"), {"splits": 6})] - - -def test_instance_and_router_symbols(): - key = det_backend.instance_key(a_mn=True, b_mn=True) - assert cake_backend.instance_symbol( - key, deterministic=False - ) == det_backend.instance_symbol(key) - for a_mn, b_mn, out_t in ( - (False, False, False), - (False, True, False), - (True, True, True), - ): - assert cake_backend.router_symbol( - a_mn, b_mn, out_t - ) == det_backend.router_symbol(a_mn, b_mn, out_t) - assert ( - cake_backend.router_symbol(False, False, False, deterministic=False) - == "router_fp32_gemm_atomic_kk" +def _mode_message(name: str) -> str: + """The whole ``ValueError`` text of ``name``'s mode rule, as a pattern (any value after ``got``).""" + package = cake_backend.__name__.rpartition(".")[0] or cake_backend.__name__ + det_entry = f"{cake_backend._det_package_name()}.cake_backend.{FORMER_TO_DET[name]}" + head = f"{name}: {package} runs the non-deterministic package only and needs deterministic=False (got " + tail = ( + "): pass deterministic=False for the non-deterministic path; " + f"the deterministic entry is {det_entry}" ) - assert ( - cake_backend.router_symbol(True, True, True, deterministic=False) - == "router_fp32_gemm_atomic_nn_t" + return "^" + re.escape(head) + ".*" + re.escape(tail) + "$" + + +def _forbid_det_package(monkeypatch) -> list: + """Every entry point of the deterministic package and every path of this package records a call, and + ``deterministic_backend`` fails: a call that raises the mode rule's ``ValueError`` reaches none of them.""" + reached = [] + + def forbidden(*args, **kwargs): + reached.append((args, kwargs)) + raise AssertionError("the mode rule let the call through") + + monkeypatch.setattr(cake_backend, "deterministic_backend", forbidden) + for name in cake_backend._DET_ENTRIES: + monkeypatch.setattr(det_backend, name, forbidden) + paths = ("_atomic_k1", "_atomic_k2", "_det_k1", "_router_gemm", "_splitk_k2") + for name in (*paths, "_k1_cached_launch", "_k2_entry"): + monkeypatch.setattr(cake_backend, name, forbidden) + return reached + + +def test_mode_calls_cover_every_entry_point(): + assert sorted(MODE_CALLS) == sorted(ENTRY_POINTS) + assert set(ENTRY_POINTS) <= set(FORMER_TO_DET) + for name in ENTRY_POINTS: + assert hasattr(det_backend, FORMER_TO_DET[name]), name + + +@pytest.mark.parametrize( + "value", NOT_FALSE, ids=lambda v: "omitted" if v is _OMITTED else repr(v) +) +@pytest.mark.parametrize("name", ENTRY_POINTS) +def test_every_value_but_false_raises_naming_the_deterministic_entry( + monkeypatch, torch_determinism, name, value +): + """An entry point takes only ``deterministic=False``: the default, ``True``, ``None``, ``0``, ``1``, ``"False"`` + and every other value raise ``ValueError`` naming the deterministic package's entry point, before anything is + planned or launched and without a call of the deterministic package; PyTorch's deterministic mode (strict or + ``warn_only``) changes nothing, and nothing warns.""" + args, kwargs = MODE_CALLS[name] + reached = _forbid_det_package(monkeypatch) + if value is not _OMITTED: + kwargs = dict(kwargs, deterministic=value) + for enabled, warn_only in ((False, False), (True, False), (True, True)): + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + with _no_warnings(), pytest.raises(ValueError, match=_mode_message(name)): + getattr(cake_backend, name)(*args, **kwargs) + assert not reached + + +def test_numpy_bools_raise(): + """A NumPy bool is not the bool ``False``: ``np.False_`` and ``np.True_`` raise like any other value.""" + np = pytest.importorskip("numpy") + for value in (np.False_, np.True_): + with pytest.raises(ValueError, match=_mode_message("projection_wgrad")): + cake_backend.projection_wgrad("G", "X", "out", deterministic=value) + + +def test_the_mode_message_names_the_deterministic_package_without_importing_it( + monkeypatch, +): + """The mode rule names the deterministic entry point from text (the bound package, else this package's + sibling): it imports nothing.""" + imported = [] + monkeypatch.setattr( + cake_jit, "deterministic_module", lambda name: imported.append(name) ) - assert cake_backend.router_symbol( - False, True, False, deterministic=False - ) == det_backend.router_symbol(False, True, False) + monkeypatch.setattr(cake_jit, "_det_package", None) + with pytest.raises(ValueError) as raised: + cake_backend.router_wgrad("G", "X", deterministic=None) + parent = cake_backend.__name__.rpartition(".")[0].rpartition(".")[0] + det_package = ( + f"{parent}.dense_projection_det_gemm" if parent else "dense_projection_det_gemm" + ) + assert f"{det_package}.cake_backend.router_det_wgrad" in str(raised.value) + assert "needs deterministic=False (got None)" in str(raised.value) + assert not imported + + +def test_router_symbol_takes_no_mode_keyword(): + """``router_symbol`` names this package's stream-K template of the ``kk`` / ``nn_t`` layout classes (a helper: + no ``deterministic`` keyword); ``kn`` (SIMT programs only) and ``nk`` have none.""" + symbol = cake_backend.router_symbol + assert "deterministic" not in inspect.signature(symbol).parameters + assert symbol(False, False, False) == "router_fp32_gemm_atomic_kk" + assert symbol(True, True, True) == "router_fp32_gemm_atomic_nn_t" + stream_k = {"router_fp32_gemm_atomic_kk", "router_fp32_gemm_atomic_nn_t"} + for arch, table in cake_jit.KERNELS.items(): + assert stream_k <= set(table), arch + for a_mn, b_mn in ((False, True), (True, False)): + with pytest.raises(ValueError, match="no stream-K program"): + cake_backend.router_symbol(a_mn, b_mn, False) + with pytest.raises(TypeError): + cake_backend.router_symbol(False, False, False, deterministic=False) @pytest.mark.parametrize("layout", ["kk", "kn", "nn"]) @@ -464,7 +569,7 @@ def test_k1_deterministic_programs_run_through_the_binding(monkeypatch, op, pdl) """``deterministic=False`` K1 calls that run the deterministic program go through this package's binding (``_det_k1``): keyed by the entry point's own tensors, with the deterministic package's convenience call as the plan (``det_args``) and as the fallback (``det_call``), the entry point's ``enable_pdl`` (and the generic entry's - ``weights_static``) passed to both; ``deterministic=True`` calls the deterministic package.""" + ``weights_static``) passed to both; ``deterministic=True`` raises ``ValueError`` and calls nothing.""" names = {"fwd": "projection_forward", "dgrad": "projection_dgrad"} det = _patch_det(monkeypatch, names.get(op, "dense_projection_gemm")) binding = Recorder() @@ -510,7 +615,9 @@ def test_k1_deterministic_programs_run_through_the_binding(monkeypatch, op, pdl) want["weights_static"] = True assert det.calls == [((A, B, out), want)] det.calls.clear() - assert call(True) is det.result and len(det.calls) == 1 and len(binding.calls) == 1 + with pytest.raises(ValueError, match="needs deterministic=False"): + call(True) + assert not det.calls and len(binding.calls) == 1 def _same_entry_views(views, a_view, split_src, flags) -> None: @@ -525,13 +632,10 @@ def test_k2_routes(monkeypatch): gradient without a split count: its SIMT windows, else the deterministic program), and so does the forward with a split count or while some router table gives the forward a verdict; without either the forward runs the deterministic program through this package's binding (the plan the planner keeps for a row without a verdict). - The split stage is the deterministic package's. Each entry point passes its own tensors as the launch views. The - forward's shipped verdicts (its SIMT windows) are cleared first.""" + Each entry point passes its own tensors as the launch views. The forward's shipped verdicts (its SIMT windows) + are cleared first.""" _without_forward_verdicts(monkeypatch) - dgrad, split = ( - _patch_det(monkeypatch, "router_dgrad"), - _patch_det(monkeypatch, "split_fp32_to_bf16x3_"), - ) + dgrad = _patch_det(monkeypatch, "router_dgrad") gemm = _patch_det(monkeypatch, "router_fp32_gemm") atomic, splitk = Recorder(), Recorder() monkeypatch.setattr(cake_backend, "_atomic_k2", atomic) @@ -555,10 +659,6 @@ def test_k2_routes(monkeypatch): ) assert cake_backend.router_dgrad(G, W, deterministic=False) is atomic.result assert not dgrad.calls - assert ( - cake_backend.split_fp32_to_bf16x3_("s", "p", "r", deterministic=False) - is split.result - ) calls = atomic.calls ( (fwd2_args, fwd2_kwargs), @@ -5522,7 +5622,7 @@ def test_k1_forward_rows_bind_their_plan(monkeypatch, arch): """``projection_forward(deterministic=False)`` binds, for a row of :data:`K1_FORWARD_RULES_BY_SM` on its part (BF16 ``out``, both operands K-major), the deterministic program the deterministic planner gives the row with the table's knobs, and the deterministic plan for a token count, row or SM count outside the table; an FP32 ``out`` - or an operand that is not K-major gets no knobs; ``deterministic=True`` calls the deterministic package.""" + or an operand that is not K-major gets no knobs; ``deterministic=True`` raises ``ValueError`` and plans nothing.""" entry_fn = _cpu_k1_launches(monkeypatch, arch) monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) bf16 = torch.bfloat16 @@ -5567,10 +5667,13 @@ def det_plan(X, W, out, knobs): monkeypatch.setattr(det_backend, "device_facts", lambda device=None, f=facts: f) assert cake_backend._k1_forward_knobs_of(X, W, out) == (), sms assert bound_plan(X, W, out) == det_plan(X, W, out, ()), sms - # deterministic=True: the deterministic package's entry point, whatever the table says + # deterministic=True (the default): ValueError naming the deterministic entry, whatever the table says det = _patch_det(monkeypatch, "projection_forward") - assert cake_backend.projection_forward(X, W, out) is det.result - assert det.calls == [((X, W, out), {})] + entry_fn.calls.clear() + cake_backend._K1_DET_CACHE.clear() + with pytest.raises(ValueError, match="projection_det_forward"): + cake_backend.projection_forward(X, W, out) + assert not det.calls and not entry_fn.calls and not cake_backend._K1_DET_CACHE def test_k1_entry_points_go_from_their_key_to_the_bound_launch(monkeypatch): @@ -7089,7 +7192,7 @@ def test_deterministic_backend_follows_the_binding_and_sys_modules(monkeypatch): ModuleType("vendored_det_pkg2"), ModuleType("vendored_det_pkg2.cake_backend"), ) - for name in cake_backend.DETERMINISTIC_NAMES.values(): + for name in cake_backend._DET_ENTRIES: setattr(fake, name, Recorder()) monkeypatch.setitem(sys.modules, "vendored_det_pkg2", fake_pkg) monkeypatch.setitem(sys.modules, "vendored_det_pkg2.cake_backend", fake) @@ -7110,7 +7213,7 @@ def test_use_deterministic_backend_drops_the_convenience_entries(monkeypatch): ModuleType("vendored_det_pkg3"), ModuleType("vendored_det_pkg3.cake_backend"), ) - for name in cake_backend.DETERMINISTIC_NAMES.values(): + for name in cake_backend._DET_ENTRIES: setattr(fake, name, Recorder()) monkeypatch.setitem(sys.modules, "vendored_det_pkg3", fake_pkg) monkeypatch.setitem(sys.modules, "vendored_det_pkg3.cake_backend", fake) @@ -7171,47 +7274,188 @@ def test_k2_misaligned_out_takes_the_deterministic_path(monkeypatch, shift_elems # --------------------------------------------------------------------------- -# PyTorch's deterministic-algorithms mode (CPU) +# PyTorch's deterministic-algorithms mode (CPU): the alert before an order-varying launch # --------------------------------------------------------------------------- -# fmt: off -ATOMIC_MODE_CALLS = { - "dense_projection_gemm": lambda: cake_backend.dense_projection_gemm(*_k1_views("nn"), deterministic=False), - "projection_forward": lambda: cake_backend.projection_forward("X", "W", deterministic=False), - "projection_wgrad": lambda: cake_backend.projection_wgrad("G", "X", "out", deterministic=False), - "router_forward": lambda: cake_backend.router_forward("X", "W", deterministic=False), - "router_dgrad": lambda: cake_backend.router_dgrad("G", "W", deterministic=False), - "router_wgrad": lambda: cake_backend.router_wgrad("G", "X", deterministic=False), - "split_fp32_to_bf16x3_": lambda: cake_backend.split_fp32_to_bf16x3_("s", "p", "r", deterministic=False), -} -# fmt: on + +def _prepared_launch(monkeypatch, k: str, order_varying: bool): + """A prepared launch of this package on CPU (nothing is loaded or launched) and the recorders of what it issues: + K1 (``"k1"``) at three splits (the partials added in the order the splits finish) or two (one add), its program; + K2 (``"k2"``) at 120 units over 40 tiles (three partials per element) or 80 (two), its split, the fill of ``out`` + and its program.""" + monkeypatch.setattr( + det_backend, + "_ffi_stream_context", + lambda index, raw=None: contextlib.nullcontext(), + ) + monkeypatch.setattr(det_backend, "_current_raw_stream", lambda index: 0) + out, program = torch.zeros(8, 8), Recorder() + if k == "k1": + plan = cake_backend.AtomicGemmPlan( + L=1, + M=4096, + N=128, + K=1000, + out_f32=True, + transposed_out=True, + block_n=128, + splits=3 if order_varying else 2, + m_tiles=16, + n_tiles=1, + sm_pairs=76, + template="k1_program", + ) + launch = cake_backend.PreparedAtomicGemm( + "k1_module", "sm_103a", plan, out, {}, _launch=program + ) + return launch, (program,) + plan = cake_backend.AtomicRouterPlan( + M=5120, + N=256, + K=192, + a_mn=True, + b_mn=True, + swapped=False, + layout="nn_t", + m_tiles=40, + n_tiles=2, + k_iters=6, + units=120 if order_varying else 80, + template="k2_program", + ) + split, fill = Recorder(), Recorder() + launch = cake_backend.PreparedAtomicRouterGemm( + "k2_module", "sm_103a", plan, out, {}, _launch=program, _split=split, _fill=fill + ) + return launch, (split, fill, program) + + +#: The words of the alert that name the order-varying launch of :func:`_prepared_launch`. +PREPARED_ALERT_TEXT = {"k1": "at 3 splits", "k2": "with 120 units over 40 tiles"} -@pytest.mark.parametrize("name", sorted(ATOMIC_MODE_CALLS)) -def test_explicit_false_under_torch_deterministic_mode_raises( - monkeypatch, torch_determinism, name +@pytest.mark.parametrize("k", ["k1", "k2"]) +def test_an_order_varying_prepared_launch_alerts_under_torch_deterministic_mode( + monkeypatch, torch_determinism, k +): + """With ``torch.use_deterministic_algorithms(True)`` a prepared launch whose result varies run to run (K1 at + three splits, K2 with three partials per element) raises PyTorch's ``RuntimeError`` before it issues anything (no + split, no fill of ``out``, no program); with ``warn_only=True`` it warns, at the caller's line, and launches; with + the mode off it launches without a warning. The alert names the launch and the deterministic package.""" + launch, recorders = _prepared_launch(monkeypatch, k, True) + assert not launch.plan.reproducible + text = PREPARED_ALERT_TEXT[k] + torch.use_deterministic_algorithms(True) + with pytest.raises(RuntimeError, match=re.escape(text)) as info: + launch.launch() + assert cake_backend._det_package_name() in str(info.value) + assert "warn_only=True" in str(info.value) + assert not any(r.calls for r in recorders) + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.warns(UserWarning, match=re.escape(text)) as record: + assert launch.launch() is launch.out + assert all(len(r.calls) == 1 for r in recorders) + (alert,) = [w for w in record if text in str(w.message)] + assert Path(alert.filename).name == Path(__file__).name + torch.use_deterministic_algorithms(False) + with _no_warnings(): + assert launch.launch() is launch.out + assert all(len(r.calls) == 2 for r in recorders) + + +@pytest.mark.parametrize("k", ["k1", "k2"]) +def test_a_fixed_order_prepared_launch_does_not_alert( + monkeypatch, torch_determinism, k ): - recorder = _patch_det(monkeypatch, name) - monkeypatch.setattr(cake_backend, "_atomic_k1", Recorder()) - monkeypatch.setattr(cake_backend, "_atomic_k2", Recorder()) + """A prepared launch whose result repeats (K1 at two splits, one add of two partials; K2 with at most two partials + per element; a SIMT program, which stores every element once) launches without an alert under + ``torch.use_deterministic_algorithms(True)``.""" + launch, recorders = _prepared_launch(monkeypatch, k, False) + assert launch.plan.reproducible + if k == "k2": + assert dataclasses.replace(launch.plan, units=0, simt=(64, 64)).reproducible torch.use_deterministic_algorithms(True) - with pytest.raises(RuntimeError, match="deterministic=True"): - ATOMIC_MODE_CALLS[name]() - assert not recorder.calls + with _no_warnings(): + assert launch.launch() is launch.out + assert all(len(r.calls) == 1 for r in recorders) -@pytest.mark.parametrize("name", sorted(ATOMIC_MODE_CALLS)) -def test_explicit_false_under_warn_only_warns_and_runs_the_deterministic_path( - monkeypatch, torch_determinism, name +@pytest.mark.parametrize("splits", [2, 3]) +def test_k1_convenience_launch_alerts_only_when_order_varying( + monkeypatch, torch_determinism, splits ): - recorder = _patch_det(monkeypatch, name) - atomic1, atomic2 = Recorder(), Recorder() - monkeypatch.setattr(cake_backend, "_atomic_k1", atomic1) - monkeypatch.setattr(cake_backend, "_atomic_k2", atomic2) - torch.use_deterministic_algorithms(True, warn_only=True) - with pytest.warns(UserWarning, match="running the deterministic path"): - assert ATOMIC_MODE_CALLS[name]() is recorder.result - assert len(recorder.calls) == 1 and not atomic1.calls and not atomic2.calls + """``projection_wgrad(deterministic=False)`` on a row this package's planner takes at ``splits`` splits (CPU + binding, nothing loaded or launched), under ``torch.use_deterministic_algorithms(True)``: at three splits the + planning call and the call from the cached key both raise before the program; with ``warn_only=True`` the call + warns and launches; at two splits both calls launch without an alert. The cached launch keeps whether its plan + varies run to run.""" + entry_fn = _cpu_k1_launches(monkeypatch, "sm_103a") + monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) + monkeypatch.setattr(cake_backend, "k1_plan", lambda **kw: (128, splits)) + monkeypatch.setattr(cake_backend, "k1_chain_plan", lambda **kw: None) + bf16, T = torch.bfloat16, 1000 + G, X = torch.empty(T, 128, dtype=bf16), torch.empty(T, 4096, dtype=bf16) + out = torch.zeros(128, 4096) + + def call(): + return cake_backend.projection_wgrad(G, X, out, deterministic=False) + + torch.use_deterministic_algorithms(True) + for _ in range(2): # the planning call, then the call from the cached key + if splits == 3: + with pytest.raises(RuntimeError, match="at 3 splits"): + call() + else: + with _no_warnings(): + assert call() is out + ((_, entry),) = cake_backend._K1_ATOMIC_CACHE.items() + assert entry.plan.splits == splits and entry.k1 is not None + assert entry.order_varying == (splits == 3) + assert len(entry_fn.calls) == (0 if splits == 3 else 2) + if splits == 3: + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.warns(UserWarning, match="at 3 splits"): + assert call() is out + assert len(entry_fn.calls) == 1 + + +@pytest.mark.parametrize("splits", [2, 3]) +def test_k2_convenience_launch_alerts_only_when_order_varying( + monkeypatch, torch_determinism, splits +): + """``router_wgrad(deterministic=False, splits=splits)`` (CPU binding, nothing loaded or launched) under + ``torch.use_deterministic_algorithms(True)``: three units per tile (three partials per element) raise before the + split, the fill of ``out`` and the program, on the planning call and from the cached key; with ``warn_only=True`` + the call warns and launches; two units per tile launch without an alert.""" + entry_fn, fill = _cpu_k2_launches(monkeypatch) + _sm103a_wgrad_waves(monkeypatch) + split = Recorder() # the torch chain's split + monkeypatch.setattr(det_backend, "split_fp32_to_bf16x3_det_", split) + G, X, out = torch.empty(192, 256), torch.empty(192, 5120), torch.zeros(256, 5120) + recorders = (entry_fn, fill, split) + + def call(): + return cake_backend.router_wgrad(G, X, out, splits=splits, deterministic=False) + + torch.use_deterministic_algorithms(True) + for _ in range(2): # the planning call, then the call from the cached key + if splits == 3: + with pytest.raises(RuntimeError, match="does not have a deterministic"): + call() + else: + with _no_warnings(): + assert call() is out + ((_, entry),) = [ + (k, e) for k, e in cake_backend._K2_ATOMIC_CACHE.items() if e.plan is not None + ] + assert entry.plan.simt is None and entry.k2 is not None + assert entry.order_varying == (splits == 3) == (not entry.plan.reproducible) + assert any(r.calls for r in recorders) == (splits == 2) + if splits == 3: + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.warns(UserWarning, match=re.escape(entry.plan.template)): + assert call() is out + assert entry_fn.calls def test_routes_table(): @@ -7344,26 +7588,32 @@ def test_no_det_token_in_generated_names_and_no_negated_token_in_the_package(): assert not DET_TOKEN.search(name) -def test_select_module_and_the_loader_forward_deterministic_names(monkeypatch): +def test_select_module_and_the_loader_take_this_package_s_programs_only(monkeypatch): + """``select_module``, the loader and ``bind_launch`` of this package serve its own programs only: a template or a + program of the deterministic package is not forwarded to that package (``select_module`` raises, the loader goes + to this package's own loading, ``bind_launch`` finds no record).""" for arch, table in det_jit.KERNELS.items(): template = sorted(table)[0] - assert cake_jit.select_module(arch, template) == det_jit.select_module( - arch, template - ) + with pytest.raises(NotImplementedError, match="not registered"): + cake_jit.select_module(arch, template) with pytest.raises(NotImplementedError, match="not registered"): cake_jit.select_module("sm_100a", "dense_proj_gemm_atomic_nn_unregistered") - loader = Recorder() - monkeypatch.setattr(det_jit, "load_cake_dense_projection_det_gemm_module", loader) + det_loader, own_loader, det_bind = Recorder(), Recorder(), Recorder() + monkeypatch.setattr( + det_jit, "load_cake_dense_projection_det_gemm_module", det_loader + ) + monkeypatch.setattr(cake_jit, "_load_module", own_loader) + monkeypatch.setattr(det_backend, "bind_launch", det_bind) name = sorted(det_jit.MODULES)[0] + assert name not in cake_jit.MODULES assert ( cake_jit.load_cake_dense_projection_gemm_module(name, "sm_100a") - is loader.result + is own_loader.result ) - assert loader.calls == [((name, "sm_100a"), {})] - bind = Recorder() - monkeypatch.setattr(det_backend, "bind_launch", bind) - assert cake_backend.bind_launch(name, {"M": 1}, (2, 1, 1), "sm_100a") is bind.result - assert bind.calls == [((name, {"M": 1}, (2, 1, 1), "sm_100a"), {})] + assert own_loader.calls == [((name, "sm_100a"), {})] + with pytest.raises(KeyError): + cake_backend.bind_launch(name, {"M": 1}, (2, 1, 1), "sm_100a") + assert not det_loader.calls and not det_bind.calls def test_clear_convenience_caches_clears_both_packages(): @@ -7487,40 +7737,6 @@ def _bf16(shape, device, generator, scale=1.0): ) -def test_gpu_default_path_is_bitwise_the_deterministic_path(): - """The former names with the default give the deterministic package's bits (K1 forward / input gradient / - weight gradient, K2 forward / input gradient / weight gradient).""" - device = _det_device_or_skip() - g = torch.Generator(device=device).manual_seed(SEED) - T, K, N = 1001, 6144, 576 - X, W, G = ( - _bf16((T, K), device, g), - _bf16((N, K), device, g, 0.02), - _bf16((T, N), device, g), - ) - cases = [ - (cake_backend.projection_forward, det_backend.projection_det_forward, (X, W)), - (cake_backend.projection_dgrad, det_backend.projection_det_dgrad, (G, W)), - (cake_backend.projection_wgrad, det_backend.projection_det_wgrad, (G, X)), - ] - T2, K2, N2 = 1001, 6144, 256 - Xf = torch.randn(T2, K2, device=device, generator=g) - Wf = torch.randn(N2, K2, device=device, generator=g) * 0.02 - Gf = torch.randn(T2, N2, device=device, generator=g) - cases += [ - (cake_backend.router_forward, det_backend.router_det_forward, (Xf, Wf)), - (cake_backend.router_dgrad, det_backend.router_det_dgrad, (Gf, Wf)), - (cake_backend.router_wgrad, det_backend.router_det_wgrad, (Gf, Xf)), - ] - for ours, theirs, args in cases: - try: - expected = theirs(*args) - except NotImplementedError as exc: - pytest.skip(f"{theirs.__name__}: {exc}") - got = ours(*args) - assert got.dtype == expected.dtype and torch.equal(got, expected), ours.__name__ - - def _rel_fro(out: torch.Tensor, ref: torch.Tensor) -> float: return float((out.double() - ref).norm() / ref.norm()) @@ -7943,19 +8159,39 @@ def test_gpu_k1_small_token_counts_guard(): _k1_prepare(G, X, out, arch, 2) -def test_gpu_torch_deterministic_mode_rule(torch_determinism): - """Under ``torch.use_deterministic_algorithms(True)`` an explicit ``deterministic=False`` raises before any - launch; with ``warn_only=True`` it warns and returns the deterministic bits.""" - device = _det_device_or_skip() - G, X, out = _k1_wgrad_inputs("indexer_hw", 4097, torch.bfloat16, device, SEED + 40) - expected = det_backend.projection_det_wgrad(G, X, torch.empty_like(out)) +def test_gpu_torch_deterministic_mode_alerts_before_an_order_varying_launch( + k1_rule, torch_determinism +): + """Under ``torch.use_deterministic_algorithms(True)`` a ``deterministic=False`` call whose launch is fixed order + (two splits: one add) runs without an alert and repeats bit for bit; a launch whose result varies run to run + (three splits) raises PyTorch's ``RuntimeError`` before it launches (``out`` untouched), or with + ``warn_only=True`` warns and runs this package's program within the accuracy rule. The setting never changes + the program that runs.""" + device, arch = _atomic_device_or_skip() + G, X, out = _k1_wgrad_inputs("indexer_k", 16231, torch.bfloat16, device, SEED + 40) + # the FP64 reference before the mode is on (its cuBLAS GEMM would need CUBLAS_WORKSPACE_CONFIG under it) + ref = G.double().t() @ X.double() + k1_rule(G, X, out, 2) torch.use_deterministic_algorithms(True) - with pytest.raises(RuntimeError, match="deterministic=True"): + with _no_warnings(): + first = cake_backend.projection_wgrad(G, X, out, deterministic=False).clone() + for _ in range(RERUNS): + again = cake_backend.projection_wgrad(G, X, out, deterministic=False) + assert torch.equal(again, first) + k1_rule(G, X, out, 3) + out.fill_(7.0) + with pytest.raises(RuntimeError, match="at 3 splits"): cake_backend.projection_wgrad(G, X, out, deterministic=False) + torch.cuda.synchronize() + assert bool((out == 7.0).all()) torch.use_deterministic_algorithms(True, warn_only=True) - with pytest.warns(UserWarning, match="running the deterministic path"): + with pytest.warns(UserWarning, match="at 3 splits"): got = cake_backend.projection_wgrad(G, X, out, deterministic=False) - assert torch.equal(got, expected) + torch.cuda.synchronize() + torch.use_deterministic_algorithms(False) + _k1_check(first, ref, "two splits") + _k1_check(got, ref, "three splits") + assert _arena_tickets_zero() def test_gpu_det_then_atomic_share_a_stream_not_a_workspace(k1_rule): From c714bec3d5ab5a3e561b5164354957d5e9362514 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 5 Nov 2026 22:01:30 -0800 Subject: [PATCH 04/13] fix(dense_projection_gemm): the cluster program is fixed-order, no alert The K1 cluster program of #8844 (K1_DSM_RANGES_BY_SM) adds its two or three splits in cluster-rank order, so its result repeats bit for bit. AtomicGemmPlan.reproducible now counts it as fixed-order, and torch.use_deterministic_algorithms(True) raises no alert before it. Only the CTA-pair split-K program at three or more splits alerts. Tests: a prepared cluster launch at three splits and the shipped 4096 -> 64 FP32 cluster row at 2048 tokens (152-SM sm_100a) run under strict mode without an alert. --- .../dense_projection_gemm/README.md | 6 ++-- .../dense_projection_gemm/cake_backend.py | 11 ++++--- .../test_cake_dense_projection_gemm.py | 32 +++++++++++++++++++ 3 files changed, 41 insertions(+), 8 deletions(-) diff --git a/flashinfer/experimental/dense_projection_gemm/README.md b/flashinfer/experimental/dense_projection_gemm/README.md index 6a03e7faf0d..8e237d21a66 100644 --- a/flashinfer/experimental/dense_projection_gemm/README.md +++ b/flashinfer/experimental/dense_projection_gemm/README.md @@ -22,7 +22,7 @@ Every compute entry point of `cake_backend` takes the keyword-only `deterministi | call | `deterministic=False` | any other value (the default, `True`, `None`, `0`, `1`, `"False"`, a NumPy bool) | under `torch.use_deterministic_algorithms(True)` (strict or `warn_only=True`) | |---|---|---|---| -| this package's compute entry points (above) | runs this package's routes (table below) | `ValueError` naming the deterministic entry, before anything is planned, allocated or launched | the same routes and programs. Right before a launch whose result can vary run to run (the K1 split-K program at three or more splits; a K2 stream-K partition that gives an element more than two partials), PyTorch's alert: `RuntimeError` (nothing launched, `out` untouched), or a `UserWarning` and then the launch under `warn_only=True`. Fixed-order launches run without an alert | +| this package's compute entry points (above) | runs this package's routes (table below) | `ValueError` naming the deterministic entry, before anything is planned, allocated or launched | the same routes and programs. Right before a launch whose result can vary run to run (the K1 split-K program of CTA pairs at three or more splits; a K2 stream-K partition that gives an element more than two partials), PyTorch's alert: `RuntimeError` (nothing launched, `out` untouched), or a `UserWarning` and then the launch under `warn_only=True`. Fixed-order launches run without an alert | | this package's helpers (`router_symbol`, `bind_launch`, the planners `plan_atomic_gemm` / `plan_atomic_router`, the rule tables, `clear_convenience_caches`, `use_deterministic_backend`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged | | `flashinfer.experimental.dense_projection_det_gemm` (`projection_det_wgrad`, `router_fp32_det_gemm`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged: every launch is fixed-order | @@ -77,8 +77,8 @@ Rules: | accuracy | the deterministic package's budgets against FP64 of the exact operands, worst over repeated runs: BF16 out `atol = rtol = 1e-2`; FP32 out relative Frobenius error `<= 5e-5` and max error `<= 1e-3 * max abs(ref)`; router `<= 2e-6` relative Frobenius error and no row whose top-8 experts differ from FP64 (outside near ties of the 8th and 9th logits). At the largest token counts some thin K1 rows miss the BF16 (FP32) budget with the deterministic program and with this package's program alike; there the row keeps its route. The deterministic program: BF16 output of `sm_100a` `4096 -> 32`, `4096 -> 48`, `4096 -> 64`, `5120 -> 32`, `5120 -> 48` and of `sm_103a` `4096 -> 32`, `4096 -> 48`, `5120 -> 32`, `5120 -> 48`, `5120 -> 64`, `5120 -> 128` (from 45058, 49153, 65536 or 65537 tokens up by row: the K1 token tables below), and FP32 output of the rows of the odd-block K1 table at an odd number of K blocks from 98305 tokens up. This package's program: BF16 output of `sm_100a` `4096 -> 64` from about 43009 to 65535 tokens and `5120 -> 32` at 47105 to 65536 tokens, of `sm_103a` `4096 -> 32` from about 49153 to 65535 tokens, and on 148-SM parts of `sm_103a` `5120 -> 48` from 49152 tokens, `5120 -> 64` at 47104 and from 49153 tokens and `5120 -> 128` from 65537 tokens | | rounding | each output element is rounded once, from its complete FP32 sum | | subnormals | kept: K1 adds the partials with IEEE round-to-nearest adds, K2 reduces with the TMA reduce-add; a K2 SIMT window (forward, input and weight gradient) accumulates with round-to-nearest FMAs and adds, merges the sums of a CTA's slices (and in a compensated form their compensations) in a fixed order through shared memory (the forward: the sums of a cluster's eight CTAs, through distributed shared memory) and stores every element of `out` once, with no reduce-add; nothing flushes to zero | -| `torch.use_deterministic_algorithms(True)` | PyTorch's alert right before an order-varying launch: K1 at three or more splits (in the shipped tables only `sm_107a` rows, see "run-to-run bits, K1"; or a caller's `splits >= 3`), K2 stream-K with more than two partials per element (the `ROUTER_RULES` partitions, the `"waves"` windows, a caller's `splits >= 3` on `nn_t`). `RuntimeError`, or a `UserWarning` with `warn_only=True` and then the launch. The alert is checked on every launch (prepared and convenience), never at plan or prepare time. The two-split K1 rows, the K2 forward, the K2 SIMT windows and every deterministic program run without an alert | -| run-to-run bits, K1 | two splits (every `sm_100a` and `sm_103a` row, `sm_107a` `8192 -> 128`): repeat bit for bit (one add of two values); one split (`sm_107a` rows marked `p = 1` below): repeat bit for bit (whole tiles); three or four splits (the other `sm_107a` rows): vary in the last bits | +| `torch.use_deterministic_algorithms(True)` | PyTorch's alert right before an order-varying launch: the K1 CTA-pair program at three or more splits (in the shipped tables only `sm_107a` rows, see "run-to-run bits, K1"; or a caller's `splits >= 3`), K2 stream-K with more than two partials per element (the `ROUTER_RULES` partitions, the `"waves"` windows, a caller's `splits >= 3` on `nn_t`). `RuntimeError`, or a `UserWarning` with `warn_only=True` and then the launch. The alert is checked on every launch (prepared and convenience), never at plan or prepare time. The two-split K1 rows, the K1 cluster program (`K1_DSM_RANGES_BY_SM`, its splits added in cluster-rank order), the K2 forward, the K2 SIMT windows and every deterministic program run without an alert | +| run-to-run bits, K1 | two splits (every `sm_100a` and `sm_103a` row, `sm_107a` `8192 -> 128`): repeat bit for bit (one add of two values); the cluster program (`K1_DSM_RANGES_BY_SM`, two or three splits): repeats bit for bit (the splits added in cluster-rank order); one split (`sm_107a` rows marked `p = 1` below): repeat bit for bit (whole tiles); three or four splits (the other `sm_107a` rows): vary in the last bits | | run-to-run bits, K2 | forward: in a SIMT window repeats bit for bit (every element stored once, the lanes and the cluster's CTAs combined in a fixed order), else the deterministic program's bits; weight gradient on a `ROUTER_RULES` partition: varies in the last bits (more than two partials per element); a SIMT window (weight or input gradient): repeats bit for bit (every element stored once, its slices merged in a fixed order); with an explicit forward `splits` of 1 or 2: repeats bit for bit (`0 + a + b == 0 + b + a`), so recomputed router logits and top-k match | | architectures | `sm_100a`, `sm_103a`, `sm_107a`: one `-gencode=arch=compute_,code=sm_` build per program, minimum CUDA 12.8 / 12.9 / 13.4 (`build.json`) | | streams and threads | planning takes a module lock; the convenience path keeps one workspace arena per (device, stream); prepared launches own their workspace | diff --git a/flashinfer/experimental/dense_projection_gemm/cake_backend.py b/flashinfer/experimental/dense_projection_gemm/cake_backend.py index ff3b7f66e5c..ea0aa3d9dd3 100644 --- a/flashinfer/experimental/dense_projection_gemm/cake_backend.py +++ b/flashinfer/experimental/dense_projection_gemm/cake_backend.py @@ -73,7 +73,7 @@ :data:`K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM ``sm_100a`` / ``sm_103a`` parts) run another plan of those programs: the single-CTA program, every tile's whole K in one CTA. * Under ``torch.use_deterministic_algorithms(True)`` a launch whose result can vary run to run -- the K1 split-K - program at three or more splits, a K2 stream-K partition that gives an element more than two partials -- raises + program of CTA pairs at three or more splits, a K2 stream-K partition that gives an element more than two partials -- raises PyTorch's alert right before it launches (:func:`_alert_order_varying`): ``RuntimeError``, or with ``warn_only=True`` a ``UserWarning``, after which the launch runs. The fixed-order launches (every other route) do not alert. The PyTorch setting never changes which program a call runs. @@ -232,8 +232,8 @@ def _warning_stacklevel() -> int: def _alert_order_varying(launch: str) -> None: """PyTorch's alert for an operation without a deterministic implementation, raised right before this package - launches ``launch``, a program whose result can vary run to run (the K1 split-K program at three or more splits, - a K2 stream-K partition that gives an element more than two partials), while PyTorch's deterministic algorithms + launches ``launch``, a program whose result can vary run to run (the K1 split-K program of CTA pairs at three or + more splits, a K2 stream-K partition that gives an element more than two partials), while PyTorch's deterministic algorithms are enabled: a ``RuntimeError``, or a ``UserWarning`` under ``torch.use_deterministic_algorithms(True, warn_only=True)``, after which the launch runs; nothing otherwise. The alert of PyTorch's own operators (PyTorch exposes no public Python entry for it). Fixed-order launches do not call it, and it never changes the program @@ -1308,8 +1308,9 @@ def items(self) -> int: @property def reproducible(self) -> bool: """Whether the result repeats bit for bit: one split stores whole tiles (its chains added in order), two add - ``own + other`` (one add of two values); three or more add the partials in the order the splits finish.""" - return self.splits <= 2 + ``own + other`` (one add of two values), the cluster program (``dsm``) adds its splits in cluster-rank order; + three or more splits of the CTA-pair program add the partials in the order the splits finish.""" + return self.dsm or self.splits <= 2 @property def grid(self) -> tuple[int, int, int]: diff --git a/tests/experimental/test_cake_dense_projection_gemm.py b/tests/experimental/test_cake_dense_projection_gemm.py index 89414a4c054..67252db9f6e 100644 --- a/tests/experimental/test_cake_dense_projection_gemm.py +++ b/tests/experimental/test_cake_dense_projection_gemm.py @@ -7657,6 +7657,38 @@ def test_a_fixed_order_prepared_launch_does_not_alert( with _no_warnings(): assert launch.launch() is launch.out assert all(len(r.calls) == 1 for r in recorders) + if k == "k1": + # the cluster program at three splits: its splits are added in cluster-rank order + launch.plan = dataclasses.replace(launch.plan, splits=3, dsm=True) + assert launch.plan.reproducible + with _no_warnings(): + assert launch.launch() is launch.out + assert all(len(r.calls) == 2 for r in recorders) + + +def test_the_cluster_program_does_not_alert(monkeypatch, torch_determinism): + """The shipped cluster row ``4096 -> 64`` FP32 at 2048 tokens on a 152-SM ``sm_100a`` part (three splits, added in + cluster-rank order) runs without an alert under ``torch.use_deterministic_algorithms(True)``: the planning call, + the call from the cached key and the prepared launch (CPU binding, nothing loaded or launched).""" + entry_fn = _cpu_k1_launches(monkeypatch, "sm_100a") + monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) + bf16, T = torch.bfloat16, 2048 + X, G = torch.empty(T, 4096, dtype=bf16), torch.empty(T, 64, dtype=bf16) + out_w = torch.empty(64, 4096) + torch.use_deterministic_algorithms(True) + with _no_warnings(): + for _ in range(2): + assert ( + cake_backend.projection_wgrad(G, X, out_w, deterministic=False) is out_w + ) + prepared = cake_backend.prepare_projection_wgrad( + G, X, out_w, deterministic=False + ) + assert prepared.launch() is out_w + ((_, atomic),) = cake_backend._K1_ATOMIC_CACHE.items() + assert atomic.plan.dsm and atomic.plan.splits == 3 and not atomic.order_varying + assert prepared.plan.dsm and prepared.plan.reproducible + assert len(entry_fn.calls) == 3 @pytest.mark.parametrize("splits", [2, 3]) From c17acc02ba48b1ddd5513d2cf61fed5c30d4f625 Mon Sep 17 00:00:00 2001 From: Avery Huang Date: Thu, 5 Nov 2026 23:40:52 -0800 Subject: [PATCH 05/13] refactor(cake_dense_gemm_det_train): move the deterministic dense GEMMs to their #8928 names The deterministic dense projection GEMM package becomes the public module flashinfer.cake_dense_gemm_det_train (flashinfer-ai/flashinfer#8928). Its backend moves from flashinfer/experimental/dense_projection_det_gemm/ to flashinfer/experimental/cake/cake_dense_gemm_det_train/. The generated sources, program and library names, kernel namespaces and templates keep their names; build.json changes only its "package" field. - The public module imports the backend on first attribute access and returns the backend's own objects. It is not an experimental API. - The entry points and types take the module name: dense_projection_det_gemm becomes cake_dense_gemm_det_train_projection, router_fp32_det_gemm becomes cake_dense_gemm_det_train_router, and the forward / dgrad / wgrad / plan / prepare entries and the four plan and launch types follow. - New support query is_cake_dense_gemm_det_train_supported(device). - The 18 aliases of the names before the det rename and the loader alias load_cake_dense_projection_gemm_module are removed. The former import path and names fail at import; there is no alias. - No entry point takes a mode argument, and none reads the PyTorch deterministic-algorithms setting. - The names test checks the mode word, the absence of a deterministic= keyword and of the former names and paths, the lazy import, and the unchanged bitwise outputs under torch.use_deterministic_algorithms. The non-deterministic package moves in a following commit of the same PR. --- .pre-commit-config.yaml | 4 +- ....py => bench_cake_dense_gemm_det_train.py} | 16 +- flashinfer/cake_dense_gemm_det_train.py | 87 +++ .../cake_dense_gemm_det_train}/.gitignore | 0 .../cake_dense_gemm_det_train}/README.md | 99 +-- .../cake_dense_gemm_det_train}/__init__.py | 10 +- .../cake_dense_gemm_det_train}/build.json | 2 +- .../cake_backend.py | 255 +++---- .../cake_dense_gemm_det_train}/cake_build.py | 0 .../cake_dense_gemm_det_train}/cake_jit.py | 3 - .../.clang-format | 0 ...t_gemm_011c85b820ec0cf6863d_tpl_binding.cu | 0 ...n_det_gemm_02803e1d2393113c0f6c_binding.cu | 0 ...on_det_gemm_02803e1d2393113c0f6c_kernel.cu | 0 ...n_det_gemm_0618ae0aec1af120afd0_binding.cu | 0 ...on_det_gemm_0618ae0aec1af120afd0_kernel.cu | 0 ...on_det_gemm_087ba4f157914949c26a_kernel.cu | 0 ...on_det_gemm_0a71a49da71c7d34ab37_kernel.cu | 0 ...n_det_gemm_0b91f29d88de00658fc5_binding.cu | 0 ...on_det_gemm_0b91f29d88de00658fc5_kernel.cu | 0 ...n_det_gemm_0bbd5d9f1d72a90ef3ea_binding.cu | 0 ...on_det_gemm_0bbd5d9f1d72a90ef3ea_kernel.cu | 0 ...t_gemm_0c72e864a4afb5347da6_tpl_binding.cu | 0 ...n_det_gemm_0d487a8f71dac54321ef_binding.cu | 0 ...on_det_gemm_0d487a8f71dac54321ef_kernel.cu | 0 ...on_det_gemm_0db65f3ddd780ed42d4d_kernel.cu | 0 ...n_det_gemm_0dd5420b570cabc144ab_binding.cu | 0 ...on_det_gemm_0dd5420b570cabc144ab_kernel.cu | 0 ...n_det_gemm_13790f025541d7d162a7_binding.cu | 0 ...on_det_gemm_13790f025541d7d162a7_kernel.cu | 0 ...on_det_gemm_13d9310f37868a18eace_kernel.cu | 0 ...on_det_gemm_14192244441c1003db74_kernel.cu | 0 ...n_det_gemm_20a2e5c04427e93572da_binding.cu | 0 ...on_det_gemm_20a2e5c04427e93572da_kernel.cu | 0 ...n_det_gemm_23305bdb787b7c6d7378_binding.cu | 0 ...on_det_gemm_23305bdb787b7c6d7378_kernel.cu | 0 ...on_det_gemm_295aef2f2e0332de2372_kernel.cu | 0 ...n_det_gemm_29d013dfb72e932eaee5_binding.cu | 0 ...on_det_gemm_29d013dfb72e932eaee5_kernel.cu | 0 ...n_det_gemm_29de51d5e9cdd6d22491_binding.cu | 0 ...on_det_gemm_29de51d5e9cdd6d22491_kernel.cu | 0 ...on_det_gemm_3088b8056e0b7f9ca32c_kernel.cu | 0 ...n_det_gemm_31be7faa868eabcedada_binding.cu | 0 ...on_det_gemm_31be7faa868eabcedada_kernel.cu | 0 ...n_det_gemm_3a87cc189b4a7c1a8d7d_binding.cu | 0 ...on_det_gemm_3a87cc189b4a7c1a8d7d_kernel.cu | 0 ...n_det_gemm_3d7b66f61c710f893b29_binding.cu | 0 ...on_det_gemm_3d7b66f61c710f893b29_kernel.cu | 0 ...on_det_gemm_438841f42a0743b63fb5_kernel.cu | 0 ...on_det_gemm_446497e750658dfb057b_kernel.cu | 0 ...on_det_gemm_49acb91191bbcced9051_kernel.cu | 0 ...n_det_gemm_4de304a13db1f2a0e850_binding.cu | 0 ...on_det_gemm_4de304a13db1f2a0e850_kernel.cu | 0 ...n_det_gemm_4eb96fb2eb1aaa14c7f7_binding.cu | 0 ...on_det_gemm_4eb96fb2eb1aaa14c7f7_kernel.cu | 0 ...on_det_gemm_505fbe05ce8b62f0a589_kernel.cu | 0 ...n_det_gemm_51c81f11e3b490b5aebb_binding.cu | 0 ...on_det_gemm_51c81f11e3b490b5aebb_kernel.cu | 0 ...on_det_gemm_5226b6c854d7db395d7e_kernel.cu | 0 ...on_det_gemm_5445e0f16ad84868fa01_kernel.cu | 0 ...n_det_gemm_552070c21b97faa67a71_binding.cu | 0 ...on_det_gemm_552070c21b97faa67a71_kernel.cu | 0 ...on_det_gemm_5568657a526e2040f0bc_kernel.cu | 0 ...on_det_gemm_56be8af43d5d0822da87_kernel.cu | 0 ...n_det_gemm_59fb66ccddd67565e592_binding.cu | 0 ...on_det_gemm_59fb66ccddd67565e592_kernel.cu | 0 ...on_det_gemm_5a06c5e2f33772a33174_kernel.cu | 0 ...n_det_gemm_5b4e87b523f799e09384_binding.cu | 0 ...on_det_gemm_5b4e87b523f799e09384_kernel.cu | 0 ...on_det_gemm_5e77ef909e88a99c4357_kernel.cu | 0 ...n_det_gemm_624eb3d7232071ddaf91_binding.cu | 0 ...on_det_gemm_624eb3d7232071ddaf91_kernel.cu | 0 ...n_det_gemm_6579637169bec3c2f1d5_binding.cu | 0 ...on_det_gemm_6579637169bec3c2f1d5_kernel.cu | 0 ...on_det_gemm_66df6446dd2cca3106d3_kernel.cu | 0 ...n_det_gemm_736ad30e83d88796883f_binding.cu | 0 ...on_det_gemm_736ad30e83d88796883f_kernel.cu | 0 ...n_det_gemm_75b1dcdaef42319f5ed4_binding.cu | 0 ...on_det_gemm_75b1dcdaef42319f5ed4_kernel.cu | 0 ...on_det_gemm_765ece06955b70a95b29_kernel.cu | 0 ...n_det_gemm_7a0ce55fd7221e354322_binding.cu | 0 ...on_det_gemm_7a0ce55fd7221e354322_kernel.cu | 0 ...on_det_gemm_7d403fa7b167188ccc24_kernel.cu | 0 ...n_det_gemm_7ef28d420de555e7a28d_binding.cu | 0 ...on_det_gemm_7ef28d420de555e7a28d_kernel.cu | 0 ...on_det_gemm_806927d101fe72fc1317_kernel.cu | 0 ...on_det_gemm_82f698f67a1cca6e4291_kernel.cu | 0 ...on_det_gemm_895f9ba7159bea07aedd_kernel.cu | 0 ...on_det_gemm_8e1f6805342f620cca8c_kernel.cu | 0 ...n_det_gemm_9269f6841d7aeb58471c_binding.cu | 0 ...on_det_gemm_9269f6841d7aeb58471c_kernel.cu | 0 ...n_det_gemm_9475731f174aaefb2e44_binding.cu | 0 ...on_det_gemm_9475731f174aaefb2e44_kernel.cu | 0 ...on_det_gemm_9aa9aa60b0c544bfab06_kernel.cu | 0 ...on_det_gemm_9cc6363566760a38b4fe_kernel.cu | 0 ...n_det_gemm_9ffc37c198f2d2f42c1f_binding.cu | 0 ...on_det_gemm_9ffc37c198f2d2f42c1f_kernel.cu | 0 ...n_det_gemm_a24a3d3fd3ceac3e1cdf_binding.cu | 0 ...on_det_gemm_a24a3d3fd3ceac3e1cdf_kernel.cu | 0 ...n_det_gemm_aa4f7c2fbd0b9a99e535_binding.cu | 0 ...on_det_gemm_aa4f7c2fbd0b9a99e535_kernel.cu | 0 ...on_det_gemm_b26ac3f3ed56a78cb052_kernel.cu | 0 ...n_det_gemm_b2a9665ee12fd6afea2b_binding.cu | 0 ...on_det_gemm_b2a9665ee12fd6afea2b_kernel.cu | 0 ...on_det_gemm_b58b50619ef5ed46e109_kernel.cu | 0 ...n_det_gemm_b773cc0e3894fd0ef2e0_binding.cu | 0 ...on_det_gemm_b773cc0e3894fd0ef2e0_kernel.cu | 0 ...n_det_gemm_bab00fb2b95c93758095_binding.cu | 0 ...on_det_gemm_bab00fb2b95c93758095_kernel.cu | 0 ...n_det_gemm_bab9e9dfd7c6ca5e17c0_binding.cu | 0 ...on_det_gemm_bab9e9dfd7c6ca5e17c0_kernel.cu | 0 ...n_det_gemm_bb0e7457db664d6f0f6a_binding.cu | 0 ...on_det_gemm_bb0e7457db664d6f0f6a_kernel.cu | 0 ...n_det_gemm_c363601373e2dfb8e9af_binding.cu | 0 ...on_det_gemm_c363601373e2dfb8e9af_kernel.cu | 0 ...n_det_gemm_c63d1400e64208d7865c_binding.cu | 0 ...on_det_gemm_c63d1400e64208d7865c_kernel.cu | 0 ...on_det_gemm_ca9a846738b1443f7318_kernel.cu | 0 ...on_det_gemm_d2a5e70bd49083888179_kernel.cu | 0 ...n_det_gemm_d2ea37e4d394bf40922a_binding.cu | 0 ...on_det_gemm_d2ea37e4d394bf40922a_kernel.cu | 0 ...on_det_gemm_d6ebaab2415ea7b00de9_kernel.cu | 0 ...on_det_gemm_ddff2b911137449c1dfc_kernel.cu | 0 ...on_det_gemm_de7492eeb6e24b31687a_kernel.cu | 0 ...ense_projection_det_gemm_device_common.cuh | 0 ...n_det_gemm_dfd65a75564d04d3826f_binding.cu | 0 ...on_det_gemm_dfd65a75564d04d3826f_kernel.cu | 0 ...n_det_gemm_e2555d533680e08880ae_binding.cu | 0 ...on_det_gemm_e2555d533680e08880ae_kernel.cu | 0 ...n_det_gemm_e33ac44bfff66daaee64_binding.cu | 0 ...on_det_gemm_e33ac44bfff66daaee64_kernel.cu | 0 ...on_det_gemm_e98b3ee7b774c63d7544_kernel.cu | 0 ...n_det_gemm_eb2bb538645fe303bcca_binding.cu | 0 ...on_det_gemm_eb2bb538645fe303bcca_kernel.cu | 0 ...n_det_gemm_ebbb505a3077e8da0a50_binding.cu | 0 ...on_det_gemm_ebbb505a3077e8da0a50_kernel.cu | 0 ...n_det_gemm_ec7277a904a7211d67ce_binding.cu | 0 ...on_det_gemm_ec7277a904a7211d67ce_kernel.cu | 0 ...n_det_gemm_ed1f3ab86ddc55345e69_binding.cu | 0 ...on_det_gemm_ed1f3ab86ddc55345e69_kernel.cu | 0 ...n_det_gemm_edf7fa4c4f889f0c4981_binding.cu | 0 ...on_det_gemm_edf7fa4c4f889f0c4981_kernel.cu | 0 ...n_det_gemm_eec6f87f8b7f25eb6aa2_binding.cu | 0 ...on_det_gemm_eec6f87f8b7f25eb6aa2_kernel.cu | 0 ...on_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu | 0 ...on_det_gemm_f31e1dcb1df00973dc0c_kernel.cu | 0 ...n_det_gemm_f589a6aff5e88318c36c_binding.cu | 0 ...on_det_gemm_f589a6aff5e88318c36c_kernel.cu | 0 ...n_det_gemm_f6c9b40184336f49390d_binding.cu | 0 ...on_det_gemm_f6c9b40184336f49390d_kernel.cu | 0 ...n_det_gemm_f993d8de6f29f9e8a1ac_binding.cu | 0 ...on_det_gemm_f993d8de6f29f9e8a1ac_kernel.cu | 0 ...on_det_gemm_fb4bcdf13915dfb96c59_kernel.cu | 0 ...n_det_gemm_fb766b1be68412e6fda7_binding.cu | 0 ...on_det_gemm_fb766b1be68412e6fda7_kernel.cu | 0 ...n_det_gemm_fdb106e6e73897521cbb_binding.cu | 0 ...on_det_gemm_fdb106e6e73897521cbb_kernel.cu | 0 ..._dense_projection_det_gemm_host_common.cuh | 0 ...ion_det_gemm_tpl_kk_n128_c1__hen_kernel.cu | 0 ...n_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu | 0 ..._tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu | 0 ...on_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu | 0 ...emm_tpl_kn_n128_f32_tma1_c1__pd2_kernel.cu | 0 ...m_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu | 0 ..._gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu | 0 ..._tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu | 0 ..._tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu | 0 ...nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu | 0 ...det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu | 0 ...tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu | 0 ...l_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu | 0 ...t_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu | 0 ..._async_bulk_gmem2smem_5198c9805c6ecc8d.cuh | 0 .../cake_elect_commit_8131fdc67daf4d23.cuh | 0 ..._commit_cg2_multicast_18e7b4ad8e69c467.cuh | 0 .../cake_elect_sync_1ee8cd91025d8932.cuh | 0 ...ake_incr_smem_desc_lo_4429b045244f74bc.cuh | 0 ...ake_make_warp_uniform_15c3664d3e11e8cd.cuh | 0 .../cake_mbarrier_arrive_71ef360c3ac78eb7.cuh | 0 ...rier_arrive_expect_tx_54ce59d489f69526.cuh | 0 .../cake_mbarrier_init_15e581aef85ee586.cuh | 0 .../cake_mbarrier_wait_c93bf7e39b218733.cuh | 0 ...ier_wait_cluster_hint_bdb29a0e36dedbee.cuh | 0 .../cake_tcgen05_mma_f16_03fde8c1cbf058c8.cuh | 0 ...cake_tma_3d_gmem2smem_7df1e07811f83fc6.cuh | 0 ...tma_3d_gmem2smem_cta2_dcd29ed146891990.cuh | 0 .../cake_tma_store_3d_a0d668161ec63906.cuh | 0 .../cake_tmem_ld_x16_a9ed61940fdba296.cuh | 0 .../csrc/cake_ffi_checks.h | 0 pyproject.toml | 5 +- scripts/check_cake_build_manifests.py | 4 +- ...m.py => test_cake_dense_gemm_det_train.py} | 686 +++++++++--------- .../test_cake_dense_gemm_det_train_names.py | 385 ++++++++++ ...st_cake_dense_projection_det_gemm_names.py | 220 ------ tests/test_helpers/cake_package.py | 2 +- 195 files changed, 1016 insertions(+), 762 deletions(-) rename benchmarks/{bench_cake_dense_projection_det_gemm.py => bench_cake_dense_gemm_det_train.py} (95%) create mode 100644 flashinfer/cake_dense_gemm_det_train.py rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/.gitignore (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/README.md (86%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/__init__.py (68%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/build.json (99%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/cake_backend.py (96%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/cake_build.py (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/cake_jit.py (99%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/.clang-format (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_011c85b820ec0cf6863d_tpl_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_087ba4f157914949c26a_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0a71a49da71c7d34ab37_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0c72e864a4afb5347da6_tpl_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0db65f3ddd780ed42d4d_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13d9310f37868a18eace_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_14192244441c1003db74_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_295aef2f2e0332de2372_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3088b8056e0b7f9ca32c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_438841f42a0743b63fb5_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_446497e750658dfb057b_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_49acb91191bbcced9051_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_505fbe05ce8b62f0a589_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5226b6c854d7db395d7e_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5445e0f16ad84868fa01_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5568657a526e2040f0bc_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_56be8af43d5d0822da87_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5a06c5e2f33772a33174_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5e77ef909e88a99c4357_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_66df6446dd2cca3106d3_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_765ece06955b70a95b29_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7d403fa7b167188ccc24_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_806927d101fe72fc1317_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_82f698f67a1cca6e4291_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_895f9ba7159bea07aedd_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_8e1f6805342f620cca8c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9aa9aa60b0c544bfab06_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9cc6363566760a38b4fe_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b26ac3f3ed56a78cb052_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b58b50619ef5ed46e109_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ca9a846738b1443f7318_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2a5e70bd49083888179_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d6ebaab2415ea7b00de9_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ddff2b911137449c1dfc_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_de7492eeb6e24b31687a_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_device_common.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e98b3ee7b774c63d7544_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f31e1dcb1df00973dc0c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb4bcdf13915dfb96c59_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_binding.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_host_common.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_c1__hen_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n128_f32_tma1_c1__pd2_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_cp_async_bulk_gmem2smem_5198c9805c6ecc8d.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_elect_commit_8131fdc67daf4d23.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_elect_commit_cg2_multicast_18e7b4ad8e69c467.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_elect_sync_1ee8cd91025d8932.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_incr_smem_desc_lo_4429b045244f74bc.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_make_warp_uniform_15c3664d3e11e8cd.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_mbarrier_arrive_71ef360c3ac78eb7.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_mbarrier_arrive_expect_tx_54ce59d489f69526.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_mbarrier_init_15e581aef85ee586.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_mbarrier_wait_c93bf7e39b218733.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_mbarrier_wait_cluster_hint_bdb29a0e36dedbee.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_tcgen05_mma_f16_03fde8c1cbf058c8.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_7df1e07811f83fc6.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_cta2_dcd29ed146891990.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_tma_store_3d_a0d668161ec63906.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_device_helpers/cake_tmem_ld_x16_a9ed61940fdba296.cuh (100%) rename flashinfer/experimental/{dense_projection_det_gemm => cake/cake_dense_gemm_det_train}/csrc/cake_ffi_checks.h (100%) rename tests/experimental/{test_cake_dense_projection_det_gemm.py => test_cake_dense_gemm_det_train.py} (94%) create mode 100644 tests/experimental/test_cake_dense_gemm_det_train_names.py delete mode 100644 tests/experimental/test_cake_dense_projection_det_gemm_names.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d921d2f5a6a..14a51d817b1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,13 +39,13 @@ files: ^(?!(?:flashinfer/moe_ep/kernel_src/blackwell_bf16_rank_major/src/|csrc/c # and sealed by cake_sm120_megamoe_manifest.json; keep their bytes. # The csrc/ trees of the Cake-exported packages (flashinfer-ai/flashinfer#6553: the rmsnorm # training and mHC training backends, the causal conv1d, ulysses and fused QK RoPE append (BF16 / FP8) backends, cake_dsa_train, cake_kda_train / cake_kda_det_train / cake_kda_det_train_kcp, cake_dsa_indexer_det, cake_mok_bf16, cake_moe_grouped_det_gemm, -# cake_lm_head_loss / cake_lm_head_det_loss, dense_projection_det_gemm, dense_projection_gemm, cake_stepfun_moe, nvfp4_attention, deepgemm_fp8_gemm, deepgemm_batched_gemm, deepgemm_mega_gate, nvfp4_sparse_mla_decode) are rendered by the Cake exporters and +# cake_lm_head_loss / cake_lm_head_det_loss, cake_dense_gemm_det_train, dense_projection_gemm, cake_stepfun_moe, nvfp4_attention, deepgemm_fp8_gemm, deepgemm_batched_gemm, deepgemm_mega_gate, nvfp4_sparse_mla_decode) are rendered by the Cake exporters and # digest-checked by their build.json manifests; keep their bytes. The package-root loaders, # cake_build.py and build.json next to them stay checked. # Keep both old and det package paths while the KDA, conv1d and LM-head renames land independently. # The Qwen 3.5 gated attention block's _cudnn_frost/ tree is cuDNN Frontend source vendored byte-faithfully # (VENDOR_MANIFEST.json digests the files as written); keep its bytes. -exclude: ^(?:csrc/kda/ptx/|flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_trtllm_mla_blackwell/generated/|cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|cake_minimax_h3_qkv_quantize_pack/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_fp4_allreduce/.*_kernel\.cu|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_union/|cake_gemm_allreduce/|cake_fused_norm_combine/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:cake_rmsnorm_det_train_backend/csrc/|cake_rmsnorm_train_backend/csrc/|cake_causal_conv1d_train_backend/csrc/|cake_causal_conv1d_det_train_backend/csrc/|experimental/cake/cake_causal_conv1d_train_backend/csrc/|cake_sage_block_sparse_attention_backend/csrc/|cake_ulysses_backend/csrc/|cake_fused_qk_rope_append_backend/csrc/|cake_fused_qk_rope_fp8_append_backend/csrc/|cake_mhc_det_train_backend/csrc/|experimental/cake/cake_mhc_train_backend/csrc/|experimental/cake/cake_lm_head_loss/csrc/|experimental/cake/cake_moe_grouped_gemm/csrc/|experimental/nvfp4_sparse_mla_decode/csrc/(?:cake_nvfp4_sparse_mla_decode/|cake_ffi_checks\.h)|experimental/(?:cake_mok_bf16|cake_mxfp8_megamoe_ep16|cake_moe_grouped_det_gemm|sm110_gqa_decode|nvfp4_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention|kimi_k3_fused_router|msa_nvfp4_decode|cake_lm_head_loss|cake_lm_head_det_loss|cake_dsa_train|cake_kda_train|cake_kda_det_train|cake_kda_det_train_kcp|cake_dsa_indexer_det|cake_msa_icp_indexer|dense_projection_det_gemm|dense_projection_gemm|cake_stepfun_moe|deepgemm_fp8_gemm|deepgemm_batched_gemm|deepgemm_mega_gate)/csrc/|moe_ep/kernel_src/(?:(?:sm100/(?:cutedsl_megamoe|next_cutedsl_megamoe)|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/(?:swapab|split)_cutedsl_megakernel)/src/|sm90/cake_bf16_megamoe/src/cake_sm90_bf16_megamoe_|sm120/cake_mxfp4_megamoe/src/cake_sm120_megamoe_(?:kernel.*\.cuh|m(?:128|64)_(?:no)?copier\.cu|geometry\.json)))|csrc/generated/dcp_alltoall/|flashinfer/experimental/qwen_3_5_block/_cudnn_frost/) +exclude: ^(?:csrc/kda/ptx/|flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_trtllm_mla_blackwell/generated/|cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|cake_minimax_h3_qkv_quantize_pack/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_fp4_allreduce/.*_kernel\.cu|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_union/|cake_gemm_allreduce/|cake_fused_norm_combine/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:cake_rmsnorm_det_train_backend/csrc/|cake_rmsnorm_train_backend/csrc/|cake_causal_conv1d_train_backend/csrc/|cake_causal_conv1d_det_train_backend/csrc/|experimental/cake/cake_causal_conv1d_train_backend/csrc/|cake_sage_block_sparse_attention_backend/csrc/|cake_ulysses_backend/csrc/|cake_fused_qk_rope_append_backend/csrc/|cake_fused_qk_rope_fp8_append_backend/csrc/|cake_mhc_det_train_backend/csrc/|experimental/cake/cake_mhc_train_backend/csrc/|experimental/cake/cake_lm_head_loss/csrc/|experimental/cake/cake_moe_grouped_gemm/csrc/|experimental/cake/cake_dense_gemm_det_train/csrc/|experimental/nvfp4_sparse_mla_decode/csrc/(?:cake_nvfp4_sparse_mla_decode/|cake_ffi_checks\.h)|experimental/(?:cake_mok_bf16|cake_mxfp8_megamoe_ep16|cake_moe_grouped_det_gemm|sm110_gqa_decode|nvfp4_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention|kimi_k3_fused_router|msa_nvfp4_decode|cake_lm_head_loss|cake_lm_head_det_loss|cake_dsa_train|cake_kda_train|cake_kda_det_train|cake_kda_det_train_kcp|cake_dsa_indexer_det|cake_msa_icp_indexer|dense_projection_gemm|cake_stepfun_moe|deepgemm_fp8_gemm|deepgemm_batched_gemm|deepgemm_mega_gate)/csrc/|moe_ep/kernel_src/(?:(?:sm100/(?:cutedsl_megamoe|next_cutedsl_megamoe)|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/(?:swapab|split)_cutedsl_megakernel)/src/|sm90/cake_bf16_megamoe/src/cake_sm90_bf16_megamoe_|sm120/cake_mxfp4_megamoe/src/cake_sm120_megamoe_(?:kernel.*\.cuh|m(?:128|64)_(?:no)?copier\.cu|geometry\.json)))|csrc/generated/dcp_alltoall/|flashinfer/experimental/qwen_3_5_block/_cudnn_frost/) repos: # Standard hooks diff --git a/benchmarks/bench_cake_dense_projection_det_gemm.py b/benchmarks/bench_cake_dense_gemm_det_train.py similarity index 95% rename from benchmarks/bench_cake_dense_projection_det_gemm.py rename to benchmarks/bench_cake_dense_gemm_det_train.py index cbe87f89f18..b60fa5a1b9b 100644 --- a/benchmarks/bench_cake_dense_projection_det_gemm.py +++ b/benchmarks/bench_cake_dense_gemm_det_train.py @@ -16,7 +16,7 @@ """Benchmark the experimental Cake dense projection GEMMs (SM100 / SM103 / SM107). -Times the GLM-5.2 training GEMMs of ``flashinfer.experimental.dense_projection_det_gemm`` +Times the GLM-5.2 training GEMMs of ``flashinfer.experimental.cake.cake_dense_gemm_det_train`` paired against ``torch.matmul`` on the same strided views: * the 11 projection rows (``o_proj``, ``q_a``, ``q_b``, ``kv_a``, ``shared_gate_up``, @@ -55,7 +55,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from flashinfer.experimental.dense_projection_det_gemm import cake_backend # noqa: E402 +from flashinfer.experimental.cake.cake_dense_gemm_det_train import cake_backend # noqa: E402 from flashinfer.testing import bench_gpu_time # noqa: E402 # GLM-5.2 projection rows: label -> (K = in features, N = out features). @@ -190,15 +190,15 @@ def rnd(*shape, scale=1.0): def cake_arm(spec, inp): """The prepared Cake launch of the row (allocation-free replays).""" if spec["family"] == "router": - prepared = cake_backend.prepare_router_fp32_det_gemm( + prepared = cake_backend.cake_dense_gemm_det_train_router_prepare( inp["A"], inp["B"], inp["out"] ) elif spec["family"] == "proj" and spec["op"] == "wgrad": - prepared = cake_backend.prepare_projection_det_wgrad( + prepared = cake_backend.cake_dense_gemm_det_train_projection_wgrad_prepare( inp["G"], inp["X"], inp["out"] ) else: - prepared = cake_backend.prepare_dense_projection_det_gemm( + prepared = cake_backend.cake_dense_gemm_det_train_projection_prepare( inp["A"], inp["B"], inp["out"] ) return prepared @@ -208,10 +208,10 @@ def eager_arm(spec, inp): """The convenience entry point of the row (the cached planned launch; ``out`` is passed, so a cache hit allocates nothing).""" if spec["family"] == "router": - return lambda: cake_backend.router_fp32_det_gemm(inp["A"], inp["B"], inp["out"]) + return lambda: cake_backend.cake_dense_gemm_det_train_router(inp["A"], inp["B"], inp["out"]) if spec["family"] == "proj" and spec["op"] == "wgrad": - return lambda: cake_backend.projection_det_wgrad(inp["G"], inp["X"], inp["out"]) - return lambda: cake_backend.dense_projection_det_gemm( + return lambda: cake_backend.cake_dense_gemm_det_train_projection_wgrad(inp["G"], inp["X"], inp["out"]) + return lambda: cake_backend.cake_dense_gemm_det_train_projection( inp["A"], inp["B"], inp["out"] ) diff --git a/flashinfer/cake_dense_gemm_det_train.py b/flashinfer/cake_dense_gemm_det_train.py new file mode 100644 index 00000000000..8cab184353d --- /dev/null +++ b/flashinfer/cake_dense_gemm_det_train.py @@ -0,0 +1,87 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Cake dense GEMMs for training, deterministic (SM100 / SM103 / SM107): the GLM-5.2 +BF16 dense projection GEMM (K1) and the FP32 router-gate GEMM through split-BF16x3 +emulation (K2). The same inputs, hardware and plan give bitwise the same outputs +on every run. The module selects the mode: no entry point takes a mode argument, +and no entry point raises or warns under ``torch.use_deterministic_algorithms``. + +The backend lives in ``flashinfer.experimental.cake.cake_dense_gemm_det_train`` +(``cake_backend.py``, the generated sources and ``build.json``). This module +imports it on the first attribute access and returns the backend's own objects, +so a call adds no Python layer. ``CONVENIENCE_CACHE_SIZE`` is a setting of the +backend module (``cake_backend.CONVENIENCE_CACHE_SIZE``), not of this module. + +Usage:: + + from flashinfer.cake_dense_gemm_det_train import ( + cake_dense_gemm_det_train_projection_forward, + cake_dense_gemm_det_train_router_wgrad, + ) + + y = cake_dense_gemm_det_train_projection_forward(X, W) # y[T, N] = X[T, K] @ W[N, K].T + dW = cake_dense_gemm_det_train_router_wgrad(G, X) # FP32 dW[N, K] = G[T, N].T @ X[T, K] +""" + +from __future__ import annotations + +import importlib +from typing import Any + +_BACKEND = "flashinfer.experimental.cake.cake_dense_gemm_det_train" + +# public name -> backend submodule that defines it +_EXPORTS: dict[str, str] = { + "ARCHES": "cake_jit", + "SUPPORTED_COMPUTE_CAPABILITIES": "cake_backend", + "is_cake_dense_gemm_det_train_supported": "cake_backend", + # K1: BF16 dense projection GEMM + "cake_dense_gemm_det_train_projection": "cake_backend", + "cake_dense_gemm_det_train_projection_forward": "cake_backend", + "cake_dense_gemm_det_train_projection_dgrad": "cake_backend", + "cake_dense_gemm_det_train_projection_wgrad": "cake_backend", + "cake_dense_gemm_det_train_projection_plan": "cake_backend", + "cake_dense_gemm_det_train_projection_prepare": "cake_backend", + "cake_dense_gemm_det_train_projection_wgrad_prepare": "cake_backend", + "CakeDenseGemmDetProjectionPlan": "cake_backend", + "CakeDenseGemmDetProjectionLaunch": "cake_backend", + # K2: FP32 router GEMM + "cake_dense_gemm_det_train_router": "cake_backend", + "cake_dense_gemm_det_train_router_forward": "cake_backend", + "cake_dense_gemm_det_train_router_dgrad": "cake_backend", + "cake_dense_gemm_det_train_router_wgrad": "cake_backend", + "cake_dense_gemm_det_train_router_plan": "cake_backend", + "cake_dense_gemm_det_train_router_prepare": "cake_backend", + "CakeDenseGemmDetRouterPlan": "cake_backend", + "CakeDenseGemmDetRouterLaunch": "cake_backend", + # the convenience path's cached plans and workspace arenas + "cake_dense_gemm_det_train_convenience_cache_clear": "cake_backend", +} + +__all__ = sorted(_EXPORTS) + + +def __getattr__(name: str) -> Any: + submodule = _EXPORTS.get(name) + if submodule is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(f"{_BACKEND}.{submodule}"), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/flashinfer/experimental/dense_projection_det_gemm/.gitignore b/flashinfer/experimental/cake/cake_dense_gemm_det_train/.gitignore similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/.gitignore rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/.gitignore diff --git a/flashinfer/experimental/dense_projection_det_gemm/README.md b/flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md similarity index 86% rename from flashinfer/experimental/dense_projection_det_gemm/README.md rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md index 9cb85750926..9c4018b78f0 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/README.md +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md @@ -1,48 +1,64 @@ -# Cake dense projection GEMM backend (SM100 / SM107) +# Cake dense GEMM training backend, deterministic (SM100 / SM103 / SM107) -Experimental backend for the GLM-5.2 training GEMMs: the BF16 dense projection -GEMM (K1 -- BF16 operands, FP32 accumulation, BF16 or FP32 output, strided -views, ragged token counts, batched head projections) and the FP32 router-gate -GEMM through split-BF16x3 tensor-core emulation with chunked FP32 promotion -(K2). Tracking issue: flashinfer-ai/flashinfer#5677. +Backend of `flashinfer.cake_dense_gemm_det_train`: the GLM-5.2 training GEMMs, the BF16 dense projection +GEMM (K1 -- BF16 operands, FP32 accumulation, BF16 or FP32 output, strided views, ragged token counts, +batched head projections) and the FP32 router-gate GEMM through split-BF16x3 tensor-core emulation with +chunked FP32 promotion (K2). Tracking issue: flashinfer-ai/flashinfer#5677. -Both the API and the backend are experimental: no compatibility guarantee. SM107 (Rubin R200) is the primary architecture; SM100 (B200 / GB200) is built from the same sources and, for five input-gradient rows, ships the knob-free eight-warp CTA form of the GEMM (one program, measured on GB200); SM103 (B300 / GB300) ships its own builds of the same programs and, for six input-gradient rows, the eight-warp CTA form in two programs (`cake_jit.KERNELS` lists the registered architectures). A device whose architecture the registry does not -list raises `NotImplementedError` naming the missing instance. The feature has -no core entry point yet (importing this package is -the explicit opt-in) and does not participate in automatic backend selection, -autotuning or trace-apply. Its libraries are loaded prebuilt from `aot/` when -the package ships them and JIT-built by FlashInfer from `build.json` otherwise -(see "Layout of this package"). +list raises `NotImplementedError` naming the missing instance; +`is_cake_dense_gemm_det_train_supported(device)` reports whether a device is a supported architecture. The +backend does not participate in automatic backend selection, autotuning or trace-apply. Its libraries are +loaded prebuilt from `aot/` when the package ships them and JIT-built by FlashInfer from `build.json` +otherwise (see "Layout of this package"). +## Names and mode -The package was formerly `flashinfer.experimental.dense_projection_gemm`; that import path and the former -entry-point names (`dense_projection_gemm`, `prepare_dense_projection_gemm`, `projection_forward`, -`router_fp32_gemm`, ...) remain as aliases of the canonical objects above. -## Entry points (`cake_backend.py`) +The public module is `flashinfer.cake_dense_gemm_det_train`. It imports this package on the first +attribute access and returns the objects of `cake_backend.py` (and `ARCHES` of `cake_jit.py`) unchanged, so +`flashinfer.cake_dense_gemm_det_train.X is cake_backend.X`. Every public name carries the mode word `det`: + +| kind | names | +|---|---| +| K1 entry points | `cake_dense_gemm_det_train_projection`, `cake_dense_gemm_det_train_projection_{forward,dgrad,wgrad}` | +| K1 plan / prepare | `cake_dense_gemm_det_train_projection_{plan,prepare,wgrad_prepare}` | +| K1 types | `CakeDenseGemmDetProjectionPlan`, `CakeDenseGemmDetProjectionLaunch` | +| K2 entry points | `cake_dense_gemm_det_train_router`, `cake_dense_gemm_det_train_router_{forward,dgrad,wgrad}` | +| K2 plan / prepare | `cake_dense_gemm_det_train_router_{plan,prepare}` | +| K2 types | `CakeDenseGemmDetRouterPlan`, `CakeDenseGemmDetRouterLaunch` | +| cache, support, constants | `cake_dense_gemm_det_train_convenience_cache_clear`, `is_cake_dense_gemm_det_train_supported`, `ARCHES`, `SUPPORTED_COMPUTE_CAPABILITIES` | + +The module selects the mode: no entry point takes a mode argument. Every program of this package runs in a +fixed order, so the same inputs, hardware and plan give bitwise the same outputs on every run. The entry +points do not read `torch.use_deterministic_algorithms`: they never raise or warn under it (strict or +`warn_only`) and run the same programs with or without it. The non-deterministic split-K programs are a +separate module, `flashinfer.cake_dense_gemm_nondet_train`. `CONVENIENCE_CACHE_SIZE` is a setting of +`cake_backend` (assign it there, not on the public module). The backend's other helpers keep their names. + +## Entry points (`flashinfer.cake_dense_gemm_det_train`, defined in `cake_backend.py`) ```python -from flashinfer.experimental.dense_projection_det_gemm import cake_backend as gemm +import flashinfer.cake_dense_gemm_det_train as gemm # K1: strided-view GEMM out[l, m, n] = sum_k A[l, m, k] B[l, k, n] -gemm.dense_projection_det_gemm(A, B, out, transposed_out=False) # convenience path (cached plan) -gemm.projection_det_forward(X, W, out=None) # out[T, N] = X[T, K] @ W[N, K].T -gemm.projection_det_dgrad(G, W, out=None, out_dtype=None) # out[T, K] = G[T, N] @ W[N, K] -gemm.projection_det_wgrad(G, X, out=None, out_dtype=None) # out[N, K] = G[T, N].T @ X[T, K] -prepared = gemm.prepare_dense_projection_det_gemm(A, B, out, transposed_out=False) -prepared = gemm.prepare_projection_det_wgrad(G, X, out) +gemm.cake_dense_gemm_det_train_projection(A, B, out, transposed_out=False) # convenience path (cached plan) +gemm.cake_dense_gemm_det_train_projection_forward(X, W, out=None) # out[T, N] = X[T, K] @ W[N, K].T +gemm.cake_dense_gemm_det_train_projection_dgrad(G, W, out=None, out_dtype=None) # out[T, K] = G[T, N] @ W[N, K] +gemm.cake_dense_gemm_det_train_projection_wgrad(G, X, out=None, out_dtype=None) # out[N, K] = G[T, N].T @ X[T, K] +prepared = gemm.cake_dense_gemm_det_train_projection_prepare(A, B, out, transposed_out=False) +prepared = gemm.cake_dense_gemm_det_train_projection_wgrad_prepare(G, X, out) prepared.launch() # allocation-free; prepared.template / .grid / .module_name / .plan # K2: FP32 router GEMM through split-BF16x3 emulation -gemm.router_fp32_det_gemm(A, B, out, splits=None) # convenience path (cached plan) -gemm.router_det_forward(X, W, out=None, splits=4) # out[T, N] = X[T, K] @ W[N, K].T -gemm.router_det_dgrad(G, W, out=None) # out[T, K] = G[T, N] @ W[N, K] -gemm.router_det_wgrad(G, X, out=None, splits=6) # out[N, K] = G[T, N].T @ X[T, K] -prepared = gemm.prepare_router_fp32_det_gemm(A, B, out, splits=None) +gemm.cake_dense_gemm_det_train_router(A, B, out, splits=None) # convenience path (cached plan) +gemm.cake_dense_gemm_det_train_router_forward(X, W, out=None, splits=4) # out[T, N] = X[T, K] @ W[N, K].T +gemm.cake_dense_gemm_det_train_router_dgrad(G, W, out=None) # out[T, K] = G[T, N] @ W[N, K] +gemm.cake_dense_gemm_det_train_router_wgrad(G, X, out=None, splits=6) # out[N, K] = G[T, N].T @ X[T, K] +prepared = gemm.cake_dense_gemm_det_train_router_prepare(A, B, out, splits=None) prepared.launch() # re-splits B, launches, reduces the split-K partials; .splits ``` @@ -51,17 +67,17 @@ prepared.launch() # re-splits B, launches, reduces the split-K partials; .spli Both paths run the same planner once per binding and bind the same generated module; they differ in who owns the plan and the workspace. -| | prepared path (`prepare_*` + `launch()`) | convenience path (`dense_projection_det_gemm`, `projection_*`, `router_*`) | +| | prepared path (`prepare_*` + `launch()`) | convenience path (`cake_dense_gemm_det_train_projection`, `projection_*`, `router_*`) | |---|---|---| | planning | once, in `prepare_*` | once per cache key, on the first call | | workspace (K1: stream-K slabs + slice counters; K2: BF16x3 stack, the torch-chain split's fp32 scratch, split-K partials) | owned by the prepared object | carved out of the workspace arena of the (device, stream): one growable buffer per kind shared by every cached entry of that device and stream, grown to the largest requirement seen, never per entry | | per call | bind nothing: launch the bound views | rebind the caller's views (view objects only), resolve the workspace views of the arena (memoised until the arena grows), launch; `out` is allocated only when the caller passes none | | device queries | none after `prepare_*` | none after the first call | -| cache | none | per process, per kind (K1 / K2), `CONVENIENCE_CACHE_SIZE` = 64 entries, least recently used first out; the arena of a (device, stream) is released with its last entry and by `clear_convenience_caches()` | +| cache | none | per process, per kind (K1 / K2), `CONVENIENCE_CACHE_SIZE` = 64 entries, least recently used first out; the arena of a (device, stream) is released with its last entry and by `cake_dense_gemm_det_train_convenience_cache_clear()` | | memory held | the prepared object's workspace | bounded by the largest geometry seen per workspace kind, however many token counts the cache has planned (a ragged T adds entries, not workspace) | | CUDA Graphs | capture `launch()` | do not capture (a later growth replaces, an eviction frees, the workspace a graph replays) | -The orientation of `projection_det_wgrad` / `prepare_projection_det_wgrad` follows the +The orientation of `cake_dense_gemm_det_train_projection_wgrad` / `cake_dense_gemm_det_train_projection_wgrad_prepare` follows the row's measured rule when the rule table keys the row in exactly one orientation (the swapped `X.T @ G` form with the transposed store, or the direct `G.T @ X`), else the launcher's padded-tile-work rule (N < 256 and N = 576 swap); @@ -86,7 +102,7 @@ the kernel resets every counter it uses (the last arriver of a tail tile's slice writes zero after the fixup), so the shared counter buffer is zero again when each launch completes -- the same contract the Cake launcher's per-device counters rely on. Threads: planning and the arena bookkeeping (creation, -growth, reference counts, insertion, eviction, `clear_convenience_caches()`) +growth, reference counts, insertion, eviction, `cake_dense_gemm_det_train_convenience_cache_clear()`) are serialised by a module lock (a miss re-checks the cache under it, so a key is planned once; a growth keeps the larger buffer); the K2 hit path holds its arena's lock from the carve of the workspace views through the split-K @@ -126,7 +142,7 @@ device index (`device_facts`). (`A` K-major, `B = W.T` K-major), input gradient `G @ W` (`B` MN-major), weight gradient `G.T @ X` (both MN-major, `K = T` ragged). A weight gradient with fewer than 256 output rows runs the swapped `X.T @ G` with - the transposed store (`projection_det_wgrad` / `prepare_projection_det_wgrad`). + the transposed store (`cake_dense_gemm_det_train_projection_wgrad` / `cake_dense_gemm_det_train_projection_wgrad_prepare`). * Numerics: BF16 x BF16 products exact in FP32, tensor-core FP32 accumulation, one rounding at the store. BF16 outputs match an FP64 reference to `atol = rtol = 1e-2` elementwise; FP32 outputs carry the @@ -143,7 +159,7 @@ stream-K split, raster group, wave working set) and the L2 size (hint gate); `prepare_*` and the convenience path read both once per device index (`device_facts`: `device_sm_count`, `device_l2_bytes` = `torch.cuda.get_device_properties(...).L2_cache_size`, falling back to -`cuDeviceGetAttribute`), `plan_dense_projection_det_gemm` takes them as inputs: +`cuDeviceGetAttribute`), `cake_dense_gemm_det_train_projection_plan` takes them as inputs: | choice | rule | |---|---| @@ -164,8 +180,8 @@ stream-K split, raster group, wave working set) and the L2 size (hint gate); | grid | `(num_full + sk_units) * 2` CTAs: one cluster of two CTAs per `2 * CTA_ROWS` x `BLOCK_N` output tile plus the stream-K units; under `cta1` (round 19) `num_full` CTAs, one per `CTA_ROWS` x `BLOCK_N` tile, launched with the module's recorded cluster dims (1, 1, 1) | | stream-K | `stream_k_plan(pair_tiles, k_blocks, sm_count // 2, sk="auto")`: only when the whole problem is one partial wave of CTA pairs whose halves fit the pairs (`pair_tiles <= pairs`, `2 * tail <= pairs`) and `k_blocks >= 16`, the tiles' linearised K steps are cut into units of `iters_per_unit = ceil(k_blocks / 2)` (two K-aligned halves per tile on an even `k_blocks`), `sk_units = ceil(pair_tiles * k_blocks / iters_per_unit)` so that every unit holds a K step, with an in-kernel deterministic fixup; otherwise every work item is a whole tile (multi-wave tails stay data-parallel). `sk=True` is the linearised-units policy (up to `pairs` units of >= 8 K steps), `sk="tiles"` the whole-tile diagnostic, `sk=False` off. A measured `sk_parts = p` rule, or the caller knob `sk_parts`, instead splits every tail tile `p` ways when `p * tail` units fit the pairs and each keeps >= 8 K steps (`sk_parts_plan`). The host asserts `tail_tiles * 16 <= 4096` slice counters, and `check_stream_k_plan` refuses a plan outside the kernel's launch contract: a unit without a K step on the `pd` 2 / `sk_slab` 2 / `stt` programs (they hand off the work in an item's last segment, so such a unit never finishes), synchronised collectors outside `[1, tail_tiles]` (`[1, ...)` on the other `sk_sync` programs), or a linearised plan with `2 * iters_per_unit < k_blocks` (more than `WS_PARTS` units on a tail tile) on a program that keys its workspace by tile (`tile_keyed_workspace`; its binding refuses the same plan) | -`plan_dense_projection_det_gemm(A, B, out, sm_count=..., l2_bytes=...)` exposes -the plan (`DetGemmPlan`, including `wave_working_set_bytes`) without a device; +`cake_dense_gemm_det_train_projection_plan(A, B, out, sm_count=..., l2_bytes=...)` exposes +the plan (`CakeDenseGemmDetProjectionPlan`, including `wave_working_set_bytes`) without a device; the prepared launch resolves the template through `cake_jit.KERNELS[arch][template]`, owns its stream-K partial slabs (`max(sk_units + tail_tiles, WS_PARTS * tail_tiles) * 2 * CTA_ROWS * BLOCK_N` @@ -181,7 +197,7 @@ launcher's. ### The swapped small-token forward -The convenience path (`dense_projection_det_gemm`) runs a batched forward row +The convenience path (`cake_dense_gemm_det_train_projection`) runs a batched forward row that the planner would swap (`swap_small_m`: `L > 1`, `M <= 256`, `N >= 2 M`) on a program that puts the weight rows in the MMA's M and the tokens in its N (the `swap_ab_` templates), when `A` and `B` are K-major, the output is @@ -281,15 +297,16 @@ registry does not list. CUDA Graphs: prepare outside capture and replay the prepared launch (`prepare_*` once, `launch` many times). The convenience entry points -(`dense_projection_det_gemm`, `projection_*`, `router_*`) reuse a cached planned +(`cake_dense_gemm_det_train_projection`, `projection_*`, `router_*`) reuse a cached planned launch per (device, stream, geometry) key and the shared workspace arena of its (device, stream); capture a prepared launch object rather than the convenience entry points (see "Host paths"). -Tests: `tests/experimental/test_cake_dense_projection_det_gemm.py` (host-side +Tests: `tests/experimental/test_cake_dense_gemm_det_train.py` and +`tests/experimental/test_cake_dense_gemm_det_train_names.py` (host-side planning tests run everywhere; the GPU tests skip without a registered generated instance -- a device of an architecture `cake_jit.KERNELS` lists). Benchmark: -`benchmarks/bench_cake_dense_projection_det_gemm.py` (the 11 GLM-5.2 projection +`benchmarks/bench_cake_dense_gemm_det_train.py` (the 11 GLM-5.2 projection rows x fwd / dgrad / wgrad x T in {16231, 16172}, the router rows and the two batched MLA rows, paired against `torch.matmul`; `--host-overhead` adds the host time per call of the prepared launch and of the convenience entry point). diff --git a/flashinfer/experimental/dense_projection_det_gemm/__init__.py b/flashinfer/experimental/cake/cake_dense_gemm_det_train/__init__.py similarity index 68% rename from flashinfer/experimental/dense_projection_det_gemm/__init__.py rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/__init__.py index 4344abcecc7..a50bca357f1 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/__init__.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/__init__.py @@ -14,11 +14,11 @@ limitations under the License. """ -# Experimental Cake backend: the GLM-5.2 dense projection GEMMs for training +# Cake backend: the GLM-5.2 dense projection GEMMs for training # (BF16 operands, FP32 accumulation, BF16 or FP32 output, strided views, ragged # token counts, batched head projections) and the FP32 router-gate GEMM through # split-BF16x3 tensor-core emulation, on SM100 / SM103 / SM107 -# (flashinfer-ai/flashinfer#5677). Host planning, argument-plan binding, the -# prepared allocation-free launches, the eager entry points, JIT registration -# and the generated sources live in this package; there is no core entry point -# yet (importing this package is the explicit opt-in). +# (flashinfer-ai/flashinfer#5677), deterministic. Host planning, argument-plan +# binding, the prepared allocation-free launches, the eager entry points, JIT +# registration and the generated sources live in this package; the public module +# is flashinfer.cake_dense_gemm_det_train. diff --git a/flashinfer/experimental/dense_projection_det_gemm/build.json b/flashinfer/experimental/cake/cake_dense_gemm_det_train/build.json similarity index 99% rename from flashinfer/experimental/dense_projection_det_gemm/build.json rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/build.json index beadb0d572f..9765ecd79e9 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/build.json +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/build.json @@ -6304,6 +6304,6 @@ ] } ], - "package": "flashinfer.experimental.dense_projection_det_gemm", + "package": "flashinfer.experimental.cake.cake_dense_gemm_det_train", "schema": "cake-build/1" } diff --git a/flashinfer/experimental/dense_projection_det_gemm/cake_backend.py b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py similarity index 96% rename from flashinfer/experimental/dense_projection_det_gemm/cake_backend.py rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py index e2ce7d1c6be..eab4407f7e6 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/cake_backend.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py @@ -39,7 +39,7 @@ input gradient ``G[T, N] @ W[N, K]`` (``B = W`` MN-major), weight gradient ``G[T, N].T @ X[T, K]`` (both MN-major, ``K = T`` ragged); a weight gradient with fewer than 256 output rows runs the swapped ``X.T @ G`` with the -transposed store (:func:`projection_det_wgrad`). +transposed store (:func:`cake_dense_gemm_det_train_projection_wgrad`). The host selects one traced kernel instance per (A layout, B layout, tile width, tile height, output kind, epilogue path, raster group, TMA L2 eviction @@ -56,7 +56,7 @@ The convenience path runs a batched forward of at most 256 tokens on the swapped program instead (the weight rows in the MMA's M, at most two waves of single-CTA tiles over the caller's strided views; :func:`swap_ab_template`), -which takes ``enable_pdl`` per call (:func:`dense_projection_det_gemm`). +which takes ``enable_pdl`` per call (:func:`cake_dense_gemm_det_train_projection`). K2 -- FP32 router GEMM through split-BF16x3 emulation ---------------------------------------------------- @@ -76,17 +76,17 @@ Two host paths share one planner and one binding: -* the prepared path -- :func:`prepare_dense_projection_det_gemm` / - :func:`prepare_router_fp32_det_gemm` plan once and own their workspace; the +* the prepared path -- :func:`cake_dense_gemm_det_train_projection_prepare` / + :func:`cake_dense_gemm_det_train_router_prepare` plan once and own their workspace; the returned launch replays with no CUDA allocation and no host synchronization for new values written into the bound tensors. Every ``launch()`` is launch-only: the generated binding encodes the operand tensor maps by value from the bound views on each call (no descriptor workspace), so a prepared launch can be captured into a CUDA Graph from its first call; -* the convenience path -- :func:`dense_projection_det_gemm`, - :func:`projection_det_forward` / :func:`projection_det_dgrad` / :func:`projection_det_wgrad`, - :func:`router_fp32_det_gemm`, :func:`router_det_forward` / :func:`router_det_dgrad` / - :func:`router_det_wgrad` -- keeps a bounded per-process cache of planned launches +* the convenience path -- :func:`cake_dense_gemm_det_train_projection`, + :func:`cake_dense_gemm_det_train_projection_forward` / :func:`cake_dense_gemm_det_train_projection_dgrad` / :func:`cake_dense_gemm_det_train_projection_wgrad`, + :func:`cake_dense_gemm_det_train_router`, :func:`cake_dense_gemm_det_train_router_forward` / :func:`cake_dense_gemm_det_train_router_dgrad` / + :func:`cake_dense_gemm_det_train_router_wgrad` -- keeps a bounded per-process cache of planned launches keyed by the device, the stream and the geometry of the views (shapes, strides, dtypes, storage offsets, alignment, flags); a call that hits the cache rebinds the caller's views and launches without planning, device @@ -96,7 +96,7 @@ owning a workspace each, so the memory the convenience path holds is bounded by the largest geometry seen, not by the number of geometries (a ragged token count). See ``CONVENIENCE_CACHE_SIZE`` and - :func:`clear_convenience_caches`. + :func:`cake_dense_gemm_det_train_convenience_cache_clear`. The device facts a plan depends on (architecture, SM count, L2 size) are read once per device index (:func:`device_facts`). @@ -131,7 +131,7 @@ # --------------------------------------------------------------------------- # K1 tile configuration (mirror of the Cake kernel module's constants; the -# Cake kernel host is ``dense_projection_det_gemm.py`` of the Cake kernel library, +# Cake kernel host is the K1 kernel module of the Cake kernel library, # cited by line below) # --------------------------------------------------------------------------- @@ -217,7 +217,7 @@ ROUTER_BLOCK_M = 128 ROUTER_BLOCK_K = 32 # fp32 elements split per stage ROUTER_BLOCK_N = 128 # output columns per CTA pair -# Production split-K counts per instance layout (router_det_forward / router_det_dgrad / router_det_wgrad). +# Production split-K counts per instance layout (cake_dense_gemm_det_train_router_forward / cake_dense_gemm_det_train_router_dgrad / cake_dense_gemm_det_train_router_wgrad). ROUTER_SPLITS = {"kk": 4, "kn": 1, "nn_t": 6} @@ -264,7 +264,7 @@ def _driver_l2_bytes(index: int) -> int: ) if err != cuda.CUresult.CUDA_SUCCESS or not l2: raise RuntimeError( - f"dense_projection_det_gemm: cannot resolve the L2 size of cuda:{index} ({err})" + f"cake_dense_gemm_det_train_projection: cannot resolve the L2 size of cuda:{index} ({err})" ) return int(l2) @@ -308,6 +308,13 @@ def require_arch(device: torch.device) -> str: return arch +def is_cake_dense_gemm_det_train_supported(device: torch.device) -> bool: + """Whether ``device`` is a CUDA device of compute capability 10.0, 10.3 or 10.7 (``sm_100a`` / ``sm_103a`` / + ``sm_107a``). Shape-level support is checked per call: an unsupported view, layout or size raises + ``ValueError``, and a geometry without a registered program raises ``NotImplementedError``.""" + return arch_for(device) is not None + + def generated_program_available( device: Optional[torch.device] = None, template: Optional[str] = None ) -> bool: @@ -354,7 +361,7 @@ def exact_waves(pair_tiles: int, units: int) -> bool: ``cta1``) in whole waves - ``((m_tiles // 2) * n_tiles * L) % sm_pairs == 0`` in the pair form. The BLOCK_N 224 weight-gradient rules were chosen for this fit on the 76-pair measuring devices (152 SMs: the 4096 x 4096 row's 152 tall tiles are two waves, the 4096 x 1024 row's 76 tiles one); the exact-wave gate of - :func:`plan_dense_projection_det_gemm` reads it for the running device.""" + :func:`cake_dense_gemm_det_train_projection_plan` reads it for the running device.""" return pair_tiles > 0 and pair_tiles % units == 0 @@ -1689,7 +1696,7 @@ def _measured_row_rule( ragged-M (N, K, None) rule, else the ragged-K (N, None, M) rule, else the parent architecture's entry (``RULE_PARENT``: sm_103a reads the sm_100a table for the rows its lanes did not measure), else empty (the lookup of Cake ``row_rule``; FlashInfer-only split so the registry fallback can drop the MEASURED rule alone, see - ``plan_dense_projection_det_gemm``).""" + ``cake_dense_gemm_det_train_projection_plan``).""" ident = (arch, bool(a_mn), bool(b_mn), bool(out_f32), bool(out_t), bool(batched)) for tail in ( (int(N), int(K), int(M)), @@ -1936,7 +1943,7 @@ def stream_k_plan( return pair_tiles - tail, tail, units, iters_per_unit -def check_stream_k_plan(plan: "DetGemmPlan", arch: Optional[str] = None) -> None: +def check_stream_k_plan(plan: "CakeDenseGemmDetProjectionPlan", arch: Optional[str] = None) -> None: """Refuse a stream-K plan outside the program's launch contract (the Cake kernel's ``requires``) before it launches. The programs that hand off the work inside an item's last segment (``pd`` 2, ``sk_slab`` 2, ``stt``) never finish an item that holds no K step: the linearised plan needs ``iters_per_unit >= 1`` and @@ -1952,7 +1959,7 @@ def check_stream_k_plan(plan: "DetGemmPlan", arch: Optional[str] = None) -> None collectors = plan.sk_collectors if collectors < 1 or (handoff and collectors > plan.tail_tiles): raise ValueError( - f"dense_projection_det_gemm: {plan.template} needs 1 <= collectors" + f"cake_dense_gemm_det_train_projection: {plan.template} needs 1 <= collectors" f"{' <= tail_tiles' if handoff else ''} (got {collectors} collectors, {plan.tail_tiles} tail tiles)" ) elif handoff and not plan.sk_exact: @@ -1961,7 +1968,7 @@ def check_stream_k_plan(plan: "DetGemmPlan", arch: Optional[str] = None) -> None or (plan.sk_units - 1) * plan.iters_per_unit >= plan.sk_iters ): raise ValueError( - f"dense_projection_det_gemm: {plan.template} needs every stream-K unit to hold a K step (got " + f"cake_dense_gemm_det_train_projection: {plan.template} needs every stream-K unit to hold a K step (got " f"{plan.sk_units} units of {plan.iters_per_unit} over {plan.sk_iters} K steps)" ) if ( @@ -1971,7 +1978,7 @@ def check_stream_k_plan(plan: "DetGemmPlan", arch: Optional[str] = None) -> None and tile_keyed_workspace(arch, plan.template) ): raise ValueError( - f"dense_projection_det_gemm: {plan.template} keys its stream-K workspace by tail tile and needs " + f"cake_dense_gemm_det_train_projection: {plan.template} keys its stream-K workspace by tail tile and needs " f"2 * iters_per_unit >= k_blocks, at most {WS_PARTS} units on a tile (got {plan.sk_units} units of " f"{plan.iters_per_unit} K steps over {plan.tail_tiles} tail tiles of {plan.k_blocks} K steps)" ) @@ -2053,7 +2060,7 @@ def operand_view( @dataclass(frozen=True) -class DetGemmPlan: +class CakeDenseGemmDetProjectionPlan: """The host plan of one dense projection GEMM launch (device-independent except ``sm_pairs``: the concurrent work items of the device - CTA pairs, or SMs under the single-CTA form ``cta1``).""" @@ -2085,7 +2092,7 @@ class DetGemmPlan: sk_units: int iters_per_unit: int template: str - # FlashInfer-only (see ``plan_dense_projection_det_gemm``): the batched small-M swap was undone / the row's + # FlashInfer-only (see ``cake_dense_gemm_det_train_projection_plan``): the batched small-M swap was undone / the row's # measured rule was dropped because the planned instance is not a generated program of this architecture # (the Cake host JIT-compiles it) swap_fallback: bool = False @@ -2240,7 +2247,7 @@ def tma_out(self) -> bool: return self.epi == "tma" or self.stt -def plan_dense_projection_det_gemm( +def cake_dense_gemm_det_train_projection_plan( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, @@ -2289,7 +2296,7 @@ def plan_dense_projection_det_gemm( _fallback: bool = True, _allow_swap: bool = True, _use_rules: bool = True, -) -> tuple[DetGemmPlan, torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[CakeDenseGemmDetProjectionPlan, torch.Tensor, torch.Tensor, torch.Tensor]: """Validate the views exactly as the Cake launcher does and plan the launch for a device with ``sm_count`` SMs and an L2 of ``l2_bytes`` (the two device facts the plan depends on: stream-K split and raster group through the CTA-pair count, the hint working-set gate @@ -2346,7 +2353,7 @@ def plan_dense_projection_det_gemm( A caller-forced ``form`` is admitted or refused by ``instance_key`` and refused on a plan with stream-K units. It changes no plan arithmetic; the launch adds the raster's four fast divisors (``fd_tiles_l`` / ``fd_tiles_g`` / ``fd_group`` / ``fd_last``), which the ten-warp programs' argument plans do not name. - [Cake ``dense_projection_det_gemm`` L1234-L1358] + [Cake K1 kernel host L1234-L1358] One FlashInfer-only deviation: the Cake host applies ``swap_small_m`` unconditionally because it compiles the swapped instance on demand, while this package can only launch the generated programs registered for @@ -2411,24 +2418,24 @@ def plan_dense_projection_det_gemm( arch=arch, ) if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: - raise ValueError("dense_projection_det_gemm: A and B must be bf16") + raise ValueError("cake_dense_gemm_det_train_projection: A and B must be bf16") if out.dtype not in (torch.bfloat16, torch.float32): - raise ValueError("dense_projection_det_gemm: out must be bf16 or fp32") + raise ValueError("cake_dense_gemm_det_train_projection: out must be bf16 or fp32") A3, B3, O3 = as_batched(A, "A"), as_batched(B, "B"), as_batched(out, "out") L, M, K = (int(s) for s in A3.shape) if tuple(B3.shape) != (L, K, int(B3.shape[2])): raise ValueError( - f"dense_projection_det_gemm: B must be [L={L}, K={K}, N], got {tuple(B3.shape)}" + f"cake_dense_gemm_det_train_projection: B must be [L={L}, K={K}, N], got {tuple(B3.shape)}" ) N = int(B3.shape[2]) expect = (L, N, M) if transposed_out else (L, M, N) if tuple(O3.shape) != expect: raise ValueError( - f"dense_projection_det_gemm: out must be {expect}, got {tuple(O3.shape)}" + f"cake_dense_gemm_det_train_projection: out must be {expect}, got {tuple(O3.shape)}" ) if N % 8 != 0 or M < 1 or K < 1 or N < 8: raise ValueError( - f"dense_projection_det_gemm: N must be a positive multiple of 8 and M, K >= 1 (M={M}, N={N}, K={K})" + f"cake_dense_gemm_det_train_projection: N must be a positive multiple of 8 and M, K >= 1 (M={M}, N={N}, K={K})" ) if ( O3.stride(2) != 1 @@ -2437,7 +2444,7 @@ def plan_dense_projection_det_gemm( or out.data_ptr() % 16 != 0 ): raise ValueError( - "dense_projection_det_gemm: out needs unit inner stride, 16-byte aligned rows / batches " + "cake_dense_gemm_det_train_projection: out needs unit inner stride, 16-byte aligned rows / batches " f"and base; strides {tuple(O3.stride())}" ) swap_eligible = swap_small_m(L, M, N, transposed_out) @@ -2521,7 +2528,7 @@ def plan_dense_projection_det_gemm( # the single-CTA form carries no tail policy (its one-K-step rows never split) and no pair-only knob; a caller # forcing one with cta1 asked for two incompatible forms [Cake launcher] raise ValueError( - f"dense_projection_det_gemm: cta1 excludes sk / sk_exact / sk_sync / htail / sk_parts / sk_max_units / batch_group " + f"cake_dense_gemm_det_train_projection: cta1 excludes sk / sk_exact / sk_sync / htail / sk_parts / sk_max_units / batch_group " f"(got sk={sk}, sk_exact={sk_exact}, sk_sync={sk_sync}, htail={htail}, sk_parts={sk_parts}, " f"sk_max_units={sk_max_units}, batch_group={batch_group})" ) @@ -2636,7 +2643,7 @@ def plan_dense_projection_det_gemm( form = str(form or "") if form not in ("", *W8_FORMS): raise ValueError( - f"dense_projection_det_gemm: form must be one of {W8_FORMS} (the eight-warp CTA form) or '', got {form!r}" + f"cake_dense_gemm_det_train_projection: form must be one of {W8_FORMS} (the eight-warp CTA form) or '', got {form!r}" ) if htail is None: htail = rule.get("htail", False) @@ -2671,7 +2678,7 @@ def plan_dense_projection_det_gemm( batch_group = int(batch_group) if L > 1 else 0 if batch_group < 0 or (batch_group > 1 and L % batch_group): raise ValueError( - f"dense_projection_det_gemm: batch_group must be 0, 1 or a divisor of the batch size {L}, got {batch_group}" + f"cake_dense_gemm_det_train_projection: batch_group must be 0, 1 or a divisor of the batch size {L}, got {batch_group}" ) if store_ef is None: store_ef = rule.get("store_ef", False) @@ -2783,7 +2790,7 @@ def plan_dense_projection_det_gemm( ) if tail_tiles * 16 > SK_DUMMY_BASE: raise ValueError( - f"dense_projection_det_gemm: {tail_tiles} stream-K tail tiles exceed the {SK_DUMMY_BASE // 16} slice-counter budget" + f"cake_dense_gemm_det_train_projection: {tail_tiles} stream-K tail tiles exceed the {SK_DUMMY_BASE // 16} slice-counter budget" ) out_f32 = out.dtype == torch.float32 mode = epi_mode(out_f32, transposed_out, K, epi, block_n, cta_rows) @@ -2855,7 +2862,7 @@ def plan_dense_projection_det_gemm( promo = str(promo) if promo not in L2_PROMOS: raise ValueError( - f"dense_projection_det_gemm: promo must be one of {L2_PROMOS}, got {promo!r}" + f"cake_dense_gemm_det_train_projection: promo must be one of {L2_PROMOS}, got {promo!r}" ) if group_m is None: group_m = rule.get( @@ -2864,7 +2871,7 @@ def plan_dense_projection_det_gemm( group_m = int(group_m) if group_m < 1 or (cg == 2 and not cgrp and (group_m < 2 or group_m % 2)): raise ValueError( - f"dense_projection_det_gemm: group_m must be an even number >= 2 (CTA pairs are adjacent row tiles; any count >= 1 under cta1, or under cgrp where it counts column tiles), got {group_m}" + f"cake_dense_gemm_det_train_projection: group_m must be an even number >= 2 (CTA pairs are adjacent row tiles; any count >= 1 under cta1, or under cgrp where it counts column tiles), got {group_m}" ) if hints is None: # round 21 (Cake W3): the hint gate over the column raster's wave working set under cgrp [Cake launcher] @@ -2895,7 +2902,7 @@ def plan_dense_projection_det_gemm( mg = False elif mg and not whole_tiles: raise ValueError( - f"dense_projection_det_gemm: mg (whole-tile tile resolve) needs a whole-tile plan under 2**21 cluster tiles " + f"cake_dense_gemm_det_train_projection: mg (whole-tile tile resolve) needs a whole-tile plan under 2**21 cluster tiles " f"(got {sk_units} stream-K units over {num_full + sk_units} cluster tiles)" ) # The drain guard of ``sw``: the ``_sw`` programs pace the bf16 TMA-store pd-2 drain and the kernel refuses the @@ -2939,7 +2946,7 @@ def plan_dense_projection_det_gemm( store_hint, store_ef, quad_store, f32_v8 = "none", False, False, False elif form and not whole_tiles: raise ValueError( - f"dense_projection_det_gemm: form (the eight-warp CTA form) needs a whole-tile kn bf16 plan on 128 x 256 pair " + f"cake_dense_gemm_det_train_projection: form (the eight-warp CTA form) needs a whole-tile kn bf16 plan on 128 x 256 pair " f"tiles under 2**21 cluster tiles (got {sk_units} stream-K units over {num_full + sk_units} cluster tiles)" ) key = instance_key( @@ -2982,7 +2989,7 @@ def plan_dense_projection_det_gemm( sw=sw, form=form, ) - plan = DetGemmPlan( + plan = CakeDenseGemmDetProjectionPlan( L=L, M=M, N=N, @@ -3046,7 +3053,7 @@ def plan_dense_projection_det_gemm( continue # the plan just made if (not allow_swap and not swapped) or (not use_rules and not rule): continue # dropping a swap not taken / an empty rule changes nothing - candidate = plan_dense_projection_det_gemm( + candidate = cake_dense_gemm_det_train_projection_plan( A, B, out, @@ -3066,7 +3073,7 @@ def plan_dense_projection_det_gemm( tried: set[tuple] = set() for allow_swap in (True, False) if swapped else (False,): try: - base = plan_dense_projection_det_gemm( + base = cake_dense_gemm_det_train_projection_plan( A, B, out, @@ -3092,7 +3099,7 @@ def plan_dense_projection_det_gemm( continue tried.add(frozen) try: - candidate = plan_dense_projection_det_gemm( + candidate = cake_dense_gemm_det_train_projection_plan( A, B, out, @@ -3113,7 +3120,7 @@ def plan_dense_projection_det_gemm( return plan, a_desc, b_desc, O3 -# Knob variants the registry fallback tries, nearest first (see ``plan_dense_projection_det_gemm``): the plain TMA L2 +# Knob variants the registry fallback tries, nearest first (see ``cake_dense_gemm_det_train_projection_plan``): the plain TMA L2 # policy (the same tiles and plan; only the eviction hints of the loads differ), the single-CTA 128-row form, the # 64-row pair family, the tall tile, the register epilogue (the canonical-T templates), then BLOCK_N = 128 with the same # forms. The TMA L2 promotion and the raster group width are launch parameters since round 11: they select no program @@ -3167,7 +3174,7 @@ def plan_dense_projection_det_gemm( "sk_first", "lw", ) -# The caller knobs the eight-warp CTA form yields to (``plan_dense_projection_det_gemm``): a value a caller forces for any +# The caller knobs the eight-warp CTA form yields to (``cake_dense_gemm_det_train_projection_plan``): a value a caller forces for any # of them shapes the pipeline, tile, epilogue, tail or raster the form cannot carry, so the row's ten-warp program runs # the plan asked for. Every name is a planner keyword whose default is None (a default call never yields). _W8_FIXED = ( @@ -3576,10 +3583,10 @@ def flat_alias(t: torch.Tensor) -> torch.Tensor: # (device index, raw stream) -- one growable buffer per kind shared by every entry of that device and stream # (:class:`_Arena`) -- so the memory held is bounded by the largest requirement seen per kind, however many # geometries (token counts) the cache has planned; an arena is released when its last entry is evicted and by -# :func:`clear_convenience_caches`. The key holds the raw stream, so launches on different streams never share +# :func:`cake_dense_gemm_det_train_convenience_cache_clear`. The key holds the raw stream, so launches on different streams never share # a workspace, and launches of one device and stream -- the only users of its arena -- are ordered by that # stream. Threads: planning and the arena bookkeeping (creation, growth, reference counts, insertion, eviction, -# ``clear_convenience_caches``) are serialised by the module lock ``_PLAN_LOCK`` (a miss re-checks the cache +# ``cake_dense_gemm_det_train_convenience_cache_clear``) are serialised by the module lock ``_PLAN_LOCK`` (a miss re-checks the cache # under it, so a key is planned once); the K2 hit path holds its arena's lock from the carve of the workspace # views through the split-K reduction, so two threads issuing router launches of different geometries on one # stream cannot interleave their split / kernel / reduction on the shared BF16x3 stack and partials; the K1 hit @@ -3690,11 +3697,11 @@ def _same_cuda_device( "k1_ws": ( torch.float32, False, - ), # stream-K partial slabs (``DetGemmPlan.ws_f32_elems``) + ), # stream-K partial slabs (``CakeDenseGemmDetProjectionPlan.ws_f32_elems``) "k1_counters": ( torch.uint32, True, - ), # stream-K slice counters (``DetGemmPlan.counters_alloc_u32``) + ), # stream-K slice counters (``CakeDenseGemmDetProjectionPlan.counters_alloc_u32``) "k2_parts": ( torch.bfloat16, False, @@ -3804,7 +3811,7 @@ def carve(self, arena: _Arena) -> dict[str, torch.Tensor]: ... _ARENAS: dict[tuple[int, int], _Arena] = {} -# Serialises planning and every mutation of the arenas and the caches (``_plan_entry``, ``clear_convenience_caches``); +# Serialises planning and every mutation of the arenas and the caches (``_plan_entry``, ``cake_dense_gemm_det_train_convenience_cache_clear``); # the hit paths (``_cache_get``, ``_entry_views``) do not take it. _PLAN_LOCK = threading.Lock() @@ -3910,7 +3917,7 @@ def _entry_views(entry: _CachedLaunch) -> dict[str, torch.Tensor]: @dataclass -class PreparedDetGemm: +class CakeDenseGemmDetProjectionLaunch: """One prepared dense projection GEMM launch. ``launch()`` writes ``out`` and returns it; it allocates nothing and never synchronizes. @@ -3920,7 +3927,7 @@ class PreparedDetGemm: module_name: str arch: str - plan: DetGemmPlan + plan: CakeDenseGemmDetProjectionPlan out: torch.Tensor tensors: dict[str, torch.Tensor] = field(repr=False) device_index: int = 0 @@ -3950,13 +3957,13 @@ def _k1_plan( *, transposed_out: bool, knobs: dict[str, Any], -) -> tuple[DetGemmPlan, torch.Tensor, torch.Tensor, torch.Tensor, str, str]: +) -> tuple[CakeDenseGemmDetProjectionPlan, torch.Tensor, torch.Tensor, torch.Tensor, str, str]: """Plan one K1 binding for its device: ``(plan, a_desc, b_desc, out3, arch, module_name)`` (the device facts - come from the per-device cache, the module from the registry; see :func:`plan_dense_projection_det_gemm`).""" + come from the per-device cache, the module from the registry; see :func:`cake_dense_gemm_det_train_projection_plan`).""" device = _same_cuda_device(A, B, out) facts = device_facts(device) arch = require_arch(device) - plan, a_desc, b_desc, O3 = plan_dense_projection_det_gemm( + plan, a_desc, b_desc, O3 = cake_dense_gemm_det_train_projection_plan( A, B, out, @@ -3970,7 +3977,7 @@ def _k1_plan( def _k1_workspace( - plan: DetGemmPlan, device: torch.device + plan: CakeDenseGemmDetProjectionPlan, device: torch.device ) -> tuple[torch.Tensor, torch.Tensor]: """The stream-K partial slabs and the slice counters of a plan (the only allocations of the K1 backend). The counters are zero-initialised once: the kernel resets every counter it uses (the last arriver of a tail @@ -4000,7 +4007,7 @@ def _out_strides(out3: torch.Tensor) -> tuple[int, int]: def _k1_values( - plan: DetGemmPlan, + plan: CakeDenseGemmDetProjectionPlan, device: torch.device, *, A: Any, @@ -4051,7 +4058,7 @@ def _k1_values( ) -def prepare_dense_projection_det_gemm( +def cake_dense_gemm_det_train_projection_prepare( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, @@ -4090,7 +4097,7 @@ def prepare_dense_projection_det_gemm( stt: Optional[bool] = None, pro: Optional[int] = None, form: Optional[str] = None, -) -> PreparedDetGemm: +) -> CakeDenseGemmDetProjectionLaunch: """Validate one binding, plan it for the device and prepare its launch (the only allocations of the K1 backend: the stream-K partial slabs and the slice counters). See the module docstring for the view contract; the keyword knobs default to the Cake launcher's choices.""" @@ -4152,7 +4159,7 @@ def prepare_dense_projection_det_gemm( out_l=out_l, ) launch = bind_launch(module_name, values, plan.grid, arch) - return PreparedDetGemm( + return CakeDenseGemmDetProjectionLaunch( module_name=module_name, arch=arch, plan=plan, @@ -4169,7 +4176,7 @@ def prepare_dense_projection_det_gemm( def _k1_swapped( - plan: DetGemmPlan, A: torch.Tensor, B: torch.Tensor, *, transposed_out: bool + plan: CakeDenseGemmDetProjectionPlan, A: torch.Tensor, B: torch.Tensor, *, transposed_out: bool ) -> bool: """Did the planner run the swapped GEMM (``B.T @ A.T`` with the transposed store) for these views? The launcher's small-M swap rule on the caller's dimensions, unless the registry fallback undid the swap @@ -4183,7 +4190,7 @@ def _k1_views( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, swapped: bool ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The operand views of a K1 launch for the caller's tensors -- ``(a_desc, b_desc, out3)`` exactly as - :func:`plan_dense_projection_det_gemm` derives them (batched, the operands exchanged and transposed when the plan + :func:`cake_dense_gemm_det_train_projection_plan` derives them (batched, the operands exchanged and transposed when the plan swapped, the ``[L, outer, inner]`` TMA spans of :func:`operand_view`) -- without planning. ``operand_view`` re-checks the stride and alignment contract of the views.""" A3, B3, O3 = as_batched(A, "A"), as_batched(B, "B"), as_batched(out, "out") @@ -4212,7 +4219,7 @@ class _CachedGemm: module_name: str arch: str - plan: DetGemmPlan + plan: CakeDenseGemmDetProjectionPlan swapped: bool entry: Callable[..., Any] = field(repr=False) template: tuple = field(repr=False) @@ -4251,8 +4258,8 @@ def _prepare_cached_gemm( for got, want in zip(rebuilt, (a_desc, b_desc, O3), strict=True) ): raise RuntimeError( - "dense_projection_det_gemm: the per-call operand views differ from the planner's; " - "the convenience path cannot cache this binding (use prepare_dense_projection_det_gemm)" + "cake_dense_gemm_det_train_projection: the per-call operand views differ from the planner's; " + "the convenience path cannot cache this binding (use cake_dense_gemm_det_train_projection_prepare)" ) ldo, out_l = _out_strides(O3) values = _k1_values( @@ -4487,7 +4494,7 @@ def _run_swap_ab( return out -def dense_projection_det_gemm( +def cake_dense_gemm_det_train_projection( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, @@ -4512,16 +4519,16 @@ def dense_projection_det_gemm( # --- the training operations as thin wrappers (the Cake entry points) --------------------- -def projection_det_forward( +def cake_dense_gemm_det_train_projection_forward( X: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None ) -> torch.Tensor: """``out[T, N] = X[T, K] @ W[N, K].T`` (bf16 out by default).""" if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.bfloat16) - return dense_projection_det_gemm(X, W.t(), out) + return cake_dense_gemm_det_train_projection(X, W.t(), out) -def projection_det_dgrad( +def cake_dense_gemm_det_train_projection_dgrad( G: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None, @@ -4533,7 +4540,7 @@ def projection_det_dgrad( out = torch.empty( G.shape[0], W.shape[1], device=G.device, dtype=out_dtype or torch.bfloat16 ) - return dense_projection_det_gemm(G, W, out) + return cake_dense_gemm_det_train_projection(G, W, out) def wgrad_views( @@ -4547,7 +4554,7 @@ def wgrad_views( when ``wgrad_swapped`` the swapped GEMM ``X.T @ G`` runs with the transposed store (the result lands directly in ``[N, K]``), otherwise ``G.T @ X``. ``arch`` / ``out_f32`` let a measured rule of that architecture and output dtype decide the orientation (see ``wgrad_swapped``). - [Cake ``projection_det_wgrad`` L1507-L1518 / ``wgrad_swapped`` L1496-L1504]""" + [Cake ``cake_dense_gemm_det_train_projection_wgrad`` L1507-L1518 / ``wgrad_swapped`` L1496-L1504]""" N, K = int(G.shape[1]), int(X.shape[1]) if wgrad_swapped(N, K, arch=arch, out_f32=out_f32): return X.t(), G, True @@ -4591,25 +4598,25 @@ def wgrad_swapped( return swapped < direct -def prepare_projection_det_wgrad( +def cake_dense_gemm_det_train_projection_wgrad_prepare( G: torch.Tensor, X: torch.Tensor, out: torch.Tensor, *, sk: bool | str = "auto", cta_rows: Optional[int] = None, -) -> PreparedDetGemm: +) -> CakeDenseGemmDetProjectionLaunch: """Prepared ``out[N, K] = G[T, N].T @ X[T, K]``; the orientation follows the row's measured rule, else the small-``N`` swap rule of the Cake launcher (``wgrad_swapped``).""" A, B, transposed = wgrad_views( G, X, arch=arch_for(out.device), out_f32=out.dtype == torch.float32 ) - return prepare_dense_projection_det_gemm( + return cake_dense_gemm_det_train_projection_prepare( A, B, out, transposed_out=transposed, sk=sk, cta_rows=cta_rows ) -def projection_det_wgrad( +def cake_dense_gemm_det_train_projection_wgrad( G: torch.Tensor, X: torch.Tensor, out: Optional[torch.Tensor] = None, @@ -4655,13 +4662,13 @@ def router_layout_class(a_mn: bool, b_mn: bool) -> str: if a_mn and b_mn: return "nn_t" raise NotImplementedError( - "router_fp32_det_gemm: an MN-major A with a K-major B has no registered instance " + "cake_dense_gemm_det_train_router: an MN-major A with a K-major B has no registered instance " f"(see {TRACKING_ISSUE})" ) @dataclass(frozen=True) -class RouterDetPlan: +class CakeDenseGemmDetRouterPlan: M: int N: int K: int @@ -4691,7 +4698,7 @@ def _router_a_layout(A: torch.Tensor) -> bool: return False if A.stride(0) == 1: return True - raise ValueError("router_fp32_det_gemm: A needs a unit stride on m or k") + raise ValueError("cake_dense_gemm_det_train_router: A needs a unit stride on m or k") def _router_b_layout(B: torch.Tensor) -> bool: @@ -4699,16 +4706,16 @@ def _router_b_layout(B: torch.Tensor) -> bool: return False if B.stride(1) == 1: # n contiguous -> MN-major parts [3, K, N] return True - raise ValueError("router_fp32_det_gemm: B needs a unit stride on k or n") + raise ValueError("cake_dense_gemm_det_train_router: B needs a unit stride on k or n") -def plan_router_fp32_det_gemm( +def cake_dense_gemm_det_train_router_plan( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] = None, -) -> tuple[RouterDetPlan, torch.Tensor, torch.Tensor]: +) -> tuple[CakeDenseGemmDetRouterPlan, torch.Tensor, torch.Tensor]: """Classify the fp32 views, apply the swap rule of the weight gradient and plan the launch. Returns ``(plan, A_eff, B_eff)``: the operands of the launched (possibly swapped) GEMM; @@ -4719,25 +4726,25 @@ def plan_router_fp32_det_gemm( or B.dtype != torch.float32 or out.dtype != torch.float32 ): - raise ValueError("router_fp32_det_gemm: A, B and out must be fp32") + raise ValueError("cake_dense_gemm_det_train_router: A, B and out must be fp32") if A.dim() != 2 or B.dim() != 2 or out.dim() != 2: - raise ValueError("router_fp32_det_gemm: A, B and out must be 2-D views") + raise ValueError("cake_dense_gemm_det_train_router: A, B and out must be 2-D views") M, K = (int(s) for s in A.shape) Kb, N = (int(s) for s in B.shape) if Kb != K: raise ValueError( - f"router_fp32_det_gemm: A is [M={M}, K={K}] but B is [{Kb}, {N}]; K must agree" + f"cake_dense_gemm_det_train_router: A is [M={M}, K={K}] but B is [{Kb}, {N}]; K must agree" ) if tuple(out.shape) != (M, N): raise ValueError( - f"router_fp32_det_gemm: out must be [{M}, {N}], got {tuple(out.shape)}" + f"cake_dense_gemm_det_train_router: out must be [{M}, {N}], got {tuple(out.shape)}" ) if M < 1 or K < 1 or N < 1: # An empty problem (a ragged batch with T = 0 reaches here as M = 0 for the # forward / input gradient or K = 0 for the weight gradient) would otherwise # produce a zero-sized grid or a TMA descriptor with a zero extent. raise ValueError( - f"router_fp32_det_gemm: M, K and N must be >= 1 (M={M}, N={N}, K={K})" + f"cake_dense_gemm_det_train_router: M, K and N must be >= 1 (M={M}, N={N}, K={K})" ) a_mn, b_mn = _router_a_layout(A), _router_b_layout(B) layout = router_layout_class(a_mn, b_mn) @@ -4752,33 +4759,33 @@ def plan_router_fp32_det_gemm( a_eff_mn, b_eff_mn = _router_a_layout(A_eff), _router_b_layout(B_eff) lda = int(A_eff.stride(1)) if a_eff_mn else int(A_eff.stride(0)) if lda % 4 != 0 or A_eff.data_ptr() % 16 != 0: - raise ValueError("router_fp32_det_gemm: A rows must be 16-byte aligned") + raise ValueError("cake_dense_gemm_det_train_router: A rows must be 16-byte aligned") if B_eff.data_ptr() % 16 != 0: raise ValueError( - "router_fp32_det_gemm: B must start at a 16-byte aligned address" + "cake_dense_gemm_det_train_router: B must start at a 16-byte aligned address" ) # the kernel's output view: [M_eff, N_eff] row-major, or [N_eff, M_eff] with the transposed store = out if out.stride(1) != 1 or out.stride(0) % 4 != 0 or N_eff % 8 != 0: raise ValueError( - "router_fp32_det_gemm: out must have unit inner stride and 16-byte rows; N a multiple of 8" + "cake_dense_gemm_det_train_router: out must have unit inner stride and 16-byte rows; N a multiple of 8" ) if splits is None: splits = ROUTER_SPLITS[layout] splits = int(splits) if splits < 1: - raise ValueError("router_fp32_det_gemm: splits must be >= 1") + raise ValueError("cake_dense_gemm_det_train_router: splits must be >= 1") # one split stores straight into ``out``: the row-major store (kk / kn) writes 16-byte vectors, the # transposed store (nn_t) one fp32 element at a time; more splits store into the workspace and reduce into ``out`` if splits == 1 and not transposed and out.data_ptr() % 16 != 0: raise ValueError( - "router_fp32_det_gemm: out must start at a 16-byte aligned address " + "cake_dense_gemm_det_train_router: out must start at a 16-byte aligned address " f"(the single-split {layout} launch stores 16-byte vectors into it)" ) m_tiles = _ceil_div(M_eff, ROUTER_BLOCK_M) m_tiles += m_tiles % CTA_GROUP n_tiles = _ceil_div(N_eff, ROUTER_BLOCK_N) k_blocks = _ceil_div(K, ROUTER_BLOCK_K) - plan = RouterDetPlan( + plan = CakeDenseGemmDetRouterPlan( M=M_eff, N=N_eff, K=K, @@ -4853,7 +4860,7 @@ def _split_grid(n: int) -> tuple[int, int, int]: @dataclass -class PreparedRouterDetGemm: +class CakeDenseGemmDetRouterLaunch: """One prepared FP32 router GEMM: the retained BF16x3 stack of the second operand (re-split from the bound fp32 view on every ``launch()`` -- one launch of the fused split program when this checkout registers it for the device, else the torch chain through the fp32 scratch @@ -4862,7 +4869,7 @@ class PreparedRouterDetGemm: module_name: str arch: str - plan: RouterDetPlan + plan: CakeDenseGemmDetRouterPlan out: torch.Tensor tensors: dict[str, torch.Tensor] = field(repr=False) device_index: int = 0 @@ -4906,16 +4913,16 @@ def launch(self) -> torch.Tensor: def _k2_plan( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] -) -> tuple[RouterDetPlan, torch.Tensor, torch.Tensor, str, str]: +) -> tuple[CakeDenseGemmDetRouterPlan, torch.Tensor, torch.Tensor, str, str]: """Plan one K2 binding for its device: ``(plan, A_eff, B_eff, arch, module_name)``.""" device = _same_cuda_device(A, B, out) arch = require_arch(device) - plan, A_eff, B_eff = plan_router_fp32_det_gemm(A, B, out, splits=splits) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) return plan, A_eff, B_eff, arch, select_module(arch, plan.template) def _k2_views( - A: torch.Tensor, B: torch.Tensor, plan: RouterDetPlan + A: torch.Tensor, B: torch.Tensor, plan: CakeDenseGemmDetRouterPlan ) -> tuple[torch.Tensor, torch.Tensor]: """``(a_view, split_src)`` of a router launch for the caller's operands: the streamed operand as the ``[1, M, K]`` / ``[1, K, M]`` TMA span and the second operand in its own layout -- the source of the BF16x3 @@ -4932,10 +4939,10 @@ def _k2_checked_views( B: torch.Tensor, A_eff: torch.Tensor, B_eff: torch.Tensor, - plan: RouterDetPlan, + plan: CakeDenseGemmDetRouterPlan, ) -> tuple[torch.Tensor, torch.Tensor]: """:func:`_k2_views` of the caller's operands, checked against the planner's operands ``(A_eff, B_eff)`` - of :func:`plan_router_fp32_det_gemm` in the launch form (the streamed ``[1, M, K]`` / ``[1, K, M]`` span under + of :func:`cake_dense_gemm_det_train_router_plan` in the launch form (the streamed ``[1, M, K]`` / ``[1, K, M]`` span under ``plan.a_mn``, the split source under ``plan.b_mn``). The convenience path rebuilds the views from the caller's tensors on every call, so a rebuild that is not the planner's view fails closed here (the K1 check of :func:`_prepare_cached_gemm`) instead of launching on a view the plan never validated.""" @@ -4944,14 +4951,14 @@ def _k2_checked_views( want_split = B_eff if plan.b_mn else B_eff.t() if not (_same_view(a_view, want_a) and _same_view(split_src, want_split)): raise RuntimeError( - "router_fp32_det_gemm: the per-call operand views differ from the planner's; " - "the convenience path cannot cache this binding (use prepare_router_fp32_det_gemm)" + "cake_dense_gemm_det_train_router: the per-call operand views differ from the planner's; " + "the convenience path cannot cache this binding (use cake_dense_gemm_det_train_router_prepare)" ) return a_view, split_src def _k2_workspace( - plan: RouterDetPlan, + plan: CakeDenseGemmDetRouterPlan, split_src: torch.Tensor, device: torch.device, *, @@ -4984,7 +4991,7 @@ def _k2_out32(out: torch.Tensor) -> torch.Tensor: def _k2_values( - plan: RouterDetPlan, + plan: CakeDenseGemmDetRouterPlan, *, A: Any, parts: Any, @@ -5009,13 +5016,13 @@ def _k2_values( ) -def prepare_router_fp32_det_gemm( +def cake_dense_gemm_det_train_router_prepare( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] = None, -) -> PreparedRouterDetGemm: +) -> CakeDenseGemmDetRouterLaunch: """Validate one fp32 binding, plan it and prepare its launch (the only allocations of the K2 backend: the BF16x3 stack, the fp32 residual scratch of the torch-chain split and the split-K workspace).""" @@ -5042,7 +5049,7 @@ def prepare_router_fp32_det_gemm( split_launch = bind_launch( split_module, _split_values(split_src, parts, n), _split_grid(n), arch ) - return PreparedRouterDetGemm( + return CakeDenseGemmDetRouterLaunch( module_name=module_name, arch=arch, plan=plan, @@ -5064,7 +5071,7 @@ class _CachedRouter: module_name: str arch: str - plan: RouterDetPlan + plan: CakeDenseGemmDetRouterPlan entry: Callable[..., Any] = field(repr=False) template: tuple = field(repr=False) arena: _Arena = field(repr=False) @@ -5089,7 +5096,7 @@ def carve(self, arena: _Arena) -> dict[str, torch.Tensor]: _K2_CACHE: "OrderedDict[tuple, _CachedRouter]" = OrderedDict() -def clear_convenience_caches() -> None: +def cake_dense_gemm_det_train_convenience_cache_clear() -> None: """Drop every cached convenience-path launch (K1 and K2) and release every workspace arena; a launch in flight keeps a reference to the buffers it was bound with until its FFI call has enqueued it, and the freed blocks are reused stream-ordered (see :class:`_Arena`).""" @@ -5212,7 +5219,7 @@ def _launch_cached_router( return out -def router_fp32_det_gemm( +def cake_dense_gemm_det_train_router( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] = None ) -> torch.Tensor: """``out[m, n] = sum_k A[m, k] B[k, n]`` for fp32 views through the emulation (the cached convenience @@ -5220,7 +5227,7 @@ def router_fp32_det_gemm( return _launch_cached_router(A, B, out, splits=splits) -def router_det_forward( +def cake_dense_gemm_det_train_router_forward( X: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None, @@ -5230,19 +5237,19 @@ def router_det_forward( """``out[T, N] = X[T, K] @ W[N, K].T`` in FP32 through the emulation.""" if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.float32) - return router_fp32_det_gemm(X, W.t(), out, splits=splits) + return cake_dense_gemm_det_train_router(X, W.t(), out, splits=splits) -def router_det_dgrad( +def cake_dense_gemm_det_train_router_dgrad( G: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None ) -> torch.Tensor: """``out[T, K] = G[T, N] @ W[N, K]`` in FP32 through the emulation.""" if out is None: out = torch.empty(G.shape[0], W.shape[1], device=G.device, dtype=torch.float32) - return router_fp32_det_gemm(G, W, out, splits=1) + return cake_dense_gemm_det_train_router(G, W, out, splits=1) -def router_det_wgrad( +def cake_dense_gemm_det_train_router_wgrad( G: torch.Tensor, X: torch.Tensor, out: Optional[torch.Tensor] = None, @@ -5252,25 +5259,5 @@ def router_det_wgrad( """``out[N, K] = G[T, N].T @ X[T, K]`` in FP32: the swapped GEMM ``X.T @ G`` with a transposed store.""" if out is None: out = torch.empty(G.shape[1], X.shape[1], device=G.device, dtype=torch.float32) - return router_fp32_det_gemm(G.t(), X, out, splits=splits) - - -# --- the former public names (before the det rename), aliases of the canonical objects ----------------------- -GemmPlan = DetGemmPlan -PreparedGemm = PreparedDetGemm -PreparedRouterGemm = PreparedRouterDetGemm -RouterPlan = RouterDetPlan -dense_projection_gemm = dense_projection_det_gemm -plan_dense_projection_gemm = plan_dense_projection_det_gemm -plan_router_fp32_gemm = plan_router_fp32_det_gemm -prepare_dense_projection_gemm = prepare_dense_projection_det_gemm -prepare_projection_wgrad = prepare_projection_det_wgrad -prepare_router_fp32_gemm = prepare_router_fp32_det_gemm -projection_dgrad = projection_det_dgrad -projection_forward = projection_det_forward -projection_wgrad = projection_det_wgrad -router_dgrad = router_det_dgrad -router_forward = router_det_forward -router_fp32_gemm = router_fp32_det_gemm -router_wgrad = router_det_wgrad -split_fp32_to_bf16x3_ = split_fp32_to_bf16x3_det_ + return cake_dense_gemm_det_train_router(G.t(), X, out, splits=splits) + diff --git a/flashinfer/experimental/dense_projection_det_gemm/cake_build.py b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_build.py similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/cake_build.py rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_build.py diff --git a/flashinfer/experimental/dense_projection_det_gemm/cake_jit.py b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py similarity index 99% rename from flashinfer/experimental/dense_projection_det_gemm/cake_jit.py rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py index 6b26c764c9b..0e7f9f3795b 100644 --- a/flashinfer/experimental/dense_projection_det_gemm/cake_jit.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py @@ -290,9 +290,6 @@ def load_cake_dense_projection_det_gemm_module(name: str, arch: str): return cake_build.load_library(library_name(name, arch)) -# The loader's former name (before the det rename), an alias of the same cached function. -load_cake_dense_projection_gemm_module = load_cake_dense_projection_det_gemm_module - def specs_for_arch(arch: str): """JIT specs of every ``build.json`` entry of ``arch`` without a prebuilt library (the diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/.clang-format b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/.clang-format similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/.clang-format rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/.clang-format diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_011c85b820ec0cf6863d_tpl_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_011c85b820ec0cf6863d_tpl_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_011c85b820ec0cf6863d_tpl_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_011c85b820ec0cf6863d_tpl_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_02803e1d2393113c0f6c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_087ba4f157914949c26a_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_087ba4f157914949c26a_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_087ba4f157914949c26a_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_087ba4f157914949c26a_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0a71a49da71c7d34ab37_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0a71a49da71c7d34ab37_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0a71a49da71c7d34ab37_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0a71a49da71c7d34ab37_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0b91f29d88de00658fc5_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0bbd5d9f1d72a90ef3ea_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0c72e864a4afb5347da6_tpl_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0c72e864a4afb5347da6_tpl_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0c72e864a4afb5347da6_tpl_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0c72e864a4afb5347da6_tpl_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0d487a8f71dac54321ef_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0db65f3ddd780ed42d4d_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0db65f3ddd780ed42d4d_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0db65f3ddd780ed42d4d_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0db65f3ddd780ed42d4d_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0dd5420b570cabc144ab_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13790f025541d7d162a7_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13d9310f37868a18eace_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13d9310f37868a18eace_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13d9310f37868a18eace_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13d9310f37868a18eace_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_14192244441c1003db74_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_14192244441c1003db74_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_14192244441c1003db74_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_14192244441c1003db74_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_20a2e5c04427e93572da_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_23305bdb787b7c6d7378_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_295aef2f2e0332de2372_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_295aef2f2e0332de2372_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_295aef2f2e0332de2372_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_295aef2f2e0332de2372_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29d013dfb72e932eaee5_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_29de51d5e9cdd6d22491_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3088b8056e0b7f9ca32c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3088b8056e0b7f9ca32c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3088b8056e0b7f9ca32c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3088b8056e0b7f9ca32c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3a87cc189b4a7c1a8d7d_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3d7b66f61c710f893b29_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_438841f42a0743b63fb5_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_438841f42a0743b63fb5_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_438841f42a0743b63fb5_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_438841f42a0743b63fb5_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_446497e750658dfb057b_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_446497e750658dfb057b_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_446497e750658dfb057b_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_446497e750658dfb057b_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_49acb91191bbcced9051_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_49acb91191bbcced9051_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_49acb91191bbcced9051_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_49acb91191bbcced9051_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4de304a13db1f2a0e850_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_4eb96fb2eb1aaa14c7f7_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_505fbe05ce8b62f0a589_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_505fbe05ce8b62f0a589_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_505fbe05ce8b62f0a589_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_505fbe05ce8b62f0a589_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_51c81f11e3b490b5aebb_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5226b6c854d7db395d7e_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5226b6c854d7db395d7e_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5226b6c854d7db395d7e_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5226b6c854d7db395d7e_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5445e0f16ad84868fa01_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5445e0f16ad84868fa01_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5445e0f16ad84868fa01_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5445e0f16ad84868fa01_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_552070c21b97faa67a71_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5568657a526e2040f0bc_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5568657a526e2040f0bc_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5568657a526e2040f0bc_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5568657a526e2040f0bc_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_56be8af43d5d0822da87_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_56be8af43d5d0822da87_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_56be8af43d5d0822da87_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_56be8af43d5d0822da87_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_59fb66ccddd67565e592_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5a06c5e2f33772a33174_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5a06c5e2f33772a33174_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5a06c5e2f33772a33174_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5a06c5e2f33772a33174_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5b4e87b523f799e09384_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5e77ef909e88a99c4357_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5e77ef909e88a99c4357_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5e77ef909e88a99c4357_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5e77ef909e88a99c4357_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_624eb3d7232071ddaf91_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_6579637169bec3c2f1d5_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_66df6446dd2cca3106d3_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_66df6446dd2cca3106d3_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_66df6446dd2cca3106d3_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_66df6446dd2cca3106d3_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_736ad30e83d88796883f_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_75b1dcdaef42319f5ed4_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_765ece06955b70a95b29_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_765ece06955b70a95b29_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_765ece06955b70a95b29_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_765ece06955b70a95b29_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7a0ce55fd7221e354322_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7d403fa7b167188ccc24_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7d403fa7b167188ccc24_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7d403fa7b167188ccc24_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7d403fa7b167188ccc24_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7ef28d420de555e7a28d_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_806927d101fe72fc1317_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_806927d101fe72fc1317_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_806927d101fe72fc1317_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_806927d101fe72fc1317_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_82f698f67a1cca6e4291_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_82f698f67a1cca6e4291_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_82f698f67a1cca6e4291_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_82f698f67a1cca6e4291_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_895f9ba7159bea07aedd_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_895f9ba7159bea07aedd_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_895f9ba7159bea07aedd_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_895f9ba7159bea07aedd_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_8e1f6805342f620cca8c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_8e1f6805342f620cca8c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_8e1f6805342f620cca8c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_8e1f6805342f620cca8c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9269f6841d7aeb58471c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9475731f174aaefb2e44_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9aa9aa60b0c544bfab06_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9aa9aa60b0c544bfab06_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9aa9aa60b0c544bfab06_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9aa9aa60b0c544bfab06_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9cc6363566760a38b4fe_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9cc6363566760a38b4fe_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9cc6363566760a38b4fe_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9cc6363566760a38b4fe_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9ffc37c198f2d2f42c1f_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_a24a3d3fd3ceac3e1cdf_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_aa4f7c2fbd0b9a99e535_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b26ac3f3ed56a78cb052_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b26ac3f3ed56a78cb052_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b26ac3f3ed56a78cb052_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b26ac3f3ed56a78cb052_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b2a9665ee12fd6afea2b_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b58b50619ef5ed46e109_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b58b50619ef5ed46e109_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b58b50619ef5ed46e109_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b58b50619ef5ed46e109_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b773cc0e3894fd0ef2e0_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab00fb2b95c93758095_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bab9e9dfd7c6ca5e17c0_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_bb0e7457db664d6f0f6a_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c363601373e2dfb8e9af_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_c63d1400e64208d7865c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ca9a846738b1443f7318_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ca9a846738b1443f7318_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ca9a846738b1443f7318_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ca9a846738b1443f7318_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2a5e70bd49083888179_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2a5e70bd49083888179_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2a5e70bd49083888179_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2a5e70bd49083888179_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2ea37e4d394bf40922a_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d6ebaab2415ea7b00de9_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d6ebaab2415ea7b00de9_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d6ebaab2415ea7b00de9_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d6ebaab2415ea7b00de9_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ddff2b911137449c1dfc_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ddff2b911137449c1dfc_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ddff2b911137449c1dfc_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ddff2b911137449c1dfc_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_de7492eeb6e24b31687a_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_de7492eeb6e24b31687a_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_de7492eeb6e24b31687a_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_de7492eeb6e24b31687a_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_device_common.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_device_common.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_device_common.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_device_common.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_dfd65a75564d04d3826f_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e2555d533680e08880ae_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e33ac44bfff66daaee64_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e98b3ee7b774c63d7544_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e98b3ee7b774c63d7544_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e98b3ee7b774c63d7544_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_e98b3ee7b774c63d7544_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eb2bb538645fe303bcca_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ebbb505a3077e8da0a50_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ec7277a904a7211d67ce_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ed1f3ab86ddc55345e69_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_edf7fa4c4f889f0c4981_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_eec6f87f8b7f25eb6aa2_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f31e1dcb1df00973dc0c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f31e1dcb1df00973dc0c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f31e1dcb1df00973dc0c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f31e1dcb1df00973dc0c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f589a6aff5e88318c36c_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f6c9b40184336f49390d_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f993d8de6f29f9e8a1ac_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb4bcdf13915dfb96c59_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb4bcdf13915dfb96c59_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb4bcdf13915dfb96c59_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb4bcdf13915dfb96c59_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb766b1be68412e6fda7_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_binding.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fdb106e6e73897521cbb_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_host_common.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_host_common.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_host_common.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_host_common.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_c1__hen_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_c1__hen_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_c1__hen_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_c1__hen_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n128_f32_tma1_c1__pd2_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n128_f32_tma1_c1__pd2_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n128_f32_tma1_c1__pd2_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n128_f32_tma1_c1__pd2_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_cp_async_bulk_gmem2smem_5198c9805c6ecc8d.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_cp_async_bulk_gmem2smem_5198c9805c6ecc8d.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_cp_async_bulk_gmem2smem_5198c9805c6ecc8d.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_cp_async_bulk_gmem2smem_5198c9805c6ecc8d.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_elect_commit_8131fdc67daf4d23.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_elect_commit_8131fdc67daf4d23.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_elect_commit_8131fdc67daf4d23.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_elect_commit_8131fdc67daf4d23.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_elect_commit_cg2_multicast_18e7b4ad8e69c467.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_elect_commit_cg2_multicast_18e7b4ad8e69c467.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_elect_commit_cg2_multicast_18e7b4ad8e69c467.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_elect_commit_cg2_multicast_18e7b4ad8e69c467.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_elect_sync_1ee8cd91025d8932.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_elect_sync_1ee8cd91025d8932.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_elect_sync_1ee8cd91025d8932.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_elect_sync_1ee8cd91025d8932.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_incr_smem_desc_lo_4429b045244f74bc.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_incr_smem_desc_lo_4429b045244f74bc.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_incr_smem_desc_lo_4429b045244f74bc.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_incr_smem_desc_lo_4429b045244f74bc.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_make_warp_uniform_15c3664d3e11e8cd.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_make_warp_uniform_15c3664d3e11e8cd.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_make_warp_uniform_15c3664d3e11e8cd.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_make_warp_uniform_15c3664d3e11e8cd.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_arrive_71ef360c3ac78eb7.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_arrive_71ef360c3ac78eb7.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_arrive_71ef360c3ac78eb7.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_arrive_71ef360c3ac78eb7.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_arrive_expect_tx_54ce59d489f69526.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_arrive_expect_tx_54ce59d489f69526.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_arrive_expect_tx_54ce59d489f69526.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_arrive_expect_tx_54ce59d489f69526.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_init_15e581aef85ee586.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_init_15e581aef85ee586.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_init_15e581aef85ee586.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_init_15e581aef85ee586.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_wait_c93bf7e39b218733.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_wait_c93bf7e39b218733.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_wait_c93bf7e39b218733.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_wait_c93bf7e39b218733.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_wait_cluster_hint_bdb29a0e36dedbee.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_wait_cluster_hint_bdb29a0e36dedbee.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_mbarrier_wait_cluster_hint_bdb29a0e36dedbee.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_mbarrier_wait_cluster_hint_bdb29a0e36dedbee.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tcgen05_mma_f16_03fde8c1cbf058c8.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tcgen05_mma_f16_03fde8c1cbf058c8.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tcgen05_mma_f16_03fde8c1cbf058c8.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tcgen05_mma_f16_03fde8c1cbf058c8.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_7df1e07811f83fc6.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_7df1e07811f83fc6.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_7df1e07811f83fc6.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_7df1e07811f83fc6.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_cta2_dcd29ed146891990.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_cta2_dcd29ed146891990.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_cta2_dcd29ed146891990.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tma_3d_gmem2smem_cta2_dcd29ed146891990.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tma_store_3d_a0d668161ec63906.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tma_store_3d_a0d668161ec63906.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tma_store_3d_a0d668161ec63906.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tma_store_3d_a0d668161ec63906.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tmem_ld_x16_a9ed61940fdba296.cuh b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tmem_ld_x16_a9ed61940fdba296.cuh similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_device_helpers/cake_tmem_ld_x16_a9ed61940fdba296.cuh rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_device_helpers/cake_tmem_ld_x16_a9ed61940fdba296.cuh diff --git a/flashinfer/experimental/dense_projection_det_gemm/csrc/cake_ffi_checks.h b/flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_ffi_checks.h similarity index 100% rename from flashinfer/experimental/dense_projection_det_gemm/csrc/cake_ffi_checks.h rename to flashinfer/experimental/cake/cake_dense_gemm_det_train/csrc/cake_ffi_checks.h diff --git a/pyproject.toml b/pyproject.toml index 27518f4b658..0c6c059baab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -185,8 +185,9 @@ exclude = [ "flashinfer.experimental.cake_moe_grouped_det_gemm" = ["csrc/**", "build.json"] # Generated sources of its non-deterministic counterpart (deterministic=False). "flashinfer.experimental.cake.cake_moe_grouped_gemm" = ["csrc/**", "build.json"] -# Generated sources for the experimental exact-SM100a/SM107a dense projection GEMM backend. -"flashinfer.experimental.dense_projection_det_gemm" = ["csrc/**", "build.json"] +# Generated sources for the deterministic SM100a/SM103a/SM107a dense GEMM training backend +# (flashinfer.cake_dense_gemm_det_train). +"flashinfer.experimental.cake.cake_dense_gemm_det_train" = ["csrc/**", "build.json"] # Generated sources for the experimental SM100a/SM103a/SM107a dense projection GEMM backend with an explicit # numerical mode (the non-deterministic split-K programs; the deterministic ones ship with dense_projection_det_gemm). "flashinfer.experimental.dense_projection_gemm" = ["csrc/**", "build.json"] diff --git a/scripts/check_cake_build_manifests.py b/scripts/check_cake_build_manifests.py index 0d0122c2ac0..48325008c59 100755 --- a/scripts/check_cake_build_manifests.py +++ b/scripts/check_cake_build_manifests.py @@ -538,8 +538,8 @@ class Package: refs_mok_bf16, ), Package( - "flashinfer.experimental.dense_projection_det_gemm", - "flashinfer/experimental/dense_projection_det_gemm", + "flashinfer.experimental.cake.cake_dense_gemm_det_train", + "flashinfer/experimental/cake/cake_dense_gemm_det_train", "cake_jit.py", refs_generic, ), diff --git a/tests/experimental/test_cake_dense_projection_det_gemm.py b/tests/experimental/test_cake_dense_gemm_det_train.py similarity index 94% rename from tests/experimental/test_cake_dense_projection_det_gemm.py rename to tests/experimental/test_cake_dense_gemm_det_train.py index 5dd9c2aad25..8a6f26b6d27 100644 --- a/tests/experimental/test_cake_dense_projection_det_gemm.py +++ b/tests/experimental/test_cake_dense_gemm_det_train.py @@ -30,9 +30,9 @@ # Backend suite: the package is named once, through ``cake_package``; nothing else # of flashinfer is imported here. -cake_backend = cake_package("dense_projection_det_gemm", "cake_backend") -cake_build = cake_package("dense_projection_det_gemm", "cake_build") -cake_jit = cake_package("dense_projection_det_gemm", "cake_jit") +cake_backend = cake_package("cake_dense_gemm_det_train", "cake_backend") +cake_build = cake_package("cake_dense_gemm_det_train", "cake_build") +cake_jit = cake_package("cake_dense_gemm_det_train", "cake_jit") ARCH_RULE_DEFAULTS = cake_backend.ARCH_RULE_DEFAULTS BLOCK_K = cake_backend.BLOCK_K BLOCK_N_CHOICES = cake_backend.BLOCK_N_CHOICES @@ -68,11 +68,11 @@ sk_sync_plan = cake_backend.sk_sync_plan check_stream_k_plan = cake_backend.check_stream_k_plan operand_view = cake_backend.operand_view -plan_dense_projection_det_gemm = cake_backend.plan_dense_projection_det_gemm -plan_router_fp32_det_gemm = cake_backend.plan_router_fp32_det_gemm -prepare_dense_projection_det_gemm = cake_backend.prepare_dense_projection_det_gemm -prepare_projection_det_wgrad = cake_backend.prepare_projection_det_wgrad -prepare_router_fp32_det_gemm = cake_backend.prepare_router_fp32_det_gemm +cake_dense_gemm_det_train_projection_plan = cake_backend.cake_dense_gemm_det_train_projection_plan +cake_dense_gemm_det_train_router_plan = cake_backend.cake_dense_gemm_det_train_router_plan +cake_dense_gemm_det_train_projection_prepare = cake_backend.cake_dense_gemm_det_train_projection_prepare +cake_dense_gemm_det_train_projection_wgrad_prepare = cake_backend.cake_dense_gemm_det_train_projection_wgrad_prepare +cake_dense_gemm_det_train_router_prepare = cake_backend.cake_dense_gemm_det_train_router_prepare router_layout_class = cake_backend.router_layout_class row_rule = cake_backend.row_rule sk_parts_plan = cake_backend.sk_parts_plan @@ -1516,7 +1516,7 @@ def test_round15_store_hint_knob(): # the round-15 sm_107a rule: the batched MLA qabs forward (N = 512, K = 192) plans the 8-head interleave with # evict_last stores; sm_100a keeps the round-13 plan without the operand v = _views("mla", "qabs", "fwd", "bf16", 1001) - r200, *_ = plan_dense_projection_det_gemm( + r200, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -1527,7 +1527,7 @@ def test_round15_store_hint_knob(): ) assert r200.store_hint == "evict_last" and r200.batch_group == 8 assert _nopro(r200.template) == "dense_proj_det_gemm_kn_n256_tma1_bg8_sol" - b200, *_ = plan_dense_projection_det_gemm( + b200, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -1582,7 +1582,7 @@ def test_round15_pd_sh_knobs(): # the round-15 rules: the sm_107a indexer_k bf16 input gradient (N = 6144, K = 128) plans the 128-row two-slot # TMA-store family with the cross-tile pipelined drain; a caller-forced tall tile drops the family knob v = _views("proj", "indexer_k", "dgrad", "bf16", 1001) - r200, *_ = plan_dense_projection_det_gemm( + r200, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -1593,7 +1593,7 @@ def test_round15_pd_sh_knobs(): ) assert r200.pd == 2 and r200.slots == 2 and r200.cta_rows == 128 assert _nopro(r200.template) == "dense_proj_det_gemm_kn_n256_tma2_pd2" - tall, *_ = plan_dense_projection_det_gemm( + tall, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -1774,7 +1774,7 @@ def test_sk_sync_rule_yields_to_caller_forced_forms(): # registrations of the ``_m256`` / ``_ov`` programs), while a caller-forced ``sk_sync`` keeps raising on the conflict v = _views("proj", "q_a", "wgrad", "bf16", 16231) kw = dict(sm_count=212, l2_bytes=L2_BYTES, arch="sm_107a", _fallback=False) - rule, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + rule, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert ( rule.sk_sync and rule.cta_rows == 128 and "_sks" in _nopro(rule.template) ) # round 25: the row's rule adds the slab path (`_sks_sb2`) @@ -1785,7 +1785,7 @@ def test_sk_sync_rule_yields_to_caller_forced_forms(): 216, ) assert rule.sk_iters == 86 and rule.sk_collectors == 20 - tall, *_ = plan_dense_projection_det_gemm( + tall, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta_rows=256, **kw ) assert ( @@ -1794,12 +1794,12 @@ def test_sk_sync_rule_yields_to_caller_forced_forms(): and "_sks" not in _nopro(tall.template) ) assert tall.sk_iters == tall.tail_tiles and tall.sk_collectors == 0 - ov, *_ = plan_dense_projection_det_gemm( + ov, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta_rows=256, ovl=True, **kw ) assert not ov.sk_sync and ov.ovl and "_ov" in _nopro(ov.template) with pytest.raises(ValueError, match="sk_sync needs the single-pass"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta_rows=256, ovl=True, sk_sync=True, **kw ) @@ -1933,7 +1933,7 @@ def test_round19_cta1_plans_one_cta_per_tile(T): # policy, and the raster group may be any count >= 1 [Cake launcher] v = _views("proj", "indexer_hw", "fwd", "bf16", T) kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a", _fallback=False) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], cta1=True, **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], cta1=True, **kw) assert plan.cta1 and _nopro(plan.template) == "dense_proj_det_gemm_kk_n128_hen_c1" assert (plan.m_tiles, plan.n_tiles, plan.pair_tiles, plan.sm_pairs) == ( 127, @@ -1945,7 +1945,7 @@ def test_round19_cta1_plans_one_cta_per_tile(T): assert plan.grid == (plan.m_tiles * plan.n_tiles, 1, 1) == (127, 1, 1) assert plan.stages == 7 and plan.ws_f32_elems == 0 and plan.sk_slab == 0 assert plan.wave_working_set_bytes == wave_working_set(127, 1, 6144, 16, 148) - pair, *_ = plan_dense_projection_det_gemm( + pair, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=False, **kw ) assert not pair.cta1 and _nopro(pair.template) == "dense_proj_det_gemm_kk_n128_hen" @@ -1957,12 +1957,12 @@ def test_round19_cta1_plans_one_cta_per_tile(T): ) assert pair.grid == (128, 1, 1) # a raster group of one is admitted under cta1 only - one, *_ = plan_dense_projection_det_gemm( + one, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=True, group_m=1, **kw ) assert one.group_m == 1 and _nopro(one.template) == _nopro(plan.template) with pytest.raises(ValueError, match="group_m must be an even number"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=False, group_m=1, **kw ) # the form carries no tail policy: a caller combining it with one asked for two incompatible forms @@ -1975,12 +1975,12 @@ def test_round19_cta1_plans_one_cta_per_tile(T): dict(sk_max_units=48), ): with pytest.raises(ValueError, match="cta1 excludes"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=True, **bad, **kw ) # a small T: one wave's panels fit the L2, so the form keeps the hint-free program small = _views("proj", "indexer_hw", "fwd", "bf16", 2049) - short, *_ = plan_dense_projection_det_gemm( + short, *_ = cake_dense_gemm_det_train_projection_plan( small["A"], small["B"], small["out"], cta1=True, **kw ) assert _nopro( @@ -2029,7 +2029,7 @@ def test_round19_rules_plan_like_the_cake_launcher(T): # streaming A operand evict_first (one column tile, the wave's panels exceed the L2) for row in ("indexer_hw", "indexer_k"): v = _views("proj", row, "fwd", "bf16", T) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert ( plan.cta1 and _nopro(plan.template) == "dense_proj_det_gemm_kk_n128_hen_c1" ), row @@ -2045,14 +2045,14 @@ def test_round19_rules_plan_like_the_cake_launcher(T): # stages (``_s4``: the single-CTA default is five), 127 x 48 tiles; a caller forcing the pair form drops the # rule's stage count (the pair default of that family is six) v = _views("proj", "indexer_hw", "dgrad", "f32", T) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert ( plan.cta1 and _nopro(plan.template) == "dense_proj_det_gemm_kn_n128_f32_tma2_s4_c1" ) assert (plan.m_tiles, plan.n_tiles, plan.stages, plan.sm_pairs) == (127, 48, 4, 148) assert plan.grid == (127 * 48, 1, 1) and plan.sk_units == 0 - pair, *_ = plan_dense_projection_det_gemm( + pair, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=False, **kw ) assert ( @@ -2066,7 +2066,7 @@ def test_round19_rules_plan_like_the_cake_launcher(T): # through the dead mainloop stages (sb2), fp32 reduce-adds into the output through the TMA-store path (sb3) v = _views("proj", "indexer_q", "wgrad", "bf16", T) assert not v["transposed"] - bf16, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + bf16, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert bf16.sk_sync and bf16.sk_slab == 2 assert _nopro(bf16.template) == "dense_proj_det_gemm_nn_n256_m256_sks_sb2" assert (bf16.pair_tiles, bf16.k_blocks, bf16.sm_pairs) == (64, -(-T // BLOCK_K), 74) @@ -2080,7 +2080,7 @@ def test_round19_rules_plan_like_the_cake_launcher(T): assert bf16.grid == (148, 1, 1) and bf16.sk_iters == 64 assert bf16.ws_f32_elems == 64 * 2 * 256 * 256 vf = _views("proj", "indexer_q", "wgrad", "f32", T) - f32, *_ = plan_dense_projection_det_gemm(vf["A"], vf["B"], vf["out"], **kw) + f32, *_ = cake_dense_gemm_det_train_projection_plan(vf["A"], vf["B"], vf["out"], **kw) assert f32.sk_sync and f32.sk_slab == 3 and f32.epi == "tma" assert _nopro(f32.template) == "dense_proj_det_gemm_nn_n256_m256_f32_tma1_sks_sb3" assert (f32.num_full, f32.tail_tiles, f32.sk_units, f32.iters_per_unit) == ( @@ -2094,34 +2094,34 @@ def test_round19_rules_plan_like_the_cake_launcher(T): # pipeline cannot stage the eight 16 KiB warp slabs, so the row keeps the plain synchronised program; a # caller-forced ``sk_slab`` keeps raising on such a conflict, and a caller-forced tile family drops the rule's # synchronised plan with its slab path - reg, *_ = plan_dense_projection_det_gemm( + reg, *_ = cake_dense_gemm_det_train_projection_plan( vf["A"], vf["B"], vf["out"], epi="reg", **kw ) assert reg.sk_sync and reg.sk_slab == 2 assert _nopro(reg.template) == "dense_proj_det_gemm_nn_n256_m256_f32_sks_sb2" - narrow, *_ = plan_dense_projection_det_gemm( + narrow, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], stages=2, **kw ) assert narrow.sk_sync and narrow.sk_slab == 0 assert _nopro(narrow.template) == "dense_proj_det_gemm_nn_n256_m256_s2_sks" with pytest.raises(ValueError, match="sk_slab=3"): - plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk_slab=3, **kw) + cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk_slab=3, **kw) with pytest.raises(ValueError, match="too small"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], stages=2, sk_slab=2, **kw ) - std, *_ = plan_dense_projection_det_gemm( + std, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta_rows=128, **kw ) assert not std.sk_sync and std.sk_slab == 0 and "_sks" not in _nopro(std.template) - off, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk=False, **kw) + off, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=False, **kw) assert ( not off.sk_sync and _nopro(off.template) == "dense_proj_det_gemm_nn_n256_m256" ) # at the test tables' T the synchronised plan is not admitted (a part under SK_MIN_ITERS): the plain tall plan for short_T in (2049, 1001): vs = _views("proj", "indexer_q", "wgrad", "bf16", short_T) - short, *_ = plan_dense_projection_det_gemm(vs["A"], vs["B"], vs["out"], **kw) + short, *_ = cake_dense_gemm_det_train_projection_plan(vs["A"], vs["B"], vs["out"], **kw) assert ( not short.sk_sync and short.sk_slab == 0 @@ -2143,21 +2143,21 @@ def test_round20_rules_plan_like_the_cake_launcher(T): # the table size is asserted by the newest round's test (round 21: 98 = 49 + 49); round 20 shipped 98 = 49 + 49 kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a", _fallback=False) v = _views("proj", "shared_down", "dgrad", "bf16", T) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert _nopro(plan.template) == "dense_proj_det_gemm_kn_n256" and not plan.cta1 assert plan.group_m == 8 assert ( default_group_m(False, True, plan.m_tiles, plan.pair_tiles, plan.sm_pairs) == 16 ) # a caller-forced raster group overrides the rule on the same program - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], group_m=16, **kw ) assert _nopro(forced.template) == _nopro(plan.template) and forced.group_m == 16 # the shared_down forward (X @ W^T, 6144 x 2048) has no ragged rule and keeps the 16-row default on ``kk_n256`` # (one token off the GB T-exact keys at T 16172 / 16231: test_gb_t_exact_rules_plan_their_programs) f = _views("proj", "shared_down", "fwd", "bf16", T + 1) - fwd, *_ = plan_dense_projection_det_gemm(f["A"], f["B"], f["out"], **kw) + fwd, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **kw) assert fwd.group_m == 16 and _nopro(fwd.template) == "dense_proj_det_gemm_kk_n256" @@ -2197,7 +2197,7 @@ def test_round21_rules_plan_like_the_cake_launcher(T): v = _views("proj", "kv_a", "dgrad", "bf16", T) # (the row's eight-warp CTA form is forced off here - its own test is # test_eight_warp_form_rules_plan_their_programs; the ten-warp program below is the row's program where it yields) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], form="", **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **kw) assert ( _nopro(plan.template) == "dense_proj_det_gemm_kn_n256_tma1_pd2" and not plan.cta1 @@ -2212,7 +2212,7 @@ def test_round21_rules_plan_like_the_cake_launcher(T): ) == ("tma", 1, 2, 12, 128, 256) # a caller forcing the former form still plans the round-20 program of the row; the rule's ``pd`` carries over # into a forced form unless the caller sets it (``pd=0``), exactly as in the Cake launcher - old, *_ = plan_dense_projection_det_gemm( + old, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], epi="reg", quad_store=True, stages=6, pd=0, **kw ) assert ( @@ -2223,16 +2223,16 @@ def test_round21_rules_plan_like_the_cake_launcher(T): # the sibling rows keep their programs: the sm_107a kv_a bf16 input gradient (quad stores + in-tile drain), the fp32 kv_a input # gradient (one-slot TMA store, no drain) and the indexer_hw bf16 input gradient (two-slot TMA store + the cross-tile drain) kw107 = dict(sm_count=212, l2_bytes=L2_BYTES, arch="sm_107a", _fallback=False) - p107, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw107) + p107, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw107) assert ( _nopro(p107.template) == "dense_proj_det_gemm_kn_n256_q_pd1" and p107.group_m == 8 ) f = _views("proj", "kv_a", "dgrad", "f32", T) - pf, *_ = plan_dense_projection_det_gemm(f["A"], f["B"], f["out"], **kw) + pf, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **kw) assert _nopro(pf.template) == "dense_proj_det_gemm_kn_n256_f32_tma1" and pf.pd == 0 h = _views("proj", "indexer_hw", "dgrad", "bf16", T) - ph, *_ = plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], **kw) + ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw) assert _nopro(ph.template) == "dense_proj_det_gemm_kn_n256_tma2_pd2" # round 21 (W3, lever B): the sm_100a shared_gate_up forward bf16 (X @ W^T, 2048 x 6144 x T) keeps its raster group 8 # and takes the column-grouped raster form (``kk_n256_cg``, instance-key field 29: the group counts column tiles); @@ -2246,18 +2246,18 @@ def test_round21_rules_plan_like_the_cake_launcher(T): "cgrp": True, } g = _views("proj", "shared_gate_up", "fwd", "bf16", T) - pg, *_ = plan_dense_projection_det_gemm(g["A"], g["B"], g["out"], **kw) + pg, *_ = cake_dense_gemm_det_train_projection_plan(g["A"], g["B"], g["out"], **kw) assert ( pg.cgrp and _nopro(pg.template) == "dense_proj_det_gemm_kk_n256_cg" and pg.group_m == 8 ) - pg107, *_ = plan_dense_projection_det_gemm( + pg107, *_ = cake_dense_gemm_det_train_projection_plan( g["A"], g["B"], g["out"], cgrp=False, **kw107 ) assert not pg107.cgrp and _nopro(pg107.template) == "dense_proj_det_gemm_kk_n256" q = _views("proj", "q_a", "fwd", "bf16", T) - pq, *_ = plan_dense_projection_det_gemm(q["A"], q["B"], q["out"], **kw) + pq, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kw) assert ( pq.cgrp and _nopro(pq.template) == "dense_proj_det_gemm_kk_n256_cg" @@ -2288,7 +2288,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): "cta1": True, } h = _views("proj", "indexer_hw", "dgrad", "bf16", T) - ph, *_ = plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], **kw107) + ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw107) # TEMPLATES_REACHED_NEW (the kernel's rule comment names kn_n256_tma2_pd3_c1: K = 32 <= K_TMA_BF16 keeps the row's two-slot # TMA-store epilogue; the K-less by-kwargs identity render kn_n256_pd3_c1 is the register-epilogue sibling, not this plan) assert _nopro(ph.template) == "dense_proj_det_gemm_kn_n256_tma2_pd3_c1" and ph.cta1 @@ -2301,7 +2301,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): 256, ) # a caller forcing the former form still plans the round-21 program of the row (pd passed explicitly) - old_h, *_ = plan_dense_projection_det_gemm( + old_h, *_ = cake_dense_gemm_det_train_projection_plan( h["A"], h["B"], h["out"], cta1=False, pd=2, **kw107 ) assert ( @@ -2310,9 +2310,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): ) # the control row indexer_k (K 128 is another key) and the sm_100a sibling keep their round-21 programs k_ = _views("proj", "indexer_k", "dgrad", "bf16", T) - pk, *_ = plan_dense_projection_det_gemm(k_["A"], k_["B"], k_["out"], **kw107) + pk, *_ = cake_dense_gemm_det_train_projection_plan(k_["A"], k_["B"], k_["out"], **kw107) assert _nopro(pk.template) == "dense_proj_det_gemm_kn_n256_tma2_pd2" and not pk.cta1 - ph100, *_ = plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], **kw) + ph100, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw) assert ( _nopro(ph100.template) == "dense_proj_det_gemm_kn_n256_tma2_pd2" and not ph100.cta1 @@ -2322,7 +2322,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): fkey = ("sm_107a", False, True, True, False, False, 6144, 32, None) assert ROW_RULES[fkey] == {"block_n": 128, "slots": 1, "group_m": 2, "cta1": True} f = _views("proj", "indexer_hw", "dgrad", "f32", T) - pf, *_ = plan_dense_projection_det_gemm(f["A"], f["B"], f["out"], **kw107) + pf, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **kw107) assert ( _nopro(pf.template) == "dense_proj_det_gemm_kn_n128_f32_tma1_c1" and pf.cta1 ) # TEMPLATES_REACHED_NEW (the kernel's rule comment names kn_n128_f32_tma1_c1) @@ -2333,7 +2333,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): 128, 128, ) - old_f, *_ = plan_dense_projection_det_gemm( + old_f, *_ = cake_dense_gemm_det_train_projection_plan( f["A"], f["B"], f["out"], cta1=False, slots=2, **kw107 ) assert ( @@ -2345,13 +2345,13 @@ def test_round24_rules_plan_like_the_cake_launcher(T): wkey = ("sm_107a", True, True, False, False, False, 2048, None, 6144) assert ROW_RULES[wkey] == {"sk_sync": True, "sk_sync_m": 38, "sk_slab": 2} w = _views("proj", "shared_down", "wgrad", "bf16", T) - pw, *_ = plan_dense_projection_det_gemm(w["A"], w["B"], w["out"], **kw107) + pw, *_ = cake_dense_gemm_det_train_projection_plan(w["A"], w["B"], w["out"], **kw107) assert ( _nopro(pw.template) == "dense_proj_det_gemm_nn_n256_sks_sb2" and pw.sk_sync and pw.sk_slab == 2 ) # TEMPLATES_REACHED_NEW (by-kwargs identity: nn_n256_sks_sb2) - old_w, *_ = plan_dense_projection_det_gemm( + old_w, *_ = cake_dense_gemm_det_train_projection_plan( w["A"], w["B"], w["out"], sk_sync=False, **kw107 ) assert ( @@ -2370,7 +2370,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): ] == {"group_m": 8, "cgrp": True} # unchanged since round 21 for row in ("q_a", "shared_gate_up"): g = _views("proj", row, "fwd", "bf16", T) - pg, *_ = plan_dense_projection_det_gemm(g["A"], g["B"], g["out"], **kw107) + pg, *_ = cake_dense_gemm_det_train_projection_plan(g["A"], g["B"], g["out"], **kw107) # TEMPLATES_REACHED_NEW / the figen tables (dense_proj_det_gemm_kk_n256_cg = the round-21 sm_100a program; the cgrp hint gate stays # ('none', 'none') at 106 pairs - wave_working_set_cgrp(128, 8, 6144, 8, 106) = 22 x 256 x 6144 x 2 B = 69 206 016 B <= L2_BYTES) assert ( @@ -2378,7 +2378,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): and _nopro(pg.template) == "dense_proj_det_gemm_kk_n256_cg" and pg.group_m == 8 ), row - old_g, *_ = plan_dense_projection_det_gemm( + old_g, *_ = cake_dense_gemm_det_train_projection_plan( g["A"], g["B"], g["out"], cgrp=False, **kw107 ) assert ( @@ -2395,7 +2395,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): "group_m": 12, } d = _views("proj", "dense_down", "fwd", "bf16", T) - pd_, *_ = plan_dense_projection_det_gemm(d["A"], d["B"], d["out"], **kw107) + pd_, *_ = cake_dense_gemm_det_train_projection_plan(d["A"], d["B"], d["out"], **kw107) # TEMPLATES_REACHED_NEW: the `_cg_or` tall program WITHOUT `_ht` (sk_parts is the row's tail policy; by-kwargs identity: kk_n256_m256_cg_or) assert ( _nopro(pd_.template) == "dense_proj_det_gemm_kk_n256_m256_cg_or" @@ -2405,7 +2405,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): and not pd_.htail ) assert (pd_.cta_rows, pd_.block_n, pd_.group_m) == (256, 256, 12) - old_d, *_ = plan_dense_projection_det_gemm( + old_d, *_ = cake_dense_gemm_det_train_projection_plan( d["A"], d["B"], d["out"], ovr=False, cgrp=False, **kw107 ) assert ( @@ -2424,7 +2424,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): "group_m": 8, } o = _views("proj", "o_proj", "fwd", "bf16", T) - po, *_ = plan_dense_projection_det_gemm(o["A"], o["B"], o["out"], **kw107) + po, *_ = cake_dense_gemm_det_train_projection_plan(o["A"], o["B"], o["out"], **kw107) # TEMPLATES_REACHED_NEW (the kernel registers dense_projection_det_gemm_kk_m256_ht_or_cg -> symbol ..._m256_ht_cg_or) assert ( _nopro(po.template) == "dense_proj_det_gemm_kk_n256_m256_ht_cg_or" @@ -2434,7 +2434,7 @@ def test_round24_rules_plan_like_the_cake_launcher(T): and not po.ovl ) assert (po.cta_rows, po.block_n, po.group_m) == (256, 256, 8) - old_o, *_ = plan_dense_projection_det_gemm( + old_o, *_ = cake_dense_gemm_det_train_projection_plan( o["A"], o["B"], o["out"], ovr=False, ovl=True, cgrp=False, **kw107 ) assert ( @@ -2443,13 +2443,13 @@ def test_round24_rules_plan_like_the_cake_launcher(T): and not old_o.ovr ) # the sm_100a dense_down / o_proj forwards carry no ovr rule and keep their round-21 programs - pd100, *_ = plan_dense_projection_det_gemm(d["A"], d["B"], d["out"], **kw) + pd100, *_ = cake_dense_gemm_det_train_projection_plan(d["A"], d["B"], d["out"], **kw) assert ( _nopro(pd100.template) == "dense_proj_det_gemm_kk_n256_m256_ov_ht" and not pd100.ovr and not pd100.cgrp ) - po100, *_ = plan_dense_projection_det_gemm(o["A"], o["B"], o["out"], **kw) + po100, *_ = cake_dense_gemm_det_train_projection_plan(o["A"], o["B"], o["out"], **kw) assert ( _nopro(po100.template) == "dense_proj_det_gemm_kk_n256_m256_ov_ht" and not po100.ovr @@ -2482,7 +2482,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): "store_ef": True, } sdf = _views("proj", "shared_down", "wgrad", "f32", T) - p_sdf, *_ = plan_dense_projection_det_gemm(sdf["A"], sdf["B"], sdf["out"], **kw107) + p_sdf, *_ = cake_dense_gemm_det_train_projection_plan(sdf["A"], sdf["B"], sdf["out"], **kw107) assert ( _nopro(p_sdf.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2" ) # ROUTE_TEMPLATE_NEW of the round-25 templates step = TEMPLATES_REACHED_NEW (nn_n256_f32_v8_ef_sks_sb2, the round-25 sweep symbol) @@ -2496,7 +2496,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): assert (p_sdf.cta_rows, p_sdf.block_n) == (128, 256) and not p_sdf.htail # a caller forcing the former tall form plans the lock-e9d8042b program of the row (the forced cta_rows leaves the rule's 128-row family, # so the planner drops the rule's epilogue / tail knobs; sk_sync and the TMA-store epilogue are passed explicitly all the same) - old_sdf, *_ = plan_dense_projection_det_gemm( + old_sdf, *_ = cake_dense_gemm_det_train_projection_plan( sdf["A"], sdf["B"], sdf["out"], cta_rows=256, sk_sync=False, epi="tma", **kw107 ) assert ( @@ -2517,7 +2517,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): } for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "f32", T) - p, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw107) + p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw107) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2", ( row ) # ROUTE_TEMPLATE_NEW of the round-25 templates step = TEMPLATES_REACHED_NEW (nn_n256_f32_v8_ef_sks_sb2) @@ -2529,7 +2529,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): and p.store_ef and p.cta_rows == 128 ), row - old, *_ = plan_dense_projection_det_gemm( + old, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta_rows=256, sk_sync=False, epi="tma", **kw107 ) assert ( @@ -2544,12 +2544,12 @@ def test_round25_rules_plan_like_the_cake_launcher(T): assert ROW_RULES[qab_key] == {"sk_sync": True, "sk_sync_m": 26, "sk_slab": 2} for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "bf16", T) - p, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw107) + p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw107) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_sks_sb2", ( row ) # ROUTE_TEMPLATE_NEW of the round-25 templates step (nn_n256_sks_sb2 = the round-24 sm_107a shared_down wgrad bf16 program, no new template) assert p.sk_sync and p.sk_slab == 2 and p.cta_rows == 128, row - old, *_ = plan_dense_projection_det_gemm( + old, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk_slab=0, **kw107 ) assert ( @@ -2562,7 +2562,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): ("sm_107a", True, True, False, False, False, 2048, None, 6144) ] == {"sk_sync": True, "sk_sync_m": 38, "sk_slab": 2} sdb107 = _views("proj", "shared_down", "wgrad", "bf16", T) - p_ctl, *_ = plan_dense_projection_det_gemm( + p_ctl, *_ = cake_dense_gemm_det_train_projection_plan( sdb107["A"], sdb107["B"], sdb107["out"], **kw107 ) assert ( @@ -2580,7 +2580,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): "sk_sync_m": 26, "sk_slab": 2, } - p_sdb, *_ = plan_dense_projection_det_gemm( + p_sdb, *_ = cake_dense_gemm_det_train_projection_plan( sdb107["A"], sdb107["B"], sdb107["out"], **kw ) assert ( @@ -2593,7 +2593,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): and p_sdb.group_m == 4 and p_sdb.cta_rows == 128 ) - old_sdb, *_ = plan_dense_projection_det_gemm( + old_sdb, *_ = cake_dense_gemm_det_train_projection_plan( sdb107["A"], sdb107["B"], sdb107["out"], sk_sync=False, sk_exact=3, **kw ) assert ( @@ -2607,12 +2607,12 @@ def test_round25_rules_plan_like_the_cake_launcher(T): assert ROW_RULES[qab100_key] == {"sk_sync": True, "sk_sync_m": 26, "sk_slab": 2} for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "bf16", T) - p, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_sks_sb2", row assert p.sk_sync and p.sk_slab == 2 and p.sk_exact == 0 and p.cta_rows == 128, ( row ) - old, *_ = plan_dense_projection_det_gemm( + old, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk_sync=False, sk_exact=3, **kw ) assert ( @@ -2632,7 +2632,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): "sk_sync_m": 26, "sk_slab": 2, } - p_sdf100, *_ = plan_dense_projection_det_gemm(sdf["A"], sdf["B"], sdf["out"], **kw) + p_sdf100, *_ = cake_dense_gemm_det_train_projection_plan(sdf["A"], sdf["B"], sdf["out"], **kw) assert ( _nopro(p_sdf100.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2" ) # ROUTE_TEMPLATE_NEW of the templates step = TEMPLATES_REACHED_NEW (predicted nn_n256_f32_v8_ef_sks_sb2 - the same new program as the sm_107a fp32 rows; confirm per route) @@ -2644,7 +2644,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): and p_sdf100.store_ef and p_sdf100.group_m == 8 ) - old_sdf100, *_ = plan_dense_projection_det_gemm( + old_sdf100, *_ = cake_dense_gemm_det_train_projection_plan( sdf["A"], sdf["B"], sdf["out"], sk_sync=False, sk_exact=3, **kw ) assert ( @@ -2667,7 +2667,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): } for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "f32", T) - p, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2", ( row ) @@ -2679,7 +2679,7 @@ def test_round25_rules_plan_like_the_cake_launcher(T): and p.store_ef and p.group_m == 32 ), row - old, *_ = plan_dense_projection_det_gemm( + old, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk_sync=False, sk_exact=3, **kw ) assert ( @@ -2716,7 +2716,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): "group_m": 12, } g = _views("proj", "dense_gate_up", "fwd", "bf16", T) - pg, *_ = plan_dense_projection_det_gemm(g["A"], g["B"], g["out"], **kw107) + pg, *_ = cake_dense_gemm_det_train_projection_plan(g["A"], g["B"], g["out"], **kw107) assert ( _nopro(pg.template) == "dense_proj_det_gemm_kk_n256_m256_ht_cg_or" and pg.ovr @@ -2725,7 +2725,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): and not pg.ovl ) # ROUTE_TEMPLATE_NEW of the templates step (predicted kk_n256_m256_ht_cg_or - the sweep symbol) assert (pg.cta_rows, pg.block_n, pg.group_m) == (256, 256, 12) - old_g, *_ = plan_dense_projection_det_gemm( + old_g, *_ = cake_dense_gemm_det_train_projection_plan( g["A"], g["B"], g["out"], @@ -2752,7 +2752,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): "store_ef": True, } wf = _views("proj", "dense_down", "wgrad", "f32", T) - pwf, *_ = plan_dense_projection_det_gemm(wf["A"], wf["B"], wf["out"], **kw107) + pwf, *_ = cake_dense_gemm_det_train_projection_plan(wf["A"], wf["B"], wf["out"], **kw107) assert ( _nopro(pwf.template) == "dense_proj_det_gemm_nn_n256_m256_f32_v8_ef_ht" and pwf.htail @@ -2764,7 +2764,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): and not pwf.sk_sync and not pwf.sk_exact ) - old_wf, *_ = plan_dense_projection_det_gemm( + old_wf, *_ = cake_dense_gemm_det_train_projection_plan( wf["A"], wf["B"], wf["out"], cta_rows=128, htail=False, epi="tma", **kw107 ) assert ( @@ -2777,14 +2777,14 @@ def test_round26_rules_plan_like_the_cake_launcher(T): wbkey = ("sm_107a", True, True, False, False, False, 12288, None, 6144) assert ROW_RULES[wbkey] == {"cta_rows": 256, "htail": True} wb = _views("proj", "dense_down", "wgrad", "bf16", T) - pwb, *_ = plan_dense_projection_det_gemm(wb["A"], wb["B"], wb["out"], **kw107) + pwb, *_ = cake_dense_gemm_det_train_projection_plan(wb["A"], wb["B"], wb["out"], **kw107) assert ( _nopro(pwb.template) == "dense_proj_det_gemm_nn_n256_m256_ht" and pwb.htail and not pwb.sk_sync ) # ROUTE_TEMPLATE_NEW of the templates step (predicted nn_n256_m256_ht - a template the previous lock already declares) assert (pwb.cta_rows, pwb.block_n) == (256, 256) - old_wb, *_ = plan_dense_projection_det_gemm( + old_wb, *_ = cake_dense_gemm_det_train_projection_plan( wb["A"], wb["B"], wb["out"], cta_rows=128, htail=False, **kw107 ) assert ( @@ -2804,7 +2804,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): "group_m": 12, } dg = _views("proj", "dense_gate_up", "dgrad", "bf16", T) - pdg, *_ = plan_dense_projection_det_gemm(dg["A"], dg["B"], dg["out"], **kw107) + pdg, *_ = cake_dense_gemm_det_train_projection_plan(dg["A"], dg["B"], dg["out"], **kw107) assert ( _nopro(pdg.template) == "dense_proj_det_gemm_kn_n256_m256_cg_or" and pdg.ovr @@ -2813,7 +2813,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): and not pdg.htail ) # ROUTE_TEMPLATE_NEW of the templates step (predicted kn_n256_m256_cg_or) assert (pdg.cta_rows, pdg.block_n, pdg.group_m) == (256, 256, 12) - old_dg, *_ = plan_dense_projection_det_gemm( + old_dg, *_ = cake_dense_gemm_det_train_projection_plan( dg["A"], dg["B"], dg["out"], cta_rows=128, ovr=False, cgrp=False, **kw107 ) assert ( @@ -2826,7 +2826,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): qkey = ("sm_107a", False, True, False, False, False, 2048, 16384, None) assert ROW_RULES[qkey] == {"cgrp": True, "group_m": 8} q = _views("proj", "q_b", "dgrad", "bf16", T) - pq, *_ = plan_dense_projection_det_gemm(q["A"], q["B"], q["out"], **kw107) + pq, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kw107) assert ( _nopro(pq.template) == "dense_proj_det_gemm_kn_n256_cg" and pq.cgrp @@ -2834,7 +2834,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): and not pq.htail ) # ROUTE_TEMPLATE_NEW of the templates step (predicted kn_n256_cg) assert (pq.cta_rows, pq.block_n, pq.group_m) == (128, 256, 8) - old_q, *_ = plan_dense_projection_det_gemm( + old_q, *_ = cake_dense_gemm_det_train_projection_plan( q["A"], q["B"], q["out"], cgrp=False, **kw107 ) assert ( @@ -2845,14 +2845,14 @@ def test_round26_rules_plan_like_the_cake_launcher(T): gwkey = ("sm_107a", True, True, True, False, False, 6144, None, 12288) assert ROW_RULES[gwkey] == {"cta_rows": 256, "htail": True} gw = _views("proj", "dense_gate_up", "wgrad", "f32", T) - pgw, *_ = plan_dense_projection_det_gemm(gw["A"], gw["B"], gw["out"], **kw107) + pgw, *_ = cake_dense_gemm_det_train_projection_plan(gw["A"], gw["B"], gw["out"], **kw107) assert ( _nopro(pgw.template) == "dense_proj_det_gemm_nn_n256_m256_f32_ht_tma1" and pgw.htail and pgw.epi == "tma" ) # ROUTE_TEMPLATE_NEW of the templates step (predicted nn_n256_m256_f32_ht_tma1 - a template the previous lock already declares; the rule names no epi, so the fp32-out tile takes the planner's default TMA-store epilogue like the sm_100a twin's nn_n256_m256_f32_tma1) assert (pgw.cta_rows, pgw.block_n) == (256, 256) and not pgw.sk_sync - old_gw, *_ = plan_dense_projection_det_gemm( + old_gw, *_ = cake_dense_gemm_det_train_projection_plan( gw["A"], gw["B"], gw["out"], @@ -2889,7 +2889,7 @@ def test_round26_rules_plan_like_the_cake_launcher(T): ), ): v = _views("proj", row, op, dt, T) - p100, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + p100, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert _nopro(p100.template) == tmpl and not p100.cgrp and not p100.ovr, ( row, op, @@ -2913,7 +2913,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): akey = ("sm_100a", False, True, True, False, False, 6144, 128, None) assert ROW_RULES[akey] == {"block_n": 128, "group_m": 2} a = _views("proj", "indexer_k", "dgrad", "f32", T) - pa, *_ = plan_dense_projection_det_gemm(a["A"], a["B"], a["out"], **kw) + pa, *_ = cake_dense_gemm_det_train_projection_plan(a["A"], a["B"], a["out"], **kw) assert ( _nopro(pa.template) == "dense_proj_det_gemm_kn_n128_f32_tma1" and pa.group_m == 2 @@ -2926,7 +2926,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): # drain; the previous lock planned the pair form dense_proj_det_gemm_kn_n128_f32_tma1 ---- bkey = ("sm_107a", False, True, True, False, False, 6144, 128, None) assert ROW_RULES[bkey] == {"block_n": 128, "cta1": True, "group_m": 2, "pd": 2} - pb, *_ = plan_dense_projection_det_gemm(a["A"], a["B"], a["out"], **kw107) + pb, *_ = cake_dense_gemm_det_train_projection_plan(a["A"], a["B"], a["out"], **kw107) assert ( _nopro(pb.template) == "dense_proj_det_gemm_kn_n128_f32_tma1_pd2_c1" and pb.cta1 @@ -2934,7 +2934,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): and pb.group_m == 2 and pb.block_n == 128 ) # ROUTE_TEMPLATE_NEW of the templates step (predicted kn_n128_f32_tma1_pd2_c1) - old_b, *_ = plan_dense_projection_det_gemm( + old_b, *_ = cake_dense_gemm_det_train_projection_plan( a["A"], a["B"], a["out"], cta1=False, pd=0, **kw107 ) assert ( @@ -2956,7 +2956,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): "pro": 15, } c = _views("proj", "kv_a", "wgrad", "bf16", T) - pc, *_ = plan_dense_projection_det_gemm( + pc, *_ = cake_dense_gemm_det_train_projection_plan( c["A"], c["B"], c["out"], transposed_out=c["transposed"], **kw107 ) assert ( @@ -2972,7 +2972,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): 128, 8, ) and pc.hints == ("evict_first", "evict_first") - old_c, *_ = plan_dense_projection_det_gemm( + old_c, *_ = cake_dense_gemm_det_train_projection_plan( c["A"], c["B"], c["out"], transposed_out=c["transposed"], stt=False, **kw107 ) assert ( @@ -2992,7 +2992,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): ) # the plan arithmetic is the pair form's # controls: the sm_100a kv_a weight gradient bf16 keeps its TMA-store epilogue program (no stt rule on sm_100a), and the fp32-out twins of the # kv_a weight gradient carry no stt on either architecture (stmatrix is .b16) - p100, *_ = plan_dense_projection_det_gemm( + p100, *_ = cake_dense_gemm_det_train_projection_plan( c["A"], c["B"], c["out"], transposed_out=c["transposed"], **kw ) assert ( @@ -3002,7 +3002,7 @@ def test_round27_rules_plan_like_the_cake_launcher(T): ) f = _views("proj", "kv_a", "wgrad", "f32", T) for kwx in (kw, kw107): - pf, *_ = plan_dense_projection_det_gemm( + pf, *_ = cake_dense_gemm_det_train_projection_plan( f["A"], f["B"], f["out"], transposed_out=f["transposed"], **kwx ) assert not pf.stt and pf.out_f32 @@ -3072,7 +3072,7 @@ def test_round27_stt_knob(T): # a row whose output is fp32 raises like the Cake launcher kw107 = dict(sm_count=212, l2_bytes=L2_BYTES, arch="sm_107a", _fallback=False) v = _views("proj", "kv_a", "wgrad", "bf16", T) - p, *_ = plan_dense_projection_det_gemm( + p, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=v["transposed"], **kw107 ) assert ( @@ -3081,7 +3081,7 @@ def test_round27_stt_knob(T): and p.epi == "reg" and _nopro(p.template).endswith("_t_sks_stt") ) - q, *_ = plan_dense_projection_det_gemm( + q, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=v["transposed"], stt=False, **kw107 ) assert ( @@ -3092,7 +3092,7 @@ def test_round27_stt_knob(T): ) f = _views("proj", "kv_a", "wgrad", "f32", T) with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( f["A"], f["B"], f["out"], transposed_out=f["transposed"], stt=True, **kw107 ) # a caller-forced form outside the staging's admission on the stt row keeps that form WITHOUT the staging (the rule's stt does not carry @@ -3101,7 +3101,7 @@ def test_round27_stt_knob(T): # pd (the pipelined TMEM drain) is not a caller form of this row at all: it needs a row-major output and a 128 / 256-column # tile, so the planner rejects it here with or without stt, like the Cake launcher with pytest.raises(ValueError, match="pd needs"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=v["transposed"], pd=2, **kw107 ) for forced in ( @@ -3109,7 +3109,7 @@ def test_round27_stt_knob(T): dict(sk_sync=True, sk_slab=2), dict(cta_rows=256), ): - r, *_ = plan_dense_projection_det_gemm( + r, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=v["transposed"], **forced, **kw107 ) assert ( @@ -3118,7 +3118,7 @@ def test_round27_stt_knob(T): and r.tma_out == (r.epi == "tma") ), forced with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -3339,30 +3339,30 @@ def test_round33_rules_plan_like_the_cake_launcher(T): # architectures; a caller-forced pro 0 is the previous lock's program with the same plan arithmetic ---- q = _views("proj", "indexer_q", "fwd", "bf16", T) for kwx in (kw, kw107): - p, *_ = plan_dense_projection_det_gemm(q["A"], q["B"], q["out"], **kwx) + p, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kwx) assert p.pro == 15 and p.template.endswith("_pro15") and not p.cta1, kwx["arch"] - old, *_ = plan_dense_projection_det_gemm(q["A"], q["B"], q["out"], pro=0, **kwx) + old, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], pro=0, **kwx) assert ( old.pro == 0 and old.template == p.template[: -len("_pro15")] and old.grid == p.grid and old.stages == p.stages ), kwx["arch"] - p107, *_ = plan_dense_projection_det_gemm(q["A"], q["B"], q["out"], **kw107) + p107, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kw107) assert ( p107.template == "dense_proj_det_gemm_kk_n256_pro15" ) # ROUTE_TEMPLATE_NEW of the templates step (predicted kk_n256_pro15) # the shared_down input gradient bf16: sm_100a pins pro 0 (the round-27 program, no suffix), sm_107a has no rule (the default pro 7) s = _views("proj", "shared_down", "dgrad", "bf16", T) - ps, *_ = plan_dense_projection_det_gemm(s["A"], s["B"], s["out"], **kw) + ps, *_ = cake_dense_gemm_det_train_projection_plan(s["A"], s["B"], s["out"], **kw) assert ( ps.pro == 0 and not re.search(r"_pro\d+$", ps.template) and ps.group_m == 8 ) # NOT `"_pro" not in`: `dense_proj` contains it - ps107, *_ = plan_dense_projection_det_gemm(s["A"], s["B"], s["out"], **kw107) + ps107, *_ = cake_dense_gemm_det_train_projection_plan(s["A"], s["B"], s["out"], **kw107) assert ps107.pro == 7 and ps107.template.endswith("_pro7") # the q_b weight gradient fp32-out on sm_107a: pro 15 on the register epilogue with the 256-bit stores (the round-27 knobs) w = _views("proj", "q_b", "wgrad", "f32", T) - pw, *_ = plan_dense_projection_det_gemm( + pw, *_ = cake_dense_gemm_det_train_projection_plan( w["A"], w["B"], w["out"], transposed_out=w["transposed"], **kw107 ) assert ( @@ -3376,7 +3376,7 @@ def test_round33_rules_plan_like_the_cake_launcher(T): pw.template == "dense_proj_det_gemm_nn_n256_f32_v8_pro15" ) # ROUTE_TEMPLATE_NEW of the templates step (predicted nn_n256_f32_v8_pro15) # a caller-forced tile family drops the rule's pro with the rule's other family-bound knobs (the launcher's filter): pro 0, no suffix - tall_w, *_ = plan_dense_projection_det_gemm( + tall_w, *_ = cake_dense_gemm_det_train_projection_plan( w["A"], w["B"], w["out"], transposed_out=w["transposed"], cta_rows=256, **kw107 ) assert ( @@ -3387,7 +3387,7 @@ def test_round33_rules_plan_like_the_cake_launcher(T): # the kv_a weight gradient fp32-out on sm_107a: the synchronised stream-K program takes the bulk slab read (sk_slab 2) and the default # pro 7; a caller-forced sk_slab 0 / pro 0 is the previous lock's program with the same plan arithmetic f = _views("proj", "kv_a", "wgrad", "f32", T) - pf, *_ = plan_dense_projection_det_gemm( + pf, *_ = cake_dense_gemm_det_train_projection_plan( f["A"], f["B"], f["out"], transposed_out=f["transposed"], **kw107 ) assert ( @@ -3400,7 +3400,7 @@ def test_round33_rules_plan_like_the_cake_launcher(T): assert ( pf.template == "dense_proj_det_gemm_nn_n192_bz64_hee_f32_t_sks_sb2_pro7" ) # ROUTE_TEMPLATE_NEW of the templates step (predicted nn_n192_bz64_hee_f32_t_sks_sb2_pro7) - old_f, *_ = plan_dense_projection_det_gemm( + old_f, *_ = cake_dense_gemm_det_train_projection_plan( f["A"], f["B"], f["out"], @@ -3416,15 +3416,15 @@ def test_round33_rules_plan_like_the_cake_launcher(T): # a single-CTA row: the indexer_hw input gradient fp32-out on sm_107a keeps its round-19 program (no pro; the merged rule's default is # zeroed), a caller-forced pro raises like the Cake launcher, and a caller-forced pair form takes the architecture default h = _views("proj", "indexer_hw", "dgrad", "f32", T) - ph, *_ = plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], **kw107) + ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw107) assert ( ph.cta1 and ph.pro == 0 and ph.template == "dense_proj_det_gemm_kn_n128_f32_tma1_c1" ) with pytest.raises(ValueError, match="needs the cluster pair"): - plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], pro=7, **kw107) - pair_h, *_ = plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], pro=7, **kw107) + pair_h, *_ = cake_dense_gemm_det_train_projection_plan( h["A"], h["B"], h["out"], cta1=False, **kw107 ) assert not pair_h.cta1 and pair_h.pro == 7 and pair_h.template.endswith("_pro7") @@ -3510,17 +3510,17 @@ def test_round33_pro_knob(T): kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a", _fallback=False) kw107 = dict(sm_count=212, l2_bytes=L2_BYTES, arch="sm_107a", _fallback=False) g = _views("proj", "dense_gate_up", "fwd", "bf16", T) - pg, *_ = plan_dense_projection_det_gemm(g["A"], g["B"], g["out"], **kw) + pg, *_ = cake_dense_gemm_det_train_projection_plan(g["A"], g["B"], g["out"], **kw) assert ( pg.pro == 7 and pg.template.endswith("_pro7") and not pg.rule_fallback ) # no sm_100a rule: the default alone w = _views("proj", "q_b", "wgrad", "f32", T) - pw, *_ = plan_dense_projection_det_gemm( + pw, *_ = cake_dense_gemm_det_train_projection_plan( w["A"], w["B"], w["out"], transposed_out=w["transposed"], **kw107 ) assert pw.pro == 15 and pw.template.endswith("_pro15") for forced in (0, 7, 8, 15): - pfx, *_ = plan_dense_projection_det_gemm( + pfx, *_ = cake_dense_gemm_det_train_projection_plan( w["A"], w["B"], w["out"], @@ -3534,7 +3534,7 @@ def test_round33_pro_knob(T): assert pfx.grid == pw.grid and pfx.stages == pw.stages, ( forced ) # the program changes, the plan arithmetic does not - nr, *_ = plan_dense_projection_det_gemm( + nr, *_ = cake_dense_gemm_det_train_projection_plan( w["A"], w["B"], w["out"], @@ -3549,14 +3549,14 @@ def test_round33_pro_knob(T): and "_f32_v8" not in nr.template ) h = _views("proj", "indexer_hw", "dgrad", "f32", T) - ph, *_ = plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], **kw107) + ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw107) assert ( ph.cta1 and ph.pro == 0 and ph.template == "dense_proj_det_gemm_kn_n128_f32_tma1_c1" ) with pytest.raises(ValueError, match="needs the cluster pair"): - plan_dense_projection_det_gemm(h["A"], h["B"], h["out"], pro=15, **kw107) + cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], pro=15, **kw107) @pytest.mark.parametrize("T", [16172, 16231]) @@ -3633,7 +3633,7 @@ def test_round21_cgrp_knob(T): # pair tiles on 74 pairs, the half-height tail wave) is the pair form's - only the program and the raster differ kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a", _fallback=False) v = _views("proj", "dense_down", "fwd", "bf16", T) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=True, group_m=12, **kw ) assert ( @@ -3655,7 +3655,7 @@ def test_round21_cgrp_knob(T): == wave_working_set_cgrp(64, 24, 12288, 12, 74) == 19 * 256 * 12288 * 2 ) - base, *_ = plan_dense_projection_det_gemm( + base, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], group_m=12, **kw ) assert ( @@ -3680,27 +3680,27 @@ def test_round21_cgrp_knob(T): ): assert getattr(plan, f) == getattr(base, f), f # the raster width defaults to the row default (16, ``default_group_m``) under cgrp as well, now as column tiles - cg16, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], cgrp=True, **kw) + cg16, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], cgrp=True, **kw) assert cg16.group_m == 16 and _nopro(cg16.template) == _nopro(plan.template) # the pair invariant holds whatever the column count, so the count has no parity constraint under cgrp: 3 (or 1) # column tiles per group are admitted where the row raster rejects an odd count (CTA pairs are adjacent row # tiles); a count < 1 is rejected on both - odd, *_ = plan_dense_projection_det_gemm( + odd, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=True, group_m=3, **kw ) assert odd.group_m == 3 and _nopro(odd.template) == _nopro(plan.template) - one, *_ = plan_dense_projection_det_gemm( + one, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=True, group_m=1, **kw ) assert one.group_m == 1 and _nopro(one.template) == _nopro(plan.template) with pytest.raises(ValueError, match="group_m must be an even number"): - plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], group_m=3, **kw) + cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], group_m=3, **kw) with pytest.raises(ValueError, match="group_m must be an even number"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=False, group_m=1, **kw ) with pytest.raises(ValueError, match="group_m must be an even number"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=True, group_m=0, **kw ) # the planner selects the gate by the raster: one 256-column tile over 25600 rows at K = 3200 (no rule) streams A @@ -3708,8 +3708,8 @@ def test_round21_cgrp_knob(T): A = torch.empty(25600, 3200, dtype=torch.bfloat16) Bt = torch.empty(256, 3200, dtype=torch.bfloat16).t() out = torch.empty(25600, 256, dtype=torch.bfloat16) - row, *_ = plan_dense_projection_det_gemm(A, Bt, out, **kw) - col, *_ = plan_dense_projection_det_gemm(A, Bt, out, cgrp=True, **kw) + row, *_ = cake_dense_gemm_det_train_projection_plan(A, Bt, out, **kw) + col, *_ = cake_dense_gemm_det_train_projection_plan(A, Bt, out, cgrp=True, **kw) assert (row.m_tiles, row.n_tiles, row.group_m, row.sm_pairs) == (200, 1, 16, 74) assert _nopro(row.template) == "dense_proj_det_gemm_kk_n256_hen" and row.hints == ( "evict_first", @@ -3777,14 +3777,14 @@ def test_round24_ovr_knob(T): # launcher's round-24 precedence rule, asserted below on the sm_107a o_proj / dense_down forwards) kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a", _fallback=False) v = _views("proj", "o_proj", "fwd", "bf16", T) - base, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + base, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert ( base.ovl and base.htail and not base.ovr and _nopro(base.template) == "dense_proj_det_gemm_kk_n256_m256_ov_ht" ) - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], ovr=True, **kw ) assert ( @@ -3812,7 +3812,7 @@ def test_round24_ovr_knob(T): ): assert getattr(forced, f) == getattr(base, f), f with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], ovr=True, ovl=True, **kw ) # (2) the reverse direction (Cake 95da40d6e8 L2745-L2746, `if ovr and ovl and not rule_ovl: ovr = False` inside the rule read of ovr): a @@ -3821,7 +3821,7 @@ def test_round24_ovr_knob(T): # without this mirror every forced-ovl planner call on these rows raises in FlashInfer while the Cake launcher plans the `_ov` program kw107 = dict(sm_count=212, l2_bytes=L2_BYTES, arch="sm_107a", _fallback=False) o107 = _views("proj", "o_proj", "fwd", "bf16", T) - forced_ovl, *_ = plan_dense_projection_det_gemm( + forced_ovl, *_ = cake_dense_gemm_det_train_projection_plan( o107["A"], o107["B"], o107["out"], ovl=True, **kw107 ) assert ( @@ -3834,7 +3834,7 @@ def test_round24_ovr_knob(T): assert _nopro(forced_ovl.template) == "dense_proj_det_gemm_kk_n256_m256_ov_ht_cg" # the same on the dense_down fwd row (rule {cta_rows 256, sk_parts 3, ovr, cgrp, group_m 12}: no htail) -> the `_ov_cg` tall program d107 = _views("proj", "dense_down", "fwd", "bf16", T) - forced_ovl_d, *_ = plan_dense_projection_det_gemm( + forced_ovl_d, *_ = cake_dense_gemm_det_train_projection_plan( d107["A"], d107["B"], d107["out"], ovl=True, **kw107 ) assert ( @@ -3846,7 +3846,7 @@ def test_round24_ovr_knob(T): ) assert _nopro(forced_ovl_d.template) == "dense_proj_det_gemm_kk_n256_m256_ov_cg" with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], ovr=True, cta_rows=128, **kw ) @@ -4006,7 +4006,7 @@ def test_instance_key_rejects_bad_configurations(): # the promotion is validated where it is resolved, by the planner v = _views("proj", "o_proj", "fwd", "bf16", 257) with pytest.raises(ValueError, match="promo"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -4119,7 +4119,7 @@ def test_check_stream_k_plan_refuses_a_plan_outside_the_launch_contract(): # the dgrad of T = 1024, N = 1032, K = 1000 on 76 CTA pairs: 16 tail tiles of 17 K steps, one partial wave E = lambda *shape: torch.empty(*shape, dtype=torch.bfloat16) # noqa: E731 kw = dict(sm_count=152, l2_bytes=L2_BYTES, _fallback=False) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( E(1024, 1032), E(1032, 1000), E(1024, 1000), pd=2, slots=2, epi="tma", **kw ) assert plan.pd == 2 and (plan.tail_tiles, plan.k_blocks) == (16, 17) @@ -4220,7 +4220,7 @@ def test_stream_k_workspace_holds_every_slab_the_program_reaches(): # program reaches slab 3 * 69 + 1 = 208, past the sk_units + tail_tiles = 146 slabs of the unit-keyed sizing assert stream_k_plan(70, 65, 76, True) == (0, 70, 76, 60) E = lambda *shape: torch.empty(*shape, dtype=torch.bfloat16) # noqa: E731 - base, *_ = plan_dense_projection_det_gemm( + base, *_ = cake_dense_gemm_det_train_projection_plan( E(1024, 1032), E(1032, 1000), E(1024, 1000), sm_count=152, l2_bytes=L2_BYTES ) review = dataclasses.replace( @@ -4238,7 +4238,7 @@ def test_stream_k_workspace_holds_every_slab_the_program_reaches(): assert _slabs_reached(review, True) == 208 and _slabs_reached(review, False) == 144 assert review.ws_f32_elems == WS_PARTS * 70 * slab # public-API plans of the planner on the four target devices: the weight gradient with the linearised policy - # (projection_det_wgrad's sk / cta_rows), and forward / input gradients with prepare's stream-K knobs; every plan + # (cake_dense_gemm_det_train_projection_wgrad's sk / cta_rows), and forward / input gradients with prepare's stream-K knobs; every plan # the planner accepts fits its workspace, and a tile-keyed program never gets more than WS_PARTS units on a tile checked, short_before, refused = 0, 0, 0 for arch, sm_count in ( @@ -4273,7 +4273,7 @@ def test_stream_k_workspace_holds_every_slab_the_program_reaches(): cases.append((E(T, N), E(N, K), E(T, K), knobs)) for A, B, out, knobs in cases: try: - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( A, B, out, sm_count=sm_count, l2_bytes=L2_BYTES, arch=arch, **knobs ) except ValueError as e: @@ -4301,7 +4301,7 @@ def test_planner_defaults_to_auto_stream_k_and_guards_the_slice_counters(): # round-9 sm_100a rule plans this row's 64-row family (48 pair tiles, which 74 pairs cannot split two ways), so # the 128-row family is forced here to keep the split under test v = _views("proj", "indexer_hw", "wgrad", "bf16", 1001) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -4321,7 +4321,7 @@ def test_planner_defaults_to_auto_stream_k_and_guards_the_slice_counters(): assert plan.grid == (96, 1, 1) and plan.sk_iters == 24 * 16 assert plan.ws_f32_elems == (48 + 24) * 2 * 128 * 128 assert plan.counters_u32 == 24 * 16 and plan.counters_alloc_u32 == 8192 - off, *_ = plan_dense_projection_det_gemm( + off, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -4350,10 +4350,10 @@ def test_planner_defaults_to_auto_stream_k_and_guards_the_slice_counters(): out = torch.empty(256, dtype=torch.bfloat16).as_strided((300 * 256, 256), (0, 1)) assert SK_DUMMY_BASE == 4096 with pytest.raises(ValueError, match="slice-counter budget"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A, W.t(), out, sm_count=1000, l2_bytes=L2_BYTES, _fallback=False, sk=True ) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( A, W.t(), out, sm_count=1000, l2_bytes=L2_BYTES, _fallback=False ) # auto: 2 * 300 > 500 pairs -> whole tiles assert (plan.num_full, plan.tail_tiles, plan.sk_units) == (300, 0, 0) @@ -4444,7 +4444,7 @@ def test_projection_rows_plan_like_the_cake_launcher(sm_count, T): for op in ("fwd", "dgrad", "wgrad"): for out_dtype in ("bf16", "f32") if op != "fwd" else ("bf16",): v = _views("proj", row, op, out_dtype, T) - plan, a_desc, b_desc, out3 = plan_dense_projection_det_gemm( + plan, a_desc, b_desc, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -4805,7 +4805,7 @@ def test_mla_rows_plan_as_batched_views(T, sm_count): for row in MLA_ROWS: for op in ("fwd", "dgrad", "wgrad"): v = _views("mla", row, op, "bf16", T) - plan, a_desc, b_desc, out3 = plan_dense_projection_det_gemm( + plan, a_desc, b_desc, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -4849,7 +4849,7 @@ def test_mla_rows_plan_as_batched_views(T, sm_count): # qabs dgrad writes the first 192 columns of a 256-wide slot of a token-major [T, H, 256] buffer: # the [H, T, D] output view has batch stride 256 (heads) and row stride 64 * 256 (tokens) v = _views("mla", "qabs", "dgrad", "bf16", 33) - plan, _, _, out3 = plan_dense_projection_det_gemm( + plan, _, _, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=148, l2_bytes=L2_BYTES ) assert (out3.stride(0), out3.stride(1), out3.stride(2)) == (256, 64 * 256, 1) @@ -4893,11 +4893,11 @@ def test_operand_view_classification_and_rejections(): A2 = torch.empty(16, 64, dtype=torch.bfloat16) B2 = torch.empty(64, 24, dtype=torch.bfloat16) with pytest.raises(ValueError, match="bf16"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A2.float(), B2, torch.empty(16, 24), sm_count=148, l2_bytes=L2_BYTES ) with pytest.raises(ValueError, match="bf16 or fp32"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A2, B2, torch.empty(16, 24, dtype=torch.float16), @@ -4905,7 +4905,7 @@ def test_operand_view_classification_and_rejections(): l2_bytes=L2_BYTES, ) with pytest.raises(ValueError, match=r"B must be \[L=1, K=64, N\]"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A2, torch.empty(32, 24, dtype=torch.bfloat16), torch.empty(16, 24), @@ -4913,11 +4913,11 @@ def test_operand_view_classification_and_rejections(): l2_bytes=L2_BYTES, ) with pytest.raises(ValueError, match="out must be"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A2, B2, torch.empty(24, 16), sm_count=148, l2_bytes=L2_BYTES ) with pytest.raises(ValueError, match="multiple of 8"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A2, torch.empty(64, 20, dtype=torch.bfloat16), torch.empty(16, 20), @@ -4925,11 +4925,11 @@ def test_operand_view_classification_and_rejections(): l2_bytes=L2_BYTES, ) with pytest.raises(ValueError, match="unit inner stride"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( A2, B2, torch.empty(24, 16).t(), sm_count=148, l2_bytes=L2_BYTES ) # transposed output: [N, M] with unit stride on m - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( A2, B2, torch.empty(24, 16), @@ -4975,7 +4975,7 @@ def test_batched_small_m_swap_and_row_rules(): out = torch.empty(H, M, N, dtype=torch.bfloat16) assert swap_small_m(H, M, N, False) and not swap_small_m(1, M, N, False) assert not swap_small_m(H, 512, 192, False) and not swap_small_m(H, M, N, True) - plan, a_desc, b_desc, out3 = plan_dense_projection_det_gemm( + plan, a_desc, b_desc, out3 = cake_dense_gemm_det_train_projection_plan( Q.permute(1, 2, 0), dL.permute(1, 0, 2), out, @@ -5018,7 +5018,7 @@ def test_batched_small_m_swap_and_row_rules(): # a rule fills only the knobs the caller left unset ROW_RULES.clear() ROW_RULES[ident + (M, None, N)] = {"cta_rows": 256, "group_m": 8} - ruled, *_ = plan_dense_projection_det_gemm( + ruled, *_ = cake_dense_gemm_det_train_projection_plan( Q.permute(1, 2, 0), dL.permute(1, 0, 2), out, @@ -5034,7 +5034,7 @@ def test_batched_small_m_swap_and_row_rules(): and "_m256" in _nopro(ruled.template) and re.search(r"_g\d+", _nopro(ruled.template)) is None ) - pinned, *_ = plan_dense_projection_det_gemm( + pinned, *_ = cake_dense_gemm_det_train_projection_plan( Q.permute(1, 2, 0), dL.permute(1, 0, 2), out, @@ -5054,7 +5054,7 @@ def test_batched_small_m_swap_and_row_rules(): "stages": 3, "epi": "reg", } - ruled, *_ = plan_dense_projection_det_gemm( + ruled, *_ = cake_dense_gemm_det_train_projection_plan( Q.permute(1, 2, 0), dL.permute(1, 0, 2), out, @@ -5069,7 +5069,7 @@ def test_batched_small_m_swap_and_row_rules(): (128, "dense_proj_det_gemm_nn_n192_t"), (64, "dense_proj_det_gemm_nn_n192_m64_t"), ): - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( Q.permute(1, 2, 0), dL.permute(1, 0, 2), out, @@ -5105,7 +5105,7 @@ def _ht_rows_without_half_fit(sm_count, arch): _fallback=False, ) try: - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk=True, **kw ) except (ValueError, NotImplementedError): @@ -5134,7 +5134,7 @@ def test_ht_program_never_plans_linearised_stream_k_units(arch, sm_count): ), (row, op, T, plan.template) # the same row with the default policy plans identically (auto never splits where the half items do not fit) v = _views("proj", row, op, "bf16", T) - auto, *_ = plan_dense_projection_det_gemm( + auto, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -5161,7 +5161,7 @@ def test_ht_program_never_plans_linearised_stream_k_units(arch, sm_count): def _plan_or_none(row, op, T, sm_count, arch, **knobs): v = _views("proj", row, op, "bf16", T) try: - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -5193,7 +5193,7 @@ def test_ht_program_with_forced_stream_k_matches_fp64_reference(): row, op, T, _ = found[0] v = _make_inputs("proj", row, op, "bf16", T, seed=SEED + 7) _require_program(_plan_template(v, sk=True)) - prepared = prepare_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk=True) + prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"], sk=True) assert "_ht" in prepared.plan.template and prepared.plan.sk_units == 0 out = prepared.launch() torch.cuda.synchronize() @@ -5282,7 +5282,7 @@ def test_sk_parts_rule(pairs): ("indexer_k", "f32", (128, 0, 48, 96, 127)), ): v = _views("proj", row, "wgrad", dt, T) - plan = plan_dense_projection_det_gemm( + plan = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -5302,7 +5302,7 @@ def test_sk_parts_rule(pairs): assert plan.num_cluster_tiles == want[1] + want[3] # an explicit caller split wins over the rule, and sm_100a keeps its own (ruleless) plan v = _views("proj", "indexer_hw", "wgrad", "bf16", T) - pinned = plan_dense_projection_det_gemm( + pinned = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -5321,7 +5321,7 @@ def test_sk_parts_rule(pairs): ) # (sm_100a has no sk_parts rule; its round-9 rule plans the 64-row family for this row, so the 128-row # family is forced to keep the auto two-way split of the 24 tiles under test) - other = plan_dense_projection_det_gemm( + other = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -5340,7 +5340,7 @@ def test_sk_parts_rule(pairs): 48, ) ) # sm_100a has no sk_parts rule: its 128-wide tiles keep the auto two-way split - forced = plan_dense_projection_det_gemm( + forced = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -5384,21 +5384,21 @@ def test_router_layout_classification_and_swap(): X = torch.empty(T, K) W = torch.empty(N, K) G = torch.empty(T, N) - plan, A_eff, B_eff = plan_router_fp32_det_gemm(X, W.t(), torch.empty(T, N)) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(X, W.t(), torch.empty(T, N)) assert (plan.layout, plan.template, plan.splits) == ( "kk", "router_fp32_det_gemm_kk", ROUTER_SPLITS["kk"], ) assert not plan.swapped and (plan.M, plan.N, plan.K) == (T, N, K) - plan, A_eff, B_eff = plan_router_fp32_det_gemm(G, W, torch.empty(T, K)) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(G, W, torch.empty(T, K)) assert (plan.layout, plan.template, plan.splits) == ( "kn", "router_fp32_det_gemm_kn", 1, ) assert (plan.M, plan.N, plan.K) == (T, K, N) and plan.b_mn - plan, A_eff, B_eff = plan_router_fp32_det_gemm(G.t(), X, torch.empty(N, K)) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(G.t(), X, torch.empty(N, K)) assert (plan.layout, plan.template, plan.splits) == ( "nn_t", "router_fp32_det_gemm_nn_t", @@ -5414,30 +5414,30 @@ def test_router_layout_classification_and_swap(): 1, 1, ) - plan, *_ = plan_router_fp32_det_gemm(X, W.t(), torch.empty(T, N), splits=2) + plan, *_ = cake_dense_gemm_det_train_router_plan(X, W.t(), torch.empty(T, N), splits=2) assert plan.splits == 2 and plan.k_iters_split == -(-plan.k_blocks // 2) with pytest.raises(NotImplementedError, match="MN-major A"): router_layout_class(True, False) with pytest.raises(ValueError, match="fp32"): - plan_router_fp32_det_gemm(X.to(torch.bfloat16), W.t(), torch.empty(T, N)) + cake_dense_gemm_det_train_router_plan(X.to(torch.bfloat16), W.t(), torch.empty(T, N)) with pytest.raises(ValueError, match="K must agree"): - plan_router_fp32_det_gemm(X, torch.empty(K + 32, N), torch.empty(T, N)) + cake_dense_gemm_det_train_router_plan(X, torch.empty(K + 32, N), torch.empty(T, N)) with pytest.raises(ValueError, match="out must be"): - plan_router_fp32_det_gemm(X, W.t(), torch.empty(N, T)) + cake_dense_gemm_det_train_router_plan(X, W.t(), torch.empty(N, T)) with pytest.raises(ValueError, match="multiple of 8"): - plan_router_fp32_det_gemm( + cake_dense_gemm_det_train_router_plan( X, torch.empty(K, 20).t().contiguous().t(), torch.empty(T, 20) ) with pytest.raises(ValueError, match="unit stride"): - plan_router_fp32_det_gemm( + cake_dense_gemm_det_train_router_plan( torch.empty(T, 2 * K)[:, ::2], W.t(), torch.empty(T, N) ) # a ragged batch with T = 0: M = 0 for the forward / input gradient, K = 0 for the # weight gradient -- rejected before a zero-sized grid or TMA extent is planned with pytest.raises(ValueError, match="must be >= 1"): - plan_router_fp32_det_gemm(torch.empty(0, K), W.t(), torch.empty(0, N)) + cake_dense_gemm_det_train_router_plan(torch.empty(0, K), W.t(), torch.empty(0, N)) with pytest.raises(ValueError, match="must be >= 1"): - plan_router_fp32_det_gemm( + cake_dense_gemm_det_train_router_plan( torch.empty(N, 0), torch.empty(0, K), torch.empty(N, K) ) @@ -5462,7 +5462,7 @@ def test_router_planner_refuses_a_misaligned_out_of_a_single_split_row_major_sto cases = dict(kk=(X, W.t(), T, N), kn=(G, W, T, K), nn_t=(G.t(), X, N, K)) for layout, (A, B, rows, cols) in cases.items(): for splits in sorted({1, 2, ROUTER_SPLITS[layout]}): - plan, *_ = plan_router_fp32_det_gemm( + plan, *_ = cake_dense_gemm_det_train_router_plan( A, B, _offset_out(rows, cols, 0), splits=splits ) assert (plan.layout, plan.splits) == (layout, splits) @@ -5470,16 +5470,16 @@ def test_router_planner_refuses_a_misaligned_out_of_a_single_split_row_major_sto out = _offset_out(rows, cols, offset) if splits == 1 and layout != "nn_t": with pytest.raises(ValueError, match="16-byte aligned address"): - plan_router_fp32_det_gemm(A, B, out, splits=splits) + cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) else: - plan, *_ = plan_router_fp32_det_gemm(A, B, out, splits=splits) + plan, *_ = cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) assert (plan.layout, plan.splits) == (layout, splits) # the production split counts: the input gradient (kn, one split) refuses, the forward (kk, 4) and the # weight gradient (nn_t, 6) plan with pytest.raises(ValueError, match="16-byte aligned address"): - plan_router_fp32_det_gemm(G, W, _offset_out(T, K, 1)) - assert plan_router_fp32_det_gemm(X, W.t(), _offset_out(T, N, 1))[0].splits == 4 - assert plan_router_fp32_det_gemm(G.t(), X, _offset_out(N, K, 1))[0].splits == 6 + cake_dense_gemm_det_train_router_plan(G, W, _offset_out(T, K, 1)) + assert cake_dense_gemm_det_train_router_plan(X, W.t(), _offset_out(T, N, 1))[0].splits == 4 + assert cake_dense_gemm_det_train_router_plan(G.t(), X, _offset_out(N, K, 1))[0].splits == 6 def test_split_fp32_to_bf16x3_is_exact_and_allocation_free_in_place(): @@ -5710,7 +5710,7 @@ def test_convenience_cache_put_reports_every_entry_that_leaves_the_cache(): def _cpu_router_entry(arena, op, T, splits=None): """A planned K2 convenience entry on CPU views (no module: a no-op entry callable).""" v = _router_inputs(op, T, seed=SEED + 8, device="cpu") - plan, A_eff, B_eff = plan_router_fp32_det_gemm( + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan( v["A"], v["B"], v["out"], splits=splits ) _, split_src = cake_backend._k2_views(v["A"], v["B"], plan) @@ -5818,7 +5818,7 @@ def test_convenience_arena_grows_to_the_largest_requirement_and_carves_views( assert fwd.views_version == arena.version # a K1 entry carves the buffers themselves (pointer parameters) v = _views("proj", "kv_a", "fwd", "bf16", 257) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a" ) arena.require("k1_ws", plan.ws_f32_elems) @@ -5836,7 +5836,7 @@ def test_convenience_arena_grows_to_the_largest_requirement_and_carves_views( def test_convenience_arena_is_released_with_its_last_entry(monkeypatch): """``_release_entry`` (the caches' ``on_evict``) drops an arena from the registry once its last cached - entry is evicted; ``clear_convenience_caches`` drops them all.""" + entry is evicted; ``cake_dense_gemm_det_train_convenience_cache_clear`` drops them all.""" monkeypatch.setattr(cake_backend, "_K1_CACHE", OrderedDict()) monkeypatch.setattr(cake_backend, "_K2_CACHE", OrderedDict()) monkeypatch.setattr(cake_backend, "_ARENAS", {}) @@ -5867,7 +5867,7 @@ def test_convenience_arena_is_released_with_its_last_entry(monkeypatch): a7b = cake_backend._arena_for(0, 7, device) cake_backend._release_entry(e1) assert cake_backend._ARENAS[(0, 7)] is a7b - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() assert cake_backend._ARENAS == {} and cache == OrderedDict() @@ -6129,7 +6129,7 @@ def chain_split(src, parts, resid): n = v["W"].numel() split_args = (n, -(-n // cake_backend.SPLIT_BLOCK_ELEMS), 1, 1) # the prepared path - prepared = prepare_router_fp32_det_gemm(v["A"], v["B"], v["out"], splits=4) + prepared = cake_dense_gemm_det_train_router_prepare(v["A"], v["B"], v["out"], splits=4) assert prepared.fused_split is fused and ("resid" in prepared.tensors) is not fused prepared.launch() if fused: @@ -6144,7 +6144,7 @@ def chain_split(src, parts, resid): # the convenience path calls.clear() chain.clear() - out = cake_backend.router_fp32_det_gemm(v["A"], v["B"], v["out"], splits=4) + out = cake_backend.cake_dense_gemm_det_train_router(v["A"], v["B"], v["out"], splits=4) assert out is v["out"] cached = next(iter(cake_backend._K2_CACHE.values())) arena = cached.arena @@ -6189,7 +6189,7 @@ def test_convenience_views_match_the_planner_views(sm_count): for family, row, op, dt, T in cases: v = _views(family, row, op, dt, T) transposed = v.get("transposed", False) - plan, a_desc, b_desc, out3 = plan_dense_projection_det_gemm( + plan, a_desc, b_desc, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -6386,7 +6386,7 @@ def test_registry_fallback_serves_every_registered_class(arch): for K in _CLASS_SWEEP_K: v = _class_views(layout, out_f32, out_t, M, N, K) try: - mirror, *_ = plan_dense_projection_det_gemm( + mirror, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -6402,7 +6402,7 @@ def test_registry_fallback_serves_every_registered_class(arch): arch, mirror.a_mn, mirror.b_mn, out_f32, out_t, False, N, K, M ): continue - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=out_t, **kw ) where = ( @@ -6494,7 +6494,7 @@ def test_registry_fallback_reaches_the_shapes_tranche_programs(monkeypatch): neighbour = next((c for c in candidates if c in registered), None) assert neighbour is not None, (M, N, K, launcher) v = _class_views(layout, out_f32, out_t, M, N, K) - mirror, *_ = plan_dense_projection_det_gemm( + mirror, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=out_t, _fallback=False, **kw ) assert mirror.template == launcher and launcher not in registered, ( @@ -6503,7 +6503,7 @@ def test_registry_fallback_reaches_the_shapes_tranche_programs(monkeypatch): K, mirror.template, ) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=out_t, **kw ) assert plan.template == neighbour, (M, N, K, plan.template, neighbour) @@ -6516,7 +6516,7 @@ def test_registry_fallback_reaches_the_shapes_tranche_programs(monkeypatch): # a caller-forced tail policy keeps the row on a registered program: a pair form admits it, the single-CTA # form (``cta1`` clears the tail knobs) stays reachable when no pair form is registered v = _class_views("kk", False, False, 16384, 128, 4096) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk=True, **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, **kw) assert plan.template in registered and plan.knob_fallback, plan.template if plan.template.endswith("_c1"): assert plan.cta1 and plan.sk_units == 0 and not plan.sk_sync @@ -6525,10 +6525,10 @@ def test_registry_fallback_reaches_the_shapes_tranche_programs(monkeypatch): registered107 = _registered("sm_107a") neighbour107 = next(c for c in nn_n128_t if c in registered107) v = _class_views("nn", False, True, 4096, 128, 16384) - mirror, *_ = plan_dense_projection_det_gemm( + mirror, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=True, _fallback=False, **kw107 ) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], transposed_out=True, **kw107 ) assert ( @@ -6556,11 +6556,11 @@ def test_registry_fallback_decodes_the_nearest_generated_program(monkeypatch): monkeypatch.setattr(cake_backend, "ROW_RULES", {}) kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a") v = _class_views("nn", False, False, 4096, 4096, 16384) - mirror, *_ = plan_dense_projection_det_gemm( + mirror, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], _fallback=False, **kw ) assert mirror.template not in fake["sm_100a"] - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert plan.template in ( "dense_proj_det_gemm_nn_n256_sks_sb2_pro7", "dense_proj_det_gemm_nn_n160_m256_bz64_pro15", @@ -6576,12 +6576,12 @@ def test_registry_fallback_decodes_the_nearest_generated_program(monkeypatch): ) # a class without a program stays unregistered (the launch binding raises, not the planner) and unflagged v = _class_views("kn", False, False, 4096, 4096, 16384) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert plan.template not in fake["sm_100a"] assert not (plan.knob_fallback or plan.rule_fallback or plan.swap_fallback) # a caller-forced knob wins over a decoded one: the forced 256-row tile never reaches the 128-row sync program v = _class_views("nn", False, False, 4096, 4096, 16384) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta_rows=256, **kw ) assert ( @@ -6995,7 +6995,7 @@ def test_measured_shape_rules_plan_their_programs(arch): kw = dict( sm_count=sm_count, l2_bytes=L2_BYTES, transposed_out=key[4], arch=arch ) - mirror, *_ = plan_dense_projection_det_gemm( + mirror, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], _fallback=False, **kw ) assert mirror.template == template, (key, T, mirror.template) @@ -7006,7 +7006,7 @@ def test_measured_shape_rules_plan_their_programs(arch): T, template, ) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) if template in registered: assert plan.template == template, (key, T, plan.template) assert not ( @@ -7108,17 +7108,17 @@ def test_sm103a_inherits_the_sm100a_rules_where_it_has_none(): direct = fallen = own_rows = 0 for label, v, extra in _owner_rows(): kw = dict(sm_count=sm, l2_bytes=L2_BYTES, **extra) - mirror, *_ = plan_dense_projection_det_gemm( + mirror, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], arch="sm_103a", _fallback=False, **kw ) if _own_measured_rule("sm_103a", mirror) is not None: own_rows += 1 # a GB300-measured row: its own rule, checked by the measured-rows test continue - parent, *_ = plan_dense_projection_det_gemm( + parent, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], arch="sm_100a", _fallback=False, **kw ) assert _plan_signature(mirror) == _plan_signature(parent), label - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], arch="sm_103a", **kw ) if parent.template in registered: @@ -7142,7 +7142,7 @@ def plans(arch): out = {} for label, v, extra in _owner_rows(): try: - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -7216,7 +7216,7 @@ def test_wgrad_swap_follows_the_measured_rule(): assert not transposed and A.shape == (512, T) and B is X A, B, transposed = wgrad_views(G, X, arch="sm_100a") assert transposed and A.shape == (4096, T) and B is G - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( A, B, torch.empty(512, 4096, dtype=torch.bfloat16), @@ -7347,7 +7347,7 @@ def test_round2_rules_plan_their_knobs(): for M, N in ((12288, 4096), (4096, 12288)): for T, k_blocks in ((16384, 256), (16231, 254)): v = _class_views("nn", False, False, M, N, T) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert ( plan.template == "dense_proj_det_gemm_nn_n256_m256_skx_skf_pro7" ), ( @@ -7375,7 +7375,7 @@ def test_round2_rules_plan_their_knobs(): and plan.ws_f32_elems > 0 ) # a caller-forced synchronised plan or single-CTA form drops the rule's order with its split - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk_sync=True, **a ) assert ( @@ -7383,7 +7383,7 @@ def test_round2_rules_plan_their_knobs(): and not forced.sk_first and "_skf" not in forced.template ) - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk=False, **a ) assert ( @@ -7392,7 +7392,7 @@ def test_round2_rules_plan_their_knobs(): and forced.template == "dense_proj_det_gemm_nn_n256_m256_pro7" ) with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk=False, sk_first=True, **a ) # LGF-R2: the 4096 squares take the half-tail column-grouped programs (group 16), the N 5120 rows group 12 @@ -7401,7 +7401,7 @@ def test_round2_rules_plan_their_knobs(): ("kn", "dense_proj_det_gemm_kn_n256_ht_cg_pro15"), ): v = _class_views(layout, False, False, 16384, 4096, 4096) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert ( plan.template == template and plan.htail @@ -7420,7 +7420,7 @@ def test_round2_rules_plan_their_knobs(): ("kk", 17408, 5120, "dense_proj_det_gemm_kk_n256_cg_pro15"), ): v = _class_views(layout, False, False, 16384, N, K) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert plan.template == template and plan.cgrp and plan.group_m == 12, ( arch, layout, @@ -7446,7 +7446,7 @@ def test_round2_rules_plan_their_knobs(): if (K_layer, N_layer) == (4096, 48) else {} ) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **forced_h, **a ) assert plan.template == template and plan.lw and plan.pd == 0, ( @@ -7463,20 +7463,20 @@ def test_round2_rules_plan_their_knobs(): plan.iters_per_unit, ) == plan_fields # the load-warp form yields to a caller-forced single-CTA form or a pro without bit 4 - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=True, **a ) assert not forced.lw and forced.template.endswith("_c1") - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], pro=3, **a ) assert not forced.lw and forced.template.endswith("_pro3") with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], pro=3, lw=True, **a ) v = _class_views("kn", False, False, 16384, 5120, 128) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma2_pd2_lw_pro15" and plan.pd == 2 @@ -7486,7 +7486,7 @@ def test_round2_rules_plan_their_knobs(): # the GB300-only round-2 picks: the 4096 -> 64 load-warp program, the l2_256b promotion of the N 512 input gradients a = dict(kw, arch="sm_103a") v = _class_views("kn", False, False, 16384, 4096, 64) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma2_pd2_lw_pro15" and plan.group_m == 32 @@ -7494,13 +7494,13 @@ def test_round2_rules_plan_their_knobs(): # (the two rows' eight-warp CTA form is forced off here - its own test is # test_eight_warp_form_rules_plan_their_programs; the ten-warp program below is the row's program where it yields) v = _class_views("kn", False, False, 16384, 4096, 512) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], form="", **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **a) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_lw_mg_pro15" and plan.promo == "l2_256b" ) # + the whole-tile tile resolve (test_mg_rules_plan_the_whole_tile_programs) v = _class_views("kn", False, False, 16384, 5120, 512) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], form="", **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **a) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_mg_sw25_pro15" and plan.promo == "l2_256b" @@ -7511,14 +7511,14 @@ def test_round2_rules_plan_their_knobs(): # single-CTA program a = dict(kw, arch="sm_100a") v = _class_views("kn", False, False, 16384, 4096, 64) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_w8r" and plan.form == "w8r" and plan.group_m == 8 and not plan.lw ) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], epi="tma", slots=2, cta1=True, pd=3, **a ) assert ( @@ -7529,7 +7529,7 @@ def test_round2_rules_plan_their_knobs(): # (the row's eight-warp CTA form is forced off here - its own test is # test_eight_warp_form_rules_plan_their_programs; the ten-warp program below is the row's program where it yields) v = _class_views("kn", False, False, 16384, 5120, 512) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], form="", **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **a) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_pro15" and plan.promo == "none" @@ -7537,15 +7537,15 @@ def test_round2_rules_plan_their_knobs(): ) # the column-grouped raster takes any group count; the pair raster still needs an even one v = _class_views("kk", False, False, 16384, 5120, 5120) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], group_m=3, **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], group_m=3, **a) assert plan.group_m == 3 and plan.cgrp with pytest.raises(ValueError): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], group_m=3, cgrp=False, **a ) # WGR2 E1: the 4096 -> 1024 bf16 weight gradient runs direct on sm_100a now (the sm_103a program) v = _class_views("nn", False, False, 1024, 4096, 16384) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) assert ( plan.template == "dense_proj_det_gemm_nn_n224_hne_cg_pro15" and plan.block_n == 224 @@ -7606,7 +7606,7 @@ def test_mg_rules_plan_the_whole_tile_programs(T): # (these three rows carry the eight-warp CTA form too: it is forced off here, the whole-tile tile resolve is # the row's program where the form yields; test_eight_warp_form_rules_plan_their_programs tests the form) v = _class_views("kn", False, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], form="", **a ) assert plan.template == template, (key, T, plan.template) @@ -7622,7 +7622,7 @@ def test_mg_rules_plan_the_whole_tile_programs(T): assert plan.num_full == plan.pair_tiles and plan.ws_f32_elems == 0 assert plan.grid == (2 * plan.pair_tiles, 1, 1) and not plan.cta1 # the twin without the resolve: the same plan, only the program differs - plain, *_ = plan_dense_projection_det_gemm( + plain, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], mg=False, **a ) assert not plain.mg and plain.template == template.replace("_mg", "") @@ -7643,17 +7643,17 @@ def test_mg_rules_plan_the_whole_tile_programs(T): == template ) # a caller forcing a form the program excludes: the rule's mg yields and the forced form plans - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=True, **a ) assert not forced.mg and forced.template.endswith("_c1"), forced.template - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=True, **a ) assert not forced.mg and forced.cgrp and "_mg" not in forced.template assert "_cg" in forced.template with pytest.raises(ValueError, match=r"mg \(whole-tile"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=True, mg=True, **a ) # the K 512 rows: ``sk=True`` admits no split (8 K steps = SK_MIN_ITERS: 36 units over the 36 tail tiles of the @@ -7661,17 +7661,17 @@ def test_mg_rules_plan_the_whole_tile_programs(T): # 18 units over the 16 tail tiles, so the rule's mg yields to the plain program, and forcing mg on that plan is # refused v = _class_views("kn", False, False, T, 4096, 512) - sk, *_ = plan_dense_projection_det_gemm( + sk, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk=True, form="", **a ) assert sk.sk_units == 0 and sk.mg assert sk.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_lw_mg_pro15" v = _class_views("kn", False, False, T, 6144, 576) - sk, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk=True, **a) + sk, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, **a) assert (sk.sk_units, sk.tail_tiles, sk.iters_per_unit) == (18, 16, 8) assert not sk.mg and sk.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_pro15" with pytest.raises(ValueError, match="whole-tile plan"): - plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk=True, mg=True, **a) + cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, mg=True, **a) # the rows without the resolve keep their program (5120 -> 1024 paces its stores), the fp32 twins and sm_100a # never carry the resolve (both rows carry the eight-warp CTA form: forced off here) for K_layer, N_layer, template in ( @@ -7679,7 +7679,7 @@ def test_mg_rules_plan_the_whole_tile_programs(T): (5120, 1024, "dense_proj_det_gemm_kn_n256_tma1_pd2_sw25_pro15"), ): v = _class_views("kn", False, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], form="", **a ) assert not plan.mg, (K_layer, N_layer) @@ -7692,10 +7692,10 @@ def test_mg_rules_plan_the_whole_tile_programs(T): ) for K_layer, N_layer, _t, _g, _p, _lw in rows: f = _class_views("kn", True, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm(f["A"], f["B"], f["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **a) assert not plan.mg and "_mg" not in plan.template, (K_layer, N_layer, "f32") v = _class_views("kn", False, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **dict(kw, arch="sm_100a") ) assert not plan.mg and "_mg" not in plan.template, (K_layer, N_layer, "sm_100a") @@ -7758,7 +7758,7 @@ def test_sw_rules_plan_the_paced_store_programs(T): # (both rows carry the eight-warp CTA form too: forced off here, the paced program is the row's program # where the form yields; test_eight_warp_form_rules_plan_their_programs tests the form) v = _class_views("kn", False, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], form="", **a ) assert plan.template == template, (key, T, plan.template) @@ -7766,7 +7766,7 @@ def test_sw_rules_plan_the_paced_store_programs(T): assert (plan.epi, plan.slots, plan.pd, plan.pro) == ("tma", 1, 2, 15), key assert plan.sk_units == 0 and not plan.cta1 # the twin without the pacing: the same plan, only the program differs - plain, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sw=0, **a) + plain, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sw=0, **a) assert plain.sw == 0 and plain.template == template.replace("_sw25", "") assert {**_plan_signature(plan), "template": None, "sw": None} == { **_plan_signature(plain), @@ -7787,24 +7787,24 @@ def test_sw_rules_plan_the_paced_store_programs(T): # a caller forcing a program outside the drain or the single-CTA form: the rule's sw yields and the forced # form plans; forcing sw with it is refused for forced_knobs in (dict(pd=0), dict(pd=3)): - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **forced_knobs, **a ) assert forced.sw == 0 and forced.pd == forced_knobs["pd"], forced.template assert "_sw" not in forced.template, forced.template with pytest.raises(ValueError, match=r"sw \(paced output stores\)"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sw=25, **forced_knobs, **a ) - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=True, **a ) assert forced.sw == 0 and forced.template.endswith("_c1"), forced.template # the fp32 twin and sm_100a never carry the pacing f = _class_views("kn", True, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm(f["A"], f["B"], f["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **a) assert plan.sw == 0 and "_sw" not in plan.template, (K_layer, N_layer, "f32") - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **dict(kw, arch="sm_100a") ) assert plan.sw == 0 and "_sw" not in plan.template, (K_layer, N_layer) @@ -7815,7 +7815,7 @@ def test_sw_rules_plan_the_paced_store_programs(T): (6144, 576, "dense_proj_det_gemm_kn_n256_tma1_pd2_mg_pro15"), ): v = _class_views("kn", False, False, T, K_layer, N_layer) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], form="", **a ) assert plan.sw == 0 and plan.template == template, (K_layer, N_layer) @@ -7916,7 +7916,7 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): rule = ROW_RULES[key] assert rule["form"] == form and rule["sk"] is False, key v = _class_views("kn", False, False, T, N_layer, K_layer) - plan, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **mirror) + plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **mirror) assert plan.template == template, (key, T, plan.template) assert plan.form == form and plan.tma_out and plan.epi == "tma" assert (plan.slots, plan.stages, plan.block_n, plan.cta_rows) == ( @@ -7933,7 +7933,7 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): assert plan.pd == 0 and plan.pro == 0 and plan.pf == 0 and plan.sh == 0, key assert plan.hints == ("none", "none") and plan.store_hint == "none", key # the registered program plans directly (no fallback) - served, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + served, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert served.template == template and template in registered, (key, T) assert not ( served.swap_fallback @@ -7942,7 +7942,7 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): or served.wave_fallback ), key # the ten-warp twin: the same plan arithmetic, the row's program with its pipeline knobs - plain, *_ = plan_dense_projection_det_gemm( + plain, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], form="", **mirror ) assert plain.form == "" and "_w8" not in plain.template, plain.template @@ -7991,37 +7991,37 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): # register epilogue or a caller-forced pipeline knob plan the row's ten-warp program; a caller-forced form is # refused on every such plan v = _class_views("kn", False, False, T, 6144, 576) - sk, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], sk=True, **mirror) + sk, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, **mirror) assert (sk.sk_units, sk.tail_tiles) == (18, 16) and sk.form == "" assert sk.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_pro15" with pytest.raises(ValueError, match="whole-tile"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sk=True, form="w8rs", **mirror ) for N_layer, K_layer in rows: v = _class_views("kn", False, False, T, N_layer, K_layer) - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cta1=True, **mirror ) assert forced.form == "" and forced.template.endswith("_c1"), forced.template - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], epi="reg", **mirror ) assert forced.form == "" and forced.epi == "reg", forced.template assert "_w8" not in forced.template, forced.template - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], pro=7, **mirror ) assert forced.form == "" and forced.pro == 7, forced.template assert forced.template.endswith("_pro7"), forced.template for bad in (dict(epi="reg"), dict(cta1=True), dict(pro=7), dict(pd=2)): with pytest.raises(ValueError, match="form"): - plan_dense_projection_det_gemm( + cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], form="w8r", **bad, **mirror ) # the fp32 twin carries no form on either architecture (the fp32 rows are other rules) f = _class_views("kn", True, False, T, N_layer, K_layer) - plan, *_ = plan_dense_projection_det_gemm(f["A"], f["B"], f["out"], **mirror) + plan, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **mirror) assert plan.form == "" and "_w8" not in plan.template, (N_layer, K_layer) # the other architectures' plans of the row: the five sm_100a rows shared with sm_103a carry the sm_103a form # (``w8rs`` / ``w8rs`` / ``w8r`` / ``w8r`` / ``w8rs``) and the sm_100a 4096 -> 64 and 5120 -> 64 rows none there @@ -8032,7 +8032,7 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): for other_arch, other_rows in EIGHT_WARP_ROW_PLANS.items(): if other_arch == arch: continue - other, *_ = plan_dense_projection_det_gemm( + other, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **dict(mirror, arch=other_arch) ) expected_other = other_rows.get((N_layer, K_layer)) @@ -8053,7 +8053,7 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): 0, other_group_m, ), (other_arch, N_layer, K_layer) - other, *_ = plan_dense_projection_det_gemm( + other, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **dict(mirror, arch="sm_107a") ) assert other.form == "" and "_w8" not in other.template, (N_layer, K_layer) @@ -8239,7 +8239,7 @@ def test_gb_t_exact_rules_plan_their_programs(arch, monkeypatch): plans = {} for t in (T - 1, T, T + 1): v = _class_views("kn" if b_mn else "kk", False, False, t, N, K) - plans[t], *_ = plan_dense_projection_det_gemm( + plans[t], *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **kw ) plan = plans[T] @@ -8252,12 +8252,12 @@ def test_gb_t_exact_rules_plan_their_programs(arch, monkeypatch): for k in _gb_t_exact_keys_of(key): m.delitem(ROW_RULES, k) v = _class_views("kn" if b_mn else "kk", False, False, T, N, K) - ragged, *_ = plan_dense_projection_det_gemm( + ragged, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **kw ) for t in (T - 1, T + 1): v = _class_views("kn" if b_mn else "kk", False, False, t, N, K) - without, *_ = plan_dense_projection_det_gemm( + without, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **kw ) assert _plan_signature(plans[t]) == _plan_signature(without), ( @@ -8317,8 +8317,8 @@ def test_gb_t_exact_rows_match_fp64_and_their_plan_without_the_key(key, monkeypa template = GB_T_EXACT_ROWS[(arch, b_mn, N, K)][1] _require_program(template) out = torch.full((T, N), float("nan"), dtype=torch.bfloat16, device=device) - cake_backend.clear_convenience_caches() - prepared = prepare_dense_projection_det_gemm(A, B, out) + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() + prepared = cake_dense_gemm_det_train_projection_prepare(A, B, out) assert prepared.template == template, prepared.template assert _whole_k(prepared.plan) prepared.launch() @@ -8328,8 +8328,8 @@ def test_gb_t_exact_rows_match_fp64_and_their_plan_without_the_key(key, monkeypa with monkeypatch.context() as m: for k in _gb_t_exact_keys_of(key): m.delitem(ROW_RULES, k) - cake_backend.clear_convenience_caches() - plain = prepare_dense_projection_det_gemm(A, B, plain_out) + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() + plain = cake_dense_gemm_det_train_projection_prepare(A, B, plain_out) assert plain.template != template and _whole_k(plain.plan), plain.template plain.launch() torch.cuda.synchronize() @@ -8365,7 +8365,7 @@ def test_half_tail_fit_gates_the_plan_not_the_program(monkeypatch): (5120, (1280, 0, 0, 1), False), ): v = _class_views("kn", False, False, 16384, K_layer, 32) - ruled, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + ruled, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert ruled.template == "dense_proj_det_gemm_kn_n256_ht_tma2_lw_pro15" assert ruled.htail and ruled.htail_plan == half assert ( @@ -8377,11 +8377,11 @@ def test_half_tail_fit_gates_the_plan_not_the_program(monkeypatch): assert ruled.sk_iters == (72 if half else 0) and ruled.ws_f32_elems == 0 with monkeypatch.context() as m: m.setattr(cake_backend, "ROW_RULES", {}) - plain, *_ = plan_dense_projection_det_gemm(v["A"], v["B"], v["out"], **kw) + plain, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) assert ( not plain.htail and not plain.htail_plan and "_ht" not in plain.template ) - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **knobs, **kw ) assert forced.template == ruled.template @@ -8394,7 +8394,7 @@ def test_half_tail_fit_gates_the_plan_not_the_program(monkeypatch): ) == fields # the fit follows the device: on 84 pairs the K 5120 tail of 20 tiles fits as 40 half items v = _class_views("kn", False, False, 16384, 5120, 32) - wide, *_ = plan_dense_projection_det_gemm( + wide, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **dict(kw, sm_count=168) ) assert wide.htail and wide.htail_plan and "_ht_" in wide.template @@ -8425,7 +8425,7 @@ def test_exact_wave_gate_follows_the_device_pair_count(): base = dict( l2_bytes=L2_BYTES, transposed_out=key[4], arch=arch, _fallback=False ) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=152, **base ) assert ( @@ -8435,7 +8435,7 @@ def test_exact_wave_gate_follows_the_device_pair_count(): tiles_256 = plan.L * (plan.m_tiles // 2) * -(-plan.N // 256) for sm_count in (148, 160): pairs = sm_count // 2 - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=sm_count, **base ) expect = not exact_waves(tiles_224, pairs) and -(-tiles_224 // pairs) > -( @@ -8447,13 +8447,13 @@ def test_exact_wave_gate_follows_the_device_pair_count(): key, T, ) # the same waves on 80 pairs: the rule's tile stays - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=148, **base ) if plan.wave_fallback: fired.add(key) assert plan.block_n == 256 and "_n256_" in plan.template, (key, T) - forced, *_ = plan_dense_projection_det_gemm( + forced, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=148, block_n=256, **base ) assert dataclasses.replace(plan, wave_fallback=False) == forced, (key, T) @@ -8465,7 +8465,7 @@ def test_exact_wave_gate_follows_the_device_pair_count(): exact_waves(tiles_224, 74) or -(-tiles_224 // 74) == -(-tiles_256 // 74) ), (key, T) - kept, *_ = plan_dense_projection_det_gemm( + kept, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=148, block_n=224, **base ) assert kept.block_n == 224 and not kept.wave_fallback, (key, T) @@ -8475,7 +8475,7 @@ def test_exact_wave_gate_follows_the_device_pair_count(): key for key, _ in rows if (key[6], key[8]) in ((4096, 4096), (4096, 1024)) }, sorted(fired, key=repr) for key in fired: - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( *(_key_views(key, 16384)[k] for k in ("A", "B", "out")), sm_count=148, l2_bytes=L2_BYTES, @@ -8504,7 +8504,7 @@ def test_batched_rows_satisfy_the_kernel_output_layout_predicate(): for op in ("fwd", "dgrad", "wgrad"): for T in (1001, 16231): v = _views("mla", row, op, "bf16", T) - plan, _, _, out3 = plan_dense_projection_det_gemm( + plan, _, _, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -8531,7 +8531,7 @@ def test_batched_rows_satisfy_the_kernel_output_layout_predicate(): assert checked == 12 # the interleaved store of the qabs input gradient (D = 192 of a 256-wide slot): the second form, with equality v = _views("mla", "qabs", "dgrad", "bf16", 33) - plan, _, _, out3 = plan_dense_projection_det_gemm( + plan, _, _, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=152, l2_bytes=L2_BYTES ) ldo, out_l = cake_backend._out_strides(out3) @@ -8542,7 +8542,7 @@ def test_batched_rows_satisfy_the_kernel_output_layout_predicate(): ) # a batch of one: the matrix span satisfies the batch-major form v = _views("proj", "q_a", "fwd", "bf16", 1001) - plan, _, _, out3 = plan_dense_projection_det_gemm( + plan, _, _, out3 = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], sm_count=152, l2_bytes=L2_BYTES ) ldo, out_l = cake_backend._out_strides(out3) @@ -8552,7 +8552,7 @@ def test_batched_rows_satisfy_the_kernel_output_layout_predicate(): def test_router_convenience_views_match_the_planner(): for op, T in (("fwd", 257), ("dgrad", 257), ("wgrad", 1001)): v = _router_inputs(op, T, seed=SEED + 5, device="cpu") - plan, A_eff, B_eff = plan_router_fp32_det_gemm(v["A"], v["B"], v["out"]) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(v["A"], v["B"], v["out"]) a_view, split_src = cake_backend._k2_views(v["A"], v["B"], plan) want_split = B_eff if plan.b_mn else B_eff.t() want_a = A_eff.t().unsqueeze(0) if plan.a_mn else A_eff.unsqueeze(0) @@ -8650,7 +8650,7 @@ def _check(actual, expected_f64, *, untouched=None): def _plan_template(v, **knobs): - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], @@ -8672,10 +8672,10 @@ def test_projection_rows_match_fp64_reference(row, op, out_dtype, T): v = _make_inputs("proj", row, op, out_dtype, T, seed=SEED) _require_program(_plan_template(v)) if op == "wgrad": - prepared = prepare_projection_det_wgrad(v["G"], v["X"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(v["G"], v["X"], v["out"]) assert prepared.plan.transposed_out == v["transposed"] else: - prepared = prepare_dense_projection_det_gemm(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"]) out = prepared.launch() torch.cuda.synchronize() assert out is v["out"] @@ -8714,13 +8714,13 @@ def test_eight_warp_form_rows_match_fp64_reference(key, T): A = torch.randn(T, K, device=device, generator=g).to(torch.bfloat16) B = (torch.randn(K, N, device=device, generator=g) * K**-0.5).to(torch.bfloat16) out = torch.empty(T, N, dtype=torch.bfloat16, device=device) - plan, *_ = plan_dense_projection_det_gemm( + plan, *_ = cake_dense_gemm_det_train_projection_plan( A, B, out, sm_count=_sm_count(), l2_bytes=_l2_bytes(), arch=arch_for(device) ) if not plan.form: pytest.skip(f"the eight-warp form is not a rule of {arch_for(device)}") _require_program(plan.template) - prepared = prepare_dense_projection_det_gemm(A, B, out) + prepared = cake_dense_gemm_det_train_projection_prepare(A, B, out) assert prepared.template == plan.template and "_w8" in prepared.template assert prepared.plan.form == ROW_RULES[key]["form"] assert prepared.plan.sk_units == 0 and prepared.plan.m_tiles % 2 == 0 @@ -8730,7 +8730,7 @@ def test_eight_warp_form_rows_match_fp64_reference(key, T): expected = torch.matmul(A.double(), B.double()) _check(out, expected) plain_out = torch.empty_like(out) - plain = prepare_dense_projection_det_gemm(A, B, plain_out, form="") + plain = cake_dense_gemm_det_train_projection_prepare(A, B, plain_out, form="") assert plain.plan.form == "" and "_w8" not in plain.template, plain.template plain.launch() torch.cuda.synchronize() @@ -8756,7 +8756,7 @@ def test_mla_batched_rows_match_fp64_reference(row, op): ) v = _make_inputs("mla", row, op, "bf16", 257, seed=SEED + 1) _require_program(_plan_template(v)) - prepared = prepare_dense_projection_det_gemm(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"]) out = prepared.launch() torch.cuda.synchronize() assert out is v["out"] and prepared.plan.L == 64 @@ -8786,7 +8786,7 @@ def test_strided_views_and_storage_offsets_match_fp64_reference(): ) out = out_buf[3:, 8 : 8 + N] _require_program(_plan_template(dict(A=A, B=W.t(), out=out))) - prepared = prepare_dense_projection_det_gemm(A, W.t(), out) + prepared = cake_dense_gemm_det_train_projection_prepare(A, W.t(), out) prepared.launch() torch.cuda.synchronize() _check(out, torch.matmul(A.double(), W.double().t())) @@ -8806,7 +8806,7 @@ def test_strided_views_and_storage_offsets_match_fp64_reference(): # small-N weight gradients to the 64-row family, so only that family's instance is exported there) v_t = dict(A=A.t(), B=Gs, out=out_t, transposed=True) knobs_t = _registered_knobs(v_t, {}, dict(cta_rows=64)) - prepared = prepare_projection_det_wgrad(Gs, A, out_t, **knobs_t) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(Gs, A, out_t, **knobs_t) assert prepared.plan.transposed_out and prepared.template == _plan_template( v_t, **knobs_t ) @@ -8843,7 +8843,7 @@ def test_prepared_gemm_launches_without_allocation_and_is_deterministic(): out = torch.empty(N, K, device=device, dtype=torch.bfloat16) v = dict(A=X.t(), B=G, out=out, transposed=True) knobs = _registered_knobs(v, dict(cta_rows=64), dict(cta_rows=128)) - prepared = prepare_projection_det_wgrad(G, X, out, **knobs) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(G, X, out, **knobs) pair_tiles = prepared.plan.M // ( 2 * prepared.plan.cta_rows ) # 24 (64-row instance) or 12 (128-row instance) @@ -8887,8 +8887,8 @@ def test_convenience_path_reuses_its_planned_launch_without_allocation(): facts = cake_backend.device_facts(device) assert facts is cake_backend.device_facts(torch.device("cuda", facts.index)) assert (facts.sm_count, facts.l2_bytes) == (_sm_count(), _l2_bytes()) - cake_backend.clear_convenience_caches() - first = cake_backend.projection_det_forward(X, W, out).clone() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() + first = cake_backend.cake_dense_gemm_det_train_projection_forward(X, W, out).clone() torch.cuda.synchronize() _check(out, torch.matmul(X.double(), W.double().t())) assert len(cake_backend._K1_CACHE) == 1 @@ -8896,17 +8896,17 @@ def test_convenience_path_reuses_its_planned_launch_without_allocation(): X2 = torch.randn(T, K, device=device, generator=g).to(torch.bfloat16) out2 = torch.empty_like(out) before = torch.cuda.memory_stats() - cake_backend.projection_det_forward(X2, W, out2) + cake_backend.cake_dense_gemm_det_train_projection_forward(X2, W, out2) torch.cuda.synchronize() after = torch.cuda.memory_stats() assert after["allocation.all.allocated"] == before["allocation.all.allocated"] assert len(cake_backend._K1_CACHE) == 1 _check(out2, torch.matmul(X2.double(), W.double().t())) - assert torch.equal(cake_backend.projection_det_forward(X, W, out), first) + assert torch.equal(cake_backend.cake_dense_gemm_det_train_projection_forward(X, W, out), first) assert len(cake_backend._ARENAS) == 1 # another stream never shares an entry (nor its workspace arena) with torch.cuda.stream(torch.cuda.Stream(device=device)): - cake_backend.projection_det_forward(X, W, out2) + cake_backend.cake_dense_gemm_det_train_projection_forward(X, W, out2) torch.cuda.synchronize() assert len(cake_backend._K1_CACHE) == 2 and torch.equal(out2, first) first_entry, second_entry = cake_backend._K1_CACHE.values() @@ -8917,7 +8917,7 @@ def test_convenience_path_reuses_its_planned_launch_without_allocation(): cake_backend._ARENAS ) # the prepared path of the same binding gives the same bits - prepared = prepare_dense_projection_det_gemm(X, W.t(), out2) + prepared = cake_dense_gemm_det_train_projection_prepare(X, W.t(), out2) assert ( prepared.template == next(iter(cake_backend._K1_CACHE.values())).plan.template ) @@ -8949,8 +8949,8 @@ def test_convenience_path_replays_a_stream_k_row_bit_exact(): transposed=True, ) knobs = _registered_knobs(v, dict(cta_rows=64), dict(cta_rows=128)) - cake_backend.clear_convenience_caches() - out1 = cake_backend.projection_det_wgrad(G, X, **knobs) # ``out`` allocated here + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() + out1 = cake_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, **knobs) # ``out`` allocated here torch.cuda.synchronize() _check(out1.t(), torch.matmul(X.t().double(), G.double())) assert len(cake_backend._K1_CACHE) == 1 @@ -8960,7 +8960,7 @@ def test_convenience_path_replays_a_stream_k_row_bit_exact(): assert entry.plan.transposed_out and not entry.swapped out2 = torch.empty_like(out1) before = torch.cuda.memory_stats() - cake_backend.projection_det_wgrad(G, X, out2, **knobs) + cake_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, out2, **knobs) torch.cuda.synchronize() after = torch.cuda.memory_stats() assert after["allocation.all.allocated"] == before["allocation.all.allocated"] @@ -9003,16 +9003,16 @@ def test_convenience_path_memory_is_bounded_across_token_counts(): W = torch.randn(256, 6144, device=device, generator=g) def held_after(T_values): - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() torch.cuda.synchronize() torch.cuda.empty_cache() base = torch.cuda.memory_allocated() for T in T_values: X = torch.randn(T, 6144, device=device, generator=g) G = torch.randn(T, 256, device=device, generator=g) - cake_backend.router_det_forward(X, W) - cake_backend.router_det_dgrad(G, W) - cake_backend.router_det_wgrad(G, X) + cake_backend.cake_dense_gemm_det_train_router_forward(X, W) + cake_backend.cake_dense_gemm_det_train_router_dgrad(G, W) + cake_backend.cake_dense_gemm_det_train_router_wgrad(G, X) del X, G torch.cuda.synchronize() torch.cuda.empty_cache() @@ -9033,21 +9033,21 @@ def held_after(T_values): X = torch.randn(T, 6144, device=device, generator=g) G = torch.randn(T, 256, device=device, generator=g) outs = [torch.empty(T, 256, device=device) for _ in range(3)] - cake_backend.router_det_forward(X, W, outs[0]) + cake_backend.cake_dense_gemm_det_train_router_forward(X, W, outs[0]) torch.cuda.synchronize() before = torch.cuda.memory_stats() - cake_backend.router_det_forward(X, W, outs[1]) + cake_backend.cake_dense_gemm_det_train_router_forward(X, W, outs[1]) torch.cuda.synchronize() after = torch.cuda.memory_stats() assert after["allocation.all.allocated"] == before["allocation.all.allocated"] - prepared = prepare_router_fp32_det_gemm(X, W.t(), outs[2], splits=4) + prepared = cake_dense_gemm_det_train_router_prepare(X, W.t(), outs[2], splits=4) prepared.launch() torch.cuda.synchronize() assert torch.equal(outs[0], outs[1]) and torch.equal(outs[0], outs[2]) - dw = cake_backend.router_det_wgrad(G, X) - prepared_w = prepare_router_fp32_det_gemm(G.t(), X, torch.empty_like(dw), splits=6) + dw = cake_backend.cake_dense_gemm_det_train_router_wgrad(G, X) + prepared_w = cake_dense_gemm_det_train_router_prepare(G.t(), X, torch.empty_like(dw), splits=6) assert torch.equal(dw, prepared_w.launch()) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() assert cake_backend._ARENAS == {} @@ -9078,7 +9078,7 @@ def test_pd2_stream_k_on_an_odd_k_completes(): dict(pd=2, slots=2, epi="tma", pro=15), dict(pd=2, slots=1, epi="tma", pro=15), ) - prepared = prepare_dense_projection_det_gemm(v["A"], v["B"], v["out"], **knobs) + prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"], **knobs) plan = prepared.plan assert plan.pd == 2 and plan.k_blocks == 17 and plan.tail_tiles * 2 <= plan.sm_pairs assert ( @@ -9113,9 +9113,9 @@ def test_router_rows_match_fp64_reference_and_rerun_bit_exact(op, T): f"(compute capability {_registered_capabilities()})" ) v = _router_inputs(op, T, seed=SEED + 4) - plan, *_ = plan_router_fp32_det_gemm(v["A"], v["B"], v["out"]) + plan, *_ = cake_dense_gemm_det_train_router_plan(v["A"], v["B"], v["out"]) _require_program(plan.template) - prepared = prepare_router_fp32_det_gemm(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_router_prepare(v["A"], v["B"], v["out"]) assert prepared.splits == ROUTER_SPLITS[plan.layout] first = prepared.launch().clone() torch.cuda.synchronize() @@ -9139,11 +9139,11 @@ def test_router_rows_match_fp64_reference_and_rerun_bit_exact(op, T): assert int((top != ref_top).any(dim=1).sum()) == 0 # the eager entry point of the row gives the same bits eager_fn = dict( - fwd=lambda: cake_backend.router_det_forward(v["X"], v["W"]), - dgrad=lambda: cake_backend.router_det_dgrad(v["G"], v["W"]), - wgrad=lambda: cake_backend.router_det_wgrad(v["G"], v["X"]), + fwd=lambda: cake_backend.cake_dense_gemm_det_train_router_forward(v["X"], v["W"]), + dgrad=lambda: cake_backend.cake_dense_gemm_det_train_router_dgrad(v["G"], v["W"]), + wgrad=lambda: cake_backend.cake_dense_gemm_det_train_router_wgrad(v["G"], v["X"]), )[op] - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() eager = eager_fn() torch.cuda.synchronize() assert torch.equal(eager, v["out"]) @@ -9169,7 +9169,7 @@ def test_router_misaligned_out_is_refused_or_matches_fp64_reference(op, T): f"(compute capability {_registered_capabilities()})" ) v = _router_inputs(op, T, seed=SEED + 22) - plan, *_ = plan_router_fp32_det_gemm(v["A"], v["B"], v["out"]) + plan, *_ = cake_dense_gemm_det_train_router_plan(v["A"], v["B"], v["out"]) _require_program(plan.template) rows, cols = (int(s) for s in v["out"].shape) expected = torch.matmul(v["A"].double(), v["B"].double()) @@ -9177,25 +9177,25 @@ def test_router_misaligned_out_is_refused_or_matches_fp64_reference(op, T): aligned = _offset_out(rows, cols, 0, device="cuda") shifted = _offset_out(rows, cols, 1, device="cuda") eager = _offset_out(rows, cols, 1, device="cuda") - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() if splits == 1 and not plan.transposed_out: with pytest.raises(ValueError, match="16-byte aligned address"): - prepare_router_fp32_det_gemm(v["A"], v["B"], shifted, splits=splits) + cake_dense_gemm_det_train_router_prepare(v["A"], v["B"], shifted, splits=splits) with pytest.raises(ValueError, match="16-byte aligned address"): - cake_backend.router_fp32_det_gemm(v["A"], v["B"], eager, splits=splits) + cake_backend.cake_dense_gemm_det_train_router(v["A"], v["B"], eager, splits=splits) torch.cuda.synchronize() # nothing was launched: the context stays healthy - prepare_router_fp32_det_gemm( + cake_dense_gemm_det_train_router_prepare( v["A"], v["B"], aligned, splits=splits ).launch() outs = [aligned] else: - prepare_router_fp32_det_gemm( + cake_dense_gemm_det_train_router_prepare( v["A"], v["B"], aligned, splits=splits ).launch() - prepare_router_fp32_det_gemm( + cake_dense_gemm_det_train_router_prepare( v["A"], v["B"], shifted, splits=splits ).launch() - cake_backend.router_fp32_det_gemm(v["A"], v["B"], eager, splits=splits) + cake_backend.cake_dense_gemm_det_train_router(v["A"], v["B"], eager, splits=splits) outs = [shifted, eager] torch.cuda.synchronize() for out in outs: @@ -9255,19 +9255,19 @@ def test_fused_split_matches_the_torch_chain_bit_exactly(): # two paths agree bit for bit _require_program("router_fp32_det_gemm_kk") X = torch.randn(257, 6144, device=device, generator=g) - prepared = prepare_router_fp32_det_gemm( + prepared = cake_dense_gemm_det_train_router_prepare( X, W.t(), torch.empty(257, 256, device=device), splits=4 ) assert prepared.fused_split and "resid" not in prepared.tensors first = prepared.launch().clone() - cake_backend.clear_convenience_caches() - eager = cake_backend.router_det_forward(X, W) + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() + eager = cake_backend.cake_dense_gemm_det_train_router_forward(X, W) cached = next(iter(cake_backend._K2_CACHE.values())) assert cached.split is not None and "k2_resid" not in cached.arena.buffers assert "resid" not in cake_backend._entry_views(cached) torch.cuda.synchronize() assert torch.equal(first, eager) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() def test_sm107a_n128_wgrad_rules_plan(): @@ -9298,7 +9298,7 @@ def test_sm107a_n128_wgrad_rules_plan(): out = torch.empty(128, K, dtype=torch.float32 if f32 else torch.bfloat16) A, B, transposed = wgrad_views(G, X, arch="sm_107a", out_f32=f32) assert transposed - plan = plan_dense_projection_det_gemm( + plan = cake_dense_gemm_det_train_projection_plan( A, B, out, @@ -9335,9 +9335,9 @@ def test_repeated_launches_are_bitwise_identical(row, op, out_dtype, T): v = _make_inputs("proj", row, op, out_dtype, T, seed=SEED + 7) _require_program(_plan_template(v)) if op == "wgrad": - prepared = prepare_projection_det_wgrad(v["G"], v["X"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(v["G"], v["X"], v["out"]) else: - prepared = prepare_dense_projection_det_gemm(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"]) first = prepared.launch().clone() torch.cuda.synchronize() assert bool(torch.isfinite(first).all()) @@ -9358,9 +9358,9 @@ def test_repeated_router_launches_are_bitwise_identical(op, T): f"(compute capability {_registered_capabilities()})" ) v = _router_inputs(op, T, seed=SEED + 8) - plan, *_ = plan_router_fp32_det_gemm(v["A"], v["B"], v["out"]) + plan, *_ = cake_dense_gemm_det_train_router_plan(v["A"], v["B"], v["out"]) _require_program(plan.template) - prepared = prepare_router_fp32_det_gemm(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_router_prepare(v["A"], v["B"], v["out"]) first = prepared.launch().clone() torch.cuda.synchronize() assert bool(torch.isfinite(first).all()) @@ -9442,7 +9442,7 @@ def test_swapped_forward_matches_fp64_on_the_grouped_views(T, token_major): template is not None ) # the grouped wo_a rows fit two waves of a Blackwell part assert cake_jit.MODULES[launch.module_name]["template"] == template - cake_backend.dense_projection_det_gemm(A, B, out) + cake_backend.cake_dense_gemm_det_train_projection(A, B, out) torch.cuda.synchronize() _check(out, torch.einsum("gtd,gdr->gtr", A.double(), B.double())) @@ -9457,7 +9457,7 @@ def test_swapped_forward_other_shapes(L, T, K, N): launch = _swap_launch(A, B, out) assert launch is not None assert cake_jit.MODULES[launch.module_name]["template"] == _swap_expected(A, B) - cake_backend.dense_projection_det_gemm(A, B, out) + cake_backend.cake_dense_gemm_det_train_projection(A, B, out) torch.cuda.synchronize() _check(out, torch.einsum("gtd,gdr->gtr", A.double(), B.double())) @@ -9489,11 +9489,11 @@ def test_swapped_forward_is_bitwise_repeatable(): _require_swap() for T in (64, 128): A, B, out = _swap_views(T, seed=SEED + 2) - cake_backend.dense_projection_det_gemm(A, B, out) + cake_backend.cake_dense_gemm_det_train_projection(A, B, out) first = out.clone() for _ in range(REPEATS): out.fill_(float("nan")) - cake_backend.dense_projection_det_gemm(A, B, out, enable_pdl=True) + cake_backend.cake_dense_gemm_det_train_projection(A, B, out, enable_pdl=True) torch.cuda.synchronize() assert torch.equal(out, first) @@ -9520,8 +9520,8 @@ def test_swapped_forward_pdl_waits_for_an_early_producer(T): A2, B2, O2 = mid.transpose(0, 1), w2.transpose(1, 2), y.transpose(0, 1) assert _swap_launch(A1, B1, O1) is not None assert _swap_launch(A2, B2, O2) is not None - cake_backend.dense_projection_det_gemm(A1, B1, O1, enable_pdl=True) - cake_backend.dense_projection_det_gemm(A2, B2, O2, enable_pdl=True) + cake_backend.cake_dense_gemm_det_train_projection(A1, B1, O1, enable_pdl=True) + cake_backend.cake_dense_gemm_det_train_projection(A2, B2, O2, enable_pdl=True) torch.cuda.synchronize() assert bool(torch.isfinite(mid).all()) and bool(torch.isfinite(y).all()) _check(O2, torch.einsum("gtd,gdr->gtr", A2.double(), B2.double())) @@ -9536,8 +9536,8 @@ def test_swapped_forward_replays_in_a_cuda_graph_with_pdl(): A2 = out # the second projection reads the first's output def chain(): - cake_backend.dense_projection_det_gemm(A, B, out, enable_pdl=True) - cake_backend.dense_projection_det_gemm(A2, B2, out2, enable_pdl=True) + cake_backend.cake_dense_gemm_det_train_projection(A, B, out, enable_pdl=True) + cake_backend.cake_dense_gemm_det_train_projection(A2, B2, out2, enable_pdl=True) chain() torch.cuda.synchronize() diff --git a/tests/experimental/test_cake_dense_gemm_det_train_names.py b/tests/experimental/test_cake_dense_gemm_det_train_names.py new file mode 100644 index 00000000000..3918f0bcc3a --- /dev/null +++ b/tests/experimental/test_cake_dense_gemm_det_train_names.py @@ -0,0 +1,385 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +# The names of the deterministic Cake dense GEMMs (flashinfer-ai/flashinfer#8928): the public module +# ``flashinfer.cake_dense_gemm_det_train`` carries the mode word in every public name, no entry point takes a +# ``deterministic=`` keyword, the former names and import paths are gone (no alias), the module imports its backend +# only on first use, the deterministic entry points never consult the PyTorch deterministic-algorithms setting, +# and the generated registry, manifest and sources spell only their own names. CPU only, apart from the one +# device test at the end. + +import importlib +import importlib.util +import inspect +import json +import re +import subprocess +import sys +import warnings +from pathlib import Path + +import pytest +import torch + +MODULE = "flashinfer.cake_dense_gemm_det_train" +BACKEND = "flashinfer.experimental.cake.cake_dense_gemm_det_train" +MODE = "det" +# the 21 public names of the module +PUBLIC = { + "ARCHES", + "SUPPORTED_COMPUTE_CAPABILITIES", + "is_cake_dense_gemm_det_train_supported", + "cake_dense_gemm_det_train_projection", + "cake_dense_gemm_det_train_projection_forward", + "cake_dense_gemm_det_train_projection_dgrad", + "cake_dense_gemm_det_train_projection_wgrad", + "cake_dense_gemm_det_train_projection_plan", + "cake_dense_gemm_det_train_projection_prepare", + "cake_dense_gemm_det_train_projection_wgrad_prepare", + "CakeDenseGemmDetProjectionPlan", + "CakeDenseGemmDetProjectionLaunch", + "cake_dense_gemm_det_train_router", + "cake_dense_gemm_det_train_router_forward", + "cake_dense_gemm_det_train_router_dgrad", + "cake_dense_gemm_det_train_router_wgrad", + "cake_dense_gemm_det_train_router_plan", + "cake_dense_gemm_det_train_router_prepare", + "CakeDenseGemmDetRouterPlan", + "CakeDenseGemmDetRouterLaunch", + "cake_dense_gemm_det_train_convenience_cache_clear", +} +# The former import paths: the package and its modules. +FORMER_MODULES = ( + "flashinfer.experimental.dense_projection_det_gemm.cake_backend", + "flashinfer.experimental.dense_projection_det_gemm.cake_jit", + "flashinfer.experimental.dense_projection_det_gemm.cake_build", +) +FORMER_PACKAGE = "flashinfer.experimental.dense_projection_det_gemm" +# Removed backend names: the former public names and the aliases of the names before the det rename. +REMOVED = { + "cake_backend": ( + "DetGemmPlan", + "PreparedDetGemm", + "RouterDetPlan", + "PreparedRouterDetGemm", + "plan_dense_projection_det_gemm", + "prepare_dense_projection_det_gemm", + "dense_projection_det_gemm", + "projection_det_forward", + "projection_det_dgrad", + "prepare_projection_det_wgrad", + "projection_det_wgrad", + "plan_router_fp32_det_gemm", + "prepare_router_fp32_det_gemm", + "router_fp32_det_gemm", + "router_det_forward", + "router_det_dgrad", + "router_det_wgrad", + "clear_convenience_caches", + "GemmPlan", + "PreparedGemm", + "PreparedRouterGemm", + "RouterPlan", + "dense_projection_gemm", + "plan_dense_projection_gemm", + "plan_router_fp32_gemm", + "prepare_dense_projection_gemm", + "prepare_projection_wgrad", + "prepare_router_fp32_gemm", + "projection_dgrad", + "projection_forward", + "projection_wgrad", + "router_dgrad", + "router_forward", + "router_fp32_gemm", + "router_wgrad", + "split_fp32_to_bf16x3_", + ), + "cake_jit": ("load_cake_dense_projection_gemm_module",), +} +# Spellings of the non-deterministic package that the deterministic registry, manifest and sources never contain. +FORMER_TOKENS = ( + "dense_projection_gemm", + "dense_proj_gemm", + "router_fp32_gemm", + "split_fp32_bf16x3", + "DENSE_PROJ_GEMM", +) +CONSTANT = re.compile(r"^[A-Z][A-Z0-9_]*$") + + +def _tokens(name: str) -> list[str]: + """The words of ``name``: split on ``_`` and on CamelCase boundaries, lower-cased.""" + words = [] + for part in name.split("_"): + words += re.findall(r"[A-Z]+(?![a-z])|[A-Z]?[a-z0-9]+", part) + return [w.lower() for w in words if w] + + +def _run(code: str) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, timeout=600 + ) + + +def test_all_lists_the_public_names(): + module = importlib.import_module(MODULE) + assert set(module.__all__) == PUBLIC + assert len(module.__all__) == len(PUBLIC) + + +@pytest.mark.parametrize("name", sorted(PUBLIC)) +def test_every_public_name_carries_the_mode_word(name): + if CONSTANT.match(name): + return # module constants: the module path carries the mode + tokens = _tokens(name) + assert MODE in tokens and "nondet" not in tokens, (name, tokens) + if name[0].isupper(): + assert name.startswith("CakeDenseGemmDet") and name[16].isupper() + else: + assert name.startswith("cake_dense_gemm_det_train_") or ( + name == "is_cake_dense_gemm_det_train_supported" + ) + assert not {"fwd", "bwd"} & set(tokens) + + +@pytest.mark.parametrize("name", sorted(n for n in PUBLIC if not CONSTANT.match(n))) +def test_no_entry_point_takes_a_mode_keyword(name): + obj = getattr(importlib.import_module(MODULE), name) + assert callable(obj) + signature = inspect.signature(obj) + with pytest.raises(TypeError): + signature.bind_partial(deterministic=False) + assert not any( + p.kind is inspect.Parameter.VAR_KEYWORD for p in signature.parameters.values() + ) + if isinstance(obj, type) and issubclass(obj, torch.autograd.Function): + with pytest.raises(TypeError): + inspect.signature(obj.forward).bind_partial(None, deterministic=False) + + +@pytest.mark.parametrize("name", sorted(PUBLIC)) +def test_the_module_returns_the_backend_object(name): + with warnings.catch_warnings(): + warnings.simplefilter("error") # no ExperimentalWarning: a stable module + obj = getattr(importlib.import_module(MODULE), name) + submodule = "cake_jit" if name == "ARCHES" else "cake_backend" + assert obj is getattr(importlib.import_module(f"{BACKEND}.{submodule}"), name) + if inspect.isfunction(obj) or inspect.isclass(obj): + assert obj.__name__ == name + assert inspect.unwrap(obj).__name__ == name + assert obj.__module__ == f"{BACKEND}.cake_backend" + + +def test_an_unknown_name_raises_attribute_error(): + module = importlib.import_module(MODULE) + with pytest.raises(AttributeError): + module.dense_projection_det_gemm # noqa: B018 + + +@pytest.mark.parametrize("submodule, names", sorted(REMOVED.items())) +def test_the_former_names_are_removed(submodule, names): + backend = importlib.import_module(f"{BACKEND}.{submodule}") + assert [n for n in names if hasattr(backend, n)] == [] + + +@pytest.mark.parametrize("name", FORMER_MODULES) +def test_the_former_modules_fail_at_import(name): + result = _run(f"import importlib; importlib.import_module({name!r})") + assert result.returncode != 0 + assert "ModuleNotFoundError" in result.stderr + + +def test_the_former_package_is_gone(): + # A checkout may keep an untracked ``aot/`` or ``__pycache__/`` in the former directory, which Python + # imports as a namespace package (no ``origin``); it holds no module. + spec = importlib.util.find_spec(FORMER_PACKAGE) + assert spec is None or spec.origin is None + + +def test_imports_are_lazy(): + code = f""" +import sys +import flashinfer +loaded = sorted(m for m in sys.modules if m.startswith("flashinfer.experimental.cake.cake_dense_gemm")) +assert loaded == [], loaded +for name in {sorted(PUBLIC)!r}: + assert name not in vars(flashinfer), name +rows = getattr(flashinfer, "_LAZY_EXPORTS", {{}}) +dense = {{k: v for k, v in rows.items() if "cake_dense_gemm" in str(v[0])}} +assert all(v[1] is None for v in dense.values()), dense +import {MODULE} as m +loaded = sorted(n for n in sys.modules if n.startswith("{BACKEND}")) +assert loaded == [], loaded +m.cake_dense_gemm_det_train_projection +assert "{BACKEND}.cake_backend" in sys.modules +print("ok") +""" + result = _run(code) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "ok" + + +def test_the_backend_never_reads_the_torch_deterministic_setting(): + """The module is the mode: the deterministic entry points neither alert nor change path under + ``torch.use_deterministic_algorithms`` (strict or ``warn_only``).""" + root = Path(importlib.import_module(f"{BACKEND}.cake_jit").__file__).parent + for path in sorted(root.glob("*.py")): + text = path.read_text() + for token in ( + "are_deterministic_algorithms_enabled", + "is_deterministic_algorithms_warn_only_enabled", + "_alert_not_deterministic", + "use_deterministic_algorithms", + ): + assert token not in text, (path.name, token) + + +def _former_lines(text: str) -> list[str]: + return [line for line in text.split("\n") if any(t in line for t in FORMER_TOKENS)] + + +def test_the_registry_spells_only_canonical_names(): + jit = importlib.import_module(f"{BACKEND}.cake_jit") + backend = importlib.import_module(f"{BACKEND}.cake_backend") + assert jit.MODULES and jit.KERNELS + # The generated names (library stems, program prefixes, kernel namespaces, templates) are build identities + # and keep their spelling across the module rename. + for name, record in jit.MODULES.items(): + assert ( + name.startswith("cake_dense_projection_det_gemm_") + and record["library"] == name + ) + assert record["kernel"] in (f"kernel_{name}",) or record["kernel"].startswith( + "cake::dense_projection_det_gemm::" + ) + assert record["template"].startswith( + ( + "dense_proj_det_gemm_", + "router_fp32_det_gemm_", + "split_fp32_det_bf16x3", + "swap_ab_", + ) + ) + assert ( + "build" not in record + or ".dense_projection_gemm" not in record["build"]["module"] + ) + templates = {t for table in jit.KERNELS.values() for t in table} + assert backend.SPLIT_TEMPLATE in templates + assert { + backend.router_symbol(a, b, t) + for a, b, t in ((False, False, False), (False, True, False), (True, True, True)) + } <= templates + assert backend._TEMPLATE_PREFIX == "dense_proj_det_gemm_" + assert all( + t.startswith(backend._TEMPLATE_PREFIX) + for t in templates + if t.startswith("dense_proj") + ) + assert _former_lines(json.dumps([jit.MODULES, jit.KERNELS])) == [] + + +def test_no_former_spelling_in_the_package(): + root = Path(importlib.import_module(f"{BACKEND}.cake_jit").__file__).parent + manifest = json.loads((root / "build.json").read_text()) + assert manifest["package"] == BACKEND + for entry in manifest["libraries"]: + assert entry["name"].startswith("cake_dense_projection_det_gemm_") + assert all( + s.startswith( + ("csrc/cake_dense_projection_det_gemm/", "csrc/cake_device_helpers/") + ) + for s in entry["sources"] + ) + files = [ + root / "build.json", + *sorted(root.glob("*.py")), + *sorted((root / "csrc").rglob("*")), + ] + found = { + str(path.relative_to(root)): lines + for path in files + if path.is_file() and (lines := _former_lines(path.read_text(errors="replace"))) + } + assert found == {} + assert not [ + p for p in (root / "csrc").rglob("*") if any(t in p.name for t in FORMER_TOKENS) + ] + + +def _supported_device(): + if not torch.cuda.is_available(): + return None + module = importlib.import_module(MODULE) + device = torch.device("cuda", torch.cuda.current_device()) + return device if module.is_cake_dense_gemm_det_train_supported(device) else None + + +def test_the_support_query(): + module = importlib.import_module(MODULE) + assert module.is_cake_dense_gemm_det_train_supported(torch.device("cpu")) is False + if torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + assert module.is_cake_dense_gemm_det_train_supported( + torch.device("cuda", torch.cuda.current_device()) + ) == ((major, minor) in module.SUPPORTED_COMPUTE_CAPABILITIES) + + +@pytest.mark.parametrize("warn_only", [False, True]) +def test_the_entry_points_run_unchanged_under_torch_deterministic_mode(warn_only): + device = _supported_device() + if device is None: + pytest.skip("needs an SM100 / SM103 / SM107 GPU") + m = importlib.import_module(MODULE) + g = torch.Generator(device=device).manual_seed(0) + T, K, N, R = 257, 6144, 576, 256 # the kv_a projection and the 256-expert router + X = torch.randn(T, K, device=device, generator=g).bfloat16() + W = (torch.randn(N, K, device=device, generator=g) * K**-0.5).bfloat16() + G = torch.randn(T, N, device=device, generator=g).bfloat16() + Xr = torch.randn(T, K, device=device, generator=g) + Wr = torch.randn(R, K, device=device, generator=g) * K**-0.5 + Gr = torch.randn(T, R, device=device, generator=g) + + def run(): + return [ + m.cake_dense_gemm_det_train_projection_forward(X, W), + m.cake_dense_gemm_det_train_projection_dgrad(G, W), + m.cake_dense_gemm_det_train_projection_wgrad(G, X), + m.cake_dense_gemm_det_train_router_forward(Xr, Wr), + m.cake_dense_gemm_det_train_router_dgrad(Gr, Wr), + m.cake_dense_gemm_det_train_router_wgrad(Gr, Xr), + ] + + try: + reference = run() + except NotImplementedError as exc: # a build without one of these programs + pytest.skip(str(exc)) + previous = ( + torch.are_deterministic_algorithms_enabled(), + torch.is_deterministic_algorithms_warn_only_enabled(), + ) + try: + torch.use_deterministic_algorithms(True, warn_only=warn_only) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + outputs = run() + torch.cuda.synchronize(device) + finally: + torch.use_deterministic_algorithms(previous[0], warn_only=previous[1]) + assert [w for w in caught if "determinis" in str(w.message)] == [] + for got, want in zip(outputs, reference, strict=True): + assert torch.equal(got, want) diff --git a/tests/experimental/test_cake_dense_projection_det_gemm_names.py b/tests/experimental/test_cake_dense_projection_det_gemm_names.py deleted file mode 100644 index 97c13ef4929..00000000000 --- a/tests/experimental/test_cake_dense_projection_det_gemm_names.py +++ /dev/null @@ -1,220 +0,0 @@ -""" -Copyright (c) 2026 by FlashInfer team. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - -http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -""" - -# The dense projection GEMM package's names: the canonical (det) names, the former names kept as aliases of the -# same objects (functions, classes), the former import path (now the package with an explicit numerical mode, -# whose default runs these canonical objects), and no former spelling left in the generated registry, manifest or -# sources. CPU only: nothing here launches a kernel. - -import importlib -import inspect -import json -import re -from pathlib import Path - -import pytest - -PACKAGE = "flashinfer.experimental.dense_projection_det_gemm" -FORMER_PACKAGE = "flashinfer.experimental.dense_projection_gemm" -# former name -> canonical name of the public entry points of ``cake_backend`` -API = { - "GemmPlan": "DetGemmPlan", - "PreparedGemm": "PreparedDetGemm", - "PreparedRouterGemm": "PreparedRouterDetGemm", - "RouterPlan": "RouterDetPlan", - "dense_projection_gemm": "dense_projection_det_gemm", - "plan_dense_projection_gemm": "plan_dense_projection_det_gemm", - "plan_router_fp32_gemm": "plan_router_fp32_det_gemm", - "prepare_dense_projection_gemm": "prepare_dense_projection_det_gemm", - "prepare_projection_wgrad": "prepare_projection_det_wgrad", - "prepare_router_fp32_gemm": "prepare_router_fp32_det_gemm", - "projection_dgrad": "projection_det_dgrad", - "projection_forward": "projection_det_forward", - "projection_wgrad": "projection_det_wgrad", - "router_dgrad": "router_det_dgrad", - "router_forward": "router_det_forward", - "router_fp32_gemm": "router_fp32_det_gemm", - "router_wgrad": "router_det_wgrad", - "split_fp32_to_bf16x3_": "split_fp32_to_bf16x3_det_", -} -LOADER = ( - "load_cake_dense_projection_gemm_module", - "load_cake_dense_projection_det_gemm_module", -) -FORMER_TOKENS = ( - "dense_projection_gemm", - "dense_proj_gemm", - "router_fp32_gemm", - "split_fp32_bf16x3", - "DENSE_PROJ_GEMM", -) -MODULES = ("cake_backend", "cake_jit", "cake_build") - - -def _module(package: str, name: str): - return importlib.import_module(f"{package}.{name}") - - -@pytest.mark.parametrize("former, canonical", sorted(API.items())) -def test_former_names_are_the_canonical_objects(former, canonical): - backend = _module(PACKAGE, "cake_backend") - obj = getattr(backend, canonical) - assert getattr(backend, former) is obj - assert obj.__name__ == canonical - assert inspect.unwrap(obj).__name__ == canonical - assert obj.__module__ == f"{PACKAGE}.cake_backend" - former_backend = _module(FORMER_PACKAGE, "cake_backend") - assert former_backend.DETERMINISTIC_NAMES[former] == canonical - if inspect.isclass(obj): - assert getattr(former_backend, former) is obj - else: - entry = getattr(former_backend, former) - assert entry is not obj - assert inspect.signature(entry).parameters["deterministic"].default is True - - -# former entry point -> its positional arguments (sentinels) -ARITY = { - "dense_projection_gemm": 3, - "plan_dense_projection_gemm": 3, - "plan_router_fp32_gemm": 3, - "prepare_dense_projection_gemm": 3, - "prepare_projection_wgrad": 3, - "prepare_router_fp32_gemm": 3, - "projection_dgrad": 2, - "projection_forward": 2, - "projection_wgrad": 2, - "router_dgrad": 2, - "router_forward": 2, - "router_fp32_gemm": 3, - "router_wgrad": 2, - "split_fp32_to_bf16x3_": 3, -} - - -@pytest.mark.parametrize("former", sorted(ARITY)) -def test_the_former_path_defaults_to_the_canonical_object(monkeypatch, former): - """The default of the former path's entry point is the canonical object with the caller's arguments: the same - program, so the same bits (the numerical mode keeps the deterministic default).""" - calls = [] - result = object() - - def canonical(*args, **kwargs): - calls.append(args) - return result - - monkeypatch.setattr(_module(PACKAGE, "cake_backend"), API[former], canonical) - args = tuple(object() for _ in range(ARITY[former])) - assert getattr(_module(FORMER_PACKAGE, "cake_backend"), former)(*args) is result - assert len(calls) == 1 and calls[0][: len(args)] == args - - -def test_the_loader_keeps_its_former_name(): - jit = _module(PACKAGE, "cake_jit") - loader = getattr(jit, LOADER[1]) - assert getattr(jit, LOADER[0]) is loader - assert inspect.unwrap(loader).__name__ == LOADER[1] - - -@pytest.mark.parametrize("name", MODULES) -def test_the_former_package_path_is_its_own_package(name): - """The former import path is the package with the numerical mode (its own modules and programs); the - canonical package never imports it.""" - assert _module(FORMER_PACKAGE, name) is not _module(PACKAGE, name) - former = importlib.import_module(FORMER_PACKAGE) - assert Path(former.__file__).parent.name == "dense_projection_gemm" - - -def _former_lines(text: str) -> list[str]: - """The lines of ``text`` spelling a former token, an alias assignment of the former names aside.""" - alias = re.compile(r"^(?P\w+) = (?P\w+)$") - out = [] - for line in text.split("\n"): - if not any(token in line for token in FORMER_TOKENS): - continue - match = alias.match(line) - known = dict(API, **{LOADER[0]: LOADER[1]}) - if match and known.get(match["former"]) == match["canonical"]: - continue - out.append(line) - return out - - -def test_the_registry_spells_only_canonical_names(): - jit = _module(PACKAGE, "cake_jit") - backend = _module(PACKAGE, "cake_backend") - assert jit.MODULES and jit.KERNELS - for name, record in jit.MODULES.items(): - assert ( - name.startswith("cake_dense_projection_det_gemm_") - and record["library"] == name - ) - assert record["kernel"] in (f"kernel_{name}",) or record["kernel"].startswith( - "cake::dense_projection_det_gemm::" - ) - assert record["template"].startswith( - ( - "dense_proj_det_gemm_", - "router_fp32_det_gemm_", - "split_fp32_det_bf16x3", - "swap_ab_", - ) - ) - assert ( - "build" not in record - or ".dense_projection_gemm" not in record["build"]["module"] - ) - templates = {t for table in jit.KERNELS.values() for t in table} - assert backend.SPLIT_TEMPLATE in templates - assert { - backend.router_symbol(a, b, t) - for a, b, t in ((False, False, False), (False, True, False), (True, True, True)) - } <= templates - assert backend._TEMPLATE_PREFIX == "dense_proj_det_gemm_" - assert all( - t.startswith(backend._TEMPLATE_PREFIX) - for t in templates - if t.startswith("dense_proj") - ) - assert _former_lines(json.dumps([jit.MODULES, jit.KERNELS])) == [] - - -def test_no_former_spelling_in_the_package(): - root = Path(_module(PACKAGE, "cake_jit").__file__).parent - manifest = json.loads((root / "build.json").read_text()) - assert manifest["package"] == PACKAGE - for entry in manifest["libraries"]: - assert entry["name"].startswith("cake_dense_projection_det_gemm_") - assert all( - s.startswith( - ("csrc/cake_dense_projection_det_gemm/", "csrc/cake_device_helpers/") - ) - for s in entry["sources"] - ) - files = [ - root / "build.json", - *sorted(root.glob("*.py")), - *sorted((root / "csrc").rglob("*")), - ] - found = { - str(path.relative_to(root)): lines - for path in files - if path.is_file() and (lines := _former_lines(path.read_text(errors="replace"))) - } - assert found == {} - assert not [ - p for p in (root / "csrc").rglob("*") if any(t in p.name for t in FORMER_TOKENS) - ] diff --git a/tests/test_helpers/cake_package.py b/tests/test_helpers/cake_package.py index 6b0efea9d4c..17cfaa33d70 100644 --- a/tests/test_helpers/cake_package.py +++ b/tests/test_helpers/cake_package.py @@ -30,7 +30,7 @@ "cake_lm_head_det_loss": "flashinfer.experimental.cake_lm_head_det_loss", # the non-deterministic package (deterministic=False); its key carries the namespace "cake.cake_lm_head_loss": "flashinfer.experimental.cake.cake_lm_head_loss", - "dense_projection_det_gemm": "flashinfer.experimental.dense_projection_det_gemm", + "cake_dense_gemm_det_train": "flashinfer.experimental.cake.cake_dense_gemm_det_train", "dense_projection_gemm": "flashinfer.experimental.dense_projection_gemm", "cake_moe_grouped_det_gemm": "flashinfer.experimental.cake_moe_grouped_det_gemm", # the suite name before the deterministic-family rename (compatibility alias) From f39592aa3303a92f32d3445513b54ea57ea7fc54 Mon Sep 17 00:00:00 2001 From: Avery Huang Date: Thu, 5 Nov 2026 23:58:45 -0800 Subject: [PATCH 06/13] test(cake_dense_gemm_det_train): run every names check on CPU; formatting - The names test checks every det entry point under torch.use_deterministic_algorithms (strict and warn_only) on CPU: no deterministic-mode error or warning, and the planners still return their plans. The GPU check (same outputs, bitwise, with and without the mode) moves to the det test. - The ExperimentalWarning check runs in a fresh process; the former-module check asserts that the missing module is the former one. - ruff format of the renamed lines; the benchmark usage, a source citation and the README name patterns name the new entries. --- benchmarks/bench_cake_dense_gemm_det_train.py | 12 +- flashinfer/cake_dense_gemm_det_train.py | 2 +- .../cake/cake_dense_gemm_det_train/README.md | 12 +- .../cake_dense_gemm_det_train/cake_backend.py | 37 +- .../cake_dense_gemm_det_train/cake_jit.py | 1 - .../test_cake_dense_gemm_det_train.py | 438 ++++++++++++++---- .../test_cake_dense_gemm_det_train_names.py | 146 ++++-- 7 files changed, 484 insertions(+), 164 deletions(-) diff --git a/benchmarks/bench_cake_dense_gemm_det_train.py b/benchmarks/bench_cake_dense_gemm_det_train.py index b60fa5a1b9b..4ad9ab95bd4 100644 --- a/benchmarks/bench_cake_dense_gemm_det_train.py +++ b/benchmarks/bench_cake_dense_gemm_det_train.py @@ -14,7 +14,7 @@ limitations under the License. """ -"""Benchmark the experimental Cake dense projection GEMMs (SM100 / SM103 / SM107). +"""Benchmark the deterministic Cake dense GEMMs (SM100 / SM103 / SM107). Times the GLM-5.2 training GEMMs of ``flashinfer.experimental.cake.cake_dense_gemm_det_train`` paired against ``torch.matmul`` on the same strided views: @@ -38,7 +38,7 @@ Usage:: - python benchmarks/bench_cake_dense_projection_det_gemm.py [--rows kv_a_fwd_bf16_t16231 ...] + python benchmarks/bench_cake_dense_gemm_det_train.py [--rows kv_a_fwd_bf16_t16231 ...] [--ops fwd,dgrad,wgrad] [--T 16231,16172] [--f32-grads] [--no-mla] [--no-router] [--steps 20] [--accuracy] [--host-overhead] [--json out.json] """ @@ -208,9 +208,13 @@ def eager_arm(spec, inp): """The convenience entry point of the row (the cached planned launch; ``out`` is passed, so a cache hit allocates nothing).""" if spec["family"] == "router": - return lambda: cake_backend.cake_dense_gemm_det_train_router(inp["A"], inp["B"], inp["out"]) + return lambda: cake_backend.cake_dense_gemm_det_train_router( + inp["A"], inp["B"], inp["out"] + ) if spec["family"] == "proj" and spec["op"] == "wgrad": - return lambda: cake_backend.cake_dense_gemm_det_train_projection_wgrad(inp["G"], inp["X"], inp["out"]) + return lambda: cake_backend.cake_dense_gemm_det_train_projection_wgrad( + inp["G"], inp["X"], inp["out"] + ) return lambda: cake_backend.cake_dense_gemm_det_train_projection( inp["A"], inp["B"], inp["out"] ) diff --git a/flashinfer/cake_dense_gemm_det_train.py b/flashinfer/cake_dense_gemm_det_train.py index 8cab184353d..25dad650695 100644 --- a/flashinfer/cake_dense_gemm_det_train.py +++ b/flashinfer/cake_dense_gemm_det_train.py @@ -84,4 +84,4 @@ def __getattr__(name: str) -> Any: def __dir__() -> list[str]: - return sorted(set(globals()) | set(__all__)) + return list(__all__) diff --git a/flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md b/flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md index 9c4018b78f0..7c9fde2fd3d 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/README.md @@ -67,12 +67,12 @@ prepared.launch() # re-splits B, launches, reduces the split-K partials; .spli Both paths run the same planner once per binding and bind the same generated module; they differ in who owns the plan and the workspace. -| | prepared path (`prepare_*` + `launch()`) | convenience path (`cake_dense_gemm_det_train_projection`, `projection_*`, `router_*`) | +| | prepared path (`*_prepare` + `launch()`) | convenience path (`cake_dense_gemm_det_train_{projection,router}` and their `_forward` / `_dgrad` / `_wgrad`) | |---|---|---| -| planning | once, in `prepare_*` | once per cache key, on the first call | +| planning | once, in `*_prepare` | once per cache key, on the first call | | workspace (K1: stream-K slabs + slice counters; K2: BF16x3 stack, the torch-chain split's fp32 scratch, split-K partials) | owned by the prepared object | carved out of the workspace arena of the (device, stream): one growable buffer per kind shared by every cached entry of that device and stream, grown to the largest requirement seen, never per entry | | per call | bind nothing: launch the bound views | rebind the caller's views (view objects only), resolve the workspace views of the arena (memoised until the arena grows), launch; `out` is allocated only when the caller passes none | -| device queries | none after `prepare_*` | none after the first call | +| device queries | none after `*_prepare` | none after the first call | | cache | none | per process, per kind (K1 / K2), `CONVENIENCE_CACHE_SIZE` = 64 entries, least recently used first out; the arena of a (device, stream) is released with its last entry and by `cake_dense_gemm_det_train_convenience_cache_clear()` | | memory held | the prepared object's workspace | bounded by the largest geometry seen per workspace kind, however many token counts the cache has planned (a ragged T adds entries, not workspace) | | CUDA Graphs | capture `launch()` | do not capture (a later growth replaces, an eviction frees, the workspace a graph replays) | @@ -156,7 +156,7 @@ hints); every rule is a verbatim copy of the Cake kernel host's (each function names its source line range). Two device facts enter the plan: the SM count (CTA pairs: stream-K split, raster group, wave working set) and the L2 size (hint gate); -`prepare_*` and the convenience path read both once per device index +`*_prepare` and the convenience path read both once per device index (`device_facts`: `device_sm_count`, `device_l2_bytes` = `torch.cuda.get_device_properties(...).L2_cache_size`, falling back to `cuDeviceGetAttribute`), `cake_dense_gemm_det_train_projection_plan` takes them as inputs: @@ -296,8 +296,8 @@ export has generated and for every template on a device whose architecture the registry does not list. CUDA Graphs: prepare outside capture and replay the prepared launch -(`prepare_*` once, `launch` many times). The convenience entry points -(`cake_dense_gemm_det_train_projection`, `projection_*`, `router_*`) reuse a cached planned +(`*_prepare` once, `launch` many times). The convenience entry points +(`cake_dense_gemm_det_train_{projection,router}` and their `_forward` / `_dgrad` / `_wgrad`) reuse a cached planned launch per (device, stream, geometry) key and the shared workspace arena of its (device, stream); capture a prepared launch object rather than the convenience entry points (see "Host paths"). diff --git a/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py index eab4407f7e6..af0f4a7e97c 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_backend.py @@ -1943,7 +1943,9 @@ def stream_k_plan( return pair_tiles - tail, tail, units, iters_per_unit -def check_stream_k_plan(plan: "CakeDenseGemmDetProjectionPlan", arch: Optional[str] = None) -> None: +def check_stream_k_plan( + plan: "CakeDenseGemmDetProjectionPlan", arch: Optional[str] = None +) -> None: """Refuse a stream-K plan outside the program's launch contract (the Cake kernel's ``requires``) before it launches. The programs that hand off the work inside an item's last segment (``pd`` 2, ``sk_slab`` 2, ``stt``) never finish an item that holds no K step: the linearised plan needs ``iters_per_unit >= 1`` and @@ -2420,7 +2422,9 @@ def cake_dense_gemm_det_train_projection_plan( if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: raise ValueError("cake_dense_gemm_det_train_projection: A and B must be bf16") if out.dtype not in (torch.bfloat16, torch.float32): - raise ValueError("cake_dense_gemm_det_train_projection: out must be bf16 or fp32") + raise ValueError( + "cake_dense_gemm_det_train_projection: out must be bf16 or fp32" + ) A3, B3, O3 = as_batched(A, "A"), as_batched(B, "B"), as_batched(out, "out") L, M, K = (int(s) for s in A3.shape) if tuple(B3.shape) != (L, K, int(B3.shape[2])): @@ -3957,7 +3961,9 @@ def _k1_plan( *, transposed_out: bool, knobs: dict[str, Any], -) -> tuple[CakeDenseGemmDetProjectionPlan, torch.Tensor, torch.Tensor, torch.Tensor, str, str]: +) -> tuple[ + CakeDenseGemmDetProjectionPlan, torch.Tensor, torch.Tensor, torch.Tensor, str, str +]: """Plan one K1 binding for its device: ``(plan, a_desc, b_desc, out3, arch, module_name)`` (the device facts come from the per-device cache, the module from the registry; see :func:`cake_dense_gemm_det_train_projection_plan`).""" device = _same_cuda_device(A, B, out) @@ -4176,7 +4182,11 @@ def cake_dense_gemm_det_train_projection_prepare( def _k1_swapped( - plan: CakeDenseGemmDetProjectionPlan, A: torch.Tensor, B: torch.Tensor, *, transposed_out: bool + plan: CakeDenseGemmDetProjectionPlan, + A: torch.Tensor, + B: torch.Tensor, + *, + transposed_out: bool, ) -> bool: """Did the planner run the swapped GEMM (``B.T @ A.T`` with the transposed store) for these views? The launcher's small-M swap rule on the caller's dimensions, unless the registry fallback undid the swap @@ -4554,7 +4564,7 @@ def wgrad_views( when ``wgrad_swapped`` the swapped GEMM ``X.T @ G`` runs with the transposed store (the result lands directly in ``[N, K]``), otherwise ``G.T @ X``. ``arch`` / ``out_f32`` let a measured rule of that architecture and output dtype decide the orientation (see ``wgrad_swapped``). - [Cake ``cake_dense_gemm_det_train_projection_wgrad`` L1507-L1518 / ``wgrad_swapped`` L1496-L1504]""" + [Cake K1 kernel host weight-gradient entry L1507-L1518 / ``wgrad_swapped`` L1496-L1504]""" N, K = int(G.shape[1]), int(X.shape[1]) if wgrad_swapped(N, K, arch=arch, out_f32=out_f32): return X.t(), G, True @@ -4698,7 +4708,9 @@ def _router_a_layout(A: torch.Tensor) -> bool: return False if A.stride(0) == 1: return True - raise ValueError("cake_dense_gemm_det_train_router: A needs a unit stride on m or k") + raise ValueError( + "cake_dense_gemm_det_train_router: A needs a unit stride on m or k" + ) def _router_b_layout(B: torch.Tensor) -> bool: @@ -4706,7 +4718,9 @@ def _router_b_layout(B: torch.Tensor) -> bool: return False if B.stride(1) == 1: # n contiguous -> MN-major parts [3, K, N] return True - raise ValueError("cake_dense_gemm_det_train_router: B needs a unit stride on k or n") + raise ValueError( + "cake_dense_gemm_det_train_router: B needs a unit stride on k or n" + ) def cake_dense_gemm_det_train_router_plan( @@ -4728,7 +4742,9 @@ def cake_dense_gemm_det_train_router_plan( ): raise ValueError("cake_dense_gemm_det_train_router: A, B and out must be fp32") if A.dim() != 2 or B.dim() != 2 or out.dim() != 2: - raise ValueError("cake_dense_gemm_det_train_router: A, B and out must be 2-D views") + raise ValueError( + "cake_dense_gemm_det_train_router: A, B and out must be 2-D views" + ) M, K = (int(s) for s in A.shape) Kb, N = (int(s) for s in B.shape) if Kb != K: @@ -4759,7 +4775,9 @@ def cake_dense_gemm_det_train_router_plan( a_eff_mn, b_eff_mn = _router_a_layout(A_eff), _router_b_layout(B_eff) lda = int(A_eff.stride(1)) if a_eff_mn else int(A_eff.stride(0)) if lda % 4 != 0 or A_eff.data_ptr() % 16 != 0: - raise ValueError("cake_dense_gemm_det_train_router: A rows must be 16-byte aligned") + raise ValueError( + "cake_dense_gemm_det_train_router: A rows must be 16-byte aligned" + ) if B_eff.data_ptr() % 16 != 0: raise ValueError( "cake_dense_gemm_det_train_router: B must start at a 16-byte aligned address" @@ -5260,4 +5278,3 @@ def cake_dense_gemm_det_train_router_wgrad( if out is None: out = torch.empty(G.shape[1], X.shape[1], device=G.device, dtype=torch.float32) return cake_dense_gemm_det_train_router(G.t(), X, out, splits=splits) - diff --git a/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py index 0e7f9f3795b..5a6122001bb 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_det_train/cake_jit.py @@ -290,7 +290,6 @@ def load_cake_dense_projection_det_gemm_module(name: str, arch: str): return cake_build.load_library(library_name(name, arch)) - def specs_for_arch(arch: str): """JIT specs of every ``build.json`` entry of ``arch`` without a prebuilt library (the ``flashinfer.jit.build_jit_specs`` warm-up).""" diff --git a/tests/experimental/test_cake_dense_gemm_det_train.py b/tests/experimental/test_cake_dense_gemm_det_train.py index 8a6f26b6d27..8ae660af156 100644 --- a/tests/experimental/test_cake_dense_gemm_det_train.py +++ b/tests/experimental/test_cake_dense_gemm_det_train.py @@ -19,6 +19,7 @@ import dataclasses import re import threading +import warnings from collections import OrderedDict from pathlib import Path from types import SimpleNamespace @@ -68,11 +69,21 @@ sk_sync_plan = cake_backend.sk_sync_plan check_stream_k_plan = cake_backend.check_stream_k_plan operand_view = cake_backend.operand_view -cake_dense_gemm_det_train_projection_plan = cake_backend.cake_dense_gemm_det_train_projection_plan -cake_dense_gemm_det_train_router_plan = cake_backend.cake_dense_gemm_det_train_router_plan -cake_dense_gemm_det_train_projection_prepare = cake_backend.cake_dense_gemm_det_train_projection_prepare -cake_dense_gemm_det_train_projection_wgrad_prepare = cake_backend.cake_dense_gemm_det_train_projection_wgrad_prepare -cake_dense_gemm_det_train_router_prepare = cake_backend.cake_dense_gemm_det_train_router_prepare +cake_dense_gemm_det_train_projection_plan = ( + cake_backend.cake_dense_gemm_det_train_projection_plan +) +cake_dense_gemm_det_train_router_plan = ( + cake_backend.cake_dense_gemm_det_train_router_plan +) +cake_dense_gemm_det_train_projection_prepare = ( + cake_backend.cake_dense_gemm_det_train_projection_prepare +) +cake_dense_gemm_det_train_projection_wgrad_prepare = ( + cake_backend.cake_dense_gemm_det_train_projection_wgrad_prepare +) +cake_dense_gemm_det_train_router_prepare = ( + cake_backend.cake_dense_gemm_det_train_router_prepare +) router_layout_class = cake_backend.router_layout_class row_rule = cake_backend.row_rule sk_parts_plan = cake_backend.sk_parts_plan @@ -1933,7 +1944,9 @@ def test_round19_cta1_plans_one_cta_per_tile(T): # policy, and the raster group may be any count >= 1 [Cake launcher] v = _views("proj", "indexer_hw", "fwd", "bf16", T) kw = dict(sm_count=148, l2_bytes=L2_BYTES, arch="sm_100a", _fallback=False) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], cta1=True, **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], cta1=True, **kw + ) assert plan.cta1 and _nopro(plan.template) == "dense_proj_det_gemm_kk_n128_hen_c1" assert (plan.m_tiles, plan.n_tiles, plan.pair_tiles, plan.sm_pairs) == ( 127, @@ -2029,7 +2042,9 @@ def test_round19_rules_plan_like_the_cake_launcher(T): # streaming A operand evict_first (one column tile, the wave's panels exceed the L2) for row in ("indexer_hw", "indexer_k"): v = _views("proj", row, "fwd", "bf16", T) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert ( plan.cta1 and _nopro(plan.template) == "dense_proj_det_gemm_kk_n128_hen_c1" ), row @@ -2080,7 +2095,9 @@ def test_round19_rules_plan_like_the_cake_launcher(T): assert bf16.grid == (148, 1, 1) and bf16.sk_iters == 64 assert bf16.ws_f32_elems == 64 * 2 * 256 * 256 vf = _views("proj", "indexer_q", "wgrad", "f32", T) - f32, *_ = cake_dense_gemm_det_train_projection_plan(vf["A"], vf["B"], vf["out"], **kw) + f32, *_ = cake_dense_gemm_det_train_projection_plan( + vf["A"], vf["B"], vf["out"], **kw + ) assert f32.sk_sync and f32.sk_slab == 3 and f32.epi == "tma" assert _nopro(f32.template) == "dense_proj_det_gemm_nn_n256_m256_f32_tma1_sks_sb3" assert (f32.num_full, f32.tail_tiles, f32.sk_units, f32.iters_per_unit) == ( @@ -2105,7 +2122,9 @@ def test_round19_rules_plan_like_the_cake_launcher(T): assert narrow.sk_sync and narrow.sk_slab == 0 assert _nopro(narrow.template) == "dense_proj_det_gemm_nn_n256_m256_s2_sks" with pytest.raises(ValueError, match="sk_slab=3"): - cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk_slab=3, **kw) + cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sk_slab=3, **kw + ) with pytest.raises(ValueError, match="too small"): cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], stages=2, sk_slab=2, **kw @@ -2114,14 +2133,18 @@ def test_round19_rules_plan_like_the_cake_launcher(T): v["A"], v["B"], v["out"], cta_rows=128, **kw ) assert not std.sk_sync and std.sk_slab == 0 and "_sks" not in _nopro(std.template) - off, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=False, **kw) + off, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sk=False, **kw + ) assert ( not off.sk_sync and _nopro(off.template) == "dense_proj_det_gemm_nn_n256_m256" ) # at the test tables' T the synchronised plan is not admitted (a part under SK_MIN_ITERS): the plain tall plan for short_T in (2049, 1001): vs = _views("proj", "indexer_q", "wgrad", "bf16", short_T) - short, *_ = cake_dense_gemm_det_train_projection_plan(vs["A"], vs["B"], vs["out"], **kw) + short, *_ = cake_dense_gemm_det_train_projection_plan( + vs["A"], vs["B"], vs["out"], **kw + ) assert ( not short.sk_sync and short.sk_slab == 0 @@ -2197,7 +2220,9 @@ def test_round21_rules_plan_like_the_cake_launcher(T): v = _views("proj", "kv_a", "dgrad", "bf16", T) # (the row's eight-warp CTA form is forced off here - its own test is # test_eight_warp_form_rules_plan_their_programs; the ten-warp program below is the row's program where it yields) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], form="", **kw + ) assert ( _nopro(plan.template) == "dense_proj_det_gemm_kn_n256_tma1_pd2" and not plan.cta1 @@ -2223,7 +2248,9 @@ def test_round21_rules_plan_like_the_cake_launcher(T): # the sibling rows keep their programs: the sm_107a kv_a bf16 input gradient (quad stores + in-tile drain), the fp32 kv_a input # gradient (one-slot TMA store, no drain) and the indexer_hw bf16 input gradient (two-slot TMA store + the cross-tile drain) kw107 = dict(sm_count=212, l2_bytes=L2_BYTES, arch="sm_107a", _fallback=False) - p107, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw107) + p107, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw107 + ) assert ( _nopro(p107.template) == "dense_proj_det_gemm_kn_n256_q_pd1" and p107.group_m == 8 @@ -2288,7 +2315,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): "cta1": True, } h = _views("proj", "indexer_hw", "dgrad", "bf16", T) - ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw107) + ph, *_ = cake_dense_gemm_det_train_projection_plan( + h["A"], h["B"], h["out"], **kw107 + ) # TEMPLATES_REACHED_NEW (the kernel's rule comment names kn_n256_tma2_pd3_c1: K = 32 <= K_TMA_BF16 keeps the row's two-slot # TMA-store epilogue; the K-less by-kwargs identity render kn_n256_pd3_c1 is the register-epilogue sibling, not this plan) assert _nopro(ph.template) == "dense_proj_det_gemm_kn_n256_tma2_pd3_c1" and ph.cta1 @@ -2310,9 +2339,13 @@ def test_round24_rules_plan_like_the_cake_launcher(T): ) # the control row indexer_k (K 128 is another key) and the sm_100a sibling keep their round-21 programs k_ = _views("proj", "indexer_k", "dgrad", "bf16", T) - pk, *_ = cake_dense_gemm_det_train_projection_plan(k_["A"], k_["B"], k_["out"], **kw107) + pk, *_ = cake_dense_gemm_det_train_projection_plan( + k_["A"], k_["B"], k_["out"], **kw107 + ) assert _nopro(pk.template) == "dense_proj_det_gemm_kn_n256_tma2_pd2" and not pk.cta1 - ph100, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw) + ph100, *_ = cake_dense_gemm_det_train_projection_plan( + h["A"], h["B"], h["out"], **kw + ) assert ( _nopro(ph100.template) == "dense_proj_det_gemm_kn_n256_tma2_pd2" and not ph100.cta1 @@ -2322,7 +2355,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): fkey = ("sm_107a", False, True, True, False, False, 6144, 32, None) assert ROW_RULES[fkey] == {"block_n": 128, "slots": 1, "group_m": 2, "cta1": True} f = _views("proj", "indexer_hw", "dgrad", "f32", T) - pf, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **kw107) + pf, *_ = cake_dense_gemm_det_train_projection_plan( + f["A"], f["B"], f["out"], **kw107 + ) assert ( _nopro(pf.template) == "dense_proj_det_gemm_kn_n128_f32_tma1_c1" and pf.cta1 ) # TEMPLATES_REACHED_NEW (the kernel's rule comment names kn_n128_f32_tma1_c1) @@ -2345,7 +2380,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): wkey = ("sm_107a", True, True, False, False, False, 2048, None, 6144) assert ROW_RULES[wkey] == {"sk_sync": True, "sk_sync_m": 38, "sk_slab": 2} w = _views("proj", "shared_down", "wgrad", "bf16", T) - pw, *_ = cake_dense_gemm_det_train_projection_plan(w["A"], w["B"], w["out"], **kw107) + pw, *_ = cake_dense_gemm_det_train_projection_plan( + w["A"], w["B"], w["out"], **kw107 + ) assert ( _nopro(pw.template) == "dense_proj_det_gemm_nn_n256_sks_sb2" and pw.sk_sync @@ -2370,7 +2407,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): ] == {"group_m": 8, "cgrp": True} # unchanged since round 21 for row in ("q_a", "shared_gate_up"): g = _views("proj", row, "fwd", "bf16", T) - pg, *_ = cake_dense_gemm_det_train_projection_plan(g["A"], g["B"], g["out"], **kw107) + pg, *_ = cake_dense_gemm_det_train_projection_plan( + g["A"], g["B"], g["out"], **kw107 + ) # TEMPLATES_REACHED_NEW / the figen tables (dense_proj_det_gemm_kk_n256_cg = the round-21 sm_100a program; the cgrp hint gate stays # ('none', 'none') at 106 pairs - wave_working_set_cgrp(128, 8, 6144, 8, 106) = 22 x 256 x 6144 x 2 B = 69 206 016 B <= L2_BYTES) assert ( @@ -2395,7 +2434,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): "group_m": 12, } d = _views("proj", "dense_down", "fwd", "bf16", T) - pd_, *_ = cake_dense_gemm_det_train_projection_plan(d["A"], d["B"], d["out"], **kw107) + pd_, *_ = cake_dense_gemm_det_train_projection_plan( + d["A"], d["B"], d["out"], **kw107 + ) # TEMPLATES_REACHED_NEW: the `_cg_or` tall program WITHOUT `_ht` (sk_parts is the row's tail policy; by-kwargs identity: kk_n256_m256_cg_or) assert ( _nopro(pd_.template) == "dense_proj_det_gemm_kk_n256_m256_cg_or" @@ -2424,7 +2465,9 @@ def test_round24_rules_plan_like_the_cake_launcher(T): "group_m": 8, } o = _views("proj", "o_proj", "fwd", "bf16", T) - po, *_ = cake_dense_gemm_det_train_projection_plan(o["A"], o["B"], o["out"], **kw107) + po, *_ = cake_dense_gemm_det_train_projection_plan( + o["A"], o["B"], o["out"], **kw107 + ) # TEMPLATES_REACHED_NEW (the kernel registers dense_projection_det_gemm_kk_m256_ht_or_cg -> symbol ..._m256_ht_cg_or) assert ( _nopro(po.template) == "dense_proj_det_gemm_kk_n256_m256_ht_cg_or" @@ -2443,13 +2486,17 @@ def test_round24_rules_plan_like_the_cake_launcher(T): and not old_o.ovr ) # the sm_100a dense_down / o_proj forwards carry no ovr rule and keep their round-21 programs - pd100, *_ = cake_dense_gemm_det_train_projection_plan(d["A"], d["B"], d["out"], **kw) + pd100, *_ = cake_dense_gemm_det_train_projection_plan( + d["A"], d["B"], d["out"], **kw + ) assert ( _nopro(pd100.template) == "dense_proj_det_gemm_kk_n256_m256_ov_ht" and not pd100.ovr and not pd100.cgrp ) - po100, *_ = cake_dense_gemm_det_train_projection_plan(o["A"], o["B"], o["out"], **kw) + po100, *_ = cake_dense_gemm_det_train_projection_plan( + o["A"], o["B"], o["out"], **kw + ) assert ( _nopro(po100.template) == "dense_proj_det_gemm_kk_n256_m256_ov_ht" and not po100.ovr @@ -2482,7 +2529,9 @@ def test_round25_rules_plan_like_the_cake_launcher(T): "store_ef": True, } sdf = _views("proj", "shared_down", "wgrad", "f32", T) - p_sdf, *_ = cake_dense_gemm_det_train_projection_plan(sdf["A"], sdf["B"], sdf["out"], **kw107) + p_sdf, *_ = cake_dense_gemm_det_train_projection_plan( + sdf["A"], sdf["B"], sdf["out"], **kw107 + ) assert ( _nopro(p_sdf.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2" ) # ROUTE_TEMPLATE_NEW of the round-25 templates step = TEMPLATES_REACHED_NEW (nn_n256_f32_v8_ef_sks_sb2, the round-25 sweep symbol) @@ -2517,7 +2566,9 @@ def test_round25_rules_plan_like_the_cake_launcher(T): } for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "f32", T) - p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw107) + p, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw107 + ) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2", ( row ) # ROUTE_TEMPLATE_NEW of the round-25 templates step = TEMPLATES_REACHED_NEW (nn_n256_f32_v8_ef_sks_sb2) @@ -2544,7 +2595,9 @@ def test_round25_rules_plan_like_the_cake_launcher(T): assert ROW_RULES[qab_key] == {"sk_sync": True, "sk_sync_m": 26, "sk_slab": 2} for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "bf16", T) - p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw107) + p, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw107 + ) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_sks_sb2", ( row ) # ROUTE_TEMPLATE_NEW of the round-25 templates step (nn_n256_sks_sb2 = the round-24 sm_107a shared_down wgrad bf16 program, no new template) @@ -2607,7 +2660,9 @@ def test_round25_rules_plan_like_the_cake_launcher(T): assert ROW_RULES[qab100_key] == {"sk_sync": True, "sk_sync_m": 26, "sk_slab": 2} for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "bf16", T) - p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + p, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_sks_sb2", row assert p.sk_sync and p.sk_slab == 2 and p.sk_exact == 0 and p.cta_rows == 128, ( row @@ -2632,7 +2687,9 @@ def test_round25_rules_plan_like_the_cake_launcher(T): "sk_sync_m": 26, "sk_slab": 2, } - p_sdf100, *_ = cake_dense_gemm_det_train_projection_plan(sdf["A"], sdf["B"], sdf["out"], **kw) + p_sdf100, *_ = cake_dense_gemm_det_train_projection_plan( + sdf["A"], sdf["B"], sdf["out"], **kw + ) assert ( _nopro(p_sdf100.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2" ) # ROUTE_TEMPLATE_NEW of the templates step = TEMPLATES_REACHED_NEW (predicted nn_n256_f32_v8_ef_sks_sb2 - the same new program as the sm_107a fp32 rows; confirm per route) @@ -2667,7 +2724,9 @@ def test_round25_rules_plan_like_the_cake_launcher(T): } for row in ("q_a", "shared_gate_up"): v = _views("proj", row, "wgrad", "f32", T) - p, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + p, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert _nopro(p.template) == "dense_proj_det_gemm_nn_n256_f32_v8_ef_sks_sb2", ( row ) @@ -2716,7 +2775,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): "group_m": 12, } g = _views("proj", "dense_gate_up", "fwd", "bf16", T) - pg, *_ = cake_dense_gemm_det_train_projection_plan(g["A"], g["B"], g["out"], **kw107) + pg, *_ = cake_dense_gemm_det_train_projection_plan( + g["A"], g["B"], g["out"], **kw107 + ) assert ( _nopro(pg.template) == "dense_proj_det_gemm_kk_n256_m256_ht_cg_or" and pg.ovr @@ -2752,7 +2813,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): "store_ef": True, } wf = _views("proj", "dense_down", "wgrad", "f32", T) - pwf, *_ = cake_dense_gemm_det_train_projection_plan(wf["A"], wf["B"], wf["out"], **kw107) + pwf, *_ = cake_dense_gemm_det_train_projection_plan( + wf["A"], wf["B"], wf["out"], **kw107 + ) assert ( _nopro(pwf.template) == "dense_proj_det_gemm_nn_n256_m256_f32_v8_ef_ht" and pwf.htail @@ -2777,7 +2840,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): wbkey = ("sm_107a", True, True, False, False, False, 12288, None, 6144) assert ROW_RULES[wbkey] == {"cta_rows": 256, "htail": True} wb = _views("proj", "dense_down", "wgrad", "bf16", T) - pwb, *_ = cake_dense_gemm_det_train_projection_plan(wb["A"], wb["B"], wb["out"], **kw107) + pwb, *_ = cake_dense_gemm_det_train_projection_plan( + wb["A"], wb["B"], wb["out"], **kw107 + ) assert ( _nopro(pwb.template) == "dense_proj_det_gemm_nn_n256_m256_ht" and pwb.htail @@ -2804,7 +2869,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): "group_m": 12, } dg = _views("proj", "dense_gate_up", "dgrad", "bf16", T) - pdg, *_ = cake_dense_gemm_det_train_projection_plan(dg["A"], dg["B"], dg["out"], **kw107) + pdg, *_ = cake_dense_gemm_det_train_projection_plan( + dg["A"], dg["B"], dg["out"], **kw107 + ) assert ( _nopro(pdg.template) == "dense_proj_det_gemm_kn_n256_m256_cg_or" and pdg.ovr @@ -2826,7 +2893,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): qkey = ("sm_107a", False, True, False, False, False, 2048, 16384, None) assert ROW_RULES[qkey] == {"cgrp": True, "group_m": 8} q = _views("proj", "q_b", "dgrad", "bf16", T) - pq, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kw107) + pq, *_ = cake_dense_gemm_det_train_projection_plan( + q["A"], q["B"], q["out"], **kw107 + ) assert ( _nopro(pq.template) == "dense_proj_det_gemm_kn_n256_cg" and pq.cgrp @@ -2845,7 +2914,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): gwkey = ("sm_107a", True, True, True, False, False, 6144, None, 12288) assert ROW_RULES[gwkey] == {"cta_rows": 256, "htail": True} gw = _views("proj", "dense_gate_up", "wgrad", "f32", T) - pgw, *_ = cake_dense_gemm_det_train_projection_plan(gw["A"], gw["B"], gw["out"], **kw107) + pgw, *_ = cake_dense_gemm_det_train_projection_plan( + gw["A"], gw["B"], gw["out"], **kw107 + ) assert ( _nopro(pgw.template) == "dense_proj_det_gemm_nn_n256_m256_f32_ht_tma1" and pgw.htail @@ -2889,7 +2960,9 @@ def test_round26_rules_plan_like_the_cake_launcher(T): ), ): v = _views("proj", row, op, dt, T) - p100, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + p100, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert _nopro(p100.template) == tmpl and not p100.cgrp and not p100.ovr, ( row, op, @@ -2926,7 +2999,9 @@ def test_round27_rules_plan_like_the_cake_launcher(T): # drain; the previous lock planned the pair form dense_proj_det_gemm_kn_n128_f32_tma1 ---- bkey = ("sm_107a", False, True, True, False, False, 6144, 128, None) assert ROW_RULES[bkey] == {"block_n": 128, "cta1": True, "group_m": 2, "pd": 2} - pb, *_ = cake_dense_gemm_det_train_projection_plan(a["A"], a["B"], a["out"], **kw107) + pb, *_ = cake_dense_gemm_det_train_projection_plan( + a["A"], a["B"], a["out"], **kw107 + ) assert ( _nopro(pb.template) == "dense_proj_det_gemm_kn_n128_f32_tma1_pd2_c1" and pb.cta1 @@ -3339,16 +3414,22 @@ def test_round33_rules_plan_like_the_cake_launcher(T): # architectures; a caller-forced pro 0 is the previous lock's program with the same plan arithmetic ---- q = _views("proj", "indexer_q", "fwd", "bf16", T) for kwx in (kw, kw107): - p, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kwx) + p, *_ = cake_dense_gemm_det_train_projection_plan( + q["A"], q["B"], q["out"], **kwx + ) assert p.pro == 15 and p.template.endswith("_pro15") and not p.cta1, kwx["arch"] - old, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], pro=0, **kwx) + old, *_ = cake_dense_gemm_det_train_projection_plan( + q["A"], q["B"], q["out"], pro=0, **kwx + ) assert ( old.pro == 0 and old.template == p.template[: -len("_pro15")] and old.grid == p.grid and old.stages == p.stages ), kwx["arch"] - p107, *_ = cake_dense_gemm_det_train_projection_plan(q["A"], q["B"], q["out"], **kw107) + p107, *_ = cake_dense_gemm_det_train_projection_plan( + q["A"], q["B"], q["out"], **kw107 + ) assert ( p107.template == "dense_proj_det_gemm_kk_n256_pro15" ) # ROUTE_TEMPLATE_NEW of the templates step (predicted kk_n256_pro15) @@ -3358,7 +3439,9 @@ def test_round33_rules_plan_like_the_cake_launcher(T): assert ( ps.pro == 0 and not re.search(r"_pro\d+$", ps.template) and ps.group_m == 8 ) # NOT `"_pro" not in`: `dense_proj` contains it - ps107, *_ = cake_dense_gemm_det_train_projection_plan(s["A"], s["B"], s["out"], **kw107) + ps107, *_ = cake_dense_gemm_det_train_projection_plan( + s["A"], s["B"], s["out"], **kw107 + ) assert ps107.pro == 7 and ps107.template.endswith("_pro7") # the q_b weight gradient fp32-out on sm_107a: pro 15 on the register epilogue with the 256-bit stores (the round-27 knobs) w = _views("proj", "q_b", "wgrad", "f32", T) @@ -3416,14 +3499,18 @@ def test_round33_rules_plan_like_the_cake_launcher(T): # a single-CTA row: the indexer_hw input gradient fp32-out on sm_107a keeps its round-19 program (no pro; the merged rule's default is # zeroed), a caller-forced pro raises like the Cake launcher, and a caller-forced pair form takes the architecture default h = _views("proj", "indexer_hw", "dgrad", "f32", T) - ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw107) + ph, *_ = cake_dense_gemm_det_train_projection_plan( + h["A"], h["B"], h["out"], **kw107 + ) assert ( ph.cta1 and ph.pro == 0 and ph.template == "dense_proj_det_gemm_kn_n128_f32_tma1_c1" ) with pytest.raises(ValueError, match="needs the cluster pair"): - cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], pro=7, **kw107) + cake_dense_gemm_det_train_projection_plan( + h["A"], h["B"], h["out"], pro=7, **kw107 + ) pair_h, *_ = cake_dense_gemm_det_train_projection_plan( h["A"], h["B"], h["out"], cta1=False, **kw107 ) @@ -3549,14 +3636,18 @@ def test_round33_pro_knob(T): and "_f32_v8" not in nr.template ) h = _views("proj", "indexer_hw", "dgrad", "f32", T) - ph, *_ = cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], **kw107) + ph, *_ = cake_dense_gemm_det_train_projection_plan( + h["A"], h["B"], h["out"], **kw107 + ) assert ( ph.cta1 and ph.pro == 0 and ph.template == "dense_proj_det_gemm_kn_n128_f32_tma1_c1" ) with pytest.raises(ValueError, match="needs the cluster pair"): - cake_dense_gemm_det_train_projection_plan(h["A"], h["B"], h["out"], pro=15, **kw107) + cake_dense_gemm_det_train_projection_plan( + h["A"], h["B"], h["out"], pro=15, **kw107 + ) @pytest.mark.parametrize("T", [16172, 16231]) @@ -3680,7 +3771,9 @@ def test_round21_cgrp_knob(T): ): assert getattr(plan, f) == getattr(base, f), f # the raster width defaults to the row default (16, ``default_group_m``) under cgrp as well, now as column tiles - cg16, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], cgrp=True, **kw) + cg16, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], cgrp=True, **kw + ) assert cg16.group_m == 16 and _nopro(cg16.template) == _nopro(plan.template) # the pair invariant holds whatever the column count, so the count has no parity constraint under cgrp: 3 (or 1) # column tiles per group are admitted where the row raster rejects an odd count (CTA pairs are adjacent row @@ -3694,7 +3787,9 @@ def test_round21_cgrp_knob(T): ) assert one.group_m == 1 and _nopro(one.template) == _nopro(plan.template) with pytest.raises(ValueError, match="group_m must be an even number"): - cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], group_m=3, **kw) + cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], group_m=3, **kw + ) with pytest.raises(ValueError, match="group_m must be an even number"): cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], cgrp=False, group_m=1, **kw @@ -5193,7 +5288,9 @@ def test_ht_program_with_forced_stream_k_matches_fp64_reference(): row, op, T, _ = found[0] v = _make_inputs("proj", row, op, "bf16", T, seed=SEED + 7) _require_program(_plan_template(v, sk=True)) - prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"], sk=True) + prepared = cake_dense_gemm_det_train_projection_prepare( + v["A"], v["B"], v["out"], sk=True + ) assert "_ht" in prepared.plan.template and prepared.plan.sk_units == 0 out = prepared.launch() torch.cuda.synchronize() @@ -5384,7 +5481,9 @@ def test_router_layout_classification_and_swap(): X = torch.empty(T, K) W = torch.empty(N, K) G = torch.empty(T, N) - plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(X, W.t(), torch.empty(T, N)) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan( + X, W.t(), torch.empty(T, N) + ) assert (plan.layout, plan.template, plan.splits) == ( "kk", "router_fp32_det_gemm_kk", @@ -5398,7 +5497,9 @@ def test_router_layout_classification_and_swap(): 1, ) assert (plan.M, plan.N, plan.K) == (T, K, N) and plan.b_mn - plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(G.t(), X, torch.empty(N, K)) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan( + G.t(), X, torch.empty(N, K) + ) assert (plan.layout, plan.template, plan.splits) == ( "nn_t", "router_fp32_det_gemm_nn_t", @@ -5414,14 +5515,20 @@ def test_router_layout_classification_and_swap(): 1, 1, ) - plan, *_ = cake_dense_gemm_det_train_router_plan(X, W.t(), torch.empty(T, N), splits=2) + plan, *_ = cake_dense_gemm_det_train_router_plan( + X, W.t(), torch.empty(T, N), splits=2 + ) assert plan.splits == 2 and plan.k_iters_split == -(-plan.k_blocks // 2) with pytest.raises(NotImplementedError, match="MN-major A"): router_layout_class(True, False) with pytest.raises(ValueError, match="fp32"): - cake_dense_gemm_det_train_router_plan(X.to(torch.bfloat16), W.t(), torch.empty(T, N)) + cake_dense_gemm_det_train_router_plan( + X.to(torch.bfloat16), W.t(), torch.empty(T, N) + ) with pytest.raises(ValueError, match="K must agree"): - cake_dense_gemm_det_train_router_plan(X, torch.empty(K + 32, N), torch.empty(T, N)) + cake_dense_gemm_det_train_router_plan( + X, torch.empty(K + 32, N), torch.empty(T, N) + ) with pytest.raises(ValueError, match="out must be"): cake_dense_gemm_det_train_router_plan(X, W.t(), torch.empty(N, T)) with pytest.raises(ValueError, match="multiple of 8"): @@ -5435,7 +5542,9 @@ def test_router_layout_classification_and_swap(): # a ragged batch with T = 0: M = 0 for the forward / input gradient, K = 0 for the # weight gradient -- rejected before a zero-sized grid or TMA extent is planned with pytest.raises(ValueError, match="must be >= 1"): - cake_dense_gemm_det_train_router_plan(torch.empty(0, K), W.t(), torch.empty(0, N)) + cake_dense_gemm_det_train_router_plan( + torch.empty(0, K), W.t(), torch.empty(0, N) + ) with pytest.raises(ValueError, match="must be >= 1"): cake_dense_gemm_det_train_router_plan( torch.empty(N, 0), torch.empty(0, K), torch.empty(N, K) @@ -5472,14 +5581,22 @@ def test_router_planner_refuses_a_misaligned_out_of_a_single_split_row_major_sto with pytest.raises(ValueError, match="16-byte aligned address"): cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) else: - plan, *_ = cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) + plan, *_ = cake_dense_gemm_det_train_router_plan( + A, B, out, splits=splits + ) assert (plan.layout, plan.splits) == (layout, splits) # the production split counts: the input gradient (kn, one split) refuses, the forward (kk, 4) and the # weight gradient (nn_t, 6) plan with pytest.raises(ValueError, match="16-byte aligned address"): cake_dense_gemm_det_train_router_plan(G, W, _offset_out(T, K, 1)) - assert cake_dense_gemm_det_train_router_plan(X, W.t(), _offset_out(T, N, 1))[0].splits == 4 - assert cake_dense_gemm_det_train_router_plan(G.t(), X, _offset_out(N, K, 1))[0].splits == 6 + assert ( + cake_dense_gemm_det_train_router_plan(X, W.t(), _offset_out(T, N, 1))[0].splits + == 4 + ) + assert ( + cake_dense_gemm_det_train_router_plan(G.t(), X, _offset_out(N, K, 1))[0].splits + == 6 + ) def test_split_fp32_to_bf16x3_is_exact_and_allocation_free_in_place(): @@ -6129,7 +6246,9 @@ def chain_split(src, parts, resid): n = v["W"].numel() split_args = (n, -(-n // cake_backend.SPLIT_BLOCK_ELEMS), 1, 1) # the prepared path - prepared = cake_dense_gemm_det_train_router_prepare(v["A"], v["B"], v["out"], splits=4) + prepared = cake_dense_gemm_det_train_router_prepare( + v["A"], v["B"], v["out"], splits=4 + ) assert prepared.fused_split is fused and ("resid" in prepared.tensors) is not fused prepared.launch() if fused: @@ -6144,7 +6263,9 @@ def chain_split(src, parts, resid): # the convenience path calls.clear() chain.clear() - out = cake_backend.cake_dense_gemm_det_train_router(v["A"], v["B"], v["out"], splits=4) + out = cake_backend.cake_dense_gemm_det_train_router( + v["A"], v["B"], v["out"], splits=4 + ) assert out is v["out"] cached = next(iter(cake_backend._K2_CACHE.values())) arena = cached.arena @@ -6516,7 +6637,9 @@ def test_registry_fallback_reaches_the_shapes_tranche_programs(monkeypatch): # a caller-forced tail policy keeps the row on a registered program: a pair form admits it, the single-CTA # form (``cta1`` clears the tail knobs) stays reachable when no pair form is registered v = _class_views("kk", False, False, 16384, 128, 4096) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sk=True, **kw + ) assert plan.template in registered and plan.knob_fallback, plan.template if plan.template.endswith("_c1"): assert plan.cta1 and plan.sk_units == 0 and not plan.sk_sync @@ -7006,7 +7129,9 @@ def test_measured_shape_rules_plan_their_programs(arch): T, template, ) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) if template in registered: assert plan.template == template, (key, T, plan.template) assert not ( @@ -7347,7 +7472,9 @@ def test_round2_rules_plan_their_knobs(): for M, N in ((12288, 4096), (4096, 12288)): for T, k_blocks in ((16384, 256), (16231, 254)): v = _class_views("nn", False, False, M, N, T) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **a + ) assert ( plan.template == "dense_proj_det_gemm_nn_n256_m256_skx_skf_pro7" ), ( @@ -7401,7 +7528,9 @@ def test_round2_rules_plan_their_knobs(): ("kn", "dense_proj_det_gemm_kn_n256_ht_cg_pro15"), ): v = _class_views(layout, False, False, 16384, 4096, 4096) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **a + ) assert ( plan.template == template and plan.htail @@ -7420,7 +7549,9 @@ def test_round2_rules_plan_their_knobs(): ("kk", 17408, 5120, "dense_proj_det_gemm_kk_n256_cg_pro15"), ): v = _class_views(layout, False, False, 16384, N, K) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **a + ) assert plan.template == template and plan.cgrp and plan.group_m == 12, ( arch, layout, @@ -7476,7 +7607,9 @@ def test_round2_rules_plan_their_knobs(): v["A"], v["B"], v["out"], pro=3, lw=True, **a ) v = _class_views("kn", False, False, 16384, 5120, 128) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **a + ) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma2_pd2_lw_pro15" and plan.pd == 2 @@ -7494,13 +7627,17 @@ def test_round2_rules_plan_their_knobs(): # (the two rows' eight-warp CTA form is forced off here - its own test is # test_eight_warp_form_rules_plan_their_programs; the ten-warp program below is the row's program where it yields) v = _class_views("kn", False, False, 16384, 4096, 512) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], form="", **a + ) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_lw_mg_pro15" and plan.promo == "l2_256b" ) # + the whole-tile tile resolve (test_mg_rules_plan_the_whole_tile_programs) v = _class_views("kn", False, False, 16384, 5120, 512) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], form="", **a + ) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_mg_sw25_pro15" and plan.promo == "l2_256b" @@ -7529,7 +7666,9 @@ def test_round2_rules_plan_their_knobs(): # (the row's eight-warp CTA form is forced off here - its own test is # test_eight_warp_form_rules_plan_their_programs; the ten-warp program below is the row's program where it yields) v = _class_views("kn", False, False, 16384, 5120, 512) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], form="", **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], form="", **a + ) assert ( plan.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_pro15" and plan.promo == "none" @@ -7537,7 +7676,9 @@ def test_round2_rules_plan_their_knobs(): ) # the column-grouped raster takes any group count; the pair raster still needs an even one v = _class_views("kk", False, False, 16384, 5120, 5120) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], group_m=3, **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], group_m=3, **a + ) assert plan.group_m == 3 and plan.cgrp with pytest.raises(ValueError): cake_dense_gemm_det_train_projection_plan( @@ -7667,11 +7808,15 @@ def test_mg_rules_plan_the_whole_tile_programs(T): assert sk.sk_units == 0 and sk.mg assert sk.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_lw_mg_pro15" v = _class_views("kn", False, False, T, 6144, 576) - sk, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, **a) + sk, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sk=True, **a + ) assert (sk.sk_units, sk.tail_tiles, sk.iters_per_unit) == (18, 16, 8) assert not sk.mg and sk.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_pro15" with pytest.raises(ValueError, match="whole-tile plan"): - cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, mg=True, **a) + cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sk=True, mg=True, **a + ) # the rows without the resolve keep their program (5120 -> 1024 paces its stores), the fp32 twins and sm_100a # never carry the resolve (both rows carry the eight-warp CTA form: forced off here) for K_layer, N_layer, template in ( @@ -7692,7 +7837,9 @@ def test_mg_rules_plan_the_whole_tile_programs(T): ) for K_layer, N_layer, _t, _g, _p, _lw in rows: f = _class_views("kn", True, False, T, K_layer, N_layer) - plan, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + f["A"], f["B"], f["out"], **a + ) assert not plan.mg and "_mg" not in plan.template, (K_layer, N_layer, "f32") v = _class_views("kn", False, False, T, K_layer, N_layer) plan, *_ = cake_dense_gemm_det_train_projection_plan( @@ -7766,7 +7913,9 @@ def test_sw_rules_plan_the_paced_store_programs(T): assert (plan.epi, plan.slots, plan.pd, plan.pro) == ("tma", 1, 2, 15), key assert plan.sk_units == 0 and not plan.cta1 # the twin without the pacing: the same plan, only the program differs - plain, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sw=0, **a) + plain, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sw=0, **a + ) assert plain.sw == 0 and plain.template == template.replace("_sw25", "") assert {**_plan_signature(plan), "template": None, "sw": None} == { **_plan_signature(plain), @@ -7802,7 +7951,9 @@ def test_sw_rules_plan_the_paced_store_programs(T): assert forced.sw == 0 and forced.template.endswith("_c1"), forced.template # the fp32 twin and sm_100a never carry the pacing f = _class_views("kn", True, False, T, K_layer, N_layer) - plan, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **a) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + f["A"], f["B"], f["out"], **a + ) assert plan.sw == 0 and "_sw" not in plan.template, (K_layer, N_layer, "f32") plan, *_ = cake_dense_gemm_det_train_projection_plan( v["A"], v["B"], v["out"], **dict(kw, arch="sm_100a") @@ -7916,7 +8067,9 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): rule = ROW_RULES[key] assert rule["form"] == form and rule["sk"] is False, key v = _class_views("kn", False, False, T, N_layer, K_layer) - plan, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **mirror) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **mirror + ) assert plan.template == template, (key, T, plan.template) assert plan.form == form and plan.tma_out and plan.epi == "tma" assert (plan.slots, plan.stages, plan.block_n, plan.cta_rows) == ( @@ -7933,7 +8086,9 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): assert plan.pd == 0 and plan.pro == 0 and plan.pf == 0 and plan.sh == 0, key assert plan.hints == ("none", "none") and plan.store_hint == "none", key # the registered program plans directly (no fallback) - served, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + served, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert served.template == template and template in registered, (key, T) assert not ( served.swap_fallback @@ -7991,7 +8146,9 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): # register epilogue or a caller-forced pipeline knob plan the row's ten-warp program; a caller-forced form is # refused on every such plan v = _class_views("kn", False, False, T, 6144, 576) - sk, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], sk=True, **mirror) + sk, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], sk=True, **mirror + ) assert (sk.sk_units, sk.tail_tiles) == (18, 16) and sk.form == "" assert sk.template == "dense_proj_det_gemm_kn_n256_tma1_pd2_pro15" with pytest.raises(ValueError, match="whole-tile"): @@ -8021,7 +8178,9 @@ def test_eight_warp_form_rules_plan_their_programs(arch, T, monkeypatch): ) # the fp32 twin carries no form on either architecture (the fp32 rows are other rules) f = _class_views("kn", True, False, T, N_layer, K_layer) - plan, *_ = cake_dense_gemm_det_train_projection_plan(f["A"], f["B"], f["out"], **mirror) + plan, *_ = cake_dense_gemm_det_train_projection_plan( + f["A"], f["B"], f["out"], **mirror + ) assert plan.form == "" and "_w8" not in plan.template, (N_layer, K_layer) # the other architectures' plans of the row: the five sm_100a rows shared with sm_103a carry the sm_103a form # (``w8rs`` / ``w8rs`` / ``w8r`` / ``w8r`` / ``w8rs``) and the sm_100a 4096 -> 64 and 5120 -> 64 rows none there @@ -8365,7 +8524,9 @@ def test_half_tail_fit_gates_the_plan_not_the_program(monkeypatch): (5120, (1280, 0, 0, 1), False), ): v = _class_views("kn", False, False, 16384, K_layer, 32) - ruled, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + ruled, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert ruled.template == "dense_proj_det_gemm_kn_n256_ht_tma2_lw_pro15" assert ruled.htail and ruled.htail_plan == half assert ( @@ -8377,7 +8538,9 @@ def test_half_tail_fit_gates_the_plan_not_the_program(monkeypatch): assert ruled.sk_iters == (72 if half else 0) and ruled.ws_f32_elems == 0 with monkeypatch.context() as m: m.setattr(cake_backend, "ROW_RULES", {}) - plain, *_ = cake_dense_gemm_det_train_projection_plan(v["A"], v["B"], v["out"], **kw) + plain, *_ = cake_dense_gemm_det_train_projection_plan( + v["A"], v["B"], v["out"], **kw + ) assert ( not plain.htail and not plain.htail_plan and "_ht" not in plain.template ) @@ -8552,7 +8715,9 @@ def test_batched_rows_satisfy_the_kernel_output_layout_predicate(): def test_router_convenience_views_match_the_planner(): for op, T in (("fwd", 257), ("dgrad", 257), ("wgrad", 1001)): v = _router_inputs(op, T, seed=SEED + 5, device="cpu") - plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan(v["A"], v["B"], v["out"]) + plan, A_eff, B_eff = cake_dense_gemm_det_train_router_plan( + v["A"], v["B"], v["out"] + ) a_view, split_src = cake_backend._k2_views(v["A"], v["B"], plan) want_split = B_eff if plan.b_mn else B_eff.t() want_a = A_eff.t().unsqueeze(0) if plan.a_mn else A_eff.unsqueeze(0) @@ -8672,10 +8837,14 @@ def test_projection_rows_match_fp64_reference(row, op, out_dtype, T): v = _make_inputs("proj", row, op, out_dtype, T, seed=SEED) _require_program(_plan_template(v)) if op == "wgrad": - prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(v["G"], v["X"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare( + v["G"], v["X"], v["out"] + ) assert prepared.plan.transposed_out == v["transposed"] else: - prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_prepare( + v["A"], v["B"], v["out"] + ) out = prepared.launch() torch.cuda.synchronize() assert out is v["out"] @@ -8806,7 +8975,9 @@ def test_strided_views_and_storage_offsets_match_fp64_reference(): # small-N weight gradients to the 64-row family, so only that family's instance is exported there) v_t = dict(A=A.t(), B=Gs, out=out_t, transposed=True) knobs_t = _registered_knobs(v_t, {}, dict(cta_rows=64)) - prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(Gs, A, out_t, **knobs_t) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare( + Gs, A, out_t, **knobs_t + ) assert prepared.plan.transposed_out and prepared.template == _plan_template( v_t, **knobs_t ) @@ -8902,7 +9073,9 @@ def test_convenience_path_reuses_its_planned_launch_without_allocation(): assert after["allocation.all.allocated"] == before["allocation.all.allocated"] assert len(cake_backend._K1_CACHE) == 1 _check(out2, torch.matmul(X2.double(), W.double().t())) - assert torch.equal(cake_backend.cake_dense_gemm_det_train_projection_forward(X, W, out), first) + assert torch.equal( + cake_backend.cake_dense_gemm_det_train_projection_forward(X, W, out), first + ) assert len(cake_backend._ARENAS) == 1 # another stream never shares an entry (nor its workspace arena) with torch.cuda.stream(torch.cuda.Stream(device=device)): @@ -8950,7 +9123,9 @@ def test_convenience_path_replays_a_stream_k_row_bit_exact(): ) knobs = _registered_knobs(v, dict(cta_rows=64), dict(cta_rows=128)) cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() - out1 = cake_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, **knobs) # ``out`` allocated here + out1 = cake_backend.cake_dense_gemm_det_train_projection_wgrad( + G, X, **knobs + ) # ``out`` allocated here torch.cuda.synchronize() _check(out1.t(), torch.matmul(X.t().double(), G.double())) assert len(cake_backend._K1_CACHE) == 1 @@ -9045,7 +9220,9 @@ def held_after(T_values): torch.cuda.synchronize() assert torch.equal(outs[0], outs[1]) and torch.equal(outs[0], outs[2]) dw = cake_backend.cake_dense_gemm_det_train_router_wgrad(G, X) - prepared_w = cake_dense_gemm_det_train_router_prepare(G.t(), X, torch.empty_like(dw), splits=6) + prepared_w = cake_dense_gemm_det_train_router_prepare( + G.t(), X, torch.empty_like(dw), splits=6 + ) assert torch.equal(dw, prepared_w.launch()) cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() assert cake_backend._ARENAS == {} @@ -9078,7 +9255,9 @@ def test_pd2_stream_k_on_an_odd_k_completes(): dict(pd=2, slots=2, epi="tma", pro=15), dict(pd=2, slots=1, epi="tma", pro=15), ) - prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"], **knobs) + prepared = cake_dense_gemm_det_train_projection_prepare( + v["A"], v["B"], v["out"], **knobs + ) plan = prepared.plan assert plan.pd == 2 and plan.k_blocks == 17 and plan.tail_tiles * 2 <= plan.sm_pairs assert ( @@ -9139,9 +9318,15 @@ def test_router_rows_match_fp64_reference_and_rerun_bit_exact(op, T): assert int((top != ref_top).any(dim=1).sum()) == 0 # the eager entry point of the row gives the same bits eager_fn = dict( - fwd=lambda: cake_backend.cake_dense_gemm_det_train_router_forward(v["X"], v["W"]), - dgrad=lambda: cake_backend.cake_dense_gemm_det_train_router_dgrad(v["G"], v["W"]), - wgrad=lambda: cake_backend.cake_dense_gemm_det_train_router_wgrad(v["G"], v["X"]), + fwd=lambda: cake_backend.cake_dense_gemm_det_train_router_forward( + v["X"], v["W"] + ), + dgrad=lambda: cake_backend.cake_dense_gemm_det_train_router_dgrad( + v["G"], v["W"] + ), + wgrad=lambda: cake_backend.cake_dense_gemm_det_train_router_wgrad( + v["G"], v["X"] + ), )[op] cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() eager = eager_fn() @@ -9180,9 +9365,13 @@ def test_router_misaligned_out_is_refused_or_matches_fp64_reference(op, T): cake_backend.cake_dense_gemm_det_train_convenience_cache_clear() if splits == 1 and not plan.transposed_out: with pytest.raises(ValueError, match="16-byte aligned address"): - cake_dense_gemm_det_train_router_prepare(v["A"], v["B"], shifted, splits=splits) + cake_dense_gemm_det_train_router_prepare( + v["A"], v["B"], shifted, splits=splits + ) with pytest.raises(ValueError, match="16-byte aligned address"): - cake_backend.cake_dense_gemm_det_train_router(v["A"], v["B"], eager, splits=splits) + cake_backend.cake_dense_gemm_det_train_router( + v["A"], v["B"], eager, splits=splits + ) torch.cuda.synchronize() # nothing was launched: the context stays healthy cake_dense_gemm_det_train_router_prepare( v["A"], v["B"], aligned, splits=splits @@ -9195,7 +9384,9 @@ def test_router_misaligned_out_is_refused_or_matches_fp64_reference(op, T): cake_dense_gemm_det_train_router_prepare( v["A"], v["B"], shifted, splits=splits ).launch() - cake_backend.cake_dense_gemm_det_train_router(v["A"], v["B"], eager, splits=splits) + cake_backend.cake_dense_gemm_det_train_router( + v["A"], v["B"], eager, splits=splits + ) outs = [shifted, eager] torch.cuda.synchronize() for out in outs: @@ -9335,9 +9526,13 @@ def test_repeated_launches_are_bitwise_identical(row, op, out_dtype, T): v = _make_inputs("proj", row, op, out_dtype, T, seed=SEED + 7) _require_program(_plan_template(v)) if op == "wgrad": - prepared = cake_dense_gemm_det_train_projection_wgrad_prepare(v["G"], v["X"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_wgrad_prepare( + v["G"], v["X"], v["out"] + ) else: - prepared = cake_dense_gemm_det_train_projection_prepare(v["A"], v["B"], v["out"]) + prepared = cake_dense_gemm_det_train_projection_prepare( + v["A"], v["B"], v["out"] + ) first = prepared.launch().clone() torch.cuda.synchronize() assert bool(torch.isfinite(first).all()) @@ -9493,7 +9688,9 @@ def test_swapped_forward_is_bitwise_repeatable(): first = out.clone() for _ in range(REPEATS): out.fill_(float("nan")) - cake_backend.cake_dense_gemm_det_train_projection(A, B, out, enable_pdl=True) + cake_backend.cake_dense_gemm_det_train_projection( + A, B, out, enable_pdl=True + ) torch.cuda.synchronize() assert torch.equal(out, first) @@ -9553,3 +9750,50 @@ def chain(): torch.cuda.synchronize() _check(out, torch.einsum("gtd,gdr->gtr", A.double(), B.double())) _check(out2, torch.einsum("gtd,gdr->gtr", A2.double(), B2.double())) + + +# The module selects the mode: torch.use_deterministic_algorithms (strict or warn_only) changes nothing. +@pytest.mark.parametrize("warn_only", [False, True]) +def test_the_entry_points_run_unchanged_under_torch_deterministic_mode(warn_only): + if arch_for(torch.device("cuda")) is None: + pytest.skip("needs an SM100 / SM103 / SM107 GPU") + device = torch.device("cuda", torch.cuda.current_device()) + m = cake_backend + g = torch.Generator(device=device).manual_seed(0) + T, K, N, R = 257, 6144, 576, 256 # the kv_a projection and the 256-expert router + X = torch.randn(T, K, device=device, generator=g).bfloat16() + W = (torch.randn(N, K, device=device, generator=g) * K**-0.5).bfloat16() + G = torch.randn(T, N, device=device, generator=g).bfloat16() + Xr = torch.randn(T, K, device=device, generator=g) + Wr = torch.randn(R, K, device=device, generator=g) * K**-0.5 + Gr = torch.randn(T, R, device=device, generator=g) + + def run(): + return [ + m.cake_dense_gemm_det_train_projection_forward(X, W), + m.cake_dense_gemm_det_train_projection_dgrad(G, W), + m.cake_dense_gemm_det_train_projection_wgrad(G, X), + m.cake_dense_gemm_det_train_router_forward(Xr, Wr), + m.cake_dense_gemm_det_train_router_dgrad(Gr, Wr), + m.cake_dense_gemm_det_train_router_wgrad(Gr, Xr), + ] + + try: + reference = run() + except NotImplementedError as exc: # a build without one of these programs + pytest.skip(str(exc)) + previous = ( + torch.are_deterministic_algorithms_enabled(), + torch.is_deterministic_algorithms_warn_only_enabled(), + ) + try: + torch.use_deterministic_algorithms(True, warn_only=warn_only) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + outputs = run() + torch.cuda.synchronize(device) + finally: + torch.use_deterministic_algorithms(previous[0], warn_only=previous[1]) + assert [w for w in caught if "determinis" in str(w.message)] == [] + for got, want in zip(outputs, reference, strict=True): + assert torch.equal(got, want) diff --git a/tests/experimental/test_cake_dense_gemm_det_train_names.py b/tests/experimental/test_cake_dense_gemm_det_train_names.py index 3918f0bcc3a..e2d1aec4ca4 100644 --- a/tests/experimental/test_cake_dense_gemm_det_train_names.py +++ b/tests/experimental/test_cake_dense_gemm_det_train_names.py @@ -173,9 +173,7 @@ def test_no_entry_point_takes_a_mode_keyword(name): @pytest.mark.parametrize("name", sorted(PUBLIC)) def test_the_module_returns_the_backend_object(name): - with warnings.catch_warnings(): - warnings.simplefilter("error") # no ExperimentalWarning: a stable module - obj = getattr(importlib.import_module(MODULE), name) + obj = getattr(importlib.import_module(MODULE), name) submodule = "cake_jit" if name == "ARCHES" else "cake_backend" assert obj is getattr(importlib.import_module(f"{BACKEND}.{submodule}"), name) if inspect.isfunction(obj) or inspect.isclass(obj): @@ -184,6 +182,22 @@ def test_the_module_returns_the_backend_object(name): assert obj.__module__ == f"{BACKEND}.cake_backend" +def test_first_access_issues_no_experimental_warning(): + """A stable module: the first access of every name, in a fresh process, issues no ExperimentalWarning.""" + code = f""" +import warnings +with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + import {MODULE} as m + for name in m.__all__: + getattr(m, name) +print(sorted({{type(w.message).__name__ for w in caught if "Experimental" in type(w.message).__name__}})) +""" + result = _run(code) + assert result.returncode == 0, result.stderr + assert result.stdout.strip().splitlines()[-1] == "[]" + + def test_an_unknown_name_raises_attribute_error(): module = importlib.import_module(MODULE) with pytest.raises(AttributeError): @@ -198,9 +212,21 @@ def test_the_former_names_are_removed(submodule, names): @pytest.mark.parametrize("name", FORMER_MODULES) def test_the_former_modules_fail_at_import(name): - result = _run(f"import importlib; importlib.import_module({name!r})") - assert result.returncode != 0 - assert "ModuleNotFoundError" in result.stderr + code = f""" +import importlib +try: + importlib.import_module({name!r}) +except ModuleNotFoundError as exc: + print("missing", exc.name) +else: + print("imported") +""" + result = _run(code) + assert result.returncode == 0, result.stderr + assert result.stdout.strip().splitlines()[-1] in ( + f"missing {name}", + f"missing {FORMER_PACKAGE}", + ) def test_the_former_package_is_gone(): @@ -321,14 +347,6 @@ def test_no_former_spelling_in_the_package(): ] -def _supported_device(): - if not torch.cuda.is_available(): - return None - module = importlib.import_module(MODULE) - device = torch.device("cuda", torch.cuda.current_device()) - return device if module.is_cake_dense_gemm_det_train_supported(device) else None - - def test_the_support_query(): module = importlib.import_module(MODULE) assert module.is_cake_dense_gemm_det_train_supported(torch.device("cpu")) is False @@ -339,47 +357,85 @@ def test_the_support_query(): ) == ((major, minor) in module.SUPPORTED_COMPUTE_CAPABILITIES) +# entry point -> its positional arguments on CPU tensors (T, K, N) = (257, 6144, 576), the kv_a projection, and the +# 256-expert router +def _cpu_calls(): + T, K, N, R = 257, 6144, 576, 256 + bf, f32 = torch.bfloat16, torch.float32 + X, W, G = ( + torch.zeros(T, K, dtype=bf), + torch.zeros(N, K, dtype=bf), + torch.zeros(T, N, dtype=bf), + ) + Xr, Wr, Gr = torch.zeros(T, K), torch.zeros(R, K), torch.zeros(T, R) + plan = dict( + arch="sm_100a", sm_count=148, l2_bytes=126 * 2**20 + ) # the B200 device facts + return { + "cake_dense_gemm_det_train_projection": ( + (X, W.t(), torch.zeros(T, N, dtype=bf)), + {}, + ), + "cake_dense_gemm_det_train_projection_forward": ((X, W), {}), + "cake_dense_gemm_det_train_projection_dgrad": ((G, W), {}), + "cake_dense_gemm_det_train_projection_wgrad": ((G, X), {}), + "cake_dense_gemm_det_train_projection_plan": ( + (X, W.t(), torch.zeros(T, N, dtype=bf)), + plan, + ), + "cake_dense_gemm_det_train_projection_prepare": ( + (X, W.t(), torch.zeros(T, N, dtype=bf)), + {}, + ), + "cake_dense_gemm_det_train_projection_wgrad_prepare": ( + (G, X, torch.zeros(N, K, dtype=f32)), + {}, + ), + "cake_dense_gemm_det_train_router": ((Xr, Wr.t(), torch.zeros(T, R)), {}), + "cake_dense_gemm_det_train_router_forward": ((Xr, Wr), {}), + "cake_dense_gemm_det_train_router_dgrad": ((Gr, Wr), {}), + "cake_dense_gemm_det_train_router_wgrad": ((Gr, Xr), {}), + "cake_dense_gemm_det_train_router_plan": ((Xr, Wr.t(), torch.zeros(T, R)), {}), + "cake_dense_gemm_det_train_router_prepare": ( + (Xr, Wr.t(), torch.zeros(T, R)), + {}, + ), + "cake_dense_gemm_det_train_convenience_cache_clear": ((), {}), + } + + +def test_every_entry_point_is_covered_by_the_mode_test(): + callables = {n for n in PUBLIC if not CONSTANT.match(n) and not n[0].isupper()} + assert set(_cpu_calls()) == callables - {"is_cake_dense_gemm_det_train_supported"} + + @pytest.mark.parametrize("warn_only", [False, True]) -def test_the_entry_points_run_unchanged_under_torch_deterministic_mode(warn_only): - device = _supported_device() - if device is None: - pytest.skip("needs an SM100 / SM103 / SM107 GPU") +@pytest.mark.parametrize("name", sorted(_cpu_calls())) +def test_no_entry_point_alerts_under_torch_deterministic_mode(name, warn_only): + """Under ``torch.use_deterministic_algorithms(True)`` (strict or ``warn_only``) a deterministic entry point raises + no deterministic-mode RuntimeError and issues no deterministic-mode warning. On CPU tensors the launches stop at + the device check (a ValueError); the planner runs (``..._projection_plan`` returns its plan).""" m = importlib.import_module(MODULE) - g = torch.Generator(device=device).manual_seed(0) - T, K, N, R = 257, 6144, 576, 256 # the kv_a projection and the 256-expert router - X = torch.randn(T, K, device=device, generator=g).bfloat16() - W = (torch.randn(N, K, device=device, generator=g) * K**-0.5).bfloat16() - G = torch.randn(T, N, device=device, generator=g).bfloat16() - Xr = torch.randn(T, K, device=device, generator=g) - Wr = torch.randn(R, K, device=device, generator=g) * K**-0.5 - Gr = torch.randn(T, R, device=device, generator=g) - - def run(): - return [ - m.cake_dense_gemm_det_train_projection_forward(X, W), - m.cake_dense_gemm_det_train_projection_dgrad(G, W), - m.cake_dense_gemm_det_train_projection_wgrad(G, X), - m.cake_dense_gemm_det_train_router_forward(Xr, Wr), - m.cake_dense_gemm_det_train_router_dgrad(Gr, Wr), - m.cake_dense_gemm_det_train_router_wgrad(Gr, Xr), - ] - - try: - reference = run() - except NotImplementedError as exc: # a build without one of these programs - pytest.skip(str(exc)) + args, kwargs = _cpu_calls()[name] previous = ( torch.are_deterministic_algorithms_enabled(), torch.is_deterministic_algorithms_warn_only_enabled(), ) + error = None try: torch.use_deterministic_algorithms(True, warn_only=warn_only) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") - outputs = run() - torch.cuda.synchronize(device) + try: + result = getattr(m, name)(*args, **kwargs) + except Exception as exc: # noqa: BLE001 + error = exc finally: torch.use_deterministic_algorithms(previous[0], warn_only=previous[1]) assert [w for w in caught if "determinis" in str(w.message)] == [] - for got, want in zip(outputs, reference, strict=True): - assert torch.equal(got, want) + assert not (isinstance(error, RuntimeError) and "determinis" in str(error)), error + if name in ( + "cake_dense_gemm_det_train_projection_plan", + "cake_dense_gemm_det_train_router_plan", + ): + assert error is None and result is not None # the planners need no device From 2230ce70240481d07dd3bfd30b8192231b44f503 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 00:08:38 -0800 Subject: [PATCH 07/13] refactor(cake_dense_gemm_nondet_train): move the non-deterministic dense GEMM backend The non-deterministic dense GEMM backend moves from flashinfer/experimental/dense_projection_gemm/ to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/, next to its deterministic sibling cake_dense_gemm_det_train (flashinfer-ai/flashinfer#8928). - The generated sources, cake_build.py and .gitignore move unchanged; the program and library names, kernel namespaces and templates keep their names. build.json changes only its "package" field. - cake_jit.DETERMINISTIC_PACKAGE names the sibling cake_dense_gemm_det_train (still resolved as a sibling, never through an absolute import). - The calls into the deterministic package use its new entry names (cake_dense_gemm_det_train_*), including its convenience cache clear. - The package-data row, the manifest checker row (with requires= on the deterministic sibling), the test helper key and the pre-commit csrc exclusion name the new package. --- .pre-commit-config.yaml | 4 +- .../cake_dense_gemm_nondet_train}/.gitignore | 0 .../cake_dense_gemm_nondet_train}/README.md | 0 .../cake_dense_gemm_nondet_train}/__init__.py | 2 +- .../cake_dense_gemm_nondet_train}/build.json | 2 +- .../cake_backend.py | 102 +++++++++--------- .../cake_build.py | 0 .../cake_dense_gemm_nondet_train}/cake_jit.py | 6 +- ...ction_gemm_16496106e8ac1c7b1048_binding.cu | 0 ...ection_gemm_16496106e8ac1c7b1048_kernel.cu | 0 ...ction_gemm_1d46f84995b7baec3e04_binding.cu | 0 ...ection_gemm_1d46f84995b7baec3e04_kernel.cu | 0 ...ction_gemm_2395f1c01b99015f5dfe_binding.cu | 0 ...ection_gemm_2395f1c01b99015f5dfe_kernel.cu | 0 ...ction_gemm_23cd6c7b07c9130bef18_binding.cu | 0 ...ection_gemm_23cd6c7b07c9130bef18_kernel.cu | 0 ...ction_gemm_243446e7407f2bed9b45_binding.cu | 0 ...ection_gemm_243446e7407f2bed9b45_kernel.cu | 0 ...ction_gemm_2f07a521e0f7ebd21021_binding.cu | 0 ...ection_gemm_2f07a521e0f7ebd21021_kernel.cu | 0 ...ction_gemm_3f431ee1064967195e90_binding.cu | 0 ...ection_gemm_3f431ee1064967195e90_kernel.cu | 0 ...ction_gemm_4b13414a327c7040abd8_binding.cu | 0 ...ection_gemm_4b13414a327c7040abd8_kernel.cu | 0 ...ction_gemm_546922f892bf12518fd7_binding.cu | 0 ...ection_gemm_546922f892bf12518fd7_kernel.cu | 0 ...ction_gemm_5aabaa7971ddc277e20f_binding.cu | 0 ...ection_gemm_5aabaa7971ddc277e20f_kernel.cu | 0 ...ction_gemm_5d34d9ca403abf798d9a_binding.cu | 0 ...ection_gemm_5d34d9ca403abf798d9a_kernel.cu | 0 ...ction_gemm_628bc3803eccd349044e_binding.cu | 0 ...ection_gemm_628bc3803eccd349044e_kernel.cu | 0 ...ction_gemm_74a2e3f388f79fe0f8ae_binding.cu | 0 ...ection_gemm_74a2e3f388f79fe0f8ae_kernel.cu | 0 ...ction_gemm_7555c5ce07ba13117bb3_binding.cu | 0 ...ection_gemm_7555c5ce07ba13117bb3_kernel.cu | 0 ...ction_gemm_7a259feabf02e3ac9e0f_binding.cu | 0 ...ection_gemm_7a259feabf02e3ac9e0f_kernel.cu | 0 ...ction_gemm_86a64ca05fce781cd7f7_binding.cu | 0 ...ection_gemm_86a64ca05fce781cd7f7_kernel.cu | 0 ...ction_gemm_95a568c7ee4ec7e19d68_binding.cu | 0 ...ection_gemm_95a568c7ee4ec7e19d68_kernel.cu | 0 ...ction_gemm_97ce792e4518a4abbdf2_binding.cu | 0 ...ection_gemm_97ce792e4518a4abbdf2_kernel.cu | 0 ...ction_gemm_9f7c65acfbefcf9442c4_binding.cu | 0 ...ection_gemm_9f7c65acfbefcf9442c4_kernel.cu | 0 ...ction_gemm_a4b267f62f75fe664995_binding.cu | 0 ...ection_gemm_a4b267f62f75fe664995_kernel.cu | 0 ...ction_gemm_bbd8e1061de96c7ebc7b_binding.cu | 0 ...ection_gemm_bbd8e1061de96c7ebc7b_kernel.cu | 0 ...ction_gemm_c8633003630e8311a4b3_binding.cu | 0 ...ection_gemm_c8633003630e8311a4b3_kernel.cu | 0 ...ction_gemm_d11c33dd107e3b89124a_binding.cu | 0 ...ection_gemm_d11c33dd107e3b89124a_kernel.cu | 0 ...ction_gemm_d2a44f7c9244b89772a2_binding.cu | 0 ...ection_gemm_d2a44f7c9244b89772a2_kernel.cu | 0 ...ction_gemm_d384571ea88c28e40e31_binding.cu | 0 ...ection_gemm_d384571ea88c28e40e31_kernel.cu | 0 ...ction_gemm_d9095ad478b0244af2f4_binding.cu | 0 ...ection_gemm_d9095ad478b0244af2f4_kernel.cu | 0 ...ction_gemm_e03279a175f68c4f2efe_binding.cu | 0 ...ection_gemm_e03279a175f68c4f2efe_kernel.cu | 0 ...ction_gemm_e3489655388220628088_binding.cu | 0 ...ection_gemm_e3489655388220628088_kernel.cu | 0 ...ction_gemm_e717e449f20f5a8edfc3_binding.cu | 0 ...ection_gemm_e717e449f20f5a8edfc3_kernel.cu | 0 ...ction_gemm_ec27209d6e178b952413_binding.cu | 0 ...ection_gemm_ec27209d6e178b952413_kernel.cu | 0 ...ction_gemm_f557f68a025cb6450dda_binding.cu | 0 ...ection_gemm_f557f68a025cb6450dda_kernel.cu | 0 .../csrc/cake_ffi_checks.h | 0 pyproject.toml | 6 +- scripts/check_cake_build_manifests.py | 12 +-- tests/test_helpers/cake_package.py | 2 +- 74 files changed, 68 insertions(+), 68 deletions(-) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/.gitignore (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/README.md (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/__init__.py (94%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/build.json (99%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/cake_backend.py (98%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/cake_build.py (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/cake_jit.py (99%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_binding.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_kernel.cu (100%) rename flashinfer/experimental/{dense_projection_gemm => cake/cake_dense_gemm_nondet_train}/csrc/cake_ffi_checks.h (100%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 373052d22e5..453af9f749c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,13 +40,13 @@ files: ^(?!(?:flashinfer/moe_ep/kernel_src/blackwell_bf16_rank_major/src/|csrc/c # and sealed by cake_sm120_megamoe_manifest.json; keep their bytes. # The csrc/ trees of the Cake-exported packages (flashinfer-ai/flashinfer#6553: the rmsnorm # training and mHC training backends, the causal conv1d, ulysses and fused QK RoPE append (BF16 / FP8) backends, cake_dsa_train, cake_kda_train / cake_kda_det_train / cake_kda_det_train_kcp, cake_dsa_indexer_det, cake_mok_bf16, cake_moe_grouped_det_gemm, -# cake_lm_head_loss / cake_lm_head_det_loss, cake_dense_gemm_det_train, dense_projection_gemm, cake_stepfun_moe, nvfp4_attention, cake_nvfp4_ptilemax_attention, deepgemm_fp8_gemm, deepgemm_batched_gemm, deepgemm_mega_gate, nvfp4_sparse_mla_decode) are rendered by the Cake exporters and +# cake_lm_head_loss / cake_lm_head_det_loss, cake_dense_gemm_det_train, cake_dense_gemm_nondet_train, cake_stepfun_moe, nvfp4_attention, cake_nvfp4_ptilemax_attention, deepgemm_fp8_gemm, deepgemm_batched_gemm, deepgemm_mega_gate, nvfp4_sparse_mla_decode) are rendered by the Cake exporters and # digest-checked by their build.json manifests; keep their bytes. The package-root loaders, # cake_build.py and build.json next to them stay checked. # Keep both old and det package paths while the KDA, conv1d and LM-head renames land independently. # The Qwen 3.5 gated attention block's _cudnn_frost/ tree is cuDNN Frontend source vendored byte-faithfully # (VENDOR_MANIFEST.json digests the files as written); keep its bytes. -exclude: ^(?:csrc/kda/ptx/|flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_trtllm_mla_blackwell/generated/|cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|cake_minimax_h3_qkv_quantize_pack/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_fp4_allreduce/.*_kernel\.cu|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_union/|cake_gemm_allreduce/|cake_fused_norm_combine/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:cake_rmsnorm_det_train_backend/csrc/|cake_rmsnorm_train_backend/csrc/|cake_causal_conv1d_train_backend/csrc/|cake_causal_conv1d_det_train_backend/csrc/|experimental/cake/cake_causal_conv1d_train_backend/csrc/|cake_sage_block_sparse_attention_backend/csrc/|cake_ulysses_backend/csrc/|cake_fused_qk_rope_append_backend/csrc/|cake_fused_qk_rope_fp8_append_backend/csrc/|cake_mhc_det_train_backend/csrc/|experimental/cake/cake_mhc_train_backend/csrc/|experimental/cake/cake_lm_head_loss/csrc/|experimental/cake/cake_moe_grouped_gemm/csrc/|experimental/cake/cake_dense_gemm_det_train/csrc/|experimental/nvfp4_sparse_mla_decode/csrc/(?:cake_nvfp4_sparse_mla_decode/|cake_ffi_checks\.h)|experimental/(?:cake_mok_bf16|cake_mxfp8_megamoe_ep16|cake_moe_grouped_det_gemm|sm110_gqa_decode|nvfp4_attention|cake_nvfp4_ptilemax_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention|kimi_k3_fused_router|msa_nvfp4_decode|cake_lm_head_loss|cake_lm_head_det_loss|cake_dsa_train|cake_kda_train|cake_kda_det_train|cake_kda_det_train_kcp|cake_dsa_indexer_det|cake_msa_icp_indexer|dense_projection_gemm|cake_stepfun_moe|deepgemm_fp8_gemm|deepgemm_batched_gemm|deepgemm_mega_gate)/csrc/|moe_ep/kernel_src/(?:(?:sm100/(?:cutedsl_megamoe|next_cutedsl_megamoe)|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/(?:swapab|split)_cutedsl_megakernel)/src/|sm90/cake_bf16_megamoe/src/cake_sm90_bf16_megamoe_|sm120/cake_mxfp4_megamoe/src/cake_sm120_megamoe_(?:kernel.*\.cuh|m(?:128|64)_(?:no)?copier\.cu|geometry\.json)))|csrc/generated/dcp_alltoall/|flashinfer/experimental/qwen_3_5_block/_cudnn_frost/) +exclude: ^(?:csrc/kda/ptx/|flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_trtllm_mla_blackwell/generated/|cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|cake_minimax_h3_qkv_quantize_pack/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_fp4_allreduce/.*_kernel\.cu|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_union/|cake_gemm_allreduce/|cake_fused_norm_combine/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:cake_rmsnorm_det_train_backend/csrc/|cake_rmsnorm_train_backend/csrc/|cake_causal_conv1d_train_backend/csrc/|cake_causal_conv1d_det_train_backend/csrc/|experimental/cake/cake_causal_conv1d_train_backend/csrc/|cake_sage_block_sparse_attention_backend/csrc/|cake_ulysses_backend/csrc/|cake_fused_qk_rope_append_backend/csrc/|cake_fused_qk_rope_fp8_append_backend/csrc/|cake_mhc_det_train_backend/csrc/|experimental/cake/cake_mhc_train_backend/csrc/|experimental/cake/cake_lm_head_loss/csrc/|experimental/cake/cake_moe_grouped_gemm/csrc/|experimental/cake/cake_dense_gemm_det_train/csrc/|experimental/cake/cake_dense_gemm_nondet_train/csrc/|experimental/nvfp4_sparse_mla_decode/csrc/(?:cake_nvfp4_sparse_mla_decode/|cake_ffi_checks\.h)|experimental/(?:cake_mok_bf16|cake_mxfp8_megamoe_ep16|cake_moe_grouped_det_gemm|sm110_gqa_decode|nvfp4_attention|cake_nvfp4_ptilemax_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention|kimi_k3_fused_router|msa_nvfp4_decode|cake_lm_head_loss|cake_lm_head_det_loss|cake_dsa_train|cake_kda_train|cake_kda_det_train|cake_kda_det_train_kcp|cake_dsa_indexer_det|cake_msa_icp_indexer|cake_stepfun_moe|deepgemm_fp8_gemm|deepgemm_batched_gemm|deepgemm_mega_gate)/csrc/|moe_ep/kernel_src/(?:(?:sm100/(?:cutedsl_megamoe|next_cutedsl_megamoe)|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/(?:swapab|split)_cutedsl_megakernel)/src/|sm90/cake_bf16_megamoe/src/cake_sm90_bf16_megamoe_|sm120/cake_mxfp4_megamoe/src/cake_sm120_megamoe_(?:kernel.*\.cuh|m(?:128|64)_(?:no)?copier\.cu|geometry\.json)))|csrc/generated/dcp_alltoall/|flashinfer/experimental/qwen_3_5_block/_cudnn_frost/) repos: # Standard hooks diff --git a/flashinfer/experimental/dense_projection_gemm/.gitignore b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/.gitignore similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/.gitignore rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/.gitignore diff --git a/flashinfer/experimental/dense_projection_gemm/README.md b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/README.md rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md diff --git a/flashinfer/experimental/dense_projection_gemm/__init__.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py similarity index 94% rename from flashinfer/experimental/dense_projection_gemm/__init__.py rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py index 7bbaa3feae9..46994a228c1 100644 --- a/flashinfer/experimental/dense_projection_gemm/__init__.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py @@ -21,7 +21,7 @@ # the keyword-only ``deterministic: bool = True`` and accepts only # ``deterministic=False``: any other value, the default included, raises # ``ValueError`` naming the entry point of the deterministic package -# ``flashinfer.experimental.dense_projection_det_gemm``. Importing this package +# ``flashinfer.experimental.cake.cake_dense_gemm_det_train``. Importing this package # imports nothing: ``cake_backend``, ``cake_jit`` and ``cake_build`` load on # first attribute access. diff --git a/flashinfer/experimental/dense_projection_gemm/build.json b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json similarity index 99% rename from flashinfer/experimental/dense_projection_gemm/build.json rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json index a479f422167..4b3f17665c8 100644 --- a/flashinfer/experimental/dense_projection_gemm/build.json +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json @@ -1315,6 +1315,6 @@ ] } ], - "package": "flashinfer.experimental.dense_projection_gemm", + "package": "flashinfer.experimental.cake.cake_dense_gemm_nondet_train", "schema": "cake-build/1" } diff --git a/flashinfer/experimental/dense_projection_gemm/cake_backend.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py similarity index 98% rename from flashinfer/experimental/dense_projection_gemm/cake_backend.py rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py index ea0aa3d9dd3..cf5f194827b 100644 --- a/flashinfer/experimental/dense_projection_gemm/cake_backend.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py @@ -16,7 +16,7 @@ Cake backend, non-deterministic: the GLM-5.2 dense projection GEMMs for training (K1) and the FP32 router-gate GEMM (K2) on SM100 / SM103 / SM107. -The deterministic implementation is the package ``dense_projection_det_gemm`` (flashinfer-ai/flashinfer#5677 and +The deterministic implementation is the package ``cake_dense_gemm_det_train`` (flashinfer-ai/flashinfer#5677 and #7043): its stream-K and split-K reductions add the partials in a fixed order, so every result is bitwise reproducible run to run. This package keeps that package's operations, view contract, dtypes, accuracy budgets and host paths (prepared and convenience) and adds programs whose split-K reductions add the partials in the order they @@ -32,7 +32,7 @@ * Any other value -- the default, ``True``, ``None``, ``0``, ``1``, ``"False"`` -- raises ``ValueError`` before anything is planned, allocated or launched. The message names the deterministic package's entry point of the - operation (``projection_wgrad`` -> ``projection_det_wgrad``, ...). The keyword never selects the deterministic + operation (``projection_wgrad`` -> ``cake_dense_gemm_det_train_projection_wgrad``, ...). The keyword never selects the deterministic package, and nothing falls back to it. * ``deterministic=False`` runs the routes of :data:`ROUTES`: @@ -139,18 +139,18 @@ #: The deterministic package's entry points this package calls on the rows that run the deterministic programs #: (see "The two packages" above); :func:`use_deterministic_backend` checks that a bound module has them. _DET_ENTRIES = ( - "dense_projection_det_gemm", - "plan_dense_projection_det_gemm", - "plan_router_fp32_det_gemm", - "prepare_dense_projection_det_gemm", - "prepare_projection_det_wgrad", - "prepare_router_fp32_det_gemm", - "projection_det_dgrad", - "projection_det_forward", - "projection_det_wgrad", - "router_fp32_det_gemm", + "cake_dense_gemm_det_train_projection", + "cake_dense_gemm_det_train_projection_plan", + "cake_dense_gemm_det_train_router_plan", + "cake_dense_gemm_det_train_projection_prepare", + "cake_dense_gemm_det_train_projection_wgrad_prepare", + "cake_dense_gemm_det_train_router_prepare", + "cake_dense_gemm_det_train_projection_dgrad", + "cake_dense_gemm_det_train_projection_forward", + "cake_dense_gemm_det_train_projection_wgrad", + "cake_dense_gemm_det_train_router", "split_fp32_to_bf16x3_det_", - "clear_convenience_caches", + "cake_dense_gemm_det_train_convenience_cache_clear", ) @@ -2250,7 +2250,7 @@ def router_partition( raise ValueError( f"router_fp32_gemm(deterministic=False): the forward takes at most {ROUTER_FORWARD_MAX_SPLITS} " f"splits (more partials per element do not repeat bit for bit under recompute), got {s}; the " - f"deterministic entry {_det_package_name()}.cake_backend.router_fp32_det_gemm takes more" + f"deterministic entry {_det_package_name()}.cake_backend.cake_dense_gemm_det_train_router takes more" ) spec = f"S{s}" else: @@ -3389,8 +3389,8 @@ def _atomic_k2( if plan is None: det = deterministic_backend() if op == "plan": - return det.plan_router_fp32_det_gemm(A, B, out, splits=splits) - return det.prepare_router_fp32_det_gemm( + return det.cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) + return det.cake_dense_gemm_det_train_router_prepare( A, B, out, splits=splits, **_pdl_kw(enable_pdl) ) if op == "plan": @@ -4612,7 +4612,7 @@ def _plan_splitk_k2( det = deterministic_backend() device = out.device arch = det.require_arch(device) - plan, A_eff, B_eff = det.plan_router_fp32_det_gemm(A, B, out, splits=splits) + plan, A_eff, B_eff = det.cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) a_view, split_src = det._k2_checked_views(A, B, A_eff, B_eff, plan) mine = _k2_views(A, B, plan) if not (det._same_view(mine[0], a_view) and det._same_view(mine[1], split_src)): @@ -4710,7 +4710,7 @@ def _launch_cached_splitk_k2( _plan_splitk_k2(A, B, out, splits, index, stream, pdl), ) if entry.k2 is None or entry.entry is None or entry.arena is None: - return det.router_fp32_det_gemm(A, B, out, splits=splits, enable_pdl=bool(pdl)) + return det.cake_dense_gemm_det_train_router(A, B, out, splits=splits, enable_pdl=bool(pdl)) a_view, split_src = _launch_views(A, B, entry.plan, views) _launch_splitk_k2_binding(det, entry, out, a_view, split_src, index, stream) return out @@ -4975,7 +4975,7 @@ def clear_convenience_caches() -> None: package's rule: a CUDA graph captures its prepared launch, not its convenience path. The graphs captured on one stream share its arena: replay them one at a time, and not concurrently with calls on that stream.""" _drop_convenience_entries() - deterministic_backend().clear_convenience_caches() + deterministic_backend().cake_dense_gemm_det_train_convenience_cache_clear() # --------------------------------------------------------------------------- @@ -5032,12 +5032,12 @@ def plan_dense_projection_gemm( (:class:`AtomicGemmPlan`; the knobs ``splits`` / ``block_n`` / ``chains`` are this package's, the others -- ``sm_count``, ``l2_bytes``, ``transposed_out``, the planner knobs, ``arch`` -- apply when the row runs the deterministic program). A row the planner keeps on the deterministic program, and every other layout, returns the - deterministic package's plan (``plan_dense_projection_det_gemm``).""" + deterministic package's plan (``cake_dense_gemm_det_train_projection_plan``).""" _require_false( - deterministic, "plan_dense_projection_gemm", "plan_dense_projection_det_gemm" + deterministic, "plan_dense_projection_gemm", "cake_dense_gemm_det_train_projection_plan" ) if _route(_k1_layout(A, B)) == "det": - return deterministic_backend().plan_dense_projection_det_gemm( + return deterministic_backend().cake_dense_gemm_det_train_projection_plan( A, B, out, **kwargs ) own, det_kwargs = _split_knobs(kwargs) @@ -5050,7 +5050,7 @@ def plan_dense_projection_gemm( own=own, arch=det_kwargs.get("arch"), sm_count=det_kwargs.get("sm_count"), - det_call=lambda: deterministic_backend().plan_dense_projection_det_gemm( + det_call=lambda: deterministic_backend().cake_dense_gemm_det_train_projection_plan( A, B, out, **det_kwargs ), ) @@ -5068,16 +5068,16 @@ def prepare_dense_projection_gemm( ``deterministic=False`` (the only value accepted) a weight-gradient layout prepares this package's split-K program (:class:`PreparedAtomicGemm`; its workspace is owned by the prepared object; ``splits`` / ``block_n`` / ``chains`` as for :func:`plan_dense_projection_gemm`). A row the planner keeps on the deterministic program, and every other - layout, prepares the deterministic package's ``prepare_dense_projection_det_gemm`` (its keywords: + layout, prepares the deterministic package's ``cake_dense_gemm_det_train_projection_prepare`` (its keywords: ``transposed_out`` and the planner knobs). ``enable_pdl`` (default: the device supports it) makes either launch a programmatic dependent launch.""" _require_false( deterministic, "prepare_dense_projection_gemm", - "prepare_dense_projection_det_gemm", + "cake_dense_gemm_det_train_projection_prepare", ) if _route(_k1_layout(A, B)) == "det": - return deterministic_backend().prepare_dense_projection_det_gemm( + return deterministic_backend().cake_dense_gemm_det_train_projection_prepare( A, B, out, **kwargs ) own, det_kwargs = _split_knobs(kwargs) @@ -5088,7 +5088,7 @@ def prepare_dense_projection_gemm( out, transposed_out=bool(det_kwargs.get("transposed_out", False)), own=own, - det_call=lambda: deterministic_backend().prepare_dense_projection_det_gemm( + det_call=lambda: deterministic_backend().cake_dense_gemm_det_train_projection_prepare( A, B, out, **det_kwargs ), enable_pdl=det_kwargs.get("enable_pdl"), @@ -5111,11 +5111,11 @@ def dense_projection_gemm( row; every other layout and row runs the deterministic program. ``enable_pdl`` (default: the device supports it) launches either program as a programmatic dependent launch; each waits for the preceding kernel before its first global access, and ``weights_static=True`` lets the deterministic swapped small-token forward start its loads of - ``B`` before the wait (see the deterministic package's ``dense_projection_det_gemm``).""" - _require_false(deterministic, "dense_projection_gemm", "dense_projection_det_gemm") + ``B`` before the wait (see the deterministic package's ``cake_dense_gemm_det_train_projection``).""" + _require_false(deterministic, "dense_projection_gemm", "cake_dense_gemm_det_train_projection") def det_call() -> Any: - return deterministic_backend().dense_projection_det_gemm( + return deterministic_backend().cake_dense_gemm_det_train_projection( A, B, out, @@ -5165,7 +5165,7 @@ def projection_forward( deterministic plan, program and bits, except on a row of :data:`K1_FORWARD_RULES_BY_SM` for the device, which runs the program the deterministic planner gives it with the table's knobs (:func:`k1_forward_knobs`; other bits, still the same on every run). ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - _require_false(deterministic, "projection_forward", "projection_det_forward") + _require_false(deterministic, "projection_forward", "cake_dense_gemm_det_train_projection_forward") if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.bfloat16) o = out @@ -5177,7 +5177,7 @@ def projection_forward( (X, W, o), ("fwd",), lambda: (X, W.t(), o, False, _k1_forward_knobs_of(X, W, o)), - lambda: deterministic_backend().projection_det_forward( + lambda: deterministic_backend().cake_dense_gemm_det_train_projection_forward( X, W, o, **_pdl_kw(enable_pdl) ), enable_pdl, @@ -5196,7 +5196,7 @@ def projection_dgrad( """``out[T, K] = G[T, N] @ W[N, K]``, ``deterministic=False`` only. The deterministic program (whole output tiles) through this package's binding of it (:func:`_det_k1`: the same plan, program and bits). ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - _require_false(deterministic, "projection_dgrad", "projection_det_dgrad") + _require_false(deterministic, "projection_dgrad", "cake_dense_gemm_det_train_projection_dgrad") if out is None: out = torch.empty( G.shape[0], @@ -5213,7 +5213,7 @@ def projection_dgrad( (G, W, o), ("dgrad",), lambda: (G, W, o, False, ()), - lambda: deterministic_backend().projection_det_dgrad( + lambda: deterministic_backend().cake_dense_gemm_det_train_projection_dgrad( G, W, o, out_dtype=out_dtype, **_pdl_kw(enable_pdl) ), enable_pdl, @@ -5232,10 +5232,10 @@ def prepare_projection_wgrad( ) -> Any: """Prepared ``out[N, K] = G[T, N].T @ X[T, K]``, ``deterministic=False`` only: this package's split-K program as the transposed-store GEMM ``X.T @ G`` (``M = K``, ``N = N``) where the planner takes the row; a row it does not - take prepares the deterministic package's ``prepare_projection_det_wgrad`` (its orientation, ``sk``, + take prepares the deterministic package's ``cake_dense_gemm_det_train_projection_wgrad_prepare`` (its orientation, ``sk``, ``cta_rows``). ``enable_pdl`` as for :func:`prepare_dense_projection_gemm`.""" _require_false( - deterministic, "prepare_projection_wgrad", "prepare_projection_det_wgrad" + deterministic, "prepare_projection_wgrad", "cake_dense_gemm_det_train_projection_wgrad_prepare" ) return _atomic_k1( "prepare", @@ -5244,7 +5244,7 @@ def prepare_projection_wgrad( out, transposed_out=True, own={}, - det_call=lambda: deterministic_backend().prepare_projection_det_wgrad( + det_call=lambda: deterministic_backend().cake_dense_gemm_det_train_projection_wgrad_prepare( G, X, out, sk=sk, cta_rows=cta_rows, **_pdl_kw(enable_pdl) ), enable_pdl=enable_pdl, @@ -5264,10 +5264,10 @@ def projection_wgrad( ) -> torch.Tensor: """``out[N, K] = G[T, N].T @ X[T, K]`` (ragged reduction length ``T``), ``deterministic=False`` only: this package's split-K program where the planner takes the row (the split partials summed in the order they finish, - the output rounded once), else the deterministic package's ``projection_det_wgrad`` program (its orientation, + the output rounded once), else the deterministic package's ``cake_dense_gemm_det_train_projection_wgrad`` program (its orientation, ``sk`` and ``cta_rows``) through this package's binding of it (:func:`_launch_cached_atomic_k1`); ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - _require_false(deterministic, "projection_wgrad", "projection_det_wgrad") + _require_false(deterministic, "projection_wgrad", "cake_dense_gemm_det_train_projection_wgrad") if out is None: out = torch.empty( int(G.shape[1]), @@ -5293,7 +5293,7 @@ def projection_wgrad( out, transposed_out=True, own={}, - det_call=lambda: deterministic_backend().projection_det_wgrad( + det_call=lambda: deterministic_backend().cake_dense_gemm_det_train_projection_wgrad( G, X, out, @@ -5315,7 +5315,7 @@ def _wgrad_det_args( sk: bool | str, cta_rows: Optional[int], ) -> tuple: - """The deterministic package's convenience call of ``projection_det_wgrad(G, X, out, sk=sk, + """The deterministic package's convenience call of ``cake_dense_gemm_det_train_projection_wgrad(G, X, out, sk=sk, cta_rows=cta_rows)``: ``(A, B, out, transposed_out, knobs)`` in its orientation (``wgrad_views``).""" det = deterministic_backend() A, B, transposed = det.wgrad_views( @@ -5335,10 +5335,10 @@ def plan_router_fp32_gemm( """Classify the fp32 views and plan the K2 launch, ``deterministic=False`` only: this package's stream-K program for the ``kk`` / ``nn_t`` classes and the SIMT programs in their windows (:class:`AtomicRouterPlan`; ``splits`` sets ``splits`` units per tile, see :func:`router_units`); a row the planner keeps on the deterministic program, - and a layout class routed ``"det"``, returns the deterministic package's plan (``plan_router_fp32_det_gemm``).""" - _require_false(deterministic, "plan_router_fp32_gemm", "plan_router_fp32_det_gemm") + and a layout class routed ``"det"``, returns the deterministic package's plan (``cake_dense_gemm_det_train_router_plan``).""" + _require_false(deterministic, "plan_router_fp32_gemm", "cake_dense_gemm_det_train_router_plan") if _route(_k2_layout(A, B)) == "det": - return deterministic_backend().plan_router_fp32_det_gemm( + return deterministic_backend().cake_dense_gemm_det_train_router_plan( A, B, out, splits=splits ) return _atomic_k2("plan", A, B, out, splits=splits) @@ -5357,14 +5357,14 @@ def prepare_router_fp32_gemm( package's stream-K program for the ``kk`` / ``nn_t`` classes (its launch splits the second operand, zeroes ``out`` on the stream when a segment adds into it, and runs the program: no partial buffers), and a SIMT program in its windows (the ``nn_t`` and ``kn`` classes: one kernel over the fp32 operands); a row the planner keeps on the - deterministic program prepares the deterministic package's ``prepare_router_fp32_det_gemm``. ``enable_pdl`` + deterministic program prepares the deterministic package's ``cake_dense_gemm_det_train_router_prepare``. ``enable_pdl`` (default: the device supports it) launches the split and the stream-K program as programmatic dependent launches.""" _require_false( - deterministic, "prepare_router_fp32_gemm", "prepare_router_fp32_det_gemm" + deterministic, "prepare_router_fp32_gemm", "cake_dense_gemm_det_train_router_prepare" ) if _route(_k2_layout(A, B)) == "det": - return deterministic_backend().prepare_router_fp32_det_gemm( + return deterministic_backend().cake_dense_gemm_det_train_router_prepare( A, B, out, splits=splits, **_pdl_kw(enable_pdl) ) return _atomic_k2("prepare", A, B, out, splits=splits, enable_pdl=enable_pdl) @@ -5404,7 +5404,7 @@ def router_fp32_gemm( (:func:`router_simt`) and the ``kk`` token counts of a forward SIMT window on the forward's cluster program (:func:`router_forward_simt`); every other row runs the deterministic program. ``enable_pdl`` as for :func:`prepare_router_fp32_gemm`.""" - _require_false(deterministic, "router_fp32_gemm", "router_fp32_det_gemm") + _require_false(deterministic, "router_fp32_gemm", "cake_dense_gemm_det_train_router") return _router_gemm(A, B, out, splits, enable_pdl) @@ -5423,7 +5423,7 @@ def router_forward( :func:`router_forward_simt`), else the deterministic program (its plan for ``router_fp32_gemm(X, W.T, out)``, through this package's binding, :func:`_launch_cached_splitk_k2`): without a split count, while no router table gives the forward a verdict, without the per-call check of the row.""" - _require_false(deterministic, "router_forward", "router_det_forward") + _require_false(deterministic, "router_forward", "cake_dense_gemm_det_train_router_forward") if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.float32) if _k2_entry("fwd", X, W, out, splits, enable_pdl) is not None: @@ -5450,7 +5450,7 @@ def router_dgrad( place, every element of ``out`` stored once), and at every other row and token count the deterministic program (one split: nothing to reduce) through this package's binding of it (:func:`_launch_cached_splitk_k2`: the same plan, program and bits). ``enable_pdl`` as for :func:`prepare_router_fp32_gemm`.""" - _require_false(deterministic, "router_dgrad", "router_det_dgrad") + _require_false(deterministic, "router_dgrad", "cake_dense_gemm_det_train_router_dgrad") if out is None: out = torch.empty(G.shape[0], W.shape[1], device=G.device, dtype=torch.float32) if _k2_entry("dgrad", G, W, out, None, enable_pdl) is not None: @@ -5472,7 +5472,7 @@ def router_wgrad( with a table of its own, in :data:`ROUTER_RULES_BY_SM` (:func:`router_partition`; a row or token count without one runs the deterministic program, through this package's binding of it without planning the row here, :func:`_launch_cached_splitk_k2`).""" - _require_false(deterministic, "router_wgrad", "router_det_wgrad") + _require_false(deterministic, "router_wgrad", "cake_dense_gemm_det_train_router_wgrad") if out is None: out = torch.empty(G.shape[1], X.shape[1], device=G.device, dtype=torch.float32) if _k2_entry("wgrad", G, X, out, splits, enable_pdl) is not None: diff --git a/flashinfer/experimental/dense_projection_gemm/cake_build.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_build.py similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/cake_build.py rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_build.py diff --git a/flashinfer/experimental/dense_projection_gemm/cake_jit.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py similarity index 99% rename from flashinfer/experimental/dense_projection_gemm/cake_jit.py rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py index cbac5023354..0503a1a511e 100644 --- a/flashinfer/experimental/dense_projection_gemm/cake_jit.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py @@ -26,7 +26,7 @@ from . import cake_build # Registry of the generated programs of this package (schema v2, the layout of the deterministic package's -# ``cake_jit``). Populated verbatim by the generated-program export (``export.families.dense_projection_gemm``); do +# ``cake_jit``). Populated verbatim by the generated-program export (the exporter family of this package); do # not edit by hand. Three literals, each wrapped in ``# fmt: off`` / ``# fmt: on`` and rendered one record per line, # so the export replaces a literal from its head line up to the next ``# fmt: on`` line: # @@ -128,10 +128,10 @@ "dense_proj_gemm_dsm_", ) -#: The deterministic package: a sibling of this package (``.dense_projection_det_gemm``), resolved on first +#: The deterministic package: a sibling of this package (``.cake_dense_gemm_det_train``), resolved on first #: use and never through an absolute ``flashinfer`` import; a vendored copy binds it explicitly #: (``cake_backend.use_deterministic_backend``). -DETERMINISTIC_PACKAGE = "dense_projection_det_gemm" +DETERMINISTIC_PACKAGE = "cake_dense_gemm_det_train" _det_package: Optional[str] = None diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_16496106e8ac1c7b1048_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_1d46f84995b7baec3e04_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2395f1c01b99015f5dfe_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_23cd6c7b07c9130bef18_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_243446e7407f2bed9b45_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_2f07a521e0f7ebd21021_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_3f431ee1064967195e90_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_4b13414a327c7040abd8_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_546922f892bf12518fd7_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5aabaa7971ddc277e20f_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_5d34d9ca403abf798d9a_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_628bc3803eccd349044e_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_74a2e3f388f79fe0f8ae_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7555c5ce07ba13117bb3_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_7a259feabf02e3ac9e0f_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_86a64ca05fce781cd7f7_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_95a568c7ee4ec7e19d68_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_97ce792e4518a4abbdf2_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_9f7c65acfbefcf9442c4_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_a4b267f62f75fe664995_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_bbd8e1061de96c7ebc7b_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_c8633003630e8311a4b3_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d11c33dd107e3b89124a_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d2a44f7c9244b89772a2_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d9095ad478b0244af2f4_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e03279a175f68c4f2efe_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e3489655388220628088_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_e717e449f20f5a8edfc3_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_ec27209d6e178b952413_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_binding.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_binding.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_binding.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_binding.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_kernel.cu b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_kernel.cu similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_kernel.cu rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_f557f68a025cb6450dda_kernel.cu diff --git a/flashinfer/experimental/dense_projection_gemm/csrc/cake_ffi_checks.h b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_ffi_checks.h similarity index 100% rename from flashinfer/experimental/dense_projection_gemm/csrc/cake_ffi_checks.h rename to flashinfer/experimental/cake/cake_dense_gemm_nondet_train/csrc/cake_ffi_checks.h diff --git a/pyproject.toml b/pyproject.toml index b3ac5e4c3c5..bf201c3a629 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -191,9 +191,9 @@ exclude = [ # Generated sources for the deterministic SM100a/SM103a/SM107a dense GEMM training backend # (flashinfer.cake_dense_gemm_det_train). "flashinfer.experimental.cake.cake_dense_gemm_det_train" = ["csrc/**", "build.json"] -# Generated sources for the experimental SM100a/SM103a/SM107a dense projection GEMM backend with an explicit -# numerical mode (the non-deterministic split-K programs; the deterministic ones ship with dense_projection_det_gemm). -"flashinfer.experimental.dense_projection_gemm" = ["csrc/**", "build.json"] +# Generated sources for the non-deterministic SM100a/SM103a/SM107a dense GEMM training backend +# (flashinfer.cake_dense_gemm_nondet_train; the deterministic programs it reuses ship with cake_dense_gemm_det_train). +"flashinfer.experimental.cake.cake_dense_gemm_nondet_train" = ["csrc/**", "build.json"] # Generated sources for the experimental exact-SM100a/SM103a StepFun fused-MoE backend (cake-build/1). "flashinfer.experimental.cake_stepfun_moe" = ["csrc/**", "build.json"] # Standalone cuDNN Frost Python sources and the native MoE adapter, compiled on first use. diff --git a/scripts/check_cake_build_manifests.py b/scripts/check_cake_build_manifests.py index 26f786e2f2f..ac04146ad38 100755 --- a/scripts/check_cake_build_manifests.py +++ b/scripts/check_cake_build_manifests.py @@ -42,8 +42,8 @@ * the ``cake_build.py`` files of all checked packages are byte-identical; * ``common.requires.tvm_ffi.python`` equals line 1 of ``requirements.txt``; * every package a package requires (``Package.requires``: a package it imports - as its sibling, as the numerical-mode ``dense_projection_gemm`` imports the - deterministic ``dense_projection_det_gemm``) is a row of this table with the + as its sibling, as the non-deterministic ``cake_dense_gemm_nondet_train`` imports + the deterministic ``cake_dense_gemm_det_train``) is a row of this table with the same parent package. Vendored gate (``--vendored-gate``; combinable with ``--package``): the form @@ -545,12 +545,12 @@ class Package: refs_generic, ), Package( - "flashinfer.experimental.dense_projection_gemm", - "flashinfer/experimental/dense_projection_gemm", + "flashinfer.experimental.cake.cake_dense_gemm_nondet_train", + "flashinfer/experimental/cake/cake_dense_gemm_nondet_train", "cake_jit.py", refs_generic, - # its deterministic path calls the deterministic package's programs - requires=("flashinfer.experimental.dense_projection_det_gemm",), + # its fixed-order rows run the deterministic package's programs + requires=("flashinfer.experimental.cake.cake_dense_gemm_det_train",), ), Package( "flashinfer.experimental.cake_stepfun_moe", diff --git a/tests/test_helpers/cake_package.py b/tests/test_helpers/cake_package.py index 1ae68494ea3..4d7ff692b12 100644 --- a/tests/test_helpers/cake_package.py +++ b/tests/test_helpers/cake_package.py @@ -31,7 +31,7 @@ # the non-deterministic package (deterministic=False); its key carries the namespace "cake.cake_lm_head_loss": "flashinfer.experimental.cake.cake_lm_head_loss", "cake_dense_gemm_det_train": "flashinfer.experimental.cake.cake_dense_gemm_det_train", - "dense_projection_gemm": "flashinfer.experimental.dense_projection_gemm", + "cake_dense_gemm_nondet_train": "flashinfer.experimental.cake.cake_dense_gemm_nondet_train", "cake_moe_grouped_det_gemm": "flashinfer.experimental.cake_moe_grouped_det_gemm", # the suite name before the deterministic-family rename (compatibility alias) "cake_moe_grouped_gemm": "flashinfer.experimental.cake_moe_grouped_det_gemm", From 417ca3dfdf68ddb8e2c4d43f8f938afce4f523e6 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 00:12:31 -0800 Subject: [PATCH 08/13] refactor(cake_dense_gemm_nondet_train): #8928 names, no deterministic= keyword The non-deterministic dense GEMMs become the public module flashinfer.cake_dense_gemm_nondet_train (flashinfer-ai/flashinfer#8928). The module selects the mode: no entry point takes a mode argument. - The entry points and types take the module name: dense_projection_gemm becomes cake_dense_gemm_nondet_train_projection, router_fp32_gemm becomes cake_dense_gemm_nondet_train_router, the forward / dgrad / wgrad / plan / prepare entries, the convenience cache clear and the four plan and launch types follow. The former names fail at import; there is no alias. - The deterministic= keyword is removed from every entry point (passing it raises TypeError). The projection plan and prepare entries list their keyword-only knobs explicitly (the deterministic planner's public knobs plus splits and chains) instead of **kwargs; only the knobs a caller gives reach the deterministic planner, so routing and plans are unchanged. - New support query is_cake_dense_gemm_nondet_train_supported(device). - The PyTorch deterministic-algorithms alert is unchanged: it fires only right before an order-varying launch and names the public modules. - The public module imports the backend on first attribute access and returns the backend's own objects; it is not an experimental API. --- flashinfer/cake_dense_gemm_nondet_train.py | 96 +++ .../cake_dense_gemm_nondet_train/__init__.py | 17 +- .../cake_backend.py | 599 ++++++++++-------- 3 files changed, 435 insertions(+), 277 deletions(-) create mode 100644 flashinfer/cake_dense_gemm_nondet_train.py diff --git a/flashinfer/cake_dense_gemm_nondet_train.py b/flashinfer/cake_dense_gemm_nondet_train.py new file mode 100644 index 00000000000..a85c53124ec --- /dev/null +++ b/flashinfer/cake_dense_gemm_nondet_train.py @@ -0,0 +1,96 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Cake dense GEMMs for training, non-deterministic (SM100 / SM103 / SM107): the +GLM-5.2 BF16 dense projection GEMM (K1) and the FP32 router-gate GEMM through +split-BF16x3 emulation (K2), with programs whose split-K reductions add the +partials in the order they finish. The results are not bitwise reproducible +across runs; the operations, view contract, dtypes and accuracy budgets are +those of ``flashinfer.cake_dense_gemm_det_train``. The module selects the mode: +no entry point takes a mode argument. + +Under ``torch.use_deterministic_algorithms(True)`` a launch whose result can vary +run to run (the K1 split-K program at three or more splits, a K2 stream-K +partition that gives an element more than two partials) raises PyTorch's alert +right before it launches: ``RuntimeError``, or with ``warn_only=True`` a +``UserWarning`` and then the launch. Fixed-order launches do not alert, and the +setting never changes which program runs. + +The backend lives in ``flashinfer.experimental.cake.cake_dense_gemm_nondet_train`` +(``cake_backend.py``, the generated sources and ``build.json``). This module +imports it on the first attribute access and returns the backend's own objects, +so a call adds no Python layer. The convenience caches are bounded by the +deterministic backend's ``cake_backend.CONVENIENCE_CACHE_SIZE`` (one setting for +both packages), not by a name of this module. + +Usage:: + + from flashinfer.cake_dense_gemm_nondet_train import ( + cake_dense_gemm_nondet_train_projection_wgrad, + cake_dense_gemm_nondet_train_router_wgrad, + ) + + dW = cake_dense_gemm_nondet_train_projection_wgrad(G, X) # dW[N, K] = G[T, N].T @ X[T, K] + dWr = cake_dense_gemm_nondet_train_router_wgrad(Gr, Xr) # FP32 dW[E, K] = G[T, E].T @ X[T, K] +""" + +from __future__ import annotations + +import importlib +from typing import Any + +_BACKEND = "flashinfer.experimental.cake.cake_dense_gemm_nondet_train" + +# public name -> backend submodule that defines it +_EXPORTS: dict[str, str] = { + "ARCHES": "cake_jit", + "is_cake_dense_gemm_nondet_train_supported": "cake_backend", + # K1: BF16 dense projection GEMM + "cake_dense_gemm_nondet_train_projection": "cake_backend", + "cake_dense_gemm_nondet_train_projection_forward": "cake_backend", + "cake_dense_gemm_nondet_train_projection_dgrad": "cake_backend", + "cake_dense_gemm_nondet_train_projection_wgrad": "cake_backend", + "cake_dense_gemm_nondet_train_projection_plan": "cake_backend", + "cake_dense_gemm_nondet_train_projection_prepare": "cake_backend", + "cake_dense_gemm_nondet_train_projection_wgrad_prepare": "cake_backend", + "CakeDenseGemmNondetProjectionPlan": "cake_backend", + "CakeDenseGemmNondetProjectionLaunch": "cake_backend", + # K2: FP32 router GEMM + "cake_dense_gemm_nondet_train_router": "cake_backend", + "cake_dense_gemm_nondet_train_router_forward": "cake_backend", + "cake_dense_gemm_nondet_train_router_dgrad": "cake_backend", + "cake_dense_gemm_nondet_train_router_wgrad": "cake_backend", + "cake_dense_gemm_nondet_train_router_plan": "cake_backend", + "cake_dense_gemm_nondet_train_router_prepare": "cake_backend", + "CakeDenseGemmNondetRouterPlan": "cake_backend", + "CakeDenseGemmNondetRouterLaunch": "cake_backend", + # the convenience path's cached plans and workspace arenas (of both packages) + "cake_dense_gemm_nondet_train_convenience_cache_clear": "cake_backend", +} + +__all__ = sorted(_EXPORTS) + + +def __getattr__(name: str) -> Any: + submodule = _EXPORTS.get(name) + if submodule is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(f"{_BACKEND}.{submodule}"), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py index 46994a228c1..9e7d9ec8f62 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py @@ -14,16 +14,15 @@ limitations under the License. """ -# Experimental Cake backend, non-deterministic: the GLM-5.2 dense projection -# GEMMs for training and the FP32 router-gate GEMM on SM100 / SM103 / SM107 +# Cake backend, non-deterministic: the GLM-5.2 dense projection GEMMs for training +# and the FP32 router-gate GEMM on SM100 / SM103 / SM107 # (flashinfer-ai/flashinfer#5677), with programs that add split-K partials in -# the order they finish. Every compute entry point of ``cake_backend`` takes -# the keyword-only ``deterministic: bool = True`` and accepts only -# ``deterministic=False``: any other value, the default included, raises -# ``ValueError`` naming the entry point of the deterministic package -# ``flashinfer.experimental.cake.cake_dense_gemm_det_train``. Importing this package -# imports nothing: ``cake_backend``, ``cake_jit`` and ``cake_build`` load on -# first attribute access. +# the order they finish. The public module is +# flashinfer.cake_dense_gemm_nondet_train; it selects the mode, so no entry +# point takes a mode argument. The rows that run fixed-order programs reuse the +# programs of the deterministic sibling package cake_dense_gemm_det_train. +# Importing this package imports nothing: ``cake_backend``, ``cake_jit`` and +# ``cake_build`` load on first attribute access. __all__ = ["cake_backend", "cake_build", "cake_jit"] diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py index cf5f194827b..48e64da3df1 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py @@ -13,92 +13,88 @@ See the License for the specific language governing permissions and limitations under the License. -Cake backend, non-deterministic: the GLM-5.2 dense projection GEMMs for training (K1) and the FP32 router-gate GEMM -(K2) on SM100 / SM103 / SM107. +Cake backend of ``flashinfer.cake_dense_gemm_nondet_train``, non-deterministic: the GLM-5.2 dense projection GEMMs +for training (K1) and the FP32 router-gate GEMM (K2) on SM100 / SM103 / SM107. The deterministic implementation is the package ``cake_dense_gemm_det_train`` (flashinfer-ai/flashinfer#5677 and -#7043): its stream-K and split-K reductions add the partials in a fixed order, so every result is bitwise -reproducible run to run. This package keeps that package's operations, view contract, dtypes, accuracy budgets and -host paths (prepared and convenience) and adds programs whose split-K reductions add the partials in the order they -finish. The products, the FP32 accumulation, the rounding points (an output is rounded once, from the complete FP32 -sum), the masks, the layouts and the dynamic shapes are the deterministic package's, and no reduction flushes a -subnormal. A result of this package carries no run-to-run bitwise guarantee on any row, also on a row that runs a -fixed-order program today. +#7043; public module ``flashinfer.cake_dense_gemm_det_train``): its stream-K and split-K reductions add the partials +in a fixed order, so every result is bitwise reproducible run to run. This package keeps that package's operations, +view contract, dtypes, accuracy budgets and host paths (prepared and convenience) and adds programs whose split-K +reductions add the partials in the order they finish. The products, the FP32 accumulation, the rounding points (an +output is rounded once, from the complete FP32 sum), the masks, the layouts and the dynamic shapes are the +deterministic package's, and no reduction flushes a subnormal. A result of this package carries no run-to-run +bitwise guarantee on any row, also on a row that runs a fixed-order program today. Mode and dispatch rules ----------------------- -Every compute entry point takes the keyword-only argument ``deterministic: bool = True`` and accepts only -``deterministic=False``: - -* Any other value -- the default, ``True``, ``None``, ``0``, ``1``, ``"False"`` -- raises ``ValueError`` before - anything is planned, allocated or launched. The message names the deterministic package's entry point of the - operation (``projection_wgrad`` -> ``cake_dense_gemm_det_train_projection_wgrad``, ...). The keyword never selects the deterministic - package, and nothing falls back to it. -* ``deterministic=False`` runs the routes of :data:`ROUTES`: - - - the K1 weight gradient (both operands MN-major, ``K = T`` ragged), launched as the transposed-store GEMM - ``out[l, n, m]``: each of ``L x m_tiles x n_tiles`` tiles is cut into ``num_splits`` K ranges; every split stores - its FP32 partial into its own workspace slot and takes a ticket, and the split that draws the last ticket adds - the other ``num_splits - 1`` partials with IEEE round-to-nearest adds, rounds once and stores (:func:`k1_plan` - reads the tile width and the split count from the measured verdicts of :data:`K1_SPLIT_RULES`; every other row - runs the deterministic program). Two splits add ``own + other``, one add of two values, so a two-split row - repeats bit for bit; three or more add in the order the splits finish; - - the K1 weight gradient of the FP32 rows of :data:`K1_DSM_RANGES_BY_SM` (thin ``d_out`` on 152-SM ``sm_100a`` - parts, in that table's token ranges), before every other K1 table: a cluster program - (:func:`k1_dsm_plan`) whose ``S`` CTAs each take one of a tile's ``S`` K ranges and meet in the cluster's - distributed shared memory -- each CTA sends the column quads another rank owns to that rank, adds the ``S`` - partials of its own quads in cluster-rank order and stores them once (no workspace, no ticket; the same bits on - every run); - - the K2 weight gradient (and a forward given an explicit split count): a stream-K partition of the (tile, K - block) steps into ``units`` spans, alike or a head/tail partition of long head spans and short tail spans - (:func:`router_partition` reads the measured verdicts and token windows of :data:`ROUTER_RULES`; every other row - and token count runs the deterministic program); a segment that covers its tile's whole K range stores the - deterministic program's one-split bits, any other adds its partial into ``out`` with a TMA reduce-add (``out`` is - zeroed on the launch's stream first). A forward partition gives every element at most two partials, - ``0 + a + b == 0 + b + a``, so a forward recomputed under activation checkpointing gives the same logits and the - same top-k; - - the K2 weight and input gradients in the SIMT token windows of their rows in the router tables - (:func:`router_rules`, :func:`router_simt`): an exact-FP32 program on the CUDA cores over the FP32 operands read - in place (compensated in the 148-SM tables' one-token window at hidden 4096), every element of ``out`` stored - once (other bits than the deterministic program's BF16x3 emulation, within the deterministic package's router - FP64 budget); the input gradient's every other row and token count runs the deterministic program (one split); - - the K2 forward in the SIMT token windows of its rows (:func:`router_forward_simt`): an exact-FP32 program on the - CUDA cores whose cluster of CTAs splits K and adds the ranks' partials in rank order, every element of ``out`` - stored once (other bits than the deterministic program's, within the deterministic package's router FP64 - budget); every other forward row and token count runs the deterministic program. - - Every other operation -- the K1 forward and input gradient, the BF16x3 split -- runs the deterministic programs: a - recomputed forward must reproduce its result, and where the deterministic plan of a K1 forward or input gradient - splits K across CTAs (its automatic two-way split), it adds the partials in a fixed order. The K1 forward rows of - :data:`K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM ``sm_100a`` / ``sm_103a`` parts) run another plan of - those programs: the single-CTA program, every tile's whole K in one CTA. -* Under ``torch.use_deterministic_algorithms(True)`` a launch whose result can vary run to run -- the K1 split-K - program of CTA pairs at three or more splits, a K2 stream-K partition that gives an element more than two partials -- raises - PyTorch's alert right before it launches (:func:`_alert_order_varying`): ``RuntimeError``, or with - ``warn_only=True`` a ``UserWarning``, after which the launch runs. The fixed-order launches (every other route) - do not alert. The PyTorch setting never changes which program a call runs. - -The helpers of this package (the planners, the rule tables, :func:`router_symbol`, :func:`bind_launch`, ...) take no -mode keyword. +The module selects the mode: no entry point takes a mode argument. Every entry point runs the routes of +:data:`ROUTES`: + +* the K1 weight gradient (both operands MN-major, ``K = T`` ragged), launched as the transposed-store GEMM + ``out[l, n, m]``: each of ``L x m_tiles x n_tiles`` tiles is cut into ``num_splits`` K ranges; every split stores + its FP32 partial into its own workspace slot and takes a ticket, and the split that draws the last ticket adds the + other ``num_splits - 1`` partials with IEEE round-to-nearest adds, rounds once and stores (:func:`k1_plan` reads the + tile width and the split count from the measured verdicts of :data:`K1_SPLIT_RULES`; every other row runs the + deterministic program). Two splits add ``own + other``, one add of two values, so a two-split row repeats bit for + bit; three or more add in the order the splits finish; +* the K1 weight gradient of the FP32 rows of :data:`K1_DSM_RANGES_BY_SM` (thin ``d_out`` on 152-SM ``sm_100a`` parts, + in that table's token ranges), before every other K1 table: a cluster program (:func:`k1_dsm_plan`) whose ``S`` + CTAs each take one of a tile's ``S`` K ranges and meet in the cluster's distributed shared memory -- each CTA sends + the column quads another rank owns to that rank, adds the ``S`` partials of its own quads in cluster-rank order and + stores them once (no workspace, no ticket; the same bits on every run); +* the K2 weight gradient (and a forward given an explicit split count): a stream-K partition of the (tile, K block) + steps into ``units`` spans, alike or a head/tail partition of long head spans and short tail spans + (:func:`router_partition` reads the measured verdicts and token windows of :data:`ROUTER_RULES`; every other row and + token count runs the deterministic program); a segment that covers its tile's whole K range stores the + deterministic program's one-split bits, any other adds its partial into ``out`` with a TMA reduce-add (``out`` is + zeroed on the launch's stream first). A forward partition gives every element at most two partials, + ``0 + a + b == 0 + b + a``, so a forward recomputed under activation checkpointing gives the same logits and the same + top-k; +* the K2 weight and input gradients in the SIMT token windows of their rows in the router tables + (:func:`router_rules`, :func:`router_simt`): an exact-FP32 program on the CUDA cores over the FP32 operands read in + place (compensated in the 148-SM tables' one-token window at hidden 4096), every element of ``out`` stored once + (other bits than the deterministic program's BF16x3 emulation, within the deterministic package's router FP64 + budget); the input gradient's every other row and token count runs the deterministic program (one split); +* the K2 forward in the SIMT token windows of its rows (:func:`router_forward_simt`): an exact-FP32 program on the CUDA + cores whose cluster of CTAs splits K and adds the ranks' partials in rank order, every element of ``out`` stored + once (other bits than the deterministic program's, within the deterministic package's router FP64 budget); every + other forward row and token count runs the deterministic program. + +Every other operation -- the K1 forward and input gradient, the BF16x3 split -- runs the deterministic programs: a +recomputed forward must reproduce its result, and where the deterministic plan of a K1 forward or input gradient +splits K across CTAs (its automatic two-way split), it adds the partials in a fixed order. The K1 forward rows of +:data:`K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM ``sm_100a`` / ``sm_103a`` parts) run another plan of those +programs: the single-CTA program, every tile's whole K in one CTA. + +Under ``torch.use_deterministic_algorithms(True)`` a launch whose result can vary run to run -- the K1 split-K program +of CTA pairs at three or more splits, a K2 stream-K partition that gives an element more than two partials -- raises +PyTorch's alert right before it launches (:func:`_alert_order_varying`): ``RuntimeError`` (nothing launched, ``out`` +untouched), or with ``warn_only=True`` a ``UserWarning``, after which the launch runs. The fixed-order launches (the +one- and two-split K1 rows, the K1 cluster program, the K2 SIMT programs, a K2 partition of at most two partials per +element, every deterministic program) do not alert. The alert is checked on every launch, never at plan or prepare +time; a launch captured into a CUDA graph is checked at capture, not at replay. The PyTorch setting never changes +which program a call runs, and no call goes to the deterministic package because of it. The two packages ---------------- This package registers only the programs of its own schedule (``cake_jit``: programs ``cake_dense_projection_gemm_``, templates ``dense_proj_gemm_atomic_*`` / ``router_fp32_gemm_*``); every -deterministic program, and every planner rule of the deterministic package, is that package's. A -``deterministic=False`` call reuses them on the rows that run the deterministic programs (above): it launches the +deterministic program, and every planner rule of the deterministic package, is that package's. A call of this +package reuses them on the rows that run the deterministic programs (above): it launches the deterministic plan and program through a binding of this package (the same bits, fewer host operations per call), or, where that binding cannot rebuild the call's launch views from its tensors (two operands that alias) and for a batched forward that the deterministic package runs on its swapped small-token program, it calls the deterministic -package's entry point of the same operation (the same program, the same bits). On such a row ``plan_*`` / -``prepare_*`` return the deterministic package's plan or prepared launch. A row that runs a program of this package +package's entry point of the same operation (the same program, the same bits). On such a row the ``*_plan`` / +``*_prepare`` entry points return the deterministic package's plan or prepared launch +(``CakeDenseGemmDetProjectionPlan``, ``CakeDenseGemmDetProjectionLaunch``, ``CakeDenseGemmDetRouterPlan``, +``CakeDenseGemmDetRouterLaunch``). A row that runs a program of this package never goes to the deterministic package. The deterministic package is imported at run time as this package's sibling, never through an absolute ``flashinfer`` import; a vendored copy binds it with :func:`use_deterministic_backend`. The deterministic package never imports this one, and this module re-exports no -name of it. :func:`clear_convenience_caches` clears the convenience caches and workspace arenas of both packages; a +name of it. :func:`cake_dense_gemm_nondet_train_convenience_cache_clear` clears the convenience caches and workspace arenas of both packages; a convenience-path call that this package launches from its arenas, captured into a CUDA graph, stays valid until then (the arenas keep every buffer a launch was bound to, through growths and cache evictions); a call it hands to the -deterministic package's entry point follows that package's rule (:func:`clear_convenience_caches`). +deterministic package's entry point follows that package's rule (:func:`cake_dense_gemm_nondet_train_convenience_cache_clear`). The tensors passed between the forward and the gradients carry no mode: recompute a checkpointed forward with the package that made the saved one (same package only). @@ -160,7 +156,7 @@ def use_deterministic_backend(module: ModuleType) -> None: under another name, without the ``flashinfer`` package around it) vendors the deterministic package too and binds its ``cake_backend`` here before the first call that needs it. Binding drops this package's convenience caches and arenas (:func:`_drop_convenience_entries`): their entries hold the plans and programs of the package - bound before. Like :func:`clear_convenience_caches`, it invalidates the CUDA graphs that captured a + bound before. Like :func:`cake_dense_gemm_nondet_train_convenience_cache_clear`, it invalidates the CUDA graphs that captured a convenience-path call of this package.""" missing = [name for name in _DET_ENTRIES if not hasattr(module, name)] if missing: @@ -191,32 +187,9 @@ def deterministic_backend() -> ModuleType: return module -def _det_package_name() -> str: - """The dotted name of the deterministic package (the bound one, else this package's sibling), as text: the mode - messages name it without importing it.""" - package = cake_jit._det_package - if package is not None: - return package - parent = __name__.rpartition(".")[0].rpartition(".")[0] - return ( - f"{parent}.{cake_jit.DETERMINISTIC_PACKAGE}" - if parent - else cake_jit.DETERMINISTIC_PACKAGE - ) - - -def _require_false(deterministic: Any, entry: str, det_entry: str) -> None: - """The mode rule of every compute entry point, its first statement: this package computes non-deterministically, - so ``entry`` takes only ``deterministic=False``. Any other value (the default ``True``, ``None``, ``0``, ``1``, - ``"False"``, a NumPy bool) raises ``ValueError`` naming the deterministic package's entry point ``det_entry``; - the keyword never selects the deterministic package.""" - if deterministic is not False: - package = __name__.rpartition(".")[0] or __name__ - raise ValueError( - f"{entry}: {package} runs the non-deterministic package only and needs deterministic=False (got " - f"{deterministic!r}): pass deterministic=False for the non-deterministic path; the deterministic entry " - f"is {_det_package_name()}.cake_backend.{det_entry}" - ) +#: The public modules of the two packages: the texts of the alert and of the errors name them. +_MODULE = "flashinfer.cake_dense_gemm_nondet_train" +_DET_MODULE = "flashinfer.cake_dense_gemm_det_train" def _warning_stacklevel() -> int: @@ -237,27 +210,27 @@ def _alert_order_varying(launch: str) -> None: are enabled: a ``RuntimeError``, or a ``UserWarning`` under ``torch.use_deterministic_algorithms(True, warn_only=True)``, after which the launch runs; nothing otherwise. The alert of PyTorch's own operators (PyTorch exposes no public Python entry for it). Fixed-order launches do not call it, and it never changes the program - that runs.""" + that runs. A launch captured into a CUDA graph calls it at capture; a replay does not.""" if not torch.are_deterministic_algorithms_enabled(): return if torch.is_deterministic_algorithms_warn_only_enabled(): warnings.warn( - f"{__name__}: {launch} does not have a deterministic implementation, but you set " + f"{_MODULE}: {launch} does not have a deterministic implementation, but you set " "'torch.use_deterministic_algorithms(True, warn_only=True)'; " - f"{_det_package_name()} is the deterministic package", + f"{_DET_MODULE} is the deterministic module", UserWarning, stacklevel=_warning_stacklevel(), ) return raise RuntimeError( - f"{__name__}: {launch} does not have a deterministic implementation, but you set " - f"'torch.use_deterministic_algorithms(True)'; {_det_package_name()} is the deterministic package, or use " + f"{_MODULE}: {launch} does not have a deterministic implementation, but you set " + f"'torch.use_deterministic_algorithms(True)'; {_DET_MODULE} is the deterministic module, or use " "the 'warn_only=True' option if that is acceptable for your application" ) def _k1_launch_text(plan: Any) -> str: - """The launch an order-varying K1 plan (:class:`AtomicGemmPlan`) names in :func:`_alert_order_varying`.""" + """The launch an order-varying K1 plan (:class:`CakeDenseGemmNondetProjectionPlan`) names in :func:`_alert_order_varying`.""" return ( f"the K1 split-K program {plan.template} at {plan.splits} splits (the partials added in the order the " "splits finish)" @@ -265,7 +238,7 @@ def _k1_launch_text(plan: Any) -> str: def _k2_launch_text(plan: Any) -> str: - """The launch an order-varying K2 plan (:class:`AtomicRouterPlan`) names in :func:`_alert_order_varying`.""" + """The launch an order-varying K2 plan (:class:`CakeDenseGemmNondetRouterPlan`) names in :func:`_alert_order_varying`.""" return ( f"the router stream-K program {plan.template} with {plan.units} units over {plan.tiles} tiles (more than " "two partials per element, added in the order they finish)" @@ -273,14 +246,13 @@ def _k2_launch_text(plan: Any) -> str: # --------------------------------------------------------------------------- -# Routing of the deterministic=False mode +# Routing # --------------------------------------------------------------------------- -#: The program family each operation runs with ``deterministic=False``: ``"atomic"`` -- this package's planner and +#: The program family each operation runs: ``"atomic"`` -- this package's planner and #: programs (the planner keeps every row without a measured verdict on the deterministic program, see #: :func:`k1_plan` / :func:`router_partition`) -- or ``"det"`` -- the deterministic package's programs (fixed order). -#: The mode keyword never reads this table: it only takes ``deterministic=False`` (:func:`_require_false`). K1 -#: operations are keyed by the operand layouts of the launched GEMM (``kk`` forward, ``kn`` input gradient, ``nn`` +#: The PyTorch deterministic-algorithms setting never reads this table. K1 operations are keyed by the operand layouts of the launched GEMM (``kk`` forward, ``kn`` input gradient, ``nn`` #: weight gradient; ``nk`` has no training use), K2 by the layout class of ``router_layout_class``. ROUTES: dict[str, str] = { "k1_kk": "det", # forward: feeds the recompute of checkpointing; whole tiles or a fixed-order two-way K split (the @@ -341,7 +313,7 @@ def _ceil_div(a: int, b: int) -> int: K1_COUNTER_WORDS = 512 # ticket words per K1 tile: 16 epilogue warp slices x 32 lanes K1_BLOCK_N_CHOICES = (64, 128, 192, 256) #: The tile widths of :data:`K1_BLOCK_N_CHOICES` this package runs as its row-major-store program, with the operands -#: exchanged (:attr:`AtomicGemmPlan.out_t` false): a CTA pair's tile is 256 rows of the transposed-store launch's ``N`` +#: exchanged (:attr:`CakeDenseGemmNondetProjectionPlan.out_t` false): a CTA pair's tile is 256 rows of the transposed-store launch's ``N`` #: by ``block_n`` columns of its ``M``. The other widths run as transposed-store programs. K1_ROW_MAJOR_BLOCK_N: tuple[int, ...] = (192,) #: Accumulation chains of a one-split K1 program: with ``c`` chains each tile's K blocks are cut into ``c`` contiguous @@ -1179,9 +1151,9 @@ def k1_block_n(N: int) -> int: "block_n": 128, "hints": ("none", "none"), } -#: The K1 forward rows of ``deterministic=False`` that run another plan of the deterministic package's programs than +#: The K1 forward rows of this package that run another plan of the deterministic package's programs than #: its own plan of the row, ``(arch, sm_count) -> {(K, N, T): knobs}``: on a part of exactly ``sm_count`` SMs, -#: :func:`projection_forward` of ``X[T, K]`` and ``W[N, K]`` (both K-major) into a BF16 ``out`` of exactly these ``K``, +#: :func:`cake_dense_gemm_nondet_train_projection_forward` of ``X[T, K]`` and ``W[N, K]`` (both K-major) into a BF16 ``out`` of exactly these ``K``, #: ``N`` and ``T`` runs the program the deterministic planner gives the row with these knobs (:func:`k1_forward_knobs`); #: every other part, row, token count, layout and output dtype runs the deterministic plan. On 152-SM ``sm_100a`` and #: ``sm_103a`` parts the deterministic plan of these rows splits K two ways over CTA pairs of 256-wide tiles (stream-K, @@ -1214,7 +1186,7 @@ def k1_block_n(N: int) -> int: def k1_forward_knobs(*, arch: str, sm_count: int, K: int, N: int, T: int) -> tuple: """The deterministic package's planner knobs of the forward ``out[T, N] = X[T, K] @ W[N, K].T`` (both operands - K-major, BF16 ``out``) with ``deterministic=False`` on a device of exactly ``sm_count`` SMs: the row's entry in + K-major, BF16 ``out``) of this package on a device of exactly ``sm_count`` SMs: the row's entry in the table of :data:`K1_FORWARD_RULES_BY_SM` for ``(arch, sm_count)`` at exactly ``T`` tokens, as sorted ``(name, value)`` pairs, or ``()`` -- the deterministic plan.""" knobs = K1_FORWARD_RULES_BY_SM.get((arch, int(sm_count)), {}).get( @@ -1265,7 +1237,7 @@ def dsm_k1_template(block_n: int, splits: int) -> str: @dataclass(frozen=True) -class AtomicGemmPlan: +class CakeDenseGemmNondetProjectionPlan: """The host plan of one K1 launch of this package: the transposed-store GEMM ``out[l, n, m] = sum_k A[l, k, m] B[l, k, n]`` (``out_t``), or the row-major-store GEMM ``out[l, m, n]`` (``out_t`` false: the widths of :data:`K1_ROW_MAJOR_BLOCK_N`), over the launch's ``[L, K, M]`` / ``[L, K, N]`` operand views, ``L x m_tiles x @@ -1364,7 +1336,7 @@ def plan_atomic_gemm( splits: Optional[int] = None, block_n: Optional[int] = None, chains: Optional[int] = None, -) -> tuple[Optional[AtomicGemmPlan], Any, Any, Any, str]: +) -> tuple[Optional[CakeDenseGemmNondetProjectionPlan], Any, Any, Any, str]: """Plan this package's K1 GEMM ``out = A @ B`` (``A [L, M, K]`` with unit stride on ``m``, ``B [L, K, N]`` with unit stride on ``n``; ``out [L, M, N]``, or ``[L, N, M]`` with ``transposed_out``). @@ -1393,12 +1365,12 @@ def plan_atomic_gemm( Lb, Kb, N = (int(s) for s in B3.shape) if Lb != L or Kb != K: raise ValueError( - f"dense_projection_gemm: A is [L={L}, M={M}, K={K}] but B is [{Lb}, {Kb}, {N}]" + f"cake_dense_gemm_nondet_train_projection: A is [L={L}, M={M}, K={K}] but B is [{Lb}, {Kb}, {N}]" ) want = (L, N, M) if transposed_out else (L, M, N) if tuple(int(s) for s in O3.shape) != want: raise ValueError( - f"dense_projection_gemm: out must be {list(want)}, got {list(O3.shape)}" + f"cake_dense_gemm_nondet_train_projection: out must be {list(want)}, got {list(O3.shape)}" ) none = (None, None, None, None) if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: @@ -1440,7 +1412,7 @@ def plan_atomic_gemm( template = dsm_k1_template(bn, p) if template not in cake_jit.KERNELS.get(arch, {}): return (*none, f"{template} is not registered for {arch}") - plan = AtomicGemmPlan( + plan = CakeDenseGemmNondetProjectionPlan( L=L, M=M_, N=N_, @@ -1483,14 +1455,14 @@ def plan_atomic_gemm( raise ValueError(f"chains is one of {K1_CHAINS_CHOICES}, got {chains}") if c > 1 and (p != 1 or not k1_chains_fit(c, K)): raise ValueError( - f"dense_projection_gemm: {c} chains run at one split (got {p}) with a 64-element K block each " + f"cake_dense_gemm_nondet_train_projection: {c} chains run at one split (got {p}) with a 64-element K block each " f"(K={K})" ) if bn not in K1_BLOCK_N_CHOICES: raise ValueError(f"block_n is one of {K1_BLOCK_N_CHOICES}, got {block_n}") if not k1_splits_fit(p, K) or p > 65536: raise ValueError( - f"dense_projection_gemm: splits must be in [1, {_ceil_div(K, K1_BLOCK_K)}] (the 64-element K blocks " + f"cake_dense_gemm_nondet_train_projection: splits must be in [1, {_ceil_div(K, K1_BLOCK_K)}] (the 64-element K blocks " f"of K={K}), got {p}" ) template = atomic_k1_template(bn, out_f32, c) @@ -1500,7 +1472,7 @@ def plan_atomic_gemm( if not out_t: # the row-major-store launch: the transposed-store launch with its operands and extents exchanged a_launch, b_launch, M_, N_ = b_launch, a_launch, N_, M_ - plan = AtomicGemmPlan( + plan = CakeDenseGemmNondetProjectionPlan( L=L, M=M_, N=N_, @@ -1520,7 +1492,10 @@ def plan_atomic_gemm( def _k1_launch_views( - A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, plan: AtomicGemmPlan + A: torch.Tensor, + B: torch.Tensor, + out: torch.Tensor, + plan: CakeDenseGemmNondetProjectionPlan, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The launch's ``[L, K, M]`` / ``[L, K, N]`` operand views and the batched output of a planned K1 launch, rebuilt from the caller's tensors (the geometry is the plan's, so its checks hold).""" @@ -1532,7 +1507,7 @@ def _k1_launch_views( def _k1_values( - plan: AtomicGemmPlan, + plan: CakeDenseGemmNondetProjectionPlan, *, A: Any, B: Any, @@ -1583,7 +1558,7 @@ class _VerdictTable(dict): it plans under them (:data:`_K2_ATOMIC_CACHE`, :data:`_K2_ENTRY_CACHE`) until one changes; a table bound in their place that is not one is read on every call. The other planner tables (the K1 tables, :data:`ROUTER_SM_COUNT`, the head/tail templates, the deterministic package's) are read when a key is planned: a cached key keeps its plan - until :func:`clear_convenience_caches`. A copy (``dict(table)``) is a plain ``dict``.""" + until :func:`cake_dense_gemm_nondet_train_convenience_cache_clear`. A copy (``dict(table)``) is a plain ``dict``.""" __slots__ = () @@ -2244,13 +2219,13 @@ def router_partition( s = int(splits) if not 1 <= s <= k_iters: raise ValueError( - f"router_fp32_gemm: splits must be in [1, {k_iters}] (the 32-element K blocks of K={K}), got {s}" + f"cake_dense_gemm_nondet_train_router: splits must be in [1, {k_iters}] (the 32-element K blocks of K={K}), got {s}" ) if layout == "kk" and s > ROUTER_FORWARD_MAX_SPLITS: raise ValueError( - f"router_fp32_gemm(deterministic=False): the forward takes at most {ROUTER_FORWARD_MAX_SPLITS} " - f"splits (more partials per element do not repeat bit for bit under recompute), got {s}; the " - f"deterministic entry {_det_package_name()}.cake_backend.cake_dense_gemm_det_train_router takes more" + f"cake_dense_gemm_nondet_train_router: the forward takes at most {ROUTER_FORWARD_MAX_SPLITS} " + f"splits (more partials per element do not repeat bit for bit under recompute), got {s}; " + f"{_DET_MODULE}.cake_dense_gemm_det_train_router takes more" ) spec = f"S{s}" else: @@ -2263,7 +2238,7 @@ def router_partition( return None if layout == "kk" and router_head_tail(spec) is not None: raise ValueError( - f"router_fp32_gemm(deterministic=False): the head/tail partition {spec!r} is the weight gradient's; " + f"cake_dense_gemm_nondet_train_router: the head/tail partition {spec!r} is the weight gradient's; " "the forward takes at most two partials per element" ) units, head_spans, head_steps = router_partition_of( @@ -2271,7 +2246,7 @@ def router_partition( ) if layout == "kk" and not router_at_most_two_partials(m_pairs, n_tiles, units): raise ValueError( - f"router_fp32_gemm(deterministic=False): the forward partition {spec!r} ({units} units over " + f"cake_dense_gemm_nondet_train_router: the forward partition {spec!r} ({units} units over " f"{m_pairs * n_tiles} tiles) gives an element more than two partials" ) if not router_partition_fits( @@ -2365,7 +2340,7 @@ def router_forward_simt( @dataclass(frozen=True) -class AtomicRouterPlan: +class CakeDenseGemmNondetRouterPlan: """The host plan of one K2 launch of this package (the launched, possibly swapped GEMM ``[M, N] = [M, K] @ [K, N]``): ``m_tiles / 2 x n_tiles`` tiles of ``256 x 128``, their ``k_iters`` K blocks cut into ``units`` stream-K spans (``2 * units`` CTAs; the first ``head_spans`` spans share the first ``head_steps`` steps on the @@ -2465,7 +2440,7 @@ def plan_atomic_router( arch: str, sm_count: int, splits: Optional[int] = None, -) -> tuple[Optional[AtomicRouterPlan], str]: +) -> tuple[Optional[CakeDenseGemmNondetRouterPlan], str]: """Plan this package's K2 GEMM ``out[M, N] = A[M, K] @ B[K, N]`` (fp32 views, the deterministic package's view contract and swap rule) of the ``kk`` / ``nn_t`` classes, and of the ``kn`` class in its SIMT windows. Returns ``(plan, reason)``: ``plan`` is ``None`` when the launch runs the deterministic program, and ``reason`` says why -- @@ -2483,18 +2458,22 @@ def plan_atomic_router( or B.dtype != torch.float32 or out.dtype != torch.float32 ): - raise ValueError("router_fp32_gemm: A, B and out must be fp32") + raise ValueError( + "cake_dense_gemm_nondet_train_router: A, B and out must be fp32" + ) if A.dim() != 2 or B.dim() != 2 or out.dim() != 2: - raise ValueError("router_fp32_gemm: A, B and out must be 2-D views") + raise ValueError( + "cake_dense_gemm_nondet_train_router: A, B and out must be 2-D views" + ) M, K = (int(s) for s in A.shape) Kb, N = (int(s) for s in B.shape) if Kb != K or tuple(out.shape) != (M, N): raise ValueError( - f"router_fp32_gemm: A [{M}, {K}], B [{Kb}, {N}] and out {list(out.shape)} do not agree" + f"cake_dense_gemm_nondet_train_router: A [{M}, {K}], B [{Kb}, {N}] and out {list(out.shape)} do not agree" ) if min(M, K, N) < 1: raise ValueError( - f"router_fp32_gemm: M, K and N must be >= 1 (M={M}, N={N}, K={K})" + f"cake_dense_gemm_nondet_train_router: M, K and N must be >= 1 (M={M}, N={N}, K={K})" ) a_mn, b_mn = _router_a_mn(A), _router_b_mn(B) if a_mn is None or b_mn is None: @@ -2542,7 +2521,7 @@ def plan_atomic_router( return None, f"{template} is not registered for {arch}" _, _, experts, tokens = forward return ( - AtomicRouterPlan( + CakeDenseGemmNondetRouterPlan( M=M_eff, N=N_eff, K=K, @@ -2592,7 +2571,7 @@ def plan_atomic_router( return None, f"{template} is not registered for {arch}" tile_n, tile_m = int(simt[0]), int(simt[1]) return ( - AtomicRouterPlan( + CakeDenseGemmNondetRouterPlan( M=M_eff, N=N_eff, K=K, @@ -2630,7 +2609,7 @@ def plan_atomic_router( if template not in cake_jit.KERNELS.get(arch, {}): return None, f"{template} is not registered for {arch}" m_tiles, n_tiles, k_iters = _router_grid(M_eff, N_eff, K) - plan = AtomicRouterPlan( + plan = CakeDenseGemmNondetRouterPlan( M=M_eff, N=N_eff, K=K, @@ -2800,7 +2779,7 @@ def _k2_key( def _k2_views( - A: torch.Tensor, B: torch.Tensor, plan: AtomicRouterPlan + A: torch.Tensor, B: torch.Tensor, plan: CakeDenseGemmNondetRouterPlan ) -> tuple[torch.Tensor, torch.Tensor]: """``(a_view, split_src)``: the streamed operand as its ``[1, M, K]`` / ``[1, K, M]`` TMA span and the source of the BF16x3 split of the second operand (the deterministic package's convention), after the swap. The swapped @@ -2821,7 +2800,7 @@ def _k2_out32(out: torch.Tensor) -> torch.Tensor: def _k2_values( - plan: AtomicRouterPlan, + plan: CakeDenseGemmNondetRouterPlan, *, A: Any, parts: Any, @@ -2850,7 +2829,7 @@ def _k2_values( def _simt_values( - plan: AtomicRouterPlan, + plan: CakeDenseGemmNondetRouterPlan, *, A: Any, B: Any, @@ -3070,7 +3049,7 @@ def _k1_dummies(device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: @dataclass -class PreparedAtomicGemm: +class CakeDenseGemmNondetProjectionLaunch: """One prepared K1 launch of this package: the bound program and its workspace (the FP32 split slots, zeroed when prepared because the last split of a partial tile reads padding columns that no split writes and discards them; the tickets, zero when prepared and zero again after every complete launch). ``launch()`` writes ``out`` @@ -3080,7 +3059,7 @@ class PreparedAtomicGemm: module_name: str arch: str - plan: AtomicGemmPlan + plan: CakeDenseGemmNondetProjectionPlan out: torch.Tensor tensors: dict[str, torch.Tensor] = field(repr=False) device_index: int = 0 @@ -3121,7 +3100,7 @@ def launch(self) -> torch.Tensor: @dataclass -class PreparedAtomicRouterGemm: +class CakeDenseGemmNondetRouterLaunch: """One prepared K2 launch of this package: the retained BF16x3 stack of the second operand (re-split from the bound fp32 view on every ``launch()`` by the deterministic package's split stage), then ``out`` zeroed on the stream (one memset, :class:`_OutFill`) when a segment adds into it, then the stream-K program. Everything but @@ -3132,7 +3111,7 @@ class PreparedAtomicRouterGemm: module_name: str arch: str - plan: AtomicRouterPlan + plan: CakeDenseGemmNondetRouterPlan out: torch.Tensor tensors: dict[str, torch.Tensor] = field(repr=False) device_index: int = 0 @@ -3196,14 +3175,14 @@ def _split_knobs(knobs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]] def _prepare_atomic_k1( - plan: AtomicGemmPlan, + plan: CakeDenseGemmNondetProjectionPlan, a_desc, b_desc, out3, out: torch.Tensor, arch: str, enable_pdl: Optional[bool], -) -> PreparedAtomicGemm: +) -> CakeDenseGemmNondetProjectionLaunch: det = deterministic_backend() device = out.device if plan.workspace: @@ -3227,7 +3206,7 @@ def _prepare_atomic_k1( enable_pdl=int(det.pdl_enabled(enable_pdl, device)), ) launch = bind_launch(module_name, values, plan.grid, arch) - return PreparedAtomicGemm( + return CakeDenseGemmNondetProjectionLaunch( module_name=module_name, arch=arch, plan=plan, @@ -3274,7 +3253,7 @@ def _atomic_k1( convenience call ``det_args()`` through this package's binding (:func:`_launch_cached_atomic_k1`; ``det_knobs``: that call's knobs). ``enable_pdl`` (default: the device supports it) launches the program as a programmatic dependent launch; - ``weights_static`` is the deterministic call's (:func:`dense_projection_gemm`).""" + ``weights_static`` is the deterministic call's (:func:`cake_dense_gemm_nondet_train_projection`).""" if op == "launch": return _launch_cached_atomic_k1( A, @@ -3300,13 +3279,13 @@ def _atomic_k1( def _prepare_atomic_k2( - plan: AtomicRouterPlan, + plan: CakeDenseGemmNondetRouterPlan, A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, arch: str, enable_pdl: Optional[bool] = None, -) -> PreparedAtomicRouterGemm: +) -> CakeDenseGemmNondetRouterLaunch: det = deterministic_backend() device = out.device pdl = int(det.pdl_enabled(enable_pdl, device)) @@ -3316,7 +3295,7 @@ def _prepare_atomic_k2( values = _simt_values( plan, A=a_view, B=split_src, out=out, ldo=out.stride(0), enable_pdl=pdl ) - return PreparedAtomicRouterGemm( + return CakeDenseGemmNondetRouterLaunch( module_name=module_name, arch=arch, plan=plan, @@ -3349,7 +3328,7 @@ def _prepare_atomic_k2( ) ) index = det._device_index(device) - return PreparedAtomicRouterGemm( + return CakeDenseGemmNondetRouterLaunch( module_name=module_name, arch=arch, plan=plan, @@ -3409,7 +3388,7 @@ def _atomic_k2( # different streams never do. The K1 tickets are zero when allocated or grown (``torch.zeros`` on the arena's # stream) and zero again after every complete launch, so a launch needs no memset; the K1 split slots are zeroed when # allocated or grown (the last split reads the never-written padding columns of a partial tile and discards them) and -# need no memset per launch either. A K1 call of ``deterministic=False`` that runs the deterministic program -- the +# need no memset per launch either. A K1 call of this package that runs the deterministic program -- the # forward, the input gradient, a layout routed ``"det"`` (cache ``_K1_DET_CACHE``, keyed by the entry point's own # tensors), a weight-gradient row the planner keeps on it (cached as such in ``_K1_ATOMIC_CACHE``) -- runs it through a # binding of this package (:func:`_plan_det_k1`): the deterministic package's plan and program, so its bits, with this @@ -3421,7 +3400,7 @@ def _atomic_k2( # so a call rebuilds only its views (each the caller's tensor, its ``unsqueeze(0)`` or one strided view, checked # against the planner's views when bound; passed as torch tensors, each converted in the FFI call) and enters its # thread's stream context of the (device, stream), made once; the projection entry points go from their key straight -# to the bound launch when it is cached (:func:`_k1_cached_launch`). A K2 call of ``deterministic=False`` that runs the +# to the bound launch when it is cached (:func:`_k1_cached_launch`). A K2 call of this package that runs the # deterministic program -- a row the planner keeps on it, a call without a split count whose row has no verdict # (:func:`_k2_row_without_verdict`, kept per device and view geometry until a router table changes, # :func:`_router_memo`; not planned here), a forward without a split count while no router table gives the forward a @@ -3434,14 +3413,14 @@ def _atomic_k2( # program's partials -- carved and converted once per arena growth), so a call passes only the caller's views (an # entry point's own tensors when they are the launch views, :func:`_launch_views`; torch tensors, each converted in # its FFI call) and fills ``out`` with the bound geometry; the K2 launches of an arena enter its stream context (made -# once) under its lock. The router entry points' calls (``router_forward`` / ``router_dgrad`` / ``router_wgrad``) +# once) under its lock. The router entry points' calls (``cake_dense_gemm_nondet_train_router_forward`` / ``cake_dense_gemm_nondet_train_router_dgrad`` / ``cake_dense_gemm_nondet_train_router_wgrad``) # are cached once more by the entry point's own tensors (``_K2_ENTRY_CACHE``, :func:`_k2_entry`): each holds the # cached launch of ``_K2_SPLITK_CACHE`` / ``_K2_ATOMIC_CACHE`` that the entry point's path takes for the key and the # router-table memo and route it was routed under, so a call goes from its key to the bound launch. After a launch -# that did not complete, :func:`clear_convenience_caches` drops the arenas (fresh tickets on the next call). +# that did not complete, :func:`cake_dense_gemm_nondet_train_convenience_cache_clear` drops the arenas (fresh tickets on the next call). # # CUDA graphs. A cached launch captured into a CUDA graph replays the workspace pointers it was bound with, so no -# buffer of an arena is freed before :func:`clear_convenience_caches` (or :func:`use_deterministic_backend`) drops the +# buffer of an arena is freed before :func:`cake_dense_gemm_nondet_train_convenience_cache_clear` (or :func:`use_deterministic_backend`) drops the # arenas: a growth retires the buffer it replaces into the arena (:class:`_AtomicArena`) instead of releasing it, and # an entry evicted from the caches leaves its arena in ``_ATOMIC_ARENAS``. A captured graph therefore stays valid # until that call, which invalidates it. The calls this path hands to the deterministic package's convenience path @@ -3476,7 +3455,7 @@ class _AtomicArena: Unlike that arena, a growth retires the buffer it replaces (``retired``) instead of releasing it: a launch captured into a CUDA graph replays the pointer it was bound with, and a launch in flight keeps using it, so every buffer of - the arena stays allocated until :func:`clear_convenience_caches` drops the arena. A retired buffer keeps its zero + the arena stays allocated until :func:`cake_dense_gemm_nondet_train_convenience_cache_clear` drops the arena. A retired buffer keeps its zero contract: the launches bound to it leave its tickets and counters zero. A growth at least doubles the buffer, so the buffers a kind retires hold fewer bytes than its live one, and a kind grows at most ``ceil(log2(largest requirement / _ATOMIC_ARENA_MIN_ELEMS))`` times.""" @@ -3935,7 +3914,7 @@ def _k1_cached_launch( sources: tuple, enable_pdl: Optional[bool], ) -> bool: - """A projection entry point's call of ``deterministic=False`` when its key is cached and bound: the key the entry + """A projection entry point's call when its key is cached and bound: the key the entry point's own path builds for the tensors ``sources`` (``(A, B, out)`` of that path) in ``cache`` -- :data:`_K1_DET_CACHE` (:func:`_det_k1`, flags ``(*head, pdl, *tail)`` = ``(flags, pdl, weights_static)``) or :data:`_K1_ATOMIC_CACHE` (:func:`_launch_cached_atomic_k1`) -- and its bound launch, without the path's call @@ -3967,7 +3946,7 @@ def _k1_cached_launch( def _bind_atomic_k1( - plan: AtomicGemmPlan, + plan: CakeDenseGemmNondetProjectionPlan, module_name: str, arg_plan: list, A: torch.Tensor, @@ -4288,7 +4267,7 @@ def _launch_atomic_k1_unbound( def _k2_binding( - plan: AtomicRouterPlan, + plan: CakeDenseGemmNondetRouterPlan, out: torch.Tensor, split_src: torch.Tensor, module_name: str, @@ -4612,7 +4591,9 @@ def _plan_splitk_k2( det = deterministic_backend() device = out.device arch = det.require_arch(device) - plan, A_eff, B_eff = det.cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits) + plan, A_eff, B_eff = det.cake_dense_gemm_det_train_router_plan( + A, B, out, splits=splits + ) a_view, split_src = det._k2_checked_views(A, B, A_eff, B_eff, plan) mine = _k2_views(A, B, plan) if not (det._same_view(mine[0], a_view) and det._same_view(mine[1], split_src)): @@ -4710,7 +4691,9 @@ def _launch_cached_splitk_k2( _plan_splitk_k2(A, B, out, splits, index, stream, pdl), ) if entry.k2 is None or entry.entry is None or entry.arena is None: - return det.cake_dense_gemm_det_train_router(A, B, out, splits=splits, enable_pdl=bool(pdl)) + return det.cake_dense_gemm_det_train_router( + A, B, out, splits=splits, enable_pdl=bool(pdl) + ) a_view, split_src = _launch_views(A, B, entry.plan, views) _launch_splitk_k2_binding(det, entry, out, a_view, split_src, index, stream) return out @@ -4773,8 +4756,8 @@ def _splitk_k2( ) -#: The cached calls of the router entry points with ``deterministic=False`` (``router_forward``, ``router_dgrad``, -#: ``router_wgrad``), keyed by the entry point's own tensors (:func:`_k1_key`), the operation, its split count and the +#: The cached calls of the router entry points (``cake_dense_gemm_nondet_train_router_forward``, ``cake_dense_gemm_nondet_train_router_dgrad``, +#: ``cake_dense_gemm_nondet_train_router_wgrad``), keyed by the entry point's own tensors (:func:`_k1_key`), the operation, its split count and the #: PDL flag (:class:`_CachedK2Entry`). _K2_ENTRY_CACHE: "OrderedDict[tuple, Any]" = OrderedDict() @@ -4902,7 +4885,7 @@ def _k2_entry( splits: Optional[int], enable_pdl: Optional[bool], ) -> Optional[torch.Tensor]: - """A router entry point's call of ``deterministic=False`` (``op`` and its tensors ``P``, ``Q`` as for + """A router entry point's call (``op`` and its tensors ``P``, ``Q`` as for :func:`_k2_entry_operands`, into ``out``) through :data:`_K2_ENTRY_CACHE`: planned once per key of the entry point's own tensors, the operation, ``splits`` and the PDL flag (``enable_pdl`` resolved for the device), and again after a router table or the route of its layout class changed (:func:`_plan_k2_entry`), then the bound launch with @@ -4962,11 +4945,11 @@ def _drop_convenience_entries() -> None: _ROUTER_MEMO = _RouterMemo(-1, None, None, False, True) -def clear_convenience_caches() -> None: +def cake_dense_gemm_nondet_train_convenience_cache_clear() -> None: """Drop every cached convenience-path launch, what the convenience path keeps from the router tables (:func:`_router_memo`) and every workspace arena of both packages (this package's launches and the deterministic package's). This invalidates every CUDA graph that captured a convenience-path call: its replay would use freed - workspace, so capture such graphs again after it. A captured ``deterministic=False`` convenience-path call that + workspace, so capture such graphs again after it. A captured convenience-path call of this package that this package launches from its arenas (its programs, and the deterministic programs it launches through its binding) stays valid until then: its arena keeps every buffer a launch was bound to, through growths and cache evictions. A call on a row that runs the deterministic program, which this package hands to the deterministic @@ -4983,6 +4966,14 @@ def clear_convenience_caches() -> None: # --------------------------------------------------------------------------- +def is_cake_dense_gemm_nondet_train_supported(device: torch.device) -> bool: + """Whether ``device`` is a CUDA device of compute capability 10.0, 10.3 or 10.7 (``sm_100a`` / ``sm_103a`` / + ``sm_107a``), the devices of the deterministic package whose programs this package reuses (its ``arch_for``). + Shape-level support is checked per call: an unsupported view, layout or size raises ``ValueError``, and a geometry + without a registered program raises ``NotImplementedError``.""" + return deterministic_backend().arch_for(device) is not None + + def _weights_kw(weights_static: bool) -> dict[str, Any]: """``weights_static`` as the keyword the deterministic entry point receives: passed only when set.""" return {"weights_static": True} if weights_static else {} @@ -5019,100 +5010,196 @@ def router_symbol(a_mn: bool, b_mn: bool, out_t: bool) -> str: ) -def plan_dense_projection_gemm( +class _Unset: + """The default of a planner knob the caller did not give: such a knob is not passed on (:func:`_given_knobs`).""" + + __slots__ = () + + def __repr__(self) -> str: + return "" + + +#: The default of every planner knob of :func:`cake_dense_gemm_nondet_train_projection_plan` and +#: :func:`cake_dense_gemm_nondet_train_projection_prepare`. +_UNSET: Any = _Unset() + + +def _given_knobs(values: dict[str, Any]) -> dict[str, Any]: + """The knobs of an entry point's call that the caller gave (``values``: the entry point's ``locals()`` on entry, + without its tensors): only these reach the deterministic planner, so it receives the same keywords as a direct call + of it, and its own defaults apply to the others.""" + return { + name: value + for name, value in values.items() + if value is not _UNSET and name not in ("A", "B", "out") + } + + +def cake_dense_gemm_nondet_train_projection_plan( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, - deterministic: bool = True, - **kwargs: Any, + splits: Optional[int] = _UNSET, + chains: Optional[int] = _UNSET, + sm_count: int = _UNSET, + l2_bytes: int = _UNSET, + transposed_out: bool = _UNSET, + sk: bool | str = _UNSET, + block_n: Optional[int] = _UNSET, + stages: Optional[int] = _UNSET, + epi: Optional[str] = _UNSET, + slots: Optional[int] = _UNSET, + cta_rows: Optional[int] = _UNSET, + pf: Optional[int] = _UNSET, + sk_max_units: Optional[int] = _UNSET, + sk_parts: Optional[int] = _UNSET, + promo: Optional[str] = _UNSET, + hints: Optional[tuple] = _UNSET, + group_m: Optional[int] = _UNSET, + f32_v8: Optional[bool] = _UNSET, + quad_store: Optional[bool] = _UNSET, + park: Optional[bool] = _UNSET, + ovl: Optional[bool] = _UNSET, + htail: Optional[bool] = _UNSET, + b_swz: Optional[int] = _UNSET, + sk_exact: Optional[int] = _UNSET, + batch_group: Optional[int] = _UNSET, + store_ef: Optional[bool] = _UNSET, + store_hint: Optional[str] = _UNSET, + pd: Optional[int] = _UNSET, + sh: Optional[int] = _UNSET, + sk_sync: Optional[bool] = _UNSET, + sk_sync_m: Optional[int] = _UNSET, + sk_slab: Optional[int] = _UNSET, + cta1: Optional[bool] = _UNSET, + cgrp: Optional[bool] = _UNSET, + ovr: Optional[bool] = _UNSET, + stt: Optional[bool] = _UNSET, + pro: Optional[int] = _UNSET, + sk_first: Optional[bool] = _UNSET, + lw: Optional[bool] = _UNSET, + mg: Optional[bool] = _UNSET, + sw: Optional[int] = _UNSET, + form: Optional[str] = _UNSET, + arch: str = _UNSET, ) -> Any: - """Validate the views and plan the K1 launch: with ``deterministic=False`` (the only value accepted) a - weight-gradient layout (both operands MN-major) is planned for this package's split-K program - (:class:`AtomicGemmPlan`; the knobs ``splits`` / ``block_n`` / ``chains`` are this package's, the others -- - ``sm_count``, ``l2_bytes``, ``transposed_out``, the planner knobs, ``arch`` -- apply when the row runs the - deterministic program). A row the planner keeps on the deterministic program, and every other layout, returns the - deterministic package's plan (``cake_dense_gemm_det_train_projection_plan``).""" - _require_false( - deterministic, "plan_dense_projection_gemm", "cake_dense_gemm_det_train_projection_plan" - ) + """Validate the views and plan the K1 launch: a weight-gradient layout (both operands MN-major) is planned for + this package's split-K program (:class:`CakeDenseGemmNondetProjectionPlan`; ``splits``, ``block_n`` and ``chains`` + are this package's knobs). A row the planner keeps on the deterministic program, and every other layout, returns + the deterministic package's plan (``cake_dense_gemm_det_train_projection_plan``, a + ``CakeDenseGemmDetProjectionPlan``). The other keyword-only knobs are the deterministic planner's public ones + (``sm_count`` and ``l2_bytes`` included): a knob the caller gives reaches that planner as given, and + ``transposed_out``, ``arch`` and ``sm_count`` also apply to this package's plan; a knob the caller does not give + takes the deterministic planner's default.""" + knobs = _given_knobs(locals()) if _route(_k1_layout(A, B)) == "det": return deterministic_backend().cake_dense_gemm_det_train_projection_plan( - A, B, out, **kwargs + A, B, out, **knobs ) - own, det_kwargs = _split_knobs(kwargs) + own, det_knobs = _split_knobs(knobs) return _atomic_k1( "plan", A, B, out, - transposed_out=bool(det_kwargs.get("transposed_out", False)), + transposed_out=bool(det_knobs.get("transposed_out", False)), own=own, - arch=det_kwargs.get("arch"), - sm_count=det_kwargs.get("sm_count"), + arch=det_knobs.get("arch"), + sm_count=det_knobs.get("sm_count"), det_call=lambda: deterministic_backend().cake_dense_gemm_det_train_projection_plan( - A, B, out, **det_kwargs + A, B, out, **det_knobs ), ) -def prepare_dense_projection_gemm( +def cake_dense_gemm_nondet_train_projection_prepare( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, - deterministic: bool = True, - **kwargs: Any, + splits: Optional[int] = _UNSET, + chains: Optional[int] = _UNSET, + transposed_out: bool = _UNSET, + sk: bool | str = _UNSET, + block_n: Optional[int] = _UNSET, + stages: Optional[int] = _UNSET, + epi: Optional[str] = _UNSET, + slots: Optional[int] = _UNSET, + cta_rows: Optional[int] = _UNSET, + pf: Optional[int] = _UNSET, + sk_max_units: Optional[int] = _UNSET, + sk_parts: Optional[int] = _UNSET, + promo: Optional[str] = _UNSET, + hints: Optional[tuple] = _UNSET, + group_m: Optional[int] = _UNSET, + f32_v8: Optional[bool] = _UNSET, + quad_store: Optional[bool] = _UNSET, + park: Optional[bool] = _UNSET, + ovl: Optional[bool] = _UNSET, + htail: Optional[bool] = _UNSET, + b_swz: Optional[int] = _UNSET, + sk_exact: Optional[int] = _UNSET, + batch_group: Optional[int] = _UNSET, + store_ef: Optional[bool] = _UNSET, + store_hint: Optional[str] = _UNSET, + pd: Optional[int] = _UNSET, + sh: Optional[int] = _UNSET, + sk_sync: Optional[bool] = _UNSET, + sk_sync_m: Optional[int] = _UNSET, + sk_slab: Optional[int] = _UNSET, + cta1: Optional[bool] = _UNSET, + cgrp: Optional[bool] = _UNSET, + ovr: Optional[bool] = _UNSET, + stt: Optional[bool] = _UNSET, + pro: Optional[int] = _UNSET, + form: Optional[str] = _UNSET, + enable_pdl: Optional[bool] = _UNSET, ) -> Any: - """Validate one K1 binding, plan it for the device and prepare its allocation-free launch: with - ``deterministic=False`` (the only value accepted) a weight-gradient layout prepares this package's split-K program - (:class:`PreparedAtomicGemm`; its workspace is owned by the prepared object; ``splits`` / ``block_n`` / ``chains`` - as for :func:`plan_dense_projection_gemm`). A row the planner keeps on the deterministic program, and every other - layout, prepares the deterministic package's ``cake_dense_gemm_det_train_projection_prepare`` (its keywords: - ``transposed_out`` and the planner knobs). ``enable_pdl`` (default: the device supports it) makes either launch - a programmatic dependent launch.""" - _require_false( - deterministic, - "prepare_dense_projection_gemm", - "cake_dense_gemm_det_train_projection_prepare", - ) + """Validate one K1 binding, plan it for the device and prepare its allocation-free launch: a weight-gradient + layout prepares this package's split-K program (:class:`CakeDenseGemmNondetProjectionLaunch`; its workspace is + owned by the prepared object; ``splits``, ``block_n`` and ``chains`` as for + :func:`cake_dense_gemm_nondet_train_projection_plan`). A row the planner keeps on the deterministic program, and + every other layout, prepares the deterministic package's ``cake_dense_gemm_det_train_projection_prepare`` (a + ``CakeDenseGemmDetProjectionLaunch``). The other keyword-only knobs are that entry point's public ones, passed + on only when the caller gives them; ``enable_pdl`` (default: the device supports it) makes either launch a + programmatic dependent launch.""" + knobs = _given_knobs(locals()) if _route(_k1_layout(A, B)) == "det": return deterministic_backend().cake_dense_gemm_det_train_projection_prepare( - A, B, out, **kwargs + A, B, out, **knobs ) - own, det_kwargs = _split_knobs(kwargs) + own, det_knobs = _split_knobs(knobs) return _atomic_k1( "prepare", A, B, out, - transposed_out=bool(det_kwargs.get("transposed_out", False)), + transposed_out=bool(det_knobs.get("transposed_out", False)), own=own, det_call=lambda: deterministic_backend().cake_dense_gemm_det_train_projection_prepare( - A, B, out, **det_kwargs + A, B, out, **det_knobs ), - enable_pdl=det_kwargs.get("enable_pdl"), + enable_pdl=det_knobs.get("enable_pdl"), ) -def dense_projection_gemm( +def cake_dense_gemm_nondet_train_projection( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, transposed_out: bool = False, - deterministic: bool = True, enable_pdl: Optional[bool] = None, weights_static: bool = False, ) -> torch.Tensor: """``out[l, m, n] = sum_k A[l, m, k] B[l, k, n]`` for strided views (the cached convenience path; see the - deterministic package for the view contract). ``deterministic=False`` (the only value accepted) runs a - weight-gradient layout (both operands MN-major) on this package's split-K program where the planner takes the - row; every other layout and row runs the deterministic program. ``enable_pdl`` (default: the device supports it) + deterministic package for the view contract). A weight-gradient layout (both operands MN-major) runs this + package's split-K program where the planner takes the row; every other layout and row runs the deterministic + program. ``enable_pdl`` (default: the device supports it) launches either program as a programmatic dependent launch; each waits for the preceding kernel before its first global access, and ``weights_static=True`` lets the deterministic swapped small-token forward start its loads of ``B`` before the wait (see the deterministic package's ``cake_dense_gemm_det_train_projection``).""" - _require_false(deterministic, "dense_projection_gemm", "cake_dense_gemm_det_train_projection") def det_call() -> Any: return deterministic_backend().cake_dense_gemm_det_train_projection( @@ -5151,21 +5238,19 @@ def det_args() -> tuple: ) -def projection_forward( +def cake_dense_gemm_nondet_train_projection_forward( X: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None, *, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, N] = X[T, K] @ W[N, K].T`` (bf16 out by default), ``deterministic=False`` only. The deterministic - package's programs (a recomputed forward must reproduce its result; where the deterministic plan splits K across - CTAs, it adds the partials in a fixed order) through this package's binding of them (:func:`_det_k1`): the - deterministic plan, program and bits, except on a row of :data:`K1_FORWARD_RULES_BY_SM` for the device, which runs - the program the deterministic planner gives it with the table's knobs (:func:`k1_forward_knobs`; other bits, still - the same on every run). ``enable_pdl`` as for :func:`dense_projection_gemm`.""" - _require_false(deterministic, "projection_forward", "cake_dense_gemm_det_train_projection_forward") + """``out[T, N] = X[T, K] @ W[N, K].T`` (bf16 out by default). The deterministic package's programs (a + recomputed forward must reproduce its result; where the deterministic plan splits K across CTAs, it adds the + partials in a fixed order) through this package's binding of them (:func:`_det_k1`): the deterministic plan, + program and bits, except on a row of :data:`K1_FORWARD_RULES_BY_SM` for the device, which runs the program the + deterministic planner gives it with the table's knobs (:func:`k1_forward_knobs`; other bits, still the same on + every run). ``enable_pdl`` as for :func:`cake_dense_gemm_nondet_train_projection`.""" if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.bfloat16) o = out @@ -5184,19 +5269,17 @@ def projection_forward( ) -def projection_dgrad( +def cake_dense_gemm_nondet_train_projection_dgrad( G: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None, *, out_dtype=None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, K] = G[T, N] @ W[N, K]``, ``deterministic=False`` only. The deterministic program (whole output tiles) - through this package's binding of it (:func:`_det_k1`: the same plan, program and bits). ``enable_pdl`` as for - :func:`dense_projection_gemm`.""" - _require_false(deterministic, "projection_dgrad", "cake_dense_gemm_det_train_projection_dgrad") + """``out[T, K] = G[T, N] @ W[N, K]``. The deterministic program (whole output tiles) through this package's + binding of it (:func:`_det_k1`: the same plan, program and bits). ``enable_pdl`` as for + :func:`cake_dense_gemm_nondet_train_projection`.""" if out is None: out = torch.empty( G.shape[0], @@ -5220,23 +5303,19 @@ def projection_dgrad( ) -def prepare_projection_wgrad( +def cake_dense_gemm_nondet_train_projection_wgrad_prepare( G: torch.Tensor, X: torch.Tensor, out: torch.Tensor, *, sk: bool | str = "auto", cta_rows: Optional[int] = None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> Any: - """Prepared ``out[N, K] = G[T, N].T @ X[T, K]``, ``deterministic=False`` only: this package's split-K program as + """Prepared ``out[N, K] = G[T, N].T @ X[T, K]``: this package's split-K program as the transposed-store GEMM ``X.T @ G`` (``M = K``, ``N = N``) where the planner takes the row; a row it does not take prepares the deterministic package's ``cake_dense_gemm_det_train_projection_wgrad_prepare`` (its orientation, ``sk``, - ``cta_rows``). ``enable_pdl`` as for :func:`prepare_dense_projection_gemm`.""" - _require_false( - deterministic, "prepare_projection_wgrad", "cake_dense_gemm_det_train_projection_wgrad_prepare" - ) + ``cta_rows``). ``enable_pdl`` as for :func:`cake_dense_gemm_nondet_train_projection_prepare`.""" return _atomic_k1( "prepare", X.t(), @@ -5251,7 +5330,7 @@ def prepare_projection_wgrad( ) -def projection_wgrad( +def cake_dense_gemm_nondet_train_projection_wgrad( G: torch.Tensor, X: torch.Tensor, out: Optional[torch.Tensor] = None, @@ -5259,15 +5338,13 @@ def projection_wgrad( out_dtype=None, sk: bool | str = "auto", cta_rows: Optional[int] = None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[N, K] = G[T, N].T @ X[T, K]`` (ragged reduction length ``T``), ``deterministic=False`` only: this + """``out[N, K] = G[T, N].T @ X[T, K]`` (ragged reduction length ``T``): this package's split-K program where the planner takes the row (the split partials summed in the order they finish, the output rounded once), else the deterministic package's ``cake_dense_gemm_det_train_projection_wgrad`` program (its orientation, ``sk`` and ``cta_rows``) through this package's binding of it (:func:`_launch_cached_atomic_k1`); ``enable_pdl`` - as for :func:`dense_projection_gemm`.""" - _require_false(deterministic, "projection_wgrad", "cake_dense_gemm_det_train_projection_wgrad") + as for :func:`cake_dense_gemm_nondet_train_projection`.""" if out is None: out = torch.empty( int(G.shape[1]), @@ -5324,19 +5401,17 @@ def _wgrad_det_args( return A, B, out, transposed, (("cta_rows", cta_rows), ("sk", sk)) -def plan_router_fp32_gemm( +def cake_dense_gemm_nondet_train_router_plan( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] = None, - deterministic: bool = True, ) -> Any: - """Classify the fp32 views and plan the K2 launch, ``deterministic=False`` only: this package's stream-K program - for the ``kk`` / ``nn_t`` classes and the SIMT programs in their windows (:class:`AtomicRouterPlan`; ``splits`` + """Classify the fp32 views and plan the K2 launch: this package's stream-K program + for the ``kk`` / ``nn_t`` classes and the SIMT programs in their windows (:class:`CakeDenseGemmNondetRouterPlan`; ``splits`` sets ``splits`` units per tile, see :func:`router_units`); a row the planner keeps on the deterministic program, and a layout class routed ``"det"``, returns the deterministic package's plan (``cake_dense_gemm_det_train_router_plan``).""" - _require_false(deterministic, "plan_router_fp32_gemm", "cake_dense_gemm_det_train_router_plan") if _route(_k2_layout(A, B)) == "det": return deterministic_backend().cake_dense_gemm_det_train_router_plan( A, B, out, splits=splits @@ -5344,25 +5419,21 @@ def plan_router_fp32_gemm( return _atomic_k2("plan", A, B, out, splits=splits) -def prepare_router_fp32_gemm( +def cake_dense_gemm_nondet_train_router_prepare( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] = None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> Any: - """Validate one fp32 binding, plan it and prepare its allocation-free launch, ``deterministic=False`` only: this + """Validate one fp32 binding, plan it and prepare its allocation-free launch: this package's stream-K program for the ``kk`` / ``nn_t`` classes (its launch splits the second operand, zeroes ``out`` on the stream when a segment adds into it, and runs the program: no partial buffers), and a SIMT program in its windows (the ``nn_t`` and ``kn`` classes: one kernel over the fp32 operands); a row the planner keeps on the deterministic program prepares the deterministic package's ``cake_dense_gemm_det_train_router_prepare``. ``enable_pdl`` (default: the device supports it) launches the split and the stream-K program as programmatic dependent launches.""" - _require_false( - deterministic, "prepare_router_fp32_gemm", "cake_dense_gemm_det_train_router_prepare" - ) if _route(_k2_layout(A, B)) == "det": return deterministic_backend().cake_dense_gemm_det_train_router_prepare( A, B, out, splits=splits, **_pdl_kw(enable_pdl) @@ -5388,42 +5459,38 @@ def _router_gemm( ) -def router_fp32_gemm( +def cake_dense_gemm_nondet_train_router( A: torch.Tensor, B: torch.Tensor, out: torch.Tensor, *, splits: Optional[int] = None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: """``out[m, n] = sum_k A[m, k] B[k, n]`` for fp32 views through the BF16x3 emulation (the cached convenience - path), ``deterministic=False`` only: the ``kk`` / ``nn_t`` rows the planner takes (:func:`router_units`) on this + path): the ``kk`` / ``nn_t`` rows the planner takes (:func:`router_units`) on this package's stream-K program, whose part tiles add into ``out`` in the order they finish, instead of the fixed-order split-K sum, the ``nn_t`` / ``kn`` token counts of a SIMT window on its exact-FP32 program (:func:`router_simt`) and the ``kk`` token counts of a forward SIMT window on the forward's cluster program (:func:`router_forward_simt`); every other row runs the deterministic program. ``enable_pdl`` as for - :func:`prepare_router_fp32_gemm`.""" - _require_false(deterministic, "router_fp32_gemm", "cake_dense_gemm_det_train_router") + :func:`cake_dense_gemm_nondet_train_router_prepare`.""" return _router_gemm(A, B, out, splits, enable_pdl) -def router_forward( +def cake_dense_gemm_nondet_train_router_forward( X: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None, *, splits: Optional[int] = None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, N] = X[T, K] @ W[N, K].T`` in FP32 through the emulation, ``deterministic=False`` only. A row runs + """``out[T, N] = X[T, K] @ W[N, K].T`` in FP32 through the emulation. A row runs this package's program only with an explicit ``splits`` (at most two partials per element: the logits repeat bit for bit) or a verdict of :data:`ROUTER_RULES` (the token windows of the forward's SIMT program, - :func:`router_forward_simt`), else the deterministic program (its plan for ``router_fp32_gemm(X, W.T, out)``, + :func:`router_forward_simt`), else the deterministic program (its plan for ``cake_dense_gemm_nondet_train_router(X, W.T, out)``, through this package's binding, :func:`_launch_cached_splitk_k2`): without a split count, while no router table gives the forward a verdict, without the per-call check of the row.""" - _require_false(deterministic, "router_forward", "cake_dense_gemm_det_train_router_forward") if out is None: out = torch.empty(X.shape[0], W.shape[0], device=X.device, dtype=torch.float32) if _k2_entry("fwd", X, W, out, splits, enable_pdl) is not None: @@ -5436,21 +5503,19 @@ def router_forward( return _router_gemm(X, W.t(), out, splits, enable_pdl, views) -def router_dgrad( +def cake_dense_gemm_nondet_train_router_dgrad( G: torch.Tensor, W: torch.Tensor, out: Optional[torch.Tensor] = None, *, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[T, K] = G[T, N] @ W[N, K]`` in FP32, ``deterministic=False`` only: in the token windows of the row's + """``out[T, K] = G[T, N] @ W[N, K]`` in FP32: in the token windows of the row's input-gradient verdict in :data:`ROUTER_RULES` (or, on a part with a table of its own, :data:`ROUTER_RULES_BY_SM`) the exact-FP32 SIMT program (:func:`router_simt`: compensated sums of the products of ``G`` and ``W`` read in place, every element of ``out`` stored once), and at every other row and token count the deterministic program (one split: nothing to reduce) through this package's binding of it (:func:`_launch_cached_splitk_k2`: the same - plan, program and bits). ``enable_pdl`` as for :func:`prepare_router_fp32_gemm`.""" - _require_false(deterministic, "router_dgrad", "cake_dense_gemm_det_train_router_dgrad") + plan, program and bits). ``enable_pdl`` as for :func:`cake_dense_gemm_nondet_train_router_prepare`.""" if out is None: out = torch.empty(G.shape[0], W.shape[1], device=G.device, dtype=torch.float32) if _k2_entry("dgrad", G, W, out, None, enable_pdl) is not None: @@ -5458,21 +5523,19 @@ def router_dgrad( return _router_gemm(G, W, out, None, enable_pdl, (G.unsqueeze(0), W, _DGRAD_VIEWS)) -def router_wgrad( +def cake_dense_gemm_nondet_train_router_wgrad( G: torch.Tensor, X: torch.Tensor, out: Optional[torch.Tensor] = None, *, splits: Optional[int] = None, - deterministic: bool = True, enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - """``out[N, K] = G[T, N].T @ X[T, K]`` in FP32, ``deterministic=False`` only: the swapped GEMM ``X.T @ G`` with a + """``out[N, K] = G[T, N].T @ X[T, K]`` in FP32: the swapped GEMM ``X.T @ G`` with a transposed store; ``splits`` defaults to the partition of the row's verdict in :data:`ROUTER_RULES` or, on a part with a table of its own, in :data:`ROUTER_RULES_BY_SM` (:func:`router_partition`; a row or token count without one runs the deterministic program, through this package's binding of it without planning the row here, :func:`_launch_cached_splitk_k2`).""" - _require_false(deterministic, "router_wgrad", "cake_dense_gemm_det_train_router_wgrad") if out is None: out = torch.empty(G.shape[1], X.shape[1], device=G.device, dtype=torch.float32) if _k2_entry("wgrad", G, X, out, splits, enable_pdl) is not None: From dc9112ddb9858ee6a2b7e3c0c4076eecc7d3f39b Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 00:20:33 -0800 Subject: [PATCH 09/13] test(cake_dense_gemm_nondet_train): names test, tests, bench, README The tests, benchmark and README of the non-deterministic dense GEMMs follow the module rename (flashinfer-ai/flashinfer#8928). - tests/experimental/test_cake_dense_projection_gemm.py becomes test_cake_dense_gemm_nondet_train.py: no deterministic= argument, the new names and package paths, the deterministic entries under their new names; the alert tests (and the fixed-order cluster program's no-alert test) stay. - New test_cake_dense_gemm_nondet_train_names.py: the mode word of every public name, checked by word (a self-contained helper); deterministic= binds to no callable and raises TypeError before anything runs; no **kwargs; the projection planner's keyword-only knobs are the deterministic planner's public ones plus splits and chains, and only the given ones reach it; the former modules fail at import and the former names are gone, with no read-through to the deterministic package; the module returns the backend's objects without a warning and imports nothing before first use. - benchmarks/bench_cake_dense_projection_gemm.py becomes bench_cake_dense_gemm_nondet_train.py. - README: the names of both modules, the alert rule per launch (the cluster program is fixed-order; CUDA graphs decide the alert at capture) and when the module returns a deterministic object. --- ... => bench_cake_dense_gemm_nondet_train.py} | 90 +- .../cake_dense_gemm_nondet_train/README.md | 196 ++- ...y => test_cake_dense_gemm_nondet_train.py} | 1509 ++++++++--------- ...test_cake_dense_gemm_nondet_train_names.py | 505 ++++++ 4 files changed, 1391 insertions(+), 909 deletions(-) rename benchmarks/{bench_cake_dense_projection_gemm.py => bench_cake_dense_gemm_nondet_train.py} (96%) rename tests/experimental/{test_cake_dense_projection_gemm.py => test_cake_dense_gemm_nondet_train.py} (91%) create mode 100644 tests/experimental/test_cake_dense_gemm_nondet_train_names.py diff --git a/benchmarks/bench_cake_dense_projection_gemm.py b/benchmarks/bench_cake_dense_gemm_nondet_train.py similarity index 96% rename from benchmarks/bench_cake_dense_projection_gemm.py rename to benchmarks/bench_cake_dense_gemm_nondet_train.py index 050bc43a898..a97f57b2ef4 100644 --- a/benchmarks/bench_cake_dense_projection_gemm.py +++ b/benchmarks/bench_cake_dense_gemm_nondet_train.py @@ -14,15 +14,17 @@ limitations under the License. """ -"""Benchmark the numerical modes of the experimental Cake dense projection GEMMs (SM100 / SM103 / SM107). +"""Benchmark the non-deterministic Cake dense GEMMs (``flashinfer.cake_dense_gemm_nondet_train``) against the +deterministic ones (``flashinfer.cake_dense_gemm_det_train``) and the non-Cake SOTA (SM100 / SM103 / SM107). Rows: every row of the deterministic family's benchmark table (flashinfer-ai/flashinfer#5677 and #7043): the projection GEMMs, the batched MLA head projections and the FP32 router x {forward, input gradient, weight gradient} x their token counts and output dtypes (294 rows: 124 weight gradients, 12 router rows, 158 forward / input-gradient rows). Arms per row: -* ``det``: the deterministic package's prepared launch (``dense_projection_det_gemm``, its production plan). -* ``atomic``: ``deterministic=False`` through ``flashinfer.experimental.dense_projection_gemm`` (prepared): this +* ``det``: the deterministic package's prepared launch (``cake_dense_gemm_det_train_projection_prepare`` and the + router and weight-gradient prepare entries, its production plan). +* ``atomic``: the non-deterministic package (``flashinfer.cake_dense_gemm_nondet_train``, prepared): this package's program where its planner takes the row, else the deterministic program (``atomic_program`` says which). On a delegated row (every forward and input gradient, the rows the planners leave) the arm shows that the dispatch costs nothing. @@ -71,7 +73,7 @@ Usage:: - python benchmarks/bench_cake_dense_projection_gemm.py [--rows indexer_hw_wgrad_bf16_t16231 ...] + python benchmarks/bench_cake_dense_gemm_nondet_train.py [--rows indexer_hw_wgrad_bf16_t16231 ...] [--families proj,mla,router] [--ops fwd,dgrad,wgrad] [--T 16231] [--only-atomic] [--steps 20] [--groups 5] [--modes eager,graph] [--accuracy] [--memory] [--composite] [--resident] [--exclusive] [--json out.json] [--csv out.csv] [--markdown out.md] @@ -95,10 +97,12 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from flashinfer.experimental.dense_projection_det_gemm import ( # noqa: E402 +from flashinfer.experimental.cake.cake_dense_gemm_det_train import ( # noqa: E402 cake_backend as det_backend, ) -from flashinfer.experimental.dense_projection_gemm import cake_backend # noqa: E402 +from flashinfer.experimental.cake.cake_dense_gemm_nondet_train import ( # noqa: E402 + cake_backend, +) from flashinfer.testing import bench_gpu_time # noqa: E402 # --------------------------------------------------------------------------- @@ -285,58 +289,69 @@ def fresh_out(out): def prepare_det(spec, inp, out): """The deterministic package's prepared launch of the row (its production plan).""" if spec["family"] == "router": - return det_backend.prepare_router_fp32_det_gemm(inp["A"], inp["B"], out) + return det_backend.cake_dense_gemm_det_train_router_prepare( + inp["A"], inp["B"], out + ) if spec["family"] == "proj" and spec["op"] == "wgrad": - return det_backend.prepare_projection_det_wgrad(inp["G"], inp["X"], out) - return det_backend.prepare_dense_projection_det_gemm(inp["A"], inp["B"], out) + return det_backend.cake_dense_gemm_det_train_projection_wgrad_prepare( + inp["G"], inp["X"], out + ) + return det_backend.cake_dense_gemm_det_train_projection_prepare( + inp["A"], inp["B"], out + ) def prepare_atomic(spec, inp, out, **knobs): - """``deterministic=False`` through this package (``knobs``: ``splits`` / ``block_n``, the control arm).""" + """This package's prepared launch of the row (``knobs``: ``splits`` / ``block_n``, the control arm).""" if spec["family"] == "router": - return cake_backend.prepare_router_fp32_gemm( - inp["A"], inp["B"], out, deterministic=False, **knobs + return cake_backend.cake_dense_gemm_nondet_train_router_prepare( + inp["A"], inp["B"], out, **knobs ) if spec["family"] == "proj" and spec["op"] == "wgrad": if not knobs: - return cake_backend.prepare_projection_wgrad( - inp["G"], inp["X"], out, deterministic=False + return cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + inp["G"], inp["X"], out ) - return cake_backend.prepare_dense_projection_gemm( + return cake_backend.cake_dense_gemm_nondet_train_projection_prepare( inp["X"].t(), inp["G"], out, transposed_out=True, - deterministic=False, **knobs, ) - return cake_backend.prepare_dense_projection_gemm( - inp["A"], inp["B"], out, deterministic=False, **knobs + return cake_backend.cake_dense_gemm_nondet_train_projection_prepare( + inp["A"], inp["B"], out, **knobs ) def convenience_det(spec, inp, out): """The deterministic package's convenience entry point of the row (cached plan, arena workspace).""" if spec["family"] == "router": - return lambda: det_backend.router_fp32_det_gemm(inp["A"], inp["B"], out) + return lambda: det_backend.cake_dense_gemm_det_train_router( + inp["A"], inp["B"], out + ) if spec["family"] == "proj" and spec["op"] == "wgrad": - return lambda: det_backend.projection_det_wgrad(inp["G"], inp["X"], out) - return lambda: det_backend.dense_projection_det_gemm(inp["A"], inp["B"], out) + return lambda: det_backend.cake_dense_gemm_det_train_projection_wgrad( + inp["G"], inp["X"], out + ) + return lambda: det_backend.cake_dense_gemm_det_train_projection( + inp["A"], inp["B"], out + ) def convenience_atomic(spec, inp, out): - """``deterministic=False`` through this package's convenience entry point of the row (cached plan, arena + """This package's convenience entry point of the row (cached plan, arena workspace).""" if spec["family"] == "router": - return lambda: cake_backend.router_fp32_gemm( - inp["A"], inp["B"], out, deterministic=False + return lambda: cake_backend.cake_dense_gemm_nondet_train_router( + inp["A"], inp["B"], out ) if spec["family"] == "proj" and spec["op"] == "wgrad": - return lambda: cake_backend.projection_wgrad( - inp["G"], inp["X"], out, deterministic=False + return lambda: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + inp["G"], inp["X"], out ) - return lambda: cake_backend.dense_projection_gemm( - inp["A"], inp["B"], out, deterministic=False + return lambda: cake_backend.cake_dense_gemm_nondet_train_projection( + inp["A"], inp["B"], out ) @@ -344,7 +359,10 @@ def is_atomic(prepared): """Whether a prepared launch runs this package's program (else the deterministic package's).""" return isinstance( prepared, - (cake_backend.PreparedAtomicGemm, cake_backend.PreparedAtomicRouterGemm), + ( + cake_backend.CakeDenseGemmNondetProjectionLaunch, + cake_backend.CakeDenseGemmNondetRouterLaunch, + ), ) @@ -352,7 +370,7 @@ def plan_record(prepared): """The plan facts of a prepared launch (and the rule-table key and value of this package's program, so a per-row verdict can be written into ``K1_SPLIT_RULES`` / ``ROUTER_RULES``).""" record = dict(template=prepared.template, grid=list(prepared.grid)) - if isinstance(prepared, cake_backend.PreparedAtomicGemm): + if isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch): p = prepared.plan # the rule key is the transposed-store launch's (M = d_in, N = d_out); a row-major plan exchanges them key_m, key_n = (p.M, p.N) if p.out_t else (p.N, p.M) @@ -373,7 +391,7 @@ def plan_record(prepared): ], rule_value=[p.block_n, p.splits], ) - elif isinstance(prepared, cake_backend.PreparedAtomicRouterGemm): + elif isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch): p = prepared.plan # the rule key's (hidden, experts) of the launched GEMM, as _router_rule_spec reads them if p.layout == "kk": @@ -1138,10 +1156,10 @@ def build_arms(spec, inp, args, libs, entry): entry["atomic_program"] = is_atomic(p) entry["atomic_plan"] = plan_record(p) control = None - if isinstance(p, cake_backend.PreparedAtomicGemm) and p.splits > 1: + if isinstance(p, cake_backend.CakeDenseGemmNondetProjectionLaunch) and p.splits > 1: control = dict(splits=1, block_n=p.plan.block_n) elif ( - isinstance(p, cake_backend.PreparedAtomicRouterGemm) + isinstance(p, cake_backend.CakeDenseGemmNondetRouterLaunch) and p.plan.units != p.plan.tiles ): control = dict(splits=1) @@ -1278,7 +1296,7 @@ def measure_row(label, spec, args, libs, resident, uuid): entry["traceback"] = traceback.format_exc() print(f"{label:34s} failed: {entry['error']}", flush=True) finally: - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() torch.cuda.empty_cache() return entry @@ -1360,12 +1378,12 @@ def _measure_row(entry, spec, args, libs, resident, uuid): ("det", convenience_det), ("atomic", convenience_atomic), ): - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() fn = convenience(spec, inp, outs[mode]) retained, peak = memory_of(lambda fn=fn: first_call(fn)) mem[f"{mode}_convenience_retained"] = retained mem[f"{mode}_convenience_peak_temp"] = peak - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() entry["memory_bytes"] = mem if args.host_overhead: entry["host_us_prepared"] = { diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md index 8e237d21a66..1f1c4433934 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md @@ -1,73 +1,98 @@ -# Cake dense projection GEMM backend, non-deterministic (SM100 / SM103 / SM107) +# Cake dense GEMM backend, non-deterministic (SM100 / SM103 / SM107) -Experimental backend for the GLM-5.2 training GEMMs -- the BF16 dense projection GEMM (K1: BF16 operands, FP32 -accumulation, BF16 or FP32 output, strided views, ragged token counts, batched head projections) and the FP32 -router-gate GEMM through split-BF16x3 tensor-core emulation (K2) -- with programs whose split-K reductions add the -partials in the order they finish. Family issue: flashinfer-ai/flashinfer#5677. +Backend of the public module `flashinfer.cake_dense_gemm_nondet_train`: the GLM-5.2 training GEMMs -- the BF16 dense projection GEMM (K1: +BF16 operands, FP32 accumulation, BF16 or FP32 output, strided views, ragged token counts, batched head projections) +and the FP32 router-gate GEMM through split-BF16x3 tensor-core emulation (K2) -- with programs whose split-K +reductions add the partials in the order they finish. Family issue: flashinfer-ai/flashinfer#5677. -The deterministic implementation is the package `flashinfer.experimental.dense_projection_det_gemm` -(flashinfer-ai/flashinfer#5677, #7043; see its README for the view contract, the host paths and the per-architecture -programs). This package keeps its operations, dtypes, accuracy budgets and host paths. Only run-to-run bitwise -reproducibility is relaxed. +The deterministic implementation is the module `flashinfer.cake_dense_gemm_det_train` (backend +`flashinfer.experimental.cake.cake_dense_gemm_det_train`; flashinfer-ai/flashinfer#5677, #7043; see its README for the view +contract, the host paths and the per-architecture programs). This package keeps its operations, dtypes, accuracy +budgets and host paths. Only run-to-run bitwise reproducibility is relaxed: a result of this module is not bitwise +reproducible across runs. -Both the API and the backend are experimental: no compatibility guarantee. +## Names + +The module selects the mode: there is no mode argument. Import the names from `flashinfer.cake_dense_gemm_nondet_train`; it loads this +backend on the first attribute access and returns the backend's own objects. + +| operation | `flashinfer.cake_dense_gemm_nondet_train` | `flashinfer.cake_dense_gemm_det_train` | +|---|---|---| +| K1 GEMM (convenience path), plan, prepared launch | `cake_dense_gemm_nondet_train_projection`, `cake_dense_gemm_nondet_train_projection_plan`, `cake_dense_gemm_nondet_train_projection_prepare` | `cake_dense_gemm_det_train_projection`, `cake_dense_gemm_det_train_projection_plan`, `cake_dense_gemm_det_train_projection_prepare` | +| K1 forward, input gradient, weight gradient | `cake_dense_gemm_nondet_train_projection_forward`, `cake_dense_gemm_nondet_train_projection_dgrad`, `cake_dense_gemm_nondet_train_projection_wgrad`, `cake_dense_gemm_nondet_train_projection_wgrad_prepare` | `cake_dense_gemm_det_train_projection_forward`, `cake_dense_gemm_det_train_projection_dgrad`, `cake_dense_gemm_det_train_projection_wgrad`, `cake_dense_gemm_det_train_projection_wgrad_prepare` | +| K2 router GEMM (convenience path), plan, prepared launch | `cake_dense_gemm_nondet_train_router`, `cake_dense_gemm_nondet_train_router_plan`, `cake_dense_gemm_nondet_train_router_prepare` | `cake_dense_gemm_det_train_router`, `cake_dense_gemm_det_train_router_plan`, `cake_dense_gemm_det_train_router_prepare` | +| K2 forward, input gradient, weight gradient | `cake_dense_gemm_nondet_train_router_forward`, `cake_dense_gemm_nondet_train_router_dgrad`, `cake_dense_gemm_nondet_train_router_wgrad` | `cake_dense_gemm_det_train_router_forward`, `cake_dense_gemm_det_train_router_dgrad`, `cake_dense_gemm_det_train_router_wgrad` | +| plan and prepared-launch types | `CakeDenseGemmNondetProjectionPlan`, `CakeDenseGemmNondetProjectionLaunch`, `CakeDenseGemmNondetRouterPlan`, `CakeDenseGemmNondetRouterLaunch` | `CakeDenseGemmDetProjectionPlan`, `CakeDenseGemmDetProjectionLaunch`, `CakeDenseGemmDetRouterPlan`, `CakeDenseGemmDetRouterLaunch` | +| convenience caches and workspace arenas | `cake_dense_gemm_nondet_train_convenience_cache_clear` (both packages') | `cake_dense_gemm_det_train_convenience_cache_clear` | +| support query | `is_cake_dense_gemm_nondet_train_supported(device)` | `is_cake_dense_gemm_det_train_supported(device)` | +| constants | `ARCHES` | `ARCHES`, `SUPPORTED_COMPUTE_CAPABILITIES` | + +*Table: the public names of the two modules, per operation.* + +- `cake_dense_gemm_nondet_train_projection_plan` and `cake_dense_gemm_nondet_train_projection_prepare` take keyword-only knobs: `splits`, `chains` and `block_n` of + this package's planner, and the public knobs of the deterministic planner of the same entry. A knob the caller + gives reaches the deterministic planner as given; a knob the caller does not give takes that planner's default. +- The convenience caches of both packages are bounded by one setting, the deterministic backend's + `cake_backend.CONVENIENCE_CACHE_SIZE` (read on every insertion); neither public module has that name. ## Mode and dispatch rules -Every compute entry point of `cake_backend` takes the keyword-only `deterministic: bool = True` and runs only with -`deterministic=False`. The entry points: `dense_projection_gemm`, `plan_dense_projection_gemm`, -`prepare_dense_projection_gemm`, `projection_forward`, `projection_dgrad`, `prepare_projection_wgrad`, -`projection_wgrad`, `plan_router_fp32_gemm`, `prepare_router_fp32_gemm`, `router_fp32_gemm`, `router_forward`, -`router_dgrad`, `router_wgrad`. +A call runs the routes of `cake_backend.ROUTES` (table below). No entry point consults the PyTorch +deterministic-algorithms setting to choose a program. Under `torch.use_deterministic_algorithms(True)` a launch whose +result can vary run to run raises PyTorch's alert right before it launches: -| call | `deterministic=False` | any other value (the default, `True`, `None`, `0`, `1`, `"False"`, a NumPy bool) | under `torch.use_deterministic_algorithms(True)` (strict or `warn_only=True`) | +| launch | rows | order of the partials | under `torch.use_deterministic_algorithms(True)` | |---|---|---|---| -| this package's compute entry points (above) | runs this package's routes (table below) | `ValueError` naming the deterministic entry, before anything is planned, allocated or launched | the same routes and programs. Right before a launch whose result can vary run to run (the K1 split-K program of CTA pairs at three or more splits; a K2 stream-K partition that gives an element more than two partials), PyTorch's alert: `RuntimeError` (nothing launched, `out` untouched), or a `UserWarning` and then the launch under `warn_only=True`. Fixed-order launches run without an alert | -| this package's helpers (`router_symbol`, `bind_launch`, the planners `plan_atomic_gemm` / `plan_atomic_router`, the rule tables, `clear_convenience_caches`, `use_deterministic_backend`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged | -| `flashinfer.experimental.dense_projection_det_gemm` (`projection_det_wgrad`, `router_fp32_det_gemm`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged: every launch is fixed-order | - -*Table: the mode rule per call and value.* - -The `ValueError` message (one line, wrapped here), for example: - -```text -projection_wgrad: flashinfer.experimental.dense_projection_gemm runs the non-deterministic package only and needs -deterministic=False (got True): pass deterministic=False for the non-deterministic path; the deterministic entry is -flashinfer.experimental.dense_projection_det_gemm.cake_backend.projection_det_wgrad -``` - -Rules: - -- The `deterministic` keyword never selects the deterministic package, and no call falls back to it. For a - deterministic result, call the deterministic package's entry point (`projection_det_wgrad`, ...) yourself. -- Every call states the mode: pass `deterministic=False` on every compute entry point call, the `plan_*` and - `prepare_*` entry points included. -- A `deterministic=False` call reuses the deterministic package's fixed-order programs on the rows the table below - marks "deterministic program". It launches them through this package's binding, or, on such a row, through the - deterministic package's entry point of the same operation where the binding cannot rebuild the call's launch views - (two operands that alias) and for a batched forward on that package's swapped small-token program: the same - program, the same bits. On those rows `plan_*` / `prepare_*` return the deterministic package's plan or prepared - launch (`DetGemmPlan`, `PreparedDetGemm`, `RouterDetPlan`, `PreparedRouterDetGemm`). A row that runs a program - of this package never goes to the deterministic package. -- A result of this package carries no run-to-run bitwise guarantee on any row, also on a row that runs a - fixed-order program today. The "run-to-run bits" rows below describe the current programs, not a contract. -- The PyTorch setting never changes which program a call runs. +| K1 CTA-pair split-K program, three or more splits | the `sm_107a` rows of `K1_SPLIT_RULES` (216 SMs) and `K1_SPLIT_RULES_BY_SM` (208 / 212 SMs); a caller's `splits >= 3` | the order the splits finish | `RuntimeError` before the launch (nothing launched, `out` untouched); with `warn_only=True` a `UserWarning`, then the launch | +| K1 CTA-pair split-K program, two splits | every `sm_100a` / `sm_103a` row, `sm_107a` `8192 -> 128` | one add of two values (repeats bit for bit) | no alert | +| K1 one split, the chain programs (`_c3`) | the one-split rows | whole tiles, chains added in order | no alert | +| K1 cluster program `dense_proj_gemm_dsm_*_s2` / `_s3` | `K1_DSM_RANGES_BY_SM` (152-SM `sm_100a` parts, FP32 out) | cluster-rank order in distributed shared memory (no global atomics, no workspace) | no alert | +| K2 stream-K program, more than two partials per element | the `ROUTER_RULES` partitions, the `"waves"` windows, a caller's `splits >= 3` on `nn_t` | the order the partials finish (TMA reduce-add) | `RuntimeError` before the split, the fill of `out` and the program; with `warn_only=True` a `UserWarning`, then the launch | +| K2 stream-K program, at most two partials per element | a forward's `splits` of 1 or 2 | `0 + a + b == 0 + b + a` | no alert | +| K2 SIMT programs (forward, input and weight gradient), `router_partials_sum` | their token windows | fixed: every element stored once | no alert | +| the deterministic programs, the BF16x3 split | the rows routed to the deterministic program (below) | fixed | no alert | + +*Table: the alert rule per launch.* + +- The alert is checked on every launch, prepared and convenience path alike, never at plan or prepare time. The + forward and input-gradient entry points never alert. +- CUDA graphs: the alert is decided when the launch is captured; a replay does not check the setting again. +- The setting never changes which program a call runs, and no call goes to the deterministic package because of it. + For a bitwise reproducible result, call `flashinfer.cake_dense_gemm_det_train`. +- The alert text starts with `flashinfer.cake_dense_gemm_nondet_train: does not have a deterministic implementation` and names + `flashinfer.cake_dense_gemm_det_train` as the deterministic module. + +On some rows this module runs the deterministic package's fixed-order programs: + +| case | where | result | +|---|---|---| +| a layout routed `"det"` (`ROUTES`), or a row the planner keeps on the deterministic program (no verdict) | `cake_dense_gemm_nondet_train_projection_plan`, `cake_dense_gemm_nondet_train_projection_prepare`, `cake_dense_gemm_nondet_train_projection_wgrad_prepare`, `cake_dense_gemm_nondet_train_router_plan`, `cake_dense_gemm_nondet_train_router_prepare` | the deterministic package's plan or prepared launch (`CakeDenseGemmDetProjectionPlan`, `CakeDenseGemmDetProjectionLaunch`, `CakeDenseGemmDetRouterPlan`, `CakeDenseGemmDetRouterLaunch`) | +| the same rows on the convenience path | the convenience entry points | the deterministic plan and program through this package's binding (the same bits, fewer host operations per call) | +| a row routed to the deterministic program whose launch views this package's binding cannot rebuild from the call's tensors (two operands that alias), and a batched forward that the deterministic package runs on its swapped small-token program | the K1 convenience entry points | the deterministic package's entry point of the same operation (the same program, the same bits) | +| tensors on different CUDA devices, or a cached entry without a binding | the K1 convenience path of a row routed to the deterministic program | the same | + +*Table: when this module returns a deterministic object or calls a deterministic entry point.* + +- A row that runs a program of this package never goes to the deterministic package: when its views cannot be + bound, it takes this package's unbound launch. +- A result of this module carries no run-to-run bitwise guarantee on any row, also on a row that runs a fixed-order + program today. The "run-to-run bits" rows below describe the current programs, not a contract. - The tensors passed between the forward and the gradients carry no mode: recompute a checkpointed forward with the - package that made the saved one (same package only). Neither package has an autograd Function or saved tape. + module that made the saved one (same package only). Neither package has an autograd Function or a saved tape. - This package imports the deterministic package at run time, as its sibling (a vendored copy binds it with - `cake_backend.use_deterministic_backend`). This package re-exports no name of the deterministic package. + `cake_backend.use_deterministic_backend`). It re-exports no name of the deterministic package. -| operation | `deterministic=False` | +| operation | route | |---|---| -| K1 forward (`projection_forward`, `kk` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order; a recomputed forward reproduces its result). The rows of `K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM `sm_100a` / `sm_103a` parts, the K1 forward table below): the deterministic package's single-CTA program, every tile's whole K in one CTA -- other bits than the deterministic package's plan, the same on every run | -| K1 input gradient (`projection_dgrad`, `kn` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order) | -| K1 weight gradient (`projection_wgrad`, `nn` layouts) | the rows of `K1_SPLIT_RULES`: split-K over the ragged token count -- every split writes its FP32 partial tile to a workspace slot; the last split of a tile to arrive (a per-tile ticket counter) adds the other slots in arrival order and rounds once into `out`, then returns the counter to zero. Before every K1 table, the FP32 rows and token ranges of `K1_DSM_RANGES_BY_SM` (152-SM `sm_100a` parts, the 152-SM cluster table below): the cluster program `dense_proj_gemm_dsm_*` -- a tile's K split over the `S` CTAs of a thread-block cluster, each CTA's FP32 partial sent to the CTAs that own its columns through distributed shared memory, the `S` partials added in cluster-rank order and stored once (no workspace, no ticket: the same bits on every run). Every other row: deterministic program | -| K2 forward (`router_forward`, `kk`) | the token windows of its `ROUTER_RULES` rows (152-SM `sm_100a` / `sm_103a` parts, 256 experts, hidden 4096, 5120 and 6144, 1 to 192 tokens; `sm_100a` hidden 4096: 1 to 160): the exact-FP32 SIMT program `router_fp32_fwd_simt_c8_*` -- `K` split over the eight CTAs of a thread-block cluster, FFMA chains of the FP32 operands read in place (up to 16 tokens in a compensated form), the eight partials added in cluster-rank order through distributed shared memory, every element of `out` stored once. Every other row, token count and part: deterministic program (its four splits); an explicit `splits` of 1 or 2 runs this package's stream-K program (at most two partials per element) | -| K2 input gradient (`router_dgrad`, `kn`) | the token windows of its `ROUTER_RULES` rows (on 212-SM `sm_107a` and 148-SM `sm_100a` / `sm_103a` parts `ROUTER_RULES_BY_SM`): the exact-FP32 SIMT program `router_fp32_dgrad_simt_*` -- one FFMA per product of the FP32 operands read in place, every element of `out` stored once. Every other row, token count and part: deterministic program (one split) | -| K2 weight gradient (`router_wgrad`, `nn_t`) | the SIMT token windows of its rows in the router tables (`ROUTER_RULES`; on 212-SM `sm_107a` and 148-SM `sm_100a` / `sm_103a` parts `ROUTER_RULES_BY_SM`): the exact-FP32 SIMT program `router_fp32_gemm_simt_*` on the CUDA cores, every element of `out` stored once; on 148-SM parts at one token at hidden 4096 its compensated form. The rows of `ROUTER_RULES` with a partition: stream-K over the K blocks, in spans alike or in long head spans and short tail spans (the head/tail program) -- whole tiles store the one-split result, tiles split between units TMA reduce-add into `out` (zeroed before the program, on the same stream). Every other row: deterministic program (its six splits) | +| K1 forward (`cake_dense_gemm_nondet_train_projection_forward`, `kk` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order; a recomputed forward reproduces its result). The rows of `K1_FORWARD_RULES_BY_SM` (exact token counts on 152-SM `sm_100a` / `sm_103a` parts, the K1 forward table below): the deterministic package's single-CTA program, every tile's whole K in one CTA -- other bits than the deterministic package's plan, the same on every run | +| K1 input gradient (`cake_dense_gemm_nondet_train_projection_dgrad`, `kn` layouts) | deterministic program (whole output tiles, or the deterministic plan's two-way split of K with the partials added in a fixed order) | +| K1 weight gradient (`cake_dense_gemm_nondet_train_projection_wgrad`, `nn` layouts) | the rows of `K1_SPLIT_RULES`: split-K over the ragged token count -- every split writes its FP32 partial tile to a workspace slot; the last split of a tile to arrive (a per-tile ticket counter) adds the other slots in arrival order and rounds once into `out`, then returns the counter to zero. Before every K1 table, the FP32 rows and token ranges of `K1_DSM_RANGES_BY_SM` (152-SM `sm_100a` parts, the 152-SM cluster table below): the cluster program `dense_proj_gemm_dsm_*` -- a tile's K split over the `S` CTAs of a thread-block cluster, each CTA's FP32 partial sent to the CTAs that own its columns through distributed shared memory, the `S` partials added in cluster-rank order and stored once (no workspace, no ticket: the same bits on every run). Every other row: deterministic program | +| K2 forward (`cake_dense_gemm_nondet_train_router_forward`, `kk`) | the token windows of its `ROUTER_RULES` rows (152-SM `sm_100a` / `sm_103a` parts, 256 experts, hidden 4096, 5120 and 6144, 1 to 192 tokens; `sm_100a` hidden 4096: 1 to 160): the exact-FP32 SIMT program `router_fp32_fwd_simt_c8_*` -- `K` split over the eight CTAs of a thread-block cluster, FFMA chains of the FP32 operands read in place (up to 16 tokens in a compensated form), the eight partials added in cluster-rank order through distributed shared memory, every element of `out` stored once. Every other row, token count and part: deterministic program (its four splits); an explicit `splits` of 1 or 2 runs this package's stream-K program (at most two partials per element) | +| K2 input gradient (`cake_dense_gemm_nondet_train_router_dgrad`, `kn`) | the token windows of its `ROUTER_RULES` rows (on 212-SM `sm_107a` and 148-SM `sm_100a` / `sm_103a` parts `ROUTER_RULES_BY_SM`): the exact-FP32 SIMT program `router_fp32_dgrad_simt_*` -- one FFMA per product of the FP32 operands read in place, every element of `out` stored once. Every other row, token count and part: deterministic program (one split) | +| K2 weight gradient (`cake_dense_gemm_nondet_train_router_wgrad`, `nn_t`) | the SIMT token windows of its rows in the router tables (`ROUTER_RULES`; on 212-SM `sm_107a` and 148-SM `sm_100a` / `sm_103a` parts `ROUTER_RULES_BY_SM`): the exact-FP32 SIMT program `router_fp32_gemm_simt_*` on the CUDA cores, every element of `out` stored once; on 148-SM parts at one token at hidden 4096 its compensated form. The rows of `ROUTER_RULES` with a partition: stream-K over the K blocks, in spans alike or in long head spans and short tail spans (the head/tail program) -- whole tiles store the one-split result, tiles split between units TMA reduce-add into `out` (zeroed before the program, on the same stream). Every other row: deterministic program (its six splits) | | BF16x3 split of the K2 second operand (inside the K2 launches) | the deterministic stage (exact) | -*Table: the routing of `deterministic=False` (`cake_backend.ROUTES`), per operation and row.* +*Table: the routing of this module (`cake_backend.ROUTES`), per operation and row.* ## Numerics, architectures and parallelism @@ -77,22 +102,22 @@ Rules: | accuracy | the deterministic package's budgets against FP64 of the exact operands, worst over repeated runs: BF16 out `atol = rtol = 1e-2`; FP32 out relative Frobenius error `<= 5e-5` and max error `<= 1e-3 * max abs(ref)`; router `<= 2e-6` relative Frobenius error and no row whose top-8 experts differ from FP64 (outside near ties of the 8th and 9th logits). At the largest token counts some thin K1 rows miss the BF16 (FP32) budget with the deterministic program and with this package's program alike; there the row keeps its route. The deterministic program: BF16 output of `sm_100a` `4096 -> 32`, `4096 -> 48`, `4096 -> 64`, `5120 -> 32`, `5120 -> 48` and of `sm_103a` `4096 -> 32`, `4096 -> 48`, `5120 -> 32`, `5120 -> 48`, `5120 -> 64`, `5120 -> 128` (from 45058, 49153, 65536 or 65537 tokens up by row: the K1 token tables below), and FP32 output of the rows of the odd-block K1 table at an odd number of K blocks from 98305 tokens up. This package's program: BF16 output of `sm_100a` `4096 -> 64` from about 43009 to 65535 tokens and `5120 -> 32` at 47105 to 65536 tokens, of `sm_103a` `4096 -> 32` from about 49153 to 65535 tokens, and on 148-SM parts of `sm_103a` `5120 -> 48` from 49152 tokens, `5120 -> 64` at 47104 and from 49153 tokens and `5120 -> 128` from 65537 tokens | | rounding | each output element is rounded once, from its complete FP32 sum | | subnormals | kept: K1 adds the partials with IEEE round-to-nearest adds, K2 reduces with the TMA reduce-add; a K2 SIMT window (forward, input and weight gradient) accumulates with round-to-nearest FMAs and adds, merges the sums of a CTA's slices (and in a compensated form their compensations) in a fixed order through shared memory (the forward: the sums of a cluster's eight CTAs, through distributed shared memory) and stores every element of `out` once, with no reduce-add; nothing flushes to zero | -| `torch.use_deterministic_algorithms(True)` | PyTorch's alert right before an order-varying launch: the K1 CTA-pair program at three or more splits (in the shipped tables only `sm_107a` rows, see "run-to-run bits, K1"; or a caller's `splits >= 3`), K2 stream-K with more than two partials per element (the `ROUTER_RULES` partitions, the `"waves"` windows, a caller's `splits >= 3` on `nn_t`). `RuntimeError`, or a `UserWarning` with `warn_only=True` and then the launch. The alert is checked on every launch (prepared and convenience), never at plan or prepare time. The two-split K1 rows, the K1 cluster program (`K1_DSM_RANGES_BY_SM`, its splits added in cluster-rank order), the K2 forward, the K2 SIMT windows and every deterministic program run without an alert | +| `torch.use_deterministic_algorithms(True)` | PyTorch's alert right before an order-varying launch: the K1 CTA-pair program at three or more splits (in the shipped tables only `sm_107a` rows, see "run-to-run bits, K1"; or a caller's `splits >= 3`), K2 stream-K with more than two partials per element (the `ROUTER_RULES` partitions, the `"waves"` windows, a caller's `splits >= 3` on `nn_t`). `RuntimeError`, or a `UserWarning` with `warn_only=True` and then the launch. The alert is checked on every launch (prepared and convenience), never at plan or prepare time. A launch captured into a CUDA graph is checked at capture, not at replay. The two-split K1 rows, the K1 cluster program (`K1_DSM_RANGES_BY_SM`, its splits added in cluster-rank order), the K2 forward, the K2 SIMT windows and every deterministic program run without an alert | | run-to-run bits, K1 | two splits (every `sm_100a` and `sm_103a` row, `sm_107a` `8192 -> 128`): repeat bit for bit (one add of two values); the cluster program (`K1_DSM_RANGES_BY_SM`, two or three splits): repeats bit for bit (the splits added in cluster-rank order); one split (`sm_107a` rows marked `p = 1` below): repeat bit for bit (whole tiles); three or four splits (the other `sm_107a` rows): vary in the last bits | | run-to-run bits, K2 | forward: in a SIMT window repeats bit for bit (every element stored once, the lanes and the cluster's CTAs combined in a fixed order), else the deterministic program's bits; weight gradient on a `ROUTER_RULES` partition: varies in the last bits (more than two partials per element); a SIMT window (weight or input gradient): repeats bit for bit (every element stored once, its slices merged in a fixed order); with an explicit forward `splits` of 1 or 2: repeats bit for bit (`0 + a + b == 0 + b + a`), so recomputed router logits and top-k match | | architectures | `sm_100a`, `sm_103a`, `sm_107a`: one `-gencode=arch=compute_,code=sm_` build per program, minimum CUDA 12.8 / 12.9 / 13.4 (`build.json`) | | streams and threads | planning takes a module lock; the convenience path keeps one workspace arena per (device, stream); prepared launches own their workspace | -| CUDA graphs | prepared launches allocate nothing on the device and capture into a CUDA graph (K2: the split, the memset of `out`, the program); no host synchronization. A `deterministic=False` convenience-path call that this package launches from its arenas (its own programs, and the deterministic programs it launches through its binding, see [Workspace](#workspace)), captured into a CUDA graph after a warm-up call on the capturing stream (so the capture plans and allocates nothing), stays valid until `clear_convenience_caches()` or `use_deterministic_backend`, which invalidate it: the arenas keep every buffer a launch was bound to, through growths and cache evictions. A call on a row that runs the deterministic program, which this package hands to the deterministic package's entry point of the same operation -- operands whose launch views the binding cannot rebuild from the call's tensors (for example two operands that alias), a batched forward on that package's swapped small-token program -- follows that package's rule (capture its prepared launch, not its convenience path). The graphs captured on one stream share its arena (`torch.cuda.graph` without `stream=` captures every graph on one shared stream): replay them one at a time, and not concurrently with calls on that stream | +| CUDA graphs | prepared launches allocate nothing on the device and capture into a CUDA graph (K2: the split, the memset of `out`, the program); no host synchronization. A convenience-path call that this package launches from its arenas (its own programs, and the deterministic programs it launches through its binding, see [Workspace](#workspace)), captured into a CUDA graph after a warm-up call on the capturing stream (so the capture plans and allocates nothing), stays valid until `cake_dense_gemm_nondet_train_convenience_cache_clear()` or `use_deterministic_backend`, which invalidate it: the arenas keep every buffer a launch was bound to, through growths and cache evictions. A call on a row that runs the deterministic program, which this package hands to the deterministic package's entry point of the same operation -- operands whose launch views the binding cannot rebuild from the call's tensors (for example two operands that alias), a batched forward on that package's swapped small-token program -- follows that package's rule (capture its prepared launch, not its convenience path). The graphs captured on one stream share its arena (`torch.cuda.graph` without `stream=` captures every graph on one shared stream): replay them one at a time, and not concurrently with calls on that stream | | other kernels on the device | the programs walk their work persistently over at most `sm_count // 2` CTA pairs and wait on no CTA that is not resident (the forward's SIMT program waits only on the CTAs of its own cluster, which the hardware schedules together); every routed row was measured next to a kernel holding 8, 16 or 32 SMs on another stream (no more than 3 % slower than the deterministic plan there) | | context parallelism (KCP) | not applicable: a local operator on the caller's CP-local shards; no collectives; the reduction of parameter gradients across ranks stays with the caller | | autograd, checkpointing, gradient accumulation | none in either package: callers compose the forward and both gradients; `out` is overwritten, not accumulated into | -*Table: the numerical, architecture and parallelism matrix of `deterministic=False`.* +*Table: the numerical, architecture and parallelism matrix of this module.* ## Planners and rule tables -A planner decides, per row and device, whether `deterministic=False` runs this package's program and with which -parameters; a row it leaves runs the deterministic program (`plan_dense_projection_gemm` / `plan_router_fp32_gemm` +A planner decides, per row and device, whether a call runs this package's program and with which +parameters; a row it leaves runs the deterministic program (`cake_dense_gemm_nondet_train_projection_plan` / `cake_dense_gemm_nondet_train_router_plan` say which). Only measured rows carry a program: a row runs this package's program where that program was faster than the deterministic package's shipped plan of the row in every measurement run: a one-sided paired sign test (p < 0.05) with a median time ratio below 0.99, both in eager launches and in CUDA graph replays, and no more than @@ -731,7 +756,7 @@ program's) and 1.53e-7 in the FFMA windows (x1.11-1.27), all values finite, and FP64 outside near ties of the 8th and 9th logits. cuBLAS FP32 (cuBLAS 13.1.1; per token count the faster of `torch.mm` without TF32 and `CUBLAS_FP32_EMULATED_BF16X9_MATH`, both within the router budget) took x0.83-2.78 of the program's time: faster than the program at one token and at some token counts from 160 to 192 at hidden 4096 and 6144. The deterministic -package's `router_forward` took x1.17-4.88 of it in CUDA graphs and x1.56-6.53 in eager launches. Above 192 tokens the +package's `cake_dense_gemm_det_train_router_forward` took x1.17-4.88 of it in CUDA graphs and x1.56-6.53 in eager launches. Above 192 tokens the rows keep the deterministic program (not measured under this rule). Other SM counts, other rows and `sm_107a` run the deterministic program.* @@ -898,7 +923,7 @@ operands (10 seeds each); the bits of the forward's program depend on `K`, the k covers both tiles of its kind. The head/tail partitions split a tile's K range over more spans than the deterministic plan's six and add them with the TMA reduce-add.* -The K1 forward rows of `K1_FORWARD_RULES_BY_SM` run, with `deterministic=False`, the deterministic package's +The K1 forward rows of `K1_FORWARD_RULES_BY_SM` run the deterministic package's single-CTA program (`K1_FORWARD_SINGLE_CTA_KNOBS`: one 128-row tile per CTA, 128-wide column tiles, no TMA cache hints; every tile's whole K in one CTA) in place of that package's plan of the row (CTA pairs of 256-wide tiles with the K range split two ways, the partials added in a fixed order). The forward stays deterministic: it repeats bit for bit @@ -912,7 +937,7 @@ slower than with the deterministic plan: the median over every run of the paired eager launches and graph replays separately, and no run above the largest ratio of the deterministic plan against itself in the same runs. The backward pass alone runs the same programs in both and is reported, not judged. Every other token count, SM count and row was not measured or missed, and runs the deterministic plan, as do the -generic `dense_projection_gemm`, `prepare_dense_projection_gemm` and `plan_dense_projection_gemm` (their caller +generic `cake_dense_gemm_nondet_train_projection`, `cake_dense_gemm_nondet_train_projection_prepare` and `cake_dense_gemm_nondet_train_projection_plan` (their caller passes the planner knobs). | K1 forward row `d_in -> d_out` (BF16 out, both operands K-major) | 152 SMs: token counts (`K1_FORWARD_RULES_BY_SM`) | measured token counts outside them | @@ -931,8 +956,8 @@ with a BF16 weight gradient 1.007 (the deterministic plan against itself at most forward + backward 1.015 (at most 1.013); `sm_103a` `5120 -> 512` 2 x forward + backward with an FP32 weight gradient 1.015 (at most 1.012).* -`prepare_*` and `plan_*` take `splits` (and, for K1, `block_n`) to force a plan, for tests and measurements; the K2 -forward takes at most two splits. `benchmarks/bench_cake_dense_projection_gemm.py` prints, per row, the +The `*_prepare` and `*_plan` entry points take `splits` (and, for K1, `block_n`) to force a plan, for tests and measurements; the K2 +forward takes at most two splits. `benchmarks/bench_cake_dense_gemm_nondet_train.py` prints, per row, the `rule_key` / `rule_value` of the plan it measured, in the format of these tables. ## Workspace @@ -942,15 +967,15 @@ zero at creation: the last split of a partial tile reads the padding columns no the ticket counters (`tiles * 512` words, zero at creation, zero again after every complete launch). A prepared launch owns both; the convenience path keeps them in an arena per (device, stream), so concurrent launches on different streams never share them. The cluster program (`K1_DSM_RANGES_BY_SM`) sums its splits in shared memory: -no workspace, no tickets, no arena. An arena never frees a buffer before `clear_convenience_caches()`: a growth (to +no workspace, no tickets, no arena. An arena never frees a buffer before `cake_dense_gemm_nondet_train_convenience_cache_clear()`: a growth (to at least twice the buffer's size) retires the buffer it replaces, and an entry evicted from the caches -(`CONVENIENCE_CACHE_SIZE`) leaves its arena, so a launch captured into a CUDA graph keeps its workspace; the retired +(the deterministic backend's `CONVENIENCE_CACHE_SIZE`) leaves its arena, so a launch captured into a CUDA graph keeps its workspace; the retired buffers of a kind hold fewer bytes than its live one, which holds less than twice the largest requirement seen or the 16-element minimum, so a kind holds at most four times that requirement or 16 elements, whichever is more. An arena -stays until `clear_convenience_caches()` even when no cached entry uses it, one per (device, stream) that made a call: +stays until `cake_dense_gemm_nondet_train_convenience_cache_clear()` even when no cached entry uses it, one per (device, stream) that made a call: a process that makes calls on many short-lived streams (`torch.cuda.ExternalStream` handles) calls it once no captured graph replays. After an aborted launch, call -`PreparedAtomicGemm.reset_counters()` (or `clear_convenience_caches()`). A K2 launch with part tiles zeroes `out` on +`CakeDenseGemmNondetProjectionLaunch.reset_counters()` (or `cake_dense_gemm_nondet_train_convenience_cache_clear()`). A K2 launch with part tiles zeroes `out` on the launch stream (one memset) between the split and the program; prepared launches of both programs allocate nothing and can be captured in a CUDA graph. A prepared K2 launch binds everything but the stream when prepared (its FFI arguments are converted once), so a call @@ -972,7 +997,7 @@ partials in this package's arena, bound once per arena growth, and the entry poi views. It issues the same launches as the deterministic package's convenience path in fewer host operations. Whether a row has a verdict is kept per device and view geometry, and whether any router table gives the forward one is kept too, until a router table changes: the router tables count their changes, so a verdict set or removed at run time is seen at the next call (a table rebound to a plain `dict` is read on every -call). The router entry points (`router_forward`, `router_dgrad`, `router_wgrad`) cache each call once more by their +call). The router entry points (`cake_dense_gemm_nondet_train_router_forward`, `cake_dense_gemm_nondet_train_router_dgrad`, `cake_dense_gemm_nondet_train_router_wgrad`) cache each call once more by their own tensors, the operation, the split count and the PDL flag: a call goes from that key to the cached launch its route takes, without building the transposed operand, routing the row or building the launch's own key; the key is routed again when a router table or `ROUTES` changes (with a table rebound to a plain `dict`, every call takes the @@ -980,34 +1005,17 @@ full path). ## Benchmark -`benchmarks/bench_cake_dense_projection_gemm.py` times every row of the deterministic family's benchmark with the -arms `det` (the deterministic package's prepared launch), `atomic` (`deterministic=False`), `atomic_p1` (the same +`benchmarks/bench_cake_dense_gemm_nondet_train.py` times every row of the deterministic family's benchmark with the +arms `det` (the deterministic package's prepared launch), `atomic` (this module), `atomic_p1` (the same program at one split), `cublas` (`torch.matmul`), `cublaslt` (a cuBLASLt heuristic sweep) and `bf16x9` (cuBLAS FP32 in `CUBLAS_FP32_EMULATED_BF16X9_MATH`, the router rows), in paired groups, eager and in CUDA graphs; `--resident` re-times `det` and `atomic` next to a kernel holding `--resident-sms` SMs on another stream. -## Names - -This package has no alias, re-export or read-through of a deterministic name. Each name below is this package's -own object. The deterministic package's names (`*_det_*`, `DetGemmPlan`, `ROW_RULES`, `stream_k_plan`, -`wgrad_views`, ...) are reached only in `flashinfer.experimental.dense_projection_det_gemm`; reading one of them -here raises `AttributeError`. - -| operation | this package (`deterministic=False` only) | deterministic package (no `deterministic` keyword) | -|---|---|---| -| K1 GEMM | `dense_projection_gemm`, `plan_dense_projection_gemm`, `prepare_dense_projection_gemm` | `dense_projection_det_gemm`, `plan_dense_projection_det_gemm`, `prepare_dense_projection_det_gemm` | -| K1 forward, input gradient, weight gradient | `projection_forward`, `projection_dgrad`, `projection_wgrad`, `prepare_projection_wgrad` | `projection_det_forward`, `projection_det_dgrad`, `projection_det_wgrad`, `prepare_projection_det_wgrad` | -| K2 GEMM | `router_fp32_gemm`, `plan_router_fp32_gemm`, `prepare_router_fp32_gemm` | `router_fp32_det_gemm`, `plan_router_fp32_det_gemm`, `prepare_router_fp32_det_gemm` | -| K2 forward, input gradient, weight gradient | `router_forward`, `router_dgrad`, `router_wgrad` | `router_det_forward`, `router_det_dgrad`, `router_det_wgrad` | -| BF16x3 split | -- (inside the K2 launches) | `split_fp32_to_bf16x3_det_` | -| plan and prepared types | `AtomicGemmPlan`, `PreparedAtomicGemm`, `AtomicRouterPlan`, `PreparedAtomicRouterGemm` (a row on the deterministic program returns that package's types, see [Mode and dispatch rules](#mode-and-dispatch-rules)) | `DetGemmPlan`, `PreparedDetGemm`, `RouterDetPlan`, `PreparedRouterDetGemm` | -| program templates | `dense_proj_gemm_atomic_nn_n<64\|128\|256>[_f32]_t`, `dense_proj_gemm_atomic_nn_n<128\|256>_f32_t_c3`, `dense_proj_gemm_atomic_nn_n192_f32[_c3]` (row-major store, `sm_107a`), `dense_proj_gemm_dsm_nn_n<32\|64\|128>_f32_t_s<2\|3>` (the K1 cluster programs, `sm_100a`), `router_fp32_gemm_atomic_{kk,nn_t,nn_t_ht_c3_nreg}`, `router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nreg`, `router_fp32_gemm_simt_*`, `router_fp32_dgrad_simt_*`, `router_fp32_fwd_simt_*`, `router_partials_sum`; `router_symbol(a_mn, b_mn, out_t)` names the stream-K template of `kk` / `nn_t` | `dense_proj_det_gemm_*`, `router_fp32_det_gemm_`, `split_fp32_det_bf16x3`, ... | - -*Table: the entry points of the two packages, per operation.* +## Helpers of this package -| name (`flashinfer.experimental.dense_projection_gemm`) | reaches | +| name (`flashinfer.experimental.cake.cake_dense_gemm_nondet_train`) | reaches | |---|---| -| `cake_backend.clear_convenience_caches` | clears the caches and arenas of both packages (invalidates the CUDA graphs that captured a convenience-path call) | +| `cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear` | clears the caches and arenas of both packages (invalidates the CUDA graphs that captured a convenience-path call) | | `cake_backend.bind_launch` | binds a program of this package (a deterministic program binds through that package's `bind_launch`) | | `cake_jit.{MODULES, KERNELS, ARG_PLANS}`, `cake_jit.select_module`, `cake_jit.load_cake_dense_projection_gemm_module` | this package's programs only (a deterministic template or program name raises; use the deterministic package's `cake_jit`) | | `cake_backend.use_deterministic_backend`, `cake_backend.deterministic_backend` | bind and return the deterministic package's `cake_backend` (the package whose fixed-order programs this package reuses) | @@ -1023,7 +1031,7 @@ the deterministic package's names there. | file | content | |---|---| | `__init__.py` | imports nothing; `cake_backend`, `cake_jit`, `cake_build` load on first attribute access | -| `cake_backend.py` | entry points, mode rules, routing, the planners, rule tables and the convenience state of this package's programs | +| `cake_backend.py` | entry points, routing, the alert, the planners, rule tables and the convenience state of this package's programs | | `cake_jit.py` | the registry of this package's programs (`ARG_PLANS`, `MODULES`, `KERNELS`) and the loader | | `cake_build.py`, `build.json`, `csrc/cake_dense_projection_gemm/` | the `cake-build/1` loader, manifest and generated sources of this package's programs (written by the Cake exporter) | diff --git a/tests/experimental/test_cake_dense_projection_gemm.py b/tests/experimental/test_cake_dense_gemm_nondet_train.py similarity index 91% rename from tests/experimental/test_cake_dense_projection_gemm.py rename to tests/experimental/test_cake_dense_gemm_nondet_train.py index 67252db9f6e..af39df8b3d7 100644 --- a/tests/experimental/test_cake_dense_projection_gemm.py +++ b/tests/experimental/test_cake_dense_gemm_nondet_train.py @@ -14,22 +14,22 @@ limitations under the License. """ -# The non-deterministic dense projection GEMM package: the mode rule of every compute entry point (only -# ``deterministic=False``; any other value raises ``ValueError`` naming the deterministic entry point), the routes of -# ``cake_backend.ROUTES``, PyTorch's deterministic-algorithms alert before an order-varying launch, the absence of -# every deterministic name (no alias, map or read-through), the K1 and -# K2 planners (the per-architecture planners, the rule tables, the token-count guard, the stream-K partition), the -# registry and file names of this package's programs and its workspace kinds. The CPU tests launch nothing. The GPU -# tests (``test_gpu_*``) need a device the deterministic package registers programs for; the ones of this package's -# programs force the split count (they do not depend on what the planner takes on the device), assert by kernel -# name that the program ran, and skip only on a checkout that registers no program of this package for the device. +# The non-deterministic Cake dense GEMMs (the public module ``flashinfer.cake_dense_gemm_nondet_train``, backend +# ``flashinfer.experimental.cake.cake_dense_gemm_nondet_train``): the routes of ``cake_backend.ROUTES``, PyTorch's +# deterministic-algorithms alert before an order-varying launch, the K1 and K2 planners (the per-architecture +# planners, the rule tables, the token-count guard, the stream-K partition), the registry and file names of this +# package's programs and its workspace kinds. The public names (the mode word, no mode keyword, the former names and +# paths gone, the lazy import) are checked by ``test_cake_dense_gemm_nondet_train_names.py``. The CPU tests launch +# nothing. The GPU tests (``test_gpu_*``) need a device the deterministic package registers programs for; the ones of +# this package's programs force the split count (they do not depend on what the planner takes on the device), assert +# by kernel name that the program ran, and skip only on a checkout that registers no program of this package for the +# device. import ast import contextlib import dataclasses import functools import gc -import inspect import json import os import re @@ -46,11 +46,11 @@ from tests.test_helpers.cake_package import cake_package # Backend suite: the packages are named through ``cake_package``; nothing else of flashinfer is imported here. -package = cake_package("dense_projection_gemm") -cake_backend = cake_package("dense_projection_gemm", "cake_backend") -cake_jit = cake_package("dense_projection_gemm", "cake_jit") -det_backend = cake_package("dense_projection_det_gemm", "cake_backend") -det_jit = cake_package("dense_projection_det_gemm", "cake_jit") +package = cake_package("cake_dense_gemm_nondet_train") +cake_backend = cake_package("cake_dense_gemm_nondet_train", "cake_backend") +cake_jit = cake_package("cake_dense_gemm_nondet_train", "cake_jit") +det_backend = cake_package("cake_dense_gemm_det_train", "cake_backend") +det_jit = cake_package("cake_dense_gemm_det_train", "cake_jit") # A vendored run imports both packages under other names: bind the deterministic package explicitly (in a # FlashInfer tree this names the sibling the package resolves by itself). cake_backend.use_deterministic_backend(det_backend) @@ -63,105 +63,8 @@ RERUNS = 20 SEED = 20261031 -# The compute entry points of ``cake_backend``: each takes the keyword-only ``deterministic: bool = True`` and accepts -# only ``deterministic=False``. -ENTRY_POINTS = ( - "dense_projection_gemm", - "plan_dense_projection_gemm", - "plan_router_fp32_gemm", - "prepare_dense_projection_gemm", - "prepare_projection_wgrad", - "prepare_router_fp32_gemm", - "projection_dgrad", - "projection_forward", - "projection_wgrad", - "router_dgrad", - "router_forward", - "router_fp32_gemm", - "router_wgrad", -) -# name of this package (or a former name) -> the deterministic package's canonical name -FORMER_TO_DET = { - "GemmPlan": "DetGemmPlan", - "PreparedGemm": "PreparedDetGemm", - "PreparedRouterGemm": "PreparedRouterDetGemm", - "RouterPlan": "RouterDetPlan", - "dense_projection_gemm": "dense_projection_det_gemm", - "plan_dense_projection_gemm": "plan_dense_projection_det_gemm", - "plan_router_fp32_gemm": "plan_router_fp32_det_gemm", - "prepare_dense_projection_gemm": "prepare_dense_projection_det_gemm", - "prepare_projection_wgrad": "prepare_projection_det_wgrad", - "prepare_router_fp32_gemm": "prepare_router_fp32_det_gemm", - "projection_dgrad": "projection_det_dgrad", - "projection_forward": "projection_det_forward", - "projection_wgrad": "projection_det_wgrad", - "router_dgrad": "router_det_dgrad", - "router_forward": "router_det_forward", - "router_fp32_gemm": "router_fp32_det_gemm", - "router_wgrad": "router_det_wgrad", - "split_fp32_to_bf16x3_": "split_fp32_to_bf16x3_det_", -} -# Every public name ``cake_backend`` defined before #8046 (FlashInfer c472a430, functions, classes and constants). -# fmt: off -FORMER_MODULE_NAMES = ( - "ARCH_RULE_DEFAULTS", "BLOCK_K", "BLOCK_M", "BLOCK_N_CHOICES", "B_SWZ_CHOICES", "CH_T_COLS", - "CONVENIENCE_CACHE_SIZE", "CTA_GROUP", "CTA_ROWS_CHOICES", "DeviceFacts", "EPI_MODES", "EPI_WARPS", - "EXACT_WAVE_BLOCK_N", "GROUP_M", "GemmPlan", "K_TMA_BF16", "K_TWO_SLOTS", "L2_HINTS", "L2_PROMOS", - "PANEL_BYTES", "PD_CHOICES", "PreparedGemm", "PreparedRouterGemm", "ROUTER_BLOCK_K", "ROUTER_BLOCK_M", - "ROUTER_BLOCK_N", "ROUTER_SPLITS", "ROW_RULES", "RULE_PARENT", "RouterPlan", "SK_COUNTERS_MIN", - "SK_DUMMY_BASE", "SK_MIN_ITERS", "SK_SLAB_CHOICES", "SLOT_BYTES", "SMEM_OPT_IN", "SMEM_OVERSIZED", - "SPLIT_BLOCK_ELEMS", "SPLIT_TEMPLATE", "SUPPORTED_COMPUTE_CAPABILITIES", "W8_FORMS", "WORK_STAGES", - "a_halves_of", "arch_for", "as_batched", "b_stage_bytes", "bind_launch", "box_rows_of", "check_stream_k_plan", - "clear_convenience_caches", "convenience_key", "default_block_n", "default_cta_rows", "default_group_m", - "default_hints", "default_hints_cgrp", "default_pf", "default_promo", "default_stages", - "dense_projection_gemm", "device_facts", "device_l2_bytes", "device_sm_count", "epi_cols", "epi_mode", - "epi_slots", "exact_waves", "flat_alias", "fused_split_module", "generated_program_available", "instance_key", - "instance_symbol", "operand_view", "plan_dense_projection_gemm", "plan_router_fp32_gemm", - "prepare_dense_projection_gemm", "prepare_projection_wgrad", "prepare_router_fp32_gemm", "projection_dgrad", - "projection_forward", "projection_wgrad", "require_arch", "router_dgrad", "router_forward", - "router_fp32_gemm", "router_layout_class", "router_symbol", "router_wgrad", "row_rule", "sk_exact_plan", - "sk_parts_plan", "sk_sync_plan", "sm_pairs", "sm_units", "smem_limit_for", "split_fp32_to_bf16x3_", - "staging_bytes", "stream_k_plan", "stt_bytes", "swap_small_m", "wave_working_set", "wave_working_set_cgrp", - "wgrad_swapped", "wgrad_views", -) -# fmt: on -# The names among FORMER_MODULE_NAMES that this package defines itself (its entry points, helpers and constants); every -# other one is absent here (no read-through to the deterministic package). -OWN_FORMER_NAMES = frozenset( - ( - *ENTRY_POINTS, - "router_symbol", - "bind_launch", - "clear_convenience_caches", - "ROUTER_BLOCK_K", - "ROUTER_BLOCK_M", - "ROUTER_BLOCK_N", - ) -) -# The names this package no longer has: the former-name map, its read-through and type tables, the mode resolver, the -# wrappers that only returned or ran a deterministic object, and the forwarding of deterministic templates and programs. -REMOVED_BACKEND_NAMES = ( - "DETERMINISTIC_NAMES", - "_DET_TYPES", - "_DET_SPELLED", - "_det_entry", - "_resolve_deterministic", - "instance_symbol", - "split_fp32_to_bf16x3_", - "__getattr__", -) -REMOVED_JIT_NAMES = ( - "is_det_template", - "is_det_program", - "DET_TEMPLATE_PREFIXES", - "DET_PROGRAM_PREFIX", -) -# The values of ``deterministic`` other than ``False``: each raises ``ValueError`` (``_OMITTED``: no keyword). -_OMITTED = object() -NOT_FALSE = (_OMITTED, True, None, 0, 1, "False", 0.0, "") +# A generated name (program, kernel, template, file) of this package never carries the token ``det``. DET_TOKEN = re.compile(r"(?:^|[^a-z])det(?:[^a-z]|$)", re.IGNORECASE) -# the mode is named by deterministic= -NEGATED_TOKEN = re.compile(r"non_?" + "det", re.IGNORECASE) class Recorder: @@ -192,9 +95,14 @@ def torch_determinism(): torch.use_deterministic_algorithms(enabled, warn_only=warn_only) -def _patch_det(monkeypatch, former: str) -> Recorder: +def _det_name(name: str) -> str: + """The deterministic package's entry point of the operation of this package's entry point ``name``.""" + return name.replace("cake_dense_gemm_nondet_train", "cake_dense_gemm_det_train") + + +def _patch_det(monkeypatch, name: str) -> Recorder: recorder = Recorder() - monkeypatch.setattr(det_backend, FORMER_TO_DET[former], recorder) + monkeypatch.setattr(det_backend, _det_name(name), recorder) return recorder @@ -241,80 +149,6 @@ def test_cake_build_is_the_pinned_loader(): assert ours.read_bytes() == theirs.read_bytes() -@pytest.mark.parametrize("name", ENTRY_POINTS) -def test_entry_points_take_keyword_only_deterministic_true(name): - parameter = inspect.signature(getattr(cake_backend, name)).parameters[ - "deterministic" - ] - assert parameter.kind is inspect.Parameter.KEYWORD_ONLY - assert parameter.default is True - - -@pytest.mark.parametrize("name", FORMER_MODULE_NAMES) -def test_no_former_module_name_reads_through(name): - """A name ``cake_backend`` had before #8046 is either this package's own object (defined in this module) or - absent: no alias and no read-through to the deterministic package.""" - if name in OWN_FORMER_NAMES: - assert name in vars(cake_backend) - obj = vars(cake_backend)[name] - if callable(obj): - assert obj.__module__ == cake_backend.__name__ - else: - assert name not in vars(cake_backend) and name not in dir(cake_backend) - with pytest.raises(AttributeError): - getattr(cake_backend, name) - - -def test_the_former_types_are_absent(): - for name in ("GemmPlan", "PreparedGemm", "RouterPlan", "PreparedRouterGemm"): - assert not hasattr(cake_backend, name), name - - -def test_the_removed_names_are_absent(): - """No former-name map, read-through (module ``__getattr__``), mode resolver, deterministic wrapper or forwarding - of deterministic templates and programs is left in this package.""" - for name in REMOVED_BACKEND_NAMES: - assert name not in vars(cake_backend), name - for name in REMOVED_BACKEND_NAMES[:-1]: - with pytest.raises(AttributeError): - getattr(cake_backend, name) - for name in REMOVED_JIT_NAMES: - assert not hasattr(cake_jit, name), name - - -@pytest.mark.parametrize( - "name", - [ - "projection_det_forward", - "DetGemmPlan", - "router_fp32_det_gemm", - "ROW_RULES", - "stream_k_plan", - "wgrad_views", - "__wrapped__", - "no_such_name", - ], -) -def test_deterministic_names_are_not_reachable(name): - with pytest.raises(AttributeError): - getattr(cake_backend, name) - - -def test_no_deterministic_object_is_reachable_through_this_package(): - """No function or class the deterministic package defines is an attribute of this package's modules (the - deterministic package is reached only through ``deterministic_backend()``).""" - det_objects = { - id(obj) - for module in (det_backend, det_jit) - for obj in vars(module).values() - if (inspect.isfunction(obj) or inspect.isclass(obj)) - and obj.__module__ == module.__name__ - } - for module in (cake_backend, cake_jit, package): - leaked = sorted(n for n, obj in vars(module).items() if id(obj) in det_objects) - assert not leaked, (module.__name__, leaked) - - def test_the_deterministic_package_resolves_as_a_sibling(monkeypatch): """In a FlashInfer tree and in a copy vendored with the deterministic package beside it (the vendored gate's layout) the deterministic package resolves as this package's sibling.""" @@ -340,16 +174,21 @@ def test_use_deterministic_backend_binds_a_vendored_copy(monkeypatch): # a forward layout (routed "det") plans with the bound package's deterministic planner A, B, out = _k1_views("kk") assert ( - cake_backend.plan_dense_projection_gemm(A, B, out, deterministic=False) - is fake.plan_dense_projection_det_gemm.result + cake_backend.cake_dense_gemm_nondet_train_projection_plan(A, B, out) + is fake.cake_dense_gemm_det_train_projection_plan.result + ) + assert fake.cake_dense_gemm_det_train_projection_plan.calls == [((A, B, out), {})] + # a knob the caller gives reaches the deterministic planner as given, and only that knob + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_plan( + A, B, out, sm_count=148, block_n=128 + ) + is fake.cake_dense_gemm_det_train_projection_plan.result + ) + assert fake.cake_dense_gemm_det_train_projection_plan.calls[1] == ( + (A, B, out), + {"sm_count": 148, "block_n": 128}, ) - assert fake.plan_dense_projection_det_gemm.calls == [((A, B, out), {})] - # the mode rule names the bound package - with pytest.raises( - ValueError, match=r"vendored_det_pkg\.cake_backend\.projection_det_forward" - ): - cake_backend.projection_forward("X", "W") - assert not fake.projection_det_forward.calls with pytest.raises(TypeError, match="lacks"): cake_backend.use_deterministic_backend(ModuleType("not_a_backend.cake_backend")) @@ -358,120 +197,11 @@ def test_use_deterministic_backend_binds_a_vendored_copy(monkeypatch): # Dispatch (CPU) # --------------------------------------------------------------------------- -# entry point -> its positional arguments (sentinels: the mode rule raises before anything reads them) and the -# keyword arguments of the call -# fmt: off -MODE_CALLS = { - "dense_projection_gemm": (("A", "B", "out"), {"transposed_out": True, "enable_pdl": True}), - "plan_dense_projection_gemm": (("A", "B", "out"), {"sm_count": 148, "transposed_out": True, "block_n": 128}), - "prepare_dense_projection_gemm": (("A", "B", "out"), {"sk": False}), - "projection_forward": (("X", "W"), {}), - "projection_dgrad": (("G", "W", "out"), {"out_dtype": "f32"}), - "projection_wgrad": (("G", "X"), {"out_dtype": "f32", "sk": True, "cta_rows": 64}), - "prepare_projection_wgrad": (("G", "X", "out"), {"cta_rows": 128}), - "plan_router_fp32_gemm": (("A", "B", "out"), {"splits": 3}), - "prepare_router_fp32_gemm": (("A", "B", "out"), {}), - "router_fp32_gemm": (("A", "B", "out"), {"splits": 2}), - "router_forward": (("X", "W"), {}), - "router_dgrad": (("G", "W", "out"), {}), - "router_wgrad": (("G", "X"), {"splits": 9}), -} -# fmt: on - - -def _mode_message(name: str) -> str: - """The whole ``ValueError`` text of ``name``'s mode rule, as a pattern (any value after ``got``).""" - package = cake_backend.__name__.rpartition(".")[0] or cake_backend.__name__ - det_entry = f"{cake_backend._det_package_name()}.cake_backend.{FORMER_TO_DET[name]}" - head = f"{name}: {package} runs the non-deterministic package only and needs deterministic=False (got " - tail = ( - "): pass deterministic=False for the non-deterministic path; " - f"the deterministic entry is {det_entry}" - ) - return "^" + re.escape(head) + ".*" + re.escape(tail) + "$" - - -def _forbid_det_package(monkeypatch) -> list: - """Every entry point of the deterministic package and every path of this package records a call, and - ``deterministic_backend`` fails: a call that raises the mode rule's ``ValueError`` reaches none of them.""" - reached = [] - - def forbidden(*args, **kwargs): - reached.append((args, kwargs)) - raise AssertionError("the mode rule let the call through") - - monkeypatch.setattr(cake_backend, "deterministic_backend", forbidden) - for name in cake_backend._DET_ENTRIES: - monkeypatch.setattr(det_backend, name, forbidden) - paths = ("_atomic_k1", "_atomic_k2", "_det_k1", "_router_gemm", "_splitk_k2") - for name in (*paths, "_k1_cached_launch", "_k2_entry"): - monkeypatch.setattr(cake_backend, name, forbidden) - return reached - - -def test_mode_calls_cover_every_entry_point(): - assert sorted(MODE_CALLS) == sorted(ENTRY_POINTS) - assert set(ENTRY_POINTS) <= set(FORMER_TO_DET) - for name in ENTRY_POINTS: - assert hasattr(det_backend, FORMER_TO_DET[name]), name - - -@pytest.mark.parametrize( - "value", NOT_FALSE, ids=lambda v: "omitted" if v is _OMITTED else repr(v) -) -@pytest.mark.parametrize("name", ENTRY_POINTS) -def test_every_value_but_false_raises_naming_the_deterministic_entry( - monkeypatch, torch_determinism, name, value -): - """An entry point takes only ``deterministic=False``: the default, ``True``, ``None``, ``0``, ``1``, ``"False"`` - and every other value raise ``ValueError`` naming the deterministic package's entry point, before anything is - planned or launched and without a call of the deterministic package; PyTorch's deterministic mode (strict or - ``warn_only``) changes nothing, and nothing warns.""" - args, kwargs = MODE_CALLS[name] - reached = _forbid_det_package(monkeypatch) - if value is not _OMITTED: - kwargs = dict(kwargs, deterministic=value) - for enabled, warn_only in ((False, False), (True, False), (True, True)): - torch.use_deterministic_algorithms(enabled, warn_only=warn_only) - with _no_warnings(), pytest.raises(ValueError, match=_mode_message(name)): - getattr(cake_backend, name)(*args, **kwargs) - assert not reached - - -def test_numpy_bools_raise(): - """A NumPy bool is not the bool ``False``: ``np.False_`` and ``np.True_`` raise like any other value.""" - np = pytest.importorskip("numpy") - for value in (np.False_, np.True_): - with pytest.raises(ValueError, match=_mode_message("projection_wgrad")): - cake_backend.projection_wgrad("G", "X", "out", deterministic=value) - - -def test_the_mode_message_names_the_deterministic_package_without_importing_it( - monkeypatch, -): - """The mode rule names the deterministic entry point from text (the bound package, else this package's - sibling): it imports nothing.""" - imported = [] - monkeypatch.setattr( - cake_jit, "deterministic_module", lambda name: imported.append(name) - ) - monkeypatch.setattr(cake_jit, "_det_package", None) - with pytest.raises(ValueError) as raised: - cake_backend.router_wgrad("G", "X", deterministic=None) - parent = cake_backend.__name__.rpartition(".")[0].rpartition(".")[0] - det_package = ( - f"{parent}.dense_projection_det_gemm" if parent else "dense_projection_det_gemm" - ) - assert f"{det_package}.cake_backend.router_det_wgrad" in str(raised.value) - assert "needs deterministic=False (got None)" in str(raised.value) - assert not imported - -def test_router_symbol_takes_no_mode_keyword(): - """``router_symbol`` names this package's stream-K template of the ``kk`` / ``nn_t`` layout classes (a helper: - no ``deterministic`` keyword); ``kn`` (SIMT programs only) and ``nk`` have none.""" +def test_router_symbol_names_the_stream_k_templates(): + """``router_symbol`` names this package's stream-K template of the ``kk`` / ``nn_t`` layout classes; ``kn`` (SIMT + programs only) and ``nk`` have none.""" symbol = cake_backend.router_symbol - assert "deterministic" not in inspect.signature(symbol).parameters assert symbol(False, False, False) == "router_fp32_gemm_atomic_kk" assert symbol(True, True, True) == "router_fp32_gemm_atomic_nn_t" stream_k = {"router_fp32_gemm_atomic_kk", "router_fp32_gemm_atomic_nn_t"} @@ -480,19 +210,17 @@ def test_router_symbol_takes_no_mode_keyword(): for a_mn, b_mn in ((False, True), (True, False)): with pytest.raises(ValueError, match="no stream-K program"): cake_backend.router_symbol(a_mn, b_mn, False) - with pytest.raises(TypeError): - cake_backend.router_symbol(False, False, False, deterministic=False) @pytest.mark.parametrize("layout", ["kk", "kn", "nn"]) def test_k1_routes(monkeypatch, layout): - """``deterministic=False``: the weight-gradient layout reaches this package's split-K path; the forward and + """The weight-gradient layout reaches this package's split-K path; the forward and input-gradient layouts run the deterministic program.""" - det = _patch_det(monkeypatch, "dense_projection_gemm") + det = _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_projection") atomic = Recorder() monkeypatch.setattr(cake_backend, "_atomic_k1", atomic) A, B, out = _k1_views(layout) - result = cake_backend.dense_projection_gemm(A, B, out, deterministic=False) + result = cake_backend.cake_dense_gemm_nondet_train_projection(A, B, out) if layout == "nn": assert result is atomic.result and not det.calls ((args, kwargs),) = atomic.calls @@ -518,21 +246,27 @@ def test_k1_routes(monkeypatch, layout): def test_k1_operation_routes(monkeypatch): forward, dgrad = ( - _patch_det(monkeypatch, "projection_forward"), - _patch_det(monkeypatch, "projection_dgrad"), + _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_projection_forward"), + _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_projection_dgrad"), ) atomic = Recorder() monkeypatch.setattr(cake_backend, "_atomic_k1", atomic) X, W, G = torch.zeros(64, 32), torch.zeros(16, 32), torch.zeros(64, 16) - assert cake_backend.projection_forward(X, W, deterministic=False) is forward.result - assert cake_backend.projection_dgrad(G, W, deterministic=False) is dgrad.result + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W) + is forward.result + ) + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_dgrad(G, W) is dgrad.result + ) assert not atomic.calls out = torch.zeros(16, 32, dtype=torch.bfloat16) assert ( - cake_backend.projection_wgrad(G, X, out, deterministic=False) is atomic.result + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is atomic.result ) assert ( - cake_backend.prepare_projection_wgrad(G, X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare(G, X, out) is atomic.result ) (launch_args, launch_kwargs), (prepare_args, prepare_kwargs) = atomic.calls @@ -543,8 +277,10 @@ def test_k1_operation_routes(monkeypatch): assert kwargs["transposed_out"] is True and kwargs["own"] == {} # a row the planner does not take runs the deterministic weight-gradient entry point with its knobs wgrad, prepare = ( - _patch_det(monkeypatch, "projection_wgrad"), - _patch_det(monkeypatch, "prepare_projection_wgrad"), + _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_projection_wgrad"), + _patch_det( + monkeypatch, "cake_dense_gemm_nondet_train_projection_wgrad_prepare" + ), ) assert launch_kwargs["det_call"]() is wgrad.result assert prepare_kwargs["det_call"]() is prepare.result @@ -566,29 +302,34 @@ def test_k1_operation_routes(monkeypatch): @pytest.mark.parametrize("pdl", [None, True]) @pytest.mark.parametrize("op", ["fwd", "dgrad", "kk", "kn"]) def test_k1_deterministic_programs_run_through_the_binding(monkeypatch, op, pdl): - """``deterministic=False`` K1 calls that run the deterministic program go through this package's binding + """K1 calls of this package that run the deterministic program go through this package's binding (``_det_k1``): keyed by the entry point's own tensors, with the deterministic package's convenience call as the plan (``det_args``) and as the fallback (``det_call``), the entry point's ``enable_pdl`` (and the generic entry's - ``weights_static``) passed to both; ``deterministic=True`` raises ``ValueError`` and calls nothing.""" - names = {"fwd": "projection_forward", "dgrad": "projection_dgrad"} - det = _patch_det(monkeypatch, names.get(op, "dense_projection_gemm")) + ``weights_static``) passed to both.""" + names = { + "fwd": "cake_dense_gemm_nondet_train_projection_forward", + "dgrad": "cake_dense_gemm_nondet_train_projection_dgrad", + } + det = _patch_det( + monkeypatch, names.get(op, "cake_dense_gemm_nondet_train_projection") + ) binding = Recorder() monkeypatch.setattr(cake_backend, "_det_k1", binding) X, W, G = torch.zeros(64, 32), torch.zeros(16, 32), torch.zeros(64, 16) if op == "fwd": - call = lambda d: cake_backend.projection_forward( # noqa: E731 - X, W, deterministic=d, enable_pdl=pdl + call = lambda: cake_backend.cake_dense_gemm_nondet_train_projection_forward( # noqa: E731 + X, W, enable_pdl=pdl ) elif op == "dgrad": - call = lambda d: cake_backend.projection_dgrad( # noqa: E731 - G, W, out_dtype=torch.float32, deterministic=d, enable_pdl=pdl + call = lambda: cake_backend.cake_dense_gemm_nondet_train_projection_dgrad( # noqa: E731 + G, W, out_dtype=torch.float32, enable_pdl=pdl ) else: A, B, out = _k1_views(op) - call = lambda d: cake_backend.dense_projection_gemm( # noqa: E731 - A, B, out, deterministic=d, enable_pdl=pdl, weights_static=bool(pdl) + call = lambda: cake_backend.cake_dense_gemm_nondet_train_projection( # noqa: E731 + A, B, out, enable_pdl=pdl, weights_static=bool(pdl) ) - assert call(False) is binding.result and not det.calls + assert call() is binding.result and not det.calls ((sources, flags, det_args, det_call, enable_pdl, *rest), kwargs) = binding.calls[0] assert not kwargs and enable_pdl is pdl pdl_kw = {} if pdl is None else {"enable_pdl": pdl} @@ -614,10 +355,7 @@ def test_k1_deterministic_programs_run_through_the_binding(monkeypatch, op, pdl) if pdl: want["weights_static"] = True assert det.calls == [((A, B, out), want)] - det.calls.clear() - with pytest.raises(ValueError, match="needs deterministic=False"): - call(True) - assert not det.calls and len(binding.calls) == 1 + assert len(binding.calls) == 1 def _same_entry_views(views, a_view, split_src, flags) -> None: @@ -628,36 +366,41 @@ def _same_entry_views(views, a_view, split_src, flags) -> None: def test_k2_routes(monkeypatch): - """``deterministic=False``: the router weight gradient and input gradient reach this package's K2 path (the input + """The router weight gradient and input gradient reach this package's K2 path (the input gradient without a split count: its SIMT windows, else the deterministic program), and so does the forward with a split count or while some router table gives the forward a verdict; without either the forward runs the deterministic program through this package's binding (the plan the planner keeps for a row without a verdict). Each entry point passes its own tensors as the launch views. The forward's shipped verdicts (its SIMT windows) are cleared first.""" _without_forward_verdicts(monkeypatch) - dgrad = _patch_det(monkeypatch, "router_dgrad") - gemm = _patch_det(monkeypatch, "router_fp32_gemm") + dgrad = _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_router_dgrad") + gemm = _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_router") atomic, splitk = Recorder(), Recorder() monkeypatch.setattr(cake_backend, "_atomic_k2", atomic) monkeypatch.setattr(cake_backend, "_splitk_k2", splitk) T, K, N = 64, 32, 16 X, W, G = torch.zeros(T, K), torch.zeros(N, K), torch.zeros(T, N) - assert cake_backend.router_forward(X, W, deterministic=False) is splitk.result + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) is splitk.result + ) ((args, kwargs),) = splitk.calls assert args[0] is X and args[1].equal(W.t()) and args[3] is None and not kwargs assert args[5] is None # enable_pdl: the device's default assert tuple(args[2].shape) == (T, N) and args[2].dtype == torch.float32 _same_entry_views(args[4], X.unsqueeze(0), W, cake_backend._FORWARD_VIEWS) assert ( - cake_backend.router_forward(X, W, deterministic=False, splits=2) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, splits=2) is atomic.result ) monkeypatch.setitem(cake_backend.ROUTER_RULES, ("sm_103a", "kk", K, N), "S2") - assert cake_backend.router_forward(X, W, deterministic=False) is atomic.result assert ( - cake_backend.router_wgrad(G, X, deterministic=False, splits=5) is atomic.result + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) is atomic.result ) - assert cake_backend.router_dgrad(G, W, deterministic=False) is atomic.result + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, splits=5) + is atomic.result + ) + assert cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W) is atomic.result assert not dgrad.calls calls = atomic.calls ( @@ -690,8 +433,8 @@ def test_k2_routes(monkeypatch): # the generic entry: every class this package's path (kn: a split count runs the deterministic program there), # without entry-point views assert ( - cake_backend.router_fp32_gemm( - G, W, torch.zeros(T, K), deterministic=False, splits=1 + cake_backend.cake_dense_gemm_nondet_train_router( + G, W, torch.zeros(T, K), splits=1 ) is atomic.result ) @@ -699,11 +442,11 @@ def test_k2_routes(monkeypatch): atomic.calls[4][1]["splits"] == 1 and atomic.calls[4][1]["enable_pdl"] is None ) assert ( - cake_backend.router_fp32_gemm(X, W.t(), torch.zeros(T, N), deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(X, W.t(), torch.zeros(T, N)) is atomic.result ) assert ( - cake_backend.router_fp32_gemm(G.t(), X, torch.zeros(N, K), deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(G.t(), X, torch.zeros(N, K)) is atomic.result ) assert all(kw["views"] is None for _, kw in atomic.calls[4:]) @@ -712,7 +455,7 @@ def test_k2_routes(monkeypatch): @pytest.mark.parametrize("pdl", [True, False]) def test_k2_routes_pass_enable_pdl(monkeypatch, pdl): - """``enable_pdl`` reaches every K2 route of ``deterministic=False``: this package's binding of the deterministic + """``enable_pdl`` reaches every K2 route of this package: this package's binding of the deterministic program (the forward without a verdict) and this package's K2 path (the input gradient, the generic ``kn`` call, the forward with a split count, the weight gradient). The forward's shipped verdicts are cleared first.""" _without_forward_verdicts(monkeypatch) @@ -721,12 +464,14 @@ def test_k2_routes_pass_enable_pdl(monkeypatch, pdl): monkeypatch.setattr(cake_backend, "_splitk_k2", splitk) T, K, N = 64, 32, 16 X, W, G = torch.zeros(T, K), torch.zeros(N, K), torch.zeros(T, N) - kw = dict(deterministic=False, enable_pdl=pdl) - cake_backend.router_forward(X, W, **kw) - cake_backend.router_dgrad(G, W, **kw) - cake_backend.router_fp32_gemm(G, W, torch.zeros(T, K), splits=1, **kw) - cake_backend.router_forward(X, W, splits=2, **kw) - cake_backend.router_wgrad(G, X, splits=5, **kw) + kw = dict(enable_pdl=pdl) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, **kw) + cake_backend.cake_dense_gemm_nondet_train_router( + G, W, torch.zeros(T, K), splits=1, **kw + ) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, splits=2, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, splits=5, **kw) assert [args[5] for args, _ in splitk.calls] == [pdl] assert [kwargs["enable_pdl"] for _, kwargs in atomic.calls] == [pdl] * 4 @@ -1849,7 +1594,10 @@ def test_k1_152_sm_one_split_ranges_plan_the_measured_launch( out = torch.zeros(d_out, 4096, dtype=torch.float32) kw = dict(transposed_out=True, arch=arch, sm_count=152) plan, *_, reason = cake_backend.plan_atomic_gemm(X.t(), G, out, **kw) - assert isinstance(plan, cake_backend.AtomicGemmPlan) and reason == "" + assert ( + isinstance(plan, cake_backend.CakeDenseGemmNondetProjectionPlan) + and reason == "" + ) explicit = cake_backend.plan_atomic_gemm(X.t(), G, out, splits=1, block_n=64, **kw) assert plan == explicit[0] assert (plan.template, plan.block_n, plan.splits, plan.chains, plan.out_t) == ( @@ -1875,7 +1623,7 @@ def test_k1_152_sm_one_split_ranges_plan_the_measured_launch( def test_k1_152_sm_one_split_ranges_bind_the_one_split_program(monkeypatch): - """With ``deterministic=False`` the weight gradient of a row at a token count of a one-split range (152-SM sm_103a + """The weight gradient of a row at a token count of a one-split range (152-SM sm_103a part, ``4096 -> 128`` FP32 at 1000 tokens) binds this package's one-split program in ``_K1_ATOMIC_CACHE``: no arena (one split: no workspace, no tickets), no deterministic plan, one launch per call.""" table = cake_backend.K1_ONE_SPLIT_RANGES_BY_SM @@ -1887,7 +1635,10 @@ def test_k1_152_sm_one_split_ranges_bind_the_one_split_program(monkeypatch): X, G = torch.empty(T, 4096, dtype=bf16), torch.empty(T, 128, dtype=bf16) out_w = torch.empty(128, 4096) for _ in range(2): - assert cake_backend.projection_wgrad(G, X, out_w, deterministic=False) is out_w + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w) + is out_w + ) assert len(entry_fn.calls) == 2 ((_, atomic),) = cake_backend._K1_ATOMIC_CACHE.items() assert atomic.plan is not None and atomic.k1 is not None @@ -2084,7 +1835,10 @@ def test_k1_152_sm_dsm_ranges_plan_the_cluster_launch(monkeypatch, d_in, d_out, for lo, hi, v in K1_152_DSM_RANGES[("sm_100a", 152)][(d_in, d_out)] if lo <= T <= hi ] - assert isinstance(plan, cake_backend.AtomicGemmPlan) and reason == "" + assert ( + isinstance(plan, cake_backend.CakeDenseGemmNondetProjectionPlan) + and reason == "" + ) assert plan.dsm and not plan.swapped and not plan.workspace assert (plan.template, plan.block_n, plan.splits, plan.chains, plan.out_t) == ( f"dense_proj_gemm_dsm_nn_n{block_n}_f32_t_s{splits}", @@ -2129,7 +1883,7 @@ def test_k1_152_sm_dsm_ranges_plan_the_cluster_launch(monkeypatch, d_in, d_out, def test_k1_152_sm_dsm_ranges_bind_the_cluster_program(monkeypatch): - """With ``deterministic=False`` the weight gradient of a row at a token count of a cluster range (152-SM sm_100a + """The weight gradient of a row at a token count of a cluster range (152-SM sm_100a part, ``4096 -> 64`` FP32 at 2048 tokens) binds this package's cluster program in ``_K1_ATOMIC_CACHE``: no arena (no workspace, no tickets), no deterministic plan, one launch per call whose arguments are the program's argument plan (the operand views, ``out``, the extents, the output strides, the PDL flag and the plan's grid of @@ -2143,7 +1897,10 @@ def test_k1_152_sm_dsm_ranges_bind_the_cluster_program(monkeypatch): X, G = torch.empty(T, 4096, dtype=bf16), torch.empty(T, 64, dtype=bf16) out_w = torch.empty(64, 4096) for _ in range(2): - assert cake_backend.projection_wgrad(G, X, out_w, deterministic=False) is out_w + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w) + is out_w + ) assert len(entry_fn.calls) == 2 ((_, atomic),) = cake_backend._K1_ATOMIC_CACHE.items() assert atomic.plan is not None and atomic.k1 is not None and atomic.plan.dsm @@ -2162,8 +1919,10 @@ def test_k1_152_sm_dsm_ranges_bind_the_cluster_program(monkeypatch): assert values["out"].data_ptr() == out_w.data_ptr() tiles = 32 * 1 assert (values["grid_x"], values["grid_y"], values["grid_z"]) == (tiles * 3, 1, 1) - prepared = cake_backend.prepare_projection_wgrad(G, X, out_w, deterministic=False) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm) + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out_w + ) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch) assert prepared.plan == atomic.plan and prepared.module_name == atomic.module_name assert prepared.grid == (tiles * 3, 1, 1) and prepared.counters_are_zero() prepared.reset_counters() # no tickets: nothing to zero @@ -2537,7 +2296,10 @@ def plan(sms, T, M=d_in, N=d_out): plan_, *_, reason = cake_backend.plan_atomic_gemm( X.t(), G, out, transposed_out=True, arch=arch, sm_count=216 ) - assert isinstance(plan_, cake_backend.AtomicGemmPlan) and reason == "" + assert ( + isinstance(plan_, cake_backend.CakeDenseGemmNondetProjectionPlan) + and reason == "" + ) assert plan_.template == "dense_proj_gemm_atomic_nn_n192_f32" and not plan_.out_t assert (plan_.M, plan_.N, plan_.block_n, plan_.splits, plan_.tiles) == ( d_out, @@ -3020,7 +2782,7 @@ def test_k1_chains_fit(monkeypatch): def test_k1_chain_plans_run_the_chain_programs(monkeypatch): """A token count of :data:`K1_CHAIN_RANGES_BY_SM` plans the one-split program of three chains of the range's width - (``AtomicGemmPlan.chains``, template ``..._c3``, no workspace and no tickets): the 128-wide and 256-wide + (``CakeDenseGemmNondetProjectionPlan.chains``, template ``..._c3``, no workspace and no tickets): the 128-wide and 256-wide transposed-store programs and the 192-wide row-major-store program, inside a 212-SM two-split range too; past the chain range the two-split range plans two splits of one chain. An explicit ``splits`` / ``block_n`` runs one chain; ``chains=3`` names the three-chain program of any row at one split; a chain count outside :data:`K1_CHAINS_CHOICES`, @@ -3043,7 +2805,10 @@ def operands(T, d_in, d_out, dtype=torch.float32): plan, *_, reason = cake_backend.plan_atomic_gemm( *operands(4097, d_in, d_out), **kw ) - assert isinstance(plan, cake_backend.AtomicGemmPlan) and reason == "" + assert ( + isinstance(plan, cake_backend.CakeDenseGemmNondetProjectionPlan) + and reason == "" + ) assert (plan.template, plan.splits, plan.chains, plan.out_t) == ( template, 1, @@ -3149,7 +2914,10 @@ def test_k1_plan_is_the_transposed_store_launch(monkeypatch): cake_backend.plan_atomic_gemm(X.t(), G, out, transposed_out=True, **kw), swapped, ): - assert isinstance(plan, cake_backend.AtomicGemmPlan) and reason == "" + assert ( + isinstance(plan, cake_backend.CakeDenseGemmNondetProjectionPlan) + and reason == "" + ) assert (plan.L, plan.M, plan.N, plan.K, plan.block_n, plan.splits) == ( 1, K, @@ -3262,7 +3030,10 @@ def test_k1_row_major_verdict_runs_the_row_major_program(monkeypatch, no_k1_chai plan, a_desc, b_desc, out3, reason = cake_backend.plan_atomic_gemm( A, B, out, transposed_out=transposed_out, **kw ) - assert isinstance(plan, cake_backend.AtomicGemmPlan) and reason == "" + assert ( + isinstance(plan, cake_backend.CakeDenseGemmNondetProjectionPlan) + and reason == "" + ) assert plan.template == "dense_proj_gemm_atomic_nn_n192_f32" assert not plan.out_t assert (plan.L, plan.M, plan.N, plan.K, plan.block_n, plan.splits) == ( @@ -4967,7 +4738,7 @@ def test_k2_values_feed_every_registered_router_program(): launches as a programmatic dependent launch); those of a forward SIMT launch every argument of every registered forward program's plan, with the call's PDL flag.""" out = torch.zeros(256, 5120) - plan = cake_backend.AtomicRouterPlan( + plan = cake_backend.CakeDenseGemmNondetRouterPlan( M=5120, N=256, K=16231, a_mn=True, b_mn=True, swapped=True, layout="nn_t", m_tiles=40, n_tiles=2, k_iters=508, units=540, template="router_fp32_gemm_atomic_nn_t_ht_c3_nreg", head_spans=54, head_steps=5080, @@ -5261,7 +5032,9 @@ def test_entry_point_views_are_the_launch_views(monkeypatch, op): out = torch.empty(A.shape[0], B.shape[1]) plans = [] try: - plan, A_eff, B_eff = det_backend.plan_router_fp32_det_gemm(A, B, out) + plan, A_eff, B_eff = det_backend.cake_dense_gemm_det_train_router_plan( + A, B, out + ) except (ValueError, NotImplementedError): assert not canonical else: @@ -5341,7 +5114,12 @@ def test_splitk_binding_is_the_deterministic_launch(monkeypatch, case, pdl): entry = cake_backend._plan_splitk_k2(A, B, out, splits, 0, 0, pdl) plan, bound, arena = entry.plan, entry.k2, entry.arena assert bound is not None and arena is not None and entry.entry is not None - assert plan == det_backend.plan_router_fp32_det_gemm(A, B, out, splits=splits)[0] + assert ( + plan + == det_backend.cake_dense_gemm_det_train_router_plan(A, B, out, splits=splits)[ + 0 + ] + ) assert plan.splits == {"fwd": 4, "dgrad": 1, "wgrad": 6}[op] a_view, split_src = cake_backend._launch_views(A, B, plan, views) split_module = det_backend._fused_split_for(arch, split_src) @@ -5744,10 +5522,10 @@ def call(): def test_k1_convenience_calls_cache_their_bindings(monkeypatch): - """The K1 entry points with ``deterministic=False`` plan once per key and launch the bound program after: the + """The K1 entry points plan once per key and launch the bound program after: the forward and the input gradient in ``_K1_DET_CACHE`` (keyed by their own tensors), the weight gradient in ``_K1_ATOMIC_CACHE`` -- a row without a verdict bound to the deterministic package's plan of - ``projection_det_wgrad`` (its orientation; another ``sk`` is another key and plan), a row with one to this + ``cake_dense_gemm_det_train_projection_wgrad`` (its orientation; another ``sk`` is another key and plan), a row with one to this package's program.""" arch = "sm_103a" entry_fn = _cpu_k1_launches(monkeypatch, arch) @@ -5765,10 +5543,17 @@ def test_k1_convenience_calls_cache_their_bindings(monkeypatch): ) for _ in range(2): assert ( - cake_backend.projection_forward(X, W, out_f, deterministic=False) is out_f + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, out_f) + is out_f + ) + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_dgrad(G, W, out_d) + is out_d + ) + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w) + is out_w ) - assert cake_backend.projection_dgrad(G, W, out_d, deterministic=False) is out_d - assert cake_backend.projection_wgrad(G, X, out_w, deterministic=False) is out_w assert len(entry_fn.calls) == 6 det_entries = list(cake_backend._K1_DET_CACHE.items()) # the flags, the PDL flag (the device's default: on) and ``weights_static`` @@ -5783,8 +5568,8 @@ def test_k1_convenience_calls_cache_their_bindings(monkeypatch): # the PDL flag is bound into the argument list: another flag is another key entry_fn.calls.clear() assert ( - cake_backend.projection_forward( - X, W, out_f, deterministic=False, enable_pdl=False + cake_backend.cake_dense_gemm_nondet_train_projection_forward( + X, W, out_f, enable_pdl=False ) is out_f ) @@ -5805,7 +5590,7 @@ def test_k1_convenience_calls_cache_their_bindings(monkeypatch): A, B, out_w, transposed_out=transposed, knobs=dict(sk="auto", cta_rows=None) )[0] ) - cake_backend.projection_wgrad(G, X, out_w, sk=False, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w, sk=False) assert len(cake_backend._K1_ATOMIC_CACHE) == 2 sk_entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) assert ( @@ -5819,7 +5604,7 @@ def test_k1_convenience_calls_cache_their_bindings(monkeypatch): monkeypatch.setitem(cake_backend.K1_SPLIT_RULES, key, (128, 2)) monkeypatch.delitem(cake_backend.K1_SPLIT_RULE_EXCLUDED_TOKENS, key, raising=False) cake_backend._K1_ATOMIC_CACHE.clear() - cake_backend.projection_wgrad(G, X, out_w, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w) ((_, atomic),) = cake_backend._K1_ATOMIC_CACHE.items() assert atomic.plan is not None and atomic.plan.splits == 2 and atomic.k1 is not None assert atomic.det_plan is None and atomic.arena is not None @@ -5864,7 +5649,7 @@ def test_k1_forward_rules_by_sm(): out = torch.empty(T, N, dtype=bf16) def plan(**kw): - return det_backend.plan_dense_projection_det_gemm( + return det_backend.cake_dense_gemm_det_train_projection_plan( X, W.t(), out, sm_count=sms, l2_bytes=GB_L2_BYTES, arch=arch, **kw )[0] @@ -5898,10 +5683,10 @@ def plan(**kw): @pytest.mark.parametrize("arch", ["sm_100a", "sm_103a"]) def test_k1_forward_rows_bind_their_plan(monkeypatch, arch): - """``projection_forward(deterministic=False)`` binds, for a row of :data:`K1_FORWARD_RULES_BY_SM` on its part + """``cake_dense_gemm_nondet_train_projection_forward()`` binds, for a row of :data:`K1_FORWARD_RULES_BY_SM` on its part (BF16 ``out``, both operands K-major), the deterministic program the deterministic planner gives the row with the table's knobs, and the deterministic plan for a token count, row or SM count outside the table; an FP32 ``out`` - or an operand that is not K-major gets no knobs; ``deterministic=True`` raises ``ValueError`` and plans nothing.""" + or an operand that is not K-major gets no knobs.""" entry_fn = _cpu_k1_launches(monkeypatch, arch) monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) bf16 = torch.bfloat16 @@ -5909,7 +5694,10 @@ def test_k1_forward_rows_bind_their_plan(monkeypatch, arch): def bound_plan(X, W, out): cake_backend._K1_DET_CACHE.clear() - assert cake_backend.projection_forward(X, W, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, out) + is out + ) ((key, entry),) = cake_backend._K1_DET_CACHE.items() assert key[2] == (("fwd",), 1, False) and entry.k1 is not None return entry.plan @@ -5946,17 +5734,10 @@ def det_plan(X, W, out, knobs): monkeypatch.setattr(det_backend, "device_facts", lambda device=None, f=facts: f) assert cake_backend._k1_forward_knobs_of(X, W, out) == (), sms assert bound_plan(X, W, out) == det_plan(X, W, out, ()), sms - # deterministic=True (the default): ValueError naming the deterministic entry, whatever the table says - det = _patch_det(monkeypatch, "projection_forward") - entry_fn.calls.clear() - cake_backend._K1_DET_CACHE.clear() - with pytest.raises(ValueError, match="projection_det_forward"): - cake_backend.projection_forward(X, W, out) - assert not det.calls and not entry_fn.calls and not cake_backend._K1_DET_CACHE def test_k1_entry_points_go_from_their_key_to_the_bound_launch(monkeypatch): - """Once a key is bound, the K1 entry points with ``deterministic=False`` launch it straight from the key + """Once a key is bound, the K1 entry points launch it straight from the key (``_k1_cached_launch``): neither ``_det_k1`` nor ``_atomic_k1`` runs, and the launch gets the arguments their launch of the key gets. An ``out`` of the same geometry is the same key and launches into itself; another geometry is another key and takes the planning path. Each thread makes the stream context of a (device, stream) once and @@ -5984,10 +5765,17 @@ def stream_context(index, raw=None): def calls(): assert ( - cake_backend.projection_forward(X, W, out_f, deterministic=False) is out_f + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, out_f) + is out_f + ) + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_dgrad(G, W, out_d) + is out_d + ) + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w) + is out_w ) - assert cake_backend.projection_dgrad(G, W, out_d, deterministic=False) is out_d - assert cake_backend.projection_wgrad(G, X, out_w, deterministic=False) is out_w def spelled(args): return [ @@ -6017,12 +5805,17 @@ def planning_path(*args, **kwargs): assert [spelled(args) for args, _ in entry_fn.calls] == planned assert len(made) == 1 other = torch.empty_like(out_f) - assert cake_backend.projection_forward(X, W, other, deterministic=False) is other + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, other) + is other + ) (args, _) = entry_fn.calls[-1] pointers = {a.data_ptr() for a in args if isinstance(a, torch.Tensor)} assert other.data_ptr() in pointers and out_f.data_ptr() not in pointers with pytest.raises(AssertionError, match="planning path"): - cake_backend.projection_forward(X[:999], W, out_f[:999], deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_forward( + X[:999], W, out_f[:999] + ) # another thread makes its own context of the same (device, stream) entry_fn.calls.clear() worker = threading.Thread(target=calls) @@ -6053,7 +5846,10 @@ def test_k1_sm107a_wide_fp32_rows_bind_per_token_count(monkeypatch, no_k1_chains out = torch.empty(d_out, d_in) for T in (T0 - 1, T0, 65536): G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (d_in, d_out, T) assert entry.k1 is not None, where @@ -6098,7 +5894,10 @@ def test_k1_212_sm_token_ranges_bind_per_token_count(monkeypatch): cases = sorted(set(cases)) for T, routed in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (d_in, d_out, dt, T) assert entry.k1 is not None, where @@ -6139,7 +5938,10 @@ def test_k1_212_sm_range_overrides_bind_per_token_count(monkeypatch): cases = sorted(set(cases)) for T, choice in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (d_in, d_out, dt, T) assert entry.k1 is not None, where @@ -6183,7 +5985,10 @@ def test_k1_212_sm_two_split_ranges_bind_per_token_count(monkeypatch, no_k1_chai cases = sorted(set(cases)) for T, choice in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (d_in, d_out, dt, T) assert entry.k1 is not None, where @@ -6235,7 +6040,10 @@ def test_k1_exact_sm_count_exclusions_bind_per_token_count(monkeypatch): torch.empty(T, d_in, dtype=bf16), ) assert ( - cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + G, X, out + ) + is out ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (arch, d_in, d_out, dt, T) @@ -6288,7 +6096,10 @@ def test_k1_odd_block_exclusions_bind_per_token_count(monkeypatch, arch, sm_coun cases = ((lo - 1, rule), (lo, None), (last, None), (last + 1, rule)) for T, choice in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (d_in, d_out, T) assert entry.k1 is not None, where @@ -6344,7 +6155,10 @@ def test_k1_chain_ranges_bind_per_token_count(monkeypatch, sm_count): assert any(c is not None and c[2] == 3 for _, c in cases) for T, choice in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) entry = next(reversed(cake_backend._K1_ATOMIC_CACHE.values())) where = (d_in, d_out, T) assert entry.k1 is not None, where @@ -6361,12 +6175,12 @@ def test_k1_chain_ranges_bind_per_token_count(monkeypatch, sm_count): def test_k1_swapped_small_token_forward_stays_on_the_deterministic_call(monkeypatch): """A batched forward the deterministic package runs on its swapped small-token program (``_swap_ab_launch``) is - not bound here: its entry stays unbound and ``dense_projection_gemm`` makes the deterministic package's call, with + not bound here: its entry stays unbound and ``cake_dense_gemm_nondet_train_projection`` makes the deterministic package's call, with ``enable_pdl`` and ``weights_static`` (part of the key); a batched forward of more tokens is bound.""" arch = "sm_103a" entry_fn = _cpu_k1_launches(monkeypatch, arch) monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) - det = _patch_det(monkeypatch, "dense_projection_gemm") + det = _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_projection") bf16, L, K, N = torch.bfloat16, 8, 4096, 1024 def views(M): @@ -6386,11 +6200,10 @@ def views(M): ) assert entry.plan is None and entry.k1 is None and entry.entry is None for _ in range(2): - result = cake_backend.dense_projection_gemm( + result = cake_backend.cake_dense_gemm_nondet_train_projection( A, B, out, - deterministic=False, enable_pdl=True, weights_static=weights_static, ) @@ -6406,9 +6219,7 @@ def views(M): A, B, out = views(1000) assert det_backend._swap_ab_launch(A, B, out, False) is None assert ( - cake_backend.dense_projection_gemm( - A, B, out, deterministic=False, enable_pdl=True - ) + cake_backend.cake_dense_gemm_nondet_train_projection(A, B, out, enable_pdl=True) is out ) assert len(det.calls) == 4 and len(entry_fn.calls) == 1 @@ -6581,8 +6392,8 @@ def test_router_memo_follows_the_router_tables(monkeypatch): """The memo of the router tables (``_router_memo``) is kept while no table changes and derived again after any change: the forward verdicts of the 152-SM table ship, and with them cleared none is seen; a forward verdict set in any table, and its removal, are seen at the next call; a table rebound in place of one is seen, and one that does not count its changes (a plain ``dict``) is read on - every call; ``clear_convenience_caches`` drops the memo.""" - cake_backend.clear_convenience_caches() + every call; ``cake_dense_gemm_nondet_train_convenience_cache_clear`` drops the memo.""" + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() memo = cake_backend._router_memo() assert memo.tracked and memo.forward_verdicts and not memo.without_verdict assert cake_backend._router_memo() is memo @@ -6638,7 +6449,7 @@ def test_router_memo_follows_the_router_tables(monkeypatch): assert memo.tracked and memo.rules is rules and cake_backend._router_memo() is memo rules[("sm_100a", "kk", 4096, 256)] = "forward" assert cake_backend._router_memo().forward_verdicts - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() assert cake_backend._router_memo() is not memo @@ -6660,7 +6471,7 @@ def test_k2_memo_is_the_planners_check(monkeypatch): cache; the memo starts over when full and still answers the same.""" parts = {0: ("sm_103a", 152), 1: ("sm_107a", 212), 2: ("sm_107a", 216)} _fake_device_facts(monkeypatch, parts) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() runs = cake_backend._k2_runs_without_verdict cases = [] for hidden in (2048, 4096, 5120, 6144): @@ -6706,7 +6517,7 @@ def check(): assert check() == kept # a full memo starts over (from an empty memo: every check is a miss) monkeypatch.setattr(cake_backend, "_ROUTER_MEMO_SIZE", 5) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() assert check() == kept assert 0 < len(cake_backend._router_memo().without_verdict) <= 5 @@ -6716,7 +6527,7 @@ def test_k2_memo_under_concurrent_table_changes(monkeypatch): tables stop changing every conclusion is the planner's for the final tables (a conclusion drawn while a table changed is not kept past that change).""" _fake_device_facts(monkeypatch, {0: ("sm_103a", 152)}) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() geometries = [] for T in (64, 128, 192, 224, 1024): G, X = torch.empty(T, 256), torch.empty(T, 4096) @@ -7031,29 +6842,31 @@ def _cpu_k2_launches(monkeypatch, arch="sm_103a"): def _router_entry_call(case): - """``(call, out)``: a router entry point's ``deterministic=False`` call of ``case`` (:data:`ROUTER_ENTRY_CASES`) + """``(call, out)``: a router entry point's call of ``case`` (:data:`ROUTER_ENTRY_CASES`) on fresh CPU tensors, and its output.""" op, T, H, splits, pad_out, pad_g, _ = ROUTER_ENTRY_CASES[case] X, W = torch.empty(T, H), torch.empty(256, H) G = torch.empty(T, 256 + 16)[:, :256] if pad_g else torch.empty(T, 256) rows, cols = {"fwd": (T, 256), "dgrad": (T, H), "wgrad": (256, H)}[op] out = torch.empty(rows, cols + 64)[:, :cols] if pad_out else torch.empty(rows, cols) - kw = ( - dict(deterministic=False) - if op == "dgrad" - else dict(deterministic=False, splits=splits) - ) + kw = dict() if op == "dgrad" else dict(splits=splits) call = { - "fwd": lambda: cake_backend.router_forward(X, W, out, **kw), - "dgrad": lambda: cake_backend.router_dgrad(G, W, out, **kw), - "wgrad": lambda: cake_backend.router_wgrad(G, X, out, **kw), + "fwd": lambda: cake_backend.cake_dense_gemm_nondet_train_router_forward( + X, W, out, **kw + ), + "dgrad": lambda: cake_backend.cake_dense_gemm_nondet_train_router_dgrad( + G, W, out, **kw + ), + "wgrad": lambda: cake_backend.cake_dense_gemm_nondet_train_router_wgrad( + G, X, out, **kw + ), }[op] return call, out @pytest.mark.parametrize("case", sorted(ROUTER_ENTRY_CASES)) def test_router_entry_points_launch_their_paths_launches(monkeypatch, case): - """With ``deterministic=False`` a router entry point's call goes through ``_K2_ENTRY_CACHE``: planned once per key + """A router entry point's call goes through ``_K2_ENTRY_CACHE``: planned once per key of its own tensors into the cached launch its own path takes (the same ``_K2_SPLITK_CACHE`` / ``_K2_ATOMIC_CACHE`` entry, under the same key), it issues that path's launches -- the same programs with the same arguments, the same split and zero fill -- and the next calls go from the key to the bound launch (no geometry, layout, key of @@ -7105,7 +6918,7 @@ def test_router_entry_points_launch_their_paths_launches(monkeypatch, case): @pytest.mark.parametrize("arch", ["sm_100a", "sm_103a"]) @pytest.mark.parametrize("T", [1, 16, 17, 128, 129, 192, 193]) def test_router_forward_entry_binds_the_simt_program(monkeypatch, arch, T): - """With ``deterministic=False`` and no split count the forward's call in a SIMT window of a 152-SM part binds the + """Without a split count the forward's call in a SIMT window of a 152-SM part binds the cluster program's launch in ``_K2_ENTRY_CACHE`` (planned once into ``_K2_ATOMIC_CACHE``): one launch per call of the program with the entry point's tensors -- ``X``, ``W`` and ``out`` read in place --, the launch's dimensions and row strides, the device's PDL flag and the grid of the cluster tiles; no split, no fill. Past the windows the @@ -7119,7 +6932,9 @@ def test_router_forward_entry_binds_the_simt_program(monkeypatch, arch, T): X = torch.empty(T, H) if spec is None else torch.empty(T, H + 8)[:, :H] W, out = torch.empty(E, H), torch.empty(T, E) for _ in range(3): - assert cake_backend.router_forward(X, W, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, out) is out + ) ((key, entry),) = cake_backend._K2_ENTRY_CACHE.items() assert key[2][0] == "fwd" and key[2][1] is None and key[2][2] == 1 assert not fill.calls and entry.launch is not None @@ -7145,7 +6960,7 @@ def test_router_forward_entry_binds_the_simt_program(monkeypatch, arch, T): @pytest.mark.parametrize("T", [1, 17, 64, 128, 129]) def test_router_dgrad_entry_binds_the_simt_program(monkeypatch, T): - """With ``deterministic=False`` the input gradient's call in a SIMT window of a 152-SM sm_100a part binds the SIMT + """The input gradient's call in a SIMT window of a 152-SM sm_100a part binds the SIMT program's launch in ``_K2_ENTRY_CACHE`` (planned once into ``_K2_ATOMIC_CACHE``): one launch per call of the program with the entry point's tensors -- ``G[None]``, ``W`` and ``out`` read in place --, the launch's dimensions and row strides, no PDL flag and the grid of ``out``'s tiles; no split, no fill. One token past the @@ -7157,7 +6972,7 @@ def test_router_dgrad_entry_binds_the_simt_program(monkeypatch, T): G, W = torch.empty(T, E), torch.empty(E, H) out = torch.empty(T, H + 64)[:, :H] for _ in range(3): - assert cake_backend.router_dgrad(G, W, out, deterministic=False) is out + assert cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out) is out ((key, entry),) = cake_backend._K2_ENTRY_CACHE.items() assert key[2][0] == "dgrad" and key[2][1] is None assert not fill.calls and not split.calls and entry.launch is not None @@ -7183,7 +6998,7 @@ def test_router_entry_cache_keys_and_routes_again(monkeypatch): the entry point's tensors, and routes a key again after a router table changes (a forward verdict set at run time runs this package's program, its removal the deterministic one) or a ``ROUTES`` value does; a call whose tensors are not on one CUDA device, or made while a router table is a plain ``dict``, takes the entry point's own path. - ``clear_convenience_caches`` drops the entries.""" + ``cake_dense_gemm_nondet_train_convenience_cache_clear`` drops the entries.""" arch, sm_count = "sm_103a", 152 X, W, out = torch.empty(64, 5120), torch.empty(256, 5120), torch.empty(64, 256) assert cake_backend._k2_entry("fwd", X, W, out, None, None) is None # CPU tensors @@ -7192,7 +7007,10 @@ def test_router_entry_cache_keys_and_routes_again(monkeypatch): def fwd(**kw): n = len(entry_fn.calls) - assert cake_backend.router_forward(X, W, out, deterministic=False, **kw) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, out, **kw) + is out + ) assert len(entry_fn.calls) > n return cache[next(reversed(cache))] @@ -7241,7 +7059,7 @@ def fwd(**kw): def wgrad(): n = len(entry_fn.calls) - assert cake_backend.router_wgrad(G, X, out, deterministic=False) is out + assert cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, out) is out assert len(entry_fn.calls) > n return cache[next(reversed(cache))] @@ -7258,9 +7076,9 @@ def wgrad(): m.setattr(cake_backend, "ROUTER_RULES", dict(cake_backend.ROUTER_RULES)) assert cake_backend._k2_entry("wgrad", G, X, out, None, None) is None n = len(entry_fn.calls) - assert cake_backend.router_wgrad(G, X, out, deterministic=False) is out + assert cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, out) is out assert len(entry_fn.calls) > n and len(cache) == size - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() assert not cache and not cake_backend._ATOMIC_ARENAS @@ -7278,14 +7096,14 @@ def wgrad(): @pytest.mark.parametrize("tracked", [True, False], ids=["verdict_tables", "dicts"]) -@pytest.mark.parametrize("path", ["entry_point", "router_fp32_gemm"]) +@pytest.mark.parametrize("path", ["entry_point", "cake_dense_gemm_nondet_train_router"]) @pytest.mark.parametrize("op", sorted(ROUTER_VERDICT_CHANGES)) def test_router_verdict_changed_at_run_time_plans_the_launch_again( monkeypatch, op, path, tracked ): """A router verdict changed between two calls of the same tensors: the next call runs the launch the planner gives under the new verdict, not the cached launch of the old one -- through the router entry point (its - ``_K2_ENTRY_CACHE`` entry and the ``_K2_ATOMIC_CACHE`` launch under it) and through ``router_fp32_gemm`` + ``_K2_ENTRY_CACHE`` entry and the ``_K2_ATOMIC_CACHE`` launch under it) and through ``cake_dense_gemm_nondet_train_router`` (``_K2_ATOMIC_CACHE``); while no table changes, a call takes its cached launch. With the router tables rebound to plain dicts (read on every call) a change is seen at the next call too. A launch with an explicit split count reads no router table: it is kept across a change.""" @@ -7314,16 +7132,22 @@ def test_router_verdict_changed_at_run_time_plans_the_launch_again( "wgrad": (G.t(), X, torch.empty(256, H)), }[op] entry_point = { - "fwd": lambda: cake_backend.router_forward(X, W, out, deterministic=False), - "dgrad": lambda: cake_backend.router_dgrad(G, W, out, deterministic=False), - "wgrad": lambda: cake_backend.router_wgrad(G, X, out, deterministic=False), + "fwd": lambda: cake_backend.cake_dense_gemm_nondet_train_router_forward( + X, W, out + ), + "dgrad": lambda: cake_backend.cake_dense_gemm_nondet_train_router_dgrad( + G, W, out + ), + "wgrad": lambda: cake_backend.cake_dense_gemm_nondet_train_router_wgrad( + G, X, out + ), }[op] def call(): """The cached launch a call ran, after checking that it launched.""" n = len(entry_fn.calls) - if path == "router_fp32_gemm": - assert cake_backend.router_fp32_gemm(A, B, out, deterministic=False) is out + if path == "cake_dense_gemm_nondet_train_router": + assert cake_backend.cake_dense_gemm_nondet_train_router(A, B, out) is out else: assert entry_point() is out assert len(entry_fn.calls) > n @@ -7331,7 +7155,7 @@ def call(): if path == "entry_point" and tracked: ((_, entry),) = cake_backend._K2_ENTRY_CACHE.items() assert entry.atomic and entry.launch is launch - else: # router_fp32_gemm, or an entry point while a router table is a plain dict + else: # cake_dense_gemm_nondet_train_router, or an entry point while a router table is a plain dict assert not cake_backend._K2_ENTRY_CACHE return launch @@ -7348,11 +7172,11 @@ def call(): assert (again is launch) == tracked and again.plan == want plans.append(want) assert plans[0] != plans[1] and plans[2] == plans[0] - if path == "router_fp32_gemm" and op != "dgrad": + if path == "cake_dense_gemm_nondet_train_router" and op != "dgrad": def with_splits(): assert ( - cake_backend.router_fp32_gemm(A, B, out, splits=2, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(A, B, out, splits=2) is out ) (kept,) = [ @@ -7407,13 +7231,17 @@ def plan(flags, bound): cake_backend._K2_SPLITK_CACHE.clear() path = Recorder() monkeypatch.setattr(cake_backend, "_atomic_k2", path) - assert cake_backend.router_dgrad(G, W, out, deterministic=False) is path.result + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out) is path.result + ) ((key, entry),) = cake_backend._K2_ENTRY_CACHE.items() assert entry.launch is None and key[2] == ("dgrad", None, 1) monkeypatch.setattr( cake_backend, "_plan_k2_entry", lambda *args: pytest.fail("planned again") ) - assert cake_backend.router_dgrad(G, W, out, deterministic=False) is path.result + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out) is path.result + ) assert len(path.calls) == 2 @@ -7421,7 +7249,7 @@ def test_router_entry_cache_eviction_keeps_the_arenas(monkeypatch): """An entry of ``_K2_ENTRY_CACHE`` uses its launch's arena, as the launch's own entry does; eviction (the cache bound) drops entries but no arena: an arena outlives its last cached entry with every buffer it holds (a CUDA graph may still replay a launch bound to them), a call on its stream plans into it again, and only - ``clear_convenience_caches`` drops it.""" + ``cake_dense_gemm_nondet_train_convenience_cache_clear`` drops it.""" _cpu_k2_launches(monkeypatch) monkeypatch.setattr(det_backend, "CONVENIENCE_CACHE_SIZE", 3) W = torch.empty(256, 6144) @@ -7433,7 +7261,7 @@ def test_router_entry_cache_eviction_keeps_the_arenas(monkeypatch): def dgrad(T): G, out = torch.empty(T, 256), torch.empty(T, 6144) - assert cake_backend.router_dgrad(G, W, out, deterministic=False) is out + assert cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out) is out arenas = cake_backend._ATOMIC_ARENAS entries = [e for c in caches for e in c.values()] assert all(arenas.get(e.arena.key) is e.arena for e in entries) @@ -7459,7 +7287,7 @@ def held(arena): monkeypatch.setattr(det_backend, "_current_raw_stream", lambda index: 7) assert dgrad(128) == {(0, 0): 4, (0, 7): 2} # planned into the same arena again assert cake_backend._ATOMIC_ARENAS[(0, 7)] is arena7 and held(arena7) == buffers7 - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() assert not cake_backend._ATOMIC_ARENAS and not any(caches) @@ -7537,14 +7365,14 @@ def test_k2_misaligned_out_takes_the_deterministic_path(monkeypatch, shift_elems X, W.t(), out, arch=arch, sm_count=152 ) assert plan is None and "16-byte" in reason - det = _patch_det(monkeypatch, "plan_router_fp32_gemm") + det = _patch_det(monkeypatch, "cake_dense_gemm_nondet_train_router_plan") monkeypatch.setattr( cake_backend, "_device_and_arch", lambda A, B, o: (torch.device("cpu"), arch, 152), ) assert ( - cake_backend.plan_router_fp32_gemm(X, W.t(), out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_plan(X, W.t(), out) is det.result ) ((args, kwargs),) = det.calls @@ -7570,7 +7398,7 @@ def _prepared_launch(monkeypatch, k: str, order_varying: bool): monkeypatch.setattr(det_backend, "_current_raw_stream", lambda index: 0) out, program = torch.zeros(8, 8), Recorder() if k == "k1": - plan = cake_backend.AtomicGemmPlan( + plan = cake_backend.CakeDenseGemmNondetProjectionPlan( L=1, M=4096, N=128, @@ -7584,11 +7412,11 @@ def _prepared_launch(monkeypatch, k: str, order_varying: bool): sm_pairs=76, template="k1_program", ) - launch = cake_backend.PreparedAtomicGemm( + launch = cake_backend.CakeDenseGemmNondetProjectionLaunch( "k1_module", "sm_103a", plan, out, {}, _launch=program ) return launch, (program,) - plan = cake_backend.AtomicRouterPlan( + plan = cake_backend.CakeDenseGemmNondetRouterPlan( M=5120, N=256, K=192, @@ -7603,7 +7431,7 @@ def _prepared_launch(monkeypatch, k: str, order_varying: bool): template="k2_program", ) split, fill = Recorder(), Recorder() - launch = cake_backend.PreparedAtomicRouterGemm( + launch = cake_backend.CakeDenseGemmNondetRouterLaunch( "k2_module", "sm_103a", plan, out, {}, _launch=program, _split=split, _fill=fill ) return launch, (split, fill, program) @@ -7627,7 +7455,10 @@ def test_an_order_varying_prepared_launch_alerts_under_torch_deterministic_mode( torch.use_deterministic_algorithms(True) with pytest.raises(RuntimeError, match=re.escape(text)) as info: launch.launch() - assert cake_backend._det_package_name() in str(info.value) + assert str(info.value).startswith("flashinfer.cake_dense_gemm_nondet_train: ") + assert "flashinfer.cake_dense_gemm_det_train is the deterministic module" in str( + info.value + ) assert "warn_only=True" in str(info.value) assert not any(r.calls for r in recorders) torch.use_deterministic_algorithms(True, warn_only=True) @@ -7670,7 +7501,10 @@ def test_the_cluster_program_does_not_alert(monkeypatch, torch_determinism): """The shipped cluster row ``4096 -> 64`` FP32 at 2048 tokens on a 152-SM ``sm_100a`` part (three splits, added in cluster-rank order) runs without an alert under ``torch.use_deterministic_algorithms(True)``: the planning call, the call from the cached key and the prepared launch (CPU binding, nothing loaded or launched).""" + table = cake_backend.K1_DSM_RANGES_BY_SM entry_fn = _cpu_k1_launches(monkeypatch, "sm_100a") + # the harness takes the cluster ranges out; this test runs the shipped one + monkeypatch.setattr(cake_backend, "K1_DSM_RANGES_BY_SM", table) monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) bf16, T = torch.bfloat16, 2048 X, G = torch.empty(T, 4096, dtype=bf16), torch.empty(T, 64, dtype=bf16) @@ -7679,10 +7513,11 @@ def test_the_cluster_program_does_not_alert(monkeypatch, torch_determinism): with _no_warnings(): for _ in range(2): assert ( - cake_backend.projection_wgrad(G, X, out_w, deterministic=False) is out_w + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out_w) + is out_w ) - prepared = cake_backend.prepare_projection_wgrad( - G, X, out_w, deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out_w ) assert prepared.launch() is out_w ((_, atomic),) = cake_backend._K1_ATOMIC_CACHE.items() @@ -7695,7 +7530,7 @@ def test_the_cluster_program_does_not_alert(monkeypatch, torch_determinism): def test_k1_convenience_launch_alerts_only_when_order_varying( monkeypatch, torch_determinism, splits ): - """``projection_wgrad(deterministic=False)`` on a row this package's planner takes at ``splits`` splits (CPU + """``cake_dense_gemm_nondet_train_projection_wgrad()`` on a row this package's planner takes at ``splits`` splits (CPU binding, nothing loaded or launched), under ``torch.use_deterministic_algorithms(True)``: at three splits the planning call and the call from the cached key both raise before the program; with ``warn_only=True`` the call warns and launches; at two splits both calls launch without an alert. The cached launch keeps whether its plan @@ -7709,7 +7544,7 @@ def test_k1_convenience_launch_alerts_only_when_order_varying( out = torch.zeros(128, 4096) def call(): - return cake_backend.projection_wgrad(G, X, out, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.use_deterministic_algorithms(True) for _ in range(2): # the planning call, then the call from the cached key @@ -7734,7 +7569,7 @@ def call(): def test_k2_convenience_launch_alerts_only_when_order_varying( monkeypatch, torch_determinism, splits ): - """``router_wgrad(deterministic=False, splits=splits)`` (CPU binding, nothing loaded or launched) under + """``cake_dense_gemm_nondet_train_router_wgrad(splits=splits)`` (CPU binding, nothing loaded or launched) under ``torch.use_deterministic_algorithms(True)``: three units per tile (three partials per element) raise before the split, the fill of ``out`` and the program, on the planning call and from the cached key; with ``warn_only=True`` the call warns and launches; two units per tile launch without an alert.""" @@ -7746,7 +7581,9 @@ def test_k2_convenience_launch_alerts_only_when_order_varying( recorders = (entry_fn, fill, split) def call(): - return cake_backend.router_wgrad(G, X, out, splits=splits, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_router_wgrad( + G, X, out, splits=splits + ) torch.use_deterministic_algorithms(True) for _ in range(2): # the planning call, then the call from the cached key @@ -7888,7 +7725,9 @@ def test_registered_templates_are_the_ones_the_planners_name(): ) -def test_no_det_token_in_generated_names_and_no_negated_token_in_the_package(): +def test_no_det_token_in_generated_names(): + """No generated name of this package -- file, program, kernel, template, library -- carries the token + ``det``.""" root = Path(cake_jit.__file__).parent files = [ p @@ -7898,7 +7737,6 @@ def test_no_det_token_in_generated_names_and_no_negated_token_in_the_package(): for path in files: rel = path.relative_to(root).as_posix() assert not DET_TOKEN.search(path.name), rel - assert not NEGATED_TOKEN.search(path.read_text(errors="replace")), rel generated = [*cake_jit.MODULES, *(r["kernel"] for r in cake_jit.MODULES.values())] generated += [t for table in cake_jit.KERNELS.values() for t in table] assert not [n for n in generated if DET_TOKEN.search(n)] @@ -7906,7 +7744,9 @@ def test_no_det_token_in_generated_names_and_no_negated_token_in_the_package(): if manifest.exists(): document = json.loads(manifest.read_text()) # the in-tree import path: a vendored copy keeps its manifest - assert document["package"] == "flashinfer.experimental.dense_projection_gemm" + assert document["package"] == ( + "flashinfer.experimental.cake.cake_dense_gemm_nondet_train" + ) for entry in document["libraries"]: assert not DET_TOKEN.search(entry["name"]) assert all( @@ -7915,8 +7755,6 @@ def test_no_det_token_in_generated_names_and_no_negated_token_in_the_package(): or s.startswith("csrc/cake_device_helpers/") for s in entry["sources"] ) - for name in ENTRY_POINTS: - assert not DET_TOKEN.search(name) def test_select_module_and_the_loader_take_this_package_s_programs_only(monkeypatch): @@ -7957,7 +7795,7 @@ def test_clear_convenience_caches_clears_both_packages(): torch.device("cpu"), (99, 99) ) try: - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() assert ( key not in det_backend._K2_CACHE and key not in cake_backend._K1_ATOMIC_CACHE @@ -8002,7 +7840,7 @@ def test_the_split_k_workspace_kinds(): def test_arena_growth_retires_the_replaced_buffers(): """A growth at least doubles a kind's buffer and retires the buffer it replaces (a CUDA graph or a launch in - flight may still use it): the retired buffers keep their storage and contents until ``clear_convenience_caches`` + flight may still use it): the retired buffers keep their storage and contents until ``cake_dense_gemm_nondet_train_convenience_cache_clear`` drops the arena, and hold fewer bytes than the live one; a requirement the live buffer meets grows nothing.""" arena = cake_backend._AtomicArena(torch.device("cpu"), (0, 0)) first = arena.require("k1_tickets", 100) @@ -8091,10 +7929,17 @@ def _profile_kernel_names(fn): ) +# The generated stem of the deterministic package's programs: its library names ``cake__``, kernel +# symbols ``kernel_cake__`` and template bodies' namespace ``cake::::``. +(_DET_STEM,) = { + name.rsplit("_", 1)[0].removeprefix("cake_") for name in det_jit.MODULES +} or {"no deterministic program registered"} + + def _det_kernel(name: str) -> bool: """A K1 / K2 GEMM kernel of the deterministic package (its programs' symbols and its template bodies' namespace); the shared BF16x3 split stage is not one.""" - return "dense_projection_det_gemm" in name and not any( + return _DET_STEM in name and not any( kernel in name for kernel in _DET_SHARED_KERNELS ) @@ -8156,16 +8001,15 @@ def _k1_prepare(G, X, out, arch, splits, block_n=None): _require_template( arch, cake_backend.atomic_k1_template(block_n, out.dtype == torch.float32) ) - prepared = cake_backend.prepare_dense_projection_gemm( + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_prepare( X.t(), G, out, transposed_out=True, - deterministic=False, splits=splits, block_n=block_n, ) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch) assert (prepared.splits, prepared.plan.block_n) == (splits, block_n) return prepared @@ -8192,7 +8036,7 @@ def _force_k1_rule(monkeypatch, G, X, out, splits, block_n=None): monkeypatch.setattr(cake_backend, "K1_CHAIN_RANGES_BY_SM", {}) monkeypatch.setattr(cake_backend, "K1_ONE_SPLIT_RANGES_BY_SM", {}) monkeypatch.setattr(cake_backend, "K1_DSM_RANGES_BY_SM", {}) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() return cake_backend.atomic_k1_template(block_n, f32) @@ -8201,10 +8045,10 @@ def no_k1_dsm(monkeypatch): """Take the cluster ranges of :data:`K1_DSM_RANGES_BY_SM` out of the planner for a test of the K1 tables it comes before (:func:`plan_atomic_gemm` reads it first; ``test_gpu_k1_152_sm_dsm_ranges_run_at_their_edges`` holds it), with the convenience caches cleared before and after: no entry planned without the table outlives the test.""" - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() monkeypatch.setattr(cake_backend, "K1_DSM_RANGES_BY_SM", {}) yield - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() @pytest.fixture @@ -8213,7 +8057,7 @@ def k1_rule(monkeypatch): cluster ranges or tables of exactly one SM count), so the convenience path runs this package's program at that split count; returns the template.""" yield functools.partial(_force_k1_rule, monkeypatch) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() def _arena_tickets_zero(): @@ -8239,7 +8083,9 @@ def test_gpu_k1_budgets_over_reruns(row, out_dtype, splits, k1_rule): device, arch = _atomic_device_or_skip() G, X, out = _k1_wgrad_inputs(row, 16231, out_dtype, device, SEED) ref = G.double().t() @ X.double() - det_out = det_backend.projection_det_wgrad(G, X, torch.empty_like(out)) + det_out = det_backend.cake_dense_gemm_det_train_projection_wgrad( + G, X, torch.empty_like(out) + ) prepared = _k1_prepare(G, X, out, arch, splits) _assert_ran(prepared.launch, prepared.module_name) k1_rule(G, X, out, splits) @@ -8255,12 +8101,12 @@ def test_gpu_k1_budgets_over_reruns(row, out_dtype, splits, k1_rule): first = out.clone() if first is None else first assert torch.equal(out, first), i out.fill_(float("nan")) - cake_backend.projection_wgrad(G, X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() _k1_check(out, ref, f"convenience rerun {i}") assert _arena_tickets_zero() _assert_ran( - lambda: cake_backend.projection_wgrad(G, X, out, deterministic=False), + lambda: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out), prepared.module_name, ) @@ -8306,16 +8152,15 @@ def test_gpu_k1_row_major_program(shape, transposed_out, k1_rule): out = torch.full((d_out, d_in), float("nan"), device=device, dtype=torch.float32) ref = G.double().t() @ X.double() A, B = (X.t(), G) if transposed_out else (G.t(), X) - prepared = cake_backend.prepare_dense_projection_gemm( + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_prepare( A, B, out, transposed_out=transposed_out, - deterministic=False, splits=splits, block_n=192, ) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch) assert prepared.template == template and not prepared.plan.out_t assert prepared.plan.swapped == transposed_out _assert_ran(prepared.launch, prepared.module_name) @@ -8333,7 +8178,7 @@ def test_gpu_k1_row_major_program(shape, transposed_out, k1_rule): assert k1_rule(G, X, out, splits, block_n=192) == template for i in range(2): out.fill_(float("nan")) - cake_backend.projection_wgrad(G, X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() _k1_check(out, ref, f"convenience rerun {i}") assert _arena_tickets_zero() @@ -8366,16 +8211,15 @@ def test_gpu_k1_chain_programs(shape, monkeypatch): g = torch.Generator(device=device).manual_seed(SEED + 11) G, X = _bf16((T, d_out), device, g), _bf16((T, d_in), device, g) out = torch.full((d_out, d_in), float("nan"), device=device, dtype=torch.float32) - prepared = cake_backend.prepare_dense_projection_gemm( + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_prepare( X.t(), G, out, transposed_out=True, - deterministic=False, block_n=block_n, chains=3, ) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch) assert prepared.template == template assert (prepared.splits, prepared.plan.chains) == (1, 3) _assert_ran(prepared.launch, prepared.module_name) @@ -8406,26 +8250,27 @@ def test_gpu_k1_chain_programs(shape, monkeypatch): {(arch, sm_count): {key: ((T, T, (block_n, 3)),)}}, ) monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() try: out.fill_(float("nan")) - cake_backend.projection_wgrad(G, X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() assert torch.equal(out, first) _assert_ran( - lambda: cake_backend.projection_wgrad(G, X, out, deterministic=False), + lambda: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + G, X, out + ), prepared.module_name, ) finally: - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() short = torch.zeros(128, d_out, device=device, dtype=torch.bfloat16) with pytest.raises(ValueError, match="at one split"): - cake_backend.prepare_dense_projection_gemm( + cake_backend.cake_dense_gemm_nondet_train_projection_prepare( torch.zeros(128, d_in, device=device, dtype=torch.bfloat16).t(), short, torch.empty(d_out, d_in, device=device), transposed_out=True, - deterministic=False, block_n=block_n, chains=3, ) @@ -8448,10 +8293,10 @@ def test_gpu_k1_batched_mla_head_views(): (H, D_in, D_out), float("nan"), device=device, dtype=torch.bfloat16 ) ref = A.double() @ B.double() - prepared = cake_backend.prepare_dense_projection_gemm( - A, B, out, deterministic=False, splits=2, block_n=256 + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_prepare( + A, B, out, splits=2, block_n=256 ) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch) assert (prepared.plan.L, prepared.splits, prepared.plan.swapped) == (H, 2, True) for _ in range(3): out.fill_(float("nan")) @@ -8461,7 +8306,7 @@ def test_gpu_k1_batched_mla_head_views(): assert prepared.counters_are_zero() # the planner keeps the batched rows on the deterministic program: the default convenience path matches it out2 = torch.empty_like(out) - cake_backend.dense_projection_gemm(A, B, out2, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection(A, B, out2) torch.cuda.synchronize() _k1_check(out2, ref) @@ -8481,12 +8326,14 @@ def test_gpu_k1_ragged_token_counts_reuse_poisoned_scratch(k1_rule): for arena in cake_backend._ATOMIC_ARENAS.values(): if "k1_slots" in arena.buffers: arena.buffers["k1_slots"].fill_(float("nan")) - cake_backend.projection_wgrad(G, X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() _k1_check(out, G.double().t() @ X.double(), f"T={T}") assert _arena_tickets_zero(), T if not cake_backend.k1_splits_fit(3, T): - assert torch.equal(out, det_backend.projection_det_wgrad(G, X)) + assert torch.equal( + out, det_backend.cake_dense_gemm_det_train_projection_wgrad(G, X) + ) assert len(cake_backend._ATOMIC_ARENAS) == 1 (arena,) = cake_backend._ATOMIC_ARENAS.values() assert arena.version >= 2 # created, then grown or reused; never shrunk @@ -8505,7 +8352,7 @@ def test_gpu_k1_small_token_counts_guard(): def test_gpu_torch_deterministic_mode_alerts_before_an_order_varying_launch( k1_rule, torch_determinism ): - """Under ``torch.use_deterministic_algorithms(True)`` a ``deterministic=False`` call whose launch is fixed order + """Under ``torch.use_deterministic_algorithms(True)`` a call whose launch is fixed order (two splits: one add) runs without an alert and repeats bit for bit; a launch whose result varies run to run (three splits) raises PyTorch's ``RuntimeError`` before it launches (``out`` untouched), or with ``warn_only=True`` warns and runs this package's program within the accuracy rule. The setting never changes @@ -8517,19 +8364,23 @@ def test_gpu_torch_deterministic_mode_alerts_before_an_order_varying_launch( k1_rule(G, X, out, 2) torch.use_deterministic_algorithms(True) with _no_warnings(): - first = cake_backend.projection_wgrad(G, X, out, deterministic=False).clone() + first = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + G, X, out + ).clone() for _ in range(RERUNS): - again = cake_backend.projection_wgrad(G, X, out, deterministic=False) + again = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + G, X, out + ) assert torch.equal(again, first) k1_rule(G, X, out, 3) out.fill_(7.0) with pytest.raises(RuntimeError, match="at 3 splits"): - cake_backend.projection_wgrad(G, X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() assert bool((out == 7.0).all()) torch.use_deterministic_algorithms(True, warn_only=True) with pytest.warns(UserWarning, match="at 3 splits"): - got = cake_backend.projection_wgrad(G, X, out, deterministic=False) + got = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() torch.use_deterministic_algorithms(False) _k1_check(first, ref, "two splits") @@ -8545,8 +8396,10 @@ def test_gpu_det_then_atomic_share_a_stream_not_a_workspace(k1_rule): k1_rule(G, X, out, 2) ref = G.double().t() @ X.double() for _ in range(3): - det = det_backend.projection_det_wgrad(G, X, torch.empty_like(out)) - atomic = cake_backend.projection_wgrad(G, X, out, deterministic=False) + det = det_backend.cake_dense_gemm_det_train_projection_wgrad( + G, X, torch.empty_like(out) + ) + atomic = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) torch.cuda.synchronize() _k1_check(det, ref) _k1_check(atomic, ref) @@ -8559,7 +8412,7 @@ def test_gpu_det_then_atomic_share_a_stream_not_a_workspace(k1_rule): def test_gpu_k1_deterministic_program_through_the_binding(): - """``deterministic=False`` K1 calls that run the deterministic program -- the forward, the input gradient (BF16 + """K1 calls of this package that run the deterministic program -- the forward, the input gradient (BF16 and FP32 out), the weight gradient of rows without a verdict (both orientations, BF16 and FP32 out, another ``sk``), the batched head views of the generic entry -- go through this package's binding and give the deterministic package's bits: dense and padded outputs (the padding untouched), fresh tensors on every call, a @@ -8568,7 +8421,7 @@ def test_gpu_k1_deterministic_program_through_the_binding(): small-token program stays unbound (that package's call); a binding the deterministic package refuses raises its error.""" device = _det_device_or_skip() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() rows = (("kv_a", 576, 6144), ("4096to128", 128, 4096), ("wide", 2048, 4096)) ran, swapped = set(), set() @@ -8580,21 +8433,21 @@ def check(T, seed, pad_out=False): _bf16((N, K), device, g, K**-0.5), _bf16((T, N), device, g), ) - # op -> (out shape, out dtype, the deterministic package's call, this package's deterministic=False call) + # op -> (out shape, out dtype, the deterministic package's call, this package's call) # fmt: off calls = { - "fwd": ((T, N), torch.bfloat16, lambda o: det_backend.projection_det_forward(X, W, o), - lambda o: cake_backend.projection_forward(X, W, o, deterministic=False)), - "dgrad": ((T, K), torch.bfloat16, lambda o: det_backend.projection_det_dgrad(G, W, o), - lambda o: cake_backend.projection_dgrad(G, W, o, deterministic=False)), - "dgrad_f32": ((T, K), torch.float32, lambda o: det_backend.projection_det_dgrad(G, W, o), - lambda o: cake_backend.projection_dgrad(G, W, o, deterministic=False)), - "wgrad": ((N, K), torch.bfloat16, lambda o: det_backend.projection_det_wgrad(G, X, o), - lambda o: cake_backend.projection_wgrad(G, X, o, deterministic=False)), - "wgrad_f32": ((N, K), torch.float32, lambda o: det_backend.projection_det_wgrad(G, X, o), - lambda o: cake_backend.projection_wgrad(G, X, o, deterministic=False)), - "wgrad_sk": ((N, K), torch.float32, lambda o: det_backend.projection_det_wgrad(G, X, o, sk=False), - lambda o: cake_backend.projection_wgrad(G, X, o, sk=False, deterministic=False)), + "fwd": ((T, N), torch.bfloat16, lambda o: det_backend.cake_dense_gemm_det_train_projection_forward(X, W, o), + lambda o: cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, o)), + "dgrad": ((T, K), torch.bfloat16, lambda o: det_backend.cake_dense_gemm_det_train_projection_dgrad(G, W, o), + lambda o: cake_backend.cake_dense_gemm_nondet_train_projection_dgrad(G, W, o)), + "dgrad_f32": ((T, K), torch.float32, lambda o: det_backend.cake_dense_gemm_det_train_projection_dgrad(G, W, o), + lambda o: cake_backend.cake_dense_gemm_nondet_train_projection_dgrad(G, W, o)), + "wgrad": ((N, K), torch.bfloat16, lambda o: det_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, o), + lambda o: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, o)), + "wgrad_f32": ((N, K), torch.float32, lambda o: det_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, o), + lambda o: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, o)), + "wgrad_sk": ((N, K), torch.float32, lambda o: det_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, o, sk=False), + lambda o: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, o, sk=False)), } # fmt: on arch = det_backend.arch_for(device) @@ -8635,11 +8488,10 @@ def check(T, seed, pad_out=False): ) if det_backend._swap_ab_launch(A, B, out, False) is not None: swapped.add((T, shape)) - want = det_backend.dense_projection_det_gemm(A, B, out).clone() + want = det_backend.cake_dense_gemm_det_train_projection(A, B, out).clone() out.fill_(float("nan")) assert ( - cake_backend.dense_projection_gemm(A, B, out, deterministic=False) - is out + cake_backend.cake_dense_gemm_nondet_train_projection(A, B, out) is out ) assert torch.equal(out, want), (T, tuple(shape)) @@ -8680,9 +8532,9 @@ def check(T, seed, pad_out=False): Xm = storage[1 : 1 + 1001 * 4096].view(1001, 4096) Xm.copy_(X) with pytest.raises(ValueError, match="16-byte aligned") as refused: - det_backend.projection_det_forward(Xm, W) + det_backend.cake_dense_gemm_det_train_projection_forward(Xm, W) with pytest.raises(ValueError, match=re.escape(str(refused.value))): - cake_backend.projection_forward(Xm, W, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_forward(Xm, W) def test_gpu_k1_bound_launch_is_the_unbound_launch(k1_rule): @@ -8697,7 +8549,10 @@ def test_gpu_k1_bound_launch_is_the_unbound_launch(k1_rule): want = prepared.out.clone() for i in range(3): G2, X2, out2 = G.clone(), X.clone(), torch.full_like(out, float("nan")) - assert cake_backend.projection_wgrad(G2, X2, out2, deterministic=False) is out2 + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G2, X2, out2) + is out2 + ) torch.cuda.synchronize() assert torch.equal(out2, want), i assert _arena_tickets_zero() @@ -8740,8 +8595,12 @@ def test_gpu_k1_pdl_waits_for_an_early_producer(splits, k1_rule): for _ in range(3): G.fill_(float("nan")) out.fill_(float("nan")) - det_backend.projection_det_forward(H, Wg, G, enable_pdl=True) - cake_backend.projection_wgrad(G, X, out, deterministic=False, enable_pdl=True) + det_backend.cake_dense_gemm_det_train_projection_forward( + H, Wg, G, enable_pdl=True + ) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + G, X, out, enable_pdl=True + ) torch.cuda.synchronize() _k1_check(out, G.double().t() @ X.double()) assert _arena_tickets_zero() @@ -8750,10 +8609,10 @@ def test_gpu_k1_pdl_waits_for_an_early_producer(splits, k1_rule): def test_gpu_k1_bound_deterministic_program_waits_for_an_early_producer(): """The deterministic program through this package's binding as a programmatic dependent launch: the deterministic forward before it (K = 8192) triggers its dependents at entry and writes ``G``; the input gradient - with ``deterministic=False`` (bound, ``enable_pdl=True``) starts while it runs and must still read every row it + (bound, ``enable_pdl=True``) starts while it runs and must still read every row it wrote -- the deterministic package's bits -- and its slice counters end zero.""" device = _det_device_or_skip() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() T, N, K = 4096, 128, 4096 g = torch.Generator(device=device).manual_seed(SEED + 61) H = _bf16((T, 8192), device, g) @@ -8761,15 +8620,19 @@ def test_gpu_k1_bound_deterministic_program_waits_for_an_early_producer(): W = _bf16((N, K), device, g, K**-0.5) G = torch.empty(T, N, device=device, dtype=torch.bfloat16) dX = torch.empty(T, K, device=device, dtype=torch.bfloat16) - det_backend.projection_det_forward(H, Wg, G, enable_pdl=False) - want = det_backend.projection_det_dgrad(G, W, enable_pdl=False).clone() + det_backend.cake_dense_gemm_det_train_projection_forward(H, Wg, G, enable_pdl=False) + want = det_backend.cake_dense_gemm_det_train_projection_dgrad( + G, W, enable_pdl=False + ).clone() for _ in range(3): G.fill_(float("nan")) dX.fill_(float("nan")) - det_backend.projection_det_forward(H, Wg, G, enable_pdl=True) + det_backend.cake_dense_gemm_det_train_projection_forward( + H, Wg, G, enable_pdl=True + ) assert ( - cake_backend.projection_dgrad( - G, W, dX, deterministic=False, enable_pdl=True + cake_backend.cake_dense_gemm_nondet_train_projection_dgrad( + G, W, dX, enable_pdl=True ) is dX ) @@ -8828,7 +8691,7 @@ def _force_router_rule(monkeypatch, op, A, B, spec): tables = dict(cake_backend.ROUTER_RULES_BY_SM) tables.pop((arch, sm_count), None) monkeypatch.setattr(cake_backend, "ROUTER_RULES_BY_SM", tables) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() @pytest.fixture @@ -8836,7 +8699,7 @@ def router_rule(monkeypatch): """``router_rule(A, B, spec)``: a verdict for the router row of ``(A, B)``, so the default plan runs this package's program with the partition ``spec`` (or the deterministic program with ``None``).""" yield functools.partial(_force_router_rule, monkeypatch) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() def _router_prepare(A, B, out, splits, arch, spec=None): @@ -8850,10 +8713,10 @@ def _router_prepare(A, B, out, splits, arch, spec=None): else "" ), ) - prepared = cake_backend.prepare_router_fp32_gemm( - A, B, out, splits=splits, deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare( + A, B, out, splits=splits ) - assert isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch) return prepared @@ -8877,7 +8740,9 @@ def _router_segment_reference(A, B, plan): min(K, k1 * cake_backend.ROUTER_BLOCK_K), ) part = torch.empty(int(A.shape[0]), int(B.shape[1]), device=A.device) - det_backend.router_fp32_det_gemm(A[:, lo:hi], B[lo:hi, :], part, splits=1) + det_backend.cake_dense_gemm_det_train_router( + A[:, lo:hi], B[lo:hi, :], part, splits=1 + ) partials[(k0, k1)] = part.t() if plan.swapped else part # the launch's [M, N] expected = torch.empty(plan.M, plan.N, device=A.device) bound = torch.zeros(plan.M, plan.N, device=A.device, dtype=torch.float64) @@ -8982,7 +8847,7 @@ def _router_gate(op, splits, spec, router_rule, T, hidden, seed): assert worst <= ROUTER_REL_FRO_MAX, worst if splits == 1: det_out = torch.empty_like(out) - det_backend.router_fp32_det_gemm(A, B, det_out, splits=1) + det_backend.cake_dense_gemm_det_train_router(A, B, det_out, splits=1) assert torch.equal(out, det_out) return plan @@ -9057,8 +8922,8 @@ def test_gpu_router_tmem_head_tail_program_in_its_windows(hidden, T): ) assert plan is None, plan assert reason == "the planner keeps the row on the deterministic program" - prepared = cake_backend.prepare_router_fp32_gemm(A, B, out, deterministic=False) - assert not isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(A, B, out) + assert not isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch) return _require_template(arch, ts) _require_template(arch, c3) @@ -9066,8 +8931,8 @@ def test_gpu_router_tmem_head_tail_program_in_its_windows(hidden, T): A, B, out, arch=arch, sm_count=sm_count ) assert plan is not None and plan.template == ts and plan.head_spans > 0, reason - prepared = cake_backend.prepare_router_fp32_gemm(A, B, out, deterministic=False) - assert isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(A, B, out) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch) assert prepared.plan == plan assert prepared.module_name == cake_jit.select_module(arch, ts) ref = A.double() @ B.double() @@ -9099,7 +8964,7 @@ def test_gpu_router_tmem_head_tail_program_in_its_windows(hidden, T): launch.launch() torch.cuda.synchronize() det_out = torch.empty_like(out) - det_backend.router_fp32_det_gemm(A, B, det_out, splits=1) + det_backend.cake_dense_gemm_det_train_router(A, B, det_out, splits=1) assert torch.equal(stores[ts], stores[c3]) and torch.equal(stores[ts], det_out) @@ -9133,7 +8998,7 @@ def test_gpu_router_simt_program_in_its_window(router_rule, tile, hidden, T): """A SIMT window (forced for the row on this device) runs the SIMT program of its tile and no other program: over ``RERUNS`` launches into a NaN-poisoned ``out`` the weight gradient is finite, within ``ROUTER_REL_FRO_MAX`` of FP64 and the same bits every time (every element stored once, its slices added in a - fixed order), and the prepared launch, the cached launch of ``router_fp32_gemm``, ``router_wgrad`` (its entry + fixed order), and the prepared launch, the cached launch of ``cake_dense_gemm_nondet_train_router``, ``cake_dense_gemm_nondet_train_router_wgrad`` (its entry cache) and a padded ``out`` (written in place, its padding untouched) give those bits.""" device, arch, template = _simt_device_or_skip(tile) sm_count = det_backend.device_facts(device).sm_count @@ -9144,8 +9009,8 @@ def test_gpu_router_simt_program_in_its_window(router_rule, tile, hidden, T): A, B, out, arch=arch, sm_count=sm_count ) assert plan is not None and plan.simt == tile and plan.template == template, reason - prepared = cake_backend.prepare_router_fp32_gemm(A, B, out, deterministic=False) - assert isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(A, B, out) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch) assert prepared.plan == plan and not prepared.fused_split assert prepared.module_name == cake_jit.select_module(arch, template) _assert_ran(prepared.launch, prepared.module_name) @@ -9163,17 +9028,17 @@ def test_gpu_router_simt_program_in_its_window(router_rule, tile, hidden, T): assert torch.equal(out, first), i for _ in range(2): cached = torch.full_like(out, float("nan")) - cake_backend.router_fp32_gemm(A, B, cached, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(A, B, cached) entry = torch.full_like(out, float("nan")) - cake_backend.router_wgrad(G, X, entry, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, entry) torch.cuda.synchronize() assert torch.equal(cached, first) and torch.equal(entry, first) _assert_ran( - lambda: cake_backend.router_wgrad(G, X, entry, deterministic=False), + lambda: cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, entry), prepared.module_name, ) padded = torch.full((256, hidden + 4), float("nan"), device=device) - cake_backend.router_wgrad(G, X, padded[:, :hidden], deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, padded[:, :hidden]) torch.cuda.synchronize() assert torch.equal(padded[:, :hidden], first) assert bool(padded[:, hidden:].isnan().all()) @@ -9181,7 +9046,7 @@ def test_gpu_router_simt_program_in_its_window(router_rule, tile, hidden, T): @pytest.mark.parametrize("tile", SIMT_TILES, ids=_simt_spec) def test_gpu_router_simt_program_replays_in_a_cuda_graph(router_rule, tile): - """``router_wgrad`` in a SIMT window captured into a CUDA graph (its entry cache, the prepared launch) replays the + """``cake_dense_gemm_nondet_train_router_wgrad`` in a SIMT window captured into a CUDA graph (its entry cache, the prepared launch) replays the eager bits into a NaN-poisoned ``out``; the activations and the gradient changed between replays are read again.""" device, arch, template = _simt_device_or_skip(tile) @@ -9190,26 +9055,26 @@ def test_gpu_router_simt_program_replays_in_a_cuda_graph(router_rule, tile): A, B, out = _router_case("wgrad", X, W, G) router_rule("wgrad", A, B, _simt_spec(tile)) eager = torch.empty_like(out) - cake_backend.router_wgrad(G, X, eager, deterministic=False) - prepared = cake_backend.prepare_router_fp32_gemm(G.t(), X, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, eager) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(G.t(), X, out) assert prepared.plan.simt == tile entry_out = torch.empty_like(out) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(2): - cake_backend.router_wgrad(G, X, entry_out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, entry_out) prepared.launch() torch.cuda.current_stream().wait_stream(stream) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - cake_backend.router_wgrad(G, X, entry_out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, entry_out) prepared.launch() for scale in (1.0, -0.5): X.mul_(scale) want = torch.empty_like(out) - cake_backend.router_wgrad(G, X, want, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, want) entry_out.fill_(float("nan")) out.fill_(float("nan")) graph.replay() @@ -9239,11 +9104,15 @@ def test_gpu_router_simt_program_waits_for_an_early_producer(router_rule): dX = torch.empty(T, hidden, device=device) dW = torch.empty(256, hidden, device=device) router_rule("wgrad", Y.t(), dX, _simt_spec(tile)) - want_y = det_backend.router_det_forward(X, W, enable_pdl=False) - want_dx = det_backend.router_det_dgrad(want_y, W, enable_pdl=False) + want_y = det_backend.cake_dense_gemm_det_train_router_forward( + X, W, enable_pdl=False + ) + want_dx = det_backend.cake_dense_gemm_det_train_router_dgrad( + want_y, W, enable_pdl=False + ) want_dw = torch.empty_like(dW) - cake_backend.router_wgrad( - want_y, want_dx, want_dw, deterministic=False, enable_pdl=False + cake_backend.cake_dense_gemm_nondet_train_router_wgrad( + want_y, want_dx, want_dw, enable_pdl=False ) torch.cuda.synchronize() assert ( @@ -9253,10 +9122,10 @@ def test_gpu_router_simt_program_waits_for_an_early_producer(router_rule): for _ in range(3): for t in (Y, dX, dW): t.fill_(float("nan")) - kw = dict(deterministic=False, enable_pdl=pdl) - cake_backend.router_forward(X, W, Y, **kw) - cake_backend.router_dgrad(Y, W, dX, **kw) - cake_backend.router_wgrad(Y, dX, dW, **kw) + kw = dict(enable_pdl=pdl) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, Y, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(Y, W, dX, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(Y, dX, dW, **kw) torch.cuda.synchronize() assert torch.equal(Y, want_y) and torch.equal(dX, want_dx), pdl assert torch.equal(dW, want_dw), pdl @@ -9285,7 +9154,7 @@ def test_gpu_router_dgrad_simt_program_in_its_window(router_rule, tile, hidden, """An input-gradient SIMT window (forced for the row on this device) runs the SIMT program of its tile and no other program: over ``RERUNS`` launches into a NaN-poisoned ``out`` the input gradient is finite (every row of a partial last row tile written), within ``ROUTER_REL_FRO_MAX`` of FP64 and the same bits every time; the prepared - launch, the cached launch of ``router_fp32_gemm``, ``router_dgrad`` (its entry cache), operands with padded rows + launch, the cached launch of ``cake_dense_gemm_nondet_train_router``, ``cake_dense_gemm_nondet_train_router_dgrad`` (its entry cache), operands with padded rows (NaN in the padding, read in place) and a padded ``out`` (its padding untouched) give those bits.""" device, arch, template = _dgrad_simt_device_or_skip(tile) sm_count = det_backend.device_facts(device).sm_count @@ -9296,8 +9165,8 @@ def test_gpu_router_dgrad_simt_program_in_its_window(router_rule, tile, hidden, A, B, out, arch=arch, sm_count=sm_count ) assert plan is not None and plan.simt == tile and plan.template == template, reason - prepared = cake_backend.prepare_router_fp32_gemm(A, B, out, deterministic=False) - assert isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(A, B, out) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch) assert prepared.plan == plan and not prepared.fused_split assert prepared.module_name == cake_jit.select_module(arch, template) _assert_ran(prepared.launch, prepared.module_name) @@ -9315,13 +9184,13 @@ def test_gpu_router_dgrad_simt_program_in_its_window(router_rule, tile, hidden, assert torch.equal(out, first), i for _ in range(2): cached = torch.full_like(out, float("nan")) - cake_backend.router_fp32_gemm(A, B, cached, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(A, B, cached) entry = torch.full_like(out, float("nan")) - cake_backend.router_dgrad(G, W, entry, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, entry) torch.cuda.synchronize() assert torch.equal(cached, first) and torch.equal(entry, first) _assert_ran( - lambda: cake_backend.router_dgrad(G, W, entry, deterministic=False), + lambda: cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, entry), prepared.module_name, ) G_pad = torch.full((T, 256 + 4), float("nan"), device=device) @@ -9329,8 +9198,8 @@ def test_gpu_router_dgrad_simt_program_in_its_window(router_rule, tile, hidden, W_pad = torch.full((256, hidden + 8), float("nan"), device=device) W_pad[:, :hidden] = W padded = torch.full((T, hidden + 4), float("nan"), device=device) - cake_backend.router_dgrad( - G_pad[:, :256], W_pad[:, :hidden], padded[:, :hidden], deterministic=False + cake_backend.cake_dense_gemm_nondet_train_router_dgrad( + G_pad[:, :256], W_pad[:, :hidden], padded[:, :hidden] ) torch.cuda.synchronize() assert torch.equal(padded[:, :hidden], first) @@ -9339,7 +9208,7 @@ def test_gpu_router_dgrad_simt_program_in_its_window(router_rule, tile, hidden, @pytest.mark.parametrize("tile", DGRAD_SIMT_TILES, ids=_simt_spec) def test_gpu_router_dgrad_simt_program_replays_in_a_cuda_graph(router_rule, tile): - """``router_dgrad`` in a SIMT window captured into a CUDA graph (its entry cache, the prepared launch) replays the + """``cake_dense_gemm_nondet_train_router_dgrad`` in a SIMT window captured into a CUDA graph (its entry cache, the prepared launch) replays the eager bits into a NaN-poisoned ``out``; the output gradient changed between replays is read again.""" device, arch, template = _dgrad_simt_device_or_skip(tile) hidden, T = 5120, 128 @@ -9347,26 +9216,26 @@ def test_gpu_router_dgrad_simt_program_replays_in_a_cuda_graph(router_rule, tile A, B, out = _router_case("dgrad", X, W, G) router_rule("dgrad", A, B, _simt_spec(tile)) eager = torch.empty_like(out) - cake_backend.router_dgrad(G, W, eager, deterministic=False) - prepared = cake_backend.prepare_router_fp32_gemm(G, W, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, eager) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(G, W, out) assert prepared.plan.simt == tile and prepared.plan.template == template entry_out = torch.empty_like(out) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(2): - cake_backend.router_dgrad(G, W, entry_out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, entry_out) prepared.launch() torch.cuda.current_stream().wait_stream(stream) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - cake_backend.router_dgrad(G, W, entry_out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, entry_out) prepared.launch() for scale in (1.0, -0.5): G.mul_(scale) want = torch.empty_like(out) - cake_backend.router_dgrad(G, W, want, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, want) entry_out.fill_(float("nan")) out.fill_(float("nan")) graph.replay() @@ -9400,26 +9269,30 @@ def test_gpu_router_dgrad_simt_program_waits_for_an_early_producer(router_rule): # The forward of this row's own route: in a forward SIMT window (64 tokens on # 152-SM parts) the cluster program, whose bits differ from the deterministic one. want_y = torch.empty_like(Y) - cake_backend.router_forward(X, W, want_y, deterministic=False, enable_pdl=False) + cake_backend.cake_dense_gemm_nondet_train_router_forward( + X, W, want_y, enable_pdl=False + ) torch.cuda.synchronize() assert _rel_fro(want_y, X.double() @ W.double().T) <= ROUTER_REL_FRO_MAX want_dx = torch.empty_like(dX) - cake_backend.router_dgrad(want_y, W, want_dx, deterministic=False, enable_pdl=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad( + want_y, W, want_dx, enable_pdl=False + ) torch.cuda.synchronize() assert _rel_fro(want_dx, want_y.double() @ W.double()) <= ROUTER_REL_FRO_MAX want_dw = torch.empty_like(dW) - cake_backend.router_wgrad( - want_y, want_dx, want_dw, deterministic=False, enable_pdl=False + cake_backend.cake_dense_gemm_nondet_train_router_wgrad( + want_y, want_dx, want_dw, enable_pdl=False ) torch.cuda.synchronize() for pdl in (True, False, None): for _ in range(3): for t in (Y, dX, dW): t.fill_(float("nan")) - kw = dict(deterministic=False, enable_pdl=pdl) - cake_backend.router_forward(X, W, Y, **kw) - cake_backend.router_dgrad(Y, W, dX, **kw) - cake_backend.router_wgrad(Y, dX, dW, **kw) + kw = dict(enable_pdl=pdl) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, Y, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(Y, W, dX, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(Y, dX, dW, **kw) torch.cuda.synchronize() assert torch.equal(Y, want_y) and torch.equal(dX, want_dx), pdl assert torch.equal(dW, want_dw), pdl @@ -9432,15 +9305,15 @@ def test_gpu_router_dgrad_simt_program_waits_for_an_early_producer(router_rule): def test_gpu_router_input_gradient_windows_run_at_their_edges(): - """Unforced (``deterministic=False``, no split count): at 1 and 17 tokens and at the first and the last token + """Unforced (no split count): at 1 and 17 tokens and at the first and the last token count of every input-gradient SIMT window of this device's router table (``router_rules`` of its architecture and - SM count), and one token past the last window, ``router_dgrad`` runs the program the planner gives -- inside a + SM count), and one token past the last window, ``cake_dense_gemm_nondet_train_router_dgrad`` runs the program the planner gives -- inside a window the SIMT program of its tile (the prepared launch's program, within the gate, the same bits on every call), outside the deterministic program's bits.""" device, arch = _atomic_device_or_skip() sm_count = det_backend.device_facts(device).sm_count rules, holds = cake_backend.router_rules(arch, sm_count) - atomic = cake_backend.PreparedAtomicRouterGemm + atomic = cake_backend.CakeDenseGemmNondetRouterLaunch ran = 0 for (rule_arch, layout, hidden, experts), rule in sorted( rules.items(), key=lambda kv: kv[0] @@ -9461,14 +9334,17 @@ def test_gpu_router_input_gradient_windows_run_at_their_edges(): hidden, T, ) - prepared = cake_backend.prepare_router_fp32_gemm( - G, W, torch.empty(T, hidden, device=device), deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare( + G, W, torch.empty(T, hidden, device=device) ) - det_out = det_backend.router_det_dgrad(G, W) + det_out = det_backend.cake_dense_gemm_det_train_router_dgrad(G, W) first = None for _ in range(2): out = torch.full((T, hidden), float("nan"), device=device) - assert cake_backend.router_dgrad(G, W, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out) + is out + ) torch.cuda.synchronize() if tile is None: assert torch.equal(out, det_out), (hidden, T) @@ -9483,8 +9359,10 @@ def test_gpu_router_input_gradient_windows_run_at_their_edges(): assert isinstance(prepared, atomic), (hidden, T) assert tuple(prepared.plan.simt) == tile, (hidden, T) _assert_ran( - lambda G=G, W=W, out=out: cake_backend.router_dgrad( - G, W, out, deterministic=False + lambda G=G, + W=W, + out=out: cake_backend.cake_dense_gemm_nondet_train_router_dgrad( + G, W, out ), prepared.module_name, ) @@ -9519,7 +9397,7 @@ def _device_simt_windows(arch, sm_count, layout) -> list: def test_gpu_router_forward_simt_windows_keep_the_fp64_rule(monkeypatch): - """Unforced (``router_forward(..., deterministic=False)``, no split count), at the first and the last token count of + """Unforced (``cake_dense_gemm_nondet_train_router_forward(...)``, no split count), at the first and the last token count of each forward SIMT window of the device's architecture (:data:`ROUTER_FWD_WINDOWS`; on a part whose SM count is not :data:`ROUTER_SM_COUNT` the measured count is set to the device's, so the program runs there too) and two seeds: the planner gives the window's cluster program, which the call launches (no program of the deterministic @@ -9563,11 +9441,11 @@ def test_gpu_router_forward_simt_windows_keep_the_fp64_rule(monkeypatch): ) assert plan is not None, (hidden, T, reason) assert plan.simt == tile and plan.template == template, (hidden, T) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() got = None for i in range(3): out.fill_(float("nan")) - cake_backend.router_forward(X, W, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, out) torch.cuda.synchronize() assert bool(torch.isfinite(out).all()), (hidden, T, i) if got is None: @@ -9576,22 +9454,24 @@ def test_gpu_router_forward_simt_windows_keep_the_fp64_rule(monkeypatch): assert torch.equal(out, got), (hidden, T, i) if seed == SEED + 250: _assert_ran( - lambda: cake_backend.router_forward( - X, W, out, deterministic=False + lambda: cake_backend.cake_dense_gemm_nondet_train_router_forward( + X, W, out ), module, ) graph = torch.cuda.CUDAGraph() out.fill_(float("nan")) with torch.cuda.graph(graph): - cake_backend.router_forward(X, W, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_forward( + X, W, out + ) out.fill_(float("nan")) graph.replay() torch.cuda.synchronize() assert torch.equal(out, got), (hidden, T) del graph ref = X.double() @ W.double().t() - base = det_backend.router_det_forward(X, W) + base = det_backend.cake_dense_gemm_det_train_router_forward(X, W) err, base_err = _rel_fro(got, ref), _rel_fro(base, ref) assert err <= ROUTER_REL_FRO_MAX, (hidden, T, seed, err) gamma = hidden * 2.0**-24 / (1 - hidden * 2.0**-24) @@ -9614,7 +9494,7 @@ def test_gpu_router_forward_simt_windows_keep_the_fp64_rule(monkeypatch): @pytest.mark.parametrize("op", ["wgrad", "dgrad"]) def test_gpu_router_shipped_simt_windows_keep_the_fp64_rule(monkeypatch, op): - """Unforced (``deterministic=False``, no split count), at the first and the last token count of each shipped + """Unforced (no split count), at the first and the last token count of each shipped SIMT window of the 148-SM tables that holds on this device (148-SM sm_100a / sm_103a parts; :func:`_device_simt_windows`) and two seeds: the planner gives the window's program; three calls into a NaN-poisoned ``out`` give finite values and the same bits; the relative Frobenius error against FP64 of the exact @@ -9643,7 +9523,11 @@ def test_gpu_router_shipped_simt_windows_keep_the_fp64_rule(monkeypatch, op): assert simt is None or len(simt) == 3, (hidden, T, simt) print(f"ROUTER_SIMT_FP64 {arch} sm_count={sm_count} {op}: no shipped window") return - entry = cake_backend.router_wgrad if op == "wgrad" else cake_backend.router_dgrad + entry = ( + cake_backend.cake_dense_gemm_nondet_train_router_wgrad + if op == "wgrad" + else cake_backend.cake_dense_gemm_nondet_train_router_dgrad + ) rules, holds = cake_backend.router_rules(arch, sm_count) assert holds ran = 0 @@ -9667,11 +9551,11 @@ def test_gpu_router_shipped_simt_windows_keep_the_fp64_rule(monkeypatch, op): ) assert plan is not None, (hidden, T, reason) assert plan.simt == tile and plan.template == template, (hidden, T) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() got = None for i in range(3): out.fill_(float("nan")) - entry(*args, out, deterministic=False) + entry(*args, out) torch.cuda.synchronize() assert bool(torch.isfinite(out).all()), (hidden, T, i) if got is None: @@ -9680,11 +9564,11 @@ def test_gpu_router_shipped_simt_windows_keep_the_fp64_rule(monkeypatch, op): assert torch.equal(out, got), (hidden, T, i) with monkeypatch.context() as m: m.setitem(rules, key, without) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() base = torch.full_like(out, float("nan")) - entry(*args, base, deterministic=False) + entry(*args, base) torch.cuda.synchronize() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() ref = A.double() @ B.double() err, base_err = _rel_fro(got, ref), _rel_fro(base, ref) print( @@ -9707,7 +9591,7 @@ def test_gpu_router_shipped_simt_windows_keep_the_fp64_rule(monkeypatch, op): def test_gpu_router_entry_points_reach_the_stream_k_program(router_rule): - """``router_forward`` / ``router_wgrad`` with ``deterministic=False`` and a verdict for each row (no forward + """``cake_dense_gemm_nondet_train_router_forward`` / ``cake_dense_gemm_nondet_train_router_wgrad`` with a verdict for each row (no forward verdict ships; the weight gradient of this hidden size runs the deterministic plan): the forward repeats bit for bit, and its top-8, like the deterministic forward's, is FP64's on every row but the near ties (the 8th and 9th logits within 16 FP32 ulps); both meet the gate.""" @@ -9715,8 +9599,8 @@ def test_gpu_router_entry_points_reach_the_stream_k_program(router_rule): X, W, G = _router_inputs(ROUTER_T, device, SEED + 21) router_rule("fwd", X, W.t(), "forward") router_rule("wgrad", G.t(), X, "pairs") - first = cake_backend.router_forward(X, W, deterministic=False) - det_fwd = det_backend.router_det_forward(X, W) + first = cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) + det_fwd = det_backend.cake_dense_gemm_det_train_router_forward(X, W) ref_fwd, ref_wgrad = X.double() @ W.double().t(), G.double().t() @ X.double() for name, out in (("atomic", first), ("det", det_fwd)): differ, ties = _top8_rows_differ(out, ref_fwd) @@ -9725,20 +9609,20 @@ def test_gpu_router_entry_points_reach_the_stream_k_program(router_rule): ) assert differ == 0 and ties <= ROUTER_T // 100, (differ, ties) for _ in range(RERUNS): - out = cake_backend.router_forward(X, W, deterministic=False) + out = cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) assert torch.equal(out, first) and _rel_fro(out, ref_fwd) <= ROUTER_REL_FRO_MAX - wgrad = cake_backend.router_wgrad(G, X, deterministic=False) + wgrad = cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X) assert _rel_fro(wgrad, ref_wgrad) <= ROUTER_REL_FRO_MAX - plan = cake_backend.plan_router_fp32_gemm(X, W.t(), first, deterministic=False) - assert isinstance(plan, cake_backend.AtomicRouterPlan) + plan = cake_backend.cake_dense_gemm_nondet_train_router_plan(X, W.t(), first) + assert isinstance(plan, cake_backend.CakeDenseGemmNondetRouterPlan) _assert_ran( - lambda: cake_backend.router_forward(X, W, first, deterministic=False), + lambda: cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, first), cake_jit.select_module(arch, plan.template), ) def test_gpu_planned_rows_run_the_programs_of_their_verdicts(): - """Unforced (``deterministic=False``, no split count): every K1 weight-gradient row with a verdict for this + """Unforced (no split count): every K1 weight-gradient row with a verdict for this architecture and every measured router row prepares the program its verdict names on this device (by kernel name, at ``T = 16231``; a row with a range of :data:`K1_DSM_RANGES_BY_SM` there the range's cluster program), and a row whose verdict is ``None`` -- or that this device's SM count leaves to the deterministic program -- prepares @@ -9759,14 +9643,22 @@ def test_gpu_planned_rows_run_the_programs_of_their_verdicts(): G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) out_dtype = torch.float32 if dt == "f32" else torch.bfloat16 out = torch.empty(N, M, device=device, dtype=out_dtype) - prepared = cake_backend.prepare_projection_wgrad(G, X, out, deterministic=False) + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) kw = dict(arch=arch, L=L, M=M, N=N, K=T, out_f32=dt == "f32", sm_count=sm_count) dsm = cake_backend.k1_dsm_plan(**kw) want = dsm or cake_backend.k1_plan(**kw) if want is None: - assert not isinstance(prepared, cake_backend.PreparedAtomicGemm), (M, N, dt) + assert not isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ), (M, N, dt) continue - assert isinstance(prepared, cake_backend.PreparedAtomicGemm), (M, N, dt) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch), ( + M, + N, + dt, + ) assert (prepared.plan.block_n, prepared.plan.splits) == want assert prepared.plan.dsm == (dsm is not None), (M, N, dt) _assert_ran(prepared.launch, prepared.module_name) @@ -9782,15 +9674,17 @@ def test_gpu_planned_rows_run_the_programs_of_their_verdicts(): X, W, G = _router_inputs(T, device, SEED + 91, K=hidden) op = {"kk": "fwd", "kn": "dgrad", "nn_t": "wgrad"}[layout] A, B, out = _router_case(op, X, W, G) - prepared = cake_backend.prepare_router_fp32_gemm(A, B, out, deterministic=False) + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare(A, B, out) plan, _ = cake_backend.plan_atomic_router( A, B, out, arch=arch, sm_count=sm_count ) if plan is None: - assert not isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + assert not isinstance( + prepared, cake_backend.CakeDenseGemmNondetRouterLaunch + ) continue assert spec is not None - assert isinstance(prepared, cake_backend.PreparedAtomicRouterGemm) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch) assert (prepared.plan.template, prepared.plan.units) == ( plan.template, plan.units, @@ -9811,7 +9705,7 @@ def test_gpu_planned_rows_run_the_programs_of_their_verdicts(): def test_gpu_k1_sm_count_overrides_run_their_programs(): - """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of + """Unforced (no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_OVERRIDES_BY_SM` (``sm_107a``: 216 SMs): every overridden row prepares and runs the override's program at the first and the last token count of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` and one token below the last, outside the part's ranges of :data:`K1_CHAIN_RANGES_BY_SM`, which come first, and its token counts of @@ -9844,10 +9738,14 @@ def test_gpu_k1_sm_count_overrides_run_their_programs(): G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) out_dtype = torch.float32 if f32 else torch.bfloat16 out = torch.full((N, M), float("nan"), device=device, dtype=out_dtype) - prepared = cake_backend.prepare_projection_wgrad( - G, X, out, deterministic=False + prepared = ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) ) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm), (key, T) + assert isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ), (key, T) assert (prepared.plan.block_n, prepared.plan.splits) == want assert prepared.template == cake_backend.atomic_k1_template(block_n, f32) _assert_ran(prepared.launch, prepared.module_name) @@ -9869,9 +9767,10 @@ def test_gpu_k1_sm_count_overrides_run_their_programs(): def test_gpu_k1_forward_rows_run_their_plan(): """On a part of exactly the SM count of a table of :data:`K1_FORWARD_RULES_BY_SM` (152-SM ``sm_100a`` / - ``sm_103a``): the forward of every row with ``deterministic=False`` binds the program the deterministic planner + ``sm_103a``): the forward of every row binds the program the deterministic planner gives the row with the table's knobs (the single-CTA program), meets the deterministic package's BF16 budget - against FP64, repeats bit for bit through the bound launch, and ``deterministic=True`` meets the budget too. A + against FP64, repeats bit for bit through the bound launch, and the deterministic package's forward meets the + budget too. A part of any other SM count skips.""" device = _det_device_or_skip() facts = det_backend.device_facts(device) @@ -9883,9 +9782,12 @@ def test_gpu_k1_forward_rows_run_their_plan(): for (K, N, T), knobs in sorted(rows.items()): X, W = _bf16((T, K), device, g), _bf16((N, K), device, g) ref = X.double() @ W.double().t() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() out = torch.full((T, N), float("nan"), device=device, dtype=torch.bfloat16) - assert cake_backend.projection_forward(X, W, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, out) + is out + ) torch.cuda.synchronize() ((_, entry),) = cake_backend._K1_DET_CACHE.items() want, *_ = det_backend._k1_plan( @@ -9895,11 +9797,11 @@ def test_gpu_k1_forward_rows_run_their_plan(): assert entry.plan.template.endswith("_kk_n128_c1"), (K, N, T) _k1_check(out, ref, f"forward {K} {N} {T}") again = torch.full_like(out, float("nan")) - cake_backend.projection_forward(X, W, again, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, again) torch.cuda.synchronize() assert torch.equal(out, again), (K, N, T) det_out = torch.full_like(out, float("nan")) - det_backend.projection_det_forward(X, W, det_out) + det_backend.cake_dense_gemm_det_train_projection_forward(X, W, det_out) torch.cuda.synchronize() _k1_check(det_out, ref, f"deterministic forward {K} {N} {T}") ran += 1 @@ -9910,7 +9812,7 @@ def test_gpu_k1_forward_rows_run_their_plan(): def test_gpu_k1_sm_count_overrides_through_the_convenience_binding(): - """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of + """Unforced (no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_OVERRIDES_BY_SM` (``sm_107a``: 216 SMs): the convenience path of every overridden row binds the override's program (its :data:`_K1_ATOMIC_CACHE` entry) at the first and the last token count of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` and one token below the last, outside the part's ranges of @@ -9939,13 +9841,18 @@ def test_gpu_k1_sm_count_overrides_through_the_convenience_binding(): G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) out_dtype = torch.float32 if f32 else torch.bfloat16 want = torch.full((N, M), float("nan"), device=device, dtype=out_dtype) - prepared = cake_backend.prepare_projection_wgrad( - G, X, want, deterministic=False + prepared = ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, want + ) ) prepared.launch() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() out = torch.full_like(want, float("nan")) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) torch.cuda.synchronize() (entry,) = [ e for e in cake_backend._K1_ATOMIC_CACHE.values() if e.plan is not None @@ -10066,7 +9973,7 @@ def _k1_with_chain_ranges(arch, sm_count, key, T, want): def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(no_k1_dsm): - """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of + """Unforced (no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` or :data:`K1_TWO_SPLIT_RANGES_BY_SM` (``sm_107a``: 212 SMs; ``sm_100a`` / ``sm_103a``: 152 SMs): at the first and the last token count of every range of a row, and one token past the first token count of a two-split range (an odd @@ -10097,10 +10004,14 @@ def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(no_k1_dsm): G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) out_dtype = torch.float32 if f32 else torch.bfloat16 out = torch.full((N, M), float("nan"), device=device, dtype=out_dtype) - prepared = cake_backend.prepare_projection_wgrad(G, X, out, deterministic=False) + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) if want_plan is not None: assert plan == want_plan, (key, T) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm), (key, T) + assert isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ), (key, T) assert (prepared.plan.block_n, prepared.plan.splits) == plan, (key, T) assert prepared.plan.chains == chains, (key, T) assert prepared.template == cake_backend.atomic_k1_template( @@ -10115,20 +10026,26 @@ def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(no_k1_dsm): if at != T: G, X = _bf16((at, N), device, g), _bf16((at, M), device, g) out = torch.full((N, M), float("nan"), device=device, dtype=out_dtype) - prepared = cake_backend.prepare_projection_wgrad( - G, X, out, deterministic=False + prepared = ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) ) - assert isinstance(prepared, cake_backend.PreparedAtomicGemm), (key, at) + assert isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ), (key, at) assert (prepared.plan.block_n, prepared.plan.splits) == plan, (key, at) _assert_ran(prepared.launch, prepared.module_name) torch.cuda.synchronize() _k1_check(out, G.double().t() @ X.double(), f"{M} {N} {dt} {at}") else: assert plan is None, (key, T) - assert not isinstance(prepared, cake_backend.PreparedAtomicGemm), (key, T) + assert not isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ), (key, T) prepared.launch() want = torch.full_like(out, float("nan")) - det_backend.projection_det_wgrad(G, X, want) + det_backend.cake_dense_gemm_det_train_projection_wgrad(G, X, want) torch.cuda.synchronize() assert torch.equal(out, want), (key, T) print( @@ -10138,7 +10055,7 @@ def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(no_k1_dsm): def test_gpu_k1_sm_count_token_ranges_through_the_convenience_binding(no_k1_dsm): - """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of + """Unforced (no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` or :data:`K1_TWO_SPLIT_RANGES_BY_SM` (``sm_107a``: 212 SMs; ``sm_100a`` / ``sm_103a``: 152 SMs): the convenience path binds the part's program @@ -10165,13 +10082,15 @@ def test_gpu_k1_sm_count_token_ranges_through_the_convenience_binding(no_k1_dsm) G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) out_dtype = torch.float32 if f32 else torch.bfloat16 want = torch.full((N, M), float("nan"), device=device, dtype=out_dtype) - prepared = cake_backend.prepare_projection_wgrad( - G, X, want, deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, want ) prepared.launch() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() out = torch.full_like(want, float("nan")) - assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) is out + ) torch.cuda.synchronize() (entry,) = list(cake_backend._K1_ATOMIC_CACHE.values()) assert entry.k1 is not None, (key, T) @@ -10192,7 +10111,7 @@ def test_gpu_k1_sm_count_token_ranges_through_the_convenience_binding(no_k1_dsm) def test_gpu_k1_152_sm_one_split_ranges_run_at_their_edges(no_k1_dsm): - """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of + """Unforced (no split count), on a part of exactly the SM count of a table of :data:`K1_ONE_SPLIT_RANGES_BY_SM` (152-SM ``sm_100a`` / ``sm_103a`` parts): at the first and the last token count of every range of a row and at 1025 tokens (an odd number of 64-token K blocks) the prepared launch runs the range's one-split program (one chain; by kernel name) and its output meets the deterministic package's budget against FP64 @@ -10222,12 +10141,16 @@ def test_gpu_k1_152_sm_one_split_ranges_run_at_their_edges(no_k1_dsm): plan = cake_backend.k1_plan( arch=arch, L=L, M=M, N=N, K=T, out_f32=f32, sm_count=sm_count ) - prepared = cake_backend.prepare_projection_wgrad( - G, X, out, deterministic=False + prepared = ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) ) if not lo <= T <= hi: assert plan != (block_n, 1), what - atomic = isinstance(prepared, cake_backend.PreparedAtomicGemm) + atomic = isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ) assert atomic == (plan is not None), what if atomic: assert (prepared.plan.block_n, prepared.plan.splits) == plan @@ -10236,7 +10159,9 @@ def test_gpu_k1_152_sm_one_split_ranges_run_at_their_edges(no_k1_dsm): _k1_check(out, ref, what) continue assert plan == (block_n, 1), what - assert isinstance(prepared, cake_backend.PreparedAtomicGemm), what + assert isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch + ), what assert (prepared.plan.block_n, prepared.plan.splits) == plan, what assert prepared.plan.chains == 1, what assert prepared.template == cake_backend.atomic_k1_template( @@ -10250,10 +10175,12 @@ def test_gpu_k1_152_sm_one_split_ranges_run_at_their_edges(no_k1_dsm): prepared.launch() torch.cuda.synchronize() _k1_check(out, ref, f"{what} run {i}") - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() conv = torch.full_like(out, float("nan")) assert ( - cake_backend.projection_wgrad(G, X, conv, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad( + G, X, conv + ) is conv ) torch.cuda.synchronize() @@ -10269,7 +10196,7 @@ def test_gpu_k1_152_sm_one_split_ranges_run_at_their_edges(no_k1_dsm): def test_gpu_k1_152_sm_dsm_ranges_run_at_their_edges(): - """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of + """Unforced (no split count), on a part of exactly the SM count of a table of :data:`K1_DSM_RANGES_BY_SM` (152-SM ``sm_100a`` parts): at the first and the last token count of every range of a row the prepared launch runs the range's cluster program (by kernel name) and its output meets the deterministic package's FP32 budget against FP64 (:func:`_k1_check`: relative Frobenius error at most 5e-5, max error at most @@ -10283,7 +10210,7 @@ def test_gpu_k1_152_sm_dsm_ranges_run_at_their_edges(): rows = cake_backend.K1_DSM_RANGES_BY_SM.get((arch, sm_count), {}) if not rows: pytest.skip(f"no K1 cluster range of {arch} on {sm_count} SMs") - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() g = torch.Generator(device=device).manual_seed(SEED + 98) ran = 0 for key, ranges in sorted(rows.items()): @@ -10303,10 +10230,14 @@ def test_gpu_k1_152_sm_dsm_ranges_run_at_their_edges(): dsm = cake_backend.k1_dsm_plan( arch=arch, L=L, M=M, N=N, K=T, out_f32=True, sm_count=sm_count ) - prepared = cake_backend.prepare_projection_wgrad( - G, X, out, deterministic=False + prepared = ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) + ) + atomic = isinstance( + prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch ) - atomic = isinstance(prepared, cake_backend.PreparedAtomicGemm) if not held: assert dsm is None, what assert not (atomic and prepared.plan.dsm), what @@ -10326,10 +10257,11 @@ def test_gpu_k1_152_sm_dsm_ranges_run_at_their_edges(): prepared.launch() torch.cuda.synchronize() _k1_check(out, ref, f"{what} run {i}") - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() conv = torch.full_like(out, float("nan")) assert ( - cake_backend.projection_wgrad(G, X, conv, deterministic=False) is conv + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, conv) + is conv ) torch.cuda.synchronize() (entry,) = list(cake_backend._K1_ATOMIC_CACHE.values()) @@ -10345,7 +10277,7 @@ def test_gpu_k1_152_sm_dsm_ranges_run_at_their_edges(): def test_gpu_router_token_windows_run_at_their_edges(): - """Unforced (``deterministic=False``, no split count): at the first and the last token count of every token window + """Unforced (no split count): at the first and the last token count of every token window of a router weight-gradient row of this architecture (the smallest measured count, 64, for a window other than SIMT that starts at one token; a SIMT window runs from its first token count) the prepared launch runs the program the planner gives -- the SIMT program of a SIMT window and @@ -10355,7 +10287,7 @@ def test_gpu_router_token_windows_run_at_their_edges(): device, arch = _atomic_device_or_skip() sm_count = det_backend.device_facts(device).sm_count rules, measured = cake_backend.router_rules(arch, sm_count) - atomic = cake_backend.PreparedAtomicRouterGemm + atomic = cake_backend.CakeDenseGemmNondetRouterLaunch ran = 0 for (rule_arch, layout, hidden, _experts), rule in sorted( rules.items(), key=lambda kv: kv[0] @@ -10383,8 +10315,8 @@ def test_gpu_router_token_windows_run_at_their_edges(): simt = cake_backend.router_simt( "nn_t", arch=arch, sm_count=sm_count, M=hidden, N=256, K=T ) - prepared = cake_backend.prepare_router_fp32_gemm( - A, B, out, deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare( + A, B, out ) tile = cake_backend.router_simt_tile(spec) if not inside: @@ -10415,21 +10347,25 @@ def test_gpu_router_token_windows_run_at_their_edges(): def test_gpu_router_rows_without_a_verdict_run_the_deterministic_program_unplanned(): - """``deterministic=False`` without a split count on a router row without a verdict (the forward; the weight + """Without a split count on a router row without a verdict (the forward; the weight gradient at a token count outside every window): the deterministic program's bits (its production split counts), through this package's binding of it (``_K2_SPLITK_CACHE``); this package plans nothing for the call.""" device, arch = _atomic_device_or_skip() for hidden in (4096, 5120, 6144): X, W, G = _router_inputs(8000, device, SEED + 23, K=hidden) - cake_backend.clear_convenience_caches() - fwd = cake_backend.router_forward(X, W, deterministic=False) - wgrad = cake_backend.router_wgrad(G, X, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() + fwd = cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) + wgrad = cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X) assert not cake_backend._K2_ATOMIC_CACHE assert len(cake_backend._K2_SPLITK_CACHE) == 2 - assert torch.equal(fwd, det_backend.router_det_forward(X, W)) - assert torch.equal(wgrad, det_backend.router_det_wgrad(G, X)) + assert torch.equal( + fwd, det_backend.cake_dense_gemm_det_train_router_forward(X, W) + ) + assert torch.equal( + wgrad, det_backend.cake_dense_gemm_det_train_router_wgrad(G, X) + ) A, B, out = _router_case("fwd", X, W, G) - assert cake_backend.router_fp32_gemm(A, B, out, deterministic=False) is out + assert cake_backend.cake_dense_gemm_nondet_train_router(A, B, out) is out assert torch.equal(out, fwd) and not cake_backend._K2_ATOMIC_CACHE @@ -10455,18 +10391,18 @@ def test_gpu_router_verdicts_set_at_run_time_are_seen(monkeypatch): layout, arch=arch, sm_count=sm_count, M=M, N=256, K=K ) assert spec is None, (layout, spec) - cake_backend.clear_convenience_caches() - want_fwd = det_backend.router_fp32_det_gemm( + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() + want_fwd = det_backend.cake_dense_gemm_det_train_router( X, W.t(), torch.empty(T, 256, device=device) ) - want_wgrad = det_backend.router_fp32_det_gemm( + want_wgrad = det_backend.cake_dense_gemm_det_train_router( G.t(), X, torch.empty(256, hidden, device=device) ) def unplanned(): for _ in range(2): - fwd = cake_backend.router_forward(X, W, deterministic=False) - wgrad = cake_backend.router_wgrad(G, X, deterministic=False) + fwd = cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) + wgrad = cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X) assert torch.equal(fwd, want_fwd) and torch.equal(wgrad, want_wgrad) assert not cake_backend._K2_ATOMIC_CACHE @@ -10476,18 +10412,18 @@ def unplanned(): ref_fwd, ref_wgrad = X.double() @ W.double().t(), G.double().t() @ X.double() def fwd(): - return cake_backend.router_forward(X, W, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W) def wgrad(): - return cake_backend.router_wgrad(G, X, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X) for op, A, B, ref, call in ( ("fwd", X, W.t(), ref_fwd, fwd), ("wgrad", G.t(), X, ref_wgrad, wgrad), ): out = torch.empty(A.shape[0], B.shape[1], device=device) - plan = cake_backend.plan_router_fp32_gemm(A, B, out, deterministic=False) - assert isinstance(plan, cake_backend.AtomicRouterPlan), op + plan = cake_backend.cake_dense_gemm_nondet_train_router_plan(A, B, out) + assert isinstance(plan, cake_backend.CakeDenseGemmNondetRouterPlan), op _assert_ran(call, cake_jit.select_module(arch, plan.template)) assert _rel_fro(call(), ref) <= ROUTER_REL_FRO_MAX, op first = fwd() @@ -10501,7 +10437,7 @@ def wgrad(): def test_gpu_router_verdict_changed_at_run_time_runs_the_new_program(monkeypatch): """A weight-gradient verdict changed from one program to another between calls of the same tensors: the next call - -- through ``router_wgrad`` and through ``router_fp32_gemm`` -- runs the program the planner gives under the new + -- through ``cake_dense_gemm_nondet_train_router_wgrad`` and through ``cake_dense_gemm_nondet_train_router`` -- runs the program the planner gives under the new verdict (the stream-K ``"pairs"`` partition, then a SIMT tile this checkout registers for the architecture, else two units per tile), within the router budget, without clearing a cache.""" device, arch = _atomic_device_or_skip() @@ -10524,15 +10460,15 @@ def test_gpu_router_verdict_changed_at_run_time_runs_the_new_program(monkeypatch None, ) calls = ( - lambda: cake_backend.router_wgrad(G, X, out, deterministic=False), - lambda: cake_backend.router_fp32_gemm(A, B, out, deterministic=False), + lambda: cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, out), + lambda: cake_backend.cake_dense_gemm_nondet_train_router(A, B, out), ) - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() plans = [] for spec in ("pairs", "S2" if simt is None else _simt_spec(simt)): monkeypatch.setitem(table, (arch, "nn_t", hidden, 256), ((1, T, spec),)) - plan = cake_backend.plan_router_fp32_gemm(A, B, out, deterministic=False) - assert isinstance(plan, cake_backend.AtomicRouterPlan), spec + plan = cake_backend.cake_dense_gemm_nondet_train_router_plan(A, B, out) + assert isinstance(plan, cake_backend.CakeDenseGemmNondetRouterPlan), spec for call in calls: _assert_ran(call, cake_jit.select_module(arch, plan.template)) assert _rel_fro(call(), ref) <= ROUTER_REL_FRO_MAX, spec @@ -10543,7 +10479,7 @@ def test_gpu_router_verdict_changed_at_run_time_runs_the_new_program(monkeypatch def test_gpu_router_deterministic_program_through_the_binding(): - """``deterministic=False`` calls that run the deterministic program -- the forward and the input gradient, the + """Calls of this package that run the deterministic program -- the forward and the input gradient, the weight gradient of a row without a verdict, the generic ``kn`` call -- go through this package's binding (``_K2_SPLITK_CACHE``; nothing in ``_K2_ATOMIC_CACHE``) and give the deterministic package's bits: dense and padded outputs (one split stores into a padded ``out`` through its flat alias, several splits reduce into it), a @@ -10553,22 +10489,22 @@ def test_gpu_router_deterministic_program_through_the_binding(): device, arch = _atomic_device_or_skip() # no router verdict on any architecture from 257 to 16171 tokens (the SIMT windows end at 256 tokens or before) hidden = 6144 - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() def check(T, seed, pad_out=False, pad_g=False): X, W, G = _router_inputs(T, device, seed, K=hidden) if pad_g: G = torch.cat([G, torch.zeros(T, 16, device=device)], 1)[:, :256] shapes = {"fwd": (T, 256), "dgrad": (T, hidden), "wgrad": (256, hidden)} - # op -> (the deterministic package's call, this package's deterministic=False call) + # op -> (the deterministic package's call, this package's call) # fmt: off calls = { - "fwd": (lambda out: det_backend.router_det_forward(X, W, out), - lambda out: cake_backend.router_forward(X, W, out, deterministic=False)), - "dgrad": (lambda out: det_backend.router_det_dgrad(G, W, out), - lambda out: cake_backend.router_dgrad(G, W, out, deterministic=False)), - "wgrad": (lambda out: det_backend.router_det_wgrad(G, X, out), - lambda out: cake_backend.router_wgrad(G, X, out, deterministic=False)), + "fwd": (lambda out: det_backend.cake_dense_gemm_det_train_router_forward(X, W, out), + lambda out: cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, out)), + "dgrad": (lambda out: det_backend.cake_dense_gemm_det_train_router_dgrad(G, W, out), + lambda out: cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out)), + "wgrad": (lambda out: det_backend.cake_dense_gemm_det_train_router_wgrad(G, X, out), + lambda out: cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, out)), } # fmt: on for op, (det_call, call) in calls.items(): @@ -10586,8 +10522,8 @@ def check(T, seed, pad_out=False, pad_g=False): if pad_out: assert bool((storage[:, cols:] == 7.0).all()), op out = torch.empty(T, hidden, device=device) - want = det_backend.router_fp32_det_gemm(G, W, out.clone()) - assert cake_backend.router_fp32_gemm(G, W, out, deterministic=False) is out + want = det_backend.cake_dense_gemm_det_train_router(G, W, out.clone()) + assert cake_backend.cake_dense_gemm_nondet_train_router(G, W, out) is out assert torch.equal(out, want) variants = ((320, False, False), (1000, True, False), (1000, False, True)) @@ -10618,9 +10554,9 @@ def check(T, seed, pad_out=False, pad_g=False): storage = torch.empty(1001 * hidden + 4, device=device) out = storage[1 : 1 + 1001 * hidden].view(1001, hidden) with pytest.raises(ValueError, match="16-byte aligned") as refused: - det_backend.router_det_dgrad(G, W, out) + det_backend.cake_dense_gemm_det_train_router_dgrad(G, W, out) with pytest.raises(ValueError, match=re.escape(str(refused.value))): - cake_backend.router_dgrad(G, W, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(G, W, out) def _special_partials(splits, rows, cols, seed): @@ -10746,16 +10682,20 @@ def test_gpu_router_partials_sum_eager_graph_and_recompute(): _require_template(arch, cake_backend.PARTIALS_SUM_TEMPLATE) hidden, T = 6144, 4096 # no router verdict at 4096 tokens on any architecture X, W, G = _router_inputs(T, device, SEED + 91, K=hidden) - cake_backend.clear_convenience_caches() - want_fwd = det_backend.router_det_forward(X, W, enable_pdl=False) - want_wgrad = det_backend.router_det_wgrad(G, X, enable_pdl=False) + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() + want_fwd = det_backend.cake_dense_gemm_det_train_router_forward( + X, W, enable_pdl=False + ) + want_wgrad = det_backend.cake_dense_gemm_det_train_router_wgrad( + G, X, enable_pdl=False + ) Y = torch.empty(T, 256, device=device) dW = torch.empty(256, hidden, device=device) stream = torch.cuda.Stream(device) def step(): - cake_backend.router_forward(X, W, Y, deterministic=False) - cake_backend.router_wgrad(G, X, dW, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, Y) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, dW) stream.wait_stream(torch.cuda.current_stream(device)) with torch.cuda.stream(stream): @@ -10796,10 +10736,16 @@ def test_gpu_router_binding_pdl_waits_for_an_early_producer(): device, arch = _atomic_device_or_skip() hidden, T = 6144, 4096 # no router verdict at 4096 tokens on any architecture X, W, _ = _router_inputs(T, device, SEED + 80, K=hidden) - cake_backend.clear_convenience_caches() - want_y = det_backend.router_det_forward(X, W, enable_pdl=False) - want_dx = det_backend.router_det_dgrad(want_y, W, enable_pdl=False) - want_dw = det_backend.router_det_wgrad(want_y, want_dx, enable_pdl=False) + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() + want_y = det_backend.cake_dense_gemm_det_train_router_forward( + X, W, enable_pdl=False + ) + want_dx = det_backend.cake_dense_gemm_det_train_router_dgrad( + want_y, W, enable_pdl=False + ) + want_dw = det_backend.cake_dense_gemm_det_train_router_wgrad( + want_y, want_dx, enable_pdl=False + ) Y = torch.empty(T, 256, device=device) dX = torch.empty(T, hidden, device=device) dW = torch.empty(256, hidden, device=device) @@ -10807,10 +10753,10 @@ def test_gpu_router_binding_pdl_waits_for_an_early_producer(): for _ in range(3): for t in (Y, dX, dW): t.fill_(float("nan")) - kw = dict(deterministic=False, enable_pdl=pdl) - cake_backend.router_forward(X, W, Y, **kw) - cake_backend.router_dgrad(Y, W, dX, **kw) - cake_backend.router_wgrad(Y, dX, dW, **kw) + kw = dict(enable_pdl=pdl) + cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, Y, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(Y, W, dX, **kw) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(Y, dX, dW, **kw) torch.cuda.synchronize() assert torch.equal(Y, want_y) and torch.equal(dX, want_dx), pdl assert torch.equal(dW, want_dw), pdl @@ -10834,19 +10780,19 @@ def test_gpu_cached_router_launch_is_the_prepared_launch(): if grown: X, W, G = _router_inputs(4096, device, SEED + 25, K=5120) A, B, out = _router_case("wgrad", X, W, G) - cake_backend.router_fp32_gemm(A, B, out, splits=4, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(A, B, out, splits=4) for (A, B, out), splits in cases: - plan = cake_backend.plan_router_fp32_gemm( - A, B, out, splits=splits, deterministic=False + plan = cake_backend.cake_dense_gemm_nondet_train_router_plan( + A, B, out, splits=splits ) - if not isinstance(plan, cake_backend.AtomicRouterPlan): + if not isinstance(plan, cake_backend.CakeDenseGemmNondetRouterPlan): continue - prepared = cake_backend.prepare_router_fp32_gemm( - A, B, out.clone(), splits=splits, deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare( + A, B, out.clone(), splits=splits ) want = prepared.launch().clone() - got = cake_backend.router_fp32_gemm( - A, B, out, splits=splits, deterministic=False + got = cake_backend.cake_dense_gemm_nondet_train_router( + A, B, out, splits=splits ) torch.cuda.synchronize() ref = A.double() @ B.double() @@ -10859,7 +10805,7 @@ def test_gpu_cached_router_launch_is_the_prepared_launch(): def test_gpu_cached_router_launch_runs_the_head_tail_windows_at_their_edges(): - """Unforced (``deterministic=False``, no split count): at the first and the last token count of every head/tail + """Unforced (no split count): at the first and the last token count of every head/tail window of this device's router table (``router_rules`` of its architecture and SM count), the convenience path's cached launch (its bound arguments, ``_K2Binding``) runs the prepared launch's plan and program, meets the gate and, for a reproducible plan, gives the prepared launch's bits; the second call of the key reuses the binding.""" @@ -10879,16 +10825,19 @@ def test_gpu_cached_router_launch_runs_the_head_tail_windows_at_their_edges(): for hidden, T in edges: X, W, G = _router_inputs(T, device, SEED + 26, K=hidden) A, B, out = _router_case("wgrad", X, W, G) - prepared = cake_backend.prepare_router_fp32_gemm( - A, B, out.clone(), deterministic=False + prepared = cake_backend.cake_dense_gemm_nondet_train_router_prepare( + A, B, out.clone() + ) + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetRouterLaunch), ( + hidden, + T, ) - assert isinstance(prepared, cake_backend.PreparedAtomicRouterGemm), (hidden, T) assert prepared.plan.head_spans > 0, (hidden, T) want = prepared.launch().clone() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() _assert_ran( - lambda A=A, B=B, out=out: cake_backend.router_fp32_gemm( - A, B, out, deterministic=False + lambda A=A, B=B, out=out: cake_backend.cake_dense_gemm_nondet_train_router( + A, B, out ), prepared.module_name, ) @@ -10904,7 +10853,7 @@ def test_gpu_cached_router_launch_runs_the_head_tail_windows_at_their_edges(): if prepared.plan.reproducible: assert torch.equal(out, want), (hidden, T) out.fill_(float("nan")) - assert cake_backend.router_fp32_gemm(A, B, out, deterministic=False) is out + assert cake_backend.cake_dense_gemm_nondet_train_router(A, B, out) is out torch.cuda.synchronize() (again,) = cake_backend._K2_ATOMIC_CACHE.values() assert again is entry and again.k2 is bound, (hidden, T) @@ -10932,8 +10881,8 @@ def test_gpu_router_padded_out_rows(op, splits): ref = A.double() @ B.double() for launch in ( prepared.launch, - lambda: cake_backend.router_fp32_gemm( - A, B, out, splits=splits, deterministic=False + lambda: cake_backend.cake_dense_gemm_nondet_train_router( + A, B, out, splits=splits ), ): for _ in range(3): @@ -10953,16 +10902,16 @@ def test_gpu_router_misaligned_out_takes_the_deterministic_path(): storage = torch.empty(1001 * 256 + 4, device=device) out = storage[1 : 1 + 1001 * 256].view(1001, 256) with pytest.raises(ValueError, match="16-byte aligned") as refused: - det_backend.plan_router_fp32_det_gemm(X, W.t(), out, splits=1) + det_backend.cake_dense_gemm_det_train_router_plan(X, W.t(), out, splits=1) message = re.escape(str(refused.value)) with pytest.raises(ValueError, match=message): - cake_backend.plan_router_fp32_gemm(X, W.t(), out, splits=1, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_plan(X, W.t(), out, splits=1) with pytest.raises(ValueError, match=message): - cake_backend.prepare_router_fp32_gemm( - X, W.t(), out, splits=1, deterministic=False + cake_backend.cake_dense_gemm_nondet_train_router_prepare( + X, W.t(), out, splits=1 ) with pytest.raises(ValueError, match=message): - cake_backend.router_fp32_gemm(X, W.t(), out, splits=1, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router(X, W.t(), out, splits=1) def _router_entry_of(op, P, Q, out, splits, stream): @@ -10982,24 +10931,24 @@ def _router_entry_bits(entry) -> bool: @pytest.mark.parametrize("hidden, T", [(5120, 64), (4096, 256), (6144, 1000)]) def test_gpu_router_entry_points_give_their_paths_results(monkeypatch, hidden, T): - """``deterministic=False`` through ``_K2_ENTRY_CACHE``: every router entry point's call -- the forward with and + """Through ``_K2_ENTRY_CACHE``: every router entry point's call -- the forward with and without a split count, the input gradient, the weight gradient with and without one; dense and padded outputs; another stream -- binds the launch its own path takes and gives that path's results on repeated calls: the deterministic program and a reproducible plan bit for bit (the deterministic program the deterministic package's bits), every launch within the gate; the padding keeps its values.""" device, arch = _atomic_device_or_skip() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() X, W, G = _router_inputs(T, device, SEED + 100 + T, K=hidden) tensors = {"fwd": (X, W), "dgrad": (G, W), "wgrad": (G, X)} calls = { - "fwd": cake_backend.router_forward, - "dgrad": cake_backend.router_dgrad, - "wgrad": cake_backend.router_wgrad, + "fwd": cake_backend.cake_dense_gemm_nondet_train_router_forward, + "dgrad": cake_backend.cake_dense_gemm_nondet_train_router_dgrad, + "wgrad": cake_backend.cake_dense_gemm_nondet_train_router_wgrad, } det_calls = { - "fwd": det_backend.router_det_forward, - "dgrad": det_backend.router_det_dgrad, - "wgrad": det_backend.router_det_wgrad, + "fwd": det_backend.cake_dense_gemm_det_train_router_forward, + "dgrad": det_backend.cake_dense_gemm_det_train_router_dgrad, + "wgrad": det_backend.cake_dense_gemm_det_train_router_wgrad, } shapes = {"fwd": (T, 256), "dgrad": (T, hidden), "wgrad": (256, hidden)} side = torch.cuda.Stream(device) @@ -11013,7 +10962,7 @@ def test_gpu_router_entry_points_give_their_paths_results(monkeypatch, hidden, T # fmt: on for op, splits, pad, on_side in cases: P, Q = tensors[op] - kw = dict(deterministic=False) + kw = dict() if op != "dgrad": kw["splits"] = splits rows, cols = shapes[op] @@ -11048,12 +10997,12 @@ def test_gpu_router_entry_points_give_their_paths_results(monkeypatch, hidden, T def test_gpu_router_entry_points_capture_in_a_cuda_graph(monkeypatch): - """The router entry points' ``deterministic=False`` calls through ``_K2_ENTRY_CACHE`` (the forward with and + """The router entry points' calls through ``_K2_ENTRY_CACHE`` (the forward with and without a split count, the input gradient, the weight gradient of a row with a verdict and of one without), their keys cached on the capturing stream, capture in a CUDA graph without planning: every replay gives the eager calls' results -- bit for bit for the deterministic program and a reproducible plan -- within the gate.""" device, arch = _atomic_device_or_skip() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() stream = torch.cuda.Stream(device) cases = [] for hidden, T in ((5120, 64), (6144, 1000)): @@ -11071,13 +11020,15 @@ def test_gpu_router_entry_points_capture_in_a_cuda_graph(monkeypatch): def run_all(): for op, P, Q, splits, out in cases: if op == "fwd": - cake_backend.router_forward( - P, Q, out, splits=splits, deterministic=False + cake_backend.cake_dense_gemm_nondet_train_router_forward( + P, Q, out, splits=splits ) elif op == "dgrad": - cake_backend.router_dgrad(P, Q, out, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_dgrad(P, Q, out) else: - cake_backend.router_wgrad(P, Q, out, splits=splits, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad( + P, Q, out, splits=splits + ) stream.wait_stream(torch.cuda.current_stream(device)) with torch.cuda.stream(stream): @@ -11127,7 +11078,7 @@ def test_gpu_concurrent_streams_keep_their_own_workspaces(k1_rule, router_rule): """Two streams issuing K1 and K2 calls of this package of different geometries concurrently: one arena per (device, stream), every result within budget, the tickets zero.""" device, arch = _atomic_device_or_skip() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() streams = [torch.cuda.Stream(device) for _ in range(2)] cases = [] for i, (row, T) in enumerate((("indexer_k", 16231), ("thin_48", 4097))): @@ -11145,8 +11096,8 @@ def test_gpu_concurrent_streams_keep_their_own_workspaces(k1_rule, router_rule): for stream, (G, X, out), rout in zip(streams, cases, router_outs, strict=True): stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - cake_backend.projection_wgrad(G, X, out, deterministic=False) - cake_backend.router_wgrad(Gf, Xf, rout, deterministic=False) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + cake_backend.cake_dense_gemm_nondet_train_router_wgrad(Gf, Xf, rout) torch.cuda.synchronize() for G, X, out in cases: _k1_check(out, G.double().t() @ X.double()) @@ -11223,8 +11174,8 @@ def test_gpu_prepared_k1_reset_counters(): GRAPH_SAFETY_KINDS = { "k1_wgrad": {"k1_slots", "k1_tickets"}, "k1_forward": {"k1_det_ws", "k1_det_counters"}, - "router_wgrad": {"k2_stack"}, - "router_forward": {"k2_stack", "k2_splitk_partials"}, + "cake_dense_gemm_nondet_train_router_wgrad": {"k2_stack"}, + "cake_dense_gemm_nondet_train_router_forward": {"k2_stack", "k2_splitk_partials"}, } GRAPH_SAFETY_CASES = tuple(GRAPH_SAFETY_KINDS) #: ``(N, K, (T_small, T_large))`` of the k1_forward case: forwards whose deterministic plan splits K across stream-K @@ -11242,7 +11193,7 @@ def _graph_safety_k1_wgrad(row, T, device, seed, monkeypatch): _force_k1_rule(monkeypatch, G, X, out, 2) def call(o): - return cake_backend.projection_wgrad(G, X, o, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, o) def check(got, want): # two splits: one add of two partials, bit for bit the eager result @@ -11258,7 +11209,7 @@ def _graph_safety_k1_forward(T, N, K, device, seed): out = torch.empty(T, N, device=device, dtype=torch.bfloat16) def call(o): - return cake_backend.projection_forward(X, W, o, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_projection_forward(X, W, o) def check(got, want): assert torch.equal(got, want) # the deterministic program's bits @@ -11282,10 +11233,10 @@ def _graph_safety_router(op, T, device, seed, monkeypatch): monkeypatch.setitem(cake_backend.ROUTER_SM_COUNT, arch, sm_count) key = (arch, "nn_t", int(X.shape[1]), int(G.shape[1])) monkeypatch.setitem(table, key, "pairs") - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() def call(o): - return cake_backend.router_wgrad(G, X, o, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_router_wgrad(G, X, o) def check(got, want): assert bool(torch.isfinite(got).all()) @@ -11295,7 +11246,7 @@ def check(got, want): out = torch.empty(T, W.shape[0], device=device) def call(o): # no forward verdict: the deterministic program and its partials - return cake_backend.router_forward(X, W, o, deterministic=False) + return cake_backend.cake_dense_gemm_nondet_train_router_forward(X, W, o) def check(got, want): assert torch.equal(got, want) @@ -11343,7 +11294,7 @@ def _graph_safety_bound(case, stream): else: entries = [e for k, e in cake_backend._K2_ENTRY_CACHE.items() if k[1] == key] assert entries and all(e.launch is not None for e in entries) - if case == "router_wgrad": + if case == "cake_dense_gemm_nondet_train_router_wgrad": assert all(e.atomic and e.launch.plan.simt is None for e in entries) else: assert all(e.launch.k2.partials_shape is not None for e in entries) @@ -11461,7 +11412,7 @@ def _graph_safety_scenario(case: str, mode: str) -> None: @pytest.mark.parametrize("mode", ["grow", "evict"]) @pytest.mark.parametrize("case", GRAPH_SAFETY_CASES) def test_gpu_captured_convenience_call_keeps_its_workspace(case, mode): - """A convenience-path call captured into a CUDA graph stays valid until ``clear_convenience_caches``: after the + """A convenience-path call captured into a CUDA graph stays valid until ``cake_dense_gemm_nondet_train_convenience_cache_clear``: after the arena grows and the captured key is bound to the grown buffers, and after cache eviction drops every entry of the capturing stream, its replay gives the eager result and writes nothing into tensors allocated since (:func:`_graph_safety_scenario`, run in a child process with a timeout: a replay on freed workspace can fault or @@ -11471,7 +11422,7 @@ def test_gpu_captured_convenience_call_keeps_its_workspace(case, mode): for call, out, _ in _graph_safety_calls(case, device, mp): call(out) torch.cuda.synchronize() - cake_backend.clear_convenience_caches() + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() root = str(Path(__file__).resolve().parents[2]) code = ( "import importlib.util, sys\n" diff --git a/tests/experimental/test_cake_dense_gemm_nondet_train_names.py b/tests/experimental/test_cake_dense_gemm_nondet_train_names.py new file mode 100644 index 00000000000..9ea2c6427b2 --- /dev/null +++ b/tests/experimental/test_cake_dense_gemm_nondet_train_names.py @@ -0,0 +1,505 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +# The names of the non-deterministic Cake dense GEMMs (flashinfer-ai/flashinfer#8928): the public module +# ``flashinfer.cake_dense_gemm_nondet_train`` carries the mode word in every public name, no entry point takes a +# ``deterministic=`` keyword (nor any ``**kwargs``), the projection planner's keyword-only knobs are the deterministic +# planner's public ones plus ``splits`` and ``chains``, the former names and import paths are gone (no alias, no +# read-through to the deterministic package), and the module imports its backend only on first use. CPU only. + +import importlib +import importlib.util +import inspect +import re +import subprocess +import sys +import warnings +from types import ModuleType + +import pytest +import torch + +MODULE = "flashinfer.cake_dense_gemm_nondet_train" +BACKEND = "flashinfer.experimental.cake.cake_dense_gemm_nondet_train" +DET_MODULE = "flashinfer.cake_dense_gemm_det_train" +DET_BACKEND = "flashinfer.experimental.cake.cake_dense_gemm_det_train" +OP = "dense_gemm" +MODE = "nondet" +# the 20 public names of the module +PUBLIC = { + "ARCHES", + "is_cake_dense_gemm_nondet_train_supported", + "cake_dense_gemm_nondet_train_projection", + "cake_dense_gemm_nondet_train_projection_forward", + "cake_dense_gemm_nondet_train_projection_dgrad", + "cake_dense_gemm_nondet_train_projection_wgrad", + "cake_dense_gemm_nondet_train_projection_plan", + "cake_dense_gemm_nondet_train_projection_prepare", + "cake_dense_gemm_nondet_train_projection_wgrad_prepare", + "CakeDenseGemmNondetProjectionPlan", + "CakeDenseGemmNondetProjectionLaunch", + "cake_dense_gemm_nondet_train_router", + "cake_dense_gemm_nondet_train_router_forward", + "cake_dense_gemm_nondet_train_router_dgrad", + "cake_dense_gemm_nondet_train_router_wgrad", + "cake_dense_gemm_nondet_train_router_plan", + "cake_dense_gemm_nondet_train_router_prepare", + "CakeDenseGemmNondetRouterPlan", + "CakeDenseGemmNondetRouterLaunch", + "cake_dense_gemm_nondet_train_convenience_cache_clear", +} +# The callables among them (functions and classes; the module constants are not) and the functions. +CALLABLES = sorted(n for n in PUBLIC if not re.fullmatch(r"[A-Z][A-Z0-9_]*", n)) +FUNCTIONS = sorted(n for n in CALLABLES if not n[0].isupper()) +# The former import paths: the package and its modules. +FORMER_MODULES = ( + "flashinfer.experimental.dense_projection_gemm.cake_backend", + "flashinfer.experimental.dense_projection_gemm.cake_jit", + "flashinfer.experimental.dense_projection_gemm.cake_build", +) +FORMER_PACKAGE = "flashinfer.experimental.dense_projection_gemm" +# Removed backend names: the 18 former public names, the former-name map, the wrappers that only returned or ran a +# deterministic object, the mode resolvers and the forwarding of deterministic templates and programs. +REMOVED = { + "cake_backend": ( + "dense_projection_gemm", + "plan_dense_projection_gemm", + "prepare_dense_projection_gemm", + "projection_forward", + "projection_dgrad", + "prepare_projection_wgrad", + "projection_wgrad", + "plan_router_fp32_gemm", + "prepare_router_fp32_gemm", + "router_fp32_gemm", + "router_forward", + "router_dgrad", + "router_wgrad", + "clear_convenience_caches", + "AtomicGemmPlan", + "PreparedAtomicGemm", + "AtomicRouterPlan", + "PreparedAtomicRouterGemm", + "DETERMINISTIC_NAMES", + "instance_symbol", + "split_fp32_to_bf16x3_", + "_DET_TYPES", + "_DET_SPELLED", + "_det_entry", + "_resolve_deterministic", + "_require_false", + "_det_package_name", + ), + "cake_jit": ( + "is_det_template", + "is_det_program", + "DET_TEMPLATE_PREFIXES", + "DET_PROGRAM_PREFIX", + ), +} +# Every public name the backend module defined before #8046 (FlashInfer c472a430: functions, classes and constants). +# The backend defines five of them itself (helpers and constants of its own programs); every other one is absent: no +# alias and no read-through to the deterministic package. +# fmt: off +FORMER_BACKEND_NAMES = ( + "ARCH_RULE_DEFAULTS", "BLOCK_K", "BLOCK_M", "BLOCK_N_CHOICES", "B_SWZ_CHOICES", "CH_T_COLS", + "CONVENIENCE_CACHE_SIZE", "CTA_GROUP", "CTA_ROWS_CHOICES", "DeviceFacts", "EPI_MODES", "EPI_WARPS", + "EXACT_WAVE_BLOCK_N", "GROUP_M", "GemmPlan", "K_TMA_BF16", "K_TWO_SLOTS", "L2_HINTS", "L2_PROMOS", + "PANEL_BYTES", "PD_CHOICES", "PreparedGemm", "PreparedRouterGemm", "ROUTER_BLOCK_K", "ROUTER_BLOCK_M", + "ROUTER_BLOCK_N", "ROUTER_SPLITS", "ROW_RULES", "RULE_PARENT", "RouterPlan", "SK_COUNTERS_MIN", "SK_DUMMY_BASE", + "SK_MIN_ITERS", "SK_SLAB_CHOICES", "SLOT_BYTES", "SMEM_OPT_IN", "SMEM_OVERSIZED", "SPLIT_BLOCK_ELEMS", + "SPLIT_TEMPLATE", "SUPPORTED_COMPUTE_CAPABILITIES", "W8_FORMS", "WORK_STAGES", "a_halves_of", "arch_for", + "as_batched", "b_stage_bytes", "bind_launch", "box_rows_of", "check_stream_k_plan", "clear_convenience_caches", + "convenience_key", "default_block_n", "default_cta_rows", "default_group_m", "default_hints", + "default_hints_cgrp", "default_pf", "default_promo", "default_stages", "dense_projection_gemm", "device_facts", + "device_l2_bytes", "device_sm_count", "epi_cols", "epi_mode", "epi_slots", "exact_waves", "flat_alias", + "fused_split_module", "generated_program_available", "instance_key", "instance_symbol", "operand_view", + "plan_dense_projection_gemm", "plan_router_fp32_gemm", "prepare_dense_projection_gemm", + "prepare_projection_wgrad", "prepare_router_fp32_gemm", "projection_dgrad", "projection_forward", + "projection_wgrad", "require_arch", "router_dgrad", "router_forward", "router_fp32_gemm", "router_layout_class", + "router_symbol", "router_wgrad", "row_rule", "sk_exact_plan", "sk_parts_plan", "sk_sync_plan", "sm_pairs", + "sm_units", "smem_limit_for", "split_fp32_to_bf16x3_", "staging_bytes", "stream_k_plan", "stt_bytes", + "swap_small_m", "wave_working_set", "wave_working_set_cgrp", "wgrad_swapped", "wgrad_views", +) +# fmt: on +OWN_FORMER_NAMES = frozenset( + { + "ROUTER_BLOCK_K", + "ROUTER_BLOCK_M", + "ROUTER_BLOCK_N", + "bind_launch", + "router_symbol", + } +) +# Names of the deterministic backend (former and #8928 spellings) and two names no module has: none is reachable here. +UNREACHABLE = ( + "projection_det_forward", + "DetGemmPlan", + "router_fp32_det_gemm", + "ROW_RULES", + "stream_k_plan", + "wgrad_views", + "cake_dense_gemm_det_train_projection_forward", + "CakeDenseGemmDetProjectionPlan", + "__wrapped__", + "no_such_name", +) +# The deterministic planners' private knobs, which this module does not pass on. +DET_PRIVATE_KNOBS = {"_fallback", "_allow_swap", "_use_rules"} +# This package's own knobs of the projection planner. +OWN_KNOBS = {"splits", "chains"} + + +def mode_word_violations(name: str, *, op: str, mode: str) -> list[str]: + """Why the public name ``name`` of the module ``flashinfer.cake___train`` breaks the #8928 naming rules; + empty when it follows them. Self-contained (it reads no state of this test module), so another family can copy + it as is. + + * A module constant (``^[A-Z][A-Z0-9_]*$``, e.g. ``ARCHES``) is exempt: the module path carries the mode. + * The check is by word: ``name`` is split on ``_`` and on CamelCase boundaries, lower-cased, and must hold the + mode word ``mode`` and no other mode word (``det`` and ``nondet`` are distinct words, so ``"det" in name`` is + never the test), nor the verbs ``fwd`` / ``bwd``. + * A function is ``cake___train_<...>`` or the support query ``is_cake___train_supported``; + a class is ``Cake`` (```` and ```` in CamelCase). + """ + if re.fullmatch(r"[A-Z][A-Z0-9_]*", name): + return [] + words = [ + w.lower() + for part in name.split("_") + for w in re.findall(r"[A-Z]+(?![a-z])|[A-Z]?[a-z0-9]+", part) + ] + problems = [] + if mode not in words: + problems.append(f"no mode word {mode!r} in {words}") + others = {"det", "nondet"} - {mode} + if others & set(words): + problems.append(f"another mode word {sorted(others & set(words))} in {words}") + if {"fwd", "bwd"} & set(words): + problems.append(f"the verbs fwd / bwd are spelled forward / backward: {words}") + prefix = f"cake_{op}_{mode}_train" + if name[0].isupper(): + camel = "Cake" + "".join(w.capitalize() for w in f"{op}_{mode}".split("_")) + if not re.fullmatch(re.escape(camel) + r"[A-Z]\w*", name): + problems.append(f"a class is {camel}") + elif not (name.startswith(prefix + "_") or name == f"is_{prefix}_supported"): + problems.append(f"a function is {prefix}_ or is_{prefix}_supported") + return problems + + +def _run(code: str) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, timeout=600 + ) + + +def _keyword_only(fn) -> set[str]: + return { + n + for n, p in inspect.signature(fn).parameters.items() + if p.kind is inspect.Parameter.KEYWORD_ONLY + } + + +def test_all_lists_the_public_names(): + module = importlib.import_module(MODULE) + assert set(module.__all__) == PUBLIC + assert len(module.__all__) == len(PUBLIC) + assert sorted(module.__all__) == module.__all__ + + +@pytest.mark.parametrize("name", sorted(PUBLIC)) +def test_every_public_name_carries_the_mode_word(name): + assert mode_word_violations(name, op=OP, mode=MODE) == [] + + +def test_the_mode_word_check_reads_words(): + """The check splits words: ``nondet`` is not ``det``, a CamelCase class is split, constants are exempt.""" + check = mode_word_violations + assert check("cake_dense_gemm_nondet_train_router", op=OP, mode="nondet") == [] + assert check("cake_dense_gemm_nondet_train_router", op=OP, mode="det") + assert check("cake_dense_gemm_det_train_router", op=OP, mode="nondet") + assert check("CakeDenseGemmNondetRouterPlan", op=OP, mode="nondet") == [] + assert check("CakeDenseGemmNondetRouterPlan", op=OP, mode="det") + assert check("CakeDenseGemmDetRouterPlan", op=OP, mode="det") == [] + assert check("RouterPlan", op=OP, mode="nondet") + assert check("cake_dense_gemm_nondet_train_projection_fwd", op=OP, mode="nondet") + assert ( + check("is_cake_dense_gemm_nondet_train_supported", op=OP, mode="nondet") == [] + ) + assert check("router_wgrad", op=OP, mode="nondet") + assert check("ARCHES", op=OP, mode="nondet") == [] + assert check("Arches", op=OP, mode="nondet") + + +@pytest.mark.parametrize("name", CALLABLES) +def test_no_entry_point_takes_a_mode_keyword(name): + """``deterministic=`` does not bind to any callable of the module (a class: its constructor), and none takes + ``**kwargs``.""" + obj = getattr(importlib.import_module(MODULE), name) + assert callable(obj) + signature = inspect.signature(obj) + with pytest.raises(TypeError): + signature.bind_partial(deterministic=False) + assert not any( + p.kind is inspect.Parameter.VAR_KEYWORD for p in signature.parameters.values() + ) + if isinstance(obj, type) and issubclass(obj, torch.autograd.Function): + with pytest.raises(TypeError): + inspect.signature(obj.forward).bind_partial(None, deterministic=False) + + +@pytest.mark.parametrize("value", [False, True, None]) +@pytest.mark.parametrize("name", FUNCTIONS) +def test_a_mode_keyword_raises_type_error_before_anything_runs( + monkeypatch, name, value +): + """A call that passes ``deterministic=`` (any value, ``False`` included) raises ``TypeError`` before anything is + planned, allocated or launched: the deterministic package and this package's paths are never reached.""" + backend = importlib.import_module(f"{BACKEND}.cake_backend") + reached = [] + + def forbidden(*args, **kwargs): + reached.append((args, kwargs)) + raise AssertionError("a call with deterministic= ran") + + monkeypatch.setattr(backend, "deterministic_backend", forbidden) + for path in ("_atomic_k1", "_atomic_k2", "_det_k1", "_router_gemm", "_splitk_k2"): + monkeypatch.setattr(backend, path, forbidden) + for path in ("_k1_cached_launch", "_k2_entry", "_drop_convenience_entries"): + monkeypatch.setattr(backend, path, forbidden) + fn = getattr(importlib.import_module(MODULE), name) + required = [ + p + for p in inspect.signature(fn).parameters.values() + if p.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + and p.default is inspect.Parameter.empty + ] + with pytest.raises(TypeError, match="deterministic"): + fn(*(object() for _ in required), deterministic=value) + assert not reached + + +@pytest.mark.parametrize( + "helper", + ["router_symbol", "bind_launch", "plan_atomic_gemm", "plan_atomic_router"], +) +def test_the_backend_helpers_take_no_mode_keyword(helper): + backend = importlib.import_module(f"{BACKEND}.cake_backend") + with pytest.raises(TypeError): + inspect.signature(getattr(backend, helper)).bind_partial(deterministic=False) + + +@pytest.mark.parametrize( + "entry, det_entry", + [ + ("projection_plan", "projection_plan"), + ("projection_prepare", "projection_prepare"), + ], +) +def test_the_projection_planner_knobs_are_the_deterministic_planners(entry, det_entry): + """The keyword-only knobs of the projection plan / prepare entries are the deterministic entry's public ones plus + ``splits`` and ``chains`` (a change of the deterministic signature shows here), each defaulting to one private + sentinel: a knob the caller does not give is not passed on.""" + fn = getattr( + importlib.import_module(MODULE), f"cake_dense_gemm_nondet_train_{entry}" + ) + det_fn = getattr( + importlib.import_module(DET_MODULE), f"cake_dense_gemm_det_train_{det_entry}" + ) + det_knobs = _keyword_only(det_fn) + assert det_knobs >= DET_PRIVATE_KNOBS or entry == "projection_prepare" + public = {k for k in det_knobs if not k.startswith("_")} + assert _keyword_only(fn) == public | OWN_KNOBS + defaults = { + p.default + for p in inspect.signature(fn).parameters.values() + if p.kind is inspect.Parameter.KEYWORD_ONLY + } + assert len(defaults) == 1 and None not in defaults + + +def test_only_the_knobs_a_caller_gives_reach_the_deterministic_planner(monkeypatch): + """On a layout routed to the deterministic program the projection plan / prepare entries call the deterministic + planner with exactly the caller's knobs (none when the caller gives none).""" + backend = importlib.import_module(f"{BACKEND}.cake_backend") + det = importlib.import_module(f"{DET_BACKEND}.cake_backend") + calls = [] + + def recorder(*args, **kwargs): + calls.append(kwargs) + return "det-object" + + monkeypatch.setattr(det, "cake_dense_gemm_det_train_projection_plan", recorder) + monkeypatch.setattr(det, "cake_dense_gemm_det_train_projection_prepare", recorder) + monkeypatch.setattr(backend, "deterministic_backend", lambda: det) + A, B, out = torch.zeros(8, 32), torch.zeros(16, 32).t(), torch.zeros(8, 16) # kk + m = importlib.import_module(MODULE) + assert m.cake_dense_gemm_nondet_train_projection_plan(A, B, out) == "det-object" + assert ( + m.cake_dense_gemm_nondet_train_projection_plan( + A, B, out, sm_count=148, l2_bytes=1 << 20, sk=False + ) + == "det-object" + ) + assert m.cake_dense_gemm_nondet_train_projection_prepare(A, B, out) == "det-object" + assert ( + m.cake_dense_gemm_nondet_train_projection_prepare(A, B, out, enable_pdl=True) + == "det-object" + ) + assert calls == [ + {}, + {"sm_count": 148, "l2_bytes": 1 << 20, "sk": False}, + {}, + {"enable_pdl": True}, + ] + + +@pytest.mark.parametrize("name", sorted(PUBLIC)) +def test_the_module_returns_the_backend_object(name): + with warnings.catch_warnings(): + warnings.simplefilter("error") # no ExperimentalWarning: a stable module + obj = getattr(importlib.import_module(MODULE), name) + submodule = "cake_jit" if name == "ARCHES" else "cake_backend" + assert obj is getattr(importlib.import_module(f"{BACKEND}.{submodule}"), name) + if inspect.isfunction(obj) or inspect.isclass(obj): + assert obj.__name__ == name + assert inspect.unwrap(obj).__name__ == name + assert obj.__module__ == f"{BACKEND}.cake_backend" + + +def test_an_unknown_name_raises_attribute_error(): + module = importlib.import_module(MODULE) + with pytest.raises(AttributeError): + module.router_fp32_gemm # noqa: B018 + # the deterministic backend's setting (it bounds both packages), not a name of this module + with pytest.raises(AttributeError): + module.CONVENIENCE_CACHE_SIZE # noqa: B018 + + +@pytest.mark.parametrize("submodule, names", sorted(REMOVED.items())) +def test_the_former_names_are_removed(submodule, names): + backend = importlib.import_module(f"{BACKEND}.{submodule}") + assert [n for n in names if hasattr(backend, n)] == [] + + +@pytest.mark.parametrize("name", FORMER_BACKEND_NAMES) +def test_no_backend_name_from_before_8046_reads_through(name): + """A name the backend had before #8046 is this package's own object (defined in its module) or absent.""" + backend = importlib.import_module(f"{BACKEND}.cake_backend") + if name in OWN_FORMER_NAMES: + obj = vars(backend)[name] + assert not callable(obj) or obj.__module__ == backend.__name__ + else: + assert name not in vars(backend) and name not in dir(backend) + with pytest.raises(AttributeError): + getattr(backend, name) + + +@pytest.mark.parametrize("name", UNREACHABLE) +def test_deterministic_names_are_not_reachable(name): + backend = importlib.import_module(f"{BACKEND}.cake_backend") + with pytest.raises(AttributeError): + getattr(backend, name) + + +def test_no_read_through_to_the_deterministic_package(): + """The backend has no module ``__getattr__`` (no read-through), and no attribute of its modules is an object of + the deterministic package (the deterministic package is reached only through ``deterministic_backend()``); a + public name of the deterministic module is not an attribute here.""" + backend = importlib.import_module(f"{BACKEND}.cake_backend") + jit = importlib.import_module(f"{BACKEND}.cake_jit") + package = importlib.import_module(BACKEND) + assert "__getattr__" not in vars(backend) and "__getattr__" not in vars(jit) + det = importlib.import_module(f"{DET_BACKEND}.cake_backend") + det_jit = importlib.import_module(f"{DET_BACKEND}.cake_jit") + det_objects = { + id(obj) + for module in (det, det_jit) + for obj in vars(module).values() + if (inspect.isfunction(obj) or inspect.isclass(obj)) + and obj.__module__.startswith(DET_BACKEND) + } + for module in (backend, jit, package): + leaked = sorted( + n + for n, obj in vars(module).items() + if id(obj) in det_objects + or (isinstance(obj, ModuleType) and obj.__name__.startswith(DET_BACKEND)) + ) + assert not leaked, (module.__name__, leaked) + det_public = importlib.import_module(DET_MODULE).__all__ + for name in det_public: + if not re.fullmatch(r"[A-Z][A-Z0-9_]*", name): + assert not hasattr(backend, name), name + + +@pytest.mark.parametrize("name", FORMER_MODULES) +def test_the_former_modules_fail_at_import(name): + result = _run(f"import importlib; importlib.import_module({name!r})") + assert result.returncode != 0 + assert "ModuleNotFoundError" in result.stderr + + +def test_the_former_package_is_gone(): + # A checkout may keep an untracked ``aot/`` or ``__pycache__/`` in the former directory, which Python + # imports as a namespace package (no ``origin``); it holds no module. + result = _run( + "import importlib.util; " + f"spec = importlib.util.find_spec({FORMER_PACKAGE!r}); " + "print(spec is None or spec.origin is None)" + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "True" + + +def test_imports_are_lazy(): + code = f""" +import sys +import flashinfer +loaded = sorted(m for m in sys.modules if "cake_dense_gemm_" in m) +assert loaded == [], loaded +for name in {sorted(PUBLIC)!r}: + assert name not in vars(flashinfer), name +rows = getattr(flashinfer, "_LAZY_EXPORTS", {{}}) +dense = {{k: v for k, v in rows.items() if "cake_dense_gemm" in str(v[0])}} +assert all(v[1] is None for v in dense.values()), dense +import {MODULE} as m +loaded = sorted(n for n in sys.modules if n.startswith("{BACKEND}")) +assert loaded == [], loaded +assert not [n for n in sys.modules if n.startswith("{DET_BACKEND}")] +m.cake_dense_gemm_nondet_train_router_wgrad +assert "{BACKEND}.cake_backend" in sys.modules +print("ok") +""" + result = _run(code) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "ok" + + +def test_the_support_query(): + module = importlib.import_module(MODULE) + det = importlib.import_module(DET_MODULE) + with pytest.raises(TypeError): + module.is_cake_dense_gemm_nondet_train_supported() # the device is required + assert ( + module.is_cake_dense_gemm_nondet_train_supported(torch.device("cpu")) is False + ) + if torch.cuda.is_available(): + device = torch.device("cuda", torch.cuda.current_device()) + assert module.is_cake_dense_gemm_nondet_train_supported( + device + ) == det.is_cake_dense_gemm_det_train_supported(device) From e9ec8101b762b854ba3e1fa1c83a506457cfe8ca Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 00:29:40 -0800 Subject: [PATCH 10/13] fix(cake_dense_gemm_nondet_train): standalone vendoring of the backend The backend declares its one dependency for a standalone build (#6553): the deterministic backend directory cake_dense_gemm_det_train, shipped next to it (build.json common.requires.packages, README "Standalone build (vendoring)" and the package docstring). - cake_jit.deterministic_module resolves the top-level sibling cake_dense_gemm_det_train when this package is vendored as a top-level package, instead of raising; a missing sibling raises ImportError that names it and use_deterministic_backend. - tests/experimental/test_cake_dense_gemm_nondet_train_standalone.py checks the build.json declaration and loads the directory as a top-level package with flashinfer blocked: with the deterministic directory beside it the first use resolves that sibling; without it the import succeeds and the first use raises the ImportError. - The facade's __dir__ lists the public names, as the deterministic facade does. --- flashinfer/cake_dense_gemm_nondet_train.py | 2 +- .../cake_dense_gemm_nondet_train/README.md | 24 ++- .../cake_dense_gemm_nondet_train/__init__.py | 28 ++- .../cake_dense_gemm_nondet_train/build.json | 3 + .../cake_backend.py | 16 +- .../cake_dense_gemm_nondet_train/cake_jit.py | 36 ++-- .../test_cake_dense_gemm_nondet_train.py | 7 +- ...cake_dense_gemm_nondet_train_standalone.py | 188 ++++++++++++++++++ 8 files changed, 263 insertions(+), 41 deletions(-) create mode 100644 tests/experimental/test_cake_dense_gemm_nondet_train_standalone.py diff --git a/flashinfer/cake_dense_gemm_nondet_train.py b/flashinfer/cake_dense_gemm_nondet_train.py index a85c53124ec..e984f267fa6 100644 --- a/flashinfer/cake_dense_gemm_nondet_train.py +++ b/flashinfer/cake_dense_gemm_nondet_train.py @@ -93,4 +93,4 @@ def __getattr__(name: str) -> Any: def __dir__() -> list[str]: - return sorted(set(globals()) | set(__all__)) + return list(__all__) diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md index 1f1c4433934..a1e9fbe8ece 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md @@ -79,8 +79,8 @@ On some rows this module runs the deterministic package's fixed-order programs: program today. The "run-to-run bits" rows below describe the current programs, not a contract. - The tensors passed between the forward and the gradients carry no mode: recompute a checkpointed forward with the module that made the saved one (same package only). Neither package has an autograd Function or a saved tape. -- This package imports the deterministic package at run time, as its sibling (a vendored copy binds it with - `cake_backend.use_deterministic_backend`). It re-exports no name of the deterministic package. +- This package imports the deterministic package at run time, as its sibling (see + [Standalone build](#standalone-build-vendoring)). It re-exports no name of the deterministic package. | operation | route | |---|---| @@ -1035,6 +1035,20 @@ the deterministic package's names there. | `cake_jit.py` | the registry of this package's programs (`ARG_PLANS`, `MODULES`, `KERNELS`) and the loader | | `cake_build.py`, `build.json`, `csrc/cake_dense_projection_gemm/` | the `cake-build/1` loader, manifest and generated sources of this package's programs (written by the Cake exporter) | -The deterministic package is resolved as this package's sibling on first use; the package never imports `flashinfer` -by absolute name. A consumer that vendors this package vendors the deterministic package too and binds it with -`cake_backend.use_deterministic_backend()`. +## Standalone build (vendoring) + +This backend directory builds and runs outside FlashInfer (flashinfer-ai/flashinfer#6553) with one declared +dependency: the deterministic backend directory `cake_dense_gemm_det_train` +(`flashinfer/experimental/cake/cake_dense_gemm_det_train`), whose fixed-order programs this package runs. `build.json` +declares it (`common.requires.packages`: `["cake_dense_gemm_det_train"]`, sibling directory names). A vendoring build +ships the two directories together, next to each other. + +- This package imports the deterministic package lazily, on the first call that needs it, as its sibling: + `.cake_dense_gemm_det_train` when this package sits in a parent package (a FlashInfer tree), the top-level + package `cake_dense_gemm_det_train` when this package is vendored as a top-level package. Importing this package + never imports the deterministic one; without it, the first use raises `ImportError` that names it. +- A tree that holds the deterministic package under another name binds it with + `cake_backend.use_deterministic_backend()` before the first call that needs it. +- The package imports nothing else outside its own directory. It imports `flashinfer` only in `cake_build`'s guarded + fallback to FlashInfer's JIT when a library is not prebuilt. +- The mode is the package name: no keyword argument and no environment variable selects it. diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py index 9e7d9ec8f62..1ccabcbec70 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/__init__.py @@ -12,17 +12,25 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. -""" -# Cake backend, non-deterministic: the GLM-5.2 dense projection GEMMs for training -# and the FP32 router-gate GEMM on SM100 / SM103 / SM107 -# (flashinfer-ai/flashinfer#5677), with programs that add split-K partials in -# the order they finish. The public module is -# flashinfer.cake_dense_gemm_nondet_train; it selects the mode, so no entry -# point takes a mode argument. The rows that run fixed-order programs reuse the -# programs of the deterministic sibling package cake_dense_gemm_det_train. -# Importing this package imports nothing: ``cake_backend``, ``cake_jit`` and -# ``cake_build`` load on first attribute access. +Cake backend, non-deterministic: the GLM-5.2 dense projection GEMMs for training +and the FP32 router-gate GEMM on SM100 / SM103 / SM107 +(flashinfer-ai/flashinfer#5677), with programs that add split-K partials in the +order they finish. The public module is ``flashinfer.cake_dense_gemm_nondet_train``. +The mode is the package name: no entry point takes a mode argument, and no +environment variable selects it. Importing this package imports nothing: +``cake_backend``, ``cake_jit`` and ``cake_build`` load on first attribute access. + +Standalone build (flashinfer-ai/flashinfer#6553): this package needs the +deterministic backend directory ``cake_dense_gemm_det_train`` shipped next to it +(``build.json``: ``common.requires.packages``). +The rows that run fixed-order programs reuse that package's programs; this +package imports it lazily, on first use, as its sibling (the top-level package +``cake_dense_gemm_det_train`` when this package is vendored as a top-level +package). It imports nothing else outside its own directory, and it imports +``flashinfer`` only in ``cake_build``'s guarded fallback to FlashInfer's JIT when +a library is not prebuilt. +""" __all__ = ["cake_backend", "cake_build", "cake_jit"] diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json index 4b3f17665c8..0ccf1c2fa48 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/build.json @@ -45,6 +45,9 @@ "cuda_toolkit": { "min": "12.8" }, + "packages": [ + "cake_dense_gemm_det_train" + ], "tvm_ffi": { "include_dirs": [ "tvm_ffi.libinfo.find_include_path()", diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py index 48e64da3df1..86c86db4f06 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py @@ -89,7 +89,8 @@ (``CakeDenseGemmDetProjectionPlan``, ``CakeDenseGemmDetProjectionLaunch``, ``CakeDenseGemmDetRouterPlan``, ``CakeDenseGemmDetRouterLaunch``). A row that runs a program of this package never goes to the deterministic package. The deterministic package is imported at run time as this package's -sibling, never through an absolute ``flashinfer`` import; a vendored copy binds it with +sibling (the top-level package ``cake_dense_gemm_det_train`` when this package is vendored as a top-level package), +never through an absolute ``flashinfer`` import; a copy under another name is bound with :func:`use_deterministic_backend`. The deterministic package never imports this one, and this module re-exports no name of it. :func:`cake_dense_gemm_nondet_train_convenience_cache_clear` clears the convenience caches and workspace arenas of both packages; a convenience-path call that this package launches from its arenas, captured into a CUDA graph, stays valid until then @@ -151,12 +152,13 @@ def use_deterministic_backend(module: ModuleType) -> None: - """Bind the deterministic package's ``cake_backend`` module explicitly. In a FlashInfer tree the deterministic - package is this package's sibling and is resolved on first use; a vendored copy of this package (imported - under another name, without the ``flashinfer`` package around it) vendors the deterministic package too and - binds its ``cake_backend`` here before the first call that needs it. Binding drops this package's convenience - caches and arenas (:func:`_drop_convenience_entries`): their entries hold the plans and programs of the package - bound before. Like :func:`cake_dense_gemm_nondet_train_convenience_cache_clear`, it invalidates the CUDA graphs that captured a + """Bind the deterministic package's ``cake_backend`` module explicitly. The deterministic package is this + package's sibling and is resolved on first use, in a FlashInfer tree and in a vendored tree that ships + ``cake_dense_gemm_det_train`` next to this package (both top-level packages, or both in one parent package); a + vendored tree that holds it under another name binds its ``cake_backend`` here before the first call that + needs it. Binding drops this package's convenience caches and arenas (:func:`_drop_convenience_entries`): their + entries hold the plans and programs of the package bound before. Like + :func:`cake_dense_gemm_nondet_train_convenience_cache_clear`, it invalidates the CUDA graphs that captured a convenience-path call of this package.""" missing = [name for name in _DET_ENTRIES if not hasattr(module, name)] if missing: diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py index 0503a1a511e..36dcb49239d 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_jit.py @@ -128,9 +128,10 @@ "dense_proj_gemm_dsm_", ) -#: The deterministic package: a sibling of this package (``.cake_dense_gemm_det_train``), resolved on first -#: use and never through an absolute ``flashinfer`` import; a vendored copy binds it explicitly -#: (``cake_backend.use_deterministic_backend``). +#: The deterministic package, which this package needs next to it: its sibling, resolved on first use and never +#: through an absolute ``flashinfer`` import -- ``.cake_dense_gemm_det_train`` when this package sits in a +#: parent package (a FlashInfer tree), the top-level package ``cake_dense_gemm_det_train`` when this package is +#: vendored as a top-level package. A copy under another name is bound with ``cake_backend.use_deterministic_backend``. DETERMINISTIC_PACKAGE = "cake_dense_gemm_det_train" _det_package: Optional[str] = None @@ -143,18 +144,25 @@ def bind_deterministic_package(package: str) -> None: def deterministic_module(name: str) -> ModuleType: - """The submodule ``name`` (``cake_backend``, ``cake_jit``, ``cake_build``) of the deterministic package.""" + """The submodule ``name`` (``cake_backend``, ``cake_jit``, ``cake_build``) of the deterministic package: the + package bound with ``cake_backend.use_deterministic_backend``, else the sibling :data:`DETERMINISTIC_PACKAGE` + (a top-level package when this package is one). Raises ``ImportError`` when that sibling is not installed.""" package = _det_package - if package is None: - parent = __name__.rpartition(".")[0].rpartition(".")[0] - if not parent: - raise RuntimeError( - f"{__name__.rpartition('.')[0]} reuses the programs of the deterministic package " - f"{DETERMINISTIC_PACKAGE!r}, which a vendored copy has no sibling of; vendor that package too and " - "bind it with cake_backend.use_deterministic_backend()" - ) - package = f"{parent}.{DETERMINISTIC_PACKAGE}" - return importlib.import_module(f"{package}.{name}") + if package is not None: + return importlib.import_module(f"{package}.{name}") + parent = __name__.rpartition(".")[0].rpartition(".")[0] + package = f"{parent}.{DETERMINISTIC_PACKAGE}" if parent else DETERMINISTIC_PACKAGE + try: + return importlib.import_module(f"{package}.{name}") + except ModuleNotFoundError as exc: + if exc.name not in (package, f"{package}.{name}"): + raise + raise ImportError( + f"{__name__.rpartition('.')[0]} reuses the programs of the deterministic package " + f"{DETERMINISTIC_PACKAGE!r} and imports it as its sibling {package!r}, which is not installed: vendor " + f"{DETERMINISTIC_PACKAGE} next to this package, or bind a copy with " + "cake_backend.use_deterministic_backend()" + ) from exc def arg_plan_key(plan) -> str: diff --git a/tests/experimental/test_cake_dense_gemm_nondet_train.py b/tests/experimental/test_cake_dense_gemm_nondet_train.py index af39df8b3d7..09559b1c401 100644 --- a/tests/experimental/test_cake_dense_gemm_nondet_train.py +++ b/tests/experimental/test_cake_dense_gemm_nondet_train.py @@ -150,10 +150,9 @@ def test_cake_build_is_the_pinned_loader(): def test_the_deterministic_package_resolves_as_a_sibling(monkeypatch): - """In a FlashInfer tree and in a copy vendored with the deterministic package beside it (the vendored gate's - layout) the deterministic package resolves as this package's sibling.""" - if not cake_jit.__name__.rpartition(".")[0].rpartition(".")[0]: - pytest.skip("a top-level copy has no parent package to hold a sibling") + """In a FlashInfer tree and in a copy vendored with the deterministic package beside it (both in one parent + package, or both top-level packages: ``test_cake_dense_gemm_nondet_train_standalone.py``) the deterministic + package resolves as this package's sibling.""" monkeypatch.setattr(cake_jit, "_det_package", None) assert cake_backend.deterministic_backend() is det_backend assert cake_jit.deterministic_module("cake_jit") is det_jit diff --git a/tests/experimental/test_cake_dense_gemm_nondet_train_standalone.py b/tests/experimental/test_cake_dense_gemm_nondet_train_standalone.py new file mode 100644 index 00000000000..bd1827ac1cb --- /dev/null +++ b/tests/experimental/test_cake_dense_gemm_nondet_train_standalone.py @@ -0,0 +1,188 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Standalone-build tests of the non-deterministic Cake dense GEMM backend +(flashinfer-ai/flashinfer#6553, #8928). + +The backend directory declares one dependency: the deterministic backend +directory ``cake_dense_gemm_det_train``, shipped next to it (``build.json``: +``common.requires.packages``). In a ``python -I`` +subprocess with ``flashinfer`` blocked, the directory loads as a top-level +package under its own name, and its ``cake_backend``, ``cake_jit`` and +``cake_build`` import without any ``flashinfer`` module. With the deterministic +directory loaded as a top-level package beside it, the first use resolves that +sibling; without it, importing still succeeds and the first use raises an +``ImportError`` that names the missing directory. CPU only: no GPU, no +generated program and no JIT build. +""" + +import json +import os +import site +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from tests.test_helpers.cake_package import cake_package + +tvm_ffi = pytest.importorskip("tvm_ffi") # cake_backend imports it at module level + +ND_DIR = Path(cake_package("cake_dense_gemm_nondet_train").__file__).resolve().parent +DET_DIR = Path(cake_package("cake_dense_gemm_det_train").__file__).resolve().parent + + +def _dependency_paths() -> list[str]: + """Where ``torch`` and ``tvm_ffi`` live, plus the interpreter's site-packages (``python -I`` drops + ``PYTHONPATH`` and the user site).""" + paths = list(site.getsitepackages()) + for module in (torch, tvm_ffi): + paths.append(os.path.dirname(os.path.dirname(os.path.abspath(module.__file__)))) + return list(dict.fromkeys(paths)) + + +_STANDALONE = r""" +import importlib +import importlib.abc +import importlib.util +import json +import os +import sys + +ND_DIR, DET_DIR = sys.argv[1], sys.argv[2] +for p in sys.argv[3:]: + if p not in sys.path: + sys.path.append(p) + + +class _Block(importlib.abc.MetaPathFinder): + def find_spec(self, name, path=None, target=None): + if name == "flashinfer" or name.startswith("flashinfer."): + raise ImportError(f"flashinfer is blocked in this test: {name}") + return None + + +sys.meta_path.insert(0, _Block()) +sys.modules["flashinfer"] = None + + +def load_top_level(directory): + name = os.path.basename(directory) + spec = importlib.util.spec_from_file_location( + name, os.path.join(directory, "__init__.py"), submodule_search_locations=[directory] + ) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def flashinfer_modules(): + return sorted( + n for n, m in sys.modules.items() if (n == "flashinfer" or n.startswith("flashinfer.")) and m is not None + ) + + +nd = load_top_level(ND_DIR) +det = load_top_level(DET_DIR) if DET_DIR else None +backend, jit, build = (importlib.import_module(f"{nd.__name__}.{sub}") for sub in ("cake_backend", "cake_jit", "cake_build")) +facts = { + "package": nd.__name__, + "submodules": sorted(m.__name__ for m in (backend, jit, build)), + "det_package": jit.DETERMINISTIC_PACKAGE, + "det_loaded_before_use": sorted(n for n in sys.modules if n.startswith(jit.DETERMINISTIC_PACKAGE + ".")), + "flashinfer_after_import": flashinfer_modules(), +} +try: + resolved = jit.deterministic_module("cake_backend") +except Exception as exc: # the case without the sibling: report the error + facts["error"] = [type(exc).__name__, str(exc)] + try: + backend.deterministic_backend() + except Exception as again: + facts["backend_error"] = [type(again).__name__, str(again)] +else: + facts["resolved"] = resolved.__name__ + facts["resolved_file"] = os.path.realpath(resolved.__file__) + facts["resolved_is_loaded_sibling"] = resolved is sys.modules.get(f"{det.__name__}.cake_backend") + facts["backend_resolves_same"] = backend.deterministic_backend() is resolved +facts["flashinfer_after_use"] = flashinfer_modules() +print(json.dumps(facts)) +""" + + +def _standalone(det_dir: str) -> dict: + result = subprocess.run( + [ + sys.executable, + "-I", + "-c", + _STANDALONE, + str(ND_DIR), + det_dir, + *_dependency_paths(), + ], + capture_output=True, + text=True, + timeout=300, + env={**os.environ, "CUDA_VISIBLE_DEVICES": ""}, + ) + assert result.returncode == 0, result.stderr + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def test_build_json_declares_the_deterministic_directory(): + """The one dependency outside the directory is declared for vendoring builds: ``common.requires.packages`` holds + the sibling directory names.""" + manifest = json.loads((ND_DIR / "build.json").read_text()) + assert ( + manifest["common"]["requires"]["packages"] + == [DET_DIR.name] + == ["cake_dense_gemm_det_train"] + ) + + +def test_the_backend_loads_as_a_top_level_package_and_resolves_its_top_level_sibling(): + facts = _standalone(str(DET_DIR)) + name = ND_DIR.name + assert facts["package"] == name + assert facts["submodules"] == sorted( + f"{name}.{sub}" for sub in ("cake_backend", "cake_build", "cake_jit") + ) + # the dependency is imported lazily: loading the package and its submodules imports none of its modules + assert facts["det_loaded_before_use"] == [] + assert facts["flashinfer_after_import"] == [] + # first use: the top-level sibling cake_dense_gemm_det_train, from the directory shipped next to it + assert facts["det_package"] == DET_DIR.name == "cake_dense_gemm_det_train" + assert facts["resolved"] == f"{DET_DIR.name}.cake_backend" + assert facts["resolved_file"].startswith(os.path.realpath(DET_DIR) + os.sep) + assert facts["resolved_is_loaded_sibling"] is True + assert facts["backend_resolves_same"] is True + assert facts["flashinfer_after_use"] == [] + + +def test_without_the_deterministic_directory_import_succeeds_and_first_use_names_it(): + facts = _standalone("") + assert facts["package"] == ND_DIR.name + assert facts["flashinfer_after_import"] == [] + assert "resolved" not in facts + for key in ("error", "backend_error"): + kind, message = facts[key] + assert kind == "ImportError", facts[key] + assert "cake_dense_gemm_det_train" in message + assert "use_deterministic_backend" in message + assert facts["flashinfer_after_use"] == [] From 86cc396d937dfef768a5850e5f76366fd254a71d Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 01:04:49 -0800 Subject: [PATCH 11/13] docs(experimental/cake): the #8928 package docstring The docstring of flashinfer/experimental/cake/__init__.py, shared byte for byte by every #8928 family PR (the first one to merge lands it). --- flashinfer/experimental/cake/__init__.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/flashinfer/experimental/cake/__init__.py b/flashinfer/experimental/cake/__init__.py index ce5aa5c1a4a..1186d96c924 100644 --- a/flashinfer/experimental/cake/__init__.py +++ b/flashinfer/experimental/cake/__init__.py @@ -13,10 +13,16 @@ See the License for the specific language governing permissions and limitations under the License. -Experimental Cake packages without a run-to-run bitwise contract: the -non-deterministic counterparts of the deterministic (``*_det_*``) Cake training -packages. Each one has a facade whose ``deterministic: bool = True`` keyword -routes to the deterministic package and whose ``deterministic=False`` runs the -package's own non-deterministic programs. Each subpackage is standalone -(import it by its full name); this module imports nothing. +Backend packages of the Cake training kernels (issue #8928). A subpackage +named ``cake___train[_kcp]`` is the backend of the top-level +module of the same name, ``flashinfer.cake___train[_kcp]``; +both modes live here. For these packages the mode is part of the module name +and no entry point takes a ``deterministic=`` keyword. A ``det`` package +returns run-to-run bitwise-identical results. A ``nondet`` package makes no +run-to-run bitwise promise; under ``torch.use_deterministic_algorithms(True)`` +it raises PyTorch's alert at its non-deterministic step (with +``warn_only=True`` it warns and runs ``nondet``) and never falls back to +``det``. A subpackage whose family is not renamed yet keeps its current entry +points until that rename. Each subpackage is standalone (import it by its +full name); this module imports nothing. """ From ffc40e893ffaf897341e31bd4886a7dc38fd31de Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 01:46:22 -0800 Subject: [PATCH 12/13] test(cake_dense_gemm_nondet_train): name the router graph-safety cases k2_wgrad / k2_forward The rename turned the case labels into the long public names, so _graph_safety_calls read the operation word 'dense' and ran the forward call for the weight-gradient cases (GPU-only failure on B200). --- tests/experimental/test_cake_dense_gemm_nondet_train.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/experimental/test_cake_dense_gemm_nondet_train.py b/tests/experimental/test_cake_dense_gemm_nondet_train.py index 09559b1c401..3a881d13c06 100644 --- a/tests/experimental/test_cake_dense_gemm_nondet_train.py +++ b/tests/experimental/test_cake_dense_gemm_nondet_train.py @@ -11173,8 +11173,8 @@ def test_gpu_prepared_k1_reset_counters(): GRAPH_SAFETY_KINDS = { "k1_wgrad": {"k1_slots", "k1_tickets"}, "k1_forward": {"k1_det_ws", "k1_det_counters"}, - "cake_dense_gemm_nondet_train_router_wgrad": {"k2_stack"}, - "cake_dense_gemm_nondet_train_router_forward": {"k2_stack", "k2_splitk_partials"}, + "k2_wgrad": {"k2_stack"}, + "k2_forward": {"k2_stack", "k2_splitk_partials"}, } GRAPH_SAFETY_CASES = tuple(GRAPH_SAFETY_KINDS) #: ``(N, K, (T_small, T_large))`` of the k1_forward case: forwards whose deterministic plan splits K across stream-K @@ -11293,7 +11293,7 @@ def _graph_safety_bound(case, stream): else: entries = [e for k, e in cake_backend._K2_ENTRY_CACHE.items() if k[1] == key] assert entries and all(e.launch is not None for e in entries) - if case == "cake_dense_gemm_nondet_train_router_wgrad": + if case == "k2_wgrad": assert all(e.atomic and e.launch.plan.simt is None for e in entries) else: assert all(e.launch.k2.partials_shape is not None for e in entries) From 4647f1ba33df966bad8ce304cd0395f7256f70f3 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 6 Nov 2026 02:06:16 -0800 Subject: [PATCH 13/13] test(cake_dense_gemm_nondet_train): the cluster program under the PyTorch deterministic mode on a GPU - test_gpu_the_cluster_program_does_not_alert: on a 152-SM sm_100a part the shipped K1_DSM_RANGES_BY_SM row 4096 -> 64 FP32 at 2048 tokens (cluster program, three splits) runs without an alert or a warning under torch.use_deterministic_algorithms(True) and with warn_only=True (planning call, cached call, prepared launch), runs the cluster program by kernel name and meets the FP32 budget against FP64. Other devices skip. - The lazy-import test checks the first attribute access, in a fresh process, for warnings (no ExperimentalWarning). - Wording: under the mode an order-varying launch raises before the program is launched (out untouched); the first convenience call may still zero its workspace. Bench description, test and section comments name the non-deterministic module instead of a numerical mode. - Note that _given_knobs(locals()) must stay the first statement. --- .../bench_cake_dense_gemm_nondet_train.py | 2 +- .../cake_dense_gemm_nondet_train/README.md | 2 +- .../cake_backend.py | 16 +++-- .../test_cake_dense_gemm_nondet_train.py | 71 ++++++++++++++++++- ...test_cake_dense_gemm_nondet_train_names.py | 7 +- 5 files changed, 87 insertions(+), 11 deletions(-) diff --git a/benchmarks/bench_cake_dense_gemm_nondet_train.py b/benchmarks/bench_cake_dense_gemm_nondet_train.py index a97f57b2ef4..7f2efb7024d 100644 --- a/benchmarks/bench_cake_dense_gemm_nondet_train.py +++ b/benchmarks/bench_cake_dense_gemm_nondet_train.py @@ -1678,7 +1678,7 @@ def image_identity(args): def main(): parser = argparse.ArgumentParser( - description="Benchmark the numerical modes of the Cake dense projection GEMMs (the module docstring has the arms, timing and reports).", + description="Benchmark the non-deterministic Cake dense GEMMs (flashinfer.cake_dense_gemm_nondet_train) against the deterministic module and the non-Cake SOTA (the module docstring has the arms, timing and reports).", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md index a1e9fbe8ece..846f9509bee 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/README.md @@ -43,7 +43,7 @@ result can vary run to run raises PyTorch's alert right before it launches: | launch | rows | order of the partials | under `torch.use_deterministic_algorithms(True)` | |---|---|---|---| -| K1 CTA-pair split-K program, three or more splits | the `sm_107a` rows of `K1_SPLIT_RULES` (216 SMs) and `K1_SPLIT_RULES_BY_SM` (208 / 212 SMs); a caller's `splits >= 3` | the order the splits finish | `RuntimeError` before the launch (nothing launched, `out` untouched); with `warn_only=True` a `UserWarning`, then the launch | +| K1 CTA-pair split-K program, three or more splits | the `sm_107a` rows of `K1_SPLIT_RULES` (216 SMs) and `K1_SPLIT_RULES_BY_SM` (208 / 212 SMs); a caller's `splits >= 3` | the order the splits finish | `RuntimeError` before the launch (the program is not launched and `out` is untouched); with `warn_only=True` a `UserWarning`, then the launch | | K1 CTA-pair split-K program, two splits | every `sm_100a` / `sm_103a` row, `sm_107a` `8192 -> 128` | one add of two values (repeats bit for bit) | no alert | | K1 one split, the chain programs (`_c3`) | the one-split rows | whole tiles, chains added in order | no alert | | K1 cluster program `dense_proj_gemm_dsm_*_s2` / `_s3` | `K1_DSM_RANGES_BY_SM` (152-SM `sm_100a` parts, FP32 out) | cluster-rank order in distributed shared memory (no global atomics, no workspace) | no alert | diff --git a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py index 86c86db4f06..694b8e53704 100644 --- a/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py +++ b/flashinfer/experimental/cake/cake_dense_gemm_nondet_train/cake_backend.py @@ -68,12 +68,12 @@ Under ``torch.use_deterministic_algorithms(True)`` a launch whose result can vary run to run -- the K1 split-K program of CTA pairs at three or more splits, a K2 stream-K partition that gives an element more than two partials -- raises -PyTorch's alert right before it launches (:func:`_alert_order_varying`): ``RuntimeError`` (nothing launched, ``out`` -untouched), or with ``warn_only=True`` a ``UserWarning``, after which the launch runs. The fixed-order launches (the -one- and two-split K1 rows, the K1 cluster program, the K2 SIMT programs, a K2 partition of at most two partials per -element, every deterministic program) do not alert. The alert is checked on every launch, never at plan or prepare -time; a launch captured into a CUDA graph is checked at capture, not at replay. The PyTorch setting never changes -which program a call runs, and no call goes to the deterministic package because of it. +PyTorch's alert right before it launches (:func:`_alert_order_varying`): ``RuntimeError`` (the program is not launched +and ``out`` is untouched), or with ``warn_only=True`` a ``UserWarning``, after which the launch runs. The fixed-order +launches (the one- and two-split K1 rows, the K1 cluster program, the K2 SIMT programs, a K2 partition of at most two +partials per element, every deterministic program) do not alert. The alert is checked on every launch, never at plan or +prepare time; a launch captured into a CUDA graph is checked at capture, not at replay. The PyTorch setting never +changes which program a call runs, and no call goes to the deterministic package because of it. The two packages ---------------- @@ -130,7 +130,7 @@ from .cake_jit import select_module as select_module # --------------------------------------------------------------------------- -# The deterministic package and the mode rules +# The deterministic package # --------------------------------------------------------------------------- #: The deterministic package's entry points this package calls on the rows that run the deterministic programs @@ -5094,6 +5094,7 @@ def cake_dense_gemm_nondet_train_projection_plan( (``sm_count`` and ``l2_bytes`` included): a knob the caller gives reaches that planner as given, and ``transposed_out``, ``arch`` and ``sm_count`` also apply to this package's plan; a knob the caller does not give takes the deterministic planner's default.""" + # must stay the first statement: ``locals()`` holds only the parameters here knobs = _given_knobs(locals()) if _route(_k1_layout(A, B)) == "det": return deterministic_backend().cake_dense_gemm_det_train_projection_plan( @@ -5166,6 +5167,7 @@ def cake_dense_gemm_nondet_train_projection_prepare( ``CakeDenseGemmDetProjectionLaunch``). The other keyword-only knobs are that entry point's public ones, passed on only when the caller gives them; ``enable_pdl`` (default: the device supports it) makes either launch a programmatic dependent launch.""" + # must stay the first statement: ``locals()`` holds only the parameters here knobs = _given_knobs(locals()) if _route(_k1_layout(A, B)) == "det": return deterministic_backend().cake_dense_gemm_det_train_projection_prepare( diff --git a/tests/experimental/test_cake_dense_gemm_nondet_train.py b/tests/experimental/test_cake_dense_gemm_nondet_train.py index 3a881d13c06..4053280efb4 100644 --- a/tests/experimental/test_cake_dense_gemm_nondet_train.py +++ b/tests/experimental/test_cake_dense_gemm_nondet_train.py @@ -55,7 +55,7 @@ # FlashInfer tree this names the sibling the package resolves by itself). cake_backend.use_deterministic_backend(det_backend) -# Accuracy gates (the deterministic package's, unchanged by the numerical mode). +# Accuracy gates (the deterministic package's budgets, which this module keeps). BF16_ATOL = BF16_RTOL = 1e-2 F32_REL_FRO_MAX = 5e-5 F32_MAX_ABS_OVER_REF_ABSMAX = 1e-3 @@ -8387,6 +8387,75 @@ def test_gpu_torch_deterministic_mode_alerts_before_an_order_varying_launch( assert _arena_tickets_zero() +def test_gpu_the_cluster_program_does_not_alert(torch_determinism): + """The shipped cluster row ``4096 -> 64`` FP32 at 2048 tokens of :data:`K1_DSM_RANGES_BY_SM` (152-SM ``sm_100a`` + parts: the cluster program at three splits, added in cluster-rank order) launches without an alert or a warning + under ``torch.use_deterministic_algorithms(True)``, and with ``warn_only=True``: the planning call, the call from + the cached key and the prepared launch run the cluster program (by kernel name), and every output meets the + deterministic package's FP32 budget against FP64 (:func:`_k1_check`). ``test_the_cluster_program_does_not_alert`` + checks the same calls on the CPU binding. A part of any other SM count or architecture skips.""" + device, arch = _atomic_device_or_skip() + sm_count = det_backend.device_facts(device).sm_count + if (arch, sm_count) != ("sm_100a", 152): + pytest.skip( + f"the cluster row is shipped for 152-SM sm_100a parts, not {arch} on {sm_count} SMs" + ) + T, N, M = 2048, 64, 4096 + dsm = cake_backend.k1_dsm_plan( + arch=arch, L=1, M=M, N=N, K=T, out_f32=True, sm_count=sm_count + ) + assert dsm == (64, 3) + _require_template(arch, cake_backend.dsm_k1_template(*dsm)) + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() + g = torch.Generator(device=device).manual_seed(SEED + 99) + G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) + # the FP64 reference before the mode is on (its cuBLAS GEMM would need CUBLAS_WORKSPACE_CONFIG under it) + ref = G.double().t() @ X.double() + out = torch.full((N, M), float("nan"), device=device, dtype=torch.float32) + got = [] + torch.use_deterministic_algorithms(True) + with _no_warnings(): + for _ in range(2): # the planning call, then the call from the cached key + out.fill_(float("nan")) + assert ( + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + is out + ) + got.append(out.clone()) + prepared = cake_backend.cake_dense_gemm_nondet_train_projection_wgrad_prepare( + G, X, out + ) + out.fill_(float("nan")) + assert prepared.launch() is out + got.append(out.clone()) + torch.cuda.synchronize() + ((_, entry),) = cake_backend._K1_ATOMIC_CACHE.items() + assert entry.plan.dsm and entry.arena is None + assert (entry.plan.block_n, entry.plan.splits) == dsm and not entry.order_varying + assert isinstance(prepared, cake_backend.CakeDenseGemmNondetProjectionLaunch) + assert prepared.plan.dsm and prepared.plan.reproducible + assert (prepared.plan.block_n, prepared.plan.splits) == dsm + # still under the mode: the calls ran the cluster program + _assert_ran( + lambda: cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out), + entry.module_name, + ) + _assert_ran(prepared.launch, prepared.module_name) + torch.use_deterministic_algorithms(True, warn_only=True) + with _no_warnings(): + out.fill_(float("nan")) + cake_backend.cake_dense_gemm_nondet_train_projection_wgrad(G, X, out) + got.append(out.clone()) + out.fill_(float("nan")) + prepared.launch() + got.append(out.clone()) + torch.cuda.synchronize() + torch.use_deterministic_algorithms(False) + for i, result in enumerate(got): + _k1_check(result, ref, f"cluster program call {i}") + cake_backend.cake_dense_gemm_nondet_train_convenience_cache_clear() + + def test_gpu_det_then_atomic_share_a_stream_not_a_workspace(k1_rule): """A deterministic call and a call of this package of one geometry on one stream: each package keeps its own arena (the deterministic package's slabs never reach this package's tickets).""" diff --git a/tests/experimental/test_cake_dense_gemm_nondet_train_names.py b/tests/experimental/test_cake_dense_gemm_nondet_train_names.py index 9ea2c6427b2..229bd00bf14 100644 --- a/tests/experimental/test_cake_dense_gemm_nondet_train_names.py +++ b/tests/experimental/test_cake_dense_gemm_nondet_train_names.py @@ -469,6 +469,7 @@ def test_the_former_package_is_gone(): def test_imports_are_lazy(): code = f""" import sys +import warnings import flashinfer loaded = sorted(m for m in sys.modules if "cake_dense_gemm_" in m) assert loaded == [], loaded @@ -481,8 +482,12 @@ def test_imports_are_lazy(): loaded = sorted(n for n in sys.modules if n.startswith("{BACKEND}")) assert loaded == [], loaded assert not [n for n in sys.modules if n.startswith("{DET_BACKEND}")] -m.cake_dense_gemm_nondet_train_router_wgrad +# the first access loads the backend: no ExperimentalWarning (a stable module), no other warning +with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + m.cake_dense_gemm_nondet_train_router_wgrad assert "{BACKEND}.cake_backend" in sys.modules +assert not caught, [(w.category.__name__, str(w.message)) for w in caught] print("ok") """ result = _run(code)