[XPU][FMHA] Add reverse-odd-heads Q-tile scheduler for causal prefill by pengzhao-intel · Pull Request #363 · vllm-project/vllm-xpu-kernels

pengzhao-intel · 2026-05-21T22:40:58Z

Motivation

Thanks cong provides the optimization idea.

In FA2 chunk-prefill on Xe2, the default XeFHMAIndividualTileScheduler dispatches Q-tiles in the same order for every head. For causal prefill this creates a per-head load imbalance:

Tile (head=h, q_blk=0) only attends the diagonal block of KV — light work.
Tile (head=h, q_blk=last) attends nearly the full KV history — heavy work.

Because all heads start from q_blk=0, the first wave is uniformly light and the last wave is uniformly heavy. XE cores idle waiting for the heaviest wave to drain.

Change

Introduce XeFHMAIndividualReverseOrderTileScheduler, a drop-in scheduler variant that reverses the Q-tile order for odd-indexed heads. Adjacent heads now mix light and heavy tiles in the same wave, smoothing per-wave occupancy.

Selection is per-call via a new bool on the args struct:

struct chunk_prefill_args_t {
  ...
  bool reverse_q_order = false;   // new
};

The FA2 op sets it from env var VLLM_XPU_FA_REVERSE_ODD_HEADS:

Value	Behavior
`unset` / `auto`	Shape-based heuristic (default)
`1` / `on` / `true`	Force enabled
`0` / `off` / `false`	Force disabled

Auto heuristic:

reverse = is_causal && !is_local && !is_sink && num_heads_q >= 2;

Performance evidence

Intel Arc B580, FA2 varlen, MQA q=2 / kv=1, head_size=128, fp16, causal. Numbers are FlashAttention_Kernel_Time(us) speedup of on over off.

Threshold sweep (bs ∈ {1, 2, 4}, per-seq qlen ∈ [2048, 32768])

Shape family	Range	Median
bs=1, qlen 8k–32k	+7 % … +50 %	~+20 %
bs=2, qlen 4k–10k	+8 % … +44 %	~+27 %
bs=4, qlen 2k–4k	-0.7 % … +22 %	~+15 %

Long-context shapes (total kvlen up to 262k)

Shape	Speedup
bs=1, qlen 32k–131k (paged & non-paged)	+1.6 % … +9.5 %
bs=2 / bs=6, qlen 10k (multi-batch short prefill)	+5 % … +22 %
bs=1 / bs=2, qlen 8k + cached_kv 57k	+2 % … +3 %

Worst observed: -0.7 % on bs=4, qlen=4k non-paged — within run-to-run noise. Every other config across 28 measured shape variants showed a speedup.

API / backwards compatibility

reverse_q_order defaults to false, so callers that construct chunk_prefill_args_t without going through fmha_xe2.cpp are unaffected.
Env-var default is auto, i.e. enabled for causal multi-head prefill. To restore strict pre-PR behavior set VLLM_XPU_FA_REVERSE_ODD_HEADS=off.

Testing

All existing chunk_prefill code paths are unchanged (head_size, block_size, mask, sink, paged, varlen, fp8 KV). The new scheduler is only instantiated when the new bool is true.
New benchmark configs are automatically included in benchmark_cutlass_flash_attn_varlen.py.
Recommended pre-merge check: run pytest tests/flash_attn/test_flash_attn_varlen_func.py with both VLLM_XPU_FA_REVERSE_ODD_HEADS=on and off to confirm numeric equivalence.

Copilot

Pull request overview

This PR adds an alternative Q-tile scheduling strategy for Xe2 FlashAttention2 chunk-prefill to reduce per-wave load imbalance in causal prefill by reversing Q-tile dispatch order for odd-indexed heads, with runtime selection controlled via a new reverse_q_order argument and an env-var/heuristic in the Xe2 entrypoint.

Changes:

Add XeFHMAIndividualReverseOrderTileScheduler that flips BlockIdxY for odd heads to mix light/heavy tiles across heads.
Add chunk_prefill_args_t::reverse_q_order and dispatch logic to select the new scheduler when enabled.
Parse VLLM_XPU_FA_REVERSE_ODD_HEADS (with an “auto” heuristic) and add additional benchmark shape configs for MQA prefill/long-context scenarios.

Reviewed changes

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

File	Description
csrc/xpu/attn/xe_2/fmha_xe2.cpp	Adds env-var parsing + heuristic to set `args.reverse_q_order` for selecting the new scheduler.
csrc/xpu/attn/xe_2/collective/chunk_prefill_scheduler.hpp	Introduces the reverse-order per-head Q-tile scheduler implementation.
csrc/xpu/attn/xe_2/chunk_prefill.hpp	Adds `reverse_q_order` flag to args and selects scheduler in `kernel_dispatch`.
benchmark/src/get_model_config.py	Adds benchmark configs covering MQA causal prefill and long-context shapes.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    static const char* env = std::getenv("VLLM_XPU_FA_REVERSE_ODD_HEADS");
+    std::string mode = "auto";
+    if (env != nullptr) {
+      mode.assign(env);
+      for (auto& c : mode) c = static_cast<char>(std::tolower(c));


+    bool reverse;
+    if (mode == "1" || mode == "on" || mode == "true" || mode == "yes") {
+      reverse = true;
+    } else if (mode == "0" || mode == "off" || mode == "false" ||
+               mode == "no") {
+      reverse = false;
+    } else {  // "auto" or anything unrecognized
+      reverse = is_causal && !is_local && !is_sink && num_heads_q >= 2;
+    }
+    args.reverse_q_order = reverse;
+  }


pengzhao-intel · 2026-05-21T22:48:24Z

@copilot fix the comment issues

Copilot · 2026-05-21T22:52:50Z

@copilot fix the comment issues

Addressed in commit e9e8ac0.

Fixed VLLM_XPU_FA_REVERSE_ODD_HEADS parsing to call std::getenv() on each invocation (removed static pointer caching) and made tolower safe via unsigned char cast.
Added test_varlen_reverse_odd_heads_env_on_off_equivalence to validate numerical equivalence when forcing env off vs on.

wuxun-zhang · 2026-05-25T01:04:21Z

I think later we can also try longest processing time scheduling (process longer request firstly from all heads), even for single head, it could help avoid load imbalance issue.

pengzhao-intel · 2026-05-25T04:00:27Z

I think later we can also try longest processing time scheduling (process longer request firstly from all heads), even for single head, it could help avoid load imbalance issue.

Yes， I think current implementation is only one way and we may provide different scheduler for different situation.

xinyu-intel · 2026-05-25T05:24:42Z

@jikunshang we may need a more proper way to handle env variables.

baodii

LGTM. @YizhouZ pls comment.

YizhouZ · 2026-05-26T05:04:23Z

Let's wait for benchmark to finally merge this PR.

pengzhao-intel · 2026-05-27T17:26:33Z

Let's wait for benchmark to finally merge this PR.

Let me know when the benchmark is done

Add XeFHMAIndividualReverseOrderTileScheduler, a variant of the XeFHMAIndividualTileScheduler that flips the Q-tile dispatch order for odd-indexed heads. This mitigates the causal-prefill load imbalance where head 0's first wave is dominated by the lightly loaded upper-triangle tiles, by letting odd heads consume the heavy lower-triangle tiles in parallel with the remaining light tiles of even heads. Selection is runtime-controlled by chunk_prefill_args_t::reverse_q_order, set from the FA2 op based on VLLM_XPU_FA_REVERSE_ODD_HEADS: unset / 'auto' : shape-based heuristic (default) '1' / 'on' / 'true' : force enabled '0' / 'off' / 'false' : force disabled The 'auto' heuristic enables the scheduler whenever the kernel is running causal prefill with num_heads_q >= 2 and no local/sink masking. An FA2 varlen MQA sweep on Arc B580 across batch in {1,2,4} and per-seq qlen in [2048, 131072] showed only a single near-noise (-0.7%) result and speedups up to +50% kernel time, so no workload-size guard is required. Also adds a set of MQA(q=2, kv=1) chunked-prefill and long-context shapes to gen_cutlass_flash_attn_varlen_perf_configs() so the new scheduler can be regression-tested via benchmark_cutlass_flash_attn_varlen. Signed-off-by: pengzhao-intel <patric.zhao@intel.com>

Signed-off-by: pengzhao-intel <patric.zhao@intel.com>

Use a generic mix of MHA / GQA-4:1 / GQA-8:1 prefill shapes (head=128, fp16, causal) at round token counts up to 16k instead of the previous MQA(q=2, kv=1) long-context set. Same intent (exercise the reverse odd-heads scheduler) but covers a broader, more representative shape mix and runs faster. Signed-off-by: pengzhao-intel <patric.zhao@intel.com>

For MHA (num_heads_q == num_heads_kv), reversing Q-tile order for odd heads scatters memory accesses across the sequence, causing L2 cache thrashing when the GPU is fully saturated. GQA/MQA workloads are immune because KV is shared across Q-heads regardless of access order. Tighten the auto-heuristic: always enable for GQA/MQA, but for MHA only enable when the GPU is under-saturated (total_wgs < 4 * sm_count). Fixes regression on Wan-AI (40,40) diffusion model config: 1624us -> ~948us. Signed-off-by: Patric Zhao <patric.zhao@intel.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: pengzhao-intel <patric.zhao@intel.com>

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

Copilot AI review requested due to automatic review settings May 21, 2026 22:40

Copilot started reviewing on behalf of pengzhao-intel May 21, 2026 22:41 View session

Copilot AI reviewed May 21, 2026

View reviewed changes

Copilot started work on behalf of pengzhao-intel May 21, 2026 22:48 View session

pengzhao-intel requested review from YizhouZ, baodii, jikunshang and xinyu-intel May 21, 2026 22:50

Copilot finished work on behalf of pengzhao-intel May 21, 2026 22:53

pengzhao-intel force-pushed the reverse_scheduler branch from 4f72116 to dd14f99 Compare May 22, 2026 00:31

xinyu-intel approved these changes May 23, 2026

View reviewed changes

wuxun-zhang approved these changes May 25, 2026

View reviewed changes

pengzhao-intel requested review from wuxun-zhang and xinyu-intel May 25, 2026 03:59

wuxun-zhang approved these changes May 25, 2026

View reviewed changes

xinyu-intel approved these changes May 25, 2026

View reviewed changes

baodii approved these changes May 25, 2026

View reviewed changes

pengzhao-intel force-pushed the reverse_scheduler branch from adf8df4 to b531241 Compare May 29, 2026 09:00

jikunshang force-pushed the reverse_scheduler branch from b531241 to 2752fbf Compare June 15, 2026 11:33

pengzhao-intel and others added 4 commits June 15, 2026 04:33

Address reviewer feedback for reverse odd-heads env handling

c01e3db

Signed-off-by: pengzhao-intel <patric.zhao@intel.com>

fix pre commit

2752fbf

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

jikunshang mentioned this pull request Jun 22, 2026

v0.1.11 release tracker #430

Closed

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[XPU][FMHA] Add reverse-odd-heads Q-tile scheduler for causal prefill#363

[XPU][FMHA] Add reverse-odd-heads Q-tile scheduler for causal prefill#363
pengzhao-intel wants to merge 5 commits into
mainfrom
reverse_scheduler

pengzhao-intel commented May 21, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

pengzhao-intel commented May 21, 2026

Uh oh!

Copilot AI commented May 21, 2026

Uh oh!

wuxun-zhang commented May 25, 2026

Uh oh!

pengzhao-intel commented May 25, 2026

Uh oh!

xinyu-intel commented May 25, 2026

Uh oh!

baodii left a comment

Uh oh!

YizhouZ commented May 26, 2026

Uh oh!

pengzhao-intel commented May 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Uh oh!

Conversation

pengzhao-intel commented May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Change

Performance evidence

Threshold sweep (bs ∈ {1, 2, 4}, per-seq qlen ∈ [2048, 32768])

Long-context shapes (total kvlen up to 262k)

API / backwards compatibility

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

pengzhao-intel commented May 21, 2026

Uh oh!

Copilot AI commented May 21, 2026

Uh oh!

wuxun-zhang commented May 25, 2026

Uh oh!

pengzhao-intel commented May 25, 2026

Uh oh!

xinyu-intel commented May 25, 2026

Uh oh!

baodii left a comment

Choose a reason for hiding this comment

Uh oh!

YizhouZ commented May 26, 2026

Uh oh!

pengzhao-intel commented May 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

pengzhao-intel commented May 21, 2026 •

edited

Loading