Perf(spmd_paged_attention): prefetch K/V and reuse block ids in AIC

[chenshengxin2026]

Summary

Performance optimization for tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention, targeting the AIC pipeline in paged_attention_parallel.cpp.

• AIC QK prefetch: issue TLOAD for the next block's K before the current QK's MTE1 -> M sync, so the next-block fetch overlaps the current step's TMOV / TMATMUL / TPUSH path on PIPE_MTE2.
• AIC PV prefetch: same prefetch pattern for V in the PV step (uses EVENT_ID1 to avoid aliasing with QK's EVENT_ID0).
• Reuse already-read block ids: hoist bt[bt_offset + i] reads up to aic_process_blocks and pass the resolved kv_block_id into the step helpers. Each block id is read from bt[] once and reused — prev_block_id feeds PV[i-1], current block_id feeds QK[i], and a one-step next_block_id lookahead drives the prefetch.

The step helpers gain three optional flags — current_loaded (skip the inline TLOAD because the previous step already prefetched it), has_next (issue the lookahead TLOAD), and next_kv_block_id. Prologue / steady state / epilogue in aic_process_blocks set these flags so:

• Prologue QK[0]: no current_loaded, prefetches K[1].
• Steady QK[i] (i>=1): current K[i] already loaded by previous step's prefetch; if i+1 < n, prefetch K[i+1].
• Steady PV[i-1]: V[i-1] already loaded for i>=2; if i < n, prefetch V[i].
• Epilogue PV[n-1]: V[n-1] already loaded by the prior PV's prefetch when n_blocks > 1.
• Degenerate n_blocks == 1: no prefetch, single QK + PV with their own loads.

No semantic / numerical change — the loaded data and matmul order are identical, only the issue timing of the next TLOAD shifts earlier within the same step.

Performance

Measured via tools/benchmark_rounds.sh on hardware.

Before (baseline on upstream/main):

================================================================
  Performance Summary (tensormap_and_ringbuffer)
================================================================

  Example                                   Elapsed (us)    Sched (us)     Orch (us)
  ----------------------------------------  ------------  ------------  ------------
  spmd_paged_attention (Case1)                    1501.8        1501.7           7.0
  spmd_paged_attention (Case2)                     752.5         752.4           6.7

================================================================
  Benchmark complete (tensormap_and_ringbuffer): 2 passed, 0 failed (2 total)
================================================================

After (this PR):

================================================================
  Performance Summary (tensormap_and_ringbuffer)
================================================================

  Example                                   Elapsed (us)    Sched (us)     Orch (us)
  ----------------------------------------  ------------  ------------  ------------
  spmd_paged_attention (Case1)                    1316.5        1316.5           6.2
  spmd_paged_attention (Case2)                     694.4         694.4           6.1

================================================================
  Benchmark complete (tensormap_and_ringbuffer): 2 passed, 0 failed (2 total)
================================================================

Delta:

Case	Before (us)	After (us)	Δ (us)	Improvement
Case1	1501.8	1316.5	-185.3	-12.3%
Case2	752.5	694.4	-58.1	-7.7%

Testing

• Simulation: pytest tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention --platform a2a3sim
• Hardware: pytest tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention --platform a2a3 --device <range>

[gemini-code-assist]

Code Review

This pull request introduces a pre-fetching mechanism in the aic_qk_step and aic_pv_step kernels to overlap data loading with computation. By passing block loading status and next block IDs, the implementation aims to improve pipeline efficiency. Feedback suggests optimizing the loop in aic_process_blocks by carrying the next block ID forward to avoid redundant global memory reads from the block table array.