diff --git a/README.md b/README.md index 8ef3e7a..bbbd8fb 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,9 @@ The implementation is clean-room and Apache-2.0 licensed. `from_pretrained()` integration. - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full dequantized weight matrix. +- A fused DiT runtime for Krea 2 (`orbitquant.runtime.krea2`): grouped INT8 + GEMMs with SwiGLU/residual epilogues and an INT8 attention kernel, 2.4x + faster than the eager path on an RTX 4060 Ti. - `torch.compile(fullgraph=True)` support: quantized forwards run behind a registered custom op with a fake implementation, so compiled models match eager output exactly. The forward is also CUDA-Graph-capturable; @@ -443,6 +446,49 @@ system, C++ runtime, and architecture. See [`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark methodology, and local package verification. +## Fused DiT Runtime (Krea 2) + +Version 0.10 speeds up every W4A4 model on CUDA without API changes: + +- Triton kernels take row counts as runtime arguments, so a new prompt length + or image size no longer recompiles them. +- Large-row W4A4 layers use an INT8 GEMM that applies the scale epilogue in + registers; the output is bit-identical to the previous `torch._int_mm` path. +- Projections that read the same tensor (attention Q/K/V, SwiGLU gate/up) + share one activation quantization. The reuse needs version counters, so run + inference under `torch.no_grad()`; `torch.inference_mode()` disables it. + +On an RTX 4060 Ti the Qwen-Image 2.1 W4A4 pipeline went from 15.0 to 10.3 s +per 1024x1024 image with pixel-identical output. + +`orbitquant.runtime.krea2` goes further for Krea 2 Turbo: each block runs one +grouped Q|K|V|gate GEMM with the sigmoid gate in its epilogue, a fused Q/K +RMSNorm + RoPE kernel, SwiGLU and gated residual updates in GEMM epilogues, +W8A8 down projections, and optionally an INT8 Q.K^T / FP16 P.V attention +kernel (the SageAttention v1 scheme). Blocks whose Q/K RMSNorm scales one +channel far above the rest keep BF16 attention. + +```python +import torch +from diffusers import Krea2Pipeline +from orbitquant.runtime.krea2 import Krea2FastRunner, install, save_fused + +pipe = Krea2Pipeline.from_pretrained(model_dir, dtype=torch.bfloat16) +fused = install(pipe.transformer, attention="int8") # on the host, before .to("cuda") +save_fused(fused, "fused_blocks.safetensors") # optional: map it next time +pipe.to("cuda") +runner = Krea2FastRunner(pipe, fused) +with torch.no_grad(): + states = runner.encode(prompt, torch.device("cuda")) + latents = runner.denoise(states, width=1024, height=1024) + image = runner.decode(latents, width=1024, height=1024) +``` + +Building the fused buffers copies the block weights into new host memory; +`install(..., fused_path="fused_blocks.safetensors")` maps a saved file instead. +On an RTX 4060 Ti a 1024x1024 eight-step image takes 10.8 s (25.9 s with the +stock pipeline), and the DiT needs 7.3 GB instead of 10 GB. + ## Validated Architecture Coverage The integration suite instantiates and inventories encoder-only, decoder-only, diff --git a/pyproject.toml b/pyproject.toml index 7ce5ccd..5d60b16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.9.8" +version = "0.10.0" 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 fd7148e..013187c 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.9.8" +__version__ = "0.10.0" register_hf_quantizers() diff --git a/src/orbitquant/kernels/triton_activation.py b/src/orbitquant/kernels/triton_activation.py new file mode 100644 index 0000000..76542d1 --- /dev/null +++ b/src/orbitquant/kernels/triton_activation.py @@ -0,0 +1,325 @@ +"""Activation-side Triton kernels for fused DiT blocks. + +* RPBH + Lloyd-Max activation quantization straight to INT8 surrogate codes, optionally fed by + a prologue that replaces the elementwise ops in front of a projection (RMSNorm with a + ``1 + weight`` scale followed by AdaLN modulation, or an elementwise product such as an + attention output gate). Every value eager code rounds to bf16 is rounded at the same point. +* per-token absmax INT8 for W8A8 projections (``Int8RowLinear`` numerics), +* per-head RMSNorm + interleaved rotary embedding for Q/K, +* packed W4 -> INT8 surrogate weight decode. +""" + +from __future__ import annotations + +import torch + +from orbitquant.kernels.triton_cuda import _load_triton, fit_int8_centroid_surrogate + +triton, tl = _load_triton() + +from triton.language.extra import libdevice # noqa: E402 + +PROLOGUE_NONE = 0 +PROLOGUE_NORM_MODULATE = 1 +PROLOGUE_MULTIPLY = 2 + + +@triton.jit +def _bf16(x): + return x.to(tl.bfloat16).to(tl.float32) + + +@triton.jit +def _prologue_kernel( + x_ptr, + x2_ptr, + w_ptr, + scale_ptr, + shift_ptr, + y_ptr, + norms_ptr, + x_stride, + x2_stride, + D: tl.constexpr, + DP2: tl.constexpr, + rms_eps, + PROLOGUE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, DP2) + mask = offs < D + x = tl.load(x_ptr + row * x_stride + offs, mask=mask, other=0.0).to(tl.float32) + if PROLOGUE == 1: + mean_sq = libdevice.div_rn(tl.sum(x * x, axis=0), D * 1.0) + inv = libdevice.rsqrt(mean_sq + rms_eps) + w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + 1.0 + normed = _bf16(x * inv * w) + scale = tl.load(scale_ptr + offs, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + offs, mask=mask, other=0.0).to(tl.float32) + y = _bf16(_bf16(_bf16(1.0 + scale) * normed) + shift) + elif PROLOGUE == 2: + x2 = tl.load(x2_ptr + row * x2_stride + offs, mask=mask, other=0.0).to(tl.float32) + y = _bf16(x * x2) + else: + y = x + tl.store(y_ptr + row * D + offs, y.to(tl.bfloat16), mask=mask) + tl.store(norms_ptr + row, tl.sqrt(tl.sum(y * y, axis=0))) + + +@triton.jit +def _fwht_local_stage(values, block_size: tl.constexpr, stage_width: tl.constexpr): + butterfly = tl.reshape(values, (block_size // (stage_width * 2), 2, stage_width)) + butterfly = tl.permute(butterfly, (0, 2, 1)) + left, right = tl.split(butterfly) + butterfly = tl.join(left + right, left - right) + butterfly = tl.permute(butterfly, (0, 2, 1)) + return tl.reshape(butterfly, (block_size,)) + + +@triton.jit +def _rpbh_int8_kernel( + y_ptr, + norms_ptr, + perm_ptr, + signs_ptr, + bounds_ptr, + codes_ptr, + q_ptr, + D: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + OB: tl.constexpr, + STAGES: tl.constexpr, + LEVELS: tl.constexpr, + eps, + inv_sqrt_block, +): + row_block = tl.program_id(0) + row = (row_block // NUM_BLOCKS).to(tl.int64) + block = row_block % NUM_BLOCKS + 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) + norm = tl.load(norms_ptr + row).to(tl.float32) + values = tl.load(y_ptr + row * D + src).to(tl.float32) + values = values * signs / (norm + eps) + if STAGES > 0: + values = _fwht_local_stage(values, OB, 1) + if STAGES > 1: + values = _fwht_local_stage(values, OB, 2) + if STAGES > 2: + values = _fwht_local_stage(values, OB, 4) + if STAGES > 3: + values = _fwht_local_stage(values, OB, 8) + if STAGES > 4: + values = _fwht_local_stage(values, OB, 16) + if STAGES > 5: + values = _fwht_local_stage(values, OB, 32) + if STAGES > 6: + values = _fwht_local_stage(values, OB, 64) + if STAGES > 7: + values = _fwht_local_stage(values, OB, 128) + if STAGES > 8: + values = _fwht_local_stage(values, OB, 256) + if STAGES > 9: + values = _fwht_local_stage(values, OB, 512) + if STAGES > 10: + values = _fwht_local_stage(values, OB, 1024) + if STAGES > 11: + values = _fwht_local_stage(values, OB, 2048) + values *= inv_sqrt_block + indices = tl.zeros((OB,), dtype=tl.int32) + for level in tl.static_range(0, LEVELS - 1): + indices += (values > tl.load(bounds_ptr + level)).to(tl.int32) + tl.store(q_ptr + row * D + cols, tl.load(codes_ptr + indices).to(tl.int8)) + + +class ActivationQuantizer: + """RPBH/Lloyd-Max INT8 surrogate activation quantizer of one OrbitQuantLinear input. + + 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. + """ + + def __init__(self, rotation, codebook, eps: float, device: torch.device): + self.dim = int(rotation.dim) + self.block = int(rotation.block_size) + self.num_blocks = int(rotation.num_blocks) + if self.block > 4096: + raise ValueError("fused activation quantization supports RPBH blocks up to 4096") + self.inv_sqrt_block = float(rotation.normalization) + self.eps = float(eps) + self.perm = rotation.permutation.to(device=device, dtype=torch.int32).contiguous() + self.signs = rotation.signs.to(device=device, dtype=torch.int8).contiguous() + self.bounds = codebook.boundaries.to(device=device, dtype=torch.float32).contiguous() + self.levels = int(codebook.centroids.numel()) + codes, self.scale = fit_int8_centroid_surrogate(codebook.centroids) + self.codes = codes.to(device=device, dtype=torch.int8).contiguous() + + def __call__( + self, + x, + *, + prologue=PROLOGUE_NONE, + x2=None, + norm_weight=None, + mod_scale=None, + mod_shift=None, + rms_eps=1e-5, + ): + """Return ``(codes int8 [rows, dim], token_norms fp32 [rows])``.""" + if x.dim() != 2 or x.stride(1) != 1: + x = x.reshape(-1, self.dim).contiguous() + rows = x.shape[0] + if x2 is not None and (x2.dim() != 2 or x2.stride(1) != 1): + x2 = x2.reshape(rows, self.dim).contiguous() + y = torch.empty((rows, self.dim), device=x.device, dtype=torch.bfloat16) + norms = torch.empty(rows, device=x.device, dtype=torch.float32) + q = torch.empty((rows, self.dim), device=x.device, dtype=torch.int8) + if rows == 0: + return q, norms + _prologue_kernel[(rows,)]( + x, + x2 if x2 is not None else x, + norm_weight if norm_weight is not None else norms, + mod_scale if mod_scale is not None else x, + mod_shift if mod_shift is not None else x, + y, + norms, + x.stride(0), + x2.stride(0) if x2 is not None else 0, + D=self.dim, + DP2=triton.next_power_of_2(self.dim), + rms_eps=float(rms_eps), + PROLOGUE=int(prologue), + num_warps=8, + ) + _rpbh_int8_kernel[(rows * self.num_blocks,)]( + y, + norms, + self.perm, + self.signs, + self.bounds, + self.codes, + q, + D=self.dim, + NUM_BLOCKS=self.num_blocks, + OB=self.block, + STAGES=self.block.bit_length() - 1, + LEVELS=self.levels, + eps=self.eps, + inv_sqrt_block=self.inv_sqrt_block, + num_warps=8 if self.block >= 512 else 4, + ) + return q, norms + + +@triton.jit +def _rows_int8_kernel(x_ptr, q_ptr, scale_ptr, D: tl.constexpr, DP2: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, DP2) + mask = offs < D + x = tl.load(x_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + scale = libdevice.div_rn(tl.maximum(tl.max(tl.abs(x), axis=0), 1e-12), 127.0) + q = libdevice.rint(libdevice.div_rn(x, scale)) + q = tl.minimum(tl.maximum(q, -127.0), 127.0) + tl.store(q_ptr + row * D + offs, q.to(tl.int8), mask=mask) + tl.store(scale_ptr + row, scale) + + +def quantize_rows_int8_with_triton(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token absmax INT8 with ``Int8RowLinear`` numerics (IEEE divisions, round-half-even).""" + dim = x.shape[-1] + x = x.reshape(-1, dim).contiguous() + rows = x.shape[0] + q = torch.empty((rows, dim), device=x.device, dtype=torch.int8) + scales = torch.empty(rows, device=x.device, dtype=torch.float32) + if rows: + _rows_int8_kernel[(rows,)]( + x, q, scales, D=dim, DP2=triton.next_power_of_2(dim), num_warps=8 + ) + return q, scales + + +@triton.jit +def _qk_norm_rope_kernel( + src_ptr, + src_stride, + out_ptr, + w_ptr, + cos_ptr, + sin_ptr, + rms_eps, + HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1) + half = tl.arange(0, HEAD_DIM // 2) + base = src_ptr + row * src_stride + head * HEAD_DIM + even = tl.load(base + 2 * half).to(tl.float32) + odd = tl.load(base + 2 * half + 1).to(tl.float32) + mean_sq = libdevice.div_rn( + tl.sum(even * even, axis=0) + tl.sum(odd * odd, axis=0), HEAD_DIM * 1.0 + ) + inv = libdevice.rsqrt(mean_sq + rms_eps) + even = _bf16(even * inv * (tl.load(w_ptr + 2 * half).to(tl.float32) + 1.0)) + odd = _bf16(odd * inv * (tl.load(w_ptr + 2 * half + 1).to(tl.float32) + 1.0)) + cos_even = tl.load(cos_ptr + row * HEAD_DIM + 2 * half) + cos_odd = tl.load(cos_ptr + row * HEAD_DIM + 2 * half + 1) + sin_even = tl.load(sin_ptr + row * HEAD_DIM + 2 * half) + sin_odd = tl.load(sin_ptr + row * HEAD_DIM + 2 * half + 1) + out_even = even * cos_even + (-odd) * sin_even + out_odd = odd * cos_odd + even * sin_odd + dst = out_ptr + (row * HEADS + head) * HEAD_DIM + tl.store(dst + 2 * half, out_even.to(tl.bfloat16)) + tl.store(dst + 2 * half + 1, out_odd.to(tl.bfloat16)) + + +def qk_norm_rope_with_triton(src, heads, head_dim, weight, cos, sin, rms_eps): + """Per-head RMSNorm (``1 + weight`` scale) then interleaved rotary embedding. + + ``src`` is a ``[rows, >= heads * head_dim]`` strided view; returns + ``[1, rows, heads, head_dim]`` bf16. Multiplies and adds are not contracted into FMAs, as + in the eager elementwise ops. + """ + rows = src.shape[0] + out = torch.empty((1, rows, heads, head_dim), device=src.device, dtype=torch.bfloat16) + if rows: + _qk_norm_rope_kernel[(rows, heads)]( + src, + src.stride(0), + out, + weight, + cos, + sin, + float(rms_eps), + HEADS=heads, + HEAD_DIM=head_dim, + num_warps=1, + enable_fp_fusion=False, + ) + return out + + +@triton.jit +def _decode_w4_kernel(packed_ptr, codes_ptr, out_ptr, total_bytes, BLOCK: tl.constexpr): + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offs < total_bytes + byte = tl.load(packed_ptr + offs, mask=mask, other=0).to(tl.int32) + low = tl.load(codes_ptr + (byte & 15)) + high = tl.load(codes_ptr + ((byte >> 4) & 15)) + pair = tl.reshape(tl.join(low, high), (2 * BLOCK,)) + out_offs = tl.program_id(0).to(tl.int64) * (2 * BLOCK) + tl.arange(0, 2 * BLOCK) + tl.store(out_ptr + out_offs, pair.to(tl.int8), mask=out_offs < 2 * total_bytes) + + +def decode_w4_to_int8_with_triton(packed, codes, out_features, in_features, out=None): + """Packed W4 indices (low nibble first) -> INT8 surrogate weight ``[out, in]``.""" + if out is None: + out = torch.empty((out_features, in_features), device=packed.device, dtype=torch.int8) + total = out_features * in_features // 2 + block = 1024 + _decode_w4_kernel[(triton.cdiv(total, block),)]( + packed.reshape(-1), codes, out, total, BLOCK=block + ) + return out diff --git a/src/orbitquant/kernels/triton_attention.py b/src/orbitquant/kernels/triton_attention.py new file mode 100644 index 0000000..69df6a5 --- /dev/null +++ b/src/orbitquant/kernels/triton_attention.py @@ -0,0 +1,234 @@ +"""Non-causal grouped-query attention with INT8 Q.K^T and FP16-accumulated P.V (Triton). + +This is the SageAttention (v1) scheme. GeForce Ada and Blackwell cards run BF16/FP16 tensor ops +with an FP32 accumulator at half the FP16-accumulator rate and INT8 at four times the BF16 +rate, so FlashAttention-style BF16 kernels leave most of the tensor throughput unused: + +* Q and K are quantized to INT8 per block of tokens and head; K first loses its per-channel + mean over the sequence, which shifts every score of a query by the same amount and leaves + the softmax unchanged. The softmax scale and ``log2(e)`` are folded into the Q scale so the + kernel exponentiates with ``exp2``. +* The online softmax runs in FP32. +* P and V are FP16; each ``BLOCK_N``-token tile of ``P @ V`` accumulates in FP16 and is added + to an FP32 accumulator. + +Tensors use the ``[batch=1, tokens, heads, head_dim]`` layout of DiT projections; rows may be +strided (e.g. a column slice of a fused Q|K|V GEMM output). +""" + +from __future__ import annotations + +import math + +import torch + +from orbitquant.kernels.triton_cuda import _load_triton +from orbitquant.kernels.triton_int8_gemm import autotune + +triton, tl = _load_triton() + +LOG2E = 1.4426950408889634 +BLOCK_Q = 128 +BLOCK_KV = 64 + + +@triton.jit +def _quantize_blocks_kernel( + x_ptr, + mean_ptr, + out_ptr, + scale_ptr, + rows, + stride_xr, + stride_xh, + stride_or, + stride_oh, + multiplier, + HAS_MEAN: tl.constexpr, + BLOCK: tl.constexpr, + D: tl.constexpr, +): + block = tl.program_id(0) + head = tl.program_id(1) + offs_r = block * BLOCK + tl.arange(0, BLOCK) + offs_d = tl.arange(0, D) + mask = offs_r[:, None] < rows + rows_x = offs_r[:, None].to(tl.int64) * stride_xr + head * stride_xh + x = tl.load(x_ptr + rows_x + offs_d[None, :], mask=mask, other=0.0).to(tl.float32) + if HAS_MEAN: + x = tl.where(mask, x - tl.load(mean_ptr + head * D + offs_d)[None, :], 0.0) + x = x * multiplier + scale = tl.maximum(tl.max(tl.max(tl.abs(x), axis=1), axis=0) / 127.0, 1e-20) + q = x / scale + q = tl.where(q >= 0, q + 0.5, q - 0.5).to(tl.int8) + rows_o = offs_r[:, None].to(tl.int64) * stride_or + head * stride_oh + tl.store(out_ptr + rows_o + offs_d[None, :], q, mask=mask) + tl.store(scale_ptr + head * tl.cdiv(rows, BLOCK) + block, scale) + + +@triton.jit +def _attention_tile( + acc, + l_i, + m_i, + q, + qk_scale, + k_ptrs, + v_ptrs, + n_valid, + MASKED: tl.constexpr, + BLOCK_N: tl.constexpr, +): + offs_n = tl.arange(0, BLOCK_N) + k = tl.load(k_ptrs, mask=offs_n[None, :] < n_valid, other=0) if MASKED else tl.load(k_ptrs) + qk = tl.dot(q, k, out_dtype=tl.int32).to(tl.float32) * qk_scale + if MASKED: + qk = tl.where(offs_n[None, :] < n_valid, qk, float("-inf")) + m_new = tl.maximum(m_i, tl.max(qk, 1)) + p = tl.math.exp2(qk - m_new[:, None]) + alpha = tl.math.exp2(m_i - m_new) + l_i = l_i * alpha + tl.sum(p, 1) + acc = acc * alpha[:, None] + v = tl.load(v_ptrs, mask=offs_n[:, None] < n_valid, other=0.0) if MASKED else tl.load(v_ptrs) + acc += tl.dot(p.to(tl.float16), v, out_dtype=tl.float16).to(tl.float32) + return acc, l_i, m_new + + +@triton.jit +def _attention_kernel( + q_ptr, + k_ptr, + v_ptr, + q_scale_ptr, + k_scale_ptr, + o_ptr, + seq, + stride_qr, + stride_qh, + stride_kr, + stride_kh, + stride_vr, + stride_vh, + stride_or, + stride_oh, + GROUPS: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + start_m = tl.program_id(0) + head = tl.program_id(1) + kv_head = head // GROUPS + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + q_rows = offs_m[:, None].to(tl.int64) * stride_qr + head * stride_qh + q = tl.load(q_ptr + q_rows + offs_d[None, :], mask=offs_m[:, None] < seq, other=0) + q_scale = tl.load(q_scale_ptr + head * tl.cdiv(seq, BLOCK_M) + start_m) + k_scales = k_scale_ptr + kv_head * tl.cdiv(seq, BLOCK_N) + k_ptrs = k_ptr + offs_n[None, :].to(tl.int64) * stride_kr + kv_head * stride_kh + k_ptrs += offs_d[:, None] + v_ptrs = v_ptr + offs_n[:, None].to(tl.int64) * stride_vr + kv_head * stride_vh + v_ptrs += offs_d[None, :] + m_i = tl.full([BLOCK_M], float("-inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, D], tl.float32) + full = seq // BLOCK_N + for block in range(0, full): + qk_scale = q_scale * tl.load(k_scales + block) + acc, l_i, m_i = _attention_tile( + acc, l_i, m_i, q, qk_scale, k_ptrs, v_ptrs, BLOCK_N, False, BLOCK_N + ) + k_ptrs += BLOCK_N * stride_kr + v_ptrs += BLOCK_N * stride_vr + if full * BLOCK_N < seq: + qk_scale = q_scale * tl.load(k_scales + full) + acc, l_i, m_i = _attention_tile( + acc, l_i, m_i, q, qk_scale, k_ptrs, v_ptrs, seq - full * BLOCK_N, True, BLOCK_N + ) + out = acc / l_i[:, None] + o_rows = offs_m[:, None].to(tl.int64) * stride_or + head * stride_oh + tl.store( + o_ptr + o_rows + offs_d[None, :], + out.to(o_ptr.dtype.element_ty), + mask=offs_m[:, None] < seq, + ) + + +_attention = autotune( + [triton.Config({}, num_warps=w, num_stages=s) for w in (4, 8) for s in (2, 3, 4)], + ["GROUPS", "D"], +)(_attention_kernel) + + +def _quantize_blocks(x, block, multiplier, mean=None): + rows, heads, dim = x.shape + out = torch.empty((rows, heads, dim), device=x.device, dtype=torch.int8) + scales = torch.empty((heads, triton.cdiv(rows, block)), device=x.device, dtype=torch.float32) + _quantize_blocks_kernel[(triton.cdiv(rows, block), heads)]( + x, + mean if mean is not None else scales, + out, + scales, + rows, + x.stride(0), + x.stride(1), + out.stride(0), + out.stride(1), + float(multiplier), + HAS_MEAN=mean is not None, + BLOCK=block, + D=dim, + num_warps=4, + ) + return out, scales + + +def int8_attention_with_triton( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + sm_scale: float | None = None, +) -> torch.Tensor: + """``softmax(q k^T * sm_scale) v`` for ``[1, S, Hq, D]`` queries and ``[1, S, Hkv, D]`` + keys/values (``Hq`` a multiple of ``Hkv``, ``D`` a power of two, last dim contiguous). + Returns ``[1, S, Hq, D]`` in the query dtype.""" + if query.dim() != 4 or query.shape[0] != 1: + raise ValueError("expected [1, tokens, heads, head_dim] tensors") + q, k, v = query[0], key[0], value[0] + seq, heads, dim = q.shape + kv_heads = k.shape[1] + if heads % kv_heads: + raise ValueError(f"{heads} query heads are not a multiple of {kv_heads} kv heads") + if dim & (dim - 1) or q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1: + raise ValueError("head_dim must be a power of two and contiguous") + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(dim) + q_codes, q_scales = _quantize_blocks(q, BLOCK_Q, sm_scale * LOG2E) + k_codes, k_scales = _quantize_blocks(k, BLOCK_KV, 1.0, k.float().mean(dim=0).contiguous()) + v16 = v.to(torch.float16) + out = torch.empty((1, seq, heads, dim), device=q.device, dtype=query.dtype) + o = out[0] + _attention[(triton.cdiv(seq, BLOCK_Q), heads)]( + q_codes, + k_codes, + v16, + q_scales, + k_scales, + o, + seq, + q_codes.stride(0), + q_codes.stride(1), + k_codes.stride(0), + k_codes.stride(1), + v16.stride(0), + v16.stride(1), + o.stride(0), + o.stride(1), + GROUPS=heads // kv_heads, + D=dim, + BLOCK_M=BLOCK_Q, + BLOCK_N=BLOCK_KV, + ) + return out diff --git a/src/orbitquant/kernels/triton_cuda.py b/src/orbitquant/kernels/triton_cuda.py index 3684558..a347355 100644 --- a/src/orbitquant/kernels/triton_cuda.py +++ b/src/orbitquant/kernels/triton_cuda.py @@ -139,7 +139,7 @@ def _weight_quantization_constants( def _row_norm_kernel( input_ptr, norms_ptr, - rows: tl.constexpr, + rows, dim: tl.constexpr, eps: tl.constexpr, block_size: tl.constexpr, @@ -160,7 +160,7 @@ def _permute_sign_normalize_activation_kernel( permutation_ptr, signs_ptr, work_ptr, - total: tl.constexpr, + total, dim: tl.constexpr, eps: tl.constexpr, block_size: tl.constexpr, @@ -184,7 +184,7 @@ def _quantize_activation_work_rescale_kernel( centroids_ptr, boundaries_ptr, output_ptr, - total: tl.constexpr, + total, dim: tl.constexpr, levels: tl.constexpr, inv_sqrt_block: tl.constexpr, @@ -226,7 +226,7 @@ def _fused_rpbh_quantize_activation_kernel( centroids_ptr, boundaries_ptr, output_ptr, - rows: tl.constexpr, + rows, dim: tl.constexpr, num_blocks: tl.constexpr, orbit_block_size: tl.constexpr, @@ -294,7 +294,7 @@ def _fused_rpbh_quantize_activation_pack_w4_kernel( signs_ptr, boundaries_ptr, output_ptr, - rows: tl.constexpr, + rows, dim: tl.constexpr, num_blocks: tl.constexpr, orbit_block_size: tl.constexpr, @@ -357,6 +357,74 @@ def _fused_rpbh_quantize_activation_pack_w4_kernel( tl.store(output_ptr + output_offsets, packed.to(tl.uint8), mask=row < rows) +@triton.jit +def _fused_rpbh_quantize_activation_int8_kernel( + input_ptr, + norms_ptr, + permutation_ptr, + signs_ptr, + boundaries_ptr, + codes_ptr, + output_ptr, + rows, + dim: tl.constexpr, + num_blocks: tl.constexpr, + orbit_block_size: tl.constexpr, + fwht_stages: tl.constexpr, + levels: tl.constexpr, + eps: tl.constexpr, + inv_sqrt_block: tl.constexpr, +): + row_block = tl.program_id(0) + row = row_block // num_blocks + orbit_block = row_block % num_blocks + local_cols = tl.arange(0, orbit_block_size) + cols = orbit_block * orbit_block_size + local_cols + mask = (row < rows) & (cols < dim) + source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int32) + signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32) + norm = tl.load(norms_ptr + row, mask=row < rows, other=0.0).to(tl.float32) + values = tl.load( + input_ptr + row * dim + source_cols, + mask=mask, + other=0.0, + ).to(tl.float32) + values = values * signs / (norm + eps) + + if fwht_stages > 0: + values = _fwht_local_stage(values, orbit_block_size, 1) + if fwht_stages > 1: + values = _fwht_local_stage(values, orbit_block_size, 2) + if fwht_stages > 2: + values = _fwht_local_stage(values, orbit_block_size, 4) + if fwht_stages > 3: + values = _fwht_local_stage(values, orbit_block_size, 8) + if fwht_stages > 4: + values = _fwht_local_stage(values, orbit_block_size, 16) + if fwht_stages > 5: + values = _fwht_local_stage(values, orbit_block_size, 32) + if fwht_stages > 6: + values = _fwht_local_stage(values, orbit_block_size, 64) + if fwht_stages > 7: + values = _fwht_local_stage(values, orbit_block_size, 128) + if fwht_stages > 8: + values = _fwht_local_stage(values, orbit_block_size, 256) + if fwht_stages > 9: + values = _fwht_local_stage(values, orbit_block_size, 512) + if fwht_stages > 10: + values = _fwht_local_stage(values, orbit_block_size, 1024) + if fwht_stages > 11: + values = _fwht_local_stage(values, orbit_block_size, 2048) + + values *= inv_sqrt_block + indices = tl.zeros((orbit_block_size,), dtype=tl.int32) + for index in tl.static_range(0, levels - 1): + boundary = tl.load(boundaries_ptr + index) + indices += (values > boundary).to(tl.int32) + quantized = tl.load(codes_ptr + indices) + tl.store(output_ptr + row * dim + cols, quantized.to(tl.int8), mask=mask) + + @triton.jit def _fused_rpbh_fwht_1024_chunk_kernel( input_ptr, @@ -364,7 +432,7 @@ def _fused_rpbh_fwht_1024_chunk_kernel( permutation_ptr, signs_ptr, work_ptr, - rows: tl.constexpr, + rows, dim: tl.constexpr, num_blocks: tl.constexpr, orbit_block_size: tl.constexpr, @@ -408,7 +476,7 @@ def _codebook_rescale_kernel( centroids_ptr, boundaries_ptr, output_ptr, - total: tl.constexpr, + total, dim: tl.constexpr, block_size: tl.constexpr, levels: tl.constexpr, @@ -431,7 +499,7 @@ def _quantize_activation_work_pack_w4_kernel( work_ptr, boundaries_ptr, output_ptr, - packed_total: tl.constexpr, + packed_total, levels: tl.constexpr, inv_sqrt_block: tl.constexpr, block_size: tl.constexpr, @@ -456,7 +524,7 @@ def _quantize_activation_work_pack_w4_kernel( @triton.jit def _fwht_stage_kernel( work_ptr, - total_pairs: tl.constexpr, + total_pairs, in_features: tl.constexpr, num_blocks: tl.constexpr, orbit_block_size: tl.constexpr, @@ -485,7 +553,7 @@ def _fwht_stage_kernel( @triton.jit def _fwht_two_stage_kernel( work_ptr, - total_quads: tl.constexpr, + total_quads, in_features: tl.constexpr, num_blocks: tl.constexpr, orbit_block_size: tl.constexpr, @@ -526,7 +594,7 @@ def _dequantize_packed_weight_kernel( row_norms_ptr, centroids_ptr, output_ptr, - total: tl.constexpr, + total, in_features: tl.constexpr, bits: tl.constexpr, block_size: tl.constexpr, @@ -555,7 +623,7 @@ def _matmul_packed_weight_kernel( centroids_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, bits: tl.constexpr, @@ -619,7 +687,7 @@ def _matmul_packed_weight_w4_kernel( centroids_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, has_bias: tl.constexpr, @@ -686,7 +754,7 @@ def _matmul_packed_adaln_int4_group_kernel( scales_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, padded_in_features: tl.constexpr, @@ -757,7 +825,7 @@ def _matmul_packed_adaln_int4_kernel( scales_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, padded_in_features: tl.constexpr, @@ -823,7 +891,7 @@ def _decode_packed_w4_codes_kernel( packed_ptr, codes_ptr, output_ptr, - total: tl.constexpr, + total, block_size: tl.constexpr, ): offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) @@ -843,7 +911,7 @@ def _scale_int32_matmul_output_chunk_kernel( bias_ptr, output_ptr, output_col_start, - rows: tl.constexpr, + rows, chunk_out_features: tl.constexpr, out_features: tl.constexpr, combined_scale: tl.constexpr, @@ -929,7 +997,7 @@ def _permute_sign_weight_kernel( permutation_ptr, signs_ptr, work_ptr, - total: tl.constexpr, + total, in_features: tl.constexpr, block_size: tl.constexpr, ): @@ -950,7 +1018,7 @@ def _permute_sign_weight_kernel( @triton.jit def _fwht_stage_weight_kernel( work_ptr, - total_pairs: tl.constexpr, + total_pairs, in_features: tl.constexpr, num_blocks: tl.constexpr, orbit_block_size: tl.constexpr, @@ -982,7 +1050,7 @@ def _quantize_rotated_weight_indices_kernel( row_norms_ptr, boundaries_ptr, indices_ptr, - total: tl.constexpr, + total, in_features: tl.constexpr, levels: tl.constexpr, eps: tl.constexpr, @@ -1140,7 +1208,7 @@ def _dequantize_adaln_weight_kernel( packed_ptr, scales_ptr, output_ptr, - total: tl.constexpr, + total, in_features: tl.constexpr, padded_in_features: tl.constexpr, num_groups: tl.constexpr, @@ -1447,6 +1515,99 @@ def quantize_activations_packed_w4_with_triton( return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape(original_shape[:-1]) +def quantize_activations_int8_with_triton( + x: torch.Tensor, + *, + rotation: RPBHRotation, + codebook: LloydMaxCodebook, + activation_codes: torch.Tensor, + eps: float, + constant_tensors: dict[str, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize RPBH activations directly to their INT8 centroid surrogates.""" + _require_triton_tensor(x, "INT8 activation quantization") + if codebook.bits != 4: + raise ValueError("INT8 activation quantization requires a 4-bit codebook") + if x.shape[-1] != rotation.dim: + raise ValueError(f"expected last dimension {rotation.dim}, got {x.shape[-1]}") + if activation_codes.numel() != 16: + raise ValueError("INT8 activation quantization requires 16 surrogate codes") + + triton, _ = _load_triton() + original_shape = x.shape + dim = int(rotation.dim) + input_contiguous = x.contiguous().reshape(-1, dim) + rows = input_contiguous.shape[0] + output = torch.empty((rows, dim), device=x.device, dtype=torch.int8) + norms = torch.empty(rows, device=x.device, dtype=torch.float32) + if output.numel() == 0: + return output.reshape(original_shape), norms.reshape(original_shape[:-1]) + + if rotation.block_size > 4096: + packed, norms = quantize_activations_packed_w4_with_triton( + x, + rotation=rotation, + codebook=codebook, + eps=eps, + constant_tensors=constant_tensors, + ) + _decode_packed_w4_codes_kernel[(triton.cdiv(output.numel(), 256),)]( + packed, + activation_codes.to(device=x.device, dtype=torch.int8).contiguous(), + output, + total=output.numel(), + block_size=256, + ) + return output.reshape(original_shape), norms + + constants = {} if constant_tensors is None else constant_tensors + permutation = ( + constants.get("permutation", rotation.permutation) + .to(device=x.device, dtype=torch.int32) + .contiguous() + ) + signs = ( + constants.get("signs", rotation.signs).to(device=x.device, dtype=torch.int8).contiguous() + ) + boundaries = ( + constants.get("boundaries", codebook.boundaries) + .to(device=x.device, dtype=torch.float32) + .contiguous() + ) + codes = activation_codes.to(device=x.device, dtype=torch.int8).contiguous() + + _row_norm_kernel[(rows,)]( + input_contiguous, + norms, + rows=rows, + dim=dim, + eps=float(eps), + block_size=triton.next_power_of_2(dim), + num_warps=8, + ) + orbit_block_size = int(rotation.block_size) + num_blocks = int(rotation.num_blocks) + _fused_rpbh_quantize_activation_int8_kernel[(rows * num_blocks,)]( + input_contiguous, + norms, + permutation, + signs, + boundaries, + codes, + output, + rows=rows, + dim=dim, + num_blocks=num_blocks, + orbit_block_size=orbit_block_size, + fwht_stages=orbit_block_size.bit_length() - 1, + levels=codebook.centroids.numel(), + eps=float(eps), + inv_sqrt_block=float(rotation.normalization), + num_warps=8 if orbit_block_size >= 512 else 4, + ) + return output.reshape(original_shape), norms.reshape(original_shape[:-1]) + + def quantize_adaln_weight_with_triton( weight: torch.Tensor, *, @@ -1867,7 +2028,7 @@ def _matmul_packed_w4a4_fused_int8_kernel( weight_codes_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, combined_scale: tl.constexpr, @@ -1942,7 +2103,7 @@ def _matmul_int8act_packed_lowbit_fused_kernel( weight_codes_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, combined_scale: tl.constexpr, @@ -2027,7 +2188,7 @@ def _matmul_packed_w2a4_fused_int8_kernel( weight_codes_ptr, bias_ptr, output_ptr, - rows: tl.constexpr, + rows, out_features: tl.constexpr, in_features: tl.constexpr, combined_scale: tl.constexpr, @@ -2474,8 +2635,13 @@ def matmul_packed_w4a4_with_int_mm( chunk_out_features: int | None = None, activations_are_int8: bool = False, decoded_weight: torch.Tensor | None = None, + fused_epilogue: bool | None = None, ) -> torch.Tensor: - """Decode active W4/A4 tiles to INT8 and use PyTorch's CUTLASS GEMM. + """Decode active W4/A4 tiles to INT8 and run an INT8 GEMM. + + ``fused_epilogue`` (default on CUDA) runs the Triton INT8 GEMM with the scale epilogue in + registers; otherwise PyTorch's cuBLASLt ``_int_mm`` writes int32 and a second kernel + scales it. Both produce identical results. ``decoded_weight`` optionally supplies a persistent INT8 surrogate weight (``decode_packed_w4_weight_to_int8``) so the per-forward decode is skipped. @@ -2554,12 +2720,14 @@ def matmul_packed_w4a4_with_int_mm( block_size=decode_block_size, ) + if fused_epilogue is None: + fused_epilogue = device.type == "cuda" and not getattr(torch.version, "hip", None) # cuBLASLt INT8 GEMM rejects some row counts (every m <= 16 and, on some # torch/CUDA stacks, shape-dependent m values that are not multiples of # 32), so run the GEMM on zero-padded rows and let the epilogue write only # the real ones. gemm_rows = max(32, -(-rows // 32) * 32) - if gemm_rows != rows: + if gemm_rows != rows and not fused_epilogue: padded_activations = torch.zeros( (gemm_rows, in_features), device=device, dtype=torch.int8 ) @@ -2611,6 +2779,22 @@ def matmul_packed_w4a4_with_int_mm( total=weight_total, block_size=decode_block_size, ) + if fused_epilogue: + # Same exact int32 sums and the same epilogue order as the int_mm path below. + from orbitquant.kernels.triton_int8_gemm import matmul_int8_scaled_with_triton + + matmul_int8_scaled_with_triton( + activation_values, + weight_storage, + token_norm_values, + row_norm_values.narrow(0, output_col_start, current_chunk), + alpha=combined_scale, + bias=bias_values.narrow(0, output_col_start, current_chunk) if has_bias else None, + out=output, + col_offset=output_col_start, + ) + del weight_storage + continue accumulator = torch._int_mm(gemm_activations, weight_storage.t()) del weight_storage epilogue_grid = ( diff --git a/src/orbitquant/kernels/triton_int8_gemm.py b/src/orbitquant/kernels/triton_int8_gemm.py new file mode 100644 index 0000000..559546d --- /dev/null +++ b/src/orbitquant/kernels/triton_int8_gemm.py @@ -0,0 +1,268 @@ +"""INT8 x INT8 GEMM with the scale epilogue fused (Triton). + +The large-row W4A4 path used ``torch._int_mm`` (int32 result in global memory) followed by a +separate scale kernel. On GeForce cards the int32 round trip costs as much as a quarter of the +GEMM; this kernel keeps the accumulator in registers and writes the scaled result in the output +dtype. Integer accumulation is exact, and the epilogue multiplies in the same order as the +kernels it replaces, so the plain epilogue is bit-identical to them: + +* ``scale_mode=0`` (OrbitQuant W4A4 surrogates): ``((acc * a_scale) * b_scale) * alpha + bias`` +* ``scale_mode=1`` (per-row INT8, ``Int8RowLinear``): ``acc * (a_scale * b_scale) + bias`` + +Optional epilogues fuse the elementwise work that follows a projection in DiT blocks. Each +rounds to bf16 exactly where the eager ops it replaces do: + +* ``EPILOGUE_SWIGLU``: rows of ``b`` interleave gate and up in chunks of ``BLOCK_N // 2``; + the kernel writes ``bf16(silu(bf16(gate))) * bf16(up)`` (half as many columns). +* ``EPILOGUE_RESIDUAL_GATE``: ``out = res + bf16(gate[n] * bf16(y))`` (in place allowed). +* ``EPILOGUE_SIGMOID_TAIL``: columns ``>= sig_from`` store ``sigmoid(bf16(y))``. +""" + +from __future__ import annotations + +import inspect + +import torch + +from orbitquant.kernels.triton_cuda import _load_triton + +triton, tl = _load_triton() + +from triton.language.extra import libdevice # noqa: E402 + +EPILOGUE_PLAIN = 0 +EPILOGUE_SWIGLU = 1 +EPILOGUE_RESIDUAL_GATE = 2 +EPILOGUE_SIGMOID_TAIL = 3 + +# The SwiGLU epilogue pairs gate/up halves inside one N tile; weights packed for it interleave +# rows in chunks of this size. +SWIGLU_INTERLEAVE = 128 + +_SHAPES = [ + (128, 128, 64, 4, 4), + (128, 128, 128, 4, 3), + (128, 128, 128, 8, 3), + (128, 256, 64, 8, 3), + (256, 128, 64, 8, 3), + (64, 256, 128, 4, 3), + (128, 64, 128, 4, 4), + (64, 128, 128, 4, 4), +] + + +def autotune(configs, key, restore_value=None): + """``triton.autotune`` that stores its choices in the Triton cache, so a restarted process + does not benchmark every configuration again (when the installed Triton supports it).""" + extra = {} + if "cache_results" in inspect.signature(triton.autotune).parameters: + extra["cache_results"] = True + if restore_value: + extra["restore_value"] = restore_value + return triton.autotune(configs=configs, key=key, **extra) + + +def _configs(block_n: int | None = None): + return [ + triton.Config( + {"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8}, + num_warps=warps, + num_stages=stages, + ) + for bm, bn, bk, warps, stages in _SHAPES + if block_n is None or bn == block_n + ] + + +@triton.jit +def _bf16(x): + return x.to(tl.bfloat16).to(tl.float32) + + +@triton.jit +def _int8_scaled_gemm_body( + a_ptr, + b_ptr, + out_ptr, + a_scale_ptr, + b_scale_ptr, + bias_ptr, + res_ptr, + gate_ptr, + M, + N, + K, + stride_om, + col_offset, + M_BUCKET, + alpha, + sig_from, + HAS_BIAS: tl.constexpr, + SCALE_MODE: tl.constexpr, + EPILOGUE: tl.constexpr, + EVEN_K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, +): + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + mask_m = offs_m < M + mask_n = offs_n < N + a_ptrs = a_ptr + offs_m[:, None].to(tl.int64) * K + offs_k[None, :] + b_ptrs = b_ptr + offs_n[None, :].to(tl.int64) * K + offs_k[:, None] + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + if EVEN_K: + a = tl.load(a_ptrs, mask=mask_m[:, None], other=0) + b = tl.load(b_ptrs, mask=mask_n[None, :], other=0) + else: + k_mask = offs_k + k * BLOCK_K < K + a = tl.load(a_ptrs, mask=mask_m[:, None] & k_mask[None, :], other=0) + b = tl.load(b_ptrs, mask=mask_n[None, :] & k_mask[:, None], other=0) + acc = tl.dot(a, b, acc, out_dtype=tl.int32) + a_ptrs += BLOCK_K + b_ptrs += BLOCK_K + + a_s = tl.load(a_scale_ptr + offs_m, mask=mask_m, other=0.0).to(tl.float32) + b_s = tl.load(b_scale_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + if SCALE_MODE == 1: + y = acc.to(tl.float32) * (a_s[:, None] * b_s[None, :]) + else: + y = acc.to(tl.float32) * a_s[:, None] * b_s[None, :] * alpha + if HAS_BIAS: + y += tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)[None, :] + rows = offs_m[:, None].to(tl.int64) * stride_om + if EPILOGUE == 1: + pairs = tl.permute(tl.reshape(_bf16(y), (BLOCK_M, 2, BLOCK_N // 2)), (0, 2, 1)) + g, u = tl.split(pairs) + h = _bf16(libdevice.div_rn(g, 1.0 + libdevice.exp(-g))) * u + offs_h = col_offset + pid_n * (BLOCK_N // 2) + tl.arange(0, BLOCK_N // 2) + tl.store( + out_ptr + rows + offs_h[None, :], + h.to(out_ptr.dtype.element_ty), + mask=mask_m[:, None] & (offs_h[None, :] < col_offset + N // 2), + ) + elif EPILOGUE == 2: + cols = (col_offset + offs_n)[None, :] + mask = mask_m[:, None] & mask_n[None, :] + gate = tl.load(gate_ptr + col_offset + offs_n, mask=mask_n, other=0.0).to(tl.float32) + res = tl.load(res_ptr + rows + cols, mask=mask, other=0.0).to(tl.float32) + out = res + _bf16(gate[None, :] * _bf16(y)) + tl.store(out_ptr + rows + cols, out.to(out_ptr.dtype.element_ty), mask=mask) + elif EPILOGUE == 3: + y = _bf16(y) + if pid_n * BLOCK_N >= sig_from: + y = libdevice.div_rn(1.0, 1.0 + libdevice.exp(-y)) + tl.store( + out_ptr + rows + (col_offset + offs_n)[None, :], + y.to(out_ptr.dtype.element_ty), + mask=mask_m[:, None] & mask_n[None, :], + ) + else: + tl.store( + out_ptr + rows + (col_offset + offs_n)[None, :], + y.to(out_ptr.dtype.element_ty), + mask=mask_m[:, None] & mask_n[None, :], + ) + + +_KEY = ["M_BUCKET", "N", "K", "SCALE_MODE", "HAS_BIAS", "EPILOGUE"] +# The residual-gate epilogue usually updates its output in place: every benchmark run of the +# autotuner would add the projection to the residual again. +_int8_scaled_gemm = autotune(_configs(), _KEY, restore_value=["out_ptr"])(_int8_scaled_gemm_body) +_int8_scaled_gemm_swiglu = autotune(_configs(block_n=2 * SWIGLU_INTERLEAVE), _KEY)( + _int8_scaled_gemm_body +) + + +def _row_bucket(rows: int) -> int: + return 1 << max(5, (rows - 1).bit_length()) + + +def matmul_int8_scaled_with_triton( + a: torch.Tensor, + b: torch.Tensor, + a_scale: torch.Tensor, + b_scale: torch.Tensor, + *, + alpha: float = 1.0, + bias: torch.Tensor | None = None, + out: torch.Tensor | None = None, + col_offset: int = 0, + output_dtype: torch.dtype = torch.bfloat16, + scale_mode: int = 0, + epilogue: int = EPILOGUE_PLAIN, + residual: torch.Tensor | None = None, + gate: torch.Tensor | None = None, + sig_from: int = 0, +) -> torch.Tensor: + """Scaled ``a @ b.T`` for INT8 ``a [M, K]`` and ``b [N, K]``. + + ``out`` may be a wider row-major buffer; the result lands in columns + ``[col_offset, col_offset + N)`` (``N // 2`` for SwiGLU) so chunked weight decodes write in + place. The residual-gate epilogue reads ``residual`` with the same layout as ``out`` and may + alias it. + """ + if a.dtype != torch.int8 or b.dtype != torch.int8: + raise ValueError("matmul_int8_scaled_with_triton expects INT8 operands") + rows, k = a.shape + n = b.shape[0] + if b.shape[1] != k: + raise ValueError(f"inner dimensions differ: {k} vs {b.shape[1]}") + if epilogue == EPILOGUE_SWIGLU and n % (2 * SWIGLU_INTERLEAVE): + raise ValueError(f"SwiGLU epilogue needs N divisible by {2 * SWIGLU_INTERLEAVE}") + if epilogue == EPILOGUE_RESIDUAL_GATE and (residual is None or gate is None): + raise ValueError("residual-gate epilogue needs residual and gate") + out_cols = n // 2 if epilogue == EPILOGUE_SWIGLU else n + if out is None: + out = torch.empty((rows, out_cols), device=a.device, dtype=output_dtype) + if rows == 0: + return out + if a.stride(1) != 1 or a.stride(0) != k or b.stride(1) != 1 or b.stride(0) != k: + a = a.contiguous() + b = b.contiguous() + if residual is not None and residual.stride() != out.stride(): + raise ValueError("residual must share the output layout") + a_scale = a_scale.to(torch.float32).contiguous() + b_scale = b_scale.to(torch.float32).contiguous() + kernel = _int8_scaled_gemm_swiglu if epilogue == EPILOGUE_SWIGLU else _int8_scaled_gemm + + def grid(meta): + return (triton.cdiv(rows, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),) + + kernel[grid]( + a, + b, + out, + a_scale, + b_scale, + bias if bias is not None else a_scale, + residual if residual is not None else out, + gate if gate is not None else a_scale, + rows, + n, + k, + out.stride(0), + int(col_offset), + _row_bucket(rows), + float(alpha), + int(sig_from), + HAS_BIAS=bias is not None, + SCALE_MODE=int(scale_mode), + EPILOGUE=int(epilogue), + EVEN_K=k % 128 == 0, + ) + return out diff --git a/src/orbitquant/layers.py b/src/orbitquant/layers.py index 9f3ecb6..afc808b 100644 --- a/src/orbitquant/layers.py +++ b/src/orbitquant/layers.py @@ -2,7 +2,9 @@ import logging import os +import threading import weakref +from collections.abc import Callable from itertools import count import torch @@ -41,6 +43,48 @@ _COMPILE_HANDLE_IDS = count(1) +_SHARED_ACTIVATION_CACHE = threading.local() + + +def _clear_shared_activation_cache() -> None: + """Drop the current thread's last prepared activation.""" + + _SHARED_ACTIVATION_CACHE.entry = None + + +def _get_or_create_shared_activation_quantization( + x: torch.Tensor, + signature: tuple[object, ...], + factory: Callable[[], tuple[torch.Tensor, torch.Tensor]], +) -> tuple[tuple[torch.Tensor, torch.Tensor], bool]: + """Reuse an unchanged activation across consecutive quantized projections. + + Attention Q/K/V (and gate) and SwiGLU gate/up projections receive the same tensor object + and use the same RPBH/codebook transform. Keeping only the most recent prepared activation + makes that common pattern cheap without an unbounded cache or pointer-based aliasing. + """ + + if torch.is_grad_enabled(): + _clear_shared_activation_cache() + return factory(), False + try: + version = x._version + except RuntimeError: + # Inference tensors hide their version counter, so an in-place mutation between two + # projections could not be detected. + _clear_shared_activation_cache() + return factory(), False + metadata = (tuple(x.shape), tuple(x.stride()), x.storage_offset(), x.dtype, x.device, version) + entry = getattr(_SHARED_ACTIVATION_CACHE, "entry", None) + if entry is not None: + input_ref, cached_signature, cached_metadata, prepared = entry + if input_ref() is x and cached_signature == signature and cached_metadata == metadata: + return prepared, True + prepared = factory() + _SHARED_ACTIVATION_CACHE.entry = (weakref.ref(x), signature, metadata, prepared) + return prepared, False + + def _compile_registry_lookup(handle: int): module = _COMPILE_REGISTRY.get(handle) if module is None: @@ -419,6 +463,7 @@ def __init__( self.last_native_w4a4_enabled = False self.last_native_w4_activation_enabled = False self.last_native_int8_activation_enabled = False + self.last_activation_cache_hit = False self._derived_constants_valid = True @classmethod @@ -1106,6 +1151,7 @@ def _forward_impl(self, x: torch.Tensor) -> torch.Tensor: self.last_native_w4a4_enabled = False self.last_native_w4_activation_enabled = False self.last_native_int8_activation_enabled = False + self.last_activation_cache_hit = False if runtime_mode == "triton_packed_matmul": self._validate_triton_packed_matmul_input(x) @@ -1211,19 +1257,48 @@ def _forward_impl(self, x: torch.Tensor) -> torch.Tensor: quantize_activations_int8_with_native_kernel, ) - int8_x, token_norms = quantize_activations_int8_with_native_kernel( - x, - activation_constants["permutation"], - activation_constants["signs"], - activation_constants["boundaries"], - activation_codes, - eps=self.activation_eps, - inv_sqrt_block=self.rotation.normalization, - threads=512 if capability == (8, 9) else 256, + threads = 512 if capability == (8, 9) else 256 + (int8_x, token_norms), self.last_activation_cache_hit = ( + _get_or_create_shared_activation_quantization( + x, + ("native_cuda_int8_surrogate", id(self.rotation), + id(self.activation_codebook), self.activation_eps, threads), + lambda: quantize_activations_int8_with_native_kernel( + x, + activation_constants["permutation"], + activation_constants["signs"], + activation_constants["boundaries"], + activation_codes, + eps=self.activation_eps, + inv_sqrt_block=self.rotation.normalization, + threads=threads, + ), + ) ) self.last_activation_kernel_backend = "native_cuda_int8_surrogate" self.last_native_w4_activation_enabled = True self.last_native_int8_activation_enabled = True + elif use_cutlass_tn and self.rotation.block_size <= 4096: + from orbitquant.kernels.triton_cuda import quantize_activations_int8_with_triton + + (int8_x, token_norms), self.last_activation_cache_hit = ( + _get_or_create_shared_activation_quantization( + x, + ("triton_cuda_int8_surrogate", id(self.rotation), + id(self.activation_codebook), self.activation_eps), + lambda: quantize_activations_int8_with_triton( + x, + rotation=self.rotation, + codebook=self.activation_codebook, + activation_codes=activation_codes, + eps=self.activation_eps, + constant_tensors=activation_constants, + ), + ) + ) + self.last_activation_kernel_backend = "triton_cuda_int8_surrogate" + self.last_native_w4_activation_enabled = True + self.last_native_int8_activation_enabled = True elif self._native_w4_activation_available(x): from orbitquant.kernels.native_packed_matmul import ( quantize_activations_packed_w4_with_native_kernel, diff --git a/src/orbitquant/runtime/__init__.py b/src/orbitquant/runtime/__init__.py new file mode 100644 index 0000000..57e9cb1 --- /dev/null +++ b/src/orbitquant/runtime/__init__.py @@ -0,0 +1 @@ +"""Model-specific fused inference runtimes built on OrbitQuant kernels.""" diff --git a/src/orbitquant/runtime/krea2.py b/src/orbitquant/runtime/krea2.py new file mode 100644 index 0000000..d74ff85 --- /dev/null +++ b/src/orbitquant/runtime/krea2.py @@ -0,0 +1,618 @@ +"""Fused inference runtime for Krea 2 with OrbitQuant W4A4 weights. + +``install(transformer)`` regroups the projection weights of every ``Krea2TransformerBlock`` +(Q|K|V|gate, out, gate/up interleaved for the SwiGLU epilogue, down) into a ``fused_blocks`` +submodule and replaces the original projection modules with placeholders. ``Krea2FastRunner`` +then runs a ``diffusers.Krea2Pipeline`` with: + +* text encoding without the 512-token padding (right-padded to a bucket; the encoder is + causal, so the valid outputs do not change) and without the unused vision tower, +* a compacted DiT sequence (padded text lanes dropped) so attention needs no mask and runs + natively grouped-query in Flash or cuDNN SDPA, or in the INT8 Q.K^T / FP16 P.V kernel + (``attention="int8"``, ``orbitquant.kernels.triton_attention``), +* text fusion, text projection and rotary tables computed once per image, +* per block: one activation quantization per shared input fused with its RMSNorm/modulation + or gate product, one grouped INT8 GEMM for Q|K|V|gate with the sigmoid gate in its + epilogue, a Q/K RMSNorm + RoPE kernel, and the SwiGLU and the gated residual updates fused + into GEMM epilogues. The W4 weights of each group are decoded to INT8 right before their + GEMM; ``prefetch=True`` runs the next decode on a side stream instead. + +``down="int8"`` stores the down projections as per-row INT8 (W8A8, ``Int8RowLinear`` +numerics) instead of the source BF16, which is 3x faster on GeForce cards and halves their +memory; ``down="bf16"`` keeps them as shipped. + +Building the fused buffers copies the block weights into new host memory. ``save_fused`` +writes them to a safetensors file once; ``install(..., fused_path=...)`` then maps that file +instead, so the fused weights stay reclaimable page cache like the rest of the checkpoint. +""" + +from __future__ import annotations + +import os +import types +from collections.abc import Callable + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors import safe_open +from safetensors.torch import save_file +from torch import nn +from torch.nn.attention import SDPBackend, sdpa_kernel + +from orbitquant.int8_head import quantize_int8_rows +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.kernels.triton_cuda import fit_int8_centroid_surrogate +from orbitquant.layers import OrbitQuantLinear + +__all__ = ["FusedKrea2Blocks", "Krea2FastRunner", "install", "save_fused"] + +FUSED_FORMAT = "orbitquant-krea2-fused" +FUSED_FORMAT_VERSION = "1" + +PROMPT_PREFIX = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|im_end|>\n" + "<|im_start|>user\n" +) +PROMPT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n" +PROMPT_PREFIX_TOKENS = 34 +PROMPT_SUFFIX_TOKENS = 5 +MAX_SEQUENCE_LENGTH = 512 +TEXT_BUCKET = 64 +DISTILLED_SHIFT_MU = 1.15 + +# INT8 Q/K loses the small channels of heads whose Q/K RMSNorm scales one channel far above +# the rest (Krea 2 block 0 scales a channel 36x the median); those blocks keep BF16 attention. +INT8_ATTENTION_MAX_NORM_SPREAD = 3.0 + +_ATTENTION_BACKENDS = { + "flash": SDPBackend.FLASH_ATTENTION, + "cudnn": SDPBackend.CUDNN_ATTENTION, + "int8": None, +} + + +def _codes_and_alpha(layer: OrbitQuantLinear) -> tuple[torch.Tensor, float]: + _, 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) + + +class FusedGroup(nn.Module): + """Row-concatenated packed W4 weights of projections that read the same activation.""" + + def __init__( + self, + packed: torch.Tensor, + row_norms: torch.Tensor, + codes: torch.Tensor, + alpha: float, + in_features: int, + ): + super().__init__() + self.in_features = int(in_features) + self.alpha = float(alpha) + self.register_buffer("packed", packed) + self.register_buffer("row_norms", row_norms) + self.register_buffer("codes", codes.to(torch.int8)) + self.out_features = int(packed.shape[0]) + + @classmethod + def from_layers( + cls, layers: list[OrbitQuantLinear], interleave: int | None = None + ) -> FusedGroup: + codes, alpha = _group_constants(layers) + packed = [layer.packed_weight_indices.reshape(layer.out_features, -1) for layer in layers] + norms = [layer.row_norms.float().reshape(-1) for layer in layers] + if interleave is not None: + (gate_p, up_p), (gate_n, up_n) = packed, norms + packed, norms = [], [] + for start in range(0, gate_p.shape[0], interleave): + end = start + interleave + packed += [gate_p[start:end], up_p[start:end]] + norms += [gate_n[start:end], up_n[start:end]] + return cls( + torch.cat(packed, dim=0).contiguous(), + torch.cat(norms, dim=0).contiguous(), + codes, + alpha, + layers[0].in_features, + ) + + @classmethod + def from_stored( + cls, layers: list[OrbitQuantLinear], packed: torch.Tensor, row_norms: torch.Tensor + ) -> FusedGroup: + codes, alpha = _group_constants(layers) + rows = sum(layer.out_features for layer in layers) + expected = (rows, layers[0].in_features // 2) + if tuple(packed.shape) != expected or tuple(row_norms.shape) != (rows,): + raise ValueError( + f"stored fused group {tuple(packed.shape)} does not match the layers {expected}" + ) + return cls(packed, row_norms, codes, alpha, layers[0].in_features) + + def decode(self, buffer: torch.Tensor) -> torch.Tensor: + view = buffer[: self.out_features * self.in_features].view( + self.out_features, self.in_features + ) + return act.decode_w4_to_int8_with_triton( + self.packed, self.codes, self.out_features, self.in_features, out=view + ) + + +def _group_constants(layers: list[OrbitQuantLinear]) -> tuple[torch.Tensor, float]: + codes, alpha = _codes_and_alpha(layers[0]) + for layer in layers: + other_codes, other_alpha = _codes_and_alpha(layer) + if ( + not isinstance(layer, OrbitQuantLinear) + or layer.weight_bits != 4 + or layer.bias is not None + or other_alpha != alpha + or not torch.equal(other_codes, codes) + ): + raise ValueError("fused groups need bias-free W4 OrbitQuant layers sharing codebooks") + return codes, alpha + + +class Int8Rows(nn.Module): + """Per-row INT8 weights of a W8A8 projection (``Int8RowLinear`` numerics).""" + + def __init__(self, q: torch.Tensor, scales: torch.Tensor): + super().__init__() + self.register_buffer("q", q) + self.register_buffer("scales", scales) + + @classmethod + def from_weight(cls, weight: torch.Tensor) -> Int8Rows: + return cls(*quantize_int8_rows(weight)) + + +class Bf16Weight(nn.Module): + def __init__(self, weight: torch.Tensor): + super().__init__() + self.register_buffer("weight", weight.detach()) + + +_GROUP_LAYERS = { + "qkvg": lambda attn, ff: [attn.to_q, attn.to_k, attn.to_v, attn.to_gate], + "out": lambda attn, ff: [attn.to_out[0]], + "gateup": lambda attn, ff: [ff.gate, ff.up], +} + + +class FusedKrea2Block(nn.Module): + def __init__(self, block, down: str, stored: Callable[[str], torch.Tensor] | None = None): + super().__init__() + attn, ff = block.attn, block.ff + for name, layers_of in _GROUP_LAYERS.items(): + layers = layers_of(attn, ff) + interleave = gemm.SWIGLU_INTERLEAVE if name == "gateup" else None + if stored is None: + group = FusedGroup.from_layers(layers, interleave=interleave) + else: + group = FusedGroup.from_stored( + layers, stored(f"{name}.packed"), stored(f"{name}.row_norms") + ) + setattr(self, name, group) + weight = ff.down.weight + if down == "int8": + if stored is None: + self.down = Int8Rows.from_weight(weight) + else: + self.down = Int8Rows(stored("down.q"), stored("down.scales")) + if self.down.q.shape != weight.shape: + raise ValueError("stored down projection does not match the layer") + elif down == "bf16": + self.down = Bf16Weight(weight if stored is None else stored("down.weight")) + else: + raise ValueError(f"unknown down projection mode {down!r}") + self.q_width = attn.num_heads * attn.head_dim + self.kv_width = attn.num_kv_heads * attn.head_dim + scales = torch.cat([attn.norm_q.weight, attn.norm_k.weight]).float().add(1.0).abs() + self.int8_attention = bool(scales.max() <= INT8_ATTENTION_MAX_NORM_SPREAD * scales.median()) + + +class FusedKrea2Blocks(nn.Module): + """Fused weights and runtime state for all transformer blocks of one Krea 2 DiT.""" + + def __init__( + self, + transformer, + *, + down: str = "int8", + attention: str = "flash", + prefetch: bool = False, + stored: Callable[[str], torch.Tensor] | None = None, + ): + super().__init__() + if attention not in _ATTENTION_BACKENDS: + raise ValueError(f"unknown attention {attention!r}") + blocks = transformer.transformer_blocks + 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 + ) + self.down_mode = down + self.attention = _ATTENTION_BACKENDS[attention] + self.prefetch = prefetch + self.rms_eps = float(blocks[0].norm1.eps) + self.fused = nn.ModuleList( + [ + FusedKrea2Block( + block, + down, + None if stored is None else (lambda key, i=index: stored(f"blocks.{i}.{key}")), + ) + for index, block in enumerate(blocks) + ] + ) + self.max_decode = ( + max(max(f.qkvg.out_features, f.gateup.out_features) for f in self.fused) + * first.in_features + ) + self.order = [ + (i, name) for i in range(len(self.fused)) for name in ("qkvg", "out", "gateup") + ] + self._reset_runtime() + + def _reset_runtime(self): + self._quantizer = None + self._buffers_ = None + self._stream = None + self._pending = {} + self._free = [None, None] + self._cursor = 0 + + 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 + ) + count = 2 if self.prefetch else 1 + self._buffers_ = [ + torch.empty(self.max_decode, device=device, dtype=torch.int8) for _ in range(count) + ] + self._stream = torch.cuda.Stream(device=device) if self.prefetch else None + self._pending = {} + self._free = [None, None] + self._cursor = 0 + return self._quantizer + + def release(self): + """Drop runtime buffers (decode double buffer, side stream) before offloading.""" + if self._stream is not None: + torch.cuda.current_stream().wait_stream(self._stream) + self._reset_runtime() + + def _issue(self, position: int): + index, name = self.order[position % len(self.order)] + slot = position % 2 + group = getattr(self.fused[index], name) + with torch.cuda.stream(self._stream): + if self._free[slot] is not None: + self._stream.wait_event(self._free[slot]) + view = group.decode(self._buffers_[slot]) + ready = torch.cuda.Event() + ready.record(self._stream) + self._pending[position] = (view, ready) + + def _weight(self, index: int, name: str) -> torch.Tensor: + group = getattr(self.fused[index], name) + if not self.prefetch: + return group.decode(self._buffers_[0]) + position = self._cursor + expected = self.order[position % len(self.order)] + if expected != (index, name): + raise RuntimeError(f"decode order mismatch: expected {expected}, got {(index, name)}") + if position not in self._pending: + self._issue(position) + view, ready = self._pending.pop(position) + torch.cuda.current_stream().wait_event(ready) + self._cursor += 1 + return view + + def _consumed(self): + # The main stream has queued the GEMM that reads the last weight: the next decode may + # start into the other buffer once the GEMM before it has finished. + if not self.prefetch: + return + done = torch.cuda.Event() + done.record(torch.cuda.current_stream()) + self._free[(self._cursor - 1) % 2] = done + self._issue(self._cursor) + + def forward_block(self, block, index: int, hidden: torch.Tensor, temb_mod: torch.Tensor, rope): + fused = self.fused[index] + quantize = self._runtime(hidden.device) + attn = block.attn + batch, seq, dim = hidden.shape + if batch != 1: + raise ValueError("the fused Krea 2 runtime runs one image at a time") + h = hidden.reshape(seq, dim) + modulation = temb_mod.unflatten(-1, (6, -1)) + block.scale_shift_table + prescale, preshift, pregate, postscale, postshift, postgate = ( + t.reshape(dim).contiguous() for t in modulation.unbind(-2) + ) + + codes, norms = quantize( + h, + prologue=act.PROLOGUE_NORM_MODULATE, + norm_weight=block.norm1.weight, + mod_scale=prescale, + mod_shift=preshift, + rms_eps=self.rms_eps, + ) + group = fused.qkvg + qkvg = gemm.matmul_int8_scaled_with_triton( + codes, + self._weight(index, "qkvg"), + norms, + group.row_norms, + alpha=group.alpha, + epilogue=gemm.EPILOGUE_SIGMOID_TAIL, + sig_from=group.out_features - dim, + ) + self._consumed() + qw, kw = fused.q_width, fused.kv_width + cos, sin = rope + query = act.qk_norm_rope_with_triton( + qkvg[:, :qw], + attn.num_heads, + attn.head_dim, + attn.norm_q.weight, + cos, + sin, + attn.norm_q.eps, + ) + key = act.qk_norm_rope_with_triton( + qkvg[:, qw : qw + kw], + attn.num_kv_heads, + attn.head_dim, + attn.norm_k.weight, + cos, + sin, + attn.norm_k.eps, + ) + value = qkvg[:, qw + kw : qw + 2 * kw].unflatten(-1, (attn.num_kv_heads, attn.head_dim))[ + None + ] + if self.attention is None and fused.int8_attention: + att = int8_attention_with_triton(query, key, value).reshape(seq, dim) + else: + with sdpa_kernel(self.attention or SDPBackend.FLASH_ATTENTION): + att = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + enable_gqa=attn.num_heads != attn.num_kv_heads, + ) + att = att.transpose(1, 2).reshape(seq, dim) + + codes, norms = quantize(att, prologue=act.PROLOGUE_MULTIPLY, x2=qkvg[:, qw + 2 * kw :]) + group = fused.out + h = gemm.matmul_int8_scaled_with_triton( + codes, + self._weight(index, "out"), + norms, + group.row_norms, + alpha=group.alpha, + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=pregate, + out=h, + ) + self._consumed() + + codes, norms = quantize( + h, + prologue=act.PROLOGUE_NORM_MODULATE, + norm_weight=block.norm2.weight, + mod_scale=postscale, + mod_shift=postshift, + rms_eps=self.rms_eps, + ) + group = fused.gateup + mid = gemm.matmul_int8_scaled_with_triton( + codes, + self._weight(index, "gateup"), + norms, + group.row_norms, + alpha=group.alpha, + epilogue=gemm.EPILOGUE_SWIGLU, + ) + self._consumed() + if self.down_mode == "int8": + codes, scales = act.quantize_rows_int8_with_triton(mid) + h = gemm.matmul_int8_scaled_with_triton( + codes, + fused.down.q, + scales, + fused.down.scales, + scale_mode=1, + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=postgate, + out=h, + ) + else: + h = h + postgate * F.linear(mid, fused.down.weight) + return h.reshape(batch, seq, dim) + + +def install( + transformer, + *, + down: str = "int8", + attention: str = "flash", + prefetch: bool = False, + fused_path: str | os.PathLike | None = None, +) -> FusedKrea2Blocks: + """Move the block projections of ``transformer`` into a ``fused_blocks`` submodule. + + Call it while the weights are on the host; the original projection modules are replaced + by placeholders. Without ``fused_path`` the grouped buffers are built from the source + tensors; with it they are the memory-mapped tensors of a ``save_fused`` file. + """ + stored = None + if fused_path is not None: + stored = _open_fused(fused_path, down, len(transformer.transformer_blocks)) + fused = FusedKrea2Blocks( + transformer, down=down, attention=attention, prefetch=prefetch, stored=stored + ) + for block in transformer.transformer_blocks: + attn, ff = block.attn, block.ff + attn.to_q = attn.to_k = attn.to_v = attn.to_gate = nn.Identity() + attn.to_out[0] = nn.Identity() + ff.gate = ff.up = ff.down = nn.Identity() + transformer.fused_blocks = fused + return fused + + +def save_fused(fused: FusedKrea2Blocks, path: str | os.PathLike) -> None: + """Write the fused block weights to a safetensors file ``install`` can map.""" + tensors = { + f"blocks.{index}.{name}": tensor.contiguous() + for index, block in enumerate(fused.fused) + for name, tensor in block.state_dict().items() + if not name.endswith(".codes") + } + metadata = { + "format": FUSED_FORMAT, + "format_version": FUSED_FORMAT_VERSION, + "down": fused.down_mode, + "blocks": str(len(fused.fused)), + "swiglu_interleave": str(gemm.SWIGLU_INTERLEAVE), + } + save_file(tensors, str(path), metadata=metadata) + + +def _open_fused(path, down: str, blocks: int) -> Callable[[str], torch.Tensor]: + handle = safe_open(str(path), framework="pt", device="cpu") + metadata = handle.metadata() or {} + expected = { + "format": FUSED_FORMAT, + "format_version": FUSED_FORMAT_VERSION, + "down": down, + "blocks": str(blocks), + "swiglu_interleave": str(gemm.SWIGLU_INTERLEAVE), + } + mismatched = {k: metadata.get(k) for k, v in expected.items() if metadata.get(k) != v} + if mismatched: + raise ValueError(f"{path} is not a matching fused Krea 2 file: {mismatched}") + return handle.get_tensor + + +class Krea2FastRunner: + """Text encode, fused denoise and decode for a ``diffusers.Krea2Pipeline`` (batch 1, no CFG). + + Components are moved between devices by the caller; the runner only computes on the + device of the module it calls. + """ + + def __init__(self, pipe, fused: FusedKrea2Blocks): + self.pipe = pipe + self.fused = fused + if getattr(pipe.text_encoder, "visual", None) is not None: + pipe.text_encoder.visual = None + + @torch.no_grad() + def encode(self, prompt: str, device: torch.device) -> torch.Tensor: + """Selected encoder hidden states of the valid tokens ``[1, tokens, layers, dim]``.""" + pipe = self.pipe + tokens = pipe.tokenizer( + [PROMPT_PREFIX + prompt], + truncation=True, + max_length=MAX_SEQUENCE_LENGTH + PROMPT_PREFIX_TOKENS - PROMPT_SUFFIX_TOKENS, + return_tensors="pt", + ).input_ids + suffix = pipe.tokenizer([PROMPT_SUFFIX], return_tensors="pt").input_ids + input_ids = torch.cat([tokens, suffix], dim=1) + length = input_ids.shape[1] + # Right padding keeps kernel shapes stable across prompts; the encoder is causal, so + # tokens after the last valid one cannot change the valid outputs. + bucket = -(-length // TEXT_BUCKET) * TEXT_BUCKET + pad_id = pipe.tokenizer.pad_token_id if pipe.tokenizer.pad_token_id is not None else 0 + padded = torch.full((1, bucket), pad_id, dtype=input_ids.dtype) + padded[:, :length] = input_ids + attention_mask = torch.zeros((1, bucket), dtype=torch.bool) + attention_mask[:, :length] = True + position_ids = torch.arange(bucket, device=device)[None].expand(3, 1, -1) + outputs = pipe.text_encoder( + input_ids=padded.to(device), + attention_mask=attention_mask.to(device), + position_ids=position_ids, + output_hidden_states=True, + ) + hidden = torch.stack( + [outputs.hidden_states[i] for i in pipe.text_encoder_select_layers], dim=2 + ) + return hidden[:, PROMPT_PREFIX_TOKENS:length] + + @torch.no_grad() + def denoise( + self, + text_states: torch.Tensor, + *, + width: int, + height: int, + steps: int = 8, + generator: torch.Generator | None = None, + callback=None, + ) -> torch.Tensor: + """Packed latents after ``steps`` Euler steps of the distilled schedule (``mu = 1.15``).""" + from diffusers.pipelines.krea2.pipeline_krea2 import retrieve_timesteps + + pipe = self.pipe + tr = pipe.transformer + device = text_states.device + text_len = text_states.shape[1] + text = tr.txt_in(tr.text_fusion(text_states, attention_mask=None)) + channels = tr.config.in_channels // (pipe.patch_size**2) + latents = pipe.prepare_latents( + 1, channels, height, width, text_states.dtype, device, generator + ) + factor = pipe.vae_scale_factor * pipe.patch_size + rope = tr.rotary_emb( + pipe.prepare_position_ids(text_len, height // factor, width // factor, device) + ) + sigmas = np.linspace(1.0, 1 / steps, steps) + timesteps, _ = retrieve_timesteps( + pipe.scheduler, steps, device, sigmas=sigmas, mu=DISTILLED_SHIFT_MU + ) + pipe.scheduler.set_begin_index(0) + try: + for index, t in enumerate(timesteps): + timestep = ( + (t / pipe.scheduler.config.num_train_timesteps).expand(1).to(latents.dtype) + ) + temb = tr.time_embed(timestep, dtype=latents.dtype) + temb_mod = tr.time_mod_proj(F.gelu(temb, approximate="tanh")) + hidden = torch.cat([text, tr.img_in(latents)], dim=1) + for block_index, block in enumerate(tr.transformer_blocks): + hidden = self.fused.forward_block(block, block_index, hidden, temb_mod, rope) + velocity = tr.final_layer(hidden[:, text_len:], temb) + latents = pipe.scheduler.step(velocity, t, latents, return_dict=False)[0] + if callback is not None: + callback(index) + finally: + self.fused.release() + return latents + + @torch.no_grad() + def decode(self, latents: torch.Tensor, *, width: int, height: int): + pipe = self.pipe + latents = pipe._unpack_latents(latents, height, width).to(pipe.vae.dtype) + z_dim = pipe.vae.config.z_dim + mean = torch.tensor(pipe.vae.config.latents_mean).view(1, z_dim, 1, 1, 1).to(latents) + inv_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view(1, z_dim, 1, 1, 1).to( + latents + ) + image = pipe.vae.decode(latents / inv_std + mean, return_dict=False)[0][:, :, 0] + return pipe.image_processor.postprocess(image, output_type="pil")[0] diff --git a/tests/test_fused_dit_kernels.py b/tests/test_fused_dit_kernels.py new file mode 100644 index 0000000..aaeec92 --- /dev/null +++ b/tests/test_fused_dit_kernels.py @@ -0,0 +1,254 @@ +import pytest +import torch + +from orbitquant.kernels import available_backends + + +def _require_cuda(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + +def _layer(in_features, out_features, *, bias=False): + from orbitquant import OrbitQuantConfig + from orbitquant.layers import OrbitQuantLinear + + torch.manual_seed(0) + linear = torch.nn.Linear(in_features, out_features, bias=bias, dtype=torch.bfloat16) + return OrbitQuantLinear.from_linear( + linear, config=OrbitQuantConfig(), module_name="probe" + ).cuda() + + +def _int_mm_reference(a, b): + rows = a.shape[0] + padded = torch.zeros( + (max(32, -(-rows // 32) * 32), a.shape[1]), device=a.device, dtype=torch.int8 + ) + padded[:rows] = a + return torch._int_mm(padded, b.t())[:rows] + + +@pytest.mark.parametrize("rows", [1, 33, 2053]) +def test_fused_epilogue_gemm_is_bit_identical_to_int_mm_path(rows): + _require_cuda() + from orbitquant.kernels import triton_cuda as tc + + layer = _layer(1024, 4096, bias=True) + x = torch.randn(rows, 1024, device="cuda", dtype=torch.bfloat16) + act_codes, act_scale, w_codes, w_scale = layer._int8_surrogate_constants(x.device) + packed, norms = tc.quantize_activations_packed_w4_with_triton( + x, rotation=layer.rotation, codebook=layer.activation_codebook, eps=layer.activation_eps + ) + common = dict( + activation_scale=act_scale, + weight_scale=w_scale, + out_features=4096, + in_features=1024, + bias=layer.bias.to(torch.bfloat16), + output_dtype=torch.bfloat16, + chunk_out_features=2048, + ) + args = (packed, layer.packed_weight_indices, norms, layer.row_norms, act_codes, w_codes) + old = tc.matmul_packed_w4a4_with_int_mm(*args, fused_epilogue=False, **common) + new = tc.matmul_packed_w4a4_with_int_mm(*args, fused_epilogue=True, **common) + assert torch.equal(old, new) + + +@pytest.mark.parametrize("dim", [1024, 3072, 6144]) +def test_direct_int8_activation_codes_match_packed_codes(dim): + _require_cuda() + from orbitquant.kernels import triton_cuda as tc + + layer = _layer(dim, 256) + x = torch.randn(77, dim, device="cuda", dtype=torch.bfloat16) + act_codes, _, _, _ = layer._int8_surrogate_constants(x.device) + packed, norms = tc.quantize_activations_packed_w4_with_triton( + x, rotation=layer.rotation, codebook=layer.activation_codebook, eps=layer.activation_eps + ) + direct, direct_norms = tc.quantize_activations_int8_with_triton( + x, + rotation=layer.rotation, + codebook=layer.activation_codebook, + activation_codes=act_codes, + eps=layer.activation_eps, + ) + codes = act_codes.cuda() + decoded = torch.stack([codes[(packed & 15).long()], codes[(packed >> 4).long()]], -1) + assert torch.equal(direct, decoded.reshape(77, dim)) + assert torch.equal(direct_norms, norms) + + +def _gemm_inputs(rows, k, n): + torch.manual_seed(1) + a = torch.randint(-127, 128, (rows, k), device="cuda", dtype=torch.int8) + b = torch.randint(-127, 128, (n, k), device="cuda", dtype=torch.int8) + a_s = torch.rand(rows, device="cuda") * 1e-3 + 1e-4 + b_s = torch.rand(n, device="cuda") + 0.5 + return a, b, a_s, b_s + + +def test_swiglu_epilogue_matches_eager_rounding(): + _require_cuda() + from orbitquant.kernels import triton_int8_gemm as gemm + + a, b, a_s, b_s = _gemm_inputs(300, 512, 1024) + alpha = 0.37 + y = (((_int_mm_reference(a, b).float() * a_s[:, None]) * b_s[None, :]) * alpha).to( + torch.bfloat16 + ) + half = gemm.SWIGLU_INTERLEAVE + chunks = y.reshape(300, -1, 2, half) + gate, up = chunks[:, :, 0].reshape(300, -1), chunks[:, :, 1].reshape(300, -1) + expected = torch.nn.functional.silu(gate) * up + actual = gemm.matmul_int8_scaled_with_triton( + a, b, a_s, b_s, alpha=alpha, epilogue=gemm.EPILOGUE_SWIGLU + ) + assert torch.equal(actual, expected) + + +def test_residual_gate_and_sigmoid_tail_epilogues_match_eager_rounding(): + _require_cuda() + from orbitquant.kernels import triton_int8_gemm as gemm + + a, b, a_s, b_s = _gemm_inputs(129, 512, 768) + y = (((_int_mm_reference(a, b).float() * a_s[:, None]) * b_s[None, :]) * 0.5).to(torch.bfloat16) + residual = torch.randn(129, 768, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(768, device="cuda", dtype=torch.bfloat16) + expected = residual + gate * y + actual = gemm.matmul_int8_scaled_with_triton( + a, + b, + a_s, + b_s, + alpha=0.5, + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=residual.clone(), + gate=gate, + ) + assert torch.equal(actual, expected) + + # Gate logits of real layers stay far from the range where the sigmoid is subnormal, which + # Triton flushes to zero and torch does not. + alpha = 0.01 + y = (((_int_mm_reference(a, b).float() * a_s[:, None]) * b_s[None, :]) * alpha).to( + torch.bfloat16 + ) + tail = gemm.matmul_int8_scaled_with_triton( + a, b, a_s, b_s, alpha=alpha, epilogue=gemm.EPILOGUE_SIGMOID_TAIL, sig_from=512 + ) + assert torch.equal(tail[:, :512], y[:, :512]) + assert torch.equal(tail[:, 512:], torch.sigmoid(y[:, 512:])) + + +def test_in_place_residual_gate_survives_autotuning(): + _require_cuda() + from orbitquant.kernels import triton_int8_gemm as gemm + + kernel = gemm._int8_scaled_gemm + if not hasattr(kernel, "cache") or not hasattr(kernel, "cache_results"): + pytest.skip("Triton autotuner internals differ") + a, b, a_s, b_s = _gemm_inputs(1500, 384, 768) + residual = torch.randn(1500, 768, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(768, device="cuda", dtype=torch.bfloat16) * 0.1 + + def call(): + h = residual.clone() + return gemm.matmul_int8_scaled_with_triton( + a, b, a_s, b_s, epilogue=gemm.EPILOGUE_RESIDUAL_GATE, residual=h, gate=gate, out=h + ) + + saved = kernel.cache_results + kernel.cache.clear() + kernel.cache_results = False + try: + tuned = call() + finally: + kernel.cache_results = saved + assert torch.equal(tuned, call()) + + +def test_rows_int8_matches_int8_row_linear(): + _require_cuda() + from orbitquant.int8_head import Int8RowLinear + from orbitquant.kernels import triton_activation as act + from orbitquant.kernels import triton_int8_gemm as gemm + + torch.manual_seed(2) + linear = torch.nn.Linear(2048, 512, bias=False, dtype=torch.bfloat16) + head = Int8RowLinear.from_linear(linear).cuda() + x = torch.randn(300, 2048, device="cuda", dtype=torch.bfloat16) + expected = head(x) + codes, scales = act.quantize_rows_int8_with_triton(x) + actual = gemm.matmul_int8_scaled_with_triton( + codes, head.weight_int8, scales, head.scales, scale_mode=1 + ) + assert torch.equal(actual, expected) + + +def test_qk_norm_rope_matches_diffusers(): + _require_cuda() + diffusers_embeddings = pytest.importorskip("diffusers.models.embeddings") + from orbitquant.kernels import triton_activation as act + + torch.manual_seed(3) + rows, heads, head_dim = 97, 4, 128 + src = torch.randn(rows, 3 * heads * head_dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(head_dim, device="cuda") * 0.1 + angles = torch.rand(rows, head_dim // 2, device="cuda", dtype=torch.float64) * 6.28 + cos = angles.cos().repeat_interleave(2, dim=1).float() + sin = angles.sin().repeat_interleave(2, dim=1).float() + view = src[:, : heads * head_dim] + q = view.unflatten(-1, (heads, head_dim))[None] + normed = torch.nn.functional.rms_norm(q.float(), (head_dim,), weight=weight + 1.0, eps=1e-5) + expected = diffusers_embeddings.apply_rotary_emb( + normed.to(torch.bfloat16), (cos, sin), sequence_dim=1 + ) + actual = act.qk_norm_rope_with_triton(view, heads, head_dim, weight, cos, sin, 1e-5) + # The RMS reduction order differs from torch's kernel; everything after it is exact. + assert (actual.float() - expected.float()).abs().max() <= 2**-6 * expected.float().abs().max() + + +def _fake_quantized_blocks(x, block, multiplier, mean=None): + x = (x.float() - (0 if mean is None else mean)) * multiplier + out = torch.empty_like(x) + for start in range(0, x.shape[0], block): + chunk = x[start : start + block] + scale = chunk.abs().amax(dim=(0, 2), keepdim=True).clamp_min(1e-20) / 127 + codes = chunk / scale + out[start : start + block] = ( + torch.where(codes >= 0, codes + 0.5, codes - 0.5).trunc() * scale + ) + return out + + +@pytest.mark.parametrize("seq", [64, 1000, 4133]) +def test_int8_attention_matches_its_quantization_and_fp32(seq): + _require_cuda() + from orbitquant.kernels import triton_attention as attention + + torch.manual_seed(4) + query = torch.randn(1, seq, 8, 128, device="cuda", dtype=torch.bfloat16) * 2 + # A shared per-channel offset in K is what the mean subtraction is for. + key = torch.randn(1, seq, 2, 128, device="cuda", dtype=torch.bfloat16) * 2 + 3 + value = torch.randn(1, seq, 2, 128, device="cuda", dtype=torch.bfloat16) + out = attention.int8_attention_with_triton(query, key, value).float() + + q = _fake_quantized_blocks(query[0], attention.BLOCK_Q, 128**-0.5 * attention.LOG2E) + k = _fake_quantized_blocks(key[0], attention.BLOCK_KV, 1.0, key[0].float().mean(dim=0)) + v = value[0].to(torch.float16).float() + logits = torch.einsum("shd,thd->hst", q, k.repeat_interleave(4, dim=1)) + weights = torch.exp2(logits - logits.amax(dim=-1, keepdim=True)) + weights = weights / weights.sum(dim=-1, keepdim=True) + emulated = torch.einsum("hst,thd->shd", weights, v.repeat_interleave(4, dim=1))[None] + exact = torch.nn.functional.scaled_dot_product_attention( + *(t.transpose(1, 2).float() for t in (query, key, value)), enable_gqa=True + ).transpose(1, 2) + + def rel(a, b): + return ((a - b).norm() / b.norm()).item() + + # Everything but the FP16 P.V tiles is emulated exactly. + assert rel(out, emulated) < 5e-3 + # INT8 Q/K itself keeps these peaked scores within a few percent. + assert rel(out, exact) < 5e-2 diff --git a/tests/test_krea2_runtime.py b/tests/test_krea2_runtime.py new file mode 100644 index 0000000..c124daa --- /dev/null +++ b/tests/test_krea2_runtime.py @@ -0,0 +1,72 @@ +import types + +import pytest +import torch + +pytest.importorskip("triton") + + +def _quantized(in_features, out_features, seed): + from orbitquant import OrbitQuantConfig + from orbitquant.layers import OrbitQuantLinear + + torch.manual_seed(seed) + linear = torch.nn.Linear(in_features, out_features, bias=False, dtype=torch.bfloat16) + return OrbitQuantLinear.from_linear(linear, config=OrbitQuantConfig(), module_name=f"l{seed}") + + +def _block(seed, dim=256, heads=2, kv_heads=1, head_dim=128, hidden=512): + attn = types.SimpleNamespace( + to_q=_quantized(dim, heads * head_dim, seed), + to_k=_quantized(dim, kv_heads * head_dim, seed + 1), + to_v=_quantized(dim, kv_heads * head_dim, seed + 2), + to_gate=_quantized(dim, heads * head_dim, seed + 3), + to_out=[_quantized(heads * head_dim, dim, seed + 4)], + num_heads=heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + norm_q=types.SimpleNamespace(weight=torch.randn(head_dim) * 0.1), + norm_k=types.SimpleNamespace(weight=torch.randn(head_dim) * 0.1), + ) + down = torch.nn.Linear(hidden, dim, bias=False, dtype=torch.bfloat16) + ff = types.SimpleNamespace( + gate=_quantized(dim, hidden, seed + 5), up=_quantized(dim, hidden, seed + 6), down=down + ) + return types.SimpleNamespace(attn=attn, ff=ff, norm1=types.SimpleNamespace(eps=1e-6)) + + +def _transformer(): + return types.SimpleNamespace(transformer_blocks=[_block(10), _block(20)]) + + +def test_fused_weights_round_trip_through_a_mapped_file(tmp_path): + from orbitquant.runtime import krea2 + + built = krea2.FusedKrea2Blocks(_transformer()) + path = tmp_path / "fused.safetensors" + krea2.save_fused(built, path) + + target = _transformer() + loaded = krea2.install(target, fused_path=path) + + assert target.fused_blocks is loaded + assert isinstance(target.transformer_blocks[0].attn.to_q, torch.nn.Identity) + expected, actual = built.state_dict(), loaded.state_dict() + assert expected.keys() == actual.keys() + for key in expected: + assert torch.equal(expected[key], actual[key]), key + for a, b in zip(built.fused, loaded.fused, strict=True): + assert (a.qkvg.alpha, a.out.alpha, a.gateup.alpha) == ( + b.qkvg.alpha, + b.out.alpha, + b.gateup.alpha, + ) + + +def test_fused_file_must_match_the_down_mode(tmp_path): + from orbitquant.runtime import krea2 + + path = tmp_path / "fused.safetensors" + krea2.save_fused(krea2.FusedKrea2Blocks(_transformer(), down="int8"), path) + with pytest.raises(ValueError, match="down"): + krea2.install(_transformer(), down="bf16", fused_path=path) diff --git a/uv.lock b/uv.lock index cc5feaf..841db14 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.9.8" +version = "0.10.0" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },