[yyihuang]

This PR changes flashinfer.experimental.cake_gdn_train_kcp, the non-deterministic context-parallel variant of Cake's Gated DeltaNet (GDN) chunked training kernel with gradients. Callers reach it through flashinfer.cake_gdn_train_kcp(deterministic=False). The change is host-side only and has three parts:

1. The forward and backward bodies now enter one tvm-ffi stream context per entry body. Before, they entered a context for every kernel launch. A launch outside every scope still enters its own context, and so does a launch after a stream switch inside a scope. Inside a CUDA graph capture, the scope binds the capture stream.
2. The static launch arguments of each launch key (record, grid, test flag, merge columns) are resolved once per context plan and kept with the context's tables. Each launch still rebuilds its public values and appends its launch record.
3. The public entry flashinfer.cake_gdn_train_kcp caches the two package modules it routes to. Before, it re-imported the module on every call. This also removes that per-call import from the deterministic=True route, whose behavior is unchanged.

Device programs, launch order, grids, streams and workspace are unchanged, and so are the checkpoint and graph-replay semantics. In every timed run (eager and graph replay, both ranks), outputs and every gradient are bitwise equal to main's non-deterministic package. Launch records (programs and launch count per phase) are also equal.

Speedup is baseline latency / new latency, so values above 1 mean the new package is faster. There are three baselines, and all of them ran in the same process as the new package:

• Deterministic Cake KCP: the deterministic Cake KCP package from main, called through its public entry.
• Strongest non-Cake, per row: the faster of FLA 0.6.0 context-parallel KCP (layout C) and a head-parallel all-to-all around the cuDNN frontend 1.32.1 FROST GDN engine. Both were qualified by comparison with the new package. FlashQLA has no sm_107a build, so these two are the applicable candidates.
• Previous non-deterministic: main's non-deterministic KCP package. It cannot share the public entry with the new package in one process, so it is loaded as a renamed copy and called through an equivalent stand-in entry. The stand-in skips the public entry's dispatch, so eager ratios vs this package are conservative by up to about 1 % (a path-only control with the new code on the stand-in entry measured 0.986-1.003 on forward+backward).

All arms ran on one exclusive R200 node (sm_107a, 212 SMs; CUDA 13.0 runtime, PyTorch 2.14.1+cu130; no power capping observed). The setup was CP2 (2 ranks, GPUs 0-1) with BF16 q/k/v, FP32 g/beta and head dim 128.

Protocol: 5 groups, each running the arms in palindromic order 2 times. Each call is 12 timed iterations after 10 warm-up, with a 256 MiB L2 flush, device sync and barrier before every iteration. Latency is the median CUDA-event time of the slowest rank, and speedup is the median of the per-group ratios [min, max]. Two runs ran in separate processes, and run 2 reverses the arm order. Eager uses default options. Graph replay uses a caller-owned workspace with graph="replay".

How to read the table:

• Backward = forward+backward − forward, per group.
• "(medians)" marks a ratio of medians.
• "unpaired" marks a non-Cake value measured with layout C for a layout-D row.
• (H, HV) = (query/key heads, value heads).
• S26 = (2, 6), 3 documents of 2049/3001/3142 tokens (8192).
• DET = (2, 6), 65536 tokens in 6 documents.
• P1 = (16, 48), 241817 tokens.
• P3 = (16, 48), 259174 tokens.
• Layout C is contiguous; layout D is zigzag per document.

Five ratio cells are classed UNSTABLE because their group-to-group spread is above the 1 % tolerance, all on S26 eager forward+backward (two in run 1, three in run 2). They are shown as measured.

Table: R200 (sm_107a, 212 SMs), CUDA 13.0 runtime, PyTorch 2.14.1+cu130, CP2; latencies in ms (median CUDA-event time of the slowest rank); speedup = median of the per-group ratios [min, max].

GPU / architecture	workload (pack, (H, HV), S tokens, CP, layout); timing mode; run	operation	deterministic Cake (ms)	strongest non-Cake (ms, implementation)	previous non-deterministic Cake (ms)	new non-deterministic Cake (ms)	speedup vs deterministic Cake	speedup vs non-Cake SOTA	speedup vs previous non-deterministic Cake
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 1	forward	0.798	3.263 (FLA KCP)	0.798	0.718	1.115 [1.099, 1.119]	4.538 [4.511, 4.590]	1.113 [1.104, 1.115]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 1	backward (= forward+backward − forward)	1.321	3.711 (FLA KCP)	1.327	1.213	1.094 [1.082, 1.103]	3.060 [3.060, 3.060] (medians)	1.096 [1.085, 1.106]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 1	forward+backward	2.118	6.974 (FLA KCP)	2.125	1.930	1.102 [1.095, 1.103]	3.616 [3.582, 3.649]	1.101 [1.095, 1.110]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 1	checkpointed 2×forward+backward	3.147	10.154 (FLA KCP)	3.158	2.888	1.089 [1.083, 1.095]	3.505 [3.485, 3.535]	1.090 [1.085, 1.098]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 1	forward	0.543	n/a (no graph-replay arm)	0.524	0.526	1.033 [1.030, 1.034]	n/a	0.997 [0.996, 1.001]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 1	backward (= forward+backward − forward)	0.854	n/a (no graph-replay arm)	0.831	0.833	1.026 [1.023, 1.028]	n/a	0.998 [0.996, 0.998]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 1	forward+backward	1.397	n/a (no graph-replay arm)	1.355	1.358	1.029 [1.027, 1.029]	n/a	0.997 [0.997, 0.999]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 1	checkpointed 2×forward+backward	1.934	n/a (no graph-replay arm)	1.887	1.890	1.024 [1.022, 1.025]	n/a	0.998 [0.997, 0.999]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 2 (reversed arm order)	forward	0.797	3.263 (FLA KCP)	0.804	0.729	1.090 [1.083, 1.105]	4.473 [4.473, 4.473] unpaired	1.105 [1.097, 1.107]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 2 (reversed arm order)	backward (= forward+backward − forward)	1.327	3.711 (FLA KCP)	1.326	1.204	1.099 [1.079, 1.104]	3.082 [3.082, 3.082] (medians)	1.096 [1.075, 1.105]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 2 (reversed arm order)	forward+backward	2.124	6.974 (FLA KCP)	2.131	1.934	1.096 [1.083, 1.099]	3.607 [3.607, 3.607] unpaired	1.098 [1.086, 1.105]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; eager; run 2 (reversed arm order)	checkpointed 2×forward+backward	3.154	10.154 (FLA KCP)	3.179	2.922	1.078 [1.076, 1.082]	3.474 [3.474, 3.474] unpaired	1.089 [1.083, 1.091]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 2 (reversed arm order)	forward	0.543	n/a (no graph-replay arm)	0.526	0.527	1.030 [1.028, 1.033]	n/a	0.998 [0.997, 1.002]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 2 (reversed arm order)	backward (= forward+backward − forward)	0.851	n/a (no graph-replay arm)	0.835	0.834	1.020 [1.017, 1.022]	n/a	1.003 [0.999, 1.004]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 2 (reversed arm order)	forward+backward	1.393	n/a (no graph-replay arm)	1.361	1.361	1.024 [1.023, 1.025]	n/a	1.001 [1.000, 1.001]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, C; graph replay; run 2 (reversed arm order)	checkpointed 2×forward+backward	1.932	n/a (no graph-replay arm)	1.895	1.893	1.020 [1.020, 1.021]	n/a	1.001 [1.000, 1.001]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 1	forward	0.880	3.263 (FLA KCP)	0.628	0.628	1.403 [1.397, 1.406]	5.194 [5.194, 5.194] unpaired	0.999 [0.997, 1.002]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 1	backward (= forward+backward − forward)	1.397	3.711 (FLA KCP)	0.997	0.868	1.616 [1.583, 1.624]	4.273 [4.273, 4.273] (medians)	1.149 [1.135, 1.160]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 1	forward+backward	2.277	6.974 (FLA KCP)	1.625	1.497	1.526 [1.508, 1.531]	4.660 [4.660, 4.660] unpaired	1.086 [1.079, 1.092]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 1	checkpointed 2×forward+backward	3.392	10.154 (FLA KCP)	2.601	2.401	1.415 [1.403, 1.422]	4.229 [4.229, 4.229] unpaired	1.082 [1.077, 1.090]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 1	forward	0.596	n/a (no graph-replay arm)	0.589	0.589	1.011 [1.010, 1.014]	n/a	1.000 [0.997, 1.000]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 1	backward (= forward+backward − forward)	0.811	n/a (no graph-replay arm)	0.809	0.808	1.004 [0.999, 1.006]	n/a	1.001 [0.997, 1.004]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 1	forward+backward	1.407	n/a (no graph-replay arm)	1.398	1.397	1.007 [1.006, 1.008]	n/a	1.001 [0.998, 1.002]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 1	checkpointed 2×forward+backward	2.036	n/a (no graph-replay arm)	2.021	2.018	1.009 [1.008, 1.010]	n/a	1.001 [1.001, 1.003]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 2 (reversed arm order)	forward	0.893	3.263 (FLA KCP)	0.627	0.627	1.425 [1.420, 1.433]	5.202 [5.202, 5.202] unpaired	1.001 [0.999, 1.002]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 2 (reversed arm order)	backward (= forward+backward − forward)	1.408	3.711 (FLA KCP)	1.021	0.886	1.581 [1.569, 1.608]	4.187 [4.187, 4.187] (medians)	1.148 [1.141, 1.164]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 2 (reversed arm order)	forward+backward	2.301	6.974 (FLA KCP)	1.648	1.513	1.517 [1.513, 1.530]	4.608 [4.608, 4.608] unpaired	1.087 [1.084, 1.095]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; eager; run 2 (reversed arm order)	checkpointed 2×forward+backward	3.415	10.154 (FLA KCP)	2.597	2.381	1.431 [1.414, 1.448]	4.264 [4.264, 4.264] unpaired	1.087 [1.077, 1.098]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 2 (reversed arm order)	forward	0.596	n/a (no graph-replay arm)	0.592	0.592	1.007 [1.006, 1.010]	n/a	1.001 [0.999, 1.002]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 2 (reversed arm order)	backward (= forward+backward − forward)	0.810	n/a (no graph-replay arm)	0.812	0.812	0.997 [0.992, 1.001]	n/a	0.997 [0.995, 1.003]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 2 (reversed arm order)	forward+backward	1.406	n/a (no graph-replay arm)	1.404	1.404	1.001 [1.000, 1.003]	n/a	0.999 [0.998, 1.001]
R200 sm_107a (212 SMs)	DET, (2, 6), 65536, CP2, D; graph replay; run 2 (reversed arm order)	checkpointed 2×forward+backward	2.034	n/a (no graph-replay arm)	2.029	2.028	1.004 [1.002, 1.004]	n/a	1.001 [1.000, 1.002]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 1	forward	4.856	22.058 (FROST head-CP)	4.795	4.794	1.013 [1.013, 1.013]	4.601 [4.601, 4.601] unpaired	1.000 [1.000, 1.000]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 1	backward (= forward+backward − forward)	8.545	32.074 (FROST head-CP)	8.542	8.549	1.000 [0.999, 1.000]	3.752 [3.752, 3.752] (medians)	0.999 [0.999, 1.000]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 1	forward+backward	13.402	54.132 (FROST head-CP)	13.336	13.344	1.004 [1.004, 1.005]	4.057 [4.057, 4.057] unpaired	0.999 [0.999, 1.000]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 1	checkpointed 2×forward+backward	15.075	70.721 (FROST head-CP)	14.936	14.950	1.008 [1.008, 1.009]	4.730 [4.730, 4.730] unpaired	0.999 [0.998, 1.000]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 2 (reversed arm order)	forward	4.847	22.058 (FROST head-CP)	4.781	4.783	1.013 [1.013, 1.014]	4.611 [4.610, 4.615]	0.999 [0.999, 1.000]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 2 (reversed arm order)	backward (= forward+backward − forward)	8.551	32.074 (FROST head-CP)	8.549	8.554	1.000 [0.999, 1.001]	3.750 [3.750, 3.750] (medians)	0.999 [0.999, 1.001]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 2 (reversed arm order)	forward+backward	13.398	54.132 (FROST head-CP)	13.330	13.337	1.005 [1.004, 1.005]	4.059 [4.058, 4.060]	0.999 [0.999, 1.000]
R200 sm_107a (212 SMs)	P1, (16, 48), 241817, CP2, C; eager; run 2 (reversed arm order)	checkpointed 2×forward+backward	15.060	70.721 (FROST head-CP)	14.926	14.930	1.009 [1.007, 1.009]	4.737 [4.733, 4.740]	1.000 [0.999, 1.001]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 1	forward	5.024	7.911 (FLA KCP)	4.947	4.949	1.015 [1.014, 1.016]	1.599 [1.599, 1.599] unpaired	1.000 [0.998, 1.000]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 1	backward (= forward+backward − forward)	9.085	51.819 (FROST head-CP forward+backward − FLA KCP forward)	9.075	9.073	1.002 [1.001, 1.002]	5.711 [5.711, 5.711] (medians)	1.000 [0.999, 1.002]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 1	forward+backward	14.109	59.730 (FROST head-CP)	14.023	14.022	1.006 [1.006, 1.007]	4.260 [4.260, 4.260] unpaired	1.000 [1.000, 1.000]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 1	checkpointed 2×forward+backward	15.823	67.948 (FLA KCP)	15.664	15.671	1.010 [1.009, 1.011]	4.336 [4.336, 4.336] unpaired	0.999 [0.999, 1.001]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 2 (reversed arm order)	forward	5.024	7.911 (FLA KCP)	4.945	4.948	1.015 [1.015, 1.016]	1.599 [1.598, 1.600]	0.999 [0.998, 1.000]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 2 (reversed arm order)	backward (= forward+backward − forward)	9.079	51.819 (FROST head-CP forward+backward − FLA KCP forward)	9.063	9.069	1.001 [1.000, 1.002]	5.714 [5.714, 5.714] (medians)	1.000 [0.997, 1.001]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 2 (reversed arm order)	forward+backward	14.103	59.730 (FROST head-CP)	14.007	14.017	1.006 [1.005, 1.007]	4.261 [4.259, 4.263]	1.000 [0.998, 1.000]
R200 sm_107a (212 SMs)	P3, (16, 48), 259174, CP2, C; eager; run 2 (reversed arm order)	checkpointed 2×forward+backward	15.815	67.948 (FLA KCP)	15.651	15.660	1.010 [1.009, 1.010]	4.339 [4.338, 4.339]	0.999 [0.999, 1.000]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 1	forward	0.525	0.840 (FLA KCP)	0.529	0.471	1.115 [1.098, 1.123]	1.785 [1.750, 1.799]	1.122 [1.110, 1.128]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 1	backward (= forward+backward − forward)	0.780	0.736 (FLA KCP)	0.772	0.706	1.109 [1.087, 1.126]	1.042 [1.042, 1.042] (medians)	1.088 [1.085, 1.121]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 1	forward+backward	1.305	1.576 (FLA KCP)	1.301	1.177	1.110 [1.095, 1.121] UNSTABLE	1.339 [1.339, 1.339] unpaired	1.102 [1.097, 1.119] UNSTABLE
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 1	checkpointed 2×forward+backward	2.223	2.557 (FLA KCP)	2.240	2.062	1.079 [1.072, 1.086]	1.236 [1.234, 1.256]	1.087 [1.081, 1.093]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 1	forward	0.306	n/a (no graph-replay arm)	0.298	0.299	1.023 [1.021, 1.024]	n/a	0.996 [0.993, 0.997]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 1	backward (= forward+backward − forward)	0.329	n/a (no graph-replay arm)	0.329	0.326	1.007 [1.004, 1.013]	n/a	1.006 [1.003, 1.010]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 1	forward+backward	0.635	n/a (no graph-replay arm)	0.627	0.626	1.014 [1.013, 1.017]	n/a	1.001 [1.000, 1.004]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 1	checkpointed 2×forward+backward	1.018	n/a (no graph-replay arm)	1.008	1.010	1.008 [1.006, 1.010]	n/a	0.998 [0.994, 1.002]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 2 (reversed arm order)	forward	0.515	0.840 (FLA KCP)	0.515	0.461	1.115 [1.106, 1.121]	1.822 [1.822, 1.822] unpaired	1.120 [1.111, 1.123]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 2 (reversed arm order)	backward (= forward+backward − forward)	0.763	0.736 (FLA KCP)	0.755	0.688	1.106 [1.087, 1.135]	1.070 [1.070, 1.070] (medians)	1.096 [1.075, 1.121]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 2 (reversed arm order)	forward+backward	1.278	1.576 (FLA KCP)	1.271	1.149	1.110 [1.098, 1.125] UNSTABLE	1.376 [1.358, 1.376] UNSTABLE	1.102 [1.093, 1.118] UNSTABLE
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; eager; run 2 (reversed arm order)	checkpointed 2×forward+backward	2.140	2.557 (FLA KCP)	2.144	1.973	1.083 [1.080, 1.089]	1.296 [1.296, 1.296] unpaired	1.085 [1.081, 1.090]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 2 (reversed arm order)	forward	0.305	n/a (no graph-replay arm)	0.297	0.298	1.023 [1.020, 1.023]	n/a	0.996 [0.994, 0.998]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 2 (reversed arm order)	backward (= forward+backward − forward)	0.327	n/a (no graph-replay arm)	0.326	0.327	0.999 [0.998, 1.001]	n/a	0.996 [0.995, 1.000]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 2 (reversed arm order)	forward+backward	0.631	n/a (no graph-replay arm)	0.623	0.625	1.010 [1.008, 1.011]	n/a	0.997 [0.995, 0.998]
R200 sm_107a (212 SMs)	S26, (2, 6), 8192, CP2, C; graph replay; run 2 (reversed arm order)	checkpointed 2×forward+backward	0.990	n/a (no graph-replay arm)	0.980	0.985	1.008 [1.001, 1.011]	n/a	0.995 [0.994, 0.998]

Reading:

• Eager, small heads: S26, DET C and DET D are 8-10 % faster than main's non-deterministic package on forward+backward and checkpointed steps. Examples: S26 forward+backward 1.102 (classed UNSTABLE), S26 checkpointed 1.085-1.087, DET C forward+backward 1.098-1.101 and DET D forward+backward 1.086-1.087. Single phases reach 1.149 (DET D backward). On forward+backward these rows are also faster than deterministic Cake: S26 1.110, DET C 1.096-1.102, DET D 1.517-1.526.
• Graph replay: graph replay runs none of the changed host code. It is at parity with main's non-deterministic package, at 0.995-1.006 on every graph cell. The smallest graph cell is S26 checkpointed: 0.998 [0.994, 1.002] in run 1 and 0.995 [0.994, 0.998] in run 2. Vs deterministic Cake, graph cells are 0.997-1.033.
• (16, 48) rows: P1 and P3 are GPU-bound and at parity with main's non-deterministic package (0.999-1.000). Vs deterministic Cake they are 1.000-1.015.
• Non-Cake: vs the strongest non-Cake implementation, the new package is faster in every eager cell. The smallest ratio is 1.042-1.070 on S26 backward. On forward+backward it is 1.339-1.376 (S26), 3.607-3.616 (DET C), 4.608-4.660 (DET D), 4.057-4.059 (P1) and 4.260-4.261 (P3). There is no non-Cake graph-replay arm.

B200 (sm_100a) and B300/GB300 (sm_103a) rows are pending and will be added when they are ready; the change is host-side only on every architecture.

The kernel's other Cake variants, all on main, are the deterministic single-GPU flashinfer.experimental.cake_gdn_det_train, the non-deterministic single-GPU flashinfer.experimental.cake.cake_gdn_train (#8157) and the deterministic context-parallel flashinfer.experimental.cake_gdn_det_train_kcp; this PR tunes the non-deterministic context-parallel package from #8172, the two context-parallel variants share the entry flashinfer.cake_gdn_train_kcp(deterministic=...), and the family ships on sm_100a, sm_103a and sm_107a.

Summary by CodeRabbit

• Performance
  • Repeated calls reuse package loading and launch setup, reducing redundant work.
  • Forward and backward operations reuse stream setup for launches on the same stream, while stream changes remain supported.
• Reliability
  • Reused launch settings preserve per-call launch records and results across supported options, precision modes, and merge configurations.
  • Stream handling is maintained across nested operations and restored when scopes exit, including after errors.

[yyihuang]

@flashinfer-bot run tests/experimental/test_cake_gdn_train_kcp.py tests/experimental/test_cake_gdn_train_kcp_api.py tests/experimental/test_cake_gdn_train_kcp_build.py tests/experimental/test_cake_gdn_train_kcp_cp1.py tests/experimental/test_cake_gdn_train_kcp_dist.py tests/experimental/test_cake_gdn_det_train_kcp_build.py tests/experimental/test_cake_gdn_det_train_kcp_api.py

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: e932a33a-250e-404f-8706-3a7633ca0cfa

Commits

Reviewing files that changed from the base of the PR and between 7e5a4a3 and 704529f.

Files selected for processing: 1

• flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py

Files skipped from review as they are similar to previous changes: 1

• flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 3 remain after this review.

Walkthrough

The PR caches routed facade imports and per-plan launch resolution. It adds backend launch scopes that bind the current Torch stream once per forward or backward body, while preserving per-launch binding when no matching scope is active. Tests cover cache equivalence, stream changes, nesting, exceptions, and CUDA execution.

Changes: Launch resolution and stream binding

Layer / File(s)	Summary
Cached launch resolution flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py, tests/experimental/test_cake_gdn_train_kcp.py	Both backends use a per-plan cache for resolved launch data. Tests compare cached resolutions and forward/backward results with uncached behavior.
Scoped stream binding flashinfer/experimental/cake_gdn_train_kcp/cake_backend.py, flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py, flashinfer/experimental/cake_gdn_train_kcp/README.md, tests/experimental/test_cake_gdn_train_kcp.py, tests/experimental/test_cake_gdn_train_kcp_cp1.py	Backends can provide a launch scope for forward and backward bodies. The kernel launcher reuses a matching stream binding and falls back to per-launch binding when no matching scope is active. Tests cover nested scopes, stream changes, and CUDA execution.

Facade route loading

Layer / File(s)	Summary
Cached route imports flashinfer/cake_gdn_train_kcp.py, tests/experimental/test_cake_gdn_train_kcp_api.py	The routed-package and deterministic-route loaders cache imported modules. A test verifies that subsequent calls perform no imports.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant EntryBody
  participant KernelBackend
  participant CakeBackend
  participant TVMFFI
  EntryBody->>KernelBackend: enter launch_scope
  KernelBackend->>CakeBackend: delegate launch_scope
  CakeBackend->>TVMFFI: bind current Torch stream
  EntryBody->>CakeBackend: launch module
  CakeBackend->>TVMFFI: call entry under matching scope

Loading

Possibly related issues

• Issue 7251: The issue covers the same package host implementation; this PR changes the KCP host launch path and caches launch data.
• Issue 4214: The issue concerns GDN cache lifecycle and cache identity; this PR adds a launch cache keyed by launch-selection inputs.
• Issue 8063: The issue requests a non-deterministic packed GDN implementation with gradients and KCP options; this PR changes the non-deterministic KCP training route.

Merge Risk: Minimal · up to 70452

No actionable merge-blocking risk was established. The change is mergeable after normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Description check		The description gives detailed technical scope, behavior, validation, and benchmark results. However, it omits the required template sections and experimental-track information, including the tracking…	Reformat the description using the repository template. Add the Description, Related Issues, Pull Request Checklist, Tests, Experimental Track, and Reviewer Notes sections as applicable. Because this PR changes code under flashinfer/experim…

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the performance-focused host-path change for the non-deterministic Cake GDN KCP package.
Docstring Coverage		Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 20 functions across 6 files.

Full details: Description check

Explanation

The description gives detailed technical scope, behavior, validation, and benchmark results. However, it omits the required template sections and experimental-track information, including the tracking issue, checklist declarations, and the required experimental-tests target block.

Resolution

Reformat the description using the repository template. Add the Description, Related Issues, Pull Request Checklist, Tests, Experimental Track, and Reviewer Notes sections as applicable. Because this PR changes code under flashinfer/experimental/, provide the experimental tracking issue and required declarations, and replace the experimental-tests placeholder with valid targets under tests/experimental/. Record the test outcomes or clearly state any pending validation.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn

[yzh119]

LGTM. Nit: _launch_key assumes resolve_launch reads nothing from the tensors except what's in the key. A comment saying so would help the next person who adds a grid input.

[flashinfer-bot]

CODEOWNERS gap attestation (not a code review).

A code owner of a changed path has approved this PR:

• @yzh119 -> flashinfer/cake_gdn_train_kcp.py
• @yzh119 -> flashinfer/experimental/cake_gdn_train_kcp/cake_backend.py
• @yzh119 -> flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py
• @yzh119 -> flashinfer/experimental/cake_gdn_train_kcp/README.md

The following files match only the catch-all rule; this approval covers that technical gap under the one-approval policy. GitHub still requires each owned path's own owner.

• tests/experimental/test_cake_gdn_train_kcp_api.py
• tests/experimental/test_cake_gdn_train_kcp_cp1.py
• tests/experimental/test_cake_gdn_train_kcp.py

Head: 7e5a4a3

[yyihuang]

Thanks, @yzh119. Added the invariant to the _launch_key docstring in 704529f (docstring only): resolve_launch derives the record, grid, test flag and merge cols_per_cta from the key's scalars alone, so a new input to resolve_launch (for example a grid-rule argument) must be added to the key as well.

[yyihuang]

@flashinfer-bot run tests/experimental/test_cake_gdn_train_kcp.py tests/experimental/test_cake_gdn_train_kcp_api.py tests/experimental/test_cake_gdn_train_kcp_build.py tests/experimental/test_cake_gdn_train_kcp_cp1.py tests/experimental/test_cake_gdn_train_kcp_dist.py tests/experimental/test_cake_gdn_det_train_kcp_build.py tests/experimental/test_cake_gdn_det_train_kcp_api.py