yyihuang · Nov 6, 2026 · Nov 6, 2026
@@ -280,6 +280,8 @@ def __init__(self, kind: str) -> None:
             self._comm = self.api.gdn_comm_bytes_kcp
             self.handle_name = "RecomputeHandleKcp"
             self.workspace_name = "WorkspaceKcp"
+        # the host module of the route's package: its ``LAST_CALL`` records and ``graph_cache_clear`` live there
+        self.kcp = importlib.import_module(f"{self.pkg.__name__}.cake_kcp")
 
     def partition(
         self, cu: torch.Tensor, world: int, layout: str, cut_points: Any
@@ -341,9 +343,7 @@ def workspace(self, ctx: Any, heads: tuple[int, int], device: torch.device) -> A
 
     def graph_cache_clear(self) -> None:
         """Release this route's captured graphs (before the next configuration and the process group's end)."""
-        fn = getattr(self.pkg, "graph_cache_clear", None) or getattr(
-            getattr(self.pkg, "cake_kcp", None), "graph_cache_clear", None
-        )
+        fn = getattr(self.kcp, "graph_cache_clear", None)
         if fn is not None:
             fn()
 
@@ -373,20 +373,19 @@ def bytes(
         }
 
     def modules(self) -> dict[str, Any]:
-        kcp = getattr(self.pkg, "cake_kcp", None)
         return {
             "facade": self.api.__name__,
             "facade_file": self.api.__file__,
             "package": self.pkg.__name__,
             "package_file": self.pkg.__file__,
-            "cake_kcp_file": getattr(kcp, "__file__", None),
+            "cake_kcp_file": self.kcp.__file__,
         }
 
     def clear_last_call(self) -> None:
-        self.pkg.LAST_CALL.clear()
+        self.kcp.LAST_CALL.clear()
 
     def last_call(self, kind: str) -> Optional[dict[str, Any]]:
-        rec = self.pkg.LAST_CALL.get(kind)
+        rec = self.kcp.LAST_CALL.get(kind)
         if not rec:
             return None
         out = {
@@ -786,7 +785,7 @@ def _probe(
     torch.cuda.synchronize(device)
     mem["after_bytes"] = torch.cuda.memory_allocated(device) - base
     mem["step_peak_bytes"] = max(mem["fwd_peak_bytes"], mem.get("bwd_peak_bytes", 0))
-    ran = sorted(kind for kind, route in routes.items() if route.pkg.LAST_CALL)
+    ran = sorted(kind for kind, route in routes.items() if route.kcp.LAST_CALL)
     if sota:
         made = {key: lib.calls[key] - before[key] for key in lib.calls}
         need = (SOTA_GDN_CALL[arm],) + (("conv",) if len(grads_out) == 2 else ())

@@ -85,15 +85,19 @@ non-deterministic one).
 
 No launch of this package is unordered, so PyTorch's deterministic algorithms
 (`torch.use_deterministic_algorithms(True)`, with or without `warn_only`)
-change nothing: no entry raises PyTorch's alert.  The deterministic packages,
-their names, defaults and aliases are unchanged.
+change nothing: no entry raises PyTorch's alert.
 
-| name | what it is now |
+The deterministic packages are named by their det names only:
+
+| name | what it is |
 |---|---|
-| `flashinfer.experimental.cake_gdn_det_train` and its aliases `chunk_gated_delta_rule`, `ChunkGatedDeltaRuleFunction`, `gdn_chunk_forward`, `gdn_chunk_backward` | unchanged: the deterministic one-rank package |
-| `flashinfer.experimental.cake_gdn_train` (the KCP package of FlashInfer v0.7.3rc1) | unchanged: a legacy alias of the deterministic KCP package `flashinfer.experimental.cake_gdn_det_train_kcp` |
-| `flashinfer.gdn_train`, `flashinfer.cake_gdn_det_train_kcp` | unchanged: deterministic KCP entry points |
-| `flashinfer.experimental.cake.cake_gdn_train` | new: this non-deterministic one-rank package |
+| `flashinfer.experimental.cake_gdn_det_train` (`chunk_gated_delta_rule_det`, `ChunkGatedDeltaRuleDetFunction`, `gdn_chunk_det_forward`, `gdn_chunk_det_backward`) | the deterministic one-rank package |
+| `flashinfer.experimental.cake_gdn_det_train_kcp`, `flashinfer.cake_gdn_det_train_kcp` | the deterministic KCP package and its entry points |
+| `flashinfer.experimental.cake.cake_gdn_train` | this non-deterministic one-rank package |
+
+The no-det spellings `chunk_gated_delta_rule`, `ChunkGatedDeltaRuleFunction`,
+`gdn_chunk_forward` and `gdn_chunk_backward` name this package's
+non-deterministic entries; a deterministic call uses the det names above.
 
 ## Chain
 

@@ -76,12 +76,10 @@ Every name of this package is the deterministic twin's name without `det`:
 | `ChunkGDNDetFunctionKcp` | `ChunkGDNFunctionKcp` |
 | `packed_qkv_views_det_kcp`, `PackedQKVSplitDetKcp` | `packed_qkv_views_kcp`, `PackedQKVSplitKcp` (package level) |
 
-The legacy names of the deterministic package keep meaning the deterministic
-objects: `flashinfer.gdn_train.chunk_gdn_kcp`, `KcpContext`, ... and the
-legacy package path `flashinfer.experimental.cake_gdn_train` (with its
-`chunk_gdn_kcp`) are aliases of `cake_gdn_det_train_kcp` objects, never of
-this package's.  The name `chunk_gdn_kcp` of this package is a different
-object from the legacy alias of the same spelling.
+The deterministic package is named by its det names (left column) only.
+This package binds none of the deterministic package's former no-det names
+(such as `kcp_partition` or `KcpContext`), and `chunk_gdn_kcp` is this
+package's own non-deterministic entry, never a deterministic object.
 
 The generated libraries are named `cake_gdn_train_<20 hex>` for the
 per-chunk programs and `cake_gdn_train_<20 hex>_kcp` for the

@@ -91,7 +91,8 @@
                              ``K K^T`` Gram, no block inverse): every other output bit for bit ``fwd_prep_l2n``'s,
                              ``akk`` read, not written; grid (NC, H_qk), two resident CTAs per SM as #1
 
-Provenance of the rules: ``flashinfer/experimental/cake_kda_train/cake_kcp_jit.py`` (records_of_program,
+Provenance of the rules: ``flashinfer/experimental/cake_kda_train/cake_kcp_jit.py`` (now
+``flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py``; records_of_program,
 select_record, launch_grid; sha256 c405e7dcd36c9cbe0ff01e3e3194fd760491504655bda28a646c2067a23b4c90) at
 FlashInfer 13d7fcf4041635f5b7f573bb0dbab23570ae1034, adapted to the GDN vocabulary; the loader glue moved to
 :mod:`.cake_backend`.

@@ -19,7 +19,8 @@
 :class:`SavedKcp`, the autograd Function and the entry points
 ``chunk_gdn_kcp`` / ``chunk_gdn_fwd_kcp`` / ``chunk_gdn_bwd_kcp``.
 
-Provenance: adapted from ``flashinfer/experimental/cake_kda_train/cake_kcp.py``
+Provenance: adapted from ``flashinfer/experimental/cake_kda_train/cake_kcp.py`` (now
+``flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py``)
 at FlashInfer ``13d7fcf4041635f5b7f573bb0dbab23570ae1034`` (the KDA KCP host
 module).  Verbatim parts: the compact pools and chains (``Indexed``,
 ``Chains``), the exchange tables (``_Tables``), the per-call state

@@ -13,11 +13,12 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 
-Provenance: package-local copy of ``flashinfer/experimental/cake_kda_train/cake_kcp_plan.py`` at FlashInfer
+Provenance: package-local copy of ``flashinfer/experimental/cake_kda_train/cake_kcp_plan.py`` (now
+``flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_plan.py``) at FlashInfer
 commit 13d7fcf4041635f5b7f573bb0dbab23570ae1034 (sha256 0f2d8b06bbea49df37b043fba3026a06799f0a0cba09d697e490d02bb747dfbf),
 adapted for the GDN KCP route: the head count of every per-head state is the value-head count ``HV``
 (``H := HV``), the plan rule is a keyword-only argument (no environment variable), and the byte queries live
-in ``cake_bytes`` (GDN saved set and exchange formulas).  ``cake_gdn_train_kcp`` never imports ``cake_kda_train``;
+in ``cake_bytes`` (GDN saved set and exchange formulas).  ``cake_gdn_train_kcp`` never imports a KDA package;
 a CPU test asserts field-identical partitions and contexts on the same inputs.
 
 Metadata layer of the context-parallel (KCP) route of the Cake packed GDN

@@ -1820,15 +1820,11 @@ def test_the_package_exports_the_det_free_names_only() -> None:
     for name in DET_NAMES:
         for module in (cake_gdn_train, cake_chain):
             assert not hasattr(module, name) and name not in module.__all__, name
-    # the deterministic package keeps its names and its aliases, distinct objects from this package's entries
-    assert (
-        cake_gdn_det_train.chunk_gated_delta_rule
-        is cake_gdn_det_train.chunk_gated_delta_rule_det
-    )
-    assert (
-        cake_gdn_train.chunk_gated_delta_rule
-        is not cake_gdn_det_train.chunk_gated_delta_rule
-    )
+    # the deterministic package's entries (its det names) are distinct objects from this package's entries
+    for name, det_name in zip(ENTRY_NAMES, DET_NAMES, strict=True):
+        assert getattr(cake_gdn_train, name) is not getattr(
+            cake_gdn_det_train, det_name
+        ), name
     assert cake_jit.PACKAGE == "flashinfer.experimental.cake.cake_gdn_train"
     assert cake_jit.LIBRARY_PREFIX == "cake_gdn_train"
 

@@ -15,8 +15,8 @@
 
 CPU tests of the GDN KCP metadata layer of the NON-DETERMINISTIC package ``flashinfer.experimental.cake_gdn_train_kcp``
 (the twin of ``cake_gdn_det_train_kcp`` with the same contract and accuracy budgets, not its bitwise reproducibility):
-the package-local plan module against ``cake_kda_train.cake_kcp_plan`` (field-identical partitions and contexts on the
-step-0 rows with ``H := HV``; the KDA package is imported HERE only, never by the package), the per-rank crossing
+the package-local plan module against ``cake_kda_det_train_kcp.cake_kcp_plan`` (field-identical partitions and contexts on
+the step-0 rows with ``H := HV``; the KDA package is imported HERE only, never by the package), the per-rank crossing
 counts and exchange bytes against the rows of record, the table builders (pool addressing, segment lists, chains),
 head ownership incl. the padded ``(HV, W) = (3, 2)`` case, the byte queries, the registry seed, an emulated multi-rank
 row (W ranks on threads), and the package's host orchestration on the fake backend
@@ -381,14 +381,14 @@ def _sha(obj) -> str:
 @ROWS
 @pytest.mark.parametrize("path", ROW_FILES, ids=ROW_IDS)
 def test_plan_parity_with_kda_on_rows(path, monkeypatch) -> None:
-    """Result 1: ``cake_gdn_train_kcp.cake_kcp_plan`` is field-identical to ``cake_kda_train.cake_kcp_plan`` on every
-    step-0 row (partition, every rank's context and head plan at ``H := HV``) and reproduces the row's recorded
+    """Result 1: ``cake_gdn_train_kcp.cake_kcp_plan`` is field-identical to ``cake_kda_det_train_kcp.cake_kcp_plan`` on
+    every step-0 row (partition, every rank's context and head plan at ``H := HV``) and reproduces the row's recorded
     fragment table.  Result 2: ``n_in`` / ``n_out`` / ``max_chain_fragments`` and the exchange bytes of
     ``gdn_comm_bytes_kcp`` equal the row JSON.  Both modules plan under the per-rank rule they share,
     ``"wave_fill1"`` (the GDN default is the row rule, the KDA default its own tie rule; each is selected explicitly
     here: the keyword-only ``p_rule`` on the GDN side, the private test hook on the KDA side)."""
     kp = importlib.import_module(
-        "flashinfer.experimental.cake_kda_train.cake_kcp_plan"
+        "flashinfer.experimental.cake_kda_det_train_kcp.cake_kcp_plan"
     )  # the test only
     # the KDA module accepts its private test hook only after its opt-in (enable_test_knobs)
     monkeypatch.setattr(kp, "_TEST_KNOBS_ENABLED", True)
@@ -398,7 +398,7 @@ def test_plan_parity_with_kda_on_rows(path, monkeypatch) -> None:
     W = int(row["cp"])
     args = _partition_inputs(row)
     part = gp.partition_kcp(**args)
-    kpart = kp.kcp_partition(**args)
+    kpart = kp.partition_det_kcp(**args)
     assert part.to_dict() == kpart.to_dict(), "partition differs from the KDA module"
     assert _plain(part.to_dict()["fragments"]) == row["fragments"], (
         "partition differs from the row of record"
@@ -419,7 +419,7 @@ def test_plan_parity_with_kda_on_rows(path, monkeypatch) -> None:
         ctx = gp.build_context_kcp(
             part, None, r, CPU, 4, sm_count=SM_COUNT, p_rule="wave_fill1"
         )  # the per-rank rule both modules share; the GDN default "row_cost" is tested on the rows below
-        kctx = kp.build_kcp_context(
+        kctx = kp.build_context_det_kcp(
             kpart,
             None,
             r,
@@ -2037,8 +2037,8 @@ def test_registry_seed() -> None:
 
 
 def test_package_imports_no_kda_and_no_environ() -> None:
-    """The package never imports ``cake_kda_train`` and reads no environment variable; the plan
-    module imports torch, torch.distributed and dataclasses only."""
+    """The package never imports a KDA package (no module path or name containing ``cake_kda``) and reads no
+    environment variable; the plan module imports torch, torch.distributed and dataclasses only."""
     for path in sorted(PACKAGE_DIR.glob("*.py")):
         tree = ast.parse(path.read_text())
         for node in ast.walk(tree):
@@ -2048,8 +2048,8 @@ def test_package_imports_no_kda_and_no_environ() -> None:
                     if isinstance(node, ast.ImportFrom) and node.module
                     else []
                 )
-                assert not any("cake_kda_train" in n for n in names), (
-                    f"{path.name} imports cake_kda_train"
+                assert not any("cake_kda" in n for n in names), (
+                    f"{path.name} imports a KDA package"
                 )
                 assert "os" not in [
                     a.name for a in node.names if isinstance(node, ast.Import)