The Dimension Convention Bug: Debugging FP4 GEMM Performance Across Two Libraries

Introduction

In the course of optimizing inference throughput for the GLM-5-NVFP4 large language model on eight NVIDIA RTX PRO 6000 Blackwell GPUs, the assistant encountered a subtle but instructive bug while benchmarking FP4 matrix multiplication performance. Message 875 in this opencode session captures a moment of diagnostic insight: the assistant realizes that two different GPU kernel libraries—FlashInfer and sgl_kernel—implement the same mathematical operation (FP4 general matrix multiplication, or GEMM) but with incompatible tensor dimension conventions. What follows is a careful fix, a re-run of the benchmark, and a set of results that reveal not only which library is faster but also why the comparison initially failed.

This message is a microcosm of the kind of debugging that defines high-performance ML engineering: it involves reading error messages, understanding API contracts, reasoning about tensor layouts, and systematically eliminating variables until the true signal emerges. The message is outwardly simple—a bash command that writes a Python script and executes it—but the reasoning embedded in it is rich with assumptions, domain knowledge, and iterative problem-solving.

The Context Leading Up to Message 875

To understand why this message was written, we must trace the investigation that preceded it. The assistant had been engaged in a multi-session effort to deploy the GLM-5-NVFP4 model—a Mixture-of-Experts (MoE) architecture with 256 experts and FP4 quantized weights—on a server with eight RTX PRO 6000 Blackwell GPUs. After resolving a cascade of infrastructure issues (driver installation, CUDA toolkit conflicts, flash-attn compilation failures, virtualization-induced PCIe bottlenecks), the assistant had finally achieved a working SGLang server and was now focused on optimizing throughput.

The critical question was: how efficiently are the FP4 GEMM kernels running on the Blackwell architecture (SM120)? The model's weights are stored in FP4 (4-bit floating point) format, which offers a theoretical 4× compute density advantage over FP8 or BF16. But theoretical peak performance is rarely achieved in practice, especially for the small matrix sizes that dominate MoE inference.

In messages 866–869, the assistant conducted a systematic microbenchmark of FlashInfer's mm_fp4 function across a range of matrix sizes relevant to GLM-5. The results were sobering. For per-expert batch sizes typical of decode (16–64 tokens), the FP4 kernels achieved only 0.8–55 TFLOPS—between 0.02% and 3% of the estimated 2,000 TOPS theoretical peak of the Blackwell GPU. Even at larger batch sizes, the kernels plateaued at around 1,200–1,300 TFLOPS, or roughly 65% of peak. The root cause was clear: the 256-expert MoE architecture scatters tokens too thinly across experts, leaving each individual GEMM operation operating on tiny matrices that cannot saturate the GPU's compute units.

In messages 870–874, the assistant shifted focus to comparing two different FP4 kernel implementations: FlashInfer's mm_fp4 (with the CUTLASS backend) and sgl_kernel's cutlass_scaled_fp4_mm. Both are open-source CUDA kernel libraries that provide FP4 matrix multiplication, but they are developed independently and may have different performance characteristics. The assistant wanted to know which one was faster for the specific matrix sizes encountered during GLM-5 inference.

The first attempt (message 871) failed with a Python NameError due to shell escaping issues. The second attempt (message 872) failed the same way. The third attempt (message 873) successfully wrote the script to a file, avoiding shell interpolation problems. But the fourth attempt (message 874) crashed at runtime with an error from sgl_kernel's cutlass_scaled_fp4_mm function.

This brings us to message 875, where the assistant diagnoses the root cause of that crash.## The Error and Its Diagnosis

The error in message 874 was opaque to someone unfamiliar with the internal conventions of these kernel libraries. The traceback showed that sgl_kernel.cutlass_scaled_fp4_mm failed during execution, but the error message was cut off in the conversation log. However, the assistant's reasoning—visible in the opening line of message 875—reveals the diagnosis: "Different dimension convention. sgl_kernel expects b as (N, K//2) transposed."

This is a classic pitfall in GPU kernel programming. Matrix multiplication libraries often differ in how they expect the weight matrix to be laid out in memory. FlashInfer's mm_fp4 follows the standard mathematical convention where matrix b has shape (K//2, N)—that is, the reduction dimension (K) comes first, and the output feature dimension (N) comes second. sgl_kernel's cutlass_scaled_fp4_mm, by contrast, expects b to be pre-transposed with shape (N, K//2). The same transposition applies to the scale tensor b_scale.

Why would a library choose a non-standard layout? The answer lies in memory access patterns. When a GPU kernel computes C = A * B, it typically loads rows of A and columns of B. If the weight matrix B is stored in a transposed layout (with the output dimension as the leading dimension), the kernel can use vectorized loads that coalesce better across threads, improving memory bandwidth utilization. This is a well-known optimization: by changing the physical layout of the data, you can change the access pattern from strided to contiguous, which is critical for achieving peak memory throughput on NVIDIA GPUs.

The assistant's diagnosis is not stated as a guess or a hypothesis. It is presented as a confident assertion: "Different dimension convention." This confidence comes from the pattern of the error—a crash inside the kernel launch, not a shape mismatch error—and from prior experience with similar libraries. The assistant has likely encountered this exact issue before, or has enough familiarity with the sgl_kernel codebase (or its documentation) to know that it uses a transposed weight layout.

The Fix: Adapting the Benchmark Script

Having identified the issue, the assistant proceeds to fix the benchmark script. The key change is in how the b and b_scale tensors are prepared for sgl_kernel:

# sgl_kernel: a=(M, K//2), b=(N, K//2), scale_a=(M, K//16), scale_b=(N, K//16)
b_fp4_sgl = b_fp4_fi.t().contiguous()
b_scale_sgl = b_scale_fi.t().contiguous()

The .t() method transposes the tensor, and .contiguous() ensures the resulting tensor is laid out contiguously in memory (transposition in PyTorch produces a view with non-contiguous strides; calling .contiguous() forces a copy with the desired layout). This is a critical detail: if the tensor is not contiguous, the CUDA kernel may silently produce incorrect results or crash with an illegal memory access.

The assistant also wraps the sgl_kernel call in a try/except block, so that if a particular matrix size still fails (perhaps due to other constraints like shared memory limits or tile size restrictions), the benchmark gracefully records "ERROR" rather than aborting the entire script. This is a pragmatic engineering decision: the goal is to compare performance across a range of sizes, and a single failure should not derail the entire experiment.

The Results and Their Interpretation

The output of the corrected benchmark is revealing:

| Description | sgl_kernel | flashinfer | ratio | |---|---|---|---| | Expert batch=16 (decode) | ERROR | 42.6 us | N/A | | Expert batch=64 | ERROR | 31.1 us | N/A | | Expert batch=256 | 27.1 us | 30.0 us | 1.10x | | Expert batch=512 | 49.3 us | 29.5 us | 0.60x | | Expert batch=1024 | ... | ... | ... |

The output is truncated in the conversation log, but the pattern is already clear. For small batch sizes (16 and 64), sgl_kernel fails entirely—the ERROR indicates that the kernel cannot handle these tiny dimensions, likely because its tile configuration requires a minimum M dimension that exceeds the batch size. For batch=256, sgl_kernel is slightly faster than FlashInfer (27.1 vs 30.0 microseconds, a 1.10× speedup). But for batch=512, the tables turn dramatically: sgl_kernel takes 49.3 microseconds while FlashInfer takes only 29.5 microseconds, making FlashInfer 1.67× faster.

This non-monotonic behavior is typical of GPU kernels with different tile size heuristics. A kernel that is optimized for medium-sized matrices may have a tile configuration (e.g., M=128, N=256, K=16) that works well for M=256 but becomes inefficient at M=512 due to register pressure, shared memory bank conflicts, or reduced occupancy. Conversely, FlashInfer's kernel may use a more conservative tile size that scales better across a wider range of dimensions.

The most important finding for the GLM-5 inference use case is that sgl_kernel fails entirely at the batch sizes most relevant to MoE decode (16–64 tokens per expert). Even if sgl_kernel were faster at larger sizes, it cannot be used for the actual inference workload without either padding the batch dimension or finding a workaround. This is a hard constraint: the kernel must work correctly for the input sizes that the model actually produces.

Assumptions Made in This Message

The assistant makes several assumptions in crafting this fix and interpreting the results:

  1. That the dimension convention is the sole cause of the crash. The assistant assumes that once the tensors are correctly transposed, the kernel will run successfully for at least some matrix sizes. This assumption is validated by the results—batch sizes 256, 512, and 1024 do work after the fix.
  2. That transposing and making contiguous is a valid transformation. The assistant assumes that b_fp4_fi.t().contiguous() produces a tensor that is semantically equivalent for the GEMM operation. This is mathematically correct: A * B = A * (B^T)^T, so if sgl_kernel computes A * B_sgl where B_sgl has shape (N, K//2), then B_sgl should equal B_fi^T for the results to be equivalent.
  3. That the benchmark is measuring the right thing. The assistant assumes that microbenchmarking individual GEMM operations in isolation is a valid proxy for end-to-end inference performance. This is a reasonable assumption for compute-bound operations, but it ignores the effects of kernel launch overhead, CUDA graph capture, and memory bandwidth contention that occur in a real serving scenario.
  4. That the random input data is representative. The assistant uses random integers for the FP4 data and random float8 values for the scales. This is a standard practice for microbenchmarking, but it may not capture edge cases in the kernel's handling of special floating-point values (NaN, inf, denorm) that could occur with real model weights.
  5. That the comparison is fair. The assistant compares FlashInfer's mm_fp4 with backend="cutlass" against sgl_kernel's cutlass_scaled_fp4_mm. Both use CUTLASS under the hood, but they may use different tile configurations, different quantization schemes, or different accumulation strategies. The comparison is apples-to-apples only at the API level, not at the kernel level.

Mistakes and Incorrect Assumptions

While the message is largely correct, there are a few potential issues:

  1. The ERROR for small batch sizes is not fully investigated. The assistant records the error and moves on, but does not attempt to understand why sgl_kernel fails for M=16 and M=64. Is it a shared memory limitation? A minimum tile size constraint? A bug in sgl_kernel's handling of small dimensions? Understanding this could inform whether a workaround (e.g., padding the batch dimension) is feasible.
  2. The truncated output leaves the story incomplete. The conversation log shows the output cut off after "Expert batch=1024 ...". We never see the full comparison for the largest batch sizes. This may be because the output was too long for the terminal, or because the assistant chose to move on after seeing the key pattern. But it means the reader (and the assistant) are working with incomplete data.
  3. The benchmark does not warm up the GPU properly for all sizes. The script warms up each kernel call with 5 iterations before timing, but the warmup uses the same tensor shapes. For GPU kernels, the first few launches may include JIT compilation overhead that is not representative of steady-state performance. However, 5 iterations is typically sufficient for CUDA kernels that have already been compiled.
  4. The comparison does not control for numerical accuracy. The assistant checks that the kernels produce non-zero results (in message 866), but does not verify that sgl_kernel and FlashInfer produce the same numerical outputs for the same inputs. If one kernel uses a different accumulation order or rounding mode, the results could differ slightly, which might affect model quality in subtle ways.

Input Knowledge Required to Understand This Message

To fully grasp what is happening in message 875, a reader needs:

  1. Familiarity with FP4 quantization. FP4 is a 4-bit floating-point format where each value is stored in 4 bits (typically 1 sign bit, 2 exponent bits, 1 mantissa bit for E2M1, or similar). Two FP4 values are packed into each byte. The scale factors are stored in a higher-precision format (float8) and are shared across blocks of values (typically 16 elements per scale).
  2. Understanding of GEMM tensor layouts. Matrix multiplication C = A * B requires that the inner dimensions match. If A has shape (M, K) and C has shape (M, N), then B must have shape (K, N). The assistant's fix involves transposing B from (K//2, N) to (N, K//2), which changes which dimension is the reduction dimension.
  3. Knowledge of GPU kernel launch mechanics. The concept of kernel launch overhead, tile sizes, shared memory limits, and occupancy are all relevant to understanding why sgl_kernel fails for small batch sizes and why the performance ratio varies with matrix dimensions.
  4. Awareness of the MoE architecture of GLM-5. The model has 256 experts, with 8 experts activated per token. This means that during decode, each expert processes roughly batch_size * 8 / 256 tokens. For a batch size of 512 concurrent requests, each expert gets about 16 tokens—hence the focus on M=16 and M=64 benchmarks.
  5. Familiarity with the opencode tooling. The message uses bash commands, heredoc syntax (<< "PYEOF"), and remote SSH execution. Understanding that the assistant is writing a script to a remote machine and executing it is essential context.

Output Knowledge Created by This Message

This message produces several concrete pieces of knowledge:

  1. sgl_kernel's dimension convention is transposed relative to FlashInfer. Anyone using cutlass_scaled_fp4_mm must ensure their weight tensor has shape (N, K//2) rather than (K//2, N). This is a documented API requirement that the assistant discovered empirically.
  2. sgl_kernel fails for small batch sizes (M ≤ 64) on Blackwell GPUs. This is a critical constraint for MoE inference, where per-expert batch sizes are typically small. The failure mode is not a graceful fallback but a hard crash, making sgl_kernel unusable for this workload without modification.
  3. FlashInfer is faster than sgl_kernel for batch=512 (the largest tested size where both succeed). At 29.5 vs 49.3 microseconds, FlashInfer achieves a 1.67× speedup. This suggests that FlashInfer's CUTLASS backend has better tile selection heuristics for medium-sized matrices on SM120.
  4. The performance ratio is not monotonic. sgl_kernel is faster at batch=256 but slower at batch=512. This non-linearity means that simple rules of thumb ("library X is faster than library Y") are unreliable—the answer depends on the specific matrix dimensions.
  5. A benchmark methodology for comparing FP4 GEMM libraries. The script itself serves as a template for future comparisons, including proper handling of dimension conventions, error handling for failed kernel launches, and warmup iterations for stable timing.

The Thinking Process Visible in the Reasoning

The assistant's reasoning in message 875 is compressed but visible. The opening sentence—"Different dimension convention. sgl_kernel expects b as (N, K//2) transposed."—is the culmination of a diagnostic chain that is not fully spelled out but can be inferred:

  1. Observation: cutlass_scaled_fp4_mm crashes when given the same tensor shapes that work for FlashInfer.
  2. Hypothesis: The crash is not due to a numerical issue or a GPU hardware problem (since FlashInfer works fine), but rather an API mismatch.
  3. Deduction: The most common API mismatch in GEMM libraries is the layout of the weight matrix. Some libraries expect the weight to be pre-transposed for memory access efficiency.
  4. Verification: The assistant checks the function signature (from the help output in message 870) and infers the expected shapes. The fix is applied and tested. The decision to wrap the sgl_kernel call in a try/except block is also a reasoned choice. Rather than letting the entire benchmark fail if one matrix size crashes, the assistant isolates the failure and continues. This is a practical engineering trade-off: the goal is to collect as much data as possible, and partial data is better than no data. The choice to transpose with .t().contiguous() rather than .permute(...) or .transpose(...) shows an understanding of PyTorch tensor mechanics. .t() is a shortcut for dimensions 0 and 1, and .contiguous() is necessary because the transposed view has non-contiguous strides that many CUDA kernels cannot handle.

Conclusion

Message 875 is a small but revealing moment in a larger optimization journey. It captures the moment when a subtle API incompatibility is diagnosed and resolved, transforming a crash into a meaningful performance comparison. The message demonstrates several qualities that define effective ML engineering: the ability to read error messages and infer root causes, the willingness to adapt experimental setups to match library conventions, and the pragmatism to handle failures gracefully without losing data.

The broader lesson is that performance optimization at this level is never straightforward. Two libraries implementing the same mathematical operation can have different tensor layouts, different tile size heuristics, and different failure modes. The engineer's job is not just to benchmark but to understand why the numbers come out the way they do—and to ensure that the benchmark is actually measuring what it claims to measure.

For the GLM-5 project, the conclusion from this message is clear: FlashInfer's FP4 GEMM is the more robust choice for MoE inference on Blackwell GPUs, because it handles the small batch sizes that dominate the decode phase, and it performs competitively or better across the range of relevant matrix dimensions. sgl_kernel's cutlass_scaled_fp4_mm may have its uses for larger batch sizes, but its failure at M ≤ 64 makes it unsuitable for this particular workload without significant modification or padding.

This message, then, is not just a bug fix—it is a decision point. It steers the optimization effort toward FlashInfer and away from sgl_kernel, saving future investigation time that would otherwise be spent debugging crashes or chasing marginal gains from a library that cannot handle the core workload. In the high-stakes world of LLM inference optimization, such decisions compound: choosing the right kernel library early can save days of engineering effort downstream.