| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FlyDSL TurboQuant decode launcher (vLLM-side). |
| 4 | |
| 5 | Drop-in replacement for ``triton_turboquant_decode_attention_soa`` for the |
| 6 | TQ decode profile HEAD_SIZE=128, MSE_BITS=4 K, VQB=4 V, N_CENTROIDS=16, |
| 7 | BLOCK_SIZE in {16, 32}. |
| 8 | |
| 9 | Per-GQA kernel dispatch: |
| 10 | * GQA group ∈ {8, 16} → canonical kernels.tq_decode (Qwen-class) |
| 11 | * GQA group == 6 → kernels.tq_decode_gqa6 sibling (MiniMax-M2.5) |
| 12 | |
| 13 | Auto-selected on gfx950 when FlyDSL is importable; falls back to the SoA Triton |
| 14 | decode otherwise (wrong arch, missing build tree). The GQA-6 sibling is imported |
| 15 | best-effort: missing it does not affect Qwen GQA-{8,16} paths. |
| 16 | |
| 17 | Architecture: |
| 18 | 1. Q rotation: ``q_rot = (query.float() @ PiT).bfloat16()`` — the same |
| 19 | launcher-side rocBLAS GEMM the SoA Triton decode uses. |
| 20 | 2. FlyDSL kernel writes per-partition outputs into |
| 21 | ``[N, Hk, P, QG, D]`` bf16 + ``[N, Hk, P, QG]`` fp32 max/sum buffers. |
| 22 | 3. A small Triton reducer combines partitions in the kernel's native layout |
| 23 | (no permute/cast), writing the final ``[N, Hq, D]`` output. |
| 24 | |
| 25 | Kernel module is built once per ``(num_kv_heads, num_partitions, |
| 26 | max_blocks_per_seq, scale)`` shape and cached. |
| 27 | """ |
| 28 | |
| 29 | from __future__ import annotations |
| 30 | |
| 31 | from typing import Any |
| 32 | |
| 33 | import torch |
| 34 | |
| 35 | from vllm.logger import init_logger |
| 36 | from vllm.triton_utils import tl, triton |
| 37 | |
| 38 | logger = init_logger(__name__) |
| 39 | |
| 40 | # -- FlyDSL import bootstrap -------------------------------------------------- |
| 41 | # FlyDSL is imported directly and must be importable (pip-installed, or on |
| 42 | # PYTHONPATH). Import failures are caught below and fall back to the SoA Triton |
| 43 | # decode, so vLLM startup is unaffected on hosts without FlyDSL. |
| 44 | |
| 45 | |
| 46 | _FLYDSL_AVAILABLE: bool | None = None |
| 47 | _TQ_MOD: Any = None # kernels.tq_decode module (Qwen GQA-{8,16}) |
| 48 | _TQ_MOD_GQA6: Any = None # kernels.tq_decode_gqa6 module (MiniMax GQA-6, optional) |
| 49 | _FLYC: Any = None # flydsl.compiler |
| 50 | _FX: Any = None # flydsl.expr |
| 51 | _TYPING_T: Any = None |
| 52 | _CC: Any = None # CompilationContext |
| 53 | _IR: Any = None |
| 54 | |
| 55 | |
| 56 | def is_flydsl_available() -> bool: |
| 57 | """Return True iff on gfx950 and FlyDSL imports + kernel module load. |
| 58 | |
| 59 | The GQA-6 sibling kernel (``kernels.tq_decode_gqa6``) is imported |
| 60 | best-effort: if it's missing (older FlyDSL checkout that pre-dates |
| 61 | MiniMax support) the canonical Qwen path stays fully functional and |
| 62 | only GQA-6 dispatches will fail with a clear error at launch time. |
| 63 | """ |
| 64 | global _FLYDSL_AVAILABLE, _TQ_MOD, _TQ_MOD_GQA6 |
| 65 | global _FLYC, _FX, _TYPING_T, _CC, _IR |
| 66 | if _FLYDSL_AVAILABLE is not None: |
| 67 | return _FLYDSL_AVAILABLE |
| 68 | try: |
| 69 | # gfx950-only kernel (CDNA4 intrinsics); gate before importing FlyDSL. |
| 70 | from vllm.platforms.rocm import on_gfx950 |
| 71 | |
| 72 | if not on_gfx950(): |
| 73 | _FLYDSL_AVAILABLE = False |
| 74 | return _FLYDSL_AVAILABLE |
| 75 | import flydsl.compiler as flyc # noqa: F401 |
| 76 | import flydsl.expr as fx # noqa: F401 |
| 77 | from flydsl._mlir import ir |
| 78 | from flydsl.compiler.kernel_function import CompilationContext |
| 79 | from flydsl.expr.typing import T # noqa: F401 |
| 80 | |
| 81 | # Kernel is vendored in-tree (ships with vLLM); FlyDSL provides only |
| 82 | # the compiler/runtime framework (imported above). |
| 83 | from vllm.v1.attention.ops.flydsl_kernels import tq_decode as tq_mod |
| 84 | |
| 85 | _FLYC = flyc |
| 86 | _FX = fx |
| 87 | _TYPING_T = T |
| 88 | _CC = CompilationContext |
| 89 | _IR = ir |
| 90 | _TQ_MOD = tq_mod |
| 91 | _FLYDSL_AVAILABLE = True |
| 92 | logger.info_once("FlyDSL TQ decode launcher: available") |
| 93 | except Exception as ex: # noqa: BLE001 |
| 94 | _FLYDSL_AVAILABLE = False |
| 95 | logger.warning_once( |
| 96 | "FlyDSL TQ decode launcher: unavailable (%s). " |
| 97 | "Falling back to SoA Triton decode.", |
| 98 | ex, |
| 99 | ) |
| 100 | return _FLYDSL_AVAILABLE |
| 101 | # Best-effort GQA-6 sibling import (does NOT gate Qwen availability). |
| 102 | try: |
| 103 | from vllm.v1.attention.ops.flydsl_kernels import ( |
| 104 | tq_decode_gqa6 as tq_mod_gqa6, |
| 105 | ) |
| 106 | |
| 107 | _TQ_MOD_GQA6 = tq_mod_gqa6 |
| 108 | logger.info_once("FlyDSL TQ decode GQA-6 sibling: available (MiniMax-class)") |
| 109 | except Exception as ex: # noqa: BLE001 |
| 110 | _TQ_MOD_GQA6 = None |
| 111 | logger.info_once( |
| 112 | "FlyDSL TQ decode GQA-6 sibling: not available (%s); " |
| 113 | "GQA-6 models will fall back to SoA Triton decode.", |
| 114 | ex, |
| 115 | ) |
| 116 | return _FLYDSL_AVAILABLE |
| 117 | |
| 118 | |
| 119 | def is_flydsl_gqa6_available() -> bool: |
| 120 | """True iff the optional GQA-6 sibling kernel module loaded. |
| 121 | |
| 122 | Used by the eligibility gate in turboquant_attn.py to decide whether |
| 123 | a layer with num_kv_groups==6 can run on FlyDSL or must fall back |
| 124 | to the SoA Triton decode. |
| 125 | """ |
| 126 | if _FLYDSL_AVAILABLE is None: |
| 127 | is_flydsl_available() |
| 128 | return _TQ_MOD_GQA6 is not None |
| 129 | |
| 130 | |
| 131 | # -- Kernel module cache ------------------------------------------------------- |
| 132 | # Compiling a FlyDSL kernel is expensive; cache by the constexprs that |
| 133 | # parameterize stride math: num_kv_heads, num_partitions, max_blocks_per_seq, |
| 134 | # scale. ``num_seqs`` is irrelevant to the kernel body so we omit it. |
| 135 | _KERN_CACHE: dict[tuple, Any] = {} |
| 136 | |
| 137 | |
| 138 | # -- One-shot log guards ------------------------------------------------------ |
| 139 | # Replaces vllm `logger.{info,warning}_once` (which hashes the format string |
| 140 | # every call to dedup) with a true zero-overhead bool guard. |
| 141 | _LOG_INVOKED_ONCE: bool = False |
| 142 | _LOG_SINKS_WARNED: bool = False |
| 143 | _LOG_NORM_WARNED: bool = False |
| 144 | |
| 145 | |
| 146 | def _detect_max_capture_B() -> int: |
| 147 | """Probe vLLM compilation config for the largest cudagraph capture size. |
| 148 | |
| 149 | Returns max(cudagraph_capture_sizes) if available, else falls back to |
| 150 | 512 (vLLM's default cap). The segm-pool auto-grows if a larger batch is |
| 151 | seen (see ``_SegmBufPool.get``). |
| 152 | |
| 153 | Used to size the segm-pool bucket once, so that all distinct B's |
| 154 | (whether captured in a graph, or hit eagerly in mixed batches) share |
| 155 | a single allocation per shape — eliminating the per-B growth that |
| 156 | previously dominated dense-capture VRAM cost. |
| 157 | """ |
| 158 | try: |
| 159 | from vllm.config import get_current_vllm_config |
| 160 | |
| 161 | cfg = get_current_vllm_config() |
| 162 | sizes = cfg.compilation_config.cudagraph_capture_sizes |
| 163 | if sizes: |
| 164 | return int(max(sizes)) |
| 165 | except Exception: # noqa: BLE001 |
| 166 | pass |
| 167 | return 512 |
| 168 | |
| 169 | |
| 170 | class _SegmBufPool: |
| 171 | """Single-bucket buffer pool for segm_out/segm_max/segm_sum/output. |
| 172 | |
| 173 | Each unique shape signature (Hk, Hq, num_partitions, QG, D, device, |
| 174 | dtype) is backed by ONE allocation sized at ``max_B`` (= max captured |
| 175 | cudagraph batch size, or env override). Per-call ``get(B, ...)`` |
| 176 | returns ``[:B]`` views — same data_ptr base, narrower N dim. The |
| 177 | kernel grid is ``(B, ...)`` so it only ever writes the first B rows |
| 178 | of segm_out / segm_max / segm_sum; downstream readers (Triton |
| 179 | reducer + caller) only consume the first B rows of output. |
| 180 | |
| 181 | Why a single bucket: cudagraph records launch pointers at capture |
| 182 | time. With per-B allocations, every captured size triggered a fresh |
| 183 | allocation, so total VRAM scaled with the number of captured sizes. |
| 184 | With single-bucket, the pool tops out at one max_B-sized buffer per |
| 185 | shape, shared across ALL captured sizes and ALL eager mixed-batch |
| 186 | calls. The data_ptr is stable from first allocation onward, so |
| 187 | cudagraph capture is happy across all captured B's. |
| 188 | |
| 189 | Eliminates per-decode-step ``cudaMalloc`` (× 4) plus the device-side |
| 190 | memset kernels behind ``torch.full(-inf)`` and ``torch.zeros``. |
| 191 | The kernel always writes the FULL ``[B, Hk, P, QG, D]`` slice and |
| 192 | every (n, kv_h, p, qg) position of segm_max / segm_sum (it stores |
| 193 | ``-inf`` and ``0`` for empty partitions itself), so uninitialized |
| 194 | bytes from the pool are safe to reuse for any B ≤ B_bucket. |
| 195 | """ |
| 196 | |
| 197 | __slots__ = ("_bufs", "_max_B") |
| 198 | |
| 199 | def __init__(self) -> None: |
| 200 | self._bufs: dict[tuple, dict[str, torch.Tensor]] = {} |
| 201 | self._max_B: int | None = None # lazy-init on first get() |
| 202 | |
| 203 | def get( |
| 204 | self, |
| 205 | B: int, |
| 206 | Hk: int, |
| 207 | Hq: int, |
| 208 | num_partitions: int, |
| 209 | QG: int, |
| 210 | D: int, |
| 211 | device: torch.device, |
| 212 | q_dtype: torch.dtype, |
| 213 | ) -> dict[str, torch.Tensor]: |
| 214 | if self._max_B is None: |
| 215 | self._max_B = _detect_max_capture_B() |
| 216 | # B may exceed the detected max if user set max-num-seqs higher |
| 217 | # than the largest captured size — still safe, we'll grow once. |
| 218 | # Note: growing AFTER cudagraph capture would invalidate captured |
| 219 | # pointers, so this should only happen during warmup or in |
| 220 | # always-eager configs. |
| 221 | B_bucket = max(self._max_B, int(B)) |
| 222 | key = ( |
| 223 | int(Hk), |
| 224 | int(Hq), |
| 225 | int(num_partitions), |
| 226 | int(QG), |
| 227 | int(D), |
| 228 | str(device), |
| 229 | q_dtype, |
| 230 | ) |
| 231 | bufs = self._bufs.get(key) |
| 232 | if bufs is None or bufs["segm_out"].shape[0] < B_bucket: |
| 233 | if bufs is not None and bufs["segm_out"].shape[0] < B_bucket: |
| 234 | # First-time grow: warn so user knows cudagraphs may be |
| 235 | # invalidated. This should only happen during warmup or in |
| 236 | # always-eager configs (a batch larger than any captured |
| 237 | # cudagraph size). |
| 238 | logger.warning_once( |
| 239 | "FlyDSL _SegmBufPool: growing bucket from %d to %d " |
| 240 | "(B=%d). If you see this AFTER cudagraph warmup, the " |
| 241 | "previously-captured graphs hold stale pointers and " |
| 242 | "will GPU-fault (max cudagraph capture size < B).", |
| 243 | bufs["segm_out"].shape[0], |
| 244 | B_bucket, |
| 245 | B, |
| 246 | ) |
| 247 | bufs = { |
| 248 | "segm_out": torch.empty( |
| 249 | (B_bucket, Hk, num_partitions, QG, D), |
| 250 | dtype=torch.bfloat16, |
| 251 | device=device, |
| 252 | ), |
| 253 | "segm_max": torch.empty( |
| 254 | (B_bucket, Hk, num_partitions, QG), |
| 255 | dtype=torch.float32, |
| 256 | device=device, |
| 257 | ), |
| 258 | "segm_sum": torch.empty( |
| 259 | (B_bucket, Hk, num_partitions, QG), |
| 260 | dtype=torch.float32, |
| 261 | device=device, |
| 262 | ), |
| 263 | "output": torch.empty( |
| 264 | (B_bucket, Hq, D), |
| 265 | dtype=q_dtype, |
| 266 | device=device, |
| 267 | ), |
| 268 | # Stable pre-allocated buffer for the rotated query |
| 269 | # tensor. On ROCm, fresh torch.empty / matmul-result allocations |
| 270 | # after HIP graph capture can land in the graph memory pool and |
| 271 | # cause GPU memory faults when passed as kernel arguments in the |
| 272 | # eager mixed-batch path. By pooling q_rot here (same pattern as |
| 273 | # segm_out/output above), the data_ptr is stable from first |
| 274 | # allocation onward — same for all captured B's and eager calls. |
| 275 | "q_rot": torch.empty( |
| 276 | (B_bucket, Hq, D), |
| 277 | dtype=q_dtype, |
| 278 | device=device, |
| 279 | ), |
| 280 | # q_float: fp32 copy of query (bf16→fp32 for mm input). |
| 281 | # Pooled to avoid fresh allocations post-capture. |
| 282 | "q_float": torch.empty( |
| 283 | (B_bucket, Hq, D), |
| 284 | dtype=torch.float32, |
| 285 | device=device, |
| 286 | ), |
| 287 | # q_rot_fp32: fp32 output of the rotation mm (before bf16 cast). |
| 288 | # Using torch.mm(..., out=this) eliminates the fresh mm-result |
| 289 | # allocation that would otherwise land in the HIP graph pool. |
| 290 | "q_rot_fp32": torch.empty( |
| 291 | (B_bucket, Hq, D), |
| 292 | dtype=torch.float32, |
| 293 | device=device, |
| 294 | ), |
| 295 | } |
| 296 | self._bufs[key] = bufs |
| 297 | self._max_B = B_bucket |
| 298 | logger.info_once( |
| 299 | "FlyDSL _SegmBufPool: allocated single-bucket " |
| 300 | "shape=(Hk=%d, Hq=%d, P=%d, QG=%d, D=%d, dtype=%s) " |
| 301 | "B_bucket=%d (covers all captured + eager B's). " |
| 302 | "VRAM = %.1f MiB / shape.", |
| 303 | Hk, |
| 304 | Hq, |
| 305 | num_partitions, |
| 306 | QG, |
| 307 | D, |
| 308 | q_dtype, |
| 309 | B_bucket, |
| 310 | sum(t.numel() * t.element_size() for t in bufs.values()) / (1 << 20), |
| 311 | ) |
| 312 | # Return [:B] views — same data_ptr base, narrower N dim. |
| 313 | return { |
| 314 | "segm_out": bufs["segm_out"][:B], |
| 315 | "segm_max": bufs["segm_max"][:B], |
| 316 | "segm_sum": bufs["segm_sum"][:B], |
| 317 | "output": bufs["output"][:B], |
| 318 | "q_rot": bufs["q_rot"][:B], |
| 319 | "q_float": bufs["q_float"][:B], |
| 320 | "q_rot_fp32": bufs["q_rot_fp32"][:B], |
| 321 | } |
| 322 | |
| 323 | def stats(self) -> dict[str, int]: |
| 324 | return { |
| 325 | "shapes": len(self._bufs), |
| 326 | "max_B": int(self._max_B or 0), |
| 327 | "bytes": sum( |
| 328 | sum(t.numel() * t.element_size() for t in d.values()) |
| 329 | for d in self._bufs.values() |
| 330 | ), |
| 331 | } |
| 332 | |
| 333 | |
| 334 | _SEGM_POOL = _SegmBufPool() |
| 335 | |
| 336 | |
| 337 | _HW_TR_CACHED: bool | None = None |
| 338 | |
| 339 | |
| 340 | def _hw_tr_enabled() -> bool: |
| 341 | """Resolve the HW V-transpose build flag. |
| 342 | |
| 343 | ON for gfx950+ (ds_read_tr16_b64 is bit-exact vs baseline and faster |
| 344 | on Qwen-class shapes), off elsewhere. Resolved from the GPU arch. |
| 345 | """ |
| 346 | global _HW_TR_CACHED |
| 347 | if _HW_TR_CACHED is not None: |
| 348 | return _HW_TR_CACHED |
| 349 | try: |
| 350 | from flydsl.runtime.device import get_rocm_arch as _arch |
| 351 | |
| 352 | a = str(_arch() or "") |
| 353 | _HW_TR_CACHED = a.startswith("gfx950") |
| 354 | except Exception: # noqa: BLE001 |
| 355 | _HW_TR_CACHED = False |
| 356 | if _HW_TR_CACHED: |
| 357 | logger.info_once("FlyDSL TQ: HW V transpose ON (default for gfx950+)") |
| 358 | else: |
| 359 | logger.info_once("FlyDSL TQ: HW V transpose OFF") |
| 360 | return _HW_TR_CACHED |
| 361 | |
| 362 | |
| 363 | _GET_KERNEL_STATS = {"hits": 0, "misses": 0, "build_total_s": 0.0} |
| 364 | |
| 365 | |
| 366 | def _get_kernel( |
| 367 | num_kv_heads: int, |
| 368 | num_partitions: int, |
| 369 | max_blocks_per_seq: int, |
| 370 | scale: float, |
| 371 | query_group_size: int, |
| 372 | kv_block_size: int, |
| 373 | use_hw_v_transpose: bool = False, |
| 374 | num_seqs_hint: int = 1, |
| 375 | tile_groups_per_partition: int = 1, |
| 376 | ): |
| 377 | # ``num_seqs_hint`` (= runtime B) is forwarded to the build for shape |
| 378 | # awareness; it does NOT participate in the cache key because the |
| 379 | # kernel body uses gpu.block_idx.x (= seq index at runtime) and is |
| 380 | # GQA-symmetric across B at build time. |
| 381 | # |
| 382 | # ``tile_groups_per_partition`` IS in the cache key — each |
| 383 | # value compiles a different unrolled K-tile loop body. |
| 384 | key = ( |
| 385 | num_kv_heads, |
| 386 | int(num_partitions), |
| 387 | int(max_blocks_per_seq), |
| 388 | round(float(scale), 8), |
| 389 | int(query_group_size), |
| 390 | int(kv_block_size), |
| 391 | bool(use_hw_v_transpose), |
| 392 | int(tile_groups_per_partition), |
| 393 | ) |
| 394 | cached = _KERN_CACHE.get(key) |
| 395 | if cached is not None: |
| 396 | _GET_KERNEL_STATS["hits"] += 1 |
| 397 | return cached |
| 398 | assert is_flydsl_available() |
| 399 | # Time the FlyDSL build path so we can quantify cudagraph capture cost. |
| 400 | import time as _t |
| 401 | |
| 402 | _build_t0 = _t.perf_counter() |
| 403 | |
| 404 | # Per-GQA dispatch: GQA-6 (MiniMax-M2.5) lives in the sibling module |
| 405 | # tq_decode_gqa6 to keep the Qwen kernel's invariants untouched. |
| 406 | # GQA-{8,16} (Qwen) keep using the canonical tq_decode kernel. |
| 407 | qg = int(query_group_size) |
| 408 | if qg == 6: |
| 409 | if _TQ_MOD_GQA6 is None: |
| 410 | raise RuntimeError( |
| 411 | "FlyDSL TQ GQA-6 sibling kernel " |
| 412 | "(vllm.v1.attention.ops.flydsl_kernels.tq_decode_gqa6) " |
| 413 | "failed to import; check that FlyDSL is installed/importable." |
| 414 | ) |
| 415 | kmod = _TQ_MOD_GQA6 |
| 416 | kfn = kmod.build_tq_decode_gqa6_module( |
| 417 | num_seqs=int(num_seqs_hint), |
| 418 | num_kv_heads=num_kv_heads, |
| 419 | num_partitions=num_partitions, |
| 420 | max_blocks_per_seq=max_blocks_per_seq, |
| 421 | softmax_scale=float(scale), |
| 422 | query_group_size=qg, |
| 423 | kv_block_size=int(kv_block_size), |
| 424 | use_hw_v_transpose=bool(use_hw_v_transpose), |
| 425 | tile_groups_per_partition=int(tile_groups_per_partition), |
| 426 | ) |
| 427 | else: |
| 428 | kmod = _TQ_MOD |
| 429 | kfn = kmod.build_tq_decode_module( |
| 430 | num_seqs=int(num_seqs_hint), |
| 431 | num_kv_heads=num_kv_heads, |
| 432 | num_partitions=num_partitions, |
| 433 | max_blocks_per_seq=max_blocks_per_seq, |
| 434 | softmax_scale=float(scale), |
| 435 | query_group_size=qg, |
| 436 | kv_block_size=int(kv_block_size), |
| 437 | use_hw_v_transpose=bool(use_hw_v_transpose), |
| 438 | tile_groups_per_partition=int(tile_groups_per_partition), |
| 439 | ) |
| 440 | al = kmod.allocator |
| 441 | block_threads = kmod.BLOCK_THREADS |
| 442 | |
| 443 | flyc = _FLYC |
| 444 | fx = _FX |
| 445 | T = _TYPING_T |
| 446 | CompilationContext = _CC |
| 447 | ir_mod = _IR |
| 448 | |
| 449 | @flyc.jit |
| 450 | def _launch( |
| 451 | out, |
| 452 | es, |
| 453 | ml, |
| 454 | q, |
| 455 | kvc, |
| 456 | cents, |
| 457 | bt, |
| 458 | sl, |
| 459 | gx: fx.Int32, # type: ignore[name-defined] |
| 460 | gy: fx.Int32, # type: ignore[name-defined] |
| 461 | gz: fx.Int32, # type: ignore[name-defined] |
| 462 | stream: fx.Stream, # type: ignore[name-defined] |
| 463 | ): |
| 464 | # Re-finalize the LDS allocator for this launch (idempotent per build). |
| 465 | al.finalized = False |
| 466 | ctx = CompilationContext.get_current() |
| 467 | with ir_mod.InsertionPoint(ctx.gpu_module_body): |
| 468 | al.finalize() |
| 469 | from flydsl.expr import arith |
| 470 | |
| 471 | grid_x = arith.index_cast(T.index, gx.ir_value()) |
| 472 | grid_y = arith.index_cast(T.index, gy.ir_value()) |
| 473 | grid_z = arith.index_cast(T.index, gz.ir_value()) |
| 474 | kfn(out, es, ml, q, kvc, cents, bt, sl).launch( |
| 475 | grid=(grid_x, grid_y, grid_z), |
| 476 | block=(block_threads, 1, 1), |
| 477 | stream=stream, |
| 478 | ) |
| 479 | |
| 480 | _KERN_CACHE[key] = _launch |
| 481 | _build_dt = _t.perf_counter() - _build_t0 |
| 482 | _GET_KERNEL_STATS["misses"] += 1 |
| 483 | _GET_KERNEL_STATS["build_total_s"] += _build_dt |
| 484 | logger.info( |
| 485 | "FlyDSL _get_kernel BUILD #%d dt=%.2fs (cumulative=%.1fs) key=%s", |
| 486 | _GET_KERNEL_STATS["misses"], |
| 487 | _build_dt, |
| 488 | _GET_KERNEL_STATS["build_total_s"], |
| 489 | key, |
| 490 | ) |
| 491 | return _launch |
| 492 | |
| 493 | |
| 494 | # -- Partition reducer (Triton, native kernel layout) ------------------------- |
| 495 | # Reduces the FA2 split-KV partials produced by the FlyDSL kernel. |
| 496 | # segm_out [N, Hk, P, QG, D] bf16 |
| 497 | # segm_max [N, Hk, P, QG] fp32 |
| 498 | # segm_sum [N, Hk, P, QG] fp32 |
| 499 | # → output [N, Hq=Hk*QG, D] of any (bf16/fp16/fp32) dtype. |
| 500 | @triton.jit |
| 501 | def _reduce_partitions( |
| 502 | output_ptr, # [N, Hq, D] in OUT_DTYPE |
| 503 | segm_out_ptr, # [N, Hk, P, QG, D] bf16 |
| 504 | segm_max_ptr, # [N, Hk, P, QG] fp32 |
| 505 | segm_sum_ptr, # [N, Hk, P, QG] fp32 |
| 506 | out_stride_n: tl.int64, # stride on N axis (in elems of OUT_DTYPE) |
| 507 | out_stride_h: tl.int64, # stride on Hq axis |
| 508 | NUM_KV_HEADS: tl.constexpr, |
| 509 | QG: tl.constexpr, |
| 510 | NUM_PARTS: tl.constexpr, |
| 511 | HEAD_SIZE: tl.constexpr, |
| 512 | ): |
| 513 | # grid = (N, Hq) |
| 514 | n = tl.program_id(0) |
| 515 | hq = tl.program_id(1) |
| 516 | kv_h = hq // QG |
| 517 | qg = hq % QG |
| 518 | |
| 519 | # Per-partition pointers for this (n, kv_h, qg) row. |
| 520 | # Strides in elems of source dtype. |
| 521 | msum_base = n * (NUM_KV_HEADS * NUM_PARTS * QG) + kv_h * (NUM_PARTS * QG) + qg |
| 522 | so_base = ( |
| 523 | n * (NUM_KV_HEADS * NUM_PARTS * QG * HEAD_SIZE) |
| 524 | + kv_h * (NUM_PARTS * QG * HEAD_SIZE) |
| 525 | + qg * HEAD_SIZE |
| 526 | ) |
| 527 | |
| 528 | p_off = tl.arange(0, NUM_PARTS) |
| 529 | d_off = tl.arange(0, HEAD_SIZE) |
| 530 | |
| 531 | # Load per-partition max/sum. |
| 532 | m_idx = msum_base + p_off * QG |
| 533 | seg_max = tl.load(segm_max_ptr + m_idx) |
| 534 | seg_sum = tl.load(segm_sum_ptr + m_idx) |
| 535 | |
| 536 | # Mask out empty partitions (max == -inf). |
| 537 | valid = seg_max > float("-inf") |
| 538 | overall_max = tl.max(tl.where(valid, seg_max, float("-inf"))) |
| 539 | |
| 540 | # Rescale exp sums. |
| 541 | rescale = tl.where(valid, tl.exp(seg_max - overall_max), 0.0) |
| 542 | seg_sum_rescaled = seg_sum * rescale |
| 543 | overall_sum = tl.sum(seg_sum_rescaled) |
| 544 | |
| 545 | # Load segment outputs and combine. |
| 546 | so_idx = so_base + p_off[:, None] * (QG * HEAD_SIZE) + d_off[None, :] |
| 547 | seg_out = tl.load(segm_out_ptr + so_idx).to(tl.float32) |
| 548 | # acc = sum_p( seg_out_p * (seg_sum_p * rescale_p) ) / overall_sum |
| 549 | weighted = seg_out * seg_sum_rescaled[:, None] |
| 550 | acc_sum = tl.sum(weighted, axis=0) |
| 551 | acc = tl.where(overall_sum > 0.0, acc_sum / overall_sum, 0.0) |
| 552 | |
| 553 | # Store. |
| 554 | out_off = n * out_stride_n + hq * out_stride_h + d_off |
| 555 | tl.store(output_ptr + out_off, acc.to(output_ptr.dtype.element_ty)) |
| 556 | |
| 557 | |
| 558 | # -- Public launcher ---------------------------------------------------------- |
| 559 | def flydsl_turboquant_decode_attention( |
| 560 | query: torch.Tensor, # [B, Hq, D] bf16/fp16 |
| 561 | kv_cache: torch.Tensor, # [num_blocks, BS, Hk, slot_size_aligned] |
| 562 | block_table: torch.Tensor, # [B, max_blocks_per_seq] int32 |
| 563 | seq_lens: torch.Tensor, # [B] int32 |
| 564 | Pi: torch.Tensor, # [D, D] fp32 |
| 565 | centroids: torch.Tensor, # [N_CENTROIDS] fp32 |
| 566 | scale: float, |
| 567 | mse_bits: int, |
| 568 | key_packed_size: int, |
| 569 | value_quant_bits: int, |
| 570 | value_packed_size: int, |
| 571 | key_fp8: bool = False, |
| 572 | norm_correction: bool = False, |
| 573 | PiT: torch.Tensor | None = None, |
| 574 | max_seq_len: int = 0, |
| 575 | mid_o_buf: Any = None, |
| 576 | output_buf: torch.Tensor | None = None, |
| 577 | lse_buf: Any = None, |
| 578 | buf_holder: Any = None, |
| 579 | max_num_kv_splits: int = 32, |
| 580 | sinks: torch.Tensor | None = None, |
| 581 | ) -> torch.Tensor: |
| 582 | """SoA-decode-compatible launcher backed by the FlyDSL decode kernel. |
| 583 | |
| 584 | Constraints: |
| 585 | * key_fp8 == False |
| 586 | * mse_bits == 4 |
| 587 | * value_quant_bits == 4 |
| 588 | * centroids.numel() == 16 |
| 589 | * D == 128 |
| 590 | * block_size in {16, 32} |
| 591 | * Hq // Hk in {6, 8, 16} |
| 592 | - 8/16 → canonical tq_decode kernel (Qwen2.5-72B / Qwen3-32B) |
| 593 | - 6 → tq_decode_gqa6 sibling kernel (MiniMax-M2.5) |
| 594 | |
| 595 | Sinks are NYI and silently ignored if set. norm_correction is honored |
| 596 | implicitly via the pre-folded stored K-norm (see footer comment). |
| 597 | """ |
| 598 | del mid_o_buf, lse_buf, key_packed_size, value_packed_size |
| 599 | if not is_flydsl_available(): |
| 600 | raise RuntimeError( |
| 601 | "FlyDSL decode requested but FlyDSL is not available (needs gfx950 " |
| 602 | "+ importable FlyDSL); the SoA Triton decode is the fallback." |
| 603 | ) |
| 604 | assert not key_fp8, "FlyDSL supports MSE-key path only" |
| 605 | assert mse_bits == 4, f"FlyDSL expects mse_bits=4, got {mse_bits}" |
| 606 | assert value_quant_bits == 4, ( |
| 607 | f"FlyDSL expects value_quant_bits=4, got {value_quant_bits}" |
| 608 | ) |
| 609 | |
| 610 | B, Hq, D = query.shape |
| 611 | Hk = kv_cache.shape[2] |
| 612 | block_size = kv_cache.shape[1] |
| 613 | QG = Hq // Hk |
| 614 | assert D == _TQ_MOD.HEAD_SIZE |
| 615 | assert block_size in (16, 32), ( |
| 616 | f"FlyDSL supports kv_block_size 16 or 32, got {block_size}" |
| 617 | ) |
| 618 | assert QG in (6, 8, 16), f"FlyDSL supports GQA factor 6, 8 or 16, got {QG}" |
| 619 | if QG == 6 and _TQ_MOD_GQA6 is None: |
| 620 | raise RuntimeError( |
| 621 | "FlyDSL launcher: GQA-6 requested (MiniMax-class) but the " |
| 622 | "tq_decode_gqa6 sibling module is not available. Update your " |
| 623 | "FlyDSL checkout (must include kernels/tq_decode_gqa6.py); " |
| 624 | "otherwise the SoA Triton decode is the fallback." |
| 625 | ) |
| 626 | assert centroids.numel() == _TQ_MOD.N_CENTROIDS, ( |
| 627 | f"centroids.numel={centroids.numel()} != {_TQ_MOD.N_CENTROIDS}" |
| 628 | ) |
| 629 | |
| 630 | # ---- Per-layer cache for PiT_f32 + contiguous centroids --- |
| 631 | # PiT and centroids are model constants (set once at layer warmup). |
| 632 | # Avoid the per-decode-step transpose+cast and `.contiguous()` no-op |
| 633 | # check by stashing the pre-cooked tensors on ``buf_holder`` (= layer). |
| 634 | PiT_f32: torch.Tensor |
| 635 | centroids_c: torch.Tensor |
| 636 | if buf_holder is not None: |
| 637 | PiT_f32 = getattr(buf_holder, "_tq_PiT_f32", None) |
| 638 | if PiT_f32 is None: |
| 639 | _PiT_src = PiT if PiT is not None else Pi.T.contiguous() |
| 640 | PiT_f32 = ( |
| 641 | _PiT_src |
| 642 | if _PiT_src.dtype == torch.float32 |
| 643 | else _PiT_src.to(torch.float32) |
| 644 | ) |
| 645 | buf_holder._tq_PiT_f32 = PiT_f32 |
| 646 | centroids_c = getattr(buf_holder, "_tq_centroids_c", None) |
| 647 | if centroids_c is None: |
| 648 | centroids_c = centroids.contiguous() |
| 649 | buf_holder._tq_centroids_c = centroids_c |
| 650 | else: |
| 651 | # Defensive fallback (unit tests may pass ``buf_holder=None``). |
| 652 | _PiT_src = PiT if PiT is not None else Pi.T.contiguous() |
| 653 | PiT_f32 = ( |
| 654 | _PiT_src if _PiT_src.dtype == torch.float32 else _PiT_src.to(torch.float32) |
| 655 | ) |
| 656 | centroids_c = centroids.contiguous() |
| 657 | |
| 658 | # ---- Pooled q_rot (stable data_ptr post-capture) ------ |
| 659 | # On ROCm, fresh tensor allocations made AFTER HIP graph capture can land |
| 660 | # in the graph memory pool and cause GPU memory faults when passed as kernel |
| 661 | # pointer arguments in the eager mixed-batch path. We pool q_rot and the |
| 662 | # fp32 intermediate the same way segm_out/output are pooled: one max-B |
| 663 | # allocation per shape, reused across all B's. |
| 664 | # |
| 665 | # Note: pool_bufs is computed BELOW after num_partitions is determined; |
| 666 | # q_rot needs pool_bufs["q_float"] and pool_bufs["q_rot"]. We compute the |
| 667 | # rotation inline after the pool is fetched (see the pooled-buffers |
| 668 | # section below). |
| 669 | |
| 670 | # ---- Partition count (FA2 split-KV) ---------------------------------- |
| 671 | # |
| 672 | # FlyDSL runs INSIDE FULL cudagraph (TurboQuantMetadataBuilder._cudagraph_ |
| 673 | # support = UNIFORM_BATCH). That means the gridDim baked at capture |
| 674 | # time MUST equal the gridDim at replay — the kernel launch parameters |
| 675 | # are recorded into the captured graph. So `num_partitions` (which |
| 676 | # becomes gridDim.z) MUST be derived from a stable source that produces |
| 677 | # the same value at capture and at runtime. |
| 678 | # |
| 679 | # We use the worst-case context length implied by the block table's |
| 680 | # allocation (block_table.shape[1] * block_size) as that source. It's |
| 681 | # bounded by `max_model_len` (vLLM allocates the block table for the |
| 682 | # configured max), and it's identical at capture and at runtime |
| 683 | # because `block_table.shape[1]` is fixed once the model is loaded. |
| 684 | # Per-tile OOB redirects + the masked FA-2 reducer ensure that for |
| 685 | # short sequences the extra partitions contribute zero to the output. |
| 686 | # |
| 687 | # We always size from the worst case (block_table.shape[1] * block_size), |
| 688 | # never from the per-step actual max_seq_len: a per-step size would make |
| 689 | # the captured gridDim mismatch runtime and GPU-fault under cudagraph. |
| 690 | kv_compute_block = _TQ_MOD.KV_COMPUTE_BLOCK |
| 691 | worst_case_max_seq_len = int(block_table.shape[1]) * int(block_size) |
| 692 | if max_seq_len <= 0: |
| 693 | max_seq_len = worst_case_max_seq_len |
| 694 | sizing_max_seq_len = worst_case_max_seq_len |
| 695 | |
| 696 | # ── Bounded num_partitions + internal tile-group looping ── |
| 697 | # The kernel previously hardcoded one partition = 256 tokens (= |
| 698 | # KV_COMPUTE_BLOCK), forcing num_partitions to scale linearly with |
| 699 | # max_model_len, which made the cudagraph capture cost grow |
| 700 | # prohibitively at long context. The new kernel takes |
| 701 | # ``tile_groups_per_partition`` (TGPP); each CTA processes |
| 702 | # ``TGPP * 16`` K-tiles = ``TGPP * KV_COMPUTE_BLOCK`` tokens with the |
| 703 | # FA-2 online-softmax state accumulating across all of them. This |
| 704 | # mirrors what the SoA Triton decode / HIP SoA-fusion already do. |
| 705 | # |
| 706 | # Strategy: cap num_partitions at ``MAX_PARTITIONS`` (default 32), |
| 707 | # then derive TGPP so that ``num_partitions * TGPP * KV_COMPUTE_BLOCK |
| 708 | # >= sizing_max_seq_len``. |
| 709 | # |
| 710 | # Examples (block_size=32 → worst_case max_bps*32): |
| 711 | # max_model_len=8K: required=32, parts=32, TGPP=1 (no waste) |
| 712 | # max_model_len=16K: required=64, parts=32, TGPP=2 (no waste) |
| 713 | # max_model_len=32K: required=128, parts=32, TGPP=4 (no waste) |
| 714 | # max_model_len=64K: required=256, parts=32, TGPP=8 (no waste) |
| 715 | # max_model_len=128K: required=512, parts=32, TGPP=16 |
| 716 | MAX_PARTITIONS = 32 |
| 717 | required_num_partitions = ( |
| 718 | sizing_max_seq_len + kv_compute_block - 1 |
| 719 | ) // kv_compute_block |
| 720 | # max_num_kv_splits acts as a parallelism floor (ensure at least this |
| 721 | # many CTAs along grid.z), but is itself capped by MAX_PARTITIONS. |
| 722 | parallelism_floor = min(MAX_PARTITIONS, max(1, max_num_kv_splits)) |
| 723 | # Cap num_partitions at MAX_PARTITIONS, but never go below the floor. |
| 724 | num_partitions_actual = max( |
| 725 | parallelism_floor, |
| 726 | min(MAX_PARTITIONS, required_num_partitions), |
| 727 | ) |
| 728 | # Round up to next power of 2 for the Triton reducer's tl.arange(0, N) |
| 729 | # constraint (Triton requires power-of-2 ≥ 2). |
| 730 | num_partitions = max(2, triton.next_power_of_2(num_partitions_actual)) |
| 731 | # Now derive tile_groups_per_partition (TGPP) so total coverage |
| 732 | # ``num_partitions * TGPP * KV_COMPUTE_BLOCK`` is >= sizing_max_seq_len. |
| 733 | # This guarantees no work is dropped at runtime regardless of seq_len. |
| 734 | # Round TGPP up to the next power-of-2 so the JIT cache key is bounded |
| 735 | # to a small set of values (e.g. {1, 2, 4, 8, 16}) — limits the number |
| 736 | # of distinct kernel binaries that need to be compiled across requests. |
| 737 | _tgpp_required = max( |
| 738 | 1, |
| 739 | (required_num_partitions + num_partitions - 1) // num_partitions, |
| 740 | ) |
| 741 | tile_groups_per_partition = int(triton.next_power_of_2(_tgpp_required)) |
| 742 | |
| 743 | # ---- Pooled buffers (no per-call cudaMalloc / memset) ---------- |
| 744 | # The FlyDSL kernel writes the FULL [B, Hk, P, QG, D] segm_out and the |
| 745 | # FULL [B, Hk, P, QG] segm_max / segm_sum (running_max=-inf, sum=0 for |
| 746 | # empty partitions are stored by the kernel itself), so uninitialized |
| 747 | # buffers from the pool are safe to reuse. |
| 748 | device = query.device |
| 749 | pool_bufs = _SEGM_POOL.get( |
| 750 | B, |
| 751 | Hk, |
| 752 | Hq, |
| 753 | num_partitions, |
| 754 | QG, |
| 755 | D, |
| 756 | device, |
| 757 | query.dtype, |
| 758 | ) |
| 759 | segm_out = pool_bufs["segm_out"] |
| 760 | segm_max = pool_bufs["segm_max"] |
| 761 | segm_sum = pool_bufs["segm_sum"] |
| 762 | if output_buf is None: |
| 763 | output = pool_bufs["output"] |
| 764 | else: |
| 765 | output = output_buf[:B] if output_buf.shape[0] != B else output_buf |
| 766 | |
| 767 | # ---- q rotation: GEMM q @ PiT via stable pooled buffers ------------- |
| 768 | _q_float = pool_bufs["q_float"] # [B, Hq, D] fp32, stable |
| 769 | _q_rot_f32 = pool_bufs["q_rot_fp32"] # [B, Hq, D] fp32, stable mm output |
| 770 | _q_rot_out = pool_bufs["q_rot"] # [B, Hq, D] query.dtype, stable |
| 771 | _q_float.copy_(query) # bf16 → fp32 in-place, no alloc |
| 772 | # mm into stable fp32 buffer via out= to avoid fresh allocations. |
| 773 | # PiT_f32 is cached on buf_holder (stable ptr since first layer warmup). |
| 774 | torch.mm( |
| 775 | _q_float.view(B * Hq, D), |
| 776 | PiT_f32, |
| 777 | out=_q_rot_f32.view(B * Hq, D), # in-place into pool buffer |
| 778 | ) |
| 779 | _q_rot_out.copy_(_q_rot_f32) # fp32 → bf16, into stable buf |
| 780 | q_for_kernel = _q_rot_out # stable ptr for kernel launch |
| 781 | |
| 782 | # ---- FlyDSL kernel launch ------------------------------------------- |
| 783 | max_bps = int(block_table.shape[1]) |
| 784 | use_hw_tr = _hw_tr_enabled() |
| 785 | launch = _get_kernel( |
| 786 | Hk, |
| 787 | num_partitions, |
| 788 | max_bps, |
| 789 | scale, |
| 790 | QG, |
| 791 | block_size, |
| 792 | use_hw_v_transpose=use_hw_tr, |
| 793 | num_seqs_hint=int(B), |
| 794 | tile_groups_per_partition=int(tile_groups_per_partition), |
| 795 | ) |
| 796 | # Zero-overhead one-shot info log (replaces logger.info_once which |
| 797 | # hashes its format string on every call to dedup). |
| 798 | global _LOG_INVOKED_ONCE, _LOG_SINKS_WARNED, _LOG_NORM_WARNED |
| 799 | if not _LOG_INVOKED_ONCE: |
| 800 | _LOG_INVOKED_ONCE = True |
| 801 | logger.info( |
| 802 | "FlyDSL launcher invoked (UNIFORM_BATCH cudagraph): " |
| 803 | "B=%d Hk=%d Hq=%d D=%d QG=%d num_partitions=%d (actual=%d, " |
| 804 | "cap=%d) TGPP=%d max_bps=%d block_size=%d max_seq_len=%d " |
| 805 | "hw_v_transpose=%s (coverage=%d tokens, worst_case=%d tokens)", |
| 806 | B, |
| 807 | Hk, |
| 808 | Hq, |
| 809 | D, |
| 810 | QG, |
| 811 | num_partitions, |
| 812 | num_partitions_actual, |
| 813 | MAX_PARTITIONS, |
| 814 | tile_groups_per_partition, |
| 815 | max_bps, |
| 816 | int(block_size), |
| 817 | int(max_seq_len), |
| 818 | use_hw_tr, |
| 819 | num_partitions * tile_groups_per_partition * kv_compute_block, |
| 820 | worst_case_max_seq_len, |
| 821 | ) |
| 822 | launch( |
| 823 | segm_out, |
| 824 | segm_sum, |
| 825 | segm_max, |
| 826 | q_for_kernel, |
| 827 | kv_cache, |
| 828 | centroids_c, |
| 829 | block_table, |
| 830 | seq_lens, |
| 831 | B, |
| 832 | Hk, |
| 833 | num_partitions, |
| 834 | torch.cuda.current_stream(), |
| 835 | ) |
| 836 | |
| 837 | # ---- Reduce partitions -> [B, Hq, D] -------------------------------- |
| 838 | _reduce_partitions[(B, Hq)]( |
| 839 | output_ptr=output, |
| 840 | segm_out_ptr=segm_out, |
| 841 | segm_max_ptr=segm_max, |
| 842 | segm_sum_ptr=segm_sum, |
| 843 | out_stride_n=output.stride(0), |
| 844 | out_stride_h=output.stride(1), |
| 845 | NUM_KV_HEADS=Hk, |
| 846 | QG=QG, |
| 847 | NUM_PARTS=num_partitions, |
| 848 | HEAD_SIZE=D, |
| 849 | ) |
| 850 | if sinks is not None and not _LOG_SINKS_WARNED: |
| 851 | _LOG_SINKS_WARNED = True |
| 852 | logger.warning( |
| 853 | "FlyDSL launcher: sinks ignored (NYI). Disable sinks if they " |
| 854 | "are required for correctness." |
| 855 | ) |
| 856 | # ── norm_correction is honored IMPLICITLY ─────────────────────────── |
| 857 | # When the model was stored with norm_correction=True (the *_nc presets |
| 858 | # turboquant_4bit_nc / k3v4_nc / 3bit_nc), the per-token K-norm scalar |
| 859 | # was pre-folded to ||k_t|| / ||c_t|| at store time by the TurboQuant |
| 860 | # store path. The decode kernel just multiplies `c_vals * stored_knorm`, |
| 861 | # which then equals `(c_vals / ||c_t||) * ||k_t||` — exactly the |
| 862 | # unit-norm-renormalized centroid times the original key norm. The |
| 863 | # Triton decode does the identical multiply, so both decoders honor |
| 864 | # norm_correction equivalently and there is nothing to "do" at decode |
| 865 | # time. The launcher's `norm_correction` arg is only kept for API |
| 866 | # parity; we emit a one-shot INFO log to make the contract explicit. |
| 867 | if norm_correction and not _LOG_NORM_WARNED: |
| 868 | _LOG_NORM_WARNED = True |
| 869 | logger.info( |
| 870 | "FlyDSL launcher: norm_correction honored implicitly via " |
| 871 | "pre-folded stored K-norm (cf. triton_turboquant_store step 3); " |
| 872 | "no decode-time work required, identical to the SoA decode behavior." |
| 873 | ) |
| 874 | return output |