Skip to content

[XPU] Add sage_attn backend#3785

Draft
xuechendi wants to merge 1 commit into
vllm-project:mainfrom
xuechendi:xpu_sage_attn
Draft

[XPU] Add sage_attn backend#3785
xuechendi wants to merge 1 commit into
vllm-project:mainfrom
xuechendi:xpu_sage_attn

Conversation

@xuechendi
Copy link
Copy Markdown
Contributor

@xuechendi xuechendi commented May 21, 2026

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Add sage_attn backend package

see 1.02% speedup per step; pure kernel see 1.26x gain so far.

Metric SAGE_ATTN FLASH_ATTN
Step 1 16.48s 16.66s
Step 2 12.94s 13.18s
Step 3 11.83s 12.09s
Step 4 11.30s 11.56s
Steps 5-17 (avg) ~10.6s ~10.88s
Step 18 (CPU offload spike) 14.01s 14.28s
Steps 19-40 (avg) ~10.6-10.7s ~10.87s
Total gen time 468.7s 480.0s

TODO:

  1. work with cache_dit will hang

Test Plan

  export SYCL_UR_USE_LEVEL_ZERO_V2=0
  export ZE_AFFINITY_MASK=0,1,2,3
  ulimit -n 1048576
  VLLM_WORKER_MULTIPROC_METHOD=spawn \
  DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
  python /workspace/vllm-omni/examples/offline_inference/text_to_video/text_to_video.py \
    --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
    --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
    --height 480 --width 832 --num-frames 81 \
    --num-inference-steps 40 \
    --boundary-ratio 0.875 --flow-shift 5.0 --fps 16 \
    --tensor-parallel-size 4 --enable-cpu-offload \
    --vae-use-slicing --vae-use-tiling \
    --enable-diffusion-pipeline-profiler \
    --output t2v_480p_tp4_cpu_offload_sage.mp4

Test Result

t2v_480p_tp4_cpu_offload_sage.mp4

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
  • The test results. Please paste the results comparison before and after, or the e2e results.
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
  • (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant