414 lines
12 KiB
Python
414 lines
12 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
|
|
|
|
|
|
# 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)
|