fix platform basic test compatibility
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user