Generalize Gemma 4 MoE Quark path to plain Quark/AWQ int4 experts

[Thiago Pereira Rocha (thpereir)]

Summary

Generalizes the Gemma 4 MoE make_moe_quark path so it also builds plain Quark/AWQ int4 experts (e.g. amd/gemma-4-26B-A4B-it-uint4-awq), not just the factored 2-bit LoRA-KD checkpoint. The previous code was hardwired for the 2-bit layout; this guards each 2-bit-specific step on the checkpoint's actual shape.

• Input transform: emit the shared prescale+rotation only when input_prescale is present. Plain Quark/AWQ experts have none, so the MoE input feeds QMoE directly.
• fc1 layout: gate the concat->interleaved fc1 row reorder on the 2-bit path (expert_weight_bits==2). AWQ int4 experts are already emitted interleaved by combine_and_repack_gate_up.
• Zero-points: emit at native dtype unless float. 2-bit carries FLOAT zero-points (cast to io_dtype); int4 carries INTEGER (uint8) zero-points (emitted without cast). The CUDA zp-omit path stays 2-bit-only.

Stacking

This PR is stacked on top of the 2-bit PR (#2482), which is itself stacked on the base Gemma 4 PR (#2473). Until those merge, the diff here shows the parent commits; the int4-specific change is the top commit. Will rebase onto main as parents land.

Test plan

• test/python/models/test_gemma4moe_builder.py - added TestGemma4MoEQuarkInt4Path: QMoE emits expert_weight_bits=4, no prescale/rotation subgraph, integer uint8 zero-points emitted without float cast, zero-points kept on CUDA int4. (14 passed total)
• ruff check clean on the test file.
• End-to-end (offline): build amd/gemma-4-26B-A4B-it-uint4-awq -> sanity generation -> wikitext PPL.

Generated with Claude Code

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

This PR extends the Gemma 4 text model builder to better support Gemma4 MoE and Quark pre-quantized expert checkpoints by (a) generalizing the make_moe_quark path to handle plain Quark/AWQ int4 experts in addition to the existing 2-bit path, and (b) adding builder/loader plumbing for int2/uint2 precision and related behaviors across the Python model-builder stack.

Changes:

• Generalize Gemma4 MoE pre-quantized Quark expert emission (input transform, fc1 reorder, zero-point dtype/omission) so the same path supports both uint2 (2-bit) and uint4/int4 (4-bit) expert layouts.
• Add int2/uint2 precision support across quant config, CLI precision selection, and Quark loader repack/unpack utilities (including split-expert uint2 handling and LoRA adapter baking for quantized loads).
• Add/extend unit tests for Gemma4 dense + MoE builder structure and for the new Quark int4 expert path expectations.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
test/python/models/test_gemma4moe_builder.py	Adds/extends Gemma4MoE builder structural tests plus new Quark int4 path assertions.
test/python/models/test_gemma4_builder.py	Adds dense Gemma4 builder tests (RoPE caches, per-layer KV geometry, attention scale/norm behaviors).
test/python/builder/test_qmoe_weights.py	Registers Gemma4Model/Gemma4MoEModel in the QMoE weight test’s builder CLI module loader.
src/python/py/models/README.md	Documents the new 2-bit (int2) pre-quantized MoE builder usage.
src/python/py/models/quantization/quant_config.py	Introduces int2/uint2 descriptors and maps --precision int2 to uint2 weight/MoE quant selection.
src/python/py/models/loaders/quark.py	Adds uint2 split-expert unpack/repack pipeline and generalizes packing/repack logic to 2-bit+4-bit.
src/python/py/models/loaders/base.py	Adds Quark factored-rotation metadata parsing, shared rotation loading, LoRA adapter attachment for quantized loads, and packed-dimension helpers.
src/python/py/models/builders/gemma.py	Adds Gemma4Model + Gemma4MoEModel builders and generalizes make_moe_quark for int4 experts.
src/python/py/models/builders/base.py	Threads precision through quant config init, aligns MoE op-type attrs, and adds factored-rotation + baked LoRA emission in matmul path.
src/python/py/models/builders/init.py	Exports Gemma4Model/Gemma4MoEModel from the builders package.
src/python/py/models/builder.py	Adds int2 CLI precision, threads true precision through extra_options, and wires Gemma4 architectures into create_model.

Suppressed comments (1)

src/python/py/models/loaders/base.py:1073

• _packed_out_factor is added as a @staticmethod with a leading underscore in the model-builder code. Per the model-builder guidelines, avoid @staticmethod helpers (and leading-underscore names); prefer an instance method (e.g., packed_out_factor) or an existing shared utility so call sites remain easy to override/mock and follow the same conventions.

    @staticmethod
    def _packed_out_factor(tensor, bits):
        """Number of logical output channels stored per element along the packed
        (column) axis: 8//bits for uint8 packing (Quark native uint2/uint4),
        32//bits for int32 packing (AWQ/GPTQ style), and 1 for unpacked floating
        tensors (per-group float scales / float zero-points)."""
        if tensor.dtype == torch.uint8:
            return 8 // bits
        if tensor.dtype == torch.int32:
            return 32 // bits
        return 1

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.