From 34132d6da50e0867426962d3681b62a03b5624b9 Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Sat, 14 Feb 2026 00:40:15 -0800 Subject: [PATCH] Kernel: optimize decoding metadata in NSA multi-spec backend with fused kernels (#17554) --- .../csrc/elementwise/fused_metadata_copy.cuh | 722 +++++++++++ .../sglang/jit_kernel/fused_metadata_copy.py | 316 +++++ .../tests/test_fused_metadata_copy.py | 1067 +++++++++++++++++ python/sglang/srt/environ.py | 2 + .../nsa/nsa_backend_mtp_precompute.py | 6 +- .../attention/nsa/nsa_mtp_verification.py | 407 +++++++ .../srt/layers/attention/nsa_backend.py | 358 +++++- 7 files changed, 2824 insertions(+), 54 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh create mode 100644 python/sglang/jit_kernel/fused_metadata_copy.py create mode 100644 python/sglang/jit_kernel/tests/test_fused_metadata_copy.py create mode 100644 python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh b/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh new file mode 100644 index 000000000..c996f6f1b --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh @@ -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 +#include + +#include + +#include + +#include // for std::min +#include + +// 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 +__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 +__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 +inline const T* unwrap_data_ptr(const tvm::ffi::TensorView& tensor, const char* name) { + using namespace host; + if (tensor.data_ptr()) { + RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); + } + return static_cast(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 +inline T* unwrap_data_ptr_mut(const tvm::ffi::TensorView& tensor, const char* name) { + using namespace host; + if (tensor.data_ptr()) { + RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); + } + return static_cast(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 +inline const T* +unwrap_optional_data_ptr(const tvm::ffi::Optional& optional_tensor, const char* name) { + using namespace host; + if (!optional_tensor.has_value()) { + return nullptr; + } + const auto& tensor = optional_tensor.value(); + RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); + return static_cast(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 +inline T* +unwrap_optional_data_ptr_mut(const tvm::ffi::Optional& optional_tensor, const char* name) { + using namespace host; + if (!optional_tensor.has_value()) { + return nullptr; + } + const auto& tensor = optional_tensor.value(); + RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); + return static_cast(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 +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 seqlens_expanded_src, + const tvm::ffi::TensorView nsa_cu_seqlens_k_src, + const tvm::ffi::Optional real_page_table_src, + const tvm::ffi::Optional flashmla_num_splits_src, + const tvm::ffi::Optional 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 seqlens_expanded_dst, + const tvm::ffi::TensorView nsa_cu_seqlens_k_dst, + const tvm::ffi::Optional real_page_table_dst, + const tvm::ffi::Optional flashmla_num_splits_dst, + const tvm::ffi::Optional 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(cache_seqlens_src, "cache_seqlens_src"), + .cu_seqlens_k = unwrap_data_ptr(cu_seqlens_k_src, "cu_seqlens_k_src"), + .page_indices = unwrap_data_ptr(page_indices_src, "page_indices_src"), + .nsa_cache_seqlens = unwrap_data_ptr(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), + .seqlens_expanded = unwrap_optional_data_ptr(seqlens_expanded_src, "seqlens_expanded_src"), + .nsa_cu_seqlens_k = unwrap_data_ptr(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), + .real_page_table = unwrap_optional_data_ptr(real_page_table_src, "real_page_table_src"), + .flashmla_num_splits = + unwrap_optional_data_ptr(flashmla_num_splits_src, "flashmla_num_splits_src"), + .flashmla_metadata = unwrap_optional_data_ptr(flashmla_metadata_src, "flashmla_metadata_src"), + }, + .dst = + { + .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst, "cache_seqlens_dst"), + .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst, "cu_seqlens_k_dst"), + .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst, "page_table_1_dst"), + .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"), + .seqlens_expanded = unwrap_optional_data_ptr_mut(seqlens_expanded_dst, "seqlens_expanded_dst"), + .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"), + .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst, "real_page_table_dst"), + .flashmla_num_splits = + unwrap_optional_data_ptr_mut(flashmla_num_splits_dst, "flashmla_num_splits_dst"), + .flashmla_metadata = + unwrap_optional_data_ptr_mut(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(page_indices_src.shape()[0]), + .page_table_1_stride = static_cast(page_table_1_dst.shape()[1]), + .real_page_table_cols = + real_page_table_src.has_value() ? static_cast(real_page_table_src.value().shape()[1]) : 0, + .real_page_table_dst_stride = + real_page_table_dst.has_value() ? static_cast(real_page_table_dst.value().stride(0)) : 0, + .flashmla_metadata_size = + flashmla_metadata_src.has_value() ? static_cast(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, 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 +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 real_page_table_src, + const tvm::ffi::Optional flashmla_num_splits_src, + const tvm::ffi::Optional 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 real_page_table_dst0, + const tvm::ffi::Optional flashmla_num_splits_dst0, + const tvm::ffi::Optional 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 real_page_table_dst1, + const tvm::ffi::Optional flashmla_num_splits_dst1, + const tvm::ffi::Optional 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 real_page_table_dst2, + const tvm::ffi::Optional flashmla_num_splits_dst2, + const tvm::ffi::Optional 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(cache_seqlens_src, "cache_seqlens_src"), + .cu_seqlens_k = unwrap_data_ptr(cu_seqlens_k_src, "cu_seqlens_k_src"), + .page_indices = unwrap_data_ptr(page_indices_src, "page_indices_src"), + .nsa_cache_seqlens = unwrap_data_ptr(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(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), + .real_page_table = unwrap_optional_data_ptr(real_page_table_src, "real_page_table_src"), + .flashmla_num_splits = + unwrap_optional_data_ptr(flashmla_num_splits_src, "flashmla_num_splits_src"), + .flashmla_metadata = unwrap_optional_data_ptr(flashmla_metadata_src, "flashmla_metadata_src"), + }, + .dst0 = + { + .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst0, "cache_seqlens_dst0"), + .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"), + .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst0, "page_table_1_dst0"), + .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"), + .seqlens_expanded = nullptr, + .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"), + .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst0, "real_page_table_dst0"), + .flashmla_num_splits = + unwrap_optional_data_ptr_mut(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"), + .flashmla_metadata = + unwrap_optional_data_ptr_mut(flashmla_metadata_dst0, "flashmla_metadata_dst0"), + }, + .dst1 = + { + .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst1, "cache_seqlens_dst1"), + .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"), + .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst1, "page_table_1_dst1"), + .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"), + .seqlens_expanded = nullptr, + .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"), + .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst1, "real_page_table_dst1"), + .flashmla_num_splits = + unwrap_optional_data_ptr_mut(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"), + .flashmla_metadata = + unwrap_optional_data_ptr_mut(flashmla_metadata_dst1, "flashmla_metadata_dst1"), + }, + .dst2 = + { + .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst2, "cache_seqlens_dst2"), + .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"), + .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst2, "page_table_1_dst2"), + .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"), + .seqlens_expanded = nullptr, + .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"), + .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst2, "real_page_table_dst2"), + .flashmla_num_splits = + unwrap_optional_data_ptr_mut(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"), + .flashmla_metadata = + unwrap_optional_data_ptr_mut(flashmla_metadata_dst2, "flashmla_metadata_dst2"), + }, + .bs = bs, + .max_len = max_len, + .seqlens_expanded_size = seqlens_expanded_size, + .page_table_1_stride = static_cast(page_table_1_dst0.shape()[1]), + .real_page_table_cols = + real_page_table_src.has_value() ? static_cast(real_page_table_src.value().shape()[1]) : 0, + .real_page_table_dst_stride = + real_page_table_dst0.has_value() ? static_cast(real_page_table_dst0.value().stride(0)) : 0, + .flashmla_metadata_size = + flashmla_metadata_src.has_value() ? static_cast(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, params); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/fused_metadata_copy.py b/python/sglang/jit_kernel/fused_metadata_copy.py new file mode 100644 index 000000000..b4d347f6a --- /dev/null +++ b/python/sglang/jit_kernel/fused_metadata_copy.py @@ -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, + ) diff --git a/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py b/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py new file mode 100644 index 000000000..9cffd8d88 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py @@ -0,0 +1,1067 @@ +""" +Comprehensive tests for JIT-compiled fused metadata copy kernels. + +This test suite verifies: +1. Single-backend fused kernel (fused_metadata_copy_cuda) - all forward modes +2. Multi-backend fused kernel (fused_metadata_copy_multi_cuda) - 3 backends at once +3. Correctness against reference implementations +4. Performance benchmarks and speedup measurements +""" + +import time + +import pytest +import torch + +# ============================================================================= +# Helper Functions +# ============================================================================= + + +def create_test_metadata( + bs: int, + max_len: int, + max_seqlen_k: int, + seqlens_expanded_size: int, + has_real_page_table: bool = False, + has_flashmla: bool = False, + device: str = "cuda", +): + """Create test metadata tensors matching NSA backend structure.""" + # Basic tensors (always present) + cache_seqlens_src = torch.randint( + 1, max_len, (bs,), dtype=torch.int32, device=device + ) + cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) + cu_seqlens_k_src[1:] = torch.cumsum(cache_seqlens_src, dim=0) + + page_indices_src = torch.randint( + 0, 1000, (bs, max_len), dtype=torch.int32, device=device + ) + nsa_cache_seqlens_src = torch.randint( + 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device + ) + seqlens_expanded_src = torch.randint( + 1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_src = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0) + + # Destination tensors + cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) + cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) + page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device) + nsa_cache_seqlens_dst = torch.zeros( + seqlens_expanded_size, dtype=torch.int32, device=device + ) + nsa_seqlens_expanded_dst = torch.zeros( + seqlens_expanded_size, dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_dst = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + + # Optional tensors + real_page_table_src = None + real_page_table_dst = None + if has_real_page_table: + real_page_table_cols = max_len // 2 + real_page_table_src = torch.randint( + 0, 1000, (bs, real_page_table_cols), dtype=torch.int32, device=device + ) + real_page_table_dst = torch.zeros( + (bs, real_page_table_cols + 8), dtype=torch.int32, device=device + ) + + flashmla_num_splits_src = None + flashmla_num_splits_dst = None + flashmla_metadata_src = None + flashmla_metadata_dst = None + if has_flashmla: + flashmla_num_splits_src = torch.randint( + 1, 10, (seqlens_expanded_size + 1,), dtype=torch.int32, device=device + ) + flashmla_num_splits_dst = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + # FlashMLA metadata is typically (num_sm_parts, TileSchedulerMetaDataSize) + # For testing, we use a simplified size + flashmla_metadata_size = 128 + flashmla_metadata_src = torch.randint( + 0, 100, (flashmla_metadata_size,), dtype=torch.int32, device=device + ) + flashmla_metadata_dst = torch.zeros( + flashmla_metadata_size, dtype=torch.int32, device=device + ) + + return { + "src": { + "cache_seqlens": cache_seqlens_src, + "cu_seqlens_k": cu_seqlens_k_src, + "page_indices": page_indices_src, + "nsa_cache_seqlens": nsa_cache_seqlens_src, + "seqlens_expanded": seqlens_expanded_src, + "nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, + "real_page_table": real_page_table_src, + "flashmla_num_splits": flashmla_num_splits_src, + "flashmla_metadata": flashmla_metadata_src, + }, + "dst": { + "cache_seqlens": cache_seqlens_dst, + "cu_seqlens_k": cu_seqlens_k_dst, + "page_table_1": page_table_1_dst, + "nsa_cache_seqlens": nsa_cache_seqlens_dst, + "nsa_seqlens_expanded": nsa_seqlens_expanded_dst, + "nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, + "real_page_table": real_page_table_dst, + "flashmla_num_splits": flashmla_num_splits_dst, + "flashmla_metadata": flashmla_metadata_dst, + }, + } + + +def reference_copy_decode(src, dst, max_len): + """Reference implementation: individual .copy_() for DECODE mode.""" + bs = src["cache_seqlens"].shape[0] + dst["cache_seqlens"].copy_(src["cache_seqlens"]) + dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) + dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) + dst["nsa_cache_seqlens"].copy_(src["nsa_cache_seqlens"]) + dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) + + if src["real_page_table"] is not None: + rows, cols = src["real_page_table"].shape + dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) + + if src["flashmla_num_splits"] is not None: + flashmla_size = bs + 1 + dst["flashmla_num_splits"][:flashmla_size].copy_( + src["flashmla_num_splits"][:flashmla_size] + ) + + if src["flashmla_metadata"] is not None: + dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) + + +def reference_copy_target_verify(src, dst, max_seqlen_k, seqlens_expanded_size): + """Reference implementation: individual .copy_() for TARGET_VERIFY mode.""" + bs = src["cache_seqlens"].shape[0] + dst["cache_seqlens"].copy_(src["cache_seqlens"]) + dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) + + rows, cols = src["page_indices"].shape + dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) + dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) + dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) + dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( + src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] + ) + + if src["real_page_table"] is not None: + rows, cols = src["real_page_table"].shape + dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) + + if src["flashmla_num_splits"] is not None: + flashmla_size = seqlens_expanded_size + 1 + dst["flashmla_num_splits"][:flashmla_size].copy_( + src["flashmla_num_splits"][:flashmla_size] + ) + + if src["flashmla_metadata"] is not None: + dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) + + +def reference_copy_draft_extend(src, dst, max_seqlen_k, seqlens_expanded_size): + """Reference implementation: individual .copy_() for DRAFT_EXTEND mode.""" + bs = src["cache_seqlens"].shape[0] + dst["cache_seqlens"].copy_(src["cache_seqlens"]) + dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) + + rows, cols = src["page_indices"].shape + dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) + dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) + dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) + dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( + src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] + ) + + if src["real_page_table"] is not None: + rows, cols = src["real_page_table"].shape + dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) + + if src["flashmla_num_splits"] is not None: + flashmla_size = seqlens_expanded_size + 1 + dst["flashmla_num_splits"][:flashmla_size].copy_( + src["flashmla_num_splits"][:flashmla_size] + ) + + if src["flashmla_metadata"] is not None: + dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) + + +# ============================================================================= +# Single-Backend Kernel Tests +# ============================================================================= + + +def test_fused_metadata_copy_dtype_validation(): + """Test that dtype validation rejects non-int32 tensors.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda + + bs = 2 + max_len = 128 + max_seqlen_k = 256 + seqlens_expanded_size = bs + device = "cuda" + + # Create tensors with WRONG dtype (int64 instead of int32) + cache_seqlens_src_wrong = torch.randint( + 1, max_len, (bs,), dtype=torch.int64, device=device + ) + cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) + page_indices_src = torch.randint( + 0, 1000, (bs, max_len), dtype=torch.int32, device=device + ) + nsa_cache_seqlens_src = torch.randint( + 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device + ) + seqlens_expanded_src = torch.randint( + 1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_src = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + + # Destination tensors (correct dtype) + cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) + cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) + page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device) + nsa_cache_seqlens_dst = torch.zeros( + seqlens_expanded_size, dtype=torch.int32, device=device + ) + nsa_seqlens_expanded_dst = torch.zeros( + seqlens_expanded_size, dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_dst = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + + # Test 1: Wrong dtype for source tensor should raise RuntimeError + with pytest.raises(RuntimeError, match="must have dtype int32"): + fused_metadata_copy_cuda( + cache_seqlens_src_wrong, # Wrong dtype: int64 + cu_seqlens_k_src, + page_indices_src, + nsa_cache_seqlens_src, + seqlens_expanded_src, + nsa_cu_seqlens_k_src, + None, # real_page_table_src + None, # flashmla_num_splits_src + None, # flashmla_metadata_src + cache_seqlens_dst, + cu_seqlens_k_dst, + page_table_1_dst, + nsa_cache_seqlens_dst, + nsa_seqlens_expanded_dst, + nsa_cu_seqlens_k_dst, + None, # real_page_table_dst + None, # flashmla_num_splits_dst + None, # flashmla_metadata_dst + 0, # forward_mode + bs, + max_len, + max_seqlen_k, + seqlens_expanded_size, + ) + + # Test 2: Wrong dtype for destination tensor should also raise RuntimeError + cache_seqlens_src = torch.randint( + 1, max_len, (bs,), dtype=torch.int32, device=device + ) + cache_seqlens_dst_wrong = torch.zeros(bs, dtype=torch.int64, device=device) + + with pytest.raises(RuntimeError, match="must have dtype int32"): + fused_metadata_copy_cuda( + cache_seqlens_src, + cu_seqlens_k_src, + page_indices_src, + nsa_cache_seqlens_src, + seqlens_expanded_src, + nsa_cu_seqlens_k_src, + None, + None, + None, + cache_seqlens_dst_wrong, # Wrong dtype: int64 + cu_seqlens_k_dst, + page_table_1_dst, + nsa_cache_seqlens_dst, + nsa_seqlens_expanded_dst, + nsa_cu_seqlens_k_dst, + None, + None, + None, + 0, + bs, + max_len, + max_seqlen_k, + seqlens_expanded_size, + ) + + +@pytest.mark.parametrize("bs", [1, 2, 4, 8]) +@pytest.mark.parametrize( + "forward_mode", [0] +) # DECODE mode only (other modes not fully tested yet) +@pytest.mark.parametrize("has_real_page_table", [False, True]) +@pytest.mark.parametrize("has_flashmla", [False, True]) +def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla): + """Test fused metadata copy kernel against reference implementation.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda + + max_len = 128 + max_seqlen_k = 256 + seqlens_expanded_size = bs if forward_mode == 0 else bs * 2 + + # Create test data + data = create_test_metadata( + bs=bs, + max_len=max_len, + max_seqlen_k=max_seqlen_k, + seqlens_expanded_size=seqlens_expanded_size, + has_real_page_table=has_real_page_table, + has_flashmla=has_flashmla, + ) + + # Create separate destination tensors for reference and fused kernel + dst_ref = {k: v.clone() if v is not None else None for k, v in data["dst"].items()} + dst_fused = { + k: v.clone() if v is not None else None for k, v in data["dst"].items() + } + + # Run reference implementation + if forward_mode == 0: # DECODE + reference_copy_decode(data["src"], dst_ref, max_len) + elif forward_mode == 1: # TARGET_VERIFY + reference_copy_target_verify( + data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size + ) + else: # DRAFT_EXTEND + reference_copy_draft_extend( + data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size + ) + + # Run fused kernel + fused_metadata_copy_cuda( + data["src"]["cache_seqlens"], + data["src"]["cu_seqlens_k"], + data["src"]["page_indices"], + data["src"]["nsa_cache_seqlens"], + data["src"]["seqlens_expanded"], + data["src"]["nsa_cu_seqlens_k"], + data["src"]["real_page_table"], + data["src"]["flashmla_num_splits"], + data["src"]["flashmla_metadata"], + dst_fused["cache_seqlens"], + dst_fused["cu_seqlens_k"], + dst_fused["page_table_1"], + dst_fused["nsa_cache_seqlens"], + dst_fused["nsa_seqlens_expanded"], + dst_fused["nsa_cu_seqlens_k"], + dst_fused["real_page_table"], + dst_fused["flashmla_num_splits"], + dst_fused["flashmla_metadata"], + forward_mode, + bs, + max_len, + max_seqlen_k, + seqlens_expanded_size, + ) + + # Compare results + assert torch.equal( + dst_ref["cache_seqlens"], dst_fused["cache_seqlens"] + ), "cache_seqlens mismatch" + assert torch.equal( + dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"] + ), "cu_seqlens_k mismatch" + assert torch.equal( + dst_ref["page_table_1"], dst_fused["page_table_1"] + ), "page_table_1 mismatch" + assert torch.equal( + dst_ref["nsa_cache_seqlens"], dst_fused["nsa_cache_seqlens"] + ), "nsa_cache_seqlens mismatch" + assert torch.equal( + dst_ref["nsa_seqlens_expanded"], dst_fused["nsa_seqlens_expanded"] + ), "nsa_seqlens_expanded mismatch" + assert torch.equal( + dst_ref["nsa_cu_seqlens_k"], dst_fused["nsa_cu_seqlens_k"] + ), "nsa_cu_seqlens_k mismatch" + + if has_real_page_table: + assert torch.equal( + dst_ref["real_page_table"], dst_fused["real_page_table"] + ), "real_page_table mismatch" + + if has_flashmla: + assert torch.equal( + dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"] + ), "flashmla_num_splits mismatch" + assert torch.equal( + dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"] + ), "flashmla_metadata mismatch" + + +@pytest.mark.parametrize("bs", [16, 32]) +def test_fused_metadata_copy_large_batch(bs): + """Test with larger batch sizes.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda + + forward_mode = 0 # DECODE + max_len = 128 + max_seqlen_k = 256 + seqlens_expanded_size = bs + + data = create_test_metadata( + bs=bs, + max_len=max_len, + max_seqlen_k=max_seqlen_k, + seqlens_expanded_size=seqlens_expanded_size, + has_real_page_table=True, + has_flashmla=True, + ) + + dst_ref = {k: v.clone() if v is not None else None for k, v in data["dst"].items()} + dst_fused = { + k: v.clone() if v is not None else None for k, v in data["dst"].items() + } + + reference_copy_decode(data["src"], dst_ref, max_len) + + fused_metadata_copy_cuda( + data["src"]["cache_seqlens"], + data["src"]["cu_seqlens_k"], + data["src"]["page_indices"], + data["src"]["nsa_cache_seqlens"], + data["src"]["seqlens_expanded"], + data["src"]["nsa_cu_seqlens_k"], + data["src"]["real_page_table"], + data["src"]["flashmla_num_splits"], + data["src"]["flashmla_metadata"], + dst_fused["cache_seqlens"], + dst_fused["cu_seqlens_k"], + dst_fused["page_table_1"], + dst_fused["nsa_cache_seqlens"], + dst_fused["nsa_seqlens_expanded"], + dst_fused["nsa_cu_seqlens_k"], + dst_fused["real_page_table"], + dst_fused["flashmla_num_splits"], + dst_fused["flashmla_metadata"], + forward_mode, + bs, + max_len, + max_seqlen_k, + seqlens_expanded_size, + ) + + # Verify all tensors match + for key in dst_ref: + if dst_ref[key] is not None: + assert torch.equal(dst_ref[key], dst_fused[key]), f"{key} mismatch" + + +# ============================================================================= +# Multi-Backend Kernel Tests +# ============================================================================= + + +def create_test_metadata_multi( + bs: int, + max_len: int, + seqlens_expanded_size: int, + has_real_page_table: bool = False, + has_flashmla: bool = False, + device: str = "cuda", +): + """Create test metadata tensors for multi-backend testing.""" + # Source tensors (precomputed metadata) + cache_seqlens_src = torch.randint( + 1, max_len, (bs,), dtype=torch.int32, device=device + ) + cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) + cu_seqlens_k_src[1:] = torch.cumsum(cache_seqlens_src, dim=0) + + page_indices_src = torch.randint( + 0, 1000, (bs, max_len), dtype=torch.int32, device=device + ) + nsa_cache_seqlens_src = torch.randint( + 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_src = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0) + + # Optional tensors + real_page_table_src = None + if has_real_page_table: + real_page_table_cols = max_len // 2 + real_page_table_src = torch.randint( + 0, 1000, (bs, real_page_table_cols), dtype=torch.int32, device=device + ) + + flashmla_num_splits_src = None + flashmla_metadata_src = None + if has_flashmla: + flashmla_num_splits_src = torch.randint( + 1, 10, (seqlens_expanded_size + 1,), dtype=torch.int32, device=device + ) + flashmla_metadata_size = 128 + flashmla_metadata_src = torch.randint( + 0, 100, (flashmla_metadata_size,), dtype=torch.int32, device=device + ) + + # Create destination tensors for 3 backends + def create_dst_tensors(): + cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) + cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) + page_table_1_dst = torch.zeros( + (bs, max_len + 16), dtype=torch.int32, device=device + ) + nsa_cache_seqlens_dst = torch.zeros( + seqlens_expanded_size, dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_dst = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + + real_page_table_dst = None + if has_real_page_table: + real_page_table_cols = max_len // 2 + real_page_table_dst = torch.zeros( + (bs, real_page_table_cols + 8), dtype=torch.int32, device=device + ) + + flashmla_num_splits_dst = None + flashmla_metadata_dst = None + if has_flashmla: + flashmla_num_splits_dst = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + flashmla_metadata_size = 128 + flashmla_metadata_dst = torch.zeros( + flashmla_metadata_size, dtype=torch.int32, device=device + ) + + return { + "cache_seqlens_int32": cache_seqlens_dst, + "cu_seqlens_k": cu_seqlens_k_dst, + "page_table_1": page_table_1_dst, + "nsa_cache_seqlens_int32": nsa_cache_seqlens_dst, + "nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, + "real_page_table": real_page_table_dst, + "flashmla_num_splits": flashmla_num_splits_dst, + "flashmla_metadata": flashmla_metadata_dst, + } + + return { + "src": { + "cache_seqlens": cache_seqlens_src, + "cu_seqlens_k": cu_seqlens_k_src, + "page_indices": page_indices_src, + "nsa_cache_seqlens": nsa_cache_seqlens_src, + "nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, + "real_page_table": real_page_table_src, + "flashmla_num_splits": flashmla_num_splits_src, + "flashmla_metadata": flashmla_metadata_src, + }, + "dst0": create_dst_tensors(), + "dst1": create_dst_tensors(), + "dst2": create_dst_tensors(), + } + + +def reference_copy_for_loop(src, dst_list, bs, max_len): + """Reference implementation: for-loop calling copy for each backend.""" + for dst in dst_list: + # Simulate what init_forward_metadata_replay_cuda_graph_from_precomputed does + dst["cache_seqlens_int32"].copy_(src["cache_seqlens"]) + dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) + dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) + dst["nsa_cache_seqlens_int32"].copy_(src["nsa_cache_seqlens"]) + dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) + + if src["real_page_table"] is not None: + rows, cols = src["real_page_table"].shape + dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) + + if src["flashmla_num_splits"] is not None: + flashmla_size = bs + 1 + dst["flashmla_num_splits"][:flashmla_size].copy_( + src["flashmla_num_splits"][:flashmla_size] + ) + + if src["flashmla_metadata"] is not None: + dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) + + +def test_fused_metadata_copy_multi_dtype_validation(): + """Test that dtype validation rejects non-int32 tensors for multi-backend kernel.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda + + bs = 2 + max_len = 128 + seqlens_expanded_size = bs + device = "cuda" + + # Create source tensors - one with WRONG dtype + cache_seqlens_src_wrong = torch.randint( + 1, max_len, (bs,), dtype=torch.int64, device=device # Wrong dtype! + ) + cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) + page_indices_src = torch.randint( + 0, 1000, (bs, max_len), dtype=torch.int32, device=device + ) + nsa_cache_seqlens_src = torch.randint( + 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device + ) + nsa_cu_seqlens_k_src = torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ) + + # Create destination tensors for 3 backends (all correct dtype) + def create_dst(): + return { + "cache_seqlens": torch.zeros(bs, dtype=torch.int32, device=device), + "cu_seqlens_k": torch.zeros(bs + 1, dtype=torch.int32, device=device), + "page_table_1": torch.zeros( + (bs, max_len + 16), dtype=torch.int32, device=device + ), + "nsa_cache_seqlens": torch.zeros( + seqlens_expanded_size, dtype=torch.int32, device=device + ), + "nsa_cu_seqlens_k": torch.zeros( + seqlens_expanded_size + 1, dtype=torch.int32, device=device + ), + } + + dst0 = create_dst() + dst1 = create_dst() + dst2 = create_dst() + + # Test: Wrong dtype for source tensor should raise RuntimeError + with pytest.raises(RuntimeError, match="must have dtype int32"): + fused_metadata_copy_multi_cuda( + cache_seqlens_src_wrong, # Wrong dtype: int64 + cu_seqlens_k_src, + page_indices_src, + nsa_cache_seqlens_src, + nsa_cu_seqlens_k_src, + None, # real_page_table_src + None, # flashmla_num_splits_src + None, # flashmla_metadata_src + # Backend 0 + dst0["cache_seqlens"], + dst0["cu_seqlens_k"], + dst0["page_table_1"], + dst0["nsa_cache_seqlens"], + dst0["nsa_cu_seqlens_k"], + None, + None, + None, + # Backend 1 + dst1["cache_seqlens"], + dst1["cu_seqlens_k"], + dst1["page_table_1"], + dst1["nsa_cache_seqlens"], + dst1["nsa_cu_seqlens_k"], + None, + None, + None, + # Backend 2 + dst2["cache_seqlens"], + dst2["cu_seqlens_k"], + dst2["page_table_1"], + dst2["nsa_cache_seqlens"], + dst2["nsa_cu_seqlens_k"], + None, + None, + None, + # Parameters + bs, + max_len, + seqlens_expanded_size, + ) + + +@pytest.mark.parametrize("bs", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("has_real_page_table", [False, True]) +@pytest.mark.parametrize("has_flashmla", [False, True]) +def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla): + """Test fused multi-backend metadata copy kernel against for-loop version.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda + + max_len = 128 + seqlens_expanded_size = bs + + # Create test data + data = create_test_metadata_multi( + bs=bs, + max_len=max_len, + seqlens_expanded_size=seqlens_expanded_size, + has_real_page_table=has_real_page_table, + has_flashmla=has_flashmla, + ) + + # Create separate destination tensors for reference (for-loop) and fused kernel + dst_ref_0 = { + k: v.clone() if v is not None else None for k, v in data["dst0"].items() + } + dst_ref_1 = { + k: v.clone() if v is not None else None for k, v in data["dst1"].items() + } + dst_ref_2 = { + k: v.clone() if v is not None else None for k, v in data["dst2"].items() + } + + dst_fused_0 = { + k: v.clone() if v is not None else None for k, v in data["dst0"].items() + } + dst_fused_1 = { + k: v.clone() if v is not None else None for k, v in data["dst1"].items() + } + dst_fused_2 = { + k: v.clone() if v is not None else None for k, v in data["dst2"].items() + } + + # Run reference implementation (for-loop) + torch.cuda.synchronize() + loop_start = time.perf_counter() + reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len) + torch.cuda.synchronize() + loop_end = time.perf_counter() + loop_time = loop_end - loop_start + + # Run fused kernel + torch.cuda.synchronize() + fused_start = time.perf_counter() + fused_metadata_copy_multi_cuda( + # Source tensors + data["src"]["cache_seqlens"], + data["src"]["cu_seqlens_k"], + data["src"]["page_indices"], + data["src"]["nsa_cache_seqlens"], + data["src"]["nsa_cu_seqlens_k"], + data["src"]["real_page_table"], + data["src"]["flashmla_num_splits"], + data["src"]["flashmla_metadata"], + # Destination tensors for backend 0 + dst_fused_0["cache_seqlens_int32"], + dst_fused_0["cu_seqlens_k"], + dst_fused_0["page_table_1"], + dst_fused_0["nsa_cache_seqlens_int32"], + dst_fused_0["nsa_cu_seqlens_k"], + dst_fused_0["real_page_table"], + dst_fused_0["flashmla_num_splits"], + dst_fused_0["flashmla_metadata"], + # Destination tensors for backend 1 + dst_fused_1["cache_seqlens_int32"], + dst_fused_1["cu_seqlens_k"], + dst_fused_1["page_table_1"], + dst_fused_1["nsa_cache_seqlens_int32"], + dst_fused_1["nsa_cu_seqlens_k"], + dst_fused_1["real_page_table"], + dst_fused_1["flashmla_num_splits"], + dst_fused_1["flashmla_metadata"], + # Destination tensors for backend 2 + dst_fused_2["cache_seqlens_int32"], + dst_fused_2["cu_seqlens_k"], + dst_fused_2["page_table_1"], + dst_fused_2["nsa_cache_seqlens_int32"], + dst_fused_2["nsa_cu_seqlens_k"], + dst_fused_2["real_page_table"], + dst_fused_2["flashmla_num_splits"], + dst_fused_2["flashmla_metadata"], + # Parameters + bs, + max_len, + seqlens_expanded_size, + ) + torch.cuda.synchronize() + fused_end = time.perf_counter() + fused_time = fused_end - fused_start + + # Compare results for all 3 backends + speedup = loop_time / fused_time if fused_time > 0 else 0 + print( + f"\n[VERIFY] bs={bs}, real_page_table={has_real_page_table}, flashmla={has_flashmla}" + ) + print( + f"[VERIFY] Fused time: {fused_time*1000:.3f}ms, Loop time: {loop_time*1000:.3f}ms, Speedup: {speedup:.2f}x" + ) + + max_diff = 0.0 + all_match = True + + for backend_idx, (dst_ref, dst_fused) in enumerate( + [ + (dst_ref_0, dst_fused_0), + (dst_ref_1, dst_fused_1), + (dst_ref_2, dst_fused_2), + ] + ): + for key in [ + "cache_seqlens_int32", + "cu_seqlens_k", + "page_table_1", + "nsa_cache_seqlens_int32", + "nsa_cu_seqlens_k", + ]: + if not torch.equal(dst_ref[key], dst_fused[key]): + diff = ( + (dst_ref[key].float() - dst_fused[key].float()).abs().max().item() + ) + max_diff = max(max_diff, diff) + all_match = False + print( + f"[ERROR] Backend {backend_idx} {key}: MISMATCH! Max diff: {diff}" + ) + + if has_real_page_table and dst_ref["real_page_table"] is not None: + if not torch.equal( + dst_ref["real_page_table"], dst_fused["real_page_table"] + ): + diff = ( + ( + dst_ref["real_page_table"].float() + - dst_fused["real_page_table"].float() + ) + .abs() + .max() + .item() + ) + max_diff = max(max_diff, diff) + all_match = False + print( + f"[ERROR] Backend {backend_idx} real_page_table: MISMATCH! Max diff: {diff}" + ) + + if has_flashmla: + if dst_ref["flashmla_num_splits"] is not None and not torch.equal( + dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"] + ): + diff = ( + ( + dst_ref["flashmla_num_splits"].float() + - dst_fused["flashmla_num_splits"].float() + ) + .abs() + .max() + .item() + ) + max_diff = max(max_diff, diff) + all_match = False + print( + f"[ERROR] Backend {backend_idx} flashmla_num_splits: MISMATCH! Max diff: {diff}" + ) + + if dst_ref["flashmla_metadata"] is not None and not torch.equal( + dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"] + ): + diff = ( + ( + dst_ref["flashmla_metadata"].float() + - dst_fused["flashmla_metadata"].float() + ) + .abs() + .max() + .item() + ) + max_diff = max(max_diff, diff) + all_match = False + print( + f"[ERROR] Backend {backend_idx} flashmla_metadata: MISMATCH! Max diff: {diff}" + ) + + if not all_match: + error_msg = ( + f"Fused metadata copy verification FAILED! " + f"Maximum difference: {max_diff}. " + f"The fused kernel produces different results than the for-loop version." + ) + print(f"[ERROR] {error_msg}") + raise AssertionError(error_msg) + + print(f"[VERIFY] Verification PASSED - all tensors match!") + + +@pytest.mark.parametrize("bs", [32, 64]) +def test_fused_metadata_copy_multi_large_batch(bs): + """Test with larger batch sizes and timing comparison.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda + + max_len = 128 + seqlens_expanded_size = bs + + data = create_test_metadata_multi( + bs=bs, + max_len=max_len, + seqlens_expanded_size=seqlens_expanded_size, + has_real_page_table=True, + has_flashmla=True, + ) + + dst_ref_0 = { + k: v.clone() if v is not None else None for k, v in data["dst0"].items() + } + dst_ref_1 = { + k: v.clone() if v is not None else None for k, v in data["dst1"].items() + } + dst_ref_2 = { + k: v.clone() if v is not None else None for k, v in data["dst2"].items() + } + + dst_fused_0 = { + k: v.clone() if v is not None else None for k, v in data["dst0"].items() + } + dst_fused_1 = { + k: v.clone() if v is not None else None for k, v in data["dst1"].items() + } + dst_fused_2 = { + k: v.clone() if v is not None else None for k, v in data["dst2"].items() + } + + # Warmup + for _ in range(5): + reference_copy_for_loop( + data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len + ) + fused_metadata_copy_multi_cuda( + data["src"]["cache_seqlens"], + data["src"]["cu_seqlens_k"], + data["src"]["page_indices"], + data["src"]["nsa_cache_seqlens"], + data["src"]["nsa_cu_seqlens_k"], + data["src"]["real_page_table"], + data["src"]["flashmla_num_splits"], + data["src"]["flashmla_metadata"], + dst_fused_0["cache_seqlens_int32"], + dst_fused_0["cu_seqlens_k"], + dst_fused_0["page_table_1"], + dst_fused_0["nsa_cache_seqlens_int32"], + dst_fused_0["nsa_cu_seqlens_k"], + dst_fused_0["real_page_table"], + dst_fused_0["flashmla_num_splits"], + dst_fused_0["flashmla_metadata"], + dst_fused_1["cache_seqlens_int32"], + dst_fused_1["cu_seqlens_k"], + dst_fused_1["page_table_1"], + dst_fused_1["nsa_cache_seqlens_int32"], + dst_fused_1["nsa_cu_seqlens_k"], + dst_fused_1["real_page_table"], + dst_fused_1["flashmla_num_splits"], + dst_fused_1["flashmla_metadata"], + dst_fused_2["cache_seqlens_int32"], + dst_fused_2["cu_seqlens_k"], + dst_fused_2["page_table_1"], + dst_fused_2["nsa_cache_seqlens_int32"], + dst_fused_2["nsa_cu_seqlens_k"], + dst_fused_2["real_page_table"], + dst_fused_2["flashmla_num_splits"], + dst_fused_2["flashmla_metadata"], + bs, + max_len, + seqlens_expanded_size, + ) + torch.cuda.synchronize() + + # Actual timing + torch.cuda.synchronize() + loop_start = time.perf_counter() + reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len) + torch.cuda.synchronize() + loop_time = time.perf_counter() - loop_start + + torch.cuda.synchronize() + fused_start = time.perf_counter() + fused_metadata_copy_multi_cuda( + data["src"]["cache_seqlens"], + data["src"]["cu_seqlens_k"], + data["src"]["page_indices"], + data["src"]["nsa_cache_seqlens"], + data["src"]["nsa_cu_seqlens_k"], + data["src"]["real_page_table"], + data["src"]["flashmla_num_splits"], + data["src"]["flashmla_metadata"], + dst_fused_0["cache_seqlens_int32"], + dst_fused_0["cu_seqlens_k"], + dst_fused_0["page_table_1"], + dst_fused_0["nsa_cache_seqlens_int32"], + dst_fused_0["nsa_cu_seqlens_k"], + dst_fused_0["real_page_table"], + dst_fused_0["flashmla_num_splits"], + dst_fused_0["flashmla_metadata"], + dst_fused_1["cache_seqlens_int32"], + dst_fused_1["cu_seqlens_k"], + dst_fused_1["page_table_1"], + dst_fused_1["nsa_cache_seqlens_int32"], + dst_fused_1["nsa_cu_seqlens_k"], + dst_fused_1["real_page_table"], + dst_fused_1["flashmla_num_splits"], + dst_fused_1["flashmla_metadata"], + dst_fused_2["cache_seqlens_int32"], + dst_fused_2["cu_seqlens_k"], + dst_fused_2["page_table_1"], + dst_fused_2["nsa_cache_seqlens_int32"], + dst_fused_2["nsa_cu_seqlens_k"], + dst_fused_2["real_page_table"], + dst_fused_2["flashmla_num_splits"], + dst_fused_2["flashmla_metadata"], + bs, + max_len, + seqlens_expanded_size, + ) + torch.cuda.synchronize() + fused_time = time.perf_counter() - fused_start + + speedup = loop_time / fused_time if fused_time > 0 else 0 + print( + f"\n[PERF] Large batch (bs={bs}): Fused={fused_time*1000:.3f}ms, Loop={loop_time*1000:.3f}ms, Speedup={speedup:.2f}x" + ) + + # Verify correctness + for backend_idx, (dst_ref, dst_fused) in enumerate( + [ + (dst_ref_0, dst_fused_0), + (dst_ref_1, dst_fused_1), + (dst_ref_2, dst_fused_2), + ] + ): + for key in dst_ref: + if dst_ref[key] is not None and dst_fused[key] is not None: + assert torch.equal( + dst_ref[key], dst_fused[key] + ), f"Backend {backend_idx} {key} mismatch" + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 37c16209c..65f81d037 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py index b9450ce09..846d276e4 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py @@ -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( diff --git a/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py b/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py new file mode 100644 index 000000000..b957d4ba8 --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py @@ -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"], + ) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 2de045c82..57193d49f 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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):