Module cublaslt 

Minimal cuBLASLt FFI for the high-efficiency GEMM path (ATLAS_CUBLAS_GEMM).

The hand-written mma.sync projection/MoE GEMMs reach only ~30% of the cuBLAS ceiling on GB10 (measured: 32 vs 85 TFLOPS bf16, 152 fp8, on the SSM-qkvz shape 3537×12288×2048). This routes those GEMMs through cuBLASLt instead. BF16 only for now — correctness-clean (no scale-format issues); native fp8 block-scaled is the follow-up once the end-to-end win is proven.

Functions

bf16_gemm_act_weight_t

Row-major out[M,N] = act[M,K] @ weight[N,K]ᵀ, all BF16 — the standard projection GEMM (activation × transposed weight). Maps to cuBLASLt’s column-major convention as D[N,M] = opT(weightᶜ[K,N]) · opN(actᶜ[K,M]).

fp8_gemm_act_weight_t_blkscaled

Native FP8 (E4M3) out[M,N] = act[M,K] @ weight[N,K]ᵀ → BF16, with the weight per-128×128-block FP32-scaled (matches Atlas’s Fp8Weight.row_scale layout exactly) and the activation cast at unit scale (post-RMSNorm acts sit in fp8’s range). ~1.8× the bf16 path (152 vs 85 TFLOPS on GB10).

fp8_gemm_act_weight_t_rowwise

Native FP8 (E4M3) out[M,N] = act[M,K] @ weight[N,K]ᵀ → BF16 with ROW-WISE scaling (OUTER_VEC): per-output-row weight scale weight_scale[N] and per-token activation scale act_scale[M]. This is the fp8 path GB10/sm_121 actually supports (128-block fp8 is B200-only). ~1.8× the bf16 path. cuBLAS folds A_scale[i]·B_scale[j] into the FP32 epilogue; with D=[N,M], i indexes weight rows (N) and j indexes tokens (M) — exactly row-wise.

prewarm

Force cuBLASLt’s one-time costs at MODEL LOAD instead of on request 1.