Add BGMV MoE CUDA kernels for multi-LoRA (#3249)

Co-authored-by: Claude

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

## 📌 Description

This PR adds fused CUDA kernels for applying multiple LoRA adapters
through MoE expert routing. When serving multiple LoRA adapters with MoE
models, each (token, expert) pair needs to be routed through the correct
LoRA adapter. These kernels fuse this operation into two efficient CUDA
kernels:
* Shrink (`bgmv_moe_shrink`): Projects input through LoRA-A matrices --
compute-bound, uses async pipeline with RANK_TILE tiling and multi-pair
blocking
* Expand(`bgmv_moe_expand`): Projects through LoRA-B matrices with
routing weights -- memory-bound, uses warp-level reduction with
atomicAdd

Key optimizations:
* Multi-pair blocking (PPB=4) for decode regime, single-pair for prefill
* 3-stage `cp.async` pipeline on SM90+ (216 KB shared memory on H100)
* RANK_TILE=8 tiling for X reuse across rank elements
* Pointer indirection for zero-copy multi-adapter access

Main changes:
* Add framework-agnostic CUDA kernel source with TVM-FFI bindings
(`csrc/bgmv_moe/`)
* Add Python wrapper (`flashinfer/fused_moe/bgmv_moe.py`)
* Add JIT loader integration using FlashInfer's TVM-FFI convention
(`flashinfer/jit/bgmv_moe.py`)
* Add correctness tests against PyTorch reference
(`tests/moe/test_bgmv_moe.py`)
* Add performance benchmark (`benchmarks/bench_bgmv_moe.py`) and
integration into `benchmarks/routines/moe.py`
* Register tests in CI (`scripts/task_jit_run_tests_part5.sh`)
* Use canonical vec_dtypes.cuh instead of a local copy

Supported configurations:
* Data types: BF16, FP16 (with FP32 accumulation for expand)
* LoRA ranks: 8, 16, 32, 64
* Hidden dimensions: 384 to 28672

Out of scope for this PR:
* vLLM integration (Expert class, oracle registration)
* FP8 quantized LoRA weights
* CUDA graph capture support

## 🔍 Related Issues
N/A

## 🚀 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.).

Tested with:
```
# Build
# Install FlashInfer (JIT compilation handles kernel build automatically)
cd flashinfer
pip install -e .

# Correctness tests
FLASHINFER_DISABLE_VERSION_CHECK=1 python -m pytest tests/moe/test_bgmv_moe.py -v

# Performance benchmark (full comparison table)
# grouped_mm_bf16 baseline requires cuDNN >= 9.18:
pip install nvidia-cudnn-cu12>=9.18 --force-reinstall
export LD_LIBRARY_PATH=$(python -c "import nvidia.cudnn as c; print(c.__path__[0] + '/lib')"):$LD_LIBRARY_PATH
FLASHINFER_DISABLE_VERSION_CHECK=1 python benchmarks/bench_bgmv_moe.py
```

Results:
Correctness Tests 123 passed

📊 Performance (H100 80GB HBM3)
3.3–6.7x faster than FlashInfer grouped_mm_bf16 (kernel only), 3.6–9.5x
faster including sort overhead, across decode and prefill regimes.

**Setup:** 8 LoRA adapters, 128 experts, rank=32, top_k=2, BF16

| Config | GG-kern (μs) | GG-sort+kern (μs) | BGMV MoE (μs) | vs GG-kern
| vs GG-sort+kern |

|--------|-------------:|-------------------:|---------------:|-----------:|----------------:|
| Decode-1tok-LargeMoE | 217.5 | 318.6 | 39.1 | 5.56x | 8.14x |
| Decode-4tok-LargeMoE | 215.5 | 316.7 | 50.0 | 4.31x | 6.33x |
| Decode-8tok-LargeMoE | 215.0 | 314.9 | 51.3 | 4.19x | 6.14x |
| Decode-32tok-LargeMoE | 240.2 | 337.3 | 37.8 | 6.36x | 8.93x |
| Prefill-256tok-LargeMoE | 443.9 | 537.2 | 109.1 | 4.07x | 4.93x |
| Prefill-512tok-LargeMoE | 695.0 | 792.0 | 175.7 | 3.96x | 4.51x |
| Prefill-1024tok-LargeMoE | 1205.9 | 1300.8 | 323.4 | 3.73x | 4.02x |
| Decode-1tok-Nemotron | 215.3 | 316.6 | 48.2 | 4.47x | 6.57x |
| Decode-4tok-Nemotron | 215.3 | 316.6 | 68.2 | 3.16x | 4.64x |
| Decode-8tok-Nemotron | 216.7 | 318.5 | 68.7 | 3.15x | 4.63x |
| Decode-32tok-Nemotron | 233.5 | 327.8 | 42.3 | 5.52x | 7.75x |
| Prefill-256tok-Nemotron | 417.2 | 513.2 | 109.1 | 3.82x | 4.70x |
| Prefill-512tok-Nemotron | 643.3 | 736.4 | 177.5 | 3.62x | 4.15x |
| Prefill-1024tok-Nemotron | 1082.3 | 1177.6 | 329.9 | 3.28x | 3.57x |


**Benchmark notes:**
* **GG-kern** = FlashInfer `grouped_mm_bf16` kernel only (cuDNN,
pre-sorted input, no sort overhead).
* **GG-sort+kern** = FlashInfer `grouped_mm_bf16` with token sorting
(sort + kernel). When combining multi-LoRA × MoE, it creates `num_loras
× num_experts` groups (e.g., 8 × 128 = 1024) with 0-1 tokens each — a
poor fit for grouped GEMM's large-group optimization.
* **BGMV MoE** = This PR's CUDA kernel. Handles arbitrary (token,
expert, lora_id) routing directly via pointer indirection — no sorting
required.
* **LargeMoE** = hidden=3072, rank=32, 128 experts (GPT-OSS-120B)
* **Nemotron** = hidden=2688, rank=32, 128 experts
(Nemotron-Nano-3-30B-A3B)
* The kernel's advantage comes from architectural differences:
multi-pair blocking (PPB=4) for decode, 3-stage `cp.async` pipeline with
explicit shared memory control, warp-level reduction (no minimum tile
size constraint), and pointer indirection for zero-copy multi-adapter
access.
* All times are median of 100 runs after 10 warmup iterations.

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

## Release Notes

* **New Features**
* Added BGMV Mixture-of-Experts CUDA kernels with support for FP16 and
BF16 data types
* Exported new Python APIs: `bgmv_moe`, `bgmv_moe_shrink`,
`bgmv_moe_expand`, and `fill_w_ptr`
  * Added JIT compilation support for MoE kernels

* **Tests**
* Added test suite for MoE operations with correctness validation and
edge case coverage
  * Added benchmarking utilities for MoE performance evaluation

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3249)
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Co-authored-by: Zihao Ye <expye@outlook.com>