Skip to content

[GPU] SDPA: use key-partitioned single-token kernels for up to 4 query rows #38996

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
osabnis wants to merge 14 commits into
base: master
Choose a base branch
from
Open

[GPU] SDPA: use key-partitioned single-token kernels for up to 4 query rows #38996

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
d974607
[GPU] SDPA: use key-partitioned single-token kernels for up to 4 quer…
osabnis Oct 19, 2026
30e2bf6
Merge branch 'master' into gpu-sdpa-few-rows-single-token
osabnis Oct 19, 2026
8d7b5d1
[GPU] SDPA: optimize the code and improve comments
osabnis Oct 20, 2026
063433c
Merge branch 'master' into gpu-sdpa-few-rows-single-token
osabnis Oct 22, 2026
8b2ad1c
Merge branch 'master' into gpu-sdpa-few-rows-single-token
osabnis Oct 23, 2026
228ee98
Merge branch 'master' into gpu-sdpa-few-rows-single-token
osabnis Nov 2, 2026
8d8db38
Merge branch 'master' into gpu-sdpa-few-rows-single-token
osabnis Nov 4, 2026
3db4909
[GPU] SDPA: keep non-exact few-row cases off the single-token kernel
osabnis Nov 4, 2026
49d071b
[GPU] SDPA: handle key-broadcast attention masks without out-of-bound…
osabnis Nov 4, 2026
175b7b0
Merge branch 'master' into gpu-sdpa-few-rows-single-token
osabnis Nov 5, 2026
5e0bf5a
[GPU] SDPA: clamp the attention mask for every mask type
osabnis Nov 9, 2026
055e02a
[GPU] SDPA: keep V wider than K on the multi-token kernels
osabnis Nov 9, 2026
5cfec25
[GPU] SDPA: honor output padding in the finalization stage
osabnis Nov 9, 2026
b6a57b3
[GPU] SDPA: test few query rows against a host reference
osabnis Nov 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 1 addition & 8 deletions src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,20 +60,12 @@ JitConstants SDPAOptGeneratorBase::get_jit_constants_base(const kernel_impl_para
GPU_DEBUG_TRACE_DETAIL << "k_head_size = " << k_head_size << ", v_head_size = " << v_head_size << "\n";

size_t data_inputs_num = get_data_inputs_num(*desc);
size_t attn_mask_idx = ScaledDotProductAttentionInputIdx::ATTN_MASK;
const bool has_attn_mask_input = sdpa_has_runtime_attn_mask_input(params);
if (desc->attn_mask_val.has_value()) {
jit.make("STATIC_SCALAR_ATTN_MASK_VALUE", desc->attn_mask_val.value());
jit.make("HAS_ATTN_MASK_INPUT", 0);
} else {
jit.make("HAS_ATTN_MASK_INPUT", has_attn_mask_input ? 1 : 0);
if (has_attn_mask_input) {
const auto& attn_mask_layout = params.get_input_layout(attn_mask_idx);
// Enable clamping only if attn_mask dtype differs from softmax_accumulator_type(f32)
if (attn_mask_layout.data_type != data_types::f32) {
jit.make("CLAMP_ATTN_MASK_INPUT", 1);
}
}
}
size_t scale_idx = ScaledDotProductAttentionInputIdx::SCALE;
if ((data_inputs_num > scale_idx) && (!desc->scale_val.has_value())) {
Expand Down Expand Up @@ -300,6 +292,7 @@ JitConstants SDPAOptGeneratorFinalization::get_jit_constants(const kernel_impl_p
auto jit = SDPAOptGeneratorBase::get_jit_constants_base(params, SDPAStage::FINALIZATION);
jit.make("SDPA_STAGE_1", 1);
jit.make("TARGET_SEQ_LEN_BLOCK_SIZE", get_target_seq_len_block_size());
jit.make("MAX_PARTITIONS_NUM", get_max_partitions_num());
return jit;
}

Expand Down
6 changes: 6 additions & 0 deletions src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,12 @@ inline size_t get_target_seq_len_block_size() {
return block_size;
}

// The finalization stage keeps one value per key partition in local memory
inline size_t get_max_partitions_num() {
constexpr size_t max_partitions_num = 128;
return max_partitions_num;
}

inline size_t get_sg_number_scale_factor(const device_info& info, size_t head_size, size_t kernel_type) {
constexpr size_t subgroup_size = 16;
const size_t optimal_scale_factor = 2;
Expand Down
51 changes: 39 additions & 12 deletions src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,34 @@ using namespace cldnn;

namespace ov::intel_gpu::ocl {

namespace {

// Decode with grouped-query heads stacked as query rows has 3-4 rows per KV head on common models.
// Single-token kernel reads the keys once per row (Arc B390, 4K keys: +0.3% at 6 rows, -9% at 8, -26% at 16).
constexpr int64_t max_single_token_rows = 4;

// Several query rows stay on the multi-token kernels where the single-token stages are not exact for them:
// per-channel compressed KV needs the micro kernel, every key partition adds the sink to its softmax sum,
// the finalization stage merges at most get_max_partitions_num() partitions, and the query's local memory is K wide
// while its loads and the output reduction use the V width.
bool single_token_supports_query_rows(const RuntimeParams& params) {
return !SDPAOpt::has_per_channel_compressed_kv(params) && !params.typed_desc<scaled_dot_product_attention>()->has_sink_input &&
get_partitions_num(params, SDPAStage::SINGLE_TOKEN) <= get_max_partitions_num() &&
SDPABase::get_effective_v_head_size(params) <= SDPABase::get_effective_k_head_size(params);
}

bool use_multi_token_stage(const RuntimeParams& params) {

@allnes allnes Nov 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

V_HEAD_SIZE > K_HEAD_SIZE also needs a fallback here. With Q=3, Q/K width 64 and V width 256, SG_SCALE_FACTOR=1 allocates query_local[64], but the unconditional query loads/stores address indices 0–255. For V width 128, the factor-2 output reduction reuses that array across indices 0–127. These accesses exceed the declared local array, although the numeric repro passes on the tested runtimes. Keep the multi-token route for V>K, or bound every query load and size scratch for both query storage and V-wide reduction; cover asymmetric widths with static and dynamic tests.

@osabnis osabnis Nov 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, fixed in “keep V wider than K on the multi-token kernels”: single_token_supports_query_rows() now also requires V ≤ K. As you said, the numbers don’t show this, so the tests check which kernel ran, for V = 128 and 256 with K = 64, static and dynamic. I also ran the grid with temporary bounds checks on the single-token kernel’s local and internal buffers. They were clean with the fix, and the V>K cases failed with it removed.

const auto query_rows = get_target_seq_len(params);
return query_rows > max_single_token_rows || (query_rows > 1 && !single_token_supports_query_rows(params));
}

// Micro single-token kernel has no key partitions, so it is used for one query row only.
bool single_query_row(const RuntimeParams& params) {
return get_target_seq_len(params) <= 1;
}

} // namespace

bool SDPAOpt::has_per_channel_compressed_kv(const kernel_impl_params& params) {
const auto desc = params.typed_desc<scaled_dot_product_attention>();
if (!desc->is_kv_compressed) {
Expand Down Expand Up @@ -90,7 +118,7 @@ class SDPAOptImpl : public SDPAImplBase {
} else {
auto is_indirect = params.typed_desc<scaled_dot_product_attention>()->indirect_axis != -1;
GPU_DEBUG_TRACE_DETAIL << "add stage for non-dynamic, is_indirect = " << is_indirect << "\n";
if (is_prefill_stage(params) || unaligned_head_size(params)) {
if (use_multi_token_stage(params) || unaligned_head_size(params)) {

@allnes allnes Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finalization allocates max_logits_u_exp_sum[128], but its load, update and output-reduction loops use the unrestricted partition count. With 256 keys per partition, K=32769 already requires 129 entries and accesses index 128 outside the array. The new Q=2–4 routing exposes this unsupported case. Keep the multi-token fallback when the partition count exceeds 128, consistently in stage construction and execution, or change finalization to reduce arbitrary partition counts without exceeding local storage.

@osabnis osabnis Nov 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed for the new route in 3db4909: 2-4 query rows with more than 128 key partitions keep the multi-token kernels. The limit now comes from one host constant (get_max_partitions_num()), passed to the finalization kernel as MAX_PARTITIONS_NUM instead of a kernel-local define.
New test: 3 query rows over 33024 keys (129 partitions).

The single-row path has the same 128-partition limit today (decode above 32K keys when micro SDPA is not used); I'd rather fix the finalization kernel for any partition count in a separate PR than grow this one.

@allnes allnes Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each single-token partition adds the sink contribution to its local softmax denominator. Finalization then combines all partition denominators, counting the same virtual key once per partition. With P partitions, the global denominator contains P sink contributions instead of one, incorrectly reducing the attention output. Include the sink exactly once in the global softmax normalization, with the correct maximum-logit adjustment, or retain the multi-token route for sink inputs.

@osabnis osabnis Nov 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, each partition adds the sink to its exp_sum. Fixed for the new route in 3db4909: with a sink input, 2-4 query rows keep the multi-token route.
New test: sdpa_gpu_sink_test with 3 query rows over 300 keys (2 partitions), which fails without this change.

As with the previous comment, the single-row path already behaves this way when it has more than one partition; I can fix the finalization (add the sink once, after the global max) in a follow-up PR.

@allnes allnes Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Q=2–4 threshold replaces supported micro prefill with the regular single-token path for all eligible inputs. This path reads K/V separately for each query row; with multiple key partitions it also needs scratch buffers and an additional finalization launch. Query length alone therefore discards the faster implementation for some supported shapes. Keep micro prefill eligible and restrict the partitioned route to shapes where its additional work is justified.

@osabnis osabnis Nov 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The description's numbers compare against master's multi-token selection on Arc B390: decode with 3-4 query rows gets faster as the context grows (+2.6% to +9.6% at depth 4096-8192 on Llama-3.2 1B, Phi-4 mini and Qwen3.5 4B) and is within run-to-run noise at depth 0. If you prefer, I can keep micro prefill when the keys fit in one partition, since that is where the partitioned kernel has no work-group advantage.

@allnes allnes Nov 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new few-row route returns NaN for a fully masked key partition. For an FP32 SDPA primitive with Q=3, KV=512 and D=64, leave the first 256 keys unmasked and set the remaining mask entries to -Inf: all 384 output elements become NaN, while the multi-token base produces finite output. The empty partition divides by zero in sdpa_opt.cl:756, and finalization preserves NaN when multiplying its partial output by zero. Handle zero-sum partitions before normalization/finalization, or exclude this precision from the new route. This is confirmed at primitive level; the standard FP32 frontend currently decomposes SDPA.

@osabnis osabnis Nov 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in “clamp the attention mask for every mask type”. The clamp to the lowest finite value was skipped for f32 masks, so -Inf reached the softmax. It now always applies. The same NaN also hit f16 SDPA with an f32 mask, and one query row too, since that path is the same kernel. sdpa_few_rows_grid covers fully masked partitions and rows with f16/f16, f16/f32 and f32/f32 data and masks, and these cases fail without the fix.

if (is_indirect) {
GPU_DEBUG_TRACE_DETAIL << "add stage for indirect non-dynamic with prefill_stage \n";
add_stage(indirect_multi_tokens, params);
Expand All @@ -115,7 +143,7 @@ class SDPAOptImpl : public SDPAImplBase {
#ifdef ENABLE_ONEDNN_FOR_GPU
const auto& gfx_ver = params.get_program().get_engine().get_device_info().gfx_ver;
bool is_ARL_H = (gfx_ver.major == 12 && gfx_ver.minor == 74);
bool can_use_micro_sdpa = SDPAOpt::supports_micro_sdpa(params) && !is_ARL_H && !is_indirect;
bool can_use_micro_sdpa = SDPAOpt::supports_micro_sdpa(params) && !is_ARL_H && !is_indirect && single_query_row(params);

@allnes allnes Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Q=2–4 now skips micro multi-token execution, while micro single-token is restricted to Q=1. However, validation still accepts per-channel compressed KV based on micro support. The selected regular kernel indexes scales by token Y with X=0, whereas [B,H,1,D] per-channel scales require Y=0 and X=channel. This produces incorrect dequantization and out-of-bounds scale reads. Preserve the micro multi-token route for per-channel KV in both stage construction and execution, or implement channel-indexed dequantization before selecting the regular path.

@osabnis osabnis Nov 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thanks. Fixed in 3db4909: with per-channel compressed KV, 2-4 query rows now stay on the multi-token route (micro), in stage construction, execution and buffer sizing alike, through one predicate (single_token_supports_query_rows()).
New tests: sdpa_gpu_compressed_kv_per_channel with 3 query rows (int4 and int8).

if (can_use_micro_sdpa) {
add_stage(regular_micro_single_token, params);
}
Expand All @@ -131,10 +159,10 @@ class SDPAOptImpl : public SDPAImplBase {
[[nodiscard]] event::ptr execute(const std::vector<event::ptr>& events, primitive_inst& instance) override {
const auto& params = *instance.get_impl_params();
auto new_params = SDPABase::requires_shape_canonicalization(params) ? SDPABase::static_canonicalize_shapes(params) : params;
bool is_prefill = is_prefill_stage(new_params);
const bool use_multi_tokens = is_prefill || unaligned_head_size(new_params);
const bool multi_token_rows = use_multi_token_stage(new_params);
const bool use_multi_tokens = multi_token_rows || unaligned_head_size(new_params);
bool is_indirect = need_indirect_load(static_cast<scaled_dot_product_attention_inst&>(instance));
GPU_DEBUG_TRACE_DETAIL << "execute indirect = " << is_indirect << ", prefill = " << is_prefill << "\n";
GPU_DEBUG_TRACE_DETAIL << "execute indirect = " << is_indirect << ", multi_token_rows = " << multi_token_rows << "\n";
update_rt_params(instance);

kernel_dump_info.clear_entries();
Expand All @@ -160,7 +188,7 @@ class SDPAOptImpl : public SDPAImplBase {
return execute_stage(events, instance, is_indirect ? indirect_multi_tokens : regular_multi_tokens);
}
#ifdef ENABLE_ONEDNN_FOR_GPU
if (has_stage(regular_micro_single_token) && !is_indirect) {
if (has_stage(regular_micro_single_token) && !is_indirect && single_query_row(new_params)) {
return execute_stage(events, instance, regular_micro_single_token);
}
#endif
Expand All @@ -176,15 +204,14 @@ class SDPAOptImpl : public SDPAImplBase {

[[nodiscard]] std::vector<BufferDescriptor> get_internal_buffer_descs(const RuntimeParams& params) const override {
std::vector<BufferDescriptor> internal_buffers;
auto desc = params.typed_desc<scaled_dot_product_attention>();
auto extended_input_q_transpose_order = extend_order_in_num_heads_dim(desc->input_q_transpose_order);
auto params_canonicalization = SDPABase::requires_shape_canonicalization(params) ? SDPABase::static_canonicalize_shapes(params) : params;
const auto head_size = get_head_size(params_canonicalization.get_input_layout(0), extended_input_q_transpose_order);
// exp_sums and max_logits hold one value per output row and partition; the output is V_HEAD_SIZE wide
const auto v_head_size = SDPABase::get_effective_v_head_size(params_canonicalization);

const auto num_of_partitions = get_partitions_num(params_canonicalization, SDPAStage::SINGLE_TOKEN);
const auto is_prefill = is_prefill_stage(params_canonicalization);
const size_t buf_elements_count = (num_of_partitions == 1 || is_prefill) ? 1 : params.output_layouts[0].count() / head_size * num_of_partitions;
const size_t tmp_out_elements_count = (num_of_partitions == 1 || is_prefill) ? 2 : params.output_layouts[0].count() * num_of_partitions;
const auto multi_token_rows = use_multi_token_stage(params_canonicalization);
const size_t buf_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 1 : params.output_layouts[0].count() / v_head_size * num_of_partitions;
const size_t tmp_out_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 2 : params.output_layouts[0].count() * num_of_partitions;

internal_buffers.emplace_back(buf_elements_count, ov::element::f32);
internal_buffers.emplace_back(buf_elements_count, ov::element::f32);
Expand Down
5 changes: 3 additions & 2 deletions src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_utils.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,10 +205,11 @@ inline ChannelName get_transposed_channel(ChannelName c, const std::vector<int64
return ChannelName::UNKNOWN;
}

inline bool is_prefill_stage(const RuntimeParams& params) {
// Number of query rows, -1 if dynamic
inline int64_t get_target_seq_len(const RuntimeParams& params) {
auto desc = params.typed_desc<cldnn::scaled_dot_product_attention>();
const auto target_seq_len = extract_dim(get_transposed_channel(ChannelName::Y, desc->input_q_transpose_order), params.input_layouts[0]);
return target_seq_len.is_static() ? target_seq_len.get_length() > 1 : false;
return target_seq_len.is_static() ? target_seq_len.get_length() : -1;
}

inline bool unaligned_head_size(const RuntimeParams& params) {
Expand Down
26 changes: 9 additions & 17 deletions src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -588,12 +588,11 @@ KERNEL(sdpa_opt)(
#endif
qk_val[seq_idx] += INPUT0_VAL_MIN;
#elif !IS_CAUSAL && HAS_ATTN_MASK_INPUT
const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, target_seq_idx + seq_idx, start_partition_idx + seq_len);
// _SAFE index macro doesn't parenthesize its arguments, so the modulo for a broadcast mask needs the whole sum
const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, (target_seq_idx + seq_idx), (start_partition_idx + seq_len));
INPUT3_COMPUTE_TYPE mask_val = DECODE_INPUT3_COMPUTE_TYPE(attn_mask[attn_mask_offset]);
#ifdef CLAMP_ATTN_MASK_INPUT
// Conditionally clamp attention mask when attention mask differs from SOFTMAX_ACCUMULATOR_TYPE(f32)
// A key partition masked only with -Inf would have an infinite maximum and turn its softmax into NaN
mask_val = INPUT3_MAX_FUNC(mask_val, INPUT3_VAL_MIN);
#endif
qk_val[seq_idx] += mask_val;
#elif defined(STATIC_SCALAR_ATTN_MASK_VALUE)
qk_val[seq_idx] += STATIC_SCALAR_ATTN_MASK_VALUE;
Expand Down Expand Up @@ -1069,14 +1068,17 @@ inline MASK_VECTOR_TYPE FUNC(load_attn_mask)(OPTIONAL_SHAPE_INFO_ARG
#endif
#if !IS_CAUSAL && HAS_ATTN_MASK_INPUT
const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, target_seq_idx, source_seq_idx);
// A mask broadcast over the keys (e.g. a scalar mask) has a single element along X; the _SAFE index above wraps
// to it, so the following keys must not step past it.
const uint mask_key_pitch = INPUT3_SIZE_X == 1 ? 0 : 1;
if (target_seq_idx >= (uint)TARGET_SEQ_LEN) {
unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) {
mask_vec[i] = NAN;
}
} else {
if (source_seq_idx + SUBGROUP_SIZE <= (uint)SOURCE_SEQ_LEN) {
unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) {
const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i];
const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i * mask_key_pitch];
mask_vec[i] = DECODE_INPUT3_COMPUTE_TYPE(mask_val);
}
} else {
Expand All @@ -1087,7 +1089,7 @@ inline MASK_VECTOR_TYPE FUNC(load_attn_mask)(OPTIONAL_SHAPE_INFO_ARG
const uint max_mask_offset = min(source_seq_idx + SUBGROUP_SIZE, (uint)SOURCE_SEQ_LEN);
for (uint i = 0; i < SUBGROUP_SIZE; i++) {
if (source_seq_idx + i < max_mask_offset) {
const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i];
const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i * mask_key_pitch];
mask_vec[i] = DECODE_INPUT3_COMPUTE_TYPE(mask_val);
} else {
mask_vec[i] = NAN;
Expand Down Expand Up @@ -2844,8 +2846,6 @@ KERNEL(sdpa_opt)(
// max_logits [batch, heads_num, q_len, partition_idx]
// tmp_out [batch, heads_num, q_len, partition_idx, head_size]

#define MAX_PARTITIONS_NUM 128

REQD_SUB_GROUP_SIZE(SUBGROUP_SIZE)
KERNEL(sdpa_opt_finalization_stage)(
OPTIONAL_SHAPE_INFO_ARG
Expand Down Expand Up @@ -2933,15 +2933,7 @@ KERNEL(sdpa_opt_finalization_stage)(
#endif
acc += out_val_f * TO_SOFTMAX_ACCUMULATOR_TYPE(max_logits_u_exp_sum[partition_idx]);
}
const uint out_offset = b0_idx * (NUM_HEADS * TARGET_SEQ_LEN * V_HEAD_SIZE) +
#ifdef OUTPUT_TRANSPOSE_ORDER_PRESENT
target_seq_idx * (NUM_HEADS * V_HEAD_SIZE) +
b1_idx * (V_HEAD_SIZE) +
#else
b1_idx * (TARGET_SEQ_LEN * V_HEAD_SIZE) +
target_seq_idx * (V_HEAD_SIZE) +
#endif
local_id;
const uint out_offset = SDPA_OUTPUT_GET_INDEX(b0_idx, b1_idx, target_seq_idx, local_id);

output[out_offset] = TO_OUTPUT_TYPE(acc / global_exp_sum);
}
Expand Down
Loading
Loading