Comfy-Org / Comfy-Org/comfy-kitchen

[Bug] Triton rms_rope kernel exceeds CUDA gridDim.z limit (65535) when seq_len >= 65536

Open
#169 0 comments 0 reactions 0 assignees View on GitHub
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.