<- Open Source

A fix to FlashInfer's Mamba selective-state-update CUDA kernels, so head dims that aren't multiples of 64 compile and give correct results.

FlashInfer is a library of GPU kernels for serving LLMs, used by inference engines such as vLLM and SGLang. Its selective state update is the decode step of a Mamba layer: each new token updates the layer’s recurrent state, head by head. The kernels assumed head dims that are multiples of 64, so models whose heads are 80 or 96 wide failed. My pull request makes those head dims work, and makes every other head dim either work or fail with a clear error, never silently. It was merged on Oct 27, 2026.

What was wrong

  • The module didn’t compile. The dispatcher split a head across 4, 2 or 1 thread blocks (CTAs) without checking that the split was even. At head dim 80, 4 and 2 CTAs give 20 and 40 rows each, neither a multiple of the 16 rows one pass handles, and nvcc stopped on the kernel’s static asserts.
  • An illegal memory access. The single-token horizontal kernel covers a head exactly only when its dim is a multiple of 64. At 80 or 96 it indexed rows past the head, and on an H100 it stopped with an illegal memory access. auto picks it on Blackwell (SM100) for 16-bit states at large batch.
  • Stale outputs. The multi-token simple kernel’s vectorized epilogue assumed each CTA’s rows were a multiple of 32. At 80, or 96 at small batch, it left the end of each row unwritten: the state was right, but the output held stale values.
  • Silent no-ops and more compile failures. Once the module compiled, a head dim such as 72, which isn’t divisible by 16, found no split that fit, so nothing launched and the output stayed unwritten. Below 64, such as 32 or 48, the single-token module didn’t compile at all.

The fix

  • Accept a CTA split only when it divides the head dim into whole 16-row passes, so 96 uses 2 CTAs per head and 80 uses 1. Head dims that already worked keep the same split.
  • Build and pick the horizontal kernels only for head dims they cover. auto falls back to the vertical or simple kernel, and an explicit request that can’t work raises a clear error.
  • Use the vectorized epilogue only when each CTA’s rows are a multiple of 32. Otherwise each lane writes every 32nd element, so every row is written.
  • Raise a clear error, instead of launching nothing, for head dims no kernel supports.

How it was checked

Head dims 80 and 96 joined the single-token and multi-token test tables, with 72 for the fallback. On an H100 the two test files give 146 passed and 0 failed (116 skipped, most of them kernels that run only on SM100). Without a Blackwell GPU, the SM100 kernels were checked for compilation: they now build at head dims 48, 64, 80 and 96, where 48 and 80 used to fail. Head dims that are multiples of 64 keep the same split, kernel choice and epilogue, so their behavior and speed don’t change.