[codex] Optimize hidden-size 512 RMSNorm dispatch (#24710)

Co-authored-by: Codex <codex@example.com>
2026-05-19 09:26:10 +08:00
parent 170fe57cf0
commit 2424303dfb
3 changed files with 33 additions and 21 deletions
@@ -88,7 +88,7 @@ def test_rmsnorm_hidden_size_support(hidden_size: int) -> None:
        (64, "RMSNormWarpKernel"),
        (128, "RMSNormWarpKernel"),
        (256, "RMSNormWarpKernel"),
-        (512, "RMSNormKernel"),
+        (512, "RMSNormHalfKernel"),
        (1536, "RMSNormKernel"),
        (2048, "RMSNormHalfKernel"),
        (2304, "RMSNormKernel"),  # NOTE: not 512 aligned