SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

feat: Add CuTe DSL grouped-gemm + combine fusion support (#2944)

## Summary

Enables https://github.com/sgl-project/sglang/pull/21877

- Add fused grouped-GEMM + combine operation to the CuTe DSL Blackwell
masked GEMM kernel, enabling the GEMM epilogue to perform weighted
scatter-reduce (combine) directly into multi-rank output buffers using
`cp.reduce.async.bulk` PTX instructions (`bf16` and `f32` variants)
- Extend `MaskedScheduler` with `is_swap_ab` support so the scheduler
correctly computes tile coordinates when A/B inputs are swapped (needed
for combine fusion where output is M-major)
- Add `barrier_flag_local`/`barrier_flag_multicast` parameters for
cross-rank synchronization via spin-lock barriers
- Add multi-GPU test covering the fused grouped-gemm + combine path

## Key changes
- **`flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py`**: New
`cp_reduce_bf16_add`/`cp_reduce_f32_add` PTX wrappers, custom
`make_fused_smem_layout_epi` for the combine path, `is_swap_ab` logic in
scheduler, and new parameters (`topk_weights`, `idx_src_info`,
`rank_src_info`, `out_ptrs`, barrier flags) threaded through the full
stack
- **`tests/comm/test_multi_gpu_cute_dsl_blockscaled_gemm_fusion.py`**:
New multi-GPU test that validates the fused gemm+combine against a
reference implementation using `mpirun`

## Test plan
- [ ] `mpirun -np 4 pytest
tests/comm/test_multi_gpu_cute_dsl_blockscaled_gemm_fusion.py` on SM100+
hardware
- [ ] Verify existing grouped-gemm tests still pass (non-fusion path
unchanged)

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

* **New Features**
* Block-scaled GEMM combine-fusion: fused weight prefetch, per-tile
async reduction stores, and optional AB swap.
* Host API extended with optional top-k weights, routing info,
output-pointer inputs, rank/count, and barrier-flag controls for
distributed fusion runs.

* **Tests**
* Added multi-GPU integration test validating the combine-fusion path,
routing, and distributed synchronization.

* **Bug Fixes**
* Persistent scheduler sizing improved by capping active clusters to
available tiles; stage computation updated for fusion prefetch buffers.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
N
Nicolas Castet committed
89bfb6882edbf21c8dd4fee7fa719f6206e9de36
Parent: 7ac3ccc
Committed by GitHub <noreply@github.com> on 5/14/2026, 7:32:29 PM