refactor(cake_dense_gemm): rename both packages to det and nondet (#8928)

.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, dense_projection_det_gemm, 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_nondet_train/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|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_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_nondet_train/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

benchmarks/bench_cake_dense_projection_det_gemm.py → benchmarks/bench_cake_dense_gemm_det_train.py

@@ -14,9 +14,9 @@
 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.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``,
@@ -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]
 """
@@ -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,14 @@ 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"]
     )
 

benchmarks/bench_cake_dense_projection_gemm.py → 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,64 +289,88 @@ 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(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: cake_backend.router_fp32_gemm(
-            inp["A"], inp["B"], out, deterministic=deterministic
+        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: cake_backend.projection_wgrad(
-            inp["G"], inp["X"], out, deterministic=deterministic
+        return lambda: det_backend.cake_dense_gemm_det_train_projection_wgrad(
+            inp["G"], inp["X"], out
         )
-    return lambda: cake_backend.dense_projection_gemm(
-        inp["A"], inp["B"], out, deterministic=deterministic
+    return lambda: det_backend.cake_dense_gemm_det_train_projection(
+        inp["A"], inp["B"], out
+    )
+
+
+def convenience_atomic(spec, inp, out):
+    """This package's convenience entry point of the row (cached plan, arena
+    workspace)."""
+    if spec["family"] == "router":
+        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.cake_dense_gemm_nondet_train_projection_wgrad(
+            inp["G"], inp["X"], out
+        )
+    return lambda: cake_backend.cake_dense_gemm_nondet_train_projection(
+        inp["A"], inp["B"], out
     )
 
 
 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,
+        ),
     )
 
 
 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)
@@ -363,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":
@@ -1128,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)
@@ -1268,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
 
@@ -1346,13 +1374,16 @@ 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)):
-            cake_backend.clear_convenience_caches()
-            fn = convenience(spec, inp, outs[mode], deterministic=deterministic)
+        for mode, convenience in (
+            ("det", convenience_det),
+            ("atomic", convenience_atomic),
+        ):
+            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"] = {