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 -->
[](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