[Hardware][XPU] Register matmul and linear batch-invariant kernels for XPU - #49209
Conversation
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
tzielinski-habana
changed the title
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
|
Documentation preview: https://vllm--49209.org.readthedocs.build/en/49209/ |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
Thanks for the work, please also test with e2e performance with / without batch invariance
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Running this benchmark: I get a ~4.4x slowdown with batch invariance. When I run a kernel-level benchmark for matmul related kernels, the average slowdowns are as follows: I'm attaching a more detailed kernel-level info in a txt file As for RL, for which we need BI, rollouts from vLLM are only a small part of the whole training process, so the slowdowns there are less impactful. When I checked on one RL benchmark, I got only a ~20% slowdown. |
| def _matmul_backward_xpu(grad, self, other, mask): | ||
| """XPU implementation of aten::matmul_backward.""" | ||
| grad_self = ( | ||
| matmul_batch_invariant(grad, other.transpose(-1, -2)) if mask[0] else None | ||
| ) | ||
| grad_other = ( | ||
| matmul_batch_invariant(self.transpose(-1, -2), grad) if mask[1] else None | ||
| ) | ||
| return grad_self, grad_other |
There was a problem hiding this comment.
Why do we need backward inside vLLM? I don't think it is a good idea
There was a problem hiding this comment.
fair point
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
| # Tensor descriptors require contiguous row-major layout. | ||
| a = a.contiguous() | ||
| # Descriptor kernel expects B in [N, K] layout (K-contiguous). | ||
| b_t = b.t().contiguous() |
There was a problem hiding this comment.
Would contiguous be safe enough? Eg. K=60、N=17?
There was a problem hiding this comment.
As far as I know it's safe. With strange dimensions we would still be correct, but slower.
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
LGTM, only one nit left, thanks for the iteration!
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #85014. |
|
@hmellor will you find some time for review, please? It looks like I need your approval as well |
|
This pull request has merge conflicts that must be resolved before it can be |
|
Some new PRs entered and reorganized the code quite significantly. I solved merge conflicts, but I'll have to reorganize the matmul config bit in my PR to match the new code. I'll probably have it ready tomorrow. |
_get_descriptor_matmul_config now sits next to _get_matmul_config in batch_invariant_configs.py, which became the home for persistent-matmul config selection in vllm-project#53247. Pure move: the returned configs are identical for every M and dtype. Also record why BLOCK_SIZE_K is derived from dtype alone: an M-dependent BLOCK_K would change the K-reduction order and break batch invariance. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
|
/ci run |
|
✅ Triggered Buildkite CI #86379 for commit |
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com>
|
/ci run |
|
✅ Triggered Buildkite CI #86600 for commit |
|
/ci retry |
|
✅ No failed, timed-out, or expired jobs need retrying: https://buildkite.com/vllm/ci/builds/86600 |
|
/ci run |
|
✅ Triggered Buildkite CI #86863 for commit |
…r XPU (vllm-project#49209) Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
…r XPU (vllm-project#49209) Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
3 participants
Purpose
This is the second (2/2) PR introducing batch invariance to Intel XPU devices. You can find the first one here.
This PR registers the missing matmul and linear kernels. It also adds a new matmul_kernel_descriptor_persistent Triton kernel for better matmul performance on XPU.
Test Plan
Tested with unit tests covering batch invariance on XPU.
Test Result
All tests pass.
Note
The code was co-developed with GitHub Copilot.
Details
Matmul and linear kernels were skipped in the previous PR, because they required more work. I tested batch invariance in Reinforcement Learning context and there I hit an error, because backward kernels were not registered for the "XPU" dispatch key. Also, I wanted to introduce a new Triton matmul kernel based on Tensor Derscriptors, because that implementation has shown ~3x speedup on XPU vs. the one based on pointers.