Module cutlass

Module cutlass 

Source
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), which build.rs sets only when CUTLASS_HOME is set at build time. Without it every public wrapper returns an error.
  • A wrapper that returns Ok got 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 by weight_scale_2.
nvfp4_grouped_down
2026-09-25: Grouped NVFP4 down projection, one kGrouped launch. a is the bf16 intermediate [M_total, K], already in expert order (no gather). packed_ptrs and sfb_ptrs hold one device pointer per expert to a [N,K/2] weight and its swizzled SFB scales; scale2_vals and expert_offsets_host are host arrays. Writes c [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_t once for gate and once for up, so the result matches nvfp4_gemm_bf16_act_weight_t on each expert’s rows.
nvfp4_grouped_gate_up_fused
2026-09-25: Grouped NVFP4 gate and up GEMMs, one GemmUniversalMode::kGrouped launch per projection over the active experts. a is 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 pass weight_scale_2 = 1.0 to nvfp4_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 on n and k, so it is built once per expert at load (MoeLayer::build_cutlass_grouped_sfb). scale_out must 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 that nvfp4_gemm_bf16_act_weight_t reads. A byte transpose: the two nibbles of each byte stay together. dst_packed must hold N * K/2 bytes.