Repository navigation
bug: Numeric divergence between the mHC layer and its Pallas kernel #5195
Description
Bug report
DeepSeek v4. The reference path had its own mapping helper that cast alpha_scale and
beta down to self.dtype before computing sigmoid(alpha_scale * h + beta),
while the Pallas path called common.compute_sigmoid_gate, which keeps the
scale and bias in float32. With dtype: bfloat16 the two branches therefore
gated on different numbers, so toggling use_mhc_pallas_kernel changed the
layer's output beyond accumulation noise.
Fixed in PR:
#5179
Logs/Output
No response
Environment Information
No response
Additional Context
No response
Activity
Kaustubh1235 commentedon Sep 12, 2026
I will take this issue. Please assign it to me.
The problem seems to be a type mismatch between the mHC layer and its Pallas kernel. The mHC layer casts
alpha_scaleandbetatoself.dtypebefore computing the sigmoid, while the Pallas kernel usesfloat32. I would first check thecompute_sigmoid_gatefunction in the Pallas kernel. A small fix would involve ensuring both paths use the same data type, likely by casting withincompute_sigmoid_gateto matchself.dtype. I will verify the fix by comparing outputs with and withoutuse_mhc_pallas_kernel.denis-mil commentedon Sep 12, 2026
Hi @Kaustubh1235, thanks for offering to look into this! This issue has actually already been addressed and fixed in PR #5179.
Kaustubh1235 commentedon Sep 20, 2026
Thanks for letting me know, Denis. I'll check out the PR to see the changes made.