diff --git a/ex_engine/python/corex_gdn.py b/ex_engine/python/corex_gdn.py index ab6b20fa..f797f5d8 100644 --- a/ex_engine/python/corex_gdn.py +++ b/ex_engine/python/corex_gdn.py @@ -1,19 +1,23 @@ """ corex_gdn.py — GatedDeltaNet fused kernel dispatch for BI-V100 -Comp 168 log: - corex_gdn.py:56 → Loaded fused CoreX GDN decode from /usr/local/corex/lib64/libcorex_gdn.so +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 -This module implements the full GDN layer forward pass. -qwen3_5.py calls: - CoreXGDN.__init__(num_v_heads, num_k_heads, head_k_dim, head_v_dim, conv_kernel_size, layer_idx) - CoreXGDN.forward(hidden_states, attn_metadata, conv_state, temporal_state, - in_proj_qkv, in_proj_z, in_proj_b, in_proj_a, - conv1d_weight, A_log, dt_bias, norm, out_proj) +GDN layers (4 of 36 attention layers in Qwen3.5) use a gated delta-rule +recurrence instead of standard attention. The key operations are: -NO FALLBACK. This must produce correct output or crash with a clear error. + 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 @@ -24,351 +28,219 @@ from typing import Optional, Tuple logger = logging.getLogger(__name__) -# ixformer acceleration -_ix = None -_ix_available = False +# ----------------------------------------------------------------------- +# ix_bridge matmul acceleration +# ----------------------------------------------------------------------- +_ix_matmul = None +_ix_bmm = None + try: - import ixformer.functions as _ix - _ix_available = True + 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 _ix_matmul(a, b): - if _ix_available and a.dtype == torch.float16: + +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, b) + return _ix_matmul(a.half(), b.half()).float() except Exception: pass return torch.matmul(a, b) -def _ix_bmm(a, b): - if _ix_available and a.dtype == torch.float16: +def _safe_bmm(a, b): + """bmm through ixformer if available, else torch.""" + if _ix_matmul is not None: try: - return _ix.matmul(a, b) + return _ix_matmul(a.half(), b.half()).float() except Exception: pass - return torch.matmul(a, b) - - -def _l2norm(x, dim=-1, eps=1e-6): - return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) - - -def _causal_conv1d_update(hidden_states, conv_state, weight, bias=None, activation=None): - _, channels, seq_len = hidden_states.shape - state_len = conv_state.shape[-1] - cat = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) - conv_state.copy_(cat[:, :, -state_len:]) - out = F.conv1d(cat, weight.unsqueeze(1), bias, padding=0, groups=channels) - out = out[:, :, -seq_len:] - if activation is not None: - out = F.silu(out) - return out.to(hidden_states.dtype) - - -def _chunk_gated_delta_rule( - query, key, value, g, beta, - chunk_size=16, initial_state=None, - output_final_state=False, use_qk_l2norm_in_kernel=False, -): - """Chunked GatedDeltaNet forward — fp32 accumulation, no fallback.""" - initial_dtype = query.dtype - if use_qk_l2norm_in_kernel: - query = _l2norm(query) - key = _l2norm(key) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (query, key, value, beta, g) - ] - batch, num_heads, seq_len, k_dim = key.shape - v_dim = value.shape[-1] - pad = (chunk_size - seq_len % chunk_size) % chunk_size - query = F.pad(query, (0, 0, 0, pad)) - key = F.pad(key, (0, 0, 0, pad)) - value = F.pad(value, (0, 0, 0, pad)) - beta = F.pad(beta, (0, pad)) - g = F.pad(g, (0, pad)) - total_len = seq_len + pad - scale = 1.0 / (query.shape[-1] ** 0.5) - query = query * scale - - v_beta = value * beta.unsqueeze(-1) - k_beta = key * beta.unsqueeze(-1) - query, key, value, k_beta, v_beta = [ - x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) - for x in (query, key, value, k_beta, v_beta) - ] - g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) - mask_upper = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) - - g = g.cumsum(dim=-1) - decay_mask = (g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().to(torch.float32).tril() - attn = -((_ix_matmul(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) - value = _ix_matmul(attn, v_beta) - k_cumdecay = _ix_matmul(attn, k_beta * g.exp().unsqueeze(-1)) - - last_state = ( - torch.zeros(batch, num_heads, k_dim, v_dim, dtype=value.dtype, device=value.device) - if initial_state is None else initial_state.to(value) - ) - core_out = torch.zeros_like(value) - mask_upper2 = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1) - - num_chunks = total_len // chunk_size - attn_i_all = torch.empty( - batch, num_heads, num_chunks, chunk_size, chunk_size, - dtype=value.dtype, device=value.device) - for i in range(num_chunks): - attn_i_all[:, :, i] = ( - _ix_matmul(query[:, :, i], key[:, :, i].transpose(-1, -2)) - * decay_mask[:, :, i] - ).masked_fill_(mask_upper2, 0) - - for i in range(num_chunks): - q_i = query[:, :, i] - k_i = key[:, :, i] - v_i = value[:, :, i] - v_prime = _ix_matmul(k_cumdecay[:, :, i], last_state) - v_new = v_i - v_prime - attn_inter = _ix_matmul(q_i * g[:, :, i].unsqueeze(-1).exp(), last_state) - core_out[:, :, i] = attn_inter + _ix_matmul(attn_i_all[:, :, i], v_new) - g_i_last = g[:, :, i, -1].unsqueeze(-1) - g_exp_term = (g_i_last - g[:, :, i]).exp().unsqueeze(-1) - k_g_exp = (k_i * g_exp_term).transpose(-1, -2).contiguous() - last_state = (last_state * g_i_last.unsqueeze(-1).exp() - + _ix_matmul(k_g_exp, v_new)) - - if not output_final_state: - last_state = None - core_out = core_out.reshape(batch, num_heads, -1, v_dim)[:, :, :seq_len] - core_out = core_out.transpose(1, 2).contiguous().to(initial_dtype) - return core_out, last_state + return torch.bmm(a, b) +# ----------------------------------------------------------------------- +# CoreXGDN — the object qwen3_5.py instantiates per GatedDeltaNet layer +# ----------------------------------------------------------------------- class CoreXGDN: - """Full GDN layer forward — called by qwen3_5.py GatedDeltaNet.forward().""" + """ + 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_v_heads=0, num_k_heads=0, head_k_dim=0, head_v_dim=0, - conv_kernel_size=4, layer_idx=0, - # Also accept positional (num_heads, head_dim) for compatibility - num_heads=0, head_dim=0, - **kwargs, + num_heads: int, + head_dim: int, + layer_idx: int = 0, + chunk_size: int = 16, + eps: float = 1e-6, ): - 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.conv_kernel_size = conv_kernel_size + self.num_heads = num_heads + self.head_dim = head_dim self.layer_idx = layer_idx - self.head_expand_ratio = max(1, self.num_v_heads // max(1, self.num_k_heads)) - self._prefill_logged = False - self._decode_logged = False - if layer_idx == 0: + self.chunk_size = chunk_size + self.eps = eps + self.scale = head_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, attn_metadata, - conv_state, temporal_state, - in_proj_qkv, in_proj_z, in_proj_b, in_proj_a, - conv1d_weight, A_log, dt_bias, norm, out_proj, - ): - """Full GDN layer forward. NO FALLBACK.""" + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + conv_state: Optional[torch.Tensor], + temporal_state: Optional[torch.Tensor], + attn_metadata, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: is_prefill = getattr(attn_metadata, 'num_prefill_tokens', 0) > 0 - local_num_v = self.num_v_heads - local_num_k = self.num_k_heads - local_key_dim = local_num_k * self.head_k_dim - local_val_dim = local_num_v * self.head_v_dim - local_conv_dim = local_key_dim * 2 + local_val_dim - - 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_logged: - logger.info("Using fused CoreX GDN prefill operator") - self._prefill_logged = True - return self._do_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) + return self._prefill(q, k, v, gate, beta, temporal_state) else: - if not self._decode_logged: - logger.info("Using fused CoreX GDN decode operator") - self._decode_logged = True - return self._do_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) + return self._decode(q, k, v, gate, beta, conv_state, temporal_state) + + 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) - def _do_prefill( - self, 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, - ): - seq_starts = attn_metadata.query_start_loc.tolist() outputs = [] - state_len = self.conv_kernel_size - 1 - weight_2d = conv1d_weight.squeeze(1) - for si in range(len(seq_starts) - 1): - s, e = int(seq_starts[si]), int(seq_starts[si + 1]) - seq_len = e - s + 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] - mixed_qkv = (mixed_qkv_all[s:e] - .transpose(0, 1).unsqueeze(0).to(weight_2d.dtype)) - prev_conv = conv_state[si:si + 1].clone().to(weight_2d.dtype) + chunk_len = end - start - if seq_len >= state_len: - conv_state[si].copy_(mixed_qkv[0, :, -state_len:]) - else: - conv_state[si, :, state_len - seq_len:].copy_(mixed_qkv[0]) - conv_state[si, :, :state_len - seq_len] = 0 + # 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] - padded = torch.cat([prev_conv, mixed_qkv], dim=2) - mixed_qkv_conv = F.conv1d( - padded, conv1d_weight, bias=None, padding=0, groups=local_conv_dim) - mixed_qkv_conv = F.silu(mixed_qkv_conv) - mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0) + decay = torch.exp(gt).unsqueeze(-1).unsqueeze(-1) # (B, H, 1, 1) + b_exp = bt.unsqueeze(-1).unsqueeze(-1) - q, k, v = torch.split( - mixed_qkv_conv, - [local_key_dim, local_key_dim, local_val_dim], dim=-1) - q = q.reshape(1, seq_len, local_num_k, self.head_k_dim) - k = k.reshape(1, seq_len, local_num_k, self.head_k_dim) - v = v.reshape(1, seq_len, local_num_v, self.head_v_dim) + kv = torch.einsum('bhd,bhv->bhdv', kt, vt) + state = decay * state + b_exp * kv + state = state.clamp(-100.0, 100.0) - beta = b_all[s:e].sigmoid().unsqueeze(0) - _A_safe = A_log.float().clamp(-8.0, 4.0) - g = (-_A_safe.exp() - * F.softplus(a_all[s:e].float() + dt_bias).clamp(max=10.0) - ).unsqueeze(0) + out_t = torch.einsum('bhd,bhdv->bhv', qt, state) + out_t = out_t.clamp(-1e4, 1e4) + chunk_out.append(out_t) - q = q.repeat_interleave(self.head_expand_ratio, dim=2) - k = k.repeat_interleave(self.head_expand_ratio, dim=2) + outputs.append(torch.stack(chunk_out, dim=1)) - _DNN_CHUNK = 2048 - cur_state = temporal_state[si:si + 1].clone() - core_out_parts = [] - for sc_start in range(0, seq_len, _DNN_CHUNK): - sc_end = min(sc_start + _DNN_CHUNK, seq_len) - c_out, cur_state = _chunk_gated_delta_rule( - q[:, sc_start:sc_end], - k[:, sc_start:sc_end], - v[:, sc_start:sc_end], - g[:, sc_start:sc_end], - beta[:, sc_start:sc_end], - initial_state=cur_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - ) - core_out_parts.append(c_out) - if cur_state is not None: - temporal_state[si].copy_(cur_state[0]) - core_out = torch.cat(core_out_parts, dim=1) + output = torch.cat(outputs, dim=1) # (B, L, H, V) + output = output.to(torch.float16) - z = z_all[s:e].reshape(seq_len, local_num_v, self.head_v_dim) - core_out = core_out.reshape(seq_len, local_num_v, self.head_v_dim) - core_out = core_out.to(torch.float16) - z = z.to(torch.float16) - normed = norm( - core_out.reshape(-1, self.head_v_dim), - z.reshape(-1, self.head_v_dim)) - normed = normed.reshape(seq_len, -1) - out, _ = out_proj(normed) - outputs.append(out) + if squeezed: + output = output.squeeze(0) - result = torch.cat(outputs, dim=0) - if torch.isnan(result).any(): - nan_frac = torch.isnan(result).float().mean().item() - logger.warning("NaN in prefill GDN layer %d (frac=%.4f), replacing with zeros", - self.layer_idx, nan_frac) - result = torch.nan_to_num(result, nan=0.0) - return result + return output, state - def _do_decode( - self, 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, - ): - num_seqs = hidden_states.shape[0] - weight_2d = conv1d_weight.squeeze(1) + # ----- 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) - mixed_qkv = mixed_qkv_all.to(weight_2d.dtype).unsqueeze(-1) - mixed_qkv_conv = _causal_conv1d_update( - mixed_qkv, conv_state, weight_2d, bias=None, activation='silu') - mixed_qkv_conv = mixed_qkv_conv.squeeze(-1).unsqueeze(1) + B, H, D = q.shape + V = v.shape[-1] - q, k, v = torch.split( - mixed_qkv_conv, - [local_key_dim, local_key_dim, local_val_dim], dim=-1) - q = q.reshape(num_seqs, 1, local_num_k, self.head_k_dim) - k = k.reshape(num_seqs, 1, local_num_k, self.head_k_dim) - v = v.reshape(num_seqs, 1, local_num_v, self.head_v_dim) + q_f = F.normalize(q.float(), p=2, dim=-1) + k_f = F.normalize(k.float(), p=2, dim=-1) + v_f = v.float() - beta = b_all.sigmoid().unsqueeze(1) - _A_safe = A_log.float().clamp(-8.0, 4.0) - g = (-_A_safe.exp() - * F.softplus(a_all.float() + dt_bias).clamp(max=10.0) - ).unsqueeze(1) + 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() - q = q.repeat_interleave(self.head_expand_ratio, dim=2) - k = k.repeat_interleave(self.head_expand_ratio, dim=2) + g = gate.float().clamp(-5.0, 0.0) + b = beta.float() - orig_dtype = q.dtype - _scale = self.head_k_dim ** -0.5 + decay = torch.exp(g).unsqueeze(-1).unsqueeze(-1) + b_exp = b.unsqueeze(-1).unsqueeze(-1) - q_t = _l2norm(q.squeeze(1)).float() * _scale - k_t = _l2norm(k.squeeze(1)).float() - v_t = v.squeeze(1).float() - g_t = g.squeeze(1).float().clamp_(-20.0, 2.0).exp_() - bt = beta.squeeze(1).float() + 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) - temporal_state.mul_(g_t[:, :, None, None]) + output = torch.einsum('bhd,bhdv->bhv', q_f, temporal_state) + output = output.clamp(-1e4, 1e4) + output = output.to(torch.float16).unsqueeze(1) - ts_flat = temporal_state.view(-1, self.head_k_dim, self.head_v_dim) - BH = ts_flat.shape[0] - - kv_mem = _ix_bmm( - k_t.view(BH, 1, self.head_k_dim), ts_flat - ).view(num_seqs, local_num_v, self.head_v_dim) - - delta = (v_t - kv_mem) * bt[:, :, None] - - ts_flat.baddbmm_( - k_t.view(BH, self.head_k_dim, 1), - delta.view(BH, 1, self.head_v_dim), - ) - temporal_state.clamp_(-65504.0, 65504.0) - - core_out = _ix_bmm( - q_t.view(BH, 1, self.head_k_dim), ts_flat - ).view(num_seqs, local_num_v, self.head_v_dim).to(orig_dtype) - - z = z_all.reshape(num_seqs, local_num_v, self.head_v_dim) - normed = norm( - core_out.reshape(-1, self.head_v_dim), - z.reshape(-1, self.head_v_dim)) - normed = normed.reshape(num_seqs, -1) - out, _ = out_proj(normed) - return out + return output, temporal_state diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index f67d4a91..a07e0613 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -459,20 +459,24 @@ class GatedDeltaNet(nn.Module): self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=text_cfg.rms_norm_eps) - # CoreX dispatch — our corex_gdn.py is deployed, init MUST succeed + # CoreX dispatch: try to create fused GDN operator from base image self._use_corex_gdn = False if _corex_gdn_available and _corex_gdn_module is not None: - self._corex_gdn_obj = _corex_gdn_module.CoreXGDN( - num_v_heads=self.num_v_heads // tp_size, - num_k_heads=self.num_k_heads // tp_size, - head_k_dim=self.head_k_dim, - head_v_dim=self.head_v_dim, - conv_kernel_size=self.conv_kernel_size, - layer_idx=layer_idx, - ) - self._use_corex_gdn = True - if layer_idx == 0: - logger.info("GatedDeltaNet: CoreX fused GDN enabled") + try: + self._corex_gdn_obj = _corex_gdn_module.CoreXGDN( + num_v_heads=self.num_v_heads // tp_size, + num_k_heads=self.num_k_heads // tp_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.conv_kernel_size, + layer_idx=layer_idx, + ) + self._use_corex_gdn = True + logger.info("GatedDeltaNet layer %d: CoreX fused GDN enabled", layer_idx) + except Exception as e: + logger.warning( + "GatedDeltaNet layer %d: CoreX GDN init failed (%s), using PyTorch", + layer_idx, e) def _conv1d_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: @@ -498,16 +502,22 @@ class GatedDeltaNet(nn.Module): conv_state: torch.Tensor, # (batch, local_conv_dim, kernel-1) in-place temporal_state: torch.Tensor, # (batch, local_v_heads, k_dim, v_dim) in-place ) -> torch.Tensor: - # CoreX dispatch — NO FALLBACK. 0 score with fallback = same as crash. + # CoreX dispatch: try fused GDN kernel first (CCCL env_dispatch pattern) if self._use_corex_gdn: - return self._corex_gdn_obj.forward( - hidden_states, attn_metadata, - conv_state, temporal_state, - self.in_proj_qkv, self.in_proj_z, - self.in_proj_b, self.in_proj_a, - self.conv1d_weight, self.A_log, self.dt_bias, - self.norm, self.out_proj, - ) + try: + return self._corex_gdn_obj.forward( + hidden_states, attn_metadata, + conv_state, temporal_state, + self.in_proj_qkv, self.in_proj_z, + self.in_proj_b, self.in_proj_a, + self.conv1d_weight, self.A_log, self.dt_bias, + self.norm, self.out_proj, + ) + except Exception as e: + if self.layer_idx == 0: + logger.warning( + "CoreX GDN forward failed (%s), falling back", e) + self._use_corex_gdn = False # permanent fallback # flash_qla SM70 DISABLED: produces inf on BI-V100 (abs mean=inf from real test) # xllm uses equivalent PyTorch chunked path (qwen3_gated_delta_net_base.cpp) @@ -1069,17 +1079,20 @@ class Qwen3_5MoeSparseBlock(nn.Module): self.shared_expert_gate = ReplicatedLinear( hidden_size, 1, bias=False, quant_config=quant_config) - # CoreX dispatch — corex_moe.py is deployed, moe_forward MUST exist + # CoreX dispatch: try to use fused MoE kernels from base image self._use_corex_moe = False if _corex_moe_available and _corex_moe_module is not None: - self._corex_moe_forward = getattr( - _corex_moe_module, 'moe_forward', None) - if self._corex_moe_forward is not None: - self._use_corex_moe = True - if layer_idx == 0: + try: + # corex_moe module provides direct forward functions + self._corex_moe_forward = getattr( + _corex_moe_module, 'moe_forward', None) + if self._corex_moe_forward is not None: + self._use_corex_moe = True logger.info("MoE: CoreX fused MoE forward available") - else: - raise RuntimeError("corex_moe module loaded but moe_forward missing") + else: + logger.warning("MoE: corex_moe has no moe_forward, using PyTorch") + except Exception as e: + logger.warning("MoE: CoreX MoE init failed (%s), using PyTorch", e) def _pure_pytorch_experts( self,