Mamba3 backward pass (mamba3_siso_bwd_kernel_dqkv) 38.7x slower on GB200 (SM100) due to ptxas C7907 eliminating autotuner configs

[tolleybot]

Summary

The mamba3_siso_bwd_kernel_dqkv Triton kernel exhibits a 38.7x performance regression on GB200 (SM100a / Blackwell) compared to H200 (SM90 / Hopper). Training achieves only 5% MFU on GB200 versus 33% MFU on H200 for an identical 1.4B parameter Mamba3 SISO model. The backward kernel alone accounts for 90.76% of total GPU time on GB200.

The root cause is an NVIDIA ptxas internal compiler error (C7907) that eliminates all num_warps=4 and num_warps=8 autotuner configurations for the backward kernel on SM100a. The 3 surviving num_warps=2 configs exhibit extreme register spilling (42 to 50 KB per thread), directly causing the regression.

Mamba2 (SSD) is unaffected and achieves consistent ~32% MFU across all GPUs.

Training throughput

Measured with d_model=2048, n_layers=48, batch=4, bf16, seq_len=4096:

GPU	Mamba2 TFLOP/s	MFU	Mamba3 TFLOP/s	MFU
A100	102.5	33%	37.5	12%
H200	239.2	24%	323.6	33%
GB200	360.4	32%	56.0	5%

Kernel level comparison

Self CUDA time over 10 timed iterations:

Kernel	H200 (ms)	GB200 (ms)	Ratio
mamba3_siso_bwd_kernel_dqkv	587	22,746	38.7x slower
mamba3_siso_fwd_kernel	602	450	0.75x (faster)
mamba3_siso_bwd_kernel_rotary_bias_angles	368	375	~1.0x
aten::mm (cuBLAS)	1,324	543	0.41x (faster)

Root Cause

When Triton compiles mamba3_siso_bwd_kernel_dqkv for SM100a, ptxas triggers C7907 on all num_warps=4 and num_warps=8 configs (5 produce trap stubs, 1 crashes outright). Only the 3 num_warps=2 configs survive, and those are capped at 32 registers with 42 to 50 KB of spill traffic per thread. With 64 threads per block, that amounts to roughly 5.8 MB of spill traffic per thread block, which dwarfs the kernel's useful data movement.

On H200, all 9 autotuner configs compile and launch, giving access to higher warp configurations with substantially lower register pressure.

Additionally, SM100's tensor memory budget independently prevents num_warps>=4 from launching on the forward kernel (required tensor memory: 544, hardware limit: 512), suggesting the Mamba3 kernels may need tile size tuning or register pressure reduction for Blackwell regardless of the C7907 fix.

Error output

Internal Triton PTX codegen error
`ptxas` stderr:
ptxas fatal   : (C7907) Internal compiler error.
ptxas fatal   : Ptx assembly aborted due to errors

Companion issue filed upstream: triton-lang/triton#9933

Environment

Component	Version
GPU	NVIDIA GB200 (Blackwell, SM100a, compute capability 10.0)
GPU (baseline)	NVIDIA H200 (Hopper, SM90, compute capability 9.0)
Driver	580.65.06
CUDA toolkit	13.1 (V13.1.115, built 2025-12-16)
Triton	3.6.0
PyTorch	2.11.0a0+eb65b36914.nv26.02
mamba_ssm	2.3.1

Workaround

None known at this time. The forward pass and other kernels are unaffected.

[tridao]

@aakashlahoti @caitWW maybe there's some tweaks we can make to the triton kernel to work around compiler issue?

[caitWW]

will take a look

[masahi]

thanks for the report, I'll take a look at this from Triton side

[caitWW]

@tolleybot which input/output dimensions are you benchmarking with to get these numbers?

[tolleybot]

Updated the issue description with the corrected root cause. The actual error is a ptxas internal compiler error (C7907) when assembling for sm_100a, not a silent codegen fallback as originally described.

@caitWW Let me see what I can find on the specific input and output dimensions.

[tolleybot]

Updated the issue with detailed root cause analysis, training throughput numbers, and environment information. The core finding is that ptxas C7907 eliminates all higher warp autotuner configs for the backward kernel on SM100a, forcing extreme register spilling in the surviving num_warps=2 configs.

@caitWW The benchmarks were run with a 1.4B parameter model: d_model=2048, n_layers=48, batch=4, bf16, seq_len=4096.

[caitWW]

@tolleybot I just created and merged a PR that adds maxnreg autotuning to bwd SISO triton kernels. I found that this allowed the compiler to allocate more regs per kernel, thus solving the reg spilling. The cap at num_warps = 2 is likely due to running out of tmem on Blackwell; we will look into ways to solve this.

[apstenku123]

Hit the same tmem/smem class of issue on the TileLang MIMO backward path — specifically mamba_mimo_bwd_bwd in mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_bwd.py — on an NVIDIA GB10 Grace Blackwell (sm_121 consumer/workstation Blackwell, 100 KB dynamic shared memory opt-in per block). Reporting here because PR #905 only touched the SISO Triton kernels and @caitWW mentioned "we will look into ways to solve this" for the sibling tmem case. Also cross-referencing #873 (@MasahiroOgawa) who hit what looks like the exact same class of blocker on RTX 4080 (sm_89) three weeks ago.

Repro

• Device: NVIDIA GB10 (cc 12.1 / sm_121). shared_memory_per_block_optin = 101376 bytes, shared_memory_per_multiprocessor = 102400 bytes (100 KB).
• Stack: torch 2.12.0.dev20260410+cu132, mamba_ssm 2.3.1, tilelang 0.1.8 (with apache-tvm-ffi==0.1.9 per AttributeError: '_NestedLoopCheckVisitor' object has no attribute '_inst' #907), Python 3.13.
• Dims: B=2 S=256 H=8 G=1 N=64 P=64 R=4 chunk_size=16 (MIMO, no Z, has D, not reduceO).
• Config: {threads: 128, num_stages: 0} (bwd_bwd autotune's smallest config).

Standalone reproducer — just calling the kernel factory, no training loop:

from mamba_ssm.ops.tilelang.mamba3 import mamba3_mimo_bwd
kernel = mamba3_mimo_bwd.mamba_mimo_bwd_bwd(
    2, 256, 8, 1, 64, 64, 4,   # B S H G N P R
    False, True, False,          # hasZ, hasD, reduceO
    16, 4, 'bfloat16',           # chunk_size, rotary_dim_divisor, dtype
)

Fails at the benchmark phase of autotune (compile succeeds, launch fails) with:

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

Both of the two bwd_bwd autotune configs (threads ∈ {128, 256}, num_stages=0) request ~140-144 KB of dynamic shared memory, exceeding the GB10 opt-in ceiling. mamba_mimo_fwd (12 configs) and mamba_mimo_bwd_fwd (12 configs) both autotune fine on the same box — 19 % speedup vs the hardcoded defaults on fwd. Only the true bwd_bwd exceeds the budget. Same blocker @MasahiroOgawa described in #873 ("Mamba3 MIMO backward fails on RTX 4080 — requires >100KB shared memory per SM").

Single-line workaround that works, zero code changes to the algorithm

TileLang ships a pass called MergeSharedMemoryAllocations (in src/transform/merge_shared_memory_allocations.cc, refactored for aggressive liveness analysis in TileLang v0.1.5 / PR #505, alignment-optimised in #693). The "aggressive" variant is off by default but can be toggled via pass_configs. Adding two flags to the existing @tilelang.jit(pass_configs={...}) decorator on both mamba_mimo_bwd_fwd and mamba_mimo_bwd_bwd drops the peak smem footprint significantly:

@tilelang.jit(
    out_idx=[],
    pass_configs={
        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
        tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
        # NEW — fixes sm_121 / sm_120 / sm_89 consumer GPUs:
        tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True,
    })
def mamba_mimo_bwd_bwd(...): ...

Measured on GB10 at the dims above (smoke-test config from mamba_mimo_bwd_bwd):

metric	before	after
peak dynamic smem (bytes)	140960	93216
peak dynamic smem (KB)	137.7 KB	91.0 KB
reduction	—	47.7 KB (33.9 %)
GB10 compile	❌ Failed to set the allowed dynamic shared memory size to 140960	✅ ok
GB10 autotune	❌ Auto-tuning failed: No configuration successfully compiled and passed benchmarking/validation	✅ 2/2 configs benchmarked
winner latency	—	{threads:128, num_stages:0} → 0.211 ms
runner-up	—	{threads:256, num_stages:0} → 0.225 ms

The merge plan (logged via TL_DEBUG_MERGE_SHARED_MEMORY_ALLOCATIONS: True) shows the pass found two large disjoint-liveness groups:

offset 68640:  lkq_masked_dtype_buf / dPsiV_combined_shared / PsiV_shared / dk_shared   (4-way alias)
offset 35872:  k_reshaped_shared / states_shared / dq_shared                            (3-way alias)
offset 52256:  k_pre_rot_shared / dangle_dk__or__dq_shared                              (2-way alias, mixed dtype)
offset 0:      dstates_shared
offset 8192:   trap_scale_shared
offset 8224:   dPhiO_shared
offset 19488:  q_shared
offset 27680:  q_pre_rot_shared
offset 44064:  k_pre_rot_shared
offset 52256:  k_shared
offset 60448:  lkq_masked__or__dkq_masked_shared
offset 76832:  dqk_from_diag_shared  (fp32, 16 KB — the single biggest buffer)
total merged: 93216 bytes

The aliased buffers are genuinely disjoint in time across the reverse-chunk loop (e.g., dq_shared is only written in the tail of each iteration, after states_shared / k_reshaped_shared are dead). The merge pass found 47 KB of these automatically from the existing source — no algorithmic change, no loop re-ordering, no phase split needed.

Numerical correctness verified — see follow-up comment below. TL;DR: ran identical inputs (seed=42) through unmerged kernel on H200 and merged kernel on GB10, all 12 output gradients match within bf16 noise (worst rel error 6.1e-6 on DQ, three orders of magnitude below bf16 round-off).

Proposed PR

Happy to send a one-line PR adding TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True to both @tilelang.jit(pass_configs=...) decorators in mamba3_mimo_bwd.py (and optionally mamba3_mimo_fwd.py for symmetry). If you'd prefer it gated behind a cppmega/sm_12x-style detection so H200 keeps the unmerged path (in case of a perf regression), I can do that too, but it'd be strictly additive either way since the flag is inert when the kernel already fits.

Also relevant: #873 (sm_89 RTX 4080, same "MIMO bwd > 100 KB" symptom) and #907 (the apache-tvm-ffi<0.1.10 pin is still needed to even get TileLang 0.1.8 importable). Fixing this blocker brings both sm_89 consumer and sm_12x consumer GPUs back onto the MIMO training path.

Thanks for the fast turnaround on #905 — the maxnreg fix for SISO was a great move. The MIMO TileLang side can be fixed even more easily (just a decorator flag flip).

[apstenku123]

Update — numerical correctness verified (H200 vs GB10 parity).

Generated identical inputs (torch.manual_seed(42), smoke shapes B=2 S=256 H=8 G=1 N=64 P=64 R=4 chunk_size=16, hasD=True, hasZ=False, reduceO=False). Ran the unmerged kernel on H200 (sm_90, fits the 228 KB budget) and the merged kernel on GB10 (sm_121, the patch I described above), then compared all 12 output tensors element-wise. Same mamba_ssm 2.3.1, same tilelang 0.1.8, same apache-tvm-ffi 0.1.9, both bf16.

TENSOR	max abs	mean abs	rel error	H200 norm	status
DK	1.5e-05	2.1e-10	3.4e-05	1.20e+00	PASS
DV	3.1e-05	1.7e-09	5.8e-05	1.52e+00	PASS
DMIMO_V	3.5e-05	8.8e-08	4.1e-05	1.49e+00	PASS
DQ	4.9e-04	6.0e-10	6.1e-06	8.21e+01	PASS
DD	8.9e-08	2.5e-08	1.2e-07	1.18e+00	PASS
DDA	5.5e-05	1.1e-06	3.9e-06	1.02e+02	PASS
DSSDA	3.0e-07	1.2e-11	1.1e-05	4.16e-02	PASS
DDA_CS_REV	8.0e-07	2.5e-09	1.3e-05	1.17e-01	PASS
DDA_CS	2.3e-05	5.6e-09	2.8e-06	8.19e+00	PASS
DFACTOR	3.1e-05	5.5e-08	2.7e-05	1.95e+00	PASS
DGAMMA_DIAG	9.3e-10	1.5e-10	7.6e-08	1.99e-01	PASS
DANGLES	3.7e-06	1.0e-09	8.5e-07	5.79e+00	PASS

Worst case anywhere is DQ rel error 6.1e-6 — about three orders of magnitude below bf16 round-off (~1e-3). All other tensors are 1-3 orders of magnitude tighter than that. Differences are consistent with normal H200-vs-GB10 bf16 rounding noise (different MMA accumulation order, different L1 sizes, different SM count) and are not correlated with the buffers the merge pass aliased — i.e., the 4-way alias group at offset 68640 (lkq_masked_dtype_buf / dPsiV_combined_shared / PsiV_shared / dk_shared) and the 3-way alias group at offset 35872 (k_reshaped_shared / states_shared / dq_shared) all produce gradients within tolerance.

So the answer to "is the aggressive merge safe?" is yes for the smoke-test config. TileLang's LinearAccessPatternFinder correctly inferred that the buffers in those groups have disjoint live ranges across the reverse-chunk loop. This crosses off the only correctness blocker I had above — happy to send the PR whenever you're ready, just let me know whether you want the flag unconditional in mamba3_mimo_bwd.py/mamba3_mimo_fwd.py or gated behind a runtime device-property check (e.g., enable only when cudaDeviceGetAttribute(cudaDevAttrMaxSharedMemoryPerBlockOptin) < 200_000). Also happy to run the same parity sweep on a sm_120 (RTX 5090) box if anyone has access — I expect the same result.

Repro scripts I used (small enough to inline): parity_make_inputs.py (creates /tmp/parity_inputs.pt with seed=42, all tensors), parity_run_kernel.py {h200|gb10} (loads inputs, calls mamba_mimo_bwd_bwd, dumps outputs to /tmp/parity_outputs_*.pt), then a 30-line numpy diff. Can paste them in if useful.

[caitWW]

cc @berlinchen7 will train w updated kernels to verify loss accuracy. thanks for the suggestions and tests!

[masahi]

@caitWW I see that the kernel uses make_tensor_descriptor. This is TMA with descriptors created in the kernel (we refer to it as "device TMA"). It is more efficient to create descriptors on host via the TensorDescriptor constructor ("host TMA"). Using device TMA when there is no dynamism just makes the kernel slower and creates issues / confusions in upstream users - see triton-lang/triton#10002 and the issues linked there.

[masahi]

Found why it's so slow on Blackwell. Turns out auto-tuner picks a config that's incompatible with Blackwell TMEM / tcgen05 MMA (num_warps = 2). So no async MMA is happening, software pipelining is disabled, and all TMA are lowered as a blocking operation. While on Hopper WGMMA instructions are correctly generated and the loop is pipelined as expected.

The auto tuner does try num_warps = 4 on Blackwell, but it fails to compile with triton.runtime.errors.OutOfResources: out of resource: tensor memory, Required: 576, Hardware limit: 512. Reducing block sizes or num_stages may help.. It means the kernel is using too much TMEM. Dumping the generated IR shows that indeed Triton creates many unique TMEM allocations in the kernel. Next I'm going to investigate if this alloc behavior is Triton's fault or inherent to the current kernel design.