library_name: kernels
license: apache-2.0
tags:
  - cuda
  - native-cuda
  - flashrt
  - moe
  - nvfp4
  - blackwell

grouped-moe-gemv

Native CUDA FlashRT grouped MoE GEMV kernels for dynamic device-side routing. Version 2 covers BF16-activation/NVFP4-weight W4A16 and fully NVFP4 W4A4.

Hardware support is backend-specific:

• SM110 (Jetson AGX Thor): native FlashRT W4A16 decode/grouped GEMV, built with the SM110-only kUnroll=2 tuning from FlashRT PR #169;
• SM120/SM121: W4A16 plus block-scaled-MMA W4A4;
• W4A4 rejects SM110 explicitly instead of dispatching an incompatible device image or silently falling back.

Available functions:

• w4a16_decode_gemv_bf16
• grouped_w4a16_gemv_bf16
• quantize_activations_nvfp4_bf16
• quantize_weights_nvfp4_bf16
• grouped_w4a4_gemv_bf16
• grouped_w4a4_gemv_from_bf16

W4A4 routing is token-major: packed activations [M,K/2], device indices [M,top_k], packed expert library [E,N,K/2], output [M,top_k,N]. Use M=routed_pairs, top_k=1 when each expert receives a distinct activation. See the repository README for static-buffer and CUDA Graph usage.

Load version 2 with both current and legacy kernels clients:

from kernels import get_kernel

try:
    moe = get_kernel(
        "flashrt/grouped-moe-gemv", version=2, trust_remote_code=True
    )
except TypeError:  # kernels==0.12.x
    moe = get_kernel("flashrt/grouped-moe-gemv", version=2)

Use quantize_activations_nvfp4_bf16 once for [M,K], followed by grouped_w4a4_gemv_bf16 for all [M,top_k] routes. The convenience function grouped_w4a4_gemv_from_bf16 performs both calls and accepts preallocated packed, sfa, and out buffers for CUDA Graph capture.

On SM110, call grouped_w4a16_gemv_bf16 with BF16 activations and the packed expert library. This path is allocation-free with a caller-provided out buffer and supports device-side expert_idx mutation during graph replay.