SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

Restore monolithic CuTe-DSL MLA decode alongside modular, gated by cute_dsl_impl= (#3296)

## Summary

PR #2805 refactored the monolithic CuTe-DSL MLA decode kernel into a
modular structure and removed the original implementation. The original
authors want it kept available because the modular path is still
maturing. This PR restores it under the existing
\`backend=\"cute-dsl\"\` user surface (no new backend name) and exposes
implementation selection via a new \`cute_dsl_impl=\` keyword argument
on \`trtllm_batch_decode_with_kv_cache_mla\`.

- **\"auto\"** (default): monolithic by default, automatically promoted
to modular when the call uses a modular-only feature (currently
\`sinks\`).
- **\"modular\"**: strict, always run the modular kernels.
- **\"monolithic\"**: strict, always run the monolithic kernels; raises
\`ValueError\` if the call uses any modular-only feature.

The dispatcher strips modular-only kwargs (\`sinks=None\`) before
forwarding to monolithic, so callers can pass \`sinks=\` unconditionally
without breaking the monolithic path.

### Sinks support on cute-dsl backend

\`trtllm_batch_decode_with_kv_cache_mla(sinks=...)\` on
\`backend=\"cute-dsl\"\` now constructs an \`AttentionWithSink\` variant
inside the modular standalone, instead of being rejected at the API
boundary. \`AttentionWithSink\` gained value-based
\`__hash__\`/\`__eq__\` (keyed on \`(type, shape, dtype)\`) so
\`@functools.cache\` on \`_compile_mla_kernel\` correctly reuses
compiled kernels across invocations with the same shape — without this,
a fresh variant per call hashed by object identity, JIT-recompiled the
kernel on every iteration, and made cuda-graph + sinks bench
measurements appear to hang.

### Layout

\`\`\`
flashinfer/cute_dsl/attention/
  monolithic/         - restored kernels (verbatim from before #2805,
                        relocated to live next to the modular code).
                        Includes the H<128 / Kimi K2.5 fix from #3235
                        backported.
  wrappers/           - existing modular standalone + wrapper.
  mla_dispatch.py     - new dispatcher in front of both impls.
\`\`\`

### Bench

\`benchmarks/bench_trtllm_gen_mla.py\` grows a focused 6-cell
\`with_sinks=True\` sub-sweep (B in {1,16,128} × S in {1024,8192} at
q_len=1, page=64, bf16) on top of the existing main sweep, instead of
doubling the full grid. Argument list deduplicated into a
\`common_kwargs\` dict so warmup and benchmark calls cannot drift.

Sinks overhead is ~free on both backends (worst case +1.9% at the
smallest cell). Cross-backend ranking does not change with sinks
enabled.

## Test plan

Existing standalone and public-API tests in
\`tests/attention/test_cute_dsl_mla_decode.py\` now parametrize over
modular/monolithic via a \`cute_dsl_impl\` fixture, doubling coverage on
the same shapes. New minimal sinks tests pin the auto/modular dispatch
branches and the monolithic+sinks \`ValueError\` contract. Wrapper sinks
numerics remain covered by the pre-existing
\`test_cute_dsl_mla_decode_attention_sink\`.

- [x] All pre-commit hooks pass on changed files (mypy, ruff check, ruff
format, EOF, whitespace, etc.).
- [x] \`pytest tests/attention/test_cute_dsl_mla_decode.py -v\` — full
sweep (544 cases incl. parametrized modular/monolithic) passes on B200.
- [x] \`pytest tests/attention/test_cute_dsl_mla_decode.py -k sinks\` —
3 new sinks integration tests pass.
- [x] \`pytest tests/attention/test_trtllm_gen_mla.py -v\` — unaffected,
passes.
- [x] H=64 / Kimi K2.5 backport on monolithic exercised via the existing
\`num_heads ∈ [128, 64]\` parametrization ×
\`cute_dsl_impl=monolithic\`.
- [x] Bench \`benchmarks/bench_trtllm_gen_mla.py --backend cute-dsl\`
and \`--backend trtllm-gen\` both run cleanly through the focused sinks
sub-sweep.

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

* **New Features**
* Optional "sinks" support for an alternate softmax path and a
cute_dsl_impl option to choose/auto-select modular vs monolithic MLA
decode implementations.
  * New monolithic CuTe-based MLA kernel targeting Blackwell hardware.

* **Performance / Reliability**
* Improved kernel caching/variant keying to enable reuse across variant
instances.
* Benchmark updated to exercise sinks-enabled and sinks-disabled paths.

* **Tests**
* Added tests for implementation selection, sinks behavior, and related
error/shape validations.

<!-- review_stack_entry_start -->

[![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/3296?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
P
Prasun Gera committed
e91ac8f15da1d64e2e61ae7373e4943389e7b768
Parent: 6f651b6
Committed by GitHub <noreply@github.com> on 5/21/2026, 8:20:49 PM