[BugFix] bound FlashMLA sparse decode intermediate tensors size

[cjackal]

Purpose

Fix #44545 - not in an ideal way though.

When FlashMLA sparse kernel is selected with FP8 KV cache dtype, FlashMLA sparse decode kernel(https://github.com/vllm-project/FlashMLA/blob/a8f794d1251cbfd88a5011445dd5582289c727e4/csrc/api/sparse_decode.h#L184) is called in mixed batch mode. This kernel internally allocates intermediate tensors (o_accum) on call, which quickly dominates the usual unallocated workspace size (~2GB per the comment in flashmla_sparse.py) for long prefill:

For 32k input on H200, the bytesize of o_accum is

# b = 1
# num_sm_parts = 1
# s_q = 32768
# h_q = 64
# d_v = 512
(b + num_sm_parts) * s_q * h_q * d_v * 4 bytes = 8 GiB

For BF16 KV cache, a different kernel (FlashMLA sparse prefill kernel - https://github.com/vllm-project/FlashMLA/blob/a8f794d1251cbfd88a5011445dd5582289c727e4/csrc/api/sparse_fwd.h#L101) is called, which does not allocate a massive VRAM on kernel call and thus does not suffer an OOM contrary to the FP8 code branch. (Note the additional d_v factor in the formular for o_accum, no other scratch tensors in FlashMLA have this factor thus negligible in terms of VRAM occupancy)

Hence we chunk the mixed-batch input into manageable bytes (VLLM_FLASHMLA_SPARSE_MAX_SCRATCH_MB) and loop over the chunks to aggregate the output attention tensor, which is in the same spirit as the indexer budget management logic in #36178.

Due to the mechanism (looping the kernel over sub-chunks instead of single big shot), it may incur a throughput penalty as a tradeoff, but I have no better idea yet.

As the context window for MLA model families keep growing, FP8 KV cache is kind of inevitable, so I think this PR can work as a bandaid.

Test Plan

No OOM for the reproducer in #44545

Test Result

pass

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[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.

[cjackal]

@MatthewBonanni ping in case if this issue is still of interest to you.

[mergify]

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

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

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