[akapoor3518]

Summary

This PR adds broadcast/batched D2/D3 Triton MAT_MUL offload support behind a new runtime deployment flag:

advanced_matmul_broadcast_offload

Supported broadcast layouts can be offloaded when enabled, while existing fallback behavior remains unchanged when disabled.

Changes

• Added advanced_matmul_broadcast_offload runtime deployment flag.
• Enabled supported broadcast/batched D2/D3 Triton MAT_MUL shapes.
• Preserved default fallback behavior when the flag is disabled.
• Added grouped D2/D3 mapping support for broadcast MAT_MUL packing.
• Preserved the new flag during deployment YAML regeneration.

Validation

• Tiny-Llama prompt output remained correct: Luna.
• Broadcast MAT_MUL OPU execution verified.
• Existing non-broadcast MAT_MUL offload behavior preserved.

Follow-up

• Additional grouped-head/repeat-head broadcast MAT_MUL patterns will be evaluated in a future PR.

TESTED on POSIX
[akapoor@wspd0 llama.cpp]$ build-posix/bin/llama-cli -p "My cat's name is" -m /proj/rel/sw/ggml/models/tinyllama-vo-5m-para.gguf --device tSavorite -c 12288 --temp 0.0 --n-predict 2 --repeat-penalty 1.5 -b 1024 --top-k 50 --top-p 0.9 --repeat-last-n 5 --no-warmup --no-conversation

TSI deploy yaml=tsavorite-model-deployment.yaml txe_count=1 multi_thread_enable=1 user_dram_size_gb=16 user_dram_size_bytes=17179869184 advanced_matmul_shape_offload=1 advanced_matmul_broadcast_offload=1 triton_matmul_small_n_transpose_opt=1
My cat's name is Tim.

llama_perf_sampler_print: sampling time = 3.83 ms / 9 runs ( 0.43 ms per token, 2346.81 tokens per second)
llama_perf_context_print: load time = 1082.19 ms
llama_perf_context_print: prompt eval time = 681.79 ms / 7 tokens ( 97.40 ms per token, 10.27 tokens per second)
llama_perf_context_print: eval time = 449.93 ms / 1 runs ( 449.93 ms per token, 2.22 tokens per second)
llama_perf_context_print: total time = 1536.75 ms / 8 tokens

=== GGML Perf Summary ===
Op Target Runs TSI_KERNEL-RUN Total us Avg us
ADD OPU 32 116 25360 792.50
MUL OPU 34 124 146173 4299.21
RMS_NORM OPU 34 34 56419 1659.38
MUL_MAT CPU 290 0 31741 109.45
MUL_MAT OPU 66 66 813995 12333.26
CONT OPU 32 0 8902 278.19
RESHAPE CPU 187 0 59 0.32
VIEW CPU 177 0 15 0.08
VIEW OPU 32 0 5 0.16
PERMUTE CPU 117 0 22 0.19
PERMUTE OPU 32 0 7 0.22
TRANSPOSE CPU 58 0 17 0.29
GET_ROWS OPU 6 0 1572 262.00
SET_ROWS CPU 124 0 144 1.16
SOFT_MAX OPU 16 0 8598 537.38
ROPE CPU 120 0 426 3.55
GLU OPU 16 58 43814 2738.38

OPU Profiling Results:
Profiler disabled

[akapoor@wspd0 llama.cpp]$
[akapoor@wspd0 llama.cpp]$
[akapoor@wspd0 llama.cpp]$
[akapoor@wspd0 llama.cpp]$ build-posix/bin/llama-cli -p "My cat's name is" -m /proj/rel/sw/ggml/models/Tiny-Llama-v0.3-FP32-1.1B-F32.gguf--device tSavorite -c 12288 --temp 0.0 --n-predict 2 --repeat-penalty 1.5 -b 1024 --top-k 50 --top-p 0.9 --repeat-last-n 5 --no-warmup --no-conversation

TSI deploy yaml=tsavorite-model-deployment.yaml txe_count=1 multi_thread_enable=1 user_dram_size_gb=16 user_dram_size_bytes=17179869184 advanced_matmul_shape_offload=1 advanced_matmul_broadcast_offload=1 triton_matmul_small_n_transpose_opt=1
error: invalid argument: tSavorite
[akapoor@wspd0 llama.cpp]$ build-posix/bin/llama-cli -p "My cat's name is" -m /proj/rel/sw/ggml/models/Tiny-Llama-v0.3-FP32-1.1B-F32.gguf --device tSavorite -c 12288 --temp 0.0 --n-predict 2 --repeat-penalty 1.5 -b 1024 --top-k 50 --top-p 0.9 --repeat-last-n 5 --no-warmup --no-conversation

TSI deploy yaml=tsavorite-model-deployment.yaml txe_count=1 multi_thread_enable=1 user_dram_size_gb=16 user_dram_size_bytes=17179869184 advanced_matmul_shape_offload=1 advanced_matmul_broadcast_offload=1 triton_matmul_small_n_transpose_opt=1
My cat's name is Luna

llama_perf_sampler_print: sampling time = 6.90 ms / 9 runs ( 0.77 ms per token, 1304.54 tokens per second)
llama_perf_context_print: load time = 148486.67 ms
llama_perf_context_print: prompt eval time = 108607.46 ms / 7 tokens (15515.35 ms per token, 0.06 tokens per second)
llama_perf_context_print: eval time = 107299.95 ms / 1 runs (107299.95 ms per token, 0.01 tokens per second)
llama_perf_context_print: total time = 255806.37 ms / 8 tokens

=== GGML Perf Summary ===
Op Target Runs TSI_KERNEL-RUN Total us Avg us
ADD OPU 88 340 111535 1267.44
MUL OPU 90 348 253851 2820.57
RMS_NORM OPU 90 90 312472 3471.91
MUL_MAT CPU 664 0 1195937 1801.11
MUL_MAT OPU 222 1586 211273483 951682.36
CONT OPU 88 0 15360 174.55
RESHAPE CPU 479 0 317 0.66
VIEW CPU 453 0 47 0.10
VIEW OPU 88 0 22 0.25
PERMUTE CPU 316 0 176 0.56
PERMUTE OPU 88 0 21 0.24
TRANSPOSE CPU 147 0 33 0.22
GET_ROWS OPU 6 0 11730 1955.00
SET_ROWS CPU 342 0 2987 8.73
SOFT_MAX OPU 44 0 87001 1977.30
ROPE CPU 350 0 9458 27.02
GLU OPU 44 170 402795 9154.43

OPU Profiling Results:
Profiler disabled

---

Summary by cubic

Extends Triton MAT_MUL offload to support grouped broadcast D2/D3 batches and raises the K cap to 12288 to widen safe OPU coverage. Addresses Linear FIR-2172 and keeps broadcast offload behind a runtime flag.

• New Features

  • Adds advanced_matmul_broadcast_offload runtime flag (read from YAML, printed at init, preserved by tsi-pkg-build.sh).
  • Requires the flag only when input batch dims differ from output; non-broadcast batched shapes stay under advanced_matmul_shape_offload.
  • Implements map_repeat_dim_i64 and updates packing loops; relaxes shape checks to allow D2/D3 that are multiples of input batch dims.
  • Increases TSAV_TRITON_MATMUL_MAX_K from 8192 to 12288.
  • Adds YAML option for broadcast offload.

• Migration

  • Default stays unchanged; broadcast offload is disabled. Enable via advanced_matmul_broadcast_offload: true in tsavorite-model-deployment.yaml.

^{Written for commit a29c9bb. Summary will update on new commits.}

[qodo-code-review]

PR Summary by Qodo

Enable Triton MAT_MUL broadcast offload with grouped D2/D3 repeat mapping

✨ Enhancement ⚙️ Configuration changes 🕐 40+ Minutes

AI Description

• Add runtime YAML flag to gate Triton broadcast/batched MAT_MUL offload.
• Expand Triton MAT_MUL K-limit to 12288 for longer-context workloads.
• Support grouped repeat mapping for D2/D3 broadcast layouts during MAT_MUL packing.

Diagram

graph TD
  Y[/"tsavorite-model-deployment.yaml"/] --> P["tsi_read_deploy_yaml()"] --> I["ensure_tsi_runtime_initialized()"] --> F["Runtime flags"] --> S{"Broadcast dims?"}
  S -->|"no"| O["Triton MAT_MUL offload"]
  S -->|"yes + flag on"| M["map_repeat_dim_i64()"] --> O
  S -->|"yes + flag off"| C["Fallback path"]

  subgraph Legend
    direction LR
    _cfg[/"Config file"/] ~~~ _fn["Function"] ~~~ _dec{"Decision"}
  end

Loading

High-Level Assessment

The following are alternative approaches to this PR:

1. Auto-enable broadcast offload for safe shapes (no explicit flag)

• ➕ Maximizes offload coverage without requiring YAML changes
• ➕ Avoids mismatched configs across environments
• ➖ Riskier rollout: broadcast mapping correctness depends on assumptions that may not hold for all models/layouts
• ➖ Harder to bisect performance/correctness regressions without an explicit kill switch

2. Generalize broadcast mapping using tensor strides/nb (full stride-aware addressing)

• ➕ More correct for arbitrary broadcast/reshape/view patterns
• ➕ Could eliminate shape divisibility heuristics
• ➖ More invasive changes across packing and validation paths
• ➖ Higher risk of subtle correctness bugs and performance regressions; longer review/validation cycle

Recommendation: Keep the PR’s explicit runtime flag + divisibility checks approach. It’s a safe, incremental expansion of Triton offload coverage, provides an immediate rollback mechanism (disable the flag), and confines correctness risk to validated broadcast/grouped-repeat layouts.

Files changed (2) +70 / -12

Enhancement (1) +63 / -11

ggml-tsavorite.cppGate broadcast MAT_MUL offload and add grouped-repeat D2/D3 mapping +63/-11

Gate broadcast MAT_MUL offload and add grouped-repeat D2/D3 mapping

• Adds a new runtime flag (advanced_matmul_broadcast_offload) parsed from the deployment YAML and printed at startup. Expands Triton MAT_MUL eligibility by raising the max supported K to 12288 and allowing batched/broadcast D2/D3 shapes when enabled, with additional validation that output dims are divisible by input dims. Updates MAT_MUL packing to use a new dimension-aware repeat mapper (map_repeat_dim_i64) for grouped-repeat layouts.

ggml/src/ggml-tsavorite/ggml-tsavorite.cpp

Other (1) +7 / -1

tsavorite-model-deployment.yamlDocument broadcast MAT_MUL offload flag and adjust default DRAM size +7/-1

Document broadcast MAT_MUL offload flag and adjust default DRAM size

• Updates the sample deployment config to set user_dram_size_gb to 16 and documents a new advanced_matmul_broadcast_offload toggle for enabling Triton broadcast/batched MAT_MUL offload.

tsavorite-model-deployment.yaml

[qodo-code-review]

Code Review by Qodo

🐞 Bugs (1) 📘 Rule violations (0) 📜 Skill insights (0)

1. Wrong repeat index mapping 🐞 Bug ≡ Correctness

Description

map_repeat_dim_i64() maps repeated batch indices using out_idx/(out_dim/in_dim) when out_dim is
divisible by in_dim, but GGML_REPEAT semantics map via out_idx % in_dim. With the new advanced shape
acceptance (e.g., A2=4, D2=32) this will select the wrong A/B batch slice during packing and can
produce incorrect MAT_MUL outputs.

Code

ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[R3730-3733]

+    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;

Relevance

●●● Strong

Likely real correctness bug in matmul broadcast/repeat mapping; team usually accepts
offload-safety/correctness fixes.

PR-#154
PR-#152

ⓘ Recommendations generated based on similar findings in past PRs

Evidence

GGML’s repeat implementation indexes the source by the inner index k while output index advances
as i*in_dim + k, which corresponds to out_idx % in_dim. The PR’s new mapping instead uses
out_idx / (out_dim/in_dim) for divisible dims, and the PR also loosens advanced-shape validation
to allow divisible batch dims specifically relying on this new mapping, so the offload can compute
with the wrong slices.

ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[3719-3739]
ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[3990-4001]
ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[5318-5327]
ggml/src/ggml-cpu/ops.cpp[1651-1672]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
`map_repeat_dim_i64()` currently uses grouped division when `out_dim % in_dim == 0`, which conflicts with GGML `GGML_OP_REPEAT`’s modulo-based indexing for repeated dimensions. This causes incorrect batch-slice selection for divisible-repeat MAT_MUL layouts newly admitted by `tsavorite_mul_mat_advanced_shape_ok()`.

### Issue Context
GGML repeat writes output indices as `(i*in_dim + k)` sourced from `k`, which implies `k = out_idx % in_dim`. The Tsavorite packing loops now use `map_repeat_dim_i64()` for A/B batch pointer selection, so the mapping must match GGML semantics.

### Fix Focus Areas
- ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[3725-3739]
- ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[3990-4001]
- ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[5318-5327]
- ggml/src/ggml-tsavorite/ggml-tsavorite.cpp[5459-5468]

### Suggested change
- For divisible repeats, map with modulo (same as `map_repeat_i64(out_idx, in_dim)`) rather than `out_idx / group`.
- If “grouped” mapping is truly required for a specific view/reshape-produced layout, introduce an explicit opt-in path and only use it when you can prove (via tensor stride/layout invariants) that the grouped interpretation matches the tensor’s memory layout.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

Context used

✅ Compliance rules (platform): 6 rules

_{ⓘ  1 issues published inline · 1 in summary — for nitpicking, see all findings}

_{To customize comments, go to the Qodo configuration screen, or learn more in the docs.}

[cubic-dev-ai]

All reported issues were addressed across 2 files

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

All reported issues were addressed across 2 files (changes from recent commits).

<sub>Reply with feedback, questions, or to request a fix.

Fix all with cubic | Re-trigger cubic</sub>

[atrivedi-tsavoritesi]

Approving it but I do have a question, does the 16GB need come from a specific model ? As I am not seeing a need to go beyond 8GB in most of the models I have run.

[atrivedi-tsavoritesi]

Also since we will be having a new build will be good if you can intercept the new build with a new SDK and update ollama if you can. Ccing @nusmangit in case you want to tag taos-app with new tag. This SDK build is already done.

[cubic-dev-ai]

All reported issues were addressed across 1 file (changes from recent commits).

<sub>Tip: Review your code locally with the cubic CLI to iterate faster.

Fix all with cubic | Re-trigger cubic</sub>

[akapoor3518]

Approving it but I do have a question, does the 16GB need come from a specific model ? As I am not seeing a need to go beyond 8GB in most of the models I have run.

I initially increased it to 16GB as a precaution for mixed-precision support, where quantized weights may be expanded to F32 for compute. I did see one memory-related issue that went away at 16GB, but I did not fully analyze the root cause. Since there is no clear evidence that 16GB is generally needed, I reverted the default back to 8GB and will revisit if we see memory-related failures.