flashinfer-gptoss-sm90-bf16-mxfp4-standalone

Standalone extraction of the FlashInfer CUTLASS SM90 MoE grouped GEMM path used by vLLM GPT-OSS for BF16 activations and MXFP4 weights.

This project uses FlashInfer/NVIDIA kernel sources directly (not a reimplementation), centered on:

• third_party/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl
• third_party/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws_mixed_dtype.h
• third_party/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu

Target numeric path:

• activations: bf16
• weights: fp4 e2m1
• scales: ue8m0
• architecture: sm90

vLLM reference path:

• vllm/vllm/model_executor/layers/quantization/mxfp4.py (SM90_FI_MXFP4_BF16, _interleave_mxfp4_cutlass_sm90)

API

Header: include/moe_cutlass_sm90.h

cudaError_t moe_cutlass_sm90_bf16_mxfp4_launch_device(
    cudaStream_t stream,
    const __nv_bfloat16* activations_bf16,
    const uint8_t* weights_fp4_e2m1_packed,
    const uint8_t* scales_ue8m0,
    const int64_t* expert_first_token_offset_device,
    int32_t num_experts,
    int32_t hidden_size,
    int32_t out_features,
    __nv_bfloat16* output_bf16);

// Compatibility wrapper (accepts host tokens_per_expert and internally builds
// device offsets before dispatching to the device API).
cudaError_t moe_cutlass_sm90_bf16_mxfp4_launch(
    cudaStream_t stream,
    const __nv_bfloat16* activations_bf16,
    const uint8_t* weights_fp4_e2m1_packed,
    const uint8_t* scales_ue8m0,
    const int32_t* tokens_per_expert_host,
    int32_t num_experts,
    int32_t hidden_size,
    int32_t out_features,
    __nv_bfloat16* output_bf16);

Input Contract

This is the exact layout contract expected by the extracted SM90 kernel.

1. activations_bf16 (const __nv_bfloat16*)

• Shape: [total_tokens, hidden_size]
• total_tokens = sum(tokens_per_expert_host)
• Row-major contiguous BF16.
• Rows must already be grouped by expert.

2. weights_fp4_e2m1_packed (const uint8_t*)

• Logical weight shape: [num_experts, out_features, hidden_size] in FP4 E2M1.
• Physical byte shape: [num_experts, out_features, hidden_size / 2].
• Packing rule for each fixed (expert, out_feature) row:
  • even k goes to low nibble (bits 3:0)
  • odd k goes to high nibble (bits 7:4)
• K is the fast axis inside each (expert, out_feature) row.

3. scales_ue8m0 (const uint8_t*)

• MXFP4 group size is 32 K-elements.
• Start from logical scales src[e, n, g] with shape [num_experts, out_features, hidden_size/32].
• Convert to SM90 interleaved layout dst[e, g_outer, n4] with shape [num_experts, hidden_size/128, out_features * 4] using:
  • g_outer = g / 4
  • g_inner = g % 4
  • n4 = n * 4 + g_inner
  • dst[e, g_outer, n4] = src[e, n, g]
• Pass dst buffer as scales_ue8m0.

4. tokens_per_expert_host (const int32_t*)

• Host memory (CPU pointer), length num_experts.
• Defines row partition of grouped tokens.
• Prefix sum gives expert row ranges:
  • expert e owns rows [offset[e], offset[e+1]) where offset[0]=0, offset[e+1]=offset[e]+tokens_per_expert_host[e].

5. output_bf16 (__nv_bfloat16*)

• Shape: [total_tokens, out_features], row-major BF16.
• Uses the same grouped row partition as activations_bf16.

Device Metadata API

For FlashInfer-style fully device-driven launches, provide:

1. expert_first_token_offset_device (const int64_t*, device)

• Length num_experts + 1
• Prefix-sum row partition:
  • offset[0] = 0
  • offset[e+1] >= offset[e]
  • total_tokens = offset[num_experts]

2. Call moe_cutlass_sm90_bf16_mxfp4_launch_device(...)

• All tensors and metadata pointers are device pointers.

Constraints

1. Device must be SM90.x (Hopper). SM100+ is not supported by this extracted path.
2. hidden_size % 32 == 0.
3. For the scale interleave mapping above, hidden_size % 128 == 0 is required.
4. num_experts > 0, out_features > 0, and all tokens_per_expert_host[e] >= 0.

Build

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build -j

If nvcc is not on PATH, pass it explicitly:

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_COMPILER=/path/to/nvcc
cmake --build build -j

Test

ctest --test-dir build --output-on-failure