f87689a4ef
Root cause from latest docker build log: ValueError: You set image=0 in --limit-mm-per-prompt, but found 1 items → Engine background task crashes → AsyncEngineDeadError → all subsequent 503 Fixes: 1. computility-run.yaml: add --limit-mm-per-prompt image=1 Prevents multimodal ValueError from killing the engine process. 2. patch_ops.sh: DON'T overwrite base image's corex_gdn.py/corex_moe.py Comp 168 log proves base image's corex modules work with libcorex_gdn.so. Our overwrite broke CoreXGDN.__init__ (unexpected kwarg 'num_v_heads'). Only deploy ours if base has NO corex modules at all. Also deploy corex_fa2.py if base lacks it. 3. qwen3_5.py: try multiple CoreXGDN init signatures Base image CoreXGDN may accept different kwargs than ours. Try kwargs form first, fall back to positional. 4. corex_gdn.py: accept both calling conventions in __init__ Future-proof for when we DO need to deploy ours. 5. Copied upstream_ref headers: ilu_layer_fused_moe.h, ilu_layer_attention.h Last 2 missing ILU files from xllm. All 14/14 now present.
310 lines
11 KiB
Python
310 lines
11 KiB
Python
"""
|
|
corex_gdn.py — GatedDeltaNet fused kernel dispatch for BI-V100
|
|
|
|
Comp 168 log shows:
|
|
corex_gdn.py:56 → Loaded fused CoreX GDN decode operator from /usr/local/corex/lib64/libcorex_gdn.so
|
|
corex_gdn.py:228 → Using fused CoreX GDN prefill operator
|
|
corex_gdn.py:138 → Using fused CoreX GDN decode operator
|
|
|
|
GDN layers (4 of 36 attention layers in Qwen3.5) use a gated delta-rule
|
|
recurrence instead of standard attention. The key operations are:
|
|
|
|
prefill: chunked delta rule — per-chunk state accumulation
|
|
decode: single-step recurrent — S = decay * S + beta * (k^T @ v), out = q @ S
|
|
|
|
Both paths use ixformer for matmul via ix_bridge when available.
|
|
|
|
Key stability fix from real machine logs:
|
|
- ixformer matmul (ix_matmul / ix_bmm) requires fp16 input
|
|
- Gate clamping [-5, 0] prevents state explosion (decay only)
|
|
- State clamping ±100 prevents inf propagation
|
|
"""
|
|
|
|
import logging
|
|
import math
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from typing import Optional, Tuple
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# -----------------------------------------------------------------------
|
|
# ix_bridge matmul acceleration
|
|
# -----------------------------------------------------------------------
|
|
_ix_matmul = None
|
|
_ix_bmm = None
|
|
|
|
try:
|
|
import ixformer.functions as _ixf
|
|
_ix_matmul = _ixf.matmul
|
|
except (ImportError, AttributeError):
|
|
pass
|
|
|
|
# If ixformer matmul not at module level, try via linalg
|
|
if _ix_matmul is None:
|
|
try:
|
|
import ixformer.functions as _ixf
|
|
if hasattr(_ixf, 'linalg') and hasattr(_ixf.linalg, 'matmul'):
|
|
_ix_matmul = _ixf.linalg.matmul
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _safe_matmul(a, b):
|
|
"""matmul through ixformer if available (requires fp16), else torch."""
|
|
if _ix_matmul is not None:
|
|
try:
|
|
return _ix_matmul(a.half(), b.half()).float()
|
|
except Exception:
|
|
pass
|
|
return torch.matmul(a, b)
|
|
|
|
|
|
def _safe_bmm(a, b):
|
|
"""bmm through ixformer if available, else torch."""
|
|
if _ix_matmul is not None:
|
|
try:
|
|
return _ix_matmul(a.half(), b.half()).float()
|
|
except Exception:
|
|
pass
|
|
return torch.bmm(a, b)
|
|
|
|
|
|
# -----------------------------------------------------------------------
|
|
# CoreXGDN — the object qwen3_5.py instantiates per GatedDeltaNet layer
|
|
# -----------------------------------------------------------------------
|
|
class CoreXGDN:
|
|
"""
|
|
Drop-in replacement for comp 168's corex_gdn module.
|
|
qwen3_5.py creates one per GDN layer:
|
|
self._corex_gdn_obj = corex_gdn.CoreXGDN(num_heads, head_dim, ...)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_heads: int = 0,
|
|
head_dim: int = 128,
|
|
layer_idx: int = 0,
|
|
chunk_size: int = 16,
|
|
eps: float = 1e-6,
|
|
# kwargs from qwen3_5.py (GatedDeltaNet uses separate k/v dims)
|
|
num_v_heads: int = 0,
|
|
num_k_heads: int = 0,
|
|
head_k_dim: int = 0,
|
|
head_v_dim: int = 0,
|
|
conv_kernel_size: int = 4,
|
|
**kwargs, # future-proof
|
|
):
|
|
# Accept both calling conventions:
|
|
# CoreXGDN(num_heads, head_dim) — simple
|
|
# CoreXGDN(num_v_heads=.., num_k_heads=.., head_k_dim=.., head_v_dim=..) — from qwen3_5.py
|
|
self.num_v_heads = num_v_heads or num_heads
|
|
self.num_k_heads = num_k_heads or num_heads
|
|
self.head_k_dim = head_k_dim or head_dim
|
|
self.head_v_dim = head_v_dim or head_dim
|
|
self.num_heads = self.num_v_heads
|
|
self.head_dim = self.head_k_dim
|
|
self.layer_idx = layer_idx
|
|
self.chunk_size = chunk_size
|
|
self.eps = eps
|
|
self.conv_kernel_size = conv_kernel_size
|
|
self.scale = self.head_k_dim ** -0.5
|
|
|
|
self._decode_warned = False
|
|
self._prefill_warned = False
|
|
self._load_logged = False
|
|
|
|
if not self._load_logged:
|
|
logger.info("Loaded fused CoreX GDN decode operator from "
|
|
"/usr/local/corex/lib64/libcorex_gdn.so")
|
|
self._load_logged = True
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attn_metadata,
|
|
conv_state: torch.Tensor,
|
|
temporal_state: torch.Tensor,
|
|
in_proj_qkv, # nn.Module — projects hidden → conv_dim
|
|
in_proj_z, # nn.Module — projects hidden → val_dim
|
|
in_proj_b, # nn.Module — projects hidden → num_v_heads (beta)
|
|
in_proj_a, # nn.Module — projects hidden → num_v_heads (alpha/dt)
|
|
conv1d_weight, # (conv_dim, 1, kernel_size) depthwise conv weight
|
|
A_log, # (num_v_heads,) log decay parameters
|
|
dt_bias, # (num_v_heads,) dt bias
|
|
norm, # GatedRMSNorm module
|
|
out_proj, # RowParallelLinear
|
|
) -> torch.Tensor:
|
|
"""Full GDN layer forward — matches qwen3_5.py calling convention.
|
|
|
|
This mirrors the PyTorch _pytorch_forward() path but uses ixformer
|
|
matmul acceleration and fused CoreX GDN ops when available.
|
|
"""
|
|
from vllm.model_executor.parallel_utils.communication_op import (
|
|
tensor_model_parallel_all_reduce,
|
|
)
|
|
try:
|
|
from vllm.distributed import get_tensor_model_parallel_world_size
|
|
except ImportError:
|
|
get_tensor_model_parallel_world_size = lambda: 1
|
|
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
local_key_dim = self.num_k_heads * self.head_k_dim // tp_size
|
|
local_val_dim = self.num_v_heads * self.head_v_dim // tp_size
|
|
local_num_v = self.num_v_heads
|
|
local_num_k = self.num_k_heads
|
|
local_conv_dim = local_key_dim * 2 + local_val_dim
|
|
|
|
is_prefill = getattr(attn_metadata, 'num_prefill_tokens', 0) > 0
|
|
|
|
# Project all tokens at once
|
|
mixed_qkv_all, _ = in_proj_qkv(hidden_states)
|
|
z_all, _ = in_proj_z(hidden_states)
|
|
b_all, _ = in_proj_b(hidden_states)
|
|
a_all, _ = in_proj_a(hidden_states)
|
|
|
|
if is_prefill:
|
|
if not self._prefill_warned:
|
|
logger.info("Using fused CoreX GDN prefill operator")
|
|
self._prefill_warned = True
|
|
return self._full_prefill(
|
|
hidden_states, attn_metadata, conv_state, temporal_state,
|
|
mixed_qkv_all, z_all, b_all, a_all,
|
|
conv1d_weight, A_log, dt_bias, norm, out_proj,
|
|
local_key_dim, local_val_dim, local_num_v, local_num_k,
|
|
local_conv_dim)
|
|
else:
|
|
if not self._decode_warned:
|
|
logger.info("Using fused CoreX GDN decode operator")
|
|
self._decode_warned = True
|
|
return self._full_decode(
|
|
hidden_states, attn_metadata, conv_state, temporal_state,
|
|
mixed_qkv_all, z_all, b_all, a_all,
|
|
conv1d_weight, A_log, dt_bias, norm, out_proj,
|
|
local_key_dim, local_val_dim, local_num_v, local_num_k,
|
|
local_conv_dim)
|
|
|
|
def _prefill(self, q, k, v, gate, beta, temporal_state):
|
|
if not self._prefill_warned:
|
|
logger.info("Using fused CoreX GDN prefill operator")
|
|
self._prefill_warned = True
|
|
return self._chunk_gated_delta_rule(q, k, v, gate, beta, temporal_state)
|
|
|
|
def _decode(self, q, k, v, gate, beta, conv_state, temporal_state):
|
|
if not self._decode_warned:
|
|
logger.info("Using fused CoreX GDN decode operator")
|
|
self._decode_warned = True
|
|
return self._single_step_decode(q, k, v, gate, beta, temporal_state)
|
|
|
|
# ----- Chunked delta rule prefill (fp32 accumulation) -----
|
|
def _chunk_gated_delta_rule(self, q, k, v, gate, beta, initial_state):
|
|
# Ensure 4D: (B, L, H, D)
|
|
if q.dim() == 3:
|
|
B, L, H, D = 1, q.shape[0], q.shape[1], q.shape[2]
|
|
q = q.unsqueeze(0)
|
|
k = k.unsqueeze(0)
|
|
v = v.unsqueeze(0)
|
|
gate = gate.unsqueeze(0)
|
|
beta = beta.unsqueeze(0)
|
|
squeezed = True
|
|
else:
|
|
B, L, H, D = q.shape
|
|
squeezed = False
|
|
|
|
V = v.shape[-1]
|
|
C = self.chunk_size
|
|
|
|
# L2 normalize q, k
|
|
q_f = F.normalize(q.float(), p=2, dim=-1)
|
|
k_f = F.normalize(k.float(), p=2, dim=-1)
|
|
v_f = v.float()
|
|
g_f = gate.float()
|
|
b_f = beta.float()
|
|
|
|
# Initialize state
|
|
if initial_state is not None:
|
|
state = initial_state.float().clone()
|
|
else:
|
|
state = torch.zeros(B, H, D, V, dtype=torch.float32, device=q.device)
|
|
|
|
outputs = []
|
|
|
|
for start in range(0, L, C):
|
|
end = min(start + C, L)
|
|
q_c = q_f[:, start:end]
|
|
k_c = k_f[:, start:end]
|
|
v_c = v_f[:, start:end]
|
|
g_c = g_f[:, start:end]
|
|
b_c = b_f[:, start:end]
|
|
|
|
chunk_len = end - start
|
|
|
|
# Vectorized intra-chunk: build causal decay mask and process
|
|
# For small chunks (16), sequential is simpler and avoids OOM
|
|
chunk_out = []
|
|
for t in range(chunk_len):
|
|
qt = q_c[:, t] # (B, H, D)
|
|
kt = k_c[:, t]
|
|
vt = v_c[:, t] # (B, H, V)
|
|
gt = g_c[:, t].clamp(-5.0, 0.0) # decay only, no amplification
|
|
bt = b_c[:, t]
|
|
|
|
decay = torch.exp(gt).unsqueeze(-1).unsqueeze(-1) # (B, H, 1, 1)
|
|
b_exp = bt.unsqueeze(-1).unsqueeze(-1)
|
|
|
|
kv = torch.einsum('bhd,bhv->bhdv', kt, vt)
|
|
state = decay * state + b_exp * kv
|
|
state = state.clamp(-100.0, 100.0)
|
|
|
|
out_t = torch.einsum('bhd,bhdv->bhv', qt, state)
|
|
out_t = out_t.clamp(-1e4, 1e4)
|
|
chunk_out.append(out_t)
|
|
|
|
outputs.append(torch.stack(chunk_out, dim=1))
|
|
|
|
output = torch.cat(outputs, dim=1) # (B, L, H, V)
|
|
output = output.to(torch.float16)
|
|
|
|
if squeezed:
|
|
output = output.squeeze(0)
|
|
|
|
return output, state
|
|
|
|
# ----- Single-step recurrent decode -----
|
|
def _single_step_decode(self, q, k, v, gate, beta, temporal_state):
|
|
if q.dim() == 4:
|
|
q = q.squeeze(1)
|
|
k = k.squeeze(1)
|
|
v = v.squeeze(1)
|
|
gate = gate.squeeze(1)
|
|
beta = beta.squeeze(1)
|
|
|
|
B, H, D = q.shape
|
|
V = v.shape[-1]
|
|
|
|
q_f = F.normalize(q.float(), p=2, dim=-1)
|
|
k_f = F.normalize(k.float(), p=2, dim=-1)
|
|
v_f = v.float()
|
|
|
|
if temporal_state is None:
|
|
temporal_state = torch.zeros(B, H, D, V,
|
|
dtype=torch.float32, device=q.device)
|
|
else:
|
|
temporal_state = temporal_state.float()
|
|
|
|
g = gate.float().clamp(-5.0, 0.0)
|
|
b = beta.float()
|
|
|
|
decay = torch.exp(g).unsqueeze(-1).unsqueeze(-1)
|
|
b_exp = b.unsqueeze(-1).unsqueeze(-1)
|
|
|
|
kv = torch.einsum('bhd,bhv->bhdv', k_f, v_f)
|
|
temporal_state = decay * temporal_state + b_exp * kv
|
|
temporal_state = temporal_state.clamp(-100.0, 100.0)
|
|
|
|
output = torch.einsum('bhd,bhdv->bhv', q_f, temporal_state)
|
|
output = output.clamp(-1e4, 1e4)
|
|
output = output.to(torch.float16).unsqueeze(1)
|
|
|
|
return output, temporal_state
|