From 126cd5cfae7afbc9c7efb2b00dd52807ff17b842 Mon Sep 17 00:00:00 2001 From: kk <43161300+kkHuang-amd@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:30:03 +0800 Subject: [PATCH] gpt-oss decode performance optimization (#20392) Co-authored-by: wunhuang --- .../srt/layers/attention/aiter_backend.py | 445 ++++++++++++++---- python/sglang/srt/layers/attention/utils.py | 186 ++++++++ .../sglang/srt/layers/quantization/unquant.py | 4 + python/sglang/srt/server_args.py | 2 + 4 files changed, 549 insertions(+), 88 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 44c868eb9..c12d6487e 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index d679025dd..7cd278c82 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -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), + ) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index a1edcf5b4..38b1e2c24 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 589ef1ce2..319a1fd61 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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"