Resolve the Blackwell TMEM over-allocation issue in the official Mamba-3 backward

Resolve the open state-spaces/mamba#904 issue by preventing the official Mamba-3 backward's Blackwell compiler path from requesting 544 Tensor Memory columns when the hardware permits only 512, thereby eliminating its compilation failure and forced slow fallback.

Background

NVIDIA Blackwell hardware imposes a 512-column Tensor Memory (TMEM) budget per warpgroup. The paper reports that the official Mamba-3 backward, when using the tensor-core matrix multiplication operation tl.dot at the normal performance setting, triggers a compiler request for 544 TMEM columns. This causes compilation failure and forces execution through a fallback reported as substantially slower on GB200/B200 systems.

The authors reproduce the failure but do not fix it upstream. They instead work around the constraint through two approaches: six Triton kernels that avoid tl.dot, and a native tcgen05 backward that uses allocation-once and relinquishment-once TMEM lifecycle management with partitioned accumulator offsets. The upstream issue therefore remains an unresolved implementation problem for the official Mamba-3 backward.

References

This constraint is the crux of the open issue state-spaces/mamba#904: the official Mamba-3 backward uses the tensor-core matmul (tl.dot), and at the normal performance setting a compiler pass requests 544 TMEM columns against the hardware's 512, so the kernel fails to compile and falls back to a path reported as $38.7\times$ slower on GB200/B200. The tensor-memory half of the failure is addressed by a merged Triton change (PR #9093) not yet shipped in a Mamba release, so a pinned pre-#9093 stack still reproduces it.

— A Contract-Grade Verifier for LLM-Generated GPU Kernels, and a Native Blackwell Backward for the Gated-Linear-Recurrence Family  (2608.12700 - Shah et al., 13 Aug 2026) in Section 2, paragraph “Blackwell tcgen05, tensor memory, and #904”; see also Section 7.3, “A faithful reproduction of #904”