Comfy-Org / Comfy-Org/comfy-kitchen
[Bug] Triton rms_rope kernel exceeds CUDA gridDim.z limit (65535) when seq_len >= 65536
- Dominant language
- Python
- Stars
- 220
- Forks
- 91
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 12
Description
### Summary
In `comfy_kitchen.backends.triton.rms_rope`, the kernel launch grid is defined as:
```python
grid = (batch, num_heads, seq_len)
```
and maps `seq_idx` to `tl.program_id(2)`.
In NVIDIA CUDA hardware specifications, the maximum grid dimensions are:
- `gridDim.x <= 2,147,483,647` (2^31 - 1)
- `gridDim.y <= 65,535`
- `gridDim.z <= 65,535`
Because `seq_len` is mapped to `grid[2]` (`gridDim.z`), **any sequence length >= 65,536 causes the CUDA driver to return `cudaErrorInvalidValue` (`RuntimeError: Triton Error [CUDA]: invalid argument`)**.
### Minimal Reproduction
This can be reproduced directly in Python:
```python
import torch
from comfy_kitchen.backends.triton import rms_rope_
device = "cuda"
head_dim = 64
num_heads = 8
# nt = 65535 succeeds; nt = 65536 crashes
for nt in [65535, 65536]:
q = torch.randn(1, nt, num_heads, head_dim, device=device, dtype=torch.bfloat16)
k = torch.randn(1, nt, num_heads, head_dim, device=device, dtype=torch.bfloat16)
q_scale = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
freqs = torch.randn(1, nt, 1, head_dim // 2, 2, 2, device=device, dtype=torch.float32)
try:
rms_rope_(q, k, freqs, q_scale)
print(f"nt={nt}: SUCCESS")
except Exception as e:
print(f"nt={nt}: FAILED with {e}")
```
**Output:**
```text
nt=65535: SUCCESS
nt=65536: FAILED with RuntimeError: Triton Error [CUDA]: invalid argument
```
### Real-world ComfyUI Failure
When running LTX-Video 2.5 decoding with `VAEDecodeTiled` (e.g., `tile_size: 512`, `temporal_size: 64`):
- In Stage 3 of `CausalDiffusionVAE` (`na_diffusion_decoder.py`), tile dimensions reach $H=64, W=64, \text{dim}=512$.
- The chunk size formula calculates: $\text{chunk} = \lfloor 2^{25} / (64 \times 64 \times 512) \rfloor = 16$.
- Number of tokens: $nt = 16 \times 64 \times 64 = 65,536$.
- The decode crashes immediately at `nodes.py:375` / `na_diffusion_decoder.py:157`.
### Proposed Fix
Map `seq_idx` to `program_id(0)` (which supports up to 2.14 billion on `gridDim.x`) and `batch_idx` to `program_id(2)` (batch sizes are typically <= 8):
```diff
--- a/comfy_kitchen/backends/triton/rms_rope.py
+++ b/comfy_kitchen/backends/triton/rms_rope.py
@@ -37,3 +37,3 @@
- batch_idx = tl.program_id(0)
+ seq_idx = tl.program_id(0)
head_idx = tl.program_id(1)
- seq_idx = tl.program_id(2)
+ batch_idx = tl.program_id(2)
@@ -130,1 +130,1 @@
- grid = (batch, num_heads, seq_len)
+ grid = (seq_len, num_heads, batch)
```
We verified this fix:
- Tested with $nt \ge 100,000$ without error.
- Verified numerical equivalence against the eager PyTorch implementation.
- Confirmed full end-to-end execution of LTX-Video 2.5 tiled VAE decode in ComfyUI.
Contributor guide
Research direction
Start with comfy_kitchen/backends/triton/rms_rope.py and inspect the kernel launch and program ID mapping described in the issue. Run the supplied Python reproduction with sequence lengths 65,535 and 65,536, then verify numerical equivalence with the eager PyTorch implementation and the LTX-Video 2.5 tiled VAE decode. Done means long sequences no longer produce the CUDA invalid-argument error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 84/100