feat: add XTuner TileLang sparse MLA backend

[CUHKSZzxy]

Summary

• add an environment-selected TileLang SparseMLA backend using XTuner's forward kernel
• use the same kernel for packed prefill and zero-copy BF16 paged-cache decode
• run native TP-local query heads, preserve DSA index production for short prefill, and reject the unsupported FP8 MLA KV-cache path

Prerequisite

Depends on #4896

Performance

Measurements used an NVIDIA H200 and the final integrated implementation.

Isolated sparse MLA forward

BF16 Q/KV, 8 TP-local query heads, dqk=576, dv=512, and topk=2048. Times are paired CUDA-event medians across five interleaved trials. FlashMLA pads query heads from 8 to 64 as required by its API; TileLang runs the native 8-head shape.

Shape	FlashMLA	TileLang	TileLang / FlashMLA
Short prefill (sq=skv=64)	0.069 ms	0.405 ms	5.83x
SWE-like prefill (sq=skv=2119)	1.057 ms	1.807 ms	1.71x

Both cases passed atol=rtol=2e-2; maximum absolute differences were 1.22e-4 and 2.44e-4, respectively.

End-to-end serving

GLM-5.2-FP8 weights with BF16 MLA KV cache, TP8, MTP5, symmetric-memory all-reduce, and the SWE 16-conversation x 5-turn workload at concurrency 16. Each backend was measured with one direct profiler-free run and no separate warmup.

Metric	FlashMLA	TileLang	TileLang vs. FlashMLA
Duration	83.42 s	89.91 s	+7.8%
Output throughput	982.00 tok/s	911.09 tok/s	-7.2%
Total throughput	20,900.17 tok/s	19,244.88 tok/s	-7.9%
Mean TTFT	3,139.53 ms	3,505.36 ms	+11.7%
Mean TPOT	12.30 ms	13.23 ms	+7.6%
Steady completion throughput	1,536.04 tok/s	1,472.31 tok/s	-4.1%

Both runs completed 80/80 requests. MTP acceptance was effectively unchanged (81.69% FlashMLA, 81.79% TileLang). The average input length differed by 0.8% because this is a live multi-turn workload, so the result should be treated as a one-trial end-to-end comparison rather than a precise kernel-only attribution.

Validation

• pre-commit hooks passed for all changed files
• pytest -q tests/pytorch/kernel/test_mla_attention.py: 13 passed
• direct comparison with InternLM/xtuner@6751212f: bitwise-equal output at 8 TP-local heads and top-k 2048, including -1-padded indices
• GLM-5.2 full GSM8K with CUDA graphs and TP8: 1267/1319 (96.0576%), with zero API errors and zero empty outputs; recent standard TP8 reference: 1268/1319 (96.1334%)

Assistance

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

[Copilot]

Pull request overview

This PR introduces an environment-selectable SparseMLA backend based on XTuner’s TileLang forward kernel, refactors/extends the existing FlashMLA sparse path, and adds bounded-memory chunking for DSA prefill score computation (with explicit runtime memory reservation). It also adjusts model/config plumbing so sparse MLA selection and DSA index reuse behavior are consistent across prefill/decode, and adds targeted tests covering these behaviors.

Changes:

• Add a TileLang SparseMLA implementation (prefill + BF16 decode) and route sparse MLA backend selection via LMDEPLOY_SPARSE_MLA_BACKEND.
• Bound DeepGEMM prefill score memory by chunking query rows under a configurable logits budget (LMDEPLOY_DSA_INDEXER_MAX_LOGITS_MB) and reserve the workspace during runtime sizing.
• Refactor sparse MLA index mapping/execution out of dense FlashMLAImpl into a dedicated sparse implementation, plus extend tests for builder selection, index mapping, and buffer/index reuse semantics.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py	Adds a regression test ensuring dense-prefill clears reused DSA indices.
tests/pytorch/nn/test_nsa.py	Adds tests for logits-budget chunking and short-prefill scoring-skip behavior; updates DeepGEMM-vs-Triton comparison coverage.
tests/pytorch/kernel/test_mla_attention.py	Adds tests for attention builder routing, index mapping, BF16/FP8 sparse decode behavior, and TileLang decode semantics.
tests/pytorch/engine/test_executor_base.py	Tests runtime sizing/workspace reservation for DSA prefill score budgeting.
lmdeploy/pytorch/nn/nsa.py	Plumbs allow_short_prefill_scoring_skip into the NSA indexer builder.
lmdeploy/pytorch/models/glm_moe_dsa.py	Updates indexer construction to optionally skip short-prefill scoring; hardens shared top-k buffer semantics against stale reuse.
lmdeploy/pytorch/models/deepseek_v32.py	Passes mla_index_topk into attention construction and aligns indexer skip behavior with backend choice.
lmdeploy/pytorch/kernels/cuda/tilelang_sparse_mla.py	Adds the TileLang SparseMLA kernel adapter and launch/validation wrapper.
lmdeploy/pytorch/envs.py	Adds env knobs for DSA logits budget and sparse MLA backend selection.
lmdeploy/pytorch/engine/executor/base.py	Reserves/logs DSA score workspace and includes it in runtime sizing.
lmdeploy/pytorch/configurations/glm_moe_dsa.py	Removes redundant KV-cache dtype override (inherits from DeepseekV32 builder).
lmdeploy/pytorch/configurations/deepseek_v32.py	Defaults DSA MLA KV cache dtype to BF16 (with explicit FP8 policy still possible elsewhere).
lmdeploy/pytorch/backends/nsa.py	Updates NSA indexer interfaces to allow returning None when scoring is intentionally skipped.
lmdeploy/pytorch/backends/cuda/nsa.py	Implements logits-budget chunking for DeepGEMM prefill scores and short-prefill scoring skip logic.
lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py	Introduces FlashMLASparseImpl/TileLangSparseMLAImpl and shared index mapping utilities.
lmdeploy/pytorch/backends/cuda/attention/mla.py	Refactors dense MLA to reject sparse indices, reshapes KV flattening behavior, and threads index_topk through metadata builder.
lmdeploy/pytorch/backends/cuda/attention/init.py	Routes attention builder to dense vs sparse FlashMLA/TileLang implementations based on mla_index_topk + env.

---

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

[grimoire]

LGTM

[lvhan028]

UT error: tests/pytorch/kernel/test_mla_attention.py .....s.....F [ 34%]