[GH-ISSUE #261] triton rms_rope kernel reads its input twice and launches 95k single-row programs #70

New issue

Open

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

bsakaguchi commented

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

Owner

Originally created by @okolenmi on GitHub (Oct 19, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/261

AI generated, but it was actually a helpful optimization.

Where

comfy_kitchen/backends/triton/rms_rope.py, rms_rope_kernel (comfy-kitchen
0.2.37). This is the implementation used when no vendor backend is available,
e.g. on Intel XPU.

Two things visible in the source

The input row is loaded twice. Line 55 loads the whole row (x_full,
head_dim elements) to compute the sum of squares; lines 66-67 then load
the same row again (x_0 / x_1) for the rotation. A 128-element row is
read as 256 elements.
One program per head row. grid = (batch, num_heads, seq_len) (line
130) with block_size = next_power_of_2(head_dim // 2) = 64 for
head_dim=128. At the shapes we use (1, 3973, 24, 128) that is 95,352
programs of 64 lanes per launch, and rms_rope_split_half launches it
twice (once for q, once for k).

Measured

1 x 3973 x 24 x 128 bf16, Intel Arc B580, torch 2.12.1+xpu:


time per `rms_rope_split_half` call	0.81 ms
traffic per call (2 launches x 95,352 rows x (256 read + 128 write) x 2 B)	~146 MB
implied bandwidth	~180 GB/s
bandwidth the same card reaches on plain fills/copies	380-570 GB/s
calls per forward in our 42-block model	44 -> 35.7 ms/forward -> 71 ms/step

Removing the second load cuts the traffic per call from ~146 MB to ~98 MB
(-33%). The 180 GB/s is also roughly half of what a plain copy achieves, so
there is more in the gap than the double read alone. I have not profiled the
kernel's own limiter (occupancy vs bandwidth), so I cannot attribute the rest
between the grid shape and the loads -- both are visible in the source.

Suggested fix

A single-pass kernel: load each row once into registers, reduce over
head_dim, normalise, rotate, store. Since the q and k calls are independent
and use an identical grid, consider one launch for both as well.

Originally created by @okolenmi on GitHub (Oct 19, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/261 AI generated, but it was actually a helpful optimization. ### Where `comfy_kitchen/backends/triton/rms_rope.py`, `rms_rope_kernel` (comfy-kitchen 0.2.37). This is the implementation used when no vendor backend is available, e.g. on Intel XPU. ### Two things visible in the source 1. **The input row is loaded twice.** Line 55 loads the whole row (`x_full`, `head_dim` elements) to compute the sum of squares; lines 66-67 then load the same row again (`x_0` / `x_1`) for the rotation. A 128-element row is read as 256 elements. 2. **One program per head row.** `grid = (batch, num_heads, seq_len)` (line 130) with `block_size = next_power_of_2(head_dim // 2)` = 64 for `head_dim=128`. At the shapes we use (`1, 3973, 24, 128`) that is **95,352 programs of 64 lanes per launch**, and `rms_rope_split_half` launches it twice (once for q, once for k). ### Measured `1 x 3973 x 24 x 128` bf16, Intel Arc B580, torch 2.12.1+xpu: | | | |---|---| | time per `rms_rope_split_half` call | 0.81 ms | | traffic per call (2 launches x 95,352 rows x (256 read + 128 write) x 2 B) | ~146 MB | | implied bandwidth | **~180 GB/s** | | bandwidth the same card reaches on plain fills/copies | 380-570 GB/s | | calls per forward in our 42-block model | 44 -> 35.7 ms/forward -> **71 ms/step** | Removing the second load cuts the traffic per call from ~146 MB to ~98 MB (-33%). The 180 GB/s is also roughly half of what a plain copy achieves, so there is more in the gap than the double read alone. I have not profiled the kernel's own limiter (occupancy vs bandwidth), so I cannot attribute the rest between the grid shape and the loads -- both are visible in the source. ### Suggested fix A single-pass kernel: load each row once into registers, reduce over `head_dim`, normalise, rotate, store. Since the q and k calls are independent and use an identical grid, consider one launch for both as well.

bsakaguchi referenced this issue

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

[PR #70] [MERGED] Add RMS fusion for RoPE #122