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

contig

Browse Source
This commit is contained in:
Dan Saunders
2025-09-19 12:47:53 -04:00
parent 336616d659
commit 0f8b921399
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -344,10 +344,15 @@ def moe_ffn_forward_grouped(
w_up_t = w_up.transpose(-2, -1).to(mm_dtype) w_up_t = w_up.transpose(-2, -1).to(mm_dtype)
w2_t = w2.transpose(-2, -1).to(mm_dtype) w2_t = w2.transpose(-2, -1).to(mm_dtype)
routed_in = routed_in.contiguous()
w_gate_t = w_gate_t.contiguous()
gate_out = torch._grouped_mm(routed_in, w_gate_t, offs=offsets) gate_out = torch._grouped_mm(routed_in, w_gate_t, offs=offsets)
torch.ops.aten.silu_(gate_out) torch.ops.aten.silu_(gate_out)
w_up_t = w_up_t.contiguous()
up_out = torch._grouped_mm(routed_in, w_up_t, offs=offsets) up_out = torch._grouped_mm(routed_in, w_up_t, offs=offsets)
gate_out.mul_(up_out) gate_out.mul_(up_out)
gate_out = gate_out.contiguous()
w2_t = w2_t.contiguous()
down_out = torch._grouped_mm(gate_out, w2_t, offs=offsets).to(expert_dtype) down_out = torch._grouped_mm(gate_out, w2_t, offs=offsets).to(expert_dtype)
weights = scores_sorted.unsqueeze(-1).to(expert_dtype) weights = scores_sorted.unsqueeze(-1).to(expert_dtype)