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

Fix wo_a: permute to groups-first layout for grouped GEMM

Browse Source 
The grouped GEMM expects mat_a to be laid out contiguously per group:
[all tokens for group0, all tokens for group1, ...]
A simple reshape of (T, G, D) → (T*G, D) gives interleaved layout
which is wrong. Fix: permute to (G, T, D) before flattening.
Same fix for output: permute (G, T, R) → (T, G, R).
This commit is contained in:
biondizzle
2026-05-19 02:41:32 +00:00
parent 77e4970d93
commit 5f5b997fc3
1 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
cutedsl/wo_a_grouped_linear.py
View File

@@ -219,7 +219,12 @@ class CuTeDSLNvfp4WoA:
# Reshape: (tokens, n_local_heads, head_dim) → (tokens, n_local_groups, group_in_features)
o_grouped = o.reshape(num_tokens, self.n_local_groups, self.group_in_features)
# Flatten for GEMM: (tokens * n_groups, group_in_features)
# CRITICAL: permute to (n_local_groups, tokens, group_in_features) before flattening.
# The grouped GEMM expects mat_a to be laid out contiguously per group:
# [all tokens for group0, all tokens for group1, ...]
# A simple reshape of (T, G, D) → (T*G, D) gives interleaved:
# [t0_g0, t0_g1, ..., t0_g7, t1_g0, ...] which is WRONG.
o_grouped = o_grouped.permute(1, 0, 2) # (G, T, D)
o_flat = o_grouped.reshape(num_tokens * self.n_local_groups, self.group_in_features)
padded_rows_per_group = cutedsl_ceil_div(num_tokens, 128) * 128
@@ -259,7 +264,10 @@ class CuTeDSLNvfp4WoA:
# Extract real outputs and reshape
out = out[:num_tokens * self.n_local_groups]
z = out.reshape(num_tokens, self.n_local_groups, self.o_lora_rank)
# GEMM output is (G*T, o_lora_rank) in groups-first order
# Reshape to (G, T, o_lora_rank) then permute to (T, G, o_lora_rank)
z = out.reshape(self.n_local_groups, num_tokens, self.o_lora_rank)
z = z.permute(1, 0, 2) # (T, G, o_lora_rank)
return z
def __call__(self, o: torch.Tensor) -> torch.Tensor: