[JIT kernel] hd=512,1024 in JIT QK norm (cta based) (#17515)

Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
2026-02-16 03:07:24 -05:00
parent 206accd15d
commit ed22720c07
3 changed files with 134 additions and 14 deletions
--- a/python/sglang/jit_kernel/tests/test_qknorm.py
+++ b/python/sglang/jit_kernel/tests/test_qknorm.py
@@ -63,7 +63,7 @@ BS_LIST = [2**n for n in range(0, 14)]
 BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
 N_K_LIST = [2, 4]
 N_Q_LIST = [8, 16]
-HEAD_DIM_LIST = [64, 128, 256]
+HEAD_DIM_LIST = [64, 128, 256, 512, 1024]
 DEVICE = "cuda"
 DTYPE = torch.bfloat16