dynamic_quant: 20% improvements in per-channel mode - #1512
Open
osavchenkox wants to merge 1 commit into
Open
Conversation
osavchenkox
requested review from
PatrykWo,
adobrzyn,
afierka-intel,
iboiko-habana,
jbyczkow,
kamil-kaczor,
ksmusz,
mgawarkiewicz-intel,
michalkuligowski and
xuechendi
as code owners
June 1, 2026 06:02
osavchenkox
had a problem deploying
to
pre-merge-approval
June 1, 2026 06:02 — with 
GitHub Actions
Error
GitHub Actions
Error
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 improves FP8 dynamic quantization performance on HPU by switching to amax(..., keepdim=True) and adds a local benchmark script (with trace/FX dump artifacts ignored by git).
Changes:
- Optimize
dynamic_quantto avoidmax(...).values+unsqueezeand enable better HPU fusion undertorch.compile. - Add an HPU benchmark script to compare BF16 vs multiple FP8 quantization schemes and dump profiler/FX artifacts.
- Ignore benchmark-generated trace directories in
.gitignore.
Reviewed changes
Copilot reviewed 2 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| vllm_gaudi/extension/ops.py | Uses amax(keepdim=True) in FP8 dynamic quant to reduce overhead and improve fusion. |
| tests/unit_tests/ops/bench_fp8_dynamic_quant.py | Adds a standalone HPU benchmark/profiling script for FP8 quantization variants. |
| .gitignore | Ignores benchmark trace/FX dump output directories. |
Comment on lines
+29
to
+39
| import os | ||
|
|
||
| # Disable torch's autoload of habana_frameworks (we import it ourselves below). | ||
| os.environ.setdefault("TORCH_DEVICE_BACKEND_AUTOLOAD", "0") | ||
| # Required for HPU profiler activity to be recorded into the chrome trace. | ||
| os.environ.setdefault("HABANA_PROFILE", "1") | ||
|
|
||
| import sys | ||
| import importlib.abc | ||
| import importlib.machinery | ||
| import torch |
Comment on lines
+73
to
+75
| sys.meta_path.insert(0, _DisableHabanaConfigPatch()) | ||
|
|
||
| import habana_frameworks.torch # noqa: E402, F401 registers torch.hpu / hpu_backend |
Comment on lines
+215
to
+228
| def profile_variant(name: str, fn, args, trace_dir: str): | ||
| for _ in range(WARMUP): | ||
| fn(*args) | ||
| torch.hpu.synchronize() | ||
|
|
||
| trace_path = os.path.join(trace_dir, f"{name}.json") | ||
| with torch.profiler.profile( | ||
| activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.HPU], | ||
| record_shapes=True, | ||
| with_stack=False, | ||
| ) as prof: | ||
| for _ in range(ITERS): | ||
| out = fn(*args) | ||
| torch.hpu.synchronize() |
Comment on lines
+936
to
+943
| # amax(keepdim=True) is preferred over max(dim=-1).values + unsqueeze: | ||
| # - aten.max.dim returns (values, indices); the indices tensor is | ||
| # computed and immediately discarded, costing an extra reduce pass | ||
| # and an i32->i64 cast on HPU. | ||
| # - keepdim=True folds the unsqueeze into the reduce kernel, allowing | ||
| # hpu_backend to fuse abs+amax into a single TPC kernel under | ||
| # torch.compile (~20% device-time win on the quant step). | ||
| scale = (torch.abs(data).amax(dim=-1, keepdim=True) + 1e-8) / FP8_MAX |
osavchenkox
force-pushed
the
dev/osavchenko/fp8_amax_quant
branch
from
June 1, 2026 12:26
b6dc287 to
0e94fd4
Compare
osavchenkox
had a problem deploying
to
pre-merge-approval
June 1, 2026 12:26 — with
GitHub Actions
Error
GitHub Actions
Error
Contributor
Author
|
Friendly bump — this PR is ready for review. It targets per-channel |
Replace max(dim=-1).values + unsqueeze with amax(dim=-1, keepdim=True). aten.max.dim builds a discarded indices tensor; amax skips it, and keepdim=True folds the unsqueeze into the reduce kernel. Add standalone bench that compares bf16/per-tensor/per-channel/static FP8 with eager + torch.compile traces and FX graph dumps. Signed-off-by: OlegX Savchenko <olegx.savchenko@intel.com>
osavchenkox
force-pushed
the
dev/osavchenko/fp8_amax_quant
branch
from
June 3, 2026 09:14
0e94fd4 to
513ae3a
Compare
osavchenkox
had a problem deploying
to
pre-merge-approval
June 3, 2026 09:14 — with
GitHub Actions
Error
GitHub Actions
Error
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 max(dim=-1).values + unsqueeze with amax(dim=-1, keepdim=True). aten.max.dim builds a discarded indices tensor; amax skips it, and keepdim=True folds the unsqueeze into the reduce kernel.
Add standalone bench that compares bf16/per-tensor/per-channel/static FP8 with eager + torch.compile traces and FX graph dumps.