Skip to content

Refactor: scope dsv4 decode+prefill kernels under auto_scope=False #585

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
zhangqi-chen wants to merge 12 commits into
base: main
Choose a base branch
from
Open

Refactor: scope dsv4 decode+prefill kernels under auto_scope=False #585

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
0a1a939
Refactor: scope decode_sparse_attn into 3 phases for memory reuse
zhangqi-chen Jun 23, 2026
27dab22
Refactor: scope decode_attention_csa to cut decode peak memory
zhangqi-chen Jun 23, 2026
f1df1ca
Refactor: scope moe + expert_routed for auto_scope=False composition
zhangqi-chen Jun 23, 2026
0ee2a6f
Refactor: scope SWA + HCA attention for auto_scope=False composition
zhangqi-chen Jun 23, 2026
0489c16
Refactor: place moe sh/ffn_out create at first use
zhangqi-chen Jun 23, 2026
aac2d36
Docs: add 'declare scratch at first use' convention (sec 8)
zhangqi-chen Jun 23, 2026
ce5af48
Refactor: scope decode_layer attention + MoE under auto_scope=False
zhangqi-chen Jun 23, 2026
51754ed
Refactor: scope prefill attention/layer/sparse_attn under auto_scope=…
zhangqi-chen Jun 23, 2026
a5bb2d7
Refactor: declare decode_fwd slices/scratch at first use
zhangqi-chen Jun 23, 2026
eeb1bd2
Refactor: scope decode_fwd per-layer attention + MoE (auto_scope=False)
zhangqi-chen Jun 23, 2026
a974cba
Fix: wire decode_attention_hca --scope-stats; tighten scratch scopes
zhangqi-chen Jun 23, 2026
5e10626
Fix: drop hidden= rebind in decode_fwd MoE scopes
zhangqi-chen Jun 23, 2026
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
23 changes: 23 additions & 0 deletions docs/pypto-coding-style.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -701,3 +701,26 @@ x_i8 = pl.cast(pl.mul(x, inv_scale), pl.INT8, mode="rint")
# pad to 32: ptoas rejects cube tiles whose cols aren't a multiple of 16
w_pad = pl.slice(w, [K, 32], [0, 0], valid_shape=[K, MIX_HC])
```

### Declare scratch at first use

Declare each `pl.create_tensor` / `pl.slice` / `pl.reshape` immediately before
the first `pl.at` / `pl.spmd` / `pl.range` / `pl.scope` (or sub-kernel call) that
uses it — not in a block at the top of the function. Tight placement keeps a
tensor's live range minimal, which lets the scope / memory-reuse passes free it
as early as possible. **Exception:** a tensor written *inside* a `pl.scope()` but
read after it must be declared **before** that scope (the scope would otherwise
free it at exit), so such scope-bridging tensors sit just before the enclosing
scope. See `pl.scope` (§1) for how scoping bounds a tensor's lifetime.

```python
# ✅ each scratch declared just before the region that first writes it
with pl.scope():
x_mixed = pl.create_tensor([T, D], dtype=pl.BF16)
hc_pre(..., x_mixed) # x_mixed dies here, freed at scope exit
x_normed = pl.create_tensor([T, D], dtype=pl.BF16)
rms_norm(x_mixed, ..., x_normed)
# attn_out is first written by sparse_attn below (body level), so declare it here
attn_out = pl.create_tensor([T, D], dtype=pl.BF16)
sparse_attn(..., attn_out)
```
370 changes: 150 additions & 220 deletions models/deepseek/v4/decode_attention_csa.py
View file
Open in desktop

Large diffs are not rendered by default.

228 changes: 97 additions & 131 deletions models/deepseek/v4/decode_attention_hca.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
# Copyright (c) PyPTO Contributors.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
Expand Down Expand Up @@ -82,7 +82,7 @@
SPARSE_ROPE_INTERLEAVE_TILE = 2 * SPARSE_ROPE_TILE


@pl.jit.inline
@pl.jit.inline(auto_scope=False)
def attention_hca(
x_hc: pl.Tensor[[T, HC_MULT, D], pl.BF16],
# hc_pre weights
Expand Down Expand Up @@ -126,144 +126,114 @@
x_out: pl.Tensor[[T, HC_MULT, D], pl.BF16],
):
"""HCA decode orchestration for compress_ratio=128."""
x_mixed = pl.create_tensor([T, D], dtype=pl.BF16)
# Frame-level: written inside the pre-attention scope below but read by
# sparse_attn / writeback / hc_post afterwards.
post_t = pl.create_tensor([T, HC_MULT], dtype=pl.FP32)
comb_t = pl.create_tensor([T, HC_MULT * HC_MULT], dtype=pl.FP32)
x_mixed = hc_pre(
x_hc,
hc_attn_fn,
hc_attn_scale,
hc_attn_base,
x_mixed,
post_t,
comb_t,
)

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)
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"):
for b in pl.range(B):
first_t = b * S
first_pos_b = pl.read(position_ids, [first_t])
cmp_offset_b = COMPRESS_RATIO - (first_pos_b % COMPRESS_RATIO)
cmp_pos_b = pl.cast(first_pos_b + cmp_offset_b - COMPRESS_RATIO, pl.INDEX)
cmp_cos_row = freqs_cos[cmp_pos_b : cmp_pos_b + 1, 0 : ROPE_HEAD_DIM // 2]
cmp_sin_row = freqs_sin[cmp_pos_b : cmp_pos_b + 1, 0 : ROPE_HEAD_DIM // 2]
cmp_cos[b : b + 1, 0 : ROPE_HEAD_DIM // 2] = pl.cast(cmp_cos_row, target_type=pl.FP32)
cmp_sin[b : b + 1, 0 : ROPE_HEAD_DIM // 2] = pl.cast(cmp_sin_row, target_type=pl.FP32)
for s in pl.range(S):
t = b * S + s
pos_b = pl.cast(pl.read(position_ids, [t]), pl.INDEX)
step_cos_row = pl.cast(freqs_cos[pos_b : pos_b + 1, 0 : ROPE_HEAD_DIM], target_type=pl.FP32)
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")

q = pl.create_tensor([T, H, HEAD_DIM], dtype=pl.BF16)
kv = pl.create_tensor([T, HEAD_DIM], dtype=pl.BF16)
qr = pl.create_tensor([T, Q_LORA], dtype=pl.INT8) # unused on HCA path
qr_scale = pl.create_tensor([T, 1], dtype=pl.FP32)
x_normed = pl.create_tensor([T, D], dtype=pl.BF16)
x_normed = rms_norm(x_mixed, attn_norm_w, x_normed)
q = qkv_proj_rope(
x_normed,
wq_a,
wq_b,
wq_b_scale,
wkv,
rope_cos_t,
rope_sin_t,
gamma_cq,
gamma_ckv,
q,
kv,
qr,
qr_scale,
)

x_normed_bsd = pl.reshape(x_normed, [B, S, D])
cmp_kv_proj = pl.create_tensor([B, S, HEAD_DIM], dtype=pl.FP32)
position_ids_bsd = pl.reshape(position_ids, [B, S])
cmp_slot_mapping_bsd = pl.reshape(cmp_slot_mapping, [B, S])
state_slot_mapping_bsd = pl.reshape(state_slot_mapping, [B, S])
cmp_kv_proj = compressor_ratio128(
x_normed_bsd,
cmp_kv_proj,
compress_state,
compress_state_block_table,
cmp_wkv,
cmp_wgate,
cmp_ape,
cmp_norm_w,
cmp_cos,
cmp_sin,
cmp_kv,
position_ids_bsd,
cmp_slot_mapping_bsd,
state_slot_mapping_bsd,
)

attn_out = pl.create_tensor([T, D], dtype=pl.BF16)
topk_all = pl.create_tensor([T, HCA_SPARSE_TOPK], dtype=pl.INT32)
with pl.at(level=pl.Level.CORE_GROUP, name_hint="hca_overlay_topk"):
for topk_b in pl.range(B):
for topk_s in pl.range(S):
topk_t = topk_b * S + topk_s
topk_abs_pos = pl.read(position_ids, [topk_t])

if topk_abs_pos >= WIN - 1:
topk_win_start = (topk_abs_pos % WIN) + 1
for topk_k in pl.range(WIN):
topk_val = (topk_win_start + topk_k) % WIN
topk_out = topk_val
for topk_os in pl.range(S):
if topk_os <= topk_s:
topk_overlay_t = topk_b * S + topk_os
topk_overlay_pos = pl.read(position_ids, [topk_overlay_t])
if topk_val == topk_overlay_pos % WIN:
topk_out = WIN + topk_os
pl.write(topk_all, [topk_t, topk_k], pl.cast(topk_out, pl.INT32))
else:
for topk_k in pl.range(WIN):
if topk_k <= topk_abs_pos:
topk_out = topk_k

# Pre-attention scope (hc_pre .. overlay-topk): scratch dead by sparse_attn,
# freed at scope exit. Sub-kernels write out-params in place (bare calls).
with pl.scope():
x_mixed = pl.create_tensor([T, D], dtype=pl.BF16)
hc_pre(
x_hc, hc_attn_fn, hc_attn_scale, hc_attn_base,
x_mixed, post_t, comb_t,
)

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"):
for b in pl.range(B):
first_t = b * S
first_pos_b = pl.read(position_ids, [first_t])
cmp_offset_b = COMPRESS_RATIO - (first_pos_b % COMPRESS_RATIO)
cmp_pos_b = pl.cast(first_pos_b + cmp_offset_b - COMPRESS_RATIO, pl.INDEX)
cmp_cos_row = freqs_cos[cmp_pos_b : cmp_pos_b + 1, 0 : ROPE_HEAD_DIM // 2]
cmp_sin_row = freqs_sin[cmp_pos_b : cmp_pos_b + 1, 0 : ROPE_HEAD_DIM // 2]
cmp_cos[b : b + 1, 0 : ROPE_HEAD_DIM // 2] = pl.cast(cmp_cos_row, target_type=pl.FP32)
cmp_sin[b : b + 1, 0 : ROPE_HEAD_DIM // 2] = pl.cast(cmp_sin_row, target_type=pl.FP32)
for s in pl.range(S):
t = b * S + s
pos_b = pl.cast(pl.read(position_ids, [t]), pl.INDEX)
step_cos_row = pl.cast(freqs_cos[pos_b : pos_b + 1, 0 : ROPE_HEAD_DIM], target_type=pl.FP32)
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")

x_normed = pl.create_tensor([T, D], dtype=pl.BF16)
rms_norm(x_mixed, attn_norm_w, x_normed)
qr = pl.create_tensor([T, Q_LORA], dtype=pl.INT8) # unused on HCA path
qr_scale = pl.create_tensor([T, 1], dtype=pl.FP32)
qkv_proj_rope(
x_normed,
wq_a, wq_b, wq_b_scale, wkv,
rope_cos_t, rope_sin_t, gamma_cq, gamma_ckv,
q, kv, qr, qr_scale,
)

cmp_kv_proj = pl.create_tensor([B, S, HEAD_DIM], dtype=pl.FP32)
x_normed_bsd = pl.reshape(x_normed, [B, S, D])
position_ids_bsd = pl.reshape(position_ids, [B, S])
cmp_slot_mapping_bsd = pl.reshape(cmp_slot_mapping, [B, S])
state_slot_mapping_bsd = pl.reshape(state_slot_mapping, [B, S])
compressor_ratio128(
x_normed_bsd, cmp_kv_proj, compress_state, compress_state_block_table,
cmp_wkv, cmp_wgate, cmp_ape, cmp_norm_w, cmp_cos, cmp_sin, cmp_kv,
position_ids_bsd, cmp_slot_mapping_bsd, state_slot_mapping_bsd,
)

with pl.at(level=pl.Level.CORE_GROUP, name_hint="hca_overlay_topk"):
for topk_b in pl.range(B):
for topk_s in pl.range(S):
topk_t = topk_b * S + topk_s
topk_abs_pos = pl.read(position_ids, [topk_t])

if topk_abs_pos >= WIN - 1:
topk_win_start = (topk_abs_pos % WIN) + 1
for topk_k in pl.range(WIN):
topk_val = (topk_win_start + topk_k) % WIN
topk_out = topk_val
for topk_os in pl.range(S):
if topk_os <= topk_s:
topk_overlay_t = topk_b * S + topk_os
topk_overlay_pos = pl.read(position_ids, [topk_overlay_t])
if topk_k == topk_overlay_pos % WIN:
if topk_val == topk_overlay_pos % WIN:
topk_out = WIN + topk_os
pl.write(topk_all, [topk_t, topk_k], pl.cast(topk_out, pl.INT32))
else:
pl.write(topk_all, [topk_t, topk_k], pl.cast(-1, pl.INT32))

topk_cmp_valid = pl.min(
HCA_TOPK_LIMIT,
pl.min((topk_abs_pos + 1) // COMPRESS_RATIO, pl.read(kv_seq_lens, [topk_b]) // COMPRESS_RATIO),
)
for topk_ck in pl.range(HCA_SPARSE_TOPK - WIN):
if topk_ck < topk_cmp_valid:
pl.write(topk_all, [topk_t, WIN + topk_ck], pl.cast(WIN + S + topk_ck, pl.INT32))
else:
pl.write(topk_all, [topk_t, WIN + topk_ck], pl.cast(-1, pl.INT32))
for topk_k in pl.range(WIN):
if topk_k <= topk_abs_pos:
topk_out = topk_k
for topk_os in pl.range(S):
if topk_os <= topk_s:
topk_overlay_t = topk_b * S + topk_os
topk_overlay_pos = pl.read(position_ids, [topk_overlay_t])
if topk_k == topk_overlay_pos % WIN:
topk_out = WIN + topk_os
pl.write(topk_all, [topk_t, topk_k], pl.cast(topk_out, pl.INT32))
else:
pl.write(topk_all, [topk_t, topk_k], pl.cast(-1, pl.INT32))

topk_cmp_valid = pl.min(
HCA_TOPK_LIMIT,
pl.min((topk_abs_pos + 1) // COMPRESS_RATIO, pl.read(kv_seq_lens, [topk_b]) // COMPRESS_RATIO),
)
for topk_ck in pl.range(HCA_SPARSE_TOPK - WIN):
if topk_ck < topk_cmp_valid:
pl.write(topk_all, [topk_t, WIN + topk_ck], pl.cast(WIN + S + topk_ck, pl.INT32))
else:
pl.write(topk_all, [topk_t, WIN + topk_ck], pl.cast(-1, pl.INT32))

attn_out = pl.create_tensor([T, D], dtype=pl.BF16)
sparse_attn_hca(
q,
kv_cache,
ori_block_table,
kv,
cmp_kv,
cmp_block_table,
topk_all,
attn_sink,
rope_cos_t,
rope_sin_t,
wo_a,
wo_b,
wo_b_scale,
attn_out,
q, kv_cache, ori_block_table, kv, cmp_kv, cmp_block_table,
topk_all, attn_sink, rope_cos_t, rope_sin_t,
wo_a, wo_b, wo_b_scale, attn_out,
)

# Commit new tokens to the cache AFTER sparse_attn reads the pre-update
Expand All @@ -279,17 +249,11 @@
if write_row >= 0:
kv_cache_flat[write_row : write_row + 1, 0 : HEAD_DIM] = kv[write_t : write_t + 1, 0 : HEAD_DIM]

x_out = hc_post(
attn_out,
x_hc,
post_t,
comb_t,
x_out,
)
hc_post(attn_out, x_hc, post_t, comb_t, x_out)
return x_out


@pl.jit
@pl.jit(auto_scope=False)
def attention_hca_test(
x_hc: pl.Tensor[[T, HC_MULT, D], pl.BF16],
hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32],
Expand Down Expand Up @@ -731,6 +695,7 @@
help="Fixture-only compatibility seed for position_ids and slot mappings; "
"otherwise use the default per-batch coverage pattern.")
parser.add_argument("--enable-l2-swimlane", action="store_true", default=False)
parser.add_argument("--scope-stats", action="store_true", default=False)

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

The --scope-stats CLI argument is defined here, but enable_scope_stats=args.scope_stats is not passed to the run_jit call in this file (unlike in decode_attention_csa.py and decode_attention_swa.py). This makes the CLI flag non-functional for this runner.

@coderabbitai coderabbitai Bot Jun 23, 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.

🎯 Functional Correctness | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify decode_attention_hca forwards --scope-stats into run_jit runtime_cfg.
# Expected: a nearby runtime_cfg entry like `enable_scope_stats=args.scope_stats`.
rg -n -C4 'scope_stats|enable_scope_stats|run_jit|runtime_cfg' models/deepseek/v4/decode_attention_hca.py

Repository: hw-native-sys/pypto-lib

Length of output: 1041

🏁 Script executed:

sed -n '705,730p' models/deepseek/v4/decode_attention_hca.py

Repository: hw-native-sys/pypto-lib

Length of output: 838

Add enable_scope_stats=args.scope_stats to the runtime_cfg dict.

The --scope-stats argument was added to the parser (line 698) but is not wired into runtime_cfg (lines 705–709). Without this entry, the flag has no effect. Add it alongside the existing enable_l2_swimlane entry to complete the wiring.

🤖 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/decode_attention_hca.py` at line 698, The --scope-stats
argument was added to the argument parser but is not connected to the
runtime_cfg dictionary. Locate the runtime_cfg dictionary initialization (around
lines 705-709) and add a new entry enable_scope_stats=args.scope_stats alongside
the existing enable_l2_swimlane entry to wire the command-line argument into the
configuration that will be used at runtime.

args = parser.parse_args()

result = run_jit(
Expand All @@ -741,6 +706,7 @@
platform=args.platform,
device_id=args.device,
enable_l2_swimlane=args.enable_l2_swimlane,
enable_scope_stats=args.scope_stats,
),
atol=1e-2,
compare_fn={
Expand Down
Loading
Loading