feat(cake_gdn_det): sm_107a (R200) libraries with grad

flashinfer/experimental/cake_gdn_det_train/README.md

@@ -1,4 +1,4 @@
-# Cake deterministic GDN chunked training backend (SM100 / SM103)
+# Cake deterministic GDN chunked training backend (SM100 / SM103 / SM107)
 
 Experimental backend for the forward and backward of packed GDN (scalar-gated
 DeltaNet) layers in training: head dimension 128, chunk 64, grouped Q/K heads
@@ -10,8 +10,9 @@ side stream; the groups write disjoint outputs, so this holds with one or two
 streams.
 
 Both the API and the backend are experimental: no compatibility guarantee.
-SM100 (B200 / GB200) and SM103 (B300 / GB300) are the supported
-architectures; another device raises `RuntimeError`.  The feature does not
+SM100 (B200 / GB200), SM103 (B300 / GB300) and SM107 (R200) are the supported
+architectures; another device raises `RuntimeError`.  The SM107 libraries need
+CUDA 13.4 or newer.  The feature does not
 participate in automatic backend selection, autotuning or trace-apply.  Its
 libraries are loaded prebuilt from `aot/` when a builder compiled them and
 JIT-built by FlashInfer from the sources and flags `build.json` records
@@ -59,7 +60,9 @@ The forward runs the prep (Q/K L2 norm, gate cumsum), the tcgen05 WY stage and
 the forward scan fused with the output (`fwd_ho`).  The prep stages the Q/K
 tiles by TMA with one CTA per (chunk, key head) (`prep_tma`), except where that
 grid needs a second, short wave: with `slots = 6 x SMs` and
-`CTAs = chunks x H`, the plain prep runs for `slots < CTAs <= 2.5 x slots`, and
+`CTAs = chunks x H`, the plain prep runs for `slots < CTAs <= 2.5 x slots`
+(`sm_100a`, `sm_103a`; `sm_107a` has no plain-prep band and runs `prep_tma`
+for every grid), and
 for Q or K not 16-byte aligned.  Both preps give the same bits.  A long sequence can be
 split into segments: a pre-scan and a stitch then give every segment its start
 state.  The backward runs the adjoint scan, either `dv_local` plus a register
@@ -115,4 +118,4 @@ captured in a CUDA graph (warm it up first, like any captured call).
 
 - Head dimension 128 and chunk 64 only; `q`, `k`, `v` BF16.
 - No context parallelism and no cross-rank state passing.
-- `sm_100a` and `sm_103a` only.
+- `sm_100a`, `sm_103a` and `sm_107a` only (`sm_107a` with CUDA 13.4 or newer).

flashinfer/experimental/cake_gdn_det_train/__init__.py

@@ -15,7 +15,7 @@
 """
 
 # Experimental Cake backend: deterministic GDN (Gated DeltaNet) chunked
-# training (head dimension 128, chunk 64) on SM100 / SM103, with FLA's
+# training (head dimension 128, chunk 64) on SM100 / SM103 / SM107, with FLA's
 # ``chunk_gated_delta_rule`` signature and no context parallelism.  The
 # public entry point is ``chunk_gated_delta_rule``; the chunk and segment
 # tables, the scan planner, the launch order and the autograd chain live in

flashinfer/experimental/cake_gdn_det_train/cake_backend.py

@@ -14,7 +14,7 @@
 limitations under the License.
 
 Loader and launch glue of the Cake deterministic GDN chunked-training kernels
-(SM100 / SM103): the device -> architecture map, the library loading through
+(SM100 / SM103 / SM107): the device -> architecture map, the library loading through
 the package's build manifest and the launch of one registry record.
 
 The registry (:mod:`.cake_jit`) is data: ``MODULES`` holds one record per
@@ -71,9 +71,10 @@
     "stage_record",
 ]
 
-# One verified build per architecture: B200 / GB200 (``sm_100a``) and B300 /
-# GB300 (``sm_103a``).  Every architecture is compiled with its exact
-# ``-gencode`` flag set (``build.json`` ``arch_flags``); no family fallback.
+# One verified build per architecture: B200 / GB200 (``sm_100a``), B300 /
+# GB300 (``sm_103a``) and R200 (``sm_107a``, CUDA 13.4 or newer).  Every
+# architecture is compiled with its exact ``-gencode`` flag set
+# (``build.json`` ``arch_flags``); no family fallback.
 ARCHES: tuple[str, ...] = tuple(cake_jit.ARCHES)
 ARCH_BY_CAPABILITY: dict[tuple[int, int], str] = dict(cake_jit.ARCH_BY_CAPABILITY)