SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

Improved `simple` mamba SSU kernel (#2962)

<!-- .github/pull_request_template.md -->

## ๐Ÿ“Œ Description

This PR upgrades the SSU MTP "simple" kernel with **cp.async state
prefetching**, vectorized loads, and a consolidated state write path,
delivering drastic performance improvements โ€” particularly in the
latency spectrum. The `async_horizontal` kernel was a temporary
development vehicle used during implementation and benchmarking; once
the optimizations were validated, the simple kernel was replaced with
the improved version and the temporary kernel was removed.

## Summary

This PR upgrades the SSU MTP "simple" kernel with **cp.async state
prefetching**, vectorized loads, and a consolidated state write path,
delivering drastic performance improvements โ€” particularly in the
latency spectrum. The `async_horizontal` kernel was a temporary
development vehicle used during implementation and benchmarking; once
the optimizations were validated, the simple kernel was replaced with
the improved version and the temporary kernel was removed.

## Key changes

### Async state prefetch (cp.async โ†’ double-buffered smem)
- Replace direct global loads of `state_in` with `cp.async` into a
double-buffered shared memory staging area (`state_in[STATE_STAGES]`).
- First pass is prefetched during the load phase; subsequent passes are
pipelined at the end of each pass loop iteration.
- Extracted into a reusable `cp_async_state_cooperative` helper
function.

### Vectorized loads
- Use `PackedAligned` for vectorized loads of B, C, and x tensors in the
load path, improving memory access efficiency when padding is not
active.

### State write path consolidation
- Precompute per-step `state_dst_slots[]` during the load phase,
eliminating redundant index recomputation per pass/dd.
- Replace three separate state-write branches (intermediate states,
per-step dst indices, final state) with a single unified `dst_slot !=
SKIP` path.
- Deduplicate encode-scale computation (was computed up to 3ร— per step).

### OOB handling cleanup
- Remove upfront shared memory zero-fill padding. Instead, zero OOB
padding columns directly in registers at load time. This eliminates an
extra `__syncthreads__` barrier.

### Latency hiding
- Hoist A and D global loads before the barrier to overlap with smem
wait.
- Move `dst_slot` prefetch earlier to hide LDS latency.
- Use `mul_f32x2` for state decode scale.

### Varlen + scaled-state support
- Remove guards that blocked the async_horizontal path from running with
`cu_seqlens` or scaled (quantized) state.
- Refactor smem layout: replace sub-tile-major `BANK_CYCLE_ELEMS` scheme
with a simpler `DSTATE_PAD` (128-byte aligned) wide tile.

### Validation
- Add mutual exclusion check: `intermediate_states_buffer` and
`dst_state_batch_indices` cannot both be provided.

### Benchmarking
- Add `bench_ssu_sweep_sol.py` โ€” SOL (speed-of-light) benchmark script
for SSU MTP mode.
<img width="3000" height="1500"
alt="sol_vs_batch_size_mtp6_bf16_NVIDIA_B200"
src="https://github.com/user-attachments/assets/33b53f4e-7e36-41c6-9360-ae0d905c5495"
/>


## ๐Ÿš€ 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.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


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

* **New Features**
* SM100-optimized selective-state-update path with new algorithms:
vertical, horizontal, and async-horizontal.
* Comprehensive benchmarks that measure runtime and "speed-of-light"
(SOL) estimates with automatic plotting.

* **Improvements**
* Improved dtype/format conversion, alignment checks, and configurable
tensor-map out-of-bounds fill.
* Input validation for mutually exclusive options, optional stochastic
rounding, and support for larger state widths (dstate=96).

* **Tests**
* Expanded tests covering SM100 capability, new algorithms,
varlen/pad-slot cases, and varied ngroups.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
I
Igor Shovkun committed
4568160cad50f234090326c45b02c8053ca21a3d
Parent: 6aba765
Committed by GitHub <noreply@github.com> on 5/13/2026, 12:09:14 AM