[None][fix] Keep hybrid KV routing for Skip Softmax

[bobboli]

Summary

1. Keep hybrid Mamba / linear-attention models on the hybrid KV cache-manager route when sparse_attention_config.algorithm is skip_softmax.
2. Preserve sparse KV cache-manager routing for non-hybrid models.
3. Fail early for sparse attention algorithms that require their own KV cache manager when combined with hybrid models.

Skip Softmax only changes attention kernel execution. It should not replace the Mamba-capable cache manager needed for recurrent state.

Testing

• pre-commit run --files tensorrt_llm/_torch/pyexecutor/_util.py
• python3 -m py_compile tensorrt_llm/_torch/pyexecutor/_util.py

Local pytest was not run because this local environment does not have pytest / torch installed.

Summary by CodeRabbit

• Bug Fixes
  • Fixed KV-cache manager selection for models using sparse attention with hybrid linear configurations to prevent errors and ensure proper cache handling.

[coderabbitai]

📝 Walkthrough

Walkthrough

get_kv_cache_manager_cls in _util.py is modified so that hybrid-linear model detection takes priority over sparse attention config. When both conditions are present, the sparse attention algorithm is validated and a ValueError is raised for anything other than "skip_softmax". The sparse-attention return path is converted from if to elif to prevent it from firing for hybrid-linear models.

Changes

KV Cache Manager Routing Fix

Layer / File(s)	Summary
Hybrid-linear priority and sparse-attention guard tensorrt_llm/_torch/pyexecutor/_util.py	Hybrid-linear check now runs before the sparse-attention check. When sparse_attention_config is present on a hybrid-linear model, algorithm is read and a ValueError is raised unless it equals "skip_softmax". The sparse-attention return path is changed to elif sparse_attn_config is not None so it applies only to non-hybrid-linear models.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~5 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: keeping hybrid KV routing for Skip Softmax algorithm, which matches the primary fix in the changeset.
Description check	✅ Passed	The description includes a clear summary of changes, testing performed, and acknowledges environment limitations; however, it lacks the structured PR template sections and doesn't formally address the PR checklist items.
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.

_{✏️ 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.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 94-99: The sparse attention algorithm validation check (that
raises ValueError for non-skip_softmax algorithms with hybrid Mamba models) is
executing before the zero-mamba fallback logic, causing degenerate hybrid
configs with zero mamba layers and algorithms like "rocket" or "dsa" to
incorrectly raise ValueError instead of being converted to non-hybrid mode. Move
the sparse attention validation block that checks sparse_attn_config and
sparse_attn_algorithm to after the zero-mamba fallback handling (currently at
line 104), so that configs with zero mamba layers can be converted to non-hybrid
mode first, and then the validation will apply only to genuine hybrid
configurations.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 79766918-b75e-453e-9969-a9d65f06e00b

📥 Commits

Reviewing files that changed from the base of the PR and between 72ecb98 and 0e9d89c.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/pyexecutor/_util.py

[bobboli]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54746 [ run ] triggered by Bot. Commit: 1d92e23 Link to invocation

[tensorrt-cicd]

PR_Github #54746 [ run ] completed with state SUCCESS. Commit: 1d92e23
/LLM/main/L0_MergeRequest_PR pipeline #43765 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[bobboli]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54905 [ run ] triggered by Bot. Commit: 88aa079 Link to invocation

[tensorrt-cicd]

PR_Github #54905 [ run ] completed with state SUCCESS. Commit: 88aa079
/LLM/main/L0_MergeRequest_PR pipeline #43909 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[bobboli]

/bot run --reuse-test

[tensorrt-cicd]

PR_Github #54930 [ run ] triggered by Bot. Commit: 88aa079 Link to invocation

[tensorrt-cicd]

PR_Github #54930 [ run ] completed with state SUCCESS. Commit: 88aa079
/LLM/main/L0_MergeRequest_PR pipeline #43932 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[bobboli]

Folded into #15417.