diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 0229805baf3a..a3ccfd64fba7 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -737,6 +737,8 @@ steps: - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py - tests/kernels/test_top_k_per_row.py + - csrc/fused_mqa_topk/ + - tests/kernels/attention/test_fused_mqa_topk.py - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py - tests/kernels/quantization/test_block_fp8.py - tests/models/quantization/test_nvfp4.py @@ -748,6 +750,7 @@ steps: - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py - pytest -v -s tests/kernels/test_top_k_per_row.py + - pytest -v -s tests/kernels/attention/test_fused_mqa_topk.py - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py diff --git a/CMakeLists.txt b/CMakeLists.txt index e9bed271b545..93d0818b6cc8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1232,6 +1232,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") message(STATUS "Building hadacore") endif() + # Fused MQA logits + top-k prefill kernel (SM100 family). It compiles against + # DeepGEMM's headers, so DeepGEMM is fetched here. + include(cmake/external_projects/deepgemm.cmake) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + cuda_archs_loose_intersection(FUSED_MQA_TOPK_ARCHS "10.0f" "${CUDA_ARCHS}") + endif() + if(FUSED_MQA_TOPK_ARCHS) + set(FUSED_MQA_TOPK_SRC "csrc/fused_mqa_topk/fused_mqa_topk.cu") + set_gencode_flags_for_srcs( + SRCS "${FUSED_MQA_TOPK_SRC}" + CUDA_ARCHS "${FUSED_MQA_TOPK_ARCHS}") + set_property(SOURCE ${FUSED_MQA_TOPK_SRC} APPEND PROPERTY INCLUDE_DIRECTORIES + "${deepgemm_SOURCE_DIR}/deep_gemm/include" + "${deepgemm_SOURCE_DIR}/third-party/cutlass/include") + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_MQA_TOPK_SRC}") + message(STATUS + "Building fused MQA top-k for archs: ${FUSED_MQA_TOPK_ARCHS}") + endif() + # if CUDA endif endif() @@ -1296,6 +1315,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_KDA_DECODE=1) endif() + if(FUSED_MQA_TOPK_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_FUSED_MQA_TOPK=1) + endif() if(FUSED_GDN_DECODE_ARCHS) target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_GDN_DECODE=1) @@ -1714,7 +1737,6 @@ endif() # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") - include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/deepselect.cmake) include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh new file mode 100644 index 000000000000..30afdec28efd --- /dev/null +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// SPDX-License-Identifier: MIT +// Copyright (c) 2025 DeepSeek +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Candidate records of the gated scan: a uint16 value (the low 16 bits of the +// fp24 score code) and an int32 word (key index in bits 0-19, the code's high +// 8 bits in bits 20-27). +#pragma once +#include +#include + +#include "common.cuh" +#include "constants.h" + +namespace fused_mqa_topk { +using CandidateValue = uint16_t; +constexpr uint32_t kCandidateIndexBits = 20; +constexpr uint32_t kCandidateIndexMask = (1u << kCandidateIndexBits) - 1u; +static_assert(kMaxKeys <= (1u << kCandidateIndexBits), + "key indices must fit the candidate index field"); + +CUTLASS_DEVICE uint32_t candidate_pack_index(const uint32_t payload, + const uint32_t kv_index) { + return (kv_index & kCandidateIndexMask) | + ((payload >> 16) << kCandidateIndexBits); +} + +CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) { + return ordered_code(value) >> 8; +} + +CUTLASS_DEVICE uint32_t candidate_load_score_code(const CandidateValue value, + const int32_t packed_idx) { + return (static_cast(packed_idx) >> kCandidateIndexBits) << 16 | + static_cast(value); +} + +CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { + return static_cast(static_cast(packed_idx) & + kCandidateIndexMask); +} + +// Stores a candidate as its fp24 code (low 16 bits in value_dst, high 8 bits +// in bits 20-27 of the index word) and key index. +CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst, + int32_t* index_dst, + const uint32_t payload, + const uint32_t kv_index) { + __stcg(value_dst, static_cast(payload)); + __stcg(index_dst, + static_cast(candidate_pack_index(payload, kv_index))); +} + +CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, + int32_t* index_dst, const float bq, + const uint32_t kv_index) { + store_candidate_record(value_dst, index_dst, candidate_fp24_code(bq), + kv_index); +} + +} // namespace fused_mqa_topk diff --git a/csrc/fused_mqa_topk/common.cuh b/csrc/fused_mqa_topk/common.cuh new file mode 100644 index 000000000000..860bed7c3dcc --- /dev/null +++ b/csrc/fused_mqa_topk/common.cuh @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Small device helpers shared by the persistent kernel, its gate sample and +// the exact rescan. +#pragma once + +#include +#include + +#include + +#include "constants.h" + +namespace fused_mqa_topk { + +// Warp roles of both pipelines (persistent kernel and dense rescan): warps 0-7 +// are the math warps (two warpgroups of 4 query rows each), warps 8-11 the +// producer / UMMA / drain warps. +constexpr uint32_t kMathThreads = 256, kMathWgThreads = 128, + kMathWarps = kMathThreads / 32; +// Query rows of a block each math warpgroup scores. +constexpr uint32_t kWgRows = kBlockRows / 2; +// Register split: ptxas gives each of the 384 threads 168 registers +// (65536 / 384, rounded down to a multiple of 8); setmaxnreg moves them from +// the producer / UMMA / drain warpgroup to the two math warpgroups. +constexpr uint32_t kMathRegs = 224, kSpecRegs = 168 * 3 - 2 * kMathRegs; + +// Named barrier ids (cutlass::arch::NamedBarrier offsets them past the ids it +// reserves). kBarMath / kBarSelect: the 256 math threads; kBarDrain: the math +// threads and the two drain warps; kBarDenseMath: the dense rescan's math +// threads. +constexpr uint32_t kBarDenseMath = 0, kBarMath = 2, kBarDrain = 4, + kBarSelect = 5; +constexpr uint32_t kDrainSyncThreads = kMathThreads + 2 * 32; + +// Barrier of the 256 math threads (warps 0-7); the compiler barriers keep +// shared-memory accesses on their side of it. +__device__ __forceinline__ void math_wg_sync() { + asm volatile("" ::: "memory"); + cutlass::arch::NamedBarrier(kMathThreads, kBarMath).sync(); + asm volatile("" ::: "memory"); +} + +// Shared-memory histograms pad one word per 32 bins (bank conflicts): +// bin b lives at padded_bin(b), a histogram of n bins takes padded_bins(n). +__host__ __device__ constexpr uint32_t padded_bins(uint32_t n) { + return n + n / 32; +} +__device__ __forceinline__ uint32_t padded_bin(uint32_t b) { + return b + b / 32; +} + +// Mask of all 32 lanes for the warp-synchronous intrinsics. +constexpr uint32_t kFullMask = 0xffffffffu; + +// Inclusive prefix sum of v over the warp's lanes. +__device__ __forceinline__ uint32_t warp_inclusive_scan(uint32_t v, + uint32_t lane) { +#pragma unroll + for (uint32_t d = 1; d < 32; d *= 2) { + const uint32_t u = __shfl_up_sync(kFullMask, v, d); + if (lane >= d) v += u; + } + return v; +} + +__device__ __forceinline__ void mbar_inval(const void* b) { + asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"( + static_cast(__cvta_generic_to_shared(b))) + : "memory"); +} + +// Reads `x` with a volatile load, so the compiler re-reads it from (shared) +// memory instead of keeping it live in a register. +template +__device__ __forceinline__ T volatile_load(const T& x) { + return *reinterpret_cast(&x); +} + +// Order-preserving uint32 code of a float (higher code = larger value) and +// its inverse. +__device__ __forceinline__ uint32_t ordered_code(float x) { + const uint32_t u = __float_as_uint(x); + return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); +} +__device__ __forceinline__ float ordered_float(uint32_t code) { + return __uint_as_float((code & 0x80000000u) ? (code ^ 0x80000000u) : ~code); +} + +// Histogram bin of a score: its fp16 value (rounded down) as an +// order-preserving 16-bit code, >> Shift. half_bin_floor is the smallest value +// in bin `b`. +template +__device__ __forceinline__ uint32_t half_bin(float x) { + const uint32_t u = __half_as_ushort(__float2half_rd(x)); + return ((u & 0x8000u) ? ((~u) & 65535u) : (u ^ 0x8000u)) >> Shift; +} +template +__device__ __forceinline__ float half_bin_floor(uint32_t b) { + const uint32_t o = b << Shift, + u = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); + return __half2float(__ushort_as_half(u)); +} + +// Weighted head reduction of one row's 32 head scores s_h in two FFMA2 chains +// seeded with sum_0 / sum_1: sum_h w_h * |s_h| (kWithLinear = false: the +// persistent kernel adds the linear part on the tensor core) or +// sum_h w_h * (s_h + |s_h|) (kWithLinear = true); relu(s) = (s + |s|) / 2. +// Adapted from DeepGEMM's sm100_fp8_mqa_logits.cuh at 88965b078186. +template +__device__ __forceinline__ float reduce_heads(const float* s, const float* w, + float2 sum_0 = {}, + float2 sum_1 = {}) { + auto term = [&](uint32_t at, float2 sum) { + const auto abs_s = make_float2(fabsf(s[at]), fabsf(s[at + 1])); + const auto x = + kWithLinear ? __fadd2_rn(make_float2(s[at], s[at + 1]), abs_s) : abs_s; + return __ffma2_rn(x, make_float2(w[at], w[at + 1]), sum); + }; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { + sum_0 = term(j, sum_0); + sum_1 = term(j + 2, sum_1); + } + const auto sum = __fadd2_rn(sum_0, sum_1); + return sum.x + sum.y; +} + +} // namespace fused_mqa_topk diff --git a/csrc/fused_mqa_topk/constants.h b/csrc/fused_mqa_topk/constants.h new file mode 100644 index 000000000000..c831bbfcfa6e --- /dev/null +++ b/csrc/fused_mqa_topk/constants.h @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Shapes and workspace layout shared by the fused MQA-logits + top-k kernel +// and its host wrapper. The Python side mirrors the workspace sizes in +// vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py (the host wrapper +// checks them). +#pragma once + +#include + +namespace fused_mqa_topk { + +// Indexer shape: 32 heads of dimension 128. +constexpr uint32_t kNumHeads = 32, kHeadDim = 128; +// Query rows per block (the unit a CTA claims), threads per CTA, keys per +// scan tile. +constexpr uint32_t kBlockRows = 8, kNumThreads = 384, kTileKeys = 256; +// Largest key range per call. +constexpr uint32_t kMaxKeys = 1u << 20; +// Per-row workspace: the gated candidate slab (a row with more candidates is +// rescanned exactly) and the exact rescan's retained set. +constexpr uint32_t kCandidatesPerRow = 16384, kRetainedPerRow = 4096; +// Candidate / retained slots (of kBlockRows rows each) per persistent CTA: one +// block in selection, one claimed ahead. +constexpr uint32_t kSlotsPerCta = 2; + +// Top-k sizes the kernel is compiled for. +constexpr bool is_supported_topk(int64_t k) { + return k == 512 || k == 1024 || k == 2048; +} + +} // namespace fused_mqa_topk diff --git a/csrc/fused_mqa_topk/deep_gemm_compat.cuh b/csrc/fused_mqa_topk/deep_gemm_compat.cuh new file mode 100644 index 000000000000..ae13a12c3d57 --- /dev/null +++ b/csrc/fused_mqa_topk/deep_gemm_compat.cuh @@ -0,0 +1,180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Adapted from DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) at +// 88965b078186ee7510ab9fc4f1d5ebc19adfa8d1, the last release before its SM100 +// MQA-logits kernels went MX-only. FP8 MQA-logits building blocks: the +// shared-storage layout with fp32 per-key scales and fp32 weights, the +// grid-stride contiguous-KV scheduler and the 2D/3D TMA load. Trimmed to what +// the fused kernel instantiates (FP8 only, no metadata schedule, no multicast, +// no paging). +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace fused_mqa_topk::dg { +namespace math = deep_gemm::math; + +// FP8 Q/K (one byte per element), fp32 per-key scales, fp32 weights: the +// DeepGEMM layout without its MX / FP4 variants. +template +struct MQALogitsSharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(sizeof(qk_dtype_t) == 1, "FP8 Q/K only"); + + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); + // Align to one 8-row Q/K swizzle tile (head_dim bytes per FP8 token). + static constexpr uint32_t kSwizzleAlignment = 8 * kHeadDim; + static constexpr uint32_t kNumQElementsPerStage = UMMA_N * kHeadDim; + static constexpr uint32_t kNumKVElementsPerStage = SPLIT_KV * kHeadDim; + // TMA destinations in shared memory must be 128-byte aligned. + static constexpr uint32_t kTmaAlignment = 128; + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightBytesPerStage = + math::constexpr_align(BLOCK_Q * kNumWeightBytesPerRow, kTmaAlignment); + static constexpr uint32_t kNumWeightElementsPerStage = + kNumWeightBytesPerStage / static_cast(sizeof(reduce_dtype_t)); + + DG_STATIC_ASSERT(kNumQElementsPerStage % kSwizzleAlignment == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(kNumKVElementsPerStage % kSwizzleAlignment == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(kSwizzleAlignment % 128 == 0, + "TMA destination must be 128-byte aligned"); + + alignas(kSwizzleAlignment) + qk_dtype_t smem_q[kNumQStages][kNumQElementsPerStage]; + alignas(kSwizzleAlignment) + qk_dtype_t smem_kv[kNumKVStages][kNumKVElementsPerStage]; + alignas(kTmaAlignment) float smem_sf_kv[kNumKVStages][SPLIT_KV]; + alignas(kTmaAlignment) + reduce_dtype_t smem_weights[kNumQStages][kNumWeightElementsPerStage]; + // Barriers require 8-byte alignment, already guaranteed by the preceding + // TMA-aligned arrays. + Barrier full_q_barriers[kNumQStages]; + Barrier empty_q_barriers[kNumQStages]; + Barrier full_kv_barriers[kNumKVStages]; + Barrier empty_kv_barriers[kNumKVStages]; + Barrier full_tmem_barriers[kNumTmemStages]; + Barrier empty_tmem_barriers[kNumTmemStages]; + uint32_t tmem_ptr_in_smem; +}; + +// Grid-stride contiguous-KV scheduler: SM `s` owns Q blocks `s`, `s + kNumSMs`, +// ..., each visiting the KV range [min k_start, max k_end) of its rows. +template +struct SM100MQALogitsScheduler { + uint32_t current_q_block_idx; + uint32_t num_q_blocks; + uint32_t num_q_tokens; + uint32_t num_kv_tokens; + const uint32_t* cu_seq_len_k_start; + const uint32_t* cu_seq_len_k_end; + + CUTLASS_DEVICE SM100MQALogitsScheduler(const uint32_t& sm_idx, + const uint32_t& num_q_tokens, + const uint32_t& num_kv_tokens, + const uint32_t* cu_seq_len_k_start, + const uint32_t* cu_seq_len_k_end) + : current_q_block_idx(sm_idx), + num_q_blocks(math::ceil_div(num_q_tokens, BLOCK_Q)), + num_q_tokens(num_q_tokens), + num_kv_tokens(num_kv_tokens), + cu_seq_len_k_start(cu_seq_len_k_start), + cu_seq_len_k_end(cu_seq_len_k_end) { + DG_STATIC_ASSERT(kNumSMs > 0, "Invalid SM count"); + } + + // NOTE: Out-of-range rows are clamped to the last valid row + CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx, + uint32_t& kv_token_base, + uint32_t& num_kv_splits) { + if (current_q_block_idx >= num_q_blocks) return false; + q_block_idx = current_q_block_idx; + current_q_block_idx += kNumSMs; + uint32_t start = cute::numeric_limits::max(); + uint32_t end = cute::numeric_limits::min(); +#pragma unroll 8 + for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++token_idx) { + const auto row_idx = + cute::min(q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1); + start = cute::min(start, + cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens)); + end = cute::max(end, cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens)); + } + kv_token_base = start / 4 * 4; + num_kv_splits = math::ceil_div(end - kv_token_base, SPLIT_KV); + return true; + } + + CUTLASS_DEVICE uint32_t + get_q_tma_token_base(const uint32_t& q_block_idx) const { + return q_block_idx * BLOCK_Q; + } + + CUTLASS_DEVICE static uint32_t get_kv_tma_offset( + const uint32_t& kv_token_base, const uint32_t& kv_split_idx) { + return kv_token_base + kv_split_idx * SPLIT_KV; + } + + CUTLASS_DEVICE static uint32_t get_logits_row(const uint32_t& q_block_idx, + const uint32_t& token_idx) { + return q_block_idx * BLOCK_Q + token_idx; + } + + CUTLASS_DEVICE static uint32_t get_logits_col( + const uint32_t& kv_token_base, const uint32_t& kv_split_idx, + const uint32_t& math_thread_idx) { + return kv_token_base + kv_split_idx * SPLIT_KV + math_thread_idx; + } +}; + +// Single-CTA TMA load of a BLOCK_OUTER x BLOCK_INNER box (3D adds a batch +// coordinate). +template +CUTLASS_DEVICE void tma_copy( + void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr, + dtype_t* smem_ptr, const uint32_t& inner_idx, const uint32_t& outer_idx, + const uint32_t& batch_idx = 0) { + constexpr auto kCacheHint = + static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL); + constexpr uint32_t kPackFactor = deep_gemm::get_smem_pack_factor(); + DG_STATIC_ASSERT(kPackFactor == 1 or sizeof(dtype_t) == 1, + "Packing expects a 1-byte storage type"); + constexpr uint32_t BLOCK_INNER_ATOM_STORAGE = + kSwizzleMode == 0 ? BLOCK_INNER / kPackFactor + : kSwizzleMode / sizeof(dtype_t); + constexpr uint32_t BLOCK_INNER_ATOM_LOGICAL = + BLOCK_INNER_ATOM_STORAGE * kPackFactor; + DG_STATIC_ASSERT(BLOCK_INNER % BLOCK_INNER_ATOM_LOGICAL == 0, + "TMA inner block must contain whole atoms"); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++i) { + if constexpr (kIs3DTMA) { + cute::SM90_TMA_LOAD_3D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx, batch_idx); + } else { + cute::SM90_TMA_LOAD_2D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx); + } + } +} + +} // namespace fused_mqa_topk::dg diff --git a/csrc/fused_mqa_topk/dense_scan.cuh b/csrc/fused_mqa_topk/dense_scan.cuh new file mode 100644 index 000000000000..395cb84f183e --- /dev/null +++ b/csrc/fused_mqa_topk/dense_scan.cuh @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Adapted from DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) at +// 88965b078186ee7510ab9fc4f1d5ebc19adfa8d1: +// deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh. +// +// Dense FP8 MQA scoring of one 8-row block over its whole KV range: DeepGEMM's +// SM100 MQA logits pipeline (contiguous KV, 32 heads, D = 128) with the logits +// store replaced by a per-score policy. Used by the persistent kernel's exact +// recovery rescan (rescan.cuh). +// +// Roles (384 threads): warps 0-7 reduce the 32 heads (4 rows per warpgroup), +// warp 8 loads Q and the weights, warp 9 the KV tiles and their scales, warp 10 +// issues the UMMAs, warp 11 is idle. +// +// Policy interface (all called by the math warps): +// kWeightScale: score = (sum_h kWeightScale * w_h * (s_h + |s_h|)) * scale_kv +// begin_block(sched, q_block, query_base): after the block's Q / weights +// landed +// begin_split(split): before the split's scores +// score(i, kv_offset, kv_lane, split, score): row query_base + i of the +// warpgroup +// end_split(split, num_splits, kv_base): after the split's scores +// end_block(sched, q_block): before the block's Q stage is released +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common.cuh" +#include "constants.h" +#include "deep_gemm_compat.cuh" + +namespace fused_mqa_topk::dense { +using namespace deep_gemm; + +constexpr uint32_t BLOCK_Q = kBlockRows, SPLIT_KV = kTileKeys; +constexpr uint32_t kNumKVStages = 2, kNumTmemStages = 2; +constexpr uint32_t UMMA_M = 128, UMMA_N = BLOCK_Q * kNumHeads, UMMA_K = 32; +// The rescan scores exactly one 8-row block: the scheduler is built with that +// block as its `sm_idx`, and its grid stride (kNumSMs) exceeds any block count +// (rows < 2^26), so next_q_block yields that block only. A smaller stride would +// also score block + stride, outside the rescan's row bounds and output rows. +constexpr uint32_t kSingleBlockStride = 1u << 31; +using Scheduler = + dg::SM100MQALogitsScheduler; +using SharedStorage = + dg::MQALogitsSharedStorage; + +// The policy's state lives in dynamic shared memory right after the pipeline +// storage. Barrier words are zeroed before init: the persistent kernel reuses +// this memory between rescans. +template +CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, + const cute::TmaDescriptor& tensor_map_q, + const cute::TmaDescriptor& tensor_map_kv, + const cute::TmaDescriptor& tensor_map_sf_kv, + const cute::TmaDescriptor& tensor_map_weights) { + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr uint32_t kSpecWarpStart = 8; + constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = BLOCK_Q * kNumHeads * kHeadDim; + constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = SPLIT_KV * kHeadDim; + constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = SPLIT_KV * sizeof(float); + constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = + BLOCK_Q * kNumHeads * sizeof(float); + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols(); + + if (warp_idx == kSpecWarpStart) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_weights); + cute::prefetch_tma_descriptor(&tensor_map_kv); + cute::prefetch_tma_descriptor(&tensor_map_sf_kv); + } + + extern __shared__ __align__(SharedStorage::kSwizzleAlignment) + uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { + const auto init = [](auto& barrier, uint32_t count) { + *reinterpret_cast(&barrier) = 0; + barrier.init(count); + }; + init(smem.full_q_barriers[0], 1); + init(smem.empty_q_barriers[0], kMathThreads + 32); // + the UMMA warp +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + init(smem.full_kv_barriers[i], 1); + init(smem.empty_kv_barriers[i], kMathThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumTmemStages; ++i) { + init(smem.full_tmem_barriers[i], 1); + init(smem.empty_tmem_barriers[i], kMathThreads); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + if (warp_idx == kSpecWarpStart + 2) + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); + __syncthreads(); + + RingPipeline<1> q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + + uint32_t q_block_idx, kv_base, num_kv_splits; + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc(); + if (cute::elect_one_sync()) { + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.empty_q_barriers[q_stage_idx].wait(q_phase ^ 1); + const uint32_t q_token_base = + scheduler.get_q_tma_token_base(q_block_idx); + dg::tma_copy( + &tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); + dg::tma_copy( + &tensor_map_weights, &smem.full_q_barriers[q_stage_idx], + smem.smem_weights[q_stage_idx], 0, q_token_base); + smem.full_q_barriers[q_stage_idx].arrive_and_expect_tx( + SMEM_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE); + } + } + __syncwarp(); + } else if (warp_idx == kSpecWarpStart + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if (cute::elect_one_sync()) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); + const uint32_t kv_tma_offset = + scheduler.get_kv_tma_offset(kv_base, kv_split_idx); + dg::tma_copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); + dg::tma_copy( + &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_sf_kv[kv_stage_idx], kv_tma_offset, 0); + smem.full_kv_barriers[kv_stage_idx].arrive_and_expect_tx( + SMEM_KV_SIZE_PER_STAGE + SMEM_SF_KV_SIZE_PER_STAGE); + } + __syncwarp(); + } + } + } else if (warp_idx == kSpecWarpStart + 2) { + cutlass::arch::warpgroup_reg_dealloc(); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); + auto instr_desc = cute::UMMA::make_instr_desc< + cutlass::float_e4m3_t, cutlass::float_e4m3_t, float, UMMA_M, UMMA_N, + cute::UMMA::Major::K, cute::UMMA::Major::K>(); + auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); + if (cute::elect_one_sync()) { + // One UMMA per math warpgroup: KV rows [128 i, 128 i + 128) x all 256 + // Q rows. +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + const uint32_t tmem_addr = tmem_stage_idx * UMMA_N; + smem.empty_tmem_barriers[tmem_stage_idx].wait(tmem_phase ^ 1); + ptx::tcgen05_after_thread_sync(); +#pragma unroll + for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { + auto a_desc = mma::sm100::make_umma_desc( + smem.smem_kv[kv_stage_idx], i * UMMA_M, k * UMMA_K); + auto b_desc = mma::sm100::make_umma_desc( + smem.smem_q[q_stage_idx], 0, k * UMMA_K); + ptx::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, tmem_addr, k, + runtime_instr_desc); + } + ptx::umma_arrive_no_elect(smem.full_tmem_barriers[tmem_stage_idx]); + } + } + __syncwarp(); + } + smem.empty_q_barriers[q_stage_idx].arrive(); + } + } else if (warp_idx == kSpecWarpStart + 3) { + cutlass::arch::warpgroup_reg_dealloc(); + } else { + cutlass::arch::warpgroup_reg_alloc(); + const uint32_t math_thread_idx = warp_idx * 32 + lane_idx; + const uint32_t query_base = (warp_idx / 4) * 4; + float weights[4][kNumHeads]; + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; ++j) + weights[i][j] = Policy::kWeightScale * + ptx::ld_shared(smem.smem_weights[q_stage_idx] + + (query_base + i) * kNumHeads + j); + policy.begin_block(scheduler, q_block_idx, query_base); + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + policy.begin_split(kv_split_idx); + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); +#pragma unroll + for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) { + const uint32_t kv_lane = + (math_thread_idx % kMathWgThreads) + kv_half * kMathWgThreads; + const uint32_t kv_offset = + scheduler.get_logits_col(kv_base, kv_split_idx, kv_lane); + const float scale_kv = + ptx::ld_shared(smem.smem_sf_kv[kv_stage_idx] + kv_lane); + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + smem.full_tmem_barriers[tmem_stage_idx].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + // Both KV half-MMAs have committed before the shared stage is freed. + if (kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + smem.empty_kv_barriers[kv_stage_idx].arrive(); + } + float accum_pair[2][kNumHeads]; +#pragma unroll + for (uint32_t pair = 0; pair < 4; pair += 2) { + // Two query rows at a time keep the live accumulator registers low. + ptx::tmem_load_32dp32b<2 * kNumHeads>( + tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads, + reinterpret_cast(accum_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair == 2) { + ptx::tcgen05_before_thread_sync(); + smem.empty_tmem_barriers[tmem_stage_idx].arrive(); + } +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) + policy.score( + pair + r, kv_offset, kv_lane, kv_split_idx, + reduce_heads(accum_pair[r], weights[pair + r]) * + scale_kv); + } + } + policy.end_split(kv_split_idx, num_kv_splits, kv_base); + } + policy.end_block(scheduler, q_block_idx); + cutlass::arch::fence_view_async_shared(); + smem.empty_q_barriers[q_stage_idx].arrive(); + } + cutlass::arch::NamedBarrier(kMathThreads, kBarDenseMath).sync(); + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +// Invalidates the pipeline's barriers so the next score_block can re-initialise +// the same shared memory (mbarrier.inval before reuse). +CUTLASS_DEVICE void invalidate_barriers(SharedStorage& smem) { + for (auto& b : smem.full_q_barriers) mbar_inval(&b); + for (auto& b : smem.empty_q_barriers) mbar_inval(&b); + for (auto& b : smem.full_kv_barriers) mbar_inval(&b); + for (auto& b : smem.empty_kv_barriers) mbar_inval(&b); + for (auto& b : smem.full_tmem_barriers) mbar_inval(&b); + for (auto& b : smem.empty_tmem_barriers) mbar_inval(&b); +} + +} // namespace fused_mqa_topk::dense diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu new file mode 100644 index 000000000000..652453494fbe --- /dev/null +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -0,0 +1,394 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Single-launch fused FP8 MQA-logits + top-k (k = 512, 1024, 2048) prefill +// selection (SM100 family): one persistent kernel (persistent_kernel.cuh, one +// CTA per SM; sampled gate, candidate emission and selection, exact rescan of +// out-of-range blocks) scores every visible key of a row and selects its top-k +// keys without materializing the [rows, keys] logits matrix. +// This file is the libtorch-stable host wrapper. All scratch memory is passed +// in by the caller (no allocation here), so the op is CUDA-graph friendly. + +#include "../torch_utils.h" + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "constants.h" +#include "persistent_kernel.cuh" + +namespace { + +using fused_mqa_topk::kBlockRows; +using fused_mqa_topk::kCandidatesPerRow; +using fused_mqa_topk::kHeadDim; +using fused_mqa_topk::kMaxKeys; +using fused_mqa_topk::kNumHeads; +using fused_mqa_topk::kNumThreads; +using fused_mqa_topk::kRetainedPerRow; +using fused_mqa_topk::kSlotsPerCta; +using fused_mqa_topk::kTileKeys; +using torch::headeronly::ScalarType; +namespace persist = fused_mqa_topk::persist; + +void check_cuda(cudaError_t err) { + STD_TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err)); +} + +// cuTensorMapEncodeTiled from the driver, resolved once through the runtime +// (thread-safe static initialisation; no link-time libcuda dependency). +PFN_cuTensorMapEncodeTiled_v12000 encode_tiled() { + static const PFN_cuTensorMapEncodeTiled_v12000 fn = [] { + void* ptr = nullptr; + cudaDriverEntryPointQueryResult status; + const cudaError_t err = cudaGetDriverEntryPointByVersion( + "cuTensorMapEncodeTiled", &ptr, 12000, cudaEnableDefault, &status); + STD_TORCH_CHECK(err == cudaSuccess && status == cudaDriverEntryPointSuccess, + "failed to resolve cuTensorMapEncodeTiled"); + return reinterpret_cast(ptr); + }(); + return fn; +} + +// SM count of `device`, queried once per device. +int sm_count(int device) { + static std::mutex mutex; + static std::map cache; + const std::lock_guard lock(mutex); + auto [it, inserted] = cache.try_emplace(device, 0); + if (inserted) { + const cudaError_t err = cudaDeviceGetAttribute( + &it->second, cudaDevAttrMultiProcessorCount, device); + if (err != cudaSuccess) cache.erase(it); + check_cuda(err); + } + return it->second; +} + +// Adapted from DeepGEMM's make_tma_2d_desc +// (csrc/jit_kernels/impls/runtime_utils.hpp at 88965b078186) +CUtensorMap make_2d(void* ptr, CUtensorMapDataType dt, int elem_size, + int gmem_inner, int gmem_outer, int smem_inner, + int smem_outer, long gmem_outer_stride, int swizzle_mode) { + if (swizzle_mode != 0) smem_inner = swizzle_mode / elem_size; + CUtensorMap tm; + const cuuint64_t gdims[2] = {(cuuint64_t)gmem_inner, (cuuint64_t)gmem_outer}; + const cuuint32_t sdims[2] = {(cuuint32_t)smem_inner, (cuuint32_t)smem_outer}; + const cuuint64_t gstrides[1] = {(cuuint64_t)(gmem_outer_stride * elem_size)}; + const cuuint32_t estrides[2] = {1, 1}; + CUtensorMapSwizzle swizzle = swizzle_mode == 128 ? CU_TENSOR_MAP_SWIZZLE_128B + : swizzle_mode == 64 ? CU_TENSOR_MAP_SWIZZLE_64B + : swizzle_mode == 32 + ? CU_TENSOR_MAP_SWIZZLE_32B + : CU_TENSOR_MAP_SWIZZLE_NONE; + CUresult r = encode_tiled()(&tm, dt, 2, ptr, gdims, gstrides, sdims, estrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + STD_TORCH_CHECK(r == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", (int)r); + return tm; +} + +// KV [keys,128] bytes viewed as {128, groups, stride} with strides +// {128*stride, 128}: box {128, 256, 1} at (0, 0, o) loads keys o, o+stride, ... +// with 128B swizzle. The view spans groups * stride keys, which can exceed +// `keys` by < stride / 2; that is safe only because the kernel loads offset +// o = stride / 2, whose last key ((groups - 1) * stride + stride / 2) is +// always < keys (launch_config sizes `groups` for that). +CUtensorMap make_strided_kv(void* ptr, uint32_t groups, uint32_t stride) { + CUtensorMap tm; + constexpr uint64_t kRowBytes = kHeadDim; + const cuuint64_t gdims[3] = {kRowBytes, groups, stride}; + const cuuint64_t gstrides[2] = {kRowBytes * stride, kRowBytes}; + const cuuint32_t box[3] = {kRowBytes, kTileKeys, 1}; + const cuuint32_t estrides[3] = {1, 1, 1}; + CUresult r = encode_tiled()( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, ptr, gdims, gstrides, box, + estrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + STD_TORCH_CHECK(r == CUDA_SUCCESS, + "strided KV cuTensorMapEncodeTiled failed: ", (int)r); + return tm; +} + +// -------- launch configuration -------- +// Keys per row above which the long-context instantiations (Interval 64, 128) +// and gate settings apply. +constexpr uint32_t kLongContextKeys = 131072; + +using PersistentKernel = + decltype(&persist::fused_mqa_topk_fp8_prefill_kernel<16, 2048>); + +// vLLM dispatches chunks of >= 32K keys per row here (fused_mqa_topk_min_keys); +// smaller key counts are supported but take the same <= 32K configuration. +template +PersistentKernel persistent_kernel(uint32_t keys) { + return keys <= 32768 ? persist::fused_mqa_topk_fp8_prefill_kernel<16, TopK> + : keys <= kLongContextKeys + ? persist::fused_mqa_topk_fp8_prefill_kernel<32, TopK> + : keys <= 524288 + ? persist::fused_mqa_topk_fp8_prefill_kernel<64, TopK> + : persist::fused_mqa_topk_fp8_prefill_kernel<128, TopK>; +} + +struct LaunchConfig { + PersistentKernel kernel; + int smem; // dynamic shared memory bytes + uint32_t phase_bytes; // selection scratch bytes past the Q / weights region + // Gate sample: sample_tiles tiles of 256 keys on a uniform grid (stride + // sample_stride). + uint32_t sample_target, sample_stride, sample_groups, sample_tiles; +}; + +// Dynamic shared memory of `kernel` on `device`: everything past its static +// shared memory up to the device's opt-in limit (232448 bytes on SM100; the +// staged selection uses it all), at least the exact rescan's layout. Sets +// the opt-in attribute; both are computed once per device and kernel. +cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) { + // The exact rescan reuses the same dynamic shared memory with its own layout. + constexpr int kRescanSmem = int(sizeof(fused_mqa_topk::dense::SharedStorage) + + sizeof(persist::RescanShared)); + constexpr int kMinSmem = int( + sizeof(persist::KernelSmem) + + std::max(sizeof(persist::HandoffRing), sizeof(persist::SampleShared))); + static std::mutex mutex; + static std::map, int> cache; + const std::lock_guard lock(mutex); + if (const auto it = cache.find({device, kernel}); it != cache.end()) { + smem = it->second; + return cudaSuccess; + } + int max_smem = 0; + if (const cudaError_t e = cudaDeviceGetAttribute( + &max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + e != cudaSuccess) + return e; + cudaFuncAttributes attr; + if (const cudaError_t e = + cudaFuncGetAttributes(&attr, reinterpret_cast(kernel)); + e != cudaSuccess) + return e; + const int bytes = std::max(kRescanSmem, max_smem - int(attr.sharedSizeBytes)); + if (bytes < kMinSmem) return cudaErrorInvalidConfiguration; + if (const cudaError_t e = cudaFuncSetAttribute( + reinterpret_cast(kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + e != cudaSuccess) + return e; + cache.emplace(std::make_pair(device, kernel), bytes); + smem = bytes; + return cudaSuccess; +} + +// Picks the kernel instantiation, its shared memory and the gate sample for +// `keys` keys and top-`topk` rows on `device` (the current device). +cudaError_t launch_config(int device, uint32_t keys, uint32_t topk, + LaunchConfig& c) { + // The gate aims for `target` candidates per row and sizes the sample so its + // quantile rank is `rank` (count error ~1/sqrt(rank)); a row whose count + // falls below the top-k is rescanned exactly. Up to 128K keys the candidate + // stream dominates, so smaller top-ks aim lower with a coarser sample. Above + // 128K the sample is clamped to 16-32 tiles (rank ~32 at 1M keys): a sample + // tile costs about two scan tiles, so a larger sample costs more than a + // smaller candidate set saves. Tuned on B300 real captures. + struct Gate { + uint32_t target, rank; + }; + constexpr Gate kGates[3][2] = { + // <= 128K > 128K + {{1024, 32}, {1536, 64}}, // top-512 + {{1536, 48}, {2048, 48}}, // top-1024 + {{3072, 64}, {4096, 64}}, // top-2048 + }; + const bool long_ctx = keys > kLongContextKeys; + const Gate gate = + kGates[topk == 512 ? 0 : (topk == 1024 ? 1 : 2)][long_ctx ? 1 : 0]; + c.sample_target = gate.target; + const uint32_t rank = gate.rank; + const uint64_t tile = kTileKeys; + const uint32_t rank_tiles = std::max( + 1u, uint32_t((uint64_t(rank) * keys + tile * c.sample_target - 1) / + (tile * c.sample_target))); + c.sample_tiles = + long_ctx ? std::min(32u, std::max(16u, rank_tiles)) : rank_tiles; + c.sample_stride = std::max(1u, uint32_t((keys + tile * c.sample_tiles - 1) / + (tile * c.sample_tiles))); + c.sample_groups = + (keys - c.sample_stride / 2 + c.sample_stride - 1) / c.sample_stride; + c.kernel = topk == 512 ? persistent_kernel<512>(keys) + : topk == 1024 ? persistent_kernel<1024>(keys) + : persistent_kernel<2048>(keys); + if (const cudaError_t e = kernel_smem(device, c.kernel, c.smem); + e != cudaSuccess) + return e; + c.phase_bytes = uint32_t(c.smem) - persist::kSelBase; + return cudaSuccess; +} + +void check_tensor(const torch::stable::Tensor& t, int32_t device, + ScalarType dtype, const char* name) { + STD_TORCH_CHECK(t.is_cuda() && t.get_device_index() == device, name, + " must be on the same CUDA device as q"); + STD_TORCH_CHECK(t.scalar_type() == dtype, name, " has the wrong dtype"); + STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +// TMA tensor maps need 16-byte aligned base addresses. +void check_aligned(const torch::stable::Tensor& t, const char* name) { + STD_TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +} // namespace + +// Inputs: q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32, +// weights [rows,32] fp32, starts/ends [rows] int32 (the visible key range +// [start,end) of each row), with 0 < rows < 2^26 and 0 < keys <= 1048576 (the +// last 8-row block may be partial). starts/ends are expected in [0, keys] with +// start <= end; they are read as uint32_t and clamped to keys, so a negative +// value means `keys` and an end before its start means an empty row. +// +// Output: out [rows,K] int32 with K (the top-k) 512, 1024 or 2048: key indices +// relative to the row's start (key - start, the top_k_per_row_prefill +// convention), unordered; -1 padded when fewer than K keys are visible. Rows +// that see at most K keys get them all, in order. Exact up to ties: the gated +// path ranks fp24 score codes (the low 8 mantissa bits dropped), the exact +// rescan ranks fp32 scores. +// +// All tensors are contiguous; q, k, scales, weights and out are 16-byte aligned +// (TMA / vector stores). +// +// Workspace (caller-owned, reused across calls), with B = ceil(rows / 8) and +// S >= 2 slots per CTA of the persistent grid (one CTA per SM): +// counter [2] int32 (zero between calls; the kernel leaves it zeroed, so a +// workspace must not be shared by launches that can run concurrently, e.g. +// on two streams), +// values [8S,16384] uint16 (low 16 bits of the fp24 candidate codes), +// indices [8S,16384] int32 (key index in bits 0-19, the code's high 8 bits +// in bits 20-27), +// retained [8S,4096] fp32, retained_indices [8S,4096] int32, +// counts [>=8B] int32, defer [>=B] int32. +void fused_mqa_topk_prefill( + torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor scales, torch::stable::Tensor weights, + torch::stable::Tensor starts, torch::stable::Tensor ends, + torch::stable::Tensor out, torch::stable::Tensor counter, + torch::stable::Tensor values, torch::stable::Tensor indices, + torch::stable::Tensor retained, torch::stable::Tensor retained_indices, + torch::stable::Tensor counts, torch::stable::Tensor defer) { + STD_TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor"); + const int32_t dev = q.get_device_index(); + check_tensor(q, dev, ScalarType::Float8_e4m3fn, "q"); + check_tensor(k, dev, ScalarType::Float8_e4m3fn, "k"); + check_tensor(scales, dev, ScalarType::Float, "scales"); + check_tensor(weights, dev, ScalarType::Float, "weights"); + check_tensor(starts, dev, ScalarType::Int, "starts"); + check_tensor(ends, dev, ScalarType::Int, "ends"); + check_tensor(out, dev, ScalarType::Int, "out"); + check_tensor(counter, dev, ScalarType::Int, "counter"); + check_tensor(values, dev, ScalarType::UInt16, "values"); + check_tensor(indices, dev, ScalarType::Int, "indices"); + check_tensor(retained, dev, ScalarType::Float, "retained"); + check_tensor(retained_indices, dev, ScalarType::Int, "retained_indices"); + check_tensor(counts, dev, ScalarType::Int, "counts"); + check_tensor(defer, dev, ScalarType::Int, "defer"); + + STD_TORCH_CHECK( + q.dim() == 3 && q.size(1) == kNumHeads && q.size(2) == kHeadDim, + "q must be [rows, 32, 128]"); + STD_TORCH_CHECK(k.dim() == 2 && k.size(1) == kHeadDim, + "k must be [keys, 128]"); + const int64_t rows = q.size(0), keys = k.size(0); + // TMA row coordinates (rows * 32 Q rows) are int32. + STD_TORCH_CHECK(rows > 0 && rows < (int64_t(1) << 26), + "rows must be in [1, 2^26)"); + STD_TORCH_CHECK(keys > 0 && keys <= kMaxKeys, "keys must be in [1, 1048576]"); + const int64_t blocks = (rows + kBlockRows - 1) / kBlockRows; + STD_TORCH_CHECK(scales.numel() == keys, "scales must be [keys]"); + STD_TORCH_CHECK(weights.dim() == 2 && weights.size(0) == rows && + weights.size(1) == kNumHeads, + "weights must be [rows, 32]"); + STD_TORCH_CHECK(starts.numel() == rows && ends.numel() == rows, + "starts/ends must be [rows]"); + STD_TORCH_CHECK(out.dim() == 2 && out.size(0) == rows && + fused_mqa_topk::is_supported_topk(out.size(1)), + "out must be [rows, K] with K in {512, 1024, 2048}"); + check_aligned(q, "q"); + check_aligned(k, "k"); + check_aligned(scales, "scales"); + check_aligned(weights, "weights"); + check_aligned(out, "out"); + + torch::stable::accelerator::DeviceGuard guard(dev); + const cudaStream_t stream = get_current_cuda_stream(dev); + const uint32_t urows = static_cast(rows), + ukeys = static_cast(keys), + ublocks = static_cast(blocks), + topk = static_cast(out.size(1)); + const auto* starts_p = static_cast(starts.data_ptr()); + const auto* ends_p = static_cast(ends.data_ptr()); + auto* out_p = static_cast(out.data_ptr()); + // Q: rows * 32 head rows of 128 bytes, a block's 256 per box; K: keys rows of + // 128 bytes, a tile's 256 per box (both 128B-swizzled); per-key scales: a + // tile's 256; weights: a block's 8 rows of 32. + const int head_rows = int(rows) * int(kNumHeads); + auto tq = make_2d(q.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, kHeadDim, + head_rows, kHeadDim, kBlockRows * kNumHeads, kHeadDim, 128); + auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, kHeadDim, + int(keys), kHeadDim, kTileKeys, kHeadDim, 128); + auto ts = make_2d(scales.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, + int(keys), 1, kTileKeys, 1, 0, 0); + auto tw = make_2d(weights.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, + kNumHeads, int(rows), kNumHeads, kBlockRows, kNumHeads, 0); + + STD_TORCH_CHECK(counter.numel() == 2, "counter must be [2]"); + constexpr int64_t kCtaRows = kSlotsPerCta * kBlockRows; + STD_TORCH_CHECK(values.dim() == 2 && values.size(1) == kCandidatesPerRow && + values.size(0) % kCtaRows == 0 && indices.dim() == 2 && + indices.size(0) == values.size(0) && + indices.size(1) == kCandidatesPerRow, + "values/indices must be [8 * slots, 16384], slots even"); + STD_TORCH_CHECK(values.size(0) > 0, "the candidate workspace is empty"); + STD_TORCH_CHECK(retained.dim() == 2 && retained.size(1) == kRetainedPerRow && + retained.size(0) == values.size(0) && + retained_indices.dim() == 2 && + retained_indices.size(0) == retained.size(0) && + retained_indices.size(1) == kRetainedPerRow, + "retained/retained_indices must be [8 * slots, 4096]"); + STD_TORCH_CHECK( + counts.numel() >= blocks * kBlockRows && defer.numel() >= blocks, + "per-row workspace too small"); + // One persistent CTA per SM (capped by the blocks and the workspace slots); + // CTA c owns candidate slots 2c and 2c + 1. + const uint32_t grid = static_cast(std::min( + {int64_t(sm_count(dev)), blocks, values.size(0) / kCtaRows})); + LaunchConfig c; + const cudaError_t config_err = launch_config(dev, ukeys, topk, c); + STD_TORCH_CHECK(config_err != cudaErrorInvalidConfiguration, + "fused_mqa_topk_prefill needs more shared memory per block " + "than this GPU allows"); + check_cuda(config_err); + auto tks = make_strided_kv(k.data_ptr(), c.sample_groups, c.sample_stride); + const fused_mqa_topk::rescan::RescanBuffers rescan{ + static_cast(retained.data_ptr()), + static_cast(retained_indices.data_ptr())}; + c.kernel<<>>( + urows, ukeys, starts_p, ends_p, out_p, + static_cast(values.data_ptr()), + static_cast(indices.data_ptr()), + static_cast(counts.data_ptr()), rescan, c.phase_bytes, + c.sample_target, c.sample_stride, c.sample_groups, c.sample_tiles, + static_cast(scales.data_ptr()), tks, tq, ts, tk, tw, + ublocks, static_cast(counter.data_ptr()), + static_cast(defer.data_ptr())); + check_cuda(cudaGetLastError()); +} diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc new file mode 100644 index 000000000000..1af6a1d2088b --- /dev/null +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -0,0 +1,742 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Body of one persistent epoch (included twice by persistent_kernel.cuh: the +// first epoch is straight-line code so the hot loops are not nested in an outer +// loop). +if (warp_idx == 9 && cute::elect_one_sync()) { + if (epoch) + for (uint32_t i = 0; i < BarIdx::total; ++i) mbar_inval(bar_words + i); + for (uint32_t i = 0; i < BarIdx::total; ++i) + reinterpret_cast(bar_words)[i] = 0; + full_q.init(1); + empty_q.init(kMathThreads + 32); // math threads + the UMMA warp + for (uint32_t i = 0; i < kKVStages; ++i) { + full_kv[i].init(kFullKvArrivals); + empty_kv[i].init(kMathThreads); + } + for (uint32_t i = 0; i < BarIdx::kTmemSlots; ++i) { + full_tmem[i].init(1); + empty_tmem[i].init(kMathWgThreads); + } + full_lin->init(32); + // Every reader lane arrives on claim_empty. + for (uint32_t i = 0; i < 2; ++i) { + claim_full[i].init(32); + claim_empty[i].init(kClaimConsumers * 32); + } + sel_done.init(kMathWarps); // one lane per math warp + for (uint32_t i = 0; i < 2 * kMathWarps; ++i) { + h_full[i].init(32); + h_empty[i].init(32); + } + cutlass::arch::fence_barrier_init(); + unsafe_rows = 0; +} +__syncwarp(); +if (warp_idx == 10) cute::TMEM::Allocator1Sm().allocate(512, &tmem_ptr); +__syncthreads(); + +if (warp_idx == 8) { + // -------- claimer + Q/weights producer + drain of WG1 rows -------- + cutlass::arch::warpgroup_reg_dealloc(); + // n_blk: blocks this CTA processed (all-short blocks excluded). + uint32_t n_blk = 0, drain_interval = 0; + for (uint32_t k = 0;; ++k) { + if (lane == 0) { + empty_q.wait((n_blk & 1) ^ 1); + } + __syncwarp(); + // Every lane waits: lanes 0-7 write the entry. + claim_empty[k & 1].wait(((k >> 1) & 1) ^ 1); + uint32_t b = 0; + if (lane == 0) + b = atomicOr(&pause_req, 0u) ? kPause : atomicAdd(counter, 1u); + b = __shfl_sync(kFullMask, b, 0); + ClaimEntry& e = entries[k & 1]; + if (b >= nblocks) { + if (lane == 0) { + const uint32_t kind = b == kPause ? kPause : kTerm; + e.vblock = kind; + epoch_kind = kind; + } + __syncwarp(); + claim_full[k & 1].arrive(); // every writer lane arrives + break; + } + // Blocks are claimed heaviest (last rows, longest context) first. + const uint32_t vblock = nblocks - 1 - b, row_base = vblock * kBlockRows; + uint32_t st = 0, en = 0; + if (lane < kBlockRows) { + const uint32_t r = row_base + lane; + if (r < rows) { + st = min(starts_g[r], keys); + en = max(min(ends_g[r], keys), st); + } + e.starts[lane] = st; + e.ends[lane] = en; + } + const bool all_short = + __ballot_sync(kFullMask, lane < kBlockRows && en - st > topk) == 0; + const bool gated = + __ballot_sync(kFullMask, + lane < kBlockRows && en - st > sample_target) != 0; + // KV range: rows past the end are clamped to the last valid row (DeepGEMM + // scheduler). + const uint32_t src = min(row_base + (lane & 7), rows - 1) - row_base; + uint32_t lo = __shfl_sync(kFullMask, st, src), + hi = __shfl_sync(kFullMask, en, src); +#pragma unroll + for (int d = 4; d >= 1; d >>= 1) { + lo = min(lo, __shfl_xor_sync(kFullMask, lo, d)); + hi = max(hi, __shfl_xor_sync(kFullMask, hi, d)); + } + const uint32_t kv_base = lo / 4 * 4, + splits = (hi - kv_base + kTileKeys - 1) / kTileKeys; + // CTA-private slot pair: block n of this CTA uses slot 2 * cta + (n & 1); + // at most two blocks (one in selection, one claimed ahead) are in flight, + // so no claim atomics. + const uint32_t slot = kSlotsPerCta * blockIdx.x + (n_blk & 1); + if (!all_short && lane < kBlockRows) row_counts[row_base + lane] = 0; + if (lane == 0) { + e.vblock = vblock; + e.all_short = all_short; + e.kv_base = kv_base; + e.splits = splits; + e.samples = gated ? sample_tiles : 0u; + e.slot = slot; + } + __syncwarp(); + claim_full[k & 1].arrive(); // every writer lane arrives + if (all_short) continue; + if (lane == 0) { + dg::tma_copy(&tq, &full_q, sm.q, 0, + row_base * kNumHeads); + dg::tma_copy(&tw, &full_q, sm.weights, 0, + row_base); + full_q.arrive_and_expect_tx(kQBytes + kWBytes); + } + __syncwarp(); + ++n_blk; + const int64_t row_delta = slot_row_delta(slot, row_base); + for (uint32_t base = 0; base < splits; base += Interval, ++drain_interval) { + for (uint32_t p = kMathWarps / 2; p < kMathWarps; ++p) + drain_handoff(hdata, h_full, h_empty, p, drain_interval, row_base + 4, + row_counts, base + Interval >= splits, + [&](uint32_t row, uint32_t slot2, uint32_t record, + uint32_t batch_base, uint32_t source) { + store_ring_record(vals, indices, row, row_delta, slot2, + record, kv_base, batch_base, + source % kMathWgThreads); + }); + } + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); + } +} else if (warp_idx == 9) { + // -------- KV producer: sample tiles then scan tiles, one ring -------- + cutlass::arch::warpgroup_reg_dealloc(); + RingPipeline kv_pipe; + uint32_t n_blk = 0; + for (uint32_t k = 0;; ++k) { + ClaimHeader claim; + if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break; + if (claim.all_short) continue; + const uint32_t kv_base = claim.kv_base, splits = claim.splits, + samples = claim.samples; + { + // The whole warp stages the strided sample scales into sf_kv[stage] + // (prefetched one tile ahead in registers); lane 0 issues the KV TMA. + if (lane == 0 && n_blk > 0) sel_done.wait((n_blk - 1) & 1); + __syncwarp(); + const uint32_t center = sample_stride / 2; + // Per-key scales of one sample tile, kScalesPerLane per lane. + constexpr uint32_t kScalesPerLane = kTileKeys / 32; + float pf[kScalesPerLane]; +#pragma unroll + for (uint32_t j = 0; j < kScalesPerLane; ++j) { + const uint32_t g = j * 32 + lane; + pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center) + : 0.0f; + } +#pragma unroll 1 + for (uint32_t t = 0; t < samples; ++t) { + // Only lane 0's ring state is maintained across blocks (the scan loop + // runs on lane 0). + CUTE_TIE_DECL(kv_pipe.advance(), s0, ph); + const uint32_t s = __shfl_sync(kFullMask, s0, 0); + // Every writing lane acquires and releases the stage itself. + empty_kv[s].wait(__shfl_sync(kFullMask, ph, 0) ^ 1); +#pragma unroll + for (uint32_t j = 0; j < kScalesPerLane; ++j) + sm.sf_kv[s][j * 32 + lane] = pf[j]; + // Generic writes before later TMA writes of this stage. + cutlass::arch::fence_view_async_shared(); + if (lane == 0) { + full_kv[s].expect_transaction(kKVBytes); + dg::tma_copy( + &tks, &full_kv[s], sm.kv[s], 0, t * kTileKeys, center); + } + full_kv[s].arrive(); +#pragma unroll + for (uint32_t j = 0; j < kScalesPerLane; ++j) { + const uint32_t g = (t + 1) * kTileKeys + j * 32 + lane; + pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center) + : 0.0f; + } + } + } + if (lane == 0) { +#pragma unroll 1 + for (uint32_t t = 0; t < splits; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), s, ph); + empty_kv[s].wait(ph ^ 1); + const uint32_t off = kv_base + t * kTileKeys; + dg::tma_copy(&tk, &full_kv[s], sm.kv[s], 0, + off); + dg::tma_copy(&ts, &full_kv[s], sm.sf_kv[s], off, 0); + mbar_arrive_expect_tx(&full_kv[s], kKVBytes + kScaleBytes, + kFullKvArrivals); + } + } + __syncwarp(); + ++n_blk; + } +} else if (warp_idx == 10) { + // -------- UMMA issuer -------- + cutlass::arch::warpgroup_reg_dealloc(); + RingPipeline<1> q_pipe; + RingPipeline kv_pipe; + RingPipeline tmem_pipe; + auto instr_desc = + cute::UMMA::make_instr_desc(); + auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + for (uint32_t k = 0;; ++k) { + ClaimHeader claim; + if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break; + if (claim.all_short) continue; + const uint32_t splits = claim.splits, samples = claim.samples; + CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase); + (void)q_stage; + full_q.wait(q_phase); + full_lin->wait(q_phase); // this block's linear-term rows (warp 11) + ptx::tcgen05_after_thread_sync(); + const uint32_t total = samples + splits; +#pragma unroll 1 + for (uint32_t t = 0; t < total; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), s, ph); + full_kv[s].wait(ph); + if (cute::elect_one_sync()) { + // Per KV half, one N=144 MMA per math warpgroup into its own TMEM slot + // (ring of 3; WG0 takes even, WG1 odd positions of the issue sequence). +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { +#pragma unroll + for (uint32_t g = 0; g < 2; ++g) { + CUTE_TIE_DECL(tmem_pipe.advance(), ts_idx, ts_ph); + const uint32_t tmem_addr = ts_idx * kSlotCols; + empty_tmem[ts_idx].wait(ts_ph ^ 1); + ptx::tcgen05_after_thread_sync(); + cutlass::float_e4m3_t* b_base = + g == 0 ? reinterpret_cast(sm.lin_lo) + : sm.q + 4 * kNumHeads * kHeadDim; // WG1's rows +#pragma unroll + for (uint32_t kk = 0; kk < 4; ++kk) { + auto a_desc = + mma::sm100::make_umma_desc( + sm.kv[s], i * 128, kk * 32); + auto b_desc = + mma::sm100::make_umma_desc( + b_base, 0, kk * 32); + ptx::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, tmem_addr, kk, + runtime_instr_desc); + } + ptx::umma_arrive_no_elect(full_tmem[ts_idx]); + } + } + } + __syncwarp(); + } + empty_q.arrive(); + } +} else if (warp_idx == 11) { + // -------- linear-term rows + drain of WG0 rows -------- + cutlass::arch::warpgroup_reg_dealloc(); + uint32_t drain_interval = 0, n_lin = 0; + for (uint32_t k = 0;; ++k) { + ClaimHeader claim; + if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break; + if (claim.all_short) continue; + const uint32_t vblock = claim.vblock, kv_base = claim.kv_base, + splits = claim.splits, slot = claim.slot; + // Q of block k lands while the math warps still select block k-1, so this + // is off the critical path except for a CTA's first block. + full_q.wait(n_lin & 1); + ++n_lin; + compute_lin_rows(sm, lane); + // Generic smem writes -> UMMA (async proxy). + cutlass::arch::fence_view_async_shared(); + full_lin->arrive(); // every lane arrives (count 32) + const uint32_t row_base = vblock * kBlockRows; + const int64_t row_delta = slot_row_delta(slot, row_base); + for (uint32_t base = 0; base < splits; base += Interval, ++drain_interval) { + for (uint32_t p = 0; p < kMathWarps / 2; ++p) + drain_handoff(hdata, h_full, h_empty, p, drain_interval, row_base, + row_counts, base + Interval >= splits, + [&](uint32_t row, uint32_t slot2, uint32_t record, + uint32_t batch_base, uint32_t source) { + store_ring_record(vals, indices, row, row_delta, slot2, + record, kv_base, batch_base, + source % kMathWgThreads); + }); + } + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); + } +} else { + // -------- math warps 0-7 -------- + cutlass::arch::warpgroup_reg_alloc(); + const uint32_t math_tid = threadIdx.x; + const uint32_t query_base = (warp_idx / 4) * kWgRows; // 4 warps per WG + RingPipeline<1> q_pipe; + RingPipeline kv_pipe; + RingPipeline tmem_pipe; + // Each warpgroup consumes every second TMEM slot position (WG1 starts at 1). + constexpr uint32_t kTmemStep = 2; + if (warp_idx >= 4) tmem_pipe.advance(1); + const uint32_t heads_col = warp_idx < 4 ? 16u : 0u; + const uint32_t lin_col = warp_idx < 4 ? 0u : 128u; + uint32_t produce_interval = 0; + for (uint32_t k = 0;; ++k) { + claim_full[k & 1].wait((k >> 1) & 1); + const ClaimEntry& e = entries[k & 1]; + const uint32_t vblock = e.vblock; + if (vblock >= kPause) break; + const uint32_t row_base = vblock * kBlockRows; + if (e.all_short) { + for (uint32_t r = 0; r < kBlockRows && row_base + r < rows; ++r) + write_all_keys(out + (uint64_t)(row_base + r) * topk, topk, + e.ends[r] - e.starts[r], math_tid, kMathThreads); + __syncwarp(); + claim_empty[k & 1].arrive(); + continue; + } + // ---- block setup: Q, weights, sample histogram ---- + { + CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase); + (void)q_stage; + full_q.wait(q_phase); + full_lin->wait(q_phase); + } + static_assert(sizeof(SampleShared) % 16 == 0); + for (uint32_t j = math_tid; j < sizeof(SampleShared) / 16; + j += kMathThreads) + reinterpret_cast(sample_smem.hist)[j] = make_uint4(0, 0, 0, 0); + float weights[kWgRows][kNumHeads]; + float row_s[kWgRows]; // pow2 row scale of the linear term + uint32_t seq_k_start[kWgRows], seq_k_end[kWgRows]; +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + seq_k_start[i] = e.starts[query_base + i]; + seq_k_end[i] = e.ends[query_base + i]; + // Weights in units of 1 / s_r (exact: s_r is a power of two). + row_s[i] = sm.lin_scale[query_base + i]; + const float inv_s = 1.0f / row_s[i]; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { + const float4 w4 = *reinterpret_cast( + sm.weights + (query_base + i) * kNumHeads + j); + weights[i][j] = w4.x * 0.5f * inv_s; + weights[i][j + 1] = w4.y * 0.5f * inv_s; + weights[i][j + 2] = w4.z * 0.5f * inv_s; + weights[i][j + 3] = w4.w * 0.5f * inv_s; + } + } + { + // Weights that are subnormal, or become subnormal or infinite once + // scaled, take the exact recovery path. + const float w = sm.weights[warp_idx * kNumHeads + lane]; // lane = head + // 2^-125: the smallest magnitude whose half is still a normal float. + constexpr float kMinScalable = 0x1p-125f; + bool bad_w = w != 0.0f && fabsf(w) < kMinScalable; + const float wr = w * 0.5f * (1.0f / sm.lin_scale[warp_idx]); + bad_w |= (wr != 0.0f && fabsf(wr) < kMinScalable) || isinf(wr); + const uint32_t bad = __ballot_sync(kFullMask, bad_w); + if (lane == 0 && bad) atomicOr(&unsafe_rows, 1u << warp_idx); + } + math_wg_sync(); + // ---- sample tiles ---- + // The first tile's per-row max is the histogram reference (see + // SampleShared): tile 0 keeps its scores until the reference is known. + float ref[kWgRows] = {}; +#pragma unroll 1 + for (uint32_t t = 0; t < e.samples; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); + full_kv[kv_stage].wait(kv_phase); + uint32_t first_code[kWgRows] = {}; +#pragma unroll + for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) { + const uint32_t kv_lane = + (math_tid % kMathWgThreads) + kv_half * kMathWgThreads; + const uint32_t group = t * kTileKeys + kv_lane; + const uint32_t kv_offset = group * sample_stride + sample_stride / 2; + const bool in_view = group < sample_groups; + // 0 outside the view (staged by the KV producer). + const float scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + kv_lane); + CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase); + full_tmem[tmem_stage].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + if (kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + empty_kv[kv_stage].arrive(); + } + // One query row (32 TMEM columns) at a time keeps the sample loop's + // live registers low enough to carry the scan weights through it. +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + float acc[kNumHeads]; + float lin[4]; + ptx::tmem_load_32dp32b<32>( + tmem_stage * kSlotCols + heads_col + i * kNumHeads, + reinterpret_cast(acc)); + ptx::tmem_load_32dp32b<4>(tmem_stage * kSlotCols + lin_col + i * 4, + reinterpret_cast(lin)); + cutlass::arch::fence_view_async_tmem_load(); + if (i == kWgRows - 1) { + ptx::tcgen05_before_thread_sync(); + empty_tmem[tmem_stage].arrive(); + } + const float bq = reduce_heads_with_lin(acc, weights[i], lin) * + (row_s[i] * scale_kv); + const bool valid = in_view && kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq); + if (t == 0) { + sample_smem.first[kv_half][i][math_tid] = + valid ? bq : -CUDART_INF_F; + if (valid) first_code[i] = max(first_code[i], ordered_code(bq)); + } else if (valid) + add_sample(sample_smem.hist[query_base + i], bq - ref[i]); + } + } + if (t == 0) { +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + const uint32_t code = __reduce_max_sync(kFullMask, first_code[i]); + if (lane == 0) atomicMax(&sample_smem.ref_code[query_base + i], code); + } + math_wg_sync(); +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + ref[i] = sample_ref(sample_smem.ref_code[query_base + i]); +#pragma unroll + for (uint32_t h = 0; h < 2; ++h) { + const float x = sample_smem.first[h][i][math_tid]; + if (isfinite(x)) + add_sample(sample_smem.hist[query_base + i], x - ref[i]); + } + } + } + } + math_wg_sync(); + finish_sample_histogram( + sample_smem.hist[warp_idx], sample_ref(sample_smem.ref_code[warp_idx]), + gate_s + warp_idx, e.ends[warp_idx] - e.starts[warp_idx], sample_target, + lane); + math_wg_sync(); + // The block index is re-read from the claim entry (volatile) instead of + // kept live across the sample loop. + const uint32_t scan_row_base = volatile_load(e.vblock) * kBlockRows; + + // ---- scan tiles ---- + const uint32_t kv_base = e.kv_base, splits = e.splits; + const int64_t cand_row_delta = slot_row_delta(e.slot, scan_row_base); + float gate_score[kWgRows]; +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + gate_score[i] = gate_s[query_base + i]; + if ((volatile_load(unsafe_rows) >> (query_base + i)) & 1u) { + // Unsafe weights: no key passes (+inf gate), and a count above the + // candidate cap sends the block to the exact rescan. + gate_score[i] = CUDART_INF_F; + if (math_tid % kMathWgThreads == 0) + atomicAdd(row_counts + scan_row_base + query_base + i, + static_cast(kCandidatesPerRow) + 1); + } + // Compare S * scale_kv against gate / row_s (row_s is a power of two -> + // exact). + gate_score[i] *= 1.0f / row_s[i]; + } + uint32_t lane_counts[kWgRows] = {}; + uint32_t flush_base = 0; +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < splits; ++kv_split_idx) { + if ((kv_split_idx & (Interval - 1)) == 0) { + h_empty[(produce_interval & 1) * kMathWarps + warp_idx].wait( + ((produce_interval >> 1) & 1) ^ 1); + __syncwarp(); + } + CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); + full_kv[kv_stage].wait(kv_phase); +#pragma unroll + for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) { + const uint32_t kv_lane = + (math_tid % kMathWgThreads) + kv_half * kMathWgThreads; + const uint32_t kv_offset = kv_base + kv_split_idx * kTileKeys + kv_lane; + const float scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + kv_lane); + CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase); + full_tmem[tmem_stage].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + if (kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + empty_kv[kv_stage].arrive(); + } + float accum_pair[2][kNumHeads]; + float lin_pair[2][4]; + float + bp4[kWgRows]; // row scores of the slot, in units of the row scale +#pragma unroll + for (uint32_t pair = 0; pair < kWgRows; pair += 2) { + const uint32_t slot_base = tmem_stage * kSlotCols; + ptx::tmem_load_32dp32b<2 * kNumHeads>( + slot_base + heads_col + pair * kNumHeads, + reinterpret_cast(accum_pair)); + ptx::tmem_load_32dp32b<8>(slot_base + lin_col + pair * 4, + reinterpret_cast(lin_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair + 2 >= kWgRows) { + ptx::tcgen05_before_thread_sync(); + empty_tmem[tmem_stage].arrive(); + } + // The linear term (sum_t term_t . k, in units of the row scale; terms + // 2, 3 in units of s * 2^-8) seeds the sums. +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) + bp4[pair + r] = reduce_heads_with_lin( + accum_pair[r], weights[pair + r], lin_pair[r]) * + scale_kv; + } + // One warp vote per slot: most keys fail every gate, so the common path + // is 4 compares + 1 vote + 1 uniform branch. + const bool any_pass = + (bp4[0] > gate_score[0]) | (bp4[1] > gate_score[1]) | + (bp4[2] > gate_score[2]) | (bp4[3] > gate_score[3]); + if (__any_sync(kFullMask, any_pass)) [[unlikely]] { +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + const float bp = bp4[i]; + // Gate first; range and finiteness only for passing keys. + if (bp > gate_score[i]) { + const float bq = bp * row_s[i]; + const uint32_t row = scan_row_base + query_base + i; + if (kv_offset >= seq_k_start[i] && kv_offset < seq_k_end[i] && + isfinite(bq)) { + const uint32_t count = lane_counts[i]; + if (count < kRingSlots) { + const uint32_t code = candidate_fp24_code(bq); + hdata + .records[produce_interval & 1][warp_idx][i][count][lane] = + (code << 8) | ((kv_split_idx - flush_base) * 2 + kv_half); + lane_counts[i] = count + 1; + } else { + const int slot = atomicAdd(row_counts + row, 1); + if (slot < static_cast(kCandidatesPerRow)) { + const uint64_t off = + candidate_offset(row, cand_row_delta, slot); + store_candidate(vals + off, indices + off, bq, kv_offset); + } + } + } + } + } + } + } + if (((kv_split_idx + 1) & (Interval - 1)) == 0 && + kv_split_idx + 1 != splits) { +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + hdata.num_records[produce_interval & 1][warp_idx][i][lane] = + lane_counts[i]; + lane_counts[i] = 0; + } + if (lane == 0) hdata.bases[produce_interval & 1][warp_idx] = flush_base; + __syncwarp(); + h_full[(produce_interval & 1) * kMathWarps + warp_idx].arrive(); + ++produce_interval; + flush_base = kv_split_idx + 1; + } + } + { + // Final interval: every math warp stores its own ring records (8 warps in + // parallel instead of 2 drain warps after the scan); the drain warps only + // cycle the barriers. + const uint32_t bank = produce_interval & 1, + rbase = scan_row_base + query_base; + static_assert(sizeof(hdata.records[0][0]) == + kWgRows * kRingSlots * 32 * sizeof(uint32_t), + "ring record layout"); + self_drain_last(&hdata.records[bank][warp_idx][0][0][0], lane_counts[0], + lane_counts[1], lane_counts[2], lane_counts[3], + row_counts, rbase, kv_base + (math_tid % kMathWgThreads), + flush_base, cand_row_delta, vals, indices); +#pragma unroll + for (uint32_t row = 0; row < kWgRows; ++row) lane_counts[row] = 0; + __syncwarp(); + h_full[bank * kMathWarps + warp_idx].arrive(); + ++produce_interval; + } + cutlass::arch::fence_view_async_shared(); + empty_q.arrive(); + // Drains of this block done. + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).sync(); + // ---- selection ---- + // Re-read as well (not kept live across the scan). + const uint32_t sel_vblock = volatile_load(e.vblock), + sel_row_base = sel_vblock * kBlockRows; + __threadfence_block(); + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + const ClaimEntry& sel_entry = entries[k & 1]; + if (math_tid == 0) { + uint32_t recover = 0, short_rows = 0; + row_offset[0] = 0; + for (uint32_t r = 0; r < kBlockRows; ++r) { + const bool is_short = sel_entry.ends[r] - sel_entry.starts[r] <= topk; + short_rows |= uint32_t(is_short) << r; + const int c = __ldcg(row_counts + sel_row_base + r); + recover |= !is_short && (c < int(topk) || c > int(kCandidatesPerRow)); + // Rows start 16-byte aligned for select_staged_row's vector loads. + row_offset[r + 1] = + row_offset[r] + (is_short ? 0 : (uint32_t(max(c, 0)) + 3u) & ~3u); + } + // 32-bit byte arithmetic (sizeof is a 64-bit size_t). + constexpr uint32_t kWordBytes = sizeof(uint32_t); + // Stage all 8 rows' candidates at once (1), in two rounds of 4 (2), or, + // when they exceed the scratch, rescan exactly. + const uint32_t half = max(row_offset[kWgRows], + row_offset[kBlockRows] - row_offset[kWgRows]); + const uint32_t staged = + recover ? 0u + : ((kBlockRows * selection::kBins + row_offset[kBlockRows]) * + kWordBytes <= + phase_bytes + ? 1u + : ((kWgRows * selection::kBins + half) * kWordBytes <= + phase_bytes + ? 2u + : 0u)); + if (staged == 0) recover = 1; + if (recover) { + defer_next[sel_vblock] = int(defer_head); + defer_head = sel_vblock; + atomicExch(&pause_req, 1u); + } + sel_recover = recover; + sel_staged = staged; + sel_short = short_rows; + unsafe_rows = 0; + } + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + const uint32_t recover = sel_recover, short_rows = sel_short; + const auto write_short = [&](uint32_t r, uint32_t t, uint32_t stride) { + if (r >= rows) return; + write_all_keys( + out + (uint64_t)r * topk, topk, + sel_entry.ends[r - sel_row_base] - sel_entry.starts[r - sel_row_base], + t, stride); + }; + if (!recover) { + const int64_t row_delta = slot_row_delta(sel_entry.slot, sel_row_base); + const uint32_t per_round = sel_staged == 1 ? kBlockRows : kWgRows; + for (uint32_t round = 0; round < kBlockRows / per_round; ++round) { + if (warp_idx < per_round) { + const uint32_t local = round * per_round + warp_idx, + first = round * per_round; + uint32_t* hist = reinterpret_cast(smem_buffer + kSelBase) + + warp_idx * selection::kBins; + uint32_t* codes = + reinterpret_cast(smem_buffer + kSelBase) + + per_round * selection::kBins + + (row_offset[local] - row_offset[first]); + const uint32_t r = sel_row_base + local; + const uint64_t base = candidate_offset(r, row_delta, 0); + if ((short_rows >> local) & 1u) + write_short(r, lane, 32); + else { + const uint32_t cnt = __ldcg(row_counts + r); + int* o = out + (uint64_t)r * topk; + if (!selection::select_staged_row(codes, cnt, vals + base, + indices + base, o, topk, hist)) + selection::select_row(codes, cnt, vals + base, indices + base, o, + topk, hist); + } + } + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + } + } + // Every row is final here: the last selection round ends with a barrier, + // and a recovered block skipped the rounds right after the barrier above. + // Short rows were written relative already. + if (!recover) + make_relative(out, topk, rows, sel_row_base, starts_g, math_tid, + kMathThreads, [&](uint32_t local) { + return sel_entry.ends[local] - sel_entry.starts[local] <= + topk; + }); + // The selection scratch aliases kv[] and sf_kv[]. Order its generic + // accesses before the next block's TMA writes, which have no proxy fence + // of their own when that block has no sample tiles. + cutlass::arch::fence_view_async_shared(); + __syncwarp(); + if (lane == 0) sel_done.arrive(); + claim_empty[k & 1].arrive(); + } +} + +// -------- exact rescan of the blocks queued by the selection -------- +// The rescan's first TMA loads overwrite the epoch's generic shared memory. +cutlass::arch::fence_view_async_shared(); +__syncthreads(); +if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, 512); +if (defer_head != kTerm) { + // The rescan expects the 224 / 56 register split of the math warps. + __syncthreads(); + for (uint32_t cur = defer_head; cur != kTerm;) { + const uint32_t vblock = cur, row_base = vblock * kBlockRows; + if (threadIdx.x < kBlockRows) { + const uint32_t r = row_base + threadIdx.x; + row_bounds[0][threadIdx.x] = r < rows ? min(starts_g[r], keys) : 0u; + row_bounds[1][threadIdx.x] = + r < rows ? max(min(ends_g[r], keys), row_bounds[0][threadIdx.x]) : 0u; + } + __syncthreads(); + const uint32_t* const starts = row_bounds[0] - row_base; + const uint32_t* const ends = row_bounds[1] - row_base; + // No block is in flight during the rescan: use the CTA's first slot. + const int64_t row_delta = + slot_row_delta(kSlotsPerCta * blockIdx.x, row_base); + const RescanBuffers block_buffers{ + rescan.retained + row_delta * kRetainedPerRow, + rescan.retained_indices + row_delta * kRetainedPerRow}; + rescan::rescan_block(dense::Scheduler(vblock, rows, keys, starts, ends), + block_buffers, out, topk, tq, tk, ts, tw); + __syncthreads(); + if (threadIdx.x == 0) + dense::invalidate_barriers( + *reinterpret_cast(smem_buffer)); + __syncthreads(); + // The rescan writes absolute indices for every row of the block. + make_relative(out, topk, rows, row_base, starts_g, threadIdx.x, kNumThreads, + [](uint32_t) { return false; }); + cur = uint32_t(__ldcg(defer_next + vblock)); + // Order the generic mbarrier.inval above before the TMA writes of the + // next rescan block or the next epoch. + cutlass::arch::fence_view_async_shared(); + __syncthreads(); + } + // The loop ran at least once and ended with a barrier. + if (threadIdx.x == 0) { + defer_head = kTerm; + pause_req = 0; + } +} // deferred blocks +more = epoch_kind != kTerm; +__syncthreads(); diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh new file mode 100644 index 000000000000..86063b08fbac --- /dev/null +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -0,0 +1,592 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Persistent fused FP8 MQA logits + per-row top-K (K = 512, 1024, 2048; SM100), +// DeepGEMM/FA4-style: one CTA per SM, barriers and TMEM initialised once, every +// warp role runs its own loop over blocks claimed (heaviest first) from a +// global counter, and the sample and scan tiles of a block flow through ONE +// continuous Q/KV/TMEM pipeline whose stage/phase state is carried across +// blocks. +// +// Roles (384 threads): +// warp 8 claimer + Q/weights TMA + drain of math WG1 rows +// warp 9 KV TMA (sample tiles via the strided map, then the scan tiles) +// warp 10 UMMA issuer +// warp 11 linear-term rows + drain of math WG0 rows +// warps 0-7 math: sample histogram -> gate -> scan/emit -> selection +// A block is rescanned exactly (rescan.cuh) when a row's candidate count +// leaves [K, kCandidatesPerRow], a row's weights are unsafe to scale, or the +// block's candidates exceed the selection scratch (~11K+ per row): the role +// loops pause, the CTA rescans its queued blocks, and the loops resume with +// freshly initialised barriers. +#pragma once +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "candidates.cuh" +#include "common.cuh" +#include "constants.h" +#include "deep_gemm_compat.cuh" +#include "rescan.cuh" +#include "select.cuh" + +namespace fused_mqa_topk::persist { +using namespace deep_gemm; +using rescan::RescanBuffers; +using rescan::RescanShared; +using TBar = cutlass::arch::ClusterTransactionBarrier; + +constexpr uint32_t kTerm = 0xffffffffu, kPause = 0xfffffffeu; +// Bytes of a block's Q, of one KV tile and its per-key scales, of a block's +// weights. +constexpr uint32_t kQBytes = kBlockRows * kNumHeads * kHeadDim, + kKVBytes = kTileKeys * kHeadDim, + kScaleBytes = kTileKeys * sizeof(float), + kWBytes = kBlockRows * kNumHeads * sizeof(float); +// KV (and per-key scale) pipeline stages. +constexpr uint32_t kKVStages = 3; +// Linear-term B rows per math warpgroup: 4 query rows x 4 FP8 terms. +constexpr uint32_t kLinRows = 16; +// Terms 2 and 3 of the linear term are stored in units of s * 2^-8. +constexpr float kLinLoScale = 0x1p-8f; +// TMEM: one UMMA per math warpgroup and KV half writes 4 rows x 32 heads plus +// the kLinRows linear-term columns (N = 144). Slots are 160 columns apart (144 +// rounded up to a multiple of 32); the 3 slots take 480 of the 512 allocated +// columns. +constexpr uint32_t kUmmaN = 4 * kNumHeads + kLinRows, kSlotCols = 160; +static_assert(kSlotCols >= kUmmaN && kSlotCols % 32 == 0); + +// Rebalanced scan: the linear-term rows of each math warpgroup bracket Q, so +// WG0's N=144 UMMA reads B = [lin WG0 | Q rows 0..127] (D = [lin | heads]) and +// WG1's reads B = [Q rows 128..255 | lin WG1] (D = [heads | lin]). +struct KernelSmem { + alignas(1024) uint8_t lin_lo[kLinRows * kHeadDim]; + cutlass::float_e4m3_t q[kQBytes]; + uint8_t lin_hi[kLinRows * kHeadDim]; + alignas(128) float weights[kBlockRows * kNumHeads]; + float lin_scale[kBlockRows]; + alignas(1024) cutlass::float_e4m3_t kv[kKVStages][kKVBytes]; + alignas(128) float sf_kv[kKVStages][kTileKeys]; +}; +static_assert(offsetof(KernelSmem, q) == kLinRows * kHeadDim && + offsetof(KernelSmem, lin_hi) == kLinRows * kHeadDim + kQBytes, + "lin rows must bracket Q"); +// Selection scratch starts after lin + Q + lin + weights + row scales, so the +// next block's Q and linear term land during this block's selection. +constexpr uint32_t kSelBase = offsetof(KernelSmem, kv); + +// Byte offset of byte `col` of row `row` in a 128B-swizzled K-major tile of +// 128-byte rows (the UMMA B layout of Q and the linear-term rows). +__device__ __forceinline__ uint32_t swizzle128_offset(uint32_t row, + uint32_t col) { + return row * 128u + (((col >> 4) ^ (row & 7u)) << 4) + (col & 15u); +} + +// q'_r = sum_h (w_rh / 2) q_rh, split into 4 FP8 e4m3 terms: terms 0,1 scale +// s_r, terms 2,3 s_r * 2^-8 (lin_scale[r]); term t of row r is B row (r % 4) * +// 4 + t of lin_lo (rows 0-3) / lin_hi (4-7), 128B-swizzled K-major like Q. One +// warp; lane l owns dims 4l..4l+3. +__device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, + uint32_t lane) { + const uint8_t* qbytes = reinterpret_cast(sm.q); +#pragma unroll 1 + for (uint32_t r = 0; r < kBlockRows; ++r) { + float acc[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll 8 + for (uint32_t h = 0; h < kNumHeads; ++h) { + const uint32_t off = swizzle128_offset(r * kNumHeads + h, lane * 4u); + const uint32_t packed = + ptx::ld_shared(reinterpret_cast(qbytes + off)); + const float wv = ptx::ld_shared(sm.weights + r * kNumHeads + h) * 0.5f; + const __half2_raw lo = __nv_cvt_fp8x2_to_halfraw2( + static_cast<__nv_fp8x2_storage_t>(packed & 0xffffu), __NV_E4M3); + const __half2_raw hi = __nv_cvt_fp8x2_to_halfraw2( + static_cast<__nv_fp8x2_storage_t>(packed >> 16), __NV_E4M3); + const float2 f0 = __half22float2(*reinterpret_cast(&lo)); + const float2 f1 = __half22float2(*reinterpret_cast(&hi)); + acc[0] = fmaf(wv, f0.x, acc[0]); + acc[1] = fmaf(wv, f0.y, acc[1]); + acc[2] = fmaf(wv, f1.x, acc[2]); + acc[3] = fmaf(wv, f1.y, acc[3]); + } + float m = fmaxf(fmaxf(fabsf(acc[0]), fabsf(acc[1])), + fmaxf(fabsf(acc[2]), fabsf(acc[3]))); +#pragma unroll + for (uint32_t o = 16; o > 0; o >>= 1) + m = fmaxf(m, __shfl_xor_sync(kFullMask, m, o)); + int e = -120; // pow2 scale s >= m / 448, clamped so 1/s stays finite + if (m > 0.f && isfinite(m)) { + frexpf(m * (1.0f / 448.0f), &e); + e = max(e, -120); + } + const float s = ldexpf(1.0f, e), inv_s = ldexpf(1.0f, -e); + float res[4]; +#pragma unroll + for (uint32_t dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s; + uint8_t* lin_base = r < kWgRows ? sm.lin_lo : sm.lin_hi; +#pragma unroll + for (uint32_t tt = 0; tt < 4; ++tt) { + const __nv_fp8x2_storage_t c01 = __nv_cvt_float2_to_fp8x2( + make_float2(res[0], res[1]), __NV_SATFINITE, __NV_E4M3); + const __nv_fp8x2_storage_t c23 = __nv_cvt_float2_to_fp8x2( + make_float2(res[2], res[3]), __NV_SATFINITE, __NV_E4M3); + const __half2_raw d01 = __nv_cvt_fp8x2_to_halfraw2(c01, __NV_E4M3); + const __half2_raw d23 = __nv_cvt_fp8x2_to_halfraw2(c23, __NV_E4M3); + const float2 g01 = + __half22float2(*reinterpret_cast(&d01)); + const float2 g23 = + __half22float2(*reinterpret_cast(&d23)); + res[0] -= g01.x; + res[1] -= g01.y; + res[2] -= g23.x; + res[3] -= g23.y; + // Terms 2, 3 in units of s * 2^-8: the residual after two terms is <= + // 0.5, so x256 <= 128 (no saturation), and the e4m3 subnormal floor drops + // from s * 2^-10 to s * 2^-18 per dim. + if (tt == 1) { + res[0] *= 1.0f / kLinLoScale; + res[1] *= 1.0f / kLinLoScale; + res[2] *= 1.0f / kLinLoScale; + res[3] *= 1.0f / kLinLoScale; + } + const uint32_t off = + swizzle128_offset((r % kWgRows) * 4u + tt, lane * 4u); + ptx::st_shared( + reinterpret_cast(lin_base + off), + static_cast(c01) | (static_cast(c23) << 16)); + } + if (lane == 0) sm.lin_scale[r] = s; + } +} + +// One row's weighted head reduction seeded with its linear-term columns +// (terms 2, 3 in units of s * 2^-8). +__device__ __forceinline__ float reduce_heads_with_lin(const float* s, + const float* w, + const float* lin) { + return reduce_heads(s, w, make_float2(lin[0], lin[1]), + __fmul2_rn(make_float2(lin[2], lin[3]), + make_float2(kLinLoScale, kLinLoScale))); +} + +// Gate sample state: each math warp histograms its row's sampled scores +// relative to the row's reference `ref` (its max over the first sample tile): +// half_bin<5>(x - ref), kSampleBins bins, 32 per octave of the distance, so +// the bins near the row's top stay fine whatever the scores' common offset +// (absolute bins are ~2 wide at |x| ~ 100, too coarse for top tails that +// narrow). +constexpr uint32_t kSampleShift = 5, kSampleBins = 65536 >> kSampleShift; +struct alignas(16) SampleShared { + uint32_t hist[kBlockRows][padded_bins(kSampleBins)]; + // Ordered code of each row's max over the first sample tile (0: none). + uint32_t ref_code[kBlockRows]; + // The first tile's scores, held until the reference is known (in registers + // they cost the hot loops spills). + float first[2][kWgRows][kMathThreads]; +}; + +__device__ __forceinline__ void add_sample(uint32_t* hist, float distance) { + atomicAdd(hist + padded_bin(half_bin(distance)), 1u); +} +__device__ __forceinline__ uint32_t sample_count(const uint32_t* hist, + uint32_t bin) { + return hist[padded_bin(bin)]; +} +// The row reference from its ordered code; 0 when the row had no valid sample +// in the first tile (then the bins are absolute). +__device__ __forceinline__ float sample_ref(uint32_t code) { + const float ref = ordered_float(code); + return isfinite(ref) ? ref : 0.0f; +} + +// The gate aims kGateSlack16 / 16 above the target count, so that few rows fall +// below the top-k (and into the exact rescan) when the sampled estimate runs +// low. On the B300 captures 17/16 cuts the rescanned blocks by ~85% at about +// the candidate counts of the coarser absolute bins. +constexpr uint32_t kGateSlack16 = 17; + +// Gate: the k-th largest sample (k = count * target / visible, count = the +// histogram total) of one row from its histogram; one warp per row. -inf when +// the row sees at most `target` keys or has no samples. +__device__ __noinline__ inline void finish_sample_histogram( + const uint32_t* hist, float ref, float* gate, uint32_t visible, + uint32_t target, uint32_t lane) { + if (visible <= target) { + if (lane == 0) *gate = -CUDART_INF_F; + __syncwarp(); + return; + } + // Lane l owns kChunksPerLane 32-bin chunks, descending with the lane: chunk + // kChunks - 1 - kChunksPerLane * l - j for j < kChunksPerLane. + constexpr uint32_t kChunks = kSampleBins / 32, kChunksPerLane = kChunks / 32; + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < kChunksPerLane; ++j) { + const uint32_t c = kChunks - 1 - kChunksPerLane * lane - j; +#pragma unroll 8 + for (uint32_t b = 0; b < 32; ++b) sum += hist[padded_bin(c * 32 + b)]; + } + const uint32_t incl = warp_inclusive_scan(sum, lane); + const uint32_t tot = __shfl_sync(kFullMask, incl, 31); + if (tot == 0) { + if (lane == 0) *gate = -CUDART_INF_F; + __syncwarp(); + return; + } + const uint32_t rank = max( + 1u, min(tot, (tot * target * kGateSlack16 / 16 + visible - 1) / visible)); + const uint32_t w = __ffs(__ballot_sync(kFullMask, incl >= rank)) - 1; + uint32_t above = __shfl_sync(kFullMask, incl - sum, w), chosen = 0; +#pragma unroll 1 + for (uint32_t cc = 0; cc < kChunksPerLane; ++cc) { + const uint32_t c = kChunks - 1 - kChunksPerLane * w - cc; + // Lane 0 = highest bin of the chunk. + const uint32_t v = + warp_inclusive_scan(sample_count(hist, c * 32 + 31 - lane), lane); + const uint32_t mask = __ballot_sync(kFullMask, above + v >= rank); + if (mask) { + chosen = c * 32 + 31 - (__ffs(mask) - 1); + break; + } + above += __shfl_sync(kFullMask, v, 31); + } + if (lane == 0) *gate = ref + half_bin_floor(chosen); + __syncwarp(); +} + +// Candidate records each lane of a math warp buffers per row and handoff +// interval before it falls back to direct global stores. +constexpr uint32_t kRingSlots = 8; +// The drain's prefix sum packs two 16-bit per-row counts into one word. +static_assert(32 * kRingSlots < 65536, + "per-warp record count must fit 16 bits"); + +// Handoff ring: [bank][math warp][row][slot][lane] records, per-lane record +// counts and the interval's first scan tile. +struct alignas(16) HandoffRing { + uint32_t records[2][kMathWarps][kWgRows][kRingSlots][32]; + uint32_t num_records[2][kMathWarps][kWgRows][32]; + uint32_t bases[2][kMathWarps]; +}; + +// samples: gate sample tiles of the block (0 when no row sees more keys than +// the gate target, so every gate is -inf). +struct ClaimHeader { + uint32_t vblock, all_short, kv_base, splits, slot, samples; +}; +struct ClaimEntry : ClaimHeader { + uint32_t starts[kBlockRows], ends[kBlockRows]; +}; + +// Claim k in a producer role (warps 9-11): waits for the entry, copies its +// header and releases the entry (every lane arrives). False once the claimer +// ends the epoch (kTerm / kPause). +__device__ __forceinline__ bool consume_claim(uint32_t k, TBar* claim_full, + TBar* claim_empty, + const ClaimEntry* entries, + ClaimHeader& claim) { + claim_full[k & 1].wait((k >> 1) & 1); + claim = entries[k & 1]; + __syncwarp(); + if (claim.vblock >= kPause) return false; + claim_empty[k & 1].arrive(); + return true; +} + +// Barrier words (static shared, never aliased with data). +struct BarIdx { + // TMEM slots, shared by the two math warpgroups (WG0 even, WG1 odd + // positions). + static constexpr uint32_t kTmemSlots = 3; + static constexpr uint32_t full_q = 0, empty_q = 1, full_kv = 2, + empty_kv = 2 + kKVStages, + full_tmem = 2 + 2 * kKVStages, + empty_tmem = full_tmem + kTmemSlots, + claim_full = empty_tmem + kTmemSlots, + claim_empty = claim_full + 2, + sel_done = claim_empty + 2, h_full = sel_done + 1, + h_empty = h_full + 2 * kMathWarps, + full_lin = h_empty + 2 * kMathWarps, + total = full_lin + 1; +}; + +// Expect `bytes` of TMA transactions and arrive `count` times on one barrier. +__device__ __forceinline__ void mbar_arrive_expect_tx(const void* b, + uint32_t bytes, + uint32_t count) { + const uint32_t a = static_cast(__cvta_generic_to_shared(b)); + asm volatile( + "mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;" ::"r"(a), + "r"(bytes) + : "memory"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;" ::"r"(a), + "r"(count) + : "memory"); +} + +// Candidate slabs: a block's 8 rows live in its CTA slot's 8 slab rows, so row +// r maps to slab row r + slot_row_delta(slot, row_base). +__device__ __forceinline__ int64_t slot_row_delta(uint32_t slot, + uint32_t row_base) { + return int64_t(slot) * kBlockRows - int64_t(row_base); +} +// Offset of entry `index` of row `row` in the candidate slabs. +__device__ __forceinline__ uint64_t candidate_offset(uint32_t row, + int64_t row_delta, + uint32_t index) { + return static_cast(int64_t(row) + row_delta) * kCandidatesPerRow + + index; +} +// Stores ring record `record` (fp24 code << 8 | tile offset * 2 + KV half, +// tile offset relative to the interval's first tile `interval_base`) of KV +// lane `kv_lane` as entry `index` of row `row`. +__device__ __forceinline__ void store_ring_record( + uint16_t* vals, int32_t* indices, uint32_t row, int64_t row_delta, + uint32_t index, uint32_t record, uint32_t kv_base, uint32_t interval_base, + uint32_t kv_lane) { + const uint32_t key = + kv_base + (interval_base * 2 + (record & 255u)) * 128 + kv_lane; + const uint64_t off = candidate_offset(row, row_delta, index); + store_candidate_record(vals + off, indices + off, record >> 8, key); +} + +// Output row `out` of a row that sees at most top-k keys: its `visible` keys +// in order (0 .. visible - 1, relative to the row start), then -1 padding. +// Threads t, t + stride, ... write the slots. +__device__ __forceinline__ void write_all_keys(int* out, uint32_t topk, + uint32_t visible, uint32_t t, + uint32_t stride) { + for (uint32_t j = t; j < topk; j += stride) + out[j] = j < visible ? int(j) : -1; +} + +// Rewrites the output rows of the block at row_base in place from absolute key +// indices to indices relative to the row's start (the top_k_per_row_prefill +// convention; -1 padding stays). Rows that start at 0 and rows for which +// skip(local row) holds are left as they are. Threads t, t + stride, ... each +// rewrite 4 slots. +template +__device__ __forceinline__ void make_relative(int* output, uint32_t topk, + uint32_t rows, uint32_t row_base, + const uint32_t* starts_g, + uint32_t t, uint32_t stride, + Skip skip) { + const uint32_t quads = topk / 4; + for (uint32_t i = t; i < kBlockRows * quads; i += stride) { + const uint32_t r = row_base + i / quads; + if (r >= rows || skip(i / quads)) continue; + const int s = int(starts_g[r]); + if (s == 0) continue; + int4* p = reinterpret_cast(output + (uint64_t)r * topk) + i % quads; + int4 v = *p; + v.x = v.x < 0 ? -1 : v.x - s; + v.y = v.y < 0 ? -1 : v.y - s; + v.z = v.z < 0 ? -1 : v.z - s; + v.w = v.w < 0 ? -1 : v.w - s; + *p = v; + } +} + +// Stores one warp's ring records of a handoff interval in its rows' candidate +// slabs. recs = [4 rows][kRingSlots records][32 lanes]; lane l holds n[row] +// records of each row. Each row's slots are reserved with one atomic on +// row_counts[row_base + row]; records past kCandidatesPerRow are dropped (the +// row's count then exceeds the cap, so its block is rescanned exactly). +// store(row, slot, record) writes one record. +// Written for a warpgroup's 4 rows (two 16-bit counts per word in the scan). +static_assert(kWgRows == 4, "store_ring_records packs 4 row counts"); +template +__device__ __forceinline__ void store_ring_records(const uint32_t* recs, + const uint32_t (&n)[4], + int* row_counts, + uint32_t row_base, + Store store) { + const uint32_t lane = threadIdx.x & 31; + uint32_t my_total = 0; +#pragma unroll + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t total = __reduce_add_sync(kFullMask, n[row]); + if (lane == row) my_total = total; + } + uint32_t my_start = 0; + if (lane < 4 && my_total) + my_start = atomicAdd(row_counts + row_base + lane, my_total); + // Two 16-bit per-row counts per word: one scan covers two rows. + uint32_t p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); +#pragma unroll + for (int dd = 1; dd < 32; dd *= 2) { + const uint32_t a = __shfl_up_sync(kFullMask, p01, dd), + b = __shfl_up_sync(kFullMask, p23, dd); + if (lane >= dd) { + p01 += a; + p23 += b; + } + } + const uint32_t prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; +#pragma unroll + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t start = __shfl_sync(kFullMask, my_start, row); + for (uint32_t j = 0; j < n[row]; ++j) { + const uint32_t slot = start + prefix[row] - n[row] + j; + const uint32_t record = recs[(row * kRingSlots + j) * 32 + lane]; + if (slot < kCandidatesPerRow) store(row_base + row, slot, record); + } + } +} + +// Self-drain, out of line: stores this warp's ring records of the last +// handoff interval. recs = records[bank][warp]. +__device__ __noinline__ inline void self_drain_last( + const uint32_t* __restrict__ recs, uint32_t n0, uint32_t n1, uint32_t n2, + uint32_t n3, int* __restrict__ row_counts, uint32_t rbase, uint32_t lane_kv, + uint32_t flush_base, int64_t row_delta, uint16_t* __restrict__ vals, + int32_t* __restrict__ indices) { + const uint32_t n[4] = {n0, n1, n2, n3}; + store_ring_records(recs, n, row_counts, rbase, + [&](uint32_t row, uint32_t slot, uint32_t record) { + store_ring_record(vals, indices, row, row_delta, slot, + record, lane_kv, flush_base, 0); + }); + __syncwarp(); +} + +// Drain warp: waits for producer warp `producer`'s handoff interval, stores its +// records (unless the producer self-drained it) and frees the ring bank. +template +__device__ __forceinline__ void drain_handoff( + HandoffRing& d, TBar* hfull, TBar* hempty, int producer, uint32_t interval, + uint32_t row_base, int* row_counts, bool skip, Store store) { + const int lane = threadIdx.x & 31, bank = interval & 1; + hfull[bank * kMathWarps + producer].wait((interval >> 1) & 1); + __syncwarp(); + if (skip) { // self-drained by the producer + hempty[bank * kMathWarps + producer].arrive(); + return; + } + const uint32_t base = d.bases[bank][producer]; + uint32_t n[4]; +#pragma unroll + for (uint32_t row = 0; row < 4; ++row) + n[row] = d.num_records[bank][producer][row][lane]; + store_ring_records(&d.records[bank][producer][0][0][0], n, row_counts, + row_base, + [&](uint32_t row, uint32_t slot, uint32_t record) { + store(row, slot, record, base, producer * 32 + lane); + }); + __syncwarp(); + hempty[bank * kMathWarps + producer].arrive(); +} + +// Interval: scan tiles per candidate handoff to the drain warps. +template +__global__ +__launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel( + uint32_t rows, uint32_t keys, const uint32_t* starts_g, + const uint32_t* ends_g, int* out, uint16_t* vals, int* indices, + int* row_counts, RescanBuffers rescan, uint32_t phase_bytes, + uint32_t sample_target, uint32_t sample_stride, uint32_t sample_groups, + uint32_t sample_tiles, const float* scales, + const __grid_constant__ cute::TmaDescriptor tks, + const __grid_constant__ cute::TmaDescriptor tq, + const __grid_constant__ cute::TmaDescriptor ts, + const __grid_constant__ cute::TmaDescriptor tk, + const __grid_constant__ cute::TmaDescriptor tw, uint32_t nblocks, + uint32_t* counter, int* defer_next) { + // Compile-time: a runtime top-k costs ~0.5% in the scan (codegen). + constexpr uint32_t topk = TopK; + static_assert(BarIdx::kTmemSlots * kSlotCols <= 512, + "TMEM slots exceed 512 columns"); + // A ring record's low byte holds (tile - interval start) * 2 + half, and the + // handoff points use `& (Interval - 1)`. + static_assert((Interval & (Interval - 1)) == 0 && Interval <= 128, + "Interval must be a power of two <= 128"); + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& sm = *reinterpret_cast(smem_buffer); + uint8_t* const region = smem_buffer + sizeof(KernelSmem); + auto& sample_smem = *reinterpret_cast(region); + auto& hdata = *reinterpret_cast(region); + + __shared__ alignas(8) uint64_t bar_words[BarIdx::total]; + __shared__ uint32_t tmem_ptr; + __shared__ ClaimEntry entries[2]; + __shared__ float gate_s[kBlockRows]; + __shared__ uint32_t sel_recover, sel_staged, sel_short, + row_offset[kBlockRows + 1]; + __shared__ uint32_t defer_head, unsafe_rows, pause_req, epoch_kind; + + TBar* const bars = reinterpret_cast(bar_words); + TBar& full_q = bars[BarIdx::full_q]; + TBar& empty_q = bars[BarIdx::empty_q]; + TBar* const full_kv = bars + BarIdx::full_kv; + // The 32 producer lanes stage the sample scales, so each arrives on full_kv. + constexpr uint32_t kFullKvArrivals = 32; + TBar* const empty_kv = bars + BarIdx::empty_kv; + TBar* const full_tmem = bars + BarIdx::full_tmem; + TBar* const empty_tmem = bars + BarIdx::empty_tmem; + TBar* const claim_full = bars + BarIdx::claim_full; + TBar* const claim_empty = bars + BarIdx::claim_empty; + TBar& sel_done = bars[BarIdx::sel_done]; + TBar* const h_full = bars + BarIdx::h_full; + TBar* const h_empty = bars + BarIdx::h_empty; + TBar* const full_lin = bars + BarIdx::full_lin; + + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane = ptx::get_lane_idx(); + constexpr uint32_t kClaimConsumers = 11; // warps 9, 10, 11, 0-7 + + if (warp_idx == 8 && lane == 0) { + cute::prefetch_tma_descriptor(&tq); + cute::prefetch_tma_descriptor(&tw); + cute::prefetch_tma_descriptor(&tk); + cute::prefetch_tma_descriptor(&ts); + cute::prefetch_tma_descriptor(&tks); + } + if (warp_idx == 9 && cute::elect_one_sync()) { + defer_head = kTerm; + pause_req = 0; + } + // Epochs: the persistent role loops run until the work runs out (kTerm) or a + // block needs the exact rescan (kPause); rare blocks are then rescanned + // exactly (rescan.cuh) and the loops resume with freshly initialised + // barriers. The first epoch is emitted outside any loop (nesting the role + // loops in an epoch loop costs ~1.5 percent in the scan: uniform address math + // is rematerialised every tile). Later epochs only run after a recovery + // pause. + __shared__ uint32_t row_bounds[2][kBlockRows]; + bool more = false; + { + constexpr uint32_t epoch = 0; +#include "persistent_epoch.inc" + } +#pragma unroll 1 + for (uint32_t epoch = 1; more; ++epoch) { +#include "persistent_epoch.inc" + } + // counter = {next block, finished CTAs}: the last CTA to finish (every claim + // is done by then) zeroes both, so the caller's workspace is ready for the + // next launch / graph replay. + if (threadIdx.x == 0) { + __threadfence(); + if (atomicAdd(counter + 1, 1u) == gridDim.x - 1) { + // Acquire: every other CTA's last claim on counter happens before the + // reset below. + __threadfence(); + atomicExch(counter, 0u); + atomicExch(counter + 1, 0u); + } + } +} + +} // namespace fused_mqa_topk::persist diff --git a/csrc/fused_mqa_topk/rescan.cuh b/csrc/fused_mqa_topk/rescan.cuh new file mode 100644 index 000000000000..80d9dc57370f --- /dev/null +++ b/csrc/fused_mqa_topk/rescan.cuh @@ -0,0 +1,350 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Exact recovery for the blocks the persistent kernel cannot select from its +// gated candidates (see persistent_kernel.cuh): the block is rescanned densely, +// and every 1024-key segment of fp32 scores is merged into a running per-row +// top set (kept <= kRetainedPerRow, compacted to K). +#pragma once +#include +#include +#include +#include +#include + +#include "common.cuh" +#include "constants.h" +#include "dense_scan.cuh" + +namespace fused_mqa_topk::rescan { +using dense::Scheduler; + +// Keys whose scores are merged into the retained sets at a time. +constexpr uint32_t kSegmentKeys = 1024; +// Per-row retained top set of the exact rescan (scores and key indices). +struct RescanBuffers { + float* retained; + int* retained_indices; +}; +// Per-row segment histograms: half_bin<6> of a score (1024 bins); the radix +// levels of compact use 256 bins. +constexpr uint32_t kSegmentBins = 1024, kRadixBins = 256; +struct alignas(16) RescanShared { + float ring[2][kBlockRows][kSegmentKeys]; + uint32_t hist[kBlockRows][padded_bins(kSegmentBins)], gates[kBlockRows], + kept[kBlockRows]; + uint32_t work[padded_bins(kRadixBins)], prefix, rank, emit_counts[kMathWarps]; + __device__ void init(uint32_t tid) { + for (uint32_t j = tid; j < sizeof(hist) / sizeof(uint32_t); + j += kNumThreads) + reinterpret_cast(hist)[j] = 0; + if (tid < kBlockRows) { + gates[tid] = __float_as_uint(-CUDART_INF_F); + kept[tid] = 0; + } + } +}; +// What compact_retained's emit pass keeps. +enum class Emit { kAboveGate, kAboveCode, kTies }; + +__device__ inline void compact_retained(RescanShared& shared, + RescanBuffers buffers, uint32_t topk, + uint32_t local, uint32_t row, + uint32_t tid, + bool compact_to_k = false) { + constexpr uint32_t kThreads = kMathThreads, + kItems = kRetainedPerRow / kThreads; + uint32_t lane = tid & 31, count = shared.kept[local]; + uint64_t base = (uint64_t)row * kRetainedPerRow; + float values[kItems]; + int indices[kItems]; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + uint32_t j = tid + i * kThreads; + values[i] = j < count ? buffers.retained[base + j] : -CUDART_INF_F; + indices[i] = j < count ? buffers.retained_indices[base + j] : -1; + } + // Emits the values >= gate, the codes above threshold, or the codes equal + // to it (ties, up to K) to the retained set from base_offset on. + auto emit = [&](float gate, uint32_t threshold, Emit mode, + uint32_t base_offset) { + const auto passes = [&](float v) { + return isfinite(v) && + (mode == Emit::kAboveGate + ? v >= gate + : (mode == Emit::kAboveCode ? ordered_code(v) > threshold + : ordered_code(v) == threshold)); + }; + uint32_t total = 0; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) total += passes(values[i]); + uint32_t warp_total = __reduce_add_sync(kFullMask, total); + if (lane == 0) shared.emit_counts[tid / 32] = warp_total; + math_wg_sync(); + uint32_t start = base_offset; + for (uint32_t w = 0; w < tid / 32; ++w) start += shared.emit_counts[w]; + uint32_t all = 0; + for (uint32_t w = 0; w < kMathWarps; ++w) all += shared.emit_counts[w]; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + const float v = values[i]; + const bool pass = passes(v); + uint32_t mask = __ballot_sync(kFullMask, pass), + pos = start + __popc(mask & ((1u << lane) - 1)); + if (pass && (mode != Emit::kTies || pos < topk)) { + buffers.retained[base + pos] = v; + buffers.retained_indices[base + pos] = indices[i]; + } + start += __popc(mask); + } + math_wg_sync(); + return all; + }; + if (!compact_to_k) { + if (tid == 0) { + shared.rank = 0; + } + math_wg_sync(); + float gate = __uint_as_float(shared.gates[local]); + uint32_t kept = 0; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) + kept += isfinite(values[i]) && values[i] >= gate; + kept = __reduce_add_sync(kFullMask, kept); + if (lane == 0) atomicAdd(&shared.rank, kept); + math_wg_sync(); + kept = shared.rank; + if (kept >= topk && kept <= kRetainedPerRow - kSegmentKeys) { + emit(gate, 0, Emit::kAboveGate, 0); + math_wg_sync(); + if (tid == 0) shared.kept[local] = kept; + math_wg_sync(); + return; + } + } + if (tid == 0) { + shared.prefix = 0; + shared.rank = topk; + } + math_wg_sync(); + for (int shift = 24; shift >= 0; shift -= 8) { + for (uint32_t j = tid; j < padded_bins(kRadixBins); j += kThreads) + shared.work[j] = 0; + math_wg_sync(); + uint32_t prefix = shared.prefix; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + uint32_t code = ordered_code(values[i]); + bool valid = isfinite(values[i]) && + (shift == 24 || (code >> (shift + 8)) == prefix); + uint32_t mask = __ballot_sync(kFullMask, valid); + if (valid) { + uint32_t bin = (code >> shift) & 255u, + peers = __match_any_sync(mask, bin); + if (lane == uint32_t(__ffs(peers) - 1)) + atomicAdd(shared.work + padded_bin(bin), __popc(peers)); + } + } + math_wg_sync(); + if (tid < 32) { + uint32_t sum = 0; + for (uint32_t j = 0; j < 8; ++j) { + uint32_t bin = 255 - lane * 8 - j; + sum += shared.work[padded_bin(bin)]; + } + const uint32_t inclusive = warp_inclusive_scan(sum, lane); + uint32_t before = inclusive - sum, rank = shared.rank, + crossing = + __ballot_sync(kFullMask, before < rank && inclusive >= rank), + chosen_lane = __ffs(crossing) - 1, chosen = 0, above = before; + if (lane == chosen_lane) { + for (uint32_t j = 0; j < 8; ++j) { + uint32_t bin = 255 - lane * 8 - j; + if (above + shared.work[padded_bin(bin)] >= rank) { + chosen = bin; + break; + } + above += shared.work[padded_bin(bin)]; + } + } + chosen = __shfl_sync(kFullMask, chosen, chosen_lane); + above = __shfl_sync(kFullMask, above, chosen_lane); + __syncwarp(); // every lane read shared.rank above + if (lane == 0) { + shared.prefix = (prefix << 8) | chosen; + shared.rank = rank - above; + } + } + math_wg_sync(); + } + uint32_t threshold = shared.prefix; + uint32_t greater = emit(0, threshold, Emit::kAboveCode, 0); + emit(0, threshold, Emit::kTies, greater); + if (tid == 0) { + shared.kept[local] = topk; + float gate = + fmaxf(__uint_as_float(shared.gates[local]), ordered_float(threshold)); + atomicExch(shared.gates + local, __float_as_uint(gate)); + } + math_wg_sync(); +} + +__device__ inline void merge_segment(RescanShared& shared, + RescanBuffers buffers, uint32_t topk, + uint32_t local, uint32_t row, + uint32_t segment, uint32_t first, + uint32_t n, uint32_t lane) { + constexpr uint32_t kItems = kSegmentKeys / 32; + const uint32_t bank = segment & 1; + float values[kItems], gate = __uint_as_float(shared.gates[local]); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + uint32_t j = lane + i * 32; + float v = j < n ? shared.ring[bank][local][j] : -CUDART_INF_F; + if (v <= gate) v = -CUDART_INF_F; + values[i] = v; + bool valid = isfinite(v); + count += valid; + uint32_t mask = __ballot_sync(kFullMask, valid); + if (valid) { + uint32_t bin = half_bin<6>(v), peers = __match_any_sync(mask, bin); + if (lane == uint32_t(__ffs(peers) - 1)) + atomicAdd(&shared.hist[local][padded_bin(bin)], __popc(peers)); + } + } + count = __reduce_add_sync(kFullMask, count); + __syncwarp(); + if (count) { + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < 32; ++j) + sum += shared.hist[local][padded_bin((31 - lane) * 32 + j)]; + uint32_t inclusive = warp_inclusive_scan(sum, lane); + uint32_t before = inclusive - sum, + mask = + __ballot_sync(kFullMask, before < topk && inclusive >= topk); + if (mask) { + uint32_t chosen = __ffs(mask) - 1, group = 31 - chosen, + rank = topk - __shfl_sync(kFullMask, before, chosen), + bin = group * 32 + 31 - lane; + sum = shared.hist[local][padded_bin(bin)]; + inclusive = warp_inclusive_scan(sum, lane); + before = inclusive - sum; + chosen = + __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - + 1; + bin = __shfl_sync(kFullMask, bin, chosen); + gate = fmaxf(gate, half_bin_floor<6>(bin)); + if (lane == 0) atomicExch(shared.gates + local, __float_as_uint(gate)); + } + uint32_t written = shared.kept[local]; + uint64_t base = (uint64_t)row * kRetainedPerRow; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + float v = values[i]; + bool pass = isfinite(v) && v >= gate; + uint32_t mask = __ballot_sync(kFullMask, pass), + pos = written + __popc(mask & ((1u << lane) - 1)); + if (pass) { + buffers.retained[base + pos] = v; + buffers.retained_indices[base + pos] = first + lane + i * 32; + } + written += __popc(mask); + } + __syncwarp(); + if (lane == 0) shared.kept[local] = written; + __syncwarp(); + } +} + +// Streams each row's scores (above its running gate) into the segment ring and +// merges every full segment into the row's retained top set; the final set is +// written to the output. +struct RescanPolicy { + static constexpr float kWeightScale = 0.5f; + static constexpr uint32_t kSplitsPerSegment = kSegmentKeys / kTileKeys; + RescanShared& shared; + RescanBuffers buffers; + int* out; // [rows, topk] + uint32_t topk, query_base, row_base, segment = 0; + uint32_t seq_k_start[kWgRows], seq_k_end[kWgRows]; + float gate[kWgRows]; + CUTLASS_DEVICE void begin_block(const Scheduler& sched, uint32_t q_block, + uint32_t base) { + query_base = base; + row_base = sched.get_logits_row(q_block, 0); +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + const uint32_t row = sched.get_logits_row(q_block, base + i); + gate[i] = -CUDART_INF_F; + seq_k_start[i] = min(sched.cu_seq_len_k_start[row], sched.num_kv_tokens); + seq_k_end[i] = min(sched.cu_seq_len_k_end[row], sched.num_kv_tokens); + } + } + CUTLASS_DEVICE void begin_split(uint32_t split) { + if (split % kSplitsPerSegment == 0) { + __syncwarp(); + for (uint32_t i = 0; i < kWgRows; ++i) + gate[i] = __uint_as_float(atomicAdd(shared.gates + query_base + i, 0u)); + } + } + CUTLASS_DEVICE void score(uint32_t i, uint32_t kv_offset, uint32_t kv_lane, + uint32_t split, float bq) { + const bool valid = kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq) && bq > gate[i]; + shared.ring[segment & 1][query_base + i] + [(split % kSplitsPerSegment) * kTileKeys + kv_lane] = + valid ? bq : -CUDART_INF_F; + } + CUTLASS_DEVICE void end_split(uint32_t split, uint32_t num_splits, + uint32_t kv_base) { + if ((split + 1) % kSplitsPerSegment != 0 && split + 1 != num_splits) return; + __threadfence_block(); + math_wg_sync(); + const uint32_t tid = threadIdx.x; + for (uint32_t local = 0; local < kBlockRows; ++local) + if (shared.kept[local] > kRetainedPerRow - kSegmentKeys) + compact_retained(shared, buffers, topk, local, row_base + local, tid); + math_wg_sync(); + const uint32_t local = tid / 32; + merge_segment(shared, buffers, topk, local, row_base + local, segment, + kv_base + (split / kSplitsPerSegment) * kSegmentKeys, + (split % kSplitsPerSegment + 1) * kTileKeys, tid % 32); + math_wg_sync(); + if (split + 1 == num_splits) + for (uint32_t l = 0; l < kBlockRows; ++l) + if (shared.kept[l] > topk) + compact_retained(shared, buffers, topk, l, row_base + l, tid, true); + math_wg_sync(); + ++segment; + } + CUTLASS_DEVICE void end_block(const Scheduler& sched, uint32_t q_block) { + const uint32_t local = threadIdx.x / 32, lane = threadIdx.x % 32; + const uint32_t row = sched.get_logits_row(q_block, local), + count = shared.kept[local]; + if (row < sched.num_q_tokens) + for (uint32_t j = lane; j < topk; j += 32) + out[(uint64_t)row * topk + j] = + j < count + ? buffers.retained_indices[(uint64_t)row * kRetainedPerRow + j] + : -1; + } +}; + +// Rescans one block exactly and writes its output rows of `out`; `buffers` +// rows are addressed by absolute row (callers offset them). Out of line: the +// rare rescan must not add register pressure to the caller's hot loops. +__device__ __noinline__ inline void rescan_block( + const Scheduler& block, RescanBuffers buffers, int* out, uint32_t topk, + const cute::TmaDescriptor& tq, const cute::TmaDescriptor& tk, + const cute::TmaDescriptor& ts, const cute::TmaDescriptor& tw) { + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& shared = *reinterpret_cast(smem_buffer + + sizeof(dense::SharedStorage)); + shared.init(threadIdx.x); + RescanPolicy policy{shared, buffers, out, topk}; + dense::score_block(block, policy, tq, tk, ts, tw); +} + +} // namespace fused_mqa_topk::rescan diff --git a/csrc/fused_mqa_topk/select.cuh b/csrc/fused_mqa_topk/select.cuh new file mode 100644 index 000000000000..759ec7fce46a --- /dev/null +++ b/csrc/fused_mqa_topk/select.cuh @@ -0,0 +1,452 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Per-row top-K selection of the persistent kernel's candidates, one warp per +// row (K = topk, the row's top-k: 512, 1024 or 2048, a runtime value). +// +// select_staged_row: bracket the K-th largest fp24 code with sampled codes, +// place the codes above the bracket directly, and resolve the bracket with +// histogram levels on register-resident codes. +// +// select_row: the fallback of select_staged_row when its sampled bracket +// misses. Candidate codes (higher code = higher score, staged once in shared +// memory) span a narrow range (gate .. row max), so each level histograms the +// range-relative code into at most kBins bins: collisions are rare, plain +// shared atomics suffice and each lane keeps kUnroll independent elements in +// flight. Levels repeat inside the boundary bin until it is one code wide; +// codes above it are emitted, then ties until K. +#pragma once +#include + +#include "candidates.cuh" +#include "common.cuh" +#include "constants.h" + +namespace fused_mqa_topk::selection { +constexpr uint32_t kBinsLog2 = 10, kBins = 1u << kBinsLog2; +// Elements each lane keeps in flight. +constexpr uint32_t kUnroll = 8; + +// Bin (descending) where the cumulative count first reaches `rank`; +// `above` = elements in higher bins. Lane l owns Bins / 32 bins, descending +// from bin Bins - 1 - l * Bins / 32. +template +__device__ __forceinline__ uint32_t find_bin_n(const uint32_t* hist, + uint32_t rank, uint32_t& above) { + const uint32_t lane = threadIdx.x & 31; + constexpr uint32_t kPerLane = Bins / 32; + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < kPerLane; ++j) + sum += hist[Bins - 1 - lane * kPerLane - j]; + const uint32_t inclusive = warp_inclusive_scan(sum, lane); + const uint32_t before = inclusive - sum; + const uint32_t winner = + __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - 1; + uint32_t chosen = 0, acc = before; + if (lane == winner) { + for (uint32_t j = 0; j < kPerLane; ++j) { + const uint32_t bin = Bins - 1 - lane * kPerLane - j, h = hist[bin]; + if (acc + h >= rank) { + chosen = bin; + break; + } + acc += h; + } + } + above = __shfl_sync(kFullMask, acc, winner); + return __shfl_sync(kFullMask, chosen, winner); +} + +// Requires count >= topk (callers recover rows with fewer candidates). +__device__ inline void select_row(uint32_t* codes, uint32_t count, + const uint16_t* values, const int32_t* packed, + int32_t* out, uint32_t topk, uint32_t* hist) { + const uint32_t lane = threadIdx.x & 31; + uint32_t lo = 0xffffffffu, hi = 0; + uint32_t j = lane; + for (; j + 32 * (kUnroll - 1) < count; j += 32 * kUnroll) { + uint16_t v[kUnroll]; + int32_t p[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + v[u] = values[j + 32 * u]; + p[u] = packed[j + 32 * u]; + } +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + const uint32_t c = candidate_load_score_code(v[u], p[u]); + codes[j + 32 * u] = c; + lo = min(lo, c); + hi = max(hi, c); + } + } + for (; j < count; j += 32) { + const uint32_t c = candidate_load_score_code(values[j], packed[j]); + codes[j] = c; + lo = min(lo, c); + hi = max(hi, c); + } + lo = __reduce_min_sync(kFullMask, lo); + hi = __reduce_max_sync(kFullMask, hi); + __syncwarp(); + + // Current bucket: codes in [lo, lo + width). rank = elements still needed + // from it. + uint32_t width = hi - lo + 1, rank = topk; +#pragma unroll 1 + while (width > 1) { + const uint32_t spread = (width - 1) >> kBinsLog2; + const uint32_t shift = spread ? 32 - __clz(spread) : 0; + for (uint32_t b = lane; b < kBins; b += 32) hist[b] = 0; + __syncwarp(); + const uint32_t top = lo + width - 1; + uint32_t i = lane; + for (; i + 32 * (kUnroll - 1) < count; i += 32 * kUnroll) { + uint32_t c[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) c[u] = codes[i + 32 * u]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) + if (c[u] >= lo && c[u] <= top) + atomicAdd(hist + ((c[u] - lo) >> shift), 1u); + } + for (; i < count; i += 32) { + const uint32_t c = codes[i]; + if (c >= lo && c <= top) atomicAdd(hist + ((c - lo) >> shift), 1u); + } + __syncwarp(); + uint32_t above; + const uint32_t chosen = find_bin_n(hist, rank, above); + rank -= above; + lo += chosen << shift; + width = min(1u << shift, top - lo + 1); + __syncwarp(); + } + + // lo is now the K-th largest code: emit codes above it, then ties until K. + const uint32_t threshold = lo, lt = (1u << lane) - 1u; + uint32_t pos = 0; +#pragma unroll 1 + for (uint32_t tie = 0; tie < 2; ++tie) { + for (uint32_t j0 = 0; j0 < count && pos < topk; j0 += 32 * kUnroll) { + bool pass[kUnroll]; + int32_t p[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + const uint32_t e = j0 + 32 * u + lane; + const uint32_t c = e < count ? codes[e] : 0; + pass[u] = e < count && (tie ? c == threshold : c > threshold); + p[u] = pass[u] ? packed[e] : 0; + } +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + const uint32_t mask = __ballot_sync(kFullMask, pass[u]), + at = pos + __popc(mask & lt); + if (pass[u] && at < topk) out[at] = candidate_decode_index(p[u]); + pos += __popc(mask); + } + } + } +} + +// Register samples per lane for the sampled bracket of select_staged_row. +constexpr uint32_t kSamplesPerLane = 16; +// Largest bracket (elements between the two sampled bounds) that +// select_staged_row resolves in registers; wider brackets fall back to +// select_row. +constexpr uint32_t kMidCap = 1024; +// Bins of the histogram levels that resolve the exact threshold inside a +// bracket (select_staged_row). +constexpr uint32_t kMidBins = 512; + +// Bracket element ids are stored as uint16 (midbuf). +static_assert(kCandidatesPerRow <= 65536, "candidate ids must fit 16 bits"); +// Per-warp scratch of select_staged_row is kBins words: radix_bracket's 768, +// or the midbuf (kMidCap uint16) plus the kMidBins-bin histogram. +static_assert(768 <= kBins && kMidCap / 2 + kMidBins <= kBins, + "selection scratch exceeds kBins words"); + +// r-th largest (1-based) of the kNumSamples * 32 register samples (kNumSamples +// per lane) for two ranks at once +// (r_hi <= r_lo), via a 256-bin hi8 level and two 256-bin lo8 levels. h8: 768 +// words of scratch. +template +__device__ __forceinline__ void radix_bracket( + const uint32_t (&smp)[kNumSamples], uint32_t r_hi, uint32_t r_lo, + uint32_t* h8, uint32_t& h_hi, uint32_t& h_lo) { + const uint32_t lane = threadIdx.x & 31; +#pragma unroll + for (uint32_t b = lane; b < 192; b += 32) + reinterpret_cast(h8)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (uint32_t u = 0; u < kNumSamples; ++u) atomicAdd(h8 + (smp[u] >> 8), 1u); + __syncwarp(); + uint32_t ab_hi, ab_lo; + const uint32_t b_hi = find_bin_n<256>(h8, r_hi, ab_hi); + const uint32_t b_lo = find_bin_n<256>(h8, r_lo, ab_lo); + uint32_t* const lh = h8 + 256; + uint32_t* const ll = h8 + 512; +#pragma unroll + for (uint32_t u = 0; u < kNumSamples; ++u) { + const uint32_t b = smp[u] >> 8; + if (b == b_hi) atomicAdd(lh + (smp[u] & 255u), 1u); + if (b == b_lo) atomicAdd(ll + (smp[u] & 255u), 1u); + } + __syncwarp(); + uint32_t d0, d1; + h_hi = b_hi << 8 | find_bin_n<256>(lh, r_hi - ab_hi, d0); + h_lo = b_lo << 8 | find_bin_n<256>(ll, r_lo - ab_lo, d1); + __syncwarp(); // h8 is reused as the output stage +} + +// Brackets the topk-th largest of the `count` staged codes between the hi16 +// codes of sample ranks expect -+ margin, where the kPerLane * 32 samples sit +// at evenly spaced quantiles and expect is the topk-th largest's expected rank +// among them. +template +__device__ __forceinline__ void sampled_bracket(const uint32_t* codes, + uint32_t count, uint32_t topk, + uint32_t margin, uint32_t* hist, + uint32_t& h_hi, + uint32_t& h_lo) { + constexpr uint32_t kSamples = kPerLane * 32; + const uint32_t lane = threadIdx.x & 31; + uint32_t smp[kPerLane]; +#pragma unroll + for (uint32_t u = 0; u < kPerLane; ++u) + smp[u] = codes[((u * 32 + lane) * 2 + 1) * count / (2 * kSamples)] >> 8; + const uint32_t expect = (topk * kSamples + count / 2) / count; + const uint32_t r_hi = expect > margin ? expect - margin : 1, + r_lo = min(expect + margin, kSamples); + radix_bracket(smp, r_hi, r_lo, hist, h_hi, h_lo); +} + +// find_bin_n with each lane's kPer contiguous bins loaded as 16-byte vectors +// and searched in registers (needs a 16-byte aligned histogram and kPer % 4 == +// 0; find_bin_n reads the bins one at a time and has neither requirement). +template +__device__ __forceinline__ uint32_t find_bin_v(const uint32_t* hist, + uint32_t rank, uint32_t& above) { + const uint32_t lane = threadIdx.x & 31; + constexpr uint32_t kPer = Bins / 32, kVec = kPer / 4; + static_assert(kPer % 4 == 0); + // lane l owns bins [Bins - (l + 1) * kPer, Bins - l * kPer), highest first. + const uint4* hv = + reinterpret_cast(hist + Bins - (lane + 1) * kPer); + uint32_t h[kPer]; +#pragma unroll + for (uint32_t v = 0; v < kVec; ++v) { + const uint4 x = hv[v]; + h[4 * v] = x.x; + h[4 * v + 1] = x.y; + h[4 * v + 2] = x.z; + h[4 * v + 3] = x.w; + } + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < kPer; ++j) sum += h[j]; + const uint32_t inclusive = warp_inclusive_scan(sum, lane); + const uint32_t before = inclusive - sum; + const uint32_t winner = + __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - 1; + // Descending walk inside the lane: bin index Bins - lane * kPer - 1 - j <-> + // h[kPer - 1 - j]. + uint32_t acc = before, chosen = 0, found = 0, acc_at = before; +#pragma unroll + for (uint32_t j = 0; j < kPer; ++j) { + const uint32_t x = h[kPer - 1 - j]; + const bool hit = !found && acc + x >= rank; + if (hit) { + chosen = Bins - lane * kPer - 1 - j; + acc_at = acc; + found = 1; + } + acc += x; + } + above = __shfl_sync(kFullMask, acc_at, winner); + return __shfl_sync(kFullMask, chosen, winner); +} + +// One row's global-memory candidates, staged to shared memory as +// fp24 codes (codes must be 16-byte aligned, i.e. row offsets rounded up to 4). +// Register budget: ptxas compiles the whole kernel to 168 regs, so every pass +// keeps <= ~32 regs of arrays live (more spills into the scan loop). +// Passes: stage (16-byte global loads of values / packed indices, two chunks of +// 8 in flight per lane, to fp24 codes); hi16 bracket (radix levels over 512 +// shared-memory samples); count above / inside the bracket (16-byte shared +// loads); collect (codes above the bracket to out, bracket element ids to +// midbuf); histogram levels on register codes; emit. +// Returns false before writing out; the caller's select_row then restages and +// runs the levels. +__device__ __noinline__ inline bool select_staged_row( + uint32_t* __restrict__ codes, uint32_t count, + const uint16_t* __restrict__ values, const int32_t* __restrict__ packed, + int32_t* __restrict__ out, uint32_t topk, uint32_t* hist) { + const uint32_t lane = threadIdx.x & 31; + if (count <= topk) return false; + const uint32_t chunks = (count + 7) / 8; + constexpr uint32_t kChunksInFlight = 2; + const uint4* v4p = reinterpret_cast(values); + const uint4* p4p = reinterpret_cast(packed); +#pragma unroll 1 + for (uint32_t c0i = lane; c0i < chunks; c0i += 32 * kChunksInFlight) { + uint4 vv[kChunksInFlight], pp0[kChunksInFlight], pp1[kChunksInFlight]; +#pragma unroll + for (uint32_t u = 0; u < kChunksInFlight; ++u) { + const uint32_t c = min(c0i + 32 * u, chunks - 1); + vv[u] = __ldcg(v4p + c); + pp0[u] = __ldcg(p4p + 2 * c); + pp1[u] = __ldcg(p4p + 2 * c + 1); + } +#pragma unroll + for (uint32_t u = 0; u < kChunksInFlight; ++u) { + const uint32_t c = c0i + 32 * u; + if (c < chunks) { + const uint32_t vw[4] = {vv[u].x, vv[u].y, vv[u].z, vv[u].w}; + const uint32_t pw[8] = {pp0[u].x, pp0[u].y, pp0[u].z, pp0[u].w, + pp1[u].x, pp1[u].y, pp1[u].z, pp1[u].w}; + uint32_t cw[8]; +#pragma unroll + for (uint32_t k = 0; k < 8; ++k) + cw[k] = candidate_load_score_code( + uint16_t(vw[k / 2] >> (16 * (k & 1))), int32_t(pw[k])); + if (8 * c + 8 <= count) { + reinterpret_cast(codes)[2 * c] = + make_uint4(cw[0], cw[1], cw[2], cw[3]); + reinterpret_cast(codes)[2 * c + 1] = + make_uint4(cw[4], cw[5], cw[6], cw[7]); + } else { +#pragma unroll + for (uint32_t k = 0; k < 8; ++k) + if (8 * c + k < count) codes[8 * c + k] = cw[k]; + } + } + } + } + __syncwarp(); + // The bracket is the code range between the sampled ranks expect -+ margin + // (expect = the K-th largest's expected rank among the samples). Rows up to + // kFewCandidates use 512 samples, +-40 ranks. Long-context rows (4-6K + // candidates) would get a bracket of ~16% of the row that way, over kMidCap + // for 10-20% of the rows at 512K-1M: 1024 samples, +-48 (~3 sd), give ~9%. + constexpr uint32_t kFewCandidates = 3584, kMargin = 40, kWideMargin = 48; + uint32_t h_hi, h_lo; + if (count <= kFewCandidates) + sampled_bracket(codes, count, topk, kMargin, hist, h_hi, + h_lo); + else + sampled_bracket<2 * kSamplesPerLane>(codes, count, topk, kWideMargin, hist, + h_hi, h_lo); + const uint4* c4p = reinterpret_cast(codes); + const uint32_t quads = (count + 3) / 4; + const uint32_t lt = (1u << lane) - 1u; + uint16_t* midbuf = reinterpret_cast(hist); + uint32_t* h512 = hist + kMidCap / 2; + // Collect with warp-ballot placement: passing lanes of one step write + // consecutive slots, so the global stores of the above elements are + // coalesced. Writes are bounded; the + // bracket check follows (out may be partially written on failure: the + // fallback select_row rewrites all K slots). + uint32_t above = 0, mid = 0; + constexpr uint32_t kQuadsInFlight = 4; +#pragma unroll 1 + for (uint32_t q0 = 0; q0 < quads; q0 += 32 * kQuadsInFlight) { + uint4 pp[kQuadsInFlight], cc[kQuadsInFlight]; +#pragma unroll + for (uint32_t u = 0; u < kQuadsInFlight; ++u) { + const uint32_t q = min(q0 + 32 * u + lane, quads - 1); + pp[u] = __ldcg(p4p + q); + cc[u] = c4p[q]; + } +#pragma unroll + for (uint32_t u = 0; u < kQuadsInFlight; ++u) { + const uint32_t q = q0 + 32 * u + lane; + const uint32_t w[4] = {cc[u].x, cc[u].y, cc[u].z, cc[u].w}; + const uint32_t p[4] = {pp[u].x, pp[u].y, pp[u].z, pp[u].w}; +#pragma unroll + for (uint32_t k = 0; k < 4; ++k) { + const uint32_t e = 4 * q + k, h = w[k] >> 8; + const bool valid = e < count; + const bool a = valid && h > h_hi, m = valid && !a && h >= h_lo; + const uint32_t ma = __ballot_sync(kFullMask, a), + mm = __ballot_sync(kFullMask, m); + const uint32_t at = above + __popc(ma & lt), mt = mid + __popc(mm & lt); + if (a && at < topk) out[at] = candidate_decode_index(int32_t(p[k])); + if (m && mt < kMidCap) midbuf[mt] = uint16_t(e); + above += __popc(ma); + mid += __popc(mm); + } + } + } + __syncwarp(); + if (above >= topk || above + mid < topk || mid > kMidCap) return false; + // Bracket codes per lane (register-resident). + constexpr uint32_t kMidPerLane = kMidCap / 32; + uint32_t mc[kMidPerLane]; +#pragma unroll + for (uint32_t u = 0; u < kMidPerLane; ++u) { + const uint32_t i = lane + 32 * u; + mc[u] = 0; + if (i < mid) mc[u] = codes[midbuf[i]]; + } + uint32_t rank = topk - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; +#pragma unroll 1 + while (top > base) { + const uint32_t width = top - base + 1, spread = (width - 1) / kMidBins; + const uint32_t shift = spread ? 32 - __clz(spread) : 0; +#pragma unroll + for (uint32_t b = lane; b < kMidBins / 4; b += 32) + reinterpret_cast(h512)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (uint32_t u = 0; u < kMidPerLane; ++u) { + if (u * 32 >= mid) break; + if (lane + 32 * u < mid && mc[u] >= base && mc[u] <= top) + atomicAdd(h512 + ((mc[u] - base) >> shift), 1u); + } + __syncwarp(); + uint32_t above_bin; + const uint32_t chosen = find_bin_v(h512, rank, above_bin); + rank -= above_bin; + base += chosen << shift; + top = min(base + (1u << shift) - 1, top); + __syncwarp(); + } + const uint32_t th = base; + { + // Ballot placement (coalesced stores): greater codes fill [above, K - + // rank), ties [K - rank, K). + uint32_t pg = above, pt = topk - rank; +#pragma unroll + for (uint32_t h = 0; h < 2; ++h) { + if (h * 16 * 32 >= mid) break; + int32_t pk[16]; +#pragma unroll + for (uint32_t j = 0; j < 16; ++j) { + const uint32_t u = h * 16 + j; + pk[j] = (lane + 32 * u < mid && mc[u] >= th) + ? __ldcg(packed + midbuf[lane + 32 * u]) + : 0; + } +#pragma unroll + for (uint32_t j = 0; j < 16; ++j) { + const uint32_t u = h * 16 + j; + if (u * 32 >= mid) break; + const bool v = lane + 32 * u < mid; + const bool g = v && mc[u] > th, tq = v && mc[u] == th; + const uint32_t mg = __ballot_sync(kFullMask, g), + mt = __ballot_sync(kFullMask, tq); + const uint32_t ag = pg + __popc(mg & lt), at = pt + __popc(mt & lt); + if (g) out[ag] = candidate_decode_index(pk[j]); + if (tq && at < topk) out[at] = candidate_decode_index(pk[j]); + pg += __popc(mg); + pt += __popc(mt); + } + } + } + return true; +} + +} // namespace fused_mqa_topk::selection diff --git a/csrc/ops.h b/csrc/ops.h index 92271860099e..802de7f550e7 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -884,3 +884,12 @@ void ngram_compute_n_gram_ids( torch::stable::Tensor& exclusive_req_len_sums, torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); + +void fused_mqa_topk_prefill( + torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor scales, torch::stable::Tensor weights, + torch::stable::Tensor starts, torch::stable::Tensor ends, + torch::stable::Tensor out, torch::stable::Tensor counter, + torch::stable::Tensor values, torch::stable::Tensor indices, + torch::stable::Tensor retained, torch::stable::Tensor retained_indices, + torch::stable::Tensor counts, torch::stable::Tensor defer); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index c0b8a4e5e85b..9969f6309b8f 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -8,6 +8,17 @@ // Note: We register under namespace "_C" so ops are accessible as // torch.ops._C. for compatibility with existing code. STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { +#ifdef VLLM_ENABLE_FUSED_MQA_TOPK + // Single-launch fused FP8 MQA-logits + top-k (512/1024/2048) prefill + // selection (SM100 family). Workspace tensors are caller-owned; counter must + // be zero. + ops.def( + "fused_mqa_topk_prefill(Tensor q, Tensor k, Tensor scales, " + "Tensor weights, Tensor starts, Tensor ends, Tensor! out, " + "Tensor! counter, Tensor! values, Tensor! indices, Tensor! retained, " + "Tensor! retained_indices, Tensor! counts, Tensor! defer) -> ()"); +#endif + // Compute per-token-group FP8 quantized tensor and scaling factor. // The dummy arguments are here so we can correctly fuse with RMSNorm. ops.def( @@ -792,6 +803,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { +#ifdef VLLM_ENABLE_FUSED_MQA_TOPK + ops.impl("fused_mqa_topk_prefill", TORCH_BOX(&fused_mqa_topk_prefill)); +#endif + // LongCat n-gram embedding index kernel. ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py new file mode 100644 index 000000000000..7042cedb6849 --- /dev/null +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Single-launch fused MQA-logits + top-k prefill vs the dense indexer path.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + fused_mqa_topk_available, + fused_mqa_topk_min_keys, + fused_mqa_topk_prefill, + get_fused_mqa_topk_workspaces, +) +from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerPrefillChunkMetadata, + DeepseekV32IndexerPrefillMetadata, +) + +requires_fused_mqa_topk = pytest.mark.skipif( + not torch.cuda.is_available() or not fused_mqa_topk_available(), + reason="Requires the SM100 fused MQA+top-k extension", +) + +TOPKS = (512, 1024, 2048) +TOPK = 2048 +LOGITS_BUDGET = 512 * 1024**2 + + +def _dense(q, k, scales, weights, starts, ends, topk=TOPK): + """Upstream dense path: budgeted DeepGEMM logits + top_k_per_row_prefill.""" + k_scale = scales + if mqa_logits_use_mx_scales(): + # Power-of-two scales as four equal UE8M0 bytes, like the K cache. + exp = torch.log2(scales).round().to(torch.int64) + 127 + k_scale = (exp * 0x01010101).to(torch.int32) + rows = q.shape[0] + out = torch.full((rows, topk), -1, device=q.device, dtype=torch.int32) + lens = ends - starts + max_len = max(1, int(lens.max())) + step = max(1, LOGITS_BUDGET // (4 * max_len)) + for i in range(0, rows, step): + j = min(rows, i + step) + q_scale, w = mqa_logits_mx_inputs(q[i:j], None, weights[i:j]) + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q[i:j], q_scale), (k, k_scale), w, starts[i:j], ends[i:j], max_len + ), + lens[i:j], + ) + # Compressed logits: row r holds keys [starts[r], ends[r]) from column 0. + ops.top_k_per_row_prefill( + logits, + torch.zeros_like(lens[i:j]), + lens[i:j], + out[i:j], + j - i, + logits.stride(0), + logits.stride(1), + topk, + ) + return out + + +def _logits(q, k, scales, weights, starts, ends): + """fp32 reference scores: sum_h w * relu(q . k) * scale, -inf off-range.""" + kf = k.float() + out = torch.zeros(q.shape[0], k.shape[0], device=q.device) + for h in range(q.shape[1]): + out += weights[:, h, None] * torch.relu(q[:, h].float() @ kf.T) + out *= scales[None] + n = torch.arange(k.shape[0], device=q.device) + off = (n[None] < starts[:, None]) | (n[None] >= ends[:, None]) + return out.masked_fill(off, -torch.inf) + + +def _check(rel, q, k, scales, weights, starts, ends, scores=True): + """Indices relative to the row start: count, bounds, uniqueness, scores. + + ``scores`` compares the selected scores with the exact fp32 top-k; the + SM100 dense path rounds its logits to bf16, so it is checked without. + """ + topk = rel.shape[1] + visible = (ends - starts).clamp(min=0) + valid = rel >= 0 + assert (valid.sum(1) == visible.clamp(max=topk)).all(), "valid count" + assert (rel[~valid] == -1).all(), "padding must be -1" + idx = rel.long() + starts[:, None].long() + assert ((~valid) | ((idx >= starts[:, None]) & (idx < ends[:, None]))).all() + srt = idx.masked_fill(~valid, -1).sort().values + assert ((srt[:, 1:] < 0) | (srt.diff(dim=1) > 0)).all(), "duplicates" + if not scores: + return + # Check scores in row blocks so the reference logits stay small. + step = max(1, LOGITS_BUDGET // (4 * k.shape[0])) + for i in range(0, q.shape[0], step): + j = min(q.shape[0], i + step) + logits = _logits(q[i:j], k, scales, weights[i:j], starts[i:j], ends[i:j]) + ref = logits.topk(topk, dim=1).values + chosen = logits.gather(1, idx[i:j].clamp(min=0)).masked_fill( + ~valid[i:j], -torch.inf + ) + torch.testing.assert_close( + chosen.sort(descending=True).values, ref, rtol=2e-4, atol=1e-4 + ) + + +def _inputs(rows, keys, seed=913): + torch.manual_seed(seed) + q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn) + k = torch.randn(keys, 128, device="cuda").to(torch.float8_e4m3fn) + scales = torch.ones(keys, device="cuda") + weights = torch.randn(rows, 32, device="cuda") / 32 + return q, k, scales, weights + + +@requires_fused_mqa_topk +@pytest.mark.parametrize("topk", TOPKS) +@pytest.mark.parametrize( + "keys,rows,case", + [ + (4096, 64, "causal"), + (16384, 256, "causal"), + (16384, 16, "early"), + (32768, 16, "bounds"), + (65536, 16, "ties"), + (131072, 16, "hot"), + (262145, 9, "padding"), + (1048576, 16, "causal"), + ], +) +def test_single_request_matches_dense(keys, rows, case, topk): + q, k, scales, weights = _inputs(rows, keys) + starts = torch.zeros(rows, device="cuda", dtype=torch.int32) + ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32) + if case == "early": # rows around each top-k's visible-key boundary + ends = torch.tensor( + [ + 0, + 1, + 511, + 512, + 513, + 1023, + 1024, + 1025, + 2047, + 2048, + 2049, + 4096, + 8192, + 12000, + 16000, + keys, + ], + device="cuda", + dtype=torch.int32, + ) + elif case == "bounds": + starts.fill_(333) + ends[::2] = keys // 2 + elif case == "ties": + weights.zero_() + elif case == "hot": + weights = weights.abs() + scales[:12288] = 4 + out = torch.empty(rows, topk, device="cuda", dtype=torch.int32) + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + _check(out, q, k, scales, weights, starts, ends) + # Same convention and valid counts as the dense path (key sets may differ + # only among equal scores at the top-k boundary). + dense = _dense(q, k, scales, weights, starts, ends, topk) + _check(dense, q, k, scales, weights, starts, ends, scores=False) + assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all() + + +@requires_fused_mqa_topk +@pytest.mark.parametrize("topk", TOPKS) +@pytest.mark.parametrize( + "lens", [[3000, 5000, 1000], [40000, 1500, 70000, 9], [600000, 200000]] +) +def test_multi_request_chunk_matches_dense(lens, topk): + """Packed requests: rows of request r see [cu[r], cu[r] + causal_end).""" + q_lens = [min(n, 1024) for n in lens] + rows, keys = sum(q_lens), sum(lens) + q, k, scales, weights = _inputs(rows, keys, seed=917) + starts, ends, base = [], [], 0 + for n, m in zip(lens, q_lens): + starts += [base] * m + ends += list(range(base + n - m + 1, base + n + 1)) + base += n + starts = torch.tensor(starts, device="cuda", dtype=torch.int32) + ends = torch.tensor(ends, device="cuda", dtype=torch.int32) + out = torch.empty(rows, topk, device="cuda", dtype=torch.int32) + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + _check(out, q, k, scales, weights, starts, ends) + dense = _dense(q, k, scales, weights, starts, ends, topk) + _check(dense, q, k, scales, weights, starts, ends, scores=False) + assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all() + # Both paths agree on (nearly) every selected key. The dense path's bf16 + # logits blur the k-th boundary, the more so for small k at long contexts + # (~98.6% overlap for top-512 over 600K keys; the fused_mqa_topk keys are + # exact). + overlap = sum( + len(set(a[a >= 0].tolist()) & set(b[b >= 0].tolist())) + for a, b in zip(out.cpu(), dense.cpu()) + ) + assert overlap >= 0.98 * int((dense >= 0).sum()) + + +@requires_fused_mqa_topk +@pytest.mark.parametrize("rows", [32768, 32770]) +def test_large_chunk_single_call(rows): + """Prefill chunks above 16384 rows (e.g. --max-num-batched-tokens 32768) + run in one launch.""" + keys = 65536 + q, k, scales, weights = _inputs(rows, keys, seed=921) + starts = torch.zeros(rows, device="cuda", dtype=torch.int32) + ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32) + out = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32) + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + _check(out, q, k, scales, weights, starts, ends) + + +@requires_fused_mqa_topk +def test_rescan_scores_only_its_block(): + """The exact rescan of a block touches only that block's 8 rows. + + Regression test: the rescan's scheduler used a grid stride of 4096 blocks, + so with more than 4096 blocks (rows > 32768) rescanning block v also scored + block v + 4096, with row bounds read outside the rescanned block and output + written without the relative-index conversion. Block 0 is forced into the + rescan (a subnormal weight takes the exact path); block 4096 belongs to a + second request (nonzero starts, so its indices are made relative) and must + stay exact. + """ + rows, keys_per_request = 33000, 8192 # 4125 blocks of 8 rows + split = rows // 2 + q, k, scales, weights = _inputs(rows, 2 * keys_per_request, seed=927) + weights[0, 0] = 1e-40 # subnormal fp32: block 0 takes the exact rescan + starts = torch.zeros(rows, device="cuda", dtype=torch.int32) + starts[split:] = keys_per_request + ends = starts + keys_per_request + out = torch.full((rows, TOPK), -7, device="cuda", dtype=torch.int32) + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + _check(out, q, k, scales, weights, starts, ends, scores=False) + # Scores of the rescanned block, the block 4096 blocks after it, and the + # last block. + check_rows = torch.cat( + [ + torch.arange(0, 8), + torch.arange(4096 * 8, 4097 * 8), + torch.arange(rows - 8, rows), + ] + ).cuda() + _check( + out[check_rows], + q[check_rows], + k, + scales, + weights[check_rows], + starts[check_rows], + ends[check_rows], + ) + + +@requires_fused_mqa_topk +def test_gate_with_common_score_offset(): + """Regression test: scores sharing a large common offset with a narrow top + tail (top ranks within ~1 of each other at ~+68) must still get a tight gate. + With absolute fp16 sample bins (~2 wide there) the gate fell a whole bin too + low, rows overflowed the candidate slab and most blocks took the exact + rescan (4.4x slower at 256K keys).""" + rows, keys, topk = 2048, 65536, TOPK + q, k, scales, weights = _inputs(rows, keys, seed=5) + # A shared component in dimension 0 adds 8 * 8 to every q.k. + q = q.float() + k = k.float() + q[:, :, 0] = 8.0 + k[:, 0] = 8.0 + q, k = q.to(torch.float8_e4m3fn), k.to(torch.float8_e4m3fn) + weights = torch.full_like(weights, 1.0 / 32) + starts = torch.zeros(rows, device="cuda", dtype=torch.int32) + ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32) + out = torch.empty(rows, topk, device="cuda", dtype=torch.int32) + workspaces = get_fused_mqa_topk_workspaces(q.device, rows) + torch.ops._C.fused_mqa_topk_prefill( + q, k, scales, weights, starts, ends, out, *workspaces + ) + _check(out, q, k, scales, weights, starts, ends) + # Per-row candidate counts the kernel left in its workspace: about the gate + # target (3072 for top-2048 at <= 128K keys), not a multiple of it. + counts = workspaces[-2][:rows] + assert (counts >= topk).all() and (counts <= 16384).all(), "rows rescanned" + assert counts.float().median() < 2 * 3072 + + +@requires_fused_mqa_topk +def test_cuda_graph_replay(): + rows, keys = 64, 65536 + q, k, scales, weights = _inputs(rows, keys) + starts = torch.zeros(rows, device="cuda", dtype=torch.int32) + ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32) + out = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32) + fn = lambda: fused_mqa_topk_prefill( # noqa: E731 + q, k, scales, weights, starts, ends, out + ) + fn() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + out.fill_(-7) + for _ in range(3): + graph.replay() + torch.accelerator.synchronize() + _check(out, q, k, scales, weights, starts, ends) + + +@requires_fused_mqa_topk +def test_op_rejects_bad_inputs(): + """The launcher does not re-check; the op enforces its contract.""" + q, k, scales, weights = _inputs(16, 4096) + starts = torch.zeros(16, device="cuda", dtype=torch.int32) + ends = torch.full((16,), 4096, device="cuda", dtype=torch.int32) + out768 = torch.empty(16, 768, device="cuda", dtype=torch.int32) + with pytest.raises(RuntimeError, match="out must be"): + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out768) + wide = torch.empty(16, 2 * TOPK, device="cuda", dtype=torch.int32) + with pytest.raises(RuntimeError, match="must be contiguous"): + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, wide[:, :TOPK]) + + +KPOOL = 4 + + +def _kpool_inputs(pool_lens=(20000, 70000), q_rows=96, tail=2, seed=919): + """Two packed GLM-5.3-Flash-style requests: K holds 4-token pools; the + last ``q_rows`` tokens of each request score the pools they can see.""" + q, k, _, weights = _inputs(q_rows * len(pool_lens), sum(pool_lens), seed) + starts, ends, positions, base = [], [], [], 0 + for n in pool_lens: + tokens = KPOOL * n + tail + pos = torch.arange(tokens - q_rows, tokens) + positions.append(pos) + starts += [base] * q_rows + ends += (base + (pos + 1) // KPOOL).tolist() + base += n + i32 = dict(device="cuda", dtype=torch.int32) + # Unit power-of-two scales as four UE8M0 bytes, like the SM100 K cache. + scales = torch.full((k.shape[0], 4), 127, device="cuda", dtype=torch.uint8) + return ( + q, + k, + scales, + weights, + torch.tensor(starts, **i32), + torch.tensor(ends, **i32), + torch.cat(positions).cuda(), + ) + + +def _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk, budget_mb=512): + """Run the GLM-5.3-Flash kpool indexer prefill on one packed chunk.""" + import vllm.models.glm5next.nvidia.sparse_indexer as kpool_indexer + + q, k, scales, weights, starts, ends, positions = inputs + rows, keys = q.shape[0], k.shape[0] + chunk = DeepseekV32IndexerPrefillChunkMetadata( + block_table=torch.zeros(2, 1, device="cuda", dtype=torch.int32), + cu_seqlen_ks=starts, + cu_seqlen_ke=ends, + cu_seq_lens=torch.zeros(3, device="cuda", dtype=torch.int32), + token_to_seq=torch.zeros(keys, device="cuda", dtype=torch.int32), + total_seq_lens=keys, + max_seqlen_k=int((ends - starts).max()), + token_start=0, + token_end=rows, + num_reqs=2, + skip_kv_gather=True, # K is already in the gather workspace + local_cu_seq_lens=torch.zeros(3, device="cuda", dtype=torch.int32), + local_total_seq_lens=keys, + max_local_total_seq_lens=keys, + ) + metadata = DeepseekV32IndexerMetadata( + seq_lens=torch.zeros(2, device="cuda", dtype=torch.int32), + max_seq_len=int(positions.max()) + 1, + slot_mapping=torch.zeros(rows, device="cuda", dtype=torch.int64), + num_decodes=0, + num_decode_tokens=0, + num_prefills=2, + num_prefill_tokens=rows, + prefill=DeepseekV32IndexerPrefillMetadata( + chunks=[chunk], max_prefill_seq_len=int(positions.max()) + 1 + ), + ) + workspace = SimpleNamespace(get_simultaneous=lambda *specs: (k, scales)) + monkeypatch.setattr( + kpool_indexer, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata={"indexer": metadata}), + ) + monkeypatch.setattr(kpool_indexer, "current_workspace_manager", lambda: workspace) + monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", fused_mqa_topk) + monkeypatch.setattr(envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", budget_mb) + calls: list[torch.Size] = [] + fused_fn = kpool_indexer.fused_mqa_topk_prefill + + def spy(*a, **kw): + calls.append(a[-1].shape) # the output buffer: [rows, top-k pools] + fused_fn(*a, **kw) + + monkeypatch.setattr(kpool_indexer, "fused_mqa_topk_prefill", spy) + out = torch.empty(rows, TOPK + KPOOL - 1, device="cuda", dtype=torch.int32) + kpool_indexer.sparse_attn_indexer_kpool( + hidden_states=torch.empty(rows, 1, device="cuda"), + k_cache_prefix="indexer", + kv_cache=torch.empty(0, device="cuda"), + q_quant=q, + q_scale=None, + k=None, + weights=weights, + quant_block_size=128, + scale_fmt="ue8m0", + topk_tokens=TOPK, + head_dim=128, + max_pool_len=keys, + total_seq_lens=keys, + topk_indices_buffer=out, + skip_k_cache_insert=True, + index_kpool=KPOOL, + positions=positions, + ) + return out, calls + + +def _row_sets(out): + return [set(r[r >= 0].tolist()) for r in out.cpu()] + + +@requires_fused_mqa_topk +def test_kpool_indexer_fused_matches_dense(monkeypatch): + """GLM-5.3-Flash kpool indexer: the fused_mqa_topk kernel ranks the top-512 pools of + a packed multi-request chunk like the dense logits + top-k path.""" + inputs = _kpool_inputs() + # 192 rows take the kernel only when dense needs more than one logits + # call, so shrink the budget below one call. + fused_mqa_topk_out, calls = _run_kpool_indexer( + monkeypatch, inputs, fused_mqa_topk=True, budget_mb=16 + ) + assert calls == [torch.Size([inputs[0].shape[0], TOPK // KPOOL])] + dense, calls = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=False) + assert calls == [] + assert ((fused_mqa_topk_out >= 0).sum(1) == (dense >= 0).sum(1)).all() + # Same pools up to the dense path's bf16 boundary blur (see above); every + # row also keeps its request's incomplete tail tokens. + overlap = sum( + len(a & b) for a, b in zip(_row_sets(fused_mqa_topk_out), _row_sets(dense)) + ) + assert overlap >= 0.98 * int((dense >= 0).sum()) + + +@requires_fused_mqa_topk +def test_kpool_indexer_short_chunk_stays_dense(monkeypatch): + """Chunks whose rows see fewer than fused_mqa_topk_min_keys pools keep the + dense path (the fused_mqa_topk kernel is not faster there).""" + inputs = _kpool_inputs(pool_lens=(9000, 20000)) + assert int((inputs[5] - inputs[4]).max()) < fused_mqa_topk_min_keys(TOPK // KPOOL) + _, calls = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=True) + assert calls == [] + + +@requires_fused_mqa_topk +def test_kpool_indexer_dense_fallback_splits_rows(monkeypatch): + """Chunks planned for fused_mqa_topk skip the logits budget; when such a + chunk takes the dense path, its rows are split to stay within the budget + and select the + same pools as one unbudgeted dense pass.""" + inputs = _kpool_inputs() + whole, _ = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=False) + split, calls = _run_kpool_indexer( + monkeypatch, inputs, fused_mqa_topk=False, budget_mb=1 + ) + assert calls == [] + # top_k_per_row_prefill breaks bf16 ties at the k-th score arbitrarily (two + # identical dense runs differ too), so compare the selected scores. + q, k, scales, weights, starts, ends, _ = inputs + lens = ends - starts + q_scale, w = mqa_logits_mx_inputs(q, None, weights) + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q, q_scale), + (k, scales.view(torch.int32).squeeze(-1)), + w, + starts, + ends, + int(lens.max()), + ), + lens, + ) + for r, (a, b) in enumerate(zip(_row_sets(split), _row_sets(whole))): + # Request-relative token ids -> pools; the appended tail is not scored. + pools = [ + sorted({t // KPOOL for t in s if t // KPOOL < int(lens[r])}) for s in (a, b) + ] + scores = [logits[r, p].sort().values for p in pools] + assert torch.equal(scores[0], scores[1]), r diff --git a/vllm/envs.py b/vllm/envs.py index 857085ee1b26..e6092b4556a9 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -59,6 +59,7 @@ VLLM_CPU_INT4_W4A8: bool = True VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache") VLLM_XLA_CHECK_RECOMPILATION: bool = False + VLLM_USE_FUSED_MQA_TOPK: bool = True VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512 VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN: int = 8192 VLLM_XPU_PP_MICROBATCH: bool = True @@ -1115,6 +1116,13 @@ def _resolve_rust_cli_path() -> str | None: ), # Enable SPMD mode for TPU backend. "VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))), + # Use the single-launch fused FP8 MQA-logits + top-k kernel for the sparse + # MLA indexer prefill chunks it supports on SM100-family GPUs (see + # use_fused_mqa_topk / fused_mqa_topk_chunk); other chunks take the dense + # path. Default: on (set to 0 for the dense path). + "VLLM_USE_FUSED_MQA_TOPK": lambda: bool( + int(os.getenv("VLLM_USE_FUSED_MQA_TOPK", "1")) + ), # Maximum size (in MB) for logits tensor in sparse MLA indexer prefill chunks. # Bounds the [M, N] float32 logits tensor to prevent CUDA OOM. # Default: 512 MB diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py new file mode 100644 index 000000000000..928e222498a7 --- /dev/null +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Single-launch fused FP8 MQA-logits + top-k prefill selection (SM100 family)""" + +import functools + +import torch + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +# Current restrictions of the fused_mqa_topk kernel: +FUSED_MQA_TOPK_HEADS = 32 +FUSED_MQA_TOPK_HEAD_DIM = 128 +FUSED_MQA_TOPK_TOPKS = (512, 1024, 2048) +# Most gathered keys (all requests of a chunk) one call accepts: candidates +# store the key index in 20 bits (kMaxKeys in csrc/fused_mqa_topk/constants.h). +# Chunks with more keys take the dense path; the planner never packs past it. +FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 +# Rows per fused_mqa_topk call. The kernel has no row limit, but prefill chunks above +# 16384 rows (--max-num-batched-tokens > 16384) are not validated end to end yet, +# so the chunk planner splits fused_mqa_topk chunks at this many rows. +FUSED_MQA_TOPK_MAX_ROWS = 16384 + + +def fused_mqa_topk_min_keys(topk: int) -> int: + """Longest-row key count from which fused_mqa_topk beats the dense path + on SM100f (shorter chunks stay dense): 32K, or 64K for top-2048.""" + return 1 << 16 if topk > 1024 else 1 << 15 + + +# Mirrors csrc/fused_mqa_topk/constants.h (the op checks the workspace shapes). +_CANDIDATES_PER_ROW = 16384 # kCandidatesPerRow +_RETAINED_PER_ROW = 4096 # kRetainedPerRow +_SLOTS_PER_SM = 2 # candidate slots per persistent CTA +_BLOCK_ROWS = 8 # kBlockRows: query rows per block, one block per CTA at a time + + +@functools.cache +def _sm_count(index: int) -> int: + return torch.cuda.get_device_properties(index).multi_processor_count + + +@functools.cache +def fused_mqa_topk_available() -> bool: + """True when the fused_mqa_topk op is compiled in and the GPU is SM100-family.""" + return ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and hasattr(torch.ops._C, "fused_mqa_topk_prefill") + ) + + +def use_fused_mqa_topk( + *, + query_shape: tuple[int, ...], + topk: int, + fp4_cache: bool, + full_context: bool, + candidate_selection: bool = False, +) -> bool: + """Whether an indexer may use the fused_mqa_topk kernel: the single + eligibility rule of the prefill chunk planner and the indexer layers (each + prefill chunk then also needs fused_mqa_topk_chunk). + + query_shape: one token's indexer queries, (heads, head_dim) = (32, 128); + topk: the per-row top-k (pools for kpool indexers) in {512, 1024, 2048}; + fp4_cache: the MXFP4 indexer cache (and its FP4 queries) stays dense; + full_context: every rank ranks its rows over the whole context (no DCP, or + PCP after the KV all-gather); pure DCP needs the per-shard top-k merge; + candidate_selection: v4.1 two-level selection needs the logits. + VLLM_BATCH_INVARIANT keeps the dense path: the fused gate depends on the + chunk's total keys, and ties at the k-th score are taken in atomic order. + """ + return ( + envs.VLLM_USE_FUSED_MQA_TOPK + and not envs.VLLM_BATCH_INVARIANT + and not fp4_cache + and full_context + and not candidate_selection + and tuple(query_shape) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM) + and topk in FUSED_MQA_TOPK_TOPKS + and fused_mqa_topk_available() + ) + + +def fused_mqa_topk_chunk( + max_seqlen_k: int, keys: int, topk: int, rows: int, dense_rows: int +) -> bool: + """Whether a prefill chunk of a fused_mqa_topk-eligible layer takes the + fused_mqa_topk kernel: its longest row sees at least fused_mqa_topk_min_keys + keys (the dense path is as fast below), its gathered keys fit the kernel, + and its rows keep the SMs busy. + + Each SM scans every key of one 8-row block at a time, so the blocks run in + waves of one per SM. A last wave that fills less than half the SMs costs a + whole wave and makes the kernel slower than dense, unless dense needs more + than one logits call (rows > dense_rows, its rows per call). + """ + sms = _sm_count(torch.accelerator.current_device_index()) + last_wave = (-(-rows // _BLOCK_ROWS) - 1) % sms + 1 # blocks, 1..sms + return ( + max_seqlen_k >= fused_mqa_topk_min_keys(topk) + and keys <= FUSED_MQA_TOPK_MAX_KEYS + and (rows > dense_rows or 2 * last_wave >= sms) + ) + + +def get_fused_mqa_topk_workspaces( + device: torch.device, rows: int +) -> list[torch.Tensor]: + """The workspaces of one fused_mqa_topk call over ``rows`` query rows, in the + op's argument order (see fused_mqa_topk.cu), all fresh: the zeroed block + counter, then the scratch.""" + index = ( + device.index + if device.index is not None + else torch.accelerator.current_device_index() + ) + blocks = -(-rows // _BLOCK_ROWS) + # One persistent CTA per SM, at most one per block; S = 2 slots per CTA. + slot_rows = _BLOCK_ROWS * _SLOTS_PER_SM * min(blocks, _sm_count(index)) + empty = functools.partial(torch.empty, device=device) + return [ + torch.zeros(2, dtype=torch.int32, device=device), # block counter + empty((slot_rows, _CANDIDATES_PER_ROW), dtype=torch.uint16), # fp24 lows + empty((slot_rows, _CANDIDATES_PER_ROW), dtype=torch.int32), # indices + empty((slot_rows, _RETAINED_PER_ROW), dtype=torch.float32), # retained + empty((slot_rows, _RETAINED_PER_ROW), dtype=torch.int32), # their indices + empty(blocks * _BLOCK_ROWS, dtype=torch.int32), # per-row candidate counts + empty(blocks, dtype=torch.int32), # blocks queued for the exact rescan + ] + + +def fused_mqa_topk_key_scales(scales: torch.Tensor) -> torch.Tensor: + """Per-key fp32 scales from the gathered indexer K-cache scales. + + scales: [keys, 4] uint8. SM100 caches store each power-of-two scale as four + equal UE8M0 bytes; the fp32 value 2**(e - 127) has bit pattern ``e << 23``. + """ + exponents = scales.view(torch.int32).squeeze(-1) & 0xFF + return (exponents << 23).view(torch.float32) + + +def fused_mqa_topk_prefill( + q: torch.Tensor, + k: torch.Tensor, + scales: torch.Tensor, + weights: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + output: torch.Tensor, +) -> None: + """Fused replacement of ``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``. + + q: [rows, 32, 128] fp8 e4m3; k: [keys, 128] fp8 e4m3 (keys <= + FUSED_MQA_TOPK_MAX_KEYS); scales: [keys] fp32; weights: [rows, 32] fp32; + starts/ends: [rows] int32, the visible key range ``[start, end)`` of each + row; output: [rows, k] int32 with k (the top-k) 512, 1024 or 2048. All + contiguous; q, k, scales, weights and output 16-byte aligned (the op checks + the contract). + + Writes the same convention as ``ops.top_k_per_row_prefill``: indices + relative to the row's ``start`` (i.e. ``key - start``), unordered, ``-1`` + padded when fewer than k keys are visible. + """ + logger.info_once( + "Using the single-launch fused MQA+top-k prefill kernel " + "(FP8, 32 heads, top-k %d).", + output.shape[1], + ) + torch.ops._C.fused_mqa_topk_prefill( + q, + k, + scales, + weights, + starts, + ends, + output, + *get_fused_mqa_topk_workspaces(q.device, q.shape[0]), + ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 25b478fe1422..89358fba8faa 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -31,6 +31,13 @@ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( select_candidate_blocks as _select_candidate_blocks, ) +from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + fused_mqa_topk_chunk, + fused_mqa_topk_key_scales, + fused_mqa_topk_prefill, + get_fused_mqa_topk_workspaces, + use_fused_mqa_topk, +) from vllm.model_executor.layers.indexer_topk import ( RADIX_TOPK_WORKSPACE_SIZE, deep_select_topk, @@ -59,6 +66,7 @@ ) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadata, + dense_logits_rows, tp_prefill_row_sharding_supported, ) from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton @@ -483,6 +491,14 @@ def sparse_attn_indexer( ) assert candidate_block_size > 0 + fused_mqa_topk = use_fused_mqa_topk( + query_shape=q_quant.shape[1:], + topk=topk_tokens, + fp4_cache=use_fp4_cache, + full_context=use_pcp or dcp_world_size == 1, + candidate_selection=candidate_blocks is not None, + ) + # DeepSelect ranks SM100's BF16 decode logits directly; other top-k # backends and the DCP merge read an FP32 copy. fp32_decode_logits = topk_backend != "deep_select" or dcp_world_size != 1 @@ -506,6 +522,10 @@ def sparse_attn_indexer( ) profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) + if fused_mqa_topk: + # Allocate the largest fused_mqa_topk workspaces so they are + # accounted for (freed again, they overlap the dense logits below). + get_fused_mqa_topk_workspaces(q_quant.device, q_quant.shape[0]) # Decode rows per step: each decode is padded to 1 + speculative tokens. # The profiling run sets the vLLM config; other dummy runs may not. @@ -516,6 +536,8 @@ def sparse_attn_indexer( cfg.scheduler_config.max_num_batched_tokens, ) * (1 + cfg.num_speculative_tokens) + # Needed with the fused kernel too: chunks fused_mqa_topk_chunk keeps + # dense take this path (row-split to stay within the budget). # Dummy allocation for the peak logits memory during inference: a # prefill chunk's, or the decode logits, which eager and piecewise # steps allocate outside the CUDA graph pools. DeepGEMM emits them in @@ -732,10 +754,29 @@ def sparse_attn_indexer( # DeepGEMM returns compressed logits: row i starts at K position # cu_seqlen_ks[i]. None means every row starts at column zero. row_starts: torch.Tensor | None = None + # Rows per dense logits call under the logits budget. + dense_rows = dense_logits_rows(chunk.max_seqlen_k) use_deep_select = False if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) + elif fused_mqa_topk and fused_mqa_topk_chunk( + chunk.max_seqlen_k, + k_quant.shape[0], + topk_tokens, + q_slice.shape[0], + dense_rows, + ): + fused_mqa_topk_prefill( + q_slice, + k_quant, + fused_mqa_topk_key_scales(k_scale), + weights[row_start:row_end], + cu_seqlen_ks, + cu_seqlen_ke, + topk_indices, + ) + logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) else: # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. @@ -749,84 +790,96 @@ def sparse_attn_indexer( k_scale_cast = k_scale.view( torch.int32 if mqa_logits_use_mx_scales() else torch.float32 ).squeeze(-1) - if current_platform.is_xpu(): - if q_scale_slice is not None: - raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") - logits = torch.ops.vllm.xpu_fp8_mqa_logits( - q_slice_cast, - k_quant_cast, - k_scale_cast, - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - ) - row_starts = cu_seqlen_ks - row_ends = cu_seqlen_ke - else: - q_scale_slice, chunk_weights = mqa_logits_mx_inputs( - q_slice_cast, q_scale_slice, weights[row_start:row_end] - ) - row_ends = cu_seqlen_ke - cu_seqlen_ks - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - chunk_weights, - cu_seqlen_ks, - cu_seqlen_ke, - chunk.max_seqlen_k, - ) - # DeepSelect ranks the bf16 logits in place of their exact fp32 - # copy; fp32 logits keep the per-row top-k. - use_deep_select = ( - topk_backend == "deep_select" - and logits.dtype == torch.bfloat16 - and dcp_world_size == 1 - and is_deep_select_supported(logits, topk_tokens) - ) - if not use_deep_select: - logits = as_fp32_logits(logits, row_ends) - num_rows = logits.shape[0] - if candidate_blocks is not None: - # Two-level selection (v4.1): the candidate source - # publishes its top blocks; later indexers mask their - # scores to them. Both before the row top-k. - chunk_candidates = candidate_blocks[ - chunk.token_start : chunk.token_end - ] - if candidate_write: - _select_candidate_blocks( + # Chunks planned for fused_mqa_topk are not bounded by the logits + # budget: split their rows (budget-planned chunks, the only ones + # under DCP, are one window, so the merge below sees all the + # rows). + for sub in range(0, q_slice.shape[0], dense_rows): + rows = slice(sub, min(sub + dense_rows, q_slice.shape[0])) + tokens = slice(row_start + rows.start, row_start + rows.stop) + if current_platform.is_xpu(): + if q_scale_slice is not None: + raise RuntimeError( + "XPU fp8_mqa_logits does not support FP4 Q" + ) + logits = torch.ops.vllm.xpu_fp8_mqa_logits( + q_slice_cast[rows], + k_quant_cast, + k_scale_cast, + weights[tokens], + cu_seqlen_ks[rows], + cu_seqlen_ke[rows], + ) + row_starts = cu_seqlen_ks[rows] + row_ends = cu_seqlen_ke[rows] + else: + q_scale_rows, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast[rows], + q_scale_slice[rows] if q_scale_slice is not None else None, + weights[tokens], + ) + row_ends = cu_seqlen_ke[rows] - cu_seqlen_ks[rows] + logits = fp8_fp4_mqa_logits( + (q_slice_cast[rows], q_scale_rows), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks[rows], + cu_seqlen_ke[rows], + chunk.max_seqlen_k, + ) + # DeepSelect ranks the bf16 logits in place of their exact + # fp32 copy; fp32 logits keep the per-row top-k. + use_deep_select = ( + topk_backend == "deep_select" + and logits.dtype == torch.bfloat16 + and dcp_world_size == 1 + and is_deep_select_supported(logits, topk_tokens) + ) + if not use_deep_select: + logits = as_fp32_logits(logits, row_ends) + num_rows = logits.shape[0] + if candidate_blocks is not None: + # Two-level selection (v4.1): the candidate source + # publishes its top blocks; later indexers mask their + # scores to them. Both before the row top-k. + chunk_candidates = candidate_blocks[tokens] + if candidate_write: + _select_candidate_blocks( + logits, + row_starts, + row_ends, + chunk_candidates.shape[1], + candidate_block_size, + chunk_candidates, + ) + else: + _apply_candidate_mask( + logits, + row_starts, + row_ends, + chunk_candidates, + candidate_block_size, + ) + if use_deep_select: + deep_select_topk( logits, - row_starts, - row_ends, - chunk_candidates.shape[1], - candidate_block_size, - chunk_candidates, + topk_tokens, + end=row_ends, + output_idx=topk_indices[rows], ) else: - _apply_candidate_mask( + ops.top_k_per_row_prefill( logits, - row_starts, + torch.zeros_like(row_ends) + if row_starts is None + else row_starts, row_ends, - chunk_candidates, - candidate_block_size, + topk_indices[rows], + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, ) - if use_deep_select: - deep_select_topk( - logits, topk_tokens, end=row_ends, output_idx=topk_indices - ) - else: - ops.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends) - if row_starts is None - else row_starts, - row_ends, - topk_indices, - num_rows, - logits.stride(0), - logits.stride(1), - topk_tokens, - ) if deinterleave_idx is None: # Under the PCP path the top-k already ran over the whole diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index 52d4f4c3b017..430886590d7c 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_mla.py +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -87,9 +87,10 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, + fused_mqa_topk_eligible: bool = False, ) -> list[tuple[slice, slice]]: - # workspace_size/max_logits_bytes are unused: no flat-gather - # workspace or dense M*N buffer is allocated. + # workspace_size/max_logits_bytes/fused_mqa_topk_eligible are unused: no + # flat-gather workspace or dense M*N buffer is allocated. num_reqs = compressed_seq_lens_cpu.shape[0] total_query_len = int(prefill_query_lens_cpu.sum().item()) return [ diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 65ba7be1fdec..0b8d4fa788f2 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -11,6 +11,13 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + fused_mqa_topk_chunk, + fused_mqa_topk_key_scales, + fused_mqa_topk_prefill, + get_fused_mqa_topk_workspaces, + use_fused_mqa_topk, +) from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.model_executor.layers.sparse_attn_indexer import ( PrefillRowShard, @@ -44,6 +51,7 @@ ) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadata, + dense_logits_rows, ) from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton from vllm.v1.worker.workspace import current_workspace_manager @@ -144,6 +152,13 @@ def sparse_attn_indexer_kpool( fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) + fused_mqa_topk = use_fused_mqa_topk( + query_shape=q_quant.shape[1:], + topk=topk_tokens // max(index_kpool, 1), + fp4_cache=use_fp4_cache, + full_context=dcp_world_size == 1, + ) + # assert isinstance(attn_metadata, dict) if not isinstance(attn_metadata, dict): # Reserve workspace for indexer during profiling run @@ -155,6 +170,10 @@ def sparse_attn_indexer_kpool( scales_spec, ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) + # Allocate the largest fused_mqa_topk workspaces so they are accounted + # for (freed again, they overlap the dense logits below). + if fused_mqa_topk: + get_fused_mqa_topk_workspaces(q_quant.device, q_quant.shape[0]) # Reserve profiler-visible memory for the worst-case decode logits, # whose shape is [B * next_n, max_pool_len]. This profiling branch @@ -355,51 +374,75 @@ def sparse_attn_indexer_kpool( ).squeeze(-1) num_rows = q_slice.shape[0] - if chunk.local_total_seq_lens == 0: - logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) - else: - q_scale_slice, chunk_weights = mqa_logits_mx_inputs( - q_slice_cast, q_scale_slice, weights[row_start:row_end] - ) - # Compressed: row i starts at K position cu_seqlen_ks[i]. - row_ends = cu_seqlen_ke - cu_seqlen_ks - logits = as_fp32_logits( - fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - chunk_weights, - cu_seqlen_ks, - cu_seqlen_ke, - chunk.max_seqlen_k, - ), - row_ends, - ) - # kpool: logits are pool-granular (compress_ratio == index_kpool), # so topk selects pools. We pick topk_tokens // kpool pools then # expand each pool back to its kpool constituent tokens. select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens if index_kpool > 1: pool_topk = torch.empty( - (num_rows, select_k), dtype=torch.int32, device=logits.device + (num_rows, select_k), dtype=torch.int32, device=q_slice.device ) topk_dst = pool_topk else: topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens] - if logits.shape[1] == 0: + # Rows per dense logits call under the logits budget. + dense_rows = dense_logits_rows(chunk.max_seqlen_k) + if chunk.local_total_seq_lens == 0: + logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) topk_dst.fill_(-1) - else: - torch.ops._C.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends), - row_ends, + elif fused_mqa_topk and fused_mqa_topk_chunk( + chunk.max_seqlen_k, k_quant.shape[0], select_k, num_rows, dense_rows + ): + fused_mqa_topk_prefill( + q_slice, + k_quant, + fused_mqa_topk_key_scales(k_scale), + weights[row_start:row_end], + cu_seqlen_ks, + cu_seqlen_ke, topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, ) + logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) + else: + # Chunks planned for fused_mqa_topk are not bounded by the logits + # budget: split their rows (budget-planned chunks, the only ones + # under DCP, are one window, so the merge below sees all the + # rows). + # One pass also for 0 rows: the merge is collective. + for sub in range(0, max(1, num_rows), dense_rows): + rows = slice(sub, min(sub + dense_rows, num_rows)) + q_scale_rows, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast[rows], + q_scale_slice[rows] if q_scale_slice is not None else None, + weights[row_start + rows.start : row_start + rows.stop], + ) + # Compressed: row i starts at K position cu_seqlen_ks[i]. + row_ends = cu_seqlen_ke[rows] - cu_seqlen_ks[rows] + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast[rows], q_scale_rows), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks[rows], + cu_seqlen_ke[rows], + chunk.max_seqlen_k, + ), + row_ends, + ) + if logits.shape[1] == 0: + topk_dst[rows].fill_(-1) + else: + torch.ops._C.top_k_per_row_prefill( + logits, + torch.zeros_like(row_ends), + row_ends, + topk_dst[rows], + logits.shape[0], + logits.stride(0), + logits.stride(1), + select_k, + ) # Local pool ids -> request-relative global pool ids. _merge_dcp_topk_global( logits, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index b6e3d1fc8c05..c03d19d98cbb 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -10,6 +10,11 @@ from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + FUSED_MQA_TOPK_MAX_KEYS, + FUSED_MQA_TOPK_MAX_ROWS, + use_fused_mqa_topk, +) from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, @@ -927,6 +932,19 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): return max_model_len * (576 * 2 // 132) * factor +def dense_logits_rows(max_seqlen_k: int, max_logits_bytes: int | None = None) -> int: + """Query rows per dense MQA-logits call whose fp32 [rows, max_seqlen_k] + logits fit max_logits_bytes (default: VLLM_SPARSE_INDEXER_MAX_LOGITS_MB). + + The prefill chunk planner sub-chunks requests with it; the indexers split + the rows of chunks planned for fused_mqa_topk (not bounded by the budget) + that take the dense path. + """ + if max_logits_bytes is None: + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + return max(1, max_logits_bytes // (4 * max(1, max_seqlen_k))) + + def _supports_varlen_paged_mqa_logits() -> bool: return ( current_platform.is_cuda() @@ -1030,6 +1048,24 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: and self.vllm_config.speculative_config.enable_adaptive_verification ) self.indexer_uses_fp4 = dsa_indexer_uses_fp4(self.vllm_config) + # Plan prefill chunks for the fused MQA+top-k kernel (not bounded by + # the logits budget). PCP + DCP has its own planner. v4.1 candidate + # selection is not excluded: its pre-candidate indexers take + # fused_mqa_topk and the others row-split their dense logits. Indexers + # that pad their heads to 32 run fused_mqa_topk on budget-planned + # chunks. + model_config = self.vllm_config.model_config + cfg = model_config.hf_text_config if model_config is not None else None + self.fused_mqa_topk_eligible = cfg is not None and use_fused_mqa_topk( + query_shape=( + getattr(cfg, "index_n_heads", 0), + getattr(cfg, "index_head_dim", 0), + ), + topk=getattr(cfg, "index_topk", 0) + // (getattr(cfg, "index_kpool", None) or 1), + fp4_cache=self.indexer_uses_fp4, + full_context=self.dcp_world_size == 1, + ) next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n @@ -1372,6 +1408,7 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, + fused_mqa_topk_eligible: bool = False, ) -> list[tuple[slice, slice]]: """Split this step's prefill requests into chunks, respecting: - N constraint: total_seq_lens <= workspace_size (existing O(N) @@ -1381,6 +1418,13 @@ def _split_indexer_prefill_chunks( When a single request-level chunk still exceeds the logits budget, sub-chunks on the query dimension (M) to bound peak memory. + With ``fused_mqa_topk_eligible`` (chunks for the fused MQA+top-k kernel, + which needs no logits buffer) chunks whose gathered keys fit + FUSED_MQA_TOPK_MAX_KEYS are bounded by FUSED_MQA_TOPK_MAX_ROWS rows + instead of the logits budget. Requests are never packed past that key + limit, so only a single request longer than it takes the logits-budget + path. Chunks the dispatch keeps dense are row-split by the indexer. + Returns list of (req_slice, query_slice) tuples. """ chunks: list[tuple[slice, slice]] = [] @@ -1397,7 +1441,15 @@ def _split_indexer_prefill_chunks( compressed_seq_lens_cpu[end].item(), ) new_m, new_n = chunk_m + q, chunk_n + s - if new_n <= workspace_size and new_m * new_n <= max_logits_elems: + if fused_mqa_topk_eligible and new_n <= FUSED_MQA_TOPK_MAX_KEYS: + fits = new_m <= FUSED_MQA_TOPK_MAX_ROWS + else: + # fused_mqa_topk planning never packs past the key limit. + fits = ( + not fused_mqa_topk_eligible + and new_m * new_n <= max_logits_elems + ) + if new_n <= workspace_size and fits: chunk_m, chunk_n = new_m, new_n end += 1 else: @@ -1414,8 +1466,12 @@ def _split_indexer_prefill_chunks( req_slice = slice(start + request_offset, end + request_offset) max_q = ( - max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) + dense_logits_rows(chunk_n, max_logits_bytes) + if chunk_n > 0 + else max(1, chunk_m) ) + if fused_mqa_topk_eligible and chunk_n <= FUSED_MQA_TOPK_MAX_KEYS: + max_q = FUSED_MQA_TOPK_MAX_ROWS for q_off in range(0, chunk_m, max_q): sub_m = min(max_q, chunk_m - q_off) chunks.append((req_slice, slice(q_off, q_off + sub_m))) @@ -1551,6 +1607,7 @@ def build( self.max_prefill_buffer_size, max_logits_bytes, request_offset=num_decodes, + fused_mqa_topk_eligible=self.fused_mqa_topk_eligible, ) chunks = [] diff --git a/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py b/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py index 77c9eba68ca4..837f7ae964e3 100644 --- a/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py +++ b/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py @@ -392,6 +392,7 @@ def _split_indexer_prefill_chunks( # type: ignore[override] workspace_size: int, max_logits_bytes: int, request_offset: int = 0, + fused_mqa_topk_eligible: bool = False, ) -> list[tuple[slice, slice]]: return split_prefill_chunks( compressed_seq_lens_cpu,