tile-ai / tile-ai/tilelang

[Question] TileLang on sm_120 (RTX 5090): shared memory overflow, cache warning flood, and autoregressive performance regression with Mamba-3 Mimo

Open
#2,201 2 comments 0 reactions 0 assignees View on GitHub
question
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.