[Refactor] Remove kernel dead code (#51051)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
@@ -116,33 +116,6 @@ inline void reshape_and_cache_fp8_amx_impl(
|
||||
// ---------------------------------------------------------------------------
|
||||
// FP8 E5M2 scalar helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Reference scalar dequant — used to verify vectorized AMX dequant.
|
||||
// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
|
||||
// Byte b → FP16 bits = b << 8 (no bias correction needed).
|
||||
inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
|
||||
const uint8_t exp_bits = (b >> 2) & 0x1F;
|
||||
const uint8_t mant_bits = b & 0x03;
|
||||
// NaN: exp=11111, mant!=00
|
||||
if (exp_bits == 0x1F && mant_bits != 0)
|
||||
return std::numeric_limits<float>::quiet_NaN();
|
||||
const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
|
||||
if (exp_bits == 0x1F)
|
||||
return sign ? -std::numeric_limits<float>::infinity()
|
||||
: std::numeric_limits<float>::infinity();
|
||||
if (exp_bits == 0) { // subnormal: (-1)^s * 2^-14 * mant/4
|
||||
if (mant_bits == 0) return 0.0f;
|
||||
float v = mant_bits * 0x1p-16f;
|
||||
return (sign ? -v : v) * scale;
|
||||
}
|
||||
// Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
|
||||
uint32_t fp32_bits = sign |
|
||||
((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
|
||||
(static_cast<uint32_t>(mant_bits) << 21);
|
||||
float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
|
||||
return val * scale;
|
||||
}
|
||||
|
||||
inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
|
||||
v *= inv_scale;
|
||||
constexpr float fp8_e5m2_max = 57344.0f;
|
||||
|
||||
@@ -186,57 +186,6 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// Grid: (num_layers, num_pairs)
|
||||
template <typename scalar_t>
|
||||
__global__ void copy_blocks_kernel(int64_t* key_cache_ptrs,
|
||||
int64_t* value_cache_ptrs,
|
||||
const int64_t* __restrict__ block_mapping,
|
||||
const int numel_per_block) {
|
||||
const int layer_idx = blockIdx.x;
|
||||
const int pair_idx = blockIdx.y;
|
||||
|
||||
scalar_t* key_cache = reinterpret_cast<scalar_t*>(key_cache_ptrs[layer_idx]);
|
||||
scalar_t* value_cache =
|
||||
reinterpret_cast<scalar_t*>(value_cache_ptrs[layer_idx]);
|
||||
int64_t src_block_number = block_mapping[2 * pair_idx];
|
||||
int64_t dst_block_number = block_mapping[2 * pair_idx + 1];
|
||||
|
||||
const int64_t src_block_offset = src_block_number * numel_per_block;
|
||||
const int64_t dst_block_offset = dst_block_number * numel_per_block;
|
||||
for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) {
|
||||
int64_t src_offset = src_block_offset + i;
|
||||
int64_t dst_offset = dst_block_offset + i;
|
||||
key_cache[dst_offset] = key_cache[src_offset];
|
||||
}
|
||||
for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) {
|
||||
int64_t src_offset = src_block_offset + i;
|
||||
int64_t dst_offset = dst_block_offset + i;
|
||||
value_cache[dst_offset] = value_cache[src_offset];
|
||||
}
|
||||
}
|
||||
|
||||
// Kernel for MLA, which works on a single joint kv_cache
|
||||
// Grid: (num_layers, num_pairs)
|
||||
template <typename scalar_t>
|
||||
__global__ void copy_blocks_mla_kernel(
|
||||
int64_t* cache_ptrs, const int64_t* __restrict__ block_mapping,
|
||||
const int mem_footprint_per_block) {
|
||||
const int layer_idx = blockIdx.x;
|
||||
const int pair_idx = blockIdx.y;
|
||||
scalar_t* cache = reinterpret_cast<scalar_t*>(cache_ptrs[layer_idx]);
|
||||
int64_t src_block = block_mapping[2 * pair_idx];
|
||||
int64_t dst_block = block_mapping[2 * pair_idx + 1];
|
||||
int64_t src_offset = src_block * mem_footprint_per_block;
|
||||
int64_t dst_offset = dst_block * mem_footprint_per_block;
|
||||
for (int i = threadIdx.x; i < mem_footprint_per_block; i += blockDim.x) {
|
||||
cache[dst_offset + i] = cache[src_offset + i];
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// Used to copy/convert one element
|
||||
template <typename OutT, typename InT, Fp8KVCacheDataType kv_dt>
|
||||
struct CopyWithScaleOp {
|
||||
|
||||
@@ -232,50 +232,6 @@ class WarpSort {
|
||||
}
|
||||
}
|
||||
|
||||
// load and merge k sorted values
|
||||
__device__ void load_sorted(T const* __restrict__ in,
|
||||
idxT const* __restrict__ in_idx, idxT start) {
|
||||
idxT idx = start + WARP_SIZE - 1 - lane_;
|
||||
for (int i = max_arr_len_ - 1; i >= 0; --i, idx += WARP_SIZE) {
|
||||
if (idx < start + k_) {
|
||||
T t = in[idx];
|
||||
bool is_better;
|
||||
if constexpr (is_stable) {
|
||||
is_better =
|
||||
is_better_than<greater>(t, val_arr_[i], in_idx[idx], idx_arr_[i]);
|
||||
} else {
|
||||
is_better = is_better_than<greater>(t, val_arr_[i]);
|
||||
}
|
||||
if (is_better) {
|
||||
val_arr_[i] = t;
|
||||
idx_arr_[i] = in_idx[idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
BitonicMerge<capacity, greater, !greater, T, idxT, is_stable>::merge(
|
||||
val_arr_, idx_arr_);
|
||||
}
|
||||
|
||||
__device__ void dump(T* __restrict__ out, idxT* __restrict__ out_idx) const {
|
||||
for (int i = 0; i < max_arr_len_; ++i) {
|
||||
idxT out_i = i * WARP_SIZE + lane_;
|
||||
if (out_i < k_) {
|
||||
out[out_i] = val_arr_[i];
|
||||
out_idx[out_i] = idx_arr_[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__device__ void dumpIdx(idxT* __restrict__ out_idx) const {
|
||||
for (int i = 0; i < max_arr_len_; ++i) {
|
||||
idxT out_i = i * WARP_SIZE + lane_;
|
||||
if (out_i < k_) {
|
||||
out_idx[out_i] = idx_arr_[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Accessors for per-lane selected value/index.
|
||||
// NOTE: For the common case `capacity == WARP_SIZE`, `max_arr_len_ == 1`
|
||||
// and callers should use `i == 0`.
|
||||
|
||||
@@ -76,15 +76,4 @@ struct ScaledQuant<quant_type_t, is_scale_inverted,
|
||||
}
|
||||
};
|
||||
|
||||
template <typename scalar_t, typename quant_type_t, bool is_scale_inverted>
|
||||
__device__ void scaled_quant_conversion(quant_type_t* __restrict__ output,
|
||||
scalar_t const* __restrict__ input,
|
||||
float const scale, int const tid,
|
||||
int const num_elements,
|
||||
int const step) {
|
||||
for (int i = tid; i < num_elements; i += step) {
|
||||
output[i] = ScaledQuant<quant_type_t, is_scale_inverted>(input[i], scale);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
@@ -151,17 +151,6 @@ class ShapeDynamicSkinnyGemm:
|
||||
options="--enable-tvm-ffi --ptxas-options -maxrregcount=64",
|
||||
)
|
||||
|
||||
def request_warmup(
|
||||
self,
|
||||
dtype: torch.dtype,
|
||||
shapes: Iterable[tuple[int, int, int]],
|
||||
) -> None:
|
||||
"""Request compilation of ``(M, N, K)`` shapes before graph capture."""
|
||||
self.request_warmup_configs(
|
||||
dtype,
|
||||
(self._config(m, n, k) for m, n, k in shapes),
|
||||
)
|
||||
|
||||
def request_warmup_configs(
|
||||
self,
|
||||
dtype: torch.dtype,
|
||||
|
||||
@@ -501,21 +501,6 @@ class AttentionProgram:
|
||||
cdna4.buffer_store(part_o, self.mid_o_ptr, mid_o_offsets, mask=valid[:, None])
|
||||
cdna4.buffer_store(part_lse, self.mid_lse_ptr, mid_lse_offsets, mask=valid)
|
||||
|
||||
@gluon.jit
|
||||
def store_output(self, acc, l_i):
|
||||
cfg = self.cfg
|
||||
offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout))
|
||||
offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout))
|
||||
q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m
|
||||
valid = (self.group_start + offs_m) < cfg.GROUP_SIZE
|
||||
acc = gl.convert_layout(acc, cfg.store_layout)
|
||||
l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout))
|
||||
output = acc * (1.0 / l_i)[:, None]
|
||||
output = output.to(self.mid_o_ptr.dtype.element_ty)
|
||||
offsets = (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.HEAD_DIM
|
||||
offsets += offs_d[None, :]
|
||||
cdna4.buffer_store(output, self.mid_o_ptr, offsets, mask=valid[:, None])
|
||||
|
||||
|
||||
# ===-----------------------------------------------------------------------===#
|
||||
# Entry Point
|
||||
@@ -625,104 +610,6 @@ def _rel_mha_decode_fp16(
|
||||
program.store_split(acc, l_i, m_i)
|
||||
|
||||
|
||||
@gluon.jit
|
||||
def _rel_mha_decode_sliding_fp16(
|
||||
q_ptr,
|
||||
rel_logits_ptr,
|
||||
k_cache_ptr,
|
||||
v_cache_ptr,
|
||||
page_table_ptr,
|
||||
cache_seqlens_ptr,
|
||||
out_ptr,
|
||||
Q_STRIDE_B: gl.constexpr,
|
||||
Q_STRIDE_H: gl.constexpr,
|
||||
Q_STRIDE_D: gl.constexpr,
|
||||
K_STRIDE_B: gl.constexpr,
|
||||
K_STRIDE_P: gl.constexpr,
|
||||
K_STRIDE_H: gl.constexpr,
|
||||
K_STRIDE_D: gl.constexpr,
|
||||
V_STRIDE_B: gl.constexpr,
|
||||
V_STRIDE_P: gl.constexpr,
|
||||
V_STRIDE_H: gl.constexpr,
|
||||
V_STRIDE_D: gl.constexpr,
|
||||
SM_SCALE: gl.constexpr,
|
||||
PAGE_TABLE_STRIDE: gl.constexpr,
|
||||
PAGE_SIZE: gl.constexpr,
|
||||
MAX_SEQLEN_Q: gl.constexpr,
|
||||
NUM_Q_HEADS: gl.constexpr,
|
||||
NUM_KV_HEADS: gl.constexpr,
|
||||
HEAD_DIM: gl.constexpr,
|
||||
BLOCK_M: gl.constexpr,
|
||||
BLOCK_N: gl.constexpr,
|
||||
IS_SLIDING: gl.constexpr,
|
||||
WINDOW_LEFT: gl.constexpr,
|
||||
REL_STRIDE_T: gl.constexpr,
|
||||
REL_STRIDE_H: gl.constexpr,
|
||||
REL_STRIDE_E: gl.constexpr,
|
||||
REL_EXTENT: gl.constexpr,
|
||||
REL_BIAS_QK_SCALE: gl.constexpr,
|
||||
IS_FP8: gl.constexpr,
|
||||
):
|
||||
cfg = AttentionConfig(
|
||||
SM_SCALE,
|
||||
PAGE_TABLE_STRIDE,
|
||||
PAGE_SIZE,
|
||||
1,
|
||||
MAX_SEQLEN_Q,
|
||||
NUM_Q_HEADS,
|
||||
NUM_KV_HEADS,
|
||||
HEAD_DIM,
|
||||
BLOCK_M,
|
||||
BLOCK_N,
|
||||
IS_SLIDING,
|
||||
WINDOW_LEFT,
|
||||
REL_EXTENT,
|
||||
REL_BIAS_QK_SCALE,
|
||||
IS_FP8,
|
||||
InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D),
|
||||
InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E),
|
||||
PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D),
|
||||
PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D),
|
||||
)
|
||||
program = AttentionProgram.create(
|
||||
cfg,
|
||||
q_ptr,
|
||||
rel_logits_ptr,
|
||||
k_cache_ptr,
|
||||
v_cache_ptr,
|
||||
page_table_ptr,
|
||||
cache_seqlens_ptr,
|
||||
out_ptr,
|
||||
out_ptr,
|
||||
)
|
||||
k_smem = gl.allocate_shared_memory(
|
||||
k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout
|
||||
)
|
||||
v_smem = gl.allocate_shared_memory(
|
||||
v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout
|
||||
)
|
||||
|
||||
q = program.load_q()
|
||||
m_i, l_i, acc = program.init_state()
|
||||
|
||||
for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N):
|
||||
physical_page = program.load_page(start_n)
|
||||
program.issue_load_k(physical_page, k_smem)
|
||||
program.issue_load_v(physical_page, v_smem)
|
||||
async_copy.wait_group(1)
|
||||
k = program.shared_load_k(k_smem)
|
||||
qk = program.compute_qk(q, k)
|
||||
qk = program.apply_rel_bias(qk, start_n)
|
||||
qk = program.apply_kv_mask(qk, start_n)
|
||||
p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc)
|
||||
|
||||
async_copy.wait_group(0)
|
||||
v = program.shared_load_v(v_smem)
|
||||
acc = program.compute_pv(p, v, acc)
|
||||
|
||||
program.store_output(acc, l_i)
|
||||
|
||||
|
||||
@gluon.jit
|
||||
def _rel_mha_decode_reduce_fp16(
|
||||
mid_o_ptr,
|
||||
|
||||
@@ -40,75 +40,6 @@ def max(input, axis=None, keep_dims=False):
|
||||
return gl.reduce(input, axis, maximum, keep_dims=keep_dims)
|
||||
|
||||
|
||||
@gluon.constexpr_function
|
||||
def attention_layouts(head_dim, block_n, is_fp8, dtype, num_warps, instr_shape):
|
||||
mfma = gl.amd.AMDMFMALayout(
|
||||
version=4,
|
||||
instr_shape=instr_shape,
|
||||
transposed=True,
|
||||
warps_per_cta=[num_warps, 1],
|
||||
)
|
||||
qk_layout = mfma
|
||||
pv_layout = mfma
|
||||
# qk_kw is derived from a 128-bit load / dtype bitwidth; pv_kw is tuned.
|
||||
qk_kw = 16 if is_fp8 else 8
|
||||
pv_kw = 8 if is_fp8 else 4
|
||||
q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw)
|
||||
k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw)
|
||||
p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw)
|
||||
v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw)
|
||||
# load_vec = elems/lane (dtype-dependent, == qk_kw); load_threads span HEAD_DIM.
|
||||
load_vec = 16 if is_fp8 else 8
|
||||
load_threads = head_dim // load_vec
|
||||
load_layout = gl.BlockedLayout(
|
||||
[1, load_vec], [64 // load_threads, load_threads], [num_warps, 1], [1, 0]
|
||||
)
|
||||
# store_vec is always 16-bit (128 / 16 = 8) regardless of input dtype.
|
||||
store_vec = 8
|
||||
store_threads = head_dim // store_vec
|
||||
store_layout = gl.BlockedLayout(
|
||||
[1, store_vec], [64 // store_threads, store_threads], [num_warps, 1], [1, 0]
|
||||
)
|
||||
# Take only the built-in's padding, not its swizzle: the swizzle scatters
|
||||
# head_dim across banks, making the LDS write stride non-constant, so it can't
|
||||
# lower through async_copy's affine [1, 0] load. Padding-only is affine and
|
||||
# DMA-legal.
|
||||
# TODO(perf): to also use the swizzle, co-design a matched load layout so the
|
||||
# DMA stays legal.
|
||||
k_api = gl.amd.cdna4.compute_efficient_padded_shared_layout(
|
||||
k_layout, [block_n, head_dim], dtype, is_k_contig=True
|
||||
)
|
||||
v_api = gl.amd.cdna4.compute_efficient_padded_shared_layout(
|
||||
v_layout, [block_n, head_dim], dtype, is_k_contig=False
|
||||
)
|
||||
assert k_api is not None and v_api is not None, (
|
||||
"no CDNA4 padded shared layout for this operand/dtype"
|
||||
)
|
||||
k_pairs = list(k_api.interval_padding_pairs)
|
||||
v_pairs = list(v_api.interval_padding_pairs)
|
||||
assert len(k_pairs) == 1 and len(v_pairs) == 1, (
|
||||
"expected a single interval padding pair from the built-in"
|
||||
)
|
||||
k_smem_layout = gl.PaddedSharedLayout.with_identity_for(
|
||||
[[int(k_pairs[0][0]), int(k_pairs[0][1])]], [block_n, head_dim], [1, 0]
|
||||
)
|
||||
v_smem_layout = gl.PaddedSharedLayout.with_identity_for(
|
||||
[[int(v_pairs[0][0]), int(v_pairs[0][1])]], [block_n, head_dim], [1, 0]
|
||||
)
|
||||
return (
|
||||
qk_layout,
|
||||
pv_layout,
|
||||
q_layout,
|
||||
k_layout,
|
||||
p_layout,
|
||||
v_layout,
|
||||
load_layout,
|
||||
store_layout,
|
||||
k_smem_layout,
|
||||
v_smem_layout,
|
||||
)
|
||||
|
||||
|
||||
@aggregate
|
||||
class InputStrides:
|
||||
stride_t: gl.constexpr
|
||||
|
||||
@@ -1,11 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
@@ -70,73 +67,6 @@ elif current_platform.is_rocm():
|
||||
reshape_and_cache_flash = ops.reshape_and_cache_flash
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FlashAttentionCuTeDSLCompileSpec:
|
||||
"""High-level FA4 compile-only request used by vLLM warmup.
|
||||
|
||||
This is not the CuTeDSL cache key. FA4 owns the selector that maps these
|
||||
serving inputs to the actual compile-static fields: tile sizes, q_stage,
|
||||
Split-KV, scheduler choice, layout-presence booleans, dtype/head dims,
|
||||
arch, and related fields.
|
||||
"""
|
||||
|
||||
q_shape: tuple[int, ...]
|
||||
k_shape: tuple[int, ...]
|
||||
v_shape: tuple[int, ...]
|
||||
q_dtype: torch.dtype
|
||||
max_seqlen_q: int
|
||||
max_seqlen_k: int
|
||||
softmax_scale: float
|
||||
causal: bool
|
||||
fa_version: int
|
||||
v_stride: tuple[int, ...] | None = None
|
||||
cu_seqlens_q_shape: tuple[int, ...] | None = None
|
||||
cu_seqlens_k_shape: tuple[int, ...] | None = None
|
||||
window_size: tuple[int, int] | None = None
|
||||
return_softmax_lse: bool = False
|
||||
num_splits: int = 0
|
||||
|
||||
def compile(self) -> None:
|
||||
assert compile_flash_attn_varlen_func_from_specs is not None
|
||||
window_size = list(self.window_size) if self.window_size is not None else None
|
||||
compile_flash_attn_varlen_func_from_specs(
|
||||
q_shape=self.q_shape,
|
||||
k_shape=self.k_shape,
|
||||
v_shape=self.v_shape,
|
||||
q_dtype=self.q_dtype,
|
||||
v_stride=self.v_stride,
|
||||
cu_seqlens_q_shape=self.cu_seqlens_q_shape,
|
||||
cu_seqlens_k_shape=self.cu_seqlens_k_shape,
|
||||
max_seqlen_q=self.max_seqlen_q,
|
||||
max_seqlen_k=self.max_seqlen_k,
|
||||
softmax_scale=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
window_size=window_size,
|
||||
return_softmax_lse=self.return_softmax_lse,
|
||||
fa_version=self.fa_version,
|
||||
num_splits=self.num_splits,
|
||||
)
|
||||
|
||||
def request_key(self) -> tuple[object, ...]:
|
||||
return (
|
||||
self.q_shape,
|
||||
self.k_shape,
|
||||
self.v_shape,
|
||||
self.q_dtype,
|
||||
self.max_seqlen_q,
|
||||
self.max_seqlen_k,
|
||||
self.softmax_scale,
|
||||
self.causal,
|
||||
self.fa_version,
|
||||
self.v_stride,
|
||||
self.cu_seqlens_q_shape,
|
||||
self.cu_seqlens_k_shape,
|
||||
self.window_size,
|
||||
self.return_softmax_lse,
|
||||
self.num_splits,
|
||||
)
|
||||
|
||||
|
||||
def get_flash_attn_version(
|
||||
requires_alibi: bool = False,
|
||||
head_size: int | None = None,
|
||||
|
||||
Reference in New Issue
Block a user