Compare commits

..

11 Commits

Author SHA1 Message Date
Ruixiang Wang 7a20b417f4 model: add MTP support for Nemotron model (#26725)
* model: add MTP support for Nemotron Nano model

* model: add mtp_flags for nemotron model

* address review comments
2026-08-10 11:25:24 +03:00
Alessandro de Oliveira Faria (A.K.A.CABELO) e23e9440eb vendor : update cpp-httplib to 0.53.0 (#26821) 2026-08-10 09:57:45 +02:00
Bar Haim 157b81fe6d model : Granite-Switch Architecture (#25107)
* granite-switch: add llama.cpp backend (POC, CPU)

New "granite-switch" architecture: a dense, all-attention Granite-4.1
model with N embedded LoRA adapters selected per-token by control tokens.

- gguf-py schema (arch, KV keys, stacked LoRA tensor names) + writer helpers
- conversion/granite.py: GraniteSwitchModel converter (stacks N adapters +
  zero base slot into per-projection A/B tensors; emits switch metadata)
- C++ arch registration (llama-arch.{h,cpp}, llama-model.{h,cpp})
- src/models/granite_switch.cpp: load + per-token switched-LoRA graph via
  ggml_mul_mat_id over stacked tensors; sticky per-token index + control-token
  substitution in llm_graph_input_switch::set_input
- llm_graph_input_switch in src/models/models.h

Runs end-to-end on CPU: convert 3b checkpoint (842 tensors, stacked dim 13)
and generate on both base and control-token paths. Sticky switch state is
single-sequence (POC); full multi-sequence machinery is a follow-up.

* granite-switch: add Mac (Metal) build + mid-sequence switch demo script

Self-contained script to build llama.cpp on Apple Silicon (Metal),
convert the composed 3b checkpoint, and run the crisp mid-sequence
adapter-switch demos verified on Vela:
  - answerability: <|answerability|> mid-seq -> "unanswerable"
  - query_rewrite: <|query_rewrite|> mid-seq -> {"rewritten_question": ...}
Each demo runs the same prompt twice, differing only by a control token
placed before the assistant turn, so the per-token switch is visible.

* granite-switch mac demo: add -no-cnv so each run is one-shot

The composed model ships a chat template, so llama-completion auto-enables
interactive conversation mode and halts at a `>` prompt after generating,
stalling the script. -no-cnv disables conversation mode: generate once from
the raw prompt and exit (also prints special tokens, making the switch visible).

* granite-switch: replace global sticky index with in-graph router attention

The POC computed the per-token adapter index on the CPU and carried it
across ubatches in ONE global `mutable int32_t poc_sticky_index`, reset
only when a ubatch contained sequence position 0. That global had two
problems:

  1. Concurrency: with multiple sequences in a batch it was last-writer-
     wins — one sequence's adapter leaked into the others.
  2. Multi-turn: an interactive `ollama run` chat continues one KV cache,
     so turn 2 never saw position 0 and the index never reset — the
     adapter stayed stuck on across turns.

Port the vLLM/HF backend mechanism faithfully: a single-head causal
"router" attention recovers the adapter index in-graph. Per token, only
dim 0 carries signal — Q[0]=1, K[0]=+gain for a control token / -gain
otherwise, V[0]=adapter slot / 0 — and the causal softmax over the single
visible control token recovers that adapter's slot (readback =
clamp(round(V[0]), 0, n_adapters)). gain=15 matches config.py and is
F16-safe (no F32 cache).

The router's K/V live in the model KV cache at an extra layer
R == hparams.router_layer (== n_layer). We bump n_layer_all to n_real+1
so the cache allocator gives the router its own per-sequence slot, and
set n_layer_nextn=1 so n_layer() stays n_real — the decoder loop and
tensor loading are untouched and never reference layer R. The router K is
exempted from the k-shift RoPE loop (its dim-0 value is a literal
magnitude, not a rotation).

Because the selection now lives in the per-sequence KV cache, CONCURRENT
requests are isolated for free (problem 1 fixed; verified by
scratch/concurrent_switch_test.cpp). set_input becomes stateless pure
per-token maps; the global is gone.

Single-switch contract / known limitation, identical to vLLM & HF: the
gain is flat (no recency), so within one sequence there is no mechanism to
revert to base mid-sequence — once an adapter fires it stays on until that
sequence ends (problem 2 is therefore NOT fixed by a faithful copy; vLLM/HF
avoid it only because each served request is a fresh sequence). A client
continuing one KV cache across turns must start a fresh sequence per turn,
or opt into a recency-biased router (a deliberate divergence, not done
here). Documented in granite_switch.cpp and asserted by
scratch/multiturn_leak_test.cpp.

Verified (CPU): both demos unchanged (answerability -> "unanswerable",
query_rewrite -> rewritten query); concurrent two-sequence isolation
passes; multi-turn carry-over matches the vLLM/HF contract.

* granite-switch: drop scratch tests and mac demo for upstream PR

Remove the local-only development artifacts that should not ship in the
upstream PR:
  - granite-switch-mac-demo.sh (local Metal build + demo driver)
  - scratch/concurrent_switch_test.cpp
  - scratch/multiturn_leak_test.cpp

Also drop the now-dangling reference to the scratch tests from the
granite_switch.cpp header comment. Leaves only the core architecture
support (conversion, gguf constants, llama-arch/model/kv-cache, and the
granite_switch graph).

* granite-switch: trim comments to match native llama.cpp style

* granite-switch: trim conversion comments to match native style

* granite-switch: drop unused adapter_ranks metadata

* granite-switch: rename arch to graniteswitch and drop obid alias

* granite-switch: fix non-ASCII comments and document router gain assumption

* granite-switch: drop section comments from constants.py to match native style

* granite-switch: add functional tensor block comments matching Granite4 Vision style

* granite-switch: clarify n_expert_used comment

State the actual constraint: mul_mat_id needs n_expert_used == 1, and
since the GGUF carries expert_count = 0 the generic loader's
n_expert == 0 => n_expert_used == 0 assertion has already passed by the
time load_arch_hparams runs, so it is forced to 1 here.

* granite-switch: note n_layer_nextn reuse has no MTP

The router carving reuses n_layer_nextn, normally the MTP/next-token
count. Clarify in the comment that it is borrowed here purely as the
trailing-layers lever and that there is no MTP head, to spare readers
the double-take.

* granite-switch: rename source file and apply review nits

* granite-switch: don't force LoRA tensors to F16, follow --outtype instead

* granite-switch: drop redundant _permute_qk wrapper, call LlamaModel.permute directly

* granite-switch: read router gain from GGUF (control_token_gain) instead of hardcoding 15.0

* granite-switch: derive n_slots()

* granite-switch: move llm_graph_input_switch into granite-switch.cpp

* granite-switch: cut AI-style narration comments

* granite-switch: collapse multi-line comments

* granite-switch: rename control_token_* maps to adapter_token_*

* granite-switch: cut noise comments

* granite-switch: rename embedded LoRA tensors to <base>.lora_a/lora_b

* granite-switch: GGML_ASSERT token input to avoid UB on embeddings

* granite-switch: TODO for raw embedding input support

* granite-switch: collapse LoRA tensor constants to .lora_a/.lora_b suffix

* granite-switch: drop n_expert_used hack, guard mul_mat_id buft probe

* granite-switch: stop forcing dense expert counts, read from config

* granite-switch: renamed control_token_gain metadata key to router_gain

* granite-switch: trim header comments to match native style

* granite-switch: collapse LoRA tensors to base name + suffix

* granite-switch: inline suffix checks in tensor op resolution

* granite-switch: drop switch-lora struct comment

* granite-switch: guard router layer index and inline n_slots

* granite-switch: group adapter metadata under {arch}.adapters.* namespace

* granite-switch: add hparams.has_rope(il) for KV-shift rope skipping

* granite-switch: skip arch in test-llama-archs (adapter fixture missing, TODO)

* granite-switch: Keys.Adapters namespace + simplify n_slots

* granite-switch: validate substitute token ids against n_vocab

* granite-switch: bound adapter count and lora rank from GGUF

* granite-switch: reject MTP context type when router_layer is set

* granite-switch: throw on bad adapter metadata instead of GGML_ASSERT

* granite-switch: use ASCII +/- in router K signal comment

* granite-switch: document n_layer_nextn repurpose and its leak points

* granite-switch: gate lora_a/lora_b op mapping on router_layer

* granite-switch: label all three preview model sizes
2026-08-10 09:53:46 +02:00
Georgi Gerganov 6ad4ab0ea0 readme : remove dev branches (#26832) 2026-08-10 09:53:26 +03:00
Aleksander Grygier 92d1bb0c99 ui: Linting & Formatting scripts (#26819) 2026-08-10 08:38:37 +02:00
Pascal 1e396e72a8 server: gate the docker tools runtime tests on a real container run (#26826)
docker info only proves the daemon answers, so the Windows CI passes
the check and then dies trying to run a linux image. The hosted
Windows runners cannot run one: GitHub states the VMs are not enabled
for nested virtualization and will not be, since they already sit one
level deep and the hypervisor does not support more levels
(https://github.com/orgs/community/discussions/25491). Probing the
image itself skips those tests there, and pulls it before the server
waits for the container id.
2026-08-10 09:32:58 +03:00
Caleb DeLeeuw 0377426cef model-saver : fix expert shared/chunk FFN length key clobber (#26693)
The saver called add_kv with LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH twice, the
second time passing n_ff_chexp. gguf_set_val_u32 removes-then-appends, so the second
call clobbers the first: the saved shared_feed_forward_length ends up as n_ff_chexp
(0 for every arch except GroveMoE), and expert_chunk_feed_forward_length is never
written at all.

So a save->load roundtrip of any MoE model with a shared expert loses n_ff_shexp. On
reload the arch falls back to n_ff for the shexp tensor shape, that no longer matches
the saved tensor, and the model FAILS to load. Hits qwen2moe, qwen3-next, granite-moe,
hunyuan-moe, ernie4.5, bailingmoe2, nemotron-h, and the other shared-expert MoEs.

Fix: the second call writes LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH.

test-llama-archs: set expert_shared_feed_forward_length to a value distinct from n_ff
in the MoE setup so the roundtrip exercises it. Without the fix the reload fails on a
shexp tensor-shape mismatch; with it, every arch roundtrips clean.
2026-08-10 09:32:01 +03:00
Eve aea252fb4a ci: fix the ctest sanitize runs (#26593)
* Update build-sanitize.yml

* make it run on pr

* fix thread

* Update build-sanitize.yml

* Update build-sanitize.yml

* just run thread on github machine
2026-08-10 09:31:28 +03:00
Masashi Yoshimura f401bb1390 ggml-webgpu : refactor several wgsl files and simplify flash_attn wgsl. (#26134) 2026-08-10 09:29:41 +03:00
Pascal 74ce15741b ui: degrade the working directory picker when file search is off (#26811)
The picker mounts whenever a cwd-aware builtin tool is enabled, so
it can open while file_glob_search is not served or was disabled by
the user. Every typed query then fired a search that could only
fail with a raw error.

Gate the debounced search on the tool state, the same way the
mention picker does, and show a message in place of the results
list that explains why search is unavailable. Manual entry with
Enter still commits a directory. The Browse button and the search
scope footer are hidden as well: Browse resolves the picked folder
name through file_glob_search, and the client-side toggle would not
stop that call.
2026-08-09 21:20:23 +02:00
Xuan-Son Nguyen 936918514c ci: add pr-draft-label (#26801) 2026-08-09 16:51:21 +02:00
582 changed files with 10894 additions and 7251 deletions
+25 -3
View File
@@ -15,6 +15,12 @@ on:
'**/*.cpp'
]
pull_request:
types: [opened, synchronize, reopened]
paths: [
'.github/workflows/build-sanitize.yml'
]
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
cancel-in-progress: true
@@ -28,19 +34,35 @@ env:
jobs:
ctest:
runs-on: [self-hosted, X64, CPU, Linux]
continue-on-error: true
strategy:
matrix:
sanitizer: [ADDRESS, THREAD, UNDEFINED]
include:
- sanitizer: ADDRESS
machine: [self-hosted, X64, Linux]
# thread doesn't run properly on some self hosted machines, so run it on Github instead
- sanitizer: THREAD
machine: ubuntu-24.04
- sanitizer: UNDEFINED
machine: [self-hosted, X64, Linux]
runs-on: ${{ matrix.machine }}
steps:
- name: Clone
id: checkout
uses: actions/checkout@v6
- name: ccache
uses: ggml-org/ccache-action@v1.2.21
if: ${{ matrix.sanitizer == 'THREAD' }}
with:
key: ctest-thread-ubuntu-24.04
variant: ccache
evict-old-files: 1d
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
# with UNDEFINED sanitizer, we have to build in Debug to avoid GCC 13 false-positive warnings
- name: Build (undefined)
id: cmake_build_undefined
+23
View File
@@ -0,0 +1,23 @@
name: Convert PR to draft
on:
pull_request_target:
types: [labeled]
permissions:
pull-requests: write
issues: write
contents: write # required for "gh pr ready" command, see https://github.com/cli/cli/issues/8910
jobs:
convert-to-draft:
if: github.event.label.name == 'draft' && github.event.pull_request.draft == false
runs-on: ubuntu-slim
steps:
- name: Convert PR to draft
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PR_URL: ${{ github.event.pull_request.html_url }}
run: |
gh pr ready --undo "$PR_URL"
gh pr edit "$PR_URL" --remove-label draft
+1 -1
View File
@@ -12,7 +12,7 @@
[![Docker](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml/badge.svg)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
[![Winget](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml/badge.svg)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)
[manifesto](https://github.com/ggml-org/llama.cpp/discussions/205) / [ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Abartowski1182%20OR%20author%3Ahipudding%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3A0cc4m%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [dev branches](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-features.md) / [compile times](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-compile-times.md) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
[manifesto](https://github.com/ggml-org/llama.cpp/discussions/205) / [ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Abartowski1182%20OR%20author%3Ahipudding%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3A0cc4m%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [compile times](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-compile-times.md) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
</div>
+1
View File
@@ -103,6 +103,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"GraniteMoeForCausalLM": "granite",
"GraniteMoeHybridForCausalLM": "granite",
"GraniteMoeSharedForCausalLM": "granite",
"GraniteSwitchForCausalLM": "granite",
"GraniteSpeechForConditionalGeneration": "granite",
"GraniteSpeechPlusForConditionalGeneration": "granite",
"Grok1ForCausalLM": "grok",
+160
View File
@@ -123,6 +123,166 @@ class GraniteMoeModel(GraniteModel):
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("GraniteSwitchForCausalLM")
class GraniteSwitchModel(GraniteMoeModel):
"""Dense, all-attention Granite with N per-token embedded LoRA adapters, stacked
over the adapter dim with a zero adapter at slot 0 (N = num_adapters + 1)."""
model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH
# permute q/k per-slice below (NORM-rope layout), not via the parent's auto-permute
undo_permute = False
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# the weightless switch reserves one cache slot: one fewer block than num_hidden_layers
self.block_count = self.block_count - 1
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
self._n_adapters = int(self.hparams["num_adapters"])
self._max_lora_rank = int(self.hparams["max_lora_rank"])
self._n_slots = self._n_adapters + 1 # +1 for the zero slot at index 0
n_head = int(self.hparams["num_attention_heads"])
n_kv_head = int(self.hparams["num_key_value_heads"])
head_dim = (
self.hparams.get("projection_head_dim")
or self.hparams.get("head_dim")
or (self.hparams["hidden_size"] // n_head)
)
self._n_head = n_head
self._n_kv_head = n_kv_head
self._head_dim = int(head_dim)
self._q_size = n_head * self._head_dim
self._kv_size = n_kv_head * self._head_dim
def set_gguf_parameters(self):
super().set_gguf_parameters()
# dense: pin expert_used_count to 0 (config carries a leftover num_experts_per_tok)
if not self.hparams.get("num_local_experts"):
self.gguf_writer.add_expert_used_count(0)
self.gguf_writer.add_adapter_count(self._n_adapters)
self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank)
self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"])
self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"])
router_gain = float(self.hparams.get("control_token_gain", 15.0))
self.gguf_writer.add_adapter_router_gain(router_gain)
logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain)
def _lora_a(self, data: Tensor) -> Tensor:
# on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in]
a = data.squeeze(1)
zero = torch.zeros_like(a[:1])
return torch.cat([zero, a], dim=0).contiguous()
def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor:
# on-disk B: [n_adapters, 1, out, max_rank] -> [n_adapters+1, out, max_rank]
b = data.squeeze(1)
if permute_n_head is not None:
# permute each adapter's B output rows to match the permuted q/k base
b = torch.stack([self.permute(b[i], permute_n_head, permute_n_head) for i in range(b.shape[0])], dim=0)
zero = torch.zeros_like(b[:1])
return torch.cat([zero, b], dim=0).contiguous()
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
T = gguf.MODEL_TENSOR
# skip the weightless switch + control-token buffers (rebuilt at load time)
bare = name.split(".")[-1]
if (
name.startswith("model.switch.") or name.startswith("switch.")
or bare in ("adapter_token_ids", "control_to_substitute_lut")
):
return
if "self_attn.qkv_proj" in name:
if name.endswith("base_layer.weight"):
# fused [q|k|v] rows: permute q/k row-blocks for ggml's NORM-rope layout
q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0)
q = self.permute(q, self._n_head, self._n_head)
k = self.permute(k, self._n_kv_head, self._n_kv_head)
fused = torch.cat([q, k, v], dim=0)
yield (self.format_tensor_name(T.ATTN_QKV, bid), fused)
return
if "lora_A_slices." in name:
slot = int(name.rsplit(".", 1)[1])
key = {0: T.ATTN_Q, 1: T.ATTN_K, 2: T.ATTN_V}[slot]
yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))
return
if "lora_B_slices." in name:
slot = int(name.rsplit(".", 1)[1])
key, ph = {
0: (T.ATTN_Q, self._n_head),
1: (T.ATTN_K, self._n_kv_head),
2: (T.ATTN_V, None),
}[slot]
yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch, ph))
return
raise ValueError(f"Unexpected qkv_proj tensor: {name}")
if "self_attn.o_proj" in name:
if name.endswith("base_layer.weight"):
yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch)
return
if name.endswith("lora_A"):
yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_a"), self._lora_a(data_torch))
return
if name.endswith("lora_B"):
yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_b"), self._lora_b(data_torch))
return
raise ValueError(f"Unexpected o_proj tensor: {name}")
if "shared_mlp.input_linear" in name:
ffn = self.hparams["shared_intermediate_size"]
if name.endswith("base_layer.weight"):
gate, up = data_torch.split([ffn, ffn], dim=0)
yield (self.format_tensor_name(T.FFN_GATE, bid), gate)
yield (self.format_tensor_name(T.FFN_UP, bid), up)
return
if "lora_A_slices." in name:
slot = int(name.rsplit(".", 1)[1])
key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]
yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))
return
if "lora_B_slices." in name:
slot = int(name.rsplit(".", 1)[1])
key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]
yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch))
return
raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}")
if "shared_mlp.output_linear" in name:
if name.endswith("base_layer.weight"):
yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch)
return
if name.endswith("lora_A"):
yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_a"), self._lora_a(data_torch))
return
if name.endswith("lora_B"):
yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_b"), self._lora_b(data_torch))
return
raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}")
if bid is not None and ".layers." in name and (
"input_layernorm" in name or "post_attention_layernorm" in name
):
key = T.ATTN_NORM if "input_layernorm" in name else T.FFN_NORM
yield (self.format_tensor_name(key, bid), data_torch)
return
if name in ("model.embed_tokens.weight", "embed_tokens.weight"):
yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch)
return
if name in ("model.norm.weight", "norm.weight"):
yield (self.format_tensor_name(T.OUTPUT_NORM), data_torch)
return
if name == "lm_head.weight":
return # tied to token_embd
raise ValueError(f"graniteswitch: unhandled tensor {name!r} (bid={bid})")
@ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM")
class GraniteHybridModel(Mamba2Model, GraniteMoeModel):
"""GraniteHybrid is a hybrid SSM + Attention model that uses Mamba2 SSM
+71 -8
View File
@@ -197,6 +197,7 @@ class NemotronHModel(GraniteHybridModel):
"""Hybrid mamba2/attention model from NVIDIA"""
model_arch = gguf.MODEL_ARCH.NEMOTRON_H
is_moe: bool = False
supports_mtp_export = True
def __init__(self, *args, **kwargs):
# We have to determine the correct model architecture (MoE vs non-MoE) before
@@ -236,6 +237,25 @@ class NemotronHModel(GraniteHybridModel):
self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"]
self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"]
# `--no-mtp` drops it entirely; `--mtp` exports only the MTP head
self._mtp_bid: int | None = None
if self.is_moe and not self.no_mtp:
n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0
if n_nextn > 0:
assert n_nextn == 1, (
"NemotronH MTP conversion currently supports num_nextn_predict_layers == 1"
)
self._mtp_bid = self.block_count
self.block_count += 1
# The folded MTP block carries both an attention sub-layer and a
# MoE sub-layer, so register it as both so the per-layer metadata arrays cover it
self._attn_layers.append(self._mtp_bid)
self._mlp_layers.append(self._mtp_bid)
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
if self.mtp_only and self._mtp_bid is None:
raise ValueError("--mtp was requested, but this model does not contain a supported MTP head")
def get_attn_layers(self):
pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")
if pattern is None:
@@ -246,6 +266,36 @@ class NemotronHModel(GraniteHybridModel):
return [i for i, val in enumerate(pattern) if val == "attention"]
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
name, gen = item
if name.startswith("mtp."):
# --no-mtp: drop the MTP head entirely
if cls.no_mtp:
return None
elif cls.mtp_only:
# --mtp: export the MTP head plus the tensors it shares with the target model
keep = name in (
"backbone.embeddings.weight",
"backbone.norm_f.weight",
"lm_head.weight",
)
if not keep:
return None
return super().filter_tensors((name, gen))
def prepare_metadata(self, vocab_only: bool):
from_dir = self.fname_out.is_dir()
super().prepare_metadata(vocab_only=vocab_only)
if not self.mtp_only or not from_dir:
return
output_type: str = self.ftype.name.partition("_")[2]
fname_default: str = gguf.naming_convention(
self.metadata.name, self.metadata.basename, self.metadata.finetune,
self.metadata.version, size_label=None, output_type=output_type, model_type=None)
self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
def set_gguf_parameters(self):
super().set_gguf_parameters()
@@ -284,6 +334,10 @@ class NemotronHModel(GraniteHybridModel):
if (latent_size := self.hparams.get("moe_latent_size")) is not None:
self.gguf_writer.add_moe_latent_size(latent_size)
# MTP head: number of trailing NextN blocks
if self._mtp_bid is not None:
self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"])
def set_vocab(self):
# The NemotronH config uses pattern characters (e.g. '-') that may not
# be supported by the installed transformers version. AutoTokenizer
@@ -350,15 +404,24 @@ class NemotronHModel(GraniteHybridModel):
if not self.is_moe:
self.gguf_writer.add_add_bos_token(True)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if self.is_moe and bid is not None:
# Skip Multi-Token Prediction (MTP) tensors. These are used for
# for speculative decoding but we don't include them in this model
# conversion. See https://github.com/ggml-org/llama.cpp/pull/18886
if name.startswith("mtp."):
logger.info(f"gguf: Skipping MTP (Speculative) layer: {name}")
return
_MTP_SPECIAL_RENAMES = {
"mtp.layers.0.enorm.weight": "model.layers.{bid}.enorm.weight",
"mtp.layers.0.hnorm.weight": "model.layers.{bid}.hnorm.weight",
"mtp.layers.0.eh_proj.weight": "model.layers.{bid}.eh_proj.weight",
"mtp.layers.1.norm.weight": "model.layers.{bid}.post_attention_layernorm.weight",
"mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",
}
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
# mtp.layers.0: NextN input fusion + attention
# mtp.layers.1: MoE + final head norm
if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")):
suffix = name.split(".", 3)[3]
bid = self._mtp_bid
renamed = self._MTP_SPECIAL_RENAMES.get(name)
name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}"
if self.is_moe and bid is not None:
if name.endswith("mixer.gate.e_score_correction.bias"):
yield from ModelBase.modify_tensors(self, data_torch, name, bid)
return
+16 -15
View File
@@ -3221,17 +3221,17 @@ class ggml_webgpu_shader_lib {
auto push_type_defines = [&](const char * prefix, ggml_type type) {
std::string s_prefix = prefix;
if (type == GGML_TYPE_F32) {
defines.push_back(s_prefix + "_F32");
defines.push_back(s_prefix + "=f32");
} else if (type == GGML_TYPE_F16) {
defines.push_back(s_prefix + "_F16");
defines.push_back(s_prefix + "=f16");
} else {
GGML_ABORT("Unsupported type for CONV_2D shader");
}
};
push_type_defines("WEIGHT", key.weight_type);
push_type_defines("INPUT", key.input_type);
push_type_defines("OUTPUT", key.output_type);
push_type_defines("WEIGHT_TYPE", key.weight_type);
push_type_defines("INPUT_TYPE", key.input_type);
push_type_defines("OUTPUT_TYPE", key.output_type);
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
@@ -3263,17 +3263,18 @@ class ggml_webgpu_shader_lib {
auto push_type_defines = [&](const char * prefix, ggml_type type) {
std::string s_prefix = prefix;
if (type == GGML_TYPE_F32) {
defines.push_back(s_prefix + "_F32");
defines.push_back(s_prefix + "=f32");
} else if (type == GGML_TYPE_F16) {
defines.push_back(s_prefix + "_F16");
defines.push_back(s_prefix + "=f16");
} else {
GGML_ABORT("Unsupported type for CONV_2D_DW shader");
GGML_ABORT("Unsupported type for CONV_2D shader");
}
};
push_type_defines("WEIGHT", key.weight_type);
push_type_defines("INPUT", key.input_type);
push_type_defines("OUTPUT", key.output_type);
push_type_defines("WEIGHT_TYPE", key.weight_type);
push_type_defines("INPUT_TYPE", key.input_type);
push_type_defines("OUTPUT_TYPE", key.output_type);
if (whcn) {
defines.push_back("WHCN");
}
@@ -3304,16 +3305,16 @@ class ggml_webgpu_shader_lib {
auto push_type_defines = [&](const char * prefix, ggml_type type) {
std::string s_prefix = prefix;
if (type == GGML_TYPE_F32) {
defines.push_back(s_prefix + "_F32");
defines.push_back(s_prefix + "=f32");
} else if (type == GGML_TYPE_F16) {
defines.push_back(s_prefix + "_F16");
defines.push_back(s_prefix + "=f16");
} else {
GGML_ABORT("Unsupported type for IM2COL shader");
}
};
push_type_defines("INPUT", key.input_type);
push_type_defines("OUTPUT", key.output_type);
push_type_defines("INPUT_TYPE", key.input_type);
push_type_defines("OUTPUT_TYPE", key.output_type);
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+12 -24
View File
@@ -930,7 +930,6 @@ static webgpu_encoded_op ggml_webgpu_solve_tri(webgpu_context & ctx,
(uint32_t) src1->ne[0],
(uint32_t) dst->ne[2],
(uint32_t) dst->ne[3],
};
std::vector<wgpu::BindGroupEntry> entries = {
@@ -1039,7 +1038,6 @@ static webgpu_encoded_op ggml_webgpu_conv_2d_dw(webgpu_context & ctx,
(uint32_t) ggml_nelements(dst),
(uint32_t) dst->ne[2],
(uint32_t) dst->ne[3],
(uint32_t) dst->ne[0],
(uint32_t) dst->ne[1],
(uint32_t) src1->ne[0],
@@ -1328,7 +1326,6 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx,
(uint32_t) src0->ne[2],
(uint32_t) src4->ne[1],
(uint32_t) src1->ne[2],
(uint32_t) src1->ne[3],
(uint32_t) ggml_nelements(src1),
};
@@ -1921,25 +1918,20 @@ static bool ggml_webgpu_flash_attn_use_vec_path(const webgpu_global_context & gl
const ggml_tensor * K,
const ggml_tensor * V) {
const size_t storage_offset_alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment;
const bool k_float_vec4_aligned = (K->type != GGML_TYPE_F16 && K->type != GGML_TYPE_F32) ||
ggml_webgpu_flash_attn_float_vec4_aligned(K, storage_offset_alignment);
const bool v_float_vec4_aligned = (V->type != GGML_TYPE_F16 && V->type != GGML_TYPE_F32) ||
ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
const bool k_vec_type_supported =
K->type == GGML_TYPE_F32 || K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0;
const bool v_vec_type_supported =
V->type == GGML_TYPE_F32 || V->type == GGML_TYPE_F16 || V->type == GGML_TYPE_Q4_0 || V->type == GGML_TYPE_Q8_0;
const uint32_t k_vec_head_align = (K->type == GGML_TYPE_F32 || K->type == GGML_TYPE_F16) ?
GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH :
(uint32_t) ggml_blck_size(K->type);
const uint32_t v_vec_head_align = (V->type == GGML_TYPE_F32 || V->type == GGML_TYPE_F16) ?
GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH :
(uint32_t) ggml_blck_size(V->type);
const bool kv_vec_head_dims_aligned = Q->ne[0] % k_vec_head_align == 0 && V->ne[0] % v_vec_head_align == 0;
const bool k_float_vec4_aligned = (K->type != GGML_TYPE_F16 && K->type != GGML_TYPE_F32) ||
ggml_webgpu_flash_attn_float_vec4_aligned(K, storage_offset_alignment);
const bool v_float_vec4_aligned = (V->type != GGML_TYPE_F16 && V->type != GGML_TYPE_F32) ||
ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
const uint32_t k_vec_head_align =
ggml_is_quantized(K->type) ? ggml_blck_size(K->type) : GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH;
const uint32_t v_vec_head_align =
ggml_is_quantized(V->type) ? ggml_blck_size(V->type) : GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH;
const bool kv_vec_head_dims_aligned = Q->ne[0] % k_vec_head_align == 0 && V->ne[0] % v_vec_head_align == 0;
return global_ctx->capabilities.supports_subgroups && (Q->ne[1] < GGML_WEBGPU_FLASH_ATTN_VEC_MAX_SEQ_LEN) &&
kv_vec_head_dims_aligned && k_vec_type_supported && v_vec_type_supported && k_float_vec4_aligned &&
v_float_vec4_aligned;
kv_vec_head_dims_aligned && k_float_vec4_aligned && v_float_vec4_aligned;
}
static ggml_webgpu_flash_attn_op ggml_webgpu_flash_attn_prepare(webgpu_context & ctx,
@@ -2514,7 +2506,6 @@ static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx,
(uint32_t) dst->ne[0],
(uint32_t) dst->ne[1],
(uint32_t) dst->ne[2],
(uint32_t) dst->ne[3],
dim,
(uint32_t) src0->ne[dim] };
@@ -2610,7 +2601,6 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_rms_norm_mul(webgpu_context
(uint32_t) dst->ne[0],
(uint32_t) dst->ne[1],
(uint32_t) dst->ne[2],
(uint32_t) dst->ne[3],
ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(rn_dst, 0)) // epsilon, treated as f32 in the shader
};
@@ -2666,7 +2656,6 @@ static webgpu_encoded_op ggml_webgpu_row_norm(webgpu_context & ctx, ggml_tensor
(uint32_t) src->ne[0],
(uint32_t) src->ne[1],
(uint32_t) src->ne[2],
(uint32_t) src->ne[3],
ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 0)) // epsilon, treated as f32 in the shader
};
@@ -2925,7 +2914,6 @@ static webgpu_encoded_op ggml_webgpu_soft_max(webgpu_context & ctx,
(uint32_t) (dst->nb[1] / ggml_type_size(dst->type)),
(uint32_t) (dst->nb[2] / ggml_type_size(dst->type)),
(uint32_t) (dst->nb[3] / ggml_type_size(dst->type)),
(uint32_t) ggml_nelements(dst),
(uint32_t) src0->ne[0],
(uint32_t) src0->ne[1],
(uint32_t) src0->ne[2],
@@ -18,7 +18,6 @@ struct Params {
ne0: u32,
ne1: u32,
ne2: u32,
ne3: u32,
dim: u32,
src0_nedim: u32
+6 -44
View File
@@ -2,25 +2,11 @@
enable f16;
@group(0) @binding(0)
#if defined(WEIGHT_F32)
var<storage, read_write> weights: array<f32>;
#elif defined(WEIGHT_F16)
var<storage, read_write> weights: array<f16>;
#endif
var<storage, read_write> weights: array<WEIGHT_TYPE>;
@group(0) @binding(1)
#if defined(INPUT_F32)
var<storage, read_write> input: array<f32>;
#elif defined(INPUT_F16)
var<storage, read_write> input: array<f16>;
#endif
var<storage, read_write> input: array<INPUT_TYPE>;
@group(0) @binding(2)
#if defined(OUTPUT_F32)
var<storage, read_write> output: array<f32>;
#elif defined(OUTPUT_F16)
var<storage, read_write> output: array<f16>;
#endif
var<storage, read_write> output: array<OUTPUT_TYPE>;
struct Params {
offset_w: u32,
@@ -50,30 +36,6 @@ struct Params {
@group(0) @binding(3)
var<uniform> params: Params;
fn load_weight(idx: u32) -> f32 {
#if defined(WEIGHT_F32)
return weights[idx];
#elif defined(WEIGHT_F16)
return f32(weights[idx]);
#endif
}
fn load_input(idx: u32) -> f32 {
#if defined(INPUT_F32)
return input[idx];
#elif defined(INPUT_F16)
return f32(input[idx]);
#endif
}
fn store_output(idx: u32, val: f32) {
#if defined(OUTPUT_F32)
output[idx] = val;
#elif defined(OUTPUT_F16)
output[idx] = f16(val);
#endif
}
fn ceil_div_u32(x: u32, y: u32) -> u32 {
return (x + y - 1) / y;
}
@@ -136,7 +98,7 @@ fn main(
// entire receptive field is out of bounds
if (kw_begin >= kw_end || kh_begin >= kh_end) {
let out_idx = params.offset_o + ow * params.so0 + oh * params.so1 + oc * params.so2 + n * params.so3;
store_output(out_idx, 0.0);
output[out_idx] = OUTPUT_TYPE(0.0);
return;
}
@@ -155,11 +117,11 @@ fn main(
let iw = u32(ow_base + i32(kw * params.d0));
let w_idx = w_row_base + kw * params.sw0;
let in_idx = in_row_base + iw * params.si0;
sum += load_weight(w_idx) * load_input(in_idx);
sum += f32(weights[w_idx]) * f32(input[in_idx]);
}
}
}
let out_idx = params.offset_o + ow * params.so0 + oh * params.so1 + oc * params.so2 + n * params.so3;
store_output(out_idx, sum);
output[out_idx] = OUTPUT_TYPE(sum);
}
@@ -6,25 +6,11 @@ enable f16;
// weight (src0) is [KW,KH,1,C]; output matches the input layout.
@group(0) @binding(0)
#if defined(WEIGHT_F32)
var<storage, read_write> weights: array<f32>;
#elif defined(WEIGHT_F16)
var<storage, read_write> weights: array<f16>;
#endif
var<storage, read_write> weights: array<WEIGHT_TYPE>;
@group(0) @binding(1)
#if defined(INPUT_F32)
var<storage, read_write> input: array<f32>;
#elif defined(INPUT_F16)
var<storage, read_write> input: array<f16>;
#endif
var<storage, read_write> input: array<INPUT_TYPE>;
@group(0) @binding(2)
#if defined(OUTPUT_F32)
var<storage, read_write> output: array<f32>;
#elif defined(OUTPUT_F16)
var<storage, read_write> output: array<f16>;
#endif
var<storage, read_write> output: array<OUTPUT_TYPE>;
struct Params {
offset_w: u32,
@@ -33,7 +19,6 @@ struct Params {
ne: u32,
channels: u32,
batches: u32,
dst_w: u32, dst_h: u32,
src_w: u32, src_h: u32,
knl_w: u32, knl_h: u32,
@@ -46,28 +31,6 @@ struct Params {
@group(0) @binding(3)
var<uniform> params: Params;
fn load_weight(idx: u32) -> f32 {
#if defined(WEIGHT_F32)
return weights[idx];
#elif defined(WEIGHT_F16)
return f32(weights[idx]);
#endif
}
fn load_input(idx: u32) -> f32 {
#if defined(INPUT_F32)
return input[idx];
#elif defined(INPUT_F16)
return f32(input[idx]);
#endif
}
fn store_output(idx: u32, val: f32) {
#if defined(OUTPUT_F32)
output[idx] = val;
#elif defined(OUTPUT_F16)
output[idx] = f16(val);
#endif
}
#if defined(WHCN)
// Input/output/kernel contiguous in [W, H, C, N] order (kernel [KW,KH,C]).
fn conv_2d_dw(idx: u32) -> f32 {
@@ -89,8 +52,8 @@ fn conv_2d_dw(idx: u32) -> f32 {
for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
let v = load_input(src_i + u32(src_y) * params.src_w + u32(src_x));
let k = load_weight(knl_i + ky * params.knl_w + kx);
let v = f32(input[src_i + u32(src_y) * params.src_w + u32(src_x)]);
let k = f32(weights[knl_i + ky * params.knl_w + kx]);
sum += v * k;
}
}
@@ -117,8 +80,8 @@ fn conv_2d_dw(idx: u32) -> f32 {
for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
let v = load_input(src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c);
let k = load_weight(params.offset_w + ky * knl_row + kx * params.channels + c);
let v = f32(input[src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c]);
let k = f32(weights[params.offset_w + ky * knl_row + kx * params.channels + c]);
sum += v * k;
}
}
@@ -133,5 +96,5 @@ fn main(
) {
let idx = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
if (idx >= params.ne) { return; }
store_output(params.offset_o + idx, conv_2d_dw(idx));
output[params.offset_o + idx] = OUTPUT_TYPE(conv_2d_dw(idx));
}
@@ -7,32 +7,18 @@ enable chromium_experimental_subgroup_matrix;
#define BYTE_HELPERS
#include "common_decls.tmpl"
#ifdef K_F32
#define K_TYPE f32
#elif defined(K_Q4_0) || defined(K_Q8_0)
#define K_TYPE u32
#else
#define K_TYPE f16
#endif
#ifdef V_F32
#define V_TYPE f32
#elif defined(V_Q4_0) || defined(V_Q8_0)
#define V_TYPE u32
#else
#define V_TYPE f16
#endif
#define FLASH_ATTN_SCALAR_KV
#include "flash_attn_decls.tmpl"
// Default values
// The actual values are defined in shader-lib.
#define HEAD_DIM_QK 64
#define HEAD_DIM_V 64
// The number of rows/columns/k in a subgroup matrix. MxK * KxN = MxN
// Note that the "K" here does not correspond to the K in attention's Q/K/V, it's just the common dimension.
#define SG_MAT_M 8
#define SG_MAT_N 8
#define SG_MAT_K 8
// Each workgroup processes one subgroup matrix of Q rows
#define Q_TILE SG_MAT_M
#define KV_TILE 16
@@ -41,104 +27,13 @@ enable chromium_experimental_subgroup_matrix;
// Number of subgroup-matrix-width blocks that span the KV tile. SG_MAT_N must divide KV_TILE.
#define KV_BLOCKS (KV_TILE / SG_MAT_N)
struct Params {
offset_q: u32,
offset_k: u32,
offset_v: u32,
offset_mask: u32,
offset_sinks: u32,
offset_dst: u32,
// shapes of Q/K/V
n_heads: u32,
seq_len_q: u32,
seq_len_kv: u32,
// strides (in elements)
stride_q1: u32,
stride_q2: u32,
stride_q3: u32,
stride_k1: u32,
stride_k2: u32,
stride_k3: u32,
stride_v1: u32,
stride_v2: u32,
stride_v3: u32,
stride_mask3: u32,
// repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
q_per_kv: u32,
// softmax params
scale: f32,
max_bias: f32,
logit_softcap: f32,
n_head_log2: f32,
m0: f32,
m1: f32,
};
@group(0) @binding(0) var<storage, read_write> Q: array<f32>;
#ifdef KV_OVERLAP
@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
#define V K
#else
@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
@group(0) @binding(2) var<storage, read_write> V: array<V_TYPE>;
#endif
#if defined(MASK) && defined(SINKS)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 4
#define PARAMS_BINDING 5
#else
@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
@group(0) @binding(4) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 5
#define PARAMS_BINDING 6
#endif
#elif defined(MASK)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
#define DST_BINDING 3
#define PARAMS_BINDING 4
#else
@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
#define DST_BINDING 4
#define PARAMS_BINDING 5
#endif
#elif defined(SINKS)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 3
#define PARAMS_BINDING 4
#else
@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 4
#define PARAMS_BINDING 5
#endif
#else
#ifdef KV_OVERLAP
#define DST_BINDING 2
#define PARAMS_BINDING 3
#else
#define DST_BINDING 3
#define PARAMS_BINDING 4
#endif
#endif
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<f32>>;
@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
// Just a very small float value.
const FLOAT_MIN: f32 = -1.0e9;
// The number of Q rows processed per workgroup
var<workgroup> q_shmem: array<f16, Q_TILE * HEAD_DIM_QK>;
#if !defined(K_DIRECT) || !defined(V_DIRECT)
#define STAGING_SHMEM kv_shmem
#define STAGING_OUT_TYPE f16
#include "flash_attn_staging.tmpl"
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
@@ -175,50 +70,6 @@ fn calc_softmax_term(kv_idx: u32, q_tile_row: u32, slope: f32) -> f32 {
return v;
}
fn load_f32x4(buf: ptr<storage, array<vec4<f32>>, read_write>, scalar_index: u32) -> vec4<f32> {
return (*buf)[scalar_index >> 2u];
}
fn load_kx4(buf: ptr<storage, array<vec4<K_TYPE>>, read_write>, scalar_index: u32) -> vec4<K_TYPE> {
return (*buf)[scalar_index >> 2u];
}
#if !defined(K_DIRECT) || !defined(V_DIRECT)
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f16
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
let k_row = elem_idx / HEAD_DIM_QK;
let k_col = elem_idx % HEAD_DIM_QK;
let global_k_row = kv_tile + k_row;
let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
kv_shmem[elem_idx] = f16(select(
0.0,
K[global_k_row_offset + k_col],
global_k_row < params.seq_len_kv && k_col < HEAD_DIM_QK));
}
}
#endif
#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
let v_row = elem_idx / HEAD_DIM_V;
let v_col = elem_idx % HEAD_DIM_V;
let global_v_row = kv_tile + v_row;
let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
kv_shmem[elem_idx] = f16(select(
0.0,
V[global_v_row_offset + v_col],
global_v_row < params.seq_len_kv && v_col < HEAD_DIM_V));
}
}
#endif
#endif
@compute @workgroup_size(WG_SIZE)
fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
@builtin(local_invocation_id) local_id: vec3<u32>,
@@ -0,0 +1,134 @@
#ifdef Q_F32
#define Q_TYPE f32
#else
#define Q_TYPE f16
#endif
#ifdef K_F32
#define K_TYPE f32
#elif defined(K_Q4_0) || defined(K_Q8_0)
#define K_TYPE u32
#else
#define K_TYPE f16
#endif
#ifdef V_F32
#define V_TYPE f32
#elif defined(V_Q4_0) || defined(V_Q8_0)
#define V_TYPE u32
#else
#define V_TYPE f16
#endif
#ifdef DST_F32
#define DST_TYPE f32
#else
#define DST_TYPE f16
#endif
#if defined(FLASH_ATTN_SCALAR_KV) || defined(K_Q4_0) || defined(K_Q8_0)
#define K_STORAGE_TYPE K_TYPE
#else
#define K_STORAGE_TYPE vec4<K_TYPE>
#endif
#if defined(FLASH_ATTN_SCALAR_KV) || defined(V_Q4_0) || defined(V_Q8_0)
#define V_STORAGE_TYPE V_TYPE
#else
#define V_STORAGE_TYPE vec4<V_TYPE>
#endif
// Just a very small float value.
const FLOAT_MIN: f32 = -1.0e9;
struct Params {
offset_q: u32,
offset_k: u32,
offset_v: u32,
offset_mask: u32,
offset_sinks: u32,
offset_dst: u32,
// shapes of Q/K/V
n_heads: u32,
seq_len_q: u32,
seq_len_kv: u32,
// strides (in elements)
stride_q1: u32,
stride_q2: u32,
stride_q3: u32,
stride_k1: u32,
stride_k2: u32,
stride_k3: u32,
stride_v1: u32,
stride_v2: u32,
stride_v3: u32,
stride_mask3: u32,
// repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
q_per_kv: u32,
// softmax params
scale: f32,
max_bias: f32,
logit_softcap: f32,
n_head_log2: f32,
m0: f32,
m1: f32,
#ifdef FLASH_ATTN_VEC_SPLIT
#ifdef BLK
blk_base: u32,
blk_nblk0: u32,
blk_nblk1: u32,
#endif
tmp_data_base: u32,
tmp_stats_base: u32,
nwg: u32,
#endif
};
@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
@group(0) @binding(1) var<storage, read_write> K: array<K_STORAGE_TYPE>;
#ifdef KV_OVERLAP
#define V K
#define MASK_BINDING 2
#else
@group(0) @binding(2) var<storage, read_write> V: array<V_STORAGE_TYPE>;
#define MASK_BINDING 3
#endif // KV_OVERLAP
#ifdef MASK
@group(0) @binding(MASK_BINDING) var<storage, read_write> mask: array<f16>;
#define SINKS_BINDING (MASK_BINDING + 1)
#else
#define SINKS_BINDING MASK_BINDING
#endif
#ifdef SINKS
@group(0) @binding(SINKS_BINDING) var<storage, read_write> sinks: array<f32>;
#define BLK_BINDING (SINKS_BINDING + 1)
#else
#define BLK_BINDING SINKS_BINDING
#endif
#ifdef FLASH_ATTN_VEC_SPLIT
#ifdef BLK
@group(0) @binding(BLK_BINDING) var<storage, read_write> blk: array<u32>;
#define TMP_BINDING (BLK_BINDING + 1)
#else
#define TMP_BINDING BLK_BINDING
#endif
@group(0) @binding(TMP_BINDING) var<storage, read_write> tmp: array<f32>;
#define DST_BINDING (TMP_BINDING + 1)
#else
#define DST_BINDING BLK_BINDING
#endif // FLASH_ATTN_VEC_SPLIT
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
#define PARAMS_BINDING (DST_BINDING + 1)
@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
@@ -1,83 +0,0 @@
#include "quant_inner_loops.tmpl"
#define BLOCK_SIZE 32
#define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
#define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
#if defined(K_Q4_0)
#define K_NQ 16
#define K_BLOCK_SIZE_BYTES 18u
#define K_BYTES_PER_THREAD 8u
#define K_BYTES_PER_INNER_LOOP 4u
#elif defined(K_Q8_0)
#define K_NQ 16
#define K_BLOCK_SIZE_BYTES 34u
#define K_BYTES_PER_THREAD 16u
#define K_BYTES_PER_INNER_LOOP 4u
#endif
#if defined(V_Q4_0)
#define V_NQ 16
#define V_BLOCK_SIZE_BYTES 18u
#define V_BYTES_PER_THREAD 8u
#define V_BYTES_PER_INNER_LOOP 4u
#elif defined(V_Q8_0)
#define V_NQ 16
#define V_BLOCK_SIZE_BYTES 34u
#define V_BYTES_PER_THREAD 16u
#define V_BYTES_PER_INNER_LOOP 4u
#endif
#if defined(K_Q4_0) || defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
let blck_idx = elem_idx / BLOCK_SIZE;
let block_offset = (elem_idx % BLOCK_SIZE) / K_NQ;
let k_row = blck_idx / BLOCKS_K;
let global_k_row = kv_tile + k_row;
let block_k = blck_idx % BLOCKS_K;
let row_offset = k_row * HEAD_DIM_QK;
let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k;
let block_byte_base = global_block_idx * K_BLOCK_SIZE_BYTES;
let d = f16_from_u16(load_k_u16_at(block_byte_base));
let thread_byte_offset = block_offset * K_BYTES_PER_THREAD;
let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
for (var j = 0u; j < K_BYTES_PER_THREAD / K_BYTES_PER_INNER_LOOP; j += 1u) {
let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * K_BYTES_PER_INNER_LOOP;
let q_packed = load_k_u32_at(q_byte_offset);
#if defined(K_Q4_0)
dequant_q4_0_packed_to_shmem(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
#elif defined(K_Q8_0)
dequant_q8_0_packed_to_shmem(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
#endif
}
}
}
#endif
#if defined(V_Q4_0) || defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
for (var elem_idx = local_x * V_NQ; elem_idx < kv_count * HEAD_DIM_V; elem_idx += WG_SIZE * V_NQ) {
let blck_idx = elem_idx / BLOCK_SIZE;
let block_offset = (elem_idx % BLOCK_SIZE) / V_NQ;
let v_row = blck_idx / BLOCKS_V;
let global_v_row = kv_tile + v_row;
let block_k = blck_idx % BLOCKS_V;
let row_offset = v_row * HEAD_DIM_V;
let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k;
let block_byte_base = global_block_idx * V_BLOCK_SIZE_BYTES;
let d = f16_from_u16(load_v_u16_at(block_byte_base));
let thread_byte_offset = block_offset * V_BYTES_PER_THREAD;
let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
for (var j = 0u; j < V_BYTES_PER_THREAD / V_BYTES_PER_INNER_LOOP; j += 1u) {
let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * V_BYTES_PER_INNER_LOOP;
let q_packed = load_v_u32_at(q_byte_offset);
#if defined(V_Q4_0)
dequant_q4_0_packed_to_shmem(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
#elif defined(V_Q8_0)
dequant_q8_0_packed_to_shmem(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
#endif
}
}
}
#endif
@@ -0,0 +1,136 @@
#if defined(K_Q4_0) || defined(K_Q8_0) || defined(V_Q4_0) || defined(V_Q8_0)
#define QUANT_SHMEM STAGING_SHMEM
#define QUANT_OUT_TYPE STAGING_OUT_TYPE
#include "quant_inner_loops.tmpl"
#undef QUANT_SHMEM
#undef QUANT_OUT_TYPE
#define BLOCK_SIZE 32
#define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
#define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
#endif
#if defined(K_Q4_0)
#define K_NQ 16
#define K_BLOCK_SIZE_BYTES 18u
#define K_BYTES_PER_THREAD 8u
#define K_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_K_PACKED_TO_SHMEM dequant_q4_0_packed_to_shmem
#elif defined(K_Q8_0)
#define K_NQ 16
#define K_BLOCK_SIZE_BYTES 34u
#define K_BYTES_PER_THREAD 16u
#define K_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_K_PACKED_TO_SHMEM dequant_q8_0_packed_to_shmem
#endif
#if defined(V_Q4_0)
#define V_NQ 16
#define V_BLOCK_SIZE_BYTES 18u
#define V_BYTES_PER_THREAD 8u
#define V_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_V_PACKED_TO_SHMEM dequant_q4_0_packed_to_shmem
#elif defined(V_Q8_0)
#define V_NQ 16
#define V_BLOCK_SIZE_BYTES 34u
#define V_BYTES_PER_THREAD 16u
#define V_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_V_PACKED_TO_SHMEM dequant_q8_0_packed_to_shmem
#endif
#ifndef K_DIRECT
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
#if defined(K_Q4_0) || defined(K_Q8_0)
for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
let blck_idx = elem_idx / BLOCK_SIZE;
let block_offset = (elem_idx % BLOCK_SIZE) / K_NQ;
let k_row = blck_idx / BLOCKS_K;
let global_k_row = kv_tile + k_row;
let block_k = blck_idx % BLOCKS_K;
let row_offset = k_row * HEAD_DIM_QK;
let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k;
let block_byte_base = global_block_idx * K_BLOCK_SIZE_BYTES;
let d = f16_from_u16(load_k_u16_at(block_byte_base));
let thread_byte_offset = block_offset * K_BYTES_PER_THREAD;
let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
for (var j = 0u; j < K_BYTES_PER_THREAD / K_BYTES_PER_INNER_LOOP; j += 1u) {
let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * K_BYTES_PER_INNER_LOOP;
let q_packed = load_k_u32_at(q_byte_offset);
DEQUANT_K_PACKED_TO_SHMEM(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
}
}
#elif defined(FLASH_ATTN_SCALAR_KV)
for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
let k_row = elem_idx / HEAD_DIM_QK;
let k_col = elem_idx % HEAD_DIM_QK;
let global_k_row = kv_tile + k_row;
let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
STAGING_SHMEM[elem_idx] = STAGING_OUT_TYPE(select(
0.0,
K[global_k_row_offset + k_col],
global_k_row < params.seq_len_kv && k_col < HEAD_DIM_QK));
}
#else
for (var vec_idx_local = local_x; vec_idx_local < kv_count * Q_CHUNKS; vec_idx_local += WG_SIZE) {
let kv_local = vec_idx_local / Q_CHUNKS;
let chunk = vec_idx_local % Q_CHUNKS;
let global_k_row = kv_tile + kv_local;
let k_vec_index = (k_head_offset + global_k_row * params.stride_k1 + chunk * 4u) >> 2u;
let k4 = K[k_vec_index];
let kv_off = kv_local * HEAD_DIM_QK + chunk * 4u;
STAGING_SHMEM[kv_off + 0u] = STAGING_OUT_TYPE(k4.x);
STAGING_SHMEM[kv_off + 1u] = STAGING_OUT_TYPE(k4.y);
STAGING_SHMEM[kv_off + 2u] = STAGING_OUT_TYPE(k4.z);
STAGING_SHMEM[kv_off + 3u] = STAGING_OUT_TYPE(k4.w);
}
#endif
}
#endif // !defined(K_DIRECT)
#ifndef V_DIRECT
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
#if defined(V_Q4_0) || defined(V_Q8_0)
for (var elem_idx = local_x * V_NQ; elem_idx < kv_count * HEAD_DIM_V; elem_idx += WG_SIZE * V_NQ) {
let blck_idx = elem_idx / BLOCK_SIZE;
let block_offset = (elem_idx % BLOCK_SIZE) / V_NQ;
let v_row = blck_idx / BLOCKS_V;
let global_v_row = kv_tile + v_row;
let block_k = blck_idx % BLOCKS_V;
let row_offset = v_row * HEAD_DIM_V;
let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k;
let block_byte_base = global_block_idx * V_BLOCK_SIZE_BYTES;
let d = f16_from_u16(load_v_u16_at(block_byte_base));
let thread_byte_offset = block_offset * V_BYTES_PER_THREAD;
let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
for (var j = 0u; j < V_BYTES_PER_THREAD / V_BYTES_PER_INNER_LOOP; j += 1u) {
let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * V_BYTES_PER_INNER_LOOP;
let q_packed = load_v_u32_at(q_byte_offset);
DEQUANT_V_PACKED_TO_SHMEM(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
}
}
#elif defined(FLASH_ATTN_SCALAR_KV)
for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
let v_row = elem_idx / HEAD_DIM_V;
let v_col = elem_idx % HEAD_DIM_V;
let global_v_row = kv_tile + v_row;
let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
STAGING_SHMEM[elem_idx] = STAGING_OUT_TYPE(select(
0.0,
V[global_v_row_offset + v_col],
global_v_row < params.seq_len_kv && v_col < HEAD_DIM_V));
}
#else
for (var vec_idx_local = local_x; vec_idx_local < kv_count * V_CHUNKS; vec_idx_local += WG_SIZE) {
let kv_local = vec_idx_local / V_CHUNKS;
let chunk = vec_idx_local % V_CHUNKS;
let global_v_row = kv_tile + kv_local;
let v_vec_index = (v_head_offset + global_v_row * params.stride_v1 + chunk * 4u) >> 2u;
let v4 = V[v_vec_index];
let kv_off = kv_local * HEAD_DIM_V + chunk * 4u;
STAGING_SHMEM[kv_off + 0u] = STAGING_OUT_TYPE(v4.x);
STAGING_SHMEM[kv_off + 1u] = STAGING_OUT_TYPE(v4.y);
STAGING_SHMEM[kv_off + 2u] = STAGING_OUT_TYPE(v4.z);
STAGING_SHMEM[kv_off + 3u] = STAGING_OUT_TYPE(v4.w);
}
#endif
}
#endif // !defined(V_DIRECT)
@@ -3,192 +3,32 @@ enable subgroups;
#define BYTE_HELPERS
#include "common_decls.tmpl"
#include "flash_attn_decls.tmpl"
#ifdef Q_F16
#define Q_TYPE f16
#else
#define Q_TYPE f32
#endif
#ifdef K_F32
#define K_TYPE f32
#elif defined(K_Q4_0) || defined(K_Q8_0)
#define K_TYPE u32
#else
#define K_TYPE f16
#endif
#ifdef V_F32
#define V_TYPE f32
#elif defined(V_Q4_0) || defined(V_Q8_0)
#define V_TYPE u32
#else
#define V_TYPE f16
#endif
#ifdef DST_F16
#define DST_TYPE f16
#else
#define DST_TYPE f32
#endif
// Default values
// The actual values are defined in shader-lib.
#define HEAD_DIM_QK 64
#define HEAD_DIM_V 64
#define Q_TILE 4
#define KV_TILE 64
#define WG_SIZE 128
#ifndef MIN_SUBGROUP_SIZE
#define MIN_SUBGROUP_SIZE MAX_SUBGROUP_SIZE
#endif
struct Params {
offset_q: u32,
offset_k: u32,
offset_v: u32,
offset_mask: u32,
offset_sinks: u32,
offset_dst: u32,
n_heads: u32,
seq_len_q: u32,
seq_len_kv: u32,
stride_q1: u32,
stride_q2: u32,
stride_q3: u32,
stride_k1: u32,
stride_k2: u32,
stride_k3: u32,
stride_v1: u32,
stride_v2: u32,
stride_v3: u32,
stride_mask3: u32,
q_per_kv: u32,
scale: f32,
max_bias: f32,
logit_softcap: f32,
n_head_log2: f32,
m0: f32,
m1: f32,
};
@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
#ifdef KV_OVERLAP
#if defined(K_Q4_0) || defined(K_Q8_0)
@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
#else
@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
#endif
#define V K
#else
#if defined(K_Q4_0) || defined(K_Q8_0)
@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
#else
@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
#endif
#if defined(V_Q4_0) || defined(V_Q8_0)
@group(0) @binding(2) var<storage, read_write> V: array<V_TYPE>;
#else
@group(0) @binding(2) var<storage, read_write> V: array<vec4<V_TYPE>>;
#endif
#endif
#if defined(MASK) && defined(SINKS)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 4
#define PARAMS_BINDING 5
#else
@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
@group(0) @binding(4) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 5
#define PARAMS_BINDING 6
#endif
#elif defined(MASK)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
#define DST_BINDING 3
#define PARAMS_BINDING 4
#else
@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
#define DST_BINDING 4
#define PARAMS_BINDING 5
#endif
#elif defined(SINKS)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 3
#define PARAMS_BINDING 4
#else
@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
#define DST_BINDING 4
#define PARAMS_BINDING 5
#endif
#else
#ifdef KV_OVERLAP
#define DST_BINDING 2
#define PARAMS_BINDING 3
#else
#define DST_BINDING 3
#define PARAMS_BINDING 4
#endif
#endif
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
const FLOAT_MIN: f32 = -1.0e9;
const Q_CHUNKS: u32 = HEAD_DIM_QK / 4u;
const V_CHUNKS: u32 = HEAD_DIM_V / 4u;
const SCORE_REGS_PER_LANE: u32 = (KV_TILE + MIN_SUBGROUP_SIZE - 1u) / MIN_SUBGROUP_SIZE;
const OUT_REGS_PER_LANE: u32 = (V_CHUNKS + MIN_SUBGROUP_SIZE - 1u) / MIN_SUBGROUP_SIZE;
#if !defined(K_DIRECT) || !defined(V_DIRECT)
#define STAGING_SHMEM kv_shmem
#define STAGING_OUT_TYPE f16
#include "flash_attn_staging.tmpl"
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
#endif
var<workgroup> q_shmem: array<Q_TYPE, Q_TILE * HEAD_DIM_QK>;
var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f16
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_Q4_0) && !defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var vec_idx_local = local_x; vec_idx_local < kv_count * Q_CHUNKS; vec_idx_local += WG_SIZE) {
let kv_local = vec_idx_local / Q_CHUNKS;
let chunk = vec_idx_local % Q_CHUNKS;
let global_k_row = kv_tile + kv_local;
let k_vec_index = (k_head_offset + global_k_row * params.stride_k1 + chunk * 4u) >> 2u;
let k4 = K[k_vec_index];
let kv_off = kv_local * HEAD_DIM_QK + chunk * 4u;
kv_shmem[kv_off + 0u] = f16(k4.x);
kv_shmem[kv_off + 1u] = f16(k4.y);
kv_shmem[kv_off + 2u] = f16(k4.z);
kv_shmem[kv_off + 3u] = f16(k4.w);
}
}
#endif
#if !defined(V_Q4_0) && !defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
for (var vec_idx_local = local_x; vec_idx_local < kv_count * V_CHUNKS; vec_idx_local += WG_SIZE) {
let kv_local = vec_idx_local / V_CHUNKS;
let chunk = vec_idx_local % V_CHUNKS;
let global_v_row = kv_tile + kv_local;
let v_vec_index = (v_head_offset + global_v_row * params.stride_v1 + chunk * 4u) >> 2u;
let v4 = V[v_vec_index];
let kv_off = kv_local * HEAD_DIM_V + chunk * 4u;
kv_shmem[kv_off + 0u] = f16(v4.x);
kv_shmem[kv_off + 1u] = f16(v4.y);
kv_shmem[kv_off + 2u] = f16(v4.z);
kv_shmem[kv_off + 3u] = f16(v4.w);
}
}
#endif
@compute @workgroup_size(WG_SIZE)
fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
@builtin(local_invocation_id) local_id: vec3<u32>,
@@ -4,200 +4,35 @@ enable subgroups;
#define BYTE_HELPERS
#include "common_decls.tmpl"
#define FLASH_ATTN_VEC_SPLIT
#include "flash_attn_decls.tmpl"
#ifdef K_F32
#define K_TYPE f32
#elif defined(K_Q4_0) || defined(K_Q8_0)
#define K_TYPE u32
#else
#define K_TYPE f16
#endif
#ifdef V_F32
#define V_TYPE f32
#elif defined(V_Q4_0) || defined(V_Q8_0)
#define V_TYPE u32
#else
#define V_TYPE f16
#endif
#ifdef Q_F16
#define Q_TYPE f16
#else
#define Q_TYPE f32
#endif
#ifdef DST_F16
#define DST_TYPE f16
#else
#define DST_TYPE f32
#endif
// Default values
// The actual values are defined in shader-lib.
#define HEAD_DIM_QK 64
#define HEAD_DIM_V 64
#define KV_GRANULARITY 8
#define KV_TILE 16
#define WG_SIZE 64
#define KV_BLOCKS (KV_TILE / KV_GRANULARITY)
struct Params {
offset_q: u32,
offset_k: u32,
offset_v: u32,
offset_mask: u32,
offset_sinks: u32,
offset_dst: u32,
// shapes of Q/K/V
n_heads: u32,
seq_len_q: u32,
seq_len_kv: u32,
// strides (in elements)
stride_q1: u32,
stride_q2: u32,
stride_q3: u32,
stride_k1: u32,
stride_k2: u32,
stride_k3: u32,
stride_v1: u32,
stride_v2: u32,
stride_v3: u32,
stride_mask3: u32,
// repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
q_per_kv: u32,
// softmax params
scale: f32,
max_bias: f32,
logit_softcap: f32,
n_head_log2: f32,
m0: f32,
m1: f32,
#ifdef BLK
blk_base: u32,
blk_nblk0: u32,
blk_nblk1: u32,
#endif
tmp_data_base: u32,
tmp_stats_base: u32,
nwg: u32,
};
@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
#ifdef KV_OVERLAP
#if defined(K_Q4_0) || defined(K_Q8_0)
@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
#else
@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
#endif
#define V K
#else
#if defined(K_Q4_0) || defined(K_Q8_0)
@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
#else
@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
#endif
#if defined(V_Q4_0) || defined(V_Q8_0)
@group(0) @binding(2) var<storage, read_write> V: array<V_TYPE>;
#else
@group(0) @binding(2) var<storage, read_write> V: array<vec4<V_TYPE>>;
#endif
#endif
#if defined(MASK) && defined(SINKS)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
#ifdef BLK
#define BLK_BINDING 4
#define TMP_BINDING 5
#define DST_BINDING 6
#define PARAMS_BINDING 7
#else
#define TMP_BINDING 4
#define DST_BINDING 5
#define PARAMS_BINDING 6
#endif
#else
@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
@group(0) @binding(4) var<storage, read_write> sinks: array<f32>;
#ifdef BLK
#define BLK_BINDING 5
#define TMP_BINDING 6
#define DST_BINDING 7
#define PARAMS_BINDING 8
#else
#define TMP_BINDING 5
#define DST_BINDING 6
#define PARAMS_BINDING 7
#endif
#endif
#elif defined(MASK)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
#ifdef BLK
#define BLK_BINDING 3
#define TMP_BINDING 4
#define DST_BINDING 5
#define PARAMS_BINDING 6
#else
#define TMP_BINDING 3
#define DST_BINDING 4
#define PARAMS_BINDING 5
#endif
#else
@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
#ifdef BLK
#define BLK_BINDING 4
#define TMP_BINDING 5
#define DST_BINDING 6
#define PARAMS_BINDING 7
#else
#define TMP_BINDING 4
#define DST_BINDING 5
#define PARAMS_BINDING 6
#endif
#endif
#elif defined(SINKS)
#ifdef KV_OVERLAP
@group(0) @binding(2) var<storage, read_write> sinks: array<f32>;
#define TMP_BINDING 3
#define DST_BINDING 4
#define PARAMS_BINDING 5
#else
@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
#define TMP_BINDING 4
#define DST_BINDING 5
#define PARAMS_BINDING 6
#endif
#else
#ifdef KV_OVERLAP
#define TMP_BINDING 2
#define DST_BINDING 3
#define PARAMS_BINDING 4
#else
#define TMP_BINDING 3
#define DST_BINDING 4
#define PARAMS_BINDING 5
#endif
#endif
#ifdef BLK
@group(0) @binding(BLK_BINDING) var<storage, read_write> blk: array<u32>;
#endif
@group(0) @binding(TMP_BINDING) var<storage, read_write> tmp: array<f32>;
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
// Just a very small float value.
const FLOAT_MIN: f32 = -1.0e9;
const Q_CHUNKS: u32 = HEAD_DIM_QK / 4u;
const V_CHUNKS: u32 = HEAD_DIM_V / 4u;
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
#if defined(K_DIRECT) || defined(V_DIRECT)
// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
// so caching it is more efficient, even on the direct path.
var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
#endif
// K/V shared memory handling
#if !defined(K_DIRECT) || !defined(V_DIRECT)
#define STAGING_SHMEM kv_shmem
#define STAGING_OUT_TYPE f32
#include "flash_attn_staging.tmpl"
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
#endif
var<workgroup> q_shmem: array<f32, HEAD_DIM_QK>;
var<workgroup> o_shmem: array<f32, HEAD_DIM_V>;
// note that we reuse the same storage for both since we only need one at a time
@@ -208,59 +43,6 @@ var<workgroup> inter_shmem: array<f32, KV_TILE>;
var<workgroup> mask_shmem: array<f32, KV_TILE>;
#endif
#if defined(K_DIRECT) || defined(V_DIRECT)
// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
// so caching it is more efficient, even on the direct path.
var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
#endif
// K/V shared memory handling
#if !defined(K_DIRECT) || !defined(V_DIRECT)
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f32
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE * 4u) {
let k_row = elem_idx / HEAD_DIM_QK;
let k_col = elem_idx % HEAD_DIM_QK;
let global_k_row = kv_tile + k_row;
let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
let in_bounds = global_k_row < params.seq_len_kv && (k_col + 3u) < HEAD_DIM_QK;
let vec_idx = (global_k_row_offset + k_col) >> 2u;
let k4 = select(vec4<K_TYPE>(0.0), K[vec_idx], in_bounds);
kv_shmem[elem_idx + 0u] = f32(k4.x);
kv_shmem[elem_idx + 1u] = f32(k4.y);
kv_shmem[elem_idx + 2u] = f32(k4.z);
kv_shmem[elem_idx + 3u] = f32(k4.w);
}
}
#endif
#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE * 4u) {
let v_row = elem_idx / HEAD_DIM_V;
let v_col = elem_idx % HEAD_DIM_V;
let global_v_row = kv_tile + v_row;
let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
let in_bounds = global_v_row < params.seq_len_kv && (v_col + 3u) < HEAD_DIM_V;
let vec_idx = (global_v_row_offset + v_col) >> 2u;
let v4 = select(vec4<V_TYPE>(0.0), V[vec_idx], in_bounds);
kv_shmem[elem_idx + 0u] = f32(v4.x);
kv_shmem[elem_idx + 1u] = f32(v4.y);
kv_shmem[elem_idx + 2u] = f32(v4.z);
kv_shmem[elem_idx + 3u] = f32(v4.w);
}
}
#endif
#endif // !defined(K_DIRECT) || !defined(V_DIRECT)
// Storage for row max and exp sum during online softmax
fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
var v = select(FLOAT_MIN,
+6 -30
View File
@@ -1,19 +1,9 @@
#include "common_decls.tmpl"
enable f16;
@group(0) @binding(0)
#if defined(INPUT_F32)
var<storage, read_write> input: array<f32>;
#elif defined(INPUT_F16)
var<storage, read_write> input: array<f16>;
#endif
var<storage, read_write> input: array<INPUT_TYPE>;
@group(0) @binding(1)
#if defined(OUTPUT_F32)
var<storage, read_write> output: array<f32>;
#elif defined(OUTPUT_F16)
var<storage, read_write> output: array<f16>;
#endif
var<storage, read_write> output: array<OUTPUT_TYPE>;
struct Params {
offset_i: u32,
@@ -38,22 +28,6 @@ struct Params {
@group(0) @binding(2)
var<uniform> params: Params;
fn load_input(idx: u32) -> f32 {
#if defined(INPUT_F32)
return input[idx];
#elif defined(INPUT_F16)
return f32(input[idx]);
#endif
}
fn store_output(idx: u32, val: f32) {
#if defined(OUTPUT_F32)
output[idx] = val;
#elif defined(OUTPUT_F16)
output[idx] = f16(val);
#endif
}
@compute @workgroup_size(WG_SIZE)
fn main(
@builtin(global_invocation_id) gid: vec3<u32>,
@@ -90,12 +64,14 @@ fn main(
let iw_i32 = i32(ow * params.s0 + kw * params.d0) - i32(params.p0);
let ih_i32 = i32(oh * params.s1 + kh * params.d1) - i32(params.p1);
let output_idx = params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3;
if (iw_i32 >= 0 && iw_i32 < i32(params.IW) && ih_i32 >= 0 && ih_i32 < i32(params.IH)) {
let iw = u32(iw_i32);
let ih = u32(ih_i32);
let in_idx = params.offset_i + iw * params.si0 + ih * params.si1 + ic * params.si2 + n * params.si3;
store_output(params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3, load_input(in_idx));
output[output_idx] = OUTPUT_TYPE(input[in_idx]);
} else {
store_output(params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3, 0.0);
output[output_idx] = OUTPUT_TYPE(0.0);
}
}
@@ -88,7 +88,6 @@ struct Params {
ne0: u32,
ne1: u32,
ne2: u32,
ne3: u32,
eps: f32
};
@@ -31,7 +31,6 @@ struct Params {
ne0: u32,
ne1: u32,
ne2: u32,
ne3: u32,
eps: f32
};
+17 -52
View File
@@ -27,7 +27,6 @@ struct Params {
stride_dst3: u32,
// shape of src0/dst
ne: u32,
ne0: u32,
ne1: u32,
ne2: u32,
@@ -43,71 +42,38 @@ struct Params {
m1: f32,
};
@group(0) @binding(0)
#define SRC_BINDING 0
@group(0) @binding(SRC_BINDING)
var<storage, read_write> src: array<f32>;
#ifdef HAS_MASK
#ifdef HAS_SINK
@group(0) @binding(1)
#define MASK_BINDING SRC_BINDING + 1
@group(0) @binding(MASK_BINDING)
var<storage, read_write> mask: array<MaskType>;
@group(0) @binding(2)
var<storage, read_write> sinks: array<f32>;
#ifdef INPLACE
@group(0) @binding(3)
var<uniform> params: Params;
#else
@group(0) @binding(3)
var<storage, read_write> dst: array<f32>;
@group(0) @binding(4)
var<uniform> params: Params;
#define MASK_BINDING SRC_BINDING
#endif
#else
@group(0) @binding(1)
var<storage, read_write> mask: array<MaskType>;
#ifdef INPLACE
@group(0) @binding(2)
var<uniform> params: Params;
#else
@group(0) @binding(2)
var<storage, read_write> dst: array<f32>;
@group(0) @binding(3)
var<uniform> params: Params;
#endif
#endif
#else
#ifdef HAS_SINK
@group(0) @binding(1)
#define SINKS_BINDING MASK_BINDING + 1
@group(0) @binding(SINKS_BINDING)
var<storage, read_write> sinks: array<f32>;
#else
#define SINKS_BINDING MASK_BINDING
#endif
#define DST_BINDING SINKS_BINDING + 1
@group(0) @binding(DST_BINDING)
var<storage, read_write> dst: array<f32>;
#ifdef INPLACE
@group(0) @binding(2)
var<uniform> params: Params;
#define PARAMS_BINDING DST_BINDING
#else
@group(0) @binding(2)
var<storage, read_write> dst: array<f32>;
@group(0) @binding(3)
var<uniform> params: Params;
#define PARAMS_BINDING (DST_BINDING + 1)
#endif
#else
#ifdef INPLACE
@group(0) @binding(1)
@group(0) @binding(PARAMS_BINDING)
var<uniform> params: Params;
#else
@group(0) @binding(1)
var<storage, read_write> dst: array<f32>;
@group(0) @binding(2)
var<uniform> params: Params;
#endif
#endif
#endif
#ifdef INPLACE
fn inter_value(i: u32) -> f32 {
@@ -242,4 +208,3 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
col += WG_SIZE;
}
}
@@ -29,7 +29,6 @@ struct Params {
k: u32,
ne2: u32,
ne3: u32,
};
@group(0) @binding(3)
@@ -39,7 +39,6 @@ struct Params {
n_head: u32,
n_group: u32,
n_seq_tokens: u32,
n_seqs: u32,
y_elems: u32,
};
+30
View File
@@ -164,6 +164,13 @@ class Keys:
NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual"
NORM_BEFORE_FC = "{arch}.norm_before_fc"
class Adapters:
COUNT = "{arch}.adapters.count"
TOKEN_IDS_ACTIVATE = "{arch}.adapters.token_ids_activate"
TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute"
LORA_RANK = "{arch}.adapters.lora_rank"
ROUTER_GAIN = "{arch}.adapters.router_gain"
class Attention:
HEAD_COUNT = "{arch}.attention.head_count"
HEAD_COUNT_KV = "{arch}.attention.head_count_kv"
@@ -527,6 +534,7 @@ class MODEL_ARCH(IntEnum):
GRANITE = auto()
GRANITE_MOE = auto()
GRANITE_HYBRID = auto()
GRANITE_SWITCH = auto()
CHAMELEON = auto()
WAVTOKENIZER_DEC = auto()
PLM = auto()
@@ -1198,6 +1206,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.GRANITE: "granite",
MODEL_ARCH.GRANITE_MOE: "granitemoe",
MODEL_ARCH.GRANITE_HYBRID: "granitehybrid",
MODEL_ARCH.GRANITE_SWITCH: "graniteswitch",
MODEL_ARCH.CHAMELEON: "chameleon",
MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
MODEL_ARCH.PLM: "plm",
@@ -3837,6 +3846,12 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN_SHEXP,
MODEL_TENSOR.FFN_UP_SHEXP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
# NextN/MTP (draft head)
MODEL_TENSOR.ATTN_POST_NORM,
MODEL_TENSOR.NEXTN_EH_PROJ,
MODEL_TENSOR.NEXTN_ENORM,
MODEL_TENSOR.NEXTN_HNORM,
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
],
MODEL_ARCH.EXAONE: [
MODEL_TENSOR.TOKEN_EMBD,
@@ -3972,6 +3987,21 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
MODEL_ARCH.GRANITE_SWITCH: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.OUTPUT,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_QKV,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.FFN_NORM,
MODEL_TENSOR.FFN_GATE,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
MODEL_ARCH.CHAMELEON: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
+15
View File
@@ -906,6 +906,21 @@ class GGUFWriter:
def add_embedding_scale(self, value: float) -> None:
self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value)
def add_adapter_count(self, count: int) -> None:
self.add_uint32(Keys.Adapters.COUNT.format(arch=self.arch), count)
def add_adapter_token_ids_activate(self, ids: Sequence[int]) -> None:
self.add_array(Keys.Adapters.TOKEN_IDS_ACTIVATE.format(arch=self.arch), ids)
def add_adapter_token_ids_substitute(self, ids: Sequence[int]) -> None:
self.add_array(Keys.Adapters.TOKEN_IDS_SUBSTITUTE.format(arch=self.arch), ids)
def add_adapter_lora_rank(self, rank: int) -> None:
self.add_uint32(Keys.Adapters.LORA_RANK.format(arch=self.arch), rank)
def add_adapter_router_gain(self, gain: float) -> None:
self.add_float32(Keys.Adapters.ROUTER_GAIN.format(arch=self.arch), gain)
def add_wkv_head_size(self, size: int) -> None:
self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size)
+1 -1
View File
@@ -5,7 +5,7 @@ import os
import sys
import subprocess
HTTPLIB_VERSION = "refs/tags/v0.52.0"
HTTPLIB_VERSION = "refs/tags/v0.53.0"
vendor = {
"https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp",
+6
View File
@@ -100,6 +100,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_GRANITE, "granite" },
{ LLM_ARCH_GRANITE_MOE, "granitemoe" },
{ LLM_ARCH_GRANITE_HYBRID, "granitehybrid" },
{ LLM_ARCH_GRANITE_SWITCH, "graniteswitch" },
{ LLM_ARCH_CHAMELEON, "chameleon" },
{ LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" },
{ LLM_ARCH_PLM, "plm" },
@@ -220,6 +221,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_TIME_DECAY_EXTRA_DIM, "%s.time_decay_extra_dim" },
{ LLM_KV_RESIDUAL_SCALE, "%s.residual_scale" },
{ LLM_KV_EMBEDDING_SCALE, "%s.embedding_scale" },
{ LLM_KV_ADAPTER_COUNT, "%s.adapters.count" },
{ LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, "%s.adapters.token_ids_activate" },
{ LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, "%s.adapters.token_ids_substitute" },
{ LLM_KV_ADAPTER_LORA_RANK, "%s.adapters.lora_rank" },
{ LLM_KV_ADAPTER_ROUTER_GAIN, "%s.adapters.router_gain" },
{ LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" },
{ LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" },
{ LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" },
+6
View File
@@ -105,6 +105,7 @@ enum llm_arch {
LLM_ARCH_GRANITE,
LLM_ARCH_GRANITE_MOE,
LLM_ARCH_GRANITE_HYBRID,
LLM_ARCH_GRANITE_SWITCH,
LLM_ARCH_CHAMELEON,
LLM_ARCH_WAVTOKENIZER_DEC,
LLM_ARCH_PLM,
@@ -225,6 +226,11 @@ enum llm_kv {
LLM_KV_TIME_DECAY_EXTRA_DIM,
LLM_KV_RESIDUAL_SCALE,
LLM_KV_EMBEDDING_SCALE,
LLM_KV_ADAPTER_COUNT,
LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE,
LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE,
LLM_KV_ADAPTER_LORA_RANK,
LLM_KV_ADAPTER_ROUTER_GAIN,
LLM_KV_TOKEN_SHIFT_COUNT,
LLM_KV_INTERLEAVE_MOE_LAYER_STEP,
LLM_KV_FULL_ATTENTION_INTERVAL,
+2 -1
View File
@@ -3602,8 +3602,9 @@ llama_context * llama_init_from_model(
model->hparams.pooling_type, params.pooling_type);
}
// router_layer >= 0 means n_layer_nextn is repurposed for a router layer, not real MTP
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
model->hparams.n_layer_nextn == 0) {
(model->hparams.n_layer_nextn == 0 || model->hparams.router_layer >= 0)) {
LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__);
return nullptr;
}
+10
View File
@@ -277,6 +277,16 @@ bool llama_hparams::has_kv(uint32_t il) const {
return true;
}
bool llama_hparams::has_rope(uint32_t il) const {
// the router layer stores adapter routing signal, not positional info,
// so it must not be RoPE-shifted
if (router_layer >= 0 && (int32_t) il == router_layer) {
return false;
}
return true;
}
uint32_t llama_hparams::n_layer() const {
return n_layer_all - n_layer_nextn;
}
+6
View File
@@ -53,6 +53,10 @@ struct llama_hparams {
uint32_t n_embd;
uint32_t n_layer_all;
uint32_t n_layer_nextn = 0;
// granite-switch: index of the single-head "router" KV layer that encodes
// per-token adapter selection. -1 when the model has no such layer.
int32_t router_layer = -1;
uint32_t n_expert = 0;
uint32_t n_expert_used = 0;
uint32_t n_rel_attn_bkts = 0;
@@ -371,6 +375,8 @@ struct llama_hparams {
bool has_kv(uint32_t il) const;
bool has_rope(uint32_t il) const;
// number of effective layers (excludes nextn layers)
uint32_t n_layer() const;
+4
View File
@@ -1931,6 +1931,10 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
for (const auto & layer : layers) {
const uint32_t il = layer.il;
if (!hparams.has_rope(il)) {
continue;
}
const int64_t n_head_kv = hparams.n_head_kv(il);
const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
+12 -12
View File
@@ -937,10 +937,11 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
} break;
case GGML_OP_MUL_MAT_ID:
{
const int n_expert_used = hparams.n_expert_used;
GGML_ASSERT(n_expert_used > 0);
ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);
// Used for either MoE expert routing or embedded adapter routing
const int n_ids_used = hparams.router_layer >= 0 ? 1 : hparams.n_expert_used;
GGML_ASSERT(n_ids_used > 0);
ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512);
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512);
op_tensor = ggml_mul_mat_id(ctx, w, b, ids);
} break;
case GGML_OP_ADD:
@@ -1123,15 +1124,14 @@ struct ggml_tensor * llama_model_loader::create_tensor(
return nullptr;
}
// tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID
// tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID;
// embedded-adapter ".lora_a"/".lora_b" tensors are always used with GGML_OP_MUL_MAT_ID
ggml_op op;
bool bias = tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0;
if (bias) {
if (info.op == GGML_OP_MUL_MAT_ID) {
op = GGML_OP_ADD_ID;
} else {
op = GGML_OP_ADD;
}
if (tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0) {
op = info.op == GGML_OP_MUL_MAT_ID ? GGML_OP_ADD_ID : GGML_OP_ADD;
} else if (hparams.router_layer >= 0 && tn.suffix != nullptr &&
(strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0)) {
op = GGML_OP_MUL_MAT_ID;
} else {
op = info.op;
}
+1 -1
View File
@@ -213,7 +213,7 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true);
add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp);
add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp);
add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res);
+9 -2
View File
@@ -234,6 +234,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_granite(params);
case LLM_ARCH_GRANITE_MOE:
return new llama_model_granite_moe(params);
case LLM_ARCH_GRANITE_SWITCH:
return new llama_model_granite_switch(params);
case LLM_ARCH_MINICPM:
return new llama_model_minicpm(params);
case LLM_ARCH_GRANITE_HYBRID:
@@ -1912,6 +1914,7 @@ void llama_model::print_info() const {
arch == LLM_ARCH_GRANITE ||
arch == LLM_ARCH_GRANITE_MOE ||
arch == LLM_ARCH_GRANITE_HYBRID ||
arch == LLM_ARCH_GRANITE_SWITCH ||
arch == LLM_ARCH_NEMOTRON_H_MOE) {
LLAMA_LOG_INFO("%s: f_embedding_scale = %f\n", __func__, hparams.f_embedding_scale);
LLAMA_LOG_INFO("%s: f_residual_scale = %f\n", __func__, hparams.f_residual_scale);
@@ -2228,6 +2231,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
const bool mtp_on_hybrid_nemotron =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE;
if (llm_arch_is_recurrent(arch)) {
res = new llama_memory_recurrent(
*this,
@@ -2238,7 +2244,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
cparams.n_seq_max,
cparams.n_rs_seq,
nullptr);
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) {
// The main difference between hybrid architectures is the
// layer filters, so pick the right one here
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
@@ -2319,7 +2325,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
};
}
if (mtp_on_hybrid_qwen) {
if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
@@ -2596,6 +2602,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_GRANITE:
case LLM_ARCH_GRANITE_MOE:
case LLM_ARCH_GRANITE_HYBRID:
case LLM_ARCH_GRANITE_SWITCH:
case LLM_ARCH_CHAMELEON:
case LLM_ARCH_BAILINGMOE:
case LLM_ARCH_NEO_BERT:
+20
View File
@@ -223,6 +223,24 @@ struct llama_layer_nextn {
struct ggml_tensor * shared_head_norm = nullptr;
};
struct llama_layer_switch_lora {
struct ggml_tensor * a_q = nullptr;
struct ggml_tensor * b_q = nullptr;
struct ggml_tensor * a_k = nullptr;
struct ggml_tensor * b_k = nullptr;
struct ggml_tensor * a_v = nullptr;
struct ggml_tensor * b_v = nullptr;
struct ggml_tensor * a_o = nullptr;
struct ggml_tensor * b_o = nullptr;
struct ggml_tensor * a_gate = nullptr;
struct ggml_tensor * b_gate = nullptr;
struct ggml_tensor * a_up = nullptr;
struct ggml_tensor * b_up = nullptr;
struct ggml_tensor * a_down = nullptr;
struct ggml_tensor * b_down = nullptr;
};
struct llama_layer {
// normalization
struct ggml_tensor * attn_norm = nullptr;
@@ -533,6 +551,8 @@ struct llama_layer {
struct llama_layer_shortconv shortconv;
struct llama_layer_nextn nextn;
struct llama_layer_switch_lora switch_lora;
};
struct llama_device {
+426
View File
@@ -0,0 +1,426 @@
#include "models.h"
#include <cmath>
void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false);
ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false);
bool rope_finetuned = true;
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
hparams.rope_finetuned = rope_finetuned;
switch (hparams.n_layer()) {
case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break;
case 64: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);
ml.get_key(LLM_KV_ADAPTER_COUNT, n_adapters);
ml.get_key(LLM_KV_ADAPTER_LORA_RANK, max_lora_rank);
ml.get_key(LLM_KV_ADAPTER_ROUTER_GAIN, router_gain, /* required */ false);
// bound counts that size tensors
if (n_adapters > 4096) {
throw std::runtime_error(format("graniteswitch: invalid adapter count %u", n_adapters));
}
if (max_lora_rank > 4096) {
throw std::runtime_error(format("graniteswitch: invalid lora rank %u", max_lora_rank));
}
std::vector<llama_token> token_ids;
std::vector<llama_token> substitute_ids;
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, token_ids);
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, substitute_ids);
if (token_ids.size() != n_adapters || substitute_ids.size() != n_adapters) {
throw std::runtime_error(format(
"graniteswitch: adapter token id arrays (%zu activate, %zu substitute) do not match adapter count %u",
token_ids.size(), substitute_ids.size(), n_adapters));
}
adapter_token_to_slot.clear();
adapter_token_to_substitute.clear();
for (uint32_t i = 0; i < n_adapters; ++i) {
// adapter i -> stacked slot i+1 (slot 0 is the base/zero delta)
adapter_token_to_slot[token_ids[i]] = (int32_t) (i + 1);
adapter_token_to_substitute[token_ids[i]] = substitute_ids[i];
}
// extra single-head attention layer at the END (index n_real) holds the router
// K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers
// keep their indices and the KV cache shift/defrag skips the router layer.
// n_layer_nextn is repurposed here (no MTP): it leaks as 1 into the
// llama_model_n_layer_nextn() getter and a re-saved nextn_predict_layers
const uint32_t n_real = hparams.n_layer();
if (n_real >= LLAMA_MAX_LAYERS) {
throw std::runtime_error(format("graniteswitch: block count %u exceeds LLAMA_MAX_LAYERS", n_real));
}
hparams.router_layer = (int32_t) n_real;
hparams.n_layer_all = n_real + 1;
hparams.n_layer_nextn = 1;
hparams.n_head_arr[n_real] = 1;
hparams.n_head_kv_arr[n_real] = 1;
hparams.n_ff_arr[n_real] = 0;
}
void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
const int64_t n_slots = (int64_t) n_adapters + 1; // slot 0 = base/zero delta
const int64_t n_rank = (int64_t) max_lora_rank;
const int64_t n_embd_q = n_embd_head_k * n_head;
const int64_t n_embd_kv = n_embd_k_gqa;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// substitute ids index tok_embd rows directly; range-check against n_vocab
for (const auto & kv : adapter_token_to_substitute) {
const llama_token sub = kv.second;
if (sub < 0 || (int64_t) sub >= n_vocab) {
throw std::runtime_error(format(
"graniteswitch: substitute token id %d out of range [0, %d)", sub, (int) n_vocab));
}
}
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
auto & sl = layer.switch_lora;
sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots}, 0);
sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots}, 0);
sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots}, 0);
sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
}
}
class llm_graph_input_switch : public llm_graph_input_i {
public:
llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {}
virtual ~llm_graph_input_switch() = default;
void set_input(const llama_ubatch * ubatch) override;
ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] adapter-substituted token ids
ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (+/-gain)
ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0)
ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0)
const llama_model_granite_switch & smodel;
};
// K dim-0 is +gain for an adapter token, -gain otherwise; the causal softmax then
// lets a single visible adapter token dominate so the readback recovers its slot.
void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) {
if (!ubatch->token) {
return;
}
const int64_t n_tokens = ubatch->n_tokens;
std::vector<int32_t> sub (n_tokens);
std::vector<float> ksig(n_tokens);
std::vector<float> vval(n_tokens);
std::vector<float> q (n_tokens, 1.0f);
for (int64_t i = 0; i < n_tokens; ++i) {
const llama_token tok = ubatch->token[i];
const auto it = smodel.adapter_token_to_slot.find(tok);
if (it != smodel.adapter_token_to_slot.end()) {
ksig[i] = +smodel.router_gain;
vval[i] = (float) it->second;
} else {
ksig[i] = -smodel.router_gain;
vval[i] = 0.0f;
}
const auto sit = smodel.adapter_token_to_substitute.find(tok);
sub[i] = (sit != smodel.adapter_token_to_substitute.end())
? (int32_t) sit->second
: (int32_t) tok;
}
ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens));
ggml_backend_tensor_set(router_ksig, ksig.data(), 0, n_tokens*ggml_element_size(router_ksig));
ggml_backend_tensor_set(router_vval, vval.data(), 0, n_tokens*ggml_element_size(router_vval));
ggml_backend_tensor_set(router_q, q.data(), 0, n_tokens*ggml_element_size(router_q));
}
std::unique_ptr<llm_graph_context> llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// per-token switched LoRA delta: B_a*(A_a*x), adapter selected per token via ids.
// cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens}
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta(
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids) {
const int64_t n_in = cur->ne[0];
const int64_t n_tokens = cur->ne[1];
ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens);
ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens);
ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens}
ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out, 1, n_tokens}
return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens);
}
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm(
ggml_tensor * w,
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids) {
ggml_tensor * base = ggml_mul_mat(ctx0, w, cur);
ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids);
return ggml_add(ctx0, base, delta);
}
llama_model_granite_switch::graph::graph(
const llama_model & model,
const llm_graph_params & params)
: llm_graph_context(params) {
const auto & smodel = static_cast<const llama_model_granite_switch &>(model);
// TODO: support raw embedding input (multimodal / pre-embedded tokens) when needed
GGML_ASSERT(ubatch.token && "granite-switch requires token input");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
auto inp_switch = std::make_unique<llm_graph_input_switch>(smodel);
inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
inp_switch->router_vval = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
inp_switch->router_q = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
ggml_set_input(inp_switch->sub_tokens);
ggml_set_input(inp_switch->router_ksig);
ggml_set_input(inp_switch->router_vval);
ggml_set_input(inp_switch->router_q);
ggml_tensor * sub_tokens = inp_switch->sub_tokens;
ggml_tensor * router_ksig = inp_switch->router_ksig;
ggml_tensor * router_vval = inp_switch->router_vval;
ggml_tensor * router_q = inp_switch->router_q;
res->add_input(std::move(inp_switch));
// embed the substituted ids directly; build_inp_embd would embed the raw tokens
ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens);
if (hparams.f_embedding_scale != 0.0f) {
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
}
cb(inpL, "inp_embd", -1);
ggml_tensor * inp_pos = nullptr;
if (hparams.rope_finetuned) {
inp_pos = build_inp_pos();
}
auto * inp_attn = build_attn_inp_kv();
// single causal head at layer R recovers the adapter index in-graph: only dim 0
// carries signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), the rest is zero-padded.
const int R = hparams.router_layer;
GGML_ASSERT(R >= 0);
auto router_lane = [&](ggml_tensor * sig1d) {
ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens);
return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0);
};
ggml_tensor * Qr = router_lane(router_q);
ggml_tensor * Kr = router_lane(router_ksig);
ggml_tensor * Vr = router_lane(router_vval);
ggml_tensor * router_out = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R);
cb(router_out, "router_out", R);
// row 0 of router_out is the attended slot; clamp+round to an I32 index
ggml_tensor * slot_f = ggml_cont(ctx0,
ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0));
slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens);
slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters);
slot_f = ggml_round(ctx0, slot_f);
ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32);
cb(adapter_ids, "adapter_ids", -1);
ggml_tensor * inp_out_ids = build_inp_out_ids();
ggml_tensor * cur;
for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
// keep adapter_ids aligned to the kept rows (2D round-trip for get_rows)
const int64_t n_out = inp_out_ids->ne[0];
adapter_ids = ggml_get_rows(ctx0,
ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids);
adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out);
}
cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il);
inpL = cur;
}
cur = inpL;
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * adapter_ids,
llm_graph_input_attn_kv * inp_attn,
const llama_model & model,
const int64_t n_embd_head,
const int il) {
const auto & layer = model.layers[il];
const auto & sl = layer.switch_lora;
const int64_t n_head = hparams.n_head(il);
const int64_t n_head_kv = hparams.n_head_kv(il);
ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur);
cb(qkv, "wqkv", il);
const int64_t n_embd_q = n_embd_head * n_head;
const int64_t n_embd_kv = n_embd_head * n_head_kv;
// slice fused qkv into Q/K/V, made contiguous so LoRA deltas can be added
ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0));
ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv)));
ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv)));
Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids));
Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids));
Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids));
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
if (hparams.rope_finetuned) {
ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
}
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
// wo = nullptr so build_attn returns concatenated heads; o-proj is switched below
ggml_tensor * attn = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(attn, "attn_pre_o", il);
cur = build_switched_lora_mm(layer.wo, sl.a_o, sl.b_o, attn, adapter_ids);
cb(cur, "attn_out", il);
return cur;
}
ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn(
ggml_tensor * cur,
ggml_tensor * inpSA,
ggml_tensor * adapter_ids,
const llama_model & model,
const int il) {
const auto & layer = model.layers[il];
const auto & sl = layer.switch_lora;
if (hparams.f_residual_scale) {
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids);
ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids);
g = ggml_silu(ctx0, g);
ggml_tensor * gu = ggml_mul(ctx0, g, u);
cur = build_switched_lora_mm(layer.ffn_down, sl.a_down, sl.b_down, gu, adapter_ids);
cb(cur, "ffn_out", il);
if (hparams.f_residual_scale) {
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
return cur;
}
+54
View File
@@ -1461,6 +1461,10 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h {
using graph = llama_model_nemotron_h::graph;
struct graph_mtp : public llm_graph_context {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
@@ -1596,6 +1600,56 @@ struct llama_model_granite_moe : public llama_model_base {
};
struct llama_model_granite_switch : public llama_model_base {
llama_model_granite_switch(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
uint32_t n_adapters = 0;
uint32_t max_lora_rank = 0;
float router_gain = 15.0f;
std::unordered_map<llama_token, int32_t> adapter_token_to_slot;
std::unordered_map<llama_token, llama_token> adapter_token_to_substitute;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
private:
ggml_tensor * build_switched_lora_delta(
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids);
ggml_tensor * build_switched_lora_mm(
ggml_tensor * w,
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids);
ggml_tensor * build_attention_layer(
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * adapter_ids,
llm_graph_input_attn_kv * inp_attn,
const llama_model & model,
const int64_t n_embd_head,
const int il);
ggml_tensor * build_layer_ffn(
ggml_tensor * cur,
ggml_tensor * inpSA,
ggml_tensor * adapter_ids,
const llama_model & model,
const int il);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_minicpm : public llama_model_base {
llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
+150
View File
@@ -1,6 +1,156 @@
#include "models.h"
std::unique_ptr<llm_graph_context> llama_model_nemotron_h_moe::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
// MTP draft head for Nemotron-H MoE
llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
const int il = hparams.n_layer();
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm);
GGML_ASSERT(layer.ffn_gate_inp);
// token embedding weights
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings");
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * tok_embd;
if (ubatch.token) {
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
} else {
tok_embd = inp->embd;
}
cb(tok_embd, "mtp_tok_embd", il);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * h_embd = inp->h;
res->add_input(std::move(inp));
ggml_tensor * inp_out_ids = build_inp_out_ids();
// attention fills KV over all tokens, but the MoE is position-wise: gather output rows before
// it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)
const bool emit_h_nextn = cparams.embeddings_nextn;
const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
cb(concat, "mtp_concat", il);
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
cb(cur, "mtp_eh_proj", il);
// dense NoPE attention sub-layer (mtp.layers.0)
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
{
auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "mtp_attn_out", il);
}
cur = ggml_add(ctx0, cur, inpSA);
cb(cur, "mtp_attn_residual", il);
// gather the output rows here so the MoE FFN below only runs on the positions we keep
if (crop_before_ffn) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
// MoE FFN sub-layer (mtp.layers.1)
ggml_tensor * ffn_residual = cur;
cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_post_norm", il);
{
ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur);
cb(router_logits, "mtp_ffn_moe_logits", il);
ggml_tensor * moe_out =
build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
nullptr, // no gate
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_RELU_SQR, hparams.expert_weights_norm,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
il,
router_logits, nullptr,
layer.ffn_up_exps_s,
nullptr, // no gate
layer.ffn_down_exps_s);
cb(moe_out, "mtp_ffn_moe_out", il);
ggml_tensor * ffn_shexp = build_ffn(cur,
layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
NULL, NULL, NULL,
layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
NULL,
LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);
cb(ffn_shexp, "mtp_ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "mtp_ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_residual);
cb(cur, "mtp_post_ffn", il);
// final head norm: the MTP head has its own LayerNorm
GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm");
cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!crop_before_ffn && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
// LM head
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
+73 -23
View File
@@ -7,13 +7,18 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// NextN/MTP: optional draft head appended as extra trailing block(s)
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
// A layer is recurrent IFF the n_head_kv value is set to 0 and
// the n_ff value is set to 0
for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0);
// the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent)
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // MTP head final_layernorm
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
@@ -30,9 +35,13 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
}
}
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
// mamba2 Mixer SSM params
// NOTE: int64_t for tensor dimensions
const int64_t d_conv = hparams.ssm_d_conv;
@@ -60,61 +69,94 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
auto & layer = layers[i];
// all blocks use the attn norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags);
if (hparams.is_recr(i)) {
// ssm layers
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags);
layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags);
// no "weight" suffix for these
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0);
layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags);
layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags);
// out_proj
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags);
} else if (hparams.n_ff(i) == 0) {
// attention layers (with optional bias)
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
} else {
if (n_expert != 0) {
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags);
// MoE branch
layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, trunk_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags);
// Shared expert branch
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, trunk_flags);
} else {
// mlp layers
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, trunk_flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, trunk_flags);
layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED);
}
}
}
// NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE
// sub-layer into a single trailing block
for (int i = n_layer; i < n_layer_all; ++i) {
auto & layer = layers[i];
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
// NextN input-fusion tensors
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags);
// attention sub-layer
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED);
// MoE sub-layer
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, mtp_flags);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, mtp_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, mtp_flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, mtp_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags);
}
}
std::unique_ptr<llm_graph_context> llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const {
@@ -153,7 +195,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cur = build_ffn_layer(cur, model, il);
}
if (il == n_layer - 1 && inp_out_ids) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -170,6 +212,14 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
// seed for the MTP/NextN draft head
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;
+4
View File
@@ -217,6 +217,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
if (moe) {
ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff);
ms.add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, n_ff / 2); // distinct from n_ff so a saver key-clobber surfaces on reload
ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, uint32_t(2));
ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(2));
ms.add_kv(LLM_KV_EXPERT_USED_COUNT, uint32_t(1));
@@ -410,6 +411,9 @@ static bool arch_supported(const llm_arch arch) {
if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
return false; // FIXME @ngxson
}
if (arch == LLM_ARCH_GRANITE_SWITCH) {
return false; // FIXME adapter fixture
}
if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
return false; // FIXME Embedding (?) models produce inconsistent results.
}
+10 -5
View File
@@ -13,6 +13,9 @@ PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..
# marker for the grep_search test to find in this file
GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search"
# image the container runtime tests run their shell in
DOCKER_IMAGE = "busybox"
@pytest.fixture(autouse=True)
def create_server():
@@ -149,14 +152,16 @@ def test_tools_builtin_cwd_header():
def _docker_unavailable_reason() -> str | None:
"""None if docker can be used to run a container, otherwise the reason it can't."""
"""None if docker can run the image these tests use, otherwise the reason it can't."""
docker_bin = shutil.which("docker")
if docker_bin is None:
return "docker is not installed"
try:
subprocess.run([docker_bin, "info"], capture_output=True, timeout=5, check=True)
# a daemon that answers `docker info` still cannot run a linux image when it serves
# windows containers, so probe the image itself, which also pulls it before the tests
subprocess.run([docker_bin, "run", "--rm", DOCKER_IMAGE, "true"], capture_output=True, timeout=60, check=True)
except Exception as e:
return f"docker daemon is not usable: {e}"
return f"docker cannot run {DOCKER_IMAGE}: {e}"
return None
@@ -167,7 +172,7 @@ def docker_container():
pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type]
proc = subprocess.run(
["docker", "run", "-d", "--rm", "busybox", "sleep", "300"],
["docker", "run", "-d", "--rm", DOCKER_IMAGE, "sleep", "300"],
capture_output=True, text=True,
)
if proc.returncode != 0:
@@ -214,7 +219,7 @@ def test_tools_builtin_docker_runtime_cleans_up_spawned_container():
pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type]
global server
server.server_tools_runtime = "docker:busybox"
server.server_tools_runtime = f"docker:{DOCKER_IMAGE}"
server.start()
# exec_shell_command runs inside the container spawned for --tools-runtime; docker sets
+50 -14
View File
@@ -1,14 +1,15 @@
// For more info, see https://github.com/storybookjs/eslint-plugin-storybook#configuration-flat-config-format
import storybook from 'eslint-plugin-storybook';
import prettier from 'eslint-config-prettier';
import svelteConfig from './svelte.config.js';
import { includeIgnoreFile } from '@eslint/compat';
import js from '@eslint/js';
import prettier from 'eslint-config-prettier';
import perfectionist from 'eslint-plugin-perfectionist';
import simpleImportSort from 'eslint-plugin-simple-import-sort';
import storybook from 'eslint-plugin-storybook';
import svelte from 'eslint-plugin-svelte';
import globals from 'globals';
import { fileURLToPath } from 'node:url';
import ts from 'typescript-eslint';
import svelteConfig from './svelte.config.js';
const gitignorePath = fileURLToPath(new URL('./.gitignore', import.meta.url));
@@ -21,32 +22,67 @@ export default ts.config(
...svelte.configs.prettier,
{
languageOptions: { globals: { ...globals.browser, ...globals.node } },
plugins: { perfectionist, 'simple-import-sort': simpleImportSort },
rules: {
// typescript-eslint strongly recommend that you do not use the no-undef lint rule on TypeScript projects.
// see: https://typescript-eslint.io/troubleshooting/faqs/eslint/#i-get-errors-from-the-no-undef-rule-about-global-variables-not-being-defined-even-though-there-are-no-typescript-errors
'no-undef': 'off',
'svelte/no-at-html-tags': 'off',
// This app uses hash-based routing (#/) where resolve() from $app/paths does not apply
'svelte/no-navigation-without-resolve': 'off',
// Snippet bodies often ignore one or more of the parent's params
// (e.g. `{#snippet children(_meta, ctx)}` when only ctx is read).
'@typescript-eslint/no-unused-vars': [
'error',
{ argsIgnorePattern: '^_', varsIgnorePattern: '^_' }
],
// Enforce empty line at end of file
'eol-last': 'error'
'eol-last': 'error',
// typescript-eslint strongly recommend that you do not use the no-undef lint rule on TypeScript projects.
// see: https://typescript-eslint.io/troubleshooting/faqs/eslint/#i-get-errors-from-the-no-undef-rule-about-global-variables-not-being-defined-even-though-there-are-no-typescript-errors
'no-undef': 'off',
'padding-line-between-statements': [
'error',
// Blank line between function/class declarations.
{ blankLine: 'always', next: ['function', 'class'], prev: ['function', 'class'] },
// Blank line around if blocks (if/else and else if stay one statement).
{ blankLine: 'always', next: '*', prev: 'if' },
{ blankLine: 'always', next: 'if', prev: '*' },
// Blank line after the last declaration in a group. Because the 'never'
// rules below are scoped per declaration kind, a const group and a let
// group get separated by a blank line, while same-kind declarations stay
// together.
{ blankLine: 'always', next: '*', prev: ['const', 'let', 'var'] },
// No blank line between consecutive declarations of the same kind (kept
// last so each takes precedence over the always rule above for matching
// declaration pairs).
{ blankLine: 'never', next: 'const', prev: 'const' },
{ blankLine: 'never', next: 'let', prev: 'let' },
{ blankLine: 'never', next: 'var', prev: 'var' },
// Blank line before a statement that follows another statement in the block
// (works for return/throw/break/continue). A blank line for a terminal
// statement that opens a block body can't be enforced here: Prettier removes
// the leading blank line of a block, so the two formatters would fight.
{ blankLine: 'always', next: ['return', 'throw', 'break', 'continue'], prev: '*' }
],
'perfectionist/sort-objects': ['error', { type: 'natural' }],
// Alphabetical order for variable declarations and object keys
'perfectionist/sort-variable-declarations': ['error', { type: 'natural' }],
// Sort imports alphabetically by module path, and sort named members within
// each statement. A single catch-all group keeps the list flat (no blank-line
// grouping); Prettier normalizes comma spacing afterwards.
'simple-import-sort/imports': ['error', { groups: [['.*']] }],
'svelte/no-at-html-tags': 'off',
// This app uses hash-based routing (#/) where resolve() from $app/paths does not apply
'svelte/no-navigation-without-resolve': 'off'
}
},
{
files: ['**/*.svelte', '**/*.svelte.ts', '**/*.svelte.js'],
languageOptions: {
parserOptions: {
projectService: true,
extraFileExtensions: ['.svelte'],
parser: ts.parser,
projectService: true,
svelteConfig
}
}
+232
View File
@@ -39,6 +39,8 @@
"dompurify": "3.4.13",
"eslint": "9.39.4",
"eslint-config-prettier": "10.1.8",
"eslint-plugin-perfectionist": "^5.10.1",
"eslint-plugin-simple-import-sort": "^14.0.0",
"eslint-plugin-storybook": "10.5.6",
"eslint-plugin-svelte": "3.19.0",
"fflate": "0.8.3",
@@ -9281,6 +9283,226 @@
"eslint": ">=7.0.0"
}
},
"node_modules/eslint-plugin-perfectionist": {
"version": "5.10.1",
"resolved": "https://registry.npmjs.org/eslint-plugin-perfectionist/-/eslint-plugin-perfectionist-5.10.1.tgz",
"integrity": "sha512-Kprsp9Us0GqAesYaAIzUViw57xYp5WBqzXrcE0Mtww++E5fexWXYBipMuuD7yvyH4vvpBH0+oJ+OMAmZ0oYXkw==",
"dev": true,
"license": "MIT",
"dependencies": {
"@typescript-eslint/utils": "^8.65.0",
"natural-orderby": "^5.0.0"
},
"engines": {
"node": "^20.0.0 || >=22.0.0"
},
"peerDependencies": {
"eslint": "^8.45.0 || ^9.0.0 || ^10.0.0"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/project-service": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/project-service/-/project-service-8.66.0.tgz",
"integrity": "sha512-7MthGPTt4BP69lSryqpqq8HQqxuzynssckL/jyDyk3+TNMQ3y2jFWkptCrktWvBrP+EH787Nl5N5Qpw7WZg+5g==",
"dev": true,
"license": "MIT",
"dependencies": {
"@typescript-eslint/tsconfig-utils": "^8.66.0",
"@typescript-eslint/types": "^8.66.0",
"debug": "^4.4.3"
},
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
},
"peerDependencies": {
"typescript": ">=4.8.4 <6.1.0"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/scope-manager": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/scope-manager/-/scope-manager-8.66.0.tgz",
"integrity": "sha512-8TGcH25j9zqJ/IULB/ppyhRvxA8QYfFEZ7nfbg6/BN9spDgb8fPWQXlE5l8TWBL50EtUx007uZ1o9VOwrq2/9g==",
"dev": true,
"license": "MIT",
"dependencies": {
"@typescript-eslint/types": "8.66.0",
"@typescript-eslint/visitor-keys": "8.66.0"
},
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/tsconfig-utils": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/tsconfig-utils/-/tsconfig-utils-8.66.0.tgz",
"integrity": "sha512-9D5gLYZG4rOjcoag8MQ/fWI8WqA9wcPDyOGyWtWFhvM1lHRbliqUSPIY5J3zqCU1tvSwzXxnnjhQhz5Ne7mJ4g==",
"dev": true,
"license": "MIT",
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
},
"peerDependencies": {
"typescript": ">=4.8.4 <6.1.0"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/types": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/types/-/types-8.66.0.tgz",
"integrity": "sha512-H6gcYaSDOyvL3AD/jHUtUFo2jqGgn/F6nuyuZSu0QTesxL+cP4dQoIMrODRofuJC09g64+WgZ6tE19Y1N2YIFQ==",
"dev": true,
"license": "MIT",
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/typescript-estree": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/typescript-estree/-/typescript-estree-8.66.0.tgz",
"integrity": "sha512-8/x4INiiQb10jGgXYD7116/zQ+OL84ZIFn0za68wwFHCanT/VLbBEroWht8RV8fn0/ZCAoazHLQgwUC0UQcDfg==",
"dev": true,
"license": "MIT",
"dependencies": {
"@typescript-eslint/project-service": "8.66.0",
"@typescript-eslint/tsconfig-utils": "8.66.0",
"@typescript-eslint/types": "8.66.0",
"@typescript-eslint/visitor-keys": "8.66.0",
"debug": "^4.4.3",
"minimatch": "^10.2.2",
"semver": "^7.7.3",
"tinyglobby": "^0.2.15",
"ts-api-utils": "^2.5.0"
},
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
},
"peerDependencies": {
"typescript": ">=4.8.4 <6.1.0"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/utils": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.66.0.tgz",
"integrity": "sha512-jasearZPolBw5NJNYGMwxzHMF83niVWmMU1VdHzG1CyfI2VS7f7nZltnKtHcg20hW+7Uo5GfK4MeDPoU3qI8EA==",
"dev": true,
"license": "MIT",
"dependencies": {
"@eslint-community/eslint-utils": "^4.9.1",
"@typescript-eslint/scope-manager": "8.66.0",
"@typescript-eslint/types": "8.66.0",
"@typescript-eslint/typescript-estree": "8.66.0"
},
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
},
"peerDependencies": {
"eslint": "^8.57.0 || ^9.0.0 || ^10.0.0",
"typescript": ">=4.8.4 <6.1.0"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/visitor-keys": {
"version": "8.66.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/visitor-keys/-/visitor-keys-8.66.0.tgz",
"integrity": "sha512-dkKR8q+lKciskj1Y3vthHktl+3cMLWGyVUP23bRiPZ5O9BRT++4EqDDV+TVeIKBL1VXVEqrJlz8MYbcnvJcAlg==",
"dev": true,
"license": "MIT",
"dependencies": {
"@typescript-eslint/types": "8.66.0",
"eslint-visitor-keys": "^5.0.0"
},
"engines": {
"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/typescript-eslint"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/balanced-match": {
"version": "4.0.4",
"resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz",
"integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==",
"dev": true,
"license": "MIT",
"engines": {
"node": "18 || 20 || >=22"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/brace-expansion": {
"version": "5.0.9",
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.9.tgz",
"integrity": "sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==",
"dev": true,
"license": "MIT",
"dependencies": {
"balanced-match": "^4.0.2"
},
"engines": {
"node": "20 || >=22"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/eslint-visitor-keys": {
"version": "5.0.1",
"resolved": "https://registry.npmjs.org/eslint-visitor-keys/-/eslint-visitor-keys-5.0.1.tgz",
"integrity": "sha512-tD40eHxA35h0PEIZNeIjkHoDR4YjjJp34biM0mDvplBe//mB+IHCqHDGV7pxF+7MklTvighcCPPZC7ynWyjdTA==",
"dev": true,
"license": "Apache-2.0",
"engines": {
"node": "^20.19.0 || ^22.13.0 || >=24"
},
"funding": {
"url": "https://opencollective.com/eslint"
}
},
"node_modules/eslint-plugin-perfectionist/node_modules/minimatch": {
"version": "10.2.6",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.6.tgz",
"integrity": "sha512-vpLQEs+VLCr1nU0BXS07maYoFwlDAH0gngQuuttxIwutDFEMHq2blX+8vpgxDdK3J1PwjCJiep77OitTZ4Ll1A==",
"dev": true,
"license": "BlueOak-1.0.0",
"dependencies": {
"brace-expansion": "^5.0.8"
},
"engines": {
"node": "18 || 20 || >=22"
},
"funding": {
"url": "https://github.com/sponsors/isaacs"
}
},
"node_modules/eslint-plugin-simple-import-sort": {
"version": "14.0.0",
"resolved": "https://registry.npmjs.org/eslint-plugin-simple-import-sort/-/eslint-plugin-simple-import-sort-14.0.0.tgz",
"integrity": "sha512-NUJO0+XFCkk+o5EsAJruTgnfMEpeWrPWeJS15UVF60GgXmqz1BJ9/3hzlvG7lkL8Bubzos5cCLptThbFfPnSMQ==",
"dev": true,
"license": "MIT",
"peerDependencies": {
"eslint": ">=5.0.0"
}
},
"node_modules/eslint-plugin-storybook": {
"version": "10.5.6",
"resolved": "https://registry.npmjs.org/eslint-plugin-storybook/-/eslint-plugin-storybook-10.5.6.tgz",
@@ -13196,6 +13418,16 @@
"dev": true,
"license": "MIT"
},
"node_modules/natural-orderby": {
"version": "5.0.0",
"resolved": "https://registry.npmjs.org/natural-orderby/-/natural-orderby-5.0.0.tgz",
"integrity": "sha512-kKHJhxwpR/Okycz4HhQKKlhWe4ASEfPgkSWNmKFHd7+ezuQlxkA5cM3+XkBPvm1gmHen3w53qsYAv+8GwRrBlg==",
"dev": true,
"license": "MIT",
"engines": {
"node": ">=18"
}
},
"node_modules/negotiator": {
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/negotiator/-/negotiator-1.0.0.tgz",
+5 -3
View File
@@ -12,7 +12,7 @@
"check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json",
"check:watch": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json --watch",
"reset": "rm -rf .svelte-kit node_modules",
"format": "prettier --write .",
"format": "eslint --fix . && prettier --write .",
"lint": "prettier --check . && eslint .",
"test": "npm run test:ui -- --run && npm run test:client -- --run && npm run test:unit -- --run && npm run test:e2e",
"test:e2e": "playwright test",
@@ -36,6 +36,7 @@
"@playwright/test": "1.56.1",
"@storybook/addon-a11y": "10.5.6",
"@storybook/addon-docs": "10.5.6",
"@storybook/addon-mcp": "0.7.0",
"@storybook/addon-svelte-csf": "5.1.2",
"@storybook/addon-vitest": "10.5.6",
"@storybook/sveltekit": "10.5.6",
@@ -57,6 +58,8 @@
"dompurify": "3.4.13",
"eslint": "9.39.4",
"eslint-config-prettier": "10.1.8",
"eslint-plugin-perfectionist": "^5.10.1",
"eslint-plugin-simple-import-sort": "^14.0.0",
"eslint-plugin-storybook": "10.5.6",
"eslint-plugin-svelte": "3.19.0",
"fflate": "0.8.3",
@@ -99,8 +102,7 @@
"vite-plugin-devtools-json": "0.2.1",
"vitest": "4.1.10",
"vitest-browser-svelte": "2.1.1",
"workbox-window": "7.4.1",
"@storybook/addon-mcp": "0.7.0"
"workbox-window": "7.4.1"
},
"overrides": {
"cookie": "1.1.1",
+14 -14
View File
@@ -1,31 +1,31 @@
import { defineConfig, devices } from '@playwright/test';
export default defineConfig({
testDir: 'tests/e2e',
testMatch: ['**/*.e2e.ts'],
timeout: 30000,
expect: {
timeout: 5000
},
fullyParallel: true,
forbidOnly: !!process.env.CI,
retries: process.env.CI ? 2 : 0,
workers: process.env.CI ? 1 : undefined,
reporter: 'line',
use: {
baseURL: 'http://localhost:8181',
trace: 'on-first-retry'
},
fullyParallel: true,
projects: [
{
name: 'chromium',
use: { ...devices['Desktop Chrome'] }
}
],
reporter: 'line',
retries: process.env.CI ? 2 : 0,
testDir: 'tests/e2e',
testMatch: ['**/*.e2e.ts'],
timeout: 30000,
use: {
baseURL: 'http://localhost:8181',
trace: 'on-first-retry'
},
webServer: {
command: 'npm run build && npx http-server ./dist -p 8181',
port: 8181,
timeout: 120000,
reuseExistingServer: !process.env.CI
}
reuseExistingServer: !process.env.CI,
timeout: 120000
},
workers: process.env.CI ? 1 : undefined
});
+9 -9
View File
@@ -1,6 +1,6 @@
import { defineConfig } from '@vite-pwa/assets-generator/config';
import { FAVICON_COLORS, PWA_ASSET_GENERATOR } from './src/lib/constants/pwa';
import { writeThemeFavicons } from './scripts/favicon-colorize';
import { FAVICON_COLORS, PWA_ASSET_GENERATOR } from './src/lib/constants/pwa';
import { defineConfig } from '@vite-pwa/assets-generator/config';
writeThemeFavicons(FAVICON_COLORS.LIGHT, FAVICON_COLORS.DARK, {
padding: PWA_ASSET_GENERATOR.FAVICON_PADDING
@@ -10,18 +10,18 @@ export default defineConfig({
headLinkOptions: {
preset: '2023'
},
images: ['static/favicon-dark.svg'],
preset: {
transparent: {
sizes: [],
favicons: [[48, 'favicon-dark.ico']],
padding: PWA_ASSET_GENERATOR.FAVICON_PADDING
apple: {
sizes: []
},
maskable: {
sizes: []
},
apple: {
transparent: {
favicons: [[48, 'favicon-dark.ico']],
padding: PWA_ASSET_GENERATOR.FAVICON_PADDING,
sizes: []
}
},
images: ['static/favicon-dark.svg']
}
});
+25 -24
View File
@@ -1,3 +1,11 @@
import { writeThemeFavicons } from './scripts/favicon-colorize';
import {
FAVICON_COLORS,
PWA_ASSET_GENERATOR,
PWA_GENERATOR_DEVICES,
THEME_COLORS
} from './src/lib/constants/pwa';
import { SplashOrientation } from './src/lib/enums/splash.enums';
import {
combinePresetAndAppleSplashScreens,
defineConfig,
@@ -5,14 +13,6 @@ import {
} from '@vite-pwa/assets-generator/config';
import { readFileSync } from 'node:fs';
import { resolve } from 'node:path';
import {
THEME_COLORS,
PWA_GENERATOR_DEVICES,
PWA_ASSET_GENERATOR,
FAVICON_COLORS
} from './src/lib/constants/pwa';
import { SplashOrientation } from './src/lib/enums/splash.enums';
import { writeThemeFavicons } from './scripts/favicon-colorize';
writeThemeFavicons(FAVICON_COLORS.LIGHT, FAVICON_COLORS.DARK, {
padding: PWA_ASSET_GENERATOR.FAVICON_PADDING
@@ -22,6 +22,7 @@ export default defineConfig({
headLinkOptions: {
preset: PWA_ASSET_GENERATOR.LINK_PRESET
},
images: ['static/favicon.svg'],
preset: combinePresetAndAppleSplashScreens(
{
...minimal2023Preset,
@@ -32,37 +33,37 @@ export default defineConfig({
}
},
{
padding: PWA_ASSET_GENERATOR.SPLASH_PADDING,
resizeOptions: {
background: THEME_COLORS.BACKGROUND_LIGHT,
fit: PWA_ASSET_GENERATOR.FIT_MODE
},
darkResizeOptions: {
background: THEME_COLORS.BACKGROUND_DARK,
fit: PWA_ASSET_GENERATOR.FIT_MODE
},
darkImageResolver: async (imageName: string) => {
if (imageName.endsWith('favicon.svg')) {
return readFileSync(resolve('static/favicon-dark.svg'));
}
},
darkResizeOptions: {
background: THEME_COLORS.BACKGROUND_DARK,
fit: PWA_ASSET_GENERATOR.FIT_MODE
},
linkMediaOptions: {
log: true,
addMediaScreen: PWA_ASSET_GENERATOR.ADD_MEDIA_SCREEN,
basePath: PWA_ASSET_GENERATOR.BASE_PATH,
log: true,
xhtml: PWA_ASSET_GENERATOR.XHTML
},
png: {
compressionLevel: PWA_ASSET_GENERATOR.PNG_COMPRESSION_LEVEL,
quality: PWA_ASSET_GENERATOR.PNG_QUALITY
},
name: (landscape, size, dark) => {
const orientation = landscape ? SplashOrientation.LANDSCAPE : SplashOrientation.PORTRAIT;
const darkPrefix = dark ? PWA_ASSET_GENERATOR.DARK_PREFIX : '';
return `apple-splash-${orientation}-${darkPrefix}${size.width}x${size.height}.png`;
},
padding: PWA_ASSET_GENERATOR.SPLASH_PADDING,
png: {
compressionLevel: PWA_ASSET_GENERATOR.PNG_COMPRESSION_LEVEL,
quality: PWA_ASSET_GENERATOR.PNG_QUALITY
},
resizeOptions: {
background: THEME_COLORS.BACKGROUND_LIGHT,
fit: PWA_ASSET_GENERATOR.FIT_MODE
}
},
PWA_GENERATOR_DEVICES
),
images: ['static/favicon.svg']
)
});
+16 -10
View File
@@ -4,12 +4,10 @@ import { fileURLToPath } from 'node:url';
const HERE = dirname(fileURLToPath(import.meta.url));
const PROJECT_ROOT = resolve(HERE, '..');
const DEFAULT_LOGO = resolve(PROJECT_ROOT, 'src/lib/assets/logo.svg');
const DEFAULT_OUT_DIR = resolve(PROJECT_ROOT, 'static');
const DEFAULT_OUT_LIGHT = resolve(DEFAULT_OUT_DIR, 'favicon.svg');
const DEFAULT_OUT_DARK = resolve(DEFAULT_OUT_DIR, 'favicon-dark.svg');
const CURRENT_COLOR = 'currentColor';
export interface ColorizedFavicon {
@@ -39,8 +37,8 @@ export function colorizeFaviconSvg(
darkColor: string
): ColorizedFavicon {
return {
light: svg.replaceAll(CURRENT_COLOR, lightColor),
dark: svg.replaceAll(CURRENT_COLOR, darkColor)
dark: svg.replaceAll(CURRENT_COLOR, darkColor),
light: svg.replaceAll(CURRENT_COLOR, lightColor)
};
}
@@ -54,33 +52,40 @@ export function padFaviconSvg(svg: string, padding: number): string {
if (!(padding > 0) || padding >= 1) return svg;
const viewBoxMatch = svg.match(/viewBox\s*=\s*["']([^"']+)["']/i);
if (!viewBoxMatch) return svg;
const parts = viewBoxMatch[1]
.trim()
.split(/[\s,]+/)
.map(Number);
if (parts.length !== 4 || parts.some((n) => !Number.isFinite(n))) return svg;
const [, , width, height] = parts;
if (width <= 0 || height <= 0) return svg;
const scale = 1 - padding;
const translateX = (padding * width) / 2;
const translateY = (padding * height) / 2;
const openTagStart = svg.search(/<svg\b/i);
if (openTagStart === -1) return svg;
const openTagEnd = svg.indexOf('>', openTagStart);
if (openTagEnd === -1) return svg;
const closeStart = svg.lastIndexOf('</svg');
if (closeStart === -1 || closeStart <= openTagEnd) return svg;
const openTag = svg.slice(0, openTagEnd + 1);
const inner = svg.slice(openTagEnd + 1, closeStart);
const closeTag = svg.slice(closeStart);
const group = `<g transform="translate(${translateX} ${translateY}) scale(${scale})">`;
return `${openTag}${group}${inner}</g>${closeTag}`;
}
@@ -93,14 +98,15 @@ export function writeThemeFavicons(
lightColor: string,
darkColor: string,
{
sourcePath = DEFAULT_LOGO,
lightOutPath = DEFAULT_OUT_LIGHT,
darkOutPath = DEFAULT_OUT_DARK,
padding = 0
lightOutPath = DEFAULT_OUT_LIGHT,
padding = 0,
sourcePath = DEFAULT_LOGO
}: WriteThemeFaviconsOptions = {}
): void {
const source = readFileSync(sourcePath, 'utf-8');
const { light, dark } = colorizeFaviconSvg(source, lightColor, darkColor);
const { dark, light } = colorizeFaviconSvg(source, lightColor, darkColor);
mkdirSync(dirname(lightOutPath), { recursive: true });
writeFileSync(lightOutPath, padFaviconSvg(light, padding));
writeFileSync(darkOutPath, padFaviconSvg(dark, padding));
+19 -17
View File
@@ -13,31 +13,28 @@
* maskable-icon and apple-touch-icon are left untouched.
*/
import sharp from 'sharp';
import fs from 'fs';
import path from 'path';
import sharp from 'sharp';
import { fileURLToPath } from 'url';
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
const STATIC_DIR = path.resolve(__dirname, '..', 'static');
const paddingPct = process.argv.reduce((acc, arg, i, args) => {
if (arg === '--padding-pct' && args[i + 1]) return parseFloat(args[i + 1]);
return acc;
}, 0);
// Scale down the source image before cropping to circle
const scalePct = process.argv.reduce((acc, arg, i, args) => {
if (arg === '--scale-pct' && args[i + 1]) return parseFloat(args[i + 1]);
return acc;
}, 85); // default 85% - icon fills 85% of the circular area
// Source for circular icons: the maskable icon (white bg, full logo)
const sourceIcon = 'maskable-icon-512x512.png';
const targetIcons = ['pwa-64x64.png', 'pwa-192x192.png', 'pwa-512x512.png'];
// maskable-icon and apple-touch-icon stay square
const untouchedIcons = ['maskable-icon-512x512.png', 'apple-touch-icon-180x180.png'];
@@ -47,10 +44,13 @@ async function makeCircle(targetFilename) {
if (!fs.existsSync(sourcePath)) {
console.log(`⏭️ ${sourceIcon} not found, skipping`);
return;
}
if (!fs.existsSync(targetPath)) {
console.log(`⏭️ ${targetFilename} not found, skipping`);
return;
}
@@ -58,16 +58,18 @@ async function makeCircle(targetFilename) {
const size = Math.max(metadata.width, metadata.height);
const radius = Math.floor((size * (1 - paddingPct / 100)) / 2);
const center = Math.floor(size / 2);
// Build circular mask as RGBA buffer: white opaque circle on transparent bg
const maskBuf = Buffer.alloc(size * size * 4, 0);
for (let y = 0; y < size; y++) {
for (let x = 0; x < size; x++) {
const dx = x - center;
const dy = y - center;
const dist = Math.sqrt(dx * dx + dy * dy);
if (dist < radius) {
const i = (y * size + x) * 4;
maskBuf[i] = 255;
maskBuf[i + 1] = 255;
maskBuf[i + 2] = 255;
@@ -77,8 +79,9 @@ async function makeCircle(targetFilename) {
}
const tmpMask = path.join(STATIC_DIR, '.mask-tmp.png');
await sharp(maskBuf, {
raw: { width: size, height: size, channels: 4 }
raw: { channels: 4, height: size, width: size }
})
.png()
.toFile(tmpMask);
@@ -87,28 +90,26 @@ async function makeCircle(targetFilename) {
const circleDiameter = Math.floor(size * (1 - paddingPct / 100));
const scaledSize = Math.floor((circleDiameter * scalePct) / 100);
const offset = Math.floor((size - scaledSize) / 2);
const scaledBuf = await sharp(sourcePath)
.resize(scaledSize, scaledSize, {
fit: 'cover',
background: { r: 255, g: 255, b: 255, alpha: 1 }
background: { alpha: 1, b: 255, g: 255, r: 255 },
fit: 'cover'
})
.ensureAlpha()
.png()
.toBuffer();
// Step 2: Composite scaled image onto white background, then apply circular mask
const output = await sharp({
create: {
width: size,
height: size,
background: { alpha: 1, b: 255, g: 255, r: 255 },
channels: 4,
background: { r: 255, g: 255, b: 255, alpha: 1 }
height: size,
width: size
}
})
.composite([
{ input: scaledBuf, top: offset, left: offset },
{ input: tmpMask, top: 0, left: 0, blend: 'dest-in' }
{ input: scaledBuf, left: offset, top: offset },
{ blend: 'dest-in', input: tmpMask, left: 0, top: 0 }
])
.png()
.toBuffer();
@@ -130,6 +131,7 @@ async function main() {
console.log('\nUnchanged:');
for (const icon of untouchedIcons) {
const fp = path.join(STATIC_DIR, icon);
console.log(` ${icon} (${fs.existsSync(fp) ? fs.statSync(fp).size + ' bytes' : 'missing'})`);
}
}
+7 -5
View File
@@ -1,7 +1,7 @@
import { writeFileSync, existsSync } from 'node:fs';
import { BUILD_CONFIG } from '../src/lib/constants/pwa';
import { existsSync, writeFileSync } from 'node:fs';
import { resolve } from 'path';
import type { Plugin } from 'vite';
import { BUILD_CONFIG } from '../src/lib/constants/pwa';
let processed = false;
@@ -15,27 +15,29 @@ const OUTPUT_DIR = process.env.LLAMA_UI_OUT_DIR ?? BUILD_CONFIG.OUTPUT_DIR;
*/
export function buildInfoPlugin(): Plugin {
return {
name: 'llamacpp:build-info',
apply: 'build',
closeBundle() {
setTimeout(() => {
try {
if (processed) return;
processed = true;
const buildNumber = process.env.LLAMA_BUILD_NUMBER || 'b0000';
const outDir = resolve(OUTPUT_DIR);
const indexPath = resolve(outDir, 'index.html');
if (!existsSync(indexPath)) return;
const buildJsonPath = resolve(outDir, 'build.json');
writeFileSync(buildJsonPath, JSON.stringify({ version: buildNumber }), 'utf-8');
console.log(`Created build.json (version: ${buildNumber})`);
} catch (error) {
console.error('Failed to write build.json:', error);
}
}, 100);
}
},
name: 'llamacpp:build-info'
};
}
+14 -12
View File
@@ -4,7 +4,6 @@ import { fileURLToPath } from 'url';
import type { Plugin } from 'vite';
const __dirname = dirname(fileURLToPath(import.meta.url));
const VENDORS_DIR = resolve(__dirname, '../src/lib/vendors');
const VIRTUAL_ID = 'virtual:nerdamer';
const RESOLVED_ID = '\0' + VIRTUAL_ID;
@@ -21,29 +20,32 @@ export function nerdamerPlugin(): Plugin {
let bundled: string | null = null;
return {
name: 'llamacpp:nerdamer',
resolveId(id) {
return id === VIRTUAL_ID ? RESOLVED_ID : undefined;
},
async load(id) {
if (id !== RESOLVED_ID) return undefined;
if (bundled === null) {
const result = await build({
entryPoints: [resolve(VENDORS_DIR, 'nerdamer-prime/all.js')],
bundle: true,
minify: true,
format: 'iife',
globalName: 'nerdamer',
alias: {
'big-integer': resolve(VENDORS_DIR, 'big-integer/BigInteger.js'),
'decimal.js': resolve(VENDORS_DIR, 'decimal.js/decimal.js')
},
write: false,
logLevel: 'silent'
bundle: true,
entryPoints: [resolve(VENDORS_DIR, 'nerdamer-prime/all.js')],
format: 'iife',
globalName: 'nerdamer',
logLevel: 'silent',
minify: true,
write: false
});
bundled = result.outputFiles[0].text;
}
return `export default ${JSON.stringify(bundled)};`;
},
name: 'llamacpp:nerdamer',
resolveId(id) {
return id === VIRTUAL_ID ? RESOLVED_ID : undefined;
}
};
}
@@ -1,7 +1,7 @@
import { readFileSync, writeFileSync, existsSync } from 'node:fs';
import { BUILD_CONFIG } from '../src/lib/constants/pwa';
import { existsSync, readFileSync, writeFileSync } from 'node:fs';
import { resolve } from 'path';
import type { Plugin } from 'vite';
import { BUILD_CONFIG } from '../src/lib/constants/pwa';
let processed = false;
@@ -11,11 +11,15 @@ function rewrite(path: string, pairs: [string, string][]): void {
if (!existsSync(path)) {
return;
}
const text = readFileSync(path, 'utf-8');
let out = text;
for (const [from, to] of pairs) {
out = out.split(from).join(to);
}
if (out !== text) {
writeFileSync(path, out, 'utf-8');
}
@@ -32,12 +36,12 @@ function rewrite(path: string, pairs: [string, string][]): void {
*/
export function relativizeBasePlugin(): Plugin {
return {
name: 'llamacpp:relativize-base',
apply: 'build',
closeBundle() {
setTimeout(() => {
try {
if (processed) return;
processed = true;
const outDir = resolve(OUTPUT_DIR);
@@ -56,6 +60,7 @@ export function relativizeBasePlugin(): Plugin {
console.error('Failed to relativize base refs:', error);
}
}, 100);
}
},
name: 'llamacpp:relativize-base'
};
}
+23 -13
View File
@@ -1,10 +1,10 @@
import { readdirSync, readFileSync, writeFileSync, existsSync } from 'node:fs';
import { NEWLINE, TAB } from '../src/lib/constants/code';
import { APPLE_DEVICES, BUILD_CONFIG, REGEX_PATTERNS, SPLASH_LINK } from '../src/lib/constants/pwa';
import { SplashOrientation } from '../src/lib/enums/splash.enums';
import type { SplashDimensions } from '../src/lib/types';
import { existsSync, readdirSync, readFileSync, writeFileSync } from 'node:fs';
import { resolve } from 'path';
import type { Plugin } from 'vite';
import { TAB, NEWLINE } from '../src/lib/constants/code';
import { APPLE_DEVICES, BUILD_CONFIG, REGEX_PATTERNS, SPLASH_LINK } from '../src/lib/constants/pwa';
import type { SplashDimensions } from '../src/lib/types';
import { SplashOrientation } from '../src/lib/enums/splash.enums';
let processed = false;
@@ -16,23 +16,26 @@ const OUTPUT_DIR = process.env.LLAMA_UI_OUT_DIR ?? BUILD_CONFIG.OUTPUT_DIR;
*/
export function generateSplashScreenLinks(outDir: string): string[] {
const files = readdirSync(outDir).filter((f) => f.match(REGEX_PATTERNS.SPLASH_FILE));
if (files.length === 0) return [];
const dimMap = new Map<string, SplashDimensions>();
for (const [dims, spec] of Object.entries(APPLE_DEVICES)) {
const [w, h] = dims.split('x').map(Number);
// logical-point dimensions
dimMap.set(`${w}x${h}`, { deviceW: spec.width, deviceH: spec.height, dpr: spec.dpr });
dimMap.set(`${h}x${w}`, { deviceW: spec.width, deviceH: spec.height, dpr: spec.dpr });
dimMap.set(`${w}x${h}`, { deviceH: spec.height, deviceW: spec.width, dpr: spec.dpr });
dimMap.set(`${h}x${w}`, { deviceH: spec.height, deviceW: spec.width, dpr: spec.dpr });
// pixel dimensions (used by actual generated splash files)
dimMap.set(`${w * spec.dpr}x${h * spec.dpr}`, {
deviceW: spec.width,
deviceH: spec.height,
deviceW: spec.width,
dpr: spec.dpr
});
dimMap.set(`${h * spec.dpr}x${w * spec.dpr}`, {
deviceW: spec.width,
deviceH: spec.height,
deviceW: spec.width,
dpr: spec.dpr
});
}
@@ -42,20 +45,23 @@ export function generateSplashScreenLinks(outDir: string): string[] {
for (const file of files) {
const match = file.match(REGEX_PATTERNS.SPLASH_FILE);
if (!match) continue;
const orientation = match[1] as SplashOrientation;
const isDark = !!match[2];
const pixelW = parseInt(match[3]);
const pixelH = parseInt(match[4]);
const key = `${pixelW}x${pixelH}`;
const spec = dimMap.get(key);
if (!spec) {
console.warn(`Unknown splash screen dimensions: ${key} (${file})`);
continue;
}
const { deviceW, deviceH, dpr } = spec;
const { deviceH, deviceW, dpr } = spec;
const media = `screen and (device-width: ${deviceW}px) and (device-height: ${deviceH}px) and (-webkit-device-pixel-ratio: ${dpr}) and (orientation: ${orientation})`;
const href = `./${file}`;
@@ -73,16 +79,17 @@ export function generateSplashScreenLinks(outDir: string): string[] {
export function splashScreenPlugin(): Plugin {
return {
name: 'llamacpp:splash-screen',
apply: 'build',
closeBundle() {
setTimeout(() => {
try {
if (processed) return;
processed = true;
const outDir = resolve(OUTPUT_DIR);
const indexPath = resolve(outDir, 'index.html');
if (!existsSync(indexPath)) return;
let content = readFileSync(indexPath, 'utf-8');
@@ -91,9 +98,11 @@ export function splashScreenPlugin(): Plugin {
// The @vite-pwa/assets-generator generates apple-splash-*.png files;
// this scans them and creates the <link> tags SvelteKit needs.
const splashLinks = generateSplashScreenLinks(outDir);
if (splashLinks.length > 0) {
console.log(`Generated ${splashLinks.length} apple-splash link tags`);
const splashHtml = splashLinks.map((l) => TAB + TAB + l).join(NEWLINE);
content = content.replace(
REGEX_PATTERNS.HEAD_CLOSE,
splashHtml + NEWLINE + TAB + TAB + '</head>'
@@ -110,6 +119,7 @@ export function splashScreenPlugin(): Plugin {
console.error('Failed to process build output:', error);
}
}, 100);
}
},
name: 'llamacpp:splash-screen'
};
}
+14 -17
View File
@@ -3,9 +3,8 @@
import 'vite-plugin-pwa/pwa-assets';
import 'vite-plugin-pwa/svelte';
import { ModelModality, ServerModelStatus, ServerRole } from '$lib/enums';
// Import chat types from dedicated module
import type {
// API types
ApiChatCompletionRequest,
@@ -13,59 +12,57 @@ import type {
ApiChatCompletionStreamChunk,
ApiChatCompletionToolCall,
ApiChatCompletionToolCallDelta,
ApiChatMessageData,
ApiChatMessageContentPart,
ApiChatMessageData,
ApiContextSizeError,
ApiErrorResponse,
ApiLlamaCppServerProps,
ApiModelDataEntry,
ApiModelListResponse,
ApiModelLoadStage,
ApiModelsSseProgress,
ApiModelsSseData,
ApiModelsSseEvent,
ApiModelListResponse,
ApiModelsSseProgress,
ApiProcessingState,
ApiRouterModelMeta,
ApiRouterModelsListResponse,
ApiRouterModelsLoadRequest,
ApiRouterModelsLoadResponse,
ApiRouterModelsStatusRequest,
ApiRouterModelsStatusResponse,
ApiRouterModelsListResponse,
ApiRouterModelsUnloadRequest,
ApiRouterModelsUnloadResponse,
// Chat types
ChatAttachmentDisplayItem,
ChatMessagePromptProgress,
ChatMessageSiblingInfo,
ChatMessageTimings,
ChatMessageType,
ChatRole,
ChatUploadedFile,
ChatMessageSiblingInfo,
ChatMessagePromptProgress,
ChatMessageTimings,
// Database types
DatabaseConversation,
DatabaseMessage,
DatabaseMessageExtra,
DatabaseMessageExtraAudioFile,
DatabaseMessageExtraVideoFile,
DatabaseMessageExtraImageFile,
DatabaseMessageExtraTextFile,
DatabaseMessageExtraPdfFile,
DatabaseMessageExtraLegacyContext,
DatabaseMessageExtraPdfFile,
DatabaseMessageExtraTextFile,
DatabaseMessageExtraVideoFile,
ExportedConversation,
ExportedConversations,
ModelLoadProgress,
// Model types
ModelModalities,
ModelOption,
ModelLoadProgress,
// Settings types
SettingsChatServiceOptions,
SettingsConfigType,
SettingsConfigValue,
SettingsFieldConfig,
SettingsConfigType
SettingsFieldConfig
} from '$lib/types';
import { ServerRole, ServerModelStatus, ModelModality } from '$lib/enums';
declare global {
// namespace App {
// interface Error {}
@@ -1,8 +1,8 @@
<script lang="ts">
import { Button, type ButtonVariant, type ButtonSize } from '$lib/components/ui/button';
import { Button, type ButtonSize, type ButtonVariant } from '$lib/components/ui/button';
import * as Tooltip from '$lib/components/ui/tooltip';
import type { Component } from 'svelte';
import { TooltipSide } from '$lib/enums';
import type { Component } from 'svelte';
interface Props {
ariaLabel?: string;
@@ -20,18 +20,18 @@
}
let {
icon,
tooltip,
variant = 'ghost',
href = '',
size = 'sm',
ariaLabel,
class: className = '',
disabled = false,
href = '',
icon,
iconSize = 'h-3 w-3',
tooltipSide = TooltipSide.TOP,
stopPropagationOnClick = false,
onclick,
ariaLabel
size = 'sm',
stopPropagationOnClick = false,
tooltip,
tooltipSide = TooltipSide.TOP,
variant = 'ghost'
}: Props = $props();
let innerWidth = $state(0);
@@ -1,8 +1,8 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Copy } from '@lucide/svelte';
import { copyToClipboard } from '$lib/utils';
import ActionIcon from './ActionIcon.svelte';
import { Copy } from '@lucide/svelte';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { copyToClipboard } from '$lib/utils';
export let ariaLabel: string = 'Copy to clipboard';
export let canCopy: boolean = true;
@@ -7,7 +7,7 @@
class?: string;
}
let { modalities, class: className = '' }: Props = $props();
let { class: className = '', modalities }: Props = $props();
</script>
{#each modalities as modality (modality)}
@@ -28,18 +28,18 @@
}
let {
class: className = '',
style = '',
activeModelId,
attachments = [],
readonly = false,
onFileRemove,
uploadedFiles = $bindable([]),
class: className = '',
// Default to small size for form previews
imageClass = '',
imageHeight = 'h-24',
imageWidth = 'w-auto',
limitToSingleRow = false,
activeModelId
onFileRemove,
readonly = false,
style = '',
uploadedFiles = $bindable([])
}: Props = $props();
let carouselRef: HorizontalScrollCarousel | undefined = $state();
@@ -48,7 +48,7 @@
let previewFocusIndex = $state(0);
let viewAllDialogOpen = $state(false);
let displayItems = $derived(getAttachmentDisplayItems({ uploadedFiles, attachments }));
let displayItems = $derived(getAttachmentDisplayItems({ attachments, uploadedFiles }));
function openPreview(item: ChatAttachmentDisplayItem, event?: MouseEvent) {
event?.stopPropagation();
@@ -2,8 +2,8 @@
import {
ChatAttachmentsListItemMcpPrompt,
ChatAttachmentsListItemMcpResource,
ChatAttachmentsListItemThumbnailImage,
ChatAttachmentsListItemThumbnailFile
ChatAttachmentsListItemThumbnailFile,
ChatAttachmentsListItemThumbnailImage
} from '$lib/components/app';
import { AttachmentType } from '$lib/enums';
import type {
@@ -49,10 +49,10 @@
return {
id,
resource: {
uri: extra.uri,
name: extra.name,
serverName: extra.serverName,
title: extra.name,
serverName: extra.serverName
uri: extra.uri
}
};
}
@@ -64,12 +64,12 @@
? (item.attachment as DatabaseMessageExtraMcpPrompt)
: item.uploadedFile?.mcpPrompt
? {
type: AttachmentType.MCP_PROMPT as const,
name: item.name,
serverName: item.uploadedFile.mcpPrompt.serverName,
promptName: item.uploadedFile.mcpPrompt.promptName,
arguments: item.uploadedFile.mcpPrompt.arguments,
content: item.textContent ?? '',
arguments: item.uploadedFile.mcpPrompt.arguments
name: item.name,
promptName: item.uploadedFile.mcpPrompt.promptName,
serverName: item.uploadedFile.mcpPrompt.serverName,
type: AttachmentType.MCP_PROMPT as const
}
: null}
{#if mcpPrompt}
@@ -1,8 +1,8 @@
<script lang="ts">
import { ChatMessageMcpPromptContent, ActionIcon } from '$lib/components/app';
import { X } from '@lucide/svelte';
import type { DatabaseMessageExtraMcpPrompt } from '$lib/types';
import { ActionIcon, ChatMessageMcpPromptContent } from '$lib/components/app';
import { McpPromptVariant } from '$lib/enums';
import type { DatabaseMessageExtraMcpPrompt } from '$lib/types';
interface Props {
class?: string;
@@ -1,11 +1,11 @@
<script lang="ts">
import { Loader2, AlertCircle } from '@lucide/svelte';
import { AlertCircle, Loader2 } from '@lucide/svelte';
import { X } from '@lucide/svelte';
import { ActionIcon } from '$lib/components/app';
import * as Tooltip from '$lib/components/ui/tooltip';
import { mcpStore } from '$lib/stores/mcp.svelte';
import type { MCPResourceAttachment } from '$lib/types';
import * as Tooltip from '$lib/components/ui/tooltip';
import { ActionIcon } from '$lib/components/app';
import { X } from '@lucide/svelte';
import { getResourceIcon, getResourceDisplayName } from '$lib/utils';
import { getResourceDisplayName, getResourceIcon } from '$lib/utils';
interface Props {
attachment: MCPResourceAttachment;
@@ -24,6 +24,7 @@
function getStatusClass(attachment: MCPResourceAttachment): string {
if (attachment.error) return 'border-red-500/50 bg-red-500/10';
if (attachment.loading) return 'border-border/50 bg-muted/30';
return 'border-border/50 bg-muted/30';
@@ -1,17 +1,17 @@
<script lang="ts">
import { Music, Video, X } from '@lucide/svelte';
import { ActionIcon } from '$lib/components/app';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { X, Music, Video } from '@lucide/svelte';
import { AttachmentType } from '$lib/enums';
import {
formatFileSize,
getFileTypeLabel,
getPreviewText,
isPdfFile,
isAudioFile,
isVideoFile,
isTextFile
isPdfFile,
isTextFile,
isVideoFile
} from '$lib/utils';
import { ActionIcon } from '$lib/components/app';
import { AttachmentType } from '$lib/enums';
interface Props {
attachment?: DatabaseMessageExtra;
@@ -31,9 +31,9 @@
attachment,
class: className = '',
id,
name,
onclick,
onRemove,
name,
readonly = false,
size,
textContent,
@@ -1,6 +1,6 @@
<script lang="ts">
import { ActionIcon } from '$lib/components/app';
import { X } from '@lucide/svelte';
import { ActionIcon } from '$lib/components/app';
interface Props {
class?: string;
@@ -20,9 +20,9 @@
height = 'h-16',
id,
imageClass = '',
name,
onclick,
onRemove,
name,
preview,
readonly = false,
width = 'w-auto'
@@ -12,12 +12,12 @@
getAttachmentDisplayItems,
getLanguageFromFilename,
isAudioFile,
isVideoFile,
isImageFile,
isMcpPrompt,
isMcpResource,
isPdfFile,
isTextFile
isTextFile,
isVideoFile
} from '$lib/utils';
interface PreviewItem {
@@ -42,21 +42,21 @@
}
let {
uploadedFiles = [],
attachments = [],
activeModelId,
attachments = [],
class: className = '',
previewFocusIndex = 0
previewFocusIndex = 0,
uploadedFiles = []
}: Props = $props();
let allItems = $derived(
getAttachmentDisplayItems({ uploadedFiles, attachments })
getAttachmentDisplayItems({ attachments, uploadedFiles })
.filter((item) => !isMcpPrompt(item) && !isMcpResource(item))
.map(
(item): PreviewItem => ({
...item,
isImage: isImageFile(item.attachment, item.uploadedFile),
isAudio: isAudioFile(item.attachment, item.uploadedFile),
isImage: isImageFile(item.attachment, item.uploadedFile),
isVideo: isVideoFile(item.attachment, item.uploadedFile)
})
)
@@ -88,10 +88,11 @@
$effect(() => {
const index = currentIndex;
setTimeout(() => {
const thumbnail = document.querySelector(`[data-thumbnail-index="${index}"]`);
thumbnail?.scrollIntoView({ behavior: 'smooth', inline: 'center', block: 'nearest' });
thumbnail?.scrollIntoView({ behavior: 'smooth', block: 'nearest', inline: 'center' });
}, 0);
});
@@ -1,12 +1,12 @@
<script lang="ts">
import type { ChatAttachmentDisplayItem } from '$lib/types';
import { Image, Music, Video, FileText, FileIcon } from '@lucide/svelte';
import ChatAttachmentsPreviewCurrentItemPdf from './ChatAttachmentsPreviewCurrentItemPdf.svelte';
import ChatAttachmentsPreviewCurrentItemImage from './ChatAttachmentsPreviewCurrentItemImage.svelte';
import ChatAttachmentsPreviewCurrentItemAudio from './ChatAttachmentsPreviewCurrentItemAudio.svelte';
import ChatAttachmentsPreviewCurrentItemVideo from './ChatAttachmentsPreviewCurrentItemVideo.svelte';
import ChatAttachmentsPreviewCurrentItemImage from './ChatAttachmentsPreviewCurrentItemImage.svelte';
import ChatAttachmentsPreviewCurrentItemPdf from './ChatAttachmentsPreviewCurrentItemPdf.svelte';
import ChatAttachmentsPreviewCurrentItemText from './ChatAttachmentsPreviewCurrentItemText.svelte';
import ChatAttachmentsPreviewCurrentItemUnavailable from './ChatAttachmentsPreviewCurrentItemUnavailable.svelte';
import ChatAttachmentsPreviewCurrentItemVideo from './ChatAttachmentsPreviewCurrentItemVideo.svelte';
import { FileIcon, FileText, Image, Music, Video } from '@lucide/svelte';
import type { ChatAttachmentDisplayItem } from '$lib/types';
interface Props {
currentItem: ChatAttachmentDisplayItem | null;
@@ -25,19 +25,19 @@
}
let {
activeModelId,
audioSrc,
currentItem,
isImage,
isAudio,
isVideo,
isPdf,
isText,
displayPreview,
displayTextContent,
audioSrc,
videoSrc,
language,
hasVisionModality,
activeModelId
isAudio,
isImage,
isPdf,
isText,
isVideo,
language,
videoSrc
}: Props = $props();
let IconComponent = $derived(
@@ -6,7 +6,7 @@
audioSrc: string | null;
}
let { currentItem, audioSrc }: Props = $props();
let { audioSrc, currentItem }: Props = $props();
</script>
<div class="flex flex-1 items-center justify-center p-8">
@@ -1,13 +1,13 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import type { ChatAttachmentDisplayItem } from '$lib/types';
import { FileText, Eye, Info } from '@lucide/svelte';
import { Button } from '$lib/components/ui/button';
import * as Alert from '$lib/components/ui/alert';
import { Eye, FileText, Info } from '@lucide/svelte';
import { SyntaxHighlightedCode } from '$lib/components/app';
import * as Alert from '$lib/components/ui/alert';
import { Button } from '$lib/components/ui/button';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { PdfViewMode } from '$lib/enums';
import type { ChatAttachmentDisplayItem } from '$lib/types';
import { getLanguageFromFilename } from '$lib/utils';
import { convertPDFToImage } from '$lib/utils/browser-only';
import { PdfViewMode } from '$lib/enums';
interface Props {
currentItem: ChatAttachmentDisplayItem | null;
@@ -17,7 +17,7 @@
activeModelId?: string;
}
let { currentItem, displayName, displayTextContent, hasVisionModality, activeModelId }: Props =
let { activeModelId, currentItem, displayName, displayTextContent, hasVisionModality }: Props =
$props();
let pdfViewMode = $state<PdfViewMode>(PdfViewMode.PAGES);
@@ -47,6 +47,7 @@
currentItem.attachment.images.length > 0
) {
pdfImages = currentItem.attachment.images;
return;
}
@@ -55,10 +56,12 @@
const base64Data = currentItem.attachment.base64Data;
const byteCharacters = atob(base64Data);
const byteNumbers = new Array(byteCharacters.length);
for (let i = 0; i < byteCharacters.length; i++) {
byteNumbers[i] = byteCharacters.charCodeAt(i);
}
const byteArray = new Uint8Array(byteNumbers);
file = new File([byteArray], displayName, { type: 'application/pdf' });
}
}
@@ -8,7 +8,7 @@
show: boolean;
}
let { onPrev, onNext, show }: Props = $props();
let { onNext, onPrev, show }: Props = $props();
</script>
{#if show}
@@ -1,7 +1,7 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Music, Video, FileText } from '@lucide/svelte';
import { FileText, Music, Video } from '@lucide/svelte';
import { HorizontalScrollCarousel } from '$lib/components/app/misc';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
interface PreviewItem {
id: string;
@@ -18,13 +18,15 @@
onNavigate: (index: number) => void;
}
let { items, currentIndex, onNavigate }: Props = $props();
let { currentIndex, items, onNavigate }: Props = $props();
function getFileExtension(name: string): string {
const parts = name.split('.');
if (parts.length > 1) {
return parts.pop()?.toUpperCase() ?? '';
}
return '';
}
</script>
@@ -1,4 +1,5 @@
<script lang="ts">
import ContextGaugePopup from './ChatFormContextGauge/ContextGaugePopup.svelte';
import {
ChatAttachmentsList,
ChatFormActions,
@@ -12,10 +13,10 @@
} from '$lib/components/app';
import {
CLIPBOARD_CONTENT_QUOTE_PREFIX,
INPUT_CLASSES,
SETTING_CONFIG_DEFAULT,
INITIAL_FILE_SIZE,
PROMPT_CONTENT_SEPARATOR
INPUT_CLASSES,
PROMPT_CONTENT_SEPARATOR,
SETTING_CONFIG_DEFAULT
} from '$lib/constants';
import {
ContentPartType,
@@ -24,20 +25,20 @@
MimeTypeText,
SpecialFileType
} from '$lib/enums';
import { config } from '$lib/stores/settings.svelte';
import ContextGaugePopup from './ChatFormContextGauge/ContextGaugePopup.svelte';
import { modelOptions, selectedModelId } from '$lib/stores/models.svelte';
import { isRouterMode } from '$lib/stores/server.svelte';
import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte';
import { chatStore } from '$lib/stores/chat.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { mcpHasResourceAttachments } from '$lib/stores/mcp-resources.svelte';
import { toolsStore } from '$lib/stores/tools.svelte';
import {
conversationsStore,
activeMessages,
activeConversation,
activeMessages,
conversationsStore,
pendingCwd
} from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { mcpHasResourceAttachments } from '$lib/stores/mcp-resources.svelte';
import { modelOptions, selectedModelId } from '$lib/stores/models.svelte';
import { isRouterMode } from '$lib/stores/server.svelte';
import { config } from '$lib/stores/settings.svelte';
import { toolsStore } from '$lib/stores/tools.svelte';
import type {
FileMentionEntry,
GetPromptResult,
@@ -56,7 +57,6 @@
parseClipboardContent,
uuid
} from '$lib/utils';
import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte';
import {
AudioRecorder,
convertToWav,
@@ -96,12 +96,6 @@
class: className = '',
disabled = false,
isLoading = false,
placeholder = 'Type a message...',
showMcpPromptButton = false,
showAddButton = true,
showModelSelector = true,
uploadedFiles = $bindable([]),
value = $bindable(''),
onAttachmentRemove,
onFilesAdd,
onStop,
@@ -109,7 +103,13 @@
onSystemPromptClick,
onUploadedFileRemove,
onUploadedFilesChange,
onValueChange
onValueChange,
placeholder = 'Type a message...',
showAddButton = true,
showMcpPromptButton = false,
showModelSelector = true,
uploadedFiles = $bindable([]),
value = $bindable('')
}: Props = $props();
// Component References
@@ -146,32 +146,35 @@
let cwd = $derived(activeConversation()?.cwd ?? pendingCwd());
const pickers = useChatFormPickers({
focusInput: refocusInput,
getCaretOffset: () => inputRef?.getCaretOffset(),
getCwd: () => cwd,
getPickersRef: () => pickersRef,
getServerHome: () => toolsStore.serverHome ?? null,
getShowModelSelector: () => showModelSelector,
getValue: () => value,
hasCwdTools: () => toolsStore.hasEnabledCwdTools,
hasPrompts: () => mcpStore.hasPromptsCapability(conversationsStore.getAllMcpServerOverrides()),
openModelSelector: () => chatFormActionsRef?.openModelSelector(),
setCaretOffset: (offset) => inputRef?.setCaretOffset(offset),
setValue: (v) => {
value = v;
onValueChange?.(v);
},
getCaretOffset: () => inputRef?.getCaretOffset(),
setCaretOffset: (offset) => inputRef?.setCaretOffset(offset),
focusInput: refocusInput,
getShowModelSelector: () => showModelSelector,
hasPrompts: () => mcpStore.hasPromptsCapability(conversationsStore.getAllMcpServerOverrides()),
hasCwdTools: () => toolsStore.hasEnabledCwdTools,
getCwd: () => cwd,
getServerHome: () => toolsStore.serverHome ?? null,
openModelSelector: () => chatFormActionsRef?.openModelSelector(),
getPickersRef: () => pickersRef
}
});
async function handleWorkingDirectoryChange(newDir: string | null) {
// Committing a directory consumes the `/cwd` token; the chip's
// clear-X path has no token to consume.
const token = findCommandToken(value);
if (token && token.name === 'cwd') {
value = '';
onValueChange?.('');
}
await conversationsStore.setCwd(newDir);
if (conversationsStore.activeConversation) {
await chatStore.recordCwdChange(newDir?.trim() || null);
}
@@ -185,6 +188,7 @@
let pasteLongTextToFileLength = $derived.by(() => {
const n = Number(currentConfig.pasteLongTextToFileLen);
return Number.isNaN(n) ? Number(SETTING_CONFIG_DEFAULT.pasteLongTextToFileLen) : n;
});
@@ -200,13 +204,16 @@
}
const selectedId = selectedModelId();
if (selectedId) {
const model = options.find((m) => m.id === selectedId);
if (model) return model.model;
}
if (conversationModel) {
const model = options.find((m) => m.model === conversationModel);
if (model) return model.model;
}
@@ -238,6 +245,7 @@
$effect(() => {
const wantContenteditable =
containsFileMentionLink(value ?? '') || containsCodeSpan(value ?? '');
if (useContenteditable === wantContenteditable) return;
if (!caretOffsetPinned) {
@@ -268,8 +276,10 @@
export function checkModelSelected(): boolean {
if (!hasModelSelected) {
chatFormActionsRef?.openModelSelector();
return false;
}
return true;
}
@@ -284,6 +294,7 @@
function handleFileRemove(fileId: string) {
if (fileId.startsWith('attachment-')) {
const index = parseInt(fileId.replace('attachment-', ''), 10);
if (!isNaN(index) && index >= 0 && index < attachments.length) {
onAttachmentRemove?.(index);
}
@@ -333,6 +344,7 @@
if (files.length > 0) {
event.preventDefault();
onFilesAdd?.(files);
return;
}
@@ -354,26 +366,27 @@
type: MimeTypeText.PLAIN
})
);
onFilesAdd?.(attachmentFiles);
}
// Handle MCP prompt attachments as ChatUploadedFile with mcpPrompt data
if (parsed.mcpPromptAttachments.length > 0) {
const mcpPromptFiles: ChatUploadedFile[] = parsed.mcpPromptAttachments.map((att) => ({
id: uuid(),
name: att.name,
size: att.content.length,
type: SpecialFileType.MCP_PROMPT,
file: new File([att.content], `${att.name}${FileExtensionText.TXT}`, {
type: MimeTypeText.PLAIN
}),
id: uuid(),
isLoading: false,
textContent: att.content,
mcpPrompt: {
serverName: att.serverName,
arguments: att.arguments,
promptName: att.promptName,
arguments: att.arguments
}
serverName: att.serverName
},
name: att.name,
size: att.content.length,
textContent: att.content,
type: SpecialFileType.MCP_PROMPT
}));
uploadedFiles = [...uploadedFiles, ...mcpPromptFiles];
@@ -412,17 +425,17 @@
const promptName = promptInfo.title || promptInfo.name;
const placeholder: ChatUploadedFile = {
id: placeholderId,
name: promptName,
size: INITIAL_FILE_SIZE,
type: SpecialFileType.MCP_PROMPT,
file: new File([], 'loading'),
id: placeholderId,
isLoading: true,
mcpPrompt: {
serverName: promptInfo.serverName,
arguments: args ? { ...args } : undefined,
promptName: promptInfo.name,
arguments: args ? { ...args } : undefined
}
serverName: promptInfo.serverName
},
name: promptName,
size: INITIAL_FILE_SIZE,
type: SpecialFileType.MCP_PROMPT
};
uploadedFiles = [...uploadedFiles, placeholder];
@@ -450,12 +463,12 @@
f.id === placeholderId
? {
...f,
isLoading: false,
textContent: promptText,
size: promptText.length,
file: new File([promptText], `${f.name}${FileExtensionText.TXT}`, {
type: MimeTypeText.PLAIN
})
}),
isLoading: false,
size: promptText.length,
textContent: promptText
}
: f
);
@@ -480,9 +493,11 @@
function handleMentionSelect(entry: FileMentionEntry) {
const cursor = inputRef?.getCaretOffset() ?? value.length;
const token = findMentionToken(value, cursor);
if (!token) return;
const built = buildMentionInsertion(entry, value, token);
if (!built) return;
// Pin the post-insertion caret BEFORE the swap effect runs;
@@ -504,6 +519,7 @@
async function handleMicClick() {
if (!audioRecorder || !recordingSupported) {
console.warn('Audio recording not supported');
return;
}
@@ -642,7 +658,7 @@
onFileUpload={handleFileUpload}
onMicClick={handleMicClick}
{onStop}
onSystemPromptClick={() => onSystemPromptClick?.({ message: value, files: uploadedFiles })}
onSystemPromptClick={() => onSystemPromptClick?.({ files: uploadedFiles, message: value })}
onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined}
onMcpResourcesClick={() => (isResourceDialogOpen = true)}
/>
@@ -1,9 +1,9 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Plus } from '@lucide/svelte';
import { Button } from '$lib/components/ui/button';
import * as Tooltip from '$lib/components/ui/tooltip';
import { ATTACHMENT_TOOLTIP_TEXT } from '$lib/constants';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
interface Props {
disabled?: boolean;
@@ -1,20 +1,20 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Plus, File, MessageSquare, Zap, FolderOpen } from '@lucide/svelte';
import { File, FolderOpen, MessageSquare, Plus, Zap } from '@lucide/svelte';
import {
ChatFormActionAddMcpServersSubmenu,
ChatFormActionAddReasoningSubmenu,
ChatFormActionAddToolsSubmenu
} from '$lib/components/app';
import { buttonVariants } from '$lib/components/ui/button';
import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
import * as Tooltip from '$lib/components/ui/tooltip';
import { buttonVariants } from '$lib/components/ui/button';
import { cn } from '$lib/components/ui/utils';
import {
ATTACHMENT_FILE_ITEMS,
ATTACHMENT_TOOLTIP_TEXT,
TOOLTIP_DELAY_DURATION
} from '$lib/constants';
import {
ChatFormActionAddToolsSubmenu,
ChatFormActionAddMcpServersSubmenu,
ChatFormActionAddReasoningSubmenu
} from '$lib/components/app';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
interface Props {
@@ -36,15 +36,15 @@
class: className = '',
disabled = false,
hasAudioModality = false,
hasVideoModality = false,
hasVisionModality = false,
hasMcpPromptsSupport = false,
hasMcpResourcesSupport = false,
hasVideoModality = false,
hasVisionModality = false,
onFileUpload,
onSystemPromptClick,
onMcpPromptClick,
onMcpResourcesClick,
onMcpSettingsClick,
onMcpResourcesClick
onSystemPromptClick
}: Props = $props();
let dropdownOpen = $state(false);
@@ -59,13 +59,13 @@
const attachmentMenu = useAttachmentMenu(
() => ({
hasVisionModality,
hasAudioModality,
hasVideoModality,
hasMcpPromptsSupport,
hasMcpResourcesSupport
hasMcpResourcesSupport,
hasVideoModality,
hasVisionModality
}),
() => ({ onFileUpload, onSystemPromptClick, onMcpPromptClick, onMcpResourcesClick }),
() => ({ onFileUpload, onMcpPromptClick, onMcpResourcesClick, onSystemPromptClick }),
() => {
dropdownOpen = false;
}
@@ -1,15 +1,15 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Settings, Plus } from '@lucide/svelte';
import { Switch } from '$lib/components/ui/switch';
import { Plus, Settings } from '@lucide/svelte';
import { goto } from '$app/navigation';
import { DropdownMenuSearchable, McpLogo, McpServerIdentity } from '$lib/components/app';
import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
import { McpLogo, DropdownMenuSearchable, McpServerIdentity } from '$lib/components/app';
import { Switch } from '$lib/components/ui/switch';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { ROUTES } from '$lib/constants/routes';
import { HealthCheckStatus } from '$lib/enums';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { HealthCheckStatus } from '$lib/enums';
import type { MCPServerSettingsEntry } from '$lib/types';
import { goto } from '$app/navigation';
import { ROUTES } from '$lib/constants/routes';
interface Props {
onMcpSettingsClick?: () => void;
@@ -24,10 +24,13 @@
let hasMcpServers = $derived(mcpServers.length > 0);
let filteredMcpServers = $derived.by(() => {
const query = mcpSearchQuery.toLowerCase().trim();
if (!query) return mcpServers;
return mcpServers.filter((s) => {
const name = getServerLabel(s).toLowerCase();
const url = s.url.toLowerCase();
return name.includes(query) || url.includes(query);
});
});
@@ -1,8 +1,8 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Lightbulb, LightbulbOff, Check, Info } from '@lucide/svelte';
import { Check, Info, Lightbulb, LightbulbOff } from '@lucide/svelte';
import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
import * as Tooltip from '$lib/components/ui/tooltip';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
const reasoning = useReasoningMenu();
@@ -1,30 +1,30 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import type { Snippet } from 'svelte';
import * as Tooltip from '$lib/components/ui/tooltip';
import * as Sheet from '$lib/components/ui/sheet';
import * as Collapsible from '$lib/components/ui/collapsible';
import { File, MessageSquare, Zap, FolderOpen } from '@lucide/svelte';
import { Switch } from '$lib/components/ui/switch';
import { Checkbox } from '$lib/components/ui/checkbox';
import { TOOLTIP_DELAY_DURATION } from '$lib/constants';
import { ATTACHMENT_FILE_ITEMS } from '$lib/constants/attachment-menu';
import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { McpLogo } from '$lib/components/app';
import { File, FolderOpen, MessageSquare, Zap } from '@lucide/svelte';
import {
PencilRuler,
Check,
ChevronDown,
ChevronRight,
Lightbulb,
LightbulbOff,
Check
PencilRuler
} from '@lucide/svelte';
import { McpLogo } from '$lib/components/app';
import { Checkbox } from '$lib/components/ui/checkbox';
import * as Collapsible from '$lib/components/ui/collapsible';
import * as Sheet from '$lib/components/ui/sheet';
import { Switch } from '$lib/components/ui/switch';
import * as Tooltip from '$lib/components/ui/tooltip';
import { TOOLTIP_DELAY_DURATION } from '$lib/constants';
import { ATTACHMENT_FILE_ITEMS } from '$lib/constants/attachment-menu';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { HealthCheckStatus } from '$lib/enums';
import { AttachmentAction } from '$lib/enums/attachment.enums';
import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import type { Snippet } from 'svelte';
interface Props {
class?: string;
@@ -45,14 +45,14 @@
class: className = '',
disabled = false,
hasAudioModality = false,
hasVisionModality = false,
hasVideoModality = false,
hasMcpPromptsSupport = false,
hasMcpResourcesSupport = false,
hasVideoModality = false,
hasVisionModality = false,
onFileUpload,
onSystemPromptClick,
onMcpPromptClick,
onMcpResourcesClick,
onSystemPromptClick,
trigger
}: Props = $props();
@@ -64,13 +64,13 @@
const attachmentMenu = useAttachmentMenu(
() => ({
hasVisionModality,
hasAudioModality,
hasVideoModality,
hasMcpPromptsSupport,
hasMcpResourcesSupport
hasMcpResourcesSupport,
hasVideoModality,
hasVisionModality
}),
() => ({ onFileUpload, onSystemPromptClick, onMcpPromptClick, onMcpResourcesClick }),
() => ({ onFileUpload, onMcpPromptClick, onMcpResourcesClick, onSystemPromptClick }),
() => {
sheetOpen = false;
}
@@ -1,14 +1,14 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { PencilRuler, ChevronDown, ChevronRight, Loader2, Info, Check } from '@lucide/svelte';
import { Check, ChevronDown, ChevronRight, Info, Loader2, PencilRuler } from '@lucide/svelte';
import { Checkbox } from '$lib/components/ui/checkbox';
import * as Collapsible from '$lib/components/ui/collapsible';
import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
import * as Tooltip from '$lib/components/ui/tooltip';
import { toolsStore } from '$lib/stores/tools.svelte';
import { CLI_FLAGS } from '$lib/constants';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { toolsStore } from '$lib/stores/tools.svelte';
const toolsPanel = useToolsPanel();
const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0);
@@ -1,8 +1,8 @@
<script lang="ts">
import { isMobile } from '$lib/stores/viewport.svelte';
import ChatFormActionAddButton from './ChatFormActionAddButton.svelte';
import ChatFormActionAddDropdown from './ChatFormActionAddDropdown.svelte';
import ChatFormActionAddSheet from './ChatFormActionAddSheet.svelte';
import ChatFormActionAddButton from './ChatFormActionAddButton.svelte';
import { isMobile } from '$lib/stores/viewport.svelte';
interface Props {
disabled?: boolean;
@@ -21,9 +21,9 @@
let {
disabled = false,
hasAudioModality = false,
hasVideoModality = false,
hasMcpPromptsSupport = false,
hasMcpResourcesSupport = false,
hasVideoModality = false,
hasVisionModality = false,
onFileUpload,
onMcpPromptClick,
@@ -1,15 +1,15 @@
<script lang="ts">
import { ModelsSelectorDropdown, ModelsSelectorSheet } from '$lib/components/app';
import { chatStore } from '$lib/stores/chat.svelte';
import { activeMessages } from '$lib/stores/conversations.svelte';
import {
modelsStore,
modelOptions,
modelsStore,
selectedModelId,
selectedModelName
} from '$lib/stores/models.svelte';
import { isRouterMode, serverError } from '$lib/stores/server.svelte';
import { ModelsSelectorDropdown, ModelsSelectorSheet } from '$lib/components/app';
import { isMobile } from '$lib/stores/viewport.svelte';
import { activeMessages } from '$lib/stores/conversations.svelte';
interface Props {
disabled?: boolean;
@@ -27,9 +27,9 @@
disabled = false,
forceForegroundText = false,
hasAudioModality = $bindable(false),
hasModelSelected = $bindable(false),
hasVideoModality = $bindable(false),
hasVisionModality = $bindable(false),
hasModelSelected = $bindable(false),
isSelectedModelInCache = $bindable(true),
submitTooltip = $bindable(''),
useGlobalSelection = false
@@ -46,6 +46,7 @@
let selectorModel = $derived.by(() => {
const storeModel = selectedModelName();
if (storeModel && storeModel !== conversationModel) {
return storeModel;
}
@@ -66,6 +67,7 @@
modelsStore.selectedModelName = null;
modelsStore.clearSelection();
}
lastSyncedConversationModel = conversationModel;
} else if (
isRouter &&
@@ -76,6 +78,7 @@
) {
lastSyncedConversationModel = null;
const first = modelOptions().find((m) => modelsStore.loadedModelIds.includes(m.model));
if (first) modelsStore.selectModelById(first.id);
}
});
@@ -1,8 +1,8 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Mic, Square } from '@lucide/svelte';
import { Button } from '$lib/components/ui/button';
import * as Tooltip from '$lib/components/ui/tooltip';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
interface Props {
class?: string;
@@ -1,28 +1,28 @@
<script lang="ts">
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { Square, SkipForward } from '@lucide/svelte';
import { Button } from '$lib/components/ui/button';
import { ChatService } from '$lib/services';
import { SkipForward, Square } from '@lucide/svelte';
import { goto } from '$app/navigation';
import { page } from '$app/state';
import {
ChatFormActionsAdd,
ChatFormActionModels,
ChatFormActionRecord,
ChatFormActionsAdd,
ChatFormActionSubmit,
ChatFormContextGauge
} from '$lib/components/app';
import { Button } from '$lib/components/ui/button';
import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
import { ROUTES } from '$lib/constants/routes';
import { FileTypeCategory, MessageRole } from '$lib/enums';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { config } from '$lib/stores/settings.svelte';
import { activeMessages, conversationsStore } from '$lib/stores/conversations.svelte';
import { ChatService } from '$lib/services';
import {
activeProcessingState,
isChatStreaming,
isLoading as chatIsLoading
} from '$lib/stores/chat.svelte';
import { activeMessages, conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { config } from '$lib/stores/settings.svelte';
import { getFileTypeCategory } from '$lib/utils';
import { goto } from '$app/navigation';
import { page } from '$app/state';
import { ROUTES } from '$lib/constants/routes';
interface Props {
canSend?: boolean;
@@ -51,15 +51,15 @@
isLoading = false,
isReasoning = false,
isRecording = false,
showAddButton = true,
showModelSelector = true,
uploadedFiles = [],
onFileUpload,
onMcpPromptClick,
onMcpResourcesClick,
onMicClick,
onStop,
onSystemPromptClick,
onMcpPromptClick,
onMcpResourcesClick
showAddButton = true,
showModelSelector = true,
uploadedFiles = []
}: Props = $props();
let currentConfig = $derived(config());
@@ -105,29 +105,39 @@
if (!page.params.id) return false;
const messages = activeMessages() as DatabaseMessage[];
let totalHistoricalTokens = 0;
for (const m of messages) {
if (m.role !== MessageRole.ASSISTANT) continue;
const timings = m.timings;
if (!timings) continue;
const agenticLlm = timings.agentic?.llm;
if (agenticLlm?.prompt_n != null || agenticLlm?.predicted_n != null) {
totalHistoricalTokens += (agenticLlm?.prompt_n ?? 0) + (agenticLlm?.predicted_n ?? 0);
} else {
totalHistoricalTokens += (timings.prompt_n ?? 0) + (timings.predicted_n ?? 0);
}
}
if (totalHistoricalTokens > 0) return true;
if (!chatIsLoading() && !isChatStreaming()) return false;
const processingState = activeProcessingState();
if (!processingState) return false;
const livePromptTokens = Math.max(
processingState.promptTokens ?? 0,
processingState.promptProgress?.processed ?? 0
);
const liveOutputTokens = processingState.outputTokensUsed ?? 0;
return livePromptTokens > 0 || liveOutputTokens > 0;
});
</script>
@@ -1,11 +1,8 @@
<script lang="ts">
import { onDestroy, onMount, untrack } from 'svelte';
import { mode } from 'mode-watcher';
import githubDarkCss from 'highlight.js/styles/github-dark.css?inline';
import githubLightCss from 'highlight.js/styles/github.css?inline';
import { isMobile } from '$lib/stores/viewport.svelte';
import { ColorMode } from '$lib/enums';
import { TRIM_LEADING_PADDING_REGEX, TRIM_TRAILING_PADDING_REGEX } from '$lib/constants';
import { ColorMode } from '$lib/enums';
import { isMobile } from '$lib/stores/viewport.svelte';
import type { ContentToken, SourceHistoryEntry } from '$lib/utils';
import {
badgeAwareWordJump,
buildFragment,
@@ -19,10 +16,13 @@
SourceHistory,
stripBlockBoundaryLineBreaks,
syncCodeBlockHatches,
tokenizeContent,
textOffsetToRange
textOffsetToRange,
tokenizeContent
} from '$lib/utils';
import type { ContentToken, SourceHistoryEntry } from '$lib/utils';
import githubLightCss from 'highlight.js/styles/github.css?inline';
import githubDarkCss from 'highlight.js/styles/github-dark.css?inline';
import { mode } from 'mode-watcher';
import { onDestroy, onMount, untrack } from 'svelte';
interface Props {
class?: string;
@@ -57,7 +57,9 @@
// serialized source, not the DOM shape.
function syncEmptyState(serialized?: string) {
if (!rootElement) return;
const source = serialized ?? serializeContent(rootElement);
rootElement.dataset.empty = source.length === 0 ? 'true' : 'false';
}
@@ -93,23 +95,24 @@
*/
function highlightCodeBlockElement(el: HTMLElement): boolean {
const segment = el.textContent ?? '';
if (highlightedSegments.get(el) === segment) return false;
const open = CODE_BLOCK_OPEN_RE.exec(segment);
if (!open) return false;
const prefix = open[0];
const language = open[1].trim().split(/\s+/)[0] ?? '';
const content = segment.slice(prefix.length, -3);
const leading = content.match(TRIM_LEADING_PADDING_REGEX)?.[0] ?? '';
const trailing = content.match(TRIM_TRAILING_PADDING_REGEX)?.[0] ?? '';
const core = content.slice(leading.length, content.length - trailing.length);
// autoDetect off: re-guessing the language on every keystroke
// costs ~38ms a call and flickers while typing
const html = core ? highlightCode(core, language || 'text', false) : '';
const tpl = document.createElement('template');
tpl.innerHTML = html;
el.replaceChildren(
@@ -118,6 +121,7 @@
document.createTextNode(trailing + '```')
);
highlightedSegments.set(el, segment);
return true;
}
@@ -136,9 +140,11 @@
if (!rootElement) return;
const range = safeRange();
if (!range) return;
let node: Node | null = range.startContainer;
if (node === rootElement) {
node = rootElement.childNodes[range.startOffset - 1] ?? null;
}
@@ -146,11 +152,14 @@
while (node && node !== rootElement) {
if (node instanceof HTMLElement && node.dataset.codeToken === 'block') {
const caret = rangeToTextOffset(rootElement, range);
if (highlightCodeBlockElement(node)) {
restoreCaret(caret);
}
return;
}
node = node.parentNode;
}
}
@@ -182,6 +191,7 @@
document.querySelectorAll('style[data-highlight-theme-preview]').forEach((s) => s.remove());
const style = document.createElement('style');
style.setAttribute('data-highlight-theme-preview', 'true');
style.textContent = isDark ? githubDarkCss : githubLightCss;
@@ -196,6 +206,7 @@
if (!rootElement) return null;
const selection = window.getSelection();
if (!selection || selection.rangeCount === 0) return null;
const range = selection.getRangeAt(0);
@@ -212,6 +223,7 @@
const target = textOffsetToRange(rootElement, offset);
const selection = window.getSelection();
if (!selection) return;
if (extend && selection.anchorNode) {
@@ -221,6 +233,7 @@
target.startContainer,
target.startOffset
);
return;
}
@@ -230,14 +243,16 @@
function resizeHeight() {
if (!rootElement) return;
rootElement.style.height = 'auto';
rootElement.style.height = `${rootElement.scrollHeight}px`;
}
function recordHistory(newGroup: boolean) {
if (!rootElement) return;
history.push(
{ value: lastEmittedValue, caret: rangeToTextOffset(rootElement, safeRange()) },
{ caret: rangeToTextOffset(rootElement, safeRange()), value: lastEmittedValue },
Date.now(),
newGroup
);
@@ -262,10 +277,12 @@
// lands on the line directly below the block.
if (inputType === 'insertLineBreak' || inputType === 'insertParagraph') {
const caret = rangeToTextOffset(rootElement, safeRange());
if (stripBlockBoundaryLineBreaks(rootElement)) {
restoreCaret(caret);
} else {
const source = serializeContent(rootElement);
let end = caret;
// the caret must end up after the inserted \n; some browsers
@@ -284,7 +301,9 @@
// \n doubles as a block's separator line (source ends with
// \n\n) or sits inside a block element.
let last = rootElement.lastChild;
while (last && last.nodeName === 'BR') last = last.previousSibling;
if (
end === source.length &&
source.endsWith('\n') &&
@@ -302,7 +321,9 @@
syncCodeBlockHatches(rootElement);
const serialized = serializeContent(rootElement);
syncEmptyState(serialized);
if (serialized === lastEmittedValue) return;
// Plain typing/deletes coalesce per time window; structural edits
@@ -316,6 +337,7 @@
// completed or broken) - the browser-owned text nodes cannot
// restyle themselves across element boundaries.
const tokens = tokenizeContent(serialized);
if (!domMatchesTokens(rootElement, tokens)) {
renderTokens(tokens);
@@ -324,6 +346,7 @@
// block element and the rebuild splits it back out, which
// synthesizes the separator newline) - keep value in sync.
const reserialized = serializeContent(rootElement);
if (reserialized !== serialized) {
lastEmittedValue = reserialized;
value = reserialized;
@@ -361,6 +384,7 @@
if (!rootElement) return;
const range = safeRange();
if (!range) return;
if (!range.collapsed) {
@@ -374,16 +398,22 @@
// a break at the very end of a code block exits the block (the
// new line belongs below it, not inside)
let exitBlock: HTMLElement | null = null;
if (container.nodeType === Node.TEXT_NODE) {
let node: Node | null = container.parentNode;
while (node && node !== rootElement) {
if (node instanceof HTMLElement && node.dataset.codeToken === 'block') {
const tail = document.createRange();
tail.setStart(container, offset);
tail.setEnd(node, node.childNodes.length);
if (tail.toString().length === 0) exitBlock = node;
break;
}
node = node.parentNode;
}
}
@@ -392,6 +422,7 @@
exitBlock.after(nl);
} else if (container.nodeType === Node.TEXT_NODE) {
const text = container as Text;
if (offset === 0) {
text.before(nl);
} else if (offset === text.length) {
@@ -405,6 +436,7 @@
const selection = window.getSelection();
const after = document.createRange();
after.setStartAfter(nl);
after.collapse(true);
selection?.removeAllRanges();
@@ -428,9 +460,11 @@
if (rootElement.firstChild?.nodeName === 'BR') return false;
const first = rootElement.firstChild;
if (!(first instanceof HTMLElement) || first.dataset.codeToken !== 'block') return false;
const range = safeRange();
if (!range || !range.collapsed) return false;
// the caret must sit inside the block: on its very first
@@ -439,16 +473,19 @@
if (!first.contains(range.startContainer)) return false;
const caret = rangeToTextOffset(rootElement, range);
if (key === 'ArrowLeft') {
if (caret !== 0) return false;
} else {
const firstLineEnd = (first.textContent ?? '').indexOf('\n');
if (firstLineEnd !== -1 && caret > firstLineEnd) return false;
}
// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the contenteditable host, never its children
rootElement.prepend(document.createElement('br'));
restoreCaret(0, extend);
return true;
}
@@ -464,14 +501,17 @@
if (!rootElement) return;
const first = rootElement.firstChild;
if (first?.nodeName !== 'BR') return;
const second = first.nextSibling;
if (!(second instanceof HTMLElement) || second.dataset.codeToken !== 'block') return;
const range = safeRange();
const onHatch =
range !== null && range.startContainer === rootElement && range.startOffset === 0;
if (!onHatch) {
first.remove();
}
@@ -491,6 +531,7 @@
*/
function handleKeydown(event: KeyboardEvent) {
const mod = event.ctrlKey || event.metaKey;
if (mod && !event.altKey && !isComposing && rootElement) {
const key = event.key.toLowerCase();
const isUndo = key === 'z' && !event.shiftKey;
@@ -499,11 +540,13 @@
if (isUndo || isRedo) {
event.preventDefault();
const current = {
value: lastEmittedValue,
caret: rangeToTextOffset(rootElement, safeRange())
caret: rangeToTextOffset(rootElement, safeRange()),
value: lastEmittedValue
};
const entry = isUndo ? history.undo(current) : history.redo(current);
if (entry) applyHistoryEntry(entry);
return;
}
}
@@ -524,6 +567,7 @@
// stuck on the old line (see insertLineBreak).
event.preventDefault();
insertLineBreak();
return;
}
@@ -543,6 +587,7 @@
// re-tokenize/re-highlight follows.
event.preventDefault();
document.execCommand('insertLineBreak');
return;
}
@@ -555,6 +600,7 @@
) {
if (moveCaretBeforeLeadingCodeBlock(event.key, event.shiftKey)) {
event.preventDefault();
return;
}
}
@@ -574,6 +620,7 @@
if (target !== null) {
event.preventDefault();
restoreCaret(target, event.shiftKey);
return;
}
}
@@ -586,6 +633,7 @@
// change as our own and does not re-render.
function applyHistoryEntry(entry: SourceHistoryEntry) {
if (!rootElement) return;
renderTokens(tokenizeContent(entry.value));
lastEmittedValue = entry.value;
value = entry.value;
@@ -602,6 +650,7 @@
*/
function handlePasteEvent(event: ClipboardEvent) {
const pasted = event.clipboardData?.getData('text/plain');
if (pasted && pasted.length > 0) {
event.preventDefault();
@@ -609,6 +658,7 @@
// element-boundary carets (e.g. right before a badge) Chromium's
// insertText can drop the preceding text node's trailing whitespace.
const range = safeRange();
if (rootElement && range && range.collapsed) {
restoreCaret(rangeToTextOffset(rootElement, range));
}
@@ -621,6 +671,7 @@
// consumes the event (files, quoted prompts, long text).
function handlePaste(event: ClipboardEvent) {
onPaste?.(event);
if (!event.defaultPrevented) {
handlePasteEvent(event);
}
@@ -634,20 +685,23 @@
if (!rootElement) return null;
const range = safeRange();
if (!range || range.collapsed) return null;
const startRange = range.cloneRange();
startRange.collapse(true);
const source = serializeContent(rootElement);
const start = rangeToTextOffset(rootElement, startRange);
const end = rangeToTextOffset(rootElement, range);
return { text: source.slice(start, end), range };
return { range, text: source.slice(start, end) };
}
function handleCopy(event: ClipboardEvent) {
const slice = selectionSourceSlice();
if (!slice) return;
event.clipboardData?.setData('text/plain', slice.text);
@@ -656,6 +710,7 @@
function handleCut(event: ClipboardEvent) {
const slice = selectionSourceSlice();
if (!slice) return;
event.clipboardData?.setData('text/plain', slice.text);
@@ -675,6 +730,7 @@
resizeHeight();
syncEmptyState();
document.addEventListener('selectionchange', handleSelectionChange);
if (!isMobile.current) {
rootElement?.focus({ preventScroll: true });
}
@@ -689,6 +745,7 @@
// browser already owns the right shape.
$effect(() => {
const incoming = value ?? '';
if (incoming === lastEmittedValue) return;
recordHistory(true); // external edit (mention insert, clear, ...): own undo step
@@ -702,6 +759,7 @@
export function getCaretOffset(): number {
if (!rootElement) return 0;
return rangeToTextOffset(rootElement, safeRange());
}
@@ -710,11 +768,13 @@
if (rootElement && rootElement !== document.activeElement) {
rootElement.focus({ preventScroll: true });
}
restoreCaret(offset);
}
export function focus() {
if (isMobile.current) return;
rootElement?.focus({ preventScroll: true });
}
@@ -1,9 +1,7 @@
<script lang="ts">
import { untrack } from 'svelte';
import { activeConversation, activeMessages } from '$lib/stores/conversations.svelte';
import { chatStore, isChatStreaming, isLoading } from '$lib/stores/chat.svelte';
import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
import ContextGaugeDial from './ContextGaugeDial.svelte';
import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
import { chatStore, isChatStreaming, isLoading } from '$lib/stores/chat.svelte';
import {
gaugeTriggerClick,
gaugeTriggerEnter,
@@ -11,22 +9,28 @@
gaugeTriggerLeave,
gaugeTriggerPointerDown
} from '$lib/stores/context-gauge-popup.svelte';
import { activeConversation, activeMessages } from '$lib/stores/conversations.svelte';
import { untrack } from 'svelte';
const gauge = useContextGauge();
$effect(() => {
const conv = activeConversation();
untrack(() => chatStore.setActiveProcessingConversation(conv?.id ?? null));
});
$effect(() => {
const conv = activeConversation();
const messages = activeMessages() as DatabaseMessage[];
if (!conv) return;
if (isLoading() || isChatStreaming()) return;
if (messages.length === 0) {
untrack(() => chatStore.clearProcessingState(conv.id));
return;
}
@@ -5,7 +5,7 @@
subtitle?: string;
}
let { label, value, subtitle }: Props = $props();
let { label, subtitle, value }: Props = $props();
</script>
<div class="grid gap-1.5">
@@ -1,8 +1,8 @@
<script lang="ts">
import ContextGaugeDetailRow from './ContextGaugeDetailRow.svelte';
import { ChevronDown } from '@lucide/svelte';
import * as Collapsible from '$lib/components/ui/collapsible';
import { STATS_UNITS } from '$lib/constants';
import ContextGaugeDetailRow from './ContextGaugeDetailRow.svelte';
interface Props {
currentRead: number;
@@ -18,15 +18,15 @@
}
let {
currentRead,
currentFresh,
currentCache,
currentOutput,
kvTotal,
cumulativeRead,
cumulativeOutput,
cumulativeCacheTotal,
averageTokensPerSecond,
cumulativeCacheTotal,
cumulativeOutput,
cumulativeRead,
currentCache,
currentFresh,
currentOutput,
currentRead,
kvTotal,
transientDetails
}: Props = $props();
@@ -8,7 +8,7 @@
size?: 'sm' | 'md';
}
let { percent, level, size = 'sm' }: Props = $props();
let { level, percent, size = 'sm' }: Props = $props();
const RADIUS = 11;
const CIRCUMFERENCE = 2 * Math.PI * RADIUS;
@@ -8,7 +8,7 @@
onLoad: () => void;
}
let { modelId, isLoading, onLoad }: Props = $props();
let { isLoading, modelId, onLoad }: Props = $props();
</script>
{#if modelId !== null && !isLoading}
@@ -1,15 +1,15 @@
<script lang="ts">
import { formatParameters } from '$lib/utils/formatters';
import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
import { colorLevelBgClass, colorLevelTextClass } from './context-gauge';
import ContextGaugeDetails from './ContextGaugeDetails.svelte';
import ContextGaugeLoadModel from './ContextGaugeLoadModel.svelte';
import { colorLevelBgClass, colorLevelTextClass } from './context-gauge';
import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
import {
gaugePopup,
gaugeCardEnter,
gaugeCardLeave,
gaugePopup,
gaugePopupClose
} from '$lib/stores/context-gauge-popup.svelte';
import { formatParameters } from '$lib/utils/formatters';
const gauge = useContextGauge();
@@ -30,13 +30,18 @@
const onPointerDown = (event: PointerEvent) => {
const target = event.target;
if (!(target instanceof Node)) return;
if (cardEl?.contains(target)) return;
if (target instanceof Element && target.closest('[data-context-gauge-trigger]')) return;
gaugePopupClose();
};
document.addEventListener('pointerdown', onPointerDown, true);
return () => document.removeEventListener('pointerdown', onPointerDown, true);
});
@@ -5,8 +5,11 @@ const CRITICAL_THRESHOLD = 95;
export function colorLevelFromPercent(percent: number | null): ColorLevel {
if (percent === null) return 'neutral';
if (percent >= CRITICAL_THRESHOLD) return 'critical';
if (percent >= WARNING_THRESHOLD) return 'warning';
return 'ok';
}
@@ -1,13 +1,13 @@
<script lang="ts">
import { mcpStore } from '$lib/stores/mcp.svelte';
import {
mcpResourceAttachments,
mcpHasResourceAttachments
} from '$lib/stores/mcp-resources.svelte';
import {
ChatAttachmentsListItemMcpResource,
HorizontalScrollCarousel
} from '$lib/components/app';
import { mcpStore } from '$lib/stores/mcp.svelte';
import {
mcpHasResourceAttachments,
mcpResourceAttachments
} from '$lib/stores/mcp-resources.svelte';
interface Props {
class?: string;
@@ -1,14 +1,14 @@
<script lang="ts">
import { FolderOpen, Sparkles } from '@lucide/svelte';
import { MODEL_SELECTOR_ICON } from '$lib/constants';
import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte';
import { ChatFormCommandAction } from '$lib/enums';
import type { ChatFormCommand } from '$lib/types';
import {
ChatFormPickerList,
ChatFormPickerListItem,
ChatFormPickerPopover
} from '$lib/components/app/chat';
import { MODEL_SELECTOR_ICON } from '$lib/constants';
import { ChatFormCommandAction } from '$lib/enums';
import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte';
import type { ChatFormCommand } from '$lib/types';
/**
* Slash-command picker; `query` (typed after `/`) filters the commands.
@@ -24,12 +24,12 @@
onSelect: (command: ChatFormCommand) => void;
}
let { class: className = '', isOpen, query, commands, onClose, onSelect }: Props = $props();
let { class: className = '', commands, isOpen, onClose, onSelect, query }: Props = $props();
const commandIcon: Record<ChatFormCommandAction, typeof Sparkles> = {
[ChatFormCommandAction.PROMPT]: Sparkles,
[ChatFormCommandAction.CWD]: FolderOpen,
[ChatFormCommandAction.MODEL]: MODEL_SELECTOR_ICON
[ChatFormCommandAction.MODEL]: MODEL_SELECTOR_ICON,
[ChatFormCommandAction.PROMPT]: Sparkles
};
const trimmedQuery = $derived((query ?? '').trim().toLowerCase());
@@ -51,20 +51,24 @@
function stepEnabled(from: number, dir: number): number {
const n = filteredCommands.length;
if (n === 0) return -1;
for (let i = 1; i <= n; i++) {
const idx = (from + dir * i + n) % n;
if (!filteredCommands[idx].disabled) return idx;
}
return -1;
}
const nav = usePickerNavigation({
isOpen: () => isOpen,
count: () => filteredCommands.length,
step: (from, dir) => (from < 0 ? firstEnabledIndex() : stepEnabled(from, dir)),
isOpen: () => isOpen,
onClose: () => onClose(),
onSelect: (index) => handleSelect(filteredCommands[index])
onSelect: (index) => handleSelect(filteredCommands[index]),
step: (from, dir) => (from < 0 ? firstEnabledIndex() : stepEnabled(from, dir))
});
$effect(() => {
@@ -76,8 +80,10 @@
$effect(() => {
if (nav.hoveredIndex < 0 || nav.hoveredIndex >= filteredCommands.length) {
nav.reset(firstEnabledIndex());
return;
}
if (filteredCommands[nav.hoveredIndex].disabled) {
nav.reset(firstEnabledIndex());
}
@@ -85,6 +91,7 @@
function handleSelect(command: ChatFormCommand) {
if (command.disabled) return;
onSelect(command);
onClose();
}
@@ -1,22 +1,22 @@
<script lang="ts">
import { File, Folder } from '@lucide/svelte';
import { abbreviateHome, runGlobSearchWithChildren, type GlobEntryResult } from '$lib/utils';
import { toolsStore } from '$lib/stores/tools.svelte';
import { BuiltInTool, FileMentionEntryType, GlobSearchType, KeyboardKey } from '$lib/enums';
import { isMobile } from '$lib/stores/viewport.svelte';
import { config } from '$lib/stores/settings.svelte';
import { ChatFormPickerList, ChatFormPickerListItem } from '$lib/components/app/chat';
import HighlightedMatch from '$lib/components/app/forms/HighlightedMatch.svelte';
import * as Popover from '$lib/components/ui/popover';
import * as Tooltip from '$lib/components/ui/tooltip';
import HighlightedMatch from '$lib/components/app/forms/HighlightedMatch.svelte';
import { ChatFormPickerList, ChatFormPickerListItem } from '$lib/components/app/chat';
import { useDebouncedSearch } from '$lib/hooks/use-debounced-search.svelte';
import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte';
import type { FileMentionEntry } from '$lib/types';
import {
FILE_GLOB_SEARCH_PICKERS_DEFAULT_SEARCH_DEPTH,
HOME_TILDE,
SEARCH_DEBOUNCE_MS
} from '$lib/constants';
import { BuiltInTool, FileMentionEntryType, GlobSearchType, KeyboardKey } from '$lib/enums';
import { useDebouncedSearch } from '$lib/hooks/use-debounced-search.svelte';
import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte';
import { config } from '$lib/stores/settings.svelte';
import { toolsStore } from '$lib/stores/tools.svelte';
import { isMobile } from '$lib/stores/viewport.svelte';
import type { FileMentionEntry } from '$lib/types';
import { abbreviateHome, type GlobEntryResult, runGlobSearchWithChildren } from '$lib/utils';
/**
* Floating file/folder mention picker. The chat input is the search
@@ -38,18 +38,18 @@
let {
class: className = '',
isOpen,
query,
customAnchor = null,
scopePath = null,
isOpen,
onClose,
onOpened,
onSelect,
onOpened
query,
scopePath = null
}: Props = $props();
const nav = usePickerNavigation({
isOpen: () => isOpen,
count: () => displayedItems.length,
isOpen: () => isOpen,
onClose: () => onClose(),
onSelect: (index) => handleSelect(displayedItems[index])
});
@@ -69,6 +69,7 @@
// would otherwise reach the server as max_depth 0 = unlimited.
const searchDepth = $derived.by(() => {
const n = Number(config().mentionSearchMaxDepth);
return Number.isInteger(n) && n > 0 ? n : FILE_GLOB_SEARCH_PICKERS_DEFAULT_SEARCH_DEPTH;
});
@@ -78,8 +79,8 @@
const MENTION_SEARCH_LIMIT = 50;
const search = useDebouncedSearch({
debounceMs: SEARCH_DEBOUNCE_MS,
canRun: () => isOpen && fileSearchEnabled,
debounceMs: SEARCH_DEBOUNCE_MS,
getQuery: () => trimmedQuery,
run: async (query, signal, isCurrent) => {
try {
@@ -91,23 +92,29 @@
searchDepth,
MENTION_SEARCH_LIMIT,
signal,
{ type: GlobSearchType.ALL, descendOnTrailingSeparator: true }
{ descendOnTrailingSeparator: true, type: GlobSearchType.ALL }
);
if (!isCurrent()) return;
if (res.error) {
searchResults = [];
searchError = res.error;
return;
}
const toEntry = (e: GlobEntryResult): FileMentionEntry => ({
path: e.path,
name: e.name,
path: e.path,
type: e.type === 'dir' ? FileMentionEntryType.DIRECTORY : FileMentionEntryType.FILE
});
searchResults = res.entries.map(toEntry);
searchError = null;
} catch (err) {
if (!isCurrent() || signal.aborted) return;
searchResults = [];
searchError = err instanceof Error ? err.message : String(err);
}
@@ -121,9 +128,11 @@
if (fileSearchKey === null) {
return 'File search is unavailable on this server (started without --tools)';
}
if (!fileSearchEnabled) {
return 'File search is disabled - enable "Search files" in Settings > Tools to use @-mentions';
}
return searchError ? `Search failed - ${searchError}` : 'No matching files or folders';
});
@@ -131,6 +140,7 @@
$effect(() => {
if (typeof window === 'undefined') return;
void toolsStore.resolveServerHome();
});
@@ -146,12 +156,15 @@
$effect(() => {
const q = (query ?? '').trim();
if (!isOpen || !q || !fileSearchEnabled) {
search.cancel();
searchResults = [];
searchError = null;
return;
}
search.setLoading(true);
search.run(q);
});
@@ -167,9 +180,11 @@
// Enter-to-submit never fires mid-search.
if (isOpen && event.key === KeyboardKey.ENTER) {
event.preventDefault();
if (nav.hoveredIndex >= 0 && displayedItems[nav.hoveredIndex]) {
handleSelect(displayedItems[nav.hoveredIndex]);
}
return true;
}
@@ -1,7 +1,7 @@
<script lang="ts">
import type { Snippet } from 'svelte';
import type { MCPServerSettingsEntry } from '$lib/types';
import { mcpStore } from '$lib/stores/mcp.svelte';
import type { MCPServerSettingsEntry } from '$lib/types';
import type { Snippet } from 'svelte';
interface Props {
server: MCPServerSettingsEntry | undefined;
@@ -12,7 +12,7 @@
subtitle?: Snippet;
}
let { server, serverLabel, title, description, titleExtra, subtitle }: Props = $props();
let { description, server, serverLabel, subtitle, title, titleExtra }: Props = $props();
let faviconUrl = $derived(server ? mcpStore.getServerFavicon(server.id) : null);
</script>
@@ -1,9 +1,9 @@
<script lang="ts" generics="T">
import type { Snippet } from 'svelte';
import { SearchInput } from '$lib/components/app';
import ScrollArea from '$lib/components/ui/scroll-area/scroll-area.svelte';
import { useScrollActiveRow } from '$lib/hooks/use-scroll-active-row.svelte';
import { CHAT_FORM_POPOVER_MAX_HEIGHT } from '$lib/constants';
import { useScrollActiveRow } from '$lib/hooks/use-scroll-active-row.svelte';
import type { Snippet } from 'svelte';
interface Props {
items: T[];
@@ -28,22 +28,22 @@
}
let {
items,
isLoading,
selectedIndex,
searchQuery = $bindable(),
showSearchInput,
searchPlaceholder = 'Search...',
emptyMessage,
autofocus = false,
inputRef = $bindable(null),
onSearchClose,
itemKey,
item,
skeleton,
skeletonCount = 6,
emptyMessage,
footer,
scrollTrigger
inputRef = $bindable(null),
isLoading,
item,
itemKey,
items,
onSearchClose,
scrollTrigger,
searchPlaceholder = 'Search...',
searchQuery = $bindable(),
selectedIndex,
showSearchInput,
skeleton,
skeletonCount = 6
}: Props = $props();
let listContainer = $state<HTMLDivElement | null>(null);
@@ -55,11 +55,11 @@
// selectedIndex/items.length are untracked so hover and result replacement
// never re-fire the scroll; keyboard nav is the only path that bumps the trigger.
useScrollActiveRow({
getTrigger: () => scrollTrigger,
dataIndex: 'picker',
getContainer: () => listContainer,
getIndex: () => selectedIndex,
getCount: () => items.length,
dataIndex: 'picker'
getIndex: () => selectedIndex,
getTrigger: () => scrollTrigger
});
</script>
@@ -12,13 +12,13 @@
}
let {
children,
class: className = '',
isSelected = false,
disabled = false,
onclick,
onmouseenter,
dataIndex,
children
disabled = false,
isSelected = false,
onclick,
onmouseenter
}: Props = $props();
</script>
@@ -4,7 +4,7 @@
showBadge?: boolean;
}
let { titleWidth = 'w-48', showBadge = false }: Props = $props();
let { showBadge = false, titleWidth = 'w-48' }: Props = $props();
</script>
<div class="flex w-full items-start gap-3 rounded-lg px-3 py-2">
@@ -1,6 +1,6 @@
<script lang="ts">
import type { Snippet } from 'svelte';
import * as Popover from '$lib/components/ui/popover';
import type { Snippet } from 'svelte';
interface Props {
class?: string;
@@ -12,12 +12,12 @@
}
let {
children,
class: className = '',
isOpen = $bindable(false),
srLabel = 'Open picker',
onClose,
onKeydown,
children
srLabel = 'Open picker'
}: Props = $props();
</script>
@@ -1,19 +1,19 @@
<script lang="ts">
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { debounce, uuid } from '$lib/utils';
import { KeyboardKey } from '$lib/enums';
import type { MCPPromptInfo, GetPromptResult, MCPServerSettingsEntry } from '$lib/types';
import { SvelteMap } from 'svelte/reactivity';
import {
ChatFormPickerPopover,
ChatFormPickerItemHeader,
ChatFormPickerList,
ChatFormPickerListItem,
ChatFormPickerItemHeader,
ChatFormPickerListItemSkeleton,
ChatFormPickerPopover,
ChatFormPromptPickerArgumentForm
} from '$lib/components/app/chat';
import Badge from '$lib/components/ui/badge/badge.svelte';
import { KeyboardKey } from '$lib/enums';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import type { GetPromptResult, MCPPromptInfo, MCPServerSettingsEntry } from '$lib/types';
import { debounce, uuid } from '$lib/utils';
import { SvelteMap } from 'svelte/reactivity';
interface Props {
class?: string;
@@ -32,11 +32,11 @@
let {
class: className = '',
isOpen = false,
searchQuery = '',
onClose,
onPromptLoadStart,
onPromptLoadComplete,
onPromptLoadError
onPromptLoadError,
onPromptLoadStart,
searchQuery = ''
}: Props = $props();
let prompts = $state<MCPPromptInfo[]>([]);
@@ -89,7 +89,6 @@
try {
const perChatOverrides = conversationsStore.getAllMcpServerOverrides();
const initialized = await mcpStore.ensureInitialized(perChatOverrides);
if (!initialized) {
@@ -118,6 +117,7 @@
requestAnimationFrame(() => {
const firstInput = document.querySelector(`#arg-${args[0].name}`) as HTMLInputElement;
if (firstInput) {
firstInput.focus();
}
@@ -131,7 +131,6 @@
promptError = null;
const placeholderId = uuid();
const nonEmptyArgs = Object.fromEntries(
Object.entries(args).filter(([, value]) => value.trim() !== '')
);
@@ -142,10 +141,12 @@
try {
const result = await mcpStore.getPrompt(prompt.serverName, prompt.name, args);
onPromptLoadComplete?.(placeholderId, result);
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : 'Unknown error executing prompt';
onPromptLoadError?.(placeholderId, errorMessage);
}
}
@@ -167,9 +168,9 @@
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) {
console.log('[ChatFormPickerMcpPrompts] Fetching completions for:', {
serverName: selectedPrompt.serverName,
promptName: selectedPrompt.name,
argName,
promptName: selectedPrompt.name,
serverName: selectedPrompt.serverName,
value
});
}
@@ -187,9 +188,9 @@
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) {
console.log('[ChatFormPickerMcpPrompts] Autocomplete result:', {
argName,
value,
result,
suggestionsCount: result?.values.length ?? 0
suggestionsCount: result?.values.length ?? 0,
value
});
}
@@ -234,6 +235,7 @@
event.preventDefault();
event.stopPropagation();
handleCancelArgumentForm();
return;
}
@@ -274,6 +276,7 @@
selectedIndex = selectedIndexBeforeArgumentForm;
selectedIndexBeforeArgumentForm = null;
}
selectedPrompt = null;
promptArgs = {};
promptError = null;
@@ -284,6 +287,7 @@
if (event.key === KeyboardKey.ESCAPE) {
event.preventDefault();
if (selectedPrompt) {
// Return to prompt selection list, keeping the selected prompt active
handleCancelArgumentForm();
@@ -296,6 +300,7 @@
if (event.key === KeyboardKey.ARROW_DOWN) {
event.preventDefault();
if (filteredPrompts.length > 0) {
selectedIndex = (selectedIndex + 1) % filteredPrompts.length;
scrollTrigger++;
@@ -306,6 +311,7 @@
if (event.key === KeyboardKey.ARROW_UP) {
event.preventDefault();
if (filteredPrompts.length > 0) {
selectedIndex = selectedIndex === 0 ? filteredPrompts.length - 1 : selectedIndex - 1;
scrollTrigger++;
@@ -316,6 +322,7 @@
if (event.key === KeyboardKey.ENTER && !selectedPrompt) {
event.preventDefault();
if (filteredPrompts[selectedIndex]) {
handlePromptClick(filteredPrompts[selectedIndex]);
}
@@ -329,14 +336,14 @@
let filteredPrompts = $derived.by(() => {
const sortedServers = mcpStore.getServers();
const serverOrderMap = new Map(sortedServers.map((server, index) => [server.id, index]));
const sortedPrompts = [...prompts].sort((a, b) => {
const orderA = serverOrderMap.get(a.serverName) ?? Number.MAX_SAFE_INTEGER;
const orderB = serverOrderMap.get(b.serverName) ?? Number.MAX_SAFE_INTEGER;
return orderA - orderB;
});
const query = (searchQuery || internalSearchQuery).toLowerCase();
if (!query) return sortedPrompts;
return sortedPrompts.filter(
@@ -1,7 +1,7 @@
<script lang="ts">
import type { MCPPromptInfo } from '$lib/types';
import ChatFormPromptPickerArgumentInput from './ChatFormPromptPickerArgumentInput.svelte';
import { Button } from '$lib/components/ui/button';
import type { MCPPromptInfo } from '$lib/types';
interface Props {
prompt: MCPPromptInfo;
@@ -21,20 +21,20 @@
}
let {
prompt,
promptArgs,
suggestions,
loadingSuggestions,
activeAutocomplete,
autocompleteIndex,
promptError,
onArgInput,
onArgKeydown,
loadingSuggestions,
onArgBlur,
onArgFocus,
onArgInput,
onArgKeydown,
onCancel,
onSelectSuggestion,
onSubmit,
onCancel
prompt,
promptArgs,
promptError,
suggestions
}: Props = $props();
</script>
@@ -1,8 +1,8 @@
<script lang="ts">
import type { MCPPromptInfo } from '$lib/types';
import { fly } from 'svelte/transition';
import { Input } from '$lib/components/ui/input';
import { Label } from '$lib/components/ui/label';
import type { MCPPromptInfo } from '$lib/types';
import { fly } from 'svelte/transition';
type PromptArgument = NonNullable<MCPPromptInfo['arguments']>[number];
@@ -22,16 +22,16 @@
let {
argument,
value = '',
suggestions = [],
isLoadingSuggestions = false,
isAutocompleteActive = false,
autocompleteIndex = 0,
onInput,
onKeydown,
isAutocompleteActive = false,
isLoadingSuggestions = false,
onBlur,
onFocus,
onSelectSuggestion
onInput,
onKeydown,
onSelectSuggestion,
suggestions = [],
value = ''
}: Props = $props();
</script>
@@ -66,7 +66,7 @@
{#if isAutocompleteActive && suggestions.length > 0}
<div
class="absolute top-full right-0 left-0 z-10 mt-1 max-h-32 overflow-y-auto rounded-lg border border-border/50 bg-background shadow-lg"
transition:fly={{ y: -5, duration: 100 }}
transition:fly={{ duration: 100, y: -5 }}
>
{#each suggestions as suggestion, i (suggestion)}
<button

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