forked from Karylab-cklius/vllm
Woosuk Kwon
@@ -564,6 +564,27 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
"in CUDA target architectures.")
|
||||
endif()
|
||||
|
||||
# DeepSeek V4 indexer top-k. Needs thread-block clusters + TMA + PDL, so
|
||||
# builds for Hopper (sm_90a) and Blackwell datacenter (sm_100/sm_103). Not
|
||||
# supported on sm_120 (consumer Blackwell, no clusters). Requires CUDA >=
|
||||
# 12.4 for the cuda::ptx mbarrier wrappers. Ported from sglang.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.4 AND DSV4_TOPK_ARCHS)
|
||||
set(DSV4_TOPK_SRC "csrc/deepseek_v4/fast_topk_v2.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${DSV4_TOPK_SRC}"
|
||||
CUDA_ARCHS "${DSV4_TOPK_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC ${DSV4_TOPK_SRC})
|
||||
message(STATUS "Building deepseek_v4 fast_topk_v2 for archs: ${DSV4_TOPK_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building deepseek_v4 fast_topk_v2 (needs CUDA >= 12.4 "
|
||||
"and a compatible Hopper+ arch).")
|
||||
endif()
|
||||
|
||||
#
|
||||
# Machete kernels
|
||||
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Microbench: fast_topk_v2 vs persistent_topk for k=512.
|
||||
|
||||
Both ops select the top-512 entries per row of a [B, L] float32 score tensor.
|
||||
The vLLM `persistent_topk` is the existing path used by the indexer
|
||||
(`sparse_attn_indexer.py`); `fast_topk_v2` is the new sm_90+ port from
|
||||
sglang that adds Hopper thread-block clusters and a fused page-table gather.
|
||||
|
||||
For fairness:
|
||||
- `persistent_topk` writes raw indices into a (B, k) tensor.
|
||||
- `fast_topk_v2` writes page-table-resolved indices. To match the same
|
||||
output (raw indices) we use ``page_size=1`` with an identity page table —
|
||||
page_to_indices becomes a no-op, so the two kernels produce comparable
|
||||
results on the hot path.
|
||||
|
||||
Timing uses **CUDA graph replay** to amortize launch overhead. We capture
|
||||
N invocations of the same kernel into a single graph, replay the graph
|
||||
many times, and divide. This isolates kernel work from per-launch host
|
||||
latency (~3–5 µs on Blackwell) which would otherwise dominate at small
|
||||
shapes.
|
||||
|
||||
Run::
|
||||
|
||||
.venv/bin/python benchmarks/kernels/benchmark_fast_topk_v2.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import statistics
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
import vllm._C # noqa: F401 ensures schemas are registered
|
||||
from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import (
|
||||
plan_topk_v2,
|
||||
workspace_ints_per_batch,
|
||||
)
|
||||
|
||||
K = 512
|
||||
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 # bytes; matches sparse_attn_indexer.py
|
||||
|
||||
|
||||
def _capture_graph(callable_fn, *, calls_per_graph: int) -> torch.cuda.CUDAGraph:
|
||||
"""Capture a CUDA graph that invokes `callable_fn` `calls_per_graph` times.
|
||||
|
||||
The kernel is run a few times outside capture (warmup + allocator priming)
|
||||
to avoid capturing one-shot setup work like cudaFuncSetAttribute.
|
||||
"""
|
||||
# Outside-capture warmup: lets cudaFuncSetAttribute's static cache fire
|
||||
# and primes any cublas/cudnn lookups the kernel might trigger.
|
||||
for _ in range(3):
|
||||
callable_fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
g = torch.cuda.CUDAGraph()
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(s):
|
||||
with torch.cuda.graph(g, stream=s):
|
||||
for _ in range(calls_per_graph):
|
||||
callable_fn()
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
return g
|
||||
|
||||
|
||||
def time_graph_us(graph: torch.cuda.CUDAGraph, *, calls_per_graph: int,
|
||||
warmup: int = 5, replays: int = 30) -> float:
|
||||
"""Median per-call latency (µs) over many graph replays."""
|
||||
for _ in range(warmup):
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
samples = []
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
for _ in range(replays):
|
||||
start.record()
|
||||
graph.replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
# ms -> us, then per-call.
|
||||
samples.append(start.elapsed_time(end) * 1000.0 / calls_per_graph)
|
||||
return statistics.median(samples)
|
||||
|
||||
|
||||
def make_inputs(batch_size: int, seq_len: int, *, seed: int = 0):
|
||||
device = torch.device("cuda")
|
||||
g = torch.Generator(device=device).manual_seed(seed)
|
||||
L = (seq_len + 3) & ~3 # TMA wants stride%4 == 0
|
||||
scores = torch.randn(batch_size, L, generator=g, dtype=torch.float32,
|
||||
device=device)
|
||||
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32,
|
||||
device=device)
|
||||
# Identity page table with page_size=1 -> page_to_indices is identity.
|
||||
page_table = (
|
||||
torch.arange(L, dtype=torch.int32, device=device)
|
||||
.unsqueeze(0).expand(batch_size, -1).contiguous()
|
||||
)
|
||||
return scores, seq_lens, page_table, L
|
||||
|
||||
|
||||
def bench_persistent_topk(scores, seq_lens, *, calls_per_graph: int) -> float:
|
||||
B = scores.shape[0]
|
||||
output = scores.new_empty((B, K), dtype=torch.int32)
|
||||
workspace = scores.new_empty((RADIX_TOPK_WORKSPACE_SIZE,),
|
||||
dtype=torch.uint8)
|
||||
max_seq_len = scores.shape[1]
|
||||
|
||||
def run():
|
||||
torch.ops._C.persistent_topk(
|
||||
scores, seq_lens, output, workspace, K, max_seq_len)
|
||||
|
||||
graph = _capture_graph(run, calls_per_graph=calls_per_graph)
|
||||
return time_graph_us(graph, calls_per_graph=calls_per_graph)
|
||||
|
||||
|
||||
def bench_fast_topk_v2(scores, seq_lens, page_table, *,
|
||||
calls_per_graph: int) -> tuple[float, float]:
|
||||
"""Returns (kernel_only_us, with_plan_us). Both are per-call medians.
|
||||
|
||||
- kernel_only: just torch.ops._C.fast_topk_v2 (plan was done once,
|
||||
metadata reused). This is the realistic decode-with-cudagraph case.
|
||||
- with_plan: plan + main kernel captured together. Pessimistic for
|
||||
cudagraph (plan would normally be done once outside the graph and
|
||||
reused, since it depends only on seq_lens shape).
|
||||
"""
|
||||
B = scores.shape[0]
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
workspace = scores.new_empty((B, workspace_ints_per_batch()),
|
||||
dtype=torch.int32)
|
||||
page_indices = scores.new_empty((B, K), dtype=torch.int32)
|
||||
|
||||
def run_kernel_only():
|
||||
torch.ops._C.fast_topk_v2(
|
||||
scores, seq_lens, page_table, page_indices, 1, workspace, metadata)
|
||||
|
||||
def run_with_plan():
|
||||
torch.ops._C.fast_topk_v2_plan(seq_lens, metadata, 0)
|
||||
torch.ops._C.fast_topk_v2(
|
||||
scores, seq_lens, page_table, page_indices, 1, workspace, metadata)
|
||||
|
||||
g_kernel = _capture_graph(run_kernel_only, calls_per_graph=calls_per_graph)
|
||||
g_with_plan = _capture_graph(run_with_plan, calls_per_graph=calls_per_graph)
|
||||
return (
|
||||
time_graph_us(g_kernel, calls_per_graph=calls_per_graph),
|
||||
time_graph_us(g_with_plan, calls_per_graph=calls_per_graph),
|
||||
)
|
||||
|
||||
|
||||
def fmt(us: float) -> str:
|
||||
return f"{us:8.2f}"
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--batch-sizes", type=int, nargs="+",
|
||||
default=[1, 4, 16, 32, 64, 128, 256])
|
||||
parser.add_argument("--seq-lens", type=int, nargs="+",
|
||||
default=[1024, 4096, 16384, 32768, 65536, 131072])
|
||||
parser.add_argument("--calls-per-graph", type=int, default=64,
|
||||
help="Invocations captured per graph (amortizes "
|
||||
"graph-replay overhead, ~3-5 µs).")
|
||||
parser.add_argument("--replays", type=int, default=30,
|
||||
help="Graph replays per measurement.")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA is required for this benchmark.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"GPU: {torch.cuda.get_device_name(0)} "
|
||||
f"(SM {torch.cuda.get_device_capability(0)})")
|
||||
print(f"k = {K}; calls_per_graph={args.calls_per_graph}, "
|
||||
f"replays={args.replays}")
|
||||
print("Per-call medians via CUDA graph replay (host launch overhead "
|
||||
"amortized).\n")
|
||||
|
||||
print(f"{'B':>4} {'L':>7} | "
|
||||
f"{'persistent_topk':>17} | "
|
||||
f"{'fast_topk_v2':>14} | "
|
||||
f"{'fast_topk_v2+plan':>19} | "
|
||||
f"{'speedup':>9} | {'path':<14}")
|
||||
print("-" * 105)
|
||||
|
||||
for B in args.batch_sizes:
|
||||
for L in args.seq_lens:
|
||||
try:
|
||||
scores, seq_lens, page_table, _ = make_inputs(B, L, seed=B * L)
|
||||
p_us = bench_persistent_topk(
|
||||
scores, seq_lens, calls_per_graph=args.calls_per_graph)
|
||||
f_us, fp_us = bench_fast_topk_v2(
|
||||
scores, seq_lens, page_table,
|
||||
calls_per_graph=args.calls_per_graph)
|
||||
speedup = p_us / f_us if f_us > 0 else float("inf")
|
||||
|
||||
if L <= 512:
|
||||
path = "trivial"
|
||||
elif L <= 4 * 4 * 1024:
|
||||
path = "register-1p"
|
||||
elif L <= 32768:
|
||||
path = "register-2p"
|
||||
elif B <= 15:
|
||||
path = "cluster-fused"
|
||||
else:
|
||||
path = "cluster-2stg"
|
||||
|
||||
print(
|
||||
f"{B:>4} {L:>7} | "
|
||||
f"{fmt(p_us):>14} us | "
|
||||
f"{fmt(f_us):>11} us | "
|
||||
f"{fmt(fp_us):>16} us | "
|
||||
f"{speedup:>6.2f}x | {path}"
|
||||
)
|
||||
except RuntimeError as e:
|
||||
print(f"{B:>4} {L:>7} | ERROR: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,566 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// DeepSeek V4 indexer top-k (k = 512). Ported from sglang's
|
||||
// jit_kernel/csrc/deepseek_v4/topk_v2.cuh.
|
||||
//
|
||||
// Combines three strategies (Register / Streaming / Cluster) dispatched per
|
||||
// row by a separate plan kernel that decides a `cluster_threshold` from the
|
||||
// observed seq_lens distribution. The host side picks one of three launch
|
||||
// shapes:
|
||||
// 1. all rows fit in the small (register) path -> single short kernel
|
||||
// 2. small batch (<= kNumClusters) with some long rows -> fused cluster
|
||||
// kernel (stage 1 + tie-break in one launch)
|
||||
// 3. larger batch -> persistent cluster stage 1 + non-cluster stage 2
|
||||
//
|
||||
// Architecture support: Hopper (sm_90a) and Blackwell datacenter (sm_100/
|
||||
// sm_103). Requires thread-block clusters, TMA bulk async copy, mbarrier,
|
||||
// and Programmatic Dependent Launch — sm_120 (consumer Blackwell) lacks
|
||||
// clusters and is not supported. The heuristic constants in `topk_plan`
|
||||
// were tuned on B200 (sglang upstream); they are functionally correct on
|
||||
// H100/H200 too but may be suboptimal until retuned.
|
||||
|
||||
#include "topk/cluster.cuh"
|
||||
#include "topk/common.cuh"
|
||||
#include "topk/register.cuh"
|
||||
#include "topk/streaming.cuh"
|
||||
#include "topk/utils.cuh"
|
||||
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <cooperative_groups.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/library.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
using Large = ClusterTopK<512>;
|
||||
using Medium = StreamingTopK<512>;
|
||||
using Small = RegisterTopK<512>;
|
||||
|
||||
using Metadata = Large::Metadata;
|
||||
constexpr uint32_t kNumClusters = 15; // hardware-capped persistent count
|
||||
constexpr uint32_t kClusterSize = Large::kClusterSize;
|
||||
constexpr uint32_t kMax2PassLength = Small::kMax2PassLength;
|
||||
constexpr uint32_t kMaxSupportedLength = Large::kMaxLength;
|
||||
|
||||
// Row 0 of the metadata tensor stores GlobalMetadata; rows [1..N+1) hold the
|
||||
// per-item Metadata entries that the persistent stage-1 consumes.
|
||||
struct alignas(16) GlobalMetadata {
|
||||
uint32_t cluster_threshold;
|
||||
uint32_t num_cluster_items;
|
||||
uint32_t reserved[2];
|
||||
};
|
||||
static_assert(sizeof(GlobalMetadata) == sizeof(Metadata),
|
||||
"metadata row 0 layout must match Metadata stride");
|
||||
|
||||
#define VLLM_SMALL_TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
|
||||
#define VLLM_LARGE_CLUSTER __cluster_dims__(1, kClusterSize, 1)
|
||||
// Stage 1 is persistent + cluster -> high smem -> occupancy 1.
|
||||
#define VLLM_LARGE_TOPK_STAGE_1 \
|
||||
__global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER
|
||||
// Stage 2 is non-cluster + small smem -> occupancy 2.
|
||||
#define VLLM_LARGE_TOPK_STAGE_2 __global__ __launch_bounds__(kBlockSize, 2)
|
||||
#define VLLM_FUSED_COMBINE_KERNEL \
|
||||
__global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER
|
||||
#define VLLM_PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1)
|
||||
|
||||
struct TopKParams {
|
||||
const uint32_t* __restrict__ seq_lens;
|
||||
const float* __restrict__ scores;
|
||||
const int32_t* __restrict__ page_table;
|
||||
int32_t* __restrict__ page_indices;
|
||||
int64_t score_stride;
|
||||
int64_t page_table_stride;
|
||||
uint8_t* __restrict__ workspace;
|
||||
const Metadata* __restrict__ metadata = nullptr;
|
||||
int64_t workspace_stride; // bytes per batch
|
||||
uint32_t batch_size;
|
||||
uint32_t page_bits;
|
||||
|
||||
VLLM_DSV4_DEVICE const float* get_scores(uint32_t batch_id) const {
|
||||
return scores + batch_id * score_stride;
|
||||
}
|
||||
VLLM_DSV4_DEVICE TransformParams get_transform(uint32_t batch_id,
|
||||
int32_t* indices) const {
|
||||
return {
|
||||
.page_table = page_table + batch_id * page_table_stride,
|
||||
.indices_in = indices,
|
||||
.indices_out = page_indices + batch_id * 512,
|
||||
.page_bits = page_bits,
|
||||
};
|
||||
}
|
||||
VLLM_DSV4_DEVICE const GlobalMetadata& get_global_metadata() const {
|
||||
return *reinterpret_cast<const GlobalMetadata*>(metadata);
|
||||
}
|
||||
VLLM_DSV4_DEVICE const Metadata& get_item_metadata(uint32_t work_id) const {
|
||||
return metadata[1 + work_id]; // skip the GlobalMetadata row
|
||||
}
|
||||
};
|
||||
|
||||
VLLM_DSV4_DEVICE uint2 partition_work(uint32_t length, uint32_t rank) {
|
||||
constexpr uint32_t kTMAAlign = 4;
|
||||
const auto total_units = (length + kTMAAlign - 1) / kTMAAlign;
|
||||
const auto base = total_units / kClusterSize;
|
||||
const auto extra = total_units % kClusterSize;
|
||||
const auto local_units = base + (rank < extra ? 1u : 0u);
|
||||
const auto offset_units = rank * base + min(rank, extra);
|
||||
const auto offset = offset_units * kTMAAlign;
|
||||
const auto finish = min(offset + local_units * kTMAAlign, length);
|
||||
return {offset, finish - offset};
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Plan kernel: decides cluster_threshold from the observed seq_lens
|
||||
// distribution and compacts items with seq_len > threshold into metadata[1..].
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
VLLM_PLAN_KERNEL void topk_plan(const uint32_t* __restrict__ seq_lens,
|
||||
Metadata* __restrict__ metadata,
|
||||
uint32_t batch_size,
|
||||
uint32_t static_cluster_threshold) {
|
||||
// (threshold, max_batch_size_for_that_threshold). Tuned on B200 by sglang.
|
||||
struct Pair {
|
||||
uint32_t threshold;
|
||||
uint32_t max_batch_size;
|
||||
};
|
||||
constexpr Pair kCandidates[] = {
|
||||
{32768, 30}, {40960, 45}, {49152, 45}, {65536, 60},
|
||||
{98304, 60}, {131072, 75}, {196608, 90}, {262144, 105},
|
||||
};
|
||||
constexpr uint32_t kNumCandidates =
|
||||
sizeof(kCandidates) / sizeof(kCandidates[0]);
|
||||
constexpr uint32_t kMinBatchSize = kCandidates[0].max_batch_size;
|
||||
static_assert(kCandidates[0].threshold == kMax2PassLength);
|
||||
static_assert(kCandidates[kNumCandidates - 1].threshold ==
|
||||
kMaxSupportedLength);
|
||||
|
||||
__shared__ uint32_t s_count;
|
||||
__shared__ uint32_t s_counts[kNumCandidates];
|
||||
__shared__ uint32_t s_threshold;
|
||||
|
||||
const auto tx = threadIdx.x;
|
||||
if (tx == 0) s_count = 0;
|
||||
if (tx < kNumCandidates) s_counts[tx] = 0;
|
||||
__syncthreads();
|
||||
|
||||
if (static_cluster_threshold > 0) {
|
||||
if (tx == 0) s_threshold = static_cluster_threshold;
|
||||
} else if (batch_size <= kMinBatchSize) {
|
||||
if (tx == 0) s_threshold = kMax2PassLength;
|
||||
} else {
|
||||
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
|
||||
const uint32_t sl = seq_lens[i];
|
||||
assert(sl <= kMaxSupportedLength);
|
||||
uint32_t count = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kNumCandidates; ++j) {
|
||||
count += (sl > kCandidates[j].threshold ? 1 : 0);
|
||||
}
|
||||
if (count > 0) {
|
||||
atomicAdd(&s_counts[count - 1], 1);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (tx == 0) {
|
||||
uint32_t accum = 0;
|
||||
uint32_t chosen = kMaxSupportedLength;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumCandidates; ++i) {
|
||||
const auto j = kNumCandidates - 1 - i;
|
||||
accum += s_counts[j];
|
||||
if (accum > kCandidates[j].max_batch_size) break;
|
||||
chosen = kCandidates[j].threshold;
|
||||
}
|
||||
s_threshold = chosen;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
const auto cluster_threshold = max(s_threshold, kMax2PassLength);
|
||||
|
||||
// Compact items with seq_len > cluster_threshold into metadata[1..N+1).
|
||||
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
|
||||
const uint32_t sl = seq_lens[i];
|
||||
if (sl > cluster_threshold) {
|
||||
const auto pos = atomicAdd(&s_count, 1);
|
||||
metadata[1 + pos] = {i, sl, false};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
const auto N = s_count;
|
||||
|
||||
// has_next chain for the persistent consumer + sentinel slots.
|
||||
for (uint32_t i = tx; i < N; i += kBlockSize) {
|
||||
if (i + kNumClusters < N) metadata[1 + i].has_next = true;
|
||||
}
|
||||
if (tx < kNumClusters && tx >= N) metadata[1 + tx] = {0, 0, false};
|
||||
if (tx == 0) {
|
||||
auto* g = reinterpret_cast<GlobalMetadata*>(metadata);
|
||||
*g = {
|
||||
.cluster_threshold = cluster_threshold,
|
||||
.num_cluster_items = N,
|
||||
.reserved = {0, 0},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Short kernel: all rows fit in the register path (max_seq_len <=
|
||||
// Small::kMax1PassLength).
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
VLLM_SMALL_TOPK_KERNEL void topk_short_transform(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[512];
|
||||
const auto batch_id = blockIdx.x;
|
||||
const auto seq_len = params.seq_lens[batch_id];
|
||||
const auto transform = params.get_transform(batch_id, s_topk_indices);
|
||||
if (seq_len <= 512) {
|
||||
trivial_transform(transform, seq_len, 512);
|
||||
} else {
|
||||
Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
|
||||
/*use_pdl=*/true);
|
||||
pdl_trigger_secondary<true>();
|
||||
Small::transform(transform);
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Persistent stage 1 (cluster). One CTA per cluster; the persistent block
|
||||
// walks `metadata[1..N]` round-robin and runs Large::stage1 per item.
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[512];
|
||||
uint32_t work_id = blockIdx.x;
|
||||
uint32_t batch_id = 0, seq_len = 0, length = 0, offset = 0;
|
||||
bool has_next = false;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
|
||||
const auto prefetch_metadata = [&] {
|
||||
const auto m = params.get_item_metadata(work_id);
|
||||
batch_id = m.batch_id;
|
||||
seq_len = m.seq_len;
|
||||
has_next = m.has_next;
|
||||
work_id += kNumClusters;
|
||||
};
|
||||
const auto launch_prologue = [&] {
|
||||
const auto partition = partition_work(seq_len, cluster_rank);
|
||||
offset = partition.x;
|
||||
length = partition.y;
|
||||
Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem);
|
||||
};
|
||||
|
||||
pdl_wait_primary<true>();
|
||||
pdl_trigger_secondary<true>();
|
||||
|
||||
prefetch_metadata();
|
||||
if (seq_len == 0) return;
|
||||
Large::stage1_init(smem);
|
||||
launch_prologue();
|
||||
while (true) {
|
||||
const auto this_length = length;
|
||||
const auto this_offset = offset;
|
||||
const auto need_prefetch = has_next;
|
||||
const auto transform = params.get_transform(batch_id, s_topk_indices);
|
||||
const auto ws = params.workspace + batch_id * params.workspace_stride;
|
||||
if (need_prefetch) prefetch_metadata();
|
||||
Large::stage1(s_topk_indices, this_length, smem, /*reuse=*/true);
|
||||
if (need_prefetch) launch_prologue();
|
||||
Large::stage1_epilogue(transform, this_offset, ws, smem);
|
||||
if (!need_prefetch) break;
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Stage 2 (non-cluster). Per-row dispatch: trivial / Small / Medium / Large.
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[512];
|
||||
const auto batch_id = blockIdx.x;
|
||||
const auto seq_len = params.seq_lens[batch_id];
|
||||
const auto cluster_threshold = params.get_global_metadata().cluster_threshold;
|
||||
const auto transform = params.get_transform(batch_id, s_topk_indices);
|
||||
if (seq_len <= 512) {
|
||||
trivial_transform(transform, seq_len, 512);
|
||||
} else if (seq_len <= kMax2PassLength) {
|
||||
if (seq_len <= Small::kMax1PassLength) {
|
||||
Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem);
|
||||
} else {
|
||||
__syncwarp();
|
||||
Small::run<true>(params.get_scores(batch_id), s_topk_indices, seq_len,
|
||||
smem);
|
||||
}
|
||||
Small::transform(transform);
|
||||
} else if (seq_len <= cluster_threshold) {
|
||||
Medium::run(params.get_scores(batch_id), seq_len, s_topk_indices, smem);
|
||||
Medium::transform(transform, smem);
|
||||
} else {
|
||||
const auto ws = params.workspace + batch_id * params.workspace_stride;
|
||||
pdl_wait_primary<true>();
|
||||
Large::transform(transform, ws, smem);
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Fused kernel for small batches. Both stage 1 and the tie-break run inside
|
||||
// the same launch; cluster rank 0 finishes the row.
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
VLLM_FUSED_COMBINE_KERNEL void topk_fused_transform(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[512];
|
||||
const auto batch_id = blockIdx.x;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
const auto seq_len = params.seq_lens[batch_id];
|
||||
const auto transform = params.get_transform(batch_id, s_topk_indices);
|
||||
if (seq_len <= 512) {
|
||||
if (cluster_rank != 0) return;
|
||||
trivial_transform(transform, seq_len, 512);
|
||||
} else if (seq_len <= Small::kMax1PassLength) {
|
||||
if (cluster_rank != 0) return;
|
||||
Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
|
||||
/*use_pdl=*/true);
|
||||
Small::transform(transform);
|
||||
} else {
|
||||
const auto partition = partition_work(seq_len, cluster_rank);
|
||||
const auto offset = partition.x;
|
||||
const auto length = partition.y;
|
||||
const auto ws = params.workspace + batch_id * params.workspace_stride;
|
||||
Large::stage1_init(smem);
|
||||
pdl_wait_primary<true>();
|
||||
Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem);
|
||||
Large::stage1(s_topk_indices, length, smem);
|
||||
Large::stage1_epilogue(transform, offset, ws, smem);
|
||||
cooperative_groups::this_cluster().sync();
|
||||
if (cluster_rank != 0) return;
|
||||
Large::transform(transform, ws, smem);
|
||||
}
|
||||
}
|
||||
|
||||
constexpr size_t kStage1SMEM = sizeof(Large::Smem) + 128;
|
||||
constexpr size_t kStage2SMEM =
|
||||
(sizeof(Small::Smem) > sizeof(Medium::Smem) ? sizeof(Small::Smem)
|
||||
: sizeof(Medium::Smem)) +
|
||||
128;
|
||||
|
||||
// Per-(kernel, smem) memoization: each instantiation has its own static. This
|
||||
// matters because cudaFuncSetAttribute is per-function and we want it to fire
|
||||
// exactly once per kernel symbol.
|
||||
template <auto* f, size_t kSmem>
|
||||
void setup_kernel_smem_once() {
|
||||
[[maybe_unused]] static const auto result = [] {
|
||||
return cudaFuncSetAttribute(reinterpret_cast<const void*>(f),
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
static_cast<int>(kSmem));
|
||||
}();
|
||||
TORCH_CHECK(result == cudaSuccess,
|
||||
"fast_topk_v2: cudaFuncSetAttribute failed: ",
|
||||
cudaGetErrorString(result));
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Host-side launchers
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_DTYPE(x, t) \
|
||||
TORCH_CHECK(x.scalar_type() == (t), #x " must be ", #t)
|
||||
#define CHECK_CONTIG(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
|
||||
void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata,
|
||||
int64_t static_cluster_threshold) {
|
||||
using namespace vllm::dsv4_topk;
|
||||
CHECK_CUDA(seq_lens);
|
||||
CHECK_CUDA(metadata);
|
||||
CHECK_DTYPE(seq_lens, torch::kInt32);
|
||||
CHECK_DTYPE(metadata, torch::kInt32);
|
||||
TORCH_CHECK(seq_lens.dim() == 1);
|
||||
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4);
|
||||
TORCH_CHECK(metadata.size(0) == seq_lens.size(0) + 1,
|
||||
"metadata must be (batch_size + 1, 4)");
|
||||
CHECK_CONTIG(seq_lens);
|
||||
CHECK_CONTIG(metadata);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(seq_lens.size(0));
|
||||
if (batch_size <= kNumClusters) return; // metadata unused in fused path
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
cudaLaunchConfig_t cfg{};
|
||||
cfg.gridDim = dim3(1);
|
||||
cfg.blockDim = dim3(kBlockSize);
|
||||
cfg.dynamicSmemBytes = 0;
|
||||
cfg.stream = stream;
|
||||
cfg.numAttrs = 0;
|
||||
TORCH_CHECK(cudaLaunchKernelEx(
|
||||
&cfg, &topk_plan,
|
||||
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
|
||||
reinterpret_cast<Metadata*>(metadata.data_ptr<int32_t>()),
|
||||
batch_size,
|
||||
static_cast<uint32_t>(static_cluster_threshold)) == cudaSuccess,
|
||||
"fast_topk_v2_plan launch failed: ",
|
||||
cudaGetErrorString(cudaGetLastError()));
|
||||
}
|
||||
|
||||
void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens,
|
||||
const torch::Tensor& page_table, torch::Tensor& page_indices,
|
||||
int64_t page_size, const torch::Tensor& workspace,
|
||||
const torch::Tensor& metadata) {
|
||||
using namespace vllm::dsv4_topk;
|
||||
CHECK_CUDA(scores);
|
||||
CHECK_CUDA(seq_lens);
|
||||
CHECK_CUDA(page_table);
|
||||
CHECK_CUDA(page_indices);
|
||||
CHECK_CUDA(workspace);
|
||||
CHECK_CUDA(metadata);
|
||||
CHECK_DTYPE(scores, torch::kFloat32);
|
||||
CHECK_DTYPE(seq_lens, torch::kInt32);
|
||||
CHECK_DTYPE(page_table, torch::kInt32);
|
||||
CHECK_DTYPE(page_indices, torch::kInt32);
|
||||
CHECK_DTYPE(workspace, torch::kInt32);
|
||||
CHECK_DTYPE(metadata, torch::kInt32);
|
||||
|
||||
TORCH_CHECK(scores.dim() == 2 && scores.stride(1) == 1,
|
||||
"scores must be 2D with last stride 1");
|
||||
TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.is_contiguous());
|
||||
TORCH_CHECK(page_table.dim() == 2 && page_table.stride(1) == 1,
|
||||
"page_table must be 2D with last stride 1");
|
||||
TORCH_CHECK(page_indices.dim() == 2 && page_indices.is_contiguous() &&
|
||||
page_indices.size(1) == 512,
|
||||
"page_indices must be (B, 512) contiguous");
|
||||
TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
|
||||
workspace.size(1) == Large::kWorkspaceInts,
|
||||
"workspace must be (B, kWorkspaceInts) with last stride 1");
|
||||
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4 &&
|
||||
metadata.is_contiguous(),
|
||||
"metadata must be (B + 1, 4) contiguous");
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(scores.size(0));
|
||||
TORCH_CHECK(seq_lens.size(0) == batch_size);
|
||||
TORCH_CHECK(page_table.size(0) == batch_size);
|
||||
TORCH_CHECK(page_indices.size(0) == batch_size);
|
||||
TORCH_CHECK(workspace.size(0) == batch_size);
|
||||
TORCH_CHECK(metadata.size(0) == batch_size + 1);
|
||||
|
||||
const auto max_seq_len = static_cast<uint32_t>(scores.size(1));
|
||||
TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0,
|
||||
"page_size must be a positive power of 2");
|
||||
TORCH_CHECK(scores.stride(0) % 4 == 0,
|
||||
"score stride must be a multiple of 4 (TMA 16-byte alignment)");
|
||||
|
||||
// page_bits = log2(page_size). __builtin_ctzll is a host-side compiler
|
||||
// builtin available under C++17 (vLLM compiles host code with C++17).
|
||||
const auto page_bits = static_cast<uint32_t>(
|
||||
__builtin_ctzll(static_cast<unsigned long long>(page_size)));
|
||||
TopKParams params{
|
||||
.seq_lens =
|
||||
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
|
||||
.scores = scores.data_ptr<float>(),
|
||||
.page_table = page_table.data_ptr<int32_t>(),
|
||||
.page_indices = page_indices.data_ptr<int32_t>(),
|
||||
.score_stride = scores.stride(0),
|
||||
.page_table_stride = page_table.stride(0),
|
||||
.workspace = reinterpret_cast<uint8_t*>(workspace.data_ptr<int32_t>()),
|
||||
.metadata =
|
||||
reinterpret_cast<const Metadata*>(metadata.data_ptr<int32_t>()),
|
||||
.workspace_stride =
|
||||
workspace.stride(0) * static_cast<int64_t>(sizeof(int32_t)),
|
||||
.batch_size = batch_size,
|
||||
.page_bits = page_bits,
|
||||
};
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
// Helper: build a cudaLaunchConfig with optional PDL + cluster attributes.
|
||||
// The attribute storage must outlive cudaLaunchKernelEx (cfg.attrs points
|
||||
// into it), so we keep it as a local in each call site.
|
||||
auto make_cfg = [&](dim3 grid, dim3 block, size_t smem,
|
||||
cudaLaunchAttribute* attrs, bool enable_cluster,
|
||||
bool enable_pdl) {
|
||||
cudaLaunchConfig_t cfg{};
|
||||
cfg.gridDim = grid;
|
||||
cfg.blockDim = block;
|
||||
cfg.dynamicSmemBytes = static_cast<unsigned>(smem);
|
||||
cfg.stream = stream;
|
||||
int n = 0;
|
||||
if (enable_pdl) {
|
||||
attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[n].val.programmaticStreamSerializationAllowed = 1;
|
||||
++n;
|
||||
}
|
||||
if (enable_cluster) {
|
||||
attrs[n].id = cudaLaunchAttributeClusterDimension;
|
||||
attrs[n].val.clusterDim = {1, kClusterSize, 1};
|
||||
++n;
|
||||
}
|
||||
cfg.numAttrs = n;
|
||||
cfg.attrs = n ? attrs : nullptr;
|
||||
return cfg;
|
||||
};
|
||||
|
||||
auto check_launch = [](cudaError_t err) {
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"fast_topk_v2 launch failed: ", cudaGetErrorString(err));
|
||||
};
|
||||
|
||||
if (max_seq_len <= Small::kMax1PassLength) {
|
||||
setup_kernel_smem_once<&topk_short_transform, kStage2SMEM>();
|
||||
cudaLaunchAttribute attrs[2];
|
||||
auto cfg = make_cfg(dim3(batch_size), dim3(kBlockSize), kStage2SMEM, attrs,
|
||||
/*cluster=*/false, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(&cfg, topk_short_transform, params));
|
||||
} else if (batch_size <= kNumClusters) {
|
||||
constexpr size_t kFusedSMEM =
|
||||
kStage1SMEM > kStage2SMEM ? kStage1SMEM : kStage2SMEM;
|
||||
setup_kernel_smem_once<&topk_fused_transform, kFusedSMEM>();
|
||||
cudaLaunchAttribute attrs[2];
|
||||
auto cfg = make_cfg(dim3(batch_size, kClusterSize), dim3(kBlockSize),
|
||||
kFusedSMEM, attrs, /*cluster=*/true, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(&cfg, topk_fused_transform, params));
|
||||
} else {
|
||||
const auto num_clusters = std::min<uint32_t>(batch_size, kNumClusters);
|
||||
setup_kernel_smem_once<&topk_combine_preprocess, kStage1SMEM>();
|
||||
cudaLaunchAttribute attrs1[2];
|
||||
auto cfg1 = make_cfg(dim3(num_clusters, kClusterSize), dim3(kBlockSize),
|
||||
kStage1SMEM, attrs1, /*cluster=*/true, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(&cfg1, topk_combine_preprocess, params));
|
||||
|
||||
setup_kernel_smem_once<&topk_combine_transform, kStage2SMEM>();
|
||||
cudaLaunchAttribute attrs2[2];
|
||||
auto cfg2 = make_cfg(dim3(batch_size), dim3(kBlockSize), kStage2SMEM,
|
||||
attrs2, /*cluster=*/false, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(&cfg2, topk_combine_transform, params));
|
||||
}
|
||||
}
|
||||
|
||||
int64_t fast_topk_v2_workspace_ints() {
|
||||
return static_cast<int64_t>(vllm::dsv4_topk::Large::kWorkspaceInts);
|
||||
}
|
||||
|
||||
// Register impls here (instead of in torch_bindings.cpp) so they only exist
|
||||
// when CMake compiles this source — i.e., when the target build has a
|
||||
// compatible Hopper / Blackwell-datacenter arch. On other configs the schema
|
||||
// remains defined but a call surfaces a clear "no impl" runtime error.
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("fast_topk_v2_plan", &fast_topk_v2_plan);
|
||||
m.impl("fast_topk_v2", &fast_topk_v2);
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CompositeExplicitAutograd, m) {
|
||||
m.impl("fast_topk_v2_workspace_ints", &fast_topk_v2_workspace_ints);
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Cluster top-k strategy for very large N. Uses Hopper thread-block clusters
|
||||
// (cooperative_groups::this_cluster) to parallelize histogram + scatter across
|
||||
// up to ``kClusterSize`` blocks per row. Each row is processed in two stages:
|
||||
// stage 1: per-block histogram, all-reduce across the cluster, threshold
|
||||
// scatter, and an epilogue that page-translates strictly-above
|
||||
// entries to global memory and stages ties into a per-row workspace.
|
||||
// stage 2: tie-break across the cluster's combined ties (run by cluster
|
||||
// rank 0 in the fused kernel, or as a separate launch otherwise).
|
||||
// Ported from
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
#include "ptx.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cooperative_groups.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
template <uint32_t K>
|
||||
struct ClusterTopK {
|
||||
static constexpr uint32_t kClusterSize = 8;
|
||||
static constexpr uint32_t kHistBits = 10;
|
||||
static constexpr uint32_t kHistBins = 1 << kHistBits;
|
||||
static constexpr uint32_t kElemPerStage = 8;
|
||||
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
|
||||
static constexpr uint32_t kNumStages = 4;
|
||||
static constexpr uint32_t kMaxLength = kClusterSize * kNumStages * kSizePerStage;
|
||||
static constexpr uint32_t kAboveBits = 11;
|
||||
|
||||
struct Smem {
|
||||
uint64_t barrier[kNumStages];
|
||||
uint32_t local_above_equal[kClusterSize];
|
||||
uint32_t prefix_above_equal;
|
||||
alignas(128) uint32_t counter_gt;
|
||||
alignas(128) uint32_t counter_eq;
|
||||
alignas(128) MatchBin match;
|
||||
alignas(128) uint32_t warp_sum[kNumWarps];
|
||||
uint32_t histogram[kHistBins];
|
||||
alignas(128) float score_buffer[kNumStages][kSizePerStage];
|
||||
Tie tie_buffer[kMaxTies];
|
||||
};
|
||||
|
||||
// Per-row metadata produced by the plan kernel and consumed by the fused /
|
||||
// stage-1 kernels. {batch_id, seq_len, has_next} arranged in an int4-sized
|
||||
// 16-byte struct so the planner can do contiguous int32x4 stores.
|
||||
struct alignas(16) Metadata {
|
||||
uint32_t batch_id;
|
||||
uint32_t seq_len;
|
||||
bool has_next;
|
||||
};
|
||||
|
||||
// Per-row workspace storing {(num_above, num_ties)} + the gathered ties.
|
||||
struct WorkSpace {
|
||||
uint2 metadata;
|
||||
Tie ties[kMaxTies];
|
||||
};
|
||||
|
||||
static constexpr uint32_t kWorkspaceInts = sizeof(WorkSpace) / sizeof(uint32_t);
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1_init(void* _smem) {
|
||||
const auto tx = threadIdx.x;
|
||||
__builtin_assume(tx < kBlockSize);
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
if (tx < kHistBins) smem->histogram[tx] = 0;
|
||||
if (tx < kNumStages) ptx::mbarrier_init(&smem->barrier[tx], 1);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1_prologue(const float* scores,
|
||||
uint32_t length, void* _smem) {
|
||||
if (threadIdx.x == 0) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
|
||||
const auto length_aligned = (length + 3u) & ~3u;
|
||||
#pragma unroll
|
||||
for (uint32_t stage = 0; stage < kNumStages; stage++) {
|
||||
if (stage >= num_stages) break;
|
||||
const auto offset = stage * kSizePerStage;
|
||||
const auto size = min(kSizePerStage, length_aligned - offset);
|
||||
const auto size_bytes = size * sizeof(float);
|
||||
const auto bar = &smem->barrier[stage];
|
||||
ptx::tma_load(smem->score_buffer[stage], scores + offset, size_bytes,
|
||||
bar);
|
||||
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1(int32_t* indices, uint32_t length,
|
||||
void* _smem, bool reuse = false) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
__builtin_assume(tx < kBlockSize);
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
// Local histogram.
|
||||
#pragma unroll
|
||||
for (uint32_t stage = 0; stage < kNumStages; stage++) {
|
||||
const auto offset = stage * kSizePerStage;
|
||||
if (offset >= length) break;
|
||||
const auto size = min(kSizePerStage, length - offset);
|
||||
if (lane_id == 0) ptx::mbarrier_wait(&smem->barrier[stage], 0);
|
||||
__syncwarp();
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemPerStage; ++i) {
|
||||
const auto idx = tx + i * kBlockSize;
|
||||
if (idx >= size) break;
|
||||
const auto score = smem->score_buffer[stage][idx];
|
||||
const auto bin = extract_coarse_bin<kHistBits>(score);
|
||||
atomicAdd(&smem->histogram[bin], 1);
|
||||
}
|
||||
}
|
||||
|
||||
static_assert(kHistBins <= kBlockSize);
|
||||
|
||||
// Two-shot all-reduce across the cluster.
|
||||
{
|
||||
auto cluster = cooperative_groups::this_cluster();
|
||||
cluster.sync();
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
const auto kLocalSize = kHistBins / kClusterSize;
|
||||
const auto offset = kLocalSize * cluster_rank;
|
||||
|
||||
const auto src_tx = tx / kClusterSize;
|
||||
const auto src_rank = tx % kClusterSize;
|
||||
|
||||
if (tx < kHistBins) {
|
||||
const auto addr = &smem->histogram[offset + src_tx];
|
||||
const auto src_addr = cluster.map_shared_rank(addr, src_rank);
|
||||
*src_addr = warp_reduce_sum<kClusterSize>(*src_addr);
|
||||
}
|
||||
cluster.sync();
|
||||
}
|
||||
|
||||
// Each block now holds the full cluster histogram. Find the threshold.
|
||||
{
|
||||
const auto value = tx < kHistBins ? smem->histogram[tx] : 0;
|
||||
const auto warp_inc = warp_inclusive_sum(lane_id, value);
|
||||
if (lane_id == kWarpThreads - 1) {
|
||||
smem->warp_sum[warp_id] = warp_inc;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
const auto tmp = smem->warp_sum[lane_id];
|
||||
const auto total_length = warp_reduce_sum(tmp);
|
||||
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
prefix_sum += warp_inc;
|
||||
const auto above = total_length - prefix_sum;
|
||||
if (tx < kHistBins && above < K && above + value >= K) {
|
||||
smem->counter_gt = smem->counter_eq = 0;
|
||||
smem->match = {
|
||||
.bin = tx,
|
||||
.above_count = above,
|
||||
.equal_count = value,
|
||||
};
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const auto thr_bin = smem->match.bin;
|
||||
|
||||
// Scatter strictly-above entries to `indices`, stash ties in tie_buffer.
|
||||
#pragma unroll
|
||||
for (uint32_t stage = 0; stage < kNumStages; stage++) {
|
||||
const auto offset = stage * kSizePerStage;
|
||||
if (offset >= length) break;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemPerStage; ++i) {
|
||||
const auto buf_idx = tx + i * kBlockSize;
|
||||
const auto global_idx = offset + buf_idx;
|
||||
if (global_idx >= length) break;
|
||||
const auto score = smem->score_buffer[stage][buf_idx];
|
||||
const auto bin = extract_coarse_bin<kHistBits>(score);
|
||||
if (bin > thr_bin) {
|
||||
indices[atomicAdd(&smem->counter_gt, 1)] = global_idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
|
||||
}
|
||||
}
|
||||
}
|
||||
if (reuse) {
|
||||
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
|
||||
if (tx < kHistBins) smem->histogram[tx] = 0;
|
||||
if (tx < num_stages) ptx::mbarrier_arrive(&smem->barrier[tx]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1_epilogue(TransformParams params,
|
||||
uint32_t offset, void* _ws,
|
||||
void* _smem) {
|
||||
auto cluster = cooperative_groups::this_cluster();
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto local_above = smem->counter_gt;
|
||||
const auto local_equal = smem->counter_eq;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
|
||||
constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
|
||||
static_assert(kAboveMask >= K);
|
||||
|
||||
static_assert(kMaxTies <= kBlockSize);
|
||||
const auto idx_above = tx < local_above ? params.indices_in[tx] : 0;
|
||||
const auto tie_value = tx < local_equal ? smem->tie_buffer[tx] : Tie{0, 0.0f};
|
||||
|
||||
// Push counts to remote shared memory to reduce inter-block latency.
|
||||
if (tx < kClusterSize) {
|
||||
const auto value = (local_equal << kAboveBits) | local_above;
|
||||
const auto dst_addr = cluster.map_shared_rank(smem->local_above_equal, tx);
|
||||
dst_addr[cluster_rank] = value;
|
||||
}
|
||||
// After this final sync, every block can read only its own smem (peer
|
||||
// ranks may have already exited), so we don't touch remote smem again.
|
||||
cluster.sync();
|
||||
if (tx < kClusterSize) {
|
||||
const auto value = tx < cluster_rank ? smem->local_above_equal[tx] : 0;
|
||||
const auto kActiveMask = (1u << kClusterSize) - 1;
|
||||
smem->prefix_above_equal = warp_reduce_sum<kClusterSize>(value, kActiveMask);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto prefix_packed = smem->prefix_above_equal;
|
||||
const auto prefix_above = prefix_packed & kAboveMask;
|
||||
const auto prefix_equal = prefix_packed >> kAboveBits;
|
||||
|
||||
// Page-translate strictly-above entries.
|
||||
if (tx < local_above) {
|
||||
params.write(tx + prefix_above, idx_above + offset);
|
||||
}
|
||||
// Stage ties into the per-row workspace (regular global writes).
|
||||
const auto ws = static_cast<WorkSpace*>(_ws);
|
||||
if (tx < local_equal && tx + prefix_equal < kMaxTies) {
|
||||
ws->ties[tx + prefix_equal] = {tie_value.idx + offset, tie_value.score};
|
||||
}
|
||||
// Last cluster rank publishes the sums into ws->metadata.
|
||||
if (cluster_rank == kClusterSize - 1 && tx == 0) {
|
||||
const auto sum_above = prefix_above + local_above;
|
||||
const auto sum_equal = prefix_equal + local_equal;
|
||||
ws->metadata = make_uint2(sum_above, sum_equal);
|
||||
}
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void transform(TransformParams params, const void* _ws,
|
||||
void* _smem) {
|
||||
const auto ws = static_cast<const WorkSpace*>(_ws);
|
||||
const auto meta = &ws->metadata;
|
||||
const auto num_above = meta->x;
|
||||
const auto num_equal = meta->y;
|
||||
if (num_above >= K || num_equal == 0) return;
|
||||
const auto clamped_ties = min(num_equal, kMaxTies);
|
||||
tie_handle_transform(ws->ties, clamped_ties, num_above, K, params, _smem);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -0,0 +1,195 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Shared types/utilities for the three DeepSeek V4 top-k strategies
|
||||
// (Register / Streaming / Cluster). Ported from sglang's
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
inline constexpr uint32_t kMaxTopK = 1024;
|
||||
inline constexpr uint32_t kBlockSize = 1024;
|
||||
inline constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
|
||||
// 1 element per thread in the tie-breaking pass.
|
||||
inline constexpr uint32_t kMaxTies = 1024;
|
||||
inline constexpr uint32_t kRadixBins = 256;
|
||||
static_assert(kMaxTopK <= kBlockSize && kMaxTies <= kBlockSize);
|
||||
|
||||
// Always vectorize global loads as float4.
|
||||
using Vec4 = AlignedVector<float, 4>;
|
||||
|
||||
// page_to_indices: convert a flat compressed-token index into a (block * page_size + offset)
|
||||
// page-table-resolved index. page_size must be a power of 2; page_bits = log2(page_size).
|
||||
VLLM_DSV4_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table,
|
||||
uint32_t i, uint32_t page_bits) {
|
||||
const uint32_t mask = (1u << page_bits) - 1u;
|
||||
return (page_table[i >> page_bits] << page_bits) | (i & mask);
|
||||
}
|
||||
|
||||
// Output-side description of the page-table fold-in: each strategy writes
|
||||
// either (a) `transform(idx)` for entries already known to be in the top-k,
|
||||
// or (b) `write(dst, src)` for entries whose final rank is determined later.
|
||||
struct TransformParams {
|
||||
const int32_t* __restrict__ page_table;
|
||||
const int32_t* __restrict__ indices_in;
|
||||
int32_t* __restrict__ indices_out;
|
||||
uint32_t page_bits;
|
||||
|
||||
VLLM_DSV4_DEVICE void transform(uint32_t idx) const {
|
||||
indices_out[idx] = page_to_indices(page_table, indices_in[idx], page_bits);
|
||||
}
|
||||
VLLM_DSV4_DEVICE void write(uint32_t dst, uint32_t src) const {
|
||||
indices_out[dst] = page_to_indices(page_table, src, page_bits);
|
||||
}
|
||||
};
|
||||
|
||||
struct alignas(16) MatchBin {
|
||||
uint32_t bin;
|
||||
uint32_t above_count;
|
||||
uint32_t equal_count;
|
||||
};
|
||||
|
||||
struct alignas(8) Tie {
|
||||
uint32_t idx;
|
||||
float score;
|
||||
};
|
||||
|
||||
// Shared-memory layout for the final tie-breaking radix pass. Reused by both
|
||||
// the streaming kernel (overlapping `score_buffer`) and the cluster kernel.
|
||||
struct TieHandleSmem {
|
||||
alignas(128) uint32_t counter;
|
||||
alignas(128) MatchBin match;
|
||||
uint32_t histogram[kRadixBins];
|
||||
uint32_t warp_sum[kNumWarps];
|
||||
};
|
||||
|
||||
// Order-preserving fp32 -> uint key, truncated to the top kBits. Used for the
|
||||
// coarse histogram pass.
|
||||
template <uint32_t kBits>
|
||||
VLLM_DSV4_DEVICE uint32_t extract_coarse_bin(float x) {
|
||||
static_assert(0 < kBits && kBits < 15);
|
||||
__half h = __float2half_rn(x);
|
||||
uint16_t bits = __half_as_ushort(h);
|
||||
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
|
||||
: static_cast<uint16_t>(bits | 0x8000);
|
||||
return key >> (16 - kBits);
|
||||
}
|
||||
|
||||
// Full 32-bit order-preserving key, used in tie-breaking.
|
||||
VLLM_DSV4_DEVICE uint32_t extract_exact_bin(float x) {
|
||||
uint32_t bits = __float_as_uint(x);
|
||||
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
|
||||
static_assert(kWarpThreads == 32);
|
||||
#pragma unroll
|
||||
for (uint32_t offset = 1; offset < 32; offset *= 2) {
|
||||
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
|
||||
if (lane_id >= offset) val += n;
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// Fast path when seq_len <= K: identity mapping, padded to K with -1.
|
||||
VLLM_DSV4_DEVICE void trivial_transform(const TransformParams& params,
|
||||
uint32_t length, uint32_t K) {
|
||||
const auto tx = threadIdx.x;
|
||||
if (tx < length) {
|
||||
params.write(tx, tx);
|
||||
} else if (tx < K) {
|
||||
params.indices_out[tx] = -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Tie-break the threshold-bin candidates that didn't fit in the strict-above
|
||||
// region. One block-wide radix pass over the full 32-bit key (fp32 bit
|
||||
// pattern, with idx as a secondary key). Writes at most `K - num_above`
|
||||
// entries via params.write(...).
|
||||
VLLM_DSV4_DEVICE void tie_handle_transform(const Tie* __restrict__ ties,
|
||||
uint32_t num_ties, uint32_t num_above,
|
||||
uint32_t K, TransformParams params,
|
||||
void* _smem) {
|
||||
auto* smem = static_cast<TieHandleSmem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
const bool has_elem = tx < num_ties;
|
||||
const auto tie = has_elem ? ties[tx] : Tie{0, 0.0f};
|
||||
const uint32_t key = extract_exact_bin(tie.score);
|
||||
const uint32_t idx = tie.idx;
|
||||
bool active = has_elem;
|
||||
uint32_t topk_remain = K - num_above;
|
||||
uint32_t write_pos = K;
|
||||
|
||||
smem->counter = 0;
|
||||
__syncthreads();
|
||||
|
||||
// 256 bins / 32 lanes = 8 warps span the histogram inter-warp prefix.
|
||||
constexpr uint32_t kRadixWarps = kRadixBins / kWarpThreads;
|
||||
|
||||
#pragma unroll
|
||||
for (int round = 0; round < 4; round++) {
|
||||
const uint32_t shift = 24 - round * 8;
|
||||
const uint32_t bin = (key >> shift) & 0xFFu;
|
||||
|
||||
// 1. Histogram.
|
||||
if (tx < kRadixBins) smem->histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
if (active) atomicAdd(&smem->histogram[bin], 1);
|
||||
__syncthreads();
|
||||
|
||||
// 2. Two-pass prefix sum across the 256 bins.
|
||||
uint32_t hist_val = 0;
|
||||
uint32_t warp_inc = 0;
|
||||
if (tx < kRadixBins) {
|
||||
hist_val = smem->histogram[tx];
|
||||
warp_inc = warp_inclusive_sum(lane_id, hist_val);
|
||||
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tx < kRadixBins) {
|
||||
const auto tmp = (lane_id < kRadixWarps) ? smem->warp_sum[lane_id] : 0;
|
||||
const auto total = warp_reduce_sum(tmp);
|
||||
const auto inter = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
const auto prefix = inter + warp_inc;
|
||||
const auto above = total - prefix;
|
||||
// 3. Find threshold bin.
|
||||
if (above < topk_remain && above + hist_val >= topk_remain) {
|
||||
smem->match = {tx, above, topk_remain - above};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto thr = smem->match.bin;
|
||||
const auto n_above = smem->match.above_count;
|
||||
|
||||
// 4. Scatter.
|
||||
if (active) {
|
||||
if (bin > thr) {
|
||||
write_pos = num_above + atomicAdd(&smem->counter, 1);
|
||||
active = false;
|
||||
} else if (bin < thr) {
|
||||
active = false;
|
||||
} else if (round == 3) {
|
||||
write_pos = K - atomicAdd(&smem->match.equal_count, -1u);
|
||||
}
|
||||
// bin == thr && round < 3: stay active for the next radix round.
|
||||
}
|
||||
|
||||
topk_remain -= n_above;
|
||||
if (topk_remain == 0) break;
|
||||
}
|
||||
|
||||
if (write_pos < K) params.write(write_pos, idx);
|
||||
}
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -0,0 +1,66 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Thin wrappers around the CUDA PTX intrinsics used by the top-k pipeline.
|
||||
// All of these require sm_90+. Ported from sglang's
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cuda/ptx>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk::ptx {
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_init(uint64_t* addr, uint32_t arrives) {
|
||||
cuda::ptx::mbarrier_init(addr, arrives);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_arrive(uint64_t* addr) {
|
||||
cuda::ptx::mbarrier_arrive(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta,
|
||||
cuda::ptx::space_shared, addr);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_arrive_expect_tx(uint64_t* addr, uint32_t tx) {
|
||||
cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
|
||||
cuda::ptx::scope_cta,
|
||||
cuda::ptx::space_shared, addr, tx);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_wait(uint64_t* addr, uint32_t phase) {
|
||||
while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
|
||||
cuda::ptx::scope_cta, addr,
|
||||
phase))
|
||||
;
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void tma_load(void* dst, const void* src, uint32_t num_bytes,
|
||||
uint64_t* mbar) {
|
||||
cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global,
|
||||
dst, src, num_bytes, mbar);
|
||||
}
|
||||
|
||||
// elect.sync: pick a single arbitrary thread out of an active mask. Used to
|
||||
// fire a single TMA load per warp without the full ``if (tx == 0)`` cost.
|
||||
VLLM_DSV4_DEVICE uint32_t elect_sync() {
|
||||
uint32_t pred = 0;
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred %%px;\n\t"
|
||||
"elect.sync _|%%px, %1;\n\t"
|
||||
"@%%px mov.s32 %0, 1;\n\t"
|
||||
"}"
|
||||
: "+r"(pred)
|
||||
: "r"(0xFFFFFFFF));
|
||||
return pred;
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE bool elect_sync_cta(uint32_t tx) {
|
||||
const auto warp_id = tx / 32;
|
||||
const auto uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
|
||||
return (uniform_warp_id == 0 && elect_sync());
|
||||
}
|
||||
|
||||
} // namespace vllm::dsv4_topk::ptx
|
||||
@@ -0,0 +1,313 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Register-resident top-k strategy for the DeepSeek V4 indexer (small N
|
||||
// fast path). One block per row; up to ``kMax2PassLength`` scores per row
|
||||
// streamed through registers, with a single 12-bit-coarse radix pass and
|
||||
// a final tie-break round. Ported from
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
#include "ptx.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
template <uint32_t K>
|
||||
struct RegisterTopK {
|
||||
static constexpr uint32_t kHistBits = 12;
|
||||
static constexpr uint32_t kHistBins = 1 << kHistBits;
|
||||
static constexpr uint32_t kVecsPerThread = 4;
|
||||
static constexpr uint32_t kMaxTolerance = 0;
|
||||
// Length covered by registers in a single pass.
|
||||
static constexpr uint32_t kMax1PassLength = kVecsPerThread * 4 * kBlockSize;
|
||||
// Extra length staged through shared memory in the 2-pass path.
|
||||
static constexpr uint32_t kMaxExtraLength = kMax1PassLength;
|
||||
static constexpr uint32_t kMax2PassLength = kMax1PassLength + kMaxExtraLength;
|
||||
|
||||
struct Smem {
|
||||
using HistVec = AlignedVector<uint32_t, kHistBins / kBlockSize>;
|
||||
alignas(128) uint32_t counter_gt;
|
||||
alignas(128) uint32_t counter_eq;
|
||||
uint64_t mbarrier; // for the cp.async.bulk in the 2-pass path
|
||||
MatchBin match;
|
||||
uint32_t warp_sum[kNumWarps];
|
||||
union {
|
||||
uint32_t histogram[kHistBins];
|
||||
HistVec histogram_vec[kBlockSize];
|
||||
Tie tie_buffer[kMaxTies];
|
||||
};
|
||||
alignas(16) float score_buffer[kMaxExtraLength];
|
||||
};
|
||||
|
||||
template <bool kIs2Pass = false>
|
||||
VLLM_DSV4_DEVICE static void run(const float* scores, int32_t* indices,
|
||||
uint32_t length, void* _smem,
|
||||
bool use_pdl = false) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
// Init histogram + counters.
|
||||
{
|
||||
typename Smem::HistVec hist_vec;
|
||||
hist_vec.fill(0);
|
||||
smem->histogram_vec[tx] = hist_vec;
|
||||
if (tx == 0) {
|
||||
smem->counter_gt = smem->counter_eq = 0;
|
||||
if constexpr (kIs2Pass) {
|
||||
ptx::mbarrier_init(&smem->mbarrier, 1);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (use_pdl) pdl_wait_primary<true>();
|
||||
|
||||
// Stream the first `kMax1PassLength` scores into registers.
|
||||
Vec4 local[kVecsPerThread];
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
const uint32_t base = (tx + v * kBlockSize) * 4;
|
||||
if (base >= length) break;
|
||||
local[v].load(scores, tx + v * kBlockSize);
|
||||
}
|
||||
|
||||
// Issue the 2-pass TMA prefetch (next chunk of scores into smem).
|
||||
if constexpr (kIs2Pass) {
|
||||
if (ptx::elect_sync_cta(tx)) {
|
||||
const auto length_aligned = (length + 3u - kMax1PassLength) & ~3u;
|
||||
const auto size_bytes = length_aligned * sizeof(float);
|
||||
ptx::tma_load(smem->score_buffer, scores + kMax1PassLength, size_bytes,
|
||||
&smem->mbarrier);
|
||||
ptx::mbarrier_arrive_expect_tx(&smem->mbarrier, size_bytes);
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// Phase 1: histogram via shared-memory atomics.
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
#pragma unroll
|
||||
for (uint32_t e = 0; e < 4; ++e) {
|
||||
if constexpr (!kIs2Pass) {
|
||||
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
|
||||
if (idx >= length) goto LABEL_ACC_FINISH;
|
||||
}
|
||||
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local[v][e])],
|
||||
1);
|
||||
}
|
||||
}
|
||||
if constexpr (kIs2Pass) {
|
||||
if (lane_id == 0) ptx::mbarrier_wait(&smem->mbarrier, 0);
|
||||
__syncwarp();
|
||||
for (uint32_t i = tx; i + kMax1PassLength < length; i += kBlockSize) {
|
||||
const auto val = smem->score_buffer[i];
|
||||
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(val)], 1);
|
||||
}
|
||||
}
|
||||
[[maybe_unused]] LABEL_ACC_FINISH:
|
||||
__syncthreads();
|
||||
|
||||
// Phase 2: prefix scan over the histogram, locate the threshold bin.
|
||||
{
|
||||
constexpr uint32_t kItems = kHistBins / kBlockSize;
|
||||
uint32_t orig[kItems];
|
||||
const auto hist_vec = smem->histogram_vec[tx];
|
||||
uint32_t tmp_local_sum = 0;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kItems; ++i) {
|
||||
orig[i] = hist_vec[i];
|
||||
tmp_local_sum += orig[i];
|
||||
}
|
||||
|
||||
const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum);
|
||||
const auto warp_exc = warp_inc - tmp_local_sum;
|
||||
if (lane_id == kWarpThreads - 1) {
|
||||
smem->warp_sum[warp_id] = warp_inc;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const auto tmp = smem->warp_sum[lane_id];
|
||||
// Exactly one bin satisfies above < K && above + count >= K.
|
||||
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
prefix_sum += warp_exc;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kItems; ++i) {
|
||||
prefix_sum += orig[i];
|
||||
const auto above = length - prefix_sum;
|
||||
if (above < K && above + orig[i] >= K) {
|
||||
smem->match = {
|
||||
.bin = tx * kItems + i,
|
||||
.above_count = above,
|
||||
.equal_count = orig[i],
|
||||
};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const auto thr_bin = smem->match.bin;
|
||||
const auto num_above = smem->match.above_count;
|
||||
const auto num_equal = smem->match.equal_count;
|
||||
|
||||
// Phase 3: Scatter.
|
||||
// - bin > thr -> write directly to output (strictly above).
|
||||
// - bin == thr -> when no tie-break is needed, admit first-come;
|
||||
// otherwise stash into tie_buffer for phase 4.
|
||||
const bool need_tiebreak = (num_equal + num_above > K + kMaxTolerance);
|
||||
const auto topk_indices = indices;
|
||||
const auto tie_buffer = smem->tie_buffer;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
#pragma unroll
|
||||
for (uint32_t e = 0; e < 4; ++e) {
|
||||
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
|
||||
if constexpr (!kIs2Pass) {
|
||||
if (idx >= length) goto LABEL_SCATTER_DONE;
|
||||
}
|
||||
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
|
||||
if (bin > thr_bin) {
|
||||
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (need_tiebreak) {
|
||||
if (pos < kMaxTies) {
|
||||
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
|
||||
}
|
||||
} else {
|
||||
if (const auto which = pos + num_above; which < K) {
|
||||
topk_indices[which] = idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// 2-pass: pull the next chunk in from the staged smem buffer.
|
||||
if constexpr (kIs2Pass) {
|
||||
local[v].load(smem->score_buffer, tx + v * kBlockSize);
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (kIs2Pass) {
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
#pragma unroll
|
||||
for (uint32_t e = 0; e < 4; ++e) {
|
||||
const uint32_t idx =
|
||||
(tx + v * kBlockSize) * 4 + e + kMax1PassLength;
|
||||
if (idx >= length) goto LABEL_SCATTER_DONE;
|
||||
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
|
||||
if (bin > thr_bin) {
|
||||
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (need_tiebreak) {
|
||||
if (pos < kMaxTies) {
|
||||
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
|
||||
}
|
||||
} else {
|
||||
if (const auto which = pos + num_above; which < K) {
|
||||
topk_indices[which] = idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
[[maybe_unused]] LABEL_SCATTER_DONE:
|
||||
if (!need_tiebreak) return;
|
||||
|
||||
// Phase 4: tie-break within the threshold bin. We assume num_ties <=
|
||||
// kBlockSize (one block of ties), so each thread takes one tied element,
|
||||
// counts the number of tied elements with strictly higher (score, -idx),
|
||||
// and writes to output if its rank is below the remaining quota.
|
||||
__syncthreads();
|
||||
static_assert(kMaxTies <= kBlockSize);
|
||||
|
||||
const uint32_t num_ties = min(num_equal, kMaxTies);
|
||||
const uint32_t topk_remain = K - num_above;
|
||||
|
||||
const auto is_greater = [](const Tie& a, const Tie& b) {
|
||||
return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
|
||||
};
|
||||
|
||||
if (num_ties <= kWarpThreads) {
|
||||
static_assert(kWarpThreads <= kNumWarps);
|
||||
if (lane_id >= num_ties || warp_id >= num_ties) return;
|
||||
const uint32_t mask = (1ull << num_ties) - 1u;
|
||||
const auto tie = tie_buffer[lane_id];
|
||||
const auto target_tie = tie_buffer[warp_id];
|
||||
const bool pred = is_greater(tie, target_tie);
|
||||
const auto rank =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(mask, pred)));
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target_tie.idx;
|
||||
}
|
||||
} else if (num_ties <= kWarpThreads * 2) {
|
||||
// 64x64 case: each thread takes 2 elements.
|
||||
const auto lane_id_1 = lane_id + kWarpThreads;
|
||||
const auto warp_id_1 = warp_id + kWarpThreads;
|
||||
const auto invalid = Tie{.idx = 0xFFFFFFFFu, .score = -FLT_MAX};
|
||||
const auto tie_0 = tie_buffer[lane_id];
|
||||
const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid;
|
||||
{
|
||||
const auto target = tie_buffer[warp_id];
|
||||
const bool pred_0 = is_greater(tie_0, target);
|
||||
const bool pred_1 = is_greater(tie_1, target);
|
||||
const auto rank_0 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
|
||||
const auto rank_1 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
|
||||
const auto rank = rank_0 + rank_1;
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target.idx;
|
||||
}
|
||||
}
|
||||
if (warp_id_1 < num_ties) {
|
||||
const auto target = tie_buffer[warp_id_1];
|
||||
const bool pred_0 = is_greater(tie_0, target);
|
||||
const bool pred_1 = is_greater(tie_1, target);
|
||||
const auto rank_0 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
|
||||
const auto rank_1 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
|
||||
const auto rank = rank_0 + rank_1;
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target.idx;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
[[unlikely]];
|
||||
// Block-wide fallback. Rarely reached.
|
||||
for (auto i = warp_id; i < num_ties; i += kNumWarps) {
|
||||
const auto target_tie = tie_buffer[i];
|
||||
uint32_t local_rank = 0;
|
||||
for (auto j = lane_id; j < num_ties; j += kWarpThreads) {
|
||||
const auto tie = tie_buffer[j];
|
||||
if (is_greater(tie, target_tie)) local_rank++;
|
||||
}
|
||||
const auto rank = warp_reduce_sum(local_rank);
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target_tie.idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void transform(TransformParams params) {
|
||||
__syncthreads();
|
||||
if (const auto tx = threadIdx.x; tx < K) params.transform(tx);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -0,0 +1,208 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Streaming top-k strategy for medium N. Uses a TMA-driven double-buffered
|
||||
// histogram pass + scatter pass over chunks of `kSizePerStage` floats.
|
||||
// Ported from
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
#include "ptx.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
template <uint32_t K>
|
||||
struct StreamingTopK {
|
||||
static constexpr uint32_t kHistBits = 12;
|
||||
static constexpr uint32_t kHistBins = 1 << kHistBits;
|
||||
static constexpr uint32_t kElemPerStage = 8;
|
||||
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
|
||||
static constexpr uint32_t kNumStages = 2; // double buffer
|
||||
|
||||
static constexpr uint32_t kHistItems = kHistBins / kBlockSize; // 4
|
||||
static_assert(kHistItems * kBlockSize == kHistBins);
|
||||
using HistVec = AlignedVector<uint32_t, kHistItems>;
|
||||
|
||||
struct Smem {
|
||||
// [phase = 0 (histogram) | 1 (scatter)] x [buffer = 0 | 1]
|
||||
uint64_t barrier[2][kNumStages];
|
||||
alignas(128) uint32_t counter_gt;
|
||||
alignas(128) uint32_t counter_eq;
|
||||
alignas(128) MatchBin match;
|
||||
alignas(128) uint32_t warp_sum[kNumWarps];
|
||||
union {
|
||||
uint32_t histogram[kHistBins];
|
||||
HistVec histogram_vec[kBlockSize];
|
||||
Tie tie_buffer[kMaxTies];
|
||||
};
|
||||
union {
|
||||
float score_buffer[kNumStages][kSizePerStage];
|
||||
TieHandleSmem stage2; // reused for the tie-handling phase
|
||||
};
|
||||
};
|
||||
|
||||
// length must be 4-aligned (caller rounds up); TMA wants 16-byte alignment.
|
||||
template <bool kIsScatter>
|
||||
VLLM_DSV4_DEVICE static void issue_tma(const float* scores, uint32_t stage,
|
||||
uint32_t length, Smem* smem) {
|
||||
const auto buf_idx = stage % kNumStages;
|
||||
const auto offset = stage * kSizePerStage;
|
||||
const auto size = min(kSizePerStage, length - offset);
|
||||
const auto size_bytes = size * sizeof(float);
|
||||
const auto bar = &smem->barrier[kIsScatter][buf_idx];
|
||||
ptx::tma_load(smem->score_buffer[buf_idx], scores + offset, size_bytes,
|
||||
bar);
|
||||
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
|
||||
}
|
||||
|
||||
// Unified streaming pass. kIsScatter=false: build histogram (phase A).
|
||||
// kIsScatter=true: scatter using the threshold bin (phase C). Each barrier
|
||||
// is reused across iterations via the reuse-arrive pattern.
|
||||
template <bool kIsScatter>
|
||||
VLLM_DSV4_DEVICE static void stream_pass(const float* scores, uint32_t length,
|
||||
uint32_t thr_bin,
|
||||
int32_t* s_topk_indices,
|
||||
Smem* smem) {
|
||||
const auto tx = threadIdx.x;
|
||||
const auto num_iters = (length + kSizePerStage - 1) / kSizePerStage;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
|
||||
const auto length_aligned = (length + 3u) & ~3u;
|
||||
if (tx == 0) {
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumStages; i++) {
|
||||
if (i >= num_iters) break;
|
||||
issue_tma<kIsScatter>(scores, i, length_aligned, smem);
|
||||
}
|
||||
}
|
||||
|
||||
for (uint32_t iter = 0; iter < num_iters; iter++) {
|
||||
const auto buf_idx = iter % kNumStages;
|
||||
const auto offset = iter * kSizePerStage;
|
||||
const auto this_size = min(kSizePerStage, length - offset);
|
||||
|
||||
if (lane_id == 1) {
|
||||
const auto phase_bit = (iter / kNumStages) & 1;
|
||||
ptx::mbarrier_wait(&smem->barrier[kIsScatter][buf_idx], phase_bit);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemPerStage; i++) {
|
||||
const auto local_idx = tx + i * kBlockSize;
|
||||
if (local_idx >= this_size) break;
|
||||
const auto score = smem->score_buffer[buf_idx][local_idx];
|
||||
const auto bin = extract_coarse_bin<kHistBits>(score);
|
||||
if constexpr (kIsScatter) {
|
||||
const auto global_idx = offset + local_idx;
|
||||
if (bin > thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_gt, 1);
|
||||
if (pos < K) s_topk_indices[pos] = global_idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
|
||||
}
|
||||
} else {
|
||||
atomicAdd(&smem->histogram[bin], 1);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
if (tx == 0) {
|
||||
if (const auto next_iter = iter + kNumStages; next_iter < num_iters) {
|
||||
issue_tma<kIsScatter>(scores, next_iter, length_aligned, smem);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase B: locate threshold bin via warp-level prefix scan.
|
||||
VLLM_DSV4_DEVICE static void find_threshold(uint32_t length, Smem* smem) {
|
||||
const auto tx = threadIdx.x;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
uint32_t orig[kHistItems];
|
||||
const auto hist_vec = smem->histogram_vec[tx];
|
||||
uint32_t local_sum = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kHistItems; ++i) {
|
||||
orig[i] = hist_vec[i];
|
||||
local_sum += orig[i];
|
||||
}
|
||||
|
||||
const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
|
||||
const auto warp_exc = warp_inc - local_sum;
|
||||
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
|
||||
__syncthreads();
|
||||
|
||||
const auto tmp = smem->warp_sum[lane_id];
|
||||
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
prefix_sum += warp_exc;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kHistItems; ++i) {
|
||||
prefix_sum += orig[i];
|
||||
const auto above = length - prefix_sum;
|
||||
if (above < K && above + orig[i] >= K) {
|
||||
smem->match = {
|
||||
.bin = tx * kHistItems + i,
|
||||
.above_count = above,
|
||||
.equal_count = orig[i],
|
||||
};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void run(const float* scores, uint32_t length,
|
||||
int32_t* topk_indices, void* _smem) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
__builtin_assume(tx < kBlockSize);
|
||||
|
||||
{
|
||||
HistVec zero;
|
||||
zero.fill(0);
|
||||
smem->histogram_vec[tx] = zero;
|
||||
if (tx < 2 * kNumStages) {
|
||||
const auto base_barrier = &smem->barrier[0][0];
|
||||
ptx::mbarrier_init(&base_barrier[tx], 1);
|
||||
}
|
||||
if (tx == 0) {
|
||||
smem->counter_gt = 0;
|
||||
smem->counter_eq = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Phase A: histogram.
|
||||
stream_pass<false>(scores, length, 0, nullptr, smem);
|
||||
|
||||
// Phase B: threshold bin.
|
||||
find_threshold(length, smem);
|
||||
|
||||
// Phase C: scatter.
|
||||
stream_pass<true>(scores, length, smem->match.bin, topk_indices, smem);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void transform(TransformParams params, void* _smem) {
|
||||
// Phase D: page-translate above entries, then refine ties.
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto num_above = smem->match.above_count;
|
||||
if (tx < num_above) params.transform(tx);
|
||||
const auto num_equal = smem->counter_eq;
|
||||
if (num_above >= K || num_equal == 0) return;
|
||||
const auto clamped_ties = min(num_equal, kMaxTies);
|
||||
tie_handle_transform(smem->tie_buffer, clamped_ties, num_above, K, params,
|
||||
&smem->stage2);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -0,0 +1,75 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Minimal device-side utilities used by the DeepSeek V4 indexer top-k port.
|
||||
// Replaces sgl_kernel/{utils,warp,vec,type}.cuh — we only need the bits the
|
||||
// top-k kernels actually touch.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
#define VLLM_DSV4_DEVICE __forceinline__ __device__
|
||||
|
||||
inline constexpr uint32_t kWarpThreads = 32u;
|
||||
inline constexpr uint32_t kFullMask = 0xffffffffu;
|
||||
|
||||
// Programmatic Dependent Launch (sm_90+). When enabled, the kernel waits for
|
||||
// the predecessor on the same stream to advance past its dependents-launch
|
||||
// trigger before doing anything memory-dependent. Used to overlap the
|
||||
// fp8_paged_mqa_logits epilogue with the first stage of top-k.
|
||||
template <bool kUsePDL>
|
||||
VLLM_DSV4_DEVICE void pdl_wait_primary() {
|
||||
if constexpr (kUsePDL) {
|
||||
asm volatile("griddepcontrol.wait;" ::: "memory");
|
||||
}
|
||||
}
|
||||
|
||||
template <bool kUsePDL>
|
||||
VLLM_DSV4_DEVICE void pdl_trigger_secondary() {
|
||||
if constexpr (kUsePDL) {
|
||||
asm volatile("griddepcontrol.launch_dependents;" :::);
|
||||
}
|
||||
}
|
||||
|
||||
// Warp-level XOR-shuffle reduce. kThreads must be a power of 2 and <= 32.
|
||||
template <uint32_t kThreads = kWarpThreads, typename T>
|
||||
VLLM_DSV4_DEVICE T warp_reduce_sum(T value, uint32_t active_mask = kFullMask) {
|
||||
#pragma unroll
|
||||
for (auto offset = kThreads >> 1; offset > 0; offset >>= 1) {
|
||||
value = value + __shfl_xor_sync(active_mask, value, offset, 32);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
// 128-bit-aligned vector of N elements of T (N must be a power of 2, total
|
||||
// size <= 16 bytes). Used for vectorized loads/stores into shared memory.
|
||||
template <typename T, std::size_t N>
|
||||
struct alignas(sizeof(T) * N) AlignedVector {
|
||||
static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of two");
|
||||
static_assert(sizeof(T) * N <= 16,
|
||||
"AlignedVector exceeds the 128-bit CUDA vector limit");
|
||||
|
||||
T data[N];
|
||||
|
||||
VLLM_DSV4_DEVICE void load(const void* ptr, std::size_t offset = 0) {
|
||||
*reinterpret_cast<AlignedVector*>(this) =
|
||||
reinterpret_cast<const AlignedVector*>(ptr)[offset];
|
||||
}
|
||||
VLLM_DSV4_DEVICE void store(void* ptr, std::size_t offset = 0) const {
|
||||
reinterpret_cast<AlignedVector*>(ptr)[offset] = *this;
|
||||
}
|
||||
VLLM_DSV4_DEVICE void fill(T value) {
|
||||
#pragma unroll
|
||||
for (std::size_t i = 0; i < N; ++i) data[i] = value;
|
||||
}
|
||||
VLLM_DSV4_DEVICE T& operator[](std::size_t i) { return data[i]; }
|
||||
VLLM_DSV4_DEVICE const T& operator[](std::size_t i) const { return data[i]; }
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
+24
@@ -125,6 +125,30 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len);
|
||||
|
||||
// DeepSeek V4 indexer top-k (k = 512). Hopper (sm_90a) and Blackwell
|
||||
// datacenter (sm_100/sm_103) — needs thread-block clusters, TMA, and PDL.
|
||||
// Two-step API:
|
||||
// 1. fast_topk_v2_plan inspects the seq_lens distribution and writes a
|
||||
// cluster_threshold + per-row Metadata into a (B+1, 4) int32 tensor. The
|
||||
// plan is amortized when cudagraph-captured: once per shape, reused across
|
||||
// layers.
|
||||
// 2. fast_topk_v2 selects the top-512 indices per row, folds the page-table
|
||||
// gather into the radix store, and writes (B, 512) int32 page indices.
|
||||
// Dispatches per row to one of three strategies (Register / Streaming /
|
||||
// Cluster) using the planned threshold.
|
||||
//
|
||||
// Returns the size in int32s of the per-row workspace required by
|
||||
// fast_topk_v2 (allocate `(B, fast_topk_v2_workspace_ints())` int32 contig).
|
||||
void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata,
|
||||
int64_t static_cluster_threshold);
|
||||
|
||||
void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens,
|
||||
const torch::Tensor& page_table, torch::Tensor& page_indices,
|
||||
int64_t page_size, const torch::Tensor& workspace,
|
||||
const torch::Tensor& metadata);
|
||||
|
||||
int64_t fast_topk_v2_workspace_ints();
|
||||
|
||||
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& weight, torch::Tensor& scale,
|
||||
double epsilon);
|
||||
|
||||
@@ -215,6 +215,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"Tensor workspace, int k, int max_seq_len) -> ()");
|
||||
ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
|
||||
|
||||
// DeepSeek V4 indexer top-k (k=512), ported from sglang's topk_v2 family.
|
||||
// Built for sm_90a (Hopper) + sm_100a/sm_103 (Blackwell datacenter).
|
||||
// Schema only here; impl is registered in csrc/deepseek_v4/fast_topk_v2.cu
|
||||
// so it's only present when CMake compiles the source for a supported arch.
|
||||
ops.def(
|
||||
"fast_topk_v2_plan(Tensor seq_lens, Tensor! metadata, "
|
||||
"int static_cluster_threshold) -> ()");
|
||||
|
||||
ops.def(
|
||||
"fast_topk_v2(Tensor scores, Tensor seq_lens, Tensor page_table, "
|
||||
"Tensor! page_indices, int page_size, Tensor workspace, Tensor metadata)"
|
||||
" -> ()");
|
||||
|
||||
ops.def("fast_topk_v2_workspace_ints() -> int");
|
||||
|
||||
// Layernorm-quant
|
||||
// Apply Root Mean Square (RMS) Normalization to the input tensor.
|
||||
ops.def(
|
||||
|
||||
@@ -0,0 +1,356 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Correctness tests for fast_topk_v2 (DeepSeek V4 indexer top-k, k=512).
|
||||
|
||||
Run::
|
||||
|
||||
.venv/bin/python -m pytest tests/kernels/test_fast_topk_v2.py -v
|
||||
|
||||
Coverage:
|
||||
- All four execution paths: trivial (sl<=512), Register (1- and 2-pass),
|
||||
Streaming, and Cluster.
|
||||
- Both launch shapes: fused (batch<=kNumClusters=15) and two-stage (>15).
|
||||
- Mixed-length batches that exercise the per-row dispatch in the stage-2
|
||||
combine kernel.
|
||||
- Page-table fold-in: parametrised across page_size in {1, 32, 64}.
|
||||
|
||||
The kernel emits page-table-resolved indices. By using
|
||||
``page_table[b, i] = i`` with ``page_size=1`` we can compare the kernel's
|
||||
output 1:1 against ``torch.topk`` on the masked scores. For other page sizes
|
||||
the test inverts the page resolution before comparing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import (
|
||||
fast_topk_v2,
|
||||
plan_topk_v2,
|
||||
workspace_ints_per_batch,
|
||||
)
|
||||
|
||||
# Match the kernel's compile-time constant.
|
||||
TOPK = 512
|
||||
|
||||
# Thresholds inside the kernel (mirrors values in topk/register.cuh,
|
||||
# topk_v2.cuh). Keep these in sync if the kernel changes.
|
||||
SMALL_1PASS = 4 * 4 * 1024 # RegisterTopK::kMax1PassLength
|
||||
SMALL_2PASS = 2 * SMALL_1PASS # RegisterTopK::kMax2PassLength = 32768
|
||||
DEFAULT_CLUSTER_THRESHOLD = SMALL_2PASS # plan picks this for batch<=30
|
||||
NUM_CLUSTERS = 15 # kNumClusters in fast_topk_v2.cu
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _max_blocks_for(seq_len: int, page_size: int) -> int:
|
||||
return (seq_len + page_size - 1) // page_size
|
||||
|
||||
|
||||
def _trivial_page_table(batch_size: int, max_blocks: int,
|
||||
device: torch.device) -> torch.Tensor:
|
||||
"""Identity page table: page_table[b, i] = i, so page_to_indices is a no-op
|
||||
when ``page_size == 1`` (page_bits == 0)."""
|
||||
return (
|
||||
torch.arange(max_blocks, dtype=torch.int32, device=device)
|
||||
.unsqueeze(0)
|
||||
.expand(batch_size, -1)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
|
||||
def _shuffled_page_table(batch_size: int, max_blocks: int, seed: int,
|
||||
device: torch.device) -> torch.Tensor:
|
||||
"""Per-row independent permutation of [0, max_blocks)."""
|
||||
g = torch.Generator(device=device).manual_seed(seed)
|
||||
rows = []
|
||||
for _ in range(batch_size):
|
||||
rows.append(torch.randperm(max_blocks, generator=g, device=device,
|
||||
dtype=torch.int32))
|
||||
return torch.stack(rows, dim=0)
|
||||
|
||||
|
||||
def _resolve(raw_idx: int, b: int, page_table: torch.Tensor,
|
||||
page_size: int) -> int:
|
||||
"""Mirror of the device-side page_to_indices."""
|
||||
block = raw_idx // page_size
|
||||
offset = raw_idx % page_size
|
||||
return int(page_table[b, block]) * page_size + offset
|
||||
|
||||
|
||||
def _invert_resolved(resolved_idx: int, b: int, page_table: torch.Tensor,
|
||||
page_size: int) -> int:
|
||||
"""Find a raw_idx in [0, max_blocks*page_size) such that
|
||||
_resolve(raw_idx, b) == resolved_idx. Used to translate kernel output
|
||||
back to raw scores for comparison with torch.topk."""
|
||||
block = resolved_idx // page_size
|
||||
offset = resolved_idx % page_size
|
||||
# Find the row in page_table[b] that holds `block`.
|
||||
matches = (page_table[b] == block).nonzero(as_tuple=False)
|
||||
assert matches.numel() == 1, (
|
||||
f"page_table row {b} is not a permutation: block {block} appears "
|
||||
f"{matches.numel()} times")
|
||||
return int(matches.item()) * page_size + offset
|
||||
|
||||
|
||||
def _reference_topk(scores: torch.Tensor, seq_lens: torch.Tensor,
|
||||
page_table: torch.Tensor, page_size: int) -> list[set[int]]:
|
||||
"""Per-row reference: page-resolved set of indices that fast_topk_v2
|
||||
should emit (excluding -1 padding)."""
|
||||
B, _ = scores.shape
|
||||
out: list[set[int]] = []
|
||||
for b in range(B):
|
||||
sl = int(seq_lens[b])
|
||||
if sl <= TOPK:
|
||||
valid = list(range(sl))
|
||||
else:
|
||||
row = scores[b, :sl]
|
||||
_, raw = torch.topk(row, TOPK)
|
||||
valid = raw.tolist()
|
||||
out.append({_resolve(i, b, page_table, page_size) for i in valid})
|
||||
return out
|
||||
|
||||
|
||||
def _check(scores: torch.Tensor, seq_lens: torch.Tensor,
|
||||
page_table: torch.Tensor, page_size: int) -> None:
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
workspace = scores.new_empty(
|
||||
(scores.shape[0], workspace_ints_per_batch()), dtype=torch.int32)
|
||||
indices = fast_topk_v2(scores, seq_lens, page_table, page_size,
|
||||
metadata=metadata, workspace=workspace)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
expected = _reference_topk(scores, seq_lens, page_table, page_size)
|
||||
B = scores.shape[0]
|
||||
for b in range(B):
|
||||
sl = int(seq_lens[b])
|
||||
valid_count = min(sl, TOPK)
|
||||
row = indices[b].tolist()
|
||||
# Padding region: -1 (only when sl < TOPK).
|
||||
if sl < TOPK:
|
||||
assert all(v == -1 for v in row[sl:]), (
|
||||
f"row {b}: expected -1 padding after position {sl}, got "
|
||||
f"{row[sl:sl + 8]}")
|
||||
got = set(row[:valid_count])
|
||||
assert -1 not in got, f"row {b}: -1 inside valid region (sl={sl})"
|
||||
assert got == expected[b], (
|
||||
f"row {b} (sl={sl}, page_size={page_size}): "
|
||||
f"missing={len(expected[b] - got)} extra={len(got - expected[b])}")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Skip non-CUDA / non-Hopper-or-later
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _supports_clusters() -> bool:
|
||||
if not current_platform.is_cuda():
|
||||
return False
|
||||
major, _ = torch.cuda.get_device_capability()
|
||||
# Thread-block clusters / TMA / PDL are sm_90+. sm_120 (consumer
|
||||
# Blackwell) is missing some of these; skip when we detect it.
|
||||
return major == 9 or major == 10
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not _supports_clusters(),
|
||||
reason="fast_topk_v2 requires sm_90 (Hopper) or sm_100 (Blackwell DC)",
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Path coverage
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens", [
|
||||
pytest.param([1], id="trivial_1"),
|
||||
pytest.param([300], id="trivial_300"),
|
||||
pytest.param([512], id="trivial_boundary_512"),
|
||||
pytest.param([513, 600, 100, 511, 512], id="trivial_mix"),
|
||||
])
|
||||
def test_trivial_path(seq_lens):
|
||||
"""sl <= 512: identity-style fill, no radix, no tie-break."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
B = len(seq_lens)
|
||||
L = max(max(seq_lens), 1024) # round up so stride is multiple of 4
|
||||
L = (L + 3) & ~3
|
||||
seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
|
||||
scores = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device)
|
||||
_check(scores, seq_lens_t, page_table, page_size=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_len", [
|
||||
pytest.param(513, id="just_above_topk"),
|
||||
pytest.param(2048, id="2k"),
|
||||
pytest.param(SMALL_1PASS - 1, id="register_1pass_max"),
|
||||
pytest.param(SMALL_1PASS, id="register_1pass_boundary"),
|
||||
pytest.param(SMALL_1PASS + 1, id="register_2pass_first"),
|
||||
pytest.param(SMALL_2PASS - 1, id="register_2pass_max"),
|
||||
])
|
||||
def test_register_path(seq_len):
|
||||
"""Register strategy (small N; both 1- and 2-pass)."""
|
||||
torch.manual_seed(seq_len)
|
||||
device = torch.device("cuda")
|
||||
B = 4
|
||||
L = (seq_len + 3) & ~3
|
||||
scores = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device)
|
||||
page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device)
|
||||
_check(scores, seq_lens, page_table, page_size=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_len", [
|
||||
pytest.param(SMALL_2PASS, id="streaming_first"),
|
||||
pytest.param(40000, id="streaming_40k"),
|
||||
pytest.param(DEFAULT_CLUSTER_THRESHOLD, id="streaming_at_cluster_thresh"),
|
||||
])
|
||||
def test_streaming_path(seq_len):
|
||||
"""Streaming strategy (medium N). With small batch and seq_len <=
|
||||
auto-picked cluster_threshold (>= 32K when batch <= 30), the per-row
|
||||
dispatch routes here."""
|
||||
torch.manual_seed(seq_len)
|
||||
device = torch.device("cuda")
|
||||
B = 4
|
||||
L = (seq_len + 3) & ~3
|
||||
scores = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device)
|
||||
page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device)
|
||||
_check(scores, seq_lens, page_table, page_size=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size,seq_len", [
|
||||
pytest.param(2, 65536, id="cluster_fused_64k"),
|
||||
pytest.param(NUM_CLUSTERS, 131072, id="cluster_fused_max_batch"),
|
||||
pytest.param(NUM_CLUSTERS + 1, 65536, id="cluster_two_stage_just_over"),
|
||||
pytest.param(32, 96000, id="cluster_two_stage_32x96k"),
|
||||
])
|
||||
def test_cluster_path(batch_size, seq_len):
|
||||
"""Large strategy (Hopper thread-block clusters). Force seq_len above
|
||||
the auto threshold by passing static_cluster_threshold=SMALL_2PASS."""
|
||||
torch.manual_seed(seq_len * batch_size)
|
||||
device = torch.device("cuda")
|
||||
L = (seq_len + 3) & ~3
|
||||
scores = torch.randn(batch_size, L, dtype=torch.float32, device=device)
|
||||
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32,
|
||||
device=device)
|
||||
page_table = _trivial_page_table(batch_size, _max_blocks_for(L, 1), device)
|
||||
metadata = plan_topk_v2(seq_lens, static_cluster_threshold=SMALL_2PASS)
|
||||
indices = fast_topk_v2(scores, seq_lens, page_table, page_size=1,
|
||||
metadata=metadata)
|
||||
torch.cuda.synchronize()
|
||||
expected = _reference_topk(scores, seq_lens, page_table, page_size=1)
|
||||
for b in range(batch_size):
|
||||
got = set(indices[b].tolist())
|
||||
assert got == expected[b], (
|
||||
f"row {b}: missing={len(expected[b] - got)} "
|
||||
f"extra={len(got - expected[b])}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", [1, 32, 64])
|
||||
def test_page_table_fold_in(page_size):
|
||||
"""page_to_indices: kernel-side fold of the page-table gather."""
|
||||
torch.manual_seed(page_size)
|
||||
device = torch.device("cuda")
|
||||
B, seq_len = 4, 6000
|
||||
L = (seq_len + 3) & ~3
|
||||
max_blocks = (L + page_size - 1) // page_size
|
||||
scores = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device)
|
||||
page_table = _shuffled_page_table(B, max_blocks, seed=page_size,
|
||||
device=device)
|
||||
_check(scores, seq_lens, page_table, page_size=page_size)
|
||||
|
||||
|
||||
def test_mixed_lengths_route_per_row():
|
||||
"""Per-row dispatch in topk_combine_transform: trivial / Register /
|
||||
Streaming / Cluster all in one batch. Use static_cluster_threshold to
|
||||
force a mix that includes the Large path."""
|
||||
torch.manual_seed(7)
|
||||
device = torch.device("cuda")
|
||||
seq_lens = [
|
||||
100, # trivial
|
||||
SMALL_1PASS - 100, # 1-pass register
|
||||
SMALL_2PASS - 100, # 2-pass register
|
||||
50000, # streaming
|
||||
40000, # streaming
|
||||
80000, # cluster (above static_cluster_threshold)
|
||||
]
|
||||
B = len(seq_lens)
|
||||
L = (max(seq_lens) + 3) & ~3
|
||||
scores = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
|
||||
page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device)
|
||||
|
||||
# Force seq_len > 49152 to take the Cluster path.
|
||||
metadata = plan_topk_v2(seq_lens_t, static_cluster_threshold=49152)
|
||||
indices = fast_topk_v2(scores, seq_lens_t, page_table, page_size=1,
|
||||
metadata=metadata)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
expected = _reference_topk(scores, seq_lens_t, page_table, page_size=1)
|
||||
for b, sl in enumerate(seq_lens):
|
||||
valid = min(sl, TOPK)
|
||||
row = indices[b].tolist()
|
||||
if sl < TOPK:
|
||||
assert all(v == -1 for v in row[sl:])
|
||||
got = set(row[:valid])
|
||||
assert got == expected[b], f"row {b} (sl={sl}) mismatched"
|
||||
|
||||
|
||||
def test_metadata_can_be_reused_across_calls():
|
||||
"""plan_topk_v2 is amortizable: same metadata reused across calls."""
|
||||
torch.manual_seed(123)
|
||||
device = torch.device("cuda")
|
||||
B, seq_len = 8, 4096
|
||||
L = (seq_len + 3) & ~3
|
||||
seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device)
|
||||
page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device)
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
|
||||
# Two independent score buffers, same metadata.
|
||||
scores_a = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
scores_b = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
out_a = fast_topk_v2(scores_a, seq_lens, page_table, page_size=1,
|
||||
metadata=metadata)
|
||||
out_b = fast_topk_v2(scores_b, seq_lens, page_table, page_size=1,
|
||||
metadata=metadata)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
expected_a = _reference_topk(scores_a, seq_lens, page_table, page_size=1)
|
||||
expected_b = _reference_topk(scores_b, seq_lens, page_table, page_size=1)
|
||||
for b in range(B):
|
||||
assert set(out_a[b].tolist()) == expected_a[b]
|
||||
assert set(out_b[b].tolist()) == expected_b[b]
|
||||
|
||||
|
||||
def test_workspace_can_be_preallocated():
|
||||
"""Workspace passed in by the caller (cudagraph-friendly path)."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
B, seq_len = 16, 70000
|
||||
L = (seq_len + 3) & ~3
|
||||
seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device)
|
||||
page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device)
|
||||
scores = torch.randn(B, L, dtype=torch.float32, device=device)
|
||||
|
||||
metadata = plan_topk_v2(seq_lens, static_cluster_threshold=SMALL_2PASS)
|
||||
workspace = scores.new_empty((B, workspace_ints_per_batch()),
|
||||
dtype=torch.int32)
|
||||
page_indices = scores.new_empty((B, TOPK), dtype=torch.int32)
|
||||
out = fast_topk_v2(scores, seq_lens, page_table, page_size=1,
|
||||
metadata=metadata, workspace=workspace,
|
||||
page_indices=page_indices)
|
||||
torch.cuda.synchronize()
|
||||
assert out.data_ptr() == page_indices.data_ptr(), (
|
||||
"kernel must write into the caller-supplied page_indices tensor")
|
||||
expected = _reference_topk(scores, seq_lens, page_table, page_size=1)
|
||||
for b in range(B):
|
||||
assert set(out[b].tolist()) == expected[b]
|
||||
@@ -7,6 +7,7 @@ from .cache_utils import (
|
||||
dequantize_and_gather_k_cache,
|
||||
quantize_and_insert_k_cache,
|
||||
)
|
||||
from .fast_topk import fast_topk_v2, plan_topk_v2, workspace_ints_per_batch
|
||||
from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
|
||||
from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
|
||||
from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
|
||||
@@ -16,8 +17,11 @@ __all__ = [
|
||||
"combine_topk_swa_indices",
|
||||
"compute_global_topk_indices_and_lens",
|
||||
"dequantize_and_gather_k_cache",
|
||||
"fast_topk_v2",
|
||||
"fused_indexer_q_rope_quant",
|
||||
"fused_inv_rope_fp8_quant",
|
||||
"fused_q_kv_rmsnorm",
|
||||
"plan_topk_v2",
|
||||
"quantize_and_insert_k_cache",
|
||||
"workspace_ints_per_batch",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""DeepSeek V4 indexer top-k (k = 512), ported from sglang's topk_v2 family.
|
||||
|
||||
The kernel is registered as ``torch.ops._C.fast_topk_v2`` (selection +
|
||||
fused page-table gather) and ``torch.ops._C.fast_topk_v2_plan`` (per-batch
|
||||
threshold and metadata). It is built for Hopper (sm_90a) and Blackwell
|
||||
datacenter (sm_100/sm_103); sm_120 (consumer Blackwell) is not supported
|
||||
because it lacks thread-block clusters.
|
||||
|
||||
Two-step usage::
|
||||
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
page_indices = fast_topk_v2(scores, seq_lens, page_table, page_size,
|
||||
metadata=metadata)
|
||||
|
||||
The plan can be amortized across same-shape forward calls (e.g. across all
|
||||
indexer layers within one cudagraph capture), since it depends only on
|
||||
``seq_lens``.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
# The metadata layout (rows of int32x4) is fixed by the kernel; the planner
|
||||
# writes one GlobalMetadata row + one row per batch entry.
|
||||
_PLAN_COLS = 4
|
||||
|
||||
# Output top-k size. Hardcoded in the kernel.
|
||||
_TOPK = 512
|
||||
|
||||
|
||||
def workspace_ints_per_batch() -> int:
|
||||
"""Number of int32s the kernel needs in `(B, _)` workspace per row."""
|
||||
return int(torch.ops._C.fast_topk_v2_workspace_ints())
|
||||
|
||||
|
||||
def plan_topk_v2(
|
||||
seq_lens: torch.Tensor,
|
||||
static_cluster_threshold: int = 0,
|
||||
metadata: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Pick a per-batch ``cluster_threshold`` and build the per-row metadata.
|
||||
|
||||
Args:
|
||||
seq_lens: int32 tensor of shape ``(B,)``. CUDA, contiguous. Each entry
|
||||
is the number of valid scored tokens for that row (i.e., the
|
||||
indexer's compressed seq_len).
|
||||
static_cluster_threshold: when nonzero, override the auto-tuned
|
||||
threshold and route any row with ``seq_len > threshold`` through
|
||||
the cluster path. 0 means auto.
|
||||
metadata: optional preallocated int32 tensor of shape ``(B + 1, 4)``.
|
||||
When provided, the planner writes into it (useful for cudagraph
|
||||
capture). When omitted, a fresh tensor is allocated.
|
||||
|
||||
Returns:
|
||||
The metadata tensor.
|
||||
"""
|
||||
assert seq_lens.dim() == 1
|
||||
assert seq_lens.dtype == torch.int32
|
||||
assert seq_lens.is_cuda and seq_lens.is_contiguous()
|
||||
|
||||
batch_size = seq_lens.size(0)
|
||||
if metadata is None:
|
||||
metadata = torch.empty(
|
||||
(batch_size + 1, _PLAN_COLS),
|
||||
dtype=torch.int32,
|
||||
device=seq_lens.device,
|
||||
)
|
||||
else:
|
||||
assert metadata.shape == (batch_size + 1, _PLAN_COLS)
|
||||
assert metadata.dtype == torch.int32
|
||||
assert metadata.is_cuda and metadata.is_contiguous()
|
||||
|
||||
torch.ops._C.fast_topk_v2_plan(
|
||||
seq_lens, metadata, int(static_cluster_threshold)
|
||||
)
|
||||
return metadata
|
||||
|
||||
|
||||
def fast_topk_v2(
|
||||
scores: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
page_table: torch.Tensor,
|
||||
page_size: int,
|
||||
*,
|
||||
metadata: Optional[torch.Tensor] = None,
|
||||
workspace: Optional[torch.Tensor] = None,
|
||||
page_indices: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Select top-512 indexer scores per row and fold the page-table gather.
|
||||
|
||||
Args:
|
||||
scores: float32 logits of shape ``(B, L)``. ``stride(-1) == 1`` and
|
||||
``stride(0) % 4 == 0`` (TMA 16-byte alignment).
|
||||
seq_lens: int32 ``(B,)``; only the first ``seq_lens[b]`` columns are
|
||||
considered for row b.
|
||||
page_table: int32 ``(B, max_blocks)`` with ``stride(-1) == 1``.
|
||||
page_size: power-of-2 page size used by the indexer KV cache.
|
||||
metadata: optional plan tensor from :func:`plan_topk_v2`. If omitted,
|
||||
it is built on the fly.
|
||||
workspace: optional preallocated workspace ``(B, workspace_ints_per_batch())``
|
||||
int32, ``stride(-1) == 1``. Used for inter-cluster tie staging in
|
||||
the large-N path. Allocated on demand if absent.
|
||||
page_indices: optional preallocated output ``(B, 512)`` int32
|
||||
contiguous. Allocated on demand if absent.
|
||||
|
||||
Returns:
|
||||
``(B, 512)`` int32 tensor of page-table-resolved indices.
|
||||
"""
|
||||
assert scores.dim() == 2
|
||||
assert scores.dtype == torch.float32
|
||||
assert scores.is_cuda
|
||||
batch_size = scores.size(0)
|
||||
|
||||
if page_indices is None:
|
||||
page_indices = scores.new_empty(
|
||||
(batch_size, _TOPK), dtype=torch.int32
|
||||
)
|
||||
if workspace is None:
|
||||
workspace = scores.new_empty(
|
||||
(batch_size, workspace_ints_per_batch()), dtype=torch.int32
|
||||
)
|
||||
if metadata is None:
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
|
||||
torch.ops._C.fast_topk_v2(
|
||||
scores,
|
||||
seq_lens,
|
||||
page_table,
|
||||
page_indices,
|
||||
int(page_size),
|
||||
workspace,
|
||||
metadata,
|
||||
)
|
||||
return page_indices
|
||||
Reference in New Issue
Block a user