Persistent SDPA kernel

[wuxun-zhang]

The new kernel implements below method, key points are:

• num of work groups are fixed to num of total XeCores
• dynamically split KV seq length from all seqs into all work groups
• each XeCore gets balanced work units

As of now there are two limitations:

• only decode support (seq_len_qo==1)
• batch_size * num_heads_q <= num of total XeCores

[pengzhao-intel]

maybe add the limitation of this algorithm in the code as well, especially for one with atomic.

[wuxun-zhang]

@pengzhao-intel Added checks here in can_implement().

[Copilot]

Pull Request Overview

This PR introduces a persistent SDPA (Scaled Dot Product Attention) kernel for decode scenarios that implements dynamic load balancing across XeCores. The key innovation is fixing the number of work groups to match total XeCores and dynamically splitting KV sequence length across all work groups for balanced workload distribution.

Key changes:

• New persistent tile scheduler (XeFHMAIndividualPersistentTileScheduler) that distributes work evenly across fixed XeCore count
• New kernel implementation (XeFMHAFwdDynamicSplitKernel) with split-K reduction for partial results
• Support infrastructure including atomic operations (atomicSub, atomicLoad) for synchronization

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
include/cutlass/gpu_generics.h	Adds atomic operations (atomicSub, atomicLoad) for synchronization primitives
examples/06_bmg_flash_attention/xe_fmha_fwd_runner.hpp	Integrates persistent kernel selection and queries hardware XeCore count
examples/06_bmg_flash_attention/CMakeLists.txt	Adds build target for persistent kernel testing
examples/06_bmg_flash_attention/06_xe_fmha_fwd.cpp	Configures persistent kernel with appropriate tile sizes and subgroup layouts
applications/flash_attention_v2/kernel/xe_tile_scheduler.hpp	Implements persistent tile scheduler with dynamic work distribution
applications/flash_attention_v2/kernel/xe_fhma_fwd_kernel.hpp	Implements dynamic split-K kernel with partial result reduction
applications/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp	Updates mainloop to use total block count for remainder masking

[Copilot]

The atomic_ref is constructed with address[0] which dereferences the pointer. This should be *address for clarity and consistency with standard atomic operations patterns.

Suggested change

	auto atm = sycl::atomic_ref<int, sycl::memory_order::relaxed, sycl::memory_scope::device, sycl::access::address_space::generic_space>(address[0]);
	auto atm = sycl::atomic_ref<int, sycl::memory_order::relaxed, sycl::memory_scope::device, sycl::access::address_space::generic_space>(*address);

[petercad]

Suggested change

	merged_res(i + size(FragA{}.shape())) = tA_max(i);
	merged_res(i + 1 + size(FragA{}.shape())) = tA_sum(i);
	merged_res(2*i + size(FragA{}.shape())) = tA_max(i);
	merged_res(2*i + 1 + size(FragA{}.shape())) = tA_sum(i);

[wuxun-zhang]

Thanks for the catch. Updated.

[petercad]

Did you mean to change the tile configuration for non-persistent mode?

Suggested change

	#define NUM_SG _16
	#define NUM_SG _8

[wuxun-zhang]

Changed back to _8 now.