127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
import vllm.model_executor.layers.fla.ops
|
|
import vllm.model_executor.layers.mamba.ops.causal_conv1d
|
|
import vllm.v1.worker.gpu.sample.gumbel
|
|
from vllm.triton_utils import HAS_TRITON, triton
|
|
from vllm.utils.math_utils import next_power_of_2
|
|
|
|
from vllm_ascend.ops.triton.fla.chunk import chunk_gated_delta_rule
|
|
from vllm_ascend.ops.triton.fla.layernorm_guard import LayerNormFn
|
|
from vllm_ascend.ops.triton.fla.sigmoid_gating import fused_recurrent_gated_delta_rule_fwd_kernel
|
|
from vllm_ascend.ops.triton.mamba.causal_conv1d import causal_conv1d_update_npu
|
|
|
|
triton.next_power_of_2 = next_power_of_2
|
|
|
|
vllm.model_executor.layers.mamba.ops.causal_conv1d.causal_conv1d_update = causal_conv1d_update_npu
|
|
vllm.model_executor.layers.fla.ops.fused_recurrent.fused_recurrent_gated_delta_rule_fwd_kernel = (
|
|
fused_recurrent_gated_delta_rule_fwd_kernel
|
|
)
|
|
vllm.model_executor.layers.fla.ops.layernorm_guard.LayerNormFn = LayerNormFn
|
|
vllm.model_executor.layers.fla.ops.chunk_gated_delta_rule = chunk_gated_delta_rule
|
|
|
|
# On NPU platforms without an active Triton backend (e.g. 310P), replace the
|
|
# Triton-based fused_post_conv_prep with a pure-PyTorch fallback so that
|
|
# qwen_gdn_linear_attn's from-import picks up the replacement before model
|
|
# load.
|
|
if not HAS_TRITON:
|
|
import torch
|
|
import torch.nn.functional as _F
|
|
|
|
def _fused_post_conv_prep_pytorch(
|
|
conv_output,
|
|
a,
|
|
b,
|
|
A_log,
|
|
dt_bias,
|
|
num_k_heads,
|
|
head_k_dim,
|
|
head_v_dim,
|
|
apply_l2norm=True,
|
|
output_g_exp=False,
|
|
):
|
|
L = conv_output.shape[0]
|
|
H, K, V = num_k_heads, head_k_dim, head_v_dim
|
|
HV = A_log.shape[0]
|
|
|
|
q = conv_output[:, : H * K].reshape(L, H, K)
|
|
k = conv_output[:, H * K : 2 * H * K].reshape(L, H, K)
|
|
v = conv_output[:, 2 * H * K :].reshape(L, HV, V)
|
|
|
|
if apply_l2norm:
|
|
# x / sqrt(sum(x^2) + eps) — matches Triton kernel, in fp32
|
|
def _l2norm(t):
|
|
t_f = t.float()
|
|
return (t_f / torch.sqrt((t_f * t_f).sum(-1, keepdim=True) + 1e-6)).to(t.dtype)
|
|
|
|
q, k = _l2norm(q), _l2norm(k)
|
|
|
|
q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
|
|
|
|
x = (a + dt_bias.unsqueeze(0)).float()
|
|
g = -torch.exp(A_log.float().unsqueeze(0)) * _F.softplus(x)
|
|
if output_g_exp:
|
|
g = torch.exp(g)
|
|
|
|
return q, k, v, g, torch.sigmoid(b.float())
|
|
|
|
vllm.model_executor.layers.fla.ops.fused_post_conv_prep = _fused_post_conv_prep_pytorch
|
|
|
|
def _fused_recurrent_packed_decode_pytorch(
|
|
mixed_qkv,
|
|
a,
|
|
b,
|
|
A_log,
|
|
dt_bias,
|
|
scale,
|
|
initial_state,
|
|
out,
|
|
ssm_state_indices,
|
|
use_qk_l2norm_in_kernel=False,
|
|
):
|
|
B = mixed_qkv.shape[0]
|
|
HV, V, K = initial_state.shape[-3:]
|
|
H = (mixed_qkv.shape[1] - HV * V) // (2 * K)
|
|
ratio = HV // H
|
|
|
|
q = mixed_qkv[:, : H * K].reshape(B, H, K)
|
|
k = mixed_qkv[:, H * K : 2 * H * K].reshape(B, H, K)
|
|
v = mixed_qkv[:, 2 * H * K :].reshape(B, HV, V)
|
|
|
|
SOFTPLUS_THRESHOLD = 20.0
|
|
x = (a + dt_bias.unsqueeze(0)).float()
|
|
softplus_x = torch.where(x <= SOFTPLUS_THRESHOLD, torch.log1p(torch.exp(x)), x)
|
|
g = -torch.exp(A_log.float().unsqueeze(0)) * softplus_x # [B, HV]
|
|
beta = torch.sigmoid(b.float()) # [B, HV]
|
|
|
|
for n in range(B):
|
|
state_idx = int(ssm_state_indices[n].item())
|
|
if state_idx <= 0:
|
|
out[n, 0] = 0
|
|
continue
|
|
|
|
h = initial_state[state_idx].float() # [HV, V, K]
|
|
q_n = q[n].float().repeat_interleave(ratio, dim=0) # [HV, K]
|
|
k_n = k[n].float().repeat_interleave(ratio, dim=0) # [HV, K]
|
|
v_n = v[n].float() # [HV, V]
|
|
|
|
if use_qk_l2norm_in_kernel:
|
|
|
|
def _l2norm(t):
|
|
t_f = t.float()
|
|
return t_f / torch.sqrt((t_f * t_f).sum(-1, keepdim=True) + 1e-6)
|
|
|
|
q_n, k_n = _l2norm(q_n), _l2norm(k_n)
|
|
q_n = q_n * scale
|
|
|
|
h = h * torch.exp(g[n]).view(HV, 1, 1)
|
|
v_n = v_n - torch.einsum("hvk,hk->hv", h, k_n)
|
|
v_n = v_n * beta[n].view(HV, 1)
|
|
h = h + torch.einsum("hv,hk->hvk", v_n, k_n)
|
|
out[n, 0] = torch.einsum("hvk,hk->hv", h, q_n).to(out.dtype)
|
|
initial_state[state_idx] = h.to(initial_state.dtype)
|
|
|
|
return out, initial_state
|
|
|
|
vllm.model_executor.layers.fla.ops.fused_recurrent.fused_recurrent_gated_delta_rule_packed_decode = (
|
|
_fused_recurrent_packed_decode_pytorch
|
|
)
|