[GH-ISSUE #169] [Bug] Triton rms_rope kernel exceeds CUDA gridDim.z limit (65535) when seq_len >= 65536 #52

New issue

Open

opened 2026-10-26 22:19:44 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:19:44 +00:00

Owner

Originally created by @trojanGoat on GitHub (Sep 10, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/169

Summary

In comfy_kitchen.backends.triton.rms_rope, the kernel launch grid is defined as:

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:

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:

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):

--- 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.

Originally created by @trojanGoat on GitHub (Sep 10, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/169 ### 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.

bsakaguchi referenced this issue

2026-10-26 22:20:27 +00:00

[PR #52] Fix crash when running quantized models through torch.compile #107