From c84789c40b506a40d6a1ec15a704d53397c564a6 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:20:19 -0400 Subject: [PATCH] [Refactor] Remove kernel dead code (#51051) Signed-off-by: yewentao256 --- csrc/cpu/cpu_attn_fp8.hpp | 27 ----- csrc/libtorch_stable/cache_kernels.cu | 51 -------- .../moe/grouped_topk_kernels.cu | 44 ------- .../fused_kernels/quant_conversions.cuh | 11 -- .../kernels/linear/cute_dsl/skinny_gemm.py | 11 -- .../amd/ops/gluon/rel_mha_decode_gfx950.py | 113 ------------------ vllm/models/inkling/amd/ops/gluon/utils.py | 69 ----------- vllm/v1/attention/backends/fa_utils.py | 70 ----------- 8 files changed, 396 deletions(-) diff --git a/csrc/cpu/cpu_attn_fp8.hpp b/csrc/cpu/cpu_attn_fp8.hpp index 764b6ed7f8..153090966d 100644 --- a/csrc/cpu/cpu_attn_fp8.hpp +++ b/csrc/cpu/cpu_attn_fp8.hpp @@ -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::quiet_NaN(); - const uint32_t sign = static_cast(b & 0x80) << 24; - if (exp_bits == 0x1F) - return sign ? -std::numeric_limits::infinity() - : std::numeric_limits::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(exp_bits) - 15 + 127) << 23) | - (static_cast(mant_bits) << 21); - float val = *reinterpret_cast(&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; diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index c35a516b5c..83716d6fc7 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -186,57 +186,6 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, namespace vllm { -// Grid: (num_layers, num_pairs) -template -__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(key_cache_ptrs[layer_idx]); - scalar_t* value_cache = - reinterpret_cast(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 -__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(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 struct CopyWithScaleOp { diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu index e329e067b2..852b2b468a 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu @@ -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(t, val_arr_[i], in_idx[idx], idx_arr_[i]); - } else { - is_better = is_better_than(t, val_arr_[i]); - } - if (is_better) { - val_arr_[i] = t; - idx_arr_[i] = in_idx[idx]; - } - } - } - - BitonicMerge::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`. diff --git a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh b/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh index dbe38092f9..a81a0b9feb 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh +++ b/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh @@ -76,15 +76,4 @@ struct ScaledQuant -__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(input[i], scale); - } -} - } // namespace vllm diff --git a/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py index 73d1b2b067..ec1e607caf 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py @@ -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, diff --git a/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py b/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py index 987508f91f..a9d80b7fbf 100644 --- a/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py +++ b/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py @@ -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, diff --git a/vllm/models/inkling/amd/ops/gluon/utils.py b/vllm/models/inkling/amd/ops/gluon/utils.py index 90f2e586bc..0aa26d75e6 100644 --- a/vllm/models/inkling/amd/ops/gluon/utils.py +++ b/vllm/models/inkling/amd/ops/gluon/utils.py @@ -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 diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index c01aa0181a..22f9e6a14d 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -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,