215 lines
7.7 KiB
Python
215 lines
7.7 KiB
Python
from typing import Type
|
|
from dataclasses import dataclass
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
from cutlass.cute.nvgpu import cpasync
|
|
from cutlass import Int32, const_expr
|
|
|
|
import sglang.jit_kernel.flash_attention.cute.utils as utils
|
|
from .cute_dsl_utils import ParamsBase
|
|
from cutlass.cute import FastDivmodDivisor
|
|
|
|
import math
|
|
|
|
|
|
@dataclass
|
|
class PagedKVManager(ParamsBase):
|
|
mPageTable: cute.Tensor
|
|
mK_paged: cute.Tensor
|
|
mV_paged: cute.Tensor
|
|
thread_idx: Int32
|
|
|
|
page_size_divmod: FastDivmodDivisor
|
|
seqlen_k: Int32
|
|
leftpad_k: Int32
|
|
n_block_size: Int32
|
|
num_threads: cutlass.Constexpr[Int32]
|
|
head_dim_padded: cutlass.Constexpr[Int32]
|
|
head_dim_v_padded: cutlass.Constexpr[Int32]
|
|
|
|
gmem_threads_per_row: cutlass.Constexpr[Int32]
|
|
page_entry_per_thread: Int32
|
|
async_copy_elems: Int32
|
|
|
|
gmem_tiled_copy_KV: cute.TiledCopy
|
|
gmem_thr_copy_KV: cute.TiledCopy
|
|
tPrPage: cute.Tensor
|
|
tPrPageOffset: cute.Tensor
|
|
tKpK: cute.Tensor
|
|
tVpV: cute.Tensor
|
|
|
|
@staticmethod
|
|
def create(
|
|
mPageTable: cute.Tensor,
|
|
mK_paged: cute.Tensor,
|
|
mV_paged: cute.Tensor,
|
|
page_size_divmod: FastDivmodDivisor,
|
|
bidb: Int32,
|
|
bidh: Int32,
|
|
thread_idx: Int32,
|
|
seqlen_k: Int32,
|
|
leftpad_k: Int32,
|
|
n_block_size: cutlass.Constexpr[Int32],
|
|
head_dim_padded: cutlass.Constexpr[Int32],
|
|
head_dim_v_padded: cutlass.Constexpr[Int32],
|
|
num_threads: cutlass.Constexpr[Int32],
|
|
dtype: Type[cutlass.Numeric],
|
|
):
|
|
universal_copy_bits = 128
|
|
async_copy_elems = universal_copy_bits // dtype.width
|
|
dtype_bytes = dtype.width // 8
|
|
gmem_k_block_size = math.gcd(
|
|
head_dim_padded,
|
|
head_dim_v_padded,
|
|
128 // dtype_bytes,
|
|
)
|
|
assert gmem_k_block_size % async_copy_elems == 0
|
|
gmem_threads_per_row = gmem_k_block_size // async_copy_elems
|
|
assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
|
|
atom_async_copy = cute.make_copy_atom(
|
|
cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
|
|
dtype,
|
|
num_bits_per_copy=universal_copy_bits,
|
|
)
|
|
thr_layout = cute.make_ordered_layout(
|
|
(num_threads // gmem_threads_per_row, gmem_threads_per_row),
|
|
order=(1, 0),
|
|
)
|
|
val_layout = cute.make_layout((1, async_copy_elems))
|
|
gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout)
|
|
gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
|
|
page_entry_per_thread = n_block_size // num_threads
|
|
|
|
tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
|
|
tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
|
|
|
|
mPageTable = mPageTable[bidb, None]
|
|
mK_paged = mK_paged[None, None, bidh, None]
|
|
mV_paged = mV_paged[None, None, bidh, None]
|
|
|
|
cK = cute.make_identity_tensor((n_block_size, head_dim_padded))
|
|
tKcK = gmem_thr_copy_KV.partition_S(cK)
|
|
tKpK = utils.predicate_k(tKcK, limit=mK_paged.shape[1])
|
|
|
|
if const_expr(head_dim_padded == head_dim_v_padded):
|
|
tVpV = tKpK
|
|
else:
|
|
cV = cute.make_identity_tensor((n_block_size, head_dim_v_padded))
|
|
tVcV = gmem_thr_copy_KV.partition_S(cV)
|
|
tVpV = utils.predicate_k(tVcV, limit=mV_paged.shape[0])
|
|
|
|
return PagedKVManager(
|
|
mPageTable,
|
|
mK_paged,
|
|
mV_paged,
|
|
thread_idx,
|
|
page_size_divmod,
|
|
seqlen_k,
|
|
leftpad_k,
|
|
n_block_size,
|
|
num_threads,
|
|
head_dim_padded,
|
|
head_dim_v_padded,
|
|
gmem_threads_per_row,
|
|
page_entry_per_thread,
|
|
async_copy_elems,
|
|
gmem_tiled_copy_KV,
|
|
gmem_thr_copy_KV,
|
|
tPrPage,
|
|
tPrPageOffset,
|
|
tKpK,
|
|
tVpV,
|
|
)
|
|
|
|
@cute.jit
|
|
def load_page_table(self, n_block: Int32):
|
|
for i in cutlass.range(self.page_entry_per_thread, unroll=1):
|
|
row = (
|
|
i * self.num_threads
|
|
+ (self.thread_idx % self.gmem_threads_per_row)
|
|
* (self.num_threads // self.gmem_threads_per_row)
|
|
+ (self.thread_idx // self.gmem_threads_per_row)
|
|
)
|
|
row_idx = n_block * self.n_block_size + row
|
|
|
|
page_idx, page_offset = divmod(row_idx + self.leftpad_k, self.page_size_divmod)
|
|
|
|
is_valid = (
|
|
(i + 1) * self.num_threads <= self.n_block_size or row < self.n_block_size
|
|
) and row_idx < self.seqlen_k
|
|
page = self.mPageTable[page_idx] if is_valid else 0
|
|
|
|
self.tPrPage[i] = page
|
|
self.tPrPageOffset[i] = page_offset
|
|
|
|
@cute.jit
|
|
def compute_X_ptr(self, K_or_V: str):
|
|
tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
|
|
for i in cutlass.range(self.page_entry_per_thread, unroll=1):
|
|
page = self.tPrPage[i]
|
|
page_offset = self.tPrPageOffset[i]
|
|
if const_expr(K_or_V == "K"):
|
|
tPrXPtr[i] = utils.elem_pointer(self.mK_paged, (page_offset, 0, page)).toint()
|
|
else:
|
|
tPrXPtr[i] = utils.elem_pointer(self.mV_paged, (0, page_offset, page)).toint()
|
|
return tPrXPtr
|
|
|
|
@cute.jit
|
|
def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
|
|
assert K_or_V in ("K", "V")
|
|
|
|
tPrXPtr = self.compute_X_ptr(K_or_V)
|
|
|
|
# Finesse sX layout to be (M, N).
|
|
sX_pi = cute.make_tensor(
|
|
sX.iterator,
|
|
cute.make_layout(
|
|
(sX.shape[0][0], (sX.shape[0][1], sX.shape[2])),
|
|
stride=(sX.stride[0][0], (sX.stride[0][1], sX.stride[2])),
|
|
),
|
|
)
|
|
|
|
if const_expr(K_or_V == "V"):
|
|
# Need to transpose V
|
|
sX_pi = cute.make_tensor(sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0]))
|
|
|
|
head_dim = self.head_dim_v_padded if const_expr(K_or_V == "V") else self.head_dim_padded
|
|
cX = cute.make_identity_tensor((self.n_block_size, head_dim))
|
|
tXsX = self.gmem_thr_copy_KV.partition_D(sX_pi)
|
|
tXcX = self.gmem_thr_copy_KV.partition_S(cX)
|
|
tXc0X = self.gmem_thr_copy_KV.get_slice(0).partition_S(cX)
|
|
|
|
seqlenk_row_limit = (
|
|
self.seqlen_k - n_block * self.n_block_size - tXcX[0][0] if n_block >= 0 else 0
|
|
)
|
|
for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
|
|
row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit
|
|
should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean)
|
|
should_load.fill(row_valid)
|
|
|
|
x_ptr_i64 = utils.shuffle_sync(
|
|
tPrXPtr[m // self.gmem_threads_per_row],
|
|
m % self.gmem_threads_per_row,
|
|
width=self.gmem_threads_per_row,
|
|
)
|
|
x_gmem_ptr = cute.make_ptr(
|
|
self.mK_paged.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
|
|
)
|
|
mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
|
|
mX_paged_cur_copy = cute.tiled_divide(mX_paged_cur, (self.async_copy_elems,))
|
|
|
|
for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
|
|
ki = tXcX[0, 0, k][1] // self.async_copy_elems
|
|
mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki]
|
|
tXsX_k = tXsX[None, m, k]
|
|
mX_paged_cur_copy_ki = cute.make_tensor(
|
|
mX_paged_cur_copy_ki.iterator, tXsX_k.layout
|
|
)
|
|
cute.copy(
|
|
self.gmem_tiled_copy_KV,
|
|
mX_paged_cur_copy_ki,
|
|
tXsX_k,
|
|
pred=should_load,
|
|
)
|