114 lines
3.2 KiB
Python
114 lines
3.2 KiB
Python
import math
|
|
from typing import Optional, Tuple, Union
|
|
|
|
import ixformer._C as ops
|
|
import torch
|
|
|
|
|
|
class BatchPrefillWithRaggedKVCacheWrapper:
|
|
def __init__(
|
|
self,
|
|
float_workspace_buffer: torch.Tensor,
|
|
kv_layout: str = "NHD",
|
|
):
|
|
pass
|
|
|
|
def plan(
|
|
self,
|
|
qo_indptr: torch.Tensor,
|
|
kv_indptr: torch.Tensor,
|
|
num_qo_heads: int,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
# custom_mask: Optional[torch.Tensor] = None,
|
|
# packed_custom_mask: Optional[torch.Tensor] = None,
|
|
causal: bool = True,
|
|
# pos_encoding_mode: str = "NONE",
|
|
# allow_fp16_qk_reduction: bool = False,
|
|
# window_left: int = -1,
|
|
# logits_soft_cap: Optional[float] = None,
|
|
sm_scale: Optional[float] = None,
|
|
# rope_scale: Optional[float] = None,
|
|
# rope_theta: Optional[float] = None,
|
|
# q_data_type: str = "float16",
|
|
) -> None:
|
|
batch_size = len(qo_indptr) - 1
|
|
if len(kv_indptr) != batch_size + 1:
|
|
raise ValueError(
|
|
"The kv_indptr length should be equal to qk_indptr length."
|
|
)
|
|
self._causal = causal
|
|
self._sm_scale = sm_scale
|
|
if sm_scale is None:
|
|
sm_scale = 1.0 / math.sqrt(head_dim)
|
|
|
|
self.cu_seqlens_q = qo_indptr
|
|
self.cu_seqlens_k = kv_indptr
|
|
self.num_qo_heads = num_qo_heads
|
|
self.num_kv_heads = num_kv_heads
|
|
self.head_dim = head_dim
|
|
self.max_seqlen_q = max_seqlen_q
|
|
self.max_seqlen_k = max_seqlen_k
|
|
|
|
begin_forward = plan
|
|
|
|
def forward(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
causal: bool = True,
|
|
# pos_encoding_mode: str = "NONE",
|
|
# allow_fp16_qk_reduction: bool = False,
|
|
# window_left: int = -1,
|
|
logits_soft_cap: Optional[float] = None,
|
|
sm_scale: Optional[float] = None,
|
|
# rope_scale: Optional[float] = None,
|
|
# rope_theta: Optional[float] = None,
|
|
) -> torch.Tensor:
|
|
r"""Warning: This function is deprecated, please use :meth:`run` instead."""
|
|
|
|
q = q.view(-1, self.num_qo_heads, self.head_dim)
|
|
k = k.view(-1, self.num_kv_heads, self.head_dim)
|
|
v = v.view(-1, self.num_kv_heads, self.head_dim)
|
|
|
|
out = torch.empty_like(q)
|
|
|
|
assert causal
|
|
assert (
|
|
logits_soft_cap is None or logits_soft_cap == 0
|
|
), f"logits_soft_cap not supported, but got logits_soft_cap={logits_soft_cap}"
|
|
|
|
ops.infer.ixinfer_flash_attn_unpad(
|
|
q,
|
|
k,
|
|
v,
|
|
out,
|
|
self.cu_seqlens_q,
|
|
self.cu_seqlens_k,
|
|
self.max_seqlen_q,
|
|
self.max_seqlen_k,
|
|
causal,
|
|
False, # need_lse =False
|
|
sm_scale,
|
|
False,
|
|
None,
|
|
)
|
|
return out
|
|
|
|
def end_forward(self) -> None:
|
|
r"""Warning: this function is deprecated and has no effect."""
|
|
pass
|
|
|
|
|
|
class BatchPrefillWithPagedKVCacheWrapper:
|
|
def __init__(
|
|
self,
|
|
float_workspace_buffer: torch.Tensor,
|
|
kv_layout: str = "NHD",
|
|
use_cuda_graph: bool = False,
|
|
) -> None:
|
|
pass
|