refactor(cake_kda): give single-rank and context-parallel KDA training deterministic names

.github/labeler.yml

@@ -83,6 +83,8 @@
           - "flashinfer/experimental/cake_gdn_det_train_kcp/**"
           - "flashinfer/experimental/cake_gdn_train/**"
           - "flashinfer/experimental/cake_kda_train/**"
+          - "flashinfer/experimental/cake_kda_det_train/**"
+          - "flashinfer/experimental/cake_kda_det_train_kcp/**"
           - "flashinfer/experimental/kda_prefill_wrapper/**"
           - "flashinfer/experimental/b12x/sequence/_shared/kda_math.py"
           - "tests/experimental/**/test_*gdn*"

.pre-commit-config.yaml

@@ -38,13 +38,13 @@ files: ^(?!(?:flashinfer/moe_ep/kernel_src/blackwell_bf16_rank_major/src/|csrc/c
 # The SM120 Cake MegaMoE device source parts, instance wrappers and geometry table are generated
 # 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_dsa_indexer, cake_mok_bf16, cake_moe_grouped_det_gemm,
+# training and mHC training backends, the causal conv1d, ulysses and fused QK RoPE append (BF16 / FP8) backends, cake_dsa_train, cake_kda_det_train, cake_kda_det_train_kcp, cake_dsa_indexer, cake_mok_bf16, cake_moe_grouped_det_gemm,
 # cake_lm_head_loss, dense_projection_det_gemm, cake_stepfun_moe, nvfp4_attention, deepgemm_fp8_gemm, deepgemm_batched_gemm, 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.
 # 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/|cake_msa_icp_indexer/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:cake_rmsnorm_det_train_backend/csrc/|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/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_dsa_train|cake_kda_train|cake_dsa_indexer|dense_projection_det_gemm|cake_stepfun_moe|deepgemm_fp8_gemm|deepgemm_batched_gemm)/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/|cake_msa_icp_indexer/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:cake_rmsnorm_det_train_backend/csrc/|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/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_dsa_train|cake_kda_det_train|cake_kda_det_train_kcp|cake_dsa_indexer|dense_projection_det_gemm|cake_stepfun_moe|deepgemm_fp8_gemm|deepgemm_batched_gemm)/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_kda_train.py → benchmarks/bench_cake_kda_det_train.py

@@ -14,7 +14,7 @@
 limitations under the License.
 
 Benchmark of the Cake packed KDA chunked training kernels
-(``flashinfer.experimental.cake_kda_train``) on three boundaries of a
+(``flashinfer.experimental.cake_kda_det_train``) on three boundaries of a
 training step, against flash-linear-attention's ``chunk_kda`` when it is
 installed:
 
@@ -32,8 +32,8 @@
 
 Example::
 
-    python benchmarks/bench_cake_kda_train.py --seqlen 16384 65536 --docs 3 --heads 4 64
-    python benchmarks/bench_cake_kda_train.py --seqlen 261393 --docs 4 --heads 4 --json out.json
+    python benchmarks/bench_cake_kda_det_train.py --seqlen 16384 65536 --docs 3 --heads 4 64
+    python benchmarks/bench_cake_kda_det_train.py --seqlen 261393 --docs 4 --heads 4 --json out.json
 """
 
 from __future__ import annotations
@@ -50,9 +50,9 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 
-from flashinfer.experimental.cake_kda_train import cake_backend, cake_jit  # noqa: E402
+from flashinfer.experimental.cake_kda_det_train import cake_backend, cake_jit  # noqa: E402
 from flashinfer.testing import bench_gpu_time  # noqa: E402
-from tests.test_helpers.cake_kda_train_reference import make_inputs  # noqa: E402
+from tests.test_helpers.cake_kda_det_train_reference import make_inputs  # noqa: E402
 
 BOUNDARIES = ("fwd", "bwd", "ckpt")
 
@@ -81,23 +81,23 @@ def _time(fn: Callable[[], object]) -> float:
 
 
 def cake_boundaries(
-    inputs: dict, ws: cake_backend.KdaWorkspace
+    inputs: dict, ws: cake_backend.KdaDetWorkspace
 ) -> dict[str, Callable[[], None]]:
     q, k, v, g, beta, cu, dout = (
         inputs[n] for n in ("q", "k", "v", "g", "beta", "cu_seqlens", "dout")
     )
 
     def fwd() -> None:
-        cake_backend.chunk_kda_fwd(
+        cake_backend.chunk_kda_det_fwd(
             q, k, v, g, beta, cu, save_for_backward=True, workspace=ws
         )
 
-    _, saved = cake_backend.chunk_kda_fwd(
+    _, saved = cake_backend.chunk_kda_det_fwd(
         q, k, v, g, beta, cu, save_for_backward=True, workspace=ws
     )
 
     def bwd() -> None:
-        cake_backend.chunk_kda_bwd(
+        cake_backend.chunk_kda_det_bwd(
             q, k, v, g, beta, cu, dout, saved=saved, workspace=ws
         )
 
@@ -108,7 +108,7 @@ def bwd() -> None:
 
     def ckpt() -> None:
         o = checkpoint(
-            lambda *a: cake_backend.chunk_kda(*a, cu, workspace=ws),
+            lambda *a: cake_backend.chunk_kda_det(*a, cu, workspace=ws),
             *leaves,
             use_reentrant=False,
         )
@@ -171,7 +171,7 @@ def run_row(
 ) -> dict:
     lengths = split_docs(seqlen, docs)
     inputs = make_inputs(lengths, heads, device)
-    ws = cake_backend.KdaWorkspace.for_problem(seqlen, docs, heads, device)
+    ws = cake_backend.KdaDetWorkspace.for_problem(seqlen, docs, heads, device)
     ws.load()
     row = {
         "seqlen": seqlen,

benchmarks/bench_cake_kda_train_kcp.py → benchmarks/bench_cake_kda_det_train_kcp.py

@@ -17,7 +17,7 @@
 process per rank (``torchrun --nproc_per_node N``), every rank timing the
 module boundary of its fragments -- the KCP causal convolution
 (``flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_kcp``) followed by
-the KCP KDA step (``flashinfer.kda_train.chunk_kda_kcp``), forward and
+the KCP KDA step (``flashinfer.cake_kda_det_train_kcp.chunk_kda_det_kcp``), forward and
 backward -- with CUDA events, and the slowest rank's time counting for the
 iteration.
 
@@ -43,7 +43,7 @@
 
 Example::
 
-    torchrun --nproc_per_node 4 benchmarks/bench_cake_kda_train_kcp.py \\
+    torchrun --nproc_per_node 4 benchmarks/bench_cake_kda_det_train_kcp.py \\
         --seqlen 65536 --docs 3 --heads 16 --channels 4096 --groups 3 --samples 8 --json kcp.json
 """
 
@@ -67,15 +67,15 @@
     cake_causal_conv1d,
     cake_causal_conv1d_kcp,
 )
-from flashinfer.kda_train import (  # noqa: E402
-    build_kcp_context,
-    chunk_kda,
-    chunk_kda_kcp,
-    kcp_comm_bytes,
-    kcp_partition,
-    kcp_saved_bytes,
-    kcp_workspace_bytes,
+from flashinfer.cake_kda_det_train_kcp import (  # noqa: E402
+    build_context_det_kcp,
+    chunk_kda_det_kcp,
+    kda_det_comm_bytes_kcp,
+    kda_det_saved_bytes_kcp,
+    kda_det_workspace_bytes_kcp,
+    partition_det_kcp,
 )
+from flashinfer.kda_train import chunk_kda  # noqa: E402
 from tests.test_helpers.cake_kda_train_reference import make_inputs  # noqa: E402
 
 ARMS = ("cake", "control", "local", "fla")
@@ -137,9 +137,9 @@ def make_problem(args, rank: int, world: int, device: torch.device) -> dict[str,
         torch.bfloat16
     )
     weight = (torch.randn(args.channels, WIDTH, generator=gen) * 0.3).to(device)
-    part = kcp_partition(cu_cpu, world, args.layout)
+    part = partition_det_kcp(cu_cpu, world, args.layout)
     group = dist.group.WORLD
-    ctx = build_kcp_context(part, group, rank, device)
+    ctx = build_context_det_kcp(part, group, rank, device)
     fragments = part.fragments[rank]
     local = {
         name: _shard(inputs[name], cu, fragments).to(device)
@@ -171,7 +171,7 @@ def _step_kcp(problem: dict[str, Any]) -> None:
         local[n].detach().requires_grad_(True) for n in ("q", "k", "v", "g", "beta")
     )
     y = cake_causal_conv1d_kcp(x, w, ctx)
-    out = chunk_kda_kcp(q, k, v, g, beta, ctx)
+    out = chunk_kda_det_kcp(q, k, v, g, beta, ctx)
     torch.autograd.backward([y, out], [local["dy"], local["dout"]])
 
 
@@ -292,9 +292,11 @@ def run(args, rank: int, world: int, device: torch.device) -> dict[str, Any]:
             "seed": args.seed,
         },
         "bytes": {
-            "workspace": kcp_workspace_bytes(ctx, H=args.heads, C=args.channels),
-            "saved": kcp_saved_bytes(ctx, H=args.heads, C=args.channels),
-            "comm": kcp_comm_bytes(ctx, H=args.heads, C=args.channels),
+            "workspace": kda_det_workspace_bytes_kcp(
+                ctx, H=args.heads, C=args.channels
+            ),
+            "saved": kda_det_saved_bytes_kcp(ctx, H=args.heads, C=args.channels),
+            "comm": kda_det_comm_bytes_kcp(ctx, H=args.heads, C=args.channels),
         },
         "arms": rows,
         "control_ratio": control,