[https://nvbugs/6465993][fix] use attention cache dtype for disaggregated transfer by chienchunhung · Pull Request #16505 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-07-16T16:20:17Z

Description

Fix disaggregated KV-cache transport for hybrid recurrent/attention models and remove the three NVBug 6465993 Nemotron 3 Super waivers.

Root cause

CppMambaHybridCacheManager keeps separate HALF recurrent-state and FP8 attention KV pools. CacheTransBufferManager selected its wire dtype and size through layer-indexed getPrimaryPool(0), which resolved to the recurrent pool while CacheFormatter transferred attention blocks. The resulting HALF transport buffer versus FP8 payload mismatch failed on the first request and surfaced as a timeout. Controls reproduced the failure with session reuse disabled and with the worker monitor bypassed, ruling both out.

Regression history

The layer-0 transport-pool assumption predated heterogeneous hybrid cache pools and remained dormant while all pools effectively used the same dtype.
#15777 added pytest MPI session reuse before the mixed-pool change; disabling reuse still reproduced the failure, so it did not create the mismatch.
#16304 correctly enabled recurrent-state and attention pools to use their actual distinct dtypes, exposing the dormant transport bug. Its CI runs #47492 and #47584 did not dispatch the affected x86 multi-GPU B200 stages, so the regression was not exercised on that PR.
#16338's worker-monitor changes landed shortly afterward, but bypassing the monitor did not prevent the timeout.
Full CI #48086 on #16444 first exercised all three affected stages and exposed the timeouts. #16495 then added the three NVBug 6465993 waivers, causing later runs to skip them until this PR removed the waivers.

Fix

Select payload pools from the actual BlockManager pools, excluding block-scale and indexer pools.
Prefer attention-pool dtype and geometry, skip recurrent layers handled by the RNN transfer path, and retain a pure-recurrent fallback.
Reject mixed attention-pool dtypes and make the formatter use the transfer manager's selected dtype.
Add hybrid-cache regression coverage and remove all three NVBug 6465993 waivers.

Verification

Local formatting, conflict, duplicate-waiver, and AST test-list validation passed; iterative TRT-LLM self-review has no remaining findings.
All three formerly waived tests passed on fixed commit f0f316bd:
- test_ctx_dp2_gen_tp4 — passed in DGX_B200-8_GPUs-PyTorch-1 (#48500).
- test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=False] — passed in DGX_B200-8_GPUs-PyTorch-2 (#48500).
- test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False] — passed in DGX_B200-8_GPUs-PyTorch-4 (#48486); it also passed independently on the preceding fixed head in #48476.
Full CI #48529 succeeded on f0f316bd.
Rebasing preserved the two PR patches without semantic changes; the current head is 94ce31a9. Subsequent full runs did not reach the B200 multi-GPU coverage: #48566 was blocked by the upstream DSpark regression repaired in #16579, while #48717 and #48756 were blocked by unrelated CI/baseline instability before multi-GPU dispatch. No reported failure implicated this PR's code.

Dev Engineer Review

KV-cache transport fix (disaggregated hybrid recurrent/attention): Updated CacheTransBufferManager to derive the wire/transfer DataType from actual KV-cache pool metadata (preferring attention-pool dtype when present, enforcing consistent attention-pool dtype, and using indexer-K pool dtype when transferIndexerKCache is enabled).
Dtype selection alignment (formatter ⇄ transfer manager): Updated CacheFormatter::unformat to allocate its temporary receive buffer using mCacheTransBufferManager->getDataType() rather than deriving dtype from the source cache’s pool-0.
Geometry/sizing correctness: In CacheTransBufferManager::computeTransferBufferSize, corrected non-indexer-K sizing to compute the byte-per-token-per-layer term per layer from that layer’s pool primary dimensions (instead of using a single pool-0-derived value). Indexer-K sizing behavior remains indexer-K dimension-based.
Attention-pool aware handling: Added logic to detect whether attention-cache pools exist and use that to gate certain linear-attention window skipping behavior.
Error handling / guards: Strengthened construction-time validation so CacheTransBufferManager throws tensorrt_llm::common::TllmException when eligible attention-pool dtypes are mixed.
API consistency: Added CacheTransBufferManager::getDataType() accessor to expose the selected transfer dtype for downstream buffer allocations.

QA Engineer Review

Test changes (unit tests):
- Added: TEST_F(KVCacheManagerTest, HybridDisaggUsesAttentionPoolDtype)
  - Verifies hybrid disaggregated behavior where recurrent window dtype (FP16) differs from attention window dtype (FP8), and asserts the transfer wire dtype comes from the attention pool (CacheTransBufferManager::getDataType(), getSendBuffer(...), and consistent buffer index handling via freeBufferIndexForSend).
- Updated: TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard)
  - Expanded the eligible-pool scan to detect mixed per-window dtypes (excluding block-scale and indexer-K cache pools) and assert constructor-time exception on mismatch.
Coverage via tests/integration/test_lists/:
- Modified: tests/integration/test_lists/waives.txt
  - Kept only accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] (nvbugs/6478726)
  - Removed other waived test_auto_dtype variants for TestNemotron3Super120B (including the use_py_transceiver=False and mtp_nextn=3-block_reuse=True-use_py_transceiver=False cases) and removed the separate accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4 waive entry (all tied to nvbugs/6465993).
Verdict: needs follow-up.

chienchunhung · 2026-07-16T16:20:26Z

/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-1, DGX_B200-8_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-4"

chienchunhung · 2026-07-16T16:27:10Z

/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-1, DGX_B200-8_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-4"

tensorrt-cicd · 2026-07-16T16:33:33Z

PR_Github #59749 [ run ] triggered by Bot. Commit: 8c2dec6 Link to invocation

tensorrt-cicd · 2026-07-16T18:49:10Z

PR_Github #59749 [ run ] completed with state SUCCESS. Commit: 8c2dec6
/LLM/main/L0_MergeRequest_PR pipeline #48175 (Partly Tested) 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

Uh oh!

Conversation

chienchunhung commented Jul 16, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Root cause

Regression history

Fix

Verification

Dev Engineer Review

QA Engineer Review

Uh oh!

chienchunhung commented Jul 16, 2026

Uh oh!

chienchunhung commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

chienchunhung commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

chienchunhung commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

chienchunhung commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

chienchunhung commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

chienchunhung commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

chienchunhung commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 18, 2026

Uh oh!

chienchunhung commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 18, 2026

Uh oh!

chienchunhung commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 18, 2026

Uh oh!

chienchunhung commented Jul 18, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

xinhe-nv commented Jul 23, 2026

Uh oh!

coderabbitai Bot commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

xinhe-nv commented Jul 23, 2026

chienchunhung commented Jul 16, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 23, 2026 •

edited

Loading