vllm.v1.attention.ops.turboquant_soa.triton_turboquant_unified_attention ¶

def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
    """Return a fresh pair-LUT for ``centroids`` on each call.

    Previously cached by ``(data_ptr, device)`` — that key is unsafe
    because the CUDA allocator can reuse freed addresses for different
    centroid tensors, yielding a stale (wrong-values) LUT. The LUT is tiny
    (N*N*2 fp32, ~2KB at N=16) so unconditional rebuild is essentially
    free compared to attention work. If this shows up on a profile,
    replace with a hash-of-values fingerprint, not data_ptr.
    """
    return build_pair_lut(centroids)

@triton.jit
def _tq_fuse_q_rotation(
    Q,  # [BLOCK_M, HEAD_SIZE_PADDED] — raw query in Q.dtype
    PiT_ptr,
    PiT_stride_0: tl.int64,
    PiT_stride_1: tl.int64,
    dim_mask,  # [HEAD_SIZE_PADDED] int1 — valid head dims
    HEAD_SIZE_PADDED: tl.constexpr,
):
    """Fused Q @ PiT prologue. Called once per program for the MSE-key path
    when the launcher has passed the raw (un-rotated) query. Returns the
    rotated Q in the original dtype.
    """
    d_offs = tl.arange(0, HEAD_SIZE_PADDED)
    pit_offsets = d_offs[:, None] * PiT_stride_0 + d_offs[None, :] * PiT_stride_1
    pit_mask = dim_mask[:, None] & dim_mask[None, :]
    PiT_tile = tl.load(PiT_ptr + pit_offsets, mask=pit_mask, other=0.0).to(tl.float32)
    # input_precision="ieee" pins both inputs to full fp32 MFMA (no TF32
    # truncation). allow_tf32 is intentionally omitted — Triton rejects
    # passing both. This matches the launcher's rocBLAS fp32 GEMM in
    # algebra; rounding-order differs, diff is <= a few fp16 ulp.
    Q_rot = tl.dot(
        Q.to(tl.float32),
        PiT_tile,
        input_precision="ieee",
    )
    return Q_rot.to(Q.dtype)

@triton.jit
def _tq_load_k_tile(
    KV_cache_ptr,
    KV_cache_u16_ptr,  # uint16 view of KV_cache (same storage)
    data_bases,  # [TILE_SIZE] int64 — byte offset to each token's DATA region
    knorm_u16_addrs,  # [TILE_SIZE] int64 — u16 element index for each token's K-norm
    d_offs,  # [HEAD_SIZE_PADDED]
    d_mask,  # [HEAD_SIZE_PADDED] int1
    tile_mask,  # [TILE_SIZE] int1
    Centroids_ptr,
    Pair_lut_ptr,
    OUT_DTYPE: tl.constexpr,  # tl.float16 or tl.bfloat16
    HEAD_DIM: tl.constexpr,
    BLOCK_D: tl.constexpr,  # HEAD_SIZE_PADDED — needed for pair-LUT reshape
    MSE_BITS: tl.constexpr,
    N_CENTROIDS: tl.constexpr,
    KEY_FP8: tl.constexpr,
    USE_PAIR_LUT: tl.constexpr,
    NORM_CORRECTION: tl.constexpr,
    FP8_E4B15: tl.constexpr,
    TILE_SIZE: tl.constexpr,
):
    """Load + dequantize a TILE_SIZE × HEAD_SIZE block of keys and return
    the transposed tile K_T : [HEAD_SIZE_PADDED, TILE_SIZE].

    Opt#3 SoA layout: packed K data is at `data_bases[t] + [0, MSE_BYTES)`
    for MSE keys (or `[0, D)` for FP8). The per-token K-norm lives in the
    per-block SoA metadata region; `knorm_u16_addrs` already encodes its
    u16 element index. For decode tiles aligned with blocks, these addresses
    are contiguous → one coalesced wide load replaces TILE_SIZE scattered
    loads (the whole point of Opt#3).
    """
    if KEY_FP8:
        k_addrs = data_bases[:, None] + d_offs[None, :]
        k_raw = tl.load(
            KV_cache_ptr + k_addrs,
            mask=tile_mask[:, None] & d_mask[None, :],
            other=0,
        )
        if FP8_E4B15:
            k_f32 = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
        else:
            k_f32 = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
        K = k_f32  # [TILE_SIZE, HEAD_SIZE_PADDED]
    else:
        # MSE path: gather packed key indices + centroid LUT.
        if MSE_BITS == 4 and USE_PAIR_LUT:
            # FLUTE pair-LUT fast path — load each packed byte once, decode
            # both nibbles, single 3-D gather returns (T[lo], T[hi]) per byte.
            HALF_D: tl.constexpr = BLOCK_D // 2
            half_offs = tl.arange(0, HALF_D)
            byte_mask = (half_offs * 2) < HEAD_DIM
            byte_addrs = data_bases[:, None] + half_offs[None, :]
            byte_raw = tl.load(
                KV_cache_ptr + byte_addrs,
                mask=tile_mask[:, None] & byte_mask[None, :],
                other=0,
            ).to(tl.int32)
            lo_idx = byte_raw & 0xF
            hi_idx = (byte_raw >> 4) & 0xF
            pair_key = lo_idx * N_CENTROIDS + hi_idx
            pair_slot = tl.arange(0, 2)
            c_pair = tl.load(
                Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],
                mask=(tile_mask[:, None, None] & byte_mask[None, :, None]),
                other=0.0,
            )
            c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])
        elif MSE_BITS == 4:
            half_idx = d_offs // 2
            nibble_shift = (d_offs % 2) * 4
            mse_addrs = data_bases[:, None] + half_idx[None, :]
            mse_raw = tl.load(
                KV_cache_ptr + mse_addrs,
                mask=tile_mask[:, None] & d_mask[None, :],
                other=0,
            ).to(tl.int32)
            mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF
            c_vals = tl.load(
                Centroids_ptr + mse_idx,
                mask=tile_mask[:, None] & d_mask[None, :],
                other=0.0,
            )
        else:
            # Generic bit extraction (3-bit, etc.)
            mse_bit_off = d_offs * MSE_BITS
            mse_byte_idx = mse_bit_off // 8
            mse_bit_shift = mse_bit_off % 8
            mse_mask_val = (1 << MSE_BITS) - 1
            mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]
            mse_raw0 = tl.load(
                KV_cache_ptr + mse_addrs0,
                mask=tile_mask[:, None] & d_mask[None, :],
                other=0,
            ).to(tl.int32)
            mse_raw1 = tl.load(
                KV_cache_ptr + mse_addrs0 + 1,
                mask=tile_mask[:, None] & d_mask[None, :],
                other=0,
            ).to(tl.int32)
            raw16 = mse_raw0 | (mse_raw1 << 8)
            mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val
            c_vals = tl.load(
                Centroids_ptr + mse_idx,
                mask=tile_mask[:, None] & d_mask[None, :],
                other=0.0,
            )

        # Opt#1: 1/||c_vec|| is pre-folded into the stored K-norm at store
        # time, so the kernel doesn't recompute norm-correction here.
        # Opt#3: K-norms for a tile are contiguous in the per-block SoA
        # region when the tile lies within one block (always true for
        # aligned decode tiles with TILE_SIZE == BLOCK_SIZE) — one coalesced
        # u16 load replaces TILE_SIZE scattered 2-byte loads.
        norm_u16 = tl.load(KV_cache_u16_ptr + knorm_u16_addrs, mask=tile_mask, other=0)
        vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)
        K = c_vals * vec_norms[:, None]  # [TILE_SIZE, HEAD_SIZE_PADDED]

    K_T = tl.trans(K.to(OUT_DTYPE))  # [HEAD_SIZE_PADDED, TILE_SIZE]
    _ = HEAD_DIM
    return K_T

@triton.jit
def _tq_load_v_tile(
    KV_cache_ptr,
    KV_cache_u16_ptr,  # uint16 view of KV_cache (same storage)
    val_bases,  # [TILE_SIZE] int64 — byte offset to each token's V-data
    vscale_u16_addrs,  # [TILE_SIZE] int64 — u16 element index for V-scale
    vzero_u16_addrs,  # [TILE_SIZE] int64 — u16 element index for V-zero
    d_offs,
    d_mask,
    tile_mask,
    OUT_DTYPE: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    VQB: tl.constexpr,
):
    """Load + dequantize a TILE_SIZE × HEAD_SIZE block of values.

    Opt#3 SoA layout: packed V data is at `val_bases[t] + [0, VAL_DATA_BYTES)`
    (V-data immediately follows K-data within the slot's data region, and
    `val_bases` is precomputed by the caller as `data_base + KEY_DATA_BYTES`).
    V-scale / V-zero live in the per-block SoA metadata region at indices
    `vscale_u16_addrs` / `vzero_u16_addrs`. For tiles aligned to block
    boundaries, those addresses are contiguous → one coalesced wide load
    per field instead of TILE_SIZE scattered 2-byte loads.
    """
    if VQB == 4:
        vb_idx = d_offs // 2
        vb_shift = (d_offs % 2) * 4
        val_addrs = val_bases[:, None] + vb_idx[None, :]
        val_raw = tl.load(
            KV_cache_ptr + val_addrs,
            mask=tile_mask[:, None] & d_mask[None, :],
            other=0,
        ).to(tl.int32)
        v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)
    else:  # VQB == 3
        val_bit_off = d_offs * 3
        val_byte_idx = val_bit_off // 8
        val_bit_shift = val_bit_off % 8
        val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
        val_raw0 = tl.load(
            KV_cache_ptr + val_addrs0,
            mask=tile_mask[:, None] & d_mask[None, :],
            other=0,
        ).to(tl.int32)
        val_raw1 = tl.load(
            KV_cache_ptr + val_addrs0 + 1,
            mask=tile_mask[:, None] & d_mask[None, :],
            other=0,
        ).to(tl.int32)
        raw16 = val_raw0 | (val_raw1 << 8)
        v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)

    # SoA scale / zero loads — coalesced on aligned decode tiles.
    scale_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addrs, mask=tile_mask, other=0)
    zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=tile_mask, other=0)
    v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)
    v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)

    V = v_idx * v_scales[:, None] + v_zeros[:, None]  # [TILE_SIZE, HEAD_SIZE_PADDED]
    _ = HEAD_DIM
    return V.to(OUT_DTYPE)

def triton_turboquant_decode_attention_soa(
    query: torch.Tensor,  # [B, Hq, D]
    kv_cache: torch.Tensor,
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
    Pi: torch.Tensor,
    centroids: torch.Tensor,
    scale: float,
    mse_bits: int,
    key_packed_size: int,
    value_quant_bits: int,
    value_packed_size: int,
    key_fp8: bool = False,
    norm_correction: bool = False,
    PiT: torch.Tensor | None = None,
    # kept for signature parity with v1/v2 decode launchers; unused here.
    max_seq_len: int = 0,
    mid_o_buf: Any = None,
    output_buf: torch.Tensor | None = None,
    lse_buf: Any = None,
    buf_holder: Any = None,
    max_num_kv_splits: int = 32,
    sinks: torch.Tensor | None = None,
) -> torch.Tensor:
    """Decode-only convenience wrapper around ``triton_turboquant_unified_attention``.

    Treats a rank-3 ``[B, Hq, D]`` query as one token per request (query_len = 1)
    and synthesizes a ``query_start_loc`` of ``[0, 1, 2, ..., B]``.
    ``max_num_kv_splits`` is forwarded to the unified launcher as the 3D
    split-KV segment count (capped per-call against ``max_seq_len``).
    ``sinks`` (optional ``[Hq]`` fp32) are forwarded to the kernel which
    folds them into the softmax denominator via the init-time trick.
    """
    del mid_o_buf, lse_buf, buf_holder
    B = query.shape[0]
    cu_seqlens_q = torch.arange(B + 1, device=query.device, dtype=seq_lens.dtype)

    out = triton_turboquant_unified_attention(
        query=query.contiguous(),
        kv_cache=kv_cache,
        block_table=block_table,
        seq_lens=seq_lens,
        query_start_loc=cu_seqlens_q,
        Pi=Pi,
        centroids=centroids,
        scale=scale,
        mse_bits=mse_bits,
        key_packed_size=key_packed_size,
        value_quant_bits=value_quant_bits,
        value_packed_size=value_packed_size,
        key_fp8=key_fp8,
        norm_correction=norm_correction,
        PiT=PiT,
        output=output_buf[:B] if output_buf is not None else None,
        max_query_len=1,
        max_seq_len=max_seq_len if max_seq_len > 0 else None,
        num_kv_splits=max_num_kv_splits,
        sinks=sinks,
    )
    return out

def triton_turboquant_unified_attention(
    query: torch.Tensor,  # [num_tokens, Hq, D] - fp16/bf16
    kv_cache: torch.Tensor,  # [num_blocks, block_size, Hk, padded_slot] uint8
    block_table: torch.Tensor,  # [num_seqs, max_num_blocks] int32
    seq_lens: torch.Tensor,  # [num_seqs] int32
    query_start_loc: torch.Tensor,  # [num_seqs+1] int32
    Pi: torch.Tensor,  # [D, D] fp32
    centroids: torch.Tensor,  # [n_centroids] fp32
    scale: float,
    mse_bits: int,
    key_packed_size: int,
    value_quant_bits: int,
    value_packed_size: int,  # unused; kept for signature parity with v2
    key_fp8: bool = False,
    norm_correction: bool = False,
    PiT: torch.Tensor | None = None,
    output: torch.Tensor | None = None,
    tile_size: int | None = None,
    max_query_len: int | None = None,
    max_seq_len: int | None = None,
    num_kv_splits: int | None = None,
    force_2d: bool = False,
    fuse_q_rot: bool = True,
    sinks: torch.Tensor | None = None,  # [Hq] float — per-head sink logits
) -> torch.Tensor:
    """Launch unified TQ attention (v3).

    ``query`` carries *raw* query vectors. For the MSE-key path the query
    has to be rotated by ``PiT`` before it can be multiplied against the
    (already-rotated) stored K. By default (``fuse_q_rot=True``) that
    rotation is done inside the attention kernel prologue as a single
    small MFMA — no extra dispatch, no HBM round-trip. Setting
    ``fuse_q_rot=False`` restores the original launcher path (fp32 rocBLAS
    GEMM + casts + .contiguous()), which is kept as an A/B toggle for
    bench harnesses; numerical results match the fused path to within a
    few ulp of fp32 round-off. The FP8-key path never rotates Q regardless
    of this flag.

    ``tile_size`` defaults to ``32`` for prefill (``max_query_len > 1``) and
    ``16`` for pure decode (``max_query_len == 1``). Callers may override.

    Dispatch rule:
      * Prefill / chunked-prefill (any query block with ``BLOCK_Q > 1``) and
        ``num_tokens > num_seqs``: 2D kernel (plenty of CTAs already).
      * Pure decode with long KV: 3D split-KV kernel + ``reduce_segments``
        for extra parallelism; this is the same pattern v2 uses.
      * ``force_2d=True``: force the 2D path regardless (used by the
        apples-to-apples bench to isolate pure TQ dequant overhead).

    ``num_kv_splits`` controls the 3D-split segment count (default ``16``).

    Returns ``output`` of shape ``[num_tokens, Hq, D]`` in ``query.dtype``.
    """
    assert query.dim() == 3, f"query must be [N, Hq, D], got {query.shape}"
    num_tokens, Hq, D = query.shape
    Hk = kv_cache.shape[2]
    block_size = kv_cache.shape[1]
    kv_group_size = Hq // Hk
    num_seqs = int(query_start_loc.shape[0] - 1)
    device = query.device

    cfg = _get_layout(D, mse_bits, value_quant_bits)
    _ = value_packed_size  # unused

    # Q-rotation strategy:
    #   * FP8-key path: no rotation at all (keys stored as fp8, no codebook).
    #   * MSE-key path + fuse_q_rot=True (default): pass the raw query to the
    #     kernel together with PiT and let the kernel's prologue do one
    #     fp32 tl.dot(Q, PiT). Eliminates the launcher-side
    #     (cast, rocBLAS GEMM, cast, contiguous) chain (~50-60 us/step at
    #     D=64, all launch-bound — see bottleneck report §10).
    #   * MSE-key path + fuse_q_rot=False: legacy launcher rotation. Kept
    #     as an A/B toggle for bench harnesses.
    #
    # PiT is always materialized as fp32 D x D contiguous; the kernel loads
    # it as fp32 (cheap, 16KB at D=64) and computes Q @ PiT with IEEE fp32
    # accumulate via tl.dot.
    if key_fp8:
        q_rot = query.contiguous()
        apply_fuse_q_rot = False
    else:
        if PiT is None:
            PiT = Pi.T.contiguous()
        apply_fuse_q_rot = bool(fuse_q_rot)
        if apply_fuse_q_rot:
            q_rot = query.contiguous()
        else:
            q_rot = (query.float() @ PiT).to(query.dtype).contiguous()

    # PiT in fp32, contiguous. For the fused path this is what the kernel
    # loads; for the legacy path it's passed through as a harmless tensor
    # (never dereferenced under the FUSE_Q_ROT constexpr guard). On FP8
    # path Pi may be unused at the call site, so we fall back to an
    # arbitrary non-null tensor (reuse centroids) to satisfy Triton's
    # non-null pointer requirement.
    if (not key_fp8) and PiT is not None:
        PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32)
        if not PiT_f32.is_contiguous():
            PiT_f32 = PiT_f32.contiguous()
        pit_stride_0 = PiT_f32.stride(0)
        pit_stride_1 = PiT_f32.stride(1)
    else:
        PiT_f32 = centroids  # harmless dummy; not dereferenced when FUSE_Q_ROT=0
        pit_stride_0 = 0
        pit_stride_1 = 0

    # Sinks: per-head fp32 logits, contiguous. The kernel only dereferences
    # this pointer when USE_SINKS=1, so when the caller passes None we bind
    # a harmless non-null tensor (centroids) to satisfy Triton's non-null
    # pointer requirement. See the sink design notes in the 2D/3D kernel
    # bodies and the v1 precedent from PR #40663 (sid==0 init trick).
    if sinks is not None:
        sinks_f32 = sinks if sinks.dtype == torch.float32 else sinks.to(torch.float32)
        if not sinks_f32.is_contiguous():
            sinks_f32 = sinks_f32.contiguous()
        assert sinks_f32.numel() == Hq, (
            f"sinks must have shape [Hq={Hq}], got numel={sinks_f32.numel()}"
        )
        use_sinks = True
    else:
        sinks_f32 = centroids  # harmless dummy; not dereferenced when USE_SINKS=0
        use_sinks = False

    if output is None:
        output = torch.empty_like(query)

    # BLOCK_M heuristic (TQ-specific; diverges from stock unified_attention).
    #
    # Stock fp16 picks BLOCK_M=16 for small GQA because fp16 is ALU-saturated
    # at any tile size and smaller BLOCK_M = better occupancy. TQ is a
    # DIFFERENT story: the dequant chain (packed-byte loads + centroid gather
    # + pair-LUT + norm/scale bitcasts) has big per-block fixed overhead, and
    # the pair-LUT gather pattern is latency-bound (NOT HBM-bandwidth-bound,
    # confirmed via rocprof: MemUnitStalled=0.01%, VALUBusy=56% at BM=16).
    #
    # For prefill/chunked (max_query_len > 1) we want BLOCK_M as large as
    # possible to (a) amortize the fixed dequant overhead over more query
    # tokens per CTA and (b) let MFMA saturate with a 128x32 tile shape.
    # Empirical sweep on MI300X across gpt-oss (D=64) and llama (D=128):
    #     prefill B=1 Q=4k     : BM=128 is 3.4x faster than BM=16
    #     chunked B=64 Q=1k C=8: BM=128 is 5.6x faster than BM=16
    # With BM=128 v3 matches or BEATS the fp16 baseline on prefill/chunked.
    #
    # For pure decode (max_query_len == 1) BLOCK_M stays small because
    # BLOCK_Q > 1 just pads rows — each sequence contributes at most 1 query
    # token and larger BLOCK_M wastes lanes on masked-out rows.
    if max_query_len is not None:
        is_prefill_like = max_query_len > 1
    else:
        is_prefill_like = num_tokens > num_seqs

    if is_prefill_like:
        # Prefill / chunked: use BLOCK_M=128 (rounded up to a multiple of
        # kv_group_size). Each CTA does very large work (BLOCK_M x TILE_SIZE
        # x KV_len MMA + dequant), so even small grids (~17 CTAs) saturate
        # MI300X just fine. We verified empirically: BM=128 wins across all
        # prefill/chunked sweep points, including small (B=1 Q=256 C=8k) and
        # large (B=64 Q=1k C=8k; 5.6x speedup vs BM=16). BLOCK_M must be a
        # multiple of kv_group_size so BLOCK_Q is an integer.
        BLOCK_M = max(128, triton.next_power_of_2(kv_group_size))
    else:
        # Decode: BLOCK_Q > 1 in decode just pads rows (at most 1 query
        # token per sequence), so there's little to amortize. Inherit
        # stock's small-BLOCK_M heuristic.
        BLOCK_M = 16 if kv_group_size <= 16 else triton.next_power_of_2(kv_group_size)
    BLOCK_Q = BLOCK_M // kv_group_size

    # Grid: at most ceil(N / BLOCK_Q) + num_seqs q-blocks total, like unified.
    total_num_q_blocks = num_tokens // BLOCK_Q + num_seqs

    # TILE_SIZE heuristic (matches stock unified_attention's _get_tile_size):
    #   prefill (max_query_len > 1): 32
    #   decode  (max_query_len == 1): 16
    if tile_size is None:
        tile_size = 32 if is_prefill_like else 16

    # Pair-LUT fast path for 4-bit MSE keys. Skipped for FP8 and non-4-bit.
    # When USE_PAIR_LUT==0 the kernel never dereferences pair_lut, but Triton
    # still requires a tensor pointer, so we fall back to reusing `centroids`.
    use_pair_lut = (not key_fp8) and (mse_bits == 4)
    pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids

    fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
    num_stages = 1 if _is_hip else 2

    # uint16-aliased view of the cache. Under the Opt#3 SoA layout, K-norm /
    # V-scale / V-zero live in a contiguous per-block metadata region and are
    # always 2-byte aligned, so single-instruction u16 loads replace the
    # original 2× uint8 + OR sequence for every per-token metadata fetch.
    kv_cache_u16 = kv_cache.view(torch.uint16)

    # Opt#3 SoA layout constants (derived locally; matches the store-side
    # computation so the launcher signature stays unchanged). Invariant:
    # data_bytes_per_slot + meta_bytes_per_slot == slot_size_aligned.
    mse_bytes = cfg["mse_bytes"]
    val_data_bytes = cfg["val_data_bytes"]
    key_data_bytes = D if key_fp8 else mse_bytes
    data_bytes_per_slot = key_data_bytes + val_data_bytes
    meta_region_offset = block_size * Hk * data_bytes_per_slot
    num_soa_fields = 2 if key_fp8 else 3
    soa_k_norm = 0  # unused for FP8; harmless constant
    soa_v_scale = 0 if key_fp8 else 1
    soa_v_zero = 1 if key_fp8 else 2

    # ------------------------------------------------------------------
    # Dispatch: 2D for prefill / chunked; 3D split-KV for pure decode with
    # long KV. Falls back to 2D if the sequences are short or if the caller
    # forces it. ``is_prefill_like`` was already computed above for the
    # BLOCK_M / tile_size heuristic; reuse it here.
    # ------------------------------------------------------------------

    # Use 3D only for pure decode where the KV history is long enough that
    # splitting meaningfully increases SM occupancy. Threshold is empirically
    # tuned on MI300X:
    #   seq <  1024 -> 2D path (short KV, reduce overhead > parallelism gain)
    #   seq >= 1024 -> 3D path (v3/v2 ratio 1.04-1.27x across shapes)
    # Below 1024 KV tokens the reduce_segments + scratch-alloc overhead
    # exceeds the parallelism gain from splitting. Above it, the single-CTA
    # 2D path is serial over the KV history and starves MI300X's 304 CUs.
    #
    # IMPORTANT: max_seq_len is expected to come from the caller (backends
    # track it from the block table). We avoid calling seq_lens.max().item()
    # here because that forces a GPU->CPU sync on every launcher call, which
    # adds ~150us of overhead to pure-decode steps.
    if max_seq_len is None:
        # Fallback: infer from block_table (host-side, no sync). This is an
        # upper bound, sufficient for the dispatch decision.
        max_seq_len_hint = int(block_table.shape[1]) * int(block_size)
    else:
        max_seq_len_hint = int(max_seq_len)
    use_3d = (not force_2d) and (not is_prefill_like) and max_seq_len_hint >= 1024

    if not use_3d:
        kernel_tq_unified_attention_2d[(total_num_q_blocks, Hk)](
            output_ptr=output,
            query_ptr=q_rot,
            KV_cache_ptr=kv_cache,
            KV_cache_u16_ptr=kv_cache_u16,
            Centroids_ptr=centroids,
            Pair_lut_ptr=pair_lut,
            PiT_ptr=PiT_f32,
            block_tables_ptr=block_table,
            seq_lens_ptr=seq_lens,
            query_start_len_ptr=query_start_loc,
            sinks_ptr=sinks_f32,
            scale=scale,
            num_query_heads=Hq,
            num_queries_per_kv=kv_group_size,
            block_table_stride=block_table.stride(0),
            query_stride_0=q_rot.stride(0),
            query_stride_1=q_rot.stride(1),
            output_stride_0=output.stride(0),
            output_stride_1=output.stride(1),
            stride_cache_block=kv_cache.stride(0),
            pit_stride_0=pit_stride_0,
            pit_stride_1=pit_stride_1,
            BLOCK_SIZE=block_size,
            TILE_SIZE=tile_size,
            HEAD_SIZE=D,
            HEAD_SIZE_PADDED=cfg["BLOCK_D"],
            BLOCK_Q=BLOCK_Q,
            BLOCK_M=BLOCK_M,
            num_seqs=num_seqs,
            MSE_BITS=mse_bits,
            MSE_BYTES=mse_bytes,
            VQB=value_quant_bits,
            VAL_DATA_BYTES=val_data_bytes,
            N_CENTROIDS=int(centroids.numel()),
            KEY_FP8=1 if key_fp8 else 0,
            USE_PAIR_LUT=1 if use_pair_lut else 0,
            NUM_KV_HEADS=Hk,
            KEY_DATA_BYTES=key_data_bytes,
            META_REGION_OFFSET=meta_region_offset,
            NUM_SOA_FIELDS=num_soa_fields,
            SOA_K_NORM=soa_k_norm,
            SOA_V_SCALE=soa_v_scale,
            SOA_V_ZERO=soa_v_zero,
            NORM_CORRECTION=1 if norm_correction else 0,
            FP8_E4B15=fp8_e4b15,
            FUSE_Q_ROT=1 if apply_fuse_q_rot else 0,
            USE_SINKS=1 if use_sinks else 0,
            num_warps=4,
            num_stages=num_stages,
        )
        return output

    # ---------------- 3D split-KV path ----------------
    # Pick segment count. 16 is a good default for MI300X (304 CUs): at
    # Hk=8, num_seqs=1, total CTAs = 1 * 8 * 16 = 128, which saturates when
    # combined with warp-level parallelism. Capped at ceil(max_seq_len /
    # TILE_SIZE) so we don't launch empty segments.
    if num_kv_splits is None:
        num_kv_splits = 16
    max_possible_splits = max(1, (max_seq_len_hint + tile_size - 1) // tile_size)
    num_segments = max(1, min(num_kv_splits, max_possible_splits))

    HEAD_SIZE_PADDED = cfg["BLOCK_D"]
    # Scratch buffers for segment partials. fp32 to match the kernel's
    # accumulator dtype. Allocated per-call for now; a production backend
    # would hoist these to persistent buffers reused across layers.
    segm_output = torch.empty(
        (num_tokens, Hq, num_segments, HEAD_SIZE_PADDED),
        dtype=torch.float32,
        device=device,
    )
    segm_max = torch.empty(
        (num_tokens, Hq, num_segments), dtype=torch.float32, device=device
    )
    segm_expsum = torch.empty(
        (num_tokens, Hq, num_segments), dtype=torch.float32, device=device
    )

    kernel_tq_unified_attention_3d[(total_num_q_blocks, Hk, num_segments)](
        segm_output_ptr=segm_output,
        segm_max_ptr=segm_max,
        segm_expsum_ptr=segm_expsum,
        query_ptr=q_rot,
        KV_cache_ptr=kv_cache,
        KV_cache_u16_ptr=kv_cache_u16,
        Centroids_ptr=centroids,
        Pair_lut_ptr=pair_lut,
        PiT_ptr=PiT_f32,
        block_tables_ptr=block_table,
        seq_lens_ptr=seq_lens,
        query_start_len_ptr=query_start_loc,
        sinks_ptr=sinks_f32,
        scale=scale,
        num_query_heads=Hq,
        num_queries_per_kv=kv_group_size,
        block_table_stride=block_table.stride(0),
        query_stride_0=q_rot.stride(0),
        query_stride_1=q_rot.stride(1),
        stride_cache_block=kv_cache.stride(0),
        pit_stride_0=pit_stride_0,
        pit_stride_1=pit_stride_1,
        BLOCK_SIZE=block_size,
        TILE_SIZE=tile_size,
        HEAD_SIZE=D,
        HEAD_SIZE_PADDED=HEAD_SIZE_PADDED,
        BLOCK_Q=BLOCK_Q,
        BLOCK_M=BLOCK_M,
        num_seqs=num_seqs,
        NUM_SEGMENTS_PER_SEQ=num_segments,
        MSE_BITS=mse_bits,
        MSE_BYTES=mse_bytes,
        VQB=value_quant_bits,
        VAL_DATA_BYTES=val_data_bytes,
        N_CENTROIDS=int(centroids.numel()),
        KEY_FP8=1 if key_fp8 else 0,
        USE_PAIR_LUT=1 if use_pair_lut else 0,
        NUM_KV_HEADS=Hk,
        KEY_DATA_BYTES=key_data_bytes,
        META_REGION_OFFSET=meta_region_offset,
        NUM_SOA_FIELDS=num_soa_fields,
        SOA_K_NORM=soa_k_norm,
        SOA_V_SCALE=soa_v_scale,
        SOA_V_ZERO=soa_v_zero,
        NORM_CORRECTION=1 if norm_correction else 0,
        FP8_E4B15=fp8_e4b15,
        FUSE_Q_ROT=1 if apply_fuse_q_rot else 0,
        USE_SINKS=1 if use_sinks else 0,
        num_warps=4,
        num_stages=num_stages,
    )

    # Reduce: merge num_segments partials per (q_token, q_head) into the
    # final output via online-softmax merge. Reusing the baseline kernel
    # (KV-format-agnostic).
    reduce_segments[(num_tokens, Hq)](
        output_ptr=output,
        segm_output_ptr=segm_output,
        segm_max_ptr=segm_max,
        segm_expsum_ptr=segm_expsum,
        seq_lens_ptr=seq_lens,
        num_seqs=num_seqs,
        num_query_heads=Hq,
        out_scale_inv=1.0,
        output_stride_0=output.stride(0),
        output_stride_1=output.stride(1),
        block_table_stride=block_table.stride(0),
        TILE_SIZE=tile_size,
        HEAD_SIZE=D,
        HEAD_SIZE_PADDED=HEAD_SIZE_PADDED,
        query_start_len_ptr=query_start_loc,
        BLOCK_Q=BLOCK_Q,
        NUM_SEGMENTS_PER_SEQ=num_segments,
        USE_FP8=False,
    )
    return output