87a19d2d00
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
2034 lines
73 KiB
Python
2034 lines
73 KiB
Python
import math
|
|
from typing import Optional, Union
|
|
|
|
import ixformer._C as ops
|
|
import ixformer._C._functions as CF
|
|
import torch
|
|
|
|
from ixformer.core import config
|
|
|
|
from .linear import linear
|
|
from .paged_attention import paged_attention as paged_attention_ixformer_impl
|
|
|
|
__all__ = [
|
|
"ref_vllm_paged_attention",
|
|
"vllm_paged_attention",
|
|
"ref_vllm_paged_attention_mla",
|
|
"vllm_paged_attention_mla",
|
|
"vllm_paged_attention_mla_fused",
|
|
"ref_vllm_paged_attention_mla_int8",
|
|
"vllm_paged_attention_mla_int8",
|
|
"ref_vllm_paged_attention_v5",
|
|
"vllm_paged_attention_v5",
|
|
"ref_vllm_paged_attention_v4",
|
|
"vllm_paged_attention_v4",
|
|
"ref_vllm_reshape_and_cache_v4",
|
|
"vllm_reshape_and_cache_v4",
|
|
"ref_vllm_reshape_and_cache",
|
|
"vllm_reshape_and_cache",
|
|
"vllm_cache_ops_reshape_and_cache",
|
|
"ref_reshape_and_cache_flash",
|
|
"reshape_and_cache_flash",
|
|
"ref_vllm_rotary_embedding",
|
|
"vllm_rotary_embedding",
|
|
"ref_vllm_rotary_embedding_phi",
|
|
"vllm_rotary_embedding_phi",
|
|
"ref_vllm_batched_rotary_embedding",
|
|
"vllm_batched_rotary_embedding",
|
|
"ref_vllm_copy_blocks",
|
|
"vllm_copy_blocks",
|
|
"ref_vllm_swap_blocks",
|
|
"vllm_swap_blocks",
|
|
"vllm_gather_cache",
|
|
"vllm_gather_cache_int8",
|
|
"ref_vllm_gather_cache_int8",
|
|
"ref_vllm_gather_cache",
|
|
"ref_vllm_concat_and_cache_mla",
|
|
"vllm_concat_and_cache_mla",
|
|
"ref_vllm_concat_and_cache_mla_int8",
|
|
"vllm_concat_and_cache_mla_int8",
|
|
"vllm_llama_mlp",
|
|
"gptq_gemm",
|
|
"vllm_gptq_shuffle",
|
|
"vllm_moe_topk_softmax",
|
|
"vllm_moe_align_block_size",
|
|
"ref_vllm_invoke_fused_moe_kernel",
|
|
"vllm_invoke_fused_moe_kernel",
|
|
"advance_step_flashattn",
|
|
"weak_ref_tensor",
|
|
# customized ops
|
|
"vllm_rotary_embedding_with_key_layer_norm",
|
|
"ref_vllm_rotary_embedding_with_key_layer_norm",
|
|
]
|
|
|
|
weak_ref_tensor = ops.infer.weak_ref_tensor
|
|
|
|
|
|
def ref_vllm_paged_attention(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
softcap: float = 0.0,
|
|
window_left: int = -1,
|
|
window_right: int = -1,
|
|
use_sqrt_alibi: bool = False,
|
|
):
|
|
assert window_right in [-1, 0]
|
|
|
|
def get_alibi_mask(num_heads, seqlen, device, dtype):
|
|
x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1)
|
|
y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1)
|
|
offsets = -(y - x).view(1, 1, seqlen)
|
|
return offsets
|
|
|
|
def ref_masked_attention(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
scale: float,
|
|
attn_mask: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
query = query * scale
|
|
dtype = query.dtype
|
|
device = query.device
|
|
query = query.to(torch.float32)
|
|
key = key.to(torch.float32)
|
|
value = value.to(torch.float32)
|
|
attn = torch.einsum("qhd,khd->hqk", query, key)
|
|
if attn_mask is not None:
|
|
attn_mask = attn_mask
|
|
attn = attn + attn_mask
|
|
attn = torch.softmax(attn, dim=-1)
|
|
out = torch.einsum("hqk,khd->qhd", attn, value)
|
|
out = out.to(device).to(dtype)
|
|
return out
|
|
|
|
head_size = query.shape[-1]
|
|
num_query_heads = query.shape[1]
|
|
num_kv_heads = value_cache.shape[1]
|
|
num_input_tokens = query.shape[0]
|
|
|
|
num_q_per_kv = num_query_heads // num_kv_heads
|
|
slopes = (
|
|
alibi_slopes.view(num_query_heads, 1, 1)
|
|
if alibi_slopes is not None
|
|
else alibi_slopes
|
|
)
|
|
|
|
for i in range(num_input_tokens):
|
|
q = query[i].unsqueeze(0)
|
|
block_table = block_tables[i]
|
|
context_len = int(context_lens[i])
|
|
|
|
keys = []
|
|
values = []
|
|
for j in range(context_len):
|
|
block_number = int(block_table[j // block_size])
|
|
block_offset = j % block_size
|
|
|
|
k = key_cache[block_number, :, block_offset, :]
|
|
keys.append(k)
|
|
|
|
v = value_cache[block_number, :, block_offset, :]
|
|
values.append(v)
|
|
keys = torch.stack(keys, dim=0)
|
|
values = torch.stack(values, dim=0)
|
|
if num_q_per_kv > 1:
|
|
keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1)
|
|
values = torch.repeat_interleave(values, num_q_per_kv, dim=1)
|
|
if alibi_slopes is not None:
|
|
offsets = get_alibi_mask(
|
|
num_query_heads, context_len, output.device, output.dtype
|
|
)
|
|
mask = offsets * slopes
|
|
mask = mask.to(output.dtype)
|
|
if window_left != -1:
|
|
index = torch.ones_like(mask, dtype=torch.int32, device=mask.device)
|
|
index[:, :, (context_len - 1 - window_left) :] = 0
|
|
index = index.bool()
|
|
mask.masked_fill_(index, float("-inf"))
|
|
else:
|
|
if window_left != -1:
|
|
mask = torch.zeros([1, 1, context_len], dtype=q.dtype, device=q.device)
|
|
index = torch.ones_like(mask, dtype=torch.int32, device=mask.device)
|
|
index[:, :, (context_len - 1 - window_left) :] = 0
|
|
index = index.bool()
|
|
mask.masked_fill_(index, float("-inf"))
|
|
else:
|
|
mask = None
|
|
|
|
out = ref_masked_attention(
|
|
q,
|
|
keys,
|
|
values,
|
|
scale,
|
|
mask,
|
|
)
|
|
out = out.view(num_query_heads, head_size)
|
|
if softcap != 0.0:
|
|
out = softcap * torch.tanh(out / softcap)
|
|
output[i].copy_(out, non_blocking=True)
|
|
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_ixinfer(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
softcap: float = 0.0,
|
|
causal: bool = True,
|
|
window_left: int = -1,
|
|
window_right: int = -1,
|
|
use_cuda_graph: bool = False,
|
|
use_sqrt_alibi: bool = False,
|
|
):
|
|
"""
|
|
Arguments:
|
|
query: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size]
|
|
key_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size]
|
|
value_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size]
|
|
num_kv_heads: int
|
|
scale: float
|
|
block_tables: [torch.int64] [num_tokens, max_num_blocks_per_seq]
|
|
context_lens: [torch.int32] [num_tokens]
|
|
block_size: int
|
|
max_context_len: int
|
|
alibi_slopes: [torch.float32] [num_heads]
|
|
softcap: float
|
|
causal: bool
|
|
window_left: int
|
|
window_right: int
|
|
use_sqrt_alibi: bool: False
|
|
Return:
|
|
output: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size]
|
|
"""
|
|
ops.infer.vllm_paged_attention(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
causal,
|
|
window_left,
|
|
window_right,
|
|
softcap,
|
|
use_cuda_graph,
|
|
use_sqrt_alibi,
|
|
)
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_ixformer(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
softcap: float = 0.0,
|
|
use_sqrt_alibi: bool = False,
|
|
need_view: bool = True,
|
|
):
|
|
"""
|
|
Args:
|
|
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
num_kv_heads: int
|
|
scale: float
|
|
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
|
context_lens_cpu: (num_tokens) torch.int32
|
|
context_lens: (num_tokens) torch.int32
|
|
block_size: int
|
|
max_context_len: int
|
|
alibi_slopes: (num_heads) torch.float32
|
|
use_sqrt_alibi: bool
|
|
Returns:
|
|
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
"""
|
|
|
|
if need_view:
|
|
num_blocks = key_cache.size(0)
|
|
head_size = query.size(-1)
|
|
key_cache = key_cache.view(num_blocks, num_kv_heads, block_size, head_size)
|
|
value_cache = value_cache.view(num_blocks, num_kv_heads, block_size, head_size)
|
|
paged_attention_ixformer_impl(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
use_sqrt_alibi,
|
|
)
|
|
return output
|
|
|
|
|
|
def ref_vllm_paged_attention_mla(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
max_context_len: int,
|
|
):
|
|
def ref_masked_attention(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
scale: float,
|
|
attn_mask: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
query = query * scale
|
|
dtype = query.dtype
|
|
device = query.device
|
|
attn = torch.einsum("qhd,khd->hqk", query, key)
|
|
if attn_mask is not None:
|
|
attn_mask = attn_mask
|
|
attn = attn + attn_mask
|
|
attn = attn.to(torch.float)
|
|
attn = torch.softmax(attn, dim=-1)
|
|
value = value.to(torch.float)
|
|
out = torch.einsum("hqk,khd->qhd", attn, value)
|
|
out = out.to(device).to(dtype)
|
|
return out
|
|
|
|
num_heads = query.shape[-2]
|
|
kv_lora_rank = output.shape[-1]
|
|
block_size = kv_cache.shape[1]
|
|
num_input_tokens = query.shape[0]
|
|
|
|
for i in range(num_input_tokens):
|
|
q = query[i].unsqueeze(0)
|
|
block_table = block_tables[i]
|
|
context_len = int(context_lens[i])
|
|
|
|
keys = []
|
|
values = []
|
|
for j in range(context_len):
|
|
block_number = int(block_table[j // block_size])
|
|
block_offset = j % block_size
|
|
|
|
k = kv_cache[block_number, block_offset, :]
|
|
keys.append(k)
|
|
|
|
v = kv_cache[block_number, block_offset, :kv_lora_rank]
|
|
values.append(v)
|
|
keys = torch.stack(keys, dim=0).unsqueeze(-2).repeat(1, num_heads, 1)
|
|
values = torch.stack(values, dim=0).unsqueeze(-2).repeat(1, num_heads, 1)
|
|
mask = None
|
|
|
|
out = ref_masked_attention(
|
|
q,
|
|
keys,
|
|
values,
|
|
scale,
|
|
mask,
|
|
)
|
|
out = out.view(num_heads, kv_lora_rank)
|
|
output[i].copy_(out, non_blocking=True)
|
|
|
|
return output
|
|
|
|
|
|
def ref_vllm_paged_attention_mla_int8(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
query_scale: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
kv_cache_scale: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
max_context_len: int,
|
|
):
|
|
def ref_masked_attention(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
scale: float,
|
|
attn_mask: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
query = query * scale
|
|
dtype = query.dtype
|
|
device = query.device
|
|
attn = torch.einsum("qhd,khd->hqk", query, key)
|
|
if attn_mask is not None:
|
|
attn_mask = attn_mask
|
|
attn = attn + attn_mask
|
|
attn = attn.to(torch.float)
|
|
attn = torch.softmax(attn, dim=-1)
|
|
value = value.to(torch.float)
|
|
out = torch.einsum("hqk,khd->qhd", attn, value)
|
|
out = out.to(device).to(dtype)
|
|
return out
|
|
|
|
# dequant q
|
|
num_heads = query.shape[-2]
|
|
kv_lora_rank = output.shape[-1]
|
|
block_size = kv_cache.shape[1]
|
|
num_input_tokens = query.shape[0]
|
|
query = query * query_scale.unsqueeze(-1)
|
|
|
|
for i in range(num_input_tokens):
|
|
q = query[i].unsqueeze(0)
|
|
block_table = block_tables[i]
|
|
context_len = int(context_lens[i])
|
|
|
|
keys = []
|
|
values = []
|
|
for j in range(context_len):
|
|
block_number = int(block_table[j // block_size])
|
|
block_offset = j % block_size
|
|
|
|
k = kv_cache[block_number, block_offset, :kv_lora_rank]
|
|
k_scale = kv_cache_scale[block_number, block_offset, 0]
|
|
k_pe = kv_cache[block_number, block_offset, kv_lora_rank:]
|
|
k_pe_scale = kv_cache_scale[block_number, block_offset, 1]
|
|
k = k * k_scale
|
|
v = k
|
|
k_pe = k_pe * k_pe_scale
|
|
k = torch.cat((k, k_pe), dim=-1)
|
|
keys.append(k)
|
|
values.append(v)
|
|
keys = torch.stack(keys, dim=0).unsqueeze(-2).repeat(1, num_heads, 1)
|
|
values = torch.stack(values, dim=0).unsqueeze(-2).repeat(1, num_heads, 1)
|
|
mask = None
|
|
|
|
out = ref_masked_attention(
|
|
q,
|
|
keys,
|
|
values,
|
|
scale,
|
|
mask,
|
|
)
|
|
out = out.view(num_heads, kv_lora_rank)
|
|
output[i].copy_(out, non_blocking=True)
|
|
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_mla(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
max_context_len: int,
|
|
use_cuda_graph: bool = False,
|
|
):
|
|
"""
|
|
Args:
|
|
output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16
|
|
query: (num_tokens, num_heads, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16
|
|
kv_cache: (num_blocks, block_size, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16
|
|
scale: float
|
|
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
|
context_lens: (num_tokens) torch.int32
|
|
max_context_len: int
|
|
use_cuda_graph: bool
|
|
Returns:
|
|
output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16
|
|
"""
|
|
ops.infer.vllm_paged_attention_mla(
|
|
output,
|
|
query,
|
|
kv_cache,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
max_context_len,
|
|
use_cuda_graph,
|
|
)
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_mla_int8(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
query_scale: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
kv_cache_scale: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
max_context_len: int,
|
|
use_cuda_graph: bool = False,
|
|
):
|
|
"""
|
|
Args:
|
|
output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16
|
|
query: (num_tokens, num_heads, kv_lora_rank+qk_rope_head_dim) torch.int8
|
|
query_scale: (num_tokens, num_heads) torch.float
|
|
kv_cache: (num_blocks, block_size, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16
|
|
kv_cache_scale: (num_blocks, block_size, 2) torch.float
|
|
scale: float
|
|
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
|
context_lens: (num_tokens) torch.int32
|
|
max_context_len: int
|
|
use_cuda_graph: bool
|
|
Returns:
|
|
output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16
|
|
"""
|
|
ops.infer.vllm_paged_attention_mla_int8(
|
|
output,
|
|
query,
|
|
query_scale,
|
|
kv_cache,
|
|
kv_cache_scale,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
max_context_len,
|
|
use_cuda_graph,
|
|
)
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_mla_fused(
|
|
output: torch.Tensor,
|
|
q_nope: torch.Tensor,
|
|
q_pe: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
max_context_len: int,
|
|
k_c_normed: torch.Tensor = None,
|
|
k_pe: torch.Tensor = None,
|
|
use_cuda_graph: bool = False,
|
|
):
|
|
"""
|
|
Args:
|
|
q_nope: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16
|
|
q_pe: (num_tokens, num_heads, qk_rope_head_dim) torch.half, torch.bfloat16
|
|
kv_cache: (num_blocks, block_size, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16
|
|
scale: float
|
|
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
|
context_lens: (num_tokens) torch.int32
|
|
max_context_len: int
|
|
k_c_normed: (num_tokens, kv_lora_rank) torch.half, torch.bfloat16
|
|
k_pe: (num_tokens, qk_rope_head_dim) torch.half, torch.bfloat16
|
|
use_cuda_graph: bool
|
|
Returns:
|
|
output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16
|
|
"""
|
|
ops.infer.vllm_paged_attention_mla_fused(
|
|
output,
|
|
q_nope,
|
|
q_pe,
|
|
kv_cache,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
max_context_len,
|
|
k_c_normed,
|
|
k_pe,
|
|
use_cuda_graph,
|
|
)
|
|
return output
|
|
|
|
|
|
def ref_vllm_paged_attention_v5(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens_cpu: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
):
|
|
def get_alibi_mask(num_heads, seqlen, device, dtype):
|
|
x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1)
|
|
y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1)
|
|
offsets = -(y - x).view(1, 1, seqlen)
|
|
return offsets
|
|
|
|
def ref_masked_attention(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
scale: float,
|
|
attn_mask: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
query = query * scale
|
|
dtype = query.dtype
|
|
device = query.device
|
|
query = query.to(torch.float32).cpu()
|
|
key = key.to(torch.float32).cpu()
|
|
value = value.to(torch.float32).cpu()
|
|
attn = torch.einsum("qhd,khd->hqk", query, key)
|
|
if attn_mask is not None:
|
|
attn_mask = attn_mask.cpu()
|
|
attn = attn + attn_mask
|
|
attn = torch.softmax(attn, dim=-1)
|
|
out = torch.einsum("hqk,khd->qhd", attn, value)
|
|
out = out.to(device).to(dtype)
|
|
return out
|
|
|
|
head_size = query.shape[-1]
|
|
num_query_heads = query.shape[1]
|
|
num_kv_heads = value_cache.shape[1]
|
|
num_input_tokens = query.shape[0]
|
|
|
|
num_q_per_kv = num_query_heads // num_kv_heads
|
|
slopes = (
|
|
alibi_slopes.view(num_query_heads, 1, 1)
|
|
if alibi_slopes is not None
|
|
else alibi_slopes
|
|
)
|
|
|
|
for i in range(num_input_tokens):
|
|
q = query[i].unsqueeze(0)
|
|
block_table = block_tables[i]
|
|
context_len = int(context_lens[i])
|
|
|
|
keys = []
|
|
values = []
|
|
for j in range(context_len):
|
|
block_number = int(block_table[j // block_size])
|
|
block_offset = j % block_size
|
|
|
|
k = key_cache[block_number, :, block_offset, :]
|
|
keys.append(k)
|
|
|
|
v = value_cache[block_number, :, block_offset, :]
|
|
values.append(v)
|
|
keys = torch.stack(keys, dim=0)
|
|
values = torch.stack(values, dim=0)
|
|
if num_q_per_kv > 1:
|
|
keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1)
|
|
values = torch.repeat_interleave(values, num_q_per_kv, dim=1)
|
|
if alibi_slopes is not None:
|
|
offsets = get_alibi_mask(
|
|
num_query_heads, context_len, output.device, output.dtype
|
|
)
|
|
mask = offsets * slopes
|
|
mask = mask.to(output.dtype)
|
|
else:
|
|
mask = None
|
|
out = ref_masked_attention(
|
|
q,
|
|
keys,
|
|
values,
|
|
scale,
|
|
mask,
|
|
)
|
|
out = out.view(num_query_heads, head_size)
|
|
output[i].copy_(out, non_blocking=True)
|
|
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_v5(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens_cpu: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
use_sqrt_alibi: bool = False,
|
|
need_view: bool = True,
|
|
):
|
|
"""
|
|
Args:
|
|
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
num_kv_heads: int
|
|
scale: float
|
|
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
|
context_lens_cpu: (num_tokens) torch.int32
|
|
context_lens: (num_tokens) torch.int32
|
|
block_size: int
|
|
max_context_len: int
|
|
alibi_slopes: (num_heads) torch.float32
|
|
use_sqrt_alibi: bool
|
|
Returns:
|
|
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
"""
|
|
if need_view:
|
|
num_blocks = key_cache.size(0)
|
|
head_size = query.size(-1)
|
|
key_cache = key_cache.view(num_blocks, num_kv_heads, block_size, head_size)
|
|
value_cache = value_cache.view(num_blocks, num_kv_heads, block_size, head_size)
|
|
ops.infer.vllm_paged_attention_v5(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
scale,
|
|
block_tables,
|
|
context_lens_cpu,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
use_sqrt_alibi,
|
|
)
|
|
return output
|
|
|
|
|
|
def ref_vllm_paged_attention_v4(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens_cpu: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
use_sqrt_alibi: bool = False,
|
|
):
|
|
def get_alibi_mask(num_heads, seqlen, device, dtype):
|
|
x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1)
|
|
y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1)
|
|
offsets = -(y - x).view(1, 1, seqlen)
|
|
return offsets
|
|
|
|
def ref_masked_attention(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
scale: float,
|
|
attn_mask: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
query = query * scale
|
|
dtype = query.dtype
|
|
device = query.device
|
|
query = query.to(torch.float32).cpu()
|
|
key = key.to(torch.float32).cpu()
|
|
value = value.to(torch.float32).cpu()
|
|
attn = torch.einsum("qhd,khd->hqk", query, key)
|
|
if attn_mask is not None:
|
|
attn_mask = attn_mask.cpu()
|
|
attn = attn + attn_mask
|
|
attn = torch.softmax(attn, dim=-1)
|
|
out = torch.einsum("hqk,khd->qhd", attn, value)
|
|
out = out.to(device).to(dtype)
|
|
return out
|
|
|
|
num_query_heads = query.shape[1]
|
|
num_kv_heads = value_cache.shape[1]
|
|
head_size = value_cache.shape[2]
|
|
block_size = value_cache.shape[3]
|
|
num_input_tokens = query.shape[0]
|
|
|
|
num_q_per_kv = num_query_heads // num_kv_heads
|
|
slopes = (
|
|
alibi_slopes.view(num_query_heads, 1, 1)
|
|
if alibi_slopes is not None
|
|
else alibi_slopes
|
|
)
|
|
|
|
for i in range(num_input_tokens):
|
|
q = query[i].unsqueeze(0)
|
|
block_table = block_tables[i]
|
|
context_len = int(context_lens[i])
|
|
|
|
keys = []
|
|
values = []
|
|
for j in range(context_len):
|
|
block_number = int(block_table[j // block_size])
|
|
block_offset = j % block_size
|
|
|
|
k = key_cache[block_number, :, :, block_offset, :]
|
|
k = k.reshape(num_kv_heads, head_size)
|
|
keys.append(k)
|
|
|
|
v = value_cache[block_number, :, :, block_offset]
|
|
values.append(v)
|
|
keys = torch.stack(keys, dim=0)
|
|
values = torch.stack(values, dim=0)
|
|
if num_q_per_kv > 1:
|
|
keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1)
|
|
values = torch.repeat_interleave(values, num_q_per_kv, dim=1)
|
|
scale = 1.0 / (head_size**0.5)
|
|
if alibi_slopes is not None:
|
|
offsets = get_alibi_mask(
|
|
num_query_heads, context_len, output.device, output.dtype
|
|
)
|
|
mask = offsets * slopes
|
|
mask = mask.to(output.dtype)
|
|
else:
|
|
mask = None
|
|
out = ref_masked_attention(
|
|
q,
|
|
keys,
|
|
values,
|
|
scale,
|
|
mask,
|
|
)
|
|
out = out.view(num_query_heads, head_size)
|
|
output[i].copy_(out, non_blocking=True)
|
|
|
|
return output
|
|
|
|
|
|
def vllm_paged_attention_v4(
|
|
output: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
num_kv_heads: torch.Tensor,
|
|
scale: float,
|
|
block_tables: torch.Tensor,
|
|
context_lens_cpu: torch.Tensor,
|
|
context_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_context_len: int,
|
|
alibi_slopes: torch.Tensor = None,
|
|
use_sqrt_alibi: bool = False,
|
|
):
|
|
"""
|
|
Args:
|
|
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
num_kv_heads: int
|
|
scale: float
|
|
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
|
context_lens_cpu: (num_tokens) torch.int32
|
|
context_lens: (num_tokens) torch.int32
|
|
block_size: int
|
|
max_context_len: int
|
|
alibi_slopes: (num_heads) torch.float32
|
|
use_sqrt_alibi: bool
|
|
Returns:
|
|
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
"""
|
|
ops.infer.vllm_paged_attention_v4(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
scale,
|
|
block_tables,
|
|
context_lens_cpu,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
use_sqrt_alibi,
|
|
)
|
|
return output
|
|
|
|
|
|
def ref_vllm_rotary_embedding(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
is_neox_style: bool = True,
|
|
):
|
|
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., ::2]
|
|
x2 = x[..., 1::2]
|
|
x = torch.stack((-x2, x1), dim=-1)
|
|
return x.flatten(-2)
|
|
|
|
query_shape = query.shape
|
|
key_shape = key.shape
|
|
B = query.shape[0]
|
|
query = query.view(B, -1, head_size)
|
|
key = key.view(B, -1, head_size)
|
|
|
|
cos_sin = cos_sin_cache[positions]
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
if is_neox_style:
|
|
cos = cos.repeat(1, 1, 2).unsqueeze(-2)
|
|
sin = sin.repeat(1, 1, 2).unsqueeze(-2)
|
|
else:
|
|
cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
|
sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
|
|
|
rotate_fn = _rotate_neox if is_neox_style else _rotate_gptj
|
|
query_rot = query * cos + rotate_fn(query) * sin
|
|
key_rot = key * cos + rotate_fn(key) * sin
|
|
|
|
query = query_rot.flatten(-2).view(query_shape)
|
|
key = key_rot.flatten(-2).view(key_shape)
|
|
return query, key
|
|
|
|
|
|
def vllm_rotary_embedding(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
is_neox_style: bool = True,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
positions: (num_tokens) torch.long
|
|
query: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16
|
|
key: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16
|
|
head_size: int
|
|
cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16
|
|
is_neox_style: bool
|
|
Returns:
|
|
None. 对query, key 做in place 操作
|
|
"""
|
|
ops.infer.vllm_rotary_embedding(
|
|
positions,
|
|
query,
|
|
key,
|
|
head_size,
|
|
cos_sin_cache,
|
|
is_neox_style,
|
|
)
|
|
|
|
|
|
def ref_vllm_rotary_embedding_phi(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
long_offset: torch.Tensor,
|
|
k: int,
|
|
offsets: torch.Tensor = None,
|
|
):
|
|
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
query_shape = query.shape
|
|
key_shape = key.shape
|
|
B = query.shape[0]
|
|
query = query.view(B, -1, head_size)
|
|
key = key.view(B, -1, head_size)
|
|
|
|
if long_offset is None:
|
|
long_offset = (
|
|
torch.any(positions > k).float() * torch.full_like(positions, k)
|
|
).long()
|
|
idx = torch.add(positions, long_offset) if long_offset is not None else positions
|
|
idx = torch.add(idx, offsets) if offsets is not None else idx
|
|
cos_sin = torch.index_select(cos_sin_cache, 0, idx)
|
|
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
cos = cos.repeat(1, 2).unsqueeze(-2)
|
|
sin = sin.repeat(1, 2).unsqueeze(-2)
|
|
|
|
query = query * cos + _rotate_neox(query) * sin
|
|
key = key * cos + _rotate_neox(key) * sin
|
|
|
|
query = query.flatten(-2).view(query_shape)
|
|
key = key.flatten(-2).view(key_shape)
|
|
|
|
return query, key
|
|
|
|
|
|
def vllm_rotary_embedding_phi(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
long_offset: torch.Tensor,
|
|
k: int,
|
|
offsets: torch.Tensor = None,
|
|
):
|
|
"""
|
|
Args:
|
|
positions: (num_tokens) torch.long
|
|
query: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16
|
|
key: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16
|
|
cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16
|
|
head_size: int
|
|
cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16
|
|
long_offset: (1,) torch.bool
|
|
k: int
|
|
offsets: (num_tokens) torch.half, torch.float, torch.bfloat16
|
|
Returns:
|
|
None. 对query, key 做in place 操作
|
|
"""
|
|
ops.infer.vllm_rotary_embedding_phi(
|
|
positions,
|
|
query,
|
|
key,
|
|
head_size,
|
|
cos_sin_cache,
|
|
long_offset,
|
|
k,
|
|
offsets,
|
|
)
|
|
|
|
|
|
def ref_vllm_rotary_embedding_with_key_layer_norm(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
key_out: torch.Tensor = None,
|
|
eps: float = 1e-5,
|
|
is_neox_style: bool = True,
|
|
):
|
|
B = key.shape[0]
|
|
query_size = query.size()
|
|
query, key = ref_vllm_rotary_embedding(
|
|
positions,
|
|
query.view(B, -1),
|
|
key.view(B, -1),
|
|
head_size,
|
|
cos_sin_cache,
|
|
is_neox_style,
|
|
)
|
|
query = query.view(query_size)
|
|
key = key.view(B, -1, head_size)
|
|
|
|
norm_key = torch.nn.functional.layer_norm(
|
|
key,
|
|
[
|
|
head_size,
|
|
],
|
|
weight,
|
|
bias,
|
|
eps,
|
|
)
|
|
if key_out is not None:
|
|
key_out.copy_(norm_key)
|
|
else:
|
|
key_out = norm_key
|
|
return query, key_out
|
|
|
|
|
|
def vllm_rotary_embedding_with_key_layer_norm(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
key_out: torch.Tensor = None,
|
|
eps: float = 1e-5,
|
|
is_neox_style: bool = True,
|
|
):
|
|
"""
|
|
Args:
|
|
positions: (num_tokens) torch.int64
|
|
query: (num_tokens, num_heads * head_size) or (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
key: (num_tokens, num_kv_heads * head_size) or (num_tokens, num_kv_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
weight: (head_size) torch.half, torch.float, torch.bfloat16
|
|
bias: (head_size) torch.half, torch.float, torch.bfloat16
|
|
head_size: int
|
|
cos_sin_cache: (max_position, rot_dim) torch.half, torch.float, torch.bfloat16
|
|
key_out: (num_tokens, num_kv_heads * head_size) or (num_tokens, num_kv_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
eps: float
|
|
is_neox_style: bool
|
|
Returns:
|
|
query: (num_tokens, num_heads * head_size) or (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
key_out: (num_tokens, num_kv_heads * head_size) or (num_tokens, num_kv_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
"""
|
|
ops.infer.vllm_rotary_embedding_with_key_layer_norm(
|
|
positions,
|
|
query,
|
|
key,
|
|
weight,
|
|
bias,
|
|
head_size,
|
|
cos_sin_cache,
|
|
key_out,
|
|
eps,
|
|
is_neox_style,
|
|
)
|
|
key_out = key if key_out is None else key_out
|
|
return query, key_out
|
|
|
|
|
|
def ref_vllm_batched_rotary_embedding(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
is_neox_style,
|
|
rotary_dim: int,
|
|
offsets: torch.Tensor,
|
|
):
|
|
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., ::2]
|
|
x2 = x[..., 1::2]
|
|
x = torch.stack((-x2, x1), dim=-1)
|
|
return x.flatten(-2)
|
|
|
|
query = query.view(*query.shape[:-1], -1, head_size)
|
|
key = key.view(*key.shape[:-1], -1, head_size)
|
|
|
|
query_rot = query[..., :rotary_dim]
|
|
key_rot = key[..., :rotary_dim]
|
|
if rotary_dim < head_size:
|
|
query_pass = query[..., rotary_dim:]
|
|
key_pass = key[..., rotary_dim:]
|
|
|
|
cos_sin = cos_sin_cache[torch.add(positions, offsets)]
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
if is_neox_style:
|
|
# NOTE(woosuk): Here we assume that the positions tensor has the
|
|
# shape [batch_size, seq_len].
|
|
cos = cos.repeat(1, 1, 2).unsqueeze(-2)
|
|
sin = sin.repeat(1, 1, 2).unsqueeze(-2)
|
|
else:
|
|
cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
|
sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
|
|
|
rotate_fn = _rotate_neox if is_neox_style else _rotate_gptj
|
|
query_rot = query_rot * cos + rotate_fn(query_rot) * sin
|
|
key_rot = key_rot * cos + rotate_fn(key_rot) * sin
|
|
|
|
if rotary_dim < head_size:
|
|
query = torch.cat((query_rot, query_pass), dim=-1)
|
|
key = torch.cat((key_rot, key_pass), dim=-1)
|
|
else:
|
|
query = query_rot
|
|
key = key_rot
|
|
query = query.flatten(-2)
|
|
key = key.flatten(-2)
|
|
return query, key
|
|
|
|
|
|
def vllm_batched_rotary_embedding(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
head_size: int,
|
|
cos_sin_cache: torch.Tensor,
|
|
is_neox_style: bool,
|
|
rotary_dim: int,
|
|
offsets: torch.Tensor,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
positions: (num_tokens) torch.long
|
|
query: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16
|
|
key: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16
|
|
head_size: int
|
|
cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16
|
|
is_neox_style: bool
|
|
rotary_dim: int
|
|
offsets: (positions, head_size) torch.int64
|
|
Returns:
|
|
None. 对query, key 做in place 操作
|
|
"""
|
|
ops.infer.vllm_batched_rotary_embedding(
|
|
positions,
|
|
query,
|
|
key,
|
|
head_size,
|
|
cos_sin_cache,
|
|
is_neox_style,
|
|
rotary_dim,
|
|
offsets,
|
|
)
|
|
|
|
|
|
def ref_vllm_reshape_and_cache_v4(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
):
|
|
num_tokens, num_heads, head_size = key.shape
|
|
x = 16 // torch.tensor([], dtype=key.dtype).element_size()
|
|
block_size = key_cache.size(3)
|
|
|
|
reshaped_key = key.reshape(num_tokens, num_heads, head_size // x, x)
|
|
for i in range(num_tokens):
|
|
block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor")
|
|
block_offset = slot_mapping[i] % block_size
|
|
key_cache[block_idx, :, :, block_offset, :] = reshaped_key[i]
|
|
value_cache[block_idx, :, :, block_offset] = value[i]
|
|
|
|
|
|
def vllm_reshape_and_cache_v4(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
key_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16
|
|
if dtype=torch.half or torch.bfloat16,key_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8)
|
|
if dtype=torch.float,key_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4)
|
|
value_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16
|
|
if dtype=torch.half or torch.bfloat16,value_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8)
|
|
if dtype=torch.float,value_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4)
|
|
slot_mapping: (num_tokens) torch.long
|
|
Returns:
|
|
None, 对key_cache,value_cache进行in place 操作
|
|
"""
|
|
ops.infer.vllm_reshape_and_cache_v4(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
key.stride(0),
|
|
value.stride(0),
|
|
)
|
|
|
|
|
|
def vllm_cache_ops_reshape_and_cache(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
key: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
value: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
|
key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16
|
|
slot_mapping: (num_tokens) torch.long
|
|
Returns:
|
|
None, 对key_cache,value_cache进行in place 操作
|
|
"""
|
|
|
|
num_tokens, num_kv_heads, head_size = key.shape
|
|
num_blocks = key_cache.size(0)
|
|
key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size)
|
|
value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size)
|
|
ops.infer.vllm_cache_ops_reshape_and_cache(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
key.stride(0),
|
|
value.stride(0),
|
|
)
|
|
|
|
|
|
def ref_vllm_reshape_and_cache(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
):
|
|
num_tokens, _, _ = key.shape
|
|
block_size = key_cache.size(2)
|
|
v_dim = value.shape[-1]
|
|
|
|
for i in range(num_tokens):
|
|
block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor")
|
|
block_offset = slot_mapping[i] % block_size
|
|
key_cache[block_idx, :, block_offset, :] = key[i]
|
|
value_cache[block_idx, :, block_offset, :v_dim] = value[i]
|
|
|
|
|
|
def vllm_reshape_and_cache(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16
|
|
value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16
|
|
slot_mapping: (num_tokens) torch.long
|
|
Returns:
|
|
None, 对key_cache,value_cache进行in place 操作
|
|
"""
|
|
ops.infer.vllm_reshape_and_cache(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
key.stride(0),
|
|
value.stride(0),
|
|
)
|
|
|
|
|
|
def ref_reshape_and_cache_flash(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
kv_cache_dtype: str,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
):
|
|
num_tokens, _, _ = key.shape
|
|
block_size = key_cache.size(2)
|
|
|
|
for i in range(num_tokens):
|
|
block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor")
|
|
block_offset = slot_mapping[i] % block_size
|
|
key_cache[block_idx, :, block_offset, :] = key[i]
|
|
value_cache[block_idx, :, block_offset, :] = value[i]
|
|
|
|
|
|
def reshape_and_cache_flash(
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
kv_cache_dtype: str,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
) -> None:
|
|
|
|
"""
|
|
Args:
|
|
key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
|
key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16
|
|
value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16
|
|
slot_mapping: (num_tokens) torch.long
|
|
kv_cache_dtype: str
|
|
k_scale: float
|
|
v_scale: float
|
|
Returns:
|
|
None, 对key_cache,value_cache进行in place 操作
|
|
"""
|
|
assert k_scale == 1 and v_scale == 1
|
|
assert kv_cache_dtype == "auto"
|
|
|
|
ops.infer.vllm_reshape_and_cache(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
key.stride(0),
|
|
value.stride(0),
|
|
)
|
|
|
|
|
|
def ref_vllm_copy_blocks(
|
|
key_caches,
|
|
value_caches,
|
|
block_mapping,
|
|
):
|
|
for k, v in zip(key_caches, value_caches):
|
|
src = block_mapping[:, 0]
|
|
dst = block_mapping[:, 1]
|
|
k[dst] = k[src]
|
|
v[dst] = v[src]
|
|
|
|
|
|
def vllm_copy_blocks(
|
|
key_caches,
|
|
value_caches,
|
|
block_mapping,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
key_caches: [(num_blocks, num_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
|
value_caches: [(num_blocks, num_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
|
block_mapping: (num_tokens, 2) torch.int64
|
|
Returns:
|
|
None, 对key_caches,value_caches进行in place 操作
|
|
"""
|
|
ops.infer.vllm_copy_blocks(
|
|
key_caches,
|
|
value_caches,
|
|
block_mapping,
|
|
)
|
|
|
|
|
|
def ref_vllm_swap_blocks(src, dst, mapping):
|
|
for item in mapping:
|
|
src_idx = item[0]
|
|
dst_idx = item[1]
|
|
dst[dst_idx] = src[src_idx].to(dst.device)
|
|
|
|
|
|
def vllm_swap_blocks(src: "torch.Tensor", dst: "torch.Tensor", mapping: "torch.Tensor"):
|
|
|
|
"""
|
|
Args:
|
|
src: [(num_blocks, num_kv_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
|
dst: [(num_blocks, num_kv_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
|
mapping: (num_tokens, 2) torch.int64
|
|
Returns:
|
|
None, 对dst进行in place 操作
|
|
"""
|
|
ops.infer.vllm_swap_blocks(src, dst, mapping)
|
|
|
|
|
|
def ref_vllm_concat_and_cache_mla(
|
|
kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale
|
|
):
|
|
num_tokens, kv_lora_rank = kv_c.shape
|
|
_, _, pe_dim = k_pe.shape
|
|
_, block_size, rope_dim = kv_cache.shape
|
|
assert kv_lora_rank + pe_dim == rope_dim
|
|
|
|
for i in range(num_tokens):
|
|
block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor")
|
|
block_offset = slot_mapping[i] % block_size
|
|
kv_cache[block_idx, block_offset, :kv_lora_rank] = kv_c[i]
|
|
kv_cache[block_idx, block_offset, kv_lora_rank:] = k_pe[i, 0]
|
|
|
|
|
|
def ref_vllm_gather_cache(
|
|
src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
|
|
dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...]
|
|
block_table: torch.Tensor, # [BATCH, BLOCK_INDICES]
|
|
cu_seq_lens: torch.Tensor, # [BATCH+1]
|
|
batch_size: int,
|
|
seq_starts: torch.Tensor = None,
|
|
):
|
|
# 验证输入张量的设备一致性
|
|
assert src_cache.device == dst.device == block_table.device == cu_seq_lens.device
|
|
if seq_starts is not None:
|
|
assert seq_starts.device == src_cache.device
|
|
|
|
# 获取基本参数
|
|
block_size = src_cache.size(1)
|
|
entry_size = src_cache.flatten(2, -1).size(2)
|
|
|
|
# 处理每个批次
|
|
for bid in range(batch_size):
|
|
seq_start = cu_seq_lens[bid]
|
|
seq_end = cu_seq_lens[bid + 1]
|
|
seq_len = seq_end - seq_start
|
|
|
|
# 计算需要的块数
|
|
tot_blocks = math.ceil(seq_len / block_size)
|
|
|
|
# 获取当前批次的块表
|
|
if seq_starts is not None:
|
|
offset = seq_starts[bid] // block_size
|
|
batch_block_table = block_table[bid, offset : offset + tot_blocks]
|
|
else:
|
|
batch_block_table = block_table[bid, :tot_blocks]
|
|
|
|
# 准备目标位置
|
|
dst_seq = dst[seq_start:seq_end]
|
|
|
|
# 处理完整块
|
|
full_blocks = seq_len // block_size
|
|
if full_blocks > 0:
|
|
# 获取所有完整块的源数据 [full_blocks, block_size, entry_size]
|
|
src_blocks = src_cache[batch_block_table[:full_blocks]]
|
|
# 展平并复制到目标位置
|
|
dst_seq[: full_blocks * block_size].copy_(src_blocks.flatten(0, 1))
|
|
|
|
# 处理部分块
|
|
partial_size = seq_len % block_size
|
|
if partial_size > 0:
|
|
last_block = src_cache[batch_block_table[full_blocks], :partial_size]
|
|
dst_seq[full_blocks * block_size :].copy_(last_block)
|
|
|
|
|
|
def vllm_gather_cache(
|
|
src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
|
|
dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...]
|
|
block_table: torch.Tensor, # [BATCH, BLOCK_INDICES]
|
|
cu_seq_lens: torch.Tensor, # [BATCH+1]
|
|
batch_size: int,
|
|
seq_starts: torch.Tensor = None,
|
|
):
|
|
"""
|
|
Args:
|
|
src_cache: [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] torch.float16, torch.bfloat16 int
|
|
dst: [TOT_TOKENS, ENTRIES...] torch.float16, torch.bfloat16
|
|
block_table: [BATCH, BLOCK_INDICES] torch.int
|
|
cu_seq_lens: [BATCH+1] torch.int
|
|
batch_size: int
|
|
seq_starts: [BATCH] or None torch.int
|
|
"""
|
|
ops.infer.vllm_gather_cache(
|
|
src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts
|
|
)
|
|
|
|
|
|
def ref_vllm_gather_cache_int8(
|
|
src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
|
|
src_cache_scale: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, 2]
|
|
kv_lora_rank: int,
|
|
dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...]
|
|
block_table: torch.Tensor, # [BATCH, BLOCK_INDICES]
|
|
cu_seq_lens: torch.Tensor, # [BATCH+1]
|
|
batch_size: int,
|
|
seq_starts: torch.Tensor = None,
|
|
):
|
|
# 验证输入张量的设备一致性
|
|
assert (
|
|
src_cache.device
|
|
== src_cache_scale.device
|
|
== dst.device
|
|
== block_table.device
|
|
== cu_seq_lens.device
|
|
)
|
|
if seq_starts is not None:
|
|
assert seq_starts.device == src_cache.device
|
|
|
|
# 获取基本参数
|
|
block_size = src_cache.size(1)
|
|
|
|
# 处理每个批次
|
|
for bid in range(batch_size):
|
|
seq_start = cu_seq_lens[bid]
|
|
seq_end = cu_seq_lens[bid + 1]
|
|
seq_len = seq_end - seq_start
|
|
|
|
# 计算需要的块数
|
|
tot_blocks = math.ceil(seq_len / block_size)
|
|
|
|
# 获取当前批次的块表
|
|
if seq_starts is not None:
|
|
offset = seq_starts[bid] // block_size
|
|
batch_block_table = block_table[bid, offset : offset + tot_blocks]
|
|
else:
|
|
batch_block_table = block_table[bid, :tot_blocks]
|
|
|
|
# 准备目标位置
|
|
dst_seq = dst[seq_start:seq_end]
|
|
|
|
# 处理完整块
|
|
full_blocks = seq_len // block_size
|
|
if full_blocks > 0:
|
|
# 获取所有完整块的源数据 [full_blocks, block_size, entry_size]
|
|
src_cache_blocks = src_cache[batch_block_table[:full_blocks]]
|
|
src_scale_blocks = src_cache_scale[batch_block_table[:full_blocks]]
|
|
src_k_cache_blocks = src_cache_blocks[
|
|
..., :kv_lora_rank
|
|
] * src_scale_blocks[..., 0].unsqueeze(-1)
|
|
src_k_pe_blocks = src_cache_blocks[..., kv_lora_rank:] * src_scale_blocks[
|
|
..., 1
|
|
].unsqueeze(-1)
|
|
src_blocks = torch.cat((src_k_cache_blocks, src_k_pe_blocks), dim=-1).to(
|
|
dst.dtype
|
|
)
|
|
# 展平并复制到目标位置
|
|
dst_seq[: full_blocks * block_size].copy_(src_blocks.flatten(0, 1))
|
|
|
|
# 处理部分块
|
|
partial_size = seq_len % block_size
|
|
if partial_size > 0:
|
|
last_block = src_cache[batch_block_table[full_blocks], :partial_size]
|
|
last_src_scale_blocks = src_cache_scale[
|
|
batch_block_table[full_blocks], :partial_size
|
|
]
|
|
last_src_k_cache_blocks = last_block[
|
|
..., :kv_lora_rank
|
|
] * last_src_scale_blocks[..., 0].unsqueeze(-1)
|
|
last_src_k_pe_blocks = last_block[
|
|
..., kv_lora_rank:
|
|
] * last_src_scale_blocks[..., 1].unsqueeze(-1)
|
|
last_block = torch.cat(
|
|
(last_src_k_cache_blocks, last_src_k_pe_blocks), dim=-1
|
|
).to(dst.dtype)
|
|
dst_seq[full_blocks * block_size :].copy_(last_block)
|
|
|
|
|
|
def vllm_gather_cache_int8(
|
|
src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
|
|
src_cache_scale: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, 2]
|
|
kv_lora_rank: int,
|
|
dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...]
|
|
block_table: torch.Tensor, # [BATCH, BLOCK_INDICES]
|
|
cu_seq_lens: torch.Tensor, # [BATCH+1]
|
|
batch_size: int,
|
|
seq_starts: torch.Tensor = None,
|
|
):
|
|
"""
|
|
Args:
|
|
src_cache: [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] torch.int8
|
|
src_cache_scale: [NUM_BLOCKS, BLOCK_SIZE, 2] torch.float32
|
|
kv_lora_rank: int
|
|
dst: [TOT_TOKENS, ENTRIES...] torch.float16, torch.bfloat16
|
|
block_table: [BATCH, BLOCK_INDICES] torch.int
|
|
cu_seq_lens: [BATCH+1] torch.int
|
|
batch_size: int
|
|
seq_starts: [BATCH] or None torch.int
|
|
"""
|
|
ops.infer.vllm_gather_cache_int8(
|
|
src_cache,
|
|
src_cache_scale,
|
|
kv_lora_rank,
|
|
dst,
|
|
block_table,
|
|
cu_seq_lens,
|
|
batch_size,
|
|
seq_starts,
|
|
)
|
|
|
|
|
|
def ref_vllm_concat_and_cache_mla_int8(
|
|
kv_c_int8: torch.Tensor,
|
|
kv_c_scale: torch.Tensor,
|
|
k_pe_int8: torch.Tensor,
|
|
k_pe_scale: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
kv_cache_scale: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
kv_cache_dtype: str,
|
|
scale: torch.Tensor,
|
|
) -> None:
|
|
|
|
num_tokens, kv_lora_rank = kv_c_int8.shape
|
|
_, block_size, _ = kv_cache_scale.shape
|
|
|
|
for i in range(num_tokens):
|
|
block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor")
|
|
block_offset = slot_mapping[i] % block_size
|
|
kv_cache[block_idx, block_offset, :kv_lora_rank] = kv_c_int8[i]
|
|
kv_cache[block_idx, block_offset, kv_lora_rank:] = k_pe_int8[i][0]
|
|
kv_cache_scale[block_idx, block_offset, 0] = kv_c_scale[i]
|
|
kv_cache_scale[block_idx, block_offset, 1] = k_pe_scale[i][0]
|
|
|
|
|
|
def vllm_concat_and_cache_mla(
|
|
kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale
|
|
):
|
|
ops.infer.vllm_concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping)
|
|
|
|
|
|
def vllm_concat_and_cache_mla_int8(
|
|
kv_c_int8,
|
|
kv_c_scale,
|
|
k_pe_int8,
|
|
k_pe_scale,
|
|
kv_cache,
|
|
kv_cache_scale,
|
|
slot_mapping,
|
|
kv_cache_dtype,
|
|
scale,
|
|
):
|
|
"""
|
|
Args:
|
|
kv_c_int8: [num_tokens, kv_lora_rank] torch.int8
|
|
kv_c_scale: [num_tokens] torch.float32
|
|
k_pe_int8: [num_tokens, n, pe_dim] torch.int8
|
|
k_pe_scale: [num_tokens, n] torch.float32
|
|
kv_cache: [num_blocks, block_size, (kv_lora_rank + pe_dim)] torch.int8
|
|
kv_cache_scale: [num_blocks, block_size, 2] torch.float32
|
|
slot_mapping: [num_tokens] torch.long
|
|
"""
|
|
|
|
ops.infer.vllm_concat_and_cache_mla_int8(
|
|
kv_c_int8,
|
|
kv_c_scale,
|
|
k_pe_int8,
|
|
k_pe_scale,
|
|
kv_cache,
|
|
kv_cache_scale,
|
|
slot_mapping,
|
|
)
|
|
|
|
|
|
class vllm_llama_mlp(CF.VllmLlamaMlp):
|
|
def __init__(
|
|
self,
|
|
gate_up_proj_weight: "torch.Tensor",
|
|
down_proj_weight: "torch.Tensor",
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
tp: int,
|
|
) -> None:
|
|
gate_up_proj_weight = gate_up_proj_weight
|
|
down_proj_weight = down_proj_weight
|
|
hidden_size = hidden_size
|
|
intermediate_size = intermediate_size
|
|
tp = tp
|
|
super().__init__(
|
|
gate_up_proj_weight,
|
|
down_proj_weight,
|
|
hidden_size,
|
|
intermediate_size,
|
|
tp,
|
|
)
|
|
|
|
def __call__(self, x: "torch.Tensor", group=None):
|
|
x1 = x
|
|
if group is None:
|
|
super().forward(x1, x1)
|
|
else:
|
|
from ixformer.distributed._distributed import _check_group
|
|
|
|
group = _check_group(group)
|
|
super().forward(x1, x1, group)
|
|
return x
|
|
|
|
|
|
def gptq_gemm(
|
|
input: torch.Tensor,
|
|
qweight: torch.Tensor,
|
|
qzeros: torch.Tensor,
|
|
scales: torch.Tensor,
|
|
g_idx: torch.Tensor,
|
|
use_exllama: bool,
|
|
weight_bits: int,
|
|
) -> torch.Tensor:
|
|
"""
|
|
use_exllama
|
|
- True uesExllama
|
|
- False GeneralGptq
|
|
g_idx.is_empty()
|
|
- True don't use g_idx
|
|
- False use g_idx
|
|
1. use_exllama == False && use g_idx
|
|
[General gptq] desc_act == True && parallel in k dimension && group_size != -1
|
|
2. use_exllama == True && use g_idx (g_idx has been argsort)
|
|
[Exllama with g_idx] desc_act == True && parallel in n dimension && group_size != -1
|
|
3. use_exllama == True && don't use g_idx
|
|
[Exllama] desc_act == False || desc_act == True && group_size == -1
|
|
|
|
Args:
|
|
input: (m, k) torch.float16, torch.bfloat16
|
|
qweight: (k // (32 / bits), n) torch.int32
|
|
qzeros: (k / group_size, n / (32 / bits)) torch.int32
|
|
scales: (k // group_size, n) torch.float16, torch.bfloat16
|
|
g_idx: (k) torch.int32
|
|
use_exllama: bool
|
|
wheather use exllama
|
|
weight_bits: int
|
|
quant bits of weight
|
|
Returns:
|
|
output: (m, n) torch.float16, torch.bfloat16
|
|
"""
|
|
bs = input.shape[0]
|
|
group_size = input.shape[1] // scales.shape[0]
|
|
|
|
# condition : without gidx or group_size == -1
|
|
ixinfer_gemm_supported = (
|
|
weight_bits == 4
|
|
and (g_idx is None or g_idx.numel() == 0)
|
|
and (scales.shape[0] == 1 or group_size in [32, 128])
|
|
)
|
|
if use_exllama:
|
|
if bs <= 8 or ixinfer_gemm_supported:
|
|
output = ops.infer.quantized_linear(
|
|
input,
|
|
qweight,
|
|
scales,
|
|
"gptq-ex",
|
|
weight_bits,
|
|
qzeros,
|
|
None,
|
|
group_size,
|
|
g_idx,
|
|
"unknown",
|
|
)
|
|
else:
|
|
# GPTQ GEMM TODO
|
|
o_dtype_str = "fp16" if input.dtype == torch.half else "bf16"
|
|
deq_w = ops.infer.quantized_weight_dequant(
|
|
qweight,
|
|
scales,
|
|
"gptq-ex",
|
|
o_dtype_str,
|
|
weight_bits,
|
|
qzeros,
|
|
group_size,
|
|
g_idx,
|
|
)
|
|
|
|
output = linear(input, deq_w.transpose(0, 1).contiguous())
|
|
else:
|
|
if bs <= 8:
|
|
output = ops.infer.quantized_linear(
|
|
input,
|
|
qweight,
|
|
scales,
|
|
"gptq",
|
|
weight_bits,
|
|
qzeros,
|
|
None,
|
|
group_size,
|
|
g_idx,
|
|
"unknown",
|
|
)
|
|
else:
|
|
# GPTQ GEMM TODO
|
|
o_dtype_str = "fp16" if input.dtype == torch.half else "bf16"
|
|
deq_w = ops.infer.quantized_weight_dequant(
|
|
qweight,
|
|
scales,
|
|
"gptq",
|
|
o_dtype_str,
|
|
weight_bits,
|
|
qzeros,
|
|
group_size,
|
|
g_idx,
|
|
)
|
|
output = linear(input, deq_w.transpose(0, 1).contiguous())
|
|
return output
|
|
|
|
|
|
def vllm_gptq_shuffle(qweights, g_idx, weight_bits):
|
|
ops.infer.vllm_gptq_shuffle(qweights, g_idx, weight_bits)
|
|
|
|
|
|
def vllm_moe_topk_softmax(
|
|
topk_weights: "torch.Tensor",
|
|
topk_ids: "torch.Tensor",
|
|
token_expert_indicies: "torch.Tensor",
|
|
gating_output: "torch.Tensor",
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
topk_weights: (num_tokens,topk) torch.float
|
|
topk_ids: (num_tokens,topk) torch.int
|
|
token_expert_indicies: (num_tokens,topk) torch.int
|
|
gating_output: (num_tokens,num_experts) torch.float
|
|
Returns:
|
|
None, 对topk_weights,topk_ids进行in place 操作
|
|
"""
|
|
assert isinstance(topk_weights, torch.Tensor)
|
|
assert gating_output.dtype == torch.float32
|
|
ops.infer.moe_topk_softmax(
|
|
topk_weights, topk_ids, token_expert_indicies, gating_output, False
|
|
)
|
|
|
|
|
|
def vllm_moe_align_block_size(
|
|
topk_ids: "torch.Tensor",
|
|
num_experts: int,
|
|
block_size: int,
|
|
sorted_ids: "torch.Tensor",
|
|
expert_ids: "torch.Tensor",
|
|
num_tokens_post_pad: "torch.Tensor",
|
|
):
|
|
"""
|
|
Args:
|
|
topk_ids: (num_tokens,topk) torch.int
|
|
num_experts: int
|
|
block_size: int
|
|
sorted_ids: (topk_ids.numel() + num_experts * (block_size - 1)) torch.int
|
|
expert_ids: (topk_ids.numel() + num_experts) torch.int
|
|
num_tokens_post_pad: (1) torch.int
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
ops.infer.moe_align_block_size(
|
|
topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad
|
|
)
|
|
|
|
|
|
def ref_vllm_invoke_fused_moe_kernel(
|
|
A: "torch.Tensor",
|
|
B: "torch.Tensor",
|
|
C: "torch.Tensor",
|
|
topk_weight: "torch.Tensor",
|
|
topk_ids: "torch.Tensor",
|
|
sorted_token_ids: "torch.Tensor",
|
|
expert_ids: "torch.Tensor",
|
|
num_tokens_post_padded: "torch.Tensor",
|
|
mul_routed_weight: bool,
|
|
top_k: int,
|
|
block_size_m: int,
|
|
persistent: bool = False,
|
|
w_scale: torch.Tensor = None,
|
|
a_scale: torch.Tensor = None,
|
|
):
|
|
|
|
expert_num, N, K = B.shape
|
|
M, topk = C.shape[:2]
|
|
|
|
clone_A = A.clone()
|
|
clone_B = B.clone()
|
|
if clone_A.shape[0] == M:
|
|
clone_A = clone_A.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
|
|
topk_ids = topk_ids.view(-1)
|
|
|
|
if A.dtype == torch.int8:
|
|
use_scale = True
|
|
clone_A = clone_A.to(torch.float32)
|
|
clone_B = clone_B.to(torch.float32)
|
|
tmp = torch.zeros(M * topk, N, dtype=torch.float32, device=C.device)
|
|
else:
|
|
use_scale = False
|
|
tmp = torch.zeros(M * topk, N, dtype=C.dtype, device=C.device)
|
|
|
|
for i in range(expert_num): # expert_num
|
|
mask = topk_ids == i
|
|
if mask.sum():
|
|
tmp[mask] = clone_A[mask] @ clone_B[i].transpose(0, 1)
|
|
if use_scale:
|
|
tmp[mask] = tmp[mask] * w_scale[i].view(1, N)
|
|
|
|
if mul_routed_weight:
|
|
tmp = tmp * topk_weight.view(-1, 1)
|
|
|
|
if use_scale:
|
|
tmp = tmp.view(M, topk, N)
|
|
tmp = tmp * a_scale.view(M, -1, 1)
|
|
|
|
C[:] = tmp.to(C.dtype).view(M, topk, N)
|
|
return C
|
|
|
|
|
|
def vllm_invoke_fused_moe_kernel(
|
|
A: "torch.Tensor",
|
|
B: "torch.Tensor",
|
|
C: "torch.Tensor",
|
|
topk_weight: "torch.Tensor",
|
|
topk_ids: "torch.Tensor",
|
|
sorted_token_ids: "torch.Tensor",
|
|
expert_ids: "torch.Tensor",
|
|
num_tokens_post_padded: "torch.Tensor",
|
|
mul_routed_weight: bool,
|
|
top_k: int,
|
|
block_size_m: int,
|
|
persistent: bool = False,
|
|
w_scale: torch.Tensor = None,
|
|
a_scale: torch.Tensor = None,
|
|
):
|
|
|
|
"""
|
|
Args:
|
|
A: (bs*seq, K) / (bs*seq*top_k, K) torch.float16, torch.bfloat16
|
|
B: (num_experts, N, K) torch.float16, torch.bfloat16
|
|
C: (bs*seq, top_k, N) torch.half,torch.bfloat16
|
|
topk_weight: (bs*seq, topk) torch.float32
|
|
topk_ids: (bs*seq, topk) torch.int32
|
|
sorted_token_ids: (topk_ids.numel() + num_experts * (block_size - 1)) torch.int32
|
|
expert_ids: (topk_ids.numel() + num_experts) torch.int32
|
|
num_tokens_post_pad:(1) torch.int32
|
|
mul_routed_weight: bool
|
|
top_k: int
|
|
block_size_m: int
|
|
Returns:
|
|
C: (bs*seq, top_k, N) torch.half,torch.bfloat16
|
|
"""
|
|
ops.infer.invoke_fused_moe_kernel(
|
|
A,
|
|
B,
|
|
C,
|
|
topk_weight,
|
|
topk_ids,
|
|
sorted_token_ids,
|
|
expert_ids,
|
|
num_tokens_post_padded,
|
|
mul_routed_weight,
|
|
top_k,
|
|
block_size_m,
|
|
persistent,
|
|
w_scale,
|
|
a_scale.view(-1) * topk_weight.view(-1)
|
|
if a_scale is not None and mul_routed_weight
|
|
else a_scale,
|
|
)
|
|
|
|
|
|
def advance_step_flashattn(
|
|
num_seqs: int,
|
|
num_queries: int,
|
|
block_size: int,
|
|
input_tokens: "torch.Tensor",
|
|
sampled_token_ids: "torch.Tensor",
|
|
input_positions: "torch.Tensor",
|
|
seq_lens: "torch.Tensor",
|
|
slot_mapping: "torch.Tensor",
|
|
block_tables: "torch.Tensor",
|
|
):
|
|
ops.infer.vllm_advance_step_flashattn(
|
|
num_seqs,
|
|
num_queries,
|
|
block_size,
|
|
input_tokens,
|
|
sampled_token_ids,
|
|
input_positions,
|
|
seq_lens,
|
|
slot_mapping,
|
|
block_tables,
|
|
)
|
|
|
|
|
|
if config.IXFORMER_PAGED_ATTENTION_ALGO == "ixformer":
|
|
print("set IXFORMER_PAGED_ATTENTION_ALGO: ixformer")
|
|
vllm_paged_attention = vllm_paged_attention_ixformer
|
|
else:
|
|
vllm_paged_attention = vllm_paged_attention_ixinfer
|