[https://nvbugs/6210714][fix] Fix mamba block calculation - #14524
Conversation
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
VALLIS-NERIA
changed the title
📝 WalkthroughWalkthroughConfiguration schema change makes ChangesMamba cache validation and sizing
Sequence DiagramsequenceDiagram
participant User as User Config
participant Schema as KvCacheConfig
participant Validator as validate_hybrid_cache_config
participant Manager as CppMambaHybridCacheManager
participant Estimator as _calculate_max_num_blocks_for_linear_attention
User->>Schema: mamba_state_cache_interval (optional)
Schema-->>Manager: pass config
Manager->>Validator: validate_hybrid_cache_config(config, tokens_per_block)
Validator-->>Manager: assert interval > 0 and aligned
Manager-->>Manager: initialize successfully
Manager->>Estimator: estimate resource requirements
Estimator->>Estimator: max_snapshots = live_state_slots + snapshot_count
Estimator->>Estimator: max_tokens = max((budget - intercept) // slope, 0)
Estimator-->>Manager: resource estimates
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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
🤖 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/resource_manager.py`:
- Around line 688-690: The estimation path and the runtime allocation both
compute snapshot slot counts differently (variables like max_snapshots,
kv_cache_config.max_tokens, linear_attention_metadata.states_snapshot_interval
and the affine token-budget model used at the other site), causing mismatched
pool sizing and possible under-allocation; refactor by extracting a single
helper (e.g., compute_snapshot_slots or derive_snapshot_slot_count) that
implements the affine formula (intercept + slope * T) and any consistent
+1/draft-slot policy, then replace the ad-hoc additive calculation that updates
max_snapshots and the separate max(...) logic in the runtime allocation with
calls to this helper so both dry-run estimate and actual pool shape are derived
from the same formula.
🪄 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: a9e214f7-5c18-4f3a-a3a3-ddaacc54feaa
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/llmapi/llm_args.py
…x (By Agent) Covers the three pieces of the fix: - validate_hybrid_cache_config: pure-Python guard. New parametric tests exercise the no-op path (reuse disabled), the happy path (positive, multiple of tokens_per_block), and the four reject branches (unset, zero, negative, non-multiple). - Dry-run recurrent-state pool sizing under block reuse: new test test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse asserts the additive live + reuse formula (was previously a max()). - Existing test_cpp_hybrid_recurrent_pool_floor_with_block_reuse now passes mamba_state_cache_interval=256 explicitly, since the default changed from 256 to None. _build_hybrid_with_mamba_layer gains mamba_state_cache_interval and is_estimating_kv_cache parameters and forwards is_estimating_kv_cache to the manager so the dry-run path is reachable from tests. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
|
PR_Github #50204 [ run ] triggered by Bot. Commit: |
|
PR_Github #50204 [ run ] completed with state
|
|
/bot run |
|
PR_Github #50250 [ run ] triggered by Bot. Commit: |
|
PR_Github #50250 [ run ] completed with state
|
|
/bot run |
|
PR_Github #50312 [ run ] triggered by Bot. Commit: |
|
PR_Github #50312 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #50421 [ run ] triggered by Bot. Commit: |
|
PR_Github #50421 [ run ] completed with state |
|
/bot run --help |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
PR_Github #50802 Bot args parsing error: usage: /bot [-h] |
Signed-off-by: xiweny <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #52013 [ run ] triggered by Bot. Commit: |
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
…eny/6210714 Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #52033 [ run ] triggered by Bot. Commit: |
|
PR_Github #52013 [ run ] completed with state |
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #52036 [ run ] triggered by Bot. Commit: |
|
PR_Github #52033 [ run ] completed with state |
|
PR_Github #52036 [ run ] completed with state
|
…gle-GPU test_single_request_chat_multiple_images[pd_disagg-qwen3_30b_a3b_fp8] in test_mm_encoder_standalone.py fails at setup on the pre-merge DGX_B200 single-GPU stage with a NIXL CacheTransceiver init assertion (status == NIXL_SUCCESS, transferAgent.cpp:614) when the pd_disagg LLM constructs its disaggregated KV-cache transfer agent. This is a fleet-wide failure on the single-GPU pre-merge stage, observed across many unrelated PRs (e.g. NVIDIA#13978, NVIDIA#13925, NVIDIA#14841, NVIDIA#14599, NVIDIA#14524, NVIDIA#14941, NVIDIA#14398); it passes only intermittently (~1/3) depending on node. Waiving until the single-GPU NIXL EPD-disagg path is fixed or the variant is gated to multi-GPU. NVBug: https://nvbugs/6269683 Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #52227 [ run ] triggered by Bot. Commit: |
|
PR_Github #52227 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52287 [ run ] triggered by Bot. Commit: |
|
PR_Github #52287 [ run ] completed with state |
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Signed-off-by: xiweny <13230610+VALLIS-NERIA@users.noreply.github.com> Signed-off-by: NVFB <186336021+NVFB@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Signed-off-by: xiweny <13230610+VALLIS-NERIA@users.noreply.github.com>
4 participants
Summary
Fixes mamba/attention hybrid pool sizing under nvbug 6210714.
Three related fixes:
resource_manager.pybudget formula — Remove the heuristic thatignored the
interceptterm whenevermamba_slope > 0(i.e. wheneverblock reuse + snapshots were enabled). The intercept captures the
fixed per-rank live-state memory
max_batch_size * pp_size * state_bytes_localand is not negligible under PP or largemax_batch_size; dropping it produced budgets larger than what thepool could actually hold, leading to OOM / undersized allocations.
Now always subtract
interceptfrom the primary budget, and warnwhen a user-specified
max_tokenscannot fit the resulting budget.resource_manager.pydry-run snapshot count — Change theestimation-dry-run path so that
max_snapshots = live_state_slots + reuse_snapshotsinstead ofmax(reuse_snapshots, live_state_slots).This mirrors the non-dry-run path (live slots and reuse snapshots
are additive), so dry-run estimates no longer undershoot the slot
count actually needed.
Hybrid cache config validation — Add
validate_hybrid_cache_configtomamba_cache_manager.pyand callit from
CppMambaHybridCacheManager.__init__. When block reuse isenabled, it asserts
mamba_state_cache_interval > 0and that it isa multiple of
tokens_per_block(required for proper blockalignment). Also relaxes the
KvCacheConfig.mamba_state_cache_intervalfield type from
PositiveInt(default 256) toOptional[int](default
None) so the validator can distinguish "unset" from "setto a specific value".
Test plan
and block reuse enabled
max_tokenslarger than the computed budgettriggers the new warning rather than silently over-allocating
enable_block_reuse=Truewith an invalidmamba_state_cache_intervalraises the new assertiontests/unittest/for mamba / hybridcache still pass
Summary by CodeRabbit
Bug Fixes
Changes