[Perf][Kernel] Single-launch fused scoring+top-k kernel for sparse-indexer prefill, without materializing logits

.buildkite/test_areas/kernels.yaml

@@ -694,6 +694,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
@@ -705,6 +707,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

CMakeLists.txt

@@ -1702,6 +1702,19 @@ endif()
 # For CUDA we also build and ship some external projects.
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
+    if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9)
+      cuda_archs_loose_intersection(FUSED_MQA_ARCHS "10.0f" "${CUDA_ARCHS}")
+      if(FUSED_MQA_ARCHS)
+        set(FUSED_MQA_SRC "csrc/fused_mqa_topk/fused_mqa_topk.cu")
+        set_gencode_flags_for_srcs(SRCS "${FUSED_MQA_SRC}" CUDA_ARCHS "${FUSED_MQA_ARCHS}")
+        set_property(SOURCE ${FUSED_MQA_SRC} APPEND PROPERTY INCLUDE_DIRECTORIES
+          "${deepgemm_SOURCE_DIR}/deep_gemm/include"
+          "${deepgemm_SOURCE_DIR}/third-party/cutlass/include")
+        target_sources(_C_stable_libtorch PRIVATE ${FUSED_MQA_SRC})
+        target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_MQA_TOPK=1)
+      endif()
+    endif()
+
     include(cmake/external_projects/deepselect.cmake)
     include(cmake/external_projects/fmha_sm100.cmake)
     include(cmake/external_projects/flashmla.cmake)

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 <cstdint>
+#include <cutlass/cutlass.h>
+
+#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<uint32_t>(packed_idx) >> kCandidateIndexBits) << 16 |
+         static_cast<uint32_t>(value);
+}
+
+CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) {
+  return static_cast<int32_t>(static_cast<uint32_t>(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<CandidateValue>(payload));
+  __stcg(index_dst,
+         static_cast<int32_t>(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

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 <cuda_fp16.h>
+#include <cutlass/arch/barrier.h>
+
+#include <cstdint>
+
+#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<uint32_t>(__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 <class T>
+__device__ __forceinline__ T volatile_load(const T& x) {
+  return *reinterpret_cast<const volatile T*>(&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 <uint32_t Shift>
+__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 <uint32_t Shift>
+__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 <bool kWithLinear>
+__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

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 <cstdint>
+
+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