SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

Issue #3047: Handle empty KV in MLA chunked-prefill (#3251)

## ๐Ÿ“Œ Description

Moves empty-KV handling for MLA chunked-prefill into the kernels.

The generic MLA path now short-circuits zero-KV tiles before invalid
negative tile access and writes the empty-attention identity directly
(`out = 0`, `lse = -inf`). The empty writer also handles partial outputs
with the same per-CTA tile offset as the normal partial-output path.

The Hopper path now guards KV prefetch for empty tiles and exports true
IEEE `-inf` LSE for empty online-softmax states.

## ๐Ÿ” Related Issues

Fixes #3047

## ๐Ÿš€ Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### โœ… Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## ๐Ÿงช Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

Validation run:

- `pre-commit run --files include/flashinfer/attention/mla.cuh
include/flashinfer/attention/mla_hopper.cuh
tests/attention/test_deepseek_mla.py`
- Blackwell FA2 targeted DeepSeek MLA varlen paged-attention test: `4320
passed, 9504 skipped`
- Hopper FA3 targeted DeepSeek MLA varlen paged-attention test on H100:
`4320 passed, 2592 skipped, 6912 deselected`

## Reviewer Notes

The empty-output helper deliberately supports both final and partial
output buffers. The partial path includes `blockIdx.x *
KTraits::CTA_TILE_Q` so split-K partial storage stays aligned with the
normal writer layout.


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

* **Bug Fixes**
* Improved handling of empty key/value ranges in MLA attention to skip
computation and produce consistent zeroed outputs and -โˆž LSEs for
affected positions
* Corrected LSE write logic to preserve -โˆž when an invalid/max state is
detected, improving numerical correctness and stability

* **Tests**
* Strengthened validation to assert LSE outputs per query row in
variable-length attention tests
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
M
Mingyang Wang committed
bb41dc1bf33cf4d56a79accb13944eab7e6bdd1b
Parent: 14f2bee
Committed by GitHub <noreply@github.com> on 5/7/2026, 8:21:24 AM