[SM120] Avoid per-token-count JIT recompilation in swapped GEMM paths

[ormandj]

Motivation

The SM120 small-M GEMM paths swap A/B and M/N before launch but pass compiled_dims through unchanged. This makes compile-time dimension labels refer to the swapped kernel axes instead of the caller's logical axes.

For the DeepSeek-V4 W_o_A shape, the logical token dimension is intended to remain dynamic. After the swap, the unchanged labels instead make the residual token count compile-time, producing a separate JIT kernel for each token count.

Observed failure

A controlled sequence sent 40 distinct, cache-busted prompts near 8K tokens through a server configured with an 8,192-token prefill chunk size.

One prompt contained 8,200 input tokens. After the first 8,192-token chunk, this left an eight-token residual. It was the first workload on that persistent cache to require the corresponding residual shape.

That request took 2.931 seconds. Nearby prompts containing 8,195–8,199 input tokens completed in approximately 1.03–1.08 seconds.

Server metrics located the additional time in prefill_forward and chunked_prefill on both TP ranks. During the request, the only newly created cache artifacts were the CUDA source and cubin for an SM120 FP8/FP4 BMM kernel.

The generated source appeared 891 ms after the request began, and the cubin completed 181 ms before the response completed. The generated CUDA specialized the dimensions as (0, 8, 4096), making the eight-token residual a compile-time dimension.

Modifications

• Remap m and n in compiled_dims when an SM120 path swaps operands.
• Apply the mapping to both FP8 BMM/einsum and FP8/FP4 GEMM.
• Leave other dimension labels and all non-swap paths unchanged.
• Add SM120 regression coverage for cold-cache kernel reuse, generated compile-time dimensions, numerical correctness, and execution through the existing sgl_deep_gemm test runner.

Direct kernel validation

A cold-cache diagnostic used the production W_o_A shapes: activation T × 4 × 4096, weight 4 × 1024 × 4096, and T=1..32.

The unpatched path generated 32 kernels, with compile-time shapes (0, T, 4096). The patched path generated one kernel with compile-time shape (1024, 0, 4096).

For token counts 3 and 11, the regression test computes independent BF16 PyTorch references:

• BMM/einsum: torch.einsum("bhr,hdr->bhd", activation, weight)
• GEMM: activation @ weight.T

The DeepGEMM paths consume the corresponding UE8M0-scaled FP8 operands and produce BF16 output.

DeepGEMM's calc_diff converts the output and reference to FP64 and computes:

1 - 2 * dot(output, reference) / (||output||² + ||reference||²)

Both regression-test paths require a value below 1e-3. The separate T=1..32 diagnostic checked outputs at T=1,8,16,32 and observed values from 0.000695 to 0.000720.

Serving validation

The same 40 saved request bodies were replayed in the same order on the patched server with an empty DeepGEMM runtime cache.

The first request incurred the expected general first-use kernel compilation. After that initialization, requests 2–40 completed in 0.869–0.897 seconds. The 8,200-token prompt that previously triggered the 2.931-second compilation completed in 0.891 seconds.

The patched cache contained the runtime-token kernel shape (1024, 0, 4096) and did not contain a (0, 8, 4096) residual-token specialization.

A previously unseen residual token count therefore no longer caused a new W_o_A kernel compilation. General first-use JIT compilation remains unchanged.

Prepared with AI assistance.