[None][fix] SA spec dec: promote accepted hybrid recurrent states in-worker#16759
[None][fix] SA spec dec: promote accepted hybrid recurrent states in-worker#16759brnguyen2 wants to merge 1 commit into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughSA speculative decoding now conditionally promotes accepted-token recurrent state for Mamba hybrid cache managers. Regression tests cover generation batches, context-only batches, non-hybrid managers, and Hopper pre-merge test-list integration. ChangesSA Hybrid State Promotion
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant SAWorker
participant MambaHybridCacheManager
participant LiveStatePools
SAWorker->>MambaHybridCacheManager: update_mamba_states(attn_metadata, num_accepted_tokens, state_indices)
MambaHybridCacheManager->>LiveStatePools: promote accepted recurrent state
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py (1)
36-58: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new test functions and helpers.
_make_metadata,_run_forward, and all three test functions are unannotated. Add precise parameter and return annotations, including-> Nonefor tests.As per coding guidelines, every function must be annotated.
Also applies to: 61-93
🤖 Prompt for 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. In `@tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py` around lines 36 - 58, Add complete parameter and return type annotations to _make_metadata and _run_forward, using appropriate types for metadata, worker, and returned values. Annotate all three test functions in this file with their precise parameter types and -> None return types, without changing test behavior.Source: Coding guidelines
🤖 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 `@tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py`:
- Around line 61-93: The tests lack coverage for mixed batches and the new test
module is not registered in test lists. Add a test alongside
test_hybrid_manager_promotes_accepted_states and
test_hybrid_manager_context_only_batch_skips_promotion using 0 < num_contexts <
num_seqs, asserting accepted-state promotion occurs appropriately; also register
test_sa_hybrid_state_promotion.py in the relevant
tests/integration/test_lists/test-db or qa list(s).
---
Nitpick comments:
In `@tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py`:
- Around line 36-58: Add complete parameter and return type annotations to
_make_metadata and _run_forward, using appropriate types for metadata, worker,
and returned values. Annotate all three test functions in this file with their
precise parameter types and -> None return types, without changing test
behavior.
🪄 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: 7420fe6f-beda-4871-9ca1-f84d916def05
📒 Files selected for processing (2)
tensorrt_llm/_torch/speculative/sa_worker.pytests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py
09bf8a8 to
9b336f8
Compare
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
SA spec-dec accepted-state promotion is confined to the SA path (SAWorker only exists in SA mode, early-returns when drafting is off, gated on generation rows + hybrid manager); normal decode unaffected, well-tested. LGTM.
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #61289 [ run ] triggered by Bot. Commit: |
|
PR_Github #61290 [ run ] triggered by Bot. Commit: |
|
PR_Github #61289 [ run ] completed with state |
|
PR_Github #61290 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61373 [ run ] triggered by Bot. Commit: |
|
PR_Github #61373 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61459 [ run ] triggered by Bot. Commit: |
|
PR_Github #61459 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61498 [ run ] triggered by Bot. Commit: |
|
PR_Github #61498 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61581 [ run ] triggered by Bot. Commit: |
|
PR_Github #61581 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61601 [ run ] triggered by Bot. Commit: |
|
PR_Github #61601 [ run ] completed with state
|
…worker SAWorker never committed accepted verification-step states for SSM/hybrid models: verification writes per-step states to the cache manager's SpeculativeState scratch buffers (live pools are read-only during verify), and nothing promoted the accepted step — silent state corruption on any hybrid model running standalone SA. Add the update_mamba_states call right after on-device acceptance, mirroring the other one-engine workers (dflash/eagle3); isinstance-gated so pure-attention models are unaffected. Add a mock-based unit test pinning the contract: promotion fires for hybrid cache managers on generation batches, and is skipped for context-only batches and pure-attention managers. Verified to fail against the pre-fix worker. Listed in the H100 L0 pre-merge set. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
b665066 to
b6dbcce
Compare
|
PR_Github #61629 [ run ] triggered by Bot. Commit: |
|
PR_Github #61629 [ run ] completed with state
|
6 participants
Description
SAWorker(standalone suffix-automaton speculative decoding) never commits accepted verification-step recurrent states for SSM/hybrid models. Verification writes per-step states into the cache manager'sSpeculativeStatescratch buffers — the live pools are read-only during verify, since a rejected draft must not corrupt them — but nothing promotes the accepted step back into the live pools afterwards. The result is silent recurrent-state corruption on any hybrid (Mamba/GDN/linear-attention) model running standalone SA spec-dec.Fix: call
update_mamba_statesright after on-device acceptance, at the same call site the other one-engine workers (dflash, eagle3, mtp_dynamic_tree) already use. The call is gated onisinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager)andnum_gens > 0, so pure-attention models and context-only batches are unaffected.Test Coverage
New
tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py(mock-based, no GPU/model needed, <1s):update_mamba_statescalled exactly once with the sampler'snum_accepted_tokensand the metadata'sstate_indicesThe promotion test fails against the pre-fix worker and passes with the fix. Added to the H100 L0 pre-merge list (
l0_h100.yml) alongside the other speculative unit tests.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.
Dev Engineer Review
tensorrt_llm/_torch/speculative/sa_worker.pyto importMambaHybridCacheManagerand, inSAWorker._forward_impl, callattn_metadata.kv_cache_manager.update_mamba_states(...)immediately after_sample_and_accept_draft_tokens(i.e., after accepted/verified tokens are determined).num_gens = batch_size - attn_metadata.num_contextsandnum_gens > 0, plusisinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager). Context-only batches (num_contexts > 0) and pure-attention (non-MambaHybridCacheManager) managers skip promotion.attn_metadata,num_accepted_tokens, andstate_indices=attn_metadata.mamba_metadata.state_indices.QA Engineer Review
tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.pywith 3 new tests:test_hybrid_manager_promotes_accepted_states: assertsupdate_mamba_statesis called once withattn_metadata,num_accepted_tokens, andstate_indices, and thatresult["new_tokens"]returns the mocked accepted tokens.test_hybrid_manager_context_only_batch_skips_promotion: assertsupdate_mamba_statesis not called whennum_contexts > 0.test_pure_attention_manager_skips_promotion: assertsupdate_mamba_statesis not called when the KV cache manager is not aMambaHybridCacheManager.unittest/_torch/speculative/test_sa_hybrid_state_promotion.pytotests/integration/test_lists/test-db/l0_h100.ymlunder thel0_h100list.