[DO_NOT_MERGE] Use l2Norm TPC kernel for HPU compile - #1457
Open
osavchenkox wants to merge 1 commit into
Open
Conversation
Contributor
🚧 CI BlockedThe main CI workflow was not started for the following reason:
|
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR refactors Q/K L2-normalization to use an HPU TPC (torch.ops.hpu.l2_norm) implementation and adjusts Q/K dtype handling in the recurrent gated delta rule path.
Changes:
- Introduce TPC- and PyTorch-based L2-norm helpers and route
_l2norm_last_dimto the HPU op. - Update
_preprocess_qk_l2normto normalize without pre-casting inputs to fp32. - Change recurrent path to avoid unconditional casting of Q/K unless L2-norm is disabled.
Comment on lines
+54
to
56
| # @torch._dynamo.disable | ||
| def _preprocess_qk_l2norm(q, k): | ||
| """L2norm in eager mode — HPU torch.compile miscompiles l2norm.""" |
| @torch._dynamo.disable | ||
| def _l2norm_tpc_last_dim(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: | ||
| # input bf16, output fp32 to avoid precision issues | ||
| return torch.ops.hpu.l2_norm(x, epsilon=eps) |
Comment on lines
624
to
+627
| # Flatten token axis. | ||
| # Compute dtype controlled by VLLM_GDN_COMPUTE_FP32 env var (default: bf16) | ||
| qf = q.reshape(-1, H, Kdim).to(_GDN_COMPUTE_DTYPE) | ||
| kf = k.reshape(-1, H, Kdim).to(_GDN_COMPUTE_DTYPE) | ||
| qf = q.reshape(-1, H, Kdim) | ||
| kf = k.reshape(-1, H, Kdim) |
Comment on lines
+633
to
+634
| qf = _l2norm_last_dim(qf) # output is fp32 | ||
| kf = _l2norm_last_dim(kf) # output is fp32 |
|
|
||
| def _l2norm_pt_last_dim(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: | ||
| x = x.to(torch.float32) | ||
| return (x / torch.sqrt(torch.sum(x * x, dim=-1, keepdim=True) + eps)).to(torch.float32) |
osavchenkox
force-pushed
the
dev/osavchenko/l2_norm
branch
from
May 19, 2026 08:35
fe99503 to
f9b2676
Compare
osavchenkox
marked this pull request as ready for review
May 19, 2026 08:36
osavchenkox
requested review from
PatrykWo,
adobrzyn,
afierka-intel,
iboiko-habana,
jbyczkow,
kamil-kaczor,
ksmusz,
mgawarkiewicz-intel,
michalkuligowski and
xuechendi
as code owners
May 19, 2026 08:36
Contributor
🚧 CI BlockedThe main CI workflow was not started for the following reason:
|
osavchenkox
changed the title
Replace eager _l2norm_last_dim with TPC l2_norm op, remove @torch._dynamo.disable to allow torch.compile, and fix dtype handling so l2norm outputs fp32 directly. Signed-off-by: OlegX Savchenko <olegx.savchenko@intel.com>
osavchenkox
force-pushed
the
dev/osavchenko/l2_norm
branch
from
June 2, 2026 17:26
f9b2676 to
ccae01f
Compare
osavchenkox
had a problem deploying
to
pre-merge-approval
June 2, 2026 17:26 — with 
GitHub Actions
Error
GitHub Actions
Error
osavchenkox
changed the title
Contributor
Author
|
Rebased on latest |
osavchenkox
changed the title
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.
Replace eager _l2norm_last_dim with TPC l2_norm op, remove @torch._dynamo.disable to allow torch.compile, and fix dtype handling so l2norm outputs fp32 directly.