793 lines
24 KiB
Python
793 lines
24 KiB
Python
#
|
|
# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
|
|
#
|
|
# This file is a part of the vllm-kunlun project.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""kunlun custom op entry"""
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
try:
|
|
import cocopod # noqa
|
|
import kunlun_ops
|
|
|
|
logger.info("Load custom ops library success!")
|
|
except ImportError as e:
|
|
logger.warning("Import error msg: %s", e.msg)
|
|
|
|
|
|
_per_token_smooth_quant = True
|
|
|
|
|
|
def is_per_token_smooth_quant():
|
|
"""is per token smooth quant"""
|
|
return _per_token_smooth_quant
|
|
|
|
|
|
class KunlunOps:
|
|
"""KunlunOps"""
|
|
|
|
# Attention ops
|
|
@staticmethod
|
|
def paged_attention_v1(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
context_lens_cpu,
|
|
is_context,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
kv_cache_dtype,
|
|
k_scale,
|
|
v_scale,
|
|
tp_rank,
|
|
blocksparse_local_blocks,
|
|
blocksparse_vert_stride,
|
|
blocksparse_block_size,
|
|
blocksparse_head_sliding_step,
|
|
alibi_sqrt=False,
|
|
):
|
|
"""PagedAttentionV1"""
|
|
# block_size = value_cache.shape[2]
|
|
kunlun_ops.paged_attention(
|
|
x=query,
|
|
k_cache=key_cache,
|
|
v_cache=value_cache,
|
|
block_tables=block_tables,
|
|
context_lens_cpu=context_lens_cpu,
|
|
context_lens_xpu=context_lens,
|
|
is_context=is_context,
|
|
is_causal=True,
|
|
out=output,
|
|
vo_head_dim=128,
|
|
)
|
|
|
|
@staticmethod
|
|
def paged_attention_v2(
|
|
output,
|
|
exp_sums,
|
|
max_logits,
|
|
tmp_output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
context_lens_cpu,
|
|
is_context,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
kv_cache_dtype,
|
|
k_scale,
|
|
v_scale,
|
|
tp_rank,
|
|
blocksparse_local_blocks,
|
|
blocksparse_vert_stride,
|
|
blocksparse_block_size,
|
|
blocksparse_head_sliding_step,
|
|
alibi_sqrt=False,
|
|
):
|
|
"""PagedAttentionV2"""
|
|
# block_size = value_cache.shape[2]
|
|
kunlun_ops.paged_attention(
|
|
x=query,
|
|
k_cache=key_cache,
|
|
v_cache=value_cache,
|
|
block_tables=block_tables,
|
|
context_lens_cpu=context_lens_cpu,
|
|
context_lens_xpu=context_lens,
|
|
is_context=is_context,
|
|
is_causal=True,
|
|
out=output,
|
|
vo_head_dim=128,
|
|
)
|
|
|
|
# Activation ops
|
|
@staticmethod
|
|
def silu_and_mul(out: torch.Tensor, x: torch.Tensor):
|
|
"""silu and mul"""
|
|
kunlun_ops.silu_and_mul(
|
|
x,
|
|
axis=-1,
|
|
turn=True,
|
|
out=out,
|
|
)
|
|
|
|
# Activation ops
|
|
@staticmethod
|
|
def quick_gelu(out: torch.Tensor, x: torch.Tensor):
|
|
"""quick gelu"""
|
|
kunlun_ops.quick_gelu(
|
|
x,
|
|
out=out,
|
|
)
|
|
|
|
# Layernorm
|
|
@staticmethod
|
|
def rms_norm(
|
|
out,
|
|
x,
|
|
weight,
|
|
epsilon,
|
|
):
|
|
"""rms_norm"""
|
|
kunlun_ops.rmsnorm(x, weight.to(torch.float32), epsilon, out=out)
|
|
|
|
@staticmethod
|
|
def fused_add_rms_norm(
|
|
x,
|
|
residual,
|
|
weight,
|
|
epsilon,
|
|
):
|
|
"""fused_add_rms_norm"""
|
|
output = torch.empty_like(x)
|
|
kunlun_ops.add_rmsnorm(
|
|
x, residual, weight.to(torch.float32), epsilon, out=output
|
|
)
|
|
fused_input = x + residual
|
|
residual.copy_(fused_input, non_blocking=True)
|
|
x.copy_(output)
|
|
|
|
# Rotary embedding
|
|
@staticmethod
|
|
def rotary_embedding(
|
|
positions, query, key, head_size, cos_sin_cache, is_neox_style
|
|
):
|
|
"""
|
|
refactor RotaryEmbedding forward function
|
|
"""
|
|
query_x = query.contiguous()
|
|
key_x = key.contiguous()
|
|
|
|
torch.ops._C.rotary_embedding(
|
|
positions, query_x, key_x, head_size, cos_sin_cache, is_neox_style
|
|
)
|
|
|
|
return query_x, key_x
|
|
|
|
# Rotary embedding
|
|
@staticmethod
|
|
def mrotary_embedding(
|
|
positions, mrope_section, query, key, head_size, cos_sin_cache, is_neox_style
|
|
):
|
|
"""
|
|
refactor RotaryEmbedding forward function
|
|
"""
|
|
query_x = query.contiguous()
|
|
key_x = key.contiguous()
|
|
assert is_neox_style
|
|
kunlun_ops.mrotary_embedding_neox(
|
|
positions, query_x, key_x, head_size, cos_sin_cache, mrope_section
|
|
)
|
|
|
|
query.data = query_x
|
|
key.data = key_x
|
|
return query, key
|
|
|
|
@staticmethod
|
|
def swap_blocks(src, dst, block_mapping):
|
|
"""swap_blocks"""
|
|
kunlun_ops.swap_blocks(src, dst, block_mapping)
|
|
|
|
@staticmethod
|
|
def copy_blocks(key_caches, value_caches, block_mapping):
|
|
"""copy_blocks"""
|
|
for i in range(len(key_caches)):
|
|
key_caches[i] = key_caches[i].contiguous()
|
|
value_caches[i] = value_caches[i].contiguous()
|
|
kunlun_ops.copy_blocks(
|
|
key_caches,
|
|
value_caches,
|
|
block_mapping,
|
|
)
|
|
|
|
@staticmethod
|
|
def reshape_and_cache(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
kv_cache_dtype,
|
|
):
|
|
"""reshape_and_cache"""
|
|
# slot_mapping_cast = slot_mapping.to(torch.int32)
|
|
kunlun_ops.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)
|
|
|
|
@staticmethod
|
|
def multi_query_kv_attention(
|
|
usual_seq_lod_xpu: torch.Tensor,
|
|
usual_seq_lod_cpu: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
**kargs,
|
|
) -> torch.Tensor:
|
|
"""
|
|
query: shape = [num_prompt_tokens, num_heads, head_size]
|
|
"""
|
|
if query.dim() == 3:
|
|
query = query.unsqueeze(0)
|
|
key = key.unsqueeze(0)
|
|
value = value.unsqueeze(0)
|
|
output = torch.empty_like(query)
|
|
|
|
B, T, Qh, Hd = query.shape
|
|
KVh = key.size(2)
|
|
if KVh != Qh:
|
|
repeat = Qh // KVh
|
|
key = key.repeat_interleave(repeat, dim=2) # [B, T, Qh, Hd]
|
|
value = value.repeat_interleave(repeat, dim=2)
|
|
kunlun_ops.attention(
|
|
q=query,
|
|
k_cache=key,
|
|
v_cache=value,
|
|
out=output,
|
|
is_causal=True,
|
|
is_prefill=True,
|
|
context_seq_lod_cpu=usual_seq_lod_cpu,
|
|
context_seq_lod_xpu=usual_seq_lod_xpu,
|
|
)
|
|
return output
|
|
|
|
@staticmethod
|
|
def quant_fusedresidual_rmsnorm_op(
|
|
x, residual, weight, bias, scale_to_int, eps, dyn_scale: bool, type: int = 1
|
|
):
|
|
"""Quantized fused residual layer normalization"""
|
|
out = torch.empty_like(x, dtype=torch.int8)
|
|
|
|
if is_per_token_smooth_quant():
|
|
out_scale = torch.empty(
|
|
x.shape[:-1], device=x.device, dtype=torch.float
|
|
).unsqueeze(-1)
|
|
else:
|
|
out_scale = torch.empty(12, device=x.device, dtype=torch.float)
|
|
|
|
kunlun_ops.quant_fusedresidual_rmsnorm(
|
|
x,
|
|
residual,
|
|
weight,
|
|
bias,
|
|
eps,
|
|
out=out,
|
|
out_scale=out_scale,
|
|
residual_tensor=residual,
|
|
)
|
|
|
|
if residual is None:
|
|
return out, out_scale
|
|
return out, out_scale, residual
|
|
|
|
@staticmethod
|
|
def quant_rmsnorm_op(
|
|
x, weight, bias, scale_to_int, eps, dyn_scale: bool, type: int = 1
|
|
):
|
|
"""Quantized RMSNorm"""
|
|
|
|
out = torch.empty_like(x, dtype=torch.int8)
|
|
if is_per_token_smooth_quant():
|
|
out_scale = torch.empty(
|
|
x.shape[:-1], device=x.device, dtype=torch.float
|
|
).unsqueeze(-1)
|
|
else:
|
|
out_scale = torch.empty(12, device=x.device, dtype=torch.float)
|
|
|
|
kunlun_ops.quant_rmsnorm(x, weight, bias, eps, out=out, out_scale=out_scale)
|
|
return out, out_scale
|
|
|
|
@staticmethod
|
|
def smooth_quant_matmul_column_row_kernels(
|
|
input_tensor,
|
|
weight,
|
|
smoother,
|
|
input_scale,
|
|
weight_scale,
|
|
perTokenScaling,
|
|
perChannelScaling,
|
|
otype,
|
|
):
|
|
"""smooth_quant_matmul_column_row_kernels"""
|
|
input_shape = input_tensor.shape
|
|
weight_shape = weight.shape
|
|
if input_tensor.dim() == 3:
|
|
input_tensor = input_tensor.reshape(-1, input_shape[-1])
|
|
out = torch.empty(
|
|
(input_shape[0] * input_shape[1], weight_shape[0]),
|
|
dtype=torch.float16,
|
|
device=weight.device,
|
|
)
|
|
output_bs_shape = [input_shape[0], input_shape[1]]
|
|
elif input_tensor.dim() == 2:
|
|
out = torch.empty(
|
|
(input_shape[0], weight_shape[0]),
|
|
dtype=torch.float16,
|
|
device=weight.device,
|
|
)
|
|
output_bs_shape = [-1]
|
|
kunlun_ops.smooth_quant_matmul_column_row_kernels(
|
|
input_tensor,
|
|
weight,
|
|
smoother,
|
|
input_scale,
|
|
weight_scale,
|
|
perTokenScaling,
|
|
perChannelScaling,
|
|
out=out,
|
|
)
|
|
|
|
out = out.view(*output_bs_shape, weight_shape[0])
|
|
|
|
return out
|
|
|
|
def _dbg(x):
|
|
if torch.is_tensor(x):
|
|
return (type(x), x.device, x.dtype, x.shape, x.is_contiguous())
|
|
return (type(x), x)
|
|
|
|
@staticmethod
|
|
def fused_moe(
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
ep_rank: int,
|
|
moe_top_k: int,
|
|
renormalize: bool,
|
|
inplace: bool = False,
|
|
use_grouped_topk: bool = False,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
w1_bias: Optional[torch.Tensor] = None,
|
|
w2_bias: Optional[torch.Tensor] = None,
|
|
scoring_func: str = "softmax",
|
|
e_score_correction_bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
"""fused_moe"""
|
|
global_num_experts, up_gate_size, _ = w1.shape
|
|
M, N = hidden_states.shape
|
|
hidden_dim = w2.shape[1]
|
|
normed_score = torch.empty(
|
|
M, moe_top_k, dtype=torch.float32, device=hidden_states.device
|
|
)
|
|
topk_ids = torch.empty(
|
|
M, moe_top_k, dtype=torch.int32, device=hidden_states.device
|
|
)
|
|
num_blocks = 12
|
|
block_statistic = torch.zeros(
|
|
num_blocks,
|
|
global_num_experts,
|
|
dtype=torch.int32,
|
|
device=hidden_states.device,
|
|
)
|
|
router_logits = router_logits.to(torch.float)
|
|
if scoring_func == "softmax":
|
|
torch.ops._C.moe_softmax_topk_norm(
|
|
x=router_logits,
|
|
normed_score=normed_score,
|
|
topk_index=topk_ids,
|
|
block_statistic=None,
|
|
stable=False,
|
|
)
|
|
elif scoring_func == "sigmoid":
|
|
torch.ops._C.moe_sigmoid_group_topk_norm(
|
|
x=router_logits,
|
|
topk_index=topk_ids,
|
|
norm_score=normed_score,
|
|
block_static=block_statistic,
|
|
bias=e_score_correction_bias,
|
|
scale=1.0,
|
|
n_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
)
|
|
|
|
if w1_bias is not None or w2_bias is not None:
|
|
# Rignt now this branch is for gpt oss
|
|
# TODO (@xyDong23): faster here using moe_fc kernel
|
|
normed_score = normed_score.to(hidden_states.dtype)
|
|
out = torch.zeros(
|
|
M * moe_top_k, N, dtype=hidden_states.dtype, device=hidden_states.device
|
|
)
|
|
repeat_x = hidden_states.repeat_interleave(moe_top_k, dim=0)
|
|
topk_ids_flat = topk_ids.flatten()
|
|
for i in range(global_num_experts):
|
|
experts_id = ep_rank * global_num_experts + i
|
|
selected_token = topk_ids_flat == experts_id
|
|
if selected_token.sum():
|
|
cur_token = repeat_x[selected_token]
|
|
up_gate = torch.empty(
|
|
selected_token.sum(),
|
|
up_gate_size // 2,
|
|
dtype=cur_token.dtype,
|
|
device=cur_token.device,
|
|
)
|
|
groupgemm1 = cur_token @ w1[i].T
|
|
# Add w13 bias
|
|
if w1_bias is not None:
|
|
groupgemm1 = groupgemm1 + w1_bias[i]
|
|
up_gate = torch.ops._C.swigluoai_and_mul(groupgemm1)
|
|
groupgemm2 = up_gate @ w2[i].T
|
|
# Add w2 bias
|
|
if w2_bias is not None:
|
|
groupgemm2 = groupgemm2 + w2_bias[i]
|
|
out[selected_token] = groupgemm2
|
|
ouput = (
|
|
(out.view(M, moe_top_k, N) * normed_score.unsqueeze(2))
|
|
.sum(dim=1)
|
|
.to(hidden_states.dtype)
|
|
)
|
|
return ouput
|
|
else:
|
|
# from vllm.forward_context import get_forward_context
|
|
# forward_context = get_forward_context()
|
|
# attn_metadata: AttentionMetadata = forward_context.attn_metadata
|
|
# prefix = "model.layers.0.linear_attn"
|
|
# if attn_metadata is not None:
|
|
# attn_metadata = attn_metadata[prefix]
|
|
|
|
# if attn_metadata is None or attn_metadata.num_prefills > 0 or :
|
|
if M * moe_top_k < 400:
|
|
sorted_tokens_idx, sorted_tokens_num_lod, moe_expand = (
|
|
torch.ops.xspeedgate_ops.moe_pre_small(
|
|
topk_ids, global_num_experts, False, False, hidden_states
|
|
)
|
|
)
|
|
experts_num_lod = torch.ops.xspeedgate_ops.moe_active_expert_balance(
|
|
topk_ids, global_num_experts, False
|
|
)
|
|
out = torch.ops.xspeedgate_ops.fused_moe(
|
|
hidden_states,
|
|
w1,
|
|
w2,
|
|
normed_score.to(hidden_states.dtype),
|
|
sorted_tokens_num_lod,
|
|
sorted_tokens_idx,
|
|
experts_num_lod,
|
|
)
|
|
return out.sum(1)
|
|
|
|
if M * moe_top_k > 768:
|
|
moe_expand = torch.empty(
|
|
(M * moe_top_k, N),
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
) # [M*top_k, N], float
|
|
expert_m = torch.zeros(
|
|
global_num_experts, dtype=torch.int32, device=hidden_states.device
|
|
) # [E]
|
|
sorted_tokens_num_lod = torch.zeros(
|
|
global_num_experts + 1,
|
|
dtype=torch.int32,
|
|
device=hidden_states.device,
|
|
) # [E+1]
|
|
sorted_tokens_idx = torch.zeros(
|
|
M * moe_top_k, dtype=torch.int32, device=hidden_states.device
|
|
)
|
|
|
|
torch.ops._C.gen_block_statistic(topk_ids, block_statistic)
|
|
|
|
torch.ops._C.moe_pre_sorted(
|
|
x=hidden_states,
|
|
topk_index=topk_ids,
|
|
block_statistic=block_statistic,
|
|
moe_expand=moe_expand,
|
|
moe_index=sorted_tokens_idx,
|
|
expert_m=expert_m,
|
|
sorted_tokens_num_lod=sorted_tokens_num_lod,
|
|
)
|
|
else:
|
|
sorted_tokens_idx, sorted_tokens_num_lod, moe_expand = (
|
|
torch.ops.xspeedgate_ops.moe_pre_small(
|
|
topk_ids,
|
|
global_num_experts,
|
|
index_have_neg=False,
|
|
sort_mode=True,
|
|
x=hidden_states,
|
|
)
|
|
)
|
|
|
|
y = torch.empty(
|
|
M,
|
|
moe_top_k,
|
|
w1.shape[1],
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
|
|
moe_expand = moe_expand.view(M * moe_top_k, hidden_dim)
|
|
|
|
if M < 1024:
|
|
torch.ops._C.moe_fc(
|
|
x=moe_expand,
|
|
weight=w1,
|
|
sorted_tokens_num_lod=sorted_tokens_num_lod,
|
|
sorted_tokens_idx=sorted_tokens_idx,
|
|
moe_topk=moe_top_k,
|
|
y=y,
|
|
)
|
|
|
|
d = y.shape[-1] // 2
|
|
output_shape = y.shape[:-1] + (d,)
|
|
out1 = torch.empty(output_shape, dtype=y.dtype, device=y.device)
|
|
torch.ops._C.silu_and_mul(out1, y)
|
|
|
|
out1 = out1.reshape(-1, out1.shape[-1])
|
|
else:
|
|
torch.ops._C.moe_fc(
|
|
x=moe_expand,
|
|
weight=w1,
|
|
sorted_tokens_num_lod=sorted_tokens_num_lod,
|
|
sorted_tokens_idx=sorted_tokens_idx,
|
|
moe_topk=moe_top_k,
|
|
y=y,
|
|
act="SWISH_GLU",
|
|
)
|
|
|
|
y = y[..., : y.shape[-1] // 2]
|
|
out1 = y.reshape(-1, y.shape[-1])
|
|
|
|
out = torch.empty(
|
|
M,
|
|
moe_top_k,
|
|
w2.shape[1],
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
|
|
torch.ops._C.moe_fc(
|
|
x=out1,
|
|
weight=w2,
|
|
sorted_tokens_num_lod=sorted_tokens_num_lod,
|
|
sorted_tokens_idx=sorted_tokens_idx,
|
|
moe_topk=moe_top_k,
|
|
y=out,
|
|
)
|
|
|
|
dequant_scale = torch.ones(
|
|
[M, moe_top_k], dtype=torch.float32, device=out.device
|
|
)
|
|
output = torch.empty(
|
|
[M, N], dtype=hidden_states.dtype, device=hidden_states.device
|
|
)
|
|
sorted_tokens_idx = sorted_tokens_idx.view(M, moe_top_k)
|
|
|
|
torch.ops._C.moe_post(
|
|
x=out,
|
|
moe_index=sorted_tokens_idx,
|
|
normed_scale=normed_score,
|
|
dequant_scale=dequant_scale,
|
|
y=output,
|
|
)
|
|
|
|
return output
|
|
|
|
@staticmethod
|
|
def fused_moe_ep(
|
|
hidden_states: torch.Tensor,
|
|
w13_weight: torch.Tensor,
|
|
w2_weight: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
linear_weights: torch.Tensor,
|
|
ep_rank: int,
|
|
top_k: int,
|
|
renormalize: bool,
|
|
inplace: bool = False,
|
|
use_grouped_topk: bool = False,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
w1_bias: Optional[torch.Tensor] = None,
|
|
w2_bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
x = hidden_states
|
|
batch, hidden_size = x.shape
|
|
num_local_experts, up_gate_size, _ = w13_weight.shape
|
|
|
|
router_logits = x.to(linear_weights.dtype) @ linear_weights.T
|
|
|
|
topk_weights = torch.empty(
|
|
batch, top_k, dtype=router_logits.dtype, device=router_logits.device
|
|
)
|
|
topk_ids = torch.empty(
|
|
batch, top_k, dtype=torch.int32, device=router_logits.device
|
|
)
|
|
block_static = torch.empty(0, dtype=torch.int32, device=router_logits.device)
|
|
torch.ops._C.moe_softmax_topk(
|
|
router_logits, topk_weights, topk_ids, block_static
|
|
)
|
|
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(1, keepdim=True)
|
|
|
|
topk_weights = topk_weights.to(x.dtype)
|
|
out = torch.zeros(batch * top_k, hidden_size, dtype=x.dtype, device=x.device)
|
|
repeat_x = x.repeat_interleave(top_k, dim=0)
|
|
topk_ids_flat = topk_ids.flatten()
|
|
for i in range(num_local_experts):
|
|
experts_id = ep_rank * num_local_experts + i
|
|
selected_token = topk_ids_flat == experts_id
|
|
if selected_token.sum():
|
|
cur_token = repeat_x[selected_token]
|
|
up_gate = torch.empty(
|
|
selected_token.sum(),
|
|
up_gate_size // 2,
|
|
dtype=cur_token.dtype,
|
|
device=cur_token.device,
|
|
)
|
|
torch.ops._C.silu_and_mul(up_gate, cur_token @ w13_weight[i].T)
|
|
out[selected_token] = up_gate @ w2_weight[i].T
|
|
output = (
|
|
(out.view(batch, top_k, hidden_size) * topk_weights.unsqueeze(2))
|
|
.sum(dim=1)
|
|
.to(x.dtype)
|
|
)
|
|
|
|
return output
|
|
|
|
@staticmethod
|
|
def fused_multi_head_latent_page_attention(
|
|
hidden_states: torch.Tensor,
|
|
q_lora_rank: int,
|
|
kv_lora_rank: int,
|
|
q_a_proj_w: torch.Tensor,
|
|
q_a_layernorm_w: torch.Tensor,
|
|
q_b_proj_w: torch.Tensor,
|
|
q_proj_w: torch.Tensor,
|
|
kv_a_proj_w: torch.Tensor,
|
|
kv_a_layernorm_w: torch.Tensor,
|
|
kv_b_proj_w: torch.Tensor,
|
|
o_proj_w: torch.Tensor,
|
|
head_num: int,
|
|
qk_nope_head_dim: int,
|
|
qk_rope_head_dim: int,
|
|
v_head_dim: int,
|
|
max_context_len: int,
|
|
layernorm_eps: float,
|
|
scale: float,
|
|
is_causal: bool,
|
|
is_context: bool,
|
|
mp_size: int,
|
|
local_rank: int,
|
|
rotary_pos_embedding: torch.Tensor,
|
|
pa_block_tables: torch.Tensor,
|
|
position: torch.Tensor,
|
|
context_lens_cpu: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
prompt_lods_cpu: torch.Tensor,
|
|
k_cache: torch.Tensor,
|
|
v_cache: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""mla pa block"""
|
|
output = torch.empty(
|
|
hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device
|
|
)
|
|
kunlun_ops.xft_multi_head_latent_page_attention_block(
|
|
hidden_states,
|
|
q_lora_rank,
|
|
kv_lora_rank,
|
|
q_a_proj_w,
|
|
q_a_layernorm_w,
|
|
q_b_proj_w,
|
|
q_proj_w,
|
|
kv_a_proj_w,
|
|
kv_a_layernorm_w,
|
|
kv_b_proj_w,
|
|
o_proj_w,
|
|
head_num,
|
|
qk_nope_head_dim,
|
|
qk_rope_head_dim,
|
|
v_head_dim,
|
|
max_context_len,
|
|
layernorm_eps,
|
|
scale,
|
|
is_causal,
|
|
is_context,
|
|
mp_size,
|
|
local_rank,
|
|
rotary_pos_embedding,
|
|
pa_block_tables,
|
|
position,
|
|
None,
|
|
context_lens_cpu,
|
|
slot_mapping,
|
|
None,
|
|
prompt_lods_cpu,
|
|
out=output,
|
|
k_cache=k_cache,
|
|
v_cache=v_cache,
|
|
)
|
|
return output
|
|
|
|
def fused_gdn_gating(
|
|
A_log: torch.Tensor,
|
|
a: torch.Tensor,
|
|
dt_bias: torch.Tensor,
|
|
beta: float = 1.0,
|
|
threshold: float = 20.0,
|
|
) -> torch.Tensor:
|
|
"""fused_gdn_gating"""
|
|
output = kunlun_ops.fused_gdn_gating(
|
|
A_log,
|
|
a,
|
|
dt_bias,
|
|
)
|
|
return output
|
|
|
|
def fused_recurrent_gated_delta_rule_fwd(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
g: torch.Tensor,
|
|
beta: torch.Tensor,
|
|
scale: float,
|
|
h0_source: torch.Tensor,
|
|
output_final_state: bool,
|
|
use_qk_l2norm_in_kernel: bool,
|
|
cu_seqlens: torch.Tensor = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Qwen3-NEXT模型中 Gated DeltaNet的核心算子, 将做完sigmoid_gating和delta_rule_update融合在一起
|
|
1. Sigmoid Gating: 对输入进行门控, 类似于 GLU (Gated Linear Unit)。
|
|
2. Delta Rule Update: 执行一个并行的状态空间模型(SSM)的递归更新, 同时结合了一个局部的注意力机制。
|
|
"""
|
|
|
|
o, final_state = kunlun_ops.fused_recurrent_gated_delta_rule_fwd(
|
|
q,
|
|
k,
|
|
v,
|
|
g,
|
|
beta,
|
|
scale,
|
|
h0_source,
|
|
output_final_state,
|
|
use_qk_l2norm_in_kernel,
|
|
cu_seqlens,
|
|
)
|
|
return (o, final_state)
|