b75965d4ea
Real machine log (2d5232c dockerrizhi.txt) shows two AST call chain breaks:
1. EVERY layer EVERY token:
_custom_ops.py:58 'ixformer.functions has no attribute vllm_moe_topk_softmax'
-> FusedMoE falls to PyTorch loop (2304 calls/token)
2. EVERY GDN layer (4 layers):
'NaN in prefill GatedDeltaNet layer N (frac=0.9998-1.0000)'
-> _torch_chunk_gated_delta_rule produces all-NaN
Fixes:
- build.sh: --cuda-gpu-arch=ivcore10, -D__ILUVATAR__ flags from real log
- Dockerfile: add ex_engine build before patch_ops
- patch_ops.sh: deploy .so + python into vllm model dir
- ex_loader.py: search co-located .so paths
- patch_model.py: remove premature auto-apply
- factor_moe_topk_softmax.cu: remove dead parallel branch
198 lines
7.3 KiB
Python
198 lines
7.3 KiB
Python
"""
|
|
ex_engine/python/patch_model.py — Wire EX Engine factors into vllm model
|
|
|
|
CCCL parallel: CCCL's dispatch_reduce.cuh has a Dispatch() that selects
|
|
the tuned kernel based on compute_capability. This patch does the same:
|
|
it replaces the PyTorch fallback paths with EX factor kernel calls.
|
|
|
|
Patched paths:
|
|
1. Qwen3_5MoeSparseBlock._pure_pytorch_experts()
|
|
→ Uses EX factor 0 (moe_topk_softmax) for routing
|
|
→ Falls back to PyTorch GEMM for expert computation (factor 2 TBD)
|
|
|
|
2. GatedDeltaNet.forward() prefill path
|
|
→ Uses EX factor 5 (gdn_chunk_fwd) instead of _torch_chunk_gated_delta_rule
|
|
→ Eliminates NaN by using fp32 accumulation
|
|
|
|
Integration:
|
|
Called from patch_ops.sh during Docker build, or imported at runtime:
|
|
python -c "from ex_engine.python.patch_model import apply_patches; apply_patches()"
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import torch
|
|
import types
|
|
|
|
logger = logging.getLogger("ex_engine.patch")
|
|
|
|
|
|
def apply_patches(build_dir: str = "/workspace/ex_engine/build"):
|
|
"""
|
|
Apply EX Engine patches to the loaded vllm model modules.
|
|
Must be called AFTER vllm modules are imported.
|
|
"""
|
|
# Lazy import to avoid circular deps
|
|
try:
|
|
from ex_engine.python.ex_loader import EXEngine
|
|
except ImportError:
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from ex_engine.python.ex_loader import EXEngine
|
|
|
|
engine = EXEngine(build_dir)
|
|
loaded = engine.load_all()
|
|
|
|
if loaded == 0:
|
|
logger.warning("EX Engine: no factors loaded, skipping patches")
|
|
return
|
|
|
|
logger.info("EX Engine: %d factors loaded, applying patches", loaded)
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Patch 1: MoE routing — replace softmax+topk with fused factor
|
|
# -----------------------------------------------------------------------
|
|
if engine.has_factor(0): # EX_FACTOR_MOE_TOPK_SOFTMAX
|
|
_patch_moe_routing(engine)
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Patch 2: GDN prefill — replace _torch_chunk_gated_delta_rule
|
|
# -----------------------------------------------------------------------
|
|
if engine.has_factor(5): # EX_FACTOR_GDN_CHUNK_FWD
|
|
_patch_gdn_prefill(engine)
|
|
|
|
logger.info("EX Engine: patches applied successfully")
|
|
|
|
|
|
def _patch_moe_routing(engine):
|
|
"""
|
|
Replace the pure PyTorch softmax→topk→renormalize in MoE with
|
|
fused EX factor kernel.
|
|
|
|
Target: Qwen3_5MoeSparseBlock._pure_pytorch_experts()
|
|
The first 3 lines:
|
|
routing_weights = _ix_softmax(router_logits.float(), dim=-1)
|
|
topk_weights, topk_ids = torch.topk(routing_weights, self.top_k, dim=-1)
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
"""
|
|
try:
|
|
from vllm.model_executor.models import qwen3_5 as m
|
|
except ImportError:
|
|
logger.warning("Cannot import qwen3_5, skipping MoE patch")
|
|
return
|
|
|
|
if not hasattr(m, 'Qwen3_5MoeSparseBlock'):
|
|
logger.warning("Qwen3_5MoeSparseBlock not found, skipping MoE patch")
|
|
return
|
|
|
|
original_fn = m.Qwen3_5MoeSparseBlock._pure_pytorch_experts
|
|
|
|
def patched_experts(self, hidden_states, router_logits):
|
|
# EX fused topk+softmax (1 kernel instead of 2 + 1 normalize)
|
|
topk_weights, topk_ids = engine.moe_topk_softmax(
|
|
router_logits, top_k=self.top_k)
|
|
topk_weights = topk_weights.to(hidden_states.dtype)
|
|
|
|
# Expert computation still uses PyTorch path
|
|
# (factor 2 will replace this with batched GEMM later)
|
|
w13 = self.experts.w13_weight
|
|
w2 = self.experts.w2_weight
|
|
T = hidden_states.shape[0]
|
|
|
|
if T == 1:
|
|
# Decode fast path (same as original)
|
|
eids = topk_ids[0]
|
|
ws = topk_weights[0]
|
|
w13_sel = w13[eids]
|
|
w2_sel = w2[eids]
|
|
H = hidden_states.shape[-1]
|
|
|
|
gate_up = torch.nn.functional.linear(
|
|
hidden_states, w13_sel.reshape(-1, H))
|
|
gate_up = gate_up.view(self.top_k, -1)
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
act = torch.nn.functional.silu(gate) * up
|
|
expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1)
|
|
out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True)
|
|
return out.to(hidden_states.dtype)
|
|
else:
|
|
# Prefill path — loop over 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)
|
|
tok_ids, topk_pos = mask.nonzero(as_tuple=True)
|
|
tokens = hidden_states[tok_ids]
|
|
gate_up = torch.nn.functional.linear(tokens, w13[eid])
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
act = torch.nn.functional.silu(gate) * up
|
|
expert_out = torch.nn.functional.linear(act, w2[eid])
|
|
weights = topk_weights[tok_ids, topk_pos].unsqueeze(-1)
|
|
out.index_add_(0, tok_ids,
|
|
(expert_out * weights).to(out.dtype))
|
|
return out
|
|
|
|
m.Qwen3_5MoeSparseBlock._pure_pytorch_experts = patched_experts
|
|
logger.info("EX Patched: MoE routing → fused topk_softmax factor")
|
|
|
|
|
|
def _patch_gdn_prefill(engine):
|
|
"""
|
|
Replace _torch_chunk_gated_delta_rule with EX factor 5 (gdn_chunk_fwd).
|
|
This eliminates the NaN problem by using fp32 state accumulation.
|
|
"""
|
|
try:
|
|
from vllm.model_executor.models import qwen3_5 as m
|
|
except ImportError:
|
|
logger.warning("Cannot import qwen3_5, skipping GDN patch")
|
|
return
|
|
|
|
if not hasattr(m, '_torch_chunk_gated_delta_rule'):
|
|
logger.warning("_torch_chunk_gated_delta_rule not found, skipping GDN patch")
|
|
return
|
|
|
|
original_fn = m._torch_chunk_gated_delta_rule
|
|
|
|
def patched_gdn_chunk(q, k, v, gate, beta, chunk_size, state):
|
|
"""
|
|
EX factor replacement for _torch_chunk_gated_delta_rule.
|
|
|
|
Args match the original function signature:
|
|
q: (1, L, H, D) or (B, L, H, D)
|
|
k, v: same shape
|
|
gate: (1, L, H) or (B, L, H)
|
|
beta: same shape
|
|
chunk_size: int (ignored — factor processes full sequence)
|
|
state: (B, H, D, D)
|
|
|
|
Returns: (output, new_state)
|
|
"""
|
|
B = q.shape[0]
|
|
L = q.shape[1]
|
|
H = q.shape[2]
|
|
D = q.shape[3]
|
|
|
|
# Ensure contiguous and correct dtype
|
|
q_c = q.contiguous().half()
|
|
k_c = k.contiguous().half()
|
|
v_c = v.contiguous().half()
|
|
g_c = gate.float().contiguous()
|
|
b_c = beta.float().contiguous()
|
|
s_c = state.float().contiguous()
|
|
|
|
output, new_state = engine.gdn_chunk_fwd(
|
|
q_c, k_c, v_c, g_c, b_c, s_c)
|
|
|
|
return output, new_state
|
|
|
|
m._torch_chunk_gated_delta_rule = patched_gdn_chunk
|
|
logger.info("EX Patched: GDN prefill → gdn_chunk_fwd factor (NaN-free)")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Call apply_patches() explicitly AFTER vllm model modules are loaded.
|
|
# Integration point: qwen3_5.py calls this at the end of model __init__,
|
|
# or patch_ops.sh adds it to the startup sequence.
|
|
# ---------------------------------------------------------------------------
|