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