[drakosha]

Purpose

FlashMLASparseImpl reserves the bf16 prefill KV workspace even when the fp8 mixed-batch path is selected and never reads it. Gate the max_model_len-sized reservation with the same path condition used by the metadata builder. Keep the three max_num_batched_tokens q/out buffers: the staging implementation in #60220 uses them.

After #61023 the prefill factor is 2 above 256K context. At max_model_len=786432, the unused bf16 buffer is 1.6875 GiB per rank. The net saving depends on the other users of the shared workspace; the earlier 4.22 GiB figure no longer describes main.

The constructor and builder share _use_fp8_mixed_batch; forward_mqa checks their decisions agree. The separate prefill/decode, HiSparse and PCP+DCP gather reservations remain covered by the constructor test. The redundant truth-table test was removed following review.

Related work: #57575 reserved these buffers before KV sizing; #60220 stages q/out in the retained batch-sized buffers; #67510 adds HiSparse token-sharded prefill. This PR gates the unused bf16 reservation on the non-HiSparse mixed-batch path. The open-PR searches for mixed batch prefill workspace found no other PR implementing that gate.

Test Plan

The existing test_fp8_mixed_batch_path_reserves_no_prefill_workspace matrix covers head counts 16/32, HiSparse on/off and PCP/DCP combinations. The existing fp8 decode-correctness parametrization includes TP8.

Full applicable pre-commit suite on the current rebase, including Ruff 0.17.2, mypy and test tethering: passed.

uvx --from pre-commit pre-commit run --files \
  vllm/v1/attention/backends/mla/flashmla_sparse.py \
  tests/v1/attention/test_sparse_mla_backends.py

Test Result

Rebased onto main 5039f736e6. The two changed files are byte-identical to the tested pre-publication revision.

Fresh A/B on nightly 615c9d6474, 4x H200 NVL, nvidia/GLM-5.3-NVFP4, TP4/DCP4/EP, MTP3, fp8_ds_mla, context 786432, batch 8192, 32 sequences, utilization 0.945; extensible KV, explicit KV memory and CPU offload disabled.

Both sides carry the same serving patches (#59367, #49249, #58209, #58985 and the narrow FULL-sampling fix from #59368), plus FlashMLA b18ed74 with #27/#28/#29. This PR is the only difference between the two rows.

	without this PR	with this PR
workspace after profile run, bytes/rank	3,238,002,688	1,745,879,040
workspace after profile run, GiB/rank	3.015625	1.6259765625
model-loading footprint, GiB/rank	113.85	112.17
GPU KV cache, tokens	1,142,016	1,251,072

Net saving: 1,492,123,648 bytes (1.3896484375 GiB) per rank, with 109,056 additional KV tokens on this stack.

The combined build also passed tool calls (streaming and non-streaming), reasoning, 230,456-token needle retrieval at depths 0.1/0.5/0.9, and three 544,235-token prompts cold and replayed with CPU offload (3/3 each). Fresh prefill was about 6.4k tok/s. These are serving smoke/load checks; no full GSM8K run or statistical throughput gain is claimed.

Earlier isolated CUDA coverage on af7f9488c: 14 unit cases passed and 56 fp8 FlashMLA decode cases passed (56 unsupported-block-size skips); with #60220 applied, 72 passed/56 skipped. That GPU unit matrix was not rerun for this publication; the current evidence is the serving run and pre-commit checks.

AI assistance (Claude Code and Codex) was used for the analysis, implementation, tests, measurements and description.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

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

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, @drakosha.

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

[lucamotz]

The gating looks right to me on main (2011b8f). Your serving numbers predate #61023, which sets the factor to 2 above 256K, so at 786K the buffer is 1.69 GiB on main, and the indexer's profile-run K-gather reservation in the same workspace keeps about 0.2 GiB of that (computed from the sizing code, not measured); could you re-measure on main?

[lucamotz]

test_fp8_mixed_batch_path_reserves_no_prefill_workspace covers the same cases and checks impl.fp8_use_mixed_batch, so I think this one can go.

Suggested change

	@pytest.mark.parametrize(
	("num_heads", "use_hisparse", "pcp_dcp_kv_gather", "expected"),
	[
	(MIN_HEADS_FOR_BF16_PREFILL - 1, False, False, True),
	(MIN_HEADS_FOR_BF16_PREFILL, False, False, False),
	(MIN_HEADS_FOR_BF16_PREFILL - 1, True, False, False),
	(MIN_HEADS_FOR_BF16_PREFILL - 1, False, True, False),
	],
	)
	def test_fp8_mixed_batch_requires_few_heads_without_hisparse_or_kv_gather(
	num_heads, use_hisparse, pcp_dcp_kv_gather, expected
	):
	assert (
	FlashMLASparseImpl._use_fp8_mixed_batch(
	num_heads, use_hisparse, pcp_dcp_kv_gather
	)
	is expected
	)

[drakosha]

@lucamotz Remeasured on nightly 615c9d6474, 4x H200, TP4/DCP4, MTP3, context 786432, batch 8192, 32 sequences, utilization 0.945; no extensible KV, explicit KV limit or CPU offload. With the common production patches held constant (including #59367 and the narrow #59368 fix), workspace after profiling falls from 3.015625 to 1.6259765625 GiB/rank and KV grows from 1,142,016 to 1,251,072 tokens. The net saving is 1.3896484375 GiB/rank on this stack. The old 4.22 GiB claim is removed from the description, and the redundant truth-table test is gone. Rebased onto 5039f736e6; full pre-commit passed. The updated description records the common patches and the serving checks.