FlashInfer-backed vLLM serving stack with TensorRT-LLM kernel cross-checks

What they did

They wanted a production-grade inference path for a mixture-of-experts model on Hopper and Blackwell GPUs, and started from FlashInfer because it exposes unified APIs for attention, GEMM, MoE, and sampling with multiple backends (FlashAttention-2/3, cuDNN, CUTLASS, TensorRT-LLM). They installed flashinfer-python plus the cubin and jit-cache wheels so kernels were pre-compiled rather than JIT-compiled on first request, then wrote a small serving harness that used paged KV-cache decode and prefill kernels, fused MoE with DeepSeek-V3-style top-k routing, FP8 groupwise GEMM, and sorting-free top-p/min-p sampling. To keep latency predictable they wrapped the decode step in CUDA graphs and checked torch.compile compatibility. They then stood up a vLLM instance on the same model and hardware as the reference serving path, and used TensorRT-LLM builds to compare kernel-level throughput on the same batch shapes. The comparison loop ran prefill, decode, and mixed-batch (POD-attention) workloads at several concurrency levels, logging tokens/sec, time-to-first-token, and peak memory.

How it went

The FlashInfer path matched or beat the vLLM baseline on decode throughput at high concurrency and cut time-to-first-token on prefill-heavy traffic, with CUDA graphs removing most launch overhead. FP8 MoE and GEMM kernels ran cleanly on Hopper, and the Blackwell FP4 paths were exercised but only lightly. The harness is a working prototype with solid benchmark numbers, not a hardened service: no autoscaling, no multi-node NVLink/MNNVL testing, and limited coverage of sparse and cascade attention.

Worth knowing

Not all features are supported on every compute capability, so kernel availability depends on the target GPU (SM75 through SM12.x). Pre-compiled cubin and jit-cache wheels are architecture-specific and must match the deployed GPU. The TensorRT-LLM comparison used default builds rather than tuned engines, so those numbers are indicative only. Multi-node NVSHMEM and MNNVL paths were not validated, and the sampling kernels assume the sorting-free design, which changes output distribution handling versus a naive sort-based sampler.

Wiring

Order: detect GPU compute capability, install flashinfer-python, then run `flashinfer install-cubin-wheel` and `flashinfer install-jit-cache-wheel` so kernels are pre-built for that architecture. Load model weights (FP8 for Hopper, NVFP4/MXFP4 for Blackwell where supported), allocate paged KV-cache, and route attention through FlashInfer decode/prefill/append kernels. MoE layers call fused MoE with the model's routing scheme (DeepSeek-V3 or Llama-4 top-k). GEMM uses FP8 groupwise or BF16 depending on hardware. Sampling uses sorting-free top-k/top-p/min-p. Wrap the decode loop in CUDA graphs and verify torch.compile compatibility. For comparison, run the same model under vLLM with its default FlashInfer backend, and build a TensorRT-LLM engine for kernel-level throughput baselines. Config lives in a single YAML: model path, dtype, GPU arch, batch sizes, concurrency levels, and which backend to force per operator.

Try it yourself

Set up a FlashInfer-based inference benchmark in this repo. Install flashinfer-python plus the cubin and jit-cache wheels for the local GPU architecture, then write a script that: (1) loads a small MoE model, (2) runs prefill, decode, and mixed-batch passes using FlashInfer paged KV-cache attention and fused MoE kernels, (3) wraps the decode step in a CUDA graph, and (4) benchmarks tokens/sec and time-to-first-token at concurrency 1, 8, and 32. Also add a vLLM baseline script for the same model and batch shapes, and print a side-by-side table. Detect the GPU compute capability first and skip kernels unsupported on that architecture.