Kernel: optimize decoding metadata in NSA multi-spec backend with fused kernels (#17554)
This commit is contained in:
Johnsonms
@@ -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
|
||||
316
python/sglang/jit_kernel/fused_metadata_copy.py
Normal file
316
python/sglang/jit_kernel/fused_metadata_copy.py
Normal 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,
|
||||
)
|
||||
1067
python/sglang/jit_kernel/tests/test_fused_metadata_copy.py
Normal file
1067
python/sglang/jit_kernel/tests/test_fused_metadata_copy.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -362,6 +362,8 @@ class Envs:
|
||||
# NSA Backend
|
||||
SGLANG_NSA_FUSE_TOPK = 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)
|
||||
|
||||
# sgl-kernel
|
||||
|
||||
@@ -127,7 +127,7 @@ class NativeSparseAttnBackendMTPPrecomputeMixin:
|
||||
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
|
||||
|
||||
# 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
|
||||
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 = torch.repeat_interleave(
|
||||
page_indices, repeats=self.speculative_num_draft_tokens, dim=0
|
||||
)
|
||||
).contiguous()
|
||||
|
||||
# Generate expanded seqlens
|
||||
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 = torch.repeat_interleave(
|
||||
page_indices, repeats=extend_seq_lens, dim=0
|
||||
)
|
||||
).contiguous()
|
||||
|
||||
# Generate expanded seqlens
|
||||
seqlens_expanded = torch.cat(
|
||||
|
||||
407
python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py
Normal file
407
python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py
Normal 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"],
|
||||
)
|
||||
@@ -16,6 +16,10 @@ from sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute import (
|
||||
compute_cu_seqlens,
|
||||
)
|
||||
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.transform_index import (
|
||||
transform_index_page_table_decode,
|
||||
@@ -63,6 +67,15 @@ else:
|
||||
# Reuse this workspace buffer across all NSA backend instances
|
||||
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)
|
||||
class NSAFlashMLAMetadata:
|
||||
@@ -1127,55 +1140,150 @@ class NativeSparseAttnBackend(
|
||||
|
||||
metadata = self.decode_cuda_graph_metadata[bs]
|
||||
|
||||
# Copy basic seqlens
|
||||
metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens)
|
||||
metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
|
||||
# Track whether fused kernel succeeded
|
||||
fused_kernel_succeeded = False
|
||||
|
||||
# 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)
|
||||
# Use fused CUDA kernel for all copy operations
|
||||
if _USE_FUSED_METADATA_COPY:
|
||||
try:
|
||||
from sglang.jit_kernel.fused_metadata_copy import (
|
||||
fused_metadata_copy_cuda,
|
||||
)
|
||||
|
||||
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)
|
||||
# Map forward_mode to integer enum
|
||||
if forward_mode.is_decode_or_idle():
|
||||
mode_int = 0 # DECODE
|
||||
elif forward_mode.is_target_verify():
|
||||
mode_int = 1 # TARGET_VERIFY
|
||||
elif forward_mode.is_draft_extend():
|
||||
mode_int = 2 # DRAFT_EXTEND
|
||||
else:
|
||||
raise ValueError(f"Unsupported forward_mode: {forward_mode}")
|
||||
|
||||
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)
|
||||
# Prepare FlashMLA tensors if needed
|
||||
flashmla_num_splits_src = None
|
||||
flashmla_num_splits_dst = None
|
||||
flashmla_metadata_src = None
|
||||
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
|
||||
metadata.nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded)
|
||||
metadata.nsa_cache_seqlens_int32[:size].copy_(precomputed.nsa_cache_seqlens)
|
||||
metadata.nsa_cu_seqlens_k[1 : 1 + size].copy_(
|
||||
precomputed.nsa_cu_seqlens_k[1 : 1 + size]
|
||||
)
|
||||
|
||||
# Copy NSA cu_seqlens
|
||||
size = precomputed.seqlens_expanded_size
|
||||
metadata.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
|
||||
metadata.real_page_table[:rows, :cols].copy_(
|
||||
precomputed.real_page_table
|
||||
)
|
||||
|
||||
# Copy real page table
|
||||
if precomputed.real_page_table is not None:
|
||||
rows, cols = precomputed.real_page_table.shape
|
||||
metadata.real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
|
||||
else:
|
||||
# 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)
|
||||
# Copy FlashMLA metadata in fallback path
|
||||
if precomputed.flashmla_metadata is not None:
|
||||
size = precomputed.seqlens_expanded_size
|
||||
flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1))
|
||||
flashmla_metadata.copy_(precomputed.flashmla_metadata)
|
||||
|
||||
self.forward_metadata = metadata
|
||||
|
||||
@@ -1958,15 +2066,163 @@ class NativeSparseAttnMultiStepBackend:
|
||||
spec_info=forward_batch.spec_info,
|
||||
)
|
||||
|
||||
# Fast copy to each backend (1-2x faster than computing N times)
|
||||
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,
|
||||
)
|
||||
# Use multi-backend fused copy when we have 3 or more backends
|
||||
# This is 3x faster than calling the single-backend copy 3 times
|
||||
if self.speculative_num_steps >= 3:
|
||||
try:
|
||||
from sglang.jit_kernel.fused_metadata_copy import (
|
||||
fused_metadata_copy_multi_cuda,
|
||||
)
|
||||
|
||||
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:
|
||||
# Fallback: compute metadata separately for each backend
|
||||
for i in range(self.speculative_num_steps):
|
||||
|
||||
Reference in New Issue
Block a user