[baominghelly]

Summary

• Add Ascend 310P workarounds for short-row fused RMSNorm, SwiGLU tiling, NeoX decode RoPE, strided KV-cache writes, and aligned dense causal prefill.
• Route other 310P FP16 attention through per-sequence PromptFlashAttentionV3, including dense/paged KV, unequal sequence lengths, causal/local masks, and zero outputs for fully masked rows. Retain the single-sequence prefill fast path and the 910C FusedInferAttention path.
• Correct cached fused RMSNorm output rebinding and batch strided RMSNorm staging with ACLNN copy.
• Centralize assert-based ACL diagnostics and consistently check ACL/ACLNN return codes in RMSNorm and RoPE, retaining parameter assertions.
• Add model-shaped layout/reuse regressions and correct the existing paged-prefill test to use 256-token pages while retaining cross-page, nontrivial block-table coverage.

Motivation

The existing Ascend backend encounters unsupported CANN paths and incomplete outputs on 310P model shapes. The changes preserve the existing 910C computation path except for the shared RMSNorm copy optimization, fused RMSNorm address correction, diagnostics, and RoPE chunk bookkeeping. RoPE uses one chunk on 910C.

Changes are split into eleven Conventional Commits so compatibility work, diagnostics, performance work, test corrections, and Ruff formatting can be reviewed separately. No related issue is claimed.

Type of Change

• feat — new feature / new operator / new platform
• fix — bug fix
• perf — performance improvement (no behavioral change)
• refactor — code restructuring without behavior change
• test — adding or fixing tests only
• docs — documentation only
• build / ci — build system or CI configuration
• chore — tooling, formatting, or other non-code changes
• Breaking change

Platforms Affected

• CPU (WITH_CPU)
• NVIDIA (WITH_NVIDIA)
• Iluvatar (WITH_ILUVATAR)
• MetaX (WITH_METAX)
• Cambricon (WITH_CAMBRICON)
• Moore (WITH_MOORE)
• Ascend (WITH_ASCEND)
• PyTorch C++ bindings (WITH_TORCH)
• Build system / CMake / CI
• Python bindings / user-facing API

Smoke Test Result

Validated tree: commit e7d34f701235fa63d1d81d6d236d78750eee3eb5, based on master 1ca4777610d77e9063fffdb975ab35da8466104d. Both machines use the same source tree, with no tracked source changes after validation. The final commit formats six changed Python tests; their parsed ASTs match the pre-format versions, and the affected regressions are rerun after formatting.

Both isolated CMake builds and installs passed with Release -O3 -UNDEBUG, WITH_ASCEND=ON, WITH_CPU=OFF, WITH_TORCH=ON, WITH_LINKED=OFF, BUILD_ASCEND_CUSTOM=ON, and Python bindings. SoC targets were Ascend310P3 on x86_64 and Ascend910_9362 on the ARM64 910C host. CANN 9.1.0, torch 2.10.0+cpu, torch_npu 2.10.0.post4; Python 3.12.13 and 3.11.15 respectively. Runtime headers/libraries and all build outputs used isolated configured prefixes; existing installations were preserved.

Exact validation entry points, with out set to the corresponding isolated validation directory:

bash "$out/build.sh" 310p  # Or: 910c
bash run_regression.sh 310p  # Or: 910c; final unchanged-operator regressions
bash run_flash.sh 310p       # Or: 910c; full FlashAttention module and smoke
# The final smoke subprocess:
python "$out/run_tests.py" -q --tb=short tests -m smoke --devices ascend \
  --junitxml="$out/flash-fix-smoke-20261022.xml"

The external validation runner disables torch_npu JIT compilation, verifies the loaded extension/library prefix, and restricts 310P collection to native implementation 0 and FP16. It does not filter dtypes on 910C. The runner and logs are validation artifacts, not repository changes.

910C full smoke selection:
95 passed, 43 skipped, 11378 deselected

310P full smoke selection:
71 passed, 26 skipped, 11419 deselected

No FlashAttention module exclusion is used. The former workspace-query abort is fixed; the full native FP16 FlashAttention selection on 310P passes, including the ten single-sequence padding cases and twelve new length/block-table reuse regressions.

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	No	N/A - not affected	N/A - not affected
Iluvatar	No	N/A - not affected	N/A - not affected
MetaX	No	N/A - not affected	N/A - not affected
Cambricon	No	N/A - not affected	N/A - not affected
Moore	No	N/A - not affected	N/A - not affected
Ascend	Yes	Both builds and smoke selections passed	Targeted native regression: 310P 294 passed / 10 skipped; 910C 475 passed / 157 skipped. Generator tests: 84 passed per machine. Full operator suite not run.

Targeted commands used separate processes per module:

python "$out/run_tests.py" -q --tb=short tests/test_rms_norm.py --devices ascend
# Same command for fused_add_rms_norm, reshape_and_cache_flash, silu_and_mul,
# rotary_embedding, flash_attn_varlen_func and causal_softmax.
# On 310P only, RoPE uses -k ascend_neox_decode. FlashAttention has no
# shape/name filter; the external runner selects native FP16 on 310P.
python "$out/run_tests.py" -q --tb=short -p tests.conftest \
  "$out/test_triangular_master.py" --devices ascend
python -m pytest --noconftest -q --tb=short \
  tests/test_generate_public_operator_header.py tests/test_generate_wrappers.py \
  tests/test_resolve_linked_ops.py

Native regression	310P passed / skipped	910C passed / skipped
RMSNorm	37 / 1	61 / 1
FusedAddRMSNorm	44 / 0	116 / 0
ReshapeAndCacheFlash	11 / 0	28 / 0
SiluAndMul	10 / 1	26 / 1
RotaryEmbedding	96 / 0	128 / 132
FlashAttnVarlenFunc	42 / 8	50 / 23
CausalSoftmax	6 / 0	18 / 0
FP16 triu/tril smoke	48 / 0	48 / 0

BF16 and general FP32 validation are not claimed for 310P. Its CANN RoPE D=8 limitation remains outside the tested supported model shapes. Only one NPU was exposed per process, so multi-device guard cases skip. The 910C NPU was idle; the 310P NPU was shared with an existing workload and only correctness tests were run for the new attention path. Smoke counts overlap targeted cases and are not added to targeted totals.

Benchmark / Performance Impact

External harness bench_rms.py: FP16 Q/K RMSNorm pairs with head dimension 128, 16 query heads, 8 key heads, batch 1; input views sliced from [S,4096]. Measure wall-clock call plus stream synchronization, five warmups, twenty samples, report median milliseconds. Output values are checked against FP32 CPU reference. TASK_QUEUE_ENABLE=0, JIT disabled, ASCEND_LAUNCH_BLOCKING unset. Both source versions use Release -O3 -UNDEBUG and the same runtime dependency. The clean-master baseline builds only native RMSNorm/Python bindings; the candidate is the full isolated validation build including Torch bindings.

bash "$out/bench.sh" 310p  # Or: 910c

Device	S	Strided output	Master pair ms	Candidate pair ms
310P3	1	No	0.119	0.126
310P3	1	Yes	0.230	0.166
310P3	36	No	3.879	0.174
310P3	36	Yes	6.599	0.189
310P3	333	No	27.444	0.358
310P3	333	Yes	54.040	0.385
910C	1	No	0.098	0.113
910C	1	Yes	0.205	0.138
910C	36	No	2.945	0.135
910C	36	Yes	5.387	0.149
910C	333	No	24.643	0.132
910C	333	Yes	47.821	0.155

These measurements support the strided-copy improvement, not an across-the-board latency claim. The contiguous single-token path is 7-15 microseconds slower in this short run; repeated statistical characterization was not performed. No new whole-model throughput claim is made. The PromptFlashAttentionV3 path includes padding/copy overhead and is not claimed as a general speedup.

Notes for Reviewers

• The installed CANN 9.1 package has no 310P FusedInferAttentionScore registration/kernel. The new fallback uses PromptFlashAttentionV3 with 256-aligned staging, explicit masks, and per-sequence synchronization; paged KV is gathered in block-table order. This is a correctness fallback with extra copies/synchronization, not a throughput optimization. The old single-sequence causal prefill path remains selected for its existing shapes.
• The 310P reference calculation uses CPU FP32 because its torch_npu runtime does not implement NanToNum. Test tolerances remain unchanged. The twelve new FP16 regressions run on both 310P and 910C and check changing lengths and block-table values on a non-default stream.
• The prior four 910C paged-prefill failures were test/contract mismatches: pages were 128 tokens while the public API requires multiples of 256. The corrected test uses lengths 256/512 and pages 256, preserving reordered, cross-page blocks; all four now pass.
• RMSNorm batching applies to all Ascend devices' non-contiguous staging. RoPE's eight-head splitting applies only to the specified 310P shapes, but its chunk bookkeeping is shared.
• ACL diagnostics remain assert-based and compile out under NDEBUG; calls still execute. Shape/stride/dtype assertions are retained.
• CANN-finalization WorkspacePool messages and existing RoPE narrowing warnings are retained; no general shutdown-lifetime or warning cleanup is included.
• clang-format 21.1.8 dry-run, Ruff 0.15.22 lint/format checks, and git diff --check pass locally. GitHub Clang Format and Ruff pass for final head e7d34f70.
• GitHub platform CI is not green: the Ascend unit job fails during custom-kernel CMake configuration because its selected SOC_VERSION=Ascend3103 is unsupported by CANN 9.1; the NVIDIA unit job fails before compilation because its build environment rejects the existing tool.scikit-build.sdist.force-include setting. Neither failure reaches operator tests. The corresponding CMake/packaging configuration is unchanged in this PR. The independently configured 310P3 and 910C builds/regressions above pass; other failed/pending platform CI jobs are not claimed as passes.
• Validation scripts, logs, build trees and installation outputs are intentionally excluded from the commits.