3c882db3ad
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
351 lines
10 KiB
Python
351 lines
10 KiB
Python
# 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, :]
|
|
)
|