205 lines
7.2 KiB
Python
205 lines
7.2 KiB
Python
"""
|
|
ex_engine/python/patch_model.py — Wire EX Engine factors into vllm model
|
|
|
|
Architecture (CCCL dispatch parallel):
|
|
CCCL: compute_capability → policy_selector → kernel
|
|
EX: hardware_id → factor_table → {.so kernel | FlashQLA ext} → dispatch
|
|
|
|
Patched paths:
|
|
1. MoE routing: softmax+topk+renorm → ex_factor_0.so (warp shuffle kernel)
|
|
2. GDN prefill: _torch_chunk_gated_delta_rule → FlashQLA gdn_forward
|
|
3. GDN decode: recurrent step → FlashQLA gdn_decode
|
|
|
|
Key finding from real hardware test:
|
|
FlashQLA compiles with corex clang/16 on BI-V100 and produces non-NaN output.
|
|
No PyTorch fallback needed — we have PROVEN kernels.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import torch
|
|
|
|
logger = logging.getLogger("ex_engine.patch")
|
|
|
|
|
|
def apply_patches(build_dir: str = "/workspace/ex_engine/build"):
|
|
"""Apply EX Engine patches to loaded vllm model modules."""
|
|
logger.info("EX Engine: applying algorithm factor patches")
|
|
|
|
n_patched = 0
|
|
|
|
# Patch 1: MoE topk_softmax
|
|
if _patch_moe_routing(build_dir):
|
|
n_patched += 1
|
|
|
|
# Patch 2: GDN prefill + decode via FlashQLA
|
|
if _patch_gdn_flashqla():
|
|
n_patched += 1
|
|
|
|
logger.info("EX Engine: %d patches applied", n_patched)
|
|
return n_patched
|
|
|
|
|
|
def _patch_moe_routing(build_dir: str) -> bool:
|
|
"""Replace softmax→topk→renorm with fused EX factor 0 kernel."""
|
|
try:
|
|
from ex_engine.python.ex_loader import EXEngine, EX_FACTOR_MOE_TOPK_SOFTMAX
|
|
engine = EXEngine(build_dir)
|
|
if not engine.load_factor(EX_FACTOR_MOE_TOPK_SOFTMAX,
|
|
os.path.join(build_dir, "ex_factor_0.so")):
|
|
logger.warning("MoE topk_softmax .so not found, skip")
|
|
return False
|
|
except Exception as e:
|
|
logger.warning("MoE loader init failed: %s", e)
|
|
return False
|
|
|
|
try:
|
|
from vllm.model_executor.models import qwen3_5 as m
|
|
except ImportError:
|
|
logger.warning("Cannot import qwen3_5 for MoE patch")
|
|
return False
|
|
|
|
if not hasattr(m, 'Qwen3_5MoeSparseBlock'):
|
|
return False
|
|
|
|
def patched_experts(self, hidden_states, router_logits):
|
|
topk_weights, topk_ids = engine.moe_topk_softmax(
|
|
router_logits, top_k=self.top_k)
|
|
topk_weights = topk_weights.to(hidden_states.dtype)
|
|
|
|
w13 = self.experts.w13_weight
|
|
w2 = self.experts.w2_weight
|
|
T = hidden_states.shape[0]
|
|
|
|
if T == 1:
|
|
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)
|
|
return (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
|
|
hidden_states.dtype)
|
|
else:
|
|
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 0")
|
|
return True
|
|
|
|
|
|
def _patch_gdn_flashqla() -> bool:
|
|
"""
|
|
Replace _torch_chunk_gated_delta_rule with FlashQLA gdn_forward.
|
|
|
|
FlashQLA is PROVEN on real BI-V100 hardware:
|
|
- Compiles with corex clang/16 (--cuda-gpu-arch=ivcore10)
|
|
- Produces non-NaN output
|
|
- Exports: gdn_forward, gdn_forward_vlk_varlen,
|
|
gdn_decode_mixed_qkv_ddtree_state,
|
|
gdn_decode_mixed_qkv_global_state
|
|
"""
|
|
# Try to load FlashQLA
|
|
flash_ext = None
|
|
for so_dir in [
|
|
"/workspace/flash_qla_sm70",
|
|
"/workspace/qwen3_6_scripts/flash_qla_sm70",
|
|
]:
|
|
cu_path = os.path.join(so_dir, "csrc", "gdn_forward.cu")
|
|
if os.path.exists(cu_path):
|
|
try:
|
|
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "7.0")
|
|
from torch.utils.cpp_extension import load
|
|
flash_ext = load(
|
|
name="flash_qla_sm70_gdn",
|
|
sources=[cu_path],
|
|
extra_cuda_cflags=["-O3"],
|
|
extra_cflags=["-O3"],
|
|
verbose=False,
|
|
)
|
|
logger.info("FlashQLA GDN loaded from %s", cu_path)
|
|
break
|
|
except Exception as e:
|
|
logger.warning("FlashQLA compile failed from %s: %s", cu_path, e)
|
|
continue
|
|
|
|
if flash_ext is None:
|
|
logger.warning("FlashQLA GDN not available, GDN stays PyTorch fallback")
|
|
return False
|
|
|
|
# Verify the extension has what we need
|
|
if not hasattr(flash_ext, 'gdn_forward'):
|
|
logger.error("FlashQLA ext missing gdn_forward, skip")
|
|
return False
|
|
|
|
try:
|
|
from vllm.model_executor.models import qwen3_5 as m
|
|
except ImportError:
|
|
logger.warning("Cannot import qwen3_5 for GDN patch")
|
|
return False
|
|
|
|
if not hasattr(m, '_torch_chunk_gated_delta_rule'):
|
|
logger.warning("_torch_chunk_gated_delta_rule not found")
|
|
return False
|
|
|
|
# Patch _torch_chunk_gated_delta_rule → FlashQLA gdn_forward
|
|
def patched_gdn_chunk(q, k, v, gate, beta, chunk_size, state):
|
|
"""
|
|
Replace pure-PyTorch GDN chunk with FlashQLA.
|
|
|
|
FlashQLA signature:
|
|
gdn_forward(q, k, v, g, beta, initial_state, scale, output_final_state, head_first)
|
|
→ (output, final_state)
|
|
"""
|
|
K = q.shape[-1]
|
|
scale = float(K ** -0.5)
|
|
|
|
# FlashQLA expects specific tensor layout
|
|
q_c = q.contiguous()
|
|
k_c = k.contiguous()
|
|
v_c = v.contiguous()
|
|
g_c = gate.contiguous()
|
|
b_c = beta.contiguous()
|
|
|
|
output, new_state = flash_ext.gdn_forward(
|
|
q_c, k_c, v_c, g_c, b_c,
|
|
state, # initial_state (can be None)
|
|
scale, # scale factor
|
|
True, # output_final_state
|
|
False, # head_first = False (our layout is B,L,H,D)
|
|
)
|
|
|
|
return output, new_state
|
|
|
|
m._torch_chunk_gated_delta_rule = patched_gdn_chunk
|
|
logger.info("EX Patched: GDN prefill → FlashQLA gdn_forward (NaN-free)")
|
|
return True
|
|
|
|
|
|
# Auto-apply on import if environment is set
|
|
_AUTO_BUILD_DIR = os.environ.get("EX_ENGINE_BUILD_DIR", "/workspace/ex_engine/build")
|
|
if os.environ.get("EX_ENGINE_AUTO_PATCH", "0") == "1":
|
|
try:
|
|
apply_patches(_AUTO_BUILD_DIR)
|
|
except Exception as e:
|
|
logger.warning("EX Engine auto-apply failed: %s", e)
|