millaguie/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

Compare commits

base: millaguie/llama.cpp:470840e01eb1249d9c1907496bc1f9e46798435d
Branches Tags
millaguie/llama.cpp:master millaguie/llama.cpp:dt3
..
compare: millaguie/llama.cpp:dt3
Branches Tags
millaguie/llama.cpp:dt3 millaguie/llama.cpp:master
35 Commits
470840e01e .. dt3
Author SHA1 Message Date
Millaguie e9616d415f cuda : gate DT3 MMQ by batch size, cuBLAS keeps the large-batch prefill 
Two integer dot products per weight cancel the 2x int8-over-fp16 tensor
core advantage, so at large batch MMQ cannot beat dequantize + fp16
cuBLAS (measured 331 vs 426 t/s pp512 on the 27B, RTX 4060 Ti, while
running at the same ~20% of its int8 ceiling as Q2_K MMQ does of its
own). MMQ still avoids the dequantization round-trip at moderate batch;
the threshold default is provisional until the crossover is measured
(GGML_CUDA_DT3_MMQ_MAX_BATCH overrides it for that measurement).
 2026-08-11 08:51:29 +02:00
Millaguie 981f439ff3 tests : judge the batched DT3 path strictly when it is MMQ
Python Type-Check / python type-check (push) Canceled after 0s
Details 
With MMQ, ncols_dst > 8 is an integer path in the same numerical regime
as MMVQ and is judged against the exact reference at 1e-5 instead of
riding the loose GEMM gate. The regime is told apart by the result
itself; a GEMM fallback (backends without DT3 MMQ) keeps the fp16
reference and the F16 GEMM bit-identity control. Also add ncols_dst=100
to exercise a wide tile with a clamped last column block.
 2026-08-11 08:33:06 +02:00
Millaguie f46e8072d8 cuda : add MMQ kernel for DT3 
Two decoded ternary planes per SRAM tile row (the planes cannot be fused
into one int8 because d1 != d2), each with its own per-chunk scales, both
multiplied against the same q8_1 y tile: sum = dB*(sumi1*dA1 + sumi2*dA2).
The load decodes each packed byte once with the same base-3 digit
iteration as the MMVQ vec_dot and turns digit bytes {0,1,2} into trit
bytes {-1,0,+1} without cross-byte borrows.

The MMA tile at I=128 takes 592 B/row: 75776 B of x tile plus the y tile,
94720 B at J=128 — fits the 99 KiB opt-in limit of Ampere-class devices
but not e.g. Turing's 64 KiB, so the runtime gate requires the MMA data
layout and enough shared memory for the narrowest tile and declines
otherwise (AMD keeps declining: no config entries select DT3).
 2026-08-11 08:33:06 +02:00
Millaguie 9622c56b0e tests : accept either accumulator precision in the DT3 GEMM bit-identity gate 
The Vulkan backend now forces fp32 accumulators for the DT3 dequant
fallback, so the DT3 GEMM is no longer bit-identical to the backend's
default-precision F16 GEMM (fp16 accumulators on fp16-capable Vulkan
devices). Run the F16 control at both the default and the F32-forced
precision and require bit-identity with either one. CUDA still matches
the default-precision control; Vulkan matches the F32 one - both with
0 mismatches.
 2026-08-11 03:31:42 +02:00
Millaguie ad6dd747d4 vulkan : fp32 accumulators for the DT3 dequant matmul fallback 
The dequant fallback runs DT3 matmuls as fp16 weights through the f16
matmul pipelines, which default to fp16 accumulators when the device
supports fp16. DT3 weights are the sum of two fp16-scaled ternary
planes and are generally not fp16-representable, so the fallback
already pays one fp16 rounding on the weights; accumulating on top of
that in fp16 measurably hurts.

Force GGML_PREC_F32 for the DT3 fallback, which selects the f32acc
pipelines - the same numerics as the CUDA GEMM fallback (fp16 inputs,
fp32 compute).

Measured on Qwen2.5-3B DT3, wiki.test, 4 chunks, --no-mmap, AMD
Radeon 890M (RADV STRIX1), CPU reference 15.0608:

  default batch, before   15.5562  (+3.29%)
  default batch, after    15.4062  (+2.29%)
  GGML_VK_DISABLE_F16=1   15.3348  (+1.82%, floor of this route)

The remaining gap is shared with the mul_mat_vec route (15.3302, which
does not move under GGML_VK_DISABLE_F16) and is under investigation
separately; DT3 dequantization itself is bit-exact vs the CPU
reference on real model tensors (214,695,936 elements, 0 mismatches).
 2026-08-11 03:01:09 +02:00
Millaguie 8ba4db150f vulkan : add DT3 dequant, get_rows and scalar mul_mat_vec 
Wires the dual-plane ternary type into the Vulkan backend through three
paths only:

- get_rows and the generic scalar mul_mat_vec use per-element decode in
  dequant_funcs.glsl: the byte and base-3 digit are located from the
  element index (regions qs[0..16), qs[16..24), qh[0..2)), the byte is
  multiplied by 3^n mod 256 and the top digit taken. w = d1*t1 + d2*t2
  is accumulated in fp32; both products are exact so the sum carries a
  single float rounding and reproduces the CPU reference bit by bit
  (verified: 0 mismatches over the 112-block synthetic test and over
  214,695,936 elements of real model tensors).
- larger matmuls fall back to dequant_dt3.comp (decode each byte once
  with q <- q*3 mod 256, fp32 sum, one rounding at the f16 write) plus
  the existing f16 matmul pipelines.

The qh bytes hold only 4 trits; their 5th base-3 digit is packing
padding that decodes to -1, so both decoders stop at 4 digits.

Deliberately NOT implemented, and declined instead of half-supported:
no coopmat/coopmat2/MMQ shaders are generated for DT3, and supports_op
answers false for MUL_MAT_ID (mul_mat_vec_id shaders are not generated
either). GET_ROWS and MUL_MAT answer true.

test-dt3-gpu accepts the Vulkan backend (and IGPU-type devices) and
passes on RADV STRIX1: dequant 0 mismatches, mul_mat n<=8 norm rel err
~1e-8, GEMM fallback bit-identical to an F16 GEMM on fp16-rounded
weights.
 2026-08-11 03:00:47 +02:00
Millaguie c01c26b56e tests : skip test-dt3-gpu on backends that do not implement DT3
Python Type-Check / python type-check (push) Canceled after 0s
Details 
The test picked the first GPU device it found and treated an unsupported
op as a failure. Vulkan and SYCL answer supports_op == false for DT3,
which is the right answer for them and not a bug to report, so the test
went red on machines that were behaving correctly. Metal is worse: it
answers true for almost any type but has no DT3 shader, so the run died
in pipeline compilation halfway through.

Pick the backend by name instead — CUDA and HIP (which reports itself as
ROCm) are the only ones implementing DT3 — and skip everything else. A
supports_op failure on those two is still a real failure.

Also document that the n <= 8 gating mirrors MMVQ_MAX_BATCH_SIZE by hand
and goes stale silently if the MMVQ dispatch changes.
 2026-08-10 23:33:14 +02:00
Millaguie 10e1fe3d3c cuda : decode DT3 bytes once in the MMVQ vec_dot 
The old vec_dot decoded every element with its own pair of multiplications
(256 inlined get_trit per block, each qs byte re-read 5 times). Decode each
byte once instead, iterating q -> (q*3) & 0xFF two bytes at a time in 16-bit
lanes, and accumulate dp4a over base-3 digits in {0, 1, 2}; one extra dp4a
with 0x01010101 per q8_1 int, shared by both planes, turns the digit sums
back into trit sums in exact integer arithmetic, so the result stays
bit-identical to the per-trit decode. The qh bytes keep their own 4-digit
path so the padding digit is never decoded.
 2026-08-10 23:33:14 +02:00
Millaguie 0f33afbe56 tests : declare the generic DT3 vec_dot weak in the parity test 
Builds without a native DT3 kernel rename the generic symbol to
ggml_vec_dot_dt3_q8_0 (arch-fallback.h), so test-dt3 failed to link on
them. With a weak declaration the test links everywhere and skips,
loudly, when there is no separate generic to compare against. MSVC has
no weak symbols, so there the test is compiled out.
 2026-08-10 23:33:14 +02:00
Millaguie 4e109bc7e6 tests : check the arch DT3 vec_dot is bit-identical to the generic 
Calls the actual ggml_vec_dot_dt3_q8_0_generic symbol against the
dispatched vec_dot and requires memcmp-equal floats. Blocks exercise
the three regions, the 79/80 and 119/120 boundaries, non-trivial qh
bytes (would expose a vectorization reading their padding 5th digit),
and scales of both and mixed signs; y reaches the full q8_0 range.

Mutation-checked: flipping one bit of a digit blend mask in the
AVX-512 kernel makes 94 of the 96 reps fail.
 2026-08-10 23:33:14 +02:00
Millaguie bf4eca0eb6 ggml : add AVX-512 DT3 vec_dot 
Decode both planes of a block with VBMI byte permutes: the *3 multiply
chain (wrapping, so it commutes with the permutation) is computed once
on the whole 56-byte block, and masked vpermb picks each element's byte
from the chain vector of its base-3 digit. The qh lanes never see 3^4,
which would read the padding 5th digit of the qh bytes. The trits reach
the integer product as xi in {0, 1, 2} via the same avg trick as
tq1_0, with VNNI dpbusd against the q8_0 bytes and sum(y) subtracted.

The per-q8_0-block sums and the float accumulation keep the exact
operation order of the generic implementation, so the result is
bit-identical to it (checked by test-dt3).

2.2x over the (autovectorized) generic on a Ryzen AI 9 HX 370.
 2026-08-10 23:33:14 +02:00
Millaguie b285eb8a4f tests : gate the DT3 GEMM fallback by bit-identity with an F16 GEMM 
The dequantize + cuBLAS fallback computes in fp16 (CUBLAS_COMPUTE_16F)
on fast-fp16 hardware and in TF32 under
GGML_CUDA_CUBLAS_COMPUTE_TYPE=f32, so no analytic tolerance separates
'correct' from 'broken' there without also tracking cuBLAS numerics.
What IS ours to guarantee: the fallback must behave exactly as if the
weights were an F16 tensor holding fp16(dequant(block)). Gate on that
bit-identity and demote the analytic GEMM errors to INFO.
 2026-08-10 23:33:14 +02:00
Millaguie e5c6656dbf tests : judge DT3 mul_mat on norm error, add fp16 reference and controls 
Elementwise max relative error explodes on cancellation whenever a true
output element is near zero, so the mul_mat checks now gate on the
relative Frobenius norm and keep the max as information. The GEMM
fallback dequantizes to fp16 on fast-fp16 hardware and DT3 weights
(d1*t1 + d2*t2) are generally not fp16-representable, so that path is
judged against a reference computed from fp16-rounded weights (taking
the better of both references so GGML_CUDA_CUBLAS_COMPUTE_TYPE=f32 also
passes). Q4_1 (same non-fp16-exact regime) and Q4_0 (fp16-exact weights,
pure GEMM error floor) run through the identical comparison as controls.
 2026-08-10 23:33:14 +02:00
Millaguie 659b1ace9e tests : add DT3 GPU vs CPU parity test 
Checks the GPU backend against the validated CPU path: GET_ROWS
dequantization must match dequantize_row_dt3 bit by bit on directed
blocks (region boundaries 79/80 and 119/120, qh elements, negative
scales) and on raw random bytes; MUL_MAT must match a double precision
reference from the dequantized weights, with activations whose q8_1
quantization is exact, plus a manual trit-sum check with non-trivial qh.
Skips cleanly when no GPU backend is available.
 2026-08-10 23:33:14 +02:00
Millaguie 0cc5e310c1 cuda : add DT3 MMVQ kernel 
vec_dot_dt3_q8_1 processes a whole 128-element DT3 block per call
(VDR_DT3_Q8_1_MMVQ = 4, QI_DT3 = 4), i.e. the 4 q8_1 chunks it spans,
with one pair of integer accumulators per chunk:

    sum_j d8[j] * (d1*sumi1[j] + d2*sumi2[j])

The trit decode reuses ggml_cuda_dt3_get_trit with fully unrolled loops,
so all indices and pow3 factors fold into constants; no __byte_perm or
other NVIDIA-only intrinsics. Enables MUL_MAT in supports_op: ncols_dst
<= 8 takes MMVQ, larger falls back to dequantization + cuBLAS (no MMQ
tile kernel yet).
 2026-08-10 23:33:14 +02:00
Millaguie 1a1f869f93 cuda : add DT3 dequantization 
Decode one packed ternary plane with the shared ggml_cuda_dt3_get_trit
helper (shifts, masks and a small pow3 table; the uint8_t wrap-around of
the intermediate product is intentional and matches the CPU reference).
The qh bytes hold only 4 trits; their 5th base-3 digit is packing padding
that always decodes to -1 and is never read.

Wires DT3 into the generic dequantize_block templates (to fp32/fp16/bf16,
contiguous and not) and into get_rows, and enables GET_ROWS in
supports_op.
 2026-08-10 23:33:14 +02:00
Millaguie a277f4c6f1 tests : check DT3 byte positions against hand-computed literals 
The previous byte-position test packed with the test's own packer on
both sides of the comparison, so it exercised none of the library code.
It now pins hand-computed byte values (43/100/127/42/124...) at the
region boundaries (79/80, 119/120) as ground truth and drives both
directions through the library: to_float must place each literal byte's
trit at the exact element, and from_float must produce the exact literal
byte, for both planes.

Also probes ggml_validate_row_data over all 256 byte values in qs and
qh positions (must accept exactly the 243/81 reachable codes), the
all-0xaa block, and a well-formed packed block.
 2026-08-10 23:33:14 +02:00
Millaguie 6d6552c862 llama : warn when quantizing to DT3 
The reference-quantizer disclaimer only existed in the code and in
llama-quantize --help; now it is also printed where the mistake would
actually be made, at the start of a quantization run targeting DT3.
 2026-08-10 23:33:14 +02:00
Millaguie 4efd061d43 ggml : harden DT3 validation and reference quantizer 
ggml_validate_row_data now rejects unreachable code bytes: the ceiling
division packing reaches only 243 of the 256 byte values in qs and 81
in qh (4 trits plus an always-zero padding digit), so corruption that
previously loaded and generated garbage silently is caught at load
time. Previously only the two fp16 scales were checked.

quantize_dt3 no longer discards quant_weights silently: an ignored
imatrix now prints a loud warning (once), otherwise an imatrix A/B on
DT3 would come out byte-identical and invite the false conclusion that
the imatrix does nothing.

The two initial trit passes of quantize_row_dt3_ref now clamp like the
refit passes do, so a NaN input cannot push an out-of-range value from
lroundf into the packer.
 2026-08-10 23:33:14 +02:00
Millaguie b300ba053d tests : add bit-level DT3 tests and Rust parity driver 
test-dt3 checks the layout against an independent packer written from
the format spec: single-trit position mapping for all 256 (plane, pos)
pairs, structural byte-position checks at the region boundaries
(79/80, 119/120), exact round-trips with negative scales, byte parity
of the in-tree quantizer on already-ternary inputs, and the vec_dot
against a hand-made sum over the known trits (catches any path that
reads the padding 5th trit of the qh bytes).

test-dt3-rust-parity.py packs known trits with ternaria's pack_dt3 and
verifies that dequantize_row_dt3 (via test-dt3 --dequant) reproduces
d1*t1 + d2*t2 bit-exactly.

Also wires DT3 into the test-quantize-fns thresholds (ternary class).
 2026-08-10 23:33:14 +02:00
Millaguie 985b0ecba2 gguf-py : add DT3 
Registers the type id, file type and block size, and implements numpy
dequantization (verified bit-exact against the C implementation with
gguf-py/tests/test_quants.py, including random byte payloads).

Quantization is intentionally left unimplemented, like the K-quants:
DT3 planes come from an external solver (PTQTP) and are packed
directly, so a from-float numpy path would only invite quantizing
models with the wrong algorithm.
 2026-08-10 23:33:14 +02:00
Millaguie 693eb7d719 llama : register the DT3 file type 
Adds LLAMA_FTYPE_MOSTLY_DT3 at the end of the ftype enum, the loader
name/guess mappings, the quantization fallbacks (same as the other
ternary types), and the llama-quantize table entry. The table entry
warns that the in-tree quantizer is only the reference one: DT3 models
with the measured quality are produced by the external PTQTP pipeline.
 2026-08-10 23:33:14 +02:00
Millaguie 698f37f40b ggml-cpu : add DT3 generic vec_dot and type traits 
The vec_dot pairs DT3 with Q8_0 (4 q8_0 blocks per DT3 block) and keeps
one integer accumulator per plane: sumf += dy * (d1*sumi1 + d2*sumi2).
Q8_0 instead of Q8_K on purpose: the planes are symmetric ternary so the
q8_K bsums are dead weight, and 32-element blocks accept any row size
that is a multiple of 128.

Trit decoding reuses unpack_plane_dt3, which reads only 4 trits per qh
byte; the 5th base-3 digit of those bytes is packer padding that always
decodes to -1 and must never be read.
 2026-08-10 23:33:14 +02:00
Millaguie 160ea6c428 ggml : add DT3 reference quantization and dequantization 
Add the dual-plane ternary DT3 type to the type registry along with its
reference row functions. Each of the two planes is packed exactly like
tq1_0 with all constants halved (block of 128 elements): qs 48 -> 24
bytes over two passes of 16 and 8 bytes, qh 4 -> 2 bytes.

The trit decoding lives in a single exported helper (unpack_plane_dt3)
so that dequantization and the upcoming CPU vec_dot share it.

The reference quantizer is a greedy two-pass (plane 1 by absolute max,
plane 2 on the residual) plus two rounds of alternating least-squares
refits. It is intentionally NOT the PTQTP solver used to produce the
published DT3 models.
 2026-08-10 23:33:14 +02:00
Millaguie 5c175d940f ggml: add block_dt3, the dual-plane ternary block 
DT3 stores w_i = d[0]*t0_i + d[1]*t1_i with t in {-1,0,+1}, two ternary
planes over a 128-element block: 56 bytes, 3.5 bpw exactly.

Each plane uses the tq1_0 base-3 packing with every constant halved for
the smaller block (qs 48->24 B, qh 4->2 B, qs passes over 16 then 8
bytes instead of 32 then 16), which tiles 128 with no leftover bytes.
Reducing tq1_0 to 128 without halving the passes does not tile: with a
24-byte qs the first pass covers nothing and the second overruns.
 2026-08-10 23:33:13 +02:00
Gaurav Garg 030ebb558a Address review comment of PR 25532 (#26852) 2026-08-11 00:02:25 +05:30
Hongqiang Wang 689e227db4 opencl: transpose the K tile in local memory for FA prefill kernels (#26428) 2026-08-10 11:09:19 -07:00
Mario Limonciello 0666ad2b2b ci : target ROCm 7.14 for build and release (#25775) 
* Switch ROCm from 7.2.1 to 7.14

ROCm 7.14 is the first production release using TheRock build system.
It can be installed using multi-arch deliverables from wheels, debs,
rpms, tarballs or runfiles.

Adjust ROCm targets for Linux and Windows to use this instead.

* ci: switch all other Windows ROCm jobs to ROCm 7.14 wheels

Move the shared windows-setup-rocm composite action from the HIP SDK PRO
Edition installer to the multi-arch ROCm wheels (rocm[libraries,devel]).
The wheel-install logic that previously lived inline in release.yml is now
in the shared action, and both build-cache.yml and release.yml call it.

Also migrate the build-cuda-windows.yml hip job to the same wheel-based
layout (cache path/key, rocm-sdk environment setup, llvm/bin compiler
paths) so it keeps working after the action's contract changed; drop its
now-unused ROCm 7.2.1 rocWMMA download and stale include path.
 2026-08-10 19:53:12 +02:00
Gaurav GargandGeorgi Gerganov dd1ea52433 llama : support multi-output backend sampling (#25532) 
* Enable backend sampling with token speculation

* Clamp the mask sum before converting it into the sampled index

* Add a numeric context parameter declaring the maximum outputs one sequence

* More fixes

* Don't reuse memory for output views.

* Match dist between CPU and GPU

* Fix CPU and backend sampling mismatches

* Simpify some of the changes

* Fix tests on Vulkan

* More test fixes

* Rebase changes

* Rebase and address review comments

* Address review comments

* Address review comments

* Update src/llama-sampler.cpp

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
 2026-08-10 16:58:56 +03:00
Hitesh Chopra d2f83055d6 ggml-cpu : fix CPU affinity mask being ignored on Android (#26838) 2026-08-10 15:13:40 +03:00
Yash Raj Pandey f8def7fe16 ggml : require contiguous src for ROLL on CUDA and Metal (#25928) 
ggml_roll only asserts nb[0] == ggml_type_size, so a permuted src is a
valid input, but the CUDA and Metal roll kernels index by ne alone and
never read the nb strides. A non-contiguous src therefore produced
silently wrong results. Neither backend declared a contiguity
requirement in supports_op, so the scheduler did not fall back to the
CPU implementation, which does handle strides correctly.

Add the requirement to both backends, matching the existing
GGML_OP_ROPE guard, and add a permuted test_roll case.
 2026-08-10 15:01:44 +03:00
Pascal 4dee52f82d ui: UI/chat form follow ups (#26743) 
* ui: split the markdown rendering setting per surface

User content and thinking get their own toggle again, so turning off
markdown for a message leaves reasoning blocks formatted. Both default
to markdown. A stored renderContentAsRawText unfolds onto the user key
and is dropped from the config.

File mentions render as badges in the raw text path too, through a
narrow pass over [name](file://path) that leaves everything else
untouched.

* ui: let the rich chat input scroll past its max height

The contenteditable renderer caps its height with max-height but had no
overflow rule, so a long buffer overflowed into the input area wrapper
and got clipped by its overflow-hidden, leaving no way to reach the
bottom of the message. The textarea renderer scrolls natively and was
never affected.

* ui: apply the new lint and format config

* ui: move the render keys unfolding into the migration service

Address review from @allozaur: the settings store no longer rewrites
persisted config on load, the raw text toggle now unfolds onto the
per-surface render keys in migration.service.ts, next to the other
config migrations. The mention scanner flag and the directory path
suffix become named constants.
 2026-08-10 13:32:51 +02:00
Sigbjørn Skjæret e5275f6f77 ci : don't specify python version in server-sanitize for broader runner compatibility (#26840) 
* don't specify python version for broader runner compatibilty

* run the workflow
 2026-08-10 13:32:22 +02:00
PascalandXuan Son Nguyen 4ae84dea27 server: add more tool isolation support (ssh remote + podman rootless) (#26774) 
* server: add an ssh transport to the tools runtime

--tools-runtime ssh:<target> runs the built-in tools on a remote host,
where target is whatever ssh already resolves, a user@host or a config
alias, so no credentials live in llama.cpp.

Only build_argv and upload differ from the docker transport: the remote
shell re-parses the command line, so the argv travels through
shell_quote_join, and files go over scp with the same quoting on the
remote path. Authentication is key-based and the host key must already
be trusted, since the tools run without a console and any prompt would
hang them.

The target is validated before use. The spec can reach us from the
x-tool-runtime header, and a leading dash would turn it into an ssh
option, which is enough to run a command back on the host.

Nothing is created and nothing is reclaimed, so an ssh spec goes
straight to the tool call instead of through the container runtime.

Note that this is remoting rather than isolation: the tools can do
whatever the target account can do, and the isolation is whatever runs
them on the far side.

* server: support podman in the tools runtime

docker and podman expose the same run, exec, cp and inspect verbs with the
same argument order, so a single implementation drives both and the engine
is carried by the spec prefix: podman:<image> and podman-container:<id> sit
next to the docker forms.

tools_io_docker becomes tools_io_container and the runtime spawner becomes
server_tools_container_runtime, both holding the client binary chosen at
parse time. A single parse_container_runtime() resolves every spec, so
adding another engine is one string in the table.

make_tools_io() now rejects the spawning forms. The spec also reaches it
from the x-tool-runtime header, which is client controlled, and only the
runtime that owns a container is allowed to create one: a tool call can
attach to a running container, nothing more.

* ./build/bin/llama-gen-docs

* server: simplify the tools runtime and drop the file copy step

A server_tools_runtime base with one virtual spec() replaces the
container runtime and the bare spec string that ssh needed next to it,
so server_tools is back to a single pointer and neither setup nor the
handler tests which of the two is set.

write_file used to spill its content into a temporary file on the host
and copy it in, because run_subprocess had no way to feed a child. It
now takes an optional stdin payload and creates the parent directory
and the file in a single round trip through a shell in the isolate.

That removes the upload virtual and both implementations: no more
container cp or scp, no second binary on the host, no sftp subsystem on
the target, no predictable temporary in a shared tmp, and none of the
content reaching an argv the remote shell re-parses. It also fixes
write_file over ssh, which never worked: scp speaks sftp and takes the
remote path literally, so quoting it kept the quotes in the file name.

Writing the payload before reading the output relies on the child
draining stdin as it goes, which holds for cat, its only user today.

* ./build/bin/llama-gen-docs

* server: harden the tools runtime against argv injection and a stdin stall

Validate the container id from x-tool-runtime and --tools-runtime the
same way the ssh target already is, so an id shaped like an option
(docker-container:--privileged) is rejected before it reaches the
engine's exec command line instead of running against a hardened
container. Feed the child's stdin after the watchdog is armed, so a
transport that stalls mid-write is terminated at the deadline rather
than blocking the request forever.

Cover both guards and fix the unknown-scheme test, which used ssh: as
its example and now names a real runtime.

* tests: exercise the tools runtime tests on podman as well as docker

Follow-up #26507. The container runtime drives docker and podman
through one implementation, so parametrize the availability helper,
the container fixture and the attach test on the engine, and cover
both engine prefixes in the container id injection test. Each engine
skips on its own when it is not installed.

The spawn cleanup test stays docker only: it recovers the spawned id
from the container hostname, which docker sets to the short id and
podman rootless does not guarantee. Podman keeps its coverage through
the attach path.

* server: release the container handle before respawning

Follow-up #26507. create() writes over the handle it is given, so a
respawn after the container died on its own leaked the pipes and the
process handle of the previous one.

* server: trim the tools runtime comments

* server: read tool output as raw bytes and harden the runtime on Windows

The stdout pipe is read with read() instead of fgets(), so a chunk
can hold any byte, including NUL, and still streams as soon as data
is available. Past the size cap the pipe keeps draining so the child
never blocks on a full pipe. Both pipe fds are forced to binary mode
on Windows, where the CRT defaults them to text mode and translates
line endings in both directions. Stdin is now always closed after
the feed: the child reads a deterministic EOF, and the Windows
docker and ssh clients stop outliving their command on a stdin pipe
that never closes.

The attach form of --tools-runtime has no lifecycle to own, so it
becomes a static target validated once at startup. This removes the
 subprocess that ran on every tool call and
serialized calls behind a mutex; a stopped container now surfaces
the engine's own error at exec time.

The cidfile path is passed as UTF-8, matching the encoding the
subprocess layer expects for the CreateProcessW command line, so
the spawn form works from a non-ASCII Windows profile.

The SIGPIPE note in server.cpp now names the tools runtime children
as well as the MCP ones.

* clean up comments

* less pollute global scope

* nits

* tests: name the container image after both engines

---------

Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
 2026-08-10 13:31:09 +02:00
Pedro CuencaYoung HanBeto de PaolaDaniel HanruanrmsXuan Son NguyenSigbjørn Skjæret
 		62bf73d25c model: Muse Glimmer Support (#26841) 
* Get started with Onyx

* Add architecture

* Skip keys handled in super()

* Loading tensors

* Shorten

* Graph

* Apply suggestion from @pcuenca

* Remove norm now embedding in transformers weights

* Add eot

* Explicit output_multiplier

* Handle post_norm_eps

* No super call; unhardcode eot.

The pattern `self._set_vocab_gpt2()` seems preferred throughout the
codebase, and it allows `set_vocab()` to be called from a different part
of the Python class hierarchy: the drafter model converter that we may
need eventually.

* Register for drafting

* DFlash: inherit rope type from the linked target.

Another option would be to store it in the gguf file itself.

* mmproj conversion

Note: some fields to be renamed after the implementation works. We are
keeping compatibility with the reference Meta gguf for testing purposes.

* "clip" header declarations

* Load mmproj

* Pre-processing

* Graph

* Go back to using delimiters.

Otherwise our generations are worse.

Transformers does not use them. We need to trace inputs to verify
whether they are equivalent.

* downsample_factor -> merge_size

* Add vision graph

lol, forgot from a previous commit

* Additional renames, align with llama.cpp / transformers

* Prefer _size instead of independent _h and _w

* Fix token layout

Co-authored-by: Young Han <younghan@fb.com>

* onyx: bring the chat parser onto the onyx branch

common/chat.cpp on this branch has no Onyx handling, so a converted model
serves malformed chat: the assistant preamble leaks into content
("to=self<|message|>...") and tool calls fail with

    HTTP 500 "The model produced output that does not match the expected
              peg-native format"

common_chat_params_init_onyx exists on onyx-fair-patch, added there by
8bb73dd3d. It was never on this branch, so this is not a regression --
the two lines developed independently.

The code here is taken verbatim from that commit. It is the clean side of
`git merge origin/onyx-fair-patch`: chat.cpp is one of the files that
merges without conflict. The full merge is not viable -- it produces 13
conflicts, including add/add on conversion/onyx.py and src/models/onyx.cpp
where the q_norm-folding and metadata-scale approaches contradict each
other, and #4/#7 are stacked on this branch's side of that.

Verified on this branch: builds with 0 errors, converts an Onyx checkpoint,
and serving it gives "4" for "What is 2+2?" plus a correct
get_weather {"city":"Paris"} tool call, where the unported branch gives the
two failures above.

No converter or runtime changes are included, so this should not interact
with the q_norm work.

Co-authored-by: Beto de Paola <betodepaola@meta.com>

* Less params, bilinear pos-emb interpolation as a graph op instead of CPU

* Map to symbolic V_MMPROJ instead of strings

* Make a couple params explicit

* Patchify via build_inp()

* No param for rope_theta

* Small cleanup

* Restore blank line

* Unpermute, to adapt to the latest transformers checkpoint

* Apply norm after token embeddings

This follows the latest transformers approach.

* Remove duplicated function

* build_vit

* onyx: use the model rope theta on sliding-window layers

* DFlash: conversion from transformers drafter

* Revert rope_type derivation from target

NOTE: this breaks compatibility with Meta's distributed DFlash GGUFs, as
the Q/K are stored in "NEOX" (rotated half) format, like in
transformers.

* Apply suggestion from @pcuenca

* Set model type

* Remove comment that will become obsolete

* Hardcode post_norm_rms_eps instead of new param

* Derive SWA+RoPE pattern from gguf array or scalar

* Fix model type <-> number of layers

* Reorder

* Rename

* Fix typo

* DFlash: seed the draft KV cache from multimodal embedding batches

`common_speculative_impl_draft_dflash::process()` returned early on any batch carrying embeddings, so an image prefill never had its target-layer features fused through the DFlash encoder and injected into the draft's KV cache. That left a hole spanning the image's positions, and the next injection at a post-image position failed to initialize its batch:

```
decoding image batch 1/1, n_tokens_batch = 256
decode: failed to initialize batch
llama_decode: failed to decode, ret = -1
process: llama_decode(ctx_dft) failed rc=-1 (n_tokens=17, offset=0)
srv decode: failed to process speculative batch
```

Every image request with `--spec-type draft-dflash` failed with HTTP 500. Text-only was unaffected, since those batches carry token ids and were let through.

Restore the earlier condition, which admits a batch that is either tokens or embeddings and skips only the degenerate neither/both cases. The rest of `process()` is already layout-agnostic -- it gathers features via `llama_get_embeddings_layer_inp()` and indexes `batch_in.pos[]` / `batch_in.seq_id[]`, none of which assume token ids -- so this is the whole fix.

Validated against `muse-glimmer-30B-bf16.gguf` + `mmproj-muse-glimmer-30B-bf16.gguf` + a DFlash draft head, on an image describe-the-shapes request:

- before: HTTP 500, `failed to process speculative batch`
- after: HTTP 200, draft acceptance 0.34012 (167 accepted / 491 generated), mean len 3.04

Output equivalence holds, which is the property that matters: at temperature 0 the drafted response is byte-identical to the same request served with no draft attached (1213/1213 chars), so the draft is drafting correctly through the image context rather than merely not crashing.

* Conversion: prefer rewrite to mapping

* Revert "Conversion: prefer rewrite to mapping"

This reverts commit a92d0ac584d315e876741e85b6dad3dbc8b23bf7.

* fix lint

* sliding_window metadata is not optional

* disable state save/load

* Apply suggestion from @pcuenca

---------

Co-authored-by: Young Han <younghan@fb.com>
Co-authored-by: Beto de Paola <betodepaola@meta.com>
Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
Co-authored-by: ruanrms <ruanslv@gmail.com>
Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
 2026-08-10 13:07:27 +02:00
98 changed files with 3782 additions and 979 deletions
Expand all files Collapse all files
.github/actions/windows-setup-rocm/action.yml
+23 -5
View File
@@ -8,8 +8,26 @@ inputs:
runs:
using: "composite"
steps:
- name: Setup ROCm
uses: ./.github/actions/install-exe
with:
url: https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ inputs.version }}-Win11-For-HIP.exe
args: -install
- name: Install ROCm with Wheels
shell: pwsh
run: |
$ErrorActionPreference = "Stop"
write-host "Setting up Python virtual environment"
# Create the venv directly at the cache location to avoid relocation issues
New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null
python -m venv C:\TheRock\build\.venv
& C:\TheRock\build\.venv\Scripts\Activate.ps1
write-host "Upgrading pip"
python -m pip install --upgrade pip
write-host "Installing ROCm wheels for multi-arch support"
# Install ROCm wheels for multi-arch support (this may take several minutes)
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ inputs.version }}"
# Pre-expand the devel tree so it is included in the cache
write-host "Initializing ROCm devel tree"
rocm-sdk init
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
write-host "Completed ROCm wheel installation to C:\TheRock\build"
.github/workflows/build-cache.yml
+5 -5
View File
@@ -123,8 +123,8 @@ jobs:
runs-on: windows-2022
env:
# Make sure this is in sync with build.yml
HIPSDK_INSTALLER_VERSION: "26.Q1"
# Make sure this is in sync with release.yml and build-cuda-windows.yml
ROCM_VERSION: "7.14.0"
steps:
- name: Clone
@@ -135,11 +135,11 @@ jobs:
uses: actions/cache@v5
id: cache-rocm
with:
path: C:\Program Files\AMD\ROCm
key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }}
path: C:\TheRock\build
key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }}
- name: Setup ROCm
if: steps.cache-rocm.outputs.cache-hit != 'true'
uses: ./.github/actions/windows-setup-rocm
with:
version: ${{ env.HIPSDK_INSTALLER_VERSION }}
version: ${{ env.ROCM_VERSION }}
.github/workflows/build-cuda-windows.yml
+46 -30
View File
@@ -83,7 +83,7 @@ jobs:
env:
# Make sure this is in sync with build-cache.yml
HIPSDK_INSTALLER_VERSION: "26.Q1"
ROCM_VERSION: "7.14.0"
strategy:
matrix:
@@ -97,36 +97,53 @@ jobs:
id: checkout
uses: actions/checkout@v6
- name: Grab rocWMMA package
id: grab_rocwmma
run: |
curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb"
7z x rocwmma.deb
7z x data.tar
- name: Use ROCm Installation Cache
- name: Cache ROCm Installation
uses: actions/cache@v5
id: cache-rocm
with:
path: C:\Program Files\AMD\ROCm
key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }}
path: C:\TheRock\build
key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }}
- name: Setup ROCm
if: steps.cache-rocm.outputs.cache-hit != 'true'
uses: ./.github/actions/windows-setup-rocm
with:
version: ${{ env.HIPSDK_INSTALLER_VERSION }}
version: ${{ env.ROCM_VERSION }}
- name: Setup ROCm Environment
run: |
$ErrorActionPreference = "Stop"
# Activate venv from cache or fresh install
& C:\TheRock\build\.venv\Scripts\Activate.ps1
# Expand the devel tree (idempotent; no-op if already done during install)
rocm-sdk init
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
# Get ROCm installation paths using the rocm-sdk CLI tool
$rocmPath = (rocm-sdk path --root)
if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" }
$rocmPath = $rocmPath.Trim()
$cmakePath = (rocm-sdk path --cmake).Trim()
$binPath = (rocm-sdk path --bin).Trim()
write-host "ROCm root: $rocmPath"
echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV
echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV
echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV
echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
echo "$binPath" >> $env:GITHUB_PATH
# Keep venv in PATH for subsequent steps
echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH
- name: Verify ROCm
id: verify
run: |
# Find and test ROCm installation
$clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1
if (-not $clangPath) {
Write-Error "ROCm installation not found"
exit 1
}
& $clangPath.FullName --version
# Test the ROCm clang shipped in the installed wheel
& "${env:HIP_PATH}\lib\llvm\bin\clang.exe" --version
- name: ccache
uses: ggml-org/ccache-action@v1.2.21
@@ -134,28 +151,27 @@ jobs:
# TODO: this build does not match the build in release.yml, so we use a different cache key
# ideally, the builds should match, similar to the CUDA build above so that we would be able
# to populate the ccache for the release with manual runs of this workflow
#key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }}
key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }}
#key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }}
key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }}
- name: Build
id: cmake_build
run: |
$env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path)
$env:CMAKE_PREFIX_PATH="${env:HIP_PATH}"
cmake -G "Unix Makefiles" -B build -S . `
-DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" `
-DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" `
-DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/" `
-DCMAKE_PREFIX_PATH="${env:HIP_PATH}" `
-DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
-DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" `
-DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
-DCMAKE_BUILD_TYPE=Release `
-DLLAMA_BUILD_BORINGSSL=ON `
-DROCM_DIR="${env:HIP_PATH}" `
-DHIP_PATH="${env:HIP_PATH}" `
-DGGML_HIP=ON `
-DGPU_TARGETS="gfx1100" `
-DGPU_TARGETS="gfx1100" `
-DGGML_RPC=ON
cmake --build build -j ${env:NUMBER_OF_PROCESSORS}
- name: ccache-clear
uses: ./.github/actions/ccache-clear
with:
#key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }}
key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }}
#key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }}
key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }}
.github/workflows/release.yml
+159 -158
View File
@@ -748,6 +748,132 @@ jobs:
path: llama-bin-win-cpu-${{ matrix.arch }}.zip
name: llama-bin-win-cpu-${{ matrix.arch }}.zip
windows-rocm:
runs-on: windows-2022
strategy:
matrix:
include:
- ROCM_VERSION: "7.14.0"
gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
build: x64
steps:
- name: Clone
id: checkout
uses: actions/checkout@v6
with:
fetch-depth: 0
- name: ccache
uses: ggml-org/ccache-action@v1.2.21
with:
key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
evict-old-files: 1d
- name: Cache ROCm Installation
id: cache-rocm
uses: actions/cache@v5
with:
path: C:\TheRock\build
key: rocm-wheels-${{ matrix.ROCM_VERSION }}-multi-arch-${{ runner.os }}
- name: Setup ROCm
if: steps.cache-rocm.outputs.cache-hit != 'true'
uses: ./.github/actions/windows-setup-rocm
with:
version: ${{ matrix.ROCM_VERSION }}
- name: Setup ROCm Environment
run: |
$ErrorActionPreference = "Stop"
# Activate venv from cache or fresh install
& C:\TheRock\build\.venv\Scripts\Activate.ps1
# Expand the devel tree (idempotent; no-op if already done during install)
rocm-sdk init
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
# Get ROCm installation paths using the rocm-sdk CLI tool
$rocmPath = (rocm-sdk path --root)
if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" }
$rocmPath = $rocmPath.Trim()
$cmakePath = (rocm-sdk path --cmake).Trim()
$binPath = (rocm-sdk path --bin).Trim()
write-host "ROCm root: $rocmPath"
write-host "CMake path: $cmakePath"
write-host "Bin path: $binPath"
echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV
echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV
echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV
echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
echo "$binPath" >> $env:GITHUB_PATH
# Keep venv in PATH for subsequent steps
echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH
- name: Build
run: |
mkdir build
cd build
cmake .. `
-G "Unix Makefiles" `
-DCMAKE_PREFIX_PATH="${env:HIP_PATH}" `
-DCMAKE_BUILD_TYPE=Release `
-DGGML_BACKEND_DL=ON `
-DGGML_NATIVE=OFF `
-DGGML_CPU=ON `
-DGGML_CPU_ALL_VARIANTS=ON `
-DGGML_HIP=ON `
-DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
-DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" `
-DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" `
-DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
-DHIP_PATH="${env:HIP_PATH}" `
-DGGML_HIP_ROCWMMA_FATTN=ON `
-DAMDGPU_TARGETS="${{ matrix.gpu_targets }}"
cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS}
- name: ccache-clear
uses: ./.github/actions/ccache-clear
with:
key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
- name: Verify HIP backend was built
run: |
$hipDll = Get-ChildItem -Path build\bin -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue
if (-not $hipDll) {
Write-Host "##[error]ggml-hip*.dll was NOT produced. The HIP backend silently failed to build."
Write-Host "Contents of build\bin:"
Get-ChildItem build\bin | Format-Table -AutoSize
exit 1
}
Write-Host "HIP backend artifact found:"
$hipDll | Format-Table FullName, Length -AutoSize
- name: Determine tag name
id: tag
uses: ./.github/actions/get-tag-name
- name: Get ROCm short version
run: |
$rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.')
echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV
- name: Pack artifacts
run: |
cp "LICENSE" "build\bin\"
7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\*
- name: Upload artifacts
uses: actions/upload-artifact@v6
with:
path: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip
name: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip
windows:
needs: [check-release]
if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -1168,8 +1294,8 @@ jobs:
strategy:
matrix:
include:
- ROCM_VERSION: "7.2.1"
gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201"
- ROCM_VERSION: "7.14.0"
gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
build: 'x64'
steps:
@@ -1201,38 +1327,36 @@ jobs:
run: |
sudo apt install -y build-essential git cmake wget
- name: Setup Legacy ROCm
if: matrix.ROCM_VERSION == '7.2.1'
id: legacy_env
run: |
sudo mkdir --parents --mode=0755 /etc/apt/keyrings
wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
sudo tee /etc/apt/sources.list.d/rocm.list << EOF
deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} jammy main
EOF
sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF
Package: *
Pin: release o=repo.radeon.com
Pin-Priority: 600
EOF
sudo apt update
sudo apt-get install -y libssl-dev rocm-hip-sdk
- name: Setup TheRock
if: matrix.ROCM_VERSION != '7.2.1'
- name: Setup TheRock with Wheels
id: therock_env
run: |
wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz
mkdir install
tar -xf *.tar.gz -C install
export ROCM_PATH=$(pwd)/install
echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV
echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV
echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV
# Create Python virtual environment
python3 -m venv .venv
source .venv/bin/activate
# Install ROCm wheels for build
# libraries = HIP runtime and CMake configs needed for linking
# devel = compilers, headers, static libs
python -m pip install --upgrade pip
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"
# Get ROCm installation paths using the rocm-sdk CLI tool
ROCM_PATH=$(rocm-sdk path --root)
CMAKE_PATH=$(rocm-sdk path --cmake)
BIN_PATH=$(rocm-sdk path --bin)
echo "ROCM_PATH=$ROCM_PATH"
echo "CMAKE_PATH=$CMAKE_PATH"
echo "BIN_PATH=$BIN_PATH"
# Set environment variables
echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
# Keep venv activated for subsequent steps
echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
- name: Build with native CMake HIP support
id: cmake_build
@@ -1276,129 +1400,6 @@ jobs:
path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
windows-hip:
needs: [check-release, get-version]
if: ${{ needs.check-release.outputs.should_release == 'true' }}
runs-on: windows-2022
permissions:
actions: write
env:
HIPSDK_INSTALLER_VERSION: "26.Q1"
strategy:
matrix:
include:
- name: "radeon"
gpu_targets: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
steps:
- name: Clone
id: checkout
uses: actions/checkout@v6
- name: Setup Node.js
uses: actions/setup-node@v6
with:
node-version: "24"
cache: "npm"
cache-dependency-path: "tools/ui/package-lock.json"
- name: Grab rocWMMA package
id: grab_rocwmma
run: |
curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb"
7z x rocwmma.deb
7z x data.tar
- name: Cache ROCm Installation
id: cache-rocm
uses: actions/cache@v5
with:
path: C:\Program Files\AMD\ROCm
key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }}
- name: ccache
uses: ggml-org/ccache-action@v1.2.21
with:
key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }}
- name: Install ROCm
if: steps.cache-rocm.outputs.cache-hit != 'true'
id: depends
run: |
$ErrorActionPreference = "Stop"
write-host "Downloading AMD HIP SDK Installer"
Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe"
write-host "Installing AMD HIP SDK"
$proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru
$completed = $proc.WaitForExit(600000)
if (-not $completed) {
Write-Error "ROCm installation timed out after 10 minutes. Killing the process"
$proc.Kill()
exit 1
}
if ($proc.ExitCode -ne 0) {
Write-Error "ROCm installation failed with exit code $($proc.ExitCode)"
exit 1
}
write-host "Completed AMD HIP SDK installation"
- name: Verify ROCm
id: verify
run: |
# Find and test ROCm installation
$clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1
if (-not $clangPath) {
Write-Error "ROCm installation not found"
exit 1
}
& $clangPath.FullName --version
- name: Build
id: cmake_build
run: |
$env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path)
$env:CMAKE_PREFIX_PATH="${env:HIP_PATH}"
cmake -G "Unix Makefiles" -B build -S . `
-DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" `
-DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" `
-DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/ -Wno-ignored-attributes -Wno-nested-anon-types" `
-DCMAKE_BUILD_TYPE=Release `
-DGGML_BACKEND_DL=ON `
-DGGML_NATIVE=OFF `
-DGGML_CPU=OFF `
-DGPU_TARGETS="${{ matrix.gpu_targets }}" `
-DGGML_HIP=ON `
-DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} `
-DLLAMA_BUILD_BORINGSSL=ON
cmake --build build --target ggml-hip -j ${env:NUMBER_OF_PROCESSORS}
md "build\bin\rocblas\library\"
md "build\bin\hipblaslt\library"
cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\"
cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\"
cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\"
cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\"
cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\"
- name: ccache-clear
uses: ./.github/actions/ccache-clear
with:
key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }}
- name: Pack artifacts
id: pack_artifacts
run: |
7z a -snl llama-bin-win-hip-${{ matrix.name }}-x64.zip .\build\bin\*
- name: Upload artifacts
uses: actions/upload-artifact@v6
with:
path: llama-bin-win-hip-${{ matrix.name }}-x64.zip
name: llama-bin-win-hip-${{ matrix.name }}-x64.zip
ios-xcode:
needs: [check-release, get-version]
if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -1572,7 +1573,7 @@ jobs:
- windows-cpu
- windows-cuda
#- windows-sycl
- windows-hip
- windows-rocm
- windows-openvino
- ubuntu-22-rocm
- ubuntu-cpu
@@ -1684,7 +1685,7 @@ jobs:
- [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
- [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
- [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
- [Ubuntu x64 (ROCm 7.2)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.2-x64.tar.gz)
- [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz)
- [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz)
- [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz)
- [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz)
@@ -1702,7 +1703,7 @@ jobs:
- [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip)
- [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip)
- [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip)
- [Windows x64 (HIP)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-hip-radeon-x64.zip)
- [Windows x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-rocm-7.14-x64.zip)
**openEuler:**
- [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23705)
.github/workflows/server-sanitize.yml
+14 -4
View File
@@ -25,6 +25,12 @@ on:
'tools/server/**.*'
]
pull_request:
types: [opened, synchronize, reopened]
paths: [
'.github/workflows/server-sanitize.yml'
]
env:
LLAMA_ARG_LOG_COLORS: 1
LLAMA_ARG_LOG_PREFIX: 1
@@ -90,15 +96,18 @@ jobs:
- name: Python setup
id: setup_python
uses: actions/setup-python@v6
with:
python-version: '3.11'
pip-install: -r tools/server/tests/requirements.txt
uses: actions/setup-python@v7
- name: Install Python dependencies
run: |
python3 -m venv .venv
.venv/bin/pip install -r tools/server/tests/requirements.txt
- name: Tests
id: server_integration_tests
if: ${{ (!matrix.disabled_on_pr || !github.event.pull_request) }}
run: |
source .venv/bin/activate
cd tools/server/tests
export ${{ matrix.extra_args }}
pytest -v -x -m "not slow"
@@ -107,6 +116,7 @@ jobs:
id: server_integration_tests_slow
if: ${{ (github.event.schedule || github.event.inputs.slow_tests == 'true') && matrix.build_type == 'Release' }}
run: |
source .venv/bin/activate
cd tools/server/tests
export ${{ matrix.extra_args }}
SLOW_TESTS=1 pytest -v -x
common/arg.cpp
+3 -2
View File
@@ -3312,8 +3312,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
{"--tools-runtime"}, "OPTION",
"experimental: run tools in a separate runtime environment (default: none, use host environment)\n"
"available options:\n"
" 'docker:<image>': spin up a new Docker container and reuse it for all invocations, clean up on server exit\n"
" 'docker-container:<id>': use an existing Docker container by ID, won't stop on server exit\n",
" 'docker:<image>', 'podman:<image>': spin up a new container and reuse it for all invocations, clean up on server exit\n"
" 'docker-container:<id>', 'podman-container:<id>': use an existing container by ID, won't stop on server exit\n"
" 'ssh:<target>': run tools on a remote POSIX host over SSH, key-based auth and a trusted host key are required\n",
[](common_params & params, const std::string & value) {
params.server_tools_runtime = value;
}
common/chat.cpp
+151
View File
@@ -3086,6 +3086,151 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
return data;
}
// An assistant turn is rendered as one or more messages, each
// "<|start|>assistant to=<recipient><|message|>{content}{END}" where END is
// <|eom|> (more messages follow) or <|eot|> (end of turn):
// - chain-of-thought: to=self, terminated by <|eom|>
// - final answer: to=user, terminated by <|eot|>
// The generation prompt is just "<|start|>assistant"; the model emits its own
// " to=...<|message|>".
static common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl,
const autoparser::generation_params & inputs) {
common_chat_params data;
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs);
data.generation_prompt = "<|start|>assistant";
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
data.preserved_tokens = {
"<|start|>", "<|message|>", "<|eom|>", "<|eot|>",
// ATEM tool-call markup emitted on " to=<tool>" turns.
"<atem:function_calls>", "<atem:invoke", "<atem:parameter", "</atem:parameter>",
"</atem:invoke>", "</atem:function_calls>",
};
data.message_delimiters = {
{ COMMON_CHAT_ROLE_ASSISTANT, "<|start|>assistant" },
{ COMMON_CHAT_ROLE_USER, "<|start|>user" },
{ COMMON_CHAT_ROLE_SYSTEM, "<|start|>system" },
{ COMMON_CHAT_ROLE_TOOL, "<|start|>tool" },
};
if (inputs.has_continuation()) {
const auto & msg = inputs.continue_msg;
data.generation_prompt = "<|start|>assistant to=self<|message|>" + msg.reasoning_content;
if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
data.generation_prompt += "<|eom|><|start|>assistant to=user<|message|>" + msg.render_content();
}
data.prompt += data.generation_prompt;
}
auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
// Constrained grammar whenever tools are offered.
auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
auto start = p.rule("start", p.literal("<|start|>assistant"));
if (!extract_reasoning && !include_grammar) {
return start + p.content(p.rest());
}
if (extract_reasoning) {
p.rule("analysis", p.literal(" to=self<|message|>") + p.reasoning(p.until("<|eom|>")) + p.literal("<|eom|>"));
} else {
p.rule("analysis", p.literal(" to=self<|message|>") + p.content(p.until("<|eom|>")) + p.literal("<|eom|>"));
}
auto analysis = p.ref("analysis");
auto recipient = p.optional(p.literal(" to=user"));
auto final_msg = p.rule("final", recipient + p.literal("<|message|>") + p.content(p.until("<|eot|>")));
if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
auto string_value = p.ac(
p.tool_arg_string_value(p.until("</atem:parameter>")) + p.tool_arg_close(p.literal("</atem:parameter>")),
"</atem:parameter>");
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
const auto & function = tool.at("function");
const std::string name = function.at("name");
auto params = function.contains("parameters") ? function.at("parameters") : json::object();
auto args = p.eps();
if (params.contains("properties") && params.at("properties").is_object() && !params.at("properties").empty()) {
auto schema_info = common_schema_info();
schema_info.resolve_refs(params);
auto arg_choice = p.choice();
for (const auto & [prop_name, prop_schema] : params.at("properties").items()) {
auto value_parser = p.eps();
if (schema_info.resolves_to_string(prop_schema)) {
value_parser = string_value;
} else {
value_parser = p.tool_arg_json_value(
p.schema(p.json(), "tool-" + name + "-arg-" + prop_name + "-schema", prop_schema, false))
+ p.tool_arg_close(p.literal("</atem:parameter>"));
}
auto arg_rule = p.tool_arg(
p.tool_arg_open(p.literal("<atem:parameter name=\"") + p.tool_arg_name(p.literal(prop_name)) + p.literal("\">")) +
value_parser);
arg_choice |= arg_rule;
}
args = p.zero_or_more(arg_choice + p.space());
}
auto tool_parser = p.tool(
p.tool_open(p.literal(" to=") + p.until("<|message|>") +
p.literal("<|message|><atem:function_calls>") + p.space() +
p.literal("<atem:invoke name=\"") + p.tool_name(p.literal(name)) + p.literal("\">") + p.space())
<< p.tool_args(args)
<< p.tool_close(p.literal("</atem:invoke>") + p.space() + p.literal("</atem:function_calls>")));
tool_choice |= p.rule("tool-" + name, tool_parser);
});
auto tool_calls = inputs.parallel_tool_calls
? p.trigger_rule("tool-call", tool_choice + p.zero_or_more(p.literal("<|eom|>") + start + tool_choice))
: p.trigger_rule("tool-call", tool_choice);
if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) {
return p.zero_or_more(start + analysis) + start + tool_calls;
}
return p.zero_or_more(start + analysis) + start + (tool_calls | final_msg);
}
return p.zero_or_more(start + analysis) + start + final_msg;
});
data.parser = parser.save();
if (include_grammar) {
data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
data.grammar = build_grammar([&](const common_grammar_builder & builder) {
foreach_function(inputs.tools, [&](const json & tool) {
const auto & function = tool.at("function");
auto schema = function.contains("parameters") ? function.at("parameters") : json::object();
builder.resolve_refs(schema);
});
parser.build_grammar(builder, data.grammar_lazy);
});
data.grammar_triggers = {
{ COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN,
"<\\|start\\|>assistant( to=(?!self<\\|message\\|>)(?!user<\\|message\\|>)[^<]*?<\\|message\\|>)" },
};
}
return data;
}
static json common_chat_extra_context() {
json ctx = json::object();
std::chrono::system_clock::time_point now = std::chrono::system_clock::now();
@@ -3114,6 +3259,12 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
return common_chat_params_init_gpt_oss(tmpl, params);
}
// Muse Glimmer format using " to=<recipient>" recipients and <|eom|>/<|eot|> message terminators.
if (src.find("<atem:function_calls>") != std::string::npos && src.find("<|eom|>") != std::string::npos) {
LOG_DBG("Using specialized template: Muse Glimmer\n");
return common_chat_params_init_muse_glimmer(tmpl, params);
}
// Functionary v3.2 - uses recipient-based format with >>>recipient\n{content}
// Detection: template has ">>>all" for content and ">>>" prefix for tool calls
if (src.find(">>>all") != std::string::npos && src.find(">>>${recipient}") != std::string::npos) {
common/common.cpp
+1
View File
@@ -1639,6 +1639,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
cparams.n_seq_max = params.n_parallel;
cparams.n_rs_seq = params.speculative.need_n_rs_seq();
cparams.n_outputs_max = std::max(params.n_outputs_max, 0);
cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0);
cparams.n_batch = params.n_batch;
cparams.n_ubatch = params.n_ubatch;
cparams.n_threads = params.cpuparams.n_threads;
common/common.h
+1
View File
@@ -447,6 +447,7 @@ struct common_params {
int32_t n_parallel = 1; // number of parallel sequences to decode
int32_t n_sequences = 1; // number of sequences to decode
int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch)
int32_t n_outputs_max_per_seq = 1; // max outputs per sequence
int32_t grp_attn_n = 1; // group-attention factor
int32_t grp_attn_w = 512; // group-attention width
int32_t n_print = -1; // print token count every n tokens (-1 = disabled)
common/llguidance.cpp
+2
View File
@@ -116,6 +116,8 @@ static llama_sampler_i llama_sampler_llg_i = {
/* .backend_accept = */ NULL,
/* .backend_apply = */ NULL,
/* .backend_set_input = */ NULL,
/* .backend_reset = */ NULL,
/* .copy_state = */ NULL,
};
static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len,
common/reasoning-budget.cpp
+2
View File
@@ -217,6 +217,8 @@ static struct llama_sampler_i common_reasoning_budget_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {
common/sampling.cpp
+20
View File
@@ -518,6 +518,26 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) {
};
}
void common_sampler_copy(const common_sampler * src, common_sampler * dst) {
if (!src || !dst || src == dst) {
return;
}
GGML_ASSERT((src->grmr == nullptr) == (dst->grmr == nullptr));
GGML_ASSERT((src->rbudget == nullptr) == (dst->rbudget == nullptr));
llama_sampler_copy(src->grmr, dst->grmr);
llama_sampler_copy(src->rbudget, dst->rbudget);
llama_sampler_copy(src->chain, dst->chain);
dst->params = src->params;
dst->prev = src->prev;
dst->cur = src->cur;
dst->cur_p = src->cur_p;
dst->cur_p.data = src->cur_p.data ? dst->cur.data() : nullptr; // re-point to dst's buffer
dst->t_total_us = src->t_total_us;
}
void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) {
// TODO: measure grammar performance
common/sampling.h
+1
View File
@@ -47,6 +47,7 @@ void common_sampler_free(struct common_sampler * gsmpl);
void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated);
void common_sampler_reset (struct common_sampler * gsmpl);
struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl);
void common_sampler_copy (const struct common_sampler * src, struct common_sampler * dst);
// arguments can be nullptr to skip printing
void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl);
common/speculative.cpp
+20 -1
View File
@@ -1032,7 +1032,14 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
return true;
}
if (batch_in.token == nullptr || batch_in.embd != nullptr) {
// Target prefill may contain token IDs or multimodal embeddings. Both
// produce the target-layer features used to seed the draft KV cache, so
// skipping the embedding batches leaves a hole in the draft's cache and
// the next injection fails to initialize.
// TODO: revisit after https://github.com/ggml-org/llama.cpp/pull/24669 is merged
const bool has_tokens = batch_in.token != nullptr;
const bool has_embeddings = batch_in.embd != nullptr;
if (has_tokens == has_embeddings) {
return true;
}
@@ -2292,6 +2299,7 @@ common_params common_base_params_to_speculative(const common_params & params) {
result.cache_type_k = params_spec.cache_type_k;
result.cache_type_v = params_spec.cache_type_v;
result.n_outputs_max = params.n_parallel;
result.n_outputs_max_per_seq = 1;
return result;
}
@@ -2377,6 +2385,17 @@ common_speculative_init_result_ptr common_speculative_init_from_params(common_pa
return std::make_unique<common_speculative_init_result>(params, model_tgt, ctx_tgt);
}
common_speculative_output_limits common_speculative_get_output_limits(
int32_t n_batch, int32_t n_parallel, int32_t n_draft) {
const int64_t per_seq = 1 + (int64_t) std::max(0, n_draft);
const int64_t total = (int64_t) n_parallel * per_seq;
return {
/* .total = */ (int32_t) std::min<int64_t>(n_batch, total),
/* .per_seq = */ (int32_t) std::min<int64_t>(n_batch, per_seq),
};
}
// initialization of the speculative decoding system
//
common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {
common/speculative.h
+9
View File
@@ -25,6 +25,15 @@ int32_t common_speculative_n_max(const common_params_speculative * spec);
common_params common_base_params_to_speculative(const common_params & params);
struct common_speculative_output_limits {
int32_t total;
int32_t per_seq;
};
// return the output limits needed for speculative decoding
common_speculative_output_limits common_speculative_get_output_limits(
int32_t n_batch, int32_t n_parallel, int32_t n_draft);
common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq);
void common_speculative_free(common_speculative * spec);
conversion/__init__.py
+3
View File
@@ -183,6 +183,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Olmo3ForCausalLM": "olmo",
"OlmoForCausalLM": "olmo",
"OlmoeForCausalLM": "olmo",
"MuseGlimmerAssistantModel": "muse_glimmer",
"MuseGlimmerForConditionalGeneration": "muse_glimmer",
"OpenELMForCausalLM": "openelm",
"OrionForCausalLM": "orion",
"PLMForCausalLM": "plm",
@@ -298,6 +300,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
"Mistral3ForConditionalGeneration": "llava",
"NemotronH_Nano_VL_V2": "nemotron",
"MuseGlimmerForConditionalGeneration": "muse_glimmer",
"PaddleOCRVisionModel": "ernie",
"Phi4ForCausalLMV": "phi",
"Qwen2AudioForConditionalGeneration": "ultravox",
conversion/muse_glimmer.py
+179
View File
@@ -0,0 +1,179 @@
from __future__ import annotations
import json
from typing import Any, Iterable, TYPE_CHECKING
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import MmprojModel, ModelBase, TextModel, gguf
def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor":
"""Invert transformers' `_permute_for_rope`: HF stores Q/K in rotate_half layout,
llama.cpp consumes the interleaved (NORM) layout."""
if tensor.ndim == 2:
dim1, dim2 = tensor.shape
return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2)
if tensor.ndim == 1:
(dim1,) = tensor.shape
return tensor.view(n_heads, 2, dim1 // n_heads // 2).transpose(1, 2).reshape(dim1)
raise ValueError(f"_unpermute_for_rope: unexpected shape {tuple(tensor.shape)}")
@ModelBase.register("MuseGlimmerForConditionalGeneration")
class MuseGlimmerModel(TextModel):
model_arch = gguf.MODEL_ARCH.MUSE_GLIMMER
def norm_shift(self, name: str) -> float:
# All four layer norms use 1, the final norm uses 0.
return 1.0 if name.endswith("layernorm.weight") else 0.0
def set_vocab(self):
self._set_vocab_gpt2()
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(self.dir_model)
eot_id = tok.convert_tokens_to_ids("<|eot|>")
if isinstance(eot_id, int) and eot_id >= 0:
self.gguf_writer.add_eot_token_id(eot_id)
def set_gguf_parameters(self):
super().set_gguf_parameters()
hparams = self.hparams
self.gguf_writer.add_final_logit_softcapping(hparams["final_logit_softcapping"])
self.gguf_writer.add_logit_scale(hparams["output_multiplier"])
self.gguf_writer.add_sliding_window(hparams["sliding_window"])
self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]])
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
shift = self.norm_shift(name)
if shift != 0.0:
data_torch = data_torch + shift
# Invert transformers' `_permute_for_rope` on Q/K, we keep ggml's NORM (interleaved) rope
if ".self_attn.q_proj." in name:
data_torch = _unpermute_for_rope(data_torch, int(self.hparams["num_attention_heads"]))
elif ".self_attn.k_proj." in name:
data_torch = _unpermute_for_rope(data_torch, int(self.hparams["num_key_value_heads"]))
# Synthesize QK-norm weights to absorb qk_scale_factor.
# MuseGlimmer implementation: scaleless RMSNorm followed by qk_scale_factor..
if bid is not None and name.endswith(f"model.layers.{bid}.self_attn.q_proj.weight"):
head_dim = self.hparams["head_dim"]
q_scale = float(self.hparams["qk_scale_factor"])
yield (
self.map_tensor_name(f"model.layers.{bid}.self_attn.q_norm.weight"),
torch.full((head_dim,), q_scale, dtype=torch.float32),
)
yield (
self.map_tensor_name(f"model.layers.{bid}.self_attn.k_norm.weight"),
torch.ones((head_dim,), dtype=torch.float32),
)
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("MuseGlimmerForConditionalGeneration")
class MuseGlimmerVisionModel(MmprojModel):
def get_vision_config(self) -> dict[str, Any] | None:
c = self.global_config.get("vision_config")
if not c:
return None
# MuseGlimmer actually uses dynamic size, initialize with nominal size
image_size = c["pos_emb_height"] * c["patch_size"] * c["merge_size"]
return {**c, "image_size": image_size}
def set_gguf_parameters(self):
super().set_gguf_parameters()
assert self.hparams_vision is not None
c = self.hparams_vision # enriched vision_config from get_vision_config()
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER)
self.gguf_writer.add_vision_attention_layernorm_eps(float(c["layer_norm_eps"]))
self.gguf_writer.add_vision_spatial_merge_size(int(c["merge_size"]))
@classmethod
def filter_tensors(cls, item):
name, gen = item
keep = ("model.vision_tower.", "model.vision_adapter.", "model.vision_projection.")
if not any(name.startswith(k) for k in keep):
return None
return super().filter_tensors((name, gen))
# 3-layer projector MLP
_MM_MLP_MAP = {
"model.vision_adapter.fc1": (gguf.MODEL_TENSOR.V_MMPROJ, 0),
"model.vision_adapter.fc2": (gguf.MODEL_TENSOR.V_MMPROJ, 1),
"model.vision_projection": (gguf.MODEL_TENSOR.V_MMPROJ, 2),
}
def modify_tensors(self, data_torch, name, bid):
assert self.hparams_vision is not None
if ".attn.q_proj." in name or ".attn.k_proj." in name:
n_heads = int(self.hparams_vision["num_attention_heads"])
data_torch = _unpermute_for_rope(data_torch, n_heads)
# Lay out the pt=2 temporal slabs of the patch embedding as a conv2d for build_inp()
if name.endswith("patch_embedder.patch_embedding.weight"):
n_embd = data_torch.shape[0]
pt = int(self.hparams_vision["patch_temporal"])
ps = int(self.hparams_vision["patch_size"])
data_torch = data_torch.view(n_embd, pt, 3, ps, ps).sum(dim=1) # (n_embd, 3, ps, ps)
stem, _, suffix = name.rpartition(".")
if stem in self._MM_MLP_MAP:
tensor_key, idx = self._MM_MLP_MAP[stem]
yield (self.format_tensor_name(tensor_key, bid=idx, suffix="." + suffix), data_torch)
return
yield (self.map_tensor_name(name), data_torch)
@ModelBase.register("MuseGlimmerAssistantModel")
class MuseGlimmerAssistantModel(TextModel):
model_arch = gguf.MODEL_ARCH.DFLASH
def set_vocab(self):
if self.target_model_dir is None:
raise ValueError(
"MuseGlimmerAssistant (DFlash drafter) requires --target-model-dir pointing to the "
"target MuseGlimmer HF directory"
)
original_dir = self.dir_model
self.dir_model = self.target_model_dir
from . import get_model_class
with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
target_arch = json.load(f)["architectures"][0]
target_cls = get_model_class(target_arch)
if target_cls is not type(self):
target_cls.set_vocab(self) # ty: ignore[unresolved-attribute]
else:
super().set_vocab()
self.dir_model = original_dir
mask_token_id = self.hparams.get("mask_token_id")
if mask_token_id is not None:
self.gguf_writer.add_mask_token_id(int(mask_token_id))
def set_gguf_parameters(self):
super().set_gguf_parameters()
h = self.hparams
self.gguf_writer.add_block_size(int(h["block_size"]))
# dflash.target_layers[k] refers to the inputs going into the ith layer, which come from the (i-1)th layer's output.
# The transformers configuration refers to the outputs being recorded.
self.gguf_writer.add_target_layers([int(x) + 1 for x in h["target_layer_ids"]])
if h.get("sliding_window") and h.get("layer_types"):
self.gguf_writer.add_sliding_window(int(h["sliding_window"]))
self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in h["layer_types"]])
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
# DFlash defaults to NEOX (rotate_half) rope, matching transformers HF layout for Q/K, QK-norms
# no permutation needed.
yield (self.map_tensor_name(name), data_torch)
docs/speculative.md
+6
View File
@@ -202,6 +202,12 @@ Example Video:
If a draft model is combined with a draftless decoding the draftless decoding has higher precedence.
### Backend Sampling
Use `--backend-sampling` to run supported target-model samplers on the model backend. Draft-model sampling uses the backend by default and can be controlled with `--spec-draft-backend-sampling` and `--no-spec-draft-backend-sampling`.
Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required.
### General Speculative Parameters
```
examples/lookup/lookup.cpp
+6
View File
@@ -3,9 +3,11 @@
#include "common.h"
#include "ngram-cache.h"
#include "sampling.h"
#include "speculative.h"
#include "log.h"
#include "llama.h"
#include <algorithm>
#include <clocale>
#include <cstdint>
#include <cstdio>
@@ -27,6 +29,10 @@ int main(int argc, char ** argv){
// max. number of additional tokens to draft if match is found
const int n_draft = params.speculative.draft.n_max;
const auto output_limits = common_speculative_get_output_limits(params.n_batch, params.n_parallel, n_draft);
params.n_outputs_max = output_limits.total;
params.n_outputs_max_per_seq = output_limits.per_seq;
// init llama.cpp
llama_backend_init();
llama_numa_init(params.numa);
examples/speculative-simple/speculative-simple.cpp
+9
View File
@@ -5,6 +5,7 @@
#include "log.h"
#include "llama.h"
#include <algorithm>
#include <clocale>
#include <cstdio>
#include <cstring>
@@ -29,6 +30,11 @@ int main(int argc, char ** argv) {
return 1;
}
const auto output_limits = common_speculative_get_output_limits(
params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
params.n_outputs_max = output_limits.total;
params.n_outputs_max_per_seq = output_limits.per_seq;
// init llama.cpp
llama_backend_init();
llama_numa_init(params.numa);
@@ -55,6 +61,9 @@ int main(int argc, char ** argv) {
auto params_dft = params;
params_dft.n_outputs_max = params.n_parallel;
params_dft.n_outputs_max_per_seq = 1;
params_dft.devices = params_spec.devices;
params_dft.model = params_spec.mparams;
params_dft.n_gpu_layers = params_spec.n_gpu_layers;
examples/speculative/speculative.cpp
+8
View File
@@ -1,6 +1,7 @@
#include "arg.h"
#include "common.h"
#include "sampling.h"
#include "speculative.h"
#include "log.h"
#include "llama.h"
@@ -57,6 +58,11 @@ int main(int argc, char ** argv) {
// max number of parallel drafting sequences (i.e. tree branches)
const int n_seq_dft = params.n_parallel;
const auto output_limits = common_speculative_get_output_limits(
params.n_batch, params.n_parallel, params.speculative.draft.n_max);
params.n_outputs_max = output_limits.total;
params.n_outputs_max_per_seq = output_limits.per_seq;
// probability threshold for splitting a draft branch (only for n_seq_dft > 1)
const float p_draft_split = params.speculative.draft.p_split;
@@ -83,6 +89,8 @@ int main(int argc, char ** argv) {
params.devices = params.speculative.draft.devices;
params.model = params.speculative.draft.mparams;
params.n_gpu_layers = params.speculative.draft.n_gpu_layers;
params.n_outputs_max = params.n_parallel;
params.n_outputs_max_per_seq = 1;
if (params.speculative.draft.cpuparams.n_threads > 0) {
params.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads;
}
ggml/src/ggml-cpu/ggml-cpu.c
+1 -1
View File
@@ -2616,7 +2616,7 @@ static bool ggml_thread_apply_priority(int32_t prio) {
return true;
}
#elif defined(__gnu_linux__)
#elif defined(__linux__)
// TODO: this may not work on BSD, to be verified
static bool ggml_thread_apply_affinity(const bool * mask) {
ggml/src/ggml-cuda/ggml-cuda.cu
+1 -1
View File
@@ -5187,7 +5187,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
return max_bias == 0.0f;
}
case GGML_OP_ROLL:
if(op->src[0]->type == GGML_TYPE_F32) {
if(op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0])) {
return true;
}
return false;
ggml/src/ggml-cuda/mmq-config-ampere.cuh
+17
View File
@@ -33,6 +33,23 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_Q2_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_DT3, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_DT3, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_DT3, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_DT3, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_DT3, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_DT3, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q4_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q4_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
ggml/src/ggml-cuda/mmq-load-tiles.cuh
+111
View File
@@ -176,6 +176,117 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_dt3(
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
// DT3: 128 elements as two ternary planes with one fp16 scale each, w = d1*t1 + d2*t2.
// The two planes cannot be fused into a single int8 value because d1 != d2, so the
// tile holds both planes decoded to int8 trits in {-1, 0, +1}, plane 2 offset by
// 2*MMQ_TILE_NE_K ints from plane 1 within each row, and 2x8 float scales per row.
// The decode is the same base-3 digit iteration as vec_dot_dt3_q8_1 (see vecdotq.cuh):
// each packed byte is decoded once with q -> (q*3) & 0xFF, two bytes at a time in the
// 16-bit lanes of one int. Digit bytes in {0, 1, 2} become trit bytes in {-1, 0, +1}
// without cross-byte borrows via ((dig | 0x80808080) - 0x01010101) ^ 0x80808080.
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
constexpr int row_stride_qs = sram_stride;
constexpr int row_stride_df = sram_stride;
int * x_qs = (int *) x_tile;
float * x_df = (float *) (x_qs + 4*MMQ_TILE_NE_K);
#else
constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_DT3, I);
constexpr int row_stride_qs = 4*MMQ_TILE_NE_K + 1;
constexpr int row_stride_df = 4*MMQ_TILE_NE_K/QI8_0 + 1;
int * x_qs = (int *) x_tile;
float * x_df = (float *) (x_qs + txs.qs);
#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
constexpr int blocks_per_iter = MMQ_ITER_K / QK_DT3;
static_assert(blocks_per_iter == 2, "DT3 load assumes 2 blocks per iteration");
// 32 threads per row: 2 blocks x 2 planes x 8 slots. Slots 0..5 decode one 4-byte
// quad of qs each (5 ints of 4 trits), slot 6 decodes the 2 qh bytes (2 ints),
// slot 7 is idle.
constexpr int threads_per_row = 32;
constexpr int nrows = warp_size / threads_per_row;
const int txi = threadIdx.x % threads_per_row;
const int kbx = txi / 16;
const int p = (txi / 8) % 2;
const int u = txi % 8;
#pragma unroll
for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
if (fallback) {
i = min(i, i_max);
}
const block_dt3 * bxi = (const block_dt3 *) x + kbx0 + i*stride + kbx;
int * dst = x_qs + i*row_stride_qs + p*(2*MMQ_TILE_NE_K) + kbx*(QK_DT3/4);
if (u < 6) {
// qs[0..16): byte m holds elements m + 16*n; qs[16..24): byte 16 + m
// holds elements 80 + m + 8*n. Either way a quad of consecutive bytes
// yields 4 consecutive elements per digit n, i.e. one tile int.
const int q32 = get_int_b4(bxi->qs[p], u);
int qa = (q32 >> 0) & 0x00FF00FF;
int qb = (q32 >> 8) & 0x00FF00FF;
#pragma unroll
for (int n = 0; n < 5; ++n) {
const int qa3 = qa*3;
const int qb3 = qb*3;
const int dig = ((qa3 >> 8) & 0x00030003) | (qb3 & 0x03000300);
qa = qa3 & 0x00FF00FF;
qb = qb3 & 0x00FF00FF;
const int idx = u < 4 ? 4*n + u : 20 + 2*n + (u - 4);
dst[idx] = ((dig | 0x80808080) - 0x01010101) ^ 0x80808080;
}
} else if (u == 6) {
// qh: byte b holds elements 120 + b + 2*n for n = 0..3 — only 4 digits
// are iterated, the 5th is packing padding and must never decode.
int q = bxi->qh[p][0] | (bxi->qh[p][1] << 16);
#pragma unroll
for (int s = 0; s < 2; ++s) {
int dig = 0;
#pragma unroll
for (int n = 0; n < 2; ++n) {
const int q3 = q*3;
dig |= (((q3 >> 8) & 0x03) | ((q3 >> 16) & 0x0300)) << (16*n);
q = q3 & 0x00FF00FF;
}
dst[30 + s] = ((dig | 0x80808080) - 0x01010101) ^ 0x80808080;
}
}
}
// 16 scale entries per row and iteration: 2 planes x 2 blocks x 4 q8_1 chunks.
// Plane 2 scales sit after the 8 plane-1 entries, matching the vec_dot indexing.
constexpr int scale_entries_per_plane = blocks_per_iter*(QK_DT3/QK8_1);
const int ksx = threadIdx.x % (2*scale_entries_per_plane);
const int ps = ksx / scale_entries_per_plane;
const int scale_block = (ksx % scale_entries_per_plane) / (QK_DT3/QK8_1);
#pragma unroll
for (int i0 = 0; i0 < I; i0 += nwarps) {
int i = i0 + threadIdx.y;
if (fallback) {
i = min(i, i_max);
}
const block_dt3 * bxi = (const block_dt3 *) x + kbx0 + i*stride + scale_block;
x_df[i*row_stride_df + ksx] = bxi->d[ps];
}
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
ggml/src/ggml-cuda/mmq-vec-dot.cuh
+134
View File
@@ -280,6 +280,140 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
}
// DT3: both decoded ternary planes of the tile are multiplied against the same y data,
// each with its own per-chunk scale: sum = dB * (sumi1*dA1 + sumi2*dA2).
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_dt3_q8_1_dp4a(
const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_DT3, I);
const int * x_qs = (const int *) x;
const float * x_df = (const float *) x_qs + txs.qs;
const int * y_qs = (const int *) y + 4;
const float * y_df = (const float *) y;
constexpr int row_stride_qs = 4*MMQ_TILE_NE_K + 1;
constexpr int row_stride_df = 4*MMQ_TILE_NE_K/QI8_0 + 1;
// #pragma unroll
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
const int k0 = k00 + k01;
#pragma unroll
for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
#pragma unroll
for (int i0 = 0; i0 < I; i0 += warp_size) {
const int i = i0 + threadIdx.x;
const int * yqs = &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K];
const float dB = y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)];
sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMQ>
(&x_qs[i*row_stride_qs + k0], yqs,
x_df[i*row_stride_df + k0/QI8_0], dB);
sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMQ>
(&x_qs[i*row_stride_qs + 2*MMQ_TILE_NE_K + k0], yqs,
x_df[i*row_stride_df + 2*MMQ_TILE_NE_K/QI8_0 + k0/QI8_0], dB);
}
}
}
}
template <ggml_type type, int J, bool fallback>
static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_dt3_q8_1_mma(
const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
// DT3 MMQ is not enabled for AMD (no config entries select it); this only has to compile.
GGML_UNUSED_VARS(x, y, sum, k00);
NO_DEVICE_CODE;
#else
typedef tile<16, 8, int> tile_A;
typedef tile< 8, 8, int> tile_B;
typedef tile<16, 8, int> tile_C;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
const int * x_qs = (const int *) x;
const float * x_df = (const float *) x_qs + 4*MMQ_TILE_NE_K;
const int * y_qs = (const int *) y + 4;
const float * y_df = (const float *) y;
tile_A A[ntx][2][MMQ_TILE_NE_K/QI8_0];
float dA[ntx][tile_C::ne/2][2][MMQ_TILE_NE_K/QI8_0];
const int i0 = (threadIdx.y/ntx)*rows_per_warp;
#pragma unroll
for (int n = 0; n < ntx; ++n) {
#pragma unroll
for (int p = 0; p < 2; ++p) {
#pragma unroll
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
const int k0 = k00 + k01;
load_ldmatrix(A[n][p][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*sram_stride + p*(2*MMQ_TILE_NE_K) + k0, sram_stride);
}
}
#pragma unroll
for (int l = 0; l < tile_C::ne/2; ++l) {
const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
#pragma unroll
for (int p = 0; p < 2; ++p) {
#pragma unroll
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
const int k0 = k00 + k01;
dA[n][l][p][k01/QI8_0] = x_df[i*sram_stride + p*(2*MMQ_TILE_NE_K/QI8_0) + k0/QI8_0];
}
}
}
}
#pragma unroll
for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
#pragma unroll
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
tile_B B;
float dB[tile_C::ne/2];
load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
#pragma unroll
for (int l = 0; l < tile_C::ne/2; ++l) {
const int j = j0 + tile_C::get_j(l);
dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
}
#pragma unroll
for (int n = 0; n < ntx; ++n) {
tile_C C1;
tile_C C2;
mma(C1, A[n][0][k01/QI8_0], B);
mma(C2, A[n][1][k01/QI8_0], B);
#pragma unroll
for (int l = 0; l < tile_C::ne; ++l) {
sum[(j0/tile_C::J + n)*tile_C::ne + l] +=
(C1.x[l]*dA[n][l/2][0][k01/QI8_0] + C2.x[l]*dA[n][l/2][1][k01/QI8_0])*dB[l%2];
}
}
}
}
#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a(
const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
ggml/src/ggml-cuda/mmq.cu
+28
View File
@@ -4,6 +4,7 @@
#include "mmid.cuh"
#include <cstdint>
#include <cstdlib>
static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
switch (args.type_x) {
@@ -13,6 +14,9 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
case GGML_TYPE_Q2_0:
mul_mat_q_case<GGML_TYPE_Q2_0>(ctx, args, stream);
break;
case GGML_TYPE_DT3:
mul_mat_q_case<GGML_TYPE_DT3>(ctx, args, stream);
break;
case GGML_TYPE_Q4_0:
mul_mat_q_case<GGML_TYPE_Q4_0>(ctx, args, stream);
break;
@@ -261,6 +265,30 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
return false;
#endif // GGML_CUDA_FORCE_CUBLAS
// DT3 keeps two decoded ternary planes per row in SRAM, roughly double the tile
// of a single-plane type: only the MMA data layout is implemented and even the
// narrowest tile needs ~76 KiB of shared memory, more than e.g. Turing offers.
if (type == GGML_TYPE_DT3) {
if (!turing_mma_available(cc)) {
return false;
}
// Two integer dot products per weight cancel the 2x int8-over-fp16 advantage
// of the tensor cores, so at large batch the dequantize + fp16 cuBLAS path
// wins; MMQ avoids the dequantization round-trip and wins below the
// crossover (measured on RTX 4060 Ti). Override for experiments with
// GGML_CUDA_DT3_MMQ_MAX_BATCH.
static const int64_t max_batch = []() {
const char * env = getenv("GGML_CUDA_DT3_MMQ_MAX_BATCH");
return env ? atoll(env) : 192;
}();
if (ne11 > max_batch) {
return false;
}
const int id = ggml_cuda_get_device();
const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
return mmq_get_nbytes_shared(ggml_cuda_mmq_get_config(GGML_TYPE_DT3, 8, true, cc), cc) <= smpbo;
}
bool mmq_supported;
switch (type) {
ggml/src/ggml-cuda/mmq.cuh
+20
View File
@@ -61,6 +61,7 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
switch (type_x) {
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_DT3:
return MMQ_Q8_1_DS_LAYOUT_D4;
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
@@ -121,6 +122,7 @@ struct tile_x_sizes {
enum ggml_cuda_mmq_sram_layout {
GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0,
GGML_CUDA_MMQ_SRAM_LAYOUT_DT3, // Two decoded ternary planes per row, each with its own per-block scales.
GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1,
GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K,
GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K,
@@ -133,6 +135,8 @@ static constexpr __host__ __device__ int ggml_cuda_mmq_get_sram_stride(ggml_cuda
switch (sram_layout) {
case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0:
return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4;
case GGML_CUDA_MMQ_SRAM_LAYOUT_DT3:
return 4*MMQ_TILE_NE_K + 4*MMQ_TILE_NE_K/QI8_0 + 4;
case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1:
return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_1 + 4;
case GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K:
@@ -151,6 +155,7 @@ static constexpr __host__ __device__ int ggml_cuda_mmq_get_sram_stride(ggml_cuda
}
static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0) % 8 == 4, "Wrong padding.");
static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_DT3) % 8 == 4, "Wrong padding.");
static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1) % 8 == 4, "Wrong padding.");
static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K) % 8 == 4, "Wrong padding.");
static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K) % 8 == 4, "Wrong padding.");
@@ -377,6 +382,7 @@ static constexpr __device__ int ggml_cuda_mmq_get_rows_per_warp(ggml_type type,
#define MMQ_DP4A_TXS_Q8_0 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_0 + I/(QI8_0/2), 0}
#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*4/QI8_0 + I/(QI8_0/4), 0}
#define MMQ_DP4A_TXS_Q8_1 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_1 + I/(QI8_1/2), 0}
#define MMQ_DP4A_TXS_DT3 tile_x_sizes{I*MMQ_TILE_NE_K*4 + I, I*MMQ_TILE_NE_K*4/QI8_0 + I, 0}
#define MMQ_DP4A_TXS_Q2_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K + I, 0}
#define MMQ_DP4A_TXS_Q3_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I, I*MMQ_TILE_NE_K/8 + I/8}
#define MMQ_DP4A_TXS_Q4_K tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_K, I*MMQ_TILE_NE_K/8 + I/8}
@@ -387,6 +393,7 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml
switch (type) {
case GGML_TYPE_Q1_0: return MMQ_DP4A_TXS_Q8_0;
case GGML_TYPE_Q2_0: return MMQ_DP4A_TXS_Q8_0;
case GGML_TYPE_DT3: return MMQ_DP4A_TXS_DT3;
case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0;
case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1;
case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0;
@@ -550,6 +557,12 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
case GGML_TYPE_DT3:
return ggml_cuda_mmq_util_funcs(
VDR_Q8_0_Q8_1_MMQ,
ggml_cuda_mmq_load_tiles_dt3<type, J, fallback>,
ggml_cuda_mmq_vec_dot_dt3_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
case GGML_TYPE_Q4_0:
return ggml_cuda_mmq_util_funcs(
VDR_Q4_0_Q8_1_MMQ,
@@ -714,6 +727,12 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
ggml_cuda_mmq_write_back_mma<type, J, fallback>);
case GGML_TYPE_DT3:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_dt3<type, J, fallback>,
ggml_cuda_mmq_vec_dot_dt3_q8_1_mma<type, J, fallback>,
ggml_cuda_mmq_write_back_mma<type, J, fallback>);
case GGML_TYPE_Q4_0:
return ggml_cuda_mmq_util_funcs(
-1,
@@ -1565,6 +1584,7 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda
extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q2_0);
extern DECL_MMQ_CASE(GGML_TYPE_DT3);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_1);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_0);
ggml/src/ggml-cuda/template-instances/generate_cu_files.py
+1
View File
@@ -37,6 +37,7 @@ SOURCE_FATTN_MMA_CASE = "DECL_FATTN_MMA_F16_CASE({head_size_kq}, {head_size_v},
TYPES_MMQ = [
"GGML_TYPE_Q1_0",
"GGML_TYPE_Q2_0",
"GGML_TYPE_DT3",
"GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0",
"GGML_TYPE_Q2_K", "GGML_TYPE_Q3_K", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", "GGML_TYPE_Q6_K",
"GGML_TYPE_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S",
ggml/src/ggml-cuda/template-instances/mmq-instance-dt3.cu
+5
View File
@@ -0,0 +1,5 @@
// This file has been autogenerated by generate_cu_files.py, do not edit manually.
#include "../mmq.cuh"
DECL_MMQ_CASE(GGML_TYPE_DT3);
ggml/src/ggml-metal/ggml-metal-device.m
+2 -1
View File
@@ -1268,8 +1268,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_OP_ARGSORT:
case GGML_OP_TOP_K:
case GGML_OP_ARANGE:
case GGML_OP_ROLL:
return true;
case GGML_OP_ROLL:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_FLASH_ATTN_EXT:
// for new head sizes, add checks here
if (op->src[0]->ne[0] != 32 &&
ggml/src/ggml-opencl/ggml-opencl.cpp
+18
View File
@@ -73,6 +73,7 @@ typedef const void * (*get_adreno_bin_kernel_func_t)(
//------------------------------------------------------------------------------
bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor);
static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor);
static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor);
static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
@@ -4629,6 +4630,23 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac
if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
opts += " -D FA_C8_NO_SG_PIN";
}
// Transposed K tile in local memory: the KV rows the QK loop walks together become
// adjacent, so a group of them is ONE 128-bit local read instead of several narrow
// ones. The QK loop is LDS-read-issue-bound (a wrong-math probe that kept every FMA/dp4a
// but removed the LDS reads ran the kernel ~40% faster), so this is worth up to +26% on
// fa=1 prefill. Output is bit-identical -- only the layout moves.
//
// DK <= 128 only. At DK=256 (gemma-3-4b) it measures 1-2% NEGATIVE and reproduces across
// rounds; padding the row stride does not recover it, so the cause is not a simple bank
// conflict and the wider tile does not want this layout.
//
// Default on within that gate; GGML_OPENCL_FA_K_LDS_T=0 restores the row-major tile.
{
const char * e = getenv("GGML_OPENCL_FA_K_LDS_T");
if ((e == nullptr || e[0] != '0') && cfg->dk <= 128) {
opts += " -D FA_K_LDS_T";
}
}
return opts;
}
ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
+46 -5
View File
@@ -211,7 +211,30 @@ __kernel void FA_TILE_NAME(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_K_LDS_T
// K tile transposed: [dk vec][kv row] instead of [kv row][dk vec].
//
// The QK loop walks 2 or 4 KV rows at a time against the same dk element. Row-major
// those are DK_VEC half4s apart, so each is its own 64-bit local read. Transposed they
// are adjacent, so a pair is one 128-bit read -- half the LDS issues for the same bytes,
// no extra registers, arithmetic untouched.
//
// This kernel looked like it should be FMA-bound (a half4 mad does ~4 ALU ops per LDS
// read, unlike the 1:1 of the dp4a loop), but it is NOT: a wrong-math probe that kept
// every FMA and removed the LDS reads ran it 38.6% faster (18.92 -> 11.62 ms/op).
// Explicitly 16-byte aligned: FA_LK_PAIR below reads two adjacent half4 as one float4,
// and the element type only obliges the compiler to align this array to 8. The indices
// are even so the offset is a multiple of 16, but the base has to be too, and relying
// on the compiler to over-align it is relying on luck.
__local KV_DATA_TYPE4 l_k[DK_VEC][BLOCK_N] __attribute__((aligned(16)));
#define FA_LK(ROW, C) l_k[C][ROW]
// Two adjacent KV rows as one 128-bit local read (half4 pair == 16 B). j is even and
// BLOCK_N is even, so &l_k[c][j] is 16 B past a 16 B-aligned base.
#define FA_LK_PAIR(C, J) as_half8(*(__local const float4 *)(&l_k[C][J]))
#else
__local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC];
#define FA_LK(ROW, C) l_k[ROW][C]
#endif
__local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
#if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE)
@@ -254,17 +277,17 @@ __kernel void FA_TILE_NAME(
#ifdef FA_K_IMG
if (use_kv_pad) {
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
} else {
const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row;
l_k[row][col] = read_imageh(k_img, k_row_px + col);
FA_LK(row, col) = read_imageh(k_img, k_row_px + col);
}
#else
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
#endif
} else {
l_k[row][col] = (KV_DATA_TYPE4)(0.0h);
FA_LK(row, col) = (KV_DATA_TYPE4)(0.0h);
}
}
for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) {
@@ -292,8 +315,15 @@ __kernel void FA_TILE_NAME(
FA_UNROLL
for (int k = 0; k < SPLIT_DK_VEC; k++) {
const ACC_TYPE4 qk = q_priv[k];
#if defined(FA_K_LDS_T)
// 2 KV rows adjacent in the transposed tile: one 128-bit local read.
const half8 kk = FA_LK_PAIR(dk_off + k, j);
ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(kk.lo);
ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(kk.hi);
#else
ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]);
ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]);
#endif
partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3;
partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3;
}
@@ -359,7 +389,7 @@ __kernel void FA_TILE_NAME(
ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
FA_UNROLL
for (int k = 0; k < SPLIT_DK_VEC; k++) {
dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc);
dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(FA_LK(j, dk_off + k)), dot_acc);
}
local_partial[j][tid] =
dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3;
@@ -452,10 +482,21 @@ __kernel void FA_TILE_NAME(
FA_UNROLL
for (int k = 0; k < DK_VEC; k++) {
const ACC_TYPE4 qk = q_priv[k];
#if defined(FA_K_LDS_T)
// 4 KV rows adjacent in the transposed tile: two 128-bit local reads
// instead of four 64-bit ones.
const half8 kk01 = FA_LK_PAIR(k, j);
const half8 kk23 = FA_LK_PAIR(k, j + 2);
dot_acc0 = mad(qk, CONVERT_KV_ACC4(kk01.lo), dot_acc0);
dot_acc1 = mad(qk, CONVERT_KV_ACC4(kk01.hi), dot_acc1);
dot_acc2 = mad(qk, CONVERT_KV_ACC4(kk23.lo), dot_acc2);
dot_acc3 = mad(qk, CONVERT_KV_ACC4(kk23.hi), dot_acc3);
#else
dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0);
dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1);
dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2);
dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3);
#endif
}
ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale;
ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
+43 -4
View File
@@ -1631,8 +1631,25 @@ __kernel void flash_attn_f32_q4_0(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_HAVE_INT_DOT
// Accessors so the staging code is layout-agnostic.
#ifdef FA_K_LDS_T
#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW]
#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW]
#else
#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX]
#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK]
#endif
#ifdef FA_K_LDS_T
// K tile transposed: the 4 KV rows the QK loop walks together become adjacent, so each
// (block, group) step is ONE 128-bit local read instead of four 32-bit ones. The QK
// loop is LDS-read-issue-bound.
__local uint l_k_packed[DK_Q4_BLOCKS_PREFILL * 8][BLOCK_N];
__local float l_k_scale [DK_Q4_BLOCKS_PREFILL][BLOCK_N];
#else
__local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8];
__local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL];
#endif
#else
__local half4 l_k[BLOCK_N][DK_VEC];
#endif
@@ -1660,17 +1677,17 @@ __kernel void flash_attn_f32_q4_0(
const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE;
const float df = (float) vload_half(0, (const global half *) blk_ptr);
const global uchar * qs = (const global uchar *)(blk_ptr + 2);
l_k_scale[row][blk] = df;
FA_K_SCALE(row, blk) = df;
uint k_packed[8];
pack_q4_0_nibbles(qs, k_packed);
#pragma unroll
for (int j = 0; j < 8; ++j) {
l_k_packed[row][blk * 8 + j] = k_packed[j];
FA_K_PACKED(row, blk * 8 + j) = k_packed[j];
}
} else {
l_k_scale[row][blk] = 0.0f;
FA_K_SCALE(row, blk) = 0.0f;
#pragma unroll
for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u;
for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u;
}
}
#else
@@ -1760,6 +1777,19 @@ __kernel void flash_attn_f32_q4_0(
for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) {
const int b = k_blk_base + b_local;
int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
#ifdef FA_K_LDS_T
// 4 KV rows are adjacent in the transposed tile: one 128-bit local
// read per (block, group) instead of four 32-bit ones.
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]);
sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0);
sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1);
sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3);
}
#else
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
@@ -1768,12 +1798,21 @@ __kernel void flash_attn_f32_q4_0(
sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3);
}
#endif
const float qd = q_d_pf[b_local];
const int q_sum = q_sum_pf[b_local];
#ifdef FA_K_LDS_T
const float4 ks4 = vload4(0, &l_k_scale[b][j]);
s0 += (float)(sum0 - 8 * q_sum) * qd * ks4.s0;
s1 += (float)(sum1 - 8 * q_sum) * qd * ks4.s1;
s2 += (float)(sum2 - 8 * q_sum) * qd * ks4.s2;
s3 += (float)(sum3 - 8 * q_sum) * qd * ks4.s3;
#else
s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b];
s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b];
s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b];
s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b];
#endif
}
#else
ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);
ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
+49 -4
View File
@@ -1393,8 +1393,31 @@ __kernel void flash_attn_f32_q8_0(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_HAVE_INT_DOT
// Accessors so the staging code is layout-agnostic.
#ifdef FA_K_LDS_T
#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW]
#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW]
#else
#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX]
#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK]
#endif
#ifdef FA_K_LDS_T
// K tile transposed: [block*8 + g][kv row] instead of [kv row][block*8 + g].
//
// The QK loop walks 4 KV rows at a time against the same (b, g), so in the original
// layout those 4 values are BLOCK_N*8 uints apart and cost 4 separate 32-bit local
// reads. Transposed they are adjacent, so they are one 128-bit read -- 4x fewer LDS
// issues for the same bytes and no extra registers. That matters because the QK loop
// is LDS-read-issue-bound: a wrong-math probe that kept every dp4a but cut the LDS
// reads ran the whole kernel 41% faster (18.51 -> 10.91 ms/op), and deleting QK
// outright only reached 10.88 -- i.e. essentially ALL of QK's cost is these reads.
__local uint l_k_packed[DK_Q8_BLOCKS_PREFILL * 8][BLOCK_N];
__local float l_k_scale [DK_Q8_BLOCKS_PREFILL][BLOCK_N];
#else
__local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8];
__local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL];
#endif
#else
__local half4 l_k[BLOCK_N][DK_VEC];
#endif
@@ -1427,7 +1450,7 @@ __kernel void flash_attn_f32_q8_0(
const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE;
const float df = (float) vload_half(0, (const global half *) blk_ptr);
const global uchar * qs = (const global uchar *)(blk_ptr + 2);
l_k_scale[row][blk] = df;
FA_K_SCALE(row, blk) = df;
#pragma unroll
for (int j = 0; j < 8; ++j) {
uint k_packed =
@@ -1435,12 +1458,12 @@ __kernel void flash_attn_f32_q8_0(
((uint) qs[j*4 + 1]) << 8 |
((uint) qs[j*4 + 2]) << 16 |
((uint) qs[j*4 + 3]) << 24;
l_k_packed[row][blk * 8 + j] = k_packed;
FA_K_PACKED(row, blk * 8 + j) = k_packed;
}
} else {
l_k_scale[row][blk] = 0.0f;
FA_K_SCALE(row, blk) = 0.0f;
#pragma unroll
for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u;
for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u;
}
}
#else
@@ -1556,6 +1579,19 @@ __kernel void flash_attn_f32_q8_0(
for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) {
const int b = k_blk_base + b_local;
int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
#if defined(FA_K_LDS_T)
// The 4 KV rows are adjacent in the transposed tile, so each (b, g)
// step is ONE 128-bit local read instead of four 32-bit ones.
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]);
sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0);
sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1);
sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3);
}
#else
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
@@ -1564,11 +1600,20 @@ __kernel void flash_attn_f32_q8_0(
sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3);
}
#endif
const float qd = q_d_pf[b_local];
#ifdef FA_K_LDS_T
const float4 ks4 = vload4(0, &l_k_scale[b][j]);
s0 += (float)sum0 * qd * ks4.s0;
s1 += (float)sum1 * qd * ks4.s1;
s2 += (float)sum2 * qd * ks4.s2;
s3 += (float)sum3 * qd * ks4.s3;
#else
s0 += (float)sum0 * qd * l_k_scale[j ][b];
s1 += (float)sum1 * qd * l_k_scale[j+1][b];
s2 += (float)sum2 * qd * l_k_scale[j+2][b];
s3 += (float)sum3 * qd * l_k_scale[j+3][b];
#endif
}
#else
ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+24 -3
View File
@@ -5220,6 +5220,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f32_f32", arr_dmmv_bf16_f32_f32_len[reduc], arr_dmmv_bf16_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f32_f32", arr_dmmv_q1_0_f32_f32_len[reduc], arr_dmmv_q1_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f32_f32", arr_dmmv_q2_0_f32_f32_len[reduc], arr_dmmv_q2_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_DT3 ][i], "mul_mat_vec_dt3_f32_f32", arr_dmmv_dt3_f32_f32_len[reduc], arr_dmmv_dt3_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32", arr_dmmv_q4_0_f32_f32_len[reduc], arr_dmmv_q4_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32", arr_dmmv_q4_1_f32_f32_len[reduc], arr_dmmv_q4_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32", arr_dmmv_q5_0_f32_f32_len[reduc], arr_dmmv_q5_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
@@ -5247,6 +5248,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f16_f32", arr_dmmv_bf16_f16_f32_len[reduc], arr_dmmv_bf16_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f16_f32", arr_dmmv_q1_0_f16_f32_len[reduc], arr_dmmv_q1_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f16_f32", arr_dmmv_q2_0_f16_f32_len[reduc], arr_dmmv_q2_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_DT3 ][i], "mul_mat_vec_dt3_f16_f32", arr_dmmv_dt3_f16_f32_len[reduc], arr_dmmv_dt3_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32", arr_dmmv_q4_0_f16_f32_len[reduc], arr_dmmv_q4_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32", arr_dmmv_q4_1_f16_f32_len[reduc], arr_dmmv_q4_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32", arr_dmmv_q5_0_f16_f32_len[reduc], arr_dmmv_q5_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
@@ -5362,6 +5364,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16", dequant_f32_len, dequant_f32_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q1_0], "dequant_q1_0", dequant_q1_0_len, dequant_q1_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 8, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_0], "dequant_q2_0", dequant_q2_0_len, dequant_q2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_DT3 ], "dequant_dt3", dequant_dt3_len, dequant_dt3_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
@@ -5390,6 +5393,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_BF16], "get_rows_bf16", get_rows_bf16_len, get_rows_bf16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q1_0], "get_rows_q1_0", get_rows_q1_0_len, get_rows_q1_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_0], "get_rows_q2_0", get_rows_q2_0_len, get_rows_q2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_DT3 ], "get_rows_dt3", get_rows_dt3_len, get_rows_dt3_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
@@ -5418,6 +5422,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_BF16], "get_rows_bf16_f32", get_rows_bf16_f32_len, get_rows_bf16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q1_0], "get_rows_q1_0_f32", get_rows_q1_0_f32_len, get_rows_q1_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_0], "get_rows_q2_0_f32", get_rows_q2_0_f32_len, get_rows_q2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_DT3 ], "get_rows_dt3_f32", get_rows_dt3_f32_len, get_rows_dt3_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
@@ -7617,6 +7622,7 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type
case GGML_TYPE_F32:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_DT3:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@@ -7760,6 +7766,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context *
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_DT3:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@@ -9171,12 +9178,18 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0;
// DT3 weights (d1*t1 + d2*t2, two fp16-scaled ternary planes) already pay
// one fp16 rounding in the dequant fallback; fp16 accumulation on top of
// it costs measurable perplexity. Force fp32 accumulators, matching the
// numerics of the CUDA GEMM fallback (fp16 inputs, fp32 compute).
const ggml_prec mm_prec = src0->type == GGML_TYPE_DT3 ? GGML_PREC_F32 : (ggml_prec)dst->op_params[0];
// Check for mmq first
vk_matmul_pipeline mmp = quantize_y ? ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, GGML_TYPE_Q8_1, (ggml_prec)dst->op_params[0]) : nullptr;
vk_matmul_pipeline mmp = quantize_y ? ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, GGML_TYPE_Q8_1, mm_prec) : nullptr;
if (mmp == nullptr) {
// Fall back to f16 dequant mul mat
mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? f16_type : src1->type, (ggml_prec)dst->op_params[0]);
mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? f16_type : src1->type, mm_prec);
quantize_y = false;
}
@@ -9185,7 +9198,7 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
if (qx_needs_dequant) {
// Fall back to dequant + f16 mulmat
mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]);
mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, mm_prec);
}
// Not implemented
@@ -17987,6 +18000,13 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
}
}
switch (src0_type) {
case GGML_TYPE_DT3:
// DT3 has dequant, get_rows and scalar mul_mat_vec shaders only:
// mul_mat_id, coopmat and MMQ are intentionally not implemented
if (op->op == GGML_OP_MUL_MAT_ID) {
return false;
}
break;
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
@@ -18097,6 +18117,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_DT3:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
ggml/src/ggml-vulkan/vulkan-shaders/dequant_dt3.comp
+47
View File
@@ -0,0 +1,47 @@
#version 450
#include "dequant_head.glsl"
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
layout (binding = 0) readonly buffer A {block_dt3 data_a[];};
layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
// Eight blocks per workgroup, 32 threads per block. Threads 0..23 decode one
// qs byte of each plane (5 trits in base 3), threads 24..25 decode one qh
// byte of each plane (4 trits — the 5th base-3 digit of a qh byte is packing
// padding that always decodes to -1, so it must not be read), threads 26..31
// idle.
void main() {
const uint ib = gl_WorkGroupID.x * 8 + gl_LocalInvocationID.x / 32;
const uint il = gl_LocalInvocationID.x % 32;
if (ib >= p.nel / 128 || il >= 26) {
return;
}
const float d1 = float(data_a[ib].d[0]);
const float d2 = float(data_a[ib].d[1]);
const uint b_idx = ib * 128;
// element covered by the first digit, distance between consecutive
// digits, and number of digits stored in this byte
const bool is_qh = il >= 24;
const uint e0 = is_qh ? 120 + (il - 24) : (il < 16 ? il : 80 + (il - 16));
const uint stride = is_qh ? 2 : (il < 16 ? 16 : 8);
const uint digits = is_qh ? 4 : 5;
uint q1 = is_qh ? uint(data_a[ib].qh[il - 24]) : uint(data_a[ib].qs[il]);
uint q2 = is_qh ? uint(data_a[ib].qh[2 + il - 24]) : uint(data_a[ib].qs[24 + il]);
// decode each byte once: take the top base-3 digit with (q*3) >> 8, then
// shift it out with q <- (q*3) mod 256
for (uint n = 0; n < digits; ++n) {
const float t1 = float(int((q1 * 3) >> 8) - 1);
const float t2 = float(int((q2 * 3) >> 8) - 1);
data_b[b_idx + e0 + n*stride] = D_TYPE(d1*t1 + d2*t2);
q1 = (q1 * 3) & 0xFF;
q2 = (q2 * 3) & 0xFF;
}
}
ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl
+50
View File
@@ -154,6 +154,49 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
}
#endif
#if defined(DATA_A_DT3)
// Dual-plane ternary: element iqs of plane p sits in a base-3 packed byte.
// Elements 0..79 use qs[m], m = iqs % 16, digit n = iqs / 16; elements
// 80..119 use qs[16 + m], m = (iqs - 80) % 8, digit n = (iqs - 80) / 8;
// elements 120..127 use qh[j], j = iqs % 2, digit n = (iqs - 120) / 2.
// A qh byte holds only 4 trits: its 5th base-3 digit is packing padding that
// always decodes to -1, never to 0, so it must not be read.
// The decode multiplies the byte by 3^n modulo 256 and takes the top digit.
float dt3_get_trit(uint ib, uint p, uint iqs, uint a_offset) {
const uint pow3[5] = {1, 3, 9, 27, 81};
uint b;
uint n;
if (iqs < 80) {
b = uint(data_a[a_offset + ib].qs[p*24 + (iqs & 15)]);
n = iqs >> 4;
} else if (iqs < 120) {
b = uint(data_a[a_offset + ib].qs[p*24 + 16 + ((iqs - 80) & 7)]);
n = (iqs - 80) >> 3;
} else {
b = uint(data_a[a_offset + ib].qh[p*2 + (iqs & 1)]);
n = (iqs - 120) >> 1;
}
const uint q = (b * pow3[n]) & 0xFF;
return float(int((q * 3) >> 8) - 1);
}
// w = d1*t1 + d2*t2; both products are exact (t in {-1,0,+1}), so the sum has
// a single float rounding and matches the CPU reference bit by bit
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
const float d1 = float(data_a[a_offset + ib].d[0]);
const float d2 = float(data_a[a_offset + ib].d[1]);
return vec2(d1*dt3_get_trit(ib, 0, iqs, a_offset) + d2*dt3_get_trit(ib, 1, iqs, a_offset),
d1*dt3_get_trit(ib, 0, iqs + 1, a_offset) + d2*dt3_get_trit(ib, 1, iqs + 1, a_offset));
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
const float d1 = float(data_a[a_offset + ib].d[0]);
const float d2 = float(data_a[a_offset + ib].d[1]);
return vec4(d1*dt3_get_trit(ib, 0, iqs, a_offset) + d2*dt3_get_trit(ib, 1, iqs, a_offset),
d1*dt3_get_trit(ib, 0, iqs + 1, a_offset) + d2*dt3_get_trit(ib, 1, iqs + 1, a_offset),
d1*dt3_get_trit(ib, 0, iqs + 2, a_offset) + d2*dt3_get_trit(ib, 1, iqs + 2, a_offset),
d1*dt3_get_trit(ib, 0, iqs + 3, a_offset) + d2*dt3_get_trit(ib, 1, iqs + 3, a_offset));
}
#endif
#if defined(DATA_A_IQ1_S)
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
const uint ib32 = iqs / 32;
@@ -571,6 +614,13 @@ vec2 get_dm(uint ib, uint a_offset) {
}
#endif
#if defined(DATA_A_DT3)
// the two scales are already applied inside dequantize/dequantize4
vec2 get_dm(uint ib, uint a_offset) {
return vec2(1, 0);
}
#endif
#if defined(DATA_A_MXFP4)
vec2 get_dm(uint ib, uint a_offset) {
return vec2(e8m0_to_fp32(data_a[a_offset + ib].e), 0);
ggml/src/ggml-vulkan/vulkan-shaders/types.glsl
+21
View File
@@ -235,6 +235,27 @@ struct block_q2_0_packed16
#define DATA_A_QUANT_LEGACY
#endif
#define QUANT_K_DT3 128
#define QUANT_R_DT3 1
// Dual-plane ternary: w = d[0]*t1 + d[1]*t2 with trits in {-1,0,+1}.
// Per plane: 24 bytes with 5 trits each in base 3 (elements 0..119), then
// 2 bytes with 4 trits each (elements 120..127). Plane p uses qs[p*24..],
// qh[p*2..] and d[p].
struct block_dt3
{
uint8_t qs[2*24];
uint8_t qh[2*2];
float16_t d[2];
};
#if defined(DATA_A_DT3)
#define QUANT_K QUANT_K_DT3
#define QUANT_R QUANT_R_DT3
#define QUANT_AUXF 1
#define A_TYPE block_dt3
#endif
#define QUANT_K_Q8_1 32
#define QUANT_R_Q8_1 1
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+15 -4
View File
@@ -51,6 +51,7 @@ const std::vector<std::string> type_names = {
"f16",
"q1_0",
"q2_0",
"dt3",
"q4_0",
"q4_1",
"q5_0",
@@ -591,6 +592,12 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
continue;
}
// DT3 has no direct matmul shaders: mul_mat goes through dequant to
// f16 + f16 matmul, and coopmat/MMQ are intentionally not implemented
if (tname == "dt3") {
continue;
}
std::string data_a_key = "DATA_A_" + to_uppercase(tname);
// For aligned matmul loads
std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16" || tname == "bf16") ? load_vec : load_vec_quant;
@@ -758,9 +765,12 @@ void process_shaders() {
}
#endif
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
// mul_mat_id is not implemented for DT3 (supports_op declines it)
if (tname != "dt3") {
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
}
// mul mat vec with integer dot product
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
@@ -1254,7 +1264,8 @@ void write_output_files() {
src << "const uint64_t arr_dmmv_" << tname << "_" << btype << "_f32_len[3] = {mul_mat_vec_" << tname << "_" << btype << "_f32_len, mul_mat_vec_" << tname << "_" << btype << "_f32_subgroup_len, mul_mat_vec_" << tname << "_" << btype << "_f32_subgroup_no_shmem_len};\n";
}
if (btype == "f16") {
if (btype == "f16" || tname == "dt3") {
// no mul_mat_vec_id shaders for DT3
continue;
}
hdr << "extern const void * arr_dmmv_id_" << tname << "_" << btype << "_f32_data[3];\n";
gguf-py/gguf/constants.py
+24 -2
View File
@@ -509,6 +509,7 @@ class MODEL_ARCH(IntEnum):
OLMO = auto()
OLMO2 = auto()
OLMOE = auto()
MUSE_GLIMMER = auto()
OPENELM = auto()
ARCTIC = auto()
DEEPSEEK = auto()
@@ -1181,6 +1182,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.OLMO: "olmo",
MODEL_ARCH.OLMO2: "olmo2",
MODEL_ARCH.OLMOE: "olmoe",
MODEL_ARCH.MUSE_GLIMMER: "muse-glimmer",
MODEL_ARCH.OPENELM: "openelm",
MODEL_ARCH.ARCTIC: "arctic",
MODEL_ARCH.DEEPSEEK: "deepseek",
@@ -1562,8 +1564,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.V_MM_UP: "mm.up",
MODEL_TENSOR.V_MM_DOWN: "mm.down",
MODEL_TENSOR.V_MM_GATE: "mm.gate",
MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1",
MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2",
MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1",
MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2",
MODEL_TENSOR.V_TOK_BOI: "v.boi",
MODEL_TENSOR.V_TOK_EOI: "v.eoi",
MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm",
@@ -3331,6 +3333,25 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_UP_EXP,
MODEL_TENSOR.FFN_DOWN_EXP,
],
MODEL_ARCH.MUSE_GLIMMER: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_Q_NORM,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_K_NORM,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.ATTN_GATE,
MODEL_TENSOR.FFN_GATE,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_POST_NORM,
MODEL_TENSOR.FFN_PRE_NORM,
MODEL_TENSOR.FFN_POST_NORM,
],
MODEL_ARCH.OPENELM: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
@@ -5168,6 +5189,7 @@ class VisionProjectorType:
MIMOVL = "mimovl"
MIMO_AUDIO = "mimo_audio"
GRANITE4_VISION = "granite4_vision"
MUSE_GLIMMER = "muse-glimmer"
# Items here are (block size, type size)
gguf-py/gguf/tensor_mapping.py
+18 -5
View File
@@ -382,7 +382,7 @@ class TensorNameMap:
),
MODEL_TENSOR.ATTN_GATE: (
"model.layers.{bid}.self_attn.gate_proj", # afmoe
"model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer
"model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5
"model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate
),
@@ -1298,10 +1298,12 @@ class TensorNameMap:
"encoder.final_layer_norm", # t5
"layer_norm", # neobert
"model.hidden_norm", # dflash
"encoder.output_norm_enc", # dflash (transformers MuseGlimmerAssistant)
),
MODEL_TENSOR.FC: (
"model.fc", # dflash
"model.fc", # dflash
"encoder.fc", # dflash (transformers MuseGlimmerAssistant)
),
MODEL_TENSOR.DSPARK_MARKOV_W1: (
@@ -1467,6 +1469,7 @@ class TensorNameMap:
"vision_tower.patch_embed.patchifier.proj", # dots.ocr
"vision_model.conv1", # Step3-VL
"model.vision_embedder.patch_dense", # gemma4 unified
"model.vision_tower.patch_embedder.patch_embedding", # muse-glimmer
),
MODEL_TENSOR.V_ENC_EMBD_NORM: (
@@ -1534,7 +1537,8 @@ class TensorNameMap:
"siglip2.vision_model.encoder.layers.{bid}.self_attn.q_proj", # youtuvl
"model.vision_model.transformer.layers.{bid}.self_attn.q_proj", # Deepseek-OCR CLIP, generated
"vision_model.model.layers.{bid}.self_attn.q_proj.linear", # gemma4
"model.qwen2_model.model.model.layers.{bid}.self_attn.q_proj" # Deepseek-OCR-2 qwen2
"model.qwen2_model.model.model.layers.{bid}.self_attn.q_proj", # Deepseek-OCR-2 qwen2
"model.vision_tower.layers.{bid}.attn.q_proj", # muse-glimmer
),
MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
@@ -1560,7 +1564,8 @@ class TensorNameMap:
"model.vision_model.transformer.layers.{bid}.self_attn.k_proj", # Deepseek-OCR CLIP, generated
"siglip2.vision_model.encoder.layers.{bid}.self_attn.k_proj",
"vision_model.model.layers.{bid}.self_attn.k_proj.linear", # gemma4
"model.qwen2_model.model.model.layers.{bid}.self_attn.k_proj" # Deepseek-OCR-2 qwen2
"model.qwen2_model.model.model.layers.{bid}.self_attn.k_proj", # Deepseek-OCR-2 qwen2
"model.vision_tower.layers.{bid}.attn.k_proj", # muse-glimmer
),
MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
@@ -1586,7 +1591,8 @@ class TensorNameMap:
"siglip2.vision_model.encoder.layers.{bid}.self_attn.v_proj",
"model.vision_model.transformer.layers.{bid}.self_attn.v_proj", # Deepseek-OCR CLIP, generated
"vision_model.model.layers.{bid}.self_attn.v_proj.linear", # gemma4
"model.qwen2_model.model.model.layers.{bid}.self_attn.v_proj" # Deepseek-OCR-2 qwen2
"model.qwen2_model.model.model.layers.{bid}.self_attn.v_proj", # Deepseek-OCR-2 qwen2
"model.vision_tower.layers.{bid}.attn.v_proj", # muse-glimmer
),
MODEL_TENSOR.V_ENC_INPUT_NORM: (
@@ -1610,6 +1616,7 @@ class TensorNameMap:
"vision_tower.blocks.{bid}.norm1", # dots.ocr
"vision_model.transformer.resblocks.{bid}.ln_1", # Step3-VL
"model.qwen2_model.model.model.layers.{bid}.input_layernorm", # Deepseek-OCR-2 qwen2
"model.vision_tower.layers.{bid}.norm1", # muse-glimmer
),
MODEL_TENSOR.V_ENC_ATTN_O: (
@@ -1635,6 +1642,7 @@ class TensorNameMap:
"vision_model.model.layers.{bid}.self_attn.o_proj.linear", # gemma4
"vision_tower.blocks.{bid}.attn.proj", # dots.ocr
"vision_model.transformer.resblocks.{bid}.attn.out_proj", # Step3-VL
"model.vision_tower.layers.{bid}.attn.proj", # muse-glimmer
),
MODEL_TENSOR.V_ENC_ATTN_SINKS: (
@@ -1663,6 +1671,7 @@ class TensorNameMap:
"vision_tower.blocks.{bid}.norm2", # dots.ocr
"vision_model.transformer.resblocks.{bid}.ln_2", # Step3-VL
"model.qwen2_model.model.model.layers.{bid}.post_attention_layernorm", # Deepseek-OCR-2 qwen2
"model.vision_tower.layers.{bid}.norm2", # muse-glimmer
),
MODEL_TENSOR.V_ENC_FFN_UP: (
@@ -1687,6 +1696,7 @@ class TensorNameMap:
"vision_model.model.layers.{bid}.mlp.up_proj", # gemma4
"vision_model.transformer.resblocks.{bid}.mlp.c_fc", # Step3-VL
"model.qwen2_model.model.model.layers.{bid}.mlp.up_proj", # Deepseek-OCR-2 qwen2
"model.vision_tower.layers.{bid}.mlp.fc1", # muse-glimmer
),
MODEL_TENSOR.V_ENC_FFN_GATE: (
@@ -1719,6 +1729,7 @@ class TensorNameMap:
"model.qwen2_model.model.model.layers.{bid}.mlp.down_proj" , # Deepseek-OCR-2 qwen2
"vision_model.model.layers.{bid}.mlp.down_proj", # gemma4
"vision_model.transformer.resblocks.{bid}.mlp.c_proj", # Step3-VL
"model.vision_tower.layers.{bid}.mlp.fc2", # muse-glimmer
),
MODEL_TENSOR.V_ENC_ATTN_POST_NORM: (
@@ -1753,6 +1764,7 @@ class TensorNameMap:
"model.vision_model.pre_layrnorm", # Deepseek-OCR CLIP
"vision_tower.patch_embed.patchifier.norm", # dots.ocr
"vision_model.ln_pre", # Step3-VL
"model.vision_tower.ln_pre", # muse-glimmer
),
MODEL_TENSOR.V_POST_NORM: (
@@ -1766,6 +1778,7 @@ class TensorNameMap:
"visual.post_layernorm", # glm4v
"siglip2.vision_model.post_layernorm",
"model.qwen2_model.model.model.norm", # Deepseek-OCR-2 qwen2
"model.vision_tower.ln_post", # muse-glimmer
),
MODEL_TENSOR.V_MM_POST_NORM: (
include/llama.h
+26 -10
View File
@@ -349,14 +349,15 @@ extern "C" {
// NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations
// https://github.com/ggml-org/llama.cpp/pull/7544
struct llama_context_params {
uint32_t n_ctx; // text context, 0 = from model
uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode
uint32_t n_ubatch; // physical maximum batch size
uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models)
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch)
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
uint32_t n_ctx; // text context, 0 = from model
uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode
uint32_t n_ubatch; // physical maximum batch size
uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models)
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch)
uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max)
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
enum llama_context_type ctx_type; // set the context type (e.g. MTP)
enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type`
@@ -1055,6 +1056,9 @@ extern "C" {
//
// Get the backend sampled token for the ith token.
// With multiple outputs, sampler state advances when the token is accepted,
// not when it is read through this function.
// When accepting multiple outputs, accept a contiguous prefix in output order.
// Returns LLAMA_TOKEN_NULL if no token was sampled.
LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i);
@@ -1271,9 +1275,12 @@ extern "C" {
// [EXPERIMENTAL]
// backend sampling interface:
// return true if the backend supports all ops needed by the sampler
// return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_max_per_seq outputs per sequence
// note: call once per sampler
bool (*backend_init)(struct llama_sampler * smpl, ggml_backend_buffer_type_t buft);
bool (*backend_init)(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq);
// call after .backend_apply()
void (*backend_accept)(
@@ -1291,6 +1298,13 @@ extern "C" {
// called before graph execution to set inputs for the current ubatch
void (*backend_set_input)(struct llama_sampler * smpl);
// called before rebuilding a sampling graph to clear any internal sampler state
void (*backend_reset)(struct llama_sampler * smpl);
// copy mutable state from src into dst while keeping dst's references to the current sampling graph
// src and dst must have the same type and configuration
void (*copy_state)(const struct llama_sampler * src, struct llama_sampler * dst);
};
struct llama_sampler {
@@ -1311,6 +1325,7 @@ extern "C" {
LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p);
LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl);
LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl);
LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst);
// important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add)
LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl);
@@ -1500,6 +1515,7 @@ extern "C" {
LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl);
/// @details Sample and accept a token from the idx-th output of the last evaluation
// For multiple outputs from one sampler, call this function in output order without gaps.
//
// Shorthand for:
// const auto * logits = llama_get_logits_ith(ctx, idx);
src/llama-arch.cpp
+1
View File
@@ -71,6 +71,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_OLMO, "olmo" },
{ LLM_ARCH_OLMO2, "olmo2" },
{ LLM_ARCH_OLMOE, "olmoe" },
{ LLM_ARCH_MUSE_GLIMMER, "muse-glimmer" },
{ LLM_ARCH_OPENELM, "openelm" },
{ LLM_ARCH_ARCTIC, "arctic" },
{ LLM_ARCH_DEEPSEEK, "deepseek" },
src/llama-arch.h
+1
View File
@@ -76,6 +76,7 @@ enum llm_arch {
LLM_ARCH_OLMO,
LLM_ARCH_OLMO2,
LLM_ARCH_OLMOE,
LLM_ARCH_MUSE_GLIMMER,
LLM_ARCH_OPENELM,
LLM_ARCH_ARCTIC,
LLM_ARCH_DEEPSEEK,
src/llama-context.cpp
+162 -147
View File
@@ -10,6 +10,7 @@
#include "llama-mmap.h"
#include "llama-model.h"
#include "llama-ext.h"
#include "llama-sampler.h"
#include "llama.h"
#include <cinttypes>
@@ -159,25 +160,6 @@ llama_context::llama_context(
}
}
// Initialize backend samplers here so they are part of the sampling graph
// before the reserve passes run later in this function. This avoids a later
// re-reserve when graph nodes change.
if (params.samplers != nullptr && params.n_samplers > 0) {
for (size_t i = 0; i < params.n_samplers; ++i) {
const auto & config = params.samplers[i];
if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {
throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");
}
if (set_sampler(config.seq_id, config.sampler)) {
const int n_samplers = llama_sampler_chain_n(config.sampler);
LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);
}
}
}
auto rope_scaling_type = params.rope_scaling_type;
if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) {
rope_scaling_type = hparams.rope_scaling_type_train;
@@ -265,6 +247,27 @@ llama_context::llama_context(
cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);
cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max;
cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ?
cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max);
// Initialize backend samplers here so they are part of the sampling graph
// before the reserve passes run later in this function. This avoids a later
// re-reserve when graph nodes change.
if (params.samplers != nullptr && params.n_samplers > 0) {
for (size_t i = 0; i < params.n_samplers; ++i) {
const auto & config = params.samplers[i];
if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {
throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");
}
if (set_sampler(config.seq_id, config.sampler)) {
const int n_samplers = llama_sampler_chain_n(config.sampler);
LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);
}
}
}
cparams.op_offload = params.op_offload;
cparams.kv_unified = params.kv_unified;
@@ -300,18 +303,19 @@ llama_context::llama_context(
}
}
LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max);
LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx);
LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq);
LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch);
LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch);
LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn);
LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type));
LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false");
LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base);
LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale);
LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq);
LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max);
LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max);
LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx);
LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq);
LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch);
LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch);
LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn);
LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type));
LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false");
LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base);
LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale);
LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq);
LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max);
LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n", __func__, cparams.n_outputs_max_per_seq);
if (cparams.n_ctx_seq < hparams.n_ctx_train) {
LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n",
@@ -1231,7 +1235,7 @@ bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) {
if (sampler && can_offload) {
auto * buft = ggml_backend_dev_buffer_type(model.dev_output());
sampler->iface->backend_init(sampler, buft);
sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq);
sampling.samplers[seq_id] = sampler;
@@ -1576,108 +1580,38 @@ int llama_context::encode(const llama_batch & batch_inp) {
return 0;
}
static std::map<llama_seq_id, uint32_t> build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) {
std::map<llama_seq_id, uint32_t> seq_to_row;
// how many output tokens we have seen so far for this ubatch.
uint32_t local = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
// skip tokens that are not output.
if (!ubatch.output[i]) {
continue;
}
const llama_seq_id seq_id = ubatch.seq_id[i][0];
// row_offset is the number of output tokens before this ubatch.
seq_to_row[seq_id] = row_offset + local;
++local;
}
return seq_to_row;
}
static void copy_tensor_async_ints(
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
const buffer_view<llama_token> & sampled,
const std::map<llama_seq_id, uint32_t> & seq_to_row,
ggml_backend_sched_t sched) {
if (!sampled.has_data()) {
return;
}
for (const auto & [seq_id, tensor] : tensor_map) {
auto it = seq_to_row.find(seq_id);
if (it == seq_to_row.end()) {
continue;
}
const uint32_t row = it->second;
GGML_ASSERT(row < sampled.size);
GGML_ASSERT(ggml_is_contiguous(tensor) && "sampled tokens tensor must be contiguous for async copy");
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
ggml_backend_tensor_get_async(backend, tensor, sampled.data + row, 0, sizeof(sampled.data[row]));
}
}
static void copy_tensor_async_floats(
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
const buffer_view<float> & dst,
template<typename T>
static void copy_tensor_async_rows(
const std::vector<ggml_tensor *> & tensors,
const buffer_view<T> & dst,
size_t stride,
std::vector<uint32_t> & counts,
const std::map<llama_seq_id, uint32_t> & seq_to_row,
ggml_backend_sched_t sched) {
uint32_t row_offset,
ggml_backend_sched_t sched,
std::vector<uint32_t> * counts = nullptr) {
if (!dst.has_data()) {
return;
}
for (const auto & [seq_id, tensor] : tensor_map) {
auto it = seq_to_row.find(seq_id);
if (it == seq_to_row.end()) {
for (size_t i = 0; i < tensors.size(); ++i) {
auto * tensor = tensors[i];
if (tensor == nullptr) {
continue;
}
const uint32_t row = it->second;
GGML_ASSERT(row < counts.size());
GGML_ASSERT(ggml_is_contiguous(tensor) && "logits/probs tensor must be contiguous for async copy");
const uint32_t row = row_offset + i;
const size_t n_elements = ggml_nelements(tensor);
GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy");
GGML_ASSERT(n_elements <= stride);
GGML_ASSERT((size_t) row * stride + n_elements <= dst.size);
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
float * row_ptr = dst.data + (size_t) row * stride;
T * row_ptr = dst.data + (size_t) row * stride;
ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
// Update the actual number of logits/probabilities that were written for this row.
counts[row] = ggml_nelements(tensor);
}
}
static void copy_tensor_async_candidates(
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
const buffer_view<llama_token> & dst,
size_t stride,
std::vector<uint32_t> & counts,
const std::map<llama_seq_id, uint32_t> & seq_to_row,
ggml_backend_sched_t sched) {
if (!dst.has_data()) {
return;
}
for (const auto & [seq_id, tensor] : tensor_map) {
auto it = seq_to_row.find(seq_id);
if (it == seq_to_row.end()) {
continue;
if (counts) {
GGML_ASSERT(row < counts->size());
(*counts)[row] = n_elements;
}
const uint32_t row = it->second;
GGML_ASSERT(row < counts.size());
GGML_ASSERT(ggml_is_contiguous(tensor) && "candidates tensor must be contiguous for async copy");
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
llama_token * row_ptr = dst.data + (size_t) row * stride;
ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
// Update the actual number of candidates that were written.
counts[row] = ggml_nelements(tensor);
}
}
@@ -1726,12 +1660,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;
// TODO: avoid this workaround in the future
if (has_samplers && batch_inp.logits) {
// embedding contexts output every token even when batch.logits is not set
if (has_samplers && (output_all || batch_inp.logits)) {
std::vector<int32_t> seq_output_count(n_seq_max, 0);
for (int32_t i = 0; i < batch_inp.n_tokens; ++i) {
if (batch_inp.logits[i] == 0) {
if (!output_all && batch_inp.logits[i] == 0) {
continue;
}
@@ -1740,10 +1674,17 @@ int llama_context::decode(const llama_batch & batch_inp) {
for (int32_t s = 0; s < ns; ++s) {
const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0;
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
continue;
}
seq_output_count[seq_id]++;
if (seq_output_count[seq_id] > 1) {
LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n",
__func__, seq_id, seq_output_count[seq_id]);
auto sampler = sampling.samplers.find(seq_id);
if (sampler != sampling.samplers.end() &&
seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) {
LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence "
"(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq,
seq_id, seq_output_count[seq_id]);
return -1;
}
}
@@ -1843,6 +1784,11 @@ int llama_context::decode(const llama_batch & batch_inp) {
return -2;
};
// start a new sampling transaction for this logical batch
for (const auto & entry : sampling.samplers) {
llama_sampler_backend_begin(entry.second);
}
int64_t n_outputs_prev = 0;
int64_t n_tokens_prev = 0;
@@ -2009,17 +1955,14 @@ int llama_context::decode(const llama_batch & batch_inp) {
}
}
// Copy backend sampling output if this ubatch produced any sampling tensors.
if (has_samplers && (!res->t_sampled.empty() || !res->t_sampled_probs.empty() || !res->t_sampled_logits.empty())) {
const auto seq_to_output_row = build_seq_to_output_row(ubatch, n_outputs_prev);
if (has_samplers) {
const auto stride = n_vocab;
// async copy the sampling data from the backend to the host
copy_tensor_async_ints(res->t_sampled, sampling.sampled, seq_to_output_row, sched.get());
copy_tensor_async_floats (res->t_sampled_logits, sampling.logits, stride, sampling.logits_count, seq_to_output_row, sched.get());
copy_tensor_async_floats (res->t_sampled_probs, sampling.probs, stride, sampling.probs_count, seq_to_output_row, sched.get());
copy_tensor_async_candidates(res->t_candidates, sampling.candidates, stride, sampling.candidates_count, seq_to_output_row, sched.get());
copy_tensor_async_rows(res->t_sampled, sampling.sampled, 1, n_outputs_prev, sched.get());
copy_tensor_async_rows(res->t_sampled_logits, sampling.logits, stride, n_outputs_prev, sched.get(), &sampling.logits_count);
copy_tensor_async_rows(res->t_sampled_probs, sampling.probs, stride, n_outputs_prev, sched.get(), &sampling.probs_count);
copy_tensor_async_rows(res->t_candidates, sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count);
}
n_outputs_prev += n_outputs;
@@ -2349,6 +2292,7 @@ void llama_context::output_reorder() {
//
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
uint32_t res;
if (model.arch == LLM_ARCH_QWEN3NEXT ||
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
@@ -2357,11 +2301,31 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
(model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
model.arch == LLM_ARCH_NANBEIGE ||
model.arch == LLM_ARCH_MINIMAX_M3) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
} else {
res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
for (const auto & lora : model.loras) {
res += lora->get_n_nodes();
}
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
for (const auto & lora : model.loras) {
res += lora->get_n_nodes();
uint32_t n_sampling_nodes = 0;
uint32_t n_sampling_nodes_max = 0;
for (const auto & [seq_id, sampler] : sampling.samplers) {
const uint32_t n_nodes = llama_sampler_backend_n_nodes(sampler);
n_sampling_nodes += n_nodes;
if (cparams.n_outputs_max_per_seq > 1) {
n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes);
}
}
const uint32_t n_sampling_outputs_max = std::min<uint64_t>(
std::min(n_tokens, cparams.n_outputs_max),
(uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq);
res += n_sampling_nodes;
if (n_sampling_outputs_max > 1) {
res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max;
}
return res;
}
@@ -2370,6 +2334,63 @@ llm_graph_result * llama_context::get_gf_res_reserve() const {
return static_cast<llm_graph_result *>(gf_res_reserve.get());
}
// pack sampler outputs into as few sequences as possible before using sequences without samplers
static void ubatch_prepare_reserve(
llama_ubatch & ubatch,
uint32_t n_outputs,
const std::map<llama_seq_id, llama_sampler *> & samplers,
uint32_t n_outputs_max_per_seq) {
const uint32_t n_seqs = ubatch.n_seqs;
const uint32_t n_seq_tokens = ubatch.n_seq_tokens;
for (uint32_t s = 0; s < n_seqs; ++s) {
for (uint32_t t = 0; t < n_seq_tokens; ++t) {
const uint32_t i = s * n_seq_tokens + t;
ubatch.n_seq_id[i] = 1;
ubatch.seq_id[i] = &ubatch.seq_id_unq[s];
}
}
// sequences with a sampler that fit in this ubatch
std::vector<uint32_t> sampler_seqs;
std::vector<bool> has_sampler(n_seqs, false);
for (const auto & entry : samplers) {
const llama_seq_id seq_id = entry.first;
if (seq_id < 0 || (uint32_t) seq_id >= n_seqs) {
continue;
}
sampler_seqs.push_back(seq_id);
has_sampler[seq_id] = true;
}
uint32_t n_outputs_set = 0;
const uint32_t n_outputs_per_seq = std::min(n_seq_tokens, n_outputs_max_per_seq);
for (uint32_t s : sampler_seqs) {
if (n_outputs_set >= n_outputs) {
break;
}
for (uint32_t t = 0; t < n_outputs_per_seq && n_outputs_set < n_outputs; ++t) {
ubatch.output[s * n_seq_tokens + t] = true;
++n_outputs_set;
}
}
// use sequences without samplers for any remaining outputs
for (uint32_t t = 0; t < n_seq_tokens && n_outputs_set < n_outputs; ++t) {
for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) {
if (has_sampler[s]) {
continue;
}
ubatch.output[s * n_seq_tokens + t] = true;
++n_outputs_set;
}
}
}
ggml_cgraph * llama_context::graph_reserve(
uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only, size_t * sizes) {
LLAMA_LOG_DEBUG("%s: reserving a graph for ubatch with n_tokens = %4u, n_seqs = %2u, n_outputs = %4u\n", __func__, n_tokens, n_seqs, n_outputs);
@@ -2394,14 +2415,7 @@ ggml_cgraph * llama_context::graph_reserve(
llama_batch_allocr balloc(model.hparams.n_pos_per_embd());
llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs);
// set one output token per sequence in order to activate all backend samplers
std::vector<llama_seq_id> seq_ids(n_seqs);
for (uint32_t i = 0; i < n_seqs; ++i) {
seq_ids[i] = i;
ubatch.n_seq_id[i] = 1;
ubatch.seq_id[i] = &seq_ids[i];
ubatch.output[i] = true;
}
ubatch_prepare_reserve(ubatch, n_outputs, sampling.samplers, cparams.n_outputs_max_per_seq);
auto * res = gf_res_reserve.get();
@@ -3488,6 +3502,7 @@ llama_context_params llama_context_default_params() {
/*.n_seq_max =*/ 1,
/*.n_rs_seq =*/ 0,
/*.n_outputs_max =*/ 0,
/*.n_outputs_max_per_seq =*/ 1,
/*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default
/*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS,
/*.ctx_type =*/ LLAMA_CONTEXT_TYPE_DEFAULT,
src/llama-cparams.h
+1
View File
@@ -15,6 +15,7 @@ struct llama_cparams {
uint32_t n_seq_max;
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback
uint32_t n_outputs_max; // max outputs supported by the context
uint32_t n_outputs_max_per_seq;
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
src/llama-graph.cpp
+95 -69
View File
@@ -4,6 +4,7 @@
#include "llama-model.h"
#include "llama-batch.h"
#include "llama-cparams.h"
#include "llama-sampler.h"
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
@@ -1353,24 +1354,24 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
}
}
}
for (auto & [seq_id, t] : t_sampled) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_sampled) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
for (auto & [seq_id, t] : t_sampled_probs) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_sampled_probs) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
for (auto & [seq_id, t] : t_sampled_logits) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_sampled_logits) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
for (auto & [seq_id, t] : t_candidates) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_candidates) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
}
@@ -3649,77 +3650,102 @@ void llm_graph_context::build_sampling() const {
auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);
res->add_input(std::move(inp_sampling));
std::map<llama_seq_id, int32_t> seq_to_logit_row;
int32_t logit_row_idx = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; i++) {
std::map<llama_seq_id, std::vector<uint32_t>> sampling_rows;
uint32_t n_rows = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
if (ubatch.output[i]) {
llama_seq_id seq_id = ubatch.seq_id[i][0];
seq_to_logit_row[seq_id] = logit_row_idx;
logit_row_idx++;
sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++);
}
}
res->t_sampled.resize(n_rows, nullptr);
res->t_sampled_probs.resize(n_rows, nullptr);
res->t_sampled_logits.resize(n_rows, nullptr);
res->t_candidates.resize(n_rows, nullptr);
// res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1)
GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor");
// add a dummy row of logits
// this trick makes the graph static, regardless of which samplers are activated
// this is important in order to minimize graph reallocations
// add a dummy row to keep the single-output graph static regardless of active samplers
// multi-output graphs can still vary with the number of output rows
ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);
for (const auto & [seq_id, sampler] : samplers) {
const auto it = seq_to_logit_row.find(seq_id);
// inactive samplers always work on the first row
const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0;
const int i_out = it != seq_to_logit_row.end() ? 1 : 0;
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);
ggml_format_name(logits_seq, "logits_seq_%d", seq_id);
struct llama_sampler_data data = {
/*.logits =*/ logits_seq,
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ nullptr,
};
assert(sampler->iface->backend_apply);
sampler->iface->backend_apply(sampler, ctx0, gf, &data);
if (data.sampled != nullptr) {
res->t_sampled[seq_id] = data.sampled;
outs[1] = data.sampled;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.probs != nullptr) {
res->t_sampled_probs[seq_id] = data.probs;
outs[1] = data.probs;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.logits != nullptr) {
res->t_sampled_logits[seq_id] = data.logits;
outs[1] = data.logits;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.candidates != nullptr) {
res->t_candidates[seq_id] = data.candidates;
outs[1] = data.candidates;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
for (const auto & entry : samplers) {
if (entry.second->iface->backend_reset) {
entry.second->iface->backend_reset(entry.second);
}
}
// TODO: Call llama_sampler_accept_ggml after all samplers have been applied.
static const std::vector<uint32_t> dummy_row = { 0 };
for (const auto & [seq_id, sampler] : samplers) {
const auto it = sampling_rows.find(seq_id);
// inactive samplers always work on the first row
const bool active = it != sampling_rows.end();
const auto & rows = active ? it->second : dummy_row;
const int i_out = active ? 1 : 0;
for (uint32_t i = 0; i < rows.size(); ++i) {
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]);
ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i);
struct llama_sampler_data data = {
/*.logits =*/ logits_seq,
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ nullptr,
};
assert(sampler->iface->backend_apply);
sampler->iface->backend_apply(sampler, ctx0, gf, &data);
if (data.sampled != nullptr) {
if (active) {
res->t_sampled[rows[i]] = data.sampled;
}
outs[1] = data.sampled;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.probs != nullptr) {
if (active) {
res->t_sampled_probs[rows[i]] = data.probs;
}
outs[1] = data.probs;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.logits != nullptr) {
if (active) {
res->t_sampled_logits[rows[i]] = data.logits;
}
outs[1] = data.logits;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.candidates != nullptr) {
if (active) {
res->t_candidates[rows[i]] = data.candidates;
}
outs[1] = data.candidates;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
}
}
// TODO: Call backend_accept after all samplers have been applied.
/*
for (const auto & [seq_id, sampler] : samplers) {
if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) {
ggml_tensor * selected_token = it->second;
if (selected_token != nullptr) {
llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token);
const auto it = sampling_rows.find(seq_id);
if (it == sampling_rows.end()) {
continue;
}
for (uint32_t row : it->second) {
ggml_tensor * selected_token = res->t_sampled[row];
if (selected_token != nullptr && sampler->iface->backend_accept) {
sampler->iface->backend_accept(sampler, ctx0, gf, selected_token);
}
}
}
src/llama-graph.h
+4 -4
View File
@@ -904,10 +904,10 @@ public:
std::vector<ggml_tensor *> t_layer_inp;
std::map<llama_seq_id, ggml_tensor *> t_sampled_logits;
std::map<llama_seq_id, ggml_tensor *> t_candidates;
std::map<llama_seq_id, ggml_tensor *> t_sampled;
std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
std::vector<ggml_tensor *> t_sampled;
std::vector<ggml_tensor *> t_sampled_probs;
std::vector<ggml_tensor *> t_sampled_logits;
std::vector<ggml_tensor *> t_candidates;
std::vector<llm_graph_input_ptr> inputs;
std::vector<llm_graph_fused_node> fused_nodes;
src/llama-model-saver.cpp
+1
View File
@@ -27,6 +27,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) {
case LLM_ARCH_APERTUS:
case LLM_ARCH_MIMO2:
case LLM_ARCH_STEP35:
case LLM_ARCH_MUSE_GLIMMER:
case LLM_ARCH_MELLUM:
case LLM_ARCH_LAGUNA:
return false;
src/llama-model.cpp
+3
View File
@@ -176,6 +176,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_olmo2(params);
case LLM_ARCH_OLMOE:
return new llama_model_olmoe(params);
case LLM_ARCH_MUSE_GLIMMER:
return new llama_model_muse_glimmer(params);
case LLM_ARCH_OPENELM:
return new llama_model_openelm(params);
case LLM_ARCH_GPTNEOX:
@@ -2599,6 +2601,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_DEEPSEEK2OCR:
case LLM_ARCH_DEEPSEEK32:
case LLM_ARCH_DEEPSEEK4:
case LLM_ARCH_MUSE_GLIMMER:
case LLM_ARCH_PLM:
case LLM_ARCH_CHATGLM:
case LLM_ARCH_GRANITE:
src/llama-sampler.cpp
+376 -93
View File
@@ -467,9 +467,11 @@ static void llama_sampler_empty_free(struct llama_sampler * smpl) {
static bool llama_sampler_empty_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
GGML_UNUSED(smpl);
GGML_UNUSED(buft);
GGML_UNUSED(n_outputs_max_per_seq);
return true;
}
@@ -511,6 +513,8 @@ static struct llama_sampler_i llama_sampler_empty_i = {
/* .backend_accept = */ llama_sampler_empty_backend_accept,
/* .backend_apply = */ llama_sampler_empty_backend_apply,
/* .backend_set_input = */ llama_sampler_empty_backend_set_input,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_empty(const char * name) {
@@ -551,6 +555,12 @@ struct llama_sampler_backend {
this->support = support;
}
// copy the state that is not tied to the current sampling graph
// samplers that hold only immutable configuration can use this as is
void copy_state(const llama_sampler_backend & src) {
GGML_UNUSED(src);
}
private:
std::string name;
std::string name_ext;
@@ -559,6 +569,71 @@ private:
bool support;
};
// .copy_state for samplers deriving from llama_sampler_backend
template<typename T>
static void llama_sampler_backend_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {
((T *) dst->ctx)->copy_state(*(const T *) src->ctx);
}
struct llama_sampler_backend_probe {
ggml_context_ptr ctx;
ggml_cgraph * gf;
};
static llama_sampler_backend_probe llama_sampler_backend_probe_graph(
llama_sampler * sampler,
int64_t n_candidates,
uint32_t max_nodes,
bool with_candidates) {
ggml_init_params params = {
/*.mem_size =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context_ptr ctx_ptr { ggml_init(params) };
if (!ctx_ptr) {
throw std::runtime_error(format("failed to create ggml context"));
}
auto * ctx = ctx_ptr.get();
auto * gf = ggml_new_graph_custom(ctx, max_nodes, false);
llama_sampler_data data = {
/*.logits =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates),
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ with_candidates ? ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates) : nullptr,
};
if (sampler->iface->backend_reset) {
sampler->iface->backend_reset(sampler);
}
sampler->iface->backend_apply(sampler, ctx, gf, &data);
for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {
if (output) {
ggml_build_forward_expand(gf, output);
}
}
if (sampler->iface->backend_reset) {
sampler->iface->backend_reset(sampler);
}
return { std::move(ctx_ptr), gf };
}
static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) {
uint32_t n_tensors = 0;
for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor;
tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) {
++n_tensors;
}
return std::max<uint32_t>(ggml_graph_n_nodes(probe.gf), n_tensors);
}
// check if all ggml ops used by the sampler are supported by the backend
static bool llama_sampler_backend_support(
llama_sampler * smpl,
@@ -569,50 +644,10 @@ static bool llama_sampler_backend_support(
return true;
}
ggml_init_params params = {
/*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true);
ggml_context_ptr ctx_ptr { ggml_init(params) };
if (!ctx_ptr) {
throw std::runtime_error(format("failed to create ggml context"));
}
ggml_context * ctx = ctx_ptr.get();
const int64_t n = 1024*1024;
llama_sampler_data data = {
/*.logits = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
/*.probs = */ nullptr,
/*.sampled = */ nullptr,
/*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
};
ggml_cgraph * gf = ggml_new_graph(ctx);
smpl->iface->backend_apply(smpl, ctx, gf, &data);
if (data.logits) {
ggml_build_forward_expand(gf, data.logits);
}
if (data.probs) {
ggml_build_forward_expand(gf, data.probs);
}
if (data.sampled) {
ggml_build_forward_expand(gf, data.sampled);
}
if (data.candidates) {
ggml_build_forward_expand(gf, data.candidates);
}
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
struct ggml_tensor * op = ggml_graph_node(gf, i);
for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) {
struct ggml_tensor * op = ggml_graph_node(probe.gf, i);
if (!ggml_backend_dev_supports_op(device, op)) {
LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
@@ -697,7 +732,8 @@ static void llama_sampler_chain_free(struct llama_sampler * smpl) {
static bool llama_sampler_chain_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * chain = (llama_sampler_chain *) smpl->ctx;
GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
@@ -705,26 +741,32 @@ static bool llama_sampler_chain_backend_init(
chain->is_init = true;
bool res = true;
bool backend_prefix = true;
for (auto & smpl : chain->samplers) {
bool res_cur = true;
bool cur_prefix = backend_prefix;
// to be able to run a sampler on the backend, it has to:
// - have the .backend_init() API implemented
// - return true during .backend_init()
if (smpl.ptr->iface->backend_init) {
if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
res_cur = false;
// - support the requested per-sequence output limit
if (cur_prefix && smpl.ptr->iface->backend_init) {
if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) {
cur_prefix = false;
}
} else {
res_cur = false;
cur_prefix = false;
}
smpl.is_backend = res_cur;
smpl.is_backend = cur_prefix;
backend_prefix = cur_prefix;
res = res && res_cur;
res = res && cur_prefix;
}
auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false);
chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe);
return res;
}
@@ -780,6 +822,36 @@ static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
}
}
static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) {
auto * chain = (llama_sampler_chain *) smpl->ctx;
for (auto & entry : chain->samplers) {
if (!entry.is_backend) {
break;
}
if (entry.ptr->iface->backend_reset) {
entry.ptr->iface->backend_reset(entry.ptr);
}
}
}
static void llama_sampler_chain_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {
const auto * src_chain = (const llama_sampler_chain *) src->ctx;
auto * dst_chain = (llama_sampler_chain *) dst->ctx;
GGML_ASSERT(src_chain->samplers.size() == dst_chain->samplers.size());
for (size_t i = 0; i < src_chain->samplers.size(); ++i) {
llama_sampler_copy(src_chain->samplers[i].ptr, dst_chain->samplers[i].ptr);
}
// note: is_init, n_nodes and is_backend belong to the current sampling graph
dst_chain->params = src_chain->params;
dst_chain->cur = src_chain->cur;
dst_chain->t_sample_us = src_chain->t_sample_us;
dst_chain->n_sample = src_chain->n_sample;
}
static struct llama_sampler_i llama_sampler_chain_i = {
/* .name = */ llama_sampler_chain_name,
/* .accept = */ llama_sampler_chain_accept,
@@ -791,22 +863,35 @@ static struct llama_sampler_i llama_sampler_chain_i = {
/* .backend_accept = */ llama_sampler_chain_backend_accept,
/* .backend_apply = */ llama_sampler_chain_backend_apply,
/* .backend_set_input = */ llama_sampler_chain_backend_set_input,
/* .backend_reset = */ llama_sampler_chain_backend_reset,
/* .copy_state = */ llama_sampler_chain_copy_state,
};
struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {
return llama_sampler_init(
/* .iface = */ &llama_sampler_chain_i,
/* .ctx = */ new llama_sampler_chain {
/* .params = */ params,
/* .is_init = */ false,
/* .samplers = */ {},
/* .cur = */ {},
/* .t_sample_us = */ 0,
/* .n_sample = */ 0,
/* .params = */ params,
/* .is_init = */ false,
/* .n_nodes = */ 0,
/* .samplers = */ {},
/* .cur = */ {},
/* .t_sample_us = */ 0,
/* .n_sample = */ 0,
}
);
}
uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) {
GGML_ASSERT(sampler != nullptr);
GGML_ASSERT(sampler->iface == &llama_sampler_chain_i);
const auto * chain = (const llama_sampler_chain *) sampler->ctx;
GGML_ASSERT(chain->is_init);
return chain->n_nodes;
}
llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx);
const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx);
@@ -816,6 +901,7 @@ llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_conte
// If a backend sampler has already sampled a token, return it.
if (sampled_token != LLAMA_TOKEN_NULL) {
LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);
llama_sampler_accept(smpl, sampled_token);
return sampled_token;
}
@@ -975,8 +1061,10 @@ static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_to
static bool llama_sampler_greedy_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_greedy *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -1012,6 +1100,8 @@ static struct llama_sampler_i llama_sampler_greedy_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_greedy_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_greedy>,
};
struct llama_sampler * llama_sampler_init_greedy() {
@@ -1031,7 +1121,25 @@ struct llama_sampler_dist : public llama_sampler_backend {
std::mt19937 rng;
ggml_tensor * inp_uniform;
// TODO: refactor + fix naming
// https://github.com/ggml-org/llama.cpp/pull/25532/changes#r3749906719
// use a temporary RNG for multi-output sampling so rejected tokens do not advance rng
bool backend_transactional;
std::mt19937 rng_backend;
size_t n_backend_draws_generated;
size_t n_backend_draws_committed;
// inputs for the current sampling graph
std::vector<ggml_tensor *> inp_uniforms;
void copy_state(const llama_sampler_dist & src) {
// note: inp_uniforms and backend_transactional belong to the current sampling graph
seed_cur = src.seed_cur;
rng = src.rng;
rng_backend = src.rng_backend;
n_backend_draws_generated = src.n_backend_draws_generated;
n_backend_draws_committed = src.n_backend_draws_committed;
}
};
static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
@@ -1050,7 +1158,11 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
cur_p->selected = 0;
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
if (cur_p->size == 1) {
// keep the RNG state aligned with backend sampling, which draws once per output
dist(ctx->rng);
cur_p->data[0].p = 1.0f;
return;
}
@@ -1075,7 +1187,6 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
// sample from the obtained probabilities and normalize the probs in a single pass
// this is ~3x faster on Mac with full gpt-oss vocab than the version below
//
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
const double rnd = dist(ctx->rng);
double sum_run = 0.0f;
@@ -1115,6 +1226,9 @@ static void llama_sampler_dist_reset(struct llama_sampler * smpl) {
auto * ctx = (llama_sampler_dist *) smpl->ctx;
ctx->seed_cur = get_rng_seed(ctx->seed);
ctx->rng.seed(ctx->seed_cur);
ctx->rng_backend = ctx->rng;
ctx->n_backend_draws_generated = 0;
ctx->n_backend_draws_committed = 0;
}
static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) {
@@ -1125,7 +1239,12 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample
{
auto * result_ctx = (llama_sampler_dist *) result->ctx;
result_ctx->rng = ctx->rng;
result_ctx->seed_cur = ctx->seed_cur;
result_ctx->rng = ctx->rng;
result_ctx->backend_transactional = ctx->backend_transactional;
result_ctx->rng_backend = ctx->rng_backend;
result_ctx->n_backend_draws_generated = ctx->n_backend_draws_generated;
result_ctx->n_backend_draws_committed = ctx->n_backend_draws_committed;
}
return result;
@@ -1137,12 +1256,17 @@ static void llama_sampler_dist_free(struct llama_sampler * smpl) {
static bool llama_sampler_dist_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
const bool res = llama_sampler_backend_support(smpl, buft);
sctx->init(res);
sctx->backend_transactional = n_outputs_max_per_seq > 1;
sctx->rng_backend = sctx->rng;
sctx->n_backend_draws_generated = 0;
sctx->n_backend_draws_committed = 0;
return res;
}
@@ -1156,9 +1280,10 @@ static void llama_sampler_dist_backend_apply(
auto * sctx = (llama_sampler_dist *) smpl->ctx;
sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
ggml_set_name (sctx->inp_uniform, "uniform");
ggml_set_input(sctx->inp_uniform);
ggml_tensor * inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
ggml_format_name(inp_uniform, "uniform_%zu", sctx->inp_uniforms.size());
ggml_set_input(inp_uniform);
sctx->inp_uniforms.push_back(inp_uniform);
// flatten
struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
@@ -1174,7 +1299,7 @@ static void llama_sampler_dist_backend_apply(
// Recall that each entry in cumsum is the cumulative probability up to that
// index so values stay negative while the cumulative total is below the
// random value, and become zero/positive once the threshold is crossed.
struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
struct ggml_tensor * diff = ggml_sub(ctx, cumsum, inp_uniform);
ggml_set_name(diff, "dist_cumsum");
// The ggml_step function produces a tensor where entries are 1 if the
@@ -1189,6 +1314,9 @@ static void llama_sampler_dist_backend_apply(
struct ggml_tensor * idxf = ggml_sum(ctx, mask);
ggml_set_name(idxf, "dist_index_f32");
// Clamp to prevent out-of-bounds access when computing the index.
idxf = ggml_clamp(ctx, idxf, 1.0f, mask->ne[0]);
// Use ggml_scale_bias to scale the index value by -1 and then add the size
// of the mask to that value so we get the correct index ((-1 * idxf) + n).
struct ggml_tensor * idx = ggml_cast(ctx, ggml_scale_bias(ctx, idxf, -1.0f, mask->ne[0]), GGML_TYPE_I32);
@@ -1210,22 +1338,52 @@ static void llama_sampler_dist_backend_apply(
static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
GGML_ASSERT(sctx->inp_uniform != nullptr);
GGML_ASSERT(!sctx->inp_uniforms.empty());
// We sample in double precision and cast to float to match rnd numbers of
// llama_dampler_dist which uses double precision (sampling from
// llama_sampler_dist which uses double precision (sampling from
// std::uniform_real_distribution<double> and
// std::uniform_real_distribution<float> with same rng will produce
// different sequences).
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
const float rnd = dist(sctx->rng);
ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
auto & rng = sctx->backend_transactional ? sctx->rng_backend : sctx->rng;
for (auto * inp_uniform : sctx->inp_uniforms) {
GGML_ASSERT(inp_uniform != nullptr);
const float rnd = dist(rng);
ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float));
if (sctx->backend_transactional) {
++sctx->n_backend_draws_generated;
}
}
}
static void llama_sampler_dist_backend_reset(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
sctx->inp_uniforms.clear();
}
static void llama_sampler_dist_accept(struct llama_sampler * smpl, llama_token token) {
GGML_UNUSED(token);
auto * sctx = (llama_sampler_dist *) smpl->ctx;
if (!sctx->backend_transactional ||
sctx->n_backend_draws_committed >= sctx->n_backend_draws_generated) {
return;
}
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
dist(sctx->rng);
++sctx->n_backend_draws_committed;
}
static struct llama_sampler_i llama_sampler_dist_i = {
/* .name = */ llama_sampler_dist_name,
/* .accept = */ nullptr,
/* .accept = */ llama_sampler_dist_accept,
/* .apply = */ llama_sampler_dist_apply,
/* .reset = */ llama_sampler_dist_reset,
/* .clone = */ llama_sampler_dist_clone,
@@ -1234,6 +1392,8 @@ static struct llama_sampler_i llama_sampler_dist_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_dist_backend_apply,
/* .backend_set_input = */ llama_sampler_dist_backend_set_input,
/* .backend_reset = */ llama_sampler_dist_backend_reset,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_dist>,
};
struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
@@ -1242,14 +1402,39 @@ struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
/* .iface = */ &llama_sampler_dist_i,
/* .ctx = */ new llama_sampler_dist {
("dist"),
/* .seed = */ seed,
/* .seed_cur = */ seed_cur,
/* .rng = */ std::mt19937(seed_cur),
/* .inp_uniform = */ nullptr,
/* .seed = */ seed,
/* .seed_cur = */ seed_cur,
/* .rng = */ std::mt19937(seed_cur),
/* .backend_transactional = */ false,
/* .rng_backend = */ std::mt19937(seed_cur),
/* .n_backend_draws_generated = */ 0,
/* .n_backend_draws_committed = */ 0,
/* .inp_uniforms = */ {},
}
);
}
void llama_sampler_backend_begin(llama_sampler * sampler) {
GGML_ASSERT(sampler != nullptr);
if (sampler->iface == &llama_sampler_chain_i) {
auto * chain = (llama_sampler_chain *) sampler->ctx;
for (auto & entry : chain->samplers) {
if (!entry.is_backend) {
break;
}
llama_sampler_backend_begin(entry.ptr);
}
} else if (sampler->iface == &llama_sampler_dist_i) {
auto * ctx = (llama_sampler_dist *) sampler->ctx;
if (ctx->backend_transactional) {
ctx->rng_backend = ctx->rng;
ctx->n_backend_draws_generated = 0;
ctx->n_backend_draws_committed = 0;
}
}
}
// top-k
struct llama_sampler_top_k : public llama_sampler_backend {
@@ -1277,8 +1462,10 @@ static void llama_sampler_top_k_free(struct llama_sampler * smpl) {
static bool llama_sampler_top_k_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_top_k *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -1325,6 +1512,8 @@ static struct llama_sampler_i llama_sampler_top_k_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_top_k_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_top_k>,
};
struct llama_sampler * llama_sampler_init_top_k(int32_t k) {
@@ -1423,8 +1612,10 @@ static void llama_sampler_top_p_free(struct llama_sampler * smpl) {
static bool llama_sampler_top_p_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_top_p *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -1521,6 +1712,8 @@ static struct llama_sampler_i llama_sampler_top_p_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_top_p_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_top_p>,
};
struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) {
@@ -1618,8 +1811,10 @@ static void llama_sampler_min_p_free(struct llama_sampler * smpl) {
static bool llama_sampler_min_p_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_min_p *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -1680,6 +1875,8 @@ static struct llama_sampler_i llama_sampler_min_p_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_min_p_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_min_p>,
};
struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
@@ -1790,6 +1987,8 @@ static struct llama_sampler_i llama_sampler_typical_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) {
@@ -1866,8 +2065,10 @@ static void llama_sampler_backend_temp_sampling(
static bool llama_sampler_temp_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_temp *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -1896,6 +2097,8 @@ static struct llama_sampler_i llama_sampler_temp_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_temp_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_temp>,
};
struct llama_sampler * llama_sampler_init_temp(float temp) {
@@ -2009,8 +2212,10 @@ static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) {
static bool llama_sampler_temp_ext_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -2095,6 +2300,8 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_temp_ext_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_temp_ext>,
};
struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) {
@@ -2202,6 +2409,8 @@ static struct llama_sampler_i llama_sampler_xtc_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
@@ -2290,7 +2499,7 @@ static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sa
// copy the state
{
auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx;
auto * result_ctx = (llama_sampler_mirostat *) result->ctx;
result_ctx->mu = ctx->mu;
result_ctx->rng = ctx->rng;
@@ -2321,6 +2530,8 @@ static struct llama_sampler_i llama_sampler_mirostat_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) {
@@ -2425,6 +2636,8 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) {
@@ -2546,6 +2759,8 @@ static struct llama_sampler_i llama_sampler_grammar_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
static struct llama_sampler * llama_sampler_init_grammar_impl(
@@ -2661,6 +2876,12 @@ struct llama_sampler_penalties : public llama_sampler_backend {
std::vector<int32_t> host_token_ids;
std::vector<int32_t> host_counts;
void copy_state(const llama_sampler_penalties & src) {
// note: inp_token_ids/inp_counts belong to the current sampling graph
prev = src.prev;
token_count = src.token_count;
}
static bool is_disabled(
int32_t penalty_last_n,
float penalty_repeat,
@@ -2790,9 +3011,15 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) {
static bool llama_sampler_penalties_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_penalties *) smpl->ctx;
if (n_outputs_max_per_seq > 1) {
sctx->init(false);
return false;
}
const bool res = llama_sampler_backend_support(smpl, buft);
sctx->init(res);
@@ -2952,6 +3179,12 @@ static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smp
ggml_backend_tensor_set(sctx->inp_counts, sctx->host_counts.data(), 0, sctx->n_max * sizeof(int32_t));
}
static void llama_sampler_penalties_backend_reset(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_penalties *) smpl->ctx;
sctx->inp_token_ids = nullptr;
sctx->inp_counts = nullptr;
}
static struct llama_sampler_i llama_sampler_penalties_i = {
/* .name = */ llama_sampler_penalties_name,
/* .accept = */ llama_sampler_penalties_accept,
@@ -2963,6 +3196,8 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_penalties_backend_apply,
/* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
/* .backend_reset = */ llama_sampler_penalties_backend_reset,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_penalties>,
};
struct llama_sampler * llama_sampler_init_penalties(
@@ -3058,6 +3293,8 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
@@ -3395,6 +3632,8 @@ static struct llama_sampler_i llama_sampler_dry_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
@@ -3614,6 +3853,8 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_adaptive_p(
@@ -3715,13 +3956,17 @@ static void llama_sampler_logit_bias_backend_apply(
const size_t n = sctx->logit_bias.size();
sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
ggml_set_name(sctx->inp_logit_bias, "logit_bias");
ggml_set_input(sctx->inp_logit_bias);
if (sctx->inp_logit_bias == nullptr) {
GGML_ASSERT(sctx->inp_logit_idxs == nullptr);
sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
ggml_set_input(sctx->inp_logit_idxs);
sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
ggml_set_name(sctx->inp_logit_bias, "logit_bias");
ggml_set_input(sctx->inp_logit_bias);
sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
ggml_set_input(sctx->inp_logit_idxs);
}
ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);
@@ -3756,10 +4001,18 @@ static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * sm
ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
}
static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
sctx->inp_logit_bias = nullptr;
sctx->inp_logit_idxs = nullptr;
}
static bool llama_sampler_logit_bias_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
GGML_UNUSED(buft);
GGML_UNUSED(n_outputs_max_per_seq);
auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
@@ -3783,6 +4036,8 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_logit_bias_backend_apply,
/* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
/* .backend_reset = */ llama_sampler_logit_bias_backend_reset,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_logit_bias>,
};
struct llama_sampler * llama_sampler_init_logit_bias(
@@ -4022,10 +4277,12 @@ static struct llama_sampler_i llama_sampler_infill_i = {
/* .reset = */ nullptr,
/* .clone = */ llama_sampler_infill_clone,
/* .free = */ llama_sampler_infill_free,
/* .backend_apply = */ nullptr,
/* .backend_accept = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_init = */ nullptr,
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
@@ -4039,6 +4296,32 @@ struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * voca
);
}
void llama_sampler_copy(const struct llama_sampler * src, struct llama_sampler * dst) {
if (!src || !dst || src == dst) {
return;
}
GGML_ASSERT(src->iface == dst->iface && "llama_sampler_copy: cannot copy between different sampler types");
if (dst->iface->copy_state) {
dst->iface->copy_state(src, dst);
return;
}
// build a temporary sampler carrying src's current state
llama_sampler * tmp = llama_sampler_clone(src);
// free dst's old state (frees dst->ctx, including children for a chain)
if (dst->iface->free) {
dst->iface->free(dst);
}
// transplant tmp's state into dst, then destroy the (now empty) temp shell
dst->ctx = tmp->ctx;
tmp->ctx = nullptr;
delete tmp;
}
// utils
uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {
src/llama-sampler.h
+5
View File
@@ -15,6 +15,8 @@ struct llama_sampler_chain {
// has .backend_init() been called?
bool is_init = false;
uint32_t n_nodes = 0;
struct info {
bool is_backend;
@@ -33,6 +35,9 @@ struct llama_sampler_chain {
mutable int32_t n_sample;
};
uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler);
void llama_sampler_backend_begin(llama_sampler * sampler);
struct llama_sampler * llama_sampler_init_dry_testing(
float dry_multiplier,
float dry_base,
src/models/models.h
+13
View File
@@ -1044,6 +1044,19 @@ struct llama_model_olmoe : public llama_model_base {
};
struct llama_model_muse_glimmer : public llama_model_base {
llama_model_muse_glimmer(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_openelm : public llama_model_base {
llama_model_openelm(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
src/models/muse-glimmer.cpp
+208
View File
@@ -0,0 +1,208 @@
#include "models.h"
void llama_model_muse_glimmer::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
uint32_t swa_period = 4;
if (ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false)) {
hparams.set_swa_pattern(swa_period);
} else {
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
}
switch (hparams.n_layer()) {
case 52: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_muse_glimmer::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
// Pre/post-attention norms (Muse Glimmer's `weight + 1` applied at conversion time).
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);
// Q/K/V/O projections.
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
// QK-norm. Weights are synthesized at conversion time to absorb `qk_scale_factor`.
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
// Attention output gate: sigmoid(gate) * attn_out before o_proj (same as afmoe).
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
// Pre/post-FFN norms (FFN_PRE_NORM is aliased to LLM_TENSOR_FFN_NORM).
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
// Dense FFN (unlike afmoe, no MoE branches).
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
llama_model_muse_glimmer::graph::graph(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
// Different to f_norm_rms_eps for post-attn / post-FFN norms
const float post_norm_eps = 1e-8f;
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
inpL = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1);
cb(inpL, "embd_norm", -1);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv_iswa();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
for (int il = 0; il < n_layer; ++il) {
// expose per-layer residual for speculative drafts (see LLM_KV_TARGET_LAYERS).
res->t_layer_inp[il] = inpL;
const float freq_base_l = model.get_rope_freq_base (cparams, il);
const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
ggml_tensor * inpSA = inpL;
// RoPE runs on the SWA layers, NoPE on full ones.
const bool use_rope = hparams.is_swa(il);
// pre-attention norm (weight+1 folded at conversion time)
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
// self-attention: attention output gate around SDPA (afmoe.cpp:147-191)
{
ggml_tensor * attn_inp = cur; // save input for gate computation
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
// gate = wqkv_gate @ attn_inp (from pre-attn hidden state)
ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);
cb(gate, "attn_gate_proj", il);
// QK-norm. attn_q_norm weight was synthesized at conversion to broadcast
// qk_scale_factor across head_dim; attn_k_norm is identity (ones).
Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
cb(Qcur, "Qcur_normed", il);
cb(Kcur, "Kcur_normed", il);
if (use_rope) {
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Qcur, "Qcur_rope", il);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Kcur, "Kcur_rope", il);
}
// SDPA. wo is deferred; the gate goes between attn_out and o_proj.
cur = build_attn(inp_attn,
NULL, NULL, NULL,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "attn_out", il);
gate = ggml_sigmoid(ctx0, gate);
cb(gate, "attn_gate_sig", il);
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
cb(cur, "attn_o_proj", il);
}
cur = ggml_rms_norm(ctx0, cur, post_norm_eps);
cur = ggml_mul(ctx0, cur, model.layers[il].attn_post_norm);
cb(cur, "attn_post_norm", il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
// pre-FFN norm
cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
// SwiGLU dense FFN
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
model.layers[il].ffn_gate, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
cur = ggml_rms_norm(ctx0, cur, post_norm_eps);
cur = ggml_mul(ctx0, cur, model.layers[il].ffn_post_norm);
cb(cur, "ffn_post_norm", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
inpL = cur;
}
cur = inpL;
// final norm
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
// lm_head, followed by output multiplier
cur = build_lora_mm(model.output, cur, model.output_s);
cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
// Final logit tanh softcap (from gemma3.cpp).
if (hparams.f_final_logit_softcapping) {
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
cur = ggml_tanh(ctx0, cur);
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
}
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
std::unique_ptr<llm_graph_context> llama_model_muse_glimmer::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
tests/test-arg-parser.cpp
+30
View File
@@ -2,7 +2,9 @@
#include "common.h"
#include "download.h"
#include "llama.h"
#include "speculative.h"
#include <limits>
#include <string>
#include <vector>
#include <sstream>
@@ -14,6 +16,34 @@
static void test(void) {
common_params params;
auto assert_output_limits = [](int32_t n_batch, int32_t n_parallel, int32_t n_draft,
int32_t total, int32_t per_seq) {
const auto limits = common_speculative_get_output_limits(n_batch, n_parallel, n_draft);
assert(limits.total == total);
assert(limits.per_seq == per_seq);
};
assert_output_limits(16, 2, 3, 8, 4);
assert_output_limits(16, 2, -1, 2, 1);
assert_output_limits( 6, 2, 3, 6, 4);
assert_output_limits( 2, 1, 3, 2, 2);
assert_output_limits(
std::numeric_limits<int32_t>::max(),
std::numeric_limits<int32_t>::max(),
std::numeric_limits<int32_t>::max(),
std::numeric_limits<int32_t>::max(),
std::numeric_limits<int32_t>::max());
{
common_params base;
base.n_parallel = 4;
base.n_outputs_max_per_seq = 8;
const auto draft = common_base_params_to_speculative(base);
assert(draft.n_outputs_max == 4);
assert(draft.n_outputs_max_per_seq == 1);
}
printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n");
for (int ex = 0; ex < LLAMA_EXAMPLE_COUNT; ex++) {
try {
tests/test-backend-ops.cpp
+11 -3
View File
@@ -6712,19 +6712,26 @@ struct test_roll : public test_case {
const int shift1;
const int shift3;
const int shift4;
const bool permute;
std::string vars() override {
return VARS_TO_STR4(shift0, shift1, shift3, shift4);
return VARS_TO_STR5(shift0, shift1, shift3, shift4, permute);
}
test_roll(int shift0 = 3, int shift1 = -2, int shift3 = 1, int shift4 = -1)
: shift0(shift0), shift1(shift1), shift3(shift3), shift4(shift4) {}
test_roll(int shift0 = 3, int shift1 = -2, int shift3 = 1, int shift4 = -1, bool permute = false)
: shift0(shift0), shift1(shift1), shift3(shift3), shift4(shift4), permute(permute) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
int64_t ne[4] = {10, 5, 4, 3};
ggml_tensor * a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
ggml_set_name(a, "a");
if (permute) {
// ggml_roll only requires nb[0] == type size, so a permuted src is valid
a = ggml_permute(ctx, a, 0, 2, 1, 3);
ggml_set_name(a, "a_permuted");
}
ggml_tensor * out = ggml_roll(ctx, a, shift0, shift1, shift3, shift4);
ggml_set_name(out, "out");
@@ -9459,6 +9466,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_pad_reflect_1d());
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 4, 1}));
test_cases.emplace_back(new test_roll());
test_cases.emplace_back(new test_roll(3, -2, 1, -1, true));
test_cases.emplace_back(new test_arange());
test_cases.emplace_back(new test_arange(GGML_TYPE_F32, 0.0f, 1048576.0f, 1.0f));
test_cases.emplace_back(new test_timestep_embedding());
tests/test-backend-sampler.cpp
+464 -33
View File
@@ -14,6 +14,7 @@
#include <fstream>
#include <functional>
#include <map>
#include <random>
#include <string>
#include <unordered_map>
#include <unordered_set>
@@ -80,7 +81,13 @@ struct test_context {
std::unordered_map<llama_seq_id, int32_t> seq_positions;
std::unordered_map<llama_seq_id, int32_t> last_batch_info;
test_context(const test_params & params, std::vector<llama_sampler_seq_config> & configs, int32_t n_seq_max = -1) {
test_context(
const test_params & params,
std::vector<llama_sampler_seq_config> & configs,
int32_t n_seq_max = -1,
uint32_t n_outputs_max = 0,
uint32_t n_ubatch = 0,
uint32_t n_outputs_max_per_seq = 1) {
auto * model = params.model.get();
GGML_ASSERT(model);
@@ -89,6 +96,11 @@ struct test_context {
llama_context_params cparams = llama_context_default_params();
cparams.n_ctx = 512;
cparams.n_batch = 512;
if (n_ubatch > 0) {
cparams.n_ubatch = n_ubatch;
}
cparams.n_outputs_max = n_outputs_max;
cparams.n_outputs_max_per_seq = n_outputs_max_per_seq;
cparams.samplers = configs.data();
cparams.n_samplers = configs.size();
cparams.kv_unified = true;
@@ -262,6 +274,66 @@ struct test_context {
}
};
struct test_single_output_backend_sampler {
bool backend_initialized = false;
uint32_t backend_outputs_max_per_seq = 0;
int backend_apply_count = 0;
int apply_count = 0;
};
static const char * test_single_output_backend_sampler_name(const llama_sampler * /*smpl*/) {
return "single-output-backend";
}
static void test_single_output_backend_sampler_apply(
llama_sampler * smpl, llama_token_data_array * /*cur_p*/) {
auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
ctx->apply_count++;
}
static void test_single_output_backend_sampler_free(llama_sampler * smpl) {
delete (test_single_output_backend_sampler *) smpl->ctx;
}
static bool test_single_output_backend_sampler_backend_init(
llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_max_per_seq) {
auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
ctx->backend_outputs_max_per_seq = n_outputs_max_per_seq;
if (n_outputs_max_per_seq > 1) {
return false;
}
ctx->backend_initialized = true;
return true;
}
static void test_single_output_backend_sampler_backend_apply(
llama_sampler * smpl, ggml_context * /*ctx*/, ggml_cgraph * /*gf*/, llama_sampler_data * /*data*/) {
auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
ctx->backend_apply_count++;
}
static llama_sampler_i test_single_output_backend_sampler_i = {
/* .name = */ test_single_output_backend_sampler_name,
/* .accept = */ nullptr,
/* .apply = */ test_single_output_backend_sampler_apply,
/* .reset = */ nullptr,
/* .clone = */ nullptr,
/* .free = */ test_single_output_backend_sampler_free,
/* .backend_init = */ test_single_output_backend_sampler_backend_init,
/* .backend_accept = */ nullptr,
/* .backend_apply = */ test_single_output_backend_sampler_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
static llama_sampler * test_single_output_backend_sampler_init(
test_single_output_backend_sampler ** sampler_ctx) {
auto * ctx = new test_single_output_backend_sampler;
*sampler_ctx = ctx;
return llama_sampler_init(&test_single_output_backend_sampler_i, ctx);
}
static void test_backend_greedy_sampling(const test_params & params) {
const int seq_id = 0;
@@ -661,7 +733,7 @@ static void test_backend_multi_sequence_sampling(const test_params & params) {
}
static void test_backend_dist_sampling(const test_params & params) {
const int seq_id = 189;
const int seq_id = 0;
const int32_t seed = 88;
struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -1527,43 +1599,398 @@ static void test_backend_cpu_mixed_batch(const test_params & params) {
printf("backend-cpu mixed batch test PASSED\n");
}
static void test_backend_max_outputs(const test_params & params) {
const int seq_id = 0;
const int32_t seed = 88;
static void test_backend_multi_output_limit(const test_params & params) {
const llama_seq_id seq_id = 0;
llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
llama_sampler_ptr backend_sampler_chain(llama_sampler_chain_init(backend_chain_params));
llama_sampler_chain_add(backend_sampler_chain.get(), llama_sampler_init_dist(seed));
std::vector<llama_sampler_seq_config> backend_sampler_configs = {{ seq_id, backend_sampler_chain.get() }};
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88));
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
test_context test_ctx(params, configs, 1, 3, 0, 2);
test_context test_ctx(params, backend_sampler_configs);
llama_batch batch = llama_batch_init(512, 0, 1);
std::string prompt = "Hello";
std::vector<llama_token> tokens;
tokens.push_back(llama_vocab_bos(test_ctx.vocab));
std::vector<llama_token> prompt_tokens(32);
int n_tokens = llama_tokenize(test_ctx.vocab, prompt.c_str(), prompt.length(),
prompt_tokens.data(), prompt_tokens.size(),
false, false);
for (int i = 0; i < n_tokens; i++) {
tokens.push_back(prompt_tokens[i]);
llama_batch batch = llama_batch_init(3, 0, 1);
for (int i = 0; i < 3; ++i) {
common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), i, { seq_id }, true);
}
for (size_t i = 0; i < tokens.size(); i++) {
// set all tokens as output to trigger error
common_batch_add(batch, tokens[i], i, { seq_id }, true);
}
printf(">>> test_max_outputs expected error start:\n");
printf(">>> test_backend_multi_output_limit expected error start:\n");
const int ret = llama_decode(test_ctx.ctx.get(), batch);
GGML_ASSERT(ret != 0 && "llama_decode should not succeed multiple outputs per sequence");
printf("<<< test_max_outputs expected error end.\n");
GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit");
printf("<<< test_backend_multi_output_limit expected error end.\n");
llama_batch_free(batch);
printf("backend max outputs test PASSED\n");
printf("backend multi-output limit test PASSED\n");
}
static void test_backend_multi_sequence_multi_output_dist(const test_params & params) {
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
const uint32_t seeds[] = { 88, 1337 };
// reduce the chance that swapped random inputs select the same token
const float temp = 10.0f;
llama_sampler_ptr chain_0(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_ptr chain_1(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_chain_add(chain_0.get(), llama_sampler_init_temp(temp));
llama_sampler_chain_add(chain_0.get(), llama_sampler_init_dist(seeds[0]));
llama_sampler_chain_add(chain_1.get(), llama_sampler_init_temp(temp));
llama_sampler_chain_add(chain_1.get(), llama_sampler_init_dist(seeds[1]));
std::vector<llama_sampler_seq_config> configs = {
{ 0, chain_0.get() },
{ 1, chain_1.get() },
};
test_context test_ctx(params, configs, 2, 4, 0, 2);
std::vector<llama_sampler_seq_config> reference_configs;
test_context reference_ctx(params, reference_configs, 2, 4);
const llama_token seq_tokens[2][2] = {
{ llama_vocab_bos(vocab), llama_vocab_eos(vocab) },
{ llama_vocab_eos(vocab), llama_vocab_bos(vocab) },
};
llama_batch batch = llama_batch_init(4, 0, 1);
for (int pos = 0; pos < 2; ++pos) {
common_batch_add(batch, seq_tokens[0][pos], pos, { 0 }, true);
common_batch_add(batch, seq_tokens[1][pos], pos, { 1 }, true);
}
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
std::mt19937 reference_rngs[] = {
std::mt19937(seeds[0]),
std::mt19937(seeds[1]),
};
std::uniform_real_distribution<double> reference_dist(0.0, 1.0);
for (int i = 0; i < batch.n_tokens; ++i) {
const llama_seq_id seq_id = batch.seq_id[i][0];
GGML_ASSERT(seq_id == 0 || seq_id == 1);
llama_sampler * chain = seq_id == 0 ? chain_0.get() : chain_1.get();
const llama_token backend_token = llama_sampler_sample(chain, test_ctx.ctx.get(), i);
const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i);
const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i);
const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i);
const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i);
const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
GGML_ASSERT(sampled_logits != nullptr);
GGML_ASSERT(sampled_probs != nullptr);
GGML_ASSERT(reference_logits != nullptr);
GGML_ASSERT(n_logits == (uint32_t) n_vocab);
GGML_ASSERT(n_probs == (uint32_t) n_vocab);
float prob_sum = 0.0f;
float cumsum_before = 0.0f;
for (llama_token token = 0; token < n_vocab; ++token) {
const float expected_logit = reference_logits[token] / temp;
const float tolerance = 1e-4f * std::max(1.0f, std::fabs(expected_logit));
GGML_ASSERT(std::fabs(sampled_logits[token] - expected_logit) <= tolerance);
GGML_ASSERT(std::isfinite(sampled_probs[token]));
GGML_ASSERT(sampled_probs[token] >= 0.0f);
prob_sum += sampled_probs[token];
if (token < backend_token) {
cumsum_before += sampled_probs[token];
}
}
GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
const float rnd = reference_dist(reference_rngs[seq_id]);
const float cumsum_sampled = cumsum_before + sampled_probs[backend_token];
GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
}
llama_batch_free(batch);
printf("backend multi-sequence multi-output dist test PASSED\n");
}
static void test_backend_multi_output_dist_transaction(const test_params & params) {
const llama_seq_id seq_id = 0;
const uint32_t seed = 95;
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_chain_add(chain.get(), llama_sampler_init_temp(10.0f));
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
test_context test_ctx(params, configs, 1, 3, 2, 3);
auto verify_random = [&](int32_t row, float rnd, bool accept = true) {
const llama_token token = accept ?
llama_sampler_sample(chain.get(), test_ctx.ctx.get(), row) :
llama_get_sampled_token_ith(test_ctx.ctx.get(), row);
const float * probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), row);
GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
GGML_ASSERT(probs != nullptr);
float cumsum_before = 0.0f;
for (llama_token i = 0; i < token; ++i) {
cumsum_before += probs[i];
}
const float cumsum_sampled = cumsum_before + probs[token];
GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
};
std::mt19937 rng(seed);
std::uniform_real_distribution<double> dist(0.0, 1.0);
float randoms[3];
for (float & rnd : randoms) {
rnd = dist(rng);
}
int32_t pos = 0;
auto decode = [&]() {
llama_batch batch = llama_batch_init(3, 0, 1);
for (int32_t i = 0; i < 3; ++i) {
common_batch_add(batch, llama_vocab_bos(vocab), pos++, { seq_id }, true);
}
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
return batch;
};
llama_batch batch = decode();
verify_random(0, randoms[0], false);
llama_batch_free(batch);
batch = decode();
verify_random(0, randoms[0]);
verify_random(1, randoms[1]);
llama_batch_free(batch);
batch = decode();
llama_sampler_ptr saved(llama_sampler_clone(chain.get()));
verify_random(0, randoms[2]);
llama_batch_free(batch);
llama_sampler_copy(saved.get(), chain.get());
batch = decode();
verify_random(0, randoms[2]);
llama_batch_free(batch);
printf("backend multi-output dist transaction test PASSED\n");
}
static void test_backend_multi_output_sampling_chain(const test_params & params) {
const llama_seq_id seq_id = 0;
const uint32_t seed = 88;
const float p = 0.9f;
const float temp = 0.8f;
const float cdf_epsilon = 1e-4f;
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
const uint32_t k = std::min<uint32_t>(512, n_vocab);
const llama_logit_bias bias = { llama_vocab_bos(vocab), -0.1f };
auto make_filter_chain = [&]() {
llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_chain_add(result.get(), llama_sampler_init_logit_bias(n_vocab, 1, &bias));
llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
llama_sampler_chain_add(result.get(), llama_sampler_init_top_p(p, 1));
llama_sampler_chain_add(result.get(), llama_sampler_init_min_p(0.01f, 1));
llama_sampler_chain_add(result.get(), llama_sampler_init_temp(temp));
return result;
};
llama_sampler_ptr chain = make_filter_chain();
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
test_context test_ctx(params, configs, 1, 2, 2, 2);
std::vector<llama_sampler_seq_config> reference_configs;
test_context reference_ctx(params, reference_configs, 1, 2, 2);
llama_sampler_ptr reference_bias(llama_sampler_init_logit_bias(n_vocab, 1, &bias));
llama_sampler_ptr reference_top_k(llama_sampler_init_top_k(k));
llama_sampler_ptr reference_top_p(llama_sampler_init_top_p(p, 1));
llama_sampler_ptr reference_min_p(llama_sampler_init_min_p(0.01f, 1));
llama_sampler_ptr reference_temp(llama_sampler_init_temp(temp));
std::vector<llama_token_data> reference_data(n_vocab);
auto make_batch = [&](int32_t pos) {
llama_batch batch = llama_batch_init(2, 0, 1);
for (int i = 0; i < 2; ++i) {
common_batch_add(batch, llama_vocab_bos(vocab), pos + i, { seq_id }, true);
}
return batch;
};
llama_batch batch = make_batch(0);
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
for (int i = 0; i < batch.n_tokens; ++i) {
const llama_token backend_token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i);
const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i);
const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i);
const llama_token * sampled_candidates = llama_get_sampled_candidates_ith(test_ctx.ctx.get(), i);
const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i);
const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i);
const uint32_t n_candidates = llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i);
const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
GGML_ASSERT(sampled_logits != nullptr);
GGML_ASSERT(sampled_probs != nullptr);
GGML_ASSERT(sampled_candidates != nullptr);
GGML_ASSERT(reference_logits != nullptr);
GGML_ASSERT(n_logits == k);
GGML_ASSERT(n_probs == n_logits);
GGML_ASSERT(n_candidates == n_logits);
for (llama_token token = 0; token < n_vocab; ++token) {
reference_data[token] = { token, reference_logits[token], 0.0f };
}
llama_token_data_array reference = {
/* .data = */ reference_data.data(),
/* .size = */ reference_data.size(),
/* .selected = */ LLAMA_TOKEN_NULL,
/* .sorted = */ false,
};
llama_sampler_apply(reference_bias.get(), &reference);
llama_sampler_apply(reference_top_k.get(), &reference);
llama_sampler_apply(reference_top_p.get(), &reference);
GGML_ASSERT(reference.size > 0);
float cdf = 0.0f;
for (size_t j = 0; j < reference.size; ++j) {
cdf += reference.data[j].p;
}
const float cdf_before = cdf - reference.data[reference.size - 1].p;
const float boundary_distance = std::min(std::fabs(cdf_before - p), std::fabs(cdf - p));
llama_sampler_apply(reference_min_p.get(), &reference);
llama_sampler_apply(reference_temp.get(), &reference);
std::unordered_map<llama_token, float> reference_by_id;
for (size_t j = 0; j < reference.size; ++j) {
reference_by_id.emplace(reference.data[j].id, reference.data[j].logit);
}
size_t n_backend_only = 0;
int32_t sampled_index = -1;
float prob_sum = 0.0f;
for (uint32_t j = 0; j < n_logits; ++j) {
GGML_ASSERT(sampled_candidates[j] >= 0 && sampled_candidates[j] < n_vocab);
GGML_ASSERT(std::isfinite(sampled_probs[j]));
GGML_ASSERT(sampled_probs[j] >= 0.0f);
prob_sum += sampled_probs[j];
if (sampled_candidates[j] == backend_token) {
sampled_index = j;
}
if (!std::isfinite(sampled_logits[j])) {
GGML_ASSERT(std::isinf(sampled_logits[j]) && sampled_logits[j] < 0.0f);
GGML_ASSERT(sampled_probs[j] == 0.0f);
continue;
}
const auto match = reference_by_id.find(sampled_candidates[j]);
if (match == reference_by_id.end()) {
++n_backend_only;
continue;
}
const float tolerance = 1e-4f * std::max(1.0f, std::fabs(match->second));
GGML_ASSERT(std::fabs(sampled_logits[j] - match->second) <= tolerance);
reference_by_id.erase(match);
}
const size_t n_reference_only = reference_by_id.size();
if (n_backend_only != 0 || n_reference_only != 0) {
GGML_ASSERT(n_backend_only <= 1);
GGML_ASSERT(n_reference_only <= 1);
GGML_ASSERT(boundary_distance <= cdf_epsilon);
}
GGML_ASSERT(sampled_index >= 0);
GGML_ASSERT(std::isfinite(sampled_logits[sampled_index]));
GGML_ASSERT(sampled_probs[sampled_index] > 0.0f);
GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
}
llama_batch_free(batch);
batch = make_batch(2);
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
llama_batch_free(batch);
printf("backend multi-output sampling chain test PASSED\n");
}
static void test_backend_multi_output_cpu_suffix(const test_params & params) {
const llama_seq_id seq_id = 0;
const int32_t k = 8;
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
auto make_chain = [&](test_single_output_backend_sampler ** sampler_ctx) {
llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
llama_sampler_chain_add(result.get(), test_single_output_backend_sampler_init(sampler_ctx));
llama_sampler_chain_add(result.get(), llama_sampler_init_dist(88));
return result;
};
{
test_single_output_backend_sampler * sampler_ctx = nullptr;
llama_sampler_ptr chain = make_chain(&sampler_ctx);
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
test_context test_ctx(params, configs, 1, 1, 0, 4);
llama_batch batch = llama_batch_init(1, 0, 1);
common_batch_add(batch, llama_vocab_bos(vocab), 0, { seq_id }, true);
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
GGML_ASSERT(sampler_ctx->backend_initialized);
GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 1);
GGML_ASSERT(sampler_ctx->backend_apply_count > 0);
GGML_ASSERT(sampler_ctx->apply_count == 0);
GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), 0) != LLAMA_TOKEN_NULL);
llama_batch_free(batch);
}
{
test_single_output_backend_sampler * sampler_ctx = nullptr;
llama_sampler_ptr chain = make_chain(&sampler_ctx);
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
test_context test_ctx(params, configs, 1, 2, 0, 0);
llama_batch batch = llama_batch_init(2, 0, 1);
for (int i = 0; i < 2; ++i) {
common_batch_add(batch, llama_vocab_bos(vocab), i, { seq_id }, true);
}
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
GGML_ASSERT(!sampler_ctx->backend_initialized);
GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 2);
GGML_ASSERT(sampler_ctx->backend_apply_count == 0);
for (int i = 0; i < batch.n_tokens; ++i) {
GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), i) == LLAMA_TOKEN_NULL);
GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
const llama_token token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i);
GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
}
GGML_ASSERT(sampler_ctx->apply_count == batch.n_tokens);
llama_batch_free(batch);
}
printf("backend multi-output CPU suffix test PASSED\n");
}
struct backend_test_case {
@@ -1583,7 +2010,11 @@ static const backend_test_case BACKEND_TESTS[] = {
{ "dist", test_backend_dist_sampling, true },
{ "dist_and_cpu", test_backend_dist_sampling_and_cpu, true },
{ "set_sampler", test_backend_set_sampler, true },
{ "max_outputs", test_backend_max_outputs, true },
{ "multi_output_limit", test_backend_multi_output_limit, true },
{ "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true },
{ "multi_output_dist_transaction", test_backend_multi_output_dist_transaction, true },
{ "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true },
{ "multi_output_cpu", test_backend_multi_output_cpu_suffix, true },
{ "mixed", test_backend_mixed_sampling, true },
{ "min_p", test_backend_min_p_sampling, true },
{ "cpu_mixed", test_backend_cpu_mixed_batch, true },
tests/test-dt3-gpu.cpp
+121 -74
View File
@@ -18,7 +18,10 @@
// the elementwise maximum is reported as information only, since it explodes
// on cancellation whenever a true output value is near zero.
//
// MUL_MAT with more destination columns than the MMVQ limit falls back to
// MUL_MAT with more destination columns than the MMVQ limit takes the MMQ
// path where the backend implements it for DT3 (CUDA on Ampere-class
// hardware and newer): integer dot products in the same numerical regime as
// MMVQ, judged just as strictly. Backends without DT3 MMQ fall back to
// dequantization + cuBLAS GEMM, which on fast-fp16 hardware rounds the
// dequantized weights to fp16. DT3 weights (d1*t1 + d2*t2, the sum of two
// fp16-scaled terms) are generally NOT fp16-representable, so that path is
@@ -33,8 +36,12 @@
// random bytes: every byte value 0..255 must decode identically on both
// sides, including values >= 243 that never come out of the packer.
//
// DT3 is implemented for CUDA and HIP only. Without one of those backends the
// test is skipped and succeeds — an unsupported backend is not a failure.
// DT3 is implemented for CUDA, HIP and Vulkan. Without one of those backends
// the test is skipped and succeeds — an unsupported backend is not a failure.
// On Vulkan the n <= 8 path is the scalar mul_mat_vec shader (fp32 dot on
// exactly decoded weights, not an integer dot), and the larger-n path is
// dequantization to fp16 + the f16 matmul pipeline; both are judged by the
// same gates as the CUDA MMVQ/GEMM paths.
#include "ggml.h"
#include "ggml-alloc.h"
@@ -260,7 +267,8 @@ static int test_mul_mat(ggml_backend_t backend, ggml_type type, const std::vecto
const std::vector<float> & ref_w, const std::vector<float> & y, bool strict) {
int num_failed = 0;
const int ncols_dst[] = {1, 2, 5, 8, 16};
// 100 exercises a wide MMQ tile with a clamped last column block
const int ncols_dst[] = {1, 2, 5, 8, 16, 100};
// the same weights as the fp16 GEMM fallback sees them
std::vector<float> ref_w16(ref_w.size());
@@ -270,6 +278,8 @@ static int test_mul_mat(ggml_backend_t backend, ggml_type type, const std::vecto
std::vector<std::vector<float>> results;
bool n16_integer = false; // whether the ncols_dst = 16 run took an integer (MMQ) path
for (int c = 0; c < (int)(sizeof(ncols_dst)/sizeof(ncols_dst[0])); ++c) {
const int n = ncols_dst[c];
@@ -320,25 +330,33 @@ static int test_mul_mat(ggml_backend_t backend, ggml_type type, const std::vecto
const mat_err err16 = compare_mat(gpu, ref16);
// n <= 8 is the MMVQ path with exact integer dot products, judged
// against the exact reference. Larger n is the dequantize + GEMM
// fallback whose numerics (fp16 or TF32 compute, depending on the
// against the exact reference (this mirrors MMVQ_MAX_BATCH_SIZE (8)
// from ggml-cuda/mmvq.cu by hand, because the constant and the
// per-arch should_use_mmvq tables are not exported). Larger n takes
// the MMQ path where the backend implements it for this type: integer
// dot products in the same numerical regime as MMVQ, judged just as
// strictly. Backends without MMQ for the type fall back to dequantize
// + GEMM, whose numerics (fp16 or TF32 compute, depending on the
// hardware and on GGML_CUDA_CUBLAS_COMPUTE_TYPE) are cuBLAS's, not
// ours: for the strict type it is gated below by bit-identity with
// the same GEMM on an F16 tensor, and only reported here.
// This mirrors MMVQ_MAX_BATCH_SIZE (8) from ggml-cuda/mmvq.cu by hand,
// because the constant and the per-arch should_use_mmvq tables are not
// exported. If upstream raises the limit, or an architecture routes a
// larger batch through MMVQ, this gating goes stale silently: n = 16
// would take the MMVQ path but still be judged as the GEMM one, which
// only loosens the check, never tightens it. Whoever touches the MMVQ
// dispatch should revisit this line.
const bool is_mmvq = n <= 8;
const bool gated = is_mmvq || !strict;
const double err_gate = is_mmvq ? err.norm_rel : (err.norm_rel < err16.norm_rel ? err.norm_rel : err16.norm_rel);
// ours: for the strict type that run is gated below by bit-identity
// with the same GEMM on an F16 tensor, and only reported here.
// The two regimes are told apart by the result itself: an integer path
// lands within float rounding of the exact reference, a fp16/TF32 GEMM
// stays orders of magnitude above it. A broken MMQ kernel cannot hide
// in the GEMM class: it would then have to be bit-identical to the F16
// GEMM control below, which an integer path never is.
const bool is_mmvq = n <= 8;
const bool integer_path = is_mmvq || err.norm_rel <= 1e-5;
if (n == 16) {
n16_integer = integer_path;
}
const bool gated = integer_path || !strict;
const double err_gate = integer_path ? err.norm_rel : (err.norm_rel < err16.norm_rel ? err.norm_rel : err16.norm_rel);
const double tol = strict ? 1e-5 : 1e-2;
const bool failed = gated && err_gate > tol;
printf("%s: %s mul_mat GPU, ncols_dst = %2d (%s): norm rel err vs exact ref = %g, vs fp16 ref = %g (max elem rel: %g)\n",
failed ? "FAILED" : gated ? "OK" : "INFO", ggml_type_name(type), n, is_mmvq ? "MMVQ" : "GEMM",
printf("%s: %s mul_mat GPU, ncols_dst = %3d (%s): norm rel err vs exact ref = %g, vs fp16 ref = %g (max elem rel: %g)\n",
failed ? "FAILED" : gated ? "OK" : "INFO", ggml_type_name(type), n,
is_mmvq ? "MMVQ" : integer_path ? "MMQ" : "GEMM",
err.norm_rel, err16.norm_rel, err.max_rel);
if (failed) {
num_failed++;
@@ -348,24 +366,26 @@ static int test_mul_mat(ggml_backend_t backend, ggml_type type, const std::vecto
ggml_free(ctx);
}
// MMVQ vs the dequantization-based path: the first 8 columns of the GEMM
// run must match the ncols_dst = 8 MMVQ run to fp16 weight rounding
// MMVQ vs the batched path: the first 8 columns of the ncols_dst = 16 run
// must match the ncols_dst = 8 MMVQ run. When the batched run took the
// integer MMQ path both sides are exact to float rounding of the
// accumulation; against a GEMM fallback the gate is fp16 weight rounding.
{
const std::vector<float> & mmvq = results[3]; // n = 8
const std::vector<float> & gemm = results[4]; // n = 16
const std::vector<float> & mmvq = results[3]; // n = 8
const std::vector<float> & batch = results[4]; // n = 16
double num = 0.0;
double den = 0.0;
for (int j = 0; j < 8; ++j) {
for (int r = 0; r < NROWS; ++r) {
const double diff = (double)mmvq[(size_t)j*NROWS + r] - (double)gemm[(size_t)j*NROWS + r];
const double diff = (double)mmvq[(size_t)j*NROWS + r] - (double)batch[(size_t)j*NROWS + r];
num += diff*diff;
den += (double)mmvq[(size_t)j*NROWS + r]*(double)mmvq[(size_t)j*NROWS + r];
}
}
const double norm_rel = sqrt(num/den);
const double tol = strict ? 5e-3 : 1e-2;
printf("%s: %s MMVQ vs GEMM path on shared columns: norm rel err = %g\n",
norm_rel <= tol ? "OK" : "FAILED", ggml_type_name(type), norm_rel);
const double tol = !strict ? 1e-2 : n16_integer ? 1e-5 : 5e-3;
printf("%s: %s MMVQ vs %s path on shared columns: norm rel err = %g\n",
norm_rel <= tol ? "OK" : "FAILED", ggml_type_name(type), n16_integer ? "MMQ" : "GEMM", norm_rel);
if (norm_rel > tol) {
num_failed++;
}
@@ -375,50 +395,73 @@ static int test_mul_mat(ggml_backend_t backend, ggml_type type, const std::vecto
// tensor holding fp16(dequant(block))": running the same GEMM with an
// F16 src0 built from the fp16-rounded reference weights must give a
// bit-identical result. This isolates our (already bit-validated)
// dequantization from cuBLAS numerics.
if (strict) {
ggml_init_params params = {
/*.mem_size =*/ ggml_tensor_overhead()*8 + ggml_graph_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context * ctx = ggml_init(params);
// dequantization from cuBLAS numerics. The backend may run the DT3
// fallback at a different accumulator precision than its default F16
// GEMM (Vulkan forces fp32 accumulators for DT3), so the F16 control is
// run at both the default and the F32-forced precision and bit-identity
// with either one passes. When the ncols_dst = 16 run took the integer
// MMQ path there is no dequantization involved and no GEMM to compare
// against — that run was already gated strictly above.
if (strict && n16_integer) {
printf("OK: %s ncols_dst = 16 took the integer MMQ path, F16 GEMM bit-identity control not applicable\n",
ggml_type_name(type));
}
if (strict && !n16_integer) {
int n_mismatch_best = -1;
double max_diff_best = 0.0;
ggml_tensor * a16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, NCOLS, NROWS);
ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, NCOLS, 16);
ggml_tensor * out = ggml_mul_mat(ctx, a16, b);
for (int force_f32_prec = 0; force_f32_prec < 2; ++force_f32_prec) {
ggml_init_params params = {
/*.mem_size =*/ ggml_tensor_overhead()*8 + ggml_graph_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context * ctx = ggml_init(params);
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
GGML_ASSERT(buf != nullptr);
std::vector<ggml_fp16_t> w16(ref_w.size());
for (size_t i = 0; i < ref_w.size(); ++i) {
w16[i] = ggml_fp32_to_fp16(ref_w[i]);
}
ggml_backend_tensor_set(a16, w16.data(), 0, w16.size()*sizeof(ggml_fp16_t));
ggml_backend_tensor_set(b, y.data(), 0, (size_t)NCOLS*16*sizeof(float));
std::vector<float> gpu16((size_t)NROWS*16);
compute_graph(backend, ctx, out, gpu16.data());
const std::vector<float> & gemm = results[4]; // n = 16
int n_mismatch = 0;
double max_diff = 0.0;
for (size_t i = 0; i < gemm.size(); ++i) {
const double diff = fabs((double)gemm[i] - (double)gpu16[i]);
max_diff = diff > max_diff ? diff : max_diff;
if (gemm[i] != gpu16[i]) {
n_mismatch++;
ggml_tensor * a16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, NCOLS, NROWS);
ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, NCOLS, 16);
ggml_tensor * out = ggml_mul_mat(ctx, a16, b);
if (force_f32_prec) {
ggml_mul_mat_set_prec(out, GGML_PREC_F32);
}
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
GGML_ASSERT(buf != nullptr);
std::vector<ggml_fp16_t> w16(ref_w.size());
for (size_t i = 0; i < ref_w.size(); ++i) {
w16[i] = ggml_fp32_to_fp16(ref_w[i]);
}
ggml_backend_tensor_set(a16, w16.data(), 0, w16.size()*sizeof(ggml_fp16_t));
ggml_backend_tensor_set(b, y.data(), 0, (size_t)NCOLS*16*sizeof(float));
std::vector<float> gpu16((size_t)NROWS*16);
compute_graph(backend, ctx, out, gpu16.data());
const std::vector<float> & gemm = results[4]; // n = 16
int n_mismatch = 0;
double max_diff = 0.0;
for (size_t i = 0; i < gemm.size(); ++i) {
const double diff = fabs((double)gemm[i] - (double)gpu16[i]);
max_diff = diff > max_diff ? diff : max_diff;
if (gemm[i] != gpu16[i]) {
n_mismatch++;
}
}
if (n_mismatch_best < 0 || n_mismatch < n_mismatch_best) {
n_mismatch_best = n_mismatch;
max_diff_best = max_diff;
}
ggml_backend_buffer_free(buf);
ggml_free(ctx);
}
printf("%s: %s GEMM path vs F16 GEMM on fp16-rounded weights: %d mismatches, max |diff| = %g\n",
n_mismatch == 0 ? "OK" : "FAILED", ggml_type_name(type), n_mismatch, max_diff);
if (n_mismatch != 0) {
printf("%s: %s GEMM path vs F16 GEMM on fp16-rounded weights (best of default/F32 prec): %d mismatches, max |diff| = %g\n",
n_mismatch_best == 0 ? "OK" : "FAILED", ggml_type_name(type), n_mismatch_best, max_diff_best);
if (n_mismatch_best != 0) {
num_failed++;
}
ggml_backend_buffer_free(buf);
ggml_free(ctx);
}
// manual sum over the trits stored by the test for row 0, column 0 —
@@ -461,21 +504,25 @@ static void build_control_data(ggml_type type, std::vector<uint8_t> & data, std:
}
int main(void) {
// Only CUDA and HIP (which reports itself as "ROCm") implement DT3. Any
// other GPU backend is skipped rather than failed: Vulkan and SYCL answer
// supports_op == false for DT3, which is the correct answer for them and
// Only CUDA, HIP (which reports itself as "ROCm") and Vulkan implement
// DT3. Any other GPU backend is skipped rather than failed: SYCL answers
// supports_op == false for DT3, which is the correct answer for it and
// not a bug to report, and Metal answers true for almost any type but has
// no DT3 shader, so it would die in pipeline compilation mid-test. Picking
// the backend by name keeps this test honest on machines we do not have.
ggml_backend_t backend = nullptr;
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) {
// IGPU is a distinct device type from GPU: an integrated Vulkan device
// with unified memory reports as IGPU, and accepting only GPU silently
// skipped the very hardware this backend is for.
const auto dt = ggml_backend_dev_type(dev);
if (dt != GGML_BACKEND_DEVICE_TYPE_GPU && dt != GGML_BACKEND_DEVICE_TYPE_IGPU) {
continue;
}
const char * name = ggml_backend_dev_name(dev);
if (strncmp(name, "CUDA", 4) != 0 && strncmp(name, "ROCm", 4) != 0) {
printf("skipping GPU backend %s: DT3 is only implemented for CUDA/HIP\n", name);
if (strncmp(name, "CUDA", 4) != 0 && strncmp(name, "ROCm", 4) != 0 && strncmp(name, "Vulkan", 6) != 0) {
printf("skipping GPU backend %s: DT3 is only implemented for CUDA/HIP/Vulkan\n", name);
continue;
}
backend = ggml_backend_dev_init(dev, nullptr);
@@ -483,7 +530,7 @@ int main(void) {
break;
}
if (backend == nullptr) {
printf("no CUDA/HIP backend available, skipping\n");
printf("no CUDA/HIP/Vulkan backend available, skipping\n");
return 0;
}
@@ -497,7 +544,7 @@ int main(void) {
// activations: integers with amax 127 in every 32-element chunk of every
// column, so their q8_1 quantization is exact
std::vector<float> y((size_t)NCOLS*16);
std::vector<float> y((size_t)NCOLS*100);
for (size_t i = 0; i < y.size(); ++i) {
y[i] = i % 32 == 0 ? 127.0f : (float)((int)(rng_next() % 255) - 127);
}
tests/test-llama-archs.cpp
+1 -1
View File
@@ -192,7 +192,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f);
// SWA pattern: every 5th layer is full attention (matches E2B layer_types)
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5));
} else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35) {
} else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER) {
std::vector<uint32_t> pattern;
pattern.reserve(n_layer);
for (uint32_t il = 0; il < n_layer; il++) {
tests/test-sampling.cpp
+31
View File
@@ -61,6 +61,35 @@ private:
std::vector<llama_token_data> cur;
};
static llama_token sample_dist(llama_sampler * sampler, const std::vector<float> & logits) {
std::vector<llama_token_data> cur;
for (llama_token token_id = 0; token_id < (llama_token) logits.size(); ++token_id) {
cur.push_back({ token_id, logits[token_id], 0.0f });
}
llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };
llama_sampler_apply(sampler, &cur_p);
GGML_ASSERT(cur_p.selected >= 0);
GGML_ASSERT((size_t) cur_p.selected < cur_p.size);
return cur_p.data[cur_p.selected].id;
}
static void test_dist_singleton_rng() {
llama_sampler * singleton = llama_sampler_init_dist(4242);
llama_sampler * control = llama_sampler_init_dist(4242);
sample_dist(singleton, { 0.0f });
sample_dist(control, { 0.0f, 0.0f });
const std::vector<float> logits(256, 0.0f);
for (int i = 0; i < 4; ++i) {
GGML_ASSERT(sample_dist(singleton, logits) == sample_dist(control, logits));
}
llama_sampler_free(singleton);
llama_sampler_free(control);
}
static void test_temp(const std::vector<float> & probs, const std::vector<float> & probs_expected, float temp) {
sampler_tester tester(probs, probs_expected);
@@ -308,6 +337,8 @@ static void test_perf() {
int main(void) {
ggml_time_init();
test_dist_singleton_rng();
test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 1.0f);
test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.0f, 0.0f, 0.0f, 1.0f}, 0.0f);
tools/cli/README.md
+1 -2
View File
@@ -54,6 +54,7 @@
| `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
| `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
| `-np, --parallel N` | number of parallel sequences to decode (default: 1)<br/>(env: LLAMA_ARG_N_PARALLEL) |
| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
| `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
| `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
@@ -84,8 +85,6 @@
| `-dr, --docker-repo [<repo>/]<model>[:quant]` | Docker Hub model repository. repo is optional, default to ai/. quant is optional, default to :latest.<br/>example: gemma3<br/>(default: unused)<br/>(env: LLAMA_ARG_DOCKER_REPO) |
| `-hf, -hfr, --hf-repo <user>/<model>[:quant]` | Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.<br/>mmproj is also downloaded automatically if available. to disable, add --no-mmproj<br/>example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M<br/>(default: unused)<br/>(env: LLAMA_ARG_HF_REPO) |
| `-hff, --hf-file FILE` | Hugging Face model file. If specified, it will override the quant in --hf-repo (default: unused)<br/>(env: LLAMA_ARG_HF_FILE) |
| `-hfv, -hfrv, --hf-repo-v <user>/<model>[:quant]` | Hugging Face model repository for the vocoder model (default: unused)<br/>(env: LLAMA_ARG_HF_REPO_V) |
| `-hffv, --hf-file-v FILE` | Hugging Face model file for the vocoder model (default: unused)<br/>(env: LLAMA_ARG_HF_FILE_V) |
| `-hft, --hf-token TOKEN` | Hugging Face access token (default: value from HF_TOKEN environment variable)<br/>(env: HF_TOKEN) |
| `--log-disable` | Log disable |
| `--log-file FNAME` | Log to file<br/>(env: LLAMA_ARG_LOG_FILE) |
tools/completion/README.md
+1 -2
View File
@@ -137,6 +137,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
| `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
| `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
| `-np, --parallel N` | number of parallel sequences to decode (default: 1)<br/>(env: LLAMA_ARG_N_PARALLEL) |
| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
| `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
| `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
@@ -167,8 +168,6 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
| `-dr, --docker-repo [<repo>/]<model>[:quant]` | Docker Hub model repository. repo is optional, default to ai/. quant is optional, default to :latest.<br/>example: gemma3<br/>(default: unused)<br/>(env: LLAMA_ARG_DOCKER_REPO) |
| `-hf, -hfr, --hf-repo <user>/<model>[:quant]` | Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.<br/>mmproj is also downloaded automatically if available. to disable, add --no-mmproj<br/>example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M<br/>(default: unused)<br/>(env: LLAMA_ARG_HF_REPO) |
| `-hff, --hf-file FILE` | Hugging Face model file. If specified, it will override the quant in --hf-repo (default: unused)<br/>(env: LLAMA_ARG_HF_FILE) |
| `-hfv, -hfrv, --hf-repo-v <user>/<model>[:quant]` | Hugging Face model repository for the vocoder model (default: unused)<br/>(env: LLAMA_ARG_HF_REPO_V) |
| `-hffv, --hf-file-v FILE` | Hugging Face model file for the vocoder model (default: unused)<br/>(env: LLAMA_ARG_HF_FILE_V) |
| `-hft, --hf-token TOKEN` | Hugging Face access token (default: value from HF_TOKEN environment variable)<br/>(env: HF_TOKEN) |
| `--log-disable` | Log disable |
| `--log-file FNAME` | Log to file<br/>(env: LLAMA_ARG_LOG_FILE) |
tools/mtmd/CMakeLists.txt
+1
View File
@@ -43,6 +43,7 @@ add_library(mtmd
models/kimivl.cpp
models/kimik25.cpp
models/nemotron-v2-vl.cpp
models/muse-glimmer.cpp
models/llama4.cpp
models/llava.cpp
models/minicpmv.cpp
tools/mtmd/clip-impl.h
+2
View File
@@ -455,6 +455,7 @@ enum projector_type {
PROJECTOR_TYPE_MIMO_AUDIO,
PROJECTOR_TYPE_QWEN3TTS_SPKENC,
PROJECTOR_TYPE_QWEN3TTS_GEN,
PROJECTOR_TYPE_MUSE_GLIMMER,
PROJECTOR_TYPE_UNKNOWN,
};
@@ -514,6 +515,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_PARAKEET, "parakeet"},
{ PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"},
{ PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"},
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
};
static projector_type clip_projector_type_from_string(const std::string & str) {
tools/mtmd/clip-model.h
+5
View File
@@ -109,6 +109,11 @@ struct clip_hparams {
int32_t downsample_query_side;
int32_t downsample_window_side;
// Muse Glimmer vision (per-block sparse-window pattern, learned pos-emb, patch-temporal)
// NOTE: these perhaps shouldn't have the architecture prefix
int32_t muse_glimmer_patch_temporal = 0;
int32_t muse_glimmer_sparse_factor = 0;
// audio
int32_t n_mel_bins = 0; // whisper preprocessor
int32_t proj_stack_factor = 0; // ultravox
tools/mtmd/clip.cpp
+91
View File
@@ -954,6 +954,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
{
builder = std::make_unique<clip_graph_minimax_m3>(ctx, img);
} break;
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
builder = std::make_unique<clip_graph_muse_glimmer>(ctx, img);
} break;
case PROJECTOR_TYPE_STEP3VL:
{
builder = std::make_unique<clip_graph_step3vl>(ctx, img);
@@ -1572,6 +1576,17 @@ struct clip_model_loader {
hparams.set_limit_image_tokens(8, 576);
hparams.set_warmup_n_tokens(16*16);
} break;
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
hparams.n_merge = 2; // pixel-shuffle downsample after the ViT
hparams.image_resize_algo = RESIZE_ALGO_LANCZOS;
hparams.rope_theta = 10000.0f;
hparams.muse_glimmer_patch_temporal = 2;
hparams.muse_glimmer_sparse_factor = 4; // 3 sparse layers + 1 global, repeating
get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
hparams.set_limit_image_tokens(1, 4096);
hparams.set_warmup_n_tokens(32*32);
} break;
case PROJECTOR_TYPE_MIMOVL:
{
hparams.n_merge = 2; // spatial_merge_size
@@ -2317,6 +2332,13 @@ struct clip_model_loader {
model.mm_merger_fc2_w = get_tensor(string_format(TN_MM_MERGER_FC2, "weight"));
model.mm_merger_fc2_b = get_tensor(string_format(TN_MM_MERGER_FC2, "bias"));
} break;
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
// 3-linear MLP: fc -> erf-GELU -> proj -> erf-GELU -> vision_proj (into LLM residual dim)
model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight"));
model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
} break;
case PROJECTOR_TYPE_STEP3VL:
{
model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight"));
@@ -3745,6 +3767,7 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) {
case PROJECTOR_TYPE_PADDLEOCR:
case PROJECTOR_TYPE_HUNYUANVL:
case PROJECTOR_TYPE_YOUTUVL:
case PROJECTOR_TYPE_MUSE_GLIMMER:
return (img->nx() / params.patch_size) / 2;
case PROJECTOR_TYPE_STEP3VL:
return img->nx() / (params.patch_size * params.n_merge);
@@ -3770,6 +3793,7 @@ int clip_n_output_tokens_y(const clip_ctx * ctx, const clip_image_f32 * img) {
case PROJECTOR_TYPE_PADDLEOCR:
case PROJECTOR_TYPE_HUNYUANVL:
case PROJECTOR_TYPE_YOUTUVL:
case PROJECTOR_TYPE_MUSE_GLIMMER:
return (img->ny() / params.patch_size) / 2;
case PROJECTOR_TYPE_STEP3VL:
return img->ny() / (params.patch_size * params.n_merge);
@@ -3848,6 +3872,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
case PROJECTOR_TYPE_MINIMAX_M3:
case PROJECTOR_TYPE_GLM4V:
case PROJECTOR_TYPE_YOUTUVL:
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
// dynamic size (2 conv, so double patch size)
int x_patch = img->nx() / (params.patch_size * 2);
@@ -4193,6 +4218,70 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
// set input per projector
switch (ctx->model.proj_type) {
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
const int grid_w = pos_w; // image_size_width / patch_size
const int grid_h = pos_h; // image_size_height / patch_size
const int n_tok = grid_w * grid_h;
const int pgrid = (int) std::sqrt((double) ctx->model.position_embeddings->ne[1]); // 32
const int f = hparams.n_merge; // downsample 2
// pixel patchify runs inside the graph via build_inp() (ggml_conv_2d);
// pos-emb bilinear interp via resize_position_embeddings().
// --- sparse window grouping (pgrid x pgrid windows) ---
const int win = pgrid;
const int nwin_h = (grid_h + win - 1) / win;
const int nwin_w = (grid_w + win - 1) / win;
std::vector<int32_t> sp_perm; sp_perm.reserve(n_tok);
std::vector<int> sp_slens;
for (int wy = 0; wy < nwin_h; wy++) {
for (int wx = 0; wx < nwin_w; wx++) {
int cnt = 0;
for (int hh = 0; hh < win; hh++) {
for (int ww = 0; ww < win; ww++) {
const int gy = wy * win + hh;
const int gx = wx * win + ww;
if (gy < grid_h && gx < grid_w) { sp_perm.push_back(gy * grid_w + gx); cnt++; }
}
}
if (cnt > 0) sp_slens.push_back(cnt);
}
}
std::vector<int32_t> rpos_w(n_tok), rpos_h(n_tok), inv_perm(n_tok);
for (int i = 0; i < n_tok; i++) {
const int orig = sp_perm[i];
rpos_w[i] = (orig % grid_w) + 1; // 1-indexed
rpos_h[i] = (orig / grid_w) + 1;
inv_perm[orig] = i;
}
set_input_i32("muse_glimmer_sp_perm", sp_perm);
set_input_i32("muse_glimmer_inv_perm", inv_perm);
set_input_i32("muse_glimmer_pos_w", rpos_w);
set_input_i32("muse_glimmer_pos_h", rpos_h);
// block-diagonal window mask (permuted order)
std::vector<float> sp_mask((size_t) n_tok * n_tok, -INFINITY);
{
int off = 0;
for (int s : sp_slens) {
for (int a = 0; a < s; a++)
for (int b = 0; b < s; b++)
sp_mask[(size_t) (off + a) * n_tok + (off + b)] = 0.0f;
off += s;
}
}
set_input_f32("muse_glimmer_sp_mask", sp_mask);
// pixel-shuffle gather (original order): f*f spatial neighbours grouped
std::vector<int32_t> dsp; dsp.reserve(n_tok);
for (int oy = 0; oy < grid_h / f; oy++)
for (int ox = 0; ox < grid_w / f; ox++)
for (int ry = 0; ry < f; ry++)
for (int rx = 0; rx < f; rx++)
dsp.push_back((oy * f + ry) * grid_w + (ox * f + rx));
set_input_i32("muse_glimmer_ds_perm", dsp);
} break;
case PROJECTOR_TYPE_MINICPMV:
{
// inspired from siglip:
@@ -5369,6 +5458,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
return ctx->model.mm_model_mlp_3_w->ne[1];
case PROJECTOR_TYPE_MINIMAX_M3:
return ctx->model.mm_merger_fc2_b->ne[0];
case PROJECTOR_TYPE_MUSE_GLIMMER:
return ctx->model.mm_2_w->ne[1];
case PROJECTOR_TYPE_QWEN2VL:
case PROJECTOR_TYPE_QWEN25VL:
case PROJECTOR_TYPE_EXAONE4_5:
tools/mtmd/models/models.h
+5
View File
@@ -365,3 +365,8 @@ private:
ggml_tensor * build_newline_row(ggml_context * ctx0);
ggml_tensor * append_rowwise_newlines(ggml_context * ctx0, ggml_tensor * tile_output);
};
struct clip_graph_muse_glimmer : clip_graph {
clip_graph_muse_glimmer(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
ggml_cgraph * build() override;
};
tools/mtmd/models/muse-glimmer.cpp
+88
View File
@@ -0,0 +1,88 @@
#include "models.h"
// MuseGlimmer vision encoder: 50-layer ViT with 2D RoPE, sparse block-diagonal
// window attention (every 4th + last layer global), pixel-shuffle downsample, then
// adapter MLP + LLM's vision_projection.
//
// Several quantities are precomputed on host and fed as named graph inputs (filled in
// clip.cpp set_input, PROJECTOR_TYPE_MUSE_GLIMMER branch):
// muse_glimmer_pos_w/_h [n_tok] i32 : 1-indexed RoPE positions (sparse-permuted order)
// muse_glimmer_sp_perm [n_tok] i32 : window grouping permutation (applied after ln_pre)
// muse_glimmer_inv_perm [n_tok] i32 : inverse of sp_perm (applied after blocks)
// muse_glimmer_ds_perm [n_tok] i32 : pixel-shuffle gather (original order)
// muse_glimmer_sp_mask [n_tok, n_tok] f32 : block-diagonal window mask (sparse layers)
ggml_cgraph * clip_graph_muse_glimmer::build() {
const int ds = hparams.n_merge; // downsample factor (2)
const int sf = hparams.muse_glimmer_sparse_factor; // 4
const int n_tok = n_patches;
const int n_out = (n_patches_x / ds) * (n_patches_y / ds);
const float rope_base = hparams.rope_theta; // 10000
auto inp_i32 = [&](const char * name, int64_t n) {
ggml_tensor * t = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n);
ggml_set_name(t, name);
ggml_set_input(t);
return t;
};
ggml_tensor * pos_w = inp_i32("muse_glimmer_pos_w", n_tok);
ggml_tensor * pos_h = inp_i32("muse_glimmer_pos_h", n_tok);
ggml_tensor * sp_perm = inp_i32("muse_glimmer_sp_perm", n_tok);
ggml_tensor * inv_perm = inp_i32("muse_glimmer_inv_perm", n_tok);
ggml_tensor * ds_perm = inp_i32("muse_glimmer_ds_perm", n_tok);
ggml_tensor * sp_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tok, n_tok);
ggml_set_name(sp_mask, "muse_glimmer_sp_mask");
ggml_set_input(sp_mask);
// patchify via build_inp (conv2d over raw pixels) + bilinear-resized learned pos-emb
ggml_tensor * x = build_inp(); // [n_embd, n_tok, 1]
x = ggml_add(ctx0, x, resize_position_embeddings(GGML_SCALE_MODE_BILINEAR));
cb(x, "after_posemb", -1);
// group patches into pgrid x pgrid windows (sparse attention order)
x = ggml_get_rows(ctx0, x, sp_perm);
cb(x, "after_sp_perm", -1);
// per-layer mask: sparse layers get sp_mask, global layers (every sf-th and last) get none
std::vector<ggml_tensor *> attn_mask_layers(n_layer);
for (int il = 0; il < n_layer; ++il) {
const bool is_global = (il == n_layer - 1) || ((il + 1) % sf == 0);
attn_mask_layers[il] = is_global ? nullptr : sp_mask;
}
// 2D RoPE: first half of head_dim uses width pos, second half uses height pos
auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
return build_rope_2d(ctx0, cur, pos_w, pos_h, rope_base, false);
};
build_vit_opts opts;
opts.attn_mask_layers = std::move(attn_mask_layers);
// pre_ln, per-layer transformer, post_ln (all inside build_vit); reference uses exact (erf) GELU
x = build_vit(x, n_tok, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr, add_pos, opts);
// un-permute back to original grid order
x = ggml_get_rows(ctx0, x, inv_perm);
cb(x, "after_inv_perm", -1);
// pixel-shuffle downsample: gather f*f spatial neighbors then concat channel-outer.
// out[c*(ds*ds)+s, o] = x[ds_perm gathered][o*(ds*ds)+s, c]
x = ggml_get_rows(ctx0, x, ds_perm); // [n_embd, n_tok], grouped
x = ggml_reshape_3d(ctx0, x, n_embd, ds * ds, n_out);// [c, s, o]
x = ggml_permute(ctx0, x, 1, 0, 2, 3); // [s, c, o]
x = ggml_cont(ctx0, x);
x = ggml_reshape_2d(ctx0, x, n_embd * ds * ds, n_out); // [6144, n_out]
cb(x, "encoder_out", -1);
// adapter (6144->4096->4096, exact GELU each) + LLM vision_projection (4096->6656)
x = build_mm(model.mm_0_w, x);
x = ggml_gelu_erf(ctx0, x);
x = build_mm(model.mm_1_w, x);
x = ggml_gelu_erf(ctx0, x);
x = build_mm(model.mm_2_w, x); // [6656, n_out]
cb(x, "projected", -1);
ggml_build_forward_expand(gf, x);
return gf;
}
tools/mtmd/mtmd-image.cpp
+62
View File
@@ -1615,3 +1615,65 @@ mtmd_image_preproc_out mtmd_image_preprocessor_granite::preprocess(const clip_im
}
return output;
}
//
// mtmd_image_preprocessor_muse_glimmer
//
// Replicates transformers' get_aspect_ratio_preserving_size
static clip_image_size muse_glimmer_grid_size(int img_w, int img_h, int patch_hw, int max_tokens) {
double i_nph = (double) img_h / patch_hw;
double i_npw = (double) img_w / patch_hw;
const double ratio = i_nph > 0.0 ? i_npw / i_nph : 1.0;
if (i_nph * i_npw > (double) max_tokens) {
i_nph = std::sqrt((double) max_tokens / ratio);
i_npw = i_nph * ratio;
}
const int hs[2] = { (int) std::floor(i_nph), (int) std::ceil(i_nph) };
const int ws[2] = { (int) std::floor(i_npw), (int) std::ceil(i_npw) };
const double target_ar = (double) img_h / (double) img_w;
int best_nph = -1;
int best_npw = -1;
double best_d = 0.0;
for (int a = 0; a < 2; ++a) {
for (int b = 0; b < 2; ++b) {
const int nph = hs[a];
const int npw = ws[b];
if (nph < 1 || npw < 1 || nph * npw > max_tokens) {
continue;
}
const double d = std::fabs((double) nph / (double) npw - target_ar);
const int n_tokens = nph * npw;
const int best_n_tokens = best_nph * best_npw;
if (best_nph < 0 || d < best_d || (d == best_d && n_tokens > best_n_tokens)) {
best_nph = nph;
best_npw = npw;
best_d = d;
}
}
}
if (best_nph < 0) { // no candidate fit under the cap: round and clamp
best_nph = std::max(1, (int) std::lround(i_nph));
best_npw = std::max(1, (int) std::lround(i_npw));
}
return clip_image_size{ best_npw * patch_hw, best_nph * patch_hw };
}
mtmd_image_preproc_out mtmd_image_preprocessor_muse_glimmer::preprocess(const clip_image_u8 & img) {
const int patch_hw = hparams.patch_size * hparams.n_merge;
const int patch_area = hparams.patch_size * hparams.patch_size * hparams.n_merge * hparams.n_merge;
GGML_ASSERT(patch_area > 0 && hparams.image_max_pixels > 0);
const int max_tokens = hparams.image_max_pixels / patch_area;
const clip_image_size original_size = img.get_size();
const clip_image_size target_size = muse_glimmer_grid_size(
original_size.width, original_size.height, patch_hw, max_tokens);
// PIL resizes directly to (target_w, target_h) -- a stretch, no padding.
clip_image_u8 resized_image;
img_tool::resize(img, resized_image, target_size, hparams.image_resize_algo, PAD_NONE);
mtmd_image_preproc_out output;
output.append(hparams, resized_image, true);
return output;
}
tools/mtmd/mtmd-image.h
+6
View File
@@ -230,3 +230,9 @@ struct mtmd_image_preprocessor_granite : mtmd_image_preprocessor_llava_uhd {
mtmd_image_preprocessor_granite(const clip_ctx * ctx) : mtmd_image_preprocessor_llava_uhd(ctx) {}
mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
};
// pick the patch grid closest to the input aspect ratio under the per-image token cap, stretch-resize.
struct mtmd_image_preprocessor_muse_glimmer : mtmd_image_preprocessor {
mtmd_image_preprocessor_muse_glimmer(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
};
tools/mtmd/mtmd.cpp
+6
View File
@@ -699,6 +699,12 @@ struct mtmd_context {
img_end = "]<]end of image[>[";
image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
} break;
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
img_beg = "<|image_start|>";
img_end = "<|image_end|>";
image_preproc = std::make_unique<mtmd_image_preprocessor_muse_glimmer>(ctx_v);
} break;
case PROJECTOR_TYPE_YOUTUVL:
{
// <|vision_start|> ... (image embeddings) ... <|vision_end|>
tools/server/README-dev.md
+1 -1
View File
@@ -201,7 +201,7 @@ Invoke a tool call, request body is a JSON object with:
Headers:
- `x-tool-cwd`: optional; if set, use as the CWD for tool; this is not part of tool's params because it's meant to be set by the runtime, not the LLM itself
- `x-tool-runtime`: optional; if set, run the tool inside this isolate instead of on the host. Only `docker-container:<id>` is supported for now, using an already-running container
- `x-tool-runtime`: optional; if set, run the tool inside this isolate instead of on the host. Either `docker-container:<id>` or `podman-container:<id>`, using an already-running container, or `ssh:<target>`, running the tool on a remote host
Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string):
tools/server/README.md
+2 -6
View File
@@ -71,6 +71,7 @@ For the full list of features, please refer to [server's changelog](https://gith
| `-ctk, --cache-type-k TYPE` | KV cache data type for K<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_K) |
| `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
| `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
| `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
| `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
@@ -101,8 +102,6 @@ For the full list of features, please refer to [server's changelog](https://gith
| `-dr, --docker-repo [<repo>/]<model>[:quant]` | Docker Hub model repository. repo is optional, default to ai/. quant is optional, default to :latest.<br/>example: gemma3<br/>(default: unused)<br/>(env: LLAMA_ARG_DOCKER_REPO) |
| `-hf, -hfr, --hf-repo <user>/<model>[:quant]` | Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.<br/>mmproj is also downloaded automatically if available. to disable, add --no-mmproj<br/>example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M<br/>(default: unused)<br/>(env: LLAMA_ARG_HF_REPO) |
| `-hff, --hf-file FILE` | Hugging Face model file. If specified, it will override the quant in --hf-repo (default: unused)<br/>(env: LLAMA_ARG_HF_FILE) |
| `-hfv, -hfrv, --hf-repo-v <user>/<model>[:quant]` | Hugging Face model repository for the vocoder model (default: unused)<br/>(env: LLAMA_ARG_HF_REPO_V) |
| `-hffv, --hf-file-v FILE` | Hugging Face model file for the vocoder model (default: unused)<br/>(env: LLAMA_ARG_HF_FILE_V) |
| `-hft, --hf-token TOKEN` | Hugging Face access token (default: value from HF_TOKEN environment variable)<br/>(env: HF_TOKEN) |
| `--log-disable` | Log disable |
| `--log-file FNAME` | Log to file<br/>(env: LLAMA_ARG_LOG_FILE) |
@@ -197,9 +196,8 @@ For the full list of features, please refer to [server's changelog](https://gith
| `--ui-config, --webui-config JSON` | JSON that provides default UI settings (overrides UI defaults)<br/>(env: LLAMA_ARG_UI_CONFIG) |
| `--ui-config-file, --webui-config-file PATH` | JSON file that provides default UI settings (overrides UI defaults)<br/>(env: LLAMA_ARG_UI_CONFIG_FILE) |
| `--ui-mcp-proxy, --webui-mcp-proxy, --no-ui-mcp-proxy, --no-webui-mcp-proxy` | experimental: whether to enable MCP CORS proxy - do not enable in untrusted environments (default: disabled)<br/>(env: LLAMA_ARG_UI_MCP_PROXY) |
| `--tools TOOL1,TOOL2,...` | experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)<br/>specify "all" to enable all tools<br/>available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_TOOLS) |
| `--tools-runtime OPTION` | experimental: run tools in a separate runtime environment (default: none, use host environment)<br/>available options:<br/> 'docker:<image>': spin up a new Docker container and reuse it for all invocations, clean up on server exit<br/> 'docker-container:<id>': use an existing Docker container by ID, won't stop on server exit<br/><br/>(env: LLAMA_ARG_TOOLS_RUNTIME) |
| `--tools TOOL1,TOOL2,...` | experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)<br/>specify "all" to enable all tools<br/>available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime, get_info<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_TOOLS) |
| `--tools-runtime OPTION` | experimental: run tools in a separate runtime environment (default: none, use host environment)<br/>available options:<br/> 'docker:<image>', 'podman:<image>': spin up a new container and reuse it for all invocations, clean up on server exit<br/> 'docker-container:<id>', 'podman-container:<id>': use an existing container by ID, won't stop on server exit<br/> 'ssh:<target>': run tools on a remote POSIX host over SSH, key-based auth and a trusted host key are required<br/><br/>(env: LLAMA_ARG_TOOLS_RUNTIME) |
| `--mcp-servers-config PATH` | experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_MCP_SERVERS_CONFIG) |
| `--mcp-servers-json JSON` | experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_MCP_SERVERS_JSON) |
| `-ag, --agent, -no-ag, --no-agent` | whether to enable CORS proxy and all built-in tools - do not enable in untrusted environments (default: disabled)<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_AGENT) |
@@ -280,8 +278,6 @@ For the full list of features, please refer to [server's changelog](https://gith
| `--spec-ngram-size-n N` | the argument has been removed. use the respective --spec-ngram-*-size-n or --spec-ngram-mod-n-match |
| `--spec-ngram-size-m N` | the argument has been removed. use the respective --spec-ngram-*-size-m |
| `--spec-ngram-min-hits N` | the argument has been removed. use the respective --spec-ngram-*-min-hits |
| `-mv, --model-vocoder FNAME` | vocoder model for audio generation (default: unused) |
| `--tts-use-guide-tokens` | Use guide tokens to improve TTS word recall |
| `--embd-gemma-default` | use default EmbeddingGemma model (note: can download weights from the internet) |
| `--fim-qwen-1.5b-default` | use default Qwen 2.5 Coder 1.5B (note: can download weights from the internet) |
| `--fim-qwen-3b-default` | use default Qwen 2.5 Coder 3B (note: can download weights from the internet) |
tools/server/server-context.cpp
+16 -20
View File
@@ -39,19 +39,18 @@ using json = nlohmann::ordered_json;
constexpr int HTTP_POLLING_SECONDS = 1;
static uint32_t server_n_outputs_max(const common_params & params) {
const uint32_t n_batch = params.n_batch;
static common_speculative_output_limits server_output_limits(const common_params & params) {
if (params.embedding ||
(params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
return n_batch;
return { params.n_batch, 1 };
}
const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(&params.speculative);
auto result = common_speculative_get_output_limits(
params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
result.total = std::max<int32_t>(1, result.total);
result.per_seq = std::max<int32_t>(1, result.per_seq);
return result;
}
// state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
@@ -1063,7 +1062,9 @@ private:
const bool is_resume = sleeping;
params_base = params;
params_base.n_outputs_max = server_n_outputs_max(params_base);
const auto output_limits = server_output_limits(params_base);
params_base.n_outputs_max = output_limits.total;
params_base.n_outputs_max_per_seq = output_limits.per_seq;
const bool has_mmproj = !params.mmproj.path.empty();
const bool has_draft = params.speculative.has_dft();
@@ -1832,18 +1833,13 @@ private:
const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs;
bool backend_sampling = true;
backend_sampling &= task.params.sampling.backend_sampling;
// TODO: speculative decoding requires multiple samples per batch - not supported yet
backend_sampling &= !(slot.can_speculate());
bool use_backend_sampling = task.params.sampling.backend_sampling;
// TODO: getting pre sampling logits is not yet supported with backend sampling
backend_sampling &= !need_pre_sample_logits;
use_backend_sampling &= !need_pre_sample_logits;
// TODO: tmp until backend sampling is fully implemented
if (backend_sampling) {
if (use_backend_sampling) {
llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
} else {
llama_set_sampler(ctx_tgt, slot.id, nullptr);
@@ -3865,7 +3861,8 @@ private:
// speculative decoding - main model sample and accept
iterate(slots, [&](server_slot & slot) {
if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || slot.spec_draft.empty()) {
if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() ||
slot.spec_draft.empty() || slot.spec_i_batch.empty()) {
return;
}
@@ -3876,7 +3873,6 @@ private:
// verify and try to accept the draft
{
// save the sampler sampler state in case we need to restore it
common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1);
@@ -3915,7 +3911,7 @@ private:
slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);
slot.prompt.tokens.keep_first(ckpt.n_tokens);
slot.smpl = std::move(smpl_save);
common_sampler_copy(smpl_save.get(), slot.smpl.get());
return;
}
tools/server/server-tools.cpp
+237 -125
View File
@@ -10,6 +10,7 @@
#include <ctime>
#include <atomic>
#include <cstring>
#include <cctype>
#include <cstdint>
#include <cstdlib>
#include <algorithm>
@@ -25,6 +26,11 @@
# define NOMINMAX
# endif
# include <windows.h>
# include <fcntl.h>
# include <io.h>
#else
# include <cerrno>
# include <unistd.h>
#endif
namespace fs = std::filesystem;
@@ -176,7 +182,7 @@ public:
const std::function<bool(const std::string &)> & on_chunk = nullptr) const = 0;
};
// shared subprocess execution helper, used by both the local and the docker-backed tools_io implementations.
// shared subprocess execution helper, used by both the local and the isolate-backed tools_io implementations.
// combine_stderr=false when the raw stdout bytes must not be tainted by stderr, e.g. reading file contents.
static tools_io::exec_result run_subprocess(
const std::vector<std::string> & args,
@@ -184,7 +190,8 @@ static tools_io::exec_result run_subprocess(
int timeout_secs,
const std::function<bool(const std::string &)> & on_chunk,
bool combine_stderr,
const std::string & cwd = "") {
const std::string & cwd = "",
const std::string * stdin_data = nullptr) {
tools_io::exec_result res;
common_subproc proc;
@@ -216,26 +223,59 @@ static tools_io::exec_result run_subprocess(
}
});
// write stdin before reading stdout, the child drains stdin as it goes
// always close stdin, a transport client waits forever if its stdin pipe stays open
if (FILE * in = proc.stdin_file()) {
if (stdin_data != nullptr && !stdin_data->empty()) {
#if defined(_WIN32)
// pipe fds default to CRT text mode: binary keeps the bytes untranslated
_setmode(_fileno(in), _O_BINARY);
#endif
// a short write is not an error by itself, the exit code below decides
fwrite(stdin_data->data(), 1, stdin_data->size(), in);
}
fflush(in);
}
proc.close_stdin();
FILE * f = proc.stdout_file();
std::string output;
bool truncated = false;
if (f) {
#if defined(_WIN32)
// pipe fds default to CRT text mode: binary keeps the bytes untranslated
_setmode(_fileno(f), _O_BINARY);
#endif
// read raw bytes, not lines: the output can hold NUL and must arrive as soon as it is ready
// keep draining past the size cap, else the child blocks on a full pipe
char buf[4096];
while (fgets(buf, sizeof(buf), f) != nullptr) {
if (!truncated) {
size_t len = strlen(buf);
if (output.size() + len <= max_output) {
output.append(buf, len);
if (on_chunk && !on_chunk(console_output_to_utf8(std::string(buf, len)))) {
proc.terminate();
break;
}
} else {
size_t remaining = max_output - output.size();
output.append(buf, remaining);
if (on_chunk && remaining > 0) on_chunk(console_output_to_utf8(std::string(buf, remaining)));
truncated = true;
for (;;) {
#if defined(_WIN32)
const int n = _read(_fileno(f), buf, (unsigned) sizeof(buf));
#else
ssize_t n = read(fileno(f), buf, sizeof(buf));
while (n < 0 && errno == EINTR) {
n = read(fileno(f), buf, sizeof(buf));
}
#endif
if (n <= 0) {
break;
}
if (truncated) {
continue;
}
const size_t len = (size_t) n;
if (output.size() + len <= max_output) {
output.append(buf, len);
if (on_chunk && !on_chunk(console_output_to_utf8(std::string(buf, len)))) {
proc.terminate();
break;
}
} else {
size_t remaining = max_output - output.size();
output.append(buf, remaining);
if (on_chunk && remaining > 0) on_chunk(console_output_to_utf8(std::string(buf, remaining)));
truncated = true;
}
}
}
@@ -473,7 +513,7 @@ private:
}
};
// timeout for auxiliary isolate calls (stat/mkdir/ls/cp helpers); exec_shell_command uses its own
// timeout for auxiliary isolate calls (stat/mkdir/ls helpers); exec_shell_command uses its own
// caller-controlled timeout instead, enforced separately in run()
static constexpr int SERVER_TOOL_ISOLATE_EXEC_TIMEOUT = 15; // seconds
static constexpr size_t SERVER_TOOL_ISOLATE_READ_FILE_MAX_SIZE = 64 * 1024 * 1024; // 64 MB
@@ -524,33 +564,12 @@ public:
}
bool write_file(const std::string & path, const std::string & content) const override {
std::string abs_path = resolve(path);
std::error_code ec;
fs::path tmp_dir = fs::temp_directory_path(ec);
if (ec) return false;
static std::atomic<uint64_t> tmp_counter{0};
fs::path tmp = tmp_dir / string_format(
"llama-tools-io-isolate-%zu-%llu.tmp",
std::hash<std::thread::id>{}(std::this_thread::get_id()),
(unsigned long long) tmp_counter.fetch_add(1));
{
std::ofstream f(tmp, std::ios::binary);
if (!f) return false;
f << content;
if (!f) return false;
}
bool ok = shell_run({"sh", "-c", "mkdir -p \"$(dirname \"$1\")\"", "_", abs_path});
if (ok) {
ok = upload(tmp.string(), abs_path);
}
std::error_code rm_ec;
fs::remove(tmp, rm_ec);
return ok;
// the content travels on stdin: no argv for the far side to re-parse, no temp file on the host
auto res = run_subprocess(
build_argv({"sh", "-c", "mkdir -p \"$(dirname \"$1\")\" && cat > \"$1\"", "_", resolve(path)},
/*needs_stdin=*/true),
4096, SERVER_TOOL_ISOLATE_EXEC_TIMEOUT, nullptr, true, "", &content);
return res.exit_code == 0 && !res.timed_out;
}
list_result list_entries(const std::string & base, int max_depth, list_kind kind) const override {
@@ -612,9 +631,6 @@ protected:
// a transport that re-parses its args in a remote shell (ssh) must join `inner` with shell_quote_join()
virtual std::vector<std::string> build_argv(const std::vector<std::string> & inner, bool needs_stdin) const = 0;
// copy a host file into the isolate, `isolate_path` is absolute and its parent already exists
virtual bool upload(const std::string & host_path, const std::string & isolate_path) const = 0;
// quote `argv` into a single string that a POSIX shell re-parses into exactly `argv`
static std::string shell_quote_join(const std::vector<std::string> & argv) {
std::string out;
@@ -634,7 +650,7 @@ protected:
private:
std::string cwd;
// set the working directory in the command itself, docker's `-w` has no equivalent on every transport
// set the working directory in the command itself, no `-w` equivalent exists on every transport
// auxiliary calls do not need this, they use the absolute paths from resolve()
std::vector<std::string> with_cwd(const std::vector<std::string> & inner) const {
if (cwd.empty()) {
@@ -697,15 +713,16 @@ private:
}
};
// an already-running docker container, driven through `docker exec` and `docker cp`
class tools_io_docker : public tools_io_isolate {
// an already-running container, driven through `<engine> exec`
// docker and podman take the same verbs and the same argument order, so one class drives both
class tools_io_container : public tools_io_isolate {
public:
tools_io_docker(std::string container_id, std::string cwd = "")
: tools_io_isolate(std::move(cwd)), container_id(std::move(container_id)) {}
tools_io_container(std::string bin, std::string container_id, std::string cwd = "")
: tools_io_isolate(std::move(cwd)), bin(std::move(bin)), container_id(std::move(container_id)) {}
protected:
std::vector<std::string> build_argv(const std::vector<std::string> & inner, bool needs_stdin) const override {
std::vector<std::string> argv = {"docker", "exec"};
std::vector<std::string> argv = {bin, "exec"};
if (needs_stdin) {
argv.push_back("-i");
}
@@ -714,30 +731,118 @@ protected:
return argv;
}
bool upload(const std::string & host_path, const std::string & isolate_path) const override {
auto res = run_subprocess(
{"docker", "cp", host_path, container_id + ":" + isolate_path},
4096, SERVER_TOOL_ISOLATE_EXEC_TIMEOUT, nullptr, true);
return res.exit_code == 0 && !res.timed_out;
}
private:
std::string bin;
std::string container_id;
};
// runtime spec used by --tools-runtime and the x-tool-runtime header
// this is the only scheme for now, ssh: and podman: can be added next to it
static const std::string SERVER_TOOL_RUNTIME_DOCKER_CONTAINER = "docker-container:";
// a remote host reached over ssh
// this is remoting, not isolation: the tools can do anything the target account can do
class tools_io_ssh : public tools_io_isolate {
public:
tools_io_ssh(std::string target, std::string cwd = "")
: tools_io_isolate(std::move(cwd)), target(std::move(target)) {}
// the target can come from a client header, and ssh reads options from its argv
// a target starting with '-' would become one, e.g. -oProxyCommand=<anything> runs on the host
static bool is_valid_target(const std::string & target) {
if (target.empty() || target[0] == '-') {
return false;
}
return std::all_of(target.begin(), target.end(), [](unsigned char c) {
return std::isalnum(c) || c == '.' || c == '-' || c == '_' || c == '@';
});
}
protected:
std::vector<std::string> build_argv(const std::vector<std::string> & inner, bool needs_stdin) const override {
// the remote shell re-parses the command line, so `inner` travels as one quoted word
std::vector<std::string> argv = ssh_argv();
if (!needs_stdin) {
argv.push_back("-n");
}
argv.push_back(target);
argv.push_back(shell_quote_join(inner));
return argv;
}
private:
std::string target;
// there is no console here, so a prompt would hang the tool call
// key-based auth only, and the admin must trust the host key beforehand
static std::vector<std::string> ssh_argv() {
return {
"ssh",
"-o", "BatchMode=yes",
"-o", "PasswordAuthentication=no",
"-o", "KbdInteractiveAuthentication=no",
"-o", "StrictHostKeyChecking=yes",
};
}
};
// "<engine>:<image>" spawns a container and owns it, "<engine>-container:<id>" attaches to one
struct container_runtime_spec {
std::string bin;
std::string arg; // image name when spawning, container id when attaching
bool attach = false;
static bool parse(const std::string & spec, container_runtime_spec & out) {
// docker and podman take the same verbs, hence a single implementation
static const char * engines[] = {"docker", "podman"};
for (const char * bin : engines) {
const std::string attach_prefix = std::string(bin) + "-container:";
if (spec.rfind(attach_prefix, 0) == 0) {
out = {bin, spec.substr(attach_prefix.size()), true};
return true;
}
const std::string spawn_prefix = std::string(bin) + ":";
if (spec.rfind(spawn_prefix, 0) == 0) {
out = {bin, spec.substr(spawn_prefix.size()), false};
return true;
}
}
return false;
}
// same risk as the ssh target: an id starting with '-' would become an engine option,
// e.g. --privileged
static bool is_valid_id(const std::string & id) {
if (id.empty() || !std::isalnum((unsigned char) id[0])) {
return false;
}
return std::all_of(id.begin(), id.end(), [](unsigned char c) {
return std::isalnum(c) || c == '.' || c == '-' || c == '_';
});
}
};
// an empty runtime runs the tools on the host
static std::unique_ptr<tools_io> make_tools_io(const json & params) {
std::string cwd = json_value(params, "cwd", std::string());
std::string runtime = json_value(params, "runtime", std::string());
if (runtime.empty()) {
// an empty runtime runs the tools on the host
return std::make_unique<tools_io_basic>(cwd);
}
if (runtime.rfind(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER, 0) == 0) {
return std::make_unique<tools_io_docker>(runtime.substr(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER.size()), cwd);
container_runtime_spec container;
if (container_runtime_spec::parse(runtime, container)) {
// spawning belongs to the runtime that owns the container, a tool call only attaches
if (!container.attach) {
throw std::runtime_error("tool runtime must name a running container: " + runtime);
}
if (!container_runtime_spec::is_valid_id(container.arg)) {
throw std::runtime_error("invalid container id: " + container.arg);
}
return std::make_unique<tools_io_container>(container.bin, container.arg, cwd);
}
const std::string ssh_prefix = "ssh:";
if (runtime.rfind(ssh_prefix, 0) == 0) {
std::string target = runtime.substr(ssh_prefix.size());
if (!tools_io_ssh::is_valid_target(target)) {
throw std::runtime_error("invalid ssh target: " + target);
}
return std::make_unique<tools_io_ssh>(target, cwd);
}
// do not fall back to the host, the caller asked for an isolate
throw std::runtime_error("unknown tool runtime: " + runtime);
@@ -1769,81 +1874,82 @@ struct server_mcp_tool : server_tool {
}
};
// owns the docker container used as the sandboxed runtime for tool invocations, as configured by
// --tools-runtime. "spawned" mode starts and stops the container itself; "existing" mode just reuses
// a container id the user already has running and never stops it.
struct server_tools_docker_runtime {
server_tools_docker_runtime(const server_tools_docker_runtime &) = delete;
// resolves --tools-runtime into the isolate that every tool call runs through
// spec() returns the runtime string make_tools_io() takes, and runs once per tool call
struct server_tools_runtime {
virtual ~server_tools_runtime() = default;
virtual std::string spec() = 0;
};
explicit server_tools_docker_runtime(const std::string & spec) {
static const std::string docker_prefix = "docker:";
if (spec.rfind(docker_prefix, 0) == 0) {
spawned = true;
image = spec.substr(docker_prefix.size());
if (image.empty()) {
throw std::runtime_error("--tools-runtime docker:<image> requires an image name");
}
spawn();
} else if (spec.rfind(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER, 0) == 0) {
spawned = false;
container_id = spec.substr(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER.size());
if (container_id.empty()) {
throw std::runtime_error("--tools-runtime docker-container:<id> requires a container id");
}
} else {
// a target that already exists and needs no lifecycle
// the spec is validated once at startup, then passed straight through
struct server_tools_static_runtime : server_tools_runtime {
explicit server_tools_static_runtime(std::string spec) : runtime_spec(std::move(spec)) {}
std::string spec() override { return runtime_spec; }
private:
std::string runtime_spec;
};
// owns the container the tools run in, as set by --tools-runtime "<engine>:<image>"
// it is spawned here and stopped when the server exits
struct server_tools_container_runtime : server_tools_runtime {
server_tools_container_runtime(const server_tools_container_runtime &) = delete;
explicit server_tools_container_runtime(const std::string & spec) {
container_runtime_spec parsed;
if (!container_runtime_spec::parse(spec, parsed)) {
throw std::runtime_error("unknown --tools-runtime option: " + spec);
}
}
~server_tools_docker_runtime() {
if (spawned && !container_id.empty()) {
// closing stdin signals the container's shell (its pid 1) to exit; --rm then removes it
proc.close_stdin();
proc.join();
bin = parsed.bin;
image = parsed.arg;
if (image.empty()) {
throw std::runtime_error("--tools-runtime " + bin + ":<image> requires an image name");
}
spawn();
}
// container id to use for the next tool call; respawns a spawned container that died on its own,
// or throws if an externally-managed one is no longer reachable
std::string get_container_id() {
~server_tools_container_runtime() override {
// closing stdin signals the container's shell (its pid 1) to exit; --rm then removes it
proc.close_stdin();
proc.join();
}
// respawns a container that died on its own, so the returned spec always names a running one
std::string spec() override {
std::lock_guard<std::mutex> lock(mutex);
if (!spawned) {
if (!is_running(container_id)) {
throw std::runtime_error(string_format(
"docker container \"%s\" is no longer running, restart it to keep using tools",
container_id.c_str()));
}
return container_id;
}
if (!proc.alive()) {
SRV_WRN("docker tools runtime container \"%s\" died, respawning\n", container_id.c_str());
SRV_WRN("%s tools runtime container \"%s\" died, respawning\n", bin.c_str(), container_id.c_str());
spawn();
}
return container_id;
return bin + "-container:" + container_id;
}
private:
bool spawned = false;
std::string image; // spawned mode only
std::string bin;
std::string image;
std::string container_id;
common_subproc proc; // spawned mode only: `docker run` client that keeps the container alive
common_subproc proc; // `<engine> run` client that keeps the container alive
std::mutex mutex;
// spawns "docker run --rm -i <image> sh" and keeps its stdin open; the shell blocks reading stdin,
// spawns "<engine> run --rm -i <image> sh" and keeps its stdin open; the shell blocks reading stdin,
// so the container stays alive until we close it (see destructor) or it is killed from the outside
void spawn() {
// create() writes over the handle it is given, so the previous one is released first
proc.join();
std::error_code ec;
fs::path cidfile = fs::temp_directory_path(ec) / string_format(
"llama-tools-runtime-cid-%zu.tmp", std::hash<std::thread::id>{}(std::this_thread::get_id()));
fs::remove(cidfile, ec);
std::vector<std::string> args = {"docker", "run", "--rm", "-i", "--cidfile", cidfile.string(), image, "sh"};
std::vector<std::string> args = {bin, "run", "--rm", "-i", "--cidfile", path_to_utf8(cidfile), image, "sh"};
int options = subprocess_option_no_window
| subprocess_option_inherit_environment
| subprocess_option_search_user_path;
if (!proc.create(args, options)) {
throw std::runtime_error("failed to spawn docker container for tools runtime (image: " + image + ")");
throw std::runtime_error("failed to spawn " + bin + " container for tools runtime (image: " + image + ")");
}
std::string cid;
@@ -1855,15 +1961,10 @@ private:
fs::remove(cidfile, ec);
if (cid.empty()) {
proc.terminate();
throw std::runtime_error("timed out waiting for docker container to start (image: " + image + ")");
throw std::runtime_error("timed out waiting for " + bin + " container to start (image: " + image + ")");
}
container_id = cid;
}
static bool is_running(const std::string & id) {
auto res = run_subprocess({"docker", "inspect", "-f", "{{.State.Running}}", id}, 16, 5, nullptr, true);
return res.exit_code == 0 && !res.timed_out && res.output.rfind("true", 0) == 0;
}
};
static server_tool & find_tool(std::vector<std::unique_ptr<server_tool>> & tools, const std::string & name, bool require_stream) {
@@ -1914,11 +2015,22 @@ static std::string get_header(const std::map<std::string, std::string> & headers
server_tools::server_tools() = default;
server_tools::~server_tools() = default;
// the "<engine>:<image>" form owns a container lifecycle
// anything else names an existing target, so only its spec is validated here at startup
static std::unique_ptr<server_tools_runtime> make_tools_runtime(const std::string & spec) {
container_runtime_spec parsed;
if (container_runtime_spec::parse(spec, parsed) && !parsed.attach) {
return std::make_unique<server_tools_container_runtime>(spec);
}
make_tools_io({{"runtime", spec}}); // nothing to own, just reject a bad spec now
return std::make_unique<server_tools_static_runtime>(spec);
}
void server_tools::setup(const std::vector<std::string> & enabled_tools,
server_mcp & mcp_mgr,
const std::string & tools_runtime) {
if (!tools_runtime.empty()) {
docker_runtime = std::make_unique<server_tools_docker_runtime>(tools_runtime);
runtime = make_tools_runtime(tools_runtime);
}
if (!enabled_tools.empty()) {
@@ -2016,11 +2128,11 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools,
if (params.contains("runtime")) {
params.erase("runtime");
}
auto runtime = get_header(req.headers, "x-tool-runtime");
if (!runtime.empty()) {
params["runtime"] = runtime;
} else if (docker_runtime) {
params["runtime"] = SERVER_TOOL_RUNTIME_DOCKER_CONTAINER + docker_runtime->get_container_id();
auto runtime_header = get_header(req.headers, "x-tool-runtime");
if (!runtime_header.empty()) {
params["runtime"] = runtime_header;
} else if (runtime) {
params["runtime"] = runtime->spec();
}
server_tool & tool = find_tool(tools, tool_name, stream);
tools/server/server-tools.h
+3 -3
View File
@@ -31,7 +31,7 @@ struct server_tool {
json to_json() const;
};
struct server_tools_docker_runtime; // impl detail, defined in server-tools.cpp
struct server_tools_runtime; // impl detail, defined in server-tools.cpp
struct server_tools {
std::vector<std::unique_ptr<server_tool>> tools;
@@ -40,8 +40,8 @@ struct server_tools {
server_response queue_res;
std::atomic<int> res_id{0};
// set when --tools-runtime is configured; owns the docker container used to run tools, if any
std::unique_ptr<server_tools_docker_runtime> docker_runtime;
// set when --tools-runtime is configured; routes every tool call through an isolate
std::unique_ptr<server_tools_runtime> runtime;
void setup(const std::vector<std::string> & enabled_tools,
server_mcp & mcp_mgr,
tools/server/server.cpp
+1 -1
View File
@@ -89,7 +89,7 @@ int llama_server(int argc, char ** argv) {
std::setlocale(LC_NUMERIC, "C");
#ifndef _WIN32
// Ignore SIGPIPE so the server does not crash if an MCP child exits while we are writing to its stdin
// Ignore SIGPIPE so the server does not crash if a child (MCP server, tools runtime) exits while we are writing to its stdin
signal(SIGPIPE, SIG_IGN);
#endif
tools/server/tests/unit/test_speculative.py
+16 -15
View File
@@ -25,33 +25,34 @@ def fixture_create_server():
def test_with_and_without_draft():
global server
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.8,
"top_k": 40,
"seed": 4242,
"n_predict": 16,
"return_tokens": True,
}
server.model_draft = None # disable draft model
server.spec_type = None
server.backend_sampling = True
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"n_predict": 16,
})
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
content_no_draft = res.body["content"]
tokens_no_draft = res.body["tokens"]
server.stop()
# create new server with draft model
create_server()
server.backend_sampling = True
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"n_predict": 16,
})
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
content_draft = res.body["content"]
tokens_draft = res.body["tokens"]
assert content_no_draft == content_draft
assert tokens_no_draft == tokens_draft
def test_different_draft_min_draft_max():
tools/server/tests/unit/test_tools_builtin.py
+59 -26
View File
@@ -14,7 +14,7 @@ PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..
GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search"
# image the container runtime tests run their shell in
DOCKER_IMAGE = "busybox"
CONTAINER_IMAGE = "busybox"
@pytest.fixture(autouse=True)
@@ -151,54 +151,59 @@ def test_tools_builtin_cwd_header():
os.remove(marker_path)
def _docker_unavailable_reason() -> str | None:
"""None if docker can run the image these tests use, otherwise the reason it can't."""
docker_bin = shutil.which("docker")
if docker_bin is None:
return "docker is not installed"
def _container_engine_unavailable_reason(engine: str) -> str | None:
"""None if `engine` can run the image these tests use, otherwise the reason it can't."""
engine_bin = shutil.which(engine)
if engine_bin is None:
return f"{engine} is not installed"
try:
# a daemon that answers `docker info` still cannot run a linux image when it serves
# windows containers, so probe the image itself, which also pulls it before the tests
subprocess.run([docker_bin, "run", "--rm", DOCKER_IMAGE, "true"], capture_output=True, timeout=60, check=True)
# a daemon that answers `info` still cannot run a linux image when it serves windows
# containers, so probe the image itself, which also pulls it before the tests
subprocess.run([engine_bin, "run", "--rm", CONTAINER_IMAGE, "true"], capture_output=True, timeout=60, check=True)
except Exception as e:
return f"docker cannot run {DOCKER_IMAGE}: {e}"
return f"{engine} cannot run {CONTAINER_IMAGE}: {e}"
return None
@pytest.fixture
def docker_container():
reason = _docker_unavailable_reason()
@pytest.fixture(params=["docker", "podman"])
def container_engine(request):
engine = request.param
reason = _container_engine_unavailable_reason(engine)
if reason is not None:
pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type]
return engine
@pytest.fixture
def container_id(container_engine: str):
proc = subprocess.run(
["docker", "run", "-d", "--rm", DOCKER_IMAGE, "sleep", "300"],
[container_engine, "run", "-d", "--rm", CONTAINER_IMAGE, "sleep", "300"],
capture_output=True, text=True,
)
if proc.returncode != 0:
pytest.skip(f"failed to start docker container: {proc.stderr.strip()}") # ty: ignore[too-many-positional-arguments, invalid-argument-type]
pytest.skip(f"failed to start {container_engine} container: {proc.stderr.strip()}") # ty: ignore[too-many-positional-arguments, invalid-argument-type]
container_id = proc.stdout.strip()
cid = proc.stdout.strip()
try:
yield container_id
yield cid
finally:
subprocess.run(["docker", "rm", "-f", container_id], capture_output=True)
subprocess.run([container_engine, "rm", "-f", cid], capture_output=True)
def test_tools_builtin_runtime_header(docker_container: str):
def test_tools_builtin_runtime_header(container_engine: str, container_id: str):
global server
server.start()
headers = {"x-tool-runtime": f"docker-container:{docker_container}", "x-tool-cwd": "/tmp"}
headers = {"x-tool-runtime": f"{container_engine}-container:{container_id}", "x-tool-cwd": "/tmp"}
write_res = call_tool("write_file", {"path": "test.log", "content": "hello docker\n"}, headers=headers)
write_res = call_tool("write_file", {"path": "test.log", "content": "hello container\n"}, headers=headers)
assert write_res["result"] == "file written successfully"
read_res = call_tool("read_file", {"path": "test.log"}, headers=headers)
assert read_res["plain_text_response"] == "hello docker\n"
assert read_res["plain_text_response"] == "hello container\n"
exec_res = call_tool("exec_shell_command", {"command": "cat test.log"}, headers=headers)
assert "hello docker" in exec_res["plain_text_response"]
assert "hello container" in exec_res["plain_text_response"]
def test_tools_builtin_runtime_header_unknown_scheme():
@@ -208,18 +213,46 @@ def test_tools_builtin_runtime_header_unknown_scheme():
# an unknown runtime must fail, never silently fall back to running on the host
res = server.make_request("POST", "/tools",
data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
headers={"x-tool-runtime": "ssh:example.com"})
headers={"x-tool-runtime": "fake:does-not-exist"})
assert res.status_code == 500, res.body
assert "unknown tool runtime" in str(res.body)
def test_tools_builtin_runtime_header_rejects_ssh_option_injection():
global server
server.start()
# ssh reads options from its argv, so a target starting with '-' must be rejected
res = server.make_request("POST", "/tools",
data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
headers={"x-tool-runtime": "ssh:-oProxyCommand=touch /tmp/pwned"})
assert res.status_code == 500, res.body
assert "invalid ssh target" in str(res.body)
@pytest.mark.parametrize("engine", ["docker", "podman"])
def test_tools_builtin_runtime_header_rejects_container_option_injection(engine: str):
global server
server.start()
# the container id lands on the `<engine> exec` command line, so an id that looks
# like an option must be rejected
res = server.make_request("POST", "/tools",
data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
headers={"x-tool-runtime": f"{engine}-container:--privileged"})
assert res.status_code == 500, res.body
assert "invalid container id" in str(res.body)
def test_tools_builtin_docker_runtime_cleans_up_spawned_container():
reason = _docker_unavailable_reason()
# docker-only: this reads the container hostname to get the spawned id, which only docker
# sets to the short id. podman is covered by the attach path above
reason = _container_engine_unavailable_reason("docker")
if reason is not None:
pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type]
global server
server.server_tools_runtime = f"docker:{DOCKER_IMAGE}"
server.server_tools_runtime = f"docker:{CONTAINER_IMAGE}"
server.start()
# exec_shell_command runs inside the container spawned for --tools-runtime; docker sets
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContenteditable.svelte
+1 -1
View File
@@ -797,7 +797,7 @@
data-placeholder={placeholder}
tabindex={disabled ? -1 : 0}
class={[
'chat-form-contenteditable text-md min-h-12 w-full whitespace-pre-wrap wrap-break-word border-0 bg-transparent p-0 leading-6 outline-none focus-visible:ring-0 focus-visible:ring-offset-0',
'chat-form-contenteditable text-md min-h-12 w-full overflow-y-auto whitespace-pre-wrap wrap-break-word border-0 bg-transparent p-0 leading-6 outline-none focus-visible:ring-0 focus-visible:ring-offset-0',
disabled && 'cursor-not-allowed'
]}
style="max-height: var(--max-message-height);"
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte
+1 -1
View File
@@ -164,7 +164,7 @@
? `max-height: ${MAX_HEIGHT}px;`
: 'max-height: none;'}
>
{#if !currentConfig.renderContentAsRawText}
{#if currentConfig.renderUserContentAsMarkdown}
<div bind:this={messageElement} class={isExpanded ? 'cursor-text' : ''}>
<MarkdownContent class="markdown-system-content" content={message.content} />
</div>
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserBubble.svelte
+5 -5
View File
@@ -1,5 +1,5 @@
<script lang="ts">
import { ChatAttachmentsList, MarkdownContent } from '$lib/components/app';
import { ChatAttachmentsList, MarkdownContent, MentionText } from '$lib/components/app';
import { Card } from '$lib/components/ui/card';
import { config } from '$lib/stores/settings.svelte';
import type { DatabaseMessageExtra } from '$lib/types/database';
@@ -64,14 +64,14 @@
data-multiline={isMultiline ? '' : undefined}
style="{maxHeightStyle} overflow-wrap: anywhere; word-break: break-word;"
>
{#if renderMarkdown && !currentConfig.renderContentAsRawText}
{#if renderMarkdown && currentConfig.renderUserContentAsMarkdown}
<div bind:this={messageElement}>
<MarkdownContent class="markdown-user-content" {content} />
</div>
{:else}
<span bind:this={messageElement} class="text-md whitespace-pre-wrap">
{content}
</span>
<span bind:this={messageElement} class="text-md whitespace-pre-wrap"
><MentionText {content} /></span
>
{/if}
</Card>
{/if}
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte
+1 -1
View File
@@ -140,7 +140,7 @@
class:is-streaming={isPending}
onscroll={handleScrollEvent}
>
{#if !currentConfig.renderContentAsRawText}
{#if currentConfig.renderThinkingAsMarkdown}
<MarkdownContent content={section.content} class="text-muted-foreground" {attachments} />
{:else}
<div
tools/ui/src/lib/components/app/content/MentionBadge.svelte
+36
View File
@@ -0,0 +1,36 @@
<script lang="ts">
import { SETTINGS_KEYS } from '$lib/constants';
import { settingsStore } from '$lib/stores/settings.svelte';
import { toolsStore } from '$lib/stores/tools.svelte';
import {
getMentionBadgeIconPaths,
getMentionBadgeLabel,
MENTION_BADGE_CLASSNAME,
MENTION_BADGE_ICON_CLASSNAME,
MENTION_BADGE_SVG_ATTRIBUTES
} from '$lib/utils';
interface Props {
name: string;
path: string;
}
let { name, path }: Props = $props();
let showFullPath = $derived(
settingsStore.getConfig(SETTINGS_KEYS.SHOW_FULL_PATH_IN_MENTIONS) as boolean
);
let label = $derived(getMentionBadgeLabel(name, path, showFullPath, toolsStore.serverHome));
</script>
<!-- The chip is a flex container, so template whitespace between its
children collapses away and the icon keeps its `gap-1` spacing. -->
<span class={MENTION_BADGE_CLASSNAME} title={path}>
<svg {...MENTION_BADGE_SVG_ATTRIBUTES} class={MENTION_BADGE_ICON_CLASSNAME}>
{#each getMentionBadgeIconPaths(path) as d (d)}
<path {d} />
{/each}
</svg>
<span class="shrink-0 truncate">{label}</span>
</span>
tools/ui/src/lib/components/app/content/MentionText.svelte
+17
View File
@@ -0,0 +1,17 @@
<script lang="ts">
import MentionBadge from './MentionBadge.svelte';
import { splitMentionSegments } from '$lib/utils';
interface Props {
content: string;
}
let { content }: Props = $props();
let segments = $derived(splitMentionSegments(content));
</script>
<!-- Segments sit in a `whitespace-pre-wrap` parent, so the markup stays
glued: any newline between the tags below would print as a space. -->
<!-- prettier-ignore -->
{#each segments as segment, index (index)}{#if segment.mention}<MentionBadge name={segment.mention.name} path={segment.mention.path} />{:else}{segment.text}{/if}{/each}
tools/ui/src/lib/components/app/content/index.ts
+14
View File
@@ -31,6 +31,20 @@
*/
export { default as MarkdownContent } from './MarkdownContent/MarkdownContent.svelte';
/**
* **MentionText** - Plain text with file mention badges
*
* Renders a message verbatim, turning only `[name](file://path)` links
* into the same badge chips the markdown path draws. Nothing else is
* interpreted, so pasted code keeps its `#` comments and underscores.
*
* @example
* ```svelte
* <span class="whitespace-pre-wrap"><MentionText content={message.content} /></span>
* ```
*/
export { default as MentionText } from './MentionText.svelte';
/**
* **SyntaxHighlightedCode** - Code syntax highlighting
*
tools/ui/src/lib/constants/mention-badge.ts
+3
View File
@@ -10,6 +10,9 @@ export const MENTION_BADGE_CLASSNAME =
export const MENTION_BADGE_ICON_CLASSNAME = 'h-3 w-3 shrink-0';
/** Regex flag that makes the mention scanner walk every link in a message instead of the first. */
export const MENTION_LINK_SCAN_FLAGS = 'g';
/**
* SVG attributes shared by the DOM-built and hast-built badge icons.
* The tokenizer applies them via `setAttribute`, the rehype plugin
tools/ui/src/lib/constants/path-display.ts
+3
View File
@@ -12,6 +12,9 @@ import { UrlProtocol } from '$lib/enums';
export const CWD_CHANGED_PREFIX = 'Set working directory to ';
export const CWD_CLEARED_TEXT = 'Working directory cleared';
/** Trailing separator that marks a path as a directory. */
export const DIRECTORY_PATH_SUFFIX = '/';
export const HOME_TILDE = '~';
export const HOME_TILDE_PREFIX = '~/'; // tilde plus path separator
tools/ui/src/lib/constants/settings-keys.ts
+2 -1
View File
@@ -41,7 +41,8 @@ export const SETTINGS_KEYS = {
// Performance
PRE_ENCODE_CONVERSATION: 'preEncodeConversation',
PRESENCE_PENALTY: 'presence_penalty',
RENDER_CONTENT_AS_RAW_TEXT: 'renderContentAsRawText',
RENDER_THINKING_AS_MARKDOWN: 'renderThinkingAsMarkdown',
RENDER_USER_CONTENT_AS_MARKDOWN: 'renderUserContentAsMarkdown',
// Penalties
REPEAT_LAST_N: 'repeat_last_n',
REPEAT_PENALTY: 'repeat_penalty',
tools/ui/src/lib/constants/settings-registry.ts
+12 -4
View File
@@ -230,10 +230,18 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
type: SettingsFieldType.CHECKBOX
},
{
defaultValue: false,
help: 'Display user, system and thinking content as plain text instead of formatted Markdown. Markdown is the default so that @-mention badges render in sent messages.',
key: SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT,
label: 'Render content as raw text',
defaultValue: true,
help: 'Render user messages using markdown formatting in the chat. Turn this off to keep a message exactly as typed; @-mention badges show either way.',
key: SETTINGS_KEYS.RENDER_USER_CONTENT_AS_MARKDOWN,
label: 'Render user content as Markdown',
section: SETTINGS_SECTION_SLUGS.DISPLAY,
type: SettingsFieldType.CHECKBOX
},
{
defaultValue: true,
help: 'Render the reasoning/thinking block content as formatted Markdown instead of plain text.',
key: SETTINGS_KEYS.RENDER_THINKING_AS_MARKDOWN,
label: 'Render thinking as Markdown',
section: SETTINGS_SECTION_SLUGS.DISPLAY,
type: SettingsFieldType.CHECKBOX
},
tools/ui/src/lib/services/migration.service.ts
+33 -1
View File
@@ -629,6 +629,37 @@ const configTypesMigration: Migration = {
console.log(`[Migration] Config types: coerced string booleans (changed=${changed})`);
}
};
const RENDER_KEYS_MIGRATION_ID = 'render-keys-unfold-v1';
const LEGACY_RENDER_RAW_TEXT_KEY = 'renderContentAsRawText';
const renderKeysMigration: Migration = {
description: 'Unfold the single raw text render toggle onto the per-surface render keys',
id: RENDER_KEYS_MIGRATION_ID,
async run(): Promise<void> {
const configRaw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY);
if (configRaw === null) return;
const config = JSON.parse(configRaw);
if (!(LEGACY_RENDER_RAW_TEXT_KEY in config)) return;
// The toggle carried user content and thinking at once and cannot say which surface
// was chosen, so it only restores the user key and thinking keeps its own default.
if (!(SETTINGS_KEYS.RENDER_USER_CONTENT_AS_MARKDOWN in config)) {
config[SETTINGS_KEYS.RENDER_USER_CONTENT_AS_MARKDOWN] =
config[LEGACY_RENDER_RAW_TEXT_KEY] !== true;
}
// Dropped rather than preserved: the two render keys and the toggle describe the same
// surfaces, so leaving it behind would let a stale value fight the restored one.
delete config[LEGACY_RENDER_RAW_TEXT_KEY];
localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config));
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG)
console.log('[Migration] Render keys: unfolded the raw text toggle');
}
};
const MCP_DEFAULT_OVERRIDES_LEGACY_KEY = `${STORAGE_APP_NAME}.mcpDefaultServerOverrides`;
const MCP_DEFAULT_OVERRIDES_MERGE_MIGRATION_ID = 'mcp-default-overrides-merge-v1';
/**
@@ -722,7 +753,8 @@ const migrations: Migration[] = [
customJsonKeyMigration,
mcpDefaultEnabledMigration,
mcpDefaultOverridesMergeMigration,
configTypesMigration
configTypesMigration,
renderKeysMigration
];
export const MigrationService = {
tools/ui/src/lib/stores/settings.svelte.ts
-26
View File
@@ -135,32 +135,6 @@ class SettingsStore {
...savedVal
};
// Migrate the legacy render keys into `renderContentAsRawText`
// (inverted semantics: the old keys opted INTO markdown). Any
// explicit raw-text preference wins when the legacy keys disagree.
const LEGACY_MARKDOWN_KEYS = ['renderUserContentAsMarkdown', 'renderThinkingAsMarkdown'];
const LEGACY_RAW_TEXT_KEY = 'renderUserContentAsRawText'; // this branch's intermediate key
const legacyKeys = [...LEGACY_MARKDOWN_KEYS, LEGACY_RAW_TEXT_KEY].filter(
(key) => key in savedVal
);
if (legacyKeys.length > 0) {
if (!(SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT in savedVal)) {
if (LEGACY_RAW_TEXT_KEY in savedVal) {
this.config[SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT] = savedVal[LEGACY_RAW_TEXT_KEY];
} else {
this.config[SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT] = LEGACY_MARKDOWN_KEYS.filter(
(key) => key in savedVal
).some((key) => savedVal[key] === false);
}
}
for (const key of legacyKeys) {
delete (this.config as Record<string, unknown>)[key];
}
this.saveConfig();
}
// Default sendOnEnter to false on mobile when the user has no saved preference
if (!(SETTINGS_KEYS.SEND_ON_ENTER in savedVal)) {
if (isMobile.current) {
tools/ui/src/lib/utils/index.ts
+1
View File
@@ -240,6 +240,7 @@ export {
MENTION_BADGE_FOLDER_ICON_PATHS,
getMentionBadgeIconPaths,
getMentionBadgeLabel,
splitMentionSegments,
buildMentionInsertion
} from './mention-badge';
tools/ui/src/lib/utils/mention-badge.ts
+41 -3
View File
@@ -1,8 +1,9 @@
import { abbreviateHome, lastPathSegment } from './path-display';
import { FILE_URI_PREFIX } from '$lib/constants';
import { DIRECTORY_PATH_SUFFIX, FILE_URI_PREFIX } from '$lib/constants';
import {
MENTION_BADGE_FILE_ICON_PATHS,
MENTION_BADGE_FOLDER_ICON_PATHS
MENTION_BADGE_FOLDER_ICON_PATHS,
MENTION_LINK_SCAN_FLAGS
} from '$lib/constants/mention-badge';
import { FileMentionEntryType } from '$lib/enums';
import type { FileMentionEntry } from '$lib/types';
@@ -48,8 +49,45 @@ export function decodeFileLinkPath(path: string): string {
}
}
export interface MentionTextSegment {
text: string;
mention: { name: string; path: string } | null;
}
/**
* Split raw text into plain runs and `[name](file://path)` mentions.
* The raw-text renderers walk these segments to draw badges without
* handing the message to the markdown parser, so a `#` stays a `#`.
*/
export function splitMentionSegments(value: string): MentionTextSegment[] {
const linkRe = fileMentionLinkRe(MENTION_LINK_SCAN_FLAGS);
const segments: MentionTextSegment[] = [];
let cursor = 0;
let match: RegExpExecArray | null;
while ((match = linkRe.exec(value)) !== null) {
if (match.index > cursor) {
segments.push({ mention: null, text: value.slice(cursor, match.index) });
}
segments.push({
mention: { name: match[1], path: decodeFileLinkPath(match[2]) },
text: match[0]
});
cursor = match.index + match[0].length;
}
if (cursor < value.length) segments.push({ mention: null, text: value.slice(cursor) });
return segments;
}
export function getMentionBadgeIconPaths(path: string): readonly string[] {
return path.endsWith('/') ? MENTION_BADGE_FOLDER_ICON_PATHS : MENTION_BADGE_FILE_ICON_PATHS;
return path.endsWith(DIRECTORY_PATH_SUFFIX)
? MENTION_BADGE_FOLDER_ICON_PATHS
: MENTION_BADGE_FILE_ICON_PATHS;
}
export function getMentionBadgeLabel(
tools/ui/tests/client/settings-raw-text-migration.svelte.test.ts
-66
View File
@@ -1,66 +0,0 @@
// Guards the legacy render-key migration: `renderUserContentAsMarkdown`
// and `renderThinkingAsMarkdown` (opt-INTO markdown) fold into the single
// `renderContentAsRawText` setting, with any explicit raw-text preference
// winning when the legacy keys disagree. Legacy keys are removed from the
// persisted config so they do not stay orphaned in localStorage.
import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants/storage';
import { config, settingsStore } from '$lib/stores/settings.svelte';
import { beforeEach, describe, expect, it } from 'vitest';
function seedConfig(stored: Record<string, unknown>) {
localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(stored));
settingsStore.initialize();
}
function persisted(): Record<string, unknown> {
return JSON.parse(localStorage.getItem(CONFIG_LOCALSTORAGE_KEY) ?? '{}');
}
describe('renderContentAsRawText migration', () => {
beforeEach(() => {
localStorage.removeItem(CONFIG_LOCALSTORAGE_KEY);
settingsStore.initialize();
});
it('maps renderUserContentAsMarkdown=false to raw text', () => {
seedConfig({ renderUserContentAsMarkdown: false });
expect(config().renderContentAsRawText).toBe(true);
});
it('maps renderUserContentAsMarkdown=true to markdown', () => {
seedConfig({ renderUserContentAsMarkdown: true });
expect(config().renderContentAsRawText).toBe(false);
});
it('maps renderThinkingAsMarkdown=false to raw text', () => {
seedConfig({ renderThinkingAsMarkdown: false });
expect(config().renderContentAsRawText).toBe(true);
});
it('lets any explicit raw-text preference win when the legacy keys disagree', () => {
seedConfig({ renderThinkingAsMarkdown: false, renderUserContentAsMarkdown: true });
expect(config().renderContentAsRawText).toBe(true);
});
it('honors the intermediate renderUserContentAsRawText key from the PR branch', () => {
seedConfig({ renderUserContentAsRawText: true });
expect(config().renderContentAsRawText).toBe(true);
});
it('keeps an already-migrated value and cleans up the legacy keys', () => {
seedConfig({ renderContentAsRawText: false, renderUserContentAsMarkdown: false });
expect(config().renderContentAsRawText).toBe(false);
const stored = persisted();
expect(stored.renderUserContentAsMarkdown).toBeUndefined();
expect(stored.renderThinkingAsMarkdown).toBeUndefined();
expect(stored.renderUserContentAsRawText).toBeUndefined();
});
it('defaults to markdown when no legacy key exists', () => {
seedConfig({});
expect(config().renderContentAsRawText).toBe(false);
});
});
tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts
+64
View File
@@ -0,0 +1,64 @@
// Guards the unfolding of `renderContentAsRawText` back onto the two
// per-surface render keys. The single toggle carried user content and
// thinking at once, so only the user key is restored from it and thinking
// returns to its own default. The toggle is removed from the persisted
// config so it does not stay orphaned in localStorage.
import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants/storage';
import { MigrationService } from '$lib/services/migration.service';
import { config, settingsStore } from '$lib/stores/settings.svelte';
import { beforeEach, describe, expect, it } from 'vitest';
const RENDER_KEYS_MIGRATION_ID = 'render-keys-unfold-v1';
async function seedConfig(stored: Record<string, unknown>) {
localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(stored));
const migration = MigrationService.getMigrations().find((m) => m.id === RENDER_KEYS_MIGRATION_ID);
await migration?.run();
settingsStore.initialize();
}
function persisted(): Record<string, unknown> {
return JSON.parse(localStorage.getItem(CONFIG_LOCALSTORAGE_KEY) ?? '{}');
}
describe('renderContentAsRawText unfolding', () => {
beforeEach(() => {
localStorage.removeItem(CONFIG_LOCALSTORAGE_KEY);
MigrationService.resetState();
settingsStore.initialize();
});
it('maps raw text to user content as plain text', async () => {
await seedConfig({ renderContentAsRawText: true });
expect(config().renderUserContentAsMarkdown).toBe(false);
});
it('maps markdown to user content as markdown', async () => {
await seedConfig({ renderContentAsRawText: false });
expect(config().renderUserContentAsMarkdown).toBe(true);
});
it('leaves thinking on its own default', async () => {
await seedConfig({ renderContentAsRawText: true });
expect(config().renderThinkingAsMarkdown).toBe(true);
});
it('keeps an explicit user preference over the toggle', async () => {
await seedConfig({ renderContentAsRawText: true, renderUserContentAsMarkdown: true });
expect(config().renderUserContentAsMarkdown).toBe(true);
});
it('drops the toggle from the persisted config', async () => {
await seedConfig({ renderContentAsRawText: true });
expect(persisted().renderContentAsRawText).toBeUndefined();
});
it('leaves both surfaces on markdown when nothing is stored', async () => {
await seedConfig({});
expect(config().renderUserContentAsMarkdown).toBe(true);
expect(config().renderThinkingAsMarkdown).toBe(true);
});
});
tools/ui/tests/unit/mention-segments.test.ts
+48
View File
@@ -0,0 +1,48 @@
import { splitMentionSegments } from '$lib/utils/mention-badge';
import { describe, expect, it } from 'vitest';
describe('splitMentionSegments', () => {
it('returns a single plain run when there is no mention', () => {
const segments = splitMentionSegments('# not a heading here');
expect(segments).toEqual([{ mention: null, text: '# not a heading here' }]);
});
it('splits text around a mention', () => {
const segments = splitMentionSegments('look at [main.c](file:///src/main.c) please');
expect(segments.map((segment) => segment.text)).toEqual([
'look at ',
'[main.c](file:///src/main.c)',
' please'
]);
expect(segments[1].mention).toEqual({ name: 'main.c', path: '/src/main.c' });
});
it('decodes percent-encoded paths', () => {
const segments = splitMentionSegments('[a b.txt](file:///tmp/a%20b.txt)');
expect(segments[0].mention?.path).toBe('/tmp/a b.txt');
});
it('keeps the directory marker so the folder icon is picked', () => {
const segments = splitMentionSegments('[src](file:///repo/src/)');
expect(segments[0].mention?.path).toBe('/repo/src/');
});
it('handles adjacent mentions with no text between them', () => {
const segments = splitMentionSegments('[a](file:///a)[b](file:///b)');
expect(segments).toHaveLength(2);
expect(segments.every((segment) => segment.mention !== null)).toBe(true);
});
it('preserves the exact source when segments are joined back', () => {
const source = 'see [a](file:///a) and [b](file:///b/) done';
expect(splitMentionSegments(source).reduce((acc, segment) => acc + segment.text, '')).toBe(
source
);
});
});