[https://nvbugs/6479324][fix] Keep KV dtype for hybrid recurrent-states pool to unstall disagg transfers

[brnguyen2]

Description

Fixes the deterministic generation-worker KV-transfer stall in the disaggregated stress tests on hybrid (Mamba/linear-attention) models, e.g. test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_5_4b_fp8_stress] (nvbug 6479324; same failure class as 6472256).

Root cause

Since #16304, CppMambaHybridCacheManager labels the recurrent-states pool with the SSM state dtype (e.g. BF16) while the KV pools stay FP8. The recurrent pool's sentinel window sorts first, making it pool 0, and the disagg KV wire moves every non-scale pool's blocks — including recurrent-states blocks — with a single canonical dtype taken from pool 0:

1. The context-side split kernel hits outputSplitBlock data type mismatch expected: 6, actual: 7 (cacheSplitConcat.cu) mid-transfer; the sender swallows the exception inside sendAndRemoveResponse.
2. The generation side then waits forever in CacheReceiver::receiveReadySignalDetailed for a ready signal that never comes.
3. The gen worker wedges with num_scheduled_requests=0 while holding max_batch_size+warmup requests; the disagg server returns ~3000 HTTP 500s, and the accuracy phase can never complete.

The existing differing-dtype guard in CacheTransBufferManager was supposed to reject this at startup, but it iterated a compacted pool count (getNumPools(false,false)) while indexing the absolute pool domain (getPrimaryPool(i)), making it vacuous for hybrid models.

The stall reproduces without request cancellation and with both UCX and NIXL backends — it is a property of the common transfer path, bisected to 999618e (#16304): parent passes 2/2, 999618e(+its own import hotfix #16401) stalls, all later commits stall.

Fix

• mamba_cache_manager.py: only override the recurrent-states pool dtype when the KV dtype is NVFP4 — the case [https://nvbugs/6374873][fix] Allow fp4 KV Cache + non-FP4 Mamba State #16304 actually addressed (FP4 container/scale semantics corrupting the byte-blob state). Other KV dtypes keep the manager dtype, restoring the pre-[https://nvbugs/6374873][fix] Allow fp4 KV Cache + non-FP4 Mamba State #16304 wire behavior.
• cacheTransBuffer.cpp/.h, cacheFormatter.cpp: derive disagg transfer-buffer sizing and the wire dtype from the first real KV pool (kvTransportPoolIdx) instead of pool 0.
• Guard fix: the differing-dtype check now iterates the absolute pool index domain and fails loudly at construction (turning the unsupported NVFP4-KV-hybrid + disagg combination into a startup error instead of a silent runtime stall).

Validation (2×B200, 1 ctx + 1 gen, TP1)

Build	Result
unfixed (any commit ≥ 999618e)	stall: gen pinned at 0/138 reqs, ctx spams 60 s KV-transfer timeouts, ~3000 HTTP 500s, test fails after ~46 min
fixed (this PR on main tip)	PASSES in ~9 min incl. accuracy phase
fixed (backported onto bcf5a20)	PASSES in ~5 min

New tests: C++ CacheTransBufferTest.TestHybridRecurrentStatesPool (canonical-pool selection + loud failure on dtype mismatch), Python test_cpp_hybrid_keeps_kv_dtype_for_recurrent_pool_on_non_nvfp4_kv_cache. Existing NVFP4 hybrid tests from #16304 still pass (4/4 locally).

@VALLIS-NERIA FYI — this preserves #16304's NVFP4 behavior; longer-term the TODO(disagg-multi-dtype) per-pool dispatch would let the SSM dtype differ on the wire too.

Test Coverage

• tests/unittest/_torch/executor/test_mamba_cache_manager.py (new + existing pool-dtype tests)
• cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp::TestHybridRecurrentStatesPool
• tests/integration/defs/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_5_4b_fp8_stress]

🤖 Generated with Claude Code

[VALLIS-NERIA]

The fix looks reasonable to me.

[VALLIS-NERIA]

The fix is supposed to be covered by #16505