[zheyuf]

Description

Bug (M3.1 ask 2a). A DSpark attention-DP gen worker in disaggregated serving raises ValueError: History length cannot exceed capacity from the draft KVCacheManagerV2.update_resources → kv_cache.resize. The crash handler then hard-kills every rank.

It shows up at gen max_batch_size 2–4 per rank: DEP4 → DEP4 B16 died in all 3 e2e runs. It has never shown up at 16 per rank, on TP gens, or in gen-only runs.

Root cause. Rolling back an attention-DP pad dummy leaks the dummy's unpaired draft mirror:

1. One rank's empty batch vetoes the forward everywhere. A rank schedules an empty batch and cannot pad it: _pad_empty_attention_dp_batch refuses once len(active_requests) >= max_num_active_requests, and requests still receiving KV reach that cap at 2–4 per rank. _can_queue then vetoes the forward on every rank.
2. A peer's dummy already holds a draft mirror. A peer had padded a pre-schedule dummy, and the V2 scheduler had scheduled it: _try_allocate_generation → _admit_unpaired_draft → admit_mirror created a draft mirror for it, at capacity 0.
3. The rollback misses that mirror. _finalize_adp_dummy_allocation(False) frees the dummy through _free_pad_dummy_kv, which frees the draft pool only when it is paired (enable_joint_kv_cache_reuse). With an unpaired pool (UNPAIRED_DRAFT_POOL, the M3.1 DSpark default on Blackwell since 15ffec1), the mirror stays in the draft kv_cache_map under the fixed dummy request ID.
4. The veto also defers the previous batch's update. On that rank, should_process_previous_batch is False, so the update of the previous batch waits. That batch held the previous dummy.
5. The deferred update hits the leaked mirror. In the next iteration the rank has real work, so it pads no new dummy and nothing clears the leak. The deferred draft update_resources of the old dummy finds the leaked mirror under the same ID and calls resize(None, 3) on a capacity-0 cache, which throws.

Small batch sizes fail only because step 1 needs a rank at the pad cap. The draft pool's size plays no part: it never filled up in any crash log.

Fix. _free_pad_dummy_kv now frees the dummy in whichever draft pool exists, paired or not. Freeing a pool that holds no cache for the dummy is a no-op, so its other two callers behave as before. Those callers handle a spec-slot failure right after a dummy is created, when an unpaired pool holds nothing for it yet.

Test Coverage

• New unit test test_adp_dummy_rollback_frees_unpaired_draft_mirror in tests/unittest/_torch/executor/test_py_executor.py. CI collects it through the unittest/_torch/executor entry (e.g. l0_cpu.yml).
  • At the branch tip: 244 passed, 6 failed. The failures are the new test plus 5 × test_forward_step_carries_context_logits_request_to_runner, whose SimpleNamespace stub lacks _loop_on_execution_stream.
  • With the fix: 245 passed, 5 failed. The new test passes; the same 5 fail as before.
  • Both runs overlaid this branch's tensorrt_llm/_torch on image 516f0342cd, which has no compiled files under _torch.
• GB300 serving validation, MiniMax-M3.1-Flash + DSpark K=3. It ran on image 516f0342cd, the one in the bug report. That image's rollback freed only the target, so it has the same leak.
  • 8k repro, 1 ctx DEP2 → 1 gen DEP4, gen max_batch_size 2 per rank, 1600 requests per run:
    • without the fix, 2/2 runs crashed, after ~480 and ~230 requests;
    • with the fix, 3/3 runs completed 1600/1600 with no errors, at 7.5 req/s;
    • instrumented runs saw the rollback free a dummy's draft mirror ~70 times per run, so the trigger kept firing and was handled.
  • AgentX e2e DEP4 ctx → DEP4 gen B16 (4 per rank), 69 sessions, the point that died 3/3 before: with the fix it ran the full 30-min window, with 3842 requests and 0 errors.
  • Those runs also carried a guard in update_resources that skips dummies. It never fired, so it is not part of this PR.

PR Checklist

• Python only: no API, config, C++ or build change.
• yapf and ruff-legacy (py_executor.py), and ruff check/format (test file), are clean at the pinned versions (yapf 0.43.0, ruff 0.9.4).

🤖 Generated with Claude Code