align platform build config with working submission
This commit is contained in:
@@ -3,9 +3,6 @@
|
||||
# Text-only (no VL, no MTP).
|
||||
|
||||
from collections import OrderedDict
|
||||
from contextlib import contextmanager
|
||||
import os
|
||||
import time
|
||||
from typing import Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
@@ -44,60 +41,20 @@ from vllm.model_executor.models.interfaces import HasInnerState, SupportsLoRA
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_ENGINEX_PROFILE_ENABLED = os.getenv("ENGINEX_PROFILE_DECODE", "0") == "1"
|
||||
_ENGINEX_PROFILE_EVERY = int(os.getenv("ENGINEX_PROFILE_EVERY", "32"))
|
||||
_ENGINEX_PROFILE_SYNC = os.getenv("ENGINEX_PROFILE_SYNC", "1") != "0"
|
||||
_ENGINEX_MOE_TINY_IMPL = os.getenv("ENGINEX_MOE_TINY_IMPL", "tokenwise")
|
||||
_ENGINEX_MOE_TINY_MAX = int(os.getenv("ENGINEX_MOE_TINY_MAX", "4"))
|
||||
_enginex_profile_stats: Dict[str, List[float]] = {}
|
||||
_enginex_profile_steps = 0
|
||||
_enginex_profile_mode = "unknown"
|
||||
|
||||
|
||||
def _enginex_profile_active() -> bool:
|
||||
return _ENGINEX_PROFILE_ENABLED and torch.cuda.is_available()
|
||||
|
||||
|
||||
def _enginex_profile_sync() -> None:
|
||||
if _ENGINEX_PROFILE_SYNC:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _enginex_profile(label: str):
|
||||
if not _enginex_profile_active():
|
||||
yield
|
||||
return
|
||||
label = f"{_enginex_profile_mode}.{label}"
|
||||
_enginex_profile_sync()
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_enginex_profile_sync()
|
||||
elapsed_ms = (time.perf_counter() - start) * 1000.0
|
||||
stat = _enginex_profile_stats.setdefault(label, [0.0, 0.0])
|
||||
stat[0] += elapsed_ms
|
||||
stat[1] += 1.0
|
||||
|
||||
|
||||
def _enginex_profile_log(mode: str) -> None:
|
||||
global _enginex_profile_steps
|
||||
if not _enginex_profile_active():
|
||||
return
|
||||
_enginex_profile_steps += 1
|
||||
if _enginex_profile_steps % max(_ENGINEX_PROFILE_EVERY, 1) != 0:
|
||||
return
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
parts = []
|
||||
for label, (total_ms, count) in sorted(
|
||||
_enginex_profile_stats.items(),
|
||||
key=lambda item: item[1][0],
|
||||
reverse=True):
|
||||
avg_ms = total_ms / max(count, 1.0)
|
||||
parts.append(f"{label}: total={total_ms:.2f}ms avg={avg_ms:.3f}ms n={int(count)}")
|
||||
logger.info("[ENGINEX_PROFILE_QWEN] rank=%d steps=%d mode=%s %s",
|
||||
tp_rank, _enginex_profile_steps, mode, " | ".join(parts))
|
||||
# ---------------------------------------------------------------------------
|
||||
# Profiling debug prints (locate where startup hangs after weight load).
|
||||
# Gate on PROF_TRACE=1 (default on, since we are diagnosing a startup hang).
|
||||
# Each print flushes; grep "[PROF]" in the server log.
|
||||
# ---------------------------------------------------------------------------
|
||||
import os as _os_prof
|
||||
import sys as _sys_prof
|
||||
import time as _time_prof
|
||||
_PROF_TRACE = _os_prof.environ.get("PROF_TRACE", "1") != "0"
|
||||
_t0 = _time_prof.time()
|
||||
def _prof(msg):
|
||||
if _PROF_TRACE:
|
||||
print(f"[PROF] {_time_prof.time()-_t0:7.2f}s {msg}", file=_sys_prof.stderr, flush=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -172,11 +129,34 @@ 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)
|
||||
_prof(f" chunk-gdn loop1 solve_triangular start shape={tuple(attn.shape)}")
|
||||
attn = torch.linalg.solve_triangular(
|
||||
-attn, _eye, upper=False, unitriangular=True)
|
||||
_prof(f" chunk-gdn loop1 solve_triangular done")
|
||||
_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))
|
||||
|
||||
@@ -190,18 +170,56 @@ 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]
|
||||
_prof(f" chunk-gdn loop2 parallel batched done NC={NC}")
|
||||
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]
|
||||
_prof(f" chunk-gdn loop2 serial scan done ({NC} iters)")
|
||||
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
|
||||
@@ -270,6 +288,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)
|
||||
@@ -417,6 +449,7 @@ class GatedDeltaNet(nn.Module):
|
||||
mixed_qkv_conv = F.silu(mixed_qkv_conv)
|
||||
# (1, seq_len, local_conv_dim)
|
||||
mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0)
|
||||
_prof(f" L{self.layer_idx} GDN conv1d done seq_len={seq_len}")
|
||||
|
||||
q, k, v = torch.split(
|
||||
mixed_qkv_conv,
|
||||
@@ -442,6 +475,7 @@ class GatedDeltaNet(nn.Module):
|
||||
_DNN_CHUNK = 4096
|
||||
cur_state = temporal_state[si:si + 1].clone()
|
||||
core_out_parts = []
|
||||
_prof(f" L{self.layer_idx} GDN chunk-scan start (DNN_CHUNK={_DNN_CHUNK}, nchunks={(seq_len + _DNN_CHUNK - 1)//_DNN_CHUNK})")
|
||||
for sc_start in range(0, seq_len, _DNN_CHUNK):
|
||||
sc_end = min(sc_start + _DNN_CHUNK, seq_len)
|
||||
c_out, cur_state = _torch_chunk_gated_delta_rule(
|
||||
@@ -455,6 +489,7 @@ class GatedDeltaNet(nn.Module):
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
core_out_parts.append(c_out)
|
||||
_prof(f" L{self.layer_idx} GDN chunk-scan done")
|
||||
if cur_state is not None:
|
||||
temporal_state[si].copy_(cur_state[0])
|
||||
# [1, seq_len, num_v_heads, head_v_dim]
|
||||
@@ -671,48 +706,46 @@ class Qwen3_5FullAttention(nn.Module):
|
||||
total_tokens = hidden_states.shape[0]
|
||||
|
||||
# q_proj output includes gate (dim doubled)
|
||||
with _enginex_profile("full_attn.qkv_proj"):
|
||||
qg, _ = self.q_proj(hidden_states) # (total, local_num_heads * head_dim * 2)
|
||||
qg = qg.view(total_tokens, self.local_num_heads, self.head_dim * 2)
|
||||
q = qg[:, :, :self.head_dim].reshape(total_tokens, -1)
|
||||
gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
|
||||
qg, _ = self.q_proj(hidden_states) # (total, local_num_heads * head_dim * 2)
|
||||
qg = qg.view(total_tokens, self.local_num_heads, self.head_dim * 2)
|
||||
q = qg[:, :, :self.head_dim].reshape(total_tokens, -1)
|
||||
gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
|
||||
|
||||
k, _ = self.k_proj(hidden_states) # (total, proj_kv_heads * head_dim)
|
||||
v, _ = self.v_proj(hidden_states)
|
||||
k, _ = self.k_proj(hidden_states) # (total, proj_kv_heads * head_dim)
|
||||
v, _ = self.v_proj(hidden_states)
|
||||
|
||||
# q_norm on local Q heads
|
||||
with _enginex_profile("full_attn.norm_rope"):
|
||||
q = self.q_norm.forward_cuda(
|
||||
q.view(total_tokens, self.local_num_heads, self.head_dim)
|
||||
.contiguous()).view(total_tokens, -1)
|
||||
q = self.q_norm.forward_cuda(
|
||||
q.view(total_tokens, self.local_num_heads, self.head_dim)
|
||||
.contiguous()).view(total_tokens, -1)
|
||||
|
||||
# GQA-aware TP: select rank-local KV head BEFORE k_norm and rope so
|
||||
# that ixformer kernels always see num_kv_heads=1 (same as 27B path).
|
||||
# Doing k_norm/rope on 2 KV heads (proj_kv_heads=2) triggers ixformer
|
||||
# paths that can produce NaN; restricting to 1 head avoids the issue.
|
||||
if self.q_per_kv_global is not None:
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
kv_idx = (tp_rank * self.local_num_heads) // self.q_per_kv_global
|
||||
k = (k.view(total_tokens, self.proj_kv_heads, self.head_dim)
|
||||
[:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head
|
||||
v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim)
|
||||
[:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head
|
||||
# GQA-aware TP: select rank-local KV head BEFORE k_norm and rope so
|
||||
# that ixformer kernels always see num_kv_heads=1 (same as 27B path).
|
||||
# Doing k_norm/rope on 2 KV heads (proj_kv_heads=2) triggers ixformer
|
||||
# paths that can produce NaN; restricting to 1 head avoids the issue.
|
||||
if self.q_per_kv_global is not None:
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
kv_idx = (tp_rank * self.local_num_heads) // self.q_per_kv_global
|
||||
k = (k.view(total_tokens, self.proj_kv_heads, self.head_dim)
|
||||
[:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head
|
||||
v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim)
|
||||
[:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head
|
||||
|
||||
# k_norm on the (now always 1) rank-local KV head
|
||||
k = self.k_norm.forward_cuda(
|
||||
k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
|
||||
.contiguous()).view(total_tokens, -1)
|
||||
# k_norm on the (now always 1) rank-local KV head
|
||||
k = self.k_norm.forward_cuda(
|
||||
k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
|
||||
.contiguous()).view(total_tokens, -1)
|
||||
|
||||
# rope: q=(T, local_num_heads*head_dim), k=(T, 1*head_dim) — mirrors 27B
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
# rope: q=(T, local_num_heads*head_dim), k=(T, 1*head_dim) — mirrors 27B
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
|
||||
with _enginex_profile("full_attn.paged_attention"):
|
||||
attn_out = self.attn(q, k, v, kv_cache, attn_metadata)
|
||||
_prof(f" L{self.layer_idx} ATTN flash start tokens={q.shape[0]}")
|
||||
attn_out = self.attn(q, k, v, kv_cache, attn_metadata)
|
||||
_prof(f" L{self.layer_idx} ATTN flash done")
|
||||
|
||||
# Multiply by sigmoid gate before output projection
|
||||
with _enginex_profile("full_attn.gate_o_proj"):
|
||||
attn_out = attn_out * torch.sigmoid(gate.float()).to(attn_out.dtype)
|
||||
output, _ = self.o_proj(attn_out)
|
||||
attn_out = attn_out * torch.sigmoid(gate.float()).to(attn_out.dtype)
|
||||
output, _ = self.o_proj(attn_out)
|
||||
return output
|
||||
|
||||
|
||||
@@ -816,13 +849,12 @@ class Qwen3_5MoeSparseBlock(nn.Module):
|
||||
Output is partial (pre-all-reduce), same contract as FusedMoE
|
||||
with reduce_results=False.
|
||||
"""
|
||||
# Routing: softmax -> topk -> renormalise
|
||||
with _enginex_profile("moe.routing_topk"):
|
||||
routing_weights = torch.softmax(router_logits.float(), dim=-1)
|
||||
topk_weights, topk_ids = torch.topk(
|
||||
routing_weights, self.top_k, dim=-1) # (T, top_k)
|
||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
topk_weights = topk_weights.to(hidden_states.dtype)
|
||||
# Routing: softmax → topk → renormalise
|
||||
routing_weights = torch.softmax(router_logits.float(), dim=-1)
|
||||
topk_weights, topk_ids = torch.topk(
|
||||
routing_weights, self.top_k, dim=-1) # (T, top_k)
|
||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
topk_weights = topk_weights.to(hidden_states.dtype)
|
||||
|
||||
w13 = self.experts.w13_weight # (E, 2*I, H)
|
||||
w2 = self.experts.w2_weight # (E, H, I)
|
||||
@@ -834,114 +866,116 @@ class Qwen3_5MoeSparseBlock(nn.Module):
|
||||
# gate_up: 1 large GEMM (1,H) × (K*2*I,H)^T → (1, K*2*I)
|
||||
# down: 1 bmm (K,H,I) @ (K,I,1) → (K,H)
|
||||
# Total: 3 kernel launches vs previous 16 (top_k*2).
|
||||
with _enginex_profile("moe.routed_decode_experts"):
|
||||
eids = topk_ids[0] # (K,)
|
||||
ws = topk_weights[0].to(hidden_states.dtype) # (K,)
|
||||
w13_sel = w13[eids] # (K, 2*I, H)
|
||||
w2_sel = w2[eids] # (K, H, I)
|
||||
eids = topk_ids[0] # (K,)
|
||||
ws = topk_weights[0].to(hidden_states.dtype) # (K,)
|
||||
w13_sel = w13[eids] # (K, 2*I, H)
|
||||
w2_sel = w2[eids] # (K, H, I)
|
||||
|
||||
H = hidden_states.shape[-1]
|
||||
H = hidden_states.shape[-1]
|
||||
|
||||
gate_up = F.linear(
|
||||
hidden_states,
|
||||
w13_sel.reshape(-1, H), # (K*2*I, H) - contiguous after indexing
|
||||
) # (1, K*2*I)
|
||||
gate_up = gate_up.view(self.top_k, -1) # (K, 2*I)
|
||||
gate, up = gate_up.chunk(2, dim=-1) # (K, I) each
|
||||
act = F.silu(gate) * up # (K, I)
|
||||
gate_up = F.linear(
|
||||
hidden_states,
|
||||
w13_sel.reshape(-1, H), # (K*2*I, H) — contiguous after indexing
|
||||
) # (1, K*2*I)
|
||||
gate_up = gate_up.view(self.top_k, -1) # (K, 2*I)
|
||||
gate, up = gate_up.chunk(2, dim=-1) # (K, I) each
|
||||
act = F.silu(gate) * up # (K, I)
|
||||
|
||||
# bmm: (K,H,I) @ (K,I,1) -> (K,H,1) -> (K,H)
|
||||
expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H)
|
||||
# bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H)
|
||||
expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H)
|
||||
|
||||
out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
|
||||
hidden_states.dtype) # (1, H)
|
||||
elif T <= _ENGINEX_MOE_TINY_MAX and _ENGINEX_MOE_TINY_IMPL == "tokenwise":
|
||||
# Fast path: tiny decode batch. Compute each token with the proven
|
||||
# T==1 large-F.linear path, avoiding the per-expert Python loop.
|
||||
# This usually beats batched bmm for very small T because the first
|
||||
# projection becomes T larger GEMMs instead of T*K tiny GEMMs.
|
||||
with _enginex_profile("moe.routed_tiny_tokenwise_experts"):
|
||||
H = hidden_states.shape[-1]
|
||||
pieces = []
|
||||
for t in range(T):
|
||||
eids = topk_ids[t]
|
||||
ws = topk_weights[t].to(hidden_states.dtype)
|
||||
w13_sel = w13[eids] # (K, 2*I, H)
|
||||
w2_sel = w2[eids] # (K, H, I)
|
||||
|
||||
gate_up = F.linear(
|
||||
hidden_states[t:t + 1],
|
||||
w13_sel.reshape(-1, H),
|
||||
).view(self.top_k, -1) # (K, 2*I)
|
||||
gate, up = gate_up.chunk(2, dim=-1)
|
||||
act = F.silu(gate) * up
|
||||
expert_out = torch.bmm(
|
||||
w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H)
|
||||
pieces.append((expert_out * ws.unsqueeze(-1)).sum(0))
|
||||
|
||||
out = torch.stack(pieces, dim=0).to(hidden_states.dtype)
|
||||
elif T <= _ENGINEX_MOE_TINY_MAX:
|
||||
# Alternative tiny decode batch implementation. Kept for A/B
|
||||
# testing via ENGINEX_MOE_TINY_IMPL=bmm.
|
||||
with _enginex_profile("moe.routed_tiny_bmm_experts"):
|
||||
H = hidden_states.shape[-1]
|
||||
flat_eids = topk_ids.reshape(-1) # (T*K,)
|
||||
flat_ws = topk_weights.reshape(-1).to(hidden_states.dtype)
|
||||
|
||||
w13_sel = w13[flat_eids] # (T*K, 2*I, H)
|
||||
w2_sel = w2[flat_eids] # (T*K, H, I)
|
||||
x = (hidden_states[:, None, :]
|
||||
.expand(T, self.top_k, H)
|
||||
.reshape(-1, 1, H)) # (T*K, 1, H)
|
||||
|
||||
gate_up = torch.bmm(
|
||||
x, w13_sel.transpose(1, 2)).squeeze(1) # (T*K, 2*I)
|
||||
gate, up = gate_up.chunk(2, dim=-1)
|
||||
act = F.silu(gate) * up # (T*K, I)
|
||||
expert_out = torch.bmm(
|
||||
w2_sel, act.unsqueeze(-1)).squeeze(-1) # (T*K, H)
|
||||
|
||||
out = (expert_out * flat_ws.unsqueeze(-1)).view(
|
||||
T, self.top_k, H).sum(1).to(hidden_states.dtype)
|
||||
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.
|
||||
with _enginex_profile("moe.routed_prefill_experts"):
|
||||
out = torch.zeros_like(hidden_states)
|
||||
unique_eids = topk_ids.view(-1).unique().tolist()
|
||||
for eid in unique_eids:
|
||||
eid = int(eid)
|
||||
mask = (topk_ids == eid) # (T, top_k)
|
||||
tok_ids, topk_pos = mask.nonzero(as_tuple=True)
|
||||
tokens = hidden_states[tok_ids] # (n, H)
|
||||
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
|
||||
gate, up = gate_up.chunk(2, dim=-1)
|
||||
act = F.silu(gate) * up # (n, I)
|
||||
expert_out = F.linear(act, w2[eid]) # (n, H)
|
||||
weights = topk_weights[tok_ids, topk_pos].unsqueeze(-1)
|
||||
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
|
||||
# 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:
|
||||
eid = int(eid)
|
||||
mask = (topk_ids == eid) # (T, top_k)
|
||||
tok_ids, topk_pos = mask.nonzero(as_tuple=True)
|
||||
tokens = hidden_states[tok_ids] # (n, H)
|
||||
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
|
||||
gate, up = gate_up.chunk(2, dim=-1)
|
||||
act = F.silu(gate) * up # (n, I)
|
||||
expert_out = F.linear(act, w2[eid]) # (n, H)
|
||||
weights = topk_weights[tok_ids, topk_pos].unsqueeze(-1)
|
||||
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
|
||||
|
||||
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:
|
||||
with _enginex_profile("moe.gate"):
|
||||
router_logits, _ = self.gate(hidden_states)
|
||||
with _enginex_profile("moe.routed_total"):
|
||||
routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
|
||||
router_logits, _ = self.gate(hidden_states)
|
||||
routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
|
||||
_prof(f" MoE routed-experts done tokens={hidden_states.shape[0]}")
|
||||
|
||||
with _enginex_profile("moe.shared_expert"):
|
||||
gate_up, _ = self.shared_expert_gate_up(hidden_states)
|
||||
shared_out = self.act_fn(gate_up)
|
||||
shared_out, _ = self.shared_expert_down(shared_out)
|
||||
# Scalar sigmoid gate (Qwen2-MoE / Qwen3.5-MoE style)
|
||||
gate_score, _ = self.shared_expert_gate(hidden_states) # (T, 1)
|
||||
shared_out = shared_out * torch.sigmoid(gate_score)
|
||||
gate_up, _ = self.shared_expert_gate_up(hidden_states)
|
||||
shared_out = self.act_fn(gate_up)
|
||||
shared_out, _ = self.shared_expert_down(shared_out)
|
||||
# Scalar sigmoid gate (Qwen2-MoE / Qwen3.5-MoE style)
|
||||
gate_score, _ = self.shared_expert_gate(hidden_states) # (T, 1)
|
||||
shared_out = shared_out * torch.sigmoid(gate_score)
|
||||
|
||||
with _enginex_profile("moe.combine"):
|
||||
out = routed_out + shared_out
|
||||
out = routed_out + shared_out
|
||||
if self.experts.tp_size > 1:
|
||||
with _enginex_profile("moe.tp_all_reduce"):
|
||||
out = tensor_model_parallel_all_reduce(out)
|
||||
_prof(f" MoE all_reduce start tp_size={self.experts.tp_size}")
|
||||
out = tensor_model_parallel_all_reduce(out)
|
||||
_prof(f" MoE all_reduce done")
|
||||
return out
|
||||
|
||||
|
||||
@@ -1001,27 +1035,27 @@ class Qwen3_5DecoderLayer(nn.Module):
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
with _enginex_profile("layer.input_norm"):
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
else:
|
||||
with _enginex_profile("layer.input_norm"):
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
|
||||
if self.layer_type == "linear_attention":
|
||||
with _enginex_profile("layer.linear_attention"):
|
||||
hidden_states = self.linear_attn(
|
||||
hidden_states, attn_metadata, conv_state, temporal_state)
|
||||
_prof(f"L{self.layer_idx} GDN-start tokens={hidden_states.shape[0]}")
|
||||
hidden_states = self.linear_attn(
|
||||
hidden_states, attn_metadata, conv_state, temporal_state)
|
||||
_prof(f"L{self.layer_idx} GDN-done")
|
||||
else:
|
||||
with _enginex_profile("layer.full_attention"):
|
||||
hidden_states = self.self_attn(
|
||||
positions, hidden_states, kv_cache, attn_metadata)
|
||||
_prof(f"L{self.layer_idx} ATTN-start tokens={hidden_states.shape[0]}")
|
||||
hidden_states = self.self_attn(
|
||||
positions, hidden_states, kv_cache, attn_metadata)
|
||||
_prof(f"L{self.layer_idx} ATTN-done")
|
||||
|
||||
with _enginex_profile("layer.post_attn_norm"):
|
||||
hidden_states, residual = self.post_attention_layernorm(
|
||||
hidden_states, residual)
|
||||
hidden_states, residual = self.post_attention_layernorm(
|
||||
hidden_states, residual)
|
||||
|
||||
with _enginex_profile("layer.mlp"):
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
_prof(f"L{self.layer_idx} MLP-start")
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
_prof(f"L{self.layer_idx} MLP-done")
|
||||
|
||||
return hidden_states, residual
|
||||
|
||||
@@ -1058,9 +1092,6 @@ class Qwen3_5Model(nn.Module):
|
||||
conv_states: torch.Tensor, # (num_linear_layers, batch, ...)
|
||||
temporal_states: torch.Tensor, # (num_linear_layers, batch, ...)
|
||||
) -> torch.Tensor:
|
||||
global _enginex_profile_mode
|
||||
mode = "prefill" if attn_metadata.num_prefill_tokens > 0 else "decode"
|
||||
_enginex_profile_mode = mode
|
||||
hidden_states = self.embed_tokens(input_ids)
|
||||
residual = None
|
||||
|
||||
@@ -1088,7 +1119,6 @@ class Qwen3_5Model(nn.Module):
|
||||
attn_idx += 1
|
||||
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
_enginex_profile_log(mode)
|
||||
return hidden_states
|
||||
|
||||
|
||||
@@ -1188,6 +1218,7 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
_prof(f"ForCausalLM.forward START input_ids={tuple(input_ids.shape)} prefill={getattr(attn_metadata,'num_prefill_tokens',0)}")
|
||||
if self.mamba_cache is None:
|
||||
if self.scheduler_config is not None:
|
||||
max_batch_size = _get_graph_batch_size(
|
||||
@@ -1270,6 +1301,7 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
|
||||
self._gdn_prefix_cache.popitem(last=False)
|
||||
# ── End save ────────────────────────────────────────────────────────────
|
||||
|
||||
_prof(f"ForCausalLM.forward END hidden={tuple(hidden_states.shape)}")
|
||||
return hidden_states
|
||||
|
||||
def compute_logits(
|
||||
|
||||
Reference in New Issue
Block a user