hip-quant 2.2.0__tar.gz → 2.2.1__tar.gz

{hip_quant-2.2.0/hip_quant.egg-info → hip_quant-2.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hip-quant
-Version: 2.2.0
+Version: 2.2.1
 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
@@ -219,6 +219,8 @@ License-File: LICENSE
 Requires-Dist: numpy
 Provides-Extra: torch
 Requires-Dist: torch>=2.0; extra == "torch"
+Provides-Extra: torch215
+Requires-Dist: torch==2.15.0a0+gitf07882e; extra == "torch215"
 Dynamic: license-file
 
 <div align="center">
@@ -237,6 +239,18 @@ Dynamic: license-file
 
 `hip-quant` is a standalone Python library and highly optimized HIP C++ backend for tensor quantization and accelerated LLM inference on AMD GPUs with zero CPU round-trips.
 
+## 🌟 What's New in 2.2.1
+
+> **2.2.1.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14). See *Custom PyTorch 2.15* below for a self-hosted install index.
+
+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.
+- **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.
+
 ## 🌟 What's New in 2.2.0
 
 > **2.2.0.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14). Use it if `2.2.0` reports a torch ABI mismatch on import.
@@ -545,19 +559,55 @@ the next replay.
 ## 📦 Installation
 
 ```powershell
-# Binary wheel with packaged ROCm 7.2.1 ctypes DLL and PyTorch extension
-pip install dist/hip_quant-1.3.0-cp312-cp312-win_amd64.whl
+# From PyPI: binary wheel with the packaged ROCm 7.2.1 ctypes DLL and PyTorch extension
+pip install hip-quant
 
-# With PyTorch optional dependency declared
+# With the PyTorch optional dependency declared
 pip install "hip-quant[torch]"
+
+# Or a local wheel
+pip install dist/hip_quant-2.2.0-cp312-cp312-win_amd64.whl
+```
+
+Two wheels are published per release, one per PyTorch ABI. Install the one that
+matches your PyTorch, so the extension loads without an ABI mismatch:
+
+| Wheel | PyTorch | ROCm |
+|---|---|---|
+| `hip-quant==2.2.1` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.1.post215` | 2.15 / TheRock | ROCm 7.14 |
+
+```powershell
+pip install "hip-quant==2.2.1.post215"   # PyTorch 2.15 / TheRock
 ```
 
+#### Custom PyTorch 2.15 (self-hosted wheels)
+
+The `post215` wheel targets a custom PyTorch 2.15 build that is not on PyPI.
+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.1.post215" `
+  --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
+```
+
+This resolves `torch==2.15.0a0+gitf07882e` (plus matching `torchvision` /
+`torchaudio`). `gpu-smi` is bundled in the wheel, so monitoring needs no extra
+step.
+
+The wheels can also be fetched directly, e.g.
+`https://dl.hipquant.download/file/Torchs/torch-2.15.0a0%2Bgitf07882e-cp312-cp312-win_amd64.whl`
+— the `+` **must** be percent-encoded as `%2B` (a literal `+` 404s on B2).
+Because wheel filenames must match their metadata version, the objects are
+named `torch-2.15.0a0+gitf07882e-…`, not `torch-2.15.0-…`.
+
 On Windows, DLL resolution order is:
 - `HIP_QUANT_DLL` or `HIP_QUANT_DLL_PATH`, if set
 - `hip_quantize_rocm721.dll`
 - `hip_quantize.dll`
 
-Runtime DLL directories include `HIP_QUANT_ROCM_BIN`, `HIP_QUANT_ROCM_HOME`, `ROCM_HOME`, `ROCM_PATH`, `HIP_PATH`, the active venv's `_rocm_sdk_core\bin`, `torch\lib`, `Scripts`, then the system ROCm 7.1 path.
+Runtime DLL directories include `HIP_QUANT_ROCM_BIN`, `HIP_QUANT_ROCM_HOME`, `ROCM_HOME`, `ROCM_PATH`, `HIP_PATH`, the active venv's `_rocm_sdk_core\bin`, `torch\lib`, `Scripts`, then the system ROCm 7.2 path.
 
 ---
 
@@ -820,6 +870,20 @@ Block-wise scaling is useful when a tensor has uneven dynamic range across its
 last dimension. It usually reduces FP8 quantization error compared with one
 global scale for the entire tensor. Existing per-tensor FP8 APIs remain unchanged.
 
+#### Dynamic Per-Token FP8 Quantization
+
+`quantize_e4m3_per_token` computes a single FP32 dequant scale for an entire row
+(`amax / 448`). This is the `activation_scheme: "dynamic"` used by
+DeepSeek/vLLM FP8 checkpoints. For an input `[..., K]` the scale tensor has
+shape `[...]` — one value per row:
+
+```python
+from hip_quant.torch_api import quantize_e4m3_per_token
+
+x_fp8, x_scale = quantize_e4m3_per_token(x)   # x_scale: [M],  real ~= fp8 * scale[m]
+```
+
+
 #### Block-scaled Linear and Adafactor Kernel Helpers
 
 Two lower-level training helpers are available for experiments and future fused
@@ -848,12 +912,51 @@ out = fp8_linear_forward_blockwise_quantized(
 )
 ```
 
-The current block-scaled linear kernel is correctness-first and intentionally
-does not use gfx12 WMMA yet. It validates the `FP8 bytes + per-block scales`
-layout and math before replacing the inner loop with a tiled/WMMA or rocBLASLt
-implementation. `Adafactor` is a full Python optimizer step with factored
-second-moment state; GPU-side row/column mean-square helpers exist, while a
-fully fused Adafactor update kernel remains a future optimization.
+#### 2-D Block-Scaled FP8 Linear (vLLM / DeepSeek)
+
+`fp8_linear_forward_blockwise_2d` consumes a safetensors `weight_scale_inv`
+directly: a `[N, K]` E4M3 weight plus one dequant scale per `block x block`
+tile (typically `128 x 128`), i.e. a `[N/block, K/block]` FP32 tensor, paired
+with a single per-token activation scale `[M]`:
+
+```text
+C[m,n] = input_scale[m] * sum_kb  W_scale[n/block, kb] * <A[m,:], B[n,:]>_kb
+```
+
+```python
+from hip_quant import fp8_linear_forward_blockwise_2d, fp8_linear_deepseek
+
+# From raw FP8 + 2-D scale tensors (already on device)
+out = fp8_linear_forward_blockwise_2d(
+    input_fp8, input_scale,        # quantize_e4m3_per_token(x)
+    weight_fp8, weight_scale_inv,  # [N, K] uint8, [N/128, K/128] float32
+    output_dtype_source=x,
+    block_size=128, backend="auto", bias=bias,
+)
+
+# One-shot: dynamically quantizes x per token, then runs the fused GEMM
+out = fp8_linear_deepseek(x, weight_fp8, weight_scale_inv, bias=bias)
+```
+
+`backend` selects the kernel at runtime:
+
+| backend | kernel | notes |
+|---|---|---|
+| `"auto"` (default) | gfx12 WMMA when the shape allows, else portable | recommended |
+| `"wmma"` | gfx12 `v_wmma_f32_16x16x16_fp8_fp8`, one weight-tile scale folded per 128-K stage | gfx12 only |
+| `"portable"` | LDS-tiled `64x64x32` microkernel, exact fp8->fp16 shared-memory decode | any AMD GPU, no matrix intrinsics |
+| `"reference"` | naive per-element kernel | slow; testing only |
+
+The per-token activation scale is constant across K, so it is factored out of
+the accumulation; only the weight tile scale is applied per 128-K block. Scales
+are dequant scales (`real ~= fp8 * scale`), matching `weight_scale_inv`, so a
+checkpoint can be fed verbatim with no re-quantization of the weights.
+
+The 1-D `fp8_linear_forward_blockwise_quantized` kernel above is correctness-first
+and does **not** use gfx12 WMMA; the 2-D path does (see the backend table).
+`Adafactor` is a full Python optimizer step with factored second-moment state;
+GPU-side row/column mean-square helpers exist, while a fully fused Adafactor
+update kernel remains a future optimization.
 
 #### Fake-FP8 Linear (autograd-safe, Phase 3)
 
@@ -964,6 +1067,9 @@ from hip_quant import (
     fp8_linear_forward_scaled,
     fp8_linear_forward_fp8_weight,
     fp8_linear_forward_blockwise,
+    fp8_linear_forward_blockwise_quantized,
+    fp8_linear_forward_blockwise_2d,
+    fp8_linear_deepseek,
     fp8_linear_backward_input,
     fp8_linear_backward_input_scaled,
     fp8_linear_backward_weight,
@@ -983,6 +1089,9 @@ grad_wt_s  = fp8_linear_backward_weight_scaled(grad_output, input, input_scale)
 
 # Correctness-first block-scaled FP8 path, no WMMA requirement
 out_block = fp8_linear_forward_blockwise(input, weight, bias, block_size=32)
+
+# 2-D (128x128 tile) block-scaled FP8 path from a safetensors weight_scale_inv
+out_2d = fp8_linear_deepseek(input, weight_fp8, weight_scale_inv, bias=bias)
 ```
 
 These functions are also used by `Fp8Linear`, `Fp8ScaledLinear`, and
@@ -1061,10 +1170,9 @@ All PyTorch extension functions are guarded against:
 
 ## 🧪 Running Tests
 
-#### Math tests (no GPU required)
+#### Extension FP8 tests (GPU required)
 ```powershell
-python tests/torch/test_math_fp8.py
-# 90/90 pass — validated against ml_dtypes reference
+& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/torch/test_fp8.py -v
 ```
 
 #### Full pipeline tests (CPU mock, no GPU required)
@@ -1128,20 +1236,27 @@ $env:HIP_QUANT_BUILD_TORCH_EXT = "1"
 Check the artifacts:
 ```powershell
 & "C:\venvs\medusa_rocm\Scripts\python.exe" -m twine check `
-  "dist\hip_quant-1.3.0-cp312-cp312-win_amd64.whl" `
-  "dist\hip_quant-1.3.0.tar.gz"
+  "dist\hip_quant-2.2.0-cp312-cp312-win_amd64.whl" `
+  "dist\hip_quant-2.2.0.tar.gz"
 ```
 
 Upload to PyPI:
 ```powershell
 & "C:\venvs\medusa_rocm\Scripts\python.exe" -m twine upload `
-  "dist\hip_quant-1.3.0-cp312-cp312-win_amd64.whl" `
-  "dist\hip_quant-1.3.0.tar.gz"
+  "dist\hip_quant-2.2.0-cp312-cp312-win_amd64.whl" `
+  "dist\hip_quant-2.2.0.tar.gz"
 ```
 
-Do not upload stale universal wheels such as `hip_quant-1.3.0-py3-none-any.whl`.
+Do not upload stale universal wheels such as `hip_quant-2.2.0-py3-none-any.whl`.
 The Windows wheel is intentionally platform-tagged because it contains DLLs.
 
+Each release publishes two versions from the same sources — `X.Y.Z` (PyTorch
+2.9.x / ROCm 7.2) and `X.Y.Z.post215` (PyTorch 2.15 / TheRock ROCm 7.14). Build
+each in its own venv: `setup_torch.py` picks the ROCm 7.14/TheRock toolchain for
+a `torch.version.hip` of `7.14.*` and `C:\Program Files\AMD\ROCm\<major.minor>`
+otherwise (see the toolchain matrix in `AGENTS.md`). Never build the 2.9.x
+artifact against TheRock, or vice versa.
+
 Suggested release order:
 - Build and run `twine check`
 - Upload to TestPyPI or PyPI
@@ -1154,9 +1269,13 @@ Suggested release order:
 
 ```
 hip_quant/
-├── __init__.py              # NumPy / ctypes offline API
+├── __init__.py              # NumPy / ctypes offline API + lazy export table
 ├── __main__.py              # CLI entry point
-├── torch_api.py             # PyTorch FP8 training API (Phases 1–4)
+├── torch_api.py             # PyTorch FP8 / GEMM / attention API (Phases 1-4)
+├── gguf.py                  # Pure-Python GGUF parser
+├── gguf_loader.py           # Streaming GGUF -> VRAM tensor loader
+├── wave_helpers.py          # Drop-in wave attention helpers
+├── smi.py                   # gpu-smi wrapper
 ├── device_info.py           # GPU/DLL compatibility probe helpers
 ├── cdna_compat.py           # CDNA feature table, build configs, CPU refs
 ├── setup_torch.py           # PyTorch C++ extension build script
@@ -1167,17 +1286,23 @@ hip_quant/
 ├── kernels/                 # Per-format offline HIP kernels (.cu)
 ├── torch_ext/               # PyTorch extension source
 │   ├── pytorch_bindings.cpp # C++ bindings (TORCH_CHECK, pybind11)
-│   ├── fp8_quant_kernels.hip# Element-wise quant/dequant kernels
-│   ├── mxfp8_kernels.hip    # MXFP8 UE8M0 quant/dequant (OCP, 32-thread warp)
-│   └── fp8_linear_kernels.hip# Tiled FP8 GEMM kernels
-└── tests/torch/             # GPU test suite (pytest)
+│   ├── fp8_quant_kernels.hip       # Element-wise + per-token + block-wise quant
+│   ├── fp8_linear_kernels.hip      # FP8 GEMM: naive + tiled portable + 2-D block scales
+│   ├── fp8_linear_kernels_v2.hip   # gfx12 LDS-staged FP8 WMMA GEMM (incl. 2-D block scales)
+│   ├── mxfp8_kernels.hip           # MXFP8 UE8M0 quant/dequant (OCP, 32-thread warp)
+│   ├── gemv_q_kernels.hip          # GGUF AOT GEMV (K-/I-quants)
+│   ├── hq2_linear_kernels.hip      # HQ2 learned-codebook formats
+│   ├── hq3_linear_kernels.hip      # HQ3 learned-codebook formats
+│   ├── ssm_kernels.hip             # SSM / DeltaNet / fused norms
+│   └── wave_attn*.hip              # WaveAttention FP8 WMMA prefill/decode/backward
+└── tests/                   # pytest suite (CPU pipeline in tests/test_pipeline.py)
 ```
 
 ---
 
 ## 📋 Architecture Notes
 
-- **RDNA4 PyTorch target** — FP8 WMMA extension kernels are compiled with `--offload-arch=gfx1200` and `--offload-arch=gfx1201`
+- **RDNA4 PyTorch target** — release `_C` builds are a multi-arch fatbin for `gfx90a`, `gfx942`, RDNA3 `gfx1100`-`gfx1103`, and RDNA4 `gfx1200`/`gfx1201` (`HIP_QUANT_ARCH` overrides this; a `gfx1201`-only build is the fast local-dev option). The gfx12 FP8 WMMA kernels require a `gfx1200`/`gfx1201` device at runtime.
 - **Default offline DLL target** — `build.ps1` compiles the portable DLL quantization kernels for `gfx90a`, `gfx942`, RDNA3 `gfx1100`-`gfx1103`, and RDNA4 `gfx1200`/`gfx1201`
 - **Current validation scope** — runtime-tested locally on `gfx1201` RX 9070 XT; `gfx1200` and CDNA code objects are build-validated and need separate hardware runtime validation
 - **BF16/FP16 PyTorch support** — FP8 quantization and linear kernels accept FP32, FP16, and BF16 tensors, accumulating in FP32 registers and storing results in the input/master dtype

{hip_quant-2.2.0 → hip_quant-2.2.1}/README.md

@@ -14,6 +14,18 @@
 
 `hip-quant` is a standalone Python library and highly optimized HIP C++ backend for tensor quantization and accelerated LLM inference on AMD GPUs with zero CPU round-trips.
 
+## 🌟 What's New in 2.2.1
+
+> **2.2.1.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14). See *Custom PyTorch 2.15* below for a self-hosted install index.
+
+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.
+- **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.
+
 ## 🌟 What's New in 2.2.0
 
 > **2.2.0.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14). Use it if `2.2.0` reports a torch ABI mismatch on import.
@@ -322,19 +334,55 @@ the next replay.
 ## 📦 Installation
 
 ```powershell
-# Binary wheel with packaged ROCm 7.2.1 ctypes DLL and PyTorch extension
-pip install dist/hip_quant-1.3.0-cp312-cp312-win_amd64.whl
+# From PyPI: binary wheel with the packaged ROCm 7.2.1 ctypes DLL and PyTorch extension
+pip install hip-quant
 
-# With PyTorch optional dependency declared
+# With the PyTorch optional dependency declared
 pip install "hip-quant[torch]"
+
+# Or a local wheel
+pip install dist/hip_quant-2.2.0-cp312-cp312-win_amd64.whl
+```
+
+Two wheels are published per release, one per PyTorch ABI. Install the one that
+matches your PyTorch, so the extension loads without an ABI mismatch:
+
+| Wheel | PyTorch | ROCm |
+|---|---|---|
+| `hip-quant==2.2.1` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.1.post215` | 2.15 / TheRock | ROCm 7.14 |
+
+```powershell
+pip install "hip-quant==2.2.1.post215"   # PyTorch 2.15 / TheRock
 ```
 
+#### Custom PyTorch 2.15 (self-hosted wheels)
+
+The `post215` wheel targets a custom PyTorch 2.15 build that is not on PyPI.
+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.1.post215" `
+  --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
+```
+
+This resolves `torch==2.15.0a0+gitf07882e` (plus matching `torchvision` /
+`torchaudio`). `gpu-smi` is bundled in the wheel, so monitoring needs no extra
+step.
+
+The wheels can also be fetched directly, e.g.
+`https://dl.hipquant.download/file/Torchs/torch-2.15.0a0%2Bgitf07882e-cp312-cp312-win_amd64.whl`
+— the `+` **must** be percent-encoded as `%2B` (a literal `+` 404s on B2).
+Because wheel filenames must match their metadata version, the objects are
+named `torch-2.15.0a0+gitf07882e-…`, not `torch-2.15.0-…`.
+
 On Windows, DLL resolution order is:
 - `HIP_QUANT_DLL` or `HIP_QUANT_DLL_PATH`, if set
 - `hip_quantize_rocm721.dll`
 - `hip_quantize.dll`
 
-Runtime DLL directories include `HIP_QUANT_ROCM_BIN`, `HIP_QUANT_ROCM_HOME`, `ROCM_HOME`, `ROCM_PATH`, `HIP_PATH`, the active venv's `_rocm_sdk_core\bin`, `torch\lib`, `Scripts`, then the system ROCm 7.1 path.
+Runtime DLL directories include `HIP_QUANT_ROCM_BIN`, `HIP_QUANT_ROCM_HOME`, `ROCM_HOME`, `ROCM_PATH`, `HIP_PATH`, the active venv's `_rocm_sdk_core\bin`, `torch\lib`, `Scripts`, then the system ROCm 7.2 path.
 
 ---
 
@@ -597,6 +645,20 @@ Block-wise scaling is useful when a tensor has uneven dynamic range across its
 last dimension. It usually reduces FP8 quantization error compared with one
 global scale for the entire tensor. Existing per-tensor FP8 APIs remain unchanged.
 
+#### Dynamic Per-Token FP8 Quantization
+
+`quantize_e4m3_per_token` computes a single FP32 dequant scale for an entire row
+(`amax / 448`). This is the `activation_scheme: "dynamic"` used by
+DeepSeek/vLLM FP8 checkpoints. For an input `[..., K]` the scale tensor has
+shape `[...]` — one value per row:
+
+```python
+from hip_quant.torch_api import quantize_e4m3_per_token
+
+x_fp8, x_scale = quantize_e4m3_per_token(x)   # x_scale: [M],  real ~= fp8 * scale[m]
+```
+
+
 #### Block-scaled Linear and Adafactor Kernel Helpers
 
 Two lower-level training helpers are available for experiments and future fused
@@ -625,12 +687,51 @@ out = fp8_linear_forward_blockwise_quantized(
 )
 ```
 
-The current block-scaled linear kernel is correctness-first and intentionally
-does not use gfx12 WMMA yet. It validates the `FP8 bytes + per-block scales`
-layout and math before replacing the inner loop with a tiled/WMMA or rocBLASLt
-implementation. `Adafactor` is a full Python optimizer step with factored
-second-moment state; GPU-side row/column mean-square helpers exist, while a
-fully fused Adafactor update kernel remains a future optimization.
+#### 2-D Block-Scaled FP8 Linear (vLLM / DeepSeek)
+
+`fp8_linear_forward_blockwise_2d` consumes a safetensors `weight_scale_inv`
+directly: a `[N, K]` E4M3 weight plus one dequant scale per `block x block`
+tile (typically `128 x 128`), i.e. a `[N/block, K/block]` FP32 tensor, paired
+with a single per-token activation scale `[M]`:
+
+```text
+C[m,n] = input_scale[m] * sum_kb  W_scale[n/block, kb] * <A[m,:], B[n,:]>_kb
+```
+
+```python
+from hip_quant import fp8_linear_forward_blockwise_2d, fp8_linear_deepseek
+
+# From raw FP8 + 2-D scale tensors (already on device)
+out = fp8_linear_forward_blockwise_2d(
+    input_fp8, input_scale,        # quantize_e4m3_per_token(x)
+    weight_fp8, weight_scale_inv,  # [N, K] uint8, [N/128, K/128] float32
+    output_dtype_source=x,
+    block_size=128, backend="auto", bias=bias,
+)
+
+# One-shot: dynamically quantizes x per token, then runs the fused GEMM
+out = fp8_linear_deepseek(x, weight_fp8, weight_scale_inv, bias=bias)
+```
+
+`backend` selects the kernel at runtime:
+
+| backend | kernel | notes |
+|---|---|---|
+| `"auto"` (default) | gfx12 WMMA when the shape allows, else portable | recommended |
+| `"wmma"` | gfx12 `v_wmma_f32_16x16x16_fp8_fp8`, one weight-tile scale folded per 128-K stage | gfx12 only |
+| `"portable"` | LDS-tiled `64x64x32` microkernel, exact fp8->fp16 shared-memory decode | any AMD GPU, no matrix intrinsics |
+| `"reference"` | naive per-element kernel | slow; testing only |
+
+The per-token activation scale is constant across K, so it is factored out of
+the accumulation; only the weight tile scale is applied per 128-K block. Scales
+are dequant scales (`real ~= fp8 * scale`), matching `weight_scale_inv`, so a
+checkpoint can be fed verbatim with no re-quantization of the weights.
+
+The 1-D `fp8_linear_forward_blockwise_quantized` kernel above is correctness-first
+and does **not** use gfx12 WMMA; the 2-D path does (see the backend table).
+`Adafactor` is a full Python optimizer step with factored second-moment state;
+GPU-side row/column mean-square helpers exist, while a fully fused Adafactor
+update kernel remains a future optimization.
 
 #### Fake-FP8 Linear (autograd-safe, Phase 3)
 
@@ -741,6 +842,9 @@ from hip_quant import (
     fp8_linear_forward_scaled,
     fp8_linear_forward_fp8_weight,
     fp8_linear_forward_blockwise,
+    fp8_linear_forward_blockwise_quantized,
+    fp8_linear_forward_blockwise_2d,
+    fp8_linear_deepseek,
     fp8_linear_backward_input,
     fp8_linear_backward_input_scaled,
     fp8_linear_backward_weight,
@@ -760,6 +864,9 @@ grad_wt_s  = fp8_linear_backward_weight_scaled(grad_output, input, input_scale)
 
 # Correctness-first block-scaled FP8 path, no WMMA requirement
 out_block = fp8_linear_forward_blockwise(input, weight, bias, block_size=32)
+
+# 2-D (128x128 tile) block-scaled FP8 path from a safetensors weight_scale_inv
+out_2d = fp8_linear_deepseek(input, weight_fp8, weight_scale_inv, bias=bias)
 ```
 
 These functions are also used by `Fp8Linear`, `Fp8ScaledLinear`, and
@@ -838,10 +945,9 @@ All PyTorch extension functions are guarded against:
 
 ## 🧪 Running Tests
 
-#### Math tests (no GPU required)
+#### Extension FP8 tests (GPU required)
 ```powershell
-python tests/torch/test_math_fp8.py
-# 90/90 pass — validated against ml_dtypes reference
+& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/torch/test_fp8.py -v
 ```
 
 #### Full pipeline tests (CPU mock, no GPU required)
@@ -905,20 +1011,27 @@ $env:HIP_QUANT_BUILD_TORCH_EXT = "1"
 Check the artifacts:
 ```powershell
 & "C:\venvs\medusa_rocm\Scripts\python.exe" -m twine check `
-  "dist\hip_quant-1.3.0-cp312-cp312-win_amd64.whl" `
-  "dist\hip_quant-1.3.0.tar.gz"
+  "dist\hip_quant-2.2.0-cp312-cp312-win_amd64.whl" `
+  "dist\hip_quant-2.2.0.tar.gz"
 ```
 
 Upload to PyPI:
 ```powershell
 & "C:\venvs\medusa_rocm\Scripts\python.exe" -m twine upload `
-  "dist\hip_quant-1.3.0-cp312-cp312-win_amd64.whl" `
-  "dist\hip_quant-1.3.0.tar.gz"
+  "dist\hip_quant-2.2.0-cp312-cp312-win_amd64.whl" `
+  "dist\hip_quant-2.2.0.tar.gz"
 ```
 
-Do not upload stale universal wheels such as `hip_quant-1.3.0-py3-none-any.whl`.
+Do not upload stale universal wheels such as `hip_quant-2.2.0-py3-none-any.whl`.
 The Windows wheel is intentionally platform-tagged because it contains DLLs.
 
+Each release publishes two versions from the same sources — `X.Y.Z` (PyTorch
+2.9.x / ROCm 7.2) and `X.Y.Z.post215` (PyTorch 2.15 / TheRock ROCm 7.14). Build
+each in its own venv: `setup_torch.py` picks the ROCm 7.14/TheRock toolchain for
+a `torch.version.hip` of `7.14.*` and `C:\Program Files\AMD\ROCm\<major.minor>`
+otherwise (see the toolchain matrix in `AGENTS.md`). Never build the 2.9.x
+artifact against TheRock, or vice versa.
+
 Suggested release order:
 - Build and run `twine check`
 - Upload to TestPyPI or PyPI
@@ -931,9 +1044,13 @@ Suggested release order:
 
 ```
 hip_quant/
-├── __init__.py              # NumPy / ctypes offline API
+├── __init__.py              # NumPy / ctypes offline API + lazy export table
 ├── __main__.py              # CLI entry point
-├── torch_api.py             # PyTorch FP8 training API (Phases 1–4)
+├── torch_api.py             # PyTorch FP8 / GEMM / attention API (Phases 1-4)
+├── gguf.py                  # Pure-Python GGUF parser
+├── gguf_loader.py           # Streaming GGUF -> VRAM tensor loader
+├── wave_helpers.py          # Drop-in wave attention helpers
+├── smi.py                   # gpu-smi wrapper
 ├── device_info.py           # GPU/DLL compatibility probe helpers
 ├── cdna_compat.py           # CDNA feature table, build configs, CPU refs
 ├── setup_torch.py           # PyTorch C++ extension build script
@@ -944,17 +1061,23 @@ hip_quant/
 ├── kernels/                 # Per-format offline HIP kernels (.cu)
 ├── torch_ext/               # PyTorch extension source
 │   ├── pytorch_bindings.cpp # C++ bindings (TORCH_CHECK, pybind11)
-│   ├── fp8_quant_kernels.hip# Element-wise quant/dequant kernels
-│   ├── mxfp8_kernels.hip    # MXFP8 UE8M0 quant/dequant (OCP, 32-thread warp)
-│   └── fp8_linear_kernels.hip# Tiled FP8 GEMM kernels
-└── tests/torch/             # GPU test suite (pytest)
+│   ├── fp8_quant_kernels.hip       # Element-wise + per-token + block-wise quant
+│   ├── fp8_linear_kernels.hip      # FP8 GEMM: naive + tiled portable + 2-D block scales
+│   ├── fp8_linear_kernels_v2.hip   # gfx12 LDS-staged FP8 WMMA GEMM (incl. 2-D block scales)
+│   ├── mxfp8_kernels.hip           # MXFP8 UE8M0 quant/dequant (OCP, 32-thread warp)
+│   ├── gemv_q_kernels.hip          # GGUF AOT GEMV (K-/I-quants)
+│   ├── hq2_linear_kernels.hip      # HQ2 learned-codebook formats
+│   ├── hq3_linear_kernels.hip      # HQ3 learned-codebook formats
+│   ├── ssm_kernels.hip             # SSM / DeltaNet / fused norms
+│   └── wave_attn*.hip              # WaveAttention FP8 WMMA prefill/decode/backward
+└── tests/                   # pytest suite (CPU pipeline in tests/test_pipeline.py)
 ```
 
 ---
 
 ## 📋 Architecture Notes
 
-- **RDNA4 PyTorch target** — FP8 WMMA extension kernels are compiled with `--offload-arch=gfx1200` and `--offload-arch=gfx1201`
+- **RDNA4 PyTorch target** — release `_C` builds are a multi-arch fatbin for `gfx90a`, `gfx942`, RDNA3 `gfx1100`-`gfx1103`, and RDNA4 `gfx1200`/`gfx1201` (`HIP_QUANT_ARCH` overrides this; a `gfx1201`-only build is the fast local-dev option). The gfx12 FP8 WMMA kernels require a `gfx1200`/`gfx1201` device at runtime.
 - **Default offline DLL target** — `build.ps1` compiles the portable DLL quantization kernels for `gfx90a`, `gfx942`, RDNA3 `gfx1100`-`gfx1103`, and RDNA4 `gfx1200`/`gfx1201`
 - **Current validation scope** — runtime-tested locally on `gfx1201` RX 9070 XT; `gfx1200` and CDNA code objects are build-validated and need separate hardware runtime validation
 - **BF16/FP16 PyTorch support** — FP8 quantization and linear kernels accept FP32, FP16, and BF16 tensors, accumulating in FP32 registers and storing results in the input/master dtype

{hip_quant-2.2.0 → hip_quant-2.2.1}/__init__.py

@@ -3,7 +3,7 @@ import numpy as np
 import os
 import sys
 
-__version__ = "2.2.0"
+__version__ = "2.2.1"
 
 
 _TORCH_EXPORTS = {
@@ -284,6 +284,13 @@ def _add_runtime_dll_dirs():
 def _shared_library_candidates():
     env_dll = os.environ.get("HIP_QUANT_DLL") or os.environ.get("HIP_QUANT_DLL_PATH")
     if env_dll:
+        # Only an absolute override is safe: a bare/relative name is resolved
+        # against the process CWD and can be shadowed by an attacker (DLL hijack).
+        if not os.path.isabs(env_dll):
+            raise ValueError(
+                "HIP_QUANT_DLL / HIP_QUANT_DLL_PATH must be an absolute path, got "
+                f"{env_dll!r}"
+            )
         yield env_dll
 
     win_names = ["hip_quantize_rocm721.dll", "hip_quantize.dll"]
@@ -389,7 +396,7 @@ class HipQuant:
             self._dequantize_tensor_to_fp8 = None
         self._dll.ggml_type_size_for.restype = ctypes.c_size_t
         self._dll.ggml_type_size_for.argtypes = [ctypes.c_int]
-        self._dll.ggml_blck_size_for.restype = ctypes.c_size_t
+        self._dll.ggml_blck_size_for.restype = ctypes.c_int
         self._dll.ggml_blck_size_for.argtypes = [ctypes.c_int]
         self._dll.ggml_row_size_for.restype = ctypes.c_size_t
         self._dll.ggml_row_size_for.argtypes = [ctypes.c_int, ctypes.c_int64]
@@ -855,6 +862,16 @@ class HipQuant:
         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))
         if isinstance(source_format, str):
             source_format = source_format.upper()