[None][feat] Fit KVCM V2 initialization constraints to GPU quota

[yizhang-nv]

@coderabbitai summary

Description

When estimation supplies avg_seq_len, the long-sequence initialization constraint can silently increase KVCM V2's GPU quota. The allocation may succeed while leaving too little memory for profiling.

Add KVCacheDesc.constraint_policy=ConstraintPolicy.FIT_TO_QUOTA. KVCM reduces the marked request's capacity and history together, preserving capacity - history_length, every other request, and all fixed constraints. The lower bound is derived as max(1, capacity - history_length); there is no separate minimum-capacity parameter. Planning uses the final buffer layout, pool rounding, and resume utilization instead of a runtime quota-to-token estimate. The hot GPU allocation stays within the supplied byte quota; infeasible fixed workloads or an unachievable lower bound raise an error. FIXED remains the default with its existing behavior.

Runtime integration adds one import, marks the long request in the existing constraint, and declares its decode headroom through history_length. All fitting stays inside KVCM. This is an independent draft based on main; #19213 is unchanged.

The initial API supports one marked request, at any position across initialization batches, and no shared system prompt in its batch. It searches the requested capacity, derived lower bound, and block-aligned capacities; it does not promise the largest token-exact fit. resolved_constraints exposes the selected workloads without changing the input config.

The resolved workloads become persistent per-pool-group minimum slot counts within that manager instance. These floors continue to constrain later pool-ratio and quota adjustments, even when typical_step supplies the target ratio. Fitting runs only at construction; a later quota below the saved floors is rejected rather than triggering another fit. The runtime destroys the estimation manager and restores the original pool-sizing configuration before constructing the final manager, so the temporary floors are not inherited. The final manager derives its own constraints from the restored configuration.

When fitting is necessary, cache the other batches' per-pool floors and the marked request's fixed peers. Each binary-search probe evaluates only the marked request, adds raw peer slots before resume-utilization rounding, and retains the existing byte-rounding and total-quota checks. The original-fit fast path and pool-ratio semantics are preserved. Both backends warn once when a constraint is reduced, including its location, old/new capacity and history, and configured/usable quota.

Test Coverage

• Repository commit hooks passed.
• Python backend: all 17 TestFitToQuota tests passed on B200 for 89fbced543. These runs used the current Python planner and existing, unchanged native copy-engine utility bindings; they do not validate the new C++ backend.
• Python backend: all 27 storage-statistics and initialization-ratio regression tests passed.
• Added backend-shared tests for bounded allocation, side buffers, resume/page rounding, SWA alignment, scratch reuse, fixed peers/envelopes, SSM placeholders, config copying, rejection, and legacy growth. Request-policy coverage includes non-first request selection, headroom values 0/1/65, and an infeasible derived lower bound.
• Added runtime coverage for byte/token budgets in temporary estimation and final managers.
• Added exhaustive allocation-oracle cases at resume utilizations 0.5, 0.97, and 1.0, covering per-request scratch rounding, coalesced SWA slots, SSM placeholders, and another batch with a shared prompt.
• Retained the [None][fix] Respect KVCM V2 initialization and warmup budgets #19213 validation cases with V2 explicitly enabled: A10 multimodal gRPC, serve, and OpenAI chat; B200 SeedOss-36B test_auto_dtype; H100 Mistral-Small-24B test_auto_dtype[forced_chunked_prefill].
• C++ backend and bindings rebuilt successfully on B200 with dynamic NVRTC linking. All 44 FIT/storage/ratio tests passed; the 17 FIT tests also passed through the production package import path. The non-first-request test now copies native descriptor values before replacing their owning config vector.
• Runtime-manager validation: all 11 selected tests passed, including byte/token budgets for estimation/final allocation, constraint declaration, ratio overrides, and external-draft allocation/resume.
• B200 TestSeedOss_36B::test_auto_dtype passed the full 1,319-sample GSM8K evaluation (92.343 accuracy, test threshold 87.597). Temporary KV usage was 87.38 GiB against a supplied 87.3936 GiB quota; profiling completed without OOM.
• A10 multimodal validation passed all 10 selected cases: gRPC 3/3, serve 2/2, and OpenAI chat 5/5, with no skips or OOM. The separate L40S-only chat parameter was deselected and is not counted. Example temporary constraints shrank from 262,144 tokens to 26,048 (gRPC), 25,632 (serve), and 89,120 (chat); serve/chat profiling completed and constructed final managers.
• H100 TestMistralSmall24B::test_auto_dtype[forced_chunked_prefill] passed all 900 MMMU samples (54.111 accuracy, test threshold 53.123). FIT shortened the temporary constraint from 131,072 to 69,472 tokens; temporary KV usage was 21.69 GiB against a 21.7105 GiB quota. Profiling completed without OOM (69.77 GiB peak), followed by final-manager construction and the complete evaluation.

Python CPU-planner microbenchmark against 91b1142ba8, on the same B200 host, Python 3.12.3, nine lifecycles, 1M requested capacity, and 32 tokens/block. Median of nine interleaved runs with warnings suppressed; all selected descriptors match the baseline. These measurements exclude GPU allocation and do not represent C++ or end-to-end inference speedups.

Batch size	Before	After	Speedup
1	1.94 ms	1.04 ms	1.86x
32	21.16 ms	3.31 ms	6.38x
256	160.64 ms	19.62 ms	8.19x
1024	636.78 ms	75.71 ms	8.41x

The already-fitting batch-256 case remains 9.41 ms before/after with the same 257 request evaluations. For a fitting batch of 256, request evaluations decrease from 4,369 to 529.

Known warmup limitation: Seed and A10 runs confirmed that the unchanged consumer fills a page-aligned capacity query and V2 then reserves one extra generation token. Graph warmup/capture skips batch sizes greater than one under these pool-limited configurations. This is a separate consumer fix already present in #19213 and is excluded from this draft's constraint-only runtime scope; the passing tests do not establish restored graph coverage or inference performance. Multi-rank profiling shapes also remain to be validated. Fixed context workloads that exceed the physical quota remain errors, including quotas derived from approximate token accounting.

Related CI OOM scope: TestQwen3_5_4B::test_mismatched_block_reuse in CI 60868 supplies 61.47 GiB explicitly during temporary estimation (50.27 GiB SSM state plus 8.125 GiB attention KV, divided by 0.95). Its original constraints fit that quota, so this policy does not shorten the long request or reduce the large fixed SSM reservation. H100 initialization reproduction with the original context configuration confirmed 61.46875 GiB actual cache, the unchanged 262,144-token constraint, skipped 1.89 GiB Mamba warmup allocations, and a fatal 384 MiB encoder-profiling allocation with 376.56 MiB free. This case remains unfixed; the existing 6793949 waiver is retained. The reproduction initializes one context worker and does not execute the full two-GPU reuse test.

The complete same-GPU test_pd_disagg_multimodal_with_block_reuse test passed on H100 with FIT and with a same-source/native-binary FIXED-policy control. V2 was explicitly selected to match CI 60801; current main's auto setting would select V1. Independent records confirm both prefill and decode workers used the C++ V2 backend. With an identical 15,686,906,675-byte decode quota, FIT allocated 15,686,696,960 bytes versus FIXED's 31,658,606,592 bytes. Prefill likewise stayed within quota under FIT and grew to approximately 29.48 GiB under FIXED. Final managers in both modes had no constraints after restoration of the user configuration. This validates the quota-growth mechanism and the new policy's effect; both modes passed, so the historical CI OOM was not reproduced and is not claimed fixed. This control is not a replay of the frozen CI binary.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.