[AutoDeploy] Re-enable SSM replay for Nemotron-Super MTP (replay kernel illegal memory access at CUDA-graph capture on Blackwell)

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

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).