661 lines
19 KiB
Python
661 lines
19 KiB
Python
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
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 sgl_kernel import concat_mla_absorb_q
|
|
|
|
|
|
@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]:
|
|
# - 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)
|
|
|
|
|
|
def mla_quantize_and_rope_for_fp8(
|
|
q_nope: torch.Tensor,
|
|
q_rope: torch.Tensor,
|
|
k_nope: torch.Tensor,
|
|
k_rope: torch.Tensor,
|
|
pos_ids: torch.Tensor,
|
|
cos_sin_cache: torch.Tensor,
|
|
is_neox: bool,
|
|
kv_lora_rank: int,
|
|
qk_rope_head_dim: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
import flashinfer.rope
|
|
|
|
"""Quantize and apply RoPE for FP8 attention path.
|
|
|
|
This function handles the FP8 quantization and RoPE application for MLA attention.
|
|
It takes separate query/key nope and rope components, applies RoPE to the rope parts,
|
|
quantizes all components to FP8, and merges the query components into a single tensor.
|
|
|
|
Args:
|
|
q_nope: Query no-position-encoding component [seq_len, num_heads, kv_lora_rank]
|
|
- expected dtype: torch.bfloat16
|
|
q_rope: Query RoPE component [seq_len, num_heads, qk_rope_head_dim]
|
|
- expected dtype: torch.bfloat16
|
|
k_nope: Key no-position-encoding component [seq_len, num_heads, kv_lora_rank]
|
|
- expected dtype: torch.bfloat16
|
|
k_rope: Key RoPE component [seq_len, num_heads, qk_rope_head_dim]
|
|
- expected dtype: torch.bfloat16
|
|
pos_ids: Position indices for each token
|
|
- expected dtype: torch.int64 or torch.int32
|
|
cos_sin_cache: Precomputed cosine/sine cache for RoPE
|
|
- expected dtype: matches q_/k_ input dtype (torch.bfloat16)
|
|
is_neox: Whether to use NeoX-style RoPE (interleaved) or GPT-style (half rotation)
|
|
kv_lora_rank: Dimension of the no-position-encoding component
|
|
qk_rope_head_dim: Dimension of the RoPE component
|
|
|
|
Returns:
|
|
tuple: (merged_q_out, k_nope_out, k_rope_out) quantized to FP8
|
|
- merged_q_out: [seq_len, num_heads, kv_lora_rank + qk_rope_head_dim], dtype=torch.float8_e4m3fn
|
|
- k_nope_out: [seq_len, num_heads, kv_lora_rank], dtype=torch.float8_e4m3fn
|
|
- k_rope_out: [seq_len, num_heads, qk_rope_head_dim], dtype=torch.float8_e4m3fn
|
|
"""
|
|
attn_dtype = torch.float8_e4m3fn
|
|
q_len, num_heads = q_rope.shape[0], q_rope.shape[1]
|
|
|
|
# Allocate output tensors with FP8 dtype
|
|
# Query output will contain merged nope + rope components
|
|
q_out = q_rope.new_empty(
|
|
q_len,
|
|
num_heads,
|
|
kv_lora_rank + qk_rope_head_dim,
|
|
dtype=attn_dtype,
|
|
)
|
|
|
|
# Key outputs maintain original shapes but with FP8 dtype
|
|
k_rope_out = k_rope.new_empty(k_rope.shape, dtype=attn_dtype)
|
|
k_nope_out = k_nope.new_empty(k_nope.shape, dtype=attn_dtype)
|
|
|
|
# Apply RoPE and quantize all components in a single fused kernel call
|
|
# This kernel handles:
|
|
# 1. RoPE application to q_rope and k_rope using cos_sin_cache and positions
|
|
# 2. Quantization of all components to FP8 format
|
|
# 3. Output placement into pre-allocated tensors
|
|
flashinfer.rope.mla_rope_quantize_fp8(
|
|
q_rope=q_rope,
|
|
k_rope=k_rope,
|
|
q_nope=q_nope,
|
|
k_nope=k_nope,
|
|
cos_sin_cache=cos_sin_cache,
|
|
pos_ids=pos_ids,
|
|
is_neox=is_neox,
|
|
quantize_dtype=attn_dtype,
|
|
# Output tensor slicing: q_out contains [nope_part, rope_part]
|
|
q_rope_out=q_out[..., kv_lora_rank:], # RoPE part goes to end
|
|
k_rope_out=k_rope_out,
|
|
q_nope_out=q_out[..., :kv_lora_rank], # Nope part goes to beginning
|
|
k_nope_out=k_nope_out,
|
|
# Quantization scales (set to 1.0 for no additional scaling)
|
|
quant_scale_q=1.0,
|
|
quant_scale_kv=1.0,
|
|
)
|
|
|
|
return q_out, k_nope_out, k_rope_out
|
|
|
|
|
|
def concat_mla_absorb_q_general(q_nope, q_rope):
|
|
if _is_cuda and q_nope.shape[-1] == 512 and q_rope.shape[-1] == 64:
|
|
return concat_mla_absorb_q(q_nope, q_rope)
|
|
else:
|
|
return torch.cat([q_nope, q_rope], dim=-1)
|
|
|
|
|
|
@triton.jit
|
|
def reshape_and_cache_flash(
|
|
key_ptr,
|
|
value_ptr,
|
|
key_cache_ptr,
|
|
value_cache_ptr,
|
|
slot_mapping_ptr,
|
|
swa_slot_mapping_ptr,
|
|
k_scale_ptr,
|
|
v_scale_ptr,
|
|
block_stride,
|
|
key_stride,
|
|
value_stride,
|
|
num_heads,
|
|
head_size,
|
|
block_size,
|
|
HEAD_BLOCK: tl.constexpr,
|
|
BLOCK_D: tl.constexpr,
|
|
HAS_SWA: tl.constexpr,
|
|
USE_SCALE: tl.constexpr,
|
|
):
|
|
"""
|
|
Triton kernel for reshaping per-token K/V tensors into paged KV cache layout.
|
|
|
|
Source layout:
|
|
key/value: [num_tokens, num_heads, head_size]
|
|
|
|
Target cache layout:
|
|
cache: [num_blocks, block_size, num_heads, head_size]
|
|
|
|
Each Triton program instance handles:
|
|
- one token (program_id(0))
|
|
- one block of heads (program_id(1))
|
|
|
|
Features:
|
|
- optional SWA slot remapping
|
|
- optional FP8 scale dequantization before cache write
|
|
|
|
Args:
|
|
key_ptr: Pointer to source key tensor.
|
|
value_ptr: Pointer to source value tensor.
|
|
key_cache_ptr: Pointer to destination key cache tensor.
|
|
value_cache_ptr: Pointer to destination value cache tensor.
|
|
slot_mapping_ptr: Maps token -> cache slot.
|
|
swa_slot_mapping_ptr: Optional second-stage slot remap for SWA mode.
|
|
k_scale_ptr: Optional key scaling factor pointer.
|
|
v_scale_ptr: Optional value scaling factor pointer.
|
|
block_stride: Stride between cache blocks.
|
|
key_stride: Stride between source key tokens.
|
|
value_stride: Stride between source value tokens.
|
|
num_heads: Number of attention heads.
|
|
head_size: Hidden dimension per head.
|
|
block_size: Number of slots per cache block.
|
|
HEAD_BLOCK: Number of heads processed per program.
|
|
BLOCK_D: Vectorized dimension size (power-of-2 padded).
|
|
HAS_SWA: Enable SWA remapping.
|
|
USE_SCALE: Enable scale division before storing.
|
|
"""
|
|
|
|
# ----------------------------------
|
|
# program ids
|
|
# pid0 = token
|
|
# pid1 = head block
|
|
# ----------------------------------
|
|
token_idx = tl.program_id(0)
|
|
head_block_idx = tl.program_id(1)
|
|
|
|
# ----------------------------------
|
|
# slot mapping
|
|
# ----------------------------------
|
|
slot_idx = tl.load(slot_mapping_ptr + token_idx)
|
|
|
|
if HAS_SWA:
|
|
slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx)
|
|
|
|
if slot_idx < 0:
|
|
return
|
|
|
|
block_idx = slot_idx // block_size
|
|
block_offset = slot_idx % block_size
|
|
|
|
# ----------------------------------
|
|
# head range
|
|
# ----------------------------------
|
|
head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)
|
|
|
|
head_mask = head_idx < num_heads
|
|
|
|
dim_idx = tl.arange(0, BLOCK_D)
|
|
|
|
# shape = [HEAD_BLOCK, BLOCK_D]
|
|
offs = head_idx[:, None] * head_size + dim_idx[None, :]
|
|
|
|
mask = head_mask[:, None] & (dim_idx[None, :] < head_size)
|
|
|
|
# ----------------------------------
|
|
# source load
|
|
# ----------------------------------
|
|
src_key = token_idx * key_stride + offs
|
|
src_value = token_idx * value_stride + offs
|
|
|
|
k = tl.load(key_ptr + src_key, mask=mask)
|
|
v = tl.load(value_ptr + src_value, mask=mask)
|
|
|
|
# ----------------------------------
|
|
# optional scale
|
|
# ----------------------------------
|
|
if USE_SCALE:
|
|
k_scale = tl.load(k_scale_ptr)
|
|
v_scale = tl.load(v_scale_ptr)
|
|
|
|
k = k / k_scale
|
|
v = v / v_scale
|
|
|
|
# ----------------------------------
|
|
# target layout
|
|
# [block_idx, block_offset, head, dim]
|
|
# ----------------------------------
|
|
tgt = block_idx * block_stride + block_offset * num_heads * head_size + offs
|
|
|
|
tl.store(key_cache_ptr + tgt, k, mask=mask)
|
|
tl.store(value_cache_ptr + tgt, v, mask=mask)
|
|
|
|
|
|
def launch_reshape_and_cache_flash(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
swa_slot_mapping=None,
|
|
k_scale=None,
|
|
v_scale=None,
|
|
):
|
|
"""
|
|
Launch wrapper for reshape_and_cache_flash Triton kernel.
|
|
|
|
This wrapper prepares launch configuration and dispatches the Triton kernel
|
|
that writes token-major K/V tensors into paged KV cache layout.
|
|
|
|
Args:
|
|
key: Source key tensor [num_tokens, num_heads, head_size]
|
|
value: Source value tensor [num_tokens, num_heads, head_size]
|
|
key_cache: Destination key cache [num_blocks, block_size, num_heads, head_size]
|
|
value_cache: Destination value cache [num_blocks, block_size, num_heads, head_size]
|
|
slot_mapping: Token-to-cache slot mapping
|
|
swa_slot_mapping: Optional SWA remapping table
|
|
k_scale: Optional key scaling factor
|
|
v_scale: Optional value scaling factor
|
|
"""
|
|
|
|
num_tokens = key.shape[0]
|
|
num_heads = key.shape[1]
|
|
head_size = key.shape[2]
|
|
|
|
HEAD_BLOCK = 4
|
|
|
|
BLOCK_D = triton.next_power_of_2(head_size)
|
|
|
|
grid = (
|
|
num_tokens,
|
|
triton.cdiv(num_heads, HEAD_BLOCK),
|
|
)
|
|
|
|
reshape_and_cache_flash[grid](
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
swa_slot_mapping if swa_slot_mapping is not None else key,
|
|
k_scale if k_scale is not None else key,
|
|
v_scale if v_scale is not None else key,
|
|
key_cache.stride(0),
|
|
key.stride(0),
|
|
value.stride(0),
|
|
num_heads,
|
|
head_size,
|
|
key_cache.shape[1],
|
|
HEAD_BLOCK=HEAD_BLOCK,
|
|
BLOCK_D=BLOCK_D,
|
|
HAS_SWA=(swa_slot_mapping is not None),
|
|
USE_SCALE=(k_scale is not None),
|
|
)
|