Fixing redundant transpose in HPUMambaMixer2 (#999)#1014
Open
ksmusz wants to merge 3 commits intoreleases/v0.15.1from
Open
Fixing redundant transpose in HPUMambaMixer2 (#999)#1014ksmusz wants to merge 3 commits intoreleases/v0.15.1from
ksmusz wants to merge 3 commits intoreleases/v0.15.1from
Conversation
We are noticing some MME transposes within the profiling, which are fixed by this change --------- Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
ksmusz
requested review from
mgawarkiewicz-intel,
piotrbocian and
wpyszka
as code owners
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes redundant transpose operations in the HPUMambaMixer2 implementation that were causing inefficient MME (Matrix Multiply Engine) operations during profiling. The fix involves creating a custom state shape calculator for HPU that eliminates the need for transpose operations.
Changes:
- Removed redundant transpose operation when accessing conv_state from the KV cache
- Added custom
hpu_mamba2_state_shapefunction to return conv_state in the correct layout - Implemented
get_state_shapemethod in HPUMambaMixer2 to use the custom state shape calculator
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| vllm_gaudi/ops/hpu_mamba_utils.py | New utility file implementing hpu_mamba2_state_shape function that returns conv_state with contiguous 'dim' axis layout |
| vllm_gaudi/ops/hpu_mamba_mixer2.py | Removed transpose operation on conv_state access and added get_state_shape method using the new HPU-specific state shape calculator |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
michalkuligowski
approved these changes
Feb 24, 2026
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.
We are noticing some MME transposes within the profiling, which are fixed by this change
cherry-pick of #999