[OMNIML-5899] Add IQ tests and release notes

[hychiang-git]

Summary

• add CPU codec, export, fused-expert, Megatron, and CUDA coverage
• cover extension loading, pack→dequantize error, canonical zero payloads, scale underflow, native dispatch, and input validation
• add the release-note entry

Scope

This is PR 5 of 5 extracted from #2381. Its diff contains tests and the changelog only. It is stacked on the recipes PR.

Validation

• 96 focused CPU tests pass locally
• GPU test modules compile as Python
• repository hooks pass for all 10 changed files
• CUDA runtime checks are left to GPU CI because the local host is macOS

Row-alignment regression coverage

• IQ1_S rejects a [512, 384] input whose total element count is divisible by 256 but whose rows are not
• IQ2_XS covers the same extension-boundary contract

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}

[hychiang-git]

/claude review

[claude]

[IMPORTANT Compatibility] These two "payload field" assertions are tautologies — they can never fail, so this half of the test validates nothing.

qh is built from two uint8 bytes, so qh ∈ [0, 65535]. Then:

• ((qh >> 12) & 0x7) is masked to 3 bits, so it is always in [0, 7] → < 8 holds unconditionally.
• (qh & 0xFFF) is masked to 12 bits, so it is always in [0, 4095] → < 0x1000 holds unconditionally.

Why it matters: the metadata word bit layout (bits 0..11 = high index bits, 12..14 = local scale, 15 = delta sign) is the single highest-risk part of the GGML-compatibility claim, and this is the only CPU test that purports to check the encoder's output bits on real data. As written, a bug that put the local-scale index in the wrong bit position, or that emitted grid indices >= 2048, would still pass.

Suggested fix — assert something the encoder could actually get wrong, e.g. that the decoded grid indices are in range and that the metadata fields are not degenerate:

    blocks = packed.reshape(-1, 50)
    qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8)
    low = blocks[:, 2:34].to(torch.int64).reshape(-1, 8, 4)
    shifts = torch.tensor([0, 3, 6, 9])
    entries = low | (((qh.unsqueeze(-1) >> shifts) & 0x7) << 8)

    # Every decoded codebook index must address the 2048-entry grid, and the
    # search must actually exercise more than one local scale / delta sign.
    assert entries.max() < 2048
    assert ((qh >> 12) & 0x7).unique().numel() > 1
    d = blocks[:, :2].contiguous().view(torch.float16).float()
    assert torch.all(d > 0)

[claude]

[IMPORTANT Compatibility] Same tautology as in test_iq1_s.py — both assertions hold for any possible byte payload:

• codes is assembled from two uint8 bytes, so codes ∈ [0, 65535].
• (codes & 0x1FF) is masked to 9 bits → always in [0, 511], so < 512 is unconditional.
• (codes >> 9) is at most 65535 >> 9 == 127, so < 128 is unconditional.

Why it matters: IQ2_XS has the more intricate layout of the two (32 uint16 codes carrying a 9-bit grid index + 7-bit sign mask, plus 16 four-bit local scales packed two-per-byte in bytes 66..73), yet the only byte-level coverage in this file is the all-zero payload. Nothing here or in test_iq2_xs_cuda.py checks the sign field or the packed scale nibbles against a hand-constructed reference the way test_iq1_s_dequantizes_ggml_metadata_bit_fields does for IQ1_S. A transposed sign bit or a swapped scale nibble would round-trip self-consistently through our own encoder/decoder pair and pass every test in this PR, while producing a checkpoint llama.cpp decodes incorrectly.

Two suggestions:

1. Replace the vacuous bounds with checks that can fail — e.g. the 7-bit sign popcount parity constraint GGML relies on, that the decoded 4-bit scale nibbles are not all identical, and that d > 0.
2. Add an IQ2_XS counterpart to test_iq1_s_dequantizes_ggml_metadata_bit_fields that builds bytes by hand (a known grid index, a known sign mask, a known scale nibble) and asserts dequantize_iq2_xs returns the value the GGML formula predicts. That is the assertion that actually pins the wire format.

[claude]

[IMPORTANT Compatibility] assert torch.equal(packed, dispatched) does not test what the test name implies — both sides come from the same CUDA extension, so this is a restatement of the determinism assertion on line 39, not a check of the native-dispatch path against a reference.

Trace: quantize_iq1_s (modelopt/torch/quantization/ggml/iq1_s.py:240-252) takes the CUDA branch whenever weight.is_cuda and an extension is loadable. _extension() has already forced the extension to load with raise_if_failed=True and get_cuda_ext_iq1_s caches on the function attribute, so quantize_iq1_s is guaranteed to hit extension.pack(...). packed and dispatched are therefore the same kernel's output twice over.

Why it matters: iq1_s.cu is an independent reimplementation of _encode_blocks — the 2048-entry codebook search, the d = (amax / 16.875) * 0.61 fp16 scale rounding, the 16-way (8 local scales × 2 delta signs) choice, and the "lowest index wins on equal error" tie-break. Divergence between the two encoders is the most likely defect in the whole feature, and no test in this PR compares them. The GPU suite would stay green with a CUDA kernel that systematically picks a different local scale than the CPU reference; only the loose normalized_mse < 0.25 bound would notice, and only for a gross error.

Suggested change (applies equally to tests/gpu/torch/quantization/test_iq2_xs_cuda.py:33-40) — compare against the reference encoder rather than against itself:

def test_iq1_s_cuda_pack_matches_reference_encoder():
    generator = torch.Generator(device="cuda").manual_seed(1234)
    weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16)

    packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50)
    reference, _ = quantize_iq1_s(weight.cpu())

    assert torch.equal(packed.cpu(), reference)

If exact byte equality is not a contract the kernel intends to hold (e.g. tie-breaks are allowed to differ), then say so in a comment and assert the weaker but still meaningful property instead — that the two encoders' reconstruction errors agree to within a tight tolerance:

    cuda_err = (dequantize_iq1_s(packed, shape).float() - weight.float()).square().mean()
    cpu_err = (dequantize_iq1_s(reference, shape.cpu()).float() - weight.cpu().float()).square().mean()
    torch.testing.assert_close(cuda_err.cpu(), cpu_err, rtol=1e-3, atol=0)

Either way, keep the existing determinism assertion — it is the only part of the current test that can fail.

[claude]

[SUGGESTION] The docstring claims parity with HF export, but the test never exercises the HF path — it only calls GPTModelExporter._name_remapping and compares against linear.weight_quantizer(linear.weight), i.e. the same quantize_iq1_s/quantize_iq2_xs encoder that _get_iq_weight_state calls. So it verifies the Megatron plumbing (packed buffer lands under the ...weight key, layer_config_dict gets the right metadata), which is worth having, but it cannot detect the Megatron and HF exporters drifting apart.

Either reword to what is actually asserted, e.g. """Megatron export writes the packed IQ payload under the plain weight key, with no scale tensors.""", or make the parity claim real by also running _export_quantized_weight on an equivalent nn.Linear and asserting the two payloads are byte-identical. A wrong docstring on a cross-backend-parity test is the kind that gets trusted later.

[claude]

[SUGGESTION] The PR description lists "input validation" as covered, but only one of the validation branches has a test. Still unexercised across both test_iq1_s.py and test_iq2_xs.py:

• validate_weight: the empty-weight ValueError, the non-floating-point TypeError (note it is TypeError, not ValueError — an easy thing to regress), and the non-finite ValueError. The last one matters most: _encode_blocks computes amax from the block, so a NaN would silently poison the fp16 scale for all 256 values if the guard were ever dropped.
• quantize_iq*: the block_chunk_size <= 0 ValueError.
• validate_packed_weights: the wrong-dtype / wrong-last-dim ValueError and the size-mismatch ValueError — the latter is the guard that stops a mismatched weight_shape from reshaping garbage during restore.
• iq*_fake_quant: the num_bits mismatch ValueError, and the search_impl != "auto" NotImplementedError. That second one is user-facing — someone setting backend_extra_args={"search_impl": "greedy"} should get a clear error, and the code also silently accepts an iq_search_impl alias that nothing tests.

Most of these are one-line pytest.raises additions in the shared style already used by test_iq1_s_requires_complete_last_dimension_blocks, and the two fake_quant guards in particular are cheap given the Quantizer stub already exists in this file.

[claude]

Claude review — 3 IMPORTANT, 2 SUGGESTION, 0 CRITICAL

Scope. 9 changed files (497 additions, 0 deletions), all tests plus one CHANGELOG.rst line. I reviewed all 8 test files and cross-read the code under test on the base branch (modelopt/torch/quantization/ggml/{iq1_s,iq2_xs,common}.py, modelopt/torch/kernels/quantization/ggml/*.{cpp,cu}, modelopt/torch/export/{unified_export_hf,unified_export_megatron,quant_utils,convert_hf_config,moe_utils}.py, extensions.py) to check the assertions against real behaviour. I could not execute pytest in this environment, so the findings below are from static tracing.

What checks out

• Payload sizes, shapes, and derived constants are all consistent with the implementation: 50/74 block bytes, 1.5625/2.3125 effective bits (convert_hf_config.py:121), (N, cols//256, bytes) packed shape, and the [*shape[:-2], shape[-2] * 256] logical-shape reconstruction in the Megatron test.
• test_iq1_s_dequantizes_ggml_metadata_bit_fields is the strongest test in the PR — I hand-verified the arithmetic: d = 2.0, local = 3 → scale = 2 * (2*3+1) = 14.0, bit 15 set → delta = -0.125, and qh = 48712 fits uint16. It matches dequantize_iq1_s exactly.
• The pytest.raises regexes match the real messages: "last weight dimension" against common.py:31, "supports float32, float64, float16, and bfloat16" against iq1_s.cu:249 / iq2_xs.cu:276, "tensor model parallel size 1" against unified_export_megatron.py:322.
• test_megatron_iq_export_rejects_tensor_parallelism reaches the guard safely: _get_quantization_format is stateless, and the raise at unified_export_megatron.py:317 precedes every use of instance state the object.__new__ shell does not set.
• test_export_registers_packed_weight_buffers monkeypatches the right target — moe_utils.py:77 imports _export_quantized_weight inside the function, so patching the unified_export_hf module attribute does take effect. The (wrapper, dtype) stub signature matches the call at moe_utils.py:218, and del wrapper.weight on an nn.Parameter followed by register_buffer correctly drives the else branch at moe_utils.py:224.
• iq2_xs_pack validates on numel, not shape, so passing the unreshaped (8, 512) weight in the CUDA test is fine and matches the (16, 256) dispatch path byte-for-byte.
• The CHANGELOG.rst entry is one sentence, user-facing, and filed under the existing *Quantization* sub-section — per CLAUDE.md guidance.

Most impactful findings

1. Four assertions across two tests are tautologies (test_iq1_s.py:89-90, test_iq2_xs.py:76-77). Every one is a comparison of a bit-masked value against its own mask bound — (qh >> 12) & 0x7 < 8, qh & 0xFFF < 0x1000, codes & 0x1FF < 512, and codes >> 9 < 128 (max 65535 >> 9 == 127). They hold for any byte payload and cannot fail. Both tests are named ..._payload_fields, so they read as covering the bit layout while asserting nothing about it.

2. The CUDA tests never compare against the CPU reference encoder (test_iq1_s_cuda.py:40, test_iq2_xs_cuda.py:40). assert torch.equal(packed, dispatched) has both operands produced by the same extension — quantize_iq* takes the CUDA branch whenever the weight is on device and the extension is cached, which _extension() guarantees. iq1_s.cu/iq2_xs.cu are independent reimplementations of _encode_blocks (codebook search, fp16 scale rounding, 16-way local-scale/delta choice, lowest-index tie-break); nothing in this PR would catch them diverging from the reference.

3. No IQ2_XS byte-layout test. IQ1_S gets a hand-constructed-bytes decode test; IQ2_XS — which has the harder layout (9-bit grid index + 7-bit sign per code, plus 16 four-bit scales packed two per byte) — gets only all-zero payloads and a self-consistent round trip. A transposed sign bit or swapped scale nibble would pass everything here while producing a checkpoint llama.cpp decodes wrongly, which defeats the "GGML-compatible" claim in the changelog entry.

Risk

Low for the shipped library — this PR adds no product code and cannot break existing checkpoints, recipes, or public API. The risk is in the coverage it is credited with: with findings 1-3 fixed the suite genuinely pins the GGML wire format and the CUDA/CPU encoder equivalence; as it stands, the two highest-risk parts of the stacked feature (bit packing and the hand-written kernels) are effectively unguarded, and the vacuous assertions make that gap invisible to a future reader.

Findings 1 and 2 are small, local edits. Finding 3 and the two SUGGESTIONs are additive and can reasonably follow later if you would rather land the stack — your call.

🤖 Generated with Claude Code

[hychiang-git]

Superseded by the four scoped PRs. The tests and changelog formerly in this PR now live with the code they cover:

• direct native-extension tests → [OMNIML-5899] Add CUDA kernels for IQ packing #2448
• CPU codec and CUDA numerical/fallback tests → [OMNIML-5899] Add IQ quantization codecs and backend #2446
• Hugging Face, Megatron, metadata, and fused-expert export tests → [OMNIML-5899] Export IQ checkpoints from HF and Megatron #2447
• changelog entry → [OMNIML-5899] Add IQ post-training quantization recipes #2449

Closing this PR without merging.