biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix: BMM batch dim alignment for wo_a

Browse Source
This commit is contained in:
biondizzle
2026-05-30 22:49:21 +00:00
parent 52df3bc26c
commit 904753f62a
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
single_shot_inference.py
View File

@@ -278,10 +278,10 @@ def forward_layer(x, w, li, cfg, rope_cos, rope_sin):
# Fallback: just try dense linear
oa_3d = oa_w_bf16.reshape(o_groups, -1, group_input_dim) if oa_w_bf16.shape[-1] == group_input_dim else oa_w_bf16.T.reshape(o_groups, -1, group_input_dim)
# Batched matmul: (1, 16, 4096) × (16, 1024, 4096)ᵀ → (1, 16, 1024)
# torch.bmm needs (B, M, K) × (B, K, N) → (B, M, N)
grouped_out = torch.bmm(attn_grouped, oa_3d.transpose(1, 2)) # (1, 16, o_rank)
grouped_flat = grouped_out.reshape(T, o_groups * o_rank) # (1, 16384)
# Batched matmul: (16, 1, 4096) × (16, 4096, 1024) → (16, 1, 1024)
attn_for_bmm = attn_grouped.permute(1, 0, 2) # (16, T=1, 4096)
grouped_out = torch.bmm(attn_for_bmm, oa_3d.transpose(1, 2)) # (16, 1, o_rank)
grouped_flat = grouped_out.permute(1, 0, 2).reshape(T, o_groups * o_rank) # (1, 16384)
attn_proj = nvfp4_linear(grouped_flat, ob_w, ob_s, ob_s2) # (1, H)