[Hardware][XPU] Register batch-invariant kernels for XPU by tzielinski-habana · Pull Request #41934 · vllm-project/vllm

tzielinski-habana · 2026-05-07T09:53:21Z

Purpose

This pull request adds partial support for batch invariance on Intel XPU devices. This is a requirement for reinforcement learning on XPU.

Test Plan

Unit tests are modified to cover the registered kernels on XPU.

Test Result

All newly added unit tests pass.

Note

The code was co-developed with GitHub Copilot.

…s for XPU Co-authored-by: GitHub Copilot Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

mergify · 2026-05-07T09:54:00Z

Documentation preview: https://vllm--41934.org.readthedocs.build/en/41934/

gemini-code-assist

Code Review

This pull request introduces experimental support for batch invariance on Intel XPU devices, covering operations such as bmm, log_softmax, softmax, and mean. It includes updated documentation, a new test suite for verifying correctness and invariance on XPU, and the necessary dispatch overrides. The reviewer identified that allow_override=True must be added to the torch.library.Library.impl calls in both the implementation and the test fixtures to prevent potential RuntimeError exceptions caused by conflicting registrations.

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com>

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

mergify · 2026-05-14T14:44:51Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tzielinski-habana.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com>

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

yewentao256

Thanks for the work!

Could you add more context for this PR? Is there a lot of demands from users to realize this in XPU?

yewentao256 · 2026-05-21T19:02:33Z

Could we try to reuse current unit test instead creating new one?

I was thinking about it before I marked the PR as ready for review. There are a couple of ways we can go and I wasn't sure which one would be most welcome by the community. Reusing existing tests would introduce some changes around the "skip_unsupported" decorator. I'd have to modify it to include checks for XPU and introduce separate checks for XPU ("skip_if_not_xpu"?) and CUDA ("skip_if_not_cuda"?). The name "skip_if_not_cuda" itself isn't very good, because it checks the condition CUDA and >= Ampere (SM80). All in all, I decided not to go this way, but if you like I can prepare a commit that refactors tests and you can decide which way works better.

Also, regarding your question about more context and the demand - the context is, we need batch-invariance for reinforcement learning and this PR is basically the first step to enable it. I don't know about the demand for RL, but I asked the team and I'll let you know when I get the answer.

@yewentao256 please see a preview of what the test refactoring would look like and let me know what you think:
tzielinski-habana#1

I have same comment as @yewentao256 , please not creating another tests/v1/determinism/test_xpu_batch_invariant.py

update to "skip_unsupported" decorator sounds reasonable to me. @wendyliu235 as well, Her team is working on migrate more cuda to run on XPU, might help to provide suggestion for the generalization work

Ok, I'll merge my PR with testing improvements: tzielinski-habana#1

yewentao256 · 2026-05-21T19:03:25Z

+def _register_matmul_overrides(lib, key: str):
+    """Register matmul overrides for batch invariance."""
+    lib.impl("aten::mm", mm_batch_invariant, key)
+    lib.impl("aten::addmm", addmm_batch_invariant, key)
+    lib.impl("aten::matmul", matmul_batch_invariant, key)
+    lib.impl("aten::linear", linear_batch_invariant, key)
+
+
+def _register_common_overrides(lib, key: str):
+    """Register batch-invariant overrides shared across CUDA and XPU."""
+    lib.impl("aten::_log_softmax", _log_softmax_batch_invariant, key)
+    lib.impl("aten::softmax", softmax_batch_invariant, key)
+    lib.impl("aten::_softmax", softmax_batch_invariant, key)
+    lib.impl("aten::mean.dim", mean_batch_invariant, key)
+    # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
+    # (torch._native.ops.bmm_outer_product), so we need allow_override
+    # to replace it at the dispatcher level.
+    lib.impl("aten::bmm", bmm_batch_invariant, key, allow_override=True)


bmm is a kind of matmul?

Right, I assume you're asking this, because I put it with "common overrides" and not "matmul overrides". There are two reasons for this:

bmm_batch_invariant has a dedicated Triton kernel (bmm_kernel) while the ops from matmul overrides all use matmul_kernel_persistent (with a small caveat - matmul_batch_invariant uses matmul_kernel_persistent in 2/3 if/else branches and bmm_kernel in the one remaining branch).

more imporantly, the common overrides are unconditional for both CUDA and XPU, while matmul overrides are only registered for cuda if this condition holds: if current_platform.is_device_capability_family(80). So we need to keep these groups separate. If you prefer some other names for these functions, can you give me some suggestions?

I don't like the refactor here, let's keep the change minimal in this file. we may use key = "cuda" or "xpu" outside, then

_batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key) directly instead of this function warpper

I significantly reduced the scope of this PR. Got rid of the refactor and removed the common ops test. I'm also skipping matmul and linear registrations for now, because they require more work. I plan to add them in another PR later.

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

mergify · 2026-06-24T15:15:18Z

Hi @tzielinski-habana, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

yewentao256 · 2026-06-26T15:40:25Z

+XPU_BACKENDS: list[str] = [
+    "TRITON_ATTN",
+]
+
+
+def skip_unsupported_xpu_backends(backend: str):
+    if current_platform.is_xpu() and backend not in XPU_BACKENDS:
+        pytest.skip(f"Backend {backend} not verified for batch invariance on XPU")


Currently all devices and backends are quite complicated, please organize them all together, design with a cleaner version.

Ok, will do

yewentao256 · 2026-06-26T15:41:30Z

 def enable_batch_invariant_mode():
+    if not (current_platform.is_cuda() or current_platform.is_xpu()):
+        raise NotImplementedError(
+            "Batch invariance is only supported on CUDA and XPU platforms, "
+            f"got {current_platform.device_name}"


Please check if Rocm is enabled

I don't have access to ROCm to check. I will remove this condition

yewentao256 · 2026-06-26T15:42:49Z

+    # Not all batch-invariant kernels are registered on XPU yet
+    # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
+    if current_platform.is_xpu():
+        pytest.xfail("Not all batch-invariant kernels registered on XPU yet")


e2e determinism is not guaranteed

We want it is guaranteed.

Right, but that means the PR needs to be much larger. I wanted to register a subset of the kernels first and add the other ones later.

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

yewentao256

Nice work! Looks much better, thanks!

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

yewentao256

All E2E tests are xfail

# Not all batch-invariant kernels are registered on XPU yet
    # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
    if current_platform.is_xpu():
        pytest.xfail("Not all batch-invariant kernels registered on XPU yet")

Are we sure this is what we want? I hope we can make sure at least one dense e2e like Qwen 3 could pass

tzielinski-habana · 2026-07-15T10:28:32Z

All E2E tests are xfail
# Not all batch-invariant kernels are registered on XPU yet
    # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
    if current_platform.is_xpu():
        pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
Are we sure this is what we want? I hope we can make sure at least one dense e2e like Qwen 3 could pass

@yewentao256 e2e tests are xfailed, because this PR is only the first step, out of two, to enable batch invariance on XPU. It's missing matmul and linear kernels registrations, because they are more problematic and require more code and a custom matmul kernel with tensor descriptors. I'm sure there will be quite a few review comments for the new code.

I wanted to break it down for two reasons:

Two smaller PRs are easier to understand and review
If something unexpected breaks, we can have a more granular bisection

I had to rebase the branch a lot of times to make CI pass, because it is extremely unstable for some reason - every other time I got some interrupts/infrastructure issues. Now I see CI has finally passed, so if you're ok with that, I would love to merge this PR and then I'm going to submit the missing kernels in the next one. If not, I can expand this PR to cover end to end scenarios, but I think the more granular approach is better and we can deliver part of the functionality sooner.

yewentao256

LGTM, thanks for the work!

[Hardware][XPU] Added registration of verified batch-invariant kernel…

c0ffa12

…s for XPU Co-authored-by: GitHub Copilot Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

mergify Bot added documentation Improvements or additions to documentation intel-gpu Related to Intel GPU v1 labels May 7, 2026

gemini-code-assist Bot reviewed May 7, 2026

View reviewed changes

Comment thread vllm/model_executor/layers/batch_invariant.py Outdated

Comment thread tests/v1/determinism/test_xpu_batch_invariant.py Outdated

tzielinski-habana and others added 4 commits May 7, 2026 13:17

Refactoring and minor fixes

974211e

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Update tests/v1/determinism/test_xpu_batch_invariant.py

498fa8b

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com>

Merge branch 'main' into xpu_rl

5024a1b

Updated docs

1f38d3b

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

mergify Bot added needs-rebase and removed needs-rebase labels May 14, 2026

Added matmul registration and guarded cublas

34eb0db

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

tzielinski-habana force-pushed the xpu_rl branch from 7108abd to 34eb0db Compare May 18, 2026 15:49

tzielinski-habana and others added 2 commits May 18, 2026 17:53

Merge branch 'main' into xpu_rl

a3c4a04

Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com>

Refactoring

3ee67e1

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

tzielinski-habana force-pushed the xpu_rl branch from cce4149 to 3ee67e1 Compare May 18, 2026 16:13

tzielinski-habana and others added 7 commits May 19, 2026 14:44

Swapped positions of common and matmul overrides

323c0b2

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Added matmul and integration tests

448e30b

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Only TRITON_ATTN works

5cf2167

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Minor test improvements

fc49b26

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Further test improvements

5b5ba56

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Test fix

bb50331

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Merge branch 'main' into xpu_rl

7e5dae1

tzielinski-habana marked this pull request as ready for review May 21, 2026 11:48

tzielinski-habana requested a review from yewentao256 as a code owner May 21, 2026 11:48

Merge branch 'main' into xpu_rl

e3c7a9a

yewentao256 reviewed May 21, 2026

View reviewed changes

Tests refactoring

75fded1

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

mergify Bot removed the needs-rebase label Jun 24, 2026

tzielinski-habana added 4 commits June 24, 2026 14:05

Merge branch 'main' into xpu_rl

3a3d5b6

Merge branch 'main' into xpu_rl

469c601

Merge branch 'main' into xpu_rl

f576920

Merge branch 'main' into xpu_rl

36869b7

tzielinski-habana added 6 commits June 24, 2026 17:16

Merge branch 'main' into xpu_rl

1fa657e

Merge branch 'main' into xpu_rl

16ca2c7

Merge branch 'main' into xpu_rl

c123747

Merge branch 'main' into xpu_rl

a53f61b

Merge branch 'main' into xpu_rl

b032b34

Merge branch 'main' into xpu_rl

194750a

yewentao256 reviewed Jun 26, 2026

View reviewed changes

tzielinski-habana and others added 2 commits July 6, 2026 10:35

Merge branch 'main' into xpu_rl

9d9c9d9

Refactored devices and backends handling in tests

3c52f32

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

tzielinski-habana commented Jul 7, 2026

View reviewed changes

Comment thread tests/v1/determinism/utils.py

yewentao256 reviewed Jul 7, 2026

View reviewed changes

Comment thread tests/v1/determinism/test_batch_invariance.py Outdated

Comment thread tests/v1/determinism/test_batch_invariance.py Outdated

tzielinski-habana and others added 4 commits July 9, 2026 19:39

Removed skip_unsupported_backend

e04f0f2

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>

Merge branch 'main' into xpu_rl

beccb29

Merge branch 'main' into xpu_rl

458a448

Merge branch 'main' into xpu_rl

94ce14c

yewentao256 reviewed Jul 14, 2026

View reviewed changes

yewentao256 approved these changes Jul 15, 2026

View reviewed changes

yewentao256 merged commit 61141ed into vllm-project:main Jul 15, 2026
93 checks passed

tzielinski-habana mentioned this pull request Jul 20, 2026

[Hardware][XPU] Register matmul and linear batch-invariant kernels for XPU #49209

Merged

yewentao256 mentioned this pull request Aug 22, 2026

[Feature]: Batch Invariant Feature and Performance Optimization #27433

Open

29 tasks

Uh oh!

Uh oh!

Conversation

tzielinski-habana commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Note

Uh oh!

mergify Bot commented May 7, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

mergify Bot commented May 14, 2026

Uh oh!

yewentao256 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Jun 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yewentao256 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yewentao256 left a comment

Choose a reason for hiding this comment

Uh oh!

tzielinski-habana commented Jul 15, 2026

Uh oh!

yewentao256 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

tzielinski-habana commented May 7, 2026 •

edited

Loading