diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index fe1f8de78..dce97babf 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -5,11 +5,13 @@ from typing import TYPE_CHECKING, Optional import numpy as np import torch -import triton -import triton.language as tl from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.triton_ops.metadata import ( + normal_decode_set_metadata, + prepare_swa_spec_page_table_triton, +) from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.layers.utils.cp_utils import ( cp_allgather_and_save_kv_cache, @@ -1256,7 +1258,6 @@ class FlashAttentionBackend(AttentionBackend): ) ) cache_seqlens = metadata.cache_seqlens_int32 - cu_seqlens_k = metadata.cu_seqlens_k max_seqlen_q = metadata.max_seq_len_q q_reshaped = q.contiguous().view( -1, layer.tp_q_head_num, layer.head_dim @@ -1991,7 +1992,6 @@ class FlashAttentionBackend(AttentionBackend): metadata_expand = None if forward_mode.is_decode_or_idle(): - if spec_info is not None: # Draft Decode if self.topk <= 1: @@ -2597,110 +2597,7 @@ class FlashAttentionBackend(AttentionBackend): metadata.swa_spec_metadata = metadata_swa -@triton.jit -def _prepare_swa_spec_page_table_kernel( - dst_ptr, - src_a_ptr, - src_b_ptr, - seq_len_a_ptr, - seq_len_b_ptr, - dst_stride_m, - dst_stride_n, - a_stride_m, - a_stride_n, - b_stride_m, - b_stride_n, - LEN_A: tl.constexpr, - LEN_B: tl.constexpr, - REPEAT_STEP: tl.constexpr, - BLOCK_N: tl.constexpr, -): - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - - idx_a = pid_m // REPEAT_STEP - idx_b = pid_m - seq_len_a = tl.load(seq_len_a_ptr + idx_a) - seq_len_b = tl.load(seq_len_b_ptr + idx_b) - - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - total_len = seq_len_a + seq_len_b - - if pid_n * BLOCK_N >= total_len: - return - - mask = offs_n < total_len - dst = dst_ptr + pid_m * dst_stride_m + offs_n * dst_stride_n - - if (pid_n + 1) * BLOCK_N < seq_len_a: - a_ptr = src_a_ptr + idx_a * a_stride_m + offs_n * a_stride_n - a_mask = mask & (offs_n < LEN_A) - val = tl.load(a_ptr, mask=a_mask, other=0) - tl.store(dst, val, mask=mask) - elif pid_n * BLOCK_N >= seq_len_a: - offs_b = offs_n - seq_len_a - b_ptr = src_b_ptr + idx_b * b_stride_m + offs_b * b_stride_n - b_mask = mask & (offs_b < LEN_B) - val = tl.load(b_ptr, mask=b_mask, other=0) - tl.store(dst, val, mask=mask) - else: - # mixed part - a_offs = offs_n - a_mask = (a_offs < seq_len_a) & (a_offs < LEN_A) - a_ptr = src_a_ptr + idx_a * a_stride_m + a_offs * a_stride_n - a_val = tl.load(a_ptr, mask=a_mask, other=0) - - b_offs = offs_n - seq_len_a - b_mask = (b_offs >= 0) & (b_offs < seq_len_b) & (b_offs < LEN_B) - b_ptr = src_b_ptr + idx_b * b_stride_m + b_offs * b_stride_n - b_val = tl.load(b_ptr, mask=b_mask, other=0) - - result = tl.where(offs_n < seq_len_a, a_val, b_val) - tl.store(dst, result, mask=mask) - - -def prepare_swa_spec_page_table_triton( - page_table_dst: torch.Tensor, - page_table_a: torch.Tensor, - page_table_b: torch.Tensor, # expand page table - seq_len_a: torch.Tensor, - seq_len_b: torch.Tensor, # expand seq lens - speculative_num_draft_tokens: int, -): - # concat page_table and expand page_table by kv seq length - bs = seq_len_a.numel() - bs_expand = seq_len_b.numel() - assert bs_expand == bs * speculative_num_draft_tokens - - LEN_A = page_table_a.shape[1] - LEN_B = page_table_b.shape[1] - LEN_OUT = LEN_A + LEN_B - REPEAT_STEP = speculative_num_draft_tokens - BLOCK_N = 256 - - grid = (bs_expand, triton.cdiv(LEN_OUT, BLOCK_N)) - _prepare_swa_spec_page_table_kernel[grid]( - page_table_dst, - page_table_a, - page_table_b, - seq_len_a, - seq_len_b, - page_table_dst.stride(0), - page_table_dst.stride(1), - page_table_a.stride(0), - page_table_a.stride(1), - page_table_b.stride(0), - page_table_b.stride(1), - LEN_A=LEN_A, - LEN_B=LEN_B, - REPEAT_STEP=REPEAT_STEP, - BLOCK_N=BLOCK_N, - num_warps=4, - ) - - class FlashAttentionMultiStepBackend: - def __init__( self, model_runner: ModelRunner, @@ -2771,310 +2668,6 @@ class FlashAttentionMultiStepBackend: ) -@triton.jit -def _fused_metadata_kernel_general( - # Input tensors - seq_lens, - seq_lens_stride_0, - req_to_token, - req_to_token_stride_0, - req_to_token_stride_1, - req_pool_indices, - req_pool_indices_stride_0, - # Output buffers - cache_seqlens_int32, - cache_seqlens_int32_stride_0, - cu_seqlens_k, - cu_seqlens_k_stride_0, - page_table, - page_table_stride_0, - page_table_stride_1, - swa_page_table, - swa_page_table_stride_0, - swa_page_table_stride_1, - full_to_swa_mapping, - full_to_swa_mapping_stride_0, - # Scalar parameters - B, - max_seq_pages, - page_size: tl.constexpr, - seq_len_delta: tl.constexpr, - use_swa: tl.constexpr, - SHIFT: tl.constexpr, - BLOCK_COLS: tl.constexpr, -): - pid_b = tl.program_id(0) # batch index - pid_c = tl.program_id(1) # column chunk index - - # 1. Prefix sum (only one block does it) - if pid_b == 0 and pid_c == 0: - acc = 0 - for idx in range(B): - seq = tl.load(seq_lens + idx * seq_lens_stride_0) - val = (seq + seq_len_delta).to(tl.int32) - tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val) - tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc) - acc += val - tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc) - - # 2. Gather for this batch and column chunk - if max_seq_pages == 0: - return - - i = pid_b - # Load row index for this batch (all threads in block have same i) - row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0) - row_offset = row_idx * req_to_token_stride_0 - - col_start = pid_c * BLOCK_COLS - col_offsets = col_start + tl.arange(0, BLOCK_COLS) - mask = col_offsets < max_seq_pages - - # Compute column indices in the source tensor (token offset) - if page_size == 1: - col_idx = col_offsets - else: - col_idx = col_offsets << SHIFT # faster than multiplication for power-of-two - - # Load page indices from req_to_token - rt_offsets = row_offset + col_idx * req_to_token_stride_1 - page_index = tl.load( - req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg" - ) - - # Compute page_table - if page_size == 1: - page_table_val = page_index - else: - page_table_val = page_index >> SHIFT - - # Store to page_table - pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1 - tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg") - - if use_swa: - swa_slot = tl.load( - full_to_swa_mapping + page_index * full_to_swa_mapping_stride_0, - mask=mask, - other=0, - cache_modifier=".cg", - ) - if page_size == 1: - swa_val = swa_slot - else: - swa_val = swa_slot >> SHIFT - swa_offsets = ( - i * swa_page_table_stride_0 + col_offsets * swa_page_table_stride_1 - ) - tl.store(swa_page_table + swa_offsets, swa_val, mask=mask, cache_modifier=".cg") - - -@triton.jit -def _fused_metadata_kernel_ps1_no_swa( - # Input tensors - seq_lens, - seq_lens_stride_0, - req_to_token, - req_to_token_stride_0, - req_to_token_stride_1, - req_pool_indices, - req_pool_indices_stride_0, - # Output buffers - cache_seqlens_int32, - cache_seqlens_int32_stride_0, - cu_seqlens_k, - cu_seqlens_k_stride_0, - page_table, - page_table_stride_0, - page_table_stride_1, - # Scalar parameters - B, - max_seq_pages, - seq_len_delta: tl.constexpr, - BLOCK_COLS: tl.constexpr, -): - pid_b = tl.program_id(0) # batch index - pid_c = tl.program_id(1) # column chunk index - - # 1. Prefix sum (only one block does it) - if pid_b == 0 and pid_c == 0: - acc = 0 - for idx in range(B): - seq = tl.load(seq_lens + idx * seq_lens_stride_0) - val = (seq + seq_len_delta).to(tl.int32) - tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val) - tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc) - acc += val - tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc) - - # 2. Gather for this batch and column chunk - if max_seq_pages == 0: - return - - i = pid_b - # Load row index for this batch (all threads in block have same i) - row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0) - row_offset = row_idx * req_to_token_stride_0 - - col_start = pid_c * BLOCK_COLS - col_offsets = col_start + tl.arange(0, BLOCK_COLS) - mask = col_offsets < max_seq_pages - - # page_size = 1: col_idx = col_offsets - rt_offsets = row_offset + col_offsets * req_to_token_stride_1 - page_index = tl.load( - req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg" - ) - - # page_table = page_index // 1 = page_index - pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1 - tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg") - - -# Fused Triton kernel implementation -def normal_decode_set_metadata( - cache_seqlens_int32: torch.Tensor, - cu_seqlens_k: torch.Tensor, - page_table: torch.Tensor, - req_to_token: torch.Tensor, - req_pool_indices: torch.Tensor, - strided_indices: torch.Tensor, - max_seq_pages: torch.Tensor, - seq_lens: torch.Tensor, - seq_len_delta: int, - page_size: int, - swa_page_table: Optional[torch.Tensor] = None, - token_to_kv_pool: Optional[SWAKVPool] = None, -): - """ - Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels: - 1. cache_seqlens = seq_lens + seq_len_delta (int64→int32 cast) - 2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum) - 3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather) - 4. page_table = page_indices // page_size (floor-divide) - 5. (optional) swa_page_table for sliding window attention - - Achieves ~5.2x speedup on H200 hardware for typical decode workloads. - """ - assert ( - page_size > 0 and (page_size & (page_size - 1)) == 0 - ), f"page_size must be a power of two, got {page_size}" - - batch_size = cache_seqlens_int32.shape[0] - device = seq_lens.device - - # Ensure contiguous memory layout for efficient Triton access - seq_lens = seq_lens.contiguous() - req_to_token = req_to_token.contiguous() - req_pool_indices = req_pool_indices.contiguous() - - # Prepare tensor strides - seq_lens_stride_0 = seq_lens.stride(0) - req_to_token_stride_0 = req_to_token.stride(0) - req_to_token_stride_1 = req_to_token.stride(1) - req_pool_indices_stride_0 = req_pool_indices.stride(0) - cache_seqlens_int32_stride_0 = cache_seqlens_int32.stride(0) - cu_seqlens_k_stride_0 = cu_seqlens_k.stride(0) - page_table_stride_0 = page_table.stride(0) - page_table_stride_1 = page_table.stride(1) - - # Check if we should use the specialized fast path for page_size=1, no SWA - use_swa = swa_page_table is not None and token_to_kv_pool is not None - - if page_size == 1 and not use_swa: - # Specialized kernel for the common case (page_size=1, no SWA) - BLOCK_COLS = 256 - if max_seq_pages == 0: - grid = (1, 1) - else: - num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS) - grid = (batch_size, num_blocks_j) - - _fused_metadata_kernel_ps1_no_swa[grid]( - seq_lens, - seq_lens_stride_0, - req_to_token, - req_to_token_stride_0, - req_to_token_stride_1, - req_pool_indices, - req_pool_indices_stride_0, - cache_seqlens_int32, - cache_seqlens_int32_stride_0, - cu_seqlens_k, - cu_seqlens_k_stride_0, - page_table, - page_table_stride_0, - page_table_stride_1, - batch_size, - max_seq_pages, - seq_len_delta, - BLOCK_COLS=BLOCK_COLS, - num_warps=8, - num_stages=3, - ) - else: - # General kernel for page_size > 1 or SWA cases - # SWA parameters - if use_swa: - assert isinstance(token_to_kv_pool, SWAKVPool) - swa_page_table = swa_page_table.contiguous() - swa_page_table_stride_0 = swa_page_table.stride(0) - swa_page_table_stride_1 = swa_page_table.stride(1) - # Extract the full_to_swa_index_mapping from token_to_kv_pool - full_to_swa_mapping = ( - token_to_kv_pool.full_to_swa_index_mapping.contiguous() - ) - full_to_swa_mapping_stride_0 = full_to_swa_mapping.stride(0) - else: - # Dummy tensors (not used) - swa_page_table = torch.empty(0, dtype=torch.int32, device=device) - swa_page_table_stride_0 = 0 - swa_page_table_stride_1 = 0 - full_to_swa_mapping = torch.empty(0, dtype=torch.int32, device=device) - full_to_swa_mapping_stride_0 = 0 - - # Kernel configuration - BLOCK_COLS = 128 - shift = (page_size).bit_length() - 1 if page_size > 1 else 0 - - if max_seq_pages == 0: - grid = (1, 1) - else: - num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS) - grid = (batch_size, num_blocks_j) - - _fused_metadata_kernel_general[grid]( - seq_lens, - seq_lens_stride_0, - req_to_token, - req_to_token_stride_0, - req_to_token_stride_1, - req_pool_indices, - req_pool_indices_stride_0, - cache_seqlens_int32, - cache_seqlens_int32_stride_0, - cu_seqlens_k, - cu_seqlens_k_stride_0, - page_table, - page_table_stride_0, - page_table_stride_1, - swa_page_table, - swa_page_table_stride_0, - swa_page_table_stride_1, - full_to_swa_mapping, - full_to_swa_mapping_stride_0, - batch_size, - max_seq_pages, - page_size, - seq_len_delta, - use_swa, - shift, - BLOCK_COLS=BLOCK_COLS, - num_warps=4, - num_stages=3, - ) - - @torch.compile(dynamic=True, backend=get_compiler_backend()) def draft_decode_set_expand_metadata( cache_seqlens_int32: torch.Tensor, # Modifies diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index d71871e5b..1a0e7afc0 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -5,11 +5,13 @@ from typing import TYPE_CHECKING, List, Optional import torch import triton -import triton.language as tl from sglang.srt.configs.model_config import AttentionArch 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.triton_ops.kv_indices import ( + create_flashinfer_kv_indices_triton, +) +from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool @@ -1447,58 +1449,6 @@ class TritonMultiStepDraftBackend: ) -@triton.jit -def get_num_kv_splits_triton( - num_kv_splits_ptr, - seq_lens_ptr, - num_seq, - num_group, - num_head, - num_kv_head, - max_kv_splits, - device_core_count, - MAX_NUM_SEQ: tl.constexpr, -): - # TODO: this method is tunable, we need more online serving data to tune it - offs_seq = tl.arange(0, MAX_NUM_SEQ) - mask_seq = offs_seq < num_seq - - seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0) - max_seq_len = tl.max(seq_lens) - seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len) - min_seq_len = tl.min(seq_lens) - if max_seq_len * 8 < min_seq_len * 10: - min_seq_len = max_seq_len - max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits) - kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1) - - # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually - ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0 - ext_device_core_count = tl.cast( - device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32 - ) - block_h, num_kv_group = 16, num_head // num_kv_head - if num_kv_group == 1: - token_grid = num_seq * num_group * num_head - else: - # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd - block_h = tl.minimum(block_h, num_kv_group) - token_grid = num_seq * num_group * tl.cdiv(num_head, block_h) - max_kv_splits_2 = tl.minimum( - tl.cdiv(ext_device_core_count, token_grid), max_kv_splits - ) - kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2) - - num_kv_splits = tl.maximum( - tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2) - ) - - offs_token = offs_seq * num_group - mask_token = offs_token < num_seq * num_group - for i in range(0, num_group): - tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token) - - def update_sliding_window_buffer( window_kv_indptr, req_to_token, diff --git a/python/sglang/srt/layers/attention/triton_ops/cache_ops.py b/python/sglang/srt/layers/attention/triton_ops/cache_ops.py new file mode 100644 index 000000000..7e1da0b53 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/cache_ops.py @@ -0,0 +1,266 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def concat_and_cast_mha_k_kernel( + k_ptr, + k_nope_ptr, + k_rope_ptr, + head_cnt: tl.constexpr, + k_stride0: tl.constexpr, + k_stride1: tl.constexpr, + nope_stride0: tl.constexpr, + nope_stride1: tl.constexpr, + rope_stride0: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, +): + pid_loc = tl.program_id(0) + head_range = tl.arange(0, head_cnt) + + k_head_ptr = k_ptr + pid_loc * k_stride0 + head_range[:, None] * k_stride1 + + nope_offs = tl.arange(0, nope_dim) + + src_nope_ptr = ( + k_nope_ptr + + pid_loc * nope_stride0 + + head_range[:, None] * nope_stride1 + + nope_offs[None, :] + ) + dst_nope_ptr = k_head_ptr + nope_offs[None, :] + + src_nope = tl.load(src_nope_ptr) + tl.store(dst_nope_ptr, src_nope) + + rope_offs = tl.arange(0, rope_dim) + src_rope_ptr = k_rope_ptr + pid_loc * rope_stride0 + rope_offs[None, :] + dst_rope_ptr = k_head_ptr + nope_dim + rope_offs[None, :] + src_rope = tl.load(src_rope_ptr) + tl.store(dst_rope_ptr, src_rope) + + +def concat_and_cast_mha_k_triton( + k: torch.Tensor, + k_nope: torch.Tensor, + k_rope: torch.Tensor, +): + # The source data type will be implicitly converted to the target data type. + assert ( + len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3 + ), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + assert ( + k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0] + ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + assert ( + k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1] + ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + assert ( + k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1] + ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + + nope_dim = k_nope.shape[-1] + rope_dim = k_rope.shape[-1] + grid = (k.shape[0],) + + concat_and_cast_mha_k_kernel[grid]( + k, + k_nope, + k_rope, + k.shape[1], + k.stride(0), + k.stride(1), + k_nope.stride(0), + k_nope.stride(1), + k_rope.stride(0), + nope_dim, + rope_dim, + ) + + +@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, + 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/attention/triton_ops/kv_indices.py b/python/sglang/srt/layers/attention/triton_ops/kv_indices.py new file mode 100644 index 000000000..aa39e0cb0 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/kv_indices.py @@ -0,0 +1,103 @@ +import triton +import triton.language as tl + +_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096 +FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE) + + +@triton.jit +def create_flashinfer_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_indptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + kv_indices_offset = tl.load(kv_indptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for i in range(num_loop): + # index into req_to_token_ptr needs to be int64 + offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE + mask = offset < kv_end - kv_start + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + offset, + mask=mask, + ) + tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask) + + +def get_num_page_per_block_flashmla(page_size: int = 64) -> int: + num_page_per_block = _FLASHMLA_CREATE_KV_BLOCK_SIZE // page_size + return num_page_per_block + + +@triton.jit +def create_flashmla_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, + kv_indices_ptr_stride: tl.constexpr, + PAGED_SIZE: tl.constexpr = 64, +): + NUM_PAGE_PER_BLOCK: tl.constexpr = ( + FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE + ) + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE) + num_pages_loop = tl.cdiv(kv_end - kv_start, FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON) + + for i in range(num_pages_loop): + # index into req_to_token_ptr needs to be int64 + paged_offset = ( + tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK + ) * PAGED_SIZE + paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK + + mask = paged_offset < num_paged * PAGED_SIZE + mask_out = paged_offset_out < num_paged + + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + paged_offset, + mask=mask, + ) + tl.store( + kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out, + data // PAGED_SIZE, + mask=mask_out, + ) diff --git a/python/sglang/srt/layers/attention/triton_ops/metadata.py b/python/sglang/srt/layers/attention/triton_ops/metadata.py new file mode 100644 index 000000000..47504a1cf --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/metadata.py @@ -0,0 +1,467 @@ +from typing import TYPE_CHECKING, Optional + +import torch +import triton +import triton.language as tl + +if TYPE_CHECKING: + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool + + +@triton.jit +def get_num_kv_splits_triton( + num_kv_splits_ptr, + seq_lens_ptr, + num_seq, + num_group, + num_head, + num_kv_head, + max_kv_splits, + device_core_count, + MAX_NUM_SEQ: tl.constexpr, +): + # TODO: this method is tunable, we need more online serving data to tune it + offs_seq = tl.arange(0, MAX_NUM_SEQ) + mask_seq = offs_seq < num_seq + + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0) + max_seq_len = tl.max(seq_lens) + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len) + min_seq_len = tl.min(seq_lens) + if max_seq_len * 8 < min_seq_len * 10: + min_seq_len = max_seq_len + max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits) + kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1) + + # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually + ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0 + ext_device_core_count = tl.cast( + device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32 + ) + block_h, num_kv_group = 16, num_head // num_kv_head + if num_kv_group == 1: + token_grid = num_seq * num_group * num_head + else: + # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd + block_h = tl.minimum(block_h, num_kv_group) + token_grid = num_seq * num_group * tl.cdiv(num_head, block_h) + max_kv_splits_2 = tl.minimum( + tl.cdiv(ext_device_core_count, token_grid), max_kv_splits + ) + kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2) + + num_kv_splits = tl.maximum( + tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2) + ) + + offs_token = offs_seq * num_group + mask_token = offs_token < num_seq * num_group + for i in range(0, num_group): + tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token) + + +@triton.jit +def _prepare_swa_spec_page_table_kernel( + dst_ptr, + src_a_ptr, + src_b_ptr, + seq_len_a_ptr, + seq_len_b_ptr, + dst_stride_m, + dst_stride_n, + a_stride_m, + a_stride_n, + b_stride_m, + b_stride_n, + LEN_A: tl.constexpr, + LEN_B: tl.constexpr, + REPEAT_STEP: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + idx_a = pid_m // REPEAT_STEP + idx_b = pid_m + seq_len_a = tl.load(seq_len_a_ptr + idx_a) + seq_len_b = tl.load(seq_len_b_ptr + idx_b) + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + total_len = seq_len_a + seq_len_b + + if pid_n * BLOCK_N >= total_len: + return + + mask = offs_n < total_len + dst = dst_ptr + pid_m * dst_stride_m + offs_n * dst_stride_n + + if (pid_n + 1) * BLOCK_N < seq_len_a: + a_ptr = src_a_ptr + idx_a * a_stride_m + offs_n * a_stride_n + a_mask = mask & (offs_n < LEN_A) + val = tl.load(a_ptr, mask=a_mask, other=0) + tl.store(dst, val, mask=mask) + elif pid_n * BLOCK_N >= seq_len_a: + offs_b = offs_n - seq_len_a + b_ptr = src_b_ptr + idx_b * b_stride_m + offs_b * b_stride_n + b_mask = mask & (offs_b < LEN_B) + val = tl.load(b_ptr, mask=b_mask, other=0) + tl.store(dst, val, mask=mask) + else: + # mixed part + a_offs = offs_n + a_mask = (a_offs < seq_len_a) & (a_offs < LEN_A) + a_ptr = src_a_ptr + idx_a * a_stride_m + a_offs * a_stride_n + a_val = tl.load(a_ptr, mask=a_mask, other=0) + + b_offs = offs_n - seq_len_a + b_mask = (b_offs >= 0) & (b_offs < seq_len_b) & (b_offs < LEN_B) + b_ptr = src_b_ptr + idx_b * b_stride_m + b_offs * b_stride_n + b_val = tl.load(b_ptr, mask=b_mask, other=0) + + result = tl.where(offs_n < seq_len_a, a_val, b_val) + tl.store(dst, result, mask=mask) + + +def prepare_swa_spec_page_table_triton( + page_table_dst: torch.Tensor, + page_table_a: torch.Tensor, + page_table_b: torch.Tensor, # expand page table + seq_len_a: torch.Tensor, + seq_len_b: torch.Tensor, # expand seq lens + speculative_num_draft_tokens: int, +): + # concat page_table and expand page_table by kv seq length + bs = seq_len_a.numel() + bs_expand = seq_len_b.numel() + assert bs_expand == bs * speculative_num_draft_tokens + + LEN_A = page_table_a.shape[1] + LEN_B = page_table_b.shape[1] + LEN_OUT = LEN_A + LEN_B + REPEAT_STEP = speculative_num_draft_tokens + BLOCK_N = 256 + + grid = (bs_expand, triton.cdiv(LEN_OUT, BLOCK_N)) + _prepare_swa_spec_page_table_kernel[grid]( + page_table_dst, + page_table_a, + page_table_b, + seq_len_a, + seq_len_b, + page_table_dst.stride(0), + page_table_dst.stride(1), + page_table_a.stride(0), + page_table_a.stride(1), + page_table_b.stride(0), + page_table_b.stride(1), + LEN_A=LEN_A, + LEN_B=LEN_B, + REPEAT_STEP=REPEAT_STEP, + BLOCK_N=BLOCK_N, + num_warps=4, + ) + + +@triton.jit +def _fused_metadata_kernel_general( + # Input tensors + seq_lens, + seq_lens_stride_0, + req_to_token, + req_to_token_stride_0, + req_to_token_stride_1, + req_pool_indices, + req_pool_indices_stride_0, + # Output buffers + cache_seqlens_int32, + cache_seqlens_int32_stride_0, + cu_seqlens_k, + cu_seqlens_k_stride_0, + page_table, + page_table_stride_0, + page_table_stride_1, + swa_page_table, + swa_page_table_stride_0, + swa_page_table_stride_1, + full_to_swa_mapping, + full_to_swa_mapping_stride_0, + # Scalar parameters + B, + max_seq_pages, + page_size: tl.constexpr, + seq_len_delta: tl.constexpr, + use_swa: tl.constexpr, + SHIFT: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + pid_b = tl.program_id(0) # batch index + pid_c = tl.program_id(1) # column chunk index + + # 1. Prefix sum (only one block does it) + if pid_b == 0 and pid_c == 0: + acc = 0 + for idx in range(B): + seq = tl.load(seq_lens + idx * seq_lens_stride_0) + val = (seq + seq_len_delta).to(tl.int32) + tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val) + tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc) + acc += val + tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc) + + # 2. Gather for this batch and column chunk + if max_seq_pages == 0: + return + + i = pid_b + # Load row index for this batch (all threads in block have same i) + row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0) + row_offset = row_idx * req_to_token_stride_0 + + col_start = pid_c * BLOCK_COLS + col_offsets = col_start + tl.arange(0, BLOCK_COLS) + mask = col_offsets < max_seq_pages + + # Compute column indices in the source tensor (token offset) + if page_size == 1: + col_idx = col_offsets + else: + col_idx = col_offsets << SHIFT # faster than multiplication for power-of-two + + # Load page indices from req_to_token + rt_offsets = row_offset + col_idx * req_to_token_stride_1 + page_index = tl.load( + req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg" + ) + + # Compute page_table + if page_size == 1: + page_table_val = page_index + else: + page_table_val = page_index >> SHIFT + + # Store to page_table + pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1 + tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg") + + if use_swa: + swa_slot = tl.load( + full_to_swa_mapping + page_index * full_to_swa_mapping_stride_0, + mask=mask, + other=0, + cache_modifier=".cg", + ) + if page_size == 1: + swa_val = swa_slot + else: + swa_val = swa_slot >> SHIFT + swa_offsets = ( + i * swa_page_table_stride_0 + col_offsets * swa_page_table_stride_1 + ) + tl.store(swa_page_table + swa_offsets, swa_val, mask=mask, cache_modifier=".cg") + + +@triton.jit +def _fused_metadata_kernel_ps1_no_swa( + # Input tensors + seq_lens, + seq_lens_stride_0, + req_to_token, + req_to_token_stride_0, + req_to_token_stride_1, + req_pool_indices, + req_pool_indices_stride_0, + # Output buffers + cache_seqlens_int32, + cache_seqlens_int32_stride_0, + cu_seqlens_k, + cu_seqlens_k_stride_0, + page_table, + page_table_stride_0, + page_table_stride_1, + # Scalar parameters + B, + max_seq_pages, + seq_len_delta: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + pid_b = tl.program_id(0) # batch index + pid_c = tl.program_id(1) # column chunk index + + # 1. Prefix sum (only one block does it) + if pid_b == 0 and pid_c == 0: + acc = 0 + for idx in range(B): + seq = tl.load(seq_lens + idx * seq_lens_stride_0) + val = (seq + seq_len_delta).to(tl.int32) + tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val) + tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc) + acc += val + tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc) + + # 2. Gather for this batch and column chunk + if max_seq_pages == 0: + return + + i = pid_b + # Load row index for this batch (all threads in block have same i) + row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0) + row_offset = row_idx * req_to_token_stride_0 + + col_start = pid_c * BLOCK_COLS + col_offsets = col_start + tl.arange(0, BLOCK_COLS) + mask = col_offsets < max_seq_pages + + # page_size = 1: col_idx = col_offsets + rt_offsets = row_offset + col_offsets * req_to_token_stride_1 + page_index = tl.load( + req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg" + ) + + # page_table = page_index // 1 = page_index + pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1 + tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg") + + +def normal_decode_set_metadata( + cache_seqlens_int32: torch.Tensor, + cu_seqlens_k: torch.Tensor, + page_table: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + strided_indices: torch.Tensor, + max_seq_pages: torch.Tensor, + seq_lens: torch.Tensor, + seq_len_delta: int, + page_size: int, + swa_page_table: Optional[torch.Tensor] = None, + token_to_kv_pool: Optional["SWAKVPool"] = None, +): + """ + Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels: + 1. cache_seqlens = seq_lens + seq_len_delta (int64->int32 cast) + 2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum) + 3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather) + 4. page_table = page_indices // page_size (floor-divide) + 5. (optional) swa_page_table for sliding window attention + + Achieves ~5.2x speedup on H200 hardware for typical decode workloads. + """ + assert ( + page_size > 0 and (page_size & (page_size - 1)) == 0 + ), f"page_size must be a power of two, got {page_size}" + + batch_size = cache_seqlens_int32.shape[0] + device = seq_lens.device + + # Ensure contiguous memory layout for efficient Triton access + seq_lens = seq_lens.contiguous() + req_to_token = req_to_token.contiguous() + req_pool_indices = req_pool_indices.contiguous() + + # Prepare tensor strides + seq_lens_stride_0 = seq_lens.stride(0) + req_to_token_stride_0 = req_to_token.stride(0) + req_to_token_stride_1 = req_to_token.stride(1) + req_pool_indices_stride_0 = req_pool_indices.stride(0) + cache_seqlens_int32_stride_0 = cache_seqlens_int32.stride(0) + cu_seqlens_k_stride_0 = cu_seqlens_k.stride(0) + page_table_stride_0 = page_table.stride(0) + page_table_stride_1 = page_table.stride(1) + + # Check if we should use the specialized fast path for page_size=1, no SWA + use_swa = swa_page_table is not None and token_to_kv_pool is not None + + if page_size == 1 and not use_swa: + # Specialized kernel for the common case (page_size=1, no SWA) + BLOCK_COLS = 256 + if max_seq_pages == 0: + grid = (1, 1) + else: + num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS) + grid = (batch_size, num_blocks_j) + + _fused_metadata_kernel_ps1_no_swa[grid]( + seq_lens, + seq_lens_stride_0, + req_to_token, + req_to_token_stride_0, + req_to_token_stride_1, + req_pool_indices, + req_pool_indices_stride_0, + cache_seqlens_int32, + cache_seqlens_int32_stride_0, + cu_seqlens_k, + cu_seqlens_k_stride_0, + page_table, + page_table_stride_0, + page_table_stride_1, + batch_size, + max_seq_pages, + seq_len_delta, + BLOCK_COLS=BLOCK_COLS, + num_warps=8, + num_stages=3, + ) + else: + # General kernel for page_size > 1 or SWA cases + # SWA parameters + if use_swa: + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool + + assert isinstance(token_to_kv_pool, SWAKVPool) + swa_page_table = swa_page_table.contiguous() + swa_page_table_stride_0 = swa_page_table.stride(0) + swa_page_table_stride_1 = swa_page_table.stride(1) + # Extract the full_to_swa_index_mapping from token_to_kv_pool + full_to_swa_mapping = ( + token_to_kv_pool.full_to_swa_index_mapping.contiguous() + ) + full_to_swa_mapping_stride_0 = full_to_swa_mapping.stride(0) + else: + # Dummy tensors (not used) + swa_page_table = torch.empty(0, dtype=torch.int32, device=device) + swa_page_table_stride_0 = 0 + swa_page_table_stride_1 = 0 + full_to_swa_mapping = torch.empty(0, dtype=torch.int32, device=device) + full_to_swa_mapping_stride_0 = 0 + + # Kernel configuration + BLOCK_COLS = 128 + shift = (page_size).bit_length() - 1 if page_size > 1 else 0 + + if max_seq_pages == 0: + grid = (1, 1) + else: + num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS) + grid = (batch_size, num_blocks_j) + + _fused_metadata_kernel_general[grid]( + seq_lens, + seq_lens_stride_0, + req_to_token, + req_to_token_stride_0, + req_to_token_stride_1, + req_pool_indices, + req_pool_indices_stride_0, + cache_seqlens_int32, + cache_seqlens_int32_stride_0, + cu_seqlens_k, + cu_seqlens_k_stride_0, + page_table, + page_table_stride_0, + page_table_stride_1, + swa_page_table, + swa_page_table_stride_0, + swa_page_table_stride_1, + full_to_swa_mapping, + full_to_swa_mapping_stride_0, + batch_size, + max_seq_pages, + page_size, + seq_len_delta, + use_swa, + shift, + BLOCK_COLS=BLOCK_COLS, + num_warps=4, + num_stages=3, + ) diff --git a/python/sglang/srt/layers/attention/triton_ops/pad.py b/python/sglang/srt/layers/attention/triton_ops/pad.py new file mode 100644 index 000000000..dfe09ce7b --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/pad.py @@ -0,0 +1,162 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def pad_sequence_with_mask_kernel( + input_ptr, # (total_tokens, hidden) + offsets_ptr, # (B,) + lengths_ptr, # (B,) + output_ptr, # (B, max_len, hidden) + mask_ptr, # (B, max_len) + max_len, + hidden_dim, + BLOCK_M: tl.constexpr, # seq block + BLOCK_D: tl.constexpr, # hidden block +): + b = tl.program_id(0) # batch index + m = tl.program_id(1) # seq block index + + offset = tl.load(offsets_ptr + b) + length = tl.load(lengths_ptr + b) + + seq_ids = m * BLOCK_M + tl.arange(0, BLOCK_M) + hid_ids = tl.arange(0, BLOCK_D) + + seq_mask = seq_ids < max_len + valid_token = seq_ids < length + + # input index + in_token = offset + seq_ids + in_ptr = input_ptr + in_token[:, None] * hidden_dim + hid_ids[None, :] + + # output index + out_ptr = ( + output_ptr + + b * max_len * hidden_dim + + seq_ids[:, None] * hidden_dim + + hid_ids[None, :] + ) + + values = tl.load( + in_ptr, + mask=valid_token[:, None] & (hid_ids[None, :] < hidden_dim), + other=0.0, + ) + + tl.store( + out_ptr, + values, + mask=seq_mask[:, None] & (hid_ids[None, :] < hidden_dim), + ) + + # attention mask + if tl.program_id(2) == 0: + mask_out_ptr = mask_ptr + b * max_len + seq_ids + tl.store(mask_out_ptr, valid_token, mask=seq_mask) + + +def pad_sequence_with_mask( + input_emb, # (total_tokens, hidden) + offsets, # (B,) + lengths, # (B,) + max_len, +): + B = offsets.shape[0] + hidden_dim = input_emb.shape[1] + + output = torch.zeros( + (B, max_len, hidden_dim), + device=input_emb.device, + dtype=input_emb.dtype, + ) + attn_mask = torch.empty( + (B * max_len), + device=input_emb.device, + dtype=torch.bool, + ) + + BLOCK_D = triton.next_power_of_2(hidden_dim) + BLOCK_M = triton.next_power_of_2(max_len) + + grid = ( + B, + triton.cdiv(max_len, BLOCK_M), + 1, + ) + + pad_sequence_with_mask_kernel[grid]( + input_emb, + offsets, + lengths, + output, + attn_mask, + max_len, + hidden_dim, + BLOCK_M=BLOCK_M, + BLOCK_D=BLOCK_D, + ) + + return B, output, attn_mask + + +@triton.jit +def seqlens_expand_kernel( + extend_seq_lens_ptr, # [N] + seq_lens_ptr, # [N] + offsets_ptr, # [N+1] + output_ptr, # [sum(extend_seq_lens)] + N, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + + if pid >= N: + return + + qo_len = tl.load(extend_seq_lens_ptr + pid) + kv_len = tl.load(seq_lens_ptr + pid) + + start = kv_len - qo_len + 1 + out_offset = tl.load(offsets_ptr + pid) + + offs = tl.arange(0, BLOCK) + mask = offs < qo_len + + values = start + offs + tl.store(output_ptr + out_offset + offs, values, mask=mask) + + +def seqlens_expand_triton( + extend_seq_lens: torch.Tensor, + seq_lens: torch.Tensor, + total_len: int, + max_q_len: int, +): + """ + extend_seq_lens: [N], int32, CUDA + seq_lens: [N], int32, CUDA + """ + assert extend_seq_lens.is_cuda + assert seq_lens.is_cuda + + N = extend_seq_lens.numel() + + offsets = torch.zeros(N + 1, device=extend_seq_lens.device, dtype=torch.int32) + offsets[1:] = torch.cumsum(extend_seq_lens, dim=0) + output = torch.empty(total_len, device=extend_seq_lens.device, dtype=torch.int32) + + BLOCK = triton.next_power_of_2(max_q_len) + grid = (N,) + + seqlens_expand_kernel[grid]( + extend_seq_lens, + seq_lens, + offsets, + output, + N, + BLOCK=BLOCK, + ) + + return output diff --git a/python/sglang/srt/layers/attention/triton_ops/rope_cache.py b/python/sglang/srt/layers/attention/triton_ops/rope_cache.py new file mode 100644 index 000000000..49b4b8ca9 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/rope_cache.py @@ -0,0 +1,736 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _get_gptj_rotated_x( + x, + x_rotated_mask, + BLOCK_D: tl.constexpr, + BLOCK_D_HALF: tl.constexpr, +): + # GPT-J rotary layout: + # Pair adjacent dimensions and apply: + # [x0, x1, x2, x3] -> [-x1, x0, -x3, x2] + + # Apply sign inversion on odd positions. + x_rotated = tl.where(x_rotated_mask, x, -x) + # Reshape into (D/2, 2) pairs. + x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2)) + # Swap each pair. + x_rotated = tl.flip(x_rotated, 1) + # Flatten back to original shape. + x_rotated = tl.reshape(x_rotated, (BLOCK_D,)) + return x_rotated + + +@triton.jit +def _get_neox_rotated_x( + x, + x_rotated_mask, + BLOCK_D: tl.constexpr, + BLOCK_D_HALF: tl.constexpr, +): + # GPT-NeoX rotary layout: + # Split head dimension into two halves: + # [x0, x1, x2, x3] -> [-x2, -x3, x0, x1] + + # Keep first half positive, second half negative. + x_rotated = tl.where(x_rotated_mask, x, -x) + # Reshape into (2, D/2). + x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF)) + # Reverse each half. + x_rotated = tl.flip(x_rotated, 1) + # Flatten and reverse full vector. + x_rotated = tl.reshape(x_rotated, (BLOCK_D,)) + x_rotated = tl.flip(x_rotated, 0) + return x_rotated + + +@triton.jit +def _unit_rope( + x_ptrs, + cos, + sin, + d_pe_offs, + IS_NEOX: tl.constexpr, + BLOCK_D_pe: tl.constexpr, + BLOCK_D_HALF_pe: tl.constexpr, +): + # Load one full attention head vector. + x_pe = tl.load(x_ptrs) + + # Stage 1: Build rotated vector according to rotary layout. + if IS_NEOX: + x_rotated_mask = d_pe_offs < BLOCK_D_HALF_pe + x_pe_rotated = _get_neox_rotated_x( + x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe + ) + else: + x_rotated_mask = d_pe_offs % 2 == 0 + x_pe_rotated = _get_gptj_rotated_x( + x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe + ) + + # Stage 2: Apply RoPE transform: + # x' = x*cos + rotate(x)*sin + x_pe = x_pe * cos + x_pe_rotated * sin + + return x_pe + + +@triton.jit +def _load_cos_sin( + cos_sin_ptr, + pos, + d_cos_offs, + stride_t, + stride_d, + freq_dim, +): + base = pos * stride_t + cos = tl.load(cos_sin_ptr + base + d_cos_offs * stride_d) + sin = tl.load(cos_sin_ptr + base + (d_cos_offs + freq_dim) * stride_d) + return cos, sin + + +@triton.jit +def _fused_qk_rope_reshape_and_cache_kernel( + q_ptr, + k_ptr, + v_ptr, + pos_ptr, + cos_sin_ptr, + offs_ptr, + key_cache_ptr, + value_cache_ptr, + slot_mapping_ptr, + swa_slot_mapping_ptr, + q_out_ptr, + k_out_ptr, + zeros_out_ptr, + T, + T_slot, + q_stride_t, + q_stride_h, + q_stride_d, + k_stride_t, + k_stride_h, + k_stride_d, + v_stride_t, + v_stride_h, + v_stride_d, + cos_sin_stride_t, + cos_sin_stride_d, + q_out_stride_t, + q_out_stride_h, + q_out_stride_d, + k_out_stride_t, + k_out_stride_h, + k_out_stride_d, + key_cache_stride_t, + key_cache_stride_h, + key_cache_stride_d, + key_cache_stride_b, + key_cache_stride_x, + value_cache_stride_t, + value_cache_stride_h, + value_cache_stride_d, + value_cache_stride_b, + value_cache_stride_slot_chunk, + value_cache_stride_x, + zeros_out_stride_t, + zeros_out_stride_h, + zeros_out_stride_d, + k_scale_ptr, + v_scale_ptr, + QH_PER_KH: tl.constexpr, + QH: tl.constexpr, + KH: tl.constexpr, + REUSE_FREQS_FRONT_PART: tl.constexpr, + IS_NEOX: tl.constexpr, + BLOCK_D_pe: tl.constexpr, + BLOCK_D_HALF_pe: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + X_SIZE: tl.constexpr, + FLASH_LAYOUT: tl.constexpr, + VALUE_SHUFFLE_LAYOUT: tl.constexpr = False, + HAVE_POS: tl.constexpr = False, + HAVE_K_SCALE: tl.constexpr = False, + HAVE_V_SCALE: tl.constexpr = False, + HAVE_ZEROS: tl.constexpr = False, + HAS_SWA: tl.constexpr = False, +): + # ============================================================ + # Stage 0: Static stride assumptions for Triton compiler + # + # These assumptions help Triton optimize pointer arithmetic and + # simplify generated address calculations. + # ============================================================ + + tl.assume(q_stride_t >= 0) + tl.assume(q_stride_h >= 0) + tl.assume(q_stride_d >= 0) + tl.assume(k_stride_t >= 0) + tl.assume(k_stride_h >= 0) + tl.assume(k_stride_d >= 0) + tl.assume(v_stride_t >= 0) + tl.assume(v_stride_h >= 0) + tl.assume(v_stride_d >= 0) + tl.assume(cos_sin_stride_t >= 0) + tl.assume(cos_sin_stride_d >= 0) + tl.assume(q_out_stride_t >= 0) + tl.assume(q_out_stride_h >= 0) + tl.assume(q_out_stride_d >= 0) + tl.assume(k_out_stride_t >= 0) + tl.assume(k_out_stride_h >= 0) + tl.assume(k_out_stride_d >= 0) + tl.assume(key_cache_stride_t >= 0) + tl.assume(key_cache_stride_h >= 0) + tl.assume(key_cache_stride_d >= 0) + tl.assume(key_cache_stride_b >= 0) + tl.assume(key_cache_stride_x >= 0) + tl.assume(value_cache_stride_t >= 0) + tl.assume(value_cache_stride_h >= 0) + tl.assume(value_cache_stride_d >= 0) + tl.assume(value_cache_stride_b >= 0) + tl.assume(value_cache_stride_slot_chunk >= 0) + tl.assume(value_cache_stride_x >= 0) + tl.assume(zeros_out_stride_t >= 0) + tl.assume(zeros_out_stride_h >= 0) + tl.assume(zeros_out_stride_d >= 0) + + # ============================================================ + # Stage 1: Program instance mapping + # + # Each program handles: + # - one (token, q_head) for Q path + # - selected KV ownership for cache write path + # + # pid layout: + # [0, T*QH) -> decode Q path + # [T*QH, extra KV) -> KV-only path + # ============================================================ + + pid = tl.program_id(0) + tl.assume(pid >= 0) + + d_pe_offs = tl.arange(0, BLOCK_D_pe).to(tl.int64) + + # ============================================================ + # Stage 2: Main decode path (Q always active) + # ============================================================ + + if pid < T * QH: + pid_t = pid // QH + pid_hq = pid % QH + + # -------------------------------------------------------- + # Stage 2.1: Compute rotary frequency offsets + # + # RoPE frequencies may be stored as: + # D/2 frequencies (shared front-half) + # D frequencies (full explicit) + # -------------------------------------------------------- + + if REUSE_FREQS_FRONT_PART: + if IS_NEOX: + d_cos_offs = d_pe_offs + d_cos_offs = tl.where( + (d_cos_offs >= BLOCK_D_HALF_pe) & (d_cos_offs < BLOCK_D_pe), + d_cos_offs - BLOCK_D_HALF_pe, + d_cos_offs, + ).to(d_cos_offs.dtype) + # d_cos_mask = d_cos_offs < BLOCK_D_pe + else: + d_cos_offs = d_pe_offs // 2 + # d_cos_mask = d_cos_offs < BLOCK_D_HALF_pe + else: + d_cos_offs = d_pe_offs + # d_cos_mask = d_cos_offs < BLOCK_D_pe + + # -------------------------------------------------------- + # Stage 2.2: Load token position and optional offset + # + # offs_ptr is used by chunked prefill / sliding-window decode. + # -------------------------------------------------------- + pos = tl.load(pos_ptr + pid_t) + if HAVE_POS: + offset = tl.load(offs_ptr + pid_t) + pos = pos + offset + + # -------------------------------------------------------- + # Stage 2.3: Load cosine / sine table + # -------------------------------------------------------- + # cos_offs = pos * cos_stride_t + d_cos_offs * cos_stride_d + # cos = tl.load(cos_ptr + cos_offs) + # sin = tl.load(sin_ptr + cos_offs) + + freq_dim = BLOCK_D_HALF_pe if REUSE_FREQS_FRONT_PART else BLOCK_D_pe + + cos, sin = _load_cos_sin( + cos_sin_ptr, + pos, + d_cos_offs, + cos_sin_stride_t, + cos_sin_stride_d, + freq_dim, + ) + + # -------------------------------------------------------- + # Stage 2.4: Apply RoPE to Q + # -------------------------------------------------------- + q_ptrs = ( + q_ptr + pid_t * q_stride_t + pid_hq * q_stride_h + d_pe_offs * q_stride_d + ) + q_pe = _unit_rope( + q_ptrs, + cos, + sin, + d_pe_offs, + IS_NEOX, + BLOCK_D_pe, + BLOCK_D_HALF_pe, + ) + + # Store rotated Q output. + q_out_ptrs = ( + q_out_ptr + + pid_t * q_out_stride_t + + pid_hq * q_out_stride_h + + d_pe_offs * q_out_stride_d + ) + tl.store(q_out_ptrs, q_pe.to(q_out_ptr.dtype.element_ty)) + + if HAVE_ZEROS: + z = tl.zeros((BLOCK_D_pe,), dtype=zeros_out_ptr.dtype.element_ty) + zeros_out_ptrs = ( + zeros_out_ptr + + pid_t * zeros_out_stride_t + + pid_hq * zeros_out_stride_h + + d_pe_offs * zeros_out_stride_d + ) + tl.store(zeros_out_ptrs, z) + + # ======================================================== + # Stage 3: KV ownership path + # + # Only one Q group leader writes KV: + # pid_hq % QH_PER_KH == 0 + # + # This prevents duplicated KV cache writes. + # ======================================================== + + if pid_hq % QH_PER_KH == 0: + # ---------------------------------------------------- + # Stage 3.1: Resolve cache slot + # ---------------------------------------------------- + pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64) + if HAS_SWA: + pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot) + + # ------------------------------------------------ + # Stage 3.2: Apply RoPE to K + # ------------------------------------------------ + if pid_slot >= 0: + pid_t_slot = pid_slot // BLOCK_SIZE + pid_b = pid_slot % BLOCK_SIZE + pid_hk = pid_hq // QH_PER_KH + if HAVE_K_SCALE: + k_scale = tl.load(k_scale_ptr) + else: + k_scale = 1 + k_ptrs = ( + k_ptr + + pid_t * k_stride_t + + pid_hk * k_stride_h + + d_pe_offs * k_stride_d + ) + k_pe = _unit_rope( + k_ptrs, + cos, + sin, + d_pe_offs, + IS_NEOX, + BLOCK_D_pe, + BLOCK_D_HALF_pe, + ) + + k_out_ptrs = ( + k_out_ptr + + pid_t * k_out_stride_t + + pid_hk * k_out_stride_h + + d_pe_offs * k_out_stride_d + ) + tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty)) + + # ------------------------------------------------ + # Stage 3.3: Optional fp8 scaling before cache + # ------------------------------------------------ + + k_scale_rcprl = 1 / k_scale + k_pe = k_pe * k_scale_rcprl + + # ------------------------------------------------ + # Stage 3.4: Write K cache + # + # Two layouts supported: + # FLASH_LAYOUT + # paged KV layout + # ------------------------------------------------ + + if FLASH_LAYOUT: + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + pid_b * key_cache_stride_b + + pid_hk * key_cache_stride_h + + d_pe_offs * key_cache_stride_d + ) + else: + k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE)) + dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64) + x_offs = tl.arange(0, X_SIZE).to(tl.int64) + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + pid_hk * key_cache_stride_h + + dx_offs[:, None] * key_cache_stride_d + + pid_b * key_cache_stride_b + + x_offs[None, :] * key_cache_stride_x + ) + + tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty)) + + # ------------------------------------------------ + # Stage 3.5: Write V cache + # + # Supports: + # normal layout + # shuffle layout + # ------------------------------------------------ + + v_ptrs = ( + v_ptr + + pid_t * v_stride_t + + pid_hk * v_stride_h + + d_pe_offs * v_stride_d + ) + if HAVE_V_SCALE: + v_scale = tl.load(v_scale_ptr) + else: + v_scale = 1 + v_scale_rcprl = 1 / v_scale + v = tl.load(v_ptrs) * v_scale_rcprl + if VALUE_SHUFFLE_LAYOUT: + slot_chunk = pid_b // X_SIZE + x_off = pid_b % X_SIZE + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + slot_chunk * value_cache_stride_slot_chunk + + d_pe_offs.to(tl.int64) * value_cache_stride_d + + x_off * value_cache_stride_x + ) + else: + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + d_pe_offs.to(tl.int64) * value_cache_stride_d + + pid_b * value_cache_stride_b + ) + tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty)) + # ============================================================ + # Stage 4: Extra KV-only path + # + # Handles tokens that only require cache update: + # T_slot > T + # + # No Q / no RoPE on Q branch. + # ============================================================ + else: + pid = pid - T * QH + T * KH + if pid < T_slot * KH: + pid_t = pid // KH + pid_hk = pid % KH + pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64) + if HAS_SWA: + pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot) + + if pid_slot >= 0: + pid_t_slot = pid_slot // BLOCK_SIZE + pid_b = pid_slot % BLOCK_SIZE + if HAVE_K_SCALE: + k_scale = tl.load(k_scale_ptr) + else: + k_scale = 1 + k_ptrs = ( + k_ptr + + pid_t * k_stride_t + + pid_hk * k_stride_h + + d_pe_offs * k_stride_d + ) + + k_pe = tl.load(k_ptrs) + + k_out_ptrs = ( + k_out_ptr + + pid_t * k_out_stride_t + + pid_hk * k_out_stride_h + + d_pe_offs * k_out_stride_d + ) + tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty)) + + k_scale_rcprl = 1 / k_scale + k_pe = k_pe * k_scale_rcprl + + if FLASH_LAYOUT: + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + d_pe_offs * key_cache_stride_d + + pid_b * key_cache_stride_b + + pid_hk * key_cache_stride_h + ) + else: + k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE)) + dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64) + x_offs = tl.arange(0, X_SIZE).to(tl.int64) + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + pid_hk * key_cache_stride_h + + dx_offs[:, None] * key_cache_stride_d + + pid_b * key_cache_stride_b + + x_offs[None, :] * key_cache_stride_x + ) + tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty)) + + v_ptrs = ( + v_ptr + + pid_t * v_stride_t + + pid_hk * v_stride_h + + d_pe_offs * v_stride_d + ) + if HAVE_V_SCALE: + v_scale = tl.load(v_scale_ptr) + else: + v_scale = 1 + v_scale_rcprl = 1 / v_scale + v = tl.load(v_ptrs) * v_scale_rcprl + if VALUE_SHUFFLE_LAYOUT: + slot_chunk = pid_b // X_SIZE + x_off = pid_b % X_SIZE + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + slot_chunk * value_cache_stride_slot_chunk + + d_pe_offs * value_cache_stride_d + + x_off * value_cache_stride_x + ) + else: + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + d_pe_offs * value_cache_stride_d + + pid_b * value_cache_stride_b + ) + tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty)) + + +def fused_qk_rope_reshape_and_cache( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + pos: torch.Tensor, + cos_sin: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + is_neox: bool, + flash_layout: bool, + apply_scale: bool = True, + offs: torch.Tensor = None, + q_out: torch.Tensor = None, + k_out: torch.Tensor = None, + output_zeros: bool = True, + zeros_out: torch.Tensor = None, + swa_slot_mapping=None, +): + """ + Perform RoPE on q and k and along the last dimension and copy k and v in to key_cache and value_cache inplace + + Key parameters: + - q: shape (T, QH, D). + - k: shape (T_slot, KH, D). + - v: shape (T_slot, KH, D). + - if flash_layout: + - key_cache: shape (T_cache, block_size, KH, D). + - value_cache: shape (T_cache, block_size, KH, D). + - else: + - key_cache: shape (T_cache, KH, D // x, block_size, x). + - value_cache: shape (T_cache, KH, D, block_size). + - slot_mapping: shape (T_slot, ). + + T is the number of decode tokens, T_cahce * block_size is the max number of tokens of kv_cache + QH must be multiple of KH + + Returns: + - q_out: same shape as input q. + - k_out: same shape as input k. + - key_cache: same shape as input key_cache (inplace). + - value_cache: same shape as input value_cache (inplace). + - zeros_out: same shape as input q. + """ + + t, qh, d = q.shape + tk, kh, dk = k.shape + tv, vh, dv = v.shape + if flash_layout: + t_cache, block_size, kh_cache, dk_cache = key_cache.shape + t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape + value_shuffle_layout = False + else: + t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape + if value_cache.ndim == 5: + # value_cache shuffle: (num_blocks, num_kv_heads, block_size // x, head_size, x) + t_cache_v, vh_cache, slot_chunk_v, dv_cache, x_v = value_cache.shape + value_shuffle_layout = True + block_size_v = slot_chunk_v * x_v + assert block_size_v == block_size and x_v == x_cache, ( + f"value_cache shuffle (T,KH,block_size//x,D,x) must match key: " + f"{block_size_v=} {block_size=} {x_v=} {x_cache=}" + ) + else: + t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape + value_shuffle_layout = False + (t_slot,) = slot_mapping.shape + + assert ( + t == tk == tv and t_slot <= tk + ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" + assert ( + block_size == block_size_v + ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" + assert ( + kh == vh == kh_cache == vh_cache + ), "KV head should be identical for k, v, key_cache, and value_cache" + assert ( + t_cache == t_cache_v + ), "Number of tokens should be identical for key_cache, and value_cache" + if flash_layout: + assert ( + d == dk == dv == dk_cache == dv_cache + ), "D dimension should be identical for q, k, and v" + else: + assert ( + d == dk == dv == dkx_cache * x_cache == dv_cache + ), "D dimension should be identical for q, k, and v" + assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2" + + assert d == triton.next_power_of_2(d), "D dimension should be power of 2" + assert block_size == triton.next_power_of_2( + block_size + ), "block_size should be power of 2" + assert qh % kh == 0, "Q heads must be multiple of H heads" + d_freq = cos_sin.shape[-1] // 2 + assert (d_freq == d // 2) or ( + d_freq == d + ), "cos/sin last dim should be the same or half of the qk last dim" + reuse_freqs_front_part = d_freq == d // 2 + + if q_out is None: + q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) + + if k_out is None: + k_out = torch.empty((tk, kh, dk), dtype=k.dtype, device=q.device) + + if zeros_out is not None: + tz, qhz, dz = zeros_out.shape + assert ( + t == tz and qh == qhz and d == dz + ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" + output_zeros = True + elif output_zeros: + zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) + else: + zeros_out = None + + n_pid = t * qh + (t_slot - t) * kh if t_slot >= t else t * qh + grid = (n_pid, 1, 1) + _fused_qk_rope_reshape_and_cache_kernel[grid]( + q, + k, + v, + pos, + cos_sin, + offs, + key_cache, + value_cache, + slot_mapping, + swa_slot_mapping, + q_out, + k_out, + zeros_out, + t, + t_slot, + *q.stride(), + *k.stride(), + *v.stride(), + cos_sin.stride(0), + cos_sin.stride(-1), + *q_out.stride(), + *k_out.stride(), + key_cache.stride(0) if not flash_layout else key_cache.stride(0), + key_cache.stride(1) if not flash_layout else key_cache.stride(2), + key_cache.stride(2) if not flash_layout else key_cache.stride(3), + key_cache.stride(3) if not flash_layout else key_cache.stride(1), + key_cache.stride(4) if not flash_layout else 0, + value_cache.stride(0) if not flash_layout else value_cache.stride(0), + value_cache.stride(1) if not flash_layout else value_cache.stride(2), + ( + value_cache.stride(3) + if (not flash_layout and value_shuffle_layout) + else (value_cache.stride(2) if not flash_layout else value_cache.stride(3)) + ), + ( + 0 + if (not flash_layout and value_shuffle_layout) + else (value_cache.stride(3) if not flash_layout else value_cache.stride(1)) + ), + value_cache.stride(2) if (not flash_layout and value_shuffle_layout) else 0, + value_cache.stride(4) if (not flash_layout and value_shuffle_layout) else 0, + zeros_out.stride(0) if zeros_out is not None else 0, + zeros_out.stride(1) if zeros_out is not None else 0, + zeros_out.stride(2) if zeros_out is not None else 0, + k_scale_ptr=k_scale, + v_scale_ptr=v_scale, + QH_PER_KH=qh // kh, + QH=qh, + KH=kh, + REUSE_FREQS_FRONT_PART=reuse_freqs_front_part, + IS_NEOX=is_neox, + BLOCK_D_pe=d, + BLOCK_D_HALF_pe=d // 2, + BLOCK_SIZE=block_size, + X_SIZE=x_cache if not flash_layout else 0, + FLASH_LAYOUT=flash_layout, + VALUE_SHUFFLE_LAYOUT=value_shuffle_layout, + HAVE_POS=(offs is not None), + HAVE_K_SCALE=(k_scale is not None and apply_scale), + HAVE_V_SCALE=(v_scale is not None and apply_scale), + HAVE_ZEROS=output_zeros, + HAS_SWA=(swa_slot_mapping is not None), + num_warps=1, + ) + + if zeros_out is not None: + return q_out.view(-1, qh * d), k_out, key_cache, value_cache, zeros_out + return q_out.view(-1, qh * d), k_out, key_cache, value_cache diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index 65328b16b..1217f93e5 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -1,353 +1,49 @@ import torch -import triton -import triton.language as tl +from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.srt.layers.attention.triton_ops.cache_ops import ( + concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel, +) +from sglang.srt.layers.attention.triton_ops.cache_ops import ( + concat_and_cast_mha_k_triton as concat_and_cast_mha_k_triton, +) +from sglang.srt.layers.attention.triton_ops.cache_ops import ( + launch_reshape_and_cache_flash as launch_reshape_and_cache_flash, +) +from sglang.srt.layers.attention.triton_ops.cache_ops import ( + reshape_and_cache_flash as reshape_and_cache_flash, +) +from sglang.srt.layers.attention.triton_ops.kv_indices import ( + create_flashinfer_kv_indices_triton as create_flashinfer_kv_indices_triton, +) +from sglang.srt.layers.attention.triton_ops.kv_indices import ( + create_flashmla_kv_indices_triton as create_flashmla_kv_indices_triton, +) +from sglang.srt.layers.attention.triton_ops.kv_indices import ( + get_num_page_per_block_flashmla as get_num_page_per_block_flashmla, +) +from sglang.srt.layers.attention.triton_ops.pad import ( + pad_sequence_with_mask as pad_sequence_with_mask, +) +from sglang.srt.layers.attention.triton_ops.pad import ( + pad_sequence_with_mask_kernel as pad_sequence_with_mask_kernel, +) +from sglang.srt.layers.attention.triton_ops.pad import ( + seqlens_expand_kernel as seqlens_expand_kernel, +) +from sglang.srt.layers.attention.triton_ops.pad import ( + seqlens_expand_triton as seqlens_expand_triton, +) +from sglang.srt.layers.attention.triton_ops.rope_cache import ( + fused_qk_rope_reshape_and_cache as fused_qk_rope_reshape_and_cache, +) from sglang.srt.utils import is_cuda -_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096 -FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE) - _is_cuda = is_cuda() if _is_cuda: from sglang.jit_kernel.concat_mla import concat_mla_absorb_q -from sglang.jit_kernel.utils import is_arch_support_pdl - - -@triton.jit -def create_flashinfer_kv_indices_triton( - req_to_token_ptr, # [max_batch, max_context_len] - req_pool_indices_ptr, - page_kernel_lens_ptr, - kv_indptr, - kv_start_idx, - kv_indices_ptr, - req_to_token_ptr_stride: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 512 - pid = tl.program_id(axis=0) - - # find the req pool idx, this is for batch to token - req_pool_index = tl.load(req_pool_indices_ptr + pid) - kv_indices_offset = tl.load(kv_indptr + pid) - - kv_start = 0 - kv_end = 0 - if kv_start_idx: - kv_start = tl.load(kv_start_idx + pid).to(tl.int32) - kv_end = kv_start - kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) - - num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) - for i in range(num_loop): - # index into req_to_token_ptr needs to be int64 - offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE - mask = offset < kv_end - kv_start - data = tl.load( - req_to_token_ptr - + req_pool_index * req_to_token_ptr_stride - + kv_start - + offset, - mask=mask, - ) - tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask) - - -def get_num_page_per_block_flashmla(page_size: int = 64) -> int: - num_page_per_block = _FLASHMLA_CREATE_KV_BLOCK_SIZE // page_size - return num_page_per_block - - -@triton.jit -def create_flashmla_kv_indices_triton( - req_to_token_ptr, # [max_batch, max_context_len] - req_pool_indices_ptr, - page_kernel_lens_ptr, - kv_start_idx, - kv_indices_ptr, - req_to_token_ptr_stride: tl.constexpr, - kv_indices_ptr_stride: tl.constexpr, - PAGED_SIZE: tl.constexpr = 64, -): - NUM_PAGE_PER_BLOCK: tl.constexpr = ( - FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE - ) - pid = tl.program_id(axis=0) - - # find the req pool idx, this is for batch to token - req_pool_index = tl.load(req_pool_indices_ptr + pid) - - kv_start = 0 - kv_end = 0 - if kv_start_idx: - kv_start = tl.load(kv_start_idx + pid).to(tl.int32) - kv_end = kv_start - - kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) - - num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE) - num_pages_loop = tl.cdiv(kv_end - kv_start, FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON) - - for i in range(num_pages_loop): - # index into req_to_token_ptr needs to be int64 - paged_offset = ( - tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK - ) * PAGED_SIZE - paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK - - mask = paged_offset < num_paged * PAGED_SIZE - mask_out = paged_offset_out < num_paged - - data = tl.load( - req_to_token_ptr - + req_pool_index * req_to_token_ptr_stride - + kv_start - + paged_offset, - mask=mask, - ) - tl.store( - kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out, - data // PAGED_SIZE, - mask=mask_out, - ) - - -@triton.jit -def concat_and_cast_mha_k_kernel( - k_ptr, - k_nope_ptr, - k_rope_ptr, - head_cnt: tl.constexpr, - k_stride0: tl.constexpr, - k_stride1: tl.constexpr, - nope_stride0: tl.constexpr, - nope_stride1: tl.constexpr, - rope_stride0: tl.constexpr, - nope_dim: tl.constexpr, - rope_dim: tl.constexpr, -): - pid_loc = tl.program_id(0) - head_range = tl.arange(0, head_cnt) - - k_head_ptr = k_ptr + pid_loc * k_stride0 + head_range[:, None] * k_stride1 - - nope_offs = tl.arange(0, nope_dim) - - src_nope_ptr = ( - k_nope_ptr - + pid_loc * nope_stride0 - + head_range[:, None] * nope_stride1 - + nope_offs[None, :] - ) - dst_nope_ptr = k_head_ptr + nope_offs[None, :] - - src_nope = tl.load(src_nope_ptr) - tl.store(dst_nope_ptr, src_nope) - - rope_offs = tl.arange(0, rope_dim) - src_rope_ptr = k_rope_ptr + pid_loc * rope_stride0 + rope_offs[None, :] - dst_rope_ptr = k_head_ptr + nope_dim + rope_offs[None, :] - src_rope = tl.load(src_rope_ptr) - tl.store(dst_rope_ptr, src_rope) - - -def concat_and_cast_mha_k_triton( - k: torch.Tensor, - k_nope: torch.Tensor, - k_rope: torch.Tensor, -): - # The source data type will be implicitly converted to the target data type. - assert ( - len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3 - ), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - - nope_dim = k_nope.shape[-1] - rope_dim = k_rope.shape[-1] - grid = (k.shape[0],) - - concat_and_cast_mha_k_kernel[grid]( - k, - k_nope, - k_rope, - k.shape[1], - k.stride(0), - k.stride(1), - k_nope.stride(0), - k_nope.stride(1), - k_rope.stride(0), - nope_dim, - rope_dim, - ) - - -@triton.jit -def pad_sequence_with_mask_kernel( - input_ptr, # (total_tokens, hidden) - offsets_ptr, # (B,) - lengths_ptr, # (B,) - output_ptr, # (B, max_len, hidden) - mask_ptr, # (B, max_len) - max_len, - hidden_dim, - BLOCK_M: tl.constexpr, # seq block - BLOCK_D: tl.constexpr, # hidden block -): - b = tl.program_id(0) # batch index - m = tl.program_id(1) # seq block index - - offset = tl.load(offsets_ptr + b) - length = tl.load(lengths_ptr + b) - - seq_ids = m * BLOCK_M + tl.arange(0, BLOCK_M) - hid_ids = tl.arange(0, BLOCK_D) - - seq_mask = seq_ids < max_len - valid_token = seq_ids < length - - # input index - in_token = offset + seq_ids - in_ptr = input_ptr + in_token[:, None] * hidden_dim + hid_ids[None, :] - - # output index - out_ptr = ( - output_ptr - + b * max_len * hidden_dim - + seq_ids[:, None] * hidden_dim - + hid_ids[None, :] - ) - - values = tl.load( - in_ptr, - mask=valid_token[:, None] & (hid_ids[None, :] < hidden_dim), - other=0.0, - ) - - tl.store( - out_ptr, - values, - mask=seq_mask[:, None] & (hid_ids[None, :] < hidden_dim), - ) - - # attention mask - if tl.program_id(2) == 0: - mask_out_ptr = mask_ptr + b * max_len + seq_ids - tl.store(mask_out_ptr, valid_token, mask=seq_mask) - - -def pad_sequence_with_mask( - input_emb, # (total_tokens, hidden) - offsets, # (B,) - lengths, # (B,) - max_len, -): - B = offsets.shape[0] - hidden_dim = input_emb.shape[1] - - output = torch.zeros( - (B, max_len, hidden_dim), - device=input_emb.device, - dtype=input_emb.dtype, - ) - attn_mask = torch.empty( - (B * max_len), - device=input_emb.device, - dtype=torch.bool, - ) - - BLOCK_D = triton.next_power_of_2(hidden_dim) - BLOCK_M = triton.next_power_of_2(max_len) - - grid = ( - B, - triton.cdiv(max_len, BLOCK_M), - 1, - ) - - pad_sequence_with_mask_kernel[grid]( - input_emb, - offsets, - lengths, - output, - attn_mask, - max_len, - hidden_dim, - BLOCK_M=BLOCK_M, - BLOCK_D=BLOCK_D, - ) - - return B, output, attn_mask - - -@triton.jit -def seqlens_expand_kernel( - extend_seq_lens_ptr, # [N] - seq_lens_ptr, # [N] - offsets_ptr, # [N+1] - output_ptr, # [sum(extend_seq_lens)] - N, - BLOCK: tl.constexpr, -): - pid = tl.program_id(0) - - if pid >= N: - return - - qo_len = tl.load(extend_seq_lens_ptr + pid) - kv_len = tl.load(seq_lens_ptr + pid) - - start = kv_len - qo_len + 1 - out_offset = tl.load(offsets_ptr + pid) - - offs = tl.arange(0, BLOCK) - mask = offs < qo_len - - values = start + offs - tl.store(output_ptr + out_offset + offs, values, mask=mask) - - -def seqlens_expand_triton( - extend_seq_lens: torch.Tensor, - seq_lens: torch.Tensor, - total_len: int, - max_q_len: int, -): - """ - extend_seq_lens: [N], int32, CUDA - seq_lens: [N], int32, CUDA - """ - assert extend_seq_lens.is_cuda - assert seq_lens.is_cuda - - N = extend_seq_lens.numel() - - offsets = torch.zeros(N + 1, device=extend_seq_lens.device, dtype=torch.int32) - offsets[1:] = torch.cumsum(extend_seq_lens, dim=0) - output = torch.empty(total_len, device=extend_seq_lens.device, dtype=torch.int32) - - BLOCK = triton.next_power_of_2(max_q_len) - grid = (N,) - - seqlens_expand_kernel[grid]( - extend_seq_lens, - seq_lens, - offsets, - output, - N, - BLOCK=BLOCK, - ) - - return output - # When num_kv_heads=1, we have tensors with degenerate strides, # For example, as below, where we have stride[-3] == stride[-2]: @@ -475,922 +171,3 @@ 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, - 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), - ) - - -@triton.jit -def _get_gptj_rotated_x( - x, - x_rotated_mask, - BLOCK_D: tl.constexpr, - BLOCK_D_HALF: tl.constexpr, -): - # GPT-J rotary layout: - # Pair adjacent dimensions and apply: - # [x0, x1, x2, x3] -> [-x1, x0, -x3, x2] - - # Apply sign inversion on odd positions. - x_rotated = tl.where(x_rotated_mask, x, -x) - # Reshape into (D/2, 2) pairs. - x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2)) - # Swap each pair. - x_rotated = tl.flip(x_rotated, 1) - # Flatten back to original shape. - x_rotated = tl.reshape(x_rotated, (BLOCK_D,)) - return x_rotated - - -@triton.jit -def _get_neox_rotated_x( - x, - x_rotated_mask, - BLOCK_D: tl.constexpr, - BLOCK_D_HALF: tl.constexpr, -): - # GPT-NeoX rotary layout: - # Split head dimension into two halves: - # [x0, x1, x2, x3] -> [-x2, -x3, x0, x1] - - # Keep first half positive, second half negative. - x_rotated = tl.where(x_rotated_mask, x, -x) - # Reshape into (2, D/2). - x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF)) - # Reverse each half. - x_rotated = tl.flip(x_rotated, 1) - # Flatten and reverse full vector. - x_rotated = tl.reshape(x_rotated, (BLOCK_D,)) - x_rotated = tl.flip(x_rotated, 0) - return x_rotated - - -@triton.jit -def _unit_rope( - x_ptrs, - cos, - sin, - d_pe_offs, - IS_NEOX: tl.constexpr, - BLOCK_D_pe: tl.constexpr, - BLOCK_D_HALF_pe: tl.constexpr, -): - # Load one full attention head vector. - x_pe = tl.load(x_ptrs) - - # Stage 1: Build rotated vector according to rotary layout. - if IS_NEOX: - x_rotated_mask = d_pe_offs < BLOCK_D_HALF_pe - x_pe_rotated = _get_neox_rotated_x( - x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe - ) - else: - x_rotated_mask = d_pe_offs % 2 == 0 - x_pe_rotated = _get_gptj_rotated_x( - x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe - ) - - # Stage 2: Apply RoPE transform: - # x' = x*cos + rotate(x)*sin - x_pe = x_pe * cos + x_pe_rotated * sin - - return x_pe - - -@triton.jit -def _load_cos_sin( - cos_sin_ptr, - pos, - d_cos_offs, - stride_t, - stride_d, - freq_dim, -): - base = pos * stride_t - cos = tl.load(cos_sin_ptr + base + d_cos_offs * stride_d) - sin = tl.load(cos_sin_ptr + base + (d_cos_offs + freq_dim) * stride_d) - return cos, sin - - -@triton.jit -def _fused_qk_rope_reshape_and_cache_kernel( - q_ptr, - k_ptr, - v_ptr, - pos_ptr, - cos_sin_ptr, - offs_ptr, - key_cache_ptr, - value_cache_ptr, - slot_mapping_ptr, - swa_slot_mapping_ptr, - q_out_ptr, - k_out_ptr, - zeros_out_ptr, - T, - T_slot, - q_stride_t, - q_stride_h, - q_stride_d, - k_stride_t, - k_stride_h, - k_stride_d, - v_stride_t, - v_stride_h, - v_stride_d, - cos_sin_stride_t, - cos_sin_stride_d, - q_out_stride_t, - q_out_stride_h, - q_out_stride_d, - k_out_stride_t, - k_out_stride_h, - k_out_stride_d, - key_cache_stride_t, - key_cache_stride_h, - key_cache_stride_d, - key_cache_stride_b, - key_cache_stride_x, - value_cache_stride_t, - value_cache_stride_h, - value_cache_stride_d, - value_cache_stride_b, - value_cache_stride_slot_chunk, - value_cache_stride_x, - zeros_out_stride_t, - zeros_out_stride_h, - zeros_out_stride_d, - k_scale_ptr, - v_scale_ptr, - QH_PER_KH: tl.constexpr, - QH: tl.constexpr, - KH: tl.constexpr, - REUSE_FREQS_FRONT_PART: tl.constexpr, - IS_NEOX: tl.constexpr, - BLOCK_D_pe: tl.constexpr, - BLOCK_D_HALF_pe: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - X_SIZE: tl.constexpr, - FLASH_LAYOUT: tl.constexpr, - VALUE_SHUFFLE_LAYOUT: tl.constexpr = False, - HAVE_POS: tl.constexpr = False, - HAVE_K_SCALE: tl.constexpr = False, - HAVE_V_SCALE: tl.constexpr = False, - HAVE_ZEROS: tl.constexpr = False, - HAS_SWA: tl.constexpr = False, -): - # ============================================================ - # Stage 0: Static stride assumptions for Triton compiler - # - # These assumptions help Triton optimize pointer arithmetic and - # simplify generated address calculations. - # ============================================================ - - tl.assume(q_stride_t >= 0) - tl.assume(q_stride_h >= 0) - tl.assume(q_stride_d >= 0) - tl.assume(k_stride_t >= 0) - tl.assume(k_stride_h >= 0) - tl.assume(k_stride_d >= 0) - tl.assume(v_stride_t >= 0) - tl.assume(v_stride_h >= 0) - tl.assume(v_stride_d >= 0) - tl.assume(cos_sin_stride_t >= 0) - tl.assume(cos_sin_stride_d >= 0) - tl.assume(q_out_stride_t >= 0) - tl.assume(q_out_stride_h >= 0) - tl.assume(q_out_stride_d >= 0) - tl.assume(k_out_stride_t >= 0) - tl.assume(k_out_stride_h >= 0) - tl.assume(k_out_stride_d >= 0) - tl.assume(key_cache_stride_t >= 0) - tl.assume(key_cache_stride_h >= 0) - tl.assume(key_cache_stride_d >= 0) - tl.assume(key_cache_stride_b >= 0) - tl.assume(key_cache_stride_x >= 0) - tl.assume(value_cache_stride_t >= 0) - tl.assume(value_cache_stride_h >= 0) - tl.assume(value_cache_stride_d >= 0) - tl.assume(value_cache_stride_b >= 0) - tl.assume(value_cache_stride_slot_chunk >= 0) - tl.assume(value_cache_stride_x >= 0) - tl.assume(zeros_out_stride_t >= 0) - tl.assume(zeros_out_stride_h >= 0) - tl.assume(zeros_out_stride_d >= 0) - - # ============================================================ - # Stage 1: Program instance mapping - # - # Each program handles: - # - one (token, q_head) for Q path - # - selected KV ownership for cache write path - # - # pid layout: - # [0, T*QH) -> decode Q path - # [T*QH, extra KV) -> KV-only path - # ============================================================ - - pid = tl.program_id(0) - tl.assume(pid >= 0) - - d_pe_offs = tl.arange(0, BLOCK_D_pe).to(tl.int64) - - # ============================================================ - # Stage 2: Main decode path (Q always active) - # ============================================================ - - if pid < T * QH: - pid_t = pid // QH - pid_hq = pid % QH - - # -------------------------------------------------------- - # Stage 2.1: Compute rotary frequency offsets - # - # RoPE frequencies may be stored as: - # D/2 frequencies (shared front-half) - # D frequencies (full explicit) - # -------------------------------------------------------- - - if REUSE_FREQS_FRONT_PART: - if IS_NEOX: - d_cos_offs = d_pe_offs - d_cos_offs = tl.where( - (d_cos_offs >= BLOCK_D_HALF_pe) & (d_cos_offs < BLOCK_D_pe), - d_cos_offs - BLOCK_D_HALF_pe, - d_cos_offs, - ).to(d_cos_offs.dtype) - # d_cos_mask = d_cos_offs < BLOCK_D_pe - else: - d_cos_offs = d_pe_offs // 2 - # d_cos_mask = d_cos_offs < BLOCK_D_HALF_pe - else: - d_cos_offs = d_pe_offs - # d_cos_mask = d_cos_offs < BLOCK_D_pe - - # -------------------------------------------------------- - # Stage 2.2: Load token position and optional offset - # - # offs_ptr is used by chunked prefill / sliding-window decode. - # -------------------------------------------------------- - pos = tl.load(pos_ptr + pid_t) - if HAVE_POS: - offset = tl.load(offs_ptr + pid_t) - pos = pos + offset - - # -------------------------------------------------------- - # Stage 2.3: Load cosine / sine table - # -------------------------------------------------------- - # cos_offs = pos * cos_stride_t + d_cos_offs * cos_stride_d - # cos = tl.load(cos_ptr + cos_offs) - # sin = tl.load(sin_ptr + cos_offs) - - freq_dim = BLOCK_D_HALF_pe if REUSE_FREQS_FRONT_PART else BLOCK_D_pe - - cos, sin = _load_cos_sin( - cos_sin_ptr, - pos, - d_cos_offs, - cos_sin_stride_t, - cos_sin_stride_d, - freq_dim, - ) - - # -------------------------------------------------------- - # Stage 2.4: Apply RoPE to Q - # -------------------------------------------------------- - q_ptrs = ( - q_ptr + pid_t * q_stride_t + pid_hq * q_stride_h + d_pe_offs * q_stride_d - ) - q_pe = _unit_rope( - q_ptrs, - cos, - sin, - d_pe_offs, - IS_NEOX, - BLOCK_D_pe, - BLOCK_D_HALF_pe, - ) - - # Store rotated Q output. - q_out_ptrs = ( - q_out_ptr - + pid_t * q_out_stride_t - + pid_hq * q_out_stride_h - + d_pe_offs * q_out_stride_d - ) - tl.store(q_out_ptrs, q_pe.to(q_out_ptr.dtype.element_ty)) - - if HAVE_ZEROS: - z = tl.zeros((BLOCK_D_pe,), dtype=zeros_out_ptr.dtype.element_ty) - zeros_out_ptrs = ( - zeros_out_ptr - + pid_t * zeros_out_stride_t - + pid_hq * zeros_out_stride_h - + d_pe_offs * zeros_out_stride_d - ) - tl.store(zeros_out_ptrs, z) - - # ======================================================== - # Stage 3: KV ownership path - # - # Only one Q group leader writes KV: - # pid_hq % QH_PER_KH == 0 - # - # This prevents duplicated KV cache writes. - # ======================================================== - - if pid_hq % QH_PER_KH == 0: - # ---------------------------------------------------- - # Stage 3.1: Resolve cache slot - # ---------------------------------------------------- - pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64) - if HAS_SWA: - pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot) - - # ------------------------------------------------ - # Stage 3.2: Apply RoPE to K - # ------------------------------------------------ - if pid_slot >= 0: - pid_t_slot = pid_slot // BLOCK_SIZE - pid_b = pid_slot % BLOCK_SIZE - pid_hk = pid_hq // QH_PER_KH - if HAVE_K_SCALE: - k_scale = tl.load(k_scale_ptr) - else: - k_scale = 1 - k_ptrs = ( - k_ptr - + pid_t * k_stride_t - + pid_hk * k_stride_h - + d_pe_offs * k_stride_d - ) - k_pe = _unit_rope( - k_ptrs, - cos, - sin, - d_pe_offs, - IS_NEOX, - BLOCK_D_pe, - BLOCK_D_HALF_pe, - ) - - k_out_ptrs = ( - k_out_ptr - + pid_t * k_out_stride_t - + pid_hk * k_out_stride_h - + d_pe_offs * k_out_stride_d - ) - tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty)) - - # ------------------------------------------------ - # Stage 3.3: Optional fp8 scaling before cache - # ------------------------------------------------ - - k_scale_rcprl = 1 / k_scale - k_pe = k_pe * k_scale_rcprl - - # ------------------------------------------------ - # Stage 3.4: Write K cache - # - # Two layouts supported: - # FLASH_LAYOUT - # paged KV layout - # ------------------------------------------------ - - if FLASH_LAYOUT: - k_out_ptrs = ( - key_cache_ptr - + pid_t_slot * key_cache_stride_t - + pid_b * key_cache_stride_b - + pid_hk * key_cache_stride_h - + d_pe_offs * key_cache_stride_d - ) - else: - k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE)) - dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64) - x_offs = tl.arange(0, X_SIZE).to(tl.int64) - k_out_ptrs = ( - key_cache_ptr - + pid_t_slot * key_cache_stride_t - + pid_hk * key_cache_stride_h - + dx_offs[:, None] * key_cache_stride_d - + pid_b * key_cache_stride_b - + x_offs[None, :] * key_cache_stride_x - ) - - tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty)) - - # ------------------------------------------------ - # Stage 3.5: Write V cache - # - # Supports: - # normal layout - # shuffle layout - # ------------------------------------------------ - - v_ptrs = ( - v_ptr - + pid_t * v_stride_t - + pid_hk * v_stride_h - + d_pe_offs * v_stride_d - ) - if HAVE_V_SCALE: - v_scale = tl.load(v_scale_ptr) - else: - v_scale = 1 - v_scale_rcprl = 1 / v_scale - v = tl.load(v_ptrs) * v_scale_rcprl - if VALUE_SHUFFLE_LAYOUT: - slot_chunk = pid_b // X_SIZE - x_off = pid_b % X_SIZE - v_out_ptrs = ( - value_cache_ptr - + pid_t_slot * value_cache_stride_t - + pid_hk * value_cache_stride_h - + slot_chunk * value_cache_stride_slot_chunk - + d_pe_offs.to(tl.int64) * value_cache_stride_d - + x_off * value_cache_stride_x - ) - else: - v_out_ptrs = ( - value_cache_ptr - + pid_t_slot * value_cache_stride_t - + pid_hk * value_cache_stride_h - + d_pe_offs.to(tl.int64) * value_cache_stride_d - + pid_b * value_cache_stride_b - ) - tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty)) - # ============================================================ - # Stage 4: Extra KV-only path - # - # Handles tokens that only require cache update: - # T_slot > T - # - # No Q / no RoPE on Q branch. - # ============================================================ - else: - pid = pid - T * QH + T * KH - if pid < T_slot * KH: - pid_t = pid // KH - pid_hk = pid % KH - pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64) - if HAS_SWA: - pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot) - - if pid_slot >= 0: - pid_t_slot = pid_slot // BLOCK_SIZE - pid_b = pid_slot % BLOCK_SIZE - if HAVE_K_SCALE: - k_scale = tl.load(k_scale_ptr) - else: - k_scale = 1 - k_ptrs = ( - k_ptr - + pid_t * k_stride_t - + pid_hk * k_stride_h - + d_pe_offs * k_stride_d - ) - - k_pe = tl.load(k_ptrs) - - k_out_ptrs = ( - k_out_ptr - + pid_t * k_out_stride_t - + pid_hk * k_out_stride_h - + d_pe_offs * k_out_stride_d - ) - tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty)) - - k_scale_rcprl = 1 / k_scale - k_pe = k_pe * k_scale_rcprl - - if FLASH_LAYOUT: - k_out_ptrs = ( - key_cache_ptr - + pid_t_slot * key_cache_stride_t - + d_pe_offs * key_cache_stride_d - + pid_b * key_cache_stride_b - + pid_hk * key_cache_stride_h - ) - else: - k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE)) - dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64) - x_offs = tl.arange(0, X_SIZE).to(tl.int64) - k_out_ptrs = ( - key_cache_ptr - + pid_t_slot * key_cache_stride_t - + pid_hk * key_cache_stride_h - + dx_offs[:, None] * key_cache_stride_d - + pid_b * key_cache_stride_b - + x_offs[None, :] * key_cache_stride_x - ) - tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty)) - - v_ptrs = ( - v_ptr - + pid_t * v_stride_t - + pid_hk * v_stride_h - + d_pe_offs * v_stride_d - ) - if HAVE_V_SCALE: - v_scale = tl.load(v_scale_ptr) - else: - v_scale = 1 - v_scale_rcprl = 1 / v_scale - v = tl.load(v_ptrs) * v_scale_rcprl - if VALUE_SHUFFLE_LAYOUT: - slot_chunk = pid_b // X_SIZE - x_off = pid_b % X_SIZE - v_out_ptrs = ( - value_cache_ptr - + pid_t_slot * value_cache_stride_t - + pid_hk * value_cache_stride_h - + slot_chunk * value_cache_stride_slot_chunk - + d_pe_offs * value_cache_stride_d - + x_off * value_cache_stride_x - ) - else: - v_out_ptrs = ( - value_cache_ptr - + pid_t_slot * value_cache_stride_t - + pid_hk * value_cache_stride_h - + d_pe_offs * value_cache_stride_d - + pid_b * value_cache_stride_b - ) - tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty)) - - -def fused_qk_rope_reshape_and_cache( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - slot_mapping: torch.Tensor, - pos: torch.Tensor, - cos_sin: torch.Tensor, - k_scale: torch.Tensor, - v_scale: torch.Tensor, - is_neox: bool, - flash_layout: bool, - apply_scale: bool = True, - offs: torch.Tensor = None, - q_out: torch.Tensor = None, - k_out: torch.Tensor = None, - output_zeros: bool = True, - zeros_out: torch.Tensor = None, - swa_slot_mapping=None, -): - """ - Perform RoPE on q and k and along the last dimension and copy k and v in to key_cache and value_cache inplace - - Key parameters: - - q: shape (T, QH, D). - - k: shape (T_slot, KH, D). - - v: shape (T_slot, KH, D). - - if flash_layout: - - key_cache: shape (T_cache, block_size, KH, D). - - value_cache: shape (T_cache, block_size, KH, D). - - else: - - key_cache: shape (T_cache, KH, D // x, block_size, x). - - value_cache: shape (T_cache, KH, D, block_size). - - slot_mapping: shape (T_slot, ). - - T is the number of decode tokens, T_cahce * block_size is the max number of tokens of kv_cache - QH must be multiple of KH - - Returns: - - q_out: same shape as input q. - - k_out: same shape as input k. - - key_cache: same shape as input key_cache (inplace). - - value_cache: same shape as input value_cache (inplace). - - zeros_out: same shape as input q. - """ - - t, qh, d = q.shape - tk, kh, dk = k.shape - tv, vh, dv = v.shape - if flash_layout: - t_cache, block_size, kh_cache, dk_cache = key_cache.shape - t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape - value_shuffle_layout = False - else: - t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape - if value_cache.ndim == 5: - # value_cache shuffle: (num_blocks, num_kv_heads, block_size // x, head_size, x) - t_cache_v, vh_cache, slot_chunk_v, dv_cache, x_v = value_cache.shape - value_shuffle_layout = True - block_size_v = slot_chunk_v * x_v - assert block_size_v == block_size and x_v == x_cache, ( - f"value_cache shuffle (T,KH,block_size//x,D,x) must match key: " - f"{block_size_v=} {block_size=} {x_v=} {x_cache=}" - ) - else: - t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape - value_shuffle_layout = False - (t_slot,) = slot_mapping.shape - - assert ( - t == tk == tv and t_slot <= tk - ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" - assert ( - block_size == block_size_v - ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" - assert ( - kh == vh == kh_cache == vh_cache - ), "KV head should be identical for k, v, key_cache, and value_cache" - assert ( - t_cache == t_cache_v - ), "Number of tokens should be identical for key_cache, and value_cache" - if flash_layout: - assert ( - d == dk == dv == dk_cache == dv_cache - ), "D dimension should be identical for q, k, and v" - else: - assert ( - d == dk == dv == dkx_cache * x_cache == dv_cache - ), "D dimension should be identical for q, k, and v" - assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2" - - assert d == triton.next_power_of_2(d), "D dimension should be power of 2" - assert block_size == triton.next_power_of_2( - block_size - ), "block_size should be power of 2" - assert qh % kh == 0, "Q heads must be multiple of H heads" - d_freq = cos_sin.shape[-1] // 2 - assert (d_freq == d // 2) or ( - d_freq == d - ), "cos/sin last dim should be the same or half of the qk last dim" - reuse_freqs_front_part = d_freq == d // 2 - - if q_out is None: - q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) - - if k_out is None: - k_out = torch.empty((tk, kh, dk), dtype=k.dtype, device=q.device) - - if zeros_out is not None: - tz, qhz, dz = zeros_out.shape - assert ( - t == tz and qh == qhz and d == dz - ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" - output_zeros = True - elif output_zeros: - zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) - else: - zeros_out = None - - n_pid = t * qh + (t_slot - t) * kh if t_slot >= t else t * qh - grid = (n_pid, 1, 1) - _fused_qk_rope_reshape_and_cache_kernel[grid]( - q, - k, - v, - pos, - cos_sin, - offs, - key_cache, - value_cache, - slot_mapping, - swa_slot_mapping, - q_out, - k_out, - zeros_out, - t, - t_slot, - *q.stride(), - *k.stride(), - *v.stride(), - cos_sin.stride(0), - cos_sin.stride(-1), - *q_out.stride(), - *k_out.stride(), - key_cache.stride(0) if not flash_layout else key_cache.stride(0), - key_cache.stride(1) if not flash_layout else key_cache.stride(2), - key_cache.stride(2) if not flash_layout else key_cache.stride(3), - key_cache.stride(3) if not flash_layout else key_cache.stride(1), - key_cache.stride(4) if not flash_layout else 0, - value_cache.stride(0) if not flash_layout else value_cache.stride(0), - value_cache.stride(1) if not flash_layout else value_cache.stride(2), - ( - value_cache.stride(3) - if (not flash_layout and value_shuffle_layout) - else (value_cache.stride(2) if not flash_layout else value_cache.stride(3)) - ), - ( - 0 - if (not flash_layout and value_shuffle_layout) - else (value_cache.stride(3) if not flash_layout else value_cache.stride(1)) - ), - value_cache.stride(2) if (not flash_layout and value_shuffle_layout) else 0, - value_cache.stride(4) if (not flash_layout and value_shuffle_layout) else 0, - zeros_out.stride(0) if zeros_out is not None else 0, - zeros_out.stride(1) if zeros_out is not None else 0, - zeros_out.stride(2) if zeros_out is not None else 0, - k_scale_ptr=k_scale, - v_scale_ptr=v_scale, - QH_PER_KH=qh // kh, - QH=qh, - KH=kh, - REUSE_FREQS_FRONT_PART=reuse_freqs_front_part, - IS_NEOX=is_neox, - BLOCK_D_pe=d, - BLOCK_D_HALF_pe=d // 2, - BLOCK_SIZE=block_size, - X_SIZE=x_cache if not flash_layout else 0, - FLASH_LAYOUT=flash_layout, - VALUE_SHUFFLE_LAYOUT=value_shuffle_layout, - HAVE_POS=(offs is not None), - HAVE_K_SCALE=(k_scale is not None and apply_scale), - HAVE_V_SCALE=(v_scale is not None and apply_scale), - HAVE_ZEROS=output_zeros, - HAS_SWA=(swa_slot_mapping is not None), - num_warps=1, - ) - - if zeros_out is not None: - return q_out.view(-1, qh * d), k_out, key_cache, value_cache, zeros_out - return q_out.view(-1, qh * d), k_out, key_cache, value_cache diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index ff304a13d..c7975c96c 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -6,10 +6,12 @@ from typing import TYPE_CHECKING, Optional import torch import triton -import triton.language as tl 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.triton_ops.kv_indices import ( + create_flashinfer_kv_indices_triton, +) +from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import get_bool_env_var, get_device_core_count @@ -22,58 +24,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -@triton.jit -def get_num_kv_splits_triton( - num_kv_splits_ptr, - seq_lens_ptr, - num_seq, - num_group, - num_head, - num_kv_head, - max_kv_splits, - device_core_count, - MAX_NUM_SEQ: tl.constexpr, -): - # TODO: this method is tunable, we need more online serving data to tune it - offs_seq = tl.arange(0, MAX_NUM_SEQ) - mask_seq = offs_seq < num_seq - - seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0) - max_seq_len = tl.max(seq_lens) - seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len) - min_seq_len = tl.min(seq_lens) - if max_seq_len * 8 < min_seq_len * 10: - min_seq_len = max_seq_len - max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits) - kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1) - - # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually - ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0 - ext_device_core_count = tl.cast( - device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32 - ) - block_h, num_kv_group = 16, num_head // num_kv_head - if num_kv_group == 1: - token_grid = num_seq * num_group * num_head - else: - # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd - block_h = tl.minimum(block_h, num_kv_group) - token_grid = num_seq * num_group * tl.cdiv(num_head, block_h) - max_kv_splits_2 = tl.minimum( - tl.cdiv(ext_device_core_count, token_grid), max_kv_splits - ) - kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2) - - num_kv_splits = tl.maximum( - tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2) - ) - - offs_token = offs_seq * num_group - mask_token = offs_token < num_seq * num_group - for i in range(0, num_group): - tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token) - - @dataclass class ForwardMetadata: attn_logits: torch.Tensor diff --git a/python/sglang/srt/layers/elementwise.py b/python/sglang/srt/layers/elementwise.py index d8f2f7e48..11095f7bf 100644 --- a/python/sglang/srt/layers/elementwise.py +++ b/python/sglang/srt/layers/elementwise.py @@ -4,74 +4,13 @@ import torch import triton import triton.language as tl +from sglang.srt.layers.triton_ops.softcap import softcap_out as fused_softcap from sglang.srt.utils import is_hip from sglang.srt.utils.custom_op import register_custom_op _is_hip = is_hip() -fused_softcap_autotune = triton.autotune( - configs=[ - triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=4), - triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8), - triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16), - triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=4), - triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8), - triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=4), - triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=8), - triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=16), - triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4), - triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8), - triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16), - triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=32), - triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=32), - triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=32), - triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32), - triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32), - triton.Config(kwargs={"BLOCK_SIZE": 32768}, num_warps=32), - ], - key=["n_ele"], -) - - -@triton.jit -def fused_softcap_kernel( - output_ptr, - input_ptr, - n_ele, - softcap_const: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_ele - x = tl.load(input_ptr + offsets, mask=mask) - fx = x.to(tl.float32) - fxs = fx / softcap_const - exped = tl.exp(2 * fxs) - top = exped - 1 - bottom = exped + 1 - output = top / bottom * softcap_const - tl.store(output_ptr + offsets, output, mask=mask) - - -fused_softcap_kernel_autotuned = fused_softcap_autotune(fused_softcap_kernel) - - -def fused_softcap(x, softcap_const, autotune=False): - output = torch.empty_like(x, dtype=torch.float32) - n_elements = output.numel() - if autotune: - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - fused_softcap_kernel_autotuned[grid](output, x, n_elements, softcap_const) - else: - fused_softcap_kernel[(triton.cdiv(n_elements, 128),)]( - output, x, n_elements, softcap_const, BLOCK_SIZE=128, num_warps=8 - ) - return output - - # cast to float + softcap class Softcap: def __init__(self, softcap_const: float): diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 771c4d9db..8fdff43bd 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -18,10 +18,7 @@ import logging from typing import Any, Dict, List, Optional, Tuple, Union import torch -import triton -import triton.language as tl from torch import nn -from triton.language.extra import libdevice from sglang.srt.distributed import ( get_tensor_model_parallel_world_size, @@ -41,6 +38,7 @@ from sglang.srt.layers.dp_attention import ( get_dp_dtype, get_dp_hidden_size, ) +from sglang.srt.layers.triton_ops.softcap import softcap_inplace_logits as fused_softcap from sglang.srt.layers.utils.logprob import ( InputLogprobsResult, get_token_ids_logprobs_chunk, @@ -208,7 +206,6 @@ class LogitsMetadata: ) def compute_dp_attention_metadata(self): - cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0) dp_rank = get_attention_dp_rank() if dp_rank == 0: @@ -1073,55 +1070,3 @@ class LogitsProcessor(nn.Module): # They should be moved to GenerationBatchResult to keep this class clean. mm_input_embeds=logits_metadata.mm_input_embeds, ) - - -@triton.jit -def fused_softcap_kernel( - full_logits_ptr, - softcapping_value, - ncols, - row_stride, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(1).to(tl.int64) - pid = tl.program_id(0).to(tl.int64) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < ncols - - # Load values - row_ptr = full_logits_ptr + row * row_stride - x = tl.load(row_ptr + offsets, mask=mask) - - # Perform operations in-place - x = x / softcapping_value - x = libdevice.tanh(x) - x = x * softcapping_value - - # Store result - tl.store(row_ptr + offsets, x, mask=mask) - - -def fused_softcap(full_logits, final_logit_softcapping): - if full_logits.is_contiguous(): - nrows, ncols = 1, full_logits.numel() - row_stride = ncols - else: - assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor" - assert ( - full_logits.stride(1) == 1 - ), "non-contiguous softcap requires contiguous columns" - nrows, ncols = full_logits.shape - row_stride = full_logits.stride(0) - - BLOCK_SIZE = 1024 - grid = ((ncols + BLOCK_SIZE - 1) // BLOCK_SIZE, nrows) - - fused_softcap_kernel[grid]( - full_logits_ptr=full_logits, - softcapping_value=final_logit_softcapping, - ncols=ncols, - row_stride=row_stride, - BLOCK_SIZE=BLOCK_SIZE, - ) - return full_logits diff --git a/python/sglang/srt/layers/triton_ops/softcap.py b/python/sglang/srt/layers/triton_ops/softcap.py new file mode 100644 index 000000000..5d39de6f8 --- /dev/null +++ b/python/sglang/srt/layers/triton_ops/softcap.py @@ -0,0 +1,120 @@ +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +softcap_out_autotune = triton.autotune( + configs=[ + triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 32768}, num_warps=32), + ], + key=["n_ele"], +) + + +@triton.jit +def softcap_out_kernel( + output_ptr, + input_ptr, + n_ele, + softcap_const: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_ele + x = tl.load(input_ptr + offsets, mask=mask) + fx = x.to(tl.float32) + fxs = fx / softcap_const + exped = tl.exp(2 * fxs) + top = exped - 1 + bottom = exped + 1 + output = top / bottom * softcap_const + tl.store(output_ptr + offsets, output, mask=mask) + + +softcap_out_kernel_autotuned = softcap_out_autotune(softcap_out_kernel) + + +def softcap_out(x, softcap_const, autotune=False): + output = torch.empty_like(x, dtype=torch.float32) + n_elements = output.numel() + if autotune: + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + + softcap_out_kernel_autotuned[grid](output, x, n_elements, softcap_const) + else: + softcap_out_kernel[(triton.cdiv(n_elements, 128),)]( + output, x, n_elements, softcap_const, BLOCK_SIZE=128, num_warps=8 + ) + return output + + +@triton.jit +def softcap_inplace_logits_kernel( + full_logits_ptr, + softcapping_value, + ncols, + row_stride, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(1).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < ncols + + # Load values + row_ptr = full_logits_ptr + row * row_stride + x = tl.load(row_ptr + offsets, mask=mask) + + # Perform operations in-place + x = x / softcapping_value + x = libdevice.tanh(x) + x = x * softcapping_value + + # Store result + tl.store(row_ptr + offsets, x, mask=mask) + + +def softcap_inplace_logits(full_logits, final_logit_softcapping): + if full_logits.is_contiguous(): + nrows, ncols = 1, full_logits.numel() + row_stride = ncols + else: + assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor" + assert ( + full_logits.stride(1) == 1 + ), "non-contiguous softcap requires contiguous columns" + nrows, ncols = full_logits.shape + row_stride = full_logits.stride(0) + + BLOCK_SIZE = 1024 + grid = ((ncols + BLOCK_SIZE - 1) // BLOCK_SIZE, nrows) + + softcap_inplace_logits_kernel[grid]( + full_logits_ptr=full_logits, + softcapping_value=final_logit_softcapping, + ncols=ncols, + row_stride=row_stride, + BLOCK_SIZE=BLOCK_SIZE, + ) + return full_logits diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 0d623f159..81378e175 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -19,13 +19,16 @@ from __future__ import annotations Page-aligned memory pool. """ + from typing import TYPE_CHECKING import torch -import triton -import triton.language as tl from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.triton_ops.allocator import ( + alloc_decode_kernel, + alloc_extend_kernel, +) from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2 if TYPE_CHECKING: @@ -92,128 +95,6 @@ def alloc_extend_naive( ).view(-1) -@triton.jit -def alloc_extend_kernel( - pre_lens_ptr, - seq_lens_ptr, - last_loc_ptr, - free_page_ptr, - out_indices, - bs_upper: tl.constexpr, - page_size: tl.constexpr, -): - pid = tl.program_id(0) - - load_offset = tl.arange(0, bs_upper) - seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid) - pre_lens = tl.load(pre_lens_ptr + load_offset, mask=load_offset <= pid) - extend_lens = seq_lens - pre_lens - - seq_len = tl.load(seq_lens_ptr + pid) - pre_len = tl.load(pre_lens_ptr + pid) - extend_len = seq_len - pre_len - - sum_extend_lens = tl.sum(extend_lens) - output_start_loc = sum_extend_lens - extend_len - - num_pages_after = (seq_lens + page_size - 1) // page_size - num_pages_before = (pre_lens + page_size - 1) // page_size - num_new_pages = num_pages_after - num_pages_before - - num_page_start_loc_self = (seq_len + page_size - 1) // page_size - ( - pre_len + page_size - 1 - ) // page_size - sum_num_new_pages = tl.sum(num_new_pages) - new_page_start_loc = sum_num_new_pages - num_page_start_loc_self - - # Part 1: fill the old partial page - last_loc = tl.load(last_loc_ptr + pid) - num_part1 = ( - min(seq_len, (pre_len + page_size - 1) // page_size * page_size) - pre_len - ) - offset_one_page = tl.arange(0, page_size) - tl.store( - out_indices + output_start_loc + offset_one_page, - last_loc + 1 + offset_one_page, - mask=offset_one_page < num_part1, - ) - if pre_len + num_part1 == seq_len: - return - - # Part 2: fill the new full pages using a dynamic blocked loop. - # The loop bound is derived from num_part2 (runtime value), so Triton - # generates a real loop instead of unrolling — no constexpr dependency - # on extend size and only one kernel compilation. - num_part2 = ( - seq_len // page_size * page_size - - (pre_len + page_size - 1) // page_size * page_size - ) - BLOCK_EXTEND: tl.constexpr = 4096 - num_blocks = (num_part2 + BLOCK_EXTEND - 1) // BLOCK_EXTEND - for block_id in range(num_blocks): - offset_in_block = tl.arange(0, BLOCK_EXTEND) - offset = block_id * BLOCK_EXTEND + offset_in_block - mask = offset < num_part2 - page_start = tl.load( - free_page_ptr + new_page_start_loc + offset // page_size, - mask=mask, - ) - tl.store( - out_indices + output_start_loc + num_part1 + offset, - page_start * page_size + offset % page_size, - mask=mask, - ) - if pre_len + num_part1 + num_part2 == seq_len: - return - - # Part 3: fill the new partial page - num_part3 = seq_len - seq_len // page_size * page_size - start_loc = tl.load( - free_page_ptr + new_page_start_loc + num_page_start_loc_self - 1 - ) - tl.store( - out_indices + output_start_loc + num_part1 + num_part2 + offset_one_page, - start_loc * page_size + offset_one_page, - mask=offset_one_page < num_part3, - ) - - -@triton.jit -def alloc_decode_kernel( - seq_lens_ptr, - last_loc_ptr, - free_page_ptr, - out_indices, - bs_upper: tl.constexpr, - page_size: tl.constexpr, -): - pid = tl.program_id(0) - - load_offset = tl.arange(0, bs_upper) - seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid) - pre_lens = tl.where(load_offset <= pid, seq_lens - 1, seq_lens) - - seq_len = tl.load(seq_lens_ptr + pid) - pre_len = seq_len - 1 - - num_pages_after = (seq_lens + page_size - 1) // page_size - num_pages_before = (pre_lens + page_size - 1) // page_size - num_new_pages = num_pages_after - num_pages_before - - num_page_start_loc_self = (seq_len + page_size - 1) // page_size - ( - pre_len + page_size - 1 - ) // page_size - sum_num_new_pages = tl.sum(num_new_pages) - new_page_start_loc = sum_num_new_pages - num_page_start_loc_self - - if num_page_start_loc_self == 0: - last_loc = tl.load(last_loc_ptr + pid) - tl.store(out_indices + pid, last_loc + 1) - else: - page = tl.load(free_page_ptr + new_page_start_loc) - tl.store(out_indices + pid, page * page_size) - - class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): """ An allocator managing the indices to kv cache data. diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 643bb4f94..3edb831d3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -5,12 +5,21 @@ from typing import TYPE_CHECKING import numpy as np import torch -import triton -import triton.language as tl from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.triton_ops.common import ( + _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, +) +from sglang.srt.mem_cache.triton_ops.common import ( + get_last_loc_kernel as get_last_loc_kernel, +) +from sglang.srt.mem_cache.triton_ops.common import ( + get_last_loc_triton, + get_last_loc_triton_safe, + write_req_to_token_pool_triton, +) from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import is_hip, support_triton from sglang.srt.utils.common import ceil_align @@ -50,57 +59,6 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -@triton.jit -def write_req_to_token_pool_triton( - req_to_token_ptr, # [max_batch, max_context_len] - req_pool_indices, - prefix_tensors, - pre_lens, - seq_lens, - extend_lens, - out_cache_loc, - req_to_token_ptr_stride: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 512 - pid = tl.program_id(0) - - req_pool_index = tl.load(req_pool_indices + pid) - pre_len = tl.load(pre_lens + pid) - seq_len = tl.load(seq_lens + pid) - prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64)) - - # write prefix - num_loop = tl.cdiv(pre_len, BLOCK_SIZE) - for i in range(num_loop): - offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE - mask = offset < pre_len - value = tl.load(prefix_tensor + offset, mask=mask) - tl.store( - req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset, - value, - mask=mask, - ) - - # NOTE: This can be slow for large bs - cumsum_start = tl.cast(0, tl.int64) - for i in range(pid): - cumsum_start += tl.load(extend_lens + i) - - num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) - for i in range(num_loop): - offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE - mask = offset < (seq_len - pre_len) - value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) - tl.store( - req_to_token_ptr - + req_pool_index * req_to_token_ptr_stride - + offset - + pre_len, - value, - mask=mask, - ) - - def write_cache_indices( out_cache_loc: torch.Tensor, req_pool_indices_tensor: torch.Tensor, @@ -192,113 +150,6 @@ def get_last_loc_torch( ) -@triton.jit -def _get_last_loc_safe_kernel( - req_to_token, - req_pool_indices_tensor, - prefix_lens_tensor, - result_i32, - num_tokens, - req_to_token_stride, - BLOCK_SIZE: tl.constexpr, - PREFIX_DTYPE_IS_I64: tl.constexpr, -): - pid = tl.program_id(0) - offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE - mask = offset < num_tokens - - if PREFIX_DTYPE_IS_I64: - prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) - req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) - token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) - else: - prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) - req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) - token_index = req_pool_indices.to(tl.int64) * req_to_token_stride + ( - prefix_lens.to(tl.int64) - 1 - ) - - token_mask = mask & (prefix_lens > 0) - tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) - # Result stays int32 (req_to_token dtype); caller promotes after return. - tl.store(result_i32 + offset, tokens, mask=mask) - - -def get_last_loc_triton_safe( - req_to_token: torch.Tensor, - req_pool_indices_tensor: torch.Tensor, - prefix_lens_tensor: torch.Tensor, -) -> torch.Tensor: - """Fused `last_loc` Triton kernel whose in-kernel result buffer is int32 - (the dtype of req_to_token). The consumer-dtype promotion happens in - torch after the kernel returns, so Triton never issues a mixed-width - store — avoiding the HIP int32->int64 store bug hit by the legacy kernel. - """ - num_tokens = prefix_lens_tensor.shape[0] - BLOCK_SIZE = 256 - result_i32 = torch.empty( - num_tokens, dtype=torch.int32, device=prefix_lens_tensor.device - ) - grid = (triton.cdiv(num_tokens, BLOCK_SIZE),) - _get_last_loc_safe_kernel[grid]( - req_to_token, - req_pool_indices_tensor, - prefix_lens_tensor, - result_i32, - num_tokens, - req_to_token.stride(0), - BLOCK_SIZE=BLOCK_SIZE, - PREFIX_DTYPE_IS_I64=(prefix_lens_tensor.dtype == torch.int64), - ) - return result_i32.to(prefix_lens_tensor.dtype) - - -@triton.jit -def get_last_loc_kernel( - req_to_token, - req_pool_indices_tensor, - prefix_lens_tensor, - result, - num_tokens, - req_to_token_stride, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE - mask = offset < num_tokens - - prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) - req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) - - token_mask = prefix_lens > 0 - token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) - tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) - - tl.store(result + offset, tokens, mask=mask) - - -def get_last_loc_triton( - req_to_token: torch.Tensor, - req_pool_indices_tensor: torch.Tensor, - prefix_lens_tensor: torch.Tensor, -) -> torch.Tensor: - BLOCK_SIZE = 256 - num_tokens = prefix_lens_tensor.shape[0] - result = torch.empty_like(prefix_lens_tensor) - grid = (triton.cdiv(num_tokens, BLOCK_SIZE),) - - get_last_loc_kernel[grid]( - req_to_token, - req_pool_indices_tensor, - prefix_lens_tensor, - result, - num_tokens, - req_to_token.stride(0), - BLOCK_SIZE, - ) - return result - - def alloc_token_slots( tree_cache: BasePrefixCache, num_tokens: int, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 8bd91fad9..b0105d7cf 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -11,11 +11,7 @@ distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. -""" -from __future__ import annotations - -""" Memory pool. SGLang has two levels of memory pool. @@ -24,6 +20,8 @@ TokenToKVPoolAllocator manages the indices to kv cache data. KVCache actually holds the physical kv cache. """ +from __future__ import annotations + import abc import dataclasses import logging @@ -33,8 +31,6 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union import numpy as np import torch -import triton -import triton.language as tl from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache from sglang.srt.configs.mamba_utils import BaseLinearStateParams @@ -48,6 +44,9 @@ from sglang.srt.layers.attention.dsa.quant_k_cache import ( 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.triton_ops.cache_move import ( + copy_all_layer_kv_cache_tiled, +) from sglang.srt.mem_cache.utils import ( get_mla_kv_buffer_triton, maybe_init_custom_mem_pool, @@ -610,11 +609,11 @@ class HybridReqToTokenPool(ReqToTokenPool): mamba_ping_pong_track_buffers.append(req.mamba_ping_pong_track_buffer) assert len(select_index) == len( mamba_indices - ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size." + ), "Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size." if self.enable_mamba_extra_buffer: assert len(select_index) == len( mamba_ping_pong_track_buffers - ), f"Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." + ), "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." mamba_index_tensor = torch.stack(mamba_indices).to(dtype=torch.int32) self.req_index_to_mamba_index_mapping[select_index] = mamba_index_tensor if self.enable_mamba_extra_buffer: @@ -795,7 +794,6 @@ class KVCache(abc.ABC): class MHATokenToKVPool(KVCache): - def __init__( self, size: int, @@ -1257,7 +1255,6 @@ class NoOpMHATokenToKVPool(MHATokenToKVPool): class MHATokenToKVPoolFP4(MHATokenToKVPool): - def _create_buffers(self): with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( @@ -1434,7 +1431,6 @@ class HybridLinearKVPool(KVCache): assert not enable_kvcache_transpose self.use_mla = use_mla if not use_mla: - TokenToKVPoolClass = MHATokenToKVPool if current_platform.is_out_of_tree(): @@ -1457,7 +1453,6 @@ class HybridLinearKVPool(KVCache): enable_memory_saver=enable_memory_saver, ) else: - TokenToKVPoolClass = MLATokenToKVPool if current_platform.is_out_of_tree(): @@ -1543,7 +1538,6 @@ class HybridLinearKVPool(KVCache): @contextmanager def _transfer_id_context(self, layer: RadixAttention): - @contextmanager def _patch_layer_id(layer): original_layer_id = layer.layer_id @@ -1863,7 +1857,6 @@ class MLATokenToKVPool(KVCache): class MLATokenToKVPoolFP4(MLATokenToKVPool): - def _create_buffers(self): with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( @@ -2012,7 +2005,6 @@ class DSATokenToKVPool(MLATokenToKVPool): end_layer: Optional[int] = None, index_buf_size: Optional[int] = None, ): - override_dim = ( kv_cache_dim if kv_cache_dim != kv_lora_rank + qk_rope_head_dim else None ) @@ -2232,39 +2224,3 @@ def move_kv_cache_native( for k_cache, v_cache in zip(k_buffer, v_buffer): k_cache[tgt_loc_flat] = k_cache[src_loc_flat] v_cache[tgt_loc_flat] = v_cache[src_loc_flat] - - -@triton.jit -def copy_all_layer_kv_cache_tiled( - data_ptrs, - strides, - tgt_loc_ptr, - src_loc_ptr, - num_locs, - num_locs_upper: tl.constexpr, - BYTES_PER_TILE: tl.constexpr, -): - """2D tiled kernel. Safe for in-place copy.""" - bid = tl.program_id(0) - tid = tl.program_id(1) - - stride = tl.load(strides + bid) - base_ptr = tl.load(data_ptrs + bid) - base_ptr = tl.cast(base_ptr, tl.pointer_type(tl.uint8)) - - byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE) - mask_byte = byte_off < stride - tl.multiple_of(byte_off, 16) - - loc_idx = tl.arange(0, num_locs_upper) - mask_loc = loc_idx < num_locs - - src = tl.load(src_loc_ptr + loc_idx, mask=mask_loc, other=0) - tgt = tl.load(tgt_loc_ptr + loc_idx, mask=mask_loc, other=0) - - src_ptr = base_ptr + src[:, None] * stride + byte_off[None, :] - tgt_ptr = base_ptr + tgt[:, None] * stride + byte_off[None, :] - - mask = mask_loc[:, None] & mask_byte[None, :] - vals = tl.load(src_ptr, mask=mask) - tl.store(tgt_ptr, vals, mask=mask) diff --git a/python/sglang/srt/mem_cache/triton_ops/__init__.py b/python/sglang/srt/mem_cache/triton_ops/__init__.py new file mode 100644 index 000000000..3e69762f9 --- /dev/null +++ b/python/sglang/srt/mem_cache/triton_ops/__init__.py @@ -0,0 +1 @@ +"""Triton kernels for memory cache operations.""" diff --git a/python/sglang/srt/mem_cache/triton_ops/allocator.py b/python/sglang/srt/mem_cache/triton_ops/allocator.py new file mode 100644 index 000000000..3a51d76e1 --- /dev/null +++ b/python/sglang/srt/mem_cache/triton_ops/allocator.py @@ -0,0 +1,124 @@ +import triton +import triton.language as tl + + +@triton.jit +def alloc_extend_kernel( + pre_lens_ptr, + seq_lens_ptr, + last_loc_ptr, + free_page_ptr, + out_indices, + bs_upper: tl.constexpr, + page_size: tl.constexpr, +): + pid = tl.program_id(0) + + load_offset = tl.arange(0, bs_upper) + seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid) + pre_lens = tl.load(pre_lens_ptr + load_offset, mask=load_offset <= pid) + extend_lens = seq_lens - pre_lens + + seq_len = tl.load(seq_lens_ptr + pid) + pre_len = tl.load(pre_lens_ptr + pid) + extend_len = seq_len - pre_len + + sum_extend_lens = tl.sum(extend_lens) + output_start_loc = sum_extend_lens - extend_len + + num_pages_after = (seq_lens + page_size - 1) // page_size + num_pages_before = (pre_lens + page_size - 1) // page_size + num_new_pages = num_pages_after - num_pages_before + + num_page_start_loc_self = (seq_len + page_size - 1) // page_size - ( + pre_len + page_size - 1 + ) // page_size + sum_num_new_pages = tl.sum(num_new_pages) + new_page_start_loc = sum_num_new_pages - num_page_start_loc_self + + # Part 1: fill the old partial page + last_loc = tl.load(last_loc_ptr + pid) + num_part1 = ( + min(seq_len, (pre_len + page_size - 1) // page_size * page_size) - pre_len + ) + offset_one_page = tl.arange(0, page_size) + tl.store( + out_indices + output_start_loc + offset_one_page, + last_loc + 1 + offset_one_page, + mask=offset_one_page < num_part1, + ) + if pre_len + num_part1 == seq_len: + return + + # Part 2: fill the new full pages using a dynamic blocked loop. + # The loop bound is derived from num_part2 (runtime value), so Triton + # generates a real loop instead of unrolling -- no constexpr dependency + # on extend size and only one kernel compilation. + num_part2 = ( + seq_len // page_size * page_size + - (pre_len + page_size - 1) // page_size * page_size + ) + BLOCK_EXTEND: tl.constexpr = 4096 + num_blocks = (num_part2 + BLOCK_EXTEND - 1) // BLOCK_EXTEND + for block_id in range(num_blocks): + offset_in_block = tl.arange(0, BLOCK_EXTEND) + offset = block_id * BLOCK_EXTEND + offset_in_block + mask = offset < num_part2 + page_start = tl.load( + free_page_ptr + new_page_start_loc + offset // page_size, + mask=mask, + ) + tl.store( + out_indices + output_start_loc + num_part1 + offset, + page_start * page_size + offset % page_size, + mask=mask, + ) + if pre_len + num_part1 + num_part2 == seq_len: + return + + # Part 3: fill the new partial page + num_part3 = seq_len - seq_len // page_size * page_size + start_loc = tl.load( + free_page_ptr + new_page_start_loc + num_page_start_loc_self - 1 + ) + tl.store( + out_indices + output_start_loc + num_part1 + num_part2 + offset_one_page, + start_loc * page_size + offset_one_page, + mask=offset_one_page < num_part3, + ) + + +@triton.jit +def alloc_decode_kernel( + seq_lens_ptr, + last_loc_ptr, + free_page_ptr, + out_indices, + bs_upper: tl.constexpr, + page_size: tl.constexpr, +): + pid = tl.program_id(0) + + load_offset = tl.arange(0, bs_upper) + seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid) + pre_lens = tl.where(load_offset <= pid, seq_lens - 1, seq_lens) + + seq_len = tl.load(seq_lens_ptr + pid) + pre_len = seq_len - 1 + + num_pages_after = (seq_lens + page_size - 1) // page_size + num_pages_before = (pre_lens + page_size - 1) // page_size + num_new_pages = num_pages_after - num_pages_before + + num_page_start_loc_self = (seq_len + page_size - 1) // page_size - ( + pre_len + page_size - 1 + ) // page_size + sum_num_new_pages = tl.sum(num_new_pages) + new_page_start_loc = sum_num_new_pages - num_page_start_loc_self + + if num_page_start_loc_self == 0: + last_loc = tl.load(last_loc_ptr + pid) + tl.store(out_indices + pid, last_loc + 1) + else: + page = tl.load(free_page_ptr + new_page_start_loc) + tl.store(out_indices + pid, page * page_size) diff --git a/python/sglang/srt/mem_cache/triton_ops/cache_move.py b/python/sglang/srt/mem_cache/triton_ops/cache_move.py new file mode 100644 index 000000000..bc6de0507 --- /dev/null +++ b/python/sglang/srt/mem_cache/triton_ops/cache_move.py @@ -0,0 +1,38 @@ +import triton +import triton.language as tl + + +@triton.jit +def copy_all_layer_kv_cache_tiled( + data_ptrs, + strides, + tgt_loc_ptr, + src_loc_ptr, + num_locs, + num_locs_upper: tl.constexpr, + BYTES_PER_TILE: tl.constexpr, +): + """2D tiled kernel. Safe for in-place copy.""" + bid = tl.program_id(0) + tid = tl.program_id(1) + + stride = tl.load(strides + bid) + base_ptr = tl.load(data_ptrs + bid) + base_ptr = tl.cast(base_ptr, tl.pointer_type(tl.uint8)) + + byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE) + mask_byte = byte_off < stride + tl.multiple_of(byte_off, 16) + + loc_idx = tl.arange(0, num_locs_upper) + mask_loc = loc_idx < num_locs + + src = tl.load(src_loc_ptr + loc_idx, mask=mask_loc, other=0) + tgt = tl.load(tgt_loc_ptr + loc_idx, mask=mask_loc, other=0) + + src_ptr = base_ptr + src[:, None] * stride + byte_off[None, :] + tgt_ptr = base_ptr + tgt[:, None] * stride + byte_off[None, :] + + mask = mask_loc[:, None] & mask_byte[None, :] + vals = tl.load(src_ptr, mask=mask) + tl.store(tgt_ptr, vals, mask=mask) diff --git a/python/sglang/srt/mem_cache/triton_ops/common.py b/python/sglang/srt/mem_cache/triton_ops/common.py new file mode 100644 index 000000000..511b67a43 --- /dev/null +++ b/python/sglang/srt/mem_cache/triton_ops/common.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def write_req_to_token_pool_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + prefix_tensors, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0) + + req_pool_index = tl.load(req_pool_indices + pid) + pre_len = tl.load(pre_lens + pid) + seq_len = tl.load(seq_lens + pid) + prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64)) + + # write prefix + num_loop = tl.cdiv(pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < pre_len + value = tl.load(prefix_tensor + offset, mask=mask) + tl.store( + req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset, + value, + mask=mask, + ) + + # NOTE: This can be slow for large bs + cumsum_start = tl.cast(0, tl.int64) + for i in range(pid): + cumsum_start += tl.load(extend_lens + i) + + num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < (seq_len - pre_len) + value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) + tl.store( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + offset + + pre_len, + value, + mask=mask, + ) + + +@triton.jit +def _get_last_loc_safe_kernel( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result_i32, + num_tokens, + req_to_token_stride, + BLOCK_SIZE: tl.constexpr, + PREFIX_DTYPE_IS_I64: tl.constexpr, +): + pid = tl.program_id(0) + offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + mask = offset < num_tokens + + if PREFIX_DTYPE_IS_I64: + prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) + req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) + token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) + else: + prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) + req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) + token_index = req_pool_indices.to(tl.int64) * req_to_token_stride + ( + prefix_lens.to(tl.int64) - 1 + ) + + token_mask = mask & (prefix_lens > 0) + tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) + # Result stays int32 (req_to_token dtype); caller promotes after return. + tl.store(result_i32 + offset, tokens, mask=mask) + + +def get_last_loc_triton_safe( + req_to_token: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + prefix_lens_tensor: torch.Tensor, +) -> torch.Tensor: + """Fused `last_loc` Triton kernel whose in-kernel result buffer is int32 + (the dtype of req_to_token). The consumer-dtype promotion happens in + torch after the kernel returns, so Triton never issues a mixed-width + store -- avoiding the HIP int32->int64 store bug hit by the legacy kernel. + """ + num_tokens = prefix_lens_tensor.shape[0] + BLOCK_SIZE = 256 + result_i32 = torch.empty( + num_tokens, dtype=torch.int32, device=prefix_lens_tensor.device + ) + grid = (triton.cdiv(num_tokens, BLOCK_SIZE),) + _get_last_loc_safe_kernel[grid]( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result_i32, + num_tokens, + req_to_token.stride(0), + BLOCK_SIZE=BLOCK_SIZE, + PREFIX_DTYPE_IS_I64=(prefix_lens_tensor.dtype == torch.int64), + ) + return result_i32.to(prefix_lens_tensor.dtype) + + +@triton.jit +def get_last_loc_kernel( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result, + num_tokens, + req_to_token_stride, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + mask = offset < num_tokens + + prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) + req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) + + token_mask = prefix_lens > 0 + token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) + tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) + + tl.store(result + offset, tokens, mask=mask) + + +def get_last_loc_triton( + req_to_token: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + prefix_lens_tensor: torch.Tensor, +) -> torch.Tensor: + BLOCK_SIZE = 256 + num_tokens = prefix_lens_tensor.shape[0] + result = torch.empty_like(prefix_lens_tensor) + grid = (triton.cdiv(num_tokens, BLOCK_SIZE),) + + get_last_loc_kernel[grid]( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result, + num_tokens, + req_to_token.stride(0), + BLOCK_SIZE, + ) + return result diff --git a/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py b/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py new file mode 100644 index 000000000..6ec9b282e --- /dev/null +++ b/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py @@ -0,0 +1,377 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from sglang.jit_kernel.utils import is_arch_support_pdl + + +@triton.jit +def set_mla_kv_buffer_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK: tl.constexpr, + USE_GDC: tl.constexpr = False, +): + pid_loc = tl.program_id(0) + pid_blk = tl.program_id(1) + + base = pid_blk * BLOCK + offs = base + tl.arange(0, BLOCK) + total_dim = nope_dim + rope_dim + mask = offs < total_dim + + if USE_GDC: + tl.extra.cuda.gdc_wait() + + loc = tl.load(loc_ptr + pid_loc).to(tl.int64) + dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs + + # Three-way branch to handle boundary correctly while preserving fast path + if base + BLOCK <= nope_dim: + # Fast path: entire block is in nope region + src = tl.load( + cache_k_nope_ptr + pid_loc * nope_stride + offs, + mask=mask, + ) + elif base >= nope_dim: + # Fast path: entire block is in rope region + offs_rope = offs - nope_dim + src = tl.load( + cache_k_rope_ptr + pid_loc * rope_stride + offs_rope, + mask=mask, + ) + else: + # Boundary case: block spans nope/rope boundary (e.g., FP8 with nope_dim=528) + # Handle each offset individually to avoid negative indexing + is_nope = offs < nope_dim + is_rope = (offs >= nope_dim) & (offs < (nope_dim + rope_dim)) + + src_nope = tl.load( + cache_k_nope_ptr + pid_loc * nope_stride + offs, + mask=mask & is_nope, + other=0, + ) + src_rope = tl.load( + cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim), + mask=mask & is_rope, + other=0, + ) + + src = tl.where(is_nope, src_nope, src_rope) + + tl.store(dst_ptr, src, mask=mask) + + if USE_GDC: + tl.extra.cuda.gdc_launch_dependents() + + +# Above this loc count the TMA bulk-store path overtakes the single-CTA-per-loc +# Triton kernel. Below it, Triton with BLOCK = next_pow2(total_dim) (one CTA +# does the whole row in one tile, no boundary fan-out) is the winning fallback. +# Tuned on GB300 with DSv4 row widths. +_TMA_BULK_STORE_MIN_LOCS = 768 + + +def set_mla_kv_buffer_triton( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, +): + """Dispatch MLA paged-KV scatter writes to the fastest available path. + + Two paths, chosen on ``n_loc``: + + - ``n_loc >= 768`` (and SM90+ with TMA-compatible row widths): JIT CUDA + kernel where each warp loads one (nope, rope) row into shared memory and + issues a single ``cp.async.bulk.global.shared::cta`` store to scatter the + row at ``kv_buffer[loc[item]]``. Wins at large bs because it packs 4-8 + items per CTA, drastically reducing the CTA count vs single-CTA-per-loc. + - Otherwise: Triton kernel with ``BLOCK = next_pow2(nope_dim + rope_dim)``, + i.e. one CTA per loc covering the entire row in one tile. Wins at small + bs because there's no per-loc CTA fan-out (5x fewer CTAs than the old + BLOCK=128 dispatch) and the row-spanning block makes the boundary branch + a one-shot per CTA. This is also the path for SM<90 and for shapes that + violate the TMA 16-byte alignment. + + Speedup vs the legacy BLOCK=128 Triton kernel on GB300 (BF16, nope=512, + rope=64): ~1.05x at bs=8, ~1.5x at bs=128, 3.5x at bs=512, **11.7x at + bs=16384**. + + Name retained for caller compatibility; the implementation is no longer + Triton-only. + """ + from sglang.jit_kernel.set_mla_kv_buffer import ( + can_use_set_mla_kv_buffer, + ) + from sglang.jit_kernel.set_mla_kv_buffer import ( + set_mla_kv_buffer as jit_set_mla_kv_buffer, + ) + + n_loc = loc.numel() + nope_bytes = cache_k_nope.shape[-1] * cache_k_nope.element_size() + rope_bytes = cache_k_rope.shape[-1] * cache_k_rope.element_size() + if ( + n_loc >= _TMA_BULK_STORE_MIN_LOCS + and is_arch_support_pdl() + and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes) + ): + jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope) + return + + # Fallback: Triton with BLOCK = next_pow2(total_dim). One CTA per loc; the + # whole row in one tile (the existing 3-way nope/rope/boundary branch in + # ``set_mla_kv_buffer_kernel`` handles the over-allocation past total_dim + # via the offsFP8 cast with paged KV write.""" + pid_loc = tl.program_id(0) + pid_blk = tl.program_id(1) + + base = pid_blk * BLOCK + offs = base + tl.arange(0, BLOCK) + total_dim = nope_dim + rope_dim + mask = offs < total_dim + + if USE_GDC: + tl.extra.cuda.gdc_wait() + + loc = tl.load(loc_ptr + pid_loc).to(tl.int64) + dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs + + if base + BLOCK <= nope_dim: + src = tl.load( + cache_k_nope_ptr + pid_loc * nope_stride + offs, + mask=mask, + other=0.0, + ) + elif base >= nope_dim: + offs_rope = offs - nope_dim + src = tl.load( + cache_k_rope_ptr + pid_loc * rope_stride + offs_rope, + mask=mask, + other=0.0, + ) + else: + is_nope = offs < nope_dim + src_nope = tl.load( + cache_k_nope_ptr + pid_loc * nope_stride + offs, + mask=mask & is_nope, + other=0.0, + ) + src_rope = tl.load( + cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim), + mask=mask & ~is_nope, + other=0.0, + ) + src = tl.where(is_nope, src_nope, src_rope) + + # Destination pointer is FP8-typed view; tl.store performs downcast. + tl.store(dst_ptr, src, mask=mask) + + if USE_GDC: + tl.extra.cuda.gdc_launch_dependents() + + +def set_mla_kv_buffer_triton_fp8_quant( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + fp8_dtype: torch.dtype, +): + """Fuse BF16/FP16 MLA K quantization with paged KV write.""" + kv_buffer_fp8 = kv_buffer.view(fp8_dtype) + + nope_dim = cache_k_nope.shape[-1] + rope_dim = cache_k_rope.shape[-1] + total_dim = nope_dim + rope_dim + BLOCK = 128 + n_loc = loc.numel() + grid = (n_loc, triton.cdiv(total_dim, BLOCK)) + + pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {} + + set_mla_kv_buffer_fp8_quant_kernel[grid]( + kv_buffer_fp8, + cache_k_nope, + cache_k_rope, + loc, + kv_buffer_fp8.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + BLOCK=BLOCK, + **pdl_kwargs, + ) + + +@triton.jit +def set_mla_kv_scale_buffer_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK: tl.constexpr, +): + pid_loc = tl.program_id(0) + pid_blk = tl.program_id(1) + + base = pid_blk * BLOCK + offs = base + tl.arange(0, BLOCK) + total_dim = nope_dim + rope_dim + mask = offs < total_dim # Make sure don't cross the boundary + + loc = tl.load(loc_ptr + pid_loc) + dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs + + # Check each offs should read 'nope' or 'rope' + is_nope = offs < nope_dim + src_nope = tl.load( + cache_k_nope_ptr + pid_loc * nope_stride + offs, mask=mask & is_nope, other=0.0 + ) + src_rope = tl.load( + cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim), + mask=mask & ~is_nope, + other=0.0, + ) + + # Combine nope + rope + src = src_nope + src_rope + tl.store(dst_ptr, src, mask=mask) + + +def set_mla_kv_scale_buffer_triton( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, +): + nope_dim = cache_k_nope.shape[-1] + rope_dim = cache_k_rope.shape[-1] + total_dim = nope_dim + rope_dim + BLOCK = 128 # Keep origin, works for smaller total_dim as well. + n_loc = loc.numel() + grid = (n_loc, triton.cdiv(total_dim, BLOCK)) + + set_mla_kv_scale_buffer_kernel[grid]( + kv_buffer, + cache_k_nope, + cache_k_rope, + loc, + kv_buffer.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + BLOCK=BLOCK, + ) + + +@triton.jit +def get_mla_kv_buffer_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, +): + pid_loc = tl.program_id(0) + loc = tl.load(loc_ptr + pid_loc).to(tl.int64) + loc_src_ptr = kv_buffer_ptr + loc * buffer_stride + + nope_offs = tl.arange(0, nope_dim) + nope_src_ptr = loc_src_ptr + nope_offs + nope_src = tl.load(nope_src_ptr) + + tl.store( + cache_k_nope_ptr + pid_loc * nope_stride + nope_offs, + nope_src, + ) + + rope_offs = tl.arange(0, rope_dim) + rope_src_ptr = loc_src_ptr + nope_dim + rope_offs + rope_src = tl.load(rope_src_ptr) + tl.store( + cache_k_rope_ptr + pid_loc * rope_stride + rope_offs, + rope_src, + ) + + +def get_mla_kv_buffer_triton( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, +): + # The source data type will be implicitly converted to the target data type. + nope_dim = cache_k_nope.shape[-1] # 512 + rope_dim = cache_k_rope.shape[-1] # 64 + n_loc = loc.numel() + grid = (n_loc,) + + get_mla_kv_buffer_kernel[grid]( + kv_buffer, + cache_k_nope, + cache_k_rope, + loc, + kv_buffer.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + ) diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index a55b63b46..9b4aaed2d 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -16,11 +16,6 @@ import hashlib from typing import Any, Callable, List, Optional, Tuple -import torch -import triton -import triton.language as tl - -from sglang.jit_kernel.utils import is_arch_support_pdl from sglang.srt.environ import envs from sglang.srt.mem_cache.evict_policy import ( EvictionStrategy, @@ -32,6 +27,30 @@ from sglang.srt.mem_cache.evict_policy import ( PriorityStrategy, SLRUStrategy, ) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + get_mla_kv_buffer_kernel as get_mla_kv_buffer_kernel, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + get_mla_kv_buffer_triton as get_mla_kv_buffer_triton, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + set_mla_kv_buffer_fp8_quant_kernel as set_mla_kv_buffer_fp8_quant_kernel, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + set_mla_kv_buffer_kernel as set_mla_kv_buffer_kernel, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + set_mla_kv_buffer_triton as set_mla_kv_buffer_triton, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + set_mla_kv_buffer_triton_fp8_quant as set_mla_kv_buffer_triton_fp8_quant, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + set_mla_kv_scale_buffer_kernel as set_mla_kv_scale_buffer_kernel, +) +from sglang.srt.mem_cache.triton_ops.mla_buffer import ( + set_mla_kv_scale_buffer_triton as set_mla_kv_scale_buffer_triton, +) _EVICTION_POLICY_FACTORIES: dict[str, Callable[[], EvictionStrategy]] = { "lru": LRUStrategy, @@ -55,376 +74,6 @@ def get_eviction_strategy(eviction_policy: str) -> EvictionStrategy: ) from None -@triton.jit -def set_mla_kv_buffer_kernel( - kv_buffer_ptr, - cache_k_nope_ptr, - cache_k_rope_ptr, - loc_ptr, - buffer_stride: tl.constexpr, - nope_stride: tl.constexpr, - rope_stride: tl.constexpr, - nope_dim: tl.constexpr, - rope_dim: tl.constexpr, - BLOCK: tl.constexpr, - USE_GDC: tl.constexpr = False, -): - pid_loc = tl.program_id(0) - pid_blk = tl.program_id(1) - - base = pid_blk * BLOCK - offs = base + tl.arange(0, BLOCK) - total_dim = nope_dim + rope_dim - mask = offs < total_dim - - if USE_GDC: - tl.extra.cuda.gdc_wait() - - loc = tl.load(loc_ptr + pid_loc).to(tl.int64) - dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs - - # Three-way branch to handle boundary correctly while preserving fast path - if base + BLOCK <= nope_dim: - # Fast path: entire block is in nope region - src = tl.load( - cache_k_nope_ptr + pid_loc * nope_stride + offs, - mask=mask, - ) - elif base >= nope_dim: - # Fast path: entire block is in rope region - offs_rope = offs - nope_dim - src = tl.load( - cache_k_rope_ptr + pid_loc * rope_stride + offs_rope, - mask=mask, - ) - else: - # Boundary case: block spans nope/rope boundary (e.g., FP8 with nope_dim=528) - # Handle each offset individually to avoid negative indexing - is_nope = offs < nope_dim - is_rope = (offs >= nope_dim) & (offs < (nope_dim + rope_dim)) - - src_nope = tl.load( - cache_k_nope_ptr + pid_loc * nope_stride + offs, - mask=mask & is_nope, - other=0, - ) - src_rope = tl.load( - cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim), - mask=mask & is_rope, - other=0, - ) - - src = tl.where(is_nope, src_nope, src_rope) - - tl.store(dst_ptr, src, mask=mask) - - if USE_GDC: - tl.extra.cuda.gdc_launch_dependents() - - -# Above this loc count the TMA bulk-store path overtakes the single-CTA-per-loc -# Triton kernel. Below it, Triton with BLOCK = next_pow2(total_dim) (one CTA -# does the whole row in one tile, no boundary fan-out) is the winning fallback. -# Tuned on GB300 with DSv4 row widths. -_TMA_BULK_STORE_MIN_LOCS = 768 - - -def set_mla_kv_buffer_triton( - kv_buffer: torch.Tensor, - loc: torch.Tensor, - cache_k_nope: torch.Tensor, - cache_k_rope: torch.Tensor, -): - """Dispatch MLA paged-KV scatter writes to the fastest available path. - - Two paths, chosen on ``n_loc``: - - - ``n_loc >= 768`` (and SM90+ with TMA-compatible row widths): JIT CUDA - kernel where each warp loads one (nope, rope) row into shared memory and - issues a single ``cp.async.bulk.global.shared::cta`` store to scatter the - row at ``kv_buffer[loc[item]]``. Wins at large bs because it packs 4-8 - items per CTA, drastically reducing the CTA count vs single-CTA-per-loc. - - Otherwise: Triton kernel with ``BLOCK = next_pow2(nope_dim + rope_dim)``, - i.e. one CTA per loc covering the entire row in one tile. Wins at small - bs because there's no per-loc CTA fan-out (5× fewer CTAs than the old - BLOCK=128 dispatch) and the row-spanning block makes the boundary branch - a one-shot per CTA. This is also the path for SM<90 and for shapes that - violate the TMA 16-byte alignment. - - Speedup vs the legacy BLOCK=128 Triton kernel on GB300 (BF16, nope=512, - rope=64): ~1.05× at bs=8, ~1.5× at bs=128, 3.5× at bs=512, **11.7× at - bs=16384**. - - Name retained for caller compatibility; the implementation is no longer - Triton-only. - """ - from sglang.jit_kernel.set_mla_kv_buffer import ( - can_use_set_mla_kv_buffer, - ) - from sglang.jit_kernel.set_mla_kv_buffer import ( - set_mla_kv_buffer as jit_set_mla_kv_buffer, - ) - - n_loc = loc.numel() - nope_bytes = cache_k_nope.shape[-1] * cache_k_nope.element_size() - rope_bytes = cache_k_rope.shape[-1] * cache_k_rope.element_size() - if ( - n_loc >= _TMA_BULK_STORE_MIN_LOCS - and is_arch_support_pdl() - and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes) - ): - jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope) - return - - # Fallback: Triton with BLOCK = next_pow2(total_dim). One CTA per loc; the - # whole row in one tile (the existing 3-way nope/rope/boundary branch in - # ``set_mla_kv_buffer_kernel`` handles the over-allocation past total_dim - # via the offsFP8 cast with paged KV write.""" - pid_loc = tl.program_id(0) - pid_blk = tl.program_id(1) - - base = pid_blk * BLOCK - offs = base + tl.arange(0, BLOCK) - total_dim = nope_dim + rope_dim - mask = offs < total_dim - - if USE_GDC: - tl.extra.cuda.gdc_wait() - - loc = tl.load(loc_ptr + pid_loc).to(tl.int64) - dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs - - if base + BLOCK <= nope_dim: - src = tl.load( - cache_k_nope_ptr + pid_loc * nope_stride + offs, - mask=mask, - other=0.0, - ) - elif base >= nope_dim: - offs_rope = offs - nope_dim - src = tl.load( - cache_k_rope_ptr + pid_loc * rope_stride + offs_rope, - mask=mask, - other=0.0, - ) - else: - is_nope = offs < nope_dim - src_nope = tl.load( - cache_k_nope_ptr + pid_loc * nope_stride + offs, - mask=mask & is_nope, - other=0.0, - ) - src_rope = tl.load( - cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim), - mask=mask & ~is_nope, - other=0.0, - ) - src = tl.where(is_nope, src_nope, src_rope) - - # Destination pointer is FP8-typed view; tl.store performs downcast. - tl.store(dst_ptr, src, mask=mask) - - if USE_GDC: - tl.extra.cuda.gdc_launch_dependents() - - -def set_mla_kv_buffer_triton_fp8_quant( - kv_buffer: torch.Tensor, - loc: torch.Tensor, - cache_k_nope: torch.Tensor, - cache_k_rope: torch.Tensor, - fp8_dtype: torch.dtype, -): - """Fuse BF16/FP16 MLA K quantization with paged KV write.""" - kv_buffer_fp8 = kv_buffer.view(fp8_dtype) - - nope_dim = cache_k_nope.shape[-1] - rope_dim = cache_k_rope.shape[-1] - total_dim = nope_dim + rope_dim - BLOCK = 128 - n_loc = loc.numel() - grid = (n_loc, triton.cdiv(total_dim, BLOCK)) - - pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {} - - set_mla_kv_buffer_fp8_quant_kernel[grid]( - kv_buffer_fp8, - cache_k_nope, - cache_k_rope, - loc, - kv_buffer_fp8.stride(0), - cache_k_nope.stride(0), - cache_k_rope.stride(0), - nope_dim, - rope_dim, - BLOCK=BLOCK, - **pdl_kwargs, - ) - - -@triton.jit -def set_mla_kv_scale_buffer_kernel( - kv_buffer_ptr, - cache_k_nope_ptr, - cache_k_rope_ptr, - loc_ptr, - buffer_stride: tl.constexpr, - nope_stride: tl.constexpr, - rope_stride: tl.constexpr, - nope_dim: tl.constexpr, - rope_dim: tl.constexpr, - BLOCK: tl.constexpr, -): - pid_loc = tl.program_id(0) - pid_blk = tl.program_id(1) - - base = pid_blk * BLOCK - offs = base + tl.arange(0, BLOCK) - total_dim = nope_dim + rope_dim - mask = offs < total_dim # Make sure don't cross the boundary - - loc = tl.load(loc_ptr + pid_loc) - dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs - - # Check each offs should read 'nope' or 'rope' - is_nope = offs < nope_dim - src_nope = tl.load( - cache_k_nope_ptr + pid_loc * nope_stride + offs, mask=mask & is_nope, other=0.0 - ) - src_rope = tl.load( - cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim), - mask=mask & ~is_nope, - other=0.0, - ) - - # Combine nope + rope - src = src_nope + src_rope - tl.store(dst_ptr, src, mask=mask) - - -def set_mla_kv_scale_buffer_triton( - kv_buffer: torch.Tensor, - loc: torch.Tensor, - cache_k_nope: torch.Tensor, - cache_k_rope: torch.Tensor, -): - nope_dim = cache_k_nope.shape[-1] - rope_dim = cache_k_rope.shape[-1] - total_dim = nope_dim + rope_dim - BLOCK = 128 # Keep origin, works for smaller total_dim as well. - n_loc = loc.numel() - grid = (n_loc, triton.cdiv(total_dim, BLOCK)) - - set_mla_kv_scale_buffer_kernel[grid]( - kv_buffer, - cache_k_nope, - cache_k_rope, - loc, - kv_buffer.stride(0), - cache_k_nope.stride(0), - cache_k_rope.stride(0), - nope_dim, - rope_dim, - BLOCK=BLOCK, - ) - - -@triton.jit -def get_mla_kv_buffer_kernel( - kv_buffer_ptr, - cache_k_nope_ptr, - cache_k_rope_ptr, - loc_ptr, - buffer_stride: tl.constexpr, - nope_stride: tl.constexpr, - rope_stride: tl.constexpr, - nope_dim: tl.constexpr, - rope_dim: tl.constexpr, -): - pid_loc = tl.program_id(0) - loc = tl.load(loc_ptr + pid_loc).to(tl.int64) - loc_src_ptr = kv_buffer_ptr + loc * buffer_stride - - nope_offs = tl.arange(0, nope_dim) - nope_src_ptr = loc_src_ptr + nope_offs - nope_src = tl.load(nope_src_ptr) - - tl.store( - cache_k_nope_ptr + pid_loc * nope_stride + nope_offs, - nope_src, - ) - - rope_offs = tl.arange(0, rope_dim) - rope_src_ptr = loc_src_ptr + nope_dim + rope_offs - rope_src = tl.load(rope_src_ptr) - tl.store( - cache_k_rope_ptr + pid_loc * rope_stride + rope_offs, - rope_src, - ) - - -def get_mla_kv_buffer_triton( - kv_buffer: torch.Tensor, - loc: torch.Tensor, - cache_k_nope: torch.Tensor, - cache_k_rope: torch.Tensor, -): - # The source data type will be implicitly converted to the target data type. - nope_dim = cache_k_nope.shape[-1] # 512 - rope_dim = cache_k_rope.shape[-1] # 64 - n_loc = loc.numel() - grid = (n_loc,) - - get_mla_kv_buffer_kernel[grid]( - kv_buffer, - cache_k_nope, - cache_k_rope, - loc, - kv_buffer.stride(0), - cache_k_nope.stride(0), - cache_k_rope.stride(0), - nope_dim, - rope_dim, - ) - - def maybe_init_custom_mem_pool( device: str, ) -> Tuple[bool, Optional[Any], Optional[str]]: diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index e39c5a811..5f0fdf433 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -34,8 +34,6 @@ from functools import total_ordering from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union import torch -import triton -import triton.language as tl from sglang.srt.distributed.parallel_state import ( get_moe_expert_parallel_world_size, @@ -57,6 +55,7 @@ from sglang.srt.layers.dp_attention import ( from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( ForwardBatchDeepSeekMHAMixin, ) +from sglang.srt.model_executor.triton_ops.position import compute_position_triton from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( is_cuda, @@ -1171,7 +1170,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): self._pad_inputs_to_size(model_runner, tokens_padded, self.batch_size) def post_forward_mlp_sync_batch(self, logits_output: LogitsProcessorOutput): - self.forward_mode = getattr(self, "_original_forward_mode", self.forward_mode) self.batch_size = getattr(self, "_original_batch_size", self.batch_size) bs = self.batch_size @@ -1284,62 +1282,6 @@ def compute_position( return positions, extend_start_loc -def compute_position_triton( - extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_sum -): - """Compute positions. It is a fused version of `compute_position_torch`.""" - batch_size = extend_seq_lens.shape[0] - has_prefix = extend_prefix_lens.shape[0] == batch_size - - positions = torch.empty( - extend_seq_lens_sum, dtype=torch.int64, device=extend_seq_lens.device - ) - extend_start_loc = torch.empty( - batch_size, dtype=torch.int32, device=extend_seq_lens.device - ) - - # Launch kernel - compute_position_kernel[(batch_size,)]( - positions, - extend_start_loc, - extend_prefix_lens, - extend_seq_lens, - has_prefix, - ) - - return positions, extend_start_loc - - -@triton.jit -def compute_position_kernel( - positions, - extend_start_loc, - extend_prefix_lens, - extend_seq_lens, - has_prefix: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 512 - pid = tl.program_id(0).to(tl.int64) - - prefix_len = tl.load(extend_prefix_lens + pid) if has_prefix else 0 - seq_len = tl.load(extend_seq_lens + pid) - - # NOTE: This can be slow for large bs - cumsum_start = tl.cast(0, tl.int64) - for i in range(pid): - cumsum_start += tl.load(extend_seq_lens + i) - - num_loop = tl.cdiv(seq_len, BLOCK_SIZE) - for i in range(num_loop): - offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE - tl.store( - positions + cumsum_start + offset, - prefix_len + offset, - mask=offset < seq_len, - ) - tl.store(extend_start_loc + pid, cumsum_start) - - def compute_position_torch( extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor ): diff --git a/python/sglang/srt/model_executor/triton_ops/position.py b/python/sglang/srt/model_executor/triton_ops/position.py new file mode 100644 index 000000000..f9196d496 --- /dev/null +++ b/python/sglang/srt/model_executor/triton_ops/position.py @@ -0,0 +1,59 @@ +import torch +import triton +import triton.language as tl + + +def compute_position_triton( + extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_sum +): + """Compute positions. It is a fused version of `compute_position_torch`.""" + batch_size = extend_seq_lens.shape[0] + has_prefix = extend_prefix_lens.shape[0] == batch_size + + positions = torch.empty( + extend_seq_lens_sum, dtype=torch.int64, device=extend_seq_lens.device + ) + extend_start_loc = torch.empty( + batch_size, dtype=torch.int32, device=extend_seq_lens.device + ) + + # Launch kernel + compute_position_kernel[(batch_size,)]( + positions, + extend_start_loc, + extend_prefix_lens, + extend_seq_lens, + has_prefix, + ) + + return positions, extend_start_loc + + +@triton.jit +def compute_position_kernel( + positions, + extend_start_loc, + extend_prefix_lens, + extend_seq_lens, + has_prefix: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0).to(tl.int64) + + prefix_len = tl.load(extend_prefix_lens + pid) if has_prefix else 0 + seq_len = tl.load(extend_seq_lens + pid) + + # NOTE: This can be slow for large bs + cumsum_start = tl.cast(0, tl.int64) + for i in range(pid): + cumsum_start += tl.load(extend_seq_lens + i) + + num_loop = tl.cdiv(seq_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + tl.store( + positions + cumsum_start + offset, + prefix_len + offset, + mask=offset < seq_len, + ) + tl.store(extend_start_loc + pid, cumsum_start) diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 88e0bd03b..6550ccbe1 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -5,8 +5,6 @@ from typing import TYPE_CHECKING, Any import torch import torch.nn.functional as F -import triton -import triton.language as tl from sglang.srt.distributed import get_tp_group from sglang.srt.layers.dp_attention import ( @@ -38,8 +36,23 @@ from sglang.srt.speculative.spec_utils import ( SIMULATE_ACC_LEN, generate_simulated_accept_index, ) +from sglang.srt.speculative.triton_ops.cache_locs import ( + assign_draft_cache_locs_page_size_1 as assign_draft_cache_locs_page_size_1, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + assign_extend_cache_locs as assign_extend_cache_locs, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + assign_extend_cache_locs_func as assign_extend_cache_locs_func, +) +from sglang.srt.speculative.triton_ops.eagle import ( + fill_accepted_out_cache_loc as fill_accepted_out_cache_loc, +) +from sglang.srt.speculative.triton_ops.eagle import ( + fill_bonus_tokens as fill_bonus_tokens, +) from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob -from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu, next_power_of_2 +from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu _is_cuda = is_cuda() _is_hip = is_hip() @@ -61,33 +74,6 @@ if is_cuda() or is_musa(): ) -@triton.jit -def assign_draft_cache_locs_page_size_1( - req_pool_indices, - req_to_token, - seq_lens, - out_cache_loc, - pool_len: tl.constexpr, - topk: tl.constexpr, - speculative_num_steps: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 128 - pid = tl.program_id(axis=0) - - copy_len = topk * speculative_num_steps - out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps - - # Copy from req_to_token to out_cache_loc - kv_start = tl.load(seq_lens + pid) - token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len - num_loop = tl.cdiv(copy_len, BLOCK_SIZE) - for i in range(num_loop): - copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE - mask = copy_offset < copy_len - data = tl.load(token_pool + kv_start + copy_offset, mask=mask) - tl.store(out_cache_ptr + copy_offset, data, mask=mask) - - @dataclass class EagleDraftInputV2Mixin: def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch): @@ -499,118 +485,3 @@ class EagleVerifyInputV2Mixin: # tensor includes the trailing/bonus token via out-of-place +1 so the # name no longer flips semantics mid-function (naming doc C2). return predict, num_correct_drafts + 1, accept_index - - -@triton.jit -def fill_bonus_tokens( - accept_tokens, - accept_lens, - bonus_tokens_ptr, - num_draft_tokens: tl.constexpr, -): - # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel - # because this kernel reads accept_lens - pid = tl.program_id(axis=0) - # `accept_lens` includes the bonus token; the last accepted slot is at -1. - accept_len = tl.load(accept_lens + pid) - - bonus_token_idx = num_draft_tokens * pid + accept_len - 1 - bonus_token = tl.load(accept_tokens + bonus_token_idx) - tl.store(bonus_tokens_ptr + pid, bonus_token) - - -@triton.jit -def fill_accepted_out_cache_loc( - accept_index, - out_cache_loc, - accepted_out_cache_loc, - size_upper: tl.constexpr, -): - pid = tl.program_id(axis=0) - offset = tl.arange(0, size_upper) - - masks = (tl.load(accept_index + offset, offset < pid, other=-1) != -1).to(tl.int64) - dst = tl.sum(masks) - src = tl.load(accept_index + pid) - if src > -1: - value = tl.load(out_cache_loc + src) - tl.store(accepted_out_cache_loc + dst, value) - - -@triton.jit -def assign_extend_cache_locs( - req_pool_indices, - req_to_token, - start_offset, - end_offset, - out_cache_loc, - pool_len: tl.constexpr, - bs_upper: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 32 - pid = tl.program_id(axis=0) - kv_start = tl.load(start_offset + pid) - kv_end = tl.load(end_offset + pid) - token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len - - length_offset = tl.arange(0, bs_upper) - start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0) - end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0) - out_offset = tl.sum(end - start, axis=0) - - out_cache_ptr = out_cache_loc + out_offset - - load_offset = tl.arange(0, BLOCK_SIZE) + kv_start - save_offset = tl.arange(0, BLOCK_SIZE) - - num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) - for _ in range(num_loop): - mask = load_offset < kv_end - data = tl.load(token_pool + load_offset, mask=mask) - tl.store(out_cache_ptr + save_offset, data, mask=mask) - load_offset += BLOCK_SIZE - save_offset += BLOCK_SIZE - - -def assign_extend_cache_locs_func( - req_pool_indices: torch.Tensor, - req_to_token: torch.Tensor, - start_offset: torch.Tensor, - end_offset: torch.Tensor, - batch_size: int, - draft_token_num: int, - device, -) -> torch.Tensor: - if _is_cuda or _is_hip or _is_musa: - out_cache_loc = torch.empty( - (batch_size * draft_token_num,), - dtype=torch.int64, - device=device, - ) - assign_extend_cache_locs[(batch_size,)]( - req_pool_indices, - req_to_token, - start_offset, - end_offset, - out_cache_loc, - req_to_token.shape[1], - next_power_of_2(batch_size), - ) - - return out_cache_loc - - elif _is_npu: - out_cache_loc = torch.empty( - (batch_size * draft_token_num,), - dtype=torch.int32, - device=device, - ) - torch.ops.npu.cache_loc_update( - req_pool_indices, - req_to_token, - start_offset, - end_offset, - out_cache_loc, - ) - - return out_cache_loc diff --git a/python/sglang/srt/speculative/multi_layer_eagle_utils.py b/python/sglang/srt/speculative/multi_layer_eagle_utils.py index f1ce9d4b0..f1db97e6c 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_utils.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_utils.py @@ -12,339 +12,22 @@ # limitations under the License. # ============================================================================== -import torch -import triton -import triton.language as tl +from sglang.srt.speculative.triton_ops.multi_layer_eagle import ( + assign_hidden_states_pool_kernel, + assign_hidden_states_pool_torch, + assign_hidden_states_pool_triton, + assign_new_state_kernel, + assign_new_state_triton, + rotate_input_ids_kernel, + rotate_input_ids_triton, +) - -@triton.jit -def rotate_input_ids_kernel( - input_ids_ptr, - extend_start_loc_ptr, - extend_seq_lens_ptr, - topk_index_ptr, - select_index_ptr, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - - start_loc = tl.load(extend_start_loc_ptr + pid) - seq_len = tl.load(extend_seq_lens_ptr + pid) - new_token = tl.load(topk_index_ptr + pid) - - num_elements_to_shift = seq_len - 1 - - for off in range(0, num_elements_to_shift, BLOCK_SIZE): - offsets = off + tl.arange(0, BLOCK_SIZE) - mask = offsets < num_elements_to_shift - - read_ptr = input_ids_ptr + start_loc + offsets + 1 - val = tl.load(read_ptr, mask=mask) - tl.debug_barrier() - - write_ptr = input_ids_ptr + start_loc + offsets - tl.store(write_ptr, val, mask=mask) - tl.debug_barrier() - - if seq_len > 0: - if select_index_ptr is not None: - last_pos_ptr = input_ids_ptr + tl.load(select_index_ptr + pid) - else: - last_pos_ptr = input_ids_ptr + start_loc + seq_len - 1 - tl.store(last_pos_ptr, new_token) - - -def rotate_input_ids_triton( - input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index=None -): - batch_size = extend_seq_lens.shape[0] - BLOCK_SIZE = 4096 if select_index is not None else 8 - grid = (batch_size,) - - rotate_input_ids_kernel[grid]( - input_ids, - extend_start_loc, - extend_seq_lens, - topk_index, - select_index, - BLOCK_SIZE=BLOCK_SIZE, - ) - return input_ids - - -@triton.jit -def assign_new_state_kernel( - # Source pointers - old_input_ids_ptr, - old_positions_ptr, - old_hidden_states_ptr, - old_out_cache_loc_ptr, - old_extend_seq_lens_ptr, - old_extend_start_loc_ptr, - # Destination pointers - input_ids_ptr, - positions_ptr, - hidden_states_ptr, - out_cache_loc_ptr, - extend_seq_lens_ptr, - extend_start_loc_ptr, - # Auxiliary data pointers - next_token_ids_ptr, - seq_lens_ptr, - padding_lens_ptr, - req_pool_indices_ptr, - req_to_token_ptr, - req_to_hidden_states_pool_ptr, - # Scalars and Strides - step, - stride_hidden_seq, - stride_hidden_dim, # hidden_states strides - stride_pool_req, - stride_pool_step, - stride_pool_dim, # pool strides - stride_req_token_0, - stride_req_token_1, # req_to_token strides - # Meta-parameters - HIDDEN_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_HID: tl.constexpr, -): - pid = tl.program_id(0) - - seq_len: tl.tensor = tl.load(seq_lens_ptr + pid) - old_extend_len = tl.load(old_extend_seq_lens_ptr + pid) - old_start = tl.load(old_extend_start_loc_ptr + pid) - new_extend_len = old_extend_len + 1 - new_start = old_start + pid - - tl.store(extend_seq_lens_ptr + pid, new_extend_len) - tl.store(extend_start_loc_ptr + pid, new_start) - - offs_seq = tl.arange(0, BLOCK_SEQ) - mask_seq = offs_seq < old_extend_len - - old_ids = tl.load(old_input_ids_ptr + old_start + offs_seq, mask=mask_seq) - tl.store(input_ids_ptr + new_start + offs_seq, old_ids, mask=mask_seq) - padding_len = tl.load(padding_lens_ptr + pid) - tl.store( - input_ids_ptr + new_start + old_extend_len - padding_len, - tl.load(next_token_ids_ptr + pid), - ) - - old_pos = tl.load(old_positions_ptr + old_start + offs_seq, mask=mask_seq) - tl.store(positions_ptr + new_start + 1 + offs_seq, old_pos, mask=mask_seq) - tl.store( - positions_ptr + new_start, max(tl.load(old_positions_ptr + old_start) - 1, 0) - ) - - old_cache = tl.load(old_out_cache_loc_ptr + old_start + offs_seq, mask=mask_seq) - tl.store(out_cache_loc_ptr + new_start + 1 + offs_seq, old_cache, mask=mask_seq) - - req_idx = tl.load(req_pool_indices_ptr + pid) - token_idx_col = seq_len - old_extend_len - 1 - if token_idx_col >= 0: - req_token_ptr_loc = ( - req_to_token_ptr - + (req_idx * stride_req_token_0) - + (token_idx_col * stride_req_token_1) - ) - last_cache_loc = tl.load(req_token_ptr_loc) - tl.store(out_cache_loc_ptr + new_start, last_cache_loc) - - pool_vec_offset_base = ((req_idx + 1) * stride_pool_req) + ( - -(step + 1) * stride_pool_step - ) - - for off_h in range(0, HIDDEN_DIM, BLOCK_HID): - offs_h = off_h + tl.arange(0, BLOCK_HID) - mask_h = offs_h < HIDDEN_DIM - - for i in range(BLOCK_SEQ): - if i < old_extend_len: - old_h_ptr = ( - old_hidden_states_ptr - + (old_start + i) * stride_hidden_seq - + (offs_h * stride_hidden_dim) - ) - new_h_ptr = ( - hidden_states_ptr - + (new_start + 1 + i) * stride_hidden_seq - + (offs_h * stride_hidden_dim) - ) - - chunk_old = tl.load(old_h_ptr, mask=mask_h) - tl.store(new_h_ptr, chunk_old, mask=mask_h) - - pool_ptrs = ( - req_to_hidden_states_pool_ptr - + pool_vec_offset_base - + (offs_h * stride_pool_dim) - ) - pool_val = tl.load(pool_ptrs, mask=mask_h) - - new_h_start_ptrs = ( - hidden_states_ptr - + (new_start * stride_hidden_seq) - + (offs_h * stride_hidden_dim) - ) - tl.store(new_h_start_ptrs, pool_val, mask=mask_h) - - -def assign_new_state_triton( - next_token_ids: torch.Tensor, - old_input_ids: torch.Tensor, - old_positions: torch.Tensor, - old_hidden_states: torch.Tensor, - old_out_cache_loc: torch.Tensor, - old_extend_seq_lens: torch.Tensor, - old_extend_start_loc: torch.Tensor, - input_ids: torch.Tensor, - positions: torch.Tensor, - hidden_states: torch.Tensor, - out_cache_loc: torch.Tensor, - extend_seq_lens: torch.Tensor, - extend_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - padding_lens: torch.Tensor, - num_seqs: int, - step: int, - req_pool_indices: torch.Tensor, - req_to_token: torch.Tensor, - req_to_hidden_states_pool: torch.Tensor, -): - """ - Wrapper function to calculate offsets and launch the Triton kernel. - """ - hidden_dim = hidden_states.shape[1] - - BLOCK_SEQ = 8 - BLOCK_HID = 64 - - grid = (num_seqs,) - - assign_new_state_kernel[grid]( - # Pointers - old_input_ids, - old_positions, - old_hidden_states, - old_out_cache_loc, - old_extend_seq_lens, - old_extend_start_loc, - input_ids, - positions, - hidden_states, - out_cache_loc, - extend_seq_lens, - extend_start_loc, - next_token_ids, - seq_lens, - padding_lens, - req_pool_indices, - req_to_token, - req_to_hidden_states_pool, - # Constants/Strides - step, - old_hidden_states.stride(0), - old_hidden_states.stride(1), - req_to_hidden_states_pool.stride(0), - req_to_hidden_states_pool.stride(1), - req_to_hidden_states_pool.stride(2), - req_to_token.stride(0), - req_to_token.stride(1), - # Meta - HIDDEN_DIM=hidden_dim, - BLOCK_SEQ=BLOCK_SEQ, - BLOCK_HID=BLOCK_HID, - ) - - -@triton.jit -def assign_hidden_states_pool_kernel( - hidden_states_ptr, - req_pool_indices_ptr, - req_to_hidden_states_pool_ptr, - extend_seq_lens_ptr, - extend_start_loc_ptr, - stride_hidden_seq, - stride_hidden_dim, - stride_pool_req, - stride_pool_step, - stride_pool_dim, - HIDDEN_DIM: tl.constexpr, - pool_size: tl.constexpr, - BLOCK_HID: tl.constexpr, -): - pid = tl.program_id(0) - - extend_len = tl.load(extend_seq_lens_ptr + pid) - start_loc = tl.load(extend_start_loc_ptr + pid) - end_loc = start_loc + extend_len - - req_idx = tl.load(req_pool_indices_ptr + pid) - pool_vec_offset_base = req_idx * stride_pool_req - - for i in range(pool_size): - for off_h in range(0, HIDDEN_DIM, BLOCK_HID): - offs_h = off_h + tl.arange(0, BLOCK_HID) - mask_h = offs_h < HIDDEN_DIM - - hid_ptr = ( - hidden_states_ptr - + (end_loc - pool_size + i) * stride_hidden_seq - + offs_h * stride_hidden_dim - ) - hid_val = tl.load(hid_ptr, mask=mask_h) - - pool_ptr = ( - req_to_hidden_states_pool_ptr - + pool_vec_offset_base - + i * stride_pool_step - + offs_h * stride_pool_dim - ) - tl.store(pool_ptr, hid_val, mask=mask_h) - - -def assign_hidden_states_pool_triton( - hidden_states: torch.Tensor, - req_pool_indices: torch.Tensor, - req_to_hidden_states_pool: torch.Tensor, - pool_size: int, - num_seqs: int, - extend_seq_lens: torch.Tensor, - extend_start_loc: torch.Tensor, -): - grid = (num_seqs,) - assign_hidden_states_pool_kernel[grid]( - hidden_states, - req_pool_indices, - req_to_hidden_states_pool, - extend_seq_lens, - extend_start_loc, - hidden_states.stride(0), - hidden_states.stride(1), - req_to_hidden_states_pool.stride(0), - req_to_hidden_states_pool.stride(1), - req_to_hidden_states_pool.stride(2), - HIDDEN_DIM=hidden_states.shape[1], - pool_size=pool_size, - BLOCK_HID=64, - ) - - -def assign_hidden_states_pool_torch( - hidden_states: torch.Tensor, - req_pool_indices: torch.Tensor, - req_to_hidden_states_pool: torch.Tensor, - pool_size: int, - num_seqs: int, - extend_seq_lens: torch.Tensor, - extend_start_loc: torch.Tensor, -): - for req in range(num_seqs): - pool_idx = req_pool_indices[req] - extend_len = extend_seq_lens[req] - start_loc = extend_start_loc[req] - end_loc = start_loc + extend_len - req_to_hidden_states_pool[pool_idx, :pool_size, :].copy_( - hidden_states[end_loc - pool_size : end_loc, :] - ) +__all__ = [ + "assign_hidden_states_pool_kernel", + "assign_hidden_states_pool_torch", + "assign_hidden_states_pool_triton", + "assign_new_state_kernel", + "assign_new_state_triton", + "rotate_input_ids_kernel", + "rotate_input_ids_triton", +] diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index e1b0f9fe8..201fc90bf 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -7,8 +7,6 @@ from contextlib import contextmanager from typing import TYPE_CHECKING, List, Optional import torch -import triton -import triton.language as tl from huggingface_hub import snapshot_download from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject @@ -20,7 +18,34 @@ from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.common import get_last_loc from sglang.srt.server_args import ServerArgs, get_global_server_args -from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, next_power_of_2 +from sglang.srt.speculative.triton_ops.cache_locs import ( + align_evict_mask_to_page_size as align_evict_mask_to_page_size, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + assign_draft_cache_locs as assign_draft_cache_locs, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + assign_req_to_token_pool as assign_req_to_token_pool, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + assign_req_to_token_pool_func as assign_req_to_token_pool_func, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + create_extend_after_decode_spec_info as create_extend_after_decode_spec_info, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + filter_finished_cache_loc_kernel as filter_finished_cache_loc_kernel, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + generate_draft_decode_kv_indices as generate_draft_decode_kv_indices, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + get_src_tgt_cache_loc as get_src_tgt_cache_loc, +) +from sglang.srt.speculative.triton_ops.cache_locs import ( + get_target_cache_loc as get_target_cache_loc, +) +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu _is_cuda = is_cuda() _is_hip = is_hip() @@ -111,403 +136,6 @@ def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool: return not server_args.enable_multi_layer_eagle -@triton.jit -def create_extend_after_decode_spec_info( - accept_tokens, - seq_lens, - accept_lens, - positions, - bonus_tokens_ptr, - bs_upper: tl.constexpr, -): - pid = tl.program_id(axis=0) - offsets = tl.arange(0, bs_upper) - seq_length = tl.load(seq_lens + pid) - # `accept_lens` includes the bonus token; load this req's value. - accept_len = tl.load(accept_lens + pid) - - accept_len_cumsum = tl.sum( - tl.load(accept_lens + offsets, mask=offsets < pid, other=0) - ) - positions_ptr = positions + accept_len_cumsum - mask = offsets < accept_len - tl.store(positions_ptr + offsets, seq_length - accept_len + offsets, mask) - - accept_len_cumsum += accept_len - 1 - bonus_token = tl.load(accept_tokens + accept_len_cumsum) - tl.store(bonus_tokens_ptr + pid, bonus_token) - - -@triton.jit -def assign_req_to_token_pool( - req_pool_indices, - req_to_token, - start_offset, - end_offset, - out_cache_loc, - pool_len: tl.constexpr, - bs_upper: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 32 - pid = tl.program_id(axis=0) - kv_start = tl.load(start_offset + pid) - kv_end = tl.load(end_offset + pid) - token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len - - length_offset = tl.arange(0, bs_upper) - start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0) - end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0) - out_offset = tl.sum(end - start, axis=0) - - out_cache_ptr = out_cache_loc + out_offset - - save_offset = tl.arange(0, BLOCK_SIZE) + kv_start - load_offset = tl.arange(0, BLOCK_SIZE) - - num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) - for _ in range(num_loop): - mask = save_offset < kv_end - data = tl.load(out_cache_ptr + load_offset, mask=mask) - tl.store(token_pool + save_offset, data, mask=mask) - save_offset += BLOCK_SIZE - load_offset += BLOCK_SIZE - - -def assign_req_to_token_pool_func( - req_pool_indices: torch.Tensor, - req_to_token: torch.Tensor, - start_offset: torch.Tensor, - end_offset: torch.Tensor, - out_cache_loc: torch.Tensor, - batch_size: int, -): - assign_req_to_token_pool[(batch_size,)]( - req_pool_indices, - req_to_token, - start_offset, - end_offset, - out_cache_loc, - req_to_token.shape[1], - next_power_of_2(batch_size), - ) - - -@triton.jit -def assign_draft_cache_locs( - req_pool_indices, - req_to_token, - seq_lens, - extend_lens, - num_new_pages_per_topk, - out_cache_loc, - source_cache_loc, - target_cache_loc, - last_page_lens_cumsum, - duplicate_cache_len: tl.constexpr, - pool_len: tl.constexpr, - topk: tl.constexpr, - speculative_num_steps: tl.constexpr, - page_size: tl.constexpr, - bs_upper: tl.constexpr, - iter_upper: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 128 - pid = tl.program_id(axis=0) - - if page_size == 1 or topk == 1: - copy_len = topk * speculative_num_steps - out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps - else: - bs_offset = tl.arange(0, bs_upper) - copy_len = tl.load(extend_lens + pid) - cum_copy_len = tl.sum(tl.load(extend_lens + bs_offset, mask=bs_offset < pid)) - out_cache_ptr = out_cache_loc + cum_copy_len - - # Part 1: Copy from out_cache_loc to req_to_token - kv_start = tl.load(seq_lens + pid) - token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len - num_loop = tl.cdiv(copy_len, BLOCK_SIZE) - for i in range(num_loop): - copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE - mask = copy_offset < copy_len - data = tl.load(out_cache_ptr + copy_offset, mask=mask) - tl.store(token_pool + kv_start + copy_offset, data, mask=mask) - # XXX (MUSA): Triton issue: chained boolean operators (A or B or C) are not supported. - if (page_size != 1 and topk != 1) and duplicate_cache_len > 0: - # Part 2: Copy indices into source_cache_loc and target_cache_loc - # Expected output: src:[8,9,10,8,9,10...] tgt:[16,17,18,24,25,26...] - prefix_len = tl.load(seq_lens + pid) - last_page_len = prefix_len % page_size - offsets = tl.arange(0, page_size) - mask = offsets < last_page_len - num_new_pages_per_topk_ = tl.load(num_new_pages_per_topk + pid) - prefix_base = token_pool + prefix_len - last_page_len - src_indices = tl.load(prefix_base + offsets, mask=mask) - last_page_lens_cumsum_ = tl.load(last_page_lens_cumsum + pid) - # Skip the first one since no copy is needed - for topk_id in range(1, topk): - tl.store( - source_cache_loc - + (topk - 1) * (last_page_lens_cumsum_ - last_page_len) - + (topk_id - 1) * last_page_len - + offsets, - src_indices, - mask=mask, - ) - tgt_indices = tl.load( - prefix_base + topk_id * num_new_pages_per_topk_ * page_size + offsets, - mask=mask, - ) - tl.store( - target_cache_loc - + (topk - 1) * (last_page_lens_cumsum_ - last_page_len) - + (topk_id - 1) * last_page_len - + offsets, - tgt_indices, - mask=mask, - ) - # Part 3: Copy and remove the used indices for duplication - # speculative_num_steps=5, page_size=4, num_new_pages_per_topk_=2, last_page_len=1 - # - xxxxx .. | - xxxxx .. | - # topk=0 topk=1 - # "-" means prefix tokens - # "x" means speculative draft tokens - # "." means padded tokens - # we only want to copy the "x" part. - iter_offset = tl.arange(0, iter_upper) - for topk_id in range(topk): - mask_upper = iter_offset < (speculative_num_steps + last_page_len) - mask_lower = iter_offset >= last_page_len - combined_mask = mask_upper & mask_lower - indices = tl.load( - prefix_base - + topk_id * num_new_pages_per_topk_ * page_size - + iter_offset, - mask=combined_mask, - other=0, - ) - # Shift from previous batches - ptr_offset = pid * speculative_num_steps * topk - # Subtract last_page_len to fill the gap of duplicated last page tokens. - # For example, token pool is (1, 2, 3, 4 ,5) and last page is 1, - # we write 2, 3, 4 to the front of out_cache_loc. - tl.store( - out_cache_loc - + ptr_offset - + topk_id * speculative_num_steps - - last_page_len - + iter_offset, - indices, - mask=combined_mask, - ) - - -@triton.jit -def generate_draft_decode_kv_indices( - req_pool_indices, - req_to_token, - paged_kernel_lens, - kv_indices, - kv_indptr, - positions, - pool_len: tl.constexpr, - kv_indices_stride: tl.constexpr, - kv_indptr_stride: tl.constexpr, - bs_upper: tl.constexpr, - iter_upper: tl.constexpr, - num_tokens_upper: tl.constexpr, - page_size: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 128 - iters = tl.program_id(axis=0) - bid = tl.program_id(axis=1) - topk_id = tl.program_id(axis=2) - - num_steps = tl.num_programs(axis=0) - num_seqs = tl.num_programs(axis=1) - topk = tl.num_programs(axis=2) - - kv_indices += kv_indices_stride * iters - kv_indptr += kv_indptr_stride * iters - iters += 1 - - load_offset = tl.arange(0, bs_upper) - seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0) - seq_len = tl.load(paged_kernel_lens + bid) - cum_seq_len = tl.sum(seq_lens) - - # Update kv_indices - kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters) - kv_ptr = kv_indices + kv_offset - token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len - - kv_offset = tl.arange(0, BLOCK_SIZE) - num_loop = tl.cdiv(seq_len, BLOCK_SIZE) - for _ in range(num_loop): - mask = kv_offset < seq_len - data = tl.load(token_pool_ptr + kv_offset, mask=mask) - tl.store(kv_ptr + kv_offset, data, mask=mask) - kv_offset += BLOCK_SIZE - - extend_offset = tl.arange(0, iter_upper) - if page_size == 1 or topk == 1: - extend_data = tl.load( - token_pool_ptr + seq_len + topk_id * num_steps + tl.arange(0, iter_upper), - mask=extend_offset < iters, - ) - else: - prefix_len = seq_len - last_page_len = prefix_len % page_size - num_new_pages_per_topk = ( - last_page_len + num_steps + page_size - 1 - ) // page_size - prefix_base = seq_len // page_size * page_size - start = ( - prefix_base + topk_id * num_new_pages_per_topk * page_size + last_page_len - ) - extend_data = tl.load( - token_pool_ptr + start + extend_offset, - mask=extend_offset < iters, - ) - - tl.store(kv_ptr + seq_len + extend_offset, extend_data, mask=extend_offset < iters) - - # Update kv_indptr - bs_offset = tl.arange(0, num_tokens_upper) - - zid = bid * topk + topk_id - if zid == 0: - zid = num_seqs * topk - positions = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0) - base = tl.sum(positions) - tl.store(kv_indptr + zid, base + zid * iters) - - -@triton.jit -def align_evict_mask_to_page_size( - seq_lens, - evict_mask, - page_size: tl.constexpr, - num_draft_tokens: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - t_range = tl.arange(0, BLOCK_SIZE) - - bid = tl.program_id(axis=0) - seq_len = tl.load(seq_lens + bid) - io_mask = t_range < num_draft_tokens - mask_row = tl.load( - evict_mask + bid * num_draft_tokens + t_range, mask=io_mask, other=0 - ) - - num_trues = tl.sum(mask_row) - num_false = num_draft_tokens - num_trues - - start = (seq_len + num_false - 1) // page_size * page_size - seq_len - for i in range(max(start, 0), min(start + page_size, num_draft_tokens)): - tl.store(evict_mask + bid * num_draft_tokens + i, False) - - -@triton.jit -def get_target_cache_loc( - tgt_cache_loc, - to_free_slots, - num_correct_drafts, - to_free_num_slots, - out_cache_loc, - num_verify_tokens: tl.constexpr, - num_verify_tokens_upper: tl.constexpr, - bs_upper: tl.constexpr, -): - bid = tl.program_id(axis=0) - offset = tl.arange(0, num_verify_tokens_upper) - bs_offset = tl.arange(0, bs_upper) - - # write the first part to tgt_cache_loc - accept_len_all = tl.load(num_correct_drafts + bs_offset, mask=bs_offset < bid) - tgt_cache_loc_start = tl.sum(accept_len_all) + bid - copy_len = tl.load(num_correct_drafts + bid) + 1 - out_cache_loc_row = tl.load( - out_cache_loc + bid * num_verify_tokens + offset, mask=offset < copy_len - ) - tl.store( - tgt_cache_loc + tgt_cache_loc_start + offset, - out_cache_loc_row, - mask=offset < copy_len, - ) - - # write the second part to to_free_num_pages - to_free_num_slots_all = tl.load(to_free_num_slots + bs_offset, mask=bs_offset < bid) - to_free_num_slots_cur = tl.load(to_free_num_slots + bid) - out_cache_loc_start = num_verify_tokens - to_free_num_slots_cur - to_free_slots_start = tl.sum(to_free_num_slots_all) - - copy_len = to_free_num_slots_cur - out_cache_loc_row = tl.load( - out_cache_loc + bid * num_verify_tokens + out_cache_loc_start + offset, - mask=offset < copy_len, - ) - tl.store( - to_free_slots + to_free_slots_start + offset, - out_cache_loc_row, - mask=offset < copy_len, - ) - - -@torch.compile(dynamic=True, disable=_is_npu) -def get_src_tgt_cache_loc( - seq_lens: torch.Tensor, - out_cache_loc: torch.Tensor, - accept_index: torch.Tensor, - num_correct_drafts: torch.Tensor, - draft_token_num: int, - page_size: int, -): - src_cache_loc = out_cache_loc[accept_index] - # zeros_like, not empty_like: any uncovered tail stays at slot 0 (padding) - # instead of caching-allocator garbage. - tgt_cache_loc = torch.zeros_like(src_cache_loc) - extended_len = seq_lens + draft_token_num - keep_len = torch.minimum( - (seq_lens + num_correct_drafts + 1 + page_size - 1) // page_size * page_size, - extended_len, - ) - to_free_num_slots = extended_len - keep_len - return src_cache_loc, tgt_cache_loc, to_free_num_slots - - -@triton.jit -def filter_finished_cache_loc_kernel( - out_cache_loc, - tgt_cache_loc, - num_correct_drafts, - num_accept_tokens_filter, - bs_upper: tl.constexpr, - num_verify_tokens_upper: tl.constexpr, -): - bid = tl.program_id(0) - bs_offset = tl.arange(0, bs_upper) - - num_correct_drafts_all = tl.load( - num_correct_drafts + bs_offset, mask=bs_offset < bid - ) - old_start = tl.sum(num_correct_drafts_all) + bid - - num_accept_tokens_filter_all = tl.load( - num_accept_tokens_filter + bs_offset, mask=bs_offset < bid - ) - new_start = tl.sum(num_accept_tokens_filter_all) - - copy_len = tl.load(num_accept_tokens_filter + bid) - copy_offset = tl.arange(0, num_verify_tokens_upper) - value = tl.load( - tgt_cache_loc + old_start + copy_offset, mask=copy_offset < copy_len - ) - tl.store( - out_cache_loc + new_start + copy_offset, value, mask=copy_offset < copy_len - ) - - @torch.compile(dynamic=True, disable=_is_npu) def create_num_accept_tokens_filter( num_correct_drafts: torch.Tensor, diff --git a/python/sglang/srt/speculative/triton_ops/cache_locs.py b/python/sglang/srt/speculative/triton_ops/cache_locs.py new file mode 100644 index 000000000..08edfb185 --- /dev/null +++ b/python/sglang/srt/speculative/triton_ops/cache_locs.py @@ -0,0 +1,515 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, next_power_of_2 + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_npu = is_npu() +_is_musa = is_musa() + + +@triton.jit +def create_extend_after_decode_spec_info( + accept_tokens, + seq_lens, + accept_lens, + positions, + bonus_tokens_ptr, + bs_upper: tl.constexpr, +): + pid = tl.program_id(axis=0) + offsets = tl.arange(0, bs_upper) + seq_length = tl.load(seq_lens + pid) + # `accept_lens` includes the bonus token; load this req's value. + accept_len = tl.load(accept_lens + pid) + + accept_len_cumsum = tl.sum( + tl.load(accept_lens + offsets, mask=offsets < pid, other=0) + ) + positions_ptr = positions + accept_len_cumsum + mask = offsets < accept_len + tl.store(positions_ptr + offsets, seq_length - accept_len + offsets, mask) + + accept_len_cumsum += accept_len - 1 + bonus_token = tl.load(accept_tokens + accept_len_cumsum) + tl.store(bonus_tokens_ptr + pid, bonus_token) + + +@triton.jit +def assign_req_to_token_pool( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + pool_len: tl.constexpr, + bs_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 32 + pid = tl.program_id(axis=0) + kv_start = tl.load(start_offset + pid) + kv_end = tl.load(end_offset + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + + length_offset = tl.arange(0, bs_upper) + start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0) + end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0) + out_offset = tl.sum(end - start, axis=0) + + out_cache_ptr = out_cache_loc + out_offset + + save_offset = tl.arange(0, BLOCK_SIZE) + kv_start + load_offset = tl.arange(0, BLOCK_SIZE) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for _ in range(num_loop): + mask = save_offset < kv_end + data = tl.load(out_cache_ptr + load_offset, mask=mask) + tl.store(token_pool + save_offset, data, mask=mask) + save_offset += BLOCK_SIZE + load_offset += BLOCK_SIZE + + +def assign_req_to_token_pool_func( + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + start_offset: torch.Tensor, + end_offset: torch.Tensor, + out_cache_loc: torch.Tensor, + batch_size: int, +): + assign_req_to_token_pool[(batch_size,)]( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + req_to_token.shape[1], + next_power_of_2(batch_size), + ) + + +@triton.jit +def assign_draft_cache_locs( + req_pool_indices, + req_to_token, + seq_lens, + extend_lens, + num_new_pages_per_topk, + out_cache_loc, + source_cache_loc, + target_cache_loc, + last_page_lens_cumsum, + duplicate_cache_len: tl.constexpr, + pool_len: tl.constexpr, + topk: tl.constexpr, + speculative_num_steps: tl.constexpr, + page_size: tl.constexpr, + bs_upper: tl.constexpr, + iter_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 128 + pid = tl.program_id(axis=0) + + if page_size == 1 or topk == 1: + copy_len = topk * speculative_num_steps + out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps + else: + bs_offset = tl.arange(0, bs_upper) + copy_len = tl.load(extend_lens + pid) + cum_copy_len = tl.sum(tl.load(extend_lens + bs_offset, mask=bs_offset < pid)) + out_cache_ptr = out_cache_loc + cum_copy_len + + # Part 1: Copy from out_cache_loc to req_to_token + kv_start = tl.load(seq_lens + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + num_loop = tl.cdiv(copy_len, BLOCK_SIZE) + for i in range(num_loop): + copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = copy_offset < copy_len + data = tl.load(out_cache_ptr + copy_offset, mask=mask) + tl.store(token_pool + kv_start + copy_offset, data, mask=mask) + # XXX (MUSA): Triton issue: chained boolean operators (A or B or C) are not supported. + if (page_size != 1 and topk != 1) and duplicate_cache_len > 0: + # Part 2: Copy indices into source_cache_loc and target_cache_loc + # Expected output: src:[8,9,10,8,9,10...] tgt:[16,17,18,24,25,26...] + prefix_len = tl.load(seq_lens + pid) + last_page_len = prefix_len % page_size + offsets = tl.arange(0, page_size) + mask = offsets < last_page_len + num_new_pages_per_topk_ = tl.load(num_new_pages_per_topk + pid) + prefix_base = token_pool + prefix_len - last_page_len + src_indices = tl.load(prefix_base + offsets, mask=mask) + last_page_lens_cumsum_ = tl.load(last_page_lens_cumsum + pid) + # Skip the first one since no copy is needed + for topk_id in range(1, topk): + tl.store( + source_cache_loc + + (topk - 1) * (last_page_lens_cumsum_ - last_page_len) + + (topk_id - 1) * last_page_len + + offsets, + src_indices, + mask=mask, + ) + tgt_indices = tl.load( + prefix_base + topk_id * num_new_pages_per_topk_ * page_size + offsets, + mask=mask, + ) + tl.store( + target_cache_loc + + (topk - 1) * (last_page_lens_cumsum_ - last_page_len) + + (topk_id - 1) * last_page_len + + offsets, + tgt_indices, + mask=mask, + ) + # Part 3: Copy and remove the used indices for duplication + # speculative_num_steps=5, page_size=4, num_new_pages_per_topk_=2, last_page_len=1 + # - xxxxx .. | - xxxxx .. | + # topk=0 topk=1 + # "-" means prefix tokens + # "x" means speculative draft tokens + # "." means padded tokens + # we only want to copy the "x" part. + iter_offset = tl.arange(0, iter_upper) + for topk_id in range(topk): + mask_upper = iter_offset < (speculative_num_steps + last_page_len) + mask_lower = iter_offset >= last_page_len + combined_mask = mask_upper & mask_lower + indices = tl.load( + prefix_base + + topk_id * num_new_pages_per_topk_ * page_size + + iter_offset, + mask=combined_mask, + other=0, + ) + # Shift from previous batches + ptr_offset = pid * speculative_num_steps * topk + # Subtract last_page_len to fill the gap of duplicated last page tokens. + # For example, token pool is (1, 2, 3, 4 ,5) and last page is 1, + # we write 2, 3, 4 to the front of out_cache_loc. + tl.store( + out_cache_loc + + ptr_offset + + topk_id * speculative_num_steps + - last_page_len + + iter_offset, + indices, + mask=combined_mask, + ) + + +@triton.jit +def assign_draft_cache_locs_page_size_1( + req_pool_indices, + req_to_token, + seq_lens, + out_cache_loc, + pool_len: tl.constexpr, + topk: tl.constexpr, + speculative_num_steps: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 128 + pid = tl.program_id(axis=0) + + copy_len = topk * speculative_num_steps + out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps + + # Copy from req_to_token to out_cache_loc + kv_start = tl.load(seq_lens + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + num_loop = tl.cdiv(copy_len, BLOCK_SIZE) + for i in range(num_loop): + copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = copy_offset < copy_len + data = tl.load(token_pool + kv_start + copy_offset, mask=mask) + tl.store(out_cache_ptr + copy_offset, data, mask=mask) + + +@triton.jit +def generate_draft_decode_kv_indices( + req_pool_indices, + req_to_token, + paged_kernel_lens, + kv_indices, + kv_indptr, + positions, + pool_len: tl.constexpr, + kv_indices_stride: tl.constexpr, + kv_indptr_stride: tl.constexpr, + bs_upper: tl.constexpr, + iter_upper: tl.constexpr, + num_tokens_upper: tl.constexpr, + page_size: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 128 + iters = tl.program_id(axis=0) + bid = tl.program_id(axis=1) + topk_id = tl.program_id(axis=2) + + num_steps = tl.num_programs(axis=0) + num_seqs = tl.num_programs(axis=1) + topk = tl.num_programs(axis=2) + + kv_indices += kv_indices_stride * iters + kv_indptr += kv_indptr_stride * iters + iters += 1 + + load_offset = tl.arange(0, bs_upper) + seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0) + seq_len = tl.load(paged_kernel_lens + bid) + cum_seq_len = tl.sum(seq_lens) + + # Update kv_indices + kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters) + kv_ptr = kv_indices + kv_offset + token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len + + kv_offset = tl.arange(0, BLOCK_SIZE) + num_loop = tl.cdiv(seq_len, BLOCK_SIZE) + for _ in range(num_loop): + mask = kv_offset < seq_len + data = tl.load(token_pool_ptr + kv_offset, mask=mask) + tl.store(kv_ptr + kv_offset, data, mask=mask) + kv_offset += BLOCK_SIZE + + extend_offset = tl.arange(0, iter_upper) + if page_size == 1 or topk == 1: + extend_data = tl.load( + token_pool_ptr + seq_len + topk_id * num_steps + tl.arange(0, iter_upper), + mask=extend_offset < iters, + ) + else: + prefix_len = seq_len + last_page_len = prefix_len % page_size + num_new_pages_per_topk = ( + last_page_len + num_steps + page_size - 1 + ) // page_size + prefix_base = seq_len // page_size * page_size + start = ( + prefix_base + topk_id * num_new_pages_per_topk * page_size + last_page_len + ) + extend_data = tl.load( + token_pool_ptr + start + extend_offset, + mask=extend_offset < iters, + ) + + tl.store(kv_ptr + seq_len + extend_offset, extend_data, mask=extend_offset < iters) + + # Update kv_indptr + bs_offset = tl.arange(0, num_tokens_upper) + + zid = bid * topk + topk_id + if zid == 0: + zid = num_seqs * topk + positions = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0) + base = tl.sum(positions) + tl.store(kv_indptr + zid, base + zid * iters) + + +@triton.jit +def align_evict_mask_to_page_size( + seq_lens, + evict_mask, + page_size: tl.constexpr, + num_draft_tokens: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + t_range = tl.arange(0, BLOCK_SIZE) + + bid = tl.program_id(axis=0) + seq_len = tl.load(seq_lens + bid) + io_mask = t_range < num_draft_tokens + mask_row = tl.load( + evict_mask + bid * num_draft_tokens + t_range, mask=io_mask, other=0 + ) + + num_trues = tl.sum(mask_row) + num_false = num_draft_tokens - num_trues + + start = (seq_len + num_false - 1) // page_size * page_size - seq_len + for i in range(max(start, 0), min(start + page_size, num_draft_tokens)): + tl.store(evict_mask + bid * num_draft_tokens + i, False) + + +@torch.compile(dynamic=True, disable=_is_npu) +def get_src_tgt_cache_loc( + seq_lens: torch.Tensor, + out_cache_loc: torch.Tensor, + accept_index: torch.Tensor, + num_correct_drafts: torch.Tensor, + draft_token_num: int, + page_size: int, +): + src_cache_loc = out_cache_loc[accept_index] + # zeros_like, not empty_like: any uncovered tail stays at slot 0 (padding) + # instead of caching-allocator garbage. + tgt_cache_loc = torch.zeros_like(src_cache_loc) + extended_len = seq_lens + draft_token_num + keep_len = torch.minimum( + (seq_lens + num_correct_drafts + 1 + page_size - 1) // page_size * page_size, + extended_len, + ) + to_free_num_slots = extended_len - keep_len + return src_cache_loc, tgt_cache_loc, to_free_num_slots + + +@triton.jit +def get_target_cache_loc( + tgt_cache_loc, + to_free_slots, + num_correct_drafts, + to_free_num_slots, + out_cache_loc, + num_verify_tokens: tl.constexpr, + num_verify_tokens_upper: tl.constexpr, + bs_upper: tl.constexpr, +): + bid = tl.program_id(axis=0) + offset = tl.arange(0, num_verify_tokens_upper) + bs_offset = tl.arange(0, bs_upper) + + # write the first part to tgt_cache_loc + accept_len_all = tl.load(num_correct_drafts + bs_offset, mask=bs_offset < bid) + tgt_cache_loc_start = tl.sum(accept_len_all) + bid + copy_len = tl.load(num_correct_drafts + bid) + 1 + out_cache_loc_row = tl.load( + out_cache_loc + bid * num_verify_tokens + offset, mask=offset < copy_len + ) + tl.store( + tgt_cache_loc + tgt_cache_loc_start + offset, + out_cache_loc_row, + mask=offset < copy_len, + ) + + # write the second part to to_free_num_pages + to_free_num_slots_all = tl.load(to_free_num_slots + bs_offset, mask=bs_offset < bid) + to_free_num_slots_cur = tl.load(to_free_num_slots + bid) + out_cache_loc_start = num_verify_tokens - to_free_num_slots_cur + to_free_slots_start = tl.sum(to_free_num_slots_all) + + copy_len = to_free_num_slots_cur + out_cache_loc_row = tl.load( + out_cache_loc + bid * num_verify_tokens + out_cache_loc_start + offset, + mask=offset < copy_len, + ) + tl.store( + to_free_slots + to_free_slots_start + offset, + out_cache_loc_row, + mask=offset < copy_len, + ) + + +@triton.jit +def filter_finished_cache_loc_kernel( + out_cache_loc, + tgt_cache_loc, + num_correct_drafts, + num_accept_tokens_filter, + bs_upper: tl.constexpr, + num_verify_tokens_upper: tl.constexpr, +): + bid = tl.program_id(0) + bs_offset = tl.arange(0, bs_upper) + + num_correct_drafts_all = tl.load( + num_correct_drafts + bs_offset, mask=bs_offset < bid + ) + old_start = tl.sum(num_correct_drafts_all) + bid + + num_accept_tokens_filter_all = tl.load( + num_accept_tokens_filter + bs_offset, mask=bs_offset < bid + ) + new_start = tl.sum(num_accept_tokens_filter_all) + + copy_len = tl.load(num_accept_tokens_filter + bid) + copy_offset = tl.arange(0, num_verify_tokens_upper) + value = tl.load( + tgt_cache_loc + old_start + copy_offset, mask=copy_offset < copy_len + ) + tl.store( + out_cache_loc + new_start + copy_offset, value, mask=copy_offset < copy_len + ) + + +@triton.jit +def assign_extend_cache_locs( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + pool_len: tl.constexpr, + bs_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 32 + pid = tl.program_id(axis=0) + kv_start = tl.load(start_offset + pid) + kv_end = tl.load(end_offset + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + + length_offset = tl.arange(0, bs_upper) + start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0) + end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0) + out_offset = tl.sum(end - start, axis=0) + + out_cache_ptr = out_cache_loc + out_offset + + load_offset = tl.arange(0, BLOCK_SIZE) + kv_start + save_offset = tl.arange(0, BLOCK_SIZE) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for _ in range(num_loop): + mask = load_offset < kv_end + data = tl.load(token_pool + load_offset, mask=mask) + tl.store(out_cache_ptr + save_offset, data, mask=mask) + load_offset += BLOCK_SIZE + save_offset += BLOCK_SIZE + + +def assign_extend_cache_locs_func( + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + start_offset: torch.Tensor, + end_offset: torch.Tensor, + batch_size: int, + draft_token_num: int, + device, +) -> torch.Tensor: + if _is_cuda or _is_hip or _is_musa: + out_cache_loc = torch.empty( + (batch_size * draft_token_num,), + dtype=torch.int64, + device=device, + ) + assign_extend_cache_locs[(batch_size,)]( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + req_to_token.shape[1], + next_power_of_2(batch_size), + ) + + return out_cache_loc + + elif _is_npu: + out_cache_loc = torch.empty( + (batch_size * draft_token_num,), + dtype=torch.int32, + device=device, + ) + torch.ops.npu.cache_loc_update( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + ) + + return out_cache_loc diff --git a/python/sglang/srt/speculative/triton_ops/eagle.py b/python/sglang/srt/speculative/triton_ops/eagle.py new file mode 100644 index 000000000..4e8df6c60 --- /dev/null +++ b/python/sglang/srt/speculative/triton_ops/eagle.py @@ -0,0 +1,38 @@ +import triton +import triton.language as tl + + +@triton.jit +def fill_bonus_tokens( + accept_tokens, + accept_lens, + bonus_tokens_ptr, + num_draft_tokens: tl.constexpr, +): + # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel + # because this kernel reads accept_lens + pid = tl.program_id(axis=0) + # `accept_lens` includes the bonus token; the last accepted slot is at -1. + accept_len = tl.load(accept_lens + pid) + + bonus_token_idx = num_draft_tokens * pid + accept_len - 1 + bonus_token = tl.load(accept_tokens + bonus_token_idx) + tl.store(bonus_tokens_ptr + pid, bonus_token) + + +@triton.jit +def fill_accepted_out_cache_loc( + accept_index, + out_cache_loc, + accepted_out_cache_loc, + size_upper: tl.constexpr, +): + pid = tl.program_id(axis=0) + offset = tl.arange(0, size_upper) + + masks = (tl.load(accept_index + offset, offset < pid, other=-1) != -1).to(tl.int64) + dst = tl.sum(masks) + src = tl.load(accept_index + pid) + if src > -1: + value = tl.load(out_cache_loc + src) + tl.store(accepted_out_cache_loc + dst, value) diff --git a/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py b/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py new file mode 100644 index 000000000..f1ce9d4b0 --- /dev/null +++ b/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py @@ -0,0 +1,350 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import torch +import triton +import triton.language as tl + + +@triton.jit +def rotate_input_ids_kernel( + input_ids_ptr, + extend_start_loc_ptr, + extend_seq_lens_ptr, + topk_index_ptr, + select_index_ptr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + + start_loc = tl.load(extend_start_loc_ptr + pid) + seq_len = tl.load(extend_seq_lens_ptr + pid) + new_token = tl.load(topk_index_ptr + pid) + + num_elements_to_shift = seq_len - 1 + + for off in range(0, num_elements_to_shift, BLOCK_SIZE): + offsets = off + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements_to_shift + + read_ptr = input_ids_ptr + start_loc + offsets + 1 + val = tl.load(read_ptr, mask=mask) + tl.debug_barrier() + + write_ptr = input_ids_ptr + start_loc + offsets + tl.store(write_ptr, val, mask=mask) + tl.debug_barrier() + + if seq_len > 0: + if select_index_ptr is not None: + last_pos_ptr = input_ids_ptr + tl.load(select_index_ptr + pid) + else: + last_pos_ptr = input_ids_ptr + start_loc + seq_len - 1 + tl.store(last_pos_ptr, new_token) + + +def rotate_input_ids_triton( + input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index=None +): + batch_size = extend_seq_lens.shape[0] + BLOCK_SIZE = 4096 if select_index is not None else 8 + grid = (batch_size,) + + rotate_input_ids_kernel[grid]( + input_ids, + extend_start_loc, + extend_seq_lens, + topk_index, + select_index, + BLOCK_SIZE=BLOCK_SIZE, + ) + return input_ids + + +@triton.jit +def assign_new_state_kernel( + # Source pointers + old_input_ids_ptr, + old_positions_ptr, + old_hidden_states_ptr, + old_out_cache_loc_ptr, + old_extend_seq_lens_ptr, + old_extend_start_loc_ptr, + # Destination pointers + input_ids_ptr, + positions_ptr, + hidden_states_ptr, + out_cache_loc_ptr, + extend_seq_lens_ptr, + extend_start_loc_ptr, + # Auxiliary data pointers + next_token_ids_ptr, + seq_lens_ptr, + padding_lens_ptr, + req_pool_indices_ptr, + req_to_token_ptr, + req_to_hidden_states_pool_ptr, + # Scalars and Strides + step, + stride_hidden_seq, + stride_hidden_dim, # hidden_states strides + stride_pool_req, + stride_pool_step, + stride_pool_dim, # pool strides + stride_req_token_0, + stride_req_token_1, # req_to_token strides + # Meta-parameters + HIDDEN_DIM: tl.constexpr, + BLOCK_SEQ: tl.constexpr, + BLOCK_HID: tl.constexpr, +): + pid = tl.program_id(0) + + seq_len: tl.tensor = tl.load(seq_lens_ptr + pid) + old_extend_len = tl.load(old_extend_seq_lens_ptr + pid) + old_start = tl.load(old_extend_start_loc_ptr + pid) + new_extend_len = old_extend_len + 1 + new_start = old_start + pid + + tl.store(extend_seq_lens_ptr + pid, new_extend_len) + tl.store(extend_start_loc_ptr + pid, new_start) + + offs_seq = tl.arange(0, BLOCK_SEQ) + mask_seq = offs_seq < old_extend_len + + old_ids = tl.load(old_input_ids_ptr + old_start + offs_seq, mask=mask_seq) + tl.store(input_ids_ptr + new_start + offs_seq, old_ids, mask=mask_seq) + padding_len = tl.load(padding_lens_ptr + pid) + tl.store( + input_ids_ptr + new_start + old_extend_len - padding_len, + tl.load(next_token_ids_ptr + pid), + ) + + old_pos = tl.load(old_positions_ptr + old_start + offs_seq, mask=mask_seq) + tl.store(positions_ptr + new_start + 1 + offs_seq, old_pos, mask=mask_seq) + tl.store( + positions_ptr + new_start, max(tl.load(old_positions_ptr + old_start) - 1, 0) + ) + + old_cache = tl.load(old_out_cache_loc_ptr + old_start + offs_seq, mask=mask_seq) + tl.store(out_cache_loc_ptr + new_start + 1 + offs_seq, old_cache, mask=mask_seq) + + req_idx = tl.load(req_pool_indices_ptr + pid) + token_idx_col = seq_len - old_extend_len - 1 + if token_idx_col >= 0: + req_token_ptr_loc = ( + req_to_token_ptr + + (req_idx * stride_req_token_0) + + (token_idx_col * stride_req_token_1) + ) + last_cache_loc = tl.load(req_token_ptr_loc) + tl.store(out_cache_loc_ptr + new_start, last_cache_loc) + + pool_vec_offset_base = ((req_idx + 1) * stride_pool_req) + ( + -(step + 1) * stride_pool_step + ) + + for off_h in range(0, HIDDEN_DIM, BLOCK_HID): + offs_h = off_h + tl.arange(0, BLOCK_HID) + mask_h = offs_h < HIDDEN_DIM + + for i in range(BLOCK_SEQ): + if i < old_extend_len: + old_h_ptr = ( + old_hidden_states_ptr + + (old_start + i) * stride_hidden_seq + + (offs_h * stride_hidden_dim) + ) + new_h_ptr = ( + hidden_states_ptr + + (new_start + 1 + i) * stride_hidden_seq + + (offs_h * stride_hidden_dim) + ) + + chunk_old = tl.load(old_h_ptr, mask=mask_h) + tl.store(new_h_ptr, chunk_old, mask=mask_h) + + pool_ptrs = ( + req_to_hidden_states_pool_ptr + + pool_vec_offset_base + + (offs_h * stride_pool_dim) + ) + pool_val = tl.load(pool_ptrs, mask=mask_h) + + new_h_start_ptrs = ( + hidden_states_ptr + + (new_start * stride_hidden_seq) + + (offs_h * stride_hidden_dim) + ) + tl.store(new_h_start_ptrs, pool_val, mask=mask_h) + + +def assign_new_state_triton( + next_token_ids: torch.Tensor, + old_input_ids: torch.Tensor, + old_positions: torch.Tensor, + old_hidden_states: torch.Tensor, + old_out_cache_loc: torch.Tensor, + old_extend_seq_lens: torch.Tensor, + old_extend_start_loc: torch.Tensor, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + out_cache_loc: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + padding_lens: torch.Tensor, + num_seqs: int, + step: int, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + req_to_hidden_states_pool: torch.Tensor, +): + """ + Wrapper function to calculate offsets and launch the Triton kernel. + """ + hidden_dim = hidden_states.shape[1] + + BLOCK_SEQ = 8 + BLOCK_HID = 64 + + grid = (num_seqs,) + + assign_new_state_kernel[grid]( + # Pointers + old_input_ids, + old_positions, + old_hidden_states, + old_out_cache_loc, + old_extend_seq_lens, + old_extend_start_loc, + input_ids, + positions, + hidden_states, + out_cache_loc, + extend_seq_lens, + extend_start_loc, + next_token_ids, + seq_lens, + padding_lens, + req_pool_indices, + req_to_token, + req_to_hidden_states_pool, + # Constants/Strides + step, + old_hidden_states.stride(0), + old_hidden_states.stride(1), + req_to_hidden_states_pool.stride(0), + req_to_hidden_states_pool.stride(1), + req_to_hidden_states_pool.stride(2), + req_to_token.stride(0), + req_to_token.stride(1), + # Meta + HIDDEN_DIM=hidden_dim, + BLOCK_SEQ=BLOCK_SEQ, + BLOCK_HID=BLOCK_HID, + ) + + +@triton.jit +def assign_hidden_states_pool_kernel( + hidden_states_ptr, + req_pool_indices_ptr, + req_to_hidden_states_pool_ptr, + extend_seq_lens_ptr, + extend_start_loc_ptr, + stride_hidden_seq, + stride_hidden_dim, + stride_pool_req, + stride_pool_step, + stride_pool_dim, + HIDDEN_DIM: tl.constexpr, + pool_size: tl.constexpr, + BLOCK_HID: tl.constexpr, +): + pid = tl.program_id(0) + + extend_len = tl.load(extend_seq_lens_ptr + pid) + start_loc = tl.load(extend_start_loc_ptr + pid) + end_loc = start_loc + extend_len + + req_idx = tl.load(req_pool_indices_ptr + pid) + pool_vec_offset_base = req_idx * stride_pool_req + + for i in range(pool_size): + for off_h in range(0, HIDDEN_DIM, BLOCK_HID): + offs_h = off_h + tl.arange(0, BLOCK_HID) + mask_h = offs_h < HIDDEN_DIM + + hid_ptr = ( + hidden_states_ptr + + (end_loc - pool_size + i) * stride_hidden_seq + + offs_h * stride_hidden_dim + ) + hid_val = tl.load(hid_ptr, mask=mask_h) + + pool_ptr = ( + req_to_hidden_states_pool_ptr + + pool_vec_offset_base + + i * stride_pool_step + + offs_h * stride_pool_dim + ) + tl.store(pool_ptr, hid_val, mask=mask_h) + + +def assign_hidden_states_pool_triton( + hidden_states: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_hidden_states_pool: torch.Tensor, + pool_size: int, + num_seqs: int, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, +): + grid = (num_seqs,) + assign_hidden_states_pool_kernel[grid]( + hidden_states, + req_pool_indices, + req_to_hidden_states_pool, + extend_seq_lens, + extend_start_loc, + hidden_states.stride(0), + hidden_states.stride(1), + req_to_hidden_states_pool.stride(0), + req_to_hidden_states_pool.stride(1), + req_to_hidden_states_pool.stride(2), + HIDDEN_DIM=hidden_states.shape[1], + pool_size=pool_size, + BLOCK_HID=64, + ) + + +def assign_hidden_states_pool_torch( + hidden_states: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_hidden_states_pool: torch.Tensor, + pool_size: int, + num_seqs: int, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, +): + for req in range(num_seqs): + pool_idx = req_pool_indices[req] + extend_len = extend_seq_lens[req] + start_loc = extend_start_loc[req] + end_loc = start_loc + extend_len + req_to_hidden_states_pool[pool_idx, :pool_size, :].copy_( + hidden_states[end_loc - pool_size : end_loc, :] + )