Expand description
2026-09-25: Host-side FFI to the CUTLASS GEMM wrappers in cuda/cutlass_*.cu:
the shared extern block and workspace here; dense BF16 and NVFP4 GEMMs in
gemm, grouped per-expert NVFP4 MoE GEMMs in grouped, and NVFP4 weight
pack, SFB swizzle and transpose in pack.
Owner: gpu-runtime. Invariants:
- The FFI and the workspace exist only under
cfg(metrale_cutlass), whichbuild.rssets only whenCUTLASS_HOMEis set at build time. Without it every public wrapper returns an error. - A wrapper that returns
Okgot status 0 from its C entry point.
Functions§
- bf16_
gemm_ act_ weight_ t - 2026-09-25: Row-major
out[M,N] = act[M,K] @ weight[N,K]^T, all BF16. - nvfp4_
gemm_ bf16_ act_ weight_ t - 2026-09-25: NVFP4 dense projection with BF16 output:
out[M,N] = quant_nvfp4(act[M,K]) @ weight[N,K]^T, scaled byweight_scale_2. - nvfp4_
grouped_ down - 2026-09-25: Grouped NVFP4 down projection, one
kGroupedlaunch.ais the bf16 intermediate[M_total, K], already in expert order (no gather).packed_ptrsandsfb_ptrshold one device pointer per expert to a[N,K/2]weight and its swizzled SFB scales;scale2_valsandexpert_offsets_hostare host arrays. Writesc[M_total, N]. - nvfp4_
grouped_ gate_ up - 2026-09-25: Per-expert NVFP4 gate and up GEMMs: for each expert with rows,
the C side calls
metrale_cutlass_nvfp4_gemm_bf16_act_weight_tonce for gate and once for up, so the result matchesnvfp4_gemm_bf16_act_weight_ton each expert’s rows. - nvfp4_
grouped_ gate_ up_ fused - 2026-09-25: Grouped NVFP4 gate and up GEMMs, one
GemmUniversalMode::kGroupedlaunch per projection over the active experts.ais packed once and shared by both. - pack_
bf16_ weight_ to_ nvfp4_ t - 2026-09-25: Pack a row-major BF16 weight
[N,K]into CUTLASS NVFP4: packed[N,K/2](K-contiguous) and E4M3 scales[K/16,N], each scale the group’s max magnitude / 6. The scales carry no second-level factor, so passweight_scale_2 = 1.0tonvfp4_gemm_bf16_act_weight_t. - pack_
weight_ sfb - 2026-09-25: Swizzle an E4M3 weight scale into the CUTLASS SFB layout
(
tile_atom_to_shape_SFB, ue4m3) that the grouped GEMM reads. The layout depends only onnandk, so it is built once per expert at load (MoeLayer::build_cutlass_grouped_sfb).scale_outmust hold the whole swizzled region. - transpose_
nvfp4_ packed_ kton - 2026-09-25: Transpose a packed NVFP4 weight from
[K/2, N]into the CUTLASS[N, K/2]layout thatnvfp4_gemm_bf16_act_weight_treads. A byte transpose: the two nibbles of each byte stay together.dst_packedmust holdN * K/2bytes.