Omni: Enhance solattn and multi-candidate support in omni xpu kernel - #669
Open
xiangyuT wants to merge 27 commits into
Open
Omni: Enhance solattn and multi-candidate support in omni xpu kernel #669xiangyuT wants to merge 27 commits into
xiangyuT wants to merge 27 commits into
Conversation
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
B580: integrate local MiniMax H3 kernel tuning
…s-20260902 fix(omni): keep fully masked Sol rows finite
Omni: complete Kitchen SolAttn controls and FP8 view correctness
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Tuning eligibility and several policy-dispatch paths currently bypass the documented safety warnings.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Enhances Omni XPU kernels with B580-aware policies, expanded Sol-Attn controls, fat BMG binaries, and fused H3 RMS modulation.
Changes:
- Adds B580 device-policy dispatch and tuning controls.
- Extends Sol-Attn with top-k, masking, tail, and block-length controls.
- Adds ComfyUI H3 RMS modulation integration and coverage.
File summaries
| File | Description |
|---|---|
omni/tests/test_comfyui_omnixpu_h3_rms_modulation.py |
Tests the H3 adapter contract. |
omni/tests/test_comfyui_omnixpu_bootstrap.py |
Covers adapter configuration. |
omni/omni_xpu_kernel/WHL_BUILD_INSTALL.md |
Documents fat BMG candidates. |
omni/omni_xpu_kernel/tests/test_platform_dispatch_source.py |
Tests native dispatch policies. |
omni/omni_xpu_kernel/tests/test_packaging.py |
Tests targets and tuning flags. |
omni/omni_xpu_kernel/tests/test_package_api.py |
Tests RMS modulation API. |
omni/omni_xpu_kernel/tests/test_norm_correctness.py |
Validates fused RMS modulation. |
omni/omni_xpu_kernel/tests/test_kitchen_rope.py |
Expands long-sequence coverage. |
omni/omni_xpu_kernel/tests/test_kitchen_fp8.py |
Tests misaligned FP8 input. |
omni/omni_xpu_kernel/tests/test_int8_correctness.py |
Restricts cache tests to XPU. |
omni/omni_xpu_kernel/tests/test_cute_sol_attn_correctness.py |
Tests new Sol-Attn semantics. |
omni/omni_xpu_kernel/tests/test_cute_sol_attn_api.py |
Tests control-aware API contracts. |
omni/omni_xpu_kernel/tests/test_cute_minimax_h3_vae_d64_correctness.py |
Expands candidate shapes. |
omni/omni_xpu_kernel/tests/test_b580_forced_sku_route_fallback.py |
Tests forced-profile fallback. |
omni/omni_xpu_kernel/setup.py |
Adds fat targets and tuning defines. |
omni/omni_xpu_kernel/README.md |
Documents policies and overrides. |
omni/omni_xpu_kernel/omni_xpu_kernel/norm/__init__.py |
Exposes fused RMS modulation. |
omni/omni_xpu_kernel/omni_xpu_kernel/device.py |
Exposes physical/effective policy data. |
omni/omni_xpu_kernel/omni_xpu_kernel/cute/sol_attn_torch.cpp |
Adds controlled Sol-Attn dispatch. |
omni/omni_xpu_kernel/omni_xpu_kernel/cute/sol_attn_prepare.cpp |
Adds top-k route preparation. |
omni/omni_xpu_kernel/omni_xpu_kernel/cute/sol_attn_mainloop.hpp |
Implements new attention controls. |
omni/omni_xpu_kernel/omni_xpu_kernel/cute/sol_attn_config.h |
Defines B580 tile policy. |
omni/omni_xpu_kernel/omni_xpu_kernel/cute/cute_fmha_torch.cpp |
Adds B580 FMHA candidates. |
omni/omni_xpu_kernel/omni_xpu_kernel/cute/__init__.py |
Extends the Python Sol-Attn API. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/svdq_fused_postproc.cpp |
Adds generic and candidate policies. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/svdq_dequant.cpp |
Adds SVDQ candidate dispatch. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/norm.cpp |
Implements fused H3 RMS modulation. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/kitchen_rope_sycl.cpp |
Adds B580 candidate routing. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/kitchen_rms_rope_sycl.cpp |
Adds B580 RMS-RoPE kernel. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/kernel_tuning_overrides.h |
Centralizes tuning defaults. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/int8_tensorwise_sycl.cpp |
Uses centralized tuning. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/int8_scaleback_esimd.cpp |
Adds policy-based scaleback. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/int8_quantize_esimd.cpp |
Adds row-quantization overrides. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/int8_convrot_quant_esimd.cpp |
Uses centralized tuning. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/int8_convrot_dequant_esimd.cpp |
Uses centralized tuning. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/group_norm_bmg.cpp |
Adds GroupNorm tuning controls. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/fp8_quant.cpp |
Adds B580 stochastic policy. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/fp8_dequant_esimd.cpp |
Supports misaligned FP8 views. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/device_utils.h |
Centralizes selection and warnings. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bmg_kernel_policy.h |
Defines generic/candidate policies. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bmg_device_warning.h |
Implements one-shot warnings. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bmg_device_policy.h |
Defines BMG identity selection. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bindings.cpp |
Exposes policies and tuning metadata. |
omni/omni_xpu_kernel/omni_xpu_kernel/csrc/adaln.cpp |
Adds B580 AdaLN candidate dispatch. |
omni/omni_xpu_kernel/CHANGELOG.md |
Records fat-target support. |
omni/ComfyUI-OmniXPU/patches/__init__.py |
Registers the H3 adapter. |
omni/ComfyUI-OmniXPU/config.py |
Adds the adapter feature flag. |
omni/ComfyUI-OmniXPU/adapters/h3_rms_modulation.py |
Implements ComfyUI fused routing. |
Review details
- Files reviewed: 50/50 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+415
to
+420
| result["performance_claim_allowed"] = | ||
| !selection.forced && | ||
| selection.b580_policy_candidate == | ||
| omni_xpu::device::B580PolicyCandidate::none && | ||
| selection.kernel_profile != | ||
| omni_xpu::device::BmgKernelProfile::generic_bmg; |
Comment on lines
+456
to
+460
| bool use_b580_h3_rms_rope(const torch::Tensor& q) { | ||
| auto& queue = utils::get_queue(q.device()); | ||
| const auto selection = device::get_bmg_selection_unwarned(queue); | ||
| return selection.physical_sku == device::BmgSku::b580 && | ||
| !selection.forced; |
Comment on lines
331
to
+332
| scale_value = q.shape[-1] ** -0.5 if scale is None else float(scale) | ||
| prepared = ops.prepare( | ||
| prepared = ops.prepare_with_controls( |
Comment on lines
+589
to
+595
| auto& queue = | ||
| c10::xpu::getCurrentXPUStream(q.device().index()).queue(); | ||
| const auto selection = | ||
| omni_xpu::device::get_bmg_selection_unwarned(queue); | ||
| if (selection.b580_policy_candidate == | ||
| omni_xpu::device::B580PolicyCandidate:: | ||
| h3_vae_d64_s1797_kv_tile) { |
Signed-off-by: xiangyuT <xiangyu.tian@intel.com>
xiangyuT
marked this pull request as ready for review
September 4, 2026 01:27
Replace the deferred RMS bridge with the reviewed source-guarded H3 forward adapter. Preserve the actual curve-checkpoint FP32-to-BF16 scale/shift conversion and the ComfyUI cast/offload lifecycle. Validated on unforced B70 0xE223: 12 native and 12 real-checkpoint XPU points are bit-exact. One 480p/5s functional workflow passed with 1000 routed blocks and 2000 fused calls. Source/interface checks: 332 passed, 11 skipped. Performance no-regression is NOT validated. No formal B/C/B was run; the owner requested this commit/push for handoff. No new clean-image or milestone acceptance claim.
Validate the supported forward API and required callables during registration. Keep unsupported norm/modulation inputs on the original local route, avoiding replayed projection/residual work and duplicate stateful norm hooks. Validated as a source overlay in the accepted B70 XPU image: 373 interface cases plus 101 subtests passed, including 45 H3 adapter cases. Supplemental actual- checkpoint XPU checks passed 48/48 with bit-exact outputs. Clean-image inclusion, canonical workflow/performance gates and full Luna milestone acceptance remain pending. No package or image version changes.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.