Get the best compiled kernel for given dimensions. When jit=True, compiles a shape-specialized variant (like DeepGEMM). When jit=False,

Create TMA descriptor with explicit swizzle parameter

Create TMA descriptor with explicit swizzle parameter

Benchmark a kernel variant, returns time in ms. Launches all iterations without per-iteration sync to allow the CUDA runtime to pipeline kern

Compile a kernel variant with specific compile-time parameters. Parameters: block_n: BLOCK_N tile size (required) shape_m:

Create TMA descriptor with oobFill=1 (zero fill for out-of-bounds)

Create TMA descriptor for K-major matrix with 128B swizzle. For K-major layout (K contiguous): - gmem_inner = K (con

Select the best BLOCK_N for given problem dimensions. Algorithm (matching DeepGEMM): 1. For each legal BLOCK_N candidate: - C

K must be aligned to BK for TMA swizzle

N stride must be aligned to 64 for TMA 128B swizzle on output

Barrier arrive on local smem

Benchmark with arbitrary M, N dimensions (K aligned to BK, C stride aligned to 64)

Benchmark our kernel with wall-clock timing (100 iterations to amortize host overhead).

Returns the relative dim3 block rank local to the cluster.

A: box = (64, 64) for each 64x64 tile load

A: box = (64, 128) for full BK x BM tile (single TMA load)

Create TMA descriptor for A (M x K), K-contiguous

B: box = (64, 256) for full BK x BN tile

B: box = (64, 128) for half tile (bidirectional multicast)

Create TMA descriptor for B_t (N x K), K-contiguous

Create TMA descriptor for store (no oobFill needed, but need proper bounds)

Create FP8 matrices and scale factors. For Level 1, all scale factors = 1.0 (UE8M0 value 127 = 0x7F). This means the GEMM output should match

Create FP8 input tensors and BF16 output tensor.