Module cutlass 

Optional CUTLASS host-wrapper FFI for de-risking GB10 GEMM replacements.

★ REFERENCE IMPLEMENTATION — A BENCHMARK TARGET, NOT A DEPENDENCY.

Atlas ships its OWN kernels. CUTLASS is wrapped here for exactly one purpose: to be the opponent we measure against and beat. Nothing in a default build or a default serve calls a single line of it.

Two independent gates keep that true, and BOTH must survive any edit:

1. COMPILE TIME — every item in this module is #[cfg(atlas_cutlass)], and build.rs sets that cfg only when CUTLASS_HOME is exported. A build without it links no CUTLASS object at all.
2. RUNTIME — the dispatch arms are opt-in behind ATLAS_CUTLASS_GEMM=1. The default is OFF.

So the honest reading of an Atlas performance number is that Atlas kernels produced it, because a default binary cannot reach this code. Export the env var and you are measuring CUTLASS — label the number that way.

★ WHY KEEP IT COMPILED-BUT-DARK. Agentic benchmarking. An optimisation claim needs a credible opponent: “faster than our own previous commit” is a far weaker statement than “faster than CUTLASS on this shape”. Keeping the wrapper one env var away lets any agent A/B a shape against the industry reference on the same box, same checkpoint, same stream — which is the only comparison worth quoting.

Do NOT promote any of this to a default path. If a CUTLASS shape beats ours, the correct response is to make OUR kernel faster and re-measure.

Split for the ≤500 LoC cap: this root holds the shared FFI extern block, the workspace Ctx, and module wiring; the public wrappers live in the gemm (dense BF16 + NVFP4), grouped (per-expert MoE), and pack (weight pack / SFB swizzle / transpose) siblings. The public API (spark_runtime::cutlass::<fn>) is preserved via the re-exports below.

Functions

bf16_gemm_act_weight_t

Row-major out[M,N] = act[M,K] @ weight[N,K]^T, all BF16.

nvfp4_gemm_bf16_act_weight_t

Native CUTLASS NVFP4 dense projection: out[M,N] = quant_nvfp4(act[M,K]) @ weight_t[N,K]^T -> BF16.

nvfp4_grouped_down

Single-launch grouped NVFP4 DOWN projection (atlas_cutlass_nvfp4_grouped_down). a is the post-SiLU bf16 intermediate [M_total, K=inter], ALREADY expert-contiguous (no gather). packed_ptrs/sfb_ptrs are device-pointer arrays into the [N=hidden,K/2] packed + swizzled-SFB down tables; scale2_vals and expert_offsets_host are HOST arrays. Writes c [M_total, N=hidden].

nvfp4_grouped_gate_up

Grouped (per-expert) NVFP4 fused gate_up GEMM — Holo MoE Phase-1 escape-hatch path. Dispatches the proven Sm120 NVFP4 collective once per active expert over its token slice; bit-faithful to nvfp4_gemm_bf16_act_weight_t (it IS that collective), at one launch per expert. Used to validate that the FP4 math integrates correctly in grouped form before the hand-rolled block-scaled mma (Phase 2).

nvfp4_grouped_gate_up_fused

Single-launch grouped (GemmUniversalMode::kGrouped) NVFP4 fused gate_up GEMM — the Phase-2 successor to nvfp4_grouped_gate_up. Replaces the per-expert collective loop with ONE grouped launch over all active experts, eliminating the N-launch overhead.

pack_bf16_weight_to_nvfp4_t

Pack BF16 row-major weight [N,K] into the native CUTLASS NVFP4 layout: packed [N,K/2] (N-major, K-contiguous — NOT the Atlas transposed [K/2,N]) and E4M3 scales [K/16,N]. weight_scale_2 is assumed to be 1.0 by the caller when feeding this into the native CUTLASS wrapper.

pack_weight_sfb

Repack an Atlas E4M3 weight scale into the CUTLASS SM120 blockscaled SFB swizzle atom (tile_atom_to_shape_SFB, ue4m3) that the grouped collective reads. M-independent (the SFB atom depends only on N,K) so this runs once per expert at load. scale_out must hold the swizzled SFB region the grouped kernel consumes.

transpose_nvfp4_packed_kton

Transpose an Atlas-packed NVFP4 weight from the checkpoint/hand-kernel [K/2, N] layout into CUTLASS’s [N, K/2] layout (the byte order the native NVFP4 GEMM consumes for the ColumnMajor B operand). Pure byte transpose; nibble pairing within each byte is preserved. dst_packed must have N * K/2 bytes.