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

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

@@ -13,6 +13,8 @@ recursive-include torch_ext *.cpp *.hip *.h
 recursive-include tests *.py
 recursive-include codebooks *.bin
 include torch_api.py
+include wave_attn_train_triton.py
+include wave_attn_train_strided_triton.py
 include wave_helpers.py
 include smi.py
 include gguf.py

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hip-quant
-Version: 2.2.2
+Version: 2.2.3
 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
@@ -239,6 +239,34 @@ 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.
 
+> ⚠️ **Numerical / experimental status — WaveAttention training kernels.**
+> The WaveAttention **training** paths are experimental. The native FP8 path
+> (`precision="native"`, the default) is an inference-grade kernel: against a
+> BF16 MATH-SDPA reference its `dQ`/`dK` relative L2 is about **8% at input
+> scale 1 and ~58% at scale 5**, so it is *not* suitable for training as-is.
+> The `precision="fp16"` native WMMA kernels (new in 2.2.3) are the accurate
+> choice: against an fp32 SDPA reference the forward/backward relative L2 is
+> ~2e-4 / ~3e-4 at the tested shapes and input scales, and an adversarial
+> NaN/Inf + extreme-range harness passes clean. The `precision="bf16"` Triton
+> kernels are kept as a fallback. **Full-model convergence and end-to-end speed
+> remain unverified** (a 3,000-update PlantUnified continuation was finite,
+> without EMA/validation). A CPU float64 operand study
+> (`benchmarks/wave_attn_train_results.md`) found that no FP8 operand mixture —
+> including Hadamard-rotated, block-scaled FP8 — reaches a sub-1% training
+> 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.3
+
+> **2.2.3.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
+
+- **🌊 Native FP16 WMMA training attention (`gfx1201`).** `wave_attn(..., precision="fp16")` now runs a purpose-built FP16 flash-attention **forward and backward** on `v_wmma_f32_16x16x16_f16` — no Triton, and no SDPA recomputation for the backward. `dQ`/`dK`/`dV` are computed on-device with a FlashAttention-2 preprocess pass for `D_i = Σ_d dO_id·O_id`. Against the AOTriton SDPA baseline (`F.scaled_dot_product_attention`, flash backend) at `B1 H16`, the **forward is faster on all eight measured shapes** (1.35–2.13×; 60.9–68.4 vs 44.8–50.9 TFLOP/s at `S4096`). The **backward is still behind AOTriton** (1.47–2.06× slower; 14–22 vs 26–38 TFLOP/s): its `dK`/`dV` block needs one wave per 16-key sub-tile and ~53 KB of shared memory, which caps it at four waves per CTA and one CTA per CU.
+- **🐛 Three training-correctness bugs found and fixed**, by a new adversarial harness (`tests/torch/test_wave_attn_f16_extreme.py`: ~70 hand-built cases plus a log-uniform fuzz against an fp64 reference):
+  - **`dS` fp16 overflow → NaN.** `|dS| ≈ 3.41·|V|·|dO|` at `D=128`, so perfectly valid fp16 inputs (`V=65504`, `dO=1`) overflowed the `dS` store and produced NaN gradients. A per-`(b,h)` scale from the preprocess pass now bounds `dS`, and the output is divided back out; the `min(1, …)` cap keeps ordinary inputs bit-identical to the unscaled kernel.
+  - **Causal mask mismatch when `Seq_K > Seq_Q`.** The FP16 forward carried a bottom-right offset while the backward — and PyTorch's `is_causal` — are strictly top-left. The offset is gone, so forward and backward now agree.
+  - **`LSE` log-term loss.** Storing `row max + log(row sum)` as a single fp32 rounded the log term away once the max exceeded ~4e7, making `dV` exactly **128× too large** and turning a cancellation residual into a spurious `dQ` of 1.44e7 where the true value is exactly 0. `LSE` is now two fp32 per row, subtracted separately.
+- **🛡️ The forward can no longer overflow fp32.** Scores stay in the raw WMMA domain and `softmax_scale` is applied *inside* the exponent as `(raw − raw)·|scale|`. Forming `sacc·scale` was the forward's only fp32 overflow site (`|sacc| ≤ D·65504² = 5.5e11`, so `scale > 6.19e26` NaNs everything downstream) — and PyTorch's own fp32 SDPA and a textbook fp32 softmax **do** NaN at exactly that threshold. The native kernel now matches fp64 exactly there, while still agreeing with SDPA to ~2e-4 in the normal range. A negative `softmax_scale` (softmax of `−s`) is normalised out, so it is unaffected.
+
 ## 🌟 What's New in 2.2.2
 
 > **2.2.2.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
@@ -310,7 +338,15 @@ 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
+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):
+qbf, kbf, vbf = q.bfloat16(), k.bfloat16(), v.bfloat16()
+out_bf16 = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16")
+# Or defer to PyTorch SDPA entirely:
+out_precise = hip_quant.wave_attn(q, k, v, precision="sdpa")
 ```
 
 **Key features:**
@@ -322,7 +358,7 @@ out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
 - 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):**
+**Historical performance — RX 9070 XT `gfx1201`, `ROCm 7.2.1` + `torch 2.9.1+rocm7.2.1`, `D=128` (measured `bench_attn.py` against its AOTriton SDPA path):**
 
 | Config (B,H,S,D) | Causal | `wave` fwd | `SDPA` fwd | fwd speedup | `wave` bwd | `SDPA` bwd | bwd speedup |
 |---|:---:|---|---|---|---|---|---|
@@ -334,9 +370,44 @@ out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
 | 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.
+The table describes that benchmark's backend and shapes; it is not a current
+speed claim for every SDPA backend or training shape. On the same RX 9070 XT
+with `B=6,H=12,S=789,D=64` BF16 inputs, the full native FP8 training API
+takes **1.87 ms** forward plus backward, while `precision="sdpa"` takes
+**1.03 ms**. Against BF16 MATH SDPA, native dQ/dK relative L2 is about 8%
+at input scale 1 and 58% at scale 5; the SDPA precision mode is below 1%.
+See `benchmarks/wave_attn_train_results.md` for methods and all metrics.
+
+> The native FP8 path remains available with `precision="native"` (the default),
+> 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.
+
+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
+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
+ended at loss 7.7051 (BF16 WaveAttention) versus 7.7122 (BF16 SDPA). A
+full-model speed win and full-epoch convergence remain unverified.
 
 ## HQ2 / HQ3 — experimental learned-codebook formats
 
@@ -574,7 +645,7 @@ pip install hip-quant
 pip install "hip-quant[torch]"
 
 # Or a local wheel
-pip install dist/hip_quant-2.2.0-cp312-cp312-win_amd64.whl
+pip install dist/hip_quant-2.2.3-cp312-cp312-win_amd64.whl
 ```
 
 Two wheels are published per release, one per PyTorch ABI. Install the one that
@@ -582,11 +653,11 @@ matches your PyTorch, so the extension loads without an ABI mismatch:
 
 | Wheel | PyTorch | ROCm |
 |---|---|---|
-| `hip-quant==2.2.2` | 2.9.x | ROCm 7.2.1 |
-| `hip-quant==2.2.2.post215` | 2.15 / TheRock | ROCm 7.14 |
+| `hip-quant==2.2.3` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.3.post215` | 2.15 / TheRock | ROCm 7.14 |
 
 ```powershell
-pip install "hip-quant==2.2.2.post215"   # PyTorch 2.15 / TheRock
+pip install "hip-quant==2.2.3.post215"   # PyTorch 2.15 / TheRock
 ```
 
 #### Custom PyTorch 2.15 (self-hosted wheels)
@@ -596,7 +667,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.2.post215" `
+pip install "hip-quant==2.2.3.post215" `
   --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
 ```
 

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

@@ -14,6 +14,34 @@
 
 `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.
 
+> ⚠️ **Numerical / experimental status — WaveAttention training kernels.**
+> The WaveAttention **training** paths are experimental. The native FP8 path
+> (`precision="native"`, the default) is an inference-grade kernel: against a
+> BF16 MATH-SDPA reference its `dQ`/`dK` relative L2 is about **8% at input
+> scale 1 and ~58% at scale 5**, so it is *not* suitable for training as-is.
+> The `precision="fp16"` native WMMA kernels (new in 2.2.3) are the accurate
+> choice: against an fp32 SDPA reference the forward/backward relative L2 is
+> ~2e-4 / ~3e-4 at the tested shapes and input scales, and an adversarial
+> NaN/Inf + extreme-range harness passes clean. The `precision="bf16"` Triton
+> kernels are kept as a fallback. **Full-model convergence and end-to-end speed
+> remain unverified** (a 3,000-update PlantUnified continuation was finite,
+> without EMA/validation). A CPU float64 operand study
+> (`benchmarks/wave_attn_train_results.md`) found that no FP8 operand mixture —
+> including Hadamard-rotated, block-scaled FP8 — reaches a sub-1% training
+> 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.3
+
+> **2.2.3.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
+
+- **🌊 Native FP16 WMMA training attention (`gfx1201`).** `wave_attn(..., precision="fp16")` now runs a purpose-built FP16 flash-attention **forward and backward** on `v_wmma_f32_16x16x16_f16` — no Triton, and no SDPA recomputation for the backward. `dQ`/`dK`/`dV` are computed on-device with a FlashAttention-2 preprocess pass for `D_i = Σ_d dO_id·O_id`. Against the AOTriton SDPA baseline (`F.scaled_dot_product_attention`, flash backend) at `B1 H16`, the **forward is faster on all eight measured shapes** (1.35–2.13×; 60.9–68.4 vs 44.8–50.9 TFLOP/s at `S4096`). The **backward is still behind AOTriton** (1.47–2.06× slower; 14–22 vs 26–38 TFLOP/s): its `dK`/`dV` block needs one wave per 16-key sub-tile and ~53 KB of shared memory, which caps it at four waves per CTA and one CTA per CU.
+- **🐛 Three training-correctness bugs found and fixed**, by a new adversarial harness (`tests/torch/test_wave_attn_f16_extreme.py`: ~70 hand-built cases plus a log-uniform fuzz against an fp64 reference):
+  - **`dS` fp16 overflow → NaN.** `|dS| ≈ 3.41·|V|·|dO|` at `D=128`, so perfectly valid fp16 inputs (`V=65504`, `dO=1`) overflowed the `dS` store and produced NaN gradients. A per-`(b,h)` scale from the preprocess pass now bounds `dS`, and the output is divided back out; the `min(1, …)` cap keeps ordinary inputs bit-identical to the unscaled kernel.
+  - **Causal mask mismatch when `Seq_K > Seq_Q`.** The FP16 forward carried a bottom-right offset while the backward — and PyTorch's `is_causal` — are strictly top-left. The offset is gone, so forward and backward now agree.
+  - **`LSE` log-term loss.** Storing `row max + log(row sum)` as a single fp32 rounded the log term away once the max exceeded ~4e7, making `dV` exactly **128× too large** and turning a cancellation residual into a spurious `dQ` of 1.44e7 where the true value is exactly 0. `LSE` is now two fp32 per row, subtracted separately.
+- **🛡️ The forward can no longer overflow fp32.** Scores stay in the raw WMMA domain and `softmax_scale` is applied *inside* the exponent as `(raw − raw)·|scale|`. Forming `sacc·scale` was the forward's only fp32 overflow site (`|sacc| ≤ D·65504² = 5.5e11`, so `scale > 6.19e26` NaNs everything downstream) — and PyTorch's own fp32 SDPA and a textbook fp32 softmax **do** NaN at exactly that threshold. The native kernel now matches fp64 exactly there, while still agreeing with SDPA to ~2e-4 in the normal range. A negative `softmax_scale` (softmax of `−s`) is normalised out, so it is unaffected.
+
 ## 🌟 What's New in 2.2.2
 
 > **2.2.2.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
@@ -85,7 +113,15 @@ 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
+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):
+qbf, kbf, vbf = q.bfloat16(), k.bfloat16(), v.bfloat16()
+out_bf16 = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16")
+# Or defer to PyTorch SDPA entirely:
+out_precise = hip_quant.wave_attn(q, k, v, precision="sdpa")
 ```
 
 **Key features:**
@@ -97,7 +133,7 @@ out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
 - 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):**
+**Historical performance — RX 9070 XT `gfx1201`, `ROCm 7.2.1` + `torch 2.9.1+rocm7.2.1`, `D=128` (measured `bench_attn.py` against its AOTriton SDPA path):**
 
 | Config (B,H,S,D) | Causal | `wave` fwd | `SDPA` fwd | fwd speedup | `wave` bwd | `SDPA` bwd | bwd speedup |
 |---|:---:|---|---|---|---|---|---|
@@ -109,9 +145,44 @@ out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
 | 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.
+The table describes that benchmark's backend and shapes; it is not a current
+speed claim for every SDPA backend or training shape. On the same RX 9070 XT
+with `B=6,H=12,S=789,D=64` BF16 inputs, the full native FP8 training API
+takes **1.87 ms** forward plus backward, while `precision="sdpa"` takes
+**1.03 ms**. Against BF16 MATH SDPA, native dQ/dK relative L2 is about 8%
+at input scale 1 and 58% at scale 5; the SDPA precision mode is below 1%.
+See `benchmarks/wave_attn_train_results.md` for methods and all metrics.
+
+> The native FP8 path remains available with `precision="native"` (the default),
+> 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.
+
+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
+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
+ended at loss 7.7051 (BF16 WaveAttention) versus 7.7122 (BF16 SDPA). A
+full-model speed win and full-epoch convergence remain unverified.
 
 ## HQ2 / HQ3 — experimental learned-codebook formats
 
@@ -349,7 +420,7 @@ pip install hip-quant
 pip install "hip-quant[torch]"
 
 # Or a local wheel
-pip install dist/hip_quant-2.2.0-cp312-cp312-win_amd64.whl
+pip install dist/hip_quant-2.2.3-cp312-cp312-win_amd64.whl
 ```
 
 Two wheels are published per release, one per PyTorch ABI. Install the one that
@@ -357,11 +428,11 @@ matches your PyTorch, so the extension loads without an ABI mismatch:
 
 | Wheel | PyTorch | ROCm |
 |---|---|---|
-| `hip-quant==2.2.2` | 2.9.x | ROCm 7.2.1 |
-| `hip-quant==2.2.2.post215` | 2.15 / TheRock | ROCm 7.14 |
+| `hip-quant==2.2.3` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.3.post215` | 2.15 / TheRock | ROCm 7.14 |
 
 ```powershell
-pip install "hip-quant==2.2.2.post215"   # PyTorch 2.15 / TheRock
+pip install "hip-quant==2.2.3.post215"   # PyTorch 2.15 / TheRock
 ```
 
 #### Custom PyTorch 2.15 (self-hosted wheels)
@@ -371,7 +442,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.2.post215" `
+pip install "hip-quant==2.2.3.post215" `
   --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
 ```
 

hip_quant-2.2.3/_C.cp312-win_amd64.pyd

@@ -0,0 +1,4 @@
+[diffend] Oversized file quarantined before diffing.
+name: hip_quant-2.2.3/_C.cp312-win_amd64.pyd
+size: 41630720 bytes
+sha256: 7f845ecb6d26ac4de5f61c9da5bd3ffaa88eb549d60a350dfe29832e50b24a36

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

@@ -3,7 +3,7 @@ import numpy as np
 import os
 import sys
 
-__version__ = "2.2.2"
+__version__ = "2.2.3"
 
 
 _TORCH_EXPORTS = {

{hip_quant-2.2.2 → hip_quant-2.2.3}/gguf.py

@@ -66,6 +66,10 @@ GGML_TYPES: Dict[int, Tuple[str, int, int]] = {
     30: ("BF16",    2,   1),
     34: ("TQ1_0",   54,  256),
     35: ("TQ2_0",   66,  256),
+    # Prism-private ternary: one FP16 scale + 26 packed trit bytes per 128
+    # weights.  The byte codec is intentionally handled by HQNAT's Prism
+    # importer rather than pretending this is upstream TQ1_0.
+    143:("PTQ1_0",  28,  128),
     39: ("MXFP4",   17,  32),
     40: ("NVFP4",   36,  64),
 }

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hip-quant
-Version: 2.2.2
+Version: 2.2.3
 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
@@ -239,6 +239,34 @@ 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.
 
+> ⚠️ **Numerical / experimental status — WaveAttention training kernels.**
+> The WaveAttention **training** paths are experimental. The native FP8 path
+> (`precision="native"`, the default) is an inference-grade kernel: against a
+> BF16 MATH-SDPA reference its `dQ`/`dK` relative L2 is about **8% at input
+> scale 1 and ~58% at scale 5**, so it is *not* suitable for training as-is.
+> The `precision="fp16"` native WMMA kernels (new in 2.2.3) are the accurate
+> choice: against an fp32 SDPA reference the forward/backward relative L2 is
+> ~2e-4 / ~3e-4 at the tested shapes and input scales, and an adversarial
+> NaN/Inf + extreme-range harness passes clean. The `precision="bf16"` Triton
+> kernels are kept as a fallback. **Full-model convergence and end-to-end speed
+> remain unverified** (a 3,000-update PlantUnified continuation was finite,
+> without EMA/validation). A CPU float64 operand study
+> (`benchmarks/wave_attn_train_results.md`) found that no FP8 operand mixture —
+> including Hadamard-rotated, block-scaled FP8 — reaches a sub-1% training
+> 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.3
+
+> **2.2.3.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
+
+- **🌊 Native FP16 WMMA training attention (`gfx1201`).** `wave_attn(..., precision="fp16")` now runs a purpose-built FP16 flash-attention **forward and backward** on `v_wmma_f32_16x16x16_f16` — no Triton, and no SDPA recomputation for the backward. `dQ`/`dK`/`dV` are computed on-device with a FlashAttention-2 preprocess pass for `D_i = Σ_d dO_id·O_id`. Against the AOTriton SDPA baseline (`F.scaled_dot_product_attention`, flash backend) at `B1 H16`, the **forward is faster on all eight measured shapes** (1.35–2.13×; 60.9–68.4 vs 44.8–50.9 TFLOP/s at `S4096`). The **backward is still behind AOTriton** (1.47–2.06× slower; 14–22 vs 26–38 TFLOP/s): its `dK`/`dV` block needs one wave per 16-key sub-tile and ~53 KB of shared memory, which caps it at four waves per CTA and one CTA per CU.
+- **🐛 Three training-correctness bugs found and fixed**, by a new adversarial harness (`tests/torch/test_wave_attn_f16_extreme.py`: ~70 hand-built cases plus a log-uniform fuzz against an fp64 reference):
+  - **`dS` fp16 overflow → NaN.** `|dS| ≈ 3.41·|V|·|dO|` at `D=128`, so perfectly valid fp16 inputs (`V=65504`, `dO=1`) overflowed the `dS` store and produced NaN gradients. A per-`(b,h)` scale from the preprocess pass now bounds `dS`, and the output is divided back out; the `min(1, …)` cap keeps ordinary inputs bit-identical to the unscaled kernel.
+  - **Causal mask mismatch when `Seq_K > Seq_Q`.** The FP16 forward carried a bottom-right offset while the backward — and PyTorch's `is_causal` — are strictly top-left. The offset is gone, so forward and backward now agree.
+  - **`LSE` log-term loss.** Storing `row max + log(row sum)` as a single fp32 rounded the log term away once the max exceeded ~4e7, making `dV` exactly **128× too large** and turning a cancellation residual into a spurious `dQ` of 1.44e7 where the true value is exactly 0. `LSE` is now two fp32 per row, subtracted separately.
+- **🛡️ The forward can no longer overflow fp32.** Scores stay in the raw WMMA domain and `softmax_scale` is applied *inside* the exponent as `(raw − raw)·|scale|`. Forming `sacc·scale` was the forward's only fp32 overflow site (`|sacc| ≤ D·65504² = 5.5e11`, so `scale > 6.19e26` NaNs everything downstream) — and PyTorch's own fp32 SDPA and a textbook fp32 softmax **do** NaN at exactly that threshold. The native kernel now matches fp64 exactly there, while still agreeing with SDPA to ~2e-4 in the normal range. A negative `softmax_scale` (softmax of `−s`) is normalised out, so it is unaffected.
+
 ## 🌟 What's New in 2.2.2
 
 > **2.2.2.post215** is the same release rebuilt against PyTorch 2.15 (TheRock / ROCm 7.14).
@@ -310,7 +338,15 @@ 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
+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):
+qbf, kbf, vbf = q.bfloat16(), k.bfloat16(), v.bfloat16()
+out_bf16 = hip_quant.wave_attn(qbf, kbf, vbf, precision="bf16")
+# Or defer to PyTorch SDPA entirely:
+out_precise = hip_quant.wave_attn(q, k, v, precision="sdpa")
 ```
 
 **Key features:**
@@ -322,7 +358,7 @@ out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
 - 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):**
+**Historical performance — RX 9070 XT `gfx1201`, `ROCm 7.2.1` + `torch 2.9.1+rocm7.2.1`, `D=128` (measured `bench_attn.py` against its AOTriton SDPA path):**
 
 | Config (B,H,S,D) | Causal | `wave` fwd | `SDPA` fwd | fwd speedup | `wave` bwd | `SDPA` bwd | bwd speedup |
 |---|:---:|---|---|---|---|---|---|
@@ -334,9 +370,44 @@ out = hip_quant.wave_attn(q, k, v)  # FP8 WMMA, 3-8x faster than SDPA
 | 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.
+The table describes that benchmark's backend and shapes; it is not a current
+speed claim for every SDPA backend or training shape. On the same RX 9070 XT
+with `B=6,H=12,S=789,D=64` BF16 inputs, the full native FP8 training API
+takes **1.87 ms** forward plus backward, while `precision="sdpa"` takes
+**1.03 ms**. Against BF16 MATH SDPA, native dQ/dK relative L2 is about 8%
+at input scale 1 and 58% at scale 5; the SDPA precision mode is below 1%.
+See `benchmarks/wave_attn_train_results.md` for methods and all metrics.
+
+> The native FP8 path remains available with `precision="native"` (the default),
+> 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.
+
+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
+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
+ended at loss 7.7051 (BF16 WaveAttention) versus 7.7122 (BF16 SDPA). A
+full-model speed win and full-epoch convergence remain unverified.
 
 ## HQ2 / HQ3 — experimental learned-codebook formats
 
@@ -574,7 +645,7 @@ pip install hip-quant
 pip install "hip-quant[torch]"
 
 # Or a local wheel
-pip install dist/hip_quant-2.2.0-cp312-cp312-win_amd64.whl
+pip install dist/hip_quant-2.2.3-cp312-cp312-win_amd64.whl
 ```
 
 Two wheels are published per release, one per PyTorch ABI. Install the one that
@@ -582,11 +653,11 @@ matches your PyTorch, so the extension loads without an ABI mismatch:
 
 | Wheel | PyTorch | ROCm |
 |---|---|---|
-| `hip-quant==2.2.2` | 2.9.x | ROCm 7.2.1 |
-| `hip-quant==2.2.2.post215` | 2.15 / TheRock | ROCm 7.14 |
+| `hip-quant==2.2.3` | 2.9.x | ROCm 7.2.1 |
+| `hip-quant==2.2.3.post215` | 2.15 / TheRock | ROCm 7.14 |
 
 ```powershell
-pip install "hip-quant==2.2.2.post215"   # PyTorch 2.15 / TheRock
+pip install "hip-quant==2.2.3.post215"   # PyTorch 2.15 / TheRock
 ```
 
 #### Custom PyTorch 2.15 (self-hosted wheels)
@@ -596,7 +667,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.2.post215" `
+pip install "hip-quant==2.2.3.post215" `
   --find-links https://dl.hipquant.download/file/Torchs/torch_index.html
 ```
 

{hip_quant-2.2.2 → hip_quant-2.2.3}/hip_quant.egg-info/SOURCES.txt

@@ -28,6 +28,8 @@ setup.py
 setup_torch.py
 smi.py
 torch_api.py
+wave_attn_train_strided_triton.py
+wave_attn_train_triton.py
 wave_helpers.py
 ./_C.cp312-win_amd64.pyd
 ./__init__.py
@@ -54,6 +56,8 @@ wave_helpers.py
 ./setup_torch.py
 ./smi.py
 ./torch_api.py
+./wave_attn_train_strided_triton.py
+./wave_attn_train_triton.py
 ./wave_helpers.py
 ./codebooks/iq1_s.bin
 ./codebooks/iq2_s.bin
@@ -144,10 +148,15 @@ wave_helpers.py
 ./tests/torch/test_q_to_fp8_torch.py
 ./tests/torch/test_quant_transpose.py
 ./tests/torch/test_wave_attn_backward_math.py
+./tests/torch/test_wave_attn_f16.py
+./tests/torch/test_wave_attn_f16_extreme.py
+./tests/torch/test_wave_attn_train_bf16.py
+./tests/torch/test_wave_helpers_cpu.py
 ./tools/.gpu-smi-version
 ./tools/fetch_gpu_smi.py
 ./tools/gpu-smi.exe
 ./torch_ext/current_stream.hip
+./torch_ext/decode_quant_kernels.hip
 ./torch_ext/dequant_q_to_fp16.hip
 ./torch_ext/fp8_linear_kernels.hip
 ./torch_ext/fp8_linear_kernels_v2.hip
@@ -156,6 +165,7 @@ wave_helpers.py
 ./torch_ext/hq2_linear_kernels.hip
 ./torch_ext/hq3_linear_kernels.hip
 ./torch_ext/hq8_g128_linear_kernels.hip
+./torch_ext/int4_wmma_kernels.hip
 ./torch_ext/iq_grid_tables.h
 ./torch_ext/kv_iu4_kernels.hip
 ./torch_ext/mxfp4_to_fp8_kernels.hip
@@ -164,13 +174,18 @@ wave_helpers.py
 ./torch_ext/pytorch_bindings_hip.cpp
 ./torch_ext/q_to_fp8_kernels.hip
 ./torch_ext/ssm_kernels.hip
+./torch_ext/ternary_gemv_kernels.hip
+./torch_ext/w4a16_gemv_kernels.hip
 ./torch_ext/wave_attn.hip
 ./torch_ext/wave_attn_backward.hip
+./torch_ext/wave_attn_backward_f16.hip
 ./torch_ext/wave_attn_common.h
 ./torch_ext/wave_attn_decode.hip
 ./torch_ext/wave_attn_diag.hip
+./torch_ext/wave_attn_f16.hip
 ./torch_ext/wave_attn_long.hip
 ./torch_ext/wave_attn_prefill.hip
+./torch_ext/wht_kernels.hip
 codebooks/iq1_s.bin
 codebooks/iq2_s.bin
 codebooks/iq2_xs.bin
@@ -282,6 +297,10 @@ tests/torch/test_mxfp8_torch.py
 tests/torch/test_q_to_fp8_torch.py
 tests/torch/test_quant_transpose.py
 tests/torch/test_wave_attn_backward_math.py
+tests/torch/test_wave_attn_f16.py
+tests/torch/test_wave_attn_f16_extreme.py
+tests/torch/test_wave_attn_train_bf16.py
+tests/torch/test_wave_helpers_cpu.py
 tools/.gpu-smi-version
 tools/fetch_gpu_smi.py
 tools/gpu-smi.exe
@@ -296,6 +315,7 @@ tools/mlp_codec/run_ladder.py
 tools/mlp_codec/safe_run.py
 tools/mlp_codec/validate_hessians.py
 torch_ext/current_stream.hip
+torch_ext/decode_quant_kernels.hip
 torch_ext/dequant_q_to_fp16.hip
 torch_ext/fp8_linear_kernels.hip
 torch_ext/fp8_linear_kernels_v2.hip
@@ -304,6 +324,7 @@ torch_ext/gemv_q_kernels.hip
 torch_ext/hq2_linear_kernels.hip
 torch_ext/hq3_linear_kernels.hip
 torch_ext/hq8_g128_linear_kernels.hip
+torch_ext/int4_wmma_kernels.hip
 torch_ext/iq_grid_tables.h
 torch_ext/kv_iu4_kernels.hip
 torch_ext/mxfp4_to_fp8_kernels.hip
@@ -312,10 +333,15 @@ torch_ext/pytorch_bindings.cpp
 torch_ext/pytorch_bindings_hip.cpp
 torch_ext/q_to_fp8_kernels.hip
 torch_ext/ssm_kernels.hip
+torch_ext/ternary_gemv_kernels.hip
+torch_ext/w4a16_gemv_kernels.hip
 torch_ext/wave_attn.hip
 torch_ext/wave_attn_backward.hip
+torch_ext/wave_attn_backward_f16.hip
 torch_ext/wave_attn_common.h
 torch_ext/wave_attn_decode.hip
 torch_ext/wave_attn_diag.hip
+torch_ext/wave_attn_f16.hip
 torch_ext/wave_attn_long.hip
-torch_ext/wave_attn_prefill.hip
+torch_ext/wave_attn_prefill.hip
+torch_ext/wht_kernels.hip

{hip_quant-2.2.2 → hip_quant-2.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hip-quant"
-version = "2.2.2"
+version = "2.2.3"
 description = "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."
 readme = "README.md"
 
@@ -88,6 +88,8 @@ include-package-data = false
     "torch_ext/pytorch_bindings.cpp",
     "torch_ext/pytorch_bindings_hip.cpp",
     "torch_api.py",
+    "wave_attn_train_triton.py",
+    "wave_attn_train_strided_triton.py",
     "gguf.py",
     "gguf_loader.py",
     "wave_helpers.py",

{hip_quant-2.2.2 → hip_quant-2.2.3}/setup.py

@@ -163,9 +163,11 @@ def _torch_extension_config():
             "torch_ext/wave_attn_prefill.hip",
             "torch_ext/wave_attn_decode.hip",
             "torch_ext/wave_attn_long.hip",
+            "torch_ext/decode_quant_kernels.hip",
             "torch_ext/gemv_q_kernels.hip",
             "torch_ext/ssm_kernels.hip",
             "torch_ext/kv_iu4_kernels.hip",
+            "torch_ext/int4_wmma_kernels.hip",
             "torch_ext/dequant_q_to_fp16.hip",
         ],
 

{hip_quant-2.2.2 → hip_quant-2.2.3}/setup_torch.py

@@ -149,12 +149,19 @@ ext = CUDAExtension(
         "torch_ext/hq2_linear_kernels.hip",
         "torch_ext/hq3_linear_kernels.hip",
         "torch_ext/hq8_g128_linear_kernels.hip",
+        "torch_ext/int4_wmma_kernels.hip",
+        "torch_ext/w4a16_gemv_kernels.hip",
+        "torch_ext/ternary_gemv_kernels.hip",
+        "torch_ext/wht_kernels.hip",
         "torch_ext/wave_attn.hip",
         "torch_ext/wave_attn_backward.hip",
+        "torch_ext/wave_attn_f16.hip",
+        "torch_ext/wave_attn_backward_f16.hip",
         "torch_ext/wave_attn_diag.hip",
         "torch_ext/wave_attn_prefill.hip",
         "torch_ext/wave_attn_decode.hip",
         "torch_ext/wave_attn_long.hip",
+        "torch_ext/decode_quant_kernels.hip",
         "torch_ext/gemv_q_kernels.hip",
         "torch_ext/ssm_kernels.hip",
         "torch_ext/kv_iu4_kernels.hip",