[Refactor] Remove kernel dead code (#51051)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-08-07 01:20:19 -07:00
committed by GitHub
parent da788334bc
commit c84789c40b
8 changed files with 0 additions and 396 deletions
-27
View File
@@ -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;
-51
View File
@@ -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
-70
View File
@@ -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,