chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

gpt-oss decode performance optimization (#20392)

Browse Source 
Co-authored-by: wunhuang <wunhuang@amd.com>
This commit is contained in:
kk
2026-03-19 13:30:03 +08:00
committed by GitHub
parent cd22aa27a9
commit 126cd5cfae
4 changed files with 549 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

445
python/sglang/srt/layers/attention/aiter_backend.py
View File

@@ -13,7 +13,10 @@ import torch
import triton
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.layers.attention.utils import (
create_flashinfer_kv_indices_triton,
create_flashmla_kv_indices_triton,
)
from sglang.srt.layers.dp_attention import (
get_attention_tp_size,
is_dp_attention_enabled,
@@ -39,14 +42,19 @@ try:
paged_attention_ragged,
)
from aiter.mla import mla_decode_fwd, mla_prefill_fwd
from aiter.ops.triton.attention.unified_attention import unified_attention
except ImportError:
print(
"aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device."
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.utils import pad_sequence_with_mask
from sglang.srt.layers.attention.utils import (
launch_reshape_and_cache_flash,
pad_sequence_with_mask,
)
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.utils import get_bool_env_var
logger = logging.getLogger(__name__)
@@ -93,6 +101,7 @@ class ForwardMetadata:
mask_indptr: Optional[torch.Tensor] = None
max_extend_len: Optional[int] = None
fp8_prefill_kv_indices: Optional[torch.Tensor] = None
swa_page_table: Optional[torch.Tensor] = None
global_workspace_buffer = None
@@ -185,6 +194,18 @@ class AiterAttnBackend(AttentionBackend):
model_runner, self
)
# sliding window attention
self.use_sliding_window_kv_pool = (
isinstance(model_runner.token_to_kv_pool, SWAKVPool)
and model_runner.token_to_kv_pool.swa_layer_nums > 0
)
if self.use_sliding_window_kv_pool:
self.token_to_kv_pool = model_runner.token_to_kv_pool
self.use_triton_unified_attention = True
else:
self.use_triton_unified_attention = False
# aiter kernel related initialization
self.max_num_partitions = (
self.max_context_len + _AITER_PARTITION_SIZE_ROCM - 1
@@ -192,7 +213,7 @@ class AiterAttnBackend(AttentionBackend):
nbyes_per_qo_elem = torch.finfo(torch.float32).bits // 8
if not self.use_mla:
if not (self.use_mla or self.use_triton_unified_attention):
self.workspace_buffer = torch.empty(
(max_bs * self.num_head * self.max_num_partitions * self.head_dim)
* nbyes_per_qo_elem
@@ -439,6 +460,17 @@ class AiterAttnBackend(AttentionBackend):
is_causal=is_causal,
)
# for page size > 1 useful conversion function
def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor:
page_size = self.page_size
if page_size == 1:
return page_table
max_seqlen_k = page_table.shape[1]
strided_indices = torch.arange(
0, max_seqlen_k, page_size, device=page_table.device, dtype=torch.int32
)
return page_table[:, strided_indices] // page_size
def _resolve_v2_num_draft_tokens(
self,
extend_seq_lens: Optional[torch.Tensor] = None,
@@ -591,6 +623,7 @@ class AiterAttnBackend(AttentionBackend):
qo_indptr = None
kv_last_page_len = None
max_q_len = None
max_kv_len = None
work_metadata = None
work_indptr = None
@@ -600,24 +633,61 @@ class AiterAttnBackend(AttentionBackend):
reduce_partial_map = None
num_kv_splits = None
# num_kv_splits_indptr = None
swa_page_table = None
if forward_batch.forward_mode.is_decode_or_idle():
if spec_info is None or forward_batch.forward_mode.is_idle():
kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self._get_kv_indices_scratch(
forward_batch.seq_lens_sum, forward_batch.seq_lens.device
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
if not self.use_triton_unified_attention:
kv_indices = self._get_kv_indices_scratch(
forward_batch.seq_lens_sum, forward_batch.seq_lens.device
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
else:
max_q_len = 1
page_size = self.page_size
max_kv_len = torch.max(forward_batch.seq_lens).item()
max_num_blocks_per_seq = (max_kv_len + page_size - 1) // page_size
kv_indices = torch.zeros(
bs, max_kv_len, dtype=torch.int32, device=self.device
)
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
None,
kv_indices,
self.req_to_token.stride(0),
max_kv_len,
1,
)
if self.use_sliding_window_kv_pool:
swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
kv_indices
)
)
kv_indices = self._transform_table_1_to_real(kv_indices)
swa_page_table = self._transform_table_1_to_real(swa_page_table)
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[1 : bs + 1] = torch.cumsum(
self.kv_last_page_len[:bs], dim=0
)
else:
kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
bs = kv_indptr.shape[0] - 1
@@ -662,7 +732,7 @@ class AiterAttnBackend(AttentionBackend):
qo_indptr,
kv_last_page_len,
max_q_len,
None,
max_kv_len,
work_metadata=work_metadata,
work_info_set=work_info_set,
work_indptr=work_indptr,
@@ -671,6 +741,7 @@ class AiterAttnBackend(AttentionBackend):
reduce_partial_map=reduce_partial_map,
num_kv_splits=num_kv_splits,
run_graph=False,
swa_page_table=swa_page_table,
)
elif forward_batch.forward_mode.is_draft_extend_v2():
@@ -1054,6 +1125,14 @@ class AiterAttnBackend(AttentionBackend):
encoder_lens=forward_batch.encoder_lens,
spec_info=None,
)
if self.use_sliding_window_kv_pool:
swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.indices_updater_prefill.kv_indices
)
)
self.forward_metadata = ForwardMetadata(
self.indices_updater_prefill.kv_indptr,
self.indices_updater_prefill.kv_indices,
@@ -1061,6 +1140,7 @@ class AiterAttnBackend(AttentionBackend):
None,
self.indices_updater_prefill.max_q_len,
self.indices_updater_prefill.max_kv_len,
swa_page_table=swa_page_table,
)
def init_cuda_graph_state(
@@ -1071,8 +1151,11 @@ class AiterAttnBackend(AttentionBackend):
):
self.cuda_graph_kv_last_page_len = torch.ones(max_bs, dtype=torch.int)
if kv_indices_buf is None:
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
self.cuda_graph_kv_indices = torch.zeros(
(max_bs * self.max_context_len),
(max_bs * max_num_blocks_per_seq),
dtype=torch.int32,
device=self.device,
)
@@ -1111,6 +1194,16 @@ class AiterAttnBackend(AttentionBackend):
self.reduce_final_map = None
self.reduce_partial_map = None
if self.use_sliding_window_kv_pool:
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
self.cuda_graph_swa_page_table = torch.zeros(
(max_bs, max_num_blocks_per_seq),
dtype=torch.int32,
device=self.device,
)
def init_forward_metadata_capture_cuda_graph(
self,
bs: int,
@@ -1133,25 +1226,67 @@ class AiterAttnBackend(AttentionBackend):
reduce_final_map = None
reduce_partial_map = None
swa_page_table = None
max_kv_len = torch.max(seq_lens).item()
if forward_mode.is_decode_or_idle():
qo_indptr = None
kv_last_page_len = None
max_q_len = None
if spec_info is None:
kv_indptr = self.kv_indptr
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
if not self.use_triton_unified_attention:
kv_indptr = self.kv_indptr
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
else:
max_q_len = 1
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
kv_indices = self.cuda_graph_kv_indices.view(
-1, max_num_blocks_per_seq
)
page_indices = self.req_to_token[req_pool_indices[:bs], :max_kv_len]
if self.use_sliding_window_kv_pool:
swa_page_indices = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_indices
)
)
page_indices = self._transform_table_1_to_real(page_indices)
swa_page_indices = self._transform_table_1_to_real(
swa_page_indices
)
new_rows = swa_page_indices.shape[0]
new_cols = swa_page_indices.shape[1]
kv_indices[:new_rows, :new_cols].copy_(page_indices)
swa_page_table = self.cuda_graph_swa_page_table
swa_page_table[:new_rows, :new_cols].copy_(swa_page_indices)
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[1 : bs + 1] = torch.cumsum(
self.cuda_graph_kv_last_page_len[:bs], dim=0
)
kv_indptr = None
else:
kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
@@ -1196,7 +1331,7 @@ class AiterAttnBackend(AttentionBackend):
qo_indptr,
kv_last_page_len,
max_q_len,
kv_indptr[-1].item(),
max_kv_len,
work_metadata=work_metadata,
work_info_set=work_info_set,
work_indptr=work_indptr,
@@ -1204,6 +1339,7 @@ class AiterAttnBackend(AttentionBackend):
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
num_kv_splits=num_kv_splits,
swa_page_table=swa_page_table,
)
elif forward_mode.is_target_verify():
@@ -1469,25 +1605,65 @@ class AiterAttnBackend(AttentionBackend):
reduce_final_map = None
reduce_partial_map = None
swa_page_table = None
max_kv_len = torch.max(seq_lens).item()
if forward_mode.is_decode_or_idle():
qo_indptr = None
kv_last_page_len = None
max_q_len = None
if spec_info is None:
kv_indptr = self.kv_indptr
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
if not self.use_triton_unified_attention:
kv_indptr = self.kv_indptr
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
else:
max_q_len = 1
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
kv_indices = self.cuda_graph_kv_indices.view(
-1, max_num_blocks_per_seq
)
page_indices = self.req_to_token[req_pool_indices[:bs], :max_kv_len]
if self.use_sliding_window_kv_pool:
swa_page_indices = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_indices
)
)
page_indices = self._transform_table_1_to_real(page_indices)
swa_page_indices = self._transform_table_1_to_real(
swa_page_indices
)
new_rows = swa_page_indices.shape[0]
new_cols = swa_page_indices.shape[1]
kv_indices[:new_rows, :new_cols].copy_(page_indices)
swa_page_table = self.cuda_graph_swa_page_table
swa_page_table[:new_rows, :new_cols].copy_(swa_page_indices)
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[1 : bs + 1] = torch.cumsum(
self.cuda_graph_kv_last_page_len[:bs], dim=0
)
kv_indptr = None
else:
kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
@@ -1526,21 +1702,23 @@ class AiterAttnBackend(AttentionBackend):
reduce_final_map = self.reduce_final_map
reduce_partial_map = self.reduce_partial_map
self.forward_metadata = ForwardMetadata(
kv_indptr,
kv_indices,
qo_indptr,
kv_last_page_len,
max_q_len,
kv_indptr[-1].item(),
work_metadata=work_metadata,
work_info_set=work_info_set,
work_indptr=work_indptr,
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
num_kv_splits=num_kv_splits,
)
self.forward_metadata = ForwardMetadata(
kv_indptr,
kv_indices,
qo_indptr,
kv_last_page_len,
max_q_len,
max_kv_len,
work_metadata=work_metadata,
work_info_set=work_info_set,
work_indptr=work_indptr,
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
num_kv_splits=num_kv_splits,
swa_page_table=swa_page_table,
# num_kv_splits_indptr=num_kv_splits_indptr,
)
elif forward_mode.is_target_verify():
bs = len(req_pool_indices)
@@ -1794,18 +1972,42 @@ class AiterAttnBackend(AttentionBackend):
save_kv_cache=True,
sinks=None,
):
self.logits_soft_cap = layer.logit_cap
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
else forward_batch.encoder_out_cache_loc
)
self.logits_soft_cap = layer.logit_cap
if k is not None:
assert v is not None
if save_kv_cache:
if self.use_mla:
if self.use_triton_unified_attention:
token_to_kv_pool = forward_batch.token_to_kv_pool
k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
layer.layer_id
)
slot_mapping_swa = token_to_kv_pool.full_to_swa_index_mapping
launch_reshape_and_cache_flash(
k.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
v.view(-1, layer.tp_v_head_num, layer.v_head_dim),
k_cache.view(
-1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim
),
v_cache.view(
-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
),
cache_loc,
(
slot_mapping_swa.long()
if layer.sliding_window_size > 0
else None
),
)
elif self.use_mla:
forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
else:
forward_batch.token_to_kv_pool.set_kv_buffer(
@@ -2132,8 +2334,14 @@ class AiterAttnBackend(AttentionBackend):
v_cache = v_cache.to(dtype)
window_size = (-1, -1)
page_table = self.forward_metadata.kv_indices
if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
window_size = (layer.sliding_window_size, -1)
# page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
# page_table
# )
page_table = self.forward_metadata.swa_page_table
o = mha_batch_prefill_func(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
@@ -2141,7 +2349,7 @@ class AiterAttnBackend(AttentionBackend):
v_cache,
self.qo_indptr[:bs0],
self.forward_metadata.kv_indptr[:bs0],
self.forward_metadata.kv_indices,
page_table,
self.forward_metadata.max_q_len,
self.forward_metadata.max_kv_len,
causal=True,
@@ -2163,6 +2371,7 @@ class AiterAttnBackend(AttentionBackend):
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
sinks=None,
):
q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
@@ -2176,10 +2385,29 @@ class AiterAttnBackend(AttentionBackend):
o = torch.empty_like(q, dtype=self.input_dtype)
if save_kv_cache:
if self.use_triton_unified_attention:
token_to_kv_pool = forward_batch.token_to_kv_pool
k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
layer.layer_id
)
slot_mapping_swa = token_to_kv_pool.full_to_swa_index_mapping
forward_batch.token_to_kv_pool.set_kv_buffer(
layer, forward_batch.out_cache_loc, k, v
)
launch_reshape_and_cache_flash(
k.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
v.view(-1, layer.tp_v_head_num, layer.v_head_dim),
k_cache.view(
-1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim
),
v_cache.view(
-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
),
forward_batch.out_cache_loc,
slot_mapping_swa.long() if layer.sliding_window_size > 0 else None,
)
else:
forward_batch.token_to_kv_pool.set_kv_buffer(
layer, forward_batch.out_cache_loc, k, v
)
if self.use_mla:
k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
@@ -2230,27 +2458,68 @@ class AiterAttnBackend(AttentionBackend):
k_cache = k_cache.to(dtype)
v_cache = v_cache.to(dtype)
paged_attention_ragged(
o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
self.workspace_buffer,
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k_cache.view(-1, 1, layer.tp_k_head_num, layer.qk_head_dim),
v_cache.view(-1, 1, layer.tp_v_head_num, layer.v_head_dim),
self.scale,
self.forward_metadata.kv_indptr,
self.forward_metadata.kv_indices,
self.kv_last_page_len,
1,
self.max_num_partitions,
None,
"auto",
"NHD",
self.logits_soft_cap,
self.k_scale,
self.v_scale,
None,
_AITER_PARTITION_SIZE_ROCM,
)
if self.use_triton_unified_attention:
bs = forward_batch.batch_size
window_size = (-1, -1)
page_table = self.forward_metadata.kv_indices
if (
layer.sliding_window_size is not None
and layer.sliding_window_size > -1
):
window_size = (layer.sliding_window_size - 1, 0)
page_table = self.forward_metadata.swa_page_table
o = torch.empty_like(q)
max_kv_len = page_table.shape[1]
unified_attention(
q=q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k=k_cache.view(
-1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim
),
v=v_cache.view(
-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
),
out=o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
cu_seqlens_q=self.forward_metadata.qo_indptr,
seqused_k=forward_batch.seq_lens,
max_seqlen_q=self.forward_metadata.max_q_len,
max_seqlen_k=max_kv_len,
softmax_scale=self.scale,
causal=True,
window_size=window_size,
block_table=page_table,
softcap=0,
q_descale=None,
k_descale=None,
v_descale=None,
sinks=sinks,
)
else:
paged_attention_ragged(
o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
self.workspace_buffer,
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k_cache.view(-1, 1, layer.tp_k_head_num, layer.qk_head_dim),
v_cache.view(-1, 1, layer.tp_v_head_num, layer.v_head_dim),
self.scale,
self.forward_metadata.kv_indptr,
self.forward_metadata.kv_indices,
self.kv_last_page_len,
1,
self.max_num_partitions,
None,
"auto",
"NHD",
self.logits_soft_cap,
self.k_scale,
self.v_scale,
None,
_AITER_PARTITION_SIZE_ROCM,
)
return o

186
python/sglang/srt/layers/attention/utils.py
View File

@@ -472,3 +472,189 @@ def concat_mla_absorb_q_general(q_nope, q_rope):
return concat_mla_absorb_q(q_nope, q_rope)
else:
return torch.cat([q_nope, q_rope], dim=-1)
@triton.jit
def reshape_and_cache_flash(
key_ptr,
value_ptr,
key_cache_ptr,
value_cache_ptr,
slot_mapping_ptr,
swa_slot_mapping_ptr,
k_scale_ptr,
v_scale_ptr,
block_stride,
key_stride,
value_stride,
num_heads,
head_size,
block_size,
HEAD_BLOCK: tl.constexpr,
BLOCK_D: tl.constexpr,
HAS_SWA: tl.constexpr,
USE_SCALE: tl.constexpr,
):
"""
Triton kernel for reshaping per-token K/V tensors into paged KV cache layout.
Source layout:
key/value: [num_tokens, num_heads, head_size]
Target cache layout:
cache: [num_blocks, block_size, num_heads, head_size]
Each Triton program instance handles:
- one token (program_id(0))
- one block of heads (program_id(1))
Features:
- optional SWA slot remapping
- optional FP8 scale dequantization before cache write
Args:
key_ptr: Pointer to source key tensor.
value_ptr: Pointer to source value tensor.
key_cache_ptr: Pointer to destination key cache tensor.
value_cache_ptr: Pointer to destination value cache tensor.
slot_mapping_ptr: Maps token -> cache slot.
swa_slot_mapping_ptr: Optional second-stage slot remap for SWA mode.
k_scale_ptr: Optional key scaling factor pointer.
v_scale_ptr: Optional value scaling factor pointer.
block_stride: Stride between cache blocks.
key_stride: Stride between source key tokens.
value_stride: Stride between source value tokens.
num_heads: Number of attention heads.
head_size: Hidden dimension per head.
block_size: Number of slots per cache block.
HEAD_BLOCK: Number of heads processed per program.
BLOCK_D: Vectorized dimension size (power-of-2 padded).
HAS_SWA: Enable SWA remapping.
USE_SCALE: Enable scale division before storing.
"""
# ----------------------------------
# program ids
# pid0 = token
# pid1 = head block
# ----------------------------------
token_idx = tl.program_id(0)
head_block_idx = tl.program_id(1)
# ----------------------------------
# slot mapping
# ----------------------------------
slot_idx = tl.load(slot_mapping_ptr + token_idx)
if HAS_SWA:
slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx)
if slot_idx < 0:
return
block_idx = slot_idx // block_size
block_offset = slot_idx % block_size
# ----------------------------------
# head range
# ----------------------------------
head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)
head_mask = head_idx < num_heads
dim_idx = tl.arange(0, BLOCK_D)
# shape = [HEAD_BLOCK, BLOCK_D]
offs = head_idx[:, None] * head_size + dim_idx[None, :]
mask = head_mask[:, None] & (dim_idx[None, :] < head_size)
# ----------------------------------
# source load
# ----------------------------------
src_key = token_idx * key_stride + offs
src_value = token_idx * value_stride + offs
k = tl.load(key_ptr + src_key, mask=mask)
v = tl.load(value_ptr + src_value, mask=mask)
# ----------------------------------
# optional scale
# ----------------------------------
if USE_SCALE:
k_scale = tl.load(k_scale_ptr)
v_scale = tl.load(v_scale_ptr)
k = k / k_scale
v = v / v_scale
# ----------------------------------
# target layout
# [block_idx, block_offset, head, dim]
# ----------------------------------
tgt = block_idx * block_stride + block_offset * num_heads * head_size + offs
tl.store(key_cache_ptr + tgt, k, mask=mask)
tl.store(value_cache_ptr + tgt, v, mask=mask)
def launch_reshape_and_cache_flash(
key,
value,
key_cache,
value_cache,
slot_mapping,
swa_slot_mapping=None,
k_scale=None,
v_scale=None,
):
"""
Launch wrapper for reshape_and_cache_flash Triton kernel.
This wrapper prepares launch configuration and dispatches the Triton kernel
that writes token-major K/V tensors into paged KV cache layout.
Args:
key: Source key tensor [num_tokens, num_heads, head_size]
value: Source value tensor [num_tokens, num_heads, head_size]
key_cache: Destination key cache [num_blocks, block_size, num_heads, head_size]
value_cache: Destination value cache [num_blocks, block_size, num_heads, head_size]
slot_mapping: Token-to-cache slot mapping
swa_slot_mapping: Optional SWA remapping table
k_scale: Optional key scaling factor
v_scale: Optional value scaling factor
"""
num_tokens = key.shape[0]
num_heads = key.shape[1]
head_size = key.shape[2]
HEAD_BLOCK = 4
BLOCK_D = triton.next_power_of_2(head_size)
grid = (
num_tokens,
triton.cdiv(num_heads, HEAD_BLOCK),
)
reshape_and_cache_flash[grid](
key,
value,
key_cache,
value_cache,
slot_mapping,
swa_slot_mapping if swa_slot_mapping is not None else key,
k_scale if k_scale is not None else key,
v_scale if v_scale is not None else key,
key_cache.stride(0),
key.stride(0),
value.stride(0),
num_heads,
head_size,
key_cache.shape[1],
HEAD_BLOCK=HEAD_BLOCK,
BLOCK_D=BLOCK_D,
HAS_SWA=(swa_slot_mapping is not None),
USE_SCALE=(k_scale is not None),
)

4
python/sglang/srt/layers/quantization/unquant.py
View File

@@ -52,6 +52,7 @@ if _use_aiter:
from aiter import ActivationType
from aiter.fused_moe import fused_moe
from aiter.ops.shuffle import shuffle_weight
from aiter.tuned_gemm import tgemm
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
@@ -150,6 +151,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
output = output.view(x_shapes[0], x_shapes[1], -1)
return output
elif _use_aiter and type(layer.weight.data) is torch.Tensor:
return tgemm.mm(x, layer.weight, bias, otype=x.dtype)
return F.linear(x, layer.weight, bias)

2
python/sglang/srt/server_args.py
View File

@@ -1626,6 +1626,8 @@ class ServerArgs:
self.attention_backend = "trtllm_mha"
elif is_sm90_supported():
self.attention_backend = "fa3"
elif is_hip():
self.attention_backend = "aiter"
else:
self.attention_backend = "triton"