icer/qwen36_01
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

profile decode path and optimize tiny batch moe

Browse Source
This commit is contained in:
the magician
2026-07-14 20:44:05 +08:00
parent a5abd66e21
commit 0597fa6c6d
10 changed files with 1633 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

268
qwen3_6_scripts/qwen3_5.py
View File

@@ -3,6 +3,9 @@
# 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
@@ -41,6 +44,59 @@ 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_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))
# ---------------------------------------------------------------------------
# Pure-PyTorch DeltaNet kernels (fallbacks from transformers 5.2.0)
@@ -613,44 +669,48 @@ class Qwen3_5FullAttention(nn.Module):
total_tokens = hidden_states.shape[0]
# q_proj output includes gate (dim doubled)
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)
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)
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
q = self.q_norm.forward_cuda(
q.view(total_tokens, self.local_num_heads, self.head_dim)
.contiguous()).view(total_tokens, -1)
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)
# 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)
attn_out = self.attn(q, k, v, kv_cache, attn_metadata)
with _enginex_profile("full_attn.paged_attention"):
attn_out = self.attn(q, k, v, kv_cache, attn_metadata)
# Multiply by sigmoid gate before output projection
attn_out = attn_out * torch.sigmoid(gate.float()).to(attn_out.dtype)
output, _ = self.o_proj(attn_out)
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)
return output
@@ -754,12 +814,13 @@ class Qwen3_5MoeSparseBlock(nn.Module):
Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False.
"""
# 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)
# 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)
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)
@@ -771,59 +832,90 @@ 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).
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)
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)
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)
out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
hidden_states.dtype) # (1, H)
elif T <= 2:
# Fast path: tiny decode batch. With max_num_seqs=2, normal decode
# often has T=2 and should not fall back to the Python per-expert
# loop used for long prefill chunks.
with _enginex_profile("moe.routed_tiny_batch_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)
else:
# General path (prefill / multi-seq): loop over unique active experts.
# At most T*top_k unique experts, always <= num_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))
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))
return out # partial, all-reduce done in forward()
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)
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)
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.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)
out = routed_out + shared_out
with _enginex_profile("moe.combine"):
out = routed_out + shared_out
if self.experts.tp_size > 1:
out = tensor_model_parallel_all_reduce(out)
with _enginex_profile("moe.tp_all_reduce"):
out = tensor_model_parallel_all_reduce(out)
return out
@@ -883,21 +975,27 @@ class Qwen3_5DecoderLayer(nn.Module):
) -> Tuple[torch.Tensor, torch.Tensor]:
if residual is None:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
with _enginex_profile("layer.input_norm"):
hidden_states = self.input_layernorm(hidden_states)
else:
hidden_states, residual = self.input_layernorm(hidden_states, residual)
with _enginex_profile("layer.input_norm"):
hidden_states, residual = self.input_layernorm(hidden_states, residual)
if self.layer_type == "linear_attention":
hidden_states = self.linear_attn(
hidden_states, attn_metadata, conv_state, temporal_state)
with _enginex_profile("layer.linear_attention"):
hidden_states = self.linear_attn(
hidden_states, attn_metadata, conv_state, temporal_state)
else:
hidden_states = self.self_attn(
positions, hidden_states, kv_cache, attn_metadata)
with _enginex_profile("layer.full_attention"):
hidden_states = self.self_attn(
positions, hidden_states, kv_cache, attn_metadata)
hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual)
with _enginex_profile("layer.post_attn_norm"):
hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual)
hidden_states = self.mlp(hidden_states)
with _enginex_profile("layer.mlp"):
hidden_states = self.mlp(hidden_states)
return hidden_states, residual
@@ -934,6 +1032,9 @@ 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
@@ -961,6 +1062,7 @@ class Qwen3_5Model(nn.Module):
attn_idx += 1
hidden_states, _ = self.norm(hidden_states, residual)
_enginex_profile_log(mode)
return hidden_states