Skip to content

[Diffusion][CPU] Adding AMX optimizations for CPU platform - #28527

Merged
mingfeima merged 7 commits into
sgl-project:mainfrom
jianan-gu:sgl_diffusion_cpu_opt1
Jul 9, 2026
Merged

[Diffusion][CPU] Adding AMX optimizations for CPU platform#28527
mingfeima merged 7 commits into
sgl-project:mainfrom
jianan-gu:sgl_diffusion_cpu_opt1

Conversation

@jianan-gu

@jianan-gu jianan-gu commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

This pr takes parts of follow-ups (mentioned in #20816) to bring key AMX based optimizations (scoping from LLM models) for CPU platforms, including:

  • Adding CPU AMX backend, replace SDPA to bring best attention perf.
  • Similar like for LLM models, adding replacement for amx weight packed linear.
  • Apply channel last 3d for VAE model

With above take Wan-AI/Wan2.2-TI2V-5B-Diffusers model as an example, the main stages show:

  • TextEncodingStage speedup by 10x (bf16 encoding, if fp32 stays no change)
  • DenoisingStage speedup by 2x
  • DecodingStage speedup by 1.57x (bf16, enable channel last)

Besides, we still find more chance to bring more speedup (but with small ratio), by adding fusions (compile or kernel based) like rmsnorm/layernorm and newGELUs, will make this in followup PRs.


CI States

Latest PR Test (Base): ❌ Run #28929873450
Latest PR Test (Extra): ❌ Run #28929873154

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

@github-actions github-actions Bot added Multi-modal multi-modal language model sgl-kernel diffusion SGLang Diffusion labels Jun 17, 2026
@jianan-gu jianan-gu changed the title [Diffusion][CPU] [Diffusion][CPU] Adding optimizations for CPU platform Jun 17, 2026
@jianan-gu jianan-gu changed the title [Diffusion][CPU] Adding optimizations for CPU platform [Diffusion][CPU] Adding AMX optimizations for CPU platform Jul 1, 2026
@jianan-gu

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci label Jul 2, 2026

@mingfeima mingfeima left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

some minor changes are needed.

@mingfeima mingfeima added intel cpu cpu backend performance optimization labels Jul 2, 2026
Comment thread test/registered/cpu/test_flash_attn.py Outdated
@jianan-gu
jianan-gu requested a review from mingfeima July 2, 2026 10:42
@mingfeima

Copy link
Copy Markdown
Collaborator

@jianan-gu rebase to fix conflict.

@jianan-gu

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

@mingfeima
mingfeima merged commit 177c048 into sgl-project:main Jul 9, 2026
97 of 153 checks passed
# Avoid unintended computation graph accumulation during inference
if isinstance(p, torch.nn.Parameter):
p.requires_grad = False
local_torch_device = get_local_torch_device()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

we have an existing quant_method.process_weights_after_loading call in L321-L330, please avoid duplicating it

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.

Thanks for pointing that, have submitted changes to avoid such duplicating #30717

@mickqian

Copy link
Copy Markdown
Collaborator

this breaks CI, reverting in #30716

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026
jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cpu cpu backend performance optimization diffusion SGLang Diffusion intel Multi-modal multi-modal language model run-ci sgl-kernel

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants