perf(ascend): accelerate packed Q/K and strided KV/SwiGLU kernels

[JoeZhang-0x000]

Summary

• Add BF16 128-dimension Q/K RMSNorm and NeoX RoPE kernels that process a token's heads together, preserve packed QKV strides, and bound launches to 64 blocks.
• Parallelize strided KV-cache writes over token/head tasks, combine contiguous heads, tile DMA through 2 KiB UB storage, and synchronize MTE2/MTE3 reuse.
• Add a BF16 SwiGLU path for padded outputs, distributing rows over eight groups of eight hidden tiles. Contiguous outputs retain ACLNN, which has lower launch overhead on the measured shapes.
• Add boundary, padding and dynamic NPUGraph regressions, plus a reproducible operator latency script. Public operator signatures are unchanged.

Motivation

Packed Q/K normalization currently stages individual rows through many device copies; RoPE gathers expanded tables and copies the vendor output back. The strided KV kernel executes all tokens and heads on one block. These costs became bottlenecks while optimizing the non-Attention vLLM path on Ascend 910B4.

The new kernels read the original layouts directly. Unsupported dtypes/shapes/layouts retain the existing ACLNN implementation; the strided KV kernel continues to support NHD/HND caches, padded heads, negative slots and unaligned tails.

Type of Change

• perf — performance improvement
• test — regression coverage
• build — AscendC source registration
• Breaking change

Platforms Affected

• Ascend (WITH_ASCEND)
• Build system / CMake — existing Ascend-only custom-kernel build
• CPU / NVIDIA / Iluvatar / MetaX / Cambricon / Moore
• PyTorch C++ backend / Python public API

Smoke Test Result

Ascend 910B4, CANN 9.1, PyTorch/TorchNPU 2.10; InfiniRT 6b256e0ab0853d9e2a7ef059bac22c847bd02251. Built with WITH_ASCEND=ON, BUILD_ASCEND_CUSTOM=ON, and the default smoke allowlist extended by rotary_embedding,reshape_and_cache_flash.

python -m pytest tests -m smoke --devices ascend -q
87 passed, 33 skipped, 4773 deselected

python -m pytest tests/test_rms_norm.py tests/test_rotary_embedding.py \
  tests/test_silu_and_mul.py tests/test_reshape_and_cache_flash.py --devices ascend -q
230 passed, 20 skipped

Exact BF16 RMSNorm/RoPE comparisons and graph replay after changing input, weights, positions and cache passed. KV checks include 8191/8192 tokens, both cache layouts, padded heads and an 8193-element head exceeding the fixed UB tile. FP16/FP32 fallback cases also passed. git diff --check, C++ formatting, and Python lint/format checks passed.

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	No	N/A — not affected	N/A
Iluvatar	No	N/A — not affected	N/A
MetaX	No	N/A — not affected	N/A
Cambricon	No	N/A — not affected	N/A
Moore	No	N/A — not affected	N/A
Ascend	Yes	Build and smoke passed	Four affected operator modules passed; unrelated full suite not run in this smoke build

Benchmark / Performance Impact

Before: InfiniOps 4c014ca57a0af78bd1501a9b9bfe98647f05b55b. After: this branch. Both packages built against the same InfiniRT/CANN, run serially on device 0 without another experiment or compilation. BF16; five warmups, five synchronized batches of 20 calls, median warm host-plus-device latency. Initialization and graph capture are excluded; these are eager operator measurements.

ASCEND_RT_VISIBLE_DEVICES=0 PYTHONPATH=/path/to/baseline/package:$PYTHONPATH \
  python scripts/bench_ascend_qk.py --output baseline.json --iterations 20 --repeats 5
ASCEND_RT_VISIBLE_DEVICES=0 PYTHONPATH=/path/to/candidate/package:$PYTHONPATH \
  python scripts/bench_ascend_qk.py --output candidate.json --iterations 20 --repeats 5

BF16 workload	Before (μs)	After (μs)
Q/K RMSNorm 64×32×128, packed input	5764.09	31.37
RoPE 64 tokens, Q32/K8×128, packed input	87.42	44.68
SwiGLU 64×3072, contiguous output (ACLNN)	24.06	24.64
SwiGLU 64×6144, contiguous output (ACLNN)	23.64	24.53
SwiGLU 64×3072, padded output	30.40	26.11
SwiGLU 64×6144, padded output	30.24	26.04
KV write 8192×8×128, NHD cache	32453.34	153.40
KV write 8192×8×128, HND cache	32493.46	1052.20

Notes for Reviewers

• The Q/K RMSNorm kernel uses FP32 square/reduction followed by square root and division, then separate normalization and weight multiplication before BF16 rounding. This matches the native Q/K path in the exact regression rather than using an approximate reciprocal square root.
• RoPE rotates in place after loading all heads for one token into UB. It reads the current packed cache and positions at execution/replay time, including cache contents changed at the same address.
• New paths are conditional on INFINI_HAS_CUSTOM_KERNELS; existing vendor paths remain available for other layouts. SwiGLU keeps contiguous output on ACLNN because measurements showed the custom kernel's launch overhead was higher there.
• This ports corresponding optimizations from the locked legacy InfiniCore C API to the current InfiniOps interfaces. The associated vllm-infinicore #8 retains the tested legacy patches. Its 30-case Qwen3-8B TP1/2/4 performance result includes plugin layout/Graph changes; it is not an isolated performance claim for this InfiniOps PR. The current InfiniOps port was validated at the operator level.