From 389ec6fa00c130749a79816590f5d3662e5f482a Mon Sep 17 00:00:00 2001 From: Anoop Kapoor Date: Sat, 8 Aug 2026 23:17:31 -0700 Subject: [PATCH 1/4] @FIR-2167 - Extend advanced Triton MAT_MUL offload coverage --- ggml/src/ggml-tsavorite/ggml-tsavorite.cpp | 74 ++++++++++++++++++---- tsavorite-model-deployment.yaml | 8 ++- 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp b/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp index 7a295301df9c..7d1adfdba6fc 100644 --- a/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp +++ b/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp @@ -125,6 +125,7 @@ struct TsavoriteRuntimeState { void **loadResult_matmul_1x8 = nullptr; void **loadResult_matmul_2x4 = nullptr; bool advanced_matmul_shape_offload = false; + bool advanced_matmul_broadcast_offload = false; bool triton_matmul_small_n_transpose_opt = false; #endif @@ -184,6 +185,7 @@ auto &loadResult_triton_add = g_rt.loadResult_triton_add; auto &loadResult_matmul_1x8 = g_rt.loadResult_matmul_1x8; auto &loadResult_matmul_2x4 = g_rt.loadResult_matmul_2x4; auto &advanced_matmul_shape_offload = g_rt.advanced_matmul_shape_offload; +auto &advanced_matmul_broadcast_offload = g_rt.advanced_matmul_broadcast_offload; auto &triton_matmul_small_n_transpose_opt = g_rt.triton_matmul_small_n_transpose_opt; #endif } // anonymous namespace @@ -282,6 +284,8 @@ struct tsi_deploy_cfg_t { #if TRITON_MAT_MUL bool advanced_matmul_shape_offload = false; bool has_advanced_matmul_shape_offload = false; + bool advanced_matmul_broadcast_offload = false; + bool has_advanced_matmul_broadcast_offload = false; bool triton_matmul_small_n_transpose_opt = false; bool has_triton_matmul_small_n_transpose_opt = false; #endif @@ -359,6 +363,15 @@ static tsi_deploy_cfg_t tsi_read_deploy_yaml(const std::string &path) { cfg.has_advanced_matmul_shape_offload = true; } } + if (t.find("advanced_matmul_broadcast_offload") != std::string::npos && + t.find(':') != std::string::npos) { + bool b = false; + if (tsi_parse_bool_after_colon(t, &b)) { + cfg.advanced_matmul_broadcast_offload = b; + cfg.has_advanced_matmul_broadcast_offload = true; + } + } + if (t.find("triton_matmul_small_n_transpose_opt") != std::string::npos && t.find(':') != std::string::npos) { bool b = false; @@ -1648,6 +1661,11 @@ static void ensure_tsi_runtime_initialized() { cfg.has_advanced_matmul_shape_offload ? cfg.advanced_matmul_shape_offload : false; + + advanced_matmul_broadcast_offload = + cfg.has_advanced_matmul_broadcast_offload ? + cfg.advanced_matmul_broadcast_offload : + false; triton_matmul_small_n_transpose_opt = cfg.has_triton_matmul_small_n_transpose_opt ? cfg.triton_matmul_small_n_transpose_opt : @@ -1680,6 +1698,8 @@ static void ensure_tsi_runtime_initialized() { #if TRITON_MAT_MUL printf(" advanced_matmul_shape_offload=%d", (int)advanced_matmul_shape_offload); + printf(" advanced_matmul_broadcast_offload=%d", + (int)advanced_matmul_broadcast_offload); printf(" triton_matmul_small_n_transpose_opt=%d", (int)triton_matmul_small_n_transpose_opt); #endif @@ -3702,6 +3722,22 @@ static inline int64_t map_repeat_i64(int64_t out_idx, int64_t in_dim) { return (r < 0) ? (r + in_dim) : r; } +static inline int64_t map_repeat_dim_i64(int64_t out_idx, int64_t out_dim, int64_t in_dim) { + if (in_dim <= 1) return 0; + if (out_dim <= 0) return 0; + if (in_dim == out_dim) return out_idx; + + if ((out_dim % in_dim) == 0) { + const int64_t group = out_dim / in_dim; + int64_t r = out_idx / group; + if (r < 0) r = 0; + if (r >= in_dim) r = in_dim - 1; + return r; + } + + return map_repeat_i64(out_idx, in_dim); +} + // ============================================================================ // ABI WRAPPERS: raw pointers -> MemRefDescriptor<4> -> call MLIR ciface @@ -3877,7 +3913,7 @@ static inline const triton_matmul_txe_shape_t &triton_matmul_select_shape( // Data type specific #define TRITON_MATMUL_F32_K_DIM 32 -#define TSAV_TRITON_MATMUL_MAX_K 8192 +#define TSAV_TRITON_MATMUL_MAX_K 12288 #define TSAV_TRITON_MATMUL_MAX_M 32768 #define TSAV_TRITON_MATMUL_MAX_N 4096 @@ -3928,6 +3964,13 @@ static inline bool tsavorite_mul_mat_advanced_shape_ok(const struct ggml_tensor const int64_t D2 = op->ne[2]; const int64_t D3 = op->ne[3]; + const bool has_broadcast_dims = + (A2 != 1 || A3 != 1 || B2 != 1 || B3 != 1 || D2 != 1 || D3 != 1); + + if (has_broadcast_dims && !advanced_matmul_broadcast_offload) { + return false; + } + if (b->ne[0] != K || op->ne[0] != M || op->ne[1] != N) { return false; } @@ -3936,15 +3979,24 @@ static inline bool tsavorite_mul_mat_advanced_shape_ok(const struct ggml_tensor return false; } + if (A2 <= 0 || A3 <= 0 || B2 <= 0 || B3 <= 0 || D2 <= 0 || D3 <= 0) { + return false; + } + if (D2 != ((A2 > B2) ? A2 : B2) || D3 != ((A3 > B3) ? A3 : B3)) { return false; } - if (!(A2 == D2 || A2 == 1) || !(B2 == D2 || B2 == 1)) { + /* + * Allow GGML repeat/broadcast layouts where an input batch dimension divides + * the output batch dimension. Runtime packing uses map_repeat_dim_i64(), so + * shapes such as A2=4, B2=32, D2=32 are valid. + */ + if ((D2 % A2) != 0 || (D2 % B2) != 0) { return false; } - if (!(A3 == D3 || A3 == 1) || !(B3 == D3 || B3 == 1)) { + if ((D3 % A3) != 0 || (D3 % B3) != 0) { return false; } @@ -5265,10 +5317,10 @@ static enum ggml_status ggml_tsavorite_run_tmu_mul_mat( for (int64_t d3 = 0; d3 < D3; ++d3) { for (int64_t d2 = 0; d2 < D2; ++d2) { - const int64_t a_d2 = map_repeat_i64(d2, A2); - const int64_t a_d3 = map_repeat_i64(d3, A3); - const int64_t b_d2 = map_repeat_i64(d2, B2); - const int64_t b_d3 = map_repeat_i64(d3, B3); + const int64_t a_d2 = map_repeat_dim_i64(d2, D2, A2); + const int64_t a_d3 = map_repeat_dim_i64(d3, D3, A3); + const int64_t b_d2 = map_repeat_dim_i64(d2, D2, B2); + const int64_t b_d3 = map_repeat_dim_i64(d3, D3, B3); char *A_ptr = A_base + a_d2 * a_nb2 + a_d3 * a_nb3; char *B_ptr = B_base + b_d2 * b_nb2 + b_d3 * b_nb3; @@ -5406,10 +5458,10 @@ static enum ggml_status ggml_tsavorite_run_tmu_mul_mat( for (int64_t d3 = 0; d3 < D3; ++d3) { for (int64_t d2 = 0; d2 < D2; ++d2) { - const int64_t a_d2 = map_repeat_i64(d2, A2); - const int64_t a_d3 = map_repeat_i64(d3, A3); - const int64_t b_d2 = map_repeat_i64(d2, B2); - const int64_t b_d3 = map_repeat_i64(d3, B3); + const int64_t a_d2 = map_repeat_dim_i64(d2, D2, A2); + const int64_t a_d3 = map_repeat_dim_i64(d3, D3, A3); + const int64_t b_d2 = map_repeat_dim_i64(d2, D2, B2); + const int64_t b_d3 = map_repeat_dim_i64(d3, D3, B3); char *A_ptr = A_base + a_d2 * a_nb2 + a_d3 * a_nb3; char *B_ptr = B_base + b_d2 * b_nb2 + b_d3 * b_nb3; diff --git a/tsavorite-model-deployment.yaml b/tsavorite-model-deployment.yaml index 0e48e49d4df2..76d1f67920d2 100755 --- a/tsavorite-model-deployment.yaml +++ b/tsavorite-model-deployment.yaml @@ -17,7 +17,7 @@ multi_thread_enable: true ## user_dram_size_gb: 8 ## user_dram_size_gb: 12 ## user_dram_size_gb: 24 -user_dram_size_gb: 8 +user_dram_size_gb: 16 # Enable additional Triton MAT_MUL shapes beyond stable baseline. # false = old behavior @@ -28,3 +28,9 @@ advanced_matmul_shape_offload: false # false = old behavior # true = for M >> N, compute swapped [N x M] and transpose copyback to [M x N] triton_matmul_small_n_transpose_opt: false + +## Enable Triton MAT_MUL broadcast/batched D2/D3 offload. +## false = keep broadcast MAT_MUL on fallback path +## true = allow advanced MAT_MUL helper to offload supported broadcast shapes + +advanced_matmul_broadcast_offload: false From e41551c01b4b15874fce9464d06c3596fc3a787f Mon Sep 17 00:00:00 2001 From: Anoop Kapoor Date: Sun, 9 Aug 2026 14:38:36 -0700 Subject: [PATCH 2/4] Address Qudo and qubic PR comments --- ggml/src/ggml-tsavorite/ggml-tsavorite.cpp | 17 ++++++++++++++++- tsi-pkg-build.sh | 12 +++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp b/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp index 7d1adfdba6fc..1f2bcc51223b 100644 --- a/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp +++ b/ggml/src/ggml-tsavorite/ggml-tsavorite.cpp @@ -3722,6 +3722,15 @@ static inline int64_t map_repeat_i64(int64_t out_idx, int64_t in_dim) { return (r < 0) ? (r + in_dim) : r; } +/* + * Map an output D2/D3 batch index to the corresponding input batch index for + * grouped broadcast MAT_MUL layouts. + * + * This is intentionally not generic GGML_OP_REPEAT modulo mapping. For grouped + * head layouts such as A2=4, D2=32, output heads 0..7 map to input head 0, + * 8..15 map to input head 1, etc. This matches the grouped-head MAT_MUL + * packing requirement used by the Triton offload path. + */ static inline int64_t map_repeat_dim_i64(int64_t out_idx, int64_t out_dim, int64_t in_dim) { if (in_dim <= 1) return 0; if (out_dim <= 0) return 0; @@ -3964,8 +3973,14 @@ static inline bool tsavorite_mul_mat_advanced_shape_ok(const struct ggml_tensor const int64_t D2 = op->ne[2]; const int64_t D3 = op->ne[3]; + /* + * Only require the broadcast flag for actual broadcast layouts where an + * input batch dimension differs from the output batch dimension. Batched + * non-broadcast layouts such as A2=B2=D2=2 should remain controlled by + * advanced_matmul_shape_offload and must not require the broadcast flag. + */ const bool has_broadcast_dims = - (A2 != 1 || A3 != 1 || B2 != 1 || B3 != 1 || D2 != 1 || D3 != 1); + (A2 != D2 || B2 != D2 || A3 != D3 || B3 != D3); if (has_broadcast_dims && !advanced_matmul_broadcast_offload) { return false; diff --git a/tsi-pkg-build.sh b/tsi-pkg-build.sh index 947ba3ecde9a..7dcd7bc5d8b5 100755 --- a/tsi-pkg-build.sh +++ b/tsi-pkg-build.sh @@ -1062,6 +1062,7 @@ update_one_tsavorite_deployment_yaml() { local deployment_yaml_path="$1" local txe_count="$2" local advanced_matmul_shape_offload="false" + local advanced_matmul_broadcast_offload="false" local triton_matmul_small_n_transpose_opt="false" local user_dram_size_gb="8" @@ -1073,12 +1074,16 @@ local user_dram_size_gb="8" local existing_user_dram_size_gb existing_advanced="$(extract_deployment_yaml_value "${deployment_yaml_path}" "advanced_matmul_shape_offload")" + existing_broadcast="$(extract_deployment_yaml_value "${deployment_yaml_path}" "advanced_matmul_broadcast_offload")" existing_small_n_opt="$(extract_deployment_yaml_value "${deployment_yaml_path}" "triton_matmul_small_n_transpose_opt")" if [ -n "${existing_advanced}" ]; then advanced_matmul_shape_offload="${existing_advanced}" fi + if [ -n "${existing_broadcast}" ]; then + advanced_matmul_broadcast_offload="${existing_broadcast}" + fi if [ -n "${existing_small_n_opt}" ]; then triton_matmul_small_n_transpose_opt="${existing_small_n_opt}" fi @@ -1116,13 +1121,18 @@ user_dram_size_gb: $user_dram_size_gb # true = new offload shapes advanced_matmul_shape_offload: ${advanced_matmul_shape_offload} +## Enable Triton MAT_MUL broadcast/batched D2/D3 offload. +## false = keep broadcast MAT_MUL on fallback path +## true = allow advanced MAT_MUL helper to offload supported broadcast shapes +advanced_matmul_broadcast_offload: ${advanced_matmul_broadcast_offload} + # Enable Triton MAT_MUL small-N transpose optimization. # false = old behavior # true = for M >> N, compute swapped [N x M] and transpose copyback to [M x N] triton_matmul_small_n_transpose_opt: ${triton_matmul_small_n_transpose_opt} EOF - echo "INFO: updated ${deployment_yaml_path} with txe_count:${txe_count}, multi_thread_enable:true; preserved advanced_matmul_shape_offload:${advanced_matmul_shape_offload}, triton_matmul_small_n_transpose_opt:${triton_matmul_small_n_transpose_opt}" + echo "INFO: updated ${deployment_yaml_path} with txe_count:${txe_count}, multi_thread_enable:true; preserved advanced_matmul_shape_offload:${advanced_matmul_shape_offload}, advanced_matmul_broadcast_offload:${advanced_matmul_broadcast_offload}, triton_matmul_small_n_transpose_opt:${triton_matmul_small_n_transpose_opt}" return 0 } From fa224ed02bb801c89943f42eff3c082e4914349b Mon Sep 17 00:00:00 2001 From: Anoop Kapoor Date: Sun, 9 Aug 2026 14:45:08 -0700 Subject: [PATCH 3/4] Address all PR comments --- tsi-pkg-build.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/tsi-pkg-build.sh b/tsi-pkg-build.sh index 7dcd7bc5d8b5..8fd1669512aa 100755 --- a/tsi-pkg-build.sh +++ b/tsi-pkg-build.sh @@ -1070,6 +1070,7 @@ local user_dram_size_gb="8" if [ -f "${deployment_yaml_path}" ]; then local existing_advanced + local existing_broadcast local existing_small_n_opt local existing_user_dram_size_gb From a29c9bbb343f1b9f649f16ae3996acbb18420f40 Mon Sep 17 00:00:00 2001 From: Anoop Kapoor Date: Sun, 9 Aug 2026 21:18:11 -0700 Subject: [PATCH 4/4] Addressed Ashish's comments --- tsavorite-model-deployment.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tsavorite-model-deployment.yaml b/tsavorite-model-deployment.yaml index 76d1f67920d2..f6daada9cc3b 100755 --- a/tsavorite-model-deployment.yaml +++ b/tsavorite-model-deployment.yaml @@ -17,7 +17,7 @@ multi_thread_enable: true ## user_dram_size_gb: 8 ## user_dram_size_gb: 12 ## user_dram_size_gb: 24 -user_dram_size_gb: 16 +user_dram_size_gb: 8 # Enable additional Triton MAT_MUL shapes beyond stable baseline. # false = old behavior