fangfangssj/enginex-vllm-bi100-qwen36
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix perforence

Browse Source
This commit is contained in:
root
2026-07-14 13:11:21 +00:00
parent 1902c81fdd
commit d0585b4a17
9 changed files with 2414 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
qwen3_6_scripts/qwen3_5.py
View File

@@ -114,11 +114,32 @@ def _torch_chunk_gated_delta_rule(
g = g.cumsum(dim=-1)
decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0)
for i in range(1, chunk_size):
row = attn[..., i, :i].clone()
sub = attn[..., :i, :i].clone()
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
# loop1: 三角求逆 (I - attn)^{-1}。默认走 solve_triangular (batched trsm kernel,
# 替代 63 步 Python 前向替换循环, 真实服务隔离 4-5x, 数值等价 max_abs<1e-7)。
# 传 -attn + unitriangular=True 等效于求解 (I-attn)@X=I → X=(I-attn)^{-1}。
# 注意: 不做运行时 attn.abs().max() 阈值检查(会引入 device→host 同步,
# 600次调用下额外耗时~1-2s)。真实模型 attn 值域远在安全范围内。
# LOOP1_NATIVE=0 可回退串行循环。
import os as _os
_loop1_done = False
if _os.environ.get("LOOP1_NATIVE", "1") != "0":
try:
_eye = torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
attn = torch.linalg.solve_triangular(
-attn, _eye, upper=False, unitriangular=True)
_loop1_done = True
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[loop1_native FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
if not _loop1_done:
for i in range(1, chunk_size):
row = attn[..., i, :i].clone()
sub = attn[..., :i, :i].clone()
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
value = attn @ v_beta
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
@@ -132,18 +153,54 @@ def _torch_chunk_gated_delta_rule(
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
diagonal=1)
for i in range(total_len // chunk_size):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
v_prime = k_cumdecay[:, :, i] @ last_state
v_new = v_i - v_prime
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
core_out[:, :, i] = attn_inter + attn_i @ v_new
last_state = (
last_state * g[:, :, i, -1, None, None].exp()
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
.transpose(-1, -2) @ v_new
)
# loop2: chunk 间 state 递推。默认走 chunk-parallel(把串行 scan 改写成
# 线性矩阵递归 S_i=A_i@S_{i-1}+B_i,重活批量化,串行阶段只对小矩阵 S 做 A@S+B)。
# 数值等价于下方串行实现(已验证 max_abs~1e-7, 3x)。CHUNK_PARALLEL=0 可回退。
import os as _os
NC = total_len // chunk_size
_done = False
if _os.environ.get("CHUNK_PARALLEL", "1") != "0":
try:
g_last = g[:, :, :, -1] # [b,h,NC]
exp_glast = g_last.exp()
attn_i_all = (query @ key.transpose(-1, -2) * decay_mask
).masked_fill(mask_upper2, 0) # [b,h,NC,cs,cs]
P_i = k_cumdecay # [b,h,NC,cs,k]
Ktil = key * (g_last[..., None, None] - g[..., None]).exp() # [b,h,NC,cs,k]
Qg = query * g.exp()[..., None] # [b,h,NC,cs,k]
eye = torch.eye(k_dim, device=query.device, dtype=query.dtype)
A = exp_glast[..., None, None] * eye - Ktil.transpose(-1, -2) @ P_i # [b,h,NC,k,k]
B = Ktil.transpose(-1, -2) @ value # [b,h,NC,k,v]
M = Qg - attn_i_all @ P_i # [b,h,NC,cs,k]
N = attn_i_all @ value # [b,h,NC,cs,v]
S = last_state
S_starts = torch.empty(batch, num_heads, NC, k_dim, v_dim,
device=query.device, dtype=query.dtype)
for i in range(NC):
S_starts[:, :, i] = S
S = A[:, :, i] @ S + B[:, :, i]
last_state = S
core_out = M @ S_starts + N # [b,h,NC,cs,v]
_done = True
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[chunk_parallel FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
if not _done:
for i in range(NC):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
v_prime = k_cumdecay[:, :, i] @ last_state
v_new = v_i - v_prime
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
core_out[:, :, i] = attn_inter + attn_i @ v_new
last_state = (
last_state * g[:, :, i, -1, None, None].exp()
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
.transpose(-1, -2) @ v_new
)
if not output_final_state:
last_state = None
@@ -212,6 +269,20 @@ class Qwen3_5RMSNormGated(nn.Module):
def forward(self, hidden_states: torch.Tensor,
gate: torch.Tensor) -> torch.Tensor:
# 默认用原生 ixformer rms_norm(替代纯 PyTorch pow/mean/rsqrt),
# 再乘 silu(gate)。已验证与下方纯 PyTorch 等价(max_abs~1e-2 fp16)。
# RMSNORM_NATIVE=0 或异常回退纯 PyTorch。
import os as _os3
if _os3.environ.get("RMSNORM_NATIVE", "1") != "0":
try:
from vllm import _custom_ops as _ops
out = torch.empty_like(hidden_states)
_ops.rms_norm(out, hidden_states.contiguous(),
self.weight, self.variance_epsilon)
return (out.to(torch.float32)
* F.silu(gate.to(torch.float32))).to(hidden_states.dtype)
except Exception:
pass
input_dtype = hidden_states.dtype
hs = hidden_states.to(torch.float32)
variance = hs.pow(2).mean(-1, keepdim=True)
@@ -792,8 +863,22 @@ class Qwen3_5MoeSparseBlock(nn.Module):
out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
hidden_states.dtype) # (1, H)
else:
# General path (prefill / multi-seq): loop over unique active experts.
# At most T*top_k unique experts, always <= num_experts.
# General path (prefill / multi-seq).
# 优化: sort-by-expert + 连续切片 + 原生 kernel(act_bias_mm/silu_and_mul/
# ixf linear),消除逐 expert 的 nonzero/index_add 循环调度开销。
# 数值等价于下方纯 PyTorch fallback(已验证 byte-for-byte 一致, 端到端1.36x)。
# 任意异常回退纯 PyTorch。可用 MOE_NATIVE=0 关闭。
import os as _os
if _os.environ.get("MOE_NATIVE", "1") != "0":
try:
return self._native_experts_sorted(
hidden_states, w13, w2, topk_weights, topk_ids)
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[moe_native FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
# ---- 纯 PyTorch fallback (原实现) ----
out = torch.zeros_like(hidden_states)
unique_eids = topk_ids.view(-1).unique().tolist()
for eid in unique_eids:
@@ -810,6 +895,46 @@ class Qwen3_5MoeSparseBlock(nn.Module):
return out # partial, all-reduce done in forward()
def _native_experts_sorted(self, hidden_states, w13, w2, topk_weights, topk_ids):
"""sort-by-expert + 原生 ixformer kernel 的 MoE expert 计算。
w13=(E,2I,H), w2=(E,H,I)。返回 partial(pre-all-reduce)。"""
import ixformer.functions as _ixf
dev = hidden_states.device
T = hidden_states.shape[0]
top_k = topk_ids.shape[1]
E = w13.shape[0]
flat_e = topk_ids.reshape(-1) # (T*K,)
flat_w = topk_weights.reshape(-1) # (T*K,)
flat_tok = torch.arange(T, device=dev).repeat_interleave(top_k) # (T*K,)
order = torch.argsort(flat_e)
se = flat_e[order]
sw = flat_w[order]
stok = flat_tok[order]
gathered = hidden_states.index_select(0, stok).contiguous() # (T*K, H) 已按 expert 连续
counts = torch.bincount(se, minlength=E)
offs = torch.cat([torch.zeros(1, device=dev, dtype=torch.long),
counts.cumsum(0)]).tolist()
# 就地把每个 expert 段的输出写回 gathered(复用显存,避免再分配 T*K×H,
# 否则 profiling 阶段峰值显存↑ → KV cache 容量不足)。
active = (counts > 0).nonzero().flatten().tolist()
for eid in active:
eid = int(eid)
s, e = int(offs[eid]), int(offs[eid + 1])
seg = gathered[s:e].contiguous() # (n, H)
gu = _ixf.act_bias_mm(seg, w13[eid], None, scale=1,
act_type="none", trans_format="TN") # (n, 2I)
act = _ixf.silu_and_mul(gu.contiguous()) # (n, I)
gathered[s:e] = _ixf.linear(act.contiguous(), w2[eid], None) # (n, H) 写回
gathered = gathered * sw.unsqueeze(-1).to(gathered.dtype)
out = torch.zeros_like(hidden_states)
out.index_add_(0, stok, gathered.to(out.dtype))
return out
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
router_logits, _ = self.gate(hidden_states)
routed_out = self._pure_pytorch_experts(hidden_states, router_logits)