@FIR-2172 - Add runtime-controlled broadcast Triton MAT_MUL offload support

Merged
AK
3mo ago
+96-12

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.