[Fix] Prevent NaN failure in MesaNet

[Undermyth]

Summary

Fix for issue #1187 .

The current implementation casts q and k after l2norm to float16 forcely, and the dtype of states (h_kk and h_kv) are assigned as the same as k, which is float16.

Such precision is not enough since v is not normalized and accumulated to h_kv, which will soon cause NaN in fwd and bwd during training. Since other ops like DeltaNet assign states in bloat16 when training under bf16 mixed precision, assigning bfloat16 states here should be suitable.

Also the precision is not enough for h_kk, since fp16 can cause a numerically non-positive definite matrix in CG solver and thus the divergence. Lifting precision to fp32 solves the problem. We've tested the op on H100 to ensure that fp32 implementation can fit in the SRAM with dim=128.

Test plan

Use the original test as MesaNet. The error bound of chunk output is slightly relaxed (from 0.006 to 0.007), seems it not likely to be achieved under h_kv with bfloat16

Benchmark / NCU (kernel changes only)

None

Breaking changes

None

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• This is not a minor/cosmetic-only PR (typo, formatting, style-only tweaks).
• Dependent tests pass locally or in CI; new behavior is covered by tests where applicable.
• Kernel changes include same-hardware before/after benchmark numbers (dense + varlen where applicable).

[Undermyth]

Sorry for the late and repeated changes since we took some time to verify the fix on models.
We've validated the fix on a 0.4B scale model with more than 5B training tokens and NaN no longer appears, so I think it's ready for merge.
Still, I'm not sure whether the fix is consistent with the original paper, because the gradient still seems not stable

[Undermyth]

@zhiyuan1i could you please take a look on this? I think it's ready for review