Skip to content

perf(dsv4 qkv): precompute q-head rope factors #648

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wangqin1723-max wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(dsv4 qkv): precompute q-head rope factors #648

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 36 additions & 17 deletions models/deepseek/v4/qkv_proj_rope.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,28 +261,47 @@ def qkv_proj_rope(
# columns [h0:h0+NOPE_DIM) and rope columns [h0+NOPE_DIM:h0+HEAD_DIM) are disjoint,
# so each task writes a clean head/row block of q.
#
# RoPE stays on the interleaved layout and rotates via a j^1 swap gather + sign mask
# (CANN A3 rotate_interleaved). The rotation index/sign and the interleave-duplicated
# cos/sin are built ENTIRELY IN-KERNEL: swap_idx (j^1), sign ([-1,+1,...]) and dup_idx
# (j>>1) from pl.arange (once per task), and cos_il/sin_il dup-gathered per tg
# (head-independent, reused across both heads). inv_rms is per-row so it factors out
# of the rotation and is folded into the writeback.
# RoPE stays on the interleaved layout and rotates via a j^1 swap gather.
# Precompute the head-invariant swap index, interleaved cos, and sign-folded
# sin once; q_head tasks only load those factors. This mirrors sparse-attn
# rope_cs and avoids scatter.
# out[j] = inv_rms * (x[j]*cos_il[j] + x[j^1]*sign[j]*sin_il[j])
q_rope_cos_il = pl.create_tensor([t_dim, ROPE_DIM], dtype=pl.FP32)
q_rope_sin_signed = pl.create_tensor([t_dim, ROPE_DIM], dtype=pl.FP32)
q_rope_swap_idx = pl.create_tensor([Q_ROPE_T_TILE, ROPE_DIM], dtype=pl.INT32)
Comment on lines +269 to +271

@coderabbitai coderabbitai Bot Jun 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Allocate precompute tensors with the static T_MAX, not the dynamic t_dim.

q_rope_cos_il and q_rope_sin_signed are allocated with t_dim = pl.tensor.dim(x, 0), a runtime-derived dimension off the dynamic T_DYN. Every other GM allocation in this kernel (x_matmul, qr_fp32, qr_i8_matmul, q_proj_fp32/i32, kv_fp32) sizes the row dimension with the compile-time T_MAX, and q_rope_swap_idx here correctly uses the static Q_ROPE_T_TILE. Sizing an allocation with a dynamic dimension breaks the static-allocation contract; keep t_dim for the views/slices only.

🛠️ Proposed fix 
-    q_rope_cos_il = pl.create_tensor([t_dim, ROPE_DIM], dtype=pl.FP32)
-    q_rope_sin_signed = pl.create_tensor([t_dim, ROPE_DIM], dtype=pl.FP32)
+    q_rope_cos_il = pl.create_tensor([T_MAX, ROPE_DIM], dtype=pl.FP32)
+    q_rope_sin_signed = pl.create_tensor([T_MAX, ROPE_DIM], dtype=pl.FP32)
     q_rope_swap_idx = pl.create_tensor([Q_ROPE_T_TILE, ROPE_DIM], dtype=pl.INT32)

Based on learnings: "avoid passing dynamic dimension variables ... to pl.create_tensor() shape arguments. Tensor allocations must use compile-time static dimension values (e.g., use the compile-time batch parameter ...)."

#!/bin/bash
# Confirm the file's create_tensor convention: row dim should be T_MAX, not t_dim.
rg -nP 'pl\.create_tensor\(\s*\[' models/deepseek/v4/qkv_proj_rope.py
# Show T_MAX definition / origin.
rg -nP '\bT_MAX\b' models/deepseek/v4/qkv_proj_rope.py
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@models/deepseek/v4/qkv_proj_rope.py` around lines 269 - 271, `q_rope_cos_il`
and `q_rope_sin_signed` in `qkv_proj_rope.py` are allocated with the
runtime-derived `t_dim`, which violates the static allocation pattern used by
this kernel. Update the `create_tensor` shapes in the `qkv_proj_rope` path to
use the compile-time `T_MAX` for the row dimension, matching the other GM
allocations like `x_matmul`, `qr_fp32`, `qr_i8_matmul`, and `q_proj_fp32/i32`,
while keeping `t_dim` only for later views or slices; `q_rope_swap_idx` already
follows the correct static sizing.

Source: Learnings

for cp in pl.spmd(1, name_hint="q_rope_cs"):
cp_r0 = cp * ROPE_HALF
cp_c0 = 2 * cp_r0
cs_col = pl.col_expand_mul(
pl.full([Q_ROPE_T_TILE, ROPE_DIM], dtype=pl.FP32, value=1.0),
pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32),
)
cs_dup_f = pl.cast(pl.cast(pl.mul(cs_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
cs_dup_idx = pl.cast(cs_dup_f, target_type=pl.INT32)
cs_lane = pl.sub(cs_col, pl.mul(cs_dup_f, 2.0))
cs_sign = pl.sub(pl.mul(cs_lane, 2.0), 1.0)
q_rope_swap_idx[0:Q_ROPE_T_TILE, 0:ROPE_DIM] = pl.cast(
pl.sub(pl.add(cs_col, 1.0), pl.mul(cs_lane, 2.0)), target_type=pl.INT32
)
for tg_idx in pl.range(t_dim // Q_ROPE_T_TILE):
tg = tg_idx * Q_ROPE_T_TILE
cs_cos = pl.cast(rope_cos_view[tg : tg + Q_ROPE_T_TILE, cp_r0 : cp_r0 + ROPE_HALF], target_type=pl.FP32)
cs_sin = pl.cast(rope_sin_view[tg : tg + Q_ROPE_T_TILE, cp_r0 : cp_r0 + ROPE_HALF], target_type=pl.FP32)
q_rope_cos_il[tg : tg + Q_ROPE_T_TILE, cp_c0 : cp_c0 + ROPE_DIM] = pl.gather(
cs_cos, dim=-1, index=cs_dup_idx
)
q_rope_sin_signed[tg : tg + Q_ROPE_T_TILE, cp_c0 : cp_c0 + ROPE_DIM] = pl.mul(
pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign
)

q_flat = pl.reshape(q, [t_dim, H * HEAD_DIM])
for hg_idx in pl.spmd(H // 2, name_hint="q_head_rms_nope_rope"):
hg = hg_idx * 2
# In-kernel A3 index/sign build (per task, reused across the inner tg/h loop).
q_ones = pl.full([Q_ROPE_T_TILE, ROPE_DIM], dtype=pl.FP32, value=1.0)
q_col = pl.col_expand_mul(q_ones, pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
q_dup_f = pl.cast(pl.cast(pl.mul(q_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
q_dup_idx = pl.cast(q_dup_f, target_type=pl.INT32) # j>>1
q_lane = pl.sub(q_col, pl.mul(q_dup_f, 2.0)) # j%2
q_swap_idx = pl.cast(pl.sub(pl.add(q_col, 1.0), pl.mul(q_lane, 2.0)), target_type=pl.INT32) # j^1
q_sign = pl.sub(pl.mul(q_lane, 2.0), 1.0) # [-1,+1,...]
q_swap_idx = q_rope_swap_idx[0:Q_ROPE_T_TILE, 0:ROPE_DIM]
for tg_idx in pl.range(t_dim // Q_ROPE_T_TILE):
tg = tg_idx * Q_ROPE_T_TILE
q_cos_il = pl.gather(pl.cast(rope_cos_view[tg : tg + Q_ROPE_T_TILE, :], target_type=pl.FP32), dim=-1, index=q_dup_idx)
q_sin_il = pl.gather(pl.cast(rope_sin_view[tg : tg + Q_ROPE_T_TILE, :], target_type=pl.FP32), dim=-1, index=q_dup_idx)
q_cos_il = q_rope_cos_il[tg : tg + Q_ROPE_T_TILE, 0:ROPE_DIM]
q_sin_signed = q_rope_sin_signed[tg : tg + Q_ROPE_T_TILE, 0:ROPE_DIM]
for h_inner in pl.range(2):
h = hg + h_inner
h0 = h * HEAD_DIM
Expand Down Expand Up @@ -310,7 +329,7 @@ def qkv_proj_rope(
# RoPE writeback on columns [h0+NOPE_DIM:h0+HEAD_DIM), inv_rms folded after.
q_rope_chunk = q_proj_fp32[tg : tg + Q_ROPE_T_TILE, h0 + NOPE_DIM : h0 + NOPE_DIM + ROPE_DIM]
q_rope_swapped = pl.gather(q_rope_chunk, dim=-1, index=q_swap_idx)
q_rope_rot = pl.add(pl.mul(q_rope_chunk, q_cos_il), pl.mul(pl.mul(q_rope_swapped, q_sign), q_sin_il))
q_rope_rot = pl.add(pl.mul(q_rope_chunk, q_cos_il), pl.mul(q_rope_swapped, q_sin_signed))
q_flat[tg : tg + Q_ROPE_T_TILE, h0 + NOPE_DIM : h0 + NOPE_DIM + ROPE_DIM] = pl.cast(
pl.row_expand_mul(q_rope_rot, q_head_inv_rms_t), target_type=pl.BF16, mode="rint"
)
Expand Down
Loading