diff --git a/.github/skills/cm_head_size_256_fix.md b/.github/skills/cm_head_size_256_fix.md new file mode 100644 index 00000000..232bda37 --- /dev/null +++ b/.github/skills/cm_head_size_256_fix.md @@ -0,0 +1,329 @@ +# Skill: HEAD_SIZE=256 Support for PA Kernels on Xe2 + +LNL Xe2 reference: 64 EUs, 2 GHz, 256 GRF (16 KB), 128 KB SLM, ~30 TFLOPS FP16 peak. +CM constants: CM_GRF_WIDTH=512, REG_N=16, REG_K=16, REG_M=8, num_P_tiles=2. +CM matrix limit: `sizeof(matrix) < 16384` bytes (strict). + +--- + +## test_kvcache_update.py (pa_kv_cache_update kernel) + +**Problem:** None. HEAD_SIZE=256 compiled and passed on the first attempt. The kvcache update kernel writes incoming K/V tokens into paged block storage; its matrix sizes are independent of HEAD_SIZE (it processes per-element, not full-head accumulation). + +**Fix:** Added test cases only (no kernel changes): +- `test_pa_kv_cache_update([32*1024], [0], ..., k_head_size=256, v_head_size=256, ...)` for fp16 and u8 +- `test_pa_kv_cache_update([32*1024], [4*1024], ..., k_head_size=256, v_head_size=256, ...)` for fp16 and u8 +- Also fixed a pre-existing bug: line 482 was missing `enable_kvcache_compress=compress_kvcache`, causing it to always default to True. + +**Performance (memory-bound, 32K tokens, 8 kv_heads, past_lens=0):** + +| Config | Data (MB) | BW (GB/s) | +|--------|-----------|-----------| +| hs128, fp16 | 268 | ~114 | +| hs256, fp16 | 537 | ~108 | +| hs128, u8 | ~203 | ~88 | +| hs256, u8 | ~405 | ~90 | + +Near-perfect scaling: 2x data with roughly the same bandwidth. No regression. + +**Suggestion:** No further work needed. This kernel is simple scatter-write with good memory access patterns. + +--- + +## test_pa_decoding.py (pa_single_token.cm: cm_sdpa_2nd + cm_sdpa_2nd_reduce) + +**Problem:** No compilation error, but severe u8 performance regression. The `#pragma unroll` on the V inner loop (`for k=0; k 128`: + ```cpp + #if HEAD_SIZE <= 128 + #pragma unroll + #endif + ``` + This lets the compiler use a runtime loop for hs256, reducing register live ranges from 16 iterations to 1. Eliminates the 2368-byte spill entirely. u8/by_token registers drop from 256 to 175. + +2. **test_pa_decoding.py line 166:** Increase `reduce_split_step` to 16 for hs256: + ```python + self.reduce_split_step = 16 if head_size >= 256 else 8 + ``` + This keeps `REDUCE_SPLIT_SIZE=16` (same as hs128), halving GWS_2 dispatch from 32 to 16 workitems. + +All 7 functional tests pass (4 hs128 + 3 hs256). Added perf benchmark extension in the `test_pa_perf_bandwidth_generate_single_subsequence_default_params` test (gated by `RUN_PA_PERF=1`). + +**Performance (memory-bound, kv_len=32769, 32 heads, 8 kv_heads):** + +*Before optimization:* + +| Config | sdpa_2nd BW (GB/s) | reduce BW (GB/s) | sdpa_2nd (ms) | reduce (ms) | Spill | +|--------|--------------------|------------------|---------------|-------------|-------| +| hs128, fp16 | 105 | 40 | 1.28 | 0.054 | 0 | +| hs128, u8/by_token | 97 | 39 | 0.70 | 0.055 | 0 | +| hs128, u8/by_channel | 97 | 39 | 0.70 | 0.055 | 0 | +| hs256, fp16 | 105 | 39 | 2.58 | 0.108 | 0 | +| hs256, u8/by_token | 58 | 41 | 2.32 | 0.103 | **2368 B** | +| hs256, u8/by_channel | 54 | 41 | 2.51 | 0.104 | **2368 B** | + +*After optimization:* + +| Config | sdpa_2nd BW (GB/s) | reduce BW (GB/s) | sdpa_2nd (ms) | reduce (ms) | Spill | +|--------|--------------------|------------------|---------------|-------------|-------| +| hs128, fp16 | 105 | 40 | 1.28 | 0.054 | 0 | +| hs128, u8/by_token | 97 | 39 | 0.70 | 0.055 | 0 | +| hs128, u8/by_channel | 97 | 39 | 0.70 | 0.055 | 0 | +| hs256, fp16 | 105 | **62** | 2.59 | **0.069** | 0 | +| hs256, u8/by_token | **101** | **63** | **1.34** | **0.068** | **0** | +| hs256, u8/by_channel | **101** | **64** | **1.34** | **0.067** | **0** | + +**End-to-end improvement (sdpa_2nd + reduce total):** +- **hs256 fp16:** 2.689 -> 2.655 ms (-1.3%, reduce-only gain) +- **hs256 u8/by_token:** 2.424 -> 1.410 ms (**-42%**) +- **hs256 u8/by_channel:** 2.614 -> 1.408 ms (**-46%**) +- **hs128:** zero regression across all configs + +**Suggestion:** +- The fp16 path (no dequant) is already optimal — 105 GB/s matches hs128's bandwidth ceiling on LNL. +- For further u8 gains, consider applying the same no-unroll technique to the K dequant loop (lines 231-236 in pa_single_token.cm) which also iterates HEAD_SIZE/REG_K times with `#pragma unroll`. + +--- + +## test_pa.py (pa_multi_token.cm -> cm_pa_xe2.hpp: pa_lsc_u8 + pa_kernel_lsc_prefetch_f16) + +**Problem:** CM compilation error. The accumulator `rO` matrix hits exactly 16384 bytes: +``` +matrix rO; +// HEAD_SIZE=256: 256/16*2=32 rows, 8*16=128 cols, 32*128*4 = 16384 bytes +// FAILS: must be strictly < 16384 +``` +This affects 4 code paths: 2 functions (`pa_lsc_u8`, `pa_kernel_lsc_prefetch_f16`) x 2 pipeline variants (optimized sparse, legacy per-step). + +**Fix:** Split `rO` into `rO_lo` + `rO_hi`, each covering half of head_size (8192 bytes each): +```cpp +constexpr int rO_half_rows = head_size / 2 / REG_N * num_P_tiles; +matrix rO_lo; // [0, head_size/2) +matrix rO_hi; // [head_size/2, head_size) +``` + +All 36 access sites updated: +1. **SLM-based PV (pa_lsc_u8):** Call `ugemm_PV0`/`ugemm_PV1` twice with adjusted SLM V offset `(head_size/2)*REG_K*sizeof(half)`. +2. **Inline DPAS PV (pa_kernel_lsc_prefetch_f16):** Split `for(k=0; k `rO_lo`, `[head_size/2, head_size)` -> `rO_hi`. +3. **Output store:** Two loops writing `rO_lo` then `rO_hi` via b2dO. +4. **Compiler flags:** Removed `-abortonspill` for HEAD_SIZE > 128 since rQ(128) + rO_lo(128) + rO_hi(128) = 384 registers > 256 GRF. Expected spill ~12-13 KB. HEAD_SIZE=128 retains zero-spill (no regression). + +Added accuracy tests (6 cases: fp16/u8 x sparse_block_sz=1/256/128) and perf benchmarks mirroring `smoke_perf_test()`. +Parameterized the roofline formula: `roofline_ms = 293.20 * (head_size/128) * (seq_len/32768)^2 * (num_heads/32)`. + +**Performance (compute-bound, seq_len=32768, 32 heads, 8 kv_heads, trunk=128 blocks):** + +| Config | sb | density | MFU (GFLOPS) | Latency (ms) | Meet | +|--------|----|---------|-------------|-------------|------| +| hs128, U8 | 1 | 1.00 | 17,188 | 512 | **0.57** | +| hs128, U8 | 256 | 1.00 | 17,883 | 496 | **0.60** | +| hs128, U8 | 256 | 0.33 | 17,970 | 163 | **0.60** | +| hs256, U8 | 1 | 1.00 | 2,255 | 7,800 | **0.08** | +| hs256, U8 | 256 | 1.00 | 2,664 | 6,655 | **0.09** | +| hs256, U8 | 256 | 0.33 | 2,660 | 2,199 | **0.09** | +| hs128, FP16 | 1 | 1.00 | 18,208 | 483 | **0.61** | +| hs128, FP16 | 256 | 1.00 | 18,109 | 490 | **0.60** | +| hs128, FP16 | 256 | 0.33 | 18,719 | 156 | **0.62** | +| hs256, FP16 | 1 | 1.00 | 3,639 | 4,835 | **0.12** | +| hs256, FP16 | 256 | 1.00 | 3,531 | 5,021 | **0.12** | +| hs256, FP16 | 256 | 0.33 | 3,789 | 1,544 | **0.13** | + +HEAD_SIZE=256 achieves meet 0.08-0.13 vs 0.57-0.62 for HEAD_SIZE=128. This is functionally correct but ~5-7x below ideal scaling. + +### Register usage breakdown (pa_kernel_lsc_prefetch_f16, hs256, Xe2 256-GRF mode) + +**Long-lived registers (persist across all kv_step iterations):** + +| Matrix | Shape | Element | Bytes | GRF regs | Lifetime | +|--------|-------|---------|-------|----------|----------| +| rQ | [16, 256] | half | 8,192 | 128 | Entire kernel (loaded once, read every KQ DPAS) | +| rO_lo | [16, 128] | float | 8,192 | 128 | Entire kernel (accumulated every kv_step) | +| rO_hi | [16, 128] | float | 8,192 | 128 | Entire kernel (accumulated every kv_step) | +| cur_max | [16] | float | 64 | 1 | Online softmax state | +| cur_sum | [16] | float | 64 | 1 | Online softmax state | +| **Total persistent** | | | **24,704** | **386** | | + +**Budget = 256 regs. Oversubscription = 130 regs (~8.3 KB spill minimum).** + +The fundamental constraint: `rQ(128) + rO_lo(128) + rO_hi(128) = 384` regs of persistent state, before any temporaries. + +**Per-iteration temporaries (inside kv_step loop):** + +| Phase | Matrix | Shape | Regs | Notes | +|-------|--------|-------|------|-------| +| KQ | Kmat | [2, 128] half | 8 | Overwritten each ri iteration | +| KQ | St | [16, 16] float | 16 | Reused in softmax | +| Softmax | max_comp | [16] float | 1 | | +| Softmax | P (Transpose output) | [16, 16] half | 8 | +8 temp inside Transpose | +| PV | Vmat | [8, 32] half | 8 | Overwritten each k iteration | +| Misc | b2dK/b2dV/addresses | — | ~14 | Descriptors, loop counters | +| **Total temporaries** | | | **~55** | | + +**Peak pressure (during PV1 phase) = 386 + 55 = ~441 regs.** The compiler spills ~12-13 KB, rotating rQ and rO tiles through scratch memory every DPAS iteration. + +For comparison, hs128: rQ(64) + rO(128) + temps(~55) = 247 regs < 256. **Zero spill.** + +**Root causes of 5-7x MFU drop:** +1. **Register spill (~12-13 KB):** Adds scratch memory traffic every DPAS iteration. HEAD_SIZE=128 has zero spill. +2. **Split rO double-read:** V data is read from SLM (or loaded via LSC) twice per kv_step — once for rO_lo, once for rO_hi — instead of being reused in a single pass. +3. **Instruction cache pressure:** `#pragma unroll` on V loops with 2x iterations (16 for hs256 vs 8 for hs128) bloats the instruction footprint. + +### Why OCL micro kernel has lower register pressure at hs256 + +The OCL `sdpa_micro__prefill` kernel peaks at ~180 regs for h256, vs CM PA's ~441. Three key design differences: + +**1. SLM decouples KQ and VS register lifetimes (max vs sum):** +- Micro kernel: Q lives in Q_slm (not GRF). KQ result S is written to S_slm after softmax. VS reads S from S_slm. Q_tile is transient (freed before VS). Peak = `max(KQ_regs, VS_regs)` = max(110, 180) = **180**. +- CM PA: rQ lives in GRF the entire kernel. rO lives in GRF the entire kernel. Peak = `KQ_regs + VS_regs` = 128 + 256 + 55 = **441**. + +**2. Head_size is tiled across 8 SG rows (not 1 thread):** +- Micro kernel VS: `wg_m_vs=8` SG rows, each accumulates `sg_tile_m=32` of 256 head dims → A_tile = 64 GRFs. +- CM PA: each thread accumulates all 256 head dims → rO = 256 GRFs. + +**3. Fewer q tokens per work unit:** +- Micro kernel h256: 2 q-tokens per lane (vs CM's 16 per thread), smaller S_tile and A_tile. + +### Lessons for CM PA optimization + +| Approach | rQ | rO | Peak | Spill | DPAS overhead | SLM cost | Occupancy | U8+FP16? | Effort | +|----------|----|----|------|-------|---------------|----------|-----------|----------|--------| +| **Current CM PA** | 128 (GRF, permanent) | 256 (split) | ~441 | ~12 KB | None | None | 2 WG | Yes | — | +| **Q from L3 cache** | 4 (transient) | 256 (split) | ~315 | ~2-3 KB | None | None | 2 WG | **Yes** | Low | +| **Move Q to SLM** | 0 (SLM) | 256 (split) | ~311 | ~2 KB | None | 128 KB | **1 WG** | **FP16 only** | Medium | +| **num_P_tiles=1** | 128 (GRF) | 128 | ~311 | ~2 KB | **+50%** | None | 2 WG | Yes | Medium | +| **Q from L3 + num_P_tiles=1** | 4 (transient) | 128 | ~187 | 0 | +50% | None | 2 WG | Yes | Medium | +| **Q to SLM + num_P_tiles=1** | 0 (SLM) | 64 | ~119 | 0 | +50% | 128 KB | 1 WG | FP16 only | Medium-High | +| **Full micro-style redesign** | SLM | 64/SG | ~180 | 0 | None | Full | 1 WG | Yes | Very High | + +### Detailed feasibility analysis of each approach + +**Move Q to SLM — SLM budget problem:** + +Each of the 16 threads has unique Q data (different q_start positions), so Q cannot be shared. SLM cost = 16 threads × 8 KB = **128 KB**. + +| Path | K_SLM (ring×4) | V_SLM (ring×4) | Q_SLM (16 threads) | Total | Limit | Fits? | +|------|---------------|----------------|---------------------|-------|-------|-------| +| U8, hs256 | 32 KB | 32 KB | 128 KB | **192 KB** | 128 KB | **No** | +| FP16, hs256 | 0 | 0 | 128 KB | 128 KB | 128 KB | Barely | + +- **U8 path: impossible** (192 KB > 128 KB SLM). +- **FP16 path: exactly fills SLM** (128 KB). But this forces 1 WG per Xe Core (128 KB SLM per core), halving occupancy from 2 WGs to 1 WG. For a compute-bound kernel, reduced occupancy hurts latency hiding. + +**num_P_tiles=1 — DPAS waste problem:** + +REG_N=16 is hardware-fixed. KQ DPAS always produces 16 output columns regardless of how many q tokens are consumed by PV. With num_P_tiles=1, each thread computes 16 KQ columns but PV only uses 8 → **50% KQ DPAS wasted**. + +System-wide DPAS to process 16 q tokens (same work as current single thread): + +| Phase | num_P_tiles=2 (1 thread) | num_P_tiles=1 (2 threads) | +|-------|--------------------------|---------------------------| +| KQ | 32 DPAS | **64 DPAS** (+100%) | +| PV | 32 DPAS | 32 DPAS (same) | +| **Total** | 64 DPAS | **96 DPAS** (+50%) | + +For a kernel already at MFU 0.08-0.13, adding 50% more DPAS is costly. + +**Q from L3 cache — the best standalone option:** + +Instead of keeping rQ permanently in GRF (128 regs), reload Q tiles from global memory per KQ iteration, relying on L3 cache: + +```cpp +// Current: rQ permanent in GRF (128 regs) +dpas(St, rQ[ri].format(), Kmat); + +// Proposed: Qtile loaded per iteration from L3 (~4 regs transient) +cm_load(Qtile, b2dQ.set_block_x(ri*REG_K)); // Q from L3 +dpas(St, Qtile.format(), Kmat); +``` + +Why L3 re-reads are near-free: +- Q is only 8 KB per thread — fits easily in 8 MB L3 +- Q is read-only, same access pattern every kv_step → stays hot in L3 +- L3 BW ~500 GB/s. Total Q re-reads at 32K kv_len: 2048 iters × 8 KB = 16 MB/thread, 32 threads = 512 MB. At 500 GB/s → ~1 ms vs ~500 ms kernel. **<0.2% overhead** +- L3 load latency (~100 cycles) overlaps with Kmat load (different addresses, pipelined) + +Advantages: works for both U8 and FP16, zero SLM cost, zero occupancy impact, zero DPAS waste, minimal code change. + +--- + +## Optimization Experiments (multi-token PA, hs256, seq=32K, 32h/8kvh) + +### Baseline performance + +| Config | Spill | MFU (GFLOPS) | Latency (ms) | Meet | +|--------|-------|-------------|-------------|------| +| hs256, FP16, sb1 | 12,736 B | 3,639 | 4,835 | 0.12 | +| hs256, U8, sb1 | 12,544 B | 2,255 | 7,800 | 0.08 | +| hs128, FP16, sb1 | 0 | 18,208 | 483 | 0.61 | +| hs128, U8, sb1 | 0 | 17,188 | 512 | 0.57 | + +hs256 achieves meet 0.08-0.12 vs 0.57-0.61 for hs128 (5-7x MFU gap). + +### Experiment 1: Q from L3 cache — FAILED, REVERTED + +**Change:** Replaced permanent `rQ[16, 256]` half (128 regs) with per-KQ-iteration `cm_load` of transient `Qtile[1, REG_K*REG_N]` (~4 regs). Added `ugemm_KQ_L3` to `cm_attention_common.hpp`. Modified both `pa_lsc_u8` and `pa_kernel_lsc_prefetch_f16`. + +| Config | Spill | MFU (GFLOPS) | Delta vs baseline | +|--------|-------|-------------|-------------------| +| hs256, FP16, sb1 | 4,928 B (-61%) | 2,635 | **-28%** | +| hs256, U8, sb1 | 4,352 B (-65%) | 1,894 | **-16%** | +| hs128, FP16 | 0 | ~18,100 | ~0% | +| hs128, U8 | 0 | ~17,800 | ~0% | + +**Why it failed:** Each KQ iteration loads 16 Q tiles from L3 via `cm_load` + `cm_mul(scale_factor)`. The per-tile L3 latency (~100 cycles) × 16 tiles = ~1600 cycles overhead per kv_step, which is **not overlapped** with DPAS execution because the Q tile is consumed immediately by the next DPAS. The compiler's spill strategy (strategic eviction/reload from scratch) is more efficient because it can overlap spill traffic with DPAS pipeline stages. + +**Lesson:** Reducing spill does not guarantee MFU improvement. Explicit per-iteration data movement can be worse than compiler-managed spill. + +### Experiment 2: Remove V loop unroll for hs256 — FAILED, REVERTED + +**Change:** Added `#if CMFLA_HEAD_SIZE <= 128` / `#pragma unroll` / `#endif` guards around 8 PV loops (4 in optimized sparse pipeline, 4 in legacy pipeline) in both `pa_lsc_u8` and `pa_kernel_lsc_prefetch_f16`. + +| Config | Spill | Delta vs baseline | +|--------|-------|-------------------| +| hs256, FP16 | **20,928 B** (+64%) | Spill **increased** | +| hs256, U8 | 12,544 B (~0%) | No change | + +**Why it failed:** The V loops in `cm_pa_xe2.hpp` are already split into two halves (rO_lo and rO_hi), each iterating only `head_size/2/REG_N = 8` times. With unroll, the compiler can interleave DPAS scheduling and register allocation across all 8 iterations efficiently. Without unroll, the compiler must handle a runtime loop, losing the ability to overlap DPAS latencies — actually **increasing** register pressure. This contrasts with `pa_single_token.cm` where the same technique worked because its V loop has 16 iterations (HEAD_SIZE/REG_N) with many more registers live simultaneously. + +**Lesson:** Unroll removal helps when loop iteration count is large (16+) and unrolling causes excessive simultaneous register liveness. It hurts when loops are short (8 iterations) and the compiler benefits from seeing all iterations for scheduling. + +### Experiment 3: Tile-based O accumulation — INFEASIBLE (not implemented) + +**Concept:** Use single `rO` (128 regs) and process head_size in two sequential passes. + +**Option A: Two passes over KV sequence.** Each pass iterates all kv_steps, computing KQ + PV for one half of head_size. KQ DPAS produces the same St in both passes → **100% KQ DPAS duplication**. Total DPAS per kv_step: (KQ(32) + PV(16)) × 2 = 96 vs current 64 → **+50% system DPAS**. For a compute-bound kernel, this is unacceptable. + +**Option B: Per-iteration rO swap through SLM/scratch.** +- SLM: 16 threads × 16 KB (rO_lo + rO_hi) = **256 KB > 128 KB SLM limit**. Infeasible. +- Global/L3 scratch: Requires host-allocated scratch buffer (new kernel arg). Adds 4 L3 operations per kv_step (2 loads + 2 stores of 8 KB each). Given that Experiment 1 showed per-iteration L3 loads degrade MFU despite reducing spill, this approach would likely also regress. + +**Not attempted** due to both variants being structurally worse than baseline. + +### Remaining suggestions (not yet tested) + +4. **Reduce V re-read via register caching (FP16 path).** Cache Vmat tiles across rO_lo/rO_hi to avoid double-loading. Adds ~8 regs pressure per Vmat but halves V load traffic. **Risk:** 8 extra regs may increase spill given already-tight 256 GRF budget. Medium priority. + +5. **U8 SLM working set tiling.** Tile head_size in `load_slm_KV` to reduce SLM working set. **Currently unnecessary** — hs256 U8 uses 64 KB SLM (half of 128 KB limit). + +6. **num_P_tiles=1 (halve q/thread).** +50% system DPAS waste. Only viable if kernel becomes memory-bound after other optimizations. + +7. **Move Q to SLM (FP16-only).** Uses all 128 KB SLM, halves occupancy, U8 incompatible. + +### Key findings and root cause analysis + +**Spill is not the primary performance bottleneck.** Q-from-L3 reduced spill by 63% but MFU regressed 16-28%. The compiler's spill strategy is surprisingly efficient — it overlaps scratch memory traffic with DPAS pipeline stages, whereas explicit data movement (L3 loads) introduces pipeline stalls. + +**The fundamental problem is register oversubscription structure:** rQ(128) + rO_lo(128) + rO_hi(128) = 384 persistent regs vs 256 available. Every approach to reduce this trades away either DPAS throughput (tile-based O, num_P_tiles=1), memory bandwidth (Q-from-L3), or occupancy (Q-to-SLM). The compiler's spill is already the least-cost option among these trade-offs. + +**hs256 vs hs128 MFU gap breakdown (estimated):** +1. **Register spill traffic:** ~2-3x slowdown (12-13 KB spill adds scratch BW overhead to every DPAS iteration) +2. **rO split double V-reads:** ~1.5x (FP16 path loads each Vmat twice; U8 path reads V from SLM twice) +3. **Instruction footprint:** ~1.2x (2x V loop iterations, 2x store loops) +4. Combined: ~3.6-5.4x ≈ observed 5-7x gap + +**Potential directions not yet explored:** +- **Compiler improvements:** Future CM compiler versions may better handle 384-reg kernels (smarter spill heuristics, better DPAS pipeline overlap). +- **Hybrid thread cooperation:** Like OCL micro kernel — tile head_size across multiple threads, communicate via SLM. Would require significant kernel redesign. +- **Hardware:** Xe3 adds 512-GRF mode (32 KB/thread) which would eliminate spill entirely for hs256. diff --git a/.github/skills/turboquant_single_token_analysis.md b/.github/skills/turboquant_single_token_analysis.md new file mode 100644 index 00000000..8f3844ef --- /dev/null +++ b/.github/skills/turboquant_single_token_analysis.md @@ -0,0 +1,289 @@ +# Skill: TurboQuant Single-Token PA — Cycle Analysis & Optimization + +Performance analysis of `pa_single_token_turboquant.cm` on LNL Xe2, with precise per-phase cycle budgets and optimization roadmap. Depends on platform constants from `lnl_xe2_pa_kernel_design.md`. + +--- + +## Kernel Overview + +File: `src/plugins/intel_gpu/src/graph/impls/cm/pa_single_token_turboquant.cm` + +Single-token (decode) paged attention with TurboQuant 4-bit KV cache compression. Each thread processes one KV partition (256 tokens) for one group of Q heads. The kernel has five phases: + +1. **Q rotation** — scalar matmul `Q_rot = Q × tq_q_t` +2. **KQ** — 4-bit centroid dequant of K + DPAS `Q × K^T` +3. **Softmax** — partition-local softmax of attention logits +4. **PV** — uint8 dequant of V + DPAS `P × V` +5. **Output** — normalize by partition sum, write partial output + LSE + +--- + +## Reference Configuration + +Qwen3-like model: HEAD_SIZE=128, 32 Q heads, 8 KV heads, batch=1, context=32K. + +| Parameter | Value | Derivation | +|-----------|-------|------------| +| Q_head_chunk_size | 4 | GQA ratio 32/8 = 4, fits ≤ 8 MaxRepeatCount | +| REG_M (RepeatCount) | 1 | Single-token decode | +| REG_N | 16 | Xe2 CM_GRF_WIDTH / 32 | +| REG_K | 16 | SystolicDepth(8) × VNNI_WIDTH(2) | +| KV_PARTITION_SIZE | 256 | = KV_BLOCK_SIZE (xattn path) | +| KV_STEP | 16 | = REG_K | +| kv_steps/partition | 16 | 256 / 16 | +| ri iterations (HEAD_SIZE/REG_K) | 8 | 128 / 16 | +| Partitions (32K) | 128 | 32768 / 256 | +| Total threads | 1024 | 1 × 8 kv_heads × 128 partitions | +| HW threads (256-GRF) | 256 | 8 cores × 8 VE × 4 threads | +| Waves | 4 | 1024 / 256 | + +### Per-Token KV Data Size + +| Cache | Per-token/head bytes | Layout | +|-------|---------------------|--------| +| K (4-bit packed + fp16 norm) | 66 B | `HEAD_SIZE*4/8 + 2 = 66` | +| V (uint8 + fp16 scale/zp) | 132 B | `HEAD_SIZE*1 + 4 = 132` | +| **Total K+V** | **198 B** | 2.6× compression vs FP16 (256+256=512 B) | + +--- + +## Per-Thread Cycle Budget (One Partition = 256 Tokens) + +### Phase 1: Q Rotation (lines 87–98) + +```cpp +for (int qi = 0; qi < Q_head_chunk_size; qi++) // 4 + for (int j = 0; j < HEAD_SIZE; j++) // 128 + for (int i = 0; i < HEAD_SIZE; i++) // 128 + acc += (float)Qmat[qi][i] * (float)tq_q_t[i * HEAD_SIZE + j]; +``` + +- **65,536 scalar MACCs** (4 × 128 × 128). +- All MACCs are scalar: `tq_q_t[i*HEAD_SIZE+j]` is pointer-indexed, no SIMD possible. +- Scalar cost = 16× SIMD → **1,048,576 SIMD-equiv cycles**. +- This is **identical across all 128 partition threads** for the same (seq, head) — pure redundancy. + +### Phase 2: KQ — K Centroid Dequant + DPAS (lines 149–191) + +Per (kv_step, ri) tile — inner dequant loop (lines 164–178): + +```cpp +for (int p = 0; p < REG_K/2; p += 2) // 4 iterations + for (int n = 0; n < REG_N; n++) // 16 iterations + // 4 scalar LUT lookups: centroid_f[nibble & 0x0f], centroid_f[nibble >> 4] × 2 rows + // 4 scalar muls: deq * knorm_vec[n] +``` + +- Per (kv_step, ri): 4 × 16 × (4 lookups + 4 muls) = **512 scalar ops × 16 = 8,192 SIMD-equiv cycles**. +- DPAS: 1 instruction per (kv_step, ri) = RepeatCount(4) × SystolicDepth(8) × REG_N(16) = 512 MACCs. +- `rS += rS_data`: 1 SIMD add (4×16 float = 4 SIMD ops). + +Total over partition (16 kv_steps × 8 ri): + +| Component | Count | SIMD-equiv cycles | +|-----------|-------|-------------------| +| Centroid LUT + knorm mul (scalar) | 16 × 8 × 512 scalar ops | **1,048,576** | +| DPAS | 128 instructions | ~3,840 XMX cycles | +| rS accumulate | 128 × 4 SIMD ops | 512 | + +### Phase 3: Softmax (lines 200–228) + +Per Q head (×4 heads): + +| Op | SIMD instructions | Latency (MACC-equiv) | +|----|-------------------|----------------------| +| `cm_mul` (scale_factor) | 16 | 16 | +| `cm_max` reduction (256 elements) | ~8 | 8 | +| sub + mul(log2e) | 2 × 16 | 32 | +| `cm_exp` | 16 | 16 × ~6 = 96 | +| fp32→fp16 convert | 16 | 16 | +| `cm_sum` (fp32) | ~8 | 8 | +| `cm_log` | 1 | ~6 | +| `cm_sum` (fp16 Pmat) | ~8 | 8 | +| **Per Q head** | | **~190** | + +Total: 4 × 190 = **760 SIMD-equiv cycles**. + +### Phase 4: PV — V Dequant + DPAS (lines 260–316) + +Per (kv_step, ri) tile — V dequant (lines 276–291): + +| Op | SIMD instructions | Notes | +|----|-------------------|-------| +| uint8→fp16 convert (`VmatNormal = Vt_quant`) | 16 | REG_K × REG_N elements / SIMD16 | +| sub zp (16 rows) | 16 | `VmatNormal[r] - temp_zp[r]` (broadcast) | +| mul scale (16 rows) | 16 | `VmatNormal[r] * temp_scale[r]` (broadcast) | +| VNNI repack (2 selects) | ~4 | Register shuffles | +| **Per (kv_step, ri)** | **52** | All SIMD — no scalar | + +DPAS: 1 instruction per (kv_step, ri) = 512 MACCs. + +Total over partition (16 kv_steps × 8 ri): + +| Component | SIMD-equiv cycles | +|-----------|-------------------| +| V dequant + repack (SIMD) | 16 × 8 × 52 = **6,656** | +| DPAS | 128 instructions = ~3,840 XMX cycles | + +### Phase 5: Output (lines 320–334) + +- `cm_div_ieee`: 4 Q_heads × 8 tiles × ~6 cycles = **192 SIMD-equiv cycles**. +- SVM writes: negligible vs compute. + +### Per-Thread Total + +| Phase | SIMD-equiv cycles | DPAS MACCs | Dominant Unit | +|-------|-------------------|------------|---------------| +| Q rotation | **1,048,576** | 0 | Scalar | +| K centroid dequant | **1,048,576** | 0 | Scalar | +| KQ DPAS | 512 | 65,536 | XMX | +| Softmax | 760 | 0 | XVE | +| V dequant + repack | 6,656 | 0 | XVE | +| PV DPAS | 0 | 65,536 | XMX | +| Output | 192 | 0 | XVE | +| **Total** | **2,105,272** | **131,072** | **Scalar (99.6%)** | + +--- + +## System-Wide Timing + +``` +256 HW threads, 4 waves. +64 VEs total (8 cores × 8 VEs), clock = 2 GHz. +4 threads per VE per wave. + +XVE/Scalar time per wave: + Per-VE: 4 threads × 2,105,272 = 8,421,088 SIMD-equiv cycles + Wall time per wave = 8,421,088 / 2 GHz = 4.21 ms + 4 waves = ~16.8 ms + +XMX time per wave: + Per-thread: 256 DPAS instructions × ~30 cycles = 7,680 cycles + Per-VE: 4 × 7,680 = 30,720 cycles → 15.4 μs + 4 waves = ~62 μs + +Memory time: + K: 32768 × 66 B × 8 heads = 16.5 MB + V: 32768 × 132 B × 8 heads = 33.0 MB + Q + tq_q_t: ~72 KB (cached) + Output (fp32): 128 × 32 × 128 × 4 = 2.0 MB + Total ≈ 51.5 MB / 102 GB/s = 0.505 ms +``` + +| Ceiling | Time | % of Wall Time | +|---------|------|----------------| +| **Memory floor** | **0.5 ms** | ~3% | +| **XMX (DPAS)** | **0.06 ms** | ~0.4% | +| **XVE/Scalar** | **~16.8 ms** | **~97%** | + +**Verdict: Massively scalar-bound.** The two scalar bottlenecks (Q rotation + centroid LUT dequant) each contribute ~1M SIMD-equiv cycles per thread, dwarfing both the memory floor (0.5 ms) and DPAS time (0.06 ms) by over 30×. + +--- + +## Optimization Roadmap + +### P0: Pre-Rotate Q in Separate Kernel + +**Problem**: Q rotation is 1M SIMD-equiv cycles/thread × 128 partitions, but the result is identical for all partitions of the same `(seq, head_num_idx)`. 128× pure redundancy. + +**Solution**: Compute `Q_rot = Q × tq_q_t` once per head group in a dedicated DPAS kernel before the attention dispatch. + +``` +Q: [4, 128] × tq_q_t: [128, 128] → standard matmul, fits DPAS perfectly +DPAS cost: 4 × (128/16) × (128/16) = 256 DPAS instructions per head group +8 head groups → 2048 DPAS total → < 50 μs on XMX +``` + +**Implementation**: +1. New small CM kernel: `tq_q_rotation.cm` — loads Q and tq_q_t, DPAS matmul, writes Q_rot. +2. Remove lines 84–99 from `pa_single_token_turboquant.cm`. +3. Pass pre-rotated Q pointer instead of raw Q + tq_q_t. + +**Impact**: Eliminates ~50% of total wall time (1M out of 2.1M SIMD-equiv cycles per thread removed, plus no tq_q_t load overhead). + +### P1: Vectorize Centroid LUT via SIMD Compare+Select + +**Problem**: `centroid_f[nibble_idx]` is scalar-indexed — each element requires an independent register gather. 512 scalar ops × 16 = 8,192 SIMD-equiv cycles per (kv_step, ri) tile. + +**Solution**: Replace scalar indexed lookup with SIMD compare+select over 16 codebook entries: + +```cpp +// Current (scalar, 256 scalar ops per p-iteration): +for (int n = 0; n < REG_N; n++) { + float deq0_lo = centroid_f[(int)(packed_idx0 & 0x0f)]; // scalar gather + ... +} + +// Proposed (SIMD, 16 SIMD ops per nibble-vector): +vector nibble = packed_row & 0x0f; // SIMD mask +vector result = centroid_f[0]; // broadcast entry 0 +#pragma unroll +for (int c = 1; c < 16; c++) { + vector mask = (nibble == c); + result = cm_sel(mask, centroid_f[c], result); // SIMD select +} +// result now has centroid_f[nibble[n]] for all n in SIMD — no scalar path +``` + +Per (kv_step, ri): 4 nibble-vectors × 16 compare+select = 64 SIMD ops + 4 SIMD knorm muls = **68 SIMD-equiv cycles** vs current 8,192. That's **120× speedup** for the K dequant inner loop. + +Total K dequant: 16 × 8 × 68 = **8,704 SIMD-equiv cycles** (down from 1,048,576). + +**Alternative (P1a)**: Precompute dequanted K in a separate kernel, write fp16 K to temp buffer. Attention kernel reads fp16 directly (like non-TQ path). Extra memory traffic: 32768 × 128 × 2 × 8 = 64 MB → 0.63 ms. Still much faster than 8.4 ms scalar cost. Simpler to implement but uses more BW. + +### P2: Increase KV_PARTITION_SIZE to 512 + +**Problem**: 256-token partitions → 128 partitions → 4 waves. Each partition re-reads the Q matrix and re-initializes rS/Pmat/Omat. Also produces 128 partial results for the reduction kernel. + +**Solution**: Double partition to 512 tokens. + +Register budget check: +``` +rS = [4, 512] float = 8,192 B (< 16,384 CM limit ✓) +Pmat = [4, 512] half = 4,096 B +Omat = [4, 128] float = 2,048 B +Qmat = [4, 128] half = 1,024 B +Total persistent ≈ 15,360 B → fits 256-GRF (16,384 B budget) +``` + +Impact: +- Partitions halve: 64 → 2 waves instead of 4. +- Reduction work halves. +- After P0, also halves the number of Q loads. + +### P3: Fuse K-Norm Multiply into SIMD Path + +After P1 vectorizes the centroid lookup, the `knorm_vec[n]` multiply (lines 174–177) should also be vectorized: + +```cpp +// After SIMD centroid lookup produces result[0..15]: +result = cm_mul(result, knorm_vec); // SIMD broadcast-multiply +Kt[row] = (half)result; // SIMD convert +``` + +This is essentially free once P1 is done — just ensuring the multiply stays in SIMD rather than falling back to scalar. + +### P4: Software Prefetch for K/V + +Once scalar bottlenecks are removed, the kernel becomes memory-bound at ~0.5 ms. Overlap K and V loads: + +- During KQ phase: `cm_prefetch` V block for current partition +- During PV phase of block[i]: `cm_prefetch` K block[i+1] (if partition > 1 block) + +### P5: fp16 Partial Output + +Output writes are fp32 (line 330): 128 × 32 × 128 × 4 = 2 MB. Switch to fp16 halves this to 1 MB. The reduction kernel applies log-sum-exp correction anyway, so fp16 partials have negligible precision impact. + +--- + +## Projected Impact + +| State | Est. Wall Time | Bottleneck | Speedup | +|-------|---------------|------------|---------| +| **Current** | **~16.8 ms** | Scalar (Q rot + centroid) | 1× | +| After P0 (pre-rotate Q) | ~8.5 ms | Scalar (centroid dequant) | 2× | +| After P0 + P1 (SIMD centroid) | ~0.55 ms | Memory | 30× | +| After P0 + P1 + P2 (512 partition) | ~0.5 ms | Memory (at BW ceiling) | 34× | +| **Memory floor** | **0.505 ms** | 51.5 MB / 102 GB/s | 33× | + +P0 + P1 together achieve a **~30× speedup**, bringing the kernel from scalar-bound to within touching distance of the memory bandwidth ceiling. P2 closes the remaining gap. diff --git a/.gitignore b/.gitignore index adae0e80..47352f5e 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ a.out .vscode ocl.venv .venv +CM_spec diff --git a/opencl/tests/pageatten/cm_pa_xe2.hpp b/opencl/tests/pageatten/cm_pa_xe2.hpp index 7d829e4f..f67e901d 100644 --- a/opencl/tests/pageatten/cm_pa_xe2.hpp +++ b/opencl/tests/pageatten/cm_pa_xe2.hpp @@ -77,7 +77,10 @@ void pa_lsc_u8( constexpr int num_P_tiles = REG_N / REG_M; matrix rQ; - matrix rO; + constexpr int rO_half_rows = head_size / 2 / REG_N * num_P_tiles; + static_assert(head_size % (2 * REG_N) == 0, "head_size must be divisible by 2*REG_N for rO split"); + matrix rO_lo; + matrix rO_hi; bool first_active = true; auto q_tokens_left = q_len; @@ -350,11 +353,14 @@ void pa_lsc_u8( matrix P; Transpose2DMatrix(St, P); + constexpr uint slm_V_hi_offset = (head_size / 2) * REG_K * sizeof(half); if (first_active) { - ugemm_PV0(slm_V, P, rO, slm_offset); + ugemm_PV0(slm_V, P, rO_lo, slm_offset); + ugemm_PV0(slm_V, P, rO_hi, slm_offset + slm_V_hi_offset); first_active = false; } else { - ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); + ugemm_PV1(slm_V, P, max_comp, rO_lo, slm_offset); + ugemm_PV1(slm_V, P, max_comp, rO_hi, slm_offset + slm_V_hi_offset); } } } @@ -541,11 +547,14 @@ void pa_lsc_u8( matrix P; Transpose2DMatrix(St, P); + constexpr uint slm_V_hi_offset_legacy = (head_size / 2) * REG_K * sizeof(half); if (first_active) { - ugemm_PV0(slm_V, P, rO, slm_offset); + ugemm_PV0(slm_V, P, rO_lo, slm_offset); + ugemm_PV0(slm_V, P, rO_hi, slm_offset + slm_V_hi_offset_legacy); first_active = false; } else { - ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); + ugemm_PV1(slm_V, P, max_comp, rO_lo, slm_offset); + ugemm_PV1(slm_V, P, max_comp, rO_hi, slm_offset + slm_V_hi_offset_legacy); } } } @@ -572,12 +581,34 @@ void pa_lsc_u8( o_pitch - 1, 0, 0); + // Store lower half of head_size from rO_lo #pragma unroll - for (int k = 0, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { + for (int k = 0, ri = 0; k < head_size / 2; k += REG_N, ri += num_P_tiles) { #pragma unroll for (int p = 0; p < num_P_tiles; p++) { - auto cO = rO[ri + p].format(); + auto cO = rO_lo[ri + p].format(); + + #pragma unroll + for (int r = 0; r < cO.n_rows(); r++) { + cur_O_f16[r + p * REG_M] = cm_mul(cO.row(r), cur_sum[r + p * REG_M]); + } + } + + b2dO.set_block_x(k); + cm_store(b2dO.set_block_y(0), + cur_O_f16.format().row(0)); + cm_store(b2dO.set_block_y(REG_M), + cur_O_f16.format().row(1)); + } + + // Store upper half of head_size from rO_hi + #pragma unroll + for (int k = head_size / 2, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { + + #pragma unroll + for (int p = 0; p < num_P_tiles; p++) { + auto cO = rO_hi[ri + p].format(); #pragma unroll for (int r = 0; r < cO.n_rows(); r++) { @@ -627,7 +658,10 @@ void pa_kernel_lsc_prefetch_f16( cur_sum = 0; constexpr int num_P_tiles = REG_N / REG_M; matrix rQ; - matrix rO; + constexpr int rO_half_rows_f16 = head_size / 2 / REG_N * num_P_tiles; + static_assert(head_size % (2 * REG_N) == 0, "head_size must be divisible by 2*REG_N for rO split"); + matrix rO_lo; + matrix rO_hi; bool first_active = true; #if SPARSE_BLOCK_SIZE > 1 @@ -766,14 +800,35 @@ void pa_kernel_lsc_prefetch_f16( b2dV.set_base_ptr((reinterpret_cast(v_cache_base)+cur_block_id*blk_stride)); b2dV.set_block_y(kv_pos%CMPA_BLOCK_SZ); if (first_active) { - // ugemm_PV0(slm_V, P, rO, slm_offset); auto P2 = P.format(); + // PV0 lower half + #pragma unroll + for(int k = 0, ri = 0; k < head_size / 2; k += REG_N, ri += num_P_tiles) { + matrix Vmat; + cm_prefetch(prefetch_V.set_block_x(k)); + cm_load(Vmat.format(), b2dV.set_block_x(k)); + if ((kv_pos + kv_step) > kv_stop) { + uint valid_rows = kv_stop - kv_pos; + uint valid_rows_vnni = (valid_rows+1)/2; + for (int r = valid_rows_vnni; r < kv_step / 2; r++) + Vmat.row(r) = 0.f; + if (valid_rows % 2 == 1) + Vmat.row(valid_rows_vnni-1).select(1) = 0.f; + } + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + rO_lo[ri + p] = cm_dpas( + 0, + Vmat.format(), + P2.row(p).format()); + } + } + // PV0 upper half #pragma unroll - for(int k = 0, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { + for(int k = head_size / 2, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { matrix Vmat; cm_prefetch(prefetch_V.set_block_x(k)); cm_load(Vmat.format(), b2dV.set_block_x(k)); - // sometimes KV cache would be filled with random Nan, so need to clean up the unused value data. if ((kv_pos + kv_step) > kv_stop) { uint valid_rows = kv_stop - kv_pos; uint valid_rows_vnni = (valid_rows+1)/2; @@ -784,7 +839,7 @@ void pa_kernel_lsc_prefetch_f16( } #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - rO[ri + p] = cm_dpas( + rO_hi[ri + p] = cm_dpas( 0, Vmat.format(), P2.row(p).format()); @@ -793,15 +848,45 @@ void pa_kernel_lsc_prefetch_f16( first_active = false; } else { - //ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); auto P2 = P.format(); + // PV1 lower half + #pragma unroll + for(int k = 0, ri=0; k < head_size / 2; k += REG_N, ri += num_P_tiles) { + matrix Vmat; + + cm_prefetch(prefetch_V.set_block_x(k)); + cm_load(Vmat.format(), b2dV.set_block_x(k)); + if ((kv_pos + kv_step) > kv_stop) { + uint valid_rows = kv_stop - kv_pos; + uint valid_rows_vnni = (valid_rows+1)/2; + for (int r = valid_rows_vnni; r < kv_step / 2; r++) + Vmat.row(r) = 0.f; + if (valid_rows % 2 == 1) + Vmat.row(valid_rows_vnni-1).select(1) = 0.f; + } + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_lo[ri + p].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + p*REG_M]); + } + + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + rO_lo[ri + p] = cm_dpas( + rO_lo[ri + p].format(), + Vmat.format(), + P2.row(p).format()); + } + } + // PV1 upper half #pragma unroll - for(int k = 0, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { + for(int k = head_size / 2, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { matrix Vmat; cm_prefetch(prefetch_V.set_block_x(k)); cm_load(Vmat.format(), b2dV.set_block_x(k)); - // sometimes KV cache would be filled with random Nan, so need to clean up the unused value data. if ((kv_pos + kv_step) > kv_stop) { uint valid_rows = kv_stop - kv_pos; uint valid_rows_vnni = (valid_rows+1)/2; @@ -810,11 +895,9 @@ void pa_kernel_lsc_prefetch_f16( if (valid_rows % 2 == 1) Vmat.row(valid_rows_vnni-1).select(1) = 0.f; } - //# compensate cur_O - // matrix rO; #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO[ri + p].format(); + auto cO = rO_hi[ri + p].format(); #pragma unroll for(int r = 0; r < REG_M; r++) cO.row(r) = cm_mul(cO.row(r), max_comp[r + p*REG_M]); @@ -822,8 +905,8 @@ void pa_kernel_lsc_prefetch_f16( #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - rO[ri + p] = cm_dpas( - rO[ri + p].format(), + rO_hi[ri + p] = cm_dpas( + rO_hi[ri + p].format(), Vmat.format(), P2.row(p).format()); } @@ -914,14 +997,35 @@ void pa_kernel_lsc_prefetch_f16( b2dV.set_base_ptr((reinterpret_cast(v_cache_base)+cur_block_id*blk_stride)); b2dV.set_block_y(kv_pos%CMPA_BLOCK_SZ); if (first_active) { - // ugemm_PV0(slm_V, P, rO, slm_offset); auto P2 = P.format(); + // PV0 lower half + #pragma unroll + for(int k = 0, ri = 0; k < head_size / 2; k += REG_N, ri += num_P_tiles) { + matrix Vmat; + cm_prefetch(prefetch_V.set_block_x(k)); + cm_load(Vmat.format(), b2dV.set_block_x(k)); + if ((kv_pos + kv_step) > kv_stop) { + uint valid_rows = kv_stop - kv_pos; + uint valid_rows_vnni = (valid_rows+1)/2; + for (int r = valid_rows_vnni; r < kv_step / 2; r++) + Vmat.row(r) = 0.f; + if (valid_rows % 2 == 1) + Vmat.row(valid_rows_vnni-1).select(1) = 0.f; + } + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + rO_lo[ri + p] = cm_dpas( + 0, + Vmat.format(), + P2.row(p).format()); + } + } + // PV0 upper half #pragma unroll - for(int k = 0, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { + for(int k = head_size / 2, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { matrix Vmat; cm_prefetch(prefetch_V.set_block_x(k)); cm_load(Vmat.format(), b2dV.set_block_x(k)); - // sometimes KV cache would be filled with random Nan, so need to clean up the unused value data. if ((kv_pos + kv_step) > kv_stop) { uint valid_rows = kv_stop - kv_pos; uint valid_rows_vnni = (valid_rows+1)/2; @@ -932,7 +1036,7 @@ void pa_kernel_lsc_prefetch_f16( } #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - rO[ri + p] = cm_dpas( + rO_hi[ri + p] = cm_dpas( 0, Vmat.format(), P2.row(p).format()); @@ -941,15 +1045,45 @@ void pa_kernel_lsc_prefetch_f16( first_active = false; } else { - //ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); auto P2 = P.format(); + // PV1 lower half + #pragma unroll + for(int k = 0, ri=0; k < head_size / 2; k += REG_N, ri += num_P_tiles) { + matrix Vmat; + + cm_prefetch(prefetch_V.set_block_x(k)); + cm_load(Vmat.format(), b2dV.set_block_x(k)); + if ((kv_pos + kv_step) > kv_stop) { + uint valid_rows = kv_stop - kv_pos; + uint valid_rows_vnni = (valid_rows+1)/2; + for (int r = valid_rows_vnni; r < kv_step / 2; r++) + Vmat.row(r) = 0.f; + if (valid_rows % 2 == 1) + Vmat.row(valid_rows_vnni-1).select(1) = 0.f; + } + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_lo[ri + p].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + p*REG_M]); + } + + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + rO_lo[ri + p] = cm_dpas( + rO_lo[ri + p].format(), + Vmat.format(), + P2.row(p).format()); + } + } + // PV1 upper half #pragma unroll - for(int k = 0, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { + for(int k = head_size / 2, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { matrix Vmat; cm_prefetch(prefetch_V.set_block_x(k)); cm_load(Vmat.format(), b2dV.set_block_x(k)); - // sometimes KV cache would be filled with random Nan, so need to clean up the unused value data. if ((kv_pos + kv_step) > kv_stop) { uint valid_rows = kv_stop - kv_pos; uint valid_rows_vnni = (valid_rows+1)/2; @@ -958,11 +1092,9 @@ void pa_kernel_lsc_prefetch_f16( if (valid_rows % 2 == 1) Vmat.row(valid_rows_vnni-1).select(1) = 0.f; } - //# compensate cur_O - // matrix rO; #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO[ri + p].format(); + auto cO = rO_hi[ri + p].format(); #pragma unroll for(int r = 0; r < REG_M; r++) cO.row(r) = cm_mul(cO.row(r), max_comp[r + p*REG_M]); @@ -970,8 +1102,8 @@ void pa_kernel_lsc_prefetch_f16( #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - rO[ri + p] = cm_dpas( - rO[ri + p].format(), + rO_hi[ri + p] = cm_dpas( + rO_hi[ri + p].format(), Vmat.format(), P2.row(p).format()); } @@ -994,15 +1126,31 @@ void pa_kernel_lsc_prefetch_f16( lsc::block_2d_desc b2dO(o_base, q_tokens_left - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); + // Store lower half of head_size from rO_lo #pragma unroll - for(int k = 0, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { + for(int k = 0, ri=0; k < head_size / 2; k += REG_N, ri += num_P_tiles) { #pragma unroll for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO[ri + p].format(); + auto cO = rO_lo[ri + p].format(); #pragma unroll for(int r = 0; r < cO.n_rows(); r++) { cur_O_f16[r + p*REG_M] = cm_mul(cO.row(r), cur_sum[r + p*REG_M]); + } + } + b2dO.set_block_x(k); + cm_store(b2dO.set_block_y(0), cur_O_f16.format().row(0)); + cm_store(b2dO.set_block_y(REG_M), cur_O_f16.format().row(1)); + } + // Store upper half of head_size from rO_hi + #pragma unroll + for(int k = head_size / 2, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_hi[ri + p].format(); + #pragma unroll + for(int r = 0; r < cO.n_rows(); r++) { + cur_O_f16[r + p*REG_M] = cm_mul(cO.row(r), cur_sum[r + p*REG_M]); } } b2dO.set_block_x(k); diff --git a/opencl/tests/pageatten/pa_single_token.cm b/opencl/tests/pageatten/pa_single_token.cm index ec072f57..06c4e505 100644 --- a/opencl/tests/pageatten/pa_single_token.cm +++ b/opencl/tests/pageatten/pa_single_token.cm @@ -488,7 +488,9 @@ extern "C" _GENX_MAIN_ void cm_sdpa_2nd( vector temp_zp = zp_vec.select(kv_pos); matrix Vt_quant; #endif + #if HEAD_SIZE <= 128 #pragma unroll + #endif #if !USE_LSC_BLOCK_2D_DESC && KV_CACHE_COMPRESSION && XE_ARCH==1 for(int k = 0, ri = 0; k < HEAD_SIZE; k += (REG_N * VALUE_TILE_NUM) * 4, ri += VALUE_TILE_NUM * 4 ) { #else diff --git a/opencl/tests/pageatten/test_kvcache_update.py b/opencl/tests/pageatten/test_kvcache_update.py index a22e0a8a..f1f0a256 100644 --- a/opencl/tests/pageatten/test_kvcache_update.py +++ b/opencl/tests/pageatten/test_kvcache_update.py @@ -563,6 +563,7 @@ def reference_kv_cache_update(kv_cache_data, cur_kv_data, past_lens, subsequence run_pa_kv_cache_update_case([128*1024], [1*1024], num_kv_heads=8, k_head_size=128, v_head_size=128, block_size=256, check_perf=True) run_pa_kv_cache_update_case([1024], [0], num_kv_heads=8, k_head_size=128, v_head_size=128, block_size=256, enable_kvcache_compress=compress_kvcache, check_perf=True) run_pa_kv_cache_update_case([129], [0], num_kv_heads=2, k_head_size=64, v_head_size=64, block_size=16, check_perf=True) + run_pa_kv_cache_update_case([32*1024], [4*1024], num_kv_heads=8, k_head_size=256, v_head_size=256, block_size=256, enable_kvcache_compress=compress_kvcache, check_perf=True) if 1: token_pairs_acc = [ diff --git a/opencl/tests/pageatten/test_pa.py b/opencl/tests/pageatten/test_pa.py index 2cc652a7..c0b9ca7e 100644 --- a/opencl/tests/pageatten/test_pa.py +++ b/opencl/tests/pageatten/test_pa.py @@ -101,8 +101,9 @@ def __init__(self, num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compr print(f"compiling {cwd} {num_heads=} {head_size=} {sparse_block_sz=}...") scale_factor = 1.0/(head_size**0.5) + abortonspill = '-Qxcm_jit_option="-abortonspill"' if head_size <= 128 else '' self.kernels = cl.kernels(src1, - (f'-cmc -Qxcm_jit_option="-abortonspill" -Qxcm_register_file_size=256 -mCM_printregusage -I{cwd}' + (f'-cmc {abortonspill} -Qxcm_register_file_size=256 -mCM_printregusage -I{cwd}' f' -DKERNEL_NAME=cm_page_attention' f" -DCMFLA_NUM_HEADS={num_heads}" f" -DCMFLA_NUM_KV_HEADS={num_kv_heads}" @@ -695,7 +696,10 @@ def test_page_attn_causal_batch1(seq_len, num_heads = 16, num_kv_heads = 16, hea else: check_close(ref, out) else: - roofline = 293.27 if compressed_kvcache != KV_CACHE_COMPRESSION_NONE else 293.20 + # Roofline latency (ms) at density=1.0 scales with (num_heads * seq_len^2 * head_size) + # Base: 293.20 ms for num_heads=32, seq_len=32768, head_size=128, fp16 + roofline_base = 293.27 if compressed_kvcache != KV_CACHE_COMPRESSION_NONE else 293.20 # base config reference + roofline = roofline_base * (head_size / 128) * (seq_len / 32768)**2 * (num_heads / 32) warmup = 5 rep = 15 latency = pa_cm.run_perf(q, k, v, approx_simple_mask, n_warmup=warmup, n_iters=rep, deterministic_block_indices=True) @@ -1100,10 +1104,31 @@ def smoke_perf_test(blocks_per_trunk = 128, compressed_kvcache = KV_CACHE_COMPRE smoke_accuracy_test(compressed_kvcache=KV_CACHE_COMPRESSION_NONE) smoke_accuracy_test(compressed_kvcache=KV_CACHE_COMPRESSION_BY_CHANNEL, sub_block_sz=DEFAULT_SUB_BLOCK_SIZE) + # HEAD_SIZE=256 accuracy tests: fp16 kvcache and compressed (by_token) kvcache + for compressed_kv in [False, True]: + test_page_attn_causal_batch1(8*1024, num_heads=2, num_kv_heads=1, head_size=256, block_sz=256, trunk_sz=128*256, compressed_kvcache=compressed_kv, sparse_block_sz=1, density=1.0, check_acc=True) + test_page_attn_causal_batch1(8*1024, num_heads=2, num_kv_heads=1, head_size=256, block_sz=256, trunk_sz=128*256, compressed_kvcache=compressed_kv, sparse_block_sz=256, density=0.33, check_acc=True) + test_page_attn_causal_batch1(8*1024, num_heads=2, num_kv_heads=1, head_size=256, block_sz=256, trunk_sz=128*256, compressed_kvcache=compressed_kv, sparse_block_sz=128, density=0.33, check_acc=True) + smoke_perf_test() smoke_perf_test(16) smoke_perf_test(compressed_kvcache=KV_CACHE_COMPRESSION_NONE) smoke_perf_test(compressed_kvcache=KV_CACHE_COMPRESSION_BY_CHANNEL, sub_block_sz=DEFAULT_SUB_BLOCK_SIZE) + # HEAD_SIZE=256 perf benchmarks (matching smoke_perf_test pattern) + def smoke_perf_test_256(blocks_per_trunk = 128, compressed_kvcache = True): + seq_len, block_sz = 32*1024, 256 + trunk_sz = blocks_per_trunk*block_sz + + test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 8, head_size = 256, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=compressed_kvcache, sparse_block_sz = 1, density=1.0, check_acc=False) + + for sparse_block_sz in [256, 128]: + for density in [1.0, 0.99, 0.66, 0.33, 0.11]: + test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 8, head_size = 256, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=compressed_kvcache, sparse_block_sz = sparse_block_sz, density=density, check_acc=False) + + smoke_perf_test_256() + smoke_perf_test_256(16) + smoke_perf_test_256(compressed_kvcache = False) + # test_ov() diff --git a/opencl/tests/pageatten/test_pa_decoding.py b/opencl/tests/pageatten/test_pa_decoding.py index 0d040d8a..4d52ae88 100644 --- a/opencl/tests/pageatten/test_pa_decoding.py +++ b/opencl/tests/pageatten/test_pa_decoding.py @@ -86,7 +86,7 @@ def __init__( self.k_partition_block_num = 1 self.kv_partition_size = int(self.block_size * self.k_partition_block_num) - self.reduce_split_step = 8 + self.reduce_split_step = 16 if head_size >= 256 else 8 max_repeat_count = 8 q_heads_per_kv_head = self.num_heads // self.num_kv_heads @@ -354,6 +354,9 @@ def _case_id(case: DecodingCase) -> str: DecodingCase(num_heads=8, num_kv_heads=2, head_size=64, block_size=256, kv_len=513, kv_cache_compression=0), DecodingCase(num_heads=8, num_kv_heads=2, head_size=64, block_size=256, kv_len=513, kv_cache_compression=1), DecodingCase(num_heads=8, num_kv_heads=2, head_size=64, block_size=256, kv_len=513, kv_cache_compression=2), + DecodingCase(num_heads=8, num_kv_heads=2, head_size=256, block_size=256, kv_len=513, kv_cache_compression=0), + DecodingCase(num_heads=8, num_kv_heads=2, head_size=256, block_size=256, kv_len=513, kv_cache_compression=1), + DecodingCase(num_heads=8, num_kv_heads=2, head_size=256, block_size=256, kv_len=513, kv_cache_compression=2), ) @@ -539,8 +542,30 @@ def test_pa_perf_bandwidth_generate_single_subsequence_default_params(): assert perf["cm_sdpa_2nd_bw_gbs"] > 0.0 assert perf["cm_sdpa_2nd_reduce_bw_gbs"] > 0.0 + # HEAD_SIZE=256 bandwidth benchmark — compare against head_size=128 + for cmpr in [0, 1, 2]: + for hs in [128, 256]: + case_cmp = DecodingCase( + num_heads=32, + num_kv_heads=8, + head_size=hs, + block_size=256, + kv_len=32769, + kv_cache_compression=cmpr, + ) + perf_cmp = _run_bandwidth_measurement(case_cmp, loop_cnt=50, warmup=5) + tag = f"hs{hs}_cmpr{cmpr}" + print( + f"[perf][{tag}] " + f"cm_sdpa_2nd_bw={perf_cmp['cm_sdpa_2nd_bw_gbs']:.3f} GB/s, " + f"cm_sdpa_2nd_reduce_bw={perf_cmp['cm_sdpa_2nd_reduce_bw_gbs']:.3f} GB/s, " + f"cm_sdpa_2nd_ms={perf_cmp['cm_sdpa_2nd_ms']:.3f}, " + f"cm_sdpa_2nd_reduce_ms={perf_cmp['cm_sdpa_2nd_reduce_ms']:.3f}" + ) + assert perf_cmp["cm_sdpa_2nd_bw_gbs"] > 0.0 + # Usage: # python -m py_compile test_pa_decoding.py # python -m pytest --collect-only -q test_pa_decoding.py # timeout 120s python -m pytest -s -q test_pa_decoding.py -vv -# timeout 120s python -m pytest -s -q test_pa_decoding.py -vv -k 'generate_only and (cmprby_token or cmprby_channel)' \ No newline at end of file +# timeout 120s python -m pytest -s -q test_pa_decoding.py -vv -k 'generate_only and (cmprby_token or cmprby_channel)'