[Feat][NVIDIA] DeepSeek-V4.1: Prefill support CP with multistream forward prepare (4/N)

python/sglang/kernels/jit/csrc/deepseek_v4/prefill_candidate_topk.cuh

@@ -0,0 +1,137 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/deepseek_v4/topk_impl.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang {
+namespace prefill_candidate {
+namespace impl = device::topk;
+
+struct Params {
+  const float* scores;
+  const int32_t* lengths;
+  const bool* candidates;
+  const int32_t* offsets;
+  int32_t* output;
+  int64_t score_stride;
+  int64_t candidate_stride;
+  uint32_t block_size;
+  uint32_t topk;
+};
+
+struct ScoreFilter {
+  const bool* blocks;
+  uint32_t block_size;
+
+  SGL_DEVICE float operator()(float value, uint32_t index) const {
+    return blocks[index / block_size] ? value : -impl::infinity_value();
+  }
+};
+
+template <bool kPDL>
+__global__ __launch_bounds__(impl::TopKConfig::kBlockSize, impl::TopKConfig::kOccupancy) void select(
+    const __grid_constant__ Params params) {
+  device::enable_smem_spilling();
+  const auto row = blockIdx.x;
+  const auto length = static_cast<uint32_t>(params.lengths[row]);
+  const auto offset = params.offsets[row];
+  const auto scores = params.scores + row * params.score_stride;
+  const auto candidates = params.candidates + row * params.candidate_stride;
+  const auto output = params.output + row * static_cast<int64_t>(params.topk);
+  const ScoreFilter filter{candidates, params.block_size};
+
+  if (length <= params.topk) {
+    device::PDLWaitPrimary<kPDL>();
+    for (uint32_t i = threadIdx.x; i < params.topk; i += blockDim.x) {
+      const bool valid = i < length && candidates[i / params.block_size] && scores[i] > -impl::infinity_value();
+      output[i] = valid ? static_cast<int32_t>(i) + offset : -1;
+    }
+    return;
+  }
+
+  using Register2 = impl::TopKRegister<2>;
+  using Register4 = impl::TopKRegister<4>;
+  using Streaming = impl::TopKStreaming;
+  const impl::TopKProblem problem{
+      .in = scores,
+      .out = output,
+      .topk = params.topk,
+      .seq_len = length,
+      .bias = impl::broadcast(offset),
+  };
+  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
+  if (length <= Register2::kMaxSeqLen) {
+    Register2::forward<kPDL>(problem, &smem, filter);
+  } else if (length <= Register4::kMaxSeqLen) {
+    Register4::forward<kPDL>(problem, &smem, filter);
+  } else {
+    Streaming::forward<kPDL>(problem, &smem, filter);
+  }
+  // Selection can underfill with -inf positions. Match mask_topk_scores in the
+  // same CTA, after every thread has finished publishing its selected indices.
+  __syncthreads();
+  for (uint32_t i = threadIdx.x; i < params.topk; i += blockDim.x) {
+    const int64_t index = static_cast<int64_t>(output[i]) - offset;
+    const bool valid = index >= 0 && index < length && candidates[index / params.block_size] &&
+                       scores[index] > -impl::infinity_value();
+    if (!valid) output[i] = -1;
+  }
+}
+
+template <bool kPDL>
+struct Kernel {
+  static void
+  run(const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView lengths,
+      const tvm::ffi::TensorView candidates,
+      const tvm::ffi::TensorView offsets,
+      const tvm::ffi::TensorView output,
+      int64_t block_size) {
+    using namespace host;
+    auto B = SymbolicSize{"rows"};
+    auto W = SymbolicSize{"width"};
+    auto S = SymbolicSize{"score_stride"};
+    auto C = SymbolicSize{"candidate_width"};
+    auto CS = SymbolicSize{"candidate_stride"};
+    auto K = SymbolicSize{"topk"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLGPU>();
+    TensorMatcher({B, W}).with_strides({S, 1}).with_dtype<float>().with_device(dev).verify(scores);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(lengths);
+    TensorMatcher({B, C}).with_strides({CS, 1}).with_device(dev).verify(candidates);
+    RuntimeCheck(
+        candidates.dtype().code == kDLBool && candidates.dtype().bits == 8 && candidates.dtype().lanes == 1,
+        "candidate mask must have dtype bool");
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(offsets);
+    TensorMatcher({B, K}).with_dtype<int32_t>().with_device(dev).verify(output);
+    RuntimeCheck(block_size > 0, "candidate block size must be positive");
+    RuntimeCheck(C.unwrap() >= (W.unwrap() + block_size - 1) / block_size, "candidate mask is too narrow");
+    RuntimeCheck(S.unwrap() % 4 == 0, "score rows must be 16-byte aligned");
+    RuntimeCheck(K.unwrap() > 0 && K.unwrap() <= impl::TopKConfig::kMaxTopK, "topk must be in (0, 2048]");
+    if (B.unwrap() == 0) return;
+    const Params params{
+        .scores = static_cast<const float*>(scores.data_ptr()),
+        .lengths = static_cast<const int32_t*>(lengths.data_ptr()),
+        .candidates = static_cast<const bool*>(candidates.data_ptr()),
+        .offsets = static_cast<const int32_t*>(offsets.data_ptr()),
+        .output = static_cast<int32_t*>(output.data_ptr()),
+        .score_stride = S.unwrap(),
+        .candidate_stride = CS.unwrap(),
+        .block_size = static_cast<uint32_t>(block_size),
+        .topk = static_cast<uint32_t>(K.unwrap()),
+    };
+    LaunchKernel(static_cast<uint32_t>(B.unwrap()), impl::TopKConfig::kBlockSize, dev.unwrap())
+        .config({.use_pdl = kPDL})
+        .launch(select<kPDL>, params);
+  }
+};
+}  // namespace prefill_candidate
+}  // namespace sglang

python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh

@@ -533,14 +533,20 @@ struct TopKRadixBase : TopKConfig {
 // wastes work on shorter sequences.
 // ---------------------------------------------------------------------------
 
+struct IdentityScoreFilter {
+  SGL_DEVICE float operator()(float value, uint32_t) const {
+    return value;
+  }
+};
+
 template <uint32_t kLocalVecs_>
 struct TopKRegister : TopKRadixBase<12> {
   static constexpr uint32_t kLocalVecs = kLocalVecs_;
   static constexpr uint32_t kMaxSeqLen = kBlockSize * kVecSize * kLocalVecs;
   using Smem = typename TopKRadixBase<12>::Smem;
 
-  template <bool kUsePDL>
-  SGL_DEVICE static void forward(const TopKProblem& problem, void* _smem) {
+  template <bool kUsePDL, typename ScoreFilter = IdentityScoreFilter>
+  SGL_DEVICE static void forward(const TopKProblem& problem, void* _smem, ScoreFilter filter = {}) {
     const auto tx = threadIdx.x;
     const auto smem = static_cast<Smem*>(_smem);
 
@@ -575,6 +581,8 @@ struct TopKRegister : TopKRadixBase<12> {
       }
 #pragma unroll
       for (uint32_t j = 0; j < kVecSize; ++j) {
+        const auto idx = vi * kVecSize + j;
+        if (idx < problem.seq_len) local_vecs[i][j] = filter(local_vecs[i][j], idx);
         atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local_vecs[i][j])], 1);
       }
     }
@@ -640,8 +648,8 @@ struct TopKStreaming : TopKRadixBase<12> {
  public:
   static constexpr uint32_t kMaxSeqLen = std::numeric_limits<uint32_t>::max();
 
-  template <bool kUsePDL>
-  SGL_DEVICE static void forward(TopKProblem problem, void* _smem) {
+  template <bool kUsePDL, typename ScoreFilter = IdentityScoreFilter>
+  SGL_DEVICE static void forward(TopKProblem problem, void* _smem, ScoreFilter filter = {}) {
     const auto tx = threadIdx.x;
     const auto smem = static_cast<Smem*>(_smem);
 
@@ -654,7 +662,8 @@ struct TopKStreaming : TopKRadixBase<12> {
     PDLWaitPrimary<kUsePDL>();
 
     // Phase 1: Load and build histogram
-    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) {
+    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+      val = filter(val, idx);
       const auto bin = extract_coarse_bin<kHistBits>(val);
       atomicAdd(&smem->histogram[bin], 1);
     });
@@ -682,6 +691,7 @@ struct TopKStreaming : TopKRadixBase<12> {
     const auto v_lo = smem->v_lo;
     const auto topk = problem.topk;
     for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+      val = filter(val, idx);
       if (val >= v_hi) {
         const auto pos = atomicAdd(&smem->count_gt, 1);
         if (pos < topk) [[likely]] {

python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py

@@ -399,6 +399,69 @@ def _store_fp4_index_k_cache_kernel(
     )
 
 
+@triton.jit
+def _gather_fp4_index_k_cache_masked_kernel(
+    cache,
+    req_to_token,
+    req_ids,
+    req_lens,
+    out_payload,
+    out_scales,
+    req_stride: tl.constexpr,
+    cache_stride: tl.constexpr,
+    width: tl.constexpr,
+    compress_ratio: tl.constexpr,
+    page_size: tl.constexpr,
+    block_rows: tl.constexpr,
+):
+    """Gather only the valid compressed-K prefix of each fixed graph slot."""
+    request = tl.program_id(0)
+    row_block = tl.program_id(1)
+    logical_k = row_block * block_rows + tl.arange(0, block_rows)
+    valid = logical_k < tl.load(req_lens + request) // compress_ratio
+
+    req_id = tl.load(req_ids + request)
+    full_slot = tl.load(
+        req_to_token + req_id * req_stride + logical_k.to(tl.int64) * compress_ratio,
+        mask=valid,
+        other=0,
+    )
+    compressed_slot = full_slot // compress_ratio
+    page = compressed_slot // page_size
+    page_offset = compressed_slot - page * page_size
+    out_row = request * width + logical_k
+
+    payload_cols = tl.arange(0, INDEX_K_PAYLOAD_BYTES)
+    payload = tl.load(
+        cache
+        + page[:, None] * cache_stride
+        + page_offset[:, None] * INDEX_K_PAYLOAD_BYTES
+        + payload_cols[None, :],
+        mask=valid[:, None],
+        other=0,
+    )
+    tl.store(
+        out_payload
+        + out_row[:, None].to(tl.int64) * INDEX_K_PAYLOAD_BYTES
+        + payload_cols[None, :],
+        payload,
+        mask=valid[:, None],
+    )
+
+    scale_cols = tl.arange(0, INDEX_K_SCALE_BYTES)
+    scale_bytes = tl.load(
+        cache
+        + page[:, None] * cache_stride
+        + page_size * INDEX_K_PAYLOAD_BYTES
+        + page_offset[:, None] * INDEX_K_SCALE_BYTES
+        + scale_cols[None, :],
+        mask=valid[:, None],
+        other=0,
+    ).to(tl.int32)
+    packed_scale = tl.sum(scale_bytes << (scale_cols[None, :] * 8), axis=1)
+    tl.store(out_scales + out_row, packed_scale, mask=valid)
+
+
 def quantize_fp4_indexer_tensor(
     x: torch.Tensor, rne: bool = False
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -461,6 +524,57 @@ def store_fp4_index_k_cache(
     )
 
 
+def gather_fp4_index_k_cache_masked(
+    cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    req_ids: torch.Tensor,
+    req_lens: torch.Tensor,
+    *,
+    width: int,
+    compress_ratio: int,
+    page_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack graph-stable contiguous FP4 K without touching padded K slots.
+
+    The result reserves ``width`` rows for every request slot because
+    DeepGEMM needs fixed offsets under CUDA graph replay. Only
+    ``req_lens[r] // compress_ratio`` rows are populated. DeepGEMM and top-k
+    consume the same real lengths, so the untouched suffix is never read.
+    """
+    assert cache.ndim == 2 and cache.element_size() == 1
+    assert cache.shape[1] == page_size * INDEX_K_SLOT_BYTES
+    assert req_to_token.ndim == 2 and req_to_token.dtype == torch.int32
+    assert req_ids.ndim == 1 and req_lens.shape == req_ids.shape
+    assert width > 0 and compress_ratio > 0
+
+    num_rows = req_ids.shape[0] * width
+    payload = torch.empty(
+        (num_rows, INDEX_K_PAYLOAD_BYTES.value),
+        dtype=torch.int8,
+        device=cache.device,
+    )
+    scales = torch.empty((num_rows,), dtype=torch.int32, device=cache.device)
+    block_rows = 32
+    _gather_fp4_index_k_cache_masked_kernel[
+        (req_ids.shape[0], triton.cdiv(width, block_rows))
+    ](
+        cache,
+        req_to_token,
+        req_ids,
+        req_lens,
+        payload,
+        scales,
+        req_to_token.stride(0),
+        cache.stride(0),
+        width,
+        compress_ratio,
+        page_size,
+        block_rows,
+        num_warps=8,
+    )
+    return payload, scales
+
+
 @triton.jit
 def _gather_fp4_index_k_kernel(
     cache,

python/sglang/kernels/ops/attention/dsv4/prefill_candidates.py

@@ -0,0 +1,42 @@
+"""Fused candidate-aware top-k for the dense FP4 prefill indexer."""
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+from .utils import make_name
+
+
+@cache_once
+def _candidate_topk_module():
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        make_name("prefill_candidate_topk"),
+        *args,
+        cuda_files=["deepseek_v4/prefill_candidate_topk.cuh"],
+        cuda_wrappers=[("run", f"prefill_candidate::Kernel<{args}>::run")],
+    )
+
+
+def topk_prefill_candidates(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_mask: torch.Tensor,
+    block_size: int,
+    out_offsets: torch.Tensor,
+    out_indices: torch.Tensor,
+) -> None:
+    """Select within each query's candidate blocks, without modifying scores.
+
+    Columns are request-local, seq_lens are nonnegative and bounded by the score
+    width. Output indices include out_offsets; invalid slots are -1. Row strides
+    must be multiples of four floats for the vectorized top-k reads.
+    """
+    _candidate_topk_module().run(
+        scores, seq_lens, block_mask, out_offsets, out_indices, block_size
+    )