[Bug] DeepSeek-V4-Flash MTP packed prefill emits NaN logits for extents > 128 tokens (real W8A8 weights)

[ndleslx]

Background

In pypto-lib, models/deepseek_v4_flash_mtp/prefill_fwd.py (l3_prefill_fwd, the packed-prefill program DeepSeek-V4-Flash serving uses) on a2a3 (8 devices) has a precision problem: any packed prefill extent above 128 tokens produces NaN logits with real W8A8 weights, and the NaN-derived routing then faults the follow-on MTP/decode program with an AICore "VEC instruction error: the ub address out of bounds", killing the request.

Reproduce with pypto-serving (hw-native-sys/pypto-serving PR #204, head a261960 — its KV-sizing ring-budget fix is required, otherwise a separate 4×ring_heap pooled-arena OOM on origin/main masks this failure first):

pypto-serving \
  --model <w8a8 checkpoint dir> \
  --prompt "<' and' repeated 150 times> Huawei is" \
  --platform a2a3 --devices <8 free device ids> \
  --dp 8 --ep 8 --block-size 128 \
  --max-model-len 1024 --max-num-seqs 8 --max-num-batched-tokens 512 \
  --long-prefill-token-threshold 4096 \
  --generate-config '{"max_new_tokens": 16, "temperature": 0.0}' \
  --num-speculative-tokens 1 --no-enable-prefix-caching \
  --ring-dep-pool 131072 --ring-task-window 131072 --ring-heap 2147483648

Checkpoint: convert HF deepseek-ai/DeepSeek-V4-Flash with pypto-serving's scripts/convert_deepseek_v4_to_w8a8.py (procedure in pypto-serving docs/dev/model/deepseek-v4.md).

Observed — deterministic on 8 devices, per packed extent:

Packed extent	Result
64, 128 tokens	correct greedy generation
129, 160, 256, 302 tokens	Sampler received non-finite logits, then AICore 507015 → empty completion

Reproduction environment:

Component	Version
pypto-lib	85e0790 (main tip; two local inert harness patches noted below)
pypto	7542c36d (main, detached)
simpler	0659e29b (pypto runtime gitlink, detached)
ptoas	0.60
pto-isa	be5ccb76 (per runtime/pto_isa.pin)
CANN	9.0.0

All components match pypto 7542c36d's pin chain. Local-only patches while reproducing (inert to the repro): --save-data flag in prefill_fwd.py's harness main(), and a GOLDEN_DEVICE_OUT_DIR device-output dump hook in golden/runner.py — neither touches kernel code or the serving path.

Diagnosis: unclear (pypto-lib kernel authoring vs toolchain) — the corruption is inside l3_prefill_fwd's multi-tile (>128-token) walk, but pypto/simpler cannot be bisected against this pin: pypto-lib 85e0790 fails to compile under ptoas 0.59 at every pypto revision tried (6863c6dd, af8c15f4, 106df194, 8cf3f058 — all with 'pto.tcvt' op requires explicit tmp for non-saturating narrowing tcvt when PlanMemory is skipped), so 85e0790 has only ever executed on ptoas v0.60 + simpler 0659e29b.

Description

Evidence collected (all on the environment above):

1. Failure chain: l3_prefill_fwd emits non-finite logits for extents > 128 (host sampler warning fires immediately after prefill); the NaN activations produce garbage routing/top-k; the next program crashes with VEC instruction error: the ub address out of bounds (0x4000000000000000); the worker lane is poisoned and the request returns empty. Device plog names the faulting kernel aicore_kernel_0_mix_aic, which is simpler's merged-binary entry macro (KERNEL_ENTRY(x) = x##_0_mix_aic), so it identifies the program, not the scope; task_id/func_id 15 maps via the generated kernel_config.py to qr_proj_seed (decode program) or qr_rms_norm_quant (mtp-prefill program, launched with a dynamic t_dim/8 block count) — both in models/deepseek_v4_flash_mtp/qkv_proj_rope.py's qr-projection/quant pipeline.
2. Ruled out: prefix caching (fails with --no-enable-prefix-caching), host-side chunking (single 302-token dispatch with --long-prefill-token-threshold 4096 fails identically — the kernel-internal 128-token tile walk alone is sufficient), --max-model-len (260 vs 1024 identical), and all PTO2 runtime timeouts (PTO2_OP_EXECUTE_TIMEOUT_US=400000000 etc. changed nothing).
3. Exact threshold: 128 tokens (one tile) passes; 129 fails — and 256 (two perfectly aligned tiles) also fails, so it is not a partial-tail-row artifact: any second tile corrupts. Requires real W8A8 weights: the standalone harness python models/deepseek_v4_flash_mtp/prefill_fwd.py -p a2a3 -d <ids> --num-tokens 160 runs clean on random weights.
4. The standalone harness cannot carry this repro — and hides a NaN of its own: with golden_fn=None it prints [RUN] PASS (validation skipped). Dumping its device outputs (the local hook above) shows NaNs even at its default 128-token config at both extents tested (128 and 160): one fully-NaN logits row (logits [2, 8, 129280], 258560 NaN), pre_hc_hidden_out token row 95 fully NaN (32768/32768 elements, both ranks), plus NaN rows in kv_cache/csa_*/hca_* states. Whatever that artifact is, the golden path currently has no way to notice.
5. This serving path has never been green on any pin: pypto-serving's main still pins pypto-lib 2a1809f, which fails at kernel-compile on current pypto (create_tensor init_value removal), so 85e0790 is the first pin that runs end-to-end — and pypto-lib's own CI skipped its serving-deepseek job at 85e0790 (its model-tests checks are red on main as of Aug 28).

Possibly related: #929 (same-tree buffer-state corruption), #1043 (different — probabilistic near-tie argmax).

Host Platform

Linux (aarch64), Ascend 910 (a2a3), 8 devices

Additional Context

Found via hw-native-sys/pypto-serving#204 CI (its k1-prefix-cache accuracy case is the only red test; all other extents in that matrix pass). Serving-side one-liner repro and full extent table in the Background. Happy to help narrow the faulting scope further on this machine (device access + the dumped standalone outputs are available).

[ndleslx]

Update: the "cannot be bisected" limitation is resolved — we've bisected the pypto stack, and the failure localizes to one commit.

Bisect range: green 42881d7b → red 57e5c315. The flip is exactly at pypto 7542c36d ("Upgrade PTOAS to v0.60 and pin runtime to 0659e29 (#2523)"). All seven intermediate commits are jointly exonerated (8cf3f058, which contains them all, passes). The earlier note that 85e0790 "fails to compile under ptoas 0.59" was an artifact of testing old pypto revisions against 0.59 — those revisions actually pin v0.57, on which 85e0790 compiles and the case passes.

Matrix (302-token single-dispatch packed prefill, pypto-lib 85e0790, pypto-serving PR #204 head a261960, a2a3 ×8):

#	pypto	simpler	ptoas	Result
1	42881d7b	799640e6	0.57	✅ correct greedy generation
2	8cf3f058 (parent)	799640e6	0.57	✅ correct greedy generation
3	8cf3f058	0659e29b	0.57	✅ correct — simpler bump exonerated (incl. the zero-fill removal)
4	8cf3f058	799640e6	0.60	✗ 'pto.tcvt' op requires explicit tmp ... (won't assemble)
5	7542c36d	799640e6	0.57	✗ expected ':' parse error (new .pto syntax)
6	7542c36d	0659e29b	0.59	✗ 'pto.tmrgsort' op format2 expects tmp.cols >= sum(src.cols) = 1024
7	7542c36d	0659e29b	0.60	🔴 assembles; NaN logits for extents > 128

So the regression lives in the inseparable pair {pypto 7542c36d explicit-tmp tcvt/tmrgsort lowering, PTOAS v0.60 assembly} — they require each other (rows 4–6), so neither half can be tested alone from a consumer repo. The runtime (simpler 799640e6→0659e29b) is ruled out (row 3).

Filed upstream as hw-native-sys/pypto#2558 with the full matrix and the reproduction environment. This issue's kernel-side analysis (multi-tile walk, qkv_proj_rope.py qr-quant pipeline, harness-hidden NaNs) still stands as the symptom-side companion.