vllm.v1.attention.ops.flydsl_turboquant_decode

FlyDSL TurboQuant decode launcher (vLLM-side).

Drop-in replacement for triton_turboquant_decode_attention_soa for the TQ decode profile HEAD_SIZE=128, MSE_BITS=4 K, VQB=4 V, N_CENTROIDS=16, BLOCK_SIZE in {16, 32}, GQA group in {6, 8, 16} (one kernel, kernels.tq_decode).

Auto-selected on gfx950 when FlyDSL is importable; falls back to the SoA Triton decode otherwise (wrong arch, missing build tree).

Architecture

1. Q rotation: q_rot = (query.float() @ PiT).bfloat16() — the same launcher-side rocBLAS GEMM the SoA Triton decode uses.
2. FlyDSL kernel writes per-partition outputs into [N, Hk, P, QG, D] bf16 + [N, Hk, P, QG] fp32 max/sum buffers.
3. A small Triton reducer combines partitions in the kernel's native layout (no permute/cast), writing the final [N, Hq, D] output.

Kernel module is built once per (num_kv_heads, num_partitions, max_blocks_per_seq, scale) shape and cached.

Functions:

• flydsl_turboquant_decode_attention –

  SoA-decode-compatible launcher backed by the FlyDSL decode kernel.

• is_flydsl_available –

  Return True iff on gfx950 and FlyDSL imports + kernel module load.

• is_flydsl_gqa6_available –

  GQA-6 runs on the same kernel; kept for the turboquant_attn.py gate.

_SegmBufPool

Single-bucket buffer pool for segm_out/segm_max/segm_sum/output.

Each unique shape signature (Hk, Hq, num_partitions, QG, D, device, dtype) is backed by ONE allocation sized at max_B (= max captured cudagraph batch size, or env override). Per-call get(B, ...) returns [:B] views — same data_ptr base, narrower N dim. The kernel grid is (B, ...) so it only ever writes the first B rows of segm_out / segm_max / segm_sum; downstream readers (Triton reducer + caller) only consume the first B rows of output.

Why a single bucket: cudagraph records launch pointers at capture time. With per-B allocations, every captured size triggered a fresh allocation, so total VRAM scaled with the number of captured sizes. With single-bucket, the pool tops out at one max_B-sized buffer per shape, shared across ALL captured sizes and ALL eager mixed-batch calls. The data_ptr is stable from first allocation onward, so cudagraph capture is happy across all captured B's.

Eliminates per-decode-step cudaMalloc (× 4) plus the device-side memset kernels behind torch.full(-inf) and torch.zeros. The kernel always writes the FULL [B, Hk, P, QG, D] slice and every (n, kv_h, p, qg) position of segm_max / segm_sum (it stores -inf and 0 for empty partitions itself), so uninitialized bytes from the pool are safe to reuse for any B ≤ B_bucket.

Source code in vllm/v1/attention/ops/flydsl_turboquant_decode.py

class _SegmBufPool:
    """Single-bucket buffer pool for segm_out/segm_max/segm_sum/output.

    Each unique shape signature (Hk, Hq, num_partitions, QG, D, device,
    dtype) is backed by ONE allocation sized at ``max_B`` (= max captured
    cudagraph batch size, or env override). Per-call ``get(B, ...)``
    returns ``[:B]`` views — same data_ptr base, narrower N dim. The
    kernel grid is ``(B, ...)`` so it only ever writes the first B rows
    of segm_out / segm_max / segm_sum; downstream readers (Triton
    reducer + caller) only consume the first B rows of output.

    Why a single bucket: cudagraph records launch pointers at capture
    time. With per-B allocations, every captured size triggered a fresh
    allocation, so total VRAM scaled with the number of captured sizes.
    With single-bucket, the pool tops out at one max_B-sized buffer per
    shape, shared across ALL captured sizes and ALL eager mixed-batch
    calls. The data_ptr is stable from first allocation onward, so
    cudagraph capture is happy across all captured B's.

    Eliminates per-decode-step ``cudaMalloc`` (× 4) plus the device-side
    memset kernels behind ``torch.full(-inf)`` and ``torch.zeros``.
    The kernel always writes the FULL ``[B, Hk, P, QG, D]`` slice and
    every (n, kv_h, p, qg) position of segm_max / segm_sum (it stores
    ``-inf`` and ``0`` for empty partitions itself), so uninitialized
    bytes from the pool are safe to reuse for any B ≤ B_bucket.
    """

    __slots__ = ("_bufs", "_max_B")

    def __init__(self) -> None:
        self._bufs: dict[tuple, dict[str, torch.Tensor]] = {}
        self._max_B: int | None = None  # lazy-init on first get()

    def get(
        self,
        B: int,
        Hk: int,
        Hq: int,
        num_partitions: int,
        QG: int,
        D: int,
        device: torch.device,
        q_dtype: torch.dtype,
    ) -> dict[str, torch.Tensor]:
        if self._max_B is None:
            self._max_B = _detect_max_capture_B()
        # B may exceed the detected max if user set max-num-seqs higher
        # than the largest captured size — still safe, we'll grow once.
        # Note: growing AFTER cudagraph capture would invalidate captured
        # pointers, so this should only happen during warmup or in
        # always-eager configs.
        B_bucket = max(self._max_B, int(B))
        key = (
            int(Hk),
            int(Hq),
            int(num_partitions),
            int(QG),
            int(D),
            str(device),
            q_dtype,
        )
        bufs = self._bufs.get(key)
        if bufs is None or bufs["segm_out"].shape[0] < B_bucket:
            if bufs is not None and bufs["segm_out"].shape[0] < B_bucket:
                # First-time grow: warn so user knows cudagraphs may be
                # invalidated. This should only happen during warmup or in
                # always-eager configs (a batch larger than any captured
                # cudagraph size).
                logger.warning_once(
                    "FlyDSL _SegmBufPool: growing bucket from %d to %d "
                    "(B=%d). If you see this AFTER cudagraph warmup, the "
                    "previously-captured graphs hold stale pointers and "
                    "will GPU-fault (max cudagraph capture size < B).",
                    bufs["segm_out"].shape[0],
                    B_bucket,
                    B,
                )
            bufs = {
                "segm_out": torch.empty(
                    (B_bucket, Hk, num_partitions, QG, D),
                    dtype=torch.bfloat16,
                    device=device,
                ),
                "segm_max": torch.empty(
                    (B_bucket, Hk, num_partitions, QG),
                    dtype=torch.float32,
                    device=device,
                ),
                "segm_sum": torch.empty(
                    (B_bucket, Hk, num_partitions, QG),
                    dtype=torch.float32,
                    device=device,
                ),
                "output": torch.empty(
                    (B_bucket, Hq, D),
                    dtype=q_dtype,
                    device=device,
                ),
                # Stable pre-allocated buffer for the rotated query
                # tensor. On ROCm, fresh torch.empty / matmul-result allocations
                # after HIP graph capture can land in the graph memory pool and
                # cause GPU memory faults when passed as kernel arguments in the
                # eager mixed-batch path. By pooling q_rot here (same pattern as
                # segm_out/output above), the data_ptr is stable from first
                # allocation onward — same for all captured B's and eager calls.
                # Always bf16: the kernel reads Q as bf16 whatever the model
                # dtype (fp16 models included).
                "q_rot": torch.empty(
                    (B_bucket, Hq, D),
                    dtype=torch.bfloat16,
                    device=device,
                ),
                # q_float: fp32 copy of query (query dtype→fp32 for mm input).
                # Pooled to avoid fresh allocations post-capture.
                "q_float": torch.empty(
                    (B_bucket, Hq, D),
                    dtype=torch.float32,
                    device=device,
                ),
                # q_rot_fp32: fp32 output of the rotation mm (before bf16 cast).
                # Using torch.mm(..., out=this) eliminates the fresh mm-result
                # allocation that would otherwise land in the HIP graph pool.
                "q_rot_fp32": torch.empty(
                    (B_bucket, Hq, D),
                    dtype=torch.float32,
                    device=device,
                ),
            }
            self._bufs[key] = bufs
            self._max_B = B_bucket
            logger.info_once(
                "FlyDSL _SegmBufPool: allocated single-bucket "
                "shape=(Hk=%d, Hq=%d, P=%d, QG=%d, D=%d, dtype=%s) "
                "B_bucket=%d (covers all captured + eager B's). "
                "VRAM = %.1f MiB / shape.",
                Hk,
                Hq,
                num_partitions,
                QG,
                D,
                q_dtype,
                B_bucket,
                sum(t.numel() * t.element_size() for t in bufs.values()) / (1 << 20),
            )
        # Return [:B] views — same data_ptr base, narrower N dim.
        return {
            "segm_out": bufs["segm_out"][:B],
            "segm_max": bufs["segm_max"][:B],
            "segm_sum": bufs["segm_sum"][:B],
            "output": bufs["output"][:B],
            "q_rot": bufs["q_rot"][:B],
            "q_float": bufs["q_float"][:B],
            "q_rot_fp32": bufs["q_rot_fp32"][:B],
        }

    def stats(self) -> dict[str, int]:
        return {
            "shapes": len(self._bufs),
            "max_B": int(self._max_B or 0),
            "bytes": sum(
                sum(t.numel() * t.element_size() for t in d.values())
                for d in self._bufs.values()
            ),
        }

_detect_max_capture_B()

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.

Source code in vllm/v1/attention/ops/flydsl_turboquant_decode.py

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

flydsl_turboquant_decode_attention(query, kv_cache, block_table, seq_lens, Pi, centroids, scale, mse_bits, key_packed_size, value_quant_bits, value_packed_size, key_fp8=False, norm_correction=False, PiT=None, max_seq_len=0, mid_o_buf=None, output_buf=None, lse_buf=None, buf_holder=None, max_num_kv_splits=32, sinks=None)

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}

Sinks are NYI and silently ignored if set. norm_correction is honored implicitly via the pre-folded stored K-norm (see footer comment).

Source code in vllm/v1/attention/ops/flydsl_turboquant_decode.py

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}

    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}"
    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] bf16, stable
    _q_float.copy_(query)  # query dtype → 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])
    launch = _get_kernel(
        Hk,
        num_partitions,
        max_bps,
        scale,
        QG,
        block_size,
        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 "
            "(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),
            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

is_flydsl_available()

Return True iff on gfx950 and FlyDSL imports + kernel module load.

Source code in vllm/v1/attention/ops/flydsl_turboquant_decode.py

def is_flydsl_available() -> bool:
    """Return True iff on gfx950 and FlyDSL imports + kernel module load."""
    global _FLYDSL_AVAILABLE, _TQ_MOD
    global _FLYC, _FX
    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

        # 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
        _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

is_flydsl_gqa6_available()

GQA-6 runs on the same kernel; kept for the turboquant_attn.py gate.

Source code in vllm/v1/attention/ops/flydsl_turboquant_decode.py

def is_flydsl_gqa6_available() -> bool:
    """GQA-6 runs on the same kernel; kept for the turboquant_attn.py gate."""
    return is_flydsl_available()