Skip to content

[rls-v3.11] fused sdpa training #4823

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
syurkevi wants to merge 19 commits into
base: rls-v3.11
Choose a base branch
from
Open

[rls-v3.11] fused sdpa training #4823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
7c94afc
common: sdpa: adds backwards sdpa primitives
syurkevi Feb 24, 2026
7624f46
common: sdpa: prepares internal primitive for bwd testing
syurkevi Feb 24, 2026
dfc8ad1
gtests: internals: adds sdpa training tests
syurkevi Feb 24, 2026
acea9d0
xe: sdpa: prepares config structs for bwd pass
syurkevi Feb 24, 2026
0cac8dc
gpu: gemm: jit: allow LLR gemms
syurkevi Mar 7, 2026
094971f
gpu: gemm: jit: enables ukernel with multiple kernel sources
syurkevi Mar 7, 2026
0905774
xe: sdpa: splits fwd/bwd gpu training primitives
syurkevi Feb 25, 2026
1ee8ade
xe: sdpa: split gemm setup for backwards pass
syurkevi Feb 25, 2026
e8a0c47
xe: sdpa: calculates forward logusmexp to ws
syurkevi Feb 25, 2026
c2cf346
xe: sdpa: updates tile_ops packed load/stores
syurkevi Feb 25, 2026
cb9dfe8
xe: sdpa: adds bwd kernel implementation
syurkevi Feb 25, 2026
fa5e7e8
xe: sdpa: adds create_sdpa_pd for backwards pass
syurkevi Feb 26, 2026
112702e
xe: sdpa: rename fwd/bwd primitives
syurkevi Mar 6, 2026
4d972d3
gtests: internals: separate sdpa from internals
syurkevi Mar 7, 2026
fa8d5f6
common: sdpa: move prop_kind param before attrs
syurkevi Mar 9, 2026
f0eadc9
xe: sdpa: enable transpose_k for training
syurkevi Mar 11, 2026
9e9721e
common: sdpa: refactors pd accessors, misc cleanup
syurkevi Mar 12, 2026
69d9cbc
xe: sdpa: adds slm for f32 softmax
syurkevi Mar 18, 2026
5dca53f
xe: sdpa: handles inf_as_zero in logsumexp
syurkevi Mar 21, 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
4 changes: 4 additions & 0 deletions src/common/memory_tracking.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,6 +325,10 @@ enum {
key_rnn_ptrs_wei_iter,
key_rnn_ptrs_wei_projection,
key_shuffle_precompute_transpose,
key_sdpa_Di,
key_sdpa_dQ_reduction,
key_sdpa_dK_reduction,
key_sdpa_dV_reduction,
key_softmax_dst_scales,
key_softmax_reduction,
key_softmax_interim_store,
Expand Down
6 changes: 6 additions & 0 deletions src/common/primitive_hashing.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -734,6 +734,7 @@ size_t get_desc_hash(const sdpa_desc_t &desc) {
size_t seed = 0;
// Kinds
seed = hash_combine(seed, static_cast<size_t>(desc.primitive_kind));
seed = hash_combine(seed, static_cast<size_t>(desc.prop_kind));
// Memory descriptors
seed = hash_combine(seed, get_md_hash(desc.q_desc));
seed = hash_combine(seed, get_md_hash(desc.k_desc));
Expand All @@ -742,7 +743,12 @@ size_t get_desc_hash(const sdpa_desc_t &desc) {
seed = hash_combine(seed, desc.kq_zero_points.get_hash());
seed = hash_combine(seed, desc.vs_scales.get_hash());
seed = hash_combine(seed, desc.vs_zero_points.get_hash());
seed = hash_combine(seed, get_md_hash(desc.dS_desc));
seed = hash_combine(seed, get_md_hash(desc.dst_desc));
seed = hash_combine(seed, get_md_hash(desc.diff_dst_desc));
seed = hash_combine(seed, get_md_hash(desc.diff_q_desc));
seed = hash_combine(seed, get_md_hash(desc.diff_k_desc));
seed = hash_combine(seed, get_md_hash(desc.diff_v_desc));
seed = hash_combine(seed, get_md_hash(desc.attn_mask_desc));
seed = hash_combine(seed, get_md_hash(desc.scale_desc));
// Scale type
Expand Down
6 changes: 6 additions & 0 deletions src/common/primitive_serialization.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -581,14 +581,20 @@ void serialize(serialization_stream_t &sstream, const sum_desc_t &desc) {
void serialize(serialization_stream_t &sstream, const sdpa_desc_t &desc) {
// Kind
sstream.append(desc.primitive_kind);
sstream.append(desc.prop_kind);
serialize(sstream, desc.q_desc);
serialize(sstream, desc.k_desc);
serialize(sstream, desc.v_desc);
desc.kq_scales.serialize(sstream);
desc.kq_zero_points.serialize(sstream);
desc.vs_scales.serialize(sstream);
desc.vs_zero_points.serialize(sstream);
serialize(sstream, desc.dS_desc);
serialize(sstream, desc.dst_desc);
serialize(sstream, desc.diff_dst_desc);
serialize(sstream, desc.diff_q_desc);
serialize(sstream, desc.diff_k_desc);
serialize(sstream, desc.diff_v_desc);
serialize(sstream, desc.attn_mask_desc);
serialize(sstream, desc.scale_desc);
sstream.append(desc.kq_acc_dt);
Expand Down
Loading
Loading