vllm.v1.attention.ops.flydsl_turboquant_decode ¶

def _detect_max_capture_B() -> int:
    """Probe vLLM compilation config for the largest cudagraph capture size.

    Returns max(cudagraph_capture_sizes) if available, else falls back to
    512 (vLLM's default cap). The segm-pool auto-grows if a larger batch is
    seen (see ``_SegmBufPool.get``).

    Used to size the segm-pool bucket once, so that all distinct B's
    (whether captured in a graph, or hit eagerly in mixed batches) share
    a single allocation per shape — eliminating the per-B growth that
    previously dominated dense-capture VRAM cost.
    """
    try:
        from vllm.config import get_current_vllm_config

        cfg = get_current_vllm_config()
        sizes = cfg.compilation_config.cudagraph_capture_sizes
        if sizes:
            return int(max(sizes))
    except Exception:  # noqa: BLE001
        pass
    return 512

def _hw_tr_enabled() -> bool:
    """Resolve the HW V-transpose build flag.

    ON for gfx950+ (ds_read_tr16_b64 is bit-exact vs baseline and faster
    on Qwen-class shapes), off elsewhere. Resolved from the GPU arch.
    """
    global _HW_TR_CACHED
    if _HW_TR_CACHED is not None:
        return _HW_TR_CACHED
    try:
        from flydsl.runtime.device import get_rocm_arch as _arch

        a = str(_arch() or "")
        _HW_TR_CACHED = a.startswith("gfx950")
    except Exception:  # noqa: BLE001
        _HW_TR_CACHED = False
    if _HW_TR_CACHED:
        logger.info_once("FlyDSL TQ: HW V transpose ON (default for gfx950+)")
    else:
        logger.info_once("FlyDSL TQ: HW V transpose OFF")
    return _HW_TR_CACHED

def flydsl_turboquant_decode_attention(
    query: torch.Tensor,  # [B, Hq, D] bf16/fp16
    kv_cache: torch.Tensor,  # [num_blocks, BS, Hk, slot_size_aligned]
    block_table: torch.Tensor,  # [B, max_blocks_per_seq] int32
    seq_lens: torch.Tensor,  # [B] 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,
    key_fp8: bool = False,
    norm_correction: bool = False,
    PiT: torch.Tensor | None = None,
    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:
    """SoA-decode-compatible launcher backed by the FlyDSL decode kernel.

    Constraints:
      * key_fp8 == False
      * mse_bits == 4
      * value_quant_bits == 4
      * centroids.numel() == 16
      * D == 128
      * block_size in {16, 32}
      * Hq // Hk in {6, 8, 16}
        - 8/16 → canonical tq_decode kernel (Qwen2.5-72B / Qwen3-32B)
        - 6    → tq_decode_gqa6 sibling kernel (MiniMax-M2.5)

    Sinks are NYI and silently ignored if set. norm_correction is honored
    implicitly via the pre-folded stored K-norm (see footer comment).
    """
    del mid_o_buf, lse_buf, key_packed_size, value_packed_size
    if not is_flydsl_available():
        raise RuntimeError(
            "FlyDSL decode requested but FlyDSL is not available (needs gfx950 "
            "+ importable FlyDSL); the SoA Triton decode is the fallback."
        )
    assert not key_fp8, "FlyDSL supports MSE-key path only"
    assert mse_bits == 4, f"FlyDSL expects mse_bits=4, got {mse_bits}"
    assert value_quant_bits == 4, (
        f"FlyDSL expects value_quant_bits=4, got {value_quant_bits}"
    )

    B, Hq, D = query.shape
    Hk = kv_cache.shape[2]
    block_size = kv_cache.shape[1]
    QG = Hq // Hk
    assert D == _TQ_MOD.HEAD_SIZE
    assert block_size in (16, 32), (
        f"FlyDSL supports kv_block_size 16 or 32, got {block_size}"
    )
    assert QG in (6, 8, 16), f"FlyDSL supports GQA factor 6, 8 or 16, got {QG}"
    if QG == 6 and _TQ_MOD_GQA6 is None:
        raise RuntimeError(
            "FlyDSL launcher: GQA-6 requested (MiniMax-class) but the "
            "tq_decode_gqa6 sibling module is not available. Update your "
            "FlyDSL checkout (must include kernels/tq_decode_gqa6.py); "
            "otherwise the SoA Triton decode is the fallback."
        )
    assert centroids.numel() == _TQ_MOD.N_CENTROIDS, (
        f"centroids.numel={centroids.numel()} != {_TQ_MOD.N_CENTROIDS}"
    )

    # ---- Per-layer cache for PiT_f32 + contiguous centroids ---
    # PiT and centroids are model constants (set once at layer warmup).
    # Avoid the per-decode-step transpose+cast and `.contiguous()` no-op
    # check by stashing the pre-cooked tensors on ``buf_holder`` (= layer).
    PiT_f32: torch.Tensor
    centroids_c: torch.Tensor
    if buf_holder is not None:
        PiT_f32 = getattr(buf_holder, "_tq_PiT_f32", None)
        if PiT_f32 is None:
            _PiT_src = PiT if PiT is not None else Pi.T.contiguous()
            PiT_f32 = (
                _PiT_src
                if _PiT_src.dtype == torch.float32
                else _PiT_src.to(torch.float32)
            )
            buf_holder._tq_PiT_f32 = PiT_f32
        centroids_c = getattr(buf_holder, "_tq_centroids_c", None)
        if centroids_c is None:
            centroids_c = centroids.contiguous()
            buf_holder._tq_centroids_c = centroids_c
    else:
        # Defensive fallback (unit tests may pass ``buf_holder=None``).
        _PiT_src = PiT if PiT is not None else Pi.T.contiguous()
        PiT_f32 = (
            _PiT_src if _PiT_src.dtype == torch.float32 else _PiT_src.to(torch.float32)
        )
        centroids_c = centroids.contiguous()

    # ---- Pooled q_rot (stable data_ptr post-capture) ------
    # On ROCm, fresh tensor allocations made AFTER HIP graph capture can land
    # in the graph memory pool and cause GPU memory faults when passed as kernel
    # pointer arguments in the eager mixed-batch path. We pool q_rot and the
    # fp32 intermediate the same way segm_out/output are pooled: one max-B
    # allocation per shape, reused across all B's.
    #
    # Note: pool_bufs is computed BELOW after num_partitions is determined;
    # q_rot needs pool_bufs["q_float"] and pool_bufs["q_rot"]. We compute the
    # rotation inline after the pool is fetched (see the pooled-buffers
    # section below).

    # ---- Partition count (FA2 split-KV) ----------------------------------
    #
    # FlyDSL runs INSIDE FULL cudagraph (TurboQuantMetadataBuilder._cudagraph_
    # support = UNIFORM_BATCH). That means the gridDim baked at capture
    # time MUST equal the gridDim at replay — the kernel launch parameters
    # are recorded into the captured graph. So `num_partitions` (which
    # becomes gridDim.z) MUST be derived from a stable source that produces
    # the same value at capture and at runtime.
    #
    # We use the worst-case context length implied by the block table's
    # allocation (block_table.shape[1] * block_size) as that source. It's
    # bounded by `max_model_len` (vLLM allocates the block table for the
    # configured max), and it's identical at capture and at runtime
    # because `block_table.shape[1]` is fixed once the model is loaded.
    # Per-tile OOB redirects + the masked FA-2 reducer ensure that for
    # short sequences the extra partitions contribute zero to the output.
    #
    # We always size from the worst case (block_table.shape[1] * block_size),
    # never from the per-step actual max_seq_len: a per-step size would make
    # the captured gridDim mismatch runtime and GPU-fault under cudagraph.
    kv_compute_block = _TQ_MOD.KV_COMPUTE_BLOCK
    worst_case_max_seq_len = int(block_table.shape[1]) * int(block_size)
    if max_seq_len <= 0:
        max_seq_len = worst_case_max_seq_len
    sizing_max_seq_len = worst_case_max_seq_len

    # ── Bounded num_partitions + internal tile-group looping ──
    # The kernel previously hardcoded one partition = 256 tokens (=
    # KV_COMPUTE_BLOCK), forcing num_partitions to scale linearly with
    # max_model_len, which made the cudagraph capture cost grow
    # prohibitively at long context. The new kernel takes
    # ``tile_groups_per_partition`` (TGPP); each CTA processes
    # ``TGPP * 16`` K-tiles = ``TGPP * KV_COMPUTE_BLOCK`` tokens with the
    # FA-2 online-softmax state accumulating across all of them. This
    # mirrors what the SoA Triton decode / HIP SoA-fusion already do.
    #
    # Strategy: cap num_partitions at ``MAX_PARTITIONS`` (default 32),
    # then derive TGPP so that ``num_partitions * TGPP * KV_COMPUTE_BLOCK
    # >= sizing_max_seq_len``.
    #
    # Examples (block_size=32 → worst_case max_bps*32):
    #   max_model_len=8K:   required=32, parts=32, TGPP=1  (no waste)
    #   max_model_len=16K:  required=64, parts=32, TGPP=2  (no waste)
    #   max_model_len=32K:  required=128, parts=32, TGPP=4 (no waste)
    #   max_model_len=64K:  required=256, parts=32, TGPP=8 (no waste)
    #   max_model_len=128K: required=512, parts=32, TGPP=16
    MAX_PARTITIONS = 32
    required_num_partitions = (
        sizing_max_seq_len + kv_compute_block - 1
    ) // kv_compute_block
    # max_num_kv_splits acts as a parallelism floor (ensure at least this
    # many CTAs along grid.z), but is itself capped by MAX_PARTITIONS.
    parallelism_floor = min(MAX_PARTITIONS, max(1, max_num_kv_splits))
    # Cap num_partitions at MAX_PARTITIONS, but never go below the floor.
    num_partitions_actual = max(
        parallelism_floor,
        min(MAX_PARTITIONS, required_num_partitions),
    )
    # Round up to next power of 2 for the Triton reducer's tl.arange(0, N)
    # constraint (Triton requires power-of-2 ≥ 2).
    num_partitions = max(2, triton.next_power_of_2(num_partitions_actual))
    # Now derive tile_groups_per_partition (TGPP) so total coverage
    # ``num_partitions * TGPP * KV_COMPUTE_BLOCK`` is >= sizing_max_seq_len.
    # This guarantees no work is dropped at runtime regardless of seq_len.
    # Round TGPP up to the next power-of-2 so the JIT cache key is bounded
    # to a small set of values (e.g. {1, 2, 4, 8, 16}) — limits the number
    # of distinct kernel binaries that need to be compiled across requests.
    _tgpp_required = max(
        1,
        (required_num_partitions + num_partitions - 1) // num_partitions,
    )
    tile_groups_per_partition = int(triton.next_power_of_2(_tgpp_required))

    # ---- Pooled buffers (no per-call cudaMalloc / memset) ----------
    # The FlyDSL kernel writes the FULL [B, Hk, P, QG, D] segm_out and the
    # FULL [B, Hk, P, QG] segm_max / segm_sum (running_max=-inf, sum=0 for
    # empty partitions are stored by the kernel itself), so uninitialized
    # buffers from the pool are safe to reuse.
    device = query.device
    pool_bufs = _SEGM_POOL.get(
        B,
        Hk,
        Hq,
        num_partitions,
        QG,
        D,
        device,
        query.dtype,
    )
    segm_out = pool_bufs["segm_out"]
    segm_max = pool_bufs["segm_max"]
    segm_sum = pool_bufs["segm_sum"]
    if output_buf is None:
        output = pool_bufs["output"]
    else:
        output = output_buf[:B] if output_buf.shape[0] != B else output_buf

    # ---- q rotation: GEMM q @ PiT via stable pooled buffers -------------
    _q_float = pool_bufs["q_float"]  # [B, Hq, D] fp32, stable
    _q_rot_f32 = pool_bufs["q_rot_fp32"]  # [B, Hq, D] fp32, stable mm output
    _q_rot_out = pool_bufs["q_rot"]  # [B, Hq, D] query.dtype, stable
    _q_float.copy_(query)  # bf16 → fp32 in-place, no alloc
    # mm into stable fp32 buffer via out= to avoid fresh allocations.
    # PiT_f32 is cached on buf_holder (stable ptr since first layer warmup).
    torch.mm(
        _q_float.view(B * Hq, D),
        PiT_f32,
        out=_q_rot_f32.view(B * Hq, D),  # in-place into pool buffer
    )
    _q_rot_out.copy_(_q_rot_f32)  # fp32 → bf16, into stable buf
    q_for_kernel = _q_rot_out  # stable ptr for kernel launch

    # ---- FlyDSL kernel launch -------------------------------------------
    max_bps = int(block_table.shape[1])
    use_hw_tr = _hw_tr_enabled()
    launch = _get_kernel(
        Hk,
        num_partitions,
        max_bps,
        scale,
        QG,
        block_size,
        use_hw_v_transpose=use_hw_tr,
        num_seqs_hint=int(B),
        tile_groups_per_partition=int(tile_groups_per_partition),
    )
    # Zero-overhead one-shot info log (replaces logger.info_once which
    # hashes its format string on every call to dedup).
    global _LOG_INVOKED_ONCE, _LOG_SINKS_WARNED, _LOG_NORM_WARNED
    if not _LOG_INVOKED_ONCE:
        _LOG_INVOKED_ONCE = True
        logger.info(
            "FlyDSL launcher invoked (UNIFORM_BATCH cudagraph): "
            "B=%d Hk=%d Hq=%d D=%d QG=%d num_partitions=%d (actual=%d, "
            "cap=%d) TGPP=%d max_bps=%d block_size=%d max_seq_len=%d "
            "hw_v_transpose=%s (coverage=%d tokens, worst_case=%d tokens)",
            B,
            Hk,
            Hq,
            D,
            QG,
            num_partitions,
            num_partitions_actual,
            MAX_PARTITIONS,
            tile_groups_per_partition,
            max_bps,
            int(block_size),
            int(max_seq_len),
            use_hw_tr,
            num_partitions * tile_groups_per_partition * kv_compute_block,
            worst_case_max_seq_len,
        )
    launch(
        segm_out,
        segm_sum,
        segm_max,
        q_for_kernel,
        kv_cache,
        centroids_c,
        block_table,
        seq_lens,
        B,
        Hk,
        num_partitions,
        torch.cuda.current_stream(),
    )

    # ---- Reduce partitions -> [B, Hq, D] --------------------------------
    _reduce_partitions[(B, Hq)](
        output_ptr=output,
        segm_out_ptr=segm_out,
        segm_max_ptr=segm_max,
        segm_sum_ptr=segm_sum,
        out_stride_n=output.stride(0),
        out_stride_h=output.stride(1),
        NUM_KV_HEADS=Hk,
        QG=QG,
        NUM_PARTS=num_partitions,
        HEAD_SIZE=D,
    )
    if sinks is not None and not _LOG_SINKS_WARNED:
        _LOG_SINKS_WARNED = True
        logger.warning(
            "FlyDSL launcher: sinks ignored (NYI). Disable sinks if they "
            "are required for correctness."
        )
    # ── norm_correction is honored IMPLICITLY ───────────────────────────
    # When the model was stored with norm_correction=True (the *_nc presets
    # turboquant_4bit_nc / k3v4_nc / 3bit_nc), the per-token K-norm scalar
    # was pre-folded to ||k_t|| / ||c_t|| at store time by the TurboQuant
    # store path. The decode kernel just multiplies `c_vals * stored_knorm`,
    # which then equals `(c_vals / ||c_t||) * ||k_t||` — exactly the
    # unit-norm-renormalized centroid times the original key norm. The
    # Triton decode does the identical multiply, so both decoders honor
    # norm_correction equivalently and there is nothing to "do" at decode
    # time. The launcher's `norm_correction` arg is only kept for API
    # parity; we emit a one-shot INFO log to make the contract explicit.
    if norm_correction and not _LOG_NORM_WARNED:
        _LOG_NORM_WARNED = True
        logger.info(
            "FlyDSL launcher: norm_correction honored implicitly via "
            "pre-folded stored K-norm (cf. triton_turboquant_store step 3); "
            "no decode-time work required, identical to the SoA decode behavior."
        )
    return output

def is_flydsl_available() -> bool:
    """Return True iff on gfx950 and FlyDSL imports + kernel module load.

    The GQA-6 sibling kernel (``kernels.tq_decode_gqa6``) is imported
    best-effort: if it's missing (older FlyDSL checkout that pre-dates
    MiniMax support) the canonical Qwen path stays fully functional and
    only GQA-6 dispatches will fail with a clear error at launch time.
    """
    global _FLYDSL_AVAILABLE, _TQ_MOD, _TQ_MOD_GQA6
    global _FLYC, _FX, _TYPING_T, _CC, _IR
    if _FLYDSL_AVAILABLE is not None:
        return _FLYDSL_AVAILABLE
    try:
        # gfx950-only kernel (CDNA4 intrinsics); gate before importing FlyDSL.
        from vllm.platforms.rocm import on_gfx950

        if not on_gfx950():
            _FLYDSL_AVAILABLE = False
            return _FLYDSL_AVAILABLE
        import flydsl.compiler as flyc  # noqa: F401
        import flydsl.expr as fx  # noqa: F401
        from flydsl._mlir import ir
        from flydsl.compiler.kernel_function import CompilationContext
        from flydsl.expr.typing import T  # noqa: F401

        # Kernel is vendored in-tree (ships with vLLM); FlyDSL provides only
        # the compiler/runtime framework (imported above).
        from vllm.v1.attention.ops.flydsl_kernels import tq_decode as tq_mod

        _FLYC = flyc
        _FX = fx
        _TYPING_T = T
        _CC = CompilationContext
        _IR = ir
        _TQ_MOD = tq_mod
        _FLYDSL_AVAILABLE = True
        logger.info_once("FlyDSL TQ decode launcher: available")
    except Exception as ex:  # noqa: BLE001
        _FLYDSL_AVAILABLE = False
        logger.warning_once(
            "FlyDSL TQ decode launcher: unavailable (%s). "
            "Falling back to SoA Triton decode.",
            ex,
        )
        return _FLYDSL_AVAILABLE
    # Best-effort GQA-6 sibling import (does NOT gate Qwen availability).
    try:
        from vllm.v1.attention.ops.flydsl_kernels import (
            tq_decode_gqa6 as tq_mod_gqa6,
        )

        _TQ_MOD_GQA6 = tq_mod_gqa6
        logger.info_once("FlyDSL TQ decode GQA-6 sibling: available (MiniMax-class)")
    except Exception as ex:  # noqa: BLE001
        _TQ_MOD_GQA6 = None
        logger.info_once(
            "FlyDSL TQ decode GQA-6 sibling: not available (%s); "
            "GQA-6 models will fall back to SoA Triton decode.",
            ex,
        )
    return _FLYDSL_AVAILABLE

def is_flydsl_gqa6_available() -> bool:
    """True iff the optional GQA-6 sibling kernel module loaded.

    Used by the eligibility gate in turboquant_attn.py to decide whether
    a layer with num_kv_groups==6 can run on FlyDSL or must fall back
    to the SoA Triton decode.
    """
    if _FLYDSL_AVAILABLE is None:
        is_flydsl_available()
    return _TQ_MOD_GQA6 is not None