[Bug] A5 onboard: inlined-kernel host-side pl.read/pl.write orchestration yields all-zero sparse_attn output

[lwDavid]

Background

In pypto-lib, models/deepseek/v4/decode_attention_csa.py on A5 (Ascend950, onboard) has a precision problem: x_out FAILs ~77% (ratio_reldiff, worst rdiff≈2.0) while kv_cache PASSES. Root cause traced via dump-tensor: the sparse_attn output attn_out ([8,4096] BF16) is ALL ZEROS on A5 when sparse_attn is inlined into the composed CSA kernel; the golden attn_out is non-zero (absmax 0.35). The same sparse_attn produces non-zero output when run standalone (decode_sparse_attn.py -p a5), and the full kernel passes on A2A3. So the kernel logic is correct; the fault is in executing the composed/inlined kernel on the A5 runtime.

Reproduce with (device access via task-submit):

task-submit --device auto --device-num 1 --run \
  "cd <pypto-lib> && PTOAS_ROOT=... PTO_ISA_ROOT=... PYTHONPATH=... \
   python models/deepseek/v4/decode_attention_csa.py -p a5 -d \$TASK_DEVICE"

Reproduction environment:

Component	Version
pypto-lib	f894b63 (branch: main)
pypto	c2edc705 (branch: pr-2044, PR #2044 wip/legalize-tile-cast)
simpler	41fc8ba8 (detached; fork yanghaoran29/simpler wip/a5-cce-simd-vf-fusion-off)
ptoas	0.48
pto-isa	83d01313
CANN	cann-9.1.T500

Platform: a5 (Ascend 950 hardware) | Runtime Variant: tensormap_and_ringbuffer | Host Platform: Linux aarch64 | Git Commit ID: 41fc8ba8 (simpler; fork branch wip/a5-cce-simd-vf-fusion-off) | CANN Version: cann-9.1.T500

Description

When sparse_attn (models/deepseek/v4/decode_sparse_attn.py, a @pl.jit.inline kernel) is inlined into the larger composed CSA kernel (decode_attention_csa.py = hc_pre + qkv_proj_rope + compressor + indexer + sparse_attn + hc_post), the A5 runtime produces all-zero attn_out, so the whole attention contribution is lost. Standalone, the identical sparse_attn produces non-zero (correct-ish) output on A5.

sparse_attn's inputs in the composed kernel are all correct (q from qkv_proj_rope — fixed by pypto-lib #792, and kv_cache from the same call PASSES in-csa; idx_topk from indexer — decode_indexer passes standalone; cmp_kv from compressor — passes standalone; kv_cache PASS, harness-validated). So the bug is in sparse_attn's execution when composed, not its inputs.

Likely culprit: sparse_attn relies heavily on host-side pl.read/pl.write control flow — valid_block_mask gating via if pl.read(valid_block_mask, [...]) > 0, the qk_wcur scalar counter pl.read/pl.write round-trips, and the qk_order planning loop (host reading device tensor elements to drive Python branching). The A5 (tensormap_and_ringbuffer) runtime appears to mishandle this host orchestration in the inlined/composed case, yielding zero attention output. (Note: allow_early_resolve on the plan scope was tested and is NOT the cause — removing it doesn't help; the attention task already lists deps=[qk_plan_tid].)

Steps to Reproduce

1. On the A5 box (pypto @ PR Name each runtime's working tensor Tensor in kernel and orchestration sources #2044 + simpler @ 41fc8ba + ptoas v0.48 + pto-isa 83d01313, all at their pins), with pypto-lib @ main (f894b63):

   task-submit --device auto --device-num 1 --run "cd <pypto-lib> && \
     PTOAS_ROOT=<ptoas-bin> PTO_ISA_ROOT=<pto-isa> PYTHONPATH=<pypto>/python:\$PYTHONPATH \
     python models/deepseek/v4/decode_attention_csa.py -p a5 -d \$TASK_DEVICE"
   

2. Observe: x_out FAIL (~77% ratio_reldiff), kv_cache PASS.
3. To confirm the zero attn_out: add pl.dump_tag(attn_out) after the sparse_attn(...) call in decode_attention_csa.py, add a --dump-tensor/enable_dump_args flag, run ... -p a5 --dump-tensor 1, and decode the dumped attn_out ([8,4096] BF16) → nonzero=0/32768.
4. Compare: python models/deepseek/v4/decode_sparse_attn.py -p a5 -d <dev> (standalone) → attn_out NON-ZERO (~10% precision error).

Expected Behavior

attn_out (and thus x_out) should be non-zero and match the torch golden within tolerance — as it does for standalone sparse_attn on A5 and for the full decode_attention_csa on A2A3.

Actual Behavior

attn_out is all zeros in the inlined/composed CSA kernel on A5 → x_out FAILs ~77%. Identical sparse_attn is non-zero standalone.

Additional Context

• Impact: explains 9 of 14 remaining A5 precision failures in DeepSeek-V4: decode/prefill_attention_{csa,hca,swa} (6, all inline sparse_attn) + decode_layer/decode_mtp/prefill_mtp (3, compose attention). The other 5 (decode_sparse_attn{,_hca,_swa}/prefill_sparse_attn ~10% drift + prefill_indexer score) are independent.
• Tooling caveat: attn_out (task-boundary output) is reliably dumpable → the all-zero finding is solid. Intra-task tensors (valid_block_mask) and reshaped inputs (idx_topk, q) could not be reliably captured (pl.dump_tag returned stale points; level-2 dump-tensor trips the AICPU collector timeout on this large kernel), so "inputs are all correct" is reasoned from each component passing standalone, not a per-tensor measurement.
• The simpler submodule here is on the fork branch wip/a5-cce-simd-vf-fusion-off (per pypto PR Name each runtime's working tensor Tensor in kernel and orchestration sources #2044). The local env also carries a temporary a5 set_allow_early_resolve patch in simpler's pto_types.h — unrelated to this bug (it only unblocks 22 set_allow_early_resolve orchestration-compile failures elsewhere).