Add Sparse MLA prefill/decode kernel

[wuxun-zhang]

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.

Purpose

This is to support DeepSeek V4 model running.

• Added new flash_mla_interface.py
• Added sparse mla prefill and fp8 decode kernel
• Added tests for correctness check

Test Plan

python -m pytest tests/flash_mla/test_mla_sparse_prefill.py -v
python -m pytest tests/flash_mla/test_mla_sparse_decode_fp8.py -v

Test Result

all passed

(Optional) Documentation Update

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

[Copilot]

Pull request overview

Adds Sparse MLA (DeepSeek V4–style) sparse attention support to the XPU FA2 extension, including configurable kernel generation, Python entrypoints, and correctness tests.

Changes:

• Introduces Sparse MLA prefill and FP8 decode kernels + runtime dispatch/registration in the FA2 extension.
• Adds build/config plumbing for sparse MLA kernel selection (SPARSE_MLA_KERNELS_ENABLED, VLLM_SPARSE_MLA_CONFIG) and new preset config files.
• Adds correctness tests for sparse MLA prefill and sparse FP8 decode, plus documentation updates for configuration.

Reviewed changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
vllm_xpu_kernels/flash_mla_interface.py	Python interface wrappers for sparse MLA prefill + FP8 decode ops.
tests/flash_mla/test_mla_sparse_prefill.py	Correctness test for sparse MLA prefill (compares to reference).
tests/flash_mla/test_mla_sparse_decode_fp8.py	Correctness test for sparse MLA FP8 decode (compares to reference).
setup.py	Enables building _vllm_fa2_C when sparse MLA kernels are enabled; forwards sparse MLA config env var.
KERNEL_CONFIGURATION.md	Documents sparse MLA kernel category, configs, presets, and troubleshooting.
csrc/xpu/attn/mla/utils.hpp	Adds MLA-specific utilities/macros used by sparse kernels.
csrc/xpu/attn/mla/sparse_mla_prefill_kernel_template.cpp.in	Template for generated prefill instantiation units.
csrc/xpu/attn/mla/sparse_mla_prefill_fwd.hpp	Sparse MLA prefill kernel implementation + gather kernel.
csrc/xpu/attn/mla/sparse_mla_fwd.cpp	Torch/XPU-facing implementations for sparse MLA prefill + decode.
csrc/xpu/attn/mla/sparse_mla_decode_fp8_kernel_template.cpp.in	Template for generated FP8 decode instantiation units.
csrc/xpu/attn/mla/sparse_mla_decode_fp8_fwd.hpp	Sparse MLA FP8 decode implementation + gather/dequant kernel.
csrc/xpu/attn/mla/sparse_mla_configure.cmake	Parses sparse MLA config and generates instantiation sources/defines.
csrc/xpu/attn/mla/flash.h	Adds shared structs/constants for sparse MLA kernels.
csrc/xpu/attn/kernel_configs/sparse_mla_full.conf	“all” preset for building all sparse MLA variants.
csrc/xpu/attn/kernel_configs/sparse_mla_default.conf	Default sparse MLA preset targeting DeepSeek V4 (2 variants).
csrc/xpu/attn/kernel_configs/README.md	Lists sparse MLA config presets (counts/usage).
csrc/xpu/attn/attn_interface.h	Declares sparse MLA C++ interfaces exposed to registration.
csrc/xpu/attn/attn_interface.cpp	Implements sparse MLA interface forwarding to MLA implementations.
csrc/flash_attn/flash_api.cpp	Registers new torch.ops for sparse MLA prefill/decode.
CMakeLists.txt	Adds SPARSE_MLA_KERNELS_ENABLED and integrates sparse MLA generation/build into FA2 extension.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[wuxun-zhang]

@xinyu-intel @pengzhao-intel please help review, thanks

[jikunshang]

no much comments for kernels part.
one open is: will we add flash mla (not sparse one)in the future?

[jikunshang]

I feel we will always build FA2? better revert this part.

[wuxun-zhang]

good catch

[jikunshang]

emm, why need this? especially we have #445