[Perf] Add pluggable backend dispatch for quantization lifecycle ops - #773
[Perf] Add pluggable backend dispatch for quantization lifecycle ops#773ishrith-gowda wants to merge 2 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughIntroduces a pluggable quantization backend abstraction ( ChangesQuantization backend abstraction
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant ForwardHelpers as "_quantize / _dequantize / _quantize_dequantize"
participant Dispatcher as get_quantization_backend
participant Backend as Active Backend
Caller->>ForwardHelpers: quantize/dequantize request (x, scale, zero_point, args)
ForwardHelpers->>Dispatcher: get_quantization_backend(x, args)
Dispatcher->>Backend: is_available(x, args)
alt backend available
Dispatcher-->>ForwardHelpers: return Backend
else backend unavailable
Dispatcher-->>ForwardHelpers: return EagerQuantizationBackend
end
ForwardHelpers->>Backend: quantize/dequantize/quantize_dequantize(...)
Backend-->>ForwardHelpers: result tensor
ForwardHelpers-->>Caller: result tensor
Related issues: None specified. Related PRs: None specified. Suggested labels: enhancement, refactor Suggested reviewers: None specified. 🐰 A backend hopped in, plug-and-play, 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
The quality checks have failed. Please run |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/test_quantization/lifecycle/test_backend.py (2)
84-114: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise custom dispatch for all backend leaf ops.
The custom backend test only proves
_quantize_dequantizedispatch. Add sentinelquantizeanddequantizemethods too, so regressions in the other changed helper paths are caught.As per path instructions,
tests/**/*.pyshould be comprehensive for quantization scenarios.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_quantization/lifecycle/test_backend.py` around lines 84 - 114, The custom backend test only covers _quantize_dequantize, so it can miss regressions in the other backend leaf helpers. Extend SentinelBackend in test_custom_backend_dispatch_and_fallback to also define sentinel quantize and dequantize methods, then add assertions that the backend-specific paths for those helpers dispatch to the sentinels and fall back to the eager/reference behavior when availability is toggled off. Use the existing QuantizationBackend.register, set_quantization_backend, and _reference_qdq-style patterns to keep the test aligned with the other helper paths.Source: Path instructions
32-44: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover
global_scalein the reference and roundtrip tests.The dispatch paths forward
global_scale, but every assertion currently exercises onlyNone, so a regression in that branch would pass.Proposed test expansion
-def _reference_qdq(x, scale, zero_point, q_min, q_max, args): +def _reference_qdq(x, scale, zero_point, q_min, q_max, args, global_scale=None): from compressed_tensors.quantization.quant_args import round_to_quantized_type_args + if global_scale is not None: + scale = scale / global_scale scaled = x / scale @@ -def test_eager_qdq_matches_reference(): +@pytest.mark.parametrize("global_scale", [None, torch.tensor(2.0)]) +def test_eager_qdq_matches_reference(global_scale): @@ - got = _quantize_dequantize(x, scale, zp, q_min, q_max, args) - ref = _reference_qdq(x, scale, zp, q_min, q_max, args) + got = _quantize_dequantize( + x, scale, zp, q_min, q_max, args, global_scale=global_scale + ) + ref = _reference_qdq( + x, scale, zp, q_min, q_max, args, global_scale=global_scale + ) assert torch.equal(got, ref) @@ -def test_quantize_then_dequantize_roundtrip(): +@pytest.mark.parametrize("global_scale", [None, torch.tensor(2.0)]) +def test_quantize_then_dequantize_roundtrip(global_scale): @@ - q = _quantize(x, scale, zp, q_min, q_max, args) - dq = _dequantize(q, scale, zp) + q = _quantize(x, scale, zp, q_min, q_max, args, global_scale=global_scale) + dq = _dequantize(q, scale, zp, global_scale=global_scale) @@ - assert torch.allclose(dq, _quantize_dequantize(x, scale, zp, q_min, q_max, args)) + assert torch.allclose( + dq, + _quantize_dequantize( + x, scale, zp, q_min, q_max, args, global_scale=global_scale + ), + )As per path instructions,
tests/**/*.pyshould cover edge cases for compression and quantization scenarios.Also applies to: 52-76
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_quantization/lifecycle/test_backend.py` around lines 32 - 44, Expand the quantization lifecycle tests to explicitly exercise the global_scale branch in both the reference helper _reference_qdq and the roundtrip assertions, since the current cases only cover None and can miss regressions in dispatch paths that forward global_scale. Add coverage in test_backend.py around the existing reference/roundtrip test helpers so they validate behavior when global_scale is provided, keeping the checks aligned with the existing quantization flow and symbols like _reference_qdq and the related roundtrip test cases.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/compressed_tensors/quantization/lifecycle/backend.py`:
- Around line 50-53: The backend availability probe in `Backend.is_available`
currently assumes `args` is always a `QuantizationArgs`, but `_dequantize` may
call it without one, causing custom backends to fail before eager fallback.
Update the `is_available` contract and all backend implementations/call sites to
allow `args=None`, and make the availability checks in the quantization
lifecycle code handle a missing args value safely without raising.
- Around line 22-34: Apply the configured import/export sorting in the backend
module so Ruff passes: separate the first-party imports from the
standard/library import and reorder them according to the project’s import
grouping rules. Also sort the entries in the module’s __all__ list in a
consistent order, keeping the exported symbols in the expected alphabetical
sequence. Use the existing import block and __all__ definition in the
QuantizationBackend module as the targets for the fix.
In `@tests/test_quantization/lifecycle/test_backend.py`:
- Around line 79-81: The test for set_quantization_backend is asserting too
broadly with pytest.raises(Exception) even though
QuantizationBackend.get_value_from_registry() raises KeyError for unknown names.
Update test_set_unknown_backend_raises to expect KeyError specifically so the
test matches the actual registry validation behavior.
---
Nitpick comments:
In `@tests/test_quantization/lifecycle/test_backend.py`:
- Around line 84-114: The custom backend test only covers _quantize_dequantize,
so it can miss regressions in the other backend leaf helpers. Extend
SentinelBackend in test_custom_backend_dispatch_and_fallback to also define
sentinel quantize and dequantize methods, then add assertions that the
backend-specific paths for those helpers dispatch to the sentinels and fall back
to the eager/reference behavior when availability is toggled off. Use the
existing QuantizationBackend.register, set_quantization_backend, and
_reference_qdq-style patterns to keep the test aligned with the other helper
paths.
- Around line 32-44: Expand the quantization lifecycle tests to explicitly
exercise the global_scale branch in both the reference helper _reference_qdq and
the roundtrip assertions, since the current cases only cover None and can miss
regressions in dispatch paths that forward global_scale. Add coverage in
test_backend.py around the existing reference/roundtrip test helpers so they
validate behavior when global_scale is provided, keeping the checks aligned with
the existing quantization flow and symbols like _reference_qdq and the related
roundtrip test cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 57071d61-bb37-42d8-a37f-aa1f74546c16
📒 Files selected for processing (3)
src/compressed_tensors/quantization/lifecycle/backend.pysrc/compressed_tensors/quantization/lifecycle/forward_helpers.pytests/test_quantization/lifecycle/test_backend.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
vllm-project/llm-compressor(manual)
The elementwise lifecycle leaf ops (_quantize, _dequantize, _quantize_dequantize) run on every weight/activation group during calibration. Profiling a W4A16 GPTQ run of Qwen2.5-0.5B shows _quantize_dequantize is called 245,760 times for a 0.5B model (billions at Kimi/DeepSeek scale), spending 11.6s of self time in a chain of small elementwise CUDA kernels. Add a QuantizationBackend registry (RegistryMixin) so these ops can dispatch to different backends. The reference PyTorch code moves into an eager backend (registered by default, bit identical to before); the leaf helpers become thin wrappers that route through the active backend; accelerated backends (torch.compile, triton kernels) register under a name and are selected per tensor with transparent eager fallback. Behavior is unchanged by default; this is the plug in point for the kernels tracked in the rest of the issue. Towards vllm-project#766 Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: ishrith-gowda <ishrithgowda@berkeley.edu>
e526fe4 to
1b9d926
Compare
brian-dellabetta
left a comment
brian-dellabetta
left a comment
There was a problem hiding this comment.
Thanks for preparing this, I understand the design. a couple questions.
| # clamp and round | ||
| quantized_value = round_to_quantized_type_args( | ||
| tensor=scaled, args=args, min=q_min, max=q_max | ||
| return get_quantization_backend(x, args).quantize( |
There was a problem hiding this comment.
since these are called potentially billions of times, can the quant backend resolution occur elsewhere a single time?
There was a problem hiding this comment.
Good call. Done in 8e3282c: the backend class is now resolved once and cached in a module global (set_quantization_backend updates the cache, first call resolves the env-configured name lazily). The hot path is now one global read plus an identity check, and the per-tensor is_available probe only runs when a non-eager backend is active. With eager active (the default) there is no registry access per call at all.
| ] | ||
|
|
||
| _DEFAULT_BACKEND = "eager" | ||
| _ACTIVE_BACKEND = os.environ.get("COMPRESSED_TENSORS_QUANT_BACKEND", _DEFAULT_BACKEND) |
There was a problem hiding this comment.
in practice, will the most suitable backend be inferred based on version / device available? Would we ever want quant backend to be different for different modules in the LLM?
There was a problem hiding this comment.
The per-tensor is_available hook already gives device-level adaptivity: a CUDA-only backend runs for GPU modules while CPU modules transparently fall back to eager, so mixed-device models do the right thing without per-module config. For auto-inference, my thinking is a follow-up "auto" mode that walks registered backends in priority order and picks the first whose is_available passes (version and device checks live in each backend). Explicit per-module backend selection could ride on QuantizationScheme later if a concrete kernel needs it, but I would hold off until the triton kernels land and we see whether they want different dispatch per scheme. Happy to adjust if you have a preference.
…able Per review feedback: - Resolve the active backend once and cache the class; the hot path now costs one global read instead of a registry lookup per call. The per-tensor is_available fallback only runs for non-eager backends. - is_available accepts args=None (dequantize carries no QuantizationArgs). - Narrow unknown-backend test assertion to KeyError. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: ishrith-gowda <ishrithgowda@berkeley.edu>
|
This pull request has merge conflicts that must be resolved before it can be |
|
Hi @ishrith-gowda! Let me know what you think of #804. I think this interface should be a little more generalizable to backending functions beyond just quantization |
|
This looks better than what I had. Generalizing to any function with One thing worth checking before this goes on the QDQ leaf: Happy to close #773 in favor of this and port over the eager path equivalence tests and the QDQ benchmark I used for the profiling numbers, if that is useful. The merge conflict on mine is moot in that case. |
@ishrith-gowda I think your
Would you be interested in opening a PR rebased on #804 to add this hot-path caching? |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
3 participants
Purpose
Towards #766, the "start by replacing the functions with dispatchers" step (design discussed in #766 (comment)).
The elementwise lifecycle leaf ops (
_quantize,_dequantize,_quantize_dequantize) run on every weight/activation group during calibration. Profiling a W4A16 GPTQ run ofQwen2.5-0.5B-Instruct(cProfile) shows_quantize_dequantizeis called 245,760 times for a 0.5B model (billions at Kimi/DeepSeek scale), spending 11.6s of self time in a chain of small elementwise CUDA kernels:This adds a
QuantizationBackendregistry: the reference PyTorch code moves into aneagerbackend (registered by default, bit identical to today), the leaf helpers become thin wrappers that route through the active backend, and accelerated backends (torch.compile, hand written triton kernels) can register under a name and are selected per tensor with transparent eager fallback. No behavior change by default; this is the plug in point for the kernels tracked in the rest of #766.Changes
quantization/lifecycle/backend.py(new):QuantizationBackend(RegistryMixin), the defaultEagerQuantizationBackend(exact original ops),get_quantization_backend()selection with eager fallback,set_quantization_backend()/COMPRESSED_TENSORS_QUANT_BACKENDenv.quantization/lifecycle/forward_helpers.py: the three leaf helpers route through the active backend; signatures and call sites unchanged.Test Plan / Test Result
pytest tests/test_quantization/lifecycle/test_backend.py(new): 5 passed. Eager output is bit identical to the reference formula; quantize/dequantize roundtrip; unknown backend raises; a custom registered backend is dispatched to, andis_available=Falsefalls back to eager.pytest tests/test_quantization/lifecycle/regression: every test that passes onmainstill passes with the dispatcher in place (2 unrelated failures are pre-existing onmain), confirming the eager path is unchanged.Built with an AI-assisted engineering workflow (Claude) under my direction: I specified the design, reviewed each iteration, and validated the profiling and test results. Co-authored-by trailer on the commit.