[nvidia] Gemma4 nvfp4 fix (#22079)

2026-04-10 08:44:34 +08:00
parent cebd9c2a1e
commit 5638d40f3a
1 changed files with 8 additions and 0 deletions
@@ -72,6 +72,14 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
                BLOCK_M, BLOCK_N = (64, 64)
            else:
                BLOCK_M, BLOCK_N = (32, 32)
+        elif _is_cuda and CUDA_CAPABILITY[0] == 10:
+            # Blackwell data-center architecture (GB200, B200, sm_100a)
+            # sm_100a has different register constraints from Hopper; Hopper block sizes
+            # cause PTX register exhaustion (>255 regs) for large head dims (Lq=512).
+            if Lq <= 256:
+                BLOCK_M, BLOCK_N = (64, 64)
+            else:
+                BLOCK_M, BLOCK_N = (16, 64)
        elif _is_cuda and CUDA_CAPABILITY[0] >= 9:
            # Hopper architecture (H100, etc.)
            if Lq <= 256: