Triton API

Triton kernels for contiguous compressed decode blocks.

UnsupportedFormatError

Bases: RuntimeError

Raised when a requested compressed attention kernel is unavailable.

TQ packed page update operations.

TQ packed page decode operations.

Fused TQ4 Flash Attention -- both K and V decompressed in the inner loop.

Both K and V tiles are decompressed inline from nibble-packed uint8 indices. The query is pre-rotated by Pi^T and the output is post-rotated by Pi outside the kernel (since K and V share the same rotation matrix). Non-power-of-two HEAD_DIM (e.g., 96) is supported via padded tl.arange + masking.

Autotune configs include BLOCK_M values 16, 32, 64, 128 to cover head_dim up to 256.

Examples:

from turboquant.triton.flash_attention_tq4_kv import (
    triton_flash_attention_tq4_kv,
)

out = triton_flash_attention_tq4_kv(
    q,
    k_packed,
    k_norms,
    v_packed,
    v_norms,
    centroids,
    rotation,
    sm_scale=None,
)

triton_flash_attention_tq4_kv

triton_flash_attention_tq4_kv(q: Tensor, k_packed: Tensor, k_norms: Tensor, v_packed: Tensor, v_norms: Tensor, centroids: Tensor, rotation: Tensor, sm_scale: float | None = None, is_causal: bool = False) -> torch.Tensor

Fused TQ4 Flash Attention with both K and V compressed.

Pre-rotates Q by rotation^T, launches the kernel that decompresses both K and V inline, then post-rotates the output by rotation to return to the original coordinate space. Non-power-of-two head dimensions (e.g., 96) are handled via padded tile loads and boundary masking inside the kernel.

Args: q: Query [batch, H_Q, seq_q, head_dim] fp16/bf16. k_packed: Nibble-packed key indices [batch, H_KV, seq_kv, D//2] uint8. k_norms: Key norms [batch, H_KV, seq_kv] or [..., 1] fp32. v_packed: Nibble-packed value indices [batch, H_KV, seq_kv, D//2] uint8. v_norms: Value norms [batch, H_KV, seq_kv] or [..., 1] fp32. centroids: Shared Lloyd-Max codebook [16] fp32. rotation: Shared orthogonal rotation [head_dim, head_dim] fp32. sm_scale: Softmax scale. Defaults to 1 / sqrt(head_dim). is_causal: Apply causal masking.

Returns: Attention output [batch, H_Q, seq_q, head_dim] in original space.

Source code in turboquant/triton/flash_attention_tq4_kv.py

def triton_flash_attention_tq4_kv(
    q: torch.Tensor,
    k_packed: torch.Tensor,
    k_norms: torch.Tensor,
    v_packed: torch.Tensor,
    v_norms: torch.Tensor,
    centroids: torch.Tensor,
    rotation: torch.Tensor,
    sm_scale: float | None = None,
    is_causal: bool = False,
) -> torch.Tensor:
    """Fused TQ4 Flash Attention with both K and V compressed.

    Pre-rotates Q by ``rotation^T``, launches the kernel that decompresses
    both K and V inline, then post-rotates the output by ``rotation`` to
    return to the original coordinate space. Non-power-of-two head
    dimensions (e.g., 96) are handled via padded tile loads and boundary
    masking inside the kernel.

    Args:
        q: Query ``[batch, H_Q, seq_q, head_dim]`` fp16/bf16.
        k_packed: Nibble-packed key indices ``[batch, H_KV, seq_kv, D//2]`` uint8.
        k_norms: Key norms ``[batch, H_KV, seq_kv]`` or ``[..., 1]`` fp32.
        v_packed: Nibble-packed value indices ``[batch, H_KV, seq_kv, D//2]`` uint8.
        v_norms: Value norms ``[batch, H_KV, seq_kv]`` or ``[..., 1]`` fp32.
        centroids: Shared Lloyd-Max codebook ``[16]`` fp32.
        rotation: Shared orthogonal rotation ``[head_dim, head_dim]`` fp32.
        sm_scale: Softmax scale. Defaults to ``1 / sqrt(head_dim)``.
        is_causal: Apply causal masking.

    Returns:
        Attention output ``[batch, H_Q, seq_q, head_dim]`` in original space.
    """
    B, H_Q, N_Q, D = q.shape
    _, H_KV, N_KV, HALF_D = k_packed.shape

    assert D % 2 == 0, f"HEAD_DIM must be even, got {D}"
    assert HALF_D == D // 2
    assert H_Q % H_KV == 0
    assert k_packed.dtype == torch.uint8
    assert v_packed.dtype == torch.uint8
    assert k_norms.dtype == torch.float32
    assert v_norms.dtype == torch.float32

    # Squeeze trailing 1 from norms
    if k_norms.dim() == 4 and k_norms.shape[-1] == 1:
        k_norms = k_norms.squeeze(-1)
    if v_norms.dim() == 4 and v_norms.shape[-1] == 1:
        v_norms = v_norms.squeeze(-1)

    if is_causal and N_Q == 1:
        is_causal = False

    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(D)

    # Pre-rotate Q
    q_rot = torch.matmul(q.float(), rotation.T).to(q.dtype)

    out_rot = torch.empty_like(q)

    def grid(META: dict) -> tuple[int, int]:
        """Compute launch grid."""
        return (triton.cdiv(N_Q, META["BLOCK_M"]), B * H_Q)

    HEAD_DIM_PAD = _next_pow2(D)
    HALF_D_PAD = _next_pow2(D // 2)
    assert HALF_D_PAD * 2 == HEAD_DIM_PAD, (
        f"Padding invariant violated: 2*HALF_D_PAD ({2 * HALF_D_PAD}) "
        f"!= HEAD_DIM_PAD ({HEAD_DIM_PAD}) — tl.join reshape requires this"
    )

    _fwd_tq4_kv_kernel[grid](
        q_rot,
        k_packed,
        k_norms,
        v_packed,
        v_norms,
        centroids,
        out_rot,
        sm_scale,
        *q_rot.stride(),
        *k_packed.stride(),
        *k_norms.stride(),
        *v_packed.stride(),
        *v_norms.stride(),
        *out_rot.stride(),
        H_Q,
        H_KV,
        N_Q,
        N_KV,
        HEAD_DIM=D,
        HEAD_DIM_PAD=HEAD_DIM_PAD,
        HALF_D_PAD=HALF_D_PAD,
        IS_CAUSAL=is_causal,
    )

    # Post-rotate: convert from rotated space back to original space
    return torch.matmul(out_rot.float(), rotation).to(q.dtype)

Fused paged TQ4 decode attention -- decompresses directly from page table.

This kernel reads TQ4-compressed blocks directly from vLLM's paged block table, decompresses in SRAM (nibble unpack -> centroid gather -> norm scale), and computes FP16 Q@K^T with online softmax in a single fused pass.

No HBM writes of decompressed cache -- HBM traffic drops from 1,160 to 136 bytes/token (8.5x reduction).

The kernel operates entirely in rotated space. The caller pre-rotates Q by Pi^T and post-rotates the output by Pi. Decompression does NOT apply rotation (matching tq4_decompress.py).

Scope: FP16/BF16 Q decode path only (USE_INT8_QK=False).

Autotune: 8 configs (BLOCK_N in {32, 64} x stages {2,3} x warps {4,8}).

Examples:

from turboquant.triton.fused_paged_tq4_attention import (
    fused_paged_tq4_decode,
)

out = fused_paged_tq4_decode(
    q,
    kv_cache,
    block_table,
    seq_lens,
    centroids,
    rotation,
    num_kv_heads=4,
    head_dim=128,
    block_size=16,
)

fused_paged_tq4_decode

fused_paged_tq4_decode(q: Tensor, kv_cache: Tensor, block_table: Tensor, seq_lens: Tensor, centroids: Tensor, rotation: Tensor, num_kv_heads: int, head_dim: int, block_size: int, sm_scale: float | None = None, out: Tensor | None = None) -> torch.Tensor

Fused paged TQ4 decode attention.

Pre-rotates Q by rotation^T, launches the fused paged kernel that decompresses TQ4 blocks in-tile from the page table, then post-rotates the output by rotation to return to original space.

Args: q: Query [num_seqs, H_Q, head_dim] fp16/bf16 (one token per seq). kv_cache: Packed paged cache [num_blocks, block_size, total_bytes] uint8. block_table: Page table [num_seqs, max_num_blocks_per_seq] int32. seq_lens: Sequence lengths [num_seqs] int32. centroids: TQ4 codebook [16] fp32. rotation: Orthogonal rotation [head_dim, head_dim] fp32. num_kv_heads: Number of KV heads. head_dim: Head dimension (e.g. 128). block_size: vLLM page size (tokens per block). sm_scale: Softmax scale. Defaults to 1 / sqrt(head_dim). out: Optional pre-allocated output [num_seqs, H_Q, head_dim].

Returns: Attention output [num_seqs, H_Q, head_dim] in original space.

Source code in turboquant/triton/fused_paged_tq4_attention.py

def fused_paged_tq4_decode(
    q: torch.Tensor,
    kv_cache: torch.Tensor,
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
    centroids: torch.Tensor,
    rotation: torch.Tensor,
    num_kv_heads: int,
    head_dim: int,
    block_size: int,
    sm_scale: float | None = None,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Fused paged TQ4 decode attention.

    Pre-rotates Q by ``rotation^T``, launches the fused paged kernel
    that decompresses TQ4 blocks in-tile from the page table, then
    post-rotates the output by ``rotation`` to return to original space.

    Args:
        q: Query ``[num_seqs, H_Q, head_dim]`` fp16/bf16 (one token per seq).
        kv_cache: Packed paged cache ``[num_blocks, block_size, total_bytes]``
            uint8.
        block_table: Page table ``[num_seqs, max_num_blocks_per_seq]`` int32.
        seq_lens: Sequence lengths ``[num_seqs]`` int32.
        centroids: TQ4 codebook ``[16]`` fp32.
        rotation: Orthogonal rotation ``[head_dim, head_dim]`` fp32.
        num_kv_heads: Number of KV heads.
        head_dim: Head dimension (e.g. 128).
        block_size: vLLM page size (tokens per block).
        sm_scale: Softmax scale. Defaults to ``1 / sqrt(head_dim)``.
        out: Optional pre-allocated output ``[num_seqs, H_Q, head_dim]``.

    Returns:
        Attention output ``[num_seqs, H_Q, head_dim]`` in original space.
    """
    num_seqs, H_Q, D = q.shape

    assert D == head_dim
    assert H_Q % num_kv_heads == 0
    assert kv_cache.dtype == torch.uint8
    assert block_table.dtype == torch.int32
    assert seq_lens.dtype == torch.int32

    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(head_dim)

    half_D = head_dim // 2

    # Byte layout constexprs
    k_norm_offset = num_kv_heads * half_D
    v_idx_offset = k_norm_offset + num_kv_heads * 4
    v_norm_offset = v_idx_offset + num_kv_heads * half_D

    # Pre-rotate Q by Pi^T (O(num_seqs), not O(cache_len))
    q_rot = torch.matmul(q.float(), rotation.T).to(q.dtype)

    out_rot = torch.empty_like(q)

    # INT8 placeholders (unused, compiled out)
    dummy = torch.empty(0, device=q.device)

    grid = (num_seqs, H_Q)

    _fused_paged_tq4_decode_kernel[grid](
        q_rot,
        kv_cache,
        block_table,
        seq_lens,
        centroids,
        out_rot,
        dummy,  # Q_scale
        dummy,  # QJL_S
        dummy,  # QJL_signs
        dummy,  # QJL_residual_norms
        q_rot.stride(0),
        q_rot.stride(1),
        q_rot.stride(2),
        kv_cache.stride(0),
        kv_cache.stride(1),
        block_table.stride(0),
        block_table.stride(1),
        out_rot.stride(0),
        out_rot.stride(1),
        out_rot.stride(2),
        sm_scale=sm_scale,
        H_Q=H_Q,
        H_KV=num_kv_heads,
        HEAD_DIM=head_dim,
        BLOCK_SIZE=block_size,
        HALF_D=half_D,
        K_NORM_OFFSET=k_norm_offset,
        V_IDX_OFFSET=v_idx_offset,
        V_NORM_OFFSET=v_norm_offset,
        USE_INT8_QK=False,
        QJL_DIM=0,
    )

    # Post-rotate: convert from rotated space back to original space
    result = torch.matmul(out_rot.float(), rotation).to(q.dtype)
    if out is not None:
        out.copy_(result)
        return out
    return result