From 8ab8cb3b2f7a095c94513b3f8aad800ee79b3ad7 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Thu, 8 Oct 2026 10:08:28 +0000 Subject: [PATCH 01/23] init Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 16 + csrc/fused_mqa_topk/candidates.cuh | 71 + csrc/fused_mqa_topk/fallback_core.cuh | 651 +++++++++ csrc/fused_mqa_topk/fallback_shared.cuh | 268 ++++ csrc/fused_mqa_topk/fast_select.cuh | 683 +++++++++ csrc/fused_mqa_topk/fused_mqa_topk.cu | 292 ++++ csrc/fused_mqa_topk/handoff.cuh | 77 + csrc/fused_mqa_topk/mqa_seeded_kernel.cuh | 423 ++++++ csrc/fused_mqa_topk/sample_core.cuh | 625 ++++++++ csrc/fused_mqa_topk/sample_shared.cuh | 105 ++ csrc/fused_mqa_topk/scan.cuh | 1106 ++++++++++++++ csrc/fused_mqa_topk/sm100_candidate_scan.cuh | 1276 +++++++++++++++++ csrc/fused_mqa_topk/upstream_select.cuh | 534 +++++++ csrc/ops.h | 10 + csrc/torch_bindings.cpp | 15 + .../kernels/attention/test_fused_mqa_topk.py | 309 ++++ vllm/envs.py | 8 + .../kernels/attention/dsa/fused_mqa_topk.py | 238 +++ .../layers/sparse_attn_indexer.py | 178 ++- vllm/v1/attention/backends/mla/indexer.py | 37 +- 20 files changed, 6863 insertions(+), 59 deletions(-) create mode 100644 csrc/fused_mqa_topk/candidates.cuh create mode 100644 csrc/fused_mqa_topk/fallback_core.cuh create mode 100644 csrc/fused_mqa_topk/fallback_shared.cuh create mode 100644 csrc/fused_mqa_topk/fast_select.cuh create mode 100644 csrc/fused_mqa_topk/fused_mqa_topk.cu create mode 100644 csrc/fused_mqa_topk/handoff.cuh create mode 100644 csrc/fused_mqa_topk/mqa_seeded_kernel.cuh create mode 100644 csrc/fused_mqa_topk/sample_core.cuh create mode 100644 csrc/fused_mqa_topk/sample_shared.cuh create mode 100644 csrc/fused_mqa_topk/scan.cuh create mode 100644 csrc/fused_mqa_topk/sm100_candidate_scan.cuh create mode 100644 csrc/fused_mqa_topk/upstream_select.cuh create mode 100644 tests/kernels/attention/test_fused_mqa_topk.py create mode 100644 vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 7694e6f8bc46..8922c022997c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1702,6 +1702,22 @@ 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") + # Single-launch fused MQA-logits + top-k prefill kernel. + set_property(SOURCE ${FUSED_MQA_SRC} APPEND PROPERTY COMPILE_OPTIONS + "-O3" "--expt-relaxed-constexpr" "--expt-extended-lambda") + 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) diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh new file mode 100644 index 000000000000..a426ee0f03e1 --- /dev/null +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025 DeepSeek +#pragma once +#include +#include +namespace fused_candidates { +using CandidateValue = uint16_t; +constexpr uint32_t kCandidateIndexBits = 20; +constexpr uint32_t kCandidateIndexMask = (1u << kCandidateIndexBits) - 1u; + +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_ordered_fp32_code(const float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); +} + +CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) { + return candidate_ordered_fp32_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 float candidate_decode_score(const CandidateValue value, + const int32_t packed_idx) { + const uint32_t code = candidate_load_score_code(value, packed_idx); + const uint32_t ordered = code << 8; + const uint32_t bits = + (ordered & 0x80000000u) ? (ordered ^ 0x80000000u) : ~ordered; + return __uint_as_float(bits); +} + +CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { + return static_cast(static_cast(packed_idx) & + kCandidateIndexMask); +} + +CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, + int32_t* index_dst, const float bq, + const uint32_t kv_index) { + const uint32_t payload = candidate_fp24_code(bq); + __stcg(value_dst, static_cast(payload)); + __stcg(index_dst, + static_cast(candidate_pack_index(payload, kv_index))); +} + +CUTLASS_DEVICE void store_candidate_payload(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_record(CandidateValue* value_dst, + int32_t* index_dst, + const uint32_t payload, + const uint32_t kv_index) { + store_candidate_payload(value_dst, index_dst, payload, kv_index); +} + +} // namespace fused_candidates diff --git a/csrc/fused_mqa_topk/fallback_core.cuh b/csrc/fused_mqa_topk/fallback_core.cuh new file mode 100644 index 000000000000..f4ae2d112ad4 --- /dev/null +++ b/csrc/fused_mqa_topk/fallback_core.cuh @@ -0,0 +1,651 @@ +// SPDX-License-Identifier: MIT +// Experimental derivative of DeepGEMM 2.8 sm100_mqa_logits.cuh. +// Persistent DeepGEMM FP8 MQA producer with fixed-threshold candidate emission. +#pragma once +#ifndef MQA_TMEM_ROWS + #define MQA_TMEM_ROWS 2 +#endif + +#include +#include "fallback_shared.cuh" +#ifndef MQA_HANDOFF_INTERVAL + #define MQA_HANDOFF_INTERVAL 64 +#endif +static_assert(MQA_HANDOFF_INTERVAL == 32 || MQA_HANDOFF_INTERVAL == 64 || + MQA_HANDOFF_INTERVAL == 128); + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "sm100_candidate_scan.cuh" + +// Shared SM100 MQA logits core plus contiguous-KV and paged entries +// Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature + +namespace seed_fallback { +using namespace deep_gemm; + +// Convert runtime valid-token count to `cute::Int` so token loops stay +// compile-time constant +template +CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, + Fn&& fn) { + if constexpr (kCandidate <= 1) { + fn(cute::Int<1>{}); + } else if (num_block_tokens >= kCandidate) { + fn(cute::Int{}); + } else { + dispatch_num_block_tokens(num_block_tokens, + static_cast(fn)); + } +} + +// Load heads using power-of-two TMEM chunks. +template +CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, + float* accum) { + constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; + if constexpr (kNumRemaining > 0) { + constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 + : kNumRemaining >= 32 ? 32 + : kNumRemaining >= 16 ? 16 + : kNumRemaining >= 8 ? 8 + : 4; + ptx::tmem_load_32dp32b( + tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); + load_tmem_heads_decomposed(tmem_col, accum); + } +} + +// Shared device core parameterized by dtype and scheduler geometry/addressing +template < + uint32_t kNumHeads, uint32_t kHeadDim, bool kIsMXSF, + bool kIsCompressedLogits, bool kCleanLogits, uint32_t BLOCK_Q, + uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages, + uint32_t kNumSMs, uint32_t kNumSpecializedThreads, uint32_t kNumMathThreads, + typename qk_dtype_t, typename logits_dtype_t, typename reduce_dtype_t, + typename MakeScheduler, uint32_t kNumMathWarpGroups = kNumMathThreads / 128> +CUTLASS_DEVICE void sm100_mqa_logits_core_impl( + const uint32_t num_q_tokens, FallbackBuffers buffers, + const cute::TmaDescriptor& tensor_map_q, + const cute::TmaDescriptor& tensor_map_sf_q, + const cute::TmaDescriptor& tensor_map_kv, + const cute::TmaDescriptor& tensor_map_sf_kv, + const cute::TmaDescriptor& tensor_map_weights, + const MakeScheduler& make_scheduler) { + constexpr bool kIsFP4 = cute::is_same_v; + + const auto sm_idx = blockIdx.x; + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; + + const auto kNegInf = -cute::numeric_limits::infinity(); + + if (warp_idx == kSpecWarpStart) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_sf_q); + cute::prefetch_tma_descriptor(&tensor_map_weights); + cute::prefetch_tma_descriptor(&tensor_map_kv); + cute::prefetch_tma_descriptor(&tensor_map_sf_kv); + } + + static constexpr uint32_t kNumTmemStages = 2; + static_assert(BLOCK_Q == 8 && kNumMathThreads == 256 && !kIsMXSF); + static constexpr uint32_t kNumUTCCPAlignedElems = 128; + static constexpr uint32_t UMMA_M = 128; + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); + static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; + static constexpr uint32_t kNumSFQ = + kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumSFKV = + kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumQKBytesPerToken = + kIsFP4 ? (kHeadDim / 2) : kHeadDim; + static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = + BLOCK_QH * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = + SPLIT_KV * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = + kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; + static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = + kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightElementsPerRow = + kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); + static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = + BLOCK_Q * kNumWeightBytesPerRow; + + DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, + "Invalid threads"); + DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and + SPLIT_KV % kNumUTCCPAlignedElems == 0, + "Invalid `SPLIT_KV`"); + DG_STATIC_ASSERT(not(kIsCompressedLogits and kCleanLogits), + "Compressed logits cannot be cleaned in-kernel"); + + using SharedStorage = layout::MQALogitsSharedStorage< + kNumHeads, kHeadDim, kIsMXSF, BLOCK_Q, SPLIT_KV, kNumQStages, + kNumKVStages, kNumTmemStages, qk_dtype_t, reduce_dtype_t>; + extern __shared__ __align__(SharedStorage::kSwizzleAlignment) + uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + constexpr uint32_t kRingSlots = 8; + constexpr uint32_t kRecordBlockMask = 255; + auto& stream = + *reinterpret_cast(smem_buffer + sizeof(SharedStorage)); + stream.init(threadIdx.x); + + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; + DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); + + if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumQStages; ++i) { + smem.full_q_barriers[i].init(1); + smem.empty_q_barriers[i].init(kNumMathThreads + 32); + } +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + smem.full_kv_barriers[i].init(1); + smem.empty_kv_barriers[i].init(kIsMXSF ? 1 : kNumMathThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumTmemStages; ++i) { + smem.full_tmem_barriers[i].init(1); + smem.empty_tmem_barriers[i].init(256); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + + if (warp_idx == kSpecWarpStart + 2) + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); + __syncthreads(); + + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + + constexpr uint32_t kNumSpecializedRegisters = 56; + constexpr uint32_t kNumMathRegisters = 224; + + cudaGridDependencySynchronize(); + + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + if (cute::elect_one_sync()) { + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + tma::copy( + &tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); + if constexpr (kIsMXSF) + tma::copy( + &tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_sf_q[q_stage_idx], 0, q_token_base); + 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_SF_Q_SIZE_PER_STAGE + + SMEM_WEIGHT_SIZE_PER_STAGE); + } + } + __syncwarp(); + } else if (warp_idx == kSpecWarpStart + 1) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + + auto scheduler = make_scheduler(sm_idx); + uint32_t cached_kv_page_base = 0; + uint32_t cached_kv_page_coord = 0; + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + cached_kv_page_base = cute::numeric_limits::max(); +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if constexpr (decltype(scheduler)::kIsPaged) { + constexpr uint32_t kPageKV = decltype(scheduler)::kPageKV; + constexpr uint32_t kNumPagesPerSplit = + decltype(scheduler)::kNumPagesPerSplit; + DG_STATIC_ASSERT(kNumPagesPerSplit <= 32, + "Split spans more pages than a warp can cache"); + + const uint32_t kv_page_base = + (kv_base + kv_split_idx) * kNumPagesPerSplit; + if (kv_page_base < cached_kv_page_base or + kv_page_base + kNumPagesPerSplit > cached_kv_page_base + 32) { + cached_kv_page_base = (kv_page_base / 32) * 32; + cached_kv_page_coord = scheduler.get_kv_page_coord_by_page_offset( + cached_kv_page_base + lane_idx); + } + + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + if (cute::elect_one_sync()) + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); + __syncwarp(); + + int page_coords[kNumPagesPerSplit]; +#pragma unroll + for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; + ++page_idx) { + const auto src_lane = + static_cast(kv_page_base - cached_kv_page_base + page_idx); + page_coords[page_idx] = + __shfl_sync(0xffffffff, cached_kv_page_coord, src_lane); + } + + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; + ++page_idx) { + tma::copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx] + + page_idx * kPageKV * kNumQKBytesPerToken, + 0, 0, 1, page_coords[page_idx]); + tma::copy( + &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_sf_kv[kv_stage_idx] + page_idx * kPageKV, 0, + page_coords[page_idx]); + } + 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 (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); + tma::copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); + 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<56>(); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); + + auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { + DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, + "Invalid aligned elements"); + uint32_t values[4]; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) + values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); + __syncwarp(); + ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], + values[3]); + }; + + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + + if constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; + utccp_required_smem_warp_transpose(smem_ptr); + } + cutlass::arch::fence_view_async_shared(); +#pragma unroll + for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + if (cute::elect_one_sync()) + cute::SM100_UTCCP_4x32dp128bit_1cta::copy( + sf_desc, kTmemStartColOfSFQ + i * 4); + __syncwarp(); + } + } + + 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 constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; + utccp_required_smem_warp_transpose(smem_ptr); + } + cutlass::arch::fence_view_async_shared(); + } + + if (cute::elect_one_sync()) { + if constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute::SM100_UTCCP_4x32dp128bit_1cta::copy( + sf_desc, kTmemStartColOfSFKV + i * 4); + } + } +#pragma unroll + for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + 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(); + + if constexpr (kIsMXSF) { + DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or + kHeadDim == 64 or kHeadDim == 128, + "Invalid head dim"); + + constexpr uint32_t kPackFactor = + get_smem_pack_factor(); + constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; + + using mma_op_t = + cute::conditional_t; + auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< + qk_dtype_t, qk_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, + UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K>(); +#pragma unroll + for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { + auto runtime_instr_desc = + mma::sm100::make_runtime_instr_desc_with_sf_id( + instr_desc, k * kPackFactor, k * kPackFactor); + 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); + mma_op_t::fma(a_desc, b_desc, tmem_addr, k, runtime_instr_desc, + kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ); + } + } else { + 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); +#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(); + if constexpr (kIsMXSF) + cutlass::arch::umma_arrive(reinterpret_cast( + &smem.empty_kv_barriers[kv_stage_idx])); + } + smem.empty_q_barriers[q_stage_idx].arrive(); + } + } else if (warp_idx == kSpecWarpStart + 3) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + } else if (warp_idx < kSpecWarpStart) { + cutlass::arch::warpgroup_reg_alloc<224>(); + + constexpr bool kLoadSeqBounds = true; + uint32_t seq_k_start[4]; + uint32_t seq_k_end[4]; + const auto math_warpgroup_idx = warp_idx / 4; + const auto math_thread_idx = warp_idx * 32 + lane_idx; + DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, + "Math warp groups exceed TMEM stages"); + const unsigned query_base = math_warpgroup_idx * 4; + + constexpr bool kIsReduceBF16 = not cute::is_same_v; + DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); + DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); + using weights_reg_dtype_t = + cute::conditional_t; + constexpr uint32_t kNumWeightsRegPerToken = + kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; + weights_reg_dtype_t weights[4][32]; + float accum[kNumHeads]; + + unsigned produce_epoch = 0; + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + + const auto process_q_block = [&](auto num_valid_tokens_t) { + constexpr uint32_t kNumValidTokens = 4; + uint32_t lane_counts[4] = {}; + uint32_t flush_base = 0; + float gate_score[4]; + +#pragma unroll + for (uint32_t i = 0; i < kNumValidTokens; ++i) { + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + // Use the raw-score threshold estimated from the distributed sample. + // This threshold remains fixed for the complete scan. + gate_score[i] = -CUDART_INF_F; + seq_k_start[i] = cute::min(scheduler.cu_seq_len_k_start[row], + scheduler.num_kv_tokens); + seq_k_end[i] = cute::min(scheduler.cu_seq_len_k_end[row], + scheduler.num_kv_tokens); + const auto smem_weights_row = + smem.smem_weights[q_stage_idx] + + (query_base + i) * kNumWeightElementsPerRow; + if constexpr (kIsReduceBF16) { + // Load two bf16 weights at a time as one packed shared u32 + const auto packed_row = + reinterpret_cast(smem_weights_row); +#pragma unroll + for (uint32_t j = 0; j < kNumHeads / 2; ++j) { + const auto packed = ptx::ld_shared(packed_row + j); + weights[i][j] = ptx::exchange( + *reinterpret_cast(&packed), 0); + } + } else { +#pragma unroll + for (unsigned j = 0; j < 32; ++j) + weights[i][j] = 0.5f * ptx::ld_shared(smem_weights_row + j); + } + } + + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if ((kv_split_idx % (FallbackSegment / 256)) == 0) { + __syncwarp(); + for (unsigned i = 0; i < 4; ++i) + gate_score[i] = + __uint_as_float(atomicAdd(stream.gates + query_base + i, 0u)); + } + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); +#pragma unroll + for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { + const unsigned kv_lane = (math_thread_idx % 128) + kv_half * 128; + const unsigned 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(); + } + + static_assert(!kIsReduceBF16, "FP32-only pair reduction"); + float accum_pair[MQA_TMEM_ROWS][kNumHeads]; +#pragma unroll + for (uint32_t pair = 0; pair < kNumValidTokens; + pair += MQA_TMEM_ROWS) { + // Load only MQA_TMEM_ROWS query rows at once to limit + // live accumulator registers during score reduction. + const uint32_t tmem_addr = + tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads; + ptx::tmem_load_32dp32b( + tmem_addr, reinterpret_cast(accum_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair + MQA_TMEM_ROWS >= kNumValidTokens) { + ptx::tcgen05_before_thread_sync(); + smem.empty_tmem_barriers[tmem_stage_idx].arrive(); + } + + float2 sum_0[MQA_TMEM_ROWS] = {}; + float2 sum_1[MQA_TMEM_ROWS] = {}; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { +#pragma unroll + for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { + if (pair + r >= kNumValidTokens) continue; + const uint32_t i = pair + r; + auto* a = accum_pair[r]; + auto transform = [&](uint32_t at, float2 sum) { + const auto x = make_float2(a[at], a[at + 1]); + const auto abs_x = + make_float2(fabsf(a[at]), fabsf(a[at + 1])); + const auto w = + make_float2(weights[i][at], weights[i][at + 1]); + return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); + }; + sum_0[r] = transform(j, sum_0[r]); + sum_1[r] = transform(j + 2, sum_1[r]); + } + } +#pragma unroll + for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { + if (pair + r >= kNumValidTokens) continue; + const uint32_t i = pair + r; + const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); + const float reduced = sum.x + sum.y; + const float bq = reduced * static_cast(scale_kv); + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + bool valid = kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq) && + bq > gate_score[i]; + unsigned position = + (kv_split_idx % (FallbackSegment / 256)) * 256 + kv_lane; + stream.ring[produce_epoch & 1][query_base + i][position] = + valid ? bq : -CUDART_INF_F; + } + } + } // KV halves; query ownership is independent of the KV half. + if (((kv_split_idx + 1) % (FallbackSegment / 256)) == 0 || + kv_split_idx + 1 == num_kv_splits) { + __threadfence_block(); + + fb_sync(); + unsigned tid = math_thread_idx; + unsigned row_base = scheduler.get_logits_row(q_block_idx, 0); + for (unsigned local = 0; local < 8; ++local) + if (stream.kept[local] > FallbackRetain - FallbackSegment) + fb_compact(stream, buffers, local, row_base + local, tid); + fb_sync(); + unsigned local = tid / 32; + fb_select(stream, buffers, local, row_base + local, produce_epoch, + kv_base + (kv_split_idx / (FallbackSegment / 256)) * + FallbackSegment, + (kv_split_idx % (FallbackSegment / 256) + 1) * 256, + tid % 32, false); + fb_sync(); + if (kv_split_idx + 1 == num_kv_splits) + for (unsigned local = 0; local < 8; ++local) + if (stream.kept[local] > 2048) + fb_compact(stream, buffers, local, row_base + local, tid, + true); + fb_sync(); + ++produce_epoch; + } + } + }; + + if constexpr (decltype(scheduler)::kHasPartialBlock) + dispatch_num_block_tokens( + scheduler.get_num_block_tokens(q_block_idx), process_q_block); + else + process_q_block(cute::Int{}); + + unsigned local = warp_idx, + row = scheduler.get_logits_row(q_block_idx, local), + count = stream.kept[local]; + if (row < scheduler.num_q_tokens) + for (unsigned j = lane_idx; j < 2048; j += 32) + buffers.output[(uint64_t)row * 2048 + j] = + j < count ? buffers.indices[(uint64_t)row * FallbackRetain + j] + : -1; + cutlass::arch::fence_view_async_shared(); + smem.empty_q_barriers[q_stage_idx].arrive(); + } + + cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +} // namespace seed_fallback diff --git a/csrc/fused_mqa_topk/fallback_shared.cuh b/csrc/fused_mqa_topk/fallback_shared.cuh new file mode 100644 index 000000000000..83a1e1674721 --- /dev/null +++ b/csrc/fused_mqa_topk/fallback_shared.cuh @@ -0,0 +1,268 @@ +#pragma once +#include +#include +#include +#include +#include +constexpr unsigned FallbackSegment = 1024, FallbackRetain = 4096; +struct FallbackBuffers { + float* ring; + float* retained; + int* indices; + int* output; +}; +struct alignas(16) FallbackShared { + cutlass::arch::ClusterBarrier full[2], empty[2]; + float ring[2][8][FallbackSegment]; + unsigned hist[8][1056], groups[8][32], gates[8], kept[8]; + unsigned work[264], prefix, rank, written, added, emit_counts[8]; + __device__ void init(unsigned tid) { + if (tid < 2) { + full[tid].init(256); + empty[tid].init(256); + } + for (unsigned j = tid; j < 8 * 1056; j += 384) + reinterpret_cast(hist)[j] = 0; + if (tid < 256) reinterpret_cast(groups)[tid] = 0; + if (tid < 8) { + gates[tid] = __float_as_uint(-CUDART_INF_F); + kept[tid] = 0; + } + cutlass::arch::fence_barrier_init(); + } +}; +__device__ __forceinline__ unsigned fb_code(float x) { + unsigned u = __float_as_uint(x); + return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); +} +__device__ __forceinline__ float fb_float(unsigned u) { + return __uint_as_float((u & 0x80000000u) ? (u ^ 0x80000000u) : ~u); +} +__device__ __forceinline__ unsigned fb_half_bin(float x) { + unsigned u = __half_as_ushort(__float2half_rd(x)); + return ((u & 0x8000u) ? ((~u) & 65535u) : (u ^ 0x8000u)) >> 6; +} +__device__ __forceinline__ float fb_bin_floor(unsigned b) { + unsigned o = b << 6, u = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); + return __half2float(__ushort_as_half(u)); +} +__device__ __forceinline__ void fb_sync() { + asm volatile("" ::: "memory"); + cutlass::arch::NamedBarrier(256, 2).sync(); + asm volatile("" ::: "memory"); +} +__device__ void fb_compact(FallbackShared& s, FallbackBuffers b, unsigned local, + unsigned row, unsigned tid, bool final = false) { + constexpr unsigned K = 2048, Full = 0xffffffffu, T = 256, + Items = FallbackRetain / T; + unsigned lane = tid & 31, count = s.kept[local]; + uint64_t base = (uint64_t)row * FallbackRetain; + float values[Items]; + int indices[Items]; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) { + unsigned j = tid + i * T; + values[i] = j < count ? b.retained[base + j] : -CUDART_INF_F; + indices[i] = j < count ? b.indices[base + j] : -1; + } + auto emit = [&](float gate, unsigned threshold, int mode, + unsigned base_offset) { + unsigned total = 0; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) { + float v = values[i]; + bool pass = + isfinite(v) && (mode == 0 ? v >= gate + : (mode == 1 ? fb_code(v) > threshold + : fb_code(v) == threshold)); + total += pass; + } + unsigned warp_total = __reduce_add_sync(Full, total); + if (lane == 0) s.emit_counts[tid / 32] = warp_total; + fb_sync(); + unsigned start = base_offset; + for (unsigned w = 0; w < tid / 32; ++w) start += s.emit_counts[w]; + unsigned all = 0; + for (unsigned w = 0; w < 8; ++w) all += s.emit_counts[w]; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) { + float v = values[i]; + bool pass = + isfinite(v) && (mode == 0 ? v >= gate + : (mode == 1 ? fb_code(v) > threshold + : fb_code(v) == threshold)); + unsigned mask = __ballot_sync(Full, pass), + pos = start + __popc(mask & ((1u << lane) - 1)); + if (pass && (mode != 2 || pos < K)) { + b.retained[base + pos] = v; + b.indices[base + pos] = indices[i]; + } + start += __popc(mask); + } + fb_sync(); + return all; + }; + if (!final) { + if (tid == 0) { + s.rank = 0; + s.written = 0; + } + fb_sync(); + float gate = __uint_as_float(s.gates[local]); + unsigned kept = 0; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) + kept += isfinite(values[i]) && values[i] >= gate; + kept = __reduce_add_sync(Full, kept); + if (lane == 0) atomicAdd(&s.rank, kept); + fb_sync(); + kept = s.rank; + if (kept >= K && kept <= FallbackRetain - FallbackSegment) { + emit(gate, 0, 0, 0); + fb_sync(); + if (tid == 0) s.kept[local] = kept; + fb_sync(); + return; + } + } + if (tid == 0) { + s.prefix = 0; + s.rank = K; + s.written = 0; + } + fb_sync(); + for (int shift = 24; shift >= 0; shift -= 8) { + for (unsigned j = tid; j < 264; j += T) s.work[j] = 0; + fb_sync(); + unsigned prefix = s.prefix; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) { + unsigned code = fb_code(values[i]); + bool valid = isfinite(values[i]) && + (shift == 24 || (code >> (shift + 8)) == prefix); + unsigned mask = __ballot_sync(Full, valid); + if (valid) { + unsigned bin = (code >> shift) & 255u, + peers = __match_any_sync(mask, bin); + if (lane == unsigned(__ffs(peers) - 1)) + atomicAdd(s.work + bin + bin / 32, __popc(peers)); + } + } + fb_sync(); + if (tid < 32) { + unsigned sum = 0; + for (unsigned j = 0; j < 8; ++j) { + unsigned bin = 255 - lane * 8 - j; + sum += s.work[bin + bin / 32]; + } + unsigned inclusive = sum; + for (unsigned d = 1; d < 32; d *= 2) { + unsigned v = __shfl_up_sync(Full, inclusive, d); + if (lane >= d) inclusive += v; + } + unsigned before = inclusive - sum, rank = s.rank, + crossing = + __ballot_sync(Full, before < rank && inclusive >= rank), + chosen_lane = __ffs(crossing) - 1, chosen = 0, above = before; + if (lane == chosen_lane) { + for (unsigned j = 0; j < 8; ++j) { + unsigned bin = 255 - lane * 8 - j; + if (above + s.work[bin + bin / 32] >= rank) { + chosen = bin; + break; + } + above += s.work[bin + bin / 32]; + } + } + chosen = __shfl_sync(Full, chosen, chosen_lane); + above = __shfl_sync(Full, above, chosen_lane); + if (lane == 0) { + s.prefix = (prefix << 8) | chosen; + s.rank = rank - above; + } + } + fb_sync(); + } + unsigned threshold = s.prefix; + unsigned greater = emit(0, threshold, 1, 0); + emit(0, threshold, 2, greater); + if (tid == 0) { + s.kept[local] = K; + float gate = fmaxf(__uint_as_float(s.gates[local]), fb_float(threshold)); + atomicExch(s.gates + local, __float_as_uint(gate)); + } + fb_sync(); +} +__device__ void fb_select(FallbackShared& s, FallbackBuffers b, unsigned local, + unsigned row, unsigned epoch, unsigned first, + unsigned n, unsigned tid, bool final) { + constexpr unsigned K = 2048, Full = 0xffffffffu, T = 32, + Items = FallbackSegment / T; + unsigned lane = tid, bank = epoch & 1; + float values[Items], gate = __uint_as_float(s.gates[local]); + unsigned count = 0; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) { + unsigned j = tid + i * T; + float v = j < n ? s.ring[bank][local][j] : -CUDART_INF_F; + if (v <= gate) v = -CUDART_INF_F; + values[i] = v; + bool valid = isfinite(v); + count += valid; + unsigned mask = __ballot_sync(Full, valid); + if (valid) { + unsigned bin = fb_half_bin(v), peers = __match_any_sync(mask, bin); + if (lane == unsigned(__ffs(peers) - 1)) + atomicAdd(&s.hist[local][bin + bin / 32], __popc(peers)); + } + } + count = __reduce_add_sync(Full, count); + __syncwarp(); + if (count) { + unsigned sum = 0; +#pragma unroll + for (unsigned j = 0; j < 32; ++j) + sum += s.hist[local][(31 - lane) * 33 + j]; + unsigned inclusive = sum; + for (unsigned d = 1; d < 32; d *= 2) { + unsigned v = __shfl_up_sync(Full, inclusive, d); + if (lane >= d) inclusive += v; + } + unsigned before = inclusive - sum, + mask = __ballot_sync(Full, before < K && inclusive >= K); + if (mask) { + unsigned chosen = __ffs(mask) - 1, group = 31 - chosen, + rank = K - __shfl_sync(Full, before, chosen), + bin = group * 32 + 31 - lane; + sum = s.hist[local][bin + bin / 32]; + inclusive = sum; + for (unsigned d = 1; d < 32; d *= 2) { + unsigned v = __shfl_up_sync(Full, inclusive, d); + if (lane >= d) inclusive += v; + } + before = inclusive - sum; + chosen = + __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; + bin = __shfl_sync(Full, bin, chosen); + gate = fmaxf(gate, fb_bin_floor(bin)); + if (lane == 0) atomicExch(s.gates + local, __float_as_uint(gate)); + } + unsigned written = s.kept[local]; + uint64_t base = (uint64_t)row * FallbackRetain; +#pragma unroll + for (unsigned i = 0; i < Items; ++i) { + float v = values[i]; + bool pass = isfinite(v) && v >= gate; + unsigned mask = __ballot_sync(Full, pass), + pos = written + __popc(mask & ((1u << lane) - 1)); + if (pass) { + b.retained[base + pos] = v; + b.indices[base + pos] = first + tid + i * T; + } + written += __popc(mask); + } + __syncwarp(); + if (lane == 0) s.kept[local] = written; + __syncwarp(); + } +} diff --git a/csrc/fused_mqa_topk/fast_select.cuh b/csrc/fused_mqa_topk/fast_select.cuh new file mode 100644 index 000000000000..f0a93a238e40 --- /dev/null +++ b/csrc/fused_mqa_topk/fast_select.cuh @@ -0,0 +1,683 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include +#include "candidates.cuh" + +// One warp selects one row's top-K from its candidate fp24 codes (higher code = +// higher score), staged once in shared memory. Candidate codes 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 U 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. +namespace fast_select { +constexpr unsigned K = 2048, Full = 0xffffffffu, kBins = 1024, U = 8; + +// Bin (descending) where the cumulative count first reaches `rank`; +// `above` = elements in higher bins. Lane l owns bins kBins-1-32l .. -31. +__device__ __forceinline__ unsigned find_bin(const unsigned* hist, + unsigned rank, unsigned& above) { + const unsigned lane = threadIdx.x & 31; + constexpr unsigned kPerLane = kBins / 32; + unsigned sum = 0; +#pragma unroll + for (unsigned j = 0; j < kPerLane; ++j) + sum += hist[kBins - 1 - lane * kPerLane - j]; + unsigned inclusive = sum; +#pragma unroll + for (unsigned d = 1; d < 32; d *= 2) { + const unsigned v = __shfl_up_sync(Full, inclusive, d); + if (lane >= d) inclusive += v; + } + const unsigned before = inclusive - sum; + const unsigned winner = + __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; + unsigned chosen = 0, acc = before; + if (lane == winner) { + for (unsigned j = 0; j < kPerLane; ++j) { + const unsigned bin = kBins - 1 - lane * kPerLane - j, h = hist[bin]; + if (acc + h >= rank) { + chosen = bin; + break; + } + acc += h; + } + } + above = __shfl_sync(Full, acc, winner); + return __shfl_sync(Full, chosen, winner); +} + +// Requires count >= K (callers recover rows with fewer candidates). +// stamps (optional): cycles for staging / histogram levels / emit. +__device__ void select_row(unsigned* codes, unsigned count, + const uint16_t* values, const int32_t* packed, + int32_t* out, unsigned* hist, + long long* stamps = nullptr) { + const unsigned lane = threadIdx.x & 31; + const long long t0 = clock64(); + unsigned lo = 0xffffffffu, hi = 0; + unsigned j = lane; + for (; j + 32 * (U - 1) < count; j += 32 * U) { + uint16_t v[U]; + int32_t p[U]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + v[u] = values[j + 32 * u]; + p[u] = packed[j + 32 * u]; + } +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + const unsigned c = + fused_candidates::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 unsigned c = + fused_candidates::candidate_load_score_code(values[j], packed[j]); + codes[j] = c; + lo = min(lo, c); + hi = max(hi, c); + } + lo = __reduce_min_sync(Full, lo); + hi = __reduce_max_sync(Full, hi); + __syncwarp(); + const long long t1 = clock64(); + + // Current bucket: codes in [lo, lo + width). rank = elements still needed + // from it. + unsigned width = hi - lo + 1, rank = K; +#pragma unroll 1 + while (width > 1) { + const unsigned spread = (width - 1) >> 10; + const unsigned shift = spread ? 32 - __clz(spread) : 0; + for (unsigned b = lane; b < kBins; b += 32) hist[b] = 0; + __syncwarp(); + const unsigned top = lo + width - 1; + unsigned i = lane; + for (; i + 32 * (U - 1) < count; i += 32 * U) { + unsigned c[U]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) c[u] = codes[i + 32 * u]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) + if (c[u] >= lo && c[u] <= top) + atomicAdd(hist + ((c[u] - lo) >> shift), 1u); + } + for (; i < count; i += 32) { + const unsigned c = codes[i]; + if (c >= lo && c <= top) atomicAdd(hist + ((c - lo) >> shift), 1u); + } + __syncwarp(); + unsigned above; + const unsigned chosen = find_bin(hist, rank, above); + rank -= above; + lo += chosen << shift; + width = min(1u << shift, top - lo + 1); + __syncwarp(); + } + const long long t2 = clock64(); + + // lo is now the K-th largest code: emit codes above it, then ties until K. + const unsigned threshold = lo, lt = (1u << lane) - 1u; + unsigned pos = 0; +#pragma unroll 1 + for (unsigned tie = 0; tie < 2; ++tie) { + for (unsigned j0 = 0; j0 < count && pos < K; j0 += 32 * U) { + bool pass[U]; + int32_t p[U]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + const unsigned e = j0 + 32 * u + lane; + const unsigned 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 (unsigned u = 0; u < U; ++u) { + const unsigned mask = __ballot_sync(Full, pass[u]), + at = pos + __popc(mask & lt); + if (pass[u] && at < K) + out[at] = fused_candidates::candidate_decode_index(p[u]); + pos += __popc(mask); + } + } + } + if (stamps && lane == 0) { + stamps[0] = t1 - t0; + stamps[1] = t2 - t1; + stamps[2] = clock64() - t2; + } +} +// Resident rows: fp24 codes live in shared memory as hi16[e] << 8 | lo8[e] for +// the row's visible keys e = 0..count-1 (key index s0 + e). Requires count >= +// K. Same levels as select_row. +__device__ __forceinline__ unsigned resident_code(const uint16_t* hi, + const uint8_t* lo, + unsigned e) { + return unsigned(hi[e]) << 8 | lo[e]; +} +struct IdentityKey { + __device__ __forceinline__ int32_t operator()(unsigned e) const { + return int32_t(e); + } +}; +template +__device__ void select_resident(const uint16_t* hi, const uint8_t* lo, + unsigned count, unsigned s0, int32_t* out, + unsigned* hist, KeyOf key_of = {}) { + const unsigned lane = threadIdx.x & 31; + unsigned lo_code = 0xffffffffu, hi_code = 0; + for (unsigned j = lane; j < count; j += 32) { + const unsigned c = resident_code(hi, lo, j); + lo_code = min(lo_code, c); + hi_code = max(hi_code, c); + } + unsigned base = __reduce_min_sync(Full, lo_code); + const unsigned top0 = __reduce_max_sync(Full, hi_code); + unsigned width = top0 - base + 1, rank = K; +#pragma unroll 1 + while (width > 1) { + const unsigned spread = (width - 1) >> 10; + const unsigned shift = spread ? 32 - __clz(spread) : 0; + for (unsigned b = lane; b < kBins; b += 32) hist[b] = 0; + __syncwarp(); + const unsigned top = base + width - 1; + unsigned i = lane; + for (; i + 32 * (U - 1) < count; i += 32 * U) { + unsigned c[U]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) c[u] = resident_code(hi, lo, i + 32 * u); +#pragma unroll + for (unsigned u = 0; u < U; ++u) + if (c[u] >= base && c[u] <= top) + atomicAdd(hist + ((c[u] - base) >> shift), 1u); + } + for (; i < count; i += 32) { + const unsigned c = resident_code(hi, lo, i); + if (c >= base && c <= top) atomicAdd(hist + ((c - base) >> shift), 1u); + } + __syncwarp(); + unsigned above; + const unsigned chosen = find_bin(hist, rank, above); + rank -= above; + base += chosen << shift; + width = min(1u << shift, top - base + 1); + __syncwarp(); + } + const unsigned threshold = base, lt = (1u << lane) - 1u; + unsigned pos = 0; +#pragma unroll 1 + for (unsigned tie = 0; tie < 2; ++tie) { + for (unsigned j0 = 0; j0 < count && pos < K; j0 += 32 * U) { + bool pass[U]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + const unsigned e = j0 + 32 * u + lane; + const unsigned c = e < count ? resident_code(hi, lo, e) : 0; + pass[u] = e < count && (tie ? c == threshold : c > threshold); + } +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + const unsigned mask = __ballot_sync(Full, pass[u]), + at = pos + __popc(mask & lt); + if (pass[u] && at < K) out[at] = key_of(s0 + j0 + 32 * u + lane); + pos += __popc(mask); + } + } + } +} +// Sample-bracketed select for a resident row (one warp, no shared atomics). 512 +// strided samples bracket the K-th largest code in [t_lo, t_hi] (+-40 sample +// ranks, ~5 sigma). One pass emits codes above t_hi straight to `out` and +// compacts the bracket into `mid` (code, offset); the exact threshold is then +// resolved inside the bracket. Returns false (caller falls back) if the bracket +// misses rank K or overflows kMidCap. +constexpr unsigned kSamplesPerLane = 16, kMidCap = 1024, + kMidBytes = kMidCap * 6; + +// Largest t with #(codes >= t) >= rank over a warp's per-lane register samples. +template +__device__ __forceinline__ unsigned sample_kth(const unsigned (&v)[N], + unsigned rank) { + unsigned t = 0; +#pragma unroll 1 + for (int bit = 23; bit >= 0; --bit) { + const unsigned cand = t | (1u << bit); + unsigned c = 0; +#pragma unroll + for (unsigned u = 0; u < N; ++u) c += v[u] >= cand; + if (__reduce_add_sync(Full, c) >= rank) t = cand; + } + return t; +} + +__device__ bool select_resident_bracket(const uint16_t* hi, const uint8_t* lo, + unsigned count, unsigned s0, + int32_t* out, unsigned char* scratch) { + const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; + constexpr unsigned S = 32 * kSamplesPerLane, kMidRegs = kMidCap / 32; + unsigned t_hi, t_lo; + { + unsigned smp[kSamplesPerLane]; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) { + const unsigned e = + unsigned((uint64_t(u * 32 + lane) * 2 + 1) * count / (2 * S)); + smp[u] = resident_code(hi, lo, e); + } + // Expected sample rank of the K-th largest code, bracketed by +-40. + const unsigned expect = unsigned((uint64_t(K) * S + count / 2) / count); + t_hi = sample_kth(smp, expect > 40 ? expect - 40 : 1); + t_lo = sample_kth(smp, min(expect + 40, S)); + } + unsigned* mid_code = reinterpret_cast(scratch); + uint16_t* mid_off = reinterpret_cast(scratch + kMidCap * 4); + unsigned above = 0, mid = 0; + // Codes are loaded U at a time before any compaction store (the stores may + // alias the loads). +#pragma unroll 1 + for (unsigned j0 = 0; j0 < count; j0 += 32 * U) { + unsigned c[U]; +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + const unsigned e = j0 + 32 * u + lane; + c[u] = e < count ? resident_code(hi, lo, e) : 0; + } +#pragma unroll + for (unsigned u = 0; u < U; ++u) { + const unsigned e = j0 + 32 * u + lane; + const bool is_above = e < count && c[u] > t_hi, + is_mid = e < count && c[u] >= t_lo && c[u] <= t_hi; + const unsigned am = __ballot_sync(Full, is_above), + mm = __ballot_sync(Full, is_mid); + if (is_above) { + const unsigned at = above + __popc(am & lt); + if (at < K) out[at] = int32_t(s0 + e); + } + if (is_mid) { + const unsigned at = mid + __popc(mm & lt); + if (at < kMidCap) { + mid_code[at] = c[u]; + mid_off[at] = uint16_t(e); + } + } + above += __popc(am); + mid += __popc(mm); + } + } + if (above >= K || above + mid < K || mid > kMidCap) return false; + __syncwarp(); + // Bracket codes in registers (0 = empty: every real code in [t_lo, t_hi] is + // >= t_lo > 0 unless t_lo == 0, where empty slots only add to rank counts at + // cand == 0, never tested). + unsigned m[kMidRegs]; +#pragma unroll + for (unsigned u = 0; u < kMidRegs; ++u) { + const unsigned i = u * 32 + lane; + m[u] = i < mid ? mid_code[i] : 0; + } + // Exact K-th code: rank r among the bracket, searching only bits where t_lo + // and t_hi differ. + const unsigned r = K - above, diff = t_lo ^ t_hi; + const int top_bit = diff ? 31 - __clz(diff) : -1; + unsigned t = top_bit >= 0 ? t_lo & ~((2u << top_bit) - 1u) : t_lo; +#pragma unroll 1 + for (int bit = top_bit; bit >= 0; --bit) { + const unsigned cand = t | (1u << bit); + unsigned n = 0; +#pragma unroll + for (unsigned u = 0; u < kMidRegs; ++u) n += m[u] >= cand; + if (__reduce_add_sync(Full, n) >= r) t = cand; + } + unsigned pos = above; +#pragma unroll 1 + for (unsigned tie = 0; tie < 2; ++tie) { +#pragma unroll + for (unsigned u = 0; u < kMidRegs; ++u) { + const unsigned i = u * 32 + lane; + const bool pass = i < mid && (tie ? m[u] == t : m[u] > t); + const unsigned mask = __ballot_sync(Full, pass), + at = pos + __popc(mask & lt); + if (pass && at < K) out[at] = int32_t(s0 + mid_off[i]); + pos += __popc(mask); + } + } + return true; +} +// v3 bracket select (one warp per resident row, ~3.5x fewer instructions than +// v2): thresholds on the hi16 part only, 8 elements per 16-byte shared load, +// per-lane counts with one warp prefix scan (no per-element ballots), output +// staged in shared memory and copied out with 16-byte stores. hi_row/lo_row are +// indexed by key index. +constexpr unsigned kMid3Cap = 1024, kScratch3Bytes = K * 4 + kMid3Cap * 2; + +__device__ __forceinline__ unsigned warp_exclusive_scan(unsigned v, + unsigned& total) { + const unsigned lane = threadIdx.x & 31; + unsigned inc = v; +#pragma unroll + for (unsigned d = 1; d < 32; d *= 2) { + const unsigned t = __shfl_up_sync(Full, inc, d); + if (lane >= d) inc += t; + } + total = __shfl_sync(Full, inc, 31); + return inc - v; +} + +template +__device__ __forceinline__ void classify8(const uint4 v, unsigned ebase, + unsigned s0, unsigned end, + unsigned h_hi, unsigned h_lo, + Fn&& fn) { + const unsigned w[4] = {v.x, v.y, v.z, v.w}; + const bool interior = ebase >= s0 && ebase + 8 <= end; +#pragma unroll + for (unsigned k = 0; k < 8; ++k) { + const unsigned h = (w[k / 2] >> (16 * (k & 1))) & 0xffffu, e = ebase + k; + const bool valid = interior || (e >= s0 && e < end); + fn(e, valid && h > h_hi, valid && h >= h_lo && h <= h_hi); + } +} + +template +__device__ bool select_resident_v3(const uint16_t* hi_row, + const uint8_t* lo_row, unsigned s0, + unsigned count, int32_t* out, + unsigned char* scratch, KeyOf key_of = {}) { + const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; + // hi16 bracket from 512 strided samples (+-40 sample ranks around the + // expected K-th). + unsigned h_hi = 0, h_lo = 0; + { + unsigned smp[kSamplesPerLane]; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) + smp[u] = hi_row[s0 + (((u * 32 + lane) * 2 + 1) * count >> 10)]; + const unsigned expect = (K * 512 + count / 2) / count; + const unsigned r_hi = expect > 40 ? expect - 40 : 1, + r_lo = min(expect + 40, 512u); +#pragma unroll 1 + for (int bit = 15; bit >= 0; --bit) { + const unsigned ch = h_hi | (1u << bit), cl = h_lo | (1u << bit); + unsigned nh = 0, nl = 0; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) { + nh += smp[u] >= ch; + nl += smp[u] >= cl; + } + if (__reduce_add_sync(Full, nh) >= r_hi) h_hi = ch; + if (__reduce_add_sync(Full, nl) >= r_lo) h_lo = cl; + } + } + const unsigned a0 = s0 & ~7u, end = s0 + count, units = (end - a0 + 7) / 8; + const uint4* hv = reinterpret_cast(hi_row + a0); + unsigned na = 0, nm = 0; +#pragma unroll 2 + for (unsigned unit = lane; unit < units; unit += 32) + classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, + [&](unsigned, bool a, bool m) { + na += a; + nm += m; + }); + unsigned above, mid; + unsigned pa = warp_exclusive_scan(na, above), + pm = warp_exclusive_scan(nm, mid); + if (above >= K || above + mid < K || mid > kMid3Cap) return false; + int32_t* stage = reinterpret_cast(scratch); + uint16_t* midbuf = reinterpret_cast(scratch + K * 4); +#pragma unroll 2 + for (unsigned unit = lane; unit < units; unit += 32) + classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, + [&](unsigned e, bool a, bool m) { + if (a) stage[pa++] = key_of(e); + if (m) midbuf[pm++] = uint16_t(e); + }); + __syncwarp(); + // Exact K-th code among the bracket (rank r), over the bits where the bracket + // bounds differ. + constexpr unsigned R = kMid3Cap / 32; + unsigned mc[R]; +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + const unsigned i = u * 32 + lane; + mc[u] = 0; + if (u * 32 < mid && i < mid) { + const unsigned e = midbuf[i]; + mc[u] = unsigned(hi_row[e]) << 8 | lo_row[e]; + } + } + const unsigned r = K - above, c_lo = h_lo << 8, c_hi = h_hi << 8 | 0xffu, + diff = c_lo ^ c_hi; + const int top_bit = diff ? 31 - __clz(diff) : -1; + unsigned t = top_bit >= 0 ? c_lo & ~((2u << top_bit) - 1u) : c_lo; +#pragma unroll 1 + for (int bit = top_bit; bit >= 0; --bit) { + const unsigned cand = t | (1u << bit); + unsigned n = 0; +#pragma unroll + for (unsigned u = 0; u < R; ++u) + if (u * 32 < mid) n += mc[u] >= cand; + if (__reduce_add_sync(Full, n) >= r) t = cand; + } + unsigned pos = above; +#pragma unroll 1 + for (unsigned tie = 0; tie < 2; ++tie) { +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + if (u * 32 >= mid || pos >= K) break; + const unsigned i = u * 32 + lane; + const bool pass = i < mid && (tie ? mc[u] == t : mc[u] > t); + const unsigned mask = __ballot_sync(Full, pass), + at = pos + __popc(mask & lt); + if (pass && at < K) stage[at] = key_of(midbuf[i]); + pos += __popc(mask); + } + } + __syncwarp(); +#pragma unroll 4 + for (unsigned j = lane * 4; j < K; j += 128) + *reinterpret_cast(out + j) = + *reinterpret_cast(stage + j); + return true; +} +// Histogram select over any code source: code_of(e) / key_of(e) for e in [0, +// count). Used for rows whose shared candidates overflowed into global memory +// (rare). +template +__device__ void select_hist_generic(CodeOf code_of, unsigned count, + int32_t* out, unsigned* hist, + KeyOf key_of) { + const unsigned lane = threadIdx.x & 31; + unsigned lo_code = 0xffffffffu, hi_code = 0; + for (unsigned j = lane; j < count; j += 32) { + const unsigned c = code_of(j); + lo_code = min(lo_code, c); + hi_code = max(hi_code, c); + } + unsigned base = __reduce_min_sync(Full, lo_code); + const unsigned top0 = __reduce_max_sync(Full, hi_code); + unsigned width = top0 - base + 1, rank = K; +#pragma unroll 1 + while (width > 1) { + const unsigned spread = (width - 1) >> 10; + const unsigned shift = spread ? 32 - __clz(spread) : 0; + for (unsigned b = lane; b < kBins; b += 32) hist[b] = 0; + __syncwarp(); + const unsigned top = base + width - 1; + for (unsigned i = lane; i < count; i += 32) { + const unsigned c = code_of(i); + if (c >= base && c <= top) atomicAdd(hist + ((c - base) >> shift), 1u); + } + __syncwarp(); + unsigned above; + const unsigned chosen = find_bin(hist, rank, above); + rank -= above; + base += chosen << shift; + width = min(1u << shift, top - base + 1); + __syncwarp(); + } + const unsigned threshold = base, lt = (1u << lane) - 1u; + unsigned pos = 0; +#pragma unroll 1 + for (unsigned tie = 0; tie < 2; ++tie) { + for (unsigned j0 = 0; j0 < count && pos < K; j0 += 32) { + const unsigned e = j0 + lane; + const unsigned c = e < count ? code_of(e) : 0; + const bool pass = e < count && (tie ? c == threshold : c > threshold); + const unsigned mask = __ballot_sync(Full, pass), + at = pos + __popc(mask & lt); + if (pass && at < K) out[at] = key_of(e); + pos += __popc(mask); + } + } +} +// v4: v3 plus (a) elements whose hi16 is 0 are empty slots (sealed; no real +// score has hi16 == 0) and are excluded, with bracket ranks taken over the +// `real` elements, and (b) the exact threshold inside the bracket is resolved +// with a 512-bin shared histogram (levels until one code wide) instead of a +// 24-step register bit search. +constexpr unsigned kMidBins = 512, + kScratch4Bytes = K * 4 + kMid3Cap * 2 + kMidBins * 4; + +template +__device__ __forceinline__ unsigned find_bin_n(const unsigned* hist, + unsigned rank, unsigned& above) { + const unsigned lane = threadIdx.x & 31; + constexpr unsigned kPerLane = Bins / 32; + unsigned sum = 0; +#pragma unroll + for (unsigned j = 0; j < kPerLane; ++j) + sum += hist[Bins - 1 - lane * kPerLane - j]; + unsigned inclusive = sum; +#pragma unroll + for (unsigned d = 1; d < 32; d *= 2) { + const unsigned v = __shfl_up_sync(Full, inclusive, d); + if (lane >= d) inclusive += v; + } + const unsigned before = inclusive - sum; + const unsigned winner = + __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; + unsigned chosen = 0, acc = before; + if (lane == winner) { + for (unsigned j = 0; j < kPerLane; ++j) { + const unsigned bin = Bins - 1 - lane * kPerLane - j, h = hist[bin]; + if (acc + h >= rank) { + chosen = bin; + break; + } + acc += h; + } + } + above = __shfl_sync(Full, acc, winner); + return __shfl_sync(Full, chosen, winner); +} + +template +__device__ bool select_resident_v4(const uint16_t* hi_row, + const uint8_t* lo_row, unsigned s0, + unsigned count, unsigned real, int32_t* out, + unsigned char* scratch, KeyOf key_of = {}) { + const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; + unsigned h_hi = 0, h_lo = 0; + { + unsigned smp[kSamplesPerLane], n_real = 0; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) { + smp[u] = hi_row[s0 + (((u * 32 + lane) * 2 + 1) * count >> 10)]; + n_real += smp[u] != 0; + } + const unsigned s_real = __reduce_add_sync(Full, n_real); + if (s_real == 0 || real <= K) return false; + // Expected sample rank of the K-th largest real element, bracketed by +-40. + const unsigned expect = (K * s_real + real / 2) / real; + const unsigned r_hi = expect > 40 ? expect - 40 : 1, + r_lo = min(expect + 40, s_real); +#pragma unroll 1 + for (int bit = 15; bit >= 0; --bit) { + const unsigned ch = h_hi | (1u << bit), cl = h_lo | (1u << bit); + unsigned nh = 0, nl = 0; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) { + nh += smp[u] >= ch; + nl += smp[u] >= cl; + } + if (__reduce_add_sync(Full, nh) >= r_hi) h_hi = ch; + if (__reduce_add_sync(Full, nl) >= r_lo) h_lo = cl; + } + h_lo = max(h_lo, 1u); + } + const unsigned a0 = s0 & ~7u, end = s0 + count, units = (end - a0 + 7) / 8; + const uint4* hv = reinterpret_cast(hi_row + a0); + unsigned na = 0, nm = 0; +#pragma unroll 2 + for (unsigned unit = lane; unit < units; unit += 32) + classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, + [&](unsigned, bool a, bool m) { + na += a; + nm += m; + }); + unsigned above, mid; + unsigned pa = warp_exclusive_scan(na, above), + pm = warp_exclusive_scan(nm, mid); + if (above >= K || above + mid < K || mid > kMid3Cap) return false; + int32_t* stage = reinterpret_cast(scratch); + uint16_t* midbuf = reinterpret_cast(scratch + K * 4); + unsigned* hist = reinterpret_cast(scratch + K * 4 + kMid3Cap * 2); +#pragma unroll 2 + for (unsigned unit = lane; unit < units; unit += 32) + classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, + [&](unsigned e, bool a, bool m) { + if (a) stage[pa++] = key_of(e); + if (m) midbuf[pm++] = uint16_t(e); + }); + __syncwarp(); + const auto code_at = [&](unsigned i) { + const unsigned e = midbuf[i]; + return unsigned(hi_row[e]) << 8 | lo_row[e]; + }; + // K-th code among the bracket (rank r): histogram levels over [h_lo << 8, + // h_hi << 8 | 0xff]. + unsigned rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; +#pragma unroll 1 + while (top > base) { + const unsigned width = top - base + 1, spread = (width - 1) / kMidBins; + const unsigned shift = spread ? 32 - __clz(spread) : 0; + for (unsigned b = lane; b < kMidBins; b += 32) hist[b] = 0; + __syncwarp(); + for (unsigned i = lane; i < mid; i += 32) { + const unsigned c = code_at(i); + if (c >= base && c <= top) atomicAdd(hist + ((c - base) >> shift), 1u); + } + __syncwarp(); + unsigned above_bin; + const unsigned chosen = find_bin_n(hist, rank, above_bin); + rank -= above_bin; + base += chosen << shift; + top = min(base + (1u << shift) - 1, top); + __syncwarp(); + } + const unsigned t = base; + unsigned pos = above; +#pragma unroll 1 + for (unsigned tie = 0; tie < 2; ++tie) { + for (unsigned j = 0; j < mid && pos < K; j += 32) { + const unsigned i = j + lane; + const unsigned c = i < mid ? code_at(i) : 0; + const bool pass = i < mid && (tie ? c == t : c > t); + const unsigned mask = __ballot_sync(Full, pass), + at = pos + __popc(mask & lt); + if (pass && at < K) stage[at] = key_of(midbuf[i]); + pos += __popc(mask); + } + } + __syncwarp(); +#pragma unroll 4 + for (unsigned j = lane * 4; j < K; j += 128) + *reinterpret_cast(out + j) = + *reinterpret_cast(stage + j); + return true; +} +} // namespace fast_select 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..f6ce592e9fe8 --- /dev/null +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Single-launch fused FP8 MQA-logits + top-2048 prefill selection (SM100 +// family). One kernel computes the indexer scores of every visible key of a +// row and selects its top-2048 keys without materializing the [rows, keys] +// logits matrix. The device code lives in the headers of this directory (kept +// verbatim from the isolated kernel; see mqa_seeded_kernel.cuh); 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 "sample_core.cuh" +#include "scan.cuh" +#include "fallback_core.cuh" +#include "upstream_select.cuh" +#include "fast_select.cuh" +#include "mqa_seeded_kernel.cuh" + +namespace { + +using torch::headeronly::ScalarType; + +void check_cuda(cudaError_t err) { + STD_TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err)); +} + +void* driver_handle() { + static void* h = nullptr; + if (!h) { + h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + STD_TORCH_CHECK(h, "failed to load libcuda.so.1"); + } + return h; +} + +CUresult enc_tiled(CUtensorMap* tm, CUtensorMapDataType dt, cuuint32_t rank, + void* addr, const cuuint64_t* dims, + const cuuint64_t* strides, const cuuint32_t* box, + const cuuint32_t* estrides, CUtensorMapInterleave il, + CUtensorMapSwizzle sw, CUtensorMapL2promotion l2, + CUtensorMapFloatOOBfill oob) { + using FT = + CUresult (*)(CUtensorMap*, CUtensorMapDataType, cuuint32_t, void*, + const cuuint64_t*, const cuuint64_t*, const cuuint32_t*, + const cuuint32_t*, CUtensorMapInterleave, CUtensorMapSwizzle, + CUtensorMapL2promotion, CUtensorMapFloatOOBfill); + static FT f = nullptr; + if (!f) { + f = reinterpret_cast(dlsym(driver_handle(), "cuTensorMapEncodeTiled")); + STD_TORCH_CHECK(f, "failed to load cuTensorMapEncodeTiled"); + } + return f(tm, dt, rank, addr, dims, strides, box, estrides, il, sw, l2, oob); +} + +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 = enc_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. +CUtensorMap make_strided_kv(void* ptr, unsigned groups, unsigned stride) { + CUtensorMap tm; + const cuuint64_t gdims[3] = {128, groups, stride}; + const cuuint64_t gstrides[2] = {128ull * stride, 128}; + const cuuint32_t box[3] = {128, 256, 1}; + const cuuint32_t estrides[3] = {1, 1, 1}; + CUresult r = enc_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; +} + +constexpr int kHeads = 32; +constexpr int kHeadDim = 128; +constexpr int kTopK = 2048; +constexpr int kCandidates = 16384; // per-row candidate slab +constexpr int kRetained = 4096; // per-row fallback slab +constexpr int kStatsPerCta = 7; + +// Gate target / sample rank of the <= 16K-key path (candidates kept in shared +// memory, kCandCap per row). +constexpr unsigned kSmallTarget = 2560; +constexpr unsigned kSmallSampleRank = 100; + +using SampleStorage = + deep_gemm::layout::MQALogitsSharedStorage<32, 128, false, 8, 256, 1, 2, 2, + cutlass::float_e4m3_t, float>; +using ScanStorage = + deep_gemm::layout::MQALogitsSharedStorage<32, 128, false, 8, 256, 1, 3, 2, + cutlass::float_e4m3_t, float>; +// Sample / fallback phases. +constexpr int kPhaseSmem = + sizeof(SampleStorage) + (sizeof(SampleShared) > sizeof(FallbackShared) + ? sizeof(SampleShared) + : sizeof(FallbackShared)); +// > 16K keys: 3-stage scan + candidate ring. +constexpr int kLongSmem = + std::max(kPhaseSmem, sizeof(ScanStorage) + sizeof(CandidateHandoff<>)); +// <= 16K keys: 2-stage scan + per-row shared candidates (5 bytes each). +constexpr int kSmallSmem = + std::max(kPhaseSmem, sizeof(SampleStorage) + kCandExt * 5); +static_assert(kSmallSmem + 2048 <= 232448, "shared candidates must fit"); + +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"); +} + +} // namespace + +// q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32, +// weights [rows,32] fp32, starts/ends [rows] int32 (visible key range +// [start,end) of each row), out [rows,2048] int32: absolute key indices +// (columns of the logits), or with relative=true indices relative to the +// row's start (key - start, the top_k_per_row_prefill convention); -1 padded +// when fewer than 2048 keys are visible. +// 0 < rows <= 16384, 0 < keys <= 1048576 (any values: the last 8-row block +// may be partial). +// Workspace (caller-owned, reused across calls), with B = ceil(rows / 8): +// slot_flags [S] int32 (all zero between calls; the kernel restores them), +// values [S*8,16384] fp16, indices [S*8,16384] int32, +// retained [S*8,4096] fp32, retained_indices [S*8,4096] int32, +// counts [>=8B] int32, gates [>=8B] fp32, stats [>=7B] int32 (per-CTA +// diagnostics). +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 slot_flags, + torch::stable::Tensor values, torch::stable::Tensor indices, + torch::stable::Tensor retained, torch::stable::Tensor retained_indices, + torch::stable::Tensor counts, torch::stable::Tensor gates, + torch::stable::Tensor stats, bool relative) { + 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(slot_flags, dev, ScalarType::Int, "slot_flags"); + check_tensor(values, dev, ScalarType::Half, "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(gates, dev, ScalarType::Float, "gates"); + check_tensor(stats, dev, ScalarType::Int, "stats"); + + STD_TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && 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); + STD_TORCH_CHECK(rows > 0 && rows <= 16384, "rows must be in [1, 16384]"); + STD_TORCH_CHECK(keys > 0 && keys <= (1 << 20), + "keys must be in [1, 1048576]"); + const int64_t blocks = (rows + 7) / 8; + STD_TORCH_CHECK(scales.numel() == keys, "scales must be [keys]"); + STD_TORCH_CHECK(weights.dim() == 2 && weights.size(0) == rows && + weights.size(1) == kHeads, + "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 && out.size(1) == kTopK, + "out must be [rows, 2048]"); + STD_TORCH_CHECK(reinterpret_cast(out.data_ptr()) % 16 == 0, + "out must be 16-byte aligned"); + + const int64_t total_slots = slot_flags.numel(); + STD_TORCH_CHECK(total_slots > 0, "slot_flags must not be empty"); + STD_TORCH_CHECK(values.dim() == 2 && values.size(0) == total_slots * 8 && + values.size(1) == kCandidates && indices.dim() == 2 && + indices.size(0) == total_slots * 8 && + indices.size(1) == kCandidates, + "values/indices must be [8 * slots, 16384]"); + STD_TORCH_CHECK(retained.dim() == 2 && retained.size(0) == total_slots * 8 && + retained.size(1) == kRetained && + retained_indices.dim() == 2 && + retained_indices.size(0) == total_slots * 8 && + retained_indices.size(1) == kRetained, + "retained/retained_indices must be [8 * slots, 4096]"); + STD_TORCH_CHECK(counts.numel() >= blocks * 8 && gates.numel() >= blocks * 8 && + stats.numel() >= blocks * kStatsPerCta, + "per-row workspace too small"); + + torch::stable::accelerator::DeviceGuard guard(dev); + const cudaStream_t stream = get_current_cuda_stream(dev); + + // Pool of 8-row slots: the kernel claims a free slot per CTA, so the + // candidate workspace is bounded independently of rows. The caller sizes it + // with >= 2 slots per SM (>= 2x resident CTAs, so a free slot always exists). + const unsigned num_slots = + static_cast(std::min(total_slots, blocks)); + + FallbackBuffers fallback{nullptr, static_cast(retained.data_ptr()), + static_cast(retained_indices.data_ptr()), + static_cast(out.data_ptr())}; + const int irows = static_cast(rows), ikeys = static_cast(keys); + auto tq = make_2d(q.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, + irows * 32, 128, 256, 128, 128); + auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, ikeys, + 128, 256, 128, 128); + auto ts = make_2d(scales.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, + ikeys, 1, 256, 1, 0, 0); + auto tw = make_2d(weights.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, 32, + irows, 32, 8, 32, 0); + // Gate sample: <= 16K keys aim for kSmallTarget candidates (shared memory), + // <= 128K for 3072, both with a sample sized so the gate's quantile rank is + // ~100 (about 10% count error). Longer contexts use 16 tiles and the + // 4096/6144 target, where candidates are already ~1% of the keys. + // (<= 4K keys keep every score in shared memory and skip the sample.) + const bool low_target = keys <= 131072; + const unsigned sample_target = + keys <= 16384 ? kSmallTarget : (low_target ? 3072u : 0u); + const unsigned sample_rank = keys <= 16384 ? kSmallSampleRank : 100u; + const unsigned sample_tiles = + low_target ? std::max(2u, (unsigned)((sample_rank * keys + + 256ull * sample_target - 1) / + (256ull * sample_target))) + : 16u; + const unsigned sample_stride = std::max( + 1u, (unsigned)((keys + 256u * sample_tiles - 1) / (256u * sample_tiles))); + const unsigned sample_groups = + (keys - sample_stride / 2 + sample_stride - 1) / sample_stride; + auto tks = make_strided_kv(k.data_ptr(), sample_groups, sample_stride); + + const int smem = keys <= int64_t(kResidentKeys) ? int(kResidentSmem) + : keys <= 16384 ? kSmallSmem + : kLongSmem; + auto kernel = keys <= int64_t(kResidentKeys) ? mqa_seeded_kernel<0> + : keys <= 8192 ? mqa_seeded_kernel<4> + : keys <= 16384 ? mqa_seeded_kernel<8> + : keys <= 32768 ? mqa_seeded_kernel<16> + : keys <= 131072 ? mqa_seeded_kernel<32> + : keys <= 524288 ? mqa_seeded_kernel<64> + : mqa_seeded_kernel<128>; + check_cuda(cudaFuncSetAttribute(reinterpret_cast(kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem)); + kernel<<(blocks), 384, smem, stream>>>( + static_cast(rows), static_cast(keys), + static_cast(starts.data_ptr()), + static_cast(ends.data_ptr()), + static_cast(gates.data_ptr()), + static_cast(values.data_ptr()), + static_cast(indices.data_ptr()), + static_cast(counts.data_ptr()), static_cast(stats.data_ptr()), + fallback, static_cast(slot_flags.data_ptr()), num_slots, + static_cast(kLongSmem), sample_target, sample_stride, + sample_groups, sample_tiles, static_cast(scales.data_ptr()), + tks, tq, ts, tk, tw, relative); + check_cuda(cudaGetLastError()); +} diff --git a/csrc/fused_mqa_topk/handoff.cuh b/csrc/fused_mqa_topk/handoff.cuh new file mode 100644 index 000000000000..1988e28a5f5f --- /dev/null +++ b/csrc/fused_mqa_topk/handoff.cuh @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include +#include + +// Each producer warp owns two banks. A single drainer visits producer warps +// in order, consuming one published batch at a time. Epochs span query blocks. +template +struct alignas(16) CandidateHandoff { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + Barrier full[2][P], empty[2][P]; + uint32_t records[2][P][R][S][32]; + uint32_t counts[2][P][R][32]; + uint32_t bases[2][P]; + __device__ void init(int tid) { + if (tid < 2 * P) { + full[tid / P][tid % P].init(32); + empty[tid / P][tid % P].init(32); + } + cutlass::arch::fence_barrier_init(); + } + __device__ void acquire(int producer, unsigned epoch) { + empty[epoch & 1][producer].wait(((epoch >> 1) & 1) ^ 1); + __syncwarp(); + } + __device__ void publish(int producer, unsigned epoch, unsigned base) { + if ((threadIdx.x & 31) == 0) bases[epoch & 1][producer] = base; + __syncwarp(); + full[epoch & 1][producer].arrive(); + } + template + __device__ void drain(int producer, unsigned epoch, unsigned rows, + unsigned row_base, int* totals, unsigned cap, + Store store) { + const int lane = threadIdx.x & 31, bank = epoch & 1; + full[bank][producer].wait((epoch >> 1) & 1); + __syncwarp(); + const unsigned base = bases[bank][producer]; + // Two 16-bit prefix sums share each shuffle. A row contributes at + // most 32*S records, so the low half cannot carry into the high half. + static_assert(R == 4 && 32 * S < 65536); + unsigned n[4]; + unsigned my_total = 0; +#pragma unroll + for (unsigned row = 0; row < 4; ++row) { + n[row] = counts[bank][producer][row][lane]; + unsigned total = __reduce_add_sync(0xffffffffu, n[row]); + if (lane == row) my_total = total; + } + unsigned my_start = 0; + if (lane < 4 && my_total) + my_start = atomicAdd(totals + row_base + lane, my_total); + unsigned p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); +#pragma unroll + for (int d = 1; d < 32; d *= 2) { + unsigned a = __shfl_up_sync(0xffffffffu, p01, d); + unsigned b = __shfl_up_sync(0xffffffffu, p23, d); + if (lane >= d) { + p01 += a; + p23 += b; + } + } + unsigned prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; +#pragma unroll + for (unsigned row = 0; row < 4; ++row) { + unsigned start = __shfl_sync(0xffffffffu, my_start, row); + for (unsigned j = 0; j < n[row]; ++j) { + unsigned slot = start + prefix[row] - n[row] + j; + unsigned record = records[bank][producer][row][j][lane]; + if (slot < cap) + store(row_base + row, slot, record, base, producer * 32 + lane); + } + } + __syncwarp(); + empty[bank][producer].arrive(); + } +}; diff --git a/csrc/fused_mqa_topk/mqa_seeded_kernel.cuh b/csrc/fused_mqa_topk/mqa_seeded_kernel.cuh new file mode 100644 index 000000000000..36a2165e06d5 --- /dev/null +++ b/csrc/fused_mqa_topk/mqa_seeded_kernel.cuh @@ -0,0 +1,423 @@ +// SPDX-License-Identifier: MIT +// Generated by experiments/integration_v3/sync_candidate.py from +// experiments/fused_clean/mqa_fixed_cuda.cu -- do not edit by hand. +// Device side of the single-launch fused MQA-logits + top-k prefill kernel. +#pragma once +// Sample tile t holds keys (t*256 + j)*S + S/2 (j < 256) through a strided +// {128, G, S} KV view: all tiles together form one uniform grid over every key +// position. +struct SampleScheduler + : deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false> { + using Base = deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false>; + unsigned sample_stride, sample_groups, sample_tiles; + const float* sample_scales; + __device__ SampleScheduler(unsigned sm, unsigned rows, unsigned keys, + const unsigned* starts, const unsigned* ends, + unsigned stride, unsigned groups, unsigned tiles, + const float* scales) + : Base(sm, rows, keys, starts, ends, nullptr), + sample_stride(stride), + sample_groups(groups), + sample_tiles(tiles), + sample_scales(scales) {} + __device__ bool next_q_block(unsigned& q, unsigned& base, unsigned& splits) { + if (!Base::next_q_block(q, base, splits)) return false; + splits = sample_tiles; + return true; + } + __device__ unsigned sample_center() const { return sample_stride / 2; } + __device__ unsigned sample_group(unsigned split, unsigned lane) const { + return split * 256 + lane; + } + __device__ unsigned get_logits_col(unsigned, unsigned split, + unsigned lane) const { + return sample_group(split, lane) * sample_stride + sample_center(); + } +}; +constexpr unsigned kResidentKeys = 4096; +// <=16K keys (handoff interval <= 8): candidates stay in shared memory, cap per +// row. +constexpr unsigned kCandCap = 3072; +// Each candidate array: 8 rows of kCandCap plus 256 per-lane trash slots. +constexpr unsigned kCandExt = 8u * kCandCap + 256u; +constexpr unsigned kResidentSmem = + ((sizeof(deep_gemm::layout::MQALogitsSharedStorage< + 32, 128, false, 8, 256, 1, 2, 2, cutlass::float_e4m3_t, float>) + + 8u * kResidentKeys * 3u) + + 1023u) / + 1024u * 1024u; +template +__global__ __launch_bounds__(384, 1) void mqa_seeded_kernel( + unsigned rows, unsigned keys, const unsigned* starts_g, + const unsigned* ends_g, float* gates, uint16_t* vals, int* indices, + int* counts, int* stats, FallbackBuffers fallback, int* slot_flags, + unsigned num_slots, unsigned phase_bytes, unsigned sample_target, + unsigned sample_stride, unsigned sample_groups, unsigned 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, bool relative) { + // Blocks run heaviest (latest rows, most keys) first so the last wave holds + // the light ones. + const unsigned vblock = gridDim.x - 1 - blockIdx.x; + unsigned row_base = vblock * 8; + bool all_short = true; + // The block's row bounds, clamped to the keys, in shared memory. Rows past + // the end (last block) see no keys; `starts`/`ends` are indexed by absolute + // row like the global arrays. + __shared__ unsigned row_bounds[2][8]; + if (threadIdx.x < 8) { + const unsigned 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 unsigned* const starts = row_bounds[0] - row_base; + const unsigned* const ends = row_bounds[1] - row_base; + // relative: outputs are key - start (top_k_per_row_prefill's convention). + // Rows with + // <= 2048 visible keys write that directly (base 0); selected rows (and every + // row of a recovered block) are rewritten in place by their block once final + // (L2-hot). + const auto finish_rows = [&](bool all_rows) { + if (!relative) return; + __syncthreads(); + for (unsigned i = threadIdx.x; i < 8 * 512; i += 384) { + const unsigned r = row_base + i / 512; + if (r >= rows || (!all_rows && ends[r] - starts[r] <= 2048u)) continue; + const int s = int(starts_g[r]); + if (s == 0) continue; + int4* p = reinterpret_cast(fallback.output + (uint64_t)r * 2048) + + i % 512; + 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; + } + }; + for (unsigned r = row_base; r < row_base + 8; ++r) + all_short &= (ends[r] - starts[r] <= 2048); + if (all_short) { + if (threadIdx.x == 0) { + for (int c = 0; c < 7; ++c) stats[blockIdx.x * 7 + c] = 0; + } + for (unsigned r = row_base; r < min(row_base + 8, rows); ++r) + for (unsigned j = threadIdx.x; j < 2048; j += 384) + fallback.output[(uint64_t)r * 2048 + j] = + j < ends[r] - starts[r] ? int((relative ? 0u : starts[r]) + j) : -1; + return; + } + if constexpr (Interval == 0) { + // Resident (keys <= 4096): scores of all 8 rows stay in shared memory as + // fp24 codes, so there is no sample, gate, candidate workspace or slot + // claim. One warp selects one row. + if (threadIdx.x == 0) { + for (int c = 0; c < 7; ++c) stats[blockIdx.x * 7 + c] = 0; + } + const long long t0 = clock64(); + using RShared = deep_gemm::layout::MQALogitsSharedStorage< + 32, 128, false, 8, 256, 1, 2, 2, cutlass::float_e4m3_t, float>; + extern __shared__ __align__(1024) unsigned char phase_smem[]; + unsigned char* const resident = phase_smem + sizeof(RShared); + auto scheduler = [&](const unsigned& sm) { + return deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false>( + vblock, rows, keys, starts, ends, nullptr); + }; + fused_mqa_scan::sm100_mqa_logits_core_impl< + 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, + cutlass::float_e4m3_t, float, float, decltype(scheduler), 2, 0>( + rows, kResidentKeys, reinterpret_cast(resident), nullptr, + nullptr, nullptr, nullptr, vals, indices, counts, 16384, 0, tq, ts, tk, + ts, tw, scheduler); + __syncthreads(); + const long long t1 = clock64(); + if (threadIdx.x < 256) cutlass::arch::warpgroup_reg_dealloc<128>(); + __syncthreads(); + if (threadIdx.x >= 256) cutlass::arch::warpgroup_reg_alloc<128>(); + const unsigned warp = threadIdx.x / 32; + if (warp < 8 && row_base + warp < rows) { + const unsigned r = row_base + warp, s0 = min(starts[r], keys), + visible = min(ends[r], keys) - s0; + int* o = fallback.output + (uint64_t)r * 2048; + if (visible <= 2048) { + for (unsigned j = threadIdx.x & 31; j < 2048; j += 32) + o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; + } else { + const uint16_t* hi = reinterpret_cast(resident) + + warp * kResidentKeys + s0; + const uint8_t* lo = + resident + 8 * kResidentKeys * 2 + warp * kResidentKeys + s0; + // Scratch (bracket buffer, or the fallback's histogram) reuses the idle + // Q/KV stage memory. + static_assert(8 * fast_select::kScratch3Bytes <= sizeof(RShared) && + fast_select::kScratch3Bytes >= fast_select::kBins * 4, + "select scratch must fit the Q/KV stages"); + unsigned char* scratch = + phase_smem + warp * fast_select::kScratch3Bytes; + if (!fast_select::select_resident_v3(hi - s0, lo - s0, s0, visible, o, + scratch)) { + __syncwarp(); + if ((threadIdx.x & 31) == 0) atomicAdd(stats + blockIdx.x * 7 + 6, 1); + fast_select::select_resident(hi, lo, visible, s0, o, + reinterpret_cast(scratch)); + } + } + } + __syncthreads(); + if (threadIdx.x == 0) { + stats[blockIdx.x * 7 + 4] = int(t1 - t0); + stats[blockIdx.x * 7 + 5] = int(clock64() - t1); + } + finish_rows(false); + return; + } else { + // Workspace pool: claim one of num_slots 8-row slots (num_slots >= 2x + // resident CTAs, so a free slot always exists and holders never wait). Base + // pointers are shifted so global row r of this block addresses slot row + // slot*8 + (r - row_base) in the cores unchanged. + __shared__ unsigned slot; + if (threadIdx.x == 0) { + unsigned s0 = blockIdx.x % num_slots; + for (unsigned i = 0;; ++i) { + const unsigned c = (s0 + i) % num_slots; + if (atomicCAS(slot_flags + c, 0, 1) == 0) { + slot = c; + break; + } + } + } + __syncthreads(); + // Slot rows are addressed as (row + row_delta): reassigning the pointer + // parameters instead costs ~10% (they stop being launch constants in the + // fused scan). + const int64_t row_delta = int64_t(slot) * 8 - int64_t(row_base); + if (threadIdx.x < 8) counts[row_base + threadIdx.x] = 0; + __syncthreads(); + auto sample_scheduler = [&](const unsigned& sm) { + return SampleScheduler(vblock, rows, keys, starts, ends, sample_stride, + sample_groups, sample_tiles, scales); + }; + SampleBuffers sample{gates, nullptr, nullptr, fallback.output, + sample_target}; + seed_sample::sm100_mqa_logits_core_impl< + 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, + cutlass::float_e4m3_t, float, float, decltype(sample_scheduler), 2>( + rows, sample, tq, ts, tks, ts, tw, sample_scheduler); + __syncthreads(); + auto scheduler = [&](const unsigned& sm) { + return deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false>( + vblock, rows, keys, starts, ends, nullptr); + }; + constexpr bool kSmemCand = Interval <= 8; + using CShared = deep_gemm::layout::MQALogitsSharedStorage< + 32, 128, false, 8, 256, 1, 2, 2, cutlass::float_e4m3_t, float>; + extern __shared__ __align__(1024) unsigned char cand_smem[]; + unsigned char* const cand = cand_smem + sizeof(CShared); + if constexpr (kSmemCand) + fused_mqa_scan::sm100_mqa_logits_core_impl< + 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, + cutlass::float_e4m3_t, float, float, decltype(scheduler), 2, Interval, + true>(rows, kCandCap, reinterpret_cast(cand), gates, nullptr, + nullptr, nullptr, vals, indices, counts, 16384, row_delta, tq, + ts, tk, ts, tw, scheduler); + else + fused_mqa_scan::sm100_mqa_logits_core_impl< + 32, 128, false, false, false, 8, 256, 1, 3, 4096, 128, 256, + cutlass::float_e4m3_t, float, float, decltype(scheduler), 2, + Interval>(rows, 0, nullptr, gates, nullptr, nullptr, nullptr, vals, + indices, counts, 16384, row_delta, tq, ts, tk, ts, tw, + scheduler); + __threadfence_block(); + __syncthreads(); + __shared__ unsigned recover, staged, short_rows; + __shared__ unsigned row_offset[9]; + if (threadIdx.x == 0) { + recover = 0; + short_rows = 0; + int lo = 2147483647, hi = 0; + row_offset[0] = 0; + for (unsigned r = row_base; r < row_base + 8; ++r) { + // Rows with <= 2048 visible keys output all of them inline; only long + // rows with a bad count recover. + const bool is_short = + min(ends[r], keys) - min(starts[r], keys) <= 2048u; + short_rows |= unsigned(is_short) << (r - row_base); + const int real = kSmemCand ? (counts[r] & ((1 << 21) - 1)) : counts[r]; + recover |= !is_short && (real < 2048 || real > 16384); + if (!is_short) { + lo = min(lo, counts[r]); + hi = max(hi, counts[r]); + } + row_offset[r - row_base + 1] = + row_offset[r - row_base] + (is_short ? 0 : max(counts[r], 0)); + } + // Fast select stages every row's codes after 8 warp-private 1024-bin + // histograms. staged: 1 = all 8 rows in one round, 2 = two rounds of 4 + // rows (long contexts), 0 = legacy. + const unsigned half = max(row_offset[4], row_offset[8] - row_offset[4]); + staged = + recover + ? 0u + : (8u * fast_select::kBins * 4u + row_offset[8] * 4u <= + phase_bytes + ? 1u + : (4u * fast_select::kBins * 4u + half * 4u <= phase_bytes + ? 2u + : 0u)); + stats[blockIdx.x * 7 + 3] = staged; + stats[blockIdx.x * 7] = recover; + stats[blockIdx.x * 7 + 1] = lo; + stats[blockIdx.x * 7 + 2] = hi; + for (int c = 4; c < 7; ++c) stats[blockIdx.x * 7 + c] = 0; + } + __syncthreads(); + if (recover) { + fallback.retained += row_delta * 4096; + fallback.indices += row_delta * 4096; + seed_fallback::sm100_mqa_logits_core_impl< + 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, + cutlass::float_e4m3_t, float, float, decltype(scheduler), 2>( + rows, fallback, tq, ts, tk, ts, tw, scheduler); + } else { + if (threadIdx.x < 256) cutlass::arch::warpgroup_reg_dealloc<128>(); + __syncthreads(); + if (threadIdx.x >= 256) cutlass::arch::warpgroup_reg_alloc<128>(); + // Same 1024-B alignment as the cores' smem_buffer so every extern array + // aliases one base (mixed alignments put the scan layout 896 B past the + // base the launch sizes). + extern __shared__ __align__(1024) unsigned char phase_smem[]; + // All visible keys of a short row, padded with -1 (warp- or + // block-strided). + auto write_short = [&](unsigned r, unsigned t, unsigned stride) { + if (r >= rows) return; + const unsigned s0 = min(starts[r], keys), + visible = min(ends[r], keys) - s0; + int* o = fallback.output + (uint64_t)r * 2048; + for (unsigned j = t; j < 2048; j += stride) + o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; + }; + if constexpr (kSmemCand) { + // One warp per row over the shared candidates; scratch reuses the idle + // Q/KV stages. + static_assert(8 * fast_select::kScratch4Bytes <= sizeof(CShared), + "select scratch must fit the Q/KV stages"); + const unsigned warp = threadIdx.x / 32; + if (warp < 8 && row_base + warp < rows) { + const unsigned r = row_base + warp; + int* o = fallback.output + (uint64_t)r * 2048; + if (short_rows >> warp & 1u) + write_short(r, threadIdx.x & 31, 32); + else { + const uint16_t* hi = + reinterpret_cast(cand) + warp * kCandCap; + const uint8_t* lo = cand + kCandExt * 2 + warp * kCandCap; + const uint16_t* idx = + reinterpret_cast(cand + kCandExt * 3) + + warp * kCandCap; + const auto key_of = [idx](unsigned e) { return int32_t(idx[e]); }; + unsigned char* scratch = + phase_smem + warp * fast_select::kScratch4Bytes; + if (counts[r] >> 21 & 1) { + // The row pool overflowed: elements are the 3072 sealed shared + // slots, then the global spills at slots 0.. (count at slot + // 16383). + const uint64_t gbase = uint64_t(int64_t(r) + row_delta) * 16384; + const uint16_t* gv = vals + gbase; + const int* gi = indices + gbase; + const unsigned total = kCandCap + unsigned(__ldcg(gi + 16383)); + __syncwarp(); + if ((threadIdx.x & 31) == 0) + atomicAdd(stats + blockIdx.x * 7 + 5, 1); + fast_select::select_hist_generic( + [&](unsigned e) { + return e < kCandCap + ? (unsigned(hi[e]) << 8 | lo[e]) + : (unsigned(__ldcg(gv + e - kCandCap)) << 8 | + (unsigned(__ldcg(gi + e - kCandCap)) >> 16)); + }, + total, o, reinterpret_cast(scratch), + [&](unsigned e) { + return e < kCandCap + ? int32_t(idx[e]) + : int32_t(__ldcg(gi + e - kCandCap) & 0xffff); + }); + } else if (!fast_select::select_resident_v4( + hi, lo, 0u, kCandCap, + unsigned(counts[r] & ((1 << 21) - 1)), o, scratch, + key_of)) { + __syncwarp(); + if ((threadIdx.x & 31) == 0) + atomicAdd(stats + blockIdx.x * 7 + 6, 1); + fast_select::select_resident(hi, lo, kCandCap, 0u, o, + reinterpret_cast(scratch), + key_of); + } + } + } + } else if (staged) { + const unsigned warp = threadIdx.x / 32, + per_round = staged == 1 ? 8u : 4u; + __shared__ long long stamp[3]; + for (unsigned round = 0; round < 8 / per_round; ++round) { + if (warp < per_round) { + const unsigned local = round * per_round + warp, + first = round * per_round; + unsigned* hist = reinterpret_cast(phase_smem) + + warp * fast_select::kBins; + unsigned* codes = reinterpret_cast(phase_smem) + + per_round * fast_select::kBins + + (row_offset[local] - row_offset[first]); + const unsigned r = row_base + local; + const uint64_t base = uint64_t(int64_t(r) + row_delta) * 16384; + if (short_rows >> local & 1u) + write_short(r, threadIdx.x & 31, 32); + else + fast_select::select_row(codes, counts[r], vals + base, + indices + base, + fallback.output + (uint64_t)r * 2048, + hist, local == 0 ? stamp : nullptr); + } + __syncthreads(); + } + if (threadIdx.x == 0) { + stats[blockIdx.x * 7 + 4] = int(stamp[0]); + stats[blockIdx.x * 7 + 5] = int(stamp[1]); + stats[blockIdx.x * 7 + 6] = int(stamp[2]); + } + } else { + float* decoded = reinterpret_cast(phase_smem + 32768); + for (unsigned r = row_base; r < row_base + 8; ++r) { + if (short_rows >> (r - row_base) & 1u) { + write_short(r, threadIdx.x, 384); + __syncthreads(); + continue; + } + const uint64_t base = uint64_t(int64_t(r) + row_delta) * 16384; + const unsigned count = counts[r]; + int* output = fallback.output + (uint64_t)r * 2048; + for (unsigned j = threadIdx.x; j < count; j += 384) + decoded[j] = fused_candidates::candidate_decode_score( + vals[base + j], indices[base + j]); + __syncthreads(); + isolated_upstream::topKPerRowJob<384, 2304, false>( + nullptr, decoded, 0, count, output, nullptr, 1, 2048); + __syncthreads(); + for (unsigned j = threadIdx.x; j < 2048; j += 384) + output[j] = fused_candidates::candidate_decode_index( + indices[base + output[j]]); + __syncthreads(); + } + } + } + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + atomicExch(slot_flags + slot, 0); + } + finish_rows(recover); + } +} diff --git a/csrc/fused_mqa_topk/sample_core.cuh b/csrc/fused_mqa_topk/sample_core.cuh new file mode 100644 index 000000000000..5a8c0b82a77c --- /dev/null +++ b/csrc/fused_mqa_topk/sample_core.cuh @@ -0,0 +1,625 @@ +// SPDX-License-Identifier: MIT +// Experimental derivative of DeepGEMM 2.8 sm100_mqa_logits.cuh. +// Persistent DeepGEMM FP8 MQA producer with fixed-threshold candidate emission. +#pragma once +#ifndef MQA_TMEM_ROWS + #define MQA_TMEM_ROWS 2 +#endif + +#include +#include "sample_shared.cuh" +#ifndef MQA_HANDOFF_INTERVAL + #define MQA_HANDOFF_INTERVAL 64 +#endif +static_assert(MQA_HANDOFF_INTERVAL == 32 || MQA_HANDOFF_INTERVAL == 64 || + MQA_HANDOFF_INTERVAL == 128); + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "sm100_candidate_scan.cuh" + +// Shared SM100 MQA logits core plus contiguous-KV and paged entries +// Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature + +namespace seed_sample { +using namespace deep_gemm; + +// Convert runtime valid-token count to `cute::Int` so token loops stay +// compile-time constant +template +CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, + Fn&& fn) { + if constexpr (kCandidate <= 1) { + fn(cute::Int<1>{}); + } else if (num_block_tokens >= kCandidate) { + fn(cute::Int{}); + } else { + dispatch_num_block_tokens(num_block_tokens, + static_cast(fn)); + } +} + +// Load heads using power-of-two TMEM chunks. +template +CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, + float* accum) { + constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; + if constexpr (kNumRemaining > 0) { + constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 + : kNumRemaining >= 32 ? 32 + : kNumRemaining >= 16 ? 16 + : kNumRemaining >= 8 ? 8 + : 4; + ptx::tmem_load_32dp32b( + tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); + load_tmem_heads_decomposed(tmem_col, accum); + } +} + +// Shared device core parameterized by dtype and scheduler geometry/addressing +template < + uint32_t kNumHeads, uint32_t kHeadDim, bool kIsMXSF, + bool kIsCompressedLogits, bool kCleanLogits, uint32_t BLOCK_Q, + uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages, + uint32_t kNumSMs, uint32_t kNumSpecializedThreads, uint32_t kNumMathThreads, + typename qk_dtype_t, typename logits_dtype_t, typename reduce_dtype_t, + typename MakeScheduler, uint32_t kNumMathWarpGroups = kNumMathThreads / 128> +CUTLASS_DEVICE void sm100_mqa_logits_core_impl( + const uint32_t num_q_tokens, SampleBuffers buffers, + const cute::TmaDescriptor& tensor_map_q, + const cute::TmaDescriptor& tensor_map_sf_q, + const cute::TmaDescriptor& tensor_map_kv, + const cute::TmaDescriptor& tensor_map_sf_kv, + const cute::TmaDescriptor& tensor_map_weights, + const MakeScheduler& make_scheduler) { + constexpr bool kIsFP4 = cute::is_same_v; + + const auto sm_idx = blockIdx.x; + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; + + const auto kNegInf = -cute::numeric_limits::infinity(); + + if (warp_idx == kSpecWarpStart) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_sf_q); + cute::prefetch_tma_descriptor(&tensor_map_weights); + cute::prefetch_tma_descriptor(&tensor_map_kv); + cute::prefetch_tma_descriptor(&tensor_map_sf_kv); + } + + static constexpr uint32_t kNumTmemStages = 2; + static_assert(BLOCK_Q == SampleQ && kNumMathThreads == 256 && !kIsMXSF); + static constexpr uint32_t kNumUTCCPAlignedElems = 128; + static constexpr uint32_t UMMA_M = 128; + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); + static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; + static constexpr uint32_t kNumSFQ = + kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumSFKV = + kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumQKBytesPerToken = + kIsFP4 ? (kHeadDim / 2) : kHeadDim; + static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = + BLOCK_QH * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = + SPLIT_KV * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = + kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; + static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = + kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightElementsPerRow = + kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); + static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = + BLOCK_Q * kNumWeightBytesPerRow; + + DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, + "Invalid threads"); + DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and + SPLIT_KV % kNumUTCCPAlignedElems == 0, + "Invalid `SPLIT_KV`"); + DG_STATIC_ASSERT(not(kIsCompressedLogits and kCleanLogits), + "Compressed logits cannot be cleaned in-kernel"); + + using SharedStorage = layout::MQALogitsSharedStorage< + kNumHeads, kHeadDim, kIsMXSF, BLOCK_Q, SPLIT_KV, kNumQStages, + kNumKVStages, kNumTmemStages, qk_dtype_t, reduce_dtype_t>; + extern __shared__ __align__(SharedStorage::kSwizzleAlignment) + uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + constexpr uint32_t kRingSlots = 8; + constexpr uint32_t kRecordBlockMask = 255; + auto& stream = + *reinterpret_cast(smem_buffer + sizeof(SharedStorage)); + stream.init(threadIdx.x); + + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; + DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); + + if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumQStages; ++i) { + smem.full_q_barriers[i].init(1); + smem.empty_q_barriers[i].init(kNumMathThreads + 32); + } +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + smem.full_kv_barriers[i].init(1); + smem.empty_kv_barriers[i].init(kIsMXSF ? 1 : kNumMathThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumTmemStages; ++i) { + smem.full_tmem_barriers[i].init(1); + smem.empty_tmem_barriers[i].init(256); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + + if (warp_idx == kSpecWarpStart + 2) + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); + __syncthreads(); + + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + + constexpr uint32_t kNumSpecializedRegisters = 56; + constexpr uint32_t kNumMathRegisters = 224; + + cudaGridDependencySynchronize(); + + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + if (cute::elect_one_sync()) { + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + tma::copy( + &tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); + if constexpr (kIsMXSF) + tma::copy( + &tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_sf_q[q_stage_idx], 0, q_token_base); + 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_SF_Q_SIZE_PER_STAGE + + SMEM_WEIGHT_SIZE_PER_STAGE); + } + } + __syncwarp(); + } else if (warp_idx == kSpecWarpStart + 1) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + + auto scheduler = make_scheduler(sm_idx); + uint32_t cached_kv_page_base = 0; + uint32_t cached_kv_page_coord = 0; + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + cached_kv_page_base = cute::numeric_limits::max(); +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if constexpr (decltype(scheduler)::kIsPaged) { + constexpr uint32_t kPageKV = decltype(scheduler)::kPageKV; + constexpr uint32_t kNumPagesPerSplit = + decltype(scheduler)::kNumPagesPerSplit; + DG_STATIC_ASSERT(kNumPagesPerSplit <= 32, + "Split spans more pages than a warp can cache"); + + const uint32_t kv_page_base = + (kv_base + kv_split_idx) * kNumPagesPerSplit; + if (kv_page_base < cached_kv_page_base or + kv_page_base + kNumPagesPerSplit > cached_kv_page_base + 32) { + cached_kv_page_base = (kv_page_base / 32) * 32; + cached_kv_page_coord = scheduler.get_kv_page_coord_by_page_offset( + cached_kv_page_base + lane_idx); + } + + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + if (cute::elect_one_sync()) + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); + __syncwarp(); + + int page_coords[kNumPagesPerSplit]; +#pragma unroll + for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; + ++page_idx) { + const auto src_lane = + static_cast(kv_page_base - cached_kv_page_base + page_idx); + page_coords[page_idx] = + __shfl_sync(0xffffffff, cached_kv_page_coord, src_lane); + } + + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; + ++page_idx) { + tma::copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx] + + page_idx * kPageKV * kNumQKBytesPerToken, + 0, 0, 1, page_coords[page_idx]); + tma::copy( + &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_sf_kv[kv_stage_idx] + page_idx * kPageKV, 0, + page_coords[page_idx]); + } + 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 (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); + + // Strided sample: groups t*256.. at offset S/2 through the {128, G, + // S} KV view. Scales are read per key by the math warps, so only KV + // bytes arrive here. + tma::copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx], 0, + scheduler.sample_group(kv_split_idx, 0), 1, + scheduler.sample_center()); + smem.full_kv_barriers[kv_stage_idx].arrive_and_expect_tx( + SMEM_KV_SIZE_PER_STAGE); + } + __syncwarp(); + } + } + } else if (warp_idx == kSpecWarpStart + 2) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); + + auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { + DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, + "Invalid aligned elements"); + uint32_t values[4]; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) + values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); + __syncwarp(); + ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], + values[3]); + }; + + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + + if constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; + utccp_required_smem_warp_transpose(smem_ptr); + } + cutlass::arch::fence_view_async_shared(); +#pragma unroll + for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + if (cute::elect_one_sync()) + cute::SM100_UTCCP_4x32dp128bit_1cta::copy( + sf_desc, kTmemStartColOfSFQ + i * 4); + __syncwarp(); + } + } + + 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 constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; + utccp_required_smem_warp_transpose(smem_ptr); + } + cutlass::arch::fence_view_async_shared(); + } + + if (cute::elect_one_sync()) { + if constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute::SM100_UTCCP_4x32dp128bit_1cta::copy( + sf_desc, kTmemStartColOfSFKV + i * 4); + } + } +#pragma unroll + for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + 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(); + + if constexpr (kIsMXSF) { + DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or + kHeadDim == 64 or kHeadDim == 128, + "Invalid head dim"); + + constexpr uint32_t kPackFactor = + get_smem_pack_factor(); + constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; + + using mma_op_t = + cute::conditional_t; + auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< + qk_dtype_t, qk_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, + UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K>(); +#pragma unroll + for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { + auto runtime_instr_desc = + mma::sm100::make_runtime_instr_desc_with_sf_id( + instr_desc, k * kPackFactor, k * kPackFactor); + 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); + mma_op_t::fma(a_desc, b_desc, tmem_addr, k, runtime_instr_desc, + kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ); + } + } else { + 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); +#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(); + if constexpr (kIsMXSF) + cutlass::arch::umma_arrive(reinterpret_cast( + &smem.empty_kv_barriers[kv_stage_idx])); + } + smem.empty_q_barriers[q_stage_idx].arrive(); + } + } else if (warp_idx == kSpecWarpStart + 3) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + } else if (warp_idx < kSpecWarpStart) { + cutlass::arch::warpgroup_reg_alloc<224>(); + + constexpr bool kLoadSeqBounds = true; + uint32_t seq_k_start[4]; + uint32_t seq_k_end[4]; + const auto math_warpgroup_idx = warp_idx / 4; + const auto math_thread_idx = warp_idx * 32 + lane_idx; + DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, + "Math warp groups exceed TMEM stages"); + const unsigned query_base = math_warpgroup_idx * (SampleQ / 2); + + constexpr bool kIsReduceBF16 = not cute::is_same_v; + DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); + DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); + using weights_reg_dtype_t = + cute::conditional_t; + constexpr uint32_t kNumWeightsRegPerToken = + kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; + weights_reg_dtype_t weights[4][32]; + float accum[kNumHeads]; + + unsigned produce_epoch = 0; + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + + const auto process_q_block = [&](auto num_valid_tokens_t) { + constexpr uint32_t kNumValidTokens = SampleQ / 2; + uint32_t lane_counts[4] = {}; + uint32_t flush_base = 0; + float gate_score[4]; + +#pragma unroll + for (uint32_t i = 0; i < kNumValidTokens; ++i) { + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + // Use the raw-score threshold estimated from the distributed sample. + // This threshold remains fixed for the complete scan. + gate_score[i] = -CUDART_INF_F; + seq_k_start[i] = cute::min(scheduler.cu_seq_len_k_start[row], + scheduler.num_kv_tokens); + seq_k_end[i] = cute::min(scheduler.cu_seq_len_k_end[row], + scheduler.num_kv_tokens); + const auto smem_weights_row = + smem.smem_weights[q_stage_idx] + + (query_base + i) * kNumWeightElementsPerRow; + if constexpr (kIsReduceBF16) { + // Load two bf16 weights at a time as one packed shared u32 + const auto packed_row = + reinterpret_cast(smem_weights_row); +#pragma unroll + for (uint32_t j = 0; j < kNumHeads / 2; ++j) { + const auto packed = ptx::ld_shared(packed_row + j); + weights[i][j] = ptx::exchange( + *reinterpret_cast(&packed), 0); + } + } else { +#pragma unroll + for (unsigned j = 0; j < 32; ++j) + weights[i][j] = 0.5f * ptx::ld_shared(smem_weights_row + j); + } + } + + 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); +#pragma unroll + for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { + const unsigned kv_lane = (math_thread_idx % 128) + kv_half * 128; + const unsigned kv_offset = + scheduler.get_logits_col(kv_base, kv_split_idx, kv_lane); + const bool in_view = scheduler.sample_group(kv_split_idx, kv_lane) < + scheduler.sample_groups; + const float scale_kv = + in_view ? __ldg(scheduler.sample_scales + kv_offset) : 0.0f; + 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(); + } + + static_assert(!kIsReduceBF16, "FP32-only pair reduction"); + float accum_pair[MQA_TMEM_ROWS][kNumHeads]; +#pragma unroll + for (uint32_t pair = 0; pair < kNumValidTokens; + pair += MQA_TMEM_ROWS) { + // Load only MQA_TMEM_ROWS query rows at once to limit + // live accumulator registers during score reduction. + const uint32_t tmem_addr = + tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads; + ptx::tmem_load_32dp32b( + tmem_addr, reinterpret_cast(accum_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair + MQA_TMEM_ROWS >= kNumValidTokens) { + ptx::tcgen05_before_thread_sync(); + smem.empty_tmem_barriers[tmem_stage_idx].arrive(); + } + + float2 sum_0[MQA_TMEM_ROWS] = {}; + float2 sum_1[MQA_TMEM_ROWS] = {}; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { +#pragma unroll + for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { + if (pair + r >= kNumValidTokens) continue; + const uint32_t i = pair + r; + auto* a = accum_pair[r]; + auto transform = [&](uint32_t at, float2 sum) { + const auto x = make_float2(a[at], a[at + 1]); + const auto abs_x = + make_float2(fabsf(a[at]), fabsf(a[at + 1])); + const auto w = + make_float2(weights[i][at], weights[i][at + 1]); + return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); + }; + sum_0[r] = transform(j, sum_0[r]); + sum_1[r] = transform(j + 2, sum_1[r]); + } + } +#pragma unroll + for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { + if (pair + r >= kNumValidTokens) continue; + const uint32_t i = pair + r; + const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); + const float reduced = sum.x + sum.y; + const float bq = reduced * static_cast(scale_kv); + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + bool valid = in_view && kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq) && + bq > gate_score[i]; + stream.scores[query_base + i][kv_lane] = + valid ? bq : -CUDART_INF_F; + } + } + } // KV halves; query ownership is independent of the KV half. + sample_sync(); + accumulate_sample(stream, math_thread_idx); + sample_sync(); + } + }; + + if constexpr (decltype(scheduler)::kHasPartialBlock) + dispatch_num_block_tokens( + scheduler.get_num_block_tokens(q_block_idx), process_q_block); + else + process_q_block(cute::Int{}); + + sample_sync(); + { + unsigned local = math_thread_idx / 32; + unsigned row = scheduler.get_logits_row(q_block_idx, local); + finish_sample( + stream, buffers, local, row, + min(scheduler.cu_seq_len_k_end[row], scheduler.num_kv_tokens) - + min(scheduler.cu_seq_len_k_start[row], scheduler.num_kv_tokens), + math_thread_idx % 32); + } + sample_sync(); + cutlass::arch::fence_view_async_shared(); + smem.empty_q_barriers[q_stage_idx].arrive(); + } + + cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +} // namespace seed_sample diff --git a/csrc/fused_mqa_topk/sample_shared.cuh b/csrc/fused_mqa_topk/sample_shared.cuh new file mode 100644 index 000000000000..b603690d4489 --- /dev/null +++ b/csrc/fused_mqa_topk/sample_shared.cuh @@ -0,0 +1,105 @@ +#pragma once +#include +#include +#include +#include +constexpr unsigned SampleQ = 8, SampleN = 8192; +// target: candidate count the gate aims for (0 = legacy 4096/6144 by visible +// length). +struct SampleBuffers { + float* ring; + float* retained; + int* indices; + int* output; + unsigned target; +}; +struct alignas(16) SampleShared { + float scores[SampleQ][256]; + unsigned hist[SampleQ][2112], counts[SampleQ], rank, cut; + __device__ void init(unsigned tid) { + for (unsigned j = tid; j < SampleQ * 2112; j += 384) + reinterpret_cast(hist)[j] = 0; + if (tid < SampleQ) counts[tid] = 0; + } +}; +__device__ __forceinline__ void sample_sync() { + asm volatile("" ::: "memory"); + cutlass::arch::NamedBarrier(256, 2).sync(); + asm volatile("" ::: "memory"); +} +__device__ __forceinline__ unsigned sample_bin(float x) { + unsigned h = __half_as_ushort(__float2half_rd(x)); + return ((h & 0x8000u) ? ((~h) & 65535u) : (h ^ 0x8000u)) >> 5; +} +__device__ __forceinline__ float sample_floor(unsigned b) { + unsigned o = b << 5, h = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); + return __half2float(__ushort_as_half(h)); +} +__device__ void finish_sample(SampleShared& s, SampleBuffers b, unsigned local, + unsigned row, unsigned visible, unsigned tid) { + const unsigned target = + b.target ? b.target : (visible <= 65536 ? 4096u : 6144u); + if (visible <= target) { + if (tid == 0) b.ring[row] = -CUDART_INF_F; + __syncwarp(); + return; + } + unsigned lane = tid % 32; + auto* hist = s.hist[local]; + unsigned count = s.counts[local]; + if (count == 0) { + if (tid == 0) b.ring[row] = -CUDART_INF_F; + __syncwarp(); + return; + } + unsigned sample_rank = + max(1u, min(count, (count * target + visible - 1) / visible)); + __syncwarp(); + if (tid < 32) { + unsigned sum = 0; + for (unsigned j = 0; j < 64; ++j) { + unsigned bin = 2047 - lane * 64 - j; + sum += hist[bin + bin / 32]; + } + unsigned inclusive = sum; + for (unsigned d = 1; d < 32; d *= 2) { + unsigned v = __shfl_up_sync(0xffffffffu, inclusive, d); + if (lane >= d) inclusive += v; + } + unsigned before = inclusive - sum, rank = sample_rank, + mask = + __ballot_sync(0xffffffffu, before < rank && inclusive >= rank), + winner = __ffs(mask) - 1, chosen = 0, above = before; + if (lane == winner) { + for (unsigned j = 0; j < 64; ++j) { + unsigned bin = 2047 - lane * 64 - j; + if (above + hist[bin + bin / 32] >= rank) { + chosen = bin; + break; + } + above += hist[bin + bin / 32]; + } + } + chosen = __shfl_sync(0xffffffffu, chosen, winner); + if (lane == 0) b.ring[row] = sample_floor(chosen); + } + __syncwarp(); +} + +__device__ void accumulate_sample(SampleShared& s, unsigned tid) { + unsigned local = tid / 32, lane = tid % 32, count = 0; +#pragma unroll + for (unsigned j = lane; j < 256; j += 32) { + float v = s.scores[local][j]; + bool valid = isfinite(v); + count += valid; + unsigned mask = __ballot_sync(0xffffffffu, valid); + if (valid) { + unsigned bin = sample_bin(v), peers = __match_any_sync(mask, bin); + if (lane == unsigned(__ffs(peers) - 1)) + atomicAdd(s.hist[local] + bin + bin / 32, __popc(peers)); + } + } + count = __reduce_add_sync(0xffffffffu, count); + if (lane == 0) s.counts[local] += count; +} diff --git a/csrc/fused_mqa_topk/scan.cuh b/csrc/fused_mqa_topk/scan.cuh new file mode 100644 index 000000000000..42dfdfed5b0b --- /dev/null +++ b/csrc/fused_mqa_topk/scan.cuh @@ -0,0 +1,1106 @@ +// SPDX-License-Identifier: MIT +// Experimental derivative of DeepGEMM 2.8 sm100_mqa_logits.cuh. +// Persistent DeepGEMM FP8 MQA producer with fixed-threshold candidate emission. +#pragma once + +#include +#include "handoff.cuh" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "candidates.cuh" + +// Shared SM100 MQA logits core plus contiguous-KV and paged entries +// Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature + +namespace fused_mqa_scan { +using namespace deep_gemm; + +// Convert runtime valid-token count to `cute::Int` so token loops stay +// compile-time constant +template +CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, + Fn&& fn) { + if constexpr (kCandidate <= 1) { + fn(cute::Int<1>{}); + } else if (num_block_tokens >= kCandidate) { + fn(cute::Int{}); + } else { + dispatch_num_block_tokens(num_block_tokens, + static_cast(fn)); + } +} + +// Load heads using power-of-two TMEM chunks. +template +CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, + float* accum) { + constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; + if constexpr (kNumRemaining > 0) { + constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 + : kNumRemaining >= 32 ? 32 + : kNumRemaining >= 16 ? 16 + : kNumRemaining >= 8 ? 8 + : 4; + ptx::tmem_load_32dp32b( + tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); + load_tmem_heads_decomposed(tmem_col, accum); + } +} + +// Warp 0 finds the descending radix bucket. Each lane sums eight adjacent +// bins, the warp locates the group containing rank, and one lane scans only +// that group's eight bins. The selected bin and preceding count are shared +// through ctl after the caller's named barrier. +CUTLASS_DEVICE void select_hist_bucket(int32_t* hist, int32_t* ctl, + uint32_t tid) { + if (tid >= 32) return; + constexpr uint32_t FULL = 0xffffffffu; + int group_count = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) + group_count += hist[255 - static_cast(tid) * 8 - j]; + int inclusive = group_count; +#pragma unroll + for (int shift = 1; shift < 32; shift <<= 1) { + const int prev = __shfl_up_sync(FULL, inclusive, shift); + if (tid >= static_cast(shift)) inclusive += prev; + } + const int before = inclusive - group_count; + const int rank = ctl[0]; + const uint32_t crossing = + __ballot_sync(FULL, before < rank && inclusive >= rank); + const int chosen_lane = __ffs(crossing) - 1; + int chosen = 0; + int above = before; + if (static_cast(tid) == chosen_lane) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int bin = 255 - static_cast(tid) * 8 - j; + if (above + hist[bin] >= rank) { + chosen = bin; + break; + } + above += hist[bin]; + } + } + chosen = __shfl_sync(FULL, chosen, chosen_lane); + above = __shfl_sync(FULL, above, chosen_lane); + if (tid == 0) { + ctl[0] = rank - above; + ctl[1] = (ctl[1] << 8) | chosen; + } +} + +// One CTA owns all KV tiles of a query block. At a flush boundary, its math +// warps compact one row's global candidate buffer in place through the now +// vacant shared ring. This keeps the live set bounded independently of ISL. +template +CUTLASS_DEVICE void compact_candidates(uint16_t* values, + int32_t* packed_indices, + int32_t* count_ptr, uint32_t* gate_ptr, + int32_t* hist, uint64_t* scratch, + uint32_t tid, uint32_t cap) { + constexpr int K = 2048; + auto sync = [] { cutlass::arch::NamedBarrier(MATH_THREADS, 1).sync(); }; + int32_t* ctl = hist + 256; + const int count = *count_ptr; + if (tid == 0) { + ctl[0] = K; + ctl[1] = 0; + ctl[2] = 0; + ctl[3] = 0; + } + sync(); + for (int shift = 16; shift >= 8; shift -= 8) { + if (tid < 256) hist[tid] = 0; + sync(); + const uint32_t prefix = static_cast(ctl[1]); + for (int j = tid; j < count; j += MATH_THREADS) { + const uint32_t code = fused_candidates::candidate_load_score_code( + values[j], packed_indices[j]); + if (shift == 16 || (code >> (shift + 8)) == prefix) + atomicAdd(hist + ((code >> shift) & 255u), 1); + } + sync(); + select_hist_bucket(hist, ctl, tid); + sync(); + } + const uint32_t prefix16 = static_cast(ctl[1]); + const int coarse_kept = K - ctl[0] + hist[prefix16 & 255u]; + if (coarse_kept <= 3072) { + const uint32_t coarse_floor = prefix16 << 8; + for (int j = tid; j < count; j += MATH_THREADS) { + const uint32_t code = fused_candidates::candidate_load_score_code( + values[j], packed_indices[j]); + if (code >= coarse_floor) { + const int pos = atomicAdd(ctl + 2, 1); + scratch[pos] = + (static_cast(static_cast(packed_indices[j])) + << 16) | + static_cast(values[j]); + } + } + sync(); + const int kept = ctl[2]; + for (int j = tid; j < kept; j += MATH_THREADS) { + const uint64_t record = scratch[j]; + values[j] = static_cast(record); + packed_indices[j] = static_cast(record >> 16); + } + sync(); + if (tid == 0) { + *count_ptr = kept; + *gate_ptr = coarse_floor == 0 ? 0 : coarse_floor - 1; + } + sync(); + return; + } + // Unusually large score ties in one high-16 bucket use the exact pass. + if (tid < 256) hist[tid] = 0; + sync(); + for (int j = tid; j < count; j += MATH_THREADS) { + const uint32_t code = fused_candidates::candidate_load_score_code( + values[j], packed_indices[j]); + if ((code >> 8) == prefix16) atomicAdd(hist + (code & 255u), 1); + } + sync(); + select_hist_bucket(hist, ctl, tid); + sync(); + const uint32_t kth = static_cast(ctl[1]); + for (int j = tid; j < count; j += MATH_THREADS) { + const uint32_t code = fused_candidates::candidate_load_score_code( + values[j], packed_indices[j]); + if (code > kth) { + const int pos = atomicAdd(ctl + 2, 1); + if (pos < K) + scratch[pos] = + (static_cast(static_cast(packed_indices[j])) + << 16) | + static_cast(values[j]); + } + } + sync(); + const int greater = ctl[2]; + for (int j = tid; j < count; j += MATH_THREADS) { + const uint32_t code = fused_candidates::candidate_load_score_code( + values[j], packed_indices[j]); + if (code == kth) { + const int pos = atomicAdd(ctl + 3, 1); + if (greater + pos < K) + scratch[greater + pos] = + (static_cast(static_cast(packed_indices[j])) + << 16) | + static_cast(values[j]); + } + } + sync(); + for (int j = tid; j < K; j += MATH_THREADS) { + const uint64_t record = scratch[j]; + values[j] = static_cast(record); + packed_indices[j] = static_cast(record >> 16); + } + sync(); + if (tid == 0) { + *count_ptr = K; + *gate_ptr = kth; + } + sync(); +} + +CUTLASS_DEVICE float candidate_gate_float(const uint32_t code) { + if (code == 0) return -cute::numeric_limits::infinity(); + const uint32_t ordered = code << 8; + const uint32_t bits = + (ordered & 0x80000000u) ? (ordered ^ 0x80000000u) : ~ordered; + return __uint_as_float(bits); +} + +// Shared device core parameterized by dtype and scheduler geometry/addressing +template < + uint32_t kNumHeads, uint32_t kHeadDim, bool kIsMXSF, + bool kIsCompressedLogits, bool kCleanLogits, uint32_t BLOCK_Q, + uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages, + uint32_t kNumSMs, uint32_t kNumSpecializedThreads, uint32_t kNumMathThreads, + typename qk_dtype_t, typename logits_dtype_t, typename reduce_dtype_t, + typename MakeScheduler, uint32_t kNumMathWarpGroups = kNumMathThreads / 128, + uint32_t kHandoffInterval = 128, bool kSmemCand = false> +CUTLASS_DEVICE void sm100_mqa_logits_core_impl( + const uint32_t num_q_tokens, const uint32_t logits_stride, + logits_dtype_t* logits, const float* origin, const float* inv_delta, + const int32_t* th_bucket, const int32_t* seed_hist, uint16_t* cand_val, + int32_t* cand_idx, int32_t* cand_cnt, const uint32_t cand_cap, + const int64_t cand_row_delta, const cute::TmaDescriptor& tensor_map_q, + const cute::TmaDescriptor& tensor_map_sf_q, + const cute::TmaDescriptor& tensor_map_kv, + const cute::TmaDescriptor& tensor_map_sf_kv, + const cute::TmaDescriptor& tensor_map_weights, + const MakeScheduler& make_scheduler) { + constexpr bool kIsFP4 = cute::is_same_v; + + const auto sm_idx = blockIdx.x; + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; + + const auto kNegInf = -cute::numeric_limits::infinity(); + + if (warp_idx == kSpecWarpStart) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_sf_q); + cute::prefetch_tma_descriptor(&tensor_map_weights); + cute::prefetch_tma_descriptor(&tensor_map_kv); + cute::prefetch_tma_descriptor(&tensor_map_sf_kv); + } + + static constexpr uint32_t kNumTmemStages = 2; + static_assert(BLOCK_Q == 8 && kNumMathThreads == 256 && !kIsMXSF); + static constexpr uint32_t kNumUTCCPAlignedElems = 128; + static constexpr uint32_t UMMA_M = 128; + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); + static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; + static constexpr uint32_t kNumSFQ = + kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumSFKV = + kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumQKBytesPerToken = + kIsFP4 ? (kHeadDim / 2) : kHeadDim; + static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = + BLOCK_QH * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = + SPLIT_KV * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = + kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; + static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = + kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightElementsPerRow = + kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); + static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = + BLOCK_Q * kNumWeightBytesPerRow; + + DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, + "Invalid threads"); + DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and + SPLIT_KV % kNumUTCCPAlignedElems == 0, + "Invalid `SPLIT_KV`"); + DG_STATIC_ASSERT(not(kIsCompressedLogits and kCleanLogits), + "Compressed logits cannot be cleaned in-kernel"); + + using SharedStorage = layout::MQALogitsSharedStorage< + kNumHeads, kHeadDim, kIsMXSF, BLOCK_Q, SPLIT_KV, kNumQStages, + kNumKVStages, kNumTmemStages, qk_dtype_t, reduce_dtype_t>; + extern __shared__ __align__(SharedStorage::kSwizzleAlignment) + uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + constexpr uint32_t kRingSlots = 8; + constexpr uint32_t kRecordBlockMask = 255; + auto& handoff = *reinterpret_cast*>(smem_buffer + + sizeof(SharedStorage)); + handoff.init(threadIdx.x); + + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; + DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); + + if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumQStages; ++i) { + smem.full_q_barriers[i].init(1); + smem.empty_q_barriers[i].init(kNumMathThreads + 32); + } +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + smem.full_kv_barriers[i].init(1); + smem.empty_kv_barriers[i].init(kIsMXSF ? 1 : kNumMathThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumTmemStages; ++i) { + smem.full_tmem_barriers[i].init(1); + smem.empty_tmem_barriers[i].init(256); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + + if (warp_idx == kSpecWarpStart + 2) + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); + __syncthreads(); + + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + + constexpr uint32_t kNumSpecializedRegisters = 56; + constexpr uint32_t kNumMathRegisters = 224; + + cudaGridDependencySynchronize(); + + // Short contexts (handoff interval <= 8, i.e. <= 16K keys): candidates are + // dense and the drain was the bottleneck, so math warps emit directly (warp + // ballot + one shared atomic per row and KV half). Longer contexts keep the + // ring + drain warps. kHandoffInterval == 0: resident mode (keys <= + // logits_stride). Every score of the block's 8 rows is stored as an fp24 code + // in shared memory (hi16 array, then lo8 array, row-major with stride + // logits_stride) at `logits`; no gate, no candidate emission. + constexpr bool kResident = kHandoffInterval == 0; + constexpr bool kDirectEmit = !kResident && kHandoffInterval <= 8; + // Two drain warps: warp +3 drains math warpgroup 0 (rows 0-3) and warp +0, + // idle after its single Q load, drains warpgroup 1 (rows 4-7). Rows are + // disjoint. + const auto drain_producers = [&](const unsigned p0) { + auto scheduler = make_scheduler(sm_idx); + uint32_t q_block_idx, kv_base, num_kv_splits; + unsigned drain_epoch = 0; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + const unsigned row_base = scheduler.get_logits_row(q_block_idx, 0); + for (unsigned base = 0; base < num_kv_splits; + base += kHandoffInterval, ++drain_epoch) { + for (unsigned p = p0; p < p0 + 4; ++p) { + handoff.drain( + p, drain_epoch, 4, row_base + (p / 4) * 4, cand_cnt, cand_cap, + [&](unsigned row, unsigned slot, unsigned record, + unsigned batch_base, unsigned source) { + const uint32_t physical = + kv_base + (batch_base * 2 + (record & 255u)) * 128 + + (source % 128); + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + cand_cap + + slot; + fused_candidates::store_candidate_record( + cand_val + off, cand_idx + off, record >> 8, physical); + }); + } + } + } + }; + + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc(); + if (cute::elect_one_sync()) { + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + tma::copy( + &tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); + if constexpr (kIsMXSF) + tma::copy( + &tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_sf_q[q_stage_idx], 0, q_token_base); + 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_SF_Q_SIZE_PER_STAGE + + SMEM_WEIGHT_SIZE_PER_STAGE); + } + } + __syncwarp(); + if constexpr (!kDirectEmit && !kResident) drain_producers(4u); + } else if (warp_idx == kSpecWarpStart + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + auto scheduler = make_scheduler(sm_idx); + uint32_t cached_kv_page_base = 0; + uint32_t cached_kv_page_coord = 0; + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + cached_kv_page_base = cute::numeric_limits::max(); +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if constexpr (decltype(scheduler)::kIsPaged) { + constexpr uint32_t kPageKV = decltype(scheduler)::kPageKV; + constexpr uint32_t kNumPagesPerSplit = + decltype(scheduler)::kNumPagesPerSplit; + DG_STATIC_ASSERT(kNumPagesPerSplit <= 32, + "Split spans more pages than a warp can cache"); + + const uint32_t kv_page_base = + (kv_base + kv_split_idx) * kNumPagesPerSplit; + if (kv_page_base < cached_kv_page_base or + kv_page_base + kNumPagesPerSplit > cached_kv_page_base + 32) { + cached_kv_page_base = (kv_page_base / 32) * 32; + cached_kv_page_coord = scheduler.get_kv_page_coord_by_page_offset( + cached_kv_page_base + lane_idx); + } + + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + if (cute::elect_one_sync()) + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); + __syncwarp(); + + int page_coords[kNumPagesPerSplit]; +#pragma unroll + for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; + ++page_idx) { + const auto src_lane = + static_cast(kv_page_base - cached_kv_page_base + page_idx); + page_coords[page_idx] = + __shfl_sync(0xffffffff, cached_kv_page_coord, src_lane); + } + + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; + ++page_idx) { + tma::copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx] + + page_idx * kPageKV * kNumQKBytesPerToken, + 0, 0, 1, page_coords[page_idx]); + tma::copy( + &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_sf_kv[kv_stage_idx] + page_idx * kPageKV, 0, + page_coords[page_idx]); + } + 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 (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); + tma::copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); + 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 utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { + DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, + "Invalid aligned elements"); + uint32_t values[4]; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) + values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); + __syncwarp(); + ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], + values[3]); + }; + + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + + if constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; + utccp_required_smem_warp_transpose(smem_ptr); + } + cutlass::arch::fence_view_async_shared(); +#pragma unroll + for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + if (cute::elect_one_sync()) + cute::SM100_UTCCP_4x32dp128bit_1cta::copy( + sf_desc, kTmemStartColOfSFQ + i * 4); + __syncwarp(); + } + } + + 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 constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; + utccp_required_smem_warp_transpose(smem_ptr); + } + cutlass::arch::fence_view_async_shared(); + } + + if (cute::elect_one_sync()) { + if constexpr (kIsMXSF) { +#pragma unroll + for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { + auto smem_ptr = + smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute::SM100_UTCCP_4x32dp128bit_1cta::copy( + sf_desc, kTmemStartColOfSFKV + i * 4); + } + } +#pragma unroll + for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + 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(); + + if constexpr (kIsMXSF) { + DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or + kHeadDim == 64 or kHeadDim == 128, + "Invalid head dim"); + + constexpr uint32_t kPackFactor = + get_smem_pack_factor(); + constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; + + using mma_op_t = + cute::conditional_t; + auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< + qk_dtype_t, qk_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, + UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K>(); +#pragma unroll + for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { + auto runtime_instr_desc = + mma::sm100::make_runtime_instr_desc_with_sf_id( + instr_desc, k * kPackFactor, k * kPackFactor); + 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); + mma_op_t::fma(a_desc, b_desc, tmem_addr, k, runtime_instr_desc, + kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ); + } + } else { + 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); +#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(); + if constexpr (kIsMXSF) + cutlass::arch::umma_arrive(reinterpret_cast( + &smem.empty_kv_barriers[kv_stage_idx])); + } + smem.empty_q_barriers[q_stage_idx].arrive(); + } + } else if (warp_idx == kSpecWarpStart + 3) { + cutlass::arch::warpgroup_reg_dealloc(); + if constexpr (!kDirectEmit && !kResident) drain_producers(0u); + } else if (warp_idx < kSpecWarpStart) { + cutlass::arch::warpgroup_reg_alloc(); + + constexpr bool kLoadSeqBounds = true; + uint32_t seq_k_start[4]; + uint32_t seq_k_end[4]; + const auto math_warpgroup_idx = warp_idx / 4; + const auto math_thread_idx = warp_idx * 32 + lane_idx; + DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, + "Math warp groups exceed TMEM stages"); + const unsigned query_base = math_warpgroup_idx * 4; + + constexpr bool kIsReduceBF16 = not cute::is_same_v; + DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); + DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); + using weights_reg_dtype_t = + cute::conditional_t; + constexpr uint32_t kNumWeightsRegPerToken = + kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; + weights_reg_dtype_t weights[4][kNumWeightsRegPerToken]; + float accum[kNumHeads]; + + unsigned produce_epoch = 0; + auto scheduler = make_scheduler(sm_idx); + // NOTES: split index for paged scheduler, token offset for contiguous-KV + // scheduler. + uint32_t q_block_idx, kv_base, num_kv_splits; + 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); + + const auto process_q_block = [&](auto num_valid_tokens_t) { + constexpr uint32_t kNumValidTokens = 4; + uint32_t lane_counts[4] = {}; + uint32_t flush_base = 0; + float gate_score[4]; + +#pragma unroll + for (uint32_t i = 0; i < kNumValidTokens; ++i) { + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + // Use the raw-score threshold estimated from the distributed sample. + // This threshold remains fixed for the complete scan. + gate_score[i] = kResident ? 0.0f : origin[row]; + seq_k_start[i] = cute::min(scheduler.cu_seq_len_k_start[row], + scheduler.num_kv_tokens); + seq_k_end[i] = cute::min(scheduler.cu_seq_len_k_end[row], + scheduler.num_kv_tokens); + const auto smem_weights_row = + smem.smem_weights[q_stage_idx] + + (query_base + i) * kNumWeightElementsPerRow; + if constexpr (kIsReduceBF16) { + // Load two bf16 weights at a time as one packed shared u32 + const auto packed_row = + reinterpret_cast(smem_weights_row); +#pragma unroll + for (uint32_t j = 0; j < kNumHeads / 2; ++j) { + const auto packed = ptx::ld_shared(packed_row + j); + weights[i][j] = ptx::exchange( + *reinterpret_cast(&packed), 0); + } + } else { + bool unsafe_weights = false; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; ++j) { + const float weight = ptx::ld_shared(smem_weights_row + j); + unsafe_weights |= weight != 0.0f && fabsf(weight) < 0x1p-125f; + // Resident: full weights, the 0.5 ReLU factor is applied per + // score. + weights[i][j] = kResident ? weight : weight * 0.5f; + } + // Hoist the ReLU factor out of the per-key reduction. + // Tiny weights can round differently when halved; use + // the existing GPU recovery with unscaled weights. + if (!kResident && unsafe_weights) { + gate_score[i] = __int_as_float(0x7f800000); + if (math_thread_idx % 128 == 0) + atomicAdd(cand_cnt + row, static_cast(cand_cap) + 1); + } + } + } + + __shared__ unsigned row_count[BLOCK_Q]; + // Shared candidates, per row [seg0][seg1][seg2][seg3][pool]: each of + // the row's 4 warps owns a static segment (no atomics; its last slot is + // a trash slot for lanes without a candidate); a warp past its segment + // takes slots from the row's shared pool (one atomic per overflowing + // ballot); past the pool, global slots row*cand_cap + j (count at + // cand_cap - 1). Unfilled slots get the lowest code. + const uint32_t pool_size = logits_stride / 6, + seg_size = (logits_stride - pool_size) / 4; + const uint32_t seg_warp = warp_idx % 4; + // Each array holds the 8 rows plus 256 per-lane trash slots (8 warps x + // 32 lanes). + const uint32_t cand_ext = BLOCK_Q * logits_stride + 256; + __shared__ unsigned pool_cnt[BLOCK_Q]; + uint32_t seg_cnt[kNumValidTokens] = {}; + if constexpr (kDirectEmit) { + if (math_thread_idx < BLOCK_Q) { + row_count[math_thread_idx] = 0; + pool_cnt[math_thread_idx] = 0; + } + cutlass::arch::NamedBarrier(kNumMathThreads, 3).sync(); + } + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if constexpr (!kDirectEmit && !kResident) + if ((kv_split_idx & (kHandoffInterval - 1)) == 0) + handoff.acquire(warp_idx, produce_epoch); + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); +#pragma unroll + for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { + const unsigned kv_lane = (math_thread_idx % 128) + kv_half * 128; + const unsigned 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(); + } + + static_assert(!kIsReduceBF16, "FP32-only pair reduction"); + float accum_pair[2][kNumHeads]; +#pragma unroll + for (uint32_t pair = 0; pair < kNumValidTokens; pair += 2) { + // Load only 2 query rows at once to limit + // live accumulator registers during score reduction. + const uint32_t tmem_addr = + tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads; + ptx::tmem_load_32dp32b<2 * kNumHeads>( + tmem_addr, reinterpret_cast(accum_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair + 2 >= kNumValidTokens) { + ptx::tcgen05_before_thread_sync(); + smem.empty_tmem_barriers[tmem_stage_idx].arrive(); + } + + float2 sum_0[2] = {}; + float2 sum_1[2] = {}; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + if (pair + r >= kNumValidTokens) continue; + const uint32_t i = pair + r; + auto* a = accum_pair[r]; + auto transform = [&](uint32_t at, float2 sum) { + const auto x = make_float2(a[at], a[at + 1]); + const auto abs_x = + make_float2(fabsf(a[at]), fabsf(a[at + 1])); + const auto w = + make_float2(weights[i][at], weights[i][at + 1]); + return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); + }; + sum_0[r] = transform(j, sum_0[r]); + sum_1[r] = transform(j + 2, sum_1[r]); + } + } +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + if (pair + r >= kNumValidTokens) continue; + const uint32_t i = pair + r; + const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); + const float reduced = sum.x + sum.y; + const float bq = + reduced * (kResident ? 0.5f * static_cast(scale_kv) + : static_cast(scale_kv)); + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + if constexpr (kResident) { + // Select reads only [start, end) of each row, so no masking + // here. + if (kv_offset < logits_stride) { + const uint32_t code = + fused_candidates::candidate_fp24_code(bq); + const uint32_t at = + (query_base + i) * logits_stride + kv_offset; + reinterpret_cast(logits)[at] = + static_cast(code >> 8); + reinterpret_cast( + logits)[BLOCK_Q * logits_stride * 2 + at] = + static_cast(code); + } + continue; + } + // NaN fails the compare; +inf passes (harmless, it would win + // the top-k anyway). + const bool pass = (kSmemCand ? kv_offset - seq_k_start[i] < + seq_k_end[i] - seq_k_start[i] + : kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && + isfinite(bq)) && + bq > gate_score[i]; + if constexpr (kSmemCand) { + // Branch-free: every lane stores; lanes without a candidate + // (or past the segment) write their own trash slot after the + // rows (never read). + const unsigned pass_mask = __ballot_sync(0xffffffffu, pass); + const unsigned slot = + seg_cnt[i] + __popc(pass_mask & ((1u << lane_idx) - 1u)); + seg_cnt[i] += __popc(pass_mask); + const uint32_t code = + fused_candidates::candidate_fp24_code(bq); + const uint32_t seg_cap = seg_size; + const uint32_t at = + pass && slot < seg_cap + ? (query_base + i) * logits_stride + + seg_warp * seg_size + slot + : BLOCK_Q * logits_stride + warp_idx * 32 + lane_idx; + auto* cbase = reinterpret_cast(logits); + reinterpret_cast(cbase)[at] = + static_cast(code >> 8); + cbase[cand_ext * 2 + at] = static_cast(code); + reinterpret_cast(cbase + cand_ext * 3)[at] = + static_cast(kv_offset); + if (seg_cnt[i] > seg_cap) { + // Warp-uniform: lanes past the segment take row-pool slots. + const bool over = pass && slot >= seg_cap; + const unsigned over_mask = __ballot_sync(0xffffffffu, over); + if (over_mask) { + unsigned pool_base = 0; + if (lane_idx == 0) + pool_base = atomicAdd(&pool_cnt[query_base + i], + __popc(over_mask)); + pool_base = __shfl_sync(0xffffffffu, pool_base, 0); + const unsigned p = + pool_base + + __popc(over_mask & ((1u << lane_idx) - 1u)); + if (over && p < pool_size) { + const uint32_t pat = + (query_base + i) * logits_stride + 4 * seg_size + p; + reinterpret_cast(cbase)[pat] = + static_cast(code >> 8); + cbase[cand_ext * 2 + pat] = static_cast(code); + reinterpret_cast(cbase + cand_ext * 3)[pat] = + static_cast(kv_offset); + } else if (over && p - pool_size < cand_cap - 1) { + // Past the pool (rare): global, val = hi16 code, idx = + // lo8 code << 16 | key. + const uint64_t off = + static_cast(int64_t(row) + + cand_row_delta) * + cand_cap + + (p - pool_size); + __stcg(cand_val + off, + static_cast(code >> 8)); + __stcg(cand_idx + off, + static_cast((code & 0xffu) << 16 | + kv_offset)); + } + } + } + } else if constexpr (kDirectEmit) { + const unsigned pass_mask = __ballot_sync(0xffffffffu, pass); + if (pass_mask) { + // Ballot + one shared atomic per row: global candidate + // slots. + const unsigned leader = __ffs(pass_mask) - 1; + unsigned base = 0; + if (lane_idx == leader) + base = atomicAdd(&row_count[query_base + i], + __popc(pass_mask)); + base = __shfl_sync(0xffffffffu, base, leader); + const unsigned slot = + base + __popc(pass_mask & ((1u << lane_idx) - 1u)); + if (pass && slot < cand_cap) { + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + cand_cap + + slot; + fused_candidates::store_candidate( + cand_val + off, cand_idx + off, bq, kv_offset); + } + } + } else if (pass) { + const uint32_t count = lane_counts[i]; + if (count < kRingSlots) { + const uint32_t pos = + ((warp_idx * BLOCK_Q + i) * kRingSlots + count) * 32u + + lane_idx; + const uint32_t code = + fused_candidates::candidate_fp24_code(bq); + handoff.records[produce_epoch & 1][warp_idx][i][count] + [lane_idx] = + (code << 8) | + ((kv_split_idx - flush_base) * 2 + kv_half); + lane_counts[i] = count + 1; + } else { + const int slot = atomicAdd(cand_cnt + row, 1); + if (slot < static_cast(cand_cap)) { + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + cand_cap + + slot; + fused_candidates::store_candidate( + cand_val + off, cand_idx + off, bq, kv_offset); + } + } + } + } + } + } // KV halves; query ownership is independent of the KV half. + if (!kDirectEmit && !kResident && + (((kv_split_idx + 1) & (kHandoffInterval - 1)) == 0 || + kv_split_idx + 1 == num_kv_splits)) { +#pragma unroll + for (unsigned i = 0; i < kNumValidTokens; ++i) { + handoff.counts[produce_epoch & 1][warp_idx][i][lane_idx] = + lane_counts[i]; + lane_counts[i] = 0; + } + handoff.publish(warp_idx, produce_epoch, flush_base); + ++produce_epoch; + flush_base = kv_split_idx + 1; + } + } + if constexpr (kSmemCand) { +// Seal each warp's segment tail with the lowest code; publish counts and +// spills. +#pragma unroll + for (uint32_t i = 0; i < kNumValidTokens; ++i) { + const uint32_t row = + scheduler.get_logits_row(q_block_idx, query_base + i); + auto* base = reinterpret_cast(logits); + for (uint32_t slot = min(seg_cnt[i], seg_size) + lane_idx; + slot < seg_size; slot += 32) { + const uint32_t at = + (query_base + i) * logits_stride + seg_warp * seg_size + slot; + reinterpret_cast(base)[at] = 0; + base[cand_ext * 2 + at] = 0; + reinterpret_cast(base + cand_ext * 3)[at] = 0; + } + if (lane_idx == 0) + atomicAdd(&row_count[query_base + i], seg_cnt[i]); + } + // After every warp's pool claims: seal each row's unused pool slots + // and record the global spill count (bit 21 of the published count + // marks a spilled row). + cutlass::arch::NamedBarrier(kNumMathThreads, 3).sync(); + for (uint32_t q = 0; q < BLOCK_Q; ++q) { + auto* base = reinterpret_cast(logits); + for (uint32_t p = pool_cnt[q] + math_thread_idx; p < pool_size; + p += kNumMathThreads) { + const uint32_t at = q * logits_stride + 4 * seg_size + p; + reinterpret_cast(base)[at] = 0; + base[cand_ext * 2 + at] = 0; + reinterpret_cast(base + cand_ext * 3)[at] = 0; + } + } + if (math_thread_idx < BLOCK_Q) { + const uint32_t row = + scheduler.get_logits_row(q_block_idx, math_thread_idx); + const uint32_t spilled = + pool_cnt[math_thread_idx] > pool_size + ? min(pool_cnt[math_thread_idx] - pool_size, cand_cap - 1) + : 0u; + __stcg(cand_idx + + static_cast(int64_t(row) + cand_row_delta) * + cand_cap + + cand_cap - 1, + static_cast(spilled)); + if (spilled) atomicOr(&row_count[math_thread_idx], 1u << 21); + } + } + if constexpr (kDirectEmit) { + // Publish row totals (added to any unsafe-weights recovery marker + // already in cand_cnt). Shared-candidate rows with < 2048 real + // candidates are pushed past 16384 (recover). + cutlass::arch::NamedBarrier(kNumMathThreads, 3).sync(); + if (math_thread_idx < BLOCK_Q) { + // Shared candidates: bits 0-20 = real count, bit 21 = some warp + // spilled. + int total = static_cast(row_count[math_thread_idx]); + atomicAdd(cand_cnt + scheduler.get_logits_row(q_block_idx, + math_thread_idx), + total); + } + } + }; + + if constexpr (decltype(scheduler)::kHasPartialBlock) + dispatch_num_block_tokens( + scheduler.get_num_block_tokens(q_block_idx), process_q_block); + else + process_q_block(cute::Int{}); + + cutlass::arch::fence_view_async_shared(); + smem.empty_q_barriers[q_stage_idx].arrive(); + } + + cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +// Unified contiguous-KV entry for FP8 / MXFP4 / MXFP8. +template +CUTLASS_GLOBAL __launch_bounds__( + kNumSpecializedThreads + kNumMathThreads, + 1) void sm100_mqa_logits(const uint32_t num_q_tokens, + const uint32_t num_kv_tokens, + const uint32_t logits_stride, + const uint32_t* cu_seq_len_k_start, + const uint32_t* cu_seq_len_k_end, + const uint32_t* schedule_meta, + logits_dtype_t* logits, const float* origin, + const float* inv_delta, const int32_t* th_bucket, + const int32_t* seed_hist, uint16_t* cand_val, + int32_t* cand_idx, int32_t* cand_cnt, + const uint32_t cand_cap, + const __grid_constant__ cute::TmaDescriptor + tensor_map_q, + const __grid_constant__ cute::TmaDescriptor + tensor_map_sf_q, + const __grid_constant__ cute::TmaDescriptor + tensor_map_kv, + const __grid_constant__ cute::TmaDescriptor + tensor_map_sf_kv, + const __grid_constant__ cute::TmaDescriptor + tensor_map_weights) { + const auto make_scheduler = [&](const uint32_t& sm_idx) { + return sched::SM100MQALogitsScheduler( + sm_idx, num_q_tokens, num_kv_tokens, cu_seq_len_k_start, + cu_seq_len_k_end, schedule_meta); + }; + + sm100_mqa_logits_core_impl< + kNumHeads, kHeadDim, kIsMXSF, kIsCompressedLogits, kCleanLogits, BLOCK_Q, + SPLIT_KV, kNumQStages, kNumKVStages, kNumSMs, kNumSpecializedThreads, + kNumMathThreads, qk_dtype_t, logits_dtype_t, reduce_dtype_t, + decltype(make_scheduler), kNumMathWarpGroups, kHandoffInterval>( + num_q_tokens, logits_stride, logits, origin, inv_delta, th_bucket, + seed_hist, cand_val, cand_idx, cand_cnt, cand_cap, 0, tensor_map_q, + tensor_map_sf_q, tensor_map_kv, tensor_map_sf_kv, tensor_map_weights, + make_scheduler); +} + +} // namespace fused_mqa_scan diff --git a/csrc/fused_mqa_topk/sm100_candidate_scan.cuh b/csrc/fused_mqa_topk/sm100_candidate_scan.cuh new file mode 100644 index 000000000000..28aaf07613fc --- /dev/null +++ b/csrc/fused_mqa_topk/sm100_candidate_scan.cuh @@ -0,0 +1,1276 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025 DeepSeek +// Derived from DeepGEMM commit 891d57b4db1071624b5c8fa0d1e51cb317fa709f; +// see LICENSE.deepseek-deepgemm. + +#pragma once + +#include + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace candidate_scan_scan { + +using namespace deep_gemm; + +inline constexpr uint32_t kEmitChunkBlocks = 256; +inline constexpr uint32_t kEmitLaneSlots = 18; +inline constexpr uint32_t kGateStride = 64; +inline constexpr uint32_t kMathRegisters = 240; +inline constexpr uint32_t kSpecializedRegisters = 24; +inline constexpr uint32_t kUmmaStages = 2; +inline constexpr uint32_t kSparseRefreshNs = 512; +inline constexpr uint32_t kSparseRefreshIdleNs = 2048; +// A ring slot can only equal this pattern if a +NaN score passed the gate, +// which both gate comparisons reject; the refresher treats it as "empty". +inline constexpr uint32_t kRingSentinel = 0xffffffffu; +// Gate updates only need to keep pace with the 65,536-token flush windows, +// so the refresher can poll orders of magnitude slower than the dynamic +// histogram daemon and stay invisible to the math warps' issue ports. +inline constexpr uint32_t kRingActiveNs = 2048; +inline constexpr uint32_t kRingIdleNs = 2048; + +// Production candidate ABI: a six-byte global record. cand_val stores the +// low 16 bits of an ascending, sign-aware FP32 high24 score code; cand_idx +// stores its high eight score bits above the exact 20-bit KV index. The local +// ring uses one uint32_t containing high24 score bits and an eight-bit +// block-in-window. +using CandidateValue = uint16_t; +constexpr uint32_t kCandidateIndexBits = 20; +constexpr uint32_t kCandidateIndexMask = (1u << kCandidateIndexBits) - 1u; + +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_ordered_fp32_code(const float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); +} + +CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) { + return candidate_ordered_fp32_code(value) >> 8; +} + +// Descending full-width ordering: smaller unsigned keys correspond to larger +// FP32 values. +CUTLASS_DEVICE uint32_t candidate_descending_fp32_code(const float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? bits : (~bits & 0x7fffffffu); +} + +// High eight bits of a sign-aware FP16 key. +// Keeping 256 coarse buckets makes the online histogram small enough to live +// beside the score pipeline while retaining much more resolution than the +// sign/exponent byte of a raw FP32 key. +CUTLASS_DEVICE uint32_t candidate_fixed_half_bucket(const float value) { + const uint32_t fp32_bits = __float_as_uint(value); + const __half half_value = __float2half(value); + uint16_t bits = __half_as_ushort(half_value) & 0x7fffu; + bits |= static_cast((fp32_bits >> 16) & 0x8000u); + bits = (bits & 0x8000u) ? bits : static_cast(~bits & 0x7fffu); + return static_cast(bits >> 8); +} + +struct OnlineFixedCandidate { + uint32_t bucket; + uint32_t payload; +}; + +CUTLASS_DEVICE OnlineFixedCandidate candidate_online_fixed(const float value) { + const uint32_t bucket = candidate_fixed_half_bucket(value); + const uint32_t raw_desc24 = candidate_descending_fp32_code(value) >> 8; + // Experimental six-byte payload: the coarse bucket is explicit, while + // the remaining 16 bits retain the raw-score order inside that bucket. + return {bucket, (bucket << 16) | (raw_desc24 & 0xffffu)}; +} + +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 float candidate_decode_score(const CandidateValue value, + const int32_t packed_idx) { + const uint32_t code = candidate_load_score_code(value, packed_idx); + const uint32_t ordered = code << 8; + const uint32_t bits = + (ordered & 0x80000000u) ? (ordered ^ 0x80000000u) : ~ordered; + return __uint_as_float(bits); +} + +CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { + return static_cast(static_cast(packed_idx) & + kCandidateIndexMask); +} + +CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, + int32_t* index_dst, const float bq, + const uint32_t kv_index) { + const uint32_t payload = candidate_fp24_code(bq); + __stcs(value_dst, static_cast(payload)); + __stcs(index_dst, + static_cast(candidate_pack_index(payload, kv_index))); +} + +CUTLASS_DEVICE void store_candidate_payload(CandidateValue* value_dst, + int32_t* index_dst, + const uint32_t payload, + const uint32_t kv_index) { + __stcs(value_dst, static_cast(payload)); + __stcs(index_dst, + static_cast(candidate_pack_index(payload, kv_index))); +} + +CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst, + int32_t* index_dst, + const uint32_t payload, + const uint32_t kv_index) { + store_candidate_payload(value_dst, index_dst, payload, kv_index); +} + +template +CUTLASS_GLOBAL __launch_bounds__( + kNumSpecializedThreads + kNumMathThreads, + 1) void sm100_candidate_scan(const uint32_t seq_len, + const uint32_t seq_len_kv, + uint32_t* cu_seq_len_k_start, + uint32_t* cu_seq_len_k_end, + const float* __restrict__ origin, // [seq_len] + const float* __restrict__ inv_delta, // [seq_len] + int32_t* __restrict__ th_bucket, // [seq_len] + int32_t* __restrict__ bcount, // [seq_len, + // num_buckets] + const uint32_t num_buckets, + const uint32_t topk, + const uint32_t refresh_every, + const uint32_t num_kv_splits, + const uint32_t + probe_group, // compacted-space group size + // (pstp-1)*64; 0 = no probe + // compaction (identity map) + const uint64_t + probe_magic, // ceil(2^42/probe_group): + // exact div via mul-shift + const uint32_t + probe_add_max, // npage*64 cap for the map + CandidateValue* __restrict__ cand_val, // [seq_len, + // cand_cap] + int32_t* __restrict__ cand_idx, // [seq_len, + // cand_cap] + int32_t* __restrict__ cand_cnt, // [seq_len] + const uint32_t cand_cap, + const __grid_constant__ cute::TmaDescriptor + tensor_map_q, + const __grid_constant__ cute::TmaDescriptor + tensor_map_kv, + // fp8: per-token fp32 dequant scales. fp4: the + // UE8M0 SF-KV int32 stream (same 1-D shape). + const __grid_constant__ cute::TmaDescriptor + tensor_map_kv_scales, + const __grid_constant__ cute::TmaDescriptor + tensor_map_weights, + // fp4 only: per-(token, head) packed UE8M0 + // SF-Q int32 stream; fp8 callers pass any map + // (unread). + const __grid_constant__ cute::TmaDescriptor + tensor_map_sf_q) { + const auto num_q_blocks = math::ceil_div(seq_len, BLOCK_Q); + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto warpgroup_idx = warp_idx / 4; + const auto lane_idx = ptx::get_lane_idx(); + constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; + constexpr uint32_t kNumMathWarps = kNumMathThreads / 32; + constexpr uint32_t kNumUmmaStages = kUmmaStages; + constexpr uint32_t kNumUmmaBuffers = kNumMathWarpGroups * kNumUmmaStages; + // fp4 replaces the per-warpgroup UMMA double buffers with a shared + // 3-stage TMEM accumulator pool (the SF columns take the freed space); + constexpr uint32_t kNumAccumBufs = kNumUmmaBuffers; + // The ring refresher reloads more often: a flooding row keeps emitting + // at the stale gate for a full stride, so the flood integral scales with + // the reload cadence. The reload is one ld.shared.v4 per stride. + constexpr uint32_t kActiveGateStride = + kOnlineFixedBuckets ? 4u : (kRingRefresh ? 8u : kGateStride); + constexpr uint32_t kActiveEmitChunkBlocks = + kOnlineFixedBuckets ? 32u : kEmitChunkBlocks; + static_assert( + !(kOnlineFixedBuckets && kStaticHotGate), + "online fixed buckets and static HOT gate are separate A/B paths"); + static_assert(!kStaticHotNoHist || kStaticHotGate, + "static HOT no-hist requires the static HOT gate"); + static_assert(!kStaticHotExactGate || kStaticHotNoHist, + "exact FP24 gate requires the static HOT no-hist path"); + static_assert( + !kRingRefresh || + (kStaticHotNoHist && !kStaticHotExactGate && !kOnlineFixedBuckets), + "ring refresher v1 covers only the plain bucket-gate no-hist writer"); + DG_STATIC_ASSERT((BLOCK_Q == 4 || BLOCK_Q == 2) && kNumMathWarps == 8, + "CandidateScan requires BLOCK_Q in {4 (H=32), 2 (H=64)} and " + "8 math warps"); + DG_STATIC_ASSERT( + BLOCK_KV == 256, + "the local ring infers the KV tile coordinate from its owner lane"); + DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, + "Invalid threads"); + + if (warp_idx == kSpecWarpStart) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_kv); + cute::prefetch_tma_descriptor(&tensor_map_kv_scales); + cute::prefetch_tma_descriptor(&tensor_map_weights); + } + + // fp4 packs two e2m1 elements per byte; SF regions are UE8M0 bytes + // packed 4-per-int32 and padded to the 128-element UTCCP granule. + static constexpr uint32_t kKvElemBytesNum = 2u; // per 2 elems + static constexpr uint32_t kNumUTCCPAlignedElems = 128; + static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = + BLOCK_Q * kNumHeads * kHeadDim * kKvElemBytesNum / 2; + static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = + BLOCK_Q * kNumHeads * sizeof(float); + static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = + BLOCK_KV * kHeadDim * kKvElemBytesNum / 2; + // fp8: one fp32 dequant scale per token. fp4: one int32 (4x UE8M0) per + // token, padded to the UTCCP granule. + static constexpr uint32_t SMEM_KV_SCALE_SIZE_PER_STAGE = + BLOCK_KV * sizeof(float); + static constexpr uint32_t SMEM_SF_Q_SIZE = 0u; + static constexpr uint32_t ALIGNED_SMEM_KV_SCALE_SIZE_PER_STAGE = + math::constexpr_align(SMEM_KV_SCALE_SIZE_PER_STAGE, 512u); + + extern __shared__ __align__(512) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_Q_SIZE_PER_STAGE % 512 == 0, "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(SMEM_WEIGHT_SIZE_PER_STAGE % 512 == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(SMEM_KV_SIZE_PER_STAGE % 512 == 0, + "Unaligned TMA swizzling"); + + constexpr uint32_t kNumAccumTmemCols = BLOCK_Q * kNumHeads * kNumAccumBufs; + constexpr uint32_t kNumTmemCols = kNumAccumTmemCols; + DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); + + auto smem_q = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + + SMEM_Q_SIZE_PER_STAGE * i); + }); + auto smem_weights = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + + SMEM_Q_SIZE_PER_STAGE * kNumQStages + + SMEM_WEIGHT_SIZE_PER_STAGE * i); + }); + auto smem_kv = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>( + smem_buffer + (SMEM_Q_SIZE_PER_STAGE * kNumQStages + + SMEM_WEIGHT_SIZE_PER_STAGE * kNumQStages + + SMEM_KV_SIZE_PER_STAGE * i)); + }); + auto smem_kv_scales = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + + SMEM_Q_SIZE_PER_STAGE * kNumQStages + + SMEM_WEIGHT_SIZE_PER_STAGE * kNumQStages + + SMEM_KV_SIZE_PER_STAGE * kNumKVStages + + ALIGNED_SMEM_KV_SCALE_SIZE_PER_STAGE * i); + }); + + auto smem_sf_q = reinterpret_cast( + reinterpret_cast(smem_kv_scales[kNumKVStages])); + auto barrier_ptr = reinterpret_cast( + reinterpret_cast(smem_sf_q) + SMEM_SF_Q_SIZE); + auto full_q_barriers = + utils::PatternVisitor([&](const uint32_t& i) { return barrier_ptr + i; }); + auto empty_q_barriers = utils::PatternVisitor( + [&](const uint32_t& i) { return barrier_ptr + (kNumQStages + i); }); + auto full_kv_barriers = utils::PatternVisitor( + [&](const uint32_t& i) { return barrier_ptr + (kNumQStages * 2 + i); }); + auto empty_kv_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return barrier_ptr + (kNumQStages * 2 + kNumKVStages + i); + }); + auto full_umma_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return barrier_ptr + (kNumQStages * 2 + kNumKVStages * 2 + i); + }); + auto empty_umma_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return barrier_ptr + + (kNumQStages * 2 + kNumKVStages * 2 + kNumAccumBufs + i); + }); + + auto sf_ready_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return barrier_ptr + + (kNumQStages * 2 + kNumKVStages * 2 + kNumAccumBufs * 2 + i); + }); + auto tmem_ptr_in_smem = reinterpret_cast( + barrier_ptr + kNumQStages * 2 + kNumKVStages * 2 + kNumAccumBufs * 2 + + kNumKVStages); + auto scan_done_flag = reinterpret_cast(tmem_ptr_in_smem + 1); + auto warpq_count = reinterpret_cast(tmem_ptr_in_smem + 4); + // Chunked emit does not use the legacy warp-queue count words. Reuse four + // of them as a CTA-local Gate4 mailbox. Values are positive float edge + // bit-patterns, so unsigned atomicMin is exactly a monotonic gate tighten. + auto sparse_gate_bits = reinterpret_cast(warpq_count); + auto emit_smem_records = reinterpret_cast( + warpq_count + (kRingRefresh ? 4u : kNumMathWarps * BLOCK_Q)); + auto smem_hist = reinterpret_cast( + emit_smem_records + kNumMathWarps * BLOCK_Q * kEmitLaneSlots * 32u); + // Ring-refresher scratch sits after the refresher's histogram (which + // reuses the dynamic-path histogram region): per-(math warp, row) flush + // sequence words, the spare warps' shadow copies, and per-lane + // consumed-slot cursors. + auto ring_seq = reinterpret_cast(smem_hist + BLOCK_Q * 256); + auto ring_shadow = ring_seq + kNumMathWarps * BLOCK_Q; + auto ring_pending = + reinterpret_cast(ring_shadow + kNumMathWarps * BLOCK_Q); + auto ring_stale = ring_pending + BLOCK_Q; + auto ring_last = ring_stale + BLOCK_Q; + auto ring_progress = reinterpret_cast(ring_last + BLOCK_Q); + + DG_STATIC_ASSERT( + kNumSpecializedThreads % 128 == 0 and kNumSpecializedThreads >= 64, + "Invalid threads"); + if (warp_idx == kSpecWarpStart and cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumQStages; ++i) { + full_q_barriers[i]->init(1); + empty_q_barriers[i]->init(kNumMathThreads + 32); + } +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + full_kv_barriers[i]->init(1); + empty_kv_barriers[i]->init(kNumMathThreads); + } + if constexpr (!kStaticHotNoHist || kRingRefresh) *scan_done_flag = 0; + cutlass::arch::fence_barrier_init(); + } + if (warp_idx == kSpecWarpStart + 1) { + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumAccumBufs; ++i) { + full_umma_barriers[i]->init(1); + empty_umma_barriers[i]->init(128); + } + cutlass::arch::fence_barrier_init(); + } + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + // One CTA owns each row. The hot-only path emits no sample seeds and scans + // the complete KV range, so its exact histogram starts at zero in SMEM. + // CANDSCAN_STATIC_HOT_NOHIST_AB removes both this clear and every update; + // a post-scan kernel rebuilds metadata from the compact candidate slab. + if constexpr (!kStaticHotNoHist) { + for (uint32_t idx = threadIdx.x; idx < BLOCK_Q * num_buckets; + idx += blockDim.x) { + smem_hist[idx] = 0; + } + } + if constexpr (kRingRefresh) { + for (uint32_t idx = threadIdx.x; + idx < kNumMathWarps * BLOCK_Q * kEmitLaneSlots * 32u; + idx += blockDim.x) { + emit_smem_records[idx] = kRingSentinel; + } + // Preload the seed sample's full bucket histogram as the daemon's + // refresh base. The sample records are genuine row records in the + // same (origin, inv) bucket space and the exact-once scan starts after + // the sampled prefix, so ring harvests never recount them; base plus + // harvested stays a subset of the row's true records and every + // publish remains a provable upper bound. + for (uint32_t idx = threadIdx.x; idx < BLOCK_Q * num_buckets; + idx += blockDim.x) { + const uint32_t r = idx / num_buckets; + const uint32_t row_q = static_cast(blockIdx.x) * BLOCK_Q + r; + smem_hist[idx] = + row_q < seq_len + ? __ldcg(bcount + static_cast(row_q) * num_buckets + + (idx - r * num_buckets)) + : 0; + } + for (uint32_t idx = threadIdx.x; idx < kNumMathWarps * BLOCK_Q; + idx += blockDim.x) { + ring_seq[idx] = 0u; + ring_shadow[idx] = 0u; + } + for (uint32_t idx = threadIdx.x; idx < BLOCK_Q; idx += blockDim.x) { + ring_pending[idx] = 0; + ring_last[idx] = 0; + ring_stale[idx] = blockIdx.x * BLOCK_Q + idx < seq_len ? 0 : 2; + } + for (uint32_t idx = threadIdx.x; idx < kNumMathWarps * BLOCK_Q * 32u; + idx += blockDim.x) { + ring_progress[idx] = 0u; + } + } + if (threadIdx.x < BLOCK_Q) { + const uint32_t row_q = + static_cast(blockIdx.x) * BLOCK_Q + threadIdx.x; + const uint32_t row = min(row_q, seq_len - 1); + const int gate = __ldcg(th_bucket + row); + if constexpr (kStaticHotExactGate) { + // Exact compaction stores the ordered-FP32 pivot edge in this + // int32 slot. Preserve its bits so the unsigned strict comparison + // admits only later records from a better high24 class. + ptx::st_shared(sparse_gate_bits + threadIdx.x, + static_cast(gate)); + } else { + ptx::st_shared(sparse_gate_bits + threadIdx.x, + __float_as_uint(static_cast(gate + 1))); + } + } + __syncthreads(); + + constexpr uint32_t kNumSpecializedRegisters = kSpecializedRegisters; + constexpr uint32_t kNumMathRegisters = kMathRegisters; + + // V1 KV-split scheduling: blockIdx.x = q-block (one per CTA), blockIdx.y = + // contiguous KV sub-window. Split boundaries are BLOCK_KV-aligned. + const uint32_t block_q_idx = blockIdx.x; + const uint32_t kv_split = blockIdx.y; + uint32_t seq_k_start[BLOCK_Q], seq_k_end[BLOCK_Q]; + const auto load_schedule = + [&](const uint32_t block_q_idx) -> cute::tuple { + uint32_t start = cute::numeric_limits::max(); + uint32_t end = cute::numeric_limits::min(); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { + const auto q_idx = min(block_q_idx * BLOCK_Q + i, seq_len - 1); + seq_k_start[i] = cu_seq_len_k_start[q_idx]; + seq_k_end[i] = cu_seq_len_k_end[q_idx]; + if (block_q_idx * BLOCK_Q + i >= seq_len) { + // Padded row of a ragged final q-block: empty, aggregation-neutral. + seq_k_start[i] = seq_len_kv; + seq_k_end[i] = 0; + } + start = min(start, min(seq_k_start[i], seq_len_kv)); + end = max(end, min(seq_k_end[i], seq_len_kv)); + } + const uint32_t total_blocks = math::ceil_div(seq_len_kv, BLOCK_KV); + const uint32_t blocks_per_split = + math::ceil_div(total_blocks, num_kv_splits); + const uint32_t split_lo = kv_split * blocks_per_split * BLOCK_KV; + const uint32_t split_hi = + min((kv_split + 1) * blocks_per_split * BLOCK_KV, seq_len_kv); + start = start / 4 * 4; // TMA alignment for SF KV + if (start < split_lo) start = split_lo; + if (end > split_hi) end = split_hi; + const uint32_t nkv = + (end > start) ? math::ceil_div(end - start, BLOCK_KV) : 0; + return {start, nkv}; + }; + + const auto get_kv_pipeline = + [&](const uint32_t& kv_block_idx) -> cute::tuple { + return {kv_block_idx % kNumKVStages, (kv_block_idx / kNumKVStages) & 1}; + }; + + constexpr uint32_t UMMA_M = 128; + constexpr uint32_t UMMA_K = 32 / sizeof(cutlass::float_e4m3_t); + constexpr uint32_t UMMA_N = BLOCK_Q * kNumHeads; + + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc(); + + if (cute::elect_one_sync()) { + if (block_q_idx < num_q_blocks) { + // Q + weights once for this q-block. Inner extents are BYTES + // for the uint8 maps, so fp4's packed rows halve them. + constexpr uint32_t kKvInnerBytes = kHeadDim * kKvElemBytesNum / 2; + tma::copy( + &tensor_map_q, full_q_barriers[0], smem_q[0], 0, + block_q_idx * BLOCK_Q * kNumHeads); + tma::copy(&tensor_map_weights, + full_q_barriers[0], smem_weights[0], 0, + block_q_idx * BLOCK_Q); + full_q_barriers[0]->arrive_and_expect_tx(SMEM_Q_SIZE_PER_STAGE + + SMEM_WEIGHT_SIZE_PER_STAGE); + + CUTE_TIE_DECL(load_schedule(block_q_idx), kv_start, num_kv_blocks); + for (uint32_t kv_block_idx = 0; kv_block_idx < num_kv_blocks; + ++kv_block_idx) { + CUTE_TIE_DECL(get_kv_pipeline(kv_block_idx), kv_stage_idx, kv_phase); + empty_kv_barriers[kv_stage_idx]->wait(kv_phase ^ 1); + + tma::copy( + &tensor_map_kv, full_kv_barriers[kv_stage_idx], + smem_kv[kv_stage_idx], 0, kv_start + kv_block_idx * BLOCK_KV); + tma::copy(&tensor_map_kv_scales, + full_kv_barriers[kv_stage_idx], + smem_kv_scales[kv_stage_idx], + kv_start + kv_block_idx * BLOCK_KV, 0); + full_kv_barriers[kv_stage_idx]->arrive_and_expect_tx( + SMEM_KV_SIZE_PER_STAGE + SMEM_KV_SCALE_SIZE_PER_STAGE); + } + } + } + } else if (warp_idx == kSpecWarpStart + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(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); + + if (block_q_idx < num_q_blocks) { + CUTE_TIE_DECL(load_schedule(block_q_idx), kv_start, num_kv_blocks); + full_q_barriers[0]->wait(0); + + for (uint32_t kv_block_idx = 0; kv_block_idx < num_kv_blocks; + ++kv_block_idx) { + const uint32_t kvg = kv_block_idx; + CUTE_TIE_DECL(get_kv_pipeline(kvg), kv_stage_idx, kv_phase); + full_kv_barriers[kv_stage_idx]->wait(kv_phase); + + DG_STATIC_ASSERT(BLOCK_KV == kNumMathThreads, "Invalid block size"); + DG_STATIC_ASSERT(kHeadDim % UMMA_K == 0, "Invalid head dim"); + // Round-robin over kNumUmmaStages accumulators. A stage is + // reused every kNumUmmaStages tiles, so its phase toggles at + // that rate, not every tile. + const uint32_t umma_stage = kvg % kNumUmmaStages; + const uint32_t umma_phase = (kvg / kNumUmmaStages) & 1; +#pragma unroll + for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { + const uint32_t buf = i * kNumUmmaStages + umma_stage; + empty_umma_barriers[buf]->wait(umma_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_kv[kv_stage_idx], i * UMMA_M, k * UMMA_K); + auto b_desc = + mma::sm100::make_umma_desc(smem_q[0], 0, k * UMMA_K); + cute::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, buf * UMMA_N, k, + runtime_instr_desc); + } + cutlass::arch::umma_arrive( + reinterpret_cast(full_umma_barriers[buf])); + } + } + empty_q_barriers[0]->arrive(); + } + } else if (warp_idx == kSpecWarpStart + 2 or warp_idx == kSpecWarpStart + 3) { + // Spare-warp threshold-refresh daemon. Keeping refresh off the math + // warps avoids placing histogram scan latency on their critical path. + cutlass::arch::warpgroup_reg_dealloc(); + + if (block_q_idx < num_q_blocks) { + const uint32_t spare_id = warp_idx - (kSpecWarpStart + 2); // 0 or 1 + const auto refresh_row = [&](const uint32_t row, + const bool publish_boundary_counts) { + if (row >= seq_len) return false; + const uint32_t local_row = row - block_q_idx * BLOCK_Q; + const int32_t* srow = + smem_hist + (row - block_q_idx * BLOCK_Q) * num_buckets; + const int current_gate = + min(max(static_cast(__uint_as_float( + ptx::ld_shared(sparse_gate_bits + local_row))) - + 1, + 0), + static_cast(num_buckets) - 1); + // Only buckets strictly below the published gate can tighten + // it. If they contain fewer than topk entries, keep the + // current gate and avoid scanning the dead upper tail. + const uint32_t search_buckets = static_cast(current_gate); + int found = current_gate; + int carry = 0; + int found_lt = 0; + int found_eq = 0; + bool done = false; + for (uint32_t base = 0; base < search_buckets && !done; base += 32) { + uint32_t b = base + lane_idx; + int v = 0; + if (b < search_buckets) { + // In the single-split fast path smem already contains + // the initial global histogram plus every scan hit. + v = srow[b]; + } + // Histogram entries are nonnegative. Most 32-bucket + // groups remain strictly below K, so reject them with + // one REDUX and reserve the five dependent prefix + // shuffles for the single group that crosses K. + const int group_sum = __reduce_add_sync(0xffffffffu, v); + if (carry + group_sum < static_cast(topk)) { + carry += group_sum; + continue; + } + int prefix = v; +#pragma unroll + for (int off = 1; off < 32; off <<= 1) { + int nsh = __shfl_up_sync(0xffffffffu, prefix, off); + if (static_cast(lane_idx) >= off) prefix += nsh; + } + int incl = carry + prefix; + bool hit = (b < search_buckets) && (incl >= static_cast(topk)) && + (incl - v < static_cast(topk)); + unsigned hm = __ballot_sync(0xffffffffu, hit); + if (hm) { + const int hit_lane = __ffs(hm) - 1; + found = static_cast(base) + hit_lane; + found_lt = __shfl_sync(0xffffffffu, incl - v, hit_lane); + found_eq = __shfl_sync(0xffffffffu, v, hit_lane); + done = true; + } else { + carry += __shfl_sync(0xffffffffu, prefix, 31); + } + } + if (!done) { + // No lower bucket reached K, so the current gate remains + // the boundary. `carry` is exactly count(bucket(found + 1)); + if (done) { + th_bucket[row] = found; + if constexpr (!kStaticHotGate) { + atomicMin(sparse_gate_bits + local_row, edge); + } + } + if (publish_boundary_counts && num_buckets >= 3) { + // Reuse the dead histogram row as selector metadata. + int32_t* meta = bcount + static_cast(row) * num_buckets; + meta[0] = ~found; + meta[1] = found_lt; + meta[2] = found_eq; + } + } + return done; + }; + // CANDSCAN_STATIC_HOT_AB: the external HOT sample provides a + // fixed scan-time gate. Avoid every intermediate prefix scan and + // publish the tight full-scan boundary certificate only once. + if constexpr (kRingRefresh) { + // Ring-readback refresher: histogram already-emitted smem + // ring records (word-atomic reads) and tighten the CTA gate + // mailbox with zero additions to the math-warp hit path. + // Counting a SUBSET of emitted records keeps every published + // edge a provable lower bound on the row's current K-th + // best, and the per-(math warp, row) flush seqlock discards + // any read that raced a drain so no record counts twice. + // + // The production daemon is always active. Quiet passes use a + // fixed 2048 ns sleep; keeping this compile-time removes the + // former launch-time daemon/pacing A/B branches. + uint32_t pass = 0; + while (*scan_done_flag == 0) { + if (ring_stale[spare_id] >= 2 && + ring_stale[BLOCK_Q == 4 ? spare_id + 2 : spare_id] >= 2) { + break; + } + const uint32_t local_row = + BLOCK_Q == 4 ? spare_id + ((pass & 1u) << 1) : spare_id; + ++pass; + uint32_t fresh_total = 0; + if (block_q_idx * BLOCK_Q + local_row < seq_len && + ring_stale[local_row] < 2) { + for (uint32_t w = 0; w < kNumMathWarps; ++w) { + const uint32_t pair = w * BLOCK_Q + local_row; + const uint32_t s0 = + *reinterpret_cast(ring_seq + pair); + if (s0 != ring_shadow[pair]) { + // New flush epoch: the drained column is + // sentinel again and refills from zero. + ring_progress[pair * 32u + lane_idx] = 0u; + __syncwarp(); + if (lane_idx == 0) ring_shadow[pair] = s0; + __syncwarp(); + } + // Deep harvest: drain the lane's whole staged + // column per visit (up to kEmitLaneSlots) instead + // of one record per pass. The flood-time daemon + // bottleneck is harvest bandwidth, not publish + // cadence; reads and smem atomics only, so the + // math warps' issue ports stay untouched. + uint32_t progress = ring_progress[pair * 32u + lane_idx]; + uint32_t took = 0; + while (progress < kEmitLaneSlots) { + const uint32_t record = *reinterpret_cast( + emit_smem_records + + (pair * kEmitLaneSlots + progress) * 32u + lane_idx); + if (record == kRingSentinel) break; + if (*reinterpret_cast(ring_seq + pair) != + s0) { + // Raced a flush drain: discard this + // read; the epoch reset above replays + // the refilled column next visit. + break; + } + // Invert the truncated fp24 code back to + // the affine bucket value. Truncation only + // rounds toward better scores and never + // crosses an integer boundary (codes and + // integer edges share the 0x100-aligned + // grid), so unshifted bins support a + // pad-free published edge. The comparison + // ladder (not fmin/fmax) routes NaN and + // out-of-range scores to the top bin, which + // the publish guard below never counts. + uint32_t fbits = record & 0xffffff00u; + fbits = (fbits & 0x80000000u) ? (fbits ^ 0x80000000u) : ~fbits; + const float value = __uint_as_float(fbits); + uint32_t bucket; + if (value < 0.0f) { + bucket = 0u; + } else if (value < static_cast(num_buckets - 1)) { + bucket = static_cast(value); + } else { + bucket = num_buckets - 1; + } + atomicAdd(smem_hist + local_row * num_buckets + bucket, 1); + ++progress; + ++took; + } + if (took != 0u) { + ring_progress[pair * 32u + lane_idx] = + static_cast(progress); + } + const uint32_t fresh = __reduce_add_sync(0xffffffffu, took); + if (lane_idx == 0 && fresh != 0u) { + ring_pending[local_row] += static_cast(fresh); + } + fresh_total += fresh; + } + __syncwarp(); + // The seed base already carries K records, so each + // topk/16 batch of fresh evidence can move the + // boundary. + const int32_t pend = ring_pending[local_row]; + if (pend - ring_last[local_row] >= + static_cast(topk) / 16) { + if (lane_idx == 0) ring_last[local_row] = pend; + __syncwarp(); + // Smallest bucket prefix holding at least K ring + // records. The truncated codes and the ordered + // codes of integers <= 256 both sit on the + // 0x100-aligned grid, so a counted record's + // true code is < ord(found + 1) exactly and the + // bare bucket upper boundary is a provably safe + // published edge — no truncation pad needed. + int carry = 0; + int found = -1; + for (uint32_t bbase = 0; bbase < num_buckets && found < 0; + bbase += 32) { + const uint32_t b = bbase + lane_idx; + const int v = b < num_buckets + ? smem_hist[local_row * num_buckets + b] + : 0; + const int group_sum = __reduce_add_sync(0xffffffffu, v); + if (carry + group_sum < static_cast(topk)) { + carry += group_sum; + continue; + } + int prefix = v; +#pragma unroll + for (int off = 1; off < 32; off <<= 1) { + const int nsh = __shfl_up_sync(0xffffffffu, prefix, off); + if (static_cast(lane_idx) >= off) prefix += nsh; + } + const bool hit = carry + prefix >= static_cast(topk); + const unsigned hm = __ballot_sync(0xffffffffu, hit); + if (hm) { + found = static_cast(bbase) + __ffs(hm) - 1; + } else { + carry += __shfl_sync(0xffffffffu, prefix, 31); + } + } + if (lane_idx == 0) { + bool published = false; + // The top bin collects NaN and clamped + // out-of-range records, so an edge may only + // be published from a strictly lower bin. + if (found >= 0 && found < static_cast(num_buckets) - 1) { + const uint32_t edge = + __float_as_uint(static_cast(found + 1)); + published = + edge < atomicMin(sparse_gate_bits + local_row, edge); + } + // Early warm-start scans that fail to + // tighten only lack fresh evidence, not + // convergence; they must not retire the + // daemon. + ring_stale[local_row] = + published ? 0 + : (pend >= static_cast(topk) + ? ring_stale[local_row] + 1 + : ring_stale[local_row]); + } + __syncwarp(); + } + } + // Quiet pacing. The CTA retires promptly once the math + // warps finish; bursts skip pacing and chain straight + // into the next pass. + if (fresh_total < 64u) { + uint32_t slept = 0; + while (slept < kRingIdleNs && *scan_done_flag == 0) { + __nanosleep(kRingActiveNs); + slept += kRingActiveNs; + } + } + } + } else if constexpr (kStaticHotNoHist) { + // Scratch path: no in-scan histogram has to be finalized, so + // both spare refresh warps can retire immediately. + } else if constexpr (kStaticHotGate) { + while (*scan_done_flag == 0) __nanosleep(kSparseRefreshIdleNs); + } else { + while (*scan_done_flag == 0) { + bool tightened = false; + for (uint32_t r = spare_id; r < BLOCK_Q; r += 2) + tightened |= refresh_row(block_q_idx * BLOCK_Q + r, false); + __nanosleep(tightened ? kSparseRefreshNs : kSparseRefreshIdleNs); + } + } + if constexpr (!kStaticHotNoHist) { + for (uint32_t r = spare_id; r < BLOCK_Q; r += 2) + refresh_row(block_q_idx * BLOCK_Q + r, true); + } + } + } else if (warp_idx < kSpecWarpStart) { + cutlass::arch::warpgroup_reg_alloc(); + + const auto math_thread_idx = warp_idx * 32 + lane_idx; + + auto tmem_load = [](auto num_elems_c, const uint32_t& tmem_addr, + float* accum) { + constexpr int N = decltype(num_elems_c)::value; + DG_STATIC_ASSERT(N == 32 or N == 64 or N == 128, + "Unsupported TMEM load size"); + using Loader = cute::conditional_t< + N == 32, cute::SM100_TMEM_LOAD_32dp32b32x, + cute::conditional_t>; + [&](cute::index_sequence) { + Loader::copy(tmem_addr, reinterpret_cast(accum)[Is]...); + }(cute::make_index_sequence{}); + cutlass::arch::fence_view_async_tmem_load(); + }; + + // Bucket comparisons use the affine score space. Raw float bit + // patterns are nonuniform across exponent ranges and cannot preserve + // this threshold contract. + float weights[BLOCK_Q][kNumHeads]; + float o_reg[BLOCK_Q], inv_reg[BLOCK_Q], vth_reg[BLOCK_Q]; + uint32_t kstart_reg[BLOCK_Q], + kspan_reg[BLOCK_Q]; // unsigned range-check trick + const unsigned FULL = 0xffffffffu; + + if (block_q_idx < num_q_blocks) { + CUTE_TIE_DECL(load_schedule(block_q_idx), kv_start, num_kv_blocks); + full_q_barriers[0]->wait(0); + +// Weights into registers (once per CTA -- the 2.5-generation win). +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; ++j) + weights[i][j] = ptx::ld_shared(smem_weights[0] + i * kNumHeads + j); + } + // Queue fill counts are warp-uniform: every lane tracks them + // redundantly in registers, so the hot emit path needs no shared + // bookkeeping or shuffle broadcast. + // Four private 8-bit counts. They reset at every fixed-size chunk, + // bounding the address dispersion of direct global record stores. + uint32_t emit_lane_counts = 0; +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { + const uint32_t rq = min(block_q_idx * BLOCK_Q + i, seq_len - 1); + if constexpr (kOnlineFixedBuckets) { + // No sample-derived affine transform in the online A/B + // path. The raw score is converted to fixed buckets in + // the epilogue below; keep these registers compile-time + // independent of the nullable origin/inv pointers. + o_reg[i] = 0.0f; + inv_reg[i] = 1.0f; + vth_reg[i] = 0.0f; + } else { + o_reg[i] = origin[rq]; + inv_reg[i] = inv_delta[rq]; + vth_reg[i] = -o_reg[i] * inv_reg[i]; + } + o_reg[i] = 0.0f; // gate closed until the first consume + kstart_reg[i] = seq_k_start[i]; + kspan_reg[i] = + seq_k_end[i] > seq_k_start[i] ? seq_k_end[i] - seq_k_start[i] : 0; + } + // Fold -inv into the register weights: the whole ReLU-weighted + // chain then accumulates directly in bucket units. 128 FMULs + // once per qb, amortized over thousands of kv blocks. + if constexpr (!kOnlineFixedBuckets) { +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; ++j) weights[i][j] *= -inv_reg[i]; + } + } + + uint32_t rs_max = 0, re_min = 0xffffffffu; +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { + rs_max = max(rs_max, kstart_reg[i]); + re_min = min(re_min, kstart_reg[i] + kspan_reg[i]); + } + + if constexpr (BLOCK_Q == 4) { + const float4 initial_gate = + ptx::ld_shared(reinterpret_cast(sparse_gate_bits)); + o_reg[0] = initial_gate.x; + o_reg[1] = initial_gate.y; + o_reg[2 < BLOCK_Q ? 2 : 0] = initial_gate.z; + o_reg[3 < BLOCK_Q ? 3 : 0] = initial_gate.w; + } else { +#pragma unroll + for (uint32_t gi = 0; gi < BLOCK_Q; ++gi) + o_reg[gi] = __uint_as_float(ptx::ld_shared(sparse_gate_bits + gi)); + } + + for (uint32_t kv_block_idx = 0; kv_block_idx < num_kv_blocks; + ++kv_block_idx) { + const uint32_t kvg = kv_block_idx; + CUTE_TIE_DECL(get_kv_pipeline(kvg), kv_stage_idx, kv_phase); + + if constexpr (!kStaticHotGate || kRingRefresh) { + if (kv_block_idx != 0 && (kv_block_idx % kActiveGateStride) == 0) { + if constexpr (BLOCK_Q == 4) { + const float4 gate = ptx::ld_shared( + reinterpret_cast(sparse_gate_bits)); + o_reg[0] = gate.x; + o_reg[1] = gate.y; + o_reg[2 < BLOCK_Q ? 2 : 0] = gate.z; + o_reg[3 < BLOCK_Q ? 3 : 0] = gate.w; + } else { +#pragma unroll + for (uint32_t gi = 0; gi < BLOCK_Q; ++gi) + o_reg[gi] = + __uint_as_float(ptx::ld_shared(sparse_gate_bits + gi)); + } + } + } + full_kv_barriers[kv_stage_idx]->wait(kv_phase); + + const float scale_kv = + ptx::ld_shared(smem_kv_scales[kv_stage_idx] + math_thread_idx); + + const uint32_t umma_buf = + warpgroup_idx * kNumUmmaStages + (kvg % kNumUmmaStages); + const uint32_t umma_wait_phase = (kvg / kNumUmmaStages) & 1; + const auto tmem_start = umma_buf * UMMA_N; + full_umma_barriers[umma_buf]->wait(umma_wait_phase); + ptx::tcgen05_after_thread_sync(); + + empty_kv_barriers[kv_stage_idx]->arrive(); + + const auto kv_offset = + kv_start + kv_block_idx * BLOCK_KV + math_thread_idx; + DG_STATIC_ASSERT(kNumHeads % 8 == 0, "Invalid head"); + + uint32_t pass_bits = 0; + float v_row[BLOCK_Q]; + uint32_t fixed_payload_row[BLOCK_Q]; + uint32_t fixed_bucket_row[BLOCK_Q]; + constexpr uint32_t kTmemRowsEff = + (128u / kNumHeads) < BLOCK_Q ? (128u / kNumHeads) : BLOCK_Q; + DG_STATIC_ASSERT(BLOCK_Q % kTmemRowsEff == 0, + "BLOCK_Q must divide evenly into TMEM load groups"); +#define CANDSCAN_SCORE_GATE(i, RC) \ + bool g; \ + if constexpr (kOnlineFixedBuckets) { \ + const float raw_score = scale_kv * (sum.x + sum.y); \ + const OnlineFixedCandidate fixed = candidate_online_fixed(raw_score); \ + fixed_payload_row[i] = fixed.payload; \ + fixed_bucket_row[i] = fixed.bucket; \ + g = fixed.bucket < static_cast(o_reg[i]); \ + } else { \ + const float bq = fmaf(scale_kv, sum.x + sum.y, vth_reg[i]); \ + if constexpr (kStaticHotExactGate) { \ + const uint32_t ordered = candidate_ordered_fp32_code(bq); \ + v_row[i] = __uint_as_float(ordered); \ + g = ordered < __float_as_uint(o_reg[i]); \ + } else { \ + v_row[i] = bq; \ + g = __float_as_int(bq) < __float_as_int(o_reg[i]); \ + } \ + } \ + if constexpr (RC) g = g and ((kv_offset - kstart_reg[i]) < kspan_reg[i]); \ + pass_bits |= g ? (1u << i) : 0u; + const uint32_t kv_base = kv_start + kv_block_idx * BLOCK_KV; + const bool interior = + (kv_base >= rs_max) && (kv_base + BLOCK_KV <= re_min); +#define CANDSCAN_SCORE_ROWS(RANGE_CHECK) \ + _Pragma("unroll") for (uint32_t pr = 0; pr < BLOCK_Q / kTmemRowsEff; ++pr) { \ + float accum2[kNumHeads * kTmemRowsEff]; \ + tmem_load(cute::Int{}, \ + tmem_start + pr * kTmemRowsEff * kNumHeads, accum2); \ + if (pr == BLOCK_Q / kTmemRowsEff - 1) { \ + ptx::tcgen05_before_thread_sync(); \ + empty_umma_barriers[umma_buf]->arrive(); \ + } \ + _Pragma("unroll") for (uint32_t k = 0; k < kTmemRowsEff; ++k) { \ + const uint32_t i = pr * kTmemRowsEff + k; \ + const float* accum = accum2 + k * kNumHeads; \ + auto sum_0 = make_float2(0, 0); \ + auto sum_1 = make_float2(0, 0); \ + const auto transform = [&](const uint32_t& j, const float2& sum) { \ + auto a = make_float2(fmaxf(accum[j], 0), fmaxf(accum[j + 1], 0)); \ + auto b = make_float2(weights[i][j], weights[i][j + 1]); \ + return __ffma2_rn(a, b, sum); \ + }; \ + _Pragma("unroll") for (uint32_t j = 0; j < kNumHeads; j += 4) { \ + sum_0 = transform(j, sum_0); \ + sum_1 = transform(j + 2, sum_1); \ + } \ + auto sum = __fadd2_rn(sum_0, sum_1); \ + CANDSCAN_SCORE_GATE(i, RANGE_CHECK) \ + } \ + } + if (interior) { + CANDSCAN_SCORE_ROWS(false) + } else { + CANDSCAN_SCORE_ROWS(true) + } +#undef CANDSCAN_SCORE_ROWS +#undef CANDSCAN_SCORE_GATE + + if (pass_bits != 0) { +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { + if ((pass_bits >> i) & 1u) { + uint32_t candidate_score_bits; + uint32_t candidate_bucket; + if constexpr (kOnlineFixedBuckets) { + candidate_score_bits = fixed_payload_row[i] << 8; + candidate_bucket = fixed_bucket_row[i]; + } else { + if constexpr (kStaticHotExactGate) { + // v_row carries the already-computed + // ordered32 key in this specialization. + // The ring's low byte belongs exclusively + // to local_block; clear the discarded FP32 + // low bits before packing that coordinate. + candidate_score_bits = + __float_as_uint(v_row[i]) & 0xffffff00u; + candidate_bucket = 0; + } else { + const float x = v_row[i]; + candidate_score_bits = candidate_fp24_code(x) << 8; + const int candidate_braw = static_cast(x); + candidate_bucket = + static_cast(max(candidate_braw, 0)); + } + } + const uint32_t count = (emit_lane_counts >> (i * 8)) & 0xffu; + if (count < kEmitLaneSlots) { + const uint32_t pos = + ((warp_idx * BLOCK_Q + i) * kEmitLaneSlots + count) * 32u + + lane_idx; + const uint32_t local_block = + kv_block_idx % kActiveEmitChunkBlocks; + const uint32_t record = candidate_score_bits | local_block; + emit_smem_records[pos] = record; + emit_lane_counts += 1u << (i * 8); + } else { + // A skewed lane can overflow its small local + // quota without losing candidates. This slow + // path writes directly to the final buffer. + const uint32_t row_q = block_q_idx * BLOCK_Q + i; + const int out = atomicAdd(cand_cnt + row_q, 1); + if (out < static_cast(cand_cap)) { + uint32_t kvo = kv_offset; + if (probe_group != 0) { + const uint32_t sup = static_cast( + (static_cast(kvo) * probe_magic) >> 42); + kvo += min((sup + 1) * 64u, probe_add_max); + } + const uint64_t out_base = + static_cast(row_q) * cand_cap; + store_candidate_record(&cand_val[out_base + out], + &cand_idx[out_base + out], + candidate_score_bits >> 8, kvo); + } + } + + // A passing positive bq is below the float edge + // (gate + 1), hence already < num_buckets. + if constexpr (!kStaticHotNoHist) { + atomicAdd(smem_hist + i * num_buckets + + static_cast(candidate_bucket), + 1); + } + } + } + } + + if (((kv_block_idx + 1) % kActiveEmitChunkBlocks) == 0 || + kv_block_idx + 1 == num_kv_blocks) { + // Reserve one contiguous final-buffer segment per + // (source warp, row, chunk), then have each lane copy its + // private records into that segment. This keeps all + // collectives and the returning global atomic off the + // ordinary-hit path while deleting the rectangular global + // chunk workspace and the post-scan compactor. + // Keep the four row reservations parallel and early so + // their L2 round trips overlap the prefix work. Replace + // four independent 32-lane scans (20 shuffles) with two + // scans whose 16-bit fields carry two rows each (10 + // shuffles). A row prefix is at most 32*18=576, so fields + // cannot carry into one another. + int my_row_base = 0; + uint32_t my_row_total = 0; +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { + const uint32_t count = (emit_lane_counts >> (i * 8)) & 0xffu; + const uint32_t total = __reduce_add_sync(FULL, count); + if (lane_idx == i) my_row_total = total; + } + if (lane_idx < BLOCK_Q) { + const uint32_t row_q = block_q_idx * BLOCK_Q + lane_idx; + if (row_q < seq_len && my_row_total != 0) { + my_row_base = + atomicAdd(cand_cnt + row_q, static_cast(my_row_total)); + } + } + + const uint32_t count0 = emit_lane_counts & 0xffu; + const uint32_t count1 = (emit_lane_counts >> 8) & 0xffu; + const uint32_t count2 = (emit_lane_counts >> 16) & 0xffu; + const uint32_t count3 = emit_lane_counts >> 24; + uint32_t inclusive01 = count0 | (count1 << 16); + uint32_t inclusive23 = count2 | (count3 << 16); +#pragma unroll + for (int delta = 1; delta < 32; delta <<= 1) { + const uint32_t other01 = __shfl_up_sync(FULL, inclusive01, delta); + const uint32_t other23 = __shfl_up_sync(FULL, inclusive23, delta); + if (lane_idx >= static_cast(delta)) { + inclusive01 += other01; + inclusive23 += other23; + } + } + + const uint32_t emit_window_kv_base = + kv_start + (kv_block_idx / kActiveEmitChunkBlocks) * + kActiveEmitChunkBlocks * BLOCK_KV; +#pragma unroll + for (uint32_t i = 0; i < BLOCK_Q; ++i) { + const uint32_t count = (emit_lane_counts >> (i * 8)) & 0xffu; + const uint32_t inclusive = + i < 2 ? ((inclusive01 >> ((i & 1u) * 16)) & 0xffffu) + : ((inclusive23 >> ((i & 1u) * 16)) & 0xffffu); + const uint32_t offset = inclusive - count; + const uint32_t row_q = block_q_idx * BLOCK_Q + i; + const int out_base = __shfl_sync(FULL, my_row_base, i); + + int copy_count = 0; + if (row_q < seq_len) { + copy_count = min(static_cast(count), + max(static_cast(cand_cap) - out_base - + static_cast(offset), + 0)); + } + for (int slot = 0; slot < copy_count; ++slot) { + const uint32_t local_pos = + ((warp_idx * BLOCK_Q + i) * kEmitLaneSlots + slot) * 32u + + lane_idx; + const uint32_t record = emit_smem_records[local_pos]; + if constexpr (kRingRefresh) { + emit_smem_records[local_pos] = kRingSentinel; + } + const int out = out_base + static_cast(offset) + slot; + const uint32_t record_payload = (record & 0xffffff00u) >> 8; + uint32_t kvo = emit_window_kv_base + + ((record & 0xffu) * BLOCK_KV) + math_thread_idx; + if (probe_group != 0) { + const uint32_t sup = static_cast( + (static_cast(kvo) * probe_magic) >> 42); + kvo += min((sup + 1) * 64u, probe_add_max); + } + const uint64_t out_pos = + static_cast(row_q) * cand_cap + out; + store_candidate_record(&cand_val[out_pos], &cand_idx[out_pos], + record_payload, kvo); + } + if constexpr (kRingRefresh) { + // A capped row copies fewer slots than it + // holds; the leftovers must not leak into + // the next flush epoch. + for (int slot = copy_count; slot < static_cast(count); + ++slot) { + emit_smem_records[((warp_idx * BLOCK_Q + i) * kEmitLaneSlots + + slot) * + 32u + + lane_idx] = kRingSentinel; + } + } + } + if constexpr (kRingRefresh) { + // Publish the flush epoch: every drain store above + // is visible before the bump, and the bump is + // visible before any next-window append. + __syncwarp(); + __threadfence_block(); + if (lane_idx < BLOCK_Q) { + ring_seq[warp_idx * BLOCK_Q + lane_idx] += 1u; + } + __threadfence_block(); + } + + emit_lane_counts = 0; + } + } + + // Every actual chunk was published at its final KV block. + empty_q_barriers[0]->arrive(); + } + + // Signal the refresh daemon, then free tensor memory. + cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); + if constexpr (!kStaticHotNoHist || kRingRefresh) { + if (threadIdx.x == 0) { + __threadfence_block(); + *scan_done_flag = 1; + if constexpr (kRingRefresh) { + // Prompt-wake a daemon parked on a stage-0 parity: odd + // count so a straggler that missed transient flips + // still sees the final parity flipped. The in-wait + // exit-flag recheck makes this an optimization, not a + // correctness requirement. + asm volatile("" ::: "memory"); +#pragma unroll + for (int poke = 0; poke < 5; ++poke) { + full_kv_barriers[0]->arrive(); + } + } + } + } + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +} // namespace candidate_scan_scan diff --git a/csrc/fused_mqa_topk/upstream_select.cuh b/csrc/fused_mqa_topk/upstream_select.cuh new file mode 100644 index 000000000000..97b424c63be2 --- /dev/null +++ b/csrc/fused_mqa_topk/upstream_select.cuh @@ -0,0 +1,534 @@ +// Isolated adaptation of vLLM sampler.cu device selection routines. +#pragma once +#include +#include +#include +namespace isolated_upstream { +template +static inline __device__ uint32_t extractBinIdx(float x) { + if constexpr (step == 0) { + __half hx = __float2half(x); + uint16_t bits = __half_as_ushort(hx); + bits = (bits & 0x8000) ? bits : ~bits & 0x7fff; + return bits >> 5; + } else { + uint32_t bits = __float_as_uint(x); + bits = (bits & 0x80000000) ? bits : ~bits & 0x7fffffff; + + if constexpr (step == 1) { + return bits >> 21; + } else if constexpr (step == 2) { + return (bits >> 10) & 0x7ff; + } else if constexpr (step == 3) { + return bits & 0x3ff; + } + } +} + +template +static inline __device__ bool isPartialMatch(float x, uint32_t pattern) { + if constexpr (shift == 0) { + return true; + } + uint32_t bits = __float_as_uint(x); + bits = (bits & 0x80000000) ? bits : ~bits & 0x7fffffff; + return (bits ^ pattern) >> shift == 0; +} + +/** + * Map a Func over the input data, using vectorized load instructions if + * possible. + * + * @tparam T element type + * @tparam IdxT indexing type + * @tparam Func void (T x, IdxT idx) + * + * @param thread_rank rank of the calling thread among all participating threads + * @param num_threads number of the threads that participate in processing + * @param in the input data + * @param len the number of elements to read + * @param f the lambda taking two arguments (T x, IdxT idx) + */ +template +__device__ void vectorized_process(size_t thread_rank, size_t num_threads, + const T* in, idxT len, Func f) { + // Use dynamic WARP_SIZE from cuda_compat.h to support both + // Wave64 (MI300X/gfx942) and Wave32 (Strix Halo/gfx1151) architectures + constexpr int kWarpSize = 32; + using WideT = float4; + if constexpr (sizeof(T) >= sizeof(WideT)) { + for (idxT i = thread_rank; i < len; i += num_threads) { + f(in[i], i); + } + } else { + static_assert(sizeof(WideT) % sizeof(T) == 0); + constexpr int items_per_scalar = sizeof(WideT) / sizeof(T); + // TODO: it's UB + union { + WideT scalar; + T array[items_per_scalar]; + } wide; + + int skip_cnt = + (reinterpret_cast(in) % sizeof(WideT)) + ? ((sizeof(WideT) - reinterpret_cast(in) % sizeof(WideT)) / + sizeof(T)) + : 0; + if (skip_cnt > len) { + skip_cnt = len; + } + const WideT* in_cast = reinterpret_cast(in + skip_cnt); + const idxT len_cast = (len - skip_cnt) / items_per_scalar; + + for (idxT i = thread_rank; i < len_cast; i += num_threads) { + wide.scalar = in_cast[i]; + const idxT real_i = skip_cnt + i * items_per_scalar; +#pragma unroll + for (int j = 0; j < items_per_scalar; ++j) { + f(wide.array[j], real_i + j); + } + } + + static_assert(kWarpSize >= items_per_scalar); + // and because items_per_scalar > skip_cnt, kWarpSize > skip_cnt + // no need to use loop + if (thread_rank < skip_cnt) { + f(in[thread_rank], thread_rank); + } + // because len_cast = (len - skip_cnt) / items_per_scalar, + // len_cast * items_per_scalar + items_per_scalar > len - skip_cnt; + // and so + // len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <= + // kWarpSize no need to use loop + const idxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank; + if (remain_i < len) { + f(in[remain_i], remain_i); + } + } +} + +template +__device__ bool processHistogramStep( + const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern, + int& thresholdBinIdx, SmemOutputType& smemOutput, int* smemThresholdBinIdx, + int* smemFinalDstIdx, int* smemFinalBinSize, int* smemFoundTopKValues, + SmemFinalType& smemFinal, int stride1, int rowStart, int topK) { + // Clear the histogram. +#pragma unroll + for (int idx = threadIdx.x; idx < kNumBins; idx += kNumThreadsPerBlock) { + smemFinal.histo.data[idx] = 0; + } + + // Make sure the histogram is ready. + __syncthreads(); + + // Update pattern + constexpr auto patternShift = step < 2 ? 0 : step == 2 ? 21 : 10; + if constexpr (step == 2) { + logitPattern = static_cast(thresholdBinIdx & 0x7ff) + << patternShift; + } else if constexpr (step == 3) { + logitPattern |= static_cast(thresholdBinIdx & 0x7ff) + << patternShift; + } + + int negativeInfinityCount = 0; + auto distributeToBins = [&](float logit, int /* idx */ = 0) { + if (isPartialMatch(logit, logitPattern)) { + uint32_t binIdx = extractBinIdx(logit); + if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) { + ++negativeInfinityCount; + } else { + atomicAdd(&smemFinal.histo.data[binIdx], 1); + } + } + }; + + // Distribute the elements to the histogram bins. + if (stride1 == 1) { + vectorized_process(threadIdx.x, kNumThreadsPerBlock, logits + rowStart, + rowEnd - rowStart, distributeToBins); + } else { + for (int idx = rowStart + threadIdx.x; idx < rowEnd; + idx += kNumThreadsPerBlock) { + float logit = logits[idx * stride1]; + distributeToBins(logit, idx); + } + } + if constexpr (useTopK512Optimization) { + if (negativeInfinityCount > 0) { + const auto bin = extractBinIdx(-INFINITY); + atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount); + } + } + // Make sure the histogram is ready. + __syncthreads(); + + // Reads the value of the starting position in the smemOutput array + int lastValue = smemFoundTopKValues[0]; + + for (int round = 0; round < kNumBins / kNumThreadsPerBlock; round++) { + // Read the values from SMEM. + int idx = threadIdx.x + kNumThreadsPerBlock * round; + int binCount{0}; + binCount = smemFinal.histo.data[idx]; + + // Make sure each thread has read its value. + __syncthreads(); + + // Compute the prefix sum. + int prefixSum{0}, totalSum{0}; + using Scan = cub::BlockScan; + Scan(smemFinal.histo.scan).ExclusiveSum(binCount, prefixSum, totalSum); + + // Update the histogram with the prefix sums. + prefixSum += lastValue; + totalSum += lastValue; + smemFinal.histo.data[idx] = prefixSum; + + // Make sure the data is in shared memory. + __syncthreads(); + + // Find the last valid bin. + bool foundThreshold = false; + if (prefixSum < topK) { + int nextPrefixSum = threadIdx.x == kNumThreadsPerBlock - 1 + ? totalSum + : smemFinal.histo.data[idx + 1]; + + if (nextPrefixSum >= topK) { + smemThresholdBinIdx[0] = idx; + smemFinalBinSize[0] = nextPrefixSum - prefixSum; + foundThreshold = true; + } + } + + // Early exit: if any thread found the threshold, we can skip remaining + // rounds + if (__syncthreads_or(foundThreshold)) { + break; + } + + lastValue = totalSum; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + // The threshold bin. + thresholdBinIdx = smemThresholdBinIdx[0]; + // Resolve a -inf cutoff with exact tie emission, keeping sort padding out. + const bool finalBinFits = + smemFinalBinSize[0] <= kNumFinalItems && + !(useTopK512Optimization && step < 3 && + isPartialMatch(-INFINITY, logitPattern) && + thresholdBinIdx == extractBinIdx(-INFINITY)); + + auto processBins = [&](float logit, int idx) { + if (isPartialMatch(logit, logitPattern)) { + uint32_t binIdx = extractBinIdx(logit); + // Only write elements with binIdx < thresholdBinIdx when: + // 1. This is step 0 and the threshold bin is small enough (no step 1) + // 2. This is step >= 1 (where pattern matching filters correctly) + // This prevents duplicates when step 0 and step 1 both run. + bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1); + if (binIdx < thresholdBinIdx && shouldWriteDirectly) { + // The element is part of the top-k selection + int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); + + if constexpr (mergeBlocks) { + smemOutput[dstIdx] = indices[idx]; + } else if constexpr (multipleBlocksPerRow) { + smemOutput[dstIdx] = idx + rowStart; + reinterpret_cast(smemOutput + topK)[dstIdx] = logit; + } else { + smemOutput[dstIdx] = idx; + } + } + if constexpr (step < 3) { + // Only fill the final items for sorting if the threshold bin fits + if (binIdx == thresholdBinIdx && finalBinFits) { + int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1); + smemFinal.items.logits[dstIdx] = logit; + if constexpr (mergeBlocks) { + smemFinal.items.indices[dstIdx] = indices[idx]; + } else if constexpr (multipleBlocksPerRow) { + smemFinal.items.indices[dstIdx] = idx + rowStart; + } else { + smemFinal.items.indices[dstIdx] = idx; + } + } + } else { + if (binIdx == thresholdBinIdx) { + // The elements in the threshold bin share the same 32 bits at step 3 + int dstIdx = atomicAdd(&smemFinal.histo.data[binIdx], 1); + if (dstIdx < topK) { + if constexpr (mergeBlocks) { + smemOutput[dstIdx] = indices[idx]; + } else if constexpr (multipleBlocksPerRow) { + smemOutput[dstIdx] = idx + rowStart; + reinterpret_cast(smemOutput + topK)[dstIdx] = logit; + } else { + smemOutput[dstIdx] = idx; + } + } + } + } + } + }; + + if (stride1 == 1) { + vectorized_process(threadIdx.x, kNumThreadsPerBlock, logits + rowStart, + rowEnd - rowStart, processBins); + } else { + for (int idx = rowStart + threadIdx.x; idx < rowEnd; + idx += kNumThreadsPerBlock) { + float logit = logits[idx * stride1]; + processBins(logit, idx); + } + } + + // Make sure the elements are in shared memory. + __syncthreads(); + + // Check if we should continue to next step + return !finalBinFits; +} + +// Follows half - 11 - 11 - 10 bit iterations +// Keep the adaptive kernel's device instantiation separate from legacy calls. +template +static __device__ void topKPerRowJob(const int* indices, const float* logits, + int rowStart, int rowEnd, int* outIndices, + float* outLogits, int stride1, int topK) { + static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks); + static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024); + // The number of slots for the final pass. + static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 1536; + // The number of elements per thread for the final sort. + static constexpr int kNumFinalItemsPerThread = + kNumFinalItems / kNumThreadsPerBlock; + // The class to sort the elements during the final pass. + using FinalSort = cub::BlockRadixSort; + using FinalSortTempStorage = + std::conditional_t; + // The class to compute the inclusive prefix-sum over the histogram. + using Scan = cub::BlockScan; + + // The structure to store the final items (for the final pass). + struct FinalItems { + // Shared memory to store the indices for the final pass. + int indices[kNumFinalItems]; + // Shared memory to store the logits for the final pass. + float logits[kNumFinalItems]; + }; + + struct Histogram { + typename Scan::TempStorage scan; + int data[kNumBins]; + }; + + // Shared memory to compute the block sort. + union FinalStorage { + FinalItems items; + FinalSortTempStorage finalSort; + Histogram histo; + }; + extern __shared__ __align__(1024) unsigned char select_workspace[]; + auto& smemFinal = *reinterpret_cast(select_workspace + 8192); + + // Shared memory to store the selected indices. + // If we are processing using multiple blocks, we need to store the logits and + // indices. + extern __shared__ __align__(1024) int32_t smemOutput[]; + + // Shared memory to store the threshold bin. + __shared__ int smemThresholdBinIdx[1]; + // Shared memory counter to register the candidates for the final phase. + __shared__ int smemFinalDstIdx[1]; + // Shared memory to determine if the threshold bin fits in the final items. + __shared__ int smemFinalBinSize[1]; + // Shared memory to keep track of the top-k values found so far by the + // previous iterations + __shared__ int smemFoundTopKValues[1]; + + // The length of the row. + int rowLen = rowEnd - rowStart; + + // Shortcut if the length of the row is smaller than Top-K. Indices are not + // sorted by their corresponding logit. + if (rowLen <= topK) { + for (int rowIt = threadIdx.x; rowIt < rowLen; + rowIt += kNumThreadsPerBlock) { + if constexpr (multipleBlocksPerRow) { + outIndices[rowIt] = rowIt + rowStart; + outLogits[rowIt] = logits[rowIt + rowStart]; + } else { + outIndices[rowIt] = rowIt; + } + } + for (int rowIt = rowLen + threadIdx.x; rowIt < topK; + rowIt += kNumThreadsPerBlock) { + outIndices[rowIt] = -1; + if constexpr (multipleBlocksPerRow) { + outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX; + } + } + + return; + } + // Initialize values + if (threadIdx.x == 0) { + smemFinalDstIdx[0] = 0; + smemFoundTopKValues[0] = 0; + } + __syncthreads(); + int thresholdBinIdx = -1; + uint32_t logitPattern = 0; + + // Step 0: Process first 11 bits of half representation + bool continueToNextStep = + processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems, + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, + smemFoundTopKValues, smemFinal, stride1, rowStart, topK); + + if (continueToNextStep) { + // Step 1: Process next 11 bits + continueToNextStep = + processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems, + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, + smemFoundTopKValues, smemFinal, stride1, rowStart, topK); + } + + if (continueToNextStep) { + // Step 2: Process next 11 bits + continueToNextStep = + processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems, + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, + smemFoundTopKValues, smemFinal, stride1, rowStart, topK); + } + + if (continueToNextStep) { + // Step 3: Process last 10 bits + processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems, + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, + smemFoundTopKValues, smemFinal, stride1, rowStart, topK); + } + + if (!continueToNextStep) { + // The histogram did not proceed to the final 10 bits, therefore we need to + // sort the final items The logits of the elements to be sorted in the final + // pass. + auto insertionSort = [&]() { + // Sorting with insertion sort + auto baseIdx = smemFoundTopKValues[0]; + for (int i = threadIdx.x; i < smemFinalDstIdx[0]; + i += kNumThreadsPerBlock) { + int outIndex = 0; + auto logit = smemFinal.items.logits[i]; + for (int j = 0; j < smemFinalDstIdx[0]; j++) { + auto otherLogit = smemFinal.items.logits[j]; + if (logit < otherLogit || (logit == otherLogit && i < j)) { + outIndex++; + } + } + // Store if outIndex is in bounds + if (outIndex + baseIdx < topK) { + smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i]; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[outIndex + baseIdx] = + smemFinal.items.logits[i]; + } + } + } + }; + if constexpr (useRadixSort) { + if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) { + insertionSort(); + } else { + // Sorting with radix sort + float finalLogits[kNumFinalItemsPerThread]; + // The indices of the elements to be sorted in the final pass. + int finalIndices[kNumFinalItemsPerThread]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX; + if constexpr (useTopK512Optimization) finalIndices[ii] = -1; + } + + // Read the elements from SMEM. +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + if (srcIdx < smemFinalDstIdx[0]) { + finalLogits[ii] = smemFinal.items.logits[srcIdx]; + finalIndices[ii] = smemFinal.items.indices[srcIdx]; + } + } + // Make sure the shared memory has been read. + __syncthreads(); + + // Sort the elements. + FinalSort(smemFinal.finalSort) + .SortDescendingBlockedToStriped(finalLogits, finalIndices); + + // Copy the data back to the shared memory storage. + int baseIdx = smemFoundTopKValues[0]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + int dstIdx = baseIdx + srcIdx; + + if (dstIdx < topK) { + smemOutput[dstIdx] = finalIndices[ii]; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[dstIdx] = + finalLogits[ii]; + } + } + } + } + } else { + insertionSort(); + } + __syncthreads(); + } + + // Store to global memory. + for (int i = threadIdx.x; i < topK; i += kNumThreadsPerBlock) { + if constexpr (multipleBlocksPerRow) { + outIndices[i] = smemOutput[i]; + outLogits[i] = reinterpret_cast(smemOutput + topK)[i]; + } else { + if (stride1 == 1) { + // stride1 == 1 will use vectorized_process, which indexes already skip + // the rowStart. + outIndices[i] = smemOutput[i]; + } else { + outIndices[i] = smemOutput[i] - rowStart; + } + } + } +} + +} // namespace isolated_upstream diff --git a/csrc/ops.h b/csrc/ops.h index 1c5ddc0beb4c..a6d19252fbda 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -870,3 +870,13 @@ 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 slot_flags, + torch::stable::Tensor values, torch::stable::Tensor indices, + torch::stable::Tensor retained, torch::stable::Tensor retained_indices, + torch::stable::Tensor counts, torch::stable::Tensor gates, + torch::stable::Tensor stats, bool relative); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 1c84b53ece2c..baec9d8aa8fe 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-2048 prefill selection (SM100 + // family). Workspace tensors are caller-owned; slot_flags must be zero. + ops.def( + "fused_mqa_topk_prefill(Tensor q, Tensor k, Tensor scales, " + "Tensor weights, Tensor starts, Tensor ends, Tensor! out, " + "Tensor! slot_flags, Tensor! values, Tensor! indices, Tensor! retained, " + "Tensor! retained_indices, Tensor! counts, Tensor! gates, " + "Tensor! stats, bool relative) -> ()"); +#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( @@ -782,6 +793,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..53fd376c1d91 --- /dev/null +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -0,0 +1,309 @@ +# 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.""" + +import pytest +import torch + +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_prefill, +) +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder + +requires_fused = pytest.mark.skipif( + not torch.cuda.is_available() or not fused_mqa_topk_available(), + reason="Requires the SM100 fused MQA+top-k extension", +) + +TOPK = 2048 +LOGITS_BUDGET = 512 * 1024**2 + + +def _dense(q, k, scales, weights, starts, ends): + """Upstream dense path: budgeted DeepGEMM logits + top_k_per_row_prefill.""" + from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + + rows = q.shape[0] + out = torch.full((rows, TOPK), -1, device=q.device, dtype=torch.int32) + step = max(1, LOGITS_BUDGET // (4 * k.shape[0])) + for i in range(0, rows, step): + j = min(rows, i + step) + logits = fp8_fp4_mqa_logits( + (q[i:j], None), + (k, scales), + weights[i:j], + starts[i:j], + ends[i:j], + clean_logits=False, + ) + ops.top_k_per_row_prefill( + logits, + starts[i:j], + ends[i:j], + out[i:j], + j - i, + logits.stride(0), + logits.stride(1), + TOPK, + ) + return out + + +def _check(rel, q, k, scales, weights, starts, ends): + """Indices relative to the row start: count, bounds, uniqueness, scores.""" + from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + + 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" + # 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 = fp8_fp4_mqa_logits( + (q[i:j], None), + (k, scales), + weights[i:j], + starts[i:j], + ends[i:j], + clean_logits=True, + ) + 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 +@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): + 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 with fewer than 2048 visible keys + ends = torch.tensor( + [ + 0, + 1, + 31, + 127, + 1024, + 2047, + 2048, + 2049, + 4096, + 6144, + 8192, + 10000, + 12000, + 14000, + 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) + assert 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) + _check(dense, q, k, scales, weights, starts, ends) + assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all() + + +@requires_fused +@pytest.mark.parametrize( + "lens", [[3000, 5000, 1000], [40000, 1500, 70000, 9], [600000, 200000]] +) +def test_multi_request_chunk_matches_dense(lens): + """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) + assert 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) + _check(dense, q, k, scales, weights, starts, ends) + assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all() + # Both paths agree on (nearly) every selected key. + 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.99 * int((dense >= 0).sum()) + + +@requires_fused +def test_rows_start_at_zero_skips_conversion(): + rows, keys = 32, 8192 + 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) + a = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32) + b = torch.empty_like(a) + assert fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, a) + assert fused_mqa_topk_prefill( + q, k, scales, weights, starts, ends, b, rows_start_at_zero=True + ) + torch.testing.assert_close(a.sort(1).values, b.sort(1).values) + + +@requires_fused +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, rows_start_at_zero=True + ) + 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 +def test_unsupported_inputs_fall_back(): + 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) + out = torch.empty(16, TOPK, device="cuda", dtype=torch.int32) + assert not fused_mqa_topk_prefill( + q[:, :16], k, scales, weights[:, :16], starts, ends, out + ) + assert not fused_mqa_topk_prefill( + q, k, scales, weights, starts, ends, out[:, :1024] + ) + q17 = torch.zeros(16385, 32, 128, device="cuda").to(torch.float8_e4m3fn) + assert not fused_mqa_topk_prefill(q17, k, scales, weights, starts, ends, out) + + +def _split(seq_lens, q_lens, fused, offset=0): + return DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( + torch.tensor(seq_lens), + torch.tensor(q_lens), + 1048576, + LOGITS_BUDGET, + request_offset=offset, + use_fused_mqa_topk=fused, + ) + + +@pytest.mark.parametrize("fused,expected", [(False, 64), (True, 1)]) +def test_planner_one_fused_call_per_16k_rows(fused, expected): + chunks = _split([524288], [16384], fused) + assert len(chunks) == expected + assert sum(q.stop - q.start for _, q in chunks) == 16384 + + +def test_planner_fused_any_length_and_short_queries(): + # No length threshold and no minimum query count on the fused path. + assert _split([524288], [512], True) == [(slice(0, 1), slice(0, 512))] + assert _split([4096], [4096], True) == [(slice(0, 1), slice(0, 4096))] + # Long single-request steps are split at 16384 query rows. + assert _split([65536], [40000], True) == [ + (slice(0, 1), slice(0, 16384)), + (slice(0, 1), slice(16384, 32768)), + (slice(0, 1), slice(32768, 40000)), + ] + + +def test_planner_packs_requests_up_to_16k_rows(): + # Four 4K requests fit one fused call (16384 rows); dense needs a split. + assert _split([4096] * 4, [4096] * 4, True) == [(slice(0, 4), slice(0, 16384))] + assert len(_split([4096] * 4, [4096] * 4, False)) > 1 + chunks = _split([524288, 262144], [16384, 9], True, offset=2) + assert chunks == [(slice(2, 3), slice(0, 16384)), (slice(3, 4), slice(0, 9))] + + +def test_planner_fused_never_packs_past_key_limit(): + # Two 600K-key continuation steps of 50 rows: dense packs them into one + # logits-budget chunk; the fused planner keeps one fused chunk per request so + # no chunk reaches the dense path (whose logits buffer is not reserved). + args = ( + torch.tensor([600_000, 600_000]), + torch.tensor([50, 50]), + 4_000_000, + LOGITS_BUDGET, + ) + dense = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks(*args) + fused = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( + *args, use_fused_mqa_topk=True + ) + assert dense == [(slice(0, 2), slice(0, 100))] + assert fused == [(slice(0, 1), slice(0, 50)), (slice(1, 2), slice(0, 50))] + + +def test_planner_keeps_budget_beyond_fused_key_limit(): + seq_lens, q_lens = [2_000_000], [256] + fused = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( + torch.tensor(seq_lens), + torch.tensor(q_lens), + 4_000_000, + LOGITS_BUDGET, + use_fused_mqa_topk=True, + ) + dense = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( + torch.tensor(seq_lens), + torch.tensor(q_lens), + 4_000_000, + LOGITS_BUDGET, + ) + assert fused == dense and len(dense) == 4 diff --git a/vllm/envs.py b/vllm/envs.py index a7bde6b336fa..f9896ee1ce36 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 = False VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512 VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN: int = 8192 VLLM_XPU_PP_MICROBATCH: bool = True @@ -1093,6 +1094,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 eligible + # sparse MLA indexer prefill chunks on SM100-family GPUs (FP8 indexer cache, + # 32 index heads, top-k 2048, no DCP/PCP). Falls back to the dense path + # otherwise. Default: off. + "VLLM_USE_FUSED_MQA_TOPK": lambda: bool( + int(os.getenv("VLLM_USE_FUSED_MQA_TOPK", "0")) + ), # 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..8ed06c69b402 --- /dev/null +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -0,0 +1,238 @@ +# 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). + +One kernel (``torch.ops._C.fused_mqa_topk_prefill``) scores every visible key +of each query row with the DeepSeek sparse-attention indexer and selects the +row's top-2048 keys, without materializing the ``[rows, keys]`` logits matrix +that the dense path (``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``) needs. + +Scratch memory is a fixed pool of 8-row slots (2 per SM, ~311 MB on a 148-SM +GPU) that is allocated once per device and reused by every call, independently +of the number of rows or keys. +""" + +from dataclasses import dataclass + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +FUSED_MQA_TOPK_HEADS = 32 +FUSED_MQA_TOPK_HEAD_DIM = 128 +FUSED_MQA_TOPK_TOPK = 2048 +# Largest query chunk / gathered key range a single call accepts. +FUSED_MQA_TOPK_MAX_ROWS = 16384 +FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 + +_CANDIDATES_PER_ROW = 16384 +_RETAINED_PER_ROW = 4096 +_SLOTS_PER_SM = 2 +_STATS_PER_CTA = 7 + + +@dataclass +class FusedMQATopKWorkspace: + """Caller-owned scratch of the fused kernel (see fused_mqa_topk.cu).""" + + slot_flags: torch.Tensor # [S] int32, zero between calls + values: torch.Tensor # [8S, 16384] fp16 + indices: torch.Tensor # [8S, 16384] int32 + retained: torch.Tensor # [8S, 4096] fp32 + retained_indices: torch.Tensor # [8S, 4096] int32 + counts: torch.Tensor # [16384] int32 + gates: torch.Tensor # [16384] fp32 + stats: torch.Tensor # [16384 / 8 * 7] int32, per-CTA diagnostics + + def as_args(self) -> tuple[torch.Tensor, ...]: + return ( + self.slot_flags, + self.values, + self.indices, + self.retained, + self.retained_indices, + self.counts, + self.gates, + self.stats, + ) + + +_WORKSPACES: dict[int, FusedMQATopKWorkspace] = {} + + +def fused_mqa_topk_available() -> bool: + """True when the fused 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 get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: + """Return the per-device workspace, allocating it on first use. + + The slot flags are zeroed once; every kernel launch releases the slots it + claimed, so the flags are zero again whenever the stream is idle. + """ + index = ( + device.index + if device.index is not None + else torch.accelerator.current_device_index() + ) + ws = _WORKSPACES.get(index) + if ws is None: + dev = torch.device("cuda", index) + num_sms = torch.cuda.get_device_properties(index).multi_processor_count + slots = _SLOTS_PER_SM * num_sms + rows = FUSED_MQA_TOPK_MAX_ROWS + ws = FusedMQATopKWorkspace( + slot_flags=torch.zeros(slots, dtype=torch.int32, device=dev), + values=torch.empty( + (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.float16, device=dev + ), + indices=torch.empty( + (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.int32, device=dev + ), + retained=torch.empty( + (8 * slots, _RETAINED_PER_ROW), dtype=torch.float32, device=dev + ), + retained_indices=torch.empty( + (8 * slots, _RETAINED_PER_ROW), dtype=torch.int32, device=dev + ), + counts=torch.empty(rows, dtype=torch.int32, device=dev), + gates=torch.empty(rows, dtype=torch.float32, device=dev), + stats=torch.empty( + rows // 8 * _STATS_PER_CTA, dtype=torch.int32, device=dev + ), + ) + nbytes = sum(t.numel() * t.element_size() for t in ws.as_args()) + logger.info( + "Allocated fused MQA+top-k workspace on cuda:%d (%d slots, %.1f MiB).", + index, + slots, + nbytes / 2**20, + ) + _WORKSPACES[index] = ws + return ws + + +def fused_mqa_topk_config_supported(q: torch.Tensor, topk_tokens: int) -> bool: + """Whether this indexer's query layout and top-k fit the fused kernel.""" + return ( + topk_tokens == FUSED_MQA_TOPK_TOPK + and q.dtype == torch.float8_e4m3fn + and q.dim() == 3 + and tuple(q.shape[1:]) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM) + and q.is_cuda + and fused_mqa_topk_available() + ) + + +def fused_mqa_topk_supported( + q: torch.Tensor, k: torch.Tensor, topk_tokens: int +) -> bool: + """Static eligibility of an indexer prefill chunk for the fused kernel.""" + return ( + k.dtype == torch.float8_e4m3fn + and k.dim() == 2 + and k.shape[1] == FUSED_MQA_TOPK_HEAD_DIM + and 0 < q.shape[0] <= FUSED_MQA_TOPK_MAX_ROWS + and 0 < k.shape[0] <= FUSED_MQA_TOPK_MAX_KEYS + and fused_mqa_topk_config_supported(q, topk_tokens) + ) + + +def fused_mqa_topk_key_scales(scales: torch.Tensor) -> torch.Tensor: + """Per-key fp32 scales from the indexer K cache. + + SM100 caches store each power-of-two scale as four equal UE8M0 bytes in an + int32; the fp32 value 2**(e - 127) has bit pattern ``e << 23``. + """ + if scales.dtype == torch.int32: + return ((scales & 0xFF) << 23).view(torch.float32) + return scales + + +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, + *, + rows_start_at_zero: bool = False, +) -> bool: + """Fused replacement of ``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``. + + q: [rows, 32, 128] fp8 e4m3; k: [keys, 128] fp8 e4m3; scales: [keys] fp32; + weights: [rows, 32] fp32; starts/ends: [rows] int32, the visible key range + ``[start, end)`` of each row; output: [rows, 2048] int32. + + 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 2048 keys are visible. ``rows_start_at_zero=True`` + promises every ``start`` is 0 (single-request chunk) and skips the + relative-index conversion. + + Returns False (and leaves ``output`` untouched) for unsupported inputs so + the caller can fall back to the dense path. + """ + rows, keys = q.shape[0], k.shape[0] + if ( + not fused_mqa_topk_supported(q, k, output.shape[1] if output.dim() == 2 else 0) + or scales.dtype != torch.float32 + or scales.numel() != keys + or weights.dtype != torch.float32 + or tuple(weights.shape) != (rows, FUSED_MQA_TOPK_HEADS) + or starts.dtype != torch.int32 + or ends.dtype != torch.int32 + or starts.numel() != rows + or ends.numel() != rows + or output.dtype != torch.int32 + or tuple(output.shape) != (rows, FUSED_MQA_TOPK_TOPK) + ): + return False + + logger.info_once( + "Using the single-launch fused MQA+top-k prefill kernel " + "(FP8, 32 heads, top-k 2048)." + ) + q = q.contiguous() + k = k.contiguous() + scales = scales.contiguous() + weights = weights.contiguous() + starts = starts.contiguous() + ends = ends.contiguous() + + # The kernel writes the output with 16-byte stores; any other layout goes + # through a temporary (never the case for the indexer's own buffers). + direct = output.is_contiguous() and output.data_ptr() % 16 == 0 + out = ( + output + if direct + else torch.empty_like(output, memory_format=torch.contiguous_format) + ) + ws = get_fused_mqa_topk_workspace(q.device) + # The kernel emits absolute key indices; the dense top-k (and the sparse + # attention consumer) use indices relative to the row start. relative=True + # converts in place (rows with start 0 are skipped on the GPU). + torch.ops._C.fused_mqa_topk_prefill( + q, + k, + scales, + weights, + starts, + ends, + out, + *ws.as_args(), + not rows_start_at_zero, + ) + if not direct: + output.copy_(out) + return True diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 2a60ca9396e5..070b89e25916 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -29,6 +29,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_MAX_KEYS, + fused_mqa_topk_config_supported, + fused_mqa_topk_key_scales, + fused_mqa_topk_prefill, + get_fused_mqa_topk_workspace, +) from vllm.model_executor.layers.indexer_topk import ( RADIX_TOPK_WORKSPACE_SIZE, get_indexer_topk, @@ -497,13 +504,29 @@ def sparse_attn_indexer( ) profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) + uses_fused = ( + envs.VLLM_USE_FUSED_MQA_TOPK + and not use_pcp + and dcp_world_size == 1 + and not use_fp4_cache + and candidate_blocks is None + and fused_mqa_topk_config_supported(q_quant, topk_tokens) + ) + if uses_fused: + # The fused prefill kernel's fixed-size workspace is persistent; + # allocate it during profiling so it is accounted for. + get_fused_mqa_topk_workspace(q_quant.device) # Dummy allocation to simulate for peak logits tensor memory during inference. - # FP8 elements so elements == bytes - max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 - _ = torch.empty( - max_logits_elems, dtype=torch.uint8, device=hidden_states.device - ) + # Skipped when the fused kernel serves every prefill chunk: the planner never + # packs a chunk past its key limit, so with max_model_len within that limit + # no chunk reaches the dense logits path. + if not (uses_fused and max_model_len <= FUSED_MQA_TOPK_MAX_KEYS): + # FP8 elements so elements == bytes + max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + _ = torch.empty( + max_logits_elems, dtype=torch.uint8, device=hidden_states.device + ) # The dummy run returns before exchange_topk executes, so reserve its # worst-case transient for the memory profile. @@ -706,70 +729,109 @@ 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( + # Single-launch fused MQA-logits + top-k (SM100, FP8, 32 heads, + # top-k 2048): no logits buffer, any context length. DCP/PCP, + # the FP4 cache and v4.1 candidate selection stay dense. + use_fused = ( + envs.VLLM_USE_FUSED_MQA_TOPK + and not use_pcp + and dcp_world_size == 1 + and not current_platform.is_xpu() + ) + if ( + use_fused + and q_scale_slice is None + and candidate_blocks is None + and fused_mqa_topk_prefill( q_slice_cast, k_quant_cast, - k_scale_cast, + fused_mqa_topk_key_scales(k_scale_cast), weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, + topk_indices, + # A single-request chunk has every row start at 0. + rows_start_at_zero=chunk.block_table.shape[0] == 1, ) - 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 = 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, + ): + continue + # Fused-planned chunks are not bounded by the logits budget, so + # the dense fallback splits their rows to stay within it. + dense_rows = q_slice_cast.shape[0] + if use_fused: + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + dense_rows = max(1, budget // (4 * max(1, chunk.max_seqlen_k))) + for sub in range(0, q_slice_cast.shape[0], dense_rows): + rows = slice(sub, min(sub + dense_rows, q_slice_cast.shape[0])) + tokens = slice(row_start + rows.start, row_start + rows.stop) + sub_ks, sub_ke = cu_seqlen_ks[rows], cu_seqlen_ke[rows] + sub_q_scale = ( + q_scale_slice[rows] if q_scale_slice is not None else None ) - 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( - logits, - row_starts, - row_ends, - chunk_candidates.shape[1], - candidate_block_size, - chunk_candidates, + if current_platform.is_xpu(): + if sub_q_scale 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], + sub_ks, + sub_ke, ) + row_starts = sub_ks + row_ends = sub_ke else: - _apply_candidate_mask( - logits, - row_starts, + sub_q_scale, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast[rows], sub_q_scale, weights[tokens] + ) + row_ends = sub_ke - sub_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast[rows], sub_q_scale), + (k_quant_cast, k_scale_cast), + chunk_weights, + sub_ks, + sub_ke, + chunk.max_seqlen_k, + ), row_ends, - chunk_candidates, - candidate_block_size, ) - 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, - ) + 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, + ) + ops.top_k_per_row_prefill( + logits, + torch.zeros_like(row_ends) if row_starts is None else row_starts, + row_ends, + topk_indices[rows], + 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/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index c353a1f2fbb8..0688c7b0ebca 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, + fused_mqa_topk_available, +) from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, @@ -1362,6 +1367,7 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, + use_fused_mqa_topk: 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) @@ -1371,8 +1377,19 @@ 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 ``use_fused_mqa_topk`` the fused MQA+top-k kernel needs no logits + buffer: chunks whose gathered keys fit its 1M-key limit are bounded by + its 16384-row limit instead of the logits budget (one fused call per + chunk). Requests are never packed past that key limit, so only a single + request longer than 1M keys takes the dense, logits-budget path (the + profiling run relies on this to skip the dense logits reservation). + Returns list of (req_slice, query_slice) tuples. """ + + def fused(n: int) -> bool: + return use_fused_mqa_topk and n <= FUSED_MQA_TOPK_MAX_KEYS + chunks: list[tuple[slice, slice]] = [] n = len(compressed_seq_lens_cpu) max_logits_elems = max_logits_bytes // 4 @@ -1387,7 +1404,11 @@ 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(new_n): + fits = new_m <= FUSED_MQA_TOPK_MAX_ROWS + else: + fits = not use_fused_mqa_topk 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: @@ -1406,6 +1427,8 @@ def _split_indexer_prefill_chunks( max_q = ( max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) ) + if fused(chunk_n): + 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))) @@ -1535,12 +1558,24 @@ def build( request_offset=num_decodes, ) else: + chunk_options = {} + if ( + envs.VLLM_USE_FUSED_MQA_TOPK + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and not self.indexer_uses_fp4 + and self.dcp_world_size == 1 + and self.pcp_world_size == 1 + and fused_mqa_topk_available() + ): + chunk_options["use_fused_mqa_topk"] = True chunk_specs = self._split_indexer_prefill_chunks( self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), prefill_query_lens_cpu, self.max_prefill_buffer_size, max_logits_bytes, request_offset=num_decodes, + **chunk_options, ) chunks = [] From 6578ea0ea6dfa1054130982edadc801934e35048 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Thu, 8 Oct 2026 11:46:40 +0000 Subject: [PATCH 02/23] cleanup Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/fallback_core.cuh | 4 +- csrc/fused_mqa_topk/sample_core.cuh | 4 +- csrc/fused_mqa_topk/sm100_candidate_scan.cuh | 1276 ------------------ 3 files changed, 6 insertions(+), 1278 deletions(-) delete mode 100644 csrc/fused_mqa_topk/sm100_candidate_scan.cuh diff --git a/csrc/fused_mqa_topk/fallback_core.cuh b/csrc/fused_mqa_topk/fallback_core.cuh index f4ae2d112ad4..95ee8619dabd 100644 --- a/csrc/fused_mqa_topk/fallback_core.cuh +++ b/csrc/fused_mqa_topk/fallback_core.cuh @@ -30,7 +30,9 @@ static_assert(MQA_HANDOFF_INTERVAL == 32 || MQA_HANDOFF_INTERVAL == 64 || #include #include #include -#include "sm100_candidate_scan.cuh" +#include +#include +#include // Shared SM100 MQA logits core plus contiguous-KV and paged entries // Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature diff --git a/csrc/fused_mqa_topk/sample_core.cuh b/csrc/fused_mqa_topk/sample_core.cuh index 5a8c0b82a77c..150f1a467297 100644 --- a/csrc/fused_mqa_topk/sample_core.cuh +++ b/csrc/fused_mqa_topk/sample_core.cuh @@ -30,7 +30,9 @@ static_assert(MQA_HANDOFF_INTERVAL == 32 || MQA_HANDOFF_INTERVAL == 64 || #include #include #include -#include "sm100_candidate_scan.cuh" +#include +#include +#include // Shared SM100 MQA logits core plus contiguous-KV and paged entries // Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature diff --git a/csrc/fused_mqa_topk/sm100_candidate_scan.cuh b/csrc/fused_mqa_topk/sm100_candidate_scan.cuh deleted file mode 100644 index 28aaf07613fc..000000000000 --- a/csrc/fused_mqa_topk/sm100_candidate_scan.cuh +++ /dev/null @@ -1,1276 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025 DeepSeek -// Derived from DeepGEMM commit 891d57b4db1071624b5c8fa0d1e51cb317fa709f; -// see LICENSE.deepseek-deepgemm. - -#pragma once - -#include - -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace candidate_scan_scan { - -using namespace deep_gemm; - -inline constexpr uint32_t kEmitChunkBlocks = 256; -inline constexpr uint32_t kEmitLaneSlots = 18; -inline constexpr uint32_t kGateStride = 64; -inline constexpr uint32_t kMathRegisters = 240; -inline constexpr uint32_t kSpecializedRegisters = 24; -inline constexpr uint32_t kUmmaStages = 2; -inline constexpr uint32_t kSparseRefreshNs = 512; -inline constexpr uint32_t kSparseRefreshIdleNs = 2048; -// A ring slot can only equal this pattern if a +NaN score passed the gate, -// which both gate comparisons reject; the refresher treats it as "empty". -inline constexpr uint32_t kRingSentinel = 0xffffffffu; -// Gate updates only need to keep pace with the 65,536-token flush windows, -// so the refresher can poll orders of magnitude slower than the dynamic -// histogram daemon and stay invisible to the math warps' issue ports. -inline constexpr uint32_t kRingActiveNs = 2048; -inline constexpr uint32_t kRingIdleNs = 2048; - -// Production candidate ABI: a six-byte global record. cand_val stores the -// low 16 bits of an ascending, sign-aware FP32 high24 score code; cand_idx -// stores its high eight score bits above the exact 20-bit KV index. The local -// ring uses one uint32_t containing high24 score bits and an eight-bit -// block-in-window. -using CandidateValue = uint16_t; -constexpr uint32_t kCandidateIndexBits = 20; -constexpr uint32_t kCandidateIndexMask = (1u << kCandidateIndexBits) - 1u; - -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_ordered_fp32_code(const float value) { - const uint32_t bits = __float_as_uint(value); - return (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); -} - -CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) { - return candidate_ordered_fp32_code(value) >> 8; -} - -// Descending full-width ordering: smaller unsigned keys correspond to larger -// FP32 values. -CUTLASS_DEVICE uint32_t candidate_descending_fp32_code(const float value) { - const uint32_t bits = __float_as_uint(value); - return (bits & 0x80000000u) ? bits : (~bits & 0x7fffffffu); -} - -// High eight bits of a sign-aware FP16 key. -// Keeping 256 coarse buckets makes the online histogram small enough to live -// beside the score pipeline while retaining much more resolution than the -// sign/exponent byte of a raw FP32 key. -CUTLASS_DEVICE uint32_t candidate_fixed_half_bucket(const float value) { - const uint32_t fp32_bits = __float_as_uint(value); - const __half half_value = __float2half(value); - uint16_t bits = __half_as_ushort(half_value) & 0x7fffu; - bits |= static_cast((fp32_bits >> 16) & 0x8000u); - bits = (bits & 0x8000u) ? bits : static_cast(~bits & 0x7fffu); - return static_cast(bits >> 8); -} - -struct OnlineFixedCandidate { - uint32_t bucket; - uint32_t payload; -}; - -CUTLASS_DEVICE OnlineFixedCandidate candidate_online_fixed(const float value) { - const uint32_t bucket = candidate_fixed_half_bucket(value); - const uint32_t raw_desc24 = candidate_descending_fp32_code(value) >> 8; - // Experimental six-byte payload: the coarse bucket is explicit, while - // the remaining 16 bits retain the raw-score order inside that bucket. - return {bucket, (bucket << 16) | (raw_desc24 & 0xffffu)}; -} - -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 float candidate_decode_score(const CandidateValue value, - const int32_t packed_idx) { - const uint32_t code = candidate_load_score_code(value, packed_idx); - const uint32_t ordered = code << 8; - const uint32_t bits = - (ordered & 0x80000000u) ? (ordered ^ 0x80000000u) : ~ordered; - return __uint_as_float(bits); -} - -CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { - return static_cast(static_cast(packed_idx) & - kCandidateIndexMask); -} - -CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, - int32_t* index_dst, const float bq, - const uint32_t kv_index) { - const uint32_t payload = candidate_fp24_code(bq); - __stcs(value_dst, static_cast(payload)); - __stcs(index_dst, - static_cast(candidate_pack_index(payload, kv_index))); -} - -CUTLASS_DEVICE void store_candidate_payload(CandidateValue* value_dst, - int32_t* index_dst, - const uint32_t payload, - const uint32_t kv_index) { - __stcs(value_dst, static_cast(payload)); - __stcs(index_dst, - static_cast(candidate_pack_index(payload, kv_index))); -} - -CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst, - int32_t* index_dst, - const uint32_t payload, - const uint32_t kv_index) { - store_candidate_payload(value_dst, index_dst, payload, kv_index); -} - -template -CUTLASS_GLOBAL __launch_bounds__( - kNumSpecializedThreads + kNumMathThreads, - 1) void sm100_candidate_scan(const uint32_t seq_len, - const uint32_t seq_len_kv, - uint32_t* cu_seq_len_k_start, - uint32_t* cu_seq_len_k_end, - const float* __restrict__ origin, // [seq_len] - const float* __restrict__ inv_delta, // [seq_len] - int32_t* __restrict__ th_bucket, // [seq_len] - int32_t* __restrict__ bcount, // [seq_len, - // num_buckets] - const uint32_t num_buckets, - const uint32_t topk, - const uint32_t refresh_every, - const uint32_t num_kv_splits, - const uint32_t - probe_group, // compacted-space group size - // (pstp-1)*64; 0 = no probe - // compaction (identity map) - const uint64_t - probe_magic, // ceil(2^42/probe_group): - // exact div via mul-shift - const uint32_t - probe_add_max, // npage*64 cap for the map - CandidateValue* __restrict__ cand_val, // [seq_len, - // cand_cap] - int32_t* __restrict__ cand_idx, // [seq_len, - // cand_cap] - int32_t* __restrict__ cand_cnt, // [seq_len] - const uint32_t cand_cap, - const __grid_constant__ cute::TmaDescriptor - tensor_map_q, - const __grid_constant__ cute::TmaDescriptor - tensor_map_kv, - // fp8: per-token fp32 dequant scales. fp4: the - // UE8M0 SF-KV int32 stream (same 1-D shape). - const __grid_constant__ cute::TmaDescriptor - tensor_map_kv_scales, - const __grid_constant__ cute::TmaDescriptor - tensor_map_weights, - // fp4 only: per-(token, head) packed UE8M0 - // SF-Q int32 stream; fp8 callers pass any map - // (unread). - const __grid_constant__ cute::TmaDescriptor - tensor_map_sf_q) { - const auto num_q_blocks = math::ceil_div(seq_len, BLOCK_Q); - using Barrier = cutlass::arch::ClusterTransactionBarrier; - - const auto warp_idx = cutlass::canonical_warp_idx_sync(); - const auto warpgroup_idx = warp_idx / 4; - const auto lane_idx = ptx::get_lane_idx(); - constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; - constexpr uint32_t kNumMathWarps = kNumMathThreads / 32; - constexpr uint32_t kNumUmmaStages = kUmmaStages; - constexpr uint32_t kNumUmmaBuffers = kNumMathWarpGroups * kNumUmmaStages; - // fp4 replaces the per-warpgroup UMMA double buffers with a shared - // 3-stage TMEM accumulator pool (the SF columns take the freed space); - constexpr uint32_t kNumAccumBufs = kNumUmmaBuffers; - // The ring refresher reloads more often: a flooding row keeps emitting - // at the stale gate for a full stride, so the flood integral scales with - // the reload cadence. The reload is one ld.shared.v4 per stride. - constexpr uint32_t kActiveGateStride = - kOnlineFixedBuckets ? 4u : (kRingRefresh ? 8u : kGateStride); - constexpr uint32_t kActiveEmitChunkBlocks = - kOnlineFixedBuckets ? 32u : kEmitChunkBlocks; - static_assert( - !(kOnlineFixedBuckets && kStaticHotGate), - "online fixed buckets and static HOT gate are separate A/B paths"); - static_assert(!kStaticHotNoHist || kStaticHotGate, - "static HOT no-hist requires the static HOT gate"); - static_assert(!kStaticHotExactGate || kStaticHotNoHist, - "exact FP24 gate requires the static HOT no-hist path"); - static_assert( - !kRingRefresh || - (kStaticHotNoHist && !kStaticHotExactGate && !kOnlineFixedBuckets), - "ring refresher v1 covers only the plain bucket-gate no-hist writer"); - DG_STATIC_ASSERT((BLOCK_Q == 4 || BLOCK_Q == 2) && kNumMathWarps == 8, - "CandidateScan requires BLOCK_Q in {4 (H=32), 2 (H=64)} and " - "8 math warps"); - DG_STATIC_ASSERT( - BLOCK_KV == 256, - "the local ring infers the KV tile coordinate from its owner lane"); - DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, - "Invalid threads"); - - if (warp_idx == kSpecWarpStart) { - cute::prefetch_tma_descriptor(&tensor_map_q); - cute::prefetch_tma_descriptor(&tensor_map_kv); - cute::prefetch_tma_descriptor(&tensor_map_kv_scales); - cute::prefetch_tma_descriptor(&tensor_map_weights); - } - - // fp4 packs two e2m1 elements per byte; SF regions are UE8M0 bytes - // packed 4-per-int32 and padded to the 128-element UTCCP granule. - static constexpr uint32_t kKvElemBytesNum = 2u; // per 2 elems - static constexpr uint32_t kNumUTCCPAlignedElems = 128; - static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = - BLOCK_Q * kNumHeads * kHeadDim * kKvElemBytesNum / 2; - static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = - BLOCK_Q * kNumHeads * sizeof(float); - static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = - BLOCK_KV * kHeadDim * kKvElemBytesNum / 2; - // fp8: one fp32 dequant scale per token. fp4: one int32 (4x UE8M0) per - // token, padded to the UTCCP granule. - static constexpr uint32_t SMEM_KV_SCALE_SIZE_PER_STAGE = - BLOCK_KV * sizeof(float); - static constexpr uint32_t SMEM_SF_Q_SIZE = 0u; - static constexpr uint32_t ALIGNED_SMEM_KV_SCALE_SIZE_PER_STAGE = - math::constexpr_align(SMEM_KV_SCALE_SIZE_PER_STAGE, 512u); - - extern __shared__ __align__(512) uint8_t smem_buffer[]; - DG_STATIC_ASSERT(SMEM_Q_SIZE_PER_STAGE % 512 == 0, "Unaligned TMA swizzling"); - DG_STATIC_ASSERT(SMEM_WEIGHT_SIZE_PER_STAGE % 512 == 0, - "Unaligned TMA swizzling"); - DG_STATIC_ASSERT(SMEM_KV_SIZE_PER_STAGE % 512 == 0, - "Unaligned TMA swizzling"); - - constexpr uint32_t kNumAccumTmemCols = BLOCK_Q * kNumHeads * kNumAccumBufs; - constexpr uint32_t kNumTmemCols = kNumAccumTmemCols; - DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); - - auto smem_q = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + - SMEM_Q_SIZE_PER_STAGE * i); - }); - auto smem_weights = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + - SMEM_Q_SIZE_PER_STAGE * kNumQStages + - SMEM_WEIGHT_SIZE_PER_STAGE * i); - }); - auto smem_kv = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast<__nv_fp8_e4m3*>( - smem_buffer + (SMEM_Q_SIZE_PER_STAGE * kNumQStages + - SMEM_WEIGHT_SIZE_PER_STAGE * kNumQStages + - SMEM_KV_SIZE_PER_STAGE * i)); - }); - auto smem_kv_scales = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + - SMEM_Q_SIZE_PER_STAGE * kNumQStages + - SMEM_WEIGHT_SIZE_PER_STAGE * kNumQStages + - SMEM_KV_SIZE_PER_STAGE * kNumKVStages + - ALIGNED_SMEM_KV_SCALE_SIZE_PER_STAGE * i); - }); - - auto smem_sf_q = reinterpret_cast( - reinterpret_cast(smem_kv_scales[kNumKVStages])); - auto barrier_ptr = reinterpret_cast( - reinterpret_cast(smem_sf_q) + SMEM_SF_Q_SIZE); - auto full_q_barriers = - utils::PatternVisitor([&](const uint32_t& i) { return barrier_ptr + i; }); - auto empty_q_barriers = utils::PatternVisitor( - [&](const uint32_t& i) { return barrier_ptr + (kNumQStages + i); }); - auto full_kv_barriers = utils::PatternVisitor( - [&](const uint32_t& i) { return barrier_ptr + (kNumQStages * 2 + i); }); - auto empty_kv_barriers = utils::PatternVisitor([&](const uint32_t& i) { - return barrier_ptr + (kNumQStages * 2 + kNumKVStages + i); - }); - auto full_umma_barriers = utils::PatternVisitor([&](const uint32_t& i) { - return barrier_ptr + (kNumQStages * 2 + kNumKVStages * 2 + i); - }); - auto empty_umma_barriers = utils::PatternVisitor([&](const uint32_t& i) { - return barrier_ptr + - (kNumQStages * 2 + kNumKVStages * 2 + kNumAccumBufs + i); - }); - - auto sf_ready_barriers = utils::PatternVisitor([&](const uint32_t& i) { - return barrier_ptr + - (kNumQStages * 2 + kNumKVStages * 2 + kNumAccumBufs * 2 + i); - }); - auto tmem_ptr_in_smem = reinterpret_cast( - barrier_ptr + kNumQStages * 2 + kNumKVStages * 2 + kNumAccumBufs * 2 + - kNumKVStages); - auto scan_done_flag = reinterpret_cast(tmem_ptr_in_smem + 1); - auto warpq_count = reinterpret_cast(tmem_ptr_in_smem + 4); - // Chunked emit does not use the legacy warp-queue count words. Reuse four - // of them as a CTA-local Gate4 mailbox. Values are positive float edge - // bit-patterns, so unsigned atomicMin is exactly a monotonic gate tighten. - auto sparse_gate_bits = reinterpret_cast(warpq_count); - auto emit_smem_records = reinterpret_cast( - warpq_count + (kRingRefresh ? 4u : kNumMathWarps * BLOCK_Q)); - auto smem_hist = reinterpret_cast( - emit_smem_records + kNumMathWarps * BLOCK_Q * kEmitLaneSlots * 32u); - // Ring-refresher scratch sits after the refresher's histogram (which - // reuses the dynamic-path histogram region): per-(math warp, row) flush - // sequence words, the spare warps' shadow copies, and per-lane - // consumed-slot cursors. - auto ring_seq = reinterpret_cast(smem_hist + BLOCK_Q * 256); - auto ring_shadow = ring_seq + kNumMathWarps * BLOCK_Q; - auto ring_pending = - reinterpret_cast(ring_shadow + kNumMathWarps * BLOCK_Q); - auto ring_stale = ring_pending + BLOCK_Q; - auto ring_last = ring_stale + BLOCK_Q; - auto ring_progress = reinterpret_cast(ring_last + BLOCK_Q); - - DG_STATIC_ASSERT( - kNumSpecializedThreads % 128 == 0 and kNumSpecializedThreads >= 64, - "Invalid threads"); - if (warp_idx == kSpecWarpStart and cute::elect_one_sync()) { -#pragma unroll - for (uint32_t i = 0; i < kNumQStages; ++i) { - full_q_barriers[i]->init(1); - empty_q_barriers[i]->init(kNumMathThreads + 32); - } -#pragma unroll - for (uint32_t i = 0; i < kNumKVStages; ++i) { - full_kv_barriers[i]->init(1); - empty_kv_barriers[i]->init(kNumMathThreads); - } - if constexpr (!kStaticHotNoHist || kRingRefresh) *scan_done_flag = 0; - cutlass::arch::fence_barrier_init(); - } - if (warp_idx == kSpecWarpStart + 1) { - if (cute::elect_one_sync()) { -#pragma unroll - for (uint32_t i = 0; i < kNumAccumBufs; ++i) { - full_umma_barriers[i]->init(1); - empty_umma_barriers[i]->init(128); - } - cutlass::arch::fence_barrier_init(); - } - cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); - } - // One CTA owns each row. The hot-only path emits no sample seeds and scans - // the complete KV range, so its exact histogram starts at zero in SMEM. - // CANDSCAN_STATIC_HOT_NOHIST_AB removes both this clear and every update; - // a post-scan kernel rebuilds metadata from the compact candidate slab. - if constexpr (!kStaticHotNoHist) { - for (uint32_t idx = threadIdx.x; idx < BLOCK_Q * num_buckets; - idx += blockDim.x) { - smem_hist[idx] = 0; - } - } - if constexpr (kRingRefresh) { - for (uint32_t idx = threadIdx.x; - idx < kNumMathWarps * BLOCK_Q * kEmitLaneSlots * 32u; - idx += blockDim.x) { - emit_smem_records[idx] = kRingSentinel; - } - // Preload the seed sample's full bucket histogram as the daemon's - // refresh base. The sample records are genuine row records in the - // same (origin, inv) bucket space and the exact-once scan starts after - // the sampled prefix, so ring harvests never recount them; base plus - // harvested stays a subset of the row's true records and every - // publish remains a provable upper bound. - for (uint32_t idx = threadIdx.x; idx < BLOCK_Q * num_buckets; - idx += blockDim.x) { - const uint32_t r = idx / num_buckets; - const uint32_t row_q = static_cast(blockIdx.x) * BLOCK_Q + r; - smem_hist[idx] = - row_q < seq_len - ? __ldcg(bcount + static_cast(row_q) * num_buckets + - (idx - r * num_buckets)) - : 0; - } - for (uint32_t idx = threadIdx.x; idx < kNumMathWarps * BLOCK_Q; - idx += blockDim.x) { - ring_seq[idx] = 0u; - ring_shadow[idx] = 0u; - } - for (uint32_t idx = threadIdx.x; idx < BLOCK_Q; idx += blockDim.x) { - ring_pending[idx] = 0; - ring_last[idx] = 0; - ring_stale[idx] = blockIdx.x * BLOCK_Q + idx < seq_len ? 0 : 2; - } - for (uint32_t idx = threadIdx.x; idx < kNumMathWarps * BLOCK_Q * 32u; - idx += blockDim.x) { - ring_progress[idx] = 0u; - } - } - if (threadIdx.x < BLOCK_Q) { - const uint32_t row_q = - static_cast(blockIdx.x) * BLOCK_Q + threadIdx.x; - const uint32_t row = min(row_q, seq_len - 1); - const int gate = __ldcg(th_bucket + row); - if constexpr (kStaticHotExactGate) { - // Exact compaction stores the ordered-FP32 pivot edge in this - // int32 slot. Preserve its bits so the unsigned strict comparison - // admits only later records from a better high24 class. - ptx::st_shared(sparse_gate_bits + threadIdx.x, - static_cast(gate)); - } else { - ptx::st_shared(sparse_gate_bits + threadIdx.x, - __float_as_uint(static_cast(gate + 1))); - } - } - __syncthreads(); - - constexpr uint32_t kNumSpecializedRegisters = kSpecializedRegisters; - constexpr uint32_t kNumMathRegisters = kMathRegisters; - - // V1 KV-split scheduling: blockIdx.x = q-block (one per CTA), blockIdx.y = - // contiguous KV sub-window. Split boundaries are BLOCK_KV-aligned. - const uint32_t block_q_idx = blockIdx.x; - const uint32_t kv_split = blockIdx.y; - uint32_t seq_k_start[BLOCK_Q], seq_k_end[BLOCK_Q]; - const auto load_schedule = - [&](const uint32_t block_q_idx) -> cute::tuple { - uint32_t start = cute::numeric_limits::max(); - uint32_t end = cute::numeric_limits::min(); - -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { - const auto q_idx = min(block_q_idx * BLOCK_Q + i, seq_len - 1); - seq_k_start[i] = cu_seq_len_k_start[q_idx]; - seq_k_end[i] = cu_seq_len_k_end[q_idx]; - if (block_q_idx * BLOCK_Q + i >= seq_len) { - // Padded row of a ragged final q-block: empty, aggregation-neutral. - seq_k_start[i] = seq_len_kv; - seq_k_end[i] = 0; - } - start = min(start, min(seq_k_start[i], seq_len_kv)); - end = max(end, min(seq_k_end[i], seq_len_kv)); - } - const uint32_t total_blocks = math::ceil_div(seq_len_kv, BLOCK_KV); - const uint32_t blocks_per_split = - math::ceil_div(total_blocks, num_kv_splits); - const uint32_t split_lo = kv_split * blocks_per_split * BLOCK_KV; - const uint32_t split_hi = - min((kv_split + 1) * blocks_per_split * BLOCK_KV, seq_len_kv); - start = start / 4 * 4; // TMA alignment for SF KV - if (start < split_lo) start = split_lo; - if (end > split_hi) end = split_hi; - const uint32_t nkv = - (end > start) ? math::ceil_div(end - start, BLOCK_KV) : 0; - return {start, nkv}; - }; - - const auto get_kv_pipeline = - [&](const uint32_t& kv_block_idx) -> cute::tuple { - return {kv_block_idx % kNumKVStages, (kv_block_idx / kNumKVStages) & 1}; - }; - - constexpr uint32_t UMMA_M = 128; - constexpr uint32_t UMMA_K = 32 / sizeof(cutlass::float_e4m3_t); - constexpr uint32_t UMMA_N = BLOCK_Q * kNumHeads; - - if (warp_idx == kSpecWarpStart) { - cutlass::arch::warpgroup_reg_dealloc(); - - if (cute::elect_one_sync()) { - if (block_q_idx < num_q_blocks) { - // Q + weights once for this q-block. Inner extents are BYTES - // for the uint8 maps, so fp4's packed rows halve them. - constexpr uint32_t kKvInnerBytes = kHeadDim * kKvElemBytesNum / 2; - tma::copy( - &tensor_map_q, full_q_barriers[0], smem_q[0], 0, - block_q_idx * BLOCK_Q * kNumHeads); - tma::copy(&tensor_map_weights, - full_q_barriers[0], smem_weights[0], 0, - block_q_idx * BLOCK_Q); - full_q_barriers[0]->arrive_and_expect_tx(SMEM_Q_SIZE_PER_STAGE + - SMEM_WEIGHT_SIZE_PER_STAGE); - - CUTE_TIE_DECL(load_schedule(block_q_idx), kv_start, num_kv_blocks); - for (uint32_t kv_block_idx = 0; kv_block_idx < num_kv_blocks; - ++kv_block_idx) { - CUTE_TIE_DECL(get_kv_pipeline(kv_block_idx), kv_stage_idx, kv_phase); - empty_kv_barriers[kv_stage_idx]->wait(kv_phase ^ 1); - - tma::copy( - &tensor_map_kv, full_kv_barriers[kv_stage_idx], - smem_kv[kv_stage_idx], 0, kv_start + kv_block_idx * BLOCK_KV); - tma::copy(&tensor_map_kv_scales, - full_kv_barriers[kv_stage_idx], - smem_kv_scales[kv_stage_idx], - kv_start + kv_block_idx * BLOCK_KV, 0); - full_kv_barriers[kv_stage_idx]->arrive_and_expect_tx( - SMEM_KV_SIZE_PER_STAGE + SMEM_KV_SCALE_SIZE_PER_STAGE); - } - } - } - } else if (warp_idx == kSpecWarpStart + 1) { - cutlass::arch::warpgroup_reg_dealloc(); - - DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(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); - - if (block_q_idx < num_q_blocks) { - CUTE_TIE_DECL(load_schedule(block_q_idx), kv_start, num_kv_blocks); - full_q_barriers[0]->wait(0); - - for (uint32_t kv_block_idx = 0; kv_block_idx < num_kv_blocks; - ++kv_block_idx) { - const uint32_t kvg = kv_block_idx; - CUTE_TIE_DECL(get_kv_pipeline(kvg), kv_stage_idx, kv_phase); - full_kv_barriers[kv_stage_idx]->wait(kv_phase); - - DG_STATIC_ASSERT(BLOCK_KV == kNumMathThreads, "Invalid block size"); - DG_STATIC_ASSERT(kHeadDim % UMMA_K == 0, "Invalid head dim"); - // Round-robin over kNumUmmaStages accumulators. A stage is - // reused every kNumUmmaStages tiles, so its phase toggles at - // that rate, not every tile. - const uint32_t umma_stage = kvg % kNumUmmaStages; - const uint32_t umma_phase = (kvg / kNumUmmaStages) & 1; -#pragma unroll - for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { - const uint32_t buf = i * kNumUmmaStages + umma_stage; - empty_umma_barriers[buf]->wait(umma_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_kv[kv_stage_idx], i * UMMA_M, k * UMMA_K); - auto b_desc = - mma::sm100::make_umma_desc(smem_q[0], 0, k * UMMA_K); - cute::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, buf * UMMA_N, k, - runtime_instr_desc); - } - cutlass::arch::umma_arrive( - reinterpret_cast(full_umma_barriers[buf])); - } - } - empty_q_barriers[0]->arrive(); - } - } else if (warp_idx == kSpecWarpStart + 2 or warp_idx == kSpecWarpStart + 3) { - // Spare-warp threshold-refresh daemon. Keeping refresh off the math - // warps avoids placing histogram scan latency on their critical path. - cutlass::arch::warpgroup_reg_dealloc(); - - if (block_q_idx < num_q_blocks) { - const uint32_t spare_id = warp_idx - (kSpecWarpStart + 2); // 0 or 1 - const auto refresh_row = [&](const uint32_t row, - const bool publish_boundary_counts) { - if (row >= seq_len) return false; - const uint32_t local_row = row - block_q_idx * BLOCK_Q; - const int32_t* srow = - smem_hist + (row - block_q_idx * BLOCK_Q) * num_buckets; - const int current_gate = - min(max(static_cast(__uint_as_float( - ptx::ld_shared(sparse_gate_bits + local_row))) - - 1, - 0), - static_cast(num_buckets) - 1); - // Only buckets strictly below the published gate can tighten - // it. If they contain fewer than topk entries, keep the - // current gate and avoid scanning the dead upper tail. - const uint32_t search_buckets = static_cast(current_gate); - int found = current_gate; - int carry = 0; - int found_lt = 0; - int found_eq = 0; - bool done = false; - for (uint32_t base = 0; base < search_buckets && !done; base += 32) { - uint32_t b = base + lane_idx; - int v = 0; - if (b < search_buckets) { - // In the single-split fast path smem already contains - // the initial global histogram plus every scan hit. - v = srow[b]; - } - // Histogram entries are nonnegative. Most 32-bucket - // groups remain strictly below K, so reject them with - // one REDUX and reserve the five dependent prefix - // shuffles for the single group that crosses K. - const int group_sum = __reduce_add_sync(0xffffffffu, v); - if (carry + group_sum < static_cast(topk)) { - carry += group_sum; - continue; - } - int prefix = v; -#pragma unroll - for (int off = 1; off < 32; off <<= 1) { - int nsh = __shfl_up_sync(0xffffffffu, prefix, off); - if (static_cast(lane_idx) >= off) prefix += nsh; - } - int incl = carry + prefix; - bool hit = (b < search_buckets) && (incl >= static_cast(topk)) && - (incl - v < static_cast(topk)); - unsigned hm = __ballot_sync(0xffffffffu, hit); - if (hm) { - const int hit_lane = __ffs(hm) - 1; - found = static_cast(base) + hit_lane; - found_lt = __shfl_sync(0xffffffffu, incl - v, hit_lane); - found_eq = __shfl_sync(0xffffffffu, v, hit_lane); - done = true; - } else { - carry += __shfl_sync(0xffffffffu, prefix, 31); - } - } - if (!done) { - // No lower bucket reached K, so the current gate remains - // the boundary. `carry` is exactly count(bucket(found + 1)); - if (done) { - th_bucket[row] = found; - if constexpr (!kStaticHotGate) { - atomicMin(sparse_gate_bits + local_row, edge); - } - } - if (publish_boundary_counts && num_buckets >= 3) { - // Reuse the dead histogram row as selector metadata. - int32_t* meta = bcount + static_cast(row) * num_buckets; - meta[0] = ~found; - meta[1] = found_lt; - meta[2] = found_eq; - } - } - return done; - }; - // CANDSCAN_STATIC_HOT_AB: the external HOT sample provides a - // fixed scan-time gate. Avoid every intermediate prefix scan and - // publish the tight full-scan boundary certificate only once. - if constexpr (kRingRefresh) { - // Ring-readback refresher: histogram already-emitted smem - // ring records (word-atomic reads) and tighten the CTA gate - // mailbox with zero additions to the math-warp hit path. - // Counting a SUBSET of emitted records keeps every published - // edge a provable lower bound on the row's current K-th - // best, and the per-(math warp, row) flush seqlock discards - // any read that raced a drain so no record counts twice. - // - // The production daemon is always active. Quiet passes use a - // fixed 2048 ns sleep; keeping this compile-time removes the - // former launch-time daemon/pacing A/B branches. - uint32_t pass = 0; - while (*scan_done_flag == 0) { - if (ring_stale[spare_id] >= 2 && - ring_stale[BLOCK_Q == 4 ? spare_id + 2 : spare_id] >= 2) { - break; - } - const uint32_t local_row = - BLOCK_Q == 4 ? spare_id + ((pass & 1u) << 1) : spare_id; - ++pass; - uint32_t fresh_total = 0; - if (block_q_idx * BLOCK_Q + local_row < seq_len && - ring_stale[local_row] < 2) { - for (uint32_t w = 0; w < kNumMathWarps; ++w) { - const uint32_t pair = w * BLOCK_Q + local_row; - const uint32_t s0 = - *reinterpret_cast(ring_seq + pair); - if (s0 != ring_shadow[pair]) { - // New flush epoch: the drained column is - // sentinel again and refills from zero. - ring_progress[pair * 32u + lane_idx] = 0u; - __syncwarp(); - if (lane_idx == 0) ring_shadow[pair] = s0; - __syncwarp(); - } - // Deep harvest: drain the lane's whole staged - // column per visit (up to kEmitLaneSlots) instead - // of one record per pass. The flood-time daemon - // bottleneck is harvest bandwidth, not publish - // cadence; reads and smem atomics only, so the - // math warps' issue ports stay untouched. - uint32_t progress = ring_progress[pair * 32u + lane_idx]; - uint32_t took = 0; - while (progress < kEmitLaneSlots) { - const uint32_t record = *reinterpret_cast( - emit_smem_records + - (pair * kEmitLaneSlots + progress) * 32u + lane_idx); - if (record == kRingSentinel) break; - if (*reinterpret_cast(ring_seq + pair) != - s0) { - // Raced a flush drain: discard this - // read; the epoch reset above replays - // the refilled column next visit. - break; - } - // Invert the truncated fp24 code back to - // the affine bucket value. Truncation only - // rounds toward better scores and never - // crosses an integer boundary (codes and - // integer edges share the 0x100-aligned - // grid), so unshifted bins support a - // pad-free published edge. The comparison - // ladder (not fmin/fmax) routes NaN and - // out-of-range scores to the top bin, which - // the publish guard below never counts. - uint32_t fbits = record & 0xffffff00u; - fbits = (fbits & 0x80000000u) ? (fbits ^ 0x80000000u) : ~fbits; - const float value = __uint_as_float(fbits); - uint32_t bucket; - if (value < 0.0f) { - bucket = 0u; - } else if (value < static_cast(num_buckets - 1)) { - bucket = static_cast(value); - } else { - bucket = num_buckets - 1; - } - atomicAdd(smem_hist + local_row * num_buckets + bucket, 1); - ++progress; - ++took; - } - if (took != 0u) { - ring_progress[pair * 32u + lane_idx] = - static_cast(progress); - } - const uint32_t fresh = __reduce_add_sync(0xffffffffu, took); - if (lane_idx == 0 && fresh != 0u) { - ring_pending[local_row] += static_cast(fresh); - } - fresh_total += fresh; - } - __syncwarp(); - // The seed base already carries K records, so each - // topk/16 batch of fresh evidence can move the - // boundary. - const int32_t pend = ring_pending[local_row]; - if (pend - ring_last[local_row] >= - static_cast(topk) / 16) { - if (lane_idx == 0) ring_last[local_row] = pend; - __syncwarp(); - // Smallest bucket prefix holding at least K ring - // records. The truncated codes and the ordered - // codes of integers <= 256 both sit on the - // 0x100-aligned grid, so a counted record's - // true code is < ord(found + 1) exactly and the - // bare bucket upper boundary is a provably safe - // published edge — no truncation pad needed. - int carry = 0; - int found = -1; - for (uint32_t bbase = 0; bbase < num_buckets && found < 0; - bbase += 32) { - const uint32_t b = bbase + lane_idx; - const int v = b < num_buckets - ? smem_hist[local_row * num_buckets + b] - : 0; - const int group_sum = __reduce_add_sync(0xffffffffu, v); - if (carry + group_sum < static_cast(topk)) { - carry += group_sum; - continue; - } - int prefix = v; -#pragma unroll - for (int off = 1; off < 32; off <<= 1) { - const int nsh = __shfl_up_sync(0xffffffffu, prefix, off); - if (static_cast(lane_idx) >= off) prefix += nsh; - } - const bool hit = carry + prefix >= static_cast(topk); - const unsigned hm = __ballot_sync(0xffffffffu, hit); - if (hm) { - found = static_cast(bbase) + __ffs(hm) - 1; - } else { - carry += __shfl_sync(0xffffffffu, prefix, 31); - } - } - if (lane_idx == 0) { - bool published = false; - // The top bin collects NaN and clamped - // out-of-range records, so an edge may only - // be published from a strictly lower bin. - if (found >= 0 && found < static_cast(num_buckets) - 1) { - const uint32_t edge = - __float_as_uint(static_cast(found + 1)); - published = - edge < atomicMin(sparse_gate_bits + local_row, edge); - } - // Early warm-start scans that fail to - // tighten only lack fresh evidence, not - // convergence; they must not retire the - // daemon. - ring_stale[local_row] = - published ? 0 - : (pend >= static_cast(topk) - ? ring_stale[local_row] + 1 - : ring_stale[local_row]); - } - __syncwarp(); - } - } - // Quiet pacing. The CTA retires promptly once the math - // warps finish; bursts skip pacing and chain straight - // into the next pass. - if (fresh_total < 64u) { - uint32_t slept = 0; - while (slept < kRingIdleNs && *scan_done_flag == 0) { - __nanosleep(kRingActiveNs); - slept += kRingActiveNs; - } - } - } - } else if constexpr (kStaticHotNoHist) { - // Scratch path: no in-scan histogram has to be finalized, so - // both spare refresh warps can retire immediately. - } else if constexpr (kStaticHotGate) { - while (*scan_done_flag == 0) __nanosleep(kSparseRefreshIdleNs); - } else { - while (*scan_done_flag == 0) { - bool tightened = false; - for (uint32_t r = spare_id; r < BLOCK_Q; r += 2) - tightened |= refresh_row(block_q_idx * BLOCK_Q + r, false); - __nanosleep(tightened ? kSparseRefreshNs : kSparseRefreshIdleNs); - } - } - if constexpr (!kStaticHotNoHist) { - for (uint32_t r = spare_id; r < BLOCK_Q; r += 2) - refresh_row(block_q_idx * BLOCK_Q + r, true); - } - } - } else if (warp_idx < kSpecWarpStart) { - cutlass::arch::warpgroup_reg_alloc(); - - const auto math_thread_idx = warp_idx * 32 + lane_idx; - - auto tmem_load = [](auto num_elems_c, const uint32_t& tmem_addr, - float* accum) { - constexpr int N = decltype(num_elems_c)::value; - DG_STATIC_ASSERT(N == 32 or N == 64 or N == 128, - "Unsupported TMEM load size"); - using Loader = cute::conditional_t< - N == 32, cute::SM100_TMEM_LOAD_32dp32b32x, - cute::conditional_t>; - [&](cute::index_sequence) { - Loader::copy(tmem_addr, reinterpret_cast(accum)[Is]...); - }(cute::make_index_sequence{}); - cutlass::arch::fence_view_async_tmem_load(); - }; - - // Bucket comparisons use the affine score space. Raw float bit - // patterns are nonuniform across exponent ranges and cannot preserve - // this threshold contract. - float weights[BLOCK_Q][kNumHeads]; - float o_reg[BLOCK_Q], inv_reg[BLOCK_Q], vth_reg[BLOCK_Q]; - uint32_t kstart_reg[BLOCK_Q], - kspan_reg[BLOCK_Q]; // unsigned range-check trick - const unsigned FULL = 0xffffffffu; - - if (block_q_idx < num_q_blocks) { - CUTE_TIE_DECL(load_schedule(block_q_idx), kv_start, num_kv_blocks); - full_q_barriers[0]->wait(0); - -// Weights into registers (once per CTA -- the 2.5-generation win). -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { -#pragma unroll - for (uint32_t j = 0; j < kNumHeads; ++j) - weights[i][j] = ptx::ld_shared(smem_weights[0] + i * kNumHeads + j); - } - // Queue fill counts are warp-uniform: every lane tracks them - // redundantly in registers, so the hot emit path needs no shared - // bookkeeping or shuffle broadcast. - // Four private 8-bit counts. They reset at every fixed-size chunk, - // bounding the address dispersion of direct global record stores. - uint32_t emit_lane_counts = 0; -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { - const uint32_t rq = min(block_q_idx * BLOCK_Q + i, seq_len - 1); - if constexpr (kOnlineFixedBuckets) { - // No sample-derived affine transform in the online A/B - // path. The raw score is converted to fixed buckets in - // the epilogue below; keep these registers compile-time - // independent of the nullable origin/inv pointers. - o_reg[i] = 0.0f; - inv_reg[i] = 1.0f; - vth_reg[i] = 0.0f; - } else { - o_reg[i] = origin[rq]; - inv_reg[i] = inv_delta[rq]; - vth_reg[i] = -o_reg[i] * inv_reg[i]; - } - o_reg[i] = 0.0f; // gate closed until the first consume - kstart_reg[i] = seq_k_start[i]; - kspan_reg[i] = - seq_k_end[i] > seq_k_start[i] ? seq_k_end[i] - seq_k_start[i] : 0; - } - // Fold -inv into the register weights: the whole ReLU-weighted - // chain then accumulates directly in bucket units. 128 FMULs - // once per qb, amortized over thousands of kv blocks. - if constexpr (!kOnlineFixedBuckets) { -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { -#pragma unroll - for (uint32_t j = 0; j < kNumHeads; ++j) weights[i][j] *= -inv_reg[i]; - } - } - - uint32_t rs_max = 0, re_min = 0xffffffffu; -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { - rs_max = max(rs_max, kstart_reg[i]); - re_min = min(re_min, kstart_reg[i] + kspan_reg[i]); - } - - if constexpr (BLOCK_Q == 4) { - const float4 initial_gate = - ptx::ld_shared(reinterpret_cast(sparse_gate_bits)); - o_reg[0] = initial_gate.x; - o_reg[1] = initial_gate.y; - o_reg[2 < BLOCK_Q ? 2 : 0] = initial_gate.z; - o_reg[3 < BLOCK_Q ? 3 : 0] = initial_gate.w; - } else { -#pragma unroll - for (uint32_t gi = 0; gi < BLOCK_Q; ++gi) - o_reg[gi] = __uint_as_float(ptx::ld_shared(sparse_gate_bits + gi)); - } - - for (uint32_t kv_block_idx = 0; kv_block_idx < num_kv_blocks; - ++kv_block_idx) { - const uint32_t kvg = kv_block_idx; - CUTE_TIE_DECL(get_kv_pipeline(kvg), kv_stage_idx, kv_phase); - - if constexpr (!kStaticHotGate || kRingRefresh) { - if (kv_block_idx != 0 && (kv_block_idx % kActiveGateStride) == 0) { - if constexpr (BLOCK_Q == 4) { - const float4 gate = ptx::ld_shared( - reinterpret_cast(sparse_gate_bits)); - o_reg[0] = gate.x; - o_reg[1] = gate.y; - o_reg[2 < BLOCK_Q ? 2 : 0] = gate.z; - o_reg[3 < BLOCK_Q ? 3 : 0] = gate.w; - } else { -#pragma unroll - for (uint32_t gi = 0; gi < BLOCK_Q; ++gi) - o_reg[gi] = - __uint_as_float(ptx::ld_shared(sparse_gate_bits + gi)); - } - } - } - full_kv_barriers[kv_stage_idx]->wait(kv_phase); - - const float scale_kv = - ptx::ld_shared(smem_kv_scales[kv_stage_idx] + math_thread_idx); - - const uint32_t umma_buf = - warpgroup_idx * kNumUmmaStages + (kvg % kNumUmmaStages); - const uint32_t umma_wait_phase = (kvg / kNumUmmaStages) & 1; - const auto tmem_start = umma_buf * UMMA_N; - full_umma_barriers[umma_buf]->wait(umma_wait_phase); - ptx::tcgen05_after_thread_sync(); - - empty_kv_barriers[kv_stage_idx]->arrive(); - - const auto kv_offset = - kv_start + kv_block_idx * BLOCK_KV + math_thread_idx; - DG_STATIC_ASSERT(kNumHeads % 8 == 0, "Invalid head"); - - uint32_t pass_bits = 0; - float v_row[BLOCK_Q]; - uint32_t fixed_payload_row[BLOCK_Q]; - uint32_t fixed_bucket_row[BLOCK_Q]; - constexpr uint32_t kTmemRowsEff = - (128u / kNumHeads) < BLOCK_Q ? (128u / kNumHeads) : BLOCK_Q; - DG_STATIC_ASSERT(BLOCK_Q % kTmemRowsEff == 0, - "BLOCK_Q must divide evenly into TMEM load groups"); -#define CANDSCAN_SCORE_GATE(i, RC) \ - bool g; \ - if constexpr (kOnlineFixedBuckets) { \ - const float raw_score = scale_kv * (sum.x + sum.y); \ - const OnlineFixedCandidate fixed = candidate_online_fixed(raw_score); \ - fixed_payload_row[i] = fixed.payload; \ - fixed_bucket_row[i] = fixed.bucket; \ - g = fixed.bucket < static_cast(o_reg[i]); \ - } else { \ - const float bq = fmaf(scale_kv, sum.x + sum.y, vth_reg[i]); \ - if constexpr (kStaticHotExactGate) { \ - const uint32_t ordered = candidate_ordered_fp32_code(bq); \ - v_row[i] = __uint_as_float(ordered); \ - g = ordered < __float_as_uint(o_reg[i]); \ - } else { \ - v_row[i] = bq; \ - g = __float_as_int(bq) < __float_as_int(o_reg[i]); \ - } \ - } \ - if constexpr (RC) g = g and ((kv_offset - kstart_reg[i]) < kspan_reg[i]); \ - pass_bits |= g ? (1u << i) : 0u; - const uint32_t kv_base = kv_start + kv_block_idx * BLOCK_KV; - const bool interior = - (kv_base >= rs_max) && (kv_base + BLOCK_KV <= re_min); -#define CANDSCAN_SCORE_ROWS(RANGE_CHECK) \ - _Pragma("unroll") for (uint32_t pr = 0; pr < BLOCK_Q / kTmemRowsEff; ++pr) { \ - float accum2[kNumHeads * kTmemRowsEff]; \ - tmem_load(cute::Int{}, \ - tmem_start + pr * kTmemRowsEff * kNumHeads, accum2); \ - if (pr == BLOCK_Q / kTmemRowsEff - 1) { \ - ptx::tcgen05_before_thread_sync(); \ - empty_umma_barriers[umma_buf]->arrive(); \ - } \ - _Pragma("unroll") for (uint32_t k = 0; k < kTmemRowsEff; ++k) { \ - const uint32_t i = pr * kTmemRowsEff + k; \ - const float* accum = accum2 + k * kNumHeads; \ - auto sum_0 = make_float2(0, 0); \ - auto sum_1 = make_float2(0, 0); \ - const auto transform = [&](const uint32_t& j, const float2& sum) { \ - auto a = make_float2(fmaxf(accum[j], 0), fmaxf(accum[j + 1], 0)); \ - auto b = make_float2(weights[i][j], weights[i][j + 1]); \ - return __ffma2_rn(a, b, sum); \ - }; \ - _Pragma("unroll") for (uint32_t j = 0; j < kNumHeads; j += 4) { \ - sum_0 = transform(j, sum_0); \ - sum_1 = transform(j + 2, sum_1); \ - } \ - auto sum = __fadd2_rn(sum_0, sum_1); \ - CANDSCAN_SCORE_GATE(i, RANGE_CHECK) \ - } \ - } - if (interior) { - CANDSCAN_SCORE_ROWS(false) - } else { - CANDSCAN_SCORE_ROWS(true) - } -#undef CANDSCAN_SCORE_ROWS -#undef CANDSCAN_SCORE_GATE - - if (pass_bits != 0) { -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { - if ((pass_bits >> i) & 1u) { - uint32_t candidate_score_bits; - uint32_t candidate_bucket; - if constexpr (kOnlineFixedBuckets) { - candidate_score_bits = fixed_payload_row[i] << 8; - candidate_bucket = fixed_bucket_row[i]; - } else { - if constexpr (kStaticHotExactGate) { - // v_row carries the already-computed - // ordered32 key in this specialization. - // The ring's low byte belongs exclusively - // to local_block; clear the discarded FP32 - // low bits before packing that coordinate. - candidate_score_bits = - __float_as_uint(v_row[i]) & 0xffffff00u; - candidate_bucket = 0; - } else { - const float x = v_row[i]; - candidate_score_bits = candidate_fp24_code(x) << 8; - const int candidate_braw = static_cast(x); - candidate_bucket = - static_cast(max(candidate_braw, 0)); - } - } - const uint32_t count = (emit_lane_counts >> (i * 8)) & 0xffu; - if (count < kEmitLaneSlots) { - const uint32_t pos = - ((warp_idx * BLOCK_Q + i) * kEmitLaneSlots + count) * 32u + - lane_idx; - const uint32_t local_block = - kv_block_idx % kActiveEmitChunkBlocks; - const uint32_t record = candidate_score_bits | local_block; - emit_smem_records[pos] = record; - emit_lane_counts += 1u << (i * 8); - } else { - // A skewed lane can overflow its small local - // quota without losing candidates. This slow - // path writes directly to the final buffer. - const uint32_t row_q = block_q_idx * BLOCK_Q + i; - const int out = atomicAdd(cand_cnt + row_q, 1); - if (out < static_cast(cand_cap)) { - uint32_t kvo = kv_offset; - if (probe_group != 0) { - const uint32_t sup = static_cast( - (static_cast(kvo) * probe_magic) >> 42); - kvo += min((sup + 1) * 64u, probe_add_max); - } - const uint64_t out_base = - static_cast(row_q) * cand_cap; - store_candidate_record(&cand_val[out_base + out], - &cand_idx[out_base + out], - candidate_score_bits >> 8, kvo); - } - } - - // A passing positive bq is below the float edge - // (gate + 1), hence already < num_buckets. - if constexpr (!kStaticHotNoHist) { - atomicAdd(smem_hist + i * num_buckets + - static_cast(candidate_bucket), - 1); - } - } - } - } - - if (((kv_block_idx + 1) % kActiveEmitChunkBlocks) == 0 || - kv_block_idx + 1 == num_kv_blocks) { - // Reserve one contiguous final-buffer segment per - // (source warp, row, chunk), then have each lane copy its - // private records into that segment. This keeps all - // collectives and the returning global atomic off the - // ordinary-hit path while deleting the rectangular global - // chunk workspace and the post-scan compactor. - // Keep the four row reservations parallel and early so - // their L2 round trips overlap the prefix work. Replace - // four independent 32-lane scans (20 shuffles) with two - // scans whose 16-bit fields carry two rows each (10 - // shuffles). A row prefix is at most 32*18=576, so fields - // cannot carry into one another. - int my_row_base = 0; - uint32_t my_row_total = 0; -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { - const uint32_t count = (emit_lane_counts >> (i * 8)) & 0xffu; - const uint32_t total = __reduce_add_sync(FULL, count); - if (lane_idx == i) my_row_total = total; - } - if (lane_idx < BLOCK_Q) { - const uint32_t row_q = block_q_idx * BLOCK_Q + lane_idx; - if (row_q < seq_len && my_row_total != 0) { - my_row_base = - atomicAdd(cand_cnt + row_q, static_cast(my_row_total)); - } - } - - const uint32_t count0 = emit_lane_counts & 0xffu; - const uint32_t count1 = (emit_lane_counts >> 8) & 0xffu; - const uint32_t count2 = (emit_lane_counts >> 16) & 0xffu; - const uint32_t count3 = emit_lane_counts >> 24; - uint32_t inclusive01 = count0 | (count1 << 16); - uint32_t inclusive23 = count2 | (count3 << 16); -#pragma unroll - for (int delta = 1; delta < 32; delta <<= 1) { - const uint32_t other01 = __shfl_up_sync(FULL, inclusive01, delta); - const uint32_t other23 = __shfl_up_sync(FULL, inclusive23, delta); - if (lane_idx >= static_cast(delta)) { - inclusive01 += other01; - inclusive23 += other23; - } - } - - const uint32_t emit_window_kv_base = - kv_start + (kv_block_idx / kActiveEmitChunkBlocks) * - kActiveEmitChunkBlocks * BLOCK_KV; -#pragma unroll - for (uint32_t i = 0; i < BLOCK_Q; ++i) { - const uint32_t count = (emit_lane_counts >> (i * 8)) & 0xffu; - const uint32_t inclusive = - i < 2 ? ((inclusive01 >> ((i & 1u) * 16)) & 0xffffu) - : ((inclusive23 >> ((i & 1u) * 16)) & 0xffffu); - const uint32_t offset = inclusive - count; - const uint32_t row_q = block_q_idx * BLOCK_Q + i; - const int out_base = __shfl_sync(FULL, my_row_base, i); - - int copy_count = 0; - if (row_q < seq_len) { - copy_count = min(static_cast(count), - max(static_cast(cand_cap) - out_base - - static_cast(offset), - 0)); - } - for (int slot = 0; slot < copy_count; ++slot) { - const uint32_t local_pos = - ((warp_idx * BLOCK_Q + i) * kEmitLaneSlots + slot) * 32u + - lane_idx; - const uint32_t record = emit_smem_records[local_pos]; - if constexpr (kRingRefresh) { - emit_smem_records[local_pos] = kRingSentinel; - } - const int out = out_base + static_cast(offset) + slot; - const uint32_t record_payload = (record & 0xffffff00u) >> 8; - uint32_t kvo = emit_window_kv_base + - ((record & 0xffu) * BLOCK_KV) + math_thread_idx; - if (probe_group != 0) { - const uint32_t sup = static_cast( - (static_cast(kvo) * probe_magic) >> 42); - kvo += min((sup + 1) * 64u, probe_add_max); - } - const uint64_t out_pos = - static_cast(row_q) * cand_cap + out; - store_candidate_record(&cand_val[out_pos], &cand_idx[out_pos], - record_payload, kvo); - } - if constexpr (kRingRefresh) { - // A capped row copies fewer slots than it - // holds; the leftovers must not leak into - // the next flush epoch. - for (int slot = copy_count; slot < static_cast(count); - ++slot) { - emit_smem_records[((warp_idx * BLOCK_Q + i) * kEmitLaneSlots + - slot) * - 32u + - lane_idx] = kRingSentinel; - } - } - } - if constexpr (kRingRefresh) { - // Publish the flush epoch: every drain store above - // is visible before the bump, and the bump is - // visible before any next-window append. - __syncwarp(); - __threadfence_block(); - if (lane_idx < BLOCK_Q) { - ring_seq[warp_idx * BLOCK_Q + lane_idx] += 1u; - } - __threadfence_block(); - } - - emit_lane_counts = 0; - } - } - - // Every actual chunk was published at its final KV block. - empty_q_barriers[0]->arrive(); - } - - // Signal the refresh daemon, then free tensor memory. - cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); - if constexpr (!kStaticHotNoHist || kRingRefresh) { - if (threadIdx.x == 0) { - __threadfence_block(); - *scan_done_flag = 1; - if constexpr (kRingRefresh) { - // Prompt-wake a daemon parked on a stage-0 parity: odd - // count so a straggler that missed transient flips - // still sees the final parity flipped. The in-wait - // exit-flag recheck makes this an optimization, not a - // correctness requirement. - asm volatile("" ::: "memory"); -#pragma unroll - for (int poke = 0; poke < 5; ++poke) { - full_kv_barriers[0]->arrive(); - } - } - } - } - if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); - } -} - -} // namespace candidate_scan_scan From 3c7036623e355d463dc15f0ec80e8254d0a4c11a Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Wed, 14 Oct 2026 09:16:59 +0000 Subject: [PATCH 03/23] further optimizations Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/candidates.cuh | 23 +- csrc/fused_mqa_topk/deep_gemm_compat.cuh | 208 +++ csrc/fused_mqa_topk/dense_scan.cuh | 284 +++++ .../{fallback_shared.cuh => fallback.cuh} | 108 +- csrc/fused_mqa_topk/fallback_core.cuh | 653 ---------- csrc/fused_mqa_topk/fast_select.cuh | 265 +--- csrc/fused_mqa_topk/fused_mqa_topk.cu | 206 ++- csrc/fused_mqa_topk/handoff.cuh | 77 -- csrc/fused_mqa_topk/mqa_seeded_kernel.cuh | 423 ------- csrc/fused_mqa_topk/persistent_epoch.inc | 1127 +++++++++++++++++ csrc/fused_mqa_topk/persistent_kernel.cuh | 550 ++++++++ csrc/fused_mqa_topk/resident_kernel.cuh | 134 ++ csrc/fused_mqa_topk/row_select.cuh | 434 +++++++ .../{sample_shared.cuh => sample.cuh} | 8 +- csrc/fused_mqa_topk/sample_core.cuh | 627 --------- csrc/fused_mqa_topk/scan.cuh | 1106 ---------------- csrc/fused_mqa_topk/upstream_select.cuh | 534 -------- csrc/ops.h | 5 +- csrc/torch_bindings.cpp | 8 +- .../kernels/attention/test_fused_mqa_topk.py | 69 +- .../kernels/attention/dsa/fused_mqa_topk.py | 32 +- 21 files changed, 3002 insertions(+), 3879 deletions(-) create mode 100644 csrc/fused_mqa_topk/deep_gemm_compat.cuh create mode 100644 csrc/fused_mqa_topk/dense_scan.cuh rename csrc/fused_mqa_topk/{fallback_shared.cuh => fallback.cuh} (66%) delete mode 100644 csrc/fused_mqa_topk/fallback_core.cuh delete mode 100644 csrc/fused_mqa_topk/handoff.cuh delete mode 100644 csrc/fused_mqa_topk/mqa_seeded_kernel.cuh create mode 100644 csrc/fused_mqa_topk/persistent_epoch.inc create mode 100644 csrc/fused_mqa_topk/persistent_kernel.cuh create mode 100644 csrc/fused_mqa_topk/resident_kernel.cuh create mode 100644 csrc/fused_mqa_topk/row_select.cuh rename csrc/fused_mqa_topk/{sample_shared.cuh => sample.cuh} (93%) delete mode 100644 csrc/fused_mqa_topk/sample_core.cuh delete mode 100644 csrc/fused_mqa_topk/scan.cuh delete mode 100644 csrc/fused_mqa_topk/upstream_select.cuh diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh index a426ee0f03e1..d19ca2496201 100644 --- a/csrc/fused_mqa_topk/candidates.cuh +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -29,15 +29,6 @@ CUTLASS_DEVICE uint32_t candidate_load_score_code(const CandidateValue value, static_cast(value); } -CUTLASS_DEVICE float candidate_decode_score(const CandidateValue value, - const int32_t packed_idx) { - const uint32_t code = candidate_load_score_code(value, packed_idx); - const uint32_t ordered = code << 8; - const uint32_t bits = - (ordered & 0x80000000u) ? (ordered ^ 0x80000000u) : ~ordered; - return __uint_as_float(bits); -} - CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { return static_cast(static_cast(packed_idx) & kCandidateIndexMask); @@ -52,20 +43,14 @@ CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, static_cast(candidate_pack_index(payload, kv_index))); } -CUTLASS_DEVICE void store_candidate_payload(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))); -} - +// Stores a candidate whose fp24 code was already computed (ring records). CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst, int32_t* index_dst, const uint32_t payload, const uint32_t kv_index) { - store_candidate_payload(value_dst, index_dst, payload, kv_index); + __stcg(value_dst, static_cast(payload)); + __stcg(index_dst, + static_cast(candidate_pack_index(payload, kv_index))); } } // namespace fused_candidates 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..2d7559ec634a --- /dev/null +++ b/csrc/fused_mqa_topk/deep_gemm_compat.cuh @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025 DeepSeek +// FP8 MQA-logits building blocks from DeepGEMM +// (github.com/deepseek-ai/DeepGEMM) as of the release before its SM100 kernels +// went MX-only: 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 (no metadata schedule, no +// multicast). +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace fused_mqa_dg { +namespace math = deep_gemm::math; + +template +struct MQALogitsSharedStorage { + static constexpr bool kIsFP4 = + cute::is_same_v; + + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using sf_dtype_t = cute::conditional_t; + + static constexpr uint32_t kNumUTCCPAlignedElems = 128; + static constexpr uint32_t kQKBytesPerElem = sizeof(qk_dtype_t); + static constexpr uint32_t kNumQKBytesPerToken = + kIsFP4 ? (kHeadDim / 2) : kHeadDim; + 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: FP4 uses head_dim / 2 bytes, FP8 use + // head_dim bytes per token + static constexpr uint32_t kSwizzleAlignment = 8 * kNumQKBytesPerToken; + static constexpr uint32_t kNumSFQ = + math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems); + static constexpr uint32_t kNumSFKV = + math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems); + static constexpr uint32_t kNumQBytesPerStage = UMMA_N * kNumQKBytesPerToken; + static constexpr uint32_t kNumKVBytesPerStage = + SPLIT_KV * kNumQKBytesPerToken; + static constexpr uint32_t kNumQElementsPerStage = + kNumQBytesPerStage / kQKBytesPerElem; + static constexpr uint32_t kNumKVElementsPerStage = + kNumKVBytesPerStage / kQKBytesPerElem; + // MX SF formats store per-block scale factors; FP8 stores one per-KV scale + // and no Q scale + static constexpr uint32_t kNumScaleQ = kIsMXSF ? kNumSFQ : 1; + static constexpr uint32_t kNumScaleKV = kIsMXSF ? kNumSFKV : SPLIT_KV; + // 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 kNumWeightElementsPerRow = + kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); + 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(kNumQBytesPerStage % kSwizzleAlignment == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(kNumKVBytesPerStage % kSwizzleAlignment == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(kSwizzleAlignment % 128 == 0, + "TMA destination must be 128-byte aligned"); + DG_STATIC_ASSERT(kTmaAlignment % 128 == 0, + "TMA destination must be 128-byte aligned"); + DG_STATIC_ASSERT(kNumWeightBytesPerStage % kTmaAlignment == 0, + "Unaligned weight TMA stage"); + + alignas(kSwizzleAlignment) + qk_dtype_t smem_q[kNumQStages][kNumQElementsPerStage]; + alignas(kSwizzleAlignment) + qk_dtype_t smem_kv[kNumKVStages][kNumKVElementsPerStage]; + alignas(kTmaAlignment) sf_dtype_t smem_sf_q[kNumQStages][kNumScaleQ]; + alignas(kTmaAlignment) sf_dtype_t smem_sf_kv[kNumKVStages][kNumScaleKV]; + 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 { + static constexpr bool kIsPaged = false; + static constexpr bool kHasPartialBlock = false; + static constexpr uint32_t kPageKV = 0; + + 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& num_tma_multicast = 1, 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_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..069bb4b8bf91 --- /dev/null +++ b/csrc/fused_mqa_topk/dense_scan.cuh @@ -0,0 +1,284 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025 DeepSeek +// 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 <= 4K resident kernel and +// by the persistent kernel's exact recovery rescan. +// +// 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, kScoreScale score = (sum_h kWeightScale * w_h * (s_h + +// |s_h|)) * kScoreScale * scale_kv begin_block(sched, q_block, query_base) +// after the block's Q / weights landed begin_split(split, query_base) 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 "deep_gemm_compat.cuh" + +namespace fused_mqa_dense { +using namespace deep_gemm; + +constexpr uint32_t kNumHeads = 32, kHeadDim = 128, BLOCK_Q = 8, SPLIT_KV = 256; +constexpr uint32_t kNumKVStages = 2, kNumTmemStages = 2; +constexpr uint32_t UMMA_M = 128, UMMA_N = BLOCK_Q * kNumHeads, UMMA_K = 32; +using Scheduler = + fused_mqa_dg::SM100MQALogitsScheduler; +using SharedStorage = fused_mqa_dg::MQALogitsSharedStorage< + kNumHeads, kHeadDim, false, BLOCK_Q, SPLIT_KV, 1, kNumKVStages, + kNumTmemStages, cutlass::float_e4m3_t, float>; + +// 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], 256 + 32); +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + init(smem.full_kv_barriers[i], 1); + init(smem.empty_kv_barriers[i], 256); + } +#pragma unroll + for (uint32_t i = 0; i < kNumTmemStages; ++i) { + init(smem.full_tmem_barriers[i], 1); + init(smem.empty_tmem_barriers[i], 256); + } + 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; + cudaGridDependencySynchronize(); + + uint32_t q_block_idx, kv_base, num_kv_splits; + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc<56>(); + 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); + fused_mqa_dg::tma_copy( + &tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); + fused_mqa_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<56>(); + 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); + fused_mqa_dg::tma_copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); + fused_mqa_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<56>(); + 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<56>(); + } else { + cutlass::arch::warpgroup_reg_alloc<224>(); + 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, query_base); + 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 % 128) + kv_half * 128; + 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(); + } + // relu(s) = (s + |s|) / 2: two FFMA2 chains per row. + float2 sum_0[2] = {}; + float2 sum_1[2] = {}; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + const float* a = accum_pair[r]; + auto transform = [&](uint32_t at, float2 sum) { + const auto x = make_float2(a[at], a[at + 1]); + const auto abs_x = + make_float2(fabsf(a[at]), fabsf(a[at + 1])); + const auto w = make_float2(weights[pair + r][at], + weights[pair + r][at + 1]); + return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); + }; + sum_0[r] = transform(j, sum_0[r]); + sum_1[r] = transform(j + 2, sum_1[r]); + } + } +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); + policy.score(pair + r, kv_offset, kv_lane, kv_split_idx, + (sum.x + sum.y) * (Policy::kScoreScale * 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(256, 0).sync(); + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +} // namespace fused_mqa_dense diff --git a/csrc/fused_mqa_topk/fallback_shared.cuh b/csrc/fused_mqa_topk/fallback.cuh similarity index 66% rename from csrc/fused_mqa_topk/fallback_shared.cuh rename to csrc/fused_mqa_topk/fallback.cuh index 83a1e1674721..607881b0a095 100644 --- a/csrc/fused_mqa_topk/fallback_shared.cuh +++ b/csrc/fused_mqa_topk/fallback.cuh @@ -1,26 +1,26 @@ +// SPDX-License-Identifier: MIT +// Exact recovery for blocks whose gated candidate count left [2048, 16384] (or +// whose weights are too small to halve exactly): the block is rescanned +// densely, and every 1024-key segment of scores is merged into a running +// per-row top set (kept <= FallbackRetain, compacted to 2048). #pragma once #include +#include "dense_scan.cuh" #include #include #include #include constexpr unsigned FallbackSegment = 1024, FallbackRetain = 4096; struct FallbackBuffers { - float* ring; float* retained; int* indices; int* output; }; struct alignas(16) FallbackShared { - cutlass::arch::ClusterBarrier full[2], empty[2]; float ring[2][8][FallbackSegment]; unsigned hist[8][1056], groups[8][32], gates[8], kept[8]; unsigned work[264], prefix, rank, written, added, emit_counts[8]; __device__ void init(unsigned tid) { - if (tid < 2) { - full[tid].init(256); - empty[tid].init(256); - } for (unsigned j = tid; j < 8 * 1056; j += 384) reinterpret_cast(hist)[j] = 0; if (tid < 256) reinterpret_cast(groups)[tid] = 0; @@ -28,7 +28,6 @@ struct alignas(16) FallbackShared { gates[tid] = __float_as_uint(-CUDART_INF_F); kept[tid] = 0; } - cutlass::arch::fence_barrier_init(); } }; __device__ __forceinline__ unsigned fb_code(float x) { @@ -176,6 +175,7 @@ __device__ void fb_compact(FallbackShared& s, FallbackBuffers b, unsigned local, } chosen = __shfl_sync(Full, chosen, chosen_lane); above = __shfl_sync(Full, above, chosen_lane); + __syncwarp(); // every lane read s.rank above if (lane == 0) { s.prefix = (prefix << 8) | chosen; s.rank = rank - above; @@ -266,3 +266,97 @@ __device__ void fb_select(FallbackShared& s, FallbackBuffers b, unsigned local, __syncwarp(); } } + +namespace fused_mqa_fallback { +using fused_mqa_dense::Scheduler; + +// 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, kScoreScale = 1.0f; + static constexpr uint32_t kSplitsPerSegment = FallbackSegment / 256; + FallbackShared& stream; + FallbackBuffers buffers; + uint32_t query_base, row_base, epoch = 0; + uint32_t seq_k_start[4], seq_k_end[4]; + float gate[4]; + CUTLASS_DEVICE void begin_block(const Scheduler& s, uint32_t q_block, + uint32_t base) { + query_base = base; + row_base = s.get_logits_row(q_block, 0); +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const uint32_t row = s.get_logits_row(q_block, base + i); + gate[i] = -CUDART_INF_F; + seq_k_start[i] = min(s.cu_seq_len_k_start[row], s.num_kv_tokens); + seq_k_end[i] = min(s.cu_seq_len_k_end[row], s.num_kv_tokens); + } + } + CUTLASS_DEVICE void begin_split(uint32_t split, uint32_t) { + if (split % kSplitsPerSegment == 0) { + __syncwarp(); + for (uint32_t i = 0; i < 4; ++i) + gate[i] = __uint_as_float(atomicAdd(stream.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]; + stream.ring[epoch & 1][query_base + i] + [(split % kSplitsPerSegment) * 256 + 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(); + fb_sync(); + const unsigned tid = threadIdx.x; + for (unsigned local = 0; local < 8; ++local) + if (stream.kept[local] > FallbackRetain - FallbackSegment) + fb_compact(stream, buffers, local, row_base + local, tid); + fb_sync(); + const unsigned local = tid / 32; + fb_select(stream, buffers, local, row_base + local, epoch, + kv_base + (split / kSplitsPerSegment) * FallbackSegment, + (split % kSplitsPerSegment + 1) * 256, tid % 32, false); + fb_sync(); + if (split + 1 == num_splits) + for (unsigned l = 0; l < 8; ++l) + if (stream.kept[l] > 2048) + fb_compact(stream, buffers, l, row_base + l, tid, true); + fb_sync(); + ++epoch; + } + CUTLASS_DEVICE void end_block(const Scheduler& s, uint32_t q_block) { + const unsigned local = threadIdx.x / 32, lane = threadIdx.x % 32; + const unsigned row = s.get_logits_row(q_block, local), + count = stream.kept[local]; + if (row < s.num_q_tokens) + for (unsigned j = lane; j < 2048; j += 32) + buffers.output[(uint64_t)row * 2048 + j] = + j < count ? buffers.indices[(uint64_t)row * FallbackRetain + j] + : -1; + } +}; + +// Rescans one block exactly; `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__ void rescan_block(const Scheduler& block, + FallbackBuffers buffers, + 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& stream = *reinterpret_cast( + smem_buffer + sizeof(fused_mqa_dense::SharedStorage)); + stream.init(threadIdx.x); + RescanPolicy policy{stream, buffers}; + fused_mqa_dense::score_block(block, policy, tq, tk, ts, tw); +} + +} // namespace fused_mqa_fallback diff --git a/csrc/fused_mqa_topk/fallback_core.cuh b/csrc/fused_mqa_topk/fallback_core.cuh deleted file mode 100644 index 95ee8619dabd..000000000000 --- a/csrc/fused_mqa_topk/fallback_core.cuh +++ /dev/null @@ -1,653 +0,0 @@ -// SPDX-License-Identifier: MIT -// Experimental derivative of DeepGEMM 2.8 sm100_mqa_logits.cuh. -// Persistent DeepGEMM FP8 MQA producer with fixed-threshold candidate emission. -#pragma once -#ifndef MQA_TMEM_ROWS - #define MQA_TMEM_ROWS 2 -#endif - -#include -#include "fallback_shared.cuh" -#ifndef MQA_HANDOFF_INTERVAL - #define MQA_HANDOFF_INTERVAL 64 -#endif -static_assert(MQA_HANDOFF_INTERVAL == 32 || MQA_HANDOFF_INTERVAL == 64 || - MQA_HANDOFF_INTERVAL == 128); - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// Shared SM100 MQA logits core plus contiguous-KV and paged entries -// Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature - -namespace seed_fallback { -using namespace deep_gemm; - -// Convert runtime valid-token count to `cute::Int` so token loops stay -// compile-time constant -template -CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, - Fn&& fn) { - if constexpr (kCandidate <= 1) { - fn(cute::Int<1>{}); - } else if (num_block_tokens >= kCandidate) { - fn(cute::Int{}); - } else { - dispatch_num_block_tokens(num_block_tokens, - static_cast(fn)); - } -} - -// Load heads using power-of-two TMEM chunks. -template -CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, - float* accum) { - constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; - if constexpr (kNumRemaining > 0) { - constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 - : kNumRemaining >= 32 ? 32 - : kNumRemaining >= 16 ? 16 - : kNumRemaining >= 8 ? 8 - : 4; - ptx::tmem_load_32dp32b( - tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); - load_tmem_heads_decomposed(tmem_col, accum); - } -} - -// Shared device core parameterized by dtype and scheduler geometry/addressing -template < - uint32_t kNumHeads, uint32_t kHeadDim, bool kIsMXSF, - bool kIsCompressedLogits, bool kCleanLogits, uint32_t BLOCK_Q, - uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages, - uint32_t kNumSMs, uint32_t kNumSpecializedThreads, uint32_t kNumMathThreads, - typename qk_dtype_t, typename logits_dtype_t, typename reduce_dtype_t, - typename MakeScheduler, uint32_t kNumMathWarpGroups = kNumMathThreads / 128> -CUTLASS_DEVICE void sm100_mqa_logits_core_impl( - const uint32_t num_q_tokens, FallbackBuffers buffers, - const cute::TmaDescriptor& tensor_map_q, - const cute::TmaDescriptor& tensor_map_sf_q, - const cute::TmaDescriptor& tensor_map_kv, - const cute::TmaDescriptor& tensor_map_sf_kv, - const cute::TmaDescriptor& tensor_map_weights, - const MakeScheduler& make_scheduler) { - constexpr bool kIsFP4 = cute::is_same_v; - - const auto sm_idx = blockIdx.x; - const auto warp_idx = cutlass::canonical_warp_idx_sync(); - const auto lane_idx = ptx::get_lane_idx(); - constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; - - const auto kNegInf = -cute::numeric_limits::infinity(); - - if (warp_idx == kSpecWarpStart) { - cute::prefetch_tma_descriptor(&tensor_map_q); - cute::prefetch_tma_descriptor(&tensor_map_sf_q); - cute::prefetch_tma_descriptor(&tensor_map_weights); - cute::prefetch_tma_descriptor(&tensor_map_kv); - cute::prefetch_tma_descriptor(&tensor_map_sf_kv); - } - - static constexpr uint32_t kNumTmemStages = 2; - static_assert(BLOCK_Q == 8 && kNumMathThreads == 256 && !kIsMXSF); - static constexpr uint32_t kNumUTCCPAlignedElems = 128; - static constexpr uint32_t UMMA_M = 128; - static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; - static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); - static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; - static constexpr uint32_t kNumSFQ = - kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kNumSFKV = - kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kNumQKBytesPerToken = - kIsFP4 ? (kHeadDim / 2) : kHeadDim; - static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = - BLOCK_QH * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = - SPLIT_KV * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = - kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; - static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = - kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); - static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( - kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); - static constexpr uint32_t kNumWeightElementsPerRow = - kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); - static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = - BLOCK_Q * kNumWeightBytesPerRow; - - DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, - "Invalid threads"); - DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and - SPLIT_KV % kNumUTCCPAlignedElems == 0, - "Invalid `SPLIT_KV`"); - DG_STATIC_ASSERT(not(kIsCompressedLogits and kCleanLogits), - "Compressed logits cannot be cleaned in-kernel"); - - using SharedStorage = layout::MQALogitsSharedStorage< - kNumHeads, kHeadDim, kIsMXSF, BLOCK_Q, SPLIT_KV, kNumQStages, - kNumKVStages, kNumTmemStages, qk_dtype_t, reduce_dtype_t>; - extern __shared__ __align__(SharedStorage::kSwizzleAlignment) - uint8_t smem_buffer[]; - auto& smem = *reinterpret_cast(smem_buffer); - constexpr uint32_t kRingSlots = 8; - constexpr uint32_t kRecordBlockMask = 255; - auto& stream = - *reinterpret_cast(smem_buffer + sizeof(SharedStorage)); - stream.init(threadIdx.x); - - constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; - constexpr uint32_t kNumTmemCols = - utils::get_num_aligned_tmem_cols(); - constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; - constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; - DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); - - if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { -#pragma unroll - for (uint32_t i = 0; i < kNumQStages; ++i) { - smem.full_q_barriers[i].init(1); - smem.empty_q_barriers[i].init(kNumMathThreads + 32); - } -#pragma unroll - for (uint32_t i = 0; i < kNumKVStages; ++i) { - smem.full_kv_barriers[i].init(1); - smem.empty_kv_barriers[i].init(kIsMXSF ? 1 : kNumMathThreads); - } -#pragma unroll - for (uint32_t i = 0; i < kNumTmemStages; ++i) { - smem.full_tmem_barriers[i].init(1); - smem.empty_tmem_barriers[i].init(256); - } - cutlass::arch::fence_barrier_init(); - } - __syncwarp(); - - if (warp_idx == kSpecWarpStart + 2) - cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); - __syncthreads(); - - RingPipeline q_pipeline; - RingPipeline kv_pipeline; - RingPipeline tmem_pipeline; - - constexpr uint32_t kNumSpecializedRegisters = 56; - constexpr uint32_t kNumMathRegisters = 224; - - cudaGridDependencySynchronize(); - - if (warp_idx == kSpecWarpStart) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - if (cute::elect_one_sync()) { - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - tma::copy( - &tensor_map_q, &smem.full_q_barriers[q_stage_idx], - smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); - if constexpr (kIsMXSF) - tma::copy( - &tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], - smem.smem_sf_q[q_stage_idx], 0, q_token_base); - 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_SF_Q_SIZE_PER_STAGE + - SMEM_WEIGHT_SIZE_PER_STAGE); - } - } - __syncwarp(); - } else if (warp_idx == kSpecWarpStart + 1) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - - auto scheduler = make_scheduler(sm_idx); - uint32_t cached_kv_page_base = 0; - uint32_t cached_kv_page_coord = 0; - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { - cached_kv_page_base = cute::numeric_limits::max(); -#pragma unroll 1 - for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; - ++kv_split_idx) { - if constexpr (decltype(scheduler)::kIsPaged) { - constexpr uint32_t kPageKV = decltype(scheduler)::kPageKV; - constexpr uint32_t kNumPagesPerSplit = - decltype(scheduler)::kNumPagesPerSplit; - DG_STATIC_ASSERT(kNumPagesPerSplit <= 32, - "Split spans more pages than a warp can cache"); - - const uint32_t kv_page_base = - (kv_base + kv_split_idx) * kNumPagesPerSplit; - if (kv_page_base < cached_kv_page_base or - kv_page_base + kNumPagesPerSplit > cached_kv_page_base + 32) { - cached_kv_page_base = (kv_page_base / 32) * 32; - cached_kv_page_coord = scheduler.get_kv_page_coord_by_page_offset( - cached_kv_page_base + lane_idx); - } - - CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); - if (cute::elect_one_sync()) - smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); - __syncwarp(); - - int page_coords[kNumPagesPerSplit]; -#pragma unroll - for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; - ++page_idx) { - const auto src_lane = - static_cast(kv_page_base - cached_kv_page_base + page_idx); - page_coords[page_idx] = - __shfl_sync(0xffffffff, cached_kv_page_coord, src_lane); - } - - if (cute::elect_one_sync()) { -#pragma unroll - for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; - ++page_idx) { - tma::copy( - &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_kv[kv_stage_idx] + - page_idx * kPageKV * kNumQKBytesPerToken, - 0, 0, 1, page_coords[page_idx]); - tma::copy( - &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_sf_kv[kv_stage_idx] + page_idx * kPageKV, 0, - page_coords[page_idx]); - } - 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 (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); - tma::copy( - &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); - 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<56>(); - DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); - - auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { - DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, - "Invalid aligned elements"); - uint32_t values[4]; -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) - values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); - __syncwarp(); - ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], - values[3]); - }; - - auto sf_desc = mma::sm100::make_sf_desc(nullptr); - - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - - if constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; - utccp_required_smem_warp_transpose(smem_ptr); - } - cutlass::arch::fence_view_async_shared(); -#pragma unroll - for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - if (cute::elect_one_sync()) - cute::SM100_UTCCP_4x32dp128bit_1cta::copy( - sf_desc, kTmemStartColOfSFQ + i * 4); - __syncwarp(); - } - } - - 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 constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; - utccp_required_smem_warp_transpose(smem_ptr); - } - cutlass::arch::fence_view_async_shared(); - } - - if (cute::elect_one_sync()) { - if constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute::SM100_UTCCP_4x32dp128bit_1cta::copy( - sf_desc, kTmemStartColOfSFKV + i * 4); - } - } -#pragma unroll - for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { - CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); - 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(); - - if constexpr (kIsMXSF) { - DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or - kHeadDim == 64 or kHeadDim == 128, - "Invalid head dim"); - - constexpr uint32_t kPackFactor = - get_smem_pack_factor(); - constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; - - using mma_op_t = - cute::conditional_t; - auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< - qk_dtype_t, qk_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, - UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K>(); -#pragma unroll - for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { - auto runtime_instr_desc = - mma::sm100::make_runtime_instr_desc_with_sf_id( - instr_desc, k * kPackFactor, k * kPackFactor); - 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); - mma_op_t::fma(a_desc, b_desc, tmem_addr, k, runtime_instr_desc, - kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ); - } - } else { - 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); -#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(); - if constexpr (kIsMXSF) - cutlass::arch::umma_arrive(reinterpret_cast( - &smem.empty_kv_barriers[kv_stage_idx])); - } - smem.empty_q_barriers[q_stage_idx].arrive(); - } - } else if (warp_idx == kSpecWarpStart + 3) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - } else if (warp_idx < kSpecWarpStart) { - cutlass::arch::warpgroup_reg_alloc<224>(); - - constexpr bool kLoadSeqBounds = true; - uint32_t seq_k_start[4]; - uint32_t seq_k_end[4]; - const auto math_warpgroup_idx = warp_idx / 4; - const auto math_thread_idx = warp_idx * 32 + lane_idx; - DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, - "Math warp groups exceed TMEM stages"); - const unsigned query_base = math_warpgroup_idx * 4; - - constexpr bool kIsReduceBF16 = not cute::is_same_v; - DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); - DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); - using weights_reg_dtype_t = - cute::conditional_t; - constexpr uint32_t kNumWeightsRegPerToken = - kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; - weights_reg_dtype_t weights[4][32]; - float accum[kNumHeads]; - - unsigned produce_epoch = 0; - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - - const auto process_q_block = [&](auto num_valid_tokens_t) { - constexpr uint32_t kNumValidTokens = 4; - uint32_t lane_counts[4] = {}; - uint32_t flush_base = 0; - float gate_score[4]; - -#pragma unroll - for (uint32_t i = 0; i < kNumValidTokens; ++i) { - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - // Use the raw-score threshold estimated from the distributed sample. - // This threshold remains fixed for the complete scan. - gate_score[i] = -CUDART_INF_F; - seq_k_start[i] = cute::min(scheduler.cu_seq_len_k_start[row], - scheduler.num_kv_tokens); - seq_k_end[i] = cute::min(scheduler.cu_seq_len_k_end[row], - scheduler.num_kv_tokens); - const auto smem_weights_row = - smem.smem_weights[q_stage_idx] + - (query_base + i) * kNumWeightElementsPerRow; - if constexpr (kIsReduceBF16) { - // Load two bf16 weights at a time as one packed shared u32 - const auto packed_row = - reinterpret_cast(smem_weights_row); -#pragma unroll - for (uint32_t j = 0; j < kNumHeads / 2; ++j) { - const auto packed = ptx::ld_shared(packed_row + j); - weights[i][j] = ptx::exchange( - *reinterpret_cast(&packed), 0); - } - } else { -#pragma unroll - for (unsigned j = 0; j < 32; ++j) - weights[i][j] = 0.5f * ptx::ld_shared(smem_weights_row + j); - } - } - - for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; - ++kv_split_idx) { - if ((kv_split_idx % (FallbackSegment / 256)) == 0) { - __syncwarp(); - for (unsigned i = 0; i < 4; ++i) - gate_score[i] = - __uint_as_float(atomicAdd(stream.gates + query_base + i, 0u)); - } - CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); - smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); -#pragma unroll - for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { - const unsigned kv_lane = (math_thread_idx % 128) + kv_half * 128; - const unsigned 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(); - } - - static_assert(!kIsReduceBF16, "FP32-only pair reduction"); - float accum_pair[MQA_TMEM_ROWS][kNumHeads]; -#pragma unroll - for (uint32_t pair = 0; pair < kNumValidTokens; - pair += MQA_TMEM_ROWS) { - // Load only MQA_TMEM_ROWS query rows at once to limit - // live accumulator registers during score reduction. - const uint32_t tmem_addr = - tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads; - ptx::tmem_load_32dp32b( - tmem_addr, reinterpret_cast(accum_pair)); - cutlass::arch::fence_view_async_tmem_load(); - if (pair + MQA_TMEM_ROWS >= kNumValidTokens) { - ptx::tcgen05_before_thread_sync(); - smem.empty_tmem_barriers[tmem_stage_idx].arrive(); - } - - float2 sum_0[MQA_TMEM_ROWS] = {}; - float2 sum_1[MQA_TMEM_ROWS] = {}; -#pragma unroll - for (uint32_t j = 0; j < kNumHeads; j += 4) { -#pragma unroll - for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { - if (pair + r >= kNumValidTokens) continue; - const uint32_t i = pair + r; - auto* a = accum_pair[r]; - auto transform = [&](uint32_t at, float2 sum) { - const auto x = make_float2(a[at], a[at + 1]); - const auto abs_x = - make_float2(fabsf(a[at]), fabsf(a[at + 1])); - const auto w = - make_float2(weights[i][at], weights[i][at + 1]); - return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); - }; - sum_0[r] = transform(j, sum_0[r]); - sum_1[r] = transform(j + 2, sum_1[r]); - } - } -#pragma unroll - for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { - if (pair + r >= kNumValidTokens) continue; - const uint32_t i = pair + r; - const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - const float reduced = sum.x + sum.y; - const float bq = reduced * static_cast(scale_kv); - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - bool valid = kv_offset >= seq_k_start[i] && - kv_offset < seq_k_end[i] && isfinite(bq) && - bq > gate_score[i]; - unsigned position = - (kv_split_idx % (FallbackSegment / 256)) * 256 + kv_lane; - stream.ring[produce_epoch & 1][query_base + i][position] = - valid ? bq : -CUDART_INF_F; - } - } - } // KV halves; query ownership is independent of the KV half. - if (((kv_split_idx + 1) % (FallbackSegment / 256)) == 0 || - kv_split_idx + 1 == num_kv_splits) { - __threadfence_block(); - - fb_sync(); - unsigned tid = math_thread_idx; - unsigned row_base = scheduler.get_logits_row(q_block_idx, 0); - for (unsigned local = 0; local < 8; ++local) - if (stream.kept[local] > FallbackRetain - FallbackSegment) - fb_compact(stream, buffers, local, row_base + local, tid); - fb_sync(); - unsigned local = tid / 32; - fb_select(stream, buffers, local, row_base + local, produce_epoch, - kv_base + (kv_split_idx / (FallbackSegment / 256)) * - FallbackSegment, - (kv_split_idx % (FallbackSegment / 256) + 1) * 256, - tid % 32, false); - fb_sync(); - if (kv_split_idx + 1 == num_kv_splits) - for (unsigned local = 0; local < 8; ++local) - if (stream.kept[local] > 2048) - fb_compact(stream, buffers, local, row_base + local, tid, - true); - fb_sync(); - ++produce_epoch; - } - } - }; - - if constexpr (decltype(scheduler)::kHasPartialBlock) - dispatch_num_block_tokens( - scheduler.get_num_block_tokens(q_block_idx), process_q_block); - else - process_q_block(cute::Int{}); - - unsigned local = warp_idx, - row = scheduler.get_logits_row(q_block_idx, local), - count = stream.kept[local]; - if (row < scheduler.num_q_tokens) - for (unsigned j = lane_idx; j < 2048; j += 32) - buffers.output[(uint64_t)row * 2048 + j] = - j < count ? buffers.indices[(uint64_t)row * FallbackRetain + j] - : -1; - cutlass::arch::fence_view_async_shared(); - smem.empty_q_barriers[q_stage_idx].arrive(); - } - - cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); - if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); - } -} - -} // namespace seed_fallback diff --git a/csrc/fused_mqa_topk/fast_select.cuh b/csrc/fused_mqa_topk/fast_select.cuh index f0a93a238e40..597733a59d9f 100644 --- a/csrc/fused_mqa_topk/fast_select.cuh +++ b/csrc/fused_mqa_topk/fast_select.cuh @@ -48,13 +48,10 @@ __device__ __forceinline__ unsigned find_bin(const unsigned* hist, } // Requires count >= K (callers recover rows with fewer candidates). -// stamps (optional): cycles for staging / histogram levels / emit. __device__ void select_row(unsigned* codes, unsigned count, const uint16_t* values, const int32_t* packed, - int32_t* out, unsigned* hist, - long long* stamps = nullptr) { + int32_t* out, unsigned* hist) { const unsigned lane = threadIdx.x & 31; - const long long t0 = clock64(); unsigned lo = 0xffffffffu, hi = 0; unsigned j = lane; for (; j + 32 * (U - 1) < count; j += 32 * U) { @@ -84,7 +81,6 @@ __device__ void select_row(unsigned* codes, unsigned count, lo = __reduce_min_sync(Full, lo); hi = __reduce_max_sync(Full, hi); __syncwarp(); - const long long t1 = clock64(); // Current bucket: codes in [lo, lo + width). rank = elements still needed // from it. @@ -118,7 +114,6 @@ __device__ void select_row(unsigned* codes, unsigned count, width = min(1u << shift, top - lo + 1); __syncwarp(); } - const long long t2 = clock64(); // lo is now the K-th largest code: emit codes above it, then ties until K. const unsigned threshold = lo, lt = (1u << lane) - 1u; @@ -145,11 +140,6 @@ __device__ void select_row(unsigned* codes, unsigned count, } } } - if (stamps && lane == 0) { - stamps[0] = t1 - t0; - stamps[1] = t2 - t1; - stamps[2] = clock64() - t2; - } } // Resident rows: fp24 codes live in shared memory as hi16[e] << 8 | lo8[e] for // the row's visible keys e = 0..count-1 (key index s0 + e). Requires count >= @@ -229,129 +219,12 @@ __device__ void select_resident(const uint16_t* hi, const uint8_t* lo, } } } -// Sample-bracketed select for a resident row (one warp, no shared atomics). 512 -// strided samples bracket the K-th largest code in [t_lo, t_hi] (+-40 sample -// ranks, ~5 sigma). One pass emits codes above t_hi straight to `out` and -// compacts the bracket into `mid` (code, offset); the exact threshold is then -// resolved inside the bracket. Returns false (caller falls back) if the bracket -// misses rank K or overflows kMidCap. -constexpr unsigned kSamplesPerLane = 16, kMidCap = 1024, - kMidBytes = kMidCap * 6; - -// Largest t with #(codes >= t) >= rank over a warp's per-lane register samples. -template -__device__ __forceinline__ unsigned sample_kth(const unsigned (&v)[N], - unsigned rank) { - unsigned t = 0; -#pragma unroll 1 - for (int bit = 23; bit >= 0; --bit) { - const unsigned cand = t | (1u << bit); - unsigned c = 0; -#pragma unroll - for (unsigned u = 0; u < N; ++u) c += v[u] >= cand; - if (__reduce_add_sync(Full, c) >= rank) t = cand; - } - return t; -} - -__device__ bool select_resident_bracket(const uint16_t* hi, const uint8_t* lo, - unsigned count, unsigned s0, - int32_t* out, unsigned char* scratch) { - const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; - constexpr unsigned S = 32 * kSamplesPerLane, kMidRegs = kMidCap / 32; - unsigned t_hi, t_lo; - { - unsigned smp[kSamplesPerLane]; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) { - const unsigned e = - unsigned((uint64_t(u * 32 + lane) * 2 + 1) * count / (2 * S)); - smp[u] = resident_code(hi, lo, e); - } - // Expected sample rank of the K-th largest code, bracketed by +-40. - const unsigned expect = unsigned((uint64_t(K) * S + count / 2) / count); - t_hi = sample_kth(smp, expect > 40 ? expect - 40 : 1); - t_lo = sample_kth(smp, min(expect + 40, S)); - } - unsigned* mid_code = reinterpret_cast(scratch); - uint16_t* mid_off = reinterpret_cast(scratch + kMidCap * 4); - unsigned above = 0, mid = 0; - // Codes are loaded U at a time before any compaction store (the stores may - // alias the loads). -#pragma unroll 1 - for (unsigned j0 = 0; j0 < count; j0 += 32 * U) { - unsigned c[U]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - const unsigned e = j0 + 32 * u + lane; - c[u] = e < count ? resident_code(hi, lo, e) : 0; - } -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - const unsigned e = j0 + 32 * u + lane; - const bool is_above = e < count && c[u] > t_hi, - is_mid = e < count && c[u] >= t_lo && c[u] <= t_hi; - const unsigned am = __ballot_sync(Full, is_above), - mm = __ballot_sync(Full, is_mid); - if (is_above) { - const unsigned at = above + __popc(am & lt); - if (at < K) out[at] = int32_t(s0 + e); - } - if (is_mid) { - const unsigned at = mid + __popc(mm & lt); - if (at < kMidCap) { - mid_code[at] = c[u]; - mid_off[at] = uint16_t(e); - } - } - above += __popc(am); - mid += __popc(mm); - } - } - if (above >= K || above + mid < K || mid > kMidCap) return false; - __syncwarp(); - // Bracket codes in registers (0 = empty: every real code in [t_lo, t_hi] is - // >= t_lo > 0 unless t_lo == 0, where empty slots only add to rank counts at - // cand == 0, never tested). - unsigned m[kMidRegs]; -#pragma unroll - for (unsigned u = 0; u < kMidRegs; ++u) { - const unsigned i = u * 32 + lane; - m[u] = i < mid ? mid_code[i] : 0; - } - // Exact K-th code: rank r among the bracket, searching only bits where t_lo - // and t_hi differ. - const unsigned r = K - above, diff = t_lo ^ t_hi; - const int top_bit = diff ? 31 - __clz(diff) : -1; - unsigned t = top_bit >= 0 ? t_lo & ~((2u << top_bit) - 1u) : t_lo; -#pragma unroll 1 - for (int bit = top_bit; bit >= 0; --bit) { - const unsigned cand = t | (1u << bit); - unsigned n = 0; -#pragma unroll - for (unsigned u = 0; u < kMidRegs; ++u) n += m[u] >= cand; - if (__reduce_add_sync(Full, n) >= r) t = cand; - } - unsigned pos = above; -#pragma unroll 1 - for (unsigned tie = 0; tie < 2; ++tie) { -#pragma unroll - for (unsigned u = 0; u < kMidRegs; ++u) { - const unsigned i = u * 32 + lane; - const bool pass = i < mid && (tie ? m[u] == t : m[u] > t); - const unsigned mask = __ballot_sync(Full, pass), - at = pos + __popc(mask & lt); - if (pass && at < K) out[at] = int32_t(s0 + mid_off[i]); - pos += __popc(mask); - } - } - return true; -} -// v3 bracket select (one warp per resident row, ~3.5x fewer instructions than -// v2): thresholds on the hi16 part only, 8 elements per 16-byte shared load, -// per-lane counts with one warp prefix scan (no per-element ballots), output -// staged in shared memory and copied out with 16-byte stores. hi_row/lo_row are -// indexed by key index. +// Register samples per lane for the sampled bracket selects below. +constexpr unsigned kSamplesPerLane = 16; +// Bracket select (one warp per resident row): thresholds on the hi16 part only, +// 8 elements per 16-byte shared load, per-lane counts with one warp prefix scan +// (no per-element ballots), output staged in shared memory and copied out with +// 16-byte stores. hi_row/lo_row are indexed by key index. constexpr unsigned kMid3Cap = 1024, kScratch3Bytes = K * 4 + kMid3Cap * 2; __device__ __forceinline__ unsigned warp_exclusive_scan(unsigned v, @@ -383,10 +256,11 @@ __device__ __forceinline__ void classify8(const uint4 v, unsigned ebase, } template -__device__ bool select_resident_v3(const uint16_t* hi_row, - const uint8_t* lo_row, unsigned s0, - unsigned count, int32_t* out, - unsigned char* scratch, KeyOf key_of = {}) { +__device__ bool select_resident_bracket(const uint16_t* hi_row, + const uint8_t* lo_row, unsigned s0, + unsigned count, int32_t* out, + unsigned char* scratch, + KeyOf key_of = {}) { const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; // hi16 bracket from 512 strided samples (+-40 sample ranks around the // expected K-th). @@ -534,13 +408,9 @@ __device__ void select_hist_generic(CodeOf code_of, unsigned count, } } } -// v4: v3 plus (a) elements whose hi16 is 0 are empty slots (sealed; no real -// score has hi16 == 0) and are excluded, with bracket ranks taken over the -// `real` elements, and (b) the exact threshold inside the bracket is resolved -// with a 512-bin shared histogram (levels until one code wide) instead of a -// 24-step register bit search. -constexpr unsigned kMidBins = 512, - kScratch4Bytes = K * 4 + kMid3Cap * 2 + kMidBins * 4; +// Bins of the histogram levels that resolve the exact threshold inside a +// bracket (row_select.cuh). +constexpr unsigned kMidBins = 512; template __device__ __forceinline__ unsigned find_bin_n(const unsigned* hist, @@ -575,109 +445,4 @@ __device__ __forceinline__ unsigned find_bin_n(const unsigned* hist, return __shfl_sync(Full, chosen, winner); } -template -__device__ bool select_resident_v4(const uint16_t* hi_row, - const uint8_t* lo_row, unsigned s0, - unsigned count, unsigned real, int32_t* out, - unsigned char* scratch, KeyOf key_of = {}) { - const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; - unsigned h_hi = 0, h_lo = 0; - { - unsigned smp[kSamplesPerLane], n_real = 0; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) { - smp[u] = hi_row[s0 + (((u * 32 + lane) * 2 + 1) * count >> 10)]; - n_real += smp[u] != 0; - } - const unsigned s_real = __reduce_add_sync(Full, n_real); - if (s_real == 0 || real <= K) return false; - // Expected sample rank of the K-th largest real element, bracketed by +-40. - const unsigned expect = (K * s_real + real / 2) / real; - const unsigned r_hi = expect > 40 ? expect - 40 : 1, - r_lo = min(expect + 40, s_real); -#pragma unroll 1 - for (int bit = 15; bit >= 0; --bit) { - const unsigned ch = h_hi | (1u << bit), cl = h_lo | (1u << bit); - unsigned nh = 0, nl = 0; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) { - nh += smp[u] >= ch; - nl += smp[u] >= cl; - } - if (__reduce_add_sync(Full, nh) >= r_hi) h_hi = ch; - if (__reduce_add_sync(Full, nl) >= r_lo) h_lo = cl; - } - h_lo = max(h_lo, 1u); - } - const unsigned a0 = s0 & ~7u, end = s0 + count, units = (end - a0 + 7) / 8; - const uint4* hv = reinterpret_cast(hi_row + a0); - unsigned na = 0, nm = 0; -#pragma unroll 2 - for (unsigned unit = lane; unit < units; unit += 32) - classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, - [&](unsigned, bool a, bool m) { - na += a; - nm += m; - }); - unsigned above, mid; - unsigned pa = warp_exclusive_scan(na, above), - pm = warp_exclusive_scan(nm, mid); - if (above >= K || above + mid < K || mid > kMid3Cap) return false; - int32_t* stage = reinterpret_cast(scratch); - uint16_t* midbuf = reinterpret_cast(scratch + K * 4); - unsigned* hist = reinterpret_cast(scratch + K * 4 + kMid3Cap * 2); -#pragma unroll 2 - for (unsigned unit = lane; unit < units; unit += 32) - classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, - [&](unsigned e, bool a, bool m) { - if (a) stage[pa++] = key_of(e); - if (m) midbuf[pm++] = uint16_t(e); - }); - __syncwarp(); - const auto code_at = [&](unsigned i) { - const unsigned e = midbuf[i]; - return unsigned(hi_row[e]) << 8 | lo_row[e]; - }; - // K-th code among the bracket (rank r): histogram levels over [h_lo << 8, - // h_hi << 8 | 0xff]. - unsigned rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; -#pragma unroll 1 - while (top > base) { - const unsigned width = top - base + 1, spread = (width - 1) / kMidBins; - const unsigned shift = spread ? 32 - __clz(spread) : 0; - for (unsigned b = lane; b < kMidBins; b += 32) hist[b] = 0; - __syncwarp(); - for (unsigned i = lane; i < mid; i += 32) { - const unsigned c = code_at(i); - if (c >= base && c <= top) atomicAdd(hist + ((c - base) >> shift), 1u); - } - __syncwarp(); - unsigned above_bin; - const unsigned chosen = find_bin_n(hist, rank, above_bin); - rank -= above_bin; - base += chosen << shift; - top = min(base + (1u << shift) - 1, top); - __syncwarp(); - } - const unsigned t = base; - unsigned pos = above; -#pragma unroll 1 - for (unsigned tie = 0; tie < 2; ++tie) { - for (unsigned j = 0; j < mid && pos < K; j += 32) { - const unsigned i = j + lane; - const unsigned c = i < mid ? code_at(i) : 0; - const bool pass = i < mid && (tie ? c == t : c > t); - const unsigned mask = __ballot_sync(Full, pass), - at = pos + __popc(mask & lt); - if (pass && at < K) stage[at] = key_of(midbuf[i]); - pos += __popc(mask); - } - } - __syncwarp(); -#pragma unroll 4 - for (unsigned j = lane * 4; j < K; j += 128) - *reinterpret_cast(out + j) = - *reinterpret_cast(stage + j); - return true; -} } // namespace fast_select diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu index f6ce592e9fe8..b7a85bb4a98b 100644 --- a/csrc/fused_mqa_topk/fused_mqa_topk.cu +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -2,12 +2,13 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // // Single-launch fused FP8 MQA-logits + top-2048 prefill selection (SM100 -// family). One kernel computes the indexer scores of every visible key of a -// row and selects its top-2048 keys without materializing the [rows, keys] -// logits matrix. The device code lives in the headers of this directory (kept -// verbatim from the isolated kernel; see mqa_seeded_kernel.cuh); 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. +// family): one kernel scores every visible key of a row and selects its +// top-2048 keys without materializing the [rows, keys] logits matrix. +// <= 4K keys: resident_kernel.cuh (all scores of a block in shared memory). +// > 4K keys: persistent_kernel.cuh (one CTA per SM; sampled gate, candidate +// emission and selection; exact rescan of out-of-range blocks). +// 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" @@ -20,12 +21,8 @@ #include #include -#include "sample_core.cuh" -#include "scan.cuh" -#include "fallback_core.cuh" -#include "upstream_select.cuh" -#include "fast_select.cuh" -#include "mqa_seeded_kernel.cuh" +#include "persistent_kernel.cuh" +#include "resident_kernel.cuh" namespace { @@ -107,32 +104,7 @@ constexpr int kHeads = 32; constexpr int kHeadDim = 128; constexpr int kTopK = 2048; constexpr int kCandidates = 16384; // per-row candidate slab -constexpr int kRetained = 4096; // per-row fallback slab -constexpr int kStatsPerCta = 7; - -// Gate target / sample rank of the <= 16K-key path (candidates kept in shared -// memory, kCandCap per row). -constexpr unsigned kSmallTarget = 2560; -constexpr unsigned kSmallSampleRank = 100; - -using SampleStorage = - deep_gemm::layout::MQALogitsSharedStorage<32, 128, false, 8, 256, 1, 2, 2, - cutlass::float_e4m3_t, float>; -using ScanStorage = - deep_gemm::layout::MQALogitsSharedStorage<32, 128, false, 8, 256, 1, 3, 2, - cutlass::float_e4m3_t, float>; -// Sample / fallback phases. -constexpr int kPhaseSmem = - sizeof(SampleStorage) + (sizeof(SampleShared) > sizeof(FallbackShared) - ? sizeof(SampleShared) - : sizeof(FallbackShared)); -// > 16K keys: 3-stage scan + candidate ring. -constexpr int kLongSmem = - std::max(kPhaseSmem, sizeof(ScanStorage) + sizeof(CandidateHandoff<>)); -// <= 16K keys: 2-stage scan + per-row shared candidates (5 bytes each). -constexpr int kSmallSmem = - std::max(kPhaseSmem, sizeof(SampleStorage) + kCandExt * 5); -static_assert(kSmallSmem + 2048 <= 232448, "shared candidates must fit"); +constexpr int kRetained = 4096; // per-row recovery slab void check_tensor(const torch::stable::Tensor& t, int32_t device, ScalarType dtype, const char* name) { @@ -150,23 +122,21 @@ void check_tensor(const torch::stable::Tensor& t, int32_t device, // (columns of the logits), or with relative=true indices relative to the // row's start (key - start, the top_k_per_row_prefill convention); -1 padded // when fewer than 2048 keys are visible. -// 0 < rows <= 16384, 0 < keys <= 1048576 (any values: the last 8-row block -// may be partial). -// Workspace (caller-owned, reused across calls), with B = ceil(rows / 8): -// slot_flags [S] int32 (all zero between calls; the kernel restores them), -// values [S*8,16384] fp16, indices [S*8,16384] int32, -// retained [S*8,4096] fp32, retained_indices [S*8,4096] int32, -// counts [>=8B] int32, gates [>=8B] fp32, stats [>=7B] int32 (per-CTA -// diagnostics). +// 0 < rows <= 16384, 0 < keys <= 1048576 (the last 8-row block may be partial). +// 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), +// values [8S,16384] fp16, indices [8S,16384] int32, +// 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 slot_flags, + 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 gates, - torch::stable::Tensor stats, bool relative) { + torch::stable::Tensor counts, torch::stable::Tensor defer, bool relative) { 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"); @@ -176,14 +146,13 @@ void fused_mqa_topk_prefill( check_tensor(starts, dev, ScalarType::Int, "starts"); check_tensor(ends, dev, ScalarType::Int, "ends"); check_tensor(out, dev, ScalarType::Int, "out"); - check_tensor(slot_flags, dev, ScalarType::Int, "slot_flags"); + check_tensor(counter, dev, ScalarType::Int, "counter"); check_tensor(values, dev, ScalarType::Half, "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(gates, dev, ScalarType::Float, "gates"); - check_tensor(stats, dev, ScalarType::Int, "stats"); + check_tensor(defer, dev, ScalarType::Int, "defer"); STD_TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim, "q must be [rows, 32, 128]"); @@ -205,88 +174,69 @@ void fused_mqa_topk_prefill( STD_TORCH_CHECK(reinterpret_cast(out.data_ptr()) % 16 == 0, "out must be 16-byte aligned"); - const int64_t total_slots = slot_flags.numel(); - STD_TORCH_CHECK(total_slots > 0, "slot_flags must not be empty"); - STD_TORCH_CHECK(values.dim() == 2 && values.size(0) == total_slots * 8 && - values.size(1) == kCandidates && indices.dim() == 2 && - indices.size(0) == total_slots * 8 && - indices.size(1) == kCandidates, - "values/indices must be [8 * slots, 16384]"); - STD_TORCH_CHECK(retained.dim() == 2 && retained.size(0) == total_slots * 8 && - retained.size(1) == kRetained && - retained_indices.dim() == 2 && - retained_indices.size(0) == total_slots * 8 && - retained_indices.size(1) == kRetained, - "retained/retained_indices must be [8 * slots, 4096]"); - STD_TORCH_CHECK(counts.numel() >= blocks * 8 && gates.numel() >= blocks * 8 && - stats.numel() >= blocks * kStatsPerCta, - "per-row workspace too small"); - torch::stable::accelerator::DeviceGuard guard(dev); const cudaStream_t stream = get_current_cuda_stream(dev); - - // Pool of 8-row slots: the kernel claims a free slot per CTA, so the - // candidate workspace is bounded independently of rows. The caller sizes it - // with >= 2 slots per SM (>= 2x resident CTAs, so a free slot always exists). - const unsigned num_slots = - static_cast(std::min(total_slots, blocks)); - - FallbackBuffers fallback{nullptr, static_cast(retained.data_ptr()), - static_cast(retained_indices.data_ptr()), - static_cast(out.data_ptr())}; - const int irows = static_cast(rows), ikeys = static_cast(keys); + const unsigned urows = static_cast(rows), + ukeys = static_cast(keys), + ublocks = static_cast(blocks); + 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()); auto tq = make_2d(q.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, - irows * 32, 128, 256, 128, 128); - auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, ikeys, - 128, 256, 128, 128); + int(rows) * 32, 128, 256, 128, 128); + auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, + int(keys), 128, 256, 128, 128); auto ts = make_2d(scales.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, - ikeys, 1, 256, 1, 0, 0); + int(keys), 1, 256, 1, 0, 0); auto tw = make_2d(weights.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, 32, - irows, 32, 8, 32, 0); - // Gate sample: <= 16K keys aim for kSmallTarget candidates (shared memory), - // <= 128K for 3072, both with a sample sized so the gate's quantile rank is - // ~100 (about 10% count error). Longer contexts use 16 tiles and the - // 4096/6144 target, where candidates are already ~1% of the keys. - // (<= 4K keys keep every score in shared memory and skip the sample.) - const bool low_target = keys <= 131072; - const unsigned sample_target = - keys <= 16384 ? kSmallTarget : (low_target ? 3072u : 0u); - const unsigned sample_rank = keys <= 16384 ? kSmallSampleRank : 100u; - const unsigned sample_tiles = - low_target ? std::max(2u, (unsigned)((sample_rank * keys + - 256ull * sample_target - 1) / - (256ull * sample_target))) - : 16u; - const unsigned sample_stride = std::max( - 1u, (unsigned)((keys + 256u * sample_tiles - 1) / (256u * sample_tiles))); - const unsigned sample_groups = - (keys - sample_stride / 2 + sample_stride - 1) / sample_stride; - auto tks = make_strided_kv(k.data_ptr(), sample_groups, sample_stride); + int(rows), 32, 8, 32, 0); + + if (ukeys <= fused_mqa_resident::kResidentKeys) { + // Every score of a block fits in shared memory: no workspace. + constexpr int smem = int(fused_mqa_resident::kSmemBytes); + check_cuda(cudaFuncSetAttribute( + reinterpret_cast(fused_mqa_topk_fp8_resident_kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + fused_mqa_topk_fp8_resident_kernel<<>>( + urows, ukeys, starts_p, ends_p, out_p, relative, tq, ts, tk, tw); + check_cuda(cudaGetLastError()); + return; + } - const int smem = keys <= int64_t(kResidentKeys) ? int(kResidentSmem) - : keys <= 16384 ? kSmallSmem - : kLongSmem; - auto kernel = keys <= int64_t(kResidentKeys) ? mqa_seeded_kernel<0> - : keys <= 8192 ? mqa_seeded_kernel<4> - : keys <= 16384 ? mqa_seeded_kernel<8> - : keys <= 32768 ? mqa_seeded_kernel<16> - : keys <= 131072 ? mqa_seeded_kernel<32> - : keys <= 524288 ? mqa_seeded_kernel<64> - : mqa_seeded_kernel<128>; - check_cuda(cudaFuncSetAttribute(reinterpret_cast(kernel), - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem)); - kernel<<(blocks), 384, smem, stream>>>( - static_cast(rows), static_cast(keys), - static_cast(starts.data_ptr()), - static_cast(ends.data_ptr()), - static_cast(gates.data_ptr()), - static_cast(values.data_ptr()), + STD_TORCH_CHECK(counter.numel() == 2, "counter must be [2]"); + STD_TORCH_CHECK(values.dim() == 2 && values.size(1) == kCandidates && + values.size(0) % 16 == 0 && indices.dim() == 2 && + indices.size(0) == values.size(0) && + indices.size(1) == kCandidates, + "values/indices must be [8 * slots, 16384], slots even"); + STD_TORCH_CHECK(retained.dim() == 2 && retained.size(1) == kRetained && + retained.size(0) == values.size(0) && + retained_indices.dim() == 2 && + retained_indices.size(0) == retained.size(0) && + retained_indices.size(1) == kRetained, + "retained/retained_indices must be [8 * slots, 4096]"); + STD_TORCH_CHECK(counts.numel() >= blocks * 8 && 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. + int num_sms = 0; + check_cuda( + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev)); + const unsigned grid = static_cast( + std::min({int64_t(num_sms), blocks, values.size(0) / 16})); + fused_persist::LaunchConfig c; + check_cuda(fused_persist::launch_config(ukeys, c)); + auto tks = make_strided_kv(k.data_ptr(), c.sample_groups, c.sample_stride); + const FallbackBuffers fallback{static_cast(retained.data_ptr()), + static_cast(retained_indices.data_ptr()), + out_p}; + c.kernel<<>>( + urows, ukeys, starts_p, ends_p, static_cast(values.data_ptr()), static_cast(indices.data_ptr()), - static_cast(counts.data_ptr()), static_cast(stats.data_ptr()), - fallback, static_cast(slot_flags.data_ptr()), num_slots, - static_cast(kLongSmem), sample_target, sample_stride, - sample_groups, sample_tiles, static_cast(scales.data_ptr()), - tks, tq, ts, tk, tw, relative); + static_cast(counts.data_ptr()), fallback, 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, + relative, ublocks, static_cast(counter.data_ptr()), + static_cast(defer.data_ptr())); check_cuda(cudaGetLastError()); } diff --git a/csrc/fused_mqa_topk/handoff.cuh b/csrc/fused_mqa_topk/handoff.cuh deleted file mode 100644 index 1988e28a5f5f..000000000000 --- a/csrc/fused_mqa_topk/handoff.cuh +++ /dev/null @@ -1,77 +0,0 @@ -// SPDX-License-Identifier: MIT -#pragma once -#include -#include - -// Each producer warp owns two banks. A single drainer visits producer warps -// in order, consuming one published batch at a time. Epochs span query blocks. -template -struct alignas(16) CandidateHandoff { - using Barrier = cutlass::arch::ClusterTransactionBarrier; - Barrier full[2][P], empty[2][P]; - uint32_t records[2][P][R][S][32]; - uint32_t counts[2][P][R][32]; - uint32_t bases[2][P]; - __device__ void init(int tid) { - if (tid < 2 * P) { - full[tid / P][tid % P].init(32); - empty[tid / P][tid % P].init(32); - } - cutlass::arch::fence_barrier_init(); - } - __device__ void acquire(int producer, unsigned epoch) { - empty[epoch & 1][producer].wait(((epoch >> 1) & 1) ^ 1); - __syncwarp(); - } - __device__ void publish(int producer, unsigned epoch, unsigned base) { - if ((threadIdx.x & 31) == 0) bases[epoch & 1][producer] = base; - __syncwarp(); - full[epoch & 1][producer].arrive(); - } - template - __device__ void drain(int producer, unsigned epoch, unsigned rows, - unsigned row_base, int* totals, unsigned cap, - Store store) { - const int lane = threadIdx.x & 31, bank = epoch & 1; - full[bank][producer].wait((epoch >> 1) & 1); - __syncwarp(); - const unsigned base = bases[bank][producer]; - // Two 16-bit prefix sums share each shuffle. A row contributes at - // most 32*S records, so the low half cannot carry into the high half. - static_assert(R == 4 && 32 * S < 65536); - unsigned n[4]; - unsigned my_total = 0; -#pragma unroll - for (unsigned row = 0; row < 4; ++row) { - n[row] = counts[bank][producer][row][lane]; - unsigned total = __reduce_add_sync(0xffffffffu, n[row]); - if (lane == row) my_total = total; - } - unsigned my_start = 0; - if (lane < 4 && my_total) - my_start = atomicAdd(totals + row_base + lane, my_total); - unsigned p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); -#pragma unroll - for (int d = 1; d < 32; d *= 2) { - unsigned a = __shfl_up_sync(0xffffffffu, p01, d); - unsigned b = __shfl_up_sync(0xffffffffu, p23, d); - if (lane >= d) { - p01 += a; - p23 += b; - } - } - unsigned prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; -#pragma unroll - for (unsigned row = 0; row < 4; ++row) { - unsigned start = __shfl_sync(0xffffffffu, my_start, row); - for (unsigned j = 0; j < n[row]; ++j) { - unsigned slot = start + prefix[row] - n[row] + j; - unsigned record = records[bank][producer][row][j][lane]; - if (slot < cap) - store(row_base + row, slot, record, base, producer * 32 + lane); - } - } - __syncwarp(); - empty[bank][producer].arrive(); - } -}; diff --git a/csrc/fused_mqa_topk/mqa_seeded_kernel.cuh b/csrc/fused_mqa_topk/mqa_seeded_kernel.cuh deleted file mode 100644 index 36a2165e06d5..000000000000 --- a/csrc/fused_mqa_topk/mqa_seeded_kernel.cuh +++ /dev/null @@ -1,423 +0,0 @@ -// SPDX-License-Identifier: MIT -// Generated by experiments/integration_v3/sync_candidate.py from -// experiments/fused_clean/mqa_fixed_cuda.cu -- do not edit by hand. -// Device side of the single-launch fused MQA-logits + top-k prefill kernel. -#pragma once -// Sample tile t holds keys (t*256 + j)*S + S/2 (j < 256) through a strided -// {128, G, S} KV view: all tiles together form one uniform grid over every key -// position. -struct SampleScheduler - : deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false> { - using Base = deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false>; - unsigned sample_stride, sample_groups, sample_tiles; - const float* sample_scales; - __device__ SampleScheduler(unsigned sm, unsigned rows, unsigned keys, - const unsigned* starts, const unsigned* ends, - unsigned stride, unsigned groups, unsigned tiles, - const float* scales) - : Base(sm, rows, keys, starts, ends, nullptr), - sample_stride(stride), - sample_groups(groups), - sample_tiles(tiles), - sample_scales(scales) {} - __device__ bool next_q_block(unsigned& q, unsigned& base, unsigned& splits) { - if (!Base::next_q_block(q, base, splits)) return false; - splits = sample_tiles; - return true; - } - __device__ unsigned sample_center() const { return sample_stride / 2; } - __device__ unsigned sample_group(unsigned split, unsigned lane) const { - return split * 256 + lane; - } - __device__ unsigned get_logits_col(unsigned, unsigned split, - unsigned lane) const { - return sample_group(split, lane) * sample_stride + sample_center(); - } -}; -constexpr unsigned kResidentKeys = 4096; -// <=16K keys (handoff interval <= 8): candidates stay in shared memory, cap per -// row. -constexpr unsigned kCandCap = 3072; -// Each candidate array: 8 rows of kCandCap plus 256 per-lane trash slots. -constexpr unsigned kCandExt = 8u * kCandCap + 256u; -constexpr unsigned kResidentSmem = - ((sizeof(deep_gemm::layout::MQALogitsSharedStorage< - 32, 128, false, 8, 256, 1, 2, 2, cutlass::float_e4m3_t, float>) + - 8u * kResidentKeys * 3u) + - 1023u) / - 1024u * 1024u; -template -__global__ __launch_bounds__(384, 1) void mqa_seeded_kernel( - unsigned rows, unsigned keys, const unsigned* starts_g, - const unsigned* ends_g, float* gates, uint16_t* vals, int* indices, - int* counts, int* stats, FallbackBuffers fallback, int* slot_flags, - unsigned num_slots, unsigned phase_bytes, unsigned sample_target, - unsigned sample_stride, unsigned sample_groups, unsigned 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, bool relative) { - // Blocks run heaviest (latest rows, most keys) first so the last wave holds - // the light ones. - const unsigned vblock = gridDim.x - 1 - blockIdx.x; - unsigned row_base = vblock * 8; - bool all_short = true; - // The block's row bounds, clamped to the keys, in shared memory. Rows past - // the end (last block) see no keys; `starts`/`ends` are indexed by absolute - // row like the global arrays. - __shared__ unsigned row_bounds[2][8]; - if (threadIdx.x < 8) { - const unsigned 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 unsigned* const starts = row_bounds[0] - row_base; - const unsigned* const ends = row_bounds[1] - row_base; - // relative: outputs are key - start (top_k_per_row_prefill's convention). - // Rows with - // <= 2048 visible keys write that directly (base 0); selected rows (and every - // row of a recovered block) are rewritten in place by their block once final - // (L2-hot). - const auto finish_rows = [&](bool all_rows) { - if (!relative) return; - __syncthreads(); - for (unsigned i = threadIdx.x; i < 8 * 512; i += 384) { - const unsigned r = row_base + i / 512; - if (r >= rows || (!all_rows && ends[r] - starts[r] <= 2048u)) continue; - const int s = int(starts_g[r]); - if (s == 0) continue; - int4* p = reinterpret_cast(fallback.output + (uint64_t)r * 2048) + - i % 512; - 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; - } - }; - for (unsigned r = row_base; r < row_base + 8; ++r) - all_short &= (ends[r] - starts[r] <= 2048); - if (all_short) { - if (threadIdx.x == 0) { - for (int c = 0; c < 7; ++c) stats[blockIdx.x * 7 + c] = 0; - } - for (unsigned r = row_base; r < min(row_base + 8, rows); ++r) - for (unsigned j = threadIdx.x; j < 2048; j += 384) - fallback.output[(uint64_t)r * 2048 + j] = - j < ends[r] - starts[r] ? int((relative ? 0u : starts[r]) + j) : -1; - return; - } - if constexpr (Interval == 0) { - // Resident (keys <= 4096): scores of all 8 rows stay in shared memory as - // fp24 codes, so there is no sample, gate, candidate workspace or slot - // claim. One warp selects one row. - if (threadIdx.x == 0) { - for (int c = 0; c < 7; ++c) stats[blockIdx.x * 7 + c] = 0; - } - const long long t0 = clock64(); - using RShared = deep_gemm::layout::MQALogitsSharedStorage< - 32, 128, false, 8, 256, 1, 2, 2, cutlass::float_e4m3_t, float>; - extern __shared__ __align__(1024) unsigned char phase_smem[]; - unsigned char* const resident = phase_smem + sizeof(RShared); - auto scheduler = [&](const unsigned& sm) { - return deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false>( - vblock, rows, keys, starts, ends, nullptr); - }; - fused_mqa_scan::sm100_mqa_logits_core_impl< - 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, - cutlass::float_e4m3_t, float, float, decltype(scheduler), 2, 0>( - rows, kResidentKeys, reinterpret_cast(resident), nullptr, - nullptr, nullptr, nullptr, vals, indices, counts, 16384, 0, tq, ts, tk, - ts, tw, scheduler); - __syncthreads(); - const long long t1 = clock64(); - if (threadIdx.x < 256) cutlass::arch::warpgroup_reg_dealloc<128>(); - __syncthreads(); - if (threadIdx.x >= 256) cutlass::arch::warpgroup_reg_alloc<128>(); - const unsigned warp = threadIdx.x / 32; - if (warp < 8 && row_base + warp < rows) { - const unsigned r = row_base + warp, s0 = min(starts[r], keys), - visible = min(ends[r], keys) - s0; - int* o = fallback.output + (uint64_t)r * 2048; - if (visible <= 2048) { - for (unsigned j = threadIdx.x & 31; j < 2048; j += 32) - o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; - } else { - const uint16_t* hi = reinterpret_cast(resident) + - warp * kResidentKeys + s0; - const uint8_t* lo = - resident + 8 * kResidentKeys * 2 + warp * kResidentKeys + s0; - // Scratch (bracket buffer, or the fallback's histogram) reuses the idle - // Q/KV stage memory. - static_assert(8 * fast_select::kScratch3Bytes <= sizeof(RShared) && - fast_select::kScratch3Bytes >= fast_select::kBins * 4, - "select scratch must fit the Q/KV stages"); - unsigned char* scratch = - phase_smem + warp * fast_select::kScratch3Bytes; - if (!fast_select::select_resident_v3(hi - s0, lo - s0, s0, visible, o, - scratch)) { - __syncwarp(); - if ((threadIdx.x & 31) == 0) atomicAdd(stats + blockIdx.x * 7 + 6, 1); - fast_select::select_resident(hi, lo, visible, s0, o, - reinterpret_cast(scratch)); - } - } - } - __syncthreads(); - if (threadIdx.x == 0) { - stats[blockIdx.x * 7 + 4] = int(t1 - t0); - stats[blockIdx.x * 7 + 5] = int(clock64() - t1); - } - finish_rows(false); - return; - } else { - // Workspace pool: claim one of num_slots 8-row slots (num_slots >= 2x - // resident CTAs, so a free slot always exists and holders never wait). Base - // pointers are shifted so global row r of this block addresses slot row - // slot*8 + (r - row_base) in the cores unchanged. - __shared__ unsigned slot; - if (threadIdx.x == 0) { - unsigned s0 = blockIdx.x % num_slots; - for (unsigned i = 0;; ++i) { - const unsigned c = (s0 + i) % num_slots; - if (atomicCAS(slot_flags + c, 0, 1) == 0) { - slot = c; - break; - } - } - } - __syncthreads(); - // Slot rows are addressed as (row + row_delta): reassigning the pointer - // parameters instead costs ~10% (they stop being launch constants in the - // fused scan). - const int64_t row_delta = int64_t(slot) * 8 - int64_t(row_base); - if (threadIdx.x < 8) counts[row_base + threadIdx.x] = 0; - __syncthreads(); - auto sample_scheduler = [&](const unsigned& sm) { - return SampleScheduler(vblock, rows, keys, starts, ends, sample_stride, - sample_groups, sample_tiles, scales); - }; - SampleBuffers sample{gates, nullptr, nullptr, fallback.output, - sample_target}; - seed_sample::sm100_mqa_logits_core_impl< - 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, - cutlass::float_e4m3_t, float, float, decltype(sample_scheduler), 2>( - rows, sample, tq, ts, tks, ts, tw, sample_scheduler); - __syncthreads(); - auto scheduler = [&](const unsigned& sm) { - return deep_gemm::sched::SM100MQALogitsScheduler<8, 256, 4096, false>( - vblock, rows, keys, starts, ends, nullptr); - }; - constexpr bool kSmemCand = Interval <= 8; - using CShared = deep_gemm::layout::MQALogitsSharedStorage< - 32, 128, false, 8, 256, 1, 2, 2, cutlass::float_e4m3_t, float>; - extern __shared__ __align__(1024) unsigned char cand_smem[]; - unsigned char* const cand = cand_smem + sizeof(CShared); - if constexpr (kSmemCand) - fused_mqa_scan::sm100_mqa_logits_core_impl< - 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, - cutlass::float_e4m3_t, float, float, decltype(scheduler), 2, Interval, - true>(rows, kCandCap, reinterpret_cast(cand), gates, nullptr, - nullptr, nullptr, vals, indices, counts, 16384, row_delta, tq, - ts, tk, ts, tw, scheduler); - else - fused_mqa_scan::sm100_mqa_logits_core_impl< - 32, 128, false, false, false, 8, 256, 1, 3, 4096, 128, 256, - cutlass::float_e4m3_t, float, float, decltype(scheduler), 2, - Interval>(rows, 0, nullptr, gates, nullptr, nullptr, nullptr, vals, - indices, counts, 16384, row_delta, tq, ts, tk, ts, tw, - scheduler); - __threadfence_block(); - __syncthreads(); - __shared__ unsigned recover, staged, short_rows; - __shared__ unsigned row_offset[9]; - if (threadIdx.x == 0) { - recover = 0; - short_rows = 0; - int lo = 2147483647, hi = 0; - row_offset[0] = 0; - for (unsigned r = row_base; r < row_base + 8; ++r) { - // Rows with <= 2048 visible keys output all of them inline; only long - // rows with a bad count recover. - const bool is_short = - min(ends[r], keys) - min(starts[r], keys) <= 2048u; - short_rows |= unsigned(is_short) << (r - row_base); - const int real = kSmemCand ? (counts[r] & ((1 << 21) - 1)) : counts[r]; - recover |= !is_short && (real < 2048 || real > 16384); - if (!is_short) { - lo = min(lo, counts[r]); - hi = max(hi, counts[r]); - } - row_offset[r - row_base + 1] = - row_offset[r - row_base] + (is_short ? 0 : max(counts[r], 0)); - } - // Fast select stages every row's codes after 8 warp-private 1024-bin - // histograms. staged: 1 = all 8 rows in one round, 2 = two rounds of 4 - // rows (long contexts), 0 = legacy. - const unsigned half = max(row_offset[4], row_offset[8] - row_offset[4]); - staged = - recover - ? 0u - : (8u * fast_select::kBins * 4u + row_offset[8] * 4u <= - phase_bytes - ? 1u - : (4u * fast_select::kBins * 4u + half * 4u <= phase_bytes - ? 2u - : 0u)); - stats[blockIdx.x * 7 + 3] = staged; - stats[blockIdx.x * 7] = recover; - stats[blockIdx.x * 7 + 1] = lo; - stats[blockIdx.x * 7 + 2] = hi; - for (int c = 4; c < 7; ++c) stats[blockIdx.x * 7 + c] = 0; - } - __syncthreads(); - if (recover) { - fallback.retained += row_delta * 4096; - fallback.indices += row_delta * 4096; - seed_fallback::sm100_mqa_logits_core_impl< - 32, 128, false, false, false, 8, 256, 1, 2, 4096, 128, 256, - cutlass::float_e4m3_t, float, float, decltype(scheduler), 2>( - rows, fallback, tq, ts, tk, ts, tw, scheduler); - } else { - if (threadIdx.x < 256) cutlass::arch::warpgroup_reg_dealloc<128>(); - __syncthreads(); - if (threadIdx.x >= 256) cutlass::arch::warpgroup_reg_alloc<128>(); - // Same 1024-B alignment as the cores' smem_buffer so every extern array - // aliases one base (mixed alignments put the scan layout 896 B past the - // base the launch sizes). - extern __shared__ __align__(1024) unsigned char phase_smem[]; - // All visible keys of a short row, padded with -1 (warp- or - // block-strided). - auto write_short = [&](unsigned r, unsigned t, unsigned stride) { - if (r >= rows) return; - const unsigned s0 = min(starts[r], keys), - visible = min(ends[r], keys) - s0; - int* o = fallback.output + (uint64_t)r * 2048; - for (unsigned j = t; j < 2048; j += stride) - o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; - }; - if constexpr (kSmemCand) { - // One warp per row over the shared candidates; scratch reuses the idle - // Q/KV stages. - static_assert(8 * fast_select::kScratch4Bytes <= sizeof(CShared), - "select scratch must fit the Q/KV stages"); - const unsigned warp = threadIdx.x / 32; - if (warp < 8 && row_base + warp < rows) { - const unsigned r = row_base + warp; - int* o = fallback.output + (uint64_t)r * 2048; - if (short_rows >> warp & 1u) - write_short(r, threadIdx.x & 31, 32); - else { - const uint16_t* hi = - reinterpret_cast(cand) + warp * kCandCap; - const uint8_t* lo = cand + kCandExt * 2 + warp * kCandCap; - const uint16_t* idx = - reinterpret_cast(cand + kCandExt * 3) + - warp * kCandCap; - const auto key_of = [idx](unsigned e) { return int32_t(idx[e]); }; - unsigned char* scratch = - phase_smem + warp * fast_select::kScratch4Bytes; - if (counts[r] >> 21 & 1) { - // The row pool overflowed: elements are the 3072 sealed shared - // slots, then the global spills at slots 0.. (count at slot - // 16383). - const uint64_t gbase = uint64_t(int64_t(r) + row_delta) * 16384; - const uint16_t* gv = vals + gbase; - const int* gi = indices + gbase; - const unsigned total = kCandCap + unsigned(__ldcg(gi + 16383)); - __syncwarp(); - if ((threadIdx.x & 31) == 0) - atomicAdd(stats + blockIdx.x * 7 + 5, 1); - fast_select::select_hist_generic( - [&](unsigned e) { - return e < kCandCap - ? (unsigned(hi[e]) << 8 | lo[e]) - : (unsigned(__ldcg(gv + e - kCandCap)) << 8 | - (unsigned(__ldcg(gi + e - kCandCap)) >> 16)); - }, - total, o, reinterpret_cast(scratch), - [&](unsigned e) { - return e < kCandCap - ? int32_t(idx[e]) - : int32_t(__ldcg(gi + e - kCandCap) & 0xffff); - }); - } else if (!fast_select::select_resident_v4( - hi, lo, 0u, kCandCap, - unsigned(counts[r] & ((1 << 21) - 1)), o, scratch, - key_of)) { - __syncwarp(); - if ((threadIdx.x & 31) == 0) - atomicAdd(stats + blockIdx.x * 7 + 6, 1); - fast_select::select_resident(hi, lo, kCandCap, 0u, o, - reinterpret_cast(scratch), - key_of); - } - } - } - } else if (staged) { - const unsigned warp = threadIdx.x / 32, - per_round = staged == 1 ? 8u : 4u; - __shared__ long long stamp[3]; - for (unsigned round = 0; round < 8 / per_round; ++round) { - if (warp < per_round) { - const unsigned local = round * per_round + warp, - first = round * per_round; - unsigned* hist = reinterpret_cast(phase_smem) + - warp * fast_select::kBins; - unsigned* codes = reinterpret_cast(phase_smem) + - per_round * fast_select::kBins + - (row_offset[local] - row_offset[first]); - const unsigned r = row_base + local; - const uint64_t base = uint64_t(int64_t(r) + row_delta) * 16384; - if (short_rows >> local & 1u) - write_short(r, threadIdx.x & 31, 32); - else - fast_select::select_row(codes, counts[r], vals + base, - indices + base, - fallback.output + (uint64_t)r * 2048, - hist, local == 0 ? stamp : nullptr); - } - __syncthreads(); - } - if (threadIdx.x == 0) { - stats[blockIdx.x * 7 + 4] = int(stamp[0]); - stats[blockIdx.x * 7 + 5] = int(stamp[1]); - stats[blockIdx.x * 7 + 6] = int(stamp[2]); - } - } else { - float* decoded = reinterpret_cast(phase_smem + 32768); - for (unsigned r = row_base; r < row_base + 8; ++r) { - if (short_rows >> (r - row_base) & 1u) { - write_short(r, threadIdx.x, 384); - __syncthreads(); - continue; - } - const uint64_t base = uint64_t(int64_t(r) + row_delta) * 16384; - const unsigned count = counts[r]; - int* output = fallback.output + (uint64_t)r * 2048; - for (unsigned j = threadIdx.x; j < count; j += 384) - decoded[j] = fused_candidates::candidate_decode_score( - vals[base + j], indices[base + j]); - __syncthreads(); - isolated_upstream::topKPerRowJob<384, 2304, false>( - nullptr, decoded, 0, count, output, nullptr, 1, 2048); - __syncthreads(); - for (unsigned j = threadIdx.x; j < 2048; j += 384) - output[j] = fused_candidates::candidate_decode_index( - indices[base + output[j]]); - __syncthreads(); - } - } - } - __syncthreads(); - if (threadIdx.x == 0) { - __threadfence(); - atomicExch(slot_flags + slot, 0); - } - finish_rows(recover); - } -} diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc new file mode 100644 index 000000000000..f76c7da8205b --- /dev/null +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -0,0 +1,1127 @@ +// 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 (unsigned i = 0; i < B::total; ++i) mbar_inval(bar_words + i); + for (unsigned i = 0; i < B::total; ++i) + reinterpret_cast(bar_words)[i] = 0; + full_q.init(1); + empty_q.init(256 + 32); + for (unsigned i = 0; i < KVS; ++i) { + full_kv[i].init(kFullKvArrivals); + empty_kv[i].init(kRel ? 1 : 256); + } + if constexpr (kRel) + for (unsigned i = 0; i < 2 * KVS; ++i) sf_full[i].init(1); + for (unsigned i = 0; i < B::NT; ++i) { + full_tmem[i].init(1); + empty_tmem[i].init(kLin ? 128 : 256); + } + if constexpr (kLin) full_lin->init(32); + for (unsigned i = 0; i < 2; ++i) { + claim_full[i].init(32); + claim_empty[i].init(kClaimConsumers * 32); + } // every reader lane arrives + sel_done.init(8); + for (unsigned i = 0; i < 16; ++i) { + h_full[i].init(32); + h_empty[i].init(32); + } + cutlass::arch::fence_barrier_init(); + unsafe_s = 0; +} +__syncwarp(); +if (warp_idx == 10) cute::TMEM::Allocator1Sm().allocate(512, &tmem_ptr); +__syncthreads(); +if (epoch == 0) cudaGridDependencySynchronize(); + +if (warp_idx == 8) { + // ---------------- claimer + Q/weights producer (+ drain of WG1 rows) + // ---------------- + cutlass::arch::warpgroup_reg_dealloc(); + unsigned n_q = 0, n_blk = 0, drain_epoch = 0; + for (unsigned k = 0;; ++k) { + if (lane == 0) { + empty_q.wait((n_q & 1) ^ 1); + } + __syncwarp(); + claim_empty[k & 1].wait(((k >> 1) & 1) ^ + 1); // every lane: lanes 0-7 write the entry + unsigned b = 0; + if (lane == 0) + b = atomicOr(&pause_req, 0u) ? kPause : atomicAdd(counter, 1u); + b = __shfl_sync(0xffffffffu, b, 0); + ClaimEntry& e = entries[k & 1]; + if (b >= nblocks) { + if (lane == 0) { + const unsigned kind = b == kPause ? kPause : kTerm; + e.vblock = kind; + epoch_kind = kind; + } + __syncwarp(); + claim_full[k & 1].arrive(); // every writer lane arrives + break; + } + const unsigned vblock = nblocks - 1 - b, row_base = vblock * 8; + unsigned st = 0, en = 0; + if (lane < 8) { + const unsigned 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(0xffffffffu, lane < 8 && en - st > 2048u) == 0; + // KV range: rows past the end are clamped to the last valid row (DeepGEMM + // scheduler). + const unsigned src = min(row_base + (lane & 7), rows - 1) - row_base; + unsigned lo = __shfl_sync(0xffffffffu, st, src), + hi = __shfl_sync(0xffffffffu, en, src); +#pragma unroll + for (int d = 4; d >= 1; d >>= 1) { + lo = min(lo, __shfl_xor_sync(0xffffffffu, lo, d)); + hi = max(hi, __shfl_xor_sync(0xffffffffu, hi, d)); + } + const unsigned kv_base = lo / 4 * 4; + if (!all_short && lane < 8) counts[row_base + lane] = 0; + if (lane == 0) { + // 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 unsigned slot = 2 * blockIdx.x + (n_blk & 1); + e.vblock = vblock; + e.all_short = all_short; + e.kv_base = kv_base; + e.splits = (hi - kv_base + 255) / 256; + e.slot = slot; + } + __syncwarp(); + claim_full[k & 1].arrive(); // every writer lane arrives + if (all_short) continue; + if (lane == 0) { + fused_mqa_dg::tma_copy<128, 256, 0>(&tq, &full_q, sm.q, 0, row_base * 32); + fused_mqa_dg::tma_copy<32, 8, 0>(&tw, &full_q, sm.weights, 0, row_base); + full_q.arrive_and_expect_tx(kQBytes + kWBytes); + } + __syncwarp(); + ++n_q; + ++n_blk; + if constexpr (!kSmem) { + const unsigned splits = (hi - kv_base + 255) / 256; + const unsigned slot = __shfl_sync(0xffffffffu, e.slot, 0); + const int64_t row_delta = int64_t(slot) * 8 - int64_t(row_base); + for (unsigned base = 0; base < splits; base += Interval, ++drain_epoch) { + for (unsigned p = 4; p < 8; ++p) + ho_drain(hdata, h_full, h_empty, p, drain_epoch, row_base + 4, counts, + kCandGlobalCap, 1 && base + Interval >= splits, + [&](unsigned row, unsigned slot2, unsigned record, + unsigned batch_base, unsigned source) { + const uint32_t physical = + kv_base + (batch_base * 2 + (record & 255u)) * 128 + + (source % 128); + const uint64_t off = + static_cast(int64_t(row) + row_delta) * + kCandGlobalCap + + slot2; + fused_candidates::store_candidate_record( + vals + off, indices + off, record >> 8, physical); + }); + } + cutlass::arch::NamedBarrier(320, 4).arrive(); + } + } +} else if (warp_idx == 9) { + // ---------------- KV producer: sample tiles then scan tiles, one ring + // ---------------- + cutlass::arch::warpgroup_reg_dealloc(); + RingPipeline kv_pipe; + [[maybe_unused]] RingPipeline<2 * KVS> sf_pipe; + unsigned n_blk = 0; + for (unsigned k = 0;; ++k) { + claim_full[k & 1].wait((k >> 1) & 1); + const ClaimEntry& e = entries[k & 1]; + const unsigned vblock = e.vblock, all_short = e.all_short, + kv_base = e.kv_base, splits = e.splits; + __syncwarp(); + if (vblock >= kPause) break; + claim_empty[k & 1].arrive(); + if (all_short) continue; + if constexpr (KVS == 3) { + // > 16K: 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 unsigned center = sample_stride / 2; + float pf[8]; +#pragma unroll + for (unsigned j = 0; j < 8; ++j) { + const unsigned g = j * 32 + lane; + pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center) + : 0.0f; + } +#pragma unroll 1 + for (unsigned t = 0; t < sample_tiles; ++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 unsigned s = __shfl_sync(0xffffffffu, s0, 0); + // Every writing lane acquires and releases the stage itself. + empty_kv[s].wait(__shfl_sync(0xffffffffu, ph, 0) ^ 1); +#pragma unroll + for (unsigned j = 0; j < 8; ++j) sm.sf_kv[s][j * 32 + lane] = pf[j]; + cutlass::arch::fence_view_async_shared(); // generic writes before + // later TMA writes of this + // stage + if (lane == 0) { + full_kv[s].expect_transaction(kKVBytes); + fused_mqa_dg::tma_copy<128, 256, 0, cutlass::float_e4m3_t, true>( + &tks, &full_kv[s], sm.kv[s], 0, t * 256, 1, center); + } + full_kv[s].arrive(); +#pragma unroll + for (unsigned j = 0; j < 8; ++j) { + const unsigned g = (t + 1) * 256 + j * 32 + lane; + pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center) + : 0.0f; + } + } + } + if (lane == 0) { + if (KVS != 3 && n_blk > 0) + sel_done.wait((n_blk - 1) & + 1); // selection scratch overlaps the KV stages + const unsigned center = sample_stride / 2; +#pragma unroll 1 + for (unsigned t = 0; t < (KVS == 3 ? 0u : sample_tiles); ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), s, ph); + empty_kv[s].wait(ph ^ 1); + fused_mqa_dg::tma_copy<128, 256, 0, cutlass::float_e4m3_t, true>( + &tks, &full_kv[s], sm.kv[s], 0, t * 256, 1, center); + mbar_arrive_expect_tx(&full_kv[s], kKVBytes, kFullKvArrivals); + } +#pragma unroll 1 + for (unsigned 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 * 256; + fused_mqa_dg::tma_copy<128, 256, 0>(&tk, &full_kv[s], sm.kv[s], 0, off); + if constexpr (kRel) { + mbar_arrive_expect_tx(&full_kv[s], kKVBytes, kFullKvArrivals); + CUTE_TIE_DECL(sf_pipe.advance(), sfs, sfph); + (void)sfph; + fused_mqa_dg::tma_copy<256, 1, 0>(&ts, &sf_full[sfs], sm.sf_kv[sfs], + off, 0); + sf_full[sfs].arrive_and_expect_tx(1024); + } else { + fused_mqa_dg::tma_copy<256, 1, 0>(&ts, &full_kv[s], sm.sf_kv[s], off, + 0); + mbar_arrive_expect_tx(&full_kv[s], kKVBytes + 1024, 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 (unsigned k = 0;; ++k) { + claim_full[k & 1].wait((k >> 1) & 1); + const ClaimEntry& e = entries[k & 1]; + const unsigned vblock = e.vblock, all_short = e.all_short, + splits = e.splits; + __syncwarp(); + if (vblock >= kPause) break; + claim_empty[k & 1].arrive(); + if (all_short) continue; + CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase); + (void)q_stage; + full_q.wait(q_phase); + if constexpr (kLin) { + full_lin->wait(q_phase); // this block's linear-term rows (warp 11) + ptx::tcgen05_after_thread_sync(); + } + const unsigned total = sample_tiles + splits; +#pragma unroll 1 + for (unsigned t = 0; t < total; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), s, ph); + full_kv[s].wait(ph); + if (cute::elect_one_sync()) { + if constexpr (kLin) { + // 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 + 128 * 128; +#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]); + } + } + } else { +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + CUTE_TIE_DECL(tmem_pipe.advance(), ts_idx, ts_ph); + const uint32_t tmem_addr = ts_idx * 256; + empty_tmem[ts_idx].wait(ts_ph ^ 1); + ptx::tcgen05_after_thread_sync(); +#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( + sm.q, 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]); + } + if constexpr (kRel) ptx::umma_arrive_no_elect(empty_kv[s]); + } + } + __syncwarp(); + } + empty_q.arrive(); + } +} else if (warp_idx == 11) { + // ---------------- drain of WG0 rows (> 16K) ---------------- + cutlass::arch::warpgroup_reg_dealloc(); + if constexpr (!kSmem) { + unsigned drain_epoch = 0, n_lin = 0; + for (unsigned k = 0;; ++k) { + claim_full[k & 1].wait((k >> 1) & 1); + const ClaimEntry& e = entries[k & 1]; + const unsigned vblock = e.vblock, all_short = e.all_short, + kv_base = e.kv_base, splits = e.splits, slot = e.slot; + __syncwarp(); + if (vblock >= kPause) break; + claim_empty[k & 1].arrive(); + if (all_short) continue; + if constexpr (kLin) { + // 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); + cutlass::arch::fence_view_async_shared(); // generic smem writes -> + // UMMA (async proxy) + full_lin->arrive(); // every lane arrives (count 32) + } + const unsigned row_base = vblock * 8; + const int64_t row_delta = int64_t(slot) * 8 - int64_t(row_base); + for (unsigned base = 0; base < splits; base += Interval, ++drain_epoch) { + for (unsigned p = 0; p < 4; ++p) + ho_drain(hdata, h_full, h_empty, p, drain_epoch, row_base, counts, + kCandGlobalCap, 1 && base + Interval >= splits, + [&](unsigned row, unsigned slot2, unsigned record, + unsigned batch_base, unsigned source) { + const uint32_t physical = + kv_base + (batch_base * 2 + (record & 255u)) * 128 + + (source % 128); + const uint64_t off = + static_cast(int64_t(row) + row_delta) * + kCandGlobalCap + + slot2; + fused_candidates::store_candidate_record( + vals + off, indices + off, record >> 8, physical); + }); + } + cutlass::arch::NamedBarrier(320, 4).arrive(); + } + } +} else { + // ---------------- math warps 0-7 ---------------- + cutlass::arch::warpgroup_reg_alloc(); + const unsigned math_tid = threadIdx.x; + const unsigned query_base = (warp_idx / 4) * 4; + RingPipeline<1> q_pipe; + RingPipeline kv_pipe; + [[maybe_unused]] RingPipeline<2 * KVS> sf_pipe; + RingPipeline tmem_pipe; + // > 16K: each warpgroup consumes every second TMEM slot position (WG1 starts + // at 1). + constexpr unsigned kTmemStep = kLin ? 2 : 1; + if (kLin && warp_idx >= 4) tmem_pipe.advance(1); + [[maybe_unused]] const uint32_t heads_col = warp_idx < 4 ? 16u : 0u; + [[maybe_unused]] const uint32_t lin_col = warp_idx < 4 ? 0u : 128u; + unsigned produce_epoch = 0; + constexpr uint32_t kRingSlots = 8; + for (unsigned k = 0;; ++k) { + claim_full[k & 1].wait((k >> 1) & 1); + const ClaimEntry& e = entries[k & 1]; + const unsigned vblock = e.vblock; + if (vblock >= kPause) break; + const unsigned row_base = vblock * 8; + if (e.all_short) { + for (unsigned r = 0; r < 8 && row_base + r < rows; ++r) { + const unsigned s0 = e.starts[r], vis = e.ends[r] - s0; + for (unsigned j = math_tid; j < 2048; j += 256) + fallback.output[(uint64_t)(row_base + r) * 2048 + j] = + j < vis ? int((relative ? 0u : s0) + j) : -1; + } + __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); + if constexpr (kLin) full_lin->wait(q_phase); + } + static_assert(SampleQ * 2112 % 4 == 0); + for (unsigned j = math_tid; j < SampleQ * 2112 / 4; j += 256) + reinterpret_cast(stream.hist)[j] = make_uint4(0, 0, 0, 0); + if (math_tid < 8) stream.counts[math_tid] = 0; + float weights[4][32]; + [[maybe_unused]] float + row_s[4]; // > 16K: pow2 row scale of the linear term + uint32_t seq_k_start[4], seq_k_end[4]; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + seq_k_start[i] = e.starts[query_base + i]; + seq_k_end[i] = e.ends[query_base + i]; + if constexpr (kLin) { + // 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 < 32; j += 4) { + const float4 w4 = *reinterpret_cast( + sm.weights + (query_base + i) * 32 + 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; + } + } else { +#pragma unroll + for (uint32_t j = 0; j < 32; j += 4) { + const float4 w4 = *reinterpret_cast( + sm.weights + (query_base + i) * 32 + j); + weights[i][j] = w4.x * 0.5f; + weights[i][j + 1] = w4.y * 0.5f; + weights[i][j + 2] = w4.z * 0.5f; + weights[i][j + 3] = w4.w * 0.5f; + } + } + } + { + const float w = sm.weights[warp_idx * 32 + lane]; + bool bad_w = w != 0.0f && fabsf(w) < 0x1p-125f; + if constexpr (kLin) { + // Scaled weights that become subnormal or infinite take the exact + // recovery path. + const float wr = w * 0.5f * (1.0f / sm.lin_scale[warp_idx]); + bad_w |= (wr != 0.0f && fabsf(wr) < 0x1p-125f) || isinf(wr); + } + const unsigned bad = __ballot_sync(0xffffffffu, bad_w); + if (lane == 0 && bad) atomicOr(&unsafe_s, 1u << warp_idx); + } + sample_sync(); + // ---- sample tiles ---- +#pragma unroll 1 + for (unsigned t = 0; t < sample_tiles; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); + if constexpr (!kRel) full_kv[kv_stage].wait(kv_phase); +#pragma unroll + for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { + const unsigned kv_lane = (math_tid % 128) + kv_half * 128; + const unsigned group = t * 256 + kv_lane; + const unsigned kv_offset = group * sample_stride + sample_stride / 2; + const bool in_view = group < sample_groups; + float scale_kv; + if constexpr (kLin) + scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + + kv_lane); // 0 outside the view + else + scale_kv = in_view ? __ldg(scales + kv_offset) : 0.0f; + CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase); + full_tmem[tmem_stage].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + if (!kRel && 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 < 4; ++i) { + float acc[32]; + [[maybe_unused]] float lin[4]; + if constexpr (kLin) { + ptx::tmem_load_32dp32b<32>( + tmem_stage * kSlotCols + heads_col + i * 32, + reinterpret_cast(acc)); + ptx::tmem_load_32dp32b<4>(tmem_stage * kSlotCols + lin_col + i * 4, + reinterpret_cast(lin)); + } else { + ptx::tmem_load_32dp32b<32>(tmem_stage * 256 + (query_base + i) * 32, + reinterpret_cast(acc)); + } + cutlass::arch::fence_view_async_tmem_load(); + if (i == 3) { + ptx::tcgen05_before_thread_sync(); + empty_tmem[tmem_stage].arrive(); + } + float2 sum_0 = {}, sum_1 = {}; + if constexpr (kLin) { // terms 2, 3 are in units of s * 2^-8 + sum_0 = make_float2(lin[0], lin[1]); + sum_1 = __fmul2_rn(make_float2(lin[2], lin[3]), + make_float2(0x1p-8f, 0x1p-8f)); + } +#pragma unroll + for (uint32_t j = 0; j < 32; j += 4) { + auto transform = [&](uint32_t at, float2 sum) { + const auto abs_x = + make_float2(fabsf(acc[at]), fabsf(acc[at + 1])); + const auto w = make_float2(weights[i][at], weights[i][at + 1]); + if constexpr (kLin) { + return __ffma2_rn( + abs_x, w, + sum); // relu = (x + |x|) / 2, linear part on the TC + } else { + const auto x = make_float2(acc[at], acc[at + 1]); + return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); + } + }; + sum_0 = transform(j, sum_0); + sum_1 = transform(j + 2, sum_1); + } + const auto sum = __fadd2_rn(sum_0, sum_1); + float bq; + if constexpr (kLin) + bq = (sum.x + sum.y) * (row_s[i] * scale_kv); + else + bq = (sum.x + sum.y) * scale_kv; + const bool valid = in_view && kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq); + if constexpr (kLin) { + if (valid) { + const unsigned bin = sample_bin(bq); + atomicAdd(stream.hist[query_base + i] + bin + bin / 32, 1u); + } + } else { + stream.scores[query_base + i][kv_lane] = valid ? bq : -CUDART_INF_F; + } + } + } + if constexpr (!kLin) { + sample_sync(); + accumulate_sample(stream, math_tid); + sample_sync(); + } + } + if constexpr (kLin) { + sample_sync(); + finish_sample_histogram(stream.hist[warp_idx], gate_s + warp_idx, + e.ends[warp_idx] - e.starts[warp_idx], + sample_target, lane); + } else { + const unsigned rb = + *reinterpret_cast(&e.vblock) * 8; + SampleBuffers sb{gate_s - rb, nullptr, nullptr, nullptr, sample_target}; + finish_sample(stream, sb, warp_idx, rb + warp_idx, + e.ends[warp_idx] - e.starts[warp_idx], lane); + } + sample_sync(); + const unsigned row_base2 = + *reinterpret_cast(&e.vblock) * 8; + + // ---- scan tiles ---- + if constexpr (kLin) { +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const float inv_s = 1.0f / row_s[i]; +#pragma unroll + for (uint32_t j = 0; j < 32; j += 4) { + const float4 w4 = *reinterpret_cast( + sm.weights + (query_base + i) * 32 + 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; + } + } + } + const unsigned kv_base = e.kv_base, splits = e.splits; + const int64_t cand_row_delta = int64_t(e.slot) * 8 - int64_t(row_base2); + float gate_score[4]; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + gate_score[i] = gate_s[query_base + i]; + if (*reinterpret_cast(&unsafe_s) >> + (query_base + i) & + 1u) { + gate_score[i] = __int_as_float(0x7f800000); + if (math_tid % 128 == 0) + atomicAdd(counts + row_base2 + query_base + i, + static_cast(kCandGlobalCap) + 1); + } + // > 16K: compare S * scale_kv against gate / row_s (row_s is a power of + // two -> exact). + if constexpr (kLin) gate_score[i] *= 1.0f / row_s[i]; + } + uint32_t lane_counts[4] = {}; + uint32_t flush_base = 0; + const uint32_t logits_stride = kCandCap; + [[maybe_unused]] const uint32_t pool_size = logits_stride / 6, + seg_size = (logits_stride - pool_size) / 4; + [[maybe_unused]] const uint32_t seg_warp = warp_idx % 4; + [[maybe_unused]] const uint32_t cand_ext = 8 * logits_stride + 256; + [[maybe_unused]] uint32_t seg_cnt[4] = {}; + if constexpr (kSmem) { + if (math_tid < 8) { + row_count[math_tid] = 0; + pool_cnt[math_tid] = 0; + } + cutlass::arch::NamedBarrier(256, 3).sync(); + } +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < splits; ++kv_split_idx) { + if constexpr (!kSmem) + if ((kv_split_idx & (Interval - 1)) == 0) { + h_empty[(produce_epoch & 1) * 8 + warp_idx].wait( + ((produce_epoch >> 1) & 1) ^ 1); + __syncwarp(); + } + CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); + unsigned sf_stage = kv_stage; + if constexpr (kRel) { + CUTE_TIE_DECL(sf_pipe.advance(), rs, rph); + sf_full[rs].wait(rph); + sf_stage = rs; + } else { + full_kv[kv_stage].wait(kv_phase); + } +#pragma unroll + for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { + const unsigned kv_lane = (math_tid % 128) + kv_half * 128; + const unsigned kv_offset = kv_base + kv_split_idx * 256 + kv_lane; + const float scale_kv = ptx::ld_shared(sm.sf_kv[sf_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 (!kRel && kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + empty_kv[kv_stage].arrive(); + } + float accum_pair[2][32]; + [[maybe_unused]] float lin_pair[2][4]; +#pragma unroll + for (uint32_t pair = 0; pair < 4; pair += 2) { + if constexpr (kLin) { + const uint32_t slot_base = tmem_stage * kSlotCols; + ptx::tmem_load_32dp32b<64>(slot_base + heads_col + pair * 32, + reinterpret_cast(accum_pair)); + ptx::tmem_load_32dp32b<8>(slot_base + lin_col + pair * 4, + reinterpret_cast(lin_pair)); + } else { + const uint32_t tmem_addr = + tmem_stage * 256 + (query_base + pair) * 32; + ptx::tmem_load_32dp32b<64>(tmem_addr, + reinterpret_cast(accum_pair)); + } + cutlass::arch::fence_view_async_tmem_load(); + if (pair + 2 >= 4) { + ptx::tcgen05_before_thread_sync(); + empty_tmem[tmem_stage].arrive(); + } + float2 sum_0[2] = {}, sum_1[2] = {}; + if constexpr (kLin) { + // The linear term (sum_t term_t . k, in units of the row scale) + // seeds the sums. +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + sum_0[r] = make_float2(lin_pair[r][0], lin_pair[r][1]); + sum_1[r] = + __fmul2_rn(make_float2(lin_pair[r][2], lin_pair[r][3]), + make_float2(0x1p-8f, 0x1p-8f)); // units s * 2^-8 + } + } +#pragma unroll + for (uint32_t j = 0; j < 32; j += 4) { +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + const uint32_t i = pair + r; + auto* a = accum_pair[r]; + auto transform = [&](uint32_t at, float2 sum) { + const auto abs_x = make_float2(fabsf(a[at]), fabsf(a[at + 1])); + const auto w = make_float2(weights[i][at], weights[i][at + 1]); + if constexpr (kLin) { + return __ffma2_rn(abs_x, w, sum); + } else { + const auto x = make_float2(a[at], a[at + 1]); + return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); + } + }; + sum_0[r] = transform(j, sum_0[r]); + sum_1[r] = transform(j + 2, sum_1[r]); + } + } + if constexpr (kLin) { +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + const uint32_t i = pair + r; + const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); + const float bp = (sum.x + sum.y) * scale_kv; + if (bp > gate_score[i]) { // gate first; range / finiteness only + // for passing keys + const float bq = bp * row_s[i]; + const uint32_t row = row_base2 + 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 = + fused_candidates::candidate_fp24_code(bq); + hdata.records[produce_epoch & 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(counts + row, 1); + if (slot < static_cast(kCandGlobalCap)) { + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + kCandGlobalCap + + slot; + fused_candidates::store_candidate( + vals + off, indices + off, bq, kv_offset); + } + } + } + } + } + } else +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) { + const uint32_t i = pair + r; + const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); + float bq; + if constexpr (kLin) + bq = (sum.x + sum.y) * (row_s[i] * scale_kv); + else + bq = (sum.x + sum.y) * scale_kv; + const uint32_t row = row_base2 + query_base + i; + const bool pass = + (kSmem ? kv_offset - seq_k_start[i] < + seq_k_end[i] - seq_k_start[i] + : kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq)) && + bq > gate_score[i]; + if constexpr (kSmem) { + const unsigned pass_mask = __ballot_sync(0xffffffffu, pass); + const unsigned slot = + seg_cnt[i] + __popc(pass_mask & ((1u << lane) - 1u)); + seg_cnt[i] += __popc(pass_mask); + const uint32_t code = fused_candidates::candidate_fp24_code(bq); + const uint32_t seg_cap = seg_size; + const uint32_t at = + pass && slot < seg_cap + ? (query_base + i) * logits_stride + + seg_warp * seg_size + slot + : 8 * logits_stride + warp_idx * 32 + lane; + auto* cbase = region; + reinterpret_cast(cbase)[at] = + static_cast(code >> 8); + cbase[cand_ext * 2 + at] = static_cast(code); + reinterpret_cast(cbase + cand_ext * 3)[at] = + static_cast(kv_offset); + if (seg_cnt[i] > seg_cap) { + const bool over = pass && slot >= seg_cap; + const unsigned over_mask = __ballot_sync(0xffffffffu, over); + if (over_mask) { + unsigned pool_base = 0; + if (lane == 0) + pool_base = atomicAdd(&pool_cnt[query_base + i], + __popc(over_mask)); + pool_base = __shfl_sync(0xffffffffu, pool_base, 0); + const unsigned p = + pool_base + __popc(over_mask & ((1u << lane) - 1u)); + if (over && p < pool_size) { + const uint32_t pat = + (query_base + i) * logits_stride + 4 * seg_size + p; + reinterpret_cast(cbase)[pat] = + static_cast(code >> 8); + cbase[cand_ext * 2 + pat] = static_cast(code); + reinterpret_cast(cbase + cand_ext * 3)[pat] = + static_cast(kv_offset); + } else if (over && p - pool_size < kCandGlobalCap - 1) { + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + kCandGlobalCap + + (p - pool_size); + __stcg(vals + off, static_cast(code >> 8)); + __stcg(indices + off, + static_cast((code & 0xffu) << 16 | + kv_offset)); + } + } + } + } else if (pass) { + const uint32_t count = lane_counts[i]; + if (count < kRingSlots) { + const uint32_t code = + fused_candidates::candidate_fp24_code(bq); + hdata.records[produce_epoch & 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(counts + row, 1); + if (slot < static_cast(kCandGlobalCap)) { + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + kCandGlobalCap + + slot; + fused_candidates::store_candidate(vals + off, indices + off, + bq, kv_offset); + } + } + } + } + } + } + if (!kSmem && + (((kv_split_idx + 1) & (Interval - 1)) == 0 || + kv_split_idx + 1 == splits) && + kv_split_idx + 1 != splits) { +#pragma unroll + for (unsigned i = 0; i < 4; ++i) { + hdata.counts[produce_epoch & 1][warp_idx][i][lane] = lane_counts[i]; + lane_counts[i] = 0; + } + if (lane == 0) hdata.bases[produce_epoch & 1][warp_idx] = flush_base; + __syncwarp(); + h_full[(produce_epoch & 1) * 8 + warp_idx].arrive(); + ++produce_epoch; + flush_base = kv_split_idx + 1; + } + } + if constexpr (!kSmem) { + // 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 unsigned bank = produce_epoch & 1, rbase = row_base2 + query_base; + static_assert(sizeof(hdata.records[0][0]) == 4 * 8 * 32 * 4, + "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], counts, + rbase, kv_base + (math_tid % 128), flush_base, + cand_row_delta, vals, indices); +#pragma unroll + for (unsigned row = 0; row < 4; ++row) lane_counts[row] = 0; + __syncwarp(); + h_full[bank * 8 + warp_idx].arrive(); + ++produce_epoch; + } + if constexpr (kSmem) { +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + auto* base = region; + for (uint32_t slot = min(seg_cnt[i], seg_size) + lane; slot < seg_size; + slot += 32) { + const uint32_t at = + (query_base + i) * logits_stride + seg_warp * seg_size + slot; + reinterpret_cast(base)[at] = 0; + base[cand_ext * 2 + at] = 0; + reinterpret_cast(base + cand_ext * 3)[at] = 0; + } + if (lane == 0) atomicAdd(&row_count[query_base + i], seg_cnt[i]); + } + cutlass::arch::NamedBarrier(256, 3).sync(); + for (uint32_t q = 0; q < 8; ++q) { + auto* base = region; + for (uint32_t p = pool_cnt[q] + math_tid; p < pool_size; p += 256) { + const uint32_t at = q * logits_stride + 4 * seg_size + p; + reinterpret_cast(base)[at] = 0; + base[cand_ext * 2 + at] = 0; + reinterpret_cast(base + cand_ext * 3)[at] = 0; + } + } + if (math_tid < 8) { + const uint32_t row = row_base2 + math_tid; + const uint32_t spilled = + pool_cnt[math_tid] > pool_size + ? min(pool_cnt[math_tid] - pool_size, kCandGlobalCap - 1) + : 0u; + __stcg(indices + + static_cast(int64_t(row) + cand_row_delta) * + kCandGlobalCap + + kCandGlobalCap - 1, + static_cast(spilled)); + if (spilled) atomicOr(&row_count[math_tid], 1u << 21); + } + cutlass::arch::NamedBarrier(256, 3).sync(); + if (math_tid < 8) + atomicAdd(counts + row_base2 + math_tid, + static_cast(row_count[math_tid])); + } + cutlass::arch::fence_view_async_shared(); + empty_q.arrive(); + if constexpr (!kSmem) + cutlass::arch::NamedBarrier(320, 4).sync(); // drains of this block done + // ---- selection ---- + const unsigned vblock3 = + *reinterpret_cast(&e.vblock), + row_base3 = vblock3 * 8; + __threadfence_block(); + cutlass::arch::NamedBarrier(256, 5).sync(); + const ClaimEntry& ce = entries[k & 1]; + if (math_tid == 0) { + unsigned recover = 0, short_rows = 0; + row_offset[0] = 0; + for (unsigned r = 0; r < 8; ++r) { + const bool is_short = ce.ends[r] - ce.starts[r] <= 2048u; + short_rows |= unsigned(is_short) << r; + const int c = + kSmem ? int(row_count[r]) : __ldcg(counts + row_base3 + r); + const int real = kSmem ? (c & ((1 << 21) - 1)) : c; + recover |= !is_short && (real < 2048 || real > 16384); + // > 16K rows start 16-byte aligned for select_staged_row's vector + // loads. + row_offset[r + 1] = + row_offset[r] + (is_short + ? 0 + : (!kSmem ? (unsigned(max(c, 0)) + 3u) & ~3u + : unsigned(max(c, 0)))); + } + unsigned staged = 1; + if (!kSmem) { + const unsigned half = max(row_offset[4], row_offset[8] - row_offset[4]); + staged = recover ? 0u + : (8u * fast_select::kBins * 4u + row_offset[8] * 4u <= + phase_bytes + ? 1u + : (4u * fast_select::kBins * 4u + half * 4u <= + phase_bytes + ? 2u + : 0u)); + if (staged == 0) + recover = 1; // candidates exceed the scratch: exact rescan + } + if (recover) { + defer_next[vblock3] = int(defer_head); + defer_head = vblock3; + atomicExch(&pause_req, 1u); + } + sel_recover = recover; + sel_staged = staged; + sel_short = short_rows; + unsafe_s = 0; + } + cutlass::arch::NamedBarrier(256, 5).sync(); + const unsigned recover = sel_recover, short_rows = sel_short; + const auto write_short = [&](unsigned r, unsigned t, unsigned stride) { + if (r >= rows) return; + const unsigned s0 = ce.starts[r - row_base3], + visible = ce.ends[r - row_base3] - s0; + int* o = fallback.output + (uint64_t)r * 2048; + for (unsigned j = t; j < 2048; j += stride) + o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; + }; + if (!recover) { + const int64_t row_delta = int64_t(ce.slot) * 8 - int64_t(row_base3); + if constexpr (kSmem) { + const unsigned warp = warp_idx; + if (row_base3 + warp < rows) { + const unsigned r = row_base3 + warp; + int* o = fallback.output + (uint64_t)r * 2048; + if (short_rows >> warp & 1u) + write_short(r, lane, 32); + else { + const uint16_t* hi = + reinterpret_cast(region) + warp * kCandCap; + const uint8_t* lo = region + kCandExt * 2 + warp * kCandCap; + const uint16_t* idx = + reinterpret_cast(region + kCandExt * 3) + + warp * kCandCap; + const auto key_of = [idx](unsigned e2) { return int32_t(idx[e2]); }; + unsigned char* scratch = smem_buffer + kSelB + warp * kScratchS; + const int c = int(row_count[warp]); + if (c >> 21 & 1) { + const uint64_t gbase = uint64_t(int64_t(r) + row_delta) * 16384; + const uint16_t* gv = vals + gbase; + const int* gi = indices + gbase; + const unsigned total = kCandCap + unsigned(__ldcg(gi + 16383)); + fast_select::select_hist_generic( + [&](unsigned e2) { + return e2 < kCandCap + ? (unsigned(hi[e2]) << 8 | lo[e2]) + : (unsigned(__ldcg(gv + e2 - kCandCap)) << 8 | + (unsigned(__ldcg(gi + e2 - kCandCap)) >> 16)); + }, + total, o, reinterpret_cast(scratch), + [&](unsigned e2) { + return e2 < kCandCap + ? int32_t(idx[e2]) + : int32_t(__ldcg(gi + e2 - kCandCap) & 0xffff); + }); + } else if (!row_select::select_shared_row( + hi, lo, idx, kCandCap, unsigned(c & ((1 << 21) - 1)), + o, scratch)) { + fast_select::select_resident(hi, lo, kCandCap, 0u, o, + reinterpret_cast(scratch), + key_of); + } + } + } + } else { + const unsigned per_round = sel_staged == 1 ? 8u : 4u; + for (unsigned round = 0; round < 8 / per_round; ++round) { + if (warp_idx < per_round) { + const unsigned local = round * per_round + warp_idx, + first = round * per_round; + unsigned* hist = reinterpret_cast(smem_buffer + kSelB) + + warp_idx * fast_select::kBins; + unsigned* codes = reinterpret_cast(smem_buffer + kSelB) + + per_round * fast_select::kBins + + (row_offset[local] - row_offset[first]); + const unsigned r = row_base3 + local; + const uint64_t base = uint64_t(int64_t(r) + row_delta) * 16384; + if (short_rows >> local & 1u) + write_short(r, lane, 32); + else { + const unsigned cnt = __ldcg(counts + r); + int* o = fallback.output + (uint64_t)r * 2048; + if (!row_select::select_staged_row(codes, cnt, vals + base, + indices + base, o, hist)) + fast_select::select_row(codes, cnt, vals + base, indices + base, + o, hist); + } + } + cutlass::arch::NamedBarrier(256, 5).sync(); + } + } + } + cutlass::arch::NamedBarrier(256, 5).sync(); + if (relative && !recover) { + for (unsigned i = math_tid; i < 8 * 512; i += 256) { + const unsigned r = row_base3 + i / 512; + if (r >= rows || ce.ends[i / 512] - ce.starts[i / 512] <= 2048u) + continue; + const int s = int(starts_g[r]); + if (s == 0) continue; + int4* p = + reinterpret_cast(fallback.output + (uint64_t)r * 2048) + + i % 512; + 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; + } + } + __syncwarp(); + if (lane == 0) sel_done.arrive(); + claim_empty[k & 1].arrive(); + } +} + +// ---------------- exact rescan of the blocks queued by the selection +// ---------------- +__syncthreads(); +if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, 512); +if (defer_head != kTerm) { + if constexpr (kMathRegs != 224) { + // The fallback core expects the 224 / 56 split. + if (warp_idx < 8) cutlass::arch::warpgroup_reg_dealloc<224>(); + __syncthreads(); + if (warp_idx >= 8) cutlass::arch::warpgroup_reg_alloc<56>(); + } + __syncthreads(); + for (unsigned cur = defer_head; cur != kTerm;) { + const unsigned vblock = cur, row_base = vblock * 8; + if (threadIdx.x < 8) { + const unsigned 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; + } + if (threadIdx.x == 0) + fslot = 2 * blockIdx.x; // no block is in flight during the rescan + __syncthreads(); + const unsigned* const starts = row_bounds[0] - row_base; + const unsigned* const ends = row_bounds[1] - row_base; + const int64_t row_delta = int64_t(fslot) * 8 - int64_t(row_base); + FallbackBuffers fb = fallback; + fb.retained += row_delta * 4096; + fb.indices += row_delta * 4096; + fused_mqa_fallback::rescan_block( + fused_mqa_dense::Scheduler(vblock, rows, keys, starts, ends), fb, tq, + tk, ts, tw); + __syncthreads(); + if (threadIdx.x == 0) { + auto& fst = + *reinterpret_cast(smem_buffer); + for (auto& bb : fst.full_q_barriers) mbar_inval(&bb); + for (auto& bb : fst.empty_q_barriers) mbar_inval(&bb); + for (auto& bb : fst.full_kv_barriers) mbar_inval(&bb); + for (auto& bb : fst.empty_kv_barriers) mbar_inval(&bb); + for (auto& bb : fst.full_tmem_barriers) mbar_inval(&bb); + for (auto& bb : fst.empty_tmem_barriers) mbar_inval(&bb); + } + __syncthreads(); + if (relative) { + for (unsigned i = threadIdx.x; i < 8 * 512; i += 384) { + const unsigned r = row_base + i / 512; + if (r >= rows) continue; + const int s = int(starts_g[r]); + if (s == 0) continue; + int4* p = + reinterpret_cast(fallback.output + (uint64_t)r * 2048) + + i % 512; + 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; + } + } + cur = unsigned(__ldcg(defer_next + vblock)); + __syncthreads(); + } + __syncthreads(); + 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..3cf85fa2b0a7 --- /dev/null +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -0,0 +1,550 @@ +// SPDX-License-Identifier: MIT +// Persistent fused FP8 MQA logits + per-row top-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, > 16K keys) +// warp 9 KV TMA (sample tiles via the strided map, then the scan tiles) +// warp 10 UMMA issuer +// warp 11 drain of math WG0 rows (> 16K keys), idle otherwise +// warps 0-7 math: sample histogram -> gate -> scan/emit -> selection +// A block whose candidate count leaves [2048, 16384] pauses the role loops; the +// CTA rescans it exactly (fallback.cuh) and resumes the loops with freshly +// initialised barriers. +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "candidates.cuh" +#include "deep_gemm_compat.cuh" +#include "fallback.cuh" +#include "fast_select.cuh" +#include "row_select.cuh" +#include "sample.cuh" + +namespace fused_persist { +using namespace deep_gemm; +using TBar = cutlass::arch::ClusterTransactionBarrier; + +constexpr unsigned kTerm = 0xffffffffu, kPause = 0xfffffffeu; +// Selection scratch starts after Q + weights so the next block's Q can land +// during selection. +constexpr unsigned kSelBase = 33792; +constexpr unsigned kScratchS = + 8192; // per-warp selection scratch (<= 16K keys) +constexpr unsigned kCandCap = 3072; +constexpr unsigned kCandExt = 8u * kCandCap + 256u; +constexpr unsigned kQBytes = 8 * 32 * 128, kKVBytes = 256 * 128, + kWBytes = 8 * 32 * 4; + +template +struct PStorage { + alignas(1024) cutlass::float_e4m3_t q[kQBytes]; + alignas(128) float weights[8 * 32]; + alignas(1024) cutlass::float_e4m3_t kv[KVS][kKVBytes]; + // <= 16K: the per-key scales have their own ring, twice as deep as the KV + // stages (see BarIdx::sf_full). + alignas(128) float sf_kv[2 * KVS][256]; +}; +// > 16K keys (rebalanced scan): the 16 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]). +template <> +struct PStorage<3> { + alignas(1024) uint8_t lin_lo[16 * 128]; + cutlass::float_e4m3_t q[kQBytes]; + uint8_t lin_hi[16 * 128]; + alignas(128) float weights[8 * 32]; + float lin_scale[8]; + alignas(1024) cutlass::float_e4m3_t kv[3][kKVBytes]; + alignas(128) float sf_kv[3][256]; +}; +static_assert(offsetof(PStorage<3>, q) == 16 * 128 && + offsetof(PStorage<3>, lin_hi) == 16 * 128 + kQBytes, + "lin rows must bracket Q"); +static_assert(offsetof(PStorage<2>, kv) == kSelBase, + "selection scratch base (<= 16K)"); +// > 16K: selection scratch starts after lin + Q + lin + weights + row scales +// (the next block's linear term is written during this block's selection). +constexpr unsigned kSelBase3 = offsetof(PStorage<3>, kv); + +// 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(PStorage<3>& sm, + unsigned lane) { + const uint8_t* qbytes = reinterpret_cast(sm.q); +#pragma unroll 1 + for (unsigned r = 0; r < 8; ++r) { + float acc[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll 8 + for (unsigned h = 0; h < 32; ++h) { + const unsigned R = r * 32 + h; + const unsigned off = R * 128u + ((((lane * 4u) >> 4) ^ (R & 7u)) << 4) + + ((lane * 4u) & 15u); + const uint32_t packed = + ptx::ld_shared(reinterpret_cast(qbytes + off)); + const float wv = ptx::ld_shared(sm.weights + r * 32 + 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 (unsigned o = 16; o > 0; o >>= 1) + m = fmaxf(m, __shfl_xor_sync(0xffffffffu, 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 (unsigned dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s; + uint8_t* lin_base = r < 4 ? sm.lin_lo : sm.lin_hi; +#pragma unroll + for (unsigned 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] *= 256.0f; + res[1] *= 256.0f; + res[2] *= 256.0f; + res[3] *= 256.0f; + } + const unsigned n = (r % 4u) * 4u + tt; + const unsigned off = n * 128u + ((((lane * 4u) >> 4) ^ (n & 7u)) << 4) + + ((lane * 4u) & 15u); + ptx::st_shared( + reinterpret_cast(lin_base + off), + static_cast(c01) | (static_cast(c23) << 16)); + } + if (lane == 0) sm.lin_scale[r] = s; + } +} + +// > 16K gate: the k-th largest sample (k = count * target / visible) of one row +// from its 2048-bin histogram; one warp per row. Same result as finish_sample +// (count = histogram total). +__device__ __noinline__ void finish_sample_histogram(const unsigned* hist, + float* gate, + unsigned visible, + unsigned target, + unsigned lane) { + if (visible <= target) { + if (lane == 0) *gate = -CUDART_INF_F; + __syncwarp(); + return; + } + // Lane l owns the 64 bins of chunks 63-2l and 62-2l (bins descending with the + // lane). + const unsigned* const h0 = hist + (63 - 2 * lane) * 33; + const unsigned* const h1 = hist + (62 - 2 * lane) * 33; + unsigned sum = 0; +#pragma unroll 8 + for (unsigned j = 0; j < 32; ++j) sum += h0[j] + h1[j]; + unsigned incl = sum; +#pragma unroll + for (unsigned dd = 1; dd < 32; dd *= 2) { + const unsigned u = __shfl_up_sync(0xffffffffu, incl, dd); + if (lane >= dd) incl += u; + } + const unsigned tot = __shfl_sync(0xffffffffu, incl, 31); + if (tot == 0) { + if (lane == 0) *gate = -CUDART_INF_F; + __syncwarp(); + return; + } + const unsigned rank = + max(1u, min(tot, (tot * target + visible - 1) / visible)); + const unsigned w = __ffs(__ballot_sync(0xffffffffu, incl >= rank)) - 1; + unsigned above = __shfl_sync(0xffffffffu, incl - sum, w), chosen = 0; +#pragma unroll 1 + for (unsigned cc = 0; cc < 2; ++cc) { + const unsigned c = 63 - 2 * w - cc; + unsigned v = hist[c * 33 + 31 - lane]; // lane 0 = highest bin of the chunk +#pragma unroll + for (unsigned dd = 1; dd < 32; dd *= 2) { + const unsigned u = __shfl_up_sync(0xffffffffu, v, dd); + if (lane >= dd) v += u; + } + const unsigned mask = __ballot_sync(0xffffffffu, above + v >= rank); + if (mask) { + chosen = c * 32 + 31 - (__ffs(mask) - 1); + break; + } + above += __shfl_sync(0xffffffffu, v, 31); + } + if (lane == 0) *gate = sample_floor(chosen); + __syncwarp(); +} + +struct alignas(16) PHandoffData { + uint32_t records[2][8][4][8][32]; + uint32_t counts[2][8][4][32]; + uint32_t bases[2][8]; +}; + +struct ClaimEntry { + unsigned vblock, all_short, kv_base, splits, slot; + unsigned starts[8], ends[8]; +}; + +// Barrier words (static shared, never aliased with data). +template +struct BarIdx { + // > 16K (KVS == 3): 3 per-warpgroup TMEM slots + the linear-term barrier. + static constexpr unsigned NT = KVS == 3 ? 3 : 2; + static constexpr unsigned full_q = 0, empty_q = 1, full_kv = 2, + empty_kv = 2 + KVS, full_tmem = 2 + 2 * KVS, + empty_tmem = full_tmem + NT, + claim_full = empty_tmem + NT, + claim_empty = claim_full + 2, + sel_done = claim_empty + 2, h_full = sel_done + 1, + h_empty = h_full + 16, full_lin = h_empty + 16, + sf_full = full_lin + (KVS == 3 ? 1 : 0), + total = sf_full + (KVS == 3 ? 0 : 2 * KVS); +}; + +__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"); +} + +// 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"); +} + +// > 16K self-drain, out of line: stores this warp's ring records of the last +// handoff interval. recs = records[bank][warp] ([4][8][32]). +__device__ __noinline__ void self_drain_last( + const uint32_t* __restrict__ recs, unsigned n0, unsigned n1, unsigned n2, + unsigned n3, int* __restrict__ counts, unsigned rbase, uint32_t lane_kv, + unsigned flush_base, int64_t row_delta, uint16_t* __restrict__ vals, + int32_t* __restrict__ indices) { + constexpr uint32_t kCap = 16384; + const unsigned lane = threadIdx.x & 31; + const unsigned n[4] = {n0, n1, n2, n3}; + unsigned my_total = 0; +#pragma unroll + for (unsigned row = 0; row < 4; ++row) { + const unsigned total = __reduce_add_sync(0xffffffffu, n[row]); + if (lane == row) my_total = total; + } + unsigned my_start = 0; + if (lane < 4 && my_total) + my_start = atomicAdd(counts + rbase + lane, my_total); + unsigned p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); +#pragma unroll + for (int dd = 1; dd < 32; dd *= 2) { + const unsigned a = __shfl_up_sync(0xffffffffu, p01, dd), + b = __shfl_up_sync(0xffffffffu, p23, dd); + if (lane >= dd) { + p01 += a; + p23 += b; + } + } + const unsigned prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; +#pragma unroll + for (unsigned row = 0; row < 4; ++row) { + const unsigned start = __shfl_sync(0xffffffffu, my_start, row); + for (unsigned j = 0; j < n[row]; ++j) { + const unsigned slot2 = start + prefix[row] - n[row] + j; + const unsigned record = recs[(row * 8 + j) * 32 + lane]; + if (slot2 < kCap) { + const uint32_t physical = + lane_kv + (flush_base * 2 + (record & 255u)) * 128; + const uint64_t off = + static_cast(int64_t(rbase + row) + row_delta) * kCap + + slot2; + fused_candidates::store_candidate_record(vals + off, indices + off, + record >> 8, physical); + } + } + } + __syncwarp(); +} + +template +__device__ __forceinline__ void ho_drain(PHandoffData& d, TBar* hfull, + TBar* hempty, int producer, + unsigned epoch, unsigned row_base, + int* totals, unsigned cap, bool skip, + Store store) { + const int lane = threadIdx.x & 31, bank = epoch & 1; + hfull[bank * 8 + producer].wait((epoch >> 1) & 1); + __syncwarp(); + if (skip) { + hempty[bank * 8 + producer].arrive(); + return; + } // self-drained by the producer + const unsigned base = d.bases[bank][producer]; + unsigned n[4]; + unsigned my_total = 0; +#pragma unroll + for (unsigned row = 0; row < 4; ++row) { + n[row] = d.counts[bank][producer][row][lane]; + unsigned total = __reduce_add_sync(0xffffffffu, n[row]); + if (lane == row) my_total = total; + } + unsigned my_start = 0; + if (lane < 4 && my_total) + my_start = atomicAdd(totals + row_base + lane, my_total); + unsigned p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); +#pragma unroll + for (int dd = 1; dd < 32; dd *= 2) { + unsigned a = __shfl_up_sync(0xffffffffu, p01, dd); + unsigned b = __shfl_up_sync(0xffffffffu, p23, dd); + if (lane >= dd) { + p01 += a; + p23 += b; + } + } + unsigned prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; +#pragma unroll + for (unsigned row = 0; row < 4; ++row) { + unsigned start = __shfl_sync(0xffffffffu, my_start, row); + for (unsigned j = 0; j < n[row]; ++j) { + unsigned slot = start + prefix[row] - n[row] + j; + unsigned record = d.records[bank][producer][row][j][lane]; + if (slot < cap) + store(row_base + row, slot, record, base, producer * 32 + lane); + } + } + __syncwarp(); + hempty[bank * 8 + producer].arrive(); +} + +template +__global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( + unsigned rows, unsigned keys, const unsigned* starts_g, + const unsigned* ends_g, uint16_t* vals, int* indices, int* counts, + FallbackBuffers fallback, unsigned phase_bytes, unsigned sample_target, + unsigned sample_stride, unsigned sample_groups, unsigned 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, bool relative, + unsigned nblocks, unsigned* counter, int* defer_next) { + constexpr bool kSmem = + Interval <= 8; // <= 16K keys: shared candidates, direct emission + constexpr uint32_t kMathRegs = kSmem ? 232 : 224, + kSpecRegs = 168 * 3 - 2 * kMathRegs; + constexpr uint32_t kCandGlobalCap = 16384; + using St = PStorage; + using B = BarIdx; + // > 16K keys: rebalanced scan (linear term on the tensor core, per-warpgroup + // N=144 MMAs). + constexpr bool kLin = !kSmem; + static_assert(kLin == (KVS == 3), + "rebalanced scan needs the 3-KV-stage lin layout"); + constexpr uint32_t kSlotCols = + 160; // 4 rows x 32 heads + 16 linear-term columns + constexpr unsigned kSelB = KVS == 3 ? kSelBase3 : kSelBase; + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& sm = *reinterpret_cast(smem_buffer); + uint8_t* const region = smem_buffer + sizeof(St); + auto& stream = *reinterpret_cast(region); + auto& hdata = *reinterpret_cast(region); + + __shared__ alignas(8) uint64_t bar_words[B::total]; + __shared__ uint32_t tmem_ptr; + __shared__ ClaimEntry entries[2]; + __shared__ float gate_s[8]; + __shared__ unsigned row_count[8], pool_cnt[8]; + __shared__ unsigned sel_recover, sel_staged, sel_short, row_offset[9]; + __shared__ unsigned defer_head, unsafe_s, pause_req, epoch_kind; + + TBar* const bars = reinterpret_cast(bar_words); + TBar& full_q = bars[B::full_q]; + TBar& empty_q = bars[B::empty_q]; + TBar* const full_kv = bars + B::full_kv; + // > 16K the 32 producer lanes stage the sample scales, so each arrives on + // full_kv. + constexpr unsigned kFullKvArrivals = KVS == 3 ? 32 : 1; + // <= 16K: the MMA warp frees each KV stage once its MMAs complete, and the + // per-key scales arrive through a separate ring (2 * KVS deep, sf_full), so + // the math warps never wait on or release KV stages. The producer leads the + // MMA by <= KVS tiles and the MMA leads the math warps by <= 1 tile (2 TMEM + // stages), so a scale stage is never overwritten while being read. + constexpr bool kRel = KVS == 2; + TBar* const sf_full = bars + B::sf_full; + TBar* const empty_kv = bars + B::empty_kv; + TBar* const full_tmem = bars + B::full_tmem; + TBar* const empty_tmem = bars + B::empty_tmem; + TBar* const claim_full = bars + B::claim_full; + TBar* const claim_empty = bars + B::claim_empty; + TBar& sel_done = bars[B::sel_done]; + TBar* const h_full = bars + B::h_full; + TBar* const h_empty = bars + B::h_empty; + TBar* const full_lin = bars + B::full_lin; // > 16K only + + const unsigned warp_idx = cutlass::canonical_warp_idx_sync(); + const unsigned lane = ptx::get_lane_idx(); + constexpr unsigned kClaimConsumers = + kSmem ? 10 : 11; // warps 9, 10, 0-7 (+11) + + 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 with + // the fallback core 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 >16K scan: uniform address math is + // rematerialised every tile). Later epochs only run after a recovery pause. + __shared__ unsigned row_bounds[2][8]; + __shared__ unsigned fslot; + bool more = false; + { + constexpr unsigned epoch = 0; +#include "persistent_epoch.inc" + } +#pragma unroll 1 + for (unsigned 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) { + atomicExch(counter, 0u); + atomicExch(counter + 1, 0u); + } + } +} + +// ---------------- host: launch configuration (> kResidentKeys keys) +// ---------------- +using PersistentKernel = decltype(&fused_mqa_topk_fp8_prefill_kernel<4, 2>); + +struct LaunchConfig { + PersistentKernel kernel; + int smem; // dynamic shared memory bytes + unsigned 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). + unsigned sample_target, sample_stride, sample_groups, sample_tiles; +}; + +// Picks the kernel instantiation, sizes its shared memory (setting the opt-in +// limit) and the gate sample for `keys` keys. +inline cudaError_t launch_config(unsigned keys, LaunchConfig& c) { + // <= 16K keys: shared candidates (cap kCandCap), so aim lower: 2560 + // candidates. Up to 128K: 3072. Both size the sample so the gate's quantile + // rank is ~100 (about 10% count error). Longer contexts aim for 4096 with + // rank 64, clamped to 16-32 sample tiles (rank ~32 at 1M keys): a sample tile + // costs about two scan tiles, so a larger sample costs more than the smaller + // candidate set saves. + const bool small = keys <= 16384, long_ctx = keys > 131072; + c.sample_target = small ? 2560u : (long_ctx ? 4096u : 3072u); + const unsigned rank = long_ctx ? 64u : 100u; + const unsigned rank_tiles = std::max( + 2u, unsigned((uint64_t(rank) * keys + 256ull * c.sample_target - 1) / + (256ull * c.sample_target))); + c.sample_tiles = + long_ctx ? std::min(32u, std::max(16u, rank_tiles)) : rank_tiles; + c.sample_stride = std::max( + 1u, (keys + 256u * c.sample_tiles - 1) / (256u * c.sample_tiles)); + c.sample_groups = + (keys - c.sample_stride / 2 + c.sample_stride - 1) / c.sample_stride; + c.kernel = keys <= 8192 ? fused_mqa_topk_fp8_prefill_kernel<4, 2> + : small ? fused_mqa_topk_fp8_prefill_kernel<8, 2> + : keys <= 32768 ? fused_mqa_topk_fp8_prefill_kernel<16, 3> + : keys <= 131072 ? fused_mqa_topk_fp8_prefill_kernel<32, 3> + : keys <= 524288 ? fused_mqa_topk_fp8_prefill_kernel<64, 3> + : fused_mqa_topk_fp8_prefill_kernel<128, 3>; + + // The recovery rescan reuses the same dynamic shared memory with the fallback + // core's layout. + constexpr int kFallbackSmem = + int(sizeof(fused_mqa_dense::SharedStorage) + sizeof(FallbackShared)); + constexpr int kSmallSmem = int(sizeof(PStorage<2>) + kCandExt * 5u); + constexpr int kLongSmem = + int(sizeof(PStorage<3>) + + std::max(sizeof(PHandoffData), sizeof(SampleShared))); + static_assert(sizeof(SampleShared) <= kCandExt * 5u, + "sample state must fit the candidate region"); + constexpr int kMaxSmem = 232448; // sm_100 opt-in limit per block + if (small) { + c.smem = std::max(kSmallSmem, kFallbackSmem); + } else { + // > 16K: the staged selection may use everything past Q + weights, up to + // the opt-in limit. + cudaFuncAttributes attr; + if (const cudaError_t e = cudaFuncGetAttributes( + &attr, reinterpret_cast(c.kernel)); + e != cudaSuccess) + return e; + c.smem = std::max(kFallbackSmem, kMaxSmem - int(attr.sharedSizeBytes)); + if (c.smem < kLongSmem) return cudaErrorInvalidConfiguration; + } + c.phase_bytes = unsigned(c.smem) - (small ? kSelBase : kSelBase3); + return cudaFuncSetAttribute(reinterpret_cast(c.kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + c.smem); +} + +} // namespace fused_persist diff --git a/csrc/fused_mqa_topk/resident_kernel.cuh b/csrc/fused_mqa_topk/resident_kernel.cuh new file mode 100644 index 000000000000..c21c0b4e8a5b --- /dev/null +++ b/csrc/fused_mqa_topk/resident_kernel.cuh @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: MIT +// <= 4K keys: every score of an 8-row block stays in shared memory as an fp24 +// code, so there is no sample, gate or candidate workspace; one warp then +// selects one row. +#pragma once + +#include "candidates.cuh" +#include "dense_scan.cuh" +#include "fast_select.cuh" + +namespace fused_mqa_resident { +using fused_mqa_dense::BLOCK_Q; +using fused_mqa_dense::Scheduler; +using fused_mqa_dense::SharedStorage; + +constexpr uint32_t kResidentKeys = 4096; +// Pipeline storage, then the block's codes: hi16 array then lo8 array, +// row-major, kResidentKeys per row. +constexpr uint32_t kSmemBytes = + (sizeof(SharedStorage) + BLOCK_Q * kResidentKeys * 3u + 1023u) / 1024u * + 1024u; + +// Stores every score as an fp24 code (selection reads only [start, end) of each +// row, so no masking). +struct CodesPolicy { + static constexpr float kWeightScale = 1.0f, kScoreScale = 0.5f; + uint8_t* codes; + uint32_t query_base; + CUTLASS_DEVICE void begin_block(const Scheduler&, uint32_t, uint32_t base) { + query_base = base; + } + CUTLASS_DEVICE void begin_split(uint32_t, uint32_t) {} + CUTLASS_DEVICE void score(uint32_t i, uint32_t kv_offset, uint32_t, uint32_t, + float score) { + if (kv_offset < kResidentKeys) { + const uint32_t code = fused_candidates::candidate_fp24_code(score); + const uint32_t at = (query_base + i) * kResidentKeys + kv_offset; + reinterpret_cast(codes)[at] = static_cast(code >> 8); + codes[BLOCK_Q * kResidentKeys * 2 + at] = static_cast(code); + } + } + CUTLASS_DEVICE void end_split(uint32_t, uint32_t, uint32_t) {} + CUTLASS_DEVICE void end_block(const Scheduler&, uint32_t) {} +}; + +} // namespace fused_mqa_resident + +// <= 4K keys: one CTA per 8-row block, heaviest (latest rows) first. +__global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_resident_kernel( + unsigned rows, unsigned keys, const unsigned* starts_g, + const unsigned* ends_g, int* output, bool relative, + const __grid_constant__ cute::TmaDescriptor tq, + const __grid_constant__ cute::TmaDescriptor ts, + const __grid_constant__ cute::TmaDescriptor tk, + const __grid_constant__ cute::TmaDescriptor tw) { + using namespace fused_mqa_resident; + const unsigned vblock = gridDim.x - 1 - blockIdx.x, row_base = vblock * 8; + // The block's row bounds clamped to the keys; rows past the end see no keys. + __shared__ unsigned row_bounds[2][8]; + if (threadIdx.x < 8) { + const unsigned 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 unsigned* const starts = row_bounds[0] - row_base; + const unsigned* const ends = row_bounds[1] - row_base; + // Rows with <= 2048 visible keys output all of them (relative indices start + // at 0). + const auto write_short = [&](unsigned r, unsigned t, unsigned stride) { + const unsigned visible = ends[r] - starts[r]; + for (unsigned j = t; j < 2048; j += stride) + output[(uint64_t)r * 2048 + j] = + j < visible ? int((relative ? 0u : starts[r]) + j) : -1; + }; + bool all_short = true; + for (unsigned r = row_base; r < row_base + 8; ++r) + all_short &= ends[r] - starts[r] <= 2048u; + if (all_short) { + for (unsigned r = row_base; r < min(row_base + 8, rows); ++r) + write_short(r, threadIdx.x, 384); + return; + } + + extern __shared__ __align__(1024) uint8_t smem[]; + uint8_t* const codes = smem + sizeof(SharedStorage); + CodesPolicy policy{codes, 0}; + fused_mqa_dense::score_block(Scheduler(vblock, rows, keys, starts, ends), + policy, tq, tk, ts, tw); + __syncthreads(); + if (threadIdx.x < 256) cutlass::arch::warpgroup_reg_dealloc<128>(); + __syncthreads(); + if (threadIdx.x >= 256) cutlass::arch::warpgroup_reg_alloc<128>(); + const unsigned warp = threadIdx.x / 32; + if (warp < 8 && row_base + warp < rows) { + const unsigned r = row_base + warp, s0 = starts[r], visible = ends[r] - s0; + if (visible <= 2048) { + write_short(r, threadIdx.x & 31, 32); + } else { + int* o = output + (uint64_t)r * 2048; + const uint16_t* hi = + reinterpret_cast(codes) + warp * kResidentKeys; + const uint8_t* lo = codes + 8 * kResidentKeys * 2 + warp * kResidentKeys; + // Scratch (bracket buffer, or the fallback's histogram) reuses the idle + // Q/KV stage memory. + static_assert(8 * fast_select::kScratch3Bytes <= sizeof(SharedStorage) && + fast_select::kScratch3Bytes >= fast_select::kBins * 4, + "select scratch must fit the Q/KV stages"); + uint8_t* scratch = smem + warp * fast_select::kScratch3Bytes; + if (!fast_select::select_resident_bracket(hi, lo, s0, visible, o, + scratch)) + fast_select::select_resident(hi + s0, lo + s0, visible, s0, o, + reinterpret_cast(scratch)); + } + } + // Relative indices: selected rows are rewritten in place (L2-hot) once final. + if (relative) { + __syncthreads(); + for (unsigned i = threadIdx.x; i < 8 * 512; i += 384) { + const unsigned r = row_base + i / 512; + if (r >= rows || ends[r] - starts[r] <= 2048u) continue; + const int s = int(starts_g[r]); + if (s == 0) continue; + int4* p = reinterpret_cast(output + (uint64_t)r * 2048) + i % 512; + 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; + } + } +} diff --git a/csrc/fused_mqa_topk/row_select.cuh b/csrc/fused_mqa_topk/row_select.cuh new file mode 100644 index 000000000000..e93037994d86 --- /dev/null +++ b/csrc/fused_mqa_topk/row_select.cuh @@ -0,0 +1,434 @@ +// SPDX-License-Identifier: MIT +// Single-warp top-2048 selection of one row's candidates for the persistent +// kernel. Both selections bracket the 2048-th largest fp24 code with 512 +// sampled codes, place the codes above the bracket directly, and resolve the +// bracket with histogram levels on register-resident codes; elements are placed +// by per-lane counts and one warp prefix scan (no per-element ballots). +#pragma once +#include "fast_select.cuh" + +namespace row_select { +using namespace fast_select; + +// r-th largest (1-based) of the 512 register samples 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 unsigned (&smp)[NS], + unsigned r_hi, unsigned r_lo, + unsigned* h8, unsigned& h_hi, + unsigned& h_lo) { + const unsigned lane = threadIdx.x & 31; +#pragma unroll + for (unsigned b = lane; b < 192; b += 32) + reinterpret_cast(h8)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (unsigned u = 0; u < NS; ++u) atomicAdd(h8 + (smp[u] >> 8), 1u); + __syncwarp(); + unsigned ab_hi, ab_lo; + const unsigned b_hi = find_bin_n<256>(h8, r_hi, ab_hi); + const unsigned b_lo = find_bin_n<256>(h8, r_lo, ab_lo); + unsigned* const lh = h8 + 256; + unsigned* const ll = h8 + 512; +#pragma unroll + for (unsigned u = 0; u < NS; ++u) { + const unsigned 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(); + unsigned 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 +} + +// find_bin_n with each lane's kPer contiguous bins loaded as 16-byte vectors +// and searched in registers. +template +__device__ __forceinline__ unsigned find_bin_v(const unsigned* hist, + unsigned rank, unsigned& above) { + const unsigned lane = threadIdx.x & 31; + constexpr unsigned 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); + unsigned h[kPer]; +#pragma unroll + for (unsigned 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; + } + unsigned sum = 0; +#pragma unroll + for (unsigned j = 0; j < kPer; ++j) sum += h[j]; + unsigned inclusive = sum; +#pragma unroll + for (unsigned d = 1; d < 32; d *= 2) { + const unsigned v = __shfl_up_sync(Full, inclusive, d); + if (lane >= d) inclusive += v; + } + const unsigned before = inclusive - sum; + const unsigned winner = + __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; + // Descending walk inside the lane: bin index Bins - lane * kPer - 1 - j <-> + // h[kPer - 1 - j]. + unsigned acc = before, chosen = 0, found = 0, acc_at = before; +#pragma unroll + for (unsigned j = 0; j < kPer; ++j) { + const unsigned 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(Full, acc_at, winner); + return __shfl_sync(Full, chosen, winner); +} +// Sampled hi16 bracket (16-step bit search over the 512 register samples). +__device__ __forceinline__ void bit_bracket( + const unsigned (&smp)[kSamplesPerLane], unsigned r_hi, unsigned r_lo, + unsigned& h_hi, unsigned& h_lo) { + h_hi = 0; + h_lo = 0; +#pragma unroll 1 + for (int bit = 15; bit >= 0; --bit) { + const unsigned ch = h_hi | (1u << bit), cl = h_lo | (1u << bit); + unsigned nh = 0, nl = 0; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) { + nh += smp[u] >= ch; + nl += smp[u] >= cl; + } + const unsigned both = + __reduce_add_sync(Full, nh | nl << 16); // counts <= 512 fit 16 bits + if ((both & 0xffffu) >= r_hi) h_hi = ch; + if ((both >> 16) >= r_lo) h_lo = cl; + } +} + +// <= 16K keys: one row's shared-memory candidates (hi16 / lo8 code arrays and +// the uint16 key array idx_row, read with 16-byte loads). Bracket keys are +// preloaded into registers before the emit. Returns false (out untouched) when +// the bracket misses; the caller then runs select_resident. +__device__ bool select_shared_row(const uint16_t* __restrict__ hi_row, + const uint8_t* __restrict__ lo_row, + const uint16_t* __restrict__ idx_row, + unsigned count, unsigned real, + int32_t* __restrict__ out, + unsigned char* scratch) { + const unsigned lane = threadIdx.x & 31; + unsigned h_hi, h_lo; + { + unsigned smp[kSamplesPerLane], n_real = 0; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) { + smp[u] = hi_row[(((u * 32 + lane) * 2 + 1) * count >> 10)]; + n_real += smp[u] != 0; + } + const unsigned s_real = __reduce_add_sync(Full, n_real); + if (s_real == 0 || real <= K) return false; + const unsigned expect = (K * s_real + real / 2) / real; + const unsigned r_hi = expect > 40 ? expect - 40 : 1, + r_lo = min(expect + 40, s_real); + bit_bracket(smp, r_hi, r_lo, h_hi, h_lo); + h_lo = max(h_lo, 1u); + } + const unsigned units = (count + 7) / 8; + const uint4* hv = reinterpret_cast(hi_row); + const uint4* iv = reinterpret_cast(idx_row); + unsigned na = 0, nm = 0; +#pragma unroll 2 + for (unsigned unit = lane; unit < units; unit += 32) + classify8(hv[unit], unit * 8, 0u, count, h_hi, h_lo, + [&](unsigned, bool a, bool m) { + na += a; + nm += m; + }); + unsigned above, mid; + unsigned pa = warp_exclusive_scan(na, above), + pm = warp_exclusive_scan(nm, mid); + if (above >= K || above + mid < K || mid > kMid3Cap) return false; + uint16_t* stage = reinterpret_cast(scratch); + uint16_t* midbuf = reinterpret_cast(scratch + K * 2); + unsigned* hist = reinterpret_cast(scratch + K * 2 + kMid3Cap * 2); +#pragma unroll 2 + for (unsigned unit = lane; unit < units; unit += 32) { + const uint4 h4 = hv[unit], k4 = iv[unit]; + const unsigned kw[4] = {k4.x, k4.y, k4.z, k4.w}; + classify8(h4, unit * 8, 0u, count, h_hi, h_lo, + [&](unsigned e, bool a, bool m) { + const unsigned k = e & 7u; + const unsigned key = (kw[k / 2] >> (16 * (k & 1))) & 0xffffu; + if (a) stage[pa++] = uint16_t(key); + if (m) midbuf[pm++] = uint16_t(e); + }); + } + __syncwarp(); + constexpr unsigned R = kMid3Cap / 32; + unsigned mc[R]; +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + const unsigned i = lane + 32 * u; + mc[u] = 0; + if (i < mid) { + const unsigned e = midbuf[i]; + mc[u] = unsigned(hi_row[e]) << 8 | lo_row[e]; + } + } + unsigned rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; +#pragma unroll 1 + while (top > base) { + const unsigned width = top - base + 1, spread = (width - 1) / kMidBins; + const unsigned shift = spread ? 32 - __clz(spread) : 0; +#pragma unroll + for (unsigned b = lane; b < kMidBins / 4; b += 32) + reinterpret_cast(hist)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + if (u * 32 >= mid) break; + if (mc[u] >= base && mc[u] <= top) + atomicAdd(hist + ((mc[u] - base) >> shift), 1u); + } + __syncwarp(); + unsigned above_bin; + const unsigned chosen = find_bin_v(hist, rank, above_bin); + rank -= above_bin; + base += chosen << shift; + top = min(base + (1u << shift) - 1, top); + __syncwarp(); + } + const unsigned th = base; + { + unsigned ng = 0, nt = 0, mk[R]; +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + ng += mc[u] > th; + nt += mc[u] == th; + mk[u] = 0; + if (u * 32 < mid && mc[u] >= th) mk[u] = idx_row[midbuf[lane + 32 * u]]; + } + unsigned wg, wt; + unsigned pg = above + warp_exclusive_scan(ng, wg), + pt = K - rank + warp_exclusive_scan(nt, wt); +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + if (u * 32 >= mid) break; + if (mc[u] > th) + stage[pg++] = uint16_t(mk[u]); + else if (mc[u] == th) { + if (pt < K) stage[pt] = uint16_t(mk[u]); + ++pt; + } + } + } + __syncwarp(); +#pragma unroll 4 + for (unsigned j = lane * 4; j < K; j += 128) { + const uint2 v = *reinterpret_cast(stage + j); + *reinterpret_cast(out + j) = make_int4( + int(v.x & 0xffffu), int(v.x >> 16), int(v.y & 0xffffu), int(v.y >> 16)); + } + return true; +} + +// > 16K keys: 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 spilled into the scan loop: 290K -> +// 860K cycles/block in a first try). +// 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__ bool select_staged_row( + unsigned* __restrict__ codes, unsigned count, + const uint16_t* __restrict__ values, const int32_t* __restrict__ packed, + int32_t* __restrict__ out, unsigned* hist) { + const unsigned lane = threadIdx.x & 31; + if (count <= K) return false; + const unsigned chunks = (count + 7) / 8; + constexpr unsigned C = 2; + const uint4* v4p = reinterpret_cast(values); + const uint4* p4p = reinterpret_cast(packed); +#pragma unroll 1 + for (unsigned c0i = lane; c0i < chunks; c0i += 32 * C) { + uint4 vv[C], pp0[C], pp1[C]; +#pragma unroll + for (unsigned u = 0; u < C; ++u) { + const unsigned 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 (unsigned u = 0; u < C; ++u) { + const unsigned c = c0i + 32 * u; + if (c < chunks) { + const unsigned vw[4] = {vv[u].x, vv[u].y, vv[u].z, vv[u].w}; + const unsigned 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}; + unsigned cw[8]; +#pragma unroll + for (unsigned k = 0; k < 8; ++k) + cw[k] = fused_candidates::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 (unsigned k = 0; k < 8; ++k) + if (8 * c + k < count) codes[8 * c + k] = cw[k]; + } + } + } + } + __syncwarp(); + unsigned h_hi, h_lo; + if (count <= 3584) { + unsigned smp[kSamplesPerLane]; +#pragma unroll + for (unsigned u = 0; u < kSamplesPerLane; ++u) + smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 10)] >> 8; + const unsigned expect = (K * 512 + count / 2) / count; + const unsigned r_hi = expect > 40 ? expect - 40 : 1, + r_lo = min(expect + 40, 512u); + radix_bracket(smp, r_hi, r_lo, hist, h_hi, h_lo); + } else { + // Long-context rows (4-6K candidates): +-40 of 512 samples gives a bracket + // of ~16% of the row, over kMid3Cap for 10-20% of the rows at 512K-1M. 1024 + // samples, +-48 (~3 sd): ~9%. + unsigned smp[2 * kSamplesPerLane]; +#pragma unroll + for (unsigned u = 0; u < 2 * kSamplesPerLane; ++u) + smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 11)] >> 8; + const unsigned expect = (K * 1024 + count / 2) / count; + const unsigned r_hi = expect > 48 ? expect - 48 : 1, + r_lo = min(expect + 48, 1024u); + radix_bracket<2 * kSamplesPerLane>(smp, r_hi, r_lo, hist, h_hi, h_lo); + } + const uint4* c4p = reinterpret_cast(codes); + const unsigned quads = (count + 3) / 4; + const unsigned lt = (1u << lane) - 1u; + uint16_t* midbuf = reinterpret_cast(hist); + unsigned* h512 = hist + kMid3Cap / 2; + // Collect with warp-ballot placement: passing lanes of one step write + // consecutive slots, so the global stores of the above elements are coalesced + // (per-lane runs were 4x slower: 32 lines/STG). Writes are bounded; the + // bracket check follows (out may be partially written on failure: the + // fallback select_row rewrites all K slots). + unsigned above = 0, mid = 0; + constexpr unsigned Q = 4; +#pragma unroll 1 + for (unsigned q0 = 0; q0 < quads; q0 += 32 * Q) { + uint4 pp[Q], cc[Q]; +#pragma unroll + for (unsigned u = 0; u < Q; ++u) { + const unsigned q = min(q0 + 32 * u + lane, quads - 1); + pp[u] = __ldcg(p4p + q); + cc[u] = c4p[q]; + } +#pragma unroll + for (unsigned u = 0; u < Q; ++u) { + const unsigned q = q0 + 32 * u + lane; + const unsigned w[4] = {cc[u].x, cc[u].y, cc[u].z, cc[u].w}; + const unsigned p[4] = {pp[u].x, pp[u].y, pp[u].z, pp[u].w}; +#pragma unroll + for (unsigned k = 0; k < 4; ++k) { + const unsigned 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 unsigned ma = __ballot_sync(Full, a), mm = __ballot_sync(Full, m); + const unsigned at = above + __popc(ma & lt), mt = mid + __popc(mm & lt); + if (a && at < K) + out[at] = fused_candidates::candidate_decode_index(int32_t(p[k])); + if (m && mt < kMid3Cap) midbuf[mt] = uint16_t(e); + above += __popc(ma); + mid += __popc(mm); + } + } + } + __syncwarp(); + if (above >= K || above + mid < K || mid > kMid3Cap) return false; + constexpr unsigned R = kMid3Cap / 32; + unsigned mc[R]; +#pragma unroll + for (unsigned u = 0; u < R; ++u) { + const unsigned i = lane + 32 * u; + mc[u] = 0; + if (i < mid) mc[u] = codes[midbuf[i]]; + } + unsigned rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; +#pragma unroll 1 + while (top > base) { + const unsigned width = top - base + 1, spread = (width - 1) / kMidBins; + const unsigned shift = spread ? 32 - __clz(spread) : 0; +#pragma unroll + for (unsigned b = lane; b < kMidBins / 4; b += 32) + reinterpret_cast(h512)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (unsigned u = 0; u < R; ++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(); + unsigned above_bin; + const unsigned chosen = find_bin_v(h512, rank, above_bin); + rank -= above_bin; + base += chosen << shift; + top = min(base + (1u << shift) - 1, top); + __syncwarp(); + } + const unsigned th = base; + { + // Ballot placement (coalesced stores): greater codes fill [above, K - + // rank), ties [K - rank, K). + unsigned pg = above, pt = K - rank; +#pragma unroll + for (unsigned h = 0; h < 2; ++h) { + if (h * 16 * 32 >= mid) break; + int32_t pk[16]; +#pragma unroll + for (unsigned j = 0; j < 16; ++j) { + const unsigned u = h * 16 + j; + pk[j] = (lane + 32 * u < mid && mc[u] >= th) + ? __ldcg(packed + midbuf[lane + 32 * u]) + : 0; + } +#pragma unroll + for (unsigned j = 0; j < 16; ++j) { + const unsigned 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 unsigned mg = __ballot_sync(Full, g), + mt = __ballot_sync(Full, tq); + const unsigned ag = pg + __popc(mg & lt), at = pt + __popc(mt & lt); + if (g) out[ag] = fused_candidates::candidate_decode_index(pk[j]); + if (tq && at < K) + out[at] = fused_candidates::candidate_decode_index(pk[j]); + pg += __popc(mg); + pt += __popc(mt); + } + } + } + return true; +} +} // namespace row_select diff --git a/csrc/fused_mqa_topk/sample_shared.cuh b/csrc/fused_mqa_topk/sample.cuh similarity index 93% rename from csrc/fused_mqa_topk/sample_shared.cuh rename to csrc/fused_mqa_topk/sample.cuh index b603690d4489..317c8c7764bc 100644 --- a/csrc/fused_mqa_topk/sample_shared.cuh +++ b/csrc/fused_mqa_topk/sample.cuh @@ -3,9 +3,8 @@ #include #include #include -constexpr unsigned SampleQ = 8, SampleN = 8192; -// target: candidate count the gate aims for (0 = legacy 4096/6144 by visible -// length). +constexpr unsigned SampleQ = 8; +// target: candidate count the gate aims for. struct SampleBuffers { float* ring; float* retained; @@ -37,8 +36,7 @@ __device__ __forceinline__ float sample_floor(unsigned b) { } __device__ void finish_sample(SampleShared& s, SampleBuffers b, unsigned local, unsigned row, unsigned visible, unsigned tid) { - const unsigned target = - b.target ? b.target : (visible <= 65536 ? 4096u : 6144u); + const unsigned target = b.target; if (visible <= target) { if (tid == 0) b.ring[row] = -CUDART_INF_F; __syncwarp(); diff --git a/csrc/fused_mqa_topk/sample_core.cuh b/csrc/fused_mqa_topk/sample_core.cuh deleted file mode 100644 index 150f1a467297..000000000000 --- a/csrc/fused_mqa_topk/sample_core.cuh +++ /dev/null @@ -1,627 +0,0 @@ -// SPDX-License-Identifier: MIT -// Experimental derivative of DeepGEMM 2.8 sm100_mqa_logits.cuh. -// Persistent DeepGEMM FP8 MQA producer with fixed-threshold candidate emission. -#pragma once -#ifndef MQA_TMEM_ROWS - #define MQA_TMEM_ROWS 2 -#endif - -#include -#include "sample_shared.cuh" -#ifndef MQA_HANDOFF_INTERVAL - #define MQA_HANDOFF_INTERVAL 64 -#endif -static_assert(MQA_HANDOFF_INTERVAL == 32 || MQA_HANDOFF_INTERVAL == 64 || - MQA_HANDOFF_INTERVAL == 128); - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// Shared SM100 MQA logits core plus contiguous-KV and paged entries -// Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature - -namespace seed_sample { -using namespace deep_gemm; - -// Convert runtime valid-token count to `cute::Int` so token loops stay -// compile-time constant -template -CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, - Fn&& fn) { - if constexpr (kCandidate <= 1) { - fn(cute::Int<1>{}); - } else if (num_block_tokens >= kCandidate) { - fn(cute::Int{}); - } else { - dispatch_num_block_tokens(num_block_tokens, - static_cast(fn)); - } -} - -// Load heads using power-of-two TMEM chunks. -template -CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, - float* accum) { - constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; - if constexpr (kNumRemaining > 0) { - constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 - : kNumRemaining >= 32 ? 32 - : kNumRemaining >= 16 ? 16 - : kNumRemaining >= 8 ? 8 - : 4; - ptx::tmem_load_32dp32b( - tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); - load_tmem_heads_decomposed(tmem_col, accum); - } -} - -// Shared device core parameterized by dtype and scheduler geometry/addressing -template < - uint32_t kNumHeads, uint32_t kHeadDim, bool kIsMXSF, - bool kIsCompressedLogits, bool kCleanLogits, uint32_t BLOCK_Q, - uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages, - uint32_t kNumSMs, uint32_t kNumSpecializedThreads, uint32_t kNumMathThreads, - typename qk_dtype_t, typename logits_dtype_t, typename reduce_dtype_t, - typename MakeScheduler, uint32_t kNumMathWarpGroups = kNumMathThreads / 128> -CUTLASS_DEVICE void sm100_mqa_logits_core_impl( - const uint32_t num_q_tokens, SampleBuffers buffers, - const cute::TmaDescriptor& tensor_map_q, - const cute::TmaDescriptor& tensor_map_sf_q, - const cute::TmaDescriptor& tensor_map_kv, - const cute::TmaDescriptor& tensor_map_sf_kv, - const cute::TmaDescriptor& tensor_map_weights, - const MakeScheduler& make_scheduler) { - constexpr bool kIsFP4 = cute::is_same_v; - - const auto sm_idx = blockIdx.x; - const auto warp_idx = cutlass::canonical_warp_idx_sync(); - const auto lane_idx = ptx::get_lane_idx(); - constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; - - const auto kNegInf = -cute::numeric_limits::infinity(); - - if (warp_idx == kSpecWarpStart) { - cute::prefetch_tma_descriptor(&tensor_map_q); - cute::prefetch_tma_descriptor(&tensor_map_sf_q); - cute::prefetch_tma_descriptor(&tensor_map_weights); - cute::prefetch_tma_descriptor(&tensor_map_kv); - cute::prefetch_tma_descriptor(&tensor_map_sf_kv); - } - - static constexpr uint32_t kNumTmemStages = 2; - static_assert(BLOCK_Q == SampleQ && kNumMathThreads == 256 && !kIsMXSF); - static constexpr uint32_t kNumUTCCPAlignedElems = 128; - static constexpr uint32_t UMMA_M = 128; - static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; - static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); - static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; - static constexpr uint32_t kNumSFQ = - kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kNumSFKV = - kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kNumQKBytesPerToken = - kIsFP4 ? (kHeadDim / 2) : kHeadDim; - static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = - BLOCK_QH * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = - SPLIT_KV * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = - kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; - static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = - kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); - static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( - kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); - static constexpr uint32_t kNumWeightElementsPerRow = - kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); - static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = - BLOCK_Q * kNumWeightBytesPerRow; - - DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, - "Invalid threads"); - DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and - SPLIT_KV % kNumUTCCPAlignedElems == 0, - "Invalid `SPLIT_KV`"); - DG_STATIC_ASSERT(not(kIsCompressedLogits and kCleanLogits), - "Compressed logits cannot be cleaned in-kernel"); - - using SharedStorage = layout::MQALogitsSharedStorage< - kNumHeads, kHeadDim, kIsMXSF, BLOCK_Q, SPLIT_KV, kNumQStages, - kNumKVStages, kNumTmemStages, qk_dtype_t, reduce_dtype_t>; - extern __shared__ __align__(SharedStorage::kSwizzleAlignment) - uint8_t smem_buffer[]; - auto& smem = *reinterpret_cast(smem_buffer); - constexpr uint32_t kRingSlots = 8; - constexpr uint32_t kRecordBlockMask = 255; - auto& stream = - *reinterpret_cast(smem_buffer + sizeof(SharedStorage)); - stream.init(threadIdx.x); - - constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; - constexpr uint32_t kNumTmemCols = - utils::get_num_aligned_tmem_cols(); - constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; - constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; - DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); - - if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { -#pragma unroll - for (uint32_t i = 0; i < kNumQStages; ++i) { - smem.full_q_barriers[i].init(1); - smem.empty_q_barriers[i].init(kNumMathThreads + 32); - } -#pragma unroll - for (uint32_t i = 0; i < kNumKVStages; ++i) { - smem.full_kv_barriers[i].init(1); - smem.empty_kv_barriers[i].init(kIsMXSF ? 1 : kNumMathThreads); - } -#pragma unroll - for (uint32_t i = 0; i < kNumTmemStages; ++i) { - smem.full_tmem_barriers[i].init(1); - smem.empty_tmem_barriers[i].init(256); - } - cutlass::arch::fence_barrier_init(); - } - __syncwarp(); - - if (warp_idx == kSpecWarpStart + 2) - cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); - __syncthreads(); - - RingPipeline q_pipeline; - RingPipeline kv_pipeline; - RingPipeline tmem_pipeline; - - constexpr uint32_t kNumSpecializedRegisters = 56; - constexpr uint32_t kNumMathRegisters = 224; - - cudaGridDependencySynchronize(); - - if (warp_idx == kSpecWarpStart) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - if (cute::elect_one_sync()) { - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - tma::copy( - &tensor_map_q, &smem.full_q_barriers[q_stage_idx], - smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); - if constexpr (kIsMXSF) - tma::copy( - &tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], - smem.smem_sf_q[q_stage_idx], 0, q_token_base); - 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_SF_Q_SIZE_PER_STAGE + - SMEM_WEIGHT_SIZE_PER_STAGE); - } - } - __syncwarp(); - } else if (warp_idx == kSpecWarpStart + 1) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - - auto scheduler = make_scheduler(sm_idx); - uint32_t cached_kv_page_base = 0; - uint32_t cached_kv_page_coord = 0; - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { - cached_kv_page_base = cute::numeric_limits::max(); -#pragma unroll 1 - for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; - ++kv_split_idx) { - if constexpr (decltype(scheduler)::kIsPaged) { - constexpr uint32_t kPageKV = decltype(scheduler)::kPageKV; - constexpr uint32_t kNumPagesPerSplit = - decltype(scheduler)::kNumPagesPerSplit; - DG_STATIC_ASSERT(kNumPagesPerSplit <= 32, - "Split spans more pages than a warp can cache"); - - const uint32_t kv_page_base = - (kv_base + kv_split_idx) * kNumPagesPerSplit; - if (kv_page_base < cached_kv_page_base or - kv_page_base + kNumPagesPerSplit > cached_kv_page_base + 32) { - cached_kv_page_base = (kv_page_base / 32) * 32; - cached_kv_page_coord = scheduler.get_kv_page_coord_by_page_offset( - cached_kv_page_base + lane_idx); - } - - CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); - if (cute::elect_one_sync()) - smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); - __syncwarp(); - - int page_coords[kNumPagesPerSplit]; -#pragma unroll - for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; - ++page_idx) { - const auto src_lane = - static_cast(kv_page_base - cached_kv_page_base + page_idx); - page_coords[page_idx] = - __shfl_sync(0xffffffff, cached_kv_page_coord, src_lane); - } - - if (cute::elect_one_sync()) { -#pragma unroll - for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; - ++page_idx) { - tma::copy( - &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_kv[kv_stage_idx] + - page_idx * kPageKV * kNumQKBytesPerToken, - 0, 0, 1, page_coords[page_idx]); - tma::copy( - &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_sf_kv[kv_stage_idx] + page_idx * kPageKV, 0, - page_coords[page_idx]); - } - 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 (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); - - // Strided sample: groups t*256.. at offset S/2 through the {128, G, - // S} KV view. Scales are read per key by the math warps, so only KV - // bytes arrive here. - tma::copy( - &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_kv[kv_stage_idx], 0, - scheduler.sample_group(kv_split_idx, 0), 1, - scheduler.sample_center()); - smem.full_kv_barriers[kv_stage_idx].arrive_and_expect_tx( - SMEM_KV_SIZE_PER_STAGE); - } - __syncwarp(); - } - } - } else if (warp_idx == kSpecWarpStart + 2) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); - - auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { - DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, - "Invalid aligned elements"); - uint32_t values[4]; -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) - values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); - __syncwarp(); - ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], - values[3]); - }; - - auto sf_desc = mma::sm100::make_sf_desc(nullptr); - - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - - if constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; - utccp_required_smem_warp_transpose(smem_ptr); - } - cutlass::arch::fence_view_async_shared(); -#pragma unroll - for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - if (cute::elect_one_sync()) - cute::SM100_UTCCP_4x32dp128bit_1cta::copy( - sf_desc, kTmemStartColOfSFQ + i * 4); - __syncwarp(); - } - } - - 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 constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; - utccp_required_smem_warp_transpose(smem_ptr); - } - cutlass::arch::fence_view_async_shared(); - } - - if (cute::elect_one_sync()) { - if constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute::SM100_UTCCP_4x32dp128bit_1cta::copy( - sf_desc, kTmemStartColOfSFKV + i * 4); - } - } -#pragma unroll - for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { - CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); - 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(); - - if constexpr (kIsMXSF) { - DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or - kHeadDim == 64 or kHeadDim == 128, - "Invalid head dim"); - - constexpr uint32_t kPackFactor = - get_smem_pack_factor(); - constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; - - using mma_op_t = - cute::conditional_t; - auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< - qk_dtype_t, qk_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, - UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K>(); -#pragma unroll - for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { - auto runtime_instr_desc = - mma::sm100::make_runtime_instr_desc_with_sf_id( - instr_desc, k * kPackFactor, k * kPackFactor); - 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); - mma_op_t::fma(a_desc, b_desc, tmem_addr, k, runtime_instr_desc, - kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ); - } - } else { - 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); -#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(); - if constexpr (kIsMXSF) - cutlass::arch::umma_arrive(reinterpret_cast( - &smem.empty_kv_barriers[kv_stage_idx])); - } - smem.empty_q_barriers[q_stage_idx].arrive(); - } - } else if (warp_idx == kSpecWarpStart + 3) { - cutlass::arch::warpgroup_reg_dealloc<56>(); - } else if (warp_idx < kSpecWarpStart) { - cutlass::arch::warpgroup_reg_alloc<224>(); - - constexpr bool kLoadSeqBounds = true; - uint32_t seq_k_start[4]; - uint32_t seq_k_end[4]; - const auto math_warpgroup_idx = warp_idx / 4; - const auto math_thread_idx = warp_idx * 32 + lane_idx; - DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, - "Math warp groups exceed TMEM stages"); - const unsigned query_base = math_warpgroup_idx * (SampleQ / 2); - - constexpr bool kIsReduceBF16 = not cute::is_same_v; - DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); - DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); - using weights_reg_dtype_t = - cute::conditional_t; - constexpr uint32_t kNumWeightsRegPerToken = - kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; - weights_reg_dtype_t weights[4][32]; - float accum[kNumHeads]; - - unsigned produce_epoch = 0; - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - - const auto process_q_block = [&](auto num_valid_tokens_t) { - constexpr uint32_t kNumValidTokens = SampleQ / 2; - uint32_t lane_counts[4] = {}; - uint32_t flush_base = 0; - float gate_score[4]; - -#pragma unroll - for (uint32_t i = 0; i < kNumValidTokens; ++i) { - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - // Use the raw-score threshold estimated from the distributed sample. - // This threshold remains fixed for the complete scan. - gate_score[i] = -CUDART_INF_F; - seq_k_start[i] = cute::min(scheduler.cu_seq_len_k_start[row], - scheduler.num_kv_tokens); - seq_k_end[i] = cute::min(scheduler.cu_seq_len_k_end[row], - scheduler.num_kv_tokens); - const auto smem_weights_row = - smem.smem_weights[q_stage_idx] + - (query_base + i) * kNumWeightElementsPerRow; - if constexpr (kIsReduceBF16) { - // Load two bf16 weights at a time as one packed shared u32 - const auto packed_row = - reinterpret_cast(smem_weights_row); -#pragma unroll - for (uint32_t j = 0; j < kNumHeads / 2; ++j) { - const auto packed = ptx::ld_shared(packed_row + j); - weights[i][j] = ptx::exchange( - *reinterpret_cast(&packed), 0); - } - } else { -#pragma unroll - for (unsigned j = 0; j < 32; ++j) - weights[i][j] = 0.5f * ptx::ld_shared(smem_weights_row + j); - } - } - - 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); -#pragma unroll - for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { - const unsigned kv_lane = (math_thread_idx % 128) + kv_half * 128; - const unsigned kv_offset = - scheduler.get_logits_col(kv_base, kv_split_idx, kv_lane); - const bool in_view = scheduler.sample_group(kv_split_idx, kv_lane) < - scheduler.sample_groups; - const float scale_kv = - in_view ? __ldg(scheduler.sample_scales + kv_offset) : 0.0f; - 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(); - } - - static_assert(!kIsReduceBF16, "FP32-only pair reduction"); - float accum_pair[MQA_TMEM_ROWS][kNumHeads]; -#pragma unroll - for (uint32_t pair = 0; pair < kNumValidTokens; - pair += MQA_TMEM_ROWS) { - // Load only MQA_TMEM_ROWS query rows at once to limit - // live accumulator registers during score reduction. - const uint32_t tmem_addr = - tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads; - ptx::tmem_load_32dp32b( - tmem_addr, reinterpret_cast(accum_pair)); - cutlass::arch::fence_view_async_tmem_load(); - if (pair + MQA_TMEM_ROWS >= kNumValidTokens) { - ptx::tcgen05_before_thread_sync(); - smem.empty_tmem_barriers[tmem_stage_idx].arrive(); - } - - float2 sum_0[MQA_TMEM_ROWS] = {}; - float2 sum_1[MQA_TMEM_ROWS] = {}; -#pragma unroll - for (uint32_t j = 0; j < kNumHeads; j += 4) { -#pragma unroll - for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { - if (pair + r >= kNumValidTokens) continue; - const uint32_t i = pair + r; - auto* a = accum_pair[r]; - auto transform = [&](uint32_t at, float2 sum) { - const auto x = make_float2(a[at], a[at + 1]); - const auto abs_x = - make_float2(fabsf(a[at]), fabsf(a[at + 1])); - const auto w = - make_float2(weights[i][at], weights[i][at + 1]); - return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); - }; - sum_0[r] = transform(j, sum_0[r]); - sum_1[r] = transform(j + 2, sum_1[r]); - } - } -#pragma unroll - for (uint32_t r = 0; r < MQA_TMEM_ROWS; ++r) { - if (pair + r >= kNumValidTokens) continue; - const uint32_t i = pair + r; - const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - const float reduced = sum.x + sum.y; - const float bq = reduced * static_cast(scale_kv); - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - bool valid = in_view && kv_offset >= seq_k_start[i] && - kv_offset < seq_k_end[i] && isfinite(bq) && - bq > gate_score[i]; - stream.scores[query_base + i][kv_lane] = - valid ? bq : -CUDART_INF_F; - } - } - } // KV halves; query ownership is independent of the KV half. - sample_sync(); - accumulate_sample(stream, math_thread_idx); - sample_sync(); - } - }; - - if constexpr (decltype(scheduler)::kHasPartialBlock) - dispatch_num_block_tokens( - scheduler.get_num_block_tokens(q_block_idx), process_q_block); - else - process_q_block(cute::Int{}); - - sample_sync(); - { - unsigned local = math_thread_idx / 32; - unsigned row = scheduler.get_logits_row(q_block_idx, local); - finish_sample( - stream, buffers, local, row, - min(scheduler.cu_seq_len_k_end[row], scheduler.num_kv_tokens) - - min(scheduler.cu_seq_len_k_start[row], scheduler.num_kv_tokens), - math_thread_idx % 32); - } - sample_sync(); - cutlass::arch::fence_view_async_shared(); - smem.empty_q_barriers[q_stage_idx].arrive(); - } - - cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); - if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); - } -} - -} // namespace seed_sample diff --git a/csrc/fused_mqa_topk/scan.cuh b/csrc/fused_mqa_topk/scan.cuh deleted file mode 100644 index 42dfdfed5b0b..000000000000 --- a/csrc/fused_mqa_topk/scan.cuh +++ /dev/null @@ -1,1106 +0,0 @@ -// SPDX-License-Identifier: MIT -// Experimental derivative of DeepGEMM 2.8 sm100_mqa_logits.cuh. -// Persistent DeepGEMM FP8 MQA producer with fixed-threshold candidate emission. -#pragma once - -#include -#include "handoff.cuh" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "candidates.cuh" - -// Shared SM100 MQA logits core plus contiguous-KV and paged entries -// Both entries use the same q / sf_q / kv / sf_kv / weights TMA signature - -namespace fused_mqa_scan { -using namespace deep_gemm; - -// Convert runtime valid-token count to `cute::Int` so token loops stay -// compile-time constant -template -CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, - Fn&& fn) { - if constexpr (kCandidate <= 1) { - fn(cute::Int<1>{}); - } else if (num_block_tokens >= kCandidate) { - fn(cute::Int{}); - } else { - dispatch_num_block_tokens(num_block_tokens, - static_cast(fn)); - } -} - -// Load heads using power-of-two TMEM chunks. -template -CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, - float* accum) { - constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; - if constexpr (kNumRemaining > 0) { - constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 - : kNumRemaining >= 32 ? 32 - : kNumRemaining >= 16 ? 16 - : kNumRemaining >= 8 ? 8 - : 4; - ptx::tmem_load_32dp32b( - tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); - load_tmem_heads_decomposed(tmem_col, accum); - } -} - -// Warp 0 finds the descending radix bucket. Each lane sums eight adjacent -// bins, the warp locates the group containing rank, and one lane scans only -// that group's eight bins. The selected bin and preceding count are shared -// through ctl after the caller's named barrier. -CUTLASS_DEVICE void select_hist_bucket(int32_t* hist, int32_t* ctl, - uint32_t tid) { - if (tid >= 32) return; - constexpr uint32_t FULL = 0xffffffffu; - int group_count = 0; -#pragma unroll - for (int j = 0; j < 8; ++j) - group_count += hist[255 - static_cast(tid) * 8 - j]; - int inclusive = group_count; -#pragma unroll - for (int shift = 1; shift < 32; shift <<= 1) { - const int prev = __shfl_up_sync(FULL, inclusive, shift); - if (tid >= static_cast(shift)) inclusive += prev; - } - const int before = inclusive - group_count; - const int rank = ctl[0]; - const uint32_t crossing = - __ballot_sync(FULL, before < rank && inclusive >= rank); - const int chosen_lane = __ffs(crossing) - 1; - int chosen = 0; - int above = before; - if (static_cast(tid) == chosen_lane) { -#pragma unroll - for (int j = 0; j < 8; ++j) { - const int bin = 255 - static_cast(tid) * 8 - j; - if (above + hist[bin] >= rank) { - chosen = bin; - break; - } - above += hist[bin]; - } - } - chosen = __shfl_sync(FULL, chosen, chosen_lane); - above = __shfl_sync(FULL, above, chosen_lane); - if (tid == 0) { - ctl[0] = rank - above; - ctl[1] = (ctl[1] << 8) | chosen; - } -} - -// One CTA owns all KV tiles of a query block. At a flush boundary, its math -// warps compact one row's global candidate buffer in place through the now -// vacant shared ring. This keeps the live set bounded independently of ISL. -template -CUTLASS_DEVICE void compact_candidates(uint16_t* values, - int32_t* packed_indices, - int32_t* count_ptr, uint32_t* gate_ptr, - int32_t* hist, uint64_t* scratch, - uint32_t tid, uint32_t cap) { - constexpr int K = 2048; - auto sync = [] { cutlass::arch::NamedBarrier(MATH_THREADS, 1).sync(); }; - int32_t* ctl = hist + 256; - const int count = *count_ptr; - if (tid == 0) { - ctl[0] = K; - ctl[1] = 0; - ctl[2] = 0; - ctl[3] = 0; - } - sync(); - for (int shift = 16; shift >= 8; shift -= 8) { - if (tid < 256) hist[tid] = 0; - sync(); - const uint32_t prefix = static_cast(ctl[1]); - for (int j = tid; j < count; j += MATH_THREADS) { - const uint32_t code = fused_candidates::candidate_load_score_code( - values[j], packed_indices[j]); - if (shift == 16 || (code >> (shift + 8)) == prefix) - atomicAdd(hist + ((code >> shift) & 255u), 1); - } - sync(); - select_hist_bucket(hist, ctl, tid); - sync(); - } - const uint32_t prefix16 = static_cast(ctl[1]); - const int coarse_kept = K - ctl[0] + hist[prefix16 & 255u]; - if (coarse_kept <= 3072) { - const uint32_t coarse_floor = prefix16 << 8; - for (int j = tid; j < count; j += MATH_THREADS) { - const uint32_t code = fused_candidates::candidate_load_score_code( - values[j], packed_indices[j]); - if (code >= coarse_floor) { - const int pos = atomicAdd(ctl + 2, 1); - scratch[pos] = - (static_cast(static_cast(packed_indices[j])) - << 16) | - static_cast(values[j]); - } - } - sync(); - const int kept = ctl[2]; - for (int j = tid; j < kept; j += MATH_THREADS) { - const uint64_t record = scratch[j]; - values[j] = static_cast(record); - packed_indices[j] = static_cast(record >> 16); - } - sync(); - if (tid == 0) { - *count_ptr = kept; - *gate_ptr = coarse_floor == 0 ? 0 : coarse_floor - 1; - } - sync(); - return; - } - // Unusually large score ties in one high-16 bucket use the exact pass. - if (tid < 256) hist[tid] = 0; - sync(); - for (int j = tid; j < count; j += MATH_THREADS) { - const uint32_t code = fused_candidates::candidate_load_score_code( - values[j], packed_indices[j]); - if ((code >> 8) == prefix16) atomicAdd(hist + (code & 255u), 1); - } - sync(); - select_hist_bucket(hist, ctl, tid); - sync(); - const uint32_t kth = static_cast(ctl[1]); - for (int j = tid; j < count; j += MATH_THREADS) { - const uint32_t code = fused_candidates::candidate_load_score_code( - values[j], packed_indices[j]); - if (code > kth) { - const int pos = atomicAdd(ctl + 2, 1); - if (pos < K) - scratch[pos] = - (static_cast(static_cast(packed_indices[j])) - << 16) | - static_cast(values[j]); - } - } - sync(); - const int greater = ctl[2]; - for (int j = tid; j < count; j += MATH_THREADS) { - const uint32_t code = fused_candidates::candidate_load_score_code( - values[j], packed_indices[j]); - if (code == kth) { - const int pos = atomicAdd(ctl + 3, 1); - if (greater + pos < K) - scratch[greater + pos] = - (static_cast(static_cast(packed_indices[j])) - << 16) | - static_cast(values[j]); - } - } - sync(); - for (int j = tid; j < K; j += MATH_THREADS) { - const uint64_t record = scratch[j]; - values[j] = static_cast(record); - packed_indices[j] = static_cast(record >> 16); - } - sync(); - if (tid == 0) { - *count_ptr = K; - *gate_ptr = kth; - } - sync(); -} - -CUTLASS_DEVICE float candidate_gate_float(const uint32_t code) { - if (code == 0) return -cute::numeric_limits::infinity(); - const uint32_t ordered = code << 8; - const uint32_t bits = - (ordered & 0x80000000u) ? (ordered ^ 0x80000000u) : ~ordered; - return __uint_as_float(bits); -} - -// Shared device core parameterized by dtype and scheduler geometry/addressing -template < - uint32_t kNumHeads, uint32_t kHeadDim, bool kIsMXSF, - bool kIsCompressedLogits, bool kCleanLogits, uint32_t BLOCK_Q, - uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages, - uint32_t kNumSMs, uint32_t kNumSpecializedThreads, uint32_t kNumMathThreads, - typename qk_dtype_t, typename logits_dtype_t, typename reduce_dtype_t, - typename MakeScheduler, uint32_t kNumMathWarpGroups = kNumMathThreads / 128, - uint32_t kHandoffInterval = 128, bool kSmemCand = false> -CUTLASS_DEVICE void sm100_mqa_logits_core_impl( - const uint32_t num_q_tokens, const uint32_t logits_stride, - logits_dtype_t* logits, const float* origin, const float* inv_delta, - const int32_t* th_bucket, const int32_t* seed_hist, uint16_t* cand_val, - int32_t* cand_idx, int32_t* cand_cnt, const uint32_t cand_cap, - const int64_t cand_row_delta, const cute::TmaDescriptor& tensor_map_q, - const cute::TmaDescriptor& tensor_map_sf_q, - const cute::TmaDescriptor& tensor_map_kv, - const cute::TmaDescriptor& tensor_map_sf_kv, - const cute::TmaDescriptor& tensor_map_weights, - const MakeScheduler& make_scheduler) { - constexpr bool kIsFP4 = cute::is_same_v; - - const auto sm_idx = blockIdx.x; - const auto warp_idx = cutlass::canonical_warp_idx_sync(); - const auto lane_idx = ptx::get_lane_idx(); - constexpr uint32_t kSpecWarpStart = kNumMathWarpGroups * 4; - - const auto kNegInf = -cute::numeric_limits::infinity(); - - if (warp_idx == kSpecWarpStart) { - cute::prefetch_tma_descriptor(&tensor_map_q); - cute::prefetch_tma_descriptor(&tensor_map_sf_q); - cute::prefetch_tma_descriptor(&tensor_map_weights); - cute::prefetch_tma_descriptor(&tensor_map_kv); - cute::prefetch_tma_descriptor(&tensor_map_sf_kv); - } - - static constexpr uint32_t kNumTmemStages = 2; - static_assert(BLOCK_Q == 8 && kNumMathThreads == 256 && !kIsMXSF); - static constexpr uint32_t kNumUTCCPAlignedElems = 128; - static constexpr uint32_t UMMA_M = 128; - static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; - static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); - static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; - static constexpr uint32_t kNumSFQ = - kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kNumSFKV = - kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kNumQKBytesPerToken = - kIsFP4 ? (kHeadDim / 2) : kHeadDim; - static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = - BLOCK_QH * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = - SPLIT_KV * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = - kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; - static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = - kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); - static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( - kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); - static constexpr uint32_t kNumWeightElementsPerRow = - kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); - static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = - BLOCK_Q * kNumWeightBytesPerRow; - - DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, - "Invalid threads"); - DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and - SPLIT_KV % kNumUTCCPAlignedElems == 0, - "Invalid `SPLIT_KV`"); - DG_STATIC_ASSERT(not(kIsCompressedLogits and kCleanLogits), - "Compressed logits cannot be cleaned in-kernel"); - - using SharedStorage = layout::MQALogitsSharedStorage< - kNumHeads, kHeadDim, kIsMXSF, BLOCK_Q, SPLIT_KV, kNumQStages, - kNumKVStages, kNumTmemStages, qk_dtype_t, reduce_dtype_t>; - extern __shared__ __align__(SharedStorage::kSwizzleAlignment) - uint8_t smem_buffer[]; - auto& smem = *reinterpret_cast(smem_buffer); - constexpr uint32_t kRingSlots = 8; - constexpr uint32_t kRecordBlockMask = 255; - auto& handoff = *reinterpret_cast*>(smem_buffer + - sizeof(SharedStorage)); - handoff.init(threadIdx.x); - - constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; - constexpr uint32_t kNumTmemCols = - utils::get_num_aligned_tmem_cols(); - constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; - constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; - DG_STATIC_ASSERT(kNumTmemCols <= 512, "Too many tensor memory"); - - if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { -#pragma unroll - for (uint32_t i = 0; i < kNumQStages; ++i) { - smem.full_q_barriers[i].init(1); - smem.empty_q_barriers[i].init(kNumMathThreads + 32); - } -#pragma unroll - for (uint32_t i = 0; i < kNumKVStages; ++i) { - smem.full_kv_barriers[i].init(1); - smem.empty_kv_barriers[i].init(kIsMXSF ? 1 : kNumMathThreads); - } -#pragma unroll - for (uint32_t i = 0; i < kNumTmemStages; ++i) { - smem.full_tmem_barriers[i].init(1); - smem.empty_tmem_barriers[i].init(256); - } - cutlass::arch::fence_barrier_init(); - } - __syncwarp(); - - if (warp_idx == kSpecWarpStart + 2) - cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); - __syncthreads(); - - RingPipeline q_pipeline; - RingPipeline kv_pipeline; - RingPipeline tmem_pipeline; - - constexpr uint32_t kNumSpecializedRegisters = 56; - constexpr uint32_t kNumMathRegisters = 224; - - cudaGridDependencySynchronize(); - - // Short contexts (handoff interval <= 8, i.e. <= 16K keys): candidates are - // dense and the drain was the bottleneck, so math warps emit directly (warp - // ballot + one shared atomic per row and KV half). Longer contexts keep the - // ring + drain warps. kHandoffInterval == 0: resident mode (keys <= - // logits_stride). Every score of the block's 8 rows is stored as an fp24 code - // in shared memory (hi16 array, then lo8 array, row-major with stride - // logits_stride) at `logits`; no gate, no candidate emission. - constexpr bool kResident = kHandoffInterval == 0; - constexpr bool kDirectEmit = !kResident && kHandoffInterval <= 8; - // Two drain warps: warp +3 drains math warpgroup 0 (rows 0-3) and warp +0, - // idle after its single Q load, drains warpgroup 1 (rows 4-7). Rows are - // disjoint. - const auto drain_producers = [&](const unsigned p0) { - auto scheduler = make_scheduler(sm_idx); - uint32_t q_block_idx, kv_base, num_kv_splits; - unsigned drain_epoch = 0; - while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { - const unsigned row_base = scheduler.get_logits_row(q_block_idx, 0); - for (unsigned base = 0; base < num_kv_splits; - base += kHandoffInterval, ++drain_epoch) { - for (unsigned p = p0; p < p0 + 4; ++p) { - handoff.drain( - p, drain_epoch, 4, row_base + (p / 4) * 4, cand_cnt, cand_cap, - [&](unsigned row, unsigned slot, unsigned record, - unsigned batch_base, unsigned source) { - const uint32_t physical = - kv_base + (batch_base * 2 + (record & 255u)) * 128 + - (source % 128); - const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - cand_cap + - slot; - fused_candidates::store_candidate_record( - cand_val + off, cand_idx + off, record >> 8, physical); - }); - } - } - } - }; - - if (warp_idx == kSpecWarpStart) { - cutlass::arch::warpgroup_reg_dealloc(); - if (cute::elect_one_sync()) { - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - tma::copy( - &tensor_map_q, &smem.full_q_barriers[q_stage_idx], - smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); - if constexpr (kIsMXSF) - tma::copy( - &tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], - smem.smem_sf_q[q_stage_idx], 0, q_token_base); - 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_SF_Q_SIZE_PER_STAGE + - SMEM_WEIGHT_SIZE_PER_STAGE); - } - } - __syncwarp(); - if constexpr (!kDirectEmit && !kResident) drain_producers(4u); - } else if (warp_idx == kSpecWarpStart + 1) { - cutlass::arch::warpgroup_reg_dealloc(); - - auto scheduler = make_scheduler(sm_idx); - uint32_t cached_kv_page_base = 0; - uint32_t cached_kv_page_coord = 0; - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { - cached_kv_page_base = cute::numeric_limits::max(); -#pragma unroll 1 - for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; - ++kv_split_idx) { - if constexpr (decltype(scheduler)::kIsPaged) { - constexpr uint32_t kPageKV = decltype(scheduler)::kPageKV; - constexpr uint32_t kNumPagesPerSplit = - decltype(scheduler)::kNumPagesPerSplit; - DG_STATIC_ASSERT(kNumPagesPerSplit <= 32, - "Split spans more pages than a warp can cache"); - - const uint32_t kv_page_base = - (kv_base + kv_split_idx) * kNumPagesPerSplit; - if (kv_page_base < cached_kv_page_base or - kv_page_base + kNumPagesPerSplit > cached_kv_page_base + 32) { - cached_kv_page_base = (kv_page_base / 32) * 32; - cached_kv_page_coord = scheduler.get_kv_page_coord_by_page_offset( - cached_kv_page_base + lane_idx); - } - - CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); - if (cute::elect_one_sync()) - smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); - __syncwarp(); - - int page_coords[kNumPagesPerSplit]; -#pragma unroll - for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; - ++page_idx) { - const auto src_lane = - static_cast(kv_page_base - cached_kv_page_base + page_idx); - page_coords[page_idx] = - __shfl_sync(0xffffffff, cached_kv_page_coord, src_lane); - } - - if (cute::elect_one_sync()) { -#pragma unroll - for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; - ++page_idx) { - tma::copy( - &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_kv[kv_stage_idx] + - page_idx * kPageKV * kNumQKBytesPerToken, - 0, 0, 1, page_coords[page_idx]); - tma::copy( - &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_sf_kv[kv_stage_idx] + page_idx * kPageKV, 0, - page_coords[page_idx]); - } - 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 (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); - tma::copy( - &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], - smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); - 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 utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { - DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, - "Invalid aligned elements"); - uint32_t values[4]; -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) - values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); - __syncwarp(); - ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], - values[3]); - }; - - auto sf_desc = mma::sm100::make_sf_desc(nullptr); - - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - - if constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; - utccp_required_smem_warp_transpose(smem_ptr); - } - cutlass::arch::fence_view_async_shared(); -#pragma unroll - for (uint32_t i = 0; i < kNumSFQ / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_q[q_stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - if (cute::elect_one_sync()) - cute::SM100_UTCCP_4x32dp128bit_1cta::copy( - sf_desc, kTmemStartColOfSFQ + i * 4); - __syncwarp(); - } - } - - 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 constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; - utccp_required_smem_warp_transpose(smem_ptr); - } - cutlass::arch::fence_view_async_shared(); - } - - if (cute::elect_one_sync()) { - if constexpr (kIsMXSF) { -#pragma unroll - for (uint32_t i = 0; i < kNumSFKV / kNumUTCCPAlignedElems; ++i) { - auto smem_ptr = - smem.smem_sf_kv[kv_stage_idx] + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute::SM100_UTCCP_4x32dp128bit_1cta::copy( - sf_desc, kTmemStartColOfSFKV + i * 4); - } - } -#pragma unroll - for (uint32_t i = 0; i < kNumMathWarpGroups; ++i) { - CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); - 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(); - - if constexpr (kIsMXSF) { - DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or - kHeadDim == 64 or kHeadDim == 128, - "Invalid head dim"); - - constexpr uint32_t kPackFactor = - get_smem_pack_factor(); - constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; - - using mma_op_t = - cute::conditional_t; - auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< - qk_dtype_t, qk_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, - UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K>(); -#pragma unroll - for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { - auto runtime_instr_desc = - mma::sm100::make_runtime_instr_desc_with_sf_id( - instr_desc, k * kPackFactor, k * kPackFactor); - 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); - mma_op_t::fma(a_desc, b_desc, tmem_addr, k, runtime_instr_desc, - kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ); - } - } else { - 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); -#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(); - if constexpr (kIsMXSF) - cutlass::arch::umma_arrive(reinterpret_cast( - &smem.empty_kv_barriers[kv_stage_idx])); - } - smem.empty_q_barriers[q_stage_idx].arrive(); - } - } else if (warp_idx == kSpecWarpStart + 3) { - cutlass::arch::warpgroup_reg_dealloc(); - if constexpr (!kDirectEmit && !kResident) drain_producers(0u); - } else if (warp_idx < kSpecWarpStart) { - cutlass::arch::warpgroup_reg_alloc(); - - constexpr bool kLoadSeqBounds = true; - uint32_t seq_k_start[4]; - uint32_t seq_k_end[4]; - const auto math_warpgroup_idx = warp_idx / 4; - const auto math_thread_idx = warp_idx * 32 + lane_idx; - DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, - "Math warp groups exceed TMEM stages"); - const unsigned query_base = math_warpgroup_idx * 4; - - constexpr bool kIsReduceBF16 = not cute::is_same_v; - DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); - DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); - using weights_reg_dtype_t = - cute::conditional_t; - constexpr uint32_t kNumWeightsRegPerToken = - kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; - weights_reg_dtype_t weights[4][kNumWeightsRegPerToken]; - float accum[kNumHeads]; - - unsigned produce_epoch = 0; - auto scheduler = make_scheduler(sm_idx); - // NOTES: split index for paged scheduler, token offset for contiguous-KV - // scheduler. - uint32_t q_block_idx, kv_base, num_kv_splits; - 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); - - const auto process_q_block = [&](auto num_valid_tokens_t) { - constexpr uint32_t kNumValidTokens = 4; - uint32_t lane_counts[4] = {}; - uint32_t flush_base = 0; - float gate_score[4]; - -#pragma unroll - for (uint32_t i = 0; i < kNumValidTokens; ++i) { - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - // Use the raw-score threshold estimated from the distributed sample. - // This threshold remains fixed for the complete scan. - gate_score[i] = kResident ? 0.0f : origin[row]; - seq_k_start[i] = cute::min(scheduler.cu_seq_len_k_start[row], - scheduler.num_kv_tokens); - seq_k_end[i] = cute::min(scheduler.cu_seq_len_k_end[row], - scheduler.num_kv_tokens); - const auto smem_weights_row = - smem.smem_weights[q_stage_idx] + - (query_base + i) * kNumWeightElementsPerRow; - if constexpr (kIsReduceBF16) { - // Load two bf16 weights at a time as one packed shared u32 - const auto packed_row = - reinterpret_cast(smem_weights_row); -#pragma unroll - for (uint32_t j = 0; j < kNumHeads / 2; ++j) { - const auto packed = ptx::ld_shared(packed_row + j); - weights[i][j] = ptx::exchange( - *reinterpret_cast(&packed), 0); - } - } else { - bool unsafe_weights = false; -#pragma unroll - for (uint32_t j = 0; j < kNumHeads; ++j) { - const float weight = ptx::ld_shared(smem_weights_row + j); - unsafe_weights |= weight != 0.0f && fabsf(weight) < 0x1p-125f; - // Resident: full weights, the 0.5 ReLU factor is applied per - // score. - weights[i][j] = kResident ? weight : weight * 0.5f; - } - // Hoist the ReLU factor out of the per-key reduction. - // Tiny weights can round differently when halved; use - // the existing GPU recovery with unscaled weights. - if (!kResident && unsafe_weights) { - gate_score[i] = __int_as_float(0x7f800000); - if (math_thread_idx % 128 == 0) - atomicAdd(cand_cnt + row, static_cast(cand_cap) + 1); - } - } - } - - __shared__ unsigned row_count[BLOCK_Q]; - // Shared candidates, per row [seg0][seg1][seg2][seg3][pool]: each of - // the row's 4 warps owns a static segment (no atomics; its last slot is - // a trash slot for lanes without a candidate); a warp past its segment - // takes slots from the row's shared pool (one atomic per overflowing - // ballot); past the pool, global slots row*cand_cap + j (count at - // cand_cap - 1). Unfilled slots get the lowest code. - const uint32_t pool_size = logits_stride / 6, - seg_size = (logits_stride - pool_size) / 4; - const uint32_t seg_warp = warp_idx % 4; - // Each array holds the 8 rows plus 256 per-lane trash slots (8 warps x - // 32 lanes). - const uint32_t cand_ext = BLOCK_Q * logits_stride + 256; - __shared__ unsigned pool_cnt[BLOCK_Q]; - uint32_t seg_cnt[kNumValidTokens] = {}; - if constexpr (kDirectEmit) { - if (math_thread_idx < BLOCK_Q) { - row_count[math_thread_idx] = 0; - pool_cnt[math_thread_idx] = 0; - } - cutlass::arch::NamedBarrier(kNumMathThreads, 3).sync(); - } - for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; - ++kv_split_idx) { - if constexpr (!kDirectEmit && !kResident) - if ((kv_split_idx & (kHandoffInterval - 1)) == 0) - handoff.acquire(warp_idx, produce_epoch); - CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); - smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); -#pragma unroll - for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { - const unsigned kv_lane = (math_thread_idx % 128) + kv_half * 128; - const unsigned 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(); - } - - static_assert(!kIsReduceBF16, "FP32-only pair reduction"); - float accum_pair[2][kNumHeads]; -#pragma unroll - for (uint32_t pair = 0; pair < kNumValidTokens; pair += 2) { - // Load only 2 query rows at once to limit - // live accumulator registers during score reduction. - const uint32_t tmem_addr = - tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads; - ptx::tmem_load_32dp32b<2 * kNumHeads>( - tmem_addr, reinterpret_cast(accum_pair)); - cutlass::arch::fence_view_async_tmem_load(); - if (pair + 2 >= kNumValidTokens) { - ptx::tcgen05_before_thread_sync(); - smem.empty_tmem_barriers[tmem_stage_idx].arrive(); - } - - float2 sum_0[2] = {}; - float2 sum_1[2] = {}; -#pragma unroll - for (uint32_t j = 0; j < kNumHeads; j += 4) { -#pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - if (pair + r >= kNumValidTokens) continue; - const uint32_t i = pair + r; - auto* a = accum_pair[r]; - auto transform = [&](uint32_t at, float2 sum) { - const auto x = make_float2(a[at], a[at + 1]); - const auto abs_x = - make_float2(fabsf(a[at]), fabsf(a[at + 1])); - const auto w = - make_float2(weights[i][at], weights[i][at + 1]); - return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); - }; - sum_0[r] = transform(j, sum_0[r]); - sum_1[r] = transform(j + 2, sum_1[r]); - } - } -#pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - if (pair + r >= kNumValidTokens) continue; - const uint32_t i = pair + r; - const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - const float reduced = sum.x + sum.y; - const float bq = - reduced * (kResident ? 0.5f * static_cast(scale_kv) - : static_cast(scale_kv)); - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - if constexpr (kResident) { - // Select reads only [start, end) of each row, so no masking - // here. - if (kv_offset < logits_stride) { - const uint32_t code = - fused_candidates::candidate_fp24_code(bq); - const uint32_t at = - (query_base + i) * logits_stride + kv_offset; - reinterpret_cast(logits)[at] = - static_cast(code >> 8); - reinterpret_cast( - logits)[BLOCK_Q * logits_stride * 2 + at] = - static_cast(code); - } - continue; - } - // NaN fails the compare; +inf passes (harmless, it would win - // the top-k anyway). - const bool pass = (kSmemCand ? kv_offset - seq_k_start[i] < - seq_k_end[i] - seq_k_start[i] - : kv_offset >= seq_k_start[i] && - kv_offset < seq_k_end[i] && - isfinite(bq)) && - bq > gate_score[i]; - if constexpr (kSmemCand) { - // Branch-free: every lane stores; lanes without a candidate - // (or past the segment) write their own trash slot after the - // rows (never read). - const unsigned pass_mask = __ballot_sync(0xffffffffu, pass); - const unsigned slot = - seg_cnt[i] + __popc(pass_mask & ((1u << lane_idx) - 1u)); - seg_cnt[i] += __popc(pass_mask); - const uint32_t code = - fused_candidates::candidate_fp24_code(bq); - const uint32_t seg_cap = seg_size; - const uint32_t at = - pass && slot < seg_cap - ? (query_base + i) * logits_stride + - seg_warp * seg_size + slot - : BLOCK_Q * logits_stride + warp_idx * 32 + lane_idx; - auto* cbase = reinterpret_cast(logits); - reinterpret_cast(cbase)[at] = - static_cast(code >> 8); - cbase[cand_ext * 2 + at] = static_cast(code); - reinterpret_cast(cbase + cand_ext * 3)[at] = - static_cast(kv_offset); - if (seg_cnt[i] > seg_cap) { - // Warp-uniform: lanes past the segment take row-pool slots. - const bool over = pass && slot >= seg_cap; - const unsigned over_mask = __ballot_sync(0xffffffffu, over); - if (over_mask) { - unsigned pool_base = 0; - if (lane_idx == 0) - pool_base = atomicAdd(&pool_cnt[query_base + i], - __popc(over_mask)); - pool_base = __shfl_sync(0xffffffffu, pool_base, 0); - const unsigned p = - pool_base + - __popc(over_mask & ((1u << lane_idx) - 1u)); - if (over && p < pool_size) { - const uint32_t pat = - (query_base + i) * logits_stride + 4 * seg_size + p; - reinterpret_cast(cbase)[pat] = - static_cast(code >> 8); - cbase[cand_ext * 2 + pat] = static_cast(code); - reinterpret_cast(cbase + cand_ext * 3)[pat] = - static_cast(kv_offset); - } else if (over && p - pool_size < cand_cap - 1) { - // Past the pool (rare): global, val = hi16 code, idx = - // lo8 code << 16 | key. - const uint64_t off = - static_cast(int64_t(row) + - cand_row_delta) * - cand_cap + - (p - pool_size); - __stcg(cand_val + off, - static_cast(code >> 8)); - __stcg(cand_idx + off, - static_cast((code & 0xffu) << 16 | - kv_offset)); - } - } - } - } else if constexpr (kDirectEmit) { - const unsigned pass_mask = __ballot_sync(0xffffffffu, pass); - if (pass_mask) { - // Ballot + one shared atomic per row: global candidate - // slots. - const unsigned leader = __ffs(pass_mask) - 1; - unsigned base = 0; - if (lane_idx == leader) - base = atomicAdd(&row_count[query_base + i], - __popc(pass_mask)); - base = __shfl_sync(0xffffffffu, base, leader); - const unsigned slot = - base + __popc(pass_mask & ((1u << lane_idx) - 1u)); - if (pass && slot < cand_cap) { - const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - cand_cap + - slot; - fused_candidates::store_candidate( - cand_val + off, cand_idx + off, bq, kv_offset); - } - } - } else if (pass) { - const uint32_t count = lane_counts[i]; - if (count < kRingSlots) { - const uint32_t pos = - ((warp_idx * BLOCK_Q + i) * kRingSlots + count) * 32u + - lane_idx; - const uint32_t code = - fused_candidates::candidate_fp24_code(bq); - handoff.records[produce_epoch & 1][warp_idx][i][count] - [lane_idx] = - (code << 8) | - ((kv_split_idx - flush_base) * 2 + kv_half); - lane_counts[i] = count + 1; - } else { - const int slot = atomicAdd(cand_cnt + row, 1); - if (slot < static_cast(cand_cap)) { - const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - cand_cap + - slot; - fused_candidates::store_candidate( - cand_val + off, cand_idx + off, bq, kv_offset); - } - } - } - } - } - } // KV halves; query ownership is independent of the KV half. - if (!kDirectEmit && !kResident && - (((kv_split_idx + 1) & (kHandoffInterval - 1)) == 0 || - kv_split_idx + 1 == num_kv_splits)) { -#pragma unroll - for (unsigned i = 0; i < kNumValidTokens; ++i) { - handoff.counts[produce_epoch & 1][warp_idx][i][lane_idx] = - lane_counts[i]; - lane_counts[i] = 0; - } - handoff.publish(warp_idx, produce_epoch, flush_base); - ++produce_epoch; - flush_base = kv_split_idx + 1; - } - } - if constexpr (kSmemCand) { -// Seal each warp's segment tail with the lowest code; publish counts and -// spills. -#pragma unroll - for (uint32_t i = 0; i < kNumValidTokens; ++i) { - const uint32_t row = - scheduler.get_logits_row(q_block_idx, query_base + i); - auto* base = reinterpret_cast(logits); - for (uint32_t slot = min(seg_cnt[i], seg_size) + lane_idx; - slot < seg_size; slot += 32) { - const uint32_t at = - (query_base + i) * logits_stride + seg_warp * seg_size + slot; - reinterpret_cast(base)[at] = 0; - base[cand_ext * 2 + at] = 0; - reinterpret_cast(base + cand_ext * 3)[at] = 0; - } - if (lane_idx == 0) - atomicAdd(&row_count[query_base + i], seg_cnt[i]); - } - // After every warp's pool claims: seal each row's unused pool slots - // and record the global spill count (bit 21 of the published count - // marks a spilled row). - cutlass::arch::NamedBarrier(kNumMathThreads, 3).sync(); - for (uint32_t q = 0; q < BLOCK_Q; ++q) { - auto* base = reinterpret_cast(logits); - for (uint32_t p = pool_cnt[q] + math_thread_idx; p < pool_size; - p += kNumMathThreads) { - const uint32_t at = q * logits_stride + 4 * seg_size + p; - reinterpret_cast(base)[at] = 0; - base[cand_ext * 2 + at] = 0; - reinterpret_cast(base + cand_ext * 3)[at] = 0; - } - } - if (math_thread_idx < BLOCK_Q) { - const uint32_t row = - scheduler.get_logits_row(q_block_idx, math_thread_idx); - const uint32_t spilled = - pool_cnt[math_thread_idx] > pool_size - ? min(pool_cnt[math_thread_idx] - pool_size, cand_cap - 1) - : 0u; - __stcg(cand_idx + - static_cast(int64_t(row) + cand_row_delta) * - cand_cap + - cand_cap - 1, - static_cast(spilled)); - if (spilled) atomicOr(&row_count[math_thread_idx], 1u << 21); - } - } - if constexpr (kDirectEmit) { - // Publish row totals (added to any unsafe-weights recovery marker - // already in cand_cnt). Shared-candidate rows with < 2048 real - // candidates are pushed past 16384 (recover). - cutlass::arch::NamedBarrier(kNumMathThreads, 3).sync(); - if (math_thread_idx < BLOCK_Q) { - // Shared candidates: bits 0-20 = real count, bit 21 = some warp - // spilled. - int total = static_cast(row_count[math_thread_idx]); - atomicAdd(cand_cnt + scheduler.get_logits_row(q_block_idx, - math_thread_idx), - total); - } - } - }; - - if constexpr (decltype(scheduler)::kHasPartialBlock) - dispatch_num_block_tokens( - scheduler.get_num_block_tokens(q_block_idx), process_q_block); - else - process_q_block(cute::Int{}); - - cutlass::arch::fence_view_async_shared(); - smem.empty_q_barriers[q_stage_idx].arrive(); - } - - cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); - if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); - } -} - -// Unified contiguous-KV entry for FP8 / MXFP4 / MXFP8. -template -CUTLASS_GLOBAL __launch_bounds__( - kNumSpecializedThreads + kNumMathThreads, - 1) void sm100_mqa_logits(const uint32_t num_q_tokens, - const uint32_t num_kv_tokens, - const uint32_t logits_stride, - const uint32_t* cu_seq_len_k_start, - const uint32_t* cu_seq_len_k_end, - const uint32_t* schedule_meta, - logits_dtype_t* logits, const float* origin, - const float* inv_delta, const int32_t* th_bucket, - const int32_t* seed_hist, uint16_t* cand_val, - int32_t* cand_idx, int32_t* cand_cnt, - const uint32_t cand_cap, - const __grid_constant__ cute::TmaDescriptor - tensor_map_q, - const __grid_constant__ cute::TmaDescriptor - tensor_map_sf_q, - const __grid_constant__ cute::TmaDescriptor - tensor_map_kv, - const __grid_constant__ cute::TmaDescriptor - tensor_map_sf_kv, - const __grid_constant__ cute::TmaDescriptor - tensor_map_weights) { - const auto make_scheduler = [&](const uint32_t& sm_idx) { - return sched::SM100MQALogitsScheduler( - sm_idx, num_q_tokens, num_kv_tokens, cu_seq_len_k_start, - cu_seq_len_k_end, schedule_meta); - }; - - sm100_mqa_logits_core_impl< - kNumHeads, kHeadDim, kIsMXSF, kIsCompressedLogits, kCleanLogits, BLOCK_Q, - SPLIT_KV, kNumQStages, kNumKVStages, kNumSMs, kNumSpecializedThreads, - kNumMathThreads, qk_dtype_t, logits_dtype_t, reduce_dtype_t, - decltype(make_scheduler), kNumMathWarpGroups, kHandoffInterval>( - num_q_tokens, logits_stride, logits, origin, inv_delta, th_bucket, - seed_hist, cand_val, cand_idx, cand_cnt, cand_cap, 0, tensor_map_q, - tensor_map_sf_q, tensor_map_kv, tensor_map_sf_kv, tensor_map_weights, - make_scheduler); -} - -} // namespace fused_mqa_scan diff --git a/csrc/fused_mqa_topk/upstream_select.cuh b/csrc/fused_mqa_topk/upstream_select.cuh deleted file mode 100644 index 97b424c63be2..000000000000 --- a/csrc/fused_mqa_topk/upstream_select.cuh +++ /dev/null @@ -1,534 +0,0 @@ -// Isolated adaptation of vLLM sampler.cu device selection routines. -#pragma once -#include -#include -#include -namespace isolated_upstream { -template -static inline __device__ uint32_t extractBinIdx(float x) { - if constexpr (step == 0) { - __half hx = __float2half(x); - uint16_t bits = __half_as_ushort(hx); - bits = (bits & 0x8000) ? bits : ~bits & 0x7fff; - return bits >> 5; - } else { - uint32_t bits = __float_as_uint(x); - bits = (bits & 0x80000000) ? bits : ~bits & 0x7fffffff; - - if constexpr (step == 1) { - return bits >> 21; - } else if constexpr (step == 2) { - return (bits >> 10) & 0x7ff; - } else if constexpr (step == 3) { - return bits & 0x3ff; - } - } -} - -template -static inline __device__ bool isPartialMatch(float x, uint32_t pattern) { - if constexpr (shift == 0) { - return true; - } - uint32_t bits = __float_as_uint(x); - bits = (bits & 0x80000000) ? bits : ~bits & 0x7fffffff; - return (bits ^ pattern) >> shift == 0; -} - -/** - * Map a Func over the input data, using vectorized load instructions if - * possible. - * - * @tparam T element type - * @tparam IdxT indexing type - * @tparam Func void (T x, IdxT idx) - * - * @param thread_rank rank of the calling thread among all participating threads - * @param num_threads number of the threads that participate in processing - * @param in the input data - * @param len the number of elements to read - * @param f the lambda taking two arguments (T x, IdxT idx) - */ -template -__device__ void vectorized_process(size_t thread_rank, size_t num_threads, - const T* in, idxT len, Func f) { - // Use dynamic WARP_SIZE from cuda_compat.h to support both - // Wave64 (MI300X/gfx942) and Wave32 (Strix Halo/gfx1151) architectures - constexpr int kWarpSize = 32; - using WideT = float4; - if constexpr (sizeof(T) >= sizeof(WideT)) { - for (idxT i = thread_rank; i < len; i += num_threads) { - f(in[i], i); - } - } else { - static_assert(sizeof(WideT) % sizeof(T) == 0); - constexpr int items_per_scalar = sizeof(WideT) / sizeof(T); - // TODO: it's UB - union { - WideT scalar; - T array[items_per_scalar]; - } wide; - - int skip_cnt = - (reinterpret_cast(in) % sizeof(WideT)) - ? ((sizeof(WideT) - reinterpret_cast(in) % sizeof(WideT)) / - sizeof(T)) - : 0; - if (skip_cnt > len) { - skip_cnt = len; - } - const WideT* in_cast = reinterpret_cast(in + skip_cnt); - const idxT len_cast = (len - skip_cnt) / items_per_scalar; - - for (idxT i = thread_rank; i < len_cast; i += num_threads) { - wide.scalar = in_cast[i]; - const idxT real_i = skip_cnt + i * items_per_scalar; -#pragma unroll - for (int j = 0; j < items_per_scalar; ++j) { - f(wide.array[j], real_i + j); - } - } - - static_assert(kWarpSize >= items_per_scalar); - // and because items_per_scalar > skip_cnt, kWarpSize > skip_cnt - // no need to use loop - if (thread_rank < skip_cnt) { - f(in[thread_rank], thread_rank); - } - // because len_cast = (len - skip_cnt) / items_per_scalar, - // len_cast * items_per_scalar + items_per_scalar > len - skip_cnt; - // and so - // len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <= - // kWarpSize no need to use loop - const idxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank; - if (remain_i < len) { - f(in[remain_i], remain_i); - } - } -} - -template -__device__ bool processHistogramStep( - const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern, - int& thresholdBinIdx, SmemOutputType& smemOutput, int* smemThresholdBinIdx, - int* smemFinalDstIdx, int* smemFinalBinSize, int* smemFoundTopKValues, - SmemFinalType& smemFinal, int stride1, int rowStart, int topK) { - // Clear the histogram. -#pragma unroll - for (int idx = threadIdx.x; idx < kNumBins; idx += kNumThreadsPerBlock) { - smemFinal.histo.data[idx] = 0; - } - - // Make sure the histogram is ready. - __syncthreads(); - - // Update pattern - constexpr auto patternShift = step < 2 ? 0 : step == 2 ? 21 : 10; - if constexpr (step == 2) { - logitPattern = static_cast(thresholdBinIdx & 0x7ff) - << patternShift; - } else if constexpr (step == 3) { - logitPattern |= static_cast(thresholdBinIdx & 0x7ff) - << patternShift; - } - - int negativeInfinityCount = 0; - auto distributeToBins = [&](float logit, int /* idx */ = 0) { - if (isPartialMatch(logit, logitPattern)) { - uint32_t binIdx = extractBinIdx(logit); - if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) { - ++negativeInfinityCount; - } else { - atomicAdd(&smemFinal.histo.data[binIdx], 1); - } - } - }; - - // Distribute the elements to the histogram bins. - if (stride1 == 1) { - vectorized_process(threadIdx.x, kNumThreadsPerBlock, logits + rowStart, - rowEnd - rowStart, distributeToBins); - } else { - for (int idx = rowStart + threadIdx.x; idx < rowEnd; - idx += kNumThreadsPerBlock) { - float logit = logits[idx * stride1]; - distributeToBins(logit, idx); - } - } - if constexpr (useTopK512Optimization) { - if (negativeInfinityCount > 0) { - const auto bin = extractBinIdx(-INFINITY); - atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount); - } - } - // Make sure the histogram is ready. - __syncthreads(); - - // Reads the value of the starting position in the smemOutput array - int lastValue = smemFoundTopKValues[0]; - - for (int round = 0; round < kNumBins / kNumThreadsPerBlock; round++) { - // Read the values from SMEM. - int idx = threadIdx.x + kNumThreadsPerBlock * round; - int binCount{0}; - binCount = smemFinal.histo.data[idx]; - - // Make sure each thread has read its value. - __syncthreads(); - - // Compute the prefix sum. - int prefixSum{0}, totalSum{0}; - using Scan = cub::BlockScan; - Scan(smemFinal.histo.scan).ExclusiveSum(binCount, prefixSum, totalSum); - - // Update the histogram with the prefix sums. - prefixSum += lastValue; - totalSum += lastValue; - smemFinal.histo.data[idx] = prefixSum; - - // Make sure the data is in shared memory. - __syncthreads(); - - // Find the last valid bin. - bool foundThreshold = false; - if (prefixSum < topK) { - int nextPrefixSum = threadIdx.x == kNumThreadsPerBlock - 1 - ? totalSum - : smemFinal.histo.data[idx + 1]; - - if (nextPrefixSum >= topK) { - smemThresholdBinIdx[0] = idx; - smemFinalBinSize[0] = nextPrefixSum - prefixSum; - foundThreshold = true; - } - } - - // Early exit: if any thread found the threshold, we can skip remaining - // rounds - if (__syncthreads_or(foundThreshold)) { - break; - } - - lastValue = totalSum; - } - - // Make sure the data is in shared memory. - __syncthreads(); - - // The threshold bin. - thresholdBinIdx = smemThresholdBinIdx[0]; - // Resolve a -inf cutoff with exact tie emission, keeping sort padding out. - const bool finalBinFits = - smemFinalBinSize[0] <= kNumFinalItems && - !(useTopK512Optimization && step < 3 && - isPartialMatch(-INFINITY, logitPattern) && - thresholdBinIdx == extractBinIdx(-INFINITY)); - - auto processBins = [&](float logit, int idx) { - if (isPartialMatch(logit, logitPattern)) { - uint32_t binIdx = extractBinIdx(logit); - // Only write elements with binIdx < thresholdBinIdx when: - // 1. This is step 0 and the threshold bin is small enough (no step 1) - // 2. This is step >= 1 (where pattern matching filters correctly) - // This prevents duplicates when step 0 and step 1 both run. - bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1); - if (binIdx < thresholdBinIdx && shouldWriteDirectly) { - // The element is part of the top-k selection - int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); - - if constexpr (mergeBlocks) { - smemOutput[dstIdx] = indices[idx]; - } else if constexpr (multipleBlocksPerRow) { - smemOutput[dstIdx] = idx + rowStart; - reinterpret_cast(smemOutput + topK)[dstIdx] = logit; - } else { - smemOutput[dstIdx] = idx; - } - } - if constexpr (step < 3) { - // Only fill the final items for sorting if the threshold bin fits - if (binIdx == thresholdBinIdx && finalBinFits) { - int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1); - smemFinal.items.logits[dstIdx] = logit; - if constexpr (mergeBlocks) { - smemFinal.items.indices[dstIdx] = indices[idx]; - } else if constexpr (multipleBlocksPerRow) { - smemFinal.items.indices[dstIdx] = idx + rowStart; - } else { - smemFinal.items.indices[dstIdx] = idx; - } - } - } else { - if (binIdx == thresholdBinIdx) { - // The elements in the threshold bin share the same 32 bits at step 3 - int dstIdx = atomicAdd(&smemFinal.histo.data[binIdx], 1); - if (dstIdx < topK) { - if constexpr (mergeBlocks) { - smemOutput[dstIdx] = indices[idx]; - } else if constexpr (multipleBlocksPerRow) { - smemOutput[dstIdx] = idx + rowStart; - reinterpret_cast(smemOutput + topK)[dstIdx] = logit; - } else { - smemOutput[dstIdx] = idx; - } - } - } - } - } - }; - - if (stride1 == 1) { - vectorized_process(threadIdx.x, kNumThreadsPerBlock, logits + rowStart, - rowEnd - rowStart, processBins); - } else { - for (int idx = rowStart + threadIdx.x; idx < rowEnd; - idx += kNumThreadsPerBlock) { - float logit = logits[idx * stride1]; - processBins(logit, idx); - } - } - - // Make sure the elements are in shared memory. - __syncthreads(); - - // Check if we should continue to next step - return !finalBinFits; -} - -// Follows half - 11 - 11 - 10 bit iterations -// Keep the adaptive kernel's device instantiation separate from legacy calls. -template -static __device__ void topKPerRowJob(const int* indices, const float* logits, - int rowStart, int rowEnd, int* outIndices, - float* outLogits, int stride1, int topK) { - static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks); - static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024); - // The number of slots for the final pass. - static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 1536; - // The number of elements per thread for the final sort. - static constexpr int kNumFinalItemsPerThread = - kNumFinalItems / kNumThreadsPerBlock; - // The class to sort the elements during the final pass. - using FinalSort = cub::BlockRadixSort; - using FinalSortTempStorage = - std::conditional_t; - // The class to compute the inclusive prefix-sum over the histogram. - using Scan = cub::BlockScan; - - // The structure to store the final items (for the final pass). - struct FinalItems { - // Shared memory to store the indices for the final pass. - int indices[kNumFinalItems]; - // Shared memory to store the logits for the final pass. - float logits[kNumFinalItems]; - }; - - struct Histogram { - typename Scan::TempStorage scan; - int data[kNumBins]; - }; - - // Shared memory to compute the block sort. - union FinalStorage { - FinalItems items; - FinalSortTempStorage finalSort; - Histogram histo; - }; - extern __shared__ __align__(1024) unsigned char select_workspace[]; - auto& smemFinal = *reinterpret_cast(select_workspace + 8192); - - // Shared memory to store the selected indices. - // If we are processing using multiple blocks, we need to store the logits and - // indices. - extern __shared__ __align__(1024) int32_t smemOutput[]; - - // Shared memory to store the threshold bin. - __shared__ int smemThresholdBinIdx[1]; - // Shared memory counter to register the candidates for the final phase. - __shared__ int smemFinalDstIdx[1]; - // Shared memory to determine if the threshold bin fits in the final items. - __shared__ int smemFinalBinSize[1]; - // Shared memory to keep track of the top-k values found so far by the - // previous iterations - __shared__ int smemFoundTopKValues[1]; - - // The length of the row. - int rowLen = rowEnd - rowStart; - - // Shortcut if the length of the row is smaller than Top-K. Indices are not - // sorted by their corresponding logit. - if (rowLen <= topK) { - for (int rowIt = threadIdx.x; rowIt < rowLen; - rowIt += kNumThreadsPerBlock) { - if constexpr (multipleBlocksPerRow) { - outIndices[rowIt] = rowIt + rowStart; - outLogits[rowIt] = logits[rowIt + rowStart]; - } else { - outIndices[rowIt] = rowIt; - } - } - for (int rowIt = rowLen + threadIdx.x; rowIt < topK; - rowIt += kNumThreadsPerBlock) { - outIndices[rowIt] = -1; - if constexpr (multipleBlocksPerRow) { - outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX; - } - } - - return; - } - // Initialize values - if (threadIdx.x == 0) { - smemFinalDstIdx[0] = 0; - smemFoundTopKValues[0] = 0; - } - __syncthreads(); - int thresholdBinIdx = -1; - uint32_t logitPattern = 0; - - // Step 0: Process first 11 bits of half representation - bool continueToNextStep = - processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks, - useTopK512Optimization>( - indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, - smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, - smemFoundTopKValues, smemFinal, stride1, rowStart, topK); - - if (continueToNextStep) { - // Step 1: Process next 11 bits - continueToNextStep = - processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks, - useTopK512Optimization>( - indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, - smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, - smemFoundTopKValues, smemFinal, stride1, rowStart, topK); - } - - if (continueToNextStep) { - // Step 2: Process next 11 bits - continueToNextStep = - processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks, - useTopK512Optimization>( - indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, - smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, - smemFoundTopKValues, smemFinal, stride1, rowStart, topK); - } - - if (continueToNextStep) { - // Step 3: Process last 10 bits - processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks, - useTopK512Optimization>( - indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, - smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, - smemFoundTopKValues, smemFinal, stride1, rowStart, topK); - } - - if (!continueToNextStep) { - // The histogram did not proceed to the final 10 bits, therefore we need to - // sort the final items The logits of the elements to be sorted in the final - // pass. - auto insertionSort = [&]() { - // Sorting with insertion sort - auto baseIdx = smemFoundTopKValues[0]; - for (int i = threadIdx.x; i < smemFinalDstIdx[0]; - i += kNumThreadsPerBlock) { - int outIndex = 0; - auto logit = smemFinal.items.logits[i]; - for (int j = 0; j < smemFinalDstIdx[0]; j++) { - auto otherLogit = smemFinal.items.logits[j]; - if (logit < otherLogit || (logit == otherLogit && i < j)) { - outIndex++; - } - } - // Store if outIndex is in bounds - if (outIndex + baseIdx < topK) { - smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i]; - if constexpr (multipleBlocksPerRow) { - reinterpret_cast(smemOutput + topK)[outIndex + baseIdx] = - smemFinal.items.logits[i]; - } - } - } - }; - if constexpr (useRadixSort) { - if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) { - insertionSort(); - } else { - // Sorting with radix sort - float finalLogits[kNumFinalItemsPerThread]; - // The indices of the elements to be sorted in the final pass. - int finalIndices[kNumFinalItemsPerThread]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX; - if constexpr (useTopK512Optimization) finalIndices[ii] = -1; - } - - // Read the elements from SMEM. -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < smemFinalDstIdx[0]) { - finalLogits[ii] = smemFinal.items.logits[srcIdx]; - finalIndices[ii] = smemFinal.items.indices[srcIdx]; - } - } - // Make sure the shared memory has been read. - __syncthreads(); - - // Sort the elements. - FinalSort(smemFinal.finalSort) - .SortDescendingBlockedToStriped(finalLogits, finalIndices); - - // Copy the data back to the shared memory storage. - int baseIdx = smemFoundTopKValues[0]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - int dstIdx = baseIdx + srcIdx; - - if (dstIdx < topK) { - smemOutput[dstIdx] = finalIndices[ii]; - if constexpr (multipleBlocksPerRow) { - reinterpret_cast(smemOutput + topK)[dstIdx] = - finalLogits[ii]; - } - } - } - } - } else { - insertionSort(); - } - __syncthreads(); - } - - // Store to global memory. - for (int i = threadIdx.x; i < topK; i += kNumThreadsPerBlock) { - if constexpr (multipleBlocksPerRow) { - outIndices[i] = smemOutput[i]; - outLogits[i] = reinterpret_cast(smemOutput + topK)[i]; - } else { - if (stride1 == 1) { - // stride1 == 1 will use vectorized_process, which indexes already skip - // the rowStart. - outIndices[i] = smemOutput[i]; - } else { - outIndices[i] = smemOutput[i] - rowStart; - } - } - } -} - -} // namespace isolated_upstream diff --git a/csrc/ops.h b/csrc/ops.h index a6d19252fbda..73f7133cec36 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -875,8 +875,7 @@ 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 slot_flags, + 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 gates, - torch::stable::Tensor stats, bool relative); + torch::stable::Tensor counts, torch::stable::Tensor defer, bool relative); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index baec9d8aa8fe..b971e7158cee 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -10,13 +10,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { #ifdef VLLM_ENABLE_FUSED_MQA_TOPK // Single-launch fused FP8 MQA-logits + top-2048 prefill selection (SM100 - // family). Workspace tensors are caller-owned; slot_flags must be zero. + // 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! slot_flags, Tensor! values, Tensor! indices, Tensor! retained, " - "Tensor! retained_indices, Tensor! counts, Tensor! gates, " - "Tensor! stats, bool relative) -> ()"); + "Tensor! counter, Tensor! values, Tensor! indices, Tensor! retained, " + "Tensor! retained_indices, Tensor! counts, Tensor! defer, " + "bool relative) -> ()"); #endif // Compute per-token-group FP8 quantized tensor and scaling factor. diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 53fd376c1d91..83f93eb26c8f 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -23,25 +23,37 @@ def _dense(q, k, scales, weights, starts, ends): """Upstream dense path: budgeted DeepGEMM logits + top_k_per_row_prefill.""" - from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + 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, + ) + 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) - step = max(1, LOGITS_BUDGET // (4 * k.shape[0])) + 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) - logits = fp8_fp4_mqa_logits( - (q[i:j], None), - (k, scales), - weights[i:j], - starts[i:j], - ends[i:j], - clean_logits=False, + 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, - starts[i:j], - ends[i:j], + torch.zeros_like(lens[i:j]), + lens[i:j], out[i:j], j - i, logits.stride(0), @@ -51,10 +63,24 @@ def _dense(q, k, scales, weights, starts, ends): return out -def _check(rel, q, k, scales, weights, starts, ends): - """Indices relative to the row start: count, bounds, uniqueness, scores.""" - from vllm.utils.deep_gemm import fp8_fp4_mqa_logits +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. + """ visible = (ends - starts).clamp(min=0) valid = rel >= 0 assert (valid.sum(1) == visible.clamp(max=TOPK)).all(), "valid count" @@ -63,18 +89,13 @@ def _check(rel, q, k, scales, weights, starts, ends): 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 = fp8_fp4_mqa_logits( - (q[i:j], None), - (k, scales), - weights[i:j], - starts[i:j], - ends[i:j], - clean_logits=True, - ) + 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 @@ -148,7 +169,7 @@ def test_single_request_matches_dense(keys, rows, case): # 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) - _check(dense, q, k, scales, weights, starts, ends) + _check(dense, q, k, scales, weights, starts, ends, scores=False) assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all() @@ -172,7 +193,7 @@ def test_multi_request_chunk_matches_dense(lens): assert 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) - _check(dense, q, k, scales, weights, starts, ends) + _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. overlap = sum( diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 8ed06c69b402..2cf90c2894fa 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -7,9 +7,10 @@ row's top-2048 keys, without materializing the ``[rows, keys]`` logits matrix that the dense path (``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``) needs. -Scratch memory is a fixed pool of 8-row slots (2 per SM, ~311 MB on a 148-SM -GPU) that is allocated once per device and reused by every call, independently -of the number of rows or keys. +Above 4K keys the kernel is persistent (one CTA per SM); each CTA owns two +8-row candidate slots. The scratch (~296 MiB on a 148-SM GPU) is allocated once +per device and reused by every call, independently of the number of rows or +keys. """ from dataclasses import dataclass @@ -31,32 +32,29 @@ _CANDIDATES_PER_ROW = 16384 _RETAINED_PER_ROW = 4096 _SLOTS_PER_SM = 2 -_STATS_PER_CTA = 7 @dataclass class FusedMQATopKWorkspace: """Caller-owned scratch of the fused kernel (see fused_mqa_topk.cu).""" - slot_flags: torch.Tensor # [S] int32, zero between calls - values: torch.Tensor # [8S, 16384] fp16 + counter: torch.Tensor # [2] int32, zero between calls (the kernel restores it) + values: torch.Tensor # [8S, 16384] fp16, S = 2 slots per SM indices: torch.Tensor # [8S, 16384] int32 retained: torch.Tensor # [8S, 4096] fp32 retained_indices: torch.Tensor # [8S, 4096] int32 - counts: torch.Tensor # [16384] int32 - gates: torch.Tensor # [16384] fp32 - stats: torch.Tensor # [16384 / 8 * 7] int32, per-CTA diagnostics + counts: torch.Tensor # [16384] int32, per-row candidate counts + defer: torch.Tensor # [16384 / 8] int32, blocks queued for the exact rescan def as_args(self) -> tuple[torch.Tensor, ...]: return ( - self.slot_flags, + self.counter, self.values, self.indices, self.retained, self.retained_indices, self.counts, - self.gates, - self.stats, + self.defer, ) @@ -75,8 +73,7 @@ def fused_mqa_topk_available() -> bool: def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: """Return the per-device workspace, allocating it on first use. - The slot flags are zeroed once; every kernel launch releases the slots it - claimed, so the flags are zero again whenever the stream is idle. + The block counter is zeroed once; every launch leaves it zeroed again. """ index = ( device.index @@ -90,7 +87,7 @@ def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: slots = _SLOTS_PER_SM * num_sms rows = FUSED_MQA_TOPK_MAX_ROWS ws = FusedMQATopKWorkspace( - slot_flags=torch.zeros(slots, dtype=torch.int32, device=dev), + counter=torch.zeros(2, dtype=torch.int32, device=dev), values=torch.empty( (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.float16, device=dev ), @@ -104,10 +101,7 @@ def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: (8 * slots, _RETAINED_PER_ROW), dtype=torch.int32, device=dev ), counts=torch.empty(rows, dtype=torch.int32, device=dev), - gates=torch.empty(rows, dtype=torch.float32, device=dev), - stats=torch.empty( - rows // 8 * _STATS_PER_CTA, dtype=torch.int32, device=dev - ), + defer=torch.empty(rows // 8, dtype=torch.int32, device=dev), ) nbytes = sum(t.numel() * t.element_size() for t in ws.as_args()) logger.info( From 61179f5692f20f7a2cc06822a129b1583842b3f1 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Wed, 14 Oct 2026 14:44:23 +0000 Subject: [PATCH 04/23] fix pre-commit Signed-off-by: LopezCastroRoberto --- vllm/model_executor/layers/sparse_attn_indexer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 070b89e25916..8a4b68d5ad55 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -824,7 +824,9 @@ def sparse_attn_indexer( ) ops.top_k_per_row_prefill( logits, - torch.zeros_like(row_ends) if row_starts is None else row_starts, + torch.zeros_like(row_ends) + if row_starts is None + else row_starts, row_ends, topk_indices[rows], num_rows, From 83a9c460130b28f40888f947d9b8c8ce0206f275 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Wed, 14 Oct 2026 15:15:44 +0000 Subject: [PATCH 05/23] fix pre-commit Signed-off-by: LopezCastroRoberto --- vllm/models/deepseek_v4/cpu/cpu_mla.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index 52d4f4c3b017..053c90700e2d 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, + use_fused_mqa_topk: 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/use_fused_mqa_topk 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 [ From d935bfc4ab83bdb72998b7afeb75137344d0cddb Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 16 Oct 2026 08:53:29 +0000 Subject: [PATCH 06/23] add topk=512/1024 support and PCP Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/fallback.cuh | 49 +++++---- csrc/fused_mqa_topk/fast_select.cuh | 17 ++-- csrc/fused_mqa_topk/fused_mqa_topk.cu | 22 +++-- csrc/fused_mqa_topk/persistent_epoch.inc | 99 ++++++++----------- csrc/fused_mqa_topk/persistent_kernel.cuh | 77 ++++++++++----- csrc/fused_mqa_topk/resident_kernel.cuh | 28 +++--- csrc/fused_mqa_topk/row_select.cuh | 39 ++------ csrc/fused_mqa_topk/sample.cuh | 86 +--------------- csrc/torch_bindings.cpp | 5 +- .../kernels/attention/test_fused_mqa_topk.py | 48 ++++----- .../kernels/attention/dsa/fused_mqa_topk.py | 19 ++-- .../layers/sparse_attn_indexer.py | 12 ++- vllm/v1/attention/backends/mla/indexer.py | 1 - 13 files changed, 216 insertions(+), 286 deletions(-) diff --git a/csrc/fused_mqa_topk/fallback.cuh b/csrc/fused_mqa_topk/fallback.cuh index 607881b0a095..824cc937f79c 100644 --- a/csrc/fused_mqa_topk/fallback.cuh +++ b/csrc/fused_mqa_topk/fallback.cuh @@ -1,8 +1,8 @@ // SPDX-License-Identifier: MIT -// Exact recovery for blocks whose gated candidate count left [2048, 16384] (or +// Exact recovery for blocks whose gated candidate count left [K, 16384] (or // whose weights are too small to halve exactly): the block is rescanned // densely, and every 1024-key segment of scores is merged into a running -// per-row top set (kept <= FallbackRetain, compacted to 2048). +// per-row top set (kept <= FallbackRetain, compacted to K). #pragma once #include #include "dense_scan.cuh" @@ -11,6 +11,7 @@ #include #include constexpr unsigned FallbackSegment = 1024, FallbackRetain = 4096; +// output: [rows, K] int32 (K = the row top-k). struct FallbackBuffers { float* retained; int* indices; @@ -50,10 +51,10 @@ __device__ __forceinline__ void fb_sync() { cutlass::arch::NamedBarrier(256, 2).sync(); asm volatile("" ::: "memory"); } -__device__ void fb_compact(FallbackShared& s, FallbackBuffers b, unsigned local, - unsigned row, unsigned tid, bool final = false) { - constexpr unsigned K = 2048, Full = 0xffffffffu, T = 256, - Items = FallbackRetain / T; +__device__ void fb_compact(FallbackShared& s, FallbackBuffers b, unsigned K, + unsigned local, unsigned row, unsigned tid, + bool final = false) { + constexpr unsigned Full = 0xffffffffu, T = 256, Items = FallbackRetain / T; unsigned lane = tid & 31, count = s.kept[local]; uint64_t base = (uint64_t)row * FallbackRetain; float values[Items]; @@ -193,11 +194,11 @@ __device__ void fb_compact(FallbackShared& s, FallbackBuffers b, unsigned local, } fb_sync(); } -__device__ void fb_select(FallbackShared& s, FallbackBuffers b, unsigned local, - unsigned row, unsigned epoch, unsigned first, - unsigned n, unsigned tid, bool final) { - constexpr unsigned K = 2048, Full = 0xffffffffu, T = 32, - Items = FallbackSegment / T; +__device__ void fb_select(FallbackShared& s, FallbackBuffers b, unsigned K, + unsigned local, unsigned row, unsigned epoch, + unsigned first, unsigned n, unsigned tid, + bool final) { + constexpr unsigned Full = 0xffffffffu, T = 32, Items = FallbackSegment / T; unsigned lane = tid, bank = epoch & 1; float values[Items], gate = __uint_as_float(s.gates[local]); unsigned count = 0; @@ -278,7 +279,7 @@ struct RescanPolicy { static constexpr uint32_t kSplitsPerSegment = FallbackSegment / 256; FallbackShared& stream; FallbackBuffers buffers; - uint32_t query_base, row_base, epoch = 0; + uint32_t topk, query_base, row_base, epoch = 0; uint32_t seq_k_start[4], seq_k_end[4]; float gate[4]; CUTLASS_DEVICE void begin_block(const Scheduler& s, uint32_t q_block, @@ -316,17 +317,17 @@ struct RescanPolicy { const unsigned tid = threadIdx.x; for (unsigned local = 0; local < 8; ++local) if (stream.kept[local] > FallbackRetain - FallbackSegment) - fb_compact(stream, buffers, local, row_base + local, tid); + fb_compact(stream, buffers, topk, local, row_base + local, tid); fb_sync(); const unsigned local = tid / 32; - fb_select(stream, buffers, local, row_base + local, epoch, + fb_select(stream, buffers, topk, local, row_base + local, epoch, kv_base + (split / kSplitsPerSegment) * FallbackSegment, (split % kSplitsPerSegment + 1) * 256, tid % 32, false); fb_sync(); if (split + 1 == num_splits) for (unsigned l = 0; l < 8; ++l) - if (stream.kept[l] > 2048) - fb_compact(stream, buffers, l, row_base + l, tid, true); + if (stream.kept[l] > topk) + fb_compact(stream, buffers, topk, l, row_base + l, tid, true); fb_sync(); ++epoch; } @@ -335,8 +336,8 @@ struct RescanPolicy { const unsigned row = s.get_logits_row(q_block, local), count = stream.kept[local]; if (row < s.num_q_tokens) - for (unsigned j = lane; j < 2048; j += 32) - buffers.output[(uint64_t)row * 2048 + j] = + for (unsigned j = lane; j < topk; j += 32) + buffers.output[(uint64_t)row * topk + j] = j < count ? buffers.indices[(uint64_t)row * FallbackRetain + j] : -1; } @@ -345,17 +346,15 @@ struct RescanPolicy { // Rescans one block exactly; `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__ void rescan_block(const Scheduler& block, - FallbackBuffers buffers, - const cute::TmaDescriptor& tq, - const cute::TmaDescriptor& tk, - const cute::TmaDescriptor& ts, - const cute::TmaDescriptor& tw) { +__device__ __noinline__ void rescan_block( + const Scheduler& block, FallbackBuffers buffers, unsigned 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& stream = *reinterpret_cast( smem_buffer + sizeof(fused_mqa_dense::SharedStorage)); stream.init(threadIdx.x); - RescanPolicy policy{stream, buffers}; + RescanPolicy policy{stream, buffers, topk}; fused_mqa_dense::score_block(block, policy, tq, tk, ts, tw); } diff --git a/csrc/fused_mqa_topk/fast_select.cuh b/csrc/fused_mqa_topk/fast_select.cuh index 597733a59d9f..fe89f4a3daed 100644 --- a/csrc/fused_mqa_topk/fast_select.cuh +++ b/csrc/fused_mqa_topk/fast_select.cuh @@ -9,9 +9,10 @@ // into at most kBins bins: collisions are rare, plain shared atomics suffice // and each lane keeps U 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. +// ties until K. K (the row's top-k: 512, 1024 or 2048) is a runtime value; +// shared-memory scratch is sized for kMaxK. namespace fast_select { -constexpr unsigned K = 2048, Full = 0xffffffffu, kBins = 1024, U = 8; +constexpr unsigned kMaxK = 2048, Full = 0xffffffffu, kBins = 1024, U = 8; // Bin (descending) where the cumulative count first reaches `rank`; // `above` = elements in higher bins. Lane l owns bins kBins-1-32l .. -31. @@ -50,7 +51,7 @@ __device__ __forceinline__ unsigned find_bin(const unsigned* hist, // Requires count >= K (callers recover rows with fewer candidates). __device__ void select_row(unsigned* codes, unsigned count, const uint16_t* values, const int32_t* packed, - int32_t* out, unsigned* hist) { + int32_t* out, unsigned K, unsigned* hist) { const unsigned lane = threadIdx.x & 31; unsigned lo = 0xffffffffu, hi = 0; unsigned j = lane; @@ -157,7 +158,7 @@ struct IdentityKey { template __device__ void select_resident(const uint16_t* hi, const uint8_t* lo, unsigned count, unsigned s0, int32_t* out, - unsigned* hist, KeyOf key_of = {}) { + unsigned K, unsigned* hist, KeyOf key_of = {}) { const unsigned lane = threadIdx.x & 31; unsigned lo_code = 0xffffffffu, hi_code = 0; for (unsigned j = lane; j < count; j += 32) { @@ -225,7 +226,7 @@ constexpr unsigned kSamplesPerLane = 16; // 8 elements per 16-byte shared load, per-lane counts with one warp prefix scan // (no per-element ballots), output staged in shared memory and copied out with // 16-byte stores. hi_row/lo_row are indexed by key index. -constexpr unsigned kMid3Cap = 1024, kScratch3Bytes = K * 4 + kMid3Cap * 2; +constexpr unsigned kMid3Cap = 1024, kScratch3Bytes = kMaxK * 4 + kMid3Cap * 2; __device__ __forceinline__ unsigned warp_exclusive_scan(unsigned v, unsigned& total) { @@ -259,7 +260,7 @@ template __device__ bool select_resident_bracket(const uint16_t* hi_row, const uint8_t* lo_row, unsigned s0, unsigned count, int32_t* out, - unsigned char* scratch, + unsigned K, unsigned char* scratch, KeyOf key_of = {}) { const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; // hi16 bracket from 512 strided samples (+-40 sample ranks around the @@ -301,7 +302,7 @@ __device__ bool select_resident_bracket(const uint16_t* hi_row, pm = warp_exclusive_scan(nm, mid); if (above >= K || above + mid < K || mid > kMid3Cap) return false; int32_t* stage = reinterpret_cast(scratch); - uint16_t* midbuf = reinterpret_cast(scratch + K * 4); + uint16_t* midbuf = reinterpret_cast(scratch + kMaxK * 4); #pragma unroll 2 for (unsigned unit = lane; unit < units; unit += 32) classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, @@ -362,7 +363,7 @@ __device__ bool select_resident_bracket(const uint16_t* hi_row, // (rare). template __device__ void select_hist_generic(CodeOf code_of, unsigned count, - int32_t* out, unsigned* hist, + int32_t* out, unsigned K, unsigned* hist, KeyOf key_of) { const unsigned lane = threadIdx.x & 31; unsigned lo_code = 0xffffffffu, hi_code = 0; diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu index b7a85bb4a98b..38c6e2d5d6f1 100644 --- a/csrc/fused_mqa_topk/fused_mqa_topk.cu +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -102,7 +102,6 @@ CUtensorMap make_strided_kv(void* ptr, unsigned groups, unsigned stride) { constexpr int kHeads = 32; constexpr int kHeadDim = 128; -constexpr int kTopK = 2048; constexpr int kCandidates = 16384; // per-row candidate slab constexpr int kRetained = 4096; // per-row recovery slab @@ -118,10 +117,10 @@ void check_tensor(const torch::stable::Tensor& t, int32_t device, // q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32, // weights [rows,32] fp32, starts/ends [rows] int32 (visible key range -// [start,end) of each row), out [rows,2048] int32: absolute key indices -// (columns of the logits), or with relative=true indices relative to the -// row's start (key - start, the top_k_per_row_prefill convention); -1 padded -// when fewer than 2048 keys are visible. +// [start,end) of each row), out [rows,K] int32 with K (the top-k) 512, 1024 or +// 2048: absolute key indices (columns of the logits), or with relative=true +// indices relative to the row's start (key - start, the top_k_per_row_prefill +// convention); -1 padded when fewer than K keys are visible. // 0 < rows <= 16384, 0 < keys <= 1048576 (the last 8-row block may be partial). // 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): @@ -169,8 +168,10 @@ void fused_mqa_topk_prefill( "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 && out.size(1) == kTopK, - "out must be [rows, 2048]"); + STD_TORCH_CHECK( + out.dim() == 2 && out.size(0) == rows && + (out.size(1) == 512 || out.size(1) == 1024 || out.size(1) == 2048), + "out must be [rows, K] with K in {512, 1024, 2048}"); STD_TORCH_CHECK(reinterpret_cast(out.data_ptr()) % 16 == 0, "out must be 16-byte aligned"); @@ -178,7 +179,8 @@ void fused_mqa_topk_prefill( const cudaStream_t stream = get_current_cuda_stream(dev); const unsigned urows = static_cast(rows), ukeys = static_cast(keys), - ublocks = static_cast(blocks); + 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()); @@ -198,7 +200,7 @@ void fused_mqa_topk_prefill( reinterpret_cast(fused_mqa_topk_fp8_resident_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); fused_mqa_topk_fp8_resident_kernel<<>>( - urows, ukeys, starts_p, ends_p, out_p, relative, tq, ts, tk, tw); + urows, ukeys, starts_p, ends_p, out_p, topk, relative, tq, ts, tk, tw); check_cuda(cudaGetLastError()); return; } @@ -225,7 +227,7 @@ void fused_mqa_topk_prefill( const unsigned grid = static_cast( std::min({int64_t(num_sms), blocks, values.size(0) / 16})); fused_persist::LaunchConfig c; - check_cuda(fused_persist::launch_config(ukeys, c)); + check_cuda(fused_persist::launch_config(ukeys, topk, c)); auto tks = make_strided_kv(k.data_ptr(), c.sample_groups, c.sample_stride); const FallbackBuffers fallback{static_cast(retained.data_ptr()), static_cast(retained_indices.data_ptr()), diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc index f76c7da8205b..3b4fa49f74cc 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -75,7 +75,9 @@ if (warp_idx == 8) { e.ends[lane] = en; } const bool all_short = - __ballot_sync(0xffffffffu, lane < 8 && en - st > 2048u) == 0; + __ballot_sync(0xffffffffu, lane < 8 && en - st > topk) == 0; + const bool gated = + __ballot_sync(0xffffffffu, lane < 8 && en - st > sample_target) != 0; // KV range: rows past the end are clamped to the last valid row (DeepGEMM // scheduler). const unsigned src = min(row_base + (lane & 7), rows - 1) - row_base; @@ -97,6 +99,7 @@ if (warp_idx == 8) { e.all_short = all_short; e.kv_base = kv_base; e.splits = (hi - kv_base + 255) / 256; + e.samples = gated ? sample_tiles : 0u; e.slot = slot; } __syncwarp(); @@ -145,7 +148,7 @@ if (warp_idx == 8) { claim_full[k & 1].wait((k >> 1) & 1); const ClaimEntry& e = entries[k & 1]; const unsigned vblock = e.vblock, all_short = e.all_short, - kv_base = e.kv_base, splits = e.splits; + kv_base = e.kv_base, splits = e.splits, samples = e.samples; __syncwarp(); if (vblock >= kPause) break; claim_empty[k & 1].arrive(); @@ -165,7 +168,7 @@ if (warp_idx == 8) { : 0.0f; } #pragma unroll 1 - for (unsigned t = 0; t < sample_tiles; ++t) { + for (unsigned 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); @@ -197,7 +200,7 @@ if (warp_idx == 8) { 1); // selection scratch overlaps the KV stages const unsigned center = sample_stride / 2; #pragma unroll 1 - for (unsigned t = 0; t < (KVS == 3 ? 0u : sample_tiles); ++t) { + for (unsigned t = 0; t < (KVS == 3 ? 0u : samples); ++t) { CUTE_TIE_DECL(kv_pipe.advance(), s, ph); empty_kv[s].wait(ph ^ 1); fused_mqa_dg::tma_copy<128, 256, 0, cutlass::float_e4m3_t, true>( @@ -242,7 +245,7 @@ if (warp_idx == 8) { claim_full[k & 1].wait((k >> 1) & 1); const ClaimEntry& e = entries[k & 1]; const unsigned vblock = e.vblock, all_short = e.all_short, - splits = e.splits; + splits = e.splits, samples = e.samples; __syncwarp(); if (vblock >= kPause) break; claim_empty[k & 1].arrive(); @@ -254,7 +257,7 @@ if (warp_idx == 8) { full_lin->wait(q_phase); // this block's linear-term rows (warp 11) ptx::tcgen05_after_thread_sync(); } - const unsigned total = sample_tiles + splits; + const unsigned total = samples + splits; #pragma unroll 1 for (unsigned t = 0; t < total; ++t) { CUTE_TIE_DECL(kv_pipe.advance(), s, ph); @@ -388,8 +391,8 @@ if (warp_idx == 8) { if (e.all_short) { for (unsigned r = 0; r < 8 && row_base + r < rows; ++r) { const unsigned s0 = e.starts[r], vis = e.ends[r] - s0; - for (unsigned j = math_tid; j < 2048; j += 256) - fallback.output[(uint64_t)(row_base + r) * 2048 + j] = + for (unsigned j = math_tid; j < topk; j += 256) + fallback.output[(uint64_t)(row_base + r) * topk + j] = j < vis ? int((relative ? 0u : s0) + j) : -1; } __syncwarp(); @@ -406,7 +409,6 @@ if (warp_idx == 8) { static_assert(SampleQ * 2112 % 4 == 0); for (unsigned j = math_tid; j < SampleQ * 2112 / 4; j += 256) reinterpret_cast(stream.hist)[j] = make_uint4(0, 0, 0, 0); - if (math_tid < 8) stream.counts[math_tid] = 0; float weights[4][32]; [[maybe_unused]] float row_s[4]; // > 16K: pow2 row scale of the linear term @@ -455,7 +457,7 @@ if (warp_idx == 8) { sample_sync(); // ---- sample tiles ---- #pragma unroll 1 - for (unsigned t = 0; t < sample_tiles; ++t) { + for (unsigned t = 0; t < e.samples; ++t) { CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); if constexpr (!kRel) full_kv[kv_stage].wait(kv_phase); #pragma unroll @@ -530,35 +532,18 @@ if (warp_idx == 8) { bq = (sum.x + sum.y) * scale_kv; const bool valid = in_view && kv_offset >= seq_k_start[i] && kv_offset < seq_k_end[i] && isfinite(bq); - if constexpr (kLin) { - if (valid) { - const unsigned bin = sample_bin(bq); - atomicAdd(stream.hist[query_base + i] + bin + bin / 32, 1u); - } - } else { - stream.scores[query_base + i][kv_lane] = valid ? bq : -CUDART_INF_F; + if (valid) { + const unsigned bin = sample_bin(bq); + atomicAdd(stream.hist[query_base + i] + bin + bin / 32, 1u); } } } - if constexpr (!kLin) { - sample_sync(); - accumulate_sample(stream, math_tid); - sample_sync(); - } - } - if constexpr (kLin) { - sample_sync(); - finish_sample_histogram(stream.hist[warp_idx], gate_s + warp_idx, - e.ends[warp_idx] - e.starts[warp_idx], - sample_target, lane); - } else { - const unsigned rb = - *reinterpret_cast(&e.vblock) * 8; - SampleBuffers sb{gate_s - rb, nullptr, nullptr, nullptr, sample_target}; - finish_sample(stream, sb, warp_idx, rb + warp_idx, - e.ends[warp_idx] - e.starts[warp_idx], lane); } sample_sync(); + finish_sample_histogram(stream.hist[warp_idx], gate_s + warp_idx, + e.ends[warp_idx] - e.starts[warp_idx], + sample_target, lane); + sample_sync(); const unsigned row_base2 = *reinterpret_cast(&e.vblock) * 8; @@ -906,12 +891,12 @@ if (warp_idx == 8) { unsigned recover = 0, short_rows = 0; row_offset[0] = 0; for (unsigned r = 0; r < 8; ++r) { - const bool is_short = ce.ends[r] - ce.starts[r] <= 2048u; + const bool is_short = ce.ends[r] - ce.starts[r] <= topk; short_rows |= unsigned(is_short) << r; const int c = kSmem ? int(row_count[r]) : __ldcg(counts + row_base3 + r); const int real = kSmem ? (c & ((1 << 21) - 1)) : c; - recover |= !is_short && (real < 2048 || real > 16384); + recover |= !is_short && (real < int(topk) || real > 16384); // > 16K rows start 16-byte aligned for select_staged_row's vector // loads. row_offset[r + 1] = @@ -950,8 +935,8 @@ if (warp_idx == 8) { if (r >= rows) return; const unsigned s0 = ce.starts[r - row_base3], visible = ce.ends[r - row_base3] - s0; - int* o = fallback.output + (uint64_t)r * 2048; - for (unsigned j = t; j < 2048; j += stride) + int* o = fallback.output + (uint64_t)r * topk; + for (unsigned j = t; j < topk; j += stride) o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; }; if (!recover) { @@ -960,7 +945,7 @@ if (warp_idx == 8) { const unsigned warp = warp_idx; if (row_base3 + warp < rows) { const unsigned r = row_base3 + warp; - int* o = fallback.output + (uint64_t)r * 2048; + int* o = fallback.output + (uint64_t)r * topk; if (short_rows >> warp & 1u) write_short(r, lane, 32); else { @@ -985,7 +970,7 @@ if (warp_idx == 8) { : (unsigned(__ldcg(gv + e2 - kCandCap)) << 8 | (unsigned(__ldcg(gi + e2 - kCandCap)) >> 16)); }, - total, o, reinterpret_cast(scratch), + total, o, topk, reinterpret_cast(scratch), [&](unsigned e2) { return e2 < kCandCap ? int32_t(idx[e2]) @@ -993,8 +978,8 @@ if (warp_idx == 8) { }); } else if (!row_select::select_shared_row( hi, lo, idx, kCandCap, unsigned(c & ((1 << 21) - 1)), - o, scratch)) { - fast_select::select_resident(hi, lo, kCandCap, 0u, o, + o, topk, scratch)) { + fast_select::select_resident(hi, lo, kCandCap, 0u, o, topk, reinterpret_cast(scratch), key_of); } @@ -1017,11 +1002,11 @@ if (warp_idx == 8) { write_short(r, lane, 32); else { const unsigned cnt = __ldcg(counts + r); - int* o = fallback.output + (uint64_t)r * 2048; + int* o = fallback.output + (uint64_t)r * topk; if (!row_select::select_staged_row(codes, cnt, vals + base, - indices + base, o, hist)) + indices + base, o, topk, hist)) fast_select::select_row(codes, cnt, vals + base, indices + base, - o, hist); + o, topk, hist); } } cutlass::arch::NamedBarrier(256, 5).sync(); @@ -1030,15 +1015,16 @@ if (warp_idx == 8) { } cutlass::arch::NamedBarrier(256, 5).sync(); if (relative && !recover) { - for (unsigned i = math_tid; i < 8 * 512; i += 256) { - const unsigned r = row_base3 + i / 512; - if (r >= rows || ce.ends[i / 512] - ce.starts[i / 512] <= 2048u) + const unsigned quads = topk / 4; + for (unsigned i = math_tid; i < 8 * quads; i += 256) { + const unsigned r = row_base3 + i / quads; + if (r >= rows || ce.ends[i / quads] - ce.starts[i / quads] <= topk) continue; const int s = int(starts_g[r]); if (s == 0) continue; int4* p = - reinterpret_cast(fallback.output + (uint64_t)r * 2048) + - i % 512; + reinterpret_cast(fallback.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; @@ -1083,8 +1069,8 @@ if (defer_head != kTerm) { fb.retained += row_delta * 4096; fb.indices += row_delta * 4096; fused_mqa_fallback::rescan_block( - fused_mqa_dense::Scheduler(vblock, rows, keys, starts, ends), fb, tq, - tk, ts, tw); + fused_mqa_dense::Scheduler(vblock, rows, keys, starts, ends), fb, topk, + tq, tk, ts, tw); __syncthreads(); if (threadIdx.x == 0) { auto& fst = @@ -1098,14 +1084,15 @@ if (defer_head != kTerm) { } __syncthreads(); if (relative) { - for (unsigned i = threadIdx.x; i < 8 * 512; i += 384) { - const unsigned r = row_base + i / 512; + const unsigned quads = topk / 4; + for (unsigned i = threadIdx.x; i < 8 * quads; i += 384) { + const unsigned r = row_base + i / quads; if (r >= rows) continue; const int s = int(starts_g[r]); if (s == 0) continue; int4* p = - reinterpret_cast(fallback.output + (uint64_t)r * 2048) + - i % 512; + reinterpret_cast(fallback.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; diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh index 3cf85fa2b0a7..6d9fa2f3c6ec 100644 --- a/csrc/fused_mqa_topk/persistent_kernel.cuh +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Persistent fused FP8 MQA logits + per-row top-2048 (SM100), +// 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 @@ -12,7 +12,7 @@ // warp 10 UMMA issuer // warp 11 drain of math WG0 rows (> 16K keys), idle otherwise // warps 0-7 math: sample histogram -> gate -> scan/emit -> selection -// A block whose candidate count leaves [2048, 16384] pauses the role loops; the +// A block whose candidate count leaves [K, 16384] pauses the role loops; the // CTA rescans it exactly (fallback.cuh) and resumes the loops with freshly // initialised barriers. #pragma once @@ -162,9 +162,9 @@ __device__ __forceinline__ void compute_lin_rows(PStorage<3>& sm, } } -// > 16K gate: the k-th largest sample (k = count * target / visible) of one row -// from its 2048-bin histogram; one warp per row. Same result as finish_sample -// (count = histogram total). +// Gate: the k-th largest sample (k = count * target / visible, count = the +// histogram total) of one row from its 2048-bin histogram; one warp per row. +// -inf when the row sees at most `target` keys or has no samples. __device__ __noinline__ void finish_sample_histogram(const unsigned* hist, float* gate, unsigned visible, @@ -224,8 +224,10 @@ struct alignas(16) PHandoffData { uint32_t bases[2][8]; }; +// samples: gate sample tiles of the block (0 when no row sees more keys than +// the gate target, so every gate is -inf). struct ClaimEntry { - unsigned vblock, all_short, kv_base, splits, slot; + unsigned vblock, all_short, kv_base, splits, slot, samples; unsigned starts[8], ends[8]; }; @@ -366,7 +368,7 @@ __device__ __forceinline__ void ho_drain(PHandoffData& d, TBar* hfull, hempty[bank * 8 + producer].arrive(); } -template +template __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( unsigned rows, unsigned keys, const unsigned* starts_g, const unsigned* ends_g, uint16_t* vals, int* indices, int* counts, @@ -385,6 +387,8 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( constexpr uint32_t kCandGlobalCap = 16384; using St = PStorage; using B = BarIdx; + // Compile-time: a runtime top-k costs ~0.5% in the > 16K scan (codegen). + constexpr unsigned topk = TopK; // > 16K keys: rebalanced scan (linear term on the tensor core, per-warpgroup // N=144 MMAs). constexpr bool kLin = !kSmem; @@ -478,7 +482,18 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( // ---------------- host: launch configuration (> kResidentKeys keys) // ---------------- -using PersistentKernel = decltype(&fused_mqa_topk_fp8_prefill_kernel<4, 2>); +using PersistentKernel = + decltype(&fused_mqa_topk_fp8_prefill_kernel<4, 2, 2048>); + +template +PersistentKernel persistent_kernel(unsigned keys) { + return keys <= 8192 ? fused_mqa_topk_fp8_prefill_kernel<4, 2, TopK> + : keys <= 16384 ? fused_mqa_topk_fp8_prefill_kernel<8, 2, TopK> + : keys <= 32768 ? fused_mqa_topk_fp8_prefill_kernel<16, 3, TopK> + : keys <= 131072 ? fused_mqa_topk_fp8_prefill_kernel<32, 3, TopK> + : keys <= 524288 ? fused_mqa_topk_fp8_prefill_kernel<64, 3, TopK> + : fused_mqa_topk_fp8_prefill_kernel<128, 3, TopK>; +} struct LaunchConfig { PersistentKernel kernel; @@ -490,19 +505,34 @@ struct LaunchConfig { }; // Picks the kernel instantiation, sizes its shared memory (setting the opt-in -// limit) and the gate sample for `keys` keys. -inline cudaError_t launch_config(unsigned keys, LaunchConfig& c) { - // <= 16K keys: shared candidates (cap kCandCap), so aim lower: 2560 - // candidates. Up to 128K: 3072. Both size the sample so the gate's quantile - // rank is ~100 (about 10% count error). Longer contexts aim for 4096 with - // rank 64, clamped to 16-32 sample tiles (rank ~32 at 1M keys): a sample tile - // costs about two scan tiles, so a larger sample costs more than the smaller - // candidate set saves. +// limit) and the gate sample for `keys` keys and top-`topk` rows. +inline cudaError_t launch_config(unsigned keys, unsigned 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. At + // <= 16K the candidates live in shared memory (cap kCandCap); top-2048 needs + // rank 100 there to keep 2560 candidates above 2048. 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 { + unsigned target, rank; + }; + constexpr Gate kGates[3][3] = { + // <= 16K <= 128K > 128K + {{2048, 50}, {1024, 32}, {1536, 64}}, // top-512 + {{2048, 50}, {1536, 48}, {2048, 48}}, // top-1024 + {{2560, 100}, {3072, 64}, {4096, 64}}, // top-2048 + }; const bool small = keys <= 16384, long_ctx = keys > 131072; - c.sample_target = small ? 2560u : (long_ctx ? 4096u : 3072u); - const unsigned rank = long_ctx ? 64u : 100u; + const Gate gate = kGates[topk == 512 ? 0 : (topk == 1024 ? 1 : 2)] + [small ? 0 : (long_ctx ? 2 : 1)]; + c.sample_target = gate.target; + const unsigned rank = gate.rank; const unsigned rank_tiles = std::max( - 2u, unsigned((uint64_t(rank) * keys + 256ull * c.sample_target - 1) / + 1u, unsigned((uint64_t(rank) * keys + 256ull * c.sample_target - 1) / (256ull * c.sample_target))); c.sample_tiles = long_ctx ? std::min(32u, std::max(16u, rank_tiles)) : rank_tiles; @@ -510,12 +540,9 @@ inline cudaError_t launch_config(unsigned keys, LaunchConfig& c) { 1u, (keys + 256u * c.sample_tiles - 1) / (256u * c.sample_tiles)); c.sample_groups = (keys - c.sample_stride / 2 + c.sample_stride - 1) / c.sample_stride; - c.kernel = keys <= 8192 ? fused_mqa_topk_fp8_prefill_kernel<4, 2> - : small ? fused_mqa_topk_fp8_prefill_kernel<8, 2> - : keys <= 32768 ? fused_mqa_topk_fp8_prefill_kernel<16, 3> - : keys <= 131072 ? fused_mqa_topk_fp8_prefill_kernel<32, 3> - : keys <= 524288 ? fused_mqa_topk_fp8_prefill_kernel<64, 3> - : fused_mqa_topk_fp8_prefill_kernel<128, 3>; + c.kernel = topk == 512 ? persistent_kernel<512>(keys) + : topk == 1024 ? persistent_kernel<1024>(keys) + : persistent_kernel<2048>(keys); // The recovery rescan reuses the same dynamic shared memory with the fallback // core's layout. diff --git a/csrc/fused_mqa_topk/resident_kernel.cuh b/csrc/fused_mqa_topk/resident_kernel.cuh index c21c0b4e8a5b..41abeeea6698 100644 --- a/csrc/fused_mqa_topk/resident_kernel.cuh +++ b/csrc/fused_mqa_topk/resident_kernel.cuh @@ -48,7 +48,7 @@ struct CodesPolicy { // <= 4K keys: one CTA per 8-row block, heaviest (latest rows) first. __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_resident_kernel( unsigned rows, unsigned keys, const unsigned* starts_g, - const unsigned* ends_g, int* output, bool relative, + const unsigned* ends_g, int* output, unsigned topk, bool relative, const __grid_constant__ cute::TmaDescriptor tq, const __grid_constant__ cute::TmaDescriptor ts, const __grid_constant__ cute::TmaDescriptor tk, @@ -66,17 +66,17 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_resident_kernel( __syncthreads(); const unsigned* const starts = row_bounds[0] - row_base; const unsigned* const ends = row_bounds[1] - row_base; - // Rows with <= 2048 visible keys output all of them (relative indices start + // Rows with <= topk visible keys output all of them (relative indices start // at 0). const auto write_short = [&](unsigned r, unsigned t, unsigned stride) { const unsigned visible = ends[r] - starts[r]; - for (unsigned j = t; j < 2048; j += stride) - output[(uint64_t)r * 2048 + j] = + for (unsigned j = t; j < topk; j += stride) + output[(uint64_t)r * topk + j] = j < visible ? int((relative ? 0u : starts[r]) + j) : -1; }; bool all_short = true; for (unsigned r = row_base; r < row_base + 8; ++r) - all_short &= ends[r] - starts[r] <= 2048u; + all_short &= ends[r] - starts[r] <= topk; if (all_short) { for (unsigned r = row_base; r < min(row_base + 8, rows); ++r) write_short(r, threadIdx.x, 384); @@ -95,10 +95,10 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_resident_kernel( const unsigned warp = threadIdx.x / 32; if (warp < 8 && row_base + warp < rows) { const unsigned r = row_base + warp, s0 = starts[r], visible = ends[r] - s0; - if (visible <= 2048) { + if (visible <= topk) { write_short(r, threadIdx.x & 31, 32); } else { - int* o = output + (uint64_t)r * 2048; + int* o = output + (uint64_t)r * topk; const uint16_t* hi = reinterpret_cast(codes) + warp * kResidentKeys; const uint8_t* lo = codes + 8 * kResidentKeys * 2 + warp * kResidentKeys; @@ -108,21 +108,23 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_resident_kernel( fast_select::kScratch3Bytes >= fast_select::kBins * 4, "select scratch must fit the Q/KV stages"); uint8_t* scratch = smem + warp * fast_select::kScratch3Bytes; - if (!fast_select::select_resident_bracket(hi, lo, s0, visible, o, + if (!fast_select::select_resident_bracket(hi, lo, s0, visible, o, topk, scratch)) - fast_select::select_resident(hi + s0, lo + s0, visible, s0, o, + fast_select::select_resident(hi + s0, lo + s0, visible, s0, o, topk, reinterpret_cast(scratch)); } } // Relative indices: selected rows are rewritten in place (L2-hot) once final. if (relative) { __syncthreads(); - for (unsigned i = threadIdx.x; i < 8 * 512; i += 384) { - const unsigned r = row_base + i / 512; - if (r >= rows || ends[r] - starts[r] <= 2048u) continue; + const unsigned quads = topk / 4; + for (unsigned i = threadIdx.x; i < 8 * quads; i += 384) { + const unsigned r = row_base + i / quads; + if (r >= rows || ends[r] - starts[r] <= topk) continue; const int s = int(starts_g[r]); if (s == 0) continue; - int4* p = reinterpret_cast(output + (uint64_t)r * 2048) + i % 512; + 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; diff --git a/csrc/fused_mqa_topk/row_select.cuh b/csrc/fused_mqa_topk/row_select.cuh index e93037994d86..824bf920ebfc 100644 --- a/csrc/fused_mqa_topk/row_select.cuh +++ b/csrc/fused_mqa_topk/row_select.cuh @@ -1,6 +1,6 @@ // SPDX-License-Identifier: MIT -// Single-warp top-2048 selection of one row's candidates for the persistent -// kernel. Both selections bracket the 2048-th largest fp24 code with 512 +// Single-warp top-K selection of one row's candidates for the persistent +// kernel. Both selections bracket the K-th largest fp24 code with 512 // sampled codes, place the codes above the bracket directly, and resolve the // bracket with histogram levels on register-resident codes; elements are placed // by per-lane counts and one warp prefix scan (no per-element ballots). @@ -93,28 +93,6 @@ __device__ __forceinline__ unsigned find_bin_v(const unsigned* hist, above = __shfl_sync(Full, acc_at, winner); return __shfl_sync(Full, chosen, winner); } -// Sampled hi16 bracket (16-step bit search over the 512 register samples). -__device__ __forceinline__ void bit_bracket( - const unsigned (&smp)[kSamplesPerLane], unsigned r_hi, unsigned r_lo, - unsigned& h_hi, unsigned& h_lo) { - h_hi = 0; - h_lo = 0; -#pragma unroll 1 - for (int bit = 15; bit >= 0; --bit) { - const unsigned ch = h_hi | (1u << bit), cl = h_lo | (1u << bit); - unsigned nh = 0, nl = 0; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) { - nh += smp[u] >= ch; - nl += smp[u] >= cl; - } - const unsigned both = - __reduce_add_sync(Full, nh | nl << 16); // counts <= 512 fit 16 bits - if ((both & 0xffffu) >= r_hi) h_hi = ch; - if ((both >> 16) >= r_lo) h_lo = cl; - } -} - // <= 16K keys: one row's shared-memory candidates (hi16 / lo8 code arrays and // the uint16 key array idx_row, read with 16-byte loads). Bracket keys are // preloaded into registers before the emit. Returns false (out untouched) when @@ -123,7 +101,7 @@ __device__ bool select_shared_row(const uint16_t* __restrict__ hi_row, const uint8_t* __restrict__ lo_row, const uint16_t* __restrict__ idx_row, unsigned count, unsigned real, - int32_t* __restrict__ out, + int32_t* __restrict__ out, unsigned K, unsigned char* scratch) { const unsigned lane = threadIdx.x & 31; unsigned h_hi, h_lo; @@ -139,7 +117,9 @@ __device__ bool select_shared_row(const uint16_t* __restrict__ hi_row, const unsigned expect = (K * s_real + real / 2) / real; const unsigned r_hi = expect > 40 ? expect - 40 : 1, r_lo = min(expect + 40, s_real); - bit_bracket(smp, r_hi, r_lo, h_hi, h_lo); + // h8 (768 words) borrows the stage area, written only after the bracket. + radix_bracket(smp, r_hi, r_lo, reinterpret_cast(scratch), h_hi, + h_lo); h_lo = max(h_lo, 1u); } const unsigned units = (count + 7) / 8; @@ -158,8 +138,9 @@ __device__ bool select_shared_row(const uint16_t* __restrict__ hi_row, pm = warp_exclusive_scan(nm, mid); if (above >= K || above + mid < K || mid > kMid3Cap) return false; uint16_t* stage = reinterpret_cast(scratch); - uint16_t* midbuf = reinterpret_cast(scratch + K * 2); - unsigned* hist = reinterpret_cast(scratch + K * 2 + kMid3Cap * 2); + uint16_t* midbuf = reinterpret_cast(scratch + kMaxK * 2); + unsigned* hist = + reinterpret_cast(scratch + kMaxK * 2 + kMid3Cap * 2); #pragma unroll 2 for (unsigned unit = lane; unit < units; unit += 32) { const uint4 h4 = hv[unit], k4 = iv[unit]; @@ -256,7 +237,7 @@ __device__ bool select_shared_row(const uint16_t* __restrict__ hi_row, __device__ __noinline__ bool select_staged_row( unsigned* __restrict__ codes, unsigned count, const uint16_t* __restrict__ values, const int32_t* __restrict__ packed, - int32_t* __restrict__ out, unsigned* hist) { + int32_t* __restrict__ out, unsigned K, unsigned* hist) { const unsigned lane = threadIdx.x & 31; if (count <= K) return false; const unsigned chunks = (count + 7) / 8; diff --git a/csrc/fused_mqa_topk/sample.cuh b/csrc/fused_mqa_topk/sample.cuh index 317c8c7764bc..1575a9084e95 100644 --- a/csrc/fused_mqa_topk/sample.cuh +++ b/csrc/fused_mqa_topk/sample.cuh @@ -1,25 +1,14 @@ +// Gate sample of the persistent kernel: each math warp histograms its row's +// sampled scores (fp16 bins >> 5, 2048 bins padded to 33 words per 32) and +// finish_sample_histogram picks the gate from them. #pragma once #include #include #include #include constexpr unsigned SampleQ = 8; -// target: candidate count the gate aims for. -struct SampleBuffers { - float* ring; - float* retained; - int* indices; - int* output; - unsigned target; -}; struct alignas(16) SampleShared { - float scores[SampleQ][256]; - unsigned hist[SampleQ][2112], counts[SampleQ], rank, cut; - __device__ void init(unsigned tid) { - for (unsigned j = tid; j < SampleQ * 2112; j += 384) - reinterpret_cast(hist)[j] = 0; - if (tid < SampleQ) counts[tid] = 0; - } + unsigned hist[SampleQ][2112]; }; __device__ __forceinline__ void sample_sync() { asm volatile("" ::: "memory"); @@ -34,70 +23,3 @@ __device__ __forceinline__ float sample_floor(unsigned b) { unsigned o = b << 5, h = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); return __half2float(__ushort_as_half(h)); } -__device__ void finish_sample(SampleShared& s, SampleBuffers b, unsigned local, - unsigned row, unsigned visible, unsigned tid) { - const unsigned target = b.target; - if (visible <= target) { - if (tid == 0) b.ring[row] = -CUDART_INF_F; - __syncwarp(); - return; - } - unsigned lane = tid % 32; - auto* hist = s.hist[local]; - unsigned count = s.counts[local]; - if (count == 0) { - if (tid == 0) b.ring[row] = -CUDART_INF_F; - __syncwarp(); - return; - } - unsigned sample_rank = - max(1u, min(count, (count * target + visible - 1) / visible)); - __syncwarp(); - if (tid < 32) { - unsigned sum = 0; - for (unsigned j = 0; j < 64; ++j) { - unsigned bin = 2047 - lane * 64 - j; - sum += hist[bin + bin / 32]; - } - unsigned inclusive = sum; - for (unsigned d = 1; d < 32; d *= 2) { - unsigned v = __shfl_up_sync(0xffffffffu, inclusive, d); - if (lane >= d) inclusive += v; - } - unsigned before = inclusive - sum, rank = sample_rank, - mask = - __ballot_sync(0xffffffffu, before < rank && inclusive >= rank), - winner = __ffs(mask) - 1, chosen = 0, above = before; - if (lane == winner) { - for (unsigned j = 0; j < 64; ++j) { - unsigned bin = 2047 - lane * 64 - j; - if (above + hist[bin + bin / 32] >= rank) { - chosen = bin; - break; - } - above += hist[bin + bin / 32]; - } - } - chosen = __shfl_sync(0xffffffffu, chosen, winner); - if (lane == 0) b.ring[row] = sample_floor(chosen); - } - __syncwarp(); -} - -__device__ void accumulate_sample(SampleShared& s, unsigned tid) { - unsigned local = tid / 32, lane = tid % 32, count = 0; -#pragma unroll - for (unsigned j = lane; j < 256; j += 32) { - float v = s.scores[local][j]; - bool valid = isfinite(v); - count += valid; - unsigned mask = __ballot_sync(0xffffffffu, valid); - if (valid) { - unsigned bin = sample_bin(v), peers = __match_any_sync(mask, bin); - if (lane == unsigned(__ffs(peers) - 1)) - atomicAdd(s.hist[local] + bin + bin / 32, __popc(peers)); - } - } - count = __reduce_add_sync(0xffffffffu, count); - if (lane == 0) s.counts[local] += count; -} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index b971e7158cee..0ee8f56130e7 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -9,8 +9,9 @@ // 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-2048 prefill selection (SM100 - // family). Workspace tensors are caller-owned; counter must be zero. + // 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, " diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 83f93eb26c8f..154d570b4865 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -17,11 +17,12 @@ 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): +def _dense(q, k, scales, weights, starts, ends, topk=TOPK): """Upstream dense path: budgeted DeepGEMM logits + top_k_per_row_prefill.""" from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits from vllm.utils.deep_gemm import ( @@ -36,7 +37,7 @@ def _dense(q, k, scales, weights, starts, ends): 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) + 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)) @@ -58,7 +59,7 @@ def _dense(q, k, scales, weights, starts, ends): j - i, logits.stride(0), logits.stride(1), - TOPK, + topk, ) return out @@ -81,9 +82,10 @@ def _check(rel, q, k, scales, weights, starts, ends, scores=True): ``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 (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() @@ -96,7 +98,7 @@ def _check(rel, q, k, scales, weights, starts, ends, scores=True): 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 + ref = logits.topk(topk, dim=1).values chosen = logits.gather(1, idx[i:j].clamp(min=0)).masked_fill( ~valid[i:j], -torch.inf ) @@ -115,6 +117,7 @@ def _inputs(rows, keys, seed=913): @requires_fused +@pytest.mark.parametrize("topk", TOPKS) @pytest.mark.parametrize( "keys,rows,case", [ @@ -128,27 +131,27 @@ def _inputs(rows, keys, seed=913): (1048576, 16, "causal"), ], ) -def test_single_request_matches_dense(keys, rows, case): +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 with fewer than 2048 visible keys + if case == "early": # rows around each top-k's visible-key boundary ends = torch.tensor( [ 0, 1, - 31, - 127, + 511, + 512, + 513, + 1023, 1024, + 1025, 2047, 2048, 2049, 4096, - 6144, 8192, - 10000, 12000, - 14000, 16000, keys, ], @@ -163,21 +166,22 @@ def test_single_request_matches_dense(keys, rows, case): elif case == "hot": weights = weights.abs() scales[:12288] = 4 - out = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32) + out = torch.empty(rows, topk, device="cuda", dtype=torch.int32) assert 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) + 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 +@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): +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) @@ -189,18 +193,20 @@ def test_multi_request_chunk_matches_dense(lens): 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) + out = torch.empty(rows, topk, device="cuda", dtype=torch.int32) assert 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) + 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. + # 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 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.99 * int((dense >= 0).sum()) + assert overlap >= 0.98 * int((dense >= 0).sum()) @requires_fused @@ -249,9 +255,7 @@ def test_unsupported_inputs_fall_back(): assert not fused_mqa_topk_prefill( q[:, :16], k, scales, weights[:, :16], starts, ends, out ) - assert not fused_mqa_topk_prefill( - q, k, scales, weights, starts, ends, out[:, :1024] - ) + assert not fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out[:, :768]) q17 = torch.zeros(16385, 32, 128, device="cuda").to(torch.float8_e4m3fn) assert not fused_mqa_topk_prefill(q17, k, scales, weights, starts, ends, out) diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 2cf90c2894fa..849a1b72e6b6 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -4,8 +4,9 @@ One kernel (``torch.ops._C.fused_mqa_topk_prefill``) scores every visible key of each query row with the DeepSeek sparse-attention indexer and selects the -row's top-2048 keys, without materializing the ``[rows, keys]`` logits matrix -that the dense path (``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``) needs. +row's top-k keys (k = 512, 1024 or 2048), without materializing the +``[rows, keys]`` logits matrix that the dense path (``fp8_fp4_mqa_logits`` + +``top_k_per_row_prefill``) needs. Above 4K keys the kernel is persistent (one CTA per SM); each CTA owns two 8-row candidate slots. The scratch (~296 MiB on a 148-SM GPU) is allocated once @@ -24,7 +25,7 @@ FUSED_MQA_TOPK_HEADS = 32 FUSED_MQA_TOPK_HEAD_DIM = 128 -FUSED_MQA_TOPK_TOPK = 2048 +FUSED_MQA_TOPK_TOPKS = (512, 1024, 2048) # Largest query chunk / gathered key range a single call accepts. FUSED_MQA_TOPK_MAX_ROWS = 16384 FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 @@ -117,7 +118,7 @@ def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: def fused_mqa_topk_config_supported(q: torch.Tensor, topk_tokens: int) -> bool: """Whether this indexer's query layout and top-k fit the fused kernel.""" return ( - topk_tokens == FUSED_MQA_TOPK_TOPK + topk_tokens in FUSED_MQA_TOPK_TOPKS and q.dtype == torch.float8_e4m3fn and q.dim() == 3 and tuple(q.shape[1:]) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM) @@ -166,11 +167,12 @@ def fused_mqa_topk_prefill( q: [rows, 32, 128] fp8 e4m3; k: [keys, 128] fp8 e4m3; scales: [keys] fp32; weights: [rows, 32] fp32; starts/ends: [rows] int32, the visible key range - ``[start, end)`` of each row; output: [rows, 2048] int32. + ``[start, end)`` of each row; output: [rows, k] int32 with k (the top-k) + 512, 1024 or 2048. 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 2048 keys are visible. ``rows_start_at_zero=True`` + padded when fewer than k keys are visible. ``rows_start_at_zero=True`` promises every ``start`` is 0 (single-request chunk) and skips the relative-index conversion. @@ -189,13 +191,14 @@ def fused_mqa_topk_prefill( or starts.numel() != rows or ends.numel() != rows or output.dtype != torch.int32 - or tuple(output.shape) != (rows, FUSED_MQA_TOPK_TOPK) + or output.shape[0] != rows ): return False logger.info_once( "Using the single-launch fused MQA+top-k prefill kernel " - "(FP8, 32 heads, top-k 2048)." + "(FP8, 32 heads, top-k %d).", + output.shape[1], ) q = q.contiguous() k = k.contiguous() diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 8a4b68d5ad55..011b1207c8b5 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -504,10 +504,11 @@ def sparse_attn_indexer( ) profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) + # Under PCP every rank ranks its rows over the whole (gathered) + # context; pure DCP needs the per-shard top-k merge, so it stays dense. uses_fused = ( envs.VLLM_USE_FUSED_MQA_TOPK - and not use_pcp - and dcp_world_size == 1 + and (use_pcp or dcp_world_size == 1) and not use_fp4_cache and candidate_blocks is None and fused_mqa_topk_config_supported(q_quant, topk_tokens) @@ -730,12 +731,13 @@ def sparse_attn_indexer( torch.int32 if mqa_logits_use_mx_scales() else torch.float32 ).squeeze(-1) # Single-launch fused MQA-logits + top-k (SM100, FP8, 32 heads, - # top-k 2048): no logits buffer, any context length. DCP/PCP, + # top-k 512/1024/2048): no logits buffer, any context length. + # It needs every row's whole context: no DCP, or PCP + DCP after + # the de-interleaving gather. Pure DCP (per-shard top-k + merge), # the FP4 cache and v4.1 candidate selection stay dense. use_fused = ( envs.VLLM_USE_FUSED_MQA_TOPK - and not use_pcp - and dcp_world_size == 1 + and (dcp_world_size == 1 or deinterleave_idx is not None) and not current_platform.is_xpu() ) if ( diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 0688c7b0ebca..7c0e8e4a0c0a 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1565,7 +1565,6 @@ def build( and current_platform.is_device_capability_family(100) and not self.indexer_uses_fp4 and self.dcp_world_size == 1 - and self.pcp_world_size == 1 and fused_mqa_topk_available() ): chunk_options["use_fused_mqa_topk"] = True From 289f49490987baa9a2fa028ef3165e5a120f3c4a Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 16 Oct 2026 08:56:24 +0000 Subject: [PATCH 07/23] set fused to default Signed-off-by: LopezCastroRoberto --- vllm/envs.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index f9896ee1ce36..bd57ebe8cd98 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -59,7 +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 = 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 @@ -1096,10 +1096,10 @@ def _resolve_rust_cli_path() -> str | None: "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 eligible # sparse MLA indexer prefill chunks on SM100-family GPUs (FP8 indexer cache, - # 32 index heads, top-k 2048, no DCP/PCP). Falls back to the dense path - # otherwise. Default: off. + # 32 index heads, top-k 512/1024/2048, no DCP/PCP). Falls back to the dense + # path otherwise. Default: on (set to 0 for the dense path). "VLLM_USE_FUSED_MQA_TOPK": lambda: bool( - int(os.getenv("VLLM_USE_FUSED_MQA_TOPK", "0")) + 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. From 55ef59dd09e06233bfb335e46e6b539430df6e5c Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Wed, 21 Oct 2026 16:05:20 +0000 Subject: [PATCH 08/23] tuning Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/persistent_epoch.inc | 70 ++++--- .../kernels/attention/test_fused_mqa_topk.py | 184 ++++++++++++++++++ .../kernels/attention/dsa/fused_mqa_topk.py | 10 + .../layers/sparse_attn_indexer.py | 22 +-- vllm/models/glm5next/nvidia/sparse_indexer.py | 135 +++++++++---- 5 files changed, 341 insertions(+), 80 deletions(-) diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc index 3b4fa49f74cc..407945c54c94 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -627,6 +627,7 @@ if (warp_idx == 8) { } float accum_pair[2][32]; [[maybe_unused]] float lin_pair[2][4]; + [[maybe_unused]] float bp4[4]; // > 16K: row scores of the slot #pragma unroll for (uint32_t pair = 0; pair < 4; pair += 2) { if constexpr (kLin) { @@ -681,36 +682,8 @@ if (warp_idx == 8) { if constexpr (kLin) { #pragma unroll for (uint32_t r = 0; r < 2; ++r) { - const uint32_t i = pair + r; const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - const float bp = (sum.x + sum.y) * scale_kv; - if (bp > gate_score[i]) { // gate first; range / finiteness only - // for passing keys - const float bq = bp * row_s[i]; - const uint32_t row = row_base2 + 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 = - fused_candidates::candidate_fp24_code(bq); - hdata.records[produce_epoch & 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(counts + row, 1); - if (slot < static_cast(kCandGlobalCap)) { - const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - kCandGlobalCap + - slot; - fused_candidates::store_candidate( - vals + off, indices + off, bq, kv_offset); - } - } - } - } + bp4[pair + r] = (sum.x + sum.y) * scale_kv; } } else #pragma unroll @@ -800,6 +773,45 @@ if (warp_idx == 8) { } } } + if constexpr (kLin) { + // 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(0xffffffffu, any_pass)) [[unlikely]] { +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const float bp = bp4[i]; + if (bp > gate_score[i]) { // gate first; range / finiteness only + // for passing keys + const float bq = bp * row_s[i]; + const uint32_t row = row_base2 + 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 = + fused_candidates::candidate_fp24_code(bq); + hdata.records[produce_epoch & 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(counts + row, 1); + if (slot < static_cast(kCandGlobalCap)) { + const uint64_t off = + static_cast(int64_t(row) + cand_row_delta) * + kCandGlobalCap + + slot; + fused_candidates::store_candidate( + vals + off, indices + off, bq, kv_offset); + } + } + } + } + } + } + } } if (!kSmem && (((kv_split_idx + 1) & (Interval - 1)) == 0 || diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 154d570b4865..c61732c377cc 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -332,3 +332,187 @@ def test_planner_keeps_budget_beyond_fused_key_limit(): LOGITS_BUDGET, ) assert fused == dense and len(dense) == 4 + + +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, fused_fn=None, budget_mb=512): + """Run the GLM-5.3-Flash kpool indexer prefill on one packed chunk.""" + from types import SimpleNamespace + + import vllm.envs as envs + import vllm.models.glm5next.nvidia.sparse_indexer as kpool_indexer + from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerPrefillChunkMetadata, + DeepseekV32IndexerPrefillMetadata, + ) + + 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) + monkeypatch.setattr(envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", budget_mb) + calls: list[torch.Size] = [] + fused_fn = fused_fn or kpool_indexer.fused_mqa_topk_prefill + + def spy(*a, **kw): + calls.append(a[-1].shape) # the output buffer: [rows, top-k pools] + return 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 +def test_kpool_indexer_fused_matches_dense(monkeypatch): + """GLM-5.3-Flash kpool indexer: the fused kernel ranks the top-512 pools of + a packed multi-request chunk like the dense logits + top-k path.""" + inputs = _kpool_inputs() + fused, calls = _run_kpool_indexer(monkeypatch, inputs, fused=True) + assert calls == [torch.Size([inputs[0].shape[0], TOPK // KPOOL])] + dense, calls = _run_kpool_indexer(monkeypatch, inputs, fused=False) + assert calls == [] + assert ((fused >= 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), _row_sets(dense))) + assert overlap >= 0.98 * int((dense >= 0).sum()) + + +@requires_fused +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 kernel is not faster there).""" + from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + fused_mqa_topk_min_keys, + ) + + 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=True) + assert calls == [] + + +@requires_fused +def test_kpool_indexer_dense_fallback_splits_rows(monkeypatch): + """Fused-planned chunks skip the logits budget; if the kernel declines, + the dense fallback splits rows to stay within it and selects the same + pools as one unbudgeted dense pass.""" + inputs = _kpool_inputs() + whole, _ = _run_kpool_indexer(monkeypatch, inputs, fused=False) + split, calls = _run_kpool_indexer( + monkeypatch, inputs, fused=True, fused_fn=lambda *a, **kw: False, budget_mb=1 + ) + assert len(calls) == 1 + # 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. + 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 + + 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/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 849a1b72e6b6..94f8f4f98662 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -30,6 +30,16 @@ FUSED_MQA_TOPK_MAX_ROWS = 16384 FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 + +def fused_mqa_topk_min_keys(topk: int) -> int: + """Keys per row below which the dense logits + top-k path is as fast or + faster, so callers keep prefill chunks whose longest row sees fewer keys on + the dense path. B300 real captures, fused vs dense: top-512 0.89x at 16K + keys and 1.06x at 32K; top-1024 0.93x / 1.01x; top-2048 0.92x at 32K and + 1.13x at 64K.""" + return 1 << 16 if topk > 1024 else 1 << 15 + + _CANDIDATES_PER_ROW = 16384 _RETAINED_PER_ROW = 4096 _SLOTS_PER_SM = 2 diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 011b1207c8b5..b8b1d0a020cc 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -30,9 +30,9 @@ select_candidate_blocks as _select_candidate_blocks, ) from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( - FUSED_MQA_TOPK_MAX_KEYS, fused_mqa_topk_config_supported, fused_mqa_topk_key_scales, + fused_mqa_topk_min_keys, fused_mqa_topk_prefill, get_fused_mqa_topk_workspace, ) @@ -519,15 +519,13 @@ def sparse_attn_indexer( get_fused_mqa_topk_workspace(q_quant.device) # Dummy allocation to simulate for peak logits tensor memory during inference. - # Skipped when the fused kernel serves every prefill chunk: the planner never - # packs a chunk past its key limit, so with max_model_len within that limit - # no chunk reaches the dense logits path. - if not (uses_fused and max_model_len <= FUSED_MQA_TOPK_MAX_KEYS): - # FP8 elements so elements == bytes - max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 - _ = torch.empty( - max_logits_elems, dtype=torch.uint8, device=hidden_states.device - ) + # Needed with the fused kernel too: chunks below fused_mqa_topk_min_keys + # keys per row take the dense path (row-split to stay within the budget). + # FP8 elements so elements == bytes + max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + _ = torch.empty( + max_logits_elems, dtype=torch.uint8, device=hidden_states.device + ) # The dummy run returns before exchange_topk executes, so reserve its # worst-case transient for the memory profile. @@ -731,7 +729,8 @@ def sparse_attn_indexer( torch.int32 if mqa_logits_use_mx_scales() else torch.float32 ).squeeze(-1) # Single-launch fused MQA-logits + top-k (SM100, FP8, 32 heads, - # top-k 512/1024/2048): no logits buffer, any context length. + # top-k 512/1024/2048): no logits buffer, for chunks whose rows + # see at least fused_mqa_topk_min_keys keys (dense is as fast below). # It needs every row's whole context: no DCP, or PCP + DCP after # the de-interleaving gather. Pure DCP (per-shard top-k + merge), # the FP4 cache and v4.1 candidate selection stay dense. @@ -742,6 +741,7 @@ def sparse_attn_indexer( ) if ( use_fused + and chunk.max_seqlen_k >= fused_mqa_topk_min_keys(topk_tokens) and q_scale_slice is None and candidate_blocks is None and fused_mqa_topk_prefill( diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 65ba7be1fdec..868b64f90915 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_config_supported, + fused_mqa_topk_key_scales, + fused_mqa_topk_min_keys, + fused_mqa_topk_prefill, + get_fused_mqa_topk_workspace, +) from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.model_executor.layers.sparse_attn_indexer import ( PrefillRowShard, @@ -155,6 +162,17 @@ def sparse_attn_indexer_kpool( scales_spec, ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) + # The fused prefill kernel's fixed-size workspace is persistent; + # allocate it during profiling so it is accounted for. + if ( + envs.VLLM_USE_FUSED_MQA_TOPK + and dcp_world_size == 1 + and not use_fp4_cache + and fused_mqa_topk_config_supported( + q_quant, topk_tokens // max(index_kpool, 1) + ) + ): + get_fused_mqa_topk_workspace(q_quant.device) # Reserve profiler-visible memory for the worst-case decode logits, # whose shape is [B * next_n, max_pool_len]. This profiling branch @@ -355,61 +373,98 @@ 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: + # Single-launch fused MQA-logits + top-k over the pools (SM100, FP8, + # 32 heads, top-k 512/1024/2048): no logits buffer, for chunks whose + # rows see at least fused_mqa_topk_min_keys(select_k) pools. It needs every + # row's whole context, so DCP (per-shard top-k + merge) stays dense. + use_fused = envs.VLLM_USE_FUSED_MQA_TOPK and dcp_world_size == 1 + if chunk.local_total_seq_lens == 0: topk_dst.fill_(-1) - else: - torch.ops._C.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends), - row_ends, + _merge_dcp_topk_global( + q_slice.new_empty((num_rows, 0), dtype=torch.float32), topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), select_k, + dcp_rank, + dcp_world_size, + pool_cp_interleave, + row_starts=None, ) - # Local pool ids -> request-relative global pool ids. - _merge_dcp_topk_global( - logits, - topk_dst, - select_k, - dcp_rank, - dcp_world_size, - pool_cp_interleave, - row_starts=None, - ) + elif not ( + use_fused + and chunk.max_seqlen_k >= fused_mqa_topk_min_keys(select_k) + and q_scale_slice is None + and fused_mqa_topk_prefill( + q_slice_cast, + k_quant_cast, + fused_mqa_topk_key_scales(k_scale_cast), + weights[row_start:row_end], + cu_seqlen_ks, + cu_seqlen_ke, + topk_dst, + # A single-request chunk has every row start at 0. + rows_start_at_zero=chunk.block_table.shape[0] == 1, + ) + ): + # Fused-planned chunks are not bounded by the logits budget, so + # the dense fallback splits their rows to stay within it. + dense_rows = max(1, num_rows) + if use_fused: + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + dense_rows = max(1, budget // (4 * max(1, chunk.max_seqlen_k))) + for sub in range(0, max(1, num_rows), dense_rows): + rows = slice(sub, min(sub + dense_rows, num_rows)) + sub_ks, sub_ke = cu_seqlen_ks[rows], cu_seqlen_ke[rows] + sub_q_scale, sub_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 = sub_ke - sub_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast[rows], sub_q_scale), + (k_quant_cast, k_scale_cast), + sub_weights, + sub_ks, + sub_ke, + chunk.max_seqlen_k, + ), + row_ends, + ) + sub_dst = topk_dst[rows] + torch.ops._C.top_k_per_row_prefill( + logits, + torch.zeros_like(row_ends), + row_ends, + sub_dst, + logits.shape[0], + logits.stride(0), + logits.stride(1), + select_k, + ) + # Local pool ids -> request-relative global pool ids. + _merge_dcp_topk_global( + logits, + sub_dst, + select_k, + dcp_rank, + dcp_world_size, + pool_cp_interleave, + row_starts=None, + ) if index_kpool > 1: pool_ids = pool_topk.to(torch.int64) From c36250bd70d29345d51e73d1bb0e45dbbb4a88fd Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Thu, 22 Oct 2026 16:23:28 +0000 Subject: [PATCH 09/23] refine Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/fused_mqa_topk.cu | 4 +- .../kernels/attention/test_fused_mqa_topk.py | 50 ++++++++++++++----- .../kernels/attention/dsa/fused_mqa_topk.py | 36 +++++++++---- .../layers/sparse_attn_indexer.py | 2 +- vllm/models/glm5next/nvidia/sparse_indexer.py | 2 +- vllm/v1/attention/backends/mla/indexer.py | 18 ++++--- 6 files changed, 80 insertions(+), 32 deletions(-) diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu index 38c6e2d5d6f1..489915c851e1 100644 --- a/csrc/fused_mqa_topk/fused_mqa_topk.cu +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -121,7 +121,7 @@ void check_tensor(const torch::stable::Tensor& t, int32_t device, // 2048: absolute key indices (columns of the logits), or with relative=true // indices relative to the row's start (key - start, the top_k_per_row_prefill // convention); -1 padded when fewer than K keys are visible. -// 0 < rows <= 16384, 0 < keys <= 1048576 (the last 8-row block may be partial). +// 0 < rows, 0 < keys <= 1048576 (the last 8-row block may be partial). // 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), @@ -158,7 +158,7 @@ void fused_mqa_topk_prefill( 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); - STD_TORCH_CHECK(rows > 0 && rows <= 16384, "rows must be in [1, 16384]"); + STD_TORCH_CHECK(rows > 0, "rows must be positive"); STD_TORCH_CHECK(keys > 0 && keys <= (1 << 20), "keys must be in [1, 1048576]"); const int64_t blocks = (rows + 7) / 8; diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index c61732c377cc..aa021ce7289f 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -7,6 +7,7 @@ from vllm import _custom_ops as ops from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + FUSED_MQA_TOPK_MAX_ROWS, fused_mqa_topk_available, fused_mqa_topk_prefill, ) @@ -209,6 +210,25 @@ def test_multi_request_chunk_matches_dense(lens, topk): assert overlap >= 0.98 * int((dense >= 0).sum()) +@requires_fused +@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; the per-row workspace grows to fit.""" + from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + get_fused_mqa_topk_workspace, + ) + + 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) + assert fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + assert get_fused_mqa_topk_workspace(q.device, rows).counts.numel() >= rows + _check(out, q, k, scales, weights, starts, ends) + + @requires_fused def test_rows_start_at_zero_skips_conversion(): rows, keys = 32, 8192 @@ -271,31 +291,37 @@ def _split(seq_lens, q_lens, fused, offset=0): ) -@pytest.mark.parametrize("fused,expected", [(False, 64), (True, 1)]) -def test_planner_one_fused_call_per_16k_rows(fused, expected): - chunks = _split([524288], [16384], fused) - assert len(chunks) == expected - assert sum(q.stop - q.start for _, q in chunks) == 16384 +@pytest.mark.parametrize("q_rows", [16384, 32768]) +@pytest.mark.parametrize("fused", [False, True]) +def test_planner_one_fused_call_per_chunk(fused, q_rows): + # One fused call per FUSED_MQA_TOPK_MAX_ROWS rows (None: per chunk); dense + # splits by the logits budget. + chunks = _split([524288], [q_rows], fused) + cap = FUSED_MQA_TOPK_MAX_ROWS or q_rows + assert len(chunks) == (-(-q_rows // cap) if fused else q_rows // 256) + assert sum(q.stop - q.start for _, q in chunks) == q_rows def test_planner_fused_any_length_and_short_queries(): # No length threshold and no minimum query count on the fused path. assert _split([524288], [512], True) == [(slice(0, 1), slice(0, 512))] assert _split([4096], [4096], True) == [(slice(0, 1), slice(0, 4096))] - # Long single-request steps are split at 16384 query rows. + # A 40000-row step: split at FUSED_MQA_TOPK_MAX_ROWS (None: one call). + cap = FUSED_MQA_TOPK_MAX_ROWS or 40000 assert _split([65536], [40000], True) == [ - (slice(0, 1), slice(0, 16384)), - (slice(0, 1), slice(16384, 32768)), - (slice(0, 1), slice(32768, 40000)), + (slice(0, 1), slice(lo, min(lo + cap, 40000))) for lo in range(0, 40000, cap) ] -def test_planner_packs_requests_up_to_16k_rows(): - # Four 4K requests fit one fused call (16384 rows); dense needs a split. +def test_planner_packs_requests_in_one_fused_call(): + # Requests share one fused call up to the key limit; dense needs a split. assert _split([4096] * 4, [4096] * 4, True) == [(slice(0, 4), slice(0, 16384))] assert len(_split([4096] * 4, [4096] * 4, False)) > 1 chunks = _split([524288, 262144], [16384, 9], True, offset=2) - assert chunks == [(slice(2, 3), slice(0, 16384)), (slice(3, 4), slice(0, 9))] + if FUSED_MQA_TOPK_MAX_ROWS is None or FUSED_MQA_TOPK_MAX_ROWS >= 16393: + assert chunks == [(slice(2, 4), slice(0, 16393))] + else: # the second request would push the call past the row cap + assert chunks == [(slice(2, 3), slice(0, 16384)), (slice(3, 4), slice(0, 9))] def test_planner_fused_never_packs_past_key_limit(): diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 94f8f4f98662..23c72bdd003d 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -26,9 +26,15 @@ FUSED_MQA_TOPK_HEADS = 32 FUSED_MQA_TOPK_HEAD_DIM = 128 FUSED_MQA_TOPK_TOPKS = (512, 1024, 2048) -# Largest query chunk / gathered key range a single call accepts. -FUSED_MQA_TOPK_MAX_ROWS = 16384 +# Largest gathered key range a single call accepts (no row limit: the per-row +# workspace is sized from the call, see get_fused_mqa_topk_workspace). FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 +# Rows per fused call. The kernel has no row limit, but prefill chunks above +# 16384 rows (--max-num-batched-tokens > 16384) are not validated end to end, +# so the chunk planner splits fused chunks at this many rows. None: one fused +# call per chunk, any size. +# TODO: check for other prefill chunk sizes. +FUSED_MQA_TOPK_MAX_ROWS: int | None = 16384 def fused_mqa_topk_min_keys(topk: int) -> int: @@ -54,8 +60,8 @@ class FusedMQATopKWorkspace: indices: torch.Tensor # [8S, 16384] int32 retained: torch.Tensor # [8S, 4096] fp32 retained_indices: torch.Tensor # [8S, 4096] int32 - counts: torch.Tensor # [16384] int32, per-row candidate counts - defer: torch.Tensor # [16384 / 8] int32, blocks queued for the exact rescan + counts: torch.Tensor # [rows] int32, per-row candidate counts + defer: torch.Tensor # [rows / 8] int32, blocks queued for the exact rescan def as_args(self) -> tuple[torch.Tensor, ...]: return ( @@ -81,11 +87,17 @@ def fused_mqa_topk_available() -> bool: ) -def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: +def get_fused_mqa_topk_workspace( + device: torch.device, rows: int +) -> FusedMQATopKWorkspace: """Return the per-device workspace, allocating it on first use. - The block counter is zeroed once; every launch leaves it zeroed again. + The per-row part (counts, defer) holds at least ``rows`` query rows and grows + when a larger call arrives; the profiling run asks for the largest prefill + chunk (max_num_batched_tokens rows) so serving never grows it. The block + counter is zeroed once; every launch leaves it zeroed again. """ + rows = (rows + 7) // 8 * 8 index = ( device.index if device.index is not None @@ -96,7 +108,6 @@ def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: dev = torch.device("cuda", index) num_sms = torch.cuda.get_device_properties(index).multi_processor_count slots = _SLOTS_PER_SM * num_sms - rows = FUSED_MQA_TOPK_MAX_ROWS ws = FusedMQATopKWorkspace( counter=torch.zeros(2, dtype=torch.int32, device=dev), values=torch.empty( @@ -122,6 +133,13 @@ def get_fused_mqa_topk_workspace(device: torch.device) -> FusedMQATopKWorkspace: nbytes / 2**20, ) _WORKSPACES[index] = ws + elif ws.counts.numel() < rows: + ws.counts = torch.empty(rows, dtype=torch.int32, device=ws.counts.device) + ws.defer = torch.empty(rows // 8, dtype=torch.int32, device=ws.defer.device) + logger.info( + "Grew the fused MQA+top-k row workspace on cuda:%d to %d rows.", index, rows + ) + _WORKSPACES[index] = ws return ws @@ -145,7 +163,7 @@ def fused_mqa_topk_supported( k.dtype == torch.float8_e4m3fn and k.dim() == 2 and k.shape[1] == FUSED_MQA_TOPK_HEAD_DIM - and 0 < q.shape[0] <= FUSED_MQA_TOPK_MAX_ROWS + and q.shape[0] > 0 and 0 < k.shape[0] <= FUSED_MQA_TOPK_MAX_KEYS and fused_mqa_topk_config_supported(q, topk_tokens) ) @@ -225,7 +243,7 @@ def fused_mqa_topk_prefill( if direct else torch.empty_like(output, memory_format=torch.contiguous_format) ) - ws = get_fused_mqa_topk_workspace(q.device) + ws = get_fused_mqa_topk_workspace(q.device, rows) # The kernel emits absolute key indices; the dense top-k (and the sparse # attention consumer) use indices relative to the row start. relative=True # converts in place (rows with start 0 are skipped on the GPU). diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index b8b1d0a020cc..e1646f4a7ccb 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -516,7 +516,7 @@ def sparse_attn_indexer( if uses_fused: # The fused prefill kernel's fixed-size workspace is persistent; # allocate it during profiling so it is accounted for. - get_fused_mqa_topk_workspace(q_quant.device) + get_fused_mqa_topk_workspace(q_quant.device, q_quant.shape[0]) # Dummy allocation to simulate for peak logits tensor memory during inference. # Needed with the fused kernel too: chunks below fused_mqa_topk_min_keys diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 868b64f90915..f95f64884487 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -172,7 +172,7 @@ def sparse_attn_indexer_kpool( q_quant, topk_tokens // max(index_kpool, 1) ) ): - get_fused_mqa_topk_workspace(q_quant.device) + get_fused_mqa_topk_workspace(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 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 7c0e8e4a0c0a..1f98924ca34f 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1378,11 +1378,12 @@ def _split_indexer_prefill_chunks( sub-chunks on the query dimension (M) to bound peak memory. With ``use_fused_mqa_topk`` the fused MQA+top-k kernel needs no logits - buffer: chunks whose gathered keys fit its 1M-key limit are bounded by - its 16384-row limit instead of the logits budget (one fused call per - chunk). Requests are never packed past that key limit, so only a single - request longer than 1M keys takes the dense, logits-budget path (the - profiling run relies on this to skip the dense logits reservation). + buffer: chunks whose gathered keys fit its 1M-key limit are not bounded + by the logits budget but by FUSED_MQA_TOPK_MAX_ROWS rows (one fused call + per chunk; None = no row limit). Requests are never packed past that key + limit, so only a single request longer than 1M keys takes the + logits-budget path. Chunks the dispatch keeps dense are row-split by the + indexer. Returns list of (req_slice, query_slice) tuples. """ @@ -1405,7 +1406,10 @@ def fused(n: int) -> bool: ) new_m, new_n = chunk_m + q, chunk_n + s if fused(new_n): - fits = new_m <= FUSED_MQA_TOPK_MAX_ROWS + fits = ( + FUSED_MQA_TOPK_MAX_ROWS is None + or new_m <= FUSED_MQA_TOPK_MAX_ROWS + ) else: fits = not use_fused_mqa_topk and new_m * new_n <= max_logits_elems if new_n <= workspace_size and fits: @@ -1428,7 +1432,7 @@ def fused(n: int) -> bool: max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) ) if fused(chunk_n): - max_q = FUSED_MQA_TOPK_MAX_ROWS + max_q = FUSED_MQA_TOPK_MAX_ROWS or max(1, chunk_m) 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))) From a3d561521860e484a2c926bf2bfb5297a3888425 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 13:19:00 +0000 Subject: [PATCH 10/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 4 +- csrc/fused_mqa_topk/candidates.cuh | 35 +- csrc/fused_mqa_topk/common.cuh | 116 ++ csrc/fused_mqa_topk/constants.h | 33 + csrc/fused_mqa_topk/deep_gemm_compat.cuh | 62 +- csrc/fused_mqa_topk/dense_scan.cuh | 114 +- csrc/fused_mqa_topk/fallback.cuh | 361 ----- csrc/fused_mqa_topk/fast_select.cuh | 449 ------- csrc/fused_mqa_topk/fused_mqa_topk.cu | 367 +++-- csrc/fused_mqa_topk/persistent_epoch.inc | 1190 ++++++----------- csrc/fused_mqa_topk/persistent_kernel.cuh | 667 +++++---- csrc/fused_mqa_topk/rescan.cuh | 338 +++++ csrc/fused_mqa_topk/resident_kernel.cuh | 136 -- csrc/fused_mqa_topk/row_select.cuh | 415 ------ csrc/fused_mqa_topk/sample.cuh | 25 - csrc/fused_mqa_topk/select.cuh | 438 ++++++ csrc/ops.h | 2 +- csrc/torch_bindings.cpp | 3 +- .../kernels/attention/test_fused_mqa_topk.py | 232 ++-- .../test_indexer_fused_mqa_topk_planner.py | 135 ++ vllm/envs.py | 8 +- .../kernels/attention/dsa/fused_mqa_topk.py | 190 ++- .../layers/sparse_attn_indexer.py | 308 +++-- vllm/models/deepseek_v4/cpu/cpu_mla.py | 4 +- vllm/models/glm5next/nvidia/sparse_indexer.py | 111 +- vllm/v1/attention/backends/mla/indexer.py | 52 +- 26 files changed, 2525 insertions(+), 3270 deletions(-) create mode 100644 csrc/fused_mqa_topk/common.cuh create mode 100644 csrc/fused_mqa_topk/constants.h delete mode 100644 csrc/fused_mqa_topk/fallback.cuh delete mode 100644 csrc/fused_mqa_topk/fast_select.cuh create mode 100644 csrc/fused_mqa_topk/rescan.cuh delete mode 100644 csrc/fused_mqa_topk/resident_kernel.cuh delete mode 100644 csrc/fused_mqa_topk/row_select.cuh delete mode 100644 csrc/fused_mqa_topk/sample.cuh create mode 100644 csrc/fused_mqa_topk/select.cuh create mode 100644 tests/v1/attention/test_indexer_fused_mqa_topk_planner.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 8922c022997c..01077db83370 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1705,14 +1705,12 @@ if (VLLM_GPU_LANG STREQUAL "CUDA") 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) + # Single-launch fused MQA-logits + top-k prefill kernel. 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") - # Single-launch fused MQA-logits + top-k prefill kernel. - set_property(SOURCE ${FUSED_MQA_SRC} APPEND PROPERTY COMPILE_OPTIONS - "-O3" "--expt-relaxed-constexpr" "--expt-extended-lambda") target_sources(_C_stable_libtorch PRIVATE ${FUSED_MQA_SRC}) target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_MQA_TOPK=1) endif() diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh index d19ca2496201..df022668af90 100644 --- a/csrc/fused_mqa_topk/candidates.cuh +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -3,10 +3,15 @@ #pragma once #include #include -namespace fused_candidates { + +#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) { @@ -14,13 +19,8 @@ CUTLASS_DEVICE uint32_t candidate_pack_index(const uint32_t payload, ((payload >> 16) << kCandidateIndexBits); } -CUTLASS_DEVICE uint32_t candidate_ordered_fp32_code(const float value) { - const uint32_t bits = __float_as_uint(value); - return (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); -} - CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) { - return candidate_ordered_fp32_code(value) >> 8; + return ordered_code(value) >> 8; } CUTLASS_DEVICE uint32_t candidate_load_score_code(const CandidateValue value, @@ -34,16 +34,8 @@ CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { kCandidateIndexMask); } -CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, - int32_t* index_dst, const float bq, - const uint32_t kv_index) { - const uint32_t payload = candidate_fp24_code(bq); - __stcg(value_dst, static_cast(payload)); - __stcg(index_dst, - static_cast(candidate_pack_index(payload, kv_index))); -} - -// Stores a candidate whose fp24 code was already computed (ring records). +// 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, @@ -53,4 +45,11 @@ CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst, static_cast(candidate_pack_index(payload, kv_index))); } -} // namespace fused_candidates +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..dfb37eb3877b --- /dev/null +++ b/csrc/fused_mqa_topk/common.cuh @@ -0,0 +1,116 @@ +// 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; +// 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; +} + +// 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(0xffffffffu, 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"); +} + +// 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. +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 index 2d7559ec634a..107e3c741def 100644 --- a/csrc/fused_mqa_topk/deep_gemm_compat.cuh +++ b/csrc/fused_mqa_topk/deep_gemm_compat.cuh @@ -4,8 +4,8 @@ // (github.com/deepseek-ai/DeepGEMM) as of the release before its SM100 kernels // went MX-only: 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 (no metadata schedule, no -// multicast). +// Trimmed to what the fused kernel instantiates (FP8 only, no metadata +// schedule, no multicast, no paging). #pragma once #include @@ -16,72 +16,46 @@ #include #include -namespace fused_mqa_dg { +namespace fused_mqa_topk::dg { namespace math = deep_gemm::math; -template struct MQALogitsSharedStorage { - static constexpr bool kIsFP4 = - cute::is_same_v; - using Barrier = cutlass::arch::ClusterTransactionBarrier; - using sf_dtype_t = cute::conditional_t; + DG_STATIC_ASSERT(sizeof(qk_dtype_t) == 1, "FP8 Q/K only"); - static constexpr uint32_t kNumUTCCPAlignedElems = 128; - static constexpr uint32_t kQKBytesPerElem = sizeof(qk_dtype_t); - static constexpr uint32_t kNumQKBytesPerToken = - kIsFP4 ? (kHeadDim / 2) : kHeadDim; 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: FP4 uses head_dim / 2 bytes, FP8 use - // head_dim bytes per token - static constexpr uint32_t kSwizzleAlignment = 8 * kNumQKBytesPerToken; - static constexpr uint32_t kNumSFQ = - math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems); - static constexpr uint32_t kNumSFKV = - math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems); - static constexpr uint32_t kNumQBytesPerStage = UMMA_N * kNumQKBytesPerToken; - static constexpr uint32_t kNumKVBytesPerStage = - SPLIT_KV * kNumQKBytesPerToken; - static constexpr uint32_t kNumQElementsPerStage = - kNumQBytesPerStage / kQKBytesPerElem; - static constexpr uint32_t kNumKVElementsPerStage = - kNumKVBytesPerStage / kQKBytesPerElem; - // MX SF formats store per-block scale factors; FP8 stores one per-KV scale - // and no Q scale - static constexpr uint32_t kNumScaleQ = kIsMXSF ? kNumSFQ : 1; - static constexpr uint32_t kNumScaleKV = kIsMXSF ? kNumSFKV : SPLIT_KV; + // 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 kNumWeightElementsPerRow = - kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); 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(kNumQBytesPerStage % kSwizzleAlignment == 0, + DG_STATIC_ASSERT(kNumQElementsPerStage % kSwizzleAlignment == 0, "Unaligned TMA swizzling"); - DG_STATIC_ASSERT(kNumKVBytesPerStage % kSwizzleAlignment == 0, + DG_STATIC_ASSERT(kNumKVElementsPerStage % kSwizzleAlignment == 0, "Unaligned TMA swizzling"); DG_STATIC_ASSERT(kSwizzleAlignment % 128 == 0, "TMA destination must be 128-byte aligned"); - DG_STATIC_ASSERT(kTmaAlignment % 128 == 0, - "TMA destination must be 128-byte aligned"); - DG_STATIC_ASSERT(kNumWeightBytesPerStage % kTmaAlignment == 0, - "Unaligned weight TMA stage"); alignas(kSwizzleAlignment) qk_dtype_t smem_q[kNumQStages][kNumQElementsPerStage]; alignas(kSwizzleAlignment) qk_dtype_t smem_kv[kNumKVStages][kNumKVElementsPerStage]; - alignas(kTmaAlignment) sf_dtype_t smem_sf_q[kNumQStages][kNumScaleQ]; - alignas(kTmaAlignment) sf_dtype_t smem_sf_kv[kNumKVStages][kNumScaleKV]; + 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 @@ -99,10 +73,6 @@ struct MQALogitsSharedStorage { // ..., each visiting the KV range [min k_start, max k_end) of its rows. template struct SM100MQALogitsScheduler { - static constexpr bool kIsPaged = false; - static constexpr bool kHasPartialBlock = false; - static constexpr uint32_t kPageKV = 0; - uint32_t current_q_block_idx; uint32_t num_q_blocks; uint32_t num_q_tokens; @@ -175,7 +145,7 @@ template (cute::TMA::CacheHintSm90::EVICT_NORMAL); constexpr uint32_t kPackFactor = deep_gemm::get_smem_pack_factor(); @@ -205,4 +175,4 @@ CUTLASS_DEVICE void tma_copy( } } -} // namespace fused_mqa_dg +} // namespace fused_mqa_topk::dg diff --git a/csrc/fused_mqa_topk/dense_scan.cuh b/csrc/fused_mqa_topk/dense_scan.cuh index 069bb4b8bf91..f0b2e2279f53 100644 --- a/csrc/fused_mqa_topk/dense_scan.cuh +++ b/csrc/fused_mqa_topk/dense_scan.cuh @@ -2,21 +2,22 @@ // Copyright (c) 2025 DeepSeek // 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 <= 4K resident kernel and -// by the persistent kernel's exact recovery rescan. +// 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, kScoreScale score = (sum_h kWeightScale * w_h * (s_h + -// |s_h|)) * kScoreScale * scale_kv begin_block(sched, q_block, query_base) -// after the block's Q / weights landed begin_split(split, query_base) 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 +// 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 @@ -31,19 +32,27 @@ #include #include +#include "common.cuh" +#include "constants.h" #include "deep_gemm_compat.cuh" -namespace fused_mqa_dense { +namespace fused_mqa_topk::dense { using namespace deep_gemm; -constexpr uint32_t kNumHeads = 32, kHeadDim = 128, BLOCK_Q = 8, SPLIT_KV = 256; +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 = - fused_mqa_dg::SM100MQALogitsScheduler; -using SharedStorage = fused_mqa_dg::MQALogitsSharedStorage< - kNumHeads, kHeadDim, false, BLOCK_Q, SPLIT_KV, 1, kNumKVStages, - kNumTmemStages, cutlass::float_e4m3_t, float>; + 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 @@ -81,16 +90,16 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, barrier.init(count); }; init(smem.full_q_barriers[0], 1); - init(smem.empty_q_barriers[0], 256 + 32); + 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], 256); + 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], 256); + init(smem.empty_tmem_barriers[i], kMathThreads); } cutlass::arch::fence_barrier_init(); } @@ -102,11 +111,10 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, RingPipeline<1> q_pipeline; RingPipeline kv_pipeline; RingPipeline tmem_pipeline; - cudaGridDependencySynchronize(); uint32_t q_block_idx, kv_base, num_kv_splits; if (warp_idx == kSpecWarpStart) { - cutlass::arch::warpgroup_reg_dealloc<56>(); + 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)) { @@ -114,10 +122,10 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, 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); - fused_mqa_dg::tma_copy( + dg::tma_copy( &tensor_map_q, &smem.full_q_barriers[q_stage_idx], smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); - fused_mqa_dg::tma_copy( + 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( @@ -126,7 +134,7 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, } __syncwarp(); } else if (warp_idx == kSpecWarpStart + 1) { - cutlass::arch::warpgroup_reg_dealloc<56>(); + cutlass::arch::warpgroup_reg_dealloc(); Scheduler scheduler = block; while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { #pragma unroll 1 @@ -137,10 +145,10 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, 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); - fused_mqa_dg::tma_copy( + dg::tma_copy( &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); - fused_mqa_dg::tma_copy( + 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( @@ -150,7 +158,7 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, } } } else if (warp_idx == kSpecWarpStart + 2) { - cutlass::arch::warpgroup_reg_dealloc<56>(); + 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, @@ -192,9 +200,9 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, smem.empty_q_barriers[q_stage_idx].arrive(); } } else if (warp_idx == kSpecWarpStart + 3) { - cutlass::arch::warpgroup_reg_dealloc<56>(); + cutlass::arch::warpgroup_reg_dealloc(); } else { - cutlass::arch::warpgroup_reg_alloc<224>(); + 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]; @@ -212,12 +220,13 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, 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, query_base); + 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 % 128) + kv_half * 128; + 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 = @@ -242,32 +251,12 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, ptx::tcgen05_before_thread_sync(); smem.empty_tmem_barriers[tmem_stage_idx].arrive(); } - // relu(s) = (s + |s|) / 2: two FFMA2 chains per row. - float2 sum_0[2] = {}; - float2 sum_1[2] = {}; #pragma unroll - for (uint32_t j = 0; j < kNumHeads; j += 4) { -#pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - const float* a = accum_pair[r]; - auto transform = [&](uint32_t at, float2 sum) { - const auto x = make_float2(a[at], a[at + 1]); - const auto abs_x = - make_float2(fabsf(a[at]), fabsf(a[at + 1])); - const auto w = make_float2(weights[pair + r][at], - weights[pair + r][at + 1]); - return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); - }; - sum_0[r] = transform(j, sum_0[r]); - sum_1[r] = transform(j + 2, sum_1[r]); - } - } -#pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - policy.score(pair + r, kv_offset, kv_lane, kv_split_idx, - (sum.x + sum.y) * (Policy::kScoreScale * scale_kv)); - } + 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); @@ -276,9 +265,20 @@ CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, cutlass::arch::fence_view_async_shared(); smem.empty_q_barriers[q_stage_idx].arrive(); } - cutlass::arch::NamedBarrier(256, 0).sync(); + cutlass::arch::NamedBarrier(kMathThreads, kBarDenseMath).sync(); if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); } } -} // namespace fused_mqa_dense +// 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/fallback.cuh b/csrc/fused_mqa_topk/fallback.cuh deleted file mode 100644 index 824cc937f79c..000000000000 --- a/csrc/fused_mqa_topk/fallback.cuh +++ /dev/null @@ -1,361 +0,0 @@ -// SPDX-License-Identifier: MIT -// Exact recovery for blocks whose gated candidate count left [K, 16384] (or -// whose weights are too small to halve exactly): the block is rescanned -// densely, and every 1024-key segment of scores is merged into a running -// per-row top set (kept <= FallbackRetain, compacted to K). -#pragma once -#include -#include "dense_scan.cuh" -#include -#include -#include -#include -constexpr unsigned FallbackSegment = 1024, FallbackRetain = 4096; -// output: [rows, K] int32 (K = the row top-k). -struct FallbackBuffers { - float* retained; - int* indices; - int* output; -}; -struct alignas(16) FallbackShared { - float ring[2][8][FallbackSegment]; - unsigned hist[8][1056], groups[8][32], gates[8], kept[8]; - unsigned work[264], prefix, rank, written, added, emit_counts[8]; - __device__ void init(unsigned tid) { - for (unsigned j = tid; j < 8 * 1056; j += 384) - reinterpret_cast(hist)[j] = 0; - if (tid < 256) reinterpret_cast(groups)[tid] = 0; - if (tid < 8) { - gates[tid] = __float_as_uint(-CUDART_INF_F); - kept[tid] = 0; - } - } -}; -__device__ __forceinline__ unsigned fb_code(float x) { - unsigned u = __float_as_uint(x); - return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); -} -__device__ __forceinline__ float fb_float(unsigned u) { - return __uint_as_float((u & 0x80000000u) ? (u ^ 0x80000000u) : ~u); -} -__device__ __forceinline__ unsigned fb_half_bin(float x) { - unsigned u = __half_as_ushort(__float2half_rd(x)); - return ((u & 0x8000u) ? ((~u) & 65535u) : (u ^ 0x8000u)) >> 6; -} -__device__ __forceinline__ float fb_bin_floor(unsigned b) { - unsigned o = b << 6, u = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); - return __half2float(__ushort_as_half(u)); -} -__device__ __forceinline__ void fb_sync() { - asm volatile("" ::: "memory"); - cutlass::arch::NamedBarrier(256, 2).sync(); - asm volatile("" ::: "memory"); -} -__device__ void fb_compact(FallbackShared& s, FallbackBuffers b, unsigned K, - unsigned local, unsigned row, unsigned tid, - bool final = false) { - constexpr unsigned Full = 0xffffffffu, T = 256, Items = FallbackRetain / T; - unsigned lane = tid & 31, count = s.kept[local]; - uint64_t base = (uint64_t)row * FallbackRetain; - float values[Items]; - int indices[Items]; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) { - unsigned j = tid + i * T; - values[i] = j < count ? b.retained[base + j] : -CUDART_INF_F; - indices[i] = j < count ? b.indices[base + j] : -1; - } - auto emit = [&](float gate, unsigned threshold, int mode, - unsigned base_offset) { - unsigned total = 0; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) { - float v = values[i]; - bool pass = - isfinite(v) && (mode == 0 ? v >= gate - : (mode == 1 ? fb_code(v) > threshold - : fb_code(v) == threshold)); - total += pass; - } - unsigned warp_total = __reduce_add_sync(Full, total); - if (lane == 0) s.emit_counts[tid / 32] = warp_total; - fb_sync(); - unsigned start = base_offset; - for (unsigned w = 0; w < tid / 32; ++w) start += s.emit_counts[w]; - unsigned all = 0; - for (unsigned w = 0; w < 8; ++w) all += s.emit_counts[w]; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) { - float v = values[i]; - bool pass = - isfinite(v) && (mode == 0 ? v >= gate - : (mode == 1 ? fb_code(v) > threshold - : fb_code(v) == threshold)); - unsigned mask = __ballot_sync(Full, pass), - pos = start + __popc(mask & ((1u << lane) - 1)); - if (pass && (mode != 2 || pos < K)) { - b.retained[base + pos] = v; - b.indices[base + pos] = indices[i]; - } - start += __popc(mask); - } - fb_sync(); - return all; - }; - if (!final) { - if (tid == 0) { - s.rank = 0; - s.written = 0; - } - fb_sync(); - float gate = __uint_as_float(s.gates[local]); - unsigned kept = 0; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) - kept += isfinite(values[i]) && values[i] >= gate; - kept = __reduce_add_sync(Full, kept); - if (lane == 0) atomicAdd(&s.rank, kept); - fb_sync(); - kept = s.rank; - if (kept >= K && kept <= FallbackRetain - FallbackSegment) { - emit(gate, 0, 0, 0); - fb_sync(); - if (tid == 0) s.kept[local] = kept; - fb_sync(); - return; - } - } - if (tid == 0) { - s.prefix = 0; - s.rank = K; - s.written = 0; - } - fb_sync(); - for (int shift = 24; shift >= 0; shift -= 8) { - for (unsigned j = tid; j < 264; j += T) s.work[j] = 0; - fb_sync(); - unsigned prefix = s.prefix; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) { - unsigned code = fb_code(values[i]); - bool valid = isfinite(values[i]) && - (shift == 24 || (code >> (shift + 8)) == prefix); - unsigned mask = __ballot_sync(Full, valid); - if (valid) { - unsigned bin = (code >> shift) & 255u, - peers = __match_any_sync(mask, bin); - if (lane == unsigned(__ffs(peers) - 1)) - atomicAdd(s.work + bin + bin / 32, __popc(peers)); - } - } - fb_sync(); - if (tid < 32) { - unsigned sum = 0; - for (unsigned j = 0; j < 8; ++j) { - unsigned bin = 255 - lane * 8 - j; - sum += s.work[bin + bin / 32]; - } - unsigned inclusive = sum; - for (unsigned d = 1; d < 32; d *= 2) { - unsigned v = __shfl_up_sync(Full, inclusive, d); - if (lane >= d) inclusive += v; - } - unsigned before = inclusive - sum, rank = s.rank, - crossing = - __ballot_sync(Full, before < rank && inclusive >= rank), - chosen_lane = __ffs(crossing) - 1, chosen = 0, above = before; - if (lane == chosen_lane) { - for (unsigned j = 0; j < 8; ++j) { - unsigned bin = 255 - lane * 8 - j; - if (above + s.work[bin + bin / 32] >= rank) { - chosen = bin; - break; - } - above += s.work[bin + bin / 32]; - } - } - chosen = __shfl_sync(Full, chosen, chosen_lane); - above = __shfl_sync(Full, above, chosen_lane); - __syncwarp(); // every lane read s.rank above - if (lane == 0) { - s.prefix = (prefix << 8) | chosen; - s.rank = rank - above; - } - } - fb_sync(); - } - unsigned threshold = s.prefix; - unsigned greater = emit(0, threshold, 1, 0); - emit(0, threshold, 2, greater); - if (tid == 0) { - s.kept[local] = K; - float gate = fmaxf(__uint_as_float(s.gates[local]), fb_float(threshold)); - atomicExch(s.gates + local, __float_as_uint(gate)); - } - fb_sync(); -} -__device__ void fb_select(FallbackShared& s, FallbackBuffers b, unsigned K, - unsigned local, unsigned row, unsigned epoch, - unsigned first, unsigned n, unsigned tid, - bool final) { - constexpr unsigned Full = 0xffffffffu, T = 32, Items = FallbackSegment / T; - unsigned lane = tid, bank = epoch & 1; - float values[Items], gate = __uint_as_float(s.gates[local]); - unsigned count = 0; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) { - unsigned j = tid + i * T; - float v = j < n ? s.ring[bank][local][j] : -CUDART_INF_F; - if (v <= gate) v = -CUDART_INF_F; - values[i] = v; - bool valid = isfinite(v); - count += valid; - unsigned mask = __ballot_sync(Full, valid); - if (valid) { - unsigned bin = fb_half_bin(v), peers = __match_any_sync(mask, bin); - if (lane == unsigned(__ffs(peers) - 1)) - atomicAdd(&s.hist[local][bin + bin / 32], __popc(peers)); - } - } - count = __reduce_add_sync(Full, count); - __syncwarp(); - if (count) { - unsigned sum = 0; -#pragma unroll - for (unsigned j = 0; j < 32; ++j) - sum += s.hist[local][(31 - lane) * 33 + j]; - unsigned inclusive = sum; - for (unsigned d = 1; d < 32; d *= 2) { - unsigned v = __shfl_up_sync(Full, inclusive, d); - if (lane >= d) inclusive += v; - } - unsigned before = inclusive - sum, - mask = __ballot_sync(Full, before < K && inclusive >= K); - if (mask) { - unsigned chosen = __ffs(mask) - 1, group = 31 - chosen, - rank = K - __shfl_sync(Full, before, chosen), - bin = group * 32 + 31 - lane; - sum = s.hist[local][bin + bin / 32]; - inclusive = sum; - for (unsigned d = 1; d < 32; d *= 2) { - unsigned v = __shfl_up_sync(Full, inclusive, d); - if (lane >= d) inclusive += v; - } - before = inclusive - sum; - chosen = - __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; - bin = __shfl_sync(Full, bin, chosen); - gate = fmaxf(gate, fb_bin_floor(bin)); - if (lane == 0) atomicExch(s.gates + local, __float_as_uint(gate)); - } - unsigned written = s.kept[local]; - uint64_t base = (uint64_t)row * FallbackRetain; -#pragma unroll - for (unsigned i = 0; i < Items; ++i) { - float v = values[i]; - bool pass = isfinite(v) && v >= gate; - unsigned mask = __ballot_sync(Full, pass), - pos = written + __popc(mask & ((1u << lane) - 1)); - if (pass) { - b.retained[base + pos] = v; - b.indices[base + pos] = first + tid + i * T; - } - written += __popc(mask); - } - __syncwarp(); - if (lane == 0) s.kept[local] = written; - __syncwarp(); - } -} - -namespace fused_mqa_fallback { -using fused_mqa_dense::Scheduler; - -// 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, kScoreScale = 1.0f; - static constexpr uint32_t kSplitsPerSegment = FallbackSegment / 256; - FallbackShared& stream; - FallbackBuffers buffers; - uint32_t topk, query_base, row_base, epoch = 0; - uint32_t seq_k_start[4], seq_k_end[4]; - float gate[4]; - CUTLASS_DEVICE void begin_block(const Scheduler& s, uint32_t q_block, - uint32_t base) { - query_base = base; - row_base = s.get_logits_row(q_block, 0); -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) { - const uint32_t row = s.get_logits_row(q_block, base + i); - gate[i] = -CUDART_INF_F; - seq_k_start[i] = min(s.cu_seq_len_k_start[row], s.num_kv_tokens); - seq_k_end[i] = min(s.cu_seq_len_k_end[row], s.num_kv_tokens); - } - } - CUTLASS_DEVICE void begin_split(uint32_t split, uint32_t) { - if (split % kSplitsPerSegment == 0) { - __syncwarp(); - for (uint32_t i = 0; i < 4; ++i) - gate[i] = __uint_as_float(atomicAdd(stream.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]; - stream.ring[epoch & 1][query_base + i] - [(split % kSplitsPerSegment) * 256 + 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(); - fb_sync(); - const unsigned tid = threadIdx.x; - for (unsigned local = 0; local < 8; ++local) - if (stream.kept[local] > FallbackRetain - FallbackSegment) - fb_compact(stream, buffers, topk, local, row_base + local, tid); - fb_sync(); - const unsigned local = tid / 32; - fb_select(stream, buffers, topk, local, row_base + local, epoch, - kv_base + (split / kSplitsPerSegment) * FallbackSegment, - (split % kSplitsPerSegment + 1) * 256, tid % 32, false); - fb_sync(); - if (split + 1 == num_splits) - for (unsigned l = 0; l < 8; ++l) - if (stream.kept[l] > topk) - fb_compact(stream, buffers, topk, l, row_base + l, tid, true); - fb_sync(); - ++epoch; - } - CUTLASS_DEVICE void end_block(const Scheduler& s, uint32_t q_block) { - const unsigned local = threadIdx.x / 32, lane = threadIdx.x % 32; - const unsigned row = s.get_logits_row(q_block, local), - count = stream.kept[local]; - if (row < s.num_q_tokens) - for (unsigned j = lane; j < topk; j += 32) - buffers.output[(uint64_t)row * topk + j] = - j < count ? buffers.indices[(uint64_t)row * FallbackRetain + j] - : -1; - } -}; - -// Rescans one block exactly; `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__ void rescan_block( - const Scheduler& block, FallbackBuffers buffers, unsigned 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& stream = *reinterpret_cast( - smem_buffer + sizeof(fused_mqa_dense::SharedStorage)); - stream.init(threadIdx.x); - RescanPolicy policy{stream, buffers, topk}; - fused_mqa_dense::score_block(block, policy, tq, tk, ts, tw); -} - -} // namespace fused_mqa_fallback diff --git a/csrc/fused_mqa_topk/fast_select.cuh b/csrc/fused_mqa_topk/fast_select.cuh deleted file mode 100644 index fe89f4a3daed..000000000000 --- a/csrc/fused_mqa_topk/fast_select.cuh +++ /dev/null @@ -1,449 +0,0 @@ -// SPDX-License-Identifier: MIT -#pragma once -#include -#include "candidates.cuh" - -// One warp selects one row's top-K from its candidate fp24 codes (higher code = -// higher score), staged once in shared memory. Candidate codes 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 U 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. K (the row's top-k: 512, 1024 or 2048) is a runtime value; -// shared-memory scratch is sized for kMaxK. -namespace fast_select { -constexpr unsigned kMaxK = 2048, Full = 0xffffffffu, kBins = 1024, U = 8; - -// Bin (descending) where the cumulative count first reaches `rank`; -// `above` = elements in higher bins. Lane l owns bins kBins-1-32l .. -31. -__device__ __forceinline__ unsigned find_bin(const unsigned* hist, - unsigned rank, unsigned& above) { - const unsigned lane = threadIdx.x & 31; - constexpr unsigned kPerLane = kBins / 32; - unsigned sum = 0; -#pragma unroll - for (unsigned j = 0; j < kPerLane; ++j) - sum += hist[kBins - 1 - lane * kPerLane - j]; - unsigned inclusive = sum; -#pragma unroll - for (unsigned d = 1; d < 32; d *= 2) { - const unsigned v = __shfl_up_sync(Full, inclusive, d); - if (lane >= d) inclusive += v; - } - const unsigned before = inclusive - sum; - const unsigned winner = - __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; - unsigned chosen = 0, acc = before; - if (lane == winner) { - for (unsigned j = 0; j < kPerLane; ++j) { - const unsigned bin = kBins - 1 - lane * kPerLane - j, h = hist[bin]; - if (acc + h >= rank) { - chosen = bin; - break; - } - acc += h; - } - } - above = __shfl_sync(Full, acc, winner); - return __shfl_sync(Full, chosen, winner); -} - -// Requires count >= K (callers recover rows with fewer candidates). -__device__ void select_row(unsigned* codes, unsigned count, - const uint16_t* values, const int32_t* packed, - int32_t* out, unsigned K, unsigned* hist) { - const unsigned lane = threadIdx.x & 31; - unsigned lo = 0xffffffffu, hi = 0; - unsigned j = lane; - for (; j + 32 * (U - 1) < count; j += 32 * U) { - uint16_t v[U]; - int32_t p[U]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - v[u] = values[j + 32 * u]; - p[u] = packed[j + 32 * u]; - } -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - const unsigned c = - fused_candidates::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 unsigned c = - fused_candidates::candidate_load_score_code(values[j], packed[j]); - codes[j] = c; - lo = min(lo, c); - hi = max(hi, c); - } - lo = __reduce_min_sync(Full, lo); - hi = __reduce_max_sync(Full, hi); - __syncwarp(); - - // Current bucket: codes in [lo, lo + width). rank = elements still needed - // from it. - unsigned width = hi - lo + 1, rank = K; -#pragma unroll 1 - while (width > 1) { - const unsigned spread = (width - 1) >> 10; - const unsigned shift = spread ? 32 - __clz(spread) : 0; - for (unsigned b = lane; b < kBins; b += 32) hist[b] = 0; - __syncwarp(); - const unsigned top = lo + width - 1; - unsigned i = lane; - for (; i + 32 * (U - 1) < count; i += 32 * U) { - unsigned c[U]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) c[u] = codes[i + 32 * u]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) - if (c[u] >= lo && c[u] <= top) - atomicAdd(hist + ((c[u] - lo) >> shift), 1u); - } - for (; i < count; i += 32) { - const unsigned c = codes[i]; - if (c >= lo && c <= top) atomicAdd(hist + ((c - lo) >> shift), 1u); - } - __syncwarp(); - unsigned above; - const unsigned chosen = find_bin(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 unsigned threshold = lo, lt = (1u << lane) - 1u; - unsigned pos = 0; -#pragma unroll 1 - for (unsigned tie = 0; tie < 2; ++tie) { - for (unsigned j0 = 0; j0 < count && pos < K; j0 += 32 * U) { - bool pass[U]; - int32_t p[U]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - const unsigned e = j0 + 32 * u + lane; - const unsigned 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 (unsigned u = 0; u < U; ++u) { - const unsigned mask = __ballot_sync(Full, pass[u]), - at = pos + __popc(mask & lt); - if (pass[u] && at < K) - out[at] = fused_candidates::candidate_decode_index(p[u]); - pos += __popc(mask); - } - } - } -} -// Resident rows: fp24 codes live in shared memory as hi16[e] << 8 | lo8[e] for -// the row's visible keys e = 0..count-1 (key index s0 + e). Requires count >= -// K. Same levels as select_row. -__device__ __forceinline__ unsigned resident_code(const uint16_t* hi, - const uint8_t* lo, - unsigned e) { - return unsigned(hi[e]) << 8 | lo[e]; -} -struct IdentityKey { - __device__ __forceinline__ int32_t operator()(unsigned e) const { - return int32_t(e); - } -}; -template -__device__ void select_resident(const uint16_t* hi, const uint8_t* lo, - unsigned count, unsigned s0, int32_t* out, - unsigned K, unsigned* hist, KeyOf key_of = {}) { - const unsigned lane = threadIdx.x & 31; - unsigned lo_code = 0xffffffffu, hi_code = 0; - for (unsigned j = lane; j < count; j += 32) { - const unsigned c = resident_code(hi, lo, j); - lo_code = min(lo_code, c); - hi_code = max(hi_code, c); - } - unsigned base = __reduce_min_sync(Full, lo_code); - const unsigned top0 = __reduce_max_sync(Full, hi_code); - unsigned width = top0 - base + 1, rank = K; -#pragma unroll 1 - while (width > 1) { - const unsigned spread = (width - 1) >> 10; - const unsigned shift = spread ? 32 - __clz(spread) : 0; - for (unsigned b = lane; b < kBins; b += 32) hist[b] = 0; - __syncwarp(); - const unsigned top = base + width - 1; - unsigned i = lane; - for (; i + 32 * (U - 1) < count; i += 32 * U) { - unsigned c[U]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) c[u] = resident_code(hi, lo, i + 32 * u); -#pragma unroll - for (unsigned u = 0; u < U; ++u) - if (c[u] >= base && c[u] <= top) - atomicAdd(hist + ((c[u] - base) >> shift), 1u); - } - for (; i < count; i += 32) { - const unsigned c = resident_code(hi, lo, i); - if (c >= base && c <= top) atomicAdd(hist + ((c - base) >> shift), 1u); - } - __syncwarp(); - unsigned above; - const unsigned chosen = find_bin(hist, rank, above); - rank -= above; - base += chosen << shift; - width = min(1u << shift, top - base + 1); - __syncwarp(); - } - const unsigned threshold = base, lt = (1u << lane) - 1u; - unsigned pos = 0; -#pragma unroll 1 - for (unsigned tie = 0; tie < 2; ++tie) { - for (unsigned j0 = 0; j0 < count && pos < K; j0 += 32 * U) { - bool pass[U]; -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - const unsigned e = j0 + 32 * u + lane; - const unsigned c = e < count ? resident_code(hi, lo, e) : 0; - pass[u] = e < count && (tie ? c == threshold : c > threshold); - } -#pragma unroll - for (unsigned u = 0; u < U; ++u) { - const unsigned mask = __ballot_sync(Full, pass[u]), - at = pos + __popc(mask & lt); - if (pass[u] && at < K) out[at] = key_of(s0 + j0 + 32 * u + lane); - pos += __popc(mask); - } - } - } -} -// Register samples per lane for the sampled bracket selects below. -constexpr unsigned kSamplesPerLane = 16; -// Bracket select (one warp per resident row): thresholds on the hi16 part only, -// 8 elements per 16-byte shared load, per-lane counts with one warp prefix scan -// (no per-element ballots), output staged in shared memory and copied out with -// 16-byte stores. hi_row/lo_row are indexed by key index. -constexpr unsigned kMid3Cap = 1024, kScratch3Bytes = kMaxK * 4 + kMid3Cap * 2; - -__device__ __forceinline__ unsigned warp_exclusive_scan(unsigned v, - unsigned& total) { - const unsigned lane = threadIdx.x & 31; - unsigned inc = v; -#pragma unroll - for (unsigned d = 1; d < 32; d *= 2) { - const unsigned t = __shfl_up_sync(Full, inc, d); - if (lane >= d) inc += t; - } - total = __shfl_sync(Full, inc, 31); - return inc - v; -} - -template -__device__ __forceinline__ void classify8(const uint4 v, unsigned ebase, - unsigned s0, unsigned end, - unsigned h_hi, unsigned h_lo, - Fn&& fn) { - const unsigned w[4] = {v.x, v.y, v.z, v.w}; - const bool interior = ebase >= s0 && ebase + 8 <= end; -#pragma unroll - for (unsigned k = 0; k < 8; ++k) { - const unsigned h = (w[k / 2] >> (16 * (k & 1))) & 0xffffu, e = ebase + k; - const bool valid = interior || (e >= s0 && e < end); - fn(e, valid && h > h_hi, valid && h >= h_lo && h <= h_hi); - } -} - -template -__device__ bool select_resident_bracket(const uint16_t* hi_row, - const uint8_t* lo_row, unsigned s0, - unsigned count, int32_t* out, - unsigned K, unsigned char* scratch, - KeyOf key_of = {}) { - const unsigned lane = threadIdx.x & 31, lt = (1u << lane) - 1u; - // hi16 bracket from 512 strided samples (+-40 sample ranks around the - // expected K-th). - unsigned h_hi = 0, h_lo = 0; - { - unsigned smp[kSamplesPerLane]; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) - smp[u] = hi_row[s0 + (((u * 32 + lane) * 2 + 1) * count >> 10)]; - const unsigned expect = (K * 512 + count / 2) / count; - const unsigned r_hi = expect > 40 ? expect - 40 : 1, - r_lo = min(expect + 40, 512u); -#pragma unroll 1 - for (int bit = 15; bit >= 0; --bit) { - const unsigned ch = h_hi | (1u << bit), cl = h_lo | (1u << bit); - unsigned nh = 0, nl = 0; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) { - nh += smp[u] >= ch; - nl += smp[u] >= cl; - } - if (__reduce_add_sync(Full, nh) >= r_hi) h_hi = ch; - if (__reduce_add_sync(Full, nl) >= r_lo) h_lo = cl; - } - } - const unsigned a0 = s0 & ~7u, end = s0 + count, units = (end - a0 + 7) / 8; - const uint4* hv = reinterpret_cast(hi_row + a0); - unsigned na = 0, nm = 0; -#pragma unroll 2 - for (unsigned unit = lane; unit < units; unit += 32) - classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, - [&](unsigned, bool a, bool m) { - na += a; - nm += m; - }); - unsigned above, mid; - unsigned pa = warp_exclusive_scan(na, above), - pm = warp_exclusive_scan(nm, mid); - if (above >= K || above + mid < K || mid > kMid3Cap) return false; - int32_t* stage = reinterpret_cast(scratch); - uint16_t* midbuf = reinterpret_cast(scratch + kMaxK * 4); -#pragma unroll 2 - for (unsigned unit = lane; unit < units; unit += 32) - classify8(hv[unit], a0 + unit * 8, s0, end, h_hi, h_lo, - [&](unsigned e, bool a, bool m) { - if (a) stage[pa++] = key_of(e); - if (m) midbuf[pm++] = uint16_t(e); - }); - __syncwarp(); - // Exact K-th code among the bracket (rank r), over the bits where the bracket - // bounds differ. - constexpr unsigned R = kMid3Cap / 32; - unsigned mc[R]; -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - const unsigned i = u * 32 + lane; - mc[u] = 0; - if (u * 32 < mid && i < mid) { - const unsigned e = midbuf[i]; - mc[u] = unsigned(hi_row[e]) << 8 | lo_row[e]; - } - } - const unsigned r = K - above, c_lo = h_lo << 8, c_hi = h_hi << 8 | 0xffu, - diff = c_lo ^ c_hi; - const int top_bit = diff ? 31 - __clz(diff) : -1; - unsigned t = top_bit >= 0 ? c_lo & ~((2u << top_bit) - 1u) : c_lo; -#pragma unroll 1 - for (int bit = top_bit; bit >= 0; --bit) { - const unsigned cand = t | (1u << bit); - unsigned n = 0; -#pragma unroll - for (unsigned u = 0; u < R; ++u) - if (u * 32 < mid) n += mc[u] >= cand; - if (__reduce_add_sync(Full, n) >= r) t = cand; - } - unsigned pos = above; -#pragma unroll 1 - for (unsigned tie = 0; tie < 2; ++tie) { -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - if (u * 32 >= mid || pos >= K) break; - const unsigned i = u * 32 + lane; - const bool pass = i < mid && (tie ? mc[u] == t : mc[u] > t); - const unsigned mask = __ballot_sync(Full, pass), - at = pos + __popc(mask & lt); - if (pass && at < K) stage[at] = key_of(midbuf[i]); - pos += __popc(mask); - } - } - __syncwarp(); -#pragma unroll 4 - for (unsigned j = lane * 4; j < K; j += 128) - *reinterpret_cast(out + j) = - *reinterpret_cast(stage + j); - return true; -} -// Histogram select over any code source: code_of(e) / key_of(e) for e in [0, -// count). Used for rows whose shared candidates overflowed into global memory -// (rare). -template -__device__ void select_hist_generic(CodeOf code_of, unsigned count, - int32_t* out, unsigned K, unsigned* hist, - KeyOf key_of) { - const unsigned lane = threadIdx.x & 31; - unsigned lo_code = 0xffffffffu, hi_code = 0; - for (unsigned j = lane; j < count; j += 32) { - const unsigned c = code_of(j); - lo_code = min(lo_code, c); - hi_code = max(hi_code, c); - } - unsigned base = __reduce_min_sync(Full, lo_code); - const unsigned top0 = __reduce_max_sync(Full, hi_code); - unsigned width = top0 - base + 1, rank = K; -#pragma unroll 1 - while (width > 1) { - const unsigned spread = (width - 1) >> 10; - const unsigned shift = spread ? 32 - __clz(spread) : 0; - for (unsigned b = lane; b < kBins; b += 32) hist[b] = 0; - __syncwarp(); - const unsigned top = base + width - 1; - for (unsigned i = lane; i < count; i += 32) { - const unsigned c = code_of(i); - if (c >= base && c <= top) atomicAdd(hist + ((c - base) >> shift), 1u); - } - __syncwarp(); - unsigned above; - const unsigned chosen = find_bin(hist, rank, above); - rank -= above; - base += chosen << shift; - width = min(1u << shift, top - base + 1); - __syncwarp(); - } - const unsigned threshold = base, lt = (1u << lane) - 1u; - unsigned pos = 0; -#pragma unroll 1 - for (unsigned tie = 0; tie < 2; ++tie) { - for (unsigned j0 = 0; j0 < count && pos < K; j0 += 32) { - const unsigned e = j0 + lane; - const unsigned c = e < count ? code_of(e) : 0; - const bool pass = e < count && (tie ? c == threshold : c > threshold); - const unsigned mask = __ballot_sync(Full, pass), - at = pos + __popc(mask & lt); - if (pass && at < K) out[at] = key_of(e); - pos += __popc(mask); - } - } -} -// Bins of the histogram levels that resolve the exact threshold inside a -// bracket (row_select.cuh). -constexpr unsigned kMidBins = 512; - -template -__device__ __forceinline__ unsigned find_bin_n(const unsigned* hist, - unsigned rank, unsigned& above) { - const unsigned lane = threadIdx.x & 31; - constexpr unsigned kPerLane = Bins / 32; - unsigned sum = 0; -#pragma unroll - for (unsigned j = 0; j < kPerLane; ++j) - sum += hist[Bins - 1 - lane * kPerLane - j]; - unsigned inclusive = sum; -#pragma unroll - for (unsigned d = 1; d < 32; d *= 2) { - const unsigned v = __shfl_up_sync(Full, inclusive, d); - if (lane >= d) inclusive += v; - } - const unsigned before = inclusive - sum; - const unsigned winner = - __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; - unsigned chosen = 0, acc = before; - if (lane == winner) { - for (unsigned j = 0; j < kPerLane; ++j) { - const unsigned bin = Bins - 1 - lane * kPerLane - j, h = hist[bin]; - if (acc + h >= rank) { - chosen = bin; - break; - } - acc += h; - } - } - above = __shfl_sync(Full, acc, winner); - return __shfl_sync(Full, chosen, winner); -} - -} // namespace fast_select diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu index 489915c851e1..194e02fa2484 100644 --- a/csrc/fused_mqa_topk/fused_mqa_topk.cu +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -1,12 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// Single-launch fused FP8 MQA-logits + top-2048 prefill selection (SM100 -// family): one kernel scores every visible key of a row and selects its -// top-2048 keys without materializing the [rows, keys] logits matrix. -// <= 4K keys: resident_kernel.cuh (all scores of a block in shared memory). -// > 4K keys: persistent_kernel.cuh (one CTA per SM; sampled gate, candidate -// emission and selection; exact rescan of out-of-range blocks). +// 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. @@ -15,14 +14,17 @@ #include #include +#include #include -#include #include #include +#include +#include +#include +#include "constants.h" #include "persistent_kernel.cuh" -#include "resident_kernel.cuh" namespace { @@ -32,32 +34,34 @@ void check_cuda(cudaError_t err) { STD_TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err)); } -void* driver_handle() { - static void* h = nullptr; - if (!h) { - h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); - STD_TORCH_CHECK(h, "failed to load libcuda.so.1"); - } - return h; +// 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; } -CUresult enc_tiled(CUtensorMap* tm, CUtensorMapDataType dt, cuuint32_t rank, - void* addr, const cuuint64_t* dims, - const cuuint64_t* strides, const cuuint32_t* box, - const cuuint32_t* estrides, CUtensorMapInterleave il, - CUtensorMapSwizzle sw, CUtensorMapL2promotion l2, - CUtensorMapFloatOOBfill oob) { - using FT = - CUresult (*)(CUtensorMap*, CUtensorMapDataType, cuuint32_t, void*, - const cuuint64_t*, const cuuint64_t*, const cuuint32_t*, - const cuuint32_t*, CUtensorMapInterleave, CUtensorMapSwizzle, - CUtensorMapL2promotion, CUtensorMapFloatOOBfill); - static FT f = nullptr; - if (!f) { - f = reinterpret_cast(dlsym(driver_handle(), "cuTensorMapEncodeTiled")); - STD_TORCH_CHECK(f, "failed to load cuTensorMapEncodeTiled"); +// 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 f(tm, dt, rank, addr, dims, strides, box, estrides, il, sw, l2, oob); + return it->second; } CUtensorMap make_2d(void* ptr, CUtensorMapDataType dt, int elem_size, @@ -74,24 +78,28 @@ CUtensorMap make_2d(void* ptr, CUtensorMapDataType dt, int elem_size, : swizzle_mode == 32 ? CU_TENSOR_MAP_SWIZZLE_32B : CU_TENSOR_MAP_SWIZZLE_NONE; - CUresult r = enc_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); + 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. -CUtensorMap make_strided_kv(void* ptr, unsigned groups, unsigned 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; - const cuuint64_t gdims[3] = {128, groups, stride}; - const cuuint64_t gstrides[2] = {128ull * stride, 128}; - const cuuint32_t box[3] = {128, 256, 1}; + constexpr uint64_t kRowBytes = fused_mqa_topk::kHeadDim; + const cuuint64_t gdims[3] = {kRowBytes, groups, stride}; + const cuuint64_t gstrides[2] = {kRowBytes * stride, kRowBytes}; + const cuuint32_t box[3] = {kRowBytes, fused_mqa_topk::kTileKeys, 1}; const cuuint32_t estrides[3] = {1, 1, 1}; - CUresult r = enc_tiled( + 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); @@ -100,10 +108,128 @@ CUtensorMap make_strided_kv(void* ptr, unsigned groups, unsigned stride) { return tm; } -constexpr int kHeads = 32; -constexpr int kHeadDim = 128; -constexpr int kCandidates = 16384; // per-row candidate slab -constexpr int kRetained = 4096; // per-row recovery slab +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; +namespace persist = fused_mqa_topk::persist; + +// -------- launch configuration -------- +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 <= 131072 ? 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. +inline cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) { + // The recovery rescan reuses the same dynamic shared memory with the rescan's + // layout. + constexpr int kFallbackSmem = + 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(kFallbackSmem, 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). +inline 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 > 131072; + 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) { @@ -113,19 +239,35 @@ void check_tensor(const torch::stable::Tensor& t, int32_t device, 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 // q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32, // weights [rows,32] fp32, starts/ends [rows] int32 (visible key range // [start,end) of each row), out [rows,K] int32 with K (the top-k) 512, 1024 or -// 2048: absolute key indices (columns of the logits), or with relative=true -// indices relative to the row's start (key - start, the top_k_per_row_prefill -// convention); -1 padded when fewer than K keys are visible. -// 0 < rows, 0 < keys <= 1048576 (the last 8-row block may be partial). +// 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. 0 < +// rows < 2^26, 0 < keys <= 1048576 (the last 8-row block may be partial). q, k, +// scales, weights and out must be 16-byte aligned (TMA / vector stores). +// 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. +// The output is exact up to ties: the gated path ranks fp24 score codes (the +// low 8 mantissa bits dropped), the exact rescan ranks fp32 scores. // 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), -// values [8S,16384] fp16, indices [8S,16384] int32, +// 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( @@ -135,7 +277,7 @@ void fused_mqa_topk_prefill( 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, bool relative) { + 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"); @@ -146,99 +288,100 @@ void fused_mqa_topk_prefill( 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::Half, "values"); + 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) == kHeads && q.size(2) == kHeadDim, - "q must be [rows, 32, 128]"); + 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); - STD_TORCH_CHECK(rows > 0, "rows must be positive"); - STD_TORCH_CHECK(keys > 0 && keys <= (1 << 20), - "keys must be in [1, 1048576]"); - const int64_t blocks = (rows + 7) / 8; + // 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) == kHeads, + 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 && - (out.size(1) == 512 || out.size(1) == 1024 || out.size(1) == 2048), - "out must be [rows, K] with K in {512, 1024, 2048}"); - STD_TORCH_CHECK(reinterpret_cast(out.data_ptr()) % 16 == 0, - "out must be 16-byte aligned"); + 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 unsigned 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()); + 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()); - auto tq = make_2d(q.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, - int(rows) * 32, 128, 256, 128, 128); - auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, 128, - int(keys), 128, 256, 128, 128); + // 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, 256, 1, 0, 0); - auto tw = make_2d(weights.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, 32, - int(rows), 32, 8, 32, 0); - - if (ukeys <= fused_mqa_resident::kResidentKeys) { - // Every score of a block fits in shared memory: no workspace. - constexpr int smem = int(fused_mqa_resident::kSmemBytes); - check_cuda(cudaFuncSetAttribute( - reinterpret_cast(fused_mqa_topk_fp8_resident_kernel), - cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); - fused_mqa_topk_fp8_resident_kernel<<>>( - urows, ukeys, starts_p, ends_p, out_p, topk, relative, tq, ts, tk, tw); - check_cuda(cudaGetLastError()); - return; - } + 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]"); - STD_TORCH_CHECK(values.dim() == 2 && values.size(1) == kCandidates && - values.size(0) % 16 == 0 && indices.dim() == 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) == kCandidates, + indices.size(1) == kCandidatesPerRow, "values/indices must be [8 * slots, 16384], slots even"); - STD_TORCH_CHECK(retained.dim() == 2 && retained.size(1) == kRetained && + 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) == kRetained, + retained_indices.size(1) == kRetainedPerRow, "retained/retained_indices must be [8 * slots, 4096]"); - STD_TORCH_CHECK(counts.numel() >= blocks * 8 && defer.numel() >= blocks, - "per-row workspace too small"); + 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. - int num_sms = 0; - check_cuda( - cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev)); - const unsigned grid = static_cast( - std::min({int64_t(num_sms), blocks, values.size(0) / 16})); - fused_persist::LaunchConfig c; - check_cuda(fused_persist::launch_config(ukeys, topk, c)); + 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 FallbackBuffers fallback{static_cast(retained.data_ptr()), - static_cast(retained_indices.data_ptr()), - out_p}; - c.kernel<<>>( - urows, ukeys, starts_p, ends_p, static_cast(values.data_ptr()), + 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()), fallback, c.phase_bytes, + 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, - relative, ublocks, static_cast(counter.data_ptr()), + 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 index 407945c54c94..40c3ca3d86a4 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -3,59 +3,57 @@ // loop). if (warp_idx == 9 && cute::elect_one_sync()) { if (epoch) - for (unsigned i = 0; i < B::total; ++i) mbar_inval(bar_words + i); - for (unsigned i = 0; i < B::total; ++i) + 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(256 + 32); - for (unsigned i = 0; i < KVS; ++i) { + 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(kRel ? 1 : 256); + empty_kv[i].init(kMathThreads); } - if constexpr (kRel) - for (unsigned i = 0; i < 2 * KVS; ++i) sf_full[i].init(1); - for (unsigned i = 0; i < B::NT; ++i) { + for (uint32_t i = 0; i < BarIdx::NT; ++i) { full_tmem[i].init(1); - empty_tmem[i].init(kLin ? 128 : 256); + empty_tmem[i].init(kMathWgThreads); } - if constexpr (kLin) full_lin->init(32); - for (unsigned i = 0; i < 2; ++i) { + 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); - } // every reader lane arrives - sel_done.init(8); - for (unsigned i = 0; i < 16; ++i) { + } + sel_done.init(8); // one lane per math warp + for (uint32_t i = 0; i < 16; ++i) { h_full[i].init(32); h_empty[i].init(32); } cutlass::arch::fence_barrier_init(); - unsafe_s = 0; + unsafe_rows = 0; } __syncwarp(); if (warp_idx == 10) cute::TMEM::Allocator1Sm().allocate(512, &tmem_ptr); __syncthreads(); -if (epoch == 0) cudaGridDependencySynchronize(); if (warp_idx == 8) { - // ---------------- claimer + Q/weights producer (+ drain of WG1 rows) - // ---------------- + // -------- claimer + Q/weights producer + drain of WG1 rows -------- cutlass::arch::warpgroup_reg_dealloc(); - unsigned n_q = 0, n_blk = 0, drain_epoch = 0; - for (unsigned k = 0;; ++k) { + // 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_q & 1) ^ 1); + empty_q.wait((n_blk & 1) ^ 1); } __syncwarp(); - claim_empty[k & 1].wait(((k >> 1) & 1) ^ - 1); // every lane: lanes 0-7 write the entry - unsigned b = 0; + // 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(0xffffffffu, b, 0); ClaimEntry& e = entries[k & 1]; if (b >= nblocks) { if (lane == 0) { - const unsigned kind = b == kPause ? kPause : kTerm; + const uint32_t kind = b == kPause ? kPause : kTerm; e.vblock = kind; epoch_kind = kind; } @@ -63,10 +61,11 @@ if (warp_idx == 8) { claim_full[k & 1].arrive(); // every writer lane arrives break; } - const unsigned vblock = nblocks - 1 - b, row_base = vblock * 8; - unsigned st = 0, en = 0; - if (lane < 8) { - const unsigned r = row_base + lane; + // 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); @@ -75,30 +74,32 @@ if (warp_idx == 8) { e.ends[lane] = en; } const bool all_short = - __ballot_sync(0xffffffffu, lane < 8 && en - st > topk) == 0; + __ballot_sync(0xffffffffu, lane < kBlockRows && en - st > topk) == 0; const bool gated = - __ballot_sync(0xffffffffu, lane < 8 && en - st > sample_target) != 0; + __ballot_sync(0xffffffffu, + lane < kBlockRows && en - st > sample_target) != 0; // KV range: rows past the end are clamped to the last valid row (DeepGEMM // scheduler). - const unsigned src = min(row_base + (lane & 7), rows - 1) - row_base; - unsigned lo = __shfl_sync(0xffffffffu, st, src), + const uint32_t src = min(row_base + (lane & 7), rows - 1) - row_base; + uint32_t lo = __shfl_sync(0xffffffffu, st, src), hi = __shfl_sync(0xffffffffu, en, src); #pragma unroll for (int d = 4; d >= 1; d >>= 1) { lo = min(lo, __shfl_xor_sync(0xffffffffu, lo, d)); hi = max(hi, __shfl_xor_sync(0xffffffffu, hi, d)); } - const unsigned kv_base = lo / 4 * 4; - if (!all_short && lane < 8) counts[row_base + lane] = 0; + 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) { - // 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 unsigned slot = 2 * blockIdx.x + (n_blk & 1); e.vblock = vblock; e.all_short = all_short; e.kv_base = kv_base; - e.splits = (hi - kv_base + 255) / 256; + e.splits = splits; e.samples = gated ? sample_tiles : 0u; e.slot = slot; } @@ -106,199 +107,134 @@ if (warp_idx == 8) { claim_full[k & 1].arrive(); // every writer lane arrives if (all_short) continue; if (lane == 0) { - fused_mqa_dg::tma_copy<128, 256, 0>(&tq, &full_q, sm.q, 0, row_base * 32); - fused_mqa_dg::tma_copy<32, 8, 0>(&tw, &full_q, sm.weights, 0, row_base); + 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_q; ++n_blk; - if constexpr (!kSmem) { - const unsigned splits = (hi - kv_base + 255) / 256; - const unsigned slot = __shfl_sync(0xffffffffu, e.slot, 0); - const int64_t row_delta = int64_t(slot) * 8 - int64_t(row_base); - for (unsigned base = 0; base < splits; base += Interval, ++drain_epoch) { - for (unsigned p = 4; p < 8; ++p) - ho_drain(hdata, h_full, h_empty, p, drain_epoch, row_base + 4, counts, - kCandGlobalCap, 1 && base + Interval >= splits, - [&](unsigned row, unsigned slot2, unsigned record, - unsigned batch_base, unsigned source) { - const uint32_t physical = - kv_base + (batch_base * 2 + (record & 255u)) * 128 + - (source % 128); - const uint64_t off = - static_cast(int64_t(row) + row_delta) * - kCandGlobalCap + - slot2; - fused_candidates::store_candidate_record( - vals + off, indices + off, record >> 8, physical); - }); - } - cutlass::arch::NamedBarrier(320, 4).arrive(); + 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 = 4; p < 8; ++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 % 128); + }); } + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); } } else if (warp_idx == 9) { - // ---------------- KV producer: sample tiles then scan tiles, one ring - // ---------------- + // -------- KV producer: sample tiles then scan tiles, one ring -------- cutlass::arch::warpgroup_reg_dealloc(); - RingPipeline kv_pipe; - [[maybe_unused]] RingPipeline<2 * KVS> sf_pipe; - unsigned n_blk = 0; - for (unsigned k = 0;; ++k) { - claim_full[k & 1].wait((k >> 1) & 1); - const ClaimEntry& e = entries[k & 1]; - const unsigned vblock = e.vblock, all_short = e.all_short, - kv_base = e.kv_base, splits = e.splits, samples = e.samples; - __syncwarp(); - if (vblock >= kPause) break; - claim_empty[k & 1].arrive(); - if (all_short) continue; - if constexpr (KVS == 3) { - // > 16K: the whole warp stages the strided sample scales into - // sf_kv[stage] (prefetched one tile ahead in registers), lane 0 issues - // the KV TMA. + 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 unsigned center = sample_stride / 2; + const uint32_t center = sample_stride / 2; float pf[8]; #pragma unroll - for (unsigned j = 0; j < 8; ++j) { - const unsigned g = j * 32 + lane; + for (uint32_t j = 0; j < 8; ++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 (unsigned t = 0; t < samples; ++t) { + 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 unsigned s = __shfl_sync(0xffffffffu, s0, 0); + const uint32_t s = __shfl_sync(0xffffffffu, s0, 0); // Every writing lane acquires and releases the stage itself. empty_kv[s].wait(__shfl_sync(0xffffffffu, ph, 0) ^ 1); #pragma unroll - for (unsigned j = 0; j < 8; ++j) sm.sf_kv[s][j * 32 + lane] = pf[j]; - cutlass::arch::fence_view_async_shared(); // generic writes before - // later TMA writes of this - // stage + for (uint32_t j = 0; j < 8; ++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); - fused_mqa_dg::tma_copy<128, 256, 0, cutlass::float_e4m3_t, true>( - &tks, &full_kv[s], sm.kv[s], 0, t * 256, 1, center); + dg::tma_copy( + &tks, &full_kv[s], sm.kv[s], 0, t * kTileKeys, center); } full_kv[s].arrive(); #pragma unroll - for (unsigned j = 0; j < 8; ++j) { - const unsigned g = (t + 1) * 256 + j * 32 + lane; + for (uint32_t j = 0; j < 8; ++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) { - if (KVS != 3 && n_blk > 0) - sel_done.wait((n_blk - 1) & - 1); // selection scratch overlaps the KV stages - const unsigned center = sample_stride / 2; -#pragma unroll 1 - for (unsigned t = 0; t < (KVS == 3 ? 0u : samples); ++t) { - CUTE_TIE_DECL(kv_pipe.advance(), s, ph); - empty_kv[s].wait(ph ^ 1); - fused_mqa_dg::tma_copy<128, 256, 0, cutlass::float_e4m3_t, true>( - &tks, &full_kv[s], sm.kv[s], 0, t * 256, 1, center); - mbar_arrive_expect_tx(&full_kv[s], kKVBytes, kFullKvArrivals); - } #pragma unroll 1 - for (unsigned t = 0; t < splits; ++t) { + 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 * 256; - fused_mqa_dg::tma_copy<128, 256, 0>(&tk, &full_kv[s], sm.kv[s], 0, off); - if constexpr (kRel) { - mbar_arrive_expect_tx(&full_kv[s], kKVBytes, kFullKvArrivals); - CUTE_TIE_DECL(sf_pipe.advance(), sfs, sfph); - (void)sfph; - fused_mqa_dg::tma_copy<256, 1, 0>(&ts, &sf_full[sfs], sm.sf_kv[sfs], - off, 0); - sf_full[sfs].arrive_and_expect_tx(1024); - } else { - fused_mqa_dg::tma_copy<256, 1, 0>(&ts, &full_kv[s], sm.sf_kv[s], off, - 0); - mbar_arrive_expect_tx(&full_kv[s], kKVBytes + 1024, kFullKvArrivals); - } + 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 ---------------- + // -------- UMMA issuer -------- cutlass::arch::warpgroup_reg_dealloc(); RingPipeline<1> q_pipe; - RingPipeline kv_pipe; - RingPipeline tmem_pipe; + RingPipeline kv_pipe; + RingPipeline tmem_pipe; auto instr_desc = cute::UMMA::make_instr_desc(); + float, 128, kUmmaN, cute::UMMA::Major::K, + cute::UMMA::Major::K>(); auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); - for (unsigned k = 0;; ++k) { - claim_full[k & 1].wait((k >> 1) & 1); - const ClaimEntry& e = entries[k & 1]; - const unsigned vblock = e.vblock, all_short = e.all_short, - splits = e.splits, samples = e.samples; - __syncwarp(); - if (vblock >= kPause) break; - claim_empty[k & 1].arrive(); - if (all_short) continue; + 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); - if constexpr (kLin) { - full_lin->wait(q_phase); // this block's linear-term rows (warp 11) - ptx::tcgen05_after_thread_sync(); - } - const unsigned total = samples + splits; + 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 (unsigned t = 0; t < total; ++t) { + 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()) { - if constexpr (kLin) { - // 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). + // 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 + 128 * 128; -#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]); - } - } - } else { + for (uint32_t i = 0; i < 2; ++i) { #pragma unroll - for (uint32_t i = 0; i < 2; ++i) { + 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 * 256; + 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 = @@ -306,13 +242,12 @@ if (warp_idx == 8) { sm.kv[s], i * 128, kk * 32); auto b_desc = mma::sm100::make_umma_desc( - sm.q, 0, kk * 32); + 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]); } - if constexpr (kRel) ptx::umma_arrive_no_elect(empty_kv[s]); } } __syncwarp(); @@ -320,81 +255,62 @@ if (warp_idx == 8) { empty_q.arrive(); } } else if (warp_idx == 11) { - // ---------------- drain of WG0 rows (> 16K) ---------------- + // -------- linear-term rows + drain of WG0 rows -------- cutlass::arch::warpgroup_reg_dealloc(); - if constexpr (!kSmem) { - unsigned drain_epoch = 0, n_lin = 0; - for (unsigned k = 0;; ++k) { - claim_full[k & 1].wait((k >> 1) & 1); - const ClaimEntry& e = entries[k & 1]; - const unsigned vblock = e.vblock, all_short = e.all_short, - kv_base = e.kv_base, splits = e.splits, slot = e.slot; - __syncwarp(); - if (vblock >= kPause) break; - claim_empty[k & 1].arrive(); - if (all_short) continue; - if constexpr (kLin) { - // 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); - cutlass::arch::fence_view_async_shared(); // generic smem writes -> - // UMMA (async proxy) - full_lin->arrive(); // every lane arrives (count 32) - } - const unsigned row_base = vblock * 8; - const int64_t row_delta = int64_t(slot) * 8 - int64_t(row_base); - for (unsigned base = 0; base < splits; base += Interval, ++drain_epoch) { - for (unsigned p = 0; p < 4; ++p) - ho_drain(hdata, h_full, h_empty, p, drain_epoch, row_base, counts, - kCandGlobalCap, 1 && base + Interval >= splits, - [&](unsigned row, unsigned slot2, unsigned record, - unsigned batch_base, unsigned source) { - const uint32_t physical = - kv_base + (batch_base * 2 + (record & 255u)) * 128 + - (source % 128); - const uint64_t off = - static_cast(int64_t(row) + row_delta) * - kCandGlobalCap + - slot2; - fused_candidates::store_candidate_record( - vals + off, indices + off, record >> 8, physical); - }); - } - cutlass::arch::NamedBarrier(320, 4).arrive(); + 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 < 4; ++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 % 128); + }); } + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); } } else { - // ---------------- math warps 0-7 ---------------- + // -------- math warps 0-7 -------- cutlass::arch::warpgroup_reg_alloc(); - const unsigned math_tid = threadIdx.x; - const unsigned query_base = (warp_idx / 4) * 4; + const uint32_t math_tid = threadIdx.x; + const uint32_t query_base = (warp_idx / 4) * 4; RingPipeline<1> q_pipe; - RingPipeline kv_pipe; - [[maybe_unused]] RingPipeline<2 * KVS> sf_pipe; - RingPipeline tmem_pipe; - // > 16K: each warpgroup consumes every second TMEM slot position (WG1 starts - // at 1). - constexpr unsigned kTmemStep = kLin ? 2 : 1; - if (kLin && warp_idx >= 4) tmem_pipe.advance(1); - [[maybe_unused]] const uint32_t heads_col = warp_idx < 4 ? 16u : 0u; - [[maybe_unused]] const uint32_t lin_col = warp_idx < 4 ? 0u : 128u; - unsigned produce_epoch = 0; - constexpr uint32_t kRingSlots = 8; - for (unsigned k = 0;; ++k) { + 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 unsigned vblock = e.vblock; + const uint32_t vblock = e.vblock; if (vblock >= kPause) break; - const unsigned row_base = vblock * 8; + const uint32_t row_base = vblock * kBlockRows; if (e.all_short) { - for (unsigned r = 0; r < 8 && row_base + r < rows; ++r) { - const unsigned s0 = e.starts[r], vis = e.ends[r] - s0; - for (unsigned j = math_tid; j < topk; j += 256) - fallback.output[(uint64_t)(row_base + r) * topk + j] = - j < vis ? int((relative ? 0u : s0) + j) : -1; - } + 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; @@ -404,78 +320,63 @@ if (warp_idx == 8) { CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase); (void)q_stage; full_q.wait(q_phase); - if constexpr (kLin) full_lin->wait(q_phase); + full_lin->wait(q_phase); } - static_assert(SampleQ * 2112 % 4 == 0); - for (unsigned j = math_tid; j < SampleQ * 2112 / 4; j += 256) - reinterpret_cast(stream.hist)[j] = make_uint4(0, 0, 0, 0); + 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[4][32]; - [[maybe_unused]] float - row_s[4]; // > 16K: pow2 row scale of the linear term + float row_s[4]; // pow2 row scale of the linear term uint32_t seq_k_start[4], seq_k_end[4]; #pragma unroll for (uint32_t i = 0; i < 4; ++i) { seq_k_start[i] = e.starts[query_base + i]; seq_k_end[i] = e.ends[query_base + i]; - if constexpr (kLin) { - // 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 < 32; j += 4) { - const float4 w4 = *reinterpret_cast( - sm.weights + (query_base + i) * 32 + 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; - } - } else { + // 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 < 32; j += 4) { - const float4 w4 = *reinterpret_cast( - sm.weights + (query_base + i) * 32 + j); - weights[i][j] = w4.x * 0.5f; - weights[i][j + 1] = w4.y * 0.5f; - weights[i][j + 2] = w4.z * 0.5f; - weights[i][j + 3] = w4.w * 0.5f; - } + for (uint32_t j = 0; j < 32; j += 4) { + const float4 w4 = *reinterpret_cast( + sm.weights + (query_base + i) * 32 + 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 * 32 + lane]; - bool bad_w = w != 0.0f && fabsf(w) < 0x1p-125f; - if constexpr (kLin) { - // Scaled weights that become subnormal or infinite take the exact - // recovery path. - const float wr = w * 0.5f * (1.0f / sm.lin_scale[warp_idx]); - bad_w |= (wr != 0.0f && fabsf(wr) < 0x1p-125f) || isinf(wr); - } - const unsigned bad = __ballot_sync(0xffffffffu, bad_w); - if (lane == 0 && bad) atomicOr(&unsafe_s, 1u << warp_idx); + // 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(0xffffffffu, bad_w); + if (lane == 0 && bad) atomicOr(&unsafe_rows, 1u << warp_idx); } - sample_sync(); + math_wg_sync(); // ---- sample tiles ---- #pragma unroll 1 - for (unsigned t = 0; t < e.samples; ++t) { + for (uint32_t t = 0; t < e.samples; ++t) { CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); - if constexpr (!kRel) full_kv[kv_stage].wait(kv_phase); + full_kv[kv_stage].wait(kv_phase); #pragma unroll - for (unsigned kv_half = 0; kv_half < 2; ++kv_half) { - const unsigned kv_lane = (math_tid % 128) + kv_half * 128; - const unsigned group = t * 256 + kv_lane; - const unsigned kv_offset = group * sample_stride + sample_stride / 2; + 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; - float scale_kv; - if constexpr (kLin) - scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + - kv_lane); // 0 outside the view - else - scale_kv = in_view ? __ldg(scales + kv_offset) : 0.0f; + // 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 (!kRel && kv_half == 1) { + if (kv_half == 1) { cutlass::arch::fence_view_async_shared(); empty_kv[kv_stage].arrive(); } @@ -484,328 +385,143 @@ if (warp_idx == 8) { #pragma unroll for (uint32_t i = 0; i < 4; ++i) { float acc[32]; - [[maybe_unused]] float lin[4]; - if constexpr (kLin) { - ptx::tmem_load_32dp32b<32>( - tmem_stage * kSlotCols + heads_col + i * 32, - reinterpret_cast(acc)); - ptx::tmem_load_32dp32b<4>(tmem_stage * kSlotCols + lin_col + i * 4, - reinterpret_cast(lin)); - } else { - ptx::tmem_load_32dp32b<32>(tmem_stage * 256 + (query_base + i) * 32, - reinterpret_cast(acc)); - } + float lin[4]; + ptx::tmem_load_32dp32b<32>( + tmem_stage * kSlotCols + heads_col + i * 32, + 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 == 3) { ptx::tcgen05_before_thread_sync(); empty_tmem[tmem_stage].arrive(); } - float2 sum_0 = {}, sum_1 = {}; - if constexpr (kLin) { // terms 2, 3 are in units of s * 2^-8 - sum_0 = make_float2(lin[0], lin[1]); - sum_1 = __fmul2_rn(make_float2(lin[2], lin[3]), - make_float2(0x1p-8f, 0x1p-8f)); - } -#pragma unroll - for (uint32_t j = 0; j < 32; j += 4) { - auto transform = [&](uint32_t at, float2 sum) { - const auto abs_x = - make_float2(fabsf(acc[at]), fabsf(acc[at + 1])); - const auto w = make_float2(weights[i][at], weights[i][at + 1]); - if constexpr (kLin) { - return __ffma2_rn( - abs_x, w, - sum); // relu = (x + |x|) / 2, linear part on the TC - } else { - const auto x = make_float2(acc[at], acc[at + 1]); - return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); - } - }; - sum_0 = transform(j, sum_0); - sum_1 = transform(j + 2, sum_1); - } - const auto sum = __fadd2_rn(sum_0, sum_1); - float bq; - if constexpr (kLin) - bq = (sum.x + sum.y) * (row_s[i] * scale_kv); - else - bq = (sum.x + sum.y) * scale_kv; + // The linear term seeds the sums (terms 2, 3 in units of s * 2^-8). + const float bq = + reduce_heads( + acc, weights[i], make_float2(lin[0], lin[1]), + __fmul2_rn(make_float2(lin[2], lin[3]), + make_float2(kLinLoScale, kLinLoScale))) * + (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 (valid) { - const unsigned bin = sample_bin(bq); - atomicAdd(stream.hist[query_base + i] + bin + bin / 32, 1u); + const uint32_t bin = half_bin<5>(bq); + atomicAdd(sample_smem.hist[query_base + i] + padded_bin(bin), 1u); } } } } - sample_sync(); - finish_sample_histogram(stream.hist[warp_idx], gate_s + warp_idx, + math_wg_sync(); + finish_sample_histogram(sample_smem.hist[warp_idx], gate_s + warp_idx, e.ends[warp_idx] - e.starts[warp_idx], sample_target, lane); - sample_sync(); - const unsigned row_base2 = - *reinterpret_cast(&e.vblock) * 8; + 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 = + *reinterpret_cast(&e.vblock) * kBlockRows; // ---- scan tiles ---- - if constexpr (kLin) { -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) { - const float inv_s = 1.0f / row_s[i]; -#pragma unroll - for (uint32_t j = 0; j < 32; j += 4) { - const float4 w4 = *reinterpret_cast( - sm.weights + (query_base + i) * 32 + 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; - } - } - } - const unsigned kv_base = e.kv_base, splits = e.splits; - const int64_t cand_row_delta = int64_t(e.slot) * 8 - int64_t(row_base2); + 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[4]; #pragma unroll for (uint32_t i = 0; i < 4; ++i) { gate_score[i] = gate_s[query_base + i]; - if (*reinterpret_cast(&unsafe_s) >> - (query_base + i) & + if ((*reinterpret_cast(&unsafe_rows) >> + (query_base + i)) & 1u) { - gate_score[i] = __int_as_float(0x7f800000); - if (math_tid % 128 == 0) - atomicAdd(counts + row_base2 + query_base + i, - static_cast(kCandGlobalCap) + 1); + // 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); } - // > 16K: compare S * scale_kv against gate / row_s (row_s is a power of - // two -> exact). - if constexpr (kLin) gate_score[i] *= 1.0f / row_s[i]; + // 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[4] = {}; uint32_t flush_base = 0; - const uint32_t logits_stride = kCandCap; - [[maybe_unused]] const uint32_t pool_size = logits_stride / 6, - seg_size = (logits_stride - pool_size) / 4; - [[maybe_unused]] const uint32_t seg_warp = warp_idx % 4; - [[maybe_unused]] const uint32_t cand_ext = 8 * logits_stride + 256; - [[maybe_unused]] uint32_t seg_cnt[4] = {}; - if constexpr (kSmem) { - if (math_tid < 8) { - row_count[math_tid] = 0; - pool_cnt[math_tid] = 0; - } - cutlass::arch::NamedBarrier(256, 3).sync(); - } #pragma unroll 1 for (uint32_t kv_split_idx = 0; kv_split_idx < splits; ++kv_split_idx) { - if constexpr (!kSmem) - if ((kv_split_idx & (Interval - 1)) == 0) { - h_empty[(produce_epoch & 1) * 8 + warp_idx].wait( - ((produce_epoch >> 1) & 1) ^ 1); - __syncwarp(); - } - CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); - unsigned sf_stage = kv_stage; - if constexpr (kRel) { - CUTE_TIE_DECL(sf_pipe.advance(), rs, rph); - sf_full[rs].wait(rph); - sf_stage = rs; - } else { - full_kv[kv_stage].wait(kv_phase); + if ((kv_split_idx & (Interval - 1)) == 0) { + h_empty[(produce_interval & 1) * 8 + 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 (unsigned kv_half = 0; kv_half < 2; ++kv_half) { - const unsigned kv_lane = (math_tid % 128) + kv_half * 128; - const unsigned kv_offset = kv_base + kv_split_idx * 256 + kv_lane; - const float scale_kv = ptx::ld_shared(sm.sf_kv[sf_stage] + kv_lane); + 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 (!kRel && kv_half == 1) { + if (kv_half == 1) { cutlass::arch::fence_view_async_shared(); empty_kv[kv_stage].arrive(); } float accum_pair[2][32]; - [[maybe_unused]] float lin_pair[2][4]; - [[maybe_unused]] float bp4[4]; // > 16K: row scores of the slot + float lin_pair[2][4]; + float bp4[4]; // row scores of the slot, in units of the row scale #pragma unroll for (uint32_t pair = 0; pair < 4; pair += 2) { - if constexpr (kLin) { - const uint32_t slot_base = tmem_stage * kSlotCols; - ptx::tmem_load_32dp32b<64>(slot_base + heads_col + pair * 32, - reinterpret_cast(accum_pair)); - ptx::tmem_load_32dp32b<8>(slot_base + lin_col + pair * 4, - reinterpret_cast(lin_pair)); - } else { - const uint32_t tmem_addr = - tmem_stage * 256 + (query_base + pair) * 32; - ptx::tmem_load_32dp32b<64>(tmem_addr, - reinterpret_cast(accum_pair)); - } + const uint32_t slot_base = tmem_stage * kSlotCols; + ptx::tmem_load_32dp32b<64>(slot_base + heads_col + pair * 32, + 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 >= 4) { ptx::tcgen05_before_thread_sync(); empty_tmem[tmem_stage].arrive(); } - float2 sum_0[2] = {}, sum_1[2] = {}; - if constexpr (kLin) { - // The linear term (sum_t term_t . k, in units of the row scale) - // seeds the sums. -#pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - sum_0[r] = make_float2(lin_pair[r][0], lin_pair[r][1]); - sum_1[r] = - __fmul2_rn(make_float2(lin_pair[r][2], lin_pair[r][3]), - make_float2(0x1p-8f, 0x1p-8f)); // units s * 2^-8 - } - } -#pragma unroll - for (uint32_t j = 0; j < 32; j += 4) { -#pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - const uint32_t i = pair + r; - auto* a = accum_pair[r]; - auto transform = [&](uint32_t at, float2 sum) { - const auto abs_x = make_float2(fabsf(a[at]), fabsf(a[at + 1])); - const auto w = make_float2(weights[i][at], weights[i][at + 1]); - if constexpr (kLin) { - return __ffma2_rn(abs_x, w, sum); - } else { - const auto x = make_float2(a[at], a[at + 1]); - return __ffma2_rn(__fadd2_rn(x, abs_x), w, sum); - } - }; - sum_0[r] = transform(j, sum_0[r]); - sum_1[r] = transform(j + 2, sum_1[r]); - } - } - if constexpr (kLin) { + // 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) { - const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - bp4[pair + r] = (sum.x + sum.y) * scale_kv; - } - } else + for (uint32_t r = 0; r < 2; ++r) + bp4[pair + r] = + reduce_heads( + accum_pair[r], weights[pair + r], + make_float2(lin_pair[r][0], lin_pair[r][1]), + __fmul2_rn(make_float2(lin_pair[r][2], lin_pair[r][3]), + make_float2(kLinLoScale, kLinLoScale))) * + 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(0xffffffffu, any_pass)) [[unlikely]] { #pragma unroll - for (uint32_t r = 0; r < 2; ++r) { - const uint32_t i = pair + r; - const auto sum = __fadd2_rn(sum_0[r], sum_1[r]); - float bq; - if constexpr (kLin) - bq = (sum.x + sum.y) * (row_s[i] * scale_kv); - else - bq = (sum.x + sum.y) * scale_kv; - const uint32_t row = row_base2 + query_base + i; - const bool pass = - (kSmem ? kv_offset - seq_k_start[i] < - seq_k_end[i] - seq_k_start[i] - : kv_offset >= seq_k_start[i] && - kv_offset < seq_k_end[i] && isfinite(bq)) && - bq > gate_score[i]; - if constexpr (kSmem) { - const unsigned pass_mask = __ballot_sync(0xffffffffu, pass); - const unsigned slot = - seg_cnt[i] + __popc(pass_mask & ((1u << lane) - 1u)); - seg_cnt[i] += __popc(pass_mask); - const uint32_t code = fused_candidates::candidate_fp24_code(bq); - const uint32_t seg_cap = seg_size; - const uint32_t at = - pass && slot < seg_cap - ? (query_base + i) * logits_stride + - seg_warp * seg_size + slot - : 8 * logits_stride + warp_idx * 32 + lane; - auto* cbase = region; - reinterpret_cast(cbase)[at] = - static_cast(code >> 8); - cbase[cand_ext * 2 + at] = static_cast(code); - reinterpret_cast(cbase + cand_ext * 3)[at] = - static_cast(kv_offset); - if (seg_cnt[i] > seg_cap) { - const bool over = pass && slot >= seg_cap; - const unsigned over_mask = __ballot_sync(0xffffffffu, over); - if (over_mask) { - unsigned pool_base = 0; - if (lane == 0) - pool_base = atomicAdd(&pool_cnt[query_base + i], - __popc(over_mask)); - pool_base = __shfl_sync(0xffffffffu, pool_base, 0); - const unsigned p = - pool_base + __popc(over_mask & ((1u << lane) - 1u)); - if (over && p < pool_size) { - const uint32_t pat = - (query_base + i) * logits_stride + 4 * seg_size + p; - reinterpret_cast(cbase)[pat] = - static_cast(code >> 8); - cbase[cand_ext * 2 + pat] = static_cast(code); - reinterpret_cast(cbase + cand_ext * 3)[pat] = - static_cast(kv_offset); - } else if (over && p - pool_size < kCandGlobalCap - 1) { - const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - kCandGlobalCap + - (p - pool_size); - __stcg(vals + off, static_cast(code >> 8)); - __stcg(indices + off, - static_cast((code & 0xffu) << 16 | - kv_offset)); - } - } - } - } else if (pass) { + for (uint32_t i = 0; i < 4; ++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 = - fused_candidates::candidate_fp24_code(bq); - hdata.records[produce_epoch & 1][warp_idx][i][count][lane] = + 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(counts + row, 1); - if (slot < static_cast(kCandGlobalCap)) { + const int slot = atomicAdd(row_counts + row, 1); + if (slot < static_cast(kCandidatesPerRow)) { const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - kCandGlobalCap + - slot; - fused_candidates::store_candidate(vals + off, indices + off, - bq, kv_offset); - } - } - } - } - } - if constexpr (kLin) { - // 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(0xffffffffu, any_pass)) [[unlikely]] { -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) { - const float bp = bp4[i]; - if (bp > gate_score[i]) { // gate first; range / finiteness only - // for passing keys - const float bq = bp * row_s[i]; - const uint32_t row = row_base2 + 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 = - fused_candidates::candidate_fp24_code(bq); - hdata.records[produce_epoch & 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(counts + row, 1); - if (slot < static_cast(kCandGlobalCap)) { - const uint64_t off = - static_cast(int64_t(row) + cand_row_delta) * - kCandGlobalCap + - slot; - fused_candidates::store_candidate( - vals + off, indices + off, bq, kv_offset); - } + candidate_offset(row, cand_row_delta, slot); + store_candidate(vals + off, indices + off, bq, kv_offset); } } } @@ -813,307 +529,173 @@ if (warp_idx == 8) { } } } - if (!kSmem && - (((kv_split_idx + 1) & (Interval - 1)) == 0 || + if ((((kv_split_idx + 1) & (Interval - 1)) == 0 || kv_split_idx + 1 == splits) && kv_split_idx + 1 != splits) { #pragma unroll - for (unsigned i = 0; i < 4; ++i) { - hdata.counts[produce_epoch & 1][warp_idx][i][lane] = lane_counts[i]; + for (uint32_t i = 0; i < 4; ++i) { + hdata.num_records[produce_interval & 1][warp_idx][i][lane] = + lane_counts[i]; lane_counts[i] = 0; } - if (lane == 0) hdata.bases[produce_epoch & 1][warp_idx] = flush_base; + if (lane == 0) hdata.bases[produce_interval & 1][warp_idx] = flush_base; __syncwarp(); - h_full[(produce_epoch & 1) * 8 + warp_idx].arrive(); - ++produce_epoch; + h_full[(produce_interval & 1) * 8 + warp_idx].arrive(); + ++produce_interval; flush_base = kv_split_idx + 1; } } - if constexpr (!kSmem) { + { // 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 unsigned bank = produce_epoch & 1, rbase = row_base2 + query_base; + const uint32_t bank = produce_interval & 1, + rbase = scan_row_base + query_base; static_assert(sizeof(hdata.records[0][0]) == 4 * 8 * 32 * 4, "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], counts, - rbase, kv_base + (math_tid % 128), flush_base, - cand_row_delta, vals, indices); + 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 (unsigned row = 0; row < 4; ++row) lane_counts[row] = 0; + for (uint32_t row = 0; row < 4; ++row) lane_counts[row] = 0; __syncwarp(); h_full[bank * 8 + warp_idx].arrive(); - ++produce_epoch; - } - if constexpr (kSmem) { -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) { - auto* base = region; - for (uint32_t slot = min(seg_cnt[i], seg_size) + lane; slot < seg_size; - slot += 32) { - const uint32_t at = - (query_base + i) * logits_stride + seg_warp * seg_size + slot; - reinterpret_cast(base)[at] = 0; - base[cand_ext * 2 + at] = 0; - reinterpret_cast(base + cand_ext * 3)[at] = 0; - } - if (lane == 0) atomicAdd(&row_count[query_base + i], seg_cnt[i]); - } - cutlass::arch::NamedBarrier(256, 3).sync(); - for (uint32_t q = 0; q < 8; ++q) { - auto* base = region; - for (uint32_t p = pool_cnt[q] + math_tid; p < pool_size; p += 256) { - const uint32_t at = q * logits_stride + 4 * seg_size + p; - reinterpret_cast(base)[at] = 0; - base[cand_ext * 2 + at] = 0; - reinterpret_cast(base + cand_ext * 3)[at] = 0; - } - } - if (math_tid < 8) { - const uint32_t row = row_base2 + math_tid; - const uint32_t spilled = - pool_cnt[math_tid] > pool_size - ? min(pool_cnt[math_tid] - pool_size, kCandGlobalCap - 1) - : 0u; - __stcg(indices + - static_cast(int64_t(row) + cand_row_delta) * - kCandGlobalCap + - kCandGlobalCap - 1, - static_cast(spilled)); - if (spilled) atomicOr(&row_count[math_tid], 1u << 21); - } - cutlass::arch::NamedBarrier(256, 3).sync(); - if (math_tid < 8) - atomicAdd(counts + row_base2 + math_tid, - static_cast(row_count[math_tid])); + ++produce_interval; } cutlass::arch::fence_view_async_shared(); empty_q.arrive(); - if constexpr (!kSmem) - cutlass::arch::NamedBarrier(320, 4).sync(); // drains of this block done + // Drains of this block done. + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).sync(); // ---- selection ---- - const unsigned vblock3 = - *reinterpret_cast(&e.vblock), - row_base3 = vblock3 * 8; + // Re-read as well (not kept live across the scan). + const uint32_t sel_vblock = + *reinterpret_cast(&e.vblock), + sel_row_base = sel_vblock * kBlockRows; __threadfence_block(); - cutlass::arch::NamedBarrier(256, 5).sync(); - const ClaimEntry& ce = entries[k & 1]; + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + const ClaimEntry& sel_entry = entries[k & 1]; if (math_tid == 0) { - unsigned recover = 0, short_rows = 0; + uint32_t recover = 0, short_rows = 0; row_offset[0] = 0; - for (unsigned r = 0; r < 8; ++r) { - const bool is_short = ce.ends[r] - ce.starts[r] <= topk; - short_rows |= unsigned(is_short) << r; - const int c = - kSmem ? int(row_count[r]) : __ldcg(counts + row_base3 + r); - const int real = kSmem ? (c & ((1 << 21) - 1)) : c; - recover |= !is_short && (real < int(topk) || real > 16384); - // > 16K rows start 16-byte aligned for select_staged_row's vector - // loads. + 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 - : (!kSmem ? (unsigned(max(c, 0)) + 3u) & ~3u - : unsigned(max(c, 0)))); - } - unsigned staged = 1; - if (!kSmem) { - const unsigned half = max(row_offset[4], row_offset[8] - row_offset[4]); - staged = recover ? 0u - : (8u * fast_select::kBins * 4u + row_offset[8] * 4u <= - phase_bytes - ? 1u - : (4u * fast_select::kBins * 4u + half * 4u <= - phase_bytes - ? 2u - : 0u)); - if (staged == 0) - recover = 1; // candidates exceed the scratch: exact rescan + row_offset[r] + (is_short ? 0 : (uint32_t(max(c, 0)) + 3u) & ~3u); } + // 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[4], row_offset[8] - row_offset[4]); + const uint32_t staged = + recover + ? 0u + : (8u * selection::kBins * 4u + row_offset[8] * 4u <= phase_bytes + ? 1u + : (4u * selection::kBins * 4u + half * 4u <= phase_bytes + ? 2u + : 0u)); + if (staged == 0) recover = 1; if (recover) { - defer_next[vblock3] = int(defer_head); - defer_head = vblock3; + 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_s = 0; + unsafe_rows = 0; } - cutlass::arch::NamedBarrier(256, 5).sync(); - const unsigned recover = sel_recover, short_rows = sel_short; - const auto write_short = [&](unsigned r, unsigned t, unsigned stride) { + 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; - const unsigned s0 = ce.starts[r - row_base3], - visible = ce.ends[r - row_base3] - s0; - int* o = fallback.output + (uint64_t)r * topk; - for (unsigned j = t; j < topk; j += stride) - o[j] = j < visible ? int((relative ? 0u : s0) + j) : -1; + 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 = int64_t(ce.slot) * 8 - int64_t(row_base3); - if constexpr (kSmem) { - const unsigned warp = warp_idx; - if (row_base3 + warp < rows) { - const unsigned r = row_base3 + warp; - int* o = fallback.output + (uint64_t)r * topk; - if (short_rows >> warp & 1u) + const int64_t row_delta = slot_row_delta(sel_entry.slot, sel_row_base); + const uint32_t per_round = sel_staged == 1 ? 8u : 4u; + for (uint32_t round = 0; round < 8 / 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 uint16_t* hi = - reinterpret_cast(region) + warp * kCandCap; - const uint8_t* lo = region + kCandExt * 2 + warp * kCandCap; - const uint16_t* idx = - reinterpret_cast(region + kCandExt * 3) + - warp * kCandCap; - const auto key_of = [idx](unsigned e2) { return int32_t(idx[e2]); }; - unsigned char* scratch = smem_buffer + kSelB + warp * kScratchS; - const int c = int(row_count[warp]); - if (c >> 21 & 1) { - const uint64_t gbase = uint64_t(int64_t(r) + row_delta) * 16384; - const uint16_t* gv = vals + gbase; - const int* gi = indices + gbase; - const unsigned total = kCandCap + unsigned(__ldcg(gi + 16383)); - fast_select::select_hist_generic( - [&](unsigned e2) { - return e2 < kCandCap - ? (unsigned(hi[e2]) << 8 | lo[e2]) - : (unsigned(__ldcg(gv + e2 - kCandCap)) << 8 | - (unsigned(__ldcg(gi + e2 - kCandCap)) >> 16)); - }, - total, o, topk, reinterpret_cast(scratch), - [&](unsigned e2) { - return e2 < kCandCap - ? int32_t(idx[e2]) - : int32_t(__ldcg(gi + e2 - kCandCap) & 0xffff); - }); - } else if (!row_select::select_shared_row( - hi, lo, idx, kCandCap, unsigned(c & ((1 << 21) - 1)), - o, topk, scratch)) { - fast_select::select_resident(hi, lo, kCandCap, 0u, o, topk, - reinterpret_cast(scratch), - key_of); - } + 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); } } - } else { - const unsigned per_round = sel_staged == 1 ? 8u : 4u; - for (unsigned round = 0; round < 8 / per_round; ++round) { - if (warp_idx < per_round) { - const unsigned local = round * per_round + warp_idx, - first = round * per_round; - unsigned* hist = reinterpret_cast(smem_buffer + kSelB) + - warp_idx * fast_select::kBins; - unsigned* codes = reinterpret_cast(smem_buffer + kSelB) + - per_round * fast_select::kBins + - (row_offset[local] - row_offset[first]); - const unsigned r = row_base3 + local; - const uint64_t base = uint64_t(int64_t(r) + row_delta) * 16384; - if (short_rows >> local & 1u) - write_short(r, lane, 32); - else { - const unsigned cnt = __ldcg(counts + r); - int* o = fallback.output + (uint64_t)r * topk; - if (!row_select::select_staged_row(codes, cnt, vals + base, - indices + base, o, topk, hist)) - fast_select::select_row(codes, cnt, vals + base, indices + base, - o, topk, hist); - } - } - cutlass::arch::NamedBarrier(256, 5).sync(); - } - } - } - cutlass::arch::NamedBarrier(256, 5).sync(); - if (relative && !recover) { - const unsigned quads = topk / 4; - for (unsigned i = math_tid; i < 8 * quads; i += 256) { - const unsigned r = row_base3 + i / quads; - if (r >= rows || ce.ends[i / quads] - ce.starts[i / quads] <= topk) - continue; - const int s = int(starts_g[r]); - if (s == 0) continue; - int4* p = - reinterpret_cast(fallback.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; + 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; + }); __syncwarp(); if (lane == 0) sel_done.arrive(); claim_empty[k & 1].arrive(); } } -// ---------------- exact rescan of the blocks queued by the selection -// ---------------- +// -------- exact rescan of the blocks queued by the selection -------- __syncthreads(); if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, 512); if (defer_head != kTerm) { - if constexpr (kMathRegs != 224) { - // The fallback core expects the 224 / 56 split. - if (warp_idx < 8) cutlass::arch::warpgroup_reg_dealloc<224>(); - __syncthreads(); - if (warp_idx >= 8) cutlass::arch::warpgroup_reg_alloc<56>(); - } + // The rescan expects the 224 / 56 register split of the math warps. __syncthreads(); - for (unsigned cur = defer_head; cur != kTerm;) { - const unsigned vblock = cur, row_base = vblock * 8; - if (threadIdx.x < 8) { - const unsigned r = row_base + threadIdx.x; + 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; } - if (threadIdx.x == 0) - fslot = 2 * blockIdx.x; // no block is in flight during the rescan __syncthreads(); - const unsigned* const starts = row_bounds[0] - row_base; - const unsigned* const ends = row_bounds[1] - row_base; - const int64_t row_delta = int64_t(fslot) * 8 - int64_t(row_base); - FallbackBuffers fb = fallback; - fb.retained += row_delta * 4096; - fb.indices += row_delta * 4096; - fused_mqa_fallback::rescan_block( - fused_mqa_dense::Scheduler(vblock, rows, keys, starts, ends), fb, topk, - tq, tk, ts, tw); + 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) { - auto& fst = - *reinterpret_cast(smem_buffer); - for (auto& bb : fst.full_q_barriers) mbar_inval(&bb); - for (auto& bb : fst.empty_q_barriers) mbar_inval(&bb); - for (auto& bb : fst.full_kv_barriers) mbar_inval(&bb); - for (auto& bb : fst.empty_kv_barriers) mbar_inval(&bb); - for (auto& bb : fst.full_tmem_barriers) mbar_inval(&bb); - for (auto& bb : fst.empty_tmem_barriers) mbar_inval(&bb); - } + if (threadIdx.x == 0) + dense::invalidate_barriers( + *reinterpret_cast(smem_buffer)); __syncthreads(); - if (relative) { - const unsigned quads = topk / 4; - for (unsigned i = threadIdx.x; i < 8 * quads; i += 384) { - const unsigned r = row_base + i / quads; - if (r >= rows) continue; - const int s = int(starts_g[r]); - if (s == 0) continue; - int4* p = - reinterpret_cast(fallback.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; - } - } - cur = unsigned(__ldcg(defer_next + vblock)); + // 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)); __syncthreads(); } __syncthreads(); diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh index 6d9fa2f3c6ec..afacbc54370b 100644 --- a/csrc/fused_mqa_topk/persistent_kernel.cuh +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -7,16 +7,20 @@ // blocks. // // Roles (384 threads): -// warp 8 claimer + Q/weights TMA (+ drain of math WG1 rows, > 16K keys) +// 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 drain of math WG0 rows (> 16K keys), idle otherwise +// warp 11 linear-term rows + drain of math WG0 rows // warps 0-7 math: sample histogram -> gate -> scan/emit -> selection -// A block whose candidate count leaves [K, 16384] pauses the role loops; the -// CTA rescans it exactly (fallback.cuh) and resumes the loops with freshly -// initialised barriers. +// 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 @@ -29,74 +33,77 @@ #include #include "candidates.cuh" +#include "common.cuh" +#include "constants.h" #include "deep_gemm_compat.cuh" -#include "fallback.cuh" -#include "fast_select.cuh" -#include "row_select.cuh" -#include "sample.cuh" +#include "rescan.cuh" +#include "select.cuh" -namespace fused_persist { +namespace fused_mqa_topk::persist { using namespace deep_gemm; +using rescan::RescanBuffers; +using rescan::RescanShared; using TBar = cutlass::arch::ClusterTransactionBarrier; -constexpr unsigned kTerm = 0xffffffffu, kPause = 0xfffffffeu; -// Selection scratch starts after Q + weights so the next block's Q can land -// during selection. -constexpr unsigned kSelBase = 33792; -constexpr unsigned kScratchS = - 8192; // per-warp selection scratch (<= 16K keys) -constexpr unsigned kCandCap = 3072; -constexpr unsigned kCandExt = 8u * kCandCap + 256u; -constexpr unsigned kQBytes = 8 * 32 * 128, kKVBytes = 256 * 128, - kWBytes = 8 * 32 * 4; +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, kept from the experiments); 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); -template -struct PStorage { - alignas(1024) cutlass::float_e4m3_t q[kQBytes]; - alignas(128) float weights[8 * 32]; - alignas(1024) cutlass::float_e4m3_t kv[KVS][kKVBytes]; - // <= 16K: the per-key scales have their own ring, twice as deep as the KV - // stages (see BarIdx::sf_full). - alignas(128) float sf_kv[2 * KVS][256]; -}; -// > 16K keys (rebalanced scan): the 16 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]). -template <> -struct PStorage<3> { - alignas(1024) uint8_t lin_lo[16 * 128]; +// 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[16 * 128]; - alignas(128) float weights[8 * 32]; - float lin_scale[8]; - alignas(1024) cutlass::float_e4m3_t kv[3][kKVBytes]; - alignas(128) float sf_kv[3][256]; + 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(PStorage<3>, q) == 16 * 128 && - offsetof(PStorage<3>, lin_hi) == 16 * 128 + kQBytes, +static_assert(offsetof(KernelSmem, q) == kLinRows * kHeadDim && + offsetof(KernelSmem, lin_hi) == kLinRows * kHeadDim + kQBytes, "lin rows must bracket Q"); -static_assert(offsetof(PStorage<2>, kv) == kSelBase, - "selection scratch base (<= 16K)"); -// > 16K: selection scratch starts after lin + Q + lin + weights + row scales -// (the next block's linear term is written during this block's selection). -constexpr unsigned kSelBase3 = offsetof(PStorage<3>, kv); +// 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(PStorage<3>& sm, - unsigned lane) { +__device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, + uint32_t lane) { const uint8_t* qbytes = reinterpret_cast(sm.q); #pragma unroll 1 - for (unsigned r = 0; r < 8; ++r) { + for (uint32_t r = 0; r < 8; ++r) { float acc[4] = {0.f, 0.f, 0.f, 0.f}; #pragma unroll 8 - for (unsigned h = 0; h < 32; ++h) { - const unsigned R = r * 32 + h; - const unsigned off = R * 128u + ((((lane * 4u) >> 4) ^ (R & 7u)) << 4) + - ((lane * 4u) & 15u); + for (uint32_t h = 0; h < 32; ++h) { + const uint32_t off = swizzle128_offset(r * 32 + h, lane * 4u); const uint32_t packed = ptx::ld_shared(reinterpret_cast(qbytes + off)); const float wv = ptx::ld_shared(sm.weights + r * 32 + h) * 0.5f; @@ -114,7 +121,7 @@ __device__ __forceinline__ void compute_lin_rows(PStorage<3>& sm, float m = fmaxf(fmaxf(fabsf(acc[0]), fabsf(acc[1])), fmaxf(fabsf(acc[2]), fabsf(acc[3]))); #pragma unroll - for (unsigned o = 16; o > 0; o >>= 1) + for (uint32_t o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); int e = -120; // pow2 scale s >= m / 448, clamped so 1/s stays finite if (m > 0.f && isfinite(m)) { @@ -124,10 +131,10 @@ __device__ __forceinline__ void compute_lin_rows(PStorage<3>& sm, const float s = ldexpf(1.0f, e), inv_s = ldexpf(1.0f, -e); float res[4]; #pragma unroll - for (unsigned dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s; + for (uint32_t dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s; uint8_t* lin_base = r < 4 ? sm.lin_lo : sm.lin_hi; #pragma unroll - for (unsigned tt = 0; tt < 4; ++tt) { + 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( @@ -146,14 +153,12 @@ __device__ __forceinline__ void compute_lin_rows(PStorage<3>& sm, // 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] *= 256.0f; - res[1] *= 256.0f; - res[2] *= 256.0f; - res[3] *= 256.0f; + res[0] *= 1.0f / kLinLoScale; + res[1] *= 1.0f / kLinLoScale; + res[2] *= 1.0f / kLinLoScale; + res[3] *= 1.0f / kLinLoScale; } - const unsigned n = (r % 4u) * 4u + tt; - const unsigned off = n * 128u + ((((lane * 4u) >> 4) ^ (n & 7u)) << 4) + - ((lane * 4u) & 15u); + const uint32_t off = swizzle128_offset((r % 4u) * 4u + tt, lane * 4u); ptx::st_shared( reinterpret_cast(lin_base + off), static_cast(c01) | (static_cast(c23) << 16)); @@ -162,97 +167,115 @@ __device__ __forceinline__ void compute_lin_rows(PStorage<3>& sm, } } +// Gate sample state: each math warp histograms its row's sampled scores +// (half_bin<5>: kSampleBins bins, padded) and finish_sample_histogram picks the +// gate from them. +constexpr uint32_t kSampleBins = 2048; +struct alignas(16) SampleShared { + uint32_t hist[kBlockRows][padded_bins(kSampleBins)]; +}; + // Gate: the k-th largest sample (k = count * target / visible, count = the // histogram total) of one row from its 2048-bin histogram; one warp per row. // -inf when the row sees at most `target` keys or has no samples. -__device__ __noinline__ void finish_sample_histogram(const unsigned* hist, - float* gate, - unsigned visible, - unsigned target, - unsigned lane) { +__device__ __noinline__ inline void finish_sample_histogram( + const uint32_t* hist, 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 the 64 bins of chunks 63-2l and 62-2l (bins descending with the - // lane). - const unsigned* const h0 = hist + (63 - 2 * lane) * 33; - const unsigned* const h1 = hist + (62 - 2 * lane) * 33; - unsigned sum = 0; + // Lane l owns the 64 bins of 32-bin chunks 63-2l and 62-2l (bins descending + // with the lane). + const uint32_t* const h0 = hist + padded_bin((63 - 2 * lane) * 32); + const uint32_t* const h1 = hist + padded_bin((62 - 2 * lane) * 32); + uint32_t sum = 0; #pragma unroll 8 - for (unsigned j = 0; j < 32; ++j) sum += h0[j] + h1[j]; - unsigned incl = sum; -#pragma unroll - for (unsigned dd = 1; dd < 32; dd *= 2) { - const unsigned u = __shfl_up_sync(0xffffffffu, incl, dd); - if (lane >= dd) incl += u; - } - const unsigned tot = __shfl_sync(0xffffffffu, incl, 31); + for (uint32_t j = 0; j < 32; ++j) sum += h0[j] + h1[j]; + const uint32_t incl = warp_inclusive_scan(sum, lane); + const uint32_t tot = __shfl_sync(0xffffffffu, incl, 31); if (tot == 0) { if (lane == 0) *gate = -CUDART_INF_F; __syncwarp(); return; } - const unsigned rank = + const uint32_t rank = max(1u, min(tot, (tot * target + visible - 1) / visible)); - const unsigned w = __ffs(__ballot_sync(0xffffffffu, incl >= rank)) - 1; - unsigned above = __shfl_sync(0xffffffffu, incl - sum, w), chosen = 0; + const uint32_t w = __ffs(__ballot_sync(0xffffffffu, incl >= rank)) - 1; + uint32_t above = __shfl_sync(0xffffffffu, incl - sum, w), chosen = 0; #pragma unroll 1 - for (unsigned cc = 0; cc < 2; ++cc) { - const unsigned c = 63 - 2 * w - cc; - unsigned v = hist[c * 33 + 31 - lane]; // lane 0 = highest bin of the chunk -#pragma unroll - for (unsigned dd = 1; dd < 32; dd *= 2) { - const unsigned u = __shfl_up_sync(0xffffffffu, v, dd); - if (lane >= dd) v += u; - } - const unsigned mask = __ballot_sync(0xffffffffu, above + v >= rank); + for (uint32_t cc = 0; cc < 2; ++cc) { + const uint32_t c = 63 - 2 * w - cc; + // Lane 0 = highest bin of the chunk. + const uint32_t v = + warp_inclusive_scan(hist[padded_bin(c * 32 + 31 - lane)], lane); + const uint32_t mask = __ballot_sync(0xffffffffu, above + v >= rank); if (mask) { chosen = c * 32 + 31 - (__ffs(mask) - 1); break; } above += __shfl_sync(0xffffffffu, v, 31); } - if (lane == 0) *gate = sample_floor(chosen); + if (lane == 0) *gate = half_bin_floor<5>(chosen); __syncwarp(); } -struct alignas(16) PHandoffData { - uint32_t records[2][8][4][8][32]; - uint32_t counts[2][8][4][32]; +// 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][8][4][kRingSlots][32]; + uint32_t num_records[2][8][4][32]; uint32_t bases[2][8]; }; // samples: gate sample tiles of the block (0 when no row sees more keys than // the gate target, so every gate is -inf). -struct ClaimEntry { - unsigned vblock, all_short, kv_base, splits, slot, samples; - unsigned starts[8], ends[8]; +struct ClaimHeader { + uint32_t vblock, all_short, kv_base, splits, slot, samples; +}; +struct ClaimEntry : ClaimHeader { + uint32_t starts[8], ends[8]; }; +// 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). -template struct BarIdx { - // > 16K (KVS == 3): 3 per-warpgroup TMEM slots + the linear-term barrier. - static constexpr unsigned NT = KVS == 3 ? 3 : 2; - static constexpr unsigned full_q = 0, empty_q = 1, full_kv = 2, - empty_kv = 2 + KVS, full_tmem = 2 + 2 * KVS, + // TMEM slots, shared by the two math warpgroups (WG0 even, WG1 odd + // positions). + static constexpr uint32_t NT = 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 + NT, claim_full = empty_tmem + NT, claim_empty = claim_full + 2, sel_done = claim_empty + 2, h_full = sel_done + 1, h_empty = h_full + 16, full_lin = h_empty + 16, - sf_full = full_lin + (KVS == 3 ? 1 : 0), - total = sf_full + (KVS == 3 ? 0 : 2 * KVS); + total = full_lin + 1; }; -__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"); -} - // 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, @@ -268,177 +291,211 @@ __device__ __forceinline__ void mbar_arrive_expect_tx(const void* b, : "memory"); } -// > 16K self-drain, out of line: stores this warp's ring records of the last -// handoff interval. recs = records[bank][warp] ([4][8][32]). -__device__ __noinline__ void self_drain_last( - const uint32_t* __restrict__ recs, unsigned n0, unsigned n1, unsigned n2, - unsigned n3, int* __restrict__ counts, unsigned rbase, uint32_t lane_kv, - unsigned flush_base, int64_t row_delta, uint16_t* __restrict__ vals, - int32_t* __restrict__ indices) { - constexpr uint32_t kCap = 16384; - const unsigned lane = threadIdx.x & 31; - const unsigned n[4] = {n0, n1, n2, n3}; - unsigned my_total = 0; +// 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 < 8 * 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][8 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. +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 (unsigned row = 0; row < 4; ++row) { - const unsigned total = __reduce_add_sync(0xffffffffu, n[row]); + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t total = __reduce_add_sync(0xffffffffu, n[row]); if (lane == row) my_total = total; } - unsigned my_start = 0; + uint32_t my_start = 0; if (lane < 4 && my_total) - my_start = atomicAdd(counts + rbase + lane, my_total); - unsigned p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); + 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 unsigned a = __shfl_up_sync(0xffffffffu, p01, dd), + const uint32_t a = __shfl_up_sync(0xffffffffu, p01, dd), b = __shfl_up_sync(0xffffffffu, p23, dd); if (lane >= dd) { p01 += a; p23 += b; } } - const unsigned prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; + const uint32_t prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; #pragma unroll - for (unsigned row = 0; row < 4; ++row) { - const unsigned start = __shfl_sync(0xffffffffu, my_start, row); - for (unsigned j = 0; j < n[row]; ++j) { - const unsigned slot2 = start + prefix[row] - n[row] + j; - const unsigned record = recs[(row * 8 + j) * 32 + lane]; - if (slot2 < kCap) { - const uint32_t physical = - lane_kv + (flush_base * 2 + (record & 255u)) * 128; - const uint64_t off = - static_cast(int64_t(rbase + row) + row_delta) * kCap + - slot2; - fused_candidates::store_candidate_record(vals + off, indices + off, - record >> 8, physical); - } + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t start = __shfl_sync(0xffffffffu, 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 ho_drain(PHandoffData& d, TBar* hfull, - TBar* hempty, int producer, - unsigned epoch, unsigned row_base, - int* totals, unsigned cap, bool skip, - Store store) { - const int lane = threadIdx.x & 31, bank = epoch & 1; - hfull[bank * 8 + producer].wait((epoch >> 1) & 1); +__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 * 8 + producer].wait((interval >> 1) & 1); __syncwarp(); - if (skip) { + if (skip) { // self-drained by the producer hempty[bank * 8 + producer].arrive(); return; - } // self-drained by the producer - const unsigned base = d.bases[bank][producer]; - unsigned n[4]; - unsigned my_total = 0; -#pragma unroll - for (unsigned row = 0; row < 4; ++row) { - n[row] = d.counts[bank][producer][row][lane]; - unsigned total = __reduce_add_sync(0xffffffffu, n[row]); - if (lane == row) my_total = total; - } - unsigned my_start = 0; - if (lane < 4 && my_total) - my_start = atomicAdd(totals + row_base + lane, my_total); - unsigned p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); -#pragma unroll - for (int dd = 1; dd < 32; dd *= 2) { - unsigned a = __shfl_up_sync(0xffffffffu, p01, dd); - unsigned b = __shfl_up_sync(0xffffffffu, p23, dd); - if (lane >= dd) { - p01 += a; - p23 += b; - } } - unsigned prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; + const uint32_t base = d.bases[bank][producer]; + uint32_t n[4]; #pragma unroll - for (unsigned row = 0; row < 4; ++row) { - unsigned start = __shfl_sync(0xffffffffu, my_start, row); - for (unsigned j = 0; j < n[row]; ++j) { - unsigned slot = start + prefix[row] - n[row] + j; - unsigned record = d.records[bank][producer][row][j][lane]; - if (slot < cap) - store(row_base + row, slot, record, base, producer * 32 + lane); - } - } + 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 * 8 + producer].arrive(); } -template -__global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( - unsigned rows, unsigned keys, const unsigned* starts_g, - const unsigned* ends_g, uint16_t* vals, int* indices, int* counts, - FallbackBuffers fallback, unsigned phase_bytes, unsigned sample_target, - unsigned sample_stride, unsigned sample_groups, unsigned sample_tiles, - const float* scales, const __grid_constant__ cute::TmaDescriptor tks, +// 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, bool relative, - unsigned nblocks, unsigned* counter, int* defer_next) { - constexpr bool kSmem = - Interval <= 8; // <= 16K keys: shared candidates, direct emission - constexpr uint32_t kMathRegs = kSmem ? 232 : 224, - kSpecRegs = 168 * 3 - 2 * kMathRegs; - constexpr uint32_t kCandGlobalCap = 16384; - using St = PStorage; - using B = BarIdx; - // Compile-time: a runtime top-k costs ~0.5% in the > 16K scan (codegen). - constexpr unsigned topk = TopK; - // > 16K keys: rebalanced scan (linear term on the tensor core, per-warpgroup - // N=144 MMAs). - constexpr bool kLin = !kSmem; - static_assert(kLin == (KVS == 3), - "rebalanced scan needs the 3-KV-stage lin layout"); - constexpr uint32_t kSlotCols = - 160; // 4 rows x 32 heads + 16 linear-term columns - constexpr unsigned kSelB = KVS == 3 ? kSelBase3 : kSelBase; + 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::NT * 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(St); - auto& stream = *reinterpret_cast(region); - auto& hdata = *reinterpret_cast(region); + 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[B::total]; + __shared__ alignas(8) uint64_t bar_words[BarIdx::total]; __shared__ uint32_t tmem_ptr; __shared__ ClaimEntry entries[2]; __shared__ float gate_s[8]; - __shared__ unsigned row_count[8], pool_cnt[8]; - __shared__ unsigned sel_recover, sel_staged, sel_short, row_offset[9]; - __shared__ unsigned defer_head, unsafe_s, pause_req, epoch_kind; + __shared__ uint32_t sel_recover, sel_staged, sel_short, row_offset[9]; + __shared__ uint32_t defer_head, unsafe_rows, pause_req, epoch_kind; TBar* const bars = reinterpret_cast(bar_words); - TBar& full_q = bars[B::full_q]; - TBar& empty_q = bars[B::empty_q]; - TBar* const full_kv = bars + B::full_kv; - // > 16K the 32 producer lanes stage the sample scales, so each arrives on - // full_kv. - constexpr unsigned kFullKvArrivals = KVS == 3 ? 32 : 1; - // <= 16K: the MMA warp frees each KV stage once its MMAs complete, and the - // per-key scales arrive through a separate ring (2 * KVS deep, sf_full), so - // the math warps never wait on or release KV stages. The producer leads the - // MMA by <= KVS tiles and the MMA leads the math warps by <= 1 tile (2 TMEM - // stages), so a scale stage is never overwritten while being read. - constexpr bool kRel = KVS == 2; - TBar* const sf_full = bars + B::sf_full; - TBar* const empty_kv = bars + B::empty_kv; - TBar* const full_tmem = bars + B::full_tmem; - TBar* const empty_tmem = bars + B::empty_tmem; - TBar* const claim_full = bars + B::claim_full; - TBar* const claim_empty = bars + B::claim_empty; - TBar& sel_done = bars[B::sel_done]; - TBar* const h_full = bars + B::h_full; - TBar* const h_empty = bars + B::h_empty; - TBar* const full_lin = bars + B::full_lin; // > 16K only + 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 unsigned warp_idx = cutlass::canonical_warp_idx_sync(); - const unsigned lane = ptx::get_lane_idx(); - constexpr unsigned kClaimConsumers = - kSmem ? 10 : 11; // warps 9, 10, 0-7 (+11) + 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); @@ -452,20 +509,20 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( 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 with - // the fallback core 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 >16K scan: uniform address math is - // rematerialised every tile). Later epochs only run after a recovery pause. - __shared__ unsigned row_bounds[2][8]; - __shared__ unsigned fslot; + // 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][8]; bool more = false; { - constexpr unsigned epoch = 0; + constexpr uint32_t epoch = 0; #include "persistent_epoch.inc" } #pragma unroll 1 - for (unsigned epoch = 1; more; ++epoch) { + for (uint32_t epoch = 1; more; ++epoch) { #include "persistent_epoch.inc" } // counter = {next block, finished CTAs}: the last CTA to finish (every claim @@ -480,98 +537,4 @@ __global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_prefill_kernel( } } -// ---------------- host: launch configuration (> kResidentKeys keys) -// ---------------- -using PersistentKernel = - decltype(&fused_mqa_topk_fp8_prefill_kernel<4, 2, 2048>); - -template -PersistentKernel persistent_kernel(unsigned keys) { - return keys <= 8192 ? fused_mqa_topk_fp8_prefill_kernel<4, 2, TopK> - : keys <= 16384 ? fused_mqa_topk_fp8_prefill_kernel<8, 2, TopK> - : keys <= 32768 ? fused_mqa_topk_fp8_prefill_kernel<16, 3, TopK> - : keys <= 131072 ? fused_mqa_topk_fp8_prefill_kernel<32, 3, TopK> - : keys <= 524288 ? fused_mqa_topk_fp8_prefill_kernel<64, 3, TopK> - : fused_mqa_topk_fp8_prefill_kernel<128, 3, TopK>; -} - -struct LaunchConfig { - PersistentKernel kernel; - int smem; // dynamic shared memory bytes - unsigned 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). - unsigned sample_target, sample_stride, sample_groups, sample_tiles; -}; - -// Picks the kernel instantiation, sizes its shared memory (setting the opt-in -// limit) and the gate sample for `keys` keys and top-`topk` rows. -inline cudaError_t launch_config(unsigned keys, unsigned 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. At - // <= 16K the candidates live in shared memory (cap kCandCap); top-2048 needs - // rank 100 there to keep 2560 candidates above 2048. 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 { - unsigned target, rank; - }; - constexpr Gate kGates[3][3] = { - // <= 16K <= 128K > 128K - {{2048, 50}, {1024, 32}, {1536, 64}}, // top-512 - {{2048, 50}, {1536, 48}, {2048, 48}}, // top-1024 - {{2560, 100}, {3072, 64}, {4096, 64}}, // top-2048 - }; - const bool small = keys <= 16384, long_ctx = keys > 131072; - const Gate gate = kGates[topk == 512 ? 0 : (topk == 1024 ? 1 : 2)] - [small ? 0 : (long_ctx ? 2 : 1)]; - c.sample_target = gate.target; - const unsigned rank = gate.rank; - const unsigned rank_tiles = std::max( - 1u, unsigned((uint64_t(rank) * keys + 256ull * c.sample_target - 1) / - (256ull * c.sample_target))); - c.sample_tiles = - long_ctx ? std::min(32u, std::max(16u, rank_tiles)) : rank_tiles; - c.sample_stride = std::max( - 1u, (keys + 256u * c.sample_tiles - 1) / (256u * 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); - - // The recovery rescan reuses the same dynamic shared memory with the fallback - // core's layout. - constexpr int kFallbackSmem = - int(sizeof(fused_mqa_dense::SharedStorage) + sizeof(FallbackShared)); - constexpr int kSmallSmem = int(sizeof(PStorage<2>) + kCandExt * 5u); - constexpr int kLongSmem = - int(sizeof(PStorage<3>) + - std::max(sizeof(PHandoffData), sizeof(SampleShared))); - static_assert(sizeof(SampleShared) <= kCandExt * 5u, - "sample state must fit the candidate region"); - constexpr int kMaxSmem = 232448; // sm_100 opt-in limit per block - if (small) { - c.smem = std::max(kSmallSmem, kFallbackSmem); - } else { - // > 16K: the staged selection may use everything past Q + weights, up to - // the opt-in limit. - cudaFuncAttributes attr; - if (const cudaError_t e = cudaFuncGetAttributes( - &attr, reinterpret_cast(c.kernel)); - e != cudaSuccess) - return e; - c.smem = std::max(kFallbackSmem, kMaxSmem - int(attr.sharedSizeBytes)); - if (c.smem < kLongSmem) return cudaErrorInvalidConfiguration; - } - c.phase_bytes = unsigned(c.smem) - (small ? kSelBase : kSelBase3); - return cudaFuncSetAttribute(reinterpret_cast(c.kernel), - cudaFuncAttributeMaxDynamicSharedMemorySize, - c.smem); -} - -} // namespace fused_persist +} // 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..0937d1997a8f --- /dev/null +++ b/csrc/fused_mqa_topk/rescan.cuh @@ -0,0 +1,338 @@ +// SPDX-License-Identifier: MIT +// 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][8][kSegmentKeys]; + uint32_t hist[8][padded_bins(kSegmentBins)], gates[8], kept[8]; + uint32_t work[padded_bins(kRadixBins)], prefix, rank, emit_counts[8]; + __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 < 8) { + gates[tid] = __float_as_uint(-CUDART_INF_F); + kept[tid] = 0; + } + } +}; +__device__ inline void compact_retained(RescanShared& shared, + RescanBuffers buffers, uint32_t K, + uint32_t local, uint32_t row, + uint32_t tid, bool final = false) { + constexpr uint32_t Full = 0xffffffffu, kThreads = kMathThreads, + Items = kRetainedPerRow / kThreads; + uint32_t lane = tid & 31, count = shared.kept[local]; + uint64_t base = (uint64_t)row * kRetainedPerRow; + float values[Items]; + int indices[Items]; +#pragma unroll + for (uint32_t i = 0; i < Items; ++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; + } + // mode 0: values >= gate; 1: codes above threshold; 2: codes equal to it. + auto emit = [&](float gate, uint32_t threshold, int mode, + uint32_t base_offset) { + const auto passes = [&](float v) { + return isfinite(v) && + (mode == 0 ? v >= gate + : (mode == 1 ? ordered_code(v) > threshold + : ordered_code(v) == threshold)); + }; + uint32_t total = 0; +#pragma unroll + for (uint32_t i = 0; i < Items; ++i) total += passes(values[i]); + uint32_t warp_total = __reduce_add_sync(Full, 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 < 8; ++w) all += shared.emit_counts[w]; +#pragma unroll + for (uint32_t i = 0; i < Items; ++i) { + const float v = values[i]; + const bool pass = passes(v); + uint32_t mask = __ballot_sync(Full, pass), + pos = start + __popc(mask & ((1u << lane) - 1)); + if (pass && (mode != 2 || pos < K)) { + buffers.retained[base + pos] = v; + buffers.retained_indices[base + pos] = indices[i]; + } + start += __popc(mask); + } + math_wg_sync(); + return all; + }; + if (!final) { + 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 < Items; ++i) + kept += isfinite(values[i]) && values[i] >= gate; + kept = __reduce_add_sync(Full, kept); + if (lane == 0) atomicAdd(&shared.rank, kept); + math_wg_sync(); + kept = shared.rank; + if (kept >= K && kept <= kRetainedPerRow - kSegmentKeys) { + emit(gate, 0, 0, 0); + math_wg_sync(); + if (tid == 0) shared.kept[local] = kept; + math_wg_sync(); + return; + } + } + if (tid == 0) { + shared.prefix = 0; + shared.rank = K; + } + 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 < Items; ++i) { + uint32_t code = ordered_code(values[i]); + bool valid = isfinite(values[i]) && + (shift == 24 || (code >> (shift + 8)) == prefix); + uint32_t mask = __ballot_sync(Full, 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(Full, 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(Full, chosen, chosen_lane); + above = __shfl_sync(Full, 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, 1, 0); + emit(0, threshold, 2, greater); + if (tid == 0) { + shared.kept[local] = K; + 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 K, + uint32_t local, uint32_t row, + uint32_t segment, uint32_t first, + uint32_t n, uint32_t tid) { + constexpr uint32_t Full = 0xffffffffu, Items = kSegmentKeys / 32; + uint32_t lane = tid, bank = segment & 1; + float values[Items], gate = __uint_as_float(shared.gates[local]); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < Items; ++i) { + uint32_t j = tid + 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(Full, 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(Full, 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(Full, before < K && inclusive >= K); + if (mask) { + uint32_t chosen = __ffs(mask) - 1, group = 31 - chosen, + rank = K - __shfl_sync(Full, 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(Full, before < rank && inclusive >= rank)) - 1; + bin = __shfl_sync(Full, 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 < Items; ++i) { + float v = values[i]; + bool pass = isfinite(v) && v >= gate; + uint32_t mask = __ballot_sync(Full, pass), + pos = written + __popc(mask & ((1u << lane) - 1)); + if (pass) { + buffers.retained[base + pos] = v; + buffers.retained_indices[base + pos] = first + tid + 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 / 256; + RescanShared& shared; + RescanBuffers buffers; + int* out; // [rows, topk] + uint32_t topk, query_base, row_base, segment = 0; + uint32_t seq_k_start[4], seq_k_end[4]; + float gate[4]; + 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 < 4; ++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 < 4; ++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) * 256 + 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 < 8; ++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) * 256, tid % 32); + math_wg_sync(); + if (split + 1 == num_splits) + for (uint32_t l = 0; l < 8; ++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/resident_kernel.cuh b/csrc/fused_mqa_topk/resident_kernel.cuh deleted file mode 100644 index 41abeeea6698..000000000000 --- a/csrc/fused_mqa_topk/resident_kernel.cuh +++ /dev/null @@ -1,136 +0,0 @@ -// SPDX-License-Identifier: MIT -// <= 4K keys: every score of an 8-row block stays in shared memory as an fp24 -// code, so there is no sample, gate or candidate workspace; one warp then -// selects one row. -#pragma once - -#include "candidates.cuh" -#include "dense_scan.cuh" -#include "fast_select.cuh" - -namespace fused_mqa_resident { -using fused_mqa_dense::BLOCK_Q; -using fused_mqa_dense::Scheduler; -using fused_mqa_dense::SharedStorage; - -constexpr uint32_t kResidentKeys = 4096; -// Pipeline storage, then the block's codes: hi16 array then lo8 array, -// row-major, kResidentKeys per row. -constexpr uint32_t kSmemBytes = - (sizeof(SharedStorage) + BLOCK_Q * kResidentKeys * 3u + 1023u) / 1024u * - 1024u; - -// Stores every score as an fp24 code (selection reads only [start, end) of each -// row, so no masking). -struct CodesPolicy { - static constexpr float kWeightScale = 1.0f, kScoreScale = 0.5f; - uint8_t* codes; - uint32_t query_base; - CUTLASS_DEVICE void begin_block(const Scheduler&, uint32_t, uint32_t base) { - query_base = base; - } - CUTLASS_DEVICE void begin_split(uint32_t, uint32_t) {} - CUTLASS_DEVICE void score(uint32_t i, uint32_t kv_offset, uint32_t, uint32_t, - float score) { - if (kv_offset < kResidentKeys) { - const uint32_t code = fused_candidates::candidate_fp24_code(score); - const uint32_t at = (query_base + i) * kResidentKeys + kv_offset; - reinterpret_cast(codes)[at] = static_cast(code >> 8); - codes[BLOCK_Q * kResidentKeys * 2 + at] = static_cast(code); - } - } - CUTLASS_DEVICE void end_split(uint32_t, uint32_t, uint32_t) {} - CUTLASS_DEVICE void end_block(const Scheduler&, uint32_t) {} -}; - -} // namespace fused_mqa_resident - -// <= 4K keys: one CTA per 8-row block, heaviest (latest rows) first. -__global__ __launch_bounds__(384, 1) void fused_mqa_topk_fp8_resident_kernel( - unsigned rows, unsigned keys, const unsigned* starts_g, - const unsigned* ends_g, int* output, unsigned topk, bool relative, - const __grid_constant__ cute::TmaDescriptor tq, - const __grid_constant__ cute::TmaDescriptor ts, - const __grid_constant__ cute::TmaDescriptor tk, - const __grid_constant__ cute::TmaDescriptor tw) { - using namespace fused_mqa_resident; - const unsigned vblock = gridDim.x - 1 - blockIdx.x, row_base = vblock * 8; - // The block's row bounds clamped to the keys; rows past the end see no keys. - __shared__ unsigned row_bounds[2][8]; - if (threadIdx.x < 8) { - const unsigned 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 unsigned* const starts = row_bounds[0] - row_base; - const unsigned* const ends = row_bounds[1] - row_base; - // Rows with <= topk visible keys output all of them (relative indices start - // at 0). - const auto write_short = [&](unsigned r, unsigned t, unsigned stride) { - const unsigned visible = ends[r] - starts[r]; - for (unsigned j = t; j < topk; j += stride) - output[(uint64_t)r * topk + j] = - j < visible ? int((relative ? 0u : starts[r]) + j) : -1; - }; - bool all_short = true; - for (unsigned r = row_base; r < row_base + 8; ++r) - all_short &= ends[r] - starts[r] <= topk; - if (all_short) { - for (unsigned r = row_base; r < min(row_base + 8, rows); ++r) - write_short(r, threadIdx.x, 384); - return; - } - - extern __shared__ __align__(1024) uint8_t smem[]; - uint8_t* const codes = smem + sizeof(SharedStorage); - CodesPolicy policy{codes, 0}; - fused_mqa_dense::score_block(Scheduler(vblock, rows, keys, starts, ends), - policy, tq, tk, ts, tw); - __syncthreads(); - if (threadIdx.x < 256) cutlass::arch::warpgroup_reg_dealloc<128>(); - __syncthreads(); - if (threadIdx.x >= 256) cutlass::arch::warpgroup_reg_alloc<128>(); - const unsigned warp = threadIdx.x / 32; - if (warp < 8 && row_base + warp < rows) { - const unsigned r = row_base + warp, s0 = starts[r], visible = ends[r] - s0; - if (visible <= topk) { - write_short(r, threadIdx.x & 31, 32); - } else { - int* o = output + (uint64_t)r * topk; - const uint16_t* hi = - reinterpret_cast(codes) + warp * kResidentKeys; - const uint8_t* lo = codes + 8 * kResidentKeys * 2 + warp * kResidentKeys; - // Scratch (bracket buffer, or the fallback's histogram) reuses the idle - // Q/KV stage memory. - static_assert(8 * fast_select::kScratch3Bytes <= sizeof(SharedStorage) && - fast_select::kScratch3Bytes >= fast_select::kBins * 4, - "select scratch must fit the Q/KV stages"); - uint8_t* scratch = smem + warp * fast_select::kScratch3Bytes; - if (!fast_select::select_resident_bracket(hi, lo, s0, visible, o, topk, - scratch)) - fast_select::select_resident(hi + s0, lo + s0, visible, s0, o, topk, - reinterpret_cast(scratch)); - } - } - // Relative indices: selected rows are rewritten in place (L2-hot) once final. - if (relative) { - __syncthreads(); - const unsigned quads = topk / 4; - for (unsigned i = threadIdx.x; i < 8 * quads; i += 384) { - const unsigned r = row_base + i / quads; - if (r >= rows || ends[r] - starts[r] <= topk) 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; - } - } -} diff --git a/csrc/fused_mqa_topk/row_select.cuh b/csrc/fused_mqa_topk/row_select.cuh deleted file mode 100644 index 824bf920ebfc..000000000000 --- a/csrc/fused_mqa_topk/row_select.cuh +++ /dev/null @@ -1,415 +0,0 @@ -// SPDX-License-Identifier: MIT -// Single-warp top-K selection of one row's candidates for the persistent -// kernel. Both selections bracket the K-th largest fp24 code with 512 -// sampled codes, place the codes above the bracket directly, and resolve the -// bracket with histogram levels on register-resident codes; elements are placed -// by per-lane counts and one warp prefix scan (no per-element ballots). -#pragma once -#include "fast_select.cuh" - -namespace row_select { -using namespace fast_select; - -// r-th largest (1-based) of the 512 register samples 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 unsigned (&smp)[NS], - unsigned r_hi, unsigned r_lo, - unsigned* h8, unsigned& h_hi, - unsigned& h_lo) { - const unsigned lane = threadIdx.x & 31; -#pragma unroll - for (unsigned b = lane; b < 192; b += 32) - reinterpret_cast(h8)[b] = make_uint4(0, 0, 0, 0); - __syncwarp(); -#pragma unroll - for (unsigned u = 0; u < NS; ++u) atomicAdd(h8 + (smp[u] >> 8), 1u); - __syncwarp(); - unsigned ab_hi, ab_lo; - const unsigned b_hi = find_bin_n<256>(h8, r_hi, ab_hi); - const unsigned b_lo = find_bin_n<256>(h8, r_lo, ab_lo); - unsigned* const lh = h8 + 256; - unsigned* const ll = h8 + 512; -#pragma unroll - for (unsigned u = 0; u < NS; ++u) { - const unsigned 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(); - unsigned 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 -} - -// find_bin_n with each lane's kPer contiguous bins loaded as 16-byte vectors -// and searched in registers. -template -__device__ __forceinline__ unsigned find_bin_v(const unsigned* hist, - unsigned rank, unsigned& above) { - const unsigned lane = threadIdx.x & 31; - constexpr unsigned 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); - unsigned h[kPer]; -#pragma unroll - for (unsigned 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; - } - unsigned sum = 0; -#pragma unroll - for (unsigned j = 0; j < kPer; ++j) sum += h[j]; - unsigned inclusive = sum; -#pragma unroll - for (unsigned d = 1; d < 32; d *= 2) { - const unsigned v = __shfl_up_sync(Full, inclusive, d); - if (lane >= d) inclusive += v; - } - const unsigned before = inclusive - sum; - const unsigned winner = - __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; - // Descending walk inside the lane: bin index Bins - lane * kPer - 1 - j <-> - // h[kPer - 1 - j]. - unsigned acc = before, chosen = 0, found = 0, acc_at = before; -#pragma unroll - for (unsigned j = 0; j < kPer; ++j) { - const unsigned 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(Full, acc_at, winner); - return __shfl_sync(Full, chosen, winner); -} -// <= 16K keys: one row's shared-memory candidates (hi16 / lo8 code arrays and -// the uint16 key array idx_row, read with 16-byte loads). Bracket keys are -// preloaded into registers before the emit. Returns false (out untouched) when -// the bracket misses; the caller then runs select_resident. -__device__ bool select_shared_row(const uint16_t* __restrict__ hi_row, - const uint8_t* __restrict__ lo_row, - const uint16_t* __restrict__ idx_row, - unsigned count, unsigned real, - int32_t* __restrict__ out, unsigned K, - unsigned char* scratch) { - const unsigned lane = threadIdx.x & 31; - unsigned h_hi, h_lo; - { - unsigned smp[kSamplesPerLane], n_real = 0; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) { - smp[u] = hi_row[(((u * 32 + lane) * 2 + 1) * count >> 10)]; - n_real += smp[u] != 0; - } - const unsigned s_real = __reduce_add_sync(Full, n_real); - if (s_real == 0 || real <= K) return false; - const unsigned expect = (K * s_real + real / 2) / real; - const unsigned r_hi = expect > 40 ? expect - 40 : 1, - r_lo = min(expect + 40, s_real); - // h8 (768 words) borrows the stage area, written only after the bracket. - radix_bracket(smp, r_hi, r_lo, reinterpret_cast(scratch), h_hi, - h_lo); - h_lo = max(h_lo, 1u); - } - const unsigned units = (count + 7) / 8; - const uint4* hv = reinterpret_cast(hi_row); - const uint4* iv = reinterpret_cast(idx_row); - unsigned na = 0, nm = 0; -#pragma unroll 2 - for (unsigned unit = lane; unit < units; unit += 32) - classify8(hv[unit], unit * 8, 0u, count, h_hi, h_lo, - [&](unsigned, bool a, bool m) { - na += a; - nm += m; - }); - unsigned above, mid; - unsigned pa = warp_exclusive_scan(na, above), - pm = warp_exclusive_scan(nm, mid); - if (above >= K || above + mid < K || mid > kMid3Cap) return false; - uint16_t* stage = reinterpret_cast(scratch); - uint16_t* midbuf = reinterpret_cast(scratch + kMaxK * 2); - unsigned* hist = - reinterpret_cast(scratch + kMaxK * 2 + kMid3Cap * 2); -#pragma unroll 2 - for (unsigned unit = lane; unit < units; unit += 32) { - const uint4 h4 = hv[unit], k4 = iv[unit]; - const unsigned kw[4] = {k4.x, k4.y, k4.z, k4.w}; - classify8(h4, unit * 8, 0u, count, h_hi, h_lo, - [&](unsigned e, bool a, bool m) { - const unsigned k = e & 7u; - const unsigned key = (kw[k / 2] >> (16 * (k & 1))) & 0xffffu; - if (a) stage[pa++] = uint16_t(key); - if (m) midbuf[pm++] = uint16_t(e); - }); - } - __syncwarp(); - constexpr unsigned R = kMid3Cap / 32; - unsigned mc[R]; -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - const unsigned i = lane + 32 * u; - mc[u] = 0; - if (i < mid) { - const unsigned e = midbuf[i]; - mc[u] = unsigned(hi_row[e]) << 8 | lo_row[e]; - } - } - unsigned rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; -#pragma unroll 1 - while (top > base) { - const unsigned width = top - base + 1, spread = (width - 1) / kMidBins; - const unsigned shift = spread ? 32 - __clz(spread) : 0; -#pragma unroll - for (unsigned b = lane; b < kMidBins / 4; b += 32) - reinterpret_cast(hist)[b] = make_uint4(0, 0, 0, 0); - __syncwarp(); -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - if (u * 32 >= mid) break; - if (mc[u] >= base && mc[u] <= top) - atomicAdd(hist + ((mc[u] - base) >> shift), 1u); - } - __syncwarp(); - unsigned above_bin; - const unsigned chosen = find_bin_v(hist, rank, above_bin); - rank -= above_bin; - base += chosen << shift; - top = min(base + (1u << shift) - 1, top); - __syncwarp(); - } - const unsigned th = base; - { - unsigned ng = 0, nt = 0, mk[R]; -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - ng += mc[u] > th; - nt += mc[u] == th; - mk[u] = 0; - if (u * 32 < mid && mc[u] >= th) mk[u] = idx_row[midbuf[lane + 32 * u]]; - } - unsigned wg, wt; - unsigned pg = above + warp_exclusive_scan(ng, wg), - pt = K - rank + warp_exclusive_scan(nt, wt); -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - if (u * 32 >= mid) break; - if (mc[u] > th) - stage[pg++] = uint16_t(mk[u]); - else if (mc[u] == th) { - if (pt < K) stage[pt] = uint16_t(mk[u]); - ++pt; - } - } - } - __syncwarp(); -#pragma unroll 4 - for (unsigned j = lane * 4; j < K; j += 128) { - const uint2 v = *reinterpret_cast(stage + j); - *reinterpret_cast(out + j) = make_int4( - int(v.x & 0xffffu), int(v.x >> 16), int(v.y & 0xffffu), int(v.y >> 16)); - } - return true; -} - -// > 16K keys: 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 spilled into the scan loop: 290K -> -// 860K cycles/block in a first try). -// 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__ bool select_staged_row( - unsigned* __restrict__ codes, unsigned count, - const uint16_t* __restrict__ values, const int32_t* __restrict__ packed, - int32_t* __restrict__ out, unsigned K, unsigned* hist) { - const unsigned lane = threadIdx.x & 31; - if (count <= K) return false; - const unsigned chunks = (count + 7) / 8; - constexpr unsigned C = 2; - const uint4* v4p = reinterpret_cast(values); - const uint4* p4p = reinterpret_cast(packed); -#pragma unroll 1 - for (unsigned c0i = lane; c0i < chunks; c0i += 32 * C) { - uint4 vv[C], pp0[C], pp1[C]; -#pragma unroll - for (unsigned u = 0; u < C; ++u) { - const unsigned 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 (unsigned u = 0; u < C; ++u) { - const unsigned c = c0i + 32 * u; - if (c < chunks) { - const unsigned vw[4] = {vv[u].x, vv[u].y, vv[u].z, vv[u].w}; - const unsigned 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}; - unsigned cw[8]; -#pragma unroll - for (unsigned k = 0; k < 8; ++k) - cw[k] = fused_candidates::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 (unsigned k = 0; k < 8; ++k) - if (8 * c + k < count) codes[8 * c + k] = cw[k]; - } - } - } - } - __syncwarp(); - unsigned h_hi, h_lo; - if (count <= 3584) { - unsigned smp[kSamplesPerLane]; -#pragma unroll - for (unsigned u = 0; u < kSamplesPerLane; ++u) - smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 10)] >> 8; - const unsigned expect = (K * 512 + count / 2) / count; - const unsigned r_hi = expect > 40 ? expect - 40 : 1, - r_lo = min(expect + 40, 512u); - radix_bracket(smp, r_hi, r_lo, hist, h_hi, h_lo); - } else { - // Long-context rows (4-6K candidates): +-40 of 512 samples gives a bracket - // of ~16% of the row, over kMid3Cap for 10-20% of the rows at 512K-1M. 1024 - // samples, +-48 (~3 sd): ~9%. - unsigned smp[2 * kSamplesPerLane]; -#pragma unroll - for (unsigned u = 0; u < 2 * kSamplesPerLane; ++u) - smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 11)] >> 8; - const unsigned expect = (K * 1024 + count / 2) / count; - const unsigned r_hi = expect > 48 ? expect - 48 : 1, - r_lo = min(expect + 48, 1024u); - radix_bracket<2 * kSamplesPerLane>(smp, r_hi, r_lo, hist, h_hi, h_lo); - } - const uint4* c4p = reinterpret_cast(codes); - const unsigned quads = (count + 3) / 4; - const unsigned lt = (1u << lane) - 1u; - uint16_t* midbuf = reinterpret_cast(hist); - unsigned* h512 = hist + kMid3Cap / 2; - // Collect with warp-ballot placement: passing lanes of one step write - // consecutive slots, so the global stores of the above elements are coalesced - // (per-lane runs were 4x slower: 32 lines/STG). Writes are bounded; the - // bracket check follows (out may be partially written on failure: the - // fallback select_row rewrites all K slots). - unsigned above = 0, mid = 0; - constexpr unsigned Q = 4; -#pragma unroll 1 - for (unsigned q0 = 0; q0 < quads; q0 += 32 * Q) { - uint4 pp[Q], cc[Q]; -#pragma unroll - for (unsigned u = 0; u < Q; ++u) { - const unsigned q = min(q0 + 32 * u + lane, quads - 1); - pp[u] = __ldcg(p4p + q); - cc[u] = c4p[q]; - } -#pragma unroll - for (unsigned u = 0; u < Q; ++u) { - const unsigned q = q0 + 32 * u + lane; - const unsigned w[4] = {cc[u].x, cc[u].y, cc[u].z, cc[u].w}; - const unsigned p[4] = {pp[u].x, pp[u].y, pp[u].z, pp[u].w}; -#pragma unroll - for (unsigned k = 0; k < 4; ++k) { - const unsigned 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 unsigned ma = __ballot_sync(Full, a), mm = __ballot_sync(Full, m); - const unsigned at = above + __popc(ma & lt), mt = mid + __popc(mm & lt); - if (a && at < K) - out[at] = fused_candidates::candidate_decode_index(int32_t(p[k])); - if (m && mt < kMid3Cap) midbuf[mt] = uint16_t(e); - above += __popc(ma); - mid += __popc(mm); - } - } - } - __syncwarp(); - if (above >= K || above + mid < K || mid > kMid3Cap) return false; - constexpr unsigned R = kMid3Cap / 32; - unsigned mc[R]; -#pragma unroll - for (unsigned u = 0; u < R; ++u) { - const unsigned i = lane + 32 * u; - mc[u] = 0; - if (i < mid) mc[u] = codes[midbuf[i]]; - } - unsigned rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; -#pragma unroll 1 - while (top > base) { - const unsigned width = top - base + 1, spread = (width - 1) / kMidBins; - const unsigned shift = spread ? 32 - __clz(spread) : 0; -#pragma unroll - for (unsigned b = lane; b < kMidBins / 4; b += 32) - reinterpret_cast(h512)[b] = make_uint4(0, 0, 0, 0); - __syncwarp(); -#pragma unroll - for (unsigned u = 0; u < R; ++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(); - unsigned above_bin; - const unsigned chosen = find_bin_v(h512, rank, above_bin); - rank -= above_bin; - base += chosen << shift; - top = min(base + (1u << shift) - 1, top); - __syncwarp(); - } - const unsigned th = base; - { - // Ballot placement (coalesced stores): greater codes fill [above, K - - // rank), ties [K - rank, K). - unsigned pg = above, pt = K - rank; -#pragma unroll - for (unsigned h = 0; h < 2; ++h) { - if (h * 16 * 32 >= mid) break; - int32_t pk[16]; -#pragma unroll - for (unsigned j = 0; j < 16; ++j) { - const unsigned u = h * 16 + j; - pk[j] = (lane + 32 * u < mid && mc[u] >= th) - ? __ldcg(packed + midbuf[lane + 32 * u]) - : 0; - } -#pragma unroll - for (unsigned j = 0; j < 16; ++j) { - const unsigned 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 unsigned mg = __ballot_sync(Full, g), - mt = __ballot_sync(Full, tq); - const unsigned ag = pg + __popc(mg & lt), at = pt + __popc(mt & lt); - if (g) out[ag] = fused_candidates::candidate_decode_index(pk[j]); - if (tq && at < K) - out[at] = fused_candidates::candidate_decode_index(pk[j]); - pg += __popc(mg); - pt += __popc(mt); - } - } - } - return true; -} -} // namespace row_select diff --git a/csrc/fused_mqa_topk/sample.cuh b/csrc/fused_mqa_topk/sample.cuh deleted file mode 100644 index 1575a9084e95..000000000000 --- a/csrc/fused_mqa_topk/sample.cuh +++ /dev/null @@ -1,25 +0,0 @@ -// Gate sample of the persistent kernel: each math warp histograms its row's -// sampled scores (fp16 bins >> 5, 2048 bins padded to 33 words per 32) and -// finish_sample_histogram picks the gate from them. -#pragma once -#include -#include -#include -#include -constexpr unsigned SampleQ = 8; -struct alignas(16) SampleShared { - unsigned hist[SampleQ][2112]; -}; -__device__ __forceinline__ void sample_sync() { - asm volatile("" ::: "memory"); - cutlass::arch::NamedBarrier(256, 2).sync(); - asm volatile("" ::: "memory"); -} -__device__ __forceinline__ unsigned sample_bin(float x) { - unsigned h = __half_as_ushort(__float2half_rd(x)); - return ((h & 0x8000u) ? ((~h) & 65535u) : (h ^ 0x8000u)) >> 5; -} -__device__ __forceinline__ float sample_floor(unsigned b) { - unsigned o = b << 5, h = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); - return __half2float(__ushort_as_half(h)); -} diff --git a/csrc/fused_mqa_topk/select.cuh b/csrc/fused_mqa_topk/select.cuh new file mode 100644 index 000000000000..43322d68dc23 --- /dev/null +++ b/csrc/fused_mqa_topk/select.cuh @@ -0,0 +1,438 @@ +// SPDX-License-Identifier: MIT +// Per-row top-K selection of the persistent kernel's candidates, one warp per +// row (K, the row's top-k: 512, 1024 or 2048, is 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 Full = 0xffffffffu, 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(Full, 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(Full, acc, winner); + return __shfl_sync(Full, chosen, winner); +} + +// Requires count >= K (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 K, 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(Full, lo); + hi = __reduce_max_sync(Full, hi); + __syncwarp(); + + // Current bucket: codes in [lo, lo + width). rank = elements still needed + // from it. + uint32_t width = hi - lo + 1, rank = K; +#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 < K; 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(Full, pass[u]), + at = pos + __popc(mask & lt); + if (pass[u] && at < K) 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 512 register samples 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 +} + +// 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(Full, 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(Full, acc_at, winner); + return __shfl_sync(Full, 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 spilled into the scan loop: 290K -> +// 860K cycles/block in a first try). +// 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 K, uint32_t* hist) { + const uint32_t lane = threadIdx.x & 31; + if (count <= K) 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) { + uint32_t smp[kSamplesPerLane]; +#pragma unroll + for (uint32_t u = 0; u < kSamplesPerLane; ++u) + smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 10)] >> 8; + const uint32_t expect = (K * 512 + count / 2) / count; + const uint32_t r_hi = expect > kMargin ? expect - kMargin : 1, + r_lo = min(expect + kMargin, 512u); + radix_bracket(smp, r_hi, r_lo, hist, h_hi, h_lo); + } else { + uint32_t smp[2 * kSamplesPerLane]; +#pragma unroll + for (uint32_t u = 0; u < 2 * kSamplesPerLane; ++u) + smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 11)] >> 8; + const uint32_t expect = (K * 1024 + count / 2) / count; + const uint32_t r_hi = expect > kWideMargin ? expect - kWideMargin : 1, + r_lo = min(expect + kWideMargin, 1024u); + radix_bracket<2 * kSamplesPerLane>(smp, r_hi, r_lo, 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 + // (per-lane runs were 4x slower: 32 lines/STG). 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(Full, a), mm = __ballot_sync(Full, m); + const uint32_t at = above + __popc(ma & lt), mt = mid + __popc(mm & lt); + if (a && at < K) 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 >= K || above + mid < K || 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 = K - 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 = K - 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(Full, g), + mt = __ballot_sync(Full, 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 < K) 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 73f7133cec36..e4e99cee6ac6 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -878,4 +878,4 @@ void fused_mqa_topk_prefill( 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, bool relative); + torch::stable::Tensor counts, torch::stable::Tensor defer); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 0ee8f56130e7..2ebfd3556220 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -16,8 +16,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "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, " - "bool relative) -> ()"); + "Tensor! retained_indices, Tensor! counts, Tensor! defer) -> ()"); #endif // Compute per-token-group FP8 quantized tensor and scaling factor. diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index aa021ce7289f..f6a3d404e9a3 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -2,16 +2,32 @@ # 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 +import vllm.models.glm5next.nvidia.sparse_indexer as kpool_indexer from vllm import _custom_ops as ops from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( - FUSED_MQA_TOPK_MAX_ROWS, fused_mqa_topk_available, + fused_mqa_topk_min_keys, fused_mqa_topk_prefill, + get_fused_mqa_topk_workspace, + use_fused_mqa_topk, +) +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, ) -from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder requires_fused = pytest.mark.skipif( not torch.cuda.is_available() or not fused_mqa_topk_available(), @@ -25,13 +41,6 @@ def _dense(q, k, scales, weights, starts, ends, topk=TOPK): """Upstream dense path: budgeted DeepGEMM logits + top_k_per_row_prefill.""" - 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, - ) - k_scale = scales if mqa_logits_use_mx_scales(): # Power-of-two scales as four equal UE8M0 bytes, like the K cache. @@ -168,7 +177,7 @@ def test_single_request_matches_dense(keys, rows, case, topk): weights = weights.abs() scales[:12288] = 4 out = torch.empty(rows, topk, device="cuda", dtype=torch.int32) - assert fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + 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). @@ -195,7 +204,7 @@ def test_multi_request_chunk_matches_dense(lens, topk): 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) - assert fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + 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) @@ -215,33 +224,55 @@ def test_multi_request_chunk_matches_dense(lens, topk): def test_large_chunk_single_call(rows): """Prefill chunks above 16384 rows (e.g. --max-num-batched-tokens 32768) run in one launch; the per-row workspace grows to fit.""" - from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( - get_fused_mqa_topk_workspace, - ) - 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) - assert fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) assert get_fused_mqa_topk_workspace(q.device, rows).counts.numel() >= rows _check(out, q, k, scales, weights, starts, ends) @requires_fused -def test_rows_start_at_zero_skips_conversion(): - rows, keys = 32, 8192 - q, k, scales, weights = _inputs(rows, keys) +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, relative=True) 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) - ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32) - a = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32) - b = torch.empty_like(a) - assert fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, a) - assert fused_mqa_topk_prefill( - q, k, scales, weights, starts, ends, b, rows_start_at_zero=True + 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], ) - torch.testing.assert_close(a.sort(1).values, b.sort(1).values) @requires_fused @@ -252,7 +283,7 @@ def test_cuda_graph_replay(): 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, rows_start_at_zero=True + q, k, scales, weights, starts, ends, out ) fn() torch.accelerator.synchronize() @@ -267,97 +298,33 @@ def test_cuda_graph_replay(): @requires_fused -def test_unsupported_inputs_fall_back(): +def test_dispatch_predicate(): + """use_fused_mqa_topk: FP8 32x128 queries, top-k 512/1024/2048, full + context, no FP4 cache / FP4 queries / candidate selection.""" + q = _inputs(16, 4096)[0] + ok = dict(fp4_cache=False, full_context=True) + assert use_fused_mqa_topk(q, None, 2048, **ok) + assert not use_fused_mqa_topk(q, None, 768, **ok) + assert not use_fused_mqa_topk(q[:, :16], None, 2048, **ok) + assert not use_fused_mqa_topk(q.float(), None, 2048, **ok) + assert not use_fused_mqa_topk(q, torch.ones(16, device="cuda"), 2048, **ok) + assert not use_fused_mqa_topk(q, None, 2048, fp4_cache=True, full_context=True) + assert not use_fused_mqa_topk(q, None, 2048, fp4_cache=False, full_context=False) + assert not use_fused_mqa_topk(q, None, 2048, candidate_selection=True, **ok) + + +@requires_fused +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) - out = torch.empty(16, TOPK, device="cuda", dtype=torch.int32) - assert not fused_mqa_topk_prefill( - q[:, :16], k, scales, weights[:, :16], starts, ends, out - ) - assert not fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out[:, :768]) - q17 = torch.zeros(16385, 32, 128, device="cuda").to(torch.float8_e4m3fn) - assert not fused_mqa_topk_prefill(q17, k, scales, weights, starts, ends, out) - - -def _split(seq_lens, q_lens, fused, offset=0): - return DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( - torch.tensor(seq_lens), - torch.tensor(q_lens), - 1048576, - LOGITS_BUDGET, - request_offset=offset, - use_fused_mqa_topk=fused, - ) - - -@pytest.mark.parametrize("q_rows", [16384, 32768]) -@pytest.mark.parametrize("fused", [False, True]) -def test_planner_one_fused_call_per_chunk(fused, q_rows): - # One fused call per FUSED_MQA_TOPK_MAX_ROWS rows (None: per chunk); dense - # splits by the logits budget. - chunks = _split([524288], [q_rows], fused) - cap = FUSED_MQA_TOPK_MAX_ROWS or q_rows - assert len(chunks) == (-(-q_rows // cap) if fused else q_rows // 256) - assert sum(q.stop - q.start for _, q in chunks) == q_rows - - -def test_planner_fused_any_length_and_short_queries(): - # No length threshold and no minimum query count on the fused path. - assert _split([524288], [512], True) == [(slice(0, 1), slice(0, 512))] - assert _split([4096], [4096], True) == [(slice(0, 1), slice(0, 4096))] - # A 40000-row step: split at FUSED_MQA_TOPK_MAX_ROWS (None: one call). - cap = FUSED_MQA_TOPK_MAX_ROWS or 40000 - assert _split([65536], [40000], True) == [ - (slice(0, 1), slice(lo, min(lo + cap, 40000))) for lo in range(0, 40000, cap) - ] - - -def test_planner_packs_requests_in_one_fused_call(): - # Requests share one fused call up to the key limit; dense needs a split. - assert _split([4096] * 4, [4096] * 4, True) == [(slice(0, 4), slice(0, 16384))] - assert len(_split([4096] * 4, [4096] * 4, False)) > 1 - chunks = _split([524288, 262144], [16384, 9], True, offset=2) - if FUSED_MQA_TOPK_MAX_ROWS is None or FUSED_MQA_TOPK_MAX_ROWS >= 16393: - assert chunks == [(slice(2, 4), slice(0, 16393))] - else: # the second request would push the call past the row cap - assert chunks == [(slice(2, 3), slice(0, 16384)), (slice(3, 4), slice(0, 9))] - - -def test_planner_fused_never_packs_past_key_limit(): - # Two 600K-key continuation steps of 50 rows: dense packs them into one - # logits-budget chunk; the fused planner keeps one fused chunk per request so - # no chunk reaches the dense path (whose logits buffer is not reserved). - args = ( - torch.tensor([600_000, 600_000]), - torch.tensor([50, 50]), - 4_000_000, - LOGITS_BUDGET, - ) - dense = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks(*args) - fused = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( - *args, use_fused_mqa_topk=True - ) - assert dense == [(slice(0, 2), slice(0, 100))] - assert fused == [(slice(0, 1), slice(0, 50)), (slice(1, 2), slice(0, 50))] - - -def test_planner_keeps_budget_beyond_fused_key_limit(): - seq_lens, q_lens = [2_000_000], [256] - fused = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( - torch.tensor(seq_lens), - torch.tensor(q_lens), - 4_000_000, - LOGITS_BUDGET, - use_fused_mqa_topk=True, - ) - dense = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( - torch.tensor(seq_lens), - torch.tensor(q_lens), - 4_000_000, - LOGITS_BUDGET, - ) - assert fused == dense and len(dense) == 4 + 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 @@ -389,18 +356,8 @@ def _kpool_inputs(pool_lens=(20000, 70000), q_rows=96, tail=2, seed=919): ) -def _run_kpool_indexer(monkeypatch, inputs, fused, fused_fn=None, budget_mb=512): +def _run_kpool_indexer(monkeypatch, inputs, fused, budget_mb=512): """Run the GLM-5.3-Flash kpool indexer prefill on one packed chunk.""" - from types import SimpleNamespace - - import vllm.envs as envs - import vllm.models.glm5next.nvidia.sparse_indexer as kpool_indexer - from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV32IndexerMetadata, - DeepseekV32IndexerPrefillChunkMetadata, - DeepseekV32IndexerPrefillMetadata, - ) - q, k, scales, weights, starts, ends, positions = inputs rows, keys = q.shape[0], k.shape[0] chunk = DeepseekV32IndexerPrefillChunkMetadata( @@ -441,11 +398,11 @@ def _run_kpool_indexer(monkeypatch, inputs, fused, fused_fn=None, budget_mb=512) monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", fused) monkeypatch.setattr(envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", budget_mb) calls: list[torch.Size] = [] - fused_fn = fused_fn or kpool_indexer.fused_mqa_topk_prefill + fused_fn = kpool_indexer.fused_mqa_topk_prefill def spy(*a, **kw): calls.append(a[-1].shape) # the output buffer: [rows, top-k pools] - return fused_fn(*a, **kw) + 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) @@ -495,10 +452,6 @@ def test_kpool_indexer_fused_matches_dense(monkeypatch): 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 kernel is not faster there).""" - from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( - fused_mqa_topk_min_keys, - ) - 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=True) @@ -507,20 +460,15 @@ def test_kpool_indexer_short_chunk_stays_dense(monkeypatch): @requires_fused def test_kpool_indexer_dense_fallback_splits_rows(monkeypatch): - """Fused-planned chunks skip the logits budget; if the kernel declines, - the dense fallback splits rows to stay within it and selects the same - pools as one unbudgeted dense pass.""" + """Fused-planned chunks 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=False) - split, calls = _run_kpool_indexer( - monkeypatch, inputs, fused=True, fused_fn=lambda *a, **kw: False, budget_mb=1 - ) - assert len(calls) == 1 + split, calls = _run_kpool_indexer(monkeypatch, inputs, fused=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. - 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 - q, k, scales, weights, starts, ends, _ = inputs lens = ends - starts q_scale, w = mqa_logits_mx_inputs(q, None, weights) diff --git a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py new file mode 100644 index 000000000000..19a4a3dbfdef --- /dev/null +++ b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prefill chunk planning and dispatch predicates of the fused MQA+top-k +indexer path (CPU only; the kernel is tested in +tests/kernels/attention/test_fused_mqa_topk.py).""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.envs as envs +from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( + FUSED_MQA_TOPK_MAX_KEYS, + FUSED_MQA_TOPK_MAX_ROWS, + fused_mqa_topk_chunk, + fused_mqa_topk_min_keys, + fused_mqa_topk_model_supported, +) +from vllm.model_executor.layers.sparse_attn_indexer import dense_logits_rows +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder + +LOGITS_BUDGET = 512 * 1024**2 + + +def _split(seq_lens, q_lens, fused, offset=0, workspace=1048576): + return DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( + torch.tensor(seq_lens), + torch.tensor(q_lens), + workspace, + LOGITS_BUDGET, + request_offset=offset, + plan_fused=fused, + ) + + +@pytest.mark.parametrize("q_rows", [16384, 32768]) +@pytest.mark.parametrize("fused", [False, True]) +def test_planner_one_fused_call_per_chunk(fused, q_rows): + # One fused call per FUSED_MQA_TOPK_MAX_ROWS rows; dense splits by the + # logits budget. + chunks = _split([524288], [q_rows], fused) + expected = -(-q_rows // FUSED_MQA_TOPK_MAX_ROWS) if fused else q_rows // 256 + assert len(chunks) == expected + assert sum(q.stop - q.start for _, q in chunks) == q_rows + + +def test_planner_fused_any_length_and_short_queries(): + # No length threshold and no minimum query count on the fused path. + assert _split([524288], [512], True) == [(slice(0, 1), slice(0, 512))] + assert _split([4096], [4096], True) == [(slice(0, 1), slice(0, 4096))] + # A 40000-row step is split at FUSED_MQA_TOPK_MAX_ROWS. + cap = FUSED_MQA_TOPK_MAX_ROWS + assert _split([65536], [40000], True) == [ + (slice(0, 1), slice(lo, min(lo + cap, 40000))) for lo in range(0, 40000, cap) + ] + + +def test_planner_packs_requests_in_one_fused_call(): + # Requests share one fused call up to the row cap; dense needs a split. + assert _split([4096] * 4, [4096] * 4, True) == [(slice(0, 4), slice(0, 16384))] + assert len(_split([4096] * 4, [4096] * 4, False)) > 1 + # The second request would push the call past the row cap. + assert _split([524288, 262144], [16384, 9], True, offset=2) == [ + (slice(2, 3), slice(0, 16384)), + (slice(3, 4), slice(0, 9)), + ] + + +def test_planner_fused_never_packs_past_key_limit(): + # Two 600K-key continuation steps of 50 rows: dense packs them into one + # logits-budget chunk; the fused planner keeps one fused chunk per request so + # no chunk reaches the dense path (whose logits buffer is not reserved). + args = ([600_000, 600_000], [50, 50], False, 0, 4_000_000) + assert _split(*args) == [(slice(0, 2), slice(0, 100))] + fused = (*args[:2], True, *args[3:]) + assert _split(*fused) == [(slice(0, 1), slice(0, 50)), (slice(1, 2), slice(0, 50))] + + +def test_planner_keeps_budget_beyond_fused_key_limit(): + # A single request longer than the fused key limit is budget-planned. + fused = _split([2_000_000], [256], True, workspace=4_000_000) + dense = _split([2_000_000], [256], False, workspace=4_000_000) + assert fused == dense and len(dense) == 4 + + +@pytest.mark.parametrize("topk", [512, 1024, 2048]) +def test_fused_chunk_thresholds(topk): + min_keys = fused_mqa_topk_min_keys(topk) + assert not fused_mqa_topk_chunk(min_keys - 1, min_keys, topk) + assert fused_mqa_topk_chunk(min_keys, min_keys, topk) + assert fused_mqa_topk_chunk(min_keys, FUSED_MQA_TOPK_MAX_KEYS, topk) + # Packed gathered keys past the kernel's limit stay dense. + assert not fused_mqa_topk_chunk(min_keys, FUSED_MQA_TOPK_MAX_KEYS + 1, topk) + + +@pytest.mark.parametrize( + "heads, head_dim, topk, kpool, supported", + [ + (32, 128, 2048, None, True), # GLM-5.2 + (32, 128, 512, None, True), # DeepSeek-V4.1-Flash + (32, 128, 2048, 4, True), # GLM-5.3-Flash: top-512 pools + (64, 128, 2048, None, False), # DeepSeek-V3.2: 64 index heads + (32, 64, 2048, None, False), + (32, 128, 256, None, False), + (32, 128, 2048, 3, False), # 682 pools + (32, 128, None, None, False), # no indexer + ], +) +def test_model_supported(heads, head_dim, topk, kpool, supported): + """The chunk planner only plans fused chunks for indexer shapes the kernel + serves; other DSA models keep the logits-budget planning.""" + config = SimpleNamespace( + index_n_heads=heads, index_head_dim=head_dim, index_topk=topk + ) + if kpool is not None: + config.index_kpool = kpool + assert fused_mqa_topk_model_supported(config) == supported + + +def test_dense_logits_rows_fits_budget(monkeypatch): + """Dense row splitting keeps each call's fp32 [rows, max_seqlen_k] logits + within the budget, and never splits a chunk the budget planner made.""" + monkeypatch.setattr(envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", 512) + for width in (1, 4096, 65536, 1 << 20): + rows = dense_logits_rows(width) + assert rows * width * 4 <= LOGITS_BUDGET + assert (rows + 1) * width * 4 > LOGITS_BUDGET + # Budget-planned chunks have rows * total_keys <= budget / 4 and every + # row's max_seqlen_k <= total_keys, so they fit in one dense call. + seq_lens = [131072, 65536, 4096] + for req_slice, q_slice in _split(seq_lens, [4096, 4096, 4096], False): + total_keys = sum(seq_lens[req_slice]) + assert q_slice.stop - q_slice.start <= dense_logits_rows(total_keys) + assert dense_logits_rows(0) == LOGITS_BUDGET // 4 diff --git a/vllm/envs.py b/vllm/envs.py index bd57ebe8cd98..4e141a953dc9 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1094,10 +1094,10 @@ 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 eligible - # sparse MLA indexer prefill chunks on SM100-family GPUs (FP8 indexer cache, - # 32 index heads, top-k 512/1024/2048, no DCP/PCP). Falls back to the dense - # path otherwise. Default: on (set to 0 for the dense path). + # 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")) ), diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 23c72bdd003d..5f702528c55e 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -8,16 +8,20 @@ ``[rows, keys]`` logits matrix that the dense path (``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``) needs. -Above 4K keys the kernel is persistent (one CTA per SM); each CTA owns two -8-row candidate slots. The scratch (~296 MiB on a 148-SM GPU) is allocated once -per device and reused by every call, independently of the number of rows or -keys. +The kernel is persistent (one CTA per SM); each CTA owns two 8-row candidate +slots. The scratch (~296 MiB on a 148-SM GPU) is allocated once per device and +reused by every call, independently of the number of rows or keys. It is tuned +for the chunks vLLM sends it (rows that see >= fused_mqa_topk_min_keys keys); +smaller key counts are supported but not optimized. """ +import functools from dataclasses import dataclass +from typing import Any import torch +import vllm.envs as envs from vllm.logger import init_logger from vllm.platforms import current_platform @@ -31,10 +35,8 @@ FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 # Rows per fused call. The kernel has no row limit, but prefill chunks above # 16384 rows (--max-num-batched-tokens > 16384) are not validated end to end, -# so the chunk planner splits fused chunks at this many rows. None: one fused -# call per chunk, any size. -# TODO: check for other prefill chunk sizes. -FUSED_MQA_TOPK_MAX_ROWS: int | None = 16384 +# so the chunk planner splits fused chunks at this many rows. +FUSED_MQA_TOPK_MAX_ROWS = 16384 def fused_mqa_topk_min_keys(topk: int) -> int: @@ -46,9 +48,10 @@ def fused_mqa_topk_min_keys(topk: int) -> int: return 1 << 16 if topk > 1024 else 1 << 15 -_CANDIDATES_PER_ROW = 16384 -_RETAINED_PER_ROW = 4096 -_SLOTS_PER_SM = 2 +# 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 @dataclass @@ -56,7 +59,7 @@ class FusedMQATopKWorkspace: """Caller-owned scratch of the fused kernel (see fused_mqa_topk.cu).""" counter: torch.Tensor # [2] int32, zero between calls (the kernel restores it) - values: torch.Tensor # [8S, 16384] fp16, S = 2 slots per SM + values: torch.Tensor # [8S, 16384] uint16 (low 16 bits of fp24 codes), S = 2/SM indices: torch.Tensor # [8S, 16384] int32 retained: torch.Tensor # [8S, 4096] fp32 retained_indices: torch.Tensor # [8S, 4096] int32 @@ -78,6 +81,7 @@ def as_args(self) -> tuple[torch.Tensor, ...]: _WORKSPACES: dict[int, FusedMQATopKWorkspace] = {} +@functools.cache def fused_mqa_topk_available() -> bool: """True when the fused op is compiled in and the GPU is SM100-family.""" return ( @@ -87,6 +91,65 @@ def fused_mqa_topk_available() -> bool: ) +def fused_mqa_topk_model_supported(hf_text_config: Any) -> bool: + """Whether a sparse-attention model's indexer shape fits the fused kernel: + 32 index heads of dimension 128 and a per-row top-k (per pool for kpool + indexers) of 512, 1024 or 2048. Lets the prefill chunk planner skip the + fused planning for models the kernel cannot serve (e.g. DeepSeek-V3.2, 64 + index heads).""" + topk = getattr(hf_text_config, "index_topk", None) + kpool = getattr(hf_text_config, "index_kpool", None) or 1 + return ( + getattr(hf_text_config, "index_n_heads", None) == FUSED_MQA_TOPK_HEADS + and getattr(hf_text_config, "index_head_dim", None) == FUSED_MQA_TOPK_HEAD_DIM + and topk is not None + and topk // kpool in FUSED_MQA_TOPK_TOPKS + ) + + +def use_fused_mqa_topk( + q: torch.Tensor, + q_scale: torch.Tensor | None, + topk_tokens: int, + *, + fp4_cache: bool, + full_context: bool, + candidate_selection: bool = False, +) -> bool: + """Whether an indexer layer call may use the fused kernel for its prefill + chunks (each chunk then also needs fused_mqa_topk_chunk). + + q / q_scale: the layer's quantized indexer queries; topk_tokens: the top-k + the kernel selects (pools for kpool indexers). 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. The FP4 cache, FP4 + queries and v4.1 candidate selection stay on the dense path. + """ + return ( + envs.VLLM_USE_FUSED_MQA_TOPK + and full_context + and not fp4_cache + and q_scale is None + and not candidate_selection + and topk_tokens in FUSED_MQA_TOPK_TOPKS + and q.dtype == torch.float8_e4m3fn + and q.dim() == 3 + and tuple(q.shape[1:]) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM) + and q.is_cuda + and fused_mqa_topk_available() + ) + + +def fused_mqa_topk_chunk(max_seqlen_k: int, keys: int, topk_tokens: int) -> bool: + """Whether a prefill chunk of a fused-eligible layer takes the fused + kernel: its longest row sees at least fused_mqa_topk_min_keys keys (the + dense path is as fast below) and its gathered keys fit the kernel.""" + return ( + max_seqlen_k >= fused_mqa_topk_min_keys(topk_tokens) + and keys <= FUSED_MQA_TOPK_MAX_KEYS + ) + + def get_fused_mqa_topk_workspace( device: torch.device, rows: int ) -> FusedMQATopKWorkspace: @@ -95,7 +158,9 @@ def get_fused_mqa_topk_workspace( The per-row part (counts, defer) holds at least ``rows`` query rows and grows when a larger call arrives; the profiling run asks for the largest prefill chunk (max_num_batched_tokens rows) so serving never grows it. The block - counter is zeroed once; every launch leaves it zeroed again. + counter is zeroed once; every launch leaves it zeroed again. The workspace is + shared by every call on the device, so calls must not run concurrently (they + are issued on one stream). """ rows = (rows + 7) // 8 * 8 index = ( @@ -111,7 +176,7 @@ def get_fused_mqa_topk_workspace( ws = FusedMQATopKWorkspace( counter=torch.zeros(2, dtype=torch.int32, device=dev), values=torch.empty( - (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.float16, device=dev + (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.uint16, device=dev ), indices=torch.empty( (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.int32, device=dev @@ -139,36 +204,9 @@ def get_fused_mqa_topk_workspace( logger.info( "Grew the fused MQA+top-k row workspace on cuda:%d to %d rows.", index, rows ) - _WORKSPACES[index] = ws return ws -def fused_mqa_topk_config_supported(q: torch.Tensor, topk_tokens: int) -> bool: - """Whether this indexer's query layout and top-k fit the fused kernel.""" - return ( - topk_tokens in FUSED_MQA_TOPK_TOPKS - and q.dtype == torch.float8_e4m3fn - and q.dim() == 3 - and tuple(q.shape[1:]) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM) - and q.is_cuda - and fused_mqa_topk_available() - ) - - -def fused_mqa_topk_supported( - q: torch.Tensor, k: torch.Tensor, topk_tokens: int -) -> bool: - """Static eligibility of an indexer prefill chunk for the fused kernel.""" - return ( - k.dtype == torch.float8_e4m3fn - and k.dim() == 2 - and k.shape[1] == FUSED_MQA_TOPK_HEAD_DIM - and q.shape[0] > 0 - and 0 < k.shape[0] <= FUSED_MQA_TOPK_MAX_KEYS - and fused_mqa_topk_config_supported(q, topk_tokens) - ) - - def fused_mqa_topk_key_scales(scales: torch.Tensor) -> torch.Tensor: """Per-key fp32 scales from the indexer K cache. @@ -188,76 +226,26 @@ def fused_mqa_topk_prefill( starts: torch.Tensor, ends: torch.Tensor, output: torch.Tensor, - *, - rows_start_at_zero: bool = False, -) -> bool: +) -> None: """Fused replacement of ``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``. - q: [rows, 32, 128] fp8 e4m3; k: [keys, 128] fp8 e4m3; 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. + 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. ``rows_start_at_zero=True`` - promises every ``start`` is 0 (single-request chunk) and skips the - relative-index conversion. - - Returns False (and leaves ``output`` untouched) for unsupported inputs so - the caller can fall back to the dense path. + padded when fewer than k keys are visible. """ - rows, keys = q.shape[0], k.shape[0] - if ( - not fused_mqa_topk_supported(q, k, output.shape[1] if output.dim() == 2 else 0) - or scales.dtype != torch.float32 - or scales.numel() != keys - or weights.dtype != torch.float32 - or tuple(weights.shape) != (rows, FUSED_MQA_TOPK_HEADS) - or starts.dtype != torch.int32 - or ends.dtype != torch.int32 - or starts.numel() != rows - or ends.numel() != rows - or output.dtype != torch.int32 - or output.shape[0] != rows - ): - return False - logger.info_once( "Using the single-launch fused MQA+top-k prefill kernel " "(FP8, 32 heads, top-k %d).", output.shape[1], ) - q = q.contiguous() - k = k.contiguous() - scales = scales.contiguous() - weights = weights.contiguous() - starts = starts.contiguous() - ends = ends.contiguous() - - # The kernel writes the output with 16-byte stores; any other layout goes - # through a temporary (never the case for the indexer's own buffers). - direct = output.is_contiguous() and output.data_ptr() % 16 == 0 - out = ( - output - if direct - else torch.empty_like(output, memory_format=torch.contiguous_format) - ) - ws = get_fused_mqa_topk_workspace(q.device, rows) - # The kernel emits absolute key indices; the dense top-k (and the sparse - # attention consumer) use indices relative to the row start. relative=True - # converts in place (rows with start 0 are skipped on the GPU). + ws = get_fused_mqa_topk_workspace(q.device, q.shape[0]) torch.ops._C.fused_mqa_topk_prefill( - q, - k, - scales, - weights, - starts, - ends, - out, - *ws.as_args(), - not rows_start_at_zero, + q, k, scales, weights, starts, ends, output, *ws.as_args() ) - if not direct: - output.copy_(out) - return True diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index e1646f4a7ccb..c13a51a13f5f 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -30,11 +30,11 @@ select_candidate_blocks as _select_candidate_blocks, ) from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( - fused_mqa_topk_config_supported, + fused_mqa_topk_chunk, fused_mqa_topk_key_scales, - fused_mqa_topk_min_keys, fused_mqa_topk_prefill, get_fused_mqa_topk_workspace, + use_fused_mqa_topk, ) from vllm.model_executor.layers.indexer_topk import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -308,6 +308,18 @@ def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.cons tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask) +def dense_logits_rows(max_seqlen_k: int) -> int: + """Query rows per dense MQA-logits call that keep its fp32 [rows, + max_seqlen_k] logits within VLLM_SPARSE_INDEXER_MAX_LOGITS_MB. + + Chunks the planner bounded by the logits budget already fit in one call; + chunks planned for the fused kernel (not bounded by it) are split when they + take the dense path. + """ + budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + return max(1, budget // (4 * max(1, max_seqlen_k))) + + def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor: """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels. @@ -331,6 +343,87 @@ def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor return out +def dense_prefill_topk( + q: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + max_seqlen_k: int, + topk_indices: torch.Tensor, + topk_tokens: int, + candidates: torch.Tensor | None, + candidate_write: bool, + candidate_block_size: int, + dcp_merge: tuple[int, int, int] | None, +) -> None: + """Dense MQA logits + per-row top-k of one row slice of a prefill chunk. + + candidates: the slice's v4.1 candidate blocks (None: no two-level + selection). dcp_merge: (dcp_rank, dcp_world_size, cp_interleave) to merge + the per-shard top-k into the global one (None under PCP, where the top-k + already ran over the whole context). + """ + # 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 + if current_platform.is_xpu(): + if q_scale is not None: + raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") + logits = torch.ops.vllm.xpu_fp8_mqa_logits( + q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke + ) + row_starts = cu_seqlen_ks + row_ends = cu_seqlen_ke + else: + q_scale, weights = mqa_logits_mx_inputs(q, q_scale, weights) + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q, q_scale), + (k, k_scale), + weights, + cu_seqlen_ks, + cu_seqlen_ke, + max_seqlen_k, + ), + row_ends, + ) + if candidates 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. + if candidate_write: + _select_candidate_blocks( + logits, + row_starts, + row_ends, + candidates.shape[1], + candidate_block_size, + candidates, + ) + else: + _apply_candidate_mask( + logits, row_starts, row_ends, candidates, candidate_block_size + ) + ops.top_k_per_row_prefill( + logits, + torch.zeros_like(row_ends) if row_starts is None else row_starts, + row_ends, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + topk_tokens, + ) + if dcp_merge is not None: + _merge_dcp_topk_global( + logits, topk_indices, topk_tokens, *dcp_merge, row_starts=row_starts + ) + + def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, @@ -485,6 +578,17 @@ def sparse_attn_indexer( ) assert candidate_block_size > 0 + # Fused MQA+top-k prefill for this layer (per chunk: fused_mqa_topk_chunk). + # Under PCP every rank ranks its rows over the whole (gathered) context. + fused = use_fused_mqa_topk( + q_quant, + q_scale, + topk_tokens, + fp4_cache=use_fp4_cache, + full_context=use_pcp or dcp_world_size == 1, + candidate_selection=candidate_blocks is not None, + ) + # assert isinstance(attn_metadata, dict) if not isinstance(attn_metadata, dict): # Reserve workspace for indexer during profiling run @@ -504,16 +608,7 @@ def sparse_attn_indexer( ) profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) - # Under PCP every rank ranks its rows over the whole (gathered) - # context; pure DCP needs the per-shard top-k merge, so it stays dense. - uses_fused = ( - envs.VLLM_USE_FUSED_MQA_TOPK - and (use_pcp or dcp_world_size == 1) - and not use_fp4_cache - and candidate_blocks is None - and fused_mqa_topk_config_supported(q_quant, topk_tokens) - ) - if uses_fused: + if fused: # The fused prefill kernel's fixed-size workspace is persistent; # allocate it during profiling so it is accounted for. get_fused_mqa_topk_workspace(q_quant.device, q_quant.shape[0]) @@ -709,145 +804,70 @@ def sparse_attn_indexer( q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] - # 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 + # Under PCP the top-k runs over the whole gathered context: no merge. + dcp_merge = ( + (dcp_rank, dcp_world_size, cp_kv_cache_interleave_size) + if deinterleave_idx is None + else None + ) 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) - else: - # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 - # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. - if use_fp4_cache: - q_slice_cast = q_slice.view(torch.int8) - k_quant_cast = k_quant.view(torch.int8) - k_scale_cast = k_scale.view(torch.int32).squeeze(-1) - else: - q_slice_cast = q_slice - k_quant_cast = k_quant - k_scale_cast = k_scale.view( - torch.int32 if mqa_logits_use_mx_scales() else torch.float32 - ).squeeze(-1) - # Single-launch fused MQA-logits + top-k (SM100, FP8, 32 heads, - # top-k 512/1024/2048): no logits buffer, for chunks whose rows - # see at least fused_mqa_topk_min_keys keys (dense is as fast below). - # It needs every row's whole context: no DCP, or PCP + DCP after - # the de-interleaving gather. Pure DCP (per-shard top-k + merge), - # the FP4 cache and v4.1 candidate selection stay dense. - use_fused = ( - envs.VLLM_USE_FUSED_MQA_TOPK - and (dcp_world_size == 1 or deinterleave_idx is not None) - and not current_platform.is_xpu() - ) - if ( - use_fused - and chunk.max_seqlen_k >= fused_mqa_topk_min_keys(topk_tokens) - and q_scale_slice is None - and candidate_blocks is None - and fused_mqa_topk_prefill( - q_slice_cast, - k_quant_cast, - fused_mqa_topk_key_scales(k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, + if dcp_merge is not None: + # Every DCP rank joins the merge, also with no local keys. + _merge_dcp_topk_global( + q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32), topk_indices, - # A single-request chunk has every row start at 0. - rows_start_at_zero=chunk.block_table.shape[0] == 1, - ) - ): - continue - # Fused-planned chunks are not bounded by the logits budget, so - # the dense fallback splits their rows to stay within it. - dense_rows = q_slice_cast.shape[0] - if use_fused: - budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 - dense_rows = max(1, budget // (4 * max(1, chunk.max_seqlen_k))) - for sub in range(0, q_slice_cast.shape[0], dense_rows): - rows = slice(sub, min(sub + dense_rows, q_slice_cast.shape[0])) - tokens = slice(row_start + rows.start, row_start + rows.stop) - sub_ks, sub_ke = cu_seqlen_ks[rows], cu_seqlen_ke[rows] - sub_q_scale = ( - q_scale_slice[rows] if q_scale_slice is not None else None - ) - if current_platform.is_xpu(): - if sub_q_scale 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], - sub_ks, - sub_ke, - ) - row_starts = sub_ks - row_ends = sub_ke - else: - sub_q_scale, chunk_weights = mqa_logits_mx_inputs( - q_slice_cast[rows], sub_q_scale, weights[tokens] - ) - row_ends = sub_ke - sub_ks - logits = as_fp32_logits( - fp8_fp4_mqa_logits( - (q_slice_cast[rows], sub_q_scale), - (k_quant_cast, k_scale_cast), - chunk_weights, - sub_ks, - sub_ke, - chunk.max_seqlen_k, - ), - 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, - ) - ops.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends) - if row_starts is None - else row_starts, - row_ends, - topk_indices[rows], - num_rows, - logits.stride(0), - logits.stride(1), topk_tokens, + *dcp_merge, ) - - if deinterleave_idx is None: - # Under the PCP path the top-k already ran over the whole - # context. - _merge_dcp_topk_global( - logits, + continue + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. + if use_fp4_cache: + q_slice_cast = q_slice.view(torch.int8) + k_quant_cast = k_quant.view(torch.int8) + k_scale_cast = k_scale.view(torch.int32).squeeze(-1) + else: + q_slice_cast = q_slice + k_quant_cast = k_quant + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) + # Single-launch fused MQA-logits + top-k: no logits buffer. + if fused and fused_mqa_topk_chunk( + chunk.max_seqlen_k, k_quant_cast.shape[0], topk_tokens + ): + fused_mqa_topk_prefill( + q_slice_cast, + k_quant_cast, + fused_mqa_topk_key_scales(k_scale_cast), + weights[row_start:row_end], + cu_seqlen_ks, + cu_seqlen_ke, topk_indices, + ) + continue + # Fused-planned chunks are not bounded by the logits budget: split + # their rows (budget-planned chunks fit in one call). + dense_rows = dense_logits_rows(chunk.max_seqlen_k) + 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) + dense_prefill_topk( + q_slice_cast[rows], + q_scale_slice[rows] if q_scale_slice is not None else None, + k_quant_cast, + k_scale_cast, + weights[tokens], + cu_seqlen_ks[rows], + cu_seqlen_ke[rows], + chunk.max_seqlen_k, + topk_indices[rows], topk_tokens, - dcp_rank, - dcp_world_size, - cp_kv_cache_interleave_size, - row_starts=row_starts, + candidate_blocks[tokens] if candidate_blocks is not None else None, + candidate_write, + candidate_block_size, + dcp_merge, ) if row_shard is not None: diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index 053c90700e2d..ed633bd6ecd2 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_mla.py +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -87,9 +87,9 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, - use_fused_mqa_topk: bool = False, + plan_fused: bool = False, ) -> list[tuple[slice, slice]]: - # workspace_size/max_logits_bytes/use_fused_mqa_topk are unused: no + # workspace_size/max_logits_bytes/plan_fused 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()) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index f95f64884487..35e0636a040e 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -12,11 +12,11 @@ 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_config_supported, + fused_mqa_topk_chunk, fused_mqa_topk_key_scales, - fused_mqa_topk_min_keys, fused_mqa_topk_prefill, get_fused_mqa_topk_workspace, + use_fused_mqa_topk, ) from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.model_executor.layers.sparse_attn_indexer import ( @@ -24,6 +24,8 @@ _merge_dcp_topk_global, _reserve_row_shard_exchange_transient, as_fp32_logits, + dense_logits_rows, + dense_prefill_topk, ) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -38,7 +40,6 @@ from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( - fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, has_deep_gemm, mqa_logits_mx_inputs, @@ -151,6 +152,16 @@ def sparse_attn_indexer_kpool( fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) + # Fused MQA+top-k prefill over the pools for this layer (per chunk: + # fused_mqa_topk_chunk). Pure DCP needs the per-shard top-k merge. + fused = use_fused_mqa_topk( + q_quant, + q_scale, + 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 @@ -164,14 +175,7 @@ def sparse_attn_indexer_kpool( ) # The fused prefill kernel's fixed-size workspace is persistent; # allocate it during profiling so it is accounted for. - if ( - envs.VLLM_USE_FUSED_MQA_TOPK - and dcp_world_size == 1 - and not use_fp4_cache - and fused_mqa_topk_config_supported( - q_quant, topk_tokens // max(index_kpool, 1) - ) - ): + if fused: get_fused_mqa_topk_workspace(q_quant.device, q_quant.shape[0]) # Reserve profiler-visible memory for the worst-case decode logits, @@ -385,27 +389,23 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens] - # Single-launch fused MQA-logits + top-k over the pools (SM100, FP8, - # 32 heads, top-k 512/1024/2048): no logits buffer, for chunks whose - # rows see at least fused_mqa_topk_min_keys(select_k) pools. It needs every - # row's whole context, so DCP (per-shard top-k + merge) stays dense. - use_fused = envs.VLLM_USE_FUSED_MQA_TOPK and dcp_world_size == 1 - if chunk.local_total_seq_lens == 0: + # Local pool ids -> request-relative global pool ids. + dcp_merge = (dcp_rank, dcp_world_size, pool_cp_interleave) + if chunk.local_total_seq_lens == 0 or num_rows == 0: topk_dst.fill_(-1) _merge_dcp_topk_global( q_slice.new_empty((num_rows, 0), dtype=torch.float32), topk_dst, select_k, - dcp_rank, - dcp_world_size, - pool_cp_interleave, + *dcp_merge, row_starts=None, ) - elif not ( - use_fused - and chunk.max_seqlen_k >= fused_mqa_topk_min_keys(select_k) - and q_scale_slice is None - and fused_mqa_topk_prefill( + # Single-launch fused MQA-logits + top-k over the pools: no logits + # buffer. It needs DCP = 1, where the pool-id merge is a no-op. + elif fused and fused_mqa_topk_chunk( + chunk.max_seqlen_k, k_quant_cast.shape[0], select_k + ): + fused_mqa_topk_prefill( q_slice_cast, k_quant_cast, fused_mqa_topk_key_scales(k_scale_cast), @@ -413,57 +413,28 @@ def sparse_attn_indexer_kpool( cu_seqlen_ks, cu_seqlen_ke, topk_dst, - # A single-request chunk has every row start at 0. - rows_start_at_zero=chunk.block_table.shape[0] == 1, ) - ): - # Fused-planned chunks are not bounded by the logits budget, so - # the dense fallback splits their rows to stay within it. - dense_rows = max(1, num_rows) - if use_fused: - budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 - dense_rows = max(1, budget // (4 * max(1, chunk.max_seqlen_k))) - for sub in range(0, max(1, num_rows), dense_rows): + else: + # Fused-planned chunks are not bounded by the logits budget: + # split their rows (budget-planned chunks fit in one call). + dense_rows = dense_logits_rows(chunk.max_seqlen_k) + for sub in range(0, num_rows, dense_rows): rows = slice(sub, min(sub + dense_rows, num_rows)) - sub_ks, sub_ke = cu_seqlen_ks[rows], cu_seqlen_ke[rows] - sub_q_scale, sub_weights = mqa_logits_mx_inputs( + dense_prefill_topk( q_slice_cast[rows], q_scale_slice[rows] if q_scale_slice is not None else None, + k_quant_cast, + k_scale_cast, weights[row_start + rows.start : row_start + rows.stop], - ) - # Compressed: row i starts at K position cu_seqlen_ks[i]. - row_ends = sub_ke - sub_ks - logits = as_fp32_logits( - fp8_fp4_mqa_logits( - (q_slice_cast[rows], sub_q_scale), - (k_quant_cast, k_scale_cast), - sub_weights, - sub_ks, - sub_ke, - chunk.max_seqlen_k, - ), - row_ends, - ) - sub_dst = topk_dst[rows] - torch.ops._C.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends), - row_ends, - sub_dst, - logits.shape[0], - logits.stride(0), - logits.stride(1), - select_k, - ) - # Local pool ids -> request-relative global pool ids. - _merge_dcp_topk_global( - logits, - sub_dst, + cu_seqlen_ks[rows], + cu_seqlen_ke[rows], + chunk.max_seqlen_k, + topk_dst[rows], select_k, - dcp_rank, - dcp_world_size, - pool_cp_interleave, - row_starts=None, + None, + False, + 0, + dcp_merge, ) if index_kpool > 1: diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 1f98924ca34f..c98fd2fc1f8f 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -14,6 +14,7 @@ FUSED_MQA_TOPK_MAX_KEYS, FUSED_MQA_TOPK_MAX_ROWS, fused_mqa_topk_available, + fused_mqa_topk_model_supported, ) from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( @@ -1031,6 +1032,18 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: else 0 ) 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). Pure DCP and the FP4 cache stay dense; PCP + DCP + # has its own planner. + model_config = self.vllm_config.model_config + self.plan_fused_mqa_topk = ( + envs.VLLM_USE_FUSED_MQA_TOPK + and not self.indexer_uses_fp4 + and self.dcp_world_size == 1 + and model_config is not None + and fused_mqa_topk_model_supported(model_config.hf_text_config) + and fused_mqa_topk_available() + ) next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n @@ -1367,7 +1380,7 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, - use_fused_mqa_topk: bool = False, + plan_fused: 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) @@ -1377,19 +1390,18 @@ 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 ``use_fused_mqa_topk`` the fused MQA+top-k kernel needs no logits - buffer: chunks whose gathered keys fit its 1M-key limit are not bounded - by the logits budget but by FUSED_MQA_TOPK_MAX_ROWS rows (one fused call - per chunk; None = no row limit). Requests are never packed past that key - limit, so only a single request longer than 1M keys takes the - logits-budget path. Chunks the dispatch keeps dense are row-split by the - indexer. + With ``plan_fused`` (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. """ def fused(n: int) -> bool: - return use_fused_mqa_topk and n <= FUSED_MQA_TOPK_MAX_KEYS + return plan_fused and n <= FUSED_MQA_TOPK_MAX_KEYS chunks: list[tuple[slice, slice]] = [] n = len(compressed_seq_lens_cpu) @@ -1406,12 +1418,10 @@ def fused(n: int) -> bool: ) new_m, new_n = chunk_m + q, chunk_n + s if fused(new_n): - fits = ( - FUSED_MQA_TOPK_MAX_ROWS is None - or new_m <= FUSED_MQA_TOPK_MAX_ROWS - ) + fits = new_m <= FUSED_MQA_TOPK_MAX_ROWS else: - fits = not use_fused_mqa_topk and new_m * new_n <= max_logits_elems + # Fused planning never packs past the key limit. + fits = not plan_fused 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 @@ -1432,7 +1442,7 @@ def fused(n: int) -> bool: max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) ) if fused(chunk_n): - max_q = FUSED_MQA_TOPK_MAX_ROWS or max(1, chunk_m) + 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))) @@ -1562,23 +1572,13 @@ def build( request_offset=num_decodes, ) else: - chunk_options = {} - if ( - envs.VLLM_USE_FUSED_MQA_TOPK - and current_platform.is_cuda() - and current_platform.is_device_capability_family(100) - and not self.indexer_uses_fp4 - and self.dcp_world_size == 1 - and fused_mqa_topk_available() - ): - chunk_options["use_fused_mqa_topk"] = True chunk_specs = self._split_indexer_prefill_chunks( self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), prefill_query_lens_cpu, self.max_prefill_buffer_size, max_logits_bytes, request_offset=num_decodes, - **chunk_options, + plan_fused=self.plan_fused_mqa_topk, ) chunks = [] From 93a0925fb26215fb9008d83f792ecaa52d53f938 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 13:48:15 +0000 Subject: [PATCH 11/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 1 - .../kernels/attention/dsa/fused_mqa_topk.py | 11 +- .../layers/sparse_attn_indexer.py | 246 ++++++++---------- vllm/models/glm5next/nvidia/sparse_indexer.py | 82 +++--- 4 files changed, 167 insertions(+), 173 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 01077db83370..0890ef665150 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1705,7 +1705,6 @@ if (VLLM_GPU_LANG STREQUAL "CUDA") 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) - # Single-launch fused MQA-logits + top-k prefill kernel. 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 diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 5f702528c55e..b13606b8ef74 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -208,14 +208,13 @@ def get_fused_mqa_topk_workspace( def fused_mqa_topk_key_scales(scales: torch.Tensor) -> torch.Tensor: - """Per-key fp32 scales from the indexer K cache. + """Per-key fp32 scales from the gathered indexer K-cache scales. - SM100 caches store each power-of-two scale as four equal UE8M0 bytes in an - int32; the fp32 value 2**(e - 127) has bit pattern ``e << 23``. + 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``. """ - if scales.dtype == torch.int32: - return ((scales & 0xFF) << 23).view(torch.float32) - return scales + exponents = scales.view(torch.int32).squeeze(-1) & 0xFF + return (exponents << 23).view(torch.float32) def fused_mqa_topk_prefill( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index c13a51a13f5f..d02ad731fb5a 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -343,87 +343,6 @@ def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor return out -def dense_prefill_topk( - q: torch.Tensor, - q_scale: torch.Tensor | None, - k: torch.Tensor, - k_scale: torch.Tensor, - weights: torch.Tensor, - cu_seqlen_ks: torch.Tensor, - cu_seqlen_ke: torch.Tensor, - max_seqlen_k: int, - topk_indices: torch.Tensor, - topk_tokens: int, - candidates: torch.Tensor | None, - candidate_write: bool, - candidate_block_size: int, - dcp_merge: tuple[int, int, int] | None, -) -> None: - """Dense MQA logits + per-row top-k of one row slice of a prefill chunk. - - candidates: the slice's v4.1 candidate blocks (None: no two-level - selection). dcp_merge: (dcp_rank, dcp_world_size, cp_interleave) to merge - the per-shard top-k into the global one (None under PCP, where the top-k - already ran over the whole context). - """ - # 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 - if current_platform.is_xpu(): - if q_scale is not None: - raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") - logits = torch.ops.vllm.xpu_fp8_mqa_logits( - q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke - ) - row_starts = cu_seqlen_ks - row_ends = cu_seqlen_ke - else: - q_scale, weights = mqa_logits_mx_inputs(q, q_scale, weights) - row_ends = cu_seqlen_ke - cu_seqlen_ks - logits = as_fp32_logits( - fp8_fp4_mqa_logits( - (q, q_scale), - (k, k_scale), - weights, - cu_seqlen_ks, - cu_seqlen_ke, - max_seqlen_k, - ), - row_ends, - ) - if candidates 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. - if candidate_write: - _select_candidate_blocks( - logits, - row_starts, - row_ends, - candidates.shape[1], - candidate_block_size, - candidates, - ) - else: - _apply_candidate_mask( - logits, row_starts, row_ends, candidates, candidate_block_size - ) - ops.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends) if row_starts is None else row_starts, - row_ends, - topk_indices, - logits.shape[0], - logits.stride(0), - logits.stride(1), - topk_tokens, - ) - if dcp_merge is not None: - _merge_dcp_topk_global( - logits, topk_indices, topk_tokens, *dcp_merge, row_starts=row_starts - ) - - def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, @@ -804,70 +723,127 @@ def sparse_attn_indexer( q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] - # Under PCP the top-k runs over the whole gathered context: no merge. - dcp_merge = ( - (dcp_rank, dcp_world_size, cp_kv_cache_interleave_size) - if deinterleave_idx is None - else None - ) + # 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 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) - if dcp_merge is not None: - # Every DCP rank joins the merge, also with no local keys. - _merge_dcp_topk_global( - q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32), - topk_indices, - topk_tokens, - *dcp_merge, - ) - continue - # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 - # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. - if use_fp4_cache: - q_slice_cast = q_slice.view(torch.int8) - k_quant_cast = k_quant.view(torch.int8) - k_scale_cast = k_scale.view(torch.int32).squeeze(-1) - else: - q_slice_cast = q_slice - k_quant_cast = k_quant - k_scale_cast = k_scale.view( - torch.int32 if mqa_logits_use_mx_scales() else torch.float32 - ).squeeze(-1) - # Single-launch fused MQA-logits + top-k: no logits buffer. - if fused and fused_mqa_topk_chunk( - chunk.max_seqlen_k, k_quant_cast.shape[0], topk_tokens + elif fused and fused_mqa_topk_chunk( + chunk.max_seqlen_k, k_quant.shape[0], topk_tokens ): + # Single-launch fused MQA-logits + top-k: no logits buffer. It + # ranks over the full context (the DCP merge below is a no-op). fused_mqa_topk_prefill( - q_slice_cast, - k_quant_cast, - fused_mqa_topk_key_scales(k_scale_cast), + q_slice, + k_quant, + fused_mqa_topk_key_scales(k_scale), weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, topk_indices, ) - continue - # Fused-planned chunks are not bounded by the logits budget: split - # their rows (budget-planned chunks fit in one call). - dense_rows = dense_logits_rows(chunk.max_seqlen_k) - 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) - dense_prefill_topk( - q_slice_cast[rows], - q_scale_slice[rows] if q_scale_slice is not None else None, - k_quant_cast, - k_scale_cast, - weights[tokens], - cu_seqlen_ks[rows], - cu_seqlen_ke[rows], - chunk.max_seqlen_k, - topk_indices[rows], + 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. + if use_fp4_cache: + q_slice_cast = q_slice.view(torch.int8) + k_quant_cast = k_quant.view(torch.int8) + k_scale_cast = k_scale.view(torch.int32).squeeze(-1) + else: + q_slice_cast = q_slice + k_quant_cast = k_quant + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) + # Fused-planned chunks 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). + dense_rows = dense_logits_rows(chunk.max_seqlen_k) + 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 = 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, + ) + 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, + ) + ops.top_k_per_row_prefill( + logits, + torch.zeros_like(row_ends) + if row_starts is None + else row_starts, + row_ends, + topk_indices[rows], + 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 + # context. + _merge_dcp_topk_global( + logits, + topk_indices, topk_tokens, - candidate_blocks[tokens] if candidate_blocks is not None else None, - candidate_write, - candidate_block_size, - dcp_merge, + dcp_rank, + dcp_world_size, + cp_kv_cache_interleave_size, + row_starts=row_starts, ) if row_shard is not None: diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 35e0636a040e..17c773826ef4 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -25,7 +25,6 @@ _reserve_row_shard_exchange_transient, as_fp32_logits, dense_logits_rows, - dense_prefill_topk, ) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -40,6 +39,7 @@ from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, has_deep_gemm, mqa_logits_mx_inputs, @@ -389,53 +389,73 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens] - # Local pool ids -> request-relative global pool ids. - dcp_merge = (dcp_rank, dcp_world_size, pool_cp_interleave) - if chunk.local_total_seq_lens == 0 or num_rows == 0: + if chunk.local_total_seq_lens == 0: + logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) topk_dst.fill_(-1) - _merge_dcp_topk_global( - q_slice.new_empty((num_rows, 0), dtype=torch.float32), - topk_dst, - select_k, - *dcp_merge, - row_starts=None, - ) - # Single-launch fused MQA-logits + top-k over the pools: no logits - # buffer. It needs DCP = 1, where the pool-id merge is a no-op. elif fused and fused_mqa_topk_chunk( - chunk.max_seqlen_k, k_quant_cast.shape[0], select_k + chunk.max_seqlen_k, k_quant.shape[0], select_k ): + # Single-launch fused MQA-logits + top-k over the pools: no + # logits buffer. It needs DCP = 1, where the merge is a no-op. fused_mqa_topk_prefill( - q_slice_cast, - k_quant_cast, - fused_mqa_topk_key_scales(k_scale_cast), + q_slice, + k_quant, + fused_mqa_topk_key_scales(k_scale), weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, topk_dst, ) + logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) else: # Fused-planned chunks are not bounded by the logits budget: - # split their rows (budget-planned chunks fit in one call). + # 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. dense_rows = dense_logits_rows(chunk.max_seqlen_k) - for sub in range(0, num_rows, dense_rows): + for sub in range(0, max(1, num_rows), dense_rows): rows = slice(sub, min(sub + dense_rows, num_rows)) - dense_prefill_topk( + 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, - k_quant_cast, - k_scale_cast, weights[row_start + rows.start : row_start + rows.stop], - cu_seqlen_ks[rows], - cu_seqlen_ke[rows], - chunk.max_seqlen_k, - topk_dst[rows], - select_k, - None, - False, - 0, - dcp_merge, ) + # 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, + topk_dst, + select_k, + dcp_rank, + dcp_world_size, + pool_cp_interleave, + row_starts=None, + ) if index_kpool > 1: pool_ids = pool_topk.to(torch.int64) From 9e3197f3ac7219578a9388ac7dff5118deb784a0 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 14:31:00 +0000 Subject: [PATCH 12/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- .../kernels/attention/test_fused_mqa_topk.py | 17 ------- .../test_indexer_fused_mqa_topk_planner.py | 45 +++++++++------- .../kernels/attention/dsa/fused_mqa_topk.py | 51 ++++++------------- .../layers/sparse_attn_indexer.py | 5 +- vllm/models/glm5next/nvidia/sparse_indexer.py | 5 +- vllm/v1/attention/backends/mla/indexer.py | 26 ++++++---- 6 files changed, 63 insertions(+), 86 deletions(-) diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index f6a3d404e9a3..8ffa9b27ca53 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -15,7 +15,6 @@ fused_mqa_topk_min_keys, fused_mqa_topk_prefill, get_fused_mqa_topk_workspace, - use_fused_mqa_topk, ) from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits from vllm.utils.deep_gemm import ( @@ -297,22 +296,6 @@ def test_cuda_graph_replay(): _check(out, q, k, scales, weights, starts, ends) -@requires_fused -def test_dispatch_predicate(): - """use_fused_mqa_topk: FP8 32x128 queries, top-k 512/1024/2048, full - context, no FP4 cache / FP4 queries / candidate selection.""" - q = _inputs(16, 4096)[0] - ok = dict(fp4_cache=False, full_context=True) - assert use_fused_mqa_topk(q, None, 2048, **ok) - assert not use_fused_mqa_topk(q, None, 768, **ok) - assert not use_fused_mqa_topk(q[:, :16], None, 2048, **ok) - assert not use_fused_mqa_topk(q.float(), None, 2048, **ok) - assert not use_fused_mqa_topk(q, torch.ones(16, device="cuda"), 2048, **ok) - assert not use_fused_mqa_topk(q, None, 2048, fp4_cache=True, full_context=True) - assert not use_fused_mqa_topk(q, None, 2048, fp4_cache=False, full_context=False) - assert not use_fused_mqa_topk(q, None, 2048, candidate_selection=True, **ok) - - @requires_fused def test_op_rejects_bad_inputs(): """The launcher does not re-check; the op enforces its contract.""" diff --git a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py index 19a4a3dbfdef..5f7f9f15569a 100644 --- a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py +++ b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py @@ -4,18 +4,17 @@ indexer path (CPU only; the kernel is tested in tests/kernels/attention/test_fused_mqa_topk.py).""" -from types import SimpleNamespace - import pytest import torch import vllm.envs as envs +import vllm.model_executor.kernels.attention.dsa.fused_mqa_topk as fused_mqa_topk from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( FUSED_MQA_TOPK_MAX_KEYS, FUSED_MQA_TOPK_MAX_ROWS, fused_mqa_topk_chunk, fused_mqa_topk_min_keys, - fused_mqa_topk_model_supported, + use_fused_mqa_topk, ) from vllm.model_executor.layers.sparse_attn_indexer import dense_logits_rows from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder @@ -97,25 +96,37 @@ def test_fused_chunk_thresholds(topk): @pytest.mark.parametrize( "heads, head_dim, topk, kpool, supported", [ - (32, 128, 2048, None, True), # GLM-5.2 - (32, 128, 512, None, True), # DeepSeek-V4.1-Flash + (32, 128, 2048, 1, True), # GLM-5.2 + (32, 128, 512, 1, True), # DeepSeek-V4.1-Flash (32, 128, 2048, 4, True), # GLM-5.3-Flash: top-512 pools - (64, 128, 2048, None, False), # DeepSeek-V3.2: 64 index heads - (32, 64, 2048, None, False), - (32, 128, 256, None, False), + (64, 128, 2048, 1, False), # DeepSeek-V3.2: 64 index heads + (32, 64, 2048, 1, False), + (32, 128, 256, 1, False), (32, 128, 2048, 3, False), # 682 pools - (32, 128, None, None, False), # no indexer ], ) -def test_model_supported(heads, head_dim, topk, kpool, supported): - """The chunk planner only plans fused chunks for indexer shapes the kernel - serves; other DSA models keep the logits-budget planning.""" - config = SimpleNamespace( - index_n_heads=heads, index_head_dim=head_dim, index_topk=topk +def test_use_fused_mqa_topk(monkeypatch, heads, head_dim, topk, kpool, supported): + """The single eligibility rule of the planner and the indexer layers.""" + monkeypatch.setattr(fused_mqa_topk, "fused_mqa_topk_available", lambda: True) + monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", True) + ok = dict( + query_shape=(heads, head_dim), + topk=topk // kpool, + fp4_cache=False, + full_context=True, ) - if kpool is not None: - config.index_kpool = kpool - assert fused_mqa_topk_model_supported(config) == supported + assert use_fused_mqa_topk(**ok) == supported + if supported: + assert not use_fused_mqa_topk(**{**ok, "fp4_cache": True}) + assert not use_fused_mqa_topk(**{**ok, "full_context": False}) + assert not use_fused_mqa_topk(**ok, candidate_selection=True) + # A flattened [tokens, heads * head_dim] query does not match. + assert not use_fused_mqa_topk(**{**ok, "query_shape": (heads * head_dim,)}) + monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", False) + assert not use_fused_mqa_topk(**ok) + monkeypatch.setattr(fused_mqa_topk, "fused_mqa_topk_available", lambda: False) + monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", True) + assert not use_fused_mqa_topk(**ok) def test_dense_logits_rows_fits_budget(monkeypatch): diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index b13606b8ef74..a5b32a04b2f1 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -17,7 +17,6 @@ import functools from dataclasses import dataclass -from typing import Any import torch @@ -91,51 +90,33 @@ def fused_mqa_topk_available() -> bool: ) -def fused_mqa_topk_model_supported(hf_text_config: Any) -> bool: - """Whether a sparse-attention model's indexer shape fits the fused kernel: - 32 index heads of dimension 128 and a per-row top-k (per pool for kpool - indexers) of 512, 1024 or 2048. Lets the prefill chunk planner skip the - fused planning for models the kernel cannot serve (e.g. DeepSeek-V3.2, 64 - index heads).""" - topk = getattr(hf_text_config, "index_topk", None) - kpool = getattr(hf_text_config, "index_kpool", None) or 1 - return ( - getattr(hf_text_config, "index_n_heads", None) == FUSED_MQA_TOPK_HEADS - and getattr(hf_text_config, "index_head_dim", None) == FUSED_MQA_TOPK_HEAD_DIM - and topk is not None - and topk // kpool in FUSED_MQA_TOPK_TOPKS - ) - - def use_fused_mqa_topk( - q: torch.Tensor, - q_scale: torch.Tensor | None, - topk_tokens: int, *, + query_shape: tuple[int, ...], + topk: int, fp4_cache: bool, full_context: bool, candidate_selection: bool = False, ) -> bool: - """Whether an indexer layer call may use the fused kernel for its prefill - chunks (each chunk then also needs fused_mqa_topk_chunk). - - q / q_scale: the layer's quantized indexer queries; topk_tokens: the top-k - the kernel selects (pools for kpool indexers). 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. The FP4 cache, FP4 - queries and v4.1 candidate selection stay on the dense path. + """The one eligibility rule of the fused kernel, shared by 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. """ return ( envs.VLLM_USE_FUSED_MQA_TOPK - and full_context and not fp4_cache - and q_scale is None + and full_context and not candidate_selection - and topk_tokens in FUSED_MQA_TOPK_TOPKS - and q.dtype == torch.float8_e4m3fn - and q.dim() == 3 - and tuple(q.shape[1:]) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM) - and q.is_cuda + 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() ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index d02ad731fb5a..61b968587be4 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -500,9 +500,8 @@ def sparse_attn_indexer( # Fused MQA+top-k prefill for this layer (per chunk: fused_mqa_topk_chunk). # Under PCP every rank ranks its rows over the whole (gathered) context. fused = use_fused_mqa_topk( - q_quant, - q_scale, - topk_tokens, + 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, diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 17c773826ef4..1d1ee85d77b2 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -155,9 +155,8 @@ def sparse_attn_indexer_kpool( # Fused MQA+top-k prefill over the pools for this layer (per chunk: # fused_mqa_topk_chunk). Pure DCP needs the per-shard top-k merge. fused = use_fused_mqa_topk( - q_quant, - q_scale, - topk_tokens // max(index_kpool, 1), + query_shape=q_quant.shape[1:], + topk=topk_tokens // max(index_kpool, 1), fp4_cache=use_fp4_cache, full_context=dcp_world_size == 1, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index c98fd2fc1f8f..360a8fc78cfe 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -13,8 +13,7 @@ from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( FUSED_MQA_TOPK_MAX_KEYS, FUSED_MQA_TOPK_MAX_ROWS, - fused_mqa_topk_available, - fused_mqa_topk_model_supported, + 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 ( @@ -1033,16 +1032,21 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: ) 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). Pure DCP and the FP4 cache stay dense; PCP + DCP - # has its own planner. + # the logits budget). PCP + DCP has its own planner. v4.1 candidate + # selection is not excluded: its pre-candidate indexers take the + # fused kernel and the others row-split their dense logits. Indexers + # that pad their heads to 32 run fused on budget-planned chunks. model_config = self.vllm_config.model_config - self.plan_fused_mqa_topk = ( - envs.VLLM_USE_FUSED_MQA_TOPK - and not self.indexer_uses_fp4 - and self.dcp_world_size == 1 - and model_config is not None - and fused_mqa_topk_model_supported(model_config.hf_text_config) - and fused_mqa_topk_available() + cfg = model_config.hf_text_config if model_config is not None else None + self.plan_fused_mqa_topk = 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 From 4b06112d2cb230a96ad6f43c44adf069a3bb396a Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 16:08:26 +0000 Subject: [PATCH 13/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- .../test_indexer_fused_mqa_topk_planner.py | 6 ++++-- .../layers/sparse_attn_indexer.py | 13 +------------ vllm/models/glm5next/nvidia/sparse_indexer.py | 2 +- vllm/v1/attention/backends/mla/indexer.py | 17 ++++++++++++++++- 4 files changed, 22 insertions(+), 16 deletions(-) diff --git a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py index 5f7f9f15569a..e2208647f7ab 100644 --- a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py +++ b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py @@ -16,8 +16,10 @@ fused_mqa_topk_min_keys, use_fused_mqa_topk, ) -from vllm.model_executor.layers.sparse_attn_indexer import dense_logits_rows -from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, + dense_logits_rows, +) LOGITS_BUDGET = 512 * 1024**2 diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 61b968587be4..d9fab3350608 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -61,6 +61,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 @@ -308,18 +309,6 @@ def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.cons tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask) -def dense_logits_rows(max_seqlen_k: int) -> int: - """Query rows per dense MQA-logits call that keep its fp32 [rows, - max_seqlen_k] logits within VLLM_SPARSE_INDEXER_MAX_LOGITS_MB. - - Chunks the planner bounded by the logits budget already fit in one call; - chunks planned for the fused kernel (not bounded by it) are split when they - take the dense path. - """ - budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 - return max(1, budget // (4 * max(1, max_seqlen_k))) - - def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor: """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels. diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 1d1ee85d77b2..2585e5923970 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -24,7 +24,6 @@ _merge_dcp_topk_global, _reserve_row_shard_exchange_transient, as_fp32_logits, - dense_logits_rows, ) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -52,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 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 360a8fc78cfe..8b015f6697a2 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -932,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 fused-planned chunks (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() @@ -1443,7 +1456,9 @@ def fused(n: int) -> bool: 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(chunk_n): max_q = FUSED_MQA_TOPK_MAX_ROWS From 696df2a6c06f8eca7613e168c25979d27c485571 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 16:54:28 +0000 Subject: [PATCH 14/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- .../kernels/attention/test_fused_mqa_topk.py | 56 +++++++++++-------- .../test_indexer_fused_mqa_topk_planner.py | 35 +++++++----- .../kernels/attention/dsa/fused_mqa_topk.py | 19 ++++--- .../layers/sparse_attn_indexer.py | 15 ++--- vllm/models/deepseek_v4/cpu/cpu_mla.py | 4 +- vllm/models/glm5next/nvidia/sparse_indexer.py | 15 ++--- vllm/v1/attention/backends/mla/indexer.py | 29 +++++----- 7 files changed, 93 insertions(+), 80 deletions(-) diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 8ffa9b27ca53..a950b3c55200 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -28,7 +28,7 @@ DeepseekV32IndexerPrefillMetadata, ) -requires_fused = pytest.mark.skipif( +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", ) @@ -125,7 +125,7 @@ def _inputs(rows, keys, seed=913): return q, k, scales, weights -@requires_fused +@requires_fused_mqa_topk @pytest.mark.parametrize("topk", TOPKS) @pytest.mark.parametrize( "keys,rows,case", @@ -185,7 +185,7 @@ def test_single_request_matches_dense(keys, rows, case, topk): assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all() -@requires_fused +@requires_fused_mqa_topk @pytest.mark.parametrize("topk", TOPKS) @pytest.mark.parametrize( "lens", [[3000, 5000, 1000], [40000, 1500, 70000, 9], [600000, 200000]] @@ -210,7 +210,8 @@ def test_multi_request_chunk_matches_dense(lens, topk): 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 keys are exact). + # (~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()) @@ -218,7 +219,7 @@ def test_multi_request_chunk_matches_dense(lens, topk): assert overlap >= 0.98 * int((dense >= 0).sum()) -@requires_fused +@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) @@ -233,7 +234,7 @@ def test_large_chunk_single_call(rows): _check(out, q, k, scales, weights, starts, ends) -@requires_fused +@requires_fused_mqa_topk def test_rescan_scores_only_its_block(): """The exact rescan of a block touches only that block's 8 rows. @@ -274,7 +275,7 @@ def test_rescan_scores_only_its_block(): ) -@requires_fused +@requires_fused_mqa_topk def test_cuda_graph_replay(): rows, keys = 64, 65536 q, k, scales, weights = _inputs(rows, keys) @@ -296,7 +297,7 @@ def test_cuda_graph_replay(): _check(out, q, k, scales, weights, starts, ends) -@requires_fused +@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) @@ -339,7 +340,7 @@ def _kpool_inputs(pool_lens=(20000, 70000), q_rows=96, tail=2, seed=919): ) -def _run_kpool_indexer(monkeypatch, inputs, fused, budget_mb=512): +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.""" q, k, scales, weights, starts, ends, positions = inputs rows, keys = q.shape[0], k.shape[0] @@ -378,7 +379,7 @@ def _run_kpool_indexer(monkeypatch, inputs, fused, budget_mb=512): lambda: SimpleNamespace(attn_metadata={"indexer": metadata}), ) monkeypatch.setattr(kpool_indexer, "current_workspace_manager", lambda: workspace) - monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", fused) + 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 @@ -415,40 +416,47 @@ def _row_sets(out): return [set(r[r >= 0].tolist()) for r in out.cpu()] -@requires_fused +@requires_fused_mqa_topk def test_kpool_indexer_fused_matches_dense(monkeypatch): - """GLM-5.3-Flash kpool indexer: the fused kernel ranks the top-512 pools of + """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() - fused, calls = _run_kpool_indexer(monkeypatch, inputs, fused=True) + fused_mqa_topk_out, calls = _run_kpool_indexer( + monkeypatch, inputs, fused_mqa_topk=True + ) assert calls == [torch.Size([inputs[0].shape[0], TOPK // KPOOL])] - dense, calls = _run_kpool_indexer(monkeypatch, inputs, fused=False) + dense, calls = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=False) assert calls == [] - assert ((fused >= 0).sum(1) == (dense >= 0).sum(1)).all() + 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), _row_sets(dense))) + 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 +@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 kernel is not faster there).""" + 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=True) + _, calls = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=True) assert calls == [] -@requires_fused +@requires_fused_mqa_topk def test_kpool_indexer_dense_fallback_splits_rows(monkeypatch): - """Fused-planned chunks skip the logits budget; when such a chunk takes the - dense path, its rows are split to stay within the budget and select the + """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=False) - split, calls = _run_kpool_indexer(monkeypatch, inputs, fused=False, budget_mb=1) + 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. diff --git a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py index e2208647f7ab..e78779e0be5d 100644 --- a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py +++ b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py @@ -24,30 +24,32 @@ LOGITS_BUDGET = 512 * 1024**2 -def _split(seq_lens, q_lens, fused, offset=0, workspace=1048576): +def _split(seq_lens, q_lens, plan_fused_mqa_topk, offset=0, workspace=1048576): return DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( torch.tensor(seq_lens), torch.tensor(q_lens), workspace, LOGITS_BUDGET, request_offset=offset, - plan_fused=fused, + plan_fused_mqa_topk=plan_fused_mqa_topk, ) @pytest.mark.parametrize("q_rows", [16384, 32768]) -@pytest.mark.parametrize("fused", [False, True]) -def test_planner_one_fused_call_per_chunk(fused, q_rows): - # One fused call per FUSED_MQA_TOPK_MAX_ROWS rows; dense splits by the +@pytest.mark.parametrize("plan_fused_mqa_topk", [False, True]) +def test_planner_one_fused_mqa_topk_call_per_chunk(plan_fused_mqa_topk, q_rows): + # One fused_mqa_topk call per FUSED_MQA_TOPK_MAX_ROWS rows; dense splits by the # logits budget. - chunks = _split([524288], [q_rows], fused) - expected = -(-q_rows // FUSED_MQA_TOPK_MAX_ROWS) if fused else q_rows // 256 + chunks = _split([524288], [q_rows], plan_fused_mqa_topk) + expected = ( + -(-q_rows // FUSED_MQA_TOPK_MAX_ROWS) if plan_fused_mqa_topk else q_rows // 256 + ) assert len(chunks) == expected assert sum(q.stop - q.start for _, q in chunks) == q_rows def test_planner_fused_any_length_and_short_queries(): - # No length threshold and no minimum query count on the fused path. + # No length threshold and no minimum query count with fused_mqa_topk. assert _split([524288], [512], True) == [(slice(0, 1), slice(0, 512))] assert _split([4096], [4096], True) == [(slice(0, 1), slice(0, 4096))] # A 40000-row step is split at FUSED_MQA_TOPK_MAX_ROWS. @@ -58,7 +60,7 @@ def test_planner_fused_any_length_and_short_queries(): def test_planner_packs_requests_in_one_fused_call(): - # Requests share one fused call up to the row cap; dense needs a split. + # Requests share one fused_mqa_topk call up to the row cap; dense needs a split. assert _split([4096] * 4, [4096] * 4, True) == [(slice(0, 4), slice(0, 16384))] assert len(_split([4096] * 4, [4096] * 4, False)) > 1 # The second request would push the call past the row cap. @@ -70,19 +72,22 @@ def test_planner_packs_requests_in_one_fused_call(): def test_planner_fused_never_packs_past_key_limit(): # Two 600K-key continuation steps of 50 rows: dense packs them into one - # logits-budget chunk; the fused planner keeps one fused chunk per request so + # logits-budget chunk; fused_mqa_topk planning keeps one chunk per request so # no chunk reaches the dense path (whose logits buffer is not reserved). args = ([600_000, 600_000], [50, 50], False, 0, 4_000_000) assert _split(*args) == [(slice(0, 2), slice(0, 100))] - fused = (*args[:2], True, *args[3:]) - assert _split(*fused) == [(slice(0, 1), slice(0, 50)), (slice(1, 2), slice(0, 50))] + planned = (*args[:2], True, *args[3:]) + assert _split(*planned) == [ + (slice(0, 1), slice(0, 50)), + (slice(1, 2), slice(0, 50)), + ] def test_planner_keeps_budget_beyond_fused_key_limit(): - # A single request longer than the fused key limit is budget-planned. - fused = _split([2_000_000], [256], True, workspace=4_000_000) + # A single request longer than the fused_mqa_topk key limit is budget-planned. + planned = _split([2_000_000], [256], True, workspace=4_000_000) dense = _split([2_000_000], [256], False, workspace=4_000_000) - assert fused == dense and len(dense) == 4 + assert planned == dense and len(dense) == 4 @pytest.mark.parametrize("topk", [512, 1024, 2048]) diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index a5b32a04b2f1..770a71d49f3e 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -32,16 +32,16 @@ # Largest gathered key range a single call accepts (no row limit: the per-row # workspace is sized from the call, see get_fused_mqa_topk_workspace). FUSED_MQA_TOPK_MAX_KEYS = 1 << 20 -# Rows per fused call. The kernel has no row limit, but prefill chunks above +# 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, -# so the chunk planner splits fused chunks at this many rows. +# 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: """Keys per row below which the dense logits + top-k path is as fast or faster, so callers keep prefill chunks whose longest row sees fewer keys on - the dense path. B300 real captures, fused vs dense: top-512 0.89x at 16K + the dense path. B300 real captures, fused_mqa_topk vs dense: top-512 0.89x at 16K keys and 1.06x at 32K; top-1024 0.93x / 1.01x; top-2048 0.92x at 32K and 1.13x at 64K.""" return 1 << 16 if topk > 1024 else 1 << 15 @@ -55,7 +55,7 @@ def fused_mqa_topk_min_keys(topk: int) -> int: @dataclass class FusedMQATopKWorkspace: - """Caller-owned scratch of the fused kernel (see fused_mqa_topk.cu).""" + """Caller-owned scratch of the fused_mqa_topk kernel (see fused_mqa_topk.cu).""" counter: torch.Tensor # [2] int32, zero between calls (the kernel restores it) values: torch.Tensor # [8S, 16384] uint16 (low 16 bits of fp24 codes), S = 2/SM @@ -82,7 +82,7 @@ def as_args(self) -> tuple[torch.Tensor, ...]: @functools.cache def fused_mqa_topk_available() -> bool: - """True when the fused op is compiled in and the GPU is SM100-family.""" + """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) @@ -98,7 +98,7 @@ def use_fused_mqa_topk( full_context: bool, candidate_selection: bool = False, ) -> bool: - """The one eligibility rule of the fused kernel, shared by the prefill + """The one eligibility rule of the fused_mqa_topk kernel, shared by the prefill chunk planner and the indexer layers (each prefill chunk then also needs fused_mqa_topk_chunk). @@ -122,9 +122,10 @@ def use_fused_mqa_topk( def fused_mqa_topk_chunk(max_seqlen_k: int, keys: int, topk_tokens: int) -> bool: - """Whether a prefill chunk of a fused-eligible layer takes the fused - kernel: its longest row sees at least fused_mqa_topk_min_keys keys (the - dense path is as fast below) and its gathered keys fit the kernel.""" + """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) and its gathered keys fit the + kernel.""" return ( max_seqlen_k >= fused_mqa_topk_min_keys(topk_tokens) and keys <= FUSED_MQA_TOPK_MAX_KEYS diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index d9fab3350608..97ee775de090 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -488,7 +488,7 @@ def sparse_attn_indexer( # Fused MQA+top-k prefill for this layer (per chunk: fused_mqa_topk_chunk). # Under PCP every rank ranks its rows over the whole (gathered) context. - fused = use_fused_mqa_topk( + fused_mqa_topk = use_fused_mqa_topk( query_shape=q_quant.shape[1:], topk=topk_tokens, fp4_cache=use_fp4_cache, @@ -515,8 +515,8 @@ def sparse_attn_indexer( ) profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) - if fused: - # The fused prefill kernel's fixed-size workspace is persistent; + if fused_mqa_topk: + # The fused_mqa_topk kernel's fixed-size workspace is persistent; # allocate it during profiling so it is accounted for. get_fused_mqa_topk_workspace(q_quant.device, q_quant.shape[0]) @@ -717,7 +717,7 @@ def sparse_attn_indexer( 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 and fused_mqa_topk_chunk( + elif fused_mqa_topk and fused_mqa_topk_chunk( chunk.max_seqlen_k, k_quant.shape[0], topk_tokens ): # Single-launch fused MQA-logits + top-k: no logits buffer. It @@ -745,9 +745,10 @@ def sparse_attn_indexer( k_scale_cast = k_scale.view( torch.int32 if mqa_logits_use_mx_scales() else torch.float32 ).squeeze(-1) - # Fused-planned chunks 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). + # 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). dense_rows = dense_logits_rows(chunk.max_seqlen_k) for sub in range(0, q_slice.shape[0], dense_rows): rows = slice(sub, min(sub + dense_rows, q_slice.shape[0])) diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index ed633bd6ecd2..e3729f68510c 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_mla.py +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -87,9 +87,9 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, - plan_fused: bool = False, + plan_fused_mqa_topk: bool = False, ) -> list[tuple[slice, slice]]: - # workspace_size/max_logits_bytes/plan_fused are unused: no + # workspace_size/max_logits_bytes/plan_fused_mqa_topk 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()) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 2585e5923970..9a7962add0b8 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -154,7 +154,7 @@ def sparse_attn_indexer_kpool( # Fused MQA+top-k prefill over the pools for this layer (per chunk: # fused_mqa_topk_chunk). Pure DCP needs the per-shard top-k merge. - fused = use_fused_mqa_topk( + 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, @@ -172,9 +172,9 @@ def sparse_attn_indexer_kpool( scales_spec, ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) - # The fused prefill kernel's fixed-size workspace is persistent; + # The fused_mqa_topk kernel's fixed-size workspace is persistent; # allocate it during profiling so it is accounted for. - if fused: + if fused_mqa_topk: get_fused_mqa_topk_workspace(q_quant.device, q_quant.shape[0]) # Reserve profiler-visible memory for the worst-case decode logits, @@ -391,7 +391,7 @@ def sparse_attn_indexer_kpool( if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) topk_dst.fill_(-1) - elif fused and fused_mqa_topk_chunk( + elif fused_mqa_topk and fused_mqa_topk_chunk( chunk.max_seqlen_k, k_quant.shape[0], select_k ): # Single-launch fused MQA-logits + top-k over the pools: no @@ -407,9 +407,10 @@ def sparse_attn_indexer_kpool( ) logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) else: - # Fused-planned chunks 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). + # 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. dense_rows = dense_logits_rows(chunk.max_seqlen_k) for sub in range(0, max(1, num_rows), dense_rows): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 8b015f6697a2..5e5ce9af3bbc 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -937,8 +937,8 @@ def dense_logits_rows(max_seqlen_k: int, max_logits_bytes: int | None = None) -> 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 fused-planned chunks (not bounded by the budget) that take the - dense path. + 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 @@ -1046,9 +1046,10 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: 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 the - # fused kernel and the others row-split their dense logits. Indexers - # that pad their heads to 32 run fused on budget-planned chunks. + # 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.plan_fused_mqa_topk = cfg is not None and use_fused_mqa_topk( @@ -1397,7 +1398,7 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, - plan_fused: bool = False, + plan_fused_mqa_topk: 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) @@ -1407,7 +1408,7 @@ 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 ``plan_fused`` (chunks for the fused MQA+top-k kernel, which needs + With ``plan_fused_mqa_topk`` (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 @@ -1416,10 +1417,6 @@ def _split_indexer_prefill_chunks( Returns list of (req_slice, query_slice) tuples. """ - - def fused(n: int) -> bool: - return plan_fused and n <= FUSED_MQA_TOPK_MAX_KEYS - chunks: list[tuple[slice, slice]] = [] n = len(compressed_seq_lens_cpu) max_logits_elems = max_logits_bytes // 4 @@ -1434,11 +1431,11 @@ def fused(n: int) -> bool: compressed_seq_lens_cpu[end].item(), ) new_m, new_n = chunk_m + q, chunk_n + s - if fused(new_n): + if plan_fused_mqa_topk and new_n <= FUSED_MQA_TOPK_MAX_KEYS: fits = new_m <= FUSED_MQA_TOPK_MAX_ROWS else: - # Fused planning never packs past the key limit. - fits = not plan_fused and new_m * new_n <= max_logits_elems + # fused_mqa_topk planning never packs past the key limit. + fits = not plan_fused_mqa_topk 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 @@ -1460,7 +1457,7 @@ def fused(n: int) -> bool: if chunk_n > 0 else max(1, chunk_m) ) - if fused(chunk_n): + if plan_fused_mqa_topk 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) @@ -1597,7 +1594,7 @@ def build( self.max_prefill_buffer_size, max_logits_bytes, request_offset=num_decodes, - plan_fused=self.plan_fused_mqa_topk, + plan_fused_mqa_topk=self.plan_fused_mqa_topk, ) chunks = [] From f698c03f01b4aa213271ff373bb63a7a805ad153 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 17:23:40 +0000 Subject: [PATCH 15/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- .../test_indexer_fused_mqa_topk_planner.py | 153 ------------------ .../kernels/attention/dsa/fused_mqa_topk.py | 36 ++--- .../layers/sparse_attn_indexer.py | 4 - vllm/models/glm5next/nvidia/sparse_indexer.py | 4 - 4 files changed, 9 insertions(+), 188 deletions(-) delete mode 100644 tests/v1/attention/test_indexer_fused_mqa_topk_planner.py diff --git a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py b/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py deleted file mode 100644 index e78779e0be5d..000000000000 --- a/tests/v1/attention/test_indexer_fused_mqa_topk_planner.py +++ /dev/null @@ -1,153 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Prefill chunk planning and dispatch predicates of the fused MQA+top-k -indexer path (CPU only; the kernel is tested in -tests/kernels/attention/test_fused_mqa_topk.py).""" - -import pytest -import torch - -import vllm.envs as envs -import vllm.model_executor.kernels.attention.dsa.fused_mqa_topk as fused_mqa_topk -from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( - FUSED_MQA_TOPK_MAX_KEYS, - FUSED_MQA_TOPK_MAX_ROWS, - fused_mqa_topk_chunk, - fused_mqa_topk_min_keys, - use_fused_mqa_topk, -) -from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV32IndexerMetadataBuilder, - dense_logits_rows, -) - -LOGITS_BUDGET = 512 * 1024**2 - - -def _split(seq_lens, q_lens, plan_fused_mqa_topk, offset=0, workspace=1048576): - return DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( - torch.tensor(seq_lens), - torch.tensor(q_lens), - workspace, - LOGITS_BUDGET, - request_offset=offset, - plan_fused_mqa_topk=plan_fused_mqa_topk, - ) - - -@pytest.mark.parametrize("q_rows", [16384, 32768]) -@pytest.mark.parametrize("plan_fused_mqa_topk", [False, True]) -def test_planner_one_fused_mqa_topk_call_per_chunk(plan_fused_mqa_topk, q_rows): - # One fused_mqa_topk call per FUSED_MQA_TOPK_MAX_ROWS rows; dense splits by the - # logits budget. - chunks = _split([524288], [q_rows], plan_fused_mqa_topk) - expected = ( - -(-q_rows // FUSED_MQA_TOPK_MAX_ROWS) if plan_fused_mqa_topk else q_rows // 256 - ) - assert len(chunks) == expected - assert sum(q.stop - q.start for _, q in chunks) == q_rows - - -def test_planner_fused_any_length_and_short_queries(): - # No length threshold and no minimum query count with fused_mqa_topk. - assert _split([524288], [512], True) == [(slice(0, 1), slice(0, 512))] - assert _split([4096], [4096], True) == [(slice(0, 1), slice(0, 4096))] - # A 40000-row step is split at FUSED_MQA_TOPK_MAX_ROWS. - cap = FUSED_MQA_TOPK_MAX_ROWS - assert _split([65536], [40000], True) == [ - (slice(0, 1), slice(lo, min(lo + cap, 40000))) for lo in range(0, 40000, cap) - ] - - -def test_planner_packs_requests_in_one_fused_call(): - # Requests share one fused_mqa_topk call up to the row cap; dense needs a split. - assert _split([4096] * 4, [4096] * 4, True) == [(slice(0, 4), slice(0, 16384))] - assert len(_split([4096] * 4, [4096] * 4, False)) > 1 - # The second request would push the call past the row cap. - assert _split([524288, 262144], [16384, 9], True, offset=2) == [ - (slice(2, 3), slice(0, 16384)), - (slice(3, 4), slice(0, 9)), - ] - - -def test_planner_fused_never_packs_past_key_limit(): - # Two 600K-key continuation steps of 50 rows: dense packs them into one - # logits-budget chunk; fused_mqa_topk planning keeps one chunk per request so - # no chunk reaches the dense path (whose logits buffer is not reserved). - args = ([600_000, 600_000], [50, 50], False, 0, 4_000_000) - assert _split(*args) == [(slice(0, 2), slice(0, 100))] - planned = (*args[:2], True, *args[3:]) - assert _split(*planned) == [ - (slice(0, 1), slice(0, 50)), - (slice(1, 2), slice(0, 50)), - ] - - -def test_planner_keeps_budget_beyond_fused_key_limit(): - # A single request longer than the fused_mqa_topk key limit is budget-planned. - planned = _split([2_000_000], [256], True, workspace=4_000_000) - dense = _split([2_000_000], [256], False, workspace=4_000_000) - assert planned == dense and len(dense) == 4 - - -@pytest.mark.parametrize("topk", [512, 1024, 2048]) -def test_fused_chunk_thresholds(topk): - min_keys = fused_mqa_topk_min_keys(topk) - assert not fused_mqa_topk_chunk(min_keys - 1, min_keys, topk) - assert fused_mqa_topk_chunk(min_keys, min_keys, topk) - assert fused_mqa_topk_chunk(min_keys, FUSED_MQA_TOPK_MAX_KEYS, topk) - # Packed gathered keys past the kernel's limit stay dense. - assert not fused_mqa_topk_chunk(min_keys, FUSED_MQA_TOPK_MAX_KEYS + 1, topk) - - -@pytest.mark.parametrize( - "heads, head_dim, topk, kpool, supported", - [ - (32, 128, 2048, 1, True), # GLM-5.2 - (32, 128, 512, 1, True), # DeepSeek-V4.1-Flash - (32, 128, 2048, 4, True), # GLM-5.3-Flash: top-512 pools - (64, 128, 2048, 1, False), # DeepSeek-V3.2: 64 index heads - (32, 64, 2048, 1, False), - (32, 128, 256, 1, False), - (32, 128, 2048, 3, False), # 682 pools - ], -) -def test_use_fused_mqa_topk(monkeypatch, heads, head_dim, topk, kpool, supported): - """The single eligibility rule of the planner and the indexer layers.""" - monkeypatch.setattr(fused_mqa_topk, "fused_mqa_topk_available", lambda: True) - monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", True) - ok = dict( - query_shape=(heads, head_dim), - topk=topk // kpool, - fp4_cache=False, - full_context=True, - ) - assert use_fused_mqa_topk(**ok) == supported - if supported: - assert not use_fused_mqa_topk(**{**ok, "fp4_cache": True}) - assert not use_fused_mqa_topk(**{**ok, "full_context": False}) - assert not use_fused_mqa_topk(**ok, candidate_selection=True) - # A flattened [tokens, heads * head_dim] query does not match. - assert not use_fused_mqa_topk(**{**ok, "query_shape": (heads * head_dim,)}) - monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", False) - assert not use_fused_mqa_topk(**ok) - monkeypatch.setattr(fused_mqa_topk, "fused_mqa_topk_available", lambda: False) - monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", True) - assert not use_fused_mqa_topk(**ok) - - -def test_dense_logits_rows_fits_budget(monkeypatch): - """Dense row splitting keeps each call's fp32 [rows, max_seqlen_k] logits - within the budget, and never splits a chunk the budget planner made.""" - monkeypatch.setattr(envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", 512) - for width in (1, 4096, 65536, 1 << 20): - rows = dense_logits_rows(width) - assert rows * width * 4 <= LOGITS_BUDGET - assert (rows + 1) * width * 4 > LOGITS_BUDGET - # Budget-planned chunks have rows * total_keys <= budget / 4 and every - # row's max_seqlen_k <= total_keys, so they fit in one dense call. - seq_lens = [131072, 65536, 4096] - for req_slice, q_slice in _split(seq_lens, [4096, 4096, 4096], False): - total_keys = sum(seq_lens[req_slice]) - assert q_slice.stop - q_slice.start <= dense_logits_rows(total_keys) - assert dense_logits_rows(0) == LOGITS_BUDGET // 4 diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 770a71d49f3e..1fa08abeb933 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -1,19 +1,6 @@ # 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). - -One kernel (``torch.ops._C.fused_mqa_topk_prefill``) scores every visible key -of each query row with the DeepSeek sparse-attention indexer and selects the -row's top-k keys (k = 512, 1024 or 2048), without materializing the -``[rows, keys]`` logits matrix that the dense path (``fp8_fp4_mqa_logits`` + -``top_k_per_row_prefill``) needs. - -The kernel is persistent (one CTA per SM); each CTA owns two 8-row candidate -slots. The scratch (~296 MiB on a 148-SM GPU) is allocated once per device and -reused by every call, independently of the number of rows or keys. It is tuned -for the chunks vLLM sends it (rows that see >= fused_mqa_topk_min_keys keys); -smaller key counts are supported but not optimized. -""" +"""Single-launch fused FP8 MQA-logits + top-k prefill selection (SM100 family)""" import functools from dataclasses import dataclass @@ -26,24 +13,23 @@ 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) -# Largest gathered key range a single call accepts (no row limit: the per-row -# workspace is sized from the call, see get_fused_mqa_topk_workspace). +# 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, +# 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: - """Keys per row below which the dense logits + top-k path is as fast or - faster, so callers keep prefill chunks whose longest row sees fewer keys on - the dense path. B300 real captures, fused_mqa_topk vs dense: top-512 0.89x at 16K - keys and 1.06x at 32K; top-1024 0.93x / 1.01x; top-2048 0.92x at 32K and - 1.13x at 64K.""" + """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 @@ -98,11 +84,7 @@ def use_fused_mqa_topk( full_context: bool, candidate_selection: bool = False, ) -> bool: - """The one eligibility rule of the fused_mqa_topk kernel, shared by 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); + """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 diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 97ee775de090..7a6d814041fe 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -486,8 +486,6 @@ def sparse_attn_indexer( ) assert candidate_block_size > 0 - # Fused MQA+top-k prefill for this layer (per chunk: fused_mqa_topk_chunk). - # Under PCP every rank ranks its rows over the whole (gathered) context. fused_mqa_topk = use_fused_mqa_topk( query_shape=q_quant.shape[1:], topk=topk_tokens, @@ -720,8 +718,6 @@ def sparse_attn_indexer( elif fused_mqa_topk and fused_mqa_topk_chunk( chunk.max_seqlen_k, k_quant.shape[0], topk_tokens ): - # Single-launch fused MQA-logits + top-k: no logits buffer. It - # ranks over the full context (the DCP merge below is a no-op). fused_mqa_topk_prefill( q_slice, k_quant, diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 9a7962add0b8..99366288cd97 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -152,8 +152,6 @@ def sparse_attn_indexer_kpool( fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) - # Fused MQA+top-k prefill over the pools for this layer (per chunk: - # fused_mqa_topk_chunk). Pure DCP needs the per-shard top-k merge. fused_mqa_topk = use_fused_mqa_topk( query_shape=q_quant.shape[1:], topk=topk_tokens // max(index_kpool, 1), @@ -394,8 +392,6 @@ def sparse_attn_indexer_kpool( elif fused_mqa_topk and fused_mqa_topk_chunk( chunk.max_seqlen_k, k_quant.shape[0], select_k ): - # Single-launch fused MQA-logits + top-k over the pools: no - # logits buffer. It needs DCP = 1, where the merge is a no-op. fused_mqa_topk_prefill( q_slice, k_quant, From 22d74f667c9e5e19f0b3469a81d21dc746bb1342 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 17:42:35 +0000 Subject: [PATCH 16/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/candidates.cuh | 25 +++++++++++++++++++++++ csrc/fused_mqa_topk/common.cuh | 1 + csrc/fused_mqa_topk/deep_gemm_compat.cuh | 18 ++++++++-------- csrc/fused_mqa_topk/dense_scan.cuh | 9 ++++++-- csrc/fused_mqa_topk/fused_mqa_topk.cu | 2 ++ csrc/fused_mqa_topk/persistent_epoch.inc | 3 +++ csrc/fused_mqa_topk/persistent_kernel.cuh | 4 +++- csrc/fused_mqa_topk/rescan.cuh | 4 +++- csrc/fused_mqa_topk/select.cuh | 4 +++- 9 files changed, 57 insertions(+), 13 deletions(-) diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh index df022668af90..409176bdb66a 100644 --- a/csrc/fused_mqa_topk/candidates.cuh +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -1,5 +1,30 @@ +// 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 diff --git a/csrc/fused_mqa_topk/common.cuh b/csrc/fused_mqa_topk/common.cuh index dfb37eb3877b..4e52408a07fa 100644 --- a/csrc/fused_mqa_topk/common.cuh +++ b/csrc/fused_mqa_topk/common.cuh @@ -91,6 +91,7 @@ __device__ __forceinline__ float half_bin_floor(uint32_t b) { } // Weighted head reduction of one row's 32 head scores s_h in two FFMA2 chains +// (adapted from DeepGEMM's sm100_fp8_mqa_logits.cuh at 88965b078186) // 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. diff --git a/csrc/fused_mqa_topk/deep_gemm_compat.cuh b/csrc/fused_mqa_topk/deep_gemm_compat.cuh index 107e3c741def..ae13a12c3d57 100644 --- a/csrc/fused_mqa_topk/deep_gemm_compat.cuh +++ b/csrc/fused_mqa_topk/deep_gemm_compat.cuh @@ -1,11 +1,13 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025 DeepSeek -// FP8 MQA-logits building blocks from DeepGEMM -// (github.com/deepseek-ai/DeepGEMM) as of the release before its SM100 kernels -// went MX-only: 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). +// 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 diff --git a/csrc/fused_mqa_topk/dense_scan.cuh b/csrc/fused_mqa_topk/dense_scan.cuh index f0b2e2279f53..395cb84f183e 100644 --- a/csrc/fused_mqa_topk/dense_scan.cuh +++ b/csrc/fused_mqa_topk/dense_scan.cuh @@ -1,5 +1,10 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025 DeepSeek +// 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 diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu index 194e02fa2484..56cf55a118ef 100644 --- a/csrc/fused_mqa_topk/fused_mqa_topk.cu +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -64,6 +64,8 @@ int sm_count(int device) { 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) { diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc index 40c3ca3d86a4..49d69330e882 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -1,3 +1,6 @@ +// 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). diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh index afacbc54370b..48a806f11daf 100644 --- a/csrc/fused_mqa_topk/persistent_kernel.cuh +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -1,4 +1,6 @@ -// SPDX-License-Identifier: MIT +// 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 diff --git a/csrc/fused_mqa_topk/rescan.cuh b/csrc/fused_mqa_topk/rescan.cuh index 0937d1997a8f..43ad8f5eafa2 100644 --- a/csrc/fused_mqa_topk/rescan.cuh +++ b/csrc/fused_mqa_topk/rescan.cuh @@ -1,4 +1,6 @@ -// SPDX-License-Identifier: MIT +// 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 diff --git a/csrc/fused_mqa_topk/select.cuh b/csrc/fused_mqa_topk/select.cuh index 43322d68dc23..c471f1fd3f80 100644 --- a/csrc/fused_mqa_topk/select.cuh +++ b/csrc/fused_mqa_topk/select.cuh @@ -1,4 +1,6 @@ -// SPDX-License-Identifier: MIT +// 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, the row's top-k: 512, 1024 or 2048, is a runtime value). // From 38139bf32241419d156296c766f72a88804a46f7 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Mon, 26 Oct 2026 18:14:22 +0000 Subject: [PATCH 17/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/candidates.cuh | 1 + csrc/fused_mqa_topk/common.cuh | 19 +- csrc/fused_mqa_topk/fused_mqa_topk.cu | 83 +++++---- csrc/fused_mqa_topk/persistent_epoch.inc | 169 +++++++++--------- csrc/fused_mqa_topk/persistent_kernel.cuh | 100 ++++++----- csrc/fused_mqa_topk/rescan.cuh | 132 +++++++------- csrc/fused_mqa_topk/select.cuh | 112 ++++++------ .../kernels/attention/test_fused_mqa_topk.py | 3 +- .../kernels/attention/dsa/fused_mqa_topk.py | 21 ++- 9 files changed, 350 insertions(+), 290 deletions(-) diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh index 409176bdb66a..30afdec28efd 100644 --- a/csrc/fused_mqa_topk/candidates.cuh +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -31,6 +31,7 @@ #include "common.cuh" #include "constants.h" + namespace fused_mqa_topk { using CandidateValue = uint16_t; constexpr uint32_t kCandidateIndexBits = 20; diff --git a/csrc/fused_mqa_topk/common.cuh b/csrc/fused_mqa_topk/common.cuh index 4e52408a07fa..860bed7c3dcc 100644 --- a/csrc/fused_mqa_topk/common.cuh +++ b/csrc/fused_mqa_topk/common.cuh @@ -17,7 +17,10 @@ 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; +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. @@ -48,12 +51,15 @@ __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(0xffffffffu, v, d); + const uint32_t u = __shfl_up_sync(kFullMask, v, d); if (lane >= d) v += u; } return v; @@ -65,6 +71,13 @@ __device__ __forceinline__ void mbar_inval(const void* 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) { @@ -91,10 +104,10 @@ __device__ __forceinline__ float half_bin_floor(uint32_t b) { } // Weighted head reduction of one row's 32 head scores s_h in two FFMA2 chains -// (adapted from DeepGEMM's sm100_fp8_mqa_logits.cuh at 88965b078186) // 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 = {}, diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu index 56cf55a118ef..652453494fbe 100644 --- a/csrc/fused_mqa_topk/fused_mqa_topk.cu +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -28,7 +28,17 @@ 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)); @@ -96,10 +106,10 @@ CUtensorMap make_2d(void* ptr, CUtensorMapDataType dt, int elem_size, // 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 = fused_mqa_topk::kHeadDim; + 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, fused_mqa_topk::kTileKeys, 1}; + 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, @@ -110,18 +120,11 @@ CUtensorMap make_strided_kv(void* ptr, uint32_t groups, uint32_t stride) { return tm; } -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; -namespace persist = fused_mqa_topk::persist; - // -------- 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>); @@ -129,8 +132,9 @@ using PersistentKernel = // 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 <= 131072 ? persist::fused_mqa_topk_fp8_prefill_kernel<32, TopK> + 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>; @@ -149,12 +153,10 @@ struct LaunchConfig { // 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. -inline cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) { - // The recovery rescan reuses the same dynamic shared memory with the rescan's - // layout. - constexpr int kFallbackSmem = - int(sizeof(fused_mqa_topk::dense::SharedStorage) + - sizeof(persist::RescanShared)); +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))); @@ -175,8 +177,7 @@ inline cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) { cudaFuncGetAttributes(&attr, reinterpret_cast(kernel)); e != cudaSuccess) return e; - const int bytes = - std::max(kFallbackSmem, max_smem - int(attr.sharedSizeBytes)); + 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), @@ -190,8 +191,8 @@ inline cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) { // Picks the kernel instantiation, its shared memory and the gate sample for // `keys` keys and top-`topk` rows on `device` (the current device). -inline cudaError_t launch_config(int device, uint32_t keys, uint32_t topk, - LaunchConfig& c) { +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 @@ -208,7 +209,7 @@ inline cudaError_t launch_config(int device, uint32_t keys, uint32_t topk, {{1536, 48}, {2048, 48}}, // top-1024 {{3072, 64}, {4096, 64}}, // top-2048 }; - const bool long_ctx = keys > 131072; + 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; @@ -249,19 +250,23 @@ void check_aligned(const torch::stable::Tensor& t, const char* name) { } // namespace -// q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32, -// weights [rows,32] fp32, starts/ends [rows] int32 (visible key range -// [start,end) of each row), 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. 0 < -// rows < 2^26, 0 < keys <= 1048576 (the last 8-row block may be partial). q, k, -// scales, weights and out must be 16-byte aligned (TMA / vector stores). -// 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. -// The output is exact up to ties: the gated path ranks fp24 score codes (the -// low 8 mantissa bits dropped), the exact rescan ranks fp32 scores. +// 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 diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc index 49d69330e882..b47323982be5 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -15,7 +15,7 @@ if (warp_idx == 9 && cute::elect_one_sync()) { full_kv[i].init(kFullKvArrivals); empty_kv[i].init(kMathThreads); } - for (uint32_t i = 0; i < BarIdx::NT; ++i) { + for (uint32_t i = 0; i < BarIdx::kTmemSlots; ++i) { full_tmem[i].init(1); empty_tmem[i].init(kMathWgThreads); } @@ -25,8 +25,8 @@ if (warp_idx == 9 && cute::elect_one_sync()) { claim_full[i].init(32); claim_empty[i].init(kClaimConsumers * 32); } - sel_done.init(8); // one lane per math warp - for (uint32_t i = 0; i < 16; ++i) { + 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); } @@ -52,7 +52,7 @@ if (warp_idx == 8) { uint32_t b = 0; if (lane == 0) b = atomicOr(&pause_req, 0u) ? kPause : atomicAdd(counter, 1u); - b = __shfl_sync(0xffffffffu, b, 0); + b = __shfl_sync(kFullMask, b, 0); ClaimEntry& e = entries[k & 1]; if (b >= nblocks) { if (lane == 0) { @@ -77,19 +77,19 @@ if (warp_idx == 8) { e.ends[lane] = en; } const bool all_short = - __ballot_sync(0xffffffffu, lane < kBlockRows && en - st > topk) == 0; + __ballot_sync(kFullMask, lane < kBlockRows && en - st > topk) == 0; const bool gated = - __ballot_sync(0xffffffffu, + __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(0xffffffffu, st, src), - hi = __shfl_sync(0xffffffffu, en, src); + 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(0xffffffffu, lo, d)); - hi = max(hi, __shfl_xor_sync(0xffffffffu, hi, d)); + 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; @@ -120,14 +120,14 @@ if (warp_idx == 8) { ++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 = 4; p < 8; ++p) + 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 % 128); + source % kMathWgThreads); }); } cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); @@ -149,9 +149,11 @@ if (warp_idx == 8) { if (lane == 0 && n_blk > 0) sel_done.wait((n_blk - 1) & 1); __syncwarp(); const uint32_t center = sample_stride / 2; - float pf[8]; + // 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 < 8; ++j) { + 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; @@ -161,11 +163,12 @@ if (warp_idx == 8) { // 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(0xffffffffu, s0, 0); + const uint32_t s = __shfl_sync(kFullMask, s0, 0); // Every writing lane acquires and releases the stage itself. - empty_kv[s].wait(__shfl_sync(0xffffffffu, ph, 0) ^ 1); + empty_kv[s].wait(__shfl_sync(kFullMask, ph, 0) ^ 1); #pragma unroll - for (uint32_t j = 0; j < 8; ++j) sm.sf_kv[s][j * 32 + lane] = pf[j]; + 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) { @@ -175,7 +178,7 @@ if (warp_idx == 8) { } full_kv[s].arrive(); #pragma unroll - for (uint32_t j = 0; j < 8; ++j) { + 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; @@ -203,7 +206,7 @@ if (warp_idx == 8) { cutlass::arch::warpgroup_reg_dealloc(); RingPipeline<1> q_pipe; RingPipeline kv_pipe; - RingPipeline tmem_pipe; + RingPipeline tmem_pipe; auto instr_desc = cute::UMMA::make_instr_desc= 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 % 128); + source % kMathWgThreads); }); } cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); @@ -294,10 +297,10 @@ if (warp_idx == 8) { // -------- math warps 0-7 -------- cutlass::arch::warpgroup_reg_alloc(); const uint32_t math_tid = threadIdx.x; - const uint32_t query_base = (warp_idx / 4) * 4; + const uint32_t query_base = (warp_idx / 4) * kWgRows; // 4 warps per WG RingPipeline<1> q_pipe; RingPipeline kv_pipe; - RingPipeline tmem_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); @@ -329,20 +332,20 @@ if (warp_idx == 8) { 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[4][32]; - float row_s[4]; // pow2 row scale of the linear term - uint32_t seq_k_start[4], seq_k_end[4]; + 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 < 4; ++i) { + 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 < 32; j += 4) { + for (uint32_t j = 0; j < kNumHeads; j += 4) { const float4 w4 = *reinterpret_cast( - sm.weights + (query_base + i) * 32 + j); + 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; @@ -352,13 +355,13 @@ if (warp_idx == 8) { { // Weights that are subnormal, or become subnormal or infinite once // scaled, take the exact recovery path. - const float w = sm.weights[warp_idx * 32 + lane]; + 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(0xffffffffu, bad_w); + const uint32_t bad = __ballot_sync(kFullMask, bad_w); if (lane == 0 && bad) atomicOr(&unsafe_rows, 1u << warp_idx); } math_wg_sync(); @@ -386,26 +389,21 @@ if (warp_idx == 8) { // 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 < 4; ++i) { - float acc[32]; + 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 * 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 == 3) { + if (i == kWgRows - 1) { ptx::tcgen05_before_thread_sync(); empty_tmem[tmem_stage].arrive(); } - // The linear term seeds the sums (terms 2, 3 in units of s * 2^-8). - const float bq = - reduce_heads( - acc, weights[i], make_float2(lin[0], lin[1]), - __fmul2_rn(make_float2(lin[2], lin[3]), - make_float2(kLinLoScale, kLinLoScale))) * - (row_s[i] * scale_kv); + 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 (valid) { @@ -422,19 +420,16 @@ if (warp_idx == 8) { 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 = - *reinterpret_cast(&e.vblock) * kBlockRows; + 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[4]; + float gate_score[kWgRows]; #pragma unroll - for (uint32_t i = 0; i < 4; ++i) { + for (uint32_t i = 0; i < kWgRows; ++i) { gate_score[i] = gate_s[query_base + i]; - if ((*reinterpret_cast(&unsafe_rows) >> - (query_base + i)) & - 1u) { + 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; @@ -446,12 +441,12 @@ if (warp_idx == 8) { // exact). gate_score[i] *= 1.0f / row_s[i]; } - uint32_t lane_counts[4] = {}; + 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) * 8 + warp_idx].wait( + h_empty[(produce_interval & 1) * kMathWarps + warp_idx].wait( ((produce_interval >> 1) & 1) ^ 1); __syncwarp(); } @@ -470,18 +465,20 @@ if (warp_idx == 8) { cutlass::arch::fence_view_async_shared(); empty_kv[kv_stage].arrive(); } - float accum_pair[2][32]; + float accum_pair[2][kNumHeads]; float lin_pair[2][4]; - float bp4[4]; // row scores of the slot, in units of the row scale + float + bp4[kWgRows]; // row scores of the slot, in units of the row scale #pragma unroll - for (uint32_t pair = 0; pair < 4; pair += 2) { + for (uint32_t pair = 0; pair < kWgRows; pair += 2) { const uint32_t slot_base = tmem_stage * kSlotCols; - ptx::tmem_load_32dp32b<64>(slot_base + heads_col + pair * 32, - reinterpret_cast(accum_pair)); + 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 >= 4) { + if (pair + 2 >= kWgRows) { ptx::tcgen05_before_thread_sync(); empty_tmem[tmem_stage].arrive(); } @@ -489,22 +486,18 @@ if (warp_idx == 8) { // 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( - accum_pair[r], weights[pair + r], - make_float2(lin_pair[r][0], lin_pair[r][1]), - __fmul2_rn(make_float2(lin_pair[r][2], lin_pair[r][3]), - make_float2(kLinLoScale, kLinLoScale))) * - scale_kv; + 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(0xffffffffu, any_pass)) [[unlikely]] { + if (__any_sync(kFullMask, any_pass)) [[unlikely]] { #pragma unroll - for (uint32_t i = 0; i < 4; ++i) { + 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]) { @@ -532,18 +525,17 @@ if (warp_idx == 8) { } } } - if ((((kv_split_idx + 1) & (Interval - 1)) == 0 || - kv_split_idx + 1 == splits) && + if (((kv_split_idx + 1) & (Interval - 1)) == 0 && kv_split_idx + 1 != splits) { #pragma unroll - for (uint32_t i = 0; i < 4; ++i) { + 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) * 8 + warp_idx].arrive(); + h_full[(produce_interval & 1) * kMathWarps + warp_idx].arrive(); ++produce_interval; flush_base = kv_split_idx + 1; } @@ -554,16 +546,17 @@ if (warp_idx == 8) { // cycle the barriers. const uint32_t bank = produce_interval & 1, rbase = scan_row_base + query_base; - static_assert(sizeof(hdata.records[0][0]) == 4 * 8 * 32 * 4, + 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 < 4; ++row) lane_counts[row] = 0; + for (uint32_t row = 0; row < kWgRows; ++row) lane_counts[row] = 0; __syncwarp(); - h_full[bank * 8 + warp_idx].arrive(); + h_full[bank * kMathWarps + warp_idx].arrive(); ++produce_interval; } cutlass::arch::fence_view_async_shared(); @@ -572,8 +565,7 @@ if (warp_idx == 8) { cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).sync(); // ---- selection ---- // Re-read as well (not kept live across the scan). - const uint32_t sel_vblock = - *reinterpret_cast(&e.vblock), + const uint32_t sel_vblock = volatile_load(e.vblock), sel_row_base = sel_vblock * kBlockRows; __threadfence_block(); cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); @@ -590,17 +582,22 @@ if (warp_idx == 8) { 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[4], row_offset[8] - row_offset[4]); + const uint32_t half = max(row_offset[kWgRows], + row_offset[kBlockRows] - row_offset[kWgRows]); const uint32_t staged = - recover - ? 0u - : (8u * selection::kBins * 4u + row_offset[8] * 4u <= phase_bytes - ? 1u - : (4u * selection::kBins * 4u + half * 4u <= phase_bytes - ? 2u - : 0u)); + 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); @@ -623,8 +620,8 @@ if (warp_idx == 8) { }; if (!recover) { const int64_t row_delta = slot_row_delta(sel_entry.slot, sel_row_base); - const uint32_t per_round = sel_staged == 1 ? 8u : 4u; - for (uint32_t round = 0; round < 8 / per_round; ++round) { + 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; @@ -701,7 +698,7 @@ if (defer_head != kTerm) { cur = uint32_t(__ldcg(defer_next + vblock)); __syncthreads(); } - __syncthreads(); + // The loop ran at least once and ended with a barrier. if (threadIdx.x == 0) { defer_head = kTerm; pause_req = 0; diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh index 48a806f11daf..43eec73856b0 100644 --- a/csrc/fused_mqa_topk/persistent_kernel.cuh +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -62,8 +62,8 @@ constexpr uint32_t kLinRows = 16; 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, kept from the experiments); the 3 slots take -// 480 of the 512 allocated columns. +// 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); @@ -101,14 +101,14 @@ __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 < 8; ++r) { + 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 < 32; ++h) { - const uint32_t off = swizzle128_offset(r * 32 + h, lane * 4u); + 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 * 32 + h) * 0.5f; + 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( @@ -124,7 +124,7 @@ __device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, fmaxf(fabsf(acc[2]), fabsf(acc[3]))); #pragma unroll for (uint32_t o = 16; o > 0; o >>= 1) - m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + 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); @@ -134,7 +134,7 @@ __device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, float res[4]; #pragma unroll for (uint32_t dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s; - uint8_t* lin_base = r < 4 ? sm.lin_lo : sm.lin_hi; + 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( @@ -160,7 +160,8 @@ __device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, res[2] *= 1.0f / kLinLoScale; res[3] *= 1.0f / kLinLoScale; } - const uint32_t off = swizzle128_offset((r % 4u) * 4u + tt, lane * 4u); + 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)); @@ -169,6 +170,16 @@ __device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, } } +// 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 // (half_bin<5>: kSampleBins bins, padded) and finish_sample_histogram picks the // gate from them. @@ -188,15 +199,17 @@ __device__ __noinline__ inline void finish_sample_histogram( __syncwarp(); return; } - // Lane l owns the 64 bins of 32-bin chunks 63-2l and 62-2l (bins descending - // with the lane). - const uint32_t* const h0 = hist + padded_bin((63 - 2 * lane) * 32); - const uint32_t* const h1 = hist + padded_bin((62 - 2 * lane) * 32); + // Lane l owns the 64 bins of 32-bin chunks kChunks-1-2l and kChunks-2-2l + // (bins descending with the lane). + constexpr uint32_t kChunks = kSampleBins / 32; + static_assert(kChunks == 2 * 32, "two 32-bin chunks per lane"); + const uint32_t* const h0 = hist + padded_bin((kChunks - 1 - 2 * lane) * 32); + const uint32_t* const h1 = hist + padded_bin((kChunks - 2 - 2 * lane) * 32); uint32_t sum = 0; #pragma unroll 8 for (uint32_t j = 0; j < 32; ++j) sum += h0[j] + h1[j]; const uint32_t incl = warp_inclusive_scan(sum, lane); - const uint32_t tot = __shfl_sync(0xffffffffu, incl, 31); + const uint32_t tot = __shfl_sync(kFullMask, incl, 31); if (tot == 0) { if (lane == 0) *gate = -CUDART_INF_F; __syncwarp(); @@ -204,20 +217,20 @@ __device__ __noinline__ inline void finish_sample_histogram( } const uint32_t rank = max(1u, min(tot, (tot * target + visible - 1) / visible)); - const uint32_t w = __ffs(__ballot_sync(0xffffffffu, incl >= rank)) - 1; - uint32_t above = __shfl_sync(0xffffffffu, incl - sum, w), chosen = 0; + 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 < 2; ++cc) { - const uint32_t c = 63 - 2 * w - cc; + const uint32_t c = kChunks - 1 - 2 * w - cc; // Lane 0 = highest bin of the chunk. const uint32_t v = warp_inclusive_scan(hist[padded_bin(c * 32 + 31 - lane)], lane); - const uint32_t mask = __ballot_sync(0xffffffffu, above + v >= rank); + const uint32_t mask = __ballot_sync(kFullMask, above + v >= rank); if (mask) { chosen = c * 32 + 31 - (__ffs(mask) - 1); break; } - above += __shfl_sync(0xffffffffu, v, 31); + above += __shfl_sync(kFullMask, v, 31); } if (lane == 0) *gate = half_bin_floor<5>(chosen); __syncwarp(); @@ -233,9 +246,9 @@ static_assert(32 * kRingSlots < 65536, // 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][8][4][kRingSlots][32]; - uint32_t num_records[2][8][4][32]; - uint32_t bases[2][8]; + 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 @@ -244,7 +257,7 @@ struct ClaimHeader { uint32_t vblock, all_short, kv_base, splits, slot, samples; }; struct ClaimEntry : ClaimHeader { - uint32_t starts[8], ends[8]; + uint32_t starts[kBlockRows], ends[kBlockRows]; }; // Claim k in a producer role (warps 9-11): waits for the entry, copies its @@ -266,15 +279,16 @@ __device__ __forceinline__ bool consume_claim(uint32_t k, TBar* claim_full, struct BarIdx { // TMEM slots, shared by the two math warpgroups (WG0 even, WG1 odd // positions). - static constexpr uint32_t NT = 3; + 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 + NT, - claim_full = empty_tmem + NT, + 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 + 16, full_lin = h_empty + 16, + h_empty = h_full + 2 * kMathWarps, + full_lin = h_empty + 2 * kMathWarps, total = full_lin + 1; }; @@ -341,7 +355,7 @@ __device__ __forceinline__ void make_relative(int* output, uint32_t topk, uint32_t t, uint32_t stride, Skip skip) { const uint32_t quads = topk / 4; - for (uint32_t i = t; i < 8 * quads; i += stride) { + 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]); @@ -357,11 +371,13 @@ __device__ __forceinline__ void make_relative(int* output, uint32_t topk, } // Stores one warp's ring records of a handoff interval in its rows' candidate -// slabs. recs = [4 rows][8 records][32 lanes]; lane l holds n[row] records of -// each row. Each row's slots are reserved with one atomic on +// 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], @@ -372,7 +388,7 @@ __device__ __forceinline__ void store_ring_records(const uint32_t* recs, uint32_t my_total = 0; #pragma unroll for (uint32_t row = 0; row < 4; ++row) { - const uint32_t total = __reduce_add_sync(0xffffffffu, n[row]); + const uint32_t total = __reduce_add_sync(kFullMask, n[row]); if (lane == row) my_total = total; } uint32_t my_start = 0; @@ -382,8 +398,8 @@ __device__ __forceinline__ void store_ring_records(const uint32_t* recs, 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(0xffffffffu, p01, dd), - b = __shfl_up_sync(0xffffffffu, p23, dd); + const uint32_t a = __shfl_up_sync(kFullMask, p01, dd), + b = __shfl_up_sync(kFullMask, p23, dd); if (lane >= dd) { p01 += a; p23 += b; @@ -392,7 +408,7 @@ __device__ __forceinline__ void store_ring_records(const uint32_t* recs, 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(0xffffffffu, my_start, 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]; @@ -424,10 +440,10 @@ __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 * 8 + producer].wait((interval >> 1) & 1); + hfull[bank * kMathWarps + producer].wait((interval >> 1) & 1); __syncwarp(); if (skip) { // self-drained by the producer - hempty[bank * 8 + producer].arrive(); + hempty[bank * kMathWarps + producer].arrive(); return; } const uint32_t base = d.bases[bank][producer]; @@ -441,7 +457,7 @@ __device__ __forceinline__ void drain_handoff( store(row, slot, record, base, producer * 32 + lane); }); __syncwarp(); - hempty[bank * 8 + producer].arrive(); + hempty[bank * kMathWarps + producer].arrive(); } // Interval: scan tiles per candidate handoff to the drain warps. @@ -461,7 +477,8 @@ __launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel( 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::NT * kSlotCols <= 512, "TMEM slots exceed 512 columns"); + 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, @@ -475,8 +492,9 @@ __launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel( __shared__ alignas(8) uint64_t bar_words[BarIdx::total]; __shared__ uint32_t tmem_ptr; __shared__ ClaimEntry entries[2]; - __shared__ float gate_s[8]; - __shared__ uint32_t sel_recover, sel_staged, sel_short, row_offset[9]; + __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); @@ -517,7 +535,7 @@ __launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel( // 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][8]; + __shared__ uint32_t row_bounds[2][kBlockRows]; bool more = false; { constexpr uint32_t epoch = 0; diff --git a/csrc/fused_mqa_topk/rescan.cuh b/csrc/fused_mqa_topk/rescan.cuh index 43ad8f5eafa2..80d9dc57370f 100644 --- a/csrc/fused_mqa_topk/rescan.cuh +++ b/csrc/fused_mqa_topk/rescan.cuh @@ -30,61 +30,68 @@ struct RescanBuffers { // levels of compact use 256 bins. constexpr uint32_t kSegmentBins = 1024, kRadixBins = 256; struct alignas(16) RescanShared { - float ring[2][8][kSegmentKeys]; - uint32_t hist[8][padded_bins(kSegmentBins)], gates[8], kept[8]; - uint32_t work[padded_bins(kRadixBins)], prefix, rank, emit_counts[8]; + 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 < 8) { + 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 K, + RescanBuffers buffers, uint32_t topk, uint32_t local, uint32_t row, - uint32_t tid, bool final = false) { - constexpr uint32_t Full = 0xffffffffu, kThreads = kMathThreads, - Items = kRetainedPerRow / kThreads; + 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[Items]; - int indices[Items]; + float values[kItems]; + int indices[kItems]; #pragma unroll - for (uint32_t i = 0; i < Items; ++i) { + 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; } - // mode 0: values >= gate; 1: codes above threshold; 2: codes equal to it. - auto emit = [&](float gate, uint32_t threshold, int mode, + // 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 == 0 ? v >= gate - : (mode == 1 ? ordered_code(v) > threshold - : ordered_code(v) == threshold)); + (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 < Items; ++i) total += passes(values[i]); - uint32_t warp_total = __reduce_add_sync(Full, total); + 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 < 8; ++w) all += shared.emit_counts[w]; + for (uint32_t w = 0; w < kMathWarps; ++w) all += shared.emit_counts[w]; #pragma unroll - for (uint32_t i = 0; i < Items; ++i) { + for (uint32_t i = 0; i < kItems; ++i) { const float v = values[i]; const bool pass = passes(v); - uint32_t mask = __ballot_sync(Full, pass), + uint32_t mask = __ballot_sync(kFullMask, pass), pos = start + __popc(mask & ((1u << lane) - 1)); - if (pass && (mode != 2 || pos < K)) { + if (pass && (mode != Emit::kTies || pos < topk)) { buffers.retained[base + pos] = v; buffers.retained_indices[base + pos] = indices[i]; } @@ -93,7 +100,7 @@ __device__ inline void compact_retained(RescanShared& shared, math_wg_sync(); return all; }; - if (!final) { + if (!compact_to_k) { if (tid == 0) { shared.rank = 0; } @@ -101,14 +108,14 @@ __device__ inline void compact_retained(RescanShared& shared, float gate = __uint_as_float(shared.gates[local]); uint32_t kept = 0; #pragma unroll - for (uint32_t i = 0; i < Items; ++i) + for (uint32_t i = 0; i < kItems; ++i) kept += isfinite(values[i]) && values[i] >= gate; - kept = __reduce_add_sync(Full, kept); + kept = __reduce_add_sync(kFullMask, kept); if (lane == 0) atomicAdd(&shared.rank, kept); math_wg_sync(); kept = shared.rank; - if (kept >= K && kept <= kRetainedPerRow - kSegmentKeys) { - emit(gate, 0, 0, 0); + 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(); @@ -117,7 +124,7 @@ __device__ inline void compact_retained(RescanShared& shared, } if (tid == 0) { shared.prefix = 0; - shared.rank = K; + shared.rank = topk; } math_wg_sync(); for (int shift = 24; shift >= 0; shift -= 8) { @@ -126,11 +133,11 @@ __device__ inline void compact_retained(RescanShared& shared, math_wg_sync(); uint32_t prefix = shared.prefix; #pragma unroll - for (uint32_t i = 0; i < Items; ++i) { + 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(Full, valid); + uint32_t mask = __ballot_sync(kFullMask, valid); if (valid) { uint32_t bin = (code >> shift) & 255u, peers = __match_any_sync(mask, bin); @@ -148,7 +155,7 @@ __device__ inline void compact_retained(RescanShared& shared, const uint32_t inclusive = warp_inclusive_scan(sum, lane); uint32_t before = inclusive - sum, rank = shared.rank, crossing = - __ballot_sync(Full, before < rank && inclusive >= rank), + __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) { @@ -160,8 +167,8 @@ __device__ inline void compact_retained(RescanShared& shared, above += shared.work[padded_bin(bin)]; } } - chosen = __shfl_sync(Full, chosen, chosen_lane); - above = __shfl_sync(Full, above, chosen_lane); + 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; @@ -171,41 +178,42 @@ __device__ inline void compact_retained(RescanShared& shared, math_wg_sync(); } uint32_t threshold = shared.prefix; - uint32_t greater = emit(0, threshold, 1, 0); - emit(0, threshold, 2, greater); + uint32_t greater = emit(0, threshold, Emit::kAboveCode, 0); + emit(0, threshold, Emit::kTies, greater); if (tid == 0) { - shared.kept[local] = K; + 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 K, + RescanBuffers buffers, uint32_t topk, uint32_t local, uint32_t row, uint32_t segment, uint32_t first, - uint32_t n, uint32_t tid) { - constexpr uint32_t Full = 0xffffffffu, Items = kSegmentKeys / 32; - uint32_t lane = tid, bank = segment & 1; - float values[Items], gate = __uint_as_float(shared.gates[local]); + 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 < Items; ++i) { - uint32_t j = tid + i * 32; + 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(Full, 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(Full, count); + count = __reduce_add_sync(kFullMask, count); __syncwarp(); if (count) { uint32_t sum = 0; @@ -214,31 +222,33 @@ __device__ inline void merge_segment(RescanShared& shared, 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(Full, before < K && inclusive >= K); + mask = + __ballot_sync(kFullMask, before < topk && inclusive >= topk); if (mask) { uint32_t chosen = __ffs(mask) - 1, group = 31 - chosen, - rank = K - __shfl_sync(Full, before, 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(Full, before < rank && inclusive >= rank)) - 1; - bin = __shfl_sync(Full, bin, 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 < Items; ++i) { + for (uint32_t i = 0; i < kItems; ++i) { float v = values[i]; bool pass = isfinite(v) && v >= gate; - uint32_t mask = __ballot_sync(Full, pass), + 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 + tid + i * 32; + buffers.retained_indices[base + pos] = first + lane + i * 32; } written += __popc(mask); } @@ -253,19 +263,19 @@ __device__ inline void merge_segment(RescanShared& shared, // written to the output. struct RescanPolicy { static constexpr float kWeightScale = 0.5f; - static constexpr uint32_t kSplitsPerSegment = kSegmentKeys / 256; + 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[4], seq_k_end[4]; - float gate[4]; + 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 < 4; ++i) { + 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); @@ -275,7 +285,7 @@ struct RescanPolicy { CUTLASS_DEVICE void begin_split(uint32_t split) { if (split % kSplitsPerSegment == 0) { __syncwarp(); - for (uint32_t i = 0; i < 4; ++i) + for (uint32_t i = 0; i < kWgRows; ++i) gate[i] = __uint_as_float(atomicAdd(shared.gates + query_base + i, 0u)); } } @@ -284,7 +294,7 @@ struct RescanPolicy { 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) * 256 + kv_lane] = + [(split % kSplitsPerSegment) * kTileKeys + kv_lane] = valid ? bq : -CUDART_INF_F; } CUTLASS_DEVICE void end_split(uint32_t split, uint32_t num_splits, @@ -293,17 +303,17 @@ struct RescanPolicy { __threadfence_block(); math_wg_sync(); const uint32_t tid = threadIdx.x; - for (uint32_t local = 0; local < 8; ++local) + 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) * 256, tid % 32); + (split % kSplitsPerSegment + 1) * kTileKeys, tid % 32); math_wg_sync(); if (split + 1 == num_splits) - for (uint32_t l = 0; l < 8; ++l) + 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(); diff --git a/csrc/fused_mqa_topk/select.cuh b/csrc/fused_mqa_topk/select.cuh index c471f1fd3f80..759ec7fce46a 100644 --- a/csrc/fused_mqa_topk/select.cuh +++ b/csrc/fused_mqa_topk/select.cuh @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // // Per-row top-K selection of the persistent kernel's candidates, one warp per -// row (K, the row's top-k: 512, 1024 or 2048, is a runtime value). +// 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 @@ -23,7 +23,7 @@ #include "constants.h" namespace fused_mqa_topk::selection { -constexpr uint32_t Full = 0xffffffffu, kBinsLog2 = 10, kBins = 1u << kBinsLog2; +constexpr uint32_t kBinsLog2 = 10, kBins = 1u << kBinsLog2; // Elements each lane keeps in flight. constexpr uint32_t kUnroll = 8; @@ -42,7 +42,7 @@ __device__ __forceinline__ uint32_t find_bin_n(const uint32_t* hist, const uint32_t inclusive = warp_inclusive_scan(sum, lane); const uint32_t before = inclusive - sum; const uint32_t winner = - __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; + __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) { @@ -54,14 +54,14 @@ __device__ __forceinline__ uint32_t find_bin_n(const uint32_t* hist, acc += h; } } - above = __shfl_sync(Full, acc, winner); - return __shfl_sync(Full, chosen, winner); + above = __shfl_sync(kFullMask, acc, winner); + return __shfl_sync(kFullMask, chosen, winner); } -// Requires count >= K (callers recover rows with fewer candidates). +// 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 K, uint32_t* hist) { + 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; @@ -87,13 +87,13 @@ __device__ inline void select_row(uint32_t* codes, uint32_t count, lo = min(lo, c); hi = max(hi, c); } - lo = __reduce_min_sync(Full, lo); - hi = __reduce_max_sync(Full, hi); + 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 = K; + uint32_t width = hi - lo + 1, rank = topk; #pragma unroll 1 while (width > 1) { const uint32_t spread = (width - 1) >> kBinsLog2; @@ -129,7 +129,7 @@ __device__ inline void select_row(uint32_t* codes, uint32_t count, uint32_t pos = 0; #pragma unroll 1 for (uint32_t tie = 0; tie < 2; ++tie) { - for (uint32_t j0 = 0; j0 < count && pos < K; j0 += 32 * kUnroll) { + for (uint32_t j0 = 0; j0 < count && pos < topk; j0 += 32 * kUnroll) { bool pass[kUnroll]; int32_t p[kUnroll]; #pragma unroll @@ -141,14 +141,15 @@ __device__ inline void select_row(uint32_t* codes, uint32_t count, } #pragma unroll for (uint32_t u = 0; u < kUnroll; ++u) { - const uint32_t mask = __ballot_sync(Full, pass[u]), + const uint32_t mask = __ballot_sync(kFullMask, pass[u]), at = pos + __popc(mask & lt); - if (pass[u] && at < K) out[at] = candidate_decode_index(p[u]); + 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 @@ -166,7 +167,8 @@ static_assert(kCandidatesPerRow <= 65536, "candidate ids must fit 16 bits"); static_assert(768 <= kBins && kMidCap / 2 + kMidBins <= kBins, "selection scratch exceeds kBins words"); -// r-th largest (1-based) of the 512 register samples for two ranks at once +// 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 @@ -199,6 +201,28 @@ __device__ __forceinline__ void radix_bracket( __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). @@ -226,7 +250,7 @@ __device__ __forceinline__ uint32_t find_bin_v(const uint32_t* hist, const uint32_t inclusive = warp_inclusive_scan(sum, lane); const uint32_t before = inclusive - sum; const uint32_t winner = - __ffs(__ballot_sync(Full, before < rank && inclusive >= rank)) - 1; + __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; @@ -241,14 +265,14 @@ __device__ __forceinline__ uint32_t find_bin_v(const uint32_t* hist, } acc += x; } - above = __shfl_sync(Full, acc_at, winner); - return __shfl_sync(Full, chosen, winner); + 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 spilled into the scan loop: 290K -> -// 860K cycles/block in a first try). +// 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 @@ -259,9 +283,9 @@ __device__ __forceinline__ uint32_t find_bin_v(const uint32_t* hist, __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 K, uint32_t* hist) { + int32_t* __restrict__ out, uint32_t topk, uint32_t* hist) { const uint32_t lane = threadIdx.x & 31; - if (count <= K) return false; + if (count <= topk) return false; const uint32_t chunks = (count + 7) / 8; constexpr uint32_t kChunksInFlight = 2; const uint4* v4p = reinterpret_cast(values); @@ -309,33 +333,20 @@ __device__ __noinline__ inline bool select_staged_row( // 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) { - uint32_t smp[kSamplesPerLane]; -#pragma unroll - for (uint32_t u = 0; u < kSamplesPerLane; ++u) - smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 10)] >> 8; - const uint32_t expect = (K * 512 + count / 2) / count; - const uint32_t r_hi = expect > kMargin ? expect - kMargin : 1, - r_lo = min(expect + kMargin, 512u); - radix_bracket(smp, r_hi, r_lo, hist, h_hi, h_lo); - } else { - uint32_t smp[2 * kSamplesPerLane]; -#pragma unroll - for (uint32_t u = 0; u < 2 * kSamplesPerLane; ++u) - smp[u] = codes[(((u * 32 + lane) * 2 + 1) * count >> 11)] >> 8; - const uint32_t expect = (K * 1024 + count / 2) / count; - const uint32_t r_hi = expect > kWideMargin ? expect - kWideMargin : 1, - r_lo = min(expect + kWideMargin, 1024u); - radix_bracket<2 * kSamplesPerLane>(smp, r_hi, r_lo, hist, 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 - // (per-lane runs were 4x slower: 32 lines/STG). Writes are bounded; the + // 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; @@ -359,9 +370,10 @@ __device__ __noinline__ inline bool select_staged_row( 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(Full, a), mm = __ballot_sync(Full, m); + 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 < K) out[at] = candidate_decode_index(int32_t(p[k])); + 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); @@ -369,7 +381,7 @@ __device__ __noinline__ inline bool select_staged_row( } } __syncwarp(); - if (above >= K || above + mid < K || mid > kMidCap) return false; + 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]; @@ -379,7 +391,7 @@ __device__ __noinline__ inline bool select_staged_row( mc[u] = 0; if (i < mid) mc[u] = codes[midbuf[i]]; } - uint32_t rank = K - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; + 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; @@ -406,7 +418,7 @@ __device__ __noinline__ inline bool select_staged_row( { // Ballot placement (coalesced stores): greater codes fill [above, K - // rank), ties [K - rank, K). - uint32_t pg = above, pt = K - rank; + uint32_t pg = above, pt = topk - rank; #pragma unroll for (uint32_t h = 0; h < 2; ++h) { if (h * 16 * 32 >= mid) break; @@ -424,11 +436,11 @@ __device__ __noinline__ inline bool select_staged_row( 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(Full, g), - mt = __ballot_sync(Full, tq); + 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 < K) out[at] = candidate_decode_index(pk[j]); + if (tq && at < topk) out[at] = candidate_decode_index(pk[j]); pg += __popc(mg); pt += __popc(mt); } diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index a950b3c55200..4f6b629046cc 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -243,7 +243,8 @@ def test_rescan_scores_only_its_block(): 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, relative=True) and must stay exact. + 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 diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 1fa08abeb933..2f25a1e87222 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -84,13 +84,16 @@ def use_fused_mqa_topk( full_context: bool, candidate_selection: bool = False, ) -> bool: - """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. + """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. """ return ( envs.VLLM_USE_FUSED_MQA_TOPK @@ -103,13 +106,13 @@ def use_fused_mqa_topk( ) -def fused_mqa_topk_chunk(max_seqlen_k: int, keys: int, topk_tokens: int) -> bool: +def fused_mqa_topk_chunk(max_seqlen_k: int, keys: int, topk: 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) and its gathered keys fit the kernel.""" return ( - max_seqlen_k >= fused_mqa_topk_min_keys(topk_tokens) + max_seqlen_k >= fused_mqa_topk_min_keys(topk) and keys <= FUSED_MQA_TOPK_MAX_KEYS ) From c0c9a054240347100d7986bfc361db7fb8f64ef8 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Tue, 27 Oct 2026 11:05:17 +0000 Subject: [PATCH 18/23] more cleanup and code refactor Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/persistent_epoch.inc | 37 ++++++++-- csrc/fused_mqa_topk/persistent_kernel.cuh | 69 +++++++++++++------ .../kernels/attention/test_fused_mqa_topk.py | 28 ++++++++ 3 files changed, 108 insertions(+), 26 deletions(-) diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc index b47323982be5..8c79c09f0c0c 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -366,10 +366,14 @@ if (warp_idx == 8) { } 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 = @@ -406,17 +410,38 @@ if (warp_idx == 8) { (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 (valid) { - const uint32_t bin = half_bin<5>(bq); - atomicAdd(sample_smem.hist[query_base + i] + padded_bin(bin), 1u); + 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], gate_s + warp_idx, - e.ends[warp_idx] - e.starts[warp_idx], - sample_target, lane); + 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. diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh index 43eec73856b0..15d1879008fd 100644 --- a/csrc/fused_mqa_topk/persistent_kernel.cuh +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -181,33 +181,62 @@ __device__ __forceinline__ float reduce_heads_with_lin(const float* s, } // Gate sample state: each math warp histograms its row's sampled scores -// (half_bin<5>: kSampleBins bins, padded) and finish_sample_histogram picks the -// gate from them. -constexpr uint32_t kSampleBins = 2048; +// 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 2048-bin histogram; one warp per row. -// -inf when the row sees at most `target` keys or has no samples. +// 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* gate, uint32_t visible, uint32_t target, - uint32_t lane) { + 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 the 64 bins of 32-bin chunks kChunks-1-2l and kChunks-2-2l - // (bins descending with the lane). - constexpr uint32_t kChunks = kSampleBins / 32; - static_assert(kChunks == 2 * 32, "two 32-bin chunks per lane"); - const uint32_t* const h0 = hist + padded_bin((kChunks - 1 - 2 * lane) * 32); - const uint32_t* const h1 = hist + padded_bin((kChunks - 2 - 2 * lane) * 32); + // 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 j = 0; j < 32; ++j) sum += h0[j] + h1[j]; + 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) { @@ -215,16 +244,16 @@ __device__ __noinline__ inline void finish_sample_histogram( __syncwarp(); return; } - const uint32_t rank = - max(1u, min(tot, (tot * target + visible - 1) / visible)); + 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 < 2; ++cc) { - const uint32_t c = kChunks - 1 - 2 * w - cc; + 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(hist[padded_bin(c * 32 + 31 - lane)], lane); + 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); @@ -232,7 +261,7 @@ __device__ __noinline__ inline void finish_sample_histogram( } above += __shfl_sync(kFullMask, v, 31); } - if (lane == 0) *gate = half_bin_floor<5>(chosen); + if (lane == 0) *gate = ref + half_bin_floor(chosen); __syncwarp(); } diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 4f6b629046cc..750f90a41a1a 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -276,6 +276,34 @@ def test_rescan_scores_only_its_block(): ) +@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) + fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out) + _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 = get_fused_mqa_topk_workspace(q.device, rows).counts[: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 From c4944cb9c8eb5fee2850f6bc8f3e07f025cd045e Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Tue, 27 Oct 2026 14:02:19 +0000 Subject: [PATCH 19/23] prefix caching heuristic Signed-off-by: LopezCastroRoberto --- .buildkite/test_areas/kernels.yaml | 3 +++ .../kernels/attention/test_fused_mqa_topk.py | 4 ++- .../kernels/attention/dsa/fused_mqa_topk.py | 26 +++++++++++++++---- .../layers/sparse_attn_indexer.py | 9 +++++-- vllm/models/glm5next/nvidia/sparse_indexer.py | 5 ++-- 5 files changed, 37 insertions(+), 10 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index a2dfc254c0da..6457ada12eb5 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 750f90a41a1a..3f58bc232ade 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -450,8 +450,10 @@ 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 + 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) diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 2f25a1e87222..b0e1caa4a026 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -37,6 +37,7 @@ def fused_mqa_topk_min_keys(topk: int) -> int: _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 @dataclass @@ -66,6 +67,11 @@ def as_args(self) -> tuple[torch.Tensor, ...]: _WORKSPACES: dict[int, FusedMQATopKWorkspace] = {} +@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.""" @@ -106,14 +112,25 @@ def use_fused_mqa_topk( ) -def fused_mqa_topk_chunk(max_seqlen_k: int, keys: int, topk: int) -> bool: +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) and its gathered keys fit the - kernel.""" + 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) ) @@ -138,8 +155,7 @@ def get_fused_mqa_topk_workspace( ws = _WORKSPACES.get(index) if ws is None: dev = torch.device("cuda", index) - num_sms = torch.cuda.get_device_properties(index).multi_processor_count - slots = _SLOTS_PER_SM * num_sms + slots = _SLOTS_PER_SM * _sm_count(index) ws = FusedMQATopKWorkspace( counter=torch.zeros(2, dtype=torch.int32, device=dev), values=torch.empty( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 7a6d814041fe..757887027c85 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -712,11 +712,17 @@ 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) 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 + chunk.max_seqlen_k, + k_quant.shape[0], + topk_tokens, + q_slice.shape[0], + dense_rows, ): fused_mqa_topk_prefill( q_slice, @@ -745,7 +751,6 @@ def sparse_attn_indexer( # budget: split their rows (budget-planned chunks, the only ones # under DCP, are one window, so the merge below sees all the # rows). - dense_rows = dense_logits_rows(chunk.max_seqlen_k) 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) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 99366288cd97..9097f4c30327 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -386,11 +386,13 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens] + # 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) elif fused_mqa_topk and fused_mqa_topk_chunk( - chunk.max_seqlen_k, k_quant.shape[0], select_k + chunk.max_seqlen_k, k_quant.shape[0], select_k, num_rows, dense_rows ): fused_mqa_topk_prefill( q_slice, @@ -408,7 +410,6 @@ def sparse_attn_indexer_kpool( # under DCP, are one window, so the merge below sees all the # rows). # One pass also for 0 rows: the merge is collective. - dense_rows = dense_logits_rows(chunk.max_seqlen_k) 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( From af3824901aa8116b4f5bf92833d36fe1c6368b88 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Thu, 29 Oct 2026 11:10:39 +0000 Subject: [PATCH 20/23] address simon comments Signed-off-by: LopezCastroRoberto --- csrc/fused_mqa_topk/persistent_epoch.inc | 9 +++++++++ csrc/fused_mqa_topk/persistent_kernel.cuh | 3 +++ .../kernels/attention/dsa/fused_mqa_topk.py | 3 +++ 3 files changed, 15 insertions(+) diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc index 8c79c09f0c0c..1af6a1d2088b 100644 --- a/csrc/fused_mqa_topk/persistent_epoch.inc +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -681,6 +681,10 @@ if (warp_idx == 8) { 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(); @@ -688,6 +692,8 @@ if (warp_idx == 8) { } // -------- 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) { @@ -721,6 +727,9 @@ if (defer_head != kTerm) { 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. diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh index 15d1879008fd..86063b08fbac 100644 --- a/csrc/fused_mqa_topk/persistent_kernel.cuh +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -580,6 +580,9 @@ __launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel( 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); } diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index b0e1caa4a026..4df90a954299 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -100,9 +100,12 @@ def use_fused_mqa_topk( 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 From 36488b2026553b47af3d8dba695d1f375aa1d90b Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 30 Oct 2026 11:30:21 +0000 Subject: [PATCH 21/23] address lucas comments - Allocate the fused_mqa_topk scratch and block counter per call from the caching allocator instead of a persistent per-device workspace: nothing stays reserved between calls, and the profiled peak is the larger of the scratch and the dense logits budget (KV cache unchanged vs dense). - Rename plan_fused_mqa_topk to fused_mqa_topk_eligible. - Build the fused source next to the other _C_stable_libtorch kernels and move the DeepGEMM include there. Co-authored-by: LopezCastroRoberto Co-authored-by: Lucas Wilkinson Co-Authored-By: Claude Opus 5.5 Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 37 ++++--- .../kernels/attention/test_fused_mqa_topk.py | 12 +- .../kernels/attention/dsa/fused_mqa_topk.py | 104 +++++------------- .../layers/sparse_attn_indexer.py | 8 +- vllm/models/deepseek_v4/cpu/cpu_mla.py | 4 +- vllm/models/glm5next/nvidia/sparse_indexer.py | 8 +- vllm/v1/attention/backends/mla/indexer.py | 27 +++-- 7 files changed, 81 insertions(+), 119 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 0890ef665150..f0aafbf2bebf 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) @@ -1701,20 +1724,6 @@ 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) diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 3f58bc232ade..89267275a62b 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -14,7 +14,7 @@ fused_mqa_topk_available, fused_mqa_topk_min_keys, fused_mqa_topk_prefill, - get_fused_mqa_topk_workspace, + get_fused_mqa_topk_workspaces, ) from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits from vllm.utils.deep_gemm import ( @@ -223,14 +223,13 @@ def test_multi_request_chunk_matches_dense(lens, 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; the per-row workspace grows to fit.""" + 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) - assert get_fused_mqa_topk_workspace(q.device, rows).counts.numel() >= rows _check(out, q, k, scales, weights, starts, ends) @@ -295,11 +294,14 @@ def test_gate_with_common_score_offset(): 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) + 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 = get_fused_mqa_topk_workspace(q.device, rows).counts[:rows] + counts = workspaces[-2][:rows] assert (counts >= topk).all() and (counts <= 16384).all(), "rows rescanned" assert counts.float().median() < 2 * 3072 diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py index 4df90a954299..928e222498a7 100644 --- a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py +++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py @@ -3,7 +3,6 @@ """Single-launch fused FP8 MQA-logits + top-k prefill selection (SM100 family)""" import functools -from dataclasses import dataclass import torch @@ -40,33 +39,6 @@ def fused_mqa_topk_min_keys(topk: int) -> int: _BLOCK_ROWS = 8 # kBlockRows: query rows per block, one block per CTA at a time -@dataclass -class FusedMQATopKWorkspace: - """Caller-owned scratch of the fused_mqa_topk kernel (see fused_mqa_topk.cu).""" - - counter: torch.Tensor # [2] int32, zero between calls (the kernel restores it) - values: torch.Tensor # [8S, 16384] uint16 (low 16 bits of fp24 codes), S = 2/SM - indices: torch.Tensor # [8S, 16384] int32 - retained: torch.Tensor # [8S, 4096] fp32 - retained_indices: torch.Tensor # [8S, 4096] int32 - counts: torch.Tensor # [rows] int32, per-row candidate counts - defer: torch.Tensor # [rows / 8] int32, blocks queued for the exact rescan - - def as_args(self) -> tuple[torch.Tensor, ...]: - return ( - self.counter, - self.values, - self.indices, - self.retained, - self.retained_indices, - self.counts, - self.defer, - ) - - -_WORKSPACES: dict[int, FusedMQATopKWorkspace] = {} - - @functools.cache def _sm_count(index: int) -> int: return torch.cuda.get_device_properties(index).multi_processor_count @@ -137,60 +109,30 @@ def fused_mqa_topk_chunk( ) -def get_fused_mqa_topk_workspace( +def get_fused_mqa_topk_workspaces( device: torch.device, rows: int -) -> FusedMQATopKWorkspace: - """Return the per-device workspace, allocating it on first use. - - The per-row part (counts, defer) holds at least ``rows`` query rows and grows - when a larger call arrives; the profiling run asks for the largest prefill - chunk (max_num_batched_tokens rows) so serving never grows it. The block - counter is zeroed once; every launch leaves it zeroed again. The workspace is - shared by every call on the device, so calls must not run concurrently (they - are issued on one stream). - """ - rows = (rows + 7) // 8 * 8 +) -> 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() ) - ws = _WORKSPACES.get(index) - if ws is None: - dev = torch.device("cuda", index) - slots = _SLOTS_PER_SM * _sm_count(index) - ws = FusedMQATopKWorkspace( - counter=torch.zeros(2, dtype=torch.int32, device=dev), - values=torch.empty( - (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.uint16, device=dev - ), - indices=torch.empty( - (8 * slots, _CANDIDATES_PER_ROW), dtype=torch.int32, device=dev - ), - retained=torch.empty( - (8 * slots, _RETAINED_PER_ROW), dtype=torch.float32, device=dev - ), - retained_indices=torch.empty( - (8 * slots, _RETAINED_PER_ROW), dtype=torch.int32, device=dev - ), - counts=torch.empty(rows, dtype=torch.int32, device=dev), - defer=torch.empty(rows // 8, dtype=torch.int32, device=dev), - ) - nbytes = sum(t.numel() * t.element_size() for t in ws.as_args()) - logger.info( - "Allocated fused MQA+top-k workspace on cuda:%d (%d slots, %.1f MiB).", - index, - slots, - nbytes / 2**20, - ) - _WORKSPACES[index] = ws - elif ws.counts.numel() < rows: - ws.counts = torch.empty(rows, dtype=torch.int32, device=ws.counts.device) - ws.defer = torch.empty(rows // 8, dtype=torch.int32, device=ws.defer.device) - logger.info( - "Grew the fused MQA+top-k row workspace on cuda:%d to %d rows.", index, rows - ) - return ws + 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: @@ -230,7 +172,13 @@ def fused_mqa_topk_prefill( "(FP8, 32 heads, top-k %d).", output.shape[1], ) - ws = get_fused_mqa_topk_workspace(q.device, q.shape[0]) torch.ops._C.fused_mqa_topk_prefill( - q, k, scales, weights, starts, ends, output, *ws.as_args() + 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 757887027c85..2ba2ed7d70f0 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -33,7 +33,7 @@ fused_mqa_topk_chunk, fused_mqa_topk_key_scales, fused_mqa_topk_prefill, - get_fused_mqa_topk_workspace, + get_fused_mqa_topk_workspaces, use_fused_mqa_topk, ) from vllm.model_executor.layers.indexer_topk import ( @@ -514,9 +514,9 @@ def sparse_attn_indexer( profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) if fused_mqa_topk: - # The fused_mqa_topk kernel's fixed-size workspace is persistent; - # allocate it during profiling so it is accounted for. - get_fused_mqa_topk_workspace(q_quant.device, q_quant.shape[0]) + # 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]) # Dummy allocation to simulate for peak logits tensor memory during inference. # Needed with the fused kernel too: chunks below fused_mqa_topk_min_keys diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index e3729f68510c..430886590d7c 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_mla.py +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -87,9 +87,9 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, - plan_fused_mqa_topk: bool = False, + fused_mqa_topk_eligible: bool = False, ) -> list[tuple[slice, slice]]: - # workspace_size/max_logits_bytes/plan_fused_mqa_topk are unused: no + # 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()) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 9097f4c30327..0b8d4fa788f2 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -15,7 +15,7 @@ fused_mqa_topk_chunk, fused_mqa_topk_key_scales, fused_mqa_topk_prefill, - get_fused_mqa_topk_workspace, + get_fused_mqa_topk_workspaces, use_fused_mqa_topk, ) from vllm.model_executor.layers.indexer_topk import get_indexer_topk @@ -170,10 +170,10 @@ def sparse_attn_indexer_kpool( scales_spec, ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) - # The fused_mqa_topk kernel's fixed-size workspace is persistent; - # allocate it during profiling so it is accounted for. + # 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_workspace(q_quant.device, q_quant.shape[0]) + 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 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 5e5ce9af3bbc..0e6c0ee77124 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1052,7 +1052,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: # chunks. model_config = self.vllm_config.model_config cfg = model_config.hf_text_config if model_config is not None else None - self.plan_fused_mqa_topk = cfg is not None and use_fused_mqa_topk( + 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), @@ -1398,7 +1398,7 @@ def _split_indexer_prefill_chunks( workspace_size: int, max_logits_bytes: int, request_offset: int = 0, - plan_fused_mqa_topk: bool = False, + 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) @@ -1408,12 +1408,12 @@ 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 ``plan_fused_mqa_topk`` (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. + 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. """ @@ -1431,11 +1431,14 @@ def _split_indexer_prefill_chunks( compressed_seq_lens_cpu[end].item(), ) new_m, new_n = chunk_m + q, chunk_n + s - if plan_fused_mqa_topk and new_n <= FUSED_MQA_TOPK_MAX_KEYS: + 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 plan_fused_mqa_topk and new_m * new_n <= max_logits_elems + 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 @@ -1457,7 +1460,7 @@ def _split_indexer_prefill_chunks( if chunk_n > 0 else max(1, chunk_m) ) - if plan_fused_mqa_topk and chunk_n <= FUSED_MQA_TOPK_MAX_KEYS: + 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) @@ -1594,7 +1597,7 @@ def build( self.max_prefill_buffer_size, max_logits_bytes, request_offset=num_decodes, - plan_fused_mqa_topk=self.plan_fused_mqa_topk, + fused_mqa_topk_eligible=self.fused_mqa_topk_eligible, ) chunks = [] From 02ab87a4f82bbf1c86b42cce2509fb503b5dfb11 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 30 Oct 2026 14:53:16 +0000 Subject: [PATCH 22/23] add co-author Co-authored-by: Ziqi Yin <62246771+Heisenberg-Yin@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Signed-off-by: LopezCastroRoberto From 69dd3ce947eeece3e6f2823f2e9b8df8cee55eb0 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 30 Oct 2026 19:07:41 +0000 Subject: [PATCH 23/23] fix rocm path on ci Signed-off-by: LopezCastroRoberto --- tests/kernels/attention/test_fused_mqa_topk.py | 3 ++- vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py index 89267275a62b..7042cedb6849 100644 --- a/tests/kernels/attention/test_fused_mqa_topk.py +++ b/tests/kernels/attention/test_fused_mqa_topk.py @@ -8,7 +8,6 @@ import torch import vllm.envs as envs -import vllm.models.glm5next.nvidia.sparse_indexer as kpool_indexer from vllm import _custom_ops as ops from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import ( fused_mqa_topk_available, @@ -373,6 +372,8 @@ def _kpool_inputs(pool_lens=(20000, 70000), q_rows=96, tail=2, seed=919): 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( 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,