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