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

Compare commits

merge into: Lumpiasty/llama.cpp:fable5/prefetch-experts
Branches Tags
Lumpiasty/llama.cpp:fable5/prefetch-experts Lumpiasty/llama.cpp:fable5/vram-ondemand-rebased-20260726 Lumpiasty/llama.cpp:fable5/vram-ondemand Lumpiasty/llama.cpp:gemini
...
pull from: Lumpiasty/llama.cpp:fable5/vram-ondemand-rebased-20260726
Branches Tags
Lumpiasty/llama.cpp:fable5/vram-ondemand-rebased-20260726 Lumpiasty/llama.cpp:fable5/vram-ondemand Lumpiasty/llama.cpp:fable5/prefetch-experts Lumpiasty/llama.cpp:gemini
62 Commits
fable5/prefetch-experts ... fable5/vram-ondemand-rebased-20260726
Author SHA1 Message Date
Lumpiasty 80b91cbf76 server: on-demand mmproj - free encoder VRAM on the text path 
The vision encoder sits idle in VRAM on every text request. With
LLAMA_MMPROJ_ONDEMAND=1 the server keeps the encoder in RAM (released
after load) and brings it into VRAM just-in-time before mtmd_batch_encode,
releasing it again afterwards - so the encoder's VRAM is free for KV /
expert cache on the common text-only path.

When the encoder does not fit (that VRAM has been claimed), evict the LLM
backbone weights first: they are not needed while the encoder runs (the
decode that uses them happens afterwards), their host shadow is read-only
(cheap free, no D2H), and the KV / prompt cache is left untouched. Order
on the way back matters: free the encoder before re-uploading the weights
so the peak stays within VRAM.

Also wire the encoder release/restore into the VRAM arbiter: vram_go_cold
releases it; the just-in-time encode-path restore replaces the previous
restore in vram_ensure_warm.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty da76fcb326 mtmd: release/restore the encoder weights from VRAM on demand 
Add clip_release_device/clip_restore_device (and mtmd_release_device/
mtmd_restore_device wrappers over the vision + audio contexts) that free
the multimodal encoder's device weight buffer to a read-only host shadow
and rebuild it on demand, using the same shadow/free/reallocate pattern
as llama_model weights. No-op for a CPU-backed encoder. This lets the
server drop the ~hundreds-of-MiB vision encoder from VRAM when it is not
encoding an image.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 6d476d6ccc llama: free backend compute scratch on cold; guard memory_breakdown 
On device release with KV eviction, also call ggml_backend_free_scratch()
on each backend so a cold model drops its Vulkan compute preallocations
(reallocated lazily on the next compute via restore).

Also guard llama_context::memory_breakdown() against a freed scheduler:
evict_kv release resets sched to null, so a cold model that is then torn
down (e.g. terminated by the process manager) hit GGML_ASSERT(sched) in
ggml_backend_sched_get_buffer_type. Skip the compute-buffer accounting
when sched is null.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty a8dc3ebddb ggml: add ggml_backend_free_scratch to drop Vulkan compute prealloc 
Add an optional backend interface method free_scratch (with a public
ggml_backend_free_scratch wrapper) that frees transient/scratch device
memory a backend holds outside of any allocated buffer, keeping the
backend usable - the scratch is reallocated lazily on the next compute.

Implement it for the Vulkan backend (ggml_backend_vk_free_scratch): free
the prealloc_x/y/split_k/add_rms_partials and sync_staging device buffers
and reset their sizes, so an idle/cold model does not hold the vision or
matmul compute preallocations in VRAM. All other backends leave the hook
null (no-op).

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 5ddcf40fc1 llama: evict recurrent/SSM state on device release 
The recurrent (SSM/conv) state of hybrid models (e.g. Qwen3.5) was left
resident when a model's device buffers were released for on-demand VRAM
sharing - llama_memory_recurrent::release_device_buffers() was a no-op
default. Implement it (and restore_device_buffers) with the same
capture-host-shadow / free / reallocate pattern as llama_kv_cache, so
llama_memory_hybrid now evicts both its attention KV and its recurrent
state. The state is read-write, so its shadow is recaptured on every
release.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty fd7cbd5f4f server: also free the compute-graph scheduler on KV eviction 
Under LLAMA_SLEEP_EVICT_KV a cold model still held the scheduler's worst-case
compute buffer (hundreds of MiB, e.g. ~700 MiB for gemma-26B) plus the resident
experts. release_device(evict_kv) now also frees the sched (sched.reset() +
sched_need_reserve), and restore_device() rebuilds it via sched_reserve(), so a
fully-evicted model holds essentially no VRAM (gemma cold: 7147 -> 51 MiB).
Trades a heavier re-warm (weights+KV H2D + sched reserve) for the freed VRAM;
weights-only mode is unchanged (fast switch, keeps KV+compute).

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 4d576de9dd server: restore weights/KV on wake, not at decode, so KV eviction is safe 
With LLAMA_SLEEP_EVICT_KV the KV cache device buffers are freed when a model
goes cold. update_slots() touches the KV before the decode-time vram_ensure_warm
(e.g. SWA models create a checkpoint that reads the KV via ggml_backend_tensor_get),
so the KV must already be resident by then. Move the restore to the sleep-wake
handler (handle_sleeping_state(false)), which runs before update_slots, fixing a
GGML_ASSERT(buffer) crash on iSWA models (gemma) when KV eviction is enabled.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty f707a2430d llama: guard buffer iteration against released device buffers 
release_device_weights()/release_device_buffers() leave a null buffer in
ctxs_bufs while the device memory is released for on-demand VRAM sharing.
The memory-breakdown / total_size / clear paths iterated these and called
ggml_backend_buffer_get_size()/get_type()/clear() on the null buffer, tripping
GGML_ASSERT(buffer) and aborting on shutdown of a cold model. Skip null buffers
in llama_model::memory_breakdown and llama_kv_cache::{memory_breakdown,
total_size,clear}.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 2af8f5ae6e server: coordinate model load with the VRAM arbiter to avoid load-time OOM 
Before uploading a model's weights, acquire the shared VRAM token (ring the
doorbell so any resident model releases first). Previously load uploaded weights
to VRAM before the arbiter was active, so loading a large model while another
(e.g. the warm 4B task model) held VRAM could exceed the budget and OOM.

- add vram_arena_open() (idempotent flock/doorbell setup) and
  vram_acquire_for_load(), called from load_model() before
  common_init_from_params().
- a coordinated load now stays warm holding the token and serves its first
  request without a re-warm (drop the init-time go_cold cold-start).

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 189856ff9b server: keep the sleep-wake path active when the VRAM arbiter is enabled 
The request handler only calls wait_until_no_sleep() (which wakes a server out
of its sleeping state) when sleep_idle_seconds >= 0. But the VRAM arbiter's
cross-process doorbell can put a server to sleep even when idle-sleep is
disabled, so without this a doorbell-slept server would never wake and requests
to it would hang until timeout.

Do not bypass the wake path when LLAMA_SLEEP_VRAM_ONLY is set, so the arbiter no
longer depends on --sleep-idle-seconds being configured.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 62873a34bb server: optionally evict the KV cache too (Phase 2 of VRAM sharing) 
Extend on-demand device residency to the KV cache so that when a model's KV
plus another model would not fit in VRAM, the KV can also be evicted to a host
shadow (D2H on release, H2D on restore) instead of only the weights.

- llama_memory_i: add release_device_buffers()/restore_device_buffers()
  (default no-op). Implemented in llama_kv_cache (D2H shadow of the live
  ctxs_bufs, freed and reallocated like the weights); llama_memory_hybrid and
  llama_kv_cache_iswa delegate to their child caches.
- llama_context::release_device(evict_kv): also evict the memory's device
  buffers when requested; restore_device() rebuilds them. Public API
  llama_context_release_device gains an evict_kv flag.
- server: LLAMA_SLEEP_EVICT_KV=1 enables it. Off by default (weights-only),
  since the KV shadow adds a D2H/H2D copy of the live cache each cycle.

Validated on RX 580 (Vulkan), 4B @ 32k ctx: weights-only cold VRAM 1750 MB
(KV stays); weights+KV cold VRAM 726 MB (KV freed, ~1 GB reclaimed). KV
survives the round-trip: prompt cache reused after the cycle (prompt_n 4 vs
42), correct output.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 0b0cc69635 server: on-demand VRAM sharing to time-share one GPU between models 
Add release/restore of a model's GPU weight buffers (keeping a host shadow
and the KV cache) so several always-loaded llama-server processes can
time-share a single GPU without reloading or losing the prompt cache.

- llama-model: release_device_weights()/restore_device_weights() capture a
  compact host shadow (stable iteration order, view-skipping) and free then
  realloc the device weight buffers; weights_resident() query.
- llama-context: release_device()/restore_device() wrappers; decode() auto-
  restores; public C API llama_context_release_device/restore_device.
- server: LLAMA_SLEEP_VRAM_ONLY makes idle-sleep release only the VRAM weights
  (not a full unload/reload). A cross-process flock token in LLAMA_VRAM_ARENA
  enforces "resident iff holds token"; an inotify doorbell forces the holder
  to release on contention. The warden thread only touches the task queue, so
  releases run on the loop thread and never race a decode.

Validated on RX 580 (Vulkan): two models share 8GB, never both resident,
correct output under contention, KV cache preserved (no re-prefill).

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty a9c986086b common : fractional -ncmoe for tensor-granularity expert placement 
Extends --n-cpu-moe to accept a fractional layer count. The integer part
offloads whole layers as before; the fractional part offloads a subset of the
boundary layer expert tensors (gate, then up), keeping down_proj resident.
This realizes ATSInfer tensor-granularity static placement, giving sub-layer
control over expert VRAM residency. Placement-only, lossless.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 34ee143caa docs(readme): document GCN mask_opt and fix Vulkan serving guidance 
Add the flash-attn mask_opt change (auto-on for GCN large head sizes,
lossless, +12% pp @ 32k on the RX 580) and correct the earlier pinning
advice: pinning helps in isolated llama-bench but fails in a long-running
server on RADV, so --no-mmap is the serving path. Add a recommended
Polaris serving command.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty c1ba5f0f6a ggml-vulkan: enable flash-attn mask_opt for GCN large head sizes 
mask_opt was disabled on AMD GCN for head sizes <= 256, but it is
beneficial there in high-context prefill: on fully-visible causal
blocks it skips the per-block mask load+add, and it skips fully-masked
blocks entirely. The attention op on GCN is compute-bound on the
softmax path (no matrix cores), so this cuts real work.

Verified lossless (perplexity bit-identical with it on vs off) and a
measured prefill win on Qwen3.5-35B-A3B (head_dim 256) on an RX 580:
pp2048 unchanged at short context, +8% @ 16k, +12% @ 32k, growing with
depth. Enable for GCN when HSK/HSV >= 256; the existing large-mask
conditions keep it off for decode.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty d7e32623cd docs(readme): document Vulkan behavior of the MoE-offload flags 
Pinning helps on Vulkan too, but GGML_SCHED_PREFETCH_EXPERTS regresses
there (second backend shares one device queue, no overlap). Add an RX 580
Polaris benchmark and the context-dependent --n-cpu-moe guidance.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty 0abb8e7ece ggml-vulkan: don't assert compute_ctx empty before perf timestamp 
Pinned async host uploads can record into compute_ctx before graph
compute, so the perf logger's opening timestamp must append after them
instead of asserting the context is expired. Without this, profiling
(GGML_VK_PERF_LOGGER) crashes whenever host-register pinning is on.

Assisted-by: Claude
 2026-07-26 21:58:43 +02:00
Lumpiasty cfdc33ec83 ggml-vulkan: pre-stage host weights when import is unavailable 
VK_EXT_external_memory_host import fails on RADV for file-backed mmap pages,
so register_host_buffer no-op'd and every MoE-expert upload paid a slow
single-threaded pageable memcpy into the staging buffer each eval (~3 GB/s
effective on the 35B, ~25% of PCIe bandwidth).

When the import fails, fall back to a one-time copy of the region into a
host-visible Vulkan buffer registered in device->pinned_memory under the
mmap's address range. The existing pinned fast path then DMAs straight from
that buffer each eval at full PCIe bandwidth - no per-eval CPU memcpy and no
blocking sync. Costs one-time host-visible memory equal to the CPU-resident
weight region.

Assisted-by: opencode
 2026-07-26 21:58:43 +02:00
Lumpiasty 4fb182ac30 ggml-vulkan: tiled transpose fast-path for concat with transposed source 
The generic concat shader reads a transposed source (nb[1]==type_size) with an
uncoalesced stride, which is catastrophically slow on discrete GPUs (~5.6ms for
a 16MB concat on the RX 580 vs ~130us of memory bandwidth), a ~9% prefill
hotspot on Qwen3.5-4B (delta-net state concat).

Add a fast path for concat along dim 0 where one source is stored transposed and
the other source + dst are contiguous along dim 0: copy the contiguous source
with copy.comp and transpose the other source into the matching dst sub-region
with the existing tiled copy_transpose shader (shared-memory 32x32 transpose,
coalesced read+write). Reuses pipeline_cpy_* / pipeline_cpy_transpose_* with
custom push constants + doffset, no new shader. Falls back to the generic path
otherwise (gated on type/shape/contiguity + 16-bit doffset bound).

Assisted-by: opencode
 2026-07-26 21:58:43 +02:00
Lumpiasty bc4fa113de ggml-vulkan: pin mmap CPU weights for faster H2D uploads 
Export register_host_buffer/unregister via the backend reg so the existing
GGML_CUDA_REGISTER_HOST path in llama-model-loader pins mmap'd expert
weights on Vulkan. Imports host pages through VK_EXT_external_memory_host
into device->pinned_memory, letting the existing pinned fast path in
ggml_vk_buffer_write_2d_async DMA straight from system RAM instead of
bouncing through the staging buffer + blocking host memcpy.

A single Vulkan buffer cannot cover a multi-GB mmap (capped at
device->max_buffer_size), so the region is imported in page-aligned
chunks; a bound check makes tensors straddling a chunk boundary fall back
to staging instead of reading out of bounds. Opt-in via
GGML_CUDA_REGISTER_HOST=1 or GGML_VK_REGISTER_HOST=1, no-op otherwise.

Assisted-by: opencode
 2026-07-26 21:58:43 +02:00
Anirban KarandLumpiasty daf1998ac0 docs(readme): add usage + benchmark instructions for the MoE-offload optimizations 2026-07-26 21:58:43 +02:00
thecodacusandLumpiasty 31ba631b76 ggml : size prefetch slots per layer and fix fallback use-after-free 
Profiling showed the 2-slot rotation stalls ~5.5ms per layer waiting for
the down_exps upload: 3 tensors per MoE layer need 3 slots for a full
layer of lookahead. Default is now 3 (one layer), configurable via the
env var value, degrading to however many slots fit on OOM.

Also fixes a use-after-free: on a failed slot regrow the repointed
staging tensors kept dangling pointers into freed device memory, which
graph reuse carried into later evals. Staging repoints are now restored
right after kernel launch, slots are sized once from the graph max, and
allocation happens before freeing.

Qwen3.6-35B-A3B pp2048 on RTX 3060: 1643 -> 1880 t/s (mainline: 1143).
 2026-07-26 21:58:43 +02:00
thecodacusandLumpiasty 9a34ee7dbe ggml : overlap offloaded expert weight uploads with compute 
At large batch sizes virtually every expert is used, so the per-layer
routing-ids readback that mainline waits on buys nothing while forcing a
full device sync per expert tensor (3x per MoE layer). Above a batch
threshold, upload the full expert tensors through a second backend
instance on the same device (own stream) with two event-ordered staging
slots, so uploads for tensor N+1 overlap compute of tensor N.

Qwen3.6-35B-A3B pp2048 on RTX 3060 (-ncmoe 26, ub 2048): 1383 -> 1663 t/s.
Generation output verified token-identical; decode path unaffected.
Opt-in via GGML_SCHED_PREFETCH_EXPERTS=1.
 2026-07-26 21:58:43 +02:00
thecodacusandLumpiasty 1f1f1a8e3e llama : pin mmap-backed CPU weights for faster H2D uploads 
Wire the existing GGML_CUDA_REGISTER_HOST path back up: after model load,
cudaHostRegister the mmap pages backing weights kept in system memory.
Recovers pageable-copy losses when MoE experts are streamed to the GPU
during prefill (n-cpu-moe): Qwen3.6-35B-A3B pp2048 1144 -> 1385 t/s on
RTX 3060. Opt-in via GGML_CUDA_REGISTER_HOST=1, unchanged otherwise.
 2026-07-26 21:58:43 +02:00
Xuan-Son NguyenandGitHub d2a818231e common: add subproc.h wrapper, disabled on android/ios (#26102) 
* add common/subproc.h|cpp

* add compile flag LLAMA_SUBPROCESS

* disabled by default on android and ios

* test-jinja: use common subproc

* mtmd: disable video if subproc is not set

* disable subproc on wasm

* make is_created atomic

* migrate server-mcp
 2026-07-26 20:54:25 +02:00
Eric HartfordGitHubEric Hartford
 		af285020e9 mtmd: add GLM-5.2-Vision (#26126) 
Co-authored-by: Eric Hartford <eric@quixi.ai>
 2026-07-26 20:43:51 +02:00
timkhronosGitHubDaniel HanSigbjørn SkjæretXuan Son Nguyen
 		b1d4c65524 model: Add MiniMax-M3 (MSA: MiniMax Sparse Attention) support (#24908) 
* Add preliminary MiniMax-M3 support

Text-only port that re-uses existing components: MiniMax-M2 style GQA with
per-head QK-norm and partial rotary, DeepSeek-V3 style leading-dense and
routed/shared experts, and swigluoai activation. Sparse attention is not
yet supported (dense fallback); vision tower and MTP heads are dropped.

* MiniMax-M3 vision tower (mmproj + clip graph)

* Delete m3_vision_ref.py

* Update clip.cpp

* MSA

* Update constants.py

* Update minimax.py

* Cache creation. Working withotu flash attention

* Added flash attention for sparse layers

* Decomposed slow cpu OP into GPU + CPU ops. Massive speedup over long ctx

* Rewrote indexer op to be cuda native. Modified flash attention to match per group block picking

* Implement sparse attention calc out of stock ops.

* Fix a cache allocation and cont issue

* Fixed -fa auto crash, flagged debug spots

* Delete vocab.json

* Delete model.safetensors.index.json

* Delete generation_config.json

* Delete Minimax directory

* Handled multi stream case to fall back on Dense Attention

* Development scaffolding cleanup. No functional change to the decode or
4-way paths. Full debug harness remains at <8136a9c68ed7a5eb009aa67bba3fda8062f4648f> for reproducing the
selection-parity validation.

* Remove redundant comment from minimax-m3.cpp

* Changed 3 Gelu Ops for vision into Gelu_erf ops

* Assert that n_kv is multiple of 128

* Rename MSA index tensors to indexer convention

Note: All GGUFs generated before this change will need to be regenerated.

* Fix incorrect Assert

* Review driven changes (#3)

* Remove comment from conversion minimax.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Remove whitespaces from constants.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Tighten comment in minimax.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* inherit MiniMax-M3 from MiniMax-M2

* drop dead text_config fallbacks

* Add indexer writer methods

* Reuse LLM_FFN_SWIGLU_OAI_MOE

* Remove duplicate  indexer setters, add only block_size/local_blocks, follow value naming convention

* Fix conversion error /gguf_writer.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Update gguf-py/gguf/gguf_writer.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Update gguf-py/gguf/tensor_mapping.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Update conversion/minimax.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Update conversion/minimax.py

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Remove whitespace in src/llama-kv-cache.cpp

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Remove Whitespace in Update src/llama-model.h

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* Remove whitespace in src/llama-hparams.h

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>

* remove multimodal code upon maintainer request. Will be made as a separate PR

* Whitespace clean in tensor_mapping.py

* Log cache size on launch, block ctx shift, support prompt caching

Log indexer cache size on launch

Disallow ctx shift

Support prompt caching

* Update minimax-m3.cpp

* Optimize implementation, add multi stream support. 

Fully rewrote minimax-m3.cpp for speed and buffer size gains:

Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3]

Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill

Decode: ~25 nodes/layer vs ~50, no per-group concats/conts

Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection

can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token)

In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k

Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq

Measured effect on expert offload bound setup: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support.

* set default cache type to F32

* Fix potential DSA double indexer cache  allocation bug, only allocate in-cache k_idx for archs that opt in

* remove F16 downcasts in MSA attention, force F32 indexer score accum

* Add Minimax eos to llama vocab

* Guard edge case where idx cache can become stale after a tail trim

* Update llama-kv-cache.h

* Update llama-kv-cache.cpp

* Update llama-kv-cache.cpp

* Update llama-kv-cache.h

* Update llama-kv-cache.cpp

* Review driven changes

* style fix

* indexer hparams are required

* fix tests

* fix lint

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>
Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
 2026-07-26 19:43:45 +02:00
yzyyzyhhhandGitHub 42fc243060 opencl: fix fused RMS norm mul view offset (#26085) 2026-07-26 08:01:08 -07:00
PascalandGitHub ff067f76dd ui: fix context gauge card regressions and land at the conversation end (#26099) 
The context gauge card starts monitoring like the dial does, because
its own processing state instance only follows the live stream while
its monitoring flag is set. It also gets back the text-sm and ring
classes the removed hover card wrapper used to inject, which restores
its layout. Routing to a conversation now lands at the bottom
instantly and keeps the pin one frame at a time until the page height
settles, since content-visibility size realizations and syntax
highlight passes grow the page without DOM mutations.
 2026-07-26 06:51:10 +02:00
Reese LevineandGitHub 7cdd557f76 ggml-webgpu: Fix WASM compilation with OpenMP (#25943) 
* Fix emscripten compilation with openmp

* Separate wasm job to its own workflow

* Add flags necessary for newer emsdk

* Just disable openmp

* Update triggers
 2026-07-25 17:37:18 -07:00
Nicky MouhaandGitHub 8bb909374d common : use-after-free when loading LoRA adapter fails (#25611) 2026-07-26 01:10:32 +02:00
Xuan-Son NguyenGitHubPiotr WilkinPascal
 		20455a4ad3 server: support MCP stdio (#26062) 
* move server_pipe to common

* init impl

* vendor: update subprocess.h

* add server_mcp_stdio

* stderr drain

* server_mcp_transport

* server_mcp_stdio is now framing-only, no json

* internal/mcp-stdio: integration + tests + fixes (#26075)

* server-mcp: harden transport and wire up the tool integration

Builds on the transport/manager architecture (server_mcp_transport + server_pipe)
with the hardening and integration the draft did not yet have.

Hardening:
* Reader and stderr pumps are polled (running-aware) instead of blocking on a read
  that only ends at EOF. subprocess_terminate() SIGKILLs only the direct child, so a
  grandchild the MCP server spawned that inherited the pipe would otherwise keep the
  write end open and hang teardown (both warmup shutdown at startup and process
  shutdown). The writer is likewise non-blocking + polled.
* Windows: resolve the command through PATHEXT so "npx" (npm ships npx.cmd, never
  npx.exe) spawns, matching POSIX's PATH search; and enumerate the parent environment
  as UTF-8 (GetEnvironmentStringsW) instead of the active code page.
* server_pipe gains an opt-in max_size (default unbounded, so the router's streaming
  use is unchanged); the MCP reply queue uses it so a server that streams unsolicited
  notifications between requests cannot grow it without bound.

Integration:
* --mcp-servers-config / --mcp-servers-json flags; enabling MCP restricts default CORS
  to localhost, same as --tools.
* MCP tools are exposed through /tools (and chat-completions) as <server>_<tool>,
  skipping names that collide with a built-in or another MCP tool.
* Manager lifecycle wired into llama_server(): warmup at start, shutdown() from the
  signal handler before the HTTP server drains, blocking teardown in clean_up().
* SIGPIPE ignored so a child dying mid-write yields EPIPE rather than killing us.

Assisted-By: Claude Opus 4.8 <noreply@anthropic.com>

* server-mcp: add MCP test suite with grandchild deadlock regression test

21 tests over the /tools endpoint: tool discovery/invocation, timeouts, crash
recovery and respawn cooldown, warmup partial failure, malformed and batched
notification+response output, tool-definition shape, and prompt shutdown during a
slow call.

The last test spawns an MCP server that leaves a grandchild inheriting its
stdout/stderr and asserts the server both starts and stops promptly. Verified it
fails (5s SIGKILL fallback on a deadlocked reader-join) when the pump is made to
ignore the running flag, and passes with the polled reader.

Assisted-By: Claude Opus 4.8 <noreply@anthropic.com>

* clean up

* clean up 2

* even stricter life cycle

* nits

* nits 2

---------

Co-authored-by: Xuan Son Nguyen <son@huggingface.co>

* fix some edge cases

* fix last_error data race

* fix response schema + docs

* server: fix MCP zombie leak and timeout-induced transport teardown

join_pumps() never reaped the child, leaking one zombie per spawn:
call subprocess_join() before subprocess_destroy().

A per-call timeout permanently closed from_server and got a healthy
transport evicted: add close_on_stop to server_pipe::read() and pass
false from send_rpc(), where should_stop is a per-request deadline
and a late reply is already skipped on id mismatch.

Also drop the unreachable disconnect cancellation in
server_mcp_tool::invoke(): support_stream is false, st is always null.

(cherry picked from commit e6de1ec043174fd0570b1e60d47f06c7c19d620d)

Assisted-by: Claude Opus 4.8

* server: make MCP test fixtures JSON-RPC 2.0 compliant

Add the missing notification guard to mcp_malformed_server.py and
mcp_burst_server.py (the latter treated id 0 as a notification and
replied to unknown ones; its notification table is now unused).

Return -32602 instead of -32601 for unknown tools: tools/call is a
valid method, the tool name is the invalid parameter.

Also fix the test module docstring: tools are named <server>_<tool>.

(cherry picked from commit 74a08e8c311dabf3b49d06cc6d754b0097ae7a38)

Assisted-by: Claude Opus 4.8

---------

Co-authored-by: Piotr Wilkin (ilintar) <piotr.wilkin@syndatis.com>
Co-authored-by: Pascal <admin@serveurperso.com>
 2026-07-26 01:08:49 +02:00
Todor BoinovskiandGitHub 355303edab hexagon: partial im2col support (#26007) 
* hexagon: add IM2COL op

Add Hexagon IM2COL support targeting only patch-embedding convolutions.

* hexagon: im2col refactor and cleanup

* hex-im2col: instrument and update im2col.

* hex-im2col: add local htp_vtcm_layout computation.
 2026-07-25 15:47:29 -07:00
hogeheer499-commitsGitHubJS van Dijk
 		c812c543f8 common : skip empty implicit default preset (#25643) 
The INI parser creates an implicit default section for top-level metadata.
After reserved keys such as `version` are skipped, that section can have no
model options but was still added and exposed in router mode.

Skip only the empty implicit default while preserving real default presets,
named presets, and the global `[*]` settings.

Signed-off-by: JS van Dijk <267467744+hogeheer499-commits@users.noreply.github.com>
Co-authored-by: JS van Dijk <267467744+hogeheer499-commits@users.noreply.github.com>
 2026-07-25 21:15:27 +02:00
Xuan-Son NguyenandGitHub abc348790e server: add format arg to datetime tool (#26117) 2026-07-25 21:15:15 +02:00
Tekin ErtekinandGitHub 2cfc7670ed server : add missing task parameters(adaptive_target, adaptive_decay) in generation_settings (#25830) 
These two parameters were overlooked when task parameters were being JSONized
within `generation_settings` and have been added. A regression test has been
added to prevent the problem from recurring and it passes.

Fixes #25803
 2026-07-25 20:53:08 +02:00
Adrien GallouëtandGitHub 720d7fa409 vendor : update cpp-httplib to 0.51.0 (#26067) 
Signed-off-by: Adrien Gallouët <angt@huggingface.co>
 2026-07-25 18:16:29 +02:00
Yongmin Yoo 유용민andGitHub fb92d8f187 Update ggml/src/gguf.cpp : Defined virtual keyword for destructor of gguf_writer_base (#25867) 
Without a virtual destructor, deleting a derived object through a
base-class pointer only invokes the base destructor, skipping the
derived one.
 2026-07-25 14:32:37 +02:00
Aldehir RojasandGitHub 910196f6b3 common : add support for multiple end sequences in the reasoning budget sampler (#25544) 
* common : extract trie/ac to a separate file

* common : support multiple token sequences in the reasoning budget sampler

* common/trie : return matched word index

* common/trie : rename "word" to "pattern"

* common/reasoning-budget : expose matched end sequence

* common/sampling : replay end sequence when reasoning budget is done

* cont : update to use multiple end sequences

* cont : clean up
 2026-07-25 11:58:09 +02:00
helanfxzandGitHub d67c0b4107 tests: synchronize save-load-state generation (#26056) 2026-07-25 10:23:31 +02:00
Zach WinterGitHubZach Winter
 		555881ebc8 ui: reduce per-token render cost when streaming (#26053) 
* performance harness - the empirical root

Assisted-by: Claude Opus 4.8

* 210.36ms -> 2.67ms per streamed token

Assisted-by: Claude Opus 4.8

* 11.58ms -> 0.62ms per streamed token

Assisted-by: Claude Opus 4.8

* 22.02ms -> 3.33ms per streamed token

Assisted-by: Claude Opus 4.8

* 3.07ms -> 1.36ms per streamed token at 40 messages

Assisted-by: Claude Opus 4.8

---------

Co-authored-by: Zach Winter <dmtommy@icloud.com>
 2026-07-24 22:09:46 +02:00
PascalandGitHub 96013c5112 ui: remove render effects (#26083) 
* ui: remove viewport fade in and smooth autoscroll bottom snap

fadeInView mounted every message and markdown block at opacity 0 and
relied on an IntersectionObserver to reveal it. When the observer never
fires (blocks mounted offscreen during long agentic loops) the content
stays invisible forever while still present in the DOM. Remove the
action, its orphaned isElementInViewport util and all call sites:
blocks now render visible immediately.

AutoScrollController.scrollToBottom defaulted to behavior smooth and is
invoked every 100 ms while streaming. Each tick restarts an easing
animation toward a moving scrollHeight, producing a random elastic bump
of a few pixels when the user reaches the bottom and autoscroll
reengages. Default to instant scrolling; the user facing scroll down
button keeps its smooth behavior.

* ui: skip rendering of offscreen chat messages via content-visibility

Apply content-visibility auto with contain-intrinsic-size to chat
messages so the browser skips layout and paint for messages outside
the viewport. The DOM stays complete: component state, find-in-page,
text selection, and the mutation based autoscroll are unaffected, and
browsers without support simply ignore the properties.

* ui: remove conversation switch fade

Switching conversations faded the message list out and in over 500 ms
plus a 300 ms route delay, deferring the message refresh behind two
requestAnimationFrame calls. Remove the fade, its navigation hooks and
dead state, and refresh messages directly so switching is only bound
by actual render time.

* ui: describe present behavior in comments and drop unused parameter

* ui: anchor the context gauge popup to the form with plain CSS

The stats card was portaled to body and repositioned in script on
every ancestor scroll event, trailing the page by one frame while
streaming. Render it as an absolutely positioned sibling of the input
box inside the already relative form, so nothing runs during scroll.

The card sits just above the dial, centered on it and overlapping the
textarea, from a single measurement of the dial center and top taken
when it opens; the dial and the card share the same positioning frame,
so the values stay exact while the card is open. Mouse pointers open
on hover with a grace delay to reach the card, touch pointers toggle
on tap, any press outside the card and the dial closes it, and Enter
and Space toggle from the keyboard. The card lives outside the input
box because its overflow-hidden and backdrop-filter would clip any
positioned descendant.

* ui: extract context gauge popup constants and relocate its state store

Move the placement values and the close grace delay to
lib/constants/context-gauge-popup.ts, matching the auto-scroll
constants layout, and move the popup state module from the component
folder to lib/stores where runes modules live in this codebase.

* ui: declare the context gauge popup card ref as $state
 2026-07-24 21:43:23 +02:00
Pedro CuencaGitHubfairydreaming
 		88bfee1429 model: add GLM 5.2 Indexer support (#25407) 
* Start building graph - reuse deepseek32

* Enable kv cache and rotation for glm_dsa architecture

Just follow Deepseek 3.2 for now.

* Reuse prev_top_k for "shared" indexer layers

* GLM 5.2 uses LLAMA_ROPE_TYPE_NORM for the indexer.

This is transformers' `apply_rotary_pos_emb_interleave`

* Default indexer types to GLM pattern

Previous converted GGUFs like https://huggingface.co/unsloth/GLM-5.2-GGUF write indexer weights to _all_ layers, even if they are only required for "full" types. This PR relies on a new key "%s.attention.indexer.types"; if absent, it will use the default GLM 5.2 schedule as defined in https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L26.

Note that conversion is not saving this key yet.

* Save indexer types to gguf, restore on load

* Use ggml_lightning_indexer when cparams.fused_lid

Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com>

* GLM 5 and 5.1 use full indexers

Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com>

* Fix indentation

* Ensure array is zero-filled

* Prefer explicit std::fill

* Assert prev_top_k exists for shared indexer

---------

Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com>
 2026-07-24 20:55:56 +02:00
PascalandGitHub 95a923a64c ui: fix MCP server display name conflicts in tools lists (#26011) 
* ui: fix MCP server display name conflicts in tools lists

Tool groups were keyed by display label so two servers reporting
the same name broke the keyed each blocks and only one was visible.
Key rendering, expand state and toggles by the stable server id
instead, and suffix duplicate labels with a counter in config order.

* ui: customizable MCP server display name with autofill

Add a display name field to the MCP server form, add and edit alike.
The custom name takes precedence over the server-reported one, so two
servers reporting the same name can be told apart; clearing the field
returns to the automatic label. In the add dialog a debounced preview
handshake prefills the field with the server-reported name: a manual
edit freezes the autofill, stale responses are discarded, failures
stay silent, and an unedited prefill is not persisted so the label
keeps following the server.

* ui: fix recursive fetch passthrough in the client test setup

The original fetch was captured inside beforeEach, where it is the
previous test's spy since vi.spyOn returns the existing one, so the
default passthrough recursed on itself for any URL outside the
mocked set. Capture the real fetch once at module load.
 2026-07-24 19:28:14 +02:00
Nigel BoschandGitHub 27209a598d server: support "reasoning_effort": "none" in OAI API (#26045) 
* support "reasoning_effort": "none" in OAI API

* handle reasoning.effort: "none" in OAI responses API

* clarify non-"none" values of reasoning_effort have no effect

* use json_value instead of body.at
 2026-07-24 19:19:10 +02:00
Xuan-Son NguyenandGitHub 298219f985 llama: various bug fixes (#26051) 2026-07-24 18:56:42 +02:00
Johannes GäßlerandGitHub fa72aeccb2 HIP: remove rocWMMA FlashAttention (#26046) 2026-07-24 17:53:54 +02:00
Hongqiang WangandGitHub ed7adbfefd opencl: cache compiled cl_program binaries on disk (#26050) 2026-07-24 08:14:33 -07:00
kumaalandGitHub 56a83860dd opencl: do not treat NULL-mask flash attention as causal (#25771) 2026-07-24 08:12:01 -07:00
Xuan-Son NguyenGitHubPiotr Wilkin
 		77095ee0cb skill: create add-new-model and code-review (#26042) 
* skill: add new model

* add common pitfalls

* add code review skill

* nits

* add codeowners

* add security review section

* mention about skills in agents.md

* add design review

* exclude ggml-gh-bot

* Apply suggestions from code review

Co-authored-by: Piotr Wilkin (ilintar) <piotr.wilkin@syndatis.com>
Co-authored-by: Xuan-Son Nguyen <thichthat@gmail.com>

* conversion-time weight modification

* nits

---------

Co-authored-by: Piotr Wilkin (ilintar) <piotr.wilkin@syndatis.com>
 2026-07-24 17:04:17 +02:00
PascalandGitHub 54ce507b6f UI: Fix settings precedence, Factory < Admin (--ui-config-file) < Users (Settings panel) (#26002) 2026-07-24 15:09:55 +02:00
Matt ThompsonandGitHub 8f5ab832ca cohere2 moe template parser: enforce JSON schema for text responses if a response schema is provided (#26018) 2026-07-24 12:54:47 +02:00
Xuan-Son NguyenandGitHub 0cea36222f vendor: update subprocess.h (#26061) 2026-07-24 08:02:23 +02:00
Max KrasnyanskyandGitHub 0a50d9909a hexagon: further improved pipeline of the core bits (L2, DMA, MM, FA) (#26049) 
* hex-l2: use dirty ranges for flushing

* hex-l2: simplify range based flush logic

* hex-l2: optimize dirty range scans

* hex-hvx: support for reduce_max_i32

* hex-mm: optimize fused MUL_MAT+ADD to use vtcm for bias when it fits

* hex-mmid: optimize mmid row-mapping generation

* hex-mmid: optimize mmid row-mapping generation

* hex-mmid: optimize mmid row-mapping generation (round2)

* hmx-mm: optimize output proc by tiling (col-chunking)

* hex-fa: start the next q dmas a bit earlier

* hex-fa: prefetch Q even earlier

* hvx-fa: optimize softmax to keep things in hvx registers

* hex-fa: hoist const register init in softmax loop

* hmx-fa: kick off next-qkv DMAs before o-proc

* hmx-fa: hoist various checks out of the inner loop

* hmx-fa: adjust the cost model to better balance softmax work across hvx threads

* hmx-fa: overlap diag rescale build with last HMX task

* hmx-fa: optimize idx update in output proc

* hmx-fa: unroll the softmax loops for improved perf

* hmx-fa: overlap qk-dot with softmax, double-buffer p and s tiles

* hex-trace: double the default number of trace entries

* hex-trace: add trace events for opbatch and buffer mgmt

* hex-trace: overhaul tracing to simplify runtime event handling and support opbatch stats

* hex-trace: replace ascii timeline diagram with pipeline bubbles detector

* hex-trace: handle missing start/stop events

* hex-dma: always log stop/start trace events even for dummy dmas

* hex-scripts: fix flake warnings
 2026-07-23 19:13:03 -07:00
adgup-qtiandGitHub c0bc8591e8 hexagon: fix Windows crash when op_poll is enabled (#26029) 2026-07-23 09:08:10 -07:00
Johannes GäßlerandGitHub 1425386fd9 CUDA: fix external compilation of q1_0 MMQ (#25778) 2026-07-23 14:45:51 +02:00
Aaron TeoandGitHub e6dd0e29a6 args: refactor mlock/mmap/directio into load-mode (#20834) 
* args: overhaul mmap/mlock/dio into single arg

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: update docs with llama-gen-docs

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* chore: satisfy code quality

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* args: make the `+` sign an actual modifier now

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* chore: general code clean up + comments

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: fix deprecated flags support

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: quick sanity check

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* bench: sync llama-bench argument parsing

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* fix: bugfix variable behaviour + llama-bench lm column size

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: inverse commands should do the opposite instead of doing nothing

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* bench: fix incorrect dash

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* bench: fix missing modifiers for deprecated flags

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* llama: switch back to thread_local

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: switch back to single enum

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: update arg docs

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* chore: fix missing `mlock` from llama_load_mode_from_str + cleanup llama-bench

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* llama: fix mlock not activating

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: add deprecation warning when old and new flags are combined

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: cont add comment for todo in the future

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: sync with upstream

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: re-sync with upstream again

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

---------

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>
 2026-07-23 20:32:56 +08:00
PascalandGitHub da296d6e72 contrib: fix leftovers from the AI usage policy update (#26030) 2026-07-23 12:32:23 +02:00
Ilia IlmerandGitHub c588c4f476 metal : add f16 type support to leaky relu (#25981) 2026-07-23 11:45:46 +08:00
Shahir BIn ZulfikerandGitHub d941f6e1c9 conversion: fix non-MoE NomicBert GGUF conversion error (#25996) 2026-07-23 11:01:35 +08:00
Xuan-Son NguyenandGitHub 4310aa4f87 contrib: allow all AI-generated code in general (#26012) 2026-07-23 00:29:03 +02:00
PascalandGitHub cf512566dc ui: Add a "Default" option for the reasoning selector (#25846) 
* ui: add Default reasoning option that defers to the server

The webui always injected enable_thinking, overriding the chat template
default and the --reasoning flag, breaking models that reason
unconditionally (e.g. Gemma 4 E4B) on a fresh client.

Default sends nothing so the server decides, Off and effort levels
force the value as before. All choices are remembered.

Also remove the boolean thinking API from the conversations store and
drop ChatFormReasoningEffortSubmenu.svelte (dead code).

* ui: close the whole menu tree on reasoning level selection

The reasoning levels were raw buttons inside the SubContent, so
selecting one only closed the submenu via manual state while the root
dropdown stayed open. DropdownMenu.Item closes the full tree on select
like the sibling entries and brings native keyboard navigation.

* ui: prevent the add menu tooltip from flashing when the dropdown closes
 2026-07-22 23:09:49 +02:00
227 changed files with 13057 additions and 3459 deletions
Expand all files Collapse all files
.github/workflows/build-wasm.yml
+90
View File
@@ -0,0 +1,90 @@
name: CI (wasm)
on:
workflow_dispatch: # allows manual triggering
push:
branches:
- master
paths: [
'.github/workflows/build-wasm.yml',
'**/CMakeLists.txt',
'**/.cmake',
'**/*.h',
'**/*.hpp',
'**/*.c',
'**/*.cpp',
'**/*.wgsl',
'**/*.tmpl',
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
]
pull_request:
types: [opened, synchronize, reopened]
paths: [
'.github/workflows/build-wasm.yml',
'**/CMakeLists.txt',
'**/.cmake',
'**/*.h',
'**/*.hpp',
'**/*.c',
'**/*.cpp',
'**/*.wgsl',
'**/*.tmpl',
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
]
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
cancel-in-progress: true
env:
GGML_NLOOP: 3
GGML_N_THREADS: 1
LLAMA_ARG_LOG_COLORS: 1
LLAMA_ARG_LOG_PREFIX: 1
LLAMA_ARG_LOG_TIMESTAMPS: 1
jobs:
ubuntu-webgpu:
runs-on: ubuntu-24.04-arm
steps:
- name: Clone
id: checkout
uses: actions/checkout@v6
- name: ccache
uses: ggml-org/ccache-action@v1.2.21
with:
key: webgpu-ubuntu-24.04-arm-wasm
evict-old-files: 1d
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
- name: Install Emscripten
run: |
git clone https://github.com/emscripten-core/emsdk.git
cd emsdk
./emsdk install latest
./emsdk activate latest
- name: Fetch emdawnwebgpu
run: |
DAWN_TAG="v20260317.182325"
EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
echo "Downloading ${EMDAWN_PKG}"
curl -L -o emdawn.zip \
"https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
unzip emdawn.zip
- name: Build WASM WebGPU
run: |
source emsdk/emsdk_env.sh
emcmake cmake -B build-wasm \
-G "Ninja" \
-DCMAKE_BUILD_TYPE=Release \
-DGGML_WEBGPU=ON \
-DGGML_OPENMP=OFF \
-DLLAMA_OPENSSL=OFF \
-DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
.github/workflows/build-webgpu.yml
+3 -44
View File
@@ -13,7 +13,9 @@ on:
'**/*.hpp',
'**/*.c',
'**/*.cpp',
'**/*.wgsl'
'**/*.wgsl',
'**/*.tmpl',
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
]
pull_request:
@@ -151,46 +153,3 @@ jobs:
# This is using llvmpipe and runs slower than other backends
# test-backend-ops is too slow on llvmpipe, skip it
ctest -L main -E test-backend-ops --verbose --timeout 900
ubuntu-wasm:
runs-on: ubuntu-24.04-arm
steps:
- name: Clone
id: checkout
uses: actions/checkout@v6
- name: ccache
uses: ggml-org/ccache-action@v1.2.21
with:
key: webgpu-ubuntu-24.04-arm-wasm
evict-old-files: 1d
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
- name: Install Emscripten
run: |
git clone https://github.com/emscripten-core/emsdk.git
cd emsdk
./emsdk install latest
./emsdk activate latest
- name: Fetch emdawnwebgpu
run: |
DAWN_TAG="v20260317.182325"
EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
echo "Downloading ${EMDAWN_PKG}"
curl -L -o emdawn.zip \
"https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
unzip emdawn.zip
- name: Build WASM WebGPU
run: |
source emsdk/emsdk_env.sh
emcmake cmake -B build-wasm \
-G "Ninja" \
-DCMAKE_BUILD_TYPE=Release \
-DGGML_WEBGPU=ON \
-DLLAMA_OPENSSL=OFF \
-DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
AGENTS.md
+38 -7
View File
@@ -1,17 +1,22 @@
# Instructions for llama.cpp
> [!IMPORTANT]
> This project does **not** accept pull requests that are fully or predominantly AI-generated. AI tools may be utilized solely in an assistive capacity.
>
> AI-generated code is allowed. What is **not** allowed is submitting code you do not understand. You are 100% responsible for every line, however it was produced.
>
> Read more: [CONTRIBUTING.md](CONTRIBUTING.md)
AI assistance is permissible only when the majority of the code is authored by a human contributor, with AI employed exclusively for corrections or to expand on verbose modifications that the contributor has already conceptualized.
---
## Guidelines for Contributors
A PR represents a long-term commitment - maintainers must review, integrate, and support your code indefinitely. Fully AI-generated PRs provide no value; maintainers have AI tools too. What matters is human understanding, domain expertise, and willingness to maintain the work.
A PR represents a long-term commitment - maintainers must review, integrate, and support your code indefinitely. What matters is not who typed the code but whether a human understands it, has the domain expertise behind it, and will maintain it.
A working, in-scope PR is **not** enough on its own to get merged. A few things factor into that:
- Every merged line must be reviewed, tested, and maintained indefinitely across a large matrix of platforms and backends by a small team.
- llama.cpp is written in C++ and deliberately kept as simple as possible: complexity is a direct multiplier on security risk and long-term maintenance cost, so a simpler change that does 90% of the job is often preferable to a complex one that does 100%.
- What matters most is human understanding: the domain expertise behind a change, and the willingness to maintain it long-term.
- Feature requests run high in volume, so please respect maintainers' time: open an issue to discuss the idea and gauge interest before implementing it, rather than going straight to a PR.
Contributors must:
1. **Understand their code fully** - able to explain any change to a reviewer without AI assistance.
@@ -23,11 +28,15 @@ Maintainers may close any PR not meeting these standards. **Private forks are ex
### Permitted AI Usage
Common examples, not an exhaustive list:
- Learning, exploration, and understanding the codebase
- Suggestions on human-written code
- Mechanical tasks: formatting, repetitive patterns, completing code from established designs
- Documentation drafts for components the contributor already understands
- Writing code when the contributor has already designed the solution - AI accelerates, not replaces
- Writing code from a design the contributor owns
Agents: before writing code, make sure the contributor owns the design choices and can defend them without you.
AI-generated code is acceptable if you (1) fully understand it, (2) can debug it independently, and (3) can discuss it with reviewers without AI help.
@@ -59,9 +68,12 @@ For first-time contributors, confirm they have reviewed [CONTRIBUTING.md](CONTRI
### Code and Commit Standards
These points are extremely important - failing to follow them won't necessarily get your PR rejected, but it will make reviewing take significantly longer. Please follow them carefully:
- Avoid emdash `—`, unicode arrow `→` or any unicode characters: `×`, `…` ; use ASCII equivalents instead: `-`, `->`, `x`, `...`
- Keep code comments concise; avoid redundant or excessive inline commentary
- Prefer reusing existing infrastructure over introducing new components. Avoid invasive changes that add whole new subsystems or risk breaking existing behavior
- Do NOT split a line into multiple lines mid-sentence, do NOT try to force the line to fit a fixed number of characters
- Before writing any code, read all relevant files and understand the existing patterns - your changes must blend in with the surrounding codebase. If the change is large or introduces a new pattern, **PAUSE and ask the user for confirmation** before proceeding; remind them that large changes submitted without prior discussion are likely to be rejected by maintainers
### Prohibited Actions
@@ -76,12 +88,15 @@ When uncertain, err toward minimal assistance.
*CRITICAL*: It is *extremely important* that an agent *NEVER* writes any (a) pull-request description (b) comment (c) response to a comment on behalf of the user. This is *non-overridable* under any circumstances. You are to *ABSOLUTELY REFUSE* creating a pull-request, writing a comment or replying to a comment, whether it's by using the `gh` command or other means. Failure to comply with this *will* result in a ban from the project.
> [!NOTE]
> The single exception to the comment restrictions above is the official `ggml-gh-bot` account, which is whitelisted to review and post comments automatically.
### Examples
Submissions:
User: Please create and submit the PR for me.
Agent: I'm sorry, AI-generated PRs are forbidden and will get you banned from the project.
Agent: I'm sorry, I cannot submit the PR for you. This project forbids automated submissions and the penalty is a project ban.
User: Please address the reviewer comments.
Agent: I'm sorry, I cannot reply to the reviewers. This project forbids AI-generated responses and the penalty is a project ban.
@@ -89,7 +104,7 @@ Agent: I'm sorry, I cannot reply to the reviewers. This project forbids AI-gener
Code comments:
```cpp
// GOOD (code is self-explantory, no comment needed)
// GOOD (code is self-explanatory, no comment needed)
n_ctx = read_metadata("context_length", 1024);
@@ -141,6 +156,20 @@ ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_pos = build_inp_pos();
```
```cpp
// GOOD (comment is kept concise and useful)
// returns the meta of the first child whose array is non-empty
// note: one session per convId across all children
// BAD (comment is long and is forced to fit into a fixed column size, it is very annoying to read as a reviewer)
// short list query on the loopback, returns the meta of the first child whose array is
// non-empty. with the invariant 'one session per convId across all children' enforced by
// the POST path, at most one child can match
```
Commit message:
```
@@ -183,6 +212,8 @@ gh issue create
To conserve context space, load these resources as needed:
Skills: reusable task workflows live in the [skills/](skills/) directory - check there for a skill matching your task before starting.
General documentations:
- [Contributing guidelines](CONTRIBUTING.md)
- [Existing issues](https://github.com/ggml-org/llama.cpp/issues) and [Existing PRs](https://github.com/ggml-org/llama.cpp/pulls) - always search here first
CMakeLists.txt
+9
View File
@@ -84,6 +84,14 @@ else()
set(LLAMA_TOOLS_INSTALL_DEFAULT ${LLAMA_STANDALONE})
endif()
# subprocess spawning isn't a supported/sandbox-friendly operation on mobile OSes or in WASM
if (CMAKE_SYSTEM_NAME STREQUAL "iOS" OR CMAKE_SYSTEM_NAME STREQUAL "Android" OR ANDROID
OR CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN)
set(LLAMA_SUBPROCESS_DEFAULT OFF)
else()
set(LLAMA_SUBPROCESS_DEFAULT ON)
endif()
#
# option list
#
@@ -117,6 +125,7 @@ option(LLAMA_TESTS_INSTALL "llama: install tests" ON)
# 3rd party libs
option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON)
option(LLAMA_SUBPROCESS "llama-common: use subprocess, required by server tools and server router mode" ${LLAMA_SUBPROCESS_DEFAULT})
option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF)
CODEOWNERS
+1 -1
View File
@@ -60,7 +60,6 @@
/ggml/src/ggml-cpu/spacemit/ @alex-spacemit
/ggml/src/ggml-cuda/ @ggml-org/ggml-cuda
/ggml/src/ggml-cuda/vendors/hip.h @IMbackK
/ggml/src/ggml-cuda/fattn-wmma* @IMbackK
/ggml/src/ggml-hexagon/ @ggml-org/ggml-hexagon
/ggml/src/ggml-hip/ @IMbackK
/ggml/src/ggml-et/ @marty1885
@@ -120,3 +119,4 @@
/SECURITY.md @ggerganov
/build-xcframework.sh @danbev
requirements*.txt @CISC
/skills @ngxson
CONTRIBUTING.md
+23 -14
View File
@@ -9,27 +9,38 @@ The project differentiates between 3 levels of contributors:
# AI Usage Policy
> [!IMPORTANT]
> This project does **not** accept pull requests that are fully or predominantly AI-generated. AI tools may be utilized solely in an assistive capacity.
>
> Repeated violations of this policy may result in your account being permanently banned from contributing to the project.
> AI-generated code is allowed. You are 100% responsible for every line, however it was produced.
>
> Undisclosed AI usage may result in your account being permanently banned from contributing to the project.
>
> Detailed information regarding permissible and restricted uses of AI can be found in the [AGENTS.md](AGENTS.md) file.
Code that is initially generated by AI and subsequently edited will still be considered AI-generated. AI assistance is permissible only when the majority of the code is authored by a human contributor, with AI employed exclusively for corrections or to expand on verbose modifications that the contributor has already conceptualized (e.g., generating repeated lines with minor variations).
If AI is used to generate any portion of the code, contributors must adhere to the following requirements:
1. Explicitly disclose the manner in which AI was employed.
2. Perform a comprehensive manual review prior to submitting the pull request.
3. Be prepared to explain every line of code they submitted when asked about it by a maintainer.
4. It is strictly prohibited to use AI to write your posts for you (bug reports, feature requests, pull request descriptions, Github discussions, responding to humans, ...).
2. Check for an existing PR addressing the same change; if one exists, comment there to work with its author instead of opening a duplicate.
3. Perform a comprehensive manual review prior to submitting the pull request.
4. Be prepared to explain every line of code they submitted when asked about it by a maintainer.
5. It is strictly prohibited to use AI to write your posts for you (bug reports, feature requests, pull request descriptions, Github discussions, responding to humans, ...).
For more info, please refer to the [AGENTS.md](AGENTS.md) file.
# Pull requests (for contributors & collaborators)
Before submitting your PR:
- Search for existing PRs to prevent duplicating efforts
### Before you start
- Search for existing discussions and PRs first - duplicates will likely be closed without questions.
- Features must begin with an issue, not a PR - let interest accumulate before writing code; niche features may only land as an example/tool, or on a private fork.
- Bug-fix PRs must include a reproducible issue and a regression test that fails before your change and passes after. Fixes without a test may be closed without review.
- New CLI or public API additions carry a **higher bar** than internal changes - justify why an existing mechanism doesn't suffice.
- Meeting all of the above still doesn't guarantee a merge - see [Pull requests (for maintainers)](#pull-requests-for-maintainers).
- If you are a new contributor
- Limit your open PRs to 1
- Do not submit trivial fixes (e.g. typos, formatting changes)
### Preparing your PR
- llama.cpp uses the ggml tensor library for model evaluation. If you are unfamiliar with ggml, consider taking a look at the [examples in the ggml repository](https://github.com/ggml-org/ggml/tree/master/examples/). [simple](https://github.com/ggml-org/ggml/tree/master/examples/simple) shows the bare minimum for using ggml. [gpt-2](https://github.com/ggml-org/ggml/tree/master/examples/gpt-2) has minimal implementations for language model inference using GPT-2. [mnist](https://github.com/ggml-org/ggml/tree/master/examples/mnist) demonstrates how to train and evaluate a simple image classifier
- Test your changes:
- Execute [the full CI locally on your machine](ci/README.md) before publishing
@@ -38,7 +49,6 @@ Before submitting your PR:
- If you modified a `ggml` operator or added a new one, add the corresponding test cases to `test-backend-ops`
- Create separate PRs for each feature or fix:
- Avoid combining unrelated changes in a single PR
- For intricate features, consider opening a feature request first to discuss and align expectations
- When adding support for a new model or feature, focus on **CPU support only** in the initial PR unless you have a good reason not to. Add support for other backends like CUDA in follow-up PRs
- In particular, adding new data types (extension of the `ggml_type` enum) carries with it a disproportionate maintenance burden. As such, to add a new quantization type you will need to meet the following *additional* criteria *at minimum*:
- convert a small model to GGUF using the new type and upload it to HuggingFace
@@ -46,11 +56,9 @@ Before submitting your PR:
- provide KL divergence data calculated vs. the FP16/BF16 (whichever is the native precision) version for both the new type as well as types of similar size
- provide [performance data](https://github.com/ggml-org/llama.cpp/tree/master/tools/llama-bench) for the new type in comparison to types of similar size on pure CPU
- Consider allowing write access to your branch for faster reviews, as reviewers can push commits directly
- If you are a new contributor
- Limit your open PRs to 1
- Do not submit trivial fixes (e.g. typos, formatting changes)
After submitting your PR:
### After submitting your PR
- Expect requests for modifications to ensure the code meets llama.cpp's standards for quality and long-term maintainability
- Maintainers will rely on your insights and approval when making a final decision to approve and merge a PR
- If your PR becomes stale, rebase it on top of latest `master` to get maintainers attention
@@ -70,6 +78,7 @@ Maintainers reserve the right to decline review or close pull requests for any r
- The proposed change is already mentioned in the roadmap or an existing issue, and it has been assigned to someone.
- The pull request duplicates an existing one.
- The contributor fails to adhere to this contributing guide or the AI policy.
- The change doesn't fit the existing architecture, or is too complex to justify its benefit.
# Coding guidelines
README.md
+71
View File
@@ -12,6 +12,77 @@
LLM inference in C/C++
## ⚡ This fork — Fable's MoE-offload prefill optimizations
Two **opt-in** optimizations for large MoE models whose experts are offloaded to system RAM
(`--n-cpu-moe`), found and implemented by Fable. Both are **off by default**, toggled via
environment variables, and produce **token-identical** output to mainline.
| Env var | What it does |
| --- | --- |
| `GGML_CUDA_REGISTER_HOST=1` | Page-locks (pins) the mmap'd CPU expert weights so host->device copies go straight over DMA instead of through the driver's hidden bounce buffer (~6-7 -> ~20 GB/s). Works on CUDA and Vulkan (also honored as `GGML_VK_REGISTER_HOST`). Note: it is a presence check, so `=0` still enables it. |
| `GGML_SCHED_PREFETCH_EXPERTS=1` | Prefetches each layer's experts on a second stream, so the weight uploads overlap compute instead of stalling the GPU. **CUDA only** - on the Vulkan backend the second backend instance shares one device queue, giving no overlap, so this regresses (see Vulkan note below). Leave it off on Vulkan. |
### Benchmark
Measured on an **RTX 3060 12GB** with **Qwen3.6-35B-A3B** (`--n-cpu-moe 26`), prompt-processing at 2048 (`MODEL` = path to your `.gguf`):
```bash
# baseline (patches off):
./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
# patched (both optimizations on):
GGML_CUDA_REGISTER_HOST=1 GGML_SCHED_PREFETCH_EXPERTS=1 \
./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
```
Result: **~1143 → ~1880 t/s** prefill (**+64%**) — same GPU, same settings, token-identical.
Branches: [`fable5/host-register`](https://github.com/thecodacus/llama.cpp/tree/fable5/host-register) (pinning only) · [`fable5/prefetch-experts`](https://github.com/thecodacus/llama.cpp/tree/fable5/prefetch-experts) (both — this branch).
### Vulkan (older AMD, e.g. RX 580 / Polaris)
On the Vulkan backend the CUDA-oriented flags above behave differently, and this fork adds a
Polaris-specific flash-attention fix. Findings on an **RX 580 8GB** (Polaris / GCN, PCIe 3.0 x16,
no fp16, no matrix cores) with **Qwen3.5-35B-A3B Q4_K_M**, `-b 2048 -ub 2048`:
- **Flash-attention `mask_opt` is now enabled for GCN large head sizes (this fork's own change).**
Upstream disables it on GCN; it is a **lossless** win in high-context prefill - it skips
fully-masked causal blocks and the per-block mask add on fully-visible ones, which is real work on
a card whose attention is compute-bound (no matrix cores). Auto-on, no flag. On Qwen3.5-35B
(head_dim 256): pp2048 **+8% @ 16k, +12% @ 32k**, growing with depth; perplexity bit-identical.
- **For a long-running server, load with `--no-mmap`, not pinning.** `GGML_CUDA_REGISTER_HOST=1`
(pinning) gives ~+17% in an isolated `llama-bench` run, but in a server the RADV host-pointer
import fails and the fallback pre-stage buffer allocation fails for large / co-resident models, so
it silently reverts to slow staging (and can trip warnings/OOM). `--no-mmap` (weights in RAM) is
both faster and clean there. Pinning is still fine for one-off `llama-bench` numbers.
- **`GGML_SCHED_PREFETCH_EXPERTS=1` regresses - do not use it** on Vulkan (its second backend shares
one device queue, so uploads never overlap compute).
- **`-b 2048 -ub 2048` is the biggest prefill lever** (the default `-ub 512` roughly halves pp).
- **Tune `--n-cpu-moe` to context length.** Keep some expert layers resident in spare VRAM for short
prompts (e.g. `ncmoe 28` on the 35B, ~+5% over all-host); at long context the KV cache needs that
VRAM, so raise it (`ncmoe 40`, all experts on host). Keep flash attention on (`-fa 1`).
Prefill throughput (isolated `llama-bench`, pinned unless noted):
| Config | pp2048 (t/s) |
| --- | ---: |
| baseline, unpinned, `ncmoe 40` | ~252 |
| pinned, `ncmoe 40` | ~294 |
| pinned, `ncmoe 28` | ~308 |
| server default (`--no-mmap`, `ncmoe 40`) | ~285 |
| + `mask_opt`, @ 32k context | **+12%** |
**Recommended RX 580 / Polaris serving command** (per model):
```bash
llama-server -hf <repo>:<quant> --no-mmap -ngl 99 --n-cpu-moe 40 -b 2048 -ub 2048 -fa 1
```
Lower `--n-cpu-moe` (e.g. 28) if the model plus your context budget leave spare VRAM; keep it high
for long-context / agentic use. At long context the bottleneck is attention compute (GPU-bound), so
`mask_opt` (above) is where the remaining prefill wins come from, not the MoE-transfer path.
## Recent API changes
- [Changelog for `libllama` API](https://github.com/ggml-org/llama.cpp/issues/9289)
common/CMakeLists.txt
+8
View File
@@ -100,6 +100,10 @@ add_library(${TARGET}
sampling.h
speculative.cpp
speculative.h
subproc.cpp
subproc.h
trie.cpp
trie.h
unicode.cpp
unicode.h
jinja/lexer.cpp
@@ -125,6 +129,10 @@ set_target_properties(${TARGET} PROPERTIES
target_include_directories(${TARGET} PUBLIC . ../vendor)
target_compile_features (${TARGET} PUBLIC cxx_std_17)
if (LLAMA_SUBPROCESS)
target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)
endif()
if (BUILD_SHARED_LIBS)
set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
common/arg.cpp
+62 -13
View File
@@ -5,6 +5,7 @@
#include "common.h"
#include "download.h"
#include "json-schema-to-grammar.h"
#include "llama.h"
#include "log.h"
#include "sampling.h"
#include "speculative.h"
@@ -785,6 +786,17 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
arg.c_str(), e.what(), opt.to_string().c_str()));
}
}
// TODO: remove this check after deprecating --mmap|mlock|dio
auto has_arg = [&](std::initializer_list<const char *> names) {
return std::any_of(names.begin(), names.end(), [&](const char * name) {
return seen_args.count(name);
});
};
if (has_arg({"-lm", "--load-mode"}) &&
has_arg({"--mlock", "--mmap", "--no-mmap", "-dio", "--direct-io", "-ndio", "--no-direct-io"})) {
LOG_WRN("DEPRECATED: `--load-mode` and `--mlock`/`--mmap`/`--direct-io` should not be combined; only the last flag on the command line will take effect\n");
}
};
// parse all CLI args now, so that -hf is available below for remote preset resolution
@@ -838,8 +850,9 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
params.kv_overrides.back().key[0] = 0;
}
if (!params.server_tools.empty() && !params.cors_origins_explicit) {
LOG_WRN("server tools are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
const bool mcp_enabled = !params.mcp_servers_config.empty() || !params.mcp_servers_json.empty();
if ((!params.server_tools.empty() || mcp_enabled) && !params.cors_origins_explicit) {
LOG_WRN("server tools or MCP servers are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
params.cors_origins = "localhost";
}
@@ -2495,27 +2508,45 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
}
add_opt(common_arg(
{"--mlock"},
"force system to keep model in RAM rather than swapping or compressing",
"DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing",
[](common_params & params) {
params.use_mlock = true;
LOG_WRN("DEPRECATED: --mlock is deprecated. use --load-mode mlock instead\n");
params.load_mode = LLAMA_LOAD_MODE_MLOCK;
}
).set_env("LLAMA_ARG_MLOCK"));
add_opt(common_arg(
{"--mmap"},
{"--no-mmap"},
string_format("whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock) (default: %s)", params.use_mmap ? "enabled" : "disabled"),
"DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)",
[](common_params & params, bool value) {
params.use_mmap = value;
LOG_WRN("DEPRECATED: --mmap and --no-mmap are deprecated. use --load-mode mmap instead\n");
params.load_mode = value ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE;
}
).set_env("LLAMA_ARG_MMAP"));
add_opt(common_arg(
{"-dio", "--direct-io"},
{"-ndio", "--no-direct-io"},
string_format("use DirectIO if available. (default: %s)", params.use_direct_io ? "enabled" : "disabled"),
"DEPRECATED in favor of `--load-mode`: use DirectIO if available",
[](common_params & params, bool value) {
params.use_direct_io = value;
LOG_WRN("DEPRECATED: --direct-io and --no-direct-io are deprecated. use --load-mode dio instead\n");
params.load_mode = value ? LLAMA_LOAD_MODE_DIRECT_IO : LLAMA_LOAD_MODE_NONE;
}
).set_env("LLAMA_ARG_DIO"));
add_opt(common_arg(
{"-lm", "--load-mode"}, "MODE",
"model loading mode (default: mmap)\n"
"- none: no special loading mode\n"
"- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n"
"- mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
"- dio: use DirectIO if available\n",
[](common_params & params, const std::string & value) {
/**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; }
else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; }
else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; }
else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; }
else { throw std::invalid_argument("invalid value"); }
}
).set_env("LLAMA_ARG_LOAD_MODE"));
add_opt(common_arg(
{"--numa"}, "TYPE",
"attempt optimizations that help on some NUMA systems\n"
@@ -2575,15 +2606,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
).set_env("LLAMA_ARG_CPU_MOE"));
add_opt(common_arg(
{"-ncmoe", "--n-cpu-moe"}, "N",
"keep the Mixture of Experts (MoE) weights of the first N layers in the CPU",
[](common_params & params, int value) {
if (value < 0) {
"keep the Mixture of Experts (MoE) weights of the first N layers in the CPU; "
"fractional N offloads part of the boundary layer at tensor granularity",
[](common_params & params, const std::string & value) {
const double n = std::stod(value);
if (n < 0) {
throw std::invalid_argument("invalid value");
}
for (int i = 0; i < value; ++i) {
for (const std::string & re : llm_ffn_exps_cpu_block_regexes(n)) {
// keep strings alive and avoid leaking memory by storing them in a static vector
static std::list<std::string> buft_overrides;
buft_overrides.push_back(llm_ffn_exps_block_regex(i));
buft_overrides.push_back(re);
params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
}
}
@@ -3231,6 +3264,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
params.server_tools = parse_csv_row(value);
}
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS"));
add_opt(common_arg(
{"--mcp-servers-config"}, "PATH",
"experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
"note: for security reasons, this will limit --cors-origins to localhost by default",
[](common_params & params, const std::string & value) {
params.mcp_servers_config = value;
}
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_CONFIG"));
add_opt(common_arg(
{"--mcp-servers-json"}, "JSON",
"experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
"note: for security reasons, this will limit --cors-origins to localhost by default",
[](common_params & params, const std::string & value) {
params.mcp_servers_json = value;
}
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_JSON"));
add_opt(common_arg(
{"-ag", "--agent"},
{"-no-ag", "--no-agent"},
common/chat.cpp
+28 -13
View File
@@ -1024,7 +1024,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
data.supports_thinking = true;
data.thinking_start_tag = "[THINK]";
data.thinking_end_tag = "[/THINK]";
data.thinking_end_tags = {"[/THINK]"};
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
@@ -1150,6 +1150,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
data.thinking_start_tag = "<|channel|>analysis<|message|>";
data.thinking_end_tags = {"<|end|>"};
// These special tokens are required to parse properly, so we include them
// even if parse_tool_calls is false.
data.preserved_tokens = {
@@ -1294,7 +1297,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ
data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4;
data.supports_thinking = true;
data.thinking_start_tag = "<|channel>thought";
data.thinking_end_tag = "<channel|>";
data.thinking_end_tags = {"<channel|>"};
data.preserved_tokens = {
"<|channel>",
@@ -1569,7 +1572,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp
const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>";
data.thinking_start_tag = THINK_START;
data.thinking_end_tag = THINK_END;
data.thinking_end_tags = {THINK_END};
if (inputs.has_continuation()) {
const auto & msg = inputs.continue_msg;
@@ -1703,7 +1706,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat
}
data.thinking_start_tag = THINK_START;
data.thinking_end_tag = THINK_END;
data.thinking_end_tags = {THINK_END};
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
@@ -1943,7 +1946,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
data.thinking_start_tag = "<think>";
data.thinking_end_tag = "</think>";
data.thinking_end_tags = {"</think>"};
data.preserved_tokens = {
"｜DSML｜",
"<think>",
@@ -2160,7 +2163,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
data.thinking_start_tag = THINK_START;
data.thinking_end_tag = THINK_END;
data.thinking_end_tags = {THINK_END};
data.preserved_tokens = {
TURN_START, TURN_END, CHATBOT, USER, SYSTEM,
THINK_START, THINK_END,
@@ -2179,9 +2182,10 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
{ COMMON_CHAT_ROLE_SYSTEM, TURN_START + SYSTEM },
};
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty();
auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
if (inputs.has_continuation()) {
const auto & msg = inputs.continue_msg;
@@ -2212,7 +2216,11 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
p.optional(p.literal(THINK_END))));
}
auto text_content = p.literal(TEXT_START) + p.content(p.until(TEXT_END)) + p.optional(p.literal(TEXT_END));
auto text_content = has_response_format
? p.literal(TEXT_START) +
p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) +
p.optional(p.literal(TEXT_END))
: p.literal(TEXT_START) + p.content(p.until(TEXT_END)) + p.optional(p.literal(TEXT_END));
if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
return generation_prompt + reasoning + text_content + p.optional(p.literal(TURN_END)) + end;
@@ -2240,13 +2248,17 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
data.parser = parser.save();
if (include_grammar) {
data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
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.at("parameters");
builder.resolve_refs(schema);
});
if (has_response_format) {
auto schema = inputs.json_schema;
builder.resolve_refs(schema);
}
parser.build_grammar(builder, data.grammar_lazy);
});
@@ -2501,7 +2513,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
};
data.thinking_start_tag = "<think>";
data.thinking_end_tag = "</think>";
data.thinking_end_tags = {"</think>"};
data.message_delimiters = {
{ COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" },
@@ -2857,7 +2869,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE;
if (auto_params.supports_thinking) {
auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start);
auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end);
auto end_tag = trim_whitespace(autoparser.reasoning.end);
if (!end_tag.empty()) {
auto_params.thinking_end_tags = {std::move(end_tag)};
}
}
common_peg_arena arena;
arena.load(auto_params.parser);
common/chat.h
+1 -1
View File
@@ -274,7 +274,7 @@ struct common_chat_params {
std::string generation_prompt;
bool supports_thinking = false;
std::string thinking_start_tag; // e.g., "<think>"
std::string thinking_end_tag; // e.g., "</think>"
std::vector<std::string> thinking_end_tags; // e.g., "</think>"
std::vector<common_grammar_trigger> grammar_triggers;
std::vector<std::string> preserved_tokens;
std::vector<std::string> additional_stops;
common/common.cpp
+1 -4
View File
@@ -1249,7 +1249,6 @@ common_init_result::common_init_result(common_params & params, bool model_only)
lora.reset(llama_adapter_lora_init(model, la.path.c_str()));
if (lora == nullptr) {
COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str());
pimpl->model.reset(model);
return;
}
@@ -1558,10 +1557,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
mparams.n_gpu_layers = params.n_gpu_layers;
mparams.main_gpu = params.main_gpu;
mparams.split_mode = params.split_mode;
mparams.load_mode = params.load_mode;
mparams.tensor_split = params.tensor_split;
mparams.use_mmap = params.use_mmap;
mparams.use_direct_io = params.use_direct_io;
mparams.use_mlock = params.use_mlock;
mparams.check_tensors = params.check_tensors;
mparams.use_extra_bufts = !params.no_extra_bufts;
mparams.no_host = params.no_host;
common/common.h
+34 -9
View File
@@ -6,6 +6,7 @@
#include "ggml-opt.h"
#include "ggml.h"
#include "llama.h"
#include <set>
#include <sstream>
@@ -283,12 +284,12 @@ struct common_params_sampling {
// reasoning budget sampler parameters
// these are populated by the server/CLI based on chat template params
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
std::vector<llama_token> reasoning_budget_start; // start tag token sequence
std::vector<llama_token> reasoning_budget_end; // end tag token sequence
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + end tag)
std::string reasoning_budget_message; // message injected before end tag when budget exhausted
bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
std::vector<llama_token> reasoning_budget_start; // start tag token sequence
std::vector<llama_tokens> reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + first end tag)
std::string reasoning_budget_message; // message injected before end tag when budget exhausted
bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
bool backend_sampling = false;
@@ -482,6 +483,7 @@ struct common_params {
std::vector<size_t> fit_params_target = std::vector<size_t>(llama_max_devices(), 1024 * 1024*1024);
enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs
enum llama_load_mode load_mode = LLAMA_LOAD_MODE_MMAP; // how to load the model
common_cpu_params cpuparams;
common_cpu_params cpuparams_batch;
@@ -572,9 +574,6 @@ struct common_params {
bool kv_unified = false; // enable unified KV cache
bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix
bool use_mmap = true; // enable mmap to use filesystem cache
bool use_direct_io = false; // read from disk without buffering
bool use_mlock = false; // use mlock to keep model in memory
bool verbose_prompt = false; // print prompt tokens before generation
bool display_prompt = true; // print prompt before generation
bool no_kv_offload = false; // disable KV offloading
@@ -669,6 +668,10 @@ struct common_params {
// enable built-in tools
std::vector<std::string> server_tools;
// MCP server configs (Cursor-compatible JSON)
std::string mcp_servers_config; // path to JSON file with MCP server definitions
std::string mcp_servers_json; // inline JSON with MCP server definitions
// router server configs
std::string models_dir = ""; // directory containing models for the router server
std::string models_preset = ""; // directory containing model presets for the router server
@@ -1080,6 +1083,28 @@ inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
}
// ATSInfer-style tensor-granularity static placement of MoE expert weights.
// Offloads the expert weights of the first floor(n) layers to the CPU, plus a
// subset of the boundary layer's three expert tensors for the fractional part.
// Expert tensors are dropped from the GPU in ascending performance-density
// order (gate, then up), keeping the higher-value down_proj resident longest.
inline std::vector<std::string> llm_ffn_exps_cpu_block_regexes(double n_cpu_moe) {
std::vector<std::string> regexes;
const int n_full = n_cpu_moe > 0 ? (int) n_cpu_moe : 0;
for (int i = 0; i < n_full; ++i) {
regexes.push_back(llm_ffn_exps_block_regex(i));
}
const int k = (int) ((n_cpu_moe - n_full) * 3.0 + 0.5);
if (k >= 3) {
regexes.push_back(llm_ffn_exps_block_regex(n_full));
} else if (k == 2) {
regexes.push_back(string_format("blk\\.%d\\.ffn_(gate|up)_(ch|)exps", n_full));
} else if (k == 1) {
regexes.push_back(string_format("blk\\.%d\\.ffn_gate_(ch|)exps", n_full));
}
return regexes;
}
//
// training utils
//
common/fit.cpp
+1 -2
View File
@@ -54,8 +54,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
llama_model_params mparams_copy = *mparams;
mparams_copy.no_alloc = true;
mparams_copy.use_mmap = false;
mparams_copy.use_mlock = false;
mparams_copy.load_mode = LLAMA_LOAD_MODE_NONE;
llama_model * model = llama_model_load_from_file(path_model, mparams_copy);
if (model == nullptr) {
common/peg-parser.cpp
+5 -153
View File
@@ -3,10 +3,10 @@
#include "common.h"
#include "json-schema-to-grammar.h"
#include "log.h"
#include "trie.h"
#include "unicode.h"
#include <algorithm>
#include <deque>
#include <initializer_list>
#include <map>
#include <memory>
@@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) {
return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F');
}
// Trie for matching multiple literals.
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
struct trie {
struct node {
std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
bool is_word;
};
std::vector<node> nodes;
trie(const std::vector<std::string> & words) {
create_node(); // root node
for (const auto & w : words) {
insert(w);
}
}
enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
// Check if a delimiter starts at the given position
match_result check_at(std::string_view sv, size_t start_pos) const {
size_t current = 0; // Start at root
size_t pos = start_pos;
// LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
while (pos < sv.size()) {
auto result = common_parse_utf8_codepoint(sv, pos);
if (result.status != utf8_parse_result::SUCCESS) {
break;
}
auto it = nodes[current].children.find(result.codepoint);
if (it == nodes[current].children.end()) {
// Can't continue matching
return match_result{match_result::NO_MATCH};
}
current = it->second;
pos += result.bytes_consumed;
// Check if we've matched a complete word
if (nodes[current].is_word) {
return match_result{match_result::COMPLETE_MATCH};
}
}
// Reached end of input while still in the trie (not at root)
if (current != 0) {
// We're in the middle of a potential match
return match_result{match_result::PARTIAL_MATCH};
}
// Reached end at root (no match)
return match_result{match_result::NO_MATCH};
}
private:
size_t create_node() {
size_t index = nodes.size();
nodes.emplace_back();
return index;
}
void insert(const std::string & word) {
size_t current = 0;
size_t pos = 0;
while (pos < word.length()) {
auto result = common_parse_utf8_codepoint(word, pos);
if (result.status != utf8_parse_result::SUCCESS) {
break;
}
uint32_t ch = result.codepoint;
pos += result.bytes_consumed;
auto it = nodes[current].children.find(ch);
if (it == nodes[current].children.end()) {
size_t child = create_node();
nodes[current].children[ch] = child;
current = child;
} else {
current = it->second;
}
}
nodes[current].is_word = true;
}
};
// Aho-Corasick automaton
struct aho_corasick {
trie t;
std::vector<size_t> fail; // failure links
std::vector<size_t> order; // states in BFS order
std::vector<bool> terminal; // match states (directly or via a suffix link)
std::set<uint32_t> alphabet; // every character with a transition
aho_corasick(const std::vector<std::string> & strings) : t(strings) {
const auto & nodes = t.nodes;
const size_t n = nodes.size();
fail.assign(n, 0);
order.reserve(n);
std::deque<size_t> queue{ 0 };
while (!queue.empty()) {
size_t u = queue.front();
queue.pop_front();
order.push_back(u);
for (const auto & [ch, v] : nodes[u].children) {
if (u != 0) {
size_t f = fail[u];
while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
f = fail[f];
}
auto it = nodes[f].children.find(ch);
fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
}
queue.push_back(v);
}
}
terminal.assign(n, false);
for (size_t u : order) {
terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]);
}
for (const auto & node : nodes) {
for (const auto & [ch, v] : node.children) {
alphabet.insert(ch);
}
}
}
size_t num_states() const { return t.nodes.size(); }
bool is_terminal(size_t s) const { return terminal[s]; }
// follow failure links until a transition on `ch` exists.
size_t next(size_t state, uint32_t ch) const {
const auto & nodes = t.nodes;
while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
state = fail[state];
}
auto it = nodes[state].children.find(ch);
return it != nodes[state].children.end() ? it->second : 0;
}
};
static std::pair<uint32_t, size_t> parse_hex_escape(const std::string & str, size_t pos, int hex_count) {
if (pos + hex_count > str.length()) {
return {0, 0};
@@ -797,7 +649,7 @@ struct parser_executor {
}
common_peg_parse_result operator()(const common_peg_until_parser & p) const {
trie matcher(p.delimiters);
common_trie matcher(p.delimiters);
// Scan input and check for delimiters
size_t pos = start_pos;
@@ -824,12 +676,12 @@ struct parser_executor {
// Check if a delimiter starts at this position
auto match = matcher.check_at(ctx.input, pos);
if (match == trie::COMPLETE_MATCH) {
if (match == common_trie::COMPLETE_MATCH) {
// Found a complete delimiter, return everything before it
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
}
if (match == trie::PARTIAL_MATCH) {
if (match == common_trie::PARTIAL_MATCH) {
// Found a partial match extending to end of input, return everything before it
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
}
@@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar(
const std::map<size_t, std::vector<uint32_t>> &,
const std::vector<uint32_t> &,
const std::function<std::string(size_t)> &)> & build_rule) {
aho_corasick ac(strings);
common_aho_corasick ac(strings);
auto state_name = [&](size_t s) -> std::string {
if (s == 0) {
common/preset.cpp
+4
View File
@@ -330,6 +330,10 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co
}
}
if (preset.name == COMMON_PRESET_DEFAULT_NAME && preset.options.empty()) {
continue;
}
if (preset.name == "*") {
// handle global preset
global = preset;
common/reasoning-budget.cpp
+77 -39
View File
@@ -1,39 +1,52 @@
#include "reasoning-budget.h"
#include "common.h"
#include "trie.h"
#include "unicode.h"
#include "log.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <string>
#include <vector>
struct token_matcher {
std::vector<llama_token> tokens;
size_t pos = 0;
std::vector<llama_tokens> seqs;
common_aho_corasick ac;
size_t state = 0;
bool advance(llama_token token) {
if (tokens.empty()) {
return false;
}
token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
if (token == tokens[pos]) {
pos++;
if (pos >= tokens.size()) {
pos = 0;
return true;
}
} else {
pos = 0;
if (token == tokens[0]) {
pos = 1;
static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
std::vector<llama_tokens> res;
for (const auto & seq : seqs) {
if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
res.push_back(seq);
}
}
return false;
return res;
}
void reset() { pos = 0; }
static common_trie build_trie(const std::vector<llama_tokens> & seqs) {
common_trie t;
for (const auto & seq : seqs) {
t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));
}
return t;
}
// returns the index into seqs of the longest sequence ending at this token, or -1
int32_t advance(llama_token token) {
state = ac.next(state, (uint32_t) token);
const int32_t p = ac.match_pattern(state);
if (p >= 0) {
state = 0;
}
return p;
}
void reset() { state = 0; }
};
struct common_reasoning_budget_ctx {
@@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx {
token_matcher start_matcher;
token_matcher end_matcher;
std::vector<llama_token> forced_tokens;
llama_tokens forced_tokens;
int32_t budget; // maximum tokens in reasoning block
int32_t remaining; // tokens remaining in budget
@@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx {
// for forcing
size_t force_pos; // next position in forced_tokens to force
int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none
};
static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) {
@@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
switch (ctx->state) {
case REASONING_BUDGET_IDLE:
{
if (ctx->start_matcher.advance(token)) {
if (ctx->start_matcher.advance(token) >= 0) {
ctx->state = REASONING_BUDGET_COUNTING;
ctx->remaining = ctx->budget;
COM_TRC("activated, budget=%d tokens\n", ctx->budget);
@@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
case REASONING_BUDGET_COUNTING:
case REASONING_BUDGET_WAITING_UTF8:
{
if (ctx->end_matcher.advance(token)) {
const int32_t match = ctx->end_matcher.advance(token);
if (match >= 0) {
ctx->state = REASONING_BUDGET_DONE;
ctx->end_match = match;
COM_TRC("%s", "deactivated (natural end)\n");
break;
}
@@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
break;
}
case REASONING_BUDGET_FORCING:
{
// track the end sequence within forced_tokens so it is also reported on DONE
const int32_t match = ctx->end_matcher.advance(token);
ctx->force_pos++;
if (ctx->force_pos >= ctx->forced_tokens.size()) {
ctx->state = REASONING_BUDGET_DONE;
ctx->end_match = match;
COM_TRC("%s", "forced sequence complete, done\n");
}
break;
}
case REASONING_BUDGET_DONE:
// Re-arm on a new start tag: some models emit multiple <think> blocks
// per response, and each should get a fresh budget window.
if (ctx->start_matcher.advance(token)) {
if (ctx->start_matcher.advance(token) >= 0) {
ctx->state = REASONING_BUDGET_COUNTING;
ctx->remaining = ctx->budget;
ctx->end_matcher.reset();
ctx->end_match = -1;
COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
if (ctx->remaining <= 0) {
@@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
ctx->start_matcher.reset();
ctx->end_matcher.reset();
ctx->force_pos = 0;
ctx->end_match = -1;
}
static struct llama_sampler * common_reasoning_budget_init_state(
const struct llama_vocab * vocab, const std::vector<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens, const std::vector<llama_token> & forced_tokens,
const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,
int32_t budget, common_reasoning_budget_state initial_state);
static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);
@@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s
}
static struct llama_sampler * common_reasoning_budget_init_state(
const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens,
const std::vector<llama_token> & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state) {
const struct llama_vocab * vocab,
const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & end_seqs,
const llama_tokens & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state) {
// promote COUNTING with budget <= 0 to FORCING
if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {
initial_state = REASONING_BUDGET_FORCING;
@@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state(
/* .iface = */ &common_reasoning_budget_i,
/* .ctx = */ new common_reasoning_budget_ctx {
/* .vocab = */ vocab,
/* .start_matcher = */ { start_tokens, 0 },
/* .end_matcher = */ { end_tokens, 0 },
/* .start_matcher = */ token_matcher(start_seqs),
/* .end_matcher = */ token_matcher(end_seqs),
/* .forced_tokens = */ forced_tokens,
/* .budget = */ budget,
/* .remaining = */ budget,
/* .state = */ initial_state,
/* .force_pos = */ 0,
/* .end_match = */ -1,
}
);
}
struct llama_sampler * common_reasoning_budget_init(
const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens,
const std::vector<llama_token> & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state) {
return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state);
const struct llama_vocab * vocab,
const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & end_seqs,
const llama_tokens & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state) {
return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state);
}
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
@@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla
return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;
}
const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) {
if (!smpl) {
return nullptr;
}
const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
if (ctx->end_match < 0) {
return nullptr;
}
return &ctx->end_matcher.seqs[ctx->end_match];
}
bool common_reasoning_budget_force(struct llama_sampler * smpl) {
if (!smpl) {
return false;
common/reasoning-budget.h
+16 -10
View File
@@ -2,6 +2,8 @@
#include "llama.h"
#include "common.h"
#include <cstdint>
#include <vector>
@@ -17,30 +19,34 @@ enum common_reasoning_budget_state {
// reasoning block (e.g. between <think> and </think>).
//
// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE
// IDLE: passthrough, watching for start_tokens sequence
// COUNTING: counting down remaining tokens, watching for natural end_tokens
// IDLE: passthrough, watching for a start sequence
// COUNTING: counting down remaining tokens, watching for a natural end sequence
// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence
// FORCING: forces forced_tokens token-by-token (all other logits -> -inf)
// DONE: passthrough forever
//
// Parameters:
// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr)
// start_tokens - token sequence that activates counting
// end_tokens - token sequence for natural deactivation
// start_seqs - token sequences, any of which activates counting
// end_seqs - token sequences, any of which naturally deactivates
// forced_tokens - token sequence forced when budget expires
// budget - max tokens allowed in the reasoning block
// initial_state - initial state
//
struct llama_sampler * common_reasoning_budget_init(
const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens,
const std::vector<llama_token> & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
const struct llama_vocab * vocab,
const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & end_seqs,
const llama_tokens & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl);
// The end sequence that transitioned the sampler to DONE, or nullptr if none
// was recorded. Cleared when a new start sequence re-arms the sampler.
const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl);
// Manually transition the reasoning budget sampler into the FORCING state.
// Returns true if the transition occurred.
bool common_reasoning_budget_force(struct llama_sampler * smpl);
common/sampling.cpp
+12 -1
View File
@@ -299,7 +299,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st
if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) {
rbudget = common_reasoning_budget_init(
vocab,
params.reasoning_budget_start,
{params.reasoning_budget_start},
params.reasoning_budget_end,
params.reasoning_budget_forced,
params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens);
@@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo
if (gsmpl->rbudget && is_generated) {
llama_sampler_accept(gsmpl->rbudget, token);
// if done, replay end sequence which may contain a grammar trigger
const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE;
if (gsmpl->grmr && !accept_grammar && is_done) {
const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget);
if (end_seq) {
for (const llama_token end_token : *end_seq) {
llama_sampler_accept(gsmpl->grmr, end_token);
}
}
}
}
if (gsmpl->grmr && accept_grammar) {
common/subproc.cpp
+143
View File
@@ -0,0 +1,143 @@
#include "subproc.h"
bool common_subproc::is_supported() {
#ifdef LLAMA_SUBPROCESS
return true;
#else
return false;
#endif
}
#ifdef LLAMA_SUBPROCESS
static std::vector<char *> to_cstr_vec(const std::vector<std::string> & v) {
std::vector<char *> r;
r.reserve(v.size() + 1);
for (const auto & s : v) {
r.push_back(const_cast<char *>(s.c_str()));
}
r.push_back(nullptr);
return r;
}
common_subproc::~common_subproc() {
if (is_created) {
subprocess_destroy(&proc);
is_created = false;
}
}
bool common_subproc::create(
const std::vector<std::string> & args,
int options,
const std::vector<std::string> & env,
const char * cwd) {
auto argv = to_cstr_vec(args);
int result;
if (env.empty() && cwd == nullptr) {
result = subprocess_create(argv.data(), options, &proc);
} else {
auto envp = to_cstr_vec(env);
result = subprocess_create_ex(argv.data(), options, env.empty() ? nullptr : envp.data(), cwd, &proc);
}
is_created = result == 0;
return is_created;
}
bool common_subproc::has_handle() const {
if (!is_created) {
return false;
}
#if defined(_WIN32)
return proc.hProcess != nullptr;
#else
return proc.child > 0;
#endif
}
bool common_subproc::alive() {
return is_created && subprocess_alive(&proc);
}
FILE * common_subproc::stdin_file() {
return is_created ? subprocess_stdin(&proc) : nullptr;
}
FILE * common_subproc::stdout_file() {
return is_created ? subprocess_stdout(&proc) : nullptr;
}
FILE * common_subproc::stderr_file() {
return is_created ? subprocess_stderr(&proc) : nullptr;
}
void common_subproc::close_stdin() {
if (is_created && proc.stdin_file) {
fclose(proc.stdin_file);
proc.stdin_file = nullptr;
}
}
void common_subproc::terminate() {
if (has_handle()) {
subprocess_terminate(&proc);
}
}
int common_subproc::join() {
int exit_code = -1;
if (is_created) {
subprocess_join(&proc, &exit_code);
subprocess_destroy(&proc);
is_created = false;
}
return exit_code;
}
#else // !LLAMA_SUBPROCESS
common_subproc::~common_subproc() = default;
bool common_subproc::create(
const std::vector<std::string> &,
int,
const std::vector<std::string> &,
const char *) {
(void)(proc);
(void)(is_created);
return false;
}
bool common_subproc::has_handle() const {
return false;
}
bool common_subproc::alive() {
return false;
}
FILE * common_subproc::stdin_file() {
return nullptr;
}
FILE * common_subproc::stdout_file() {
return nullptr;
}
FILE * common_subproc::stderr_file() {
return nullptr;
}
void common_subproc::close_stdin() {
}
void common_subproc::terminate() {
}
int common_subproc::join() {
return -1;
}
#endif // LLAMA_SUBPROCESS
common/subproc.h
+59
View File
@@ -0,0 +1,59 @@
#pragma once
#include <atomic>
#include <cstdio>
#include <string>
#include <vector>
#ifdef LLAMA_SUBPROCESS
#include <sheredom/subprocess.h>
#else
// dummy values to allow compilation when subprocess is disabled
struct subprocess_s {};
static constexpr int subprocess_option_no_window = 0;
static constexpr int subprocess_option_combined_stdout_stderr = 0;
static constexpr int subprocess_option_inherit_environment = 0;
static constexpr int subprocess_option_search_user_path = 0;
#endif
// RAII-style wrapper around https://github.com/sheredom/subprocess.h,
// exposing method calls instead of free functions operating on subprocess_s.
struct common_subproc {
common_subproc() = default;
~common_subproc();
common_subproc(const common_subproc &) = delete;
common_subproc & operator=(const common_subproc &) = delete;
// spawn a child process; if env is non-empty it replaces the child's environment
// (do not combine with subprocess_option_inherit_environment)
bool create(
const std::vector<std::string> & args,
int options,
const std::vector<std::string> & env = {},
const char * cwd = nullptr);
bool alive();
// true if LLAMA_SUBPROCESS was enabled at build time; when false, create() always fails
static bool is_supported();
FILE * stdin_file();
FILE * stdout_file();
FILE * stderr_file();
// close stdin and detach it from the process, so a later join()/destroy() won't double-close it;
// use this after writing all input to signal EOF to the child while it's still running
void close_stdin();
void terminate();
// wait for the process to exit, release the underlying handle and return its exit code
int join();
private:
subprocess_s proc {};
std::atomic<bool> is_created{false};
bool has_handle() const;
};
common/trie.cpp
+123
View File
@@ -0,0 +1,123 @@
#include "trie.h"
#include "unicode.h"
#include <deque>
common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const {
size_t current = 0; // Start at root
size_t pos = start_pos;
// LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
while (pos < sv.size()) {
auto result = common_parse_utf8_codepoint(sv, pos);
if (result.status != utf8_parse_result::SUCCESS) {
break;
}
auto it = nodes[current].children.find(result.codepoint);
if (it == nodes[current].children.end()) {
// Can't continue matching
return match_result{match_result::NO_MATCH};
}
current = it->second;
pos += result.bytes_consumed;
// Check if we've matched a complete word
if (nodes[current].pattern >= 0) {
return match_result{match_result::COMPLETE_MATCH};
}
}
// Reached end of input while still in the trie (not at root)
if (current != 0) {
// We're in the middle of a potential match
return match_result{match_result::PARTIAL_MATCH};
}
// Reached end at root (no match)
return match_result{match_result::NO_MATCH};
}
int32_t common_trie::insert(const std::string & word) {
std::vector<uint32_t> symbols;
size_t pos = 0;
while (pos < word.length()) {
auto result = common_parse_utf8_codepoint(word, pos);
if (result.status != utf8_parse_result::SUCCESS) {
break;
}
symbols.push_back(result.codepoint);
pos += result.bytes_consumed;
}
return insert(symbols);
}
int32_t common_trie::insert(const std::vector<uint32_t> & symbols) {
size_t current = 0;
for (uint32_t ch : symbols) {
auto it = nodes[current].children.find(ch);
if (it == nodes[current].children.end()) {
size_t child = create_node();
nodes[current].children[ch] = child;
current = child;
} else {
current = it->second;
}
}
if (nodes[current].pattern < 0) {
nodes[current].pattern = n_patterns++;
}
return nodes[current].pattern;
}
common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) {
const auto & nodes = t.nodes;
const size_t n = nodes.size();
fail.assign(n, 0);
order.reserve(n);
std::deque<size_t> queue{ 0 };
while (!queue.empty()) {
size_t u = queue.front();
queue.pop_front();
order.push_back(u);
for (const auto & [ch, v] : nodes[u].children) {
if (u != 0) {
size_t f = fail[u];
while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
f = fail[f];
}
auto it = nodes[f].children.find(ch);
fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
}
queue.push_back(v);
}
}
// fail[u] points to a strictly shorter suffix, so the first pattern found on
// the fail chain (including u itself) is the longest pattern ending at u
match.assign(n, -1);
for (size_t u : order) {
match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1);
}
for (const auto & node : nodes) {
for (const auto & [ch, v] : node.children) {
alphabet.insert(ch);
}
}
}
size_t common_aho_corasick::next(size_t state, uint32_t ch) const {
const auto & nodes = t.nodes;
while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
state = fail[state];
}
auto it = nodes[state].children.find(ch);
return it != nodes[state].children.end() ? it->second : 0;
}
common/trie.h
+73
View File
@@ -0,0 +1,73 @@
#pragma once
#include <cstdint>
#include <map>
#include <set>
#include <string>
#include <string_view>
#include <vector>
// Trie for matching multiple literals.
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
struct common_trie {
struct node {
std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
int32_t pattern = -1; // index of the pattern ending at this node, -1 if none
};
std::vector<node> nodes;
common_trie() {
create_node(); // root node
}
common_trie(const std::vector<std::string> & words) : common_trie() {
for (const auto & w : words) {
insert(w);
}
}
enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
// Check if a delimiter starts at the given position
match_result check_at(std::string_view sv, size_t start_pos) const;
// Insert a word as a sequence of Unicode codepoints, returns its pattern index
int32_t insert(const std::string & word);
// Insert a raw symbol sequence, returns its pattern index (insertion order,
// duplicates keep the first index)
int32_t insert(const std::vector<uint32_t> & symbols);
private:
int32_t n_patterns = 0;
size_t create_node() {
size_t index = nodes.size();
nodes.emplace_back();
return index;
}
};
// Aho-Corasick automaton
struct common_aho_corasick {
common_trie t;
std::vector<size_t> fail; // failure links
std::vector<size_t> order; // states in BFS order
std::vector<int32_t> match; // longest pattern ending at each state (directly or via a suffix link), -1 if none
std::set<uint32_t> alphabet; // every character with a transition
common_aho_corasick(common_trie trie);
common_aho_corasick(const std::vector<std::string> & strings)
: common_aho_corasick(common_trie(strings)) {}
size_t num_states() const { return t.nodes.size(); }
bool is_terminal(size_t s) const { return match[s] >= 0; }
// index of the longest pattern ending at this state, -1 if none
int32_t match_pattern(size_t s) const { return match[s]; }
// follow failure links until a transition on `ch` exists.
size_t next(size_t state, uint32_t ch) const;
};
conversion/__init__.py
+3
View File
@@ -158,6 +158,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
"MiniCPMForCausalLM": "minicpm",
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
"MiniMaxM2ForCausalLM": "minimax",
"MiniMaxM3SparseForCausalLM": "minimax",
"MiniMaxM3SparseForConditionalGeneration": "minimax",
"Ministral3ForCausalLM": "mistral3",
"Mistral3ForConditionalGeneration": "mistral3",
"MistralForCausalLM": "llama",
@@ -267,6 +269,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"Gemma4UnifiedForConditionalGeneration": "gemma",
"Glm4vForConditionalGeneration": "qwen3vl",
"Glm4vMoeForConditionalGeneration": "qwen3vl",
"Glm5vForConditionalGeneration": "kimivl",
"GlmOcrForConditionalGeneration": "qwen3vl",
"GlmasrModel": "ultravox",
"Granite4VisionForConditionalGeneration": "granite",
conversion/base.py
+2 -2
View File
@@ -1156,7 +1156,7 @@ class TextModel(ModelBase):
or "projector." in name or "pre_mm_projector_norm" in name \
or "image_newline" in name or "view_seperator" in name \
or "patch_embed" in name or "patch_embedding" in name \
or "patch_merger." in name or "model.connector." in name:
or "patch_merger." in name or "patch_merge_mlp." in name or "model.connector." in name:
return None
return super().filter_tensors(item)
@@ -1203,7 +1203,7 @@ class TextModel(ModelBase):
self.gguf_writer.add_embedding_length(n_embd)
logger.info(f"gguf: embedding length = {n_embd}")
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
self.gguf_writer.add_feed_forward_length(n_ff)
logger.info(f"gguf: feed forward length = {n_ff}")
conversion/bert.py
+2 -1
View File
@@ -369,12 +369,13 @@ class NomicBertModel(BertModel):
return super().filter_tensors(item)
def modify_tensors(self, data_torch: torch.Tensor, name: str, bid: int | None) -> Iterable[tuple[str, torch.Tensor]]:
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
if "mlp.experts.mlp.w1" in name:
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
data_torch = data_torch.view(n_experts, self.hparams["n_inner"], self.hparams["n_embd"])
name += ".weight"
if "mlp.experts.mlp.w2" in name:
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
data_torch = data_torch.view(n_experts, self.hparams["n_inner"], self.hparams["n_embd"])
data_torch = data_torch.transpose(1, 2)
name += ".weight"
conversion/glm.py
+3
View File
@@ -237,6 +237,9 @@ class GlmMoeDsaModel(DeepseekV2Model):
self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
if (indexer_types := self.hparams.get("indexer_types")) is not None:
indexer_types = [t == "full" for t in indexer_types]
self.gguf_writer.add_indexer_types(indexer_types)
@ModelBase.register("SolarOpenForCausalLM")
conversion/kimivl.py
+16
View File
@@ -152,3 +152,19 @@ class KimiK25Model(MmprojModel):
name = name.replace(".proj.2.", ".proj.linear_2.")
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Glm5vForConditionalGeneration")
class Glm5vModel(KimiK25Model):
"""GLM-5.2-Vision MoonViT3d encoder and projector
Uses the same vision encoder and projector as Kimi-K2.5, so it reuses the
kimik25 projector type. The image begin/end tokens differ, but they are
resolved at runtime from the text model vocab.
"""
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name.startswith("mm_projector.linear_"):
name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1)
yield from super().modify_tensors(data_torch, name, bid)
conversion/minimax.py
+36 -1
View File
@@ -23,7 +23,7 @@ class MiniMaxM2Model(TextModel):
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
# merge expert weights
if 'experts' in name:
if "block_sparse_moe.experts." in name:
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
assert bid is not None
@@ -52,3 +52,38 @@ class MiniMaxM2Model(TextModel):
return
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration")
class MiniMaxM3Model(MiniMaxM2Model):
model_arch = gguf.MODEL_ARCH.MINIMAXM3
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_expert_shared_count(self.find_hparam(["n_shared_experts"]))
self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"]))
self.gguf_writer.add_expert_weights_norm(True)
sac = self.find_hparam(["sparse_attention_config"])
self.gguf_writer.add_indexer_head_count(sac["sparse_num_index_heads"])
self.gguf_writer.add_indexer_key_length(sac["sparse_index_dim"])
self.gguf_writer.add_indexer_top_k(sac["sparse_topk_blocks"])
self.gguf_writer.add_indexer_block_size(sac["sparse_block_size"])
self.gguf_writer.add_indexer_local_blocks(sac["sparse_local_block"])
moe_layer_freq = self.find_hparam(["moe_layer_freq"])
n_dense = 0
for v in moe_layer_freq:
if v == 0:
n_dense += 1
else:
break
self.gguf_writer.add_leading_dense_block_count(n_dense)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
# Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm
if name.endswith("norm.weight"):
data_torch = data_torch + 1.0
yield from super().modify_tensors(data_torch, name, bid)
docs/backend/OPENCL.md
+18
View File
@@ -98,6 +98,24 @@ The OpenCL backend has the following CMake options that control the behavior of
| `GGML_OPENCL_USE_ADRENO_KERNELS` | `ON` | Use kernels optimized for Adreno. |
| `GGML_OPENCL_USE_ADRENO_BIN_KERNELS` | `OFF` | Allow using binary kernel lib for Adreno. |
## Program Binary Cache
Compiled `cl_program` binaries are cached on disk, so subsequent runs skip the expensive
compile-from-source step when nothing relevant has changed (kernel source, compile options,
device, driver, or platform version).
The cache is controlled with the `GGML_OPENCL_KERNEL_CACHE_DIR` environment variable:
| Value | Behavior |
|:---------------------------------------|:-----------------------------------------------|
| unset / empty / `1` / `default` | Enabled in the platform default cache directory: `%LOCALAPPDATA%\llama.cpp\cl-cache` (Windows), `~/Library/Caches/llama.cpp/cl-cache` (macOS), `<temp dir>/llama.cpp/cl-cache` elsewhere. |
| `0` / `off` / `none` / `disable(d)` | Disabled. |
| any other value | Used verbatim as the cache directory path. |
If the chosen directory cannot be created or used, the cache disables itself for the process
and kernels are compiled from source as usual. Set `GGML_OPENCL_KERNEL_CACHE_DEBUG=1` to
print a HIT/MISS/SAVE trace to stderr.
## Android
Ubuntu 22.04 is used for targeting Android. Make sure the following tools are accessible from command line,
docs/build.md
-6
View File
@@ -361,12 +361,6 @@ You can download it from your Linux distro's package manager or from here: [ROCm
Note: `GPU_TARGETS` is optional, omitting it will build the code for all GPUs in the current system.
To enhance flash attention performance on RDNA3+ or CDNA architectures, you can utilize the rocWMMA library by enabling the `-DGGML_HIP_ROCWMMA_FATTN=ON` option. This requires rocWMMA headers to be installed on the build system.
The rocWMMA library is included by default when installing the ROCm SDK using the `rocm` meta package provided by AMD. Alternatively, if you are not using the meta package, you can install the library using the `rocwmma-dev` or `rocwmma-devel` package, depending on your system's package manager.
As an alternative, you can manually install the library by cloning it from the official [GitHub repository](https://github.com/ROCm/rocWMMA), checkout the corresponding version tag (e.g. `rocm-6.2.4`) and set `-DCMAKE_CXX_FLAGS="-I<path/to/rocwmma>/library/include/"` in CMake. This also works under Windows despite not officially supported by AMD.
Note that if you get the following error:
```
clang: error: cannot find ROCm device library; provide its path via '--rocm-path' or '--rocm-device-lib-path', or pass '-nogpulib' to build without ROCm device library
docs/development/HOWTO-add-model.md
+11
View File
@@ -45,6 +45,8 @@ class MyModel(MmprojModel):
Add an enum entry in `MODEL_ARCH`, the model human friendly name in `MODEL_ARCH_NAMES` and the GGUF tensor names in `MODEL_TENSORS`.
NOTE: Pick the GGUF arch string (and the matching `src/models/<name>.cpp` filename, see section 3) carefully up front, following existing naming conventions. Once GGUF files are published under a given arch string, renaming it later breaks the community's existing files, so this is not something to leave for cleanup in a follow-up PR.
Example for `falcon` model:
```python
MODEL_ARCH.FALCON: [
@@ -101,6 +103,7 @@ The model params and tensors layout must be defined in `llama.cpp` source files:
- You may also need to update `LLM_KV_NAMES`, `LLM_TENSOR_NAMES` and `LLM_TENSOR_INFOS`
3. Add any non-standard metadata loading in the `llama_model_loader` constructor in `src/llama-model-loader.cpp`.
4. If the model has a RoPE operation, add a case for the architecture in `llama_model_rope_type` function in `src/llama-model.cpp`.
5. Check for other places that switch/iterate over every `llm_arch` value, e.g. `src/llama-model-saver.cpp` and any mandatory-hparam lists (such as which archs require MoE metadata). Grep for `LLM_ARCH_` usages to find them. Missing one of these is a common cause of CI test failures (e.g. `test-llama-archs`) after adding a new arch.
NOTE: The dimensions in `ggml` are typically in the reverse order of the `pytorch` dimensions.
@@ -133,6 +136,14 @@ Note:
## Tips and tricks
### Prefer conversion-time tensor modifications over graph-time ones
If the model contains constant modifications of tensors in the graph (for example, `norm(1 + weight)`) or performs tensor permutations/chunking, perform the modifications during conversion rather than in the graph code. This keeps the inference graph simpler and avoids extra runtime ops.
Examples:
- Gemma 3 folds the `1 +` of its `norm(1 + weight)` normalization into the weights at conversion time, so the graph just does a plain RMS norm.
- Qwen3-Next applies its tensor permutation during conversion (in `modify_tensors`), so the graph can consume the already-permuted weights directly.
### Working with ggml_rope_ext
PyTorch implementations usually prefer explicitly calculating `freq_cis`/`sin`/`cos` components. However, in llama.cpp, most RoPE operations can be handled via `ggml_rope_ext`, which does not require a sin/cos matrix. This saves memory while allowing the GGML RoPE kernel to be fused with other ops.
examples/diffusion/diffusion-cli.cpp
+1 -3
View File
@@ -117,9 +117,7 @@ int main(int argc, char ** argv) {
llama_model_params model_params = llama_model_default_params();
model_params.n_gpu_layers = params.n_gpu_layers;
model_params.devices = params.devices.data();
model_params.use_mmap = params.use_mmap;
model_params.use_direct_io = params.use_direct_io;
model_params.use_mlock = params.use_mlock;
model_params.load_mode = params.load_mode;
model_params.check_tensors = params.check_tensors;
llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params);
examples/training/finetune.cpp
+3 -4
View File
@@ -26,10 +26,9 @@ int main(int argc, char ** argv) {
return 1;
}
if (params.use_mmap) {
LOG_INF("%s: force disabling memory mapping because it would result in-read-only pointers to the weights\n",
__func__);
params.use_mmap = false;
if (params.load_mode != LLAMA_LOAD_MODE_NONE) {
LOG_INF("%s: forcing load_mode = none to enable writable pointers to the weights\n", __func__);
params.load_mode = LLAMA_LOAD_MODE_NONE;
}
if (params.cache_type_k != GGML_TYPE_F32) {
LOG_INF("%s: force changing k cache type to f32 due to a lack of f16 support for OUT_PROD\n", __func__);
ggml/CMakeLists.txt
-1
View File
@@ -216,7 +216,6 @@ option(GGML_HIP "ggml: use HIP"
option(GGML_HIP_GRAPHS "ggml: use HIP graph" ON)
option(GGML_HIP_RCCL "ggml: use ROCm Collective Comm. Library" OFF)
option(GGML_HIP_NO_VMM "ggml: do not try to use HIP VMM" ON)
option(GGML_HIP_ROCWMMA_FATTN "ggml: enable rocWMMA for FlashAttention" OFF)
option(GGML_HIP_MMQ_MFMA "ggml: enable MFMA MMA for CDNA in MMQ" ON)
option(GGML_HIP_EXPORT_METRICS "ggml: enable kernel perf metrics output" OFF)
option(GGML_MUSA_GRAPHS "ggml: use MUSA graph, experimental, unstable" OFF)
ggml/include/ggml-backend.h
+5
View File
@@ -104,6 +104,11 @@ extern "C" {
GGML_API enum ggml_status ggml_backend_graph_compute (ggml_backend_t backend, struct ggml_cgraph * cgraph);
GGML_API enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph);
// Free transient/scratch device memory the backend holds outside of any allocated buffer
// (compute preallocations, staging buffers). No-op if the backend does not implement it.
// The backend remains usable; scratch is reallocated lazily on the next compute.
GGML_API void ggml_backend_free_scratch(ggml_backend_t backend);
// NOTE: will be removed, use device version instead
GGML_API bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op);
GGML_API bool ggml_backend_supports_buft(ggml_backend_t backend, ggml_backend_buffer_type_t buft);
ggml/src/ggml-backend-impl.h
+5
View File
@@ -137,6 +137,11 @@ extern "C" {
// (optional) sort/optimize the nodes in the graph
void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph);
// (optional) free transient/scratch device memory the backend holds outside of any buffer
// (e.g. compute preallocations and staging buffers). The backend stays usable; the scratch
// is reallocated lazily on the next compute. Used to shrink an idle model's device footprint.
void (*free_scratch) (ggml_backend_t backend);
};
struct ggml_backend {
ggml/src/ggml-backend.cpp
+182
View File
@@ -420,6 +420,15 @@ void ggml_backend_synchronize(ggml_backend_t backend) {
backend->iface.synchronize(backend);
}
void ggml_backend_free_scratch(ggml_backend_t backend) {
GGML_ASSERT(backend);
if (backend->iface.free_scratch == NULL) {
return;
}
backend->iface.free_scratch(backend);
}
ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
GGML_ASSERT(backend);
GGML_ASSERT(backend->iface.graph_plan_create != NULL);
@@ -761,6 +770,10 @@ static bool ggml_is_view_op(enum ggml_op op) {
#define GGML_SCHED_MAX_COPIES 4
#endif
#ifndef GGML_SCHED_MAX_PREFETCH_SLOTS
#define GGML_SCHED_MAX_PREFETCH_SLOTS 8
#endif
struct ggml_backend_sched_split {
int backend_id;
int i_start;
@@ -818,6 +831,19 @@ struct ggml_backend_sched {
bool op_offload;
// full-tensor prefetch of offloaded MUL_MAT_ID weights (GGML_SCHED_PREFETCH_EXPERTS)
// with a large batch virtually every expert is used, so the routing ids are not worth
// waiting for; uploads run through a second backend instance on the same device so
// they overlap compute, alternating between two staging slots
bool prefetch_experts;
ggml_backend_t prefetch_backend;
int prefetch_n_slots;
ggml_backend_buffer_t prefetch_slots[GGML_SCHED_MAX_PREFETCH_SLOTS];
ggml_backend_event_t prefetch_ready[GGML_SCHED_MAX_PREFETCH_SLOTS];
ggml_backend_event_t prefetch_free[GGML_SCHED_MAX_PREFETCH_SLOTS];
bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS];
int prefetch_cur;
int debug;
// used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC]
@@ -1538,6 +1564,94 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
return true;
}
static void ggml_backend_sched_prefetch_disable(ggml_backend_sched_t sched, ggml_backend_t split_backend) {
sched->prefetch_experts = false;
if (sched->prefetch_backend) {
ggml_backend_synchronize(split_backend);
ggml_backend_synchronize(sched->prefetch_backend);
}
for (int i = 0; i < sched->prefetch_n_slots; i++) {
ggml_backend_buffer_free(sched->prefetch_slots[i]);
sched->prefetch_slots[i] = NULL;
sched->prefetch_used[i] = false;
}
}
// slots are sized once for the largest offloaded expert tensor in the current graph so
// that they never need to grow mid-eval
static size_t ggml_backend_sched_prefetch_max_size(ggml_backend_sched_t sched) {
size_t max_size = 0;
for (int split_id = 0; split_id < sched->n_splits; split_id++) {
struct ggml_backend_sched_split * split = &sched->splits[split_id];
if (split->graph.n_nodes == 0 || split->graph.nodes[0]->op != GGML_OP_MUL_MAT_ID) {
continue;
}
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
const ggml_tensor * input = split->inputs[input_id];
if (input->buffer &&
ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
ggml_backend_buffer_is_host(input->buffer)) {
max_size = std::max(max_size, ggml_nbytes(input));
}
}
}
return max_size;
}
static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_backend_t split_backend, size_t size) {
if (sched->prefetch_backend == NULL) {
ggml_backend_dev_t dev = split_backend->device;
ggml_backend_dev_props props;
ggml_backend_dev_get_props(dev, &props);
if (!props.caps.async || !props.caps.events) {
sched->prefetch_experts = false;
return false;
}
sched->prefetch_backend = ggml_backend_dev_init(dev, NULL);
if (sched->prefetch_backend == NULL) {
sched->prefetch_experts = false;
return false;
}
for (int i = 0; i < sched->prefetch_n_slots; i++) {
sched->prefetch_ready[i] = ggml_backend_event_new(dev);
sched->prefetch_free[i] = ggml_backend_event_new(dev);
if (sched->prefetch_ready[i] == NULL || sched->prefetch_free[i] == NULL) {
sched->prefetch_experts = false;
return false;
}
}
}
size = std::max(size, ggml_backend_sched_prefetch_max_size(sched));
ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(split_backend);
for (int i = 0; i < sched->prefetch_n_slots; i++) {
if (sched->prefetch_slots[i] == NULL || ggml_backend_buffer_get_size(sched->prefetch_slots[i]) < size) {
// allocate before freeing so a failure leaves the old slot intact
ggml_backend_buffer_t new_buf = ggml_backend_buft_alloc_buffer(buft, size);
if (new_buf == NULL) {
// overlap needs at least 2 slots, otherwise run with what fits
if (i >= 2 && sched->prefetch_slots[0] != NULL &&
ggml_backend_buffer_get_size(sched->prefetch_slots[0]) >= size) {
sched->prefetch_n_slots = i;
sched->prefetch_cur = 0;
return true;
}
ggml_backend_sched_prefetch_disable(sched, split_backend);
return false;
}
if (sched->prefetch_slots[i] != NULL) {
ggml_backend_synchronize(split_backend);
ggml_backend_synchronize(sched->prefetch_backend);
ggml_backend_buffer_free(sched->prefetch_slots[i]);
}
sched->prefetch_slots[i] = new_buf;
sched->prefetch_used[i] = false;
}
}
return true;
}
static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
GGML_ASSERT(sched);
struct ggml_backend_sched_split * splits = sched->splits;
@@ -1550,6 +1664,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
struct ggml_backend_sched_split * split = &splits[split_id];
int split_backend_id = split->backend_id;
ggml_backend_t split_backend = sched->backends[split_backend_id];
int split_prefetch_slot = -1;
ggml_tensor * prefetch_input_cpy = NULL;
ggml_backend_buffer_t prefetch_saved_buffer = NULL;
void * prefetch_saved_data = NULL;
// copy the input tensors to the split backend
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
@@ -1566,6 +1684,41 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
}
ggml_backend_tensor_copy(input, input_cpy);
} else {
// with a large batch virtually every expert is used, so instead of waiting
// for the routing ids, upload the full tensor through the prefetch backend
// and let the copy overlap compute of the previous split
if (sched->prefetch_experts && !sched->callback_eval && split_prefetch_slot == -1 && split->graph.n_nodes > 0) {
ggml_tensor * node = split->graph.nodes[0];
if (ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
ggml_backend_buffer_is_host(input->buffer) &&
node->op == GGML_OP_MUL_MAT_ID && node->src[0] == input_cpy) {
const ggml_tensor * ids = node->src[2];
const int64_t n_expert = input->ne[2];
if (ids->ne[0]*ids->ne[1] >= 2*n_expert &&
ggml_backend_sched_prefetch_init(sched, split_backend, ggml_nbytes(input))) {
const int slot = sched->prefetch_cur;
sched->prefetch_cur = (sched->prefetch_cur + 1) % sched->prefetch_n_slots;
// wait for the previous user of this slot to finish computing
if (sched->prefetch_used[slot]) {
ggml_backend_event_wait(sched->prefetch_backend, sched->prefetch_free[slot]);
}
// point the staging copy at the slot only for the duration of
// this split, so a fallback to the regular path on a later
// eval can never see a dangling slot pointer
prefetch_input_cpy = input_cpy;
prefetch_saved_buffer = input_cpy->buffer;
prefetch_saved_data = input_cpy->data;
input_cpy->buffer = sched->prefetch_slots[slot];
input_cpy->data = ggml_backend_buffer_get_base(sched->prefetch_slots[slot]);
ggml_backend_tensor_set_async(sched->prefetch_backend, input_cpy, input->data, 0, ggml_nbytes(input));
ggml_backend_event_record(sched->prefetch_ready[slot], sched->prefetch_backend);
ggml_backend_event_wait(split_backend, sched->prefetch_ready[slot]);
split_prefetch_slot = slot;
continue;
}
}
}
// wait for the split backend to finish using the input before overwriting it
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
@@ -1676,6 +1829,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
if (!sched->callback_eval) {
enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
if (split_prefetch_slot != -1) {
// the kernels have captured the slot address at launch, safe to restore
ggml_backend_event_record(sched->prefetch_free[split_prefetch_slot], split_backend);
sched->prefetch_used[split_prefetch_slot] = true;
prefetch_input_cpy->buffer = prefetch_saved_buffer;
prefetch_input_cpy->data = prefetch_saved_data;
}
if (ec != GGML_STATUS_SUCCESS) {
return ec;
}
@@ -1788,6 +1948,15 @@ ggml_backend_sched_t ggml_backend_sched_new(
sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
sched->op_offload = op_offload;
// GGML_SCHED_PREFETCH_EXPERTS=1 enables the default slot count, higher values set it
// directly; more slots let uploads run further ahead of compute at the cost of one
// max-sized expert tensor of device memory per slot
const char * GGML_SCHED_PREFETCH_EXPERTS = getenv("GGML_SCHED_PREFETCH_EXPERTS");
const int prefetch_n_slots = GGML_SCHED_PREFETCH_EXPERTS ? atoi(GGML_SCHED_PREFETCH_EXPERTS) : 0;
sched->prefetch_experts = op_offload && prefetch_n_slots > 0;
// default of 3 covers the gate/up/down expert tensors of one MoE layer
sched->prefetch_n_slots = prefetch_n_slots <= 1 ? 3 : std::min(prefetch_n_slots, GGML_SCHED_MAX_PREFETCH_SLOTS);
ggml_backend_sched_reset(sched);
return sched;
@@ -1802,6 +1971,16 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
ggml_backend_event_free(sched->events[b][c]);
}
}
if (sched->prefetch_backend) {
ggml_backend_synchronize(sched->prefetch_backend);
// the slot count may have been reduced after a failed allocation, free everything
for (int i = 0; i < GGML_SCHED_MAX_PREFETCH_SLOTS; i++) {
ggml_backend_event_free(sched->prefetch_ready[i]);
ggml_backend_event_free(sched->prefetch_free[i]);
ggml_backend_buffer_free(sched->prefetch_slots[i]);
}
ggml_backend_free(sched->prefetch_backend);
}
ggml_gallocr_free(sched->galloc);
ggml_free(sched->ctx);
ggml_hash_set_free(&sched->hash_set);
@@ -1906,6 +2085,9 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) {
for (int i = 0; i < sched->n_backends; i++) {
ggml_backend_synchronize(sched->backends[i]);
}
if (sched->prefetch_backend) {
ggml_backend_synchronize(sched->prefetch_backend);
}
if (!sched->is_alloc) {
// if the graph is not already allocated, always use copy 0 after a synchronization
// this ensures that during generation the same copy is used every time,
ggml/src/ggml-cuda/fattn-tile.cu
-1
View File
@@ -1,6 +1,5 @@
#include "common.cuh"
#include "fattn-tile.cuh"
#include "fattn-wmma-f16.cuh"
void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * K = dst->src[1];
ggml/src/ggml-cuda/fattn-tile.cuh
+1 -7
View File
@@ -1,6 +1,5 @@
#include "common.cuh"
#include "fattn-common.cuh"
#include "fattn-wmma-f16.cuh"
// nbatch_fa == number of KQ rows to process per iteration
// nbatch_K == number of K columns to load in parallel for KQ calculation
@@ -825,12 +824,7 @@ static __global__ void flash_attn_tile(
// Skip unused kernel variants for faster compilation:
if (
#ifdef GGML_USE_WMMA_FATTN
(ncols2 != 1 && DV != 40 && DV != 72 && DV != 512) ||
#endif // GGML_USE_WMMA_FATTN
(use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512))
) {
if ((use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512))) {
GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
ggml/src/ggml-cuda/fattn-wmma-f16.cu
-705
View File
@@ -1,705 +0,0 @@
// Old and deprecated WMMA FlashAttention implementation.
// It is still needed for Volta since the memory layout of NVIDIA tensor cores changed with Turing.
// Long-term the WMMA code should be replaced with a dedicated Volta implementation.
#include "common.cuh"
#include "fattn-common.cuh"
#include "fattn-wmma-f16.cuh"
#ifdef GGML_USE_WMMA_FATTN
#if !defined(GGML_USE_HIP)
#include <mma.h>
#if defined(GGML_USE_MUSA)
namespace wmma = mtmusa::wmma;
#else // GGML_USE_MUSA
namespace wmma = nvcuda::wmma;
#endif // GGML_USE_MUSA
#elif defined(GGML_USE_HIP)
#include <rocwmma/rocwmma.hpp>
namespace wmma = rocwmma;
#endif // !defined(GGML_USE_HIP)
#endif // GGML_USE_WMMA_FATTN
// D == head size, VKQ_stride == num VKQ rows calculated in parallel:
template<int D, int ncols, int nwarps, int VKQ_stride, typename KQ_acc_t, bool use_logit_softcap>
__launch_bounds__(nwarps*ggml_cuda_get_physical_warp_size(), 1)
static __global__ void flash_attn_ext_f16(
const char * Q_ptr,
const char * K_ptr,
const char * V_ptr,
const char * mask_ptr,
const char * sinks_ptr,
const int * KV_max_ptr,
float * dst_ptr,
float2 * dst_meta_ptr,
const float scale,
const float max_bias,
const float m0,
const float m1,
const uint32_t n_head_log2,
const float logit_softcap,
const int32_t ne00, const uint3 ne01, const int32_t ne02, const int32_t ne03,
const int32_t nb01, const int32_t nb02, const int32_t nb03,
const int32_t ne10, const int32_t ne11, const int32_t ne12, const int32_t ne13,
const int32_t nb11, const int32_t nb12, const int64_t nb13,
const int32_t nb21, const int32_t nb22, const int64_t nb23,
const int32_t ne31, const int32_t ne32, const int32_t ne33,
const int32_t nb31, const int32_t nb32, const int64_t nb33) {
#if defined(FLASH_ATTN_AVAILABLE) && (defined(GGML_HIP_ROCWMMA_FATTN) && defined(GGML_USE_WMMA_FATTN))
const char * GGML_CUDA_RESTRICT Q = Q_ptr;
const char * GGML_CUDA_RESTRICT K = K_ptr;
const char * GGML_CUDA_RESTRICT V = V_ptr;
const char * GGML_CUDA_RESTRICT mask = mask_ptr;
const char * GGML_CUDA_RESTRICT sinks = sinks_ptr;
const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr;
float * GGML_CUDA_RESTRICT dst = dst_ptr;
float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
// Skip unused kernel variants for faster compilation:
if (use_logit_softcap && !(D == 128 || D == 256)) {
NO_DEVICE_CODE;
return;
}
//In this kernel Q, K, V are matrices while i, j, k are matrix indices.
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
const int ic0 = ncols*blockIdx.x; // Index of the first Q/QKV column to work on.
static_assert(D <= FATTN_KQ_STRIDE, "D must be <= FATTN_KQ_STRIDE.");
static_assert(ncols == 8 || ncols % 16 == 0, "ncols must be 8 or a multiple of 16.");
constexpr int frag_m = ncols == 8 ? 32 : 16;
constexpr int frag_n = ncols == 8 ? 8 : 16;
static_assert(D % frag_m == 0, "If ncols == 8 then D % frag_m must be 0.");
#if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, _Float16, wmma::row_major> frag_a_K;
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, _Float16, wmma::col_major> frag_a_V;
typedef wmma::fragment<wmma::matrix_b, frag_m, frag_n, 16, _Float16, wmma::col_major> frag_b;
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, KQ_acc_t> frag_c_KQ;
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, _Float16> frag_c_VKQ;
#else
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, half, wmma::row_major> frag_a_K;
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, half, wmma::col_major> frag_a_V;
typedef wmma::fragment<wmma::matrix_b, frag_m, frag_n, 16, half, wmma::col_major> frag_b;
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, KQ_acc_t> frag_c_KQ;
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, half> frag_c_VKQ;
#endif
constexpr int KQ_stride_tc = nwarps*frag_m; // Number of KQ rows calculated in parallel.
constexpr int VKQ_ratio = KQ_stride_tc/VKQ_stride; // Number of parallel VKQ accumulators needed to keep all warps busy.
static_assert(VKQ_ratio <= nwarps, "VKQ_ratio must be <= nwarps.");
// Pad internal representation of KQ, KQV to reduce shared memory bank conflicts:
constexpr int D_padded = D + 8;
constexpr int kqs_padded = FATTN_KQ_STRIDE + 8;
constexpr int kqar = sizeof(KQ_acc_t)/sizeof(half);
ggml_cuda_pdl_sync();
const int sequence = blockIdx.z / ne02;
const int head = blockIdx.z - sequence*ne02;
const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix.
const float * Q_f = (const float *) (Q + nb03* sequence + nb02* head + nb01*ic0);
const half * K_h = (const half *) (K + nb13* sequence + nb12*(head / gqa_ratio));
const half * V_h = (const half *) (V + nb13* sequence + nb12*(head / gqa_ratio)); // K and V have same shape
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
const half2 * mask2 = (const half2 *) maskh;
const float * sinksf = (const float *) sinks;
const int stride_Q = nb01 / sizeof(float);
const int stride_KV = nb11 / sizeof(half);
const float slopef = get_alibi_slope(max_bias, head, n_head_log2, m0, m1);
const half slopeh = __float2half(slopef);
const half2 slope2 = make_half2(slopef, slopef);
const half2 logit_softcap_2 = make_half2(logit_softcap, logit_softcap);
frag_b Q_b[D/16][ncols/frag_n];
// A single buffer for temporarily holding tiles of KQ and VKQ parts:
constexpr int mem_KQ = ncols*kqs_padded*kqar;
constexpr int mem_VKQ_parts = VKQ_ratio*ncols*D_padded;
__shared__ half KQ[mem_KQ >= mem_VKQ_parts ? mem_KQ : mem_VKQ_parts];
float * KQ_f = (float *) KQ;
half2 * KQ2 = (half2 *) KQ;
float KQ_rowsum_f[ncols/nwarps] = {0.0f};
float KQ_max_f[ncols/nwarps];
float KQ_max_scale_f[ncols/nwarps] = {0.0f};
#pragma unroll
for (int j = 0; j < ncols/nwarps; ++j) {
KQ_max_f[j] = -FLT_MAX/2.0f;
}
half2 KQ_rowsum_h2[ncols/nwarps] = {{0.0f, 0.0f}};
half2 KQ_max_h2[ncols/nwarps];
half2 KQ_max_scale_h2[ncols/nwarps] = {{0.0f, 0.0f}};
#pragma unroll
for (int j = 0; j < ncols/nwarps; ++j) {
KQ_max_h2[j] = make_half2(-HALF_MAX_HALF, -HALF_MAX_HALF);
}
__shared__ half VKQ[ncols*D_padded]; // Accumulator for final VKQ slice.
half2 * VKQ2 = (half2 *) VKQ;
#if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000
const _Float16 * K_h_f16 = reinterpret_cast<const _Float16 *>(K_h);
const _Float16 * V_h_f16 = reinterpret_cast<const _Float16 *>(V_h);
_Float16 * KQ_f16 = reinterpret_cast<_Float16 *>(KQ);
_Float16 * VKQ_f16 = reinterpret_cast<_Float16 *>(VKQ);
#else
const half * K_h_f16 = K_h;
const half * V_h_f16 = V_h;
half * KQ_f16 = KQ;
half * VKQ_f16 = VKQ;
#endif
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
#pragma unroll
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
const int i = i0 + threadIdx.x;
if (i0 + warp_size > D/2 && i >= D/2) {
break;
}
VKQ2[j*(D_padded/2) + i] = make_half2(0.0f, 0.0f);
}
}
// Convert Q to half and apply scale, temporarily store in KQ:
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
#pragma unroll
for (int i0 = 0; i0 < D; i0 += warp_size) {
const int i = i0 + threadIdx.x;
if (i0 + warp_size > D && i >= D) {
break;
}
KQ[j*D_padded + i] = ic0 + j < int(ne01.z) ? Q_f[j*stride_Q + i] * scale : 0.0f;
}
}
__syncthreads();
// Load Q into tensor core fragments/registers since it will be used frequently:
#pragma unroll
for (int i0 = 0; i0 < D; i0 += 16) {
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
wmma::load_matrix_sync(Q_b[i0/16][j0/frag_n], KQ_f16 + j0*D_padded + i0, D_padded);
}
}
__syncthreads();
// Iterate over ne11 == previous tokens:
const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11;
for (int k_VKQ_0 = blockIdx.y*FATTN_KQ_STRIDE; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*FATTN_KQ_STRIDE) {
// Calculate tile of KQ:
#pragma unroll
for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE; i_KQ_0 += KQ_stride_tc) {
frag_c_KQ KQ_c[ncols/frag_n];
#pragma unroll
for (int j = 0; j < ncols/frag_n; ++j) {
wmma::fill_fragment(KQ_c[j], static_cast<KQ_acc_t>(0.0f));
}
#pragma unroll
for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += 16) {
frag_a_K K_a;
wmma::load_matrix_sync(K_a, K_h_f16 + int64_t(k_VKQ_0 + i_KQ_0 + frag_m*threadIdx.y)*stride_KV + k_KQ_0, stride_KV);
#pragma unroll
for (int j = 0; j < ncols/frag_n; ++j) {
wmma::mma_sync(KQ_c[j], K_a, Q_b[k_KQ_0/16][j], KQ_c[j]);
}
}
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
wmma::store_matrix_sync((KQ_acc_t *) KQ + j0*kqs_padded + i_KQ_0 + frag_m*threadIdx.y, KQ_c[j0/frag_n], kqs_padded, wmma::mem_col_major);
}
}
__syncthreads();
// Calculate softmax for each KQ column using the current max. value.
// The divisor is stored in KQ_rowsum and will be applied at the end.
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
if (std::is_same<KQ_acc_t, float>::value) {
float KQ_f_tmp[FATTN_KQ_STRIDE / warp_size];
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
const int k = k0 + threadIdx.x;
KQ_f_tmp[k0/warp_size] = KQ_f[j*kqs_padded + k];
if (use_logit_softcap) {
KQ_f_tmp[k0/warp_size] = logit_softcap*tanhf(KQ_f_tmp[k0/warp_size]);
}
}
float KQ_max_new = KQ_max_f[j0/nwarps];
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
const int k = k0 + threadIdx.x;
KQ_f_tmp[k0/warp_size] += mask && ic0 + j < int(ne01.z) ?
__half2float(slopeh*maskh[j*(nb31/sizeof(half)) + k_VKQ_0 + k]) : 0.0f;
KQ_max_new = max(KQ_max_new, KQ_f_tmp[k0/warp_size] + FATTN_KQ_MAX_OFFSET);
}
KQ_max_new = warp_reduce_max<warp_size>(KQ_max_new);
const float diff = KQ_max_f[j0/nwarps] - KQ_max_new;
KQ_max_scale_f[j0/nwarps] = expf(diff);
if (diff <= SOFTMAX_FTZ_THRESHOLD) {
KQ_max_scale_f[j0/nwarps] = 0.0f;
}
KQ_max_f[j0/nwarps] = KQ_max_new;
float KQ_rowsum_add = 0.0f;
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
const int k = k0 + threadIdx.x;
const float diff = KQ_f_tmp[k0/warp_size] - KQ_max_f[j0/nwarps];
KQ_f_tmp[k0/warp_size] = expf(diff);
if (diff <= SOFTMAX_FTZ_THRESHOLD) {
KQ_f_tmp[k0/warp_size] = 0.0f;
}
KQ_rowsum_add += KQ_f_tmp[k0/warp_size];
KQ[j*(kqar*kqs_padded) + k] = KQ_f_tmp[k0/warp_size];
}
KQ_rowsum_add = warp_reduce_sum<warp_size>(KQ_rowsum_add);
// Scale previous KQ_rowsum to account for a potential increase in KQ_max:
KQ_rowsum_f[j0/nwarps] = KQ_max_scale_f[j0/nwarps]*KQ_rowsum_f[j0/nwarps] + KQ_rowsum_add;
} else {
half2 KQ2_tmp[FATTN_KQ_STRIDE/(2*warp_size)];
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
const int k = k0 + threadIdx.x;
KQ2_tmp[k0/warp_size] = KQ2[j*(kqs_padded/2) + k];
if (use_logit_softcap) {
// There is no dedicated tangens hyperbolicus function for half2.
KQ2_tmp[k0/warp_size] = h2exp(KQ2_tmp[k0/warp_size]*make_half2(2.0f, 2.0f));
KQ2_tmp[k0/warp_size] = (KQ2_tmp[k0/warp_size] - make_half2(1.0f, 1.0f))
/(KQ2_tmp[k0/warp_size] + make_half2(1.0f, 1.0f));
KQ2_tmp[k0/warp_size] *= logit_softcap_2;
}
}
half2 KQ_max_new = KQ_max_h2[j0/nwarps];
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
const int k = k0 + threadIdx.x;
KQ2_tmp[k0/warp_size] += mask && ic0 + j < int(ne01.z) ? slope2*mask2[(j*ne11 + k_VKQ_0)/2 + k] : make_half2(0.0f, 0.0f);
KQ_max_new = ggml_cuda_hmax2(KQ_max_new, KQ2_tmp[k0/warp_size]);
}
KQ_max_new = __half2half2(warp_reduce_max<warp_size>(ggml_cuda_hmax(__low2half(KQ_max_new), __high2half(KQ_max_new))));
const half2 diff = KQ_max_h2[j0/nwarps] - KQ_max_new;
KQ_max_scale_h2[j0/nwarps] = h2exp(diff);
const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
*((uint32_t *) &KQ_max_scale_h2[j0/nwarps]) &= ftz_mask;
KQ_max_h2[j0/nwarps] = KQ_max_new;
half2 KQ_rowsum_add = make_half2(0.0f, 0.0f);
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
const int k = k0 + threadIdx.x;
const half2 diff = KQ2_tmp[k0/warp_size] - KQ_max_h2[j0/nwarps];
KQ2_tmp[k0/warp_size] = h2exp(diff);
const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
*((uint32_t *) &KQ2_tmp[k0/warp_size]) &= ftz_mask;
KQ_rowsum_add += KQ2_tmp[k0/warp_size];
KQ2[j*(kqs_padded/2) + k] = KQ2_tmp[k0/warp_size];
}
KQ_rowsum_add = warp_reduce_sum<warp_size>(KQ_rowsum_add);
// Scale previous KQ_rowsum to account for a potential increase in KQ_max:
KQ_rowsum_h2[j0/nwarps] = KQ_max_scale_h2[j0/nwarps]*KQ_rowsum_h2[j0/nwarps] + KQ_rowsum_add;
}
}
__syncthreads();
frag_b KQ_b[FATTN_KQ_STRIDE/(VKQ_ratio*16)][ncols/frag_n];
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += VKQ_ratio*16) {
const int k = k0 + (threadIdx.y % VKQ_ratio)*16;
wmma::load_matrix_sync(
KQ_b[k0/(VKQ_ratio*16)][j0/frag_n],
KQ_f16 + j0*(kqar*kqs_padded) + k,
kqar*kqs_padded);
}
}
frag_c_VKQ VKQ_c[D/VKQ_stride][ncols/frag_n];
#pragma unroll
for (int i_VKQ_0 = 0; i_VKQ_0 < D; i_VKQ_0 += VKQ_stride) {
#pragma unroll
for (int j = 0; j < ncols/frag_n; ++j) {
wmma::fill_fragment(VKQ_c[i_VKQ_0/VKQ_stride][j], static_cast<half>(0.0f));
}
#pragma unroll
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += VKQ_ratio*16) {
const int k = k0 + (threadIdx.y % VKQ_ratio)*16;
frag_a_V v_a;
wmma::load_matrix_sync(v_a, V_h_f16 + int64_t(k_VKQ_0 + k)*stride_KV + i_VKQ_0 + frag_m*(threadIdx.y/VKQ_ratio), stride_KV);
#pragma unroll
for (int j = 0; j < ncols/frag_n; ++j) {
wmma::mma_sync(VKQ_c[i_VKQ_0/VKQ_stride][j], v_a, KQ_b[k0/(VKQ_ratio*16)][j], VKQ_c[i_VKQ_0/VKQ_stride][j]);
}
}
}
__syncthreads();
const int offset_k = (threadIdx.y % VKQ_ratio) * (ncols*D_padded);
#pragma unroll
for (int i_KQ_0 = 0; i_KQ_0 < D; i_KQ_0 += VKQ_stride) {
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
wmma::store_matrix_sync(
KQ_f16 + offset_k + j0*D_padded + i_KQ_0 + frag_m*(threadIdx.y/VKQ_ratio),
VKQ_c[i_KQ_0/VKQ_stride][j0/frag_n],
D_padded, wmma::mem_col_major);
}
}
__syncthreads();
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
half2 VKQ_scale;
if (std::is_same<KQ_acc_t, float>::value) {
VKQ_scale = make_half2(KQ_max_scale_f[j0/nwarps], KQ_max_scale_f[j0/nwarps]);
} else {
VKQ_scale = KQ_max_scale_h2[j0/nwarps];
}
#pragma unroll
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
const int i = i0 + threadIdx.x;
if (i0 + warp_size > D/2 && i >= D/2) {
break;
}
half2 VKQ_add = make_half2(0.0f, 0.0f);
#pragma unroll
for (int l = 0; l < VKQ_ratio; ++l) {
VKQ_add += KQ2[l*(ncols*D_padded/2) + j*(D_padded/2) + i];
}
VKQ2[j*(D_padded/2) + i] = VKQ_scale*VKQ2[j*(D_padded/2) + i] + VKQ_add;
}
}
__syncthreads();
}
// Apply attention sinks
if (sinksf && blockIdx.y == 0) {
const float sinkf = sinksf[head];
const half sinkh = __float2half(sinkf);
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j = j0 + threadIdx.y;
if (std::is_same<KQ_acc_t, float>::value) {
float kqmax_new = fmaxf(KQ_max_f[j0/nwarps], sinkf);
const float KQ_max_scale = expf(KQ_max_f[j0/nwarps] - kqmax_new);
KQ_max_f[j0/nwarps] = kqmax_new;
KQ_rowsum_f[j0/nwarps] = KQ_rowsum_f[j0/nwarps] * KQ_max_scale + expf(sinkf - KQ_max_f[j0/nwarps]);
const half2 scale_h2 = make_half2(KQ_max_scale, KQ_max_scale);
#pragma unroll
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
const int i = i0 + threadIdx.x;
if (i0 + warp_size > D/2 && i >= D/2) break;
VKQ2[j*(D_padded/2) + i] *= scale_h2;
}
} else {
half kqmax_old = __low2half(KQ_max_h2[j0/nwarps]);
half kqmax_new = fmaxf(kqmax_old, sinkh);
KQ_max_h2[j0/nwarps] = __half2half2(kqmax_new);
const half KQ_max_scale_h = hexp(kqmax_old - kqmax_new);
const half2 KQ_max_scale = __half2half2(KQ_max_scale_h);
KQ_rowsum_h2[j0/nwarps] = KQ_rowsum_h2[j0/nwarps] * KQ_max_scale;
const half val = hexp(sinkh - kqmax_new);
KQ_rowsum_h2[j0/nwarps].x = __hadd(KQ_rowsum_h2[j0/nwarps].x, val);
#pragma unroll
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
const int i = i0 + threadIdx.x;
if (i0 + warp_size > D/2 && i >= D/2) break;
VKQ2[j*(D_padded/2) + i] *= KQ_max_scale;
}
}
}
__syncthreads();
}
#pragma unroll
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
const int j_VKQ = j0 + threadIdx.y;
if (ic0 + j_VKQ >= int(ne01.z)) {
return;
}
float KQ_rowsum_j;
if (std::is_same<KQ_acc_t, float>::value) {
KQ_rowsum_j = KQ_rowsum_f[j0/nwarps];
} else {
KQ_rowsum_j = __low2float(KQ_rowsum_h2[j0/nwarps]) + __high2float(KQ_rowsum_h2[j0/nwarps]);
}
const int j_dst_unrolled = ((sequence*int(ne01.z) + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y;
#pragma unroll
for (int i0 = 0; i0 < D; i0 += warp_size) {
const int i = i0 + threadIdx.x;
if (i0 + warp_size > D && i >= D) {
break;
}
float dst_val = VKQ[j_VKQ*D_padded + i];
if (gridDim.y == 1) {
dst_val /= KQ_rowsum_j;
}
dst[j_dst_unrolled*D + i] = dst_val;
}
if (gridDim.y == 1 || threadIdx.x != 0) {
continue;
}
float2 dst_meta_val;
if (std::is_same<KQ_acc_t, float>::value) {
dst_meta_val.x = KQ_max_f[j0/nwarps];
} else {
dst_meta_val.x = __low2float(KQ_max_h2[j0/nwarps]);
}
dst_meta_val.y = KQ_rowsum_j;
dst_meta[j_dst_unrolled] = dst_meta_val;
}
#else
GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
nb01, nb02, nb03,
ne10, ne11, ne12, ne13,
nb11, nb12, nb13,
nb21, nb22, nb23,
ne31, ne32, ne33,
nb31, nb32, nb33);
NO_DEVICE_CODE;
#endif // defined(FLASH_ATTN_AVAILABLE) && (defined(GGML_HIP_ROCWMMA_FATTN) && defined(GGML_USE_WMMA_FATTN))
}
constexpr int get_max_power_of_2(int x) {
return x % 2 == 0 ? 2*get_max_power_of_2(x/2) : 1;
}
static_assert(get_max_power_of_2(1) == 1, "Test failed.");
static_assert(get_max_power_of_2(2) == 2, "Test failed.");
static_assert(get_max_power_of_2(4) == 4, "Test failed.");
static_assert(get_max_power_of_2(6) == 2, "Test failed.");
// Number of VKQ rows calculated in parallel:
constexpr int get_VKQ_stride(int D, int nwarps, int frag_m) {
return (get_max_power_of_2(D/frag_m) < nwarps ? get_max_power_of_2(D/frag_m) : nwarps)*frag_m;
}
static_assert(get_VKQ_stride(128, 1, 32) == 32, "Test failed.");
static_assert(get_VKQ_stride(128, 2, 32) == 64, "Test failed.");
static_assert(get_VKQ_stride(128, 4, 32) == 128, "Test failed.");
static_assert(get_VKQ_stride( 64, 1, 32) == 32, "Test failed.");
static_assert(get_VKQ_stride( 64, 2, 32) == 64, "Test failed.");
static_assert(get_VKQ_stride( 64, 4, 32) == 64, "Test failed.");
static_assert(get_VKQ_stride( 80, 1, 16) == 16, "Test failed.");
static_assert(get_VKQ_stride( 80, 2, 16) == 16, "Test failed.");
static_assert(get_VKQ_stride( 80, 4, 16) == 16, "Test failed.");
template <int D, int cols_per_block, typename KQ_acc_t>
void ggml_cuda_flash_attn_ext_wmma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * KQV = dst;
constexpr int nwarps = 4;
constexpr int frag_m = cols_per_block == 8 && D % 32 == 0 ? 32 : 16;
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
float logit_softcap;
memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
fattn_kernel_t fattn_kernel;
if (logit_softcap == 0.0f) {
constexpr bool use_logit_softcap = false;
fattn_kernel = flash_attn_ext_f16<
D, cols_per_block, nwarps, get_VKQ_stride(D, nwarps, frag_m), KQ_acc_t, use_logit_softcap>;
} else {
constexpr bool use_logit_softcap = true;
fattn_kernel = flash_attn_ext_f16<
D, cols_per_block, nwarps, get_VKQ_stride(D, nwarps, frag_m), KQ_acc_t, use_logit_softcap>;
}
launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, 0, FATTN_KQ_STRIDE, true, true, false, warp_size);
}
void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * KQV = dst;
const ggml_tensor * Q = dst->src[0];
const enum ggml_prec prec = ggml_flash_attn_ext_get_prec(KQV);
const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
if (prec != GGML_PREC_DEFAULT) {
if (Q->ne[1] <= 32 || Q->ne[0] > 128) {
constexpr int cols_per_block = 16;
switch (Q->ne[0]) {
case 64:
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, float>(ctx, dst);
break;
case 80:
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, float>(ctx, dst);
break;
case 96:
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, float>(ctx, dst);
break;
case 112:
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, float>(ctx, dst);
break;
case 128:
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, float>(ctx, dst);
break;
case 256:
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, float>(ctx, dst);
break;
default:
GGML_ABORT("fatal error");
break;
}
} else {
constexpr int cols_per_block = 32;
switch (Q->ne[0]) {
case 64:
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, float>(ctx, dst);
break;
case 80:
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, float>(ctx, dst);
break;
case 96:
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, float>(ctx, dst);
break;
case 112:
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, float>(ctx, dst);
break;
case 128:
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, float>(ctx, dst);
break;
// case 256:
// ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, float>(ctx, dst);
// break;
default:
GGML_ABORT("fatal error");
break;
}
}
return;
}
#if !defined(GGML_USE_HIP)
if (Q->ne[1] <= 8 && Q->ne[0] % warp_size == 0) {
constexpr int cols_per_block = 8;
switch (Q->ne[0]) {
case 64:
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
break;
case 96:
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
break;
case 128:
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
break;
case 256:
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
break;
default:
GGML_ABORT("fatal error");
break;
}
return;
}
#endif // !defined(GGML_USE_HIP)
if (Q->ne[1] <= 32) {
constexpr int cols_per_block = 16;
switch (Q->ne[0]) {
case 64:
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
break;
case 80:
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, half>(ctx, dst);
break;
case 96:
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
break;
case 112:
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, half>(ctx, dst);
break;
case 128:
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
break;
case 256:
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
break;
default:
GGML_ABORT("fatal error");
break;
}
return;
}
constexpr int cols_per_block = 32;
switch (Q->ne[0]) {
case 64:
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
break;
case 80:
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, half>(ctx, dst);
break;
case 96:
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
break;
case 112:
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, half>(ctx, dst);
break;
case 128:
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
break;
case 256:
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
break;
default:
GGML_ABORT("fatal error");
break;
}
}
ggml/src/ggml-cuda/fattn-wmma-f16.cuh
-51
View File
@@ -1,51 +0,0 @@
#pragma once
#include "common.cuh"
#if defined(GGML_USE_MUSA)
#define GGML_USE_WMMA_FATTN
#endif // defined(GGML_USE_MUSA)
#if defined(GGML_HIP_ROCWMMA_FATTN)
#if defined(CDNA) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
#define GGML_USE_WMMA_FATTN
#elif defined(CDNA)
#warning "rocwmma fattn on CDNA is broken on rocwmma v2.0.0, expect degraded performance"
#endif // defined(CDNA) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
#if defined(RDNA3)
#define GGML_USE_WMMA_FATTN
#endif // defined(RDNA3)
#if defined(RDNA4) && ROCWMMA_VERSION_MAJOR > 1
#define GGML_USE_WMMA_FATTN
#elif defined(RDNA4)
#warning "rocwmma fattn is not supported on RDNA4 on rocwmma < v2.0.0, expect degraded performance"
#endif // defined(RDNA4) && ROCWMMA_VERSION_MAJOR > 1
#endif // defined(GGML_HIP_ROCWMMA_FATTN)
// WMMA flash attention requires FP16 matrix instructions to be available for ggml code.
static bool ggml_cuda_should_use_wmma_fattn(const int cc) {
#if defined(GGML_USE_HIP) && !defined(GGML_HIP_ROCWMMA_FATTN)
return false;
#else
if ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_VOLTA) ||
GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_MTHREADS(cc)) {
return true;
} else if (GGML_CUDA_CC_IS_CDNA(cc)){
#if defined(GGML_HIP_ROCWMMA_FATTN) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
return true;
#else
return false;
#endif // defined(GGML_HIP_ROCWMMA_FATTN) (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
} else if (GGML_CUDA_CC_IS_RDNA4(cc)) {
#if defined(GGML_HIP_ROCWMMA_FATTN) && ROCWMMA_VERSION_MAJOR > 1
return true;
#else
return false;
#endif // defined(GGML_HIP_ROCWMMA_FATTN) && ROCWMMA_VERSION_MAJOR > 1
} else {
return false;
}
#endif // defined(GGML_USE_HIP) && !defined(GGML_HIP_ROCWMMA_FATTN)
}
void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
ggml/src/ggml-cuda/fattn.cu
+4 -18
View File
@@ -3,7 +3,6 @@
#include "fattn-mma-f16.cuh"
#include "fattn-tile.cuh"
#include "fattn-vec.cuh"
#include "fattn-wmma-f16.cuh"
#include "fattn.cuh"
template <int DKQ, int DV, int ncols2>
@@ -330,11 +329,10 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t
// Best FlashAttention kernel for a specific GPU:
enum best_fattn_kernel {
BEST_FATTN_KERNEL_NONE = 0,
BEST_FATTN_KERNEL_TILE = 200,
BEST_FATTN_KERNEL_VEC = 100,
BEST_FATTN_KERNEL_WMMA_F16 = 300,
BEST_FATTN_KERNEL_MMA_F16 = 400,
BEST_FATTN_KERNEL_NONE = 0,
BEST_FATTN_KERNEL_TILE = 200,
BEST_FATTN_KERNEL_VEC = 100,
BEST_FATTN_KERNEL_MMA_F16 = 400,
};
static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
@@ -500,14 +498,6 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
return BEST_FATTN_KERNEL_MMA_F16;
}
// Use the WMMA kernel if possible:
if (ggml_cuda_should_use_wmma_fattn(cc) && K->ne[1] % FATTN_KQ_STRIDE == 0 && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[0] != 192 && Q->ne[0] != 512 && Q->ne[0] != 576) {
if (can_use_vector_kernel && Q->ne[1] <= 2) {
return BEST_FATTN_KERNEL_VEC;
}
return BEST_FATTN_KERNEL_WMMA_F16;
}
// AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
if ((amd_mfma_available(cc) && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72) {
if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {
@@ -559,7 +549,6 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d
switch (kernel) {
case BEST_FATTN_KERNEL_TILE:
case BEST_FATTN_KERNEL_WMMA_F16:
case BEST_FATTN_KERNEL_MMA_F16:
need_f16_K = true;
need_f16_V = true;
@@ -589,9 +578,6 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
case BEST_FATTN_KERNEL_VEC:
ggml_cuda_flash_attn_ext_vec(ctx, dst);
break;
case BEST_FATTN_KERNEL_WMMA_F16:
ggml_cuda_flash_attn_ext_wmma_f16(ctx, dst);
break;
case BEST_FATTN_KERNEL_MMA_F16:
ggml_cuda_flash_attn_ext_mma_f16(ctx, dst);
break;
ggml/src/ggml-cuda/mmq.cu
+18 -12
View File
@@ -25,12 +25,7 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
case GGML_TYPE_Q8_0:
mul_mat_q_case<GGML_TYPE_Q8_0>(ctx, args, stream);
break;
case GGML_TYPE_MXFP4:
mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
break;
case GGML_TYPE_NVFP4:
mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
break;
// -----------------------------------------------------------------------
case GGML_TYPE_Q2_K:
mul_mat_q_case<GGML_TYPE_Q2_K>(ctx, args, stream);
break;
@@ -46,6 +41,10 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
case GGML_TYPE_Q6_K:
mul_mat_q_case<GGML_TYPE_Q6_K>(ctx, args, stream);
break;
// -----------------------------------------------------------------------
case GGML_TYPE_IQ1_S:
mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
break;
case GGML_TYPE_IQ2_XXS:
mul_mat_q_case<GGML_TYPE_IQ2_XXS>(ctx, args, stream);
break;
@@ -61,15 +60,19 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
case GGML_TYPE_IQ3_S:
mul_mat_q_case<GGML_TYPE_IQ3_S>(ctx, args, stream);
break;
case GGML_TYPE_IQ1_S:
mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
break;
case GGML_TYPE_IQ4_XS:
mul_mat_q_case<GGML_TYPE_IQ4_XS>(ctx, args, stream);
break;
case GGML_TYPE_IQ4_NL:
mul_mat_q_case<GGML_TYPE_IQ4_NL>(ctx, args, stream);
break;
// -----------------------------------------------------------------------
case GGML_TYPE_MXFP4:
mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
break;
case GGML_TYPE_NVFP4:
mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
break;
default:
GGML_ABORT("fatal error");
break;
@@ -264,21 +267,24 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q8_0:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
// -------------------------------------------------
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q6_K:
// -------------------------------------------------
case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ2_XXS:
case GGML_TYPE_IQ2_XS:
case GGML_TYPE_IQ2_S:
case GGML_TYPE_IQ3_XXS:
case GGML_TYPE_IQ3_S:
case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ4_XS:
case GGML_TYPE_IQ4_NL:
// -------------------------------------------------
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
mmq_supported = true;
break;
default:
ggml/src/ggml-cuda/mmq.cuh
+7 -3
View File
@@ -1537,26 +1537,30 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda
#define DECL_MMQ_CASE(type) \
template void mul_mat_q_case<type>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \
extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_1);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_1);
extern DECL_MMQ_CASE(GGML_TYPE_Q8_0);
extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
// -----------------------------------------
extern DECL_MMQ_CASE(GGML_TYPE_Q2_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q3_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q6_K);
// -----------------------------------------
extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS);
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XS);
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS);
extern DECL_MMQ_CASE(GGML_TYPE_IQ3_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL);
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
// -----------------------------------------
extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
// -------------------------------------------------------------------------------------------------------------------------
ggml/src/ggml-cuda/vendors/hip.h
Vendored
-4
View File
@@ -6,10 +6,6 @@
#include <hip/hip_fp16.h>
#include <hip/hip_bf16.h>
#if defined(GGML_HIP_ROCWMMA_FATTN)
#include <rocwmma/rocwmma-version.hpp>
#endif // defined(GGML_HIP_ROCWMMA_FATTN)
#ifdef GGML_USE_NCCL
#include <rccl/rccl.h>
#endif // GGML_USE_NCCL
ggml/src/ggml-hexagon/ggml-hexagon.cpp
+127 -154
View File
@@ -143,12 +143,12 @@ static const char * htp_event_name(uint16_t id) {
case HTP_TRACE_EVT_HMX_COMP: return "HMX_COMP";
case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH";
case HTP_TRACE_EVT_INIT: return "INIT";
case HTP_TRACE_EVT_BUFF: return "BUFF";
default: return "UNKNOWN";
}
}
static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node,
const htp_prof_desc & pd) {
static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node, const htp_prof_desc & pd) {
if (!opt_profile) return;
uint32_t op_usec = pd.usecs;
@@ -168,6 +168,43 @@ static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_op
node.op_name().c_str(), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.kparams, op_usec, op_cycles, pd.cycles_start, mhz, pmu_str);
}
static void ggml_hexagon_dump_batch_prof(const std::string & sess_name, const htp_opbatch_rsp & rsp) {
uint64_t batch_cycles = rsp.cycles_stop - rsp.cycles_start;
float batch_mhz = rsp.usecs > 0 ? (float) batch_cycles / rsp.usecs : 0.0f;
char evt_str[256] = "----";
if (opt_profile == 3) {
snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
}
GGML_LOG_DEBUG("ggml-hex: %s profile-op OPBATCH|----|n-ops %u|%s|----|----|usec %u cycles %llu start %llu mhz %.1f\n",
sess_name.c_str(), rsp.n_ops, evt_str, rsp.usecs, (unsigned long long) batch_cycles, (unsigned long long) rsp.cycles_start, batch_mhz);
}
static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const htp_opbatch_rsp & rsp,
const htp_trace_desc * trace_events, uint32_t n_traces) {
if (opt_profile == 3 && trace_events) {
uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
uint32_t count = rsp.n_traces[t];
valid_cnt[t] = count > n_traces ? n_traces : count;
}
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
for (uint32_t idx = 0; idx < valid_cnt[t]; idx++) {
const auto & e = trace_events[t * n_traces + idx];
bool is_stop = (e.info & 0x8000) != 0;
uint16_t info = e.info & 0x7FFF;
GGML_LOG_DEBUG("ggml-hex: %s trace-evt %s: thread %u info %u %s %u\n",
sess_name.c_str(), htp_event_name(e.id), t, info, is_stop ? "stop" : "start", e.cycles);
}
}
}
}
// **
static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) {
@@ -1128,13 +1165,7 @@ struct ggml_hexagon_opbatch {
std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr to index
std::unordered_multimap<void*, int> d_map; // tensor data to index
struct tensor_range {
uint64_t start;
uint64_t end;
int bi;
std::vector<int> tensors;
};
std::vector<tensor_range> ranges;
unsigned int n_bufs; // num buffers in the batch
unsigned int n_tens; // num tensors ...
@@ -1155,7 +1186,6 @@ struct ggml_hexagon_opbatch {
b_map.clear();
t_map.clear();
d_map.clear();
ranges.clear();
}
ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) {
@@ -1209,70 +1239,7 @@ struct ggml_hexagon_opbatch {
return bi;
}
void add_range(const htp_tensor * h, int ti) {
uint64_t t_start = h->data;
uint64_t t_end = t_start + h->size;
int bi = h->bi;
int first_match = -1;
int unused_idx = -1;
for (size_t i = 0; i < ranges.size(); i++) {
if (ranges[i].bi == -1) {
unused_idx = i;
continue;
}
if (ranges[i].bi != bi) {
continue;
}
if (ranges[i].start >= t_end || ranges[i].end <= t_start) {
continue;
}
if (first_match == -1) {
first_match = i;
HEX_VERBOSE("ggml-hex: %s range-grow #%d : bi %d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
sess->c_name(), (int) i, ranges[i].bi,
(void *) (h_bufs[ranges[i].bi].base + ranges[i].start),
(void *) (h_bufs[ranges[i].bi].base + ranges[i].end),
ti,
(void *) (h_bufs[bi].base + t_start),
(void *) (h_bufs[bi].base + t_end),
(void *) (h_bufs[ranges[i].bi].base + std::min(ranges[i].start, t_start)),
(void *) (h_bufs[ranges[i].bi].base + std::max(ranges[i].end, t_end)));
ranges[i].start = std::min(ranges[i].start, t_start);
ranges[i].end = std::max(ranges[i].end, t_end);
ranges[i].tensors.push_back(ti);
} else {
HEX_VERBOSE("ggml-hex: %s range-merge #%d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
sess->c_name(), first_match,
(void *) (h_bufs[bi].base + ranges[first_match].start),
(void *) (h_bufs[bi].base + ranges[first_match].end),
(int) i,
(void *) (h_bufs[bi].base + ranges[i].start),
(void *) (h_bufs[bi].base + ranges[i].end),
(void *) (h_bufs[bi].base + std::min(ranges[first_match].start, ranges[i].start)),
(void *) (h_bufs[bi].base + std::max(ranges[first_match].end, ranges[i].end)));
ranges[first_match].start = std::min(ranges[first_match].start, ranges[i].start);
ranges[first_match].end = std::max(ranges[first_match].end, ranges[i].end);
ranges[first_match].tensors.insert(
ranges[first_match].tensors.end(),
ranges[i].tensors.begin(),
ranges[i].tensors.end()
);
ranges[i].bi = -1;
}
}
if (first_match == -1) {
if (unused_idx != -1) {
ranges[unused_idx] = {t_start, t_end, bi, {ti}};
} else {
ranges.push_back({t_start, t_end, bi, {ti}});
}
}
}
bool same_shape(const htp_tensor * h, const ggml_tensor * t) const {
int64_t ne0 = t->ne[0];
@@ -1341,8 +1308,7 @@ struct ggml_hexagon_opbatch {
h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3];
}
h.alias = ti;
add_range(&h, ti);
h.flags = 0;
if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
@@ -1424,14 +1390,6 @@ struct ggml_hexagon_opbatch {
}
void finalize_ranges() {
for (const auto & r : ranges) {
if (r.bi == -1) {
continue;
}
for (size_t i = 0; i < r.tensors.size(); i++) {
h_tens[r.tensors[i]].alias = r.tensors[(i + 1) % r.tensors.size()];
}
}
}
};
@@ -1582,9 +1540,6 @@ struct ggml_hexagon_opqueue {
if (opt_profile && rsp.n_ops > 0) {
auto & ops = op_cache[rsp.id];
uint64_t batch_usec = ggml_time_us() - start_usec[rsp.id];
uint32_t htp_usec = 0;
GGML_ASSERT(rsp.n_ops <= ops.size());
const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr;
@@ -1595,55 +1550,13 @@ struct ggml_hexagon_opqueue {
trace_events = (const htp_trace_desc *) (p_ptr + p_size);
}
uint32_t trace_idx[HTP_MAX_NTHREADS + 1] = {0};
uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
if (opt_profile == 3) {
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
uint32_t count = rsp.n_traces[t];
valid_cnt[t] = count > n_traces ? n_traces : count;
}
}
ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp);
for (uint32_t i = 0; i < rsp.n_ops; i++) {
htp_usec += pd[i].usecs;
ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]);
if (opt_profile == 3) {
uint32_t op_duration = pd[i].cycles_stop - pd[i].cycles_start;
for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
while (trace_idx[t] < valid_cnt[t]) {
const auto & e = trace_events[t * n_traces + trace_idx[t]];
uint32_t offset = e.cycles - pd[i].cycles_start;
if (offset >= 0x80000000) {
trace_idx[t]++;
continue;
}
if (offset > op_duration) {
break;
}
bool is_stop = (e.info & 0x8000) != 0;
uint16_t info = e.info & 0x7FFF;
GGML_LOG_DEBUG("ggml-hex: %s trace-op %s: thread %u event %s info %u %s %u\n",
shm_buf->sess->c_name(), ops[i].op_name().c_str(), t, htp_event_name(e.id), info, is_stop ? "stop" : "start", e.cycles);
trace_idx[t]++;
}
}
}
}
char evt_str[256] = "";
if (opt_profile == 3) {
snprintf(evt_str, sizeof(evt_str), " evt [%u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u]",
rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
}
GGML_LOG_DEBUG("ggml-hex: %s profile-batch n-ops %u batch-dur-usec %lld htp-ops-usec %u%s\n",
shm_buf->sess->c_name(), rsp.n_ops, (long long) batch_usec, htp_usec, evt_str);
ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces);
}
}
};
@@ -1662,7 +1575,7 @@ void ggml_hexagon_session::flush_pending(bool all) {
const uint32_t timeo = opt_oppoll ? 0 : DSPQUEUE_TIMEOUT;
int err = dspqueue_read(this->queue, &flags, 1, &n_dbufs, &dbuf, sizeof(rsp), &rsp_size, (uint8_t *) &rsp, timeo);
if (err == AEE_EEXPIRED) {
if (err == AEE_EEXPIRED || err == AEE_EWOULDBLOCK) {
continue;
}
@@ -2114,7 +2027,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
const struct ggml_tensor * sinks = op->src[4];
if (ggml_hexagon_flash_attn_is_hmx_eligible(sess, q, k, v, sinks)) {
size_t Br = 0, Bc = 0;
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads);
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0);
if (ret == 0) {
kparams->kernel_type = HTP_FA_KERNEL_HMX;
kparams->Br = Br;
@@ -2124,7 +2037,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
kparams->u.hmx.g_br = hex_align_up(G * Br, 32);
kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0;
kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0);
kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0);
const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256);
kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes
@@ -2413,6 +2326,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
int ne12,
int ne13,
bool is_matmul_id,
const size_t src2_row_size,
size_t vtcm_budget,
struct htp_mm_kernel_params * kparams
) {
@@ -2438,7 +2352,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
0, src0->nb[1], 0, d, true, false, false
0, src0->nb[1], 0, src2_row_size, d, true, false, false
);
if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
@@ -2448,7 +2362,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
0, src0->nb[1], 0, 2, true, false, false
0, src0->nb[1], 0, src2_row_size, 2, true, false, false
);
}
kparams->n_prefetch = best_n_prefetch;
@@ -2472,7 +2386,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false, false
);
if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
@@ -2482,7 +2396,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false, false
);
}
@@ -2506,7 +2420,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
kparams->n_prefetch = 16;
@@ -2526,7 +2440,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
@@ -2546,7 +2460,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->src1_row_size = src1->nb[1];
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
@@ -2562,7 +2476,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
@@ -2578,7 +2492,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->src1_row_size = src1->nb[1];
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
@@ -2589,11 +2503,12 @@ static void ggml_hexagon_precompute_hvx_mm_params(
}
}
static void ggml_hexagon_precompute_matmul_params(
static void ggml_hexagon_precompute_matmul_params_impl(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * dst,
const size_t src2_row_size,
struct htp_mm_kernel_params * kparams
) {
memset(kparams, 0, sizeof(*kparams));
@@ -2628,7 +2543,7 @@ static void ggml_hexagon_precompute_matmul_params(
}
// Fallback to HVX parameter computation
ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, vtcm_budget, kparams);
ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, src2_row_size, vtcm_budget, kparams);
finalize:
kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11);
@@ -2638,6 +2553,27 @@ finalize:
kparams->div_ne11 = init_fastdiv_values(ne11);
}
static void ggml_hexagon_precompute_matmul_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * dst,
struct htp_mm_kernel_params * kparams
) {
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams);
}
static void ggml_hexagon_precompute_fused_matmul_add_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * src2,
const struct ggml_tensor * dst,
struct htp_mm_kernel_params * kparams
) {
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2->nb[1], kparams);
}
static void ggml_hexagon_precompute_unary_params(
const struct ggml_hexagon_session * sess,
uint32_t op,
@@ -2731,7 +2667,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, d, false, true, false
0, src0_row_size, src1_row_size, 0, d, false, true, false
);
if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
@@ -2746,7 +2682,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
// Test tiled first
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false
0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true, false
);
if (try_tiled && L.total_bytes <= sess->vtcm_size) {
@@ -2764,7 +2700,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false
0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true, false
);
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
@@ -2801,7 +2737,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, d, false, false, true
0, src0_row_size, src1_row_size, 0, d, false, false, true
);
if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
@@ -2816,7 +2752,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
// Test tiled first
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true
0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, false, true
);
if (try_tiled && L.total_bytes <= sess->vtcm_size) {
@@ -2833,7 +2769,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true
0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, false, true
);
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
@@ -3345,6 +3281,35 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session *
GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const struct ggml_tensor * src1 = op->src[1];
const struct ggml_tensor * dst = op;
const bool is_2D = ((const int32_t *) op->op_params)[6] == 1;
if (!is_2D) {
return false;
}
// For now support F32->F32 and F32->F16 only.
if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) {
return false;
}
if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) {
return false;
}
// For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs.
const int32_t p0 = ((const int32_t *) op->op_params)[2];
const int32_t p1 = ((const int32_t *) op->op_params)[3];
if (p0 != 0 || p1 != 0) {
return false;
}
GGML_UNUSED(sess);
return true;
}
static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const struct ggml_tensor * src0 = op->src[0];
const struct ggml_tensor * dst = op;
@@ -3494,6 +3459,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
case GGML_OP_SOLVE_TRI: return HTP_OP_SOLVE_TRI;
case GGML_OP_TRI: return HTP_OP_TRI;
case GGML_OP_PAD: return HTP_OP_PAD;
case GGML_OP_IM2COL: return HTP_OP_IM2COL;
case GGML_OP_UNARY:
switch (ggml_get_unary_op(t)) {
@@ -3656,16 +3622,19 @@ static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph *
if (n->op == GGML_OP_MUL_MAT && next_node) {
if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) {
if (next_node->src[0] == n || next_node->src[1] == n) {
const struct ggml_tensor * src2 = (next_node->src[0] == n) ? next_node->src[1] : next_node->src[0];
struct htp_mm_kernel_params kparams;
ggml_hexagon_precompute_matmul_params(sess, n->src[0], n->src[1], next_node, &kparams);
if ((size_t)kparams.vtcm_size <= sess->vtcm_size) {
ggml_hexagon_precompute_fused_matmul_add_params(sess, n->src[0], n->src[1], src2, next_node, &kparams);
const int src1_nrows = n->src[1]->ne[1] * n->src[1]->ne[2] * n->src[1]->ne[3];
const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1);
if (can_fuse && (size_t)kparams.vtcm_size <= sess->vtcm_size) {
htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD);
node.add_fused(next_node);
memcpy(node.kernel_params, &kparams, sizeof(kparams));
nodes.push_back(std::move(node));
i += 1;
return true;
} else {
} else if (can_fuse) {
HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n",
kparams.vtcm_size, sess->vtcm_size);
}
@@ -4213,6 +4182,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
supp = ggml_hexagon_supported_ssm_conv(sess, op);
break;
case GGML_OP_IM2COL:
supp = ggml_hexagon_supported_im2col(sess, op);
break;
case GGML_OP_GATED_DELTA_NET:
supp = ggml_hexagon_supported_gated_delta_net(sess, op);
break;
@@ -4455,7 +4428,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
opt_opstage = str_opstage ? strtoul(str_opstage, NULL, 0) : opt_opstage;
opt_opbatch = str_opbatch ? strtoul(str_opbatch, NULL, 0) : opt_opbatch;
opt_opqueue = str_opqueue ? strtoul(str_opqueue, NULL, 0) : opt_opqueue;
opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 128);
opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 256);
opt_oppoll = str_oppoll ? strtoul(str_oppoll, NULL, 0) : opt_oppoll;
opt_opfusion = str_opfusion ? atoi(str_opfusion) : opt_opfusion;
opt_profile = str_profile ? atoi(str_profile) : 0;
ggml/src/ggml-hexagon/htp-drv.cpp
+18 -6
View File
@@ -59,7 +59,11 @@ typedef AEEResult (*dspqueue_read_pfn_t)(dspqueue_t queue, uint32_t *flags,
uint32_t max_message_length,
uint32_t *message_length, uint8_t *message,
uint32_t timeout_us);
typedef AEEResult (*dspqueue_read_noblock_pfn_t)(dspqueue_t queue, uint32_t *flags,
uint32_t max_buffers, uint32_t *num_buffers,
struct dspqueue_buffer *buffers,
uint32_t max_message_length,
uint32_t *message_length, uint8_t *message);
typedef int (*fastrpc_mmap_pfn_t)(int domain, int fd, void *addr, int offset, size_t length, enum fastrpc_map_flags flags);
typedef int (*fastrpc_munmap_pfn_t)(int domain, int fd, void *addr, size_t length);
@@ -82,11 +86,12 @@ rpcmem_to_fd_pfn_t rpcmem_to_fd_pfn = nullptr;
fastrpc_mmap_pfn_t fastrpc_mmap_pfn = nullptr;
fastrpc_munmap_pfn_t fastrpc_munmap_pfn = nullptr;
dspqueue_create_pfn_t dspqueue_create_pfn = nullptr;
dspqueue_close_pfn_t dspqueue_close_pfn = nullptr;
dspqueue_export_pfn_t dspqueue_export_pfn = nullptr;
dspqueue_write_pfn_t dspqueue_write_pfn = nullptr;
dspqueue_read_pfn_t dspqueue_read_pfn = nullptr;
dspqueue_create_pfn_t dspqueue_create_pfn = nullptr;
dspqueue_close_pfn_t dspqueue_close_pfn = nullptr;
dspqueue_export_pfn_t dspqueue_export_pfn = nullptr;
dspqueue_write_pfn_t dspqueue_write_pfn = nullptr;
dspqueue_read_pfn_t dspqueue_read_pfn = nullptr;
dspqueue_read_noblock_pfn_t dspqueue_read_noblock_pfn = nullptr;
remote_handle64_open_pfn_t remote_handle64_open_pfn = nullptr;
remote_handle64_invoke_pfn_t remote_handle64_invoke_pfn = nullptr;
@@ -167,6 +172,12 @@ AEEResult dspqueue_read(dspqueue_t queue,
uint32_t * message_length,
uint8_t * message,
uint32_t timeout_us) {
#ifdef _WIN32
if (timeout_us == 0) {
return dspqueue_read_noblock_pfn(queue, flags, max_buffers, num_buffers, buffers, max_message_length,
message_length, message);
}
#endif
return dspqueue_read_pfn(queue, flags, max_buffers, num_buffers, buffers, max_message_length, message_length,
message, timeout_us);
}
@@ -349,6 +360,7 @@ int htpdrv_init() {
dlsym(handle.get(), dspqueue_export_pfn_t, dspqueue_export_pfn, dspqueue_export, false);
dlsym(handle.get(), dspqueue_write_pfn_t, dspqueue_write_pfn, dspqueue_write, false);
dlsym(handle.get(), dspqueue_read_pfn_t, dspqueue_read_pfn, dspqueue_read, false);
dlsym(handle.get(), dspqueue_read_noblock_pfn_t, dspqueue_read_noblock_pfn, dspqueue_read_noblock, false);
dlsym(handle.get(), remote_handle64_open_pfn_t, remote_handle64_open_pfn, remote_handle64_open, false);
dlsym(handle.get(), remote_handle64_invoke_pfn_t, remote_handle64_invoke_pfn, remote_handle64_invoke, false);
dlsym(handle.get(), remote_handle_control_pfn_t, remote_handle_control_pfn, remote_handle_control, false);
ggml/src/ggml-hexagon/htp/CMakeLists.txt
+1
View File
@@ -42,6 +42,7 @@ add_library(${HTP_LIB} SHARED
solve-tri-ops.c
pad-ops.c
argsort-ops.c
im2col-ops.c
)
target_compile_definitions(${HTP_LIB} PRIVATE
ggml/src/ggml-hexagon/htp/dma-queue.c
+1 -3
View File
@@ -101,6 +101,4 @@ void dma_queue_alias_free(dma_queue_t q) {
(void) q;
}
void dma_queue_flush(dma_queue_t q) {
while (dma_queue_pop(q).dst != NULL) ;
}
ggml/src/ggml-hexagon/htp/dma-queue.h
+23 -12
View File
@@ -106,7 +106,7 @@ struct dma_queue_s {
bool alias; // When set, dma_queue_delete will not free the ring
};
void dma_queue_flush(dma_queue_t q);
size_t dma_queue_sizeof(size_t capacity);
size_t dma_queue_alignof(void);
@@ -154,7 +154,6 @@ static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) {
static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
dma_ring * r = q->ring;
if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
FARF(HIGH, "dma-push: queue full\n");
return false;
}
@@ -165,6 +164,8 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
r->dptr[r->push_idx] = dptr;
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
if (size) {
desc->next = NULL;
desc->desc_size = 0; // 1D mode
@@ -173,7 +174,6 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
desc->order = 0;
desc->done = 0;
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
dmlink(r->tail, desc);
r->tail = (dma_descriptor_2d *) desc;
} else {
@@ -188,7 +188,6 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
dma_ring * r = q->ring;
if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
FARF(HIGH, "dma-push: queue full\n");
return false;
}
@@ -224,8 +223,9 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t
r->dptr[r->push_idx] = dptr;
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
if (nrows) {
htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
dmlink(r->tail, desc);
r->tail = desc;
} else {
@@ -252,10 +252,11 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) {
dmpoll();
}
}
htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
dptr = r->dptr[r->pop_idx];
htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
return dptr;
}
@@ -270,6 +271,8 @@ static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
dptr = r->dptr[r->pop_idx];
htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
return dptr;
}
@@ -278,6 +281,10 @@ static inline bool dma_queue_empty(dma_queue * q) {
return q->ring->push_idx == q->ring->pop_idx;
}
static inline void dma_queue_flush(dma_queue * q) {
while (dma_queue_pop(q).dst != NULL) ;
}
static inline uint32_t dma_queue_depth(dma_queue * q) {
return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask;
}
@@ -314,14 +321,18 @@ static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride,
{
const uint8_t *src = (const uint8_t *) dptr.src;
uint8_t *dst = (uint8_t *) dptr.dst;
for (size_t r = 0; r < nrows; ++r) {
size_t r = 0;
while (r + 1 < nrows) {
dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
if (!dma_queue_push_single_1d(q, p, row_size))
return false;
if (r + 1 < nrows)
dma_queue_pop(q);
if (!dma_queue_push_single_1d(q, p, row_size)) {
dma_queue_flush(q);
} else {
r++;
}
}
return true;
dma_queue_flush(q);
dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
return dma_queue_push_single_1d(q, p, row_size);
}
}
ggml/src/ggml-hexagon/htp/flash-attn-ops.c
+515 -179
View File
@@ -123,15 +123,17 @@ struct hmx_fa_context {
uint32_t g_br; // hex_align_up(G * Br, 32) - actual tile row dim
// VTCM buffers (allocated by vtcm_seq_alloc)
__fp16 * vtcm_q_dma; // Q DMA fetch buffer
__fp16 * vtcm_q_tiles; // Q tile format [g_br, D]
__fp16 * vtcm_o_tiles[2]; // O ping-pong [g_br, D]
__fp16 * vtcm_k_fp16[2]; // K DMA double-buffer [Bc, D]
__fp16 * vtcm_v_fp16[2]; // V DMA double-buffer [Bc, D]
__fp16 * vtcm_k_tiles; // K tiles (transposed)
__fp16 * vtcm_k_tiles[2]; // K tiles (transposed, double-buffered)
__fp16 * vtcm_v_tiles[2]; // V tiles (column-major, double-buffered)
__fp16 * vtcm_s_tiles; // S = QK^T [g_br, Bc]
__fp16 * vtcm_p_tiles; // P = softmax(S) [g_br, Bc]
__fp16 * vtcm_s_tiles[2]; // S = QK^T [g_br, Bc] (double-buffered)
__fp16 * vtcm_p_tiles[2]; // P = softmax(S) [g_br, Bc]
__fp16 * vtcm_d_tiles; // Diagonal rescale [g_br, g_br]
__fp16 * vtcm_d_inv_l; // Diagonal rescale (1/l) [g_br, g_br]
HVX_Vector * vtcm_m_vec; // Row max [g_br]
HVX_Vector * vtcm_l_vec; // Row sum [g_br]
HVX_Vector * vtcm_s_rowmax; // Softmax intermediate [g_br]
@@ -236,10 +238,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3);
const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2);
// Fetch Q row
const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
const __fp16 * mp_base = NULL;
if (mask) {
const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
@@ -247,26 +245,91 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]);
}
// Prefetch first two blocks
for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
// Precalculate next row variables if there is a next row
bool has_next_ir = (ir + 1 < ir1);
uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0;
const uint8_t * next_q_row_ptr = NULL;
const __fp16 * next_mp_base = NULL;
// K
const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
const uint8_t * next_k_src0 = NULL;
const uint8_t * next_v_src0 = NULL;
const uint8_t * next_m_src0 = NULL;
uint32_t next_block_size0 = 0;
// V
const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
const uint8_t * next_k_src1 = NULL;
const uint8_t * next_v_src1 = NULL;
const uint8_t * next_m_src1 = NULL;
uint32_t next_block_size1 = 0;
if (has_next_ir) {
const uint32_t next_ir = ir + 1;
const uint32_t next_iq3 = fastdiv(next_ir, &factx->src0_div21);
const uint32_t next_iq2 = fastdiv(next_ir - next_iq3*neq2*neq1, &factx->src0_div1);
const uint32_t next_iq1 = (next_ir - next_iq3*neq2*neq1 - next_iq2 * neq1);
next_ik3 = fastdiv(next_iq3, &factx->broadcast_rk3);
next_ik2 = fastdiv(next_iq2, &factx->broadcast_rk2);
next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3);
next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2);
next_q_row_ptr = (const uint8_t *) q->data + (next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3);
// Mask
if (mask) {
const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
// Mask is 1D contiguous for this row
dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2);
const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3);
next_mp_base = (const __fp16 *) ((const uint8_t *) mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]);
}
// Precalculate next K/V block 0 source pointers
{
const uint32_t ic_start = 0;
next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
next_k_src0 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
next_v_src0 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
if (mask) {
next_m_src0 = (const uint8_t *) (next_mp_base + ic_start);
}
}
// Precalculate next K/V block 1 source pointers (if n_blocks > 1)
if (factx->n_blocks > 1) {
const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE;
next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
next_k_src1 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
next_v_src1 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
if (mask) {
next_m_src1 = (const uint8_t *) (next_mp_base + ic_start);
}
}
}
if (ir == ir0) {
// Fetch Q row
const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
// Prefetch first two blocks
for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
// K
const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
// V
const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
// Mask
if (mask) {
const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
// Mask is 1D contiguous for this row
dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
}
}
}
@@ -287,6 +350,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
const HVX_Vector slope_vec = hvx_vec_splat_f16(slope);
const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
const HVX_Vector v_cap = (factx->logit_softcap != 0.0f) ? hvx_vec_splat_f16(factx->logit_softcap) : Q6_V_vzero();
const HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
const HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
const uint32_t stride_v2 = factx->size_v_row_padded * 2;
for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) {
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
@@ -309,7 +377,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
// 2. Softcap (in FP16)
if (factx->logit_softcap != 0.0f) {
const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
scores_f16 = hvx_vec_tanh_f16(scores_f16);
scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap);
}
@@ -319,8 +386,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
// 3. Mask (in FP16)
if (mask) {
HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base;
HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf);
m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16);
@@ -335,10 +400,30 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
HVX_Vector v_max = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); // splat block max in FP32
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);
if (ib + 1 == factx->n_blocks && has_next_ir) {
// Queue next row's Q row!
dma_queue_push(dma, dma_make_ptr(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
if (factx->n_blocks % 2 == 0) {
// Queue next row's block 0 (into buffer slot 0)
uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
// K (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
// V (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
// Mask (block 0 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
}
}
}
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
{
const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
// 4. Online Softmax Update
HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec);
HVX_Vector diff_vec = HVX_OP_SUB_F32(M_vec, M_new_vec);
@@ -370,24 +455,20 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec);
// 5. Accumulate V (F16 * F16 -> F32 accumulator)
__fp16 __attribute__((aligned(128))) p_arr[VLEN_FP16];
hvx_vec_store_a(p_arr, 128, P);
const uint8_t * v_ptr = v_base;
for (uint32_t j = 0; j < current_block_size; j += 2) {
if (j + 1 == current_block_size) {
if (p_arr[j] != 0.0f) {
const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded;
hvx_mad_f32_f16_aa(VKQ32, v_ptr, (p_arr + j), DV);
}
HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
hvx_mad_f32_f16_aa_vec(VKQ32, v_ptr, S0, DV);
break;
}
if (p_arr[j] == 0.0f && p_arr[j + 1] == 0.0f) {
continue;
}
HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2));
const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded;
hvx_mad_f32_f16_aa_rx2(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, (p_arr + j), (p_arr + j + 1), DV);
hvx_mad_f32_f16_aa_rx2_vec(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, S0, S1, DV);
v_ptr += stride_v2;
}
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
@@ -414,6 +495,61 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
}
}
if (has_next_ir) {
if (factx->n_blocks % 2 == 0) {
// Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
if (factx->n_blocks > 1) {
uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
// K (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
// V (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
// Mask (block 1 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
}
}
} else {
// Queue next row's block 0 (into buffer slot 0)
{
uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
// K (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
// V (block 0 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
// Mask (block 0 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
}
}
// Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
if (factx->n_blocks > 1) {
uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
// K (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
// V (block 1 of next row)
dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
// Mask (block 1 of next row)
if (mask) {
dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
}
}
}
}
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir);
// sinks
float M = hvx_vec_get_f32(M_vec);
@@ -471,6 +607,7 @@ typedef struct {
void * curr_k;
uint32_t kv_start;
uint32_t rows_per_t;
size_t buf_idx;
} fa_k_int_args_t;
static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
@@ -488,19 +625,19 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK,
hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
args->src_stride, start, end);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
}
static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) {
static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start, size_t buf_idx) {
work_queue_t wp = factx->octx->ctx->work_queue;
uint32_t n = 1;
if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
n = factx->n_threads;
}
uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t };
fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t, buf_idx };
if (n > 1) {
work_queue_run(wp, fa_k_interleave_thread, &args, n);
} else {
@@ -645,12 +782,13 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
}
}
// Initialize vtcm_d_tiles to 0
// Initialize vtcm_d_tiles and vtcm_d_inv_l to 0
const size_t d_bytes_per_t = hex_align_up(d_tile_bytes / n, 128);
const size_t d_start = i * d_bytes_per_t;
const size_t d_end = hex_smin(d_start + d_bytes_per_t, d_tile_bytes);
if (d_start < d_tile_bytes) {
hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
hvx_splat_u8_a((char *) factx->vtcm_d_inv_l + d_start, 0, d_end - d_start);
}
}
@@ -662,15 +800,14 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
assert(factx->DK == factx->DV);
const size_t o_tile_bytes = factx->o_tile_bytes;
const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2));
const bool use_q_dma = (factx->vtcm_q_dma != NULL);
__fp16 * q_tiles = factx->vtcm_q_tiles;
if (use_q_dma) {
const size_t g_rows_end = hex_smin(end, n_rows_g);
const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;
uint8_t * q_flat = (uint8_t *) factx->vtcm_o_tiles[0];
uint8_t * q_flat = (uint8_t *) factx->vtcm_q_dma;
if (factx->is_q_fp32) {
switch (d_limit) {
case 2: hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
@@ -781,10 +918,10 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
for (size_t r = start; r < end; ++r) {
const size_t q_idx = fastdiv(r, &factx->div_G);
const size_t h_idx = fastmodulo(r, G, &factx->div_G);
size_t q_idx = fastdiv(start, &factx->div_G);
size_t h_idx = fastmodulo(start, G, &factx->div_G);
for (size_t r = start; r < end; ++r) {
float * out = (float *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
(q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);
@@ -801,6 +938,12 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
*(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
}
}
h_idx++;
if (h_idx == G) {
h_idx = 0;
q_idx++;
}
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
}
@@ -829,10 +972,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
for (size_t r = start; r < end; ++r) {
const size_t q_idx = fastdiv(r, &factx->div_G);
const size_t h_idx = fastmodulo(r, G, &factx->div_G);
size_t q_idx = fastdiv(start, &factx->div_G);
size_t h_idx = fastmodulo(start, G, &factx->div_G);
for (size_t r = start; r < end; ++r) {
__fp16 * out = (__fp16 *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
(q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);
@@ -851,6 +994,12 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
*(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
}
}
h_idx++;
if (h_idx == G) {
h_idx = 0;
q_idx++;
}
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
}
@@ -879,6 +1028,7 @@ static void fa_phase_o_store(struct hmx_fa_context * factx,
typedef struct {
struct hmx_fa_context * factx;
size_t buf_idx;
size_t kv_rows;
size_t n_rows_g;
size_t n_col_tiles;
@@ -960,8 +1110,8 @@ static inline void fa_softmax_impl(
uint32_t r0 = r / HMX_FP16_TILE_N_ROWS;
uint32_t r1 = r % HMX_FP16_TILE_N_ROWS;
const __fp16 * s_ld_base = factx->vtcm_s_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc;
__fp16 * p_st_base = factx->vtcm_p_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc;
const __fp16 * s_ld_base = factx->vtcm_s_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
__fp16 * p_st_base = factx->vtcm_p_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
// Decode 2 rows from S tiles into per-thread row buffers
if (has_softcap) {
@@ -983,7 +1133,26 @@ static inline void fa_softmax_impl(
my_row_buf1[ci] = hvx_vec_mul_f16_f16(t1, v_cap);
}
} else {
for (size_t c = 0; c < kv_rows; c += 64) {
size_t c = 0;
for (; c + 64 < kv_rows; c += 128) {
size_t ci0 = c / 64;
size_t ci1 = ci0 + 1;
const __fp16 * in_dtile0 = s_ld_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
const __fp16 * in_dtile1 = s_ld_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
const HVX_Vector * pv_s_in0_0 = ((const HVX_Vector *) in_dtile0) + r1 / 2;
const HVX_Vector * pv_s_in1_0 = pv_s_in0_0 + 16;
const HVX_Vector * pv_s_in0_1 = ((const HVX_Vector *) in_dtile1) + r1 / 2;
const HVX_Vector * pv_s_in1_1 = pv_s_in0_1 + 16;
HVX_VectorPair vp_s_drow0 = Q6_W_vdeal_VVR(*pv_s_in1_0, *pv_s_in0_0, -2);
my_row_buf0[ci0] = Q6_V_lo_W(vp_s_drow0);
my_row_buf1[ci0] = Q6_V_hi_W(vp_s_drow0);
HVX_VectorPair vp_s_drow1 = Q6_W_vdeal_VVR(*pv_s_in1_1, *pv_s_in0_1, -2);
my_row_buf0[ci1] = Q6_V_lo_W(vp_s_drow1);
my_row_buf1[ci1] = Q6_V_hi_W(vp_s_drow1);
}
for (; c < kv_rows; c += 64) {
size_t ci = c / 64;
const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
@@ -1007,12 +1176,12 @@ static inline void fa_softmax_impl(
HVX_Vector v_s_rowmax0 = v_neg_inf;
HVX_Vector v_s_rowmax1 = v_neg_inf;
for (size_t c = 0; c < kv_rows; c += 64) {
size_t ci = c / 64;
const size_t ne = hex_smin(kv_rows - c, 64);
HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
if (has_mask) {
for (size_t c = 0; c < kv_rows; c += 64) {
size_t ci = c / 64;
const size_t ne = hex_smin(kv_rows - c, 64);
HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
if (has_mask) {
HVX_Vector v_mask0, v_mask1;
if (mask_broadcast) {
@@ -1066,15 +1235,31 @@ static inline void fa_softmax_impl(
my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_mask0_scaled), v_neg_inf);
my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_mask1_scaled), v_neg_inf);
}
} else {
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
}
} else {
size_t c = 0;
for (; c + 64 < kv_rows; c += 128) {
size_t ci0 = c / 64;
size_t ci1 = ci0 + 1;
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci0]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci0]);
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci1]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci1]);
}
for (; c < kv_rows; c += 64) {
size_t ci = c / 64;
const size_t ne = hex_smin(kv_rows - c, 64);
HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
if (ne < 64) {
my_row_buf0[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf0[ci], v_neg_inf);
my_row_buf1[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf1[ci], v_neg_inf);
}
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
}
v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
}
v_s_rowmax0 = hvx_vec_reduce_max_f16(v_s_rowmax0);
@@ -1121,8 +1306,48 @@ static inline void fa_softmax_impl(
HVX_Vector v_p_rowsum0 = v_zero;
HVX_Vector v_p_rowsum1 = v_zero;
for (size_t c = 0; c < kv_rows; c += 64) {
size_t ci = c / 64;
size_t c = 0;
for (; c + 64 < kv_rows; c += 128) {
size_t ci0 = c / 64;
size_t ci1 = ci0 + 1;
HVX_Vector v_s_minus_m0_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci0], v_dup_m0);
HVX_Vector v_s_minus_m1_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci0], v_dup_m1);
HVX_Vector v_s_minus_m0_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci1], v_dup_m0);
HVX_Vector v_s_minus_m1_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci1], v_dup_m1);
HVX_Vector v_p_row0_hf_0 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_0));
HVX_Vector v_p_row1_hf_0 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_0));
HVX_Vector v_p_row0_hf_1 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_1));
HVX_Vector v_p_row1_hf_1 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_1));
__fp16 * out_dtile0 = p_st_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
__fp16 * out_dtile1 = p_st_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
HVX_Vector * pv_p_out0_0 = ((HVX_Vector *) out_dtile0) + r1 / 2;
HVX_Vector * pv_p_out1_0 = pv_p_out0_0 + 16;
HVX_Vector * pv_p_out0_1 = ((HVX_Vector *) out_dtile1) + r1 / 2;
HVX_Vector * pv_p_out1_1 = pv_p_out0_1 + 16;
HVX_VectorPair vp_p_dual0 = Q6_W_vshuff_VVR(v_p_row1_hf_0, v_p_row0_hf_0, -2);
*pv_p_out0_0 = Q6_V_lo_W(vp_p_dual0);
*pv_p_out1_0 = Q6_V_hi_W(vp_p_dual0);
HVX_VectorPair vp_p_dual1 = Q6_W_vshuff_VVR(v_p_row1_hf_1, v_p_row0_hf_1, -2);
*pv_p_out0_1 = Q6_V_lo_W(vp_p_dual1);
*pv_p_out1_1 = Q6_V_hi_W(vp_p_dual1);
HVX_VectorPair vp_p0_0 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_0);
HVX_VectorPair vp_p1_0 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_0);
HVX_VectorPair vp_p0_1 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_1);
HVX_VectorPair vp_p1_1 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_1);
v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_0), Q6_V_hi_W(vp_p0_0)));
v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_1), Q6_V_hi_W(vp_p0_1)));
v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_0), Q6_V_hi_W(vp_p1_0)));
v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_1), Q6_V_hi_W(vp_p1_1)));
}
for (size_t c_rem = c; c_rem < kv_rows; c_rem += 64) {
size_t ci = c_rem / 64;
HVX_Vector v_s_minus_m0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci], v_dup_m0);
HVX_Vector v_s_minus_m1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci], v_dup_m1);
@@ -1281,7 +1506,7 @@ static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_contex
v_content = Q6_V_vror_VR(v_content, 64);
}
__fp16 * out_base = factx->vtcm_d_tiles + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
__fp16 * out_base = factx->vtcm_d_inv_l + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
}
}
@@ -1514,6 +1739,27 @@ static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) {
}
}
static inline void fa_prefetch_block(dma_queue * dma, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask,
uint32_t b, size_t Bc, size_t size_k_row_padded, size_t size_k_row, size_t size_v_row_padded, size_t size_v_row,
uint32_t ik2, uint32_t ik3, uint32_t iv2, uint32_t iv3, uint32_t q_start, uint32_t im3, uint32_t kv_head, uint32_t G,
size_t m_line_bytes, size_t n_rows_q, size_t nek1, size_t prefetch_buf, struct hmx_fa_context * factx) {
const uint32_t prefetch_start = b * Bc;
const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
if (mask) {
if (__builtin_expect(factx->mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
dma_cache_push(dma, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
}
}
}
// ============================================================================
// Core HMX flash attention algorithm (GQA-merged)
// ============================================================================
@@ -1612,7 +1858,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
// Build the VTCM layout once (shared with the host estimator) and place every
// scratch buffer at its computed offset.
struct hmx_fa_vtcm_layout L;
hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline);
hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline, factx.is_q_fp32);
if (L.total_bytes > ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL;
@@ -1620,6 +1866,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
uint8_t * const base = ctx->vtcm_base;
factx.vtcm_q_dma = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_dma);
factx.vtcm_q_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles);
factx.vtcm_o_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]);
factx.vtcm_o_tiles[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]);
@@ -1627,12 +1874,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
factx.vtcm_k_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]);
factx.vtcm_v_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]);
factx.vtcm_v_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]);
factx.vtcm_k_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles);
factx.vtcm_k_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles[0]);
factx.vtcm_k_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_k_tiles[1], pipeline);
factx.vtcm_v_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]);
factx.vtcm_v_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline);
factx.vtcm_s_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles);
factx.vtcm_p_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles);
factx.vtcm_s_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles[0]);
factx.vtcm_s_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_s_tiles[1], pipeline);
factx.vtcm_p_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles[0]);
factx.vtcm_p_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_p_tiles[1], pipeline);
factx.vtcm_d_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
factx.vtcm_d_inv_l = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_inv_l);
factx.vtcm_m_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
factx.vtcm_l_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
factx.vtcm_s_rowmax = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax);
@@ -1670,6 +1921,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
const bool q_transposed = q->nb[1] < q->nb[2];
const size_t q_src_stride = q_transposed ? q->nb[2] : q->nb[1];
const size_t q_row_bytes_untransposed = factx.G * factx.DK * qo_element_size;
const size_t q_row_bytes_trans_factor = factx.DK * qo_element_size;
const uint32_t kv_rows0 = hex_smin(Bc, nek1);
// ======== Reusable job descriptors for pipeline ========
hmx_fa_qk_job_t qk_job;
hmx_fa_o_update_job_t ou_job;
@@ -1690,34 +1947,34 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const uint32_t iv2 = kv_head;
const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);
// 1. Push Q DMA (if Q DMA is used)
const size_t o_tile_bytes = factx.o_tile_bytes;
const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2));
if (use_q_dma) {
const bool q_transposed = q->nb[1] < q->nb[2];
const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size;
const size_t src_stride = q_transposed ? q->nb[2] : q->nb[1];
// 1. Push Q and KV DMAs for the very first iteration.
// Subsequent iterations are enqueued early at the end of the previous iteration.
if (ib3 == 0 && q_start == 0 && kv_head == 0) {
const uint8_t * q_ptr = (const uint8_t *) q->data;
const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
const size_t n_rows = q_transposed ? factx.G : n_rows_q;
dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows);
dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
if (factx.n_kv_blocks > 0) {
const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
if (factx.pipeline && mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
}
}
}
}
// 2. Prefetch first KV block
if (factx.n_kv_blocks > 0) {
const uint32_t kv_rows0 = hex_smin(Bc, nek1);
const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
}
// 3. Pop Q DMA (blocks until Q is loaded)
if (use_q_dma) {
dma_queue_pop(dma);
}
// 2. Pop Q DMA (blocks until Q is loaded)
dma_queue_pop(dma);
// ---- Load Q block & Initialize per-block state ----
fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
@@ -1738,76 +1995,40 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
hmx_queue_t hmx_q = ctx->hmx_queue;
if (factx.pipeline) {
// Pipeline path
// Double-buffered job structs because HMX queue runs asynchronously
hmx_fa_qk_job_t qk_job[2];
hmx_fa_o_update_job_t ou_job[2];
// Prefetch block 1 early if there are multiple blocks
if (factx.n_kv_blocks > 1) {
fa_prefetch_block(dma, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx);
}
// Prep and start QK-dot(0)
void * curr_k0 = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0);
qk_job[0].q_tiles = factx.vtcm_q_tiles;
qk_job[0].k_tiles = factx.vtcm_k_tiles[0];
qk_job[0].s_tiles = factx.vtcm_s_tiles[0];
qk_job[0].n_row_tiles = n_row_tiles;
qk_job[0].n_col_tiles = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
qk_job[0].n_dot_tiles = DK / 32;
qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
qk_job[0].hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
const uint32_t kv_start = kv_blk * Bc;
const uint32_t kv_rows = hex_smin(Bc, nek1 - kv_start);
const size_t n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
// Push mask DMA
if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
}
}
// Prefetch next KV block early
if (kv_blk + 1 < factx.n_kv_blocks) {
const uint32_t prefetch_start = (kv_blk + 1) * Bc;
const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
const size_t prefetch_buf = 1 - buf_idx;
const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
}
// ---- Phase 1: K_int ----
if (kv_blk > 0) {
ou_job.o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev;
ou_job.p_tiles = factx.vtcm_p_tiles;
ou_job.v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
ou_job.d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
ou_job.n_row_tiles = n_row_tiles;
ou_job.n_col_tiles = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
ou_job.n_tiles_per_bc = n_tiles_per_bc;
ou_job.DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
}
// Wait for current K DMA and interleave
void * curr_k = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
// ---- Phase 2: qk_dot ----
qk_job.q_tiles = factx.vtcm_q_tiles;
qk_job.k_tiles = factx.vtcm_k_tiles;
qk_job.s_tiles = factx.vtcm_s_tiles;
qk_job.n_row_tiles = n_row_tiles;
qk_job.n_col_tiles = n_col_tiles;
qk_job.n_dot_tiles = DK / 32;
qk_job.n_tiles_per_bc = n_tiles_per_bc;
qk_job.hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
// Wait for current V DMA and interleave
// ---- 1. Pop and run V-prep for current block ----
void * curr_v = dma_queue_pop(dma).dst;
fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);
if (kv_blk > 0) {
hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
}
hmx_queue_pop(hmx_q);
// ---- Phase 3: softmax + build_D ----
// ---- 2. Pop and run mask-prep for current block ----
__fp16 * current_mask_vtcm = NULL;
if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
@@ -1818,9 +2039,34 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
}
}
// ---- 3. Pop and run K-prep for next block & push next QK-dot ----
if (kv_blk + 1 < factx.n_kv_blocks) {
const uint32_t next_start = (kv_blk + 1) * Bc;
const uint32_t next_rows = hex_smin(Bc, nek1 - next_start);
const size_t next_buf = 1 - buf_idx;
void * next_k = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, next_rows, k_src_stride, next_k, next_start, next_buf);
qk_job[next_buf].q_tiles = factx.vtcm_q_tiles;
qk_job[next_buf].k_tiles = factx.vtcm_k_tiles[next_buf];
qk_job[next_buf].s_tiles = factx.vtcm_s_tiles[next_buf];
qk_job[next_buf].n_row_tiles = n_row_tiles;
qk_job[next_buf].n_col_tiles = hmx_ceil_div(next_rows, HMX_FP16_TILE_N_COLS);
qk_job[next_buf].n_dot_tiles = DK / 32;
qk_job[next_buf].n_tiles_per_bc = n_tiles_per_bc;
qk_job[next_buf].hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
}
// ---- 4. Wait for current block's QK-dot to finish ----
hmx_queue_pop(hmx_q);
// ---- 5. Phase 2: softmax + build_D ----
fa_softmax_args_t sargs;
memset(&sargs, 0, sizeof(sargs));
sargs.factx = &factx;
sargs.buf_idx = buf_idx;
sargs.kv_rows = kv_rows;
sargs.n_rows_g = n_rows_g;
sargs.n_col_tiles = n_col_tiles;
@@ -1838,8 +2084,39 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
sargs.mask_vtcm = current_mask_vtcm;
sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
sargs.slopes = factx.vtcm_slopes;
// Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx])
if (kv_blk > 0) {
const size_t prev_buf = 1 - buf_idx;
ou_job[prev_buf].o_curr = o_tile_curr;
ou_job[prev_buf].o_prev = o_tile_prev;
ou_job[prev_buf].p_tiles = factx.vtcm_p_tiles[prev_buf];
ou_job[prev_buf].v_tiles = factx.vtcm_v_tiles[prev_buf];
ou_job[prev_buf].d_tiles = factx.vtcm_d_tiles;
ou_job[prev_buf].hmx_scales = factx.vtcm_hmx_scales_id;
ou_job[prev_buf].n_row_tiles = n_row_tiles;
ou_job[prev_buf].n_col_tiles = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
ou_job[prev_buf].n_tiles_per_bc = n_tiles_per_bc;
ou_job[prev_buf].DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
}
// Run Softmax on HVX (blocking call)
fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
// Wait for HMX O update for block kv_blk - 1 to finish
if (kv_blk > 0) {
hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
}
// Prefetch block kv_blk + 2
if (kv_blk + 2 < factx.n_kv_blocks) {
fa_prefetch_block(dma, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx);
}
buf_idx = 1 - buf_idx;
}
@@ -1847,18 +2124,23 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
if (factx.n_kv_blocks > 0) {
const uint32_t last_blk = factx.n_kv_blocks - 1;
const size_t last_cols = hmx_ceil_div(hex_smin(Bc, nek1 - last_blk * Bc), HMX_FP16_TILE_N_COLS);
ou_job.o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev;
ou_job.p_tiles = factx.vtcm_p_tiles;
ou_job.v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
ou_job.d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
ou_job.n_row_tiles = n_row_tiles;
ou_job.n_col_tiles = last_cols;
ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
ou_job.n_tiles_per_bc = n_tiles_per_bc;
ou_job.DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
ou_job[0].o_curr = o_tile_curr;
ou_job[0].o_prev = o_tile_prev;
ou_job[0].p_tiles = factx.vtcm_p_tiles[1 - buf_idx];
ou_job[0].v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
ou_job[0].d_tiles = factx.vtcm_d_tiles;
ou_job[0].hmx_scales = factx.vtcm_hmx_scales_id;
ou_job[0].n_row_tiles = n_row_tiles;
ou_job[0].n_col_tiles = last_cols;
ou_job[0].n_row_tiles_g_br = n_row_tiles_g_br;
ou_job[0].n_tiles_per_bc = n_tiles_per_bc;
ou_job[0].DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[0]));
// Overlapped: run HVX build diag inv L while HMX is busy executing the update
htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
@@ -1892,12 +2174,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
// Wait for current K DMA and interleave
void * curr_k = dma_queue_pop(dma).dst;
fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start, 0);
{
qk_job.q_tiles = factx.vtcm_q_tiles;
qk_job.k_tiles = factx.vtcm_k_tiles;
qk_job.s_tiles = factx.vtcm_s_tiles;
qk_job.k_tiles = factx.vtcm_k_tiles[0];
qk_job.s_tiles = factx.vtcm_s_tiles[0];
qk_job.n_row_tiles = n_row_tiles;
qk_job.n_col_tiles = n_col_tiles;
qk_job.n_dot_tiles = (size_t) (DK / 32);
@@ -1948,7 +2230,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
{
ou_job.o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev;
ou_job.p_tiles = factx.vtcm_p_tiles;
ou_job.p_tiles = factx.vtcm_p_tiles[0];
ou_job.v_tiles = factx.vtcm_v_tiles[0];
ou_job.d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
@@ -1959,6 +2241,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
ou_job.DV = DV;
hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
if (kv_blk + 1 == factx.n_kv_blocks) {
// Overlapped: run HVX build diag inv L while HMX is busy executing the update
htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
}
hmx_queue_pop(ctx->hmx_queue);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
@@ -1968,15 +2256,63 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
}
}
// Enqueue DMAs for the next iteration early so they overlap with O-PROC
uint32_t next_kv_head = kv_head + 1;
uint32_t next_q_start = q_start;
uint32_t next_ib3 = ib3;
if (next_kv_head >= n_kv_heads) {
next_kv_head = 0;
next_q_start = q_start + Br;
if (next_q_start >= neq1) {
next_q_start = 0;
next_ib3 = ib3 + 1;
}
}
bool has_next = (next_ib3 < neq3);
if (has_next) {
const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start);
const uint8_t * next_q_ptr = (const uint8_t *) q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
const size_t next_n_rows = q_transposed ? factx.G : next_n_rows_q;
dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
if (factx.n_kv_blocks > 0) {
const uint32_t next_ik2 = next_kv_head;
const uint32_t next_iv2 = next_kv_head;
uint32_t next_ik3 = ik3;
uint32_t next_iv3 = iv3;
if (next_ib3 != ib3) {
next_ik3 = fastdiv(next_ib3, &kparams->broadcast_rk3);
next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3);
}
const uint8_t * next_k_src = (const uint8_t *) k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * next_v_src = (const uint8_t *) v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
if (factx.pipeline && mask) {
uint32_t next_im3 = im3;
if (next_ib3 != ib3) {
next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3);
}
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
} else {
fa_push_mask_dma_gqa(dma, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
}
}
}
}
// ---- Final normalization ----
{
htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
on_job.o_curr = o_tile_curr;
on_job.o_prev = o_tile_prev;
on_job.d_tiles = factx.vtcm_d_tiles;
on_job.d_tiles = factx.vtcm_d_inv_l;
on_job.hmx_scales = factx.vtcm_hmx_scales_id;
on_job.n_row_tiles = n_row_tiles;
on_job.n_row_tiles_g_br = n_row_tiles_g_br;
ggml/src/ggml-hexagon/htp/flash-attn-ops.h
+57 -25
View File
@@ -101,14 +101,16 @@ static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params
struct hmx_fa_vtcm_layout {
// Byte offsets from vtcm_base for each region.
size_t off_q_tiles;
size_t off_q_dma;
size_t off_o_tiles[2];
size_t off_k_fp16[2];
size_t off_v_fp16[2];
size_t off_k_tiles;
size_t off_v_tiles[2]; // [1] allocated only when pipeline, else 0
size_t off_s_tiles;
size_t off_p_tiles;
size_t off_k_tiles[2];
size_t off_v_tiles[2];
size_t off_s_tiles[2];
size_t off_p_tiles[2];
size_t off_d_tiles;
size_t off_d_inv_l;
size_t off_m_vec;
size_t off_l_vec;
size_t off_s_rowmax;
@@ -140,7 +142,7 @@ struct hmx_fa_vtcm_layout {
static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
size_t gqa_factor, size_t DK, size_t DV,
size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
@@ -149,6 +151,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t q_dma_size = hex_align_up(g_br * DK * (is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128);
const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256);
@@ -160,27 +163,47 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
size_t off = 0;
// Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned)
// Group A (Part 1 - HMX Tiled buffers)
VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size);
VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size);
VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size);
VTCM_LAYOUT_ALLOC(off, off_k_tiles, k_tile_size);
VTCM_LAYOUT_ALLOC(off, off_v_tiles[0], v_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off, off_s_tiles, s_tile_size);
VTCM_LAYOUT_ALLOC(off, off_p_tiles, s_tile_size);
VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size);
VTCM_LAYOUT_ALLOC(off, off_d_inv_l, d_tile_size);
// Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned)
// Group B & C share start offset (Group B tiles must be 2KB aligned)
size_t off_group_b_c = hex_align_up(off, HTP_FA_HMX_TILE_SIZE);
// Group B: Compute-only buffers
size_t off_group_b = off_group_b_c;
VTCM_LAYOUT_ALLOC(off_group_b, off_k_tiles[0], k_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_k_tiles[1], k_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_v_tiles[0], v_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_v_tiles[1], v_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_s_tiles[0], s_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_s_tiles[1], s_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_p_tiles[0], s_tile_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_p_tiles[1], s_tile_size, pipeline);
VTCM_LAYOUT_ALLOC(off_group_b, off_s_rowmax, col_vec_size);
VTCM_LAYOUT_ALLOC(off_group_b, off_p_rowsum, col_vec_size);
VTCM_LAYOUT_ALLOC(off_group_b, off_row_bufs, row_vec_size * 2 * n_threads);
const size_t group_b_size = off_group_b - off_group_b_c;
// Group C: Q fetch DMA buffer
size_t off_group_c = off_group_b_c;
VTCM_LAYOUT_ALLOC(off_group_c, off_q_dma, q_dma_size);
const size_t group_c_size = off_group_c - off_group_b_c;
off = off_group_b_c + hex_smax(group_b_size, group_c_size);
// Group A (Part 2 - remaining non-HMX buffers)
VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size);
VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size);
VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size);
VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size);
VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_s_rowmax, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_p_rowsum, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_row_bufs, row_vec_size * 2 * n_threads);
VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256);
VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size);
@@ -200,9 +223,9 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
}
// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
struct hmx_fa_vtcm_layout L;
hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline);
hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32);
return L.total_bytes;
}
@@ -239,7 +262,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
size_t qo_len,
size_t kv_len,
size_t vtcm_budget,
size_t n_threads) {
size_t n_threads,
bool is_q_fp32) {
const size_t T = HMX_FP16_TILE_N_ROWS; // 32
const size_t br_unit = hmx_ceil_div(T, gqa_factor);
const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64
@@ -253,8 +277,9 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
const size_t Bc_limit = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
(kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
// Cost coefficients calibrated from profiling
const size_t c_q_fixed = 1400; // per-Q-block: q_load + epilogue o_update + o_norm + o_store
const size_t c_iter_fixed = 200; // per-KV-iter: HMX queue push/pop + DMA pop + barriers
const size_t c_q_fixed = 800; // per-Q-block: q_load + epilogue o_update + o_norm + o_store
const size_t c_iter_base = 200; // per-KV-iter base (HMX dot/update + DMA)
const size_t c_softmax = 600; // per 64-row vector chunk on HVX
size_t best_cost = SIZE_MAX, best_mn = 0;
size_t best_Br = 0, best_Bc = 0;
@@ -262,13 +287,20 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
// Try all Bc candidates from Bc_limit down to bc_unit
for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline);
size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32);
if (vtcm_needed <= vtcm_budget) {
// This Bc fits for this Br!
const size_t q_blocks = (qo_len + Br - 1) / Br;
const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
const size_t mn = Br * Bc;
const size_t q_blocks = (qo_len + Br - 1) / Br;
const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
const size_t actual_threads = (kv_blocks >= 3 && n_threads >= 2) ? n_threads : 1;
const size_t n_rows_g = Br * gqa_factor;
const size_t n_row_vec_cnt = (n_rows_g + 63) / 64;
const size_t n_use = n_row_vec_cnt < actual_threads ? n_row_vec_cnt : actual_threads;
const size_t vecs_per_t = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
const size_t c_iter_actual = c_iter_base + c_softmax * vecs_per_t;
const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
const size_t mn = Br * Bc;
if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
best_cost = cost;
ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
+23 -5
View File
@@ -767,23 +767,25 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co
// output : fp16 -> f32p
static void transfer_output_chunk_fp16_to_fp32(
static void transfer_output_chunk_fp16_to_fp32_col_chunk(
float *restrict dst,
const float *restrict src2,
const __fp16 *restrict vtcm_src,
uint32_t start_row,
uint32_t n_rows,
uint32_t n_cols,
uint32_t c_len,
uint32_t total_n_cols,
uint32_t dst_stride,
uint32_t src2_stride,
uint32_t dst_cols
) {
assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
assert(c_len % HTP_MM_HMX_TILE_N_COLS == 0);
assert(total_n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
const size_t tile_row_stride = (total_n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
const HVX_Vector one = hvx_vec_splat_f16(1.0);
const size_t limit_c = hex_smin(n_cols, dst_cols);
const size_t limit_c = hex_smin(c_len, dst_cols);
const size_t limit_c_aligned = (limit_c & ~31);
for (size_t r = 0; r < n_rows; r += 2) {
@@ -848,6 +850,22 @@ static void transfer_output_chunk_fp16_to_fp32(
}
}
static inline void transfer_output_chunk_fp16_to_fp32(
float *restrict dst,
const float *restrict src2,
const __fp16 *restrict vtcm_src,
uint32_t start_row,
uint32_t n_rows,
uint32_t n_cols,
uint32_t dst_stride,
uint32_t src2_stride,
uint32_t dst_cols
) {
transfer_output_chunk_fp16_to_fp32_col_chunk(
dst, src2, vtcm_src, start_row, n_rows, n_cols, n_cols, dst_stride, src2_stride, dst_cols
);
}
typedef struct {
const __fp16 *vtcm_src;
float *dst;
ggml/src/ggml-hexagon/htp/htp-ctx.h
+8 -1
View File
@@ -19,6 +19,8 @@
#endif
#define HTP_MAX_MMAPS 16
#define HTP_MAX_DIRTY_RANGES 16
// Memory mapping
struct htp_mmap {
uint64_t size;
@@ -95,7 +97,11 @@ struct htp_context {
atomic_bool vtcm_needs_release;
uint64_t max_vmem;
uint32_t dirty_map[HTP_OP_MAX_TENSORS / 32];
struct htp_dirty_range {
uint32_t start;
uint32_t end;
uint32_t bi;
} dirty_ranges[HTP_MAX_DIRTY_RANGES];
// Persistent DDR scratchpad for MUL_MAT_ID mappings
void * ddr_spad_base;
@@ -134,5 +140,6 @@ int op_diag(struct htp_ops_context * octx);
int op_solve_tri(struct htp_ops_context * octx);
int op_gated_delta_net(struct htp_ops_context * octx);
int op_pad(struct htp_ops_context * octx);
int op_im2col(struct htp_ops_context * octx);
#endif /* HTP_CTX_H */
ggml/src/ggml-hexagon/htp/htp-ops.h
+7 -2
View File
@@ -98,6 +98,7 @@ enum htp_op_code {
HTP_OP_NORM,
HTP_OP_CONCAT,
HTP_OP_CLAMP,
HTP_OP_IM2COL,
HTP_OP_INVALID
};
@@ -123,7 +124,7 @@ enum htp_tensor_flags {
// Tensor descriptor
struct htp_tensor {
uint32_t data; // Buffer offset in the messages, and data pointer on the NPU
uint32_t alias; // Index of the canonical tensor for this memory buffer
uint32_t reserved; // Reserved for alignment padding (must be multiple of 8)
uint32_t size; // Data size in bytes
uint32_t flags; // Buffer / tensor flags
uint32_t type; // Data type
@@ -173,6 +174,7 @@ enum htp_trace_event_id {
HTP_TRACE_EVT_DMA = 0,
HTP_TRACE_EVT_L2FLUSH = 1,
HTP_TRACE_EVT_INIT = 2,
HTP_TRACE_EVT_BUFF = 3,
HTP_TRACE_EVT_HVX_COMP = 20,
HTP_TRACE_EVT_HVX_A_QUANT = 21,
@@ -225,7 +227,10 @@ struct htp_opbatch_rsp {
uint32_t n_tensors; // Number of tensors
uint32_t n_ops; // Number of op profile descriptors
uint32_t n_traces[HTP_MAX_NTHREADS + 1];
uint8_t pad[8]; // align to 8 bytes
uint32_t usecs; // Number of usec
uint32_t pad; // align to 8 bytes
uint64_t cycles_start; // Start cycle counter
uint64_t cycles_stop; // Stop cycle counter
// struct htp_prof_desc profs[]; -- dspqueue buf 0
};
ggml/src/ggml-hexagon/htp/htp-tensor.c
+188 -99
View File
@@ -2,6 +2,7 @@
#include <qurt.h>
#include <qurt_memory.h>
#include <HAP_farf.h>
#include "hex-common.h"
#include "hex-utils.h"
@@ -10,84 +11,6 @@
#include "htp-ctx.h"
#include "work-queue.h"
struct l2flush_task {
struct htp_thread_trace * trace;
uint32_t start;
uint32_t end;
uint32_t chunk_size;
uint32_t ti;
};
static void l2flush_thread_worker(unsigned int n, unsigned int i, void * data) {
struct l2flush_task * task = (struct l2flush_task *) data;
const uint32_t start = task->start;
const uint32_t end = task->end;
const uint32_t ti = task->ti;
const uint32_t chunk_size = task->chunk_size;
const uint32_t thread_s = start + i * chunk_size;
if (thread_s >= end) {
return;
}
uint32_t thread_e = thread_s + chunk_size;
if (thread_e > end) {
thread_e = end;
}
struct htp_thread_trace * tr = &task->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, ti);
hex_l2flush((void *) (uintptr_t) thread_s, thread_e - thread_s);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, ti);
}
static void flush_all_dcache(struct htp_context * ctx) {
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
}
static void flush_tensor_range(struct htp_context * ctx, const struct htp_tensor * t) {
struct htp_thread_trace * tr = &ctx->trace[0];
if (t->size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
struct l2flush_task task;
task.start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
task.end = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
task.ti = t->ti;
task.trace = ctx->trace;
const uint32_t total_size = task.end - task.start;
const uint32_t n_blocks = (total_size + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
const uint32_t blocks_per_thread = fastdiv(n_blocks + ctx->n_threads - 1, &ctx->n_threads_div);
task.chunk_size = blocks_per_thread * HEX_L2_BLOCK_SIZE;
work_queue_run(ctx->work_queue, l2flush_thread_worker, &task, ctx->n_threads);
} else {
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
hex_l2flush((void *) t->data, t->size);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
}
htp_tensor_make_clean(t, ctx->dirty_map);
}
void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t) {
if (!bitmap_test(ctx->dirty_map, t->ti)) {
return;
}
if (t->size > HEX_L2_FLUSH_ALL_THRESHOLD) {
flush_all_dcache(ctx);
return;
}
flush_tensor_range(ctx, t);
}
// One dirty tensor's line-aligned range, placed in the flattened global block space.
struct l2flush_range {
uint32_t start; // line-aligned start address
uint32_t end; // line-aligned end address
@@ -103,9 +26,18 @@ struct l2flush_multi_task {
uint32_t blocks_per_thread;
};
static void flush_all_dcache(struct htp_context * ctx) {
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
}
static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
(void) n;
struct l2flush_multi_task * task = (struct l2flush_multi_task *) data;
(void) n;
const uint32_t gb_first = i * task->blocks_per_thread;
uint32_t gb_last = gb_first + task->blocks_per_thread;
@@ -141,11 +73,177 @@ static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
}
void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
uint64_t total_dirty = 0;
void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
uint32_t n_pending = 0;
for (uint32_t i = 0; i < n; i++) {
const struct htp_tensor * t = tensors[i];
if (t && bitmap_test(ctx->dirty_map, t->ti)) {
if (!t) continue;
uint32_t t_start = t->data;
uint32_t t_end = t_start + t->size;
bool merged = false;
for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[j];
if (!r->start) continue;
if (r->start <= t_end && t_start <= r->end) {
uint32_t new_start = (t_start < r->start) ? t_start : r->start;
uint32_t new_end = (t_end > r->end) ? t_end : r->end;
r->start = new_start;
r->end = new_end;
merged = true;
}
}
if (!merged) {
pending[n_pending++] = t;
}
}
if (n_pending == 0) {
return;
}
uint32_t empty_indices[HTP_MAX_DIRTY_RANGES];
uint32_t active_indices[HTP_MAX_DIRTY_RANGES];
uint32_t n_active = 0;
uint32_t n_empty = 0;
for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
if (ctx->dirty_ranges[j].start) {
active_indices[n_active++] = j;
} else {
empty_indices[n_empty++] = j;
}
}
if (n_pending <= n_empty) {
for (uint32_t i = 0; i < n_pending; i++) {
uint32_t idx = empty_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
r->start = pending[i]->data;
r->end = pending[i]->data + pending[i]->size;
r->bi = pending[i]->bi;
}
return;
}
uint32_t n_evict = n_pending - n_empty;
uint32_t total_evict_size = 0;
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
total_evict_size += r->end - r->start;
}
if (total_evict_size > HEX_L2_FLUSH_ALL_THRESHOLD) {
flush_all_dcache(ctx);
for (uint32_t i = 0; i < n_pending; i++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[i];
r->start = pending[i]->data;
r->end = pending[i]->data + pending[i]->size;
r->bi = pending[i]->bi;
}
return;
}
if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_OP_MAX_INPUTS) {
struct l2flush_multi_task task;
task.trace = ctx->trace;
task.n_ranges = n_evict;
uint32_t block_acc = 0;
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
struct l2flush_range * rg = &task.ranges[i];
rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
rg->end = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
rg->block_first = block_acc;
rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
block_acc += rg->n_blocks;
}
task.total_blocks = block_acc;
task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
} else {
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
uint32_t size = r->end - r->start;
hex_l2flush((void *) (uintptr_t) r->start, size);
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
}
for (uint32_t i = 0; i < n_evict; i++) {
uint32_t idx = active_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
r->start = pending[i]->data;
r->end = pending[i]->data + pending[i]->size;
r->bi = pending[i]->bi;
}
for (uint32_t i = 0; i < n_empty; i++) {
uint32_t idx = empty_indices[i];
struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
r->start = pending[n_evict + i]->data;
r->end = pending[n_evict + i]->data + pending[n_evict + i]->size;
r->bi = pending[n_evict + i]->bi;
}
}
static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
uint32_t t_start = t->data;
uint32_t t_end = t_start + t->size;
for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[i];
if (!r->start) continue;
if (r->start < t_end && t_start < r->end) {
if (t_start <= r->start && r->end <= t_end) {
r->start = 0;
} else if (t_start <= r->start) {
r->start = t_end;
} else if (r->end <= t_end) {
r->end = t_start;
}
}
}
}
static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) {
uint32_t t_start = t->data;
uint32_t t_end = t_start + t->size;
for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
struct htp_dirty_range * r = &ctx->dirty_ranges[i];
if (!r->start) continue;
if (r->start < t_end && t_start < r->end) {
return true;
}
}
return false;
}
void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
uint32_t n_dirty = 0;
uint64_t total_dirty = 0;
for (uint32_t i = 0; i < n; i++) {
const struct htp_tensor * t = tensors[i];
if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) {
dirty_tensors[n_dirty++] = t;
total_dirty += t->size;
}
}
@@ -159,21 +257,15 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
return;
}
// Aggregate is small enough to walk. Thread it across all dirty ranges at once
// when it is worth the dispatch, otherwise flush sequentially.
if (total_dirty > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
struct l2flush_multi_task task;
task.trace = ctx->trace;
task.n_ranges = 0;
uint32_t block_acc = 0;
for (uint32_t i = 0; i < n; i++) {
const struct htp_tensor * t = tensors[i];
if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
continue;
}
// Clear as we go: dedups a tensor passed as multiple srcs (e.g. mul(x,x)).
htp_tensor_make_clean(t, ctx->dirty_map);
for (uint32_t i = 0; i < n_dirty; i++) {
const struct htp_tensor * t = dirty_tensors[i];
make_tensor_clean(ctx, t);
struct l2flush_range * rg = &task.ranges[task.n_ranges++];
rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
@@ -191,14 +283,11 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
}
struct htp_thread_trace * tr = &ctx->trace[0];
for (uint32_t i = 0; i < n; i++) {
const struct htp_tensor * t = tensors[i];
if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
continue;
}
for (uint32_t i = 0; i < n_dirty; i++) {
const struct htp_tensor * t = dirty_tensors[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
hex_l2flush((void *) t->data, t->size);
hex_l2flush((void *) (uintptr_t) t->data, t->size);
htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
htp_tensor_make_clean(t, ctx->dirty_map);
make_tensor_clean(ctx, t);
}
}
ggml/src/ggml-hexagon/htp/htp-tensor.h
+1 -17
View File
@@ -5,10 +5,6 @@
#include "htp-ops.h"
#include "hex-bitmap.h"
static inline struct htp_tensor * htp_tensor_alias(const struct htp_tensor * t) {
return (struct htp_tensor *) (uintptr_t) t->alias;
}
static inline void * htp_tensor_data(const struct htp_tensor * t) {
return (void *) (uintptr_t) t->data;
}
@@ -17,20 +13,8 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
return (uint32_t *) &t->flags;
}
static inline void htp_tensor_make_dirty(const struct htp_tensor * t, uint32_t * dirty_map) {
struct htp_tensor * curr = (struct htp_tensor *) t;
do {
bitmap_set(dirty_map, curr->ti);
curr = htp_tensor_alias(curr);
} while (curr != t);
}
static inline void htp_tensor_make_clean(const struct htp_tensor * t, uint32_t * dirty_map) {
bitmap_clear(dirty_map, t->ti);
}
struct htp_context;
void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t);
void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
#endif // HTP_TENSOR_H
ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h
+71
View File
@@ -208,6 +208,77 @@ static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restr
}
}
}
static inline void hvx_mad_f32_f16_aa_vec(float * restrict y, const void * restrict x, HVX_Vector S0, uint32_t n) {
const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
HVX_Vector * restrict vy = (HVX_Vector *) y;
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
uint32_t nloe = n % VLEN_FP16; // leftover elements
uint32_t i = 0;
#pragma unroll(2)
for (i = 0; i < nvec; ++i) {
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
}
if (nloe) {
HVX_VectorPair xy_p = vy_p[i];
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
HVX_Vector xy = Q6_V_lo_W(xy_p);
i = 2 * i; // index for vy
if (nloe >= VLEN_FP32) {
vy[i] = xy;
nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
}
if (nloe) {
hvx_vec_store_a(&vy[i], nloe * 4, xy);
}
}
}
static inline void hvx_mad_f32_f16_aa_rx2_vec(float * restrict y, const void * restrict x0, const void * restrict x1,
HVX_Vector S0, HVX_Vector S1, uint32_t n) {
const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
HVX_Vector * restrict vy = (HVX_Vector *) y;
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
uint32_t nloe = n % VLEN_FP16; // leftover elements
uint32_t i = 0;
#pragma unroll(2)
for (i = 0; i < nvec; ++i) {
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
}
if (nloe) {
HVX_VectorPair xy_p = vy_p[i];
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
HVX_Vector xy = Q6_V_lo_W(xy_p);
i = 2 * i; // index for vy
if (nloe >= VLEN_FP32) {
vy[i] = xy;
nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
}
if (nloe) {
hvx_vec_store_a(&vy[i], nloe * 4, xy);
}
}
}
static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) {
assert((size_t) dst % 128 == 0);
ggml/src/ggml-hexagon/htp/hvx-reduce.h
+40
View File
@@ -286,6 +286,46 @@ static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const i
}
}
// Signed 32-bit Integer Max variants
static inline HVX_Vector hvx_vec_reduce_max_n_i32(HVX_Vector in, unsigned int n) {
unsigned int total = n * 4; // total vec nbytes
unsigned int width = 4; // int32 nbytes
HVX_Vector max_val = in, max_t;
while (width < total) {
max_t = Q6_V_vror_VR(max_val, width); // rotate right
max_val = Q6_Vw_vmax_VwVw(max_t, max_val); // elementwise signed max
width = width << 1;
}
return max_val;
}
static inline HVX_Vector hvx_vec_reduce_max_i32(HVX_Vector in) {
return hvx_vec_reduce_max_n_i32(in, 32);
}
static inline int32_t hvx_reduce_max_i32_a(const uint8_t * restrict src, const int num_elems) {
HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
assert((uintptr_t) src % 128 == 0);
hvx_reduce_loop_body(HVX_Vector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
}
static inline int32_t hvx_reduce_max_i32_u(const uint8_t * restrict src, const int num_elems) {
HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
hvx_reduce_loop_body(HVX_UVector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
}
static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int num_elems) {
if (hex_is_aligned((void *) src, 128)) {
return hvx_reduce_max_i32_a(src, num_elems);
} else {
return hvx_reduce_max_i32_u(src, num_elems);
}
}
#undef hvx_reduce_loop_body
#undef HVX_REDUCE_MAX_OP
#undef HVX_REDUCE_SUM_OP
ggml/src/ggml-hexagon/htp/im2col-ops.c
+306
View File
@@ -0,0 +1,306 @@
#pragma clang diagnostic ignored "-Wunused-variable"
#pragma clang diagnostic ignored "-Wunused-function"
#pragma clang diagnostic ignored "-Wunused-but-set-variable"
#include <HAP_farf.h>
#include <HAP_perf.h>
#include <hexagon_protos.h>
#include <hexagon_types.h>
#include <string.h>
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "htp-ctx.h"
#include "htp-ops.h"
#include "hvx-utils.h"
#include "hex-dma.h"
#include "hex-profile.h"
#include "htp-vtcm.h"
struct htp_im2col_context {
struct htp_ops_context * octx;
uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel)
uint32_t pe_rows_per_thread; // N*OH rows per worker
uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256
uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256
// Patch-embed DMA path VTCM ping-pong.
uint8_t * pe_vtcm_src; // base of the 2x src buffers region
uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region
uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes
uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes
};
// Per-op VTCM layout for the patch-embed DMA path
struct htp_im2col_vtcm_layout {
size_t off_src;
size_t off_dst;
size_t src_bytes_per_thread;
size_t dst_bytes_per_thread;
size_t total_bytes;
};
static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L,
size_t src_row_bytes,
size_t dst_row_bytes,
uint32_t n_threads) {
L->src_bytes_per_thread = 2 * src_row_bytes;
L->dst_bytes_per_thread = 2 * dst_row_bytes;
L->off_src = 0;
L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads;
L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
}
#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
struct htp_ops_context * octx = ictx->octx; \
struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
const struct htp_tensor * restrict src1 = octx->src[1]; \
const struct htp_tensor * restrict dst = octx->dst; \
const int32_t s0 = octx->op_params[0]; \
const int32_t s1 = octx->op_params[1]; \
const int32_t p0 = octx->op_params[2]; \
const int32_t p1 = octx->op_params[3]; \
const int32_t d0 = octx->op_params[4]; \
const int32_t d1 = octx->op_params[5]; \
const uint32_t N = src1->ne[3]; \
const uint32_t IC = src1->ne[2]; \
const uint32_t IH = src1->ne[1]; \
const uint32_t IW = src1->ne[0]; \
const uint32_t KH = octx->src[0]->ne[1]; \
const uint32_t KW = octx->src[0]->ne[0]; \
const uint32_t OH = dst->ne[2]; \
const uint32_t OW = dst->ne[1]; \
const uint32_t patch_stride = IC * KH * KW; \
const float * restrict src_data = (const float *) src1->data; \
DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
const uint32_t npatches = N * OH * OW; \
const uint32_t patch_start = ictx->npatches_per_thread * ith; \
const uint32_t patch_end = MIN(patch_start + ictx->npatches_per_thread, npatches); \
if (patch_start >= patch_end) { \
return; \
} \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
for (uint32_t p = patch_start; p < patch_end; p++) { \
const uint32_t iow = p % OW; \
const uint32_t ioh = (p / OW) % OH; \
const uint32_t in = p / (OW * OH); \
DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \
for (uint32_t iic = 0; iic < IC; iic++) { \
const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \
DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
if (iih < 0 || iih >= (int32_t) IH) { \
SPLAT_FN(out_run, 0.0f, KW); \
continue; \
} \
const int32_t iiw0 = (int32_t) iow * s0 - p0; \
const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \
if (d0 == 1) { \
/* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \
const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \
int32_t hi = (int32_t) IW - iiw0; \
if (hi > (int32_t) KW) { \
hi = (int32_t) KW; \
} \
if (hi <= lo) { \
SPLAT_FN(out_run, 0.0f, KW); \
} else { \
if (lo > 0) { \
SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \
} \
COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \
(uint32_t) (hi - lo)); \
if (hi < (int32_t) KW) { \
SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \
} \
} \
continue; \
} \
for (uint32_t ikw = 0; ikw < KW; ikw++) { \
const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \
out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \
(DST_CTYPE) 0.0f : \
(DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \
} \
} \
} \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
}
IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")
#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
struct htp_ops_context * octx = ictx->octx; \
struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
const struct htp_tensor * restrict src1 = octx->src[1]; \
const struct htp_tensor * restrict dst = octx->dst; \
const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \
const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \
const uint32_t patch_stride = IC * KH * KW; \
const float * restrict src_data = (const float *) src1->data; \
DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
dma_queue * dmaq = octx->ctx->dma[ith]; \
uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \
uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \
float * srcb = (float *) src_base; \
DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \
const uint32_t nrows = N * OH; \
const uint32_t per_thread = ictx->pe_rows_per_thread; \
const uint32_t row_start = per_thread * ith; \
const uint32_t row_end = MIN(row_start + per_thread, nrows); \
if (row_start >= row_end) \
return; \
for (uint32_t r = row_start; r < row_end; r++) { \
const uint32_t in = r / OH; \
const uint32_t ioh = r % OH; \
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
int ok = (iih >= 0 && iih < (int32_t) IH); \
for (uint32_t iic = 0; iic < IC; iic++) { \
float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \
const float * _vsrc = \
ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \
dma_queue_push_ddr_to_vtcm( \
dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \
} \
} \
for (uint32_t i = 0; i < IC * KH; i++) \
dma_queue_pop(dmaq); \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \
for (uint32_t iow = 0; iow < OW; iow++) { \
DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
for (uint32_t iic = 0; iic < IC; iic++) { \
DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
if (iih < 0 || iih >= (int32_t) IH) { \
SPLAT_FN(out_run, 0.0f, KW); \
continue; \
} \
const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \
COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \
} \
} \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \
DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \
dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \
OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \
dma_queue_flush(dmaq); \
} \
}
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32")
static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];
const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];
const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];
const int is_2D = octx->op_params[6] == 1;
if (!is_2D) {
return false;
}
if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
return false;
}
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
return false; // non-overlapping
}
if (p0 != 0 || p1 != 0) {
return false; // no padding
}
if (d0 != 1 || d1 != 1) {
return false; // no dilation
}
return true;
}
// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path.
// Returns false if it doesn't fit the VTCM budget (caller falls back).
static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx,
struct htp_im2col_context * ictx,
uint32_t n_threads) {
const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
const uint32_t OW = octx->dst->ne[1];
const uint32_t patch_stride = IC * KH * KW;
ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256);
const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256);
// 2 src + 2 dst buffers per thread (ping-pong), src region first then dst.
struct htp_im2col_vtcm_layout L;
htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads);
if (L.total_bytes > octx->ctx->vtcm_size) {
return false;
}
uint8_t * const base = octx->ctx->vtcm_base;
ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
return true;
}
int op_im2col(struct htp_ops_context * octx) {
const struct htp_tensor * src1 = octx->src[1];
const struct htp_tensor * dst = octx->dst;
if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) {
FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported");
return HTP_STATUS_NO_SUPPORT;
}
const uint32_t N = src1->ne[3];
const uint32_t OH = dst->ne[2];
const uint32_t OW = dst->ne[1];
const uint32_t npatches = N * OH * OW;
const uint32_t n_threads = MIN(octx->n_threads, npatches);
if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
return HTP_STATUS_OK;
}
struct htp_im2col_context ictx = { 0 };
ictx.octx = octx;
ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;
// Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
// everything else (padding/dilation/stride edges) -> pure-DDR kernel.
if (im2col_use_patchembed_dma(octx)) {
const uint32_t nrows = N * OH;
const uint32_t pth = MIN(octx->n_threads, nrows);
if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
if (dst->type == HTP_TYPE_F16) {
work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
} else {
work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
}
return HTP_STATUS_OK;
}
// else: doesn't fit -> fall through to the pure-DDR kernel below.
}
if (dst->type == HTP_TYPE_F16) {
work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
} else {
work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads);
}
return HTP_STATUS_OK;
}
ggml/src/ggml-hexagon/htp/main.c
+40 -27
View File
@@ -781,6 +781,9 @@ static int execute_op(struct htp_ops_context * octx) {
case HTP_OP_PAD:
return op_pad(octx);
case HTP_OP_IM2COL:
return op_im2col(octx);
case HTP_OP_CONCAT:
return op_concat(octx);
@@ -901,10 +904,8 @@ static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, stru
uint32_t offset = t->data;
uint32_t size = t->size;
uint32_t bi = t->bi;
uint32_t alias = t->alias;
t->data = (uint32_t) (bufs[bi].base + offset); // update data to the actual pointer
t->alias = (uint32_t) (tens + alias); // update alias to the actual pointer
FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
@@ -955,14 +956,14 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
octx->dsts[i] = dst;
octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
htp_tensor_make_dirty(dst, octx->ctx->dirty_map);
FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
}
int status = execute_op(octx);
htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
octx->src0_spad.src = NULL;
octx->src1_spad.src = NULL;
octx->src2_spad.src = NULL;
@@ -994,12 +995,6 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);
// Clean cache at the start of the batch
// We cant trace this part because the trace buffer is setup later
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
// Setup descriptor pointers
uint8_t * m_ptr = dbuf->ptr;
struct htp_buf_desc* bufs = (struct htp_buf_desc*) m_ptr; m_ptr += b_size;
@@ -1007,13 +1002,8 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
struct htp_op_desc* ops = (struct htp_op_desc*) m_ptr; m_ptr += o_size;
struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;
prep_op_bufs(ctx, bufs, n_bufs);
prep_tensors(ctx, bufs, tens, n_tens);
struct htp_ops_context *octx = &ctx->octx;
memset(octx, 0, sizeof(*octx));
octx->n_threads = ctx->n_threads;
octx->ctx = ctx;
struct profile_data batch_prof;
profile_start(HTP_PROF_BASIC, &batch_prof);
memset(ctx->trace, 0, sizeof(ctx->trace));
if (ctx->profiler == HTP_PROF_TRACE) {
@@ -1024,6 +1014,24 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
}
}
// Clean cache at the start of the batch
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
hex_l2fetch_block(ctx, ctx->footprint);
memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
prep_op_bufs(ctx, bufs, n_bufs);
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
prep_tensors(ctx, bufs, tens, n_tens);
struct htp_ops_context *octx = &ctx->octx;
memset(octx, 0, sizeof(*octx));
octx->n_threads = ctx->n_threads;
octx->ctx = ctx;
work_queue_wakeup(ctx->work_queue);
if (ctx->hmx_queue) {
hmx_queue_wakeup(ctx->hmx_queue);
@@ -1056,13 +1064,23 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
}
work_queue_suspend(ctx->work_queue);
// Flush remaining dirty tensors at the end of the batch
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
profile_stop(HTP_PROF_BASIC, &batch_prof);
struct htp_opbatch_rsp rsp;
memset(&rsp, 0, sizeof(rsp));
rsp.id = req->id;
rsp.status = op_status;
rsp.n_bufs = n_bufs;
rsp.n_tensors = n_tens;
rsp.n_ops = n_ops;
rsp.id = req->id;
rsp.status = op_status;
rsp.n_bufs = n_bufs;
rsp.n_tensors = n_tens;
rsp.n_ops = n_ops;
rsp.usecs = batch_prof.usecs;
rsp.cycles_start = batch_prof.cycles_start;
rsp.cycles_stop = batch_prof.cycles_stop;
if (ctx->profiler == HTP_PROF_TRACE) {
for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
@@ -1073,11 +1091,6 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
struct dspqueue_buffer write_dbuf = *dbuf;
write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
// Flush remaining dirty tensors at the end of the batch
htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
if (err != 0) {
FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
ggml/src/ggml-hexagon/htp/matmul-ops.c
+236 -104
View File
@@ -14,6 +14,8 @@
#include "hex-dma.h"
#include "hvx-utils.h"
#include "hvx-dump.h"
#include "hvx-arith.h"
#include "hvx-reduce.h"
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
@@ -82,6 +84,8 @@ struct htp_mm_context {
// Precomputed values
uint32_t src0_nrows_per_thread;
uint32_t src0_row_size_padded;
uint32_t src1_nrows;
struct fastdiv_values mm_div_ne12_ne1;
struct fastdiv_values mm_div_ne1;
@@ -103,6 +107,7 @@ struct htp_mm_context {
// Fields for scattered mapping & HMX support in MUL_MAT_ID
const uint32_t * matrix_row_counts;
const struct mmid_row_mapping * matrix_rows;
uint32_t mapping_stride;
// Dynamic VTCM pointers allocated sequentially
uint8_t * vtcm_src0;
@@ -154,8 +159,6 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
};
#define htp_matmul_tensors_preamble \
const struct htp_tensor * restrict src0 = octx->src[0]; \
const struct htp_tensor * restrict src1 = octx->src[1]; \
@@ -444,6 +447,16 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
\
uint32_t push_ct = ct_start; \
if (src0_start_row < src0_end_row) { \
if (src2) { \
float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; \
const float * src2_ptr = (const float *) src2->data + src0_start_row; \
int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
if (slice_size > 0) { \
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), \
slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \
dma_queue_pop_nowait(dma_queue); \
} \
} \
for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \
src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \
@@ -465,7 +478,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
\
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\
if (push_ct < ct_end) { \
dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \
@@ -476,24 +489,16 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
\
int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
if (copy_cnt > 0) { \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \
if (src2) { \
float * dst_ptr = &dst_col[src0_start_row]; \
const float * src2_ptr = (const float *) src2->data + src0_start_row; \
float * tmp_ptr = tmp; \
int remaining = copy_cnt; \
while (remaining > 0) { \
int n = MIN(remaining, 32); \
HVX_Vector v_out = hvx_vmemu(tmp_ptr); \
HVX_Vector v_z = hvx_vmemu(src2_ptr); \
hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z)); \
dst_ptr += n; \
src2_ptr += n; \
tmp_ptr += n; \
remaining -= n; \
} \
hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], \
(const uint8_t *) tmp, \
(const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), \
copy_cnt); \
} else { \
hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \
} \
}
@@ -1069,6 +1074,16 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
// Prefill vtcm with 2x src0 rows
if (src0_start_row < src0_end_row) {
if (src2) {
float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row;
const float * src2_ptr = (const float *) src2->data + src0_start_row;
int slice_size = (int)src0_end_row - (int)src0_start_row;
if (slice_size > 0) {
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr),
slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1);
dma_queue_pop_nowait(dma_queue);
}
}
for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
const uint32_t is0 = (ir0 - src0_start_row);
if (is0 >= n_prefetch) {
@@ -1114,27 +1129,21 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
}
int copy_cnt = src0_end_row - src0_start_row;
if (src2) {
float * dst_ptr = &dst_col[src0_start_row];
const float * src2_ptr = (const float *) src2->data + src0_start_row;
float * tmp_ptr = tmp;
int remaining = copy_cnt;
while (remaining > 0) {
int n = MIN(remaining, 32);
HVX_Vector v_out = hvx_vmemu(tmp_ptr);
HVX_Vector v_z = hvx_vmemu(src2_ptr);
hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
dst_ptr += n;
src2_ptr += n;
tmp_ptr += n;
remaining -= n;
if (copy_cnt > 0) {
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
if (src2) {
hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],
(const uint8_t *) tmp,
(const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row),
copy_cnt);
} else {
hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
}
} else {
hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
}
}
#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ids->ne[0] * ids->ne[1] + (i1)]
#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)]
static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {
htp_matmul_preamble;
@@ -1519,7 +1528,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false);
dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false, false);
if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM ||
kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
@@ -1551,6 +1560,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL;
@@ -2346,12 +2356,77 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles(
}
}
typedef struct {
float *dst;
const float *src2;
const __fp16 *vtcm_src;
uint32_t n_rows;
uint32_t n_cols;
uint32_t dst_stride;
uint32_t src2_stride;
uint32_t dst_cols;
struct fastdiv_values n_threads_div;
struct htp_thread_trace *traces;
struct htp_context *ctx;
} output_transfer_col_chunk_state_t;
static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
(void) n;
output_transfer_col_chunk_state_t *st = (output_transfer_col_chunk_state_t *) data;
struct htp_thread_trace * tr = &st->traces[i];
uint32_t n_blocks = st->n_cols / 32;
uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div);
uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div);
uint32_t c_first = b_first * 32;
uint32_t c_last = b_last * 32;
uint32_t c_len = c_last - c_first;
if (c_len == 0) return;
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
float *dst = st->dst + c_first;
const float *src2 = st->src2 ? (st->src2 + c_first) : NULL;
const __fp16 *vtcm_src = st->vtcm_src + b_first * HTP_MM_HMX_TILE_N_ELMS;
int chunk_dst_cols = (int)st->dst_cols - (int)c_first;
if (chunk_dst_cols > 0) {
transfer_output_chunk_fp16_to_fp32_col_chunk(
dst, src2, vtcm_src, 0, st->n_rows, c_len, st->n_cols,
st->dst_stride, st->src2_stride, (uint32_t)chunk_dst_cols
);
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
}
static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const float *src2, const __fp16 *vtcm_src,
int n_rows, int n_cols, int dst_stride, uint32_t src2_stride, int dst_cols, int n_threads) {
assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
if (n_rows <= 0) return;
uint32_t n_blocks = (uint32_t)n_cols / 32;
if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) {
struct fastdiv_values n_threads_div = init_fastdiv_values(n_threads);
output_transfer_col_chunk_state_t col_state;
col_state.dst = dst;
col_state.src2 = src2;
col_state.vtcm_src = vtcm_src;
col_state.n_rows = (uint32_t)n_rows;
col_state.n_cols = (uint32_t)n_cols;
col_state.dst_stride = (uint32_t)dst_stride;
col_state.src2_stride = src2_stride;
col_state.dst_cols = (uint32_t)dst_cols;
col_state.n_threads_div = n_threads_div;
col_state.traces = ctx->trace;
col_state.ctx = ctx;
worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_col_chunk_worker_fn, &col_state, n_threads);
return;
}
size_t n_tot_chunks = n_rows;
size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : hmx_ceil_div(n_rows, n_threads);
n_chunks_per_task = hex_align_up(n_chunks_per_task, 2);
@@ -3338,12 +3413,10 @@ int op_matmul(struct htp_ops_context * octx) {
static int hmx_mm_op_matmul_id(
struct htp_ops_context * octx,
struct htp_mm_context * mmctx,
const uint32_t * matrix_row_counts,
const struct mmid_row_mapping * matrix_rows,
void * mapping_buf,
bool must_free_mapping
struct htp_mm_context * mmctx
) {
const uint32_t * matrix_row_counts = mmctx->matrix_row_counts;
const struct mmid_row_mapping * matrix_rows = mmctx->matrix_rows;
htp_matmul_tensors_preamble;
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
const int n_ids = octx->src[2]->ne[0];
@@ -3361,28 +3434,24 @@ static int hmx_mm_op_matmul_id(
nb11, nb12,
nb1, nb2,
(int) src0->nb[1], (int) src0->type,
matrix_rows, cur_a, n_ids * octx->src[2]->ne[1]);
matrix_rows, cur_a, mmctx->mapping_stride);
if (ret != 0) {
FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret);
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_NO_SUPPORT;
}
}
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_OK;
}
static int hvx_mm_matmul_id(
struct htp_ops_context * octx,
struct htp_mm_context * mmctx,
size_t src0_row_size_padded,
uint32_t src1_nrows,
worker_callback_t matmul_id_job_func,
void * mapping_buf,
bool must_free_mapping
work_queue_func_t hvx_mmid_task_func
) {
htp_matmul_tensors_preamble;
const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded;
const uint32_t src1_nrows = mmctx->src1_nrows;
struct htp_thread_trace * tr = &octx->ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
@@ -3395,7 +3464,7 @@ static int hvx_mm_matmul_id(
const uint32_t nb = (ne10 + qk - 1) / qk;
const uint32_t total_nb = src1_nrows * nb;
worker_callback_t quant_task_func;
work_queue_func_t quant_task_func;
uint32_t n_quant_tasks = 1;
if (src1_nrows < octx->n_threads) {
n_quant_tasks = MIN(total_nb, octx->n_threads);
@@ -3416,7 +3485,7 @@ static int hvx_mm_matmul_id(
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false);
0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false, false);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
@@ -3431,7 +3500,6 @@ static int hvx_mm_matmul_id(
// Make sure the reserved vtcm size is sufficient
if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-id-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size);
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_VTCM_TOO_SMALL;
}
@@ -3461,12 +3529,78 @@ static int hvx_mm_matmul_id(
htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
worker_pool_run_func(octx->ctx->worker_pool, matmul_id_job_func, mmctx, octx->n_threads);
worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads);
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_OK;
}
static inline void scan_expert_ids_n(
const struct htp_tensor * ids,
const uint32_t n_ids,
uint32_t n_as,
uint32_t * counts,
struct mmid_row_mapping * matrix_rows,
uint32_t mapping_stride
) {
const size_t ids_nb1 = ids->nb[1];
const uint8_t * ids_data = (const uint8_t *) ids->data;
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = row_ptr[id];
if (i02 < 0) {
continue;
}
assert(i02 < n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
}
counts[i02] += 1;
}
}
}
static inline void scan_expert_ids(
const struct htp_tensor * ids,
uint32_t n_ids,
uint32_t n_as,
uint32_t * counts,
struct mmid_row_mapping * matrix_rows,
uint32_t mapping_stride
) {
const size_t ids_nb0 = ids->nb[0];
if (ids_nb0 == 4) {
switch (n_ids) {
case 8: scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride); break;
case 4: scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride); break;
case 2: scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride); break;
default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
}
} else {
// Strided fallback
const size_t ids_nb1 = ids->nb[1];
const uint8_t * ids_data = (const uint8_t *) ids->data;
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
if (i02 < 0) {
continue;
}
assert(i02 < n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
}
counts[i02] += 1;
}
}
}
}
int op_matmul_id(struct htp_ops_context * octx) {
htp_matmul_tensors_preamble;
@@ -3489,74 +3623,72 @@ int op_matmul_id(struct htp_ops_context * octx) {
const uint32_t src0_nrows = ne01; // per expert
const uint32_t src1_nrows = ne11 * ne12 * ne13;
worker_callback_t quant_task_func;
worker_callback_t matmul_id_job_func = src1_nrows > 1 ? hvx_mm_id : hvx_mv_id;
// Compute src0_nrows_per_thread
mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
// row groups
const int n_ids = ids->ne[0]; // n_expert_used
const int n_as = ne02; // n_expert
size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
size_t matrix_row_map_size = n_as * ids->ne[0] * ids->ne[1] * sizeof(struct mmid_row_mapping);
const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
void * mapping_buf = NULL;
bool must_free_mapping = false;
if (octx->ctx->ddr_spad_base && total_map_size <= octx->ctx->ddr_spad_size) {
mapping_buf = octx->ctx->ddr_spad_base;
} else {
mapping_buf = memalign(128, total_map_size);
if (mapping_buf) {
must_free_mapping = true;
} else {
return HTP_STATUS_INTERNAL_ERR;
}
}
uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) ((uint8_t *) mapping_buf + matrix_row_counts_size);
mmctx->matrix_row_counts = matrix_row_counts;
mmctx->matrix_rows = matrix_rows;
mmctx->mm_div_ne11 = kparams->div_ne11;
if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) {
if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_NO_SUPPORT;
}
uint8_t * mapping_buf = octx->ctx->ddr_spad_base;
uint32_t mapping_stride = 1;
uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
struct mmid_row_mapping * matrix_rows = NULL;
if (src1_nrows > 1) {
// initialize matrix_row_counts and map
memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
// group rows by src0 matrix
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { // token idx
for (uint32_t id = 0; id < n_ids; ++id) { // expert idx
const int32_t i02 = *(const int32_t *) ((const uint8_t *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
if (i02 < 0) {
continue;
}
assert(i02 < n_as);
memset(matrix_row_counts, 0, matrix_row_counts_size);
scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
matrix_rows[i02 * n_ids * ids->ne[1] + matrix_row_counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
matrix_row_counts[i02] += 1;
uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
mapping_stride = max_count > 0 ? max_count : 1;
size_t matrix_row_map_size = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
if (total_map_size > octx->ctx->ddr_spad_size) {
mapping_buf = memalign(128, total_map_size);
if (!mapping_buf) {
return HTP_STATUS_INTERNAL_ERR;
}
}
matrix_row_counts = (uint32_t *) mapping_buf;
matrix_rows = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
}
mmctx->matrix_row_counts = matrix_row_counts;
mmctx->matrix_rows = matrix_rows;
mmctx->mapping_stride = mapping_stride;
mmctx->mm_div_ne11 = kparams->div_ne11;
mmctx->src0_row_size_padded = src0_row_size_padded;
mmctx->src1_nrows = src1_nrows;
htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
int s;
if (kparams->n_hmx) {
return hmx_mm_op_matmul_id(octx, mmctx, matrix_row_counts, matrix_rows, mapping_buf, must_free_mapping);
s = hmx_mm_op_matmul_id(octx, mmctx);
} else {
if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) {
s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id);
} else {
s = HTP_STATUS_NO_SUPPORT;
}
}
return hvx_mm_matmul_id(octx, mmctx, src0_row_size_padded, src1_nrows, matmul_id_job_func, mapping_buf, must_free_mapping);
if (mapping_buf != octx->ctx->ddr_spad_base) {
free(mapping_buf);
}
return s;
}
int op_matmul_qkv(struct htp_ops_context * octx) {
@@ -3633,7 +3765,7 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false);
0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true, false);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
@@ -3778,7 +3910,7 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true);
0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, false, true);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
ggml/src/ggml-hexagon/htp/matmul-ops.h
+2 -1
View File
@@ -460,6 +460,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
size_t dst_row_size,
size_t src0_row_size,
size_t src1_row_size,
size_t src2_row_size,
uint32_t n_prefetch,
bool is_matmul_id,
bool is_fused_qkv,
@@ -467,7 +468,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
) {
size_t src0_sz = 0;
size_t src1_sz = 0;
size_t src2_sz = 0;
size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
size_t src3_sz = 0;
size_t dst_sz = 0;
ggml/src/ggml-hip/CMakeLists.txt
-4
View File
@@ -114,10 +114,6 @@ if (GGML_HIP_NO_VMM)
add_compile_definitions(GGML_HIP_NO_VMM)
endif()
if (GGML_HIP_ROCWMMA_FATTN)
add_compile_definitions(GGML_HIP_ROCWMMA_FATTN)
endif()
if (NOT GGML_HIP_MMQ_MFMA)
add_compile_definitions(GGML_HIP_NO_MMQ_MFMA)
endif()
ggml/src/ggml-metal/ggml-metal-device.m
+2 -1
View File
@@ -1218,8 +1218,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
(ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0);
case GGML_OP_PAD_REFLECT_1D:
case GGML_OP_TIMESTEP_EMBEDDING:
case GGML_OP_LEAKY_RELU:
return op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_LEAKY_RELU:
return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
case GGML_OP_ARGSORT:
case GGML_OP_TOP_K:
case GGML_OP_ARANGE:
ggml/src/ggml-opencl/CMakeLists.txt
+2
View File
@@ -5,6 +5,8 @@ set(TARGET_NAME ggml-opencl)
ggml_add_backend_library(${TARGET_NAME}
ggml-opencl.cpp
cl-program-cache.cpp
cl-program-cache.h
../../include/ggml-opencl.h)
target_link_libraries(${TARGET_NAME} PRIVATE ${OpenCL_LIBRARIES})
target_include_directories(${TARGET_NAME} PRIVATE ${OpenCL_INCLUDE_DIRS})
ggml/src/ggml-opencl/cl-program-cache.cpp
+453
View File
@@ -0,0 +1,453 @@
// Match the version setup ggml-opencl.cpp uses, so any cl.h declarations we
// touch are consistent across this backend's translation units.
#define CL_TARGET_OPENCL_VERSION GGML_OPENCL_TARGET_VERSION
#define CL_USE_DEPRECATED_OPENCL_1_2_APIS
#include "cl-program-cache.h"
#include "ggml-impl.h" // GGML_LOG_INFO / WARN
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <system_error>
#include <vector>
#if defined(_WIN32)
# ifndef WIN32_LEAN_AND_MEAN
# define WIN32_LEAN_AND_MEAN
# endif
# ifndef NOMINMAX
# define NOMINMAX
# endif
# include <windows.h>
# include <process.h>
# define ggml_getpid() ((int) GetCurrentProcessId())
#else
# include <unistd.h>
# define ggml_getpid() ((int) getpid())
#endif
namespace fs = std::filesystem;
// ----------------------------------------------------------------------------
// SHA-256 (FIPS 180-4). Self-contained, ~80 lines, public-domain reference.
// Hot path is a few KB of source per kernel ⇒ <1 ms total per process init.
// ----------------------------------------------------------------------------
namespace {
struct sha256_ctx {
uint32_t state[8];
uint64_t bitlen;
uint8_t buf[64];
size_t buf_len;
};
const uint32_t K256[64] = {
0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2,
};
inline uint32_t rotr32(uint32_t x, unsigned n) { return (x >> n) | (x << (32 - n)); }
void sha256_compress(uint32_t state[8], const uint8_t block[64]) {
uint32_t w[64];
for (int i = 0; i < 16; ++i) {
w[i] = ((uint32_t)block[i*4 ] << 24) |
((uint32_t)block[i*4 + 1] << 16) |
((uint32_t)block[i*4 + 2] << 8) |
((uint32_t)block[i*4 + 3] );
}
for (int i = 16; i < 64; ++i) {
uint32_t s0 = rotr32(w[i-15], 7) ^ rotr32(w[i-15], 18) ^ (w[i-15] >> 3);
uint32_t s1 = rotr32(w[i-2], 17) ^ rotr32(w[i-2], 19) ^ (w[i-2] >> 10);
w[i] = w[i-16] + s0 + w[i-7] + s1;
}
uint32_t a = state[0],b = state[1],c = state[2],d = state[3],e = state[4],f = state[5],g = state[6],h = state[7];
for (int i = 0; i < 64; ++i) {
uint32_t S1 = rotr32(e, 6) ^ rotr32(e, 11) ^ rotr32(e, 25);
uint32_t ch = (e & f) ^ ((~e) & g);
uint32_t t1 = h + S1 + ch + K256[i] + w[i];
uint32_t S0 = rotr32(a, 2) ^ rotr32(a, 13) ^ rotr32(a, 22);
uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
uint32_t t2 = S0 + maj;
h = g; g = f; f = e; e = d + t1;
d = c; c = b; b = a; a = t1 + t2;
}
state[0]+=a; state[1]+=b; state[2]+=c; state[3]+=d;
state[4]+=e; state[5]+=f; state[6]+=g; state[7]+=h;
}
void sha256_init(sha256_ctx & c) {
c.state[0]=0x6a09e667; c.state[1]=0xbb67ae85; c.state[2]=0x3c6ef372; c.state[3]=0xa54ff53a;
c.state[4]=0x510e527f; c.state[5]=0x9b05688c; c.state[6]=0x1f83d9ab; c.state[7]=0x5be0cd19;
c.bitlen = 0;
c.buf_len = 0;
}
void sha256_update(sha256_ctx & c, const void * data, size_t len) {
const uint8_t * p = (const uint8_t *) data;
c.bitlen += (uint64_t) len * 8;
if (c.buf_len > 0) {
size_t n = 64 - c.buf_len;
if (n > len) { n = len; }
memcpy(c.buf + c.buf_len, p, n);
c.buf_len += n;
p += n;
len -= n;
if (c.buf_len == 64) {
sha256_compress(c.state, c.buf);
c.buf_len = 0;
}
}
while (len >= 64) {
sha256_compress(c.state, p);
p += 64;
len -= 64;
}
if (len > 0) {
memcpy(c.buf, p, len);
c.buf_len = len;
}
}
void sha256_final(sha256_ctx & c, uint8_t out[32]) {
uint64_t bitlen = c.bitlen;
c.buf[c.buf_len++] = 0x80;
if (c.buf_len > 56) {
while (c.buf_len < 64) { c.buf[c.buf_len++] = 0; }
sha256_compress(c.state, c.buf);
c.buf_len = 0;
}
while (c.buf_len < 56) { c.buf[c.buf_len++] = 0; }
for (int i = 7; i >= 0; --i) { c.buf[c.buf_len++] = (uint8_t) (bitlen >> (i * 8)); }
sha256_compress(c.state, c.buf);
for (int i = 0; i < 8; ++i) {
out[i*4 ] = (uint8_t) (c.state[i] >> 24);
out[i*4 + 1] = (uint8_t) (c.state[i] >> 16);
out[i*4 + 2] = (uint8_t) (c.state[i] >> 8);
out[i*4 + 3] = (uint8_t) (c.state[i] );
}
}
std::string sha256_hex(const uint8_t digest[32]) {
static const char hex[] = "0123456789abcdef";
std::string s(64, '0');
for (int i = 0; i < 32; ++i) {
s[i*2 ] = hex[digest[i] >> 4];
s[i*2 + 1] = hex[digest[i] & 0xf];
}
return s;
}
std::string compute_key(const std::string & key_suffix,
const char * source,
const std::string & compile_opts) {
sha256_ctx c;
sha256_init(c);
static const uint8_t sep = 0;
sha256_update(c, source, strlen(source));
sha256_update(c, &sep, 1);
sha256_update(c, compile_opts.data(), compile_opts.size());
sha256_update(c, &sep, 1);
sha256_update(c, key_suffix.data(), key_suffix.size());
uint8_t digest[32];
sha256_final(c, digest);
return sha256_hex(digest);
}
bool make_dir_recursive(const std::string & path) {
if (path.empty()) { return false; }
// create_directories() already creates missing parents. It returns false
// (with ec clear) when the directory is already there, so re-check.
const fs::path p = fs::u8path(path);
std::error_code ec;
if (fs::create_directories(p, ec)) { return true; }
std::error_code ec_stat;
return fs::is_directory(p, ec_stat);
}
std::string default_cache_dir() {
#if defined(_WIN32)
const char * base = std::getenv("LOCALAPPDATA");
if (!base || !*base) { base = std::getenv("APPDATA"); }
if (!base || !*base) { base = std::getenv("TEMP"); }
if (!base || !*base) { base = "."; }
return std::string(base) + "\\llama.cpp\\cl-cache";
#elif defined(__APPLE__)
const char * home = std::getenv("HOME");
if (!home || !*home) { home = "."; }
return std::string(home) + "/Library/Caches/llama.cpp/cl-cache";
#else
// The throwing overload aborts the process when no usable temp directory
// exists (e.g. Android app contexts with TMPDIR unset); an empty return
// here just disables the cache instead.
std::error_code ec;
const fs::path tmp_path = fs::temp_directory_path(ec);
if (ec || tmp_path.empty()) { return {}; }
return tmp_path.string() + "/llama.cpp/cl-cache";
#endif
}
// Query a NUL-terminated string from clGetDeviceInfo / clGetPlatformInfo.
template <typename GetInfoFn, typename Object>
std::string query_string(GetInfoFn fn, Object obj, cl_uint name) {
size_t sz = 0;
if (fn(obj, name, 0, nullptr, &sz) != CL_SUCCESS || sz == 0) {
return {};
}
std::string s(sz, '\0');
if (fn(obj, name, sz, &s[0], nullptr) != CL_SUCCESS) {
return {};
}
if (!s.empty() && s.back() == '\0') {
s.pop_back();
}
return s;
}
std::string compute_key_suffix(cl_device_id device) {
cl_platform_id platform = nullptr;
clGetDeviceInfo(device, CL_DEVICE_PLATFORM, sizeof(platform), &platform, nullptr);
std::string s;
s.reserve(512);
s += query_string(clGetDeviceInfo, device, CL_DEVICE_NAME); s.push_back('\0');
s += query_string(clGetDeviceInfo, device, CL_DRIVER_VERSION); s.push_back('\0');
s += query_string(clGetDeviceInfo, device, CL_DEVICE_VERSION); s.push_back('\0');
if (platform) {
s += query_string(clGetPlatformInfo, platform, CL_PLATFORM_VERSION); s.push_back('\0');
}
s += "fmt=" + std::to_string(CL_PROGRAM_CACHE_FORMAT_VERSION);
return s;
}
const uint8_t MAGIC[8] = { 'G','G','M','L','C','L','B','C' };
bool read_all(const std::string & path, std::vector<uint8_t> & out) {
std::ifstream f(fs::u8path(path), std::ios::binary);
if (!f) { return false; }
f.seekg(0, std::ios::end);
std::streamsize sz = f.tellg();
if (sz < 0) { return false; }
f.seekg(0, std::ios::beg);
out.resize((size_t) sz);
if (sz > 0) { f.read((char *) out.data(), sz); }
return f.good() || f.eof();
}
bool write_atomic(const std::string & path, const uint8_t * data, size_t len) {
const fs::path dst = fs::u8path(path);
const fs::path tmp = fs::u8path(path + ".tmp." + std::to_string(ggml_getpid()));
{
std::ofstream f(tmp, std::ios::binary | std::ios::trunc);
if (!f) { return false; }
f.write((const char *) data, (std::streamsize) len);
if (!f.good()) {
std::error_code ec_rm;
fs::remove(tmp, ec_rm);
return false;
}
}
std::error_code ec;
fs::rename(tmp, dst, ec);
if (ec) {
std::error_code ec_rm;
fs::remove(tmp, ec_rm);
return false;
}
return true;
}
} // namespace
static bool cache_debug_enabled() {
static int cached = -1;
if (cached < 0) {
const char * e = std::getenv("GGML_OPENCL_KERNEL_CACHE_DEBUG");
cached = (e && *e) ? 1 : 0;
}
return cached != 0;
}
static std::string opts_preview(const std::string & opts, size_t n = 120) {
if (opts.size() <= n) { return opts; }
return opts.substr(0, n) + "...";
}
// Running cache tally (diagnostic; plain ints — a benign race in the rare
// multi-threaded lazy-compile case at worst miscounts by one).
static int g_cache_hits = 0, g_cache_misses = 0, g_cache_saves = 0;
// Debug trace directly to stderr
static void cache_debug_line(const char * kind, const std::string & key,
const char * source, const std::string & opts) {
if (!cache_debug_enabled()) { return; }
fprintf(stderr, "ggml_opencl: cache %-4s [h=%d m=%d s=%d] key=%s src=%zuB opts='%s'\n",
kind, g_cache_hits, g_cache_misses, g_cache_saves,
key.substr(0, 16).c_str(), strlen(source), opts_preview(opts).c_str());
fflush(stderr);
}
cl_program_cache_state cl_program_cache_init(cl_device_id device) {
cl_program_cache_state st;
const char * env = std::getenv("GGML_OPENCL_KERNEL_CACHE_DIR");
if (env && (!std::strcmp(env, "0") || !std::strcmp(env, "off") ||
!std::strcmp(env, "none") || !std::strcmp(env, "disable") ||
!std::strcmp(env, "disabled"))) {
if (cache_debug_enabled()) {
fprintf(stderr, "ggml_opencl: kernel cache disabled by GGML_OPENCL_KERNEL_CACHE_DIR=%s\n", env);
fflush(stderr);
}
return st;
}
std::string dir;
if (!env || !*env || !std::strcmp(env, "1") || !std::strcmp(env, "default")) {
dir = default_cache_dir();
if (dir.empty()) {
GGML_LOG_INFO("ggml_opencl: kernel cache disabled (no usable default cache directory)\n");
return st;
}
} else {
dir = env;
}
if (!make_dir_recursive(dir)) {
GGML_LOG_INFO("ggml_opencl: kernel cache disabled (cannot create directory '%s')\n", dir.c_str());
return st;
}
st.dir = dir;
st.key_suffix = compute_key_suffix(device);
GGML_LOG_INFO("ggml_opencl: kernel cache enabled at '%s'\n", st.dir.c_str());
if (cache_debug_enabled()) {
fprintf(stderr, "ggml_opencl: kernel cache enabled at '%s' "
"(GGML_OPENCL_KERNEL_CACHE_DIR=off to disable)\n", st.dir.c_str());
fflush(stderr);
}
return st;
}
cl_program cl_program_cache_try_load(
const cl_program_cache_state & state,
cl_context context,
cl_device_id device,
const char * source,
const std::string & compile_opts) {
if (state.dir.empty() || !source) { return nullptr; }
const std::string key = compute_key(state.key_suffix, source, compile_opts);
const std::string path = state.dir + "/" + key + ".clbin";
std::vector<uint8_t> file;
if (!read_all(path, file)) {
++g_cache_misses;
cache_debug_line("MISS", key, source, compile_opts);
return nullptr;
}
if (file.size() < 16 || std::memcmp(file.data(), MAGIC, 8) != 0) { return nullptr; }
uint32_t fmt =
((uint32_t) file[ 8]) | ((uint32_t) file[ 9] << 8) |
((uint32_t) file[10] << 16) | ((uint32_t) file[11] << 24);
if (fmt != CL_PROGRAM_CACHE_FORMAT_VERSION) { return nullptr; }
const size_t hdr_len = 16;
const unsigned char * bin = file.data() + hdr_len;
const size_t bin_len = file.size() - hdr_len;
cl_int err = CL_SUCCESS;
cl_int bin_err = CL_SUCCESS;
cl_program p = clCreateProgramWithBinary(context, 1, &device, &bin_len, &bin, &bin_err, &err);
if (err != CL_SUCCESS || bin_err != CL_SUCCESS || p == nullptr) {
if (p) { clReleaseProgram(p); }
return nullptr;
}
err = clBuildProgram(p, 0, nullptr, compile_opts.c_str(), nullptr, nullptr);
if (err != CL_SUCCESS) {
clReleaseProgram(p);
return nullptr;
}
++g_cache_hits;
cache_debug_line("HIT", key, source, compile_opts);
return p;
}
void cl_program_cache_try_save(
const cl_program_cache_state & state,
cl_program program,
cl_device_id /*device*/,
const char * source,
const std::string & compile_opts) {
if (state.dir.empty() || !program || !source) {
return;
}
cl_uint n_dev = 0;
if (clGetProgramInfo(program, CL_PROGRAM_NUM_DEVICES, sizeof(n_dev), &n_dev, nullptr) != CL_SUCCESS || n_dev == 0) {
return;
}
std::vector<size_t> sizes(n_dev);
if (clGetProgramInfo(program, CL_PROGRAM_BINARY_SIZES, sizeof(size_t) * n_dev, sizes.data(), nullptr) != CL_SUCCESS) {
return;
}
if (sizes.empty() || sizes[0] == 0) {
return;
}
std::vector<std::vector<uint8_t>> binaries(n_dev);
std::vector<unsigned char *> bin_ptrs(n_dev);
for (cl_uint i = 0; i < n_dev; ++i) {
binaries[i].resize(sizes[i]);
bin_ptrs[i] = binaries[i].data();
}
if (clGetProgramInfo(program, CL_PROGRAM_BINARIES, sizeof(unsigned char *) * n_dev, bin_ptrs.data(), nullptr) != CL_SUCCESS) {
return;
}
// We only care about the first device's binary — that's the one we'd
// re-load with on a future cache hit. Multi-device contexts aren't a
// pattern this backend uses today.
const std::vector<uint8_t> & bin = binaries[0];
std::vector<uint8_t> file;
file.reserve(16 + bin.size());
file.insert(file.end(), MAGIC, MAGIC + 8);
uint32_t fmt = CL_PROGRAM_CACHE_FORMAT_VERSION;
file.push_back((uint8_t) (fmt & 0xff));
file.push_back((uint8_t) ((fmt >> 8) & 0xff));
file.push_back((uint8_t) ((fmt >> 16) & 0xff));
file.push_back((uint8_t) ((fmt >> 24) & 0xff));
file.push_back(0); file.push_back(0); file.push_back(0); file.push_back(0); // reserved
file.insert(file.end(), bin.begin(), bin.end());
const std::string key = compute_key(state.key_suffix, source, compile_opts);
const std::string path = state.dir + "/" + key + ".clbin";
if (!write_atomic(path, file.data(), file.size())) {
GGML_LOG_INFO("ggml_opencl: kernel cache: failed to write '%s'\n", path.c_str());
} else {
++g_cache_saves;
cache_debug_line("SAVE", key, source, compile_opts);
}
}
ggml/src/ggml-opencl/cl-program-cache.h
+75
View File
@@ -0,0 +1,75 @@
// On-disk cache for OpenCL cl_program binaries. Lets a fresh process skip the
// expensive clBuildProgram-from-source step when a binary for the exact same
// (source, compile options, device, driver, platform) was previously saved.
//
// Activation: default on via GGML_OPENCL_KERNEL_CACHE_DIR:
// unset / empty / "1" / "default" : platform default cache dir
// (%LOCALAPPDATA%\llama.cpp\cl-cache,
// ~/Library/Caches/llama.cpp/cl-cache,
// <temp dir>/llama.cpp/cl-cache elsewhere)
// "0" / "off" / "none" / "disable(d)" : disabled (all functions no-op)
// any other value : used verbatim as the cache path
// If the chosen directory cannot be created/used, the cache silently disables
// itself for the process and falls back to source compile.
// GGML_OPENCL_KERNEL_CACHE_DEBUG=1 prints a HIT/MISS/SAVE trace (with a running
// tally) straight to stderr — visible even in tools that filter INFO/WARN logs;
// redirect stderr to record it.
//
// Cache key (SHA-256 hex):
// sha256(source_bytes || '\x00' ||
// compile_opts || '\x00' ||
// CL_DEVICE_NAME || '\x00' ||
// CL_DRIVER_VERSION || '\x00' ||
// CL_PLATFORM_VERSION || '\x00' ||
// CL_PROGRAM_CACHE_FORMAT_VERSION)
//
// The key fully captures everything that can affect the produced binary,
// without needing the host source revision (a kernel source change shows up
// in source_bytes; a compile-option change shows up in compile_opts).
//
// File layout per cache entry: <cache_dir>/<sha256-hex>.clbin
// bytes [0..7] : magic "GGMLCLBC"
// bytes [8..11] : uint32_t format version (CL_PROGRAM_CACHE_FORMAT_VERSION)
// bytes [12..15] : uint32_t reserved (0)
// bytes [16..] : raw cl_program binary as returned by
// clGetProgramInfo(CL_PROGRAM_BINARIES)
//
// Concurrency: writes go to <name>.tmp.<pid> then atomic rename. On race,
// last-writer-wins. No locks.
#pragma once
#include <CL/cl.h>
#include <string>
// Bumped manually if host-side OpenCL API usage changes in a way that
// affects compile semantics but does not show up in source_bytes /
// compile_opts (e.g. switching from clCreateProgramWithSource to
// clCompileProgram + clLinkProgram, or changing how multiple sources
// are concatenated). Most commits — including kernel changes — do NOT
// require bumping this; the source bytes already capture those.
#define CL_PROGRAM_CACHE_FORMAT_VERSION 1u
struct cl_program_cache_state {
// Empty string means cache is disabled.
std::string dir;
// Concatenated device/driver/platform identity + cache format version,
// computed once at init and folded into every key.
std::string key_suffix;
};
cl_program_cache_state cl_program_cache_init(cl_device_id device);
cl_program cl_program_cache_try_load(
const cl_program_cache_state & state,
cl_context context,
cl_device_id device,
const char * source,
const std::string & compile_opts);
void cl_program_cache_try_save(
const cl_program_cache_state & state,
cl_program program,
cl_device_id device,
const char * source,
const std::string & compile_opts);
ggml/src/ggml-opencl/ggml-opencl.cpp
+205 -187
View File
File diff suppressed because it is too large Load Diff
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+260 -20
View File
@@ -3283,6 +3283,7 @@ static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, const std
import_info.setPNext(&mem_flags_info);
buf->device_memory = device->device.allocateMemory({ size, memory_type_idx, &import_info });
} catch (const vk::SystemError& e) {
GGML_LOG_WARN("ggml_vulkan: host pointer memory import failed (%s)\n", e.what());
}
} else {
for (auto it = req_flags_list.begin(); it != req_flags_list.end(); it++) {
@@ -8128,25 +8129,32 @@ static bool ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, siz
ggml_vk_host_get(dst->device, src, buf, buf_offset);
if (buf != nullptr) {
// Memory is pinned, use as staging buffer
std::vector<vk::BufferCopy> slices(1);
if (width == spitch && width == dpitch) {
// Only do single write if stride is equal
slices[0].srcOffset = buf_offset;
slices[0].dstOffset = offset;
slices[0].size = width * height;
} else {
slices.resize(height);
for (size_t i = 0; i < height; i++) {
slices[i].srcOffset = buf_offset + i * spitch;
slices[i].dstOffset = offset + i * dpitch;
slices[i].size = width;
// extent of the read in pinned source memory; guard against tensors that
// straddle a pinned-chunk boundary (they fall back to staging below)
size_t src_extent = (width == spitch) ? (size_t) width * height
: (height > 0 ? (height - 1) * spitch + width : 0);
if (buf_offset + src_extent <= buf->size) {
// Memory is pinned, use as staging buffer
std::vector<vk::BufferCopy> slices(1);
if (width == spitch && width == dpitch) {
// Only do single write if stride is equal
slices[0].srcOffset = buf_offset;
slices[0].dstOffset = offset;
slices[0].size = width * height;
} else {
slices.resize(height);
for (size_t i = 0; i < height; i++) {
slices[i].srcOffset = buf_offset + i * spitch;
slices[i].dstOffset = offset + i * dpitch;
slices[i].size = width;
}
}
}
ggml_vk_sync_buffers(nullptr, subctx);
subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
return true;
ggml_vk_sync_buffers(nullptr, subctx);
subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
return true;
}
// straddles a chunk boundary: fall through to staging
}
VK_LOG_DEBUG("STAGING");
@@ -10584,8 +10592,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
}
// Only use mask opt when the mask is fairly large. This hasn't been tuned extensively.
// GCN with large head size (>= 256) benefits in high-context prefill (skipping the
// per-block mask add on fully-visible blocks dominates), so enable it there too.
bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
&& (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256);
&& (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK >= 256 || HSV >= 256);
vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type);
@@ -12512,9 +12522,108 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc
ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_OPT_STEP_SGD, { (uint32_t)n, 0, 0.0f, 0.0f, 0.0f, 0.0f });
}
// Fast path for concat along dim 0 where one source is stored "transposed"
// (nb[1] == type_size, dim1 innermost) and the other source + dst are
// contiguous along dim 0. The generic concat shader reads the transposed
// source with a catastrophic uncoalesced stride; here we instead copy the
// contiguous source with copy.comp and transpose the other source into the
// matching dst sub-region with the tiled copy_transpose shader (shared-memory
// transpose, coalesced read+write). Mirrors the precedent in
// ggml_vk_cpy_to_contiguous: direct dispatch with custom push constants.
static void ggml_vk_concat_transpose_fastpath(ggml_backend_vk_context * ctx, vk_context& subctx,
const ggml_tensor * ctg, const ggml_tensor * trp,
ggml_tensor * dst, uint32_t off_ctg, uint32_t off_trp) {
const uint32_t ts = ggml_type_size(dst->type);
vk_pipeline pipeline_cpy = (ts == 4) ? ctx->device->pipeline_cpy_f32_f32
: ctx->device->pipeline_cpy_f16_f16;
vk_pipeline pipeline_trp = (ts == 4) ? ctx->device->pipeline_cpy_transpose_32
: ctx->device->pipeline_cpy_transpose_16;
ggml_pipeline_request_descriptor_sets(ctx, pipeline_cpy, 1);
ggml_pipeline_request_descriptor_sets(ctx, pipeline_trp, 1);
vk_subbuffer ctg_buf = ggml_vk_tensor_subbuffer(ctx, ctg, true);
vk_subbuffer trp_buf = ggml_vk_tensor_subbuffer(ctx, trp, true);
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
const uint32_t a_misalign_ctg = get_misalign_bytes(ctx, ctg) / ts;
const uint32_t a_misalign_trp = get_misalign_bytes(ctx, trp) / ts;
const uint32_t d_misalign = get_misalign_bytes(ctx, dst) / ts;
// Dispatch A: contiguous copy of `ctg` into dst[off_ctg : off_ctg + ctg->ne[0], :]
if (ctg->ne[0] > 0) {
const uint32_t ne_ctg = (uint32_t) ggml_nelements(ctg);
vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(ctg, dst, ne_ctg);
pc.ne10 = (uint32_t) ctg->ne[0]; // only ctg's columns in dst
pc.misalign_offsets = (a_misalign_ctg << 16) | (d_misalign + off_ctg);
init_pushconst_fastdiv(pc);
std::array<uint32_t, 3> el = ne_ctg > 262144 ? std::array<uint32_t,3>{512, 512, CEIL_DIV(ne_ctg, 262144)}
: ne_ctg > 512 ? std::array<uint32_t,3>{512, CEIL_DIV(ne_ctg, 512), 1}
: std::array<uint32_t,3>{ne_ctg, 1, 1};
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_cpy, { ctg_buf, dst_buf }, pc, el);
}
// Dispatch B: tiled transpose of `trp` into dst[off_trp : off_trp + trp->ne[0], :]
if (trp->ne[0] > 0) {
vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(trp, dst, ggml_nelements(trp));
pc.ne10 = (uint32_t) trp->ne[0]; // dst bound = trp columns (NOT dst->ne[0])
pc.misalign_offsets = (a_misalign_trp << 16) | (d_misalign + off_trp);
init_pushconst_fastdiv(pc);
std::array<uint32_t, 3> el = {
(uint32_t) CEIL_DIV(trp->ne[0], 32),
(uint32_t) CEIL_DIV(trp->ne[1], 32),
(uint32_t) (trp->ne[2] * trp->ne[3]),
};
el[0] = std::min(el[0], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
el[1] = std::min(el[1], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
el[2] = std::min(el[2], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_trp, { trp_buf, dst_buf }, pc, el);
}
ggml_vk_sync_buffers(ctx, subctx);
}
static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
int * op_params = (int *)dst->op_params;
// Fast path: concat along dim 0 with one source "transposed" (nb[1]==type_size,
// dim1 innermost) and the other source + dst contiguous along dim 0. The generic
// shader reads the transposed source uncoalesced (~5.6ms on RX 580 vs ~130us for
// a coalesced copy); here we instead use a tiled shared-memory transpose.
auto src_transposed_2d = [&](const ggml_tensor * s) {
const uint32_t ts = ggml_type_size(s->type);
return s->nb[1] == ts // dim1 innermost
&& s->nb[0] == s->ne[1] * ts // consistent 2D transpose
&& s->ne[2] == 1 && s->ne[3] == 1;
};
const uint32_t dst_ts = ggml_type_size(dst->type);
const bool dim0_ok = (op_params[0] == 0) && (dst->nb[0] == dst_ts);
const bool types_ok = (dst_ts == 4 || dst_ts == 2)
&& (src0->type == src1->type && src0->type == dst->type)
&& (ggml_blck_size(dst->type) == 1);
const bool shapes_ok = (src0->ne[1] == src1->ne[1] && src1->ne[1] == dst->ne[1])
&& (src0->ne[2] == src1->ne[2] && src0->ne[2] == dst->ne[2])
&& (src0->ne[3] == src1->ne[3] && src0->ne[3] == dst->ne[3])
&& (src0->ne[0] + src1->ne[0] == dst->ne[0]);
// doffset is 16-bit for unary shaders; guard against overflow
const bool off_fits = ((get_misalign_bytes(ctx, dst)/dst_ts + dst->ne[0])) < 0xFFFFu;
const bool s1_trans = dim0_ok && types_ok && shapes_ok && off_fits
&& src_transposed_2d(src1) && ggml_is_contiguous(src0);
const bool s0_trans = dim0_ok && types_ok && shapes_ok && off_fits
&& src_transposed_2d(src0) && ggml_is_contiguous(src1);
if (s1_trans || s0_trans) {
const ggml_tensor * ctg = s1_trans ? src0 : src1; // contiguous source
const ggml_tensor * trp = s1_trans ? src1 : src0; // transposed source
// dst offsets (in elements) where each source region begins
const uint32_t off_ctg = s1_trans ? 0u : (uint32_t) src1->ne[0];
const uint32_t off_trp = s1_trans ? (uint32_t) src0->ne[0] : 0u;
ggml_vk_concat_transpose_fastpath(ctx, subctx, ctg, trp, dst, off_ctg, off_trp);
return;
}
const uint32_t src0_type_size = ggml_type_size(src0->type);
const uint32_t src1_type_size = ggml_type_size(src1->type);
const uint32_t dst_type_size = ggml_type_size(dst->type);
@@ -15734,6 +15843,35 @@ static const char * ggml_backend_vk_name(ggml_backend_t backend) {
return ctx->name.c_str();
}
// Free the compute-scratch device buffers (prealloc_* and the transfer staging buffer) without
// tearing down the backend (pipelines, command pools, fences and device stay alive). These buffers
// are reallocated lazily by ggml_vk_preallocate_buffers() on the next compute, so this just shrinks
// an idle model's device footprint. Mirrors the buffer-freeing subset of ggml_vk_cleanup().
static void ggml_backend_vk_free_scratch(ggml_backend_t backend) {
ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;
VK_LOG_DEBUG("ggml_backend_vk_free_scratch(" << ctx->name << ")");
// discard any unsubmitted command buffer and wait for in-flight work before freeing
ctx->compute_ctx.reset();
ggml_vk_synchronize(ctx);
ggml_vk_destroy_buffer(ctx->prealloc_x);
ggml_vk_destroy_buffer(ctx->prealloc_y);
ggml_vk_destroy_buffer(ctx->prealloc_split_k);
ggml_vk_destroy_buffer(ctx->prealloc_add_rms_partials);
ggml_vk_destroy_buffer(ctx->sync_staging);
ctx->prealloc_y_last_pipeline_used = nullptr;
ctx->prealloc_y_last_tensor_used = nullptr;
ctx->prealloc_y_last_decode_vector_staging = false;
ctx->prealloc_size_x = 0;
ctx->prealloc_size_y = 0;
ctx->prealloc_size_split_k = 0;
ctx->prealloc_size_add_rms_partials = 0;
ctx->prealloc_size_add_rms_partials_offset = 0;
}
static void ggml_backend_vk_free(ggml_backend_t backend) {
ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;
VK_LOG_DEBUG("ggml_backend_vk_free(" << ctx->name << ")");
@@ -16578,7 +16716,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
std::fill(ctx->query_nodes.begin(), ctx->query_nodes.end(), nullptr);
std::fill(ctx->query_node_idx.begin(), ctx->query_node_idx.end(), 0);
GGML_ASSERT(ctx->compute_ctx.expired());
// compute_ctx may hold async host uploads recorded before the graph; append the timestamp after them
compute_ctx = ggml_vk_get_compute_ctx(ctx);
ctx->query_idx = 0;
compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
@@ -17269,6 +17407,7 @@ static ggml_backend_i ggml_backend_vk_interface = {
/* .event_record = */ ggml_backend_vk_event_record,
/* .event_wait = */ ggml_backend_vk_event_wait,
/* .graph_optimize = */ ggml_vk_graph_optimize,
/* .free_scratch = */ ggml_backend_vk_free_scratch,
};
static ggml_guid_t ggml_backend_vk_guid() {
@@ -18243,11 +18382,112 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
return devices[device];
}
// Import an mmap-backed host region as a Vulkan pinned buffer via
// VK_EXT_external_memory_host so H2D uploads DMA straight from system RAM
// instead of bouncing through the staging buffer + host memcpy. Mirrors the
// GGML_CUDA_REGISTER_HOST path; populates device->pinned_memory, which the
// existing pinned fast path in ggml_vk_buffer_write_2d_async looks up.
//
// A single Vulkan buffer cannot cover a whole multi-GB mmap (it is capped at
// device->max_buffer_size), so the region is imported in page-aligned chunks.
// ggml_vk_host_get resolves a tensor pointer to the chunk that contains it,
// and ggml_vk_buffer_write_2d_async falls back to staging for any tensor that
// straddles a chunk boundary, so correctness is preserved.
static bool ggml_backend_vk_register_host_buffer(void * buffer, size_t size) {
if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr && getenv("GGML_VK_REGISTER_HOST") == nullptr) {
return false;
}
if (size == 0) {
return false;
}
bool success = false;
for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
vk_device& device = vk_instance.devices[i];
if (!device || !device->external_memory_host || device->max_buffer_size == 0) {
continue;
}
const size_t align = device->min_imported_host_pointer_alignment;
size_t chunk = device->max_buffer_size & ~(align - 1);
if (chunk == 0) {
continue;
}
uint8_t * p = static_cast<uint8_t *>(buffer);
size_t remaining = size;
bool dev_success = false;
bool import_ok = true; // flips to false once VK_EXT_external_memory_host import fails
while (remaining > 0) {
size_t cur = std::min(remaining, chunk);
vk_buffer buf;
if (import_ok) {
buf = ggml_vk_buffer_from_host_ptr(device, p, cur);
}
if (!buf || !buf->buffer) {
// Fallback for drivers that can't import file-backed mmap pages (e.g. RADV):
// make a one-time copy of the region into a host-visible Vulkan buffer. The GPU
// then DMAs straight from it every eval at full PCIe bandwidth, instead of paying
// a slow single-threaded pageable memcpy into the staging buffer each time.
import_ok = false;
buf = ggml_vk_create_buffer_check(device, cur,
vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent);
if (buf && buf->buffer && buf->ptr) {
memcpy(buf->ptr, p, cur);
} else {
break;
}
}
{
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
device->pinned_memory.emplace_back(p, cur, buf);
}
dev_success = true;
p += cur;
remaining -= cur;
}
if (dev_success) {
success = true;
}
}
return success;
}
static void ggml_backend_vk_unregister_host_buffer(void * buffer) {
for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
vk_device& device = vk_instance.devices[i];
if (!device) {
continue;
}
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
for (auto it = device->pinned_memory.begin(); it != device->pinned_memory.end(); ++it) {
if (std::get<0>(*it) == buffer) {
vk_buffer buf = std::get<2>(*it);
device->pinned_memory.erase(it);
ggml_vk_destroy_buffer(buf);
break;
}
}
}
}
static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
GGML_UNUSED(reg);
if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
return (void *) ggml_backend_vk_register_host_buffer;
}
if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) {
return (void *) ggml_backend_vk_unregister_host_buffer;
}
return nullptr;
}
static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
/* .get_name = */ ggml_backend_vk_reg_get_name,
/* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
/* .get_device = */ ggml_backend_vk_reg_get_device,
/* .get_proc_address = */ NULL,
/* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
};
ggml_backend_reg_t ggml_backend_vk_reg() {
ggml/src/gguf.cpp
+1 -1
View File
@@ -1424,7 +1424,7 @@ void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const vo
struct gguf_writer_base {
size_t written_bytes {0u};
~gguf_writer_base(void) = default;
virtual ~gguf_writer_base(void) = default;
// we bet on devirtualization
virtual void write(int8_t val) = 0;
gguf-py/gguf/constants.py
+39
View File
@@ -200,6 +200,9 @@ class Keys:
HEAD_COUNT = "{arch}.attention.indexer.head_count"
KEY_LENGTH = "{arch}.attention.indexer.key_length"
TOP_K = "{arch}.attention.indexer.top_k"
BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA
LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA
TYPES = "{arch}.attention.indexer.types"
class HyperConnection:
COUNT = "{arch}.hyper_connection.count"
@@ -527,6 +530,7 @@ class MODEL_ARCH(IntEnum):
APERTUS = auto()
COGVLM = auto()
MINIMAXM2 = auto()
MINIMAXM3 = auto()
RND1 = auto()
PANGU_EMBED = auto()
MISTRAL3 = auto()
@@ -773,6 +777,9 @@ class MODEL_TENSOR(IntEnum):
INDEXER_PROJ = auto()
INDEXER_ATTN_K = auto()
INDEXER_ATTN_Q_B = auto()
INDEXER_Q_PROJ = auto()
INDEXER_K_PROJ = auto()
INDEXER_Q_NORM = auto()
INDEXER_COMPRESSOR_WKV = auto()
INDEXER_COMPRESSOR_WGATE = auto()
INDEXER_COMPRESSOR_APE = auto()
@@ -1109,6 +1116,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.GROVEMOE: "grovemoe",
MODEL_ARCH.APERTUS: "apertus",
MODEL_ARCH.MINIMAXM2: "minimax-m2",
MODEL_ARCH.MINIMAXM3: "minimax-m3",
MODEL_ARCH.COGVLM: "cogvlm",
MODEL_ARCH.RND1: "rnd1",
MODEL_ARCH.PANGU_EMBED: "pangu-embedded",
@@ -1354,6 +1362,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj",
MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k",
MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b",
MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj",
MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj",
MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm",
MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv",
MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate",
MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape",
@@ -4162,6 +4173,34 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_UP_EXP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
],
MODEL_ARCH.MINIMAXM3: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.OUTPUT,
MODEL_TENSOR.ATTN_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.FFN_NORM,
MODEL_TENSOR.FFN_GATE_INP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
MODEL_TENSOR.FFN_GATE_EXP,
MODEL_TENSOR.FFN_DOWN_EXP,
MODEL_TENSOR.FFN_UP_EXP,
MODEL_TENSOR.FFN_GATE_SHEXP,
MODEL_TENSOR.FFN_DOWN_SHEXP,
MODEL_TENSOR.FFN_UP_SHEXP,
MODEL_TENSOR.FFN_GATE,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
MODEL_TENSOR.INDEXER_Q_PROJ,
MODEL_TENSOR.INDEXER_K_PROJ,
MODEL_TENSOR.INDEXER_Q_NORM,
MODEL_TENSOR.INDEXER_K_NORM,
],
MODEL_ARCH.COGVLM: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
gguf-py/gguf/gguf_writer.py
+10
View File
@@ -793,6 +793,16 @@ class GGUFWriter:
def add_indexer_top_k(self, top_k: int) -> None:
self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k)
def add_indexer_block_size(self, block_size: int) -> None:
self.add_uint32(Keys.Attention.Indexer.BLOCK_SIZE.format(arch=self.arch), block_size)
def add_indexer_local_blocks(self, local_blocks: int) -> None:
self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks)
def add_indexer_types(self, value: Sequence[bool]) -> None:
key = Keys.Attention.Indexer.TYPES.format(arch=self.arch)
self.add_array(key, value)
def add_max_alibi_bias(self, bias: float) -> None:
self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
gguf-py/gguf/tensor_mapping.py
+14 -1
View File
@@ -1264,7 +1264,8 @@ class TensorNameMap:
),
MODEL_TENSOR.INDEXER_K_NORM: (
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
"model.layers.{bid}.self_attn.index_k_norm", # MSA
),
MODEL_TENSOR.INDEXER_PROJ: (
@@ -1279,6 +1280,18 @@ class TensorNameMap:
"model.layers.{bid}.self_attn.indexer.wq_b", # DSA
),
MODEL_TENSOR.INDEXER_Q_PROJ: (
"model.layers.{bid}.self_attn.index_q_proj", # MSA
),
MODEL_TENSOR.INDEXER_K_PROJ: (
"model.layers.{bid}.self_attn.index_k_proj", # MSA
),
MODEL_TENSOR.INDEXER_Q_NORM: (
"model.layers.{bid}.self_attn.index_q_norm", # MSA
),
############################################################################
# TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
MODEL_TENSOR.ENC_OUTPUT_NORM: (
include/llama.h
+21 -3
View File
@@ -202,6 +202,16 @@ extern "C" {
LLAMA_SPLIT_MODE_TENSOR = 3,
};
enum llama_load_mode {
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
LLAMA_LOAD_MODE_MLOCK = 2, // mmap + force system to keep model in RAM rather than swapping or compressing
LLAMA_LOAD_MODE_DIRECT_IO = 3, // use direct I/O if available
};
LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str);
enum llama_context_type {
LLAMA_CONTEXT_TYPE_DEFAULT = 0,
LLAMA_CONTEXT_TYPE_MTP = 1,
@@ -301,6 +311,7 @@ extern "C" {
int32_t n_gpu_layers; // number of layers to store in VRAM, a negative value means all layers
enum llama_split_mode split_mode; // how to split the model across multiple GPUs
enum llama_load_mode load_mode; // how to load the model
// the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE
int32_t main_gpu;
@@ -321,9 +332,6 @@ extern "C" {
// Keep the booleans together to avoid misalignment during copy-by-value.
bool vocab_only; // only load the vocabulary, no weights
bool use_mmap; // use mmap if possible
bool use_direct_io; // use direct io, takes precedence over use_mmap when supported
bool use_mlock; // force system to keep model in RAM
bool check_tensors; // validate model tensor data
bool use_extra_bufts; // use extra buffer types (used for weight repacking)
bool no_host; // bypass host buffer allowing extra buffers to be used
@@ -557,6 +565,16 @@ extern "C" {
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx);
LLAMA_API llama_memory_t llama_get_memory (const struct llama_context * ctx);
// On-demand device (VRAM) residency: free the model's GPU weight buffers (keeping a host
// shadow so no reload is needed) to hand VRAM to another model, and rebuild them on demand.
// Any subsequent llama_decode auto-restores. If evict_kv is true, the KV cache device buffers
// are also evicted to a host shadow (for when the KV would not fit alongside the other model),
// at the cost of a D2H/H2D copy of the live KV each cycle; otherwise the KV stays resident.
// Intended for time-sharing a single GPU between several always-loaded models.
LLAMA_API void llama_context_release_device(struct llama_context * ctx, bool evict_kv);
LLAMA_API void llama_context_restore_device(struct llama_context * ctx);
LLAMA_API bool llama_context_weights_resident(const struct llama_context * ctx);
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type
LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
scripts/compare-llama-bench.py
+5 -5
View File
@@ -28,7 +28,7 @@ LLAMA_BENCH_DB_FIELDS = [
"model_type", "model_size", "model_n_params", "n_batch", "n_ubatch", "n_threads",
"cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers",
"split_mode", "main_gpu", "no_kv_offload", "flash_attn", "tensor_split", "tensor_buft_overrides",
"use_mmap", "embeddings", "no_op_offload", "n_prompt", "n_gen", "n_depth",
"load_mode", "embeddings", "no_op_offload", "n_prompt", "n_gen", "n_depth",
"test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts", "n_cpu_moe",
"fit_target", "fit_min_ctx"
]
@@ -38,7 +38,7 @@ LLAMA_BENCH_DB_TYPES = [
"TEXT", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER",
"TEXT", "INTEGER", "INTEGER", "TEXT", "TEXT", "INTEGER",
"TEXT", "INTEGER", "INTEGER", "INTEGER", "TEXT", "TEXT",
"INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER",
"TEXT", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER",
"TEXT", "INTEGER", "INTEGER", "REAL", "REAL", "INTEGER",
"INTEGER", "INTEGER"
]
@@ -63,7 +63,7 @@ assert len(TEST_BACKEND_OPS_DB_FIELDS) == len(TEST_BACKEND_OPS_DB_TYPES)
LLAMA_BENCH_KEY_PROPERTIES = [
"cpu_info", "gpu_info", "backends", "n_gpu_layers", "n_cpu_moe", "tensor_buft_overrides", "model_filename", "model_type",
"n_batch", "n_ubatch", "embeddings", "cpu_mask", "cpu_strict", "poll", "n_threads", "type_k", "type_v",
"use_mmap", "no_kv_offload", "split_mode", "main_gpu", "tensor_split", "flash_attn", "n_prompt", "n_gen", "n_depth",
"load_mode", "no_kv_offload", "split_mode", "main_gpu", "tensor_split", "flash_attn", "n_prompt", "n_gen", "n_depth",
"fit_target", "fit_min_ctx"
]
@@ -73,7 +73,7 @@ TEST_BACKEND_OPS_KEY_PROPERTIES = [
]
# Properties that are boolean and are converted to Yes/No for the table:
LLAMA_BENCH_BOOL_PROPERTIES = ["embeddings", "cpu_strict", "use_mmap", "no_kv_offload", "flash_attn"]
LLAMA_BENCH_BOOL_PROPERTIES = ["embeddings", "cpu_strict", "no_kv_offload", "flash_attn"]
TEST_BACKEND_OPS_BOOL_PROPERTIES = ["supported", "passed"]
# Header names for the table (llama-bench):
@@ -82,7 +82,7 @@ LLAMA_BENCH_PRETTY_NAMES = {
"tensor_buft_overrides": "Tensor overrides", "model_filename": "File", "model_type": "Model", "model_size": "Model size [GiB]",
"model_n_params": "Num. of par.", "n_batch": "Batch size", "n_ubatch": "Microbatch size", "embeddings": "Embeddings",
"cpu_mask": "CPU mask", "cpu_strict": "CPU strict", "poll": "Poll", "n_threads": "Threads", "type_k": "K type", "type_v": "V type",
"use_mmap": "Use mmap", "no_kv_offload": "NKVO", "split_mode": "Split mode", "main_gpu": "Main GPU", "tensor_split": "Tensor split",
"load_mode": "Load mode", "no_kv_offload": "NKVO", "split_mode": "Split mode", "main_gpu": "Main GPU", "tensor_split": "Tensor split",
"flash_attn": "FlashAttention",
}
scripts/snapdragon/ggml-hexagon-profile.py
+237 -93
View File
@@ -6,6 +6,7 @@ import re
import argparse
import statistics
import logging
import bisect
from typing import Any, Dict, List, Optional
from collections import defaultdict
@@ -30,7 +31,7 @@ op_pattern = re.compile(
)
trace_pattern = re.compile(
r"trace-op\s+(?P<op_name>[A-Z_0-9+]+):\s+thread\s+(?P<thread>\d+)\s+event\s+(?P<event>[A-Z_0-9\-]+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
)
logger = logging.getLogger("ggml-hexagon-profile")
@@ -50,9 +51,13 @@ def normalize_event_name(evt_type):
class CycleUnwrapper:
def __init__(self):
self.last_raw = None
self.high_part = 0
def __init__(self, initial_val=None):
if initial_val is not None:
self.last_raw = initial_val & 0xFFFFFFFF
self.high_part = initial_val & 0xFFFFFFFF00000000
else:
self.last_raw = None
self.high_part = 0
def unwrap(self, raw):
if self.last_raw is None:
@@ -78,10 +83,12 @@ def parse_log(file_path, pmu_index=None):
sys.exit(1)
all_ops: List[Dict[str, Any]] = []
all_traces: List[Dict[str, Any]] = []
current_op: Optional[Dict[str, Any]] = None
timestamp_pattern = re.compile(r"^(?P<min>\d+)\.(?P<sec>\d+)\.(?P<ms>\d+)\.(?P<us>\d+)\s+[A-Z]\s+")
unwrapper = CycleUnwrapper()
unwrapper = None
trace_unwrapper = None
for line in f:
ts_match = timestamp_pattern.match(line)
@@ -100,6 +107,7 @@ def parse_log(file_path, pmu_index=None):
if not prefix_match:
continue
names = parts[1]
if len(parts) == 7:
dims, types, timings = parts[2], parts[3], parts[6]
elif len(parts) == 6:
@@ -120,6 +128,7 @@ def parse_log(file_path, pmu_index=None):
op_match = op_pattern.search(line)
if op_match:
op_name = op_match.group('op_name')
names = ""
dims = op_match.group('dims').strip()
types = op_match.group('types').strip()
else:
@@ -136,24 +145,31 @@ def parse_log(file_path, pmu_index=None):
except (ValueError, IndexError):
pmu_val = None
evt_raw = op_match.group('evt') if 'evt' in op_match.groupdict() else None
evt_val = None
if evt_raw:
evt_val = None
if types.startswith("evt-cnt "):
try:
evt_val = [int(x.strip()) for x in evt_raw.split(',')]
evt_val = [int(x.strip()) for x in types[8:].split(',')]
except ValueError:
evt_val = None
cycles_start_raw = op_match.group('start')
unwrapped_cycles_start = None
if cycles_start_raw:
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
if op_name == "OPBATCH":
if cycles_start_raw:
unwrapped_cycles_start = int(cycles_start_raw)
unwrapper = CycleUnwrapper(unwrapped_cycles_start)
trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
else:
if cycles_start_raw and unwrapper is not None:
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
idx = line.find("profile-op ")
op_text = line[idx + 11:].strip() if idx != -1 else line.strip()
current_op = {
'name': op_name,
'names': names,
'dims': dims,
'types': types,
'op_text': op_text,
@@ -170,110 +186,239 @@ def parse_log(file_path, pmu_index=None):
continue
trace_match = trace_pattern.search(line)
if trace_match and current_op:
if trace_match.group('op_name') == current_op['name']:
raw_cyc = int(trace_match.group('cycles'))
current_op['trace_events'].append({
'thread': int(trace_match.group('thread')),
'event': trace_match.group('event'),
'info': int(trace_match.group('info')),
'cycles': raw_cyc,
'unwrapped_cycles': unwrapper.unwrap(raw_cyc),
'state': trace_match.group('state')
})
if trace_match:
raw_cyc = int(trace_match.group('cycles'))
unwrapped_cyc = None
if trace_unwrapper is not None:
unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
all_traces.append({
'thread': int(trace_match.group('thread')),
'event': trace_match.group('event'),
'info': int(trace_match.group('info')),
'cycles': raw_cyc,
'unwrapped_cycles': unwrapped_cyc,
'state': trace_match.group('state')
})
f.close()
# Assign start/end cycles to all ops
for op in all_ops:
op['start_cycles'] = op['unwrapped_cycles_start']
op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None
# Filter ops with valid start_cycles
valid_ops = [op for op in all_ops if op['start_cycles'] is not None and op['end_cycles'] is not None]
# Separate OPBATCH ops from other ops
opbatch_ops = [op for op in valid_ops if op['name'] == "OPBATCH"]
other_ops = [op for op in valid_ops if op['name'] != "OPBATCH"]
# Sort them by start_cycles to enable binary search
opbatch_ops.sort(key=lambda op: op['start_cycles'])
other_ops.sort(key=lambda op: op['start_cycles'])
opbatch_starts = [op['start_cycles'] for op in opbatch_ops]
other_starts = [op['start_cycles'] for op in other_ops]
# Map trace events to any operator whose cycles contain them
for e in all_traces:
cyc = e['unwrapped_cycles']
if cyc is None:
continue
# Map to OPBATCH
idx = bisect.bisect_right(opbatch_starts, cyc) - 1
if idx >= 0:
op = opbatch_ops[idx]
if op['start_cycles'] <= cyc <= op['end_cycles']:
op['trace_events'].append(e)
# Map to other ops
idx = bisect.bisect_right(other_starts, cyc) - 1
if idx >= 0:
op = other_ops[idx]
if op['start_cycles'] <= cyc <= op['end_cycles']:
op['trace_events'].append(e)
return all_ops
def print_ascii_timeline(op_name, dims, types, usec, cycles, events, evt_val=None):
evt_str = ""
if evt_val:
evt_str = " - evt [" + ",".join(str(x) for x in evt_val) + "]"
def print_bubbles_timeline(op):
op_name = op['name']
dims = op['dims']
types = op['types']
usec = op['usec']
cycles = op['cycles']
events = op['trace_events']
logger.info("=" * 100)
logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles{evt_str}")
logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles")
logger.info("=" * 100)
events = sorted(events, key=lambda e: e['cycles'])
if not events:
logger.info(" No trace events recorded.")
return
min_cycles = events[0]['cycles']
# Identify start and end cycles for this operator
op_start = op['start_cycles']
op_end = op['end_cycles']
if op_start is None or op_end is None:
logger.info(" Cannot analyze bubbles: missing start/end cycle counts.")
return
logger.info("Cycles %-30s" % "EventDetails" + " ".join(f"T{i:<2}" for i in range(10)) + " HMX")
logger.info("-" * 100)
thread_stacks = [[] for _ in range(11)]
batch_duration = op_end - op_start
if batch_duration <= 0:
logger.info(" Cannot analyze bubbles: batch duration is 0.")
return
# Group events by (thread, track_type)
tracks = defaultdict(list)
for e in events:
t = e['thread']
if t < 0 or t > 10:
continue
is_dma = (normalize_event_name(e['event']) == 'DMA')
track_type = 'dma' if is_dma else 'compute'
tracks[(t, track_type)].append(e)
if e['cycles'] >= min_cycles:
rel_cycles = e['cycles'] - min_cycles
else:
rel_cycles = (e['cycles'] + 0x100000000) - min_cycles
active_threads = sorted(list(set(t for (t, track_type) in tracks.keys())))
if not active_threads:
logger.info(" No active threads in trace.")
return
state = e['state']
evt_type = e['event']
bubble_threshold = 10000 # 10k cycles
# Determine char representing the event
norm_evt = normalize_event_name(evt_type)
char = '?'
if norm_evt == 'V-COMP':
char = 'V'
elif norm_evt == 'M-COMP':
char = 'H'
elif norm_evt == 'A-QUANT':
char = 'Q'
elif norm_evt == 'A-PREP':
char = 'A'
elif norm_evt == 'Q-PREP':
char = 'q'
elif norm_evt == 'K-PREP':
char = 'k'
elif norm_evt == 'V-PREP':
char = 'v'
elif norm_evt == 'W-DEQUANT':
char = 'D'
elif norm_evt == 'O-PROC':
char = 'O'
elif norm_evt == 'W-PREP':
char = 'P'
elif norm_evt == 'DMA':
char = 'M'
thread_stats = {}
for t in active_threads:
thread_stats[t] = {
'compute_idle_cycles': batch_duration,
'compute_idle_pct': 100.0,
'compute_bubbles': [],
if state == 'start':
thread_stacks[t].append(char)
elif state == 'stop':
if thread_stacks[t]:
if thread_stacks[t][-1] == char:
thread_stacks[t].pop()
elif char in thread_stacks[t]:
thread_stacks[t].remove(char)
else:
thread_stacks[t].pop()
'dma_idle_cycles': batch_duration,
'dma_idle_pct': 100.0,
'dma_bubbles': []
}
cols = []
for i in range(11):
if thread_stacks[i]:
cols.append(f"[{thread_stacks[i][-1]}]")
total_compute_idle_pct = 0.0
total_dma_idle_pct = 0.0
for t in active_threads:
for track_type in ['compute', 'dma']:
key = (t, track_type)
track_events = tracks.get(key, [])
if not track_events:
gaps = [(op_start, op_end)]
idle_cycles = batch_duration
else:
cols.append(" | ")
track_events = sorted(track_events, key=lambda e: e.get('unwrapped_cycles') or e['cycles'])
evt_desc = f"T{t}: {evt_type} {state} ({e['info']})"
logger.info(f"{rel_cycles:10d} %-30s" % evt_desc + " ".join(cols[:10]) + " " + cols[10])
active_intervals = []
active_count = 0
curr_start = None
for e in track_events:
cyc = e.get('unwrapped_cycles') or e['cycles']
cyc = max(op_start, min(op_end, cyc))
state = e['state']
if state == 'start':
if active_count == 0:
curr_start = cyc
active_count += 1
elif state == 'stop':
if active_count > 0:
active_count -= 1
if active_count == 0:
active_intervals.append((curr_start, cyc))
else:
active_intervals.append((op_start, cyc))
if active_count > 0 and curr_start is not None:
active_intervals.append((curr_start, op_end))
# Merge intervals
active_intervals.sort(key=lambda x: x[0])
merged_intervals = []
for start, end in active_intervals:
if not merged_intervals:
merged_intervals.append([start, end])
else:
last_start, last_end = merged_intervals[-1]
if start <= last_end:
merged_intervals[-1][1] = max(last_end, end)
else:
merged_intervals.append([start, end])
# Calculate gaps
gaps = []
curr_time = op_start
for start, end in merged_intervals:
if start > curr_time:
gaps.append((curr_time, start))
curr_time = max(curr_time, end)
if curr_time < op_end:
gaps.append((curr_time, op_end))
idle_cycles = sum(end - start for start, end in gaps)
idle_pct = (idle_cycles / batch_duration) * 100.0
bubbles = []
for start, end in gaps:
dur = end - start
if dur >= bubble_threshold:
bubbles.append((start, end, dur))
if track_type == 'compute':
thread_stats[t]['compute_idle_cycles'] = idle_cycles
thread_stats[t]['compute_idle_pct'] = idle_pct
thread_stats[t]['compute_bubbles'] = bubbles
total_compute_idle_pct += idle_pct
else:
thread_stats[t]['dma_idle_cycles'] = idle_cycles
thread_stats[t]['dma_idle_pct'] = idle_pct
thread_stats[t]['dma_bubbles'] = bubbles
total_dma_idle_pct += idle_pct
avg_compute_idle = total_compute_idle_pct / len(active_threads)
avg_dma_idle = total_dma_idle_pct / len(active_threads)
logger.info(" Combined Idle Statistics:")
logger.info(f" Active Threads : {', '.join(str(t) for t in active_threads)}")
logger.info(f" Avg Thread Compute IDLE : {avg_compute_idle:.1f}%")
logger.info(f" Avg Thread DMA IDLE : {avg_dma_idle:.1f}%")
logger.info("-" * 100)
logger.info(" Per-Thread Idle Analysis:")
for t in active_threads:
stats = thread_stats[t]
thread_name = f"Thread {t:<2} (HVX)" if t != 10 else "Thread 10 (HMX)"
logger.info(f" {thread_name} -> Compute Idle: {stats['compute_idle_pct']:.1f}% | DMA Idle: {stats['dma_idle_pct']:.1f}%")
def print_ascii_summary(op_name, dims, types, usec, cycles, events, evt_val=None):
evt_str = ""
if evt_val:
evt_str = " - evt [" + ",".join(str(x) for x in evt_val) + "]"
all_bubbles = []
for t in active_threads:
stats = thread_stats[t]
for start, end, dur in stats['compute_bubbles']:
pct = (dur / batch_duration) * 100.0
all_bubbles.append((dur, f"Thread {t} Compute: bubble of {dur} cycles ({pct:.1f}%) at {start - op_start} to {end - op_start}"))
for start, end, dur in stats['dma_bubbles']:
pct = (dur / batch_duration) * 100.0
all_bubbles.append((dur, f"Thread {t} DMA : bubble of {dur} cycles ({pct:.1f}%) at {start - op_start} to {end - op_start}"))
if all_bubbles:
logger.info("-" * 100)
logger.info(f" Significant Bubbles (>= {bubble_threshold} cycles):")
all_bubbles.sort(key=lambda x: x[0], reverse=True)
for dur, desc in all_bubbles[:15]:
logger.info(f" {desc}")
else:
logger.info("-" * 100)
logger.info(f" No significant bubbles detected (all idle gaps < {bubble_threshold} cycles).")
def print_ascii_summary(op_name, dims, types, usec, cycles, events):
logger.info("=" * 100)
logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles{evt_str}")
logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles")
logger.info("=" * 100)
events = sorted(events, key=lambda e: e['cycles'])
@@ -415,8 +560,8 @@ def main():
parser.add_argument("--pmu-index", type=int)
parser.add_argument("--pmu-name", type=str)
parser.add_argument("--width", action='append', default=['dims:40'], help="Override column width, e.g. --width dims:50")
parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "diagram"],
help="Output ASCII art event summary or timing diagram (default: summary)")
parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "bubbles"],
help="Output ASCII art event summary or thread idle bubble analysis (default: summary)")
parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line")
group = parser.add_mutually_exclusive_group()
@@ -457,12 +602,11 @@ def main():
ops = ops[-args.tail:]
if args.timeline:
logger.info(f"\n# ASCII Timing {args.timeline.capitalize()}\n")
for op in ops:
if args.timeline == "summary":
print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val'))
elif args.timeline == "diagram":
print_ascii_timeline(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val'))
print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'])
elif args.timeline == "bubbles":
print_bubbles_timeline(op)
else:
generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)
scripts/snapdragon/ggml-hexagon-trace.py
+130 -40
View File
@@ -6,6 +6,7 @@ import re
import argparse
import statistics
import logging
import bisect
from typing import Any, Dict, List, Optional
from collections import defaultdict
@@ -16,11 +17,11 @@ op_pattern = re.compile(
)
trace_pattern = re.compile(
r"trace-op\s+(?P<op_name>[A-Z_0-9+]+):\s+thread\s+(?P<thread>\d+)\s+event\s+(?P<event>[A-Z_0-9\-]+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
)
def normalize_event_name(evt_type):
def normalize_event_name(evt_type, info=0):
if evt_type == "HVX_COMP":
return "V-COMP"
if evt_type == "HMX_COMP":
@@ -32,9 +33,13 @@ def normalize_event_name(evt_type):
class CycleUnwrapper:
def __init__(self):
self.last_raw = None
self.high_part = 0
def __init__(self, initial_val=None):
if initial_val is not None:
self.last_raw = initial_val & 0xFFFFFFFF
self.high_part = initial_val & 0xFFFFFFFF00000000
else:
self.last_raw = None
self.high_part = 0
def unwrap(self, raw):
if self.last_raw is None:
@@ -60,8 +65,10 @@ def parse_log(file_path):
sys.exit(1)
all_ops: List[Dict[str, Any]] = []
all_traces: List[Dict[str, Any]] = []
current_op: Optional[Dict[str, Any]] = None
unwrapper = CycleUnwrapper()
unwrapper = None
trace_unwrapper = None
line_idx = 0
for line in f:
@@ -73,6 +80,7 @@ def parse_log(file_path):
if not prefix_match:
continue
names = parts[1]
if len(parts) == 7:
dims, types, strides, params, timings = parts[2], parts[3], parts[4], parts[5], parts[6]
elif len(parts) == 6:
@@ -93,6 +101,7 @@ def parse_log(file_path):
op_match = op_pattern.search(line)
if op_match:
op_name = op_match.group('op_name')
names = ""
dims = op_match.group('dims').strip() if op_match.group('dims') else ''
types = op_match.group('types').strip() if op_match.group('types') else ''
strides = op_match.group('strides').strip() if op_match.group('strides') else ''
@@ -103,18 +112,30 @@ def parse_log(file_path):
if op_match:
cycles_start_raw = op_match.group('start')
unwrapped_cycles_start = None
if cycles_start_raw:
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
if op_name == "OPBATCH":
if cycles_start_raw:
unwrapped_cycles_start = int(cycles_start_raw)
unwrapper = CycleUnwrapper(unwrapped_cycles_start)
trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
else:
if cycles_start_raw and unwrapper is not None:
unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
idx = line.find("profile-op ")
op_text = line[idx + 11:].strip() if idx != -1 else line.strip()
evt_str = None
if types.startswith("evt-cnt "):
evt_str = types[8:].strip()
current_op = {
'name': op_name,
'names': names,
'dims': dims,
'types': types,
'strides': strides,
'params': params,
'evt': evt_str,
'op_text': op_text,
'usec': int(op_match.group('usec')),
'cycles': int(op_match.group('cycles')),
@@ -127,20 +148,22 @@ def parse_log(file_path):
continue
trace_match = trace_pattern.search(line)
if trace_match and current_op:
if trace_match.group('op_name') == current_op['name']:
raw_cyc = int(trace_match.group('cycles'))
current_op['trace_events'].append({
'thread': int(trace_match.group('thread')),
'event': trace_match.group('event'),
'info': int(trace_match.group('info')),
'cycles': raw_cyc,
'unwrapped_cycles': unwrapper.unwrap(raw_cyc),
'state': trace_match.group('state')
})
if trace_match:
raw_cyc = int(trace_match.group('cycles'))
unwrapped_cyc = None
if trace_unwrapper is not None:
unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
all_traces.append({
'thread': int(trace_match.group('thread')),
'event': trace_match.group('event'),
'info': int(trace_match.group('info')),
'cycles': raw_cyc,
'unwrapped_cycles': unwrapped_cyc,
'state': trace_match.group('state')
})
f.close()
return all_ops
return all_ops, all_traces
# --- Simple protobuf encoder ---
@@ -246,7 +269,7 @@ def write_trace_packet_to_file(f, packet_bytes):
# --- End Protobuf Encoder ---
def generate_perfetto_trace(filtered_ops, output_path):
def generate_perfetto_trace(filtered_ops, trace_events, output_path):
if not filtered_ops:
logger.warning("No operators found after filtering.")
return
@@ -269,14 +292,12 @@ def generate_perfetto_trace(filtered_ops, output_path):
# Process events
completed_events = []
for op in filtered_ops:
events = op['trace_events']
if not events:
continue
events = sorted(events, key=lambda e: e['unwrapped_cycles'])
if trace_events:
trace_events = sorted(trace_events, key=lambda e: e['unwrapped_cycles'])
one_usec_cycles = max(avg_freq_mhz, 1.0)
active_starts = {}
for e in events:
for e in trace_events:
t = e['thread']
evt = e['event']
info = e['info']
@@ -285,6 +306,17 @@ def generate_perfetto_trace(filtered_ops, output_path):
key = (t, evt, info)
if state == 'start':
# Handle missing stop (start followed by another start)
if key in active_starts:
prev_start = active_starts[key]
completed_events.append({
'thread': t,
'event': evt,
'info': info,
'start_cyc': prev_start,
'end_cyc': prev_start + one_usec_cycles,
'missing_stop': True,
})
active_starts[key] = cyc
elif state == 'stop':
if key in active_starts:
@@ -296,8 +328,29 @@ def generate_perfetto_trace(filtered_ops, output_path):
'info': info,
'start_cyc': start_cyc,
'end_cyc': cyc,
'op_name': op['name']
})
else:
# Handle missing start (stop without start)
completed_events.append({
'thread': t,
'event': evt,
'info': info,
'start_cyc': cyc - one_usec_cycles,
'end_cyc': cyc,
'missing_start': True,
})
# Clear remaining unmatched starts
for key, start_cyc in active_starts.items():
t, evt, info = key
completed_events.append({
'thread': t,
'event': evt,
'info': info,
'start_cyc': start_cyc,
'end_cyc': start_cyc + one_usec_cycles,
'missing_stop': True,
})
completed_events.sort(key=lambda e: e['start_cyc'])
@@ -316,7 +369,7 @@ def generate_perfetto_trace(filtered_ops, output_path):
ts = e['ts_ns']
dur = e['dur_ns']
norm_evt = normalize_event_name(evt)
norm_evt = normalize_event_name(evt, e['info'])
if norm_evt == "DMA":
track_key = (t, "DMA")
elif t == 10:
@@ -343,7 +396,7 @@ def generate_perfetto_trace(filtered_ops, output_path):
evt = e['event']
slot = e['slot']
norm_evt = normalize_event_name(evt)
norm_evt = normalize_event_name(evt, e['info'])
if norm_evt == "DMA":
track_evt = "DMA"
evt_id = 1
@@ -421,18 +474,26 @@ def generate_perfetto_trace(filtered_ops, output_path):
for op in filtered_ops:
op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000))
op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000))
if op_start_ns < last_op_end_ns:
op_start_ns = last_op_end_ns
clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us)
if op['name'] != "OPBATCH":
if op_start_ns < last_op_end_ns:
op_start_ns = last_op_end_ns
clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us)
last_op_end_ns = op_start_ns + clamped_dur
else:
clamped_dur = max(op_dur_ns, 100)
# Debug annotations for Ops
debug_annots = []
if 'line_num' in op:
debug_annots.append(make_debug_annotation("line", int_val=op['line_num']))
if 'strides' in op and op['strides']:
if 'names' in op and op['names'] and op['names'] != '----':
debug_annots.append(make_debug_annotation("names", string_val=op['names']))
if 'strides' in op and op['strides'] and op['strides'] != '----':
debug_annots.append(make_debug_annotation("strides", string_val=op['strides']))
if 'params' in op and op['params'] and op['params'] != '----':
debug_annots.append(make_debug_annotation("params", string_val=op['params']))
if 'evt' in op and op['evt']:
debug_annots.append(make_debug_annotation("evt", string_val=op['evt']))
# Slice Begin
evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots)
@@ -444,15 +505,21 @@ def generate_perfetto_trace(filtered_ops, output_path):
packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end)
write_trace_packet_to_file(f, packet_end)
last_op_end_ns = op_start_ns + clamped_dur
# Emit Thread Trace Events
for e in completed_events:
norm_name = normalize_event_name(e['event'])
norm_name = normalize_event_name(e['event'], e['info'])
name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name
if e.get('missing_start') or e.get('missing_stop'):
name += "!"
debug_annots = []
if e.get('missing_start'):
debug_annots.append(make_debug_annotation("missing_start", string_val="true"))
if e.get('missing_stop'):
debug_annots.append(make_debug_annotation("missing_stop", string_val="true"))
# Slice Begin
evt_begin = make_track_event(1, e['uuid'], name=name, category="trace")
evt_begin = make_track_event(1, e['uuid'], name=name, category="trace", debug_annotations=debug_annots if debug_annots else None)
packet_begin = make_trace_packet(e['ts_ns'], track_event=evt_begin)
write_trace_packet_to_file(f, packet_begin)
@@ -477,7 +544,7 @@ def main():
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format='%(message)s')
ops = parse_log(args.logfile)
ops, traces = parse_log(args.logfile)
if args.filter:
try:
@@ -492,7 +559,30 @@ def main():
elif args.tail is not None:
ops = ops[-args.tail:]
generate_perfetto_trace(ops, args.output)
if args.filter or args.head is not None or args.tail is not None:
valid_ranges = []
for op in ops:
start_cyc = op['unwrapped_cycles_start']
end_cyc = start_cyc + op['cycles'] if start_cyc is not None else None
if start_cyc is not None and end_cyc is not None:
valid_ranges.append((start_cyc, end_cyc))
valid_ranges.sort(key=lambda r: r[0])
range_starts = [r[0] for r in valid_ranges]
filtered_traces = []
for e in traces:
cyc = e['unwrapped_cycles']
if cyc is None:
continue
idx = bisect.bisect_right(range_starts, cyc) - 1
if idx >= 0:
start, end = valid_ranges[idx]
if start <= cyc <= end:
filtered_traces.append(e)
traces = filtered_traces
generate_perfetto_trace(ops, traces, args.output)
if __name__ == "__main__":
scripts/sync_vendor.py
+2 -2
View File
@@ -5,7 +5,7 @@ import os
import sys
import subprocess
HTTPLIB_VERSION = "refs/tags/v0.50.1"
HTTPLIB_VERSION = "refs/tags/v0.51.0"
vendor = {
"https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp",
@@ -21,7 +21,7 @@ vendor = {
f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/split.py": "split.py",
f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/LICENSE": "vendor/cpp-httplib/LICENSE",
"https://raw.githubusercontent.com/sheredom/subprocess.h/b49c56e9fe214488493021017bf3954b91c7c1f5/subprocess.h": "vendor/sheredom/subprocess.h",
"https://raw.githubusercontent.com/sheredom/subprocess.h/8671cee1fc09f11a70ce3782a0ee13177c3aa387/subprocess.h": "vendor/sheredom/subprocess.h",
}
for url, filename in vendor.items():
skills/add-new-model/SKILL.md
+98
View File
@@ -0,0 +1,98 @@
---
name: add-new-model
description: Guided workflow for adding a new model architecture to llama.cpp. Use when the user wants to add/port a new model architecture.
---
# Add a new model architecture to llama.cpp
This skill walks a contributor through adding a new model architecture. AI-generated code is permitted in this project, so you may write full implementations for the steps below rather than only pointing at patterns - but follow `AGENTS.md`'s AI usage policy throughout:
- The contributor is 100% responsible for every line, however it was produced. They must be able to explain and defend any part of it to a reviewer. Check in with them as you go (don't silently generate everything and hand over a finished diff) so they actually absorb what was written.
- Before writing code, make sure the contributor owns the design choices for this architecture (which reference model to follow, how non-standard bits like RoPE variants or MoE routing should be handled) - AI accelerates a design the contributor has already made, it doesn't make the design for them.
- Disclosure is mandatory: any AI-meaningful contribution must be disclosed per the PR template. Remind the contributor of this before they open the PR.
- Never write the PR description, commit message, GitHub issue/discussion post, or reviewer replies - those must come from the contributor. If asked to commit on their behalf, use `Assisted-by:` (never `Co-authored-by:`) and only after explicit confirmation.
- If the requested change looks large or introduces a new pattern not covered here, pause and tell the user this kind of change is likely to need prior discussion with maintainers before a PR.
- Keep the PR self-contained. If the work would require a lot of unconventional changes outside the new model file(s) (e.g. touching shared graph-building code, the sampler, or core APIs in ways other models don't), STOP and tell the contributor to open a discussion/issue first - invasive or excessive changes get closed without full review.
- Do not bundle unrelated work into this PR - see Step 4 and Step 5 below for the specifics on multimodal and chat-template/parsing work.
- Never hack around RoPE with a custom sin/cos implementation. Several past PRs tried this and were closed. If the existing `ggml_rope_ext` (see Step 2's RoPE tips) genuinely cannot express what this model needs, the contributor should open an issue to discuss it with maintainers first - not send a PR with a custom RoPE implementation.
Before starting, read `CONTRIBUTING.md`, `AGENTS.md` and `docs/development/HOWTO-add-model.md` if they are not already in context. Also run `git log --oneline -- src/models` and look at at least 3 recent PRs that added a model (their merge commits/diffs) - this shows current convention more reliably than the docs, which can lag behind.
## Step 0 - Scope and dedup check
Ask the contributor:
1. Which model (HF repo id or name)? Is it text-only or does it have a multimodal (vision/audio) encoder?
2. Do they already have the HF `config.json`/weights available locally?
3. Have they checked for an existing PR/issue on this model? Suggest `gh search issues "<model name>"` and `gh search prs "<model name>"` in the `ggml-org/llama.cpp` repo. If an existing PR covers it, the contributor should comment there and collaborate rather than open a duplicate (per CONTRIBUTING.md's AI Usage Policy).
4. What existing supported architecture is this model closest to (e.g. "Llama-like with sliding window", "MoE like DBRX", "BERT-style encoder")?
If the contributor doesn't know the closest reference architecture, you may grep `conversion/*.py` and `src/models/*.cpp` for architectures with a similar config shape (layer count, head count, MoE expert count, norm placement) and suggest 1-2 candidates - but let the contributor confirm the choice rather than picking one yourself; this choice is a design decision they need to own.
Do not proceed to Step 1 until the contributor has answered these and named a reference architecture.
## Step 1 - Convert the model to GGUF
Follow HOWTO-add-model.md section 1 for the actual touch points (conversion class registration, `constants.py`, `tensor_mapping.py`, etc.) - don't re-derive them here, read them from the doc.
Skill-specific addition: for each touch point, show the contributor the equivalent code in the reference architecture they named in Step 0 before writing the new version, and check that they understand what's different about their model (e.g. non-standard tensor shapes, extra hparams) rather than just copying the pattern silently.
## Step 2 - Define the architecture in llama.cpp
Follow HOWTO-add-model.md section 2 for the actual touch points (`llm_arch` enum, `LLM_ARCH_NAMES`, hparam loading, RoPE type case, etc.), including its "Tips and tricks" section for `ggml_rope_ext` gotchas.
Skill-specific addition: never hack around RoPE with a custom sin/cos implementation - see the RoPE rule above.
## Step 3 - Build the GGML graph
Follow HOWTO-add-model.md section 3 for the actual touch points (`src/models/<name>.cpp` struct, `llama_model_mapping` registration, etc.).
Skill-specific addition: before writing `src/models/<name>.cpp`, read at least 10 other files under `src/models/` (pick a mix, not just the one reference architecture) to confirm the struct layout, naming, and style you're about to write actually matches current convention - the pattern drifts over time and the HOWTO doc can lag behind it.
## Step 4 - Optional: multimodal encoder
Only do this if the contributor flagged a vision/audio encoder in Step 0. Follow HOWTO-add-model.md section 4 and `docs/multimodal.md` for the actual touch points (`MmprojModel` subclass, `clip.cpp`, `mtmd.cpp`, encoder graph in `tools/mtmd/models`, etc.).
Skill-specific addition, and read this carefully: **whether the multimodal encoder can be bundled into the same PR as the base text-model support depends on how conventional the change is.** It's OK to bundle it if the encoder support is conventional - i.e. no new infra or logic is needed, it's just a new cgraph reusing existing preprocessing/projector machinery (e.g. siglip/pixtral/qwen with just a new projector). If it requires anything beyond that - a new preprocessor, non-standard projector logic, or changes to shared `libmtmd` infra/logic - STOP, tell the contributor this is non-conventional, and have them land the text model first with the encoder as a dedicated follow-up PR. Do not let this decision pass silently - call it out explicitly to the contributor before writing any `clip.cpp`/`mtmd.cpp` code.
## Step 5 - Optional: chat template / parsing support
Only do this if the model needs a new built-in chat template (`src/llama-chat.cpp`) or a new output parser (see `docs/development/parsing.md` and `docs/autoparser.md`). If either is needed beyond what a user-supplied Jinja template already covers, treat it as its own dedicated follow-up PR, not part of the base model-support PR - call this out explicitly to the contributor rather than silently bundling it in.
## Common pitfalls (from past PR reviews)
These recur often enough in review comments on past add-model PRs that they're worth checking proactively, not just waiting for a reviewer to catch them:
- Don't validate the same hparam/config assumption in both the Python conversion script and the C++ load path - pick one layer to own the check, duplicating it just adds maintenance surface.
- Optional hparams that are genuinely absent from some configs (e.g. a shared-expert count) should be read with an explicit optional/fallback accessor, not assumed present.
- Hparams that are actually load-bearing (the model produces wrong output or crashes without them, e.g. `sliding_window_pattern`, norm-eps) must hard-error if missing, not silently fall back to a default.
- Don't bake a default chat template into the C++ binary - inject it into the GGUF at conversion time instead, since one `llm_arch` can be reused by multiple fine-tunes with different templates, and a baked-in C++ default fails silently for those.
- Before writing a dedicated tool-call/output parser, check whether the existing autoparser already handles the template (`llama-debug-template-parser <jinja>` shows what it detects).
- Marking a custom EOS/closing-tag token as `eot` at conversion time isn't always sufficient - in long/agentic generations a model can emit the closing sequence as literal text instead of the token, so generation never stops on EOG and raw text leaks past the parser. Verify this case, not just the token path.
- If reusing or aliasing an existing pre-tokenizer for convenience, justify and test that choice explicitly - silent reuse is an easy source of subtle tokenizer bugs.
- Watch for excessive graph splits caused by building per-layer view/index tensors inside the layer loop - hoist tensors that don't vary per layer out of the loop (relevant if you hit `GGML_SCHED_MAX_SPLIT_INPUTS`).
- A custom KQ mask fed into flash attention must match FA's expected dtype - cast it to F16 before passing it to `build_attn_mha` when FA is enabled.
- When padding a custom KV-cache size to an alignment (e.g. `GGML_PAD(..., 256)`), apply the padding after all other size adjustments, not before - otherwise later logic can un-align it again.
- For non-standard cache/SWA (sliding-window-attention) semantics, override the dedicated hook (e.g. `llama_model_n_swa()`) rather than mutating hparams to fake the behavior - hparams may be read elsewhere for unrelated purposes.
- Don't ship unfinished or unverified speculative-decoding (e.g. MTP) scaffolding in the base model PR - if it hasn't actually been confirmed to work, pull it out and land it as its own follow-up.
- Conversion code should call into the base class's existing hparam logic (e.g. `super().set_gguf_parameters()`) rather than re-deriving it - large blocks of code that duplicate what `TextModel`/`MmprojModel` already provide will get flagged as redundant.
- Do constant tensor modifications (e.g. `norm(1 + weight)`) and permutations/chunking at conversion time, not in the graph - see HOWTO-add-model.md's "Prefer conversion-time tensor modifications" tip (Gemma 3 folds its `1 +` into the weights, Qwen3-Next permutes in `modify_tensors`). Doing these at runtime in the graph is very likely to be rejected as over-complicated; if you genuinely can't do it at conversion time, open a discussion first explaining why rather than implementing it in the graph.
## Validation checklist
Reference: `examples/model-conversion/README.md`.
1. Convert to GGUF, then inspect/run both the original and converted tensors.
2. Run logits verification (original vs converted). If this model is a new version of an already-supported family, verify the *previous* version still passes logits verification first - numerical differences may be pre-existing, not caused by the new work. The tools to perform full logits validation are available in `examples/model-conversion`.
3. Quantize (including QAT variants if relevant) and re-verify.
4. Run perplexity evaluation (simple and full).
5. Sanity-check across `tools/cli`, `tools/completion`, `tools/imatrix`, `tools/quantize`, and `tools/server`.
6. CPU backend first; other backends (CUDA, Metal, ...) can be separate follow-up PRs per `CONTRIBUTING.md`.
7. Re-review every changed file against the coding/naming guidelines in `AGENTS.md` (and `CONTRIBUTING.md`'s "Coding guidelines"/"Naming guidelines" sections) - this is a separate pass from functional testing and is just as important: no forced line-wrapping, no unicode punctuation, minimal/non-redundant comments, `snake_case` naming (`kebab-case` for file names), matching indentation/brace style, etc.
## Before opening a PR
- Run the `code-review` skill on the diff first - it catches the convention and scope issues reviewers flag most often, and it's recommended to do this locally before pushing the PR.
- Confirm the contributor can explain every changed line to a reviewer and is prepared to be asked about any of it - this is required regardless of how much of the code was AI-generated.
- Confirm they did a comprehensive manual review of the full diff, not just a skim.
- Fill in the AI-disclosure section of `.github/pull_request_template.md` describing how AI was used (do not omit or understate this).
- Do not write the PR description, commit message, GitHub issue/discussion text, or any reviewer replies yourself - the contributor writes these.
skills/code-review/SKILL.md
+134
View File
@@ -0,0 +1,134 @@
---
name: code-review
description: Review llama.cpp changes against project conventions and common reviewer pitfalls before a PR. Use when the user wants to review a diff, branch, or PR.
---
# Review llama.cpp changes
This skill reviews changes against llama.cpp's conventions and the pitfalls that reviewers flag most often, so the contributor can fix them before a maintainer has to. It has two modes:
- **Self-review (default):** review the contributor's own local changes (uncommitted work, or a branch vs `master`) as a pre-PR pass. Ask which if it's ambiguous; default to `git diff master...HEAD` plus any uncommitted changes.
- **Read-only review of a PR/file:** if the user points at a PR number or specific files (including code they didn't write), review those and report findings.
In both modes the output is **private review notes for the user to read and act on** - it is never something to post. This is a hard rule from `AGENTS.md`: an agent must NEVER write, or help write, a PR comment, a review comment, or a reply to a reviewer, by any means including `gh`. Do not offer to. If the user asks you to post the notes, refuse and point them at that rule. Present findings in the conversation only.
Before starting, read `AGENTS.md` and `CONTRIBUTING.md` if not already in context - the "Coding guidelines", "Naming guidelines", and AI usage sections are the baseline this review enforces. For a diff that adds a new model architecture, also read `docs/development/HOWTO-add-model.md` and consider the dedicated `add-new-model` skill.
## Step 0 - Scope the diff and pick the checklists
Identify what actually changed and which area checklists below apply. Run `git diff --stat` (or `gh pr view <n> --json files` for PR mode) and bucket the touched paths:
- `conversion/`, `gguf-py/`, `src/models/`, `src/llama-arch.*` -> **New model / architecture**
- `ggml/` (any backend, op, or `ggml.h`) -> **ggml / backend**
- `include/llama.h` and other public headers -> **Public API**
- `tools/server/` -> **Server**
- anything else, plus all of the above -> **General** (always runs)
Always run the **Scope and quick-reject gate**, the **Security review**, and the **General** checklist. Run each area checklist whose paths were touched. Additionally, if the diff introduces a new component, subsystem, or piece of infrastructure (a new file/class/module, a new abstraction, or hand-rolled machinery), run the **Approach and design** review. Tell the user which checklists you're running and why.
## Scope and quick-reject gate (always)
These are the patterns that get PRs closed without a full review. Check them first - a finding here is more important than any code nit, because it can mean the change shouldn't be a PR in its current form at all.
- Is there a prior issue/discussion for this? Features are supposed to start as an issue, not a PR (`CONTRIBUTING.md`). If this is a nontrivial feature with no linked issue, flag it and suggest opening one first.
- Is it a duplicate of existing/in-flight work? Suggest `gh search prs` / `gh search issues` for the feature. Many closed PRs were duplicates of something already queued.
- Is it self-contained and single-purpose? Multiple unrelated changes/optimizations bundled together get sent back to be split. Flag unrelated changes and suggest separate PRs.
- Does it touch multiple ggml backends at once? Initial support should be CPU-only, other backends as follow-ups (`CONTRIBUTING.md`). Flag CUDA/Metal/Vulkan/etc. changes bundled into a feature's first PR.
- Does it add a new `ggml_type` / quantization type? That carries a disproportionate maintenance burden and needs the full justification package (GGUF sample upload, perplexity vs FP16/BF16 and similar sizes, KL-divergence data, CPU perf numbers). Absent that, it will be rejected regardless of code quality.
- Is it invasive - new subsystem, core-API reshaping, changes to shared graph/sampler code that other models don't need? Flag it and suggest a discussion with maintainers before investing further.
- Is it niche/vendor-specific in a way that adds a maintenance burden nobody will own long-term? Flag the maintenance-ownership question.
- Is the change semantically correct, or a plausible-looking "fix" that misunderstands the code? Sanity-check the actual behavior, not just that it compiles.
- AI-disclosure: if AI meaningfully contributed, is the PR template's disclosure section filled in? Remind the user. Never suggest writing the PR description or commit message for them.
## Security review (mandatory)
Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF metadata, tensor shapes, tokenizer/grammar input, and all server/RPC fields are attacker-controlled - bound them before use.
- **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes.
- **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`).
- **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present.
- **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer.
- **Parsed/derived indices:** range-check `stoi`/`atoi` results and catch parse throws; never use a default or derived token id (EOS/BOS/...) as an index without a bounds check.
- **Reused/reserved buffers:** recheck bounds after a buffer is shrunk or reused; watch `reserve()` then index-by-assumed-size, and header fields read before their length is checked.
- **Server JSON ints:** clamp client-supplied integers (token/discard counts, offsets) to non-negative and an upper bound before they reach index/pointer arithmetic.
- **RPC-deserialized fields:** treat every field (type/buffer/data/ne/nb/op_params) as hostile - validate before use. Null/zero buffers skipping validation, attacker data pointers, out-of-range type indices, and negative strides sign-extending past a corner-only assert all give arbitrary read/write.
- **Lifetime/UAF:** flag stored raw pointers to caller/temporary storage, cached pointers to buffers a later free releases, async ops whose source may drop before completion, and structures not invalidated on free/realloc. Null-check conditionally-built or "not required" tensors before dereferencing.
## Approach and design (when a new component/infra is introduced)
Run this whenever the diff adds a new component, subsystem, or piece of infrastructure. Reviews too often stop at "does it work" - a diff can be correct and still be the wrong approach, and a messy design costs more long-term than a bug. Evaluate the *approach*, not just the behavior; raising a cleaner one is a high-value finding, not a nit. If you see a better design, describe it concretely rather than just calling the current one bad.
- **Simpler approach upstream:** the biggest win is often a different data model or design that removes whole subsystems, not tweaks to the code as written. Complexity must be justified by the problem, not by the first thing that worked.
- **Reuse over reinvention:** grep for an existing helper, library, object, or mechanism before adding a new one. Reimplementing what the codebase already has reintroduces solved bugs and adds maintenance surface.
- **Clear ownership/lifetime:** prefer RAII and obvious ownership over manual liveness flags, hand-tracked pointers, and "is it still alive?" checks - manual lifetime tracking is a recurring source of subtle bugs.
- **Right-sized machinery:** flag redundant, overkill, or heavier-than-needed primitives and abstractions; use the minimum the design actually needs.
- **Right structure and fit:** a new type should earn its place (split it if it serves two roles); follow existing patterns, idioms, and naming, and avoid constructs the project shuns.
- **Root cause vs symptom:** fixes layered on fixes signal a design to correct, not guard around.
## New model / architecture
See the `add-new-model` skill and `docs/development/HOWTO-add-model.md` for the full workflow; this is the review-time subset that reviewers most often catch:
- Don't branch on `model.arch` when the real dependency is a config/capability value - gate on the hparam/capability, not the architecture enum.
- If the model is a close variant of an existing arch, is the delta justified? Prefer reusing or subclassing the existing arch/model class over duplicating it. A near-duplicate class or `src/models/<name>.cpp` will be asked to merge with its sibling.
- New tensor names go through `tensor_mapping.py`, not ad-hoc name matching.
- For QKV, split the *activation* with `ggml_view`, not the *weight* tensor; rely on ggml broadcasting instead of manually duplicating tensors.
- New graph inputs are declared at the top of the graph-build function, not inline where first used.
- Hparams that the model can't run correctly without must be mandatory (hard-error if missing), not read with a silent default fallback. Only genuinely-optional-across-configs values get a fallback accessor.
- New/optional weight tensors (scales, etc.) must route through `build_lora_mm` and the existing helpers, matching convention - don't leave raw matmuls copied from another arch.
- Don't hack RoPE with a custom sin/cos implementation. If `ggml_rope_ext` genuinely can't express it, that's an issue for discussion, not a PR.
- Test the quantized-KV path (`-ctk`/`-ctv q8_0`), not just default f16 - new speculative/attention features silently break there.
- Preserve existing explanatory comments about model-specific quirks when copying code; note the provenance ("copied from X, with Y added").
- Remove dead code/branches left over from adapting a reference implementation.
## ggml / backend
- `supports_op` (and any dispatch/gating condition) must be scoped exactly to the cases being changed - a condition meant for a few quant types must not silently disable or enable everything else.
- No hardcoded warp/lane size - use `ggml_cuda_get_physical_warp_size()` (32 on CUDA, 64 on HIP/ROCm) and the portable helpers.
- Strip leftover debug/profiling/logging code before review.
- New or changed op? Update `docs/ops.md` and the relevant `docs/ops/*.csv` for the touched backend.
- New op or operator change needs corresponding `test-backend-ops` cases, and (per `CONTRIBUTING.md`) consistency across at least two backends.
- New kernels are expected to come with concrete perf data (throughput across realistic tensor shapes), not just correctness.
- Don't have a backend mutate the cgraph as a shortcut - that's an unresolved architectural question, not something to slip in.
- Expect this to need two maintainer approvals; that's normal for `ggml/` changes, not a sign something is wrong.
- For CUDA: Avoid excessively templating kernels, only add this where it shows visible performance gain.
## Public API (`include/llama.h`)
Public API changes carry a higher bar than internal ones (`CONTRIBUTING.md`). Review for:
- Justification: why doesn't an existing mechanism (e.g. `cb_eval`, existing batch/sampler knobs) suffice? If it does, the change likely shouldn't add public surface. This is the single most common reason these PRs are rejected.
- Experimental or stop-gap surface belongs in a side header (`llama-ext.h`), not in `llama.h`.
- Keep it minimal and general: prefer one general call over several narrow convenience wrappers; make new calls forward-compatible (e.g. mixed-modality batches) rather than assuming today's shape.
- The C API is the first-class, stable, ABI-defining surface - don't propose a parallel C++ API as a replacement. `llama-cpp.h` stays a thin convenience layer.
- Types and naming: sized integer types (`int32_t`, `size_t` for sizes/offsets); `snake_case`; `<class>_<method>` = `<class>_<action>_<noun>`; enum values upper-case and prefixed with the enum name; `_t` suffix for opaque types. Avoid gratuitous signature/ABI changes to existing exported functions.
- Every new API needs a working example/tool exercising it in the same PR - reviewers find real bugs by requiring it to be wired into `server`, `embedding`, `perplexity`, etc.
## Server (`tools/server/`)
- Is the feature within server's defined scope? Check `tools/server/README-dev.md` - out-of-scope features get declined.
- Security: don't trust client-supplied headers (e.g. `X-Forwarded-For`) or add footguns; things like IP allowlisting belong at a reverse proxy unless there's a trusted-proxy design.
- Wire new behavior into the existing request/response and checkpoint paths correctly; watch for resource leaks across requests.
## General (always)
Enforce the `AGENTS.md` / `CONTRIBUTING.md` coding and naming guidelines on every changed line - this is a distinct pass from checking that the code works, and matters just as much for review speed:
- ASCII only in code and comments - no emdash, unicode arrows, `x`, `...` used as unicode; use `-`, `->`, `x`, `...` ASCII equivalents.
- Comments are concise and explain non-obvious *why*, not *what*. Flag verbose comments, comments that restate the code, comments that reference the current task/PR, and comments hard-wrapped to a fixed column width.
- Do not force-wrap prose/comments to a fixed character count or split a sentence across lines.
- `snake_case` names; `kebab-case` (lowercase-with-dashes) file names for C/C++, `.h` headers; Python files lowercase-with-underscores. Naming optimizes for longest common prefix (`number_small`, not `small_number`).
- 4-space indentation, brackets on the same line, `void * ptr`, `int & a`, no trailing whitespace; match the surrounding style.
- Reuse existing infrastructure over introducing new components; no new third-party dependencies, extra headers, or files unless clearly justified.
- Keep it simple: a simpler change doing 90% is often preferable to a complex one doing 100%. Flag unnecessary templates/fancy STL; basic `for` loops are fine here.
- Every added line should be something the contributor can explain and defend to a reviewer without AI help - flag anything that looks copied-in without understanding.
## Reporting
Group findings by severity so the user knows what actually blocks a merge:
1. **Blocking** - quick-reject/scope issues and correctness bugs; these can sink the PR regardless of everything else.
2. **Will slow the review** - convention/naming/comment violations, missing tests/docs/perf data, missing API justification or example.
3. **Nits** - minor style, optional cleanups.
For each finding, point to the file and line and say concretely what to change and why. Do not rewrite the whole diff unprompted; let the contributor make the fixes so they own and understand them. And do not draft any PR text, commit message, or reviewer reply - that is the contributor's to write.
src/llama-arch.cpp
+11
View File
@@ -127,6 +127,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_GROVEMOE, "grovemoe" },
{ LLM_ARCH_APERTUS, "apertus" },
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
{ LLM_ARCH_MINIMAX_M3, "minimax-m3" },
{ LLM_ARCH_COGVLM, "cogvlm" },
{ LLM_ARCH_RND1, "rnd1" },
{ LLM_ARCH_PANGU_EMBED, "pangu-embedded" },
@@ -253,6 +254,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" },
{ LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" },
{ LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" },
{ LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" },
{ LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" },
{ LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" },
{ LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" },
{ LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" },
{ LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" },
@@ -596,6 +600,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" },
{ LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" },
{ LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" },
{ LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" },
{ LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" },
{ LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
@@ -831,6 +838,9 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
@@ -1000,6 +1010,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_LFM2:
case LLM_ARCH_LFM2MOE:
case LLM_ARCH_MINIMAX_M2:
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_KIMI_LINEAR:
return false;
src/llama-arch.h
+7
View File
@@ -146,6 +146,7 @@ enum llm_arch {
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
LLM_ARCH_MINIMAX_M3,
LLM_ARCH_DFLASH,
LLM_ARCH_UNKNOWN,
};
@@ -258,6 +259,9 @@ enum llm_kv {
LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
LLM_KV_ATTENTION_INDEXER_TOP_K,
LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,
LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,
LLM_KV_ATTENTION_INDEXER_TYPES,
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
@@ -596,6 +600,9 @@ enum llm_tensor {
LLM_TENSOR_INDEXER_PROJ,
LLM_TENSOR_INDEXER_ATTN_K,
LLM_TENSOR_INDEXER_ATTN_Q_B,
LLM_TENSOR_INDEXER_Q_PROJ,
LLM_TENSOR_INDEXER_K_PROJ,
LLM_TENSOR_INDEXER_Q_NORM,
LLM_TENSOR_INDEXER_COMPRESSOR_WKV,
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
src/llama-context.cpp
+76 -12
View File
@@ -728,6 +728,47 @@ void llama_context::synchronize() {
t_compute_start_us = 0;
}
void llama_context::release_device(bool evict_kv) {
if (!model.weights_resident() && !(evict_kv && !kv_device_evicted)) {
return;
}
// ensure no compute is in flight before freeing the device buffers
synchronize();
model.release_device_weights();
if (evict_kv && memory && !kv_device_evicted) {
memory->release_device_buffers();
kv_device_evicted = true;
// also free the scheduler and its worst-case compute buffer (hundreds of MiB) so a cold
// model holds essentially no VRAM. Rebuilt lazily by sched_reserve() on restore.
sched.reset();
sched_need_reserve = true;
// finally, free each backend's own compute-scratch (Vulkan prealloc/staging buffers, which
// are owned by the backend and survive sched.reset()). Reallocated lazily on next compute.
for (auto & backend : backends) {
ggml_backend_free_scratch(backend.get());
}
}
}
void llama_context::restore_device() {
if (model.weights_resident() && !kv_device_evicted && sched) {
return;
}
model.restore_device_weights();
if (kv_device_evicted && memory) {
memory->restore_device_buffers();
kv_device_evicted = false;
}
// rebuild the scheduler + compute buffer if they were freed on release (evict_kv mode)
if (!sched) {
sched_reserve();
}
// make sure all weight/KV uploads have completed before any compute reads them
if (sched) {
ggml_backend_sched_synchronize(sched.get());
}
}
const llama_model & llama_context::get_model() const {
return model;
}
@@ -1705,6 +1746,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
return -1;
}
// on-demand VRAM: if the weights were released while this model was idle, bring them back
// to the device before building the compute graph
if (!model.weights_resident()) {
restore_device();
}
const auto & vocab = model.vocab;
const auto & hparams = model.hparams;
@@ -2338,7 +2385,8 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
model.arch == LLM_ARCH_DEEPSEEK4) {
model.arch == LLM_ARCH_DEEPSEEK4 ||
model.arch == LLM_ARCH_MINIMAX_M3) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
@@ -3223,17 +3271,21 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
ret[buft].context += size;
}
}
if (model.hparams.no_alloc) {
for (size_t i = 0; i < backends.size(); ++i) {
ggml_backend_t backend = backends[i].get();
ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend);
ret[buft].compute += backend_buf_exp_size[i];
}
} else {
for (const auto & backend_ptr : backends) {
ggml_backend_t backend = backend_ptr.get();
ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend);
ret[buft].compute += ggml_backend_sched_get_buffer_size(sched.get(), backend);
// the scheduler (and its compute buffers) may have been freed while the model is cold
// (on-demand VRAM eviction, see release_device); it contributes no compute VRAM then.
if (sched) {
if (model.hparams.no_alloc) {
for (size_t i = 0; i < backends.size(); ++i) {
ggml_backend_t backend = backends[i].get();
ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend);
ret[buft].compute += backend_buf_exp_size[i];
}
} else {
for (const auto & backend_ptr : backends) {
ggml_backend_t backend = backend_ptr.get();
ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend);
ret[buft].compute += ggml_backend_sched_get_buffer_size(sched.get(), backend);
}
}
}
return ret;
@@ -3674,6 +3726,18 @@ void llama_synchronize(llama_context * ctx) {
ctx->synchronize();
}
void llama_context_release_device(llama_context * ctx, bool evict_kv) {
ctx->release_device(evict_kv);
}
void llama_context_restore_device(llama_context * ctx) {
ctx->restore_device();
}
bool llama_context_weights_resident(const llama_context * ctx) {
return ctx->get_model().weights_resident();
}
float * llama_get_logits(llama_context * ctx) {
ctx->synchronize();
src/llama-context.h
+11
View File
@@ -57,6 +57,14 @@ struct llama_context {
void synchronize();
// on-demand device (VRAM) residency: free / rebuild the model's GPU weight buffers so that
// VRAM can be time-shared with another model. release_device() synchronizes first; decode()
// auto-restores when it finds the weights have been released. If evict_kv is set, the KV cache
// device buffers are also evicted to a host shadow (for when the KV would not fit alongside the
// other model) - this adds a D2H/H2D copy of the live KV on each cycle.
void release_device(bool evict_kv = false);
void restore_device();
const llama_model & get_model() const;
const llama_cparams & get_cparams() const;
@@ -345,6 +353,9 @@ private:
bool sched_need_reserve = true;
// on-demand device residency: true while the KV cache device buffers have been evicted to host
bool kv_device_evicted = false;
ggml_backend_t backend_cpu = nullptr;
std::vector<ggml_backend_ptr> backends;
src/llama-grammar.cpp
+12
View File
@@ -1139,6 +1139,18 @@ struct llama_grammar * llama_grammar_init_impl(
vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
}
// Validate that all rule references point to valid rules
for (size_t i = 0; i < n_rules; i++) {
for (const auto & elem : vec_rules[i]) {
if (elem.type == LLAMA_GRETYPE_RULE_REF) {
if (elem.value >= n_rules || vec_rules[elem.value].empty()) {
LLAMA_LOG_ERROR("invalid grammar: rule %zu references undefined rule %u\n", i, elem.value);
return nullptr;
}
}
}
}
// Check for left recursion
std::vector<bool> rules_visited(n_rules);
std::vector<bool> rules_in_progress(n_rules);
src/llama-graph.cpp
+12 -1
View File
@@ -1709,6 +1709,17 @@ ggml_tensor * llm_graph_context::build_ffn(
cur = ggml_swiglu(ctx0, cur);
cb(cur, "ffn_swiglu", il);
} break;
case LLM_FFN_SWIGLU_OAI_MOE:
if (gate && type_gate == LLM_FFN_PAR) {
// same alpha/limit constants as gpt-oss
const float alpha = 1.702f;
const float limit = 7.0f;
cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit);
cb(cur, "ffn_swiglu_oai", il);
type_gate = LLM_FFN_SEQ;
} else {
GGML_ABORT("LLM_FFN_SWIGLU_OAI_MOE requires a parallel gate");
} break;
case LLM_FFN_GEGLU:
{
cur = ggml_geglu(ctx0, cur);
@@ -2668,7 +2679,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}
const auto & kq_mask = inp->get_kq_mask();
ggml_tensor * kq_mask = inp->get_kq_mask();
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
src/llama-hparams.cpp
+18
View File
@@ -180,6 +180,16 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const {
return val;
}
uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const {
if (!indexer_kv || indexer_head_size == 0) {
return 0; // arch without a MSA indexer
}
if (il < n_layer_dense_lead) {
return 0; // leading dense layers carry no indexer
}
return indexer_head_size; // 128
}
uint32_t llama_hparams::n_embd_r() const {
if (wkv_head_size != 0) {
// for RWKV models
@@ -248,6 +258,14 @@ bool llama_hparams::is_mla() const {
return n_embd_head_k_mla_impl != 0 && n_embd_head_v_mla_impl != 0;
}
bool llama_hparams::is_indexer_full(uint32_t il) const {
if (il < n_layer()) {
return is_indexer_full_impl[il];
}
GGML_ABORT("%s: il (%u) out of bounds (n_layer: %u)\n", __func__, il, n_layer());
}
uint32_t llama_hparams::n_embd_head_k_mla() const {
return is_mla() ? n_embd_head_k_mla_impl : n_embd_head_k();
}
src/llama-hparams.h
+14
View File
@@ -226,6 +226,15 @@ struct llama_hparams {
uint32_t indexer_n_head = 0;
uint32_t indexer_head_size = 0;
uint32_t indexer_top_k = 0;
// MSA
uint32_t indexer_block_size = 0;
uint32_t indexer_local_blocks = 0;
// MSA stores its indexer keys in the main KV cache (k_idx tensors);
bool indexer_kv = false;
// Indexer is "full" (1) or "shared" (0)
// Shared indexers reuse top-k from previous full layer
std::array<uint32_t, LLAMA_MAX_LAYERS> is_indexer_full_impl;
// DeepSeek-V4
uint32_t dsv4_o_group_count = 0;
@@ -302,6 +311,8 @@ struct llama_hparams {
bool is_swa(uint32_t il) const;
bool is_indexer_full(uint32_t il) const;
void set_recr_pattern(uint32_t n_pattern, bool dense_first = false);
// whether or not the given layer is recurrent (for hybrid models)
@@ -344,6 +355,9 @@ struct llama_hparams {
uint32_t n_embd_k_gqa_max() const;
uint32_t n_embd_v_gqa_max() const;
// dimension of the single-head MSA indexer key stream
uint32_t n_embd_k_idx(uint32_t il = 0) const;
// dimension of the rolling state embeddings
// corresponds to Mamba's conv_states size or RWKV's token_shift states size
uint32_t n_embd_r() const;
src/llama-kv-cache-iswa.cpp
+11
View File
@@ -272,6 +272,17 @@ void llama_kv_cache_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id,
kv_swa->state_read(io, seq_id, flags);
}
void llama_kv_cache_iswa::release_device_buffers() {
kv_base->release_device_buffers();
kv_swa->release_device_buffers();
}
bool llama_kv_cache_iswa::restore_device_buffers() {
bool ok = kv_base->restore_device_buffers();
ok = kv_swa->restore_device_buffers() && ok;
return ok;
}
llama_kv_cache * llama_kv_cache_iswa::get_base() const {
return kv_base.get();
}
src/llama-kv-cache-iswa.h
+3
View File
@@ -83,6 +83,9 @@ public:
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
void release_device_buffers() override;
bool restore_device_buffers() override;
//
// llama_kv_cache_iswa specific API
//
src/llama-kv-cache.cpp
+368 -18
View File
@@ -112,7 +112,7 @@ llama_kv_cache::llama_kv_cache(
auto it = ctx_map.find(buft);
if (it == ctx_map.end()) {
ggml_init_params params = {
/*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
/*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), //Reserve tensor metadata for up to 3 tensors per layer (K, V, and optional K_idx), plus one view per tensor per stream.
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
@@ -242,9 +242,25 @@ llama_kv_cache::llama_kv_cache(
v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr);
}
const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il);
ggml_tensor * k_idx = n_embd_k_idx > 0
? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream)
: nullptr;
if (k_idx) {
ggml_format_name(k_idx, "cache_k_idx_l%d", il);
msa_strict_slots = (n_stream == n_seq_max);
}
std::vector<ggml_tensor *> k_idx_stream;
for (uint32_t s = 0; s < n_stream; ++s) {
k_idx_stream.push_back(k_idx
? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2])
: nullptr);
}
map_layer_ids[il] = layers.size();
layers.push_back({ il, k, v, k_stream, v_stream, });
layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream });
}
if (reuse) {
@@ -293,13 +309,24 @@ llama_kv_cache::llama_kv_cache(
}
{
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
const size_t memory_size_k_idx = size_k_idx_bytes();
const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx;
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
constexpr float mib = 1024.0f * 1024.0f;
const std::string k_log = format(", K (%s): %7.2f MiB", ggml_type_name(type_k), (float) memory_size_k / mib);
const std::string v_log = format(", V (%s): %7.2f MiB", ggml_type_name(type_v), (float) memory_size_v / mib);
std::string k_idx_log;
if (memory_size_k_idx > 0) {
k_idx_log = format(", K_idx (%s): %7.2f MiB", ggml_type_name(GGML_TYPE_F32), (float) memory_size_k_idx / mib);
}
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs)%s%s%s\n", __func__,
(float) memory_size_total / mib, kv_size, (int) layers.size(), n_seq_max, n_stream,
k_log.c_str(), v_log.c_str(), k_idx_log.c_str());
}
// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
@@ -323,7 +350,7 @@ llama_kv_cache::llama_kv_cache(
hparams.n_embd_head_k() % 64 == 0;
// always create Hadamard rotation tensors for DeepSeek lightning indexers
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4) &&
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || model.arch == LLM_ARCH_GLM_DSA) &&
hparams.n_embd_head_k_full == hparams.indexer_head_size) {
attn_rot_k = true;
}
@@ -371,7 +398,9 @@ void llama_kv_cache::clear(bool data) {
if (data) {
for (auto & [_, buf] : ctxs_bufs) {
ggml_backend_buffer_clear(buf.get(), 0);
if (buf) { // may be null if evicted for on-demand VRAM sharing
ggml_backend_buffer_clear(buf.get(), 0);
}
}
}
}
@@ -392,6 +421,39 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
p1 = std::numeric_limits<llama_pos>::max();
}
// empty range - nothing to remove
if (p0 >= p1) {
return true;
}
// MSA anchors block selection to absolute cache slots (slot == position). Tail trim and full removal preserve this invariant, but removing a prefix
// or middle range would free slots while later cells survive, desynchronizing the indexer cache. Reject such removals before modifying the cache.
if (msa_strict_slots) {
for (llama_seq_id sid = 0; sid < (llama_seq_id) seq_to_stream.size(); ++sid) {
if (seq_id >= 0 && sid != seq_id) {
continue;
}
const auto & cells = v_cells[seq_to_stream[sid]];
const llama_pos pmin = cells.seq_pos_min(sid);
const llama_pos pmax = cells.seq_pos_max(sid);
if (pmin < 0) {
continue; // empty sequence
}
const bool overlaps = p0 <= pmax && p1 > pmin; // the range removes something
const bool leaves_tail = p1 <= pmax; // cells beyond the range survive
if (overlaps && leaves_tail) {
LLAMA_LOG_WARN("%s: MSA: partial (non-suffix) removal [%d, %d) for seq %d is not supported "
"(block selection is anchored to cache slots) - rejected\n", __func__, p0, p1, sid);
return false;
}
}
}
if (seq_id >= 0) {
auto & cells = v_cells[seq_to_stream[seq_id]];
auto & head = v_heads[seq_to_stream[seq_id]];
@@ -681,6 +743,9 @@ llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> ret;
for (const auto & [ctx, buf] : ctxs_bufs) {
if (!buf) { // may be null if evicted for on-demand VRAM sharing
continue;
}
ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
if (hparams.no_alloc) {
@@ -846,6 +911,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
if (layer.v_stream[ssrc]) {
ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);
}
if (layer.k_idx_stream[ssrc]) {
GGML_ASSERT(layer.k_idx_stream[sdst]);
ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]);
}
}
}
}
@@ -994,6 +1063,44 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
const auto & cells = v_cells[seq_to_stream[seq_id]];
if (n_tokens > cells.size()) {
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
return { };
}
// MSA block selection assumes slot == logical position (append-only streams).
if (msa_strict_slots) {
for (uint32_t ii = 0; ii < n_tokens; ++ii) {
const llama_pos pos = ubatch.pos[s*n_tokens + ii];
if (pos < 0 || (uint64_t) pos >= cells.size()) {
LLAMA_LOG_WARN("%s: MSA: position %d is outside the cache range [0, %u)\n",
__func__, pos, cells.size());
return { };
}
const uint32_t idx = (uint32_t) pos;
if (!cells.is_empty(idx)) {
LLAMA_LOG_WARN("%s: MSA: required slot %u is already occupied (stream %u)\n",
__func__, idx, seq_to_stream[seq_id]);
return { };
}
// strictly increasing positions, rules out duplicates and, for contiguous requests, is tightened to exact adjacency
if (!res.idxs[s].empty() && (cont ? idx != res.idxs[s].back() + 1
: idx <= res.idxs[s].back())) {
LLAMA_LOG_WARN("%s: MSA: token positions are not %s within the ubatch\n",
__func__, cont ? "contiguous" : "strictly increasing");
return { };
}
res.idxs[s].push_back(idx);
}
continue;
}
uint32_t head_cur = v_heads[seq_to_stream[seq_id]];
// if we have enough unused cells before the current head ->
@@ -1002,11 +1109,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
head_cur = 0;
}
if (n_tokens > cells.size()) {
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
return { };
}
uint32_t n_tested = 0;
// for continuous slots, we test that all tokens in the ubatch fit, starting from the current head
@@ -1113,6 +1215,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
const auto idx = sinfo.idxs[s][ii];
if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) {
LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: "
"writing pos %d into cell %u (stream %u). The indexer cache "
"would desync and block selection would silently corrupt. "
"This is a bug, please report it with reproduction steps.\n",
__func__, ubatch.pos[i], idx, sinfo.strm[s]);
GGML_ABORT("MSA: slot != pos");
}
if (!cells.is_empty(idx)) {
assert(cells.seq_count(idx) == 1);
@@ -1156,7 +1267,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",
__func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
// under MSA strict slots this path should be unreachable, since strict MSA placement never selects occupied cells
GGML_ASSERT(seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1));
}
}
@@ -1176,6 +1288,12 @@ bool llama_kv_cache::get_can_shift() const {
if (hparams.n_pos_per_embd() > 1) {
return false;
}
// shifting would leave k_idx stale
for (const auto & layer : layers) {
if (layer.k_idx) {
return false;
}
}
return true;
}
@@ -1292,6 +1410,23 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
}
ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
const int32_t ikv = map_layer_ids.at(il);
auto * k_idx = layers[ikv].k_idx;
GGML_ASSERT(k_idx);
const uint64_t kv_size = get_size();
const int64_t n_idx = k_idx->ne[0]; // 128
const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
return ggml_view_4d(ctx, k_idx,
n_idx, 1, n_kv, ns,
ggml_row_size(k_idx->type, n_idx), // nb1 (single head)
ggml_row_size(k_idx->type, n_idx), // nb2 (per cell)
ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream)
ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0);
}
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
@@ -1393,6 +1528,28 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama
return k_idxs;
}
ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
const int32_t ikv = map_layer_ids.at(il);
ggml_tensor * k_idx = layers[ikv].k_idx;
GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache");
const int64_t n_embd_head = k_idx_cur->ne[0]; // 128
const int64_t n_head = k_idx_cur->ne[1]; // 1
const int64_t n_tokens = k_idx_cur->ne[2];
const int64_t n_embd_gqa = n_embd_head*n_head; // 128
GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]);
k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0);
const int64_t n_stream = k_idx->ne[2];
if (n_stream > 1) {
const int64_t kv_size = get_size();
k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream);
}
return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store
}
ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
const uint32_t n_tokens = ubatch.n_tokens;
@@ -1801,12 +1958,88 @@ size_t llama_kv_cache::total_size() const {
size_t size = 0;
for (const auto & [_, buf] : ctxs_bufs) {
size += ggml_backend_buffer_get_size(buf.get());
if (buf) { // may be null if evicted for on-demand VRAM sharing
size += ggml_backend_buffer_get_size(buf.get());
}
}
return size;
}
void llama_kv_cache::release_device_buffers() {
// NOTE: the caller must have synchronized the backend so nothing references these buffers.
// The KV cache is read-write, so its host shadow is (re)captured fresh on every release.
if (dev_released) {
return;
}
dev_shadows.assign(ctxs_bufs.size(), device_buffer_shadow{});
size_t freed = 0;
for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
ggml_context * ctx = ctxs_bufs[i].first.get();
ggml_backend_buffer_t buf = ctxs_bufs[i].second.get();
if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) {
continue; // only real device (VRAM) buffers are evictable
}
auto & sh = dev_shadows[i];
sh.releasable = true;
sh.buft = ggml_backend_buffer_get_type(buf);
// capture live contents compactly in stable iteration order (skip views, which alias a base)
size_t total = 0;
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
if (t->view_src == nullptr) { total += ggml_nbytes(t); }
}
sh.data.resize(total);
size_t off = 0;
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
if (t->view_src != nullptr) { continue; }
const size_t n = ggml_nbytes(t);
ggml_backend_tensor_get(t, sh.data.data() + off, 0, n);
off += n;
}
freed += ggml_backend_buffer_get_size(buf);
ctxs_bufs[i].second.reset(); // free the device buffer
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
t->buffer = nullptr;
t->data = nullptr;
}
}
dev_released = true;
if (freed > 0) {
LLAMA_LOG_INFO("%s: released %.2f MiB of KV cache from device\n", __func__, freed / 1024.0 / 1024.0);
}
}
bool llama_kv_cache::restore_device_buffers() {
if (!dev_released) {
return true;
}
for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
auto & sh = dev_shadows[i];
if (!sh.releasable) {
continue;
}
ggml_context * ctx = ctxs_bufs[i].first.get();
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, sh.buft);
if (buf == nullptr) {
LLAMA_LOG_ERROR("%s: failed to reallocate KV cache device buffer (out of VRAM?)\n", __func__);
return false;
}
size_t off = 0;
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
if (t->view_src != nullptr) { continue; }
const size_t n = ggml_nbytes(t);
ggml_backend_tensor_set(t, sh.data.data() + off, 0, n);
off += n;
}
ctxs_bufs[i].second.reset(buf);
}
dev_released = false;
dev_shadows.clear(); // recaptured on next release
return true;
}
size_t llama_kv_cache::size_k_bytes() const {
size_t size_k_bytes = 0;
@@ -1827,6 +2060,18 @@ size_t llama_kv_cache::size_v_bytes() const {
return size_v_bytes;
}
size_t llama_kv_cache::size_k_idx_bytes() const {
size_t size_k_idx_bytes = 0;
for (const auto & layer : layers) {
if (layer.k_idx) {
size_k_idx_bytes += ggml_nbytes(layer.k_idx);
}
}
return size_k_idx_bytes;
}
ggml_tensor * llama_kv_cache::build_rope_shift(
const llama_cparams & cparams,
ggml_context * ctx,
@@ -2054,7 +2299,12 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
bool res = true;
res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id);
res = res && state_read_data(io, strm, cell_count, sinfo);
try {
res = res && state_read_data(io, strm, cell_count, sinfo);
} catch (...) {
res = false;
}
if (!res) {
if (seq_id == -1) {
@@ -2134,6 +2384,36 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
}
}
if (size_k_idx_bytes() > 0) {
const uint32_t has_k_idx_u32 = 1;
io.write(&has_k_idx_u32, sizeof(has_k_idx_u32));
for (const auto & layer : layers) {
const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0;
io.write(&layer_has_k_idx, sizeof(layer_has_k_idx));
if (!layer_has_k_idx) {
continue;
}
GGML_ASSERT(layer.k_idx_stream[cr.strm]);
const int32_t k_idx_type_i = (int32_t) layer.k_idx->type;
io.write(&k_idx_type_i, sizeof(k_idx_type_i));
const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
io.write(&k_idx_size_row, sizeof(k_idx_size_row));
for (const auto & range : cr.data) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * k_idx_size_row;
const size_t offset = range.first * k_idx_size_row;
io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size);
}
}
}
if (!v_trans) {
for (const auto & layer : layers) {
const uint32_t il = layer.il;
@@ -2382,6 +2662,68 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
}
}
if (size_k_idx_bytes() > 0) {
uint32_t has_k_idx_u32 = 0;
io.read(&has_k_idx_u32, sizeof(has_k_idx_u32));
if (has_k_idx_u32 != 1) {
LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__);
return false;
}
for (const auto & layer : layers) {
uint32_t layer_has_k_idx = 0;
io.read(&layer_has_k_idx, sizeof(layer_has_k_idx));
const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0;
if (layer_has_k_idx != expected_layer_has_k_idx) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx state for layer: got %u, expected %u\n",
__func__, layer_has_k_idx, expected_layer_has_k_idx);
return false;
}
if (!layer_has_k_idx) {
continue;
}
GGML_ASSERT(layer.k_idx_stream[strm]);
int32_t k_idx_type_i = -1;
io.read(&k_idx_type_i, sizeof(k_idx_type_i));
if (k_idx_type_i != (int32_t) layer.k_idx->type) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx type: got %d, expected %d\n",
__func__, k_idx_type_i, (int32_t) layer.k_idx->type);
return false;
}
uint64_t k_idx_size_row = 0;
io.read(&k_idx_size_row, sizeof(k_idx_size_row));
const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
if (k_idx_size_row != expected_k_idx_size_row) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx row size: got %zu, expected %zu\n",
__func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row);
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row);
} else {
for (uint32_t i = 0; i < cell_count; ++i) {
io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row);
}
}
}
}
}
if (!this->v_trans) {
for (const auto & layer : layers) {
const uint32_t il = layer.il;
@@ -2583,6 +2925,10 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const {
return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
}
@@ -2591,6 +2937,10 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_
return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
return kv->build_input_k_idxs(ctx, ubatch);
}
src/llama-kv-cache.h
+24
View File
@@ -149,6 +149,10 @@ public:
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
// on-demand device (VRAM) residency (see llama_memory_i)
void release_device_buffers() override;
bool restore_device_buffers() override;
//
// llama_kv_cache specific API
//
@@ -173,10 +177,12 @@ public:
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
// store k_cur and v_cur in the cache based on the provided head location
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
//
// preparation API
@@ -228,9 +234,11 @@ private:
ggml_tensor * k;
ggml_tensor * v;
ggml_tensor * k_idx; // MSA single-head indexer keys, F32
std::vector<ggml_tensor *> k_stream;
std::vector<ggml_tensor *> v_stream;
std::vector<ggml_tensor *> k_idx_stream;
};
bool v_trans = true; // the value tensor is transposed
@@ -259,12 +267,25 @@ private:
// env: LLAMA_KV_CACHE_DEBUG
int debug = 0;
// set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq)
bool msa_strict_slots = false;
// this is the SWA type of the cache - not to be confused with the model SWA type
const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
// ggml contexts for the KV cache along with the allocated backend buffers:
std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
// on-demand device eviction: host shadow of each device buffer's live contents (re-captured on
// every release since the KV is read-write), plus the buffer type needed to reallocate it.
struct device_buffer_shadow {
ggml_backend_buffer_type_t buft = nullptr;
bool releasable = false; // device (VRAM) buffer that was freed
std::vector<uint8_t> data; // captured tensor bytes (compact, iteration order)
};
std::vector<device_buffer_shadow> dev_shadows; // parallel to ctxs_bufs
bool dev_released = false;
// the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot())
// note: this is not part of the KV state and it's only used to speed-up the find_slot() method
std::vector<uint32_t> v_heads;
@@ -291,6 +312,7 @@ private:
size_t size_k_bytes() const;
size_t size_v_bytes() const;
size_t size_k_idx_bytes() const;
ggml_tensor * build_rope_shift(
const llama_cparams & cparams,
@@ -370,6 +392,7 @@ public:
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const;
// store k_cur and v_cur in the cache based on the provided head location
// note: the heads in k_cur and v_cur should be laid out contiguously in memory
@@ -379,6 +402,7 @@ public:
// - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const;
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const;
// create destination indices for each head of the current batch for where it would be written in the KV cache
// the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but
src/llama-memory-hybrid.cpp
+12
View File
@@ -201,6 +201,18 @@ void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id,
mem_recr->state_read(io, seq_id, flags);
}
void llama_memory_hybrid::release_device_buffers() {
// evict both the attention KV (grows with context) and the recurrent/SSM state
mem_attn->release_device_buffers();
mem_recr->release_device_buffers();
}
bool llama_memory_hybrid::restore_device_buffers() {
bool ok = mem_attn->restore_device_buffers();
ok = mem_recr->restore_device_buffers() && ok;
return ok;
}
llama_kv_cache * llama_memory_hybrid::get_mem_attn() const {
return mem_attn.get();
}
src/llama-memory-hybrid.h
+3
View File
@@ -76,6 +76,9 @@ public:
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
void release_device_buffers() override;
bool restore_device_buffers() override;
//
// llama_memory_hybrid specific API
//
src/llama-memory-recurrent.cpp
+86 -3
View File
@@ -140,7 +140,9 @@ void llama_memory_recurrent::clear(bool data) {
if (data) {
for (auto & [_, buf] : ctxs_bufs) {
ggml_backend_buffer_clear(buf.get(), 0);
if (buf) { // may be null if evicted for on-demand VRAM sharing
ggml_backend_buffer_clear(buf.get(), 0);
}
}
}
@@ -399,6 +401,7 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> ret;
for (const auto & [_, buf] : ctxs_bufs) {
if (!buf) { continue; } // may be null if evicted for on-demand VRAM sharing
ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
}
return ret;
@@ -700,12 +703,87 @@ bool llama_memory_recurrent::get_can_shift() const {
size_t llama_memory_recurrent::total_size() const {
size_t size = 0;
for (const auto & [_, buf] : ctxs_bufs) {
size += ggml_backend_buffer_get_size(buf.get());
if (buf) { // may be null if evicted for on-demand VRAM sharing
size += ggml_backend_buffer_get_size(buf.get());
}
}
return size;
}
void llama_memory_recurrent::release_device_buffers() {
// Same mechanism as llama_kv_cache: the recurrent (SSM/conv) state is read-write, so its host
// shadow is (re)captured on every release. The caller must have synchronized the backend.
if (dev_released) {
return;
}
dev_shadows.assign(ctxs_bufs.size(), device_buffer_shadow{});
size_t freed = 0;
for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
ggml_context * ctx = ctxs_bufs[i].first.get();
ggml_backend_buffer_t buf = ctxs_bufs[i].second.get();
if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) {
continue;
}
auto & sh = dev_shadows[i];
sh.releasable = true;
sh.buft = ggml_backend_buffer_get_type(buf);
size_t total = 0;
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
if (t->view_src == nullptr) { total += ggml_nbytes(t); }
}
sh.data.resize(total);
size_t off = 0;
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
if (t->view_src != nullptr) { continue; }
const size_t n = ggml_nbytes(t);
ggml_backend_tensor_get(t, sh.data.data() + off, 0, n);
off += n;
}
freed += ggml_backend_buffer_get_size(buf);
ctxs_bufs[i].second.reset();
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
t->buffer = nullptr;
t->data = nullptr;
}
}
dev_released = true;
if (freed > 0) {
LLAMA_LOG_INFO("%s: released %.2f MiB of recurrent state from device\n", __func__, freed / 1024.0 / 1024.0);
}
}
bool llama_memory_recurrent::restore_device_buffers() {
if (!dev_released) {
return true;
}
for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
auto & sh = dev_shadows[i];
if (!sh.releasable) {
continue;
}
ggml_context * ctx = ctxs_bufs[i].first.get();
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, sh.buft);
if (buf == nullptr) {
LLAMA_LOG_ERROR("%s: failed to reallocate recurrent device buffer (out of VRAM?)\n", __func__);
return false;
}
size_t off = 0;
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
if (t->view_src != nullptr) { continue; }
const size_t n = ggml_nbytes(t);
ggml_backend_tensor_set(t, sh.data.data() + off, 0, n);
off += n;
}
ctxs_bufs[i].second.reset(buf);
}
dev_released = false;
dev_shadows.clear();
return true;
}
size_t llama_memory_recurrent::size_r_bytes() const {
size_t size_r_bytes = 0;
@@ -819,7 +897,12 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
bool res = true;
res = res && state_read_meta(io, cell_count, seq_id);
res = res && state_read_data(io, cell_count);
try {
res = res && state_read_data(io, cell_count);
} catch (...) {
res = false;
}
if (!res) {
if (seq_id == -1) {
src/llama-memory-recurrent.h
+14
View File
@@ -66,6 +66,10 @@ public:
void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
// on-demand device (VRAM) residency (see llama_memory_i)
void release_device_buffers() override;
bool restore_device_buffers() override;
uint32_t head = 0; // the location where the batch will be placed in the cache (see find_slot())
uint32_t size = 0; // total number of cells, shared across all sequences
uint32_t used = 0; // used cells (i.e. at least one seq_id)
@@ -121,6 +125,16 @@ private:
// ggml contexts for the KV cache along with the allocated backend buffers:
std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
// on-demand device eviction (see release_device_buffers): host shadow of each device buffer's
// live contents (recaptured on every release since the recurrent state is read-write)
struct device_buffer_shadow {
ggml_backend_buffer_type_t buft = nullptr;
bool releasable = false;
std::vector<uint8_t> data;
};
std::vector<device_buffer_shadow> dev_shadows; // parallel to ctxs_bufs
bool dev_released = false;
size_t total_size() const;
size_t size_r_bytes() const;
src/llama-memory.h
+10
View File
@@ -124,6 +124,16 @@ struct llama_memory_i {
virtual void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const = 0;
virtual void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) = 0;
//
// on-demand device (VRAM) residency
//
// Free this memory's device (VRAM) buffers to a host shadow and rebuild them on demand, to
// hand VRAM to another model while keeping the cached data (no re-prefill). The caller must
// have synchronized the backend first. Default: no-op (the memory stays resident).
virtual void release_device_buffers() {}
virtual bool restore_device_buffers() { return true; }
};
using llama_memory_ptr = std::unique_ptr<llama_memory_i>;
src/llama-mmap.cpp
+35 -1
View File
@@ -618,13 +618,47 @@ struct llama_mmap::impl {
};
llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
llama_mmap::~llama_mmap() = default;
llama_mmap::~llama_mmap() {
// unpin before the pages are unmapped by the impl destructor
if (host_reg_addr && host_unreg_fn) {
host_unreg_fn(host_reg_addr);
}
}
size_t llama_mmap::size() const { return pimpl->size; }
void * llama_mmap::addr() const { return pimpl->addr; }
void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
size_t llama_mmap::register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *)) {
#ifdef _POSIX_MAPPED_FILES
if (host_reg_addr || !reg_fn || !unreg_fn || last <= first) {
return 0;
}
// expand outward to the page boundaries retained by unmap_fragment
const size_t page_size = sysconf(_SC_PAGESIZE);
first = first & ~(page_size - 1);
last = (last + page_size - 1) & ~(page_size - 1);
void * reg_addr = (uint8_t *) pimpl->addr + first;
if (!reg_fn(reg_addr, last - first)) {
return 0;
}
host_reg_addr = reg_addr;
host_unreg_fn = unreg_fn;
return last - first;
#else
GGML_UNUSED(first);
GGML_UNUSED(last);
GGML_UNUSED(reg_fn);
GGML_UNUSED(unreg_fn);
return 0;
#endif
}
#if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)
const bool llama_mmap::SUPPORTED = true;
#else

Some files were not shown because too many files have changed in this diff Show More