[Question] TileLang on sm_120 (RTX 5090): shared memory overflow, cache warning flood, and autoregressive performance regression with Mamba-3 Mimo
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 742
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 104
Description
### Required prerequisites
- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) that this hasn't already been reported. (comment there if it has.)
### Questions
## TileLang Issues on sm_120 (NVIDIA RTX 5090 Blackwell)
### Environment
- GPU: NVIDIA GeForce RTX 5090 (sm_120, Blackwell)
- CUDA: 13.0
- PyTorch: 2.9.0+cu130
- Python: 3.13
- tilelang: 0.1.8 → 0.1.9
- mamba-ssm: 2.3.2.post1
---
### Issue 1: Dynamic shared memory overflow on sm_120 (tilelang 0.1.8)
**Description**: Mamba-3 Mimo kernel from mamba-ssm fails to launch on RTX 5090 due to dynamic shared memory allocation exceeding hardware limits.
**Error**:
```
tvm.error.InternalError: Failed to set the allowed dynamic shared memory size to 168128
```
**Context**: RTX 5090 max shared memory per block (opt-in): 101,376 bytes (99 KB). The Mamba-3 Mimo kernel requests 168,128 bytes (164 KB), which is ~1.66× the hardware limit.
**Reproduction**:
```python
from mamba_ssm import Mamba3
import torch
m = Mamba3(d_model=896, d_state=128, headdim=64, is_mimo=True,
mimo_rank=4, chunk_size=16, dtype=torch.bfloat16).cuda()
x = torch.randn(1, 128, 896, device='cuda', dtype=torch.bfloat16)
y = m(x) # Fails on tilelang 0.1.8, OK on 0.1.9
```
**Resolution**: Upgrading to tilelang 0.1.9 resolved the shared memory overflow. The kernel now compiles and launches successfully.
**Status**: ✅ Fixed in 0.1.9
---
### Issue 2: Excessive cache warnings during autoregressive generation
**Description**: After upgrading to tilelang 0.1.9, the Mamba-3 Mimo kernel compiles successfully but produces massive warning output during autoregressive (token-by-token) generation. Each step in the generate loop triggers:
```
[TileLang:tilelang.cache.kernel_cache:WARNING] Found kernel 'mamba_mimo_fwd_kernel'
in memory cache. For better performance, consider using `@tilelang.jit` instead
of direct kernel caching.
```
**Impact**: During S2S inference (autoregressive generation of ~200 tokens through 14 decoder layers), this warning fires ~2800 times (200 × 14), filling the terminal output and significantly slowing down kernel retrieval.
**Workaround**: Set environment variable `TILELANG_DISABLE_JIT_WARNING=1` to suppress warnings.
**Root cause**: The mamba-ssm package uses `T.prim_func` (tilelang internal API) rather than `@tilelang.jit` decorator for kernel compilation. The caching mechanism works correctly but produces the warning on every cache hit.
**Suggested fix**:
1. TileLang side: reduce warning from per-call to per-compilation, or add a rate limit
2. mamba-ssm side: use `@tilelang.jit` decorator instead of `T.prim_func` for the Mimo kernel
3. Kernel: `mamba_mimo_fwd_kernel` in `mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py:70`
**Status**: ⚠️ Workaround available; root fix needed in tilelang or mamba-ssm
---
### Issue 3: Autoregressive generation performance with tilelang kernels
**Description**: Mamba-3 Mimo's autoregressive generation (token-by-token inference) is significantly slower than Mamba-2 (causal-conv1d CUDA kernel). Each step in the generate loop triggers kernel cache lookup overhead.
**Comparison** (same hardware, same model size):
| Model | Self-Attention | Autoregressive speed | Training speed |
|-------|---------------|---------------------|---------------|
| Mamba-2 (causal-conv1d) | CUDA kernel | Fast | Fast |
| Mamba-3 Mimo (tilelang) | tilelang kernel | **Very slow** | Acceptable |
**Root cause**: The tilelang kernel cache lookup per token adds overhead that compounds over hundreds of autoregressive steps. Mamba-2's causal-conv1d CUDA kernel has no such overhead.
**Suggested fix**: Optimize tilelang kernel cache retrieval for repeated calls with identical shapes (common in autoregressive generation).
**Status**: ❌ Not suitable for autoregressive S2S inference; Mamba-2 used as fallback
---
### Summary
| Issue | Severity | tilelang Version | Status |
|-------|----------|-----------------|--------|
| Shared memory overflow | Critical | 0.1.8 | ✅ Fixed in 0.1.9 |
| Cache warning flood | Medium | 0.1.9 | ⚠️ Workaround exists |
| Autoregressive speed | High | 0.1.9 | ❌ Not usable |
**Final configuration**: Mamba-2-Hybrid (causal-conv1d CUDA) as default. Mamba-3 Mimo reserved for future tilelang optimization.
Contributor guide
Research direction
Start by reproducing the Mamba-3 Mimo generation path on the listed RTX 5090 environment and inspect the cache-warning behavior around tilelang.cache.kernel_cache. The payload points to mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py:70 and the T.prim_func-based kernel. Done means eliminating repeated cache warnings and reducing lookup overhead for repeated autoregressive calls; the shared-memory failure is already fixed in 0.1.9.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100