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

Fix HcHead: use FP32 for RMSNorm + linear (matches HF reference)

Browse Source
This commit is contained in:
biondizzle
2026-05-31 21:13:21 +00:00
parent 274ea13251
commit 23f1cf4065
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

@@ -343,13 +343,13 @@ class HcHead:
# Flatten: (T, n_hc * H)
X_flat = X_L.reshape(T, self.K).bfloat16()
# Unweighted RMSNorm (same as in mHC)
# Unweighted RMSNorm on flattened residual (FP32 for numerical stability)
X_f = X_flat.float()
rms_inv = X_f.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
X_normed = (X_f * rms_inv).bfloat16()
X_normed = X_f * rms_inv # Keep FP32 for the linear
# Linear projection: (T, K) @ (4, K).T → (T, 4)
mixes = torch.nn.functional.linear(X_normed, self.fn.bfloat16()).float()
# Linear projection: (T, K) @ (4, K).T → (T, 4) in FP32
mixes = torch.nn.functional.linear(X_normed, self.fn.float()).float()
# Apply scale + bias + sigmoid + eps
pre = torch.sigmoid(mixes * self.scale + self.base.float().unsqueeze(0)) + self.eps