From 95181ac48449af57dff12e903b1d63584daff67a Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Fri, 12 Jun 2026 17:28:29 +0800 Subject: [PATCH 01/16] =?UTF-8?q?[CM=20SDPA]=20Items=2013-16:=20KV=20block?= =?UTF-8?q?ing=20+=20log2e=20fold=20+=20tree=20reduction=20+=20half-width?= =?UTF-8?q?=20transpose=20=E2=80=94=209.25=20ms=20=E2=86=92=207.654=20ms?= =?UTF-8?q?=20(=E2=88=9217%)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply optimizations from aboutSHW commit 0621405 (River Li): - Item 13: fold log2e into Q pre-scale (qscale = scale_factor * log2e); removes 16 mul/tile from softmax critical path, St lands in log2 domain - Item 14: KV blocking (KV_BLK=2, default); rO rescale amortized over 2 tiles — valid amortization unlike Q-row doubling because exp count is unchanged - Item 15: online_softmax_update_tree; balanced binary tree max/sum reduction, depth log2(BLK_ROWS)=5 vs linear chain depth 31 for BLK_ROWS=32 - Item 16: transpose_St_to_P_half; cast float->half before GRF shuffle so the 4-pass select network runs at half data-path width Also includes prerequisite items already in aboutSHW but absent here: - V prefetch moved to K phase (item 2) — gives V full K-DPAS lead time - Pre-loop warm-up prefetch (item 4) — cold-start for kv_pos=0 - Unified PV path (item 5) — eliminates kv_pos==0 branch via max_comp=0 trick Combined result on 2-seq x 3432 (PTL 4xe, 16h, d=64): 9.25 ms -> 7.654 ms (-17%). Co-Authored-By: Claude Sonnet 4.6 (cherry picked from commit 334e3d9d79ea1ac25a2de652ef4da46b3a2d4361) --- .../impls/cm/include/cm_attention_common.hpp | 4 + .../graph/impls/cm/include/cm_sdpa_common.hpp | 201 ++++++++++++------ 2 files changed, 137 insertions(+), 68 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp index 9ab9b86930e235..6a1f7b4b860a3f 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp @@ -302,7 +302,11 @@ vector online_softmax_update(matrix_ref St, vector_r // Pt = torch.exp(St - new_max) constexpr float log2e = 1.4426950408889634f; +#ifdef ABLATE_NO_EXP + for(int r = 0; r < St.n_rows(); r++) St[r] = (St[r] - new_max_t)*log2e; +#else for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp((St[r] - new_max_t)*log2e); +#endif vector row_sum_t; row_sum_t = cm_add(St[0], St[1]); diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index d9c368406014c8..9e945192d03213 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -15,6 +15,69 @@ *******************************************************************************/ #include "cm_attention_common.hpp" +// Number of kv sub-tiles (kv_step rows each) processed per online-softmax update in +// sdpa_kernel_lsc_prefetch. Larger amortizes the rO rescale + softmax bookkeeping over +// more DPAS work at the cost of GRF pressure. Overridable via -DCMFLA_KV_BLK=N. +#ifndef CMFLA_KV_BLK +#define CMFLA_KV_BLK 2 +#endif + +// Flashattn-local online-softmax with tree reductions over the kv (row) dimension. +// Identical math to online_softmax_update in cm_attention_common.hpp, but the max and +// sum reductions are balanced binary trees (dependency depth log2(rows)=4 instead of the +// linear chain's depth 15), which shortens the loop-carried softmax critical path. +template +CM_INLINE vector online_softmax_update_tree(matrix_ref St, + vector_ref cur_max, + vector_ref cur_sum) { + static_assert((rows & (rows - 1)) == 0, "tree reduction needs power-of-two rows"); + vector new_max_t; + { + matrix 1 ? rows/2 : 1), cols> t; + #pragma unroll + for (int r = 0; r < rows/2; r++) t.row(r) = cm_max(St[r], St[r + rows/2]); + #pragma unroll + for (int stride = rows/4; stride > 0; stride >>= 1) + #pragma unroll + for (int r = 0; r < stride; r++) + t.row(r) = cm_max(t.row(r), t.row(r + stride)); + new_max_t = cm_max(t.row(0), cur_max); + } + + #pragma unroll + for (int r = 0; r < rows; r++) St[r] = cm_exp(St[r] - new_max_t); + + vector row_sum_t; + { + matrix 1 ? rows/2 : 1), cols> t; + #pragma unroll + for (int r = 0; r < rows/2; r++) t.row(r) = cm_add(St[r], St[r + rows/2]); + #pragma unroll + for (int stride = rows/4; stride > 0; stride >>= 1) + #pragma unroll + for (int r = 0; r < stride; r++) + t.row(r) = cm_add(t.row(r), t.row(r + stride)); + row_sum_t = t.row(0); + } + + vector max_comp; + max_comp = cm_exp(cur_max - new_max_t); + cur_sum = cm_mul(cur_sum, max_comp); + cur_sum = cm_add(cur_sum, row_sum_t); + cur_max = new_max_t; + return max_comp; +} + +// Transpose a [16,16] float score tile (kv x q) into a half [16,16] P tile (q x kv) for +// the P@V matmul. Casting float->half first lets the 16x16 shuffle network run at half +// width — roughly halving the mov count vs transposing the float tile directly. +CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { + matrix Sh; + #pragma unroll + for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); + Transpose_16x16(Sh.select<16,1,16,1>(0,0), P); +} + #ifdef CM_HAS_LSC_UNTYPED_2D //@prefetch_u8 would have duplicated decompress perf issue. comments out for now. // template @@ -442,8 +505,6 @@ void sdpa_kernel_lsc_prefetch( constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); - // round up head_size to multiple of 16 - // block_2d_desc will automatically handle the tailing block constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; vector cur_max; @@ -455,7 +516,7 @@ void sdpa_kernel_lsc_prefetch( matrix rQ; matrix rO; - auto q_tokens_left = q_len;// - q_start; + auto q_tokens_left = q_len; static_assert(q_step == REG_N); static_assert(kv_step == REG_K); @@ -463,11 +524,14 @@ void sdpa_kernel_lsc_prefetch( if (q_tokens_left > q_step) q_tokens_left = q_step; if (q_tokens_left > 0) { + // Fold log2(e) into Q pre-scale so St = K@Q^T lands in log2 domain; cm_exp + // then needs no per-element *log2e in the softmax critical path. + constexpr float qscale = scale_factor * 1.4426950408889634f; lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); #pragma unroll for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)qscale); } } @@ -480,107 +544,108 @@ void sdpa_kernel_lsc_prefetch( int causal_left = q_start; - for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step, - k_base += kv_step * kv_pitch, - v_base += kv_step * kv_pitch) { - //# St = k @ Qt - matrix St; // = ugemm_KQ(slm_K, rQ, slm_offset); - { - constexpr int num_K = kv_step/REG_M; - auto St2 = St.format(); + constexpr int KV_BLK = CMFLA_KV_BLK; + constexpr int num_K = kv_step/REG_M; + constexpr int BLK_ROWS = KV_BLK*kv_step; + + // Warm up first K and V tiles before the loop so kv_pos=0 finds them in cache. + prefetch_K.set_block_y(wg_local_id); + prefetch_V.set_block_y(wg_local_id); + cm_prefetch(prefetch_K.set_block_x(0)); + cm_prefetch(prefetch_V.set_block_x(0)); + #pragma unroll + for(int ri = 1; ri < padded_head_size/REG_K; ri++) { + cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); + cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); + } + for(int kv_base = 0; kv_base < kv_stop; kv_base += BLK_ROWS) { + //# St = K @ Q^T for KV_BLK stacked sub-tiles -> [KV_BLK*kv_step, q_step] + matrix St; + auto St2 = St.format(); + #pragma unroll + for(int b = 0; b < KV_BLK; b++) { + int kv_pos = kv_base + b*kv_step; matrix Kmat; - //cm_slm_block_read(slm_K, GENX_NONE, slm_offset, Kmat.format()); - prefetch_K.set_block_y(wg_local_id + kv_pos + kv_step); + prefetch_K.set_block_y(wg_local_id + kv_pos + BLK_ROWS); + prefetch_V.set_block_y(wg_local_id + kv_pos); cm_prefetch(prefetch_K.set_block_x(0)); + cm_prefetch(prefetch_V.set_block_x(0)); b2dK.set_block_y(kv_pos); cm_load(Kmat.format(), b2dK.set_block_x(0)); #pragma unroll for(int k = 0; k < num_K; k++) - St2.row(k) = cm_dpas( + St2.row(b*num_K + k) = cm_dpas( 0, rQ[0].format(), Kmat[k].format()); #pragma unroll for(int ri = 1; ri < padded_head_size/REG_K; ri++) { - //cm_slm_block_read(slm_K, GENX_NONE, slm_offset + ri * Kmat.n_elems() * sizeof(half), Kmat.format()); cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); + cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); cm_load(Kmat.format(), b2dK.set_block_x(ri*REG_K)); #pragma unroll for(int k = 0; k < num_K; k++) { - St2.row(k) = cm_dpas( - St2.row(k), + St2.row(b*num_K + k) = cm_dpas( + St2.row(b*num_K + k), rQ[ri].format(), Kmat[k].format()); } } } + + // ---- mask per sub-tile (causal or kv-tail) ---- if constexpr (use_causal_mask) { - // since kv_step == q_step == 16, causal_left is n*kv_step - if (causal_left == 0) { - apply_causal_mask<1>(St); - } else if (causal_left < 0) { - St = -3.4e38f; + #pragma unroll + for(int b = 0; b < KV_BLK; b++) { + auto Stb = St.select(b*kv_step, 0); + int cl = causal_left - b*kv_step; + if (cl == 0) { + apply_causal_mask<1>(Stb); + } else if (cl < 0) { + Stb = -3.4e38f; + } } - causal_left -= kv_step; + causal_left -= BLK_ROWS; } else { - int kv_tokens = kv_stop - kv_pos; - // LSC ensures no overflow-access, but mask off k-tails attn-score is still required - for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; + int kv_tokens = kv_stop - kv_base; + for(int p = kv_tokens; p < BLK_ROWS; p++) St[p] = -3.4e38f; } - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); - - matrix P; - Transpose2DMatrix(St, P); + // ---- one online-softmax update over the whole block ---- + auto max_comp = online_softmax_update_tree(St, cur_max, cur_sum); - b2dV.set_block_y(kv_pos); - prefetch_V.set_block_y(wg_local_id +kv_pos + kv_step); - if (kv_pos == 0) { - // ugemm_PV0(slm_V, P, rO, slm_offset); - auto P2 = P.format(); + // ---- rescale rO ONCE for the whole block (amortized over KV_BLK tiles) ---- + // For kv_base=0 cur_max was -3e38 so max_comp=exp(-inf)=0 -> zeroes rO (== acc=0). + #pragma unroll + for(int t = 0; t < padded_head_size/REG_N*num_P_tiles; t++) { + auto cO = rO[t].format(); #pragma unroll - for(int k = 0, ri = 0; k < padded_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)); - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - rO[ri + p] = cm_dpas( - 0, - Vmat.format(), - P2.row(p).format()); - } - } + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + (t % num_P_tiles)*REG_M]); } - else { - //ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); + + // ---- transpose each sub-tile and accumulate P@V into rO ---- + constexpr int num_Vchunks = padded_head_size/REG_N; + #pragma unroll + for(int b = 0; b < KV_BLK; b++) { + matrix P; + transpose_St_to_P_half(St.select(b*kv_step, 0), P); auto P2 = P.format(); + b2dV.set_block_y(kv_base + b*kv_step); + matrix Vmat; #pragma unroll - for(int k = 0, ri=0; k < padded_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)); - - //# compensate cur_O - // matrix rO; - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO[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]); - } - + for(int ci = 0; ci < num_Vchunks; ci++) + cm_load(Vmat.select(ci*(REG_K/2),0).format(), b2dV.set_block_x(ci*REG_N)); + #pragma unroll + for(int ci = 0, ri = 0; ci < num_Vchunks; ci++, ri += num_P_tiles) { #pragma unroll for(int p = 0; p < num_P_tiles; p++) { rO[ri + p] = cm_dpas( rO[ri + p].format(), - Vmat.format(), + Vmat.select(ci*(REG_K/2),0).format(), P2.row(p).format()); } } From 831cc6c28a81eb496a9b08b077727c423be4c134 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Fri, 12 Jun 2026 20:14:44 +0800 Subject: [PATCH 02/16] some fix. opt3 (cherry picked from commit a8a10c4cabfcd906da85908370a8b330e4aff06e) --- src/plugins/intel_cpu/CMakeLists.txt | 6 +- .../graph/impls/cm/include/cm_sdpa_common.hpp | 259 +++++++++++++++++- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 3 + .../intel_gpu/thirdparty/CMakeLists.txt | 2 +- 4 files changed, 256 insertions(+), 14 deletions(-) diff --git a/src/plugins/intel_cpu/CMakeLists.txt b/src/plugins/intel_cpu/CMakeLists.txt index 9df556712cf3cc..484a0f6879bc37 100644 --- a/src/plugins/intel_cpu/CMakeLists.txt +++ b/src/plugins/intel_cpu/CMakeLists.txt @@ -456,6 +456,6 @@ if(OV_CPU_WITH_ACL) endif() endif() -if(ENABLE_TESTS) - add_subdirectory(tests) -endif() +#if(ENABLE_TESTS) +# add_subdirectory(tests) +#endif() diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index 9e945192d03213..63eb790a63952b 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -14,6 +14,7 @@ * limitations under the License. *******************************************************************************/ #include "cm_attention_common.hpp" +#define CM_HAS_LSC_UNTYPED_2D 1 // Number of kv sub-tiles (kv_step rows each) processed per online-softmax update in // sdpa_kernel_lsc_prefetch. Larger amortizes the rO rescale + softmax bookkeeping over @@ -31,6 +32,9 @@ CM_INLINE vector online_softmax_update_tree(matrix_ref cur_max, vector_ref cur_sum) { static_assert((rows & (rows - 1)) == 0, "tree reduction needs power-of-two rows"); + // Reduce into rows/2 scratch (not a full copy) so register pressure stays low for + // large kv blocks. St itself is not modified during the max reduction (its values + // are still needed for the exp), so we fold pairs into scratch first. vector new_max_t; { matrix 1 ? rows/2 : 1), cols> t; @@ -44,9 +48,14 @@ CM_INLINE vector online_softmax_update_tree(matrix_ref(t.row(0), cur_max); } + // St is already in the log2 domain (Q was pre-scaled by scale_factor*log2e at load), + // so cm_exp (== exp2) needs no per-element *log2e here. cur_max/cur_sum stay in the + // same domain, so the running-max comparison and rescale are exact. #pragma unroll for (int r = 0; r < rows; r++) St[r] = cm_exp(St[r] - new_max_t); + // Sum reduction can fold St in place (values already consumed into exp form and the + // per-row sum is all we need afterwards). vector row_sum_t; { matrix 1 ? rows/2 : 1), cols> t; @@ -69,8 +78,9 @@ CM_INLINE vector online_softmax_update_tree(matrix_refhalf first lets the 16x16 shuffle network run at half -// width — roughly halving the mov count vs transposing the float tile directly. +// the P@V matmul. Casting float->half first (one vectorized cm_mul-free copy) lets the +// 16x16 shuffle network run at half width — roughly halving the mov count vs transposing +// the float tile directly through Transpose_16x16. CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { matrix Sh; #pragma unroll @@ -505,6 +515,8 @@ void sdpa_kernel_lsc_prefetch( constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); + // round up head_size to multiple of 16 + // block_2d_desc will automatically handle the tailing block constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; vector cur_max; @@ -516,7 +528,7 @@ void sdpa_kernel_lsc_prefetch( matrix rQ; matrix rO; - auto q_tokens_left = q_len; + auto q_tokens_left = q_len;// - q_start; static_assert(q_step == REG_N); static_assert(kv_step == REG_K); @@ -524,8 +536,9 @@ void sdpa_kernel_lsc_prefetch( if (q_tokens_left > q_step) q_tokens_left = q_step; if (q_tokens_left > 0) { - // Fold log2(e) into Q pre-scale so St = K@Q^T lands in log2 domain; cm_exp - // then needs no per-element *log2e in the softmax critical path. + // Fold log2(e) into the Q pre-scale so St = K@Q^T lands in the log2 domain; the + // online softmax then uses cm_exp (== exp2) directly, dropping a *log2e per St + // element (16 muls/tile off the softmax critical path). Math is identical. constexpr float qscale = scale_factor * 1.4426950408889634f; lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); #pragma unroll @@ -544,6 +557,10 @@ void sdpa_kernel_lsc_prefetch( int causal_left = q_start; + // KV-blocking factor: number of kv tiles (each kv_step rows) processed per online- + // softmax update. The per-iteration rO rescale (REG_M*num_rO_tiles muls) and the + // softmax bookkeeping amortize over KV_BLK tiles, cutting the ALU-bound consume cost. + // The K@Q^T and the St->P transpose are still per sub-tile (they don't amortize). constexpr int KV_BLK = CMFLA_KV_BLK; constexpr int num_K = kv_step/REG_M; constexpr int BLK_ROWS = KV_BLK*kv_step; @@ -567,6 +584,8 @@ void sdpa_kernel_lsc_prefetch( for(int b = 0; b < KV_BLK; b++) { int kv_pos = kv_base + b*kv_step; matrix Kmat; + // Prefetch K of the matching tile one block ahead and V of this tile, so V + // is warm by the PV phase and the next block's K is warm by its K phase. prefetch_K.set_block_y(wg_local_id + kv_pos + BLK_ROWS); prefetch_V.set_block_y(wg_local_id + kv_pos); cm_prefetch(prefetch_K.set_block_x(0)); @@ -611,6 +630,7 @@ void sdpa_kernel_lsc_prefetch( causal_left -= BLK_ROWS; } else { int kv_tokens = kv_stop - kv_base; + // LSC ensures no overflow-access, but mask off k-tail attn-score is still required for(int p = kv_tokens; p < BLK_ROWS; p++) St[p] = -3.4e38f; } @@ -635,17 +655,15 @@ void sdpa_kernel_lsc_prefetch( transpose_St_to_P_half(St.select(b*kv_step, 0), P); auto P2 = P.format(); b2dV.set_block_y(kv_base + b*kv_step); - matrix Vmat; - #pragma unroll - for(int ci = 0; ci < num_Vchunks; ci++) - cm_load(Vmat.select(ci*(REG_K/2),0).format(), b2dV.set_block_x(ci*REG_N)); #pragma unroll for(int ci = 0, ri = 0; ci < num_Vchunks; ci++, ri += num_P_tiles) { + matrix Vmat; + cm_load(Vmat.format(), b2dV.set_block_x(ci*REG_N)); #pragma unroll for(int p = 0; p < num_P_tiles; p++) { rO[ri + p] = cm_dpas( rO[ri + p].format(), - Vmat.select(ci*(REG_K/2),0).format(), + Vmat.format(), P2.row(p).format()); } } @@ -675,6 +693,227 @@ void sdpa_kernel_lsc_prefetch( } } +// Item 9: process 2 Q-rows per thread — amortize per-kv softmax cost over 2x DPAS work. +// q_len is the number of Q tokens for the FIRST row; the second row is at q_base + q_pitch. +// The caller must ensure at most 2*q_step tokens are dispatched to this thread. +template +void sdpa_kernel_lsc_prefetch_q2( + int wg_local_id, + int q_start, + int kv_stop, + int q_len, // tokens for this thread, up to 2*q_step + int kv_len, + svmptr_t q_base [[type("svmptr_t")]], + svmptr_t k_base [[type("svmptr_t")]], + svmptr_t v_base [[type("svmptr_t")]], + svmptr_t o_base [[type("svmptr_t")]]) { + + constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); + constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; + constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); + constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; + + // Two independent softmax states for the two Q rows. + vector cur_max_a, cur_sum_a; + vector cur_max_b, cur_sum_b; + cur_max_a = -3e38f; cur_sum_a = 0; + cur_max_b = -3e38f; cur_sum_b = 0; + + constexpr int num_P_tiles = REG_N / REG_M; + // Two rQ halves: rQ_a for first row, rQ_b for second row + matrix rQ_a; + matrix rQ_b; + // Two rO halves + matrix rO_a; + matrix rO_b; + + int q_len_a = q_len; + if (q_len_a < 0) q_len_a = 0; + if (q_len_a > q_step) q_len_a = q_step; + + int q_len_b = q_len - q_step; + if (q_len_b < 0) q_len_b = 0; + if (q_len_b > q_step) q_len_b = q_step; + + // Load rQ_a (first Q row) + if (q_len_a > 0) { + lsc::block_2d_desc b2dQa(reinterpret_cast(q_base), q_len_a - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); + #pragma unroll + for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { + cm_load(rQ_a[ri].format(), b2dQa.set_block_x(k)); + rQ_a[ri].format() = cm_mul(rQ_a[ri].format(), (half)scale_factor); + } + } + // Load rQ_b (second Q row block: q_step tokens further = q_step * q_pitch bytes) + if (q_len_b > 0) { + lsc::block_2d_desc b2dQb(reinterpret_cast(q_base + (svmptr_t)q_step * q_pitch), q_len_b - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); + #pragma unroll + for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { + cm_load(rQ_b[ri].format(), b2dQb.set_block_x(k)); + rQ_b[ri].format() = cm_mul(rQ_b[ri].format(), (half)scale_factor); + } + } + + lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + + static_assert(wg_local_size == 16); + lsc::block_2d_desc prefetch_K(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc prefetch_V(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + + int causal_left = q_start; + + // Warm-up prefetch for kv_pos=0 K+V + prefetch_K.set_block_y(wg_local_id); + prefetch_V.set_block_y(wg_local_id); + cm_prefetch(prefetch_K.set_block_x(0)); + cm_prefetch(prefetch_V.set_block_x(0)); + #pragma unroll + for(int ri = 1; ri < padded_head_size/REG_K; ri++) { + cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); + cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); + } + + for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step, + k_base += kv_step * kv_pitch, + v_base += kv_step * kv_pitch) { + + // K phase: compute St_a = K × rQ_a and St_b = K × rQ_b + matrix St_a; + matrix St_b; + { + constexpr int num_K = kv_step/REG_M; + auto St_a2 = St_a.format(); + auto St_b2 = St_b.format(); + + matrix Kmat; + + prefetch_K.set_block_y(wg_local_id + kv_pos + kv_step); + prefetch_V.set_block_y(wg_local_id + kv_pos); + cm_prefetch(prefetch_K.set_block_x(0)); + cm_prefetch(prefetch_V.set_block_x(0)); + + b2dK.set_block_y(kv_pos); + cm_load(Kmat.format(), b2dK.set_block_x(0)); + #pragma unroll + for(int k = 0; k < num_K; k++) { + St_a2.row(k) = cm_dpas( + 0, rQ_a[0].format(), Kmat[k].format()); + St_b2.row(k) = cm_dpas( + 0, rQ_b[0].format(), Kmat[k].format()); + } + + #pragma unroll + for(int ri = 1; ri < padded_head_size/REG_K; ri++) { + cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); + cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); + cm_load(Kmat.format(), b2dK.set_block_x(ri*REG_K)); + #pragma unroll + for(int k = 0; k < num_K; k++) { + St_a2.row(k) = cm_dpas( + St_a2.row(k), rQ_a[ri].format(), Kmat[k].format()); + St_b2.row(k) = cm_dpas( + St_b2.row(k), rQ_b[ri].format(), Kmat[k].format()); + } + } + } + + // Tail masking + { + int kv_tokens = kv_stop - kv_pos; + for(int p = kv_tokens; p < kv_step; p++) { + St_a[p] = -3.4e38f; + St_b[p] = -3.4e38f; + } + } + + auto max_comp_a = online_softmax_update(St_a, cur_max_a, cur_sum_a); + auto max_comp_b = online_softmax_update(St_b, cur_max_b, cur_sum_b); + + matrix P_a, P_b; + Transpose2DMatrix(St_a, P_a); + Transpose2DMatrix(St_b, P_b); + + // V phase: unified path (max_comp=0 at kv_pos=0 zeroes rO) + b2dV.set_block_y(kv_pos); + { + auto Pa2 = P_a.format(); + auto Pb2 = P_b.format(); + #pragma unroll + for(int k = 0, ri = 0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { + matrix Vmat; + cm_load(Vmat.format(), b2dV.set_block_x(k)); + // Rescale and DPAS for rO_a + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_a[ri + p].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp_a[r + p*REG_M]); + } + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + rO_a[ri + p] = cm_dpas( + rO_a[ri + p].format(), Vmat.format(), Pa2.row(p).format()); + } + // Rescale and DPAS for rO_b + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_b[ri + p].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp_b[r + p*REG_M]); + } + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + rO_b[ri + p] = cm_dpas( + rO_b[ri + p].format(), Vmat.format(), Pb2.row(p).format()); + } + } + } + } + + // Store output for first Q row + if (q_len_a > 0) { + matrix cur_O_f16; + cur_sum_a = cm_inv(cur_sum_a); + lsc::block_2d_desc b2dO(o_base, q_len_a - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); + #pragma unroll + for(int k = 0, ri=0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_a[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_a[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 output for second Q row + if (q_len_b > 0) { + matrix cur_O_f16; + cur_sum_b = cm_inv(cur_sum_b); + lsc::block_2d_desc b2dO(o_base + (svmptr_t)q_step * o_pitch, q_len_b - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); + #pragma unroll + for(int k = 0, ri=0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { + #pragma unroll + for(int p = 0; p < num_P_tiles; p++) { + auto cO = rO_b[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_b[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)); + } + } +} + #else // CM_HAS_LSC_UNTYPED_2D template diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index e4c12f182d0044..553a0e1db4f6dc 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -56,6 +56,8 @@ class VLSDPAGenerator : public KernelGenerator { const size_t num_q_heads = query_shape[query_shape.size() - 3].get_length(); const size_t num_kv_heads = key_shape[key_shape.size() - 3].get_length(); const float scale_factor = 1.0f / std::sqrt(static_cast(head_size)); + // tradeoff between register usage and parallelism: blk=1 allows more flexible scheduling and better parallelism + const size_t CMFLA_KV_BLK = head_size <= 64 ? 2 : 1; GPU_DEBUG_TRACE_DETAIL << "VLSDPA query_shape " << query_shape << ", q_transpose_order " << PartialShape(desc->input_q_transpose_order) << ", key_shape " << key_shape << ", k_transpose_order " << PartialShape(desc->input_k_transpose_order) @@ -67,6 +69,7 @@ class VLSDPAGenerator : public KernelGenerator { make_jit_constant("CMFLA_NUM_KV_HEADS", num_kv_heads), make_jit_constant("CMFLA_HEAD_SIZE", head_size), make_jit_constant("CMFLA_SCALE_FACTOR", scale_factor), + make_jit_constant("CMFLA_KV_BLK", CMFLA_KV_BLK), }); return jit; diff --git a/src/plugins/intel_gpu/thirdparty/CMakeLists.txt b/src/plugins/intel_gpu/thirdparty/CMakeLists.txt index 22db5e9c392cb7..4ad6554ddf9476 100644 --- a/src/plugins/intel_gpu/thirdparty/CMakeLists.txt +++ b/src/plugins/intel_gpu/thirdparty/CMakeLists.txt @@ -151,7 +151,7 @@ if(ENABLE_ONEDNN_FOR_GPU) "-DDNNL_ENABLE_WORKLOAD=INFERENCE" "-DDNNL_ENABLE_JIT_PROFILING=${BUILD_SHARED_LIBS}" "-DDNNL_ENABLE_ITT_TASKS=${BUILD_SHARED_LIBS}" - "-DDNNL_BUILD_TESTS=OFF" + "-DDNNL_BUILD_TESTS=ON" "-DDNNL_BUILD_EXAMPLES=OFF" "-DDNNL_BLAS_VENDOR=NONE" "-DDNNL_LIBRARY_TYPE=STATIC" From 1a1316e75e7f303955aca1fd5c59147eadd178e9 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 16 Jun 2026 18:34:45 +0800 Subject: [PATCH 03/16] CM PA/SDPA: sync aboutSHW item15 + log2e-fold to OV kernels cm_attention_common.hpp: - Add constexpr log2e and q_scale_factor (fold log2e into Q pre-scale, removing per-element mul from softmax critical path) - Add online_softmax_update_tree (tree-reduction variant, depth log2 vs linear) - Add transpose_St_to_P_half (float->half before GRF shuffle, halves mov width) cm_pa_xe2.hpp: - Add CMPA_USE_TREE_SOFTMAX macro + pa_softmax_update dispatch macro - Switch Q scale from scale_factor to q_scale_factor (both U8 and FP16 paths) - Replace online_softmax_update + Transpose2DMatrix call sites with pa_softmax_update + transpose_St_to_P_half at all 4 locations cm_sdpa_common.hpp: - Remove now-duplicate online_softmax_update_tree and transpose_St_to_P_half (inherited from cm_attention_common.hpp via existing include) paged_attention_gen.cpp: - Emit CMPA_USE_TREE_SOFTMAX=1 for INT8, =0 for FP16: tree softmax wins 1-5% on INT8 (dequant amortises scratch), regresses ~9% on FP16 (register pressure from scratch matrix) Co-Authored-By: Claude Sonnet 4.6 (cherry picked from commit 92dc15757c032affd1d252973efe972b698c327f) --- .../impls/cm/include/cm_attention_common.hpp | 65 ++++++++++++++++-- .../src/graph/impls/cm/include/cm_pa_xe2.hpp | 32 ++++++--- .../graph/impls/cm/include/cm_sdpa_common.hpp | 66 +------------------ .../graph/impls/cm/paged_attention_gen.cpp | 3 + 4 files changed, 87 insertions(+), 79 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp index 6a1f7b4b860a3f..712ba20c42d614 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp @@ -24,6 +24,8 @@ static_assert(__cplusplus >= 201703L); #define q_step REG_N constexpr float scale_factor = CMFLA_SCALE_FACTOR; +constexpr float log2e = 1.4426950408889634f; +constexpr float q_scale_factor = scale_factor * log2e; static_assert(q_step == 16 || q_step == 8); static_assert(kv_step == 16); @@ -301,11 +303,11 @@ vector online_softmax_update(matrix_ref St, vector_r new_max_t = cm_max(new_max_t, cur_max); // Pt = torch.exp(St - new_max) - constexpr float log2e = 1.4426950408889634f; + constexpr float _log2e = 1.4426950408889634f; #ifdef ABLATE_NO_EXP - for(int r = 0; r < St.n_rows(); r++) St[r] = (St[r] - new_max_t)*log2e; + for(int r = 0; r < St.n_rows(); r++) St[r] = (St[r] - new_max_t)*_log2e; #else - for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp((St[r] - new_max_t)*log2e); + for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp((St[r] - new_max_t)*_log2e); #endif vector row_sum_t; @@ -313,13 +315,68 @@ vector online_softmax_update(matrix_ref St, vector_r for(int r = 2; r < St.n_rows(); r++) row_sum_t = cm_add(row_sum_t, St[r]); vector max_comp; - max_comp = cm_exp((cur_max - new_max_t)*log2e); + max_comp = cm_exp((cur_max - new_max_t)*_log2e); cur_sum = cm_mul(cur_sum, max_comp); cur_sum = cm_add(cur_sum, row_sum_t); cur_max = new_max_t; return max_comp; } +// Tree-reduction variant: max/sum reductions use a balanced binary tree (depth log2(rows)) +// instead of a linear chain (depth rows-1), shortening the loop-carried dependency chain. +// Requires rows to be a power of two; the PA kv_step=16 and KV_BLK*kv_step satisfy this. +template +CM_INLINE vector online_softmax_update_tree(matrix_ref St, + vector_ref cur_max, + vector_ref cur_sum) { + static_assert((rows & (rows - 1)) == 0, "tree reduction needs power-of-two rows"); + vector new_max_t; + { + matrix 1 ? rows/2 : 1), cols> t; + #pragma unroll + for (int r = 0; r < rows/2; r++) t.row(r) = cm_max(St[r], St[r + rows/2]); + #pragma unroll + for (int stride = rows/4; stride > 0; stride >>= 1) + #pragma unroll + for (int r = 0; r < stride; r++) + t.row(r) = cm_max(t.row(r), t.row(r + stride)); + new_max_t = cm_max(t.row(0), cur_max); + } + #pragma unroll + for (int r = 0; r < rows; r++) St[r] = cm_exp(St[r] - new_max_t); + + vector row_sum_t; + { + matrix 1 ? rows/2 : 1), cols> t; + #pragma unroll + for (int r = 0; r < rows/2; r++) t.row(r) = cm_add(St[r], St[r + rows/2]); + #pragma unroll + for (int stride = rows/4; stride > 0; stride >>= 1) + #pragma unroll + for (int r = 0; r < stride; r++) + t.row(r) = cm_add(t.row(r), t.row(r + stride)); + row_sum_t = t.row(0); + } + + vector max_comp; + max_comp = cm_exp(cur_max - new_max_t); + cur_sum = cm_mul(cur_sum, max_comp); + cur_sum = cm_add(cur_sum, row_sum_t); + cur_max = new_max_t; + return max_comp; +} + +// Transpose a [16,16] float score tile (kv x q) into a half [16,16] P tile (q x kv) for +// the P@V matmul. Casting float->half first (one vectorized cm_mul-free copy) lets the +// 16x16 shuffle network run at half width — roughly halving the mov count vs transposing +// the float tile directly through Transpose_16x16. +CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { + matrix Sh; + #pragma unroll + for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); + Transpose_16x16(Sh.select<16,1,16,1>(0,0), P); +} + #ifdef CM_HAS_LSC_UNTYPED_2D #define cm_load_normal cm_load #define cm_load_transpose cm_load diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp index 56b011f6f1bb40..7bf48dbbeb347d 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp @@ -10,6 +10,18 @@ #define CMPA_WG_SEQ_LEN 0 #endif +// Item 15: use balanced binary-tree reduction in online_softmax_update (log2 depth vs linear). +// Set to 0 to disable (use linear chain), 1 to enable. +#ifndef CMPA_USE_TREE_SOFTMAX +#define CMPA_USE_TREE_SOFTMAX 1 +#endif + +#if CMPA_USE_TREE_SOFTMAX +#define pa_softmax_update(St, cur_max, cur_sum) online_softmax_update_tree(St, cur_max, cur_sum) +#else +#define pa_softmax_update(St, cur_max, cur_sum) online_softmax_update(St, cur_max, cur_sum) +#endif + #ifndef SPARSE_BLOCK_SIZE #error "SPARSE_BLOCK_SIZE must be defined" #endif @@ -105,7 +117,7 @@ void pa_lsc_u8( #pragma unroll for (int k = 0, ri = 0; k < head_size / 2; k += REG_K / 2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_scale_factor); } } @@ -342,10 +354,10 @@ void pa_lsc_u8( } int kv_tokens = kv_stop - kv_pos; for (int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = pa_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); if (first_active) { ugemm_PV0(slm_V, P, rO, slm_offset); @@ -528,10 +540,10 @@ void pa_lsc_u8( } int kv_tokens = kv_stop - kv_pos; for (int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = pa_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); if (first_active) { ugemm_PV0(slm_V, P, rO, slm_offset); @@ -654,7 +666,7 @@ void pa_kernel_lsc_prefetch_f16( #pragma unroll for(int k = 0, ri = 0; k < head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_scale_factor); } lsc::block_2d_desc b2dK(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(half) - 1, k_pitch - 1, 0, 0); @@ -744,10 +756,10 @@ void pa_kernel_lsc_prefetch_f16( // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = pa_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); prefetch_V.set_base_ptr((reinterpret_cast(v_cache_base)+prefetch_block_id*blk_stride)); prefetch_V.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); @@ -885,10 +897,10 @@ void pa_kernel_lsc_prefetch_f16( // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = pa_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); prefetch_V.set_base_ptr((reinterpret_cast(v_cache_base)+prefetch_block_id*blk_stride)); prefetch_V.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index 63eb790a63952b..aa17f73271b354 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -14,7 +14,6 @@ * limitations under the License. *******************************************************************************/ #include "cm_attention_common.hpp" -#define CM_HAS_LSC_UNTYPED_2D 1 // Number of kv sub-tiles (kv_step rows each) processed per online-softmax update in // sdpa_kernel_lsc_prefetch. Larger amortizes the rO rescale + softmax bookkeeping over @@ -23,70 +22,7 @@ #define CMFLA_KV_BLK 2 #endif -// Flashattn-local online-softmax with tree reductions over the kv (row) dimension. -// Identical math to online_softmax_update in cm_attention_common.hpp, but the max and -// sum reductions are balanced binary trees (dependency depth log2(rows)=4 instead of the -// linear chain's depth 15), which shortens the loop-carried softmax critical path. -template -CM_INLINE vector online_softmax_update_tree(matrix_ref St, - vector_ref cur_max, - vector_ref cur_sum) { - static_assert((rows & (rows - 1)) == 0, "tree reduction needs power-of-two rows"); - // Reduce into rows/2 scratch (not a full copy) so register pressure stays low for - // large kv blocks. St itself is not modified during the max reduction (its values - // are still needed for the exp), so we fold pairs into scratch first. - vector new_max_t; - { - matrix 1 ? rows/2 : 1), cols> t; - #pragma unroll - for (int r = 0; r < rows/2; r++) t.row(r) = cm_max(St[r], St[r + rows/2]); - #pragma unroll - for (int stride = rows/4; stride > 0; stride >>= 1) - #pragma unroll - for (int r = 0; r < stride; r++) - t.row(r) = cm_max(t.row(r), t.row(r + stride)); - new_max_t = cm_max(t.row(0), cur_max); - } - - // St is already in the log2 domain (Q was pre-scaled by scale_factor*log2e at load), - // so cm_exp (== exp2) needs no per-element *log2e here. cur_max/cur_sum stay in the - // same domain, so the running-max comparison and rescale are exact. - #pragma unroll - for (int r = 0; r < rows; r++) St[r] = cm_exp(St[r] - new_max_t); - - // Sum reduction can fold St in place (values already consumed into exp form and the - // per-row sum is all we need afterwards). - vector row_sum_t; - { - matrix 1 ? rows/2 : 1), cols> t; - #pragma unroll - for (int r = 0; r < rows/2; r++) t.row(r) = cm_add(St[r], St[r + rows/2]); - #pragma unroll - for (int stride = rows/4; stride > 0; stride >>= 1) - #pragma unroll - for (int r = 0; r < stride; r++) - t.row(r) = cm_add(t.row(r), t.row(r + stride)); - row_sum_t = t.row(0); - } - - vector max_comp; - max_comp = cm_exp(cur_max - new_max_t); - cur_sum = cm_mul(cur_sum, max_comp); - cur_sum = cm_add(cur_sum, row_sum_t); - cur_max = new_max_t; - return max_comp; -} - -// Transpose a [16,16] float score tile (kv x q) into a half [16,16] P tile (q x kv) for -// the P@V matmul. Casting float->half first (one vectorized cm_mul-free copy) lets the -// 16x16 shuffle network run at half width — roughly halving the mov count vs transposing -// the float tile directly through Transpose_16x16. -CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { - matrix Sh; - #pragma unroll - for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); - Transpose_16x16(Sh.select<16,1,16,1>(0,0), P); -} +// online_softmax_update_tree and transpose_St_to_P_half are defined in cm_attention_common.hpp #ifdef CM_HAS_LSC_UNTYPED_2D //@prefetch_u8 would have duplicated decompress perf issue. comments out for now. diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp index 633187798c2f7a..ceed4e0cff2d2e 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp @@ -381,6 +381,9 @@ JitConstants PagedAttentionGeneratorMultiToken::get_jit_constants(const kernel_i jit.make("CMPA_BLOCK_SZ", PA_KV_CACHE_BLOCK_SIZE_LEGACY); } jit.make("CMPA_WG_SEQ_LEN", get_wg_seq_len(params)); + // Tree softmax: wins on INT8 (dequant amortises scratch, shorter dep-chain helps), + // regresses ~9% on FP16 (extra scratch matrix increases register pressure). + jit.make("CMPA_USE_TREE_SOFTMAX", get_kv_compressed(params) ? 1 : 0); return jit; } From aefc561114e2e6280d4c4de030c3bf6e2367e890 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 16 Jun 2026 23:19:00 +0800 Subject: [PATCH 04/16] gpu(vl_sdpa): remove host sync and read cu_seqlens via mem_lock (cherry picked from commit cbe4e6034e2e1d526d34a95b86a21ece55c9b22c) --- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 64 ++++++------------- .../src/graph/include/vl_sdpa_inst.h | 9 ++- src/plugins/intel_gpu/src/graph/vl_sdpa.cpp | 16 ----- 3 files changed, 25 insertions(+), 64 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index 553a0e1db4f6dc..6d13c37a7b0fa6 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -20,10 +20,6 @@ constexpr auto get_vlsdpa_build_options() { return " -cmc -Qxcm_register_file_size=256"; } -struct VLSDPARuntimeParams : public ImplRuntimeParams { - std::vector cu_seqlens; -}; - class VLSDPAGenerator : public KernelGenerator { public: VLSDPAGenerator() : KernelGenerator("cm_sdpa_vlen") {} @@ -94,7 +90,7 @@ class VLSDPAGenerator : public KernelGenerator { } [[nodiscard]] DispatchDataFunc get_dispatch_data_func() const override { - return DispatchDataFunc{[](const RuntimeParams& params, KernelData& kd, ImplRuntimeParams* rt_params) { + return DispatchDataFunc{[](const RuntimeParams& params, KernelData& kd, ImplRuntimeParams*) { assert(!params.is_dynamic()); auto desc = params.typed_desc(); @@ -110,28 +106,27 @@ class VLSDPAGenerator : public KernelGenerator { return transposed_pshape; }; const auto query_shape = transpose_pshape(params.get_input_layout(0).get_shape(), desc->input_q_transpose_order); - const size_t num_q_heads = query_shape[query_shape.size() - 3]; // TODO: make it to be configuration of primitive_inst + const size_t num_q_heads = query_shape[query_shape.size() - 3]; - const auto& vlsdpa_rt_params = static_cast(*rt_params); - const auto& cu_seqlens = vlsdpa_rt_params.cu_seqlens; + // Read cu_seqlens via mem_lock — zero-cost on USM-host (PC), no GPU sync stall. + const auto cu_seqlens_mem = params.memory_deps.at(params.input_layouts.size() - 1); + mem_lock cu_seqlens_lock(cu_seqlens_mem, *params.strm); size_t max_seq_len = 0; - for (size_t i = 1; i < cu_seqlens.size(); i++) { - auto start_idx = cu_seqlens[i - 1]; - auto end_idx = cu_seqlens[i]; + for (size_t i = 1; i < cu_seqlens_lock.size(); i++) { + auto start_idx = cu_seqlens_lock[i - 1]; + auto end_idx = cu_seqlens_lock[i]; max_seq_len = std::max(max_seq_len, static_cast(end_idx - start_idx)); } const auto& info = params.get_device_info(); const size_t CM_GRF_WIDTH = (info.arch <= gpu_arch::xe_hpc) ? 256 : 512; - const size_t q_step = static_cast(std::floor(CM_GRF_WIDTH / 32)); // or 8 on Xe1 + const size_t q_step = static_cast(std::floor(CM_GRF_WIDTH / 32)); size_t wg_size = static_cast(std::floor((max_seq_len + q_step - 1) / q_step)); int32_t need_wg_mapping = 0; if (wg_size > 16) { - // # seq_len is too big to fit into a single work-group - // # will use fixed work-group size 16, process 16*16 (or 16*8 on xe1) - // # part of sequence, in this case, kernel needs to figure-out which part - // # it needs to handle + // seq_len is too large for a single work-group; use wg_size=16 and let + // the kernel's while-loop scan cu_seqlens to find its sequence/block. need_wg_mapping = 1; wg_size = 16; } @@ -140,13 +135,13 @@ class VLSDPAGenerator : public KernelGenerator { if (need_wg_mapping) { wg_count = 0; const auto wg_seq_len = wg_size * q_step; - for (size_t i = 1; i < cu_seqlens.size(); i++) { - auto start_idx = cu_seqlens[i - 1]; - auto end_idx = cu_seqlens[i]; - wg_count += static_cast(std::floor((end_idx - start_idx + wg_seq_len - 1) / wg_seq_len)); + for (size_t i = 1; i < cu_seqlens_lock.size(); i++) { + auto start_idx = cu_seqlens_lock[i - 1]; + auto end_idx = cu_seqlens_lock[i]; + wg_count += static_cast((end_idx - start_idx + wg_seq_len - 1) / wg_seq_len); } } else { - wg_count = cu_seqlens.size() - 1; + wg_count = cu_seqlens_lock.size() - 1; } auto& wgs = kd.params.workGroups; @@ -156,10 +151,10 @@ class VLSDPAGenerator : public KernelGenerator { std::vector scalars{need_wg_mapping}; kd.params.scalars.clear(); for (auto i : scalars) { - scalar_desc desc; - desc.t = scalar_desc::Types::INT32; - desc.v.s32 = static_cast(i); - kd.params.scalars.push_back(desc); + scalar_desc s; + s.t = scalar_desc::Types::INT32; + s.v.s32 = static_cast(i); + kd.params.scalars.push_back(s); } }}; } @@ -171,9 +166,7 @@ class VLSDPAOptImpl : public PrimitiveImplCM { Stage::Ptr vl_sdpa = make_stage(); - VLSDPAOptImpl() : PrimitiveImplOCL(VLSDPAOptImplementationManager::get_type_info_static()) { - m_rt_params = std::make_unique(); - } + VLSDPAOptImpl() : PrimitiveImplOCL(VLSDPAOptImplementationManager::get_type_info_static()) {} VLSDPAOptImpl(const program_node& node, const RuntimeParams& params) : VLSDPAOptImpl() { add_stage(vl_sdpa, params); } @@ -181,21 +174,6 @@ class VLSDPAOptImpl : public PrimitiveImplCM { [[nodiscard]] std::unique_ptr clone() const override { return make_deep_copy(this); } - - void update_rt_params(const cldnn::primitive_inst& instance) override { - update_stages_flags(instance); - - auto rt_params = static_cast(m_rt_params.get()); - auto& vlsdpa_instance = dynamic_cast&>(instance); - vlsdpa_instance.get_mask_seqlens_from_memory(rt_params->cu_seqlens); - } - - cldnn::event::ptr execute(const std::vector& events, cldnn::primitive_inst& instance) override { - if (m_rt_params == nullptr) { - m_rt_params = std::make_unique(); - } - return PrimitiveImplCM::execute(events, instance); - } }; } // namespace diff --git a/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h b/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h index 42eaf052eec78b..4a86ade1283bef 100644 --- a/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h +++ b/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h @@ -18,7 +18,10 @@ struct typed_program_node : public typed_program_node_base { using parent::parent; program_node& input(size_t index = 0) const { return get_dependency(index); } - std::vector get_shape_infer_dependencies() const override { return {}; } + std::vector get_shape_infer_dependencies() const override { + // cu_seq_lens is input index 3; dispatch_data_func reads it via mem_lock. + return {3}; + } }; using vl_sdpa_node = typed_program_node; @@ -39,10 +42,6 @@ class typed_primitive_inst : public typed_primitive_inst_base static std::string to_string(const vl_sdpa_node& node); typed_primitive_inst(network& network, const vl_sdpa_node& node); - - void get_mask_seqlens_from_memory(std::vector& cu_seqlens) const; - - memory::ptr cu_seqlens_memory_ptr() const { return dep_memory_ptr(3); } }; using vl_sdpa_inst = typed_primitive_inst; diff --git a/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp b/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp index 60674bfe70bfe9..227a351264e3a2 100644 --- a/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp +++ b/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp @@ -5,7 +5,6 @@ #include "vl_sdpa_inst.h" #include "primitive_type_base.h" #include "json_object.h" -#include "intel_gpu/runtime/utils.hpp" namespace cldnn { GPU_DEFINE_PRIMITIVE_TYPE_ID(vl_sdpa); @@ -30,19 +29,4 @@ std::string vl_sdpa_inst::to_string(const vl_sdpa_node& node) { vl_sdpa_inst::typed_primitive_inst(network& network, const vl_sdpa_node& node) : parent(network, node) {} -void vl_sdpa_inst::get_mask_seqlens_from_memory(std::vector& cu_seqlens) const { - cldnn::stream& stream = get_network().get_stream(); - stream.finish(); - - const auto cu_seqlens_mem = cu_seqlens_memory_ptr(); - - auto size = cu_seqlens_mem->count(); - OPENVINO_ASSERT(cu_seqlens_mem->get_layout().data_type == ov::element::i32 && size > 0); - - cu_seqlens.resize(size); - cu_seqlens_mem->copy_to(stream, cu_seqlens.data(), 0, 0, size * sizeof(int32_t), true); - - GPU_DEBUG_TRACE_DETAIL << " get_mask_seqlens_from_memory " << cu_seqlens << std::endl; -} - } // namespace cldnn From 6f2b4f2d165d0fdafb04159e55d816b68253721b Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 7 Jul 2026 00:13:26 +0800 Subject: [PATCH 05/16] gate rescale rO (skip first iter only when KV blocking amortizes the branch) (cherry picked from commit f11f6c46d81ae092fcff83cbc6af83dd611ee388) --- .../graph/impls/cm/include/cm_sdpa_common.hpp | 37 +++++++++++++++---- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 20 ++++++++-- 2 files changed, 46 insertions(+), 11 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index aa17f73271b354..5626b940e20132 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -19,7 +19,7 @@ // sdpa_kernel_lsc_prefetch. Larger amortizes the rO rescale + softmax bookkeeping over // more DPAS work at the cost of GRF pressure. Overridable via -DCMFLA_KV_BLK=N. #ifndef CMFLA_KV_BLK -#define CMFLA_KV_BLK 2 +#define CMFLA_KV_BLK 1 #endif // online_softmax_update_tree and transpose_St_to_P_half are defined in cm_attention_common.hpp @@ -303,6 +303,7 @@ void sdpa_kernel_lsc( constexpr int num_P_tiles = REG_N / REG_M; matrix rQ; matrix rO; + rO = 0.0f; // Zero the accumulator: the first softmax block scales rO by max_comp==0, and 0*NaN==NaN if the GRF holds stale NaN/Inf bits. auto q_tokens_left = q_len; static_assert(q_step == REG_N); @@ -463,6 +464,7 @@ void sdpa_kernel_lsc_prefetch( constexpr int num_P_tiles = REG_N / REG_M; matrix rQ; matrix rO; + rO = 0.0f; auto q_tokens_left = q_len;// - q_start; static_assert(q_step == REG_N); @@ -573,14 +575,30 @@ void sdpa_kernel_lsc_prefetch( // ---- one online-softmax update over the whole block ---- auto max_comp = online_softmax_update_tree(St, cur_max, cur_sum); - // ---- rescale rO ONCE for the whole block (amortized over KV_BLK tiles) ---- - // For kv_base=0 cur_max was -3e38 so max_comp=exp(-inf)=0 -> zeroes rO (== acc=0). - #pragma unroll - for(int t = 0; t < padded_head_size/REG_N*num_P_tiles; t++) { - auto cO = rO[t].format(); + // ---- rescale rO (skip first iter only when KV blocking amortizes the branch) ---- + // The kv_base=0 rescale is mathematically redundant for every head size: max_comp=0 + // and rO is still zero. The skip, however, adds a runtime kv_base branch to every + // outer KV block. Enable it only for KV_BLK>=2, where fewer outer iterations make + // the branch cost small enough to be profitable; KV_BLK=1 paths such as Omni HD=72 + // keep the straight-line rescale because forcing the branch measured as a regression. + if constexpr (KV_BLK >= 2) { + if (kv_base > 0) { + #pragma unroll + for(int t = 0; t < padded_head_size/REG_N*num_P_tiles; t++) { + auto cO = rO[t].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + (t % num_P_tiles)*REG_M]); + } + } + } else { #pragma unroll - for(int r = 0; r < REG_M; r++) - cO.row(r) = cm_mul(cO.row(r), max_comp[r + (t % num_P_tiles)*REG_M]); + for(int t = 0; t < padded_head_size/REG_N*num_P_tiles; t++) { + auto cO = rO[t].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + (t % num_P_tiles)*REG_M]); + } } // ---- transpose each sub-tile and accumulate P@V into rO ---- @@ -662,6 +680,8 @@ void sdpa_kernel_lsc_prefetch_q2( // Two rO halves matrix rO_a; matrix rO_b; + rO_a = 0.0f; // See rO note above: avoid 0*NaN from stale GRF bits. + rO_b = 0.0f; // See rO note above: avoid 0*NaN from stale GRF bits. int q_len_a = q_len; if (q_len_a < 0) q_len_a = 0; @@ -926,6 +946,7 @@ void sdpa_kernel( constexpr int num_P_tiles = REG_N / REG_M; matrix rO; + rO = 0.0f; // Zero the accumulator: the first softmax block scales rO by max_comp==0, and 0*NaN==NaN if the GRF holds stale NaN/Inf bits. int causal_left = q_start; constexpr uint slm_buff_size = kv_step * padded_head_size * sizeof(half); diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index 6d13c37a7b0fa6..f422be4ea1429a 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -20,6 +20,21 @@ constexpr auto get_vlsdpa_build_options() { return " -cmc -Qxcm_register_file_size=256"; } +constexpr size_t round_up_to_tile(size_t value, size_t tile_size) { + return (value + tile_size - 1) / tile_size * tile_size; +} + +constexpr size_t get_default_kv_blk(size_t head_size) { + const size_t padded_head_size = round_up_to_tile(head_size, 16); + const size_t tail_size = padded_head_size - head_size; + + if (padded_head_size <= 64 || tail_size <= 1) { + return 2; + } + + return 1; +} + class VLSDPAGenerator : public KernelGenerator { public: VLSDPAGenerator() : KernelGenerator("cm_sdpa_vlen") {} @@ -52,8 +67,7 @@ class VLSDPAGenerator : public KernelGenerator { const size_t num_q_heads = query_shape[query_shape.size() - 3].get_length(); const size_t num_kv_heads = key_shape[key_shape.size() - 3].get_length(); const float scale_factor = 1.0f / std::sqrt(static_cast(head_size)); - // tradeoff between register usage and parallelism: blk=1 allows more flexible scheduling and better parallelism - const size_t CMFLA_KV_BLK = head_size <= 64 ? 2 : 1; + const size_t cmfla_kv_blk = get_default_kv_blk(head_size); GPU_DEBUG_TRACE_DETAIL << "VLSDPA query_shape " << query_shape << ", q_transpose_order " << PartialShape(desc->input_q_transpose_order) << ", key_shape " << key_shape << ", k_transpose_order " << PartialShape(desc->input_k_transpose_order) @@ -65,7 +79,7 @@ class VLSDPAGenerator : public KernelGenerator { make_jit_constant("CMFLA_NUM_KV_HEADS", num_kv_heads), make_jit_constant("CMFLA_HEAD_SIZE", head_size), make_jit_constant("CMFLA_SCALE_FACTOR", scale_factor), - make_jit_constant("CMFLA_KV_BLK", CMFLA_KV_BLK), + make_jit_constant("CMFLA_KV_BLK", cmfla_kv_blk), }); return jit; From 4b7d163dec55c9a2e73e1fa78b0922a4ce12c9f4 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 7 Jul 2026 13:32:37 +0800 Subject: [PATCH 06/16] roll back unexpected files. --- src/plugins/intel_cpu/CMakeLists.txt | 6 +++--- src/plugins/intel_gpu/thirdparty/CMakeLists.txt | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/plugins/intel_cpu/CMakeLists.txt b/src/plugins/intel_cpu/CMakeLists.txt index 484a0f6879bc37..9df556712cf3cc 100644 --- a/src/plugins/intel_cpu/CMakeLists.txt +++ b/src/plugins/intel_cpu/CMakeLists.txt @@ -456,6 +456,6 @@ if(OV_CPU_WITH_ACL) endif() endif() -#if(ENABLE_TESTS) -# add_subdirectory(tests) -#endif() +if(ENABLE_TESTS) + add_subdirectory(tests) +endif() diff --git a/src/plugins/intel_gpu/thirdparty/CMakeLists.txt b/src/plugins/intel_gpu/thirdparty/CMakeLists.txt index 4ad6554ddf9476..22db5e9c392cb7 100644 --- a/src/plugins/intel_gpu/thirdparty/CMakeLists.txt +++ b/src/plugins/intel_gpu/thirdparty/CMakeLists.txt @@ -151,7 +151,7 @@ if(ENABLE_ONEDNN_FOR_GPU) "-DDNNL_ENABLE_WORKLOAD=INFERENCE" "-DDNNL_ENABLE_JIT_PROFILING=${BUILD_SHARED_LIBS}" "-DDNNL_ENABLE_ITT_TASKS=${BUILD_SHARED_LIBS}" - "-DDNNL_BUILD_TESTS=ON" + "-DDNNL_BUILD_TESTS=OFF" "-DDNNL_BUILD_EXAMPLES=OFF" "-DDNNL_BLAS_VENDOR=NONE" "-DDNNL_LIBRARY_TYPE=STATIC" From 21ab814a3deba30d1ff0281045c9f373e6071a44 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Mon, 13 Jul 2026 14:13:24 +0800 Subject: [PATCH 07/16] clean up --- .../graph/impls/cm/include/cm_sdpa_common.hpp | 472 ------------------ 1 file changed, 472 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index 5626b940e20132..a130ca31ffe1e4 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -25,255 +25,6 @@ // online_softmax_update_tree and transpose_St_to_P_half are defined in cm_attention_common.hpp #ifdef CM_HAS_LSC_UNTYPED_2D -//@prefetch_u8 would have duplicated decompress perf issue. comments out for now. -// template -// void sdpa_kernel_lsc_prefetch_u8( -// int wg_local_id, -// int q_start, -// int kv_stop, // -// int q_len, //q_step -// int kv_len, //not used for now -// svmptr_t q_base [[type("svmptr_t")]], -// svmptr_t k_cache_base [[type("svmptr_t")]], -// svmptr_t v_cache_base [[type("svmptr_t")]], -// svmptr_t o_base [[type("svmptr_t")]], -// int32_t past_lens, -// int32_t* block_indices [[type("svmptr_t")]]) { -// constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); -// constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; -// //[block_num, kv_heads, block_size, head_size] -// constexpr uint kv_pitch = head_size * sizeof(uint8_t); - -// vector cur_max; -// vector cur_sum; - -// cur_max = -3e38f; -// cur_sum = 0; -// constexpr int num_P_tiles = REG_N / REG_M; -// matrix rQ; -// matrix rO; - -// auto q_tokens_left = q_len;// - q_start; -// static_assert(q_step == REG_N); -// static_assert(kv_step == REG_K); - -// if (q_tokens_left < 0) q_tokens_left = 0; -// if (q_tokens_left > q_step) q_tokens_left = q_step; - -// if (q_tokens_left > 0) { -// lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); -// #pragma unroll -// for(int k = 0, ri = 0; k < head_size/2; k += REG_K/2, ri++) { -// cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); -// rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); -// } -// } - -// lsc::block_2d_desc b2dKV(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0); - -// static_assert(wg_local_size == 16); -// lsc::block_2d_desc b2dKV_prefetch(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0); -// // constexpr int blk_stride = CMFLA_NUM_KV_HEADS * CMFLA_HEAD_SIZE * CMPA_BLOCK_SZ; -// constexpr int quan_blk_stride = CMFLA_NUM_KV_HEADS * (CMFLA_HEAD_SIZE+4) * CMPA_BLOCK_SZ * sizeof(uint8_t); - - - -// int causal_left = q_start+past_lens; -// for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step) { -// auto cur_block_id = block_indices[kv_pos / CMPA_BLOCK_SZ]; -// //For the last step, duplicate prefetch here. -// uint32_t prefetch_kv_pos = (kv_pos+kv_step) >= kv_stop ? kv_pos : (kv_pos+kv_step); -// auto prefetch_block_id = block_indices[prefetch_kv_pos / CMPA_BLOCK_SZ]; -// uint32_t dscale_offset = cur_block_id*quan_blk_stride + CMPA_BLOCK_SZ * head_size * sizeof(uint8_t) + kv_pos%CMPA_BLOCK_SZ*sizeof(half); - -// vector dscale; -// vector zp; -// cm_svm_block_read(reinterpret_cast( k_cache_base + dscale_offset), dscale); -// cm_svm_block_read(reinterpret_cast( k_cache_base + dscale_offset + CMPA_BLOCK_SZ*sizeof(half)), zp); - -// // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { -// // show(dscale.format()); -// // } -// //# St = k @ Qt - -// matrix St; // = ugemm_KQ(slm_K, rQ, slm_offset); -// { -// constexpr int num_K = kv_step/REG_M; -// auto St2 = St.format(); -// matrix Kmat; -// auto quan_Kmat = Kmat.format().row(1).format(); -// auto dq_Kmat = Kmat.format(); -// //cm_slm_block_read(slm_K, GENX_NONE, slm_offset, Kmat.format()); - -// b2dKV_prefetch.set_base_ptr(reinterpret_cast(k_cache_base+prefetch_block_id*quan_blk_stride)); -// b2dKV_prefetch.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); -// cm_prefetch(b2dKV_prefetch.set_block_x(0)); - -// b2dKV.set_base_ptr(reinterpret_cast(k_cache_base+cur_block_id*quan_blk_stride)); -// b2dKV.set_block_y(kv_pos%CMPA_BLOCK_SZ); - -// cm_load(quan_Kmat.format(), b2dKV.set_block_x(0)); -// // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { -// // show(quan_Kmat.format(), false); -// // } -// #pragma unroll -// for(int r = 0; r < kv_step; r++) { -// dq_Kmat[r] = quan_Kmat[r] - zp[r]; -// dq_Kmat[r] = cm_mul(dq_Kmat[r], dscale[r]); -// } - -// #pragma unroll -// for(int k = 0; k < num_K; k++) -// St2.row(k) = cm_dpas( -// 0, -// rQ[0].format(), -// Kmat[k].format()); - -// #pragma unroll -// for(int ri = 1; ri < head_size/REG_K; ri++) { -// cm_prefetch(b2dKV_prefetch.set_block_x(ri*REG_K)); -// //cm_load(Kmat.format(), b2dKV.set_block_x(ri*REG_K)); -// cm_load(quan_Kmat.format(), b2dKV.set_block_x(ri*REG_K)); -// #pragma unroll -// for(int r = 0; r < kv_step; r++) { -// dq_Kmat[r] = quan_Kmat[r] - zp[r]; -// dq_Kmat[r] = cm_mul(dq_Kmat[r], dscale[r]); -// } -// #pragma unroll -// for(int k = 0; k < num_K; k++) { -// St2.row(k) = cm_dpas( -// St2.row(k), -// rQ[ri].format(), -// Kmat[k].format()); -// } -// } -// } -// if constexpr (use_causal_mask) { -// // since kv_step == q_step == 16, causal_left is n*kv_step -// if (causal_left == 0) { -// apply_causal_mask<1>(St); -// } else if (causal_left < 0) { -// St = -3.4e38f; -// } -// causal_left -= kv_step; -// } else { -// int kv_tokens = kv_stop - kv_pos; -// // LSC ensures no overflow-access, but mask off k-tails attn-score is still required -// for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; -// } - -// //show(St); -// auto max_comp = online_softmax_update(St, cur_max, cur_sum); - -// matrix P; -// Transpose2DMatrix(St, P); - -// b2dKV_prefetch.set_base_ptr(reinterpret_cast(v_cache_base+prefetch_block_id*quan_blk_stride)); -// b2dKV_prefetch.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); - -// b2dKV.set_base_ptr(reinterpret_cast(v_cache_base+cur_block_id*quan_blk_stride)); -// b2dKV.set_block_y(kv_pos%CMPA_BLOCK_SZ); - - -// cm_svm_block_read(reinterpret_cast(v_cache_base+dscale_offset), dscale); -// cm_svm_block_read(reinterpret_cast(v_cache_base+dscale_offset+CMPA_BLOCK_SZ*sizeof(half)), zp); - -// { -// matrix VmatVNNI2; -// matrix Vmat; -// auto quanVmat = Vmat.format().row(1).format(); -// int kv_tokens = kv_stop - kv_pos; -// if (kv_pos == 0) { -// // ugemm_PV0(slm_V, P, rO, slm_offset); -// auto P2 = P.format(); -// #pragma unroll -// for(int k = 0, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { -// cm_prefetch(b2dKV_prefetch.set_block_x(k)); -// cm_load(quanVmat.format(), b2dKV.set_block_x(k)); -// #pragma unroll -// for(int r = 0; r < kv_step;r++) { -// Vmat[r] = quanVmat[r]-zp[r]; -// Vmat[r] = cm_mul(Vmat[r], dscale[r]); -// } -// for(int r = kv_step-1; r>=kv_tokens;r--) { -// Vmat[r] = 0; -// } - -// prepackAsVNNIWidth2(Vmat, VmatVNNI2); - -// #pragma unroll -// for(int p = 0; p < num_P_tiles; p++) { -// rO[ri + p] = cm_dpas( -// 0, -// VmatVNNI2.format(), -// P2.row(p).format()); -// } -// } -// } -// else { -// //ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); -// auto P2 = P.format(); -// #pragma unroll -// for(int k = 0, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { -// cm_prefetch(b2dKV_prefetch.set_block_x(k)); -// cm_load(quanVmat.format(), b2dKV.set_block_x(k)); - -// #pragma unroll -// for(int r = 0; r < kv_step;r++) { -// Vmat[r] = quanVmat[r]-zp[r]; -// Vmat[r] = cm_mul(Vmat[r], dscale[r]); -// } -// for(int r = kv_step-1; r>=kv_tokens;r--) { -// Vmat[r] = 0; -// } - -// prepackAsVNNIWidth2(Vmat, VmatVNNI2); -// //# compensate cur_O -// // matrix rO; -// #pragma unroll -// for(int p = 0; p < num_P_tiles; p++) { -// auto cO = rO[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[ri + p] = cm_dpas( -// rO[ri + p].format(), -// VmatVNNI2.format(), -// P2.row(p).format()); -// } -// } -// } -// } -// } -// if (q_tokens_left == 0) return; - -// //# save cur_O/cur_sum.transpose(0, 1) -// matrix cur_O_f16; -// cur_sum = cm_inv(cur_sum); - -// lsc::block_2d_desc b2dO(o_base, q_tokens_left - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); - -// #pragma unroll -// for(int k = 0, 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[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)); -// } -// } - template void sdpa_kernel_lsc( uint slm_K, @@ -647,229 +398,6 @@ void sdpa_kernel_lsc_prefetch( } } -// Item 9: process 2 Q-rows per thread — amortize per-kv softmax cost over 2x DPAS work. -// q_len is the number of Q tokens for the FIRST row; the second row is at q_base + q_pitch. -// The caller must ensure at most 2*q_step tokens are dispatched to this thread. -template -void sdpa_kernel_lsc_prefetch_q2( - int wg_local_id, - int q_start, - int kv_stop, - int q_len, // tokens for this thread, up to 2*q_step - int kv_len, - svmptr_t q_base [[type("svmptr_t")]], - svmptr_t k_base [[type("svmptr_t")]], - svmptr_t v_base [[type("svmptr_t")]], - svmptr_t o_base [[type("svmptr_t")]]) { - - constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); - constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; - - // Two independent softmax states for the two Q rows. - vector cur_max_a, cur_sum_a; - vector cur_max_b, cur_sum_b; - cur_max_a = -3e38f; cur_sum_a = 0; - cur_max_b = -3e38f; cur_sum_b = 0; - - constexpr int num_P_tiles = REG_N / REG_M; - // Two rQ halves: rQ_a for first row, rQ_b for second row - matrix rQ_a; - matrix rQ_b; - // Two rO halves - matrix rO_a; - matrix rO_b; - rO_a = 0.0f; // See rO note above: avoid 0*NaN from stale GRF bits. - rO_b = 0.0f; // See rO note above: avoid 0*NaN from stale GRF bits. - - int q_len_a = q_len; - if (q_len_a < 0) q_len_a = 0; - if (q_len_a > q_step) q_len_a = q_step; - - int q_len_b = q_len - q_step; - if (q_len_b < 0) q_len_b = 0; - if (q_len_b > q_step) q_len_b = q_step; - - // Load rQ_a (first Q row) - if (q_len_a > 0) { - lsc::block_2d_desc b2dQa(reinterpret_cast(q_base), q_len_a - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); - #pragma unroll - for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { - cm_load(rQ_a[ri].format(), b2dQa.set_block_x(k)); - rQ_a[ri].format() = cm_mul(rQ_a[ri].format(), (half)scale_factor); - } - } - // Load rQ_b (second Q row block: q_step tokens further = q_step * q_pitch bytes) - if (q_len_b > 0) { - lsc::block_2d_desc b2dQb(reinterpret_cast(q_base + (svmptr_t)q_step * q_pitch), q_len_b - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); - #pragma unroll - for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { - cm_load(rQ_b[ri].format(), b2dQb.set_block_x(k)); - rQ_b[ri].format() = cm_mul(rQ_b[ri].format(), (half)scale_factor); - } - } - - lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - - static_assert(wg_local_size == 16); - lsc::block_2d_desc prefetch_K(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc prefetch_V(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - - int causal_left = q_start; - - // Warm-up prefetch for kv_pos=0 K+V - prefetch_K.set_block_y(wg_local_id); - prefetch_V.set_block_y(wg_local_id); - cm_prefetch(prefetch_K.set_block_x(0)); - cm_prefetch(prefetch_V.set_block_x(0)); - #pragma unroll - for(int ri = 1; ri < padded_head_size/REG_K; ri++) { - cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); - cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); - } - - for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step, - k_base += kv_step * kv_pitch, - v_base += kv_step * kv_pitch) { - - // K phase: compute St_a = K × rQ_a and St_b = K × rQ_b - matrix St_a; - matrix St_b; - { - constexpr int num_K = kv_step/REG_M; - auto St_a2 = St_a.format(); - auto St_b2 = St_b.format(); - - matrix Kmat; - - prefetch_K.set_block_y(wg_local_id + kv_pos + kv_step); - prefetch_V.set_block_y(wg_local_id + kv_pos); - cm_prefetch(prefetch_K.set_block_x(0)); - cm_prefetch(prefetch_V.set_block_x(0)); - - b2dK.set_block_y(kv_pos); - cm_load(Kmat.format(), b2dK.set_block_x(0)); - #pragma unroll - for(int k = 0; k < num_K; k++) { - St_a2.row(k) = cm_dpas( - 0, rQ_a[0].format(), Kmat[k].format()); - St_b2.row(k) = cm_dpas( - 0, rQ_b[0].format(), Kmat[k].format()); - } - - #pragma unroll - for(int ri = 1; ri < padded_head_size/REG_K; ri++) { - cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); - cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); - cm_load(Kmat.format(), b2dK.set_block_x(ri*REG_K)); - #pragma unroll - for(int k = 0; k < num_K; k++) { - St_a2.row(k) = cm_dpas( - St_a2.row(k), rQ_a[ri].format(), Kmat[k].format()); - St_b2.row(k) = cm_dpas( - St_b2.row(k), rQ_b[ri].format(), Kmat[k].format()); - } - } - } - - // Tail masking - { - int kv_tokens = kv_stop - kv_pos; - for(int p = kv_tokens; p < kv_step; p++) { - St_a[p] = -3.4e38f; - St_b[p] = -3.4e38f; - } - } - - auto max_comp_a = online_softmax_update(St_a, cur_max_a, cur_sum_a); - auto max_comp_b = online_softmax_update(St_b, cur_max_b, cur_sum_b); - - matrix P_a, P_b; - Transpose2DMatrix(St_a, P_a); - Transpose2DMatrix(St_b, P_b); - - // V phase: unified path (max_comp=0 at kv_pos=0 zeroes rO) - b2dV.set_block_y(kv_pos); - { - auto Pa2 = P_a.format(); - auto Pb2 = P_b.format(); - #pragma unroll - for(int k = 0, ri = 0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { - matrix Vmat; - cm_load(Vmat.format(), b2dV.set_block_x(k)); - // Rescale and DPAS for rO_a - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO_a[ri + p].format(); - #pragma unroll - for(int r = 0; r < REG_M; r++) - cO.row(r) = cm_mul(cO.row(r), max_comp_a[r + p*REG_M]); - } - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - rO_a[ri + p] = cm_dpas( - rO_a[ri + p].format(), Vmat.format(), Pa2.row(p).format()); - } - // Rescale and DPAS for rO_b - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO_b[ri + p].format(); - #pragma unroll - for(int r = 0; r < REG_M; r++) - cO.row(r) = cm_mul(cO.row(r), max_comp_b[r + p*REG_M]); - } - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - rO_b[ri + p] = cm_dpas( - rO_b[ri + p].format(), Vmat.format(), Pb2.row(p).format()); - } - } - } - } - - // Store output for first Q row - if (q_len_a > 0) { - matrix cur_O_f16; - cur_sum_a = cm_inv(cur_sum_a); - lsc::block_2d_desc b2dO(o_base, q_len_a - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); - #pragma unroll - for(int k = 0, ri=0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO_a[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_a[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 output for second Q row - if (q_len_b > 0) { - matrix cur_O_f16; - cur_sum_b = cm_inv(cur_sum_b); - lsc::block_2d_desc b2dO(o_base + (svmptr_t)q_step * o_pitch, q_len_b - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); - #pragma unroll - for(int k = 0, ri=0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO_b[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_b[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)); - } - } -} - #else // CM_HAS_LSC_UNTYPED_2D template From ae0ec396e188d4a843beda58fa0fff5f248fc7f9 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 14 Jul 2026 13:26:01 +0800 Subject: [PATCH 08/16] PA: CMFLA_USE_TREE_SOFTMAX and CMFLA_Q_SCALED_BY_LOG2 decouple --- .../impls/cm/include/cm_attention_common.hpp | 158 ++++++++---------- .../src/graph/impls/cm/include/cm_pa_xe1.hpp | 10 +- .../src/graph/impls/cm/include/cm_pa_xe2.hpp | 26 +-- 3 files changed, 77 insertions(+), 117 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp index 712ba20c42d614..7dc8c6ca0d27ab 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp @@ -7,8 +7,6 @@ //# CM-compiler is C++17 static_assert(__cplusplus >= 201703L); -//# static_assert(__cplusplus >= 202002L); -//# static_assert(__cplusplus >= 202302L); #define SystolicDepth 8 #define RepeatCount 8 @@ -23,15 +21,41 @@ static_assert(__cplusplus >= 201703L); #define kv_step REG_K #define q_step REG_N +// Q pre-scale math constants: +// scale_factor = 1 / sqrt(head_size) (provided by CMFLA_SCALE_FACTOR at JIT time) +// log2e = log2(e) (fixed) constexpr float scale_factor = CMFLA_SCALE_FACTOR; constexpr float log2e = 1.4426950408889634f; -constexpr float q_scale_factor = scale_factor * log2e; + +// JIT const: chooses the domain the softmax score St = K @ Q^T lives in. +// 1 => Q is pre-multiplied by log2e so the online softmax can call cm_exp (== exp2) +// directly (Item 13 in OPTIMIZE_PLAN, drops one *log2e off the softmax critical path). +// 0 => Q keeps the natural-log domain and the softmax multiplies by log2e internally. +// Set at JIT time (typically via -DCMFLA_Q_SCALED_BY_LOG2=... from the host Python script); +// defaults to 1. The value MUST match how the kernel actually pre-scales Q -- both the Q +// load (via q_prescale) and the softmax templates read this macro, so keeping them in +// sync is enough to prevent log2e from being applied zero or two times. +#ifndef CMFLA_Q_SCALED_BY_LOG2 +#define CMFLA_Q_SCALED_BY_LOG2 1 +#endif + +// Single Q pre-scale constant used by every SDPA/PA kernel at Q load time. +constexpr float q_prescale = CMFLA_Q_SCALED_BY_LOG2 ? (scale_factor * log2e) : scale_factor; + +// JIT const: 1 => use the balanced binary-tree reduction (depth log2(rows)) in the online +// softmax; requires rows to be a power of two. +// 0 => use the linear-chain reduction (default). +// Applies uniformly to the SDPA (flashattn) and PA (pageatten) kernel families via the +// `cm_online_softmax_update` dispatch macro defined below. +#ifndef CMFLA_USE_TREE_SOFTMAX +#define CMFLA_USE_TREE_SOFTMAX 0 +#endif static_assert(q_step == 16 || q_step == 8); static_assert(kv_step == 16); static_assert(CM_HAS_DPAS); -#define DEBUG_SHOW 1 +#define DEBUG_SHOW 0 #if !DEBUG_SHOW template void show(const matrix mat, bool isfloat=true) { @@ -219,9 +243,6 @@ inline matrix ugemm_KQ(uint slm_K, matrix_ref Kmat; cm_slm_block_read(slm_K, GENX_NONE, slm_offset, Kmat.format()); - // if (cm_local_id(2) == 3 && cm_group_id(2) == 0) { - // show(Kmat.format()); - // } #pragma unroll for(int k = 0; k < num_K; k++) St2.row(k) = cm_dpas(0, Qt[0].format(), Kmat[k].format()); @@ -246,16 +267,12 @@ inline void ugemm_PV0(uint slm_V, matrix_ref P, matrix_ref Vmat; cm_slm_block_read(slm_V, GENX_NONE, slm_offset + REG_K*k*sizeof(half), Vmat.format()); - // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { - // show(Vmat.format()); - // } #pragma unroll for(int p = 0; p < num_P_tiles; p++) { rO[ri + p] = cm_dpas( 0, Vmat.format(), P2.row(p).format()); - //show(rO[ri + p].format()); } } } @@ -270,9 +287,6 @@ inline void ugemm_PV1(uint slm_V, matrix_ref P, vector_ref Vmat; cm_slm_block_read(slm_V, GENX_NONE, slm_offset + REG_K*k*sizeof(half), Vmat.format()); - // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { - // show(Vmat.format()); - // } #pragma unroll for(int p = 0; p < num_P_tiles; p++) { auto cO = rO[ri + p].format(); @@ -281,7 +295,6 @@ inline void ugemm_PV1(uint slm_V, matrix_ref P, vector_ref(cO.row(r), max_comp[r + p*REG_M]); } - //show(rO[ri].format()); #pragma unroll for(int p = 0; p < num_P_tiles; p++) { @@ -289,12 +302,11 @@ inline void ugemm_PV1(uint slm_V, matrix_ref P, vector_ref(), Vmat.format(), P2.row(p).format()); - //if (kv_pos == args_verbose) show(rO[ri + p].format()); } - // if (kv_pos == args_verbose) show(cur_O.format()); } } +// Online-softmax update (linear-chain max/sum reduction). template vector online_softmax_update(matrix_ref St, vector_ref cur_max, vector_ref cur_sum) { vector new_max_t; @@ -303,11 +315,10 @@ vector online_softmax_update(matrix_ref St, vector_r new_max_t = cm_max(new_max_t, cur_max); // Pt = torch.exp(St - new_max) - constexpr float _log2e = 1.4426950408889634f; -#ifdef ABLATE_NO_EXP - for(int r = 0; r < St.n_rows(); r++) St[r] = (St[r] - new_max_t)*_log2e; +#if CMFLA_Q_SCALED_BY_LOG2 + for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp(St[r] - new_max_t); #else - for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp((St[r] - new_max_t)*_log2e); + for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp((St[r] - new_max_t)*log2e); #endif vector row_sum_t; @@ -315,7 +326,11 @@ vector online_softmax_update(matrix_ref St, vector_r for(int r = 2; r < St.n_rows(); r++) row_sum_t = cm_add(row_sum_t, St[r]); vector max_comp; - max_comp = cm_exp((cur_max - new_max_t)*_log2e); +#if CMFLA_Q_SCALED_BY_LOG2 + max_comp = cm_exp(cur_max - new_max_t); +#else + max_comp = cm_exp((cur_max - new_max_t)*log2e); +#endif cur_sum = cm_mul(cur_sum, max_comp); cur_sum = cm_add(cur_sum, row_sum_t); cur_max = new_max_t; @@ -342,8 +357,13 @@ CM_INLINE vector online_softmax_update_tree(matrix_ref(t.row(r), t.row(r + stride)); new_max_t = cm_max(t.row(0), cur_max); } +#if CMFLA_Q_SCALED_BY_LOG2 #pragma unroll for (int r = 0; r < rows; r++) St[r] = cm_exp(St[r] - new_max_t); +#else + #pragma unroll + for (int r = 0; r < rows; r++) St[r] = cm_exp((St[r] - new_max_t) * log2e); +#endif vector row_sum_t; { @@ -359,88 +379,44 @@ CM_INLINE vector online_softmax_update_tree(matrix_ref max_comp; +#if CMFLA_Q_SCALED_BY_LOG2 max_comp = cm_exp(cur_max - new_max_t); +#else + max_comp = cm_exp((cur_max - new_max_t) * log2e); +#endif cur_sum = cm_mul(cur_sum, max_comp); cur_sum = cm_add(cur_sum, row_sum_t); cur_max = new_max_t; return max_comp; } -// Transpose a [16,16] float score tile (kv x q) into a half [16,16] P tile (q x kv) for -// the P@V matmul. Casting float->half first (one vectorized cm_mul-free copy) lets the -// 16x16 shuffle network run at half width — roughly halving the mov count vs transposing -// the float tile directly through Transpose_16x16. +// Dispatch macro used by every SDPA/PA kernel; selects linear vs tree reduction. +#if CMFLA_USE_TREE_SOFTMAX +#define cm_online_softmax_update(St, cur_max, cur_sum) online_softmax_update_tree(St, cur_max, cur_sum) +#else +#define cm_online_softmax_update(St, cur_max, cur_sum) online_softmax_update(St, cur_max, cur_sum) +#endif + +// Transpose a float score tile (kv x q) into a half P tile (q x kv) for the P@V matmul. +// Casting float->half first (one vectorized cm_mul-free copy) lets the shuffle network +// run at half width -- roughly halving the mov count vs transposing the float tile +// directly through Transpose_*x*. +// +// Two overloads cover both Xe generations (kv_step is always 16; q_step = REG_N): +// * [16,16] -> [16,16] for Xe2 (q_step = 16) +// * [16, 8] -> [ 8,16] for Xe1 (q_step = 8) CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { matrix Sh; #pragma unroll for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); Transpose_16x16(Sh.select<16,1,16,1>(0,0), P); } - -#ifdef CM_HAS_LSC_UNTYPED_2D - #define cm_load_normal cm_load - #define cm_load_transpose cm_load - #define cm_load_vnni cm_load - #define cm_store_normal cm_store - // // simulation of LSC API using SVM API - // template - // inline void cm_load_normal(vector_ref Res, const lsc::block_2d_desc &Desc, int16_t Pred = 1) { - // static_assert(NBlocks == 1); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_read(base + i * pitch, Res.select(i*BlockW)); - // } - // } - - // template - // inline void cm_load_transpose(vector_ref Res, const lsc::block_2d_desc &Desc, int16_t Pred = 1) { - // static_assert(NBlocks == 1); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // matrix temp; - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_read(base + i * pitch, temp[i]); - // } - // Transpose2DMatrix(temp, Res.format()); - // } - - // // in VNNI case, NBlocks is increasing along X dimension (increase cache-line usage) - // template - // inline void cm_load_vnni(vector_ref Res, const lsc::block_2d_desc &Desc, int16_t Pred = 1) { - // static_assert(NBlocks == 1 || NBlocks == 2); - // // each block must be a full XMX B matrix - // static_assert(BlockH == REG_K); - // static_assert(BlockW == REG_N); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // matrix temp; - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_read(base + i * pitch, temp[i]); - // } - - // auto out_vnni = Res.format(); - // #pragma unroll - // for(int i = 0; i < NBlocks; i ++) { - // out_vnni.select(i*(BlockH/2), 0) = temp.select(0, i*BlockW); - // out_vnni.select(i*(BlockH/2), 1) = temp.select(1, i*BlockW); - // } - // } - - // template - // inline void cm_store_normal(const lsc::block_2d_desc &Desc, vector_ref Res) { - // static_assert(NBlocks == 1); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_write(base + i * pitch, Res.select(i*BlockW)); - // } - // } -#endif +CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { + matrix Sh; + #pragma unroll + for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); + Transpose2DMatrix(Sh.select<16,1,8,1>(0,0), P); +} //=============================================================================================== template diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp index ad98fc979928e8..6eec4e4a8f8635 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp @@ -97,7 +97,7 @@ void pa_lsc_u8( CacheHint::Cached, CacheHint::Cached>(q_gather, gather_offsets, gather_pred); rQ[ri].format() = gathered; - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } #if KV_CACHE_COMPRESSION == 1 @@ -273,7 +273,6 @@ void pa_lsc_u8( cm_fence(CM_LOCAL_BARRIER); cm_sbarrier(0); - //if (kv_pos > 1024000) if (kv_pos + kv_step < kv_stop) cm_sbarrier(1); load_slm_KV(kv_pos + kv_step*2); @@ -300,7 +299,7 @@ void pa_lsc_u8( // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; } - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); @@ -480,7 +479,7 @@ void pa_kernel_lsc_prefetch_f16( CacheHint::Cached, CacheHint::Cached>(q_gather, gather_offsets, gather_pred); rQ[ri].format() = gathered; - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } constexpr int blk_stride = CMFLA_NUM_KV_HEADS * CMFLA_HEAD_SIZE*CMPA_BLOCK_SZ; @@ -568,8 +567,7 @@ void pa_kernel_lsc_prefetch_f16( for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; } - // show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp index 082cf7013929aa..bc60e3077ed5a6 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp @@ -10,18 +10,6 @@ #define CMPA_WG_SEQ_LEN 0 #endif -// Item 15: use balanced binary-tree reduction in online_softmax_update (log2 depth vs linear). -// Set to 0 to disable (use linear chain), 1 to enable. -#ifndef CMPA_USE_TREE_SOFTMAX -#define CMPA_USE_TREE_SOFTMAX 1 -#endif - -#if CMPA_USE_TREE_SOFTMAX -#define pa_softmax_update(St, cur_max, cur_sum) online_softmax_update_tree(St, cur_max, cur_sum) -#else -#define pa_softmax_update(St, cur_max, cur_sum) online_softmax_update(St, cur_max, cur_sum) -#endif - #ifndef SPARSE_BLOCK_SIZE #error "SPARSE_BLOCK_SIZE must be defined" #endif @@ -138,7 +126,7 @@ void pa_lsc_u8( #pragma unroll for (int k = 0, ri = 0; k < process_head_size / 2; k += REG_K / 2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(worker_offset / 2 + k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } @@ -410,7 +398,7 @@ void pa_lsc_u8( } int kv_tokens = kv_stop - kv_pos; for (int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - auto max_comp = pa_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; transpose_St_to_P_half(St, P); @@ -644,7 +632,7 @@ void pa_lsc_u8( } int kv_tokens = kv_stop - kv_pos; for (int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - auto max_comp = pa_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; transpose_St_to_P_half(St, P); @@ -859,7 +847,7 @@ void pa_kernel_lsc_prefetch_f16( #pragma unroll for(int k = 0, ri = 0; k < process_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(worker_offset / 2 + k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } @@ -979,8 +967,7 @@ void pa_kernel_lsc_prefetch_f16( int kv_tokens = kv_stop - kv_pos; // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - //show(St); - auto max_comp = pa_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; transpose_St_to_P_half(St, P); @@ -1283,8 +1270,7 @@ void pa_kernel_lsc_prefetch_f16( int kv_tokens = kv_stop - kv_pos; // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - //show(St); - auto max_comp = pa_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; transpose_St_to_P_half(St, P); From aecbb8fc2e650689d116dc9fbdfa6c5b6a21849f Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 14 Jul 2026 13:36:20 +0800 Subject: [PATCH 09/16] Change vlsdpa kernel signature --- .../src/graph/impls/cm/cm_sdpa_vlen.cm | 69 +++++++++++---- .../graph/impls/cm/include/cm_sdpa_common.hpp | 85 ++++++++----------- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 35 +++++++- 3 files changed, 121 insertions(+), 68 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm b/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm index 87ac3c56222d2e..75e15ba1d4b081 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm +++ b/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm @@ -15,7 +15,7 @@ * limitations under the License. *******************************************************************************/ namespace KERNEL_NAME { - + #include "cm_sdpa_common.hpp" #ifdef CM_HAS_LSC_UNTYPED_2D @@ -36,10 +36,32 @@ _GENX_MAIN_ void KERNEL_NAME( SurfaceIndex query [[type("buffer_t")]], SurfaceIndex key [[type("buffer_t")]], SurfaceIndex value [[type("buffer_t")]], - SurfaceIndex output [[type("buffer_t")]], + SurfaceIndex output [[type("buffer_t")]], #endif int* cu_seqlens [[type("svmptr_t")]], - int need_wg_mapping + int need_wg_mapping, + // Memory-layout contract for Q/K/V inputs (must be preserved by upstream passes): + // Only the token (outer) axis is allowed to be discontinuous. The kernel treats + // each input as a 3-D tensor [tokens, heads, head_size] and hardcodes the two + // inner strides: + // - per-head stride == head_size (heads contiguous within a token) + // - per-element stride == 1 (head_size dim is contiguous) + // Only the outer per-token stride is layout-driven via {q,k,v}_token_pitch below. + // If prepare_buffer_fusing ever propagates padding onto the head or head_size axis + // of a vl_sdpa input layout, this kernel will silently address the wrong memory. + // + // Element-count offsets from the SVM base to the slice start of each input, taken + // from input_layout.get_linear_offset(). Non-zero when an in-place crop (e.g. the + // TransposeSplitMatcher axis=1 pattern) leaves the SVM base pointing at the packed + // parent buffer and shifts the logical view via layout padding on the token axis. + int token_offset_q, + int token_offset_k, + int token_offset_v, + // Per-input token strides in elements, from input_layout.get_pitches()[0]. Allows + // mixed layouts (e.g. Q/K contiguous post-RoPE, V still viewing the padded parent). + int q_token_pitch, + int k_token_pitch, + int v_token_pitch ) { constexpr int num_heads = CMFLA_NUM_HEADS; constexpr int head_size = CMFLA_HEAD_SIZE; @@ -107,8 +129,12 @@ _GENX_MAIN_ void KERNEL_NAME( kv_start = cu_seqlens[wg_id]; kv_seq_len = cu_seqlens[wg_id + 1] - kv_start; } - uint qo_offset = (q_start*num_heads + h)*head_size; - uint kv_offset = (kv_start*num_kv_heads + hkv)*head_size; + // Input offsets are computed with per-input runtime token pitches. + // Output is always contiguous [q_len, num_heads, head_size]. + uint q_in_offset = q_start * q_token_pitch + h * head_size; + uint k_in_offset = kv_start * k_token_pitch + hkv * head_size; + uint v_in_offset = kv_start * v_token_pitch + hkv * head_size; + uint o_offset = (q_start * num_heads + h) * head_size; #if USE_LSC == 1 #if USE_LSC_PREFETCH @@ -118,10 +144,13 @@ _GENX_MAIN_ void KERNEL_NAME( kv_seq_len, q_len, kv_seq_len, //kv_len, - reinterpret_cast(query + qo_offset), - reinterpret_cast(key + kv_offset), - reinterpret_cast(value + kv_offset), - reinterpret_cast(output + qo_offset)); + q_token_pitch * sizeof(half), + k_token_pitch * sizeof(half), + v_token_pitch * sizeof(half), + reinterpret_cast((query + token_offset_q) + q_in_offset), + reinterpret_cast((key + token_offset_k) + k_in_offset), + reinterpret_cast((value + token_offset_v) + v_in_offset), + reinterpret_cast(output + o_offset)); #else sdpa_kernel_lsc( slm_K, @@ -132,10 +161,13 @@ _GENX_MAIN_ void KERNEL_NAME( kv_seq_len, //kv_stop, q_len, //q_len, kv_seq_len, //kv_len, - reinterpret_cast(query + qo_offset), - reinterpret_cast(key + kv_offset), - reinterpret_cast(value + kv_offset), - reinterpret_cast(output + qo_offset)); + q_token_pitch * sizeof(half), + k_token_pitch * sizeof(half), + v_token_pitch * sizeof(half), + reinterpret_cast((query + token_offset_q) + q_in_offset), + reinterpret_cast((key + token_offset_k) + k_in_offset), + reinterpret_cast((value + token_offset_v) + v_in_offset), + reinterpret_cast(output + o_offset)); #endif #else sdpa_kernel( @@ -151,10 +183,13 @@ _GENX_MAIN_ void KERNEL_NAME( key, value, output, - qo_offset * sizeof(half), - kv_offset * sizeof(half), - kv_offset * sizeof(half), - qo_offset * sizeof(half) + q_token_pitch * sizeof(half), + k_token_pitch * sizeof(half), + v_token_pitch * sizeof(half), + (token_offset_q + q_in_offset) * sizeof(half), + (token_offset_k + k_in_offset) * sizeof(half), + (token_offset_v + v_in_offset) * sizeof(half), + o_offset * sizeof(half) ); #endif } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index a130ca31ffe1e4..02a13f033d4555 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -35,14 +35,15 @@ void sdpa_kernel_lsc( int kv_stop, int q_len, int kv_len, + uint q_pitch_bytes, + uint k_pitch_bytes, + uint v_pitch_bytes, svmptr_t q_base [[type("svmptr_t")]], svmptr_t k_base [[type("svmptr_t")]], svmptr_t v_base [[type("svmptr_t")]], svmptr_t o_base [[type("svmptr_t")]]) { constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); // round up head_size to multiple of 16 // block_2d_desc will automatically handle the tailing block constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; @@ -64,16 +65,16 @@ void sdpa_kernel_lsc( if (q_tokens_left > q_step) q_tokens_left = q_step; if (q_tokens_left > 0) { - lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); + lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch_bytes - 1, 0, 0); #pragma unroll for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } - lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, k_pitch_bytes - 1, 0, 0); + lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, v_pitch_bytes - 1, 0, 0); int causal_left = q_start; @@ -109,8 +110,8 @@ void sdpa_kernel_lsc( cm_sbarrier(1); for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step, - k_base += kv_step * kv_pitch, - v_base += kv_step * kv_pitch, + k_base += kv_step * k_pitch_bytes, + v_base += kv_step * v_pitch_bytes, slm_buff_id_read ++) { // load0, load1, signal1, @@ -126,7 +127,6 @@ void sdpa_kernel_lsc( cm_fence(CM_LOCAL_BARRIER); cm_sbarrier(0); - //if (kv_pos > 1024000) // for debugging if (kv_pos + kv_step < kv_stop) cm_sbarrier(1); @@ -151,8 +151,7 @@ void sdpa_kernel_lsc( for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; } - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); @@ -195,14 +194,15 @@ void sdpa_kernel_lsc_prefetch( int kv_stop, int q_len, int kv_len, + uint q_pitch_bytes, + uint k_pitch_bytes, + uint v_pitch_bytes, svmptr_t q_base [[type("svmptr_t")]], svmptr_t k_base [[type("svmptr_t")]], svmptr_t v_base [[type("svmptr_t")]], svmptr_t o_base [[type("svmptr_t")]]) { constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); // round up head_size to multiple of 16 // block_2d_desc will automatically handle the tailing block constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; @@ -225,24 +225,20 @@ void sdpa_kernel_lsc_prefetch( if (q_tokens_left > q_step) q_tokens_left = q_step; if (q_tokens_left > 0) { - // Fold log2(e) into the Q pre-scale so St = K@Q^T lands in the log2 domain; the - // online softmax then uses cm_exp (== exp2) directly, dropping a *log2e per St - // element (16 muls/tile off the softmax critical path). Math is identical. - constexpr float qscale = scale_factor * 1.4426950408889634f; - lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); + lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch_bytes - 1, 0, 0); #pragma unroll for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)qscale); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } - lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, k_pitch_bytes - 1, 0, 0); + lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, v_pitch_bytes - 1, 0, 0); static_assert(wg_local_size == 16); - lsc::block_2d_desc prefetch_K(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc prefetch_V(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc prefetch_K(k_base, kv_stop - 1, head_size*sizeof(half) - 1, k_pitch_bytes - 1, 0, 0); + lsc::block_2d_desc prefetch_V(v_base, kv_stop - 1, head_size*sizeof(half) - 1, v_pitch_bytes - 1, 0, 0); int causal_left = q_start; @@ -324,7 +320,7 @@ void sdpa_kernel_lsc_prefetch( } // ---- one online-softmax update over the whole block ---- - auto max_comp = online_softmax_update_tree(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); // ---- rescale rO (skip first iter only when KV blocking amortizes the branch) ---- // The kv_base=0 rescale is mathematically redundant for every head size: max_comp=0 @@ -414,14 +410,15 @@ void sdpa_kernel( SurfaceIndex key [[type("buffer_t")]], SurfaceIndex value [[type("buffer_t")]], SurfaceIndex output [[type("buffer_t")]], + uint q_pitch_bytes, + uint k_pitch_bytes, + uint v_pitch_bytes, uint q_off, uint k_off, uint v_off, uint o_off) { constexpr int padded_head_size = (head_size + 16 - 1) / 16 *16; constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); vector cur_max; vector cur_sum; @@ -442,11 +439,11 @@ void sdpa_kernel( // load as many as possible given one address if constexpr (head_size == 128 || head_size == 64) { matrix QmatI32; - cm_load_2d(QmatI32, query, q_off, q_pitch); + cm_load_2d(QmatI32, query, q_off, q_pitch_bytes); #pragma unroll for(int k = 0, ri = 0; k < head_size/2; k += REG_K/2, ri++) { Transpose2DMatrix(QmatI32.select(0, k), rQ[ri].format()); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } else { constexpr int num_full_blocks = head_size / REG_K; @@ -455,9 +452,9 @@ void sdpa_kernel( for(; i < num_full_blocks; i++) { matrix QmatI32; int k = i * REG_K; - cm_load_2d(QmatI32, query, q_off + k * sizeof(uint) / 2, q_pitch); + cm_load_2d(QmatI32, query, q_off + k * sizeof(uint) / 2, q_pitch_bytes); Transpose2DMatrix(QmatI32, rQ[i].format()); - rQ[i].format() = cm_mul(rQ[i].format(), (half)scale_factor); + rQ[i].format() = cm_mul(rQ[i].format(), (half)q_prescale); } // if with tail, load with head_size_tail @@ -465,9 +462,9 @@ void sdpa_kernel( if constexpr (head_size % REG_K > 0) { int k = num_full_blocks * REG_K; matrix QmatI32; - cm_load_2d_with_tail(QmatI32, query, q_off + k * sizeof(half), q_pitch); + cm_load_2d_with_tail(QmatI32, query, q_off + k * sizeof(half), q_pitch_bytes); Transpose2DMatrix(QmatI32, rQ[num_full_blocks].format()); - rQ[num_full_blocks].format() = cm_mul(rQ[num_full_blocks].format(), (half)scale_factor); + rQ[num_full_blocks].format() = cm_mul(rQ[num_full_blocks].format(), (half)q_prescale); } } } @@ -482,7 +479,6 @@ void sdpa_kernel( int slm_buff_id_read = 0; auto load_slm_KV = [&](int kv_pos) { - //if (kv_pos < 1024000) return; int kv_tokens = kv_stop - kv_pos; if (kv_tokens <= 0) return; // Calculate valid rows for this block (used to zero out garbage data) @@ -492,14 +488,13 @@ void sdpa_kernel( // non-tail branch is faster if (wg_local_id < local_size/2) { - //if (kv_pos > 1024000) { matrix temp; // head_size is split into blocks of constexpr int num_full_blocks = head_size / REG_K; int i = wg_local_id; for (; i < num_full_blocks; i += local_size / 2) { int k = i * REG_K; - cm_load_2d(temp, key, k_off + k * sizeof(half), kv_pitch); + cm_load_2d(temp, key, k_off + k * sizeof(half), k_pitch_bytes); // Zero out unused K rows to prevent NaN from garbage data in KV cache // (Similar approach as PA kernel: cm_pa_common.hpp) for (int r = kv_valid_rows; r < kv_step; r++) @@ -512,7 +507,7 @@ void sdpa_kernel( // following code will be optimized out when head_size_tail = 0 if constexpr (head_size % REG_K > 0) { int k = num_full_blocks * REG_K; - cm_load_2d_with_tail<2*REG_M, REG_K, head_size % REG_K>(temp, key, k_off + k * sizeof(half), kv_pitch); + cm_load_2d_with_tail<2*REG_M, REG_K, head_size % REG_K>(temp, key, k_off + k * sizeof(half), k_pitch_bytes); // Zero out unused K rows for (int r = kv_valid_rows; r < kv_step; r++) temp.row(r) = 0; @@ -521,20 +516,18 @@ void sdpa_kernel( temp.format()); } } else { - //if (kv_pos > 1024000) { // read 16x16 XMX-B matrix (1x REG_N in Xe2, 2x REG_N in Xe1) constexpr int VK_STEP = 16; static_assert((VK_STEP % REG_N) == 0); matrix temp2; matrix temp_vnni; - //b2dV.set_block_y(kv_pos); // head_size is split into blocks of constexpr int num_full_blocks = head_size / VK_STEP; int i = wg_local_id - local_size / 2; for (; i < num_full_blocks; i += local_size / 2) { int k = i * VK_STEP; - cm_load_2d(temp2, value, v_off + k * sizeof(half), kv_pitch); + cm_load_2d(temp2, value, v_off + k * sizeof(half), v_pitch_bytes); // Zero out unused V rows to prevent NaN from garbage data in KV cache // (Similar approach as PA kernel: cm_pa_common.hpp) for (int r = kv_valid_rows; r < kv_step; r++) @@ -551,7 +544,7 @@ void sdpa_kernel( // following code will be optimized out when head_size_tail = 0 if constexpr (head_size % VK_STEP > 0) { int k = num_full_blocks * VK_STEP; - cm_load_2d_with_tail(temp2, value, v_off + k * sizeof(half), kv_pitch); + cm_load_2d_with_tail(temp2, value, v_off + k * sizeof(half), v_pitch_bytes); // Zero out unused V rows for (int r = kv_valid_rows; r < kv_step; r++) temp2.row(r) = 0; @@ -564,9 +557,8 @@ void sdpa_kernel( } } } - k_off += kv_step * kv_pitch; - v_off += kv_step * kv_pitch; - // printf(" diff= %lu\n", get_clock() - clk0); + k_off += kv_step * k_pitch_bytes; + v_off += kv_step * v_pitch_bytes; }; load_slm_KV(0); @@ -597,10 +589,8 @@ void sdpa_kernel( if (kv_pos + kv_step < kv_stop) cm_sbarrier(1); - //if (kv_pos < 1024000) continue; uint slm_offset = (slm_buff_id_read & 3) * slm_buff_size; - //=========================================================== 1807 ~ 3247 //# St = k @ Qt matrix St = ugemm_KQ(slm_K, rQ, slm_offset); @@ -618,7 +608,6 @@ void sdpa_kernel( St[p] = cm_add(St[p], cmask); if (v < q_step - 1) v++; } - //if (wg_local_id == 0) show(St);return; } causal_left -= kv_step; } @@ -627,8 +616,7 @@ void sdpa_kernel( int kv_tokens = kv_stop - kv_pos; for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); @@ -654,7 +642,6 @@ void sdpa_kernel( cur_O_f16[r + p*REG_M] = cm_mul(cO.row(r), cur_sum[r + p*REG_M]); } } - // if (i == args_verbose) show(cur_O_f16); cm_store_2d(cur_O_f16, output, o_off + k*sizeof(half), o_pitch); } } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index f422be4ea1429a..34a16180412458 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -98,7 +98,17 @@ class VLSDPAGenerator : public KernelGenerator { args.push_back({ArgumentDescriptor::Types::INPUT, static_cast(params.input_layouts.size() - 1)}); // input: cu_seq_lens - args.push_back({ArgumentDescriptor::Types::SCALAR, 0}); + args.push_back({ArgumentDescriptor::Types::SCALAR, 0}); // need_wg_mapping + + // Per-input token-axis discontinuity descriptors (populated in the dispatch lambda + // from each input layout's get_linear_offset() and get_pitches()[0]). See + // cm_sdpa_vlen.cm for the exact memory-layout contract with the kernel. + args.push_back({ArgumentDescriptor::Types::SCALAR, 1}); // token_offset_q + args.push_back({ArgumentDescriptor::Types::SCALAR, 2}); // token_offset_k + args.push_back({ArgumentDescriptor::Types::SCALAR, 3}); // token_offset_v + args.push_back({ArgumentDescriptor::Types::SCALAR, 4}); // q_token_pitch + args.push_back({ArgumentDescriptor::Types::SCALAR, 5}); // k_token_pitch + args.push_back({ArgumentDescriptor::Types::SCALAR, 6}); // v_token_pitch return args; } @@ -162,7 +172,28 @@ class VLSDPAGenerator : public KernelGenerator { wgs.global = {num_q_heads, wg_count * wg_size, 1}; wgs.local = {1, wg_size, 1}; - std::vector scalars{need_wg_mapping}; + // Element-count offsets and per-input token strides taken directly from the + // input layouts. Non-zero token_offset_x arises when an in-place crop + // (e.g. TransposeSplitMatcher axis=1) leaves the SVM base pointing at the + // packed parent buffer and expresses the slice view via layout padding on + // the token axis; the kernel shifts the Q/K/V pointers accordingly. + // Contract with the CM kernel (see cm_sdpa_vlen.cm): only the token (outer) + // axis of each input layout may be padded — the head and head_size strides + // are hardcoded in the kernel and are NOT parameterized here. + int32_t token_offset_q = static_cast(params.input_layouts[0].get_linear_offset()); + int32_t token_offset_k = static_cast(params.input_layouts[1].get_linear_offset()); + int32_t token_offset_v = static_cast(params.input_layouts[2].get_linear_offset()); + const auto q_token_pitch = static_cast(params.input_layouts[0].get_pitches()[0]); + const auto k_token_pitch = static_cast(params.input_layouts[1].get_pitches()[0]); + const auto v_token_pitch = static_cast(params.input_layouts[2].get_pitches()[0]); + + std::vector scalars{need_wg_mapping, + token_offset_q, + token_offset_k, + token_offset_v, + q_token_pitch, + k_token_pitch, + v_token_pitch}; kd.params.scalars.clear(); for (auto i : scalars) { scalar_desc s; From 71db11db7f22ce13b0ce17c7fde20fbdf81945e4 Mon Sep 17 00:00:00 2001 From: cecilia peng Date: Tue, 14 Jul 2026 14:19:57 +0800 Subject: [PATCH 10/16] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp index 710c68f878e3d6..27778bab936e04 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp @@ -383,7 +383,7 @@ JitConstants PagedAttentionGeneratorMultiToken::get_jit_constants(const kernel_i jit.make("CMPA_WG_SEQ_LEN", get_wg_seq_len(params)); // Tree softmax: wins on INT8 (dequant amortises scratch, shorter dep-chain helps), // regresses ~9% on FP16 (extra scratch matrix increases register pressure). - jit.make("CMPA_USE_TREE_SOFTMAX", get_kv_compressed(params) ? 1 : 0); + jit.make("CMFLA_USE_TREE_SOFTMAX", get_kv_compressed(params) ? 1 : 0); return jit; } From 941752ffd2639a4bb3f39591439dbdc10fe5a6dc Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 14 Jul 2026 14:47:07 +0800 Subject: [PATCH 11/16] safe guard to cu_seqlens and q/k/v offsets/pitches. --- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index 34a16180412458..ee60f9b9602502 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -134,10 +134,15 @@ class VLSDPAGenerator : public KernelGenerator { // Read cu_seqlens via mem_lock — zero-cost on USM-host (PC), no GPU sync stall. const auto cu_seqlens_mem = params.memory_deps.at(params.input_layouts.size() - 1); + OPENVINO_ASSERT(cu_seqlens_mem->get_layout().data_type == ov::element::i32, + "VLSDPA expects cu_seqlens input to be i32"); mem_lock cu_seqlens_lock(cu_seqlens_mem, *params.strm); + const auto cu_seqlens_sz = cu_seqlens_lock.size(); + OPENVINO_ASSERT(cu_seqlens_sz >= 2 && cu_seqlens_sz % 2 == 0, + "VLSDPA expects cu_seqlens to contain at least two elements (start/end) and be even-sized"); size_t max_seq_len = 0; - for (size_t i = 1; i < cu_seqlens_lock.size(); i++) { + for (size_t i = 1; i < cu_seqlens_sz; i++) { auto start_idx = cu_seqlens_lock[i - 1]; auto end_idx = cu_seqlens_lock[i]; max_seq_len = std::max(max_seq_len, static_cast(end_idx - start_idx)); @@ -159,13 +164,13 @@ class VLSDPAGenerator : public KernelGenerator { if (need_wg_mapping) { wg_count = 0; const auto wg_seq_len = wg_size * q_step; - for (size_t i = 1; i < cu_seqlens_lock.size(); i++) { + for (size_t i = 1; i < cu_seqlens_sz; i++) { auto start_idx = cu_seqlens_lock[i - 1]; auto end_idx = cu_seqlens_lock[i]; wg_count += static_cast((end_idx - start_idx + wg_seq_len - 1) / wg_seq_len); } } else { - wg_count = cu_seqlens_lock.size() - 1; + wg_count = cu_seqlens_sz - 1; } auto& wgs = kd.params.workGroups; @@ -180,12 +185,22 @@ class VLSDPAGenerator : public KernelGenerator { // Contract with the CM kernel (see cm_sdpa_vlen.cm): only the token (outer) // axis of each input layout may be padded — the head and head_size strides // are hardcoded in the kernel and are NOT parameterized here. - int32_t token_offset_q = static_cast(params.input_layouts[0].get_linear_offset()); - int32_t token_offset_k = static_cast(params.input_layouts[1].get_linear_offset()); - int32_t token_offset_v = static_cast(params.input_layouts[2].get_linear_offset()); - const auto q_token_pitch = static_cast(params.input_layouts[0].get_pitches()[0]); - const auto k_token_pitch = static_cast(params.input_layouts[1].get_pitches()[0]); - const auto v_token_pitch = static_cast(params.input_layouts[2].get_pitches()[0]); + constexpr size_t i32_max = 0x7fffffff; + const auto token_offset_q_sz = params.input_layouts[0].get_linear_offset(); + const auto token_offset_k_sz = params.input_layouts[1].get_linear_offset(); + const auto token_offset_v_sz = params.input_layouts[2].get_linear_offset(); + const auto q_token_pitch_sz = params.input_layouts[0].get_pitches()[0]; + const auto k_token_pitch_sz = params.input_layouts[1].get_pitches()[0]; + const auto v_token_pitch_sz = params.input_layouts[2].get_pitches()[0]; + OPENVINO_ASSERT(token_offset_q_sz <= i32_max && token_offset_k_sz <= i32_max && token_offset_v_sz <= i32_max && + q_token_pitch_sz <= i32_max && k_token_pitch_sz <= i32_max && v_token_pitch_sz <= i32_max, + "VLSDPA token offsets/pitches must fit into int32"); + const int32_t token_offset_q = static_cast(token_offset_q_sz); + const int32_t token_offset_k = static_cast(token_offset_k_sz); + const int32_t token_offset_v = static_cast(token_offset_v_sz); + const int32_t q_token_pitch = static_cast(q_token_pitch_sz); + const int32_t k_token_pitch = static_cast(k_token_pitch_sz); + const int32_t v_token_pitch = static_cast(v_token_pitch_sz); std::vector scalars{need_wg_mapping, token_offset_q, From 6cf27e94af5d46eda57a856afafe8bf744c338fc Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 14 Jul 2026 15:17:59 +0800 Subject: [PATCH 12/16] fix --- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 30 ++-- .../tests/unit/test_cases/vlsdpa_gpu_test.cpp | 156 ++++++++++++++++++ 2 files changed, 168 insertions(+), 18 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index ee60f9b9602502..fc58129edbc86b 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -139,7 +139,7 @@ class VLSDPAGenerator : public KernelGenerator { mem_lock cu_seqlens_lock(cu_seqlens_mem, *params.strm); const auto cu_seqlens_sz = cu_seqlens_lock.size(); - OPENVINO_ASSERT(cu_seqlens_sz >= 2 && cu_seqlens_sz % 2 == 0, + OPENVINO_ASSERT(cu_seqlens_sz >= 2, "VLSDPA expects cu_seqlens to contain at least two elements (start/end) and be even-sized"); size_t max_seq_len = 0; for (size_t i = 1; i < cu_seqlens_sz; i++) { @@ -152,7 +152,7 @@ class VLSDPAGenerator : public KernelGenerator { const size_t CM_GRF_WIDTH = (info.arch <= gpu_arch::xe_hpc) ? 256 : 512; const size_t q_step = static_cast(std::floor(CM_GRF_WIDTH / 32)); size_t wg_size = static_cast(std::floor((max_seq_len + q_step - 1) / q_step)); - int32_t need_wg_mapping = 0; + size_t need_wg_mapping = 0; if (wg_size > 16) { // seq_len is too large for a single work-group; use wg_size=16 and let // the kernel's while-loop scan cu_seqlens to find its sequence/block. @@ -186,23 +186,17 @@ class VLSDPAGenerator : public KernelGenerator { // axis of each input layout may be padded — the head and head_size strides // are hardcoded in the kernel and are NOT parameterized here. constexpr size_t i32_max = 0x7fffffff; - const auto token_offset_q_sz = params.input_layouts[0].get_linear_offset(); - const auto token_offset_k_sz = params.input_layouts[1].get_linear_offset(); - const auto token_offset_v_sz = params.input_layouts[2].get_linear_offset(); - const auto q_token_pitch_sz = params.input_layouts[0].get_pitches()[0]; - const auto k_token_pitch_sz = params.input_layouts[1].get_pitches()[0]; - const auto v_token_pitch_sz = params.input_layouts[2].get_pitches()[0]; - OPENVINO_ASSERT(token_offset_q_sz <= i32_max && token_offset_k_sz <= i32_max && token_offset_v_sz <= i32_max && - q_token_pitch_sz <= i32_max && k_token_pitch_sz <= i32_max && v_token_pitch_sz <= i32_max, + const auto token_offset_q = params.input_layouts[0].get_linear_offset(); + const auto token_offset_k = params.input_layouts[1].get_linear_offset(); + const auto token_offset_v = params.input_layouts[2].get_linear_offset(); + const size_t q_token_pitch = static_cast(params.input_layouts[0].get_pitches()[0]); + const size_t k_token_pitch = static_cast(params.input_layouts[1].get_pitches()[0]); + const size_t v_token_pitch = static_cast(params.input_layouts[2].get_pitches()[0]); + OPENVINO_ASSERT(token_offset_q <= i32_max && token_offset_k <= i32_max && token_offset_v <= i32_max && + q_token_pitch <= i32_max && k_token_pitch <= i32_max && v_token_pitch <= i32_max, "VLSDPA token offsets/pitches must fit into int32"); - const int32_t token_offset_q = static_cast(token_offset_q_sz); - const int32_t token_offset_k = static_cast(token_offset_k_sz); - const int32_t token_offset_v = static_cast(token_offset_v_sz); - const int32_t q_token_pitch = static_cast(q_token_pitch_sz); - const int32_t k_token_pitch = static_cast(k_token_pitch_sz); - const int32_t v_token_pitch = static_cast(v_token_pitch_sz); - - std::vector scalars{need_wg_mapping, + + std::vector scalars{need_wg_mapping, token_offset_q, token_offset_k, token_offset_v, diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp index c07786b1ab53f1..e45ca319ca8f66 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include #include "scaled_dot_product_attention_inst.h" #include "vl_sdpa_inst.h" @@ -101,6 +103,69 @@ struct vlsdpa_gpu_test : public ::testing::TestWithParam { return topo; } + // Build a Split(axis=1, num_splits=3) crop chain fanning a fused [L, 3*H, S] + // parent into named "q"/"k"/"v" slices. Enabling optimize_data on the resulting + // topology causes prepare_buffer_fusing to keep the crops in-place, so each + // slice's runtime layout carries a non-zero linear_offset and the parent's row + // pitch (3 * num_heads * head_size). This is exactly the discontinuous-token-axis + // input shape that the CM vl_sdpa kernel's token_offset_{q,k,v} / {q,k,v}_token_pitch + // scalars are designed to handle (see cm_sdpa_vlen.cm memory-layout contract). + void add_qkv_split(topology& topo, + const primitive_id& fused_id, + const primitive_id& axis_data_id, + int32_t num_heads, + int32_t head_size) { + const auto slice_shape = tensor(batch(1), feature(num_heads), spatial(head_size, 1)); + topo.add(crop("q", { input_info(fused_id), input_info(axis_data_id) }, + slice_shape, tensor(feature(0 * num_heads)), + crop_ngraph_op_mode::split, 0, /*axis*/1, /*num_splits*/3)); + topo.add(crop("k", { input_info(fused_id), input_info(axis_data_id) }, + slice_shape, tensor(feature(1 * num_heads)), + crop_ngraph_op_mode::split, 1, /*axis*/1, /*num_splits*/3)); + topo.add(crop("v", { input_info(fused_id), input_info(axis_data_id) }, + slice_shape, tensor(feature(2 * num_heads)), + crop_ngraph_op_mode::split, 2, /*axis*/1, /*num_splits*/3)); + } + + topology create_ref_topology_fused(layout fused_layout, + layout attn_mask_layout, + memory::ptr axis_data_mem, + int32_t num_heads, + int32_t head_size) { + topology topo; + topo.add(input_layout("fused", fused_layout)); + topo.add(data("split_axis_ref", axis_data_mem)); + add_qkv_split(topo, "fused", "split_axis_ref", num_heads, head_size); + + topo.add(input_layout("attn", attn_mask_layout)); // "attention_mask" + + auto sdpa_prim = scaled_dot_product_attention("sdpa", {input_info("q"), input_info("k"), input_info("v"), input_info("attn")}, + false, -1, order_q, order_k, order_v, {0, 1, 2}, {}, false); + topo.add(sdpa_prim); + topo.add(permute("permute_o", input_info("sdpa"), order_o2)); + topo.add(reorder("result", input_info("permute_o"), format::bfyx, data_types::f16)); + return topo; + } + + topology create_topology_fused(layout fused_layout, + layout cu_seqlen_layout, + memory::ptr axis_data_mem, + int32_t num_heads, + int32_t head_size) { + topology topo; + topo.add(input_layout("fused", fused_layout)); + topo.add(data("split_axis_opt", axis_data_mem)); + add_qkv_split(topo, "fused", "split_axis_opt", num_heads, head_size); + + topo.add(input_layout("attn", cu_seqlen_layout)); // "cu_seqlens" + + auto vlsdpa_prim = vl_sdpa("vlsdpa", {input_info("q"), input_info("k"), input_info("v"), input_info("attn")}, + order_q, order_k, order_v, order_o); + topo.add(vlsdpa_prim); + topo.add(reorder("result", input_info("vlsdpa"), format::bfyx, data_types::f16)); + return topo; + } + std::tuple run_network(topology &topo, cldnn::memory::ptr q_mem, cldnn::memory::ptr k_mem, @@ -185,6 +250,78 @@ struct vlsdpa_gpu_test : public ::testing::TestWithParam { } } + void execute_fused_qkv(vlsdpa_test_params& p, const bool is_caching_test = false) { + const auto head_size = p.head_size; + const auto num_heads = p.num_heads; + const auto cu_seqlens = p.cu_seqlens; + const auto cumsum = cu_seqlens.back(); + const auto seq_length_q = cumsum; + const auto seq_length_kv = cumsum; + + assert(cu_seqlens.front() == 0); + + auto& engine = get_test_engine(); + + // Parent-layout dims: {L, 3*H, S}. After optimize_data-driven buffer fusing, + // the three Split(axis=1, num_splits=3) crops become in-place views of this + // parent, producing the discontinuous token-axis layout the CM kernel expects. + const int32_t fused_feature = 3 * num_heads; + cldnn::layout fused_dyn_layout({-1, fused_feature, head_size}, data_types::f16, format::bfyx); + cldnn::layout attn_mask_dyn_layout({1, -1, -1}, data_types::f16, format::bfyx); + cldnn::layout cu_seqlens_dyn_layout({-1}, data_types::i32, format::bfyx); + + // Split-axis constant shared by both topologies (Split expects axis as i64 scalar). + auto axis_data_mem = engine.allocate_memory({{}, data_types::i64, format::bfyx}); + set_values(axis_data_mem, {1}); + + topology ref_topo = create_ref_topology_fused(fused_dyn_layout, attn_mask_dyn_layout, + axis_data_mem, num_heads, head_size); + topology opt_topo = create_topology_fused(fused_dyn_layout, cu_seqlens_dyn_layout, + axis_data_mem, num_heads, head_size); + + // Concrete input allocations. The fused parent is filled once and reused by both + // networks, so any output divergence isolates the vl_sdpa kernel's handling of + // the non-zero token_offset_{q,k,v} / {q,k,v}_token_pitch scalar path. + cldnn::layout fused_static_layout({seq_length_q, fused_feature, head_size}, data_types::f16, format::bfyx); + cldnn::layout attn_mask_static_layout({1, seq_length_q, seq_length_kv}, data_types::f16, format::bfyx); + cldnn::layout cu_seqlens_static_layout({static_cast(cu_seqlens.size())}, data_types::i32, format::bfyx); + + auto fused_mem = engine.allocate_memory(fused_static_layout); + auto attn_mask_mem = engine.allocate_memory(attn_mask_static_layout); + auto cu_seqlens_mem = engine.allocate_memory(cu_seqlens_static_layout); + + load_input(fused_mem, 0); + get_attention_mask(attn_mask_mem, cu_seqlens); + get_cu_seqlens(cu_seqlens_mem, cu_seqlens); + + auto run = [&](topology& topo, cldnn::memory::ptr attn_mem, bool optimize) { + ExecutionConfig config = get_test_default_config(engine); + config.set_property(ov::intel_gpu::allow_new_shape_infer(true)); + // optimize=false keeps the reference crops materialized (contiguous Q/K/V). + // optimize=true lets prepare_buffer_fusing fold the vl_sdpa crops in-place, + // which is the whole point of this test. + config.set_property(ov::intel_gpu::optimize_data(optimize)); + + cldnn::network::ptr net = get_network(engine, topo, config, get_test_stream_ptr(), is_caching_test); + net->set_input_data("fused", fused_mem); + net->set_input_data("attn", attn_mem); + + auto outputs = net->execute(); + return outputs.at("result").get_memory(); + }; + + auto mem_ref_ptr = run(ref_topo, attn_mask_mem, /*optimize=*/false); + auto mem_opt_ptr = run(opt_topo, cu_seqlens_mem, /*optimize=*/true); + + cldnn::mem_lock ref_data(mem_ref_ptr, get_test_stream()); + cldnn::mem_lock opt_data(mem_opt_ptr, get_test_stream()); + for (size_t idx = 0; idx < ref_data.size(); idx++) { + ASSERT_FALSE(std::isnan(opt_data[idx]) || std::isnan(ref_data[idx])) << "NaN found at index " << idx; + } + auto ret = cosineSimilarity(ref_data, opt_data); + ASSERT_GE(ret, 0.95f); + } + static std::string PrintToStringParamName(const testing::TestParamInfo& info) { std::string result = "vlsdpa_gpu_test_" + std::to_string(info.param.head_size) + "_" + @@ -227,6 +364,25 @@ TEST_P(vlsdpa_gpu_test, basic_caching) { execute(p, true); } +// Exercises the token_offset_{q,k,v} / {q,k,v}_token_pitch runtime scalars that +// the CM vl_sdpa kernel gained to handle discontinuous Q/K/V slices of a fused +// [L, 3*H, S] parent buffer (Split(axis=1, num_splits=3) with in-place buffer fusing). +TEST_P(vlsdpa_gpu_test, discontinuous_qkv_fused) { + if (!check_vlsdpa_available()) + GTEST_SKIP(); + + auto p = GetParam(); + execute_fused_qkv(p); +} + +TEST_P(vlsdpa_gpu_test, discontinuous_qkv_fused_caching) { + if (!check_vlsdpa_available()) + GTEST_SKIP(); + + auto p = GetParam(); + execute_fused_qkv(p, true); +} + INSTANTIATE_TEST_SUITE_P(smoke_vlsdpa_gpu_test, vlsdpa_gpu_test, ::testing::Values( From 8ca994a11773d74674e175b29a4e6d6ad5b62f05 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 14 Jul 2026 16:35:26 +0800 Subject: [PATCH 13/16] [GPU] Fix VLSDPA CM kernel GPU hang when num_head==1 Bug --- Running on GPU hangs and dies with: clWaitForEvents, error code: -14 CL_EXEC_STATUS_ERROR_FOR_EVENTS_IN_WAIT_LIST clFinish, error code: -5 CL_OUT_OF_RESOURCES All 9 test cases with are affected; passes. here means the CM kernel performs out-of-bounds SVM addressing. Root cause ---------- The kernel argument dump for the first failing case (num_head=1, head_size=16, cu_seqlens={0,16}, Q/K/V buffers = 512 B each) shows: scalar 9 (q_token_pitch) = 256 <-- expected 16 scalar 10 (k_token_pitch) = 256 <-- expected 16 scalar 11 (v_token_pitch) = 256 <-- expected 16 computes in half-elements, so with (instead of 16) it walks 16x past the end of the 256-element Q buffer as soon as . The wrong value comes from , which unconditionally reads the per-token stride from (the pitch of layout dim 0, ). The kernel contract expects the memory to be with the token axis on dim 0, which is only true when the surrounding ops are fused into (giving non-empty ). In this pipeline the fusion did not run: 1. (core, early) creates with **empty** orders and leaves the four surrounding ops in place. 2. -> rewrites those Transposes to because the permuted dim has size 1. 3. (plugin, late) requires the exact pattern , no longer matches, and does not fuse -> stays empty. VLSDPA's input layout is then the post-Reshape view = . The token axis is , so the correct per-token stride is , but the impl reads . Silent, layout-dependent OOB -> GPU hang. Fix --- 1) Fuse the surrounding at VLSDPA-creation time in , so the fusion cannot be undone by any later pass (in particular ). When all four adjacent transposes are single-consumer , peel them and build VLSDPA(q, k, v, cu_seqlens, {1,0,2}, {1,0,2}, {1,0,2}, {1,0,2}) then . When any of the four is missing or has a different order, fall back to the previous behavior (empty orders, no fusion). 2) Defense in depth: add in the CM VLSDPA impl, invoked from (main thread, before async kernel compilation), , and the dispatch data-func lambda. Any future path that reintroduces an unfused now fails fast at with a message naming the token-stride contract, instead of a silent GPU hang at inference time. Testing ------- - : 18/18 pass (previously 9x cases aborted with ). - / updated to assert the fused reference model (VLSDPA carries orders; no surrounding transposes). - Verified the assert path by temporarily disabling (1): the failing case surfaces as an at model-compile time with the intended message, no GPU hang. Files ----- - src/core/src/pass/sdpa_to_vlsdpa.cpp - src/common/transformations/tests/common_optimizations/ sdpa_to_vlsdpa_test.cpp - src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp --- .../sdpa_to_vlsdpa_test.cpp | 21 ++--- src/core/src/pass/sdpa_to_vlsdpa.cpp | 76 +++++++++++++++++-- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 27 +++++++ 3 files changed, 106 insertions(+), 18 deletions(-) diff --git a/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp b/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp index 4732009cde454e..078e06fd8f44aa 100644 --- a/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp +++ b/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp @@ -58,24 +58,19 @@ std::shared_ptr build_target_model(const string& mask_name) { k->set_friendly_name("k"); v->set_friendly_name("v"); - auto transpose_q = std::make_shared(q, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - auto transpose_k = std::make_shared(k, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - auto transpose_v = std::make_shared(v, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - transpose_q->set_friendly_name("transpose_q"); - transpose_k->set_friendly_name("transpose_k"); - transpose_v->set_friendly_name("transpose_v"); - auto cuseq_mask = std::make_shared(element::i32, PartialShape{-1}); cuseq_mask->set_friendly_name(mask_name); cuseq_mask->get_output_tensor(0).set_names({mask_name}); - auto vlsdpa = - std::make_shared(OutputVector{transpose_q, transpose_k, transpose_v, cuseq_mask}); - - auto transpose_o = std::make_shared(vlsdpa, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - transpose_o->set_friendly_name("transpose_o"); + const std::vector order{1, 0, 2}; + auto vlsdpa = std::make_shared(OutputVector{q, k, v, cuseq_mask}, + order, + order, + order, + order); + vlsdpa->set_friendly_name("transpose_o"); - return std::make_shared(OutputVector{transpose_o}, ParameterVector{q, k, v, cuseq_mask}); + return std::make_shared(OutputVector{vlsdpa}, ParameterVector{q, k, v, cuseq_mask}); } }; // namespace diff --git a/src/core/src/pass/sdpa_to_vlsdpa.cpp b/src/core/src/pass/sdpa_to_vlsdpa.cpp index 9988ed4bd2ff1c..214eeeecfb101b 100644 --- a/src/core/src/pass/sdpa_to_vlsdpa.cpp +++ b/src/core/src/pass/sdpa_to_vlsdpa.cpp @@ -9,7 +9,11 @@ #include #include "openvino/cc/pass/itt.hpp" +#include "openvino/core/graph_util.hpp" +#include "openvino/core/rt_info.hpp" +#include "openvino/op/constant.hpp" #include "openvino/op/scaled_dot_product_attention.hpp" +#include "openvino/op/transpose.hpp" #include "openvino/op/util/node_util.hpp" #include "openvino/pass/manager.hpp" #include "ov_ops/vl_sdpa.hpp" @@ -85,14 +89,76 @@ bool SDPAToVLSDPA::run_on_model(const std::shared_ptr& model) { const auto sdpa_consumers = sdpa->get_output_target_inputs(0); const auto new_args = sdpa->input_values(); - OutputVector inputs{new_args.at(0), new_args.at(1), new_args.at(2), cu_seqlens_param}; + + // Try to fuse the {1,0,2} Transposes surrounding SDPA into the VLSDPA op + // at op-creation time. Doing it here (rather than as a later plugin-side + // matcher on `Transpose->VLSDPA->Transpose`) is race-free w.r.t. common + // optimizations that may convert Transpose->Reshape when one of the + // permuted dims degenerates to size 1 (e.g. num_head==1). Constraints + // mirror `VLSDPA::validate_and_infer_types`: all four orders must be + // equal and, when non-empty, equal to {1,0,2}. + const std::vector expected_order{1, 0, 2}; + auto try_peel_input_transpose = + [&expected_order](ov::Output& value) -> std::shared_ptr { + auto transpose = ov::as_type_ptr(value.get_node_shared_ptr()); + if (!transpose) + return nullptr; + if (transpose->get_output_target_inputs(0).size() != 1) + return nullptr; + auto order_const = + ov::as_type_ptr(transpose->input_value(1).get_node_shared_ptr()); + if (!order_const || ov::shape_size(order_const->get_shape()) != expected_order.size()) + return nullptr; + if (order_const->cast_vector() != expected_order) + return nullptr; + value = transpose->input_value(0); + return transpose; + }; + + ov::Output q_in = new_args.at(0); + ov::Output k_in = new_args.at(1); + ov::Output v_in = new_args.at(2); + auto tq = try_peel_input_transpose(q_in); + auto tk = try_peel_input_transpose(k_in); + auto tv = try_peel_input_transpose(v_in); + + std::shared_ptr to; + if (sdpa_consumers.size() == 1) { + auto candidate_out = ov::as_type_ptr( + sdpa_consumers.begin()->get_node()->shared_from_this()); + if (candidate_out) { + auto out_order_const = ov::as_type_ptr( + candidate_out->input_value(1).get_node_shared_ptr()); + if (out_order_const && + ov::shape_size(out_order_const->get_shape()) == expected_order.size() && + out_order_const->cast_vector() == expected_order) { + to = candidate_out; + } + } + } + + const bool fuse_transposes = tq && tk && tv && to; std::shared_ptr vl_sdpa; - vl_sdpa = std::make_shared(inputs); - vl_sdpa->set_friendly_name(sdpa->get_friendly_name()); + if (fuse_transposes) { + OutputVector inputs{q_in, k_in, v_in, cu_seqlens_param}; + vl_sdpa = std::make_shared(inputs, + expected_order, + expected_order, + expected_order, + expected_order); + vl_sdpa->set_friendly_name(to->get_friendly_name()); + ov::copy_runtime_info({tq, tk, tv, sdpa->shared_from_this(), to}, vl_sdpa); + ov::replace_node(to, vl_sdpa); + } else { + OutputVector inputs{new_args.at(0), new_args.at(1), new_args.at(2), cu_seqlens_param}; + vl_sdpa = std::make_shared(inputs); + vl_sdpa->set_friendly_name(sdpa->get_friendly_name()); + ov::copy_runtime_info(sdpa->shared_from_this(), vl_sdpa); - for (auto& consumer : sdpa_consumers) - consumer.replace_source_output(vl_sdpa); + for (auto& consumer : sdpa_consumers) + consumer.replace_source_output(vl_sdpa); + } } } } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index fc58129edbc86b..9204f6ffb74a6c 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -35,6 +35,27 @@ constexpr size_t get_default_kv_blk(size_t head_size) { return 1; } +// Enforce the CM kernel memory-layout contract documented in cm_sdpa_vlen.cm: +// Q/K/V/output layouts must all carry the fused transpose order {1,0,2}, i.e. the +// original SDPA-facing view [H, L, S] is expressed as physical memory [L, H, S] with +// the token (outer) axis at dim 0. When this contract is violated (empty or non-{1,0,2} +// orders, typically because TransposeVLSDPA fusion did not fire — e.g. surrounding +// Transposes were converted to Reshape by TransposeToReshape when a permuted dim +// degenerates to size 1), `get_pitches()[0]` no longer represents the per-token stride +// and the kernel silently walks out of bounds -> GPU hang (CL_OUT_OF_RESOURCES). +// Fail fast at compile / dispatch time with an actionable message. +inline void ensure_fused_transpose_order(const vl_sdpa& desc) { + const std::vector expected_order{1, 0, 2}; + OPENVINO_ASSERT(desc.input_q_transpose_order == expected_order && + desc.input_k_transpose_order == expected_order && + desc.input_v_transpose_order == expected_order && + desc.output_transpose_order == expected_order, + "VLSDPA CM impl requires the surrounding Transpose{1,0,2} to be fused into " + "VLSDPA (all four transpose orders must be {1,0,2}). If this fires, the " + "SDPAToVLSDPA / TransposeVLSDPAMatcher fusion did not run for this pattern; " + "kernel arg computation (`get_pitches()[0]` as the token stride) is unsafe."); +} + class VLSDPAGenerator : public KernelGenerator { public: VLSDPAGenerator() : KernelGenerator("cm_sdpa_vlen") {} @@ -60,6 +81,7 @@ class VLSDPAGenerator : public KernelGenerator { }; auto desc = params.typed_desc(); + ensure_fused_transpose_order(*desc); const auto query_shape = transpose_pshape(params.get_input_layout(0).get_partial_shape(), desc->input_q_transpose_order); const auto key_shape = transpose_pshape(params.get_input_layout(1).get_partial_shape(), desc->input_k_transpose_order); @@ -117,6 +139,7 @@ class VLSDPAGenerator : public KernelGenerator { return DispatchDataFunc{[](const RuntimeParams& params, KernelData& kd, ImplRuntimeParams*) { assert(!params.is_dynamic()); auto desc = params.typed_desc(); + ensure_fused_transpose_order(*desc); // transpose shape into BHLS(4D), or HLS(3D) auto transpose_pshape = [](const ov::Shape& pshape, const std::vector& order) { @@ -234,6 +257,10 @@ class VLSDPAOptImpl : public PrimitiveImplCM { std::unique_ptr VLSDPAOptImplementationManager::create_impl(const program_node& node, const RuntimeParams& params) const { assert(node.is_type()); + // Validate the fused-transpose invariant on the main thread, before triggering async + // kernel compilation — the same check inside get_jit_constants would be swallowed by + // the kernels cache and surface only as "Kernel for {vlsdpa:sdpa} is not found". + ensure_fused_transpose_order(*params.typed_desc()); return std::make_unique(node, params); } From 48df96b61a099144ef61b4765976c71bbb2a4ced Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Tue, 14 Jul 2026 16:43:22 +0800 Subject: [PATCH 14/16] clang format --- .../sdpa_to_vlsdpa_test.cpp | 7 +-- src/core/src/pass/sdpa_to_vlsdpa.cpp | 14 +++--- .../src/graph/impls/cm/vl_sdpa_opt.cpp | 46 ++++++++----------- .../impls/ocl_v2/sdpa/paged_attention_opt.cpp | 4 +- 4 files changed, 28 insertions(+), 43 deletions(-) diff --git a/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp b/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp index 078e06fd8f44aa..f3ef31dff99ec3 100644 --- a/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp +++ b/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp @@ -63,11 +63,8 @@ std::shared_ptr build_target_model(const string& mask_name) { cuseq_mask->get_output_tensor(0).set_names({mask_name}); const std::vector order{1, 0, 2}; - auto vlsdpa = std::make_shared(OutputVector{q, k, v, cuseq_mask}, - order, - order, - order, - order); + auto vlsdpa = + std::make_shared(OutputVector{q, k, v, cuseq_mask}, order, order, order, order); vlsdpa->set_friendly_name("transpose_o"); return std::make_shared(OutputVector{vlsdpa}, ParameterVector{q, k, v, cuseq_mask}); diff --git a/src/core/src/pass/sdpa_to_vlsdpa.cpp b/src/core/src/pass/sdpa_to_vlsdpa.cpp index 214eeeecfb101b..8f5f90a205dc0d 100644 --- a/src/core/src/pass/sdpa_to_vlsdpa.cpp +++ b/src/core/src/pass/sdpa_to_vlsdpa.cpp @@ -105,8 +105,7 @@ bool SDPAToVLSDPA::run_on_model(const std::shared_ptr& model) { return nullptr; if (transpose->get_output_target_inputs(0).size() != 1) return nullptr; - auto order_const = - ov::as_type_ptr(transpose->input_value(1).get_node_shared_ptr()); + auto order_const = ov::as_type_ptr(transpose->input_value(1).get_node_shared_ptr()); if (!order_const || ov::shape_size(order_const->get_shape()) != expected_order.size()) return nullptr; if (order_const->cast_vector() != expected_order) @@ -124,13 +123,12 @@ bool SDPAToVLSDPA::run_on_model(const std::shared_ptr& model) { std::shared_ptr to; if (sdpa_consumers.size() == 1) { - auto candidate_out = ov::as_type_ptr( - sdpa_consumers.begin()->get_node()->shared_from_this()); + auto candidate_out = + ov::as_type_ptr(sdpa_consumers.begin()->get_node()->shared_from_this()); if (candidate_out) { - auto out_order_const = ov::as_type_ptr( - candidate_out->input_value(1).get_node_shared_ptr()); - if (out_order_const && - ov::shape_size(out_order_const->get_shape()) == expected_order.size() && + auto out_order_const = + ov::as_type_ptr(candidate_out->input_value(1).get_node_shared_ptr()); + if (out_order_const && ov::shape_size(out_order_const->get_shape()) == expected_order.size() && out_order_const->cast_vector() == expected_order) { to = candidate_out; } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index 9204f6ffb74a6c..f375368d2f3a9a 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -46,10 +46,8 @@ constexpr size_t get_default_kv_blk(size_t head_size) { // Fail fast at compile / dispatch time with an actionable message. inline void ensure_fused_transpose_order(const vl_sdpa& desc) { const std::vector expected_order{1, 0, 2}; - OPENVINO_ASSERT(desc.input_q_transpose_order == expected_order && - desc.input_k_transpose_order == expected_order && - desc.input_v_transpose_order == expected_order && - desc.output_transpose_order == expected_order, + OPENVINO_ASSERT(desc.input_q_transpose_order == expected_order && desc.input_k_transpose_order == expected_order && + desc.input_v_transpose_order == expected_order && desc.output_transpose_order == expected_order, "VLSDPA CM impl requires the surrounding Transpose{1,0,2} to be fused into " "VLSDPA (all four transpose orders must be {1,0,2}). If this fires, the " "SDPAToVLSDPA / TransposeVLSDPAMatcher fusion did not run for this pattern; " @@ -157,17 +155,15 @@ class VLSDPAGenerator : public KernelGenerator { // Read cu_seqlens via mem_lock — zero-cost on USM-host (PC), no GPU sync stall. const auto cu_seqlens_mem = params.memory_deps.at(params.input_layouts.size() - 1); - OPENVINO_ASSERT(cu_seqlens_mem->get_layout().data_type == ov::element::i32, - "VLSDPA expects cu_seqlens input to be i32"); + OPENVINO_ASSERT(cu_seqlens_mem->get_layout().data_type == ov::element::i32, "VLSDPA expects cu_seqlens input to be i32"); mem_lock cu_seqlens_lock(cu_seqlens_mem, *params.strm); const auto cu_seqlens_sz = cu_seqlens_lock.size(); - OPENVINO_ASSERT(cu_seqlens_sz >= 2, - "VLSDPA expects cu_seqlens to contain at least two elements (start/end) and be even-sized"); + OPENVINO_ASSERT(cu_seqlens_sz >= 2, "VLSDPA expects cu_seqlens to contain at least two elements (start/end) and be even-sized"); size_t max_seq_len = 0; for (size_t i = 1; i < cu_seqlens_sz; i++) { auto start_idx = cu_seqlens_lock[i - 1]; - auto end_idx = cu_seqlens_lock[i]; + auto end_idx = cu_seqlens_lock[i]; max_seq_len = std::max(max_seq_len, static_cast(end_idx - start_idx)); } @@ -189,7 +185,7 @@ class VLSDPAGenerator : public KernelGenerator { const auto wg_seq_len = wg_size * q_step; for (size_t i = 1; i < cu_seqlens_sz; i++) { auto start_idx = cu_seqlens_lock[i - 1]; - auto end_idx = cu_seqlens_lock[i]; + auto end_idx = cu_seqlens_lock[i]; wg_count += static_cast((end_idx - start_idx + wg_seq_len - 1) / wg_seq_len); } } else { @@ -208,24 +204,18 @@ class VLSDPAGenerator : public KernelGenerator { // Contract with the CM kernel (see cm_sdpa_vlen.cm): only the token (outer) // axis of each input layout may be padded — the head and head_size strides // are hardcoded in the kernel and are NOT parameterized here. - constexpr size_t i32_max = 0x7fffffff; - const auto token_offset_q = params.input_layouts[0].get_linear_offset(); - const auto token_offset_k = params.input_layouts[1].get_linear_offset(); - const auto token_offset_v = params.input_layouts[2].get_linear_offset(); - const size_t q_token_pitch = static_cast(params.input_layouts[0].get_pitches()[0]); - const size_t k_token_pitch = static_cast(params.input_layouts[1].get_pitches()[0]); - const size_t v_token_pitch = static_cast(params.input_layouts[2].get_pitches()[0]); - OPENVINO_ASSERT(token_offset_q <= i32_max && token_offset_k <= i32_max && token_offset_v <= i32_max && - q_token_pitch <= i32_max && k_token_pitch <= i32_max && v_token_pitch <= i32_max, - "VLSDPA token offsets/pitches must fit into int32"); - - std::vector scalars{need_wg_mapping, - token_offset_q, - token_offset_k, - token_offset_v, - q_token_pitch, - k_token_pitch, - v_token_pitch}; + constexpr size_t i32_max = 0x7fffffff; + const auto token_offset_q = params.input_layouts[0].get_linear_offset(); + const auto token_offset_k = params.input_layouts[1].get_linear_offset(); + const auto token_offset_v = params.input_layouts[2].get_linear_offset(); + const size_t q_token_pitch = static_cast(params.input_layouts[0].get_pitches()[0]); + const size_t k_token_pitch = static_cast(params.input_layouts[1].get_pitches()[0]); + const size_t v_token_pitch = static_cast(params.input_layouts[2].get_pitches()[0]); + OPENVINO_ASSERT(token_offset_q <= i32_max && token_offset_k <= i32_max && token_offset_v <= i32_max && q_token_pitch <= i32_max && + k_token_pitch <= i32_max && v_token_pitch <= i32_max, + "VLSDPA token offsets/pitches must fit into int32"); + + std::vector scalars{need_wg_mapping, token_offset_q, token_offset_k, token_offset_v, q_token_pitch, k_token_pitch, v_token_pitch}; kd.params.scalars.clear(); for (auto i : scalars) { scalar_desc s; diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp index 5e41219aabe04f..0915210090772f 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp @@ -1799,8 +1799,8 @@ class PagedAttentionOptImpl : public SDPAImplBase { total_matrix_elements += evictable_size * evictable_size; total_vector_elements += evictable_size; } - GPU_DEBUG_TRACE_DETAIL << "Adaptive RKV: Allocating dynamic buffers - " - << "matrix: " << total_matrix_elements << ", vector: " << total_vector_elements << std::endl; + GPU_DEBUG_TRACE_DETAIL << "Adaptive RKV: Allocating dynamic buffers - " << "matrix: " << total_matrix_elements + << ", vector: " << total_vector_elements << std::endl; } else { // Fallback: use maximum size (512) if runtime values not available const size_t max_evictable_size = 512; From 6605746af2dafc00c1faa8c2cbf1570b49690c21 Mon Sep 17 00:00:00 2001 From: cecilia peng Date: Wed, 15 Jul 2026 11:08:08 +0800 Subject: [PATCH 15/16] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index f375368d2f3a9a..1b565836714b16 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -159,7 +159,7 @@ class VLSDPAGenerator : public KernelGenerator { mem_lock cu_seqlens_lock(cu_seqlens_mem, *params.strm); const auto cu_seqlens_sz = cu_seqlens_lock.size(); - OPENVINO_ASSERT(cu_seqlens_sz >= 2, "VLSDPA expects cu_seqlens to contain at least two elements (start/end) and be even-sized"); + OPENVINO_ASSERT(cu_seqlens_sz >= 2, "VLSDPA expects cu_seqlens to contain at least two elements (start/end)"); size_t max_seq_len = 0; for (size_t i = 1; i < cu_seqlens_sz; i++) { auto start_idx = cu_seqlens_lock[i - 1]; From 66149a6fde1291cd92b22836423d746fbc15fea3 Mon Sep 17 00:00:00 2001 From: ceciliapeng2011 Date: Wed, 15 Jul 2026 11:30:34 +0800 Subject: [PATCH 16/16] license header --- .../graph/impls/cm/include/cm_sdpa_common.hpp | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index 02a13f033d4555..78fa3f98aa039b 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -1,18 +1,7 @@ -/******************************************************************************* - * Copyright (c) 2018-2026 Intel Corporation - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - *******************************************************************************/ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + #include "cm_attention_common.hpp" // Number of kv sub-tiles (kv_step rows each) processed per online-softmax update in