perf(cake_gdn_kcp): with grad, graph-replay checkpoint recompute hands its tapes to the backward

flashinfer/cake_gdn_train_kcp.py

@@ -34,7 +34,10 @@
 queries of either route describe the same plan.  A saved set is read only by
 the backward of the route that produced it, and an options instance is
 converted to the routed package's options type (mappings and keywords work for
-both).  Thin experimental entry points: the partition / context plan, the host
+both; on the deterministic route an ``OptionsKcp`` with the default
+``reuse_recompute_tapes="auto"`` takes that route's own default for the call's
+final ``graph``: ``False`` under replay, whether replay comes from the instance,
+the ``graph`` keyword or a knob).  Thin experimental entry points: the partition / context plan, the host
 orchestration, the saved set, the autograd Function and the generated kernels
 live in the routed package; this module forwards every keyword.  Every switch
 is a keyword argument; no environment variable is read.
@@ -132,13 +135,24 @@ def _deterministic(deterministic: bool, entry: str) -> bool:
 def _options_for(det: bool, options: Any) -> Any:
     """``options`` for the routed package: an options instance of the other route becomes the mapping of its fields
     (the routed package validates it); a mapping, ``None`` or the routed package's own instance pass through.  Both
-    routes take every field, ``use_qk_l2norm_in_kernel`` included."""
+    routes take every field, ``use_qk_l2norm_in_kernel`` included.  On the deterministic route a non-deterministic
+    instance whose ``reuse_recompute_tapes`` is the default ``"auto"`` maps that field to ``None`` whatever its own
+    ``graph`` field says: the deterministic package resolves it to its own default with the call's final ``graph``
+    (the instance's, the ``graph`` keyword's or a knob's): ``"auto"`` under ``"off"``, ``False`` under replay.  An
+    explicit ``True`` passes and is refused there under replay."""
     if options is None or isinstance(options, dict):
         return options
     own = _det().OptionsDetKcp if det else _pkg().OptionsKcp
     if isinstance(options, own) or not dataclasses.is_dataclass(options):
         return options
-    return {f.name: getattr(options, f.name) for f in dataclasses.fields(options)}
+    fields = {f.name: getattr(options, f.name) for f in dataclasses.fields(options)}
+    if (
+        det
+        and isinstance(fields.get("reuse_recompute_tapes"), str)
+        and fields["reuse_recompute_tapes"] == "auto"
+    ):
+        fields["reuse_recompute_tapes"] = None
+    return fields
 
 
 @flashinfer_experimental_api(feature=_FEATURE)
@@ -267,8 +281,8 @@ def chunk_gdn_fwd_kcp(
         Keyword-only options (``save_mode`` ``"light"`` | ``"full"``, ``boundary_mode`` ``"saved"`` | ``"replay"``,
         ``exchange`` ``"a2a"`` | ``"ring"``, ``save_transfer``, ``offset_class`` ``"fla"`` | ``"fp32"``, ``layout`` ``"tma"`` |
         ``"copy"``, ``group_reduce`` ``"fused"`` | ``"host"``, ``graph`` ``"off"`` | ``"replay"``,
-        ``reuse_recompute_tapes`` ``"auto"`` (default) | ``True`` | ``False`` (``False`` with ``graph="replay"``; see
-        :func:`chunk_gdn_kcp`),
+        ``reuse_recompute_tapes`` ``"auto"`` (default) | ``True`` | ``False`` (the deterministic route: ``False`` with
+        ``graph="replay"``; see :func:`chunk_gdn_kcp`),
         ``recompute_output`` ``"auto"`` | ``"always"`` (an autograd-node switch; this forward-only entry never skips a
         launch), ...).
 
@@ -425,8 +439,15 @@ def chunk_gdn_kcp(
     ``recompute_output``) and takes the replay form otherwise -- a region this node closes keeps today's speed and
     memory; ``True`` always keeps them, ``False`` never (the replay form).  One kept set is
     live at a time: in a checkpoint segment holding several GDN layers only the last layer skips the replay (the
-    earlier layers replay as before).  Not available with ``graph="replay"``: there the default resolves to
-    ``False`` and an explicit ``True`` / ``"auto"`` is refused.
+    earlier layers replay as before).  On the deterministic route it is not available with ``graph="replay"`` (the
+    default resolves to ``False``, so does an ``OptionsKcp``'s default ``"auto"``, and an explicit ``True`` /
+    ``"auto"`` is refused); on the non-deterministic route a caller-held workspace keeps its own tapes under
+    ``graph="replay"`` and the backward binds them only while the workspace made no other call in between (else the
+    replay form); the kept tapes add no memory there, while the recompute's own hand-off graph holds its static
+    results like every graph key (one more saved set per checkpointed layer and, when that recompute launches the
+    output program -- ``recompute_output="always"``, or ``True`` with a later saver -- a full static output); there
+    ``"auto"`` decides only from ``cake_kcp.GRAPH_RECOMPUTE_MIN_ROWS`` token x value-head rows per rank (a smaller
+    recompute replays the original forward's graph: its host-bound step gains nothing from the hand-off).
 
     ``recompute_output`` (keyword-only, ``"auto"`` default | ``"always"``): torch's recompute of a checkpointed
     forward returns nothing, so when this node is the last operation of the checkpointed region that saves a tensor,
@@ -435,9 +456,11 @@ def chunk_gdn_kcp(
     early stop is on, and the torch release -- ``torch.__version__`` without its local ``+cuNNN`` build tag -- is one
     of the package's verified releases, ``cake_kcp.RECOMPUTE_OUTPUT_TORCH_VERSIONS``) and launches it as before in every other state -- a later operation
     of the region that saves the output (another layer, a norm), a reentrant checkpoint, a forward outside a
-    recompute, a non-executable node, another torch build (``cake_kcp.LAST_CALL["fwd"]["fwd_o_reason"]`` records the
-    decision and its reason).  The gradients and the retained bytes are unchanged either way; ``"always"`` is the
-    unconditional launch.
+    recompute, a non-executable node, another torch build, and under ``graph="replay"`` with the default
+    ``reuse_recompute_tapes="auto"`` a rank below ``cake_kcp.GRAPH_RECOMPUTE_MIN_ROWS`` (reason ``"small"``; an
+    explicit ``True`` / ``False`` reads the proof at every size) (``cake_kcp.LAST_CALL["fwd"]["fwd_o_reason"]``
+    records the decision and its reason).  The gradients and the retained bytes are unchanged either way;
+    ``"always"`` is the unconditional launch.
 
     ``recompute_handle`` (keyword-only, default ``None``; a :class:`RecomputeHandleKcp`): the checkpoint recompute
     of this forward runs without cross-rank communication.  The original forward exchanges as always and stores

flashinfer/experimental/cake_gdn_train_kcp/README.md

@@ -262,9 +262,9 @@ eager calls too.
 | `precision_mode` | **`0`** (TF32 hi/lo), `1`, `2` (test-only) | the state-transfer products of the transfer walk: `0` is the shipped two-MMA hi/lo form (the TF32x3 class); `1` is its three-MMA twin (bitwise the same results) and `2` plain TF32 (the negative control), both test-only programs of `build.test.json` |
 | `transfer_walk` | **`"tf32_hilo"`** | the transfer-walk program; the vocabulary also names `"bf16x3"` (BF16x3 state products), which this package does not build: the validator refuses it as not built |
 | `layout` | **`"tma"`**, `"copy"` | strided input read in place, or a timed contiguous copy |
-| `graph` | **`"off"`**, `"replay"` | CUDA-graph replay of the steady state (a static pool per key) |
-| `reuse_recompute_tapes` | **`"auto"`**, `True`, `False` | under `torch.utils.checkpoint(use_reentrant=False)` the light backward binds the forward tapes torch's recompute just wrote instead of replaying #1 and the #2 tape walk (bitwise the same gradients, the retention unchanged; a no-op outside that backward; `False` restores the replay form). `"auto"` keeps the tapes only when the GDN node is the last operation of the checkpointed region that saves a tensor (the same checkpoint-frame proof as `recompute_output="auto"`): then its backward is the region's first and the kept tapes die before any other backward runs; with a later saver in the region (an output projection, an MLP) the kept set would live through those operations' backward and raise the region's backward peak by the full tape set (`gdn_saved_bytes_kcp(save_mode="full")` minus the saved states), so `"auto"` takes the replay form there. Measured on one GB300 (torch 2.14.1), 261,987 tokens, `(H_qk, HV) = (2, 6)`: a region of GDN + a 768x768 output projection + a 768->2x2048->768 SwiGLU has a backward peak of 9250.54 MiB with the kept tapes (`True`) and 6927.31 MiB with the replay form (`False` and `"auto"`), at 21.29 vs 22.70 ms per training step of the region; a region the GDN node closes keeps the tapes under `"auto"` (five 240k-262k-token (16, 48) packs: training step and steady peak unchanged against `True`). `True` always keeps the tapes, `False` never. With `graph="replay"`, the one mode that excludes the kept tapes, the default is `False` and an explicit `True` / `"auto"` is refused. One live kept set: in a checkpoint segment with several GDN layers only the last layer skips the replay, the earlier layers replay as before |
-| `recompute_output` | **`"auto"`**, `"always"` | under `torch.utils.checkpoint(use_reentrant=False)` torch re-runs the forward inside `.backward()` only to re-save its tensors and discards that run's output; when the GDN node is the last operation of the checkpointed region that saves a tensor, `"auto"` does not launch the output program (#5) of that re-run. The proof is read from torch's checkpoint frame (the node is executable, the remaining saved-tensor holders are exactly this node's, early stop is on, and the torch release -- `torch.__version__` without its local `+cuNNN` build tag (`cake_kcp.torch_release`) -- is one of the verified releases in `cake_kcp.RECOMPUTE_OUTPUT_TORCH_VERSIONS`); any other state -- a later saver in the region such as the next layer or a norm on the output, a non-executable node, a reentrant checkpoint, `graph="replay"`, another torch build -- launches it as before, and `cake_kcp.LAST_CALL["fwd"]["fwd_o_reason"]` records why. Gradients and retained bytes are the same either way; `"always"` is the unconditional launch |
+| `graph` | **`"off"`**, `"replay"` | CUDA-graph replay of the steady state (a static pool per key). The process keeps up to `cake_kcp.GRAPH_CACHE_MAX` = 256 keys, least recently used first out: a plain layer holds one forward and one backward key, a checkpointed layer three (the original forward, the recompute in its hand-off form, the kept backward; two below the size rule of `reuse_recompute_tapes`), so 85 checkpointed layers replay without recapture; each key's graph holds its static results (out, saved set, gradients) until it leaves the cache or `graph_cache_clear()` |
+| `reuse_recompute_tapes` | **`"auto"`**, `True`, `False` | under `torch.utils.checkpoint(use_reentrant=False)` the light backward binds the forward tapes torch's recompute just wrote instead of replaying #1 and the #2 tape walk (bitwise the same gradients, the retention unchanged; a no-op outside that backward; `False` restores the replay form). `"auto"` keeps the tapes only when the GDN node is the last operation of the checkpointed region that saves a tensor (the same checkpoint-frame proof as `recompute_output="auto"`): then its backward is the region's first and the kept tapes die before any other backward runs; with a later saver in the region (an output projection, an MLP) the kept set would live through those operations' backward and raise the region's backward peak by the full tape set (`gdn_saved_bytes_kcp(save_mode="full")` minus the saved states), so `"auto"` takes the replay form there. Measured on one GB300 (torch 2.14.1), 261,987 tokens, `(H_qk, HV) = (2, 6)`: a region of GDN + a 768x768 output projection + a 768->2x2048->768 SwiGLU has a backward peak of 9250.54 MiB with the kept tapes (`True`) and 6927.31 MiB with the replay form (`False` and `"auto"`), at 21.29 vs 22.70 ms per training step of the region; a region the GDN node closes keeps the tapes under `"auto"` (five 240k-262k-token (16, 48) packs: training step and steady peak unchanged against `True`). `True` always keeps the tapes, `False` never. With `graph="replay"` (same default) the kept tapes are the caller-held workspace's own arena tapes (a captured graph bakes their addresses, so nothing is unbound and the kept tapes add no memory); the recompute captures its hand-off form as a graph of its own and the kept backward its own graph (three graph keys per checkpointed layer, see `graph`), so a checkpointed step replays the eager step's programs. Like every key, the recompute's hand-off graph holds its own static results: one more static saved set per checkpointed layer (`gdn_saved_bytes_kcp` of the call) and, when that recompute launches #5 (`recompute_output="always"`, or `True` with a later saver in the region), a full static output (`T_r * HV * 256` bytes) per checkpointed layer. Size rule of `"auto"` under `graph="replay"` (caller-held workspace): a rank with fewer token x value-head rows (`T_r * HV`) than `cake_kcp.GRAPH_RECOMPUTE_MIN_ROWS` of its architecture (61440 on sm_100a; 0 on sm_103a and sm_107a) decides nothing -- the recompute replays the original forward's graph (#5 included, nothing kept, two graph keys per checkpointed layer) and the backward replays #1 and the tape walk -- because such a checkpointed step is host-bound: the decision and the hand-off add host time per layer (the checkpoint-frame read, the stamp and the backward's checks) that the GPU work they save does not repay there (one checkpointed layer, previous / hand-off step time: B200 0.953 at 49152 rows, 1.016 at 61440 rows; GB300 1.013 and R200 1.07 at 49152 rows); `True` / `False` and `graph="off"` decide at every size. The hand-off holds only while the workspace has made no other call since the recompute, still binds those tapes and the backward runs on the recompute's stream; otherwise the backward replays #1 and the tape walk (same bits; `cake_kcp.KEPT_STALE_STATS` counts why). One live kept set: in a checkpoint segment with several GDN layers only the last layer skips the replay, the earlier layers replay as before |
+| `recompute_output` | **`"auto"`**, `"always"` | under `torch.utils.checkpoint(use_reentrant=False)` torch re-runs the forward inside `.backward()` only to re-save its tensors and discards that run's output; when the GDN node is the last operation of the checkpointed region that saves a tensor, `"auto"` does not launch the output program (#5) of that re-run. The proof is read from torch's checkpoint frame (the node is executable, the remaining saved-tensor holders are exactly this node's, early stop is on, and the torch release -- `torch.__version__` without its local `+cuNNN` build tag (`cake_kcp.torch_release`) -- is one of the verified releases in `cake_kcp.RECOMPUTE_OUTPUT_TORCH_VERSIONS`); any other state -- a later saver in the region such as the next layer or a norm on the output, a non-executable node, a reentrant checkpoint, another torch build -- launches it as before (with `graph="replay"` the proof is read only inside a backward pass and the decision is part of the recompute's CUDA-graph key; the left-out output is a one-element broadcast view there, so the captured graph holds no output buffer; on a caller-held workspace with `reuse_recompute_tapes="auto"`, a rank below `cake_kcp.GRAPH_RECOMPUTE_MIN_ROWS`, the size rule of `reuse_recompute_tapes`, decides nothing and launches #5 with reason `"small"`, while an explicit `True` / `False` reads the proof at every size), and `cake_kcp.LAST_CALL["fwd"]["fwd_o_reason"]` records why. Gradients and retained bytes are the same either way; `"always"` is the unconditional launch |
 | `segments` | **`None`**, an int | how many intra-rank segments a fragment is walked as (`P`): `None` takes the context's plan, an int overrides it within `[1, min(64, chunks of the rank's longest fragment)]` (the byte queries' `P=`). The plan (test knob `p_rule` names its policy) is a function of metadata and the SM count only. The shipped policy `fill_small_guard` starts from the deterministic twin's `row_cost_lpt` plan, which costs the three per-rank programs over every rank's fragment table with a crossing-aware model and switches the whole row away from the per-rank default only when an alternative is cheaper by at least its margin (2 % for the `occ` program; 0 % for the `one` program, which has no segment composition and takes a model tie), then lowers `P` as described under "Numerical mode" (faster schedule); the same topology gives the same `P` on every rank |
 | `offset_class` | **`"fla"`**, `"fp32"` | the class of the forward offset `B_f`: the walk from zero with the reference rounding points, or an FP32 walk (`B_f` only) |
 | `use_qk_l2norm_in_kernel` | **`False`**, `True` | the per-token L2 norm of every 128-channel Q / K head inside the programs (FLA's `use_qk_l2norm_in_kernel`; see "In-kernel Q / K L2 norm" below); a literal keyword of the three entry points as well |
@@ -469,7 +469,12 @@ gdn_workspace_bytes_kcp(ctx, *, H_qk, HV, save_mode="light", boundary_mode="save
     so the forward is every forward pool and tape it writes (no tape with save_mode="full" and a saved set) and the
     backward every backward pool, the ten light tapes (none with save_mode="full" and a saved set) and every
     backward tape; with use_qk_l2norm_in_kernel both add qn / kn and the arena qk_rstd where #1 writes it (a
-    forward without a saved set, a backward that re-runs #1).
+    forward without a saved set, a backward that re-runs #1).  There `reuse_recompute_tapes=True` gives the arena
+    hand-off's figure exactly: the recompute's tapes, M_seg^T plane and qn / kn stay arena tensors, so it is the
+    replay form's less the backward's arena qk_rstd (`bwd.kept_tapes` 0). The default (`"auto"`) gives the replay
+    form's figure at every size, an upper bound of the call: equal when the call takes the replay form (below
+    cake_kcp.GRAPH_RECOMPUTE_MIN_ROWS, or with a later saver in the region), larger by the arena qk_rstd (with
+    use_qk_l2norm_in_kernel; else equal) when it keeps the tapes.
     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
@@ -519,8 +524,10 @@ other reader (#2, #3, #5, #2+5, #7) is the unchanged program bound to `qn` / `kn
 The saved set gains `qk_rstd` (FP32 `[2, T_r, H_qk]`, `8 T_r H_qk` bytes;
 `KCP_SAVED_VERSION` 2); the normalized rows are never saved
 (`2 x 256 T_r H_qk` bytes: workspace in the forward's tape walk and every
-backward phase; in the kept-tape backward they are hand-off tensors,
-`bwd.kept_tapes`).  The byte queries take `use_qk_l2norm_in_kernel=`;
+backward phase; in the kept-tape backward under `graph="off"` they are
+hand-off tensors, `bwd.kept_tapes`; under `graph="replay"` the arena hand-off
+keeps them as arena rows: `bwd.qn` / `bwd.kn` at full size, `bwd.kept_tapes`
+0).  The byte queries take `use_qk_l2norm_in_kernel=`;
 `LAST_CALL[...]` records `qk_l2norm`, `qk_source`, `rstd_source` and
 `arena_tapes`.  The recompute handle's entry key includes the flag.  Until
 the registry renders the three programs the kernel backend refuses the