vllm.v1.attention.ops.turboquant_soa.triton_turboquant_decode_v2 ¶

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

    The LUT is tiny (N*N*2 fp32, e.g. 2KB for 4-bit MSE) so the build cost
    is negligible compared to attention. We avoid caching by data_ptr()
    because CUDA allocator memory reuse across different centroid tensors
    can silently return a stale LUT (subtle correctness bug). If this ever
    shows up on a profile, cache by a hash-of-values fingerprint instead.
    """
    return build_pair_lut(centroids)

def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
    """Build vectorized pair lookup table.

    For N centroids, returns a [N, N, 2] float32 tensor where
    pair_lut[i, j] = (centroids[i], centroids[j]).
    Flattened to [N*N, 2] for kernel indexing: pair_lut[i*N + j].

    For 4-bit MSE (N=16): 16*16*2*4 = 2048 bytes — fits in L1/smem.
    For 3-bit MSE (N=8):   8*8*2*4  = 512 bytes.
    """
    N = centroids.shape[0]
    # pair_lut[i,j,0] = centroids[i], pair_lut[i,j,1] = centroids[j]
    c = centroids.float()
    lut = torch.empty(N, N, 2, dtype=torch.float32, device=centroids.device)
    lut[:, :, 0] = c[:, None]
    lut[:, :, 1] = c[None, :]
    return lut.reshape(N * N, 2).contiguous()

def triton_turboquant_decode_attention_v2(
    query: torch.Tensor,
    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,
    max_seq_len: int = 0,  # unused; kept for backward compatibility
    key_fp8: bool = False,
    norm_correction: bool = False,
    PiT: torch.Tensor | None = None,
    # Pre-allocated buffers (optional, required for CUDA-graph stability).
    mid_o_buf: torch.Tensor | None = None,
    output_buf: torch.Tensor | None = None,
    lse_buf: torch.Tensor | None = None,
    buf_holder: Any = None,
    # Fixed split count — MUST be a compile-time constant across iterations
    # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract.
    max_num_kv_splits: int = 32,
) -> torch.Tensor:
    """Launch optimized TQ decode attention (v2 stage1 + shared stage2).

    Follows the same buffer-reuse + fixed-grid contract as the v1 launcher
    (triton_turboquant_decode_attention) so the backend can capture a CUDA
    graph across both versions.
    """
    B, Hq, D = query.shape
    Hk = kv_cache.shape[2]
    block_size = kv_cache.shape[1]
    padded_slot = kv_cache.shape[3]
    max_num_blocks = block_table.shape[1]
    n_centroids = centroids.shape[0]
    kv_group_size = Hq // Hk
    device = query.device
    del max_seq_len  # no longer used: splits is fixed via max_num_kv_splits

    cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)

    # Opt#3 SoA layout constants (match store-side computation).
    key_data_bytes = D if key_fp8 else cfg["mse_bytes"]
    data_bytes_per_slot = key_data_bytes + cfg["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
    soa_v_scale = 0 if key_fp8 else 1
    soa_v_zero = 1 if key_fp8 else 2
    kv_cache_u16 = kv_cache.view(torch.uint16)

    # Compute q_rot = q @ Pi.T
    if key_fp8:
        q_rot = query.float().contiguous()
    else:
        q_float = query.float()
        if PiT is None:
            PiT = Pi.T.contiguous()
        q_rot = (q_float @ PiT).contiguous()

    # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores
    BLOCK_M = max(16, triton.next_power_of_2(kv_group_size))

    # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration
    TILE_SIZE = 16

    # Fixed split count — must be constant across calls for cudagraph replay.
    NUM_KV_SPLITS = max_num_kv_splits

    # --- mid_o buffer reuse (same pattern as v1 launcher) ---
    if (
        mid_o_buf is not None
        and mid_o_buf.shape[0] >= B
        and mid_o_buf.shape[2] >= NUM_KV_SPLITS
    ):
        mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]
    else:
        mid_o = torch.empty(
            B,
            Hq,
            NUM_KV_SPLITS,
            D + 1,
            dtype=torch.float32,
            device=device,
        )
        if buf_holder is not None:
            buf_holder._tq_mid_o_buf = mid_o

    fp8_e4b15 = _use_fp8_e4b15(device.index or 0)

    # Build pair LUT for 4-bit MSE (FLUTE §3.2)
    use_pair_lut = mse_bits == 4 and not key_fp8
    pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids

    # Platform-dependent pipelining depth: ROCm prefers num_stages=1.
    stage1_num_stages = 1 if _is_hip else 2

    # --- Stage 1: v2 kernel ---
    # Grid over KV heads (not Q heads) — each program handles all Q heads
    grid = (B, Hk, NUM_KV_SPLITS)
    _tq_decode_stage1_v2[grid](
        q_rot,
        kv_cache,
        kv_cache_u16,
        block_table,
        seq_lens,
        centroids,
        pair_lut,
        mid_o,
        q_rot.stride(0),
        q_rot.stride(1),
        kv_cache.stride(0),
        block_table.stride(0),
        mid_o.stride(0),
        mid_o.stride(1),
        mid_o.stride(2),
        NUM_Q_HEADS=Hq,
        NUM_KV_HEADS=Hk,
        HEAD_DIM=D,
        BLOCK_SIZE=block_size,
        PADDED_SLOT=padded_slot,
        MAX_NUM_BLOCKS=max_num_blocks,
        NUM_KV_SPLITS=NUM_KV_SPLITS,
        KV_GROUP_SIZE=kv_group_size,
        MSE_BITS=mse_bits,
        MSE_BYTES=cfg["mse_bytes"],
        VQB=value_quant_bits,
        VAL_DATA_BYTES=cfg["val_data_bytes"],
        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,
        N_CENTROIDS=n_centroids,
        ATTN_SCALE=scale,
        BLOCK_D=cfg["BLOCK_D"],
        TILE_SIZE=TILE_SIZE,
        BLOCK_M=BLOCK_M,
        KEY_FP8=1 if key_fp8 else 0,
        NORM_CORRECTION=1 if norm_correction else 0,
        FP8_E4B15=fp8_e4b15,
        USE_PAIR_LUT=1 if use_pair_lut else 0,
        USE_BF16_DOT=1 if _is_hip else 0,
        num_warps=4,
        num_stages=stage1_num_stages,
    )

    # --- output / lse buffer reuse (same pattern as v1 launcher) ---
    if output_buf is not None and output_buf.shape[0] >= B:
        output = output_buf[:B, :Hq, :D]
    else:
        output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)
        if buf_holder is not None:
            buf_holder._tq_output_buf = output
    if lse_buf is not None and lse_buf.shape[0] >= B:
        lse = lse_buf[:B, :Hq]
    else:
        lse = torch.empty(B, Hq, dtype=torch.float32, device=device)
        if buf_holder is not None:
            buf_holder._tq_lse_buf = lse

    # --- Stage 2: reduce across KV splits (unchanged) ---
    grid2 = (B, Hq)
    _fwd_kernel_stage2[grid2](
        mid_o,
        output,
        lse,
        seq_lens,
        mid_o.stride(0),
        mid_o.stride(1),
        mid_o.stride(2),
        output.stride(0),
        output.stride(1),
        lse.stride(0),
        NUM_KV_SPLITS=NUM_KV_SPLITS,
        BLOCK_DV=cfg["BLOCK_D"],
        Lv=D,
        num_warps=4,
        num_stages=2,
    )

    return output.to(query.dtype)