| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared utilities for flash-maxsim kernels: padding, autotune configs, |
| 4 | and pruning.""" |
| 5 | |
| 6 | from vllm.triton_utils import triton |
| 7 | |
| 8 | |
| 9 | def _next_pow2(x: int) -> int: |
| 10 | return 1 << (x - 1).bit_length() |
| 11 | |
| 12 | |
| 13 | def _get_configs() -> list: |
| 14 | """Autotune grid over (BLOCK_Q, BLOCK_D); `_prune_configs` filters by |
| 15 | SMEM budget and Lq at runtime, so one list is safe across GPUs.""" |
| 16 | return [ |
| 17 | triton.Config({"BLOCK_Q": 16, "BLOCK_D": 16}, num_warps=2, num_stages=2), |
| 18 | triton.Config({"BLOCK_Q": 16, "BLOCK_D": 32}, num_warps=2, num_stages=2), |
| 19 | triton.Config({"BLOCK_Q": 32, "BLOCK_D": 16}, num_warps=2, num_stages=2), |
| 20 | triton.Config({"BLOCK_Q": 32, "BLOCK_D": 32}, num_warps=4, num_stages=2), |
| 21 | triton.Config({"BLOCK_Q": 32, "BLOCK_D": 64}, num_warps=4, num_stages=2), |
| 22 | triton.Config({"BLOCK_Q": 32, "BLOCK_D": 128}, num_warps=8, num_stages=2), |
| 23 | triton.Config({"BLOCK_Q": 64, "BLOCK_D": 64}, num_warps=4, num_stages=2), |
| 24 | triton.Config({"BLOCK_Q": 64, "BLOCK_D": 128}, num_warps=8, num_stages=1), |
| 25 | triton.Config({"BLOCK_Q": 128, "BLOCK_D": 64}, num_warps=8, num_stages=1), |
| 26 | triton.Config({"BLOCK_Q": 128, "BLOCK_D": 128}, num_warps=8, num_stages=1), |
| 27 | ] |
| 28 | |
| 29 | |
| 30 | def _smem_budget() -> int: |
| 31 | """Per-block opt-in shared-memory budget of the active device (bytes): |
| 32 | Triton driver query, else a conservative capability table.""" |
| 33 | try: |
| 34 | driver = triton.runtime.driver.active |
| 35 | props = driver.utils.get_device_properties(driver.get_current_device()) |
| 36 | return int(props["max_shared_mem"]) |
| 37 | except Exception: |
| 38 | try: |
| 39 | from vllm.platforms import current_platform |
| 40 | |
| 41 | cap = current_platform.get_device_capability() |
| 42 | if cap is not None: |
| 43 | if cap.major >= 9: |
| 44 | return 232_448 # Hopper+: 227 KB opt-in |
| 45 | if cap.major == 8: |
| 46 | return 101_376 # worst sm_8x (86/89); A100 has more |
| 47 | except Exception: |
| 48 | pass |
| 49 | return 98_304 # conservative pre-Ampere default |
| 50 | |
| 51 | |
| 52 | def _prune_configs(configs, named_args, **kwargs): |
| 53 | """Reject configs that exceed the GPU's shared-memory budget. |
| 54 | |
| 55 | BLOCK_Q and BLOCK_D load fp16 tiles of size [BLOCK_*, d_pad] — the |
| 56 | kernels index with ``tl.arange(0, d_pad)``, so the estimate must use |
| 57 | the padded dimension, not the real one (d=513 pads to 1024 and |
| 58 | nearly doubles the tile) — and the [BLOCK_Q, BLOCK_D] score tile is |
| 59 | fp32. Total: |
| 60 | (BLOCK_Q*d_pad + BLOCK_D*d_pad)*2 + BLOCK_Q*BLOCK_D*4 bytes. |
| 61 | """ |
| 62 | Lq = named_args.get("Lq", 32) |
| 63 | d_pad = named_args.get("d_pad") |
| 64 | if d_pad is None: |
| 65 | d_pad = _next_pow2(max(named_args.get("d", 128), 16)) |
| 66 | smem_limit = _smem_budget() |
| 67 | |
| 68 | def _est(cfg): |
| 69 | bq = cfg.kwargs["BLOCK_Q"] |
| 70 | bd = cfg.kwargs["BLOCK_D"] |
| 71 | return (bq * d_pad + bd * d_pad) * 2 + bq * bd * 4 |
| 72 | |
| 73 | pruned = [] |
| 74 | for cfg in configs: |
| 75 | if cfg.kwargs["BLOCK_Q"] > Lq * 2: |
| 76 | continue |
| 77 | if _est(cfg) > smem_limit: |
| 78 | continue |
| 79 | pruned.append(cfg) |
| 80 | if pruned: |
| 81 | return pruned |
| 82 | # Nothing fits the estimate (large d_pad on a small-SMEM device). |
| 83 | # Return the single smallest-footprint config — a deterministic |
| 84 | # least-bad launch attempt — never arbitrary rejected configs. |
| 85 | return [min(configs, key=_est)] |