SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

Add DeepSeek V4 sparse MLA TRTLLM-GEN kernels (#3269)

## Summary
- Add DeepSeek V4 sparse MLA TRTLLM-GEN decode support for BF16 and
per-tensor FP8 paths.
- Plumb SWA and compressed KV pools, concatenated sparse indices, and
per-query sparse top-k lengths through FlashInfer.
- Add DeepSeek V4 sparse MLA tests covering SWA-only and compressed
top-k cases with variable Q/KV lengths.

## Tests
- `python -m pytest -q --tb=short -k 'not xqa and not cute and not
trtllm-native' tests/attention/test_trtllm_gen_sparse_mla_dsv4.py`: 57
passed
- `python -m pytest -q --tb=short -k 'not xqa and not cute and not
trtllm-native' tests/attention/test_attention_sink_blackwell.py`: 144
passed
- `python -m pytest -q --tb=short -k 'not xqa and not cute and not
trtllm-native' tests/attention/test_trtllm_gen_mla.py`: 7686 passed,
12672 deselected
- `python -m pytest -q --tb=short -n 8 -k 'not xqa and not cute and not
trtllm-native' tests/attention/test_trtllm_gen_attention.py`: 75736
passed, 30800 skipped


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added DeepSeek V4 sparse MLA decode support with variable top-k
behavior and sliding-window KV cache integration.
* Enhanced kernel selection with dynamic token-per-page support for
improved performance flexibility.

* **Tests**
* Added comprehensive test suite for DeepSeek V4 sparse MLA decode
across multiple configurations.

* **Chores**
  * Updated environment variable priority for CUBIN directory selection.
  * Added backward-compatibility alias for MLA decode function.

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3269)
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Perkz Zheng <perkzz@users.noreply.github.com>
P
Perkz Zheng committed
9c76c994b6d64fcfc071214ed5a29af124e680ee
Parent: a0f61ad
Committed by GitHub <noreply@github.com> on 5/21/2026, 7:52:55 AM