From c5aa9c50ed0f1810f8b629d5824f40a42f04f6c8 Mon Sep 17 00:00:00 2001 From: Siyuan Feng <25500082+Hzfengsy@users.noreply.github.com> Date: Sat, 20 Jun 2026 11:17:15 +0800 Subject: [PATCH] feat(deepseek/v4): gather-free split-half RoPE for the decode path (decode-only) Convert decode RoPE from interleaved (GPT-J, gather-based) to split-half (GPT-NeoX, gather-free): forward in qkv_proj_rope + decode compressors, inverse in sparse_attn hca/swa/csa. Removes the j^1 swap-gather, j>>1 dup-gather, and the rope_cs pre-pass; the rotation partner is now a contiguous lo/hi half-slice. Bit-exact on a2a3 (HCA/SWA standalone+composed+layer, CSA standalone+composed). L2 swimlane: rope compute -70.5%, HCA attention-module wall-clock -9.6%; all VGATHERs eliminated. Decode-only: qkv_proj_rope is shared with prefill, which is left interleaved (latently half-converted) -- prefill conversion is a tracked follow-up. A deployment-time offline weight-permutation (q-proj/k_pe/wo_a rope columns) is required for real checkpoints; synthetic tests need none. Co-Authored-By: Claude Opus 4.8 --- models/deepseek/v4/decode_attention_csa.py | 20 ++- models/deepseek/v4/decode_attention_hca.py | 23 ++- models/deepseek/v4/decode_attention_swa.py | 23 ++- .../deepseek/v4/decode_compressor_ratio128.py | 48 +++--- .../deepseek/v4/decode_compressor_ratio4.py | 48 +++--- models/deepseek/v4/decode_layer_dp_ep.py | 4 + models/deepseek/v4/decode_sparse_attn.py | 145 ++++++++---------- models/deepseek/v4/decode_sparse_attn_hca.py | 142 ++++++++--------- models/deepseek/v4/decode_sparse_attn_swa.py | 142 ++++++++--------- models/deepseek/v4/qkv_proj_rope.py | 101 ++++++------ 10 files changed, 343 insertions(+), 353 deletions(-) diff --git a/models/deepseek/v4/decode_attention_csa.py b/models/deepseek/v4/decode_attention_csa.py index 03f7a54d..87c60858 100644 --- a/models/deepseek/v4/decode_attention_csa.py +++ b/models/deepseek/v4/decode_attention_csa.py @@ -185,6 +185,11 @@ def attention_csa( rope_cos_t = pl.create_tensor([T, ROPE_HEAD_DIM], dtype=pl.BF16) rope_sin_t = pl.create_tensor([T, ROPE_HEAD_DIM], dtype=pl.BF16) + # Half-width unsigned inverse-RoPE snapshots for the split-half sparse_attn: the first + # HALF columns of the per-token cos/sin (one value per frequency), cast BF16->FP32 once per + # token so sparse_attn's per-head rope loop rotates with no in-loop cast and no gather. + rope_cos_half_t = pl.create_tensor([T, HALF_ROPE], dtype=pl.FP32) + rope_sin_half_t = pl.create_tensor([T, HALF_ROPE], dtype=pl.FP32) step_cos = pl.create_tensor([B, HALF_ROPE], dtype=pl.FP32) step_sin = pl.create_tensor([B, HALF_ROPE], dtype=pl.FP32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="csa_rope_step"): @@ -199,6 +204,10 @@ def attention_csa( sin_row = pl.cast(pl.slice(freqs_sin, [1, ROPE_HEAD_DIM], [pos_b, 0]), target_type=pl.FP32) rope_cos_t = pl.assemble(rope_cos_t, pl.cast(cos_row, target_type=pl.BF16), [t, 0]) rope_sin_t = pl.assemble(rope_sin_t, pl.cast(sin_row, target_type=pl.BF16), [t, 0]) + rope_cos_half_t = pl.assemble( + rope_cos_half_t, pl.cast(pl.slice(freqs_cos, [1, HALF_ROPE], [pos_b, 0]), target_type=pl.FP32), [t, 0]) + rope_sin_half_t = pl.assemble( + rope_sin_half_t, pl.cast(pl.slice(freqs_sin, [1, HALF_ROPE], [pos_b, 0]), target_type=pl.FP32), [t, 0]) step_cos = pl.assemble(step_cos, pl.cast(pl.slice(freqs_cos, [1, HALF_ROPE], [step_pos_b, 0]), target_type=pl.FP32), [b, 0]) step_sin = pl.assemble(step_sin, pl.cast(pl.slice(freqs_sin, [1, HALF_ROPE], [step_pos_b, 0]), target_type=pl.FP32), [b, 0]) @@ -357,8 +366,8 @@ def attention_csa( cmp_block_table, cmp_sparse_indices, attn_sink, - rope_cos_t, - rope_sin_t, + rope_cos_half_t, + rope_sin_half_t, wo_a, wo_b, wo_b_scale, @@ -523,6 +532,9 @@ def golden_attention_csa(tensors): freqs_sin = tensors["freqs_sin"] rope_cos_t = freqs_cos[position_ids].contiguous() rope_sin_t = freqs_sin[position_ids].contiguous() + # Half-width unsigned inverse-RoPE tables (first HALF columns, FP32) for the split-half sparse_attn golden. + rope_cos_half_t = freqs_cos[position_ids, :HALF_ROPE].float().contiguous() + rope_sin_half_t = freqs_sin[position_ids, :HALF_ROPE].float().contiguous() first_pos = position_ids.reshape(B, S)[:, 0] step_cos = freqs_cos[first_pos, :HALF_ROPE].float().contiguous() step_sin = freqs_sin[first_pos, :HALF_ROPE].float().contiguous() @@ -644,8 +656,8 @@ def golden_attention_csa(tensors): "cmp_block_table": cmp_block_table, "cmp_sparse_indices": sparse_topk, "attn_sink": tensors["attn_sink"], - "freqs_cos": rope_cos_t, - "freqs_sin": rope_sin_t, + "rope_cos_half": rope_cos_half_t, + "rope_sin_half": rope_sin_half_t, "wo_a": tensors["wo_a"], "wo_b": tensors["wo_b"], "wo_b_scale": tensors["wo_b_scale"], diff --git a/models/deepseek/v4/decode_attention_hca.py b/models/deepseek/v4/decode_attention_hca.py index 9b3f2463..2322a15b 100644 --- a/models/deepseek/v4/decode_attention_hca.py +++ b/models/deepseek/v4/decode_attention_hca.py @@ -141,6 +141,12 @@ def attention_hca( rope_cos_t = pl.create_tensor([T, ROPE_HEAD_DIM], dtype=pl.BF16) rope_sin_t = pl.create_tensor([T, ROPE_HEAD_DIM], dtype=pl.BF16) + # Half-width unsigned inverse-RoPE snapshots for the split-half sparse_attn_hca: the first + # HALF columns of the per-token cos/sin (one value per frequency), cast BF16->FP32 once per + # token so the per-head rope loop rotates with no in-loop cast and no gather. Same first-HALF + # columns qkv_proj_rope's forward rope consumes -> a single rope profile, no separate il table. + rope_cos_half_t = pl.create_tensor([T, ROPE_HEAD_DIM // 2], dtype=pl.FP32) + rope_sin_half_t = pl.create_tensor([T, ROPE_HEAD_DIM // 2], dtype=pl.FP32) cmp_cos = pl.create_tensor([B, ROPE_HEAD_DIM // 2], dtype=pl.FP32) cmp_sin = pl.create_tensor([B, ROPE_HEAD_DIM // 2], dtype=pl.FP32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="hca_rope"): @@ -160,6 +166,10 @@ def attention_hca( step_sin_row = pl.cast(freqs_sin[pos_b : pos_b + 1, 0 : ROPE_HEAD_DIM], target_type=pl.FP32) rope_cos_t[t : t + 1, 0 : ROPE_HEAD_DIM] = pl.cast(step_cos_row, target_type=pl.BF16, mode="rint") rope_sin_t[t : t + 1, 0 : ROPE_HEAD_DIM] = pl.cast(step_sin_row, target_type=pl.BF16, mode="rint") + rope_cos_half_t[t : t + 1, 0 : ROPE_HEAD_DIM // 2] = pl.cast( + freqs_cos[pos_b : pos_b + 1, 0 : ROPE_HEAD_DIM // 2], target_type=pl.FP32) + rope_sin_half_t[t : t + 1, 0 : ROPE_HEAD_DIM // 2] = pl.cast( + freqs_sin[pos_b : pos_b + 1, 0 : ROPE_HEAD_DIM // 2], target_type=pl.FP32) q = pl.create_tensor([T, H, HEAD_DIM], dtype=pl.BF16) kv = pl.create_tensor([T, HEAD_DIM], dtype=pl.BF16) @@ -262,8 +272,8 @@ def attention_hca( cmp_block_table, topk_all, attn_sink, - rope_cos_t, - rope_sin_t, + rope_cos_half_t, + rope_sin_half_t, wo_a, wo_b, wo_b_scale, @@ -384,10 +394,15 @@ def golden_attention_hca(tensors): freqs_sin = tensors["freqs_sin"] rope_cos_T = torch.empty(T, rd, dtype=freqs_cos.dtype) rope_sin_T = torch.empty(T, rd, dtype=freqs_sin.dtype) + # Half-width unsigned inverse-RoPE tables for the split-half sparse_attn_hca golden. + rope_cos_half_T = torch.empty(T, rd // 2, dtype=torch.float32) + rope_sin_half_T = torch.empty(T, rd // 2, dtype=torch.float32) for t in range(T): pos = int(position_ids[t].item()) rope_cos_T[t] = freqs_cos[pos] rope_sin_T[t] = freqs_sin[pos] + rope_cos_half_T[t] = freqs_cos[pos, : rd // 2].float() + rope_sin_half_T[t] = freqs_sin[pos, : rd // 2].float() # q + win kv (W8A8 q_proj) q = torch.zeros(T, H, HEAD_DIM, dtype=torch.bfloat16) @@ -480,8 +495,8 @@ def golden_attention_hca(tensors): "cmp_block_table": cmp_block_table, "cmp_sparse_indices": topk_all, "attn_sink": tensors["attn_sink"], - "freqs_cos": rope_cos_T, - "freqs_sin": rope_sin_T, + "rope_cos_half": rope_cos_half_T, + "rope_sin_half": rope_sin_half_T, "wo_a": tensors["wo_a"], "wo_b": tensors["wo_b"], "wo_b_scale": tensors["wo_b_scale"], diff --git a/models/deepseek/v4/decode_attention_swa.py b/models/deepseek/v4/decode_attention_swa.py index c90a5634..580cc055 100644 --- a/models/deepseek/v4/decode_attention_swa.py +++ b/models/deepseek/v4/decode_attention_swa.py @@ -111,6 +111,12 @@ def attention_swa( rope_cos_t = pl.create_tensor([T, ROPE_HEAD_DIM], dtype=pl.BF16) rope_sin_t = pl.create_tensor([T, ROPE_HEAD_DIM], dtype=pl.BF16) + # Half-width unsigned inverse-RoPE snapshots for the split-half sparse_attn_swa: the first + # HALF columns of the per-token cos/sin (one value per frequency), cast BF16->FP32 once per + # token so the per-head rope loop rotates with no in-loop cast and no gather. Same first-HALF + # columns qkv_proj_rope's forward rope consumes -> a single rope profile, no separate il table. + rope_cos_half_t = pl.create_tensor([T, ROPE_HEAD_DIM // 2], dtype=pl.FP32) + rope_sin_half_t = pl.create_tensor([T, ROPE_HEAD_DIM // 2], dtype=pl.FP32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="swa_rope_step"): for b in pl.parallel(B): for s_idx in pl.range(S): @@ -120,6 +126,10 @@ def attention_swa( sin_row = pl.cast(pl.slice(freqs_sin, [1, ROPE_HEAD_DIM], [pos_b, 0]), target_type=pl.FP32) rope_cos_t = pl.assemble(rope_cos_t, pl.cast(cos_row, target_type=pl.BF16, mode="rint"), [t, 0]) rope_sin_t = pl.assemble(rope_sin_t, pl.cast(sin_row, target_type=pl.BF16, mode="rint"), [t, 0]) + rope_cos_half_t = pl.assemble( + rope_cos_half_t, cos_row[0 : 1, 0 : ROPE_HEAD_DIM // 2], [t, 0]) + rope_sin_half_t = pl.assemble( + rope_sin_half_t, sin_row[0 : 1, 0 : ROPE_HEAD_DIM // 2], [t, 0]) q = pl.create_tensor([T, H, HEAD_DIM], dtype=pl.BF16) kv = pl.create_tensor([T, HEAD_DIM], dtype=pl.BF16) @@ -185,8 +195,8 @@ def attention_swa( cmp_block_table, sparse_topk, attn_sink, - rope_cos_t, - rope_sin_t, + rope_cos_half_t, + rope_sin_half_t, wo_a, wo_b, wo_b_scale, @@ -300,10 +310,15 @@ def golden_attention_swa(tensors): freqs_sin = tensors["freqs_sin"] rope_cos_T = torch.empty(T, rd, dtype=freqs_cos.dtype) rope_sin_T = torch.empty(T, rd, dtype=freqs_sin.dtype) + # Half-width unsigned inverse-RoPE tables for the split-half sparse_attn_swa golden. + rope_cos_half_T = torch.empty(T, rd // 2, dtype=torch.float32) + rope_sin_half_T = torch.empty(T, rd // 2, dtype=torch.float32) for t in range(T): pos = int(position_ids[t].item()) rope_cos_T[t] = freqs_cos[pos] rope_sin_T[t] = freqs_sin[pos] + rope_cos_half_T[t] = freqs_cos[pos, : rd // 2].float() + rope_sin_half_T[t] = freqs_sin[pos, : rd // 2].float() # q + win kv (model.py:495-504) q = torch.zeros(T, H, HEAD_DIM, dtype=torch.bfloat16) @@ -363,8 +378,8 @@ def golden_attention_swa(tensors): "cmp_block_table": cmp_block_table_dummy, "cmp_sparse_indices": sparse_topk_all, "attn_sink": tensors["attn_sink"], - "freqs_cos": rope_cos_T, - "freqs_sin": rope_sin_T, + "rope_cos_half": rope_cos_half_T, + "rope_sin_half": rope_sin_half_T, "wo_a": tensors["wo_a"], "wo_b": tensors["wo_b"], "wo_b_scale": tensors["wo_b_scale"], diff --git a/models/deepseek/v4/decode_compressor_ratio128.py b/models/deepseek/v4/decode_compressor_ratio128.py index e3837cf7..5bb05fcd 100644 --- a/models/deepseek/v4/decode_compressor_ratio128.py +++ b/models/deepseek/v4/decode_compressor_ratio128.py @@ -194,26 +194,22 @@ def compressor_ratio128( kv_rope_norm = pooled_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM : HEAD_DIM] gamma_rope = norm_w_2d[:, NOPE_HEAD_DIM : HEAD_DIM] - # A3 interleaved swap-gather (same form as kv_rope_fused in qkv_proj_rope), - # replacing the de-interleave gather + rotate + re-interleave scatter. gamma+inv_rms - # are folded into rope_normed BEFORE the swap, so the swapped lane n[j^1] correctly - # carries gamma[j^1]; inv_rms is per-row so it commutes. swap_idx (j^1), sign - # ([-1,+1,...]) and dup_idx (j>>1) are built IN-KERNEL from pl.arange; cos_il/sin_il - # are dup-gathered from the per-batch cos/sin rows. normed_kv is FP32 -> write directly. - # out[j] = n[j]*cos_il[j] + n[j^1]*sign[j]*sin_il[j] - rope_normed = pl.col_expand_mul(pl.row_expand_mul(kv_rope_norm, inv_rms), gamma_rope) - rope_ones = pl.full([RMS_TILE, ROPE_HEAD_DIM], dtype=pl.FP32, value=1.0) - rope_col = pl.col_expand_mul(rope_ones, pl.cast(pl.arange(0, [1, ROPE_HEAD_DIM], dtype=pl.INT32), target_type=pl.FP32)) - rope_dup_f = pl.cast(pl.cast(pl.mul(rope_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32) - rope_dup_idx = pl.cast(rope_dup_f, target_type=pl.INT32) # j>>1 - rope_lane = pl.sub(rope_col, pl.mul(rope_dup_f, 2.0)) # j%2 - rope_swap_idx = pl.cast(pl.sub(pl.add(rope_col, 1.0), pl.mul(rope_lane, 2.0)), target_type=pl.INT32) # j^1 - rope_sign = pl.sub(pl.mul(rope_lane, 2.0), 1.0) # [-1,+1,...] - cos_il = pl.gather(cos_b, dim=-1, index=rope_dup_idx) - sin_il = pl.gather(sin_b, dim=-1, index=rope_dup_idx) - swapped = pl.gather(rope_normed, dim=-1, index=rope_swap_idx) - rope_rot = pl.add(pl.mul(rope_normed, cos_il), pl.mul(pl.mul(swapped, rope_sign), sin_il)) - normed_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM : HEAD_DIM] = rope_rot + # Split-half (NeoX) forward RoPE, gather-free. rope segment = [x_lo | x_hi] = dims + # [0:rh | rh:ROPE_HEAD_DIM]; partner of lane k is k+rh (contiguous slice, no j^1 gather + # and no j>>1 dup-gather -- cos_b/sin_b are already half-width). gamma is per-column so + # it does NOT commute with the rotation: fold gamma_lo onto x_lo, gamma_hi onto x_hi + # BEFORE rotating; inv_rms is per-row and commutes. + # out_lo = x_lo*cos - x_hi*sin ; out_hi = x_lo*sin + x_hi*cos + gamma_lo = gamma_rope[0:1, 0 : ROPE_HEAD_DIM // 2] + gamma_hi = gamma_rope[0:1, ROPE_HEAD_DIM // 2 : ROPE_HEAD_DIM] + kv_lo = kv_rope_norm[0 : RMS_TILE, 0 : ROPE_HEAD_DIM // 2] + kv_hi = kv_rope_norm[0 : RMS_TILE, ROPE_HEAD_DIM // 2 : ROPE_HEAD_DIM] + lo_n = pl.col_expand_mul(pl.row_expand_mul(kv_lo, inv_rms), gamma_lo) + hi_n = pl.col_expand_mul(pl.row_expand_mul(kv_hi, inv_rms), gamma_hi) + out_lo = pl.sub(pl.mul(lo_n, cos_b), pl.mul(hi_n, sin_b)) # x_lo*c - x_hi*s + out_hi = pl.add(pl.mul(lo_n, sin_b), pl.mul(hi_n, cos_b)) # x_lo*s + x_hi*c + normed_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM : NOPE_HEAD_DIM + ROPE_HEAD_DIM // 2] = out_lo + normed_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM + ROPE_HEAD_DIM // 2 : HEAD_DIM] = out_hi kv_flat = pl.reshape(kv, [bs, HEAD_DIM]) cmp_flat_rows = cmp_block_num * BLOCK_SIZE @@ -367,13 +363,15 @@ def rmsnorm(x, w): continue kv_b = rmsnorm(pooled[b : b + 1], norm_w) - x_pair = kv_b[..., -rd:].unflatten(-1, (-1, 2)) - x0, x1 = x_pair[..., 0], x_pair[..., 1] + # Split-half (NeoX) forward RoPE: lo = first rd/2 rope dims, hi = last rd/2. + rope = kv_b[..., -rd:] + x_lo = rope[..., : rd // 2] + x_hi = rope[..., rd // 2 :] cos_v, sin_v = cos[b].view(-1), sin[b].view(-1) - y0 = x0 * cos_v - x1 * sin_v - y1 = x0 * sin_v + x1 * cos_v + y_lo = x_lo * cos_v - x_hi * sin_v + y_hi = x_lo * sin_v + x_hi * cos_v - kv_b = torch.cat([kv_b[..., :-rd], torch.stack([y0, y1], dim=-1).flatten(-2)], dim=-1) + kv_b = torch.cat([kv_b[..., :-rd], y_lo, y_hi], dim=-1) # Kernel writes pooled result only to kv[:, 0, :]; leave kv[:, 1:, :] = 0. tensors["kv"][b : b + 1, 0:1, :] = kv_b diff --git a/models/deepseek/v4/decode_compressor_ratio4.py b/models/deepseek/v4/decode_compressor_ratio4.py index 2cf436f2..ac2ccf1f 100644 --- a/models/deepseek/v4/decode_compressor_ratio4.py +++ b/models/deepseek/v4/decode_compressor_ratio4.py @@ -195,26 +195,22 @@ def compressor_ratio4( kv_rope_norm = pooled_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM : HEAD_DIM] gamma_rope = norm_w_2d[:, NOPE_HEAD_DIM : HEAD_DIM] - # A3 interleaved swap-gather (same form as kv_rope_fused in qkv_proj_rope), - # replacing the de-interleave gather + rotate + re-interleave scatter. gamma+inv_rms - # are folded into rope_normed BEFORE the swap, so the swapped lane n[j^1] correctly - # carries gamma[j^1]; inv_rms is per-row so it commutes. swap_idx (j^1), sign - # ([-1,+1,...]) and dup_idx (j>>1) are built IN-KERNEL from pl.arange; cos_il/sin_il - # are dup-gathered from the per-batch cos/sin rows. normed_kv is FP32 -> write directly. - # out[j] = n[j]*cos_il[j] + n[j^1]*sign[j]*sin_il[j] - rope_normed = pl.col_expand_mul(pl.row_expand_mul(kv_rope_norm, inv_rms), gamma_rope) - rope_ones = pl.full([RMS_TILE, ROPE_HEAD_DIM], dtype=pl.FP32, value=1.0) - rope_col = pl.col_expand_mul(rope_ones, pl.cast(pl.arange(0, [1, ROPE_HEAD_DIM], dtype=pl.INT32), target_type=pl.FP32)) - rope_dup_f = pl.cast(pl.cast(pl.mul(rope_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32) - rope_dup_idx = pl.cast(rope_dup_f, target_type=pl.INT32) # j>>1 - rope_lane = pl.sub(rope_col, pl.mul(rope_dup_f, 2.0)) # j%2 - rope_swap_idx = pl.cast(pl.sub(pl.add(rope_col, 1.0), pl.mul(rope_lane, 2.0)), target_type=pl.INT32) # j^1 - rope_sign = pl.sub(pl.mul(rope_lane, 2.0), 1.0) # [-1,+1,...] - cos_il = pl.gather(cos_b, dim=-1, index=rope_dup_idx) - sin_il = pl.gather(sin_b, dim=-1, index=rope_dup_idx) - swapped = pl.gather(rope_normed, dim=-1, index=rope_swap_idx) - rope_rot = pl.add(pl.mul(rope_normed, cos_il), pl.mul(pl.mul(swapped, rope_sign), sin_il)) - normed_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM : HEAD_DIM] = rope_rot + # Split-half (NeoX) forward RoPE, gather-free. rope segment = [x_lo | x_hi] = dims + # [0:rh | rh:ROPE_HEAD_DIM]; partner of lane k is k+rh (contiguous slice, no j^1 gather + # and no j>>1 dup-gather -- cos_b/sin_b are already half-width). gamma is per-column so + # it does NOT commute with the rotation: fold gamma_lo onto x_lo, gamma_hi onto x_hi + # BEFORE rotating; inv_rms is per-row and commutes. + # out_lo = x_lo*cos - x_hi*sin ; out_hi = x_lo*sin + x_hi*cos + gamma_lo = gamma_rope[0:1, 0 : ROPE_HEAD_DIM // 2] + gamma_hi = gamma_rope[0:1, ROPE_HEAD_DIM // 2 : ROPE_HEAD_DIM] + kv_lo = kv_rope_norm[0 : RMS_TILE, 0 : ROPE_HEAD_DIM // 2] + kv_hi = kv_rope_norm[0 : RMS_TILE, ROPE_HEAD_DIM // 2 : ROPE_HEAD_DIM] + lo_n = pl.col_expand_mul(pl.row_expand_mul(kv_lo, inv_rms), gamma_lo) + hi_n = pl.col_expand_mul(pl.row_expand_mul(kv_hi, inv_rms), gamma_hi) + out_lo = pl.sub(pl.mul(lo_n, cos_b), pl.mul(hi_n, sin_b)) # x_lo*c - x_hi*s + out_hi = pl.add(pl.mul(lo_n, sin_b), pl.mul(hi_n, cos_b)) # x_lo*s + x_hi*c + normed_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM : NOPE_HEAD_DIM + ROPE_HEAD_DIM // 2] = out_lo + normed_kv[batch_base : batch_base + RMS_TILE, NOPE_HEAD_DIM + ROPE_HEAD_DIM // 2 : HEAD_DIM] = out_hi for batch_base_idx in pl.spmd(B // RMS_TILE, name_hint="kv_and_cache_write"): batch_base = batch_base_idx * RMS_TILE @@ -361,13 +357,15 @@ def rmsnorm(x, w): boundary_s = ratio - 1 - (first_pos % ratio) kv_b = rmsnorm(pooled[b : b + 1], norm_w) - x_pair = kv_b[..., -rd:].unflatten(-1, (-1, 2)) - x0, x1 = x_pair[..., 0], x_pair[..., 1] + # Split-half (NeoX) forward RoPE: lo = first rd/2 rope dims, hi = last rd/2. + rope = kv_b[..., -rd:] + x_lo = rope[..., : rd // 2] + x_hi = rope[..., rd // 2 :] cos_v, sin_v = cos[b].view(-1), sin[b].view(-1) - y0 = x0 * cos_v - x1 * sin_v - y1 = x0 * sin_v + x1 * cos_v + y_lo = x_lo * cos_v - x_hi * sin_v + y_hi = x_lo * sin_v + x_hi * cos_v - kv_b = torch.cat([kv_b[..., :-rd], torch.stack([y0, y1], dim=-1).flatten(-2)], dim=-1) + kv_b = torch.cat([kv_b[..., :-rd], y_lo, y_hi], dim=-1) # Kernel writes pooled result only to kv[:, 0, :]; leave kv[:, 1:, :] = 0 # so the golden matches its [B, S, HEAD_DIM] zero-init. diff --git a/models/deepseek/v4/decode_layer_dp_ep.py b/models/deepseek/v4/decode_layer_dp_ep.py index 9275ae4a..bbd84278 100644 --- a/models/deepseek/v4/decode_layer_dp_ep.py +++ b/models/deepseek/v4/decode_layer_dp_ep.py @@ -652,6 +652,10 @@ def build_tensor_specs(start_pos=None, layer_id=10): "csa": csa_specs, }[attention_kind] + # Split-half (NeoX) RoPE: the sparse_attn kernels read the half-width unsigned cos/sin + # straight from the first HALF columns of freqs_cos/freqs_sin (the caller slices them), so + # there is no separate interleaved/sign-folded table -- one rope profile for qkv + sparse_attn. + replicated_attention = { "hc_attn_fn", "hc_attn_scale", diff --git a/models/deepseek/v4/decode_sparse_attn.py b/models/deepseek/v4/decode_sparse_attn.py index 6c7a8717..f3159790 100644 --- a/models/deepseek/v4/decode_sparse_attn.py +++ b/models/deepseek/v4/decode_sparse_attn.py @@ -56,8 +56,6 @@ # (its [64,128] softmax and co-resident QK+PV L0C accumulators overflow Vec/L0C). QK_M_TILE = 32 ATTN_K_TILE = 128 -ROPE_TILE = 16 -ROPE_INTERLEAVE_TILE = 2 * ROPE_TILE A_T_TILE = 32 A_K_TILE = 128 A_N_TILE = 128 @@ -114,8 +112,13 @@ def sparse_attn( cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32], cmp_sparse_indices: pl.Tensor[[T, TOPK], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], - freqs_cos: pl.Tensor[[T, ROPE_DIM], pl.BF16], - freqs_sin: pl.Tensor[[T, ROPE_DIM], pl.BF16], + # Half-width unsigned inverse-RoPE tables for the gather-free split-half rotate: + # rope_cos_half = [c0,c1,...,c_{HALF-1}] rope_sin_half = [s0,s1,...] (one per freq) + # The rope segment is split-half ([x_lo | x_hi] = dims [0:HALF | HALF:ROPE_DIM]), + # so the rotation partner of lane k is lane k+HALF -- a contiguous slice, not a + # j^1 swap gather. FP32 so the rope loop reads them straight into its FP32 rotate. + rope_cos_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], + rope_sin_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], wo_b_scale: pl.Tensor[[D], pl.FP32], @@ -295,46 +298,26 @@ def sparse_attn( n_col = n_hh * HEAD_DIM o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + NOPE_DIM] = n_out[n_hi : n_hi + 1, 0 : NOPE_DIM] - # Precompute the head-invariant interleaved cos and sign*sin once: they depend - # only on (token, column), not head, so building them per head would repeat the - # same dup-gather H times on the bottleneck Vec engine. sign is folded into sin - # (multiply by +/-1). The conjugate (inverse) rotation is: - # out[j] = x[j]*cos_il[j] + x[j^1]*sign[j]*sin_il[j] - rope_cos_il = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32) - rope_sin_signed = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32) - for cp in pl.spmd(HALF_ROPE // ROPE_TILE, name_hint="rope_cs"): - cp_r0 = cp * ROPE_TILE - cp_c0 = 2 * cp_r0 - cs_col = pl.col_expand_mul( - pl.full([T, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0), - pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], 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) # j>>1 - cs_lane = pl.sub(cs_col, pl.mul(cs_dup_f, 2.0)) # j%2 - cs_sign = pl.neg(pl.sub(pl.mul(cs_lane, 2.0), 1.0)) # [+1,-1,...] (conjugate) - cs_cos = pl.cast(freqs_cos[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32) - cs_sin = pl.cast(freqs_sin[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32) - rope_cos_il[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx) - rope_sin_signed[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.mul( - pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign) + # Gather-free conjugate (inverse) split-half RoPE. The rope segment is laid out + # split-half ([x_lo | x_hi] = dims [0:HALF_ROPE | HALF_ROPE:ROPE_DIM]), so the + # rotation partner of lane k is lane k+HALF_ROPE -- a contiguous slice, never a + # j^1 swap gather. For k in [0, HALF_ROPE): + # out_lo[k] = x_lo[k]*cos[k] + x_hi[k]*sin[k] + # out_hi[k] = x_hi[k]*cos[k] - x_lo[k]*sin[k] + # cos/sin are half-width (one value per frequency); no in-kernel index build. # Inverse RoPE fused with the rope-column pack: each task rotates its heads' - # rope segments and stores them straight into o_packed's strided rope columns, - # dropping a separate rope_pack stage and GM round-trip. cos_il / sign*sin come - # from the pre-pass above; only swap_idx (j^1) is rebuilt per task. + # rope segments and stores them straight into o_packed's rope columns, dropping + # a separate rope_pack stage and GM round-trip. The whole ROPE_DIM is rotated in + # one pass as a [ROPE_OUT_TOK_TILE, HALF_ROPE] lo tile + matching hi tile. attn_rope_stage_3d = pl.reshape(attn_rope_stage, [T, H, ROPE_DIM]) for rp_idx in pl.spmd((H // 4) * 2, name_hint="rope"): rp_hg = rp_idx // 2 rp_tt = rp_idx - rp_hg * 2 rp_t0 = rp_tt * ROPE_OUT_TOK_TILE - # Head-invariant swap index (j^1), built once and reused across the head - # group -- the only per-head input is this head's strided rope slice. - sp_col = pl.col_expand_mul( - pl.full([ROPE_OUT_TOK_TILE, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0), - pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], dtype=pl.INT32), target_type=pl.FP32)) - sp_dup_f = pl.cast(pl.cast(pl.mul(sp_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32) - sp_lane = pl.sub(sp_col, pl.mul(sp_dup_f, 2.0)) # j%2 - sp_swap_idx = pl.cast(pl.sub(pl.add(sp_col, 1.0), pl.mul(sp_lane, 2.0)), target_type=pl.INT32) # j^1 + # Head-invariant half-width tables, hoisted once per task. + r_cos = rope_cos_half[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] + r_sin = rope_sin_half[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] for rp_hl in pl.range(0, 4): rp_gh = rp_hg * 4 + rp_hl @@ -342,21 +325,19 @@ def sparse_attn( rp_hh = rp_gh - rp_g * HEADS_PER_GROUP rp_col = rp_hh * HEAD_DIM + NOPE_DIM rp_o0 = rp_g * T + rp_t0 - for r_r0 in pl.range(0, HALF_ROPE, ROPE_TILE): - c0 = 2 * r_r0 - # This head's rope rows for this token tile (stride H); gather needs - # FP/INT, so cast the strided BF16 slice to FP32 first. - r_tile_fp32 = pl.cast( - pl.reshape(attn_rope_stage_3d[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, rp_gh : rp_gh + 1, c0 : c0 + ROPE_INTERLEAVE_TILE], [ROPE_OUT_TOK_TILE, ROPE_INTERLEAVE_TILE]), - target_type=pl.FP32) - r_cos_il = rope_cos_il[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, c0 : c0 + ROPE_INTERLEAVE_TILE] - r_sin_signed = rope_sin_signed[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, c0 : c0 + ROPE_INTERLEAVE_TILE] - r_swapped = pl.gather(r_tile_fp32, dim=-1, index=sp_swap_idx) - r_rot = pl.add(pl.mul(r_tile_fp32, r_cos_il), pl.mul(r_swapped, r_sin_signed)) - # Store BF16-rounded rotated values straight into o_packed's rope - # columns for this head (golden also rounds inverse-RoPE to bf16). - r_rot = pl.cast(r_rot, target_type=pl.BF16, mode="rint") - o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col + c0 : rp_col + c0 + ROPE_INTERLEAVE_TILE] = r_rot + # This head's full ROPE_DIM rope rows for this token tile (stride H); the + # rotate math is FP32, so cast the strided BF16 slice to FP32 first. + r_tile = pl.cast( + pl.reshape(attn_rope_stage_3d[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, rp_gh : rp_gh + 1, 0 : ROPE_DIM], [ROPE_OUT_TOK_TILE, ROPE_DIM]), + target_type=pl.FP32) + r_lo = r_tile[0 : ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] # x_lo (contiguous, no gather) + r_hi = r_tile[0 : ROPE_OUT_TOK_TILE, HALF_ROPE : ROPE_DIM] # x_hi (contiguous, no gather) + out_lo = pl.add(pl.mul(r_lo, r_cos), pl.mul(r_hi, r_sin)) # x_lo*c + x_hi*s + out_hi = pl.sub(pl.mul(r_hi, r_cos), pl.mul(r_lo, r_sin)) # x_hi*c - x_lo*s + # Two contiguous BF16-rounded stores into o_packed's rope columns + # (golden also rounds inverse-RoPE to bf16). + o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col : rp_col + HALF_ROPE] = pl.cast(out_lo, target_type=pl.BF16, mode="rint") + o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col + HALF_ROPE : rp_col + ROPE_DIM] = pl.cast(out_hi, target_type=pl.BF16, mode="rint") # Grouped BF16 projection `o_packed @ wo_a^T` -> `o_r`. Vec post-process # (BF16 store + per-row amax) is T-tiled to keep the fused AIV side from @@ -450,8 +431,8 @@ def sparse_attn_test( cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32], cmp_sparse_indices: pl.Tensor[[T, TOPK], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], - freqs_cos: pl.Tensor[[T, ROPE_DIM], pl.BF16], - freqs_sin: pl.Tensor[[T, ROPE_DIM], pl.BF16], + rope_cos_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], + rope_sin_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], wo_b_scale: pl.Tensor[[D], pl.FP32], @@ -466,8 +447,8 @@ def sparse_attn_test( cmp_block_table, cmp_sparse_indices, attn_sink, - freqs_cos, - freqs_sin, + rope_cos_half, + rope_sin_half, wo_a, wo_b, wo_b_scale, @@ -512,8 +493,10 @@ def golden_sparse_attn(tensors): cmp_block_table = tensors["cmp_block_table"] cmp_sparse_indices = tensors["cmp_sparse_indices"] attn_sink = tensors["attn_sink"].float() - cos = tensors["freqs_cos"].float() - sin = tensors["freqs_sin"].float() + # Half-width unsigned inverse-RoPE tables (one value per frequency), read + # directly by the split-half reference rotation below. + cos = tensors["rope_cos_half"].float() + sin = tensors["rope_sin_half"].float() wo_a = tensors["wo_a"].float() wo_b_i8 = tensors["wo_b"] wo_b_scale = tensors["wo_b_scale"].float() @@ -592,14 +575,15 @@ def golden_sparse_attn(tensors): denom = li + torch.exp(attn_sink.unsqueeze(-1) - score_max) o[t] = oi_num / denom - rope_pair = o[..., NOPE_DIM:].unflatten(-1, (-1, 2)) - rope_even = rope_pair[..., 0] - rope_odd = rope_pair[..., 1] - cos_half = cos[:, :HALF_ROPE].unsqueeze(1) - sin_half = sin[:, :HALF_ROPE].unsqueeze(1) - inv_even = (rope_even * cos_half + rope_odd * sin_half).to(torch.bfloat16).float() - inv_odd = (rope_odd * cos_half - rope_even * sin_half).to(torch.bfloat16).float() - o_rope = torch.stack([inv_even, inv_odd], dim=-1).flatten(-2) + # Split-half inverse RoPE: lo = dims [:HALF_ROPE], hi = dims [HALF_ROPE:]. + rope_seg = o[..., NOPE_DIM:] + x_lo = rope_seg[..., :HALF_ROPE] + x_hi = rope_seg[..., HALF_ROPE:] + cos_h = cos.unsqueeze(1) # [T, 1, HALF_ROPE] broadcast over heads + sin_h = sin.unsqueeze(1) + inv_lo = (x_lo * cos_h + x_hi * sin_h).to(torch.bfloat16).float() + inv_hi = (x_hi * cos_h - x_lo * sin_h).to(torch.bfloat16).float() + o_rope = torch.cat([inv_lo, inv_hi], dim=-1) o = torch.cat([o[..., :NOPE_DIM], o_rope], dim=-1).to(torch.bfloat16) seq_per_batch = T // B @@ -623,15 +607,8 @@ def build_tensor_specs( """Build deterministic demo tensors for the merged standalone harness.""" import torch from golden import TensorSpec - from rope_tables import build_deepseek_v4_rope_tables, materialize_token_rope_tables cmp_valid = get_standalone_cmp_valid(compress_ratio) - shared_freqs_cos, shared_freqs_sin = build_deepseek_v4_rope_tables(M, compress_ratio, dtype=torch.bfloat16) - shared_rope_cos, shared_rope_sin = materialize_token_rope_tables( - shared_freqs_cos, - shared_freqs_sin, - torch.arange(T, dtype=torch.int32), - ) def seeded_uniform(shape, seed): """Create a deterministic centered uniform tensor for repeatable tests.""" @@ -716,13 +693,19 @@ def init_cmp_sparse_indices(): indices[0, WIN - 1] = WIN - 1 return indices - def init_cos(): - """Build the split-half cosine table used by the inverse-RoPE reference.""" - return shared_rope_cos.clone() + # Half-width unsigned cos/sin (one value per frequency), rounded to BF16 then back to + # FP32 so the FP32 kernel input holds exactly the bf16 values (bit-exact with the reference). + angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 + rope_cos_half_tbl = torch.cos(angles).to(torch.bfloat16).float() + rope_sin_half_tbl = torch.sin(angles).to(torch.bfloat16).float() + + def init_rope_cos_half(): + """Build the half-width cosine table [c0,c1,...] read directly by the split-half inverse-RoPE kernel.""" + return rope_cos_half_tbl.clone() - def init_sin(): - """Build the split-half sine table used by the inverse-RoPE reference.""" - return shared_rope_sin.clone() + def init_rope_sin_half(): + """Build the half-width sine table [s0,s1,...] read directly by the split-half inverse-RoPE kernel.""" + return rope_sin_half_tbl.clone() def init_wo_a(): """Initialize the grouped first-stage output-projection weights.""" @@ -748,8 +731,8 @@ def init_wo_b_scale(): TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("cmp_sparse_indices", [T, TOPK], torch.int32, init_value=init_cmp_sparse_indices), TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), - TensorSpec("freqs_cos", [T, ROPE_DIM], torch.bfloat16, init_value=init_cos), - TensorSpec("freqs_sin", [T, ROPE_DIM], torch.bfloat16, init_value=init_sin), + TensorSpec("rope_cos_half", [T, HALF_ROPE], torch.float32, init_value=init_rope_cos_half), + TensorSpec("rope_sin_half", [T, HALF_ROPE], torch.float32, init_value=init_rope_sin_half), TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), diff --git a/models/deepseek/v4/decode_sparse_attn_hca.py b/models/deepseek/v4/decode_sparse_attn_hca.py index 51e36b4b..386d1be9 100644 --- a/models/deepseek/v4/decode_sparse_attn_hca.py +++ b/models/deepseek/v4/decode_sparse_attn_hca.py @@ -56,8 +56,6 @@ # (its [64,128] softmax and co-resident QK+PV L0C accumulators overflow Vec/L0C). QK_M_TILE = 32 ATTN_K_TILE = 128 -ROPE_TILE = 16 -ROPE_INTERLEAVE_TILE = 2 * ROPE_TILE A_T_TILE = 32 A_K_TILE = 128 A_N_TILE = 128 @@ -112,8 +110,13 @@ def sparse_attn_hca( cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32], cmp_sparse_indices: pl.Tensor[[T, TOPK], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], - freqs_cos: pl.Tensor[[T, ROPE_DIM], pl.BF16], - freqs_sin: pl.Tensor[[T, ROPE_DIM], pl.BF16], + # Half-width unsigned inverse-RoPE tables for the gather-free split-half rotate: + # rope_cos_half = [c0,c1,...,c_{HALF-1}] rope_sin_half = [s0,s1,...] (one per freq) + # The rope segment is split-half ([x_lo | x_hi] = dims [0:HALF | HALF:ROPE_DIM]), + # so the rotation partner of lane k is lane k+HALF -- a contiguous slice, not a + # j^1 swap gather. FP32 so the rope loop reads them straight into its FP32 rotate. + rope_cos_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], + rope_sin_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], wo_b_scale: pl.Tensor[[D], pl.FP32], @@ -293,46 +296,26 @@ def sparse_attn_hca( n_col = n_hh * HEAD_DIM o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + NOPE_DIM] = n_out[n_hi : n_hi + 1, 0 : NOPE_DIM] - # Precompute the head-invariant interleaved cos and sign*sin once: they depend - # only on (token, column), not head, so building them per head would repeat the - # same dup-gather H times on the bottleneck Vec engine. sign is folded into sin - # (multiply by +/-1). The conjugate (inverse) rotation is: - # out[j] = x[j]*cos_il[j] + x[j^1]*sign[j]*sin_il[j] - rope_cos_il = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32) - rope_sin_signed = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32) - for cp in pl.spmd(HALF_ROPE // ROPE_TILE, name_hint="rope_cs"): - cp_r0 = cp * ROPE_TILE - cp_c0 = 2 * cp_r0 - cs_col = pl.col_expand_mul( - pl.full([T, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0), - pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], 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) # j>>1 - cs_lane = pl.sub(cs_col, pl.mul(cs_dup_f, 2.0)) # j%2 - cs_sign = pl.neg(pl.sub(pl.mul(cs_lane, 2.0), 1.0)) # [+1,-1,...] (conjugate) - cs_cos = pl.cast(freqs_cos[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32) - cs_sin = pl.cast(freqs_sin[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32) - rope_cos_il[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx) - rope_sin_signed[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.mul( - pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign) + # Gather-free conjugate (inverse) split-half RoPE. The rope segment is laid out + # split-half ([x_lo | x_hi] = dims [0:HALF_ROPE | HALF_ROPE:ROPE_DIM]), so the + # rotation partner of lane k is lane k+HALF_ROPE -- a contiguous slice, never a + # j^1 swap gather. For k in [0, HALF_ROPE): + # out_lo[k] = x_lo[k]*cos[k] + x_hi[k]*sin[k] + # out_hi[k] = x_hi[k]*cos[k] - x_lo[k]*sin[k] + # cos/sin are half-width (one value per frequency); no in-kernel index build. # Inverse RoPE fused with the rope-column pack: each task rotates its heads' - # rope segments and stores them straight into o_packed's strided rope columns, - # dropping a separate rope_pack stage and GM round-trip. cos_il / sign*sin come - # from the pre-pass above; only swap_idx (j^1) is rebuilt per task. + # rope segments and stores them straight into o_packed's rope columns, dropping + # a separate rope_pack stage and GM round-trip. The whole ROPE_DIM is rotated in + # one pass as a [ROPE_OUT_TOK_TILE, HALF_ROPE] lo tile + matching hi tile. attn_rope_stage_3d = pl.reshape(attn_rope_stage, [T, H, ROPE_DIM]) for rp_idx in pl.spmd((H // 4) * 2, name_hint="rope"): rp_hg = rp_idx // 2 rp_tt = rp_idx - rp_hg * 2 rp_t0 = rp_tt * ROPE_OUT_TOK_TILE - # Head-invariant swap index (j^1), built once and reused across the head - # group -- the only per-head input is this head's strided rope slice. - sp_col = pl.col_expand_mul( - pl.full([ROPE_OUT_TOK_TILE, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0), - pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], dtype=pl.INT32), target_type=pl.FP32)) - sp_dup_f = pl.cast(pl.cast(pl.mul(sp_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32) - sp_lane = pl.sub(sp_col, pl.mul(sp_dup_f, 2.0)) # j%2 - sp_swap_idx = pl.cast(pl.sub(pl.add(sp_col, 1.0), pl.mul(sp_lane, 2.0)), target_type=pl.INT32) # j^1 + # Head-invariant half-width tables, hoisted once per task. + r_cos = rope_cos_half[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] + r_sin = rope_sin_half[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] for rp_hl in pl.range(0, 4): rp_gh = rp_hg * 4 + rp_hl @@ -340,21 +323,19 @@ def sparse_attn_hca( rp_hh = rp_gh - rp_g * HEADS_PER_GROUP rp_col = rp_hh * HEAD_DIM + NOPE_DIM rp_o0 = rp_g * T + rp_t0 - for r_r0 in pl.range(0, HALF_ROPE, ROPE_TILE): - c0 = 2 * r_r0 - # This head's rope rows for this token tile (stride H); gather needs - # FP/INT, so cast the strided BF16 slice to FP32 first. - r_tile_fp32 = pl.cast( - pl.reshape(attn_rope_stage_3d[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, rp_gh : rp_gh + 1, c0 : c0 + ROPE_INTERLEAVE_TILE], [ROPE_OUT_TOK_TILE, ROPE_INTERLEAVE_TILE]), - target_type=pl.FP32) - r_cos_il = rope_cos_il[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, c0 : c0 + ROPE_INTERLEAVE_TILE] - r_sin_signed = rope_sin_signed[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, c0 : c0 + ROPE_INTERLEAVE_TILE] - r_swapped = pl.gather(r_tile_fp32, dim=-1, index=sp_swap_idx) - r_rot = pl.add(pl.mul(r_tile_fp32, r_cos_il), pl.mul(r_swapped, r_sin_signed)) - # Store BF16-rounded rotated values straight into o_packed's rope - # columns for this head (golden also rounds inverse-RoPE to bf16). - r_rot = pl.cast(r_rot, target_type=pl.BF16, mode="rint") - o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col + c0 : rp_col + c0 + ROPE_INTERLEAVE_TILE] = r_rot + # This head's full ROPE_DIM rope rows for this token tile (stride H); the + # rotate math is FP32, so cast the strided BF16 slice to FP32 first. + r_tile = pl.cast( + pl.reshape(attn_rope_stage_3d[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, rp_gh : rp_gh + 1, 0 : ROPE_DIM], [ROPE_OUT_TOK_TILE, ROPE_DIM]), + target_type=pl.FP32) + r_lo = r_tile[0 : ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] # x_lo (contiguous, no gather) + r_hi = r_tile[0 : ROPE_OUT_TOK_TILE, HALF_ROPE : ROPE_DIM] # x_hi (contiguous, no gather) + out_lo = pl.add(pl.mul(r_lo, r_cos), pl.mul(r_hi, r_sin)) # x_lo*c + x_hi*s + out_hi = pl.sub(pl.mul(r_hi, r_cos), pl.mul(r_lo, r_sin)) # x_hi*c - x_lo*s + # Two contiguous BF16-rounded stores into o_packed's rope columns + # (golden also rounds inverse-RoPE to bf16). + o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col : rp_col + HALF_ROPE] = pl.cast(out_lo, target_type=pl.BF16, mode="rint") + o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col + HALF_ROPE : rp_col + ROPE_DIM] = pl.cast(out_hi, target_type=pl.BF16, mode="rint") # Grouped BF16 projection `o_packed @ wo_a^T` -> `o_r`. Vec post-process # (BF16 store + per-row amax) is T-tiled to keep the fused AIV side from @@ -448,8 +429,8 @@ def sparse_attn_test( cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32], cmp_sparse_indices: pl.Tensor[[T, TOPK], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], - freqs_cos: pl.Tensor[[T, ROPE_DIM], pl.BF16], - freqs_sin: pl.Tensor[[T, ROPE_DIM], pl.BF16], + rope_cos_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], + rope_sin_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], wo_b_scale: pl.Tensor[[D], pl.FP32], @@ -464,8 +445,8 @@ def sparse_attn_test( cmp_block_table, cmp_sparse_indices, attn_sink, - freqs_cos, - freqs_sin, + rope_cos_half, + rope_sin_half, wo_a, wo_b, wo_b_scale, @@ -510,8 +491,10 @@ def golden_sparse_attn(tensors): cmp_block_table = tensors["cmp_block_table"] cmp_sparse_indices = tensors["cmp_sparse_indices"] attn_sink = tensors["attn_sink"].float() - cos = tensors["freqs_cos"].float() - sin = tensors["freqs_sin"].float() + # Half-width unsigned inverse-RoPE tables (one value per frequency), read + # directly by the split-half reference rotation below. + cos = tensors["rope_cos_half"].float() + sin = tensors["rope_sin_half"].float() wo_a = tensors["wo_a"].float() wo_b_i8 = tensors["wo_b"] wo_b_scale = tensors["wo_b_scale"].float() @@ -590,14 +573,15 @@ def golden_sparse_attn(tensors): denom = li + torch.exp(attn_sink.unsqueeze(-1) - score_max) o[t] = oi_num / denom - rope_pair = o[..., NOPE_DIM:].unflatten(-1, (-1, 2)) - rope_even = rope_pair[..., 0] - rope_odd = rope_pair[..., 1] - cos_half = cos[:, :HALF_ROPE].unsqueeze(1) - sin_half = sin[:, :HALF_ROPE].unsqueeze(1) - inv_even = (rope_even * cos_half + rope_odd * sin_half).to(torch.bfloat16).float() - inv_odd = (rope_odd * cos_half - rope_even * sin_half).to(torch.bfloat16).float() - o_rope = torch.stack([inv_even, inv_odd], dim=-1).flatten(-2) + # Split-half inverse RoPE: lo = dims [:HALF_ROPE], hi = dims [HALF_ROPE:]. + rope_seg = o[..., NOPE_DIM:] + x_lo = rope_seg[..., :HALF_ROPE] + x_hi = rope_seg[..., HALF_ROPE:] + cos_h = cos.unsqueeze(1) # [T, 1, HALF_ROPE] broadcast over heads + sin_h = sin.unsqueeze(1) + inv_lo = (x_lo * cos_h + x_hi * sin_h).to(torch.bfloat16).float() + inv_hi = (x_hi * cos_h - x_lo * sin_h).to(torch.bfloat16).float() + o_rope = torch.cat([inv_lo, inv_hi], dim=-1) o = torch.cat([o[..., :NOPE_DIM], o_rope], dim=-1).to(torch.bfloat16) seq_per_batch = T // B @@ -707,17 +691,19 @@ def init_cmp_sparse_indices(): indices[0, WIN - 1] = WIN - 1 return indices - def init_cos(): - """Build the split-half cosine table used by the inverse-RoPE reference.""" - angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 - cos_half = torch.cos(angles) - return torch.cat([cos_half, cos_half], dim=-1) + # Half-width unsigned cos/sin (one value per frequency), rounded to BF16 then back to + # FP32 so the FP32 kernel input holds exactly the bf16 values (bit-exact with the reference). + angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 + rope_cos_half_tbl = torch.cos(angles).to(torch.bfloat16).float() + rope_sin_half_tbl = torch.sin(angles).to(torch.bfloat16).float() - def init_sin(): - """Build the split-half sine table used by the inverse-RoPE reference.""" - angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 - sin_half = torch.sin(angles) - return torch.cat([sin_half, sin_half], dim=-1) + def init_rope_cos_half(): + """Build the half-width cosine table [c0,c1,...] read directly by the split-half inverse-RoPE kernel.""" + return rope_cos_half_tbl.clone() + + def init_rope_sin_half(): + """Build the half-width sine table [s0,s1,...] read directly by the split-half inverse-RoPE kernel.""" + return rope_sin_half_tbl.clone() def init_wo_a(): """Initialize the grouped first-stage output-projection weights.""" @@ -743,8 +729,8 @@ def init_wo_b_scale(): TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("cmp_sparse_indices", [T, TOPK], torch.int32, init_value=init_cmp_sparse_indices), TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), - TensorSpec("freqs_cos", [T, ROPE_DIM], torch.bfloat16, init_value=init_cos), - TensorSpec("freqs_sin", [T, ROPE_DIM], torch.bfloat16, init_value=init_sin), + TensorSpec("rope_cos_half", [T, HALF_ROPE], torch.float32, init_value=init_rope_cos_half), + TensorSpec("rope_sin_half", [T, HALF_ROPE], torch.float32, init_value=init_rope_sin_half), TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), diff --git a/models/deepseek/v4/decode_sparse_attn_swa.py b/models/deepseek/v4/decode_sparse_attn_swa.py index 4802e139..da56b8ac 100644 --- a/models/deepseek/v4/decode_sparse_attn_swa.py +++ b/models/deepseek/v4/decode_sparse_attn_swa.py @@ -56,8 +56,6 @@ # (its [64,128] softmax and co-resident QK+PV L0C accumulators overflow Vec/L0C). QK_M_TILE = 32 ATTN_K_TILE = 128 -ROPE_TILE = 16 -ROPE_INTERLEAVE_TILE = 2 * ROPE_TILE A_T_TILE = 32 A_K_TILE = 128 A_N_TILE = 128 @@ -112,8 +110,13 @@ def sparse_attn_swa( cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32], cmp_sparse_indices: pl.Tensor[[T, TOPK], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], - freqs_cos: pl.Tensor[[T, ROPE_DIM], pl.BF16], - freqs_sin: pl.Tensor[[T, ROPE_DIM], pl.BF16], + # Half-width unsigned inverse-RoPE tables for the gather-free split-half rotate: + # rope_cos_half = [c0,c1,...,c_{HALF-1}] rope_sin_half = [s0,s1,...] (one per freq) + # The rope segment is split-half ([x_lo | x_hi] = dims [0:HALF | HALF:ROPE_DIM]), + # so the rotation partner of lane k is lane k+HALF -- a contiguous slice, not a + # j^1 swap gather. FP32 so the rope loop reads them straight into its FP32 rotate. + rope_cos_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], + rope_sin_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], wo_b_scale: pl.Tensor[[D], pl.FP32], @@ -293,46 +296,26 @@ def sparse_attn_swa( n_col = n_hh * HEAD_DIM o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + NOPE_DIM] = n_out[n_hi : n_hi + 1, 0 : NOPE_DIM] - # Precompute the head-invariant interleaved cos and sign*sin once: they depend - # only on (token, column), not head, so building them per head would repeat the - # same dup-gather H times on the bottleneck Vec engine. sign is folded into sin - # (multiply by +/-1). The conjugate (inverse) rotation is: - # out[j] = x[j]*cos_il[j] + x[j^1]*sign[j]*sin_il[j] - rope_cos_il = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32) - rope_sin_signed = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32) - for cp in pl.spmd(HALF_ROPE // ROPE_TILE, name_hint="rope_cs"): - cp_r0 = cp * ROPE_TILE - cp_c0 = 2 * cp_r0 - cs_col = pl.col_expand_mul( - pl.full([T, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0), - pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], 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) # j>>1 - cs_lane = pl.sub(cs_col, pl.mul(cs_dup_f, 2.0)) # j%2 - cs_sign = pl.neg(pl.sub(pl.mul(cs_lane, 2.0), 1.0)) # [+1,-1,...] (conjugate) - cs_cos = pl.cast(freqs_cos[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32) - cs_sin = pl.cast(freqs_sin[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32) - rope_cos_il[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx) - rope_sin_signed[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.mul( - pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign) + # Gather-free conjugate (inverse) split-half RoPE. The rope segment is laid out + # split-half ([x_lo | x_hi] = dims [0:HALF_ROPE | HALF_ROPE:ROPE_DIM]), so the + # rotation partner of lane k is lane k+HALF_ROPE -- a contiguous slice, never a + # j^1 swap gather. For k in [0, HALF_ROPE): + # out_lo[k] = x_lo[k]*cos[k] + x_hi[k]*sin[k] + # out_hi[k] = x_hi[k]*cos[k] - x_lo[k]*sin[k] + # cos/sin are half-width (one value per frequency); no in-kernel index build. # Inverse RoPE fused with the rope-column pack: each task rotates its heads' - # rope segments and stores them straight into o_packed's strided rope columns, - # dropping a separate rope_pack stage and GM round-trip. cos_il / sign*sin come - # from the pre-pass above; only swap_idx (j^1) is rebuilt per task. + # rope segments and stores them straight into o_packed's rope columns, dropping + # a separate rope_pack stage and GM round-trip. The whole ROPE_DIM is rotated in + # one pass as a [ROPE_OUT_TOK_TILE, HALF_ROPE] lo tile + matching hi tile. attn_rope_stage_3d = pl.reshape(attn_rope_stage, [T, H, ROPE_DIM]) for rp_idx in pl.spmd((H // 4) * 2, name_hint="rope"): rp_hg = rp_idx // 2 rp_tt = rp_idx - rp_hg * 2 rp_t0 = rp_tt * ROPE_OUT_TOK_TILE - # Head-invariant swap index (j^1), built once and reused across the head - # group -- the only per-head input is this head's strided rope slice. - sp_col = pl.col_expand_mul( - pl.full([ROPE_OUT_TOK_TILE, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0), - pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], dtype=pl.INT32), target_type=pl.FP32)) - sp_dup_f = pl.cast(pl.cast(pl.mul(sp_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32) - sp_lane = pl.sub(sp_col, pl.mul(sp_dup_f, 2.0)) # j%2 - sp_swap_idx = pl.cast(pl.sub(pl.add(sp_col, 1.0), pl.mul(sp_lane, 2.0)), target_type=pl.INT32) # j^1 + # Head-invariant half-width tables, hoisted once per task. + r_cos = rope_cos_half[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] + r_sin = rope_sin_half[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] for rp_hl in pl.range(0, 4): rp_gh = rp_hg * 4 + rp_hl @@ -340,21 +323,19 @@ def sparse_attn_swa( rp_hh = rp_gh - rp_g * HEADS_PER_GROUP rp_col = rp_hh * HEAD_DIM + NOPE_DIM rp_o0 = rp_g * T + rp_t0 - for r_r0 in pl.range(0, HALF_ROPE, ROPE_TILE): - c0 = 2 * r_r0 - # This head's rope rows for this token tile (stride H); gather needs - # FP/INT, so cast the strided BF16 slice to FP32 first. - r_tile_fp32 = pl.cast( - pl.reshape(attn_rope_stage_3d[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, rp_gh : rp_gh + 1, c0 : c0 + ROPE_INTERLEAVE_TILE], [ROPE_OUT_TOK_TILE, ROPE_INTERLEAVE_TILE]), - target_type=pl.FP32) - r_cos_il = rope_cos_il[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, c0 : c0 + ROPE_INTERLEAVE_TILE] - r_sin_signed = rope_sin_signed[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, c0 : c0 + ROPE_INTERLEAVE_TILE] - r_swapped = pl.gather(r_tile_fp32, dim=-1, index=sp_swap_idx) - r_rot = pl.add(pl.mul(r_tile_fp32, r_cos_il), pl.mul(r_swapped, r_sin_signed)) - # Store BF16-rounded rotated values straight into o_packed's rope - # columns for this head (golden also rounds inverse-RoPE to bf16). - r_rot = pl.cast(r_rot, target_type=pl.BF16, mode="rint") - o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col + c0 : rp_col + c0 + ROPE_INTERLEAVE_TILE] = r_rot + # This head's full ROPE_DIM rope rows for this token tile (stride H); the + # rotate math is FP32, so cast the strided BF16 slice to FP32 first. + r_tile = pl.cast( + pl.reshape(attn_rope_stage_3d[rp_t0 : rp_t0 + ROPE_OUT_TOK_TILE, rp_gh : rp_gh + 1, 0 : ROPE_DIM], [ROPE_OUT_TOK_TILE, ROPE_DIM]), + target_type=pl.FP32) + r_lo = r_tile[0 : ROPE_OUT_TOK_TILE, 0 : HALF_ROPE] # x_lo (contiguous, no gather) + r_hi = r_tile[0 : ROPE_OUT_TOK_TILE, HALF_ROPE : ROPE_DIM] # x_hi (contiguous, no gather) + out_lo = pl.add(pl.mul(r_lo, r_cos), pl.mul(r_hi, r_sin)) # x_lo*c + x_hi*s + out_hi = pl.sub(pl.mul(r_hi, r_cos), pl.mul(r_lo, r_sin)) # x_hi*c - x_lo*s + # Two contiguous BF16-rounded stores into o_packed's rope columns + # (golden also rounds inverse-RoPE to bf16). + o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col : rp_col + HALF_ROPE] = pl.cast(out_lo, target_type=pl.BF16, mode="rint") + o_packed[rp_o0 : rp_o0 + ROPE_OUT_TOK_TILE, rp_col + HALF_ROPE : rp_col + ROPE_DIM] = pl.cast(out_hi, target_type=pl.BF16, mode="rint") # Grouped BF16 projection `o_packed @ wo_a^T` -> `o_r`. Vec post-process # (BF16 store + per-row amax) is T-tiled to keep the fused AIV side from @@ -448,8 +429,8 @@ def sparse_attn_test( cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32], cmp_sparse_indices: pl.Tensor[[T, TOPK], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], - freqs_cos: pl.Tensor[[T, ROPE_DIM], pl.BF16], - freqs_sin: pl.Tensor[[T, ROPE_DIM], pl.BF16], + rope_cos_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], + rope_sin_half: pl.Tensor[[T, HALF_ROPE], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], wo_b_scale: pl.Tensor[[D], pl.FP32], @@ -464,8 +445,8 @@ def sparse_attn_test( cmp_block_table, cmp_sparse_indices, attn_sink, - freqs_cos, - freqs_sin, + rope_cos_half, + rope_sin_half, wo_a, wo_b, wo_b_scale, @@ -510,8 +491,10 @@ def golden_sparse_attn(tensors): cmp_block_table = tensors["cmp_block_table"] cmp_sparse_indices = tensors["cmp_sparse_indices"] attn_sink = tensors["attn_sink"].float() - cos = tensors["freqs_cos"].float() - sin = tensors["freqs_sin"].float() + # Half-width unsigned inverse-RoPE tables (one value per frequency), read + # directly by the split-half reference rotation below. + cos = tensors["rope_cos_half"].float() + sin = tensors["rope_sin_half"].float() wo_a = tensors["wo_a"].float() wo_b_i8 = tensors["wo_b"] wo_b_scale = tensors["wo_b_scale"].float() @@ -590,14 +573,15 @@ def golden_sparse_attn(tensors): denom = li + torch.exp(attn_sink.unsqueeze(-1) - score_max) o[t] = oi_num / denom - rope_pair = o[..., NOPE_DIM:].unflatten(-1, (-1, 2)) - rope_even = rope_pair[..., 0] - rope_odd = rope_pair[..., 1] - cos_half = cos[:, :HALF_ROPE].unsqueeze(1) - sin_half = sin[:, :HALF_ROPE].unsqueeze(1) - inv_even = (rope_even * cos_half + rope_odd * sin_half).to(torch.bfloat16).float() - inv_odd = (rope_odd * cos_half - rope_even * sin_half).to(torch.bfloat16).float() - o_rope = torch.stack([inv_even, inv_odd], dim=-1).flatten(-2) + # Split-half inverse RoPE: lo = dims [:HALF_ROPE], hi = dims [HALF_ROPE:]. + rope_seg = o[..., NOPE_DIM:] + x_lo = rope_seg[..., :HALF_ROPE] + x_hi = rope_seg[..., HALF_ROPE:] + cos_h = cos.unsqueeze(1) # [T, 1, HALF_ROPE] broadcast over heads + sin_h = sin.unsqueeze(1) + inv_lo = (x_lo * cos_h + x_hi * sin_h).to(torch.bfloat16).float() + inv_hi = (x_hi * cos_h - x_lo * sin_h).to(torch.bfloat16).float() + o_rope = torch.cat([inv_lo, inv_hi], dim=-1) o = torch.cat([o[..., :NOPE_DIM], o_rope], dim=-1).to(torch.bfloat16) seq_per_batch = T // B @@ -707,17 +691,19 @@ def init_cmp_sparse_indices(): indices[0, WIN - 1] = WIN - 1 return indices - def init_cos(): - """Build the split-half cosine table used by the inverse-RoPE reference.""" - angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 - cos_half = torch.cos(angles) - return torch.cat([cos_half, cos_half], dim=-1) + # Half-width unsigned cos/sin (one value per frequency), rounded to BF16 then back to + # FP32 so the FP32 kernel input holds exactly the bf16 values (bit-exact with the reference). + angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 + rope_cos_half_tbl = torch.cos(angles).to(torch.bfloat16).float() + rope_sin_half_tbl = torch.sin(angles).to(torch.bfloat16).float() - def init_sin(): - """Build the split-half sine table used by the inverse-RoPE reference.""" - angles = torch.arange(T * HALF_ROPE).reshape(T, HALF_ROPE) * 1e-3 - sin_half = torch.sin(angles) - return torch.cat([sin_half, sin_half], dim=-1) + def init_rope_cos_half(): + """Build the half-width cosine table [c0,c1,...] read directly by the split-half inverse-RoPE kernel.""" + return rope_cos_half_tbl.clone() + + def init_rope_sin_half(): + """Build the half-width sine table [s0,s1,...] read directly by the split-half inverse-RoPE kernel.""" + return rope_sin_half_tbl.clone() def init_wo_a(): """Initialize the grouped first-stage output-projection weights.""" @@ -743,8 +729,8 @@ def init_wo_b_scale(): TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("cmp_sparse_indices", [T, TOPK], torch.int32, init_value=init_cmp_sparse_indices), TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), - TensorSpec("freqs_cos", [T, ROPE_DIM], torch.bfloat16, init_value=init_cos), - TensorSpec("freqs_sin", [T, ROPE_DIM], torch.bfloat16, init_value=init_sin), + TensorSpec("rope_cos_half", [T, HALF_ROPE], torch.float32, init_value=init_rope_cos_half), + TensorSpec("rope_sin_half", [T, HALF_ROPE], torch.float32, init_value=init_rope_sin_half), TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), diff --git a/models/deepseek/v4/qkv_proj_rope.py b/models/deepseek/v4/qkv_proj_rope.py index b9aa7a16..6120a7a2 100644 --- a/models/deepseek/v4/qkv_proj_rope.py +++ b/models/deepseek/v4/qkv_proj_rope.py @@ -226,33 +226,21 @@ def qkv_proj_rope( q_normed, target_type=pl.BF16, mode="rint" ) - # Per-head RoPE (CANN A3 rotate_interleaved): stay on the interleaved layout and - # rotate via an i^1 swap gather + sign mask, dropping the de-interleave gather + - # re-interleave scatter. The rotation indices/sign and the interleave-duplicated - # cos/sin are built ENTIRELY IN-KERNEL (no host inputs): swap_idx (j^1), sign - # ([-1,+1,...]) and dup_idx (j>>1) come from pl.arange per task, and cos_il/sin_il - # are dup-gathered from rope_cos/rope_sin in-task. The prior in-task index-tile - # tail clobber (-> tgather UB-OOB -> 507018 hang) that once forced these to be - # kernel inputs is resolved. inv_rms is per-row so it factors out of the rotation - # and is applied after; the writeback into q_flat is folded in (no FP32 GM stage). - # swapped = gather(x, j^1) = [x1,x0,x3,x2,...]; sign = [-1,+1,...] - # out = inv_rms * (x*cos_il + swapped*sign*sin_il) + # Per-head split-half (NeoX) RoPE: the rope segment is laid out [x_lo | x_hi] = + # dims [0:ROPE_HALF | ROPE_HALF:ROPE_DIM], so the rotation partner of lane k is + # lane k+ROPE_HALF -- a contiguous slice, not a j^1 swap gather. No in-kernel + # index build and no dup-gather: the half-width cos/sin are read straight from the + # first ROPE_HALF columns of the (duplicated) tables. inv_rms is per-row so it + # factors out of the rotation and is applied after; the q_flat writeback is folded + # in (no FP32 GM stage). + # out_lo = x_lo*cos - x_hi*sin ; out_hi = x_lo*sin + x_hi*cos for hg_idx in pl.spmd(H // 2, name_hint="q_head_rope_fused"): 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,...] - # tg-outer / head-inner: cos_il/sin_il are head-independent, so dup-gather once - # per tg and reuse for both heads. + # tg-outer / head-inner: cos/sin are head-independent, so read once per tg. 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 = pl.cast(rope_cos_view[tg : tg + Q_ROPE_T_TILE, 0 : ROPE_HALF], target_type=pl.FP32) + q_sin = pl.cast(rope_sin_view[tg : tg + Q_ROPE_T_TILE, 0 : ROPE_HALF], target_type=pl.FP32) for h_inner in pl.range(2): h = hg + h_inner h0 = h * HEAD_DIM @@ -260,10 +248,15 @@ def qkv_proj_rope( q_head_inv_rms_all[h : h + 1, tg : tg + Q_ROPE_T_TILE], [Q_ROPE_T_TILE, 1] ) 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_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_rope_inv_rms_chunk), target_type=pl.BF16, mode="rint" + q_lo = q_rope_chunk[0 : Q_ROPE_T_TILE, 0 : ROPE_HALF] # x_lo (contiguous, no gather) + q_hi = q_rope_chunk[0 : Q_ROPE_T_TILE, ROPE_HALF : ROPE_DIM] # x_hi (contiguous, no gather) + q_out_lo = pl.sub(pl.mul(q_lo, q_cos), pl.mul(q_hi, q_sin)) # x_lo*c - x_hi*s + q_out_hi = pl.add(pl.mul(q_lo, q_sin), pl.mul(q_hi, q_cos)) # x_lo*s + x_hi*c + q_flat[tg : tg + Q_ROPE_T_TILE, h0 + NOPE_DIM : h0 + NOPE_DIM + ROPE_HALF] = pl.cast( + pl.row_expand_mul(q_out_lo, q_rope_inv_rms_chunk), target_type=pl.BF16, mode="rint" + ) + q_flat[tg : tg + Q_ROPE_T_TILE, h0 + NOPE_DIM + ROPE_HALF : h0 + NOPE_DIM + ROPE_DIM] = pl.cast( + pl.row_expand_mul(q_out_hi, q_rope_inv_rms_chunk), target_type=pl.BF16, mode="rint" ) kv_fp32 = pl.create_tensor([t_dim, HEAD_DIM], dtype=pl.FP32) @@ -311,35 +304,35 @@ def qkv_proj_rope( # task's chunk keeps Vec UB well under the 192 KB cap. for tg_idx in pl.spmd(t_dim // KV_ROPE_T_TILE, name_hint="kv_rope_fused"): tg = tg_idx * KV_ROPE_T_TILE + # Split-half (NeoX) RoPE, gather-free. gamma is per-column so it does NOT + # commute with the rotation -- fold gamma_lo onto x_lo and gamma_hi onto x_hi + # (matching the golden rms_norm-then-rope, which scales every column by gamma + # before rotating). inv_rms is per-row and commutes; fold it in too. + # out_lo = x_lo*cos - x_hi*sin ; out_hi = x_lo*sin + x_hi*cos gamma_rope = pl.reshape( pl.cast(gamma_ckv[NOPE_DIM : NOPE_DIM + ROPE_DIM], target_type=pl.FP32), [1, ROPE_DIM], ) + gamma_lo = gamma_rope[0:1, 0 : ROPE_HALF] + gamma_hi = gamma_rope[0:1, ROPE_HALF : ROPE_DIM] kv_rope_inv_rms_chunk = pl.reshape( kv_inv_rms_tensor[0:1, tg : tg + KV_ROPE_T_TILE], [KV_ROPE_T_TILE, 1] ) kv_rope_chunk = kv_fp32[tg : tg + KV_ROPE_T_TILE, NOPE_DIM : NOPE_DIM + ROPE_DIM] - # A3 interleaved swap-gather (same form as q_head_rope_fused), built in-kernel. - # gamma is folded into kv_rope_norm_chunk BEFORE the swap so the swapped lane - # n[j^1] correctly carries gamma[j^1] (gamma is per-column, does NOT commute); - # inv_rms is per-row so it commutes. out[j] = n[j]*cos_il[j] + n[j^1]*sign[j]*sin_il[j]. - kv_rope_norm_chunk = pl.col_expand_mul(pl.row_expand_mul(kv_rope_chunk, kv_rope_inv_rms_chunk), gamma_rope) - kv_ones = pl.full([KV_ROPE_T_TILE, ROPE_DIM], dtype=pl.FP32, value=1.0) - kv_col = pl.col_expand_mul(kv_ones, pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32)) - kv_dup_f = pl.cast(pl.cast(pl.mul(kv_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32) - kv_dup_idx = pl.cast(kv_dup_f, target_type=pl.INT32) # j>>1 - kv_lane = pl.sub(kv_col, pl.mul(kv_dup_f, 2.0)) # j%2 - kv_swap_idx = pl.cast(pl.sub(pl.add(kv_col, 1.0), pl.mul(kv_lane, 2.0)), target_type=pl.INT32) # j^1 - kv_sign = pl.sub(pl.mul(kv_lane, 2.0), 1.0) # [-1,+1,...] - kv_cos_il = pl.gather(pl.cast(rope_cos_view[tg : tg + KV_ROPE_T_TILE, :], target_type=pl.FP32), dim=-1, index=kv_dup_idx) - kv_sin_il = pl.gather(pl.cast(rope_sin_view[tg : tg + KV_ROPE_T_TILE, :], target_type=pl.FP32), dim=-1, index=kv_dup_idx) - kv_swapped = pl.gather(kv_rope_norm_chunk, dim=-1, index=kv_swap_idx) - kv_rope_rot = pl.add( - pl.mul(kv_rope_norm_chunk, kv_cos_il), - pl.mul(pl.mul(kv_swapped, kv_sign), kv_sin_il), + kv_lo = kv_rope_chunk[0 : KV_ROPE_T_TILE, 0 : ROPE_HALF] + kv_hi = kv_rope_chunk[0 : KV_ROPE_T_TILE, ROPE_HALF : ROPE_DIM] + # Fold inv_rms (per-row) + gamma (per-column) into each half BEFORE the rotation. + kv_lo_n = pl.col_expand_mul(pl.row_expand_mul(kv_lo, kv_rope_inv_rms_chunk), gamma_lo) + kv_hi_n = pl.col_expand_mul(pl.row_expand_mul(kv_hi, kv_rope_inv_rms_chunk), gamma_hi) + kv_cos = pl.cast(rope_cos_view[tg : tg + KV_ROPE_T_TILE, 0 : ROPE_HALF], target_type=pl.FP32) + kv_sin = pl.cast(rope_sin_view[tg : tg + KV_ROPE_T_TILE, 0 : ROPE_HALF], target_type=pl.FP32) + kv_out_lo = pl.sub(pl.mul(kv_lo_n, kv_cos), pl.mul(kv_hi_n, kv_sin)) # x_lo*c - x_hi*s + kv_out_hi = pl.add(pl.mul(kv_lo_n, kv_sin), pl.mul(kv_hi_n, kv_cos)) # x_lo*s + x_hi*c + kv_view[tg : tg + KV_ROPE_T_TILE, NOPE_DIM : NOPE_DIM + ROPE_HALF] = pl.cast( + kv_out_lo, target_type=pl.BF16, mode="rint" ) - kv_view[tg : tg + KV_ROPE_T_TILE, NOPE_DIM : NOPE_DIM + ROPE_DIM] = pl.cast( - kv_rope_rot, target_type=pl.BF16, mode="rint" + kv_view[tg : tg + KV_ROPE_T_TILE, NOPE_DIM + ROPE_HALF : NOPE_DIM + ROPE_DIM] = pl.cast( + kv_out_hi, target_type=pl.BF16, mode="rint" ) return q @@ -421,17 +414,17 @@ def matmul_bf16_input_fp32(a, b): return torch.matmul(a_fp32, b_fp32).float() def apply_rope(x_rope, cos, sin): - # x_rope: [T, ..., ROPE_DIM] with interleaved even/odd rotary pairs. - x_pair = x_rope.unflatten(-1, (-1, 2)) - x_even, x_odd = x_pair[..., 0], x_pair[..., 1] + # x_rope: [T, ..., ROPE_DIM] split-half layout (lo = [:ROPE_HALF], hi = [ROPE_HALF:]). + x_lo = x_rope[..., :ROPE_HALF] + x_hi = x_rope[..., ROPE_HALF:] cos_v = cos[..., :ROPE_HALF] sin_v = sin[..., :ROPE_HALF] - while cos_v.ndim < x_even.ndim: + while cos_v.ndim < x_lo.ndim: cos_v = cos_v.unsqueeze(-2) sin_v = sin_v.unsqueeze(-2) - y_even = (x_even * cos_v - x_odd * sin_v).to(torch.bfloat16) - y_odd = (x_even * sin_v + x_odd * cos_v).to(torch.bfloat16) - return torch.stack([y_even, y_odd], dim=-1).flatten(-2) + y_lo = (x_lo * cos_v - x_hi * sin_v).to(torch.bfloat16) + y_hi = (x_lo * sin_v + x_hi * cos_v).to(torch.bfloat16) + return torch.cat([y_lo, y_hi], dim=-1) t_dim = x.shape[0] token_x = x.view(t_dim, D)