[AutoDeploy] Re-enable SSM replay for Nemotron-Super MTP (replay kernel illegal memory access at CUDA-graph capture on Blackwell)
Speculative DecodingAutoDeployCUDA Graph
## Summary
In AutoDeploy, the SSM **replay** kernel used by Nemotron-Super MTP speculative decoding
crashes with an illegal memory access during **CUDA-graph capture** on Blackwell. As a
workaround, `ssm_replay` has been disabled in
`examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml` (see the fix for
nvbug 6316981). This issue tracks **fixing the replay kernel so `ssm_replay` can be
re-enabled** for Nemotron-Super (it provides the Mamba state "fast-forward" optimization
for the MTP extend path).
## Symptom
`RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered`, raised during
the cudagraph capture warmup forward (`torch_cudagraph.py::capture_graph`), i.e. at engine
build time — not during graph replay and not during eval.
## Affected configurations
| dtype | `ssm_replay: true` + cudagraph (capture) |
|-------|------------------------------------------|
| fp8 | OK |
| bf16 | **CRASH** (deterministic) |
| nvfp4 | **CRASH** |
Mamba state stays bf16 in all three checkpoints, so why fp8 survives while bf16/nvfp4 fail is
not yet understood — likely an element-size-sensitive path in the kernel (TMA descriptor /
shared-memory sizing).
## Faulting kernel
```
flashinfer_cached_ssm (tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py:245)
-> replay_selective_state_update (tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py:4508)
-> launch_persistent_main -> _persistent_main_kernel[grid] (:4232)
```
Reached via `insert_cached_ssm_attention: { backend: flashinfer_ssm, ssm_replay: true }`.
Pinpointed with `CUDA_LAUNCH_BLOCKING=1` (without it, the async error surfaces misleadingly in
the downstream Mamba gated RMSNorm kernel).
## Reproduction
On 4x Blackwell (e.g. B200):
```
LLM_MODELS_ROOT=<...> pytest "tests/integration/defs/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm]"
# or [nvfp4_ws4_80gb-trtllm]
```
Crashes at engine build (capture). With `ssm_replay: false` the same config builds, captures,
and generates correctly (verified: bf16 acceptance 65.91%).
## Current state / workaround
`ssm_replay: false` in `super_v3_mtp.yaml` keeps CUDA graphs + the overlap scheduler enabled
for all dtypes; it only drops the replay fast-forward optimization.
## Ask
Root-cause and fix the OOB in `_persistent_main_kernel` (replay SSM), then re-enable
`ssm_replay: true` for Nemotron-Super.
## Related
- nvbug 6316981 (Blackwell GB200/GB300), nvbug 6336682 (DGX_H100).
0 条评论