Mamba3 MIMO backward cannot launch on a lower-shared-memory Blackwell SKU (220 KB requested vs 99 KB limit)

Summary

mamba3_mimo_bwd.py cannot launch the MIMO backward path on this Blackwell device: the generated kernels require more dynamic shared memory than the GPU permits.

• device opt-in shared-memory limit: 101,376 B (~99 KB)
• minimum generated requirement observed: 152,768 B (~149 KB)
• other tested shapes: 220,352 B (~215 KB)

The backward kernel fails before gradients can be produced, so Mamba-3 MIMO fine-tuning is blocked entirely on this SKU.

This appears distinct from the existing Blackwell reports:

• #938 — Mamba3 SISO exploding gradients on B200
• #904 — Mamba3 SISO backward performance / compiler pathology on GB200
• #997 — SISO forward corruption fixed by restricting Blackwell to num_stages=1

Here the MIMO backward kernel is not launchable at the tested geometry — not merely slow, and not numerically unstable. There is no gradient to be wrong.

Environment

Literal error

tvm.error.InternalError: Failed to set the allowed dynamic shared memory size to 220352

raised from:

mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py:184 in backward
  → mamba_mimo_bwd_combined(...)
  → mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_bwd.py:1513 in mamba_mimo_bwd_combined
  → bwd_bwd_kernel(...)
  → tilelang/jit/kernel.py:202 __call__
  → tvm runtime/cuda/cuda_module.cc:271 CUDAWrappedFunc::operator()

Repro geometry

Fused MIMO invocation (Mamba3 with is_mimo=True, mimo_rank=4, chunk_size=16):

batch      = 1
sequence   = 17
d_state    = 64
headdim    = 32
mimo_rank  = 4
chunk_size = 16
dtype      = fp32

Boundary tensors at mamba3_mimo_combined:

Q        [1,17,4,1,64]     K        [1,17,4,1,64]
V        [1,17,16,32]      Z        [1,17,16,32]
ADT      [1,16,17]         DT       [1,16,17]      Trap  [1,16,17]
Q_bias   [16,4,64]         K_bias   [16,4,64]      D     [16]
MIMO_V   [16,4,32]         MIMO_Z   [16,4,32]      MIMO_Out [16,4,32]
Angles   [1,17,16,16]

Forward runs and produces [1,17,16,32]. torch.autograd.grad through the same call raises the shared-memory error above.

Sweep: reducing outer width does not help

The requirement is dominated by the internal recurrence geometry (d_state, rank, chunk/tile shapes) rather than outer model width, so shrinking d_model does not bring it under the limit. The three smallest configurations would require headdim=16, which the kernel itself refuses as untested.

Expected behavior

At least one MIMO backward configuration should remain launchable within the device's shared-memory budget — analogous to the conservative num_stages=1 fallback added for SISO forward in #997. Potential directions: smaller shared-memory tile geometry, reduced buffer duplication / more buffer reuse, splitting the fused backward kernel, architecture-aware config pruning, or a lower-resource fallback path.

Isolation

A frozen scan corpus was captured directly at the mamba3_mimo_combined fused boundary, containing the exact inputs and output of one invocation with SHA-256 hashes per tensor. The failure is therefore isolated to the fused MIMO path rather than in_proj, Q/K normalization, or upstream dt/A/trap/angle construction — all of which are ordinary PyTorch and execute correctly.

I can attach the corpus (15 tensors, ~200 KB total) if useful as a regression fixture.

Impact

• no gradient audit can be completed on this device;
• no fine-tuning smoke test can be interpreted;
• MIMO training is blocked on this GPU class/configuration.

The forward path remains capturable and validatable; only backward is affected.