hip-quant 2.2.3__tar.gz → 2.2.4__tar.gz

{hip_quant-2.2.3 → hip_quant-2.2.4}/MANIFEST.in

@@ -1,6 +1,5 @@
 include hip_quantize.dll
 include hip_quantize_rocm721.dll
-include _C*.pyd
 include _C.cp312-win_amd64.pyd
 include hip_quant_types.h
 include hip_quant_util.h
@@ -10,6 +9,9 @@ include build.ps1
 include build.sh
 recursive-include kernels *.cu
 recursive-include torch_ext *.cpp *.hip *.h
+exclude torch_ext/ple_directstorage.cpp
+exclude torch_ext/ple_directstorage.h
+exclude torch_ext/ple_directstorage_hip.hip
 recursive-include tests *.py
 recursive-include codebooks *.bin
 include torch_api.py

{hip_quant-2.2.3/hip_quant.egg-info → hip_quant-2.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hip-quant
-Version: 2.2.3
+Version: 2.2.4
 Summary: High-performance GPU tensor quantization and inference engine for AMD ROCm. Native AOT GEMV (K-Quants & I-Quants), State Space Model (SSM/DeltaNet) kernels, FP8 WMMA WaveAttention, and portable HIP C++ kernels for RDNA4, RDNA3, and CDNA.
 Author-email: Thenukegun10x <conormorris84@gmail.com>
 License:                                  Apache License
@@ -256,6 +256,14 @@ Dynamic: license-file
 > target; f16 WMMA operands are the accurate choice. **Validate accuracy and
 > convergence on your own workload before training with any of these paths.**
 
+## 🌟 What's New in 2.2.4
+
+> **2.2.4.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
+
+- **🛡️ Safer imatrix handling.** The Python quantization entry points validate full-matrix and llama-style per-column imatrices before native calls. The three GGML I-Quants that require importance data warn once when the native unweighted fallback is used.
+- **🧵 Thread-safe HIP initialization.** Quantization selects the cached HIP device on each calling thread and serializes device/codebook initialization, avoiding worker-thread launches against an unintended device.
+- **📦 Native PLE reads on Windows.** The optional IOCP reader fills caller-owned pinned memory for large row batches. The DirectStorage backend is not included in this release.
+
 ## 🌟 What's New in 2.2.3
 
 > **2.2.3.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
@@ -283,7 +291,7 @@ Security, correctness and robustness fixes from an internal audit:
 
 - **Kernels**: fixed a missing CTA barrier in `wave_attn`/`wave_attn_long` (V overwrote the K tile while waves were still reading it), corrected the causal K-tile clamp (dropped the diagonal tile at exact boundaries), added `Dim <= 128` guards to the wave forward entry points, made the `gemv_tune` dispatch selector thread-local and exception-safe, and 64-bit-corrected FP8 GEMM index arithmetic.
 - **Parsers (untrusted input)**: bounded GGUF string/array reads and the legacy HQ2 manifest (small crafted files could force multi-GiB allocations), rejected `general.alignment <= 0`, and replaced `assert`-based shape validation with explicit raises.
-- **Offline API**: `quantize_from_fp8` now validates the `imatrix` element count (was a host + device out-of-bounds read), `ggml_blck_size_for` restype corrected, `math.prod` for element counts.
+- **Offline API**: quantize_from_fp8 validates the imatrix element count (was a host + device out-of-bounds read), ggml_blck_size_for restype corrected, math.prod for element counts.
 - **Hardening**: `_C` is loaded by absolute path (no `sys.path` mutation), `HIP_QUANT_DLL`/`_PATH` must be absolute, and the mxfp4 LUT init is `std::call_once`.
 - **gpu-smi** is now fetched from its GitHub Releases at build time (checksum-verified) and bundled in the wheel; `smi.py` discovers it dynamically (env → bundled → console script → PATH) instead of hardcoding a path.
 
@@ -342,9 +350,11 @@ out = hip_quant.wave_attn(q, k, v)  # native FP8 WMMA (inference-grade)
 # Native FP16 WMMA training kernels on gfx1201 (forward + backward, no Triton):
 qh, kh, vh = q.half(), k.half(), v.half()
 out_fp16 = hip_quant.wave_attn(qh, kh, vh, precision="fp16")
-# BF16 Triton fallback (requires ROCm Triton):
+# Native BF16 WMMA training kernels on gfx1201:
 qbf, kbf, vbf = q.bfloat16(), k.bfloat16(), v.bfloat16()
 out_bf16 = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16")
+# Deprecated Triton compatibility path:
+out_bf16_triton = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16_triton")
 # Or defer to PyTorch SDPA entirely:
 out_precise = hip_quant.wave_attn(q, k, v, precision="sdpa")
 ```
@@ -382,27 +392,33 @@ See `benchmarks/wave_attn_train_results.md` for methods and all metrics.
 > including for prequantized `uint8` inputs. The SDPA precision mode is for
 > floating-point training inputs and returns a float32 output.
 
-**BF16 training mode (experimental):** `precision="bf16"` uses Triton kernels
-that compile to gfx1201 `v_wmma_f32_16x16x16_bf16` instructions. It accepts
-contiguous, equal-shaped BF16 Q/K/V with head dimension 64 or 128; the output
-and input gradients are BF16. A separate FP32 output is saved for backward so
-large-score gradients retain accuracy. The full public forward/backward call
-beats default BF16 SDPA on the RX 9070 XT at `B=6,H=12,S=789`: 0.75 vs 1.02 ms
-for D=64 noncausal, 0.35 vs 0.98 ms causal; 2.24 vs 2.85 ms for D=128
-noncausal, 1.41 vs 1.96 ms causal (100 GPU-event repeats, medians of three
-rounds). Output and gradient relative L2 against FP32 MATH stay below 0.26%
-at input scales 1 and 5. Small attention calls can be slower than SDPA, so
-select this mode for measured training workloads. See
-`benchmarks/check_wave_attn_train_gate.py`,
-`benchmarks/check_plant_wave_continuation.py`, and
-`benchmarks/wave_attn_train_results.md` for the checks.
+**BF16 training mode:** `precision="bf16"` uses native GFX12 BF16 WMMA forward
+and backward kernels. Q/K/V must share batch, head count, and head dimension;
+K and V must have equal sequence lengths, while query and key lengths may differ.
+The supported head dimensions are 64 and 128. The path returns BF16 output and
+gradients, and saves FP32 output and softmax state for backward. Backward uses
+the incoming BF16 gradient directly, avoiding an FP32 round trip and a duplicate
+BF16 gradient buffer. The accuracy
+and guard cases in `tests/torch/test_wave_attn_train_bf16.py` cover distinct
+query/key lengths and invalid native-binding arguments. A bounded GPU graph
+replay at `B=1,H=8,S=512,D=64` measured native BF16 at 0.98x native FP16 time
+for noncausal and causal forward plus backward; this is a shape-specific
+comparison, not a universal speed claim.
+
+The previous BF16 performance table and 3,000-update checkpoint run below used
+the Triton implementation. It is retained temporarily as
+`precision="bf16_triton"`, which emits a `FutureWarning`; use `precision="bf16"`
+for the native path. The Triton path requires ROCm Triton. See
+`benchmarks/check_wave_attn_train_gate.py` and
+`benchmarks/wave_attn_train_results.md` for the historical Triton measurements.
 
 To route timm's SDPA calls through the BF16 path, set
 `HIP_QUANT_WAVE_ATTN_PRECISION=bf16` before calling `patch_sdpa()` (or before a
 trainer's `--wave-attention` flag). The helper makes timm's transposed Q/K/V
 views contiguous and returns BF16 outputs; unsupported attention shapes still
 use SDPA. A BF16 kernel error propagates instead of silently falling back.
-From a real PlantUnified checkpoint at step 131,780, the BF16 path completed
+From a real PlantUnified checkpoint at step 131,780, the legacy
+`precision="bf16_triton"` path completed
 3,000 more training updates with finite loss and gradients, using 24 attention
 calls per update across the 12 DINOv3 blocks and their checkpoint recomputations.
 The bounded run did not perform EMA or validation; the eight-step matched A/B
@@ -510,13 +526,13 @@ Non-linear quants that preserve quality at extreme low bits:
 - `IQ1_S`, `IQ1_M`, `IQ2_XXS`, `IQ2_XS`, `IQ2_S`, `IQ3_XXS`, `IQ3_S`, `IQ4_NL`, `IQ4_XS`
 
 imatrix contract: llama.cpp defines **one float per input column, shared
-across rows** (size `ne[0]`). hip_quant kernels index `imatrix + row*n_per_row`,
-so `quantize_numpy` requires `imatrix.shape == arr.shape` — tile a raw llama
-`.dat` column vector across rows first. Weighting math matches llama
+across rows** (size `ne[0]`). The Python APIs accept that `(n_per_row,)`
+vector or a full matrix and tile vectors internally because the current
+kernels index `imatrix + row*n_per_row`. Weighting math matches llama
 (`qw[i] * sqrt(sigma2 + x[i]^2)`); all-ones imatrix is the identity and is
-what the byte-exact tests use today. `quantize_from_fp8` does not yet validate
-imatrix shape (known gap — pass a full matrix). See `IMATRIX_PLAN.md` for the
-per-column-accept + non-trivial-imatrix test rollout.
+what the byte-exact tests use today. The three llama-required I-Quants warn
+once per process when called without an imatrix. `.dat` loading and
+non-trivial-imatrix byte-exact coverage remain in `IMATRIX_PLAN.md`.
 
 ### ⚖️ Ternary Quants
 For models trained to be ternary (BitNet, TriLM):
@@ -543,6 +559,27 @@ x = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
 q, s = hq.quantize_mxfp8_e4m3(x)   # q:uint8 [4096,4096], s:uint8 UE8M0 [4096,128]
 y = hq.dequantize_mxfp8_e4m3(q, s) # float32
 q2, s2 = hq.quantize_mxfp8_e5m2(x) # E5M2 variant for gradients
+
+# Direct OCP MXFP8 GEMM: no FP32 dequantization buffer.
+w = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+wq, ws = hq.quantize_mxfp8_e4m3(w)
+y = hq.mxfp8_linear_forward_ue8m0(q, s, wq, ws, x)
+```
+
+The gfx12 path applies both UE8M0 scales to each K=32 dot-product partial
+while the E4M3 operands stay packed for WMMA.
+
+#### Experimental signed MXINT4 (custom format)
+
+This is a hip_quant format, not an OCP v1.0 MX format: signed symmetric int4
+values in `[-7,7]`, with one UE8M0 scale per 32 contiguous values. Quantized
+operands stay nibble-packed and the GEMM uses RDNA4's native IU4 WMMA
+instruction; software applies the block scales to each K=32 partial.
+
+```python
+a4, a4_scales = hq.quantize_mxint4(x)
+w4, w4_scales = hq.quantize_mxint4(w)
+y4 = hq.mxint4_wmma_linear(a4, a4_scales, w4, w4_scales, x)
 ```
 
 CPU reference in `tests/test_mxfp8_numerical.py` (11 cases: `zero`, `1.0→ue119 q0x78`, `448→ue127`, `500→ue128`, `NaN/Inf`, `subnormal`) and GPU `tests/torch/test_mxfp8_torch.py` (10 warp-correct cases, `tail 33/31`, `F16/BF16`, `vs FP32-scale 1.4×`).
@@ -645,7 +682,7 @@ pip install hip-quant
 pip install "hip-quant[torch]"
 
 # Or a local wheel
-pip install dist/hip_quant-2.2.3-cp312-cp312-win_amd64.whl
+pip install dist/hip_quant-2.2.4-cp312-cp312-win_amd64.whl
 ```
 
 Two wheels are published per release, one per PyTorch ABI. Install the one that
@@ -653,11 +690,11 @@ matches your PyTorch, so the extension loads without an ABI mismatch:
 
 | Wheel | PyTorch | ROCm |
 |---|---|---|
-| `hip-quant==2.2.3` | 2.9.x | ROCm 7.2.1 |
-| `hip-quant==2.2.3.post215` | 2.15 / TheRock | ROCm 7.14 |
+| `hip-quant==2.2.4` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.4.post215` | 2.15 / TheRock | ROCm 7.14 |
 
 ```powershell
-pip install "hip-quant==2.2.3.post215"   # PyTorch 2.15 / TheRock
+pip install "hip-quant==2.2.4.post215"   # PyTorch 2.15 / TheRock
 ```
 
 #### Custom PyTorch 2.15 (self-hosted wheels)
@@ -667,7 +704,7 @@ Its wheels live in a public Backblaze B2 bucket; point pip at the index page
 with `--find-links`:
 
 ```powershell
-pip install "hip-quant==2.2.3.post215" `
+pip install "hip-quant==2.2.4.post215" `
   --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
 ```
 
@@ -1241,9 +1278,8 @@ All PyTorch extension functions are guarded against:
 - **Gated-norm weight layout** — `fast_rms_norm_gated_forward` dispatches by
   `w.numel()`: per-head `[H,D]` or shared `[D]` (stride 0). Anything else raises
   instead of reading out of bounds (regression: Qwen3.5 shared `ssm_norm`)
-- **imatrix shape** — `quantize_numpy` requires `imatrix.shape == arr.shape`
-  (full matrix; tile llama per-column vectors first). `quantize_from_fp8`
-  does not validate yet — pass a full matrix there (see `IMATRIX_PLAN.md`)
+- **imatrix shape** — all Python quantization entry points accept only a
+  matching full matrix or a per-column vector and validate before native calls.
 
 ---
 

{hip_quant-2.2.3 → hip_quant-2.2.4}/README.md

@@ -31,6 +31,14 @@
 > target; f16 WMMA operands are the accurate choice. **Validate accuracy and
 > convergence on your own workload before training with any of these paths.**
 
+## 🌟 What's New in 2.2.4
+
+> **2.2.4.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
+
+- **🛡️ Safer imatrix handling.** The Python quantization entry points validate full-matrix and llama-style per-column imatrices before native calls. The three GGML I-Quants that require importance data warn once when the native unweighted fallback is used.
+- **🧵 Thread-safe HIP initialization.** Quantization selects the cached HIP device on each calling thread and serializes device/codebook initialization, avoiding worker-thread launches against an unintended device.
+- **📦 Native PLE reads on Windows.** The optional IOCP reader fills caller-owned pinned memory for large row batches. The DirectStorage backend is not included in this release.
+
 ## 🌟 What's New in 2.2.3
 
 > **2.2.3.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
@@ -58,7 +66,7 @@ Security, correctness and robustness fixes from an internal audit:
 
 - **Kernels**: fixed a missing CTA barrier in `wave_attn`/`wave_attn_long` (V overwrote the K tile while waves were still reading it), corrected the causal K-tile clamp (dropped the diagonal tile at exact boundaries), added `Dim <= 128` guards to the wave forward entry points, made the `gemv_tune` dispatch selector thread-local and exception-safe, and 64-bit-corrected FP8 GEMM index arithmetic.
 - **Parsers (untrusted input)**: bounded GGUF string/array reads and the legacy HQ2 manifest (small crafted files could force multi-GiB allocations), rejected `general.alignment <= 0`, and replaced `assert`-based shape validation with explicit raises.
-- **Offline API**: `quantize_from_fp8` now validates the `imatrix` element count (was a host + device out-of-bounds read), `ggml_blck_size_for` restype corrected, `math.prod` for element counts.
+- **Offline API**: quantize_from_fp8 validates the imatrix element count (was a host + device out-of-bounds read), ggml_blck_size_for restype corrected, math.prod for element counts.
 - **Hardening**: `_C` is loaded by absolute path (no `sys.path` mutation), `HIP_QUANT_DLL`/`_PATH` must be absolute, and the mxfp4 LUT init is `std::call_once`.
 - **gpu-smi** is now fetched from its GitHub Releases at build time (checksum-verified) and bundled in the wheel; `smi.py` discovers it dynamically (env → bundled → console script → PATH) instead of hardcoding a path.
 
@@ -117,9 +125,11 @@ out = hip_quant.wave_attn(q, k, v)  # native FP8 WMMA (inference-grade)
 # Native FP16 WMMA training kernels on gfx1201 (forward + backward, no Triton):
 qh, kh, vh = q.half(), k.half(), v.half()
 out_fp16 = hip_quant.wave_attn(qh, kh, vh, precision="fp16")
-# BF16 Triton fallback (requires ROCm Triton):
+# Native BF16 WMMA training kernels on gfx1201:
 qbf, kbf, vbf = q.bfloat16(), k.bfloat16(), v.bfloat16()
 out_bf16 = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16")
+# Deprecated Triton compatibility path:
+out_bf16_triton = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16_triton")
 # Or defer to PyTorch SDPA entirely:
 out_precise = hip_quant.wave_attn(q, k, v, precision="sdpa")
 ```
@@ -157,27 +167,33 @@ See `benchmarks/wave_attn_train_results.md` for methods and all metrics.
 > including for prequantized `uint8` inputs. The SDPA precision mode is for
 > floating-point training inputs and returns a float32 output.
 
-**BF16 training mode (experimental):** `precision="bf16"` uses Triton kernels
-that compile to gfx1201 `v_wmma_f32_16x16x16_bf16` instructions. It accepts
-contiguous, equal-shaped BF16 Q/K/V with head dimension 64 or 128; the output
-and input gradients are BF16. A separate FP32 output is saved for backward so
-large-score gradients retain accuracy. The full public forward/backward call
-beats default BF16 SDPA on the RX 9070 XT at `B=6,H=12,S=789`: 0.75 vs 1.02 ms
-for D=64 noncausal, 0.35 vs 0.98 ms causal; 2.24 vs 2.85 ms for D=128
-noncausal, 1.41 vs 1.96 ms causal (100 GPU-event repeats, medians of three
-rounds). Output and gradient relative L2 against FP32 MATH stay below 0.26%
-at input scales 1 and 5. Small attention calls can be slower than SDPA, so
-select this mode for measured training workloads. See
-`benchmarks/check_wave_attn_train_gate.py`,
-`benchmarks/check_plant_wave_continuation.py`, and
-`benchmarks/wave_attn_train_results.md` for the checks.
+**BF16 training mode:** `precision="bf16"` uses native GFX12 BF16 WMMA forward
+and backward kernels. Q/K/V must share batch, head count, and head dimension;
+K and V must have equal sequence lengths, while query and key lengths may differ.
+The supported head dimensions are 64 and 128. The path returns BF16 output and
+gradients, and saves FP32 output and softmax state for backward. Backward uses
+the incoming BF16 gradient directly, avoiding an FP32 round trip and a duplicate
+BF16 gradient buffer. The accuracy
+and guard cases in `tests/torch/test_wave_attn_train_bf16.py` cover distinct
+query/key lengths and invalid native-binding arguments. A bounded GPU graph
+replay at `B=1,H=8,S=512,D=64` measured native BF16 at 0.98x native FP16 time
+for noncausal and causal forward plus backward; this is a shape-specific
+comparison, not a universal speed claim.
+
+The previous BF16 performance table and 3,000-update checkpoint run below used
+the Triton implementation. It is retained temporarily as
+`precision="bf16_triton"`, which emits a `FutureWarning`; use `precision="bf16"`
+for the native path. The Triton path requires ROCm Triton. See
+`benchmarks/check_wave_attn_train_gate.py` and
+`benchmarks/wave_attn_train_results.md` for the historical Triton measurements.
 
 To route timm's SDPA calls through the BF16 path, set
 `HIP_QUANT_WAVE_ATTN_PRECISION=bf16` before calling `patch_sdpa()` (or before a
 trainer's `--wave-attention` flag). The helper makes timm's transposed Q/K/V
 views contiguous and returns BF16 outputs; unsupported attention shapes still
 use SDPA. A BF16 kernel error propagates instead of silently falling back.
-From a real PlantUnified checkpoint at step 131,780, the BF16 path completed
+From a real PlantUnified checkpoint at step 131,780, the legacy
+`precision="bf16_triton"` path completed
 3,000 more training updates with finite loss and gradients, using 24 attention
 calls per update across the 12 DINOv3 blocks and their checkpoint recomputations.
 The bounded run did not perform EMA or validation; the eight-step matched A/B
@@ -285,13 +301,13 @@ Non-linear quants that preserve quality at extreme low bits:
 - `IQ1_S`, `IQ1_M`, `IQ2_XXS`, `IQ2_XS`, `IQ2_S`, `IQ3_XXS`, `IQ3_S`, `IQ4_NL`, `IQ4_XS`
 
 imatrix contract: llama.cpp defines **one float per input column, shared
-across rows** (size `ne[0]`). hip_quant kernels index `imatrix + row*n_per_row`,
-so `quantize_numpy` requires `imatrix.shape == arr.shape` — tile a raw llama
-`.dat` column vector across rows first. Weighting math matches llama
+across rows** (size `ne[0]`). The Python APIs accept that `(n_per_row,)`
+vector or a full matrix and tile vectors internally because the current
+kernels index `imatrix + row*n_per_row`. Weighting math matches llama
 (`qw[i] * sqrt(sigma2 + x[i]^2)`); all-ones imatrix is the identity and is
-what the byte-exact tests use today. `quantize_from_fp8` does not yet validate
-imatrix shape (known gap — pass a full matrix). See `IMATRIX_PLAN.md` for the
-per-column-accept + non-trivial-imatrix test rollout.
+what the byte-exact tests use today. The three llama-required I-Quants warn
+once per process when called without an imatrix. `.dat` loading and
+non-trivial-imatrix byte-exact coverage remain in `IMATRIX_PLAN.md`.
 
 ### ⚖️ Ternary Quants
 For models trained to be ternary (BitNet, TriLM):
@@ -318,6 +334,27 @@ x = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
 q, s = hq.quantize_mxfp8_e4m3(x)   # q:uint8 [4096,4096], s:uint8 UE8M0 [4096,128]
 y = hq.dequantize_mxfp8_e4m3(q, s) # float32
 q2, s2 = hq.quantize_mxfp8_e5m2(x) # E5M2 variant for gradients
+
+# Direct OCP MXFP8 GEMM: no FP32 dequantization buffer.
+w = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+wq, ws = hq.quantize_mxfp8_e4m3(w)
+y = hq.mxfp8_linear_forward_ue8m0(q, s, wq, ws, x)
+```
+
+The gfx12 path applies both UE8M0 scales to each K=32 dot-product partial
+while the E4M3 operands stay packed for WMMA.
+
+#### Experimental signed MXINT4 (custom format)
+
+This is a hip_quant format, not an OCP v1.0 MX format: signed symmetric int4
+values in `[-7,7]`, with one UE8M0 scale per 32 contiguous values. Quantized
+operands stay nibble-packed and the GEMM uses RDNA4's native IU4 WMMA
+instruction; software applies the block scales to each K=32 partial.
+
+```python
+a4, a4_scales = hq.quantize_mxint4(x)
+w4, w4_scales = hq.quantize_mxint4(w)
+y4 = hq.mxint4_wmma_linear(a4, a4_scales, w4, w4_scales, x)
 ```
 
 CPU reference in `tests/test_mxfp8_numerical.py` (11 cases: `zero`, `1.0→ue119 q0x78`, `448→ue127`, `500→ue128`, `NaN/Inf`, `subnormal`) and GPU `tests/torch/test_mxfp8_torch.py` (10 warp-correct cases, `tail 33/31`, `F16/BF16`, `vs FP32-scale 1.4×`).
@@ -420,7 +457,7 @@ pip install hip-quant
 pip install "hip-quant[torch]"
 
 # Or a local wheel
-pip install dist/hip_quant-2.2.3-cp312-cp312-win_amd64.whl
+pip install dist/hip_quant-2.2.4-cp312-cp312-win_amd64.whl
 ```
 
 Two wheels are published per release, one per PyTorch ABI. Install the one that
@@ -428,11 +465,11 @@ matches your PyTorch, so the extension loads without an ABI mismatch:
 
 | Wheel | PyTorch | ROCm |
 |---|---|---|
-| `hip-quant==2.2.3` | 2.9.x | ROCm 7.2.1 |
-| `hip-quant==2.2.3.post215` | 2.15 / TheRock | ROCm 7.14 |
+| `hip-quant==2.2.4` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.4.post215` | 2.15 / TheRock | ROCm 7.14 |
 
 ```powershell
-pip install "hip-quant==2.2.3.post215"   # PyTorch 2.15 / TheRock
+pip install "hip-quant==2.2.4.post215"   # PyTorch 2.15 / TheRock
 ```
 
 #### Custom PyTorch 2.15 (self-hosted wheels)
@@ -442,7 +479,7 @@ Its wheels live in a public Backblaze B2 bucket; point pip at the index page
 with `--find-links`:
 
 ```powershell
-pip install "hip-quant==2.2.3.post215" `
+pip install "hip-quant==2.2.4.post215" `
   --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
 ```
 
@@ -1016,9 +1053,8 @@ All PyTorch extension functions are guarded against:
 - **Gated-norm weight layout** — `fast_rms_norm_gated_forward` dispatches by
   `w.numel()`: per-head `[H,D]` or shared `[D]` (stride 0). Anything else raises
   instead of reading out of bounds (regression: Qwen3.5 shared `ssm_norm`)
-- **imatrix shape** — `quantize_numpy` requires `imatrix.shape == arr.shape`
-  (full matrix; tile llama per-column vectors first). `quantize_from_fp8`
-  does not validate yet — pass a full matrix there (see `IMATRIX_PLAN.md`)
+- **imatrix shape** — all Python quantization entry points accept only a
+  matching full matrix or a per-column vector and validate before native calls.
 
 ---
 

hip_quant-2.2.4/_C.cp312-win_amd64.pyd

@@ -0,0 +1,4 @@
+[diffend] Oversized file quarantined before diffing.
+name: hip_quant-2.2.4/_C.cp312-win_amd64.pyd
+size: 48742400 bytes
+sha256: 3cc14589a6ad04ba4981381c14660b79cc4104b851fffbef5fc0c512158e5be7

{hip_quant-2.2.3 → hip_quant-2.2.4}/__init__.py

@@ -2,8 +2,9 @@ import ctypes
 import numpy as np
 import os
 import sys
+import warnings
 
-__version__ = "2.2.3"
+__version__ = "2.2.4"
 
 
 _TORCH_EXPORTS = {
@@ -46,6 +47,14 @@ _TORCH_EXPORTS = {
     "dequantize_e5m2",
     "dequantize_e4m3_blockwise",
     "dequantize_e5m2_blockwise",
+    "quantize_mxfp8_e4m3",
+    "quantize_mxfp8_e5m2",
+    "dequantize_mxfp8_e4m3",
+    "dequantize_mxfp8_e5m2",
+    "mxfp8_linear_forward_ue8m0",
+    "quantize_mxint4",
+    "mxint4_wmma_linear",
+    "MXINT4_BLOCK_SIZE",
     # ── Attention ──────────────────────────────────────────
     "wave_attn",
     # ── Linear / GEMM ─────────────────────────────────────
@@ -54,6 +63,7 @@ _TORCH_EXPORTS = {
     "fp8_linear_forward_fp8_weight",
     "fp8_linear_forward_fp8_input",
     "fp8_linear_forward_fp8_input_weight",
+    "fp8_grouped_expert_linear_forward_fp8_input_weight",
     "pack_fp8_weight_for_wmma",
     "fp8_linear_forward_fp8_input_weight_packed",
     "fp8_linear_forward_blockwise",
@@ -160,6 +170,46 @@ GGML_TYPE = {
     "F8_E5M2": 37,
 }
 
+_IMATRIX_REQUIRED_TYPES = frozenset(
+    (GGML_TYPE["IQ2_XXS"], GGML_TYPE["IQ2_XS"], GGML_TYPE["IQ1_S"])
+)
+_MISSING_IMATRIX_WARNED = False
+
+
+def _as_nonempty_2d(array, dtype, name):
+    array = np.ascontiguousarray(array, dtype=dtype)
+    if array.ndim == 1:
+        array = array.reshape(1, -1)
+    if array.ndim != 2:
+        raise ValueError(f"{name} must be a 1-D or 2-D array")
+    if array.shape[0] == 0 or array.shape[1] == 0:
+        raise ValueError(f"{name} must contain at least one row and one element per row")
+    return array
+
+
+def _prepare_imatrix(imatrix, shape, type_num):
+    global _MISSING_IMATRIX_WARNED
+    if imatrix is None:
+        if int(type_num) in _IMATRIX_REQUIRED_TYPES and not _MISSING_IMATRIX_WARNED:
+            warnings.warn(
+                "IQ2_XXS, IQ2_XS, and IQ1_S require an imatrix for GGML-compatible "
+                "quantization; using the native unweighted fallback",
+                UserWarning,
+                stacklevel=3,
+            )
+            _MISSING_IMATRIX_WARNED = True
+        return None
+
+    imatrix = np.asarray(imatrix, dtype=np.float32)
+    if imatrix.shape == (shape[1],):
+        # llama.cpp stores one weight per input column, shared by every row.
+        imatrix = np.broadcast_to(imatrix, shape)
+    elif imatrix.shape != shape:
+        raise ValueError(
+            f"imatrix shape {imatrix.shape} must be ({shape[1]},) or {shape}"
+        )
+    return np.ascontiguousarray(imatrix, dtype=np.float32)
+
 GGML_TYPE_BLOCK_SIZE = {
     2: 32,
     3: 32,
@@ -527,9 +577,7 @@ class HipQuant:
         """
         if self._quantize_stochastic_e5m2 is None:
             raise RuntimeError("Stochastic E5M2 not available in loaded DLL")
-        arr = np.ascontiguousarray(arr, dtype=np.float32)
-        if arr.ndim == 1:
-            arr = arr.reshape(1, -1)
+        arr = _as_nonempty_2d(arr, np.float32, "arr")
         nrows, n_per_row = arr.shape
         dst = np.empty(nrows * n_per_row, dtype=np.uint8)
         result = self._quantize_stochastic_e5m2(
@@ -657,7 +705,7 @@ class HipQuant:
         Args:
             arr: 2-D float32 numpy array (nrows, n_per_row) or 1-D.
             type_num: GGML type number (use GGML_TYPE dict).
-            imatrix: Optional importance matrix (same shape as arr).
+            imatrix: Optional full matrix or llama-style per-column vector.
             hq2_iterations: Lloyd iterations for HQ2 (1-16; default 4).
             aq2_iterations: Lloyd iterations for AQ2 (1-16; default 8).
                 AQ2 should normally use an attention-derived imatrix.
@@ -665,9 +713,7 @@ class HipQuant:
         Returns:
             np.uint8 array of quantized bytes.
         """
-        arr = np.ascontiguousarray(arr, dtype=np.float32)
-        if arr.ndim == 1:
-            arr = arr.reshape(1, -1)
+        arr = _as_nonempty_2d(arr, np.float32, "arr")
         nrows, n_per_row = arr.shape
         blck = self.blck_size(type_num)
         if blck <= 0:
@@ -680,12 +726,11 @@ class HipQuant:
         dst = np.empty(out_nbytes, dtype=np.uint8)
         src_ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
         dst_ptr = dst.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
-        im_ptr = None
-        if imatrix is not None:
-            imatrix = np.ascontiguousarray(imatrix, dtype=np.float32)
-            if imatrix.shape != arr.shape:
-                raise ValueError(f"imatrix shape {imatrix.shape} != arr shape {arr.shape}")
-            im_ptr = imatrix.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+        imatrix = _prepare_imatrix(imatrix, arr.shape, type_num)
+        im_ptr = (
+            imatrix.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+            if imatrix is not None else None
+        )
         if int(type_num) in {GGML_TYPE["AQ2"], GGML_TYPE["AQ2_QK"], GGML_TYPE["AQ2_VO"]}:
             if not 1 <= int(aq2_iterations) <= 16:
                 raise ValueError("aq2_iterations must be in [1, 16]")
@@ -729,11 +774,7 @@ class HipQuant:
             raise RuntimeError(
                 "Loaded hip_quantize library does not support kernel benchmarking; rebuild it"
             )
-        arr = np.ascontiguousarray(arr, dtype=np.float32)
-        if arr.ndim == 1:
-            arr = arr.reshape(1, -1)
-        if arr.ndim != 2:
-            raise ValueError("arr must be a 1-D or 2-D float32 array")
+        arr = _as_nonempty_2d(arr, np.float32, "arr")
         nrows, n_per_row = arr.shape
         block_size = self.blck_size(type_num)
         if block_size <= 0:
@@ -764,10 +805,12 @@ class HipQuant:
 
     def quantize_numpy_to(self, arr, type_num, dst, imatrix=None, hq2_iterations=4,
                           aq2_iterations=8):
-        """Quantize into a pre-allocated uint8 buffer. Modifies dst in-place."""
-        arr = np.ascontiguousarray(arr, dtype=np.float32)
-        if arr.ndim == 1:
-            arr = arr.reshape(1, -1)
+        """Quantize into a pre-allocated uint8 buffer. Modifies dst in-place.
+
+        imatrix may be a full matrix matching arr or a per-column vector
+        shared by every row, as used by llama.cpp.
+        """
+        arr = _as_nonempty_2d(arr, np.float32, "arr")
         nrows, n_per_row = arr.shape
         block_size = self.blck_size(type_num)
         if block_size <= 0:
@@ -787,12 +830,11 @@ class HipQuant:
             )
         src_ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
         dst_ptr = dst.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
-        im_ptr = None
-        if imatrix is not None:
-            imatrix = np.ascontiguousarray(imatrix, dtype=np.float32)
-            if imatrix.shape != arr.shape:
-                raise ValueError(f"imatrix shape {imatrix.shape} != arr shape {arr.shape}")
-            im_ptr = imatrix.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+        imatrix = _prepare_imatrix(imatrix, arr.shape, type_num)
+        im_ptr = (
+            imatrix.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+            if imatrix is not None else None
+        )
         if int(type_num) in {GGML_TYPE["AQ2"], GGML_TYPE["AQ2_QK"], GGML_TYPE["AQ2_VO"]}:
             if not 1 <= int(aq2_iterations) <= 16:
                 raise ValueError("aq2_iterations must be in [1, 16]")
@@ -837,16 +879,14 @@ class HipQuant:
                      or the output of quantize_numpy(..., GGML_TYPE["F8_E4M3"])
                      / quantize_numpy(..., GGML_TYPE["F8_E5M2"]).
             type_num: GGML type number for the output format.
-            imatrix: Optional float32 importance matrix (same logical shape).
+            imatrix: Optional full matrix or per-column vector shared by rows.
             source_format: "E4M3"/GGML_TYPE["F8_E4M3"] or
                            "E5M2"/GGML_TYPE["F8_E5M2"].
 
         Returns:
             np.uint8 array of quantized bytes.
         """
-        arr_fp8 = np.ascontiguousarray(arr_fp8, dtype=np.uint8)
-        if arr_fp8.ndim == 1:
-            arr_fp8 = arr_fp8.reshape(1, -1)
+        arr_fp8 = _as_nonempty_2d(arr_fp8, np.uint8, "arr_fp8")
         nrows, n_per_row = arr_fp8.shape
         blck = self.blck_size(type_num)
         if blck <= 0:
@@ -859,20 +899,11 @@ class HipQuant:
         dst = np.empty(out_nbytes, dtype=np.uint8)
         src_ptr = arr_fp8.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
         dst_ptr = dst.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
-        im_ptr = None
-        if imatrix is not None:
-            imatrix = np.ascontiguousarray(imatrix, dtype=np.float32)
-            # The native kernel indexes imatrix + row*n_per_row for every row,
-            # i.e. it reads exactly arr_fp8.size floats. Require at least that
-            # many (a short buffer is a host hipMemcpy + device OOB read), then
-            # reshape so the row indexing is well defined.
-            if imatrix.size != arr_fp8.size:
-                raise ValueError(
-                    f"imatrix has {imatrix.size} elements but FP8 input has "
-                    f"{arr_fp8.size} ({arr_fp8.shape})"
-                )
-            imatrix = imatrix.reshape(arr_fp8.shape)
-            im_ptr = imatrix.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+        imatrix = _prepare_imatrix(imatrix, arr_fp8.shape, type_num)
+        im_ptr = (
+            imatrix.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+            if imatrix is not None else None
+        )
         if isinstance(source_format, str):
             source_format = source_format.upper()
         if source_format in ("E4M3", "F8_E4M3", GGML_TYPE["F8_E4M3"]):