Metadata-Version: 2.4
Name: hip-quant
Version: 2.0.0.post215
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
                                   Version 2.0, January 2004
                                http://www.apache.org/licenses/
        
           TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
        
           1. Definitions.
        
              "License" shall mean the terms and conditions for use, reproduction,
              and distribution as defined by Sections 1 through 9 of this document.
        
              "Licensor" shall mean the copyright owner or entity authorized by
              the copyright owner that is granting the License.
        
              "Legal Entity" shall mean the union of the acting entity and all
              other entities that control, are controlled by, or are under common
              control with that entity. For the purposes of this definition,
              "control" means (i) the power, direct or indirect, to cause the
              direction or management of such entity, whether by contract or
              otherwise, or (ii) ownership of fifty percent (50%) or more of the
              outstanding shares, or (iii) beneficial ownership of such entity.
        
              "You" (or "Your") shall mean an individual or Legal Entity
              exercising permissions granted by this License.
        
              "Source" form shall mean the preferred form for making modifications,
              including but not limited to software source code, documentation
              source, and configuration files.
        
              "Object" form shall mean any form resulting from mechanical
              transformation or translation of a Source form, including but
              not limited to compiled object code, generated documentation,
              and conversions to other media types.
        
              "Work" shall mean the work of authorship, whether in Source or
              Object form, made available under the License, as indicated by a
              copyright notice that is included in or attached to the work
              (an example is provided in the Appendix below).
        
              "Derivative Works" shall mean any work, whether in Source or Object
              form, that is based on (or derived from) the Work and for which the
              editorial revisions, annotations, elaborations, or other modifications
              represent, as a whole, an original work of authorship. For the purposes
              of this License, Derivative Works shall not include works that remain
              separable from, or merely link (or bind by name) to the interfaces of,
              the Work and Derivative Works thereof.
        
              "Contribution" shall mean any work of authorship, including
              the original version of the Work and any modifications or additions
              to that Work or Derivative Works thereof, that is intentionally
              submitted to Licensor for inclusion in the Work by the copyright owner
              or by an individual or Legal Entity authorized to submit on behalf of
              the copyright owner. For the purposes of this definition, "submitted"
              means any form of electronic, verbal, or written communication sent
              to the Licensor or its representatives, including but not limited to
              communication on electronic mailing lists, source code control systems,
              and issue tracking systems that are managed by, or on behalf of, the
              Licensor for the purpose of discussing and improving the Work, but
              excluding communication that is conspicuously marked or otherwise
              designated in writing by the copyright owner as "Not a Contribution."
        
              "Contributor" shall mean Licensor and any individual or Legal Entity
              on behalf of whom a Contribution has been received by Licensor and
              subsequently incorporated within the Work.
        
           2. Grant of Copyright License. Subject to the terms and conditions of
              this License, each Contributor hereby grants to You a perpetual,
              worldwide, non-exclusive, no-charge, royalty-free, irrevocable
              copyright license to reproduce, prepare Derivative Works of,
              publicly display, publicly perform, sublicense, and distribute the
              Work and such Derivative Works in Source or Object form.
        
           3. Grant of Patent License. Subject to the terms and conditions of
              this License, each Contributor hereby grants to You a perpetual,
              worldwide, non-exclusive, no-charge, royalty-free, irrevocable
              (except as stated in this section) patent license to make, have made,
              use, offer to sell, sell, import, and otherwise transfer the Work,
              where such license applies only to those patent claims licensable
              by such Contributor that are necessarily infringed by their
              Contribution(s) alone or by combination of their Contribution(s)
              with the Work to which such Contribution(s) was submitted. If You
              institute patent litigation against any entity (including a
              cross-claim or counterclaim in a lawsuit) alleging that the Work
              or a Contribution incorporated within the Work constitutes direct
              or contributory patent infringement, then any patent licenses
              granted to You under this License for that Work shall terminate
              as of the date such litigation is filed.
        
           4. Redistribution. You may reproduce and distribute copies of the
              Work or Derivative Works thereof in any medium, with or without
              modifications, and in Source or Object form, provided that You
              meet the following conditions:
        
              (a) You must give any other recipients of the Work or
                  Derivative Works a copy of this License; and
        
              (b) You must cause any modified files to carry prominent notices
                  stating that You changed the files; and
        
              (c) You must retain, in the Source form of any Derivative Works
                  that You distribute, all copyright, patent, trademark, and
                  attribution notices from the Source form of the Work,
                  excluding those notices that do not pertain to any part of
                  the Derivative Works; and
        
              (d) If the Work includes a "NOTICE" text file as part of its
                  distribution, then any Derivative Works that You distribute must
                  include a readable copy of the attribution notices contained
                  within such NOTICE file, excluding those notices that do not
                  pertain to any part of the Derivative Works, in at least one
                  of the following places: within a NOTICE text file distributed
                  as part of the Derivative Works; within the Source form or
                  documentation, if provided along with the Derivative Works; or,
                  within a display generated by the Derivative Works, if and
                  wherever such third-party notices normally appear. The contents
                  of the NOTICE file are for informational purposes only and
                  do not modify the License. You may add Your own attribution
                  notices within Derivative Works that You distribute, alongside
                  or as an addendum to the NOTICE text from the Work, provided
                  that such additional attribution notices cannot be construed
                  as modifying the License.
        
              You may add Your own copyright statement to Your modifications and
              may provide additional or different license terms and conditions
              for use, reproduction, or distribution of Your modifications, or
              for any such Derivative Works as a whole, provided Your use,
              reproduction, and distribution of the Work otherwise complies with
              the conditions stated in this License.
        
           5. Submission of Contributions. Unless You explicitly state otherwise,
              any Contribution intentionally submitted for inclusion in the Work
              by You to the Licensor shall be under the terms and conditions of
              this License, without any additional terms or conditions.
              Notwithstanding the above, nothing herein shall supersede or modify
              the terms of any separate license agreement you may have executed
              with Licensor regarding such Contributions.
        
           6. Trademarks. This License does not grant permission to use the trade
              names, trademarks, service marks, or product names of the Licensor,
              except as required for reasonable and customary use in describing the
              origin of the Work and reproducing the content of the NOTICE file.
        
           7. Disclaimer of Warranty. Unless required by applicable law or
              agreed to in writing, Licensor provides the Work (and each
              Contributor provides its Contributions) on an "AS IS" BASIS,
              WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
              implied, including, without limitation, any warranties or conditions
              of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
              PARTICULAR PURPOSE. You are solely responsible for determining the
              appropriateness of using or redistributing the Work and assume any
              risks associated with Your exercise of permissions under this License.
        
           8. Limitation of Liability. In no event and under no legal theory,
              whether in tort (including negligence), contract, or otherwise,
              unless required by applicable law (such as deliberate and grossly
              negligent acts) or agreed to in writing, shall any Contributor be
              liable to You for damages, including any direct, indirect, special,
              incidental, or consequential damages of any character arising as a
              result of this License or out of the use or inability to use the
              Work (including but not limited to damages for loss of goodwill,
              work stoppage, computer failure or malfunction, or any and all
              other commercial damages or losses), even if such Contributor
              has been advised of the possibility of such damages.
        
           9. Accepting Warranty or Additional Liability. While redistributing
              the Work or Derivative Works thereof, You may choose to offer,
              and charge a fee for, acceptance of support, warranty, indemnity,
              or other liability obligations and/or rights consistent with this
              License. However, in accepting such obligations, You may act only
              on Your own behalf and on Your sole responsibility, not on behalf
              of any other Contributor, and only if You agree to indemnify,
              defend, and hold each Contributor harmless for any liability
              incurred by, or claims asserted against, such Contributor by reason
              of your accepting any such warranty or additional liability.
        
           END OF TERMS AND CONDITIONS
        
           Copyright 2026 Thenukegun10x
        
           Licensed under the Apache License, Version 2.0 (the "License");
           you may not use this file except in compliance with the License.
           You may obtain a copy of the License at
        
               http://www.apache.org/licenses/LICENSE-2.0
        
           Unless required by applicable law or agreed to in writing, software
           distributed under the License is distributed on an "AS IS" BASIS,
           WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
           See the License for the specific language governing permissions and
           limitations under the License.
Project-URL: Homepage, https://github.com/Thenukegun10x/HIP-Quant
Project-URL: Documentation, https://github.com/Thenukegun10x/HIP-Quant/blob/main/README.md
Project-URL: Source, https://github.com/Thenukegun10x/HIP-Quant
Project-URL: Issues, https://github.com/Thenukegun10x/HIP-Quant/issues
Keywords: quantization,gpu,hip,rocm,amd,fp8,wmma,attention,pytorch,ggml
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: C++
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == "torch"
Dynamic: license-file

<div align="center">
  <h1>🚀 hip-quant 2.0</h1>
  <p><b>High-Performance On-Device Tensor Quantization & Inference Engine for AMD GPUs</b></p>
  <p>
    <img alt="ROCm 7.2.1" src="https://img.shields.io/badge/ROCm-7.2.1-ED1C24?logo=amd"/>
    <img alt="RDNA4" src="https://img.shields.io/badge/RDNA4-gfx1200%20%7C%20gfx1201-blue"/>
    <img alt="RDNA3" src="https://img.shields.io/badge/RDNA3-gfx1100%20%7C%20gfx1101%20%7C%20gfx1102%20%7C%20gfx1103-0096FF"/>
    <img alt="CDNA" src="https://img.shields.io/badge/CDNA-gfx90a%20%7C%20gfx942-purple"/>
    <img alt="License" src="https://img.shields.io/badge/license-Apache%202.0-blue.svg"/>
    <img alt="Python 3.8+" src="https://img.shields.io/badge/python-3.8+-3776AB?logo=python&logoColor=white"/>
    <img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-2.9%2BROCm-EE4C2C?logo=pytorch"/>
  </p>
</div>

`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.0.0

`hip-quant` 2.0.0 expands the project from an offline quantization tool into a **complete high-performance ROCm inference and kernel acceleration suite** for AMD GPUs (specifically targeting RDNA4 `gfx1201`/`gfx1200`, RDNA3, and CDNA):

- **⚡ Native AOT GEMV Engine (400+ GB/s)**: Complete hardware-accelerated GEMV for GGUF formats (`Q4_0`, `Q8_0`, `Q2_K`..`Q6_K`, `IQ1_M`, `IQ2_XXS`, `IQ2_XS`, `IQ2_S`, `IQ3_XXS`, `IQ3_S`, `IQ4_XS`). Zero register spills with single-launch long-K paths.
- **🧠 Native State Space Model (SSM / DeltaNet) Suite**: Fused gating, fast 1D causal convolution, fused Delta-Net decode, and gated RMSNorm supporting shared and per-head layouts.
- **🌊 Full WaveAttention FP8 WMMA Suite**: Native GFX12 FP8 WMMA attention calling `v_wmma_f32_16x16x16_fp8_fp8` across prefill, decode, long context, and native autograd backward.
- **📦 Streaming Zero-Dependency GGUF Loader**: Native pure Python GGUF parser (`gguf.py`) and fast GPU tensor loader (`gguf_loader.py`) streaming gigabyte-scale weights directly to VRAM in seconds.
- **📊 Bundled `gpu-smi` v1.2.0**: Zero-overhead single-binary monitor for dedicated VRAM, OS pinned memory, VRAM junction temperatures, and per-process memory tracking.
- **📄 Apache 2.0 License**: Officially licensed under Apache 2.0.

It ships **two independent APIs** that can be used together or separately:


| API | Purpose | Requires |
|---|---|---|
| **NumPy / ctypes** (offline) | Offline GGUF-format quantization via packaged DLL | ROCm runtime, numpy |
| **PyTorch extension** (training) | GPU-resident FP8 training ops with full autograd | PyTorch 2.x + ROCm, built `_C` extension |

## WaveAttention: Native FP8 Flash Attention for RDNA4

`wave_attn` is a FP8 WMMA flash attention kernel that bypasses AOTriton entirely.
It calls `v_wmma_f32_16x16x16_fp8_fp8` directly — the native RDNA4 FP8 matrix
multiply instruction — instead of the 1200-instruction software FP8 decode path
that AOTriton v2 uses.

```python
import torch
import torch_api as hip_quant

q = torch.randn(4, 8, 256, 64, device="cuda")
k = torch.randn(4, 8, 256, 64, device="cuda")
v = torch.randn(4, 8, 256, 64, device="cuda")

out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
```

**Key features:**
- Native `v_wmma_f32_16x16x16_fp8_fp8` (real FP8 hardware, not emulated)
- `fast_expf` via `__builtin_exp2f` (replaces software exp, ~6x faster)
- Parallel K/V tile loads (split threads, 2x load throughput)
- Adaptive tile selection (Q_TILE 16-128, K_TILE 64-128)
- Online softmax with warp-level shuffle reduction
- Multi-split parallelism for low-occupancy configs
- Native INT4 WMMA path (`wave_attn_int4`) — Q/K 4-bit, V FP8, LUT unpack (new in 1.3.0)

**Performance — RX 9070 XT `gfx1201`, `ROCm 7.2.1` + `torch 2.9.1+rocm7.2.1`, `D=128` (measured `bench_attn.py` vs `F.scaled_dot_product_attention`/AOTriton):**

| Config (B,H,S,D) | Causal | `wave` fwd | `SDPA` fwd | fwd speedup | `wave` bwd | `SDPA` bwd | bwd speedup |
|---|:---:|---|---|---|---|---|---|
| 1,32,128,128 | N | 0.079 ms | 0.175 ms | **2.23×** | 0.177 ms | 0.561 ms | **3.18×** |
| 1,32,256,128 | N | 0.134 ms | 0.533 ms | **3.97×** | 0.318 ms | 1.517 ms | **4.78×** |
| 1,32,512,128 | N | 0.308 ms | 1.702 ms | **5.52×** | 0.845 ms | 5.108 ms | **6.04×** |
| 1,32,1024,128 | N | 1.148 ms | 6.431 ms | **5.60×** | 2.902 ms | 17.963 ms | **6.19×** |
| 1,32,2048,128 | N | 3.917 ms | 22.520 ms | **5.75×** | 10.686 ms | 68.469 ms | **6.41×** |
| 1,32,256,128 | Y | 0.139 ms | 0.370 ms | **2.66×** | 0.298 ms | 1.158 ms | **3.89×** |
| 1,32,1024,128 | Y | 0.792 ms | 4.422 ms | **5.59×** | 2.878 ms | 10.632 ms | **3.69×** |

Accuracy: FP8 `cos ~0.9986-0.999` vs FP32 reference, backward `dQ 0.9981 / dK 0.9986 / dV 0.9999`, stable across `S`.

> **Why not FA3?** FlashAttention-3 is `Hopper` (`sm90`) CUDA-only — no ROCm/HIP port and no `gfx12` codegen. The fair ROCm baseline is `PyTorch SDPA` which on ROCm routes to `AOTriton` (`Efficient Attention`). `wave_attn` beats that path by specializing for `gfx1200/1201` `WMMA` directly.

## HQ2 / HQ3 — experimental learned-codebook formats

`hq2` is a backend-neutral Python API for the HQ-family formats. The shipped
HQ2 stores each 256-weight block as four learned FP16 centroids plus a two-bit
selector per weight (72 bytes, **2.25 bpw**); HQ3 stores eight FP16 centroids
plus 3-bit selectors (112 bytes, **3.5 bpw**). Blocks have one stable
little-endian layout, so a tensor quantized on one backend decodes on another.
This 2.25 bpw scalar layout is the original format; the 2.06 bpw `HQ2V`
successor is described below and is not yet wired into this module.

```python
import numpy as np
import hq2

weights = np.random.randn(4096, 4096).astype(np.float32)
packed = hq2.quantize(weights, importance=np.abs(weights), backend="cpu")
assert packed.bits_per_weight == 2.25
restored = packed.dequantize()
```

> **Status: experimental — not competitive yet.** HQ2 is a blockwise *scalar*
> quantizer that sits within 0.02 dB of its Lloyd-Max optimum, so its gap to
> the vector-quantized GGML IQ family is structural, not a bug. In a controlled
> MLP-only study it **loses to `iq2_xxs` while spending more bits** (2.25 vs
> 2.0625 bpw, Δln PPL −0.179), and HQ3 (3.5 bpw) loses to `iq3_s`/`q3_k`
> (3.44 bpw). A full mixed HQ2/HQ3 policy also loses to Unsloth `IQ2_M` at
> equal size. See `Own Quant/HQ_LADDER_FINDINGS.md` and
> `Own Quant/HQ2_FINDINGS.md` for the measurements.

> **The packed kernels are a functional baseline, not a speed win.** The direct
> packed-weight ROCm kernels (`HQ2Linear`, `HQ3Linear`) decode correctly
> on-device, but the current one-output-row GEMV measures ~14–75× slower than
> decoded BF16 `F.linear` (1.8 tok/s on a hybrid Gemma 4 package; see
> `Own Quant/experimental/inference_speed/RESULTS.md`).

The encoder-side research that *does* pass the GGML rate–distortion frontier —
`HQ2V` (a 256-entry 4-D codebook, **2.0625 bpw**) and `HQ2VL` (HQ2V plus
LDLQ/GPTQ-style Hessian error feedback at **zero bit cost**) — is not yet
shipped as a codec: it needs a GPU encoder, a sequential layer-by-layer driver,
and held-out scoring before it is shippable (`Own Quant/HQ_LADDER_FINDINGS.md`).

The API and format remain useful for archive inspection and reference decode,
with backends `cpu`, `torch`, and `rocm`:

| Backend | Input | State |
|---|---|---|
| `cpu` | NumPy array | Portable reference; no GPU dependency. |
| `torch` | Torch tensor | Functional Torch implementation (CPU/CUDA/ROCm), preserves device residency. |
| `rocm` | NumPy array | Native HIP quantizer. |
| `vulkan` | — | Reserved name only; fails clearly instead of silently using another backend. |

Inspect the actual availability with `hq2.backend_status()`.

## Hardware Status

Runtime validation is currently on RDNA4. The PyTorch FP8 WMMA kernels target `gfx1200` and `gfx1201`; `gfx1200` is treated as the cut-down `gfx1201` die with the same relevant FP8 WMMA capabilities.

CDNA support is included for the offline NumPy/DLL quantization path and compatibility tooling. The default DLL build now emits one all-target DLL for `gfx90a`, `gfx942`, RDNA3 `gfx1100`-`gfx1103`, and RDNA4 `gfx1200`/`gfx1201`. The gfx12 WMMA FP8 GEMM test is intentionally disabled on CDNA; CDNA can support FP8/BF16 through MFMA/rocBLASLt-style paths, but not this RDNA4-specific gfx12 WMMA builtin path.

### gfx12 FP8 WMMA (Windows RDNA4)

The PyTorch FP8 training API (`Fp8Linear`, `Fp8ScaledLinear`, and
`Fp8ShadowLinear`) and the ctypes micro-GEMM use the gfx12 Wave32 FP8/BF8
WMMA intrinsics. They are enabled by default on a `gfx1200`/`gfx1201` device
with ROCm/HIP 7.2 or newer. Both the Python and native entry points reject
other architectures or older runtimes before launch.

Windows wheels prefer the packaged ROCm 7.2.1-built
`hip_quantize_rocm721.dll`; set `HIP_QUANT_DLL_VARIANT=legacy` only when the
legacy DLL is deliberately required. To explicitly suppress the custom WMMA
path on a compatible system:
```powershell
$env:HIP_QUANT_DISABLE_WMMA = "1"
```

Validated local test system:
- GPU: AMD Radeon RX 9070 XT, `gfx1201`, 16 GB VRAM
- CPU: AMD Ryzen 7 7800X3D, 8 cores / 16 threads
- RAM: 32 GB system memory
- OS/toolchain: Windows, Visual Studio 2022 Build Tools, ROCm installed at `C:\Program Files\AMD\ROCm\7.1`
- PyTorch venv: `C:\venvs\medusa_rocm\Scripts\python.exe`
- PyTorch: `2.9.1+rocm7.2.1`, HIP runtime: `7.2.53211-158bd99533`
- Full WMMA diagnostic/stress suite verified: native diagnostic, padded and uneven-K GEMMs, 512²/1024² matrices, and repeated launches. Every numeric check compares the GEMM result with the decoded E4M3 operands actually consumed by WMMA.

> **Note:** The offline NumPy/DLL quantization path does not use WMMA. The
> custom PyTorch and micro-GEMM WMMA paths require gfx12 and ROCm/HIP 7.2+.

---

## ⚡ Supported Quantization Formats

### 🔢 Standard & K-Quants (offline API)
- **Legacy:** `Q4_0`, `Q4_1`, `Q5_0`, `Q5_1`, `Q8_0`, `Q8_1`
- **K-Quants:** `Q2_K`, `Q3_K`, `Q4_K`, `Q5_K`, `Q6_K`, `Q8_K`
- **Binary / low-bit:** `Q1_0` (1.125 bpw), `Q2_0` (2.25 bpw)

### 🧠 I-Quants (Importance Matrix)
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
(`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.

### ⚖️ Ternary Quants
For models trained to be ternary (BitNet, TriLM):
- `TQ1_0` (1.69 bpw), `TQ2_0` (2.06 bpw)

### 🎯 Attention-calibrated & HQ (experimental)
- `AQ2`, `AQ2_QK`, `AQ2_VO` (2.25 bpw; HQ2 physical block, fit with
  attention-derived saliency)
- `HQ2` (2.25 bpw) — experimental; see the HQ2/HQ3 section above

### 🧪 FP8 / BF16 Formats (both APIs)
| Format | Layout | Use case |
|---|---|---|
| `F8_E4M3` | 1s·4e·3m, bias=7, max=±448, NaN only | Forward activations & weights |
| `F8_E5M2` | 1s·5e·2m, bias=15, max=±57344, ±Inf+NaN | Backward gradients |
| `BF16` | 1s·8e·7m IEEE bfloat16 | Checkpoint / weight storage (cast) |

#### 🧊 MXFP8 OCP Microscaling — true UE8M0 (new in 1.2.0, INT4 WMMA in 1.3.0)
`MXFP8_E4M3` / `MXFP8_E5M2` are OCP v1.0 microscaling: **1× UE8M0 (`2^(b-127)` power-of-two, `0xFF`=NaN) + 32× E4M3/E5M2** `8.25 bpw`, 1 warp (32 threads) per scale block, shuffle-reduced `amax → ceil(log2(amax/448|57344))+127`, RNE. No FP32-scale emulation — native `MX` wire `real = fp8 * 2^(scale-127)`. Optimized for `gfx12 w32` (`hip_quant_util.h`).

```python
import torch, hip_quant.torch_api as hq
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
```

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×`).

Default FP8 quantization uses OCP standard semantics with round-to-nearest-even. The PyTorch extension also exposes opt-in stochastic E5M2 rounding for backward gradients.

---

## 🛠️ Build

### Offline DLL (NumPy API)
Default build emits one DLL for CDNA, RDNA3, and RDNA4 targets. By default it uses `C:\Program Files\AMD\ROCm\7.1\bin\hipcc.exe`; pass `-RocmBin` to use a ROCm/PyTorch venv toolchain:
```powershell
.\build.ps1

# Build the packaged ROCm 7.2.1 DLL from a PyTorch ROCm venv
.\build.ps1 -Output hip_quantize_rocm721.dll -RocmBin "C:\venvs\medusa_rocm\Scripts"

# Custom target set
.\build.ps1 -Arch "gfx942,gfx1200,gfx1201"
```

The build script adds `-mno-wavefrontsize64` so gfx12 `w32` WMMA code is compiled as Wave32.

### PyTorch Extension (`_C`)
Requires PyTorch with ROCm support (`torch 2.x+rocm`) and an **x64** MSVC
toolchain (`Hostx64\x64\link.exe`). Do not build from an x86 Developer shell —
that produces `temp.win32` objects and fails to link a 64-bit `_C.pyd`.

```powershell
# From "x64 Native Tools Command Prompt for VS 2022", or after vcvars64.bat:
& "C:\venvs\medusa_rocm\Scripts\python.exe" setup_torch.py build_ext --inplace
```

Expect `build\temp.win-amd64-cpython-312` and `rc.exe` /
`link.exe` from the Windows SDK x64 and MSVC Hostx64 paths.

To build a PyPI wheel that includes the compiled `_C.pyd` extension, build with
the extension flag from the ROCm/PyTorch environment:
```powershell
$env:HIP_QUANT_BUILD_TORCH_EXT = "1"
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m build --wheel --no-isolation
```

Without `HIP_QUANT_BUILD_TORCH_EXT=1`, `python -m build` creates a Windows wheel
that packages the ctypes DLLs but does not include `_C.pyd`. The PyTorch
extension can still be built locally with `setup_torch.py build_ext --inplace`.

---

### Runtime codebooks and GPU execution features

I-Quant lookup data is shipped as five versioned binary files under
`codebooks/`, rather than being compiled into the DLL. The loader finds them
next to the DLL; set `HIP_QUANT_CODEBOOK_DIR` when deliberately using a DLL
outside its package directory. Regenerate and verify the checked-in assets with:

```powershell
& "C:\venvs\medusa_rocm\Scripts\python.exe" tools\export_iq_codebooks.py --check
```

The PyTorch extension also exposes `quantize_e4m3_transpose()` and
`quantize_e5m2_transpose()` for a rank-2 tensor, avoiding a materialised
full-precision transpose. For stable-shape inference, `capture_hip_graph()`
captures a callable through PyTorch's ROCm `CUDAGraph` interface (HIP Graph):

```python
runner = hip_quant.capture_hip_graph(model, example_tokens)
logits = runner.replay(next_tokens)
```

`replay()` returns independent outputs by default; use
`clone_output=False` only when the graph-owned output will be consumed before
the next replay.

### Current release capabilities

- WMMA diagnostics use Wave32-correct launch shapes, cooperative-group tiled
  barriers for LDS staging, and decoded-FP8 GEMM references; the custom path
  is enabled on supported gfx12 ROCm/HIP 7.2+ systems.
- I-Quant quantizers are parallelized to 256 threads, and legacy quantizers
  use Wave32-safe shuffle reductions. I-/T-Quant dequantization to E4M3/E5M2
  and direct FP8-to-quant fused kernels are covered by byte-level tests.
- The PyTorch extension uses PyTorch's current HIP stream, supports seeded
  stochastic E5M2 conversion, non-temporal FP8 memory hints where supported,
  and fused FP8 quantize-plus-transpose kernels.
- I-Quant codebooks are external, versioned package assets verified by CRC;
  the captured HIP Graph runner provides stable-shape replay without Python
  launch overhead.

---

## 📦 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

# With PyTorch optional dependency declared
pip install "hip-quant[torch]"
```

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.

---

## 🐍 Usage

### Offline NumPy API

```python
import numpy as np
from hip_quant import quantize

weights = np.random.randn(4096, 4096).astype(np.float32)

# Quantize directly to Q4_K on the GPU — byte-exact match to llama.cpp
q4k_bytes = quantize(weights, type_num=12)  # 12 = Q4_K
```

#### FP8 (offline)
```python
from hip_quant import GGML_TYPE, get_hip_quant

hq = get_hip_quant()
x    = np.random.randn(4096, 4096).astype(np.float32)
grad = (np.random.randn(4096, 4096) * 128).astype(np.float32)

x_e4m3    = hq.quantize_numpy(x,    GGML_TYPE["F8_E4M3"])  # forward
grad_e5m2 = hq.quantize_numpy(grad, GGML_TYPE["F8_E5M2"])  # backward
```

#### BF16 checkpoint tensors (CPU decode)

The offline HIP DLL has GPU dequantizers for supported packed GGML Q formats;
BF16 is not one of those formats and does not need a HIP kernel. Use the
vectorized NumPy helper at the checkpoint/loader boundary instead. It converts
raw IEEE BF16 bit patterns to FP32 exactly, without loading the HIP DLL or the
PyTorch extension. The surrounding model loader remains responsible for its
usual BF16 weight placement/casting policy.

```python
import numpy as np
from hip_quant import bf16_to_fp32

# BF16 tensor bytes read from a checkpoint, in little-endian IEEE BF16 order.
raw_bf16 = np.frombuffer(checkpoint_bytes, dtype="<u2")
weights_fp32 = bf16_to_fp32(raw_bf16, shape=(4096, 4096))
```

Continue to use `HipQuant.dequantize_to_fp8(...)` for packed GGML Q tensors;
the BF16 helper intentionally does not route standard BF16 storage through the
quantized HIP dequantization path.

#### Q → FP8 dequantization (offline)
`dequantize_to_fp8` expands packed GGML Q blocks **directly** to raw FP8 bytes on
the GPU. Each thread reconstructs a scalar from the Q block and immediately
encodes it as E4M3 or E5M2 in the same kernel, so no float32 buffer is
allocated or transferred. Supported source types: legacy `Q4_0/Q4_1/Q5_0/Q5_1/
Q8_0/Q8_1`, K-quants `Q2_K` through `Q6_K`, I-quants `IQ1_S` through
`IQ4_XS` (including `IQ4_NL`), and ternary `TQ1_0`/`TQ2_0`.

```python
from hip_quant import GGML_TYPE, get_hip_quant

hq = get_hip_quant()
w  = np.random.randn(4096, 4096).astype(np.float32)

# First quantize to a narrow Q type, then expand straight to FP8 bytes
q4k = hq.quantize_numpy(w, GGML_TYPE["Q4_K"])          # 4-bit K-quant
e4m3 = hq.dequantize_to_e4m3(q4k, GGML_TYPE["Q4_K"], 4096)  # uint8 (4096,4096)
e5m2 = hq.dequantize_to_e5m2(q4k, GGML_TYPE["Q4_K"], 4096)  # uint8 (4096,4096)

# Generic form: pick output format at runtime
e4m3 = hq.dequantize_to_fp8(q4k, GGML_TYPE["Q4_K"], 4096, output_format="E4M3")
```

- The output array shape is `(nrows, n_per_row)` with one FP8 byte per logical
  element — same layout as `quantize_numpy(..., GGML_TYPE["F8_E4M3"])`.
- If the source type matches the requested FP8 format (`F8_E4M3`/`F8_E5M2`), the
  path short-circuits to a host byte copy and skips the GPU altogether.

#### CLI
```powershell
hip-quant --help
python -m hip_quant --help
```

---

### PyTorch Training API

> **Requires:** `python setup_torch.py build_ext --inplace` first.

#### Q → FP8 dequantization
`dequantize_q_to_fp8` expands packed GGML Q blocks **directly** to raw FP8 bytes on the GPU. It reads packed PyTorch GPU bytes and writes `torch.uint8` directly into another GPU tensor, eliminating expensive PCI-e transfers and host CPU decoding.

The PyTorch extension currently supports legacy `Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q8_1`
and K-quants `Q2_K` through `Q6_K`; I-quants and ternary quants remain
offline-DLL-only for this particular API.

```python
import torch
import hip_quant.torch_api as hq
from hip_quant import GGML_TYPE

# Packed Q4_K bytes, residing on the GPU
packed_q = torch.load("q4k_tensor.pt").cuda()

# Direct GPU-to-GPU expansion -> torch.uint8
e4m3 = hq.dequantize_q_to_fp8(packed_q, GGML_TYPE["Q4_K"], n_per_row=4096, e5m2=False)

# Shortcuts
e4m3 = hq.dequantize_q_to_e4m3(packed_q, GGML_TYPE["Q4_K"], 4096)
e5m2 = hq.dequantize_q_to_e5m2(packed_q, GGML_TYPE["Q4_K"], 4096)
```

The output shape is `[nrows, n_per_row]`.

#### MXFP4 compatibility bridge (FP8 hipBLASLt execution)

`dequantize_mxfp4_to_fp8` and `mxfp4_linear_forward` accept OCP MXFP4 E2M1
weights: two values per byte (even index in the low nibble) and one raw UE8M0
scale byte per contiguous block of 32 values. They keep the source weights in
that packed representation, expand the current layer to E4M3 on the GPU, and
use PyTorch's hipBLASLt-backed `torch._scaled_mm` path when available.

> **Warning:** This is an MXFP4 *emulation* path. RDNA4 does not execute
> native FP4 instructions here. The API emits a `RuntimeWarning` on first use;
> expect slower decode and higher transient VRAM use than native MXFP4. Set
> `HIP_QUANT_SUPPRESS_EMULATED_MXFP4_WARNING=1` only after acknowledging it.

```python
from hip_quant.torch_api import mxfp4_linear_forward

# values: [out_features, in_features // 2] uint8, two E2M1 nibbles per byte
# scales: [out_features, in_features // 32] uint8, UE8M0 exponents
output = mxfp4_linear_forward(activations, values, scales, bias)
```

#### Native CDNA4 MXFP4 (hipBLASLt)

`native_mxfp4_linear_forward` is a separate **true MXFP4** path for CDNA4
`gfx950` (MI350-class) hardware. It passes packed `HIP_R_4F_E2M1` operands
and their raw `HIP_R_8F_UE8M0` Vec32 scale bytes directly to hipBLASLt: there
is no FP8 expansion, software decoder, or custom WMMA kernel. It is therefore
not a Radeon fallback and will deliberately raise on RDNA GPUs (including the
RX 9070 XT) and earlier CDNA devices such as MI300.

```python
import torch
from hip_quant.torch_api import (
    native_mxfp4_capability,
    native_mxfp4_linear_forward,
)

status = native_mxfp4_capability()
if not status["available"]:
    raise RuntimeError(status["reason"])

# A values: [M, K // 2], B values: [N, K // 2], uint8 packed E2M1.
# Low nibble is the even logical value. Scales are raw UE8M0 bytes, one per
# contiguous 32 logical values: [M, K // 32] and [N, K // 32].
output = native_mxfp4_linear_forward(
    a_values, a_scales, b_values, b_scales, output_dtype=torch.bfloat16
)  # output is [M, N] = A @ B.T
```

The native ABI is intentionally narrow because it maps to hipBLASLt's native
MX contract: `M` and `N` must be multiples of 16, `K` a multiple of 128,
batch count is one, and fused bias/activation epilogues are unavailable. Add
bias or activation after this call. `native_mxfp4_contract()` exposes these
requirements as a testable dictionary, while `native_mxfp4_capability()` is
safe to call in CI without a CDNA GPU. The actual GPU test is opt-in:

```powershell
$env:HIP_QUANT_TEST_NATIVE_MXFP4='1'  # only on a gfx950 runner
& 'C:\venvs\medusa_rocm\Scripts\python.exe' -m pytest tests\torch\test_mxfp4_native_contract_torch.py -q
```

---

#### Element-wise FP8 quant / dequant (Phase 1 & 2)

```python
import torch
from hip_quant.torch_api import quantize_e4m3, dequantize_e4m3
from hip_quant.torch_api import quantize_e5m2, quantize_e5m2_stochastic, dequantize_e5m2

x = torch.randn(1024, 1024, device="cuda")  # stays on GPU the whole time

x_fp8  = quantize_e4m3(x)          # torch.uint8, same shape, same device
x_back = dequantize_e4m3(x_fp8)    # torch.float32, no CPU transfer

g_fp8 = quantize_e5m2_stochastic(x, seed=1234)  # reproducible stochastic E5M2
```

#### Stochastic E5M2 Backward Gradients

E5M2 has the range needed for backward gradients, but only two mantissa bits.
For tiny gradients, deterministic round-to-nearest-even can repeatedly flush or
bias values. `quantize_e5m2_stochastic()` rounds between adjacent E5M2 bins with
probability proportional to the input value's distance between those bins, using
a stateless per-element hash of `(seed, element_index)`.

Use it directly:
```python
from hip_quant.torch_api import quantize_e5m2_stochastic, dequantize_e5m2

grad_fp8 = quantize_e5m2_stochastic(grad, seed=42)
grad_sim = dequantize_e5m2(grad_fp8)
```

Enable stochastic E5M2 for FP8 linear backward `grad_output` quantization:
```powershell
$env:HIP_QUANT_STOCHASTIC_E5M2 = "1"

# Optional deterministic base seed for reproducible experiments
$env:HIP_QUANT_STOCHASTIC_E5M2_SEED = "1234"
```

This path is opt-in. It stochastic-quantizes `grad_output` once, dequantizes
those exact FP8 choices back to the training dtype, then reuses the existing
hipBLASLt/custom backward matrix kernels.

#### Block-wise FP8 Scaling

The PyTorch extension also exposes block-wise FP8 quantization. Values are stored
as raw FP8 bytes plus one FP32 dequant scale per block along the last dimension:

```text
real_value ~= fp8_value * fp32_block_scale
```

For an input shape `[..., K]`, the scale tensor has shape
`[..., ceil(K / block_size)]`.

```python
from hip_quant.torch_api import (
    quantize_e4m3_blockwise,
    quantize_e5m2_blockwise,
    quantize_e5m2_blockwise_stochastic,
    dequantize_e4m3_blockwise,
    refresh_fp8_blockwise_shadow,
)

x = torch.randn(8, 4096, device="cuda", dtype=torch.bfloat16)
grad = torch.randn_like(x)
weight = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)

# Forward activations/weights: E4M3 + FP32 per-block scales
x_fp8, x_scales = quantize_e4m3_blockwise(x, block_size=32)
x_back = dequantize_e4m3_blockwise(x_fp8, x_scales, block_size=32)

# Backward gradients: E5M2 + stochastic rounding + FP32 per-block scales
g_fp8, g_scales = quantize_e5m2_blockwise_stochastic(grad, block_size=32, seed=1234)

# Master weight -> block-wise FP8 shadow buffers
weight_fp8, weight_scales = refresh_fp8_blockwise_shadow(weight, block_size=32)
```

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.

#### Block-scaled Linear and Adafactor Kernel Helpers

Two lower-level training helpers are available for experiments and future fused
training paths:

```python
from hip_quant.torch_api import (
    adafactor_row_col_mean_square,
    fp8_linear_forward_blockwise,
    fp8_linear_forward_blockwise_quantized,
)

# GPU-side Adafactor 2-D statistics
row_ms, col_ms = adafactor_row_col_mean_square(grad_2d, eps=1e-30)

# Convenience path: quantize input/weight block-wise, then run block-scaled FP8 linear
out = fp8_linear_forward_blockwise(input, weight, bias=bias, block_size=32)

# Pre-quantized path: consumes FP8 bytes + FP32 scale tensors directly
out = fp8_linear_forward_blockwise_quantized(
    input_fp8, input_scales,
    weight_fp8, weight_scales,
    output_dtype_source=input,
    block_size=32,
    bias=bias,
)
```

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.

#### Fake-FP8 Linear (autograd-safe, Phase 3)

`Fp8LinearFunction` uses **E4M3** for forward activations/weights and **E5M2** for backward gradients. It accepts `torch.float32`, `torch.float16`, and `torch.bfloat16` inputs/weights. It also implements **Activation Compression**, saving `uint8` tensors in the autograd graph to cut activation VRAM by 4× versus FP32, and 2× versus FP16/BF16.

BF16/FP16 support applies to:
- `quantize_e4m3()` and `quantize_e5m2()` inputs
- `Fp8LinearFunction` forward/backward
- `Fp8Linear`, `Fp8ScaledLinear`, `Fp8ShadowLinear`, `Fp8Conv1d`, and `Fp8Conv2d` module parameters and gradients
- `Fp8ShadowLinear` master weights, so user-selected BF16/FP16 master weights reduce persistent parameter and gradient VRAM versus FP32

```python
from hip_quant.torch_api import convert_to_fp8, Adafactor

# Drop-in replacement for all nn.Linear layers in a model
model = MySmallLM(...)

# shadow=True: replaces nn.Linear with Fp8ShadowLinear
# Weights are stored as uint8 in memory, forward pass decompresses on the fly
# Cuts weight VRAM by 4×
convert_to_fp8(model, shadow=True, skip_names={"lm_head"})
model.cuda()

# Adafactor optimizer: adaptive learning rates with sublinear memory cost
# Cuts optimizer state VRAM by ~1000× compared to AdamW
opt = Adafactor(model.parameters(), relative_step=True)

# One nonfinite gradient skips the entire optimizer step (no state poison).
loss.backward()
opt.step()
if opt.last_step_skipped:
    # reduce loss scale / skip weight update for this step
    pass
```

Training-path numerical guards:
- hipBLASLt backward applies quant scales as *pre-quant multipliers* and uses
  the reciprocal as the GEMM dequant scale (avoids exploding `grad_input`).
- `Fp8TensorMeta.update` ignores NaN/Inf amax samples and keeps the last valid
  scale (`found_nonfinite` sticky flag).
- `Adafactor.step` preflights all gradients and sets `last_step_skipped=True`
  instead of mutating state when any gradient is nonfinite.
- Weight-gradient kernels can retain FP32 accumulation for master weights.
- `Fp8ShadowLinear` caches the FP8 weight shadow between optimizer updates;
  hipBLASLt reuses pre-quantized E4M3 bytes instead of re-casting masters.
- On CUDA/HIP tensors, `Fp8TensorMeta.quantize_e4m3_delayed(x)` streams the
  next amax reduction while emitting E4M3 activation bytes; that observation
  only selects the next call's scale. Shadow refresh writes directly into the
  persistent `weight_fp8` buffer, without a temporary uint8 tensor.
- `fp8_grouped_linear_forward_fp8_input(shared_input_fp8, weights_fp8, ...)`
  runs multiple projections (for example Q/K/V) from one shared raw-E4M3
  activation buffer. The custom backward materializes each E5M2 operand once
  and reuses those raw bytes across WMMA tiles.

#### FP8 Conv1d / Conv2d

`fp8_conv1d`, `fp8_conv2d`, `Fp8Conv1d`, and `Fp8Conv2d` lower convolution to
an unfold/im2col matrix multiply and reuse the same FP8 scaled linear backend.
That means hipBLASLt via PyTorch `torch._scaled_mm` is used first when
available, while the custom gfx12 WMMA path remains the fallback/testing path.

```python
import torch
from hip_quant.torch_api import Fp8Conv1d, Fp8Conv2d, fp8_conv1d, fp8_conv2d

x1 = torch.randn(8, 16, 1024, device="cuda", dtype=torch.bfloat16)
conv1 = Fp8Conv1d(16, 32, kernel_size=3, padding=1,
                  device="cuda", dtype=torch.bfloat16)
y1 = conv1(x1)
y1_func = fp8_conv1d(x1, conv1.weight, conv1.bias, padding=1)

x = torch.randn(8, 3, 224, 224, device="cuda", dtype=torch.bfloat16)
conv = Fp8Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
                 device="cuda", dtype=torch.bfloat16)

y = conv(x)

# Functional form mirrors torch.nn.functional.conv2d for numeric parameters.
y2 = fp8_conv2d(x, conv.weight, conv.bias, stride=2, padding=1)
```

Supported convolution options: numeric `stride`, `padding`, `dilation`, and
`groups` with zero padding mode. Inputs and weights must be CUDA/HIP tensors.

**Combined VRAM savings for a 500M-param LLM:**
Before: ~7.6 GB (Weights 2GB, Acts 1.6GB, AdamW 4GB)
After: ~0.9 GB (Weights 0.5GB, Acts 0.4GB, Adafactor 4MB)

#### Direct autograd.Function

```python
from hip_quant.torch_api import Fp8LinearFunction

out = Fp8LinearFunction.apply(input, weight, bias)  # bias optional
```

#### Fused FP8 Linear (gfx12 WMMA kernels)

The high-level `Fp8Linear`, `Fp8ScaledLinear`, `Fp8ShadowLinear`, `Fp8Conv1d`,
and `Fp8Conv2d` APIs try the hipBLASLt-backed PyTorch `_scaled_mm` route first.
These direct custom WMMA entry points are the fallback/testing path on a
compatible gfx12 ROCm/HIP 7.2+ device. Set `HIP_QUANT_DISABLE_WMMA=1` to
explicitly disable them.

```python
from hip_quant import (
    fp8_linear_forward,
    fp8_linear_forward_scaled,
    fp8_linear_forward_fp8_weight,
    fp8_linear_forward_blockwise,
    fp8_linear_backward_input,
    fp8_linear_backward_input_scaled,
    fp8_linear_backward_weight,
    fp8_linear_backward_weight_scaled,
)

# [M,K] @ [N,K].T = [M,N]
# forward: E4M3 x E4M3 WMMA, backward: E5M2/BF8 x E5M2/BF8 WMMA
out        = fp8_linear_forward(input, weight, bias=None)
grad_in    = fp8_linear_backward_input(grad_output, weight)
grad_wt    = fp8_linear_backward_weight(grad_output, input)

# Scaled path used by Fp8ScaledLinear and Fp8ShadowLinear
out_scaled = fp8_linear_forward_scaled(input, weight, bias, input_scale, weight_scale)
grad_in_s  = fp8_linear_backward_input_scaled(grad_output, weight, weight_scale)
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)
```

These functions are also used by `Fp8Linear`, `Fp8ScaledLinear`, and
`Fp8ShadowLinear` after the extension is built.

#### gfx1201 FP8/BF16 Microbenchmark

Measured on the validated local RX 9070 XT `gfx1201` system with PyTorch
`2.9.1+rocm7.2.1`:

```powershell
& "C:\venvs\medusa_rocm\Scripts\python.exe" tests\torch\bench_fp8.py
```

```text
Elementwise FP8 ops, shape=(4096, 4096), dtype=bf16
quantize_e4m3:   0.243 ms
quantize_e5m2:   0.205 ms
dequantize_e4m3: 0.218 ms
dequantize_e5m2: 0.206 ms
Fp8ShadowLinear, batch=32, in=4096, out=4096, dtype=bf16
forward:          2.553 ms
forward+backward: 5.622 ms
total wall time: 0.86 s
```

The benchmark is available at `tests/torch/bench_fp8.py`. It measures custom
WMMA automatically on a compatible device and reports the runtime guard reason
otherwise.

The FP8/BF16 kernel path is primarily a speed and memory-bandwidth
improvement: it reuses pre-quantized FP8 activations/gradients, skips redundant
output zeroing, fuses bias stores, vectorizes elementwise FP8 kernels, and caches
offline FP8 temporary buffers. Persistent VRAM savings are still mainly provided
by `Fp8ShadowLinear` FP8 weight shadows and activation compression; this path
reduces transient allocations and extra memory passes around those features.

RDNA3 (`gfx11`) and CDNA devices are rejected for this specific builtin path.
CDNA FP8/BF16 GEMM should use an MFMA/rocBLASLt implementation instead.

#### Scale / amax tracking (Phase 4 scaffold)

```python
from hip_quant.torch_api import Fp8TensorMeta

meta = Fp8TensorMeta(history_len=16, device="cuda")
meta.update(x)                   # records amax, updates scale/inv_scale

x_fp8  = meta.quantize_e4m3(x)  # scaled, then quantized
x_back = meta.dequantize_e4m3(x_fp8)  # dequantized, then rescaled
```

---

## 🔒 Memory Safety

All PyTorch extension functions are guarded against:
- Non-CUDA tensors (`TORCH_CHECK(is_cuda)`)
- Non-contiguous layout (`TORCH_CHECK(is_contiguous)`)
- Wrong dtype (`float32` / `float16` / `bfloat16` for floating inputs, `uint8` for FP8 buffers)
- Dimension mismatch for GEMM
- **`int64 → int` narrowing** — explicit `checked_int()` with `TORCH_CHECK`
- **Hardware grid limit** — `gridDim.y ≤ 65535` validated before launch
- **Cross-device pointers** — `input.device() == weight.device()` checked
- **Empty tensors** — `numel == 0` early-return before `dim3(0)` (UB in HIP)
- **Positive finite scales** — invalid FP8 scales raise before launch / GEMM
- **Nonfinite training step** — delayed scales and Adafactor refuse to poison state
- **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`)

---

## 🧪 Running Tests

#### Math tests (no GPU required)
```powershell
python tests/torch/test_math_fp8.py
# 90/90 pass — validated against ml_dtypes reference
```

#### Full pipeline tests (CPU mock, no GPU required)
```powershell
# Preferred: pure unittest (no GPU init)
& "C:\venvs\medusa_rocm\Scripts\python.exe" -c "import unittest, tests.test_pipeline as t; unittest.main(module=t, exit=True)"

# Or pytest
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/test_pipeline.py -q
```

`Fp8TensorMeta` and scaled/shadow modules keep delayed-scale metadata on the
parameter device (CPU by default). That avoids accidental ROCm init during the
mocked CPU suite. If a GPU-enabled run still stalls on process exit under
Windows, force CPU visibility:

```powershell
$env:CUDA_VISIBLE_DEVICES = ""
$env:HIP_VISIBLE_DEVICES = ""
```

#### PyTorch GPU tests
```powershell
# Build extension first (x64 VS toolchain)
& "C:\venvs\medusa_rocm\Scripts\python.exe" setup_torch.py build_ext --inplace
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/torch/test_fp8.py -v
```

#### Compatibility Tests (CPU + DLL)
```powershell
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/test_compat.py -v

# Device/compat reports
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m hip_quant --info
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m hip_quant --compat
```

#### Optional gfx12 FP8 WMMA Stress Test

Requires a ROCm/HIP 7.2+ gfx12 system.
```powershell
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests\test_wmma_diag.py -v

# Include 512²/1024² and repeated-large coverage
$env:HIP_QUANT_WMMA_STRESS = "1"
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests\test_wmma_diag.py -v
```

The release DLL was checked on `gfx1201` with HIP runtime `70253211`.

---

## 📤 Release / PyPI Upload

Build the distributables:
```powershell
$env:HIP_QUANT_BUILD_TORCH_EXT = "1"
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m build --no-isolation
```

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"
```

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"
```

Do not upload stale universal wheels such as `hip_quant-1.3.0-py3-none-any.whl`.
The Windows wheel is intentionally platform-tagged because it contains DLLs.

Suggested release order:
- Build and run `twine check`
- Upload to TestPyPI or PyPI
- Install the uploaded package in a clean venv and verify `HipQuant().dll_path` resolves to `hip_quantize_rocm721.dll`
- Commit/tag the exact source and DLL used for the PyPI upload

---

## 🗂️ Project Structure

```
hip_quant/
├── __init__.py              # NumPy / ctypes offline API
├── __main__.py              # CLI entry point
├── torch_api.py             # PyTorch FP8 training API (Phases 1–4)
├── 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
├── build.ps1                # DLL build script (hipcc)
├── hip_quantize.cpp         # Offline quantization kernels (DLL source)
├── hip_quant_util.h         # Shared FP8 / FP16 device helpers
├── hip_quant_types.h        # GGML block type definitions
├── 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)
```

---

## 📋 Architecture Notes

- **RDNA4 PyTorch target** — FP8 WMMA extension kernels are compiled with `--offload-arch=gfx1200` and `--offload-arch=gfx1201`
- **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
- **Device-resident kernels** — FP8 tensor data stays on device through `tensor.data_ptr()`. hipBLASLt training paths keep delayed scales device-resident where possible; legacy custom WMMA launchers still take scalar float scales. Block-wise FP8 metadata stays in device FP32 scale tensors. All extension launches use PyTorch's current HIP stream, including the Windows/MSVC build.
- **Phase 4 GEMM** includes gfx12 WMMA per-tensor-scale paths, packed-weight WMMA variants, and a correctness-first block-scaled FP8 linear path. Large training shapes prefer hipBLASLt via `torch._scaled_mm`; custom WMMA remains the fallback/small-shape path.
- **Adafactor** provides a complete optimizer step in Python with nonfinite step skipping. GPU-side row/column mean-square reductions for 2-D gradients exist; a fully fused Adafactor update kernel is a future optimization target.
- **Offline API unchanged** — the NumPy/ctypes path is untouched; both APIs coexist cleanly

---

## 📄 License
This project is licensed under the **Apache License, Version 2.0**. See the [LICENSE](LICENSE) file for the full license text.

