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
python/sglang/srt/arg_groups/deepseek_v4_hook.py
+3 -3
View File
@@ -54,13 +54,13 @@ def apply_deepseek_v4_defaults(server_args: "ServerArgs", model_arch: str) -> No
def validate_deepseek_v4_cp(server_args: "ServerArgs") -> None:
"""Validate DeepSeek V4 context-parallel configuration."""
if not server_args.enable_nsa_prefill_context_parallel:
if not server_args.enable_dsa_prefill_context_parallel:
return
if server_args.nsa_prefill_cp_mode != "round-robin-split":
if server_args.dsa_prefill_cp_mode != "round-robin-split":
raise ValueError(
f"DeepSeekV4 only supports round-robin-split CP mode, "
f"got {server_args.nsa_prefill_cp_mode}"
f"got {server_args.dsa_prefill_cp_mode}"
)
server_args.enable_dp_attention = True
python/sglang/srt/arg_groups/hisparse_hook.py
+12 -12
View File
@@ -20,13 +20,13 @@ def _hisparse_default_backend(kv_cache_dtype: str) -> str:
return "flashmla_kv" if kv_cache_dtype == "fp8_e4m3" else "flashmla_sparse"
def apply_hisparse_nsa_backend_defaults(
def apply_hisparse_dsa_backend_defaults(
server_args: "ServerArgs",
user_set_prefill: bool,
user_set_decode: bool,
kv_cache_dtype: str,
) -> bool:
"""Pick NSA backends for --enable-hisparse based on KV dtype.
"""Pick DSA backends for --enable-hisparse based on KV dtype.
BF16 KV -> flashmla_sparse, FP8 KV -> flashmla_kv. Returns True if hisparse
handled backend selection (caller should skip its own default logic).
@@ -36,29 +36,29 @@ def apply_hisparse_nsa_backend_defaults(
backend = _hisparse_default_backend(kv_cache_dtype)
if not user_set_prefill:
server_args.nsa_prefill_backend = backend
server_args.dsa_prefill_backend = backend
if not user_set_decode:
server_args.nsa_decode_backend = backend
server_args.dsa_decode_backend = backend
logger.warning(
f"HiSparse enabled ({kv_cache_dtype}): using NSA backends "
f"prefill={server_args.nsa_prefill_backend}, decode={server_args.nsa_decode_backend}."
f"HiSparse enabled ({kv_cache_dtype}): using DSA backends "
f"prefill={server_args.dsa_prefill_backend}, decode={server_args.dsa_decode_backend}."
)
return True
def validate_hisparse(server_args: "ServerArgs") -> None:
"""Validate --enable-hisparse constraints (model class, radix cache, NSA backend)."""
"""Validate --enable-hisparse constraints (model class, radix cache, DSA backend)."""
if not server_args.enable_hisparse:
return
from sglang.srt.configs.model_config import (
is_deepseek_nsa,
is_deepseek_dsa,
is_deepseek_v4,
)
hf_config = server_args.get_model_config().hf_config
is_v4_hisparse = is_deepseek_v4(hf_config)
assert is_deepseek_nsa(hf_config) or is_v4_hisparse, (
assert is_deepseek_dsa(hf_config) or is_v4_hisparse, (
"--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) "
"models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. "
)
@@ -83,13 +83,13 @@ def validate_hisparse(server_args: "ServerArgs") -> None:
server_args.kv_cache_dtype, {"flashmla_sparse", "flashmla_kv"}
)
for attr, label in [
("nsa_prefill_backend", "prefill"),
("nsa_decode_backend", "decode"),
("dsa_prefill_backend", "prefill"),
("dsa_decode_backend", "decode"),
]:
backend = getattr(server_args, attr)
if backend is not None and backend not in allowed_backends:
raise ValueError(
f"HiSparse with --kv-cache-dtype={server_args.kv_cache_dtype} requires "
f"--nsa-{label}-backend in {sorted(allowed_backends)}, "
f"--dsa-{label}-backend in {sorted(allowed_backends)}, "
f"but got {backend}."
)
python/sglang/srt/configs/model_config.py
+12 -12
View File
@@ -99,7 +99,7 @@ def _hf_attr(config, name):
return getattr(config, name, None)
def is_deepseek_nsa(config) -> bool:
def is_deepseek_dsa(config) -> bool:
return (
_hf_arch(config)
in (
@@ -121,31 +121,31 @@ def is_deepseek_v4(config) -> bool:
)
def get_nsa_index_head_dim(config: PretrainedConfig) -> int:
assert is_deepseek_nsa(config) or is_deepseek_v4(config)
def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
assert is_deepseek_dsa(config) or is_deepseek_v4(config)
return config.index_head_dim
def get_nsa_index_topk(config: PretrainedConfig) -> int:
assert is_deepseek_nsa(config)
def get_dsa_index_topk(config: PretrainedConfig) -> int:
assert is_deepseek_dsa(config)
return config.index_topk
def get_nsa_index_n_heads(config: PretrainedConfig) -> int:
assert is_deepseek_nsa(config)
def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
assert is_deepseek_dsa(config)
return config.index_n_heads
def get_num_indexer_layers(config) -> int:
"""Layer count for the global indexer-topk capturer's host buffer.
NSA models (V3.2) instantiate an Indexer on every transformer layer.
DSA models (V3.2) instantiate an Indexer on every transformer layer.
With index_topk_freq > 1 some layers reuse prev layer's topk; those still
get a slot (mirrored at the MLA call site). DSv4 has C4 indexers only on
layers whose compress_ratio == 4. Other architectures: set
num_indexer_layers on hf_text_config; 0 disables the capturer.
"""
if is_deepseek_nsa(config):
if is_deepseek_dsa(config):
return config.num_hidden_layers
if is_deepseek_v4(config):
compress_ratios = getattr(config, "compress_ratios", None) or []
@@ -329,7 +329,7 @@ class ModelConfig:
self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False)
self.is_piecewise_cuda_graph_disabled_model = (
is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
or is_deepseek_nsa(self.hf_text_config)
or is_deepseek_dsa(self.hf_text_config)
)
self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
@@ -622,8 +622,8 @@ class ModelConfig:
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_text_config.v_head_dim
self.index_head_dim = (
get_nsa_index_head_dim(self.hf_text_config)
if is_deepseek_nsa(self.hf_text_config)
get_dsa_index_head_dim(self.hf_text_config)
if is_deepseek_dsa(self.hf_text_config)
else None
)
# Handle rope scaling
python/sglang/srt/disaggregation/base/conn.py
+1 -1
View File
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
class StateType(str, enum.Enum):
MAMBA = "mamba"
SWA = "swa"
NSA = "nsa"
DSA = "dsa"
@dataclasses.dataclass
python/sglang/srt/disaggregation/decode.py
+3 -3
View File
@@ -954,7 +954,7 @@ class DecodePreallocQueue:
window_kv_indices_swa.cpu().numpy(), page_size
)
def _nsa_payload():
def _dsa_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, :seq_len
]
@@ -971,8 +971,8 @@ class DecodePreallocQueue:
state_indices.append(_mamba_payload())
elif st == StateType.SWA:
state_indices.append(_swa_payload())
elif st == StateType.NSA:
state_indices.append(_nsa_payload())
elif st == StateType.DSA:
state_indices.append(_dsa_payload())
else:
state_indices.append(None)
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+1 -1
View File
@@ -238,7 +238,7 @@ class DecodeKVCacheOffloadManager:
kv_committed_len = req.pop_committed_kv_cache()
start = start_offset
end = kv_committed_len
# Free the incremental part of the request (NSA-aware)
# Free the incremental part of the request (DSA-aware)
kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end]
self.token_to_kv_pool_allocator.free(kv_indices)
python/sglang/srt/disaggregation/mooncake/conn.py
+1 -1
View File
@@ -1012,7 +1012,7 @@ class MooncakeKVManager(CommonKVManager):
)
or rc
)
elif st in (StateType.SWA, StateType.NSA):
elif st in (StateType.SWA, StateType.DSA):
if (
target_rank_registration_info is not None
and not self.is_mla_backend
python/sglang/srt/disaggregation/mori/conn.py
+3 -3
View File
@@ -960,8 +960,8 @@ class MoriKVManager(CommonKVManager):
return self._send_mamba_state(
peer_info, src_state_indices, dst_state_indices
)
elif state_type in ("swa", "nsa"):
return self._send_swa_nsa_state(
elif state_type in ("swa", "dsa"):
return self._send_swa_dsa_state(
peer_info, src_state_indices, dst_state_indices, state_type
)
else:
@@ -1056,7 +1056,7 @@ class MoriKVManager(CommonKVManager):
return statuses
def _send_swa_nsa_state(
def _send_swa_dsa_state(
self,
peer_info: KVArgsRegisterInfo,
src_state_indices: npt.NDArray[np.int32],
python/sglang/srt/disaggregation/nixl/conn.py
+1 -1
View File
@@ -1541,7 +1541,7 @@ class NixlKVManager(CommonKVManager):
dst_gpu_id,
comp_notif,
)
elif st in (StateType.SWA, StateType.NSA):
elif st in (StateType.SWA, StateType.DSA):
if not self.is_mla_backend and self.attn_tp_size != decode_tp_size:
raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet."
python/sglang/srt/disaggregation/prefill.py
+3 -3
View File
@@ -818,7 +818,7 @@ class SchedulerDisaggregationPrefillMixin:
window_kv_indices_swa.cpu().numpy(), page_size
)
def _nsa_payload():
def _dsa_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :seq_len
]
@@ -833,8 +833,8 @@ class SchedulerDisaggregationPrefillMixin:
state_indices.append(_mamba_payload())
elif st == StateType.SWA:
state_indices.append(_swa_payload())
elif st == StateType.NSA:
state_indices.append(_nsa_payload())
elif st == StateType.DSA:
state_indices.append(_dsa_payload())
else:
state_indices.append(None)
python/sglang/srt/disaggregation/utils.py
+4 -4
View File
@@ -567,7 +567,7 @@ def setup_state_kv_args(
from sglang.srt.disaggregation.base.conn import StateType
from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool
kv_args.state_types = []
kv_args.state_data_ptrs = []
@@ -593,9 +593,9 @@ def setup_state_kv_args(
append_state_component(
kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim
)
elif isinstance(token_to_kv_pool, (NSATokenToKVPool, NPUMLATokenToKVPool)):
elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)):
if draft_token_to_kv_pool is not None and isinstance(
draft_token_to_kv_pool, NSATokenToKVPool
draft_token_to_kv_pool, DSATokenToKVPool
):
(
draft_data_ptrs,
@@ -612,7 +612,7 @@ def setup_state_kv_args(
kv_args.total_kv_layers = total_kv_layers
else:
append_state_component(
kv_args, StateType.NSA, data_ptrs, data_lens, item_lens
kv_args, StateType.DSA, data_ptrs, data_lens, item_lens
)
if (
python/sglang/srt/environ.py
+46 -4
View File
@@ -134,6 +134,41 @@ class EnvInt(EnvField):
raise ValueError(f'"{value}" is not a valid integer value')
class _DeprecatedEnvFallback:
"""Mixin for EnvField subclasses: if the canonical env var is not set,
check *deprecated_name* and emit DeprecationWarning before reading it.
Usage:
SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK")
"""
def __init__(self, default: Any, deprecated_name: str):
super().__init__(default)
self.deprecated_name = deprecated_name
def get(self) -> Any:
if os.getenv(self.name) is None:
fallback = os.getenv(self.deprecated_name)
if fallback is not None:
warnings.warn(
f"Environment variable '{self.deprecated_name}' is deprecated; "
f"use '{self.name}' instead. "
"The alias will be removed in a future release.",
DeprecationWarning,
stacklevel=2,
)
os.environ[self.name] = fallback
return super().get()
class EnvBoolWithAlias(_DeprecatedEnvFallback, EnvBool):
pass
class EnvIntWithAlias(_DeprecatedEnvFallback, EnvInt):
pass
class EnvFloat(EnvField):
def parse(self, value: str) -> float:
try:
@@ -428,11 +463,18 @@ class Envs:
SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False)
SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# NSA Backend
SGLANG_NSA_FUSE_TOPK = EnvBool(True)
SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True)
# DSA Backend (canonical names; fall back to SGLANG_NSA_* with deprecation warning)
SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK")
SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBoolWithAlias(
True, deprecated_name="SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA"
)
SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvIntWithAlias(
2048, deprecated_name="SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD"
)
SGLANG_DSA_HIP_DISABLE_PRESHUFFLE = EnvBoolWithAlias(
False, deprecated_name="SGLANG_NSA_HIP_DISABLE_PRESHUFFLE"
)
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvInt(2048)
# sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+2 -2
View File
@@ -20,7 +20,7 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_kv_cache
@@ -932,7 +932,7 @@ class AscendAttnBackend(AttentionBackend):
if (
is_prefill
and is_nsa_enable_prefill_cp()
and is_dsa_enable_prefill_cp()
and forward_batch.attn_cp_metadata is not None
):
attn_out = self.do_cp_balance_attn(
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py
+2 -2
View File
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.configs.model_config import is_deepseek_nsa
from sglang.srt.configs.model_config import is_deepseek_dsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
@@ -59,7 +59,7 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
)
def _replay(self, forward_batch: ForwardBatch):
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
self.bs - self.raw_bs
)
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
+2 -2
View File
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Dict, Union
import numpy as np
import torch
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
EAGLEDraftCudaGraphRunner,
@@ -96,7 +96,7 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
def _replay(self, forward_batch: ForwardBatch):
self.update_attr_name = self._get_update_attr_name()
self.update_attr_type = self._get_update_attr_type()
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
seq_lens_for_each_draft_step = []
for speculative_step_id in range(self.speculative_num_steps - 1):
seq_lens_cpu = forward_batch.seq_lens_cpu + speculative_step_id + 1
python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
+2 -2
View File
@@ -26,7 +26,7 @@ import numpy as np
import torch
import sglang
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.environ import envs
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
@@ -188,7 +188,7 @@ class NPUGraphRunner(CudaGraphRunner):
self.update_attr_name = self._get_update_attr_name()
self.update_attr_type = self._get_update_attr_type()
# Replay
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
if forward_batch.forward_mode.is_target_verify():
seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
+9 -9
View File
@@ -11,9 +11,9 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
is_fia_nz,
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full
from sglang.srt.layers.attention.nsa.utils import (
nsa_use_prefill_cp,
from sglang.srt.layers.attention.dsa.dsa_indexer import scattered_to_tp_attn_full
from sglang.srt.layers.attention.dsa.utils import (
dsa_use_prefill_cp,
)
from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
@@ -43,9 +43,9 @@ def forward_mha_prepare_npu(
)
)
# NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
# DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
if m.use_nsa:
if m.use_dsa:
q_lora = m.q_a_layernorm(q)
q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
_ = m.indexer(
@@ -206,7 +206,7 @@ def forward_mla_prepare_npu(
k_nope = m.kv_a_layernorm(k_nope)
# q_lora needed by indexer
if m.use_nsa:
if m.use_dsa:
q_lora = q
k_nope = k_nope.unsqueeze(1)
@@ -226,7 +226,7 @@ def forward_mla_prepare_npu(
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
# support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
@@ -359,7 +359,7 @@ def forward_dsa_prepare_npu(
if q_event is not None:
torch.npu.current_stream().wait_event(q_event)
else:
if fused_qkv_a_proj_out.shape[0] < 65535 and not nsa_use_prefill_cp(
if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp(
forward_batch
):
q_lora, k_nope, k_pe = fused_split_qk_norm(
@@ -398,7 +398,7 @@ def forward_dsa_prepare_npu(
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
# support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
python/sglang/srt/layers/attention/attention_registry.py
+16 -4
View File
@@ -1,4 +1,5 @@
import logging
import warnings
from typing import TYPE_CHECKING
from sglang.srt.configs.linear_attn_model_registry import (
@@ -96,11 +97,22 @@ def create_ascend_backend(runner):
return AscendAttnBackend(runner)
@register_attention_backend("nsa")
def create_nsa_backend(runner):
from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend
@register_attention_backend("dsa")
def create_dsa_backend(runner):
from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
return NativeSparseAttnBackend(runner)
return DeepseekSparseAttnBackend(runner)
@register_attention_backend("nsa")
def _create_nsa_compat(runner):
warnings.warn(
"attention-backend='nsa' is deprecated; use 'dsa' instead. "
"The alias will be removed in a future release.",
DeprecationWarning,
stacklevel=2,
)
return create_dsa_backend(runner)
@register_attention_backend("dsv4")
python/sglang/srt/layers/attention/base_attn_backend.py
+1 -1
View File
@@ -9,7 +9,7 @@ from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.common import is_npu
if TYPE_CHECKING:
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.speculative.spec_info import SpecInput
python/sglang/srt/layers/attention/dsa/dequant_k_cache.py
+289
View File
@@ -0,0 +1,289 @@
import torch
import triton
import triton.language as tl
def dequantize_k_cache(quant_k_cache):
return _dequantize_k_cache_fast_wrapped(quant_k_cache)
def _dequantize_k_cache_ref(
quant_k_cache: torch.Tensor, # (num_blocks, block_size, 1, bytes_per_token)
dv: int = 512,
tile_size: int = 128,
d: int = 576,
) -> torch.Tensor:
"""
De-quantize the k-cache
"""
assert dv % tile_size == 0
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_tiles = dv // tile_size
num_blocks, block_size, h_k, _ = quant_k_cache.shape
assert h_k == 1
result = torch.empty(
(num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
)
quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
input_nope = quant_k_cache[..., :dv]
input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
result[..., dv:] = input_rope
for tile_idx in range(0, num_tiles):
cur_nope = input_nope[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].to(torch.float32)
cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_nope * cur_scales
)
if original_ndim == 3:
return result.view(num_blocks, 1, -1)
else:
return result.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast_wrapped(
quant_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_blocks, block_size, _, dim_quant = quant_k_cache.shape
assert dv == 512
assert dim_quant == 656
assert tile_size == 128
quant_k_cache = quant_k_cache.view((-1, dim_quant))
output = _dequantize_k_cache_fast(quant_k_cache)
if original_ndim == 3:
return output.view(num_blocks, 1, -1)
else:
return output.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
num_tokens, dim_quant = quant_k_cache.shape
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size
assert dim_quant == 656
output = torch.empty(
(num_tokens, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_fast_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q
ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
def dequantize_k_cache_paged(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
group_size: int = 128,
) -> torch.Tensor:
"""
De-quantize the k-cache with paged layout
Args:
quant_k_cache: [total_num_tokens, 1, dim_quant] or [num_blocks, block_size, 1, dim_quant], the quantized k-cache in paged layout
page_table_1_flattened: [num_tokens], the flattened page_table_1 with the page indices in each requests concatenated together
Returns:
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
"""
dim_quant = quant_k_cache.shape[-1]
assert (
dim_quant == 656
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
quant_k_cache = quant_k_cache.view((-1, dim_quant))
# num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
# Index bounds validated in dsa_backend.init_forward_metadata
num_tokens = page_table_1_flattened.shape[0]
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size # 512 // 128 = 4
output = torch.empty(
(num_tokens, 1, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
# cdiv(512 + 64, 128) = 5
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
# [:, 512:512+4*4] = [:, 512:528]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
# [:, 528:]
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_paged_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
page_table_1_flattened,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_paged_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
page_table_1_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q
ptr_s = (
input_nope_s_ptr
+ token_id_paged * input_nope_s_stride_0
+ effective_block_id
)
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
raise Exception("UT is in quant_k_cache.py")
python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py
+331
View File
@@ -0,0 +1,331 @@
"""Multi-step precompute utilities for Native Sparse Attention backend.
This module provides optimization utilities for multi-step speculative decoding
by precomputing shared metadata once and copying it to multiple backend instances.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.attention.dsa.utils import compute_dsa_seqlens
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.speculative.spec_info import SpecInput
@dataclass
class PrecomputedMetadata:
"""Precomputed metadata shared across multiple backend instances.
Used for multi-step speculative decoding where multiple backends
need identical metadata. Precomputing once and copying N times
is much faster than computing N times.
"""
# Basic seqlens
cache_seqlens: torch.Tensor # int32, [bs]
cu_seqlens_k: torch.Tensor # int32, [bs+1]
# Page table
page_indices: torch.Tensor # int32, [bs, max_len] or [expanded_bs, max_len]
real_page_table: Optional[torch.Tensor] # int32, transformed version
# DSA seqlens
seqlens_expanded: torch.Tensor # int32, [expanded_size]
dsa_cache_seqlens: torch.Tensor # int32, [expanded_size]
dsa_cu_seqlens_k: torch.Tensor # int32, [expanded_size+1]
seqlens_expanded_size: int
# Dimensions
max_len: int # for decode/draft_extend
max_seqlen_k: int # for target_verify
# FlashMLA (optional)
flashmla_metadata: Optional[torch.Tensor] = None
def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
"""Compute cumulative sequence lengths with padding."""
assert seqlens.dtype == torch.int32
return torch.nn.functional.pad(
torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
)
class DeepseekSparseAttnBackendMTPPrecomputeMixin:
"""Mixin class providing metadata precomputation for multi-step speculative decoding.
This mixin provides the _precompute_replay_metadata method and its helpers,
which are used to optimize CUDA graph replay in multi-step scenarios.
"""
def _precompute_replay_metadata(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
forward_mode: "ForwardMode",
spec_info: Optional["SpecInput"],
) -> PrecomputedMetadata:
"""Precompute all shared metadata for multi-step backends.
This function extracts and computes all operations that are
identical across different backend instances in multi-step
speculative decoding.
Args:
bs: Batch size
req_pool_indices: Request pool indices [bs]
seq_lens: Sequence lengths [bs]
seq_lens_cpu: Sequence lengths on CPU [bs]
forward_mode: Forward mode (decode/target_verify/draft_extend)
spec_info: Speculative decoding info (for draft_extend mode)
Returns:
PrecomputedMetadata containing all shared intermediate results
"""
# Slice inputs to batch size
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
# Dispatch to mode-specific precomputation
if forward_mode.is_decode_or_idle():
return self._precompute_decode_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_target_verify():
return self._precompute_target_verify_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_draft_extend():
return self._precompute_draft_extend_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu, spec_info
)
else:
raise ValueError(f"Unsupported forward mode: {forward_mode}")
def _precompute_decode_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for normal decode mode."""
max_len = int(seq_lens_cpu.max().item())
# Convert to int32 and compute cumsum
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Get page indices from cache
page_indices = self.req_to_token[req_pool_indices, :max_len].contiguous()
# Compute DSA seqlens
dsa_cache_seqlens = compute_dsa_seqlens(
cache_seqlens, dsa_index_topk=self.dsa_index_topk
)
seqlens_expanded = cache_seqlens
seqlens_expanded_size = seqlens_expanded.shape[0]
# Compute DSA cumsum
dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
# Transform page table if needed
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None # Will use page_indices directly
# Compute FlashMLA metadata if needed
flashmla_metadata = None
if self.dsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=dsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
dsa_cache_seqlens=dsa_cache_seqlens,
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_len,
max_seqlen_k=max_len,
flashmla_metadata=flashmla_metadata,
)
def _precompute_target_verify_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for target verify mode."""
max_seqlen_k = int(
seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
)
# Cache seqlens with draft tokens
cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Page indices (repeated for each draft token)
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
seqlens_int32_cpu = [
self.speculative_num_draft_tokens + kv_len
for kv_len in seq_lens_cpu.tolist()
]
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seqlens_int32_cpu,
strict=True,
)
]
)
# Compute DSA seqlens
dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# DSA cumsum
dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.dsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=dsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
dsa_cache_seqlens=dsa_cache_seqlens,
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=-1, # Not used in this mode
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
def _precompute_draft_extend_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
spec_info: "SpecInput",
) -> PrecomputedMetadata:
"""Precompute metadata for draft extend mode."""
max_seqlen_k = int(seq_lens_cpu.max().item())
# Cache seqlens
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Extend seqlens from spec_info: num_accept_tokens already includes
# the bonus token (drafts + 1).
extend_seq_lens = spec_info.num_accept_tokens[:bs]
extend_seq_lens_cpu = extend_seq_lens.tolist()
# Page indices (repeated per accept length)
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(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seq_lens_cpu.tolist(),
strict=True,
)
]
)
# Compute DSA seqlens
dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# DSA cumsum
dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.dsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=dsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
dsa_cache_seqlens=dsa_cache_seqlens,
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_seqlen_k,
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
# Backward-compat alias
DeepseekSparseAttnBackendMTPPrecomputeMixin = (
DeepseekSparseAttnBackendMTPPrecomputeMixin
)
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+1746
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py
+407
View File
@@ -0,0 +1,407 @@
"""
Verification utilities for DSA 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 DSA 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_dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32.clone()
fused_dsa_seqlens_expanded = metadata.dsa_seqlens_expanded.clone()
fused_dsa_cu_seqlens_k = metadata.dsa_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_dsa_cache_seqlens = torch.zeros_like(metadata.dsa_cache_seqlens_int32)
ref_dsa_seqlens_expanded = torch.zeros_like(metadata.dsa_seqlens_expanded)
ref_dsa_cu_seqlens_k = torch.zeros_like(metadata.dsa_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_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
elif forward_mode.is_target_verify():
# Target verify mode
ref_page_table_1[:, : precomputed.max_seqlen_k].copy_(precomputed.page_indices)
ref_dsa_seqlens_expanded.copy_(precomputed.seqlens_expanded)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_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_dsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded)
ref_dsa_cache_seqlens[:size].copy_(precomputed.dsa_cache_seqlens)
# Copy DSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_dsa_cu_seqlens_k[1 : 1 + size].copy_(precomputed.dsa_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 dsa_cache_seqlens only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"dsa_cache_seqlens",
fused_dsa_cache_seqlens,
ref_dsa_cache_seqlens,
)
else: # TARGET_VERIFY or DRAFT_EXTEND
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_cache_seqlens",
fused_dsa_cache_seqlens[:size],
ref_dsa_cache_seqlens[:size],
)
# Compare dsa_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(
"dsa_seqlens_expanded",
fused_dsa_seqlens_expanded[:size],
ref_dsa_seqlens_expanded[:size],
)
# Compare dsa_cu_seqlens_k only for the region that was updated
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_cu_seqlens_k",
fused_dsa_cu_seqlens_k[: 1 + size],
ref_dsa_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 DSA metadata object for backend 0
metadata1: The DSA metadata object for backend 1
metadata2: The DSA 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_dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32.clone()
fused_dsa_cu_seqlens_k = metadata.dsa_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,
"dsa_cache_seqlens": fused_dsa_cache_seqlens,
"dsa_cu_seqlens_k": fused_dsa_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_dsa_cache_seqlens = torch.zeros_like(metadata.dsa_cache_seqlens_int32)
ref_dsa_cu_seqlens_k = torch.zeros_like(metadata.dsa_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_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
# Copy DSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_dsa_cu_seqlens_k[1 : 1 + size].copy_(
precomputed.dsa_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,
"dsa_cache_seqlens": ref_dsa_cache_seqlens,
"dsa_cu_seqlens_k": ref_dsa_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,
"dsa_cache_seqlens",
fused["dsa_cache_seqlens"],
ref["dsa_cache_seqlens"],
)
# DECODE mode uses bs for dsa_cu_seqlens_k size
check_tensor_equal(
idx,
"dsa_cu_seqlens_k",
fused["dsa_cu_seqlens_k"][: bs + 1],
ref["dsa_cu_seqlens_k"][: bs + 1],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
idx,
"real_page_table",
fused["real_page_table"][:rows, :cols],
ref["real_page_table"][:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
# DECODE mode uses bs + 1 for flashmla_num_splits
check_tensor_equal(
idx,
"flashmla_num_splits",
fused["flashmla_num_splits"][: bs + 1],
ref["flashmla_num_splits"][: bs + 1],
)
check_tensor_equal(
idx,
"flashmla_metadata",
fused["flashmla_metadata"],
ref["flashmla_metadata"],
)
python/sglang/srt/layers/attention/dsa/index_buf_accessor.py
+814
View File
@@ -0,0 +1,814 @@
from typing import TYPE_CHECKING
import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import get_bool_env_var, is_hip
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# aiter cp_gather kernel with preshuffle=True is only valid when the indexer
# uses the page_size=64 preshuffle layout (i.e. when the matching MQA gluon path
# is also enabled).
_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
if _use_aiter_preshuffle:
from aiter.ops.cache import cp_gather_indexer_k_quant_cache
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
"""
k: data, 128 item per token, fp8
s: scale, 1 item per token, fp32
"""
class GetK:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
index_k_fp8 = torch.empty(
(seq_len_, pool.index_head_dim),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][: pool.page_size * pool.index_head_dim].view(-1, pool.index_head_dim)
return index_k_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim), uint8
"""
# can handle per 128B instead of per element
# page_indices: (num_pages,), element := a page index
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_page = pool.page_size * pool.index_head_dim
num_k_bytes_per_token = pool.index_head_dim
# buf: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4), uint8
# flat_buf: (whatever,), uint8
flat_buf = buf.flatten()
# flat_indices: (num_pages, num_k_bytes_per_page), int32, element := an index into flat_buf that we want to access
flat_indices = (page_indices * buf_numel_per_page)[:, None] + torch.arange(
num_k_bytes_per_page, dtype=torch.int32, device="cuda"
)[None, :]
flat_indices = flat_indices.flatten()[: seq_len * num_k_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 128)
@classmethod
def triton(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering K data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, index_head_dim), uint8
"""
return _get_k_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetS:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
assert pool.index_head_dim // pool.quant_block_size == 1
index_k_scale_fp8 = torch.empty(
(seq_len_, 4),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_scale_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][pool.page_size * pool.index_head_dim :].view(-1, 4)
return index_k_scale_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim // quant_block_size), uint8
"""
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_page = buf.shape[1] - pool.page_size * pool.index_head_dim
num_s_bytes_per_token = pool.index_head_dim // pool.quant_block_size * 4
s_offset_in_page = pool.page_size * pool.index_head_dim
flat_buf = buf.flatten()
flat_indices = (
(page_indices * buf_numel_per_page)[:, None]
+ torch.arange(num_s_bytes_per_page, dtype=torch.int32, device="cuda")[
None, :
]
+ s_offset_in_page
)
flat_indices = flat_indices.flatten()[: seq_len * num_s_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 4)
@classmethod
def triton(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering S (scale) data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, 4), uint8
"""
return _get_s_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetKAndS:
@classmethod
def execute(cls, *args, **kwargs):
# The aiter path uses cp_gather_indexer_k_quant_cache(preshuffle=True),
# which only matches the layout produced when the rest of the indexer
# is on the page_size=64 preshuffle path. Otherwise fall back to the
# triton implementation (which works on the page_size=1 legacy layout).
if _use_aiter_preshuffle:
return cls.aiter(*args, **kwargs)
return cls.triton(*args, **kwargs)
@classmethod
def aiter(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
page_size = pool.page_size
index_head_dim = pool.index_head_dim
quant_block_size = pool.quant_block_size
scale_elems = index_head_dim // quant_block_size
kv_cache = buf.view(-1, page_size, index_head_dim + scale_elems * 4).view(
fp8_dtype
)
dst_k = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
dst_scale = torch.empty(
(seq_len_sum, scale_elems * 4), dtype=torch.uint8, device=buf.device
)
cu_seq_lens = torch.zeros(
seq_len_tensor.shape[0] + 1, dtype=torch.int32, device=buf.device
)
torch.cumsum(seq_len_tensor.to(torch.int32), dim=0, out=cu_seq_lens[1:])
cp_gather_indexer_k_quant_cache(
kv_cache,
dst_k.view(fp8_dtype),
dst_scale,
page_indices.to(torch.int32),
cu_seq_lens,
preshuffle=True,
)
return dst_k, dst_scale
@classmethod
def triton(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
"""
Triton implementation for gathering both K and S data from paged buffer in a single call.
:param page_indices: (num_pages,), int32/int64
:param seq_len_tensor: (num_pages,), int32/int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of all sequence len, int32
:return: tuple of (k_fp8, k_scale) where
k_fp8: (seq_len, index_head_dim), uint8
k_scale: (seq_len, 4), uint8
"""
return _get_k_and_s_triton(
buf=buf,
page_indices=page_indices,
seq_lens=seq_len_tensor,
seq_len_sum=seq_len_sum,
max_seq_len=max_seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class SetK:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
buf[
page_index,
offset * pool.index_head_dim : (offset + 1) * pool.index_head_dim,
] = index_k[i].view(torch.uint8)
@classmethod
def torch_fast(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_token = pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ (loc_token_offset_in_page * num_k_bytes_per_token)[:, None]
+ torch.arange(num_k_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
num_k_bytes_total = num_tokens_to_write * num_k_bytes_per_token
flat_indices = flat_indices.flatten()[:num_k_bytes_total]
flat_buf[flat_indices] = index_k.view(torch.uint8).flatten()
class SetS:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
start = pool.page_size * pool.index_head_dim
buf[page_index, start + offset * 4 : start + (offset + 1) * 4] = (
index_k_scale[i].view(torch.uint8)
)
@classmethod
def torch_fast(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_token = 4
s_offset_in_page = pool.page_size * pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ s_offset_in_page
+ (loc_token_offset_in_page * num_s_bytes_per_token)[:, None]
+ torch.arange(num_s_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
number_s_bytes_total = num_tokens_to_write * num_s_bytes_per_token
flat_indices = flat_indices.flatten()[:number_s_bytes_total]
flat_buf[flat_indices] = index_k_scale.view(torch.uint8).flatten()
class SetKAndS:
@classmethod
def execute(cls, *args, buf, **kwargs):
if 0:
# print("SetK, SetS comparison test")
buf_cloned = buf.clone()
cls.vanilla(*args, **kwargs, buf=buf)
cls.triton(*args, **kwargs, buf=buf_cloned)
def _clear_token_0(target):
target[0, :128] = target[0, 64 * 128 : 64 * 128 + 4] = 0
_clear_token_0(buf)
_clear_token_0(buf_cloned)
assert torch.all(
buf == buf_cloned
), f"{buf=} {buf_cloned=} {kwargs['loc'].to_list()=}"
return
cls.triton(*args, **kwargs, buf=buf)
@classmethod
def vanilla(cls, pool, buf, loc, index_k, index_k_scale):
SetK.execute(pool=pool, buf=buf, loc=loc, index_k=index_k)
SetS.execute(pool=pool, buf=buf, loc=loc, index_k_scale=index_k_scale)
@classmethod
def triton(cls, pool, buf, loc, index_k, index_k_scale):
loc = loc.to(torch.int64)
_set_k_and_s_triton(
buf=buf,
loc=loc,
index_k=index_k,
index_k_scale=index_k_scale,
page_size=pool.page_size,
)
def _set_k_and_s_triton(
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
index_k_scale: torch.Tensor,
page_size: int,
):
"""
:param buf: (num_pages, page_size 64 * (128B data + 4B scale)), uint8
:param loc: (num_tokens_to_write,), int, element := the token index to write to
:param index_k: (num_tokens_to_write, 128 elem), fp8
:param index_k_scale: (num_tokens_to_write, 1 elem), fp32
:return:
"""
num_pages, buf_numel_per_page = buf.shape
(num_tokens_to_write,) = loc.shape
num_tokens_to_write_, index_head_dim = index_k.shape
# Handle both 1D (num_tokens,) and 2D (num_tokens, 1) shapes for index_k_scale
if index_k_scale.ndim == 1:
num_tokens_to_write__ = index_k_scale.shape[0]
scale_dim = 1
elif index_k_scale.ndim == 2:
num_tokens_to_write__, scale_dim = index_k_scale.shape
else:
raise ValueError(
f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}"
)
assert buf_numel_per_page == page_size * (128 + 4)
assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__
assert index_head_dim == 128
assert scale_dim == 1
if _is_hip:
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert page_size == 64
assert buf.dtype == torch.uint8
assert loc.dtype == torch.int64, f"{loc.dtype=}" # can be int32
if _is_fp8_fnuz:
assert index_k.dtype == torch.float8_e4m3fnuz
else:
assert index_k.dtype == torch.float8_e4m3fn
assert index_k_scale.dtype == torch.float32
assert buf.is_contiguous()
assert loc.is_contiguous()
assert index_k.is_contiguous()
assert index_k_scale.is_contiguous()
if _is_fp8_fnuz:
buf_fp8 = buf.view(torch.float8_e4m3fnuz)
else:
buf_fp8 = buf.view(torch.float8_e4m3fn)
buf_fp32 = buf.view(torch.float32)
_set_k_and_s_triton_kernel[(num_tokens_to_write,)](
buf_fp8,
buf_fp32,
loc,
index_k,
index_k_scale,
index_k.stride(0),
PAGE_SIZE=page_size,
BUF_NUMEL_PER_PAGE=buf_numel_per_page,
NUM_K_ELEMS_PER_TOKEN=index_head_dim,
S_OFFSET_NBYTES_IN_PAGE=page_size * index_head_dim,
)
@triton.jit
def _set_k_and_s_triton_kernel(
buf_fp8_ptr,
buf_fp32_ptr,
loc_ptr,
index_k_ptr,
index_k_scale_ptr,
index_k_ptr_stride_0,
PAGE_SIZE: tl.constexpr,
BUF_NUMEL_PER_PAGE: tl.constexpr,
NUM_K_ELEMS_PER_TOKEN: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
):
token_id = tl.program_id(0)
loc = tl.load(loc_ptr + token_id)
in_k_offsets = token_id * index_k_ptr_stride_0 + tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
# no need for `mask`, since we read 128B for k and 4B for scale, both pow of 2
k = tl.load(index_k_ptr + in_k_offsets)
k_scale = tl.load(index_k_scale_ptr + token_id)
loc_page_index = loc // PAGE_SIZE
loc_token_offset_in_page = loc % PAGE_SIZE
out_k_offsets = (
loc_page_index * BUF_NUMEL_PER_PAGE
+ loc_token_offset_in_page * NUM_K_ELEMS_PER_TOKEN
+ tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
)
# "//4" b/c it is fp32 instead of uint8
out_s_offset = (
loc_page_index * BUF_NUMEL_PER_PAGE // 4
+ S_OFFSET_NBYTES_IN_PAGE // 4
+ loc_token_offset_in_page
)
tl.store(buf_fp8_ptr + out_k_offsets, k)
tl.store(buf_fp32_ptr + out_s_offset, k_scale)
def _get_k_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather K (key) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, index_head_dim), uint8
"""
num_pages, buf_numel_per_page = buf.shape
# Allocate output
out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_k_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
index_head_dim,
BLOCK_SIZE=128,
)
return out
@triton.jit
def _get_k_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 128 bytes from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# buf[page_index, token_offset_in_page * index_head_dim : ...]
src_base_offset = (
page_index * buf_numel_per_page + token_offset_in_page * index_head_dim
)
# Load 128 bytes (index_head_dim elements)
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < index_head_dim
data = tl.load(buf_ptr + src_base_offset + offsets, mask=mask)
# Store to output
dst_offset = token_id * index_head_dim
tl.store(out_ptr + dst_offset + offsets, data, mask=mask)
def _get_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather S (scale) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, 4), uint8 (representing fp32 scale)
"""
num_pages, buf_numel_per_page = buf.shape
s_offset_in_page = page_size * index_head_dim # Scales start after K data
# Allocate output
out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_s_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
s_offset_in_page,
)
return out
@triton.jit
def _get_s_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
s_offset_in_page: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 4 bytes (fp32 scale) from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# Scales are stored after K data: page_size * index_head_dim offset
# buf[page_index, s_offset_in_page + token_offset_in_page * 4 : ...]
src_base_offset = (
page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4
)
# Load 4 bytes (fp32 scale)
offsets = tl.arange(0, 4)
data = tl.load(buf_ptr + src_base_offset + offsets)
# Store to output
dst_offset = token_id * 4
tl.store(out_ptr + dst_offset + offsets, data)
def _get_k_and_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Fused gather of both K (key) and S (scale) data from paged buffer using Triton.
This is more efficient than calling GetK and GetS separately.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_lens: tensor of sequence lens, int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of sequence len, int32
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: tuple of (k_out, s_out) where
k_out: (seq_len, index_head_dim), uint8
s_out: (seq_len, 4), uint8
"""
# Allocate outputs
k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device=buf.device)
_, buf_numel_per_page = buf.shape
_, page_indice_batch_offset = page_indices.shape
s_offset_in_page = page_size * index_head_dim
# Launch kernel with one thread per token
BLOCK_SIZE = 256
BLOCK_SIZE_K = 128
num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
seq_num = seq_lens.shape[0]
grid = (seq_num, num_token_blocks, num_k_threads)
seq_num_pow2 = 1
while seq_num_pow2 < seq_num:
seq_num_pow2 *= 2
_get_k_and_s_triton_kernel[grid](
buf_ptr=buf,
page_indices_ptr=page_indices,
k_out_ptr=k_out,
s_out_ptr=s_out,
seq_len_ptr=seq_lens,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
return k_out, s_out
@triton.jit
def _get_k_and_s_triton_kernel(
buf_ptr,
page_indices_ptr,
k_out_ptr,
s_out_ptr,
seq_len_ptr,
seq_len_num_pow: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
s_offset_in_page: tl.constexpr,
page_indice_batch_offset,
BLOCK_SIZE: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
"""
Fused kernel that gathers both K and S data in a single pass.
Each program handles one token (seq_len tokens total).
Loads 128 bytes (K) + 4 bytes (S) from the appropriate page.
"""
batch_id = tl.program_id(0)
block_token_start = tl.program_id(1) * BLOCK_SIZE
thread_idx = tl.program_id(2)
# Define the token range within the block and the K dimension range handled by the thread.
token_ids_in_block = tl.arange(0, BLOCK_SIZE)
token_ids = block_token_start + token_ids_in_block
k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
seq_len = tl.load(seq_len_ptr + batch_id)
token_valid_mask = token_ids < seq_len
pre_batch_idx = tl.arange(0, seq_len_num_pow)
mask_pre_batch_idx = pre_batch_idx < batch_id
prev_seq_lens = tl.load(seq_len_ptr + pre_batch_idx, mask=mask_pre_batch_idx)
batch_token_offset = tl.sum(prev_seq_lens)
# Batch calculate the page index and in-page offset of each token.
page_idx = token_ids // page_size
token_offset_in_page = token_ids % page_size
page_indices_base = batch_id * page_indice_batch_offset
page_idx_valid_mask = page_idx < page_indice_batch_offset
page_index = tl.load(
page_indices_ptr + page_idx + page_indices_base,
mask=token_valid_mask & page_idx_valid_mask,
)
# ===== Load K data =====
# The address calculation logic for K: page_index * total number of elements in a single page + K offset of the token within the page.
k_src_token_offset = token_offset_in_page * index_head_dim
k_src_base_offset = page_index * buf_numel_per_page + k_src_token_offset
k_load_addr = buf_ptr + k_src_base_offset[:, None] + k_offsets[None, :]
k_dim_mask = k_offsets[None, :] < index_head_dim
k_mask = token_valid_mask[:, None] & k_dim_mask
k_data = tl.load(k_load_addr, mask=k_mask, other=0)
# Store K to output
k_dst_token_offset = batch_token_offset + token_ids
k_dst_base_offset = k_dst_token_offset * index_head_dim
k_store_addr = k_out_ptr + k_dst_base_offset[:, None] + k_offsets[None, :]
tl.store(k_store_addr, k_data, mask=k_mask)
# ===== Load S data =====
# The address calculation logic for S: page_index * total number of elements in a single page + starting offset of S within the page + offset of token within S in the page
s_src_token_offset = s_offset_in_page + token_offset_in_page * 4
s_src_base_offset = page_index * buf_numel_per_page + s_src_token_offset
s_offsets = tl.arange(0, 4)
s_load_addr = buf_ptr + s_src_base_offset[:, None] + s_offsets[None, :]
s_mask = token_valid_mask[:, None] & (s_offsets[None, :] < 4)
s_data = tl.load(s_load_addr, mask=s_mask, other=0)
# Store S to output
s_dst_token_offset = batch_token_offset + token_ids
s_dst_base_offset = s_dst_token_offset * 4
s_store_addr = s_out_ptr + s_dst_base_offset[:, None] + s_offsets[None, :]
tl.store(s_store_addr, s_data, mask=s_mask)
python/sglang/srt/layers/attention/dsa/quant_k_cache.py
+449
View File
@@ -0,0 +1,449 @@
import torch
import triton
import triton.language as tl
def quantize_k_cache(cache_k):
return _quantize_k_cache_fast_wrapped(cache_k)
def quantize_k_cache_separate(
k_nope: torch.Tensor,
k_rope: torch.Tensor,
tile_size: int = 128,
):
"""
Quantize k_nope and k_rope separately without concat, returns two tensors.
This avoids the concat operation and enables direct reuse of set_mla_kv_buffer_triton
by returning two separate byte tensors for the nope and rope parts.
Args:
k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
Must have dim_nope=512 for FP8 MLA quantization
k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
Must have dim_rope=64 for FP8 MLA quantization
tile_size: quantization tile size (default 128)
Returns:
Tuple of (nope_part, rope_part) where:
- nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
- rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
"""
# Squeeze middle dimension if present
k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
num_tokens = k_nope_2d.shape[0]
dim_nope = k_nope_2d.shape[1]
dim_rope = k_rope_2d.shape[1]
# Validate dimensions for FP8 MLA
if dim_nope != 512:
raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
if dim_rope != 64:
raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
if k_rope_2d.shape[0] != num_tokens:
raise ValueError(
f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
)
return _quantize_k_cache_fast_separate(
k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
)
# Copied from original
def _quantize_k_cache_ref(
input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d)
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
"""
Quantize the k-cache
Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size()
For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md
"""
assert dv % tile_size == 0
num_tiles = dv // tile_size
num_blocks, block_size, h_k, d = input_k_cache.shape
assert h_k == 1
input_k_cache = input_k_cache.squeeze(2) # [num_blocks, block_size, d]
input_elem_size = input_k_cache.element_size()
result = torch.empty(
(num_blocks, block_size, dv + num_tiles * 4 + input_elem_size * (d - dv)),
dtype=torch.float8_e4m3fn,
device=input_k_cache.device,
)
result_k_nope_part = result[..., :dv]
result_k_scale_factor = result[..., dv : dv + num_tiles * 4].view(torch.float32)
result_k_rope_part = result[..., dv + num_tiles * 4 :].view(input_k_cache.dtype)
result_k_rope_part[:] = input_k_cache[..., dv:]
for tile_idx in range(0, num_tiles):
cur_scale_factors_inv = (
torch.abs(
input_k_cache[..., tile_idx * tile_size : (tile_idx + 1) * tile_size]
)
.max(dim=-1)
.values
/ 448.0
) # [num_blocks, block_size]
result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv
cur_scale_factors_inv.unsqueeze_(-1) # [num_blocks, block_size, 1]
cur_quantized_nope = (
input_k_cache[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].float()
/ cur_scale_factors_inv.float()
).to(torch.float8_e4m3fn)
result_k_nope_part[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_quantized_nope
)
result = result.view(num_blocks, block_size, 1, -1)
return result
def _quantize_k_cache_fast_wrapped(
input_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
# TODO the final API may be 2D instead of 4D, thus we convert them here
num_blocks, block_size, _, dim_nope_and_rope = input_k_cache.shape
assert dv == 512
assert dim_nope_and_rope == 512 + 64
assert tile_size == 128
input_k_cache = input_k_cache.view((-1, dim_nope_and_rope))
# TODO deliberately split into two tensors, then upstream can provide the two tensors instead of concat into one
k_nope = input_k_cache[:, :dv]
k_rope = input_k_cache[:, dv:]
output = _quantize_k_cache_fast(k_nope=k_nope, k_rope=k_rope)
return output.view(num_blocks, block_size, 1, -1)
def _quantize_k_cache_fast(k_nope, k_rope, group_size: int = 128):
"""
:param k_nope: (num_tokens, dim_nope 512)
:param k_rope: (num_tokens, dim_rope 64)
"""
assert k_nope.dtype == torch.bfloat16
assert k_rope.dtype == torch.bfloat16
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_
assert dim_nope == 512
assert dim_rope == 64
assert k_nope.dtype == k_rope.dtype
num_tiles = dim_nope // group_size
assert k_nope.stride(1) == 1
assert k_rope.stride(1) == 1
output = torch.empty(
(num_tokens, dim_nope + num_tiles * 4 + k_rope.element_size() * dim_rope),
dtype=torch.float8_e4m3fn,
device=k_nope.device,
)
output_nope_q = output[..., :dim_nope]
output_nope_s = output[..., dim_nope : dim_nope + num_tiles * 4].view(torch.float32)
output_rope = output[..., dim_nope + num_tiles * 4 :].view(torch.bfloat16)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
NUM_NOPE_BLOCKS = dim_nope // group_size
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output_nope_q,
output_nope_s,
output_rope,
k_nope,
k_rope,
output_nope_q.stride(0),
output_nope_s.stride(0),
output_rope.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
return output
def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
"""
Quantize k_nope and k_rope in a single Triton kernel, directly outputting two separate tensors.
This avoids packing/unpacking and enables direct use with set_mla_kv_buffer_triton.
:param k_nope: (num_tokens, dim_nope 512) bfloat16
:param k_rope: (num_tokens, dim_rope 64) bfloat16
:param group_size: quantization tile size (default 128, kernel is tuned for this value)
:return: Tuple of (nope_part_u8, rope_part_u8)
- nope_part_u8: (num_tokens, 1, nope_part_bytes) uint8, layout [nope_fp8(dim_nope) | scales(num_tiles*4)]
- rope_part_u8: (num_tokens, 1, rope_part_bytes) uint8, layout [rope_bf16_bytes(dim_rope*2)]
"""
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_, f"k_nope and k_rope must have same num_tokens"
# Ensure contiguous tensors for kernel
k_nope = k_nope.contiguous()
k_rope = k_rope.contiguous()
num_tiles = dim_nope // group_size
# Calculate byte sizes based on validated dimensions
# nope_part: [FP8 quantized data (dim_nope bytes)] + [FP32 scales (num_tiles * 4 bytes)]
# rope_part: [BF16 raw data (dim_rope * 2 bytes)]
nope_part_bytes = (
dim_nope + num_tiles * 4
) # e.g., 512 + 4*4 = 528 for dim_nope=512, group_size=128
rope_part_bytes = (
dim_rope * k_rope.element_size()
) # e.g., 64 * 2 = 128 for dim_rope=64, BF16
# Allocate two separate output buffers (as uint8 for direct byte-level access)
nope_part_u8 = torch.empty(
(num_tokens, nope_part_bytes), dtype=torch.uint8, device=k_nope.device
)
rope_part_u8 = torch.empty(
(num_tokens, rope_part_bytes), dtype=torch.uint8, device=k_rope.device
)
# Create typed views for the kernel to write into
# Fixed byte layout for nope_part: [nope_fp8 (dim_nope bytes) | scales_fp32 (num_tiles*4 bytes)]
# Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
rope_view = rope_part_u8.view(torch.bfloat16)
# Kernel launch parameters
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
NUM_NOPE_BLOCKS = dim_nope // group_size
# Use the same kernel as _quantize_k_cache_fast (reuse existing implementation)
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
nope_q_view,
nope_s_view,
rope_view,
k_nope,
k_rope,
nope_q_view.stride(0),
nope_s_view.stride(0),
rope_view.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
# Add middle dimension for compatibility with set_mla_kv_buffer_triton
return nope_part_u8.unsqueeze(1), rope_part_u8.unsqueeze(1)
@triton.jit
def _quantize_k_cache_fast_kernel(
output_nope_q_ptr,
output_nope_s_ptr,
output_rope_ptr,
k_nope_ptr,
k_rope_ptr,
output_nope_q_stride_0: int,
output_nope_s_stride_0: int,
output_rope_stride_0: int,
k_nope_stride_0: int,
k_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
FP8_MIN: tl.constexpr,
FP8_MAX: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. quant nope
effective_block_id = raw_block_id
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_NOPE
ptr = k_nope_ptr + token_id * k_nope_stride_0 + offs
y = tl.load(ptr, mask=mask, other=0.0).to(tl.float32)
# the ref impl do not have a `tl.maximum(... eps)`, so we remove it here
y_s = tl.max(tl.abs(y)) / FP8_MAX
y_s_inv = 1.0 / y_s
y_q = tl.clamp(y * y_s_inv, FP8_MIN, FP8_MAX).to(
output_nope_q_ptr.dtype.element_ty
)
dst_q_ptr = output_nope_q_ptr + token_id * output_nope_q_stride_0 + offs
dst_s_ptr = (
output_nope_s_ptr + token_id * output_nope_s_stride_0 + effective_block_id
)
tl.store(dst_q_ptr, y_q, mask=mask)
tl.store(dst_s_ptr, y_s)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = k_rope_ptr + token_id * k_rope_stride_0 + offs
dst_ptr = output_rope_ptr + token_id * output_rope_stride_0 + offs
data = tl.load(src_ptr, mask=mask)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
import dequant_k_cache
for num_blocks, block_size in [
(1, 1),
(10, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
ref_quant = _quantize_k_cache_ref(input_k_cache)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
ref_ref_dequant = dequant_k_cache._dequantize_k_cache_slow(ref_quant)
ref_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(ref_quant)
actual_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(
actual_quant
)
print(f"{ref_ref_dequant=}")
print(f"{actual_actual_dequant=}")
print(f"{actual_actual_dequant - ref_ref_dequant=}")
print(f"{torch.mean(ref_ref_dequant - actual_actual_dequant)=}")
# TODO too different?
torch.testing.assert_close(
ref_ref_dequant, ref_actual_dequant, atol=0.2, rtol=0.2
)
torch.testing.assert_close(
ref_ref_dequant, actual_actual_dequant, atol=0.2, rtol=0.2
)
# test dequant_k_cache_paged
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
actual_dequant_paged = dequant_k_cache.dequantize_k_cache_paged(
actual_quant, page_table_1
).reshape(actual_actual_dequant.shape)
print(f"{torch.mean(actual_actual_dequant - actual_dequant_paged)=}")
torch.testing.assert_close(
ref_ref_dequant, actual_dequant_paged, atol=0.2, rtol=0.2
)
print("Passed")
# Test quantize_k_cache_separate: verify output matches concat path
print("\nTesting quantize_k_cache_separate...")
for num_tokens in [64, 100]:
dim_nope = 512
dim_rope = 64
k_nope = torch.randn(
num_tokens, 1, dim_nope, dtype=torch.bfloat16, device="cuda"
)
k_rope = torch.randn(
num_tokens, 1, dim_rope, dtype=torch.bfloat16, device="cuda"
)
# Old path: concat then quantize
k_concat = torch.cat([k_nope, k_rope], dim=-1).squeeze(1) # (num_tokens, 576)
old_output = quantize_k_cache(k_concat.unsqueeze(1).unsqueeze(1)) # 4D input
old_output = old_output.squeeze(1).squeeze(1) # Back to (num_tokens, 656)
# New path: quantize separately
nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
new_bytes = torch.cat([nope_part.squeeze(1), rope_part.squeeze(1)], dim=-1)
# Compare byte-level equality
old_bytes = old_output.view(torch.uint8)
if old_bytes.shape != new_bytes.shape:
raise RuntimeError(
f"Shape mismatch: {old_bytes.shape} vs {new_bytes.shape}"
)
diff_bytes = (old_bytes != new_bytes).sum().item()
if diff_bytes > 0:
max_diff = (old_bytes.float() - new_bytes.float()).abs().max().item()
raise RuntimeError(
f"quantize_k_cache_separate output doesn't match concat path: "
f"{diff_bytes} differing bytes, max_diff={max_diff}"
)
print(f" num_tokens={num_tokens}: PASSED (outputs match byte-wise)")
print("quantize_k_cache_separate tests passed!")
print("\nDo benchmark...")
for num_blocks, block_size in [
(1, 64),
(64, 64),
(128, 64),
(256, 64),
(512, 64),
(1024, 64),
(2048, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
def run_ans():
return dequant_k_cache.dequantize_k_cache_paged(actual_quant, page_table_1)
ans_time: float = triton.testing.do_bench(run_ans, warmup=10, rep=20) / 1000 # type: ignore
print(f"seq_kv: {num_blocks * block_size}, time: {ans_time * 1e6: 4.0f} us")
python/sglang/srt/layers/attention/dsa/tilelang_kernel.py
+2589
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/dsa/transform_index.py
+144
View File
@@ -0,0 +1,144 @@
from typing import List, Optional
import torch
import triton
import triton.language as tl
def transform_index_page_table_prefill(**kwargs):
return transform_index_page_table_prefill_ref(**kwargs)
def transform_index_page_table_decode(**kwargs):
return transform_index_page_table_decode_ref(**kwargs)
@triton.jit
def transform_index_page_table_decode_kernel(
page_table_ptr: torch.Tensor,
topk_indices_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_size: tl.constexpr,
max_seqlen_k: tl.constexpr,
):
TOPK: tl.constexpr = 2048
req_id = tl.program_id(0)
page_table_ptr = page_table_ptr + req_id * max_seqlen_k
topk_indices_ptr = topk_indices_ptr + req_id * TOPK
result_ptr = result_ptr + req_id * TOPK
offset = tl.arange(0, TOPK) # topk should be 2048
loaded_topk_indices = tl.load(topk_indices_ptr + offset)
mask = loaded_topk_indices >= 0
loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask)
tl.store(result_ptr + offset, loaded_kv_indices, mask=mask)
tl.store(result_ptr + offset, -1, mask=~mask)
def transform_index_page_table_decode_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
"""
Transform the page table according to topk indices for sparse topk attention.
Args:
page_table: [qo_len, max_seqlen_k], the original page table
topk_indices: [qo_len, topk], the topk indices for each query position
Returns:
transformed_page_table: [qo_len, topk], the transformed page table
For out-of-bound indices in topk_indices, this should be filled with -1.
"""
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
assert topk_indices.shape[1] == 2048
qo_len = topk_indices.shape[0]
max_seqlen_k = page_table.shape[1]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
# Launch triton kernel
grid = (qo_len,)
transform_index_page_table_decode_kernel[grid](
page_table,
topk_indices,
result,
page_size,
max_seqlen_k=max_seqlen_k,
)
return result
def transform_index_page_table_prefill_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
# TODO(baizhou): can be implemented with another triton kernel
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_fast(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
def transform_index_page_table_decode_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert result.shape == topk_indices.shape
torch.gather(
page_table.to(result.dtype),
dim=1,
index=topk_indices.clamp(min=0),
out=result,
)
result[topk_indices < 0] = -1
return result
def transform_index_page_table_prefill_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_ref(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
if __name__ == "__main__":
bs, topk, max_seqlen = 10, 2048, 3000
page_table = torch.randint(0, 100, (bs, max_seqlen), device="cuda")
topk_indices = torch.full((bs, topk), -1, device="cuda")
topk_indices[:, :1600] = torch.arange(1600).unsqueeze(0).repeat(bs, 1)
ref_result = transform_index_page_table_decode_ref(page_table, topk_indices)
result = transform_index_page_table_decode_fast(page_table, topk_indices)
assert torch.all(result == ref_result)
print("Passed")
python/sglang/srt/layers/attention/dsa/triton_kernel.py
+196
View File
@@ -0,0 +1,196 @@
from typing import Optional, Tuple
import torch
import triton
import triton.language as tl
# Triton implementation
@triton.jit
def _act_quant_kernel(
X_ptr,
Y_ptr,
S_ptr,
M,
N,
group_size: tl.constexpr,
round_scale: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
"""
Triton kernel for activation quantization.
Each block processes BLOCK_M rows and group_size columns.
"""
# Get block IDs
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
# FP8 constants
fp8_min = -448.0
fp8_max = 448.0
fp8_max_inv = 1.0 / fp8_max
# Calculate row and column offsets
row_start = pid_m * BLOCK_M
col_start = pid_n * group_size
# Create offset arrays
rows = row_start + tl.arange(0, BLOCK_M)
cols = col_start + tl.arange(0, BLOCK_N)
# Mask for valid rows and columns
row_mask = rows < M
col_mask = cols < N
mask = row_mask[:, None] & col_mask[None, :]
# Load input data
x_ptrs = X_ptr + rows[:, None] * N + cols[None, :]
x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
# Compute absolute max along columns (group_size dimension) for each row
x_abs = tl.abs(x)
amax = tl.max(x_abs, axis=1) # Shape: (BLOCK_M,)
# Clamp amax to avoid division by zero
amax = tl.maximum(amax, 1e-4)
# Compute scale
if round_scale:
# Fast round scale using bit manipulation approximation
# This is a simplified version - the exact bit manipulation is harder in Triton
# Using log2 + ceil + pow2 as approximation
log_val = tl.log2(amax * fp8_max_inv)
log_ceil = tl.ceil(log_val)
scale = tl.exp2(log_ceil)
else:
scale = amax * fp8_max_inv
# Quantize: y = clamp(x / scale, fp8_min, fp8_max)
scale_broadcast = scale[:, None]
y = x / scale_broadcast
y = tl.minimum(tl.maximum(y, fp8_min), fp8_max)
# Store quantized output
y_ptrs = Y_ptr + rows[:, None] * N + cols[None, :]
tl.store(y_ptrs, y, mask=mask)
# Store scales
s_cols = pid_n
s_ptrs = S_ptr + rows * (N // group_size) + s_cols
s_mask = row_mask
tl.store(s_ptrs, scale, mask=s_mask)
def act_quant(
x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantizes the input tensor `x` using block-wise quantization with Triton.
Args:
x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
scale_fmt (Optional[str], optional): The format of the scale. Default is None.
Returns:
Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
- The quantized tensor with dtype `torch.float8_e4m3fn`.
- A tensor of scaling factors with dtype `torch.float32`.
"""
assert x.is_contiguous(), "Input tensor must be contiguous"
assert (
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
# Flatten all dims except last
N = x.size(-1)
x_flat = x.view(-1, N)
M = x_flat.size(0)
# Allocate output tensors
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
y_flat = y.view(-1, N)
s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
s_flat = s.view(-1, N // block_size)
# Launch kernel
BLOCK_M = 32
BLOCK_N = block_size
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, block_size))
round_scale = scale_fmt is not None
_act_quant_kernel[grid](
x_flat,
y_flat,
s_flat,
M,
N,
group_size=block_size,
round_scale=round_scale,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
num_stages=0 if round_scale else 2,
)
return y, s
@triton.jit
def _get_valid_kv_indices_kernel(
page_table_ptr, # [bs, topk]
kv_indptr_ptr, # [bs + 1]
kv_indices_ptr, # [bs * topk] output buffer
bs: tl.constexpr,
topk: tl.constexpr,
):
"""
Extract valid indices (non -1) from page_table into kv_indices.
Each program handles one batch.
"""
batch_id = tl.program_id(0)
# Get the start position for this batch in kv_indices
dst_start = tl.load(kv_indptr_ptr + batch_id)
# Load all topk indices for this batch
src_offset = batch_id * topk
offsets = tl.arange(0, topk)
indices = tl.load(page_table_ptr + src_offset + offsets)
# Count valid indices and compact them
mask = indices != -1
# Use prefix sum to compute destination positions for valid elements
# For each position, count how many valid elements are before it
prefix_sum = tl.cumsum(mask.to(tl.int32), axis=0) - 1
# Store valid indices to their compacted positions
dst_positions = dst_start + prefix_sum
tl.store(kv_indices_ptr + dst_positions, indices, mask=mask)
def get_valid_kv_indices(
page_table_1: torch.Tensor,
kv_indptr: torch.Tensor,
kv_indices: torch.Tensor,
bs: int,
):
"""
Extract valid indices from page_table_1 into kv_indices buffer.
Args:
page_table_1: [bs, topk] page table with -1 as invalid
kv_indptr: [bs + 1] cumulative count of valid indices per batch
kv_indices: [bs * topk] pre-allocated output buffer
bs: batch size
"""
topk = page_table_1.shape[1]
grid = (bs,)
_get_valid_kv_indices_kernel[grid](
page_table_1,
kv_indptr,
kv_indices,
bs,
topk,
)
python/sglang/srt/layers/attention/dsa/utils.py
+271
View File
@@ -0,0 +1,271 @@
from functools import lru_cache
from typing import TYPE_CHECKING, List, Tuple, Union
import torch
import triton
import triton.language as tl
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.common import ceil_align, ceil_div
@lru_cache(maxsize=1)
def aiter_can_use_preshuffle_paged_mqa() -> bool:
"""Whether aiter's preshuffle paged MQA / cache kernels can be used on this runtime.
aiter's ``deepgemm_fp8_paged_mqa_logits`` only supports ``KVBlockSize > 1`` and
``Preshuffle=True`` on its gluon kernel path. The gluon path is enabled when
Triton >= 3.5.0, OR when ``AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1`` is set
(which additionally requires that the AOT gluon kernel artifacts ship inside
the aiter wheel/image). Otherwise aiter asserts ``KVBlockSize == 1`` and
refuses ``Preshuffle=True``.
sglang's DSA indexer uses this single decision to pick:
* ``page_size``: 64 (preshuffle) vs 1 (legacy) on ROCm
* ``Preshuffle`` / ``preshuffle`` flags on the aiter MQA + cache kernels
* ``get_page_table_64`` vs ``get_page_table_1`` on the metadata
* whether ``GetKAndS.execute`` uses the aiter or the triton implementation
The result is cached so the cost is paid once per process.
Set ``SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1`` to force the legacy path even when
the gluon kernel would otherwise be available (useful for CI bisection).
``SGLANG_NSA_HIP_DISABLE_PRESHUFFLE`` is a deprecated alias.
"""
if not is_hip():
return False
if not get_bool_env_var("SGLANG_USE_AITER"):
return False
if envs.SGLANG_DSA_HIP_DISABLE_PRESHUFFLE.get():
return False
if get_bool_env_var("AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS"):
return True
try:
from packaging.version import Version
return Version(Version(triton.__version__).base_version) >= Version("3.5.0")
except Exception:
return False
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
return original_seq_lens.clamp(max=dsa_index_topk)
def is_dsa_enable_prefill_cp():
return get_global_server_args().enable_dsa_prefill_context_parallel
def is_dsa_prefill_cp_in_seq_split():
return (
is_dsa_enable_prefill_cp()
and get_global_server_args().dsa_prefill_cp_mode == "in-seq-split"
)
def is_dsa_prefill_cp_round_robin_split():
return (
is_dsa_enable_prefill_cp()
and get_global_server_args().dsa_prefill_cp_mode == "round-robin-split"
)
def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_attention_cp_size()
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return (
is_dsa_prefill_cp_round_robin_split()
and seq_len > 0
and seq_len >= cp_size
and cp_size > 1
)
def dsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
"""
# for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
| +-----------before split------------+|
| token0, token1, token2, token3, token4, token5, token6, token7, ...
|
| +--------------result-------------------+
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
if isinstance(input_, (tuple, list)):
indices = range(cp_rank, len(input_), cp_size)
return input_[indices]
tokens = len(input_)
if tokens % cp_size != 0:
cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
if cur_len == 0:
return input_.new_empty(0, *input_.shape[1:])
indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
return input_[indices]
# for torch device tensor
return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
forward_batch.is_extend_in_batch, global_num_tokens
)
if dp_padding_mode.is_max_len():
tokens = max(global_num_tokens)
elif len(global_num_tokens) > 1:
tokens = global_num_tokens[get_attention_dp_rank()]
else:
tokens = global_num_tokens[0]
if can_dsa_prefill_cp_round_robin_split(forward_batch):
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
attn_cp_size = get_attention_cp_size()
needs_cp_pad = attn_cp_size > 1 and can_dsa_prefill_cp_round_robin_split(
forward_batch
)
needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
if not needs_cp_pad and not needs_dp_pad:
return dsa_cache_seqlens
tokens = cal_padded_tokens(forward_batch)
pad_len = tokens - dsa_cache_seqlens.shape[0]
if pad_len > 0:
dsa_cache_seqlens = torch.cat(
[
dsa_cache_seqlens,
dsa_cache_seqlens.new_zeros(pad_len, *dsa_cache_seqlens.shape[1:]),
]
)
return dsa_cache_seqlens
def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch):
if is_dsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
seq_len % cp_size == 0
), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when dsa_prefill_cp_mode is round-robin-split"
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
# the seq data needs to be divided and recombined at twice the size of cp_size.
cur_cp_seq_len = seq_len // (cp_size * 2)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_dsa
and forward_batch.forward_mode.is_context_parallel_extend()
and is_dsa_enable_prefill_cp()
and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
):
return True
else:
return False
@triton.jit
def dsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
out_seqs_ptr,
bs_idx_ptr,
tokens: tl.constexpr,
cp_size: tl.constexpr,
cp_rank: tl.constexpr,
):
extra_seq = 0
bs_idx = 0
for bs in range(tokens):
cur_len = tl.load(in_seqs_ptr + bs)
cur_len += extra_seq
cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank)
if cur_seq > 0:
tl.store(bs_idx_ptr + bs_idx, bs)
tl.store(out_seqs_ptr + bs_idx, cur_seq)
bs_idx += 1
extra_seq = cur_len - cur_seq * cp_size
def dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_seqs.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
q_seqs = [q_len for q_len in q_seqs if q_len > 0]
return q_seqs, bs_idx
def dsa_cp_round_robin_split_q_seqs(
extend_seqs_cpu, extend_seqs
) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
"""
round-robin-split distributes tokens across ranks based on token_idx % cp_size.
Return:
ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
for each sequence after distribution across cp ranks.
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_q_lens = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
)
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
grid = (1,)
dsa_cp_round_robin_split_q_seqs_kernel[grid](
extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
def dsa_use_prefill_cp(forward_batch, dsa_enable_prefill_cp=None):
if dsa_enable_prefill_cp is None:
dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if (
forward_batch.attn_cp_metadata is not None
and dsa_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
):
return True
else:
return False
python/sglang/srt/layers/attention/dsa_backend.py
+2528
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/dsv4/compress_hip.py
+1 -1
View File
@@ -10,8 +10,8 @@ import triton
import triton.language as tl
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation
from sglang.srt.layers.attention.dsv4.compressor import Compressor as _CompressorBase
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
from sglang.srt.layers.deepseek_v4_rope import (
apply_rotary_emb_triton,
fused_norm_rope_inplace_triton,
python/sglang/srt/layers/attention/dsv4/compressor.py
+4 -4
View File
@@ -15,11 +15,11 @@ from sglang.jit_kernel.deepseek_v4 import (
)
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.attention.dsv4.quant_k_cache import (
quant_to_nope_fp8_rope_bf16_pack_triton,
)
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
from sglang.srt.layers.dp_attention import get_attention_cp_size
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
@@ -71,7 +71,7 @@ class CompressorBackendMixin:
compress_ratio: int,
is_paged: bool = False,
) -> torch.Tensor:
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation
assert compress_ratio in (
4,
@@ -358,7 +358,7 @@ class Compressor(nn.Module):
kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
# CUDA path: delegate to backend
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
kv_score = cp_all_gather_rerange_output(
kv_score,
get_attention_cp_size(),
python/sglang/srt/layers/attention/dsv4/indexer.py
+1 -1
View File
@@ -368,7 +368,7 @@ class C4IndexerBackendMixin:
assert len(weights.shape) == 3
weights = weights.squeeze(2)
if envs.SGLANG_OPT_USE_TILELANG_INDEXER.get():
from sglang.srt.layers.attention.nsa.tilelang_kernel import (
from sglang.srt.layers.attention.dsa.tilelang_kernel import (
tilelang_fp8_paged_mqa_logits as fn,
)
elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get():
python/sglang/srt/layers/attention/hip_flash_mla.py
+1 -1
View File
@@ -12,7 +12,7 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs):
if is_hip():
import os
from sglang.srt.layers.attention.nsa.tilelang_kernel import (
from sglang.srt.layers.attention.dsa.tilelang_kernel import (
dpsk_v4_fp8_attention_fwd,
)
python/sglang/srt/layers/attention/hybrid_attn_backend.py
+1 -1
View File
@@ -3,7 +3,7 @@ from typing import Optional
import torch
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner
python/sglang/srt/layers/attention/nsa/__init__.py
+11
View File
@@ -0,0 +1,11 @@
# [Deprecated] attention/nsa/ is a thin re-export shim for backward compatibility.
# Use attention/dsa/ instead. This directory will be removed in a future release.
import warnings
warnings.warn(
"sglang.srt.layers.attention.nsa is deprecated; "
"use sglang.srt.layers.attention.dsa instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/dequant_k_cache.py
+9 -288
View File
@@ -1,289 +1,10 @@
import torch
import triton
import triton.language as tl
# [Deprecated] Re-export shim for backward compatibility. Use dsa.dequant_k_cache instead.
import warnings
def dequantize_k_cache(quant_k_cache):
return _dequantize_k_cache_fast_wrapped(quant_k_cache)
def _dequantize_k_cache_ref(
quant_k_cache: torch.Tensor, # (num_blocks, block_size, 1, bytes_per_token)
dv: int = 512,
tile_size: int = 128,
d: int = 576,
) -> torch.Tensor:
"""
De-quantize the k-cache
"""
assert dv % tile_size == 0
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_tiles = dv // tile_size
num_blocks, block_size, h_k, _ = quant_k_cache.shape
assert h_k == 1
result = torch.empty(
(num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
)
quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
input_nope = quant_k_cache[..., :dv]
input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
result[..., dv:] = input_rope
for tile_idx in range(0, num_tiles):
cur_nope = input_nope[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].to(torch.float32)
cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_nope * cur_scales
)
if original_ndim == 3:
return result.view(num_blocks, 1, -1)
else:
return result.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast_wrapped(
quant_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_blocks, block_size, _, dim_quant = quant_k_cache.shape
assert dv == 512
assert dim_quant == 656
assert tile_size == 128
quant_k_cache = quant_k_cache.view((-1, dim_quant))
output = _dequantize_k_cache_fast(quant_k_cache)
if original_ndim == 3:
return output.view(num_blocks, 1, -1)
else:
return output.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
num_tokens, dim_quant = quant_k_cache.shape
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size
assert dim_quant == 656
output = torch.empty(
(num_tokens, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_fast_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q
ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
def dequantize_k_cache_paged(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
group_size: int = 128,
) -> torch.Tensor:
"""
De-quantize the k-cache with paged layout
Args:
quant_k_cache: [total_num_tokens, 1, dim_quant] or [num_blocks, block_size, 1, dim_quant], the quantized k-cache in paged layout
page_table_1_flattened: [num_tokens], the flattened page_table_1 with the page indices in each requests concatenated together
Returns:
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
"""
dim_quant = quant_k_cache.shape[-1]
assert (
dim_quant == 656
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
quant_k_cache = quant_k_cache.view((-1, dim_quant))
# num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
# Index bounds validated in nsa_backend.init_forward_metadata
num_tokens = page_table_1_flattened.shape[0]
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size # 512 // 128 = 4
output = torch.empty(
(num_tokens, 1, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
# cdiv(512 + 64, 128) = 5
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
# [:, 512:512+4*4] = [:, 512:528]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
# [:, 528:]
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_paged_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
page_table_1_flattened,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_paged_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
page_table_1_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q
ptr_s = (
input_nope_s_ptr
+ token_id_paged * input_nope_s_stride_0
+ effective_block_id
)
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
raise Exception("UT is in quant_k_cache.py")
warnings.warn(
"sglang.srt.layers.attention.nsa.dequant_k_cache is deprecated; "
"use sglang.srt.layers.attention.dsa.dequant_k_cache instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.dequant_k_cache import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/index_buf_accessor.py
+10 -814
View File
@@ -1,814 +1,10 @@
from typing import TYPE_CHECKING
import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.nsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import get_bool_env_var, is_hip
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# aiter cp_gather kernel with preshuffle=True is only valid when the indexer
# uses the page_size=64 preshuffle layout (i.e. when the matching MQA gluon path
# is also enabled).
_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
if _use_aiter_preshuffle:
from aiter.ops.cache import cp_gather_indexer_k_quant_cache
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
"""
k: data, 128 item per token, fp8
s: scale, 1 item per token, fp32
"""
class GetK:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
index_k_fp8 = torch.empty(
(seq_len_, pool.index_head_dim),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][: pool.page_size * pool.index_head_dim].view(-1, pool.index_head_dim)
return index_k_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim), uint8
"""
# can handle per 128B instead of per element
# page_indices: (num_pages,), element := a page index
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_page = pool.page_size * pool.index_head_dim
num_k_bytes_per_token = pool.index_head_dim
# buf: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4), uint8
# flat_buf: (whatever,), uint8
flat_buf = buf.flatten()
# flat_indices: (num_pages, num_k_bytes_per_page), int32, element := an index into flat_buf that we want to access
flat_indices = (page_indices * buf_numel_per_page)[:, None] + torch.arange(
num_k_bytes_per_page, dtype=torch.int32, device="cuda"
)[None, :]
flat_indices = flat_indices.flatten()[: seq_len * num_k_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 128)
@classmethod
def triton(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering K data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, index_head_dim), uint8
"""
return _get_k_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetS:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
assert pool.index_head_dim // pool.quant_block_size == 1
index_k_scale_fp8 = torch.empty(
(seq_len_, 4),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_scale_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][pool.page_size * pool.index_head_dim :].view(-1, 4)
return index_k_scale_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim // quant_block_size), uint8
"""
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_page = buf.shape[1] - pool.page_size * pool.index_head_dim
num_s_bytes_per_token = pool.index_head_dim // pool.quant_block_size * 4
s_offset_in_page = pool.page_size * pool.index_head_dim
flat_buf = buf.flatten()
flat_indices = (
(page_indices * buf_numel_per_page)[:, None]
+ torch.arange(num_s_bytes_per_page, dtype=torch.int32, device="cuda")[
None, :
]
+ s_offset_in_page
)
flat_indices = flat_indices.flatten()[: seq_len * num_s_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 4)
@classmethod
def triton(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering S (scale) data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, 4), uint8
"""
return _get_s_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetKAndS:
@classmethod
def execute(cls, *args, **kwargs):
# The aiter path uses cp_gather_indexer_k_quant_cache(preshuffle=True),
# which only matches the layout produced when the rest of the indexer
# is on the page_size=64 preshuffle path. Otherwise fall back to the
# triton implementation (which works on the page_size=1 legacy layout).
if _use_aiter_preshuffle:
return cls.aiter(*args, **kwargs)
return cls.triton(*args, **kwargs)
@classmethod
def aiter(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
page_size = pool.page_size
index_head_dim = pool.index_head_dim
quant_block_size = pool.quant_block_size
scale_elems = index_head_dim // quant_block_size
kv_cache = buf.view(-1, page_size, index_head_dim + scale_elems * 4).view(
fp8_dtype
)
dst_k = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
dst_scale = torch.empty(
(seq_len_sum, scale_elems * 4), dtype=torch.uint8, device=buf.device
)
cu_seq_lens = torch.zeros(
seq_len_tensor.shape[0] + 1, dtype=torch.int32, device=buf.device
)
torch.cumsum(seq_len_tensor.to(torch.int32), dim=0, out=cu_seq_lens[1:])
cp_gather_indexer_k_quant_cache(
kv_cache,
dst_k.view(fp8_dtype),
dst_scale,
page_indices.to(torch.int32),
cu_seq_lens,
preshuffle=True,
)
return dst_k, dst_scale
@classmethod
def triton(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
"""
Triton implementation for gathering both K and S data from paged buffer in a single call.
:param page_indices: (num_pages,), int32/int64
:param seq_len_tensor: (num_pages,), int32/int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of all sequence len, int32
:return: tuple of (k_fp8, k_scale) where
k_fp8: (seq_len, index_head_dim), uint8
k_scale: (seq_len, 4), uint8
"""
return _get_k_and_s_triton(
buf=buf,
page_indices=page_indices,
seq_lens=seq_len_tensor,
seq_len_sum=seq_len_sum,
max_seq_len=max_seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class SetK:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
buf[
page_index,
offset * pool.index_head_dim : (offset + 1) * pool.index_head_dim,
] = index_k[i].view(torch.uint8)
@classmethod
def torch_fast(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_token = pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ (loc_token_offset_in_page * num_k_bytes_per_token)[:, None]
+ torch.arange(num_k_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
num_k_bytes_total = num_tokens_to_write * num_k_bytes_per_token
flat_indices = flat_indices.flatten()[:num_k_bytes_total]
flat_buf[flat_indices] = index_k.view(torch.uint8).flatten()
class SetS:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
start = pool.page_size * pool.index_head_dim
buf[page_index, start + offset * 4 : start + (offset + 1) * 4] = (
index_k_scale[i].view(torch.uint8)
)
@classmethod
def torch_fast(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_token = 4
s_offset_in_page = pool.page_size * pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ s_offset_in_page
+ (loc_token_offset_in_page * num_s_bytes_per_token)[:, None]
+ torch.arange(num_s_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
number_s_bytes_total = num_tokens_to_write * num_s_bytes_per_token
flat_indices = flat_indices.flatten()[:number_s_bytes_total]
flat_buf[flat_indices] = index_k_scale.view(torch.uint8).flatten()
class SetKAndS:
@classmethod
def execute(cls, *args, buf, **kwargs):
if 0:
# print("SetK, SetS comparison test")
buf_cloned = buf.clone()
cls.vanilla(*args, **kwargs, buf=buf)
cls.triton(*args, **kwargs, buf=buf_cloned)
def _clear_token_0(target):
target[0, :128] = target[0, 64 * 128 : 64 * 128 + 4] = 0
_clear_token_0(buf)
_clear_token_0(buf_cloned)
assert torch.all(
buf == buf_cloned
), f"{buf=} {buf_cloned=} {kwargs['loc'].to_list()=}"
return
cls.triton(*args, **kwargs, buf=buf)
@classmethod
def vanilla(cls, pool, buf, loc, index_k, index_k_scale):
SetK.execute(pool=pool, buf=buf, loc=loc, index_k=index_k)
SetS.execute(pool=pool, buf=buf, loc=loc, index_k_scale=index_k_scale)
@classmethod
def triton(cls, pool, buf, loc, index_k, index_k_scale):
loc = loc.to(torch.int64)
_set_k_and_s_triton(
buf=buf,
loc=loc,
index_k=index_k,
index_k_scale=index_k_scale,
page_size=pool.page_size,
)
def _set_k_and_s_triton(
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
index_k_scale: torch.Tensor,
page_size: int,
):
"""
:param buf: (num_pages, page_size 64 * (128B data + 4B scale)), uint8
:param loc: (num_tokens_to_write,), int, element := the token index to write to
:param index_k: (num_tokens_to_write, 128 elem), fp8
:param index_k_scale: (num_tokens_to_write, 1 elem), fp32
:return:
"""
num_pages, buf_numel_per_page = buf.shape
(num_tokens_to_write,) = loc.shape
num_tokens_to_write_, index_head_dim = index_k.shape
# Handle both 1D (num_tokens,) and 2D (num_tokens, 1) shapes for index_k_scale
if index_k_scale.ndim == 1:
num_tokens_to_write__ = index_k_scale.shape[0]
scale_dim = 1
elif index_k_scale.ndim == 2:
num_tokens_to_write__, scale_dim = index_k_scale.shape
else:
raise ValueError(
f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}"
)
assert buf_numel_per_page == page_size * (128 + 4)
assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__
assert index_head_dim == 128
assert scale_dim == 1
if _is_hip:
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert page_size == 64
assert buf.dtype == torch.uint8
assert loc.dtype == torch.int64, f"{loc.dtype=}" # can be int32
if _is_fp8_fnuz:
assert index_k.dtype == torch.float8_e4m3fnuz
else:
assert index_k.dtype == torch.float8_e4m3fn
assert index_k_scale.dtype == torch.float32
assert buf.is_contiguous()
assert loc.is_contiguous()
assert index_k.is_contiguous()
assert index_k_scale.is_contiguous()
if _is_fp8_fnuz:
buf_fp8 = buf.view(torch.float8_e4m3fnuz)
else:
buf_fp8 = buf.view(torch.float8_e4m3fn)
buf_fp32 = buf.view(torch.float32)
_set_k_and_s_triton_kernel[(num_tokens_to_write,)](
buf_fp8,
buf_fp32,
loc,
index_k,
index_k_scale,
index_k.stride(0),
PAGE_SIZE=page_size,
BUF_NUMEL_PER_PAGE=buf_numel_per_page,
NUM_K_ELEMS_PER_TOKEN=index_head_dim,
S_OFFSET_NBYTES_IN_PAGE=page_size * index_head_dim,
)
@triton.jit
def _set_k_and_s_triton_kernel(
buf_fp8_ptr,
buf_fp32_ptr,
loc_ptr,
index_k_ptr,
index_k_scale_ptr,
index_k_ptr_stride_0,
PAGE_SIZE: tl.constexpr,
BUF_NUMEL_PER_PAGE: tl.constexpr,
NUM_K_ELEMS_PER_TOKEN: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
):
token_id = tl.program_id(0)
loc = tl.load(loc_ptr + token_id)
in_k_offsets = token_id * index_k_ptr_stride_0 + tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
# no need for `mask`, since we read 128B for k and 4B for scale, both pow of 2
k = tl.load(index_k_ptr + in_k_offsets)
k_scale = tl.load(index_k_scale_ptr + token_id)
loc_page_index = loc // PAGE_SIZE
loc_token_offset_in_page = loc % PAGE_SIZE
out_k_offsets = (
loc_page_index * BUF_NUMEL_PER_PAGE
+ loc_token_offset_in_page * NUM_K_ELEMS_PER_TOKEN
+ tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
)
# "//4" b/c it is fp32 instead of uint8
out_s_offset = (
loc_page_index * BUF_NUMEL_PER_PAGE // 4
+ S_OFFSET_NBYTES_IN_PAGE // 4
+ loc_token_offset_in_page
)
tl.store(buf_fp8_ptr + out_k_offsets, k)
tl.store(buf_fp32_ptr + out_s_offset, k_scale)
def _get_k_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather K (key) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, index_head_dim), uint8
"""
num_pages, buf_numel_per_page = buf.shape
# Allocate output
out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_k_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
index_head_dim,
BLOCK_SIZE=128,
)
return out
@triton.jit
def _get_k_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 128 bytes from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# buf[page_index, token_offset_in_page * index_head_dim : ...]
src_base_offset = (
page_index * buf_numel_per_page + token_offset_in_page * index_head_dim
)
# Load 128 bytes (index_head_dim elements)
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < index_head_dim
data = tl.load(buf_ptr + src_base_offset + offsets, mask=mask)
# Store to output
dst_offset = token_id * index_head_dim
tl.store(out_ptr + dst_offset + offsets, data, mask=mask)
def _get_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather S (scale) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, 4), uint8 (representing fp32 scale)
"""
num_pages, buf_numel_per_page = buf.shape
s_offset_in_page = page_size * index_head_dim # Scales start after K data
# Allocate output
out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_s_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
s_offset_in_page,
)
return out
@triton.jit
def _get_s_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
s_offset_in_page: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 4 bytes (fp32 scale) from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# Scales are stored after K data: page_size * index_head_dim offset
# buf[page_index, s_offset_in_page + token_offset_in_page * 4 : ...]
src_base_offset = (
page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4
)
# Load 4 bytes (fp32 scale)
offsets = tl.arange(0, 4)
data = tl.load(buf_ptr + src_base_offset + offsets)
# Store to output
dst_offset = token_id * 4
tl.store(out_ptr + dst_offset + offsets, data)
def _get_k_and_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Fused gather of both K (key) and S (scale) data from paged buffer using Triton.
This is more efficient than calling GetK and GetS separately.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_lens: tensor of sequence lens, int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of sequence len, int32
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: tuple of (k_out, s_out) where
k_out: (seq_len, index_head_dim), uint8
s_out: (seq_len, 4), uint8
"""
# Allocate outputs
k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device=buf.device)
_, buf_numel_per_page = buf.shape
_, page_indice_batch_offset = page_indices.shape
s_offset_in_page = page_size * index_head_dim
# Launch kernel with one thread per token
BLOCK_SIZE = 256
BLOCK_SIZE_K = 128
num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
seq_num = seq_lens.shape[0]
grid = (seq_num, num_token_blocks, num_k_threads)
seq_num_pow2 = 1
while seq_num_pow2 < seq_num:
seq_num_pow2 *= 2
_get_k_and_s_triton_kernel[grid](
buf_ptr=buf,
page_indices_ptr=page_indices,
k_out_ptr=k_out,
s_out_ptr=s_out,
seq_len_ptr=seq_lens,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
return k_out, s_out
@triton.jit
def _get_k_and_s_triton_kernel(
buf_ptr,
page_indices_ptr,
k_out_ptr,
s_out_ptr,
seq_len_ptr,
seq_len_num_pow: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
s_offset_in_page: tl.constexpr,
page_indice_batch_offset,
BLOCK_SIZE: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
"""
Fused kernel that gathers both K and S data in a single pass.
Each program handles one token (seq_len tokens total).
Loads 128 bytes (K) + 4 bytes (S) from the appropriate page.
"""
batch_id = tl.program_id(0)
block_token_start = tl.program_id(1) * BLOCK_SIZE
thread_idx = tl.program_id(2)
# Define the token range within the block and the K dimension range handled by the thread.
token_ids_in_block = tl.arange(0, BLOCK_SIZE)
token_ids = block_token_start + token_ids_in_block
k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
seq_len = tl.load(seq_len_ptr + batch_id)
token_valid_mask = token_ids < seq_len
pre_batch_idx = tl.arange(0, seq_len_num_pow)
mask_pre_batch_idx = pre_batch_idx < batch_id
prev_seq_lens = tl.load(seq_len_ptr + pre_batch_idx, mask=mask_pre_batch_idx)
batch_token_offset = tl.sum(prev_seq_lens)
# Batch calculate the page index and in-page offset of each token.
page_idx = token_ids // page_size
token_offset_in_page = token_ids % page_size
page_indices_base = batch_id * page_indice_batch_offset
page_idx_valid_mask = page_idx < page_indice_batch_offset
page_index = tl.load(
page_indices_ptr + page_idx + page_indices_base,
mask=token_valid_mask & page_idx_valid_mask,
)
# ===== Load K data =====
# The address calculation logic for K: page_index * total number of elements in a single page + K offset of the token within the page.
k_src_token_offset = token_offset_in_page * index_head_dim
k_src_base_offset = page_index * buf_numel_per_page + k_src_token_offset
k_load_addr = buf_ptr + k_src_base_offset[:, None] + k_offsets[None, :]
k_dim_mask = k_offsets[None, :] < index_head_dim
k_mask = token_valid_mask[:, None] & k_dim_mask
k_data = tl.load(k_load_addr, mask=k_mask, other=0)
# Store K to output
k_dst_token_offset = batch_token_offset + token_ids
k_dst_base_offset = k_dst_token_offset * index_head_dim
k_store_addr = k_out_ptr + k_dst_base_offset[:, None] + k_offsets[None, :]
tl.store(k_store_addr, k_data, mask=k_mask)
# ===== Load S data =====
# The address calculation logic for S: page_index * total number of elements in a single page + starting offset of S within the page + offset of token within S in the page
s_src_token_offset = s_offset_in_page + token_offset_in_page * 4
s_src_base_offset = page_index * buf_numel_per_page + s_src_token_offset
s_offsets = tl.arange(0, 4)
s_load_addr = buf_ptr + s_src_base_offset[:, None] + s_offsets[None, :]
s_mask = token_valid_mask[:, None] & (s_offsets[None, :] < 4)
s_data = tl.load(s_load_addr, mask=s_mask, other=0)
# Store S to output
s_dst_token_offset = batch_token_offset + token_ids
s_dst_base_offset = s_dst_token_offset * 4
s_store_addr = s_out_ptr + s_dst_base_offset[:, None] + s_offsets[None, :]
tl.store(s_store_addr, s_data, mask=s_mask)
# [Deprecated] Re-export shim for backward compatibility. Use dsa.index_buf_accessor instead.
import warnings
warnings.warn(
"sglang.srt.layers.attention.nsa.index_buf_accessor is deprecated; "
"use sglang.srt.layers.attention.dsa.index_buf_accessor instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.index_buf_accessor import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py
+9 -324
View File
@@ -1,325 +1,10 @@
"""Multi-step precompute utilities for Native Sparse Attention backend.
# [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_backend_mtp_precompute instead.
import warnings
This module provides optimization utilities for multi-step speculative decoding
by precomputing shared metadata once and copying it to multiple backend instances.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.attention.nsa.utils import compute_nsa_seqlens
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.speculative.spec_info import SpecInput
@dataclass
class PrecomputedMetadata:
"""Precomputed metadata shared across multiple backend instances.
Used for multi-step speculative decoding where multiple backends
need identical metadata. Precomputing once and copying N times
is much faster than computing N times.
"""
# Basic seqlens
cache_seqlens: torch.Tensor # int32, [bs]
cu_seqlens_k: torch.Tensor # int32, [bs+1]
# Page table
page_indices: torch.Tensor # int32, [bs, max_len] or [expanded_bs, max_len]
real_page_table: Optional[torch.Tensor] # int32, transformed version
# NSA seqlens
seqlens_expanded: torch.Tensor # int32, [expanded_size]
nsa_cache_seqlens: torch.Tensor # int32, [expanded_size]
nsa_cu_seqlens_k: torch.Tensor # int32, [expanded_size+1]
seqlens_expanded_size: int
# Dimensions
max_len: int # for decode/draft_extend
max_seqlen_k: int # for target_verify
# FlashMLA (optional)
flashmla_metadata: Optional[torch.Tensor] = None
def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
"""Compute cumulative sequence lengths with padding."""
assert seqlens.dtype == torch.int32
return torch.nn.functional.pad(
torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
)
class NativeSparseAttnBackendMTPPrecomputeMixin:
"""Mixin class providing metadata precomputation for multi-step speculative decoding.
This mixin provides the _precompute_replay_metadata method and its helpers,
which are used to optimize CUDA graph replay in multi-step scenarios.
"""
def _precompute_replay_metadata(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
forward_mode: "ForwardMode",
spec_info: Optional["SpecInput"],
) -> PrecomputedMetadata:
"""Precompute all shared metadata for multi-step backends.
This function extracts and computes all operations that are
identical across different backend instances in multi-step
speculative decoding.
Args:
bs: Batch size
req_pool_indices: Request pool indices [bs]
seq_lens: Sequence lengths [bs]
seq_lens_cpu: Sequence lengths on CPU [bs]
forward_mode: Forward mode (decode/target_verify/draft_extend)
spec_info: Speculative decoding info (for draft_extend mode)
Returns:
PrecomputedMetadata containing all shared intermediate results
"""
# Slice inputs to batch size
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
# Dispatch to mode-specific precomputation
if forward_mode.is_decode_or_idle():
return self._precompute_decode_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_target_verify():
return self._precompute_target_verify_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_draft_extend():
return self._precompute_draft_extend_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu, spec_info
)
else:
raise ValueError(f"Unsupported forward mode: {forward_mode}")
def _precompute_decode_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for normal decode mode."""
max_len = int(seq_lens_cpu.max().item())
# Convert to int32 and compute cumsum
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Get page indices from cache
page_indices = self.req_to_token[req_pool_indices, :max_len].contiguous()
# Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens(
cache_seqlens, nsa_index_topk=self.nsa_index_topk
)
seqlens_expanded = cache_seqlens
seqlens_expanded_size = seqlens_expanded.shape[0]
# Compute NSA cumsum
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens)
# Transform page table if needed
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None # Will use page_indices directly
# Compute FlashMLA metadata if needed
flashmla_metadata = None
if self.nsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
nsa_cache_seqlens=nsa_cache_seqlens,
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_len,
max_seqlen_k=max_len,
flashmla_metadata=flashmla_metadata,
)
def _precompute_target_verify_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for target verify mode."""
max_seqlen_k = int(
seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
)
# Cache seqlens with draft tokens
cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Page indices (repeated for each draft token)
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
seqlens_int32_cpu = [
self.speculative_num_draft_tokens + kv_len
for kv_len in seq_lens_cpu.tolist()
]
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seqlens_int32_cpu,
strict=True,
)
]
)
# Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens(seqlens_expanded, self.nsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# NSA cumsum
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.nsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
nsa_cache_seqlens=nsa_cache_seqlens,
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=-1, # Not used in this mode
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
def _precompute_draft_extend_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
spec_info: "SpecInput",
) -> PrecomputedMetadata:
"""Precompute metadata for draft extend mode."""
max_seqlen_k = int(seq_lens_cpu.max().item())
# Cache seqlens
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Extend seqlens from spec_info: num_accept_tokens already includes
# the bonus token (drafts + 1).
extend_seq_lens = spec_info.num_accept_tokens[:bs]
extend_seq_lens_cpu = extend_seq_lens.tolist()
# Page indices (repeated per accept length)
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(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seq_lens_cpu.tolist(),
strict=True,
)
]
)
# Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens(seqlens_expanded, self.nsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# NSA cumsum
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.nsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
nsa_cache_seqlens=nsa_cache_seqlens,
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_seqlen_k,
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
warnings.warn(
"sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute is deprecated; "
"use sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+8 -1744
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py
+9 -406
View File
@@ -1,407 +1,10 @@
"""
Verification utilities for NSA backend fused metadata copy operations.
# [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_mtp_verification instead.
import warnings
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"],
)
warnings.warn(
"sglang.srt.layers.attention.nsa.nsa_mtp_verification is deprecated; "
"use sglang.srt.layers.attention.dsa.dsa_mtp_verification instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.dsa_mtp_verification import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/quant_k_cache.py
+10 -449
View File
@@ -1,449 +1,10 @@
import torch
import triton
import triton.language as tl
def quantize_k_cache(cache_k):
return _quantize_k_cache_fast_wrapped(cache_k)
def quantize_k_cache_separate(
k_nope: torch.Tensor,
k_rope: torch.Tensor,
tile_size: int = 128,
):
"""
Quantize k_nope and k_rope separately without concat, returns two tensors.
This avoids the concat operation and enables direct reuse of set_mla_kv_buffer_triton
by returning two separate byte tensors for the nope and rope parts.
Args:
k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
Must have dim_nope=512 for FP8 MLA quantization
k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
Must have dim_rope=64 for FP8 MLA quantization
tile_size: quantization tile size (default 128)
Returns:
Tuple of (nope_part, rope_part) where:
- nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
- rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
"""
# Squeeze middle dimension if present
k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
num_tokens = k_nope_2d.shape[0]
dim_nope = k_nope_2d.shape[1]
dim_rope = k_rope_2d.shape[1]
# Validate dimensions for FP8 MLA
if dim_nope != 512:
raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
if dim_rope != 64:
raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
if k_rope_2d.shape[0] != num_tokens:
raise ValueError(
f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
)
return _quantize_k_cache_fast_separate(
k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
)
# Copied from original
def _quantize_k_cache_ref(
input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d)
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
"""
Quantize the k-cache
Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size()
For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md
"""
assert dv % tile_size == 0
num_tiles = dv // tile_size
num_blocks, block_size, h_k, d = input_k_cache.shape
assert h_k == 1
input_k_cache = input_k_cache.squeeze(2) # [num_blocks, block_size, d]
input_elem_size = input_k_cache.element_size()
result = torch.empty(
(num_blocks, block_size, dv + num_tiles * 4 + input_elem_size * (d - dv)),
dtype=torch.float8_e4m3fn,
device=input_k_cache.device,
)
result_k_nope_part = result[..., :dv]
result_k_scale_factor = result[..., dv : dv + num_tiles * 4].view(torch.float32)
result_k_rope_part = result[..., dv + num_tiles * 4 :].view(input_k_cache.dtype)
result_k_rope_part[:] = input_k_cache[..., dv:]
for tile_idx in range(0, num_tiles):
cur_scale_factors_inv = (
torch.abs(
input_k_cache[..., tile_idx * tile_size : (tile_idx + 1) * tile_size]
)
.max(dim=-1)
.values
/ 448.0
) # [num_blocks, block_size]
result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv
cur_scale_factors_inv.unsqueeze_(-1) # [num_blocks, block_size, 1]
cur_quantized_nope = (
input_k_cache[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].float()
/ cur_scale_factors_inv.float()
).to(torch.float8_e4m3fn)
result_k_nope_part[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_quantized_nope
)
result = result.view(num_blocks, block_size, 1, -1)
return result
def _quantize_k_cache_fast_wrapped(
input_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
# TODO the final API may be 2D instead of 4D, thus we convert them here
num_blocks, block_size, _, dim_nope_and_rope = input_k_cache.shape
assert dv == 512
assert dim_nope_and_rope == 512 + 64
assert tile_size == 128
input_k_cache = input_k_cache.view((-1, dim_nope_and_rope))
# TODO deliberately split into two tensors, then upstream can provide the two tensors instead of concat into one
k_nope = input_k_cache[:, :dv]
k_rope = input_k_cache[:, dv:]
output = _quantize_k_cache_fast(k_nope=k_nope, k_rope=k_rope)
return output.view(num_blocks, block_size, 1, -1)
def _quantize_k_cache_fast(k_nope, k_rope, group_size: int = 128):
"""
:param k_nope: (num_tokens, dim_nope 512)
:param k_rope: (num_tokens, dim_rope 64)
"""
assert k_nope.dtype == torch.bfloat16
assert k_rope.dtype == torch.bfloat16
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_
assert dim_nope == 512
assert dim_rope == 64
assert k_nope.dtype == k_rope.dtype
num_tiles = dim_nope // group_size
assert k_nope.stride(1) == 1
assert k_rope.stride(1) == 1
output = torch.empty(
(num_tokens, dim_nope + num_tiles * 4 + k_rope.element_size() * dim_rope),
dtype=torch.float8_e4m3fn,
device=k_nope.device,
)
output_nope_q = output[..., :dim_nope]
output_nope_s = output[..., dim_nope : dim_nope + num_tiles * 4].view(torch.float32)
output_rope = output[..., dim_nope + num_tiles * 4 :].view(torch.bfloat16)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
NUM_NOPE_BLOCKS = dim_nope // group_size
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output_nope_q,
output_nope_s,
output_rope,
k_nope,
k_rope,
output_nope_q.stride(0),
output_nope_s.stride(0),
output_rope.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
return output
def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
"""
Quantize k_nope and k_rope in a single Triton kernel, directly outputting two separate tensors.
This avoids packing/unpacking and enables direct use with set_mla_kv_buffer_triton.
:param k_nope: (num_tokens, dim_nope 512) bfloat16
:param k_rope: (num_tokens, dim_rope 64) bfloat16
:param group_size: quantization tile size (default 128, kernel is tuned for this value)
:return: Tuple of (nope_part_u8, rope_part_u8)
- nope_part_u8: (num_tokens, 1, nope_part_bytes) uint8, layout [nope_fp8(dim_nope) | scales(num_tiles*4)]
- rope_part_u8: (num_tokens, 1, rope_part_bytes) uint8, layout [rope_bf16_bytes(dim_rope*2)]
"""
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_, f"k_nope and k_rope must have same num_tokens"
# Ensure contiguous tensors for kernel
k_nope = k_nope.contiguous()
k_rope = k_rope.contiguous()
num_tiles = dim_nope // group_size
# Calculate byte sizes based on validated dimensions
# nope_part: [FP8 quantized data (dim_nope bytes)] + [FP32 scales (num_tiles * 4 bytes)]
# rope_part: [BF16 raw data (dim_rope * 2 bytes)]
nope_part_bytes = (
dim_nope + num_tiles * 4
) # e.g., 512 + 4*4 = 528 for dim_nope=512, group_size=128
rope_part_bytes = (
dim_rope * k_rope.element_size()
) # e.g., 64 * 2 = 128 for dim_rope=64, BF16
# Allocate two separate output buffers (as uint8 for direct byte-level access)
nope_part_u8 = torch.empty(
(num_tokens, nope_part_bytes), dtype=torch.uint8, device=k_nope.device
)
rope_part_u8 = torch.empty(
(num_tokens, rope_part_bytes), dtype=torch.uint8, device=k_rope.device
)
# Create typed views for the kernel to write into
# Fixed byte layout for nope_part: [nope_fp8 (dim_nope bytes) | scales_fp32 (num_tiles*4 bytes)]
# Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
rope_view = rope_part_u8.view(torch.bfloat16)
# Kernel launch parameters
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
NUM_NOPE_BLOCKS = dim_nope // group_size
# Use the same kernel as _quantize_k_cache_fast (reuse existing implementation)
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
nope_q_view,
nope_s_view,
rope_view,
k_nope,
k_rope,
nope_q_view.stride(0),
nope_s_view.stride(0),
rope_view.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
# Add middle dimension for compatibility with set_mla_kv_buffer_triton
return nope_part_u8.unsqueeze(1), rope_part_u8.unsqueeze(1)
@triton.jit
def _quantize_k_cache_fast_kernel(
output_nope_q_ptr,
output_nope_s_ptr,
output_rope_ptr,
k_nope_ptr,
k_rope_ptr,
output_nope_q_stride_0: int,
output_nope_s_stride_0: int,
output_rope_stride_0: int,
k_nope_stride_0: int,
k_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
FP8_MIN: tl.constexpr,
FP8_MAX: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. quant nope
effective_block_id = raw_block_id
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_NOPE
ptr = k_nope_ptr + token_id * k_nope_stride_0 + offs
y = tl.load(ptr, mask=mask, other=0.0).to(tl.float32)
# the ref impl do not have a `tl.maximum(... eps)`, so we remove it here
y_s = tl.max(tl.abs(y)) / FP8_MAX
y_s_inv = 1.0 / y_s
y_q = tl.clamp(y * y_s_inv, FP8_MIN, FP8_MAX).to(
output_nope_q_ptr.dtype.element_ty
)
dst_q_ptr = output_nope_q_ptr + token_id * output_nope_q_stride_0 + offs
dst_s_ptr = (
output_nope_s_ptr + token_id * output_nope_s_stride_0 + effective_block_id
)
tl.store(dst_q_ptr, y_q, mask=mask)
tl.store(dst_s_ptr, y_s)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = k_rope_ptr + token_id * k_rope_stride_0 + offs
dst_ptr = output_rope_ptr + token_id * output_rope_stride_0 + offs
data = tl.load(src_ptr, mask=mask)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
import dequant_k_cache
for num_blocks, block_size in [
(1, 1),
(10, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
ref_quant = _quantize_k_cache_ref(input_k_cache)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
ref_ref_dequant = dequant_k_cache._dequantize_k_cache_slow(ref_quant)
ref_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(ref_quant)
actual_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(
actual_quant
)
print(f"{ref_ref_dequant=}")
print(f"{actual_actual_dequant=}")
print(f"{actual_actual_dequant - ref_ref_dequant=}")
print(f"{torch.mean(ref_ref_dequant - actual_actual_dequant)=}")
# TODO too different?
torch.testing.assert_close(
ref_ref_dequant, ref_actual_dequant, atol=0.2, rtol=0.2
)
torch.testing.assert_close(
ref_ref_dequant, actual_actual_dequant, atol=0.2, rtol=0.2
)
# test dequant_k_cache_paged
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
actual_dequant_paged = dequant_k_cache.dequantize_k_cache_paged(
actual_quant, page_table_1
).reshape(actual_actual_dequant.shape)
print(f"{torch.mean(actual_actual_dequant - actual_dequant_paged)=}")
torch.testing.assert_close(
ref_ref_dequant, actual_dequant_paged, atol=0.2, rtol=0.2
)
print("Passed")
# Test quantize_k_cache_separate: verify output matches concat path
print("\nTesting quantize_k_cache_separate...")
for num_tokens in [64, 100]:
dim_nope = 512
dim_rope = 64
k_nope = torch.randn(
num_tokens, 1, dim_nope, dtype=torch.bfloat16, device="cuda"
)
k_rope = torch.randn(
num_tokens, 1, dim_rope, dtype=torch.bfloat16, device="cuda"
)
# Old path: concat then quantize
k_concat = torch.cat([k_nope, k_rope], dim=-1).squeeze(1) # (num_tokens, 576)
old_output = quantize_k_cache(k_concat.unsqueeze(1).unsqueeze(1)) # 4D input
old_output = old_output.squeeze(1).squeeze(1) # Back to (num_tokens, 656)
# New path: quantize separately
nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
new_bytes = torch.cat([nope_part.squeeze(1), rope_part.squeeze(1)], dim=-1)
# Compare byte-level equality
old_bytes = old_output.view(torch.uint8)
if old_bytes.shape != new_bytes.shape:
raise RuntimeError(
f"Shape mismatch: {old_bytes.shape} vs {new_bytes.shape}"
)
diff_bytes = (old_bytes != new_bytes).sum().item()
if diff_bytes > 0:
max_diff = (old_bytes.float() - new_bytes.float()).abs().max().item()
raise RuntimeError(
f"quantize_k_cache_separate output doesn't match concat path: "
f"{diff_bytes} differing bytes, max_diff={max_diff}"
)
print(f" num_tokens={num_tokens}: PASSED (outputs match byte-wise)")
print("quantize_k_cache_separate tests passed!")
print("\nDo benchmark...")
for num_blocks, block_size in [
(1, 64),
(64, 64),
(128, 64),
(256, 64),
(512, 64),
(1024, 64),
(2048, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
def run_ans():
return dequant_k_cache.dequantize_k_cache_paged(actual_quant, page_table_1)
ans_time: float = triton.testing.do_bench(run_ans, warmup=10, rep=20) / 1000 # type: ignore
print(f"seq_kv: {num_blocks * block_size}, time: {ans_time * 1e6: 4.0f} us")
# [Deprecated] Re-export shim for backward compatibility. Use dsa.quant_k_cache instead.
import warnings
warnings.warn(
"sglang.srt.layers.attention.nsa.quant_k_cache is deprecated; "
"use sglang.srt.layers.attention.dsa.quant_k_cache instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.quant_k_cache import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/tilelang_kernel.py
+8 -2587
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/nsa/transform_index.py
+9 -143
View File
@@ -1,144 +1,10 @@
from typing import List, Optional
# [Deprecated] Re-export shim for backward compatibility. Use dsa.transform_index instead.
import warnings
import torch
import triton
import triton.language as tl
def transform_index_page_table_prefill(**kwargs):
return transform_index_page_table_prefill_ref(**kwargs)
def transform_index_page_table_decode(**kwargs):
return transform_index_page_table_decode_ref(**kwargs)
@triton.jit
def transform_index_page_table_decode_kernel(
page_table_ptr: torch.Tensor,
topk_indices_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_size: tl.constexpr,
max_seqlen_k: tl.constexpr,
):
TOPK: tl.constexpr = 2048
req_id = tl.program_id(0)
page_table_ptr = page_table_ptr + req_id * max_seqlen_k
topk_indices_ptr = topk_indices_ptr + req_id * TOPK
result_ptr = result_ptr + req_id * TOPK
offset = tl.arange(0, TOPK) # topk should be 2048
loaded_topk_indices = tl.load(topk_indices_ptr + offset)
mask = loaded_topk_indices >= 0
loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask)
tl.store(result_ptr + offset, loaded_kv_indices, mask=mask)
tl.store(result_ptr + offset, -1, mask=~mask)
def transform_index_page_table_decode_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
"""
Transform the page table according to topk indices for sparse topk attention.
Args:
page_table: [qo_len, max_seqlen_k], the original page table
topk_indices: [qo_len, topk], the topk indices for each query position
Returns:
transformed_page_table: [qo_len, topk], the transformed page table
For out-of-bound indices in topk_indices, this should be filled with -1.
"""
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
assert topk_indices.shape[1] == 2048
qo_len = topk_indices.shape[0]
max_seqlen_k = page_table.shape[1]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
# Launch triton kernel
grid = (qo_len,)
transform_index_page_table_decode_kernel[grid](
page_table,
topk_indices,
result,
page_size,
max_seqlen_k=max_seqlen_k,
)
return result
def transform_index_page_table_prefill_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
# TODO(baizhou): can be implemented with another triton kernel
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_fast(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
def transform_index_page_table_decode_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert result.shape == topk_indices.shape
torch.gather(
page_table.to(result.dtype),
dim=1,
index=topk_indices.clamp(min=0),
out=result,
)
result[topk_indices < 0] = -1
return result
def transform_index_page_table_prefill_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_ref(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
if __name__ == "__main__":
bs, topk, max_seqlen = 10, 2048, 3000
page_table = torch.randint(0, 100, (bs, max_seqlen), device="cuda")
topk_indices = torch.full((bs, topk), -1, device="cuda")
topk_indices[:, :1600] = torch.arange(1600).unsqueeze(0).repeat(bs, 1)
ref_result = transform_index_page_table_decode_ref(page_table, topk_indices)
result = transform_index_page_table_decode_fast(page_table, topk_indices)
assert torch.all(result == ref_result)
print("Passed")
warnings.warn(
"sglang.srt.layers.attention.nsa.transform_index is deprecated; "
"use sglang.srt.layers.attention.dsa.transform_index instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.transform_index import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/triton_kernel.py
+9 -195
View File
@@ -1,196 +1,10 @@
from typing import Optional, Tuple
# [Deprecated] Re-export shim for backward compatibility. Use dsa.triton_kernel instead.
import warnings
import torch
import triton
import triton.language as tl
# Triton implementation
@triton.jit
def _act_quant_kernel(
X_ptr,
Y_ptr,
S_ptr,
M,
N,
group_size: tl.constexpr,
round_scale: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
"""
Triton kernel for activation quantization.
Each block processes BLOCK_M rows and group_size columns.
"""
# Get block IDs
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
# FP8 constants
fp8_min = -448.0
fp8_max = 448.0
fp8_max_inv = 1.0 / fp8_max
# Calculate row and column offsets
row_start = pid_m * BLOCK_M
col_start = pid_n * group_size
# Create offset arrays
rows = row_start + tl.arange(0, BLOCK_M)
cols = col_start + tl.arange(0, BLOCK_N)
# Mask for valid rows and columns
row_mask = rows < M
col_mask = cols < N
mask = row_mask[:, None] & col_mask[None, :]
# Load input data
x_ptrs = X_ptr + rows[:, None] * N + cols[None, :]
x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
# Compute absolute max along columns (group_size dimension) for each row
x_abs = tl.abs(x)
amax = tl.max(x_abs, axis=1) # Shape: (BLOCK_M,)
# Clamp amax to avoid division by zero
amax = tl.maximum(amax, 1e-4)
# Compute scale
if round_scale:
# Fast round scale using bit manipulation approximation
# This is a simplified version - the exact bit manipulation is harder in Triton
# Using log2 + ceil + pow2 as approximation
log_val = tl.log2(amax * fp8_max_inv)
log_ceil = tl.ceil(log_val)
scale = tl.exp2(log_ceil)
else:
scale = amax * fp8_max_inv
# Quantize: y = clamp(x / scale, fp8_min, fp8_max)
scale_broadcast = scale[:, None]
y = x / scale_broadcast
y = tl.minimum(tl.maximum(y, fp8_min), fp8_max)
# Store quantized output
y_ptrs = Y_ptr + rows[:, None] * N + cols[None, :]
tl.store(y_ptrs, y, mask=mask)
# Store scales
s_cols = pid_n
s_ptrs = S_ptr + rows * (N // group_size) + s_cols
s_mask = row_mask
tl.store(s_ptrs, scale, mask=s_mask)
def act_quant(
x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantizes the input tensor `x` using block-wise quantization with Triton.
Args:
x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
scale_fmt (Optional[str], optional): The format of the scale. Default is None.
Returns:
Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
- The quantized tensor with dtype `torch.float8_e4m3fn`.
- A tensor of scaling factors with dtype `torch.float32`.
"""
assert x.is_contiguous(), "Input tensor must be contiguous"
assert (
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
# Flatten all dims except last
N = x.size(-1)
x_flat = x.view(-1, N)
M = x_flat.size(0)
# Allocate output tensors
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
y_flat = y.view(-1, N)
s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
s_flat = s.view(-1, N // block_size)
# Launch kernel
BLOCK_M = 32
BLOCK_N = block_size
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, block_size))
round_scale = scale_fmt is not None
_act_quant_kernel[grid](
x_flat,
y_flat,
s_flat,
M,
N,
group_size=block_size,
round_scale=round_scale,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
num_stages=0 if round_scale else 2,
)
return y, s
@triton.jit
def _get_valid_kv_indices_kernel(
page_table_ptr, # [bs, topk]
kv_indptr_ptr, # [bs + 1]
kv_indices_ptr, # [bs * topk] output buffer
bs: tl.constexpr,
topk: tl.constexpr,
):
"""
Extract valid indices (non -1) from page_table into kv_indices.
Each program handles one batch.
"""
batch_id = tl.program_id(0)
# Get the start position for this batch in kv_indices
dst_start = tl.load(kv_indptr_ptr + batch_id)
# Load all topk indices for this batch
src_offset = batch_id * topk
offsets = tl.arange(0, topk)
indices = tl.load(page_table_ptr + src_offset + offsets)
# Count valid indices and compact them
mask = indices != -1
# Use prefix sum to compute destination positions for valid elements
# For each position, count how many valid elements are before it
prefix_sum = tl.cumsum(mask.to(tl.int32), axis=0) - 1
# Store valid indices to their compacted positions
dst_positions = dst_start + prefix_sum
tl.store(kv_indices_ptr + dst_positions, indices, mask=mask)
def get_valid_kv_indices(
page_table_1: torch.Tensor,
kv_indptr: torch.Tensor,
kv_indices: torch.Tensor,
bs: int,
):
"""
Extract valid indices from page_table_1 into kv_indices buffer.
Args:
page_table_1: [bs, topk] page table with -1 as invalid
kv_indptr: [bs + 1] cumulative count of valid indices per batch
kv_indices: [bs * topk] pre-allocated output buffer
bs: batch size
"""
topk = page_table_1.shape[1]
grid = (bs,)
_get_valid_kv_indices_kernel[grid](
page_table_1,
kv_indptr,
kv_indices,
bs,
topk,
)
warnings.warn(
"sglang.srt.layers.attention.nsa.triton_kernel is deprecated; "
"use sglang.srt.layers.attention.dsa.triton_kernel instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.triton_kernel import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/utils.py
+8 -267
View File
@@ -1,269 +1,10 @@
from functools import lru_cache
from typing import TYPE_CHECKING, List, Tuple, Union
# [Deprecated] Re-export shim for backward compatibility. Use dsa.utils instead.
import warnings
import torch
import triton
import triton.language as tl
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
warnings.warn(
"sglang.srt.layers.attention.nsa.utils is deprecated; "
"use sglang.srt.layers.attention.dsa.utils instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.common import ceil_align, ceil_div
@lru_cache(maxsize=1)
def aiter_can_use_preshuffle_paged_mqa() -> bool:
"""Whether aiter's preshuffle paged MQA / cache kernels can be used on this runtime.
aiter's ``deepgemm_fp8_paged_mqa_logits`` only supports ``KVBlockSize > 1`` and
``Preshuffle=True`` on its gluon kernel path. The gluon path is enabled when
Triton >= 3.5.0, OR when ``AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1`` is set
(which additionally requires that the AOT gluon kernel artifacts ship inside
the aiter wheel/image). Otherwise aiter asserts ``KVBlockSize == 1`` and
refuses ``Preshuffle=True``.
sglang's NSA indexer uses this single decision to pick:
* ``page_size``: 64 (preshuffle) vs 1 (legacy) on ROCm
* ``Preshuffle`` / ``preshuffle`` flags on the aiter MQA + cache kernels
* ``get_page_table_64`` vs ``get_page_table_1`` on the metadata
* whether ``GetKAndS.execute`` uses the aiter or the triton implementation
The result is cached so the cost is paid once per process.
Set ``SGLANG_NSA_HIP_DISABLE_PRESHUFFLE=1`` to force the legacy path even when
the gluon kernel would otherwise be available (useful for CI bisection).
"""
if not is_hip():
return False
if not get_bool_env_var("SGLANG_USE_AITER"):
return False
if get_bool_env_var("SGLANG_NSA_HIP_DISABLE_PRESHUFFLE"):
return False
if get_bool_env_var("AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS"):
return True
try:
from packaging.version import Version
return Version(Version(triton.__version__).base_version) >= Version("3.5.0")
except Exception:
return False
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def compute_nsa_seqlens(original_seq_lens, nsa_index_topk: int):
return original_seq_lens.clamp(max=nsa_index_topk)
def is_nsa_enable_prefill_cp():
return get_global_server_args().enable_nsa_prefill_context_parallel
def is_nsa_prefill_cp_in_seq_split():
return (
is_nsa_enable_prefill_cp()
and get_global_server_args().nsa_prefill_cp_mode == "in-seq-split"
)
def is_nsa_prefill_cp_round_robin_split():
return (
is_nsa_enable_prefill_cp()
and get_global_server_args().nsa_prefill_cp_mode == "round-robin-split"
)
def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_attention_cp_size()
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return (
is_nsa_prefill_cp_round_robin_split()
and seq_len > 0
and seq_len >= cp_size
and cp_size > 1
)
def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
"""
# for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
| +-----------before split------------+|
| token0, token1, token2, token3, token4, token5, token6, token7, ...
|
| +--------------result-------------------+
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
if isinstance(input_, (tuple, list)):
indices = range(cp_rank, len(input_), cp_size)
return input_[indices]
tokens = len(input_)
if tokens % cp_size != 0:
cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
if cur_len == 0:
return input_.new_empty(0, *input_.shape[1:])
indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
return input_[indices]
# for torch device tensor
return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
forward_batch.is_extend_in_batch, global_num_tokens
)
if dp_padding_mode.is_max_len():
tokens = max(global_num_tokens)
elif len(global_num_tokens) > 1:
tokens = global_num_tokens[get_attention_dp_rank()]
else:
tokens = global_num_tokens[0]
if can_nsa_prefill_cp_round_robin_split(forward_batch):
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
attn_cp_size = get_attention_cp_size()
needs_cp_pad = attn_cp_size > 1 and can_nsa_prefill_cp_round_robin_split(
forward_batch
)
needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
if not needs_cp_pad and not needs_dp_pad:
return nsa_cache_seqlens
tokens = cal_padded_tokens(forward_batch)
pad_len = tokens - nsa_cache_seqlens.shape[0]
if pad_len > 0:
nsa_cache_seqlens = torch.cat(
[
nsa_cache_seqlens,
nsa_cache_seqlens.new_zeros(pad_len, *nsa_cache_seqlens.shape[1:]),
]
)
return nsa_cache_seqlens
def can_nsa_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
seq_len % cp_size == 0
), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split"
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
# the seq data needs to be divided and recombined at twice the size of cp_size.
cur_cp_seq_len = seq_len // (cp_size * 2)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_nsa
and forward_batch.forward_mode.is_context_parallel_extend()
and is_nsa_enable_prefill_cp()
and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
):
return True
else:
return False
@triton.jit
def nsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
out_seqs_ptr,
bs_idx_ptr,
tokens: tl.constexpr,
cp_size: tl.constexpr,
cp_rank: tl.constexpr,
):
extra_seq = 0
bs_idx = 0
for bs in range(tokens):
cur_len = tl.load(in_seqs_ptr + bs)
cur_len += extra_seq
cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank)
if cur_seq > 0:
tl.store(bs_idx_ptr + bs_idx, bs)
tl.store(out_seqs_ptr + bs_idx, cur_seq)
bs_idx += 1
extra_seq = cur_len - cur_seq * cp_size
def nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_seqs.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
q_seqs = [q_len for q_len in q_seqs if q_len > 0]
return q_seqs, bs_idx
def nsa_cp_round_robin_split_q_seqs(
extend_seqs_cpu, extend_seqs
) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
"""
round-robin-split distributes tokens across ranks based on token_idx % cp_size.
Return:
ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
for each sequence after distribution across cp ranks.
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_q_lens = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
)
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
grid = (1,)
nsa_cp_round_robin_split_q_seqs_kernel[grid](
extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None):
if nsa_enable_prefill_cp is None:
nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if (
forward_batch.attn_cp_metadata is not None
and nsa_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
):
return True
else:
return False
from sglang.srt.layers.attention.dsa.utils import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa_backend.py
+20 -2517
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/communicator.py
+20 -20
View File
@@ -33,9 +33,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
from sglang.srt.layers.attention.dsa.utils import (
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.dp_attention import (
attn_tp_all_gather_into_tensor,
@@ -202,7 +202,7 @@ class ScatterMode(Enum):
@staticmethod
def model_input_output():
"""The scatter mode for model forward pass input and output data"""
if is_nsa_enable_prefill_cp():
if is_dsa_enable_prefill_cp():
return ScatterMode.SCATTERED
return ScatterMode.TP_ATTN_FULL
@@ -256,15 +256,15 @@ class AttnTpContext:
self.allow_input_scattered = False
self.input_scattered_ = False
self.attn_inputs_: Optional[AttentionInputs] = None
self.is_nsa = False
self.is_dsa = False
def init_context(self, q_lora_rank, is_nsa):
self.is_nsa = is_nsa
def init_context(self, q_lora_rank, is_dsa):
self.is_dsa = is_dsa
self.allow_input_scattered = (
get_global_server_args().enable_attn_tp_input_scattered
and (_is_cuda or _is_npu)
and q_lora_rank is not None
and not is_nsa
and not is_dsa
and get_tensor_model_parallel_world_size() > 1
and not is_dp_attention_enabled()
and get_moe_a2a_backend().is_none()
@@ -379,8 +379,8 @@ class LayerScatterModes:
or should_use_flashinfer_cutlass_moe_fp4_allgather()
):
return ScatterMode.SCATTERED
# NSA CP doesn't support MOE_FULL yet; fall back to FULL
if is_enable_moe_cp_allgather() and not is_nsa_enable_prefill_cp():
# DSA CP doesn't support MOE_FULL yet; fall back to FULL
if is_enable_moe_cp_allgather() and not is_dsa_enable_prefill_cp():
return ScatterMode.MOE_FULL
return ScatterMode.FULL
else:
@@ -551,10 +551,10 @@ class LayerCommunicator:
)
elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"):
# aiter (ROCm gfx95) fused RMSNorm + FP8 group quant.
# When NSA is active, also preserve the unquantized bf16
# output as a 3-tuple (fp8, scale, bf16) so the NSA
# When DSA is active, also preserve the unquantized bf16
# output as a 3-tuple (fp8, scale, bf16) so the DSA
# indexer can skip redundant FP8 dequantization.
_nsa_needs_bf16 = get_attn_tp_context().is_nsa
_dsa_needs_bf16 = get_attn_tp_context().is_dsa
hidden_states, _unq_bf16, _, _res = fused_rms_fp8_group_quant(
hidden_states,
self.input_layernorm.weight,
@@ -565,9 +565,9 @@ class LayerCommunicator:
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=_nsa_needs_bf16,
output_unquantized_inp1=_dsa_needs_bf16,
)
if _nsa_needs_bf16:
if _dsa_needs_bf16:
hidden_states = (
hidden_states[0],
hidden_states[1],
@@ -596,9 +596,9 @@ class LayerCommunicator:
)
elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"):
# aiter (ROCm gfx95) fused RMSNorm + FP8 group quant
# with residual addition. When NSA is active, pack
# with residual addition. When DSA is active, pack
# the unquantized bf16 as a 3-tuple (fp8, scale, bf16).
_nsa_needs_bf16 = get_attn_tp_context().is_nsa
_dsa_needs_bf16 = get_attn_tp_context().is_dsa
hidden_states, _unq_bf16, _, residual = (
fused_rms_fp8_group_quant(
hidden_states,
@@ -610,10 +610,10 @@ class LayerCommunicator:
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=residual,
output_unquantized_inp1=_nsa_needs_bf16,
output_unquantized_inp1=_dsa_needs_bf16,
)
)
if _nsa_needs_bf16:
if _dsa_needs_bf16:
hidden_states = (
hidden_states[0],
hidden_states[1],
@@ -709,7 +709,7 @@ class LayerCommunicator:
return True
if forward_batch.dp_padding_mode.is_max_len():
return True
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
return True
if get_attn_tp_context().input_scattered and not self.is_last_layer:
return True
python/sglang/srt/layers/communicator_nsa_cp.py → python/sglang/srt/layers/communicator_dsa_cp.py
+19 -19
View File
@@ -18,9 +18,9 @@ from typing import Callable, Optional
import torch
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
from sglang.srt.layers.attention.dsa.utils import (
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.communicator import (
CommunicateContext,
@@ -40,14 +40,14 @@ from sglang.srt.layers.dp_attention import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def nsa_enable_prefill_cp():
def dsa_enable_prefill_cp():
# After using cp, the communication mode of this part changes.
# The three parts of prepare_attn, prepare_mlp, and postprocess_layer
# no longer require additional communication for reduce, scatter, etc.
return is_nsa_enable_prefill_cp()
return is_dsa_enable_prefill_cp()
class NSACPLayerCommunicator(LayerCommunicator):
class DSACPLayerCommunicator(LayerCommunicator):
def __init__(
self,
layer_scatter_modes: LayerScatterModes,
@@ -73,19 +73,19 @@ class NSACPLayerCommunicator(LayerCommunicator):
assert (
self._context.attn_dp_size == 1
), f"dp_size should be 1 when moe_runner_backend is none"
self._communicate_simple_fn = NSACPCommunicateSimpleFn.get_fn(
self._communicate_simple_fn = DSACPCommunicateSimpleFn.get_fn(
input_mode=ScatterMode.SCATTERED,
output_mode=ScatterMode.SCATTERED,
context=self._context,
)
self._communicate_with_all_reduce_and_layer_norm_fn = NSACPCommunicateWithAllReduceAndLayerNormFn.get_fn(
self._communicate_with_all_reduce_and_layer_norm_fn = DSACPCommunicateWithAllReduceAndLayerNormFn.get_fn(
hidden_states_input_mode=ScatterMode.SCATTERED,
residual_input_mode=ScatterMode.SCATTERED,
hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL
residual_output_mode=ScatterMode.SCATTERED,
context=self._context,
)
self._communicate_summable_tensor_pair_fn = NSACPCommunicateSummableTensorPairFn.get_fn(
self._communicate_summable_tensor_pair_fn = DSACPCommunicateSummableTensorPairFn.get_fn(
hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL
residual_input_mode=ScatterMode.SCATTERED,
output_mode=ScatterMode.SCATTERED,
@@ -93,7 +93,7 @@ class NSACPLayerCommunicator(LayerCommunicator):
)
class NSACPCommunicateSimpleFn(CommunicateSimpleFn):
class DSACPCommunicateSimpleFn(CommunicateSimpleFn):
@staticmethod
def get_fn(
input_mode: ScatterMode,
@@ -101,12 +101,12 @@ class NSACPCommunicateSimpleFn(CommunicateSimpleFn):
context: CommunicateContext,
):
if context.is_same_group_size(input_mode, output_mode):
return NSACPCommunicateSimpleFn._trivial
return DSACPCommunicateSimpleFn._trivial
raise NotImplementedError(f"{input_mode=} {output_mode=}")
class NSACPCommunicateWithAllReduceAndLayerNormFn(
class DSACPCommunicateWithAllReduceAndLayerNormFn(
CommunicateWithAllReduceAndLayerNormFn
):
"""Besides communication, needs to
@@ -126,11 +126,11 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
assert residual_input_mode == ScatterMode.SCATTERED
assert residual_output_mode == ScatterMode.SCATTERED
if hidden_states_output_mode == ScatterMode.SCATTERED:
return NSACPCommunicateWithAllReduceAndLayerNormFn._simple
return DSACPCommunicateWithAllReduceAndLayerNormFn._simple
if hidden_states_output_mode == ScatterMode.FULL:
return partial(
NSACPCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual,
DSACPCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual,
residual_input_mode=residual_input_mode,
)
@@ -152,7 +152,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
hidden_states, residual = layernorm(hidden_states, residual)
# for prefill: attn tp scattered -> full
# for decode: attn tp full -> full
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
assert context.attn_dp_size == 1
hidden_states, local_hidden_states = (
get_local_dp_buffer(get_attention_cp_group()),
@@ -165,7 +165,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
return hidden_states, residual
class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
class DSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
"""It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed."""
@staticmethod
@@ -184,12 +184,12 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
and (residual_input_mode == ScatterMode.SCATTERED)
and (output_mode == ScatterMode.SCATTERED)
):
return NSACPCommunicateSummableTensorPairFn._scatter_hidden_states
return DSACPCommunicateSummableTensorPairFn._scatter_hidden_states
if context.is_same_group_size(
hidden_states_input_mode, output_mode
) and context.is_same_group_size(residual_input_mode, output_mode):
return NSACPCommunicateSummableTensorPairFn._trivial
return DSACPCommunicateSummableTensorPairFn._trivial
raise NotImplementedError(
f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}"
@@ -205,7 +205,7 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
):
# for prefill: full -> attn tp scattered
# for decode: full -> attn tp full
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
assert context.attn_dp_size == 1
input_hidden_states = hidden_states
hidden_states = hidden_states.tensor_split(context.attn_cp_size)[
python/sglang/srt/layers/dp_attention.py
+1 -1
View File
@@ -72,7 +72,7 @@ class DpPaddingMode(IntEnum):
# When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding
# overhead from uneven token distribution.
# For dp_size=1, max_len equals sum_len, so prefer MAX_LEN mode
# to enable symmetric memory optimization (needed for NSA CP, etc.).
# to enable symmetric memory optimization (needed for DSA CP, etc.).
if is_extend_in_batch and dp_size > 1:
return DpPaddingMode.SUM_LEN
python/sglang/srt/layers/mhc.py
+2 -2
View File
@@ -8,7 +8,7 @@ import torch
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_round_robin_split
from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_round_robin_split
from sglang.srt.layers.utils.common import strict_contiguous
tilelang.set_log_level("WARNING")
@@ -880,7 +880,7 @@ def mhc_post(
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor:
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
x = strict_contiguous(x)
residual = strict_contiguous(residual)
post_layer_mix = strict_contiguous(post_layer_mix)
python/sglang/srt/layers/utils/cp_utils.py
+20 -20
View File
@@ -66,17 +66,17 @@ def can_cp_split(seq_len: int, cp_size: int, forward_batch):
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
nsa_cp_round_robin_split_data,
from sglang.srt.layers.attention.dsa.utils import (
dsa_cp_round_robin_split_data,
is_dsa_prefill_cp_round_robin_split,
)
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert (
input_.shape[0] % cp_size == 0
), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
return nsa_cp_round_robin_split_data(input_)
return dsa_cp_round_robin_split_data(input_)
input_list = list(
torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
@@ -88,18 +88,18 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
nsa_cp_round_robin_split_data,
from sglang.srt.layers.attention.dsa.utils import (
dsa_cp_round_robin_split_data,
is_dsa_prefill_cp_round_robin_split,
)
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert positions.shape[0] % cp_size == 0, (
f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
f"cp size {cp_size}"
)
return nsa_cp_round_robin_split_data(positions)
return dsa_cp_round_robin_split_data(positions)
position_id_list = list(
torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
@@ -238,11 +238,11 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
| token0, token1, token2, token3, token4, token5, token6, token7, ...
| +-------------------------+
"""
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
from sglang.srt.layers.attention.dsa.utils import (
is_dsa_prefill_cp_round_robin_split,
)
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
@@ -395,11 +395,11 @@ def prepare_context_parallel_metadata(
cp_size,
seqs_len,
):
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_prefill_cp_round_robin_split,
from sglang.srt.layers.attention.dsa.utils import (
is_dsa_prefill_cp_round_robin_split,
)
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
return ContextParallelMetadata()
"""prepare_input_dp_with_cp_dsa-zigzag index
@@ -505,16 +505,16 @@ def prepare_context_parallel_metadata(
# TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues
# Prefix offset is critical when radix cache hits (prefix_len > 0).
# For non-NSA CP (e.g. qwen3-moe), consumers use these values directly as
# For non-DSA CP (e.g. qwen3-moe), consumers use these values directly as
# FlashAttention cache_seqlens, so the prefix must be baked in here.
# For NSA CP, `_get_topk_ragged_with_cp` re-adds the cached-prefix offset
# For DSA CP, `_get_topk_ragged_with_cp` re-adds the cached-prefix offset
# from (seq_lens_cpu - extend_seq_lens_cpu); baking prefix_len in here
# would silently drop it whenever the scheduler packs multiple requests
# into a single CP extend (len(seqs_len) != 1 -> prefix_len falls back
# to 0), corrupting the indexer's ke_offset on prefix-cache hits.
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
if is_nsa_enable_prefill_cp():
if is_dsa_enable_prefill_cp():
kv_len_prev = prefix_sum_list[cp_rank]
kv_len_next = prefix_sum_list[cp_size * 2 - cp_rank - 1]
else:
python/sglang/srt/managers/hisparse_coordinator.py
+2 -2
View File
@@ -9,7 +9,7 @@ from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
DeepSeekV4SingleKVPoolHost,
HiSparseNSATokenToKVPool,
HiSparseDSATokenToKVPool,
HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
@@ -79,7 +79,7 @@ class HiSparseCoordinator:
assert isinstance(
self.token_to_kv_pool_allocator, HiSparseTokenToKVPoolAllocator
)
self.mem_pool_device: HiSparseNSATokenToKVPool = (
self.mem_pool_device: HiSparseDSATokenToKVPool = (
self.token_to_kv_pool_allocator.get_kvcache()
)
self.mem_pool_host = MLATokenToKVPoolHost(
python/sglang/srt/managers/schedule_policy.py
+3 -3
View File
@@ -34,7 +34,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
import torch
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.mem_cache.base_prefix_cache import (
@@ -467,7 +467,7 @@ class PrefillAdder:
self.priority_scheduling_preemption_threshold = (
priority_scheduling_preemption_threshold
)
self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split()
self.dsa_prefill_cp_in_seq_split = is_dsa_prefill_cp_in_seq_split()
self.max_running_requests = max_running_requests
self.prefill_context_parallel_enabled = is_prefill_context_parallel_enabled()
self.prefill_max_requests = prefill_max_requests
@@ -826,7 +826,7 @@ class PrefillAdder:
# Enabling context parallelism currently presents precision issues;
# therefore, the prefill-batch setting is temporarily set to 1.
if (
self.nsa_prefill_cp_in_seq_split or self.prefill_context_parallel_enabled
self.dsa_prefill_cp_in_seq_split or self.prefill_context_parallel_enabled
) and len(self.can_run_list) >= 1:
return AddReqResult.OTHER
python/sglang/srt/managers/scheduler_pp_mixin.py
+1 -1
View File
@@ -523,7 +523,7 @@ class SchedulerPPMixin:
self.pp_loop_size: int = self.ps.pp_size + self.server_args.pp_async_batch_depth
# In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation.
self.require_attn_tp_allgather = (
not self.server_args.enable_nsa_prefill_context_parallel
not self.server_args.enable_dsa_prefill_context_parallel
)
self.mbs = [None] * self.pp_loop_size
self.last_mbs = [None] * self.pp_loop_size
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+1 -1
View File
@@ -9,11 +9,11 @@ import torch
from sglang.jit_kernel.deepseek_v4 import fused_k_norm_rope_flashmla, fused_store_cache
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa import index_buf_accessor
from sglang.srt.layers.attention.dsv4 import (
index_buf_accessor as dsv4_index_buf_accessor,
)
from sglang.srt.layers.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack
from sglang.srt.layers.attention.nsa import index_buf_accessor
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
from sglang.srt.mem_cache.memory_pool import KVCache
python/sglang/srt/mem_cache/hiradix_cache.py
+8 -10
View File
@@ -36,12 +36,12 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController,
)
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
attach_hybrid_nsa_pool_to_hiradix_cache,
attach_hybrid_dsa_pool_to_hiradix_cache,
)
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
MHATokenToKVPool,
MLATokenToKVPool,
NSATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
@@ -82,8 +82,8 @@ class HiRadixCache(RadixCache):
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
)
elif isinstance(self.kv_cache, NSATokenToKVPool):
# Filled by attach_hybrid_nsa_pool_to_hiradix_cache after storage extra_config is parsed.
elif isinstance(self.kv_cache, DSATokenToKVPool):
# Filled by attach_hybrid_dsa_pool_to_hiradix_cache after storage extra_config is parsed.
self.token_to_kv_pool_host = None
elif isinstance(self.kv_cache, MLATokenToKVPool):
self.token_to_kv_pool_host = MLATokenToKVPoolHost(
@@ -95,9 +95,7 @@ class HiRadixCache(RadixCache):
allocator_type=server_args.hicache_storage_backend,
)
else:
raise ValueError(
"HiRadixCache only supports MHA, MLA, and NSA (DSA) models"
)
raise ValueError("HiRadixCache only supports MHA, MLA, and DSA models")
self.tp_group = params.tp_cache_group
self.attn_cp_group = params.attn_cp_cache_group
@@ -122,8 +120,8 @@ class HiRadixCache(RadixCache):
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
self.load_cache_event = threading.Event()
if isinstance(self.kv_cache, NSATokenToKVPool):
attach_hybrid_nsa_pool_to_hiradix_cache(
if isinstance(self.kv_cache, DSATokenToKVPool):
attach_hybrid_dsa_pool_to_hiradix_cache(
self,
params,
server_args,
@@ -643,7 +641,7 @@ class HiRadixCache(RadixCache):
def _get_extra_pools(self) -> dict:
if not isinstance(self.cache_controller, HybridCacheController):
return {}
if isinstance(self.kv_cache, NSATokenToKVPool):
if isinstance(self.kv_cache, DSATokenToKVPool):
pool = PoolTransfer(
name=PoolName.INDEXER,
hit_policy=PoolHitPolicy.ALL_PAGES,
python/sglang/srt/mem_cache/hisparse_memory_pool.py
+3 -3
View File
@@ -16,7 +16,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DeepSeekV4TokenToKVPool,
HiSparseC4DevicePool,
)
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import HiSparseHostPoolMixin
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils.common import get_num_new_pages
@@ -37,7 +37,7 @@ else:
)
class HiSparseNSATokenToKVPool(NSATokenToKVPool):
class HiSparseDSATokenToKVPool(DSATokenToKVPool):
def __init__(
self,
size: int,
@@ -143,7 +143,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
page_size: int,
dtype: torch.dtype,
device: torch.device,
kvcache: HiSparseNSATokenToKVPool,
kvcache: HiSparseDSATokenToKVPool,
need_sort: bool,
host_to_device_ratio: int = 2,
):
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+11 -11
View File
@@ -10,12 +10,12 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
from sglang.srt.mem_cache.memory_pool_host import (
DeepSeekV4PagedHostPool,
DeepSeekV4StateHostPool,
DSAIndexerPoolHost,
HostPoolGroup,
LogicalHostPool,
MambaPoolHost,
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSAIndexerPoolHost,
PoolEntry,
)
@@ -656,9 +656,9 @@ def attach_hybrid_pool_to_unified_cache(
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
HybridLinearKVPool,
MLATokenToKVPool,
NSATokenToKVPool,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.unified_cache_components import ComponentType
@@ -667,7 +667,7 @@ def attach_hybrid_pool_to_unified_cache(
kvcache = params.token_to_kv_pool_allocator.get_kvcache()
swa_stack = isinstance(kvcache, SWAKVPool)
mamba_stack = isinstance(kvcache, HybridLinearKVPool)
nsa_stack = isinstance(kvcache, NSATokenToKVPool)
dsa_stack = isinstance(kvcache, DSATokenToKVPool)
deepseek_v4_stack = isinstance(kvcache, DeepSeekV4TokenToKVPool)
if deepseek_v4_stack:
@@ -820,7 +820,7 @@ def attach_hybrid_pool_to_unified_cache(
cache.swa_kv_pool_host
)
transfer_layer_num = len(full_layer_mapping | swa_layer_mapping)
elif nsa_stack:
elif dsa_stack:
full_layer_mapping = {
layer_id: layer_id for layer_id in range(full_kv_pool.layer_num)
}
@@ -838,7 +838,7 @@ def attach_hybrid_pool_to_unified_cache(
storage_backend=None,
use_mla=use_mla,
override_kv_cache_dim=full_kv_pool.kv_cache_dim,
sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost(
sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost(
full_kv_pool,
kv_host_pool,
server_args.hicache_mem_layout,
@@ -897,7 +897,7 @@ def attach_hybrid_pool_to_unified_cache(
pools_desc = "KV + MAMBA"
elif swa_stack:
pools_desc = "KV + SWA"
elif nsa_stack:
elif dsa_stack:
pools_desc = "KV + INDEXER"
else:
pools_desc = "KV"
@@ -911,7 +911,7 @@ def attach_hybrid_pool_to_unified_cache(
raise
def attach_hybrid_nsa_pool_to_hiradix_cache(
def attach_hybrid_dsa_pool_to_hiradix_cache(
radix_cache: HiRadixCache,
params: CacheInitParams,
server_args: ServerArgs,
@@ -925,7 +925,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
) -> None:
"""Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache.
This entrypoint is currently intended only for HiRadixCache's NSA path.
This entrypoint is currently intended only for HiRadixCache's DSA path.
"""
try:
kv = radix_cache.kv_cache
@@ -945,7 +945,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
use_mla=True,
override_kv_cache_dim=kv.kv_cache_dim,
prefetch_threshold=prefetch_threshold,
sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost(
sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost(
kv,
kv_host_pool,
server_args.hicache_mem_layout,
@@ -961,12 +961,12 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
radix_cache.token_to_kv_pool_host = host_pool_group
radix_cache.cache_controller = cache_controller
logger.info(
"Attached hybrid NSA pool stack to HiRadixCache: pools=KV + INDEXER, "
"Attached hybrid DSA pool stack to HiRadixCache: pools=KV + INDEXER, "
"transfer_layer_num=%s",
len(layer_mapping),
)
except Exception:
logger.exception("attach_hybrid_nsa_pool_to_hiradix_cache failed")
logger.exception("attach_hybrid_dsa_pool_to_hiradix_cache failed")
raise
python/sglang/srt/mem_cache/memory_pool.py
+22 -22
View File
@@ -40,12 +40,12 @@ from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
from sglang.srt.configs.mamba_utils import BaseLinearStateParams
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import index_buf_accessor
from sglang.srt.layers.attention.nsa.quant_k_cache import (
from sglang.srt.layers.attention.dsa import index_buf_accessor
from sglang.srt.layers.attention.dsa.quant_k_cache import (
quantize_k_cache,
quantize_k_cache_separate,
)
from sglang.srt.layers.attention.nsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.utils import (
@@ -1618,7 +1618,7 @@ class MLATokenToKVPool(KVCache):
enable_memory_saver: bool,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
use_nsa: bool = False,
use_dsa: bool = False,
override_kv_cache_dim: Optional[int] = None,
):
super().__init__(
@@ -1634,17 +1634,17 @@ class MLATokenToKVPool(KVCache):
self.kv_lora_rank = kv_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
self.use_nsa = use_nsa
self.nsa_kv_cache_store_fp8 = (
use_nsa
self.use_dsa = use_dsa
self.dsa_kv_cache_store_fp8 = (
use_dsa
and dtype == torch.float8_e4m3fn
and override_kv_cache_dim is not None
)
# When override_kv_cache_dim is provided with nsa model, we assume the
# When override_kv_cache_dim is provided with dsa model, we assume the
# override kv cache dim is correct and use it directly.
self.kv_cache_dim = (
override_kv_cache_dim
if self.nsa_kv_cache_store_fp8
if self.dsa_kv_cache_store_fp8
else (kv_lora_rank + qk_rope_head_dim)
)
@@ -1655,8 +1655,8 @@ class MLATokenToKVPool(KVCache):
dtype=torch.uint64,
device=self.device,
)
if not use_nsa:
# NSA will allocate indexer KV cache later and then log the total size
if not use_dsa:
# DSA will allocate indexer KV cache later and then log the total size
self._finalize_allocation_log(size)
def _create_buffers(self):
@@ -1726,7 +1726,7 @@ class MLATokenToKVPool(KVCache):
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
assert not self.nsa_kv_cache_store_fp8
assert not self.dsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
@@ -1746,7 +1746,7 @@ class MLATokenToKVPool(KVCache):
):
layer_id = layer.layer_id
if _is_hip and self.use_nsa and self.dtype == fp8_dtype:
if _is_hip and self.use_dsa and self.dtype == fp8_dtype:
# HIP FP8 path uses raw MLA KV layout (nope + rope) without per-block scales.
# Fuse BF16/FP16 -> FP8 cast with paged KV write.
set_mla_kv_buffer_triton_fp8_quant(
@@ -1756,7 +1756,7 @@ class MLATokenToKVPool(KVCache):
cache_k_rope,
fp8_dtype,
)
elif self.nsa_kv_cache_store_fp8:
elif self.dsa_kv_cache_store_fp8:
# OPTIMIZATION: Quantize k_nope and k_rope separately to avoid concat overhead
# This also enables reuse of set_mla_kv_buffer_triton two-tensor write path
# quantize_k_cache_separate returns (nope_part, rope_part) as uint8 bytes
@@ -1905,7 +1905,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
assert not self.nsa_kv_cache_store_fp8
assert not self.dsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype:
from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil
@@ -1930,7 +1930,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
):
layer_id = layer.layer_id
if self.nsa_kv_cache_store_fp8:
if self.dsa_kv_cache_store_fp8:
# original cache_k: (num_tokens, num_heads 1, hidden 576); we unsqueeze the page_size=1 dim here
# TODO no need to cat
cache_k = torch.cat([cache_k_nope, cache_k_rope], dim=-1)
@@ -1968,7 +1968,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
)
class NSATokenToKVPool(MLATokenToKVPool):
class DSATokenToKVPool(MLATokenToKVPool):
quant_block_size = 128
index_k_with_scale_buffer_dtype = torch.uint8
rope_storage_dtype = torch.bfloat16 # rope is always stored in bf16
@@ -2005,7 +2005,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
enable_memory_saver,
start_layer,
end_layer,
use_nsa=True,
use_dsa=True,
override_kv_cache_dim=override_dim,
)
# self.index_k_dtype = torch.float8_e4m3fn
@@ -2013,7 +2013,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
self.index_head_dim = index_head_dim
if index_buf_size is None:
index_buf_size = size
# num head == 1 and head dim == 128 for index_k in NSA
# num head == 1 and head dim == 128 for index_k in DSA
assert index_head_dim == 128
if _is_hip:
@@ -2024,7 +2024,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
else:
assert (
self.page_size == 1
), f"HIP legacy NSA path requires page_size == 1, got {self.page_size}"
), f"HIP legacy DSA path requires page_size == 1, got {self.page_size}"
else:
assert self.page_size == 64
with (
@@ -2133,11 +2133,11 @@ class NSATokenToKVPool(MLATokenToKVPool):
)
def get_cpu_copy(self, indices):
# NSA keeps a page-indexed index_k_with_scale_buffer alongside kv_buffer.
# DSA keeps a page-indexed index_k_with_scale_buffer alongside kv_buffer.
# Retract frees the slots/pages and they get reused by other reqs'
# set_index_k_scale_buffer, so we must offload it here too -- otherwise
# resume restores kv_buffer but leaves foreign index/scale in place and
# NSA attention reads garbage at those token positions.
# DSA attention reads garbage at those token positions.
kv_cache_cpu = super().get_cpu_copy(indices)
page_indices = indices[:: self.page_size] // self.page_size
python/sglang/srt/mem_cache/memory_pool_host.py
+9 -9
View File
@@ -31,11 +31,11 @@ from sglang.jit_kernel.hicache import (
transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
)
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
KVCache,
MambaPool,
MHATokenToKVPool,
MLATokenToKVPool,
NSATokenToKVPool,
)
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
@@ -2608,14 +2608,14 @@ class HostPoolGroup:
)
class NSAIndexerPoolHost(HostKVCache):
"""Host-side NSA index buffers only. Slot layout matches the anchor MLA host pool."""
class DSAIndexerPoolHost(HostKVCache):
"""Host-side DSA index buffers only. Slot layout matches the anchor MLA host pool."""
device_pool: NSATokenToKVPool
device_pool: DSATokenToKVPool
def __init__(
self,
device_pool: NSATokenToKVPool,
device_pool: DSATokenToKVPool,
anchor_host: MLATokenToKVPoolHost,
layout: str,
pin_memory: bool = True,
@@ -2635,7 +2635,7 @@ class NSAIndexerPoolHost(HostKVCache):
self.index_head_dim = device_pool.index_head_dim
self.indexer_quant_block_size = device_pool.quant_block_size
self.indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype
self.indexer_dtype = DSATokenToKVPool.index_k_with_scale_buffer_dtype
self.indexer_size_per_token = (
self.index_head_dim
+ self.index_head_dim // self.indexer_quant_block_size * 4
@@ -2658,12 +2658,12 @@ class NSAIndexerPoolHost(HostKVCache):
available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
if requested_bytes > available_bytes:
raise ValueError(
f"Not enough host memory for NSA indexer hierarchical cache. "
f"Not enough host memory for DSA indexer hierarchical cache. "
f"Requesting {requested_bytes / 1e9:.2f} GB but only have "
f"{available_bytes / 1e9:.2f} GB free."
)
logger.info(
"Allocating %.2f GB host memory for NSA indexer (layout=%s).",
"Allocating %.2f GB host memory for DSA indexer (layout=%s).",
requested_bytes / 1e9,
layout,
)
@@ -2726,7 +2726,7 @@ class NSAIndexerPoolHost(HostKVCache):
return host_indices, device_indices
if host_indices.numel() % self.page_size != 0:
raise ValueError(
"Index buffer transfer expects page-aligned indices for NSA."
"Index buffer transfer expects page-aligned indices for DSA."
)
host_page_indices = (
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
python/sglang/srt/mem_cache/sparsity/__init__.py
+2 -2
View File
@@ -1,7 +1,7 @@
from sglang.srt.mem_cache.sparsity.algorithms import (
BaseSparseAlgorithm,
BaseSparseAlgorithmImpl,
DeepSeekNSAAlgorithm,
DeepSeekDSAAlgorithm,
QuestAlgorithm,
)
from sglang.srt.mem_cache.sparsity.backend import BackendAdaptor, FlashAttentionAdaptor
@@ -17,7 +17,7 @@ __all__ = [
"BaseSparseAlgorithm",
"BaseSparseAlgorithmImpl",
"QuestAlgorithm",
"DeepSeekNSAAlgorithm",
"DeepSeekDSAAlgorithm",
"BackendAdaptor",
"FlashAttentionAdaptor",
"SparseConfig",
python/sglang/srt/mem_cache/sparsity/algorithms/__init__.py
+2 -2
View File
@@ -2,12 +2,12 @@ from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
BaseSparseAlgorithm,
BaseSparseAlgorithmImpl,
)
from sglang.srt.mem_cache.sparsity.algorithms.deepseek_nsa import DeepSeekNSAAlgorithm
from sglang.srt.mem_cache.sparsity.algorithms.deepseek_dsa import DeepSeekDSAAlgorithm
from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm
__all__ = [
"BaseSparseAlgorithm",
"BaseSparseAlgorithmImpl",
"DeepSeekNSAAlgorithm",
"DeepSeekDSAAlgorithm",
"QuestAlgorithm",
]
python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_nsa.py → python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_dsa.py
+4 -4
View File
@@ -7,12 +7,12 @@ from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
)
class DeepSeekNSAAlgorithm(BaseSparseAlgorithmImpl):
class DeepSeekDSAAlgorithm(BaseSparseAlgorithmImpl):
"""
Sparse attention algorithm for DeepSeek NSA.
Sparse attention algorithm for DeepSeek DSA.
This algorithm uses NSA's native indexer for TopK retrieval.
Overrides all parent methods as NSA has its own specialized flow.
This algorithm uses DSA's native indexer for TopK retrieval.
Overrides all parent methods as DSA has its own specialized flow.
"""
def __init__(self, config, device: torch.device, **kwargs):
python/sglang/srt/mem_cache/sparsity/backend/__init__.py
+2 -2
View File
@@ -1,7 +1,7 @@
from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import (
BackendAdaptor,
DSABackendAdaptor,
FlashAttentionAdaptor,
NSABackendAdaptor,
)
__all__ = ["BackendAdaptor", "FlashAttentionAdaptor", "NSABackendAdaptor"]
__all__ = ["BackendAdaptor", "FlashAttentionAdaptor", "DSABackendAdaptor"]
python/sglang/srt/mem_cache/sparsity/backend/backend_adaptor.py
+4 -4
View File
@@ -46,8 +46,8 @@ class BackendAdaptor(ABC):
pass
class NSABackendAdaptor(BackendAdaptor):
"""Adaptor for NSA (Native Sparse Attention) backend."""
class DSABackendAdaptor(BackendAdaptor):
"""Adaptor for DSA (DeepSeek Sparse Attention) backend."""
def __init__(
self,
@@ -70,9 +70,9 @@ class NSABackendAdaptor(BackendAdaptor):
**kwargs,
) -> Optional[torch.Tensor]:
"""
Transform logical page indices to physical device indices for NSA backend.
Transform logical page indices to physical device indices for DSA backend.
"""
# TODO: Implement NSA backend adaptor logic
# TODO: Implement DSA backend adaptor logic
pass
python/sglang/srt/mem_cache/sparsity/factory.py
+5 -5
View File
@@ -5,11 +5,11 @@ from typing import Optional
import torch
from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import BaseSparseAlgorithm
from sglang.srt.mem_cache.sparsity.algorithms.deepseek_nsa import DeepSeekNSAAlgorithm
from sglang.srt.mem_cache.sparsity.algorithms.deepseek_dsa import DeepSeekDSAAlgorithm
from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm
from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import (
DSABackendAdaptor,
FlashAttentionAdaptor,
NSABackendAdaptor,
)
from sglang.srt.mem_cache.sparsity.core.sparse_coordinator import (
SparseConfig,
@@ -22,7 +22,7 @@ _global_sparse_coordinator: Optional[SparseCoordinator] = None
_ALGORITHM_REGISTRY = {
"quest": lambda config, device, **kw: QuestAlgorithm(config, device, **kw),
"deepseek_nsa": lambda config, device, **kw: DeepSeekNSAAlgorithm(
"deepseek_dsa": lambda config, device, **kw: DeepSeekDSAAlgorithm(
config, device, **kw
),
}
@@ -49,8 +49,8 @@ def _create_backend_adaptor(
req_to_token_pool,
):
"""Create backend adaptor."""
if isinstance(sparse_algorithm, DeepSeekNSAAlgorithm):
return NSABackendAdaptor(device, req_to_token_pool)
if isinstance(sparse_algorithm, DeepSeekDSAAlgorithm):
return DSABackendAdaptor(device, req_to_token_pool)
if backend in ["fa3", "flashattention"]:
return FlashAttentionAdaptor(device)
python/sglang/srt/model_executor/cuda_graph_runner.py
+6 -6
View File
@@ -43,7 +43,7 @@ from sglang.srt.distributed.parallel_state import (
)
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_size,
@@ -285,7 +285,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
seq_len_fill_value: int,
require_gathered_buffer: bool,
num_tokens_per_bs: int,
nsa_enable_prefill_cp: bool,
dsa_enable_prefill_cp: bool,
enable_num_token_non_padded_flag: bool,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
):
@@ -354,7 +354,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
if enable_num_token_non_padded_flag:
if require_gathered_buffer and not nsa_enable_prefill_cp:
if require_gathered_buffer and not dsa_enable_prefill_cp:
num_tokens_per_dp = bs * num_tokens_per_bs
local = compute_local_num_token_non_padded(
global_num_token_non_padded=forward_batch.num_token_non_padded,
@@ -588,7 +588,7 @@ class CudaGraphRunner:
self.attn_tp_size = get_attention_tp_size()
self.attn_tp_rank = get_attention_tp_rank()
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
@@ -946,7 +946,7 @@ class CudaGraphRunner:
if (
enable_num_token_non_padded()
and self.require_gathered_buffer
and not self.nsa_enable_prefill_cp
and not self.dsa_enable_prefill_cp
):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=buffers.num_token_non_padded,
@@ -1211,7 +1211,7 @@ class CudaGraphRunner:
seq_len_fill_value=self.seq_len_fill_value,
require_gathered_buffer=self.require_gathered_buffer,
num_tokens_per_bs=self.num_tokens_per_bs,
nsa_enable_prefill_cp=self.nsa_enable_prefill_cp,
dsa_enable_prefill_cp=self.dsa_enable_prefill_cp,
enable_num_token_non_padded_flag=enable_num_token_non_padded(),
pp_proxy_tensors=pp_proxy_tensors,
)
python/sglang/srt/model_executor/model_runner.py
+5 -4
View File
@@ -110,7 +110,7 @@ from sglang.srt.layers.attention.attention_registry import (
ATTENTION_BACKENDS,
attn_backend_wrapper,
)
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
@@ -246,7 +246,8 @@ MLA_ATTENTION_BACKENDS = [
"trtllm_mla",
"tokenspeed_mla",
"ascend",
"nsa",
"dsa",
"nsa", # Deprecated alias for "dsa"
"intel_xpu",
]
@@ -3061,7 +3062,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
# In DP Attention, IDLE batches are padded (batch_size > 0) for MLP sync.
# in this case, we need to reinit the forward metadata, otherwise the stale
# metadata causes batch_size mismatch in attention kernel(e.g. NSA Indexer).
# metadata causes batch_size mismatch in attention kernel(e.g. DSA Indexer).
if forward_batch.batch_size > 0:
self.attn_backend.init_forward_metadata(forward_batch)
@@ -3237,7 +3238,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
forward_batch.num_token_non_padded is not None
and forward_batch.global_num_tokens_gpu is not None
and require_gathered_buffer(self.server_args)
and not is_nsa_enable_prefill_cp()
and not is_dsa_enable_prefill_cp()
):
forward_batch.adjust_num_token_non_padded_for_attn_tp(
server_args=self.server_args,
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+28 -28
View File
@@ -6,8 +6,8 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.configs.model_config import (
get_nsa_index_head_dim,
is_deepseek_nsa,
get_dsa_index_head_dim,
is_deepseek_dsa,
is_deepseek_v4,
)
from sglang.srt.distributed.parallel_state import get_world_group
@@ -20,10 +20,11 @@ from sglang.srt.mem_cache.allocator import (
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
HiSparseNSATokenToKVPool,
HiSparseDSATokenToKVPool,
HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
HybridLinearKVPool,
HybridReqToTokenPool,
MHATokenToKVPool,
@@ -31,7 +32,6 @@ from sglang.srt.mem_cache.memory_pool import (
MLATokenToKVPool,
MLATokenToKVPoolFP4,
NoOpMHATokenToKVPool,
NSATokenToKVPool,
ReqToTokenPool,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
@@ -137,36 +137,36 @@ class ModelRunnerKVCacheMixin:
return total_rest_memory - mamba_state_memory
def calculate_mla_kv_cache_dim(self: ModelRunner) -> int:
is_nsa_model = is_deepseek_nsa(self.model_config.hf_config)
is_dsa_model = is_deepseek_dsa(self.model_config.hf_config)
kv_cache_dtype = self.kv_cache_dtype
kv_lora_rank = self.model_config.kv_lora_rank
qk_rope_head_dim = self.model_config.qk_rope_head_dim
kv_cache_dim = kv_lora_rank + qk_rope_head_dim # default mla kv cache dim
# For non-NSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim
if not is_nsa_model:
# For non-DSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim
if not is_dsa_model:
return kv_cache_dim
# TRTLLM backend does not override kv_cache_dim for MLA kv cache
# Assuming nsa prefill and decode backends are the same when using trtllm MLA backend,
# Assuming dsa prefill and decode backends are the same when using trtllm MLA backend,
# since it is not compatible for trtllm and other mla attn backend due to the different
# kv cache layout.
if (
self.server_args.nsa_prefill_backend == "trtllm"
or self.server_args.nsa_decode_backend == "trtllm"
self.server_args.dsa_prefill_backend == "trtllm"
or self.server_args.dsa_decode_backend == "trtllm"
):
return kv_cache_dim
# On HIP with TileLang backend, keep the default MLA KV cache dimension.
# FP8 attention uses the nope(512 fp8) + rope(64 fp8) layout, without extra per-block scales.
if _is_hip and (
self.server_args.nsa_prefill_backend == "tilelang"
or self.server_args.nsa_decode_backend == "tilelang"
self.server_args.dsa_prefill_backend == "tilelang"
or self.server_args.dsa_decode_backend == "tilelang"
):
return kv_cache_dim
quant_block_size = NSATokenToKVPool.quant_block_size
rope_storage_dtype = NSATokenToKVPool.rope_storage_dtype
quant_block_size = DSATokenToKVPool.quant_block_size
rope_storage_dtype = DSATokenToKVPool.rope_storage_dtype
# Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout (excluding TRTLLM and HIP+TileLang).
# kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage
# Note: rope dimension is stored in original dtype (bf16), not quantized to fp8
@@ -199,7 +199,7 @@ class ModelRunnerKVCacheMixin:
def _validate_prefill_only_disable_kv_cache_pool_family(
self: ModelRunner,
is_nsa_model: bool,
is_dsa_model: bool,
is_dsv4_model: bool,
current_platform,
):
@@ -215,8 +215,8 @@ class ModelRunnerKVCacheMixin:
self.server_args.attention_backend == "ascend" and not self.mambaish_config
):
unsupported_pool_family = "NPU/Ascend KV pool"
elif self.use_mla_backend and is_nsa_model:
unsupported_pool_family = "NSA/MLA KV pool"
elif self.use_mla_backend and is_dsa_model:
unsupported_pool_family = "DSA/MLA KV pool"
elif self.use_mla_backend and not self.mambaish_config:
unsupported_pool_family = "MLA KV pool"
elif self.is_hybrid_swa:
@@ -328,14 +328,14 @@ class ModelRunnerKVCacheMixin:
assert self.is_draft_worker
# Initialize token_to_kv_pool
is_nsa_model = is_deepseek_nsa(self.model_config.hf_config)
is_dsa_model = is_deepseek_dsa(self.model_config.hf_config)
is_dsv4_model = is_deepseek_v4(self.model_config.hf_config)
# Out-of-tree platform plugin system — used by elif below
from sglang.srt.platforms import current_platform
self._validate_prefill_only_disable_kv_cache_pool_family(
is_nsa_model, is_dsv4_model, current_platform
is_dsa_model, is_dsv4_model, current_platform
)
if is_dsv4_model:
@@ -375,8 +375,8 @@ class ModelRunnerKVCacheMixin:
enable_hisparse=self.enable_hisparse,
)
elif current_platform.is_out_of_tree() and not self.mambaish_config:
if self.use_mla_backend and is_nsa_model:
PoolCls = current_platform.get_nsa_kv_pool_cls()
if self.use_mla_backend and is_dsa_model:
PoolCls = current_platform.get_dsa_kv_pool_cls()
self.token_to_kv_pool = PoolCls(
self.max_total_num_tokens,
page_size=self.page_size,
@@ -389,7 +389,7 @@ class ModelRunnerKVCacheMixin:
enable_memory_saver=self.server_args.enable_memory_saver,
start_layer=self.start_layer,
end_layer=self.end_layer,
index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config),
index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config),
)
elif self.use_mla_backend:
PoolCls = current_platform.get_mla_kv_pool_cls()
@@ -400,7 +400,7 @@ class ModelRunnerKVCacheMixin:
kv_lora_rank=self.model_config.kv_lora_rank,
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
index_head_dim=(
self.model_config.index_head_dim if is_nsa_model else None
self.model_config.index_head_dim if is_dsa_model else None
),
layer_num=self.num_effective_layers,
device=self.device,
@@ -472,7 +472,7 @@ class ModelRunnerKVCacheMixin:
kv_lora_rank=self.model_config.kv_lora_rank,
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
index_head_dim=(
self.model_config.index_head_dim if is_nsa_model else None
self.model_config.index_head_dim if is_dsa_model else None
),
layer_num=self.num_effective_layers,
device=self.device,
@@ -499,9 +499,9 @@ class ModelRunnerKVCacheMixin:
start_layer=self.start_layer,
end_layer=self.end_layer,
)
elif self.use_mla_backend and is_nsa_model:
elif self.use_mla_backend and is_dsa_model:
PoolCls = (
HiSparseNSATokenToKVPool if self.enable_hisparse else NSATokenToKVPool
HiSparseDSATokenToKVPool if self.enable_hisparse else DSATokenToKVPool
)
pool_kwargs = {}
if self.enable_hisparse:
@@ -522,11 +522,11 @@ class ModelRunnerKVCacheMixin:
enable_memory_saver=self.server_args.enable_memory_saver,
start_layer=self.start_layer,
end_layer=self.end_layer,
index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config),
index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config),
**pool_kwargs,
)
elif self.use_mla_backend and not self.mambaish_config:
assert not is_nsa_model
assert not is_dsa_model
if is_float4_e2m1fn_x2(self.kv_cache_dtype):
self.token_to_kv_pool = MLATokenToKVPoolFP4(
self.max_total_num_tokens,
python/sglang/srt/model_executor/pool_configurator.py
+9 -9
View File
@@ -20,14 +20,14 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.configs.model_config import (
get_nsa_index_head_dim,
is_deepseek_nsa,
get_dsa_index_head_dim,
is_deepseek_dsa,
is_deepseek_v4,
)
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.utils.common import is_float4_e2m1fn_x2
@@ -84,7 +84,7 @@ class MemoryPoolConfigurator:
class DefaultPoolConfigurator(MemoryPoolConfigurator):
"""Configurator for standard models: MHA, MLA, NSA, FP4.
"""Configurator for standard models: MHA, MLA, DSA, FP4.
coeff = cell_size (bytes per token across all layers)
bias = 0
@@ -149,15 +149,15 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
* kv_size
)
# Add indexer KV cache overhead for NSA models (DeepSeek V3.2)
if is_deepseek_nsa(model_config.hf_config):
index_head_dim = get_nsa_index_head_dim(model_config.hf_config)
# Add indexer KV cache overhead for DSA models (DeepSeek V3.2)
if is_deepseek_dsa(model_config.hf_config):
index_head_dim = get_dsa_index_head_dim(model_config.hf_config)
indexer_size_per_token = (
index_head_dim
+ index_head_dim // NSATokenToKVPool.quant_block_size * 4
+ index_head_dim // DSATokenToKVPool.quant_block_size * 4
)
element_size = torch._utils._element_size(
NSATokenToKVPool.index_k_with_scale_buffer_dtype
DSATokenToKVPool.index_k_with_scale_buffer_dtype
)
cell_size += indexer_size_per_token * num_layers * element_size
else:
python/sglang/srt/models/deepseek_common/attention_backend_handler.py
+6 -3
View File
@@ -147,9 +147,9 @@ def handle_attention_aiter(attn, forward_batch):
return AttnForwardMethod.MLA
def handle_attention_nsa(attn, forward_batch):
def handle_attention_dsa(attn, forward_batch):
"""
Dispatch logic is centralized in NativeSparseAttnBackend.set_nsa_prefill_impl and executed
Dispatch logic is centralized in DeepseekSparseAttnBackend.set_dsa_prefill_impl and executed
in init_forward_metadata. Read the decision from backend.use_mha.
"""
@@ -191,6 +191,9 @@ AttentionBackendRegistry.register("fa4", handle_attention_fa4)
AttentionBackendRegistry.register("trtllm_mla", handle_attention_trtllm_mla)
AttentionBackendRegistry.register("tokenspeed_mla", handle_attention_tokenspeed_mla)
AttentionBackendRegistry.register("aiter", handle_attention_aiter)
AttentionBackendRegistry.register("nsa", handle_attention_nsa)
AttentionBackendRegistry.register("dsa", handle_attention_dsa)
AttentionBackendRegistry.register(
"nsa", handle_attention_dsa
) # Deprecated alias; use "dsa"
AttentionBackendRegistry.register("triton", handle_attention_triton)
AttentionBackendRegistry.register("intel_xpu", handle_attention_intel_xpu)
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+11 -11
View File
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged
from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_paged
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton
from sglang.srt.layers.communicator import get_attn_tp_context
@@ -122,10 +122,10 @@ class DeepseekMHAForwardMixin:
)
)
# NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
# DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
if self.use_nsa:
# NSA requires unquantized q_lora for the indexer. When q_b_proj is FP8
if self.use_dsa:
# DSA requires unquantized q_lora for the indexer. When q_b_proj is FP8
# on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
# the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
# tuple.
@@ -230,15 +230,15 @@ class DeepseekMHAForwardMixin:
and sum(forward_batch.extend_prefix_lens_cpu) != 0
):
if (
self.use_nsa
self.use_dsa
and self.kv_cache_dtype == "fp8_e4m3"
and (
not get_global_server_args().nsa_decode_backend == "trtllm"
or not get_global_server_args().nsa_prefill_backend == "trtllm"
not get_global_server_args().dsa_decode_backend == "trtllm"
or not get_global_server_args().dsa_prefill_backend == "trtllm"
)
):
# FP8 path: dequantize NSA-specific FP8 format to BF16
kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_nsa(forward_batch)
# FP8 path: dequantize DSA-specific FP8 format to BF16
kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch)
else:
# BF16/FP16 path: directly fetch from cache
kv_a, k_pe = self._get_mla_kv_buffer(
@@ -471,12 +471,12 @@ class DeepseekMHAForwardMixin:
kv_a = kv_a.squeeze(1).contiguous()
return kv_a, k_pe
def _get_mla_kv_buffer_from_fp8_for_nsa(
def _get_mla_kv_buffer_from_fp8_for_dsa(
self: DeepseekV2AttentionMLA,
forward_batch: ForwardBatch,
):
"""
Dequantize FP8 KV cache to BF16 for MLA attention (NSA-specific format).
Dequantize FP8 KV cache to BF16 for MLA attention (DSA-specific format).
Returns: (kv_a, k_pe) both in BF16
"""
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+33 -26
View File
@@ -6,7 +6,7 @@ import torch
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.quantization.fp8_kernel import (
fp8_dtype,
@@ -124,7 +124,6 @@ if _use_aiter_gfx95:
class DeepseekMLAForwardMixin:
def init_mla_forward(self: DeepseekV2AttentionMLA):
self.flashinfer_mla_disable_ragged = (
get_global_server_args().flashinfer_mla_disable_ragged
@@ -178,7 +177,7 @@ class DeepseekMLAForwardMixin:
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
if self.use_nsa:
if self.use_dsa:
q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
@@ -220,7 +219,7 @@ class DeepseekMLAForwardMixin:
k_nope = self.kv_a_layernorm(k_nope)
# q_lora needed by indexer
if self.use_nsa:
if self.use_dsa:
if q_lora is None:
q_lora = q
@@ -281,9 +280,13 @@ class DeepseekMLAForwardMixin:
k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
if self.use_deep_gemm_bmm:
q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = (
per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1))
)
(
q_nope_val,
q_nope_scale,
masked_m,
expected_m,
aligned_m,
) = per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1))
q_nope_out = q_nope.new_empty(
(self.num_local_heads, aligned_m, self.kv_lora_rank)
)
@@ -362,18 +365,18 @@ class DeepseekMLAForwardMixin:
if is_kv_b_lora_active(self):
q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out)
skip_rope_for_nsa_tilelang_fused = self._skip_rope_for_nsa_tilelang_fused()
skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused()
skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla()
if (
self.rotary_emb is not None
and (not self._fuse_rope_for_trtllm_mla(forward_batch))
and (not skip_rope_for_nsa_tilelang_fused)
and (not skip_rope_for_dsa_tilelang_fused)
and (not skip_rope_for_aiter_fused_mla)
and (not _use_aiter or not _is_gfx95_supported or self.use_nsa)
and (not _use_aiter or not _is_gfx95_supported or self.use_dsa)
):
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
# support allgather+rerrange
k_nope, k_pe = self.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
@@ -406,7 +409,7 @@ class DeepseekMLAForwardMixin:
save_kv_cache = True
if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
if self._skip_rope_for_nsa_tilelang_fused() and self.rotary_emb is not None:
if self._skip_rope_for_dsa_tilelang_fused() and self.rotary_emb is not None:
cos = self.rotary_emb.cos_cache
sin = self.rotary_emb.sin_cache
kv_cache_dtype = (
@@ -430,14 +433,14 @@ class DeepseekMLAForwardMixin:
)
save_kv_cache = False
# On decode, pass q_cat directly to attn_mqa with q_rope=None so
# nsa_backend.forward_decode reuses q_cat as a zero-copy view
# dsa_backend.forward_decode reuses q_cat as a zero-copy view
# (`q.contiguous().view(...)` fast-path) instead of running the
# redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)`
# that would otherwise rebuild a tensor byte-identical to q_cat.
# On ROCm tilelang decode, this eliminates the
# `CatArrayBatchedCopy<OpaqueType<1u>, ...>` kernel that used to
# fire once per layer per decode step (~2.6 us / layer saved).
# Prefill keeps the split form because nsa_backend.forward_extend
# Prefill keeps the split form because dsa_backend.forward_extend
# asserts `q_rope is not None`.
if forward_batch.forward_mode.is_decode_or_idle():
if llama_4_scaling is not None:
@@ -545,10 +548,14 @@ class DeepseekMLAForwardMixin:
attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
if self.use_deep_gemm_bmm:
attn_output_val, attn_output_scale, masked_m, expected_m, aligned_m = (
per_token_group_quant_mla_deep_gemm_masked_fp8(
attn_output.transpose(0, 1)
)
(
attn_output_val,
attn_output_scale,
masked_m,
expected_m,
aligned_m,
) = per_token_group_quant_mla_deep_gemm_masked_fp8(
attn_output.transpose(0, 1)
)
attn_bmm_output = attn_output.new_empty(
(self.num_local_heads, aligned_m, self.v_head_dim)
@@ -683,10 +690,10 @@ class DeepseekMLAForwardMixin:
"""
Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path.
"""
if self.current_attention_backend == "nsa":
if self.current_attention_backend in ("dsa", "nsa"):
return (
get_global_server_args().nsa_decode_backend == "trtllm"
or get_global_server_args().nsa_prefill_backend == "trtllm"
get_global_server_args().dsa_decode_backend == "trtllm"
or get_global_server_args().dsa_prefill_backend == "trtllm"
) and forward_batch.attn_backend.kv_cache_dtype == torch.float8_e4m3fn
return (
@@ -698,17 +705,17 @@ class DeepseekMLAForwardMixin:
and forward_batch.attn_backend.data_type == torch.float8_e4m3fn
)
def _skip_rope_for_nsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool:
"""
Check if we should skip rope and use fused rope+cache path for TileLang NSA on gfx95.
Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95.
"""
server_args = get_global_server_args()
return (
_use_aiter_gfx95
and self.current_attention_backend == "nsa"
and self.current_attention_backend in ("dsa", "nsa")
and (
server_args.nsa_decode_backend == "tilelang"
or server_args.nsa_prefill_backend == "tilelang"
server_args.dsa_decode_backend == "tilelang"
or server_args.dsa_prefill_backend == "tilelang"
)
)
python/sglang/srt/models/deepseek_common/utils.py
+2 -1
View File
@@ -57,7 +57,8 @@ NVFP4_CKPT_FP8_ATTN_QUANT_MODULES = ["q_b_proj"]
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [
"fa3",
"nsa",
"dsa",
"nsa", # Deprecated alias for "dsa"
"flashinfer",
"cutlass_mla",
"trtllm_mla",
python/sglang/srt/models/deepseek_nextn.py
+15 -15
View File
@@ -23,14 +23,14 @@ from safetensors.torch import load_file
from torch import nn
from transformers import PretrainedConfig
from sglang.srt.configs.model_config import is_deepseek_nsa
from sglang.srt.configs.model_config import is_deepseek_dsa
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.attention.nsa.utils import (
can_nsa_cp_split,
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.dp_attention import (
get_attention_cp_rank,
@@ -148,8 +148,8 @@ class DeepseekModelNextN(nn.Module):
self.shared_head = nn.Module()
self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_size = get_attention_cp_size()
else:
self.cp_size = None
@@ -193,7 +193,7 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.eh_proj(eh_input)
if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
if dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp):
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
residual = None
@@ -212,7 +212,7 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.shared_head.norm(hidden_states)
if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
if dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp):
# allgather + rerrange
hidden_states = cp_all_gather_rerange_output(
hidden_states,
@@ -248,9 +248,9 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
# if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
self.pp_group = get_pp_group()
self.determine_num_fused_shared_experts("DeepseekV3ForCausalLMNextN")
self.use_nsa = is_deepseek_nsa(config)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.use_dsa = is_deepseek_dsa(config)
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_attention_cp_rank()
self.cp_size = get_attention_cp_size()
else:
@@ -289,9 +289,9 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
forward_batch: ForwardBatch,
) -> torch.Tensor:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
if self.nsa_enable_prefill_cp:
if can_nsa_cp_split(
len(input_ids), self.cp_size, self.use_nsa, forward_batch
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(
len(input_ids), self.cp_size, self.use_dsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
python/sglang/srt/models/deepseek_v2.py
+37 -37
View File
@@ -40,10 +40,10 @@ from sglang.srt.batch_overlap.two_batch_overlap import (
)
from sglang.srt.configs.model_config import (
compute_mla_mscale_scaling,
get_nsa_index_head_dim,
get_nsa_index_n_heads,
get_nsa_index_topk,
is_deepseek_nsa,
get_dsa_index_head_dim,
get_dsa_index_n_heads,
get_dsa_index_topk,
is_deepseek_dsa,
)
from sglang.srt.distributed import (
divide,
@@ -59,11 +59,11 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.amx_utils import PackWeightMethod
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
from sglang.srt.layers.attention.nsa.utils import (
can_nsa_cp_split,
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.communicator import (
LayerCommunicator,
@@ -71,7 +71,7 @@ from sglang.srt.layers.communicator import (
enable_moe_dense_fully_dp,
get_attn_tp_context,
)
from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator
from sglang.srt.layers.communicator_dsa_cp import DSACPLayerCommunicator
from sglang.srt.layers.dp_attention import (
get_attention_cp_rank,
get_attention_cp_size,
@@ -368,7 +368,7 @@ class MoEGate(nn.Module):
self.e_score_correction_bias = None
if _is_cpu and _is_cpu_amx_available:
self.quant_method = PackWeightMethod(weight_names=["weight"])
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
def forward(
self,
@@ -390,7 +390,7 @@ class MoEGate(nn.Module):
if (
not self.is_deepseek_v4
and forward_batch is not None
and nsa_use_prefill_cp(forward_batch)
and dsa_use_prefill_cp(forward_batch)
):
logits = F.linear(hidden_states, self.weight, None)
else:
@@ -1352,12 +1352,12 @@ class DeepseekV2AttentionMLA(
self.quant_config = quant_config
attn_tp_rank = get_attention_tp_rank()
attn_tp_size = get_attention_tp_size()
self.use_nsa = is_deepseek_nsa(config)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
assert self.use_nsa, "CP currently only supports deepseek v3.2 model"
self.use_dsa = is_deepseek_dsa(config)
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
assert self.use_dsa, "CP currently only supports deepseek v3.2 model"
# cp reuse the attn_tp comm group but need to duplicate the weights
if self.nsa_enable_prefill_cp and self.use_nsa:
if self.dsa_enable_prefill_cp and self.use_dsa:
self.cp_size = get_attention_cp_size()
self.num_heads = num_heads
assert num_heads % attn_tp_size == 0
@@ -1410,14 +1410,14 @@ class DeepseekV2AttentionMLA(
self.skip_topk = None
self.next_skip_topk = None
if self.use_nsa:
if self.use_dsa:
is_neox_style = not getattr(config, "indexer_rope_interleave", False)
self.indexer = Indexer(
hidden_size=hidden_size,
index_n_heads=get_nsa_index_n_heads(config),
index_head_dim=get_nsa_index_head_dim(config),
index_n_heads=get_dsa_index_n_heads(config),
index_head_dim=get_dsa_index_head_dim(config),
rope_head_dim=qk_rope_head_dim,
index_topk=get_nsa_index_topk(config),
index_topk=get_dsa_index_topk(config),
q_lora_rank=q_lora_rank,
max_position_embeddings=max_position_embeddings,
rope_theta=rope_theta,
@@ -1582,7 +1582,7 @@ class DeepseekV2AttentionMLA(
def op_core(self, state):
result = self.forward_core(state.pop("attn_intermediate_state"))
# forward_core may return (hidden_states, topk_indices) for NSA models
# forward_core may return (hidden_states, topk_indices) for DSA models
# with index cache enabled. In the TBO path, topk_indices is not
# propagated between layers, so we discard it here.
if isinstance(result, tuple):
@@ -1808,7 +1808,7 @@ class DeepseekV2DecoderLayer(nn.Module):
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
get_global_server_args().speculative_algorithm
)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.layer_id = layer_id
self.is_nextn = is_nextn
self.self_attn = DeepseekV2AttentionMLA(
@@ -1881,8 +1881,8 @@ class DeepseekV2DecoderLayer(nn.Module):
self._gfx95_quant_format = self._detect_gfx95_quant_format()
if self.nsa_enable_prefill_cp:
self.layer_communicator = NSACPLayerCommunicator(
if self.dsa_enable_prefill_cp:
self.layer_communicator = DSACPLayerCommunicator(
layer_scatter_modes=self.layer_scatter_modes,
input_layernorm=self.input_layernorm,
post_attention_layernorm=self.post_attention_layernorm,
@@ -1996,7 +1996,7 @@ class DeepseekV2DecoderLayer(nn.Module):
gemm_output_zero_allocator,
)
if not self.nsa_enable_prefill_cp and should_allreduce_fusion:
if not self.dsa_enable_prefill_cp and should_allreduce_fusion:
hidden_states._sglang_needs_allreduce_fusion = True
if not should_allreduce_fusion:
@@ -2093,8 +2093,8 @@ class DeepseekV2Model(nn.Module):
self.vocab_size = config.vocab_size
self.first_k_dense_replace = config.first_k_dense_replace
self.pp_group = get_pp_group()
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_size = get_attention_cp_size()
else:
self.cp_size = None
@@ -2253,7 +2253,7 @@ class DeepseekV2Model(nn.Module):
else None
)
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
if self.pp_group.is_first_rank:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
@@ -2338,7 +2338,7 @@ class DeepseekV2Model(nn.Module):
else:
hidden_states, _ = self.norm(hidden_states, residual)
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
if self.pp_group.is_last_rank and dsa_use_prefill_cp(forward_batch):
# allgather + rerrange
hidden_states = cp_all_gather_rerange_output(
hidden_states,
@@ -2385,7 +2385,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.tp_size = get_tensor_model_parallel_world_size()
self.quant_config = quant_config
self.determine_num_fused_shared_experts()
self.use_nsa = is_deepseek_nsa(config)
self.use_dsa = is_deepseek_dsa(config)
self.model = DeepseekV2Model(
config, quant_config, prefix=add_prefix("model", prefix)
)
@@ -2415,15 +2415,15 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
)
self.capture_aux_hidden_states = False
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_attention_cp_rank()
self.cp_size = get_attention_cp_size()
else:
self.cp_rank = self.cp_size = None
q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None
get_attn_tp_context().init_context(q_lora_rank, is_deepseek_nsa(config))
get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config))
@property
def routed_experts_weights_of_layer(self):
@@ -2498,9 +2498,9 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
input_embeds: torch.Tensor = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
if self.nsa_enable_prefill_cp:
if can_nsa_cp_split(
len(input_ids), self.cp_size, self.use_nsa, forward_batch
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(
len(input_ids), self.cp_size, self.use_dsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
python/sglang/srt/models/deepseek_v4.py
+26 -26
View File
@@ -33,14 +33,14 @@ from sglang.srt.distributed import (
)
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.attention.dsv4.compressor import Compressor
from sglang.srt.layers.attention.dsv4.indexer import C4Indexer
from sglang.srt.layers.attention.nsa.utils import (
can_nsa_cp_split,
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_round_robin_split,
nsa_use_prefill_cp,
)
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.dp_attention import (
_DpGatheredBufferWrapper,
@@ -177,8 +177,8 @@ class MQALayer(nn.Module):
super().__init__()
self.tp_rank = attn_tp_rank = get_attention_tp_rank()
self.tp_size = attn_tp_size = get_attention_tp_size()
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_size = get_attention_cp_size()
self.tp_rank = attn_tp_rank = 0
self.tp_size = attn_tp_size = 1
@@ -390,7 +390,7 @@ class MQALayer(nn.Module):
) -> None:
"""Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache.
Replaces the bf16-kv-intermediate path. Used everywhere except the NSA
Replaces the bf16-kv-intermediate path. Used everywhere except the DSA
prefill-CP case (which needs bf16 kv for the cross-rank all-gather).
"""
if qkv_a is not None:
@@ -416,7 +416,7 @@ class MQALayer(nn.Module):
positions: torch.Tensor,
qkv_a: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Bf16-kv path used by the NSA prefill-CP case (needs all-gather)."""
"""Bf16-kv path used by the DSA prefill-CP case (needs all-gather)."""
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
else:
@@ -508,10 +508,10 @@ class MQALayer(nn.Module):
q_lora = self.q_norm(q_lora)
q = self._compute_q_b(q_lora, positions, q_out)
use_cp = self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch)
use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
kv: Optional[torch.Tensor]
if use_cp:
# NSA CP: keep bf16 kv around for the cross-rank all-gather, then
# DSA CP: keep bf16 kv around for the cross-rank all-gather, then
# write to the FlashMLA cache after gather.
kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a)
kv = cp_all_gather_rerange_output(
@@ -567,7 +567,7 @@ class MQALayer(nn.Module):
and self.alt_streams is not None
and get_is_capture_mode()
and x.shape[0] <= self._multi_stream_bs_limit
and not (self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch))
and not (self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch))
)
tp_slice, q_padded, q_out = slice(None), None, None
@@ -591,7 +591,7 @@ class MQALayer(nn.Module):
# The cache write is always fused / already done by _forward_prepare* --
# tell the backend to skip its own store_cache. When `kv is None`
# (no NSA-CP), pass `q` as a sentinel for the `k is v` assert; the
# (no DSA-CP), pass `q` as a sentinel for the `k is v` assert; the
# attention path doesn't read it once `save_kv_cache=False`.
attn_k = kv if kv is not None else q
o = attn_backend.forward(
@@ -694,7 +694,7 @@ class DeepseekV4DecoderLayer(nn.Module):
self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
self.rms_norm_eps = config.rms_norm_eps
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
def hc_pre(
self,
@@ -869,7 +869,7 @@ class DeepseekV4DecoderLayer(nn.Module):
if not norm_fused:
hidden_states = self.post_attention_layernorm(hidden_states)
_use_cp = self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch)
_use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
_use_tp_moe_gather = (
not _use_cp
and get_attention_dp_size() > 1
@@ -979,8 +979,8 @@ class DeepseekV4Model(nn.Module):
self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32))
self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32))
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_size = get_attention_cp_size()
def hc_head(
@@ -1040,7 +1040,7 @@ class DeepseekV4Model(nn.Module):
else:
input_ids_global = input_ids
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
if self.pp_group.is_first_rank:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
@@ -1060,7 +1060,7 @@ class DeepseekV4Model(nn.Module):
)
# CP all-gather only on the last PP rank; PP IPC carries CP-split tensors.
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
if self.pp_group.is_last_rank and dsa_use_prefill_cp(forward_batch):
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
@@ -1113,7 +1113,7 @@ class DeepseekV4ForCausalLM(nn.Module):
self.lm_head = PPMissingLayer()
self.logits_processor = LogitsProcessor(config)
self.capture_aux_hidden_states = False
get_attn_tp_context().init_context(config.q_lora_rank, is_nsa=True)
get_attn_tp_context().init_context(config.q_lora_rank, is_dsa=True)
self._routed_experts_weights_of_layer = LazyValue(
lambda: {
@@ -1129,8 +1129,8 @@ class DeepseekV4ForCausalLM(nn.Module):
self.start_layer = self.model.start_layer
self.end_layer = self.model.end_layer
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_attention_cp_rank()
self.cp_size = get_attention_cp_size()
@@ -1159,15 +1159,15 @@ class DeepseekV4ForCausalLM(nn.Module):
input_embeds: Optional[torch.Tensor] = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
if self.nsa_enable_prefill_cp:
if can_nsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
)
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
metadata = forward_batch.attn_backend.forward_metadata
core_meta = metadata.core_attn_metadata
core_meta.apply_cp_reindex()
python/sglang/srt/models/deepseek_v4_nextn.py
+14 -14
View File
@@ -7,11 +7,11 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
from sglang.srt.layers.attention.nsa.utils import (
can_nsa_cp_split,
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_round_robin_split,
nsa_use_prefill_cp,
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.dp_attention import (
_DpGatheredBufferWrapper,
@@ -104,8 +104,8 @@ class DeepseekV4ModelNextN(nn.Module):
compress_ratio_override=COMPRESS_RATIO_NEXTN_LAYER,
)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_size = get_attention_cp_size()
else:
self.cp_size = None
@@ -165,7 +165,7 @@ class DeepseekV4ModelNextN(nn.Module):
else:
input_ids_global = input_ids
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
@@ -177,7 +177,7 @@ class DeepseekV4ModelNextN(nn.Module):
input_ids_global=input_ids_global,
)
if nsa_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch):
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
@@ -209,8 +209,8 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
self.pp_group = get_pp_group()
self.quant_config = quant_config
self.determine_num_fused_shared_experts()
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_attention_cp_rank()
self.cp_size = get_attention_cp_size()
else:
@@ -236,15 +236,15 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
positions: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
if self.nsa_enable_prefill_cp:
if can_nsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
)
if is_nsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_round_robin_split():
metadata = forward_batch.attn_backend.forward_metadata
core_meta = metadata.core_attn_metadata
core_meta.apply_cp_reindex()
python/sglang/srt/models/glm4_moe_lite.py
+6 -6
View File
@@ -30,7 +30,7 @@ from sglang.srt.distributed import (
get_tensor_model_parallel_world_size,
)
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.communicator import (
LayerCommunicator,
LayerScatterModes,
@@ -341,7 +341,7 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer):
nn.Module.__init__(self)
self.hidden_size = config.hidden_size
self.config = config
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
rope_theta, rope_scaling = get_rope_config(config)
max_position_embeddings = getattr(config, "max_position_embeddings", 202752)
self.layer_id = layer_id
@@ -433,8 +433,8 @@ class Glm4MoeLiteModel(DeepseekV2Model):
self.pp_group = get_pp_group()
# DeepseekV2Model.forward expects these attributes to exist.
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.cp_size = get_attention_tp_size() if self.nsa_enable_prefill_cp else None
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.cp_size = get_attention_tp_size() if self.dsa_enable_prefill_cp else None
self.gemm_output_zero_allocator_size = 0
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
@@ -503,8 +503,8 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
)
self.capture_aux_hidden_states = False
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_attention_tp_rank()
self.cp_size = get_attention_tp_size()
else:
python/sglang/srt/models/mistral_large_3_eagle.py
+2 -2
View File
@@ -8,7 +8,7 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import RowParallelLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -35,7 +35,7 @@ class MistralLarge3EagleModel(DeepseekV2Model):
self.vocab_size = config.vocab_size
assert get_pp_group().world_size == 1
self.pp_group = get_pp_group()
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
python/sglang/srt/models/sarvam_moe.py
+2 -6
View File
@@ -104,7 +104,8 @@ class AttnForwardMethod(IntEnum):
SEPARATE_ROPE_BACKENDS = frozenset(
["fa3", "flashinfer", "nsa", "cutlass_mla", "trtllm_mla"]
["fa3", "flashinfer", "dsa", "nsa", "cutlass_mla", "trtllm_mla"]
# "nsa" is a deprecated alias for "dsa"
)
CONCAT_ROPE_BACKENDS = frozenset(["flashmla", "triton"])
@@ -667,7 +668,6 @@ class SarvamMoEMLAAttention(nn.Module):
k_pe: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
q[..., self.qk_nope_head_dim :] = q_pe
@@ -989,7 +989,6 @@ class SarvamMoEMLAAttention(nn.Module):
class SarvamMoEMLADecoderLayer(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -1139,7 +1138,6 @@ class SarvamMoEMLADecoderLayer(nn.Module):
class SarvamMLAModel(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -1223,7 +1221,6 @@ class SarvamMLAModel(nn.Module):
class SarvamMLAForCausalLM(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -1475,7 +1472,6 @@ class SarvamMLAForCausalLM(nn.Module):
class SarvamMoEForCausalLM(BailingMoEForCausalLM):
@torch.no_grad()
def forward_split_prefill(
self,
python/sglang/srt/platforms/interface.py
+2 -2
View File
@@ -61,8 +61,8 @@ class SRTPlatform(DeviceMixin):
"""Return the MLA KV pool class for this platform."""
raise NotImplementedError
def get_nsa_kv_pool_cls(self) -> type:
"""Return the NSA KV pool class for this platform (DeepSeek V3.2)."""
def get_dsa_kv_pool_cls(self) -> type:
"""Return the DSA KV pool class for this platform (DeepSeek V3.2)."""
raise NotImplementedError
def get_paged_allocator_cls(self) -> type:
python/sglang/srt/server_args.py
+106 -62
View File
@@ -155,7 +155,8 @@ ATTENTION_BACKEND_CHOICES = [
"triton",
"torch_native",
"flex_attention",
"nsa",
"dsa",
"nsa", # Deprecated alias for "dsa"
"dsv4",
"compressed", # Deprecated alias for "dsv4"
# NVIDIA specific
@@ -246,13 +247,14 @@ LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"]
ENCODER_TRANSFER_BACKEND_CHOICES = ["zmq_to_scheduler", "zmq_to_tokenizer", "mooncake"]
NSA_PREFILL_CP_SPLIT_CHOICES = ["in-seq-split", "round-robin-split"]
DSA_PREFILL_CP_SPLIT_CHOICES = ["in-seq-split", "round-robin-split"]
NSA_PREFILL_CP_SPLIT_CHOICES = DSA_PREFILL_CP_SPLIT_CHOICES # deprecated alias
PREFILL_CP_SPLIT_CHOICES = ["in-seq-split"]
DEFAULT_LORA_EVICTION_POLICY = "lru"
NSA_CHOICES = [
DSA_CHOICES = [
"flashmla_sparse",
"flashmla_kv",
"flashmla_auto",
@@ -261,6 +263,7 @@ NSA_CHOICES = [
"aiter",
"trtllm",
]
NSA_CHOICES = DSA_CHOICES # deprecated alias
MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"]
@@ -535,10 +538,10 @@ class ServerArgs:
mm_attention_backend: Optional[str] = None
fp8_gemm_runner_backend: str = "auto"
fp4_gemm_runner_backend: str = "auto"
nsa_prefill_backend: Optional[str] = (
dsa_prefill_backend: Optional[str] = (
None # None = auto-detect based on hardware/kv_cache_dtype
)
nsa_decode_backend: Optional[str] = (
dsa_decode_backend: Optional[str] = (
None # auto-detect based on hardware/kv_cache_dtype
)
disable_flashinfer_autotune: bool = False
@@ -741,8 +744,8 @@ class ServerArgs:
enable_attn_tp_input_scattered: bool = False
gc_threshold: Optional[List[int]] = None
# Context parallelism used in the long sequence prefill phase of DeepSeek v3.2
enable_nsa_prefill_context_parallel: bool = False
nsa_prefill_cp_mode: str = "round-robin-split"
enable_dsa_prefill_context_parallel: bool = False
dsa_prefill_cp_mode: str = "round-robin-split"
enable_fused_qk_norm_rope: bool = False
enable_precise_embedding_interpolation: bool = False
enable_fused_moe_sum_all_reduce: bool = False
@@ -1647,15 +1650,15 @@ class ServerArgs:
return capture_sizes
def _set_default_nsa_kv_cache_dtype(self, major: int, quantization: str) -> str:
user_set_prefill = self.nsa_prefill_backend is not None
user_set_decode = self.nsa_decode_backend is not None
def _set_default_dsa_kv_cache_dtype(self, major: int, quantization: str) -> str:
user_set_prefill = self.dsa_prefill_backend is not None
user_set_decode = self.dsa_decode_backend is not None
# If user specified a backend but didn't explicitly set kv_cache_dtype,
# suggest them to be explicit about kv_cache_dtype to avoid surprises
if (user_set_prefill or user_set_decode) and self.kv_cache_dtype == "auto":
logger.warning(
"When specifying --nsa-prefill-backend or --nsa-decode-backend, "
"When specifying --dsa-prefill-backend or --dsa-decode-backend, "
"you should also explicitly set --kv-cache-dtype (e.g., 'fp8_e4m3' or 'bfloat16'). "
"DeepSeek V3.2 defaults to FP8 KV cache which may not be compatible with all backends."
)
@@ -1675,56 +1678,56 @@ class ServerArgs:
"fp8_e4m3",
], "DeepSeek DSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype"
def _set_default_nsa_backends(self, kv_cache_dtype: str, major: int) -> str:
def _set_default_dsa_backends(self, kv_cache_dtype: str, major: int) -> str:
from sglang.srt.arg_groups.hisparse_hook import (
apply_hisparse_nsa_backend_defaults,
apply_hisparse_dsa_backend_defaults,
)
user_set_prefill = self.nsa_prefill_backend is not None
user_set_decode = self.nsa_decode_backend is not None
user_set_prefill = self.dsa_prefill_backend is not None
user_set_decode = self.dsa_decode_backend is not None
if apply_hisparse_nsa_backend_defaults(
if apply_hisparse_dsa_backend_defaults(
self, user_set_prefill, user_set_decode, kv_cache_dtype
):
return
if not user_set_prefill and not user_set_decode and is_hip():
self.nsa_prefill_backend = "tilelang"
self.nsa_decode_backend = "tilelang"
self.dsa_prefill_backend = "tilelang"
self.dsa_decode_backend = "tilelang"
elif kv_cache_dtype == "fp8_e4m3":
if major >= 10:
if not user_set_prefill:
self.nsa_prefill_backend = "trtllm"
self.dsa_prefill_backend = "trtllm"
if not user_set_decode:
self.nsa_decode_backend = "trtllm"
self.dsa_decode_backend = "trtllm"
else:
# Hopper FP8 defaults to flashmla_kv for both prefill and decode.
if not user_set_prefill:
self.nsa_prefill_backend = "flashmla_kv"
self.dsa_prefill_backend = "flashmla_kv"
if not user_set_decode:
self.nsa_decode_backend = "flashmla_kv"
self.dsa_decode_backend = "flashmla_kv"
else:
# set prefill/decode backends based on hardware architecture.
if major >= 10:
if not user_set_prefill:
self.nsa_prefill_backend = "flashmla_sparse"
self.dsa_prefill_backend = "flashmla_sparse"
if not user_set_decode:
self.nsa_decode_backend = "trtllm"
self.dsa_decode_backend = "trtllm"
else:
# Hopper defaults for bfloat16
if not user_set_prefill:
self.nsa_prefill_backend = "flashmla_sparse"
self.dsa_prefill_backend = "flashmla_sparse"
if not user_set_decode:
self.nsa_decode_backend = "fa3"
self.dsa_decode_backend = "fa3"
logger.warning(
f"Set NSA backends for {self.kv_cache_dtype} KV Cache: prefill={self.nsa_prefill_backend}, decode={self.nsa_decode_backend}."
f"Set DSA backends for {self.kv_cache_dtype} KV Cache: prefill={self.dsa_prefill_backend}, decode={self.dsa_decode_backend}."
)
def _handle_model_specific_adjustments(self):
from sglang.srt.configs.model_config import (
get_mimo_v2_fused_qkv_expected_tp_size,
is_deepseek_nsa,
is_deepseek_dsa,
)
if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
@@ -1765,37 +1768,37 @@ class ServerArgs:
"GlmMoeDsaForCausalLM",
]:
# Set attention backend for DeepSeek
if is_deepseek_nsa(hf_config): # DeepSeek 3.2/GLM 5
if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5
if model_arch == "GlmMoeDsaForCausalLM" and is_blackwell_supported():
envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0)
envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0)
logger.warning(
"Force NSA prefill to use sparse MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell."
"Force DSA prefill to use sparse MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell."
)
else:
if envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set():
if envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set():
logger.warning(
f"Dense attention kv len threshold is manually set to {envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model."
f"Dense attention kv len threshold is manually set to {envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model."
)
else:
# When threshold is not manually set, set it to the index topk of model
from sglang.srt.configs.model_config import get_nsa_index_topk
from sglang.srt.configs.model_config import get_dsa_index_topk
envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(
get_nsa_index_topk(hf_config)
envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(
get_dsa_index_topk(hf_config)
)
logger.warning(
f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA."
f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA."
)
if self.is_attention_backend_not_set():
self.attention_backend = "nsa"
logger.info("Use nsa attention backend for DeepSeek with DSA.")
self.attention_backend = "dsa"
logger.info("Use dsa attention backend for DeepSeek with DSA.")
if not is_npu() and not is_xpu(): # CUDA or ROCm GPU
if self.enable_nsa_prefill_context_parallel:
if self.enable_dsa_prefill_context_parallel:
logger.warning(
"Context parallel feature is still under experiment. It has only been verified on Hopper platform."
)
if self.nsa_prefill_cp_mode == "in-seq-split":
if self.dsa_prefill_cp_mode == "in-seq-split":
# TODO Supports moe_dense_tp_size != 1, kv cache dtype = "fp8",moe_a2a_backend non-deepep and cross-machine operation .
self.enable_dp_attention = True
self.moe_dense_tp_size = 1
@@ -1819,7 +1822,7 @@ class ServerArgs:
f"Enable Context Parallel opt for deeeseekv3.2-DSA, Setting dp_size == {self.dp_size} and moe_dense_tp_size == {self.moe_dense_tp_size}, ep_size == {self.ep_size}, tp_size == {self.tp_size}, kv_cache_dtype == {self.kv_cache_dtype}, moe_a2a_backend {self.moe_a2a_backend} "
)
else:
# Pure TP and partial DP Attention mode is active for NSA, logging a warning
# Pure TP and partial DP Attention mode is active for DSA, logging a warning
if self.dp_size < self.tp_size:
logger.warning(
f"DSA with TP mode is active, dp_size={self.dp_size}, tp_size={self.tp_size}, "
@@ -1827,15 +1830,15 @@ class ServerArgs:
)
# Deferred import to avoid a circular import at module-load
# time (nsa.utils imports get_global_server_args).
from sglang.srt.layers.attention.nsa.utils import (
# time (dsa.utils imports get_global_server_args).
from sglang.srt.layers.attention.dsa.utils import (
aiter_can_use_preshuffle_paged_mqa,
)
if is_hip() and not aiter_can_use_preshuffle_paged_mqa():
# Legacy ROCm NSA path: aiter's gluon paged-MQA kernel is
# Legacy ROCm DSA path: aiter's gluon paged-MQA kernel is
# unavailable (Triton<3.5 and AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS
# not set, or SGLANG_NSA_HIP_DISABLE_PRESHUFFLE=1 / SGLANG_USE_AITER=0).
# not set, or SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1 / SGLANG_USE_AITER=0).
self.page_size = 1
logger.warning(
"Setting page size to 1 for DeepSeek DSA on ROCm "
@@ -1849,13 +1852,13 @@ class ServerArgs:
import torch
major, _ = torch.cuda.get_device_capability()
self._set_default_nsa_kv_cache_dtype(major, self.quantization)
self._set_default_nsa_backends(self.kv_cache_dtype, major)
self._set_default_dsa_kv_cache_dtype(major, self.quantization)
self._set_default_dsa_backends(self.kv_cache_dtype, major)
if self.enable_nsa_prefill_context_parallel:
if self.enable_dsa_prefill_context_parallel:
assert (
self.disaggregation_mode != "decode"
), "CP is only supported for prefill when PD disaggregation, please remove --enable-nsa-prefill-context-parallel."
), "CP is only supported for prefill when PD disaggregation, please remove --enable-dsa-prefill-context-parallel."
else:
# DeepSeek V3/R1/V3.1
@@ -3408,7 +3411,7 @@ class ServerArgs:
"the paged cache, which the no-op pool does not support."
)
# HiSparse selects a different pool class (HiSparseNSATokenToKVPool /
# HiSparse selects a different pool class (HiSparseDSATokenToKVPool /
# HiSparseTokenToKVPoolAllocator) that is not the no-op pool.
if self.enable_hisparse:
raise ValueError(
@@ -5349,18 +5352,40 @@ class ServerArgs:
help="Set multimodal attention backend.",
)
parser.add_argument(
"--nsa-prefill-backend",
default=ServerArgs.nsa_prefill_backend,
"--dsa-prefill-backend",
dest="dsa_prefill_backend",
default=ServerArgs.dsa_prefill_backend,
type=str,
choices=NSA_CHOICES,
help="NSA prefill backend. If not specified, auto-detects based on hardware and kv_cache_dtype.",
choices=DSA_CHOICES,
help="DSA (DeepSeek Sparse Attention) prefill backend. If not specified, auto-detects based on hardware and kv_cache_dtype.",
)
parser.add_argument(
"--nsa-prefill-backend",
dest="dsa_prefill_backend",
action=DeprecatedAliasStoreAction,
new_flag="--dsa-prefill-backend",
default=argparse.SUPPRESS,
type=str,
choices=DSA_CHOICES,
help="[Deprecated] Use --dsa-prefill-backend instead.",
)
parser.add_argument(
"--dsa-decode-backend",
dest="dsa_decode_backend",
default=ServerArgs.dsa_decode_backend,
type=str,
choices=DSA_CHOICES,
help="DSA (DeepSeek Sparse Attention) decode backend. If not specified, auto-detects based on hardware and kv_cache_dtype.",
)
parser.add_argument(
"--nsa-decode-backend",
default=ServerArgs.nsa_decode_backend,
dest="dsa_decode_backend",
action=DeprecatedAliasStoreAction,
new_flag="--dsa-decode-backend",
default=argparse.SUPPRESS,
type=str,
choices=NSA_CHOICES,
help="NSA decode backend. If not specified, auto-detects based on hardware and kv_cache_dtype.",
choices=DSA_CHOICES,
help="[Deprecated] Use --dsa-decode-backend instead.",
)
parser.add_argument(
"--fp8-gemm-backend",
@@ -6396,15 +6421,34 @@ class ServerArgs:
help="Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent.",
)
parser.add_argument(
"--enable-nsa-prefill-context-parallel",
"--enable-dsa-prefill-context-parallel",
dest="enable_dsa_prefill_context_parallel",
action="store_true",
help="Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2.",
)
parser.add_argument(
"--nsa-prefill-cp-mode",
"--enable-nsa-prefill-context-parallel",
dest="enable_dsa_prefill_context_parallel",
action=DeprecatedStoreTrueAction,
new_flag="--enable-dsa-prefill-context-parallel",
help="[Deprecated] Use --enable-dsa-prefill-context-parallel instead.",
)
parser.add_argument(
"--dsa-prefill-cp-mode",
dest="dsa_prefill_cp_mode",
type=str,
default=ServerArgs.nsa_prefill_cp_mode,
choices=NSA_PREFILL_CP_SPLIT_CHOICES,
default=ServerArgs.dsa_prefill_cp_mode,
choices=DSA_PREFILL_CP_SPLIT_CHOICES,
help="Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism.",
)
parser.add_argument(
"--nsa-prefill-cp-mode",
dest="dsa_prefill_cp_mode",
action=DeprecatedAliasStoreAction,
new_flag="--dsa-prefill-cp-mode",
default=argparse.SUPPRESS,
type=str,
choices=DSA_PREFILL_CP_SPLIT_CHOICES,
help="Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: 'round-robin-split'(default), 'in-seq-split' "
"'round-robin-split' distributes tokens across ranks based on token_idx %% cp_size. It supports multi-batch prefill, fused MoE, and FP8 KV cache.",
)
python/sglang/srt/speculative/draft_utils.py
+11 -9
View File
@@ -54,7 +54,8 @@ class DraftBackendFactory:
"trtllm_mha": self._create_trtllm_mha_decode_backend,
"trtllm_mla": self._create_trtllm_mla_decode_backend,
"tokenspeed_mla": self._create_tokenspeed_mla_decode_backend,
"nsa": self._create_nsa_decode_backend,
"dsa": self._create_dsa_decode_backend,
"nsa": self._create_dsa_decode_backend, # Deprecated alias for "dsa"
"ascend": self._create_ascend_decode_backend,
"fa4": self._create_fa4_decode_backend,
"dsv4": self._create_dsv4_decode_backend,
@@ -81,7 +82,8 @@ class DraftBackendFactory:
"trtllm_mha": self._create_trtllm_mha_prefill_backend,
"trtllm_mla": self._create_trtllm_mla_prefill_backend,
"tokenspeed_mla": self._create_tokenspeed_mla_prefill_backend,
"nsa": self._create_nsa_prefill_backend,
"dsa": self._create_dsa_prefill_backend,
"nsa": self._create_dsa_prefill_backend, # Deprecated alias for "dsa"
"ascend": self._create_ascend_prefill_backend,
"fa4": self._create_fa4_prefill_backend,
"dsv4": self._create_dsv4_prefill_backend,
@@ -97,19 +99,19 @@ class DraftBackendFactory:
"EAGLE is not supported in attention backend {backend_type}",
)
def _create_nsa_decode_backend(self):
from sglang.srt.layers.attention.nsa_backend import (
NativeSparseAttnMultiStepBackend,
def _create_dsa_decode_backend(self):
from sglang.srt.layers.attention.dsa_backend import (
DeepseekSparseAttnMultiStepBackend,
)
return NativeSparseAttnMultiStepBackend(
return DeepseekSparseAttnMultiStepBackend(
self.draft_model_runner, self.topk, self.speculative_num_steps
)
def _create_nsa_prefill_backend(self):
from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend
def _create_dsa_prefill_backend(self):
from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
return NativeSparseAttnBackend(self.draft_model_runner, skip_prefill=False)
return DeepseekSparseAttnBackend(self.draft_model_runner, skip_prefill=False)
def _create_flashinfer_decode_backend(self):
if not get_global_server_args().use_mla_backend:
python/sglang/test/nightly_utils.py
+1 -1
View File
@@ -64,7 +64,7 @@ class NightlyBenchmarkRunner:
Args:
model_path: Path to the model (e.g., "deepseek-ai/DeepSeek-V3.1")
variant: Optional variant suffix (e.g., "basic", "mtp", "nsa")
variant: Optional variant suffix (e.g., "basic", "mtp", "dsa")
Returns:
Tuple of (profile_path_prefix, json_output_file)