f1384f5293
Co-authored-by: Duyi-Wang <duyi.wang@amd.com> Co-authored-by: billishyahao <bill.he@amd.com> Co-authored-by: HaiShaw <hixiao@gmail.com>
315 lines
8.6 KiB
Python
315 lines
8.6 KiB
Python
import torch
|
|
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,
|
|
)
|
|
|
|
|
|
@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
|
|
|
|
|
|
# When num_kv_heads=1, we have tensors with degenerate strides,
|
|
# For example, as below, where we have stride[-3] == stride[-2]:
|
|
# - shape: [num_pages, 1, 64, 128]
|
|
# - stride: [8192, 128, 128, 1]
|
|
# This will cause TMA desc validation fail in flashinfer (trtllm-mha backend).
|
|
#
|
|
# See: https://github.com/flashinfer-ai/flashinfer/issues/2232
|
|
def canonicalize_stride(tensor: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Adjust degenerate strides for a tensor, make it canonical.
|
|
"""
|
|
sizes = tensor.size()
|
|
strides = tensor.stride()
|
|
ndim = tensor.dim()
|
|
|
|
need_fix = any(
|
|
sizes[i] == 1 and strides[i] == strides[i + 1] for i in range(ndim - 1)
|
|
)
|
|
|
|
if not need_fix:
|
|
return tensor
|
|
|
|
# canonicalize the stride
|
|
# Example:
|
|
# - shape: [num_pages, 1, 64, 128]
|
|
# - stride: [8192, 128, 128, 1] (wrong!)
|
|
# Gives new stride: [8192, 8192, 128 ,1] (correct!)
|
|
new_strides = [0] * ndim
|
|
new_strides[-1] = 1
|
|
for i in range(ndim - 2, -1, -1):
|
|
new_strides[i] = new_strides[i + 1] * sizes[i + 1]
|
|
|
|
return tensor.as_strided(sizes, new_strides)
|