fix: bound DSA prefill score memory

[CUHKSZzxy]

Supersedes #4892 by carrying its sparse FlashMLA refactor together with the DSA memory and short-prefill optimizations.

Summary

• Separate dense MLA and sparse DSA FlashMLA execution, with implementation selection owned by the CUDA attention builder.
• Bound DeepGEMM DSA prefill scores by chunking query rows under a configurable 512 MiB budget, reusing the flattened indexer K cache, and reserving the same workspace before KV-cache sizing, addressing the same [M, N] FP32 logits OOM as vLLM #36178.
• When the full prefill context fits within DSA top-k, skip index scoring and sparse mapping and use equivalent dense attention, including FA3 when available. Indexer K-cache writes are preserved, while MTP still computes the indices required by later draft iterations.
• Default DSA MLA KV cache to BF16; an explicit FP8 cache policy still selects the packed FP8 MLA layout.

Validation

• Pre-commit checks passed.
• Focused attention, DSA indexer, MTP, configuration, and cache-engine tests passed (60 passed).
• GLM-5.2 TP8 SWE serving completed 80/80 requests with FP8 KV cache, MTP5, prefix caching, and symmetric-memory all-reduce: 914.67 tok/s overall output and 1425.27 tok/s steady completion throughput.
• The SWE prompts use 16K+ contexts, so this run validates the existing long-context sparse path and absence of a serving regression; it does not measure the short-prefill dense-routing speedup.

Assistance

Assisted with Codex + GPT-5.6-Sol xHigh, reviewed manually

[Copilot]

Pull request overview

This PR separates dense and sparse FlashMLA paths while bounding DSA prefill score memory and adding short-prefill optimizations.

Changes:

• Adds chunked DSA scoring and reserved workspace.
• Adds dense routing for short prefills with MTP index safeguards.
• Supports BF16 and FP8 MLA cache policies.

Review findings:

• Critical (1 vote): sparse_mla.py:230 has an invalid BF16 dense fallback when FA3 is unavailable.
• Moderate (3 votes): nsa.py:391 can exceed the memory budget through bitonic_topk temporaries.
• Moderate (2 votes): nsa.py:417 can allocate the full score tensor when DeepGEMM is unavailable.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Summary
tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py	Tests MTP index-buffer behavior.
tests/pytorch/nn/test_nsa.py	Tests score chunking and short-prefill behavior.
tests/pytorch/kernel/test_mla_attention.py	Tests MLA routing and mapping.
tests/pytorch/engine/test_executor_base.py	Tests workspace reservation.
lmdeploy/pytorch/nn/nsa.py	Exposes indexer and prefill configuration.
lmdeploy/pytorch/models/glm_moe_dsa.py	Handles dense prefills and MTP indices.
lmdeploy/pytorch/models/deepseek_v32.py	Integrates indexer and MLA configuration.
lmdeploy/pytorch/envs.py	Adds score-memory configuration.
lmdeploy/pytorch/engine/executor/base.py	Reserves DSA score workspace.
lmdeploy/pytorch/configurations/glm_moe_dsa.py	Updates inherited cache defaults.
lmdeploy/pytorch/configurations/deepseek_v32.py	Defines BF16/FP8 MLA cache policy.
lmdeploy/pytorch/backends/nsa.py	Updates NSA interfaces.
lmdeploy/pytorch/backends/cuda/nsa.py	Implements bounded scoring and skip logic.
lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py	Implements sparse FlashMLA execution.
lmdeploy/pytorch/backends/cuda/attention/mla.py	Implements dense FlashMLA execution.
lmdeploy/pytorch/backends/cuda/attention/__init__.py	Selects dense or sparse implementations.

Suppressed comments (2)

lmdeploy/pytorch/backends/cuda/nsa.py:48

• The max(1, ...) clamp means the claimed budget is exceeded whenever one row of FP32 logits is larger than max_logits_bytes (for example, a 1 MiB budget with a context wider than 262,144 tokens still allocates at least 4 MiB). Since this is configurable, either reject such a configuration with a clear error or document/account for the minimum per-row allocation instead of claiming the tensor is bounded.

    _fp32_bytes = 4
    return max(1, max_logits_bytes // (max_kv_seqlen * _fp32_bytes))

lmdeploy/pytorch/models/deepseek_v32.py:198

• This call leaves allow_short_prefill_scoring_skip at its default False for the regular Deepseek-V3.2 DSA model. As a result, prefill contexts at or below index_topk still run _score_and_select and materialize the dense logits instead of taking the newly implemented dense route; only the GLM-MoE-DSA call site enables this optimization. Enable the flag here as well so short prefills preserve the K-cache write but skip scoring.

            allow_short_prefill_scoring_skip=layer_idx < config.num_hidden_layers,

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[grimoire]

LGTM