orbitquant 0.2.2__tar.gz → 0.3.0__tar.gz

{orbitquant-0.2.2 → orbitquant-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orbitquant
-Version: 0.2.2
+Version: 0.3.0
 Summary: Calibration-free OrbitQuant for transformer linear projections
 Project-URL: Homepage, https://github.com/iamwavecut/OrbitQuant
 Project-URL: Repository, https://github.com/iamwavecut/OrbitQuant
@@ -260,7 +260,7 @@ pipe = load_quantized_pipeline_from_artifact(
 
 | Device | Dispatch |
 | --- | --- |
-| CUDA | Native packed CUDA package, then Triton packed matmul |
+| CUDA | Native activation kernel plus packed W4A4 tensor-core path; native or Triton packed fallback |
 | MPS | Native packed Metal package |
 | CPU | PyTorch reference path |
 
@@ -274,6 +274,18 @@ Use the explicit reference path for compatibility or numerical debugging:
 config = orbitquant.recipe("w4a4", runtime_mode="dequant_bf16")
 ```
 
+On CUDA compute capability 8.0 or newer, the W4A4 fast path fuses token norm,
+RPBH/FWHT, and codebook assignment in the native package, decodes only a bounded
+output-channel chunk of packed weights to INT8, and uses the Torch CUTLASS
+tensor-core matmul. It never materializes the full BF16/FP16 weight matrix.
+The existing direct packed CUDA MMA kernel remains the fallback for compatible
+W4A4 shapes when CUTLASS INT8 matmul is unavailable.
+
+The optimized CUDA path maps the fixed Lloyd-Max centroids to a symmetric INT8
+surrogate plus one scalar per codebook. Packed checkpoint indices and artifact
+size are unchanged. Use `dequant_bf16` when exact Lloyd-Max centroid evaluation
+is required.
+
 Build the ABI3 native package locally without Kernel Hub:
 
 ```bash
@@ -282,6 +294,13 @@ nix --option sandbox relaxed run .#build-and-copy -L
 export PYTHONPATH="$PWD/build/<matching-torch-backend-platform-variant>:$PYTHONPATH"
 ```
 
+PyTorch 2.9 CUDA users can reduce allocator reservation during native diffusion
+inference by setting the allocator before Python starts:
+
+```bash
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py
+```
+
 The variant must match the Torch minor version, CUDA or Metal backend, C++ ABI,
 architecture, and operating system. See
 [`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark

{orbitquant-0.2.2 → orbitquant-0.3.0}/README.md

@@ -216,7 +216,7 @@ pipe = load_quantized_pipeline_from_artifact(
 
 | Device | Dispatch |
 | --- | --- |
-| CUDA | Native packed CUDA package, then Triton packed matmul |
+| CUDA | Native activation kernel plus packed W4A4 tensor-core path; native or Triton packed fallback |
 | MPS | Native packed Metal package |
 | CPU | PyTorch reference path |
 
@@ -230,6 +230,18 @@ Use the explicit reference path for compatibility or numerical debugging:
 config = orbitquant.recipe("w4a4", runtime_mode="dequant_bf16")
 ```
 
+On CUDA compute capability 8.0 or newer, the W4A4 fast path fuses token norm,
+RPBH/FWHT, and codebook assignment in the native package, decodes only a bounded
+output-channel chunk of packed weights to INT8, and uses the Torch CUTLASS
+tensor-core matmul. It never materializes the full BF16/FP16 weight matrix.
+The existing direct packed CUDA MMA kernel remains the fallback for compatible
+W4A4 shapes when CUTLASS INT8 matmul is unavailable.
+
+The optimized CUDA path maps the fixed Lloyd-Max centroids to a symmetric INT8
+surrogate plus one scalar per codebook. Packed checkpoint indices and artifact
+size are unchanged. Use `dequant_bf16` when exact Lloyd-Max centroid evaluation
+is required.
+
 Build the ABI3 native package locally without Kernel Hub:
 
 ```bash
@@ -238,6 +250,13 @@ nix --option sandbox relaxed run .#build-and-copy -L
 export PYTHONPATH="$PWD/build/<matching-torch-backend-platform-variant>:$PYTHONPATH"
 ```
 
+PyTorch 2.9 CUDA users can reduce allocator reservation during native diffusion
+inference by setting the allocator before Python starts:
+
+```bash
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py
+```
+
 The variant must match the Torch minor version, CUDA or Metal backend, C++ ABI,
 architecture, and operating system. See
 [`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark

orbitquant-0.3.0/docs/flux2-klein-9b-sdnq-vs-orbitquant.md

@@ -0,0 +1,143 @@
+# FLUX.2 Klein 9B: OrbitQuant W4A4 vs SDNQ UINT4
+
+This report compares OrbitQuant and SDNQ on the same FLUX.2 Klein 9B source,
+prompts and generation settings. Runtime rows use the same L40S hardware and
+software environment. It covers checkpoint size, load and generation latency,
+VRAM, paired visual output and the measured OrbitQuant kernel path.
+
+## Checkpoints
+
+| Variant | Repository | Revision | Quantized components |
+| --- | --- | --- | --- |
+| BF16 | [`black-forest-labs/FLUX.2-klein-9B`](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) | `92196c8e11f7b6cf2b7493e037d8c5345c559216` | None |
+| SDNQ UINT4 | [`WaveCut/FLUX.2-klein-9B-SDNQ-uint4-static`](https://huggingface.co/WaveCut/FLUX.2-klein-9B-SDNQ-uint4-static) | `ed71b3f19ce640e88b66a2a743aabb8a613adeac` | Transformer and Qwen3 text encoder |
+| OrbitQuant W4A4 | [`WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4`](https://huggingface.co/WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4) | `ee3a38f7767ae199818d746c840be0f1837887bf` | Transformer and Qwen3 text encoder |
+
+The public SDNQ checkpoint was selected instead of the AI Farm LoRA/heretic-text-encoder
+variant so that quantizer choice is the main experimental difference. SDNQ uses UINT4
+weights with BF16 activations. OrbitQuant uses 4-bit codebook weights and 4-bit codebook
+activations.
+
+OrbitQuant quantized 144 transformer projections and 252 text-encoder projections. Three
+transformer modulation projections use INT4 RTN with BF16 activations, and the text
+encoder `lm_head` remains BF16. Quantizing the text encoder is a deliberate universal
+adapter extension for this comparison; the OrbitQuant paper leaves text encoders in BF16.
+
+## Protocol
+
+- GPU: NVIDIA L40S 48 GB (`sm_89`)
+- Torch: 2.9.1+cu128
+- CUDA: 12.8
+- Diffusers: 0.39.0
+- Transformers: 5.13.0
+- OrbitQuant: 0.3.0
+- SDNQ: 0.1.8
+- Arithmetic: BF16
+- CPU offload: disabled
+- Output: 1024x1024
+- Steps: 4
+- Guidance: 1.0
+- Seed: 0
+- Batch size: 1
+- Ten identical prompts per variant
+- Each variant loaded and ran in a separate process
+- `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`
+
+The prompt pack stresses micro-detail, exact counting, nested spatial composition,
+fictional authorial style, abstract material separation, English fine print, Russian,
+Japanese and Chinese typography, reflections, occlusion and a dense panoramic scene.
+
+## Artifact Size
+
+The table counts model weight payloads only, excluding cards and comparison images.
+
+| Variant | Transformer | Text encoder | VAE | Total weights |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 18.157 GB | 16.382 GB | 0.168 GB | 34.707 GB |
+| SDNQ UINT4 | 5.616 GB | 6.397 GB | 0.168 GB | 12.181 GB |
+| OrbitQuant W4A4 | 4.705 GB | 5.966 GB | 0.168 GB | 10.839 GB |
+
+OrbitQuant's weight payload is 68.8% smaller than BF16 and 11.0% smaller than the
+controlled SDNQ artifact. The complete loadable OrbitQuant pipeline before its model-card
+matrix is 10.85 GB.
+
+## Runtime
+
+The OrbitQuant row used a locally built ABI3 CUDA package from
+`native-kernels/orbitquant-packed-matmul`; Kernel Hub was not involved. All 396
+packed projections selected `native_packed_matmul` and the optimized W4A4 path.
+
+| Variant | Load | Cold image | Hot mean | Hot median | Hot p95 |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| SDNQ UINT4 | 5.918 s | 16.956 s | 2.0885 s | 2.0875 s | 2.0966 s |
+| OrbitQuant W4A4 | 2.543 s | 4.193 s | 2.0907 s | 2.0920 s | 2.0988 s |
+
+| Variant | Load NVML peak | Hot NVML peak | CUDA allocated peak | CUDA reserved peak |
+| --- | ---: | ---: | ---: | ---: |
+| SDNQ UINT4 | 13.383 GB | 17.564 GB | 14.844 GB | 16.377 GB |
+| OrbitQuant W4A4 | 11.959 GB | 15.731 GB | 13.942 GB | 14.544 GB |
+
+OrbitQuant's hot mean was 0.11% slower than SDNQ, which is practical parity for
+this run. OrbitQuant loaded 57.0% faster, used 0.902 GB less peak CUDA allocated
+memory, and used 1.833 GB less peak CUDA reserved and NVML memory. The packed
+weight payload also remains 11.0% smaller.
+
+The selected CUDA path performs native token norm, RPBH/FWHT and codebook-bin
+selection, emits an INT8 surrogate of the 4-bit activation codebook, decodes
+only a bounded output-channel chunk of the packed W4 weights, calls the
+CUTLASS-backed `torch._int_mm`, and applies norms, scales and bias in a Triton
+epilogue. It does not materialize a full BF16/FP16 weight matrix. The direct
+packed CUDA MMA implementation remains the fallback for unsupported shapes.
+
+The optimized path adds a small runtime approximation beyond the paper
+equation: fixed Lloyd-Max centroids are represented by symmetric INT8 codes and
+one scalar per codebook. Packed indices and the checkpoint are unchanged;
+`runtime_mode="dequant_bf16"` remains the exact-centroid reference. Local build
+and verification instructions are in [the kernel audit](kernel-audit.md#local-native-package).
+
+## Paired Visual Comparison
+
+The matrix is a separate controlled A40 visual run using full 1024x1024 tiles
+and WebP quality 95. Every column uses the same prompt, seed and pipeline
+settings. The L40S production-default OrbitQuant rerun produced ten finite
+images that matched the previously validated optimized W4A4 outputs byte for
+byte.
+
+![BF16, SDNQ UINT4 and OrbitQuant W4A4 across ten difficult prompts](assets/flux2-klein-9b-sdnq-vs-orbitquant.webp)
+
+## Visual Assessment
+
+- **No collapse:** all thirty outputs are finite, coherent and detailed. OrbitQuant did not
+  produce blank, noisy or structurally broken images.
+- **Micro-detail and materials:** all three variants preserve gears, filigree, architectural
+  interiors, paper grain, metal, resin and reflected surfaces. OrbitQuant remains competitive
+  with BF16 and SDNQ in these cases.
+- **Dense composition:** all variants retain foreground/background separation and the main
+  hierarchy in the architectural cutaway and orbital-banquet prompts. Individual requested
+  objects move or disappear because quantization changes the denoising trajectory.
+- **Counting:** none of the variants reliably renders exactly nine performers or every exact
+  repeated motif. This is a base-model limitation in the tested setting rather than an
+  OrbitQuant-only collapse.
+- **English typography:** SDNQ is strongest on the small four-line specification table.
+  OrbitQuant preserves the headline and subtitle but misspells or truncates some fine print.
+- **Russian typography:** all variants render the large headline, subtitle and archive stamp
+  well; small contents text contains errors in every column.
+- **Japanese and Chinese typography:** visual glyph quality is plausible, but exact requested
+  strings are not reliably reproduced by any variant.
+- **Trajectory fidelity:** SDNQ often stays visually closer to BF16. OrbitQuant also quantizes
+  activations and therefore produces larger compositional changes at the same seed.
+
+This assessment is subjective and paired. It demonstrates non-collapse and exposes concrete
+failure modes; it is not a substitute for GenEval or another task-specific objective metric.
+
+## Result
+
+OrbitQuant produces the smaller complete 4-bit pipeline, additionally quantizes
+activations without calibration data, and reaches SDNQ hot-generation parity on
+the tested L40S while using materially less runtime memory. The visual matrix
+shows preserved complex structure without collapse; SDNQ remains closer to BF16
+on some typography cases. This is a controlled result for FLUX.2 Klein 9B, not a
+universal speed claim for every model or GPU.
+
+Machine-readable metrics and the exact ten prompts are included with the
+[OrbitQuant checkpoint](https://huggingface.co/WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4).

{orbitquant-0.2.2 → orbitquant-0.3.0}/docs/kernel-audit.md

@@ -10,7 +10,7 @@ when explicitly selected.
 
 | Device | Default dispatch | Required support |
 | --- | --- | --- |
-| CUDA | Native packed matmul, then Triton packed matmul | An importable local native package or Triton |
+| CUDA | Native activation kernel plus packed W4A4 tensor-core path; native or Triton packed fallback | An importable local native package and, for the CUTLASS path, Triton |
 | MPS | Native packed matmul | An importable local Metal package |
 | CPU | Reference matmul | PyTorch |
 
@@ -26,7 +26,7 @@ Other explicit modes are `native_packed_matmul`, `triton_packed_matmul`,
 
 | Backend | Status | Implemented path |
 | --- | --- | --- |
-| CUDA | Optimized packed inference | Native CUDA packed matmul; Triton activation norm, RPBH/FWHT, codebook lookup, rescale, and packed matmul fallback |
+| CUDA | Optimized packed inference | Native RPBH/quantization, chunked packed-weight decode plus CUTLASS INT8 matmul, direct packed CUDA MMA fallback, and generic Triton packed fallback |
 | MPS/Metal | Optimized packed inference | Native Metal packed matmul and Metal activation quantization stages |
 | CPU | Reference | PyTorch activation quantization, weight dequantization, and linear matmul |
 | ROCm | Unsupported | No release backend |
@@ -60,6 +60,27 @@ TORCH_CUDA_ARCH_LIST="8.9" CUDACXX=/usr/local/cuda/bin/nvcc \
 export PYTHONPATH="$PWD/build/<matching-cuda-variant>:$PYTHONPATH"
 ```
 
+For a locally built Metal variant that remains loadable on macOS 15 and newer:
+
+```bash
+cargo install --git https://github.com/huggingface/kernels hf-kernel-builder
+cd native-kernels/orbitquant-packed-matmul
+kernel-builder check-config .
+kernel-builder create-pyproject -f .
+MACOSX_DEPLOYMENT_TARGET=15.0 \
+  CMAKE_ARGS="-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0" \
+  python setup.py build_kernel
+kernel-builder check-abi --macos 15.0 --python-abi 3.9 .
+export PYTHONPATH="$PWD/build/<matching-metal-variant>:$PYTHONPATH"
+```
+
+For PyTorch 2.9 CUDA workloads, enable expandable allocator segments before the
+Python process starts to minimize reserved/NVML memory:
+
+```bash
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py
+```
+
 The generated `setup.py` and CMake files come from `kernel-builder` and must not
 be committed. This development build targets the current host toolchain. Use the
 Nix build for a redistributable variant and run `kernel-builder check-abi` before
@@ -96,37 +117,67 @@ dequantization or `F.linear` fallback in optimized mode.
 MPS verification passed on Apple Silicon with Torch 2.12.1. It covered the
 native Metal package, inline shader stages, `auto_fused` dispatch, and a real
 3072x3072 projection restored from the published FLUX.2 W4A4 artifact. The
-packed and reference outputs were finite and numerically close.
+macOS 15 deployment-target build passed `kernel-builder check-abi` for the
+Python 3.9 stable ABI, and the packed and reference outputs were finite and
+numerically close.
 
 The Metal package also passed an ABI3 build matrix for Torch 2.11, 2.12, and
 2.13. A quantized tiny GPT-2 run exercised all eight wrapped projections during
 prefill and cached decode through `native_packed_matmul`, with finite outputs.
 
-The native CUDA package was also built and tested on an NVIDIA RTX 4090
-(`sm_89`) with Torch 2.9.1+cu128, CUDA 12.8, and kernel-builder 0.17.0-dev0.
-All 34 applicable package tests passed for W2/W3/W4/W6, FP16/BF16, bias and
-no-bias paths, partial output tiles, short rows, and tensor-core rows.
-
-The same local package build passed its 34 CUDA tests on an NVIDIA A40
-(`sm_86`) with Torch 2.9.1+cu128 and CUDA 12.8. A full FLUX.2 Klein 9B W4A4
-pipeline exercised 396 packed linears across the transformer and Qwen3 text
-encoder. Native 1024x1024 generation at four steps used
-`native_packed_matmul` for every packed linear and the Triton activation path;
-no full weight matrix was materialized.
-
-Representative A40 W4A4 layer timings from that pipeline are:
-
-| Projection | Rows | Shape | RPBH + activation quantization | Native packed matmul | Full layer |
-| --- | ---: | ---: | ---: | ---: | ---: |
-| Double-stream Q | 4096 | 4096 -> 4096 | 0.698 ms | 2.607 ms | 3.279 ms |
-| Single-stream fused input | 4608 | 4096 -> 36864 | 0.770 ms | 25.962 ms | 26.677 ms |
-| Single-stream output | 4608 | 16384 -> 4096 | 4.257 ms | 11.244 ms | 15.514 ms |
-
-The corresponding full-pipeline hot-generation mean was 7.418 seconds with a
-17.66 GB NVML peak. On the same A40 and prompt pack, BF16 measured 3.911
-seconds and 40.83 GB, while SDNQ UINT4 measured 3.546 seconds and 17.55 GB.
-These results establish the memory reduction and the remaining throughput gap;
-they do not support a model-level speedup claim for OrbitQuant on A40.
+The final CUDA W4A4 stack was built and tested on an NVIDIA L40S (`sm_89`) with
+Torch 2.9.1+cu128 and CUDA 12.8. The native package suite passed 49 CUDA tests;
+the ten skipped cases were Metal-only. Coverage includes W2/W3/W4/W6 generic
+packed matmul, FP16/BF16, bias and no-bias paths, partial output tiles, direct
+packed W4A4 MMA, native packed-A4 activation quantization, and native INT8
+activation quantization for full-block dimensions and the 12288/4096 blocked
+RPBH case.
+
+For W4A4 on compute capability 8.0 or newer, the selected path is:
+
+1. A native CUDA launch computes token norms, applies RPBH/FWHT, selects the
+   fixed Lloyd-Max bins, and emits INT8 surrogate codes.
+2. Packed row-major W4 indices are decoded one bounded output-channel chunk at
+   a time; the complete floating-point weight matrix is never materialized.
+3. `torch._int_mm` dispatches the INT8 matrix product to CUTLASS tensor cores.
+4. A Triton epilogue applies token norms, BF16 row norms, both surrogate scales,
+   and bias.
+
+The direct packed CUDA MMA implementation remains available for unsupported
+CUTLASS shapes. It includes asynchronous packed loads and SM89-specific tile
+selection. The checkpoint keeps the original row-major four-bit payload; no
+repacked duplicate weights are stored.
+
+The selected production dispatch was also profiled on an NVIDIA GeForce RTX
+4090 (`sm_89`) with Torch 2.9.1+cu128 and CUDA 12.8. For a representative
+FLUX.2 fused-input projection with 4608 activation rows, 4096 input channels,
+and 36864 output channels, ten post-warmup calls measured 4.317 ms median,
+4.352 ms mean, and 0.867 GB peak allocated memory. The output was finite and
+the dispatch reported `native_packed_matmul` with
+`native_cuda_int8_surrogate` activation quantization.
+
+Nsight Systems attributed 77.8% of GPU kernel time to the CUTLASS INT8 GEMM,
+10.0% to the fused scale/norm/bias epilogue, 8.2% to bounded packed-W4 decode,
+and 4.0% to native token norm, RPBH/FWHT, and codebook assignment. The same
+shape measured 4.809 ms median on the L40S. Nsight Compute performance counters
+were unavailable on the hosted 4090 because the provider disabled GPU counter
+access (`ERR_NVGPUCTRPERM`); the Systems trace and CUDA event timings do not
+depend on those counters.
+
+A full FLUX.2 Klein 9B W4A4 pipeline exercised all 396 packed projections
+across the transformer and Qwen3 text encoder. The controlled native run used
+1024x1024 output, four steps, guidance 1.0, seed 0, and ten identical prompts:
+
+| Runtime | Load | Hot mean | Hot median | CUDA allocated peak | CUDA reserved peak | NVML peak |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| SDNQ UINT4 | 5.918 s | 2.0885 s | 2.0875 s | 14.844 GB | 16.377 GB | 17.564 GB |
+| OrbitQuant W4A4 | 2.543 s | 2.0907 s | 2.0920 s | 13.942 GB | 14.544 GB | 15.731 GB |
+
+OrbitQuant was within 0.11% of SDNQ hot mean while using 0.902 GB less CUDA
+allocated memory and 1.833 GB less reserved/NVML memory. Every projection
+reported `native_cuda_int8_surrogate`; no full-weight dequantization path was
+entered. The ten deterministic outputs were finite and matched the separately
+validated cumulative W4A4 run byte for byte.
 
 Measured W4 BF16 operator latency for `in_features=768` and
 `out_features=2304`:
@@ -202,9 +253,12 @@ python scripts/verify_hf_kernel_model_artifact.py \
 ## Performance Claims
 
 Packed execution reduces weight-side materialization and runtime memory for the
-validated image pipelines. Throughput depends on model shapes, device, Torch,
-offload policy, and backend. Wan with CPU offload did not show a throughput or
-peak-memory improvement in the recorded native run. OrbitQuant does not claim a universal speedup.
+validated image pipelines. On the controlled L40S FLUX.2 Klein 9B comparison,
+the optimized W4A4 path reached practical SDNQ hot-generation parity with lower
+allocated, reserved, and NVML memory. Throughput still depends on model shapes,
+device, Torch, offload policy, and backend. Wan with CPU offload did not show a
+throughput or peak-memory improvement in the recorded native run. OrbitQuant
+does not claim a universal speedup.
 
 Synthetic operator benchmarks are diagnostics. Results above compare packed
 execution with both weight materialization plus `F.linear` and, where stated,

{orbitquant-0.2.2 → orbitquant-0.3.0}/docs/paper-methodology-audit.md

@@ -65,7 +65,7 @@ Status legend:
 | Embeddings, timestep MLPs, final projection/unpatchify heads, text encoders, VAE, scheduler, safety/image processors remain unquantized by default. | Pass for configured transformer components | `src/orbitquant/policies/generic_dit.py`, `tests/test_target_policies.py`; inventory summary below | Text encoders and VAE are outside the transformer component and are not passed into the default quantization helper. Artifact manifests still need per-artifact cross-checks before final publication. |
 | Native settings match paper for FLUX.1-schnell, Z-Image-Turbo, and Wan 2.1-1.3B. | Pass for encoded settings | `src/orbitquant/eval/native_settings.py`, `README.md`, `src/orbitquant/artifacts/model_card.py` | Native artifact-readiness evidence is separate from release-grade metric tables. Full metric runs are required before metric-table or paper-reproduction claims. |
 | FLUX.2 Klein is separated from paper-reproduction targets. | Pass | `src/orbitquant/eval/native_settings.py`, `src/orbitquant/artifacts/model_card.py` | It is treated as an additional target using paper-style native settings. |
-| Runtime acceleration claims match implemented kernels. | Partial | `src/orbitquant/kernels/dispatch.py`, `src/orbitquant/kernels/triton_cuda.py`, `src/orbitquant/kernels/mps.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py` | Default `auto_fused` requires packed low-bit matmul on CUDA/MPS and fails loudly when kernels are missing. CUDA and MPS avoid full weight materialization; throughput varies by model and is reported only from measured benchmarks. |
+| Runtime acceleration claims match implemented kernels. | Pass for the measured FLUX.2 Klein 9B W4A4 configuration; partial across other models | `src/orbitquant/kernels/dispatch.py`, `src/orbitquant/kernels/triton_cuda.py`, `src/orbitquant/kernels/mps.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `docs/kernel-audit.md` | Default `auto_fused` requires packed low-bit matmul on CUDA/MPS and fails loudly when kernels are missing. CUDA and MPS avoid full floating-point weight materialization. The CUDA INT8-surrogate fast path is an explicitly documented runtime approximation; exact Lloyd-Max centroid evaluation remains available through `dequant_bf16`. |
 | Release-grade GenEval/VBench metrics are available for paper target claims. | Blocked for metric claims | `src/orbitquant/hub.py`, `src/orbitquant/eval/` | Missing metrics block only paper metric/reproduction claims. |
 
 ## Model Policy Evidence
@@ -126,7 +126,7 @@ metrics; they do not by themselves claim GenEval or VBench scores.
 | Backend | Status | Evidence | Claim boundary |
 | --- | --- | --- | --- |
 | CPU | Pass as reference | `src/orbitquant/kernels/dispatch.py`, `src/orbitquant/functional.py` | Correctness baseline only; no optimized CPU kernel claim. |
-| CUDA/Triton | Pass for packed fallback | `src/orbitquant/kernels/triton_cuda.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `docs/kernel-audit.md` | Covers activation norm/RPBH/lookup/rescale, low-bit pack/unpack, offline weight quantization, AdaLN RTN quant/dequant, and packed matmul. Default `auto_fused` selects native packed matmul first, then Triton packed matmul when available. |
+| CUDA/native/Triton | Pass for optimized W4A4 and packed fallback | `native-kernels/orbitquant-packed-matmul`, `src/orbitquant/kernels/native_packed_matmul.py`, `src/orbitquant/kernels/triton_cuda.py`, `tests/test_native_packed_matmul.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `docs/kernel-audit.md` | The selected W4A4 path fuses norm/RPBH/FWHT/codebook assignment to INT8 surrogate activations, decodes bounded W4 chunks, uses CUTLASS INT8 matmul, and applies a fused epilogue. Direct packed CUDA MMA and generic Triton packed matmul remain fallbacks. No path selected by `auto_fused` materializes a full BF16/FP16 weight matrix. |
 | MPS/Metal | Pass for native packed inference | `src/orbitquant/kernels/mps.py`, `src/orbitquant/kernels/dispatch.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `docs/kernel-audit.md` | A fused Metal shader performs activation norm, RPBH/FWHT, codebook lookup, and rescale. The native package performs packed matmul for generic leading dimensions, including short decode rows and partial matrix tiles, without full weight materialization. Offline weight and AdaLN quantization remain reference paths on MPS. |
 | ROCm | Blocked for backend claim | No implementation in current tree | Do not claim ROCm optimization. |
 | XPU | Blocked for backend claim | No implementation in current tree | Do not claim XPU optimization. |
@@ -136,6 +136,11 @@ metrics; they do not by themselves claim GenEval or VBench scores.
 - `runtime_mode="auto_fused"` is the default optimized policy. It avoids silent
   CUDA/MPS fallback to full dequantized BF16 weight materialization. Explicit
   `runtime_mode="dequant_bf16"` remains the compatibility/debug reference path.
+- The CUDA W4A4 tensor-core path retains the paper's packed nearest-centroid
+  indices but approximates each fixed Lloyd-Max centroid with a symmetric INT8
+  code and one scalar per codebook. Measured codebook relative RMSE is
+  0.21-0.28% for the FLUX.2 dimensions. This extra approximation is not part of
+  the paper equation and is never presented as the exact reference path.
 
 ## Native Eval And Claim Policy
 
@@ -172,7 +177,8 @@ not silently substitute a different quantization method.
 | --- | --- | --- |
 | Explicit `dequant_bf16` runtime uses dequantized BF16 matmul. | Accepted reference path | It is kept for compatibility and debugging. Do not claim it as low-bit fused inference. |
 | Zero weight rows use an epsilon guard for direction quantization. | Accepted implementation guard | The paper defines weight directions as `w' / ||w'||` for nonzero rows. The implementation divides by `max(||w'||, ε)` only when choosing codebook indices, stores the raw BF16 row norm, and dequantizes zero rows back to exactly zero. |
-| Full-model speedup is not yet a release claim. | Accepted claim boundary | `auto_fused`, `native_packed_matmul`, and `triton_packed_matmul` use packed matmul paths, but model-specific benchmark artifacts are still required before broad acceleration claims. |
+| The optimized CUDA W4A4 path evaluates an INT8 surrogate of each Lloyd-Max codebook. | Documented runtime deviation | Packed nearest-centroid indices, row norms, token norms, and the artifact remain unchanged. The surrogate adds 0.21-0.28% codebook relative RMSE for the measured FLUX.2 dimensions. `dequant_bf16` evaluates the stored Lloyd-Max centroids directly and is the exact methodology reference. |
+| Full-model speedup is configuration-specific. | Accepted claim boundary | FLUX.2 Klein 9B on L40S reached practical SDNQ hot-generation parity with lower memory. That result does not establish universal speedup for other models, shapes, GPUs, or offload policies. |
 | The paper's block-size enumeration omits `h=256`, although its stated largest-power-of-two-divisor rule gives `h=256` for Z-Image `d=3840` and Wan `d=8960` projections. | Paper inconsistency | The implementation follows the formal rule. The selected target dimensions produce `h` in `{256, 512, 1024, 2048, 4096}`. |
 | Published checkpoints use converged Lloyd-Max codebook version 2 and `activation_eps=1e-10`. | Pass | All 14 canonical FLUX.2, FLUX.1-schnell, Z-Image-Turbo, and Wan2.1 artifacts were regenerated and validated. Legacy version 1 artifacts remain loadable by the library but are no longer the published release checkpoints. |
 | Full config-derived inventories are audit artifacts, not committed source files. | Accepted artifact hygiene choice | Inventory summaries are recorded above; raw JSON may remain unpublished to avoid turning the repository into an artifact store. |

{orbitquant-0.2.2 → orbitquant-0.3.0}/native-kernels/orbitquant-packed-matmul/CARD.md

@@ -36,6 +36,21 @@ Inputs:
 `x` may be `float32`, `float16`, or `bfloat16`. The output has shape
 `[..., out_features]` and the same dtype as `x`.
 
+The CUDA package also exports the operations used by OrbitQuant's W4A4 runtime:
+
+- `quantize_activations_int8`: token norm, RPBH/FWHT, nearest-codebook
+  assignment, and INT8-surrogate output in one native launch.
+- `quantize_activations_packed_w4`: the same activation path with packed 4-bit
+  output for the direct packed matmul fallback.
+- `matmul_packed_w4a4_int8`: direct packed A4/W4 CUDA MMA with fused token norm,
+  row norm, surrogate scales, and bias epilogue.
+
+On CUDA compute capability 8.0 or newer, OrbitQuant normally combines
+`quantize_activations_int8` with chunked packed-weight decode and Torch's
+CUTLASS-backed INT8 matmul. The direct packed MMA operation remains available
+when that path is unsupported. Neither path materializes a complete BF16/FP16
+weight matrix.
+
 ## Build And Test
 
 ```bash
@@ -61,6 +76,18 @@ TORCH_CUDA_ARCH_LIST="8.9" CUDACXX=/usr/local/cuda/bin/nvcc \
   python setup.py build_kernel
 ```
 
+For a local Metal build compatible with macOS 15 and newer:
+
+```bash
+cargo install --git https://github.com/huggingface/kernels hf-kernel-builder
+kernel-builder check-config .
+kernel-builder create-pyproject -f .
+MACOSX_DEPLOYMENT_TARGET=15.0 \
+  CMAKE_ARGS="-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0" \
+  python setup.py build_kernel
+kernel-builder check-abi --macos 15.0 --python-abi 3.9 .
+```
+
 This generated project is for local testing and must not be committed or
 distributed without a successful `kernel-builder check-abi`. Use the Nix build
 for redistributable variants.
@@ -74,6 +101,13 @@ export PYTHONPATH="/path/to/build/torch212-metal-aarch64-darwin:$PYTHONPATH"
 python -c "import orbitquant_packed_matmul; print(orbitquant_packed_matmul)"
 ```
 
+For PyTorch 2.9 CUDA inference, set the allocator before starting Python when
+minimum reserved memory is important:
+
+```bash
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py
+```
+
 OrbitQuant detects that importable package before trying any Hub loader. For
 Hugging Face `kernels` local loading instead, set `LOCAL_KERNELS` to the same
 built variant directory containing `metadata.json`, not to the source package
@@ -118,3 +152,6 @@ It also reports storage accounting for the packed weight path:
 `packed_weight_path_vs_materialized_weight_ratio`. These values describe only
 the weight-side storage used by this operator; they are not end-to-end model
 VRAM measurements.
+
+End-to-end FLUX.2 Klein 9B measurements and the SDNQ comparison are recorded in
+[`docs/flux2-klein-9b-sdnq-vs-orbitquant.md`](../../docs/flux2-klein-9b-sdnq-vs-orbitquant.md).