[https://nvbugs/6261164][fix] AutoDeploy: Don't allocate speculative caches when speculation is off

[tensorrt-cicd]

fixes: #15120

Summary

• Root cause: SSM/causal-conv backends unconditionally register intermediate speculative-decoding state handlers; with spec_config is None they're never bound by the Mamba cache manager and fall through to a full per-layer unmanaged allocation (~650 MiB each), OOMing the 44 GiB L40S.
• Fix: In the kvcache insert transform (_InsertCachedOperator._apply), when cm._spec_config is None, drop SpecSSMResourceHandler/SpecCausalConvResourceHandler to the existing None sentinel so no buffer is registered/allocated; removed the waiver.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6261164

Summary by CodeRabbit

• New Features

  • KV-cache now suppresses speculative-only resources when speculative decoding is disabled, reducing unnecessary allocations.

• Bug Fixes

  • Fixed unnecessary memory buffer allocation in KV-cache handling for standard inference.
  • Restored accuracy test for Nemotron V2 FP8 (waiver removed).

• Tests

  • Added validation tests ensuring replay-only SSM requires speculative decoding and unit tests for speculative-resource gating.

[coderabbitai]

📝 Walkthrough

Walkthrough

Add SpeculativeOnly trait and new Intermediate* speculative handler types; update backends and shim to use them; KV-cache transform now converts speculative-only handlers to None when speculative decoding is disabled; add validator enforcing ssm_replay requires speculative config; update tests and remove one waiver.

Changes

Speculative resource handler allocation fix

Layer / File(s)	Summary
Spec trait and intermediate speculative handlers tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py	Add SpeculativeOnly mixin, introduce IntermediateSSMStateHandler and IntermediateConvStateHandler (with from_base), and mark replay-related handlers as speculative-only.
Backend imports and cache initializers tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/..., tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py	Replace imports of prior speculative handlers with Intermediate* handlers and use Intermediate*.from_base(...) for intermediate-state cache initializers.
Shim: identify/manage intermediate speculative resources tensorrt_llm/_torch/auto_deploy/shim/interface.py	Update imports and _identify_managed_state_resources to construct/check IntermediateSSMStateHandler/IntermediateConvStateHandler instead of old Spec* handlers; adjust cache-statistics checks.
KV-cache transform: suppress speculative handlers when no spec tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py	Import ResourceHandler and SpeculativeOnly, add _suppress_speculative_when_no_spec helper, remove redundant local import, and apply helper during pass-1 registration so speculative-only handlers become None when cm._spec_config is None.
Validator and tests tensorrt_llm/_torch/auto_deploy/llm_args.py, tests/*	Add LlmArgs.validate_ssm_replay_requires_spec; update unit tests to use Intermediate* handlers and add tests that assert KV-cache speculative gating behavior; remove one waived integration test entry.

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#13725: Touches SSM replay/speculative handler paths used by cached attention and relates to replay buffer/handler usage.

Suggested labels

KV-Cache Management

Suggested reviewers

• hnover-nv
• govind-ramnarayan
• galagam
• suyoggupta

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description clearly explains the root cause, fix, test coverage, and includes required metadata (bug link, GitHub issue reference).
Title check	✅ Passed	The title directly describes the main fix: preventing speculative cache allocation when speculation is disabled, which aligns with the core change in the KV-cache transform.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[suyoggupta]

@hnover-nv , @govind-ramnarayan . PTAL. Would it make sense to add a unit test too?

[govind-ramnarayan]

@hnover-nv , @govind-ramnarayan . PTAL. Would it make sense to add a unit test too?

I have a slightly expanded version of this that I'll push to the branch. The main issue with this is that it only skips over assigning the (default) "Rewind" speculative buffers, but should ideally also skip over the "Replay" ones too. When speculation is off but "replay" is on in the FlashInfer kernel, which is possible right now, these will fall through unmanaged.

[galagam]

LGTM

[govind-ramnarayan]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52837 [ run ] triggered by Bot. Commit: e034795 Link to invocation

[govind-ramnarayan]

Additionally verified this test now passes on L40s.

[tensorrt-cicd]

PR_Github #52837 [ run ] completed with state SUCCESS. Commit: e034795
/LLM/main/L0_MergeRequest_PR pipeline #42090 completed with status: 'SUCCESS'

CI Report

Link to invocation