ITADN

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

#15565Opengovind-ramnarayan 创建于 2026-06-24
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 条评论