Gramnarayan/pr15063 mtp adp nan debug

[govind-ramnarayan]

@coderabbitai summary

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[govind-ramnarayan]

Here is the simple local diff for the current branch HEAD (8ef9470ad852) against its parent (HEAD^). The GitHub compare is noisy because of the rebase.

Validation from the rerun of tests/integration/defs/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm]:

• Result: PASSED in 467.00s
• ADP enabled: enable_attention_dp=True
• Main graph sharding: EP=40
• Draft graph sharding: EP=0, with EP reverted/replicated under attention-DP
• Torch cudagraph captured batch sizes: 8, 4, 2, 1
• GSM8K average accuracy: 92.68
• Evaluated accuracy: 92.684, threshold 89.407
• Spec-dec acceptance: 53.28% (16744/31428 tokens, 1001 speculative iterations)

diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml
index b46ff60e6f..04ec1fc316 100644
--- a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml
+++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml
@@ -21,6 +21,9 @@ attention_dp_config:
   batching_wait_iters: 10
   timeout_iters: 16
 transforms:
+  compile_model:
+    # Piecewise CUDA graph capture is disabled for MTP until MTP supports it.
+    piecewise_enabled: false
   detect_sharding:
     allreduce_strategy: NCCL
     enable_attention_dp: true
diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
index 950b908160..55be655918 100644
--- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
@@ -459,23 +459,16 @@ class PythonMambaCacheManager(BaseResourceManager):
         ssm_state_shape = (nheads, head_dim, d_state)
 
         # create mamba conv and ssm states
-        # Zero-initialize (not torch.empty): an un-prefilled cache slot (e.g. an
-        # attention-DP idle/dummy generation request that is decoded without a
-        # prefill) is read by the causal-conv / SSM ops before anything writes it.
-        # Uninitialized HBM can be NaN/Inf, which the causal-conv then propagates
-        # into the residual stream (and the SSM replay path latches into the
-        # persistent recurrent state, spreading it across ranks via the MoE
-        # all-to-all). Whether the garbage is NaN vs finite depends on the
-        # allocation layout, which is exactly why the failure only surfaced under
-        # {ssm_replay=true × attention-DP × MTP}. Zeros are a valid empty state.
-        conv_states = torch.zeros(
+        conv_states = torch.full(
             size=(num_local_layers, max_batch_size) + conv_state_shape,
+            fill_value=float("nan"),
             dtype=dtype,
             device=device,
         )
 
-        ssm_states = torch.zeros(
+        ssm_states = torch.full(
             size=(num_local_layers, max_batch_size) + ssm_state_shape,
+            fill_value=float("nan"),
             dtype=self.mamba_ssm_cache_dtype,
             device=device,
         )