diff --git a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w8_blockscale.py b/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w8_blockscale.py index ef146ed865c..fa741571dd5 100644 --- a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w8_blockscale.py +++ b/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: diff --git a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8_blockscale.py b/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8_blockscale.py index 7f15dc913e8..cb8d7b638bc 100644 --- a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8_blockscale.py +++ b/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) diff --git a/aiter/ops/triton/gemm/basic/gemm_a16w8_blockscale.py b/aiter/ops/triton/gemm/basic/gemm_a16w8_blockscale.py index c30e67d6ba8..4e6da13f13b 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a16w8_blockscale.py +++ b/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( diff --git a/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py b/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py index 5914c58412d..fc6e416c979 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py +++ b/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), diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a16w8_blockscale.py b/op_tests/triton_tests/gemm/basic/test_gemm_a16w8_blockscale.py index 58458d0b211..68136a7f52b 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a16w8_blockscale.py +++ b/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) diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py index c2fd8a2dba9..f1f14d08820 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py +++ b/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