diff --git a/README.md b/README.md index bb13847..9e0a6e3 100644 --- a/README.md +++ b/README.md @@ -20,10 +20,11 @@ The implementation is clean-room and Apache-2.0 licensed. - Model-specific policies for paper-sensitive AdaLN and output-layer handling. - Optional GPTQ rounding (`orbitquant.gptq`): calibrated codes in the same codebook and artifact format, for models that lose too much with nearest rounding. -- Optional quality knobs outside the paper recipe: a second RPBH round for narrow - Hadamard blocks, per-token absmax INT8 activations (`recipe("w4a8")`), MSE-fitted - AdaLN scales, a Lloyd-Max gain correction, and BF16 side columns for massive - activation channels (see [Quality options](#quality-options-opt-in)). +- Measured defaults on top of the paper recipe: a second RPBH round where the paper's + Hadamard block is narrower than 512, and per-token absmax INT8 activations at 8 bits; + `recipe(name, paper=True)` restores the paper fields. Opt-in knobs: MSE-fitted AdaLN + scales, a Lloyd-Max gain correction, and BF16 side columns for massive activation + channels (see [Quality options](#quality-options)). - Compact `safetensors` artifacts and Hugging Face `save_pretrained()` / `from_pretrained()` integration. - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full @@ -100,27 +101,52 @@ On the MOSS-Image 4-bit transformer, calibrated on 10 prompts at 12 sampling ste LPIPS of 8 generated images against BF16 fell from 0.371 (nearest rounding) to 0.209; 6-bit nearest rounding gives 0.166. -## Quality options (opt-in) +## Quality options -Every option below defaults to the paper recipe and is stored in the artifact config. +The paper recipe is the baseline. A field defaults to another value only where that value +measured better in quality, speed or memory; every value is stored in the artifact config. +`recipe(name, paper=True)` and `orbitquant.recipes.PAPER_FIELDS` restore the paper fields +(`rotation_rounds=1`, `activation_quantizer="lloyd_max"`, `output_gain_correction=False`). +Artifacts saved before a default changed load with the value they were quantized with. | Field | Values | Effect | Runtime cost | | --- | --- | --- | --- | -| `rotation_rounds` | `1` (paper), `2`, `"auto"` | An independent first RPBH round in front of the paper round. The paper block `h` is the largest power of two dividing `d`: `d = 3360` gets `h = 32`, `d = 3840` and `d = 8960` get `h = 256`, so one round spreads an outlier over few coordinates. `"auto"` adds the round where `h < 512`. | One extra RPBH pass per activation, done in PyTorch: on an RTX 4090 it adds 0.7-0.8 ms per call at 256 tokens and 2.6-3.4 ms at 4096 tokens for `d = 3072`/`3840` (shared across Q/K/V of one input); not available on the fused block runtime. | -| `activation_quantizer` | `"lloyd_max"` (paper), `"absmax_int8"` | With `activation_bits=8`, per-token absmax INT8 of the rotated input, the numerics of the fused runtime's 8-bit activations. `recipe("w4a8")` selects it. No activation codebook is built or stored: exact 8-bit Lloyd-Max tables take minutes per width and do not converge for `d >= 7680`, so `"lloyd_max"` A8 fails on such inputs. | CUDA runs an INT8 GEMM against the packed W2/W4 weights, the same GEMM class as W4A4. | +| `rotation_rounds` | `"auto"` (default), `1` (paper), `2` | An independent first RPBH round in front of the paper round. The paper block `h` is the largest power of two dividing `d`: `d = 3360` gets `h = 32`, `d = 3840` and `d = 8960` get `h = 256`, so one round spreads an outlier over few coordinates. `"auto"` adds the round where `h < 512`. | One extra RPBH pass per activation, shared across the projections of one input. CUDA runs it as one Triton kernel: on an RTX 4090 at `d = 3840` it adds 40-50 µs per call at 256 tokens and 70-75 µs at 4096 tokens (the PyTorch path added 0.5 ms and 3.4 ms). MPS runs it as a Metal kernel (+0.06-0.18 ms at 256 tokens, +1.4-1.6 ms at 4096 on an M2 Max). The fused runtime, Krea 2 and the ExecuTorch Vulkan export (blocks up to 1024) support it. | +| `activation_quantizer` | `"auto"` (default), `"lloyd_max"` (paper), `"absmax_int8"` | `"auto"` resolves to `"absmax_int8"` at `activation_bits=8` and to `"lloyd_max"` below. `"absmax_int8"` is per-token absmax INT8 of the rotated input, the numerics of the fused runtime's 8-bit activations; the artifact stores no activation codebook. On Z-Image W4A8 it matched 8-bit Lloyd-Max under the same rotation (LPIPS 0.331 against 0.330, paired difference -0.007 to +0.007). Codebook version 3 builds 8-bit Lloyd-Max tables in about 0.1 s per width, including `d >= 7680`, where version 2 does not converge. | CUDA runs absmax A8 as an INT8 GEMM against the packed W2/W4 weights, the same GEMM class as W4A4. On an RTX 4090 at 4096 tokens an absmax W4A8 projection takes 3.4-6.6x less time than Lloyd-Max A8 (for example 1.01 ms against 6.71 ms with 12288 inputs and 3072 outputs). | | `adaln_scale` | `"absmax"` (paper implementation), `"mse"` | Per INT4 group, the BF16 scale with the lowest squared error from `absmax x [0.6, 1.15] / {7, 7.5, 8}`. | None: offline only, same packed format. Not recommended: it lowers the weight error but raised the AdaLN output error on 59 of 60 CogVideoX-2B projections (median 10x) and 17 of 32 Z-Image projections. | -| `output_gain_correction` | `False`, `True` | Divides the BF16 row norms by `E` of the weight and activation codebooks (`0.9905` per side at 4 bits, `0.8826` at 2). This helps outputs whose weight rows and inputs are aligned and hurts uncorrelated ones; compare generated samples before adopting it. | None. | +| `output_gain_correction` | `False` (default, paper), `True` | Divides the BF16 row norms by `E` of the weight and activation codebooks (`0.9905` per side at 4 bits, `0.8826` at 2). This helps outputs whose weight rows and inputs are aligned and hurts uncorrelated ones. On Z-Image it raised LPIPS against BF16 at W2A4 (+0.016, 95% interval +0.004 to +0.029, worse on 24 of 30 pairs) and did not lower it at W4A4, W4A8 or W3A3, so no setting enables it by default. | None. | | `outlier_channels` | `{module name: [channels]}` | Input channels kept out of the quantizer and applied as a BF16 side GEMM over their weight columns. `orbitquant.outliers.find_outlier_channels(model, run)` finds channels above `ratio x` the median token RMS from calibration forwards. | A `k`-column BF16 GEMM per projection; not available on the fused block runtime or in checkpoint-streaming conversion. | `outlier_channels` uses calibration forwards; the other options use no data. -Measured on Z-Image-Turbo (1024², 10 steps, 10 prompts x 2 seeds, LPIPS against BF16 on an -RTX 4090, 2026-10-29): W4A4 0.434, W4A4 + `rotation_rounds="auto"` 0.344 (better on 18 of 20 -pairs), W4A8 absmax 0.375, W4A8 absmax + `"auto"` 0.327, W3A3 0.535, W3A3 + `"auto"` 0.501. -`output_gain_correction`, `adaln_scale="mse"` and `outlier_channels` changed LPIPS by less than -the pair-to-pair spread (95% intervals include zero). `outlier_channels` cut the geometric-mean -output error of the 34 Z-Image FFN `w2` projections (input peak at a median 35000x the token -RMS) by 67%. +Measured on Z-Image-Turbo (1024², 10 steps, 10 prompts x 3 seeds, LPIPS against BF16 with the +`dequant_bf16` reference runtime on an RTX 4090, 2026-10-29; intervals are 95% paired +bootstrap): + +| Bits | Paper fields | Defaults | Defaults minus paper | Pairs better | +| --- | ---: | ---: | ---: | ---: | +| W4A4 | 0.438 | 0.348 | -0.090 (-0.124 to -0.060) | 27 of 30 | +| W2A4 | 0.643 | 0.632 | -0.011 (-0.029 to +0.012) | 22 of 30 | + +W4A8 with the default absmax INT8 activations reaches 0.331, 0.018 below the default W4A4 +(-0.032 to -0.004). W3A3 with the defaults reaches 0.505; an earlier 20-pair run measured 0.535 +with one round. W2A4 samples are visibly broken with either setting. `adaln_scale="mse"` and +`outlier_channels` changed LPIPS by less than the pair-to-pair spread (95% intervals include +zero). `outlier_channels` cut the geometric-mean output error of the 34 Z-Image FFN `w2` +projections (input peak at a median 35000x the token RMS) by 67%. + +FLUX.1-dev keeps one round everywhere (every block is at least 1024), so its default W4A4 is the +paper W4A4. On a GenEval subset (96 prompts, 16 per task, x 2 seeds, 1024², official evaluator, +RTX 6000 Ada, 2026-10-29) it scored 0.641 against 0.615 for BF16 (+0.026, 95% paired interval +-0.020 to +0.074), took 14.1 s per image instead of 21.5 s and peaked at 8.7 GiB instead of +25.1 GiB. A partial W4A8 run (35 images, single- and two-object prompts only) lowered LPIPS +against BF16 by 0.022 relative to W4A4 on the same images (-0.034 to -0.011). + +CogVideoX-2B (`d = 1920`, `h = 128`, so `"auto"` adds the round) shows no gain from it at W4A6: +on 12 VBench prompts the default and paper settings differ by less than the pair-to-pair spread +on all eight custom-input dimensions and on frame LPIPS against BF16 (0.339 against 0.326, +-0.022 to +0.052), while the default took 7% longer (1656 s against 1547 s) and peaked +1.2 GiB higher. Use `rotation_rounds=1` for such models when speed matters. ## Install @@ -214,8 +240,9 @@ path instead of the much slower FP32 compatibility path. An explicit dtype is always honored, and on-the-fly source conversion keeps the source dtype. Named recipes are `w4a4`, `w3a3`, `w2a4`, `w2a3`, `w4a6`, and `w4a8` (absmax INT8 -activations). They create a normal `OrbitQuantConfig`, so every field can be -overridden. +activations). They create a normal `OrbitQuantConfig` with the library defaults, so +every field can be overridden; `recipe(name, paper=True)` pins the paper fields listed under +[Quality options](#quality-options). For an unknown architecture resolved to `target_policy="universal"`, W2 uses mixed-bit protection by default: out-of-block and the first/last four indexed @@ -259,6 +286,28 @@ or patch embedder stays BF16, whatever the attribute is called. Configs saved before this rule carry no `target_policy_version` and load as version 1, the name-token classification they were quantized with. +On the published transformer configs (built on the meta device, 2026-10-29) version 2 +classifies: + +| Model | Linears | OrbitQuant | AdaLN INT4 | BF16 | AdaLN INT4 under version 1 | +| --- | ---: | ---: | ---: | ---: | ---: | +| FLUX.1-dev | 504 | 418 | 76 | 10 | 76 | +| FLUX.1-schnell | 502 | 418 | 76 | 8 | 76 | +| FLUX.2-dev | 203 | 192 | 3 | 8 | 3 | +| Z-Image-Turbo | 276 | 238 | 32 | 6 | 32 | +| Wan 2.1 T2V 1.3B | 306 | 300 | 0 | 6 | 0 | +| CogVideoX-2B | 245 | 180 | 60 | 5 | 0 | +| CogVideoX-5B | 341 | 252 | 84 | 5 | 0 | +| HunyuanVideo | 547 | 440 | 80 | 27 | 0 | +| SD3.5 Medium | 392 | 337 | 48 | 7 | 0 | +| SD3.5 Large | 536 | 453 | 76 | 7 | 0 | +| PixArt-Sigma XL | 286 | 280 | 1 | 5 | 0 | + +The modulation linears outside AdaLN INT4 are the final `norm_out.linear` +(`AdaLayerNormContinuous`) and the two HunyuanVideo token-refiner `norm_out.linear`; they stay +BF16. Wan 2.1 modulates through parameter tables and a shared time projection, so it has no +per-block AdaLN linear. + Qwen2.5-VL is supported by the universal policy, including W4A6 text and vision attention/FFN projections. Token embeddings, the language output head and the vision patch embedding remain unquantized. Packed loading is tested with both @@ -660,8 +709,8 @@ The implementation follows the paper's shared data-agnostic basis: GPTQ rounding (`orbitquant.gptq`) is not part of the paper's method: it uses calibration data to choose codes within the same codebook and format. The default quantizer never calls it. -The [quality options](#quality-options-opt-in) are also outside the paper recipe; their -defaults keep it. +The [quality options](#quality-options) are outside the paper recipe; +`recipe(name, paper=True)` keeps it. The detailed requirement matrix is in [`docs/paper-methodology-audit.md`](docs/paper-methodology-audit.md). diff --git a/docs/kernel-audit.md b/docs/kernel-audit.md index 2e5f455..72cd134 100644 --- a/docs/kernel-audit.md +++ b/docs/kernel-audit.md @@ -36,7 +36,7 @@ Other explicit modes are `native_packed_matmul`, `triton_packed_matmul`, | 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 | Hardware-verified native source and CI wheel builds; wheel publication pending | Runtime ISA dispatch across scalar, AVX2/FMA, AVX-512/BF16, and ARM64 NEON; exact packed activation, packed low-bit matmul, and packed INT4 group-64 AdaLN | -| Vulkan | Experimental source path; hardware-verified on AMD Cezanne/RADV Linux, Windows unverified | ExecuTorch whole-graph delegate adapter with exact packed W4A4 activation and matmul shaders | +| Vulkan | Experimental source path; W4A4 hardware-verified on AMD Cezanne/RADV Linux, W4A8 verified on Mesa lavapipe only, Windows unverified | ExecuTorch whole-graph delegate adapter with exact packed W4A4 Lloyd-Max and W4A8 absmax INT8 activation and matmul shaders | | ROCm | Experimental source candidate; supported-hardware proof pending | Exact Triton activation, low-bit pack/unpack, packed matmul, weight quantization, and packed INT4 AdaLN; explicit-only | | XPU | Experimental source candidate; Intel hardware proof pending | Exact Triton activation, low-bit pack/unpack, packed matmul, weight quantization, and packed INT4 AdaLN on `torch.xpu`; explicit-only | @@ -151,9 +151,11 @@ tuple. OrbitQuant rejects incompatible packages instead of loading them. The Vulkan path targets the ExecuTorch whole-graph delegate because desktop PyTorch does not expose a Vulkan tensor backend. It accepts exported FP16 or FP32 activations, keeps W4 weights and activation indices packed, and uses a -256-entry exact pair-product LUT. It does not create an FP16/BF16 weight -matrix. The source path currently implements W4A4 only and is not selected by -`auto_fused`. +256-entry exact pair-product LUT for W4A4. W4A8 with absmax INT8 activations +packs four signed activation codes per word and accumulates INT8 code times +FP32 weight centroid. Neither path creates an FP16/BF16 weight matrix. The +source path implements W4A4 with Lloyd-Max activations and W4A8 with absmax +INT8 activations and is not selected by `auto_fused`. Desktop Vulkan support requires an ExecuTorch source revision containing the June/July 2026 desktop delegate changes. The source below was verified at @@ -195,12 +197,17 @@ cmake --build .local-artifacts/orbitquant-vulkan-build \ ``` `prepare_executorch_vulkan_w4a4_model()` converts already-loaded -`OrbitQuantLinear` modules without requantizing or changing the artifact. -`ExecuTorchVulkanW4A4Linear` and -`register_executorch_vulkan_w4a4()` in -`orbitquant.kernels.executorch_vulkan` provide the Python export boundary. The -native package supplies token norm, RPBH/FWHT plus exact code assignment, and -direct packed pair-LUT matmul dispatches. +`OrbitQuantLinear` modules without requantizing or changing the artifact; it +selects `ExecuTorchVulkanW4A4Linear` for A4 Lloyd-Max layers and +`ExecuTorchVulkanW4A8Linear` for A8 absmax INT8 layers, and refuses other +activation quantizers. `register_executorch_vulkan_w4a4()` and +`register_executorch_vulkan_w4a8()` in `orbitquant.kernels.executorch_vulkan` +provide the Python export boundary. A two-round artifact exports its +pre-rotation, up to 1024-wide blocks, as a dense block-Hadamard matmul in +front of the custom operator. For W4A4 the native package supplies token norm, +RPBH/FWHT plus exact code assignment, and direct packed pair-LUT matmul +dispatches. For W4A8 it supplies RPBH/FWHT with per-block absmax, per-token +scale and INT8 packing, and a packed W4 by INT8 matmul with FP32 accumulation. The same build and an exported `.pte` were executed on an AMD Ryzen 5 5600G Cezanne integrated Radeon through Mesa 25.2.8 RADV, Vulkan 1.4.318, ExecuTorch @@ -254,9 +261,22 @@ packed prepack times were 2.925/0.725 ms at 1536 and 10.454/2.273 ms at 3072, versus 5.645/2.101 ms and 17.969/8.692 ms for FP16 upload. No full floating weight is created by the packed delegate. -The earlier llvmpipe run remains build and validation evidence only. Windows -Vulkan and other AMD architectures are unverified, the path is W4A4-only, and -it is not selected by `auto_fused`; these limitations keep Vulkan experimental. +The W4A8 shaders were built at the same ExecuTorch commit with Vulkan SDK +1.4.350.1 `glslc` on Ubuntu 24.04 and run on Mesa lavapipe, the CPU Vulkan +driver, because the GPU container exposed no NVIDIA graphics capability. All +12 C++ cases passed against the independent host reference: the seven W4A4 +cases above and five W4A8 cases covering FP32/FP16, bias/no-bias, a 9x24 +partial tile, and 2x1536 and 8x1536 DiT-width rows. The W4A8 cases use +absolute and relative tolerances of 8e-3 for FP16 and 5e-4 for FP32 against an +FP32 host reference. Python export of a +W4A8 layer matches the PyTorch absmax layer within 1e-5 and places +`orbitquant_vulkan.linear_w4a8_absmax` in the graph. + +The llvmpipe and lavapipe runs are build and correctness evidence only; their +timings are not GPU performance. W4A8 has no RADV hardware run and no `.pte` +runner run yet. Windows Vulkan and other AMD architectures are unverified, and +the path is not selected by `auto_fused`; these limitations keep Vulkan +experimental. Windows x86_64 source portability is covered by the path-filtered `Windows Vulkan build` workflow. At OrbitQuant commit `f41ae5e`, run diff --git a/docs/paper-methodology-audit.md b/docs/paper-methodology-audit.md index 705d72f..7d80c1c 100644 --- a/docs/paper-methodology-audit.md +++ b/docs/paper-methodology-audit.md @@ -69,7 +69,7 @@ Status legend: | Paper requirement | Current status | Evidence | Notes | | --- | --- | --- | --- | | No calibration data, prompt statistics, timestep ranges, or generated-image statistics are used to construct the quantizer. | Pass | `src/orbitquant/codebooks/lloyd_max.py`, `src/orbitquant/rotations/rpbh.py`, `src/orbitquant/functional.py`, `src/orbitquant/layers.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `tests/test_artifact_writer.py` | Codebooks depend on `(dim, bits, algorithm_version)`. Rotations depend on `(dim, seed, block_size)`. Activations use only runtime token norms, and serialized layer/artifact state excludes activation calibration/range tensors. | -| One converged Lloyd-Max scalar codebook is built offline per input dimension and bit width. | Pass for codebook version 2 | `get_codebook(dim, bits, algorithm_version)` in `src/orbitquant/codebooks/lloyd_max.py`; `tests/test_codebooks.py`; `tests/test_paper_reference_oracle.py`; `tests/test_artifact_writer.py` | Version 2 evaluates the exact beta marginal and iterates to the Lloyd-Max centroid condition. Independent numerical quadrature checks the centroid condition at 2, 3, 4, and 6 bits. Version 1 remains read-only compatibility for existing packed indices. Persistent cache keys include algorithm version, dimension, bits, checksum, and structural validation. | +| One converged Lloyd-Max scalar codebook is built offline per input dimension and bit width. | Pass for codebook versions 2 and 3 | `get_codebook(dim, bits, algorithm_version)` in `src/orbitquant/codebooks/lloyd_max.py`; `tests/test_codebooks.py`; `tests/test_paper_reference_oracle.py`; `tests/test_artifact_writer.py` | Version 2 evaluates the exact beta marginal and iterates to the Lloyd-Max centroid condition. Version 3, the default for new artifacts, solves the same fixed point with Newton steps on its tridiagonal Jacobian from a Panter-Dite start; it agrees with version 2 within one float32 ulp where version 2 converges (tested at 2, 4 and 6 bits up to `d = 3072`) and also converges at 8 bits for `d >= 7680`, where version 2 does not. Version 2 stays frozen because artifact validation compares stored tables with the declared version. Independent numerical quadrature checks the centroid condition at 2, 3, 4, and 6 bits. Version 1 remains read-only compatibility for existing packed indices. Persistent cache keys include algorithm version, dimension, bits, checksum, and structural validation. | | Lloyd-Max target distribution is the coordinate marginal of a random unit vector in dimension `d`. | Pass | `_coordinate_density()` in `src/orbitquant/codebooks/lloyd_max.py`; `tests/test_paper_reference_oracle.py` | The implementation uses the paper density up to normalization: `(1 - z^2)^((d - 3) / 2)` on `[-1, 1]`. The independent oracle integrates this density directly rather than calling the production incomplete-beta routines. | | Quantization uses nearest fixed centroids, with no zero-point, learned scale, per-channel range, or timestep/prompt range. | Pass | `LloydMaxCodebook.quantize_indices()` and `quantize()` in `src/orbitquant/codebooks/lloyd_max.py`; `tests/test_paper_reference_oracle.py` | Production uses `torch.bucketize` against midpoint boundaries. The independent Algorithm 1 oracle uses exhaustive absolute-distance `argmin` and obtains the same weight indices and layer outputs. | | RPBH uses uniform random permutation, Rademacher signs, per-block Walsh-Hadamard transform, and `1 / sqrt(block_size)` normalization. | Pass | `src/orbitquant/rotations/rpbh.py`, `src/orbitquant/rotations/fwht.py`, `tests/test_rpbh.py`, `tests/test_paper_reference_oracle.py` | The implementation applies permutation first, then signs, then block FWHT and normalization. The independent oracle constructs the dense matrix `P.T @ D @ H_block` from Equation (9) and checks the production transform and orthogonality. | @@ -173,16 +173,19 @@ from a paper projection to an unrelated module fails the gate. ## Options Outside The Paper Recipe -These config fields default to the paper recipe; a non-default value is a -documented deviation recorded in the artifact config (and in the manifest for -`rotation_rounds` and `activation_quantizer`). +A non-paper value of these fields is a documented deviation recorded in the +artifact config (and in the manifest for `rotation_rounds` and +`activation_quantizer`). `rotation_rounds` and `activation_quantizer` default to +`"auto"`, which deviates from the paper where it measured better; configs saved +without the field load the paper value, and `recipe(name, paper=True)` pins the +paper fields. | Field | Paper value | Deviation | | --- | --- | --- | -| `rotation_rounds` | `1` | `2`/`"auto"` composes an independent RPBH (seed `rotation_seed + 0x9E3779B9`) in front of the paper rotation. The product is orthogonal, so Equations (4)-(8) hold in the composed basis; the first round is folded into the weights and applied to activations before the paper-round kernels. | -| `activation_quantizer` | `"lloyd_max"` | `"absmax_int8"` replaces Equation (7)'s 8-bit codebook with per-token absmax INT8 of the rotated input. The artifact then carries no activation codebook. The exact 8-bit Lloyd-Max generator (`codebooks/lloyd_max.py`, tolerance `1e-13`, `6 * 4**8` iterations) does not converge for `d = 7680`, `10240` (measured 2026-10-29), so paper-recipe A8 is unavailable on those widths. | +| `rotation_rounds` | `1` | Default `"auto"` (a second round where the paper block is below 512; on Z-Image-Turbo W4A4 it lowered LPIPS against BF16 from 0.438 to 0.348, better on 27 of 30 prompt/seed pairs); `2`/`"auto"` composes an independent RPBH (seed `rotation_seed + 0x9E3779B9`) in front of the paper rotation. The product is orthogonal, so Equations (4)-(8) hold in the composed basis; the first round is folded into the weights and applied to activations before the paper-round kernels. | +| `activation_quantizer` | `"lloyd_max"` | Default `"auto"`: `"absmax_int8"` at 8 bits, `"lloyd_max"` otherwise. `"absmax_int8"` replaces Equation (7)'s 8-bit codebook with per-token absmax INT8 of the rotated input. The artifact then carries no activation codebook. The version 2 generator (tolerance `1e-13`, `6 * 4**8` iterations) does not converge at 8 bits for `d = 7680`, `10240`; version 3 builds those tables, so paper-recipe A8 is available. | | `adaln_scale` | not fixed by the paper | `"mse"` replaces the absmax/7 implementation contract with a per-group squared-error scale search. It minimises the weight error, not the AdaLN output error; on 2026-10-29 it raised the output error on 59 of 60 CogVideoX-2B AdaLN projections. | -| `output_gain_correction` | `False` | Divides the row norms of Equation (6) by the codebooks' retained energy `E`. | +| `output_gain_correction` | `False` | Divides the row norms of Equation (6) by the codebooks' retained energy `E`. Stays off by default: on Z-Image-Turbo (30 prompt/seed pairs, 2026-10-29) it raised LPIPS against BF16 at W2A4 by 0.016 (95% interval +0.004 to +0.029) and lowered it at none of W4A4, W4A8, W3A3. | | `outlier_channels` | empty | Keeps listed input channels out of Equations (7)-(8); this uses calibration forwards, so it is not data-agnostic. | ## Native Setting Provenance @@ -248,6 +251,17 @@ from compact summaries are rejected by the HF artifact audit. Full metric runs are required only before saying that an artifact reproduces the paper's GenEval or VBench numbers. +Subset runs from 2026-10-29 (not full-benchmark reproductions): + +- FLUX.1-dev, GenEval subset of 96 prompts (16 per task) x 2 seeds, official + evaluator: BF16 0.615, W4A4 0.641 (paired difference +0.026, 95% interval + -0.020 to +0.074). The default W4A4 equals the paper W4A4 on FLUX because + every block is at least 1024. +- CogVideoX-2B, 12 VBench prompts, eight custom-input dimensions: paper W4A6 + and default W4A6 (`rotation_rounds="auto"` adds a round at `d = 1920`) + differ by less than the pair-to-pair spread on every dimension and on frame + LPIPS against BF16; the default took 7% longer. W4A4 was not run. + Automatic coverage outside the paper models does not imply that W4A4 or another profile preserves task quality. Real GPT-2 and DeiT checks confirm that loading, rotation folding, quantization, execution, and serialization work, while their diff --git a/native-kernels/orbitquant-vulkan/OrbitLinear.cpp b/native-kernels/orbitquant-vulkan/OrbitLinear.cpp index bf42247..cf43db6 100644 --- a/native-kernels/orbitquant-vulkan/OrbitLinear.cpp +++ b/native-kernels/orbitquant-vulkan/OrbitLinear.cpp @@ -234,10 +234,143 @@ void orbitquant_linear_w4a4_exact( PushConstantDataInfo(&apply_bias, sizeof(apply_bias))})); } +void orbitquant_linear_w4a8_absmax( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef input = args.at(idx++); + const ValueRef packed_weight_data = args.at(idx++); + const ValueRef row_norms_data = args.at(idx++); + const ValueRef permutation_data = args.at(idx++); + const ValueRef signs_data = args.at(idx++); + const ValueRef weight_centroids_data = args.at(idx++); + const ValueRef block_size_ref = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx); + + const vkapi::ScalarType input_dtype = graph.dtype_of(input); + VK_CHECK_COND( + input_dtype == vkapi::kFloat || input_dtype == vkapi::kHalf, + "OrbitQuant Vulkan W4A8 requires float32 or float16 activations"); + VK_CHECK_COND(graph.dtype_of(output) == input_dtype); + VK_CHECK_COND(graph.dtype_of(packed_weight_data) == vkapi::kByte); + VK_CHECK_COND(graph.dtype_of(permutation_data) == vkapi::kInt); + VK_CHECK_COND(graph.dtype_of(signs_data) == vkapi::kInt); + check_float_constant(graph, row_norms_data); + check_float_constant(graph, weight_centroids_data); + + const std::vector input_sizes = graph.sizes_of(input); + VK_CHECK_COND(!input_sizes.empty()); + const int64_t K = input_sizes.back(); + const int64_t M = graph.numel_of(input) / K; + const int64_t N = logical_numel(graph, row_norms_data); + const int64_t block_size = graph.extract_scalar(block_size_ref); + + VK_CHECK_COND(K > 0 && M > 0 && N > 0); + VK_CHECK_COND(K % 8 == 0, "OrbitQuant Vulkan W4A8 requires K divisible by 8"); + VK_CHECK_COND( + block_size >= 8 && block_size <= 4096 && + (block_size & (block_size - 1)) == 0 && K % block_size == 0, + "OrbitQuant Vulkan W4A8 requires a power-of-two block_size in " + "[8, 4096] dividing K"); + VK_CHECK_COND( + logical_numel(graph, packed_weight_data) == N * K / 2, + "packed W4 weight length does not match N*K/2"); + VK_CHECK_COND(logical_numel(graph, permutation_data) == K); + VK_CHECK_COND(logical_numel(graph, signs_data) == K); + VK_CHECK_COND(logical_numel(graph, weight_centroids_data) == 16); + VK_CHECK_COND(graph.numel_of(output) == M * N); + + const ValueRef packed_weight = + prepack_w4_transposed(graph, packed_weight_data, N, K); + const ValueRef row_norms = prepack_buffer(graph, row_norms_data); + const ValueRef permutation = prepack_buffer(graph, permutation_data); + const ValueRef signs = prepack_buffer(graph, signs_data); + const ValueRef weight_centroids = + prepack_buffer(graph, weight_centroids_data); + + const int64_t blocks = K / block_size; + TmpTensor rotated( + &graph, {M * K}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + TmpTensor block_absmax( + &graph, {M * blocks}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + TmpTensor token_scales( + &graph, {M}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + TmpTensor packed_activations( + &graph, + {M * K / 4}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + + const utils::ivec4 problem_sizes = { + utils::safe_downcast(M), + utils::safe_downcast(N), + utils::safe_downcast(K), + utils::safe_downcast(block_size)}; + + const float inv_sqrt_block = 1.0f / std::sqrt(static_cast(block_size)); + std::string rpbh_kernel = + "orbitquant_rpbh_a8_b" + std::to_string(block_size); + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(dtype_shader_name(rpbh_kernel, input_dtype)), + {utils::safe_downcast(blocks) * kRpbhWorkers, + utils::safe_downcast(M), + 1u}, + {kRpbhWorkers, 1u, 1u}, + {{{rotated.vref, block_absmax.vref}, vkapi::kWrite}, + {{input, permutation, signs}, vkapi::kRead}}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes)), + PushConstantDataInfo(&inv_sqrt_block, sizeof(inv_sqrt_block))})); + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR("orbitquant_absmax_a8"), + {kNormWorkers, utils::safe_downcast(M), 1u}, + {kNormWorkers, 1u, 1u}, + {{{packed_activations.vref, token_scales.vref}, vkapi::kWrite}, + {{rotated.vref, block_absmax.vref}, vkapi::kRead}}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes))})); + + TmpTensor dummy_bias( + &graph, {}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + ValueRef bias = dummy_bias.vref; + int32_t apply_bias = 0; + if (graph.val_is_not_none(bias_data)) { + check_float_constant(graph, bias_data); + VK_CHECK_COND(logical_numel(graph, bias_data) == N); + bias = prepack_buffer(graph, bias_data); + apply_bias = 1; + } + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(dtype_shader_name("orbitquant_w4a8", input_dtype)), + {utils::safe_downcast(N), utils::safe_downcast(M), 1u}, + {64u, 1u, 1u}, + {{output, vkapi::kWrite}, + {{packed_activations.vref, + packed_weight, + token_scales.vref, + row_norms, + weight_centroids, + bias}, + vkapi::kRead}}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes)), + PushConstantDataInfo(&apply_bias, sizeof(apply_bias))})); +} + REGISTER_OPERATORS { VK_REGISTER_OP( orbitquant_vulkan.linear_w4a4_exact.default, orbitquant_linear_w4a4_exact); + VK_REGISTER_OP( + orbitquant_vulkan.linear_w4a8_absmax.default, + orbitquant_linear_w4a8_absmax); } } // namespace vkcompute diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_absmax_a8.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_absmax_a8.glsl new file mode 100644 index 0000000..5d5476d --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_absmax_a8.glsl @@ -0,0 +1,67 @@ +#version 450 core + +${define_required_extensions("buffer", "int")} +${define_active_storage_type("buffer")} + +#define PRECISION highp + +layout(std430) buffer; + +// Per-token absmax INT8 of the rotated values: scale = max(absmax, 1e-12) / 127 and +// code = clamp(roundEven(value / scale), -127, 127), four codes per word, word-major. +${layout_declare_tensor(B, "w", "t_packed", "int", "buffer")} +${layout_declare_tensor(B, "w", "t_scales", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_rotated", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_block_absmax", "float", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +shared float maxima[64]; + +void main() { + const uint row = gl_GlobalInvocationID.y; + const uint lane = gl_LocalInvocationID.x; + const uint rows = uint(problem_sizes.x); + const uint K = uint(problem_sizes.z); + const uint blocks = K / uint(problem_sizes.w); + + float largest = 0.0; + if (row < rows) { + for (uint block = lane; block < blocks; block += 64u) { + largest = max(largest, t_block_absmax[row * blocks + block]); + } + } + maxima[lane] = largest; + memoryBarrierShared(); + barrier(); + + for (uint stride = 32u; stride > 0u; stride >>= 1u) { + if (lane < stride) { + maxima[lane] = max(maxima[lane], maxima[lane + stride]); + } + memoryBarrierShared(); + barrier(); + } + if (row >= rows) { + return; + } + + const float scale = max(maxima[0], 1.0e-12) / 127.0; + if (lane == 0u) { + t_scales[row] = scale; + } + const uint words = K >> 2u; + for (uint word = lane; word < words; word += 64u) { + uint packed = 0u; + for (uint byte = 0u; byte < 4u; ++byte) { + const float value = t_rotated[row * K + word * 4u + byte]; + const int code = int(clamp(roundEven(value / scale), -127.0, 127.0)); + packed |= (uint(code) & 255u) << (byte * 8u); + } + t_packed[word * rows + row] = int(packed); + } +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_absmax_a8.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_absmax_a8.yaml new file mode 100644 index 0000000..3483a6d --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_absmax_a8.yaml @@ -0,0 +1,4 @@ +orbitquant_absmax_a8: + parameter_names_with_default_values: {} + shader_variants: + - NAME: orbitquant_absmax_a8 diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_a8.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_a8.glsl new file mode 100644 index 0000000..a3bd61e --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_a8.glsl @@ -0,0 +1,89 @@ +#version 450 core + +${define_required_extensions("buffer", DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} +#define BLOCK_SIZE ${BLOCK_SIZE} + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +// One RPBH block of one token: the rotated FP32 values and the block's absmax. +// Absmax INT8 activations keep the token scale, so the input is not normalized. +${layout_declare_tensor(B, "w", "t_rotated", "float", "buffer")} +${layout_declare_tensor(B, "w", "t_block_absmax", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_input", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_permutation", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_signs", "int", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; + float inv_sqrt_block; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +shared float values[BLOCK_SIZE]; +shared float maxima[64]; + +void main() { + const uint block = gl_WorkGroupID.x; + const uint row = gl_WorkGroupID.y; + const uint lane = gl_LocalInvocationID.x; + const uint rows = uint(problem_sizes.x); + const uint K = uint(problem_sizes.z); + const uint block_size = uint(problem_sizes.w); + if (row >= rows || block_size != uint(BLOCK_SIZE)) { + return; + } + + const uint block_offset = block * block_size; + for (uint i = lane; i < block_size; i += 64u) { + const uint rotated_index = block_offset + i; + const uint source_index = uint(t_permutation[rotated_index]); + const float sign = float(t_signs[rotated_index]); + values[i] = float(t_input[row * K + source_index]) * sign; + } + memoryBarrierShared(); + barrier(); + + for (uint span = 1u; span < block_size; span <<= 1u) { + const uint butterflies = block_size >> 1u; + for (uint i = lane; i < butterflies; i += 64u) { + const uint group = i / span; + const uint offset = i - group * span; + const uint left = group * (span << 1u) + offset; + const uint right = left + span; + const float a = values[left]; + const float b = values[right]; + values[left] = a + b; + values[right] = a - b; + } + memoryBarrierShared(); + barrier(); + } + + float largest = 0.0; + for (uint i = lane; i < block_size; i += 64u) { + const float value = values[i] * inv_sqrt_block; + t_rotated[row * K + block_offset + i] = value; + largest = max(largest, abs(value)); + } + maxima[lane] = largest; + memoryBarrierShared(); + barrier(); + + for (uint stride = 32u; stride > 0u; stride >>= 1u) { + if (lane < stride) { + maxima[lane] = max(maxima[lane], maxima[lane + stride]); + } + memoryBarrierShared(); + barrier(); + } + + if (lane == 0u) { + t_block_absmax[row * (K / block_size) + block] = maxima[0]; + } +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_a8.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_a8.yaml new file mode 100644 index 0000000..0a0d390 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_a8.yaml @@ -0,0 +1,29 @@ +orbitquant_rpbh_a8: + parameter_names_with_default_values: + BLOCK_SIZE: 8 + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: orbitquant_rpbh_a8_b8 + BLOCK_SIZE: 8 + - NAME: orbitquant_rpbh_a8_b16 + BLOCK_SIZE: 16 + - NAME: orbitquant_rpbh_a8_b32 + BLOCK_SIZE: 32 + - NAME: orbitquant_rpbh_a8_b64 + BLOCK_SIZE: 64 + - NAME: orbitquant_rpbh_a8_b128 + BLOCK_SIZE: 128 + - NAME: orbitquant_rpbh_a8_b256 + BLOCK_SIZE: 256 + - NAME: orbitquant_rpbh_a8_b512 + BLOCK_SIZE: 512 + - NAME: orbitquant_rpbh_a8_b1024 + BLOCK_SIZE: 1024 + - NAME: orbitquant_rpbh_a8_b2048 + BLOCK_SIZE: 2048 + - NAME: orbitquant_rpbh_a8_b4096 + BLOCK_SIZE: 4096 diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a8.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a8.glsl new file mode 100644 index 0000000..ed4db9f --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a8.glsl @@ -0,0 +1,63 @@ +#version 450 core + +${define_required_extensions("buffer", DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +// out[m, n] = scale[m] * row_norm[n] * sum_k code[m, k] * centroid[w[n, k]] (+ bias[n]), +// with FP32 accumulation of the signed INT8 activation codes. +${layout_declare_tensor(B, "w", "t_output", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_activations", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_weight", "uint", "buffer")} +${layout_declare_tensor(B, "r", "t_scales", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_row_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_centroids", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_bias", "float", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; + int apply_bias; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const uint n = gl_GlobalInvocationID.x; + const uint m = gl_GlobalInvocationID.y; + const uint M = uint(problem_sizes.x); + const uint N = uint(problem_sizes.y); + const uint K = uint(problem_sizes.z); + if (m >= M || n >= N) { + return; + } + + float centroids[16]; + for (uint index = 0u; index < 16u; ++index) { + centroids[index] = t_centroids[index]; + } + + float accumulator = 0.0; + const uint weight_words = K >> 3u; + for (uint word = 0u; word < weight_words; ++word) { + const uint weight_codes = t_weight[word * N + n]; + const int low = t_activations[(word * 2u) * M + m]; + const int high = t_activations[(word * 2u + 1u) * M + m]; + for (uint nibble = 0u; nibble < 8u; ++nibble) { + const int activations = nibble < 4u ? low : high; + const int code = bitfieldExtract(activations, int((nibble & 3u) * 8u), 8); + accumulator += + float(code) * centroids[(weight_codes >> (nibble * 4u)) & 15u]; + } + } + + float value = accumulator * t_scales[m] * t_row_norms[n]; + if (apply_bias != 0) { + value += t_bias[n]; + } + t_output[m * N + n] = T(value); +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a8.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a8.yaml new file mode 100644 index 0000000..0c6d31b --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a8.yaml @@ -0,0 +1,9 @@ +orbitquant_w4a8: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: orbitquant_w4a8 diff --git a/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp b/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp index 618c79e..3fb6029 100644 --- a/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp +++ b/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp @@ -186,6 +186,90 @@ TestCase make_test_case( return test_case; } +TestCase make_w4a8_test_case( + const TestShape shape, + const vkapi::ScalarType input_dtype, + const bool has_bias) { + const std::string dtype_name = + input_dtype == vkapi::kHalf ? "f16" : "f32"; + TestCase test_case( + std::string("w4a8_") + shape.name + "_" + dtype_name + + (has_bias ? "_bias" : "_no_bias")); + test_case.set_operator_name("orbitquant_vulkan.linear_w4a8_absmax.default"); + const float tolerance = input_dtype == vkapi::kHalf ? 8.0e-3f : 5.0e-4f; + test_case.set_abs_tolerance(tolerance); + test_case.set_rel_tolerance(tolerance); + test_case.set_op_invocations_per_execute(1); + + // Unnormalized inputs with one large channel: absmax INT8 keeps the token scale. + ValueSpec input = + tensor_spec({shape.rows, shape.input_dim}, input_dtype, false); + for (int32_t row = 0; row < shape.rows; ++row) { + for (int32_t k = 0; k < shape.input_dim; ++k) { + float value = std::sin(float(row * 19 + k * 7) * 0.13f) * (1.0f + row); + if (k == 5) { + value *= 40.0f; + } + const int32_t index = row * shape.input_dim + k; + if (input_dtype == vkapi::kHalf) { + input.get_half_data()[index] = float_to_half(value); + } else { + input.get_float_data()[index] = value; + } + } + } + + ValueSpec packed_weight = + tensor_spec({shape.outputs, shape.input_dim / 2}, vkapi::kByte, true); + for (int32_t n = 0; n < shape.outputs; ++n) { + for (int32_t k = 0; k < shape.input_dim; k += 2) { + const uint8_t low = static_cast((n * 3 + k * 5) & 15); + const uint8_t high = + static_cast((n * 3 + (k + 1) * 5) & 15); + packed_weight.get_uint8_data()[n * (shape.input_dim / 2) + k / 2] = + static_cast(low | (high << 4)); + } + } + + ValueSpec row_norms = tensor_spec({shape.outputs}, vkapi::kFloat, true); + for (int32_t n = 0; n < shape.outputs; ++n) { + row_norms.get_float_data()[n] = 0.55f + float(n % 17) * 0.007f; + } + + ValueSpec permutation = tensor_spec({shape.input_dim}, vkapi::kInt, true); + ValueSpec signs = tensor_spec({shape.input_dim}, vkapi::kInt, true); + for (int32_t k = 0; k < shape.input_dim; ++k) { + permutation.get_int32_data()[k] = (k * 5) % shape.input_dim; + signs.get_int32_data()[k] = (k % 3 == 0) ? -1 : 1; + } + + ValueSpec weight_centroids = tensor_spec({16}, vkapi::kFloat, true); + for (int32_t index = 0; index < 16; ++index) { + weight_centroids.get_float_data()[index] = -0.9f + float(index) * 0.12f; + } + + ValueSpec bias = tensor_spec({shape.outputs}, vkapi::kFloat, true); + if (has_bias) { + for (int32_t n = 0; n < shape.outputs; ++n) { + bias.get_float_data()[n] = float(n - 3) * 0.025f; + } + } else { + bias.set_none(true); + } + + test_case.add_input_spec(input); + test_case.add_input_spec(packed_weight); + test_case.add_input_spec(row_norms); + test_case.add_input_spec(permutation); + test_case.add_input_spec(signs); + test_case.add_input_spec(weight_centroids); + test_case.add_input_spec(ValueSpec(shape.block_size)); + test_case.add_input_spec(bias); + test_case.add_output_spec( + tensor_spec({shape.rows, shape.outputs}, input_dtype, false)); + return test_case; +} + TestCase make_fp16_linear_test_case(const TestShape shape) { TestCase test_case(std::string("fp16_linear_") + shape.name); test_case.set_operator_name("aten.linear.default"); @@ -236,6 +320,11 @@ std::vector generate_test_cases() { make_test_case(kDit1536Shape, vkapi::kFloat, false), make_test_case(kTiledPartialShape, vkapi::kHalf, true), make_test_case(kDit1536TiledShape, vkapi::kHalf, true), + make_w4a8_test_case(kSmallShape, vkapi::kFloat, false), + make_w4a8_test_case(kSmallShape, vkapi::kHalf, true), + make_w4a8_test_case(kDit1536Shape, vkapi::kFloat, true), + make_w4a8_test_case(kTiledPartialShape, vkapi::kHalf, false), + make_w4a8_test_case(kDit1536TiledShape, vkapi::kHalf, true), }; if (profile_enabled()) { const std::string filter = profile_case_filter(); @@ -291,6 +380,78 @@ void fwht(std::vector& values, const int32_t block_size) { } } +// Mirrors the shaders: FP32 rotation in the same butterfly order, the per-token +// absmax scale, round-half-to-even INT8 codes and FP32 accumulation. +void reference_compute_w4a8(TestCase& test_case) { + const ValueSpec& input_spec = test_case.inputs().at(0); + const auto& packed_weight = test_case.inputs().at(1).get_uint8_data(); + const auto& row_norms = test_case.inputs().at(2).get_float_data(); + const auto& permutation = test_case.inputs().at(3).get_int32_data(); + const auto& signs = test_case.inputs().at(4).get_int32_data(); + const auto& centroids = test_case.inputs().at(5).get_float_data(); + const int32_t block_size = test_case.inputs().at(6).get_int_value(); + const ValueSpec& bias_spec = test_case.inputs().at(7); + const auto& bias = bias_spec.get_float_data(); + const int32_t M = static_cast(input_spec.get_tensor_sizes().at(0)); + const int32_t K = static_cast(input_spec.get_tensor_sizes().at(1)); + const int32_t N = static_cast(row_norms.size()); + auto& output = test_case.outputs().at(0).get_ref_float_data(); + output.assign(M * N, 0.0f); + + const auto input_value = [&input_spec](const int32_t index) { + if (input_spec.dtype == vkapi::kHalf) { + return half_to_float(input_spec.get_half_data().at(index)); + } + return input_spec.get_float_data().at(index); + }; + + const float inv_sqrt_block = 1.0f / std::sqrt(float(block_size)); + std::vector codes(M * K); + std::vector scales(M); + for (int32_t row = 0; row < M; ++row) { + std::vector rotated(K); + float largest = 0.0f; + for (int32_t block = 0; block < K / block_size; ++block) { + std::vector values(block_size); + for (int32_t k = 0; k < block_size; ++k) { + const int32_t rotated_index = block * block_size + k; + values[k] = input_value(row * K + permutation[rotated_index]) * + float(signs[rotated_index]); + } + fwht(values, block_size); + for (int32_t k = 0; k < block_size; ++k) { + const float value = values[k] * inv_sqrt_block; + rotated[block * block_size + k] = value; + largest = std::max(largest, std::fabs(value)); + } + } + const float scale = std::max(largest, 1.0e-12f) / 127.0f; + scales[row] = scale; + for (int32_t k = 0; k < K; ++k) { + const float code = std::nearbyint(rotated[k] / scale); + codes[row * K + k] = + static_cast(std::min(127.0f, std::max(-127.0f, code))); + } + } + + for (int32_t row = 0; row < M; ++row) { + for (int32_t n = 0; n < N; ++n) { + float accumulator = 0.0f; + for (int32_t k = 0; k < K; ++k) { + const uint8_t packed = packed_weight[n * (K / 2) + k / 2]; + const uint8_t weight_code = + (k & 1) == 0 ? packed & 15 : (packed >> 4) & 15; + accumulator += float(codes[row * K + k]) * centroids[weight_code]; + } + float value = accumulator * scales[row] * row_norms[n]; + if (!bias_spec.is_none()) { + value += bias[n]; + } + output[row * N + n] = value; + } + } +} + void reference_compute(TestCase& test_case) { if (test_case.operator_name() == "aten.linear.default") { const ValueSpec& input = test_case.inputs().at(0); @@ -318,6 +479,12 @@ void reference_compute(TestCase& test_case) { return; } + if (test_case.operator_name() == + "orbitquant_vulkan.linear_w4a8_absmax.default") { + reference_compute_w4a8(test_case); + return; + } + const ValueSpec& input_spec = test_case.inputs().at(0); const auto& packed_weight = test_case.inputs().at(1).get_uint8_data(); const auto& row_norms = test_case.inputs().at(2).get_float_data(); @@ -413,7 +580,7 @@ int main() { const TestResult result = execute_test_cases( generate_test_cases, flop_count, - "OrbitQuant exact W4A4", + "OrbitQuant W4A4 and W4A8", profile ? 10 : 1, profile ? 31 : 3, reference_compute); diff --git a/src/orbitquant/benchmarks.py b/src/orbitquant/benchmarks.py index 271ed2a..4aa293a 100644 --- a/src/orbitquant/benchmarks.py +++ b/src/orbitquant/benchmarks.py @@ -252,10 +252,12 @@ def source_forward() -> torch.Tensor: return source(x) def activation_quant() -> torch.Tensor: + # The forward applies the first round of a two-round rotation before quantization. + rotated = x if quantized.pre_rotation is None else quantized._pre_rotate(x) if quantized.activation_codebook is None: - return quantized._absmax_int8_activations(x) + return quantized._absmax_int8_activations(rotated) return quantize_activations_kernel( - x, + rotated, rotation=quantized.rotation, codebook=quantized.activation_codebook, eps=quantized.activation_eps, diff --git a/src/orbitquant/codebooks/lloyd_max.py b/src/orbitquant/codebooks/lloyd_max.py index 1a6c081..3a6eb60 100644 --- a/src/orbitquant/codebooks/lloyd_max.py +++ b/src/orbitquant/codebooks/lloyd_max.py @@ -11,8 +11,8 @@ import torch from safetensors.torch import load_file, save_file -_ALGORITHM_VERSION = 2 -_SUPPORTED_ALGORITHM_VERSIONS = {1, 2} +_ALGORITHM_VERSION = 3 +_SUPPORTED_ALGORITHM_VERSIONS = {1, 2, 3} _DISABLE_CACHE_VALUES = {"1", "true", "yes", "on"} @@ -340,7 +340,7 @@ def _regularized_incomplete_beta(a: float, b: float, x: float) -> float: return 1.0 - factor * _beta_continued_fraction(b, a, 1.0 - x) / b -def _positive_coordinate_quantile(dim: int, probability: float) -> float: +def _positive_coordinate_quantile(dim: float, probability: float) -> float: beta_b = (dim - 1) / 2 lower = 0.0 upper = 1.0 @@ -471,6 +471,122 @@ def _generate_exact_v2_codebook(dim: int, bits: int) -> LloydMaxCodebook: ) +def _positive_interval_mass(dim: int, left: float, right: float) -> float: + beta_b = (dim - 1) / 2 + if right >= 1.0: + return 0.5 * _regularized_incomplete_beta(beta_b, 0.5, 1.0 - left * left) + return 0.5 * ( + _regularized_incomplete_beta(0.5, beta_b, right * right) + - _regularized_incomplete_beta(0.5, beta_b, left * left) + ) + + +def _positive_density(dim: int, value: float) -> float: + if value >= 1.0: + return 0.0 + beta_b = (dim - 1) / 2 + log_normalization = math.lgamma(dim / 2) - 0.5 * math.log(math.pi) - math.lgamma(beta_b) + return math.exp(log_normalization + (beta_b - 1.0) * math.log1p(-(value * value))) + + +def _centroid_residual( + dim: int, inner: list[float] +) -> tuple[list[float], list[float], list[float]]: + edges = [0.0, *inner, 1.0] + centroids = [ + _positive_interval_centroid(dim, edges[index], edges[index + 1]) + for index in range(len(edges) - 1) + ] + residual = [ + inner[index] - (centroids[index] + centroids[index + 1]) / 2 + for index in range(len(inner)) + ] + return edges, centroids, residual + + +def _generate_newton_v3_codebook(dim: int, bits: int) -> LloydMaxCodebook: + """The v2 fixed point (exact f_d centroids and midpoint boundaries), found by Newton. + + f_d is log-concave, so the Lloyd-Max quantizer is unique and v3 targets the same codes + as v2. v2's relaxed Lloyd iteration converges linearly: 8 bits take minutes and do not + converge for d >= 7680. The centroid c of [l, r] with mass m moves by f(l)(c - l)/m + and f(r)(r - c)/m, so the Jacobian of b_i - (c_i + c_(i+1))/2 is tridiagonal and + Newton converges in tens of steps from the Panter-Dite start (boundaries at + quantiles of f_d^(1/3), the same family at dim (d + 6) / 3). Where v2 stopped short + of the fixed point, the float32 tables differ from v2 by one ulp.""" + positive_levels = 2**bits // 2 + if positive_levels == 1: + positive_centroids = [_positive_interval_centroid(dim, 0.0, 1.0)] + else: + inner = [ + _positive_coordinate_quantile((dim + 6) / 3, index / positive_levels) + for index in range(1, positive_levels) + ] + edges, positive_centroids, residual = _centroid_residual(dim, inner) + norm = max(abs(value) for value in residual) + count = len(inner) + for _ in range(200): + moves_left, moves_right = [], [] + for index, centroid in enumerate(positive_centroids): + left, right = edges[index], edges[index + 1] + mass = _positive_interval_mass(dim, left, right) + moves_left.append(_positive_density(dim, left) * (centroid - left) / mass) + moves_right.append(_positive_density(dim, right) * (right - centroid) / mass) + # Thomas algorithm on rows sub * b_(i-1) + diagonal * b_i + sup * b_(i+1). + factors, solved = [0.0] * count, [0.0] * count + for index in range(count): + sub = -moves_left[index] / 2 if index else 0.0 + diagonal = 1.0 - (moves_right[index] + moves_left[index + 1]) / 2 + denominator = diagonal - (sub * factors[index - 1] if index else 0.0) + if index < count - 1: + factors[index] = -moves_right[index + 1] / 2 / denominator + previous_solved = solved[index - 1] if index else 0.0 + solved[index] = (residual[index] - sub * previous_solved) / denominator + step = [0.0] * count + for index in reversed(range(count)): + following = step[index + 1] if index < count - 1 else 0.0 + step[index] = solved[index] - factors[index] * following + scale = 1.0 + while True: + candidate = [inner[index] - scale * step[index] for index in range(count)] + bounds = [0.0, *candidate, 1.0] + if all(bounds[index] < bounds[index + 1] for index in range(count + 1)): + trial = _centroid_residual(dim, candidate) + trial_norm = max(abs(value) for value in trial[2]) + if trial_norm < norm or scale < 1e-3: + break + scale /= 2 + largest_step = scale * max(abs(value) for value in step) + previous_norm = norm + inner = candidate + edges, positive_centroids, residual = trial + norm = trial_norm + if largest_step <= 4e-16 * inner[-1] or norm >= previous_norm: + break + # Interval masses near the tail are differences of CDFs close to 1, which puts the + # noise floor near 1e-12 of the boundary scale; 2**-32 keeps the residual 256 times + # below the float32 resolution of the tables. + if norm > 2.0**-32 * inner[-1]: + raise RuntimeError( + f"Lloyd-Max Newton solve did not converge for dim={dim}, bits={bits}: " + f"residual {norm:.3e}" + ) + + centroids = torch.tensor( + [-value for value in reversed(positive_centroids)] + positive_centroids, + dtype=torch.float64, + device="cpu", + ) + boundaries = (centroids[:-1] + centroids[1:]) / 2 + return LloydMaxCodebook( + dim=dim, + bits=bits, + centroids=centroids.to(torch.float32), + boundaries=boundaries.to(torch.float32), + algorithm_version=3, + ) + + def _generate_codebook( dim: int, bits: int, algorithm_version: int = _ALGORITHM_VERSION ) -> LloydMaxCodebook: @@ -487,6 +603,8 @@ def _generate_codebook( ) if algorithm_version == 2: return _generate_exact_v2_codebook(dim, bits) + if algorithm_version == 3: + return _generate_newton_v3_codebook(dim, bits) raise ValueError( f"algorithm_version must be one of {sorted(_SUPPORTED_ALGORITHM_VERSIONS)}" ) diff --git a/src/orbitquant/config.py b/src/orbitquant/config.py index 1ec8c6e..d781ff6 100644 --- a/src/orbitquant/config.py +++ b/src/orbitquant/config.py @@ -79,12 +79,13 @@ class OrbitQuantConfig(QuantizationConfigMixin): rotation_seed: int = 0 block_size: int | str = "paper" # 1 is the paper's single RPBH round. 2 adds an independent first RPBH round - # in front of it on every projection; "auto" adds it only where the block is - # narrower than 512 (d = 3360 gives h = 32, d = 3840 gives h = 256). See - # orbitquant.rotations.rpbh. Not supported by the fused block runtime. - rotation_rounds: int | str = 1 + # in front of it on every projection; "auto" (default) adds it only where the + # block is narrower than 512 (d = 3360 gives h = 32, d = 3840 gives h = 256). + # On Z-Image W4A4 "auto" lowers LPIPS against BF16 from 0.438 to 0.348. See + # orbitquant.rotations.rpbh. Loaded artifacts without the field are read as 1. + rotation_rounds: int | str = "auto" codebook: str = "lloyd_max" - codebook_version: int = 2 + codebook_version: int = 3 codebook_dtype: str = "float32" row_norm_dtype: str = "bfloat16" activation_norm_dtype: str = "float32" @@ -92,8 +93,11 @@ class OrbitQuantConfig(QuantizationConfigMixin): # "lloyd_max" is the paper's activation codebook. "absmax_int8" (8-bit # activations only) quantizes the rotated input to INT8 with one absmax scale # per token, the numerics of the fused runtime's 8-bit activations, on every - # backend; CUDA runs it as an INT8 GEMM against the packed weights. - activation_quantizer: str = "lloyd_max" + # backend; CUDA runs it as an INT8 GEMM against the packed weights. "auto" + # (default) resolves to "absmax_int8" at 8 bits, where CUDA has no Lloyd-Max + # INT8 GEMM, and to "lloyd_max" below. Loaded artifacts without the field + # are read as "lloyd_max". + activation_quantizer: str = "auto" weight_pack_dtype: str = "uint8" target_policy: str = "auto" # Layer-classification revision. 1 detects AdaLN projections by module-name @@ -114,7 +118,8 @@ class OrbitQuantConfig(QuantizationConfigMixin): # activation codebooks. Lloyd-Max codes shrink a unit vector by that factor # (0.9905 per side at 4 bits, 0.9655 at 3, 0.8826 at 2), which scales an # aligned output by its product; for uncorrelated weight and input rows the - # uncorrected output is already the conditional mean. + # uncorrected output is already the conditional mean. On Z-Image it raised + # LPIPS against BF16 at W2A4 and lowered it at no tested width, so it stays off. output_gain_correction: bool = False modules_to_convert: list[str] = field(default_factory=list) modules_to_use_adaln: list[str] = field(default_factory=list) @@ -171,14 +176,18 @@ def __post_init__(self) -> None: raise ValueError("rotation_rounds must be 1, 2 or 'auto'") if self.codebook != "lloyd_max": raise ValueError("only Lloyd-Max codebooks are implemented") - if self.codebook_version not in {1, 2}: - raise ValueError("codebook_version must be 1 or 2") + if self.codebook_version not in {1, 2, 3}: + raise ValueError("codebook_version must be 1, 2 or 3") if self.codebook_dtype.lower() not in _SUPPORTED_CODEBOOK_DTYPES: raise ValueError("codebook_dtype must be 'float32'") if self.row_norm_dtype.lower() not in _SUPPORTED_ROW_NORM_DTYPES: raise ValueError("row_norm_dtype must be 'bfloat16'") if self.activation_norm_dtype.lower() not in _SUPPORTED_ACTIVATION_NORM_DTYPES: raise ValueError("activation_norm_dtype must be 'float32'") + if self.activation_quantizer == "auto": + self.activation_quantizer = ( + "absmax_int8" if self.activation_bits == 8 else "lloyd_max" + ) if self.activation_quantizer not in _SUPPORTED_ACTIVATION_QUANTIZERS: raise ValueError( "activation_quantizer must be one of " @@ -284,4 +293,8 @@ def from_dict(cls, data: dict[str, Any]) -> OrbitQuantConfig: # Artifacts written before parent-class AdaLN detection were classified # by module-name tokens only. values.setdefault("target_policy_version", 1) + # Artifacts written before these options existed used the paper's single + # RPBH round and Lloyd-Max activation codes. + values.setdefault("rotation_rounds", 1) + values.setdefault("activation_quantizer", "lloyd_max") return cls(**values) diff --git a/src/orbitquant/fused/groups.py b/src/orbitquant/fused/groups.py index f2ff170..445dacd 100644 --- a/src/orbitquant/fused/groups.py +++ b/src/orbitquant/fused/groups.py @@ -224,6 +224,7 @@ def signature(layer): layer.activation_bits, layer.weight_codebook.algorithm_version, layer.rotation, + layer.pre_rotation, ) for layer in layers: @@ -234,11 +235,6 @@ def signature(layer): "the fused runtime quantizes every input channel; projections with " "outlier_channels run on the per-layer OrbitQuantLinear path" ) - if layer.pre_rotation is not None: - raise ValueError( - "the fused runtime applies one RPBH round; projections quantized with " - "rotation_rounds=2 run on the per-layer OrbitQuantLinear path" - ) if signature(layer) != signature(first): raise ValueError("projections of one group must share width and rotation") segments = cls.source_segments(layers, sources, swiglu=swiglu) diff --git a/src/orbitquant/fused/install.py b/src/orbitquant/fused/install.py index 9296f18..89d33e9 100644 --- a/src/orbitquant/fused/install.py +++ b/src/orbitquant/fused/install.py @@ -126,6 +126,8 @@ def fuse( modules[name] = group_cls.from_sources( block, list(spec.sources), swiglu=spec.swiglu, activation_bits=widths.get(name) ) + if group_cls is PackedGroup: + runtime.check_pre_rotation(_resolve(block, spec.sources[0].path)) prepare = getattr(adapter, "prepare", None) if prepare is not None: prepare(block, kind) diff --git a/src/orbitquant/fused/linear.py b/src/orbitquant/fused/linear.py index ba3b607..d92fd28 100644 --- a/src/orbitquant/fused/linear.py +++ b/src/orbitquant/fused/linear.py @@ -55,6 +55,7 @@ def fuse_linears(model: nn.Module, *, activation_bits: int = 8, config=None) -> group = PackedGroup.from_sources( holder, [RowSource("proj")], activation_bits=activation_bits ) + runtime.check_pre_rotation(module) parent_path, _, child = name.rpartition(".") parent = model.get_submodule(parent_path) if parent_path else model fused = FusedLinear(group, runtime) diff --git a/src/orbitquant/fused/runtime.py b/src/orbitquant/fused/runtime.py index 372d8e0..1981570 100644 --- a/src/orbitquant/fused/runtime.py +++ b/src/orbitquant/fused/runtime.py @@ -24,7 +24,7 @@ from orbitquant.kernels import triton_activation as act from orbitquant.kernels import triton_int8_gemm as gemm from orbitquant.kernels.triton_attention import int8_attention_with_triton -from orbitquant.rotations import get_rpbh_rotation +from orbitquant.rotations import get_pre_rotation, get_rpbh_rotation ATTENTION_BACKENDS = ("int8", "flash", "cudnn", "sdpa") @@ -68,11 +68,13 @@ def __init__( codebook_version: int, activation_eps: float, attention: str = "int8", + rotation_rounds: int | str = 1, ): if attention not in ATTENTION_BACKENDS: raise ValueError(f"unknown attention backend {attention!r}") self.rotation_seed = rotation_seed self.block_size = block_size + self.rotation_rounds = rotation_rounds self.codebook_version = codebook_version self.activation_eps = activation_eps self.attention_backend = attention @@ -81,17 +83,13 @@ def __init__( @classmethod def from_config(cls, config, *, attention: str = "int8") -> FusedRuntime: - if getattr(config, "rotation_rounds", 1) != 1: - raise ValueError( - "the fused runtime applies one RPBH round; rotation_rounds=" - f"{config.rotation_rounds!r} needs the per-layer OrbitQuantLinear path" - ) return cls( rotation_seed=config.rotation_seed, block_size=config.block_size, codebook_version=config.codebook_version, activation_eps=config.activation_eps, attention=attention, + rotation_rounds=getattr(config, "rotation_rounds", 1), ) def quantizer(self, width: int, bits: int, device: torch.device) -> act.ActivationQuantizer: @@ -103,10 +101,42 @@ def quantizer(self, width: int, bits: int, device: torch.device) -> act.Activati ) # 8-bit activations are per-token absmax INT8 of the rotated input, no codebook. codebook = None if bits == 8 else get_codebook(width, bits, self.codebook_version) - quantizer = act.ActivationQuantizer(rotation, codebook, self.activation_eps, device) + pre_rotation = get_pre_rotation( + width, + seed=self.rotation_seed, + block_size=self.block_size, + rotation_rounds=self.rotation_rounds, + ) + quantizer = act.ActivationQuantizer( + rotation, codebook, self.activation_eps, device, pre_rotation=pre_rotation + ) self._quantizers[key] = quantizer return quantizer + def check_pre_rotation(self, layer) -> None: + """Refuse a projection whose first RPBH round is not the one this runtime applies to + its width: its weights would meet differently rotated activations.""" + expected = get_pre_rotation( + layer.in_features, + seed=self.rotation_seed, + block_size=self.block_size, + rotation_rounds=self.rotation_rounds, + ) + actual = layer.pre_rotation + if expected is None and actual is None: + return + if ( + expected is None + or actual is None + or expected.block_size != actual.block_size + or not torch.equal(expected.permutation, actual.permutation) + or not torch.equal(expected.signs, actual.signs) + ): + raise ValueError( + "the projection's rotation_rounds pre-rotation does not match the runtime " + f"config (rotation_rounds={self.rotation_rounds!r})" + ) + def release(self) -> None: """Drop device buffers before the weights leave the GPU.""" self._quantizers.clear() diff --git a/src/orbitquant/gptq.py b/src/orbitquant/gptq.py index cd690de..da585b9 100644 --- a/src/orbitquant/gptq.py +++ b/src/orbitquant/gptq.py @@ -97,7 +97,7 @@ def gptq_round( hessian: torch.Tensor, *, bits: int, - codebook_version: int = 2, + codebook_version: int = 3, damp: float = 0.01, block: int = 128, act_order: bool = True, diff --git a/src/orbitquant/kernels/executorch_vulkan.py b/src/orbitquant/kernels/executorch_vulkan.py index 5a072d3..1c4b0a1 100644 --- a/src/orbitquant/kernels/executorch_vulkan.py +++ b/src/orbitquant/kernels/executorch_vulkan.py @@ -11,6 +11,9 @@ if TYPE_CHECKING: from orbitquant.layers import OrbitQuantLinear +# A dense Hadamard is pre_block_size**2 FP32 values per layer. +_MAX_PRE_ROTATION_BLOCK = 1024 + def _linear_w4a4_exact_reference( x: torch.Tensor, @@ -113,32 +116,113 @@ def _linear_w4a4_exact_fake( return x.new_empty((*x.shape[:-1], row_norms.numel())) -class ExecuTorchVulkanW4A4Linear(nn.Module): - """Export-only exact W4A4 OrbitLinear for the ExecuTorch Vulkan delegate.""" +def _linear_w4a8_absmax_reference( + x: torch.Tensor, + packed_weight: torch.Tensor, + row_norms: torch.Tensor, + permutation: torch.Tensor, + signs: torch.Tensor, + weight_centroids: torch.Tensor, + block_size: int, + bias: torch.Tensor | None, +) -> torch.Tensor: + # The numerics of OrbitQuantLinear's absmax_int8 path: per-token absmax INT8 of the + # rotated input (no token normalization), then the INT8 codes against W4 centroids. + in_features = x.shape[-1] + out_features = row_norms.numel() + work = x.to(dtype=torch.float32) + rotated = work.index_select(-1, permutation.to(dtype=torch.int64)) + rotated = rotated * signs.to(dtype=torch.float32) + rotated = rotated.reshape(*rotated.shape[:-1], in_features // block_size, block_size) + rotated = fwht(rotated) * (block_size**-0.5) + rotated = rotated.reshape(-1, in_features) + scale = rotated.abs().amax(dim=-1, keepdim=True).clamp_min(1e-12) / 127.0 + codes = torch.round(rotated / scale).clamp_(-127.0, 127.0) + + weight_indices = unpack_lowbit( + packed_weight, + bits=4, + length=out_features * in_features, + ).to(dtype=torch.int64) + weight = weight_centroids.to(dtype=torch.float32)[ + weight_indices.reshape(out_features, in_features) + ] + output = (codes @ weight.T) * scale * row_norms.to(dtype=torch.float32) + if bias is not None: + output = output + bias.to(dtype=torch.float32) + output = output.reshape(*x.shape[:-1], out_features) + return output.to(dtype=x.dtype) + + +@torch.library.custom_op("orbitquant_vulkan::linear_w4a8_absmax", mutates_args=()) +def linear_w4a8_absmax( + x: torch.Tensor, + packed_weight: torch.Tensor, + row_norms: torch.Tensor, + permutation: torch.Tensor, + signs: torch.Tensor, + weight_centroids: torch.Tensor, + block_size: int, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + return _linear_w4a8_absmax_reference( + x, + packed_weight, + row_norms, + permutation, + signs, + weight_centroids, + block_size, + bias, + ) + + +@linear_w4a8_absmax.register_fake +def _linear_w4a8_absmax_fake( + x: torch.Tensor, + packed_weight: torch.Tensor, + row_norms: torch.Tensor, + permutation: torch.Tensor, + signs: torch.Tensor, + weight_centroids: torch.Tensor, + block_size: int, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + del packed_weight, permutation, signs, weight_centroids, block_size, bias + return x.new_empty((*x.shape[:-1], row_norms.numel())) + + +class _ExecuTorchVulkanLinear(nn.Module): + """Packed W4 weights, the RPBH tables and the optional first RPBH round, shared by the + export-only Vulkan linears.""" def __init__(self, layer: OrbitQuantLinear) -> None: super().__init__() - if layer.weight_bits != 4 or layer.activation_bits != 4: - raise ValueError("ExecuTorch Vulkan currently implements exact W4A4 only") + if layer.weight_bits != 4: + raise ValueError("ExecuTorch Vulkan currently implements W4 weights only") + if layer.outlier_channels is not None: + raise ValueError("ExecuTorch Vulkan quantizes every input channel; no outlier_channels") + pre_rotation = layer.pre_rotation + if pre_rotation is not None and pre_rotation.block_size > _MAX_PRE_ROTATION_BLOCK: + raise ValueError( + "ExecuTorch Vulkan applies the rotation_rounds pre-rotation as a dense Hadamard " + f"matmul, up to {_MAX_PRE_ROTATION_BLOCK}-wide blocks" + ) if layer.packed_weight_indices is None or layer.row_norms is None: raise ValueError("OrbitQuantLinear is missing packed weights or row norms") if layer.in_features % 8 != 0: - raise ValueError("ExecuTorch Vulkan W4A4 requires in_features divisible by 8") + raise ValueError("ExecuTorch Vulkan requires in_features divisible by 8") block_size = int(layer.rotation.block_size) if block_size < 8 or block_size > 4096: - raise ValueError( - "ExecuTorch Vulkan W4A4 currently requires an RPBH block size in [8, 4096]" - ) + raise ValueError("ExecuTorch Vulkan currently requires an RPBH block size in [8, 4096]") if block_size & (block_size - 1) or layer.in_features % block_size != 0: raise ValueError( - "ExecuTorch Vulkan W4A4 requires a power-of-two RPBH block size " - "dividing in_features" + "ExecuTorch Vulkan requires a power-of-two RPBH block size dividing in_features" ) self.in_features = layer.in_features self.out_features = layer.out_features self.block_size = block_size - self.eps = float(layer.activation_eps) self.register_buffer( "packed_weight", layer.packed_weight_indices.detach().to(device="cpu", dtype=torch.uint8).contiguous(), @@ -155,6 +239,65 @@ def __init__(self, layer: OrbitQuantLinear) -> None: "signs", layer.rotation.signs.detach().to(device="cpu", dtype=torch.int32).contiguous(), ) + # The first round of a two-round RPBH runs as ordinary exportable ops in front of the + # custom op: a gather, the signs and one Hadamard matmul per block, in FP32 like + # RPBHRotation.apply_to_activations. + self.pre_block_size = 0 if pre_rotation is None else int(pre_rotation.block_size) + if pre_rotation is None: + for name in ("pre_permutation", "pre_signs", "pre_hadamard"): + self.register_buffer(name, None) + else: + self.register_buffer( + "pre_permutation", + pre_rotation.permutation.detach().to(device="cpu", dtype=torch.int64).contiguous(), + ) + self.register_buffer( + "pre_signs", + pre_rotation.signs.detach().to(device="cpu", dtype=torch.float32).contiguous(), + ) + hadamard = fwht(torch.eye(self.pre_block_size, dtype=torch.float32)) + self.register_buffer( + "pre_hadamard", (hadamard * pre_rotation.normalization).contiguous() + ) + if layer.bias is None: + self.register_buffer("bias", None) + else: + self.register_buffer( + "bias", + layer.bias.detach().to(device="cpu", dtype=torch.float32).contiguous(), + ) + + def _prepare_input(self, x: torch.Tensor) -> torch.Tensor: + if x.dtype not in {torch.float16, torch.float32}: + raise TypeError( + "ExecuTorch Vulkan accepts float16 or float32 activations; " + f"got {x.dtype}. Convert the exported model explicitly." + ) + if x.shape[-1] != self.in_features: + raise ValueError(f"expected last dimension {self.in_features}, got {x.shape[-1]}") + if self.pre_hadamard is None: + return x + rotated = x.to(torch.float32).index_select(-1, self.pre_permutation) * self.pre_signs + rotated = rotated.reshape( + *x.shape[:-1], self.in_features // self.pre_block_size, self.pre_block_size + ) + return torch.matmul(rotated, self.pre_hadamard).reshape(x.shape).to(x.dtype) + + +class ExecuTorchVulkanW4A4Linear(_ExecuTorchVulkanLinear): + """Export-only exact W4A4 OrbitLinear for the ExecuTorch Vulkan delegate.""" + + def __init__(self, layer: OrbitQuantLinear) -> None: + if ( + layer.weight_bits != 4 + or layer.activation_bits != 4 + or layer.activation_quantizer != "lloyd_max" + ): + raise ValueError( + "ExecuTorchVulkanW4A4Linear implements exact W4A4 with Lloyd-Max activations" + ) + super().__init__(layer) + self.eps = float(layer.activation_eps) activation_centroids = layer.activation_codebook.centroids.detach().to( device="cpu", dtype=torch.float32 ) @@ -171,22 +314,9 @@ def __init__(self, layer: OrbitQuantLinear) -> None: "pair_lut", (activation_centroids[:, None] * weight_centroids[None, :]).flatten().contiguous(), ) - if layer.bias is None: - self.register_buffer("bias", None) - else: - self.register_buffer( - "bias", - layer.bias.detach().to(device="cpu", dtype=torch.float32).contiguous(), - ) def forward(self, x: torch.Tensor) -> torch.Tensor: - if x.dtype not in {torch.float16, torch.float32}: - raise TypeError( - "ExecuTorch Vulkan accepts float16 or float32 activations; " - f"got {x.dtype}. Convert the exported model explicitly." - ) - if x.shape[-1] != self.in_features: - raise ValueError(f"expected last dimension {self.in_features}, got {x.shape[-1]}") + x = self._prepare_input(x) return linear_w4a4_exact( x, self.packed_weight, @@ -201,20 +331,69 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) +class ExecuTorchVulkanW4A8Linear(_ExecuTorchVulkanLinear): + """Export-only W4A8 OrbitLinear with per-token absmax INT8 activations for the ExecuTorch + Vulkan delegate.""" + + def __init__(self, layer: OrbitQuantLinear) -> None: + if ( + layer.weight_bits != 4 + or layer.activation_bits != 8 + or layer.activation_quantizer != "absmax_int8" + ): + raise ValueError( + "ExecuTorchVulkanW4A8Linear implements W4A8 with absmax INT8 activations" + ) + super().__init__(layer) + self.register_buffer( + "weight_centroids", + layer.weight_codebook.centroids.detach() + .to(device="cpu", dtype=torch.float32) + .contiguous(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self._prepare_input(x) + return linear_w4a8_absmax( + x, + self.packed_weight, + self.row_norms, + self.permutation, + self.signs, + self.weight_centroids, + self.block_size, + self.bias, + ) + + +def executorch_vulkan_linear(layer: OrbitQuantLinear) -> _ExecuTorchVulkanLinear: + """The export-only Vulkan module for one loaded OrbitQuantLinear: exact W4A4 with + Lloyd-Max activations, or W4A8 with absmax INT8 activations.""" + + if layer.activation_bits == 8 and layer.activation_quantizer == "absmax_int8": + return ExecuTorchVulkanW4A8Linear(layer) + if layer.activation_bits == 4 and layer.activation_quantizer == "lloyd_max": + return ExecuTorchVulkanW4A4Linear(layer) + raise ValueError( + "ExecuTorch Vulkan implements W4A4 with Lloyd-Max activations and W4A8 with absmax " + f"INT8 activations; got A{layer.activation_bits} {layer.activation_quantizer}" + ) + + def prepare_executorch_vulkan_w4a4_model(model: nn.Module) -> nn.Module: - """Replace loaded OrbitQuant W4A4 linears with export-only Vulkan modules.""" + """Replace loaded OrbitQuant W4A4 and W4A8 linears with export-only Vulkan modules.""" from orbitquant.layers import OrbitQuantLinear if isinstance(model, OrbitQuantLinear): - return ExecuTorchVulkanW4A4Linear(model) + return executorch_vulkan_linear(model) def replace_children(module: nn.Module, prefix: str) -> None: for child_name, child in list(module.named_children()): qualified_name = f"{prefix}.{child_name}" if prefix else child_name if isinstance(child, OrbitQuantLinear): try: - replacement = ExecuTorchVulkanW4A4Linear(child) + replacement = executorch_vulkan_linear(child) except (TypeError, ValueError) as exc: raise ValueError( f"cannot prepare {qualified_name!r} for ExecuTorch Vulkan: {exc}" @@ -227,9 +406,7 @@ def replace_children(module: nn.Module, prefix: str) -> None: return model -def register_executorch_vulkan_w4a4() -> object: - """Register the custom op with the current ExecuTorch Vulkan partitioner.""" - +def _register_executorch_vulkan_op(name: str) -> object: try: from executorch.backends.vulkan import utils from executorch.backends.vulkan.op_registry import ( @@ -244,7 +421,7 @@ def register_executorch_vulkan_w4a4() -> object: "desktop Vulkan support (upstream after 2026-07-01)." ) from exc - edge_op = exir_ops.edge.orbitquant_vulkan.linear_w4a4_exact.default + edge_op = getattr(exir_ops.edge.orbitquant_vulkan, name).default if edge_op not in vulkan_supported_ops: @update_features(edge_op) @@ -260,9 +437,25 @@ def register_features() -> OpFeatures: return edge_op +def register_executorch_vulkan_w4a4() -> object: + """Register the exact W4A4 custom op with the current ExecuTorch Vulkan partitioner.""" + + return _register_executorch_vulkan_op("linear_w4a4_exact") + + +def register_executorch_vulkan_w4a8() -> object: + """Register the absmax W4A8 custom op with the current ExecuTorch Vulkan partitioner.""" + + return _register_executorch_vulkan_op("linear_w4a8_absmax") + + __all__ = [ "ExecuTorchVulkanW4A4Linear", + "ExecuTorchVulkanW4A8Linear", + "executorch_vulkan_linear", "linear_w4a4_exact", + "linear_w4a8_absmax", "prepare_executorch_vulkan_w4a4_model", "register_executorch_vulkan_w4a4", + "register_executorch_vulkan_w4a8", ] diff --git a/src/orbitquant/kernels/mps.py b/src/orbitquant/kernels/mps.py index 2539282..73b7f32 100644 --- a/src/orbitquant/kernels/mps.py +++ b/src/orbitquant/kernels/mps.py @@ -119,6 +119,84 @@ } } +// First round of a two-round RPBH (rotation_rounds=2): x[perm] * signs, block FWHT, scaled, +// in FP32 with the stage order of RPBHRotation.apply_to_activations, stored as scalar_t. +template +inline void orbitquant_rpbh_rotate_impl( + const device scalar_t* input, + const device long* permutation, + const device char* signs, + device scalar_t* output, + constant int& rows, + constant int& dim, + constant int& block_size, + constant int& num_blocks, + constant float& normalization, + threadgroup float* values, + uint group, + uint thread_index, + uint thread_count) { + const uint row = group / uint(num_blocks); + if (row >= uint(rows)) { + return; + } + const uint block = group - row * uint(num_blocks); + const uint block_start = block * uint(block_size); + const uint row_offset = row * uint(dim); + + for (uint col = thread_index; col < uint(block_size); col += thread_count) { + const uint rotated_col = block_start + col; + values[col] = float(input[row_offset + uint(permutation[rotated_col])]) * + float(signs[rotated_col]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 1; stride < uint(block_size); stride <<= 1) { + const uint pairs = uint(block_size) >> 1; + for (uint pair = thread_index; pair < pairs; pair += thread_count) { + const uint pair_group = pair / stride; + const uint pair_offset = pair - pair_group * stride; + const uint left = pair_group * (stride << 1) + pair_offset; + const uint right = left + stride; + const float left_value = values[left]; + const float right_value = values[right]; + values[left] = left_value + right_value; + values[right] = left_value - right_value; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint col = thread_index; col < uint(block_size); col += thread_count) { + output[row_offset + block_start + col] = scalar_t(values[col] * normalization); + } +} + +#define ORBITQUANT_RPBH_ROTATE_KERNEL(NAME, SCALAR_T) \ +kernel void NAME( \ + const device SCALAR_T* input [[buffer(0)]], \ + const device long* permutation [[buffer(1)]], \ + const device char* signs [[buffer(2)]], \ + device SCALAR_T* output [[buffer(3)]], \ + constant int& rows [[buffer(4)]], \ + constant int& dim [[buffer(5)]], \ + constant int& block_size [[buffer(6)]], \ + constant int& num_blocks [[buffer(7)]], \ + constant float& normalization [[buffer(8)]], \ + uint3 group_id [[threadgroup_position_in_grid]], \ + uint thread_index [[thread_index_in_threadgroup]], \ + uint3 group_size [[threads_per_threadgroup]]) { \ + threadgroup float values[orbitquant_max_rpbh_block_size]; \ + orbitquant_rpbh_rotate_impl( \ + input, permutation, signs, output, rows, dim, block_size, num_blocks, \ + normalization, values, group_id.x, thread_index, group_size.x); \ +} + +ORBITQUANT_RPBH_ROTATE_KERNEL(orbitquant_rpbh_rotate_float, float) +ORBITQUANT_RPBH_ROTATE_KERNEL(orbitquant_rpbh_rotate_half, half) +ORBITQUANT_RPBH_ROTATE_KERNEL(orbitquant_rpbh_rotate_bfloat16, bfloat) + +#undef ORBITQUANT_RPBH_ROTATE_KERNEL + #define ORBITQUANT_ROW_NORM_KERNEL(NAME, SCALAR_T, THREADS) \ kernel void NAME( \ const device SCALAR_T* input [[buffer(0)]], \ @@ -363,6 +441,62 @@ def quantize_activations_with_mps( return output.reshape_as(contiguous) +_ROTATE_KERNEL_NAMES = { + torch.float32: "orbitquant_rpbh_rotate_float", + torch.float16: "orbitquant_rpbh_rotate_half", + torch.bfloat16: "orbitquant_rpbh_rotate_bfloat16", +} + + +def rotate_activations_with_mps( + x: torch.Tensor, + rotation: RPBHRotation, + *, + permutation: torch.Tensor | None = None, + signs: torch.Tensor | None = None, +) -> torch.Tensor: + """``rotation.apply_to_activations`` in one Metal kernel, rounded to ``x.dtype`` once.""" + if x.device.type != "mps": + raise RuntimeError("mps backend requires MPS tensors") + if x.shape[-1] != rotation.dim: + raise ValueError(f"expected last dimension {rotation.dim}, got {x.shape[-1]}") + if rotation.block_size > _MAX_FUSED_RPBH_BLOCK_SIZE: + raise RuntimeError( + f"MPS RPBH rotation supports block sizes up to {_MAX_FUSED_RPBH_BLOCK_SIZE}" + ) + try: + name = _ROTATE_KERNEL_NAMES[x.dtype] + except KeyError as exc: + raise TypeError( + f"MPS RPBH rotation supports float32, float16 and bfloat16; got {x.dtype}" + ) from exc + contiguous = x.contiguous() + flat = contiguous.reshape(-1, rotation.dim) + rows = flat.shape[0] + output = torch.empty_like(flat) + if rows == 0: + return output.reshape_as(contiguous) + permutation = (rotation.permutation if permutation is None else permutation).to( + device=x.device, dtype=torch.int64 + ) + signs = (rotation.signs if signs is None else signs).to(device=x.device, dtype=torch.int8) + threads = 256 + getattr(_mps_shader(_fused_block_bucket(rotation.block_size)), name)( + flat, + permutation.contiguous(), + signs.contiguous(), + output, + rows, + rotation.dim, + rotation.block_size, + rotation.num_blocks, + rotation.normalization, + threads=[rows * rotation.num_blocks * threads, 1, 1], + group_size=[threads, 1, 1], + ) + return output.reshape_as(contiguous) + + def quantize_rotated_activations_with_mps( rotated: torch.Tensor, norms: torch.Tensor, diff --git a/src/orbitquant/kernels/triton_activation.py b/src/orbitquant/kernels/triton_activation.py index 5e8aadb..bc4f1c4 100644 --- a/src/orbitquant/kernels/triton_activation.py +++ b/src/orbitquant/kernels/triton_activation.py @@ -218,6 +218,90 @@ def _rotated_chunks( return v0, v1, v2, v3 +@triton.jit +def _rpbh_rotate_kernel( + x_ptr, + perm_ptr, + signs_ptr, + out_ptr, + D: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + OB: tl.constexpr, + STAGES: tl.constexpr, + inv_sqrt_block, +): + # One (row, RPBH block) per program, in FP32 like RPBHRotation.apply_to_activations; the + # store rounds to the output dtype once. + row_block = tl.program_id(0) + row = (row_block // NUM_BLOCKS).to(tl.int64) + block = row_block % NUM_BLOCKS + if OB <= _CHUNK: + cols = block * OB + tl.arange(0, OB) + src = tl.load(perm_ptr + cols).to(tl.int32) + signs = tl.load(signs_ptr + cols).to(tl.float32) + values = tl.load(x_ptr + row * D + src).to(tl.float32) * signs + values = _fwht(values, OB, STAGES) * inv_sqrt_block + tl.store(out_ptr + row * D + cols, values.to(out_ptr.dtype.element_ty)) + else: + start = block * OB + v0, v1, v2, v3 = _rotated_chunks( + x_ptr, + perm_ptr, + signs_ptr, + row, + start, + 1.0, + D, + OB // _CHUNK, + STAGES - _CHUNK_STAGES, + False, + ) + out = out_ptr + row * D + start + tl.arange(0, _CHUNK) + tl.store(out, (v0 * inv_sqrt_block).to(out_ptr.dtype.element_ty)) + tl.store(out + _CHUNK, (v1 * inv_sqrt_block).to(out_ptr.dtype.element_ty)) + if OB // _CHUNK == 4: + tl.store(out + 2 * _CHUNK, (v2 * inv_sqrt_block).to(out_ptr.dtype.element_ty)) + tl.store(out + 3 * _CHUNK, (v3 * inv_sqrt_block).to(out_ptr.dtype.element_ty)) + + +class RPBHRotator: + """Device constants of one RPBH rotation applied as a standalone activation transform. + + ``rotation_rounds=2`` puts a first RPBH round in front of the paper round; this applies + that round to an input in one kernel (the same FP32 arithmetic and stage order as + ``RPBHRotation.apply_to_activations``, rounded to the input dtype once). + """ + + def __init__(self, rotation, device: torch.device): + self.dim = int(rotation.dim) + self.block = int(rotation.block_size) + self.num_blocks = int(rotation.num_blocks) + if self.block > 16384: + raise ValueError("fused RPBH rotation supports blocks up to 16384") + self.inv_sqrt_block = float(rotation.normalization) + self.perm = rotation.permutation.to(device=device, dtype=torch.int32).contiguous() + self.signs = rotation.signs.to(device=device, dtype=torch.int8).contiguous() + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + shape = x.shape + x = x.reshape(-1, self.dim).contiguous() + out = torch.empty_like(x) + if x.shape[0]: + _rpbh_rotate_kernel[(x.shape[0] * self.num_blocks,)]( + x, + self.perm, + self.signs, + out, + D=self.dim, + NUM_BLOCKS=self.num_blocks, + OB=self.block, + STAGES=self.block.bit_length() - 1, + inv_sqrt_block=self.inv_sqrt_block, + num_warps=16 if self.block >= 8192 else 8 if self.block >= 512 else 4, + ) + return out.reshape(shape) + + @triton.jit def _absmax_int8_store(values, q_ptr, row, start, scale, D: tl.constexpr): q = libdevice.rint(libdevice.div_rn(values, scale)) @@ -389,10 +473,12 @@ class ActivationQuantizer: Projections that share an input (Q/K/V/gate, gate/up) share one quantizer: the codes and token norms are identical, so they are produced once. Without a codebook the rotated input is quantized to INT8 with one absmax scale per token (8-bit activations) and the call - returns that scale instead of the token norm. + returns that scale instead of the token norm. ``pre_rotation`` is the first round of a + two-round RPBH (``rotation_rounds=2``): it rotates the prologue output before the codes. """ - def __init__(self, rotation, codebook, eps: float, device: torch.device): + def __init__(self, rotation, codebook, eps: float, device: torch.device, pre_rotation=None): + self.pre_rotator = None if pre_rotation is None else RPBHRotator(pre_rotation, device) self.dim = int(rotation.dim) self.block = int(rotation.block_size) self.num_blocks = int(rotation.num_blocks) @@ -437,6 +523,39 @@ def __call__( ``mod`` override its norm and modulation. With ``mod_index`` the modulation vectors are rows ``mod_index[row]`` of the ``mod_scale``/``mod_shift`` tables. """ + fields = dict( + prologue=prologue, + norm=norm, + mod=mod, + x2=x2, + norm_weight=norm_weight, + mod_scale=mod_scale, + mod_shift=mod_shift, + mod_index=mod_index, + rms_eps=rms_eps, + ) + if self.pre_rotator is None or prologue_only: + return self._quantize(x, prologue_only=prologue_only, **fields) + # The eager model rounds the prologue output to bf16, then the layer rotates it. + plain = prologue == PROLOGUE_NONE and norm in (None, NORM_NONE) and mod in (None, MOD_NONE) + y = x if plain else self._quantize(x, prologue_only=True, **fields)[0] + return self._quantize(self.pre_rotator(y)) + + def _quantize( + self, + x, + *, + prologue=PROLOGUE_NONE, + norm=None, + mod=None, + x2=None, + norm_weight=None, + mod_scale=None, + mod_shift=None, + mod_index=None, + rms_eps=1e-5, + prologue_only=False, + ): if x.dim() != 2 or x.stride(1) != 1: x = x.reshape(-1, self.dim).contiguous() rows = x.shape[0] diff --git a/src/orbitquant/layers.py b/src/orbitquant/layers.py index d7f5243..a62b018 100644 --- a/src/orbitquant/layers.py +++ b/src/orbitquant/layers.py @@ -581,6 +581,7 @@ def __init__( self._bias_cache: tuple[tuple[str, torch.dtype], torch.Tensor] | None = None self._int8_weight_cache: tuple[str, torch.Tensor] | None = None self._absmax_quantizer: tuple[str, object] | None = None + self._pre_rotator: tuple[str, object] | None = None self._compile_handle = _compile_registry_register(self) self.last_effective_runtime_mode: str | None = None self.last_activation_kernel_backend: str | None = None @@ -1405,11 +1406,53 @@ def _absmax_int8_cuda_forward(self, x: torch.Tensor) -> torch.Tensor: ) return output.reshape(*x.shape[:-1], self.out_features) + def _triton_pre_rotator(self, x: torch.Tensor): + """The Triton kernel of the first RPBH round on CUDA, or None for the PyTorch path.""" + if not x.is_cuda or _torch_uses_hip() or self.pre_rotation.block_size > 16384: + return None + + def probe() -> bool: + try: + from orbitquant.kernels import triton_activation # noqa: F401 + except (ImportError, RuntimeError): + return False + return True + + if not self._cached_runtime_probe("triton_pre_rotation", probe): + return None + device_key = str(x.device) + if self._pre_rotator is None or self._pre_rotator[0] != device_key: + from orbitquant.kernels import triton_activation + + self._pre_rotator = ( + device_key, + triton_activation.RPBHRotator(self.pre_rotation, x.device), + ) + return self._pre_rotator[1] + def _pre_rotate(self, x: torch.Tensor) -> torch.Tensor: """First RPBH round of a two-round rotation, shared across projections of one input.""" rotation = self.pre_rotation def factory() -> torch.Tensor: + rotator = self._triton_pre_rotator(x) + if rotator is not None: + # Same FP32 arithmetic and stage order in one kernel. + return rotator(x) + if x.device.type == "mps" and rotation.block_size <= 4096: + from orbitquant.kernels.mps import mps_metal_available, rotate_activations_with_mps + + if mps_metal_available() and x.dtype in ( + torch.float32, + torch.float16, + torch.bfloat16, + ): + return rotate_activations_with_mps( + x, + rotation, + permutation=self._pre_rotation_permutation, + signs=self._pre_rotation_signs, + ) rotated = rotation.apply_to_activations( x.to(torch.float32), permutation=self._pre_rotation_permutation, diff --git a/src/orbitquant/recipes.py b/src/orbitquant/recipes.py index d784e5f..36b3e0b 100644 --- a/src/orbitquant/recipes.py +++ b/src/orbitquant/recipes.py @@ -20,8 +20,20 @@ } -def recipe(name: str, **overrides: Any) -> OrbitQuantConfig: - """Create a universal OrbitQuant config from a named bit profile.""" +# The paper's method (arXiv 2607.02461) where the defaults go beyond it: one RPBH round, +# Lloyd-Max activation codes at every width, no output gain correction. +PAPER_FIELDS: dict[str, Any] = { + "rotation_rounds": 1, + "activation_quantizer": "lloyd_max", + "output_gain_correction": False, +} + + +def recipe(name: str, *, paper: bool = False, **overrides: Any) -> OrbitQuantConfig: + """Create a universal OrbitQuant config from a named bit profile. + + The defaults include the measured improvements over the paper; ``paper=True`` gives the + paper's method as the baseline to compare against.""" normalized_name = name.lower().replace("-", "") try: @@ -32,10 +44,10 @@ def recipe(name: str, **overrides: Any) -> OrbitQuantConfig: "weight_bits": weight_bits, "activation_bits": activation_bits, "target_policy": "universal", - **_RECIPE_FIELDS.get(normalized_name, {}), + **(PAPER_FIELDS if paper else _RECIPE_FIELDS.get(normalized_name, {})), } values.update(overrides) return OrbitQuantConfig(**values) -__all__ = ["recipe"] +__all__ = ["PAPER_FIELDS", "recipe"] diff --git a/src/orbitquant/runtime/krea2.py b/src/orbitquant/runtime/krea2.py index 7f2b0f8..3c30db9 100644 --- a/src/orbitquant/runtime/krea2.py +++ b/src/orbitquant/runtime/krea2.py @@ -76,9 +76,12 @@ def _codes_and_alpha(layer: OrbitQuantLinear) -> tuple[torch.Tensor, float]: - if layer.activation_codebook is None: - raise ValueError("the Krea 2 runtime requires Lloyd-Max activation codebooks") - _, act_scale = fit_int8_centroid_surrogate(layer.activation_codebook.centroids) + if layer.outlier_channels is not None: + raise ValueError("outlier_channels run on the per-layer OrbitQuantLinear path") + # Absmax INT8 activations carry their absolute scale per token: a_scale = 1. + act_scale = 1.0 + if layer.activation_codebook is not None: + _, act_scale = fit_int8_centroid_surrogate(layer.activation_codebook.centroids) codes, w_scale = fit_int8_centroid_surrogate(layer.weight_codebook.centroids) return codes, float(act_scale * w_scale) @@ -238,7 +241,10 @@ def __init__( first = blocks[0].attn.to_q # Plain attribute, not a submodule: only the RPBH/codebook constants are needed. self.__dict__["act_spec"] = types.SimpleNamespace( - rotation=first.rotation, codebook=first.activation_codebook, eps=first.activation_eps + rotation=first.rotation, + pre_rotation=first.pre_rotation, + codebook=first.activation_codebook, + eps=first.activation_eps, ) self.down_mode = down self.attention = _ATTENTION_BACKENDS[attention] @@ -275,7 +281,7 @@ def _runtime(self, device): if self._quantizer is None or self._buffers_[0].device != device: spec = self.act_spec self._quantizer = act.ActivationQuantizer( - spec.rotation, spec.codebook, spec.eps, device + spec.rotation, spec.codebook, spec.eps, device, pre_rotation=spec.pre_rotation ) count = 2 if self.prefetch else 1 self._buffers_ = [ diff --git a/tests/test_absmax_activations.py b/tests/test_absmax_activations.py index 9cfa020..fca7e1c 100644 --- a/tests/test_absmax_activations.py +++ b/tests/test_absmax_activations.py @@ -95,10 +95,13 @@ def relative_error(**fields): @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS required") def test_absmax_int8_runs_on_mps_with_reference_numerics(): + from orbitquant.kernels.native_packed_matmul import load_native_packed_matmul_kernel + try: - import orbitquant_packed_matmul + # Attaches the provisioned prebuilt variant, as the runtime does. + orbitquant_packed_matmul = load_native_packed_matmul_kernel() except Exception: - pytest.skip("native packed matmul kernel package is not importable") + pytest.skip("native packed matmul kernel package is not available") if not orbitquant_packed_matmul.supports_device("mps"): pytest.skip("the importable native packed matmul variant has no MPS backend") torch.manual_seed(0) diff --git a/tests/test_artifact_writer.py b/tests/test_artifact_writer.py index afb12da..fecd2cf 100644 --- a/tests/test_artifact_writer.py +++ b/tests/test_artifact_writer.py @@ -114,7 +114,7 @@ def test_save_orbitquant_artifact_writes_manifest_readme_weights_and_checksums(t assert model_index["weight_quantization_backend"] == summary.weight_quantization_backend assert model_index["quantization_staging_mode"] == summary.quantization_staging_mode assert model_index["activation_eps"] == 1e-10 - assert model_index["codebook_version"] == 2 + assert model_index["codebook_version"] == 3 assert "README.md" in {path.name for path in tmp_path.iterdir()} assert "model_index.json" in {path.name for path in tmp_path.iterdir()} assert "SHA256SUMS" in {path.name for path in tmp_path.iterdir()} @@ -168,7 +168,7 @@ def test_save_orbitquant_artifact_writes_manifest_readme_weights_and_checksums(t assert any("量子轨道" in item["prompt"] for item in prompts["prompts"]) assert benchmark_summary["status"] == "not_run" assert benchmark_summary["source_model_id"] == "example/model" - assert benchmark_summary["codebook_version"] == 2 + assert benchmark_summary["codebook_version"] == 3 assert benchmark_summary["quantization_device"] == summary.quantization_device assert benchmark_summary["weight_quantization_backend"] == summary.weight_quantization_backend assert benchmark_summary["quantization_staging_mode"] == summary.quantization_staging_mode @@ -283,6 +283,11 @@ def test_saved_artifact_contains_no_activation_calibration_state_and_deduplicate "dim8_seed0_block4.normalization", "dim8_seed0_block4.permutation", "dim8_seed0_block4.signs", + # rotation_rounds="auto": a 4-wide block gets the first RPBH round. + "pre_dim8_seed2654435769_block4.inverse_permutation", + "pre_dim8_seed2654435769_block4.normalization", + "pre_dim8_seed2654435769_block4.permutation", + "pre_dim8_seed2654435769_block4.signs", } @@ -689,7 +694,7 @@ def test_validate_orbitquant_artifact_rejects_wrong_versioned_codebook(tmp_path) try: validate_orbitquant_artifact(tmp_path) except RuntimeError as exc: - assert "do not match codebook version 2" in str(exc) + assert "do not match codebook version 3" in str(exc) else: raise AssertionError("validator accepted centroids from a different algorithm") diff --git a/tests/test_codebooks.py b/tests/test_codebooks.py index ea85ffa..a8ff219 100644 --- a/tests/test_codebooks.py +++ b/tests/test_codebooks.py @@ -1,3 +1,6 @@ +import time + +import pytest import torch from safetensors.torch import load_file, save_file @@ -128,6 +131,47 @@ def test_v2_codebook_satisfies_lloyd_max_centroid_condition(monkeypatch): torch.testing.assert_close(positive, conditional_means, atol=1e-7, rtol=0) +def _float32_ulps(left, right): + return int((left.view(torch.int32) - right.view(torch.int32)).abs().max()) + + +@pytest.mark.parametrize(("dim", "bits"), [(64, 4), (128, 2), (3072, 4), (128, 6)]) +def test_v3_newton_codebook_reaches_the_v2_fixed_point(monkeypatch, dim, bits): + monkeypatch.setenv("ORBITQUANT_DISABLE_CODEBOOK_DISK_CACHE", "1") + clear_codebook_cache() + v2 = get_codebook(dim=dim, bits=bits, algorithm_version=2) + v3 = get_codebook(dim=dim, bits=bits, algorithm_version=3) + + # v2 stops its linear iteration short of the fixed point by up to one float32 ulp. + assert v3.algorithm_version == 3 + assert _float32_ulps(v3.centroids, v2.centroids) <= 1 + assert _float32_ulps(v3.boundaries, v2.boundaries) <= 1 + + +def test_v3_builds_eight_bit_codebooks_where_v2_does_not_converge(monkeypatch): + # v2 needs 393216 iterations and still fails for d = 7680 at 8 bits. + monkeypatch.setenv("ORBITQUANT_DISABLE_CODEBOOK_DISK_CACHE", "1") + clear_codebook_cache() + dim = 7680 + start = time.perf_counter() + codebook = get_codebook(dim=dim, bits=8, algorithm_version=3) + elapsed = time.perf_counter() - start + + positive = codebook.centroids[codebook.centroids.numel() // 2 :].double() + edges = [0.0, *((positive[:-1] + positive[1:]) / 2).tolist(), 1.0] + conditional_means = torch.tensor( + [ + lloyd_max_module._positive_interval_centroid(dim, edges[index], edges[index + 1]) + for index in range(positive.numel()) + ], + dtype=torch.float64, + ) + assert codebook.centroids.numel() == 256 + # Two float32 epsilons: the stored tables are float32 roundings of the fixed point. + torch.testing.assert_close(positive, conditional_means, atol=0, rtol=2 * 2**-23) + assert elapsed < 10.0 + + def test_lloyd_max_quantization_error_decreases_with_bits(): torch.manual_seed(0) values = torch.randn(2048) diff --git a/tests/test_config.py b/tests/test_config.py index fb69ef3..cda98dd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -64,7 +64,7 @@ def test_orbit_quant_config_default_epsilon_matches_paper_equation(): def test_orbit_quant_config_defaults_to_converged_codebook_version(): - assert OrbitQuantConfig().codebook_version == 2 + assert OrbitQuantConfig().codebook_version == 3 def test_orbit_quant_config_reads_unversioned_artifacts_as_legacy_codebooks(): @@ -219,3 +219,31 @@ def test_orbit_quant_config_rejects_unimplemented_serialized_knobs(): assert next(iter(kwargs)) in str(exc) else: raise AssertionError(f"unsupported OrbitQuantConfig kwargs accepted: {kwargs}") + + +def test_defaults_take_the_measured_improvements_over_the_paper(): + config = OrbitQuantConfig() + assert config.rotation_rounds == "auto" + assert config.codebook_version == 3 + assert config.activation_quantizer == "lloyd_max" + assert OrbitQuantConfig(activation_bits=8).activation_quantizer == "absmax_int8" + assert OrbitQuantConfig(activation_bits=6).activation_quantizer == "lloyd_max" + + +def test_artifacts_without_the_new_fields_load_with_the_paper_values(): + restored = OrbitQuantConfig.from_dict({"weight_bits": 4, "activation_bits": 8}) + assert restored.rotation_rounds == 1 + assert restored.activation_quantizer == "lloyd_max" + round_trip = OrbitQuantConfig.from_dict(OrbitQuantConfig(activation_bits=8).to_dict()) + assert round_trip.rotation_rounds == "auto" + assert round_trip.activation_quantizer == "absmax_int8" + + +def test_paper_recipe_is_the_paper_method(): + from orbitquant import recipe + + paper = recipe("w4a8", paper=True) + assert (paper.rotation_rounds, paper.activation_quantizer) == (1, "lloyd_max") + assert paper.output_gain_correction is False + assert recipe("w4a8").activation_quantizer == "absmax_int8" + assert recipe("w4a4").rotation_rounds == "auto" diff --git a/tests/test_fused_family_kernels.py b/tests/test_fused_family_kernels.py index 2ff098f..c830a3e 100644 --- a/tests/test_fused_family_kernels.py +++ b/tests/test_fused_family_kernels.py @@ -302,6 +302,94 @@ def test_lloyd_max_activation_codes_match_the_rotated_reference(dim): assert (codes != expected).float().mean().item() < 1e-3 +@pytest.mark.parametrize( + ("dim", "block_size"), + [(3360, "paper"), (3840, "paper"), (3072, 1024), (8192, 8192), (16384, 16384)], +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_rpbh_rotator_matches_the_pytorch_rotation_bit_for_bit(dim, block_size, dtype): + _require_cuda() + from orbitquant.kernels import triton_activation as act + from orbitquant.rotations import get_rpbh_rotation + + torch.manual_seed(9) + rotation = get_rpbh_rotation(dim=dim, seed=3, block_size=block_size) + x = torch.randn(2, 19, dim, device="cuda", dtype=dtype) + x[..., 7] *= 300 # a massive channel + expected = rotation.apply_to_activations(x.float()).to(dtype) + rotated = act.RPBHRotator(rotation, torch.device("cuda"))(x) + assert rotated.shape == x.shape and rotated.dtype == dtype + assert torch.equal(rotated, expected) + + +@pytest.mark.parametrize("codebook_bits", [4, None]) +def test_activation_quantizer_applies_the_pre_rotation_after_the_prologue(codebook_bits): + _require_cuda() + from orbitquant.codebooks import get_codebook + from orbitquant.kernels import triton_activation as act + from orbitquant.rotations import get_pre_rotation, get_rpbh_rotation + + torch.manual_seed(10) + dim = 3840 + rotation = get_rpbh_rotation(dim=dim, seed=0, block_size="paper") + pre = get_pre_rotation(dim, seed=0, block_size="paper", rotation_rounds="auto") + codebook = None if codebook_bits is None else get_codebook(dim, codebook_bits, 3) + device = torch.device("cuda") + plain = act.ActivationQuantizer(rotation, codebook, 1e-10, device) + two_round = act.ActivationQuantizer(rotation, codebook, 1e-10, device, pre_rotation=pre) + x = torch.randn(23, dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16) + scale = torch.randn(dim, device="cuda", dtype=torch.bfloat16) + shift = torch.randn(dim, device="cuda", dtype=torch.bfloat16) + kwargs = dict( + prologue=act.PROLOGUE_NORM_MODULATE, norm_weight=weight, mod_scale=scale, mod_shift=shift + ) + + prologue, _ = plain(x, prologue_only=True, **kwargs) + expected_codes, expected_norms = plain(act.RPBHRotator(pre, device)(prologue)) + codes, norms = two_round(x, **kwargs) + assert torch.equal(codes, expected_codes) + assert torch.equal(norms, expected_norms) + # BF16 groups reading the same input still get the unrotated prologue output. + assert torch.equal(two_round(x, prologue_only=True, **kwargs)[0], prologue) + + +@pytest.mark.parametrize("activation_bits", [4, 8]) +def test_fused_two_round_projection_matches_the_per_layer_path(activation_bits): + _require_cuda() + from orbitquant.config import OrbitQuantConfig + from orbitquant.fused.groups import PackedGroup, RowSource + from orbitquant.fused.runtime import FusedRuntime + from orbitquant.layers import OrbitQuantLinear + + torch.manual_seed(11) + config = OrbitQuantConfig( + rotation_rounds="auto", + activation_bits=activation_bits, + activation_quantizer="absmax_int8" if activation_bits == 8 else "lloyd_max", + ) + root = torch.nn.Module() + root.proj = OrbitQuantLinear.from_linear( + torch.nn.Linear(3840, 512, bias=False).to(torch.bfloat16), config=config, module_name="p" + ).cuda() + assert root.proj.pre_rotation is not None + runtime = FusedRuntime.from_config(config) + group = PackedGroup.from_sources(root, [RowSource("proj")]).cuda() + runtime.check_pre_rotation(root.proj) + x = torch.randn(64, 3840, device="cuda", dtype=torch.bfloat16) + x[:, 11] *= 200 + with torch.no_grad(): + fused = runtime.linear(group, runtime.input(x, 3840)) + reference = root.proj(x) + error = (fused.float() - reference.float()).norm() / reference.float().norm() + assert error.item() < 2e-2 + + with pytest.raises(ValueError, match="pre-rotation does not match"): + FusedRuntime.from_config(OrbitQuantConfig(rotation_rounds=1)).check_pre_rotation( + root.proj + ) + + def test_krea2_batch_one_forward_runs_on_the_valid_text_rows(): _require_cuda() from orbitquant.fused.families import krea2 diff --git a/tests/test_kernels.py b/tests/test_kernels.py index 637d1d4..f507213 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -11,7 +11,11 @@ quantize_activations_kernel, select_backend, ) -from orbitquant.kernels.executorch_vulkan import ExecuTorchVulkanW4A4Linear +from orbitquant.kernels.executorch_vulkan import ( + ExecuTorchVulkanW4A4Linear, + ExecuTorchVulkanW4A8Linear, + prepare_executorch_vulkan_w4a4_model, +) from orbitquant.layers import OrbitQuantLinear from orbitquant.modeling import quantize_linear_modules from orbitquant.packing import pack_lowbit, unpack_lowbit @@ -44,6 +48,75 @@ def test_executorch_vulkan_export_keeps_exact_w4a4_as_one_custom_op(): assert torch.ops.orbitquant_vulkan.linear_w4a4_exact.default in call_targets +def test_executorch_vulkan_export_runs_the_pre_rotation_as_plain_ops(): + torch.manual_seed(0) + source = torch.nn.Linear(96, 40) + layer = OrbitQuantLinear.from_linear( + source, + config=OrbitQuantConfig(rotation_rounds="auto", runtime_mode="dequant_bf16"), + module_name="probe", + ) + assert layer.pre_rotation is not None + export_layer = ExecuTorchVulkanW4A4Linear(layer) + x = torch.randn(5, 96) + x[:, 3] *= 50 + + with torch.no_grad(): + torch.testing.assert_close(export_layer(x), layer(x), rtol=1e-5, atol=1e-5) + exported = torch.export.export(export_layer, (x,), strict=True) + call_targets = {node.target for node in exported.graph.nodes if node.op == "call_function"} + assert torch.ops.orbitquant_vulkan.linear_w4a4_exact.default in call_targets + assert torch.ops.aten.matmul.default in call_targets + + +def test_executorch_vulkan_export_refuses_outlier_channels(): + layer = OrbitQuantLinear.from_linear( + torch.nn.Linear(96, 40), + config=OrbitQuantConfig( + runtime_mode="dequant_bf16", outlier_channels={"probe": [3]} + ), + module_name="probe", + ) + with pytest.raises(ValueError, match="outlier_channels"): + ExecuTorchVulkanW4A4Linear(layer) + + +@pytest.mark.parametrize(("in_features", "rotation_rounds"), [(96, "auto"), (1536, 1)]) +def test_executorch_vulkan_w4a8_export_matches_the_absmax_layer(in_features, rotation_rounds): + torch.manual_seed(0) + layer = OrbitQuantLinear.from_linear( + torch.nn.Linear(in_features, 40), + config=OrbitQuantConfig( + activation_bits=8, rotation_rounds=rotation_rounds, runtime_mode="dequant_bf16" + ), + module_name="probe", + ) + assert layer.activation_quantizer == "absmax_int8" + prepared = prepare_executorch_vulkan_w4a4_model(torch.nn.Sequential(layer)) + export_layer = prepared[0] + assert isinstance(export_layer, ExecuTorchVulkanW4A8Linear) + x = torch.randn(5, in_features) + x[:, 3] *= 50 + + with torch.no_grad(): + torch.testing.assert_close(export_layer(x), layer(x), rtol=1e-5, atol=1e-5) + exported = torch.export.export(export_layer, (x,), strict=True) + call_targets = {node.target for node in exported.graph.nodes if node.op == "call_function"} + assert torch.ops.orbitquant_vulkan.linear_w4a8_absmax.default in call_targets + + +def test_executorch_vulkan_export_refuses_lloyd_max_a8(): + layer = OrbitQuantLinear.from_linear( + torch.nn.Linear(96, 40), + config=OrbitQuantConfig( + activation_bits=8, activation_quantizer="lloyd_max", runtime_mode="dequant_bf16" + ), + module_name="probe", + ) + with pytest.raises(ValueError, match="absmax INT8"): + prepare_executorch_vulkan_w4a4_model(layer) + + def test_executorch_vulkan_export_rejects_unsupported_bfloat16_activations(): layer = _make_executorch_vulkan_export_layer() diff --git a/tests/test_manifest.py b/tests/test_manifest.py index 99a6c1b..13ac864 100644 --- a/tests/test_manifest.py +++ b/tests/test_manifest.py @@ -32,7 +32,7 @@ def test_manifest_records_source_and_quantization_settings(): assert data["block_size"] == 128 assert data["block_size_policy"] == "explicit" assert data["codebook"] == "lloyd_max" - assert data["codebook_version"] == 2 + assert data["codebook_version"] == 3 assert data["activation_eps"] == 1e-10 assert data["quantized_modules"] == ["transformer_blocks.0.attn.to_q"] @@ -40,7 +40,7 @@ def test_manifest_records_source_and_quantization_settings(): assert restored.rotation_seed == 17 assert restored.block_size == 128 assert restored.block_size_policy == "explicit" - assert restored.codebook_version == 2 + assert restored.codebook_version == 3 assert restored.activation_eps == 1e-10 diff --git a/tests/test_model_card.py b/tests/test_model_card.py index 1c001e5..57a4f88 100644 --- a/tests/test_model_card.py +++ b/tests/test_model_card.py @@ -71,7 +71,7 @@ def test_model_card_renders_rotation_and_codebook_metadata(): assert "- Rotation seed: `9`" in card assert "- Block size: `paper`" in card assert "- Block size policy: `largest_power_of_two_dividing_dim`" in card - assert "- Codebook version: `2`" in card + assert "- Codebook version: `3`" in card assert "- Target policy: `z_image`" in card assert "- Quantized transformer modules: `1`" in card assert "## Native Settings" in card diff --git a/tests/test_orbit_linear.py b/tests/test_orbit_linear.py index a0bcb39..b05f6d4 100644 --- a/tests/test_orbit_linear.py +++ b/tests/test_orbit_linear.py @@ -78,7 +78,9 @@ def test_orbit_linear_shares_codebooks_by_dimension_and_bits_not_module_name(): def test_orbit_linear_weight_indices_quantize_rotated_unit_directions(): torch.manual_seed(2) source = torch.nn.Linear(16, 7) - config = OrbitQuantConfig(weight_bits=4, activation_bits=4, rotation_seed=11, block_size=8) + config = OrbitQuantConfig( + weight_bits=4, activation_bits=4, rotation_seed=11, rotation_rounds=1, block_size=8 + ) quantized = OrbitQuantLinear.from_linear(source, config=config, module_name="block.ff.linear") @@ -145,6 +147,7 @@ def test_orbit_linear_quantized_forward_matches_manual_paper_equation(monkeypatc weight_bits=3, activation_bits=3, rotation_seed=13, + rotation_rounds=1, block_size="paper", activation_kernel_backend="cpu", ) diff --git a/tests/test_paper_reference_oracle.py b/tests/test_paper_reference_oracle.py index 2568a73..e4ec93b 100644 --- a/tests/test_paper_reference_oracle.py +++ b/tests/test_paper_reference_oracle.py @@ -154,6 +154,7 @@ def test_reference_runtime_matches_independent_algorithm_1_oracle( weight_bits=weight_bits, activation_bits=activation_bits, rotation_seed=29, + rotation_rounds=1, block_size="paper", runtime_mode=runtime_mode, activation_kernel_backend="cpu", @@ -212,6 +213,7 @@ def test_executorch_vulkan_w4a4_export_op_matches_independent_algorithm_1_oracle weight_bits=4, activation_bits=4, rotation_seed=31, + rotation_rounds=1, block_size="paper", runtime_mode="dequant_bf16", ) @@ -273,6 +275,7 @@ def test_triton_rocm_w4a4_matches_independent_algorithm_1_oracle(): weight_bits=4, activation_bits=4, rotation_seed=37, + rotation_rounds=1, block_size="paper", runtime_mode="triton_packed_matmul", activation_kernel_backend="triton_rocm", @@ -343,6 +346,7 @@ def test_triton_xpu_w4a4_matches_independent_algorithm_1_oracle(): weight_bits=4, activation_bits=4, rotation_seed=41, + rotation_rounds=1, block_size="paper", runtime_mode="triton_packed_matmul", activation_kernel_backend="triton_xpu", diff --git a/tests/test_rotation_rounds.py b/tests/test_rotation_rounds.py index 4546b8a..8a2803e 100644 --- a/tests/test_rotation_rounds.py +++ b/tests/test_rotation_rounds.py @@ -170,23 +170,6 @@ def test_two_round_artifact_validates_and_restores_the_same_outputs(tmp_path): torch.testing.assert_close(layer(x), source.transformer_blocks[0]["attn"]["to_q"](x)) -def test_fused_runtime_refuses_two_round_rotations(): - pytest.importorskip("triton") - from orbitquant.fused.groups import PackedGroup, RowSource - from orbitquant.fused.runtime import FusedRuntime - - with pytest.raises(ValueError, match="one RPBH round"): - FusedRuntime.from_config(OrbitQuantConfig(rotation_rounds=2)) - holder = torch.nn.Module() - holder.proj = OrbitQuantLinear.from_linear( - torch.nn.Linear(96, 32), - config=OrbitQuantConfig(rotation_rounds=2, runtime_mode="dequant_bf16"), - module_name="proj", - ) - with pytest.raises(ValueError, match="one RPBH round"): - PackedGroup.from_sources(holder, [RowSource("proj")]) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_two_round_layer_on_cuda_matches_the_cpu_reference(): torch.manual_seed(0) @@ -207,3 +190,40 @@ def test_two_round_layer_on_cuda_matches_the_cpu_reference(): error = (output - expected).pow(2).sum() / expected.pow(2).sum() assert float(error) < 5e-3 + + +def _mps_shader_available(): + from orbitquant.kernels.mps import mps_metal_available + + return mps_metal_available() + + +@pytest.mark.skipif(not _mps_shader_available(), reason="MPS Metal shaders required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize(("dim", "block_size"), [(3840, "paper"), (3360, "paper"), (4096, 4096)]) +def test_metal_rpbh_rotation_matches_the_pytorch_rotation_bit_for_bit(dim, block_size, dtype): + from orbitquant.kernels.mps import rotate_activations_with_mps + from orbitquant.rotations import get_rpbh_rotation + + torch.manual_seed(0) + rotation = get_rpbh_rotation(dim, 3, block_size) + x = torch.randn(7, 3, dim, dtype=dtype) + x[..., 5] *= 300 + expected = rotation.apply_to_activations(x.float()).to(dtype) + assert torch.equal(rotate_activations_with_mps(x.to("mps"), rotation).cpu(), expected) + + +@pytest.mark.skipif(not _mps_shader_available(), reason="MPS Metal shaders required") +def test_two_round_layer_on_mps_matches_the_cpu_reference(): + torch.manual_seed(0) + linear = torch.nn.Linear(3360, 64) + x = _outlier_tokens(3360) + config = OrbitQuantConfig(rotation_rounds="auto", runtime_mode="dequant_bf16") + reference = OrbitQuantLinear.from_linear(linear, config=config, module_name="proj") + layer = OrbitQuantLinear.from_linear(linear, config=config, module_name="proj").to("mps") + + with torch.no_grad(): + expected = reference(x) + output = layer(x.to("mps")).cpu() + + torch.testing.assert_close(output, expected, rtol=1e-3, atol=1e-3)