From 1c6a0e91e1257619dc5d5c0693e7f937b4be3227 Mon Sep 17 00:00:00 2001 From: Wu Jiangming Date: Wed, 29 Jul 2026 22:36:48 +0800 Subject: [PATCH] fix mqa preshuffle layout issue for deepseek v4 (#31563) --- .../csrc/deepseek_v4/fused_norm_rope_v2.cuh | 27 +++++++++++++--- .../ops/attention/dsa/index_buf_accessor.py | 32 +++++++++++++++---- .../kernels/ops/attention/dsv4/compress.py | 12 ++++++- .../kernels/ops/kvcache/triton_store_cache.py | 25 ++++++++++++++- .../sglang/srt/layers/attention/dsa/utils.py | 6 ++++ .../srt/layers/attention/dsv4/indexer.py | 3 +- 6 files changed, 92 insertions(+), 13 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh index 9374b4c47..5c1b31096 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh @@ -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 +template 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(); - 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(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(scale_ptr)[0] = scale; } @@ -499,7 +511,14 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR } } -template +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 static constexpr auto select_kernel() { if constexpr (kIsIndexer) { - return fused_norm_rope_indexer; + return fused_norm_rope_indexer; } else { return fused_norm_rope_flashmla; } diff --git a/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py b/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py index 494ab948f..2cd4dd539 100644 --- a/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py +++ b/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py @@ -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 = ( diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index 148e23a10..bb3540076 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -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: diff --git a/python/sglang/kernels/ops/kvcache/triton_store_cache.py b/python/sglang/kernels/ops/kvcache/triton_store_cache.py index 9824989f9..509171e58 100644 --- a/python/sglang/kernels/ops/kvcache/triton_store_cache.py +++ b/python/sglang/kernels/ops/kvcache/triton_store_cache.py @@ -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, diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index d5b0b72b3..73d10839d 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 76f5954e0..bf46609e1 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -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