fabolous005
d1ccc5c294
Measured on the live 16k agent setup (A/B with runtime flags, greedy answers vs reference): decode 3.49 -> 5.24 tok/s, first token 4.3 -> 3.3 s, 2.3k-token prompt 27.5 -> 20.6 s, same answers. Dropped after measuring: route stats on the CPU, a Triton router GEMV, pregate prefetch from disk (all within noise or slower). The GEMV keeps one config instead of autotuning: Triton benchmarks with a 256 MiB buffer that torch keeps cached, which left 8 MiB of VRAM and made lazy kernel loads fail (hipErrorNoBinaryForGpu). prain also returns torch cached blocks after materializing and whenever the driver has < 128 MiB free. quality.py score --greedy-only skips prompt logprobs, which OOM the 16k config.
91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
"""Decode GEMMs (at most MAX_ROWS rows) for gfx10 GPUs as a Triton GEMV.
|
|
|
|
vLLM sends unquantized linear layers on ROCm through rocBLAS, except on gfx9/gfx11/gfx12 where it
|
|
has skinny-GEMM kernels. On an RX 6600 XT (gfx1032) rocBLAS runs a bf16 GEMV at 13-16 GB/s
|
|
(RDNA2 has no bf16 dot instructions), so the dense part of a decode step reads its weights at
|
|
6% of the card's bandwidth. This kernel reads them at ~210 GB/s (measured, lm_head 262144 x 2816:
|
|
93 ms -> 6.7 ms). install() wraps vLLM's dispatch_unquantized_gemm; PRAIN_GEMV=0 turns it off.
|
|
"""
|
|
|
|
import functools
|
|
import os
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
from prain.kernels import gfx10
|
|
|
|
MAX_ROWS = 4
|
|
# ponytail: one config, the one autotuning picked for every Gemma 4 dense shape on gfx1032.
|
|
# Triton's autotuner benchmarks with a 256 MiB buffer that torch keeps cached, and an 8 GB card
|
|
# has no 256 MiB to spare (lazy kernel loads then fail). Re-measure if the GPU or model changes.
|
|
BLOCK_N, BLOCK_K, NUM_WARPS = 4, 512, 8
|
|
|
|
|
|
@triton.jit
|
|
def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
|
|
BLOCK_K: tl.constexpr):
|
|
rows = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
|
|
ms = tl.arange(0, BLOCK_M)
|
|
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
|
|
for k0 in range(0, K, BLOCK_K):
|
|
cols = k0 + tl.arange(0, BLOCK_K)
|
|
w = tl.load(w_ptr + rows[:, None] * K + cols[None, :],
|
|
mask=(rows[:, None] < N) & (cols[None, :] < K), other=0.0)
|
|
x = tl.load(x_ptr + ms[:, None] * K + cols[None, :],
|
|
mask=(ms[:, None] < M) & (cols[None, :] < K), other=0.0)
|
|
acc += tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], axis=2)
|
|
tl.store(y_ptr + ms[:, None] * N + rows[None, :], acc.to(y_ptr.dtype.element_ty),
|
|
mask=(ms[:, None] < M) & (rows[None, :] < N))
|
|
|
|
|
|
@functools.cache
|
|
def fast_device(index):
|
|
"""gfx10 (RDNA2) GPUs, where vLLM has no skinny-GEMM kernels."""
|
|
return gfx10(index) and os.environ.get("PRAIN_GEMV", "1") != "0"
|
|
|
|
|
|
def _small(x, weight, bias):
|
|
return (bias is None and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)
|
|
and weight.dtype == x.dtype and weight.is_contiguous()
|
|
and x.numel() <= MAX_ROWS * x.shape[-1] and fast_device(x.device.index or 0))
|
|
|
|
|
|
def linear(x, weight, bias=None):
|
|
"""F.linear(x, weight, bias); the Triton GEMV for up to MAX_ROWS rows on gfx10."""
|
|
if not _small(x, weight, bias):
|
|
return F.linear(x, weight, bias)
|
|
n, k = weight.shape
|
|
x2 = x.reshape(-1, k).contiguous()
|
|
m = x2.shape[0]
|
|
y = torch.empty(m, n, dtype=x.dtype, device=x.device)
|
|
_gemv[(triton.cdiv(n, BLOCK_N),)](weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m),
|
|
BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=NUM_WARPS)
|
|
return y.reshape(*x.shape[:-1], n)
|
|
|
|
|
|
def wrap(dispatch):
|
|
"""vLLM's dispatch_unquantized_gemm(), with decode shapes sent to linear()."""
|
|
def wrapped(*args, **kwargs):
|
|
fallback = dispatch(*args, **kwargs)
|
|
|
|
def gemm(layer, x, weight, bias=None):
|
|
if _small(x, weight, bias):
|
|
return linear(x, weight)
|
|
return fallback(layer, x, weight, bias)
|
|
return gemm
|
|
return wrapped
|
|
|
|
|
|
def install():
|
|
"""Before the model is built: every module that looks up dispatch_unquantized_gemm."""
|
|
from vllm.model_executor.layers import linear as linear_layers, utils
|
|
from vllm.model_executor.layers import vocab_parallel_embedding
|
|
from vllm.model_executor.layers.fused_moe.runner import moe_runner
|
|
|
|
wrapped = wrap(utils.dispatch_unquantized_gemm)
|
|
for module in (linear_layers, vocab_parallel_embedding, moe_runner):
|
|
module.dispatch_unquantized_gemm = wrapped
|