[viraatc]

Description

Port of #17154 (closed unmerged against main) onto feat/k3_26.08.

DFlashWorker._lazy_init_ctx_buffers and DSv4DSparkWorker._lazy_init allocate their per-request slot pools once and reuse them for every later batch shape. Both sized the pool from spec_metadata.max_num_requests. That value is not a stable bound:

• create_cuda_graph_metadata narrows it to the captured graph bucket, so a first call on a small bucket pins the pool to that bucket.
• Under disaggregated attention DP with overlap headroom, the seq-slot pool (compute_max_num_sequences) is 2 * max_batch_size, which max_num_requests never covers.

feat/k3_26.08 already passes num_seq_slots into DFlashSpecMetadata and DSparkSpecMetadata through get_spec_metadata. This PR:

• sizes both worker pools from spec_metadata.num_seq_slots;
• normalizes its default (num_seq_slots or max_num_requests) in both metadata __post_init__s, so directly built metadata and CUDA-graph copies keep the full capacity.

The tests from #17154 do not apply here: the DSpark worker on this branch uses manager-owned draft pages instead of _kv_windows. New tests are written against this branch's layout.

Memory: when the pool is 2 * max_batch_size, the private DFlash context cache (used when the draft KV is not in the unified pool) doubles with it. The DSpark progress buffers are small.

Test Coverage

• test_dspark_worker.py::test_metadata_preserves_seq_slot_pool_in_graph_copy
• test_dspark_worker.py::test_worker_graph_bucket_uses_full_seq_slot_pool[pool-eq-batch|disagg-adp-2x-pool]
• test_dflash_worker.py::test_context_cache_spans_full_seq_slot_pool_on_graph_bucket
• Two existing SimpleNamespace metadata stubs gain num_seq_slots.
• Run on GB300 against feat/k3_26.08 @ e8a5275c4f (both test files): 37 passed, 3 failed with this PR; 33 passed, 3 failed on the unpatched base. The 3 failures are the same tests on both, so they predate this PR: test_dspark_worker.py::test_prepare_keeps_cuda_graph_slots_across_replays, ::test_position_ceiling_covers_the_config_cap_and_grows_after_kv_estimation and ::test_position_bound_holds_across_cuda_graph_replays. The 4 new test cases pass. test_freqs_table_prefers_the_runtime_ceiling_over_the_config_cap needs pytest-mock.
• CI-equivalent run on a GB300 node, in the feat/k3_26.08 @ e8a5275c4f image, with [https://nvbugs/6535874][fix] Size DFlash/DSpark worker slot pools from num_seq_slots #21305, [None][fix] Disaggregated generation: one-hot draft probabilities for placeholder drafts at the handoff #21496, [None][perf] Kimi K3: split-K FP8-KV MLA context FMHA for SM100/SM103 #21529 and [None][fix] Port #17850 to feat/k3_26.08: keep the KV cache estimation managers off the host tier #21530 applied together:
  • Suites: _torch/executor, _torch/speculative/hw_agnostic, _torch/attention/test_attention_mla.py and test_fp4_mla.py, run from tests/unittest with its pytest config.
  • 2,812 tests: no failure that the unpatched base does not also have. The base has the same 189 failures and 10 errors.
  • The new tests pass.

PR Checklist

• Commits signed off (DCO).
• GPU unit tests run (GB300; the 3 failures above also fail on the base).

[viraatc]

/bot run

[tensorrt-cicd]

PR_Github #83641 [ run ] triggered by Bot. Commit: 2064b5d Link to invocation

[tensorrt-cicd]

PR_Github #83641 [ run ] completed with state FAILURE. Commit: 2064b5d
/LLM/main/L0_MergeRequest_PR pipeline #69151 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

CI Agent Failure Analysis

Link to invocation