fix platform basic test compatibility

This commit is contained in:
2026-07-15 14:25:52 +08:00
parent e374b14bf9
commit 2b7880efa7
6 changed files with 153 additions and 18 deletions

View File

@@ -940,13 +940,31 @@ class Qwen3_5MoeSparseBlock(nn.Module):
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)
act = None
if _os_prof.environ.get("MOE_DECODE_NATIVE", "0") == "1":
try:
import ixformer.functions as _ixf
gate_up = _ixf.act_bias_mm(
hidden_states.contiguous(),
w13_sel.reshape(-1, H).contiguous(),
None,
scale=1,
act_type="none",
trans_format="TN",
) # (1, K*2*I)
gate_up = gate_up.view(self.top_k, -1).contiguous()
act = _ixf.silu_and_mul(gate_up) # (K, I)
except Exception:
act = None
if act is None:
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)