[Core] Move FlashInfer MLA decode workspaces into WorkspaceManager

[aoshen02]

Overview

Move the FlashInfer MLA decode workspaces (sparse trtllm-gen/cute-dsl, dense MLA, DSV4/DSV4.1 sparse) from module-global torch.zeros tensors into WorkspaceManager scratch. Part of #59055; with #59156 the memory is then released by /sleep.

Claims

• The 394 MiB per GPU FlashInfer sparse-MLA workspace (GLM-5.3 / GLM-5.3-Flash) and the 128/256 MiB dense-MLA and DSV4 workspaces no longer stay resident while the engine sleeps once [Feature] Release WorkspaceManager scratch on sleep #59156 lands, and are no longer additive to the shared scratch while awake.
• No behavior change for the kernels: the buffers are requested at the same size, reserved before CUDA graph capture, and re-requested on every call (same pattern as flashmla_sparse.py, prefill/trtllm_ragged.py).
• The DCP sizing (_required_workspace_bytes) and VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE override are unchanged.

Validation

• Read FlashInfer 0.7.0 (mla/_core.py, data/csrc/trtllm_fmha_kernel_launcher.cu, cute_dsl/attention/{mla_decode.py,monolithic/mla_decode.py}); 0.7.0.post1 (the pinned version) differs from 0.7.0 only in aot.py, jit/gemm/core.py and tests. On every decode path these backends use, the workspace holds only per-launch scratch (trtllm-gen: softmax-stats slab + split-KV scratch, semaphores live in the separate multi_ctas_kv_counter_buffer; cute-dsl: split-KV acc_o/acc_lse partials consumed by the reduction kernel of the same call; SM120 sparse: mid_out/mid_lse partials), so no content survives between calls and no zero-init is needed. The prefill backends already hand trtllm-gen an uninitialized WorkspaceManager view.
• No other WorkspaceManager view is an input/output of these decode kernels on the paths read (inputs come from the persistent topk_indices_buffer, the index group, Triton converters, _mtp_hidden_buffer).
• Unit tests (H200, FlashInfer 0.6.12 venv): tests/v1/attention/test_flashinfer_sparse_mla_workspace.py, test_sparse_mla_backends.py, test_flashinfer_mla_dcp.py, test_dspark_noncausal_sparse_mla.py, tests/kernels/attention/test_flashinfer_mla_decode.py: 164 passed, 671 skipped (SM100-only), 18 failed; the 18 failures are identical on unmodified main in that environment (16 × cp_gather_and_upconvert_fp8_kv_cache signature mismatch of the local _C build in FlashMLA tests, 1 × FlashInfer 0.6.12 SM90 JIT build failure, 1 × test_index_group_convert_during_piecewise_capture).
• pre-commit (ruff, ruff format, mypy, etc.) passes on the commit.
• End-to-end, GLM-5.3-Flash TP4 with MTP, CUDA graphs and --enable-sleep-mode on 4× GB300, same image and node, on top of [Feature] Release WorkspaceManager scratch on sleep #59156 + [Feature] Offload KV-init runtime state on sleep #59158 + [Core] Release the FlashInfer allreduce workspace (MNNVL or TRT-LLM) with --enable-nccl-comm-suspend #59360 (so the workspace tag is released on sleep):
  • Memory per GPU after /sleep (nvidia-smi): 5945 MiB without this PR, 5551 MiB with it (-394 MiB, exactly the sparse-MLA workspace). Torch-allocated memory drops by 352 MiB while awake (263465 to 263113 MiB per worker) because the workspace now fits inside the existing 1520 MiB workspace slot; KV cache grows by 42 MiB.
  • Post-wake parity 0/0/0/0 MiB in both runs; 3 sleep/wake cycles (full, selective, full), 195/195 outputs identical to the pre-sleep reference in both runs; no Workspace is locked errors.
• Not validated: DCP, HiSparse, the dense flashinfer_mla.py path and the DSV4/DSV4.1 FlashInfer paths on a real model (V4-Flash in this image uses FlashMLA; forcing FLASHINFER_MLA_SPARSE_DSV4 rejects its packed KV layout).

Details

• flashinfer_mla_sparse.py: _get_workspace_buffer() now requests ((bytes,), int8) from the manager; a module-level int keeps the largest size requested so the DCP builder can still raise it in its __init__, before the first forward/capture. FlashInferMLASparseImpl and the SM120 impl reserve in __init__ and re-request in _run_mqa_kernel; _prepare_mqa_kernel loses its device argument.
• flashinfer_mla.py: same, sized by need_to_return_lse_for_decode (known at construction).
• deepseek_v4(1)/nvidia/flashinfer_sparse.py: _get_flashinfer_dsv4_workspace() requests ((128 MiB,), uint8); the SM100 class reserves it in the dummy run (like the FlashMLA DSV4 path), the SM120 class already did via _reserve_empty_forward_workspace.
• Each (ubatch, lane) slot that runs these kernels grows to at least the workspace size; with DBO that is ×2, as for every other scratch user.
• Two tests that assigned impl._workspace_buffer directly now monkeypatch _get_workspace_buffer; the autotune test installs a WorkspaceManager.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

AI assistance: written with Claude Code (Claude Fable 5.1), reviewed and tested by the human submitter.

Re-check on current main (bdd31c31d), GB300 (rebased; adds a WorkspaceManager to the SM100-only DSV4.1 dspark test, which this change otherwise breaks)

Re-checked on main bdd31c3 (rebased; only an import-line conflict in
test_sparse_mla_backends.py), 4x GB300 (SM103, so the SM100-only tests run),
FlashInfer 0.7.0.post1:

• The GB300 run found that test_dsv41_flashinfer_dspark_window_matches_reference (12
  cases) called DeepseekV4FlashInferMLAAttention._forward without a WorkspaceManager
  and failed with this PR; fixed in the test, which now installs one (as the autotune
  test does). test_dspark_noncausal_sparse_mla.py + test_flashinfer_sparse_mla_workspace.py
  • test_flashinfer_mla_decode.py: 59 passed, 9 skipped, 3 runs.
• DeepSeek-V4-Flash TP4 + MTP keeps FlashMLA on SM100/SM103 (and the FlashInfer DSV4
  backend still rejects its packed KV), so it does not exercise this change: as a
  regression check, asleep memory, wake time and outputs (3 cycles, 28/28 greedy and
  28/28 seeded) are unchanged. The GLM-5.3-Flash numbers above were not re-measured.

🤖 Generated with Claude Code

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork