Phase-split (prefill/decode) quantisation for Llama and Qwen3 - #323
Open
Shreyas8612 wants to merge 25 commits into
Open
Shreyas8612 wants to merge 25 commits into
Shreyas8612 wants to merge 25 commits into
Conversation
added 25 commits
July 6, 2026 01:03
Block scale becomes floor(log2(max)) - emax_element for MXFP and ceil(log2(max/qmax)) for MXINT, so the block maximum sits at the top of the element range instead of clipping or flushing small elements to zero. MX blocks are formed per row with zero padding, so a block can no longer straddle a row boundary. Adds zero-block protection and the full E8M0 signed range. Replaces quantile(1.0) with amax and drops a dead encode/decode round trip, bit-identical and 1.5-1.6x faster.
RMSNorm, softmax, SiLU with its gated product, RoPE and the residual adds now round at the operation boundary rather than quantising only the input, which makes the vector-format setting a real parameter. Adds the Qwen3 decoder layer and per-head Q/K normalisation before RoPE.
Two symbols removed from the transform package were still imported by two parents, which broke collection for the whole suite. Decode routing now rejects the inconsistent hardware MXINT path, which is retained only as a compatibility adapter.
A stray print of Q.shape fired on every call to random_hadamard_matrix, flooding stdout during rotation search. Restores the file to upstream.
Llama decoder layers had no quantised counterpart, so whole-layer replacement silently skipped them while Qwen3 layers were replaced normally. Adds the wrapper and registers it end to end: exported from the llama module package, entered in quantized_llama_module_map as llama_decoder_layer_minifloat, and mapped in both llama_prefix_map and from_self_prefix_map so the modify helper resolves LlamaDecoderLayer to the llama_decoder_layer prefix.
The four Llama attention variants each hardcoded a family-specific KV quantiser (kv_cache_mxfp, kv_cache_mxint, and their rotate variants), so a profile mixing element families with the KV format took whichever quantiser its attention class was built around rather than the one the profile requested. Twelve of thirty-six profiles failed this way. All four sites now call the format-dispatching kv_cache_mx, with the rotate variants passing rotate=self.kv_cache_use_rotate instead of selecting a separate function.
Rounding in fake-quantised decode is launch-bound: each vector, matrix and MX call expands into a chain of tiny elementwise kernels. The pure rounding functions are now compiled once per distinct constant tuple (bit widths, block size, block dim, rounding mode) so shapes stay dynamic and no recompile fires per generated token. MASE_COMPILE_ROUNDING=0 restores the eager path, and any dynamo failure falls back to eager.
The optional torch.compile fusion of the rounding primitives was enabled by default (MASE_COMPILE_ROUNDING had to be set to 0 to turn it off), and importing chop.nn.quantizers._compile mutated global torch._dynamo.config (cache limits, suppress_errors) as an import side effect. Its module docstring also pointed at a verification script that does not live in this repository. Flip the default: fusion is now enabled only when MASE_COMPILE_ROUNDING=1 is set before chop is imported; with the variable unset or any other value maybe_compile() returns the function it was given and compiled_variant() returns the builder's function as-is, so the eager path is exactly the code that ran before. Move the dynamo configuration into a helper that runs once, from the first real torch.compile call, so importing the package no longer touches torch._dynamo. Reword the docstring so it stands on its own. Add test/nn/quantizers/test_compile_rounding.py, which runs a fixed workload in fresh interpreters with the variable unset, "0", "yes" and "1" and checks that (a) the eager settings leave _dynamo.config untouched and cache no variants and (b) every setting, including the compiled one, produces tensors bit-identical to the eager implementations for the MXFP, MXINT, minifloat, accumulator and vector-format rounding paths. Verified with pytest on test/nn/quantizers/test_compile_rounding.py, test/passes/module/transforms/quantize/test_decode_phase_quantize.py and test/passes/module/transforms/quantize/test_quantize_module.py.
There was a problem hiding this comment.
🟡 Changes recommended
Decoder output contracts, concurrent phase overrides, legacy minifloat defaults, and validator KV handoff behavior contain unresolved correctness issues.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds phase-aware prefill/decode quantisation across Llama and Qwen3, including separate weight banks, KV-cache formats, GPTQ/rotation integration, and MX/vector rounding support.
Changes:
- Adds normalized phase configuration and runtime dispatch.
- Extends quantized modules, GPTQ, rotation search, and MX quantizers.
- Adds validation and CPU test coverage.
File summaries
| File | Description |
|---|---|
src/chop/nn/quantized/modules/phase_config.py |
Normalizes phase-specific configurations and KV handoff rules. |
src/chop/nn/quantized/modules/phase_context.py |
Tracks runtime phase during model execution. |
src/chop/nn/quantized/modules/linear.py |
Implements separate prefill/decode weight banks. |
src/chop/nn/quantized/modules/llama/attention.py |
Adds phase-aware Llama attention. |
src/chop/nn/quantized/modules/llama/decoder_layer.py |
Adds phase-aware Llama decoder replacement. |
src/chop/nn/quantized/modules/llama/mlp.py |
Adds phase-aware Llama MLP behavior. |
src/chop/nn/quantized/modules/llama/rms_norm.py |
Adds phase-aware RMSNorm policies. |
src/chop/nn/quantized/modules/qwen3/attention.py |
Adds phase-aware Qwen3 attention. |
src/chop/nn/quantized/modules/qwen3/decoder_layer.py |
Adds phase-aware Qwen3 decoder replacement. |
src/chop/nn/quantized/modules/qwen3/mlp.py |
Adds phase-aware Qwen3 MLP behavior. |
src/chop/nn/quantized/functional/{attention,kvcache,matrix,rope,silu,softmax,vector}.py |
Adds vector rounding, matrix boundaries, and format-dispatched operations. |
src/chop/nn/quantizers/{mxint,mxfp,_minifloat_mx,rotation}/ |
Updates MX quantization and rotation support. |
src/chop/passes/module/module_modify_helper.py |
Preserves phase banks during replacement. |
src/chop/passes/module/transforms/gptq/{checkpoint,run}.py |
Adds phase-targeted GPTQ persistence and restoration. |
src/chop/passes/module/transforms/quantize/{quantize,rotation_search}.py |
Integrates phase validation and rotation search. |
scripts/validate_decode_phase_quant.py |
Adds end-to-end validation workflow. |
test/passes/module/transforms/quantize/test_decode_phase_quantize.py |
Tests phase configuration, banks, GPTQ, and model behavior. |
test/nn/quantizers/test_compile_rounding.py |
Tests eager and compiled rounding paths. |
Review details
- Files reviewed: 49/50 changed files
- Comments generated: 6
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+199
to
+205
| "config": { | ||
| "name": "mxint", | ||
| # kv_cache_handoff defaults to "decode_format": prefill KV | ||
| # is quantised on write into the decode chip's HBM format. | ||
| "prefill": {"bypass": True}, | ||
| "decode": attn_decode, | ||
| } |
Comment on lines
+101
to
+102
| is_finite=cfg.get("data_in_is_finite", False), | ||
| round_mode=cfg.get("data_in_round_mode", "rn"), |
Comment on lines
+66
to
+70
| return ( | ||
| policy.residual_add(residual, hidden_states) | ||
| if policy.enabled | ||
| else residual + hidden_states | ||
| ) |
| # Process-wide (not a ContextVar): an evaluation-time override must also be | ||
| # visible to forwards running in worker threads (e.g. nn.DataParallel), which | ||
| # do not inherit the caller's ContextVar state. | ||
| _PHASE_OVERRIDE: Phase | None = None |
Comment on lines
+66
to
+70
| return ( | ||
| policy.residual_add(residual, hidden_states) | ||
| if policy.enabled | ||
| else residual + hidden_states | ||
| ) |
Comment on lines
+44
to
+46
| if name == "mxint_hardware" and "decode" in config: | ||
| raise ValueError( | ||
| "decode quantization must select 'mxint', which uses " |
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.
Adds phase-aware post-training quantisation: a config layer (phase_config.py,
phase_context.py) that lets prefill and decode carry different weight,
activation and KV formats; separate prefill/decode weight banks preserved
across module replacement; phase-aware Llama/Qwen3 attention, MLP, RMSNorm
and decoder layers; GPTQ and rotation-search phase targeting; rounding at
every vector-op boundary in decode with format-dispatched KV-cache
quantisation; and OCP-conformant MX scaling with row-local blocking.
An optional torch.compile fusion of the rounding kernels is included but off
by default: set MASE_COMPILE_ROUNDING=1 before importing chop to enable it.
Importing the package no longer touches torch._dynamo.config.
Testing: test/passes/module/transforms/quantize/test_decode_phase_quantize.py
(44 tests) plus test_quantize_module.py pass on CPU;
test/nn/quantizers/test_compile_rounding.py checks the compiled and eager
paths are bit-identical and that the eager default has no side effects.
pyproject.toml and uv.lock are unchanged.
Provenance: phase_config.py and phase_context.py extend the versions
Yuxuan Han opened in #320 (all of their functions are kept; roughly half the
lines are verbatim). The atomic-write hunk in gptq/checkpoint.py (commit
3e25f10) is identical to Yuxuan Han's commit bb2d147 on
yx/qwen3-moe-gptq. I suggest #320 and this PR be reconciled so only one of
them lands on main.