Repository navigation
feat: add ReplaySSM state modes for GDN and KDA - #1690
Merged
Merged
Conversation
hiworldwzj
force-pushed
the
support/replayssm-clean
branch
from
November 12, 2026 08:48
a114f93 to
fe650a4
Compare
Drop the shared SSM executor hierarchy so each algorithm keeps its own replay implementation and the native paths stay easier to review. Co-authored-by: Cursor <cursoragent@cursor.com>
Native linear backends only run fused recurrent now. Replay keeps its own state modules and attention backends, selected when ssm_state_mode is replay. Co-authored-by: Cursor <cursoragent@cursor.com>
Replay inference no longer subclasses the native linear backends. Each algorithm owns its entry point and configures its own prefill kernel. Co-authored-by: Cursor <cursoragent@cursor.com>
History runtimes keep request storage, materialization, and copies. The inference step owns kernel selection and the per-layer replay. Co-authored-by: Cursor <cursoragent@cursor.com>
The separate runtime objects only duplicated the canonical SSM. History buffers, cursors, and accepted-state materialization now live on the state. Co-authored-by: Cursor <cursoragent@cursor.com>
Prefill writes the canonical SSM directly. Replay keeps its kernel config local and no longer shares a slot stride or an executor alias. Co-authored-by: Cursor <cursoragent@cursor.com>
Native slot stride is mtp_step + 1 at each call site. New GDN and KDA constructors take their replay arguments as ordinary parameters. Co-authored-by: Cursor <cursoragent@cursor.com>
GDN and KDA no longer split page copies and accepted snapshots into separate modules. Co-authored-by: Cursor <cursoragent@cursor.com>
The shared acceptance wrapper only unpacked tensors the state already receives. Private helpers stay at the end of the replay state. Co-authored-by: Cursor <cursoragent@cursor.com>
The step modules only launch one decode replay. Decode makes that role visible next to the attention backends. Co-authored-by: Cursor <cursoragent@cursor.com>
run_config is Triton launch metadata, so it no longer sits ahead of state initialization. Co-authored-by: Cursor <cursoragent@cursor.com>
Autotune no longer wraps an existing SSM through a public constructor. The replay kernel entry only consumes that layer's tensors. Co-authored-by: Cursor <cursoragent@cursor.com>
Native GDN and KDA did not need a method reorder. Replay decode now calls prepare and accept directly, and config selection stays in the autotuner. Co-authored-by: Cursor <cursoragent@cursor.com>
Decode appends into a single cursor-length buffer and compresses the whole prefix only when the next verify would exceed capacity, unless force_compress is set. Co-authored-by: Cursor <cursoragent@cursor.com>
Decide whether the accepted prefix fits the next verify on device, so CUDA graph replay keeps the same rule without a host-side branch. Co-authored-by: Cursor <cursoragent@cursor.com>
Record each accepted decode token with the cursor kernel so CUDA graph capture does not synchronize on request indices. Tune the update on a one-layer scratch, and let a real scratch allocation failure surface instead of being remembered as an OOM key. Co-authored-by: Cursor <cursoragent@cursor.com>
Explicit compute_mode=precompute now projects a width-1 verify step, while the inline default stays. Single-token acceptance and materialize share the existing cursor and snapshot paths. Co-authored-by: Cursor <cursoragent@cursor.com>
The replay branch should not change autotuner.py for whitespace alone. Co-authored-by: Cursor <cursoragent@cursor.com>
Settings, the last-layer check, index upload, and prefill kernel imports only wrapped a single call site. Co-authored-by: Cursor <cursoragent@cursor.com>
Readers now meet the decode kernel, autotune resolution, and accepted-tensor type before the helpers they call. Decode init still requires exactly one replay state. Co-authored-by: Cursor <cursoragent@cursor.com>
Replay reuses the existing dynamic verify helper, so this file does not need its kernel moved. Co-authored-by: Cursor <cursoragent@cursor.com>
The old fold file only advanced the cursor, while the state update lived under materialize. File and function names now match those steps. Co-authored-by: Cursor <cursoragent@cursor.com>
Use tf32x3 dots and an fp16 k/d cache for bf16 activations, and give KDA the same windowed reconstruction. Folding stays an fp32 pass over raw k/v, and history loads stay inside the prefix. Co-authored-by: Cursor <cursoragent@cursor.com>
The separate fold cleared the shared cursor while other layers were still reading it, so accepted history never reached the SSM and GSM8K trajectories collapsed. Record the write base during replay and advance the cursor only after every layer has finished. Co-authored-by: Cursor <cursoragent@cursor.com>
Decode and the separate fold kernel each walked the same history loop. One device function keeps those two paths on the same update. Co-authored-by: Cursor <cursoragent@cursor.com>
Acceptance and the decode kernel already keep the single prefix cursor. History setup stays in the constructor, and tests attach an existing SSM themselves. Co-authored-by: Cursor <cursoragent@cursor.com>
Decode reads the prefix cursor inside the kernel. Tests that need the length read it directly. Co-authored-by: Cursor <cursoragent@cursor.com>
Mode checks are asserts. The state stores the caller's config unchanged, and replay and fold supply missing BV, warps, and stages when the kernel starts. Co-authored-by: Cursor <cursoragent@cursor.com>
Page save, PD, and DP run before any replay history exists. They copy the accepted conv window and canonical SSM directly. Co-authored-by: Cursor <cursoragent@cursor.com>
Prefill writes the canonical SSM directly and leaves the cursor at zero. The snapshot helpers that only wrapped that empty prefix are gone. Co-authored-by: Cursor <cursoragent@cursor.com>
MTP verify trims the rejected tail from that cursor, so decode no longer records a per-step write base. Co-authored-by: Cursor <cursoragent@cursor.com>
Folding rewrites shared raw keys and gates at slot 0. Each launch row now waits until its value tiles have finished reading the old prefix, so a late tile cannot fold the new token into the SSM. Co-authored-by: Cursor <cursoragent@cursor.com>
Width-1 decode now builds its cumulative lengths while initializing indices. The last layer advances the cursor with that tensor instead of allocating one when the argument is missing. Co-authored-by: Cursor <cursoragent@cursor.com>
The index tensor is always allocated, and verification trims the replay cursor through commit_ssm_prefix. Co-authored-by: Cursor <cursoragent@cursor.com>
The state no longer wraps cursor updates or exposes a snapshot fold. Tests that need a read-only result fold a clone. Co-authored-by: Cursor <cursoragent@cursor.com>
…back Decode calls the replay kernel directly, and the external fold addresses canonical SSM with its real strides. Co-authored-by: Cursor <cursoragent@cursor.com>
The autotune cache key uses the same values under clearer field names, so old entries are left unused. Co-authored-by: Cursor <cursoragent@cursor.com>
Decode reads every request window from cu_seqlens, and q, k, v must already be [tokens, heads, dim]. Autotune times that same layout, so old cache entries are left unused. Co-authored-by: Cursor <cursoragent@cursor.com>
The shared tuner only forwarded both algorithms through one entry. Each replay kernel now owns its config search and returns only the attention output. Co-authored-by: Cursor <cursoragent@cursor.com>
A pinned BV and warp count is an input to the launch. Serving state no longer carries a flag that overrides the autotuner. Co-authored-by: Cursor <cursoragent@cursor.com>
Each algorithm keeps its own native and replay class tables, and the chunk-index host sync is marked for removal. Co-authored-by: Cursor <cursoragent@cursor.com>
The decode kernel receives the layer SSM, history, and cursor directly, so the operator no longer indexes the serving state. Co-authored-by: Cursor <cursoragent@cursor.com>
New calls use named arguments and type annotations, and production paths stay free of test-only APIs. Co-authored-by: Cursor <cursoragent@cursor.com>
Decode launches use the kernel parameter names, and a pinned tile config is an argument instead of a field on the test state. Co-authored-by: Cursor <cursoragent@cursor.com>
The helper is only called after the module loads, so it sits with the other startup checks instead of above the launch path. Co-authored-by: Cursor <cursoragent@cursor.com>
A restored slot only receives the canonical Conv and SSM, so its cursor must be cleared or the previous request's history is folded in. Deltas can exceed the fp16 range, so keep them in fp32 while normalized keys stay fp16 for bf16 activations. Co-authored-by: Cursor <cursoragent@cursor.com>
Grouped GDN configs read one RawKeys buffer from several value heads. The fold barrier now counts that whole key group, and the sync mark is shaped to the key heads. Co-authored-by: Cursor <cursoragent@cursor.com>
Neighboring programs now walk the V tiles and then the value heads of one request. Those heads sit in adjacent SSM slots and share the key-head fold barrier. Co-authored-by: Cursor <cursoragent@cursor.com>
hiworldwzj
force-pushed
the
support/replayssm-clean
branch
from
November 12, 2026 09:05
fe650a4 to
59adc05
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
为 GDN、KDA 增加 native、compact、replay SSM 状态模式，涵盖 Conv/SSM 状态管理、MTP 接受与回滚、checkpoint、PD/DP 状态交换、replay kernel 和 autotune。
恢复 checkpoint、PD 或 DP 状态时清零目标 replay cursor，避免复用请求槽时读入旧 history。delta history 使用 FP32，避免 BF16 输入的大 correction 写入 FP16 后溢出。当前 overflow 在 decode kernel 内融合 fold，通过共享 key 的 value heads 和 V tiles 计数同步，grid 为 (NV, HV, requests)。
最新提交 083ac58 仅修复测试：Sensenova 使用 get_layer_view() 检查实际状态映射，并在 gdn 调用处替换 CPU mock；动态窗口按 token 轴检查 padding；恢复 grouped-head 自动 rollover 的 16 组压力测试，手动 fold 对照使用与融合 kernel 相同的 tile/warp 配置。
基于 59adc05 加上述测试修复，在 H200 单卡完成 44 项针对性回归：Sensenova CPU/GPU 12 项、动态窗口 16 项、grouped-head rollover 16 项，全部通过。覆盖 GDN/KDA、inline/precompute、BV 32/128、verify width 1/3、batch 64、多层、重复 HOLD 和 CUDA Graph。Black、Flake8 和 diff 检查通过。
此前 625 项结果属于 a114f93 的独立 fold kernel 版本，不作为当前融合实现的验证结果。此次未改跨 CTA barrier 或 GDN precompute 的非有限值隔离逻辑；单卡用例通过不证明 barrier 对所有调度都有前进保证。完整服务请求生命周期、跨进程 PD/DP 传输及性能影响尚未验证。