diff --git a/README.md b/README.md index 9e0a6e3..a8b814e 100644 --- a/README.md +++ b/README.md @@ -111,7 +111,7 @@ Artifacts saved before a default changed load with the value they were quantized | Field | Values | Effect | Runtime cost | | --- | --- | --- | --- | -| `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. | +| `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 round per activation, shared across the projections of one input. On CUDA the absmax INT8 activation kernel applies it in the same pass for inputs up to 16384 wide: on an RTX 5090 at 4096-9216 tokens the two-round quantizer takes 0.22-0.49 ms at `d = 9536` and 0.05-0.10 ms at `d = 3840`, at most 4% longer than the one-round quantizer of 0.16.0 (0.16.0 ran the round as a separate kernel: 0.57-1.31 ms and 0.12-0.28 ms). At `d = 3360` (`h = 32`) and 9216 tokens it takes 16-19% longer. A norm or modulation prologue keeps its own kernel: with one, the two-round quantizer takes 15-22% longer than the 0.16.0 one-round quantizer at `d = 9536` and `d = 13568`. Lloyd-Max activations get the round from a separate Triton kernel: +0.03-0.06 ms per call at `d = 3840` and +0.07-0.19 ms at `d = 9536` on the same GPU and token counts. 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` (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. | @@ -145,7 +145,7 @@ 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 +-0.022 to +0.052), while the default took 7% longer with 0.16.0 (1656 s against 1547 s) and peaked 1.2 GiB higher. Use `rotation_rounds=1` for such models when speed matters. ## Install diff --git a/docs/paper-methodology-audit.md b/docs/paper-methodology-audit.md index 7d80c1c..7c5f79d 100644 --- a/docs/paper-methodology-audit.md +++ b/docs/paper-methodology-audit.md @@ -182,7 +182,7 @@ paper fields. | Field | Paper value | Deviation | | --- | --- | --- | -| `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. | +| `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 ahead of the paper round (on CUDA, inside the 8-bit absmax activation kernel). | | `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`. 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. | diff --git a/pyproject.toml b/pyproject.toml index 993fabf..85e80b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.16.0" +version = "0.16.1" description = "Calibration-free OrbitQuant for transformer linear projections" readme = "README.md" requires-python = ">=3.11" diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index 374f4ae..e75458b 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -24,7 +24,7 @@ from orbitquant.quantizer import register_hf_quantizers from orbitquant.recipes import recipe -__version__ = "0.16.0" +__version__ = "0.16.1" register_hf_quantizers() diff --git a/src/orbitquant/kernels/triton_activation.py b/src/orbitquant/kernels/triton_activation.py index bc4f1c4..b31d66a 100644 --- a/src/orbitquant/kernels/triton_activation.py +++ b/src/orbitquant/kernels/triton_activation.py @@ -218,6 +218,59 @@ def _rotated_chunks( return v0, v1, v2, v3 +# tl.gather (Triton >= 3.2) permutes a register tensor through shared memory. The row kernels +# use it for rows up to GATHER_MAX_ROW wide (a 16384-wide bf16 row takes 32 KiB). +HAS_GATHER = hasattr(tl, "gather") +GATHER_MAX_ROW = 16384 + + +def _row_warps(row_width: int) -> int: + # Measured on an RTX 5090 for one row per program. + return 16 if row_width >= 16384 else 8 if row_width >= 8192 else 4 + + +def _absmax_warps(row_width: int, block: int, pre_rotate: bool) -> int: + # Measured on an RTX 5090 for the absmax row kernel: rows up to 2048 wide run 14-18% faster + # on 2 warps than on 4, 256-wide RPBH blocks 6-31% faster on 4 warps than on 8 or 16, and + # their pre-round on 4096-wide rows 10% faster on 2. + if row_width <= 2048: + return 2 + if block == 256: + return 2 if pre_rotate and row_width <= 4096 else 4 + return _row_warps(row_width) + + +@triton.jit +def _signed_gather(row, perm_ptr, signs_ptr, cols, mask): + # ``row[perm] * signs`` in FP32 for a row held in registers, zero past D. + src = tl.load(perm_ptr + cols, mask=mask, other=0).to(tl.int32) + signs = tl.load(signs_ptr + cols, mask=mask, other=0).to(tl.float32) + return tl.where(mask, tl.gather(row, src, 0).to(tl.float32) * signs, 0.0) + + +@triton.jit +def _rpbh_rotate_rows_kernel( + x_ptr, + perm_ptr, + signs_ptr, + out_ptr, + D: tl.constexpr, + DP2: tl.constexpr, + STAGES: tl.constexpr, + inv_sqrt_block, +): + # One row per program, permuted in shared memory: narrow blocks make the per-block programs + # below tiny (0.37 ms vs 0.10 ms for 4096 rows of 149 64-wide blocks on an RTX 5090). Same + # FP32 arithmetic and stage order. + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, DP2) + mask = cols < D + values = tl.load(x_ptr + row * D + cols, mask=mask, other=0.0) + values = _signed_gather(values, perm_ptr, signs_ptr, cols, mask) + values = _fwht(values, DP2, STAGES) * inv_sqrt_block + tl.store(out_ptr + row * D + cols, values.to(out_ptr.dtype.element_ty), mask=mask) + + @triton.jit def _rpbh_rotate_kernel( x_ptr, @@ -281,12 +334,28 @@ def __init__(self, rotation, device: torch.device): 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() + # The blocks below 512 that rotation_rounds="auto" adds a round for. + self.row_width = triton.next_power_of_2(self.dim) + self.rows_kernel = HAS_GATHER and self.block < 512 and self.row_width <= GATHER_MAX_ROW 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]: + # The shared-memory copy of a row stays within that of a 16384-wide bf16 row. + if x.shape[0] and self.rows_kernel and self.row_width * x.element_size() <= 32768: + _rpbh_rotate_rows_kernel[(x.shape[0],)]( + x, + self.perm, + self.signs, + out, + D=self.dim, + DP2=self.row_width, + STAGES=self.block.bit_length() - 1, + inv_sqrt_block=self.inv_sqrt_block, + num_warps=_row_warps(self.row_width), + ) + elif x.shape[0]: _rpbh_rotate_kernel[(x.shape[0] * self.num_blocks,)]( x, self.perm, @@ -314,24 +383,42 @@ def _rpbh_absmax_int8_kernel( y_ptr, perm_ptr, signs_ptr, + pre_perm_ptr, + pre_signs_ptr, q_ptr, scale_ptr, + y_stride, D: tl.constexpr, DP2: tl.constexpr, STAGES: tl.constexpr, + PRE_STAGES: tl.constexpr, CHUNKS: tl.constexpr, inv_sqrt_block, + pre_inv_sqrt_block, + PRE: tl.constexpr, + GATHER: tl.constexpr, ): # One row per program: the INT8 scale needs the absolute maximum of the whole rotated row. # D is a multiple of the RPBH block, so the butterflies of the padded row never pair a # channel with another block. Blocks wider than a chunk imply D == DP2 (no padding). + # GATHER loads the row coalesced and permutes it in shared memory instead of loading + # scattered elements. PRE (requires GATHER) applies the first round of a two-round RPBH in + # front of the paper round, in FP32 and rounded to bf16 once like RPBHRotator's output. row = tl.program_id(0).to(tl.int64) + base = y_ptr + row * y_stride if CHUNKS == 1: cols = tl.arange(0, DP2) mask = cols < D - src = tl.load(perm_ptr + cols, mask=mask, other=0).to(tl.int32) - signs = tl.load(signs_ptr + cols, mask=mask, other=0).to(tl.float32) - values = tl.load(y_ptr + row * D + src, mask=mask, other=0.0).to(tl.float32) * signs + if GATHER: + values = tl.load(base + cols, mask=mask, other=0.0) + if PRE: + values = _signed_gather(values, pre_perm_ptr, pre_signs_ptr, cols, mask) + values = (_fwht(values, DP2, PRE_STAGES) * pre_inv_sqrt_block).to(tl.bfloat16) + values = _signed_gather(values, perm_ptr, signs_ptr, cols, mask) + else: + src = tl.load(perm_ptr + cols, mask=mask, other=0).to(tl.int32) + signs = tl.load(signs_ptr + cols, mask=mask, other=0).to(tl.float32) + values = tl.load(base + src, mask=mask, other=0.0).to(tl.float32) * signs values = _fwht(values, DP2, STAGES) * inv_sqrt_block scale = libdevice.div_rn(tl.maximum(tl.max(tl.abs(values), axis=0), 1e-12), 127.0) q = libdevice.rint(libdevice.div_rn(values, scale)) @@ -339,7 +426,7 @@ def _rpbh_absmax_int8_kernel( tl.store(q_ptr + row * D + cols, q.to(tl.int8), mask=mask) else: v0, v1, v2, v3 = _rotated_chunks( - y_ptr, perm_ptr, signs_ptr, row, 0, 1.0, D, CHUNKS, STAGES - _CHUNK_STAGES, False + base, perm_ptr, signs_ptr, 0, 0, 1.0, D, CHUNKS, STAGES - _CHUNK_STAGES, False ) v0 *= inv_sqrt_block v1 *= inv_sqrt_block @@ -475,6 +562,8 @@ class ActivationQuantizer: is quantized to INT8 with one absmax scale per token (8-bit activations) and the call 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. + The absmax kernel applies it in the same pass where the padded row is at most + ``GATHER_MAX_ROW`` wide (``fused_pre_rotation``); otherwise ``RPBHRotator`` runs in front. """ def __init__(self, rotation, codebook, eps: float, device: torch.device, pre_rotation=None): @@ -489,12 +578,23 @@ def __init__(self, rotation, codebook, eps: float, device: torch.device, pre_rot self.perm = rotation.permutation.to(device=device, dtype=torch.int32).contiguous() self.signs = rotation.signs.to(device=device, dtype=torch.int8).contiguous() self.absmax = codebook is None + self.fused_pre_rotation = False if self.absmax: if self.block > FWHT_CHUNK and self.dim not in (2 * FWHT_CHUNK, 4 * FWHT_CHUNK): raise ValueError( "8-bit activations with RPBH blocks above 4096 need an 8192 or 16384 wide input" ) self.levels, self.scale = 256, 1.0 + self.gather = ( + HAS_GATHER + and self.block <= FWHT_CHUNK + and triton.next_power_of_2(self.dim) <= GATHER_MAX_ROW + ) + self.fused_pre_rotation = ( + self.pre_rotator is not None + and self.gather + and self.pre_rotator.block <= FWHT_CHUNK + ) return self.bounds = codebook.boundaries.to(device=device, dtype=torch.float32).contiguous() self.levels = int(codebook.centroids.numel()) @@ -536,8 +636,11 @@ def __call__( ) 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. + # The eager model rounds the prologue output to bf16, then the layer rotates it. The + # fused kernel reads that bf16 output; RPBHRotator keeps a plain wider input's dtype. plain = prologue == PROLOGUE_NONE and norm in (None, NORM_NONE) and mod in (None, MOD_NONE) + if self.fused_pre_rotation and (not plain or x.dtype == torch.bfloat16): + return self._quantize(x, pre_rotate=True, **fields) y = x if plain else self._quantize(x, prologue_only=True, **fields)[0] return self._quantize(self.pre_rotator(y)) @@ -555,6 +658,7 @@ def _quantize( mod_index=None, rms_eps=1e-5, prologue_only=False, + pre_rotate=False, ): if x.dim() != 2 or x.stride(1) != 1: x = x.reshape(-1, self.dim).contiguous() @@ -568,6 +672,11 @@ def _quantize( norm_mode = norm if mod is not None: mod_mode = mod + plain = not multiply and norm_mode == NORM_NONE and mod_mode == MOD_NONE + if self.absmax and plain and not prologue_only and x.dtype == torch.bfloat16: + # The absmax kernel reads a plain bf16 row at any row stride; the token norms of a + # prologue pass would only be overwritten by the scales. + return self._absmax(x, pre_rotate) if x2 is not None and (x2.dim() != 2 or x2.stride(1) != 1): x2 = x2.reshape(rows, self.dim).contiguous() if mod_mode != MOD_NONE: @@ -577,12 +686,7 @@ def _quantize( mod_shift = mod_shift.reshape(-1, self.dim).contiguous() # Without a prologue a bf16 input row is its own prologue output; only norms are written. passthrough = ( - not multiply - and norm_mode == NORM_NONE - and mod_mode == MOD_NONE - and not prologue_only - and x.dtype == torch.bfloat16 - and x.stride(0) == self.dim + plain and not prologue_only and x.dtype == torch.bfloat16 and x.stride(0) == self.dim ) y = x if not passthrough: @@ -619,21 +723,8 @@ def _quantize( if prologue_only: return y, norms if self.absmax: - row_width = triton.next_power_of_2(self.dim) - _rpbh_absmax_int8_kernel[(rows,)]( - y, - self.perm, - self.signs, - q, - norms, - D=self.dim, - DP2=row_width, - STAGES=self.block.bit_length() - 1, - CHUNKS=row_width // FWHT_CHUNK if self.block > FWHT_CHUNK else 1, - inv_sqrt_block=self.inv_sqrt_block, - num_warps=16 if row_width >= 16384 else 8, - ) - return q, norms + # The scales overwrite the token norms. + return self._absmax(y, pre_rotate, q, norms) _rpbh_int8_kernel[(rows * self.num_blocks,)]( y, norms, @@ -653,6 +744,45 @@ def _quantize( ) return q, norms + def _absmax(self, y, pre_rotate, q=None, scales=None): + """INT8 codes and per-token scales of a bf16 ``[rows, dim]`` row view (unit column + stride), with the first RPBH round in front when ``pre_rotate``.""" + rows = y.shape[0] + if q is None: + q = torch.empty((rows, self.dim), device=y.device, dtype=torch.int8) + scales = torch.empty(rows, device=y.device, dtype=torch.float32) + if rows == 0: + return q, scales + pre = self.pre_rotator if pre_rotate else None + row_width = triton.next_power_of_2(self.dim) + _rpbh_absmax_int8_kernel[(rows,)]( + y, + self.perm, + self.signs, + self.perm if pre is None else pre.perm, + self.signs if pre is None else pre.signs, + q, + scales, + y.stride(0), + D=self.dim, + DP2=row_width, + STAGES=self.block.bit_length() - 1, + PRE_STAGES=0 if pre is None else pre.block.bit_length() - 1, + CHUNKS=row_width // FWHT_CHUNK if self.block > FWHT_CHUNK else 1, + inv_sqrt_block=self.inv_sqrt_block, + pre_inv_sqrt_block=1.0 if pre is None else pre.inv_sqrt_block, + PRE=pre is not None, + GATHER=self.gather, + num_warps=( + _absmax_warps(row_width, self.block, pre is not None) + if self.gather + else 16 + if row_width >= 16384 + else 8 + ), + ) + return q, scales + @triton.jit def _rows_int8_kernel(x_ptr, q_ptr, scale_ptr, D: tl.constexpr, DP2: tl.constexpr): diff --git a/src/orbitquant/layers.py b/src/orbitquant/layers.py index a62b018..bfe94d7 100644 --- a/src/orbitquant/layers.py +++ b/src/orbitquant/layers.py @@ -1346,7 +1346,8 @@ def probe() -> bool: return self._cached_runtime_probe("triton_absmax_int8", probe) def _absmax_int8_cuda_forward(self, x: torch.Tensor) -> torch.Tensor: - """INT8 GEMM of absmax INT8 activations against the packed W2/W4 weight codes.""" + """INT8 GEMM of absmax INT8 activations against the packed W2/W4 weight codes. ``x`` is + the unrotated input: the quantizer applies both RPBH rounds.""" from orbitquant.kernels import triton_activation from orbitquant.kernels.triton_cuda import ( matmul_int8_activations_packed_lowbit_fused_with_triton, @@ -1356,14 +1357,15 @@ def _absmax_int8_cuda_forward(self, x: torch.Tensor) -> torch.Tensor: device_key = str(x.device) if self._absmax_quantizer is None or self._absmax_quantizer[0] != device_key: quantizer = triton_activation.ActivationQuantizer( - self.rotation, None, self.activation_eps, x.device + self.rotation, None, self.activation_eps, x.device, pre_rotation=self.pre_rotation ) self._absmax_quantizer = (device_key, quantizer) quantizer = self._absmax_quantizer[1] rows_x = x.reshape(-1, self.in_features) + signature = ("triton_absmax_int8", id(self.rotation), id(self.pre_rotation)) (codes, token_scales), self.last_activation_cache_hit = ( _get_or_create_shared_activation_quantization( - x, ("triton_absmax_int8", id(self.rotation)), lambda: quantizer(rows_x) + x, signature, lambda: quantizer(rows_x) ) ) _, _, weight_codes, weight_scale = self._int8_surrogate_constants(x.device) @@ -1492,8 +1494,6 @@ def factory() -> torch.Tensor: def _forward_quantized(self, x: torch.Tensor) -> torch.Tensor: self._ensure_derived_constants(x.device) - if self.pre_rotation is not None: - x = self._pre_rotate(x) runtime_mode = ( self._resolve_auto_fused_runtime(x) if self.runtime_mode == "auto_fused" @@ -1511,6 +1511,16 @@ def _forward_quantized(self, x: torch.Tensor) -> torch.Tensor: elif runtime_mode == "native_packed_matmul": self._validate_native_packed_matmul_input(x) + absmax_int8 = self.activation_quantizer == "absmax_int8" and runtime_mode not in { + "debug_no_quant", + "debug_no_activation_quant", + } + if absmax_int8 and self._absmax_int8_cuda_available(x, runtime_mode): + # The INT8 activation kernel applies the first RPBH round itself. + return self._absmax_int8_cuda_forward(x) + if self.pre_rotation is not None: + x = self._pre_rotate(x) + if runtime_mode == "native_packed_matmul" and self._fused_w2a4_available(x): activation_constants = self._activation_kernel_constant_tensors(x.device) activation_codes, activation_scale, weight_codes, weight_scale = ( @@ -1797,13 +1807,6 @@ def _forward_quantized(self, x: torch.Tensor) -> torch.Tensor: weight_k_major=False, ) - absmax_int8 = self.activation_quantizer == "absmax_int8" and runtime_mode not in { - "debug_no_quant", - "debug_no_activation_quant", - } - if absmax_int8 and self._absmax_int8_cuda_available(x, runtime_mode): - return self._absmax_int8_cuda_forward(x) - if absmax_int8: self.last_activation_kernel_backend = "torch_absmax_int8" rotated_x = self._absmax_int8_activations(x) diff --git a/tests/test_absmax_activations.py b/tests/test_absmax_activations.py index fca7e1c..1abb7d6 100644 --- a/tests/test_absmax_activations.py +++ b/tests/test_absmax_activations.py @@ -201,3 +201,31 @@ def test_cuda_absmax_int8_gemm_matches_the_reference(weight_bits, rows): # The GEMM uses the INT8 surrogate of the weight centroids (docs/kernel-audit.md). error = (output - expected).pow(2).sum() / expected.pow(2).sum() assert float(error) < 2e-3 + + +@pytest.mark.skipif(not _cuda_triton_available(), reason="CUDA with Triton required") +@pytest.mark.parametrize("in_features", [3840, 9536]) +def test_cuda_absmax_int8_layer_applies_the_pre_rotation_once(in_features): + torch.manual_seed(1) + config = OrbitQuantConfig( + runtime_mode="triton_packed_matmul", + activation_bits=8, + activation_quantizer="absmax_int8", + rotation_rounds="auto", + ) + linear = torch.nn.Linear(in_features, 256).to(torch.bfloat16) + layer = OrbitQuantLinear.from_linear(linear, config=config, module_name="p").to("cuda") + assert layer.pre_rotation is not None + x = torch.randn(80, in_features, device="cuda", dtype=torch.bfloat16) + x[:, 11] *= 200 + + with torch.no_grad(): + output = layer(x) + assert layer.last_activation_kernel_backend == "triton_cuda_absmax_int8" + # v0.16 order: the first round as its own pass, then the one-round quantizer. + rotated = layer._pre_rotate(x) + layer.pre_rotation = None + layer._absmax_quantizer = None + expected = layer(rotated) + + assert torch.equal(output, expected) diff --git a/tests/test_fused_family_kernels.py b/tests/test_fused_family_kernels.py index c830a3e..5d1f2fa 100644 --- a/tests/test_fused_family_kernels.py +++ b/tests/test_fused_family_kernels.py @@ -262,8 +262,20 @@ def test_mixed_width_group_is_one_gemm_matching_its_projections(): assert torch.equal(again, fused) -@pytest.mark.parametrize("dim", [3360, 13568, 6144, 8192, 16384]) -def test_absmax_int8_activations_match_rotated_reference(dim): +def _absmax_reference(rotation, x): + """Per-token absmax INT8 of the FP32 rotation with the kernel's numerics: IEEE divisions + (tensor by tensor; PyTorch multiplies by the reciprocal of a Python scalar) and round half + to even.""" + rotated = rotation.apply_to_activations(x.float()) + peak = rotated.abs().amax(-1).clamp_min(1e-12) + scales = peak / torch.full_like(peak, 127.0) + codes = torch.round(rotated / scales[:, None]).clamp(-127, 127).to(torch.int8) + return codes, scales + + +@pytest.mark.parametrize("dim", [1920, 3360, 13568, 6144, 8192, 16384]) +@pytest.mark.parametrize("strided", [False, True]) +def test_absmax_int8_activations_match_rotated_reference(dim, strided): _require_cuda() from orbitquant.kernels import triton_activation as act from orbitquant.rotations import get_rpbh_rotation @@ -271,15 +283,12 @@ def test_absmax_int8_activations_match_rotated_reference(dim): torch.manual_seed(7) rotation = get_rpbh_rotation(dim=dim, seed=0, block_size="paper") quantizer = act.ActivationQuantizer(rotation, None, 1e-10, torch.device("cuda")) - x = torch.randn(37, dim, device="cuda", dtype=torch.bfloat16) + x = torch.randn(37, dim + 64 * strided, device="cuda", dtype=torch.bfloat16)[:, :dim] x[:, 5] *= 40 # an outlier channel codes, scales = quantizer(x) - rotated = rotation.apply_to_activations(x.float()) - expected_scales = rotated.abs().amax(-1) / 127 - torch.testing.assert_close(scales, expected_scales, rtol=1e-5, atol=0) - expected = torch.round(rotated / expected_scales[:, None]).clamp(-127, 127) - assert (codes.float() - expected).abs().max().item() <= 1 - assert (codes.float() != expected).float().mean().item() < 1e-3 + expected_codes, expected_scales = _absmax_reference(rotation, x) + assert torch.equal(scales, expected_scales) + assert torch.equal(codes, expected_codes) @pytest.mark.parametrize("dim", [3360, 8192, 16384]) @@ -304,7 +313,14 @@ def test_lloyd_max_activation_codes_match_the_rotated_reference(dim): @pytest.mark.parametrize( ("dim", "block_size"), - [(3360, "paper"), (3840, "paper"), (3072, 1024), (8192, 8192), (16384, 16384)], + [ + (3360, "paper"), + (3840, "paper"), + (9536, "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): @@ -354,6 +370,88 @@ def test_activation_quantizer_applies_the_pre_rotation_after_the_prologue(codebo assert torch.equal(two_round(x, prologue_only=True, **kwargs)[0], prologue) +def _prologue_case(case, rows, dim, device): + """An input with a massive channel and the quantizer arguments of one prologue variant.""" + from orbitquant.kernels import triton_activation as act + + def vector(spread=0.1, offset=0.0, count=None): + shape = (dim,) if count is None else (count, dim) + return (torch.randn(shape, device=device) * spread + offset).to(torch.bfloat16) + + wide = torch.randn(rows, dim + 64, device=device, dtype=torch.bfloat16) + x = wide[:, :dim] if case == "strided" else wide[:, :dim].contiguous() + x[:, 3] *= 300 + index = torch.randint(0, 3, (rows,), device=device, dtype=torch.int32) + kwargs = { + "plain": {}, + "strided": {}, + "fp32": {}, + "norm_modulate": dict( + prologue=act.PROLOGUE_NORM_MODULATE, + norm_weight=vector(), + mod_scale=vector(), + mod_shift=vector(), + ), + "rms_scale_indexed": dict( + norm=act.NORM_RMS, + mod=act.MOD_SCALE, + mod_index=index, + norm_weight=vector(offset=1.0), + mod_scale=vector(0.2, 1.0, count=3), + rms_eps=1e-6, + ), + "rms_plain_shift_indexed": dict( + norm=act.NORM_RMS_PLAIN, + mod=act.MOD_SHIFT_SCALE, + mod_index=index, + mod_scale=vector(count=3), + mod_shift=vector(count=3), + ), + "layer_one_plus": dict(norm=act.NORM_LAYER, mod=act.MOD_ONE_PLUS_SCALE, mod_scale=vector()), + "rms_torch": dict(norm=act.NORM_RMS_TORCH, norm_weight=vector(offset=1.0)), + "multiply": dict(prologue=act.PROLOGUE_MULTIPLY, x2=vector(1.0, count=rows)), + }[case] + return (x.float() if case == "fp32" else x), kwargs + + +@pytest.mark.parametrize("dim", [1920, 3840, 9536, 16640]) +@pytest.mark.parametrize( + "case", + [ + "plain", + "strided", + "fp32", + "norm_modulate", + "rms_scale_indexed", + "rms_plain_shift_indexed", + "layer_one_plus", + "rms_torch", + "multiply", + ], +) +def test_absmax_quantizer_fuses_the_pre_rotation_bit_for_bit(dim, case): + _require_cuda() + from orbitquant.kernels import triton_activation as act + from orbitquant.rotations import get_pre_rotation, get_rpbh_rotation + + torch.manual_seed(12) + device = torch.device("cuda") + rotation = get_rpbh_rotation(dim=dim, seed=0, block_size="paper") + pre = get_pre_rotation(dim, seed=0, block_size="paper", rotation_rounds="auto") + one_round = act.ActivationQuantizer(rotation, None, 1e-10, device) + two_round = act.ActivationQuantizer(rotation, None, 1e-10, device, pre_rotation=pre) + # d = 16640 (h = 256) pads to a 32768-wide row: RPBHRotator runs in front of the kernel. + assert two_round.fused_pre_rotation == (act.HAS_GATHER and dim < 16384) + x, kwargs = _prologue_case(case, 45, dim, device) + + # The v0.16 order: bf16 prologue output, RPBHRotator rounded to its dtype, one round. + prologue = x if not kwargs else one_round(x, prologue_only=True, **kwargs)[0] + expected_codes, expected_scales = one_round(act.RPBHRotator(pre, device)(prologue)) + codes, scales = two_round(x, **kwargs) + assert torch.equal(scales, expected_scales) + assert torch.equal(codes, expected_codes) + + @pytest.mark.parametrize("activation_bits", [4, 8]) def test_fused_two_round_projection_matches_the_per_layer_path(activation_bits): _require_cuda() diff --git a/uv.lock b/uv.lock index 04e1260..edf5c00 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.16.0" +version = "0.16.1" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },