Repository navigation
Description
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
- On the A5 box (pypto @ PR Name each runtime's working tensor
Tensorin 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" - Observe:
x_outFAIL (~77% ratio_reldiff),kv_cachePASS. - To confirm the zero
attn_out: addpl.dump_tag(attn_out)after thesparse_attn(...)call indecode_attention_csa.py, add a--dump-tensor/enable_dump_argsflag, run... -p a5 --dump-tensor 1, and decode the dumpedattn_out([8,4096] BF16) →nonzero=0/32768. - Compare:
python models/deepseek/v4/decode_sparse_attn.py -p a5 -d <dev>(standalone) →attn_outNON-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 inlinesparse_attn) +decode_layer/decode_mtp/prefill_mtp(3, compose attention). The other 5 (decode_sparse_attn{,_hca,_swa}/prefill_sparse_attn~10% drift +prefill_indexerscore) 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_tagreturned stale points; level-2dump-tensortrips 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 tensorTensorin kernel and orchestration sources #2044). The local env also carries a temporary a5set_allow_early_resolvepatch in simpler'spto_types.h— unrelated to this bug (it only unblocks 22set_allow_early_resolveorchestration-compile failures elsewhere).
Metadata
Metadata
Assignees
Labels
Type
Projects
- StatusShow more project fieldsDone
Activity