minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Kernel: optimize decoding metadata in NSA multi-spec backend with fused kernels (#17554)

Browse Source
This commit is contained in:
Johnsonms
2026-02-14 16:40:15 +08:00
committed by GitHub
parent 38473f8ee0
commit 34132d6da5
7 changed files with 2824 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh
+722
View File
@@ -0,0 +1,722 @@
/*
* Fused metadata copy kernel for NSA backend CUDA graph replay.
* JIT-compiled version for python/sglang/jit_kernel.
*
* OVERVIEW:
* This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k,
* page_table, nsa metadata, and optional FlashMLA metadata) into single kernel
* launches, significantly reducing kernel launch overhead and improving CUDA
* graph replay performance during inference.
*
* PERFORMANCE BENEFITS:
* - Single kernel launch vs. multiple separate copies (3-10x faster)
* - Optimized memory coalescing and SM utilization
* - __grid_constant__ parameter passing via constant memory
* - Especially beneficial in CUDA graph replay scenarios
*
* DESIGN:
* - Unified kernel supporting all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND)
* - Structured parameter passing (SourcePointers/DestinationPointers) for clarity
* - Template parameters (HAS_REAL_PAGE_TABLE, HAS_FLASHMLA) for compile-time optimization
* - Multi-backend variant copies to 3 destinations in one kernel (for speculative decoding)
*
* USAGE:
* This header is included by JIT compilation system. The FusedMetadataCopyKernel
* and FusedMetadataCopyMultiKernel wrapper structs provide the Python-accessible interface.
*/
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <algorithm> // for std::min
#include <cuda_runtime.h>
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/nsa_backend.py)
enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
/**
* Source pointers for metadata copy operations.
* Groups all source tensor pointers for cleaner parameter passing.
* Some pointers may be nullptr depending on forward mode and feature flags.
*/
struct SourcePointers {
const int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache
const int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths
const int32_t* __restrict__ page_indices; // page table indices
const int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths
const int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
const int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths
const int32_t* __restrict__ real_page_table; // optional real page table
const int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts
const int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata
};
/**
* Destination pointers for metadata copy operations.
* Groups all destination tensor pointers for cleaner parameter passing.
* Layout matches SourcePointers for consistency.
*/
struct DestinationPointers {
int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache
int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths
int32_t* __restrict__ page_table_1; // page table (note: different name from source)
int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths
int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths
int32_t* __restrict__ real_page_table; // optional real page table
int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts
int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata
};
/**
* Parameter structure for single-backend fused metadata copy kernel.
* Passed via __grid_constant__ for efficient constant memory access.
*/
struct FusedMetadataCopyParams {
SourcePointers src; // Source tensor pointers
DestinationPointers dst; // Destination tensor pointers
// Kernel parameters
int forward_mode; // 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
int bs; // Batch size
int max_len; // Max length for DECODE mode
int max_seqlen_k; // Max sequence length for TARGET_VERIFY/DRAFT_EXTEND
int seqlens_expanded_size; // Size of expanded sequence lengths
int page_indices_rows; // Number of rows in page_indices
int page_table_1_stride; // Stride for page_table_1
int real_page_table_cols; // Columns in real_page_table
int real_page_table_dst_stride; // Stride for destination real_page_table
int flashmla_metadata_size; // Size of FlashMLA metadata
};
/**
* Parameter structure for multi-backend fused metadata copy kernel.
* Enables copying from one source to three destinations in a single kernel launch.
* Used for speculative decoding with multiple draft backends.
*/
struct FusedMetadataCopyMultiParams {
SourcePointers src; // Source pointers (shared across all backends)
DestinationPointers dst0; // Backend 0 destination pointers
DestinationPointers dst1; // Backend 1 destination pointers
DestinationPointers dst2; // Backend 2 destination pointers
// Kernel parameters
int bs; // Batch size
int max_len; // Max length (DECODE mode only)
int seqlens_expanded_size; // Size of expanded sequence lengths
int page_table_1_stride; // Stride for page_table_1
int real_page_table_cols; // Columns in real_page_table
int real_page_table_dst_stride; // Stride for destination real_page_table
int flashmla_metadata_size; // Size of FlashMLA metadata
};
/**
* Unified kernel for all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND).
* Uses runtime branches for mode selection, with template parameters for
* compile-time optimization of optional features.
*
* DESIGN:
* - Runtime branches (forward_mode) handle mode-specific logic
* - Template parameters (HAS_*) eliminate unused feature code at compile time
* - Structured parameters (SourcePointers/DestinationPointers) passed via constant memory
*
* Used by FusedMetadataCopyKernel for single-backend metadata copy.
*
* @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support
* @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support
*/
template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
__global__ void fused_metadata_copy_kernel(const FusedMetadataCopyParams __grid_constant__ params) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = gridDim.x * blockDim.x;
// Unpack parameters for readability
const auto& src = params.src;
const auto& dst = params.dst;
const int forward_mode = params.forward_mode;
const int bs = params.bs;
const int max_len = params.max_len;
const int max_seqlen_k = params.max_seqlen_k;
const int seqlens_expanded_size = params.seqlens_expanded_size;
const int page_indices_rows = params.page_indices_rows;
const int page_table_1_stride = params.page_table_1_stride;
const int real_page_table_cols = params.real_page_table_cols;
const int real_page_table_dst_stride = params.real_page_table_dst_stride;
const int flashmla_metadata_size = params.flashmla_metadata_size;
// Copy cache_seqlens (bs elements) - common to all modes
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.cache_seqlens[i] = src.cache_seqlens[i];
}
// Copy cu_seqlens_k (skip first element) - common to all modes
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.cu_seqlens_k[i + 1] = src.cu_seqlens_k[i + 1];
}
// Branch 1: page_table copy (different dimensions per mode)
if (forward_mode == 0) { // DECODE
int page_table_elements = bs * max_len;
#pragma unroll 4
for (int i = tid; i < page_table_elements; i += total_threads) {
int row = i / max_len;
int col = i % max_len;
dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i];
}
} else { // TARGET_VERIFY or DRAFT_EXTEND
int page_table_elements = page_indices_rows * max_seqlen_k;
#pragma unroll 4
for (int i = tid; i < page_table_elements; i += total_threads) {
int row = i / max_seqlen_k;
int col = i % max_seqlen_k;
dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i];
}
}
// Branch 2: seqlens_expanded copy (only for TARGET_VERIFY/DRAFT_EXTEND)
if (forward_mode != 0) { // TARGET_VERIFY or DRAFT_EXTEND
#pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.seqlens_expanded[i] = src.seqlens_expanded[i];
}
}
// Branch 3: NSA metadata copy (different loop sizes per mode)
if (forward_mode == 0) { // DECODE
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i];
}
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1];
}
} else { // TARGET_VERIFY or DRAFT_EXTEND
#pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i];
}
#pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1];
}
}
// Copy real page table - compile-time branch
if constexpr (HAS_REAL_PAGE_TABLE) {
int real_table_elements = (forward_mode == 0 ? bs : page_indices_rows) * real_page_table_cols;
#pragma unroll 2
for (int i = tid; i < real_table_elements; i += total_threads) {
int row = i / real_page_table_cols;
int col = i % real_page_table_cols;
dst.real_page_table[row * real_page_table_dst_stride + col] =
src.real_page_table[row * real_page_table_cols + col];
}
}
// Branch 4: FlashMLA metadata copy (different sizes per mode)
if constexpr (HAS_FLASHMLA) {
int flashmla_size = (forward_mode == 0) ? (bs + 1) : (seqlens_expanded_size + 1);
if (forward_mode == 0) {
#pragma unroll 8
for (int i = tid; i < flashmla_size; i += total_threads) {
dst.flashmla_num_splits[i] = src.flashmla_num_splits[i];
}
} else {
#pragma unroll 4
for (int i = tid; i < flashmla_size; i += total_threads) {
dst.flashmla_num_splits[i] = src.flashmla_num_splits[i];
}
}
#pragma unroll 2
for (int i = tid; i < flashmla_metadata_size; i += total_threads) {
dst.flashmla_metadata[i] = src.flashmla_metadata[i];
}
}
}
/**
* Multi-backend kernel for DECODE mode.
* Copies from one source to THREE destinations in a single kernel launch.
*
* PERFORMANCE: 3x faster than three separate kernel launches due to:
* - Reduced kernel launch overhead (1 launch instead of 3)
* - Improved memory coalescing (source read once, written to 3 destinations)
* - Better instruction-level parallelism
*
* Used by FusedMetadataCopyMultiKernel for speculative decoding scenarios.
*
* @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support
* @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support
*/
template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
__global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiParams __grid_constant__ params) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = gridDim.x * blockDim.x;
// Unpack parameters for readability
const auto& src = params.src;
const auto& dst0 = params.dst0;
const auto& dst1 = params.dst1;
const auto& dst2 = params.dst2;
const int bs = params.bs;
const int max_len = params.max_len;
const int seqlens_expanded_size = params.seqlens_expanded_size;
const int page_table_1_stride = params.page_table_1_stride;
const int real_page_table_cols = params.real_page_table_cols;
const int real_page_table_dst_stride = params.real_page_table_dst_stride;
const int flashmla_metadata_size = params.flashmla_metadata_size;
// Copy cache_seqlens to all 3 backends
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.cache_seqlens[i];
dst0.cache_seqlens[i] = val;
dst1.cache_seqlens[i] = val;
dst2.cache_seqlens[i] = val;
}
// Copy cu_seqlens_k to all 3 backends (skip first element)
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.cu_seqlens_k[i + 1];
dst0.cu_seqlens_k[i + 1] = val;
dst1.cu_seqlens_k[i + 1] = val;
dst2.cu_seqlens_k[i + 1] = val;
}
// DECODE mode: copy page_table_1 to all 3 backends
int page_table_elements = bs * max_len;
#pragma unroll 4
for (int i = tid; i < page_table_elements; i += total_threads) {
int row = i / max_len;
int col = i % max_len;
int32_t val = src.page_indices[i];
dst0.page_table_1[row * page_table_1_stride + col] = val;
dst1.page_table_1[row * page_table_1_stride + col] = val;
dst2.page_table_1[row * page_table_1_stride + col] = val;
}
// Copy nsa_cache_seqlens to all 3 backends
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.nsa_cache_seqlens[i];
dst0.nsa_cache_seqlens[i] = val;
dst1.nsa_cache_seqlens[i] = val;
dst2.nsa_cache_seqlens[i] = val;
}
// Copy NSA cu_seqlens to all 3 backends
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.nsa_cu_seqlens_k[i + 1];
dst0.nsa_cu_seqlens_k[i + 1] = val;
dst1.nsa_cu_seqlens_k[i + 1] = val;
dst2.nsa_cu_seqlens_k[i + 1] = val;
}
// Copy real page table to all 3 backends
if (src.real_page_table != nullptr && dst0.real_page_table != nullptr) {
int real_table_elements = bs * real_page_table_cols;
#pragma unroll 2
for (int i = tid; i < real_table_elements; i += total_threads) {
int row = i / real_page_table_cols;
int col = i % real_page_table_cols;
int src_idx = row * real_page_table_cols + col;
int dst_idx = row * real_page_table_dst_stride + col;
int32_t val = src.real_page_table[src_idx];
dst0.real_page_table[dst_idx] = val;
dst1.real_page_table[dst_idx] = val;
dst2.real_page_table[dst_idx] = val;
}
}
// Copy FlashMLA metadata to all 3 backends
if constexpr (HAS_FLASHMLA) {
int flashmla_size = bs + 1;
#pragma unroll 8
for (int i = tid; i < flashmla_size; i += total_threads) {
int32_t val = src.flashmla_num_splits[i];
dst0.flashmla_num_splits[i] = val;
dst1.flashmla_num_splits[i] = val;
dst2.flashmla_num_splits[i] = val;
}
#pragma unroll 2
for (int i = tid; i < flashmla_metadata_size; i += total_threads) {
int32_t val = src.flashmla_metadata[i];
dst0.flashmla_metadata[i] = val;
dst1.flashmla_metadata[i] = val;
dst2.flashmla_metadata[i] = val;
}
}
}
// ============================================================================
// Host-side launcher wrappers for JIT compilation
// ============================================================================
namespace {
// Launch configuration constants
constexpr int THREADS_PER_BLOCK = 256;
constexpr int MAX_GRID_SIZE = 1024; // Limit to prevent excessive resource usage
/**
* Helper function to extract a typed data pointer from a TensorView.
* Performs runtime type checking and returns the properly cast pointer.
*
* @tparam T The expected element type (e.g., int32_t)
* @param tensor The TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed pointer to the tensor data
*/
template <typename T>
inline const T* unwrap_data_ptr(const tvm::ffi::TensorView& tensor, const char* name) {
using namespace host;
if (tensor.data_ptr()) {
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
}
return static_cast<const T*>(tensor.data_ptr());
}
/**
* Helper function to extract a typed mutable data pointer from a TensorView.
* Performs runtime type checking and returns the properly cast pointer.
*
* @tparam T The expected element type (e.g., int32_t)
* @param tensor The TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed mutable pointer to the tensor data
*/
template <typename T>
inline T* unwrap_data_ptr_mut(const tvm::ffi::TensorView& tensor, const char* name) {
using namespace host;
if (tensor.data_ptr()) {
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
}
return static_cast<T*>(tensor.data_ptr());
}
/**
* Helper function to extract a typed data pointer from an Optional TensorView.
* Returns nullptr if the optional has no value, otherwise performs type checking.
*
* @tparam T The expected element type (e.g., int32_t)
* @param optional_tensor The Optional TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed pointer to the tensor data, or nullptr if optional has no value
*/
template <typename T>
inline const T*
unwrap_optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& optional_tensor, const char* name) {
using namespace host;
if (!optional_tensor.has_value()) {
return nullptr;
}
const auto& tensor = optional_tensor.value();
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
return static_cast<const T*>(tensor.data_ptr());
}
/**
* Helper function to extract a typed mutable data pointer from an Optional TensorView.
* Returns nullptr if the optional has no value, otherwise performs type checking.
*
* @tparam T The expected element type (e.g., int32_t)
* @param optional_tensor The Optional TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed mutable pointer to the tensor data, or nullptr if optional has no value
*/
template <typename T>
inline T*
unwrap_optional_data_ptr_mut(const tvm::ffi::Optional<tvm::ffi::TensorView>& optional_tensor, const char* name) {
using namespace host;
if (!optional_tensor.has_value()) {
return nullptr;
}
const auto& tensor = optional_tensor.value();
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
return static_cast<T*>(tensor.data_ptr());
}
/**
* Calculate kernel launch configuration.
*
* @param total_work Total number of work items
* @param threads_per_block Threads per block (default: THREADS_PER_BLOCK)
* @return Grid dimension for kernel launch
*/
inline dim3 get_launch_config(int total_work, int threads_per_block = THREADS_PER_BLOCK) {
int num_blocks = (total_work + threads_per_block - 1) / threads_per_block;
// Limit grid size to prevent excessive resource usage while ensuring coverage
num_blocks = std::min(num_blocks, MAX_GRID_SIZE);
return dim3(num_blocks);
}
/**
* JIT wrapper for single-backend fused metadata copy kernel.
*
* This struct provides a unified interface for launching the fused metadata copy
* kernel with different forward modes. It constructs the parameter struct and
* launches the unified kernel.
*
* IMPLEMENTATION:
* - Extracts raw pointers from TensorView objects
* - Constructs FusedMetadataCopyParams with nested SourcePointers/DestinationPointers
* - Calculates grid configuration based on maximum work size
* - Launches fused_metadata_copy_kernel with __grid_constant__ parameters
*
* @tparam FORWARD_MODE Forward mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
* @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present
* @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present
*/
template <int FORWARD_MODE, bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
struct FusedMetadataCopyKernel {
static_assert(
FORWARD_MODE >= 0 && FORWARD_MODE <= 2,
"FORWARD_MODE must be 0 (DECODE), 1 (TARGET_VERIFY), or 2 (DRAFT_EXTEND)");
static void
run(const tvm::ffi::TensorView cache_seqlens_src,
const tvm::ffi::TensorView cu_seqlens_k_src,
const tvm::ffi::TensorView page_indices_src,
const tvm::ffi::TensorView nsa_cache_seqlens_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst,
const tvm::ffi::TensorView cu_seqlens_k_dst,
const tvm::ffi::TensorView page_table_1_dst,
const tvm::ffi::TensorView nsa_cache_seqlens_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_dst,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst,
int bs,
int max_len,
int max_seqlen_k,
int seqlens_expanded_size) {
using namespace host;
// Build parameter struct with nested source/destination pointers
// unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation
const auto params = FusedMetadataCopyParams{
.src =
{
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
.seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"),
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
.flashmla_metadata = unwrap_optional_data_ptr<int32_t>(flashmla_metadata_src, "flashmla_metadata_src"),
},
.dst =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"),
.seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"),
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst, "flashmla_metadata_dst"),
},
.forward_mode = FORWARD_MODE,
.bs = bs,
.max_len = max_len,
.max_seqlen_k = max_seqlen_k,
.seqlens_expanded_size = seqlens_expanded_size,
.page_indices_rows = static_cast<int>(page_indices_src.shape()[0]),
.page_table_1_stride = static_cast<int>(page_table_1_dst.shape()[1]),
.real_page_table_cols =
real_page_table_src.has_value() ? static_cast<int>(real_page_table_src.value().shape()[1]) : 0,
.real_page_table_dst_stride =
real_page_table_dst.has_value() ? static_cast<int>(real_page_table_dst.value().stride(0)) : 0,
.flashmla_metadata_size =
flashmla_metadata_src.has_value() ? static_cast<int>(flashmla_metadata_src.value().numel()) : 0,
};
// Calculate grid configuration
int max_elements = std::max(
{bs,
params.page_indices_rows * max_seqlen_k,
seqlens_expanded_size,
HAS_FLASHMLA ? (seqlens_expanded_size + 1) : 0,
HAS_FLASHMLA ? params.flashmla_metadata_size : 0});
dim3 grid = get_launch_config(max_elements);
dim3 block(THREADS_PER_BLOCK);
DLDevice device = cache_seqlens_src.device();
// Launch unified kernel with params struct
host::LaunchKernel(grid, block, device)(fused_metadata_copy_kernel<HAS_REAL_PAGE_TABLE, HAS_FLASHMLA>, params);
}
};
/**
* JIT wrapper for multi-backend fused metadata copy kernel.
*
* This kernel optimizes the common case where metadata needs to be copied from
* one source to THREE destination backends in a single kernel launch. This is
* 3x faster than launching three separate kernels due to:
* - Reduced kernel launch overhead (1 launch instead of 3)
* - Improved memory coalescing (source read once, written to 3 destinations)
* - Better GPU occupancy and instruction-level parallelism
*
* USAGE: Primarily for speculative decoding with multiple draft models, where
* the same source metadata needs to be replicated to multiple backend contexts.
*
* LIMITATION: Currently only supports DECODE mode, which is the most frequently
* used mode in speculative decoding scenarios.
*
* IMPLEMENTATION:
* - Constructs FusedMetadataCopyMultiParams with 1 SourcePointers + 3 DestinationPointers
* - Launches fused_metadata_copy_multi_kernel with __grid_constant__ parameters
*
* @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present
* @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present
*/
template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
struct FusedMetadataCopyMultiKernel {
static void
run(const tvm::ffi::TensorView cache_seqlens_src,
const tvm::ffi::TensorView cu_seqlens_k_src,
const tvm::ffi::TensorView page_indices_src,
const tvm::ffi::TensorView nsa_cache_seqlens_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst0,
const tvm::ffi::TensorView cu_seqlens_k_dst0,
const tvm::ffi::TensorView page_table_1_dst0,
const tvm::ffi::TensorView nsa_cache_seqlens_dst0,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst0,
const tvm::ffi::TensorView cache_seqlens_dst1,
const tvm::ffi::TensorView cu_seqlens_k_dst1,
const tvm::ffi::TensorView page_table_1_dst1,
const tvm::ffi::TensorView nsa_cache_seqlens_dst1,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst1,
const tvm::ffi::TensorView cache_seqlens_dst2,
const tvm::ffi::TensorView cu_seqlens_k_dst2,
const tvm::ffi::TensorView page_table_1_dst2,
const tvm::ffi::TensorView nsa_cache_seqlens_dst2,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2,
int bs,
int max_len,
int seqlens_expanded_size) {
using namespace host;
// Build parameter struct with nested source/destination pointers
// unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation
const auto params = FusedMetadataCopyMultiParams{
.src =
{
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
.seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
.flashmla_metadata = unwrap_optional_data_ptr<int32_t>(flashmla_metadata_src, "flashmla_metadata_src"),
},
.dst0 =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst0, "flashmla_metadata_dst0"),
},
.dst1 =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst1, "flashmla_metadata_dst1"),
},
.dst2 =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst2, "flashmla_metadata_dst2"),
},
.bs = bs,
.max_len = max_len,
.seqlens_expanded_size = seqlens_expanded_size,
.page_table_1_stride = static_cast<int>(page_table_1_dst0.shape()[1]),
.real_page_table_cols =
real_page_table_src.has_value() ? static_cast<int>(real_page_table_src.value().shape()[1]) : 0,
.real_page_table_dst_stride =
real_page_table_dst0.has_value() ? static_cast<int>(real_page_table_dst0.value().stride(0)) : 0,
.flashmla_metadata_size =
flashmla_metadata_src.has_value() ? static_cast<int>(flashmla_metadata_src.value().numel()) : 0,
};
dim3 grid = get_launch_config(bs * max_len);
dim3 block(THREADS_PER_BLOCK);
DLDevice device = cache_seqlens_src.device();
// Launch multi-backend kernel with params struct
host::LaunchKernel(grid, block, device)(
fused_metadata_copy_multi_kernel<HAS_REAL_PAGE_TABLE, HAS_FLASHMLA>, params);
}
};
} // namespace
python/sglang/jit_kernel/fused_metadata_copy.py
+316
View File
@@ -0,0 +1,316 @@
"""
Fused metadata copy kernel for NSA backend CUDA graph replay.
This module provides JIT-compiled CUDA kernels for fusing multiple tensor
copy operations into single kernel launches, reducing kernel launch overhead
and improving CUDA graph replay performance.
The kernels are compiled on-demand using TVM FFI and cached for subsequent use.
"""
from __future__ import annotations
import logging
from typing import Optional
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
logger = logging.getLogger(__name__)
# ============================================================================
# JIT Module Compilation
# ============================================================================
@cache_once
def _jit_fused_metadata_copy_module(
forward_mode: int, has_real_page_table: bool, has_flashmla: bool
):
"""Compile JIT module for single-backend fused metadata copy.
Args:
forward_mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
has_real_page_table: Whether real_page_table tensors are used
has_flashmla: Whether FlashMLA metadata tensors are used
"""
args = make_cpp_args(forward_mode, has_real_page_table, has_flashmla)
try:
return load_jit(
"fused_metadata_copy",
*args,
cuda_files=["elementwise/fused_metadata_copy.cuh"],
cuda_wrappers=[
(
"fused_metadata_copy",
f"FusedMetadataCopyKernel<{args}>::run",
)
],
)
except Exception as e:
logger.error(
f"Failed to compile JIT fused metadata copy kernel "
f"(forward_mode={forward_mode}, has_real_page_table={has_real_page_table}, "
f"has_flashmla={has_flashmla}): {e}"
)
raise
@cache_once
def _jit_fused_metadata_copy_multi_module(
has_real_page_table: bool, has_flashmla: bool
):
"""Compile JIT module for multi-backend fused metadata copy (DECODE mode only).
Args:
has_real_page_table: Whether real_page_table tensors are used
has_flashmla: Whether FlashMLA metadata tensors are used
"""
args = make_cpp_args(has_real_page_table, has_flashmla)
try:
return load_jit(
"fused_metadata_copy_multi",
*args,
cuda_files=["elementwise/fused_metadata_copy.cuh"],
cuda_wrappers=[
(
"fused_metadata_copy_multi",
f"FusedMetadataCopyMultiKernel<{args}>::run",
)
],
)
except Exception as e:
logger.error(
f"Failed to compile JIT fused metadata copy multi kernel "
f"(has_real_page_table={has_real_page_table}, has_flashmla={has_flashmla}): {e}"
)
raise
# ============================================================================
# Public API
# ============================================================================
def fused_metadata_copy_cuda(
cache_seqlens_src: torch.Tensor,
cu_seqlens_k_src: torch.Tensor,
page_indices_src: torch.Tensor,
nsa_cache_seqlens_src: torch.Tensor,
seqlens_expanded_src: Optional[torch.Tensor],
nsa_cu_seqlens_k_src: torch.Tensor,
real_page_table_src: Optional[torch.Tensor],
flashmla_num_splits_src: Optional[torch.Tensor],
flashmla_metadata_src: Optional[torch.Tensor],
cache_seqlens_dst: torch.Tensor,
cu_seqlens_k_dst: torch.Tensor,
page_table_1_dst: torch.Tensor,
nsa_cache_seqlens_dst: torch.Tensor,
seqlens_expanded_dst: Optional[torch.Tensor],
nsa_cu_seqlens_k_dst: torch.Tensor,
real_page_table_dst: Optional[torch.Tensor],
flashmla_num_splits_dst: Optional[torch.Tensor],
flashmla_metadata_dst: Optional[torch.Tensor],
forward_mode: int,
bs: int,
max_len: int,
max_seqlen_k: int,
seqlens_expanded_size: int,
) -> None:
"""
Fused metadata copy kernel for NSA backend CUDA graph replay.
This function fuses multiple tensor copy operations into a single kernel launch,
reducing kernel launch overhead and improving performance.
Args:
cache_seqlens_src: Source cache sequence lengths [bs]
cu_seqlens_k_src: Source cumulative sequence lengths [bs+1]
page_indices_src: Source page indices [rows, max_len]
nsa_cache_seqlens_src: Source NSA cache sequence lengths [size]
seqlens_expanded_src: Optional source expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND)
nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [size+1]
real_page_table_src: Optional source real page table [rows, cols]
flashmla_num_splits_src: Optional source FlashMLA num_splits [size+1]
flashmla_metadata_src: Optional source FlashMLA metadata tensor
cache_seqlens_dst: Destination cache sequence lengths [bs]
cu_seqlens_k_dst: Destination cumulative sequence lengths [bs+1]
page_table_1_dst: Destination page table [rows, stride]
nsa_cache_seqlens_dst: Destination NSA cache sequence lengths [size]
seqlens_expanded_dst: Optional destination expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND)
nsa_cu_seqlens_k_dst: Destination NSA cumulative sequence lengths [size+1]
real_page_table_dst: Optional destination real page table [rows, cols]
flashmla_num_splits_dst: Optional destination FlashMLA num_splits [size+1]
flashmla_metadata_dst: Optional destination FlashMLA metadata tensor
forward_mode: Forward mode (0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND)
bs: Batch size
max_len: Maximum length for decode/draft_extend mode
max_seqlen_k: Maximum sequence length for target_verify mode
seqlens_expanded_size: Size of expanded sequence lengths
"""
# Determine template parameters for kernel specialization
has_real_page_table = real_page_table_src is not None
has_flashmla = flashmla_num_splits_src is not None
# Get JIT-compiled module for this configuration (cached after first use)
module = _jit_fused_metadata_copy_module(
forward_mode, has_real_page_table, has_flashmla
)
# Ensure all required source tensors are contiguous (required for kernel's linear indexing)
# This matches the CHECK_INPUT checks in the verified sgl-kernel implementation
cache_seqlens_src = cache_seqlens_src.contiguous()
cu_seqlens_k_src = cu_seqlens_k_src.contiguous()
page_indices_src = page_indices_src.contiguous()
nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous()
if seqlens_expanded_src is not None:
seqlens_expanded_src = seqlens_expanded_src.contiguous()
nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous()
# Call JIT-compiled kernel (None values are passed as Optional with no value)
module.fused_metadata_copy(
cache_seqlens_src,
cu_seqlens_k_src,
page_indices_src,
nsa_cache_seqlens_src,
seqlens_expanded_src,
nsa_cu_seqlens_k_src,
real_page_table_src,
flashmla_num_splits_src,
flashmla_metadata_src,
cache_seqlens_dst,
cu_seqlens_k_dst,
page_table_1_dst,
nsa_cache_seqlens_dst,
seqlens_expanded_dst,
nsa_cu_seqlens_k_dst,
real_page_table_dst,
flashmla_num_splits_dst,
flashmla_metadata_dst,
bs,
max_len,
max_seqlen_k,
seqlens_expanded_size,
)
def fused_metadata_copy_multi_cuda(
cache_seqlens_src: torch.Tensor,
cu_seqlens_k_src: torch.Tensor,
page_indices_src: torch.Tensor,
nsa_cache_seqlens_src: torch.Tensor,
nsa_cu_seqlens_k_src: torch.Tensor,
real_page_table_src: Optional[torch.Tensor],
flashmla_num_splits_src: Optional[torch.Tensor],
flashmla_metadata_src: Optional[torch.Tensor],
cache_seqlens_dst0: torch.Tensor,
cu_seqlens_k_dst0: torch.Tensor,
page_table_1_dst0: torch.Tensor,
nsa_cache_seqlens_dst0: torch.Tensor,
nsa_cu_seqlens_k_dst0: torch.Tensor,
real_page_table_dst0: Optional[torch.Tensor],
flashmla_num_splits_dst0: Optional[torch.Tensor],
flashmla_metadata_dst0: Optional[torch.Tensor],
cache_seqlens_dst1: torch.Tensor,
cu_seqlens_k_dst1: torch.Tensor,
page_table_1_dst1: torch.Tensor,
nsa_cache_seqlens_dst1: torch.Tensor,
nsa_cu_seqlens_k_dst1: torch.Tensor,
real_page_table_dst1: Optional[torch.Tensor],
flashmla_num_splits_dst1: Optional[torch.Tensor],
flashmla_metadata_dst1: Optional[torch.Tensor],
cache_seqlens_dst2: torch.Tensor,
cu_seqlens_k_dst2: torch.Tensor,
page_table_1_dst2: torch.Tensor,
nsa_cache_seqlens_dst2: torch.Tensor,
nsa_cu_seqlens_k_dst2: torch.Tensor,
real_page_table_dst2: Optional[torch.Tensor],
flashmla_num_splits_dst2: Optional[torch.Tensor],
flashmla_metadata_dst2: Optional[torch.Tensor],
bs: int,
max_len: int,
seqlens_expanded_size: int,
) -> None:
"""
Multi-backend fused metadata copy kernel for NSA backend CUDA graph replay.
This function copies metadata from one source to THREE destinations in a single
kernel launch, eliminating the overhead of 3 separate kernel calls. Currently
only supports DECODE mode, which is the most common case.
Args:
cache_seqlens_src: Source cache sequence lengths [bs]
cu_seqlens_k_src: Source cumulative sequence lengths [bs+1]
page_indices_src: Source page indices [bs, max_len]
nsa_cache_seqlens_src: Source NSA cache sequence lengths [bs]
nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [bs+1]
real_page_table_src: Optional source real page table [bs, cols]
flashmla_num_splits_src: Optional source FlashMLA num_splits [bs+1]
flashmla_metadata_src: Optional source FlashMLA metadata tensor
cache_seqlens_dst0-2: Destination cache sequence lengths for backends 0-2
cu_seqlens_k_dst0-2: Destination cumulative sequence lengths for backends 0-2
page_table_1_dst0-2: Destination page tables for backends 0-2
nsa_cache_seqlens_dst0-2: Destination NSA cache sequence lengths for backends 0-2
nsa_cu_seqlens_k_dst0-2: Destination NSA cumulative sequence lengths for backends 0-2
real_page_table_dst0-2: Optional destination real page tables for backends 0-2
flashmla_num_splits_dst0-2: Optional destination FlashMLA num_splits for backends 0-2
flashmla_metadata_dst0-2: Optional destination FlashMLA metadata tensors for backends 0-2
bs: Batch size
max_len: Maximum length for decode mode
seqlens_expanded_size: Size of expanded sequence lengths
"""
# Determine template parameters for kernel specialization
has_real_page_table = real_page_table_src is not None
has_flashmla = flashmla_num_splits_src is not None
# Get JIT-compiled module for this configuration (cached after first use)
module = _jit_fused_metadata_copy_multi_module(has_real_page_table, has_flashmla)
# Ensure all source tensors are contiguous (required for kernel's linear indexing)
# This matches the CHECK_INPUT checks in the verified sgl-kernel implementation
cache_seqlens_src = cache_seqlens_src.contiguous()
cu_seqlens_k_src = cu_seqlens_k_src.contiguous()
page_indices_src = page_indices_src.contiguous()
nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous()
nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous()
# Call JIT-compiled kernel (None values are passed as Optional with no value)
module.fused_metadata_copy_multi(
cache_seqlens_src,
cu_seqlens_k_src,
page_indices_src,
nsa_cache_seqlens_src,
nsa_cu_seqlens_k_src,
real_page_table_src,
flashmla_num_splits_src,
flashmla_metadata_src,
cache_seqlens_dst0,
cu_seqlens_k_dst0,
page_table_1_dst0,
nsa_cache_seqlens_dst0,
nsa_cu_seqlens_k_dst0,
real_page_table_dst0,
flashmla_num_splits_dst0,
flashmla_metadata_dst0,
cache_seqlens_dst1,
cu_seqlens_k_dst1,
page_table_1_dst1,
nsa_cache_seqlens_dst1,
nsa_cu_seqlens_k_dst1,
real_page_table_dst1,
flashmla_num_splits_dst1,
flashmla_metadata_dst1,
cache_seqlens_dst2,
cu_seqlens_k_dst2,
page_table_1_dst2,
nsa_cache_seqlens_dst2,
nsa_cu_seqlens_k_dst2,
real_page_table_dst2,
flashmla_num_splits_dst2,
flashmla_metadata_dst2,
bs,
max_len,
seqlens_expanded_size,
)
python/sglang/jit_kernel/tests/test_fused_metadata_copy.py
+1067
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/environ.py
+2
View File
@@ -362,6 +362,8 @@ class Envs:
# NSA Backend # NSA Backend
SGLANG_NSA_FUSE_TOPK = EnvBool(True) SGLANG_NSA_FUSE_TOPK = EnvBool(True)
SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True) SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True)
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_VERIFY_FUSED_METADATA_COPY = EnvBool(False)
SGLANG_NSA_FORCE_MLA = EnvBool(False) SGLANG_NSA_FORCE_MLA = EnvBool(False)
# sgl-kernel # sgl-kernel
python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py
+3 -3
View File
@@ -127,7 +127,7 @@ class NativeSparseAttnBackendMTPPrecomputeMixin:
cu_seqlens_k = compute_cu_seqlens(cache_seqlens) cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Get page indices from cache # Get page indices from cache
page_indices = self.req_to_token[req_pool_indices, :max_len] page_indices = self.req_to_token[req_pool_indices, :max_len].contiguous()
# Compute NSA seqlens # Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens( nsa_cache_seqlens = compute_nsa_seqlens(
@@ -187,7 +187,7 @@ class NativeSparseAttnBackendMTPPrecomputeMixin:
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave( page_indices = torch.repeat_interleave(
page_indices, repeats=self.speculative_num_draft_tokens, dim=0 page_indices, repeats=self.speculative_num_draft_tokens, dim=0
) ).contiguous()
# Generate expanded seqlens # Generate expanded seqlens
extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs
@@ -269,7 +269,7 @@ class NativeSparseAttnBackendMTPPrecomputeMixin:
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave( page_indices = torch.repeat_interleave(
page_indices, repeats=extend_seq_lens, dim=0 page_indices, repeats=extend_seq_lens, dim=0
) ).contiguous()
# Generate expanded seqlens # Generate expanded seqlens
seqlens_expanded = torch.cat( seqlens_expanded = torch.cat(
python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py
+407
View File
@@ -0,0 +1,407 @@
"""
Verification utilities for NSA backend fused metadata copy operations.
This module contains verification code to ensure that fused metadata copy kernels
produce the same results as individual copy operations.
"""
import torch
def verify_single_backend_fused_metadata_copy(
metadata,
precomputed,
forward_mode,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
flashmla_num_splits_dst=None,
flashmla_metadata_dst=None,
):
"""
Verify that the fused metadata copy kernel produces the same results as individual copies.
Args:
metadata: The NSA metadata object containing destination tensors
precomputed: The precomputed metadata containing source tensors
forward_mode: The forward mode (decode, target_verify, or draft_extend)
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
flashmla_num_splits_dst: Destination FlashMLA num_splits tensor (optional)
flashmla_metadata_dst: Destination FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_nsa_cache_seqlens = metadata.nsa_cache_seqlens_int32.clone()
fused_nsa_seqlens_expanded = metadata.nsa_seqlens_expanded.clone()
fused_nsa_cu_seqlens_k = metadata.nsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = flashmla_num_splits_dst.clone()
fused_flashmla_metadata = flashmla_metadata_dst.clone()
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_nsa_cache_seqlens = torch.zeros_like(metadata.nsa_cache_seqlens_int32)
ref_nsa_seqlens_expanded = torch.zeros_like(metadata.nsa_seqlens_expanded)
ref_nsa_cu_seqlens_k = torch.zeros_like(metadata.nsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(flashmla_num_splits_dst)
ref_flashmla_metadata = torch.zeros_like(flashmla_metadata_dst)
# Run individual copy operations (reference implementation)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
if forward_mode.is_decode_or_idle():
# Decode mode
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_nsa_cache_seqlens.copy_(precomputed.nsa_cache_seqlens)
elif forward_mode.is_target_verify():
# Target verify mode
ref_page_table_1[:, : precomputed.max_seqlen_k].copy_(precomputed.page_indices)
ref_nsa_seqlens_expanded.copy_(precomputed.seqlens_expanded)
ref_nsa_cache_seqlens.copy_(precomputed.nsa_cache_seqlens)
elif forward_mode.is_draft_extend():
# Draft extend mode
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
ref_page_table_1[:rows, :cols].copy_(precomputed.page_indices)
size = precomputed.seqlens_expanded_size
ref_nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded)
ref_nsa_cache_seqlens[:size].copy_(precomputed.nsa_cache_seqlens)
# Copy NSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_nsa_cu_seqlens_k[1 : 1 + size].copy_(precomputed.nsa_cu_seqlens_k[1 : 1 + size])
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
ref_flashmla_num_splits[: size + 1].copy_(flashmla_num_splits_src[: size + 1])
ref_flashmla_metadata.copy_(flashmla_metadata_src)
# Compare results and crash if inconsistent
def check_tensor_equal(name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Forward mode: {forward_mode}, bs={bs}\n"
f"The fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors (only compare the slices that were actually updated)
check_tensor_equal("cache_seqlens", fused_cache_seqlens, ref_cache_seqlens)
check_tensor_equal("cu_seqlens_k", fused_cu_seqlens_k, ref_cu_seqlens_k)
# Compare page_table_1 only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_len],
ref_page_table_1[:, : precomputed.max_len],
)
elif forward_mode.is_target_verify():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_seqlen_k],
ref_page_table_1[:, : precomputed.max_seqlen_k],
)
elif forward_mode.is_draft_extend():
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
check_tensor_equal(
"page_table_1",
fused_page_table_1[:rows, :cols],
ref_page_table_1[:rows, :cols],
)
# Compare nsa_cache_seqlens only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"nsa_cache_seqlens",
fused_nsa_cache_seqlens,
ref_nsa_cache_seqlens,
)
else: # TARGET_VERIFY or DRAFT_EXTEND
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"nsa_cache_seqlens",
fused_nsa_cache_seqlens[:size],
ref_nsa_cache_seqlens[:size],
)
# Compare nsa_seqlens_expanded only for TARGET_VERIFY and DRAFT_EXTEND
if forward_mode.is_target_verify() or forward_mode.is_draft_extend():
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"nsa_seqlens_expanded",
fused_nsa_seqlens_expanded[:size],
ref_nsa_seqlens_expanded[:size],
)
# Compare nsa_cu_seqlens_k only for the region that was updated
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"nsa_cu_seqlens_k",
fused_nsa_cu_seqlens_k[: 1 + size],
ref_nsa_cu_seqlens_k[: 1 + size],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
"real_page_table",
fused_real_page_table[:rows, :cols],
ref_real_page_table[:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"flashmla_num_splits",
fused_flashmla_num_splits[: size + 1],
ref_flashmla_num_splits[: size + 1],
)
check_tensor_equal(
"flashmla_metadata",
fused_flashmla_metadata,
ref_flashmla_metadata,
)
def verify_multi_backend_fused_metadata_copy(
metadata0,
metadata1,
metadata2,
precomputed,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
):
"""
Verify that the multi-backend fused metadata copy kernel produces the same results
as individual copies for all three backends.
Args:
metadata0: The NSA metadata object for backend 0
metadata1: The NSA metadata object for backend 1
metadata2: The NSA metadata object for backend 2
precomputed: The precomputed metadata containing source tensors
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_results = []
for idx, metadata in enumerate([metadata0, metadata1, metadata2]):
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_nsa_cache_seqlens = metadata.nsa_cache_seqlens_int32.clone()
fused_nsa_cu_seqlens_k = metadata.nsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = metadata.flashmla_metadata.num_splits.clone()
fused_flashmla_metadata = (
metadata.flashmla_metadata.flashmla_metadata.clone()
)
fused_results.append(
{
"cache_seqlens": fused_cache_seqlens,
"cu_seqlens_k": fused_cu_seqlens_k,
"page_table_1": fused_page_table_1,
"nsa_cache_seqlens": fused_nsa_cache_seqlens,
"nsa_cu_seqlens_k": fused_nsa_cu_seqlens_k,
"real_page_table": fused_real_page_table,
"flashmla_num_splits": fused_flashmla_num_splits,
"flashmla_metadata": fused_flashmla_metadata,
}
)
# Run individual copy operations for each backend (reference implementation)
ref_results = []
for idx in range(3):
metadata = [metadata0, metadata1, metadata2][idx]
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_nsa_cache_seqlens = torch.zeros_like(metadata.nsa_cache_seqlens_int32)
ref_nsa_cu_seqlens_k = torch.zeros_like(metadata.nsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(
metadata.flashmla_metadata.num_splits
)
ref_flashmla_metadata = torch.zeros_like(
metadata.flashmla_metadata.flashmla_metadata
)
# Copy operations (decode mode)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_nsa_cache_seqlens.copy_(precomputed.nsa_cache_seqlens)
# Copy NSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_nsa_cu_seqlens_k[1 : 1 + size].copy_(
precomputed.nsa_cu_seqlens_k[1 : 1 + size]
)
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits[: size + 1].copy_(
flashmla_num_splits_src[: size + 1]
)
ref_flashmla_metadata.copy_(flashmla_metadata_src)
ref_results.append(
{
"cache_seqlens": ref_cache_seqlens,
"cu_seqlens_k": ref_cu_seqlens_k,
"page_table_1": ref_page_table_1,
"nsa_cache_seqlens": ref_nsa_cache_seqlens,
"nsa_cu_seqlens_k": ref_nsa_cu_seqlens_k,
"real_page_table": ref_real_page_table,
"flashmla_num_splits": ref_flashmla_num_splits,
"flashmla_metadata": ref_flashmla_metadata,
}
)
# Compare results for all 3 backends
def check_tensor_equal(backend_idx, name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"MULTI-BACKEND FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Backend: {backend_idx}\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Batch size: {bs}\n"
f"The multi-backend fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors for all 3 backends (multi-backend is DECODE mode only)
for idx in range(3):
fused = fused_results[idx]
ref = ref_results[idx]
check_tensor_equal(
idx,
"cache_seqlens",
fused["cache_seqlens"],
ref["cache_seqlens"],
)
check_tensor_equal(
idx,
"cu_seqlens_k",
fused["cu_seqlens_k"],
ref["cu_seqlens_k"],
)
# Multi-backend is DECODE mode only, so compare only [:, :max_len]
check_tensor_equal(
idx,
"page_table_1",
fused["page_table_1"][:, : precomputed.max_len],
ref["page_table_1"][:, : precomputed.max_len],
)
check_tensor_equal(
idx,
"nsa_cache_seqlens",
fused["nsa_cache_seqlens"],
ref["nsa_cache_seqlens"],
)
# DECODE mode uses bs for nsa_cu_seqlens_k size
check_tensor_equal(
idx,
"nsa_cu_seqlens_k",
fused["nsa_cu_seqlens_k"][: bs + 1],
ref["nsa_cu_seqlens_k"][: bs + 1],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
idx,
"real_page_table",
fused["real_page_table"][:rows, :cols],
ref["real_page_table"][:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
# DECODE mode uses bs + 1 for flashmla_num_splits
check_tensor_equal(
idx,
"flashmla_num_splits",
fused["flashmla_num_splits"][: bs + 1],
ref["flashmla_num_splits"][: bs + 1],
)
check_tensor_equal(
idx,
"flashmla_metadata",
fused["flashmla_metadata"],
ref["flashmla_metadata"],
)
python/sglang/srt/layers/attention/nsa_backend.py
+307 -51
View File
@@ -16,6 +16,10 @@ from sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute import (
compute_cu_seqlens, compute_cu_seqlens,
) )
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.attention.nsa.nsa_mtp_verification import (
verify_multi_backend_fused_metadata_copy,
verify_single_backend_fused_metadata_copy,
)
from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache
from sglang.srt.layers.attention.nsa.transform_index import ( from sglang.srt.layers.attention.nsa.transform_index import (
transform_index_page_table_decode, transform_index_page_table_decode,
@@ -63,6 +67,15 @@ else:
# Reuse this workspace buffer across all NSA backend instances # Reuse this workspace buffer across all NSA backend instances
global_workspace_buffer = None global_workspace_buffer = None
# Control whether to use fused metadata copy kernel (default: enabled)
# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable
_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get()
# Control whether to verify fused metadata copy against individual copies (default: disabled)
# Set SGLANG_VERIFY_FUSED_METADATA_COPY=1 or true to enable verification
# This will crash with detailed error message if any inconsistency is detected
_VERIFY_FUSED_METADATA_COPY = envs.SGLANG_VERIFY_FUSED_METADATA_COPY.get()
@dataclass(frozen=True) @dataclass(frozen=True)
class NSAFlashMLAMetadata: class NSAFlashMLAMetadata:
@@ -1127,55 +1140,150 @@ class NativeSparseAttnBackend(
metadata = self.decode_cuda_graph_metadata[bs] metadata = self.decode_cuda_graph_metadata[bs]
# Copy basic seqlens # Track whether fused kernel succeeded
metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens) fused_kernel_succeeded = False
metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
# Mode-specific copy logic # Use fused CUDA kernel for all copy operations
if forward_mode.is_decode_or_idle(): if _USE_FUSED_METADATA_COPY:
# Decode mode try:
metadata.page_table_1[:, : precomputed.max_len].copy_( from sglang.jit_kernel.fused_metadata_copy import (
precomputed.page_indices fused_metadata_copy_cuda,
) )
metadata.nsa_cache_seqlens_int32.copy_(precomputed.nsa_cache_seqlens)
# seqlens_expanded is same as cache_seqlens (already copied)
elif forward_mode.is_target_verify(): # Map forward_mode to integer enum
# Target verify mode if forward_mode.is_decode_or_idle():
metadata.page_table_1[:, : precomputed.max_seqlen_k].copy_( mode_int = 0 # DECODE
precomputed.page_indices elif forward_mode.is_target_verify():
) mode_int = 1 # TARGET_VERIFY
metadata.nsa_seqlens_expanded.copy_(precomputed.seqlens_expanded) elif forward_mode.is_draft_extend():
metadata.nsa_cache_seqlens_int32.copy_(precomputed.nsa_cache_seqlens) mode_int = 2 # DRAFT_EXTEND
else:
raise ValueError(f"Unsupported forward_mode: {forward_mode}")
elif forward_mode.is_draft_extend(): # Prepare FlashMLA tensors if needed
# Draft extend mode flashmla_num_splits_src = None
rows = precomputed.page_indices.shape[0] flashmla_num_splits_dst = None
cols = precomputed.max_seqlen_k flashmla_metadata_src = None
metadata.page_table_1[:rows, :cols].copy_(precomputed.page_indices) flashmla_metadata_dst = None
if precomputed.flashmla_metadata is not None:
flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits
flashmla_num_splits_dst = metadata.flashmla_metadata.num_splits
flashmla_metadata_src = (
precomputed.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst = metadata.flashmla_metadata.flashmla_metadata
# Call fused kernel
fused_metadata_copy_cuda(
# Source tensors
precomputed.cache_seqlens,
precomputed.cu_seqlens_k,
precomputed.page_indices,
precomputed.nsa_cache_seqlens,
precomputed.seqlens_expanded,
precomputed.nsa_cu_seqlens_k,
precomputed.real_page_table,
flashmla_num_splits_src,
flashmla_metadata_src,
# Destination tensors
metadata.cache_seqlens_int32,
metadata.cu_seqlens_k,
metadata.page_table_1,
metadata.nsa_cache_seqlens_int32,
metadata.nsa_seqlens_expanded,
metadata.nsa_cu_seqlens_k,
(
metadata.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst,
flashmla_metadata_dst,
# Parameters
mode_int,
bs,
precomputed.max_len,
precomputed.max_seqlen_k,
precomputed.seqlens_expanded_size,
)
# Successfully used fused kernel
fused_kernel_succeeded = True
# Verification: compare fused kernel results against individual copies
if _VERIFY_FUSED_METADATA_COPY:
verify_single_backend_fused_metadata_copy(
metadata=metadata,
precomputed=precomputed,
forward_mode=forward_mode,
bs=bs,
flashmla_num_splits_src=flashmla_num_splits_src,
flashmla_metadata_src=flashmla_metadata_src,
flashmla_num_splits_dst=flashmla_num_splits_dst,
flashmla_metadata_dst=flashmla_metadata_dst,
)
except ImportError:
print(
"Warning: Fused metadata copy kernel not available, falling back to individual copies."
)
except Exception as e:
print(
f"Warning: Fused metadata copy kernel failed with error: {e}, falling back to individual copies."
)
# Fallback to individual copy operations if fused kernel disabled or failed
if not fused_kernel_succeeded:
# Copy basic seqlens
metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens)
metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
# Mode-specific copy logic
if forward_mode.is_decode_or_idle():
# Decode mode
metadata.page_table_1[:, : precomputed.max_len].copy_(
precomputed.page_indices
)
metadata.nsa_cache_seqlens_int32.copy_(precomputed.nsa_cache_seqlens)
# seqlens_expanded is same as cache_seqlens (already copied)
elif forward_mode.is_target_verify():
# Target verify mode
metadata.page_table_1[:, : precomputed.max_seqlen_k].copy_(
precomputed.page_indices
)
metadata.nsa_seqlens_expanded.copy_(precomputed.seqlens_expanded)
metadata.nsa_cache_seqlens_int32.copy_(precomputed.nsa_cache_seqlens)
elif forward_mode.is_draft_extend():
# Draft extend mode
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
metadata.page_table_1[:rows, :cols].copy_(precomputed.page_indices)
size = precomputed.seqlens_expanded_size
metadata.nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded)
metadata.nsa_cache_seqlens_int32[:size].copy_(
precomputed.nsa_cache_seqlens
)
# Copy NSA cu_seqlens
size = precomputed.seqlens_expanded_size size = precomputed.seqlens_expanded_size
metadata.nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded) metadata.nsa_cu_seqlens_k[1 : 1 + size].copy_(
metadata.nsa_cache_seqlens_int32[:size].copy_(precomputed.nsa_cache_seqlens) precomputed.nsa_cu_seqlens_k[1 : 1 + size]
)
# Copy NSA cu_seqlens # Copy real page table
size = precomputed.seqlens_expanded_size if precomputed.real_page_table is not None:
metadata.nsa_cu_seqlens_k[1 : 1 + size].copy_( rows, cols = precomputed.real_page_table.shape
precomputed.nsa_cu_seqlens_k[1 : 1 + size] metadata.real_page_table[:rows, :cols].copy_(
) precomputed.real_page_table
)
# Copy real page table # Copy FlashMLA metadata in fallback path
if precomputed.real_page_table is not None: if precomputed.flashmla_metadata is not None:
rows, cols = precomputed.real_page_table.shape size = precomputed.seqlens_expanded_size
metadata.real_page_table[:rows, :cols].copy_(precomputed.real_page_table) flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1))
else: flashmla_metadata.copy_(precomputed.flashmla_metadata)
# real_page_table is same as page_table_1 (already copied)
pass
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1))
flashmla_metadata.copy_(precomputed.flashmla_metadata)
self.forward_metadata = metadata self.forward_metadata = metadata
@@ -1958,15 +2066,163 @@ class NativeSparseAttnMultiStepBackend:
spec_info=forward_batch.spec_info, spec_info=forward_batch.spec_info,
) )
# Fast copy to each backend (1-2x faster than computing N times) # Use multi-backend fused copy when we have 3 or more backends
for i in range(self.speculative_num_steps): # This is 3x faster than calling the single-backend copy 3 times
self.attn_backends[ if self.speculative_num_steps >= 3:
i try:
].init_forward_metadata_replay_cuda_graph_from_precomputed( from sglang.jit_kernel.fused_metadata_copy import (
bs=bs, fused_metadata_copy_multi_cuda,
precomputed=precomputed, )
forward_mode=ForwardMode.DECODE,
) metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs]
metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs]
metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs]
# Set nsa_prefill_impl for first 3 backends (required by the method)
for i in range(3):
self.attn_backends[i].set_nsa_prefill_impl(forward_batch=None)
# Prepare FlashMLA tensors if needed
flashmla_num_splits_src = None
flashmla_metadata_src = None
flashmla_num_splits_dst0 = None
flashmla_num_splits_dst1 = None
flashmla_num_splits_dst2 = None
flashmla_metadata_dst0 = None
flashmla_metadata_dst1 = None
flashmla_metadata_dst2 = None
if precomputed.flashmla_metadata is not None:
flashmla_num_splits_src = (
precomputed.flashmla_metadata.num_splits
)
flashmla_metadata_src = (
precomputed.flashmla_metadata.flashmla_metadata
)
flashmla_num_splits_dst0 = (
metadata0.flashmla_metadata.num_splits
)
flashmla_num_splits_dst1 = (
metadata1.flashmla_metadata.num_splits
)
flashmla_num_splits_dst2 = (
metadata2.flashmla_metadata.num_splits
)
flashmla_metadata_dst0 = (
metadata0.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst1 = (
metadata1.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst2 = (
metadata2.flashmla_metadata.flashmla_metadata
)
# Call the multi-backend fused kernel for first 3 backends
fused_metadata_copy_multi_cuda(
# Source tensors
precomputed.cache_seqlens,
precomputed.cu_seqlens_k,
precomputed.page_indices,
precomputed.nsa_cache_seqlens,
precomputed.nsa_cu_seqlens_k,
precomputed.real_page_table,
flashmla_num_splits_src,
flashmla_metadata_src,
# Destination tensors for backend 0
metadata0.cache_seqlens_int32,
metadata0.cu_seqlens_k,
metadata0.page_table_1,
metadata0.nsa_cache_seqlens_int32,
metadata0.nsa_cu_seqlens_k,
(
metadata0.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst0,
flashmla_metadata_dst0,
# Destination tensors for backend 1
metadata1.cache_seqlens_int32,
metadata1.cu_seqlens_k,
metadata1.page_table_1,
metadata1.nsa_cache_seqlens_int32,
metadata1.nsa_cu_seqlens_k,
(
metadata1.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst1,
flashmla_metadata_dst1,
# Destination tensors for backend 2
metadata2.cache_seqlens_int32,
metadata2.cu_seqlens_k,
metadata2.page_table_1,
metadata2.nsa_cache_seqlens_int32,
metadata2.nsa_cu_seqlens_k,
(
metadata2.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst2,
flashmla_metadata_dst2,
# Parameters
bs,
precomputed.max_len,
precomputed.seqlens_expanded_size,
)
# Verification: compare fused kernel results against individual copies
if _VERIFY_FUSED_METADATA_COPY:
verify_multi_backend_fused_metadata_copy(
metadata0=metadata0,
metadata1=metadata1,
metadata2=metadata2,
precomputed=precomputed,
bs=bs,
flashmla_num_splits_src=flashmla_num_splits_src,
flashmla_metadata_src=flashmla_metadata_src,
)
# Copy remaining backends one by one (if > 3 backends)
for i in range(3, self.speculative_num_steps):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
)
except (ImportError, Exception) as e:
# Fallback to loop if multi-backend kernel not available or fails
if isinstance(e, ImportError):
print(
"Warning: Multi-backend fused metadata copy kernel not available, falling back to loop."
)
else:
print(
f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop."
)
for i in range(self.speculative_num_steps):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
)
else:
# Less than 3 backends: copy to each backend individually
for i in range(self.speculative_num_steps):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
)
else: else:
# Fallback: compute metadata separately for each backend # Fallback: compute metadata separately for each backend
for i in range(self.speculative_num_steps): for i in range(self.speculative_num_steps):