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

[Refactor] Rename NSA → DSA: user-facing aliases, file/class/import rename (#25821)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-20 00:18:04 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent da6d549ab2
commit 8131641bc6
162 changed files with 11298 additions and 10740 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh → python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh
+46 -46
View File
@@ -1,10 +1,10 @@
/*
* Fused metadata copy kernel for NSA backend CUDA graph replay.
* Fused metadata copy kernel for DSA 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
* page_table, dsa metadata, and optional FlashMLA metadata) into single kernel
* launches, significantly reducing kernel launch overhead and improving CUDA
* graph replay performance during inference.
*
@@ -37,7 +37,7 @@
#include <algorithm> // for std::min
#include <cuda_runtime.h>
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/nsa_backend.py)
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/dsa_backend.py)
enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
/**
@@ -49,9 +49,9 @@ 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__ dsa_cache_seqlens; // DSA-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__ dsa_cu_seqlens_k; // DSA 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
@@ -66,9 +66,9 @@ 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__ dsa_cache_seqlens; // DSA-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__ dsa_cu_seqlens_k; // DSA 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
@@ -189,26 +189,26 @@ __global__ void fused_metadata_copy_kernel(const FusedMetadataCopyParams __grid_
}
}
// Branch 3: NSA metadata copy (different loop sizes per mode)
// Branch 3: DSA 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];
dst.dsa_cache_seqlens[i] = src.dsa_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];
dst.dsa_cu_seqlens_k[i + 1] = src.dsa_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];
dst.dsa_cache_seqlens[i] = src.dsa_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];
dst.dsa_cu_seqlens_k[i + 1] = src.dsa_cu_seqlens_k[i + 1];
}
}
@@ -309,22 +309,22 @@ __global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiPar
dst2.page_table_1[row * page_table_1_stride + col] = val;
}
// Copy nsa_cache_seqlens to all 3 backends
// Copy dsa_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;
int32_t val = src.dsa_cache_seqlens[i];
dst0.dsa_cache_seqlens[i] = val;
dst1.dsa_cache_seqlens[i] = val;
dst2.dsa_cache_seqlens[i] = val;
}
// Copy NSA cu_seqlens to all 3 backends
// Copy DSA 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;
int32_t val = src.dsa_cu_seqlens_k[i + 1];
dst0.dsa_cu_seqlens_k[i + 1] = val;
dst1.dsa_cu_seqlens_k[i + 1] = val;
dst2.dsa_cu_seqlens_k[i + 1] = val;
}
// Copy real page table to all 3 backends
@@ -493,18 +493,18 @@ struct FusedMetadataCopyKernel {
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 dsa_cache_seqlens_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src,
const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst,
const tvm::ffi::TensorView cu_seqlens_k_dst,
const tvm::ffi::TensorView page_table_1_dst,
const tvm::ffi::TensorView nsa_cache_seqlens_dst,
const tvm::ffi::TensorView dsa_cache_seqlens_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_dst,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst,
const tvm::ffi::TensorView dsa_cu_seqlens_k_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst,
@@ -522,9 +522,9 @@ struct FusedMetadataCopyKernel {
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
.dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
.seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"),
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
.dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
@@ -535,9 +535,9 @@ struct FusedMetadataCopyKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"),
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst, "dsa_cache_seqlens_dst"),
.seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"),
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"),
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k_dst"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"),
@@ -605,32 +605,32 @@ struct FusedMetadataCopyMultiKernel {
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::TensorView dsa_cache_seqlens_src,
const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst0,
const tvm::ffi::TensorView cu_seqlens_k_dst0,
const tvm::ffi::TensorView page_table_1_dst0,
const tvm::ffi::TensorView nsa_cache_seqlens_dst0,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst0,
const tvm::ffi::TensorView dsa_cache_seqlens_dst0,
const tvm::ffi::TensorView dsa_cu_seqlens_k_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst0,
const tvm::ffi::TensorView cache_seqlens_dst1,
const tvm::ffi::TensorView cu_seqlens_k_dst1,
const tvm::ffi::TensorView page_table_1_dst1,
const tvm::ffi::TensorView nsa_cache_seqlens_dst1,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst1,
const tvm::ffi::TensorView dsa_cache_seqlens_dst1,
const tvm::ffi::TensorView dsa_cu_seqlens_k_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst1,
const tvm::ffi::TensorView cache_seqlens_dst2,
const tvm::ffi::TensorView cu_seqlens_k_dst2,
const tvm::ffi::TensorView page_table_1_dst2,
const tvm::ffi::TensorView nsa_cache_seqlens_dst2,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst2,
const tvm::ffi::TensorView dsa_cache_seqlens_dst2,
const tvm::ffi::TensorView dsa_cu_seqlens_k_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2,
@@ -647,9 +647,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
.dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
.seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
.dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
@@ -660,9 +660,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"),
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst0, "dsa_cache_seqlens_dst0"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"),
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst0, "dsa_cu_seqlens_k_dst0"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"),
@@ -674,9 +674,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"),
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst1, "dsa_cache_seqlens_dst1"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"),
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst1, "dsa_cu_seqlens_k_dst1"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"),
@@ -688,9 +688,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"),
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst2, "dsa_cache_seqlens_dst2"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"),
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst2, "dsa_cu_seqlens_k_dst2"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"),
python/sglang/jit_kernel/fused_metadata_copy.py
+39 -39
View File
@@ -1,5 +1,5 @@
"""
Fused metadata copy kernel for NSA backend CUDA graph replay.
Fused metadata copy kernel for DSA 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
@@ -98,18 +98,18 @@ 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,
dsa_cache_seqlens_src: torch.Tensor,
seqlens_expanded_src: Optional[torch.Tensor],
nsa_cu_seqlens_k_src: torch.Tensor,
dsa_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,
dsa_cache_seqlens_dst: torch.Tensor,
seqlens_expanded_dst: Optional[torch.Tensor],
nsa_cu_seqlens_k_dst: torch.Tensor,
dsa_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],
@@ -120,7 +120,7 @@ def fused_metadata_copy_cuda(
seqlens_expanded_size: int,
) -> None:
"""
Fused metadata copy kernel for NSA backend CUDA graph replay.
Fused metadata copy kernel for DSA backend CUDA graph replay.
This function fuses multiple tensor copy operations into a single kernel launch,
reducing kernel launch overhead and improving performance.
@@ -129,18 +129,18 @@ def fused_metadata_copy_cuda(
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]
dsa_cache_seqlens_src: Source DSA 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]
dsa_cu_seqlens_k_src: Source DSA 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]
dsa_cache_seqlens_dst: Destination DSA 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]
dsa_cu_seqlens_k_dst: Destination DSA 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
@@ -164,28 +164,28 @@ def fused_metadata_copy_cuda(
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()
dsa_cache_seqlens_src = dsa_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()
dsa_cu_seqlens_k_src = dsa_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,
dsa_cache_seqlens_src,
seqlens_expanded_src,
nsa_cu_seqlens_k_src,
dsa_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,
dsa_cache_seqlens_dst,
seqlens_expanded_dst,
nsa_cu_seqlens_k_dst,
dsa_cu_seqlens_k_dst,
real_page_table_dst,
flashmla_num_splits_dst,
flashmla_metadata_dst,
@@ -200,32 +200,32 @@ 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,
dsa_cache_seqlens_src: torch.Tensor,
dsa_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,
dsa_cache_seqlens_dst0: torch.Tensor,
dsa_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,
dsa_cache_seqlens_dst1: torch.Tensor,
dsa_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,
dsa_cache_seqlens_dst2: torch.Tensor,
dsa_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],
@@ -234,7 +234,7 @@ def fused_metadata_copy_multi_cuda(
seqlens_expanded_size: int,
) -> None:
"""
Multi-backend fused metadata copy kernel for NSA backend CUDA graph replay.
Multi-backend fused metadata copy kernel for DSA 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
@@ -244,16 +244,16 @@ def fused_metadata_copy_multi_cuda(
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]
dsa_cache_seqlens_src: Source DSA cache sequence lengths [bs]
dsa_cu_seqlens_k_src: Source DSA 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
dsa_cache_seqlens_dst0-2: Destination DSA cache sequence lengths for backends 0-2
dsa_cu_seqlens_k_dst0-2: Destination DSA 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
@@ -273,40 +273,40 @@ def fused_metadata_copy_multi_cuda(
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()
dsa_cache_seqlens_src = dsa_cache_seqlens_src.contiguous()
dsa_cu_seqlens_k_src = dsa_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,
dsa_cache_seqlens_src,
dsa_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,
dsa_cache_seqlens_dst0,
dsa_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,
dsa_cache_seqlens_dst1,
dsa_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,
dsa_cache_seqlens_dst2,
dsa_cu_seqlens_k_dst2,
real_page_table_dst2,
flashmla_num_splits_dst2,
flashmla_metadata_dst2,
python/sglang/jit_kernel/fused_store_index_cache.py
+8 -8
View File
@@ -28,7 +28,7 @@ logger = logging.getLogger(__name__)
@cache_once
def _jit_nsa_fused_store_module(
def _jit_dsa_fused_store_module(
key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int
) -> Module:
"""
@@ -39,13 +39,13 @@ def _jit_nsa_fused_store_module(
return load_jit(
"fused_store_index_k_cache",
*args,
cuda_files=["nsa/fused_store_index_cache.cuh"],
cuda_files=["dsa/fused_store_index_cache.cuh"],
cuda_wrappers=[
(
"fused_store_index_k_cache",
# - Float = bf16_t (sgl_kernel/type.cuh)
# - IndicesT = int64_t (out_cache_loc is int64 in SGLang SetKAndS)
# - kPageSize = 64 (CUDA NSA)
# - kPageSize = 64 (CUDA DSA)
f"FusedStoreCacheIndexerKernel<{args}>::run",
)
],
@@ -53,15 +53,15 @@ def _jit_nsa_fused_store_module(
@cache_once
def can_use_nsa_fused_store(
def can_use_dsa_fused_store(
key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int
) -> bool:
logger = logging.getLogger(__name__)
try:
_jit_nsa_fused_store_module(key_dtype, indices_dtype, page_size)
_jit_dsa_fused_store_module(key_dtype, indices_dtype, page_size)
return True
except Exception as e:
logger.warning(f"Failed to load nsa fused store JIT kernel: {e}")
logger.warning(f"Failed to load dsa fused store JIT kernel: {e}")
return False
@@ -73,7 +73,7 @@ def fused_store_index_k_cache(
page_size: int = 64,
) -> None:
"""
Fused: quantize bf16 key (N,128) -> fp8 + fp32 scale and write into NSATokenToKVPool.index_k_with_scale_buffer.
Fused: quantize bf16 key (N,128) -> fp8 + fp32 scale and write into DSATokenToKVPool.index_k_with_scale_buffer.
key: (num_tokens, 128) bf16 (or reshapeable to it)
index_k_with_scale: (num_pages, 64*(128+4)) uint8
@@ -101,5 +101,5 @@ def fused_store_index_k_cache(
if not index_k_with_scale.is_contiguous():
index_k_with_scale = index_k_with_scale.contiguous()
module = _jit_nsa_fused_store_module(key.dtype, out_cache_loc.dtype, page_size)
module = _jit_dsa_fused_store_module(key.dtype, out_cache_loc.dtype, page_size)
module.fused_store_index_k_cache(key, index_k_with_scale, out_cache_loc)
python/sglang/jit_kernel/tests/test_fused_metadata_copy.py
+102 -102
View File
@@ -33,7 +33,7 @@ def create_test_metadata(
has_flashmla: bool = False,
device: str = "cuda",
):
"""Create test metadata tensors matching NSA backend structure."""
"""Create test metadata tensors matching DSA backend structure."""
# Basic tensors (always present)
cache_seqlens_src = torch.randint(
1, max_len, (bs,), dtype=torch.int32, device=device
@@ -44,28 +44,28 @@ def create_test_metadata(
page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device
)
nsa_cache_seqlens_src = torch.randint(
dsa_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(
dsa_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)
dsa_cu_seqlens_k_src[1:] = torch.cumsum(dsa_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(
dsa_cache_seqlens_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device
)
nsa_seqlens_expanded_dst = torch.zeros(
dsa_seqlens_expanded_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device
)
nsa_cu_seqlens_k_dst = torch.zeros(
dsa_cu_seqlens_k_dst = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device
)
@@ -107,9 +107,9 @@ def create_test_metadata(
"cache_seqlens": cache_seqlens_src,
"cu_seqlens_k": cu_seqlens_k_src,
"page_indices": page_indices_src,
"nsa_cache_seqlens": nsa_cache_seqlens_src,
"dsa_cache_seqlens": dsa_cache_seqlens_src,
"seqlens_expanded": seqlens_expanded_src,
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_src,
"dsa_cu_seqlens_k": dsa_cu_seqlens_k_src,
"real_page_table": real_page_table_src,
"flashmla_num_splits": flashmla_num_splits_src,
"flashmla_metadata": flashmla_metadata_src,
@@ -118,9 +118,9 @@ def create_test_metadata(
"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,
"dsa_cache_seqlens": dsa_cache_seqlens_dst,
"dsa_seqlens_expanded": dsa_seqlens_expanded_dst,
"dsa_cu_seqlens_k": dsa_cu_seqlens_k_dst,
"real_page_table": real_page_table_dst,
"flashmla_num_splits": flashmla_num_splits_dst,
"flashmla_metadata": flashmla_metadata_dst,
@@ -134,8 +134,8 @@ def reference_copy_decode(src, dst, max_len):
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])
dst["dsa_cache_seqlens"].copy_(src["dsa_cache_seqlens"])
dst["dsa_cu_seqlens_k"][1 : bs + 1].copy_(src["dsa_cu_seqlens_k"][1 : bs + 1])
if src["real_page_table"] is not None:
rows, cols = src["real_page_table"].shape
@@ -159,10 +159,10 @@ def reference_copy_target_verify(src, dst, max_seqlen_k, seqlens_expanded_size):
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]
dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"])
dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1]
)
if src["real_page_table"] is not None:
@@ -187,10 +187,10 @@ def reference_copy_draft_extend(src, dst, max_seqlen_k, seqlens_expanded_size):
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]
dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"])
dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1]
)
if src["real_page_table"] is not None:
@@ -233,13 +233,13 @@ def test_fused_metadata_copy_dtype_validation():
page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device
)
nsa_cache_seqlens_src = torch.randint(
dsa_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(
dsa_cu_seqlens_k_src = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device
)
@@ -247,13 +247,13 @@ def test_fused_metadata_copy_dtype_validation():
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(
dsa_cache_seqlens_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device
)
nsa_seqlens_expanded_dst = torch.zeros(
dsa_seqlens_expanded_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device
)
nsa_cu_seqlens_k_dst = torch.zeros(
dsa_cu_seqlens_k_dst = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device
)
@@ -263,18 +263,18 @@ def test_fused_metadata_copy_dtype_validation():
cache_seqlens_src_wrong, # Wrong dtype: int64
cu_seqlens_k_src,
page_indices_src,
nsa_cache_seqlens_src,
dsa_cache_seqlens_src,
seqlens_expanded_src,
nsa_cu_seqlens_k_src,
dsa_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,
dsa_cache_seqlens_dst,
dsa_seqlens_expanded_dst,
dsa_cu_seqlens_k_dst,
None, # real_page_table_dst
None, # flashmla_num_splits_dst
None, # flashmla_metadata_dst
@@ -296,18 +296,18 @@ def test_fused_metadata_copy_dtype_validation():
cache_seqlens_src,
cu_seqlens_k_src,
page_indices_src,
nsa_cache_seqlens_src,
dsa_cache_seqlens_src,
seqlens_expanded_src,
nsa_cu_seqlens_k_src,
dsa_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,
dsa_cache_seqlens_dst,
dsa_seqlens_expanded_dst,
dsa_cu_seqlens_k_dst,
None,
None,
None,
@@ -369,18 +369,18 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"],
data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"],
data["src"]["dsa_cache_seqlens"],
data["src"]["seqlens_expanded"],
data["src"]["nsa_cu_seqlens_k"],
data["src"]["dsa_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["dsa_cache_seqlens"],
dst_fused["dsa_seqlens_expanded"],
dst_fused["dsa_cu_seqlens_k"],
dst_fused["real_page_table"],
dst_fused["flashmla_num_splits"],
dst_fused["flashmla_metadata"],
@@ -402,14 +402,14 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
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"
dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]
), "dsa_cache_seqlens mismatch"
assert torch.equal(
dst_ref["nsa_seqlens_expanded"], dst_fused["nsa_seqlens_expanded"]
), "nsa_seqlens_expanded mismatch"
dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"]
), "dsa_seqlens_expanded mismatch"
assert torch.equal(
dst_ref["nsa_cu_seqlens_k"], dst_fused["nsa_cu_seqlens_k"]
), "nsa_cu_seqlens_k mismatch"
dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]
), "dsa_cu_seqlens_k mismatch"
if has_real_page_table:
assert torch.equal(
@@ -458,18 +458,18 @@ def test_fused_metadata_copy_large_batch(bs):
data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"],
data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"],
data["src"]["dsa_cache_seqlens"],
data["src"]["seqlens_expanded"],
data["src"]["nsa_cu_seqlens_k"],
data["src"]["dsa_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["dsa_cache_seqlens"],
dst_fused["dsa_seqlens_expanded"],
dst_fused["dsa_cu_seqlens_k"],
dst_fused["real_page_table"],
dst_fused["flashmla_num_splits"],
dst_fused["flashmla_metadata"],
@@ -510,13 +510,13 @@ def create_test_metadata_multi(
page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device
)
nsa_cache_seqlens_src = torch.randint(
dsa_cache_seqlens_src = torch.randint(
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
)
nsa_cu_seqlens_k_src = torch.zeros(
dsa_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)
dsa_cu_seqlens_k_src[1:] = torch.cumsum(dsa_cache_seqlens_src, dim=0)
# Optional tensors
real_page_table_src = None
@@ -544,10 +544,10 @@ def create_test_metadata_multi(
page_table_1_dst = torch.zeros(
(bs, max_len + 16), dtype=torch.int32, device=device
)
nsa_cache_seqlens_dst = torch.zeros(
dsa_cache_seqlens_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device
)
nsa_cu_seqlens_k_dst = torch.zeros(
dsa_cu_seqlens_k_dst = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device
)
@@ -573,8 +573,8 @@ def create_test_metadata_multi(
"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,
"dsa_cache_seqlens_int32": dsa_cache_seqlens_dst,
"dsa_cu_seqlens_k": dsa_cu_seqlens_k_dst,
"real_page_table": real_page_table_dst,
"flashmla_num_splits": flashmla_num_splits_dst,
"flashmla_metadata": flashmla_metadata_dst,
@@ -585,8 +585,8 @@ def create_test_metadata_multi(
"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,
"dsa_cache_seqlens": dsa_cache_seqlens_src,
"dsa_cu_seqlens_k": dsa_cu_seqlens_k_src,
"real_page_table": real_page_table_src,
"flashmla_num_splits": flashmla_num_splits_src,
"flashmla_metadata": flashmla_metadata_src,
@@ -604,8 +604,8 @@ def reference_copy_for_loop(src, dst_list, bs, max_len):
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])
dst["dsa_cache_seqlens_int32"].copy_(src["dsa_cache_seqlens"])
dst["dsa_cu_seqlens_k"][1 : bs + 1].copy_(src["dsa_cu_seqlens_k"][1 : bs + 1])
if src["real_page_table"] is not None:
rows, cols = src["real_page_table"].shape
@@ -641,10 +641,10 @@ def test_fused_metadata_copy_multi_dtype_validation():
page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device
)
nsa_cache_seqlens_src = torch.randint(
dsa_cache_seqlens_src = torch.randint(
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
)
nsa_cu_seqlens_k_src = torch.zeros(
dsa_cu_seqlens_k_src = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device
)
@@ -656,10 +656,10 @@ def test_fused_metadata_copy_multi_dtype_validation():
"page_table_1": torch.zeros(
(bs, max_len + 16), dtype=torch.int32, device=device
),
"nsa_cache_seqlens": torch.zeros(
"dsa_cache_seqlens": torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device
),
"nsa_cu_seqlens_k": torch.zeros(
"dsa_cu_seqlens_k": torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device
),
}
@@ -674,8 +674,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
cache_seqlens_src_wrong, # Wrong dtype: int64
cu_seqlens_k_src,
page_indices_src,
nsa_cache_seqlens_src,
nsa_cu_seqlens_k_src,
dsa_cache_seqlens_src,
dsa_cu_seqlens_k_src,
None, # real_page_table_src
None, # flashmla_num_splits_src
None, # flashmla_metadata_src
@@ -683,8 +683,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
dst0["cache_seqlens"],
dst0["cu_seqlens_k"],
dst0["page_table_1"],
dst0["nsa_cache_seqlens"],
dst0["nsa_cu_seqlens_k"],
dst0["dsa_cache_seqlens"],
dst0["dsa_cu_seqlens_k"],
None,
None,
None,
@@ -692,8 +692,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
dst1["cache_seqlens"],
dst1["cu_seqlens_k"],
dst1["page_table_1"],
dst1["nsa_cache_seqlens"],
dst1["nsa_cu_seqlens_k"],
dst1["dsa_cache_seqlens"],
dst1["dsa_cu_seqlens_k"],
None,
None,
None,
@@ -701,8 +701,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
dst2["cache_seqlens"],
dst2["cu_seqlens_k"],
dst2["page_table_1"],
dst2["nsa_cache_seqlens"],
dst2["nsa_cu_seqlens_k"],
dst2["dsa_cache_seqlens"],
dst2["dsa_cu_seqlens_k"],
None,
None,
None,
@@ -772,8 +772,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
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"]["dsa_cache_seqlens"],
data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"],
@@ -781,8 +781,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
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["dsa_cache_seqlens_int32"],
dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"],
@@ -790,8 +790,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
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["dsa_cache_seqlens_int32"],
dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"],
@@ -799,8 +799,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
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["dsa_cache_seqlens_int32"],
dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"],
@@ -836,8 +836,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
"cache_seqlens_int32",
"cu_seqlens_k",
"page_table_1",
"nsa_cache_seqlens_int32",
"nsa_cu_seqlens_k",
"dsa_cache_seqlens_int32",
"dsa_cu_seqlens_k",
]:
if not torch.equal(dst_ref[key], dst_fused[key]):
diff = (
@@ -965,32 +965,32 @@ def test_fused_metadata_copy_multi_large_batch(bs):
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"]["dsa_cache_seqlens"],
data["src"]["dsa_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["dsa_cache_seqlens_int32"],
dst_fused_0["dsa_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["dsa_cache_seqlens_int32"],
dst_fused_1["dsa_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["dsa_cache_seqlens_int32"],
dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"],
@@ -1013,32 +1013,32 @@ def test_fused_metadata_copy_multi_large_batch(bs):
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"]["dsa_cache_seqlens"],
data["src"]["dsa_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["dsa_cache_seqlens_int32"],
dst_fused_0["dsa_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["dsa_cache_seqlens_int32"],
dst_fused_1["dsa_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["dsa_cache_seqlens_int32"],
dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"],
python/sglang/jit_kernel/tests/test_fused_store_index_cache.py
+3 -3
View File
@@ -26,7 +26,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
try:
from sglang.jit_kernel.fused_store_index_cache import (
can_use_nsa_fused_store,
can_use_dsa_fused_store,
fused_store_index_k_cache,
)
@@ -70,7 +70,7 @@ def _skip_if_unavailable(page_size: int = PAGE_SIZE):
pytest.skip("torch.float8_e4m3fn not available")
if not HAS_FUSED:
pytest.skip("fused_store_index_cache not importable")
if not can_use_nsa_fused_store(torch.bfloat16, torch.int64, page_size):
if not can_use_dsa_fused_store(torch.bfloat16, torch.int64, page_size):
pytest.skip("JIT kernel unavailable / failed to compile")
@@ -187,7 +187,7 @@ def _reference_quantize_and_store(
def _import_act_quant():
try:
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
return act_quant
except Exception:
python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py
+1 -1
View File
@@ -75,7 +75,7 @@ def test_set_mla_kv_buffer_loc_dtypes(loc_dtype):
def test_set_mla_kv_buffer_uint8_byte_layout():
"""FP8 NSA byte-layout: cache_k_nope is uint8 with [fp8(512) | scales(16)] = 528,
"""FP8 DSA byte-layout: cache_k_nope is uint8 with [fp8(512) | scales(16)] = 528,
cache_k_rope is uint8 [128]; total payload = 656 bytes."""
nope_bytes, rope_bytes = 528, 128
batch_size = 64