High-perf SPMD paged attention (GQA) matching CANN IFA performance

[learning-chip]

Reaches ~1 TB/s on A2 for GQA shapes, equal to the perf of CANN's torch_npu.npu_incre_flash_attention (for some shapes even 1.1x faster)

@ChaoWao @chenshengxin2026

Reproduce

cd simpler/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention_highperf/kernel
bash ./compile.sh
python ./test_pa_accuracy.py
python ./bench_pa_performance.py

TODOs:

Current version is 2000+ lines of raw CCE & host-side launch. Remaining items:

• Use PTO-ISA instead of raw intrinsics
• Port cross-core sync flags to PTO-ISA tpush/tpop APIs
• Register to aicpu run-time same as Add: SPMD paged attention with TPUSH/TPOP MIX kernel pipeline #596
  • @SergioMartin86 @georgebisbas please help on this item

Acknowledgement

I took the tiling scheme from ATB PA tiling, but rewrote it in numpy instead of original C++, to more trivially interface with python launcher.

[gemini-code-assist]

Code Review

This pull request adds a standalone Paged Attention kernel for NPU, including tiling logic, benchmarking, and accuracy tests. Critical issues were identified in the tiling logic where memory offsets for various tensors are calculated as element offsets instead of byte offsets, and the workspace size for split-KV reduction is missing the batch dimension. Additionally, the sequence rank assignment is incorrect, the kernel entry point fails to handle certain tiling key configurations, and the performance runner lacks necessary workspace zero-initialization between launches.

[gemini-code-assist]

The offsets addr_q and addr_o are calculated as element offsets, but the CCE kernel entry point in pa_entry.cce receives these tensors as uint8_t*. If the kernel adds these offsets directly to the base pointers without scaling by the element size, it will access incorrect memory locations. These should be converted to byte offsets using np.dtype(...).itemsize to avoid hardcoded values.

Suggested change

	addr_q += num_heads * head_dim * q_seqlen
	addr_o += num_heads * head_dim_v * q_seqlen
	addr_q += num_heads * head_dim * q_seqlen * np.dtype(np.float16 if dtype in (torch.float16, torch.bfloat16) else np.float32).itemsize
	addr_o += num_heads * head_dim_v * q_seqlen * np.dtype(np.float16 if dtype in (torch.float16, torch.bfloat16) else np.float32).itemsize

References

1. To calculate the size in bytes of a numpy array, use np.dtype().itemsize instead of a hardcoded value. This is more self-documenting and robust if the data type changes in the future.

[gemini-code-assist]

Similar to the Q/O offsets, addr_l and addr_ofd are element offsets being passed to a kernel that treats the corresponding GM pointers as uint8_t*. These must be byte offsets to ensure correct indexing in the kernel. Use np.dtype(...).itemsize instead of hardcoded values for robustness.

Suggested change

	addr_l += kvCN * num_heads * q_seqlen
	addr_ofd += num_heads * head_dim * q_seqlen # embeddingSize for oFd
	addr_l += kvCN * num_heads * q_seqlen * np.dtype(np.float32).itemsize
	addr_ofd += num_heads * head_dim * q_seqlen * np.dtype(np.float32).itemsize

References

1. To calculate the size in bytes of a numpy array, use np.dtype().itemsize instead of a hardcoded value. This is more self-documenting and robust if the data type changes in the future.

[gemini-code-assist]

The assignment tiling[base + 13] = indices[seq_idx] appears to be incorrect. indices is a list of original sequence indices sorted by length, so indices[seq_idx] returns the original index of the sequence at rank seq_idx. Since the loop is already iterating over the original index seq_idx, this assignment does not provide the rank of the current sequence. It should instead store the rank (sorted position) of the current seq_idx to allow the kernel to process sequences in the optimized order.

Suggested change

	tiling[base + 13] = indices[seq_idx] # sorted position → original index
	tiling[base + 13] = indices.index(seq_idx) # Store the rank of sequence seq_idx

[gemini-code-assist]

The workspace size calculations for o_core and l_size are missing the batch dimension. These buffers are used to store partial results for split-KV reduction across the entire batch. Without the batch factor, the allocated workspace will be insufficient for any batch size greater than 1. Additionally, use np.dtype(...).itemsize instead of hardcoded values.

Suggested change

	o_core = int(block_dim * SPLITKV_RATIO) * num_heads * block_dim * head_dim * 4
	l_size = int(block_dim * SPLITKV_RATIO) * num_heads * block_dim * 4
	o_core = batch * int(block_dim * SPLITKV_RATIO) * num_heads * head_dim * np.dtype(np.float32).itemsize
	l_size = batch * int(block_dim * SPLITKV_RATIO) * num_heads * np.dtype(np.float32).itemsize

References

1. To calculate the size in bytes of a numpy array, use np.dtype().itemsize instead of a hardcoded value. This is more self-documenting and robust if the data type changes in the future.

[gemini-code-assist]

The CustomPARunner does not zero-initialize the workspace buffers before each kernel launch. If the kernel uses these buffers for atomic operations or partial sum reductions (which is common in split-KV implementations), stale data from previous iterations will corrupt the results. While torch.zeros is used in __init__, the buffers must be cleared in __call__ if they are modified by the kernel.

[gemini-code-assist]

The kernel entry point does not handle tiling_key_val values that include the is_split_block bit (bit 7). According to pa_tiling.py, this bit is set when block_size >= 128 and head_dim == 256. If a user provides such a configuration, the kernel will silently skip execution as none of the if/else if conditions will match.

[zouzias]

Could it be that as_ptr has a non-negligible kernel launch overhead here?

[zouzias]

Accessing the class members here from Python class CustomPARunner might have a non-negligible overhead.