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

fix mqa preshuffle layout issue for deepseek v4 (#31563)

Browse Source
This commit is contained in:
Wu Jiangming
2026-07-29 07:36:48 -07:00
committed by GitHub
parent 22151edca1
commit 1c6a0e91e1
6 changed files with 92 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+23 -4
View File
@@ -64,7 +64,7 @@ enum class ForwardMode : bool {
// Each warp's 32 lanes cover the full 128-elem head_dim (kVecSize = 4 each).
// Cache layout: 132 bytes/token (128 fp8 nope + 4 fp32 scale).
// ----------------------------------------------------------------------------
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, int32_t kPreshuffleSize = 0>
INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRopeStoreParams params) {
using namespace device;
using enum ForwardMode;
@@ -213,7 +213,19 @@ INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRop
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
PDLTriggerSecondary<kUsePDL>();
result.store(value_ptr, lane_id);
if constexpr (kPreshuffleSize != 0) {
constexpr int32_t kTile = kPreshuffleSize;
const int32_t dim_base = lane_id * kVecSize;
const int32_t token_tile_id = offset / kTile;
const int32_t token_in_tile = offset % kTile;
const int32_t col_tile_id = dim_base / kTile;
const int32_t col_in_tile = dim_base % kTile;
const int32_t value_offset = token_tile_id * (kTile * static_cast<int32_t>(kHeadDim)) +
col_tile_id * (kTile * kTile) + token_in_tile * kTile + col_in_tile;
result.store(page_ptr + value_offset, 0);
} else {
result.store(value_ptr, lane_id);
}
// The single fp32 scale is identical across all lanes -- write from any lane.
if (lane_id == 0) reinterpret_cast<float*>(scale_ptr)[0] = scale;
}
@@ -499,7 +511,14 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
}
}
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL, bool kBf16Store = false>
template <
typename DType,
int64_t kHeadDim,
int64_t kRopeDim,
uint32_t kPageSize,
bool kUsePDL,
int32_t kPreshuffleSize = 0,
bool kBf16Store = false>
struct FusedNormRopeKernel {
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
static constexpr bool kIsIndexer = (kHeadDim == 128);
@@ -516,7 +535,7 @@ struct FusedNormRopeKernel {
template <ForwardMode kMode>
static constexpr auto select_kernel() {
if constexpr (kIsIndexer) {
return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL>;
return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL, kPreshuffleSize>;
} else {
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
}
python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py
+26 -6
View File
@@ -5,7 +5,10 @@ import triton
import triton.language as tl
from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.attention.dsa.utils import (
INDEXER_K_CACHE_PRESHUFFLE_TILE,
aiter_can_use_preshuffle_paged_mqa,
)
from sglang.srt.utils import get_bool_env_var, is_hip
_is_hip = is_hip()
@@ -341,6 +344,7 @@ def _set_k_and_s_triton(
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,
PRESHUFFLE_TILE=INDEXER_K_CACHE_PRESHUFFLE_TILE if _use_aiter_preshuffle else 0,
)
@@ -356,6 +360,7 @@ def _set_k_and_s_triton_kernel(
BUF_NUMEL_PER_PAGE: tl.constexpr,
NUM_K_ELEMS_PER_TOKEN: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
PRESHUFFLE_TILE: tl.constexpr,
):
token_id = tl.program_id(0)
@@ -370,11 +375,26 @@ def _set_k_and_s_triton_kernel(
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)
)
k_range = tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
if PRESHUFFLE_TILE:
tile = PRESHUFFLE_TILE
token_tile_id = loc_token_offset_in_page // tile
token_in_tile = loc_token_offset_in_page % tile
col_tile_id = k_range // tile
col_in_tile = k_range % tile
out_k_offsets = (
loc_page_index * BUF_NUMEL_PER_PAGE
+ token_tile_id * (tile * NUM_K_ELEMS_PER_TOKEN)
+ col_tile_id * (tile * tile)
+ token_in_tile * tile
+ col_in_tile
)
else:
out_k_offsets = (
loc_page_index * BUF_NUMEL_PER_PAGE
+ loc_token_offset_in_page * NUM_K_ELEMS_PER_TOKEN
+ k_range
)
# "//4" b/c it is fp32 instead of uint8
out_s_offset = (
python/sglang/kernels/ops/attention/dsv4/compress.py
+11 -1
View File
@@ -10,6 +10,10 @@ from sglang.kernels.jit.utils import (
load_jit,
make_cpp_args,
)
from sglang.srt.layers.attention.dsa.utils import (
INDEXER_K_CACHE_PRESHUFFLE_TILE,
aiter_can_use_preshuffle_paged_mqa,
)
from sglang.srt.utils import is_hip, is_xpu
from .utils import make_name
@@ -47,7 +51,13 @@ def _jit_compress_norm_rope_module(
bf16_store: bool = False,
) -> Module:
args = make_cpp_args(
dtype, head_dim, rope_dim, page_size, is_arch_support_pdl(), bf16_store
dtype,
head_dim,
rope_dim,
page_size,
is_arch_support_pdl(),
INDEXER_K_CACHE_PRESHUFFLE_TILE if aiter_can_use_preshuffle_paged_mqa() else 0,
bf16_store,
)
cuda_wrappers = [("forward", f"FusedNormRopeKernel<{args}>::forward")]
if head_dim == 128:
python/sglang/kernels/ops/kvcache/triton_store_cache.py
+24 -1
View File
@@ -5,6 +5,10 @@ import triton
import triton.language as tl
from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.attention.dsa.utils import (
INDEXER_K_CACHE_PRESHUFFLE_TILE,
aiter_can_use_preshuffle_paged_mqa,
)
_FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn
_FP8_INFO = torch.finfo(_FP8_DTYPE)
@@ -156,6 +160,7 @@ def _triton_fused_store_indexer_kernel(
BYTES_PER_PAGE_F32: tl.constexpr,
SCALE_PAGE_OFFSET_F32: tl.constexpr,
HEAD_DIM: tl.constexpr,
PRESHUFFLE_TILE: tl.constexpr,
FP8_MIN: tl.constexpr,
FP8_MAX: tl.constexpr,
EPS: tl.constexpr,
@@ -179,7 +184,20 @@ def _triton_fused_store_indexer_kernel(
cache_fp8_ptr.dtype.element_ty
)
fp8_offset = page * BYTES_PER_PAGE + slot * HEAD_DIM + lane
if PRESHUFFLE_TILE:
token_tile_id = slot // PRESHUFFLE_TILE
token_in_tile = slot % PRESHUFFLE_TILE
col_tile_id = lane // PRESHUFFLE_TILE
col_in_tile = lane % PRESHUFFLE_TILE
fp8_offset = (
page * BYTES_PER_PAGE
+ token_tile_id * (PRESHUFFLE_TILE * HEAD_DIM)
+ col_tile_id * (PRESHUFFLE_TILE * PRESHUFFLE_TILE)
+ token_in_tile * PRESHUFFLE_TILE
+ col_in_tile
)
else:
fp8_offset = page * BYTES_PER_PAGE + slot * HEAD_DIM + lane
tl.store(cache_fp8_ptr + fp8_offset, x_fp8)
f32_offset = page * BYTES_PER_PAGE_F32 + SCALE_PAGE_OFFSET_F32 + slot
@@ -216,6 +234,11 @@ def triton_fused_store_indexer(
BYTES_PER_PAGE_F32=bytes_per_page_f32,
SCALE_PAGE_OFFSET_F32=scale_page_offset_f32,
HEAD_DIM=_INDEXER_HEAD_DIM,
PRESHUFFLE_TILE=(
INDEXER_K_CACHE_PRESHUFFLE_TILE
if aiter_can_use_preshuffle_paged_mqa()
else 0
),
FP8_MIN=_FP8_INFO.min,
FP8_MAX=_FP8_INFO.max,
EPS=1e-8,
python/sglang/srt/layers/attention/dsa/utils.py
+6
View File
@@ -56,6 +56,12 @@ def aiter_can_use_preshuffle_paged_mqa() -> bool:
return False
# Tile size for the indexer FP8 K-cache preshuffle layout. Store and gather
# kernels reorganize each page into (tile x tile) blocks so the aiter preshuffle
# paged-MQA gather can consume the cache directly.
INDEXER_K_CACHE_PRESHUFFLE_TILE = 16
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
python/sglang/srt/layers/attention/dsv4/indexer.py
+2 -1
View File
@@ -26,6 +26,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.attention.dsv4.compressor import Compressor
from sglang.srt.layers.attention.dsv4.metadata import (
NonPagedIndexerPlan,
@@ -164,7 +165,7 @@ def _aiter_fp8_paged_mqa_logits(
page_table.to(torch.int32),
max_seq_len,
KVBlockSize=kv_block_size,
Preshuffle=True,
Preshuffle=aiter_can_use_preshuffle_paged_mqa(),
)
return logits