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