[Triton/Gluon] [Bugfix] Stop the split-K tail at K in the blockscale GEMMs

aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w8_blockscale.py

@@ -24,6 +24,7 @@
         "NUM_KSPLIT",
         "SPLITK_BLOCK_SIZE",
         "EVEN_K",
+        "EVEN_SPLITK",
         "GRID_MN",
         "PREQUANT",
     ],
@@ -33,6 +34,8 @@
 @triton.heuristics(
     {
         "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0,
+        "EVEN_SPLITK": lambda args: args["NUM_KSPLIT"] * args["SPLITK_BLOCK_SIZE"]
+        == args["K"],
         "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"])
         * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]),
     }
@@ -71,6 +74,7 @@ def _gemm_a16w8_blockscale_kernel(
     NUM_KSPLIT: tl.constexpr,
     SPLITK_BLOCK_SIZE: tl.constexpr,
     EVEN_K: tl.constexpr,
+    EVEN_SPLITK: tl.constexpr,
     GRID_MN: tl.constexpr,
     PREQUANT: tl.constexpr,
     DTYPE_MAX: tl.constexpr,
@@ -155,9 +159,11 @@ def _gemm_a16w8_blockscale_kernel(
         acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32
         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)
 
-        for k in tl.range(
-            pid_k * num_k_iter, (pid_k + 1) * num_k_iter, num_stages=num_stages
-        ):
+        k_end = (pid_k + 1) * num_k_iter
+        if not EVEN_SPLITK:
+            # SPLITK_BLOCK_SIZE is rounded up, so the last split can run past K
+            k_end = tl.minimum(k_end, tl.cdiv(K, BLOCK_SIZE_K))
+        for k in tl.range(pid_k * num_k_iter, k_end, num_stages=num_stages):
             # Load the next block of A and B, generate a mask by checking the K dimension.
             # If it is out of bounds, set it to 0.
             if EVEN_K:
@@ -223,6 +229,7 @@ def _gemm_a16w8_blockscale_kernel(
         "NUM_KSPLIT",
         "SPLITK_BLOCK_SIZE",
         "EVEN_K",
+        "EVEN_SPLITK",
         "GRID_MN",
         "PREQUANT",
     ],
@@ -232,6 +239,8 @@ def _gemm_a16w8_blockscale_kernel(
 @triton.heuristics(
     {
         "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0,
+        "EVEN_SPLITK": lambda args: args["NUM_KSPLIT"] * args["SPLITK_BLOCK_SIZE"]
+        == args["K"],
         "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"])
         * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]),
     }
@@ -270,6 +279,7 @@ def _gemm_a16w8_blockscale_preshuffle_kernel(
     NUM_KSPLIT: tl.constexpr,
     SPLITK_BLOCK_SIZE: tl.constexpr,
     EVEN_K: tl.constexpr,
+    EVEN_SPLITK: tl.constexpr,
     GRID_MN: tl.constexpr,
     PREQUANT: tl.constexpr,
     DTYPE_MAX: tl.constexpr,
@@ -362,7 +372,11 @@ def _gemm_a16w8_blockscale_preshuffle_kernel(
         acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32
         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)
 
-        for k in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):
+        k_end = (pid_k + 1) * num_k_iter
+        if not EVEN_SPLITK:
+            # SPLITK_BLOCK_SIZE is rounded up, so the last split can run past K
+            k_end = tl.minimum(k_end, tl.cdiv(K, BLOCK_SIZE_K))
+        for k in range(pid_k * num_k_iter, k_end):
             # Load the next block of A and B, generate a mask by checking the K dimension.
             # If it is out of bounds, set it to 0.
             if EVEN_K:

aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8_blockscale.py

@@ -20,6 +20,7 @@
         "NUM_KSPLIT",
         "SPLITK_BLOCK_SIZE",
         "EVEN_K",
+        "EVEN_SPLITK",
         "cache_modifier",
     ],
 )
@@ -28,6 +29,8 @@
 @triton.heuristics(
     {
         "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0,
+        "EVEN_SPLITK": lambda args: args["NUM_KSPLIT"] * args["SPLITK_BLOCK_SIZE"]
+        == args["K"],
     }
 )
 @triton.jit(repr=_gemm_a8w8_blockscale_repr)
@@ -67,6 +70,7 @@ def _gemm_a8w8_blockscale_kernel(
     NUM_KSPLIT: tl.constexpr,
     SPLITK_BLOCK_SIZE: tl.constexpr,
     EVEN_K: tl.constexpr,
+    EVEN_SPLITK: tl.constexpr,
     cache_modifier: tl.constexpr,
     num_stages: tl.constexpr,
 ):
@@ -154,9 +158,11 @@ def _gemm_a8w8_blockscale_kernel(
         acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32
         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)
 
-        for k in tl.range(
-            pid_k * num_k_iter, (pid_k + 1) * num_k_iter, num_stages=num_stages
-        ):
+        k_end = (pid_k + 1) * num_k_iter
+        if not EVEN_SPLITK:
+            # SPLITK_BLOCK_SIZE is rounded up, so the last split can run past K
+            k_end = tl.minimum(k_end, tl.cdiv(K, BLOCK_SIZE_K))
+        for k in tl.range(pid_k * num_k_iter, k_end, num_stages=num_stages):
             # Load the next block of A and B, generate a mask by checking the K dimension.
             # If it is out of bounds, set it to 0.
             if EVEN_K:
@@ -210,6 +216,7 @@ def _gemm_a8w8_blockscale_kernel(
         "NUM_KSPLIT",
         "SPLITK_BLOCK_SIZE",
         "EVEN_K",
+        "EVEN_SPLITK",
         "cache_modifier",
     ],
 )
@@ -218,6 +225,8 @@ def _gemm_a8w8_blockscale_kernel(
 @triton.heuristics(
     {
         "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0,
+        "EVEN_SPLITK": lambda args: args["NUM_KSPLIT"] * args["SPLITK_BLOCK_SIZE"]
+        == args["K"],
     }
 )
 @triton.jit(repr=_gemm_a8w8_blockscale_preshuffle_repr)
@@ -257,6 +266,7 @@ def _gemm_a8w8_blockscale_preshuffle_kernel(
     NUM_KSPLIT: tl.constexpr,
     SPLITK_BLOCK_SIZE: tl.constexpr,
     EVEN_K: tl.constexpr,
+    EVEN_SPLITK: tl.constexpr,
     cache_modifier: tl.constexpr,
 ):
     """
@@ -353,7 +363,11 @@ def _gemm_a8w8_blockscale_preshuffle_kernel(
         acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32
         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)
 
-        for k in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):
+        k_end = (pid_k + 1) * num_k_iter
+        if not EVEN_SPLITK:
+            # SPLITK_BLOCK_SIZE is rounded up, so the last split can run past K
+            k_end = tl.minimum(k_end, tl.cdiv(K, BLOCK_SIZE_K))
+        for k in range(pid_k * num_k_iter, k_end):
             # Load the per-K-tile block scales up-front
             a_scale = tl.load(a_scale_ptrs)
             b_scale = tl.load(b_scale_ptrs)

aiter/ops/triton/gemm/basic/gemm_a16w8_blockscale.py

@@ -62,6 +62,9 @@ def gemm_a16w8_blockscale(
     if config is None:
         config, _ = _get_config(M, N, K)
 
+    # Normalizing can lower NUM_KSPLIT, so size the outputs afterwards.
+    compute_splitk_params(config, K)
+
     return_y_pp = config["NUM_KSPLIT"] > 1 and skip_reduce
 
     if config["NUM_KSPLIT"] > 1:
@@ -74,8 +77,6 @@ def gemm_a16w8_blockscale(
     if y is None and not return_y_pp:
         y = torch.empty((M, N), dtype=dtype, device=x.device)
 
-    compute_splitk_params(config, K)
-
     # Scale block sizes
     # TODO: need a better way to pass scale block sizes around
     config["GROUP_K"] = triton.next_power_of_2(triton.cdiv(K, w_scale.shape[0]))
@@ -194,6 +195,9 @@ def gemm_a16w8_blockscale_preshuffle(
     if config is None:
         config, _ = _get_config(M, N, K, True)
 
+    # Normalizing can lower NUM_KSPLIT, so size the outputs afterwards.
+    compute_splitk_params(config, K)
+
     return_y_pp = config["NUM_KSPLIT"] > 1 and skip_reduce
 
     if config["NUM_KSPLIT"] > 1:
@@ -206,8 +210,6 @@ def gemm_a16w8_blockscale_preshuffle(
     if y is None and not return_y_pp:
         y = torch.empty((M, N), dtype=dtype, device=x.device)
 
-    compute_splitk_params(config, K)
-
     # Scale block sizes
     # TODO: need a better way to pass scale block sizes around
     config["GROUP_K"] = triton.next_power_of_2(

aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py

@@ -99,12 +99,12 @@ def gemm_a8w8_blockscale(
     if config is None:
         config, _ = _get_config(M, N, K, backend=backend)
 
+    # Normalizing can lower NUM_KSPLIT, so size the outputs afterwards.
+    compute_splitk_params(config, K)
+
     if y is None and (config["NUM_KSPLIT"] == 1 or not skip_reduce):
         y = torch.empty((M, N), dtype=dtype, device=x.device)
 
-    config["SPLITK_BLOCK_SIZE"] = triton.cdiv(
-        K, config["NUM_KSPLIT"]
-    )  # How big each split_k partition is
     if config["NUM_KSPLIT"] > 1:
         y_pp = torch.empty(
             (config["NUM_KSPLIT"], M, N),
@@ -114,8 +114,6 @@ def gemm_a8w8_blockscale(
     else:
         y_pp = None
 
-    compute_splitk_params(config, K)
-
     # Scale block sizes
     # TODO: need a better way to pass scale block sizes around
     config["GROUP_K"] = triton.next_power_of_2(
@@ -314,12 +312,12 @@ def gemm_a8w8_blockscale_preshuffle(
     if kernel_type_from_config is not None:
         kernel_type = kernel_type_from_config
 
+    # Normalizing can lower NUM_KSPLIT, so size the outputs afterwards.
+    compute_splitk_params(config, K)
+
     if y is None and (config["NUM_KSPLIT"] == 1 or not skip_reduce):
         y = torch.empty((M, N), dtype=dtype, device=x.device)
 
-    config["SPLITK_BLOCK_SIZE"] = triton.cdiv(
-        K, config["NUM_KSPLIT"]
-    )  # How big each split_k partition is
     if config["NUM_KSPLIT"] > 1:
         y_pp = torch.empty(
             (config["NUM_KSPLIT"], M, N),

op_tests/triton_tests/gemm/basic/test_gemm_a16w8_blockscale.py

@@ -53,6 +53,9 @@ def get_x_vals():
     # GPT-OSS-120B attention projections
     x_vals += [(2**i, 2880, 4096) for i in range(5, 9)]  # output projection
     x_vals += [(v, 106496, 16384) for v in (256, 4096)]  # LL3 405B FC1
+    # gfx950's default config splits these K into 256-wide partitions, so the
+    # last one runs past K (K=640 -> 3 x 256)
+    x_vals += [(v, 6144, K) for v in (1, 16, 64) for K in (640, 896)]
     return x_vals
 
 
@@ -146,3 +149,21 @@ def test_gemm(dtype, M, N, K, output, shuffle):
     b = run_triton(impl, x, weight_triton, w_scale, prequant, dtype, y)
 
     triton.testing.assert_close(a, b, atol=0.1, rtol=0.1)
+
+
+@pytest.mark.parametrize("shuffle", [False, True])
+def test_gemm_skip_reduce(shuffle):
+    # gfx950's default config asks for 8 splits here and normalizes to 3; the
+    # returned partials must be exactly the ones the kernel wrote
+    M, N, K = 16, 6144, 640
+    x, weight, weight_triton, w_scale, _ = generate_gemm_a16w8_blockscale_inputs(
+        M, N, K, *block_shape, shuffle=shuffle
+    )
+    impl = gemm_a16w8_blockscale_preshuffle if shuffle else gemm_a16w8_blockscale
+
+    a = run_torch(x, weight, w_scale)
+    b = impl(x, weight_triton, w_scale, torch.bfloat16, skip_reduce=True)
+    if b.dim() == 3:
+        b = b.sum(dim=0).to(torch.bfloat16)
+
+    triton.testing.assert_close(a, b, atol=0.1, rtol=0.1)

op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py

@@ -71,6 +71,9 @@ def get_x_vals():
     x_vals += [(v, 7168, 16384) for v in (1, 32, 64, 128, 256, 1024)]
     x_vals += [(v, 6144, 7168) for v in (1, 32, 64, 128, 256, 1024)]
     x_vals += [(v, 7168, 3072) for v in (1, 32, 64, 128, 256, 1024)]
+    # gfx942's tuned preshuffle config splits K=11008 into four 2816-wide
+    # partitions, so the last one runs 256 past K
+    x_vals += [(16, 4096, 11008)]
     return x_vals