Build the GGML IQ packing kernels as a single CUDA extension

[cjluo-nv]

What does this PR do?

Type of change: Code refactoring

#2448 added the GGML IQ packing kernels as two torch extensions,
modelopt_cuda_ext_iq1_s and modelopt_cuda_ext_iq2_xs. This merges them into one,
modelopt_cuda_ext_ggml.

The existing per-extension split in extensions.py exists for reasons that don't apply
to the IQ formats: get_cuda_ext gates on CUDA >=11 while _fp8/_mx gate on
>=11.8, and _mx needs --use_fast_math, which must not reach the base tensor_quant
kernels. get_cuda_ext_iq1_s and get_cuda_ext_iq2_xs differed in none of that —
same >=11.8 gate, same -O3 flags, same common.cuh — so the split only compiled the
shared header twice, ran nvcc twice, and grew the loader, __getattr__, and
precompile() once per format. With IQ2_XXS / IQ3_S / IQ4_NL plausibly following, that
scales badly.

Changes:

• New ggml/ggml.cpp holds both host-side validation wrappers and the single
  PYBIND11_MODULE, binding iq1_s_pack and iq2_xs_pack (previously each module
  exported a bare pack). Deletes ggml/iq1_s.cpp and ggml/iq2_xs.cpp; the
  validation logic and docstrings carry over unchanged.
• get_cuda_ext_iq1_s + get_cuda_ext_iq2_xs → get_cuda_ext_ggml, which builds
  ggml.cpp, iq1_s.cu, and iq2_xs.cu together. The retry-on-raise_if_failed
  semantics of the old getters are preserved.
• Each format keeps its kernels in its own translation unit, so adding a format is a new
  .cu plus one module.def — no new extension, loader, or precompile() line.

No caller outside extensions.py and its tests referenced the old getters on main, so
nothing else changes. Note for the follow-up PRs in the #2448 series
(#2446/#2447/#2449): the codec layer should call
get_cuda_ext_ggml().iq1_s_pack(...) / .iq2_xs_pack(...) instead of
get_cuda_ext_iq1_s().pack(...) / get_cuda_ext_iq2_xs().pack(...).

Usage

from modelopt.torch.quantization.extensions import get_cuda_ext_ggml

ext = get_cuda_ext_ggml(raise_if_failed=True)
iq1_s_payload = ext.iq1_s_pack(weight, iq1s_grid)            # uint8 [numel / 256, 50]
iq2_xs_payload = ext.iq2_xs_pack(weight, iq2xs_grid, scales) # uint8 [numel / 256, 74]

Testing

Ran on a single H200 NVL (TRT-LLM 1.3.0rc27.dev202609170000 container), building the
merged extension from scratch:

• pytest tests/gpu/_extensions/test_torch_extensions.py — 24 passed (6:44). This
  is the full existing IQ suite (zero-block layout, encode, dtype rejection,
  row-straddling rejection, invalid/negative-zero scales, byte-exact dtype equivalence,
  and the brute-force optimality round-trip) reparametrized onto the merged module,
  plus the untouched modelopt_cuda_ext / _fp8 / _mx load tests.
• Verified precompile() loads all four extensions and that the merged module exports
  exactly iq1_s_pack and iq2_xs_pack with the expected arities.
• Off-GPU: compiled the three sources directly and linked them into one .so to confirm
  no duplicate-symbol collisions between the two .cu translation units.
• pre-commit run --files ... passes on all changed files (ruff, mypy, clang-format,
  bandit, license headers).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — the removed getters were added in #2448
  (merged today, unreleased) and have no callers outside this file's own tests.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A — no new code or dependencies; the moved wrappers keep their original attribution.
• Did you write any new necessary tests?: ✅ — existing coverage reparametrized onto the merged module; no behavior change to test.
• Did you update Changelog?: N/A — internal refactor of an unreleased, not-yet-wired-up API.
• Did you get Claude approval on this PR?: ❌ — not yet run.

Additional Information

Follow-up to #2448. Merge before the remaining PRs in that series (#2446, #2447, #2449)
land, so the codec layer is written against get_cuda_ext_ggml and no rename is needed
afterwards.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added IQ1_S packing support through the GGML CUDA extension.
  • Added a unified GGML extension loader for IQ1_S and IQ2_XS packing.
  • Improved extension loading reliability when a cached extension is unavailable.

• Changes

  • Renamed the IQ2_XS packing binding from pack to iq2_xs_pack.
  • Consolidated IQ1_S and IQ2_XS extension access under the shared GGML loader.
  • Updated GPU validation and coverage to use the unified extension interface.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 52cfd1bf-51ee-4b56-b280-4cc351e354ca

📥 Commits

Reviewing files that changed from the base of the PR and between ad1bad7 and d30edd6.

📒 Files selected for processing (4)

• modelopt/torch/kernels/quantization/ggml/ggml.cpp
• modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
• modelopt/torch/quantization/extensions.py
• tests/gpu/_extensions/test_torch_extensions.py

💤 Files with no reviewable changes (1)

• modelopt/torch/kernels/quantization/ggml/iq1_s.cpp

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

The change consolidates IQ1_S and IQ2_XS packing in one GGML CUDA extension. It moves IQ1_S binding logic into the shared C++ module, renames the IQ2_XS binding, and updates GPU tests to use the shared loader.

Changes

GGML IQ packer consolidation

Layer / File(s)	Summary
Shared GGML binding API modelopt/torch/kernels/quantization/ggml/ggml.cpp, modelopt/torch/kernels/quantization/ggml/iq1_s.cpp	The shared module adds validated iq1_s_pack handling, exports iq1_s_pack, and renames the IQ2_XS binding from pack to iq2_xs_pack. The standalone IQ1_S wrapper is removed.
Combined GGML extension loader modelopt/torch/quantization/extensions.py	The loader builds and caches one GGML extension from the shared, IQ1_S, and IQ2_XS sources. Lazy loading and precompilation use cuda_ext_ggml.
Shared extension test coverage tests/gpu/_extensions/test_torch_extensions.py	GPU tests use the shared extension and named packer attributes. Existing encoding, validation, dtype, scale, optimality, and non-finite-input coverage is rewired accordingly.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant Test
  participant GGMLLoader
  participant GGMLExtension
  participant CUDAPacker
  Test->>GGMLLoader: request get_cuda_ext_ggml()
  GGMLLoader->>GGMLExtension: build or return cached extension
  Test->>GGMLExtension: call iq1_s_pack or iq2_xs_pack
  GGMLExtension->>CUDAPacker: execute CUDA packer
  CUDAPacker-->>Test: return packed tensor

Loading

Merge Risk: ⚪ Minimal · up to d30ed

The consolidated extension retains device validation and has no remaining merge-blocking issue.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 55.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 20 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: consolidating the GGML IQ packing kernels into a single CUDA extension.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	PASS. The authoritative PR diff changes only the GGML CUDA sources, modelopt/torch/quantization/extensions.py, and GPU tests. Added Python lines introduce no torch.load(..., weights_only=False), `…

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

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

[meenchen]

Bot review (claude-opus-5) — DM the bot to share feedback.

LGTM — merging the two IQ extensions into modelopt_cuda_ext_ggml is a clean, mechanical consolidation; the wrappers, validation, docstrings and retry semantics carry over unchanged, and the GPU suite is reparametrized onto the merged module with coverage intact.

Verified: iq1_s.cu/iq2_xs.cu keep all internals in anonymous namespaces and only export the distinct *_pack_cuda entry points, and common.cuh is header-only (inline/templates), so the three-TU build has no ODR or duplicate-symbol risk. package-data globs pick up the renamed ggml.cpp automatically.

Needs action:

• Sign off as a human on the test edit: test_cuda_ext_iq1_s/test_cuda_ext_iq2_xs were collapsed into one test_cuda_ext_ggml and fmt.get_extension replaced by fmt.packer, which the single-extension merge justifies — but a deleted test needs owner confirmation.
• Update the stacked PRs (#2446/#2447/#2449) to call get_cuda_ext_ggml().iq1_s_pack(...) / .iq2_xs_pack(...) before they land, as the PR body notes.

No action needed:

• Optional: _IQ_EXTENSIONS in tests/gpu/_extensions/test_torch_extensions.py now holds formats, not extensions; renaming it would match the new _IqFormat docstring.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-09-17 20:46 UTC

[codecov]

Codecov Report

❌ Patch coverage is 85.71429% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 78.96%. Comparing base (216f28a) to head (d30edd6).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/extensions.py	85.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2462      +/-   ##
==========================================
+ Coverage   71.48%   78.96%   +7.48%     
==========================================
  Files         590      590              
  Lines       64766    64779      +13     
==========================================
+ Hits        46299    51154    +4855     
+ Misses      18467    13625    -4842     

Flag	Coverage Δ
examples-diffusers	20.88% <28.57%> (+<0.01%)	⬆️
examples-gpt-oss	13.40% <28.57%> (+<0.01%)	⬆️
examples-hf_ptq	22.50% <71.42%> (-0.03%)	⬇️
examples-llm_distill	13.46% <28.57%> (+<0.01%)	⬆️
examples-llm_eval	17.38% <28.57%> (+<0.01%)	⬆️
examples-llm_qat	17.67% <28.57%> (-0.01%)	⬇️
examples-llm_sparsity	15.93% <28.57%> (+<0.01%)	⬆️
examples-megatron_bridge	26.28% <28.57%> (-0.12%)	⬇️
examples-specdec_bench	13.15% <28.57%> (+<0.01%)	⬆️
examples-speculative_decoding	17.80% <28.57%> (-0.06%)	⬇️
examples-torch_onnx	21.90% <28.57%> (+<0.01%)	⬆️
examples-torch_trt	15.23% <28.57%> (+<0.01%)	⬆️
gpu	58.36% <85.71%> (+25.92%)	⬆️
regression	15.16% <28.57%> (+0.30%)	⬆️
unit	57.85% <28.57%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.