tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

remove contig

Browse Source
This commit is contained in:
Dan Saunders
2025-09-18 11:55:15 -04:00
parent 5c197275ad
commit ad4cd39bcd
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/axolotl/kernels/moe/torch_grouped.py
View File

@@ -289,9 +289,9 @@ def moe_ffn_forward_grouped(
zero = torch.zeros_like(x_flat) zero = torch.zeros_like(x_flat)
return zero.view(bsz, seqlen, hdim), router_logits return zero.view(bsz, seqlen, hdim), router_logits
w_gate_t = w_gate[active_idx].transpose(-2, -1).contiguous() w_gate_t = w_gate[active_idx].transpose(-2, -1)
w_up_t = w_up[active_idx].transpose(-2, -1).contiguous() w_up_t = w_up[active_idx].transpose(-2, -1)
w2_t = w2[active_idx].transpose(-2, -1).contiguous() w2_t = w2[active_idx].transpose(-2, -1)
routed_in = routed_input.to(expert_dtype) routed_in = routed_input.to(expert_dtype)
gate_out = torch._grouped_mm(routed_in, w_gate_t, offs=offsets) gate_out = torch._grouped_mm(routed_in, w_gate_t, offs=offsets)