[BBaditi]

Description

Adds an opt-in compact speculative replay mode to the com.microsoft::Mamba2Verify CUDA kernel (#33969), selected by the session config entry ep.cuda.mamba2verify_compact_replay (default "0"; only "0" and "1" are accepted).

With "1", in-place calls (ssm_snapshots aliases ssm_state, slots >= 2) with 1 <= R <= min(8, slots) keep the committed FP32 state in slot 0 and the window's (x, B, dt) in a per-head 32 KiB slice of slot 1, then replay the accepted rows in registers on the next call instead of writing R full FP32 snapshots every verify step (8 x 4 MiB x 40 layers = 1.25 GiB per step on Nemotron-3.5-Bolt).

Other in-place calls in this mode, including the large-R prefill from #34164, run the unchanged default kernels after a small per-head pre-pass. The pre-pass replays any pending window into slot state_row and clears the slot 1 header, so those kernels start from the same snapshot "0" would have left there. Callers that read or snapshot/restore intermediate or committed recurrent slots (for example prefix-cache recurrent-state capture/restore) must keep "0", because "1" keeps the committed base SSM state in slot 0 and pending step inputs in slot 1 rather than materializing slot state_row. A reset must rewrite the whole state tensors, not only slot 0 (onnxruntime-genai's NemotronBoltState::Reset() zeroes them). Nonaliased calls run the default kernels and do not read slot 1.

With "0" nothing changes: the 9 existing kernel instantiations compile to identical PTX under both the plugin and non-plugin build flags, and there is no schema change. The compact kernels pin the FMA pairing nvcc emits for the default recurrence, fma(st, dA, dtx * B), so "0" and "1" produce bitwise-identical outputs.

Testing

This branch is rebased onto main f4cd337911c2; none of its 6 files changed since d2f77cd96dfd. On the pushed commit (plugin build), onnxruntime_provider_test --gtest_filter=*FusedStateSpace*:*Mamba* passes 18/18. The gates below ran on main d2f77cd96dfd + this PR, NVIDIA A100-SXM4-80GB, driver 580.159.04, CUDA 13.0, cuDNN 9.12, -DCMAKE_CUDA_ARCHITECTURES=80, plugin and non-plugin CUDA EP builds.

• onnxruntime_provider_test --gtest_filter=*FusedStateSpace*:*Mamba*: 18/18 in both builds. New: MambaCompactReplayMultiRoundAndPrefillHandoffParity ("0" vs "1" bitwise at 4 and 64 heads: multi-round R=8 partial acceptance, R=37 prefill into the 8-slot bank handing off through state_row=7, a later R=37 chunk resuming from a pending window, nonaliased fallback, invalid selectors), MambaCompactReplayCudaGraphCaptureAndReplay, MambaInvalidCompactReplayPolicyRejected.
• FusedStateSpaceKernelTest: 7/7 through CUDA_EP_Unittest.All (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON) and standalone. New: MambaCompactReplayMatchesDefaultAcrossPrefillHandoffAndResets (64 heads, 8 and 16 slots, R in {1, 5, 7, 8, 9, 12, 16, 37} including a fold into slot 1, two passes with a full state reset in between; Y, xbc, readout and conv snapshots compared bitwise after every call, the full SSM state after every default-kernel call).
• That test fails when the pre-pass is removed or only clears the header: the first failure is the first R > 8 call after a pending window (8 slots: R=37, state_row=2; 16 slots: R=12, state_row=6). Earlier steps, including R=37 prefill followed by R=8 verify at state_row=7, pass either way.
• compute-sanitizer: memcheck 0 errors (18 provider tests, 7 kernel tests); racecheck 0 hazards, synccheck 0 errors, initcheck 0 errors on the compact tests.
• ptxas: 0 spill bytes for every kernel; compact kernels 64-66 registers, pre-pass 60.
• pytest onnxruntime/test/python/transformers/test_gguf_packed_ops.py (plugin build): 50 passed.
• clang-format 20.1.8 and git diff --check clean.

Performance

Kernel microbenchmark, A100-SXM4-80GB, Nemotron-3.5-Bolt geometry (128 heads, 8 groups, 8 slots), 8 rotating state banks (cold L2), 200 interleaved AB/BA pairs, per layer (medians; ratio of means with bootstrap 95% CI):

Call	"0"	"1"	Ratio
R=8 verify, CUDA Graph, state_row 0	43.01 us	28.54 us	1.532x [1.523, 1.541]
R=8 verify, CUDA Graph, state_row 3	37.25 us	26.62 us	1.399x [1.390, 1.407]
R=8 verify, CUDA Graph, state_row 7	37.50 us	30.46 us	1.230x [1.225, 1.234]
R=8 verify, eager, state_row 3	40.96 us	30.08 us	1.346x [1.306, 1.379]
R=37 / 512 / 4096 chunk, no pending window	84.86 / 535.42 / 4099.39 us	88.06 / 538.88 / 4102.53 us	pre-pass +3.2 / +3.5 / +3.1 us
Same, first chunk after a compact window	84.86 / 535.94 / 4099.78 us	97.28 / 547.46 / 4113.54 us	+12.4 / +11.5 / +13.8 us

The same binary checks Bolt-geometry parity: 13/13 calls bitwise identical over an R = 4096, 8, 37, 1, 512 sequence.

End-to-end on this branch (main f4cd337911c2 + this PR, plugin CUDA EP) with onnxruntime-genai 75fa2250c9af: fixed [1, 8] 74B package with compact replay set only through genai_config.json, 100 MMLU-Pro prompts (greedy, up to 1024 new tokens) in interleaved ABBA blocks on one A100. 100/100 identical token sequences and speculative counters, decode 193.42 -> 198.60 tok/s (+2.73%, 95% CI [+2.68%, +2.77%]), mean verify plus draft round latency -2.65%. Setting "2" the same way fails model creation with this kernel's error, so the entry reaches the kernel.

Not covered

B > 1, GPUs other than A100, multi-turn continuation through onnxruntime-genai, and banks with fewer slots than rows (R <= 8 but R > slots), which take the pre-pass plus default path and are not separately tested.

[Copilot]

🟡 Changes recommended

A critical metadata race can send warps down inconsistent barrier paths and hang the kernel; it needs an unconditional block barrier after metadata reads and a regression test for the compact-path rejection case.

1 open issue

• Synchronize metadata reads before the compact-path rejection branch · New

What changed in this PR

Adds opt-in compact speculative replay to ONNX Runtime’s CUDA Mamba2Verify operator, reducing recurrent-state writes while preserving the default mode and supporting large-prefill handoffs.

Changes:

• Adds a validated session option and compact replay kernels with a pre-pass for fallback calls.
• Adds parity, reset, invalid-selector, and CUDA Graph tests.
• Documents compact-state layout and caller restrictions.

File	Description
onnxruntime/​test/​contrib_ops/​fused_state_space_op_test.cc	Adds session-level parity, option validation, and graph-capture tests.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​fused_state_space_test.cc	Tests replay, prefill handoffs, resets, and invalid selectors.
onnxruntime/​contrib_ops/​cuda/​sequence/​fused_state_space.cc	Validates the session option and selects the write mode.
onnxruntime/​contrib_ops/​cuda/​sequence/​fused_state_space_impl.h	Defines replay modes and extends the launch interface.
onnxruntime/​contrib_ops/​cuda/​sequence/​fused_state_space_impl.cu	Implements compact replay storage, kernels, and fallback resolution.
docs/​contrib_ops/​cuda/​mamba2_verify.md	Documents configuration, state semantics, and usage restrictions.

_{Review effort: Balanced}