refactor(cake_gdn_kcp): the GDN KCP facade runs the non-deterministic package only

benchmarks/bench_cake_gdn_train_kcp.py

@@ -20,8 +20,8 @@
 
 Arms (``--arms``; every one on the same packed batch, the same partition and the same input storage):
 
-* ``det``: the deterministic route, the objects of ``flashinfer.cake_gdn_det_train_kcp`` (what
-  ``deterministic=True`` dispatches to), with its own partition and context;
+* ``det``: the deterministic route, the entries of ``flashinfer.cake_gdn_det_train_kcp`` called directly, with its
+  own partition and context;
 * ``nondet``: ``flashinfer.cake_gdn_train_kcp`` with ``deterministic=False`` (the package
   ``flashinfer.experimental.cake_gdn_train_kcp``), default options;
 * ``det_graph`` / ``nondet_graph``: ``det`` / ``nondet`` with a caller-owned workspace (``WorkspaceDetKcp`` /
@@ -246,9 +246,9 @@ def _query(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
 
 
 class Route:
-    """The entry points of one numerical route.  ``det``: the objects of ``flashinfer.cake_gdn_det_train_kcp`` (what
-    ``deterministic=True`` dispatches to); ``nondet``: ``flashinfer.cake_gdn_train_kcp`` with
-    ``deterministic=False`` (the package ``flashinfer.experimental.cake_gdn_train_kcp``)."""
+    """The entry points of one numerical route.  ``det``: the entries of ``flashinfer.cake_gdn_det_train_kcp`` (called
+    directly; they take no mode keyword); ``nondet``: ``flashinfer.cake_gdn_train_kcp`` with ``deterministic=False``
+    (the package ``flashinfer.experimental.cake_gdn_train_kcp``)."""
 
     def __init__(self, kind: str) -> None:
         self.kind = kind
@@ -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 ())
@@ -1827,7 +1826,8 @@ def main() -> int:
     args = parse_args()
     if torch.are_deterministic_algorithms_enabled():
         raise SystemExit(
-            "torch deterministic algorithms are enabled: deterministic=False would run the deterministic route"
+            "torch deterministic algorithms are enabled: this bench measures with them off (they can change or "
+            "refuse the kernels of other arms); neither GDN route depends on them"
         )
     rank, world, device = _init()
     arms = args.arm_list

flashinfer/experimental/cake/cake_gdn_train/README.md

@@ -24,6 +24,26 @@ libraries are loaded prebuilt from `aot/` when a builder compiled them and
 JIT-built by FlashInfer from the sources and flags `build.json` records
 otherwise.  The package reads no environment variable.
 
+## Mode and dispatch rules
+
+| call | `deterministic=False` | any other value (the default, `True`, `None`, `0`, `1`, `"False"`) | under `torch.use_deterministic_algorithms(True)` (strict or `warn_only=True`) |
+|---|---|---|---|
+| this package (`chunk_gated_delta_rule`, `ChunkGatedDeltaRuleFunction.apply`, `gdn_chunk_forward`, `gdn_chunk_backward`) | runs | `ValueError` naming the deterministic entry, before any work | runs; no alert (every launch is fixed-order) |
+| `flashinfer.experimental.cake_gdn_det_train` (`chunk_gated_delta_rule_det`, ...) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged: every launch is fixed-order |
+
+* This package is the non-deterministic route only.  Nothing routes a call
+  to the deterministic package and nothing falls back to it: a caller who
+  wants the deterministic route calls `flashinfer.experimental.cake_gdn_det_train`.
+* Every call states the mode with `deterministic=False`.
+* PyTorch's deterministic algorithms do not change the route.  A launch whose
+  result can vary run to run raises PyTorch's alert immediately before it runs
+  (`RuntimeError`; a `UserWarning` and then the launch under
+  `warn_only=True`).  Every launch of this version is fixed-order, so no call
+  raises it.
+* The context-parallel packages follow the same rules
+  (`flashinfer.cake_gdn_train_kcp` / `flashinfer.experimental.cake_gdn_train_kcp`
+  and `flashinfer.cake_gdn_det_train_kcp`).
+
 ## Opt-in
 
 Import the entry point explicitly; nothing in `flashinfer` imports this
@@ -79,21 +99,24 @@ package implicitly.  Every entry of this package (`chunk_gated_delta_rule`,
 non-deterministic route fails instead of running it.  The autograd function
 records the forward's mode, and `gdn_chunk_backward` refuses (`ValueError`) a
 forward tape of the other mode, such as one of the deterministic package.
-The KCP facade `flashinfer.cake_gdn_train_kcp` is the entry that dispatches on
-`deterministic` (`True`: the deterministic package; `False`: the
-non-deterministic one).
+The KCP facade `flashinfer.cake_gdn_train_kcp` follows the same rule: it runs
+the non-deterministic KCP package with `deterministic=False` only.
 
 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
 

flashinfer/experimental/cake/cake_gdn_train/cake_chain.py

@@ -167,15 +167,24 @@
 _MODE_CHECKED = object()
 
 
+# Each entry's deterministic counterpart in ``flashinfer.experimental.cake_gdn_det_train`` (named in the refusal).
+_DET_ENTRIES = {
+    "chunk_gated_delta_rule": "chunk_gated_delta_rule_det",
+    "ChunkGatedDeltaRuleFunction": "ChunkGatedDeltaRuleDetFunction",
+    "gdn_chunk_forward": "gdn_chunk_det_forward",
+    "gdn_chunk_backward": "gdn_chunk_det_backward",
+}
+
+
 def _require_nondeterministic(deterministic: Any, entry: str) -> None:
     """The mode contract of every entry: the package is the non-deterministic route (its results match the
     deterministic package within that package's accuracy budgets, not bitwise), so an entry takes
-    ``deterministic=False`` and raises ``ValueError`` for any other value (the deterministic package is
-    ``flashinfer.experimental.cake_gdn_det_train``)."""
+    ``deterministic=False`` and raises ``ValueError`` for any other value, naming the entry's deterministic
+    counterpart in ``flashinfer.experimental.cake_gdn_det_train``."""
     if deterministic is not False:
         raise ValueError(f"{entry} is the non-deterministic route and takes deterministic=False (got "
-                         f"{deterministic!r}); flashinfer.experimental.cake_gdn_det_train is the deterministic "
-                         "package")  # fmt: skip
+                         f"{deterministic!r}): pass deterministic=False for the non-deterministic path; the "
+                         f"deterministic entry is flashinfer.experimental.cake_gdn_det_train.{_DET_ENTRIES[entry]}")  # fmt: skip
 
 
 # Kernel parameter names of every stage in call-site order; a delivered record whose parameters are these names in

flashinfer/experimental/cake_gdn_train_kcp/README.md

@@ -28,6 +28,31 @@ the twin) and reads no environment variable.
 
 The numbers of this file (shapes, bytes, limits) are the contract.
 
+## Mode and dispatch rules
+
+| call | `deterministic=False` | any other value (the default, `True`, `None`, `0`, `1`, `"False"`) | under `torch.use_deterministic_algorithms(True)` (strict or `warn_only=True`) |
+|---|---|---|---|
+| `flashinfer.cake_gdn_train_kcp.*` (every entry, the byte queries and `is_gdn_train_kcp_supported` included) | runs this package | `ValueError` naming the deterministic entry, before any work | runs this package; no reroute, no mode or fallback warning (each entry's one-time `ExperimentalWarning` is unrelated to the mode) |
+| this package's compute entries (`chunk_gdn_kcp`, `chunk_gdn_fwd_kcp`, `chunk_gdn_bwd_kcp`) | runs | `ValueError` naming the deterministic entry | runs; no alert (every launch is fixed-order) |
+| `flashinfer.cake_gdn_det_train_kcp.*_det_*` (the deterministic route) | no `deterministic` keyword (`TypeError`) | no `deterministic` keyword | unchanged: every launch is fixed-order |
+
+* This package and its facade `flashinfer.cake_gdn_train_kcp` are the
+  non-deterministic route only.  Nothing routes a call to the deterministic
+  package and nothing falls back to it: a caller who wants the deterministic
+  route calls `flashinfer.cake_gdn_det_train_kcp` (`partition_det_kcp`,
+  `build_context_det_kcp`, `chunk_gdn_det_kcp`, ...).
+* Every call states the mode: pass `deterministic=False` on every facade
+  call, the partition, the context and the byte queries included.
+* PyTorch's deterministic algorithms do not change the route.  A launch whose
+  result can vary run to run raises PyTorch's alert immediately before it runs
+  (`RuntimeError`; a `UserWarning` and then the launch under
+  `warn_only=True`).  Every launch of this version is fixed-order, so no call
+  raises it.
+* The partition / context / byte-query functions of this package
+  (`partition_kcp`, `build_context_kcp`, `gdn_*_bytes_kcp`) take no mode
+  keyword; their facade copies take `deterministic=False` like every facade
+  entry.
+
 ## Public entry points (`flashinfer/cake_gdn_train_kcp.py`)
 
 ```python
@@ -38,15 +63,16 @@ from flashinfer.cake_gdn_train_kcp import (
     is_gdn_train_kcp_supported,
 )
 
-part = partition_kcp(cu_seqlens_global_cpu, world_size, layout="zigzag_per_document")   # once per batch layout
-ctx = build_context_kcp(part, group, rank, device, conv_kernel_size=4)                 # tables and collectives of this rank
+nd = dict(deterministic=False)                                                         # required on every call
+part = partition_kcp(cu_seqlens_global_cpu, world_size, layout="zigzag_per_document", **nd)  # once per batch layout
+ctx = build_context_kcp(part, group, rank, device, conv_kernel_size=4, **nd)          # tables and collectives of this rank
 
-out = chunk_gdn_kcp(q, k, v, g, beta, ctx)                                             # autograd
-out, saved = chunk_gdn_fwd_kcp(q, k, v, g, beta, ctx, save_for_backward=True)
-dq, dk, dv, dbeta, dg = chunk_gdn_bwd_kcp(q, k, v, g, beta, ctx, dout, saved=saved)
-dq, dk, dv, dbeta, dg = chunk_gdn_bwd_kcp(q, k, v, g, beta, ctx, dout)                 # saved=None: checkpoint replay
+out = chunk_gdn_kcp(q, k, v, g, beta, ctx, **nd)                                       # autograd
+out, saved = chunk_gdn_fwd_kcp(q, k, v, g, beta, ctx, save_for_backward=True, **nd)
+dq, dk, dv, dbeta, dg = chunk_gdn_bwd_kcp(q, k, v, g, beta, ctx, dout, saved=saved, **nd)
+dq, dk, dv, dbeta, dg = chunk_gdn_bwd_kcp(q, k, v, g, beta, ctx, dout, **nd)           # saved=None: checkpoint replay
 
-ws = gdn_workspace_bytes_kcp(ctx, H_qk=16, HV=48)                                      # per rank, pure CPU: a fresh workspace's first call
+ws = gdn_workspace_bytes_kcp(ctx, H_qk=16, HV=48, **nd)                                # per rank, pure CPU: a fresh workspace's first call
 ```
 
 These thirteen names are the whole public surface (the `__all__` of
@@ -76,12 +102,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
@@ -100,15 +124,11 @@ The packed GDN training routes come in four combinations:
 | deterministic | `flashinfer.experimental.cake_gdn_det_train` | `flashinfer.cake_gdn_det_train_kcp` -> `flashinfer.experimental.cake_gdn_det_train_kcp` |
 | non-deterministic | `flashinfer.experimental.cake.cake_gdn_train` | `flashinfer.cake_gdn_train_kcp` with `deterministic=False` -> this package |
 
-* The facade `flashinfer.cake_gdn_train_kcp` takes `deterministic: bool = True`
-  on every entry: `True` runs the deterministic package's objects, `False`
-  runs this package.  With `torch.use_deterministic_algorithms(True)`,
-  `deterministic=False` runs the deterministic package and warns once per
-  process.
-* This package's compute entries (`chunk_gdn_kcp`, `chunk_gdn_fwd_kcp`,
-  `chunk_gdn_bwd_kcp`) take `deterministic: bool = True` as well and raise
-  `ValueError` unless it is `False`: the package holds no deterministic
-  program and never falls back.
+* The facade `flashinfer.cake_gdn_train_kcp` and this package's compute
+  entries take `deterministic: bool = True` on every entry and raise
+  `ValueError` unless it is `False` (see "Mode and dispatch rules"): the
+  package holds no deterministic program, never routes to the deterministic
+  package and never falls back.
 * Faster schedule: the default segment plan of this package (`fill_small_guard`)
   walks fewer intra-rank segments than the deterministic package's plan where
   segment merges cost more than they save: at `HV <= 12`, when no fragment
@@ -140,8 +160,7 @@ The packed GDN training routes come in four combinations:
 * The mode is not stored in the context, the partition or the options: both
   routes build the same plan objects.  The autograd node records its
   forward's mode and its backward refuses another; `chunk_gdn_bwd_kcp`
-  refuses a saved set that is not this package's (`SavedKcp`), and the
-  facade's deterministic backward refuses this package's saved set.
+  refuses a saved set that is not this package's (`SavedKcp`).
 * Accuracy: every output stays within the deterministic route's budgets
   against the FP64 reference (`rel-RMS <= 1e-2`, `max-abs <= 2e-2` of
   `max|ref|`; the per-tensor KCP budgets of the deterministic tests).  The
@@ -350,7 +369,7 @@ a `RecomputeHandleKcp` to the checkpointed function and to `chunk_gdn_kcp`:
 ```python
 def layer(x, handle):
     ...                                   # convolution, split, gates
-    return chunk_gdn_kcp(q, k, v, g, beta, ctx, recompute_handle=handle)
+    return chunk_gdn_kcp(q, k, v, g, beta, ctx, recompute_handle=handle, deterministic=False)
 
 y = torch.utils.checkpoint.checkpoint(layer, x, RecomputeHandleKcp(), use_reentrant=False)
 ```
@@ -495,7 +514,8 @@ gdn_workspace_bytes_kcp(ctx, *, H_qk, HV, save_mode="light", boundary_mode="save
     At P > 1 the c1 send pool also holds the M_f^T tiles (`c1_send_mf` = s * n_out(r)) and, with the fused walks, the
     carry plane max(n_in(r), n_out(r)) more slots (`fused_walks=0` is the test knob's split form); graph=,
     save_for_backward= and fused_walks= are keywords of this package's query only (the deterministic twin's query
-    has none: through the facade they need deterministic=False with torch's deterministic algorithms off); the b1
+    has none; through the facade they need deterministic=False, like every facade call, and
+    torch.use_deterministic_algorithms does not change the route); the b1
     pools take the larger split form; a backward on a caller-held workspace also sizes the forward exchange pools
     (c1_send, c2_send, c1_recv unless saved); with_conv=True (default) adds the layer convolution's halo rows, which
     the entry points do not allocate (with_conv=False: what a call binds); boundary_mode="replay" gives the same

flashinfer/experimental/cake_gdn_train_kcp/cake_bytes.py

@@ -559,7 +559,8 @@ def gdn_workspace_layout_kcp(
     (``c1_send_mf``, ``n_out`` states; ``c1_send`` stays the collective's bytes) follow it; the split terms do not
     depend on ``P``.  ``fused_walks`` is the call's test knob (1, the shipped form).  ``graph``,
     ``save_for_backward`` and ``fused_walks`` are keywords of this package's query only (the deterministic twin's has
-    none of them): through the facade they need ``deterministic=False`` with torch's deterministic algorithms off.
+    none of them).  Through the facade they need ``deterministic=False``, like every facade call;
+    ``torch.use_deterministic_algorithms`` does not change the route.
     The b1 pools are counted at the larger of their two split forms, as the host sizes them.  A backward on a
     caller-held workspace sizes the forward exchange pools too (``c1_send`` / ``c1_send_mf``, ``c2_send`` and
     ``c1_recv`` unless the saved set holds it; a call without ``workspace=`` skips them after a saved forward).