From 03d5d098ba8ab3c5d13000b76bcf7b7a3e4c525c Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Mon, 9 Nov 2026 16:16:43 +0100 Subject: [PATCH] INT8 attention: scale large values per head before the FP16 P.V tiles (0.17.1) A 64-key FP16 tile of P.V reaches 64 * max|v|. Vimoo 1.0 Lite reaches |v| = 7136 in block 38 late in sampling, so the tile overflowed to inf and the video turned to NaN. Heads above |v| = 512 are now divided by a power of two before the FP16 conversion and the output is multiplied back; the scale is exact, so heads within the limit give the same bits as before. --- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/kernels/triton_attention.py | 21 ++++++++++++++++++--- tests/test_fused_dit_kernels.py | 22 ++++++++++++++++++++++ uv.lock | 2 +- 5 files changed, 43 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ea0595a..c4d8a3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.17.0" +version = "0.17.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 9b08c45..b73d06c 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.17.0" +__version__ = "0.17.1" register_hf_quantizers() diff --git a/src/orbitquant/kernels/triton_attention.py b/src/orbitquant/kernels/triton_attention.py index 7f854fb..cdbe0b0 100644 --- a/src/orbitquant/kernels/triton_attention.py +++ b/src/orbitquant/kernels/triton_attention.py @@ -11,7 +11,10 @@ * 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. V is converted once up front: converting each tile inside the loop - costs about 10% of the kernel time. + costs about 10% of the kernel time. A tile sum reaches ``BLOCK_N * max|v|``, so a head whose + values exceed ``V_LIMIT`` is divided by a power of two before the conversion and the output is + multiplied back. The power-of-two scale is exact: heads within the limit give the same bits as + without it, and the others no longer overflow FP16 (Vimoo 1.0 Lite reaches |v| = 7136). 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). Keys and values may be a different @@ -32,6 +35,8 @@ LOG2E = 1.4426950408889634 BLOCK_Q = 128 BLOCK_KV = 64 +# Largest |v| that a head keeps unscaled: BLOCK_KV * V_LIMIT = 32768 leaves FP16 headroom. +V_LIMIT = 512.0 @triton.jit @@ -103,6 +108,7 @@ def _attention_kernel( v_ptr, q_scale_ptr, k_scale_ptr, + v_scale_ptr, o_ptr, seq, seq_kv, @@ -149,7 +155,7 @@ def _attention_kernel( acc, l_i, m_i = _attention_tile( acc, l_i, m_i, q, qk_scale, k_ptrs, v_ptrs, seq_kv - full * BLOCK_N, True, BLOCK_N ) - out = acc / l_i[:, None] + out = acc / l_i[:, None] * tl.load(v_scale_ptr + kv_head) o_rows = offs_m[:, None].to(tl.int64) * stride_or + head * stride_oh tl.store( o_ptr + o_rows + offs_d[None, :], @@ -164,6 +170,14 @@ def _attention_kernel( )(_attention_kernel) +def _scaled_fp16_values(v): + """FP16 copy of ``[tokens, heads, dim]`` values whose heads stay within ``V_LIMIT``, and the + power-of-two factor of each head (1 for heads already within it).""" + peak = torch.linalg.vector_norm(v, ord=math.inf, dim=(0, 2), dtype=torch.float32) + scales = torch.exp2(torch.ceil(torch.log2(peak / V_LIMIT)).clamp_(min=0)) + return (v / scales[:, None].to(v.dtype)).to(torch.float16), scales + + def _query_block(dim: int) -> int: # The FP32 accumulator is BLOCK_M x D: wide heads take shorter query tiles. return BLOCK_Q if dim <= 128 else 64 @@ -217,7 +231,7 @@ def int8_attention_with_triton( block_q = _query_block(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.mean(dim=0, dtype=torch.float32)) - v16 = v if v.dtype == torch.float16 else v.to(torch.float16) + v16, v_scales = _scaled_fp16_values(v) out = torch.empty((1, seq, heads, dim), device=q.device, dtype=query.dtype) o = out[0] _attention[(triton.cdiv(seq, block_q), heads)]( @@ -226,6 +240,7 @@ def int8_attention_with_triton( v16, q_scales, k_scales, + v_scales, o, seq, seq_kv, diff --git a/tests/test_fused_dit_kernels.py b/tests/test_fused_dit_kernels.py index aaeec92..e81fa74 100644 --- a/tests/test_fused_dit_kernels.py +++ b/tests/test_fused_dit_kernels.py @@ -252,3 +252,25 @@ def rel(a, b): assert rel(out, emulated) < 5e-3 # INT8 Q/K itself keeps these peaked scores within a few percent. assert rel(out, exact) < 5e-2 + + +def test_int8_attention_keeps_large_values_finite(): + # Equal scores give every key of a tile p = 1, so one FP16 tile of P.V reaches 64 * |v|: + # 64 * 4000 overflows FP16 unless the values of that head are scaled down first. + _require_cuda() + from orbitquant.kernels import triton_attention as attention + + torch.manual_seed(5) + seq = 1000 + query = torch.zeros(1, seq, 4, 128, device="cuda", dtype=torch.bfloat16) + key = torch.randn(1, seq, 2, 128, device="cuda", dtype=torch.bfloat16) + value = torch.randn(1, seq, 2, 128, device="cuda", dtype=torch.bfloat16) + value[0, :, 0, 0] = 4000.0 # kv head 0 needs a scale; kv head 1 stays within V_LIMIT + out = attention.int8_attention_with_triton(query, key, value).float() + 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) + assert torch.isfinite(out).all() + torch.testing.assert_close(out[0, :, :2, 0], exact[0, :, :2, 0], rtol=2e-3, atol=0) + small, small_exact = out[0, :, 2:], exact[0, :, 2:] + assert ((small - small_exact).norm() / small_exact.norm()).item() < 5e-3 diff --git a/uv.lock b/uv.lock index 09b5581..ea791e1 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.17.0" +version = "0.17.1" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },