Skip to content

Commit 8104075

Browse files
YunjiQinYunjiQin
authored
feat(dsv4/v3_2): mark L2 orch outputs with correct direction (pl.Out / pl.InOut) (#659)
## What pypto requires orchestration-entry outputs to declare a direction (hw-native-sys/pypto#1901). An output left as a plain `pl.Tensor` is treated as `In`, so its device→host copy-back is skipped and the tensor **silently reads back as all-zeros on the host**, failing golden. This PR annotates every golden-compared output on its orchestration entry with the correct direction, decided by the golden `TensorSpec`: | spec | direction | annotation | |------|-----------|------------| | `is_output=True` **+** `init_value` | inout (read-modify-write) | **`pl.InOut`** | | `is_output=True`, no `init_value` | pure output | **`pl.Out`** | `pl.InOut` round-trips correctly as of **hw-native-sys/pypto#1918** (the specializer previously dropped the wrapper → "missing type annotation"). pypto-lib CI builds against pypto `main`, which now includes that fix. ## Scope / rules - Direction lives on the **orchestration entry only** — the `@pl.jit` entry, its `@pl.jit.host` L3 driver, or the `@pl.function(type=Opaque)` method. `@pl.jit.inline` sub-kernels are left as upstream (their direction tag is stripped at splice time). - Normalizes **pre-existing** `pl.Out`-on-inout entries (prefill_attention_*, decode_compressor_*, decode_indexer*, prefill_indexer*) to `pl.InOut` so the whole codebase is consistent. ## Changes (21 files, +36/−36 — 33 `pl.InOut`, 3 `pl.Out`) - **inout → `pl.InOut`**: decode/prefill attention `kv_cache`; compressor `kv_state`/`score_state`/`compress_state`/`cmp_kv`/`cmp_kv_cache`; indexer `idx_kv_cache`; decode_layer / decode_fwd / prefill_fwd `kv_cache` (entry + L3 host); v3_2 decode_front `kv_cache`/`pe_cache`/`k_cache_idx`/`dispatch_buf`. - **pure output → `pl.Out`**: v3_2 back `out`; prefill_front_draft `dispatch_buf` (x_out / x_next / logits already `pl.Out`). ## Verification - AST audit: every `is_output` spec (incl. `@pl.jit.host` params and conditional `is_output=name==…` specs) maps to a correctly-directioned entry param; no `@pl.jit.inline` param carries `pl.InOut`. - **Device (a2a3) golden PASS**: decode/prefill attention (csa/hca/swa), compressor ratio4/128, indexer(_compressor), decode_compressor, decode_indexer, decode_layer — all compile and pass (kv_cache / kv_state / score_state / cmp_kv / idx_kv_cache / x_out). > Note: some cases fail their **second output** (`x_out`/`x_next`) on the **simulator** (a2a3sim/a5sim) — a pre-existing numerical-precision limitation, tracked as daily-CI known failures, independent of this annotation change (fails identically on upstream regardless of direction).
1 parent cbb345d commit 8104075

21 files changed

Lines changed: 36 additions & 36 deletions

‎models/deepseek/v3_2/deepseek_v3_2_decode_back.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ def deepseek_v3_2_decode_back_layer(
7070
w_gate: pl.Tensor[[HIDDEN_SIZE, INTER_SIZE], pl.BF16],
7171
w_up: pl.Tensor[[HIDDEN_SIZE, INTER_SIZE], pl.BF16],
7272
w_down: pl.Tensor[[INTER_SIZE, HIDDEN_SIZE], pl.BF16],
73-
out: pl.Tensor[[BATCH_SIZE, HIDDEN_SIZE], pl.BF16],
73+
out: pl.Out[pl.Tensor[[BATCH_SIZE, HIDDEN_SIZE], pl.BF16]],
7474
) -> pl.Tensor[[BATCH_SIZE, HIDDEN_SIZE], pl.BF16]:
7575
# Scope: output projection + residual + post-rms + MLP + residual.
7676
node_id = pl.cast(pl.tensor.read(node_id_t, [0]), pl.INDEX)

‎models/deepseek/v3_2/deepseek_v3_2_decode_front.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -106,10 +106,10 @@ def deepseek_v3_2_decode_front_scope1234(
106106
layer_id_t: pl.Tensor[[1], pl.INT32],
107107
w_q_nope_to_latent: pl.Tensor[[NUM_HEADS, QK_NOPE_HEAD_DIM, KV_LORA_RANK], pl.BF16],
108108
w_latent_to_v: pl.Tensor[[NUM_HEADS, KV_LORA_RANK, V_HEAD_DIM], pl.BF16],
109-
kv_cache: pl.Tensor[[CACHE_ROWS, KV_LORA_RANK], pl.BF16],
110-
pe_cache: pl.Tensor[[CACHE_ROWS, QK_ROPE_HEAD_DIM], pl.BF16],
111-
k_cache_idx: pl.Tensor[[CACHE_ROWS, INDEX_HEAD_DIM], pl.BF16],
112-
dispatch_buf: pl.Tensor[[EP_NODES, BATCH, ATTN_OUT], pl.BF16],
109+
kv_cache: pl.InOut[pl.Tensor[[CACHE_ROWS, KV_LORA_RANK], pl.BF16]],
110+
pe_cache: pl.InOut[pl.Tensor[[CACHE_ROWS, QK_ROPE_HEAD_DIM], pl.BF16]],
111+
k_cache_idx: pl.InOut[pl.Tensor[[CACHE_ROWS, INDEX_HEAD_DIM], pl.BF16]],
112+
dispatch_buf: pl.InOut[pl.Tensor[[EP_NODES, BATCH, ATTN_OUT], pl.BF16]],
113113
) -> pl.Tensor[[EP_NODES, BATCH, ATTN_OUT], pl.BF16]:
114114
# ===== scope1: MLA front path (RMSNorm + Q/KV projection + RoPE + cache writeback) =====
115115
# Outputs:

‎models/deepseek/v3_2/deepseek_v3_2_prefill_back.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ def deepseek_v3_2_prefill_back_layer(
6969
w_gate: pl.Tensor[[HIDDEN_CFG, INTER_CFG], pl.BF16],
7070
w_up: pl.Tensor[[HIDDEN_CFG, INTER_CFG], pl.BF16],
7171
w_down: pl.Tensor[[INTER_CFG, HIDDEN_CFG], pl.BF16],
72-
out: pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16],
72+
out: pl.Out[pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16]],
7373
) -> pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16]:
7474
for b in pl.parallel(0, BATCH_CFG, 1):
7575
seq_len_b = pl.tensor.read(seq_lens, [b])

‎models/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def deepseek_v3_2_prefill_front_layer(
137137
kv_norm_weight: pl.Tensor[[1, KV_LORA_RANK_CFG], pl.FP32],
138138
w_q_nope_to_latent: pl.Tensor[[NUM_HEADS_CFG, QK_NOPE_HEAD_DIM_CFG, KV_LORA_RANK_CFG], pl.BF16],
139139
w_latent_to_v: pl.Tensor[[NUM_HEADS_CFG, KV_LORA_RANK_CFG, V_HEAD_DIM_CFG], pl.BF16],
140-
dispatch_buf: pl.Tensor[[EP_NODES_CFG, BATCH_CFG, MAX_SEQ_CFG, ATTN_OUT_CFG], pl.BF16],
140+
dispatch_buf: pl.Out[pl.Tensor[[EP_NODES_CFG, BATCH_CFG, MAX_SEQ_CFG, ATTN_OUT_CFG], pl.BF16]],
141141
) -> pl.Tensor[[EP_NODES_CFG, BATCH_CFG, MAX_SEQ_CFG, ATTN_OUT_CFG], pl.BF16]:
142142
with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.auto_chunk]):
143143
layer_id = pl.tensor.read(layer_id_t, [0])

‎models/deepseek/v4/decode_attention_csa.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ def attention_csa_test(
343343
inner_norm_w: pl.Tensor[[IDX_HEAD_DIM], pl.BF16],
344344
inner_compress_state: pl.Tensor[[INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], pl.FP32],
345345
inner_compress_state_block_table: pl.Tensor[[B, INNER_STATE_MAX_BLOCKS], pl.INT32],
346-
kv_cache: pl.Tensor[[ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
346+
kv_cache: pl.InOut[pl.Tensor[[ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
347347
ori_block_table: pl.Tensor[[B, ORI_MAX_BLOCKS], pl.INT32],
348348
cmp_kv: pl.Tensor[[CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
349349
cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32],

‎models/deepseek/v4/decode_attention_hca.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ def attention_hca_test(
295295
cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16],
296296
compress_state: pl.Tensor[[COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32],
297297
compress_state_block_table: pl.Tensor[[B, COMPRESS_STATE_MAX_BLOCKS], pl.INT32],
298-
kv_cache: pl.Tensor[[ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
298+
kv_cache: pl.InOut[pl.Tensor[[ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
299299
ori_block_table: pl.Tensor[[B, ORI_MAX_BLOCKS], pl.INT32],
300300
cmp_kv: pl.Tensor[[CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
301301
cmp_block_table: pl.Tensor[[B, CMP_MAX_BLOCKS], pl.INT32],

‎models/deepseek/v4/decode_attention_swa.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,7 @@ def attention_swa_test(
234234
freqs_cos: pl.Tensor[[MAX_SEQ_LEN, ROPE_HEAD_DIM], pl.BF16],
235235
freqs_sin: pl.Tensor[[MAX_SEQ_LEN, ROPE_HEAD_DIM], pl.BF16],
236236
# KV cache (sliding-window only: [0, WIN) ori; no cmp portion)
237-
kv_cache: pl.Tensor[[B * ORI_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
237+
kv_cache: pl.InOut[pl.Tensor[[B * ORI_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
238238
block_table: pl.Tensor[[B, ORI_MAX_BLOCKS], pl.INT32],
239239
ori_slot_mapping: pl.Tensor[[T], pl.INT64],
240240
position_ids: pl.Tensor[[T], pl.INT32],

‎models/deepseek/v4/decode_compressor_ratio128.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -290,15 +290,15 @@ def compressor_ratio128(
290290
def compressor_test(
291291
x: pl.Tensor[[B_DYN, S_DYN, D], pl.BF16],
292292
kv: pl.Out[pl.Tensor[[B_DYN, S_DYN, HEAD_DIM], pl.FP32]],
293-
compress_state: pl.Out[pl.Tensor[[COMPRESS_STATE_BLOCK_NUM_DYN, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32]],
293+
compress_state: pl.InOut[pl.Tensor[[COMPRESS_STATE_BLOCK_NUM_DYN, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32]],
294294
compress_state_block_table: pl.Tensor[[B_DYN, COMPRESS_STATE_MAX_BLOCKS_DYN], pl.INT32],
295295
wkv: pl.Tensor[[OUT_DIM, D], pl.BF16],
296296
wgate: pl.Tensor[[OUT_DIM, D], pl.BF16],
297297
ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32],
298298
norm_w: pl.Tensor[[HEAD_DIM], pl.BF16],
299299
cos: pl.Tensor[[B_DYN, ROPE_HEAD_DIM // 2], pl.FP32],
300300
sin: pl.Tensor[[B_DYN, ROPE_HEAD_DIM // 2], pl.FP32],
301-
cmp_kv_cache: pl.Out[pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
301+
cmp_kv_cache: pl.InOut[pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
302302
position_ids: pl.Tensor[[B_DYN, S_DYN], pl.INT32],
303303
cmp_slot_mapping: pl.Tensor[[B_DYN, S_DYN], pl.INT64],
304304
state_slot_mapping: pl.Tensor[[B_DYN, S_DYN], pl.INT64],

‎models/deepseek/v4/decode_compressor_ratio4.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -268,15 +268,15 @@ def compressor_ratio4(
268268
def compressor_test(
269269
x: pl.Tensor[[B, S, D], pl.BF16],
270270
kv: pl.Out[pl.Tensor[[B, S, HEAD_DIM], pl.FP32]],
271-
compress_state: pl.Out[pl.Tensor[[COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32]],
271+
compress_state: pl.InOut[pl.Tensor[[COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32]],
272272
compress_state_block_table: pl.Tensor[[B, COMPRESS_STATE_MAX_BLOCKS], pl.INT32],
273273
wkv: pl.Tensor[[OUT_DIM, D], pl.BF16],
274274
wgate: pl.Tensor[[OUT_DIM, D], pl.BF16],
275275
ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32],
276276
norm_w: pl.Tensor[[HEAD_DIM], pl.BF16],
277277
cos: pl.Tensor[[B, ROPE_HEAD_DIM // 2], pl.FP32],
278278
sin: pl.Tensor[[B, ROPE_HEAD_DIM // 2], pl.FP32],
279-
cmp_kv_cache: pl.Out[pl.Tensor[[CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
279+
cmp_kv_cache: pl.InOut[pl.Tensor[[CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
280280
position_ids: pl.Tensor[[B, S], pl.INT32],
281281
cmp_slot_mapping: pl.Tensor[[B, S], pl.INT64],
282282
state_slot_mapping: pl.Tensor[[B, S], pl.INT64],

‎models/deepseek/v4/decode_fwd.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ def decode_fwd(
153153
wkv: pl.Tensor[[FWD_NUM_LAYERS * D, HEAD_DIM], pl.BF16],
154154
gamma_cq: pl.Tensor[[FWD_NUM_LAYERS * Q_LORA], pl.BF16],
155155
gamma_ckv: pl.Tensor[[FWD_NUM_LAYERS * HEAD_DIM], pl.BF16],
156-
kv_cache: pl.Tensor[[FWD_NUM_LAYERS * B * ORI_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
156+
kv_cache: pl.InOut[pl.Tensor[[FWD_NUM_LAYERS * B * ORI_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
157157
attn_sink: pl.Tensor[[FWD_NUM_LAYERS * H], pl.FP32],
158158
wo_a: pl.Tensor[[FWD_NUM_LAYERS * O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16],
159159
wo_b: pl.Tensor[[FWD_NUM_LAYERS * D, O_GROUPS * O_LORA], pl.INT8],
@@ -632,7 +632,7 @@ def l3_decode_fwd(
632632
wkv: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * D, HEAD_DIM], pl.BF16],
633633
gamma_cq: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * Q_LORA], pl.BF16],
634634
gamma_ckv: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * HEAD_DIM], pl.BF16],
635-
kv_cache: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * B * ORI_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16],
635+
kv_cache: pl.InOut[pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * B * ORI_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]],
636636
attn_sink: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * H], pl.FP32],
637637
wo_a: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16],
638638
wo_b: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * D, O_GROUPS * O_LORA], pl.INT8],

0 commit comments

Comments
 (0)