Repository navigation
lean: text in, text out engine; lean-chat and lean-core crates; Llama 3 RoPE - #6
Merged
Merged
Conversation
lean no longer keeps a conversation (ChatSession, chatGenerate, chatReset): generate() now takes the whole message list every call and returns one reply, with context handling left to the caller. Internally the engine still keeps the previous call's KV cache and reuses whatever prefix the new call's ids agree with (PrefixCache), invisible to the caller and scoped to a single client. Adds a typed LeanError::ContextOverflow (needed/available token counts) in place of a string-only error, surfaced in JS as an Error with name === "ContextOverflow" and numeric needed/available properties. Overflow leaves the engine's cache state untouched. Adds render_prompt/token_count over a message list, and a native examples/chat.rs terminal client (the same client-side context-overflow trim web/lean-chat-worker.js now uses) as the reference for both engines. Low-level token-id APIs (forward_prefill, forward_decode_step_argmax, tokenize, prefillTokens, decodeStepArgmax, decodeGreedy, load) are unchanged.
web/lean-chat-worker.js now holds the conversation itself (HISTORY) and calls the new generate(messages, params, onToken, abort); on a ContextOverflow error it drops the oldest user/assistant pair and retries, same trim rule as the native chat example. onToken now carries text only, no token id. Bumps the engine build tag on both loading URLs.
New tests/fixture_parity_generate.rs: runs the HF-reference fixture's prompt-bearing cases through chat::gpu_generate/cpu_generate (the new message-level path) and checks the reply text against the fixture's greedy continuation, closing the gap left by the existing fixture_parity tests (which exercise forward_prefill/forward_decode_step directly, not the new generate wrapper). web/lean-chat-worker.js: logs when generateWithTrim drops a turn after a ContextOverflow, so a successful trim (which otherwise reaches the page as a normal reply, with nothing visibly different) is observable in the browser console.
RUSTFLAGS overrides .cargo/config.toml target rustflags entirely rather than merging with it; this repo compensated by manually repeating +simd128 in every RUSTFLAGS string, which worked but was fragile and would silently drop simd128 again on the next change to .cargo/config.toml. Pass the remap-path-prefix (and, for the pkg-mt build, the atomics/bulk-memory/shared-memory link flags) through cargo --config instead, which merges with the config file. wasm-pack forwards EXTRA_OPTIONS straight to cargo build with no -- separator; passing --config after a literal -- made cargo reject it. Verified byte-identical lean_bg.wasm/lean_bg_mt.wasm output before and after this change (sha256 match), and confirmed +simd128 and the threads-build features are present in both target_features sections.
Lean is the engine's single entry point, with the same surface in Rust
and JavaScript: Lean::load(LoadOptions { gguf, tokenizer, max_ctx,
backend }) picks WebGPU, then the CPU forward pass (threads or one
thread), by capability; generate(text, params, on_token, abort) runs one
turn over the text exactly as given; token_count(text) tokenizes only.
The engine no longer applies a chat template or takes messages: the
message-level generate, renderPrompt, tokenCount(messages) and the
template-rendering tokenize are gone from the JS classes, and
LeanEngine/LeanEngineCpu keep only the token-level calls. Prefix KV reuse
is automatic and documented as serving one client per engine.
ContextOverflow { needed, available } is returned before anything is
touched; in JS it is an Error with name ContextOverflow, code
CONTEXT_OVERFLOW and both counts. AbortFlag is one Arc<AtomicBool> type
shared by Rust and JS.
The chat page and examples/chat.rs render their history with the chat
template themselves (chatRender in JS) and pass the text to generate.
Tests: fixture_parity_generate runs the HF fixtures through
Lean::generate on GPU and CPU; engine_api covers reuse equivalence
(growing and diverging prefixes) and overflow counts and recovery.
New workspace crate crates/lean-chat (no dependency on lean): Message (role, content, tool_calls, name, tool_call_id), render(template, messages, tools, add_generation_prompt) and trim_oldest_exchange. render matches transformers' apply_chat_template(..., tokenize=False), tools included: trim_blocks and lstrip_blocks on, tojson replaced with Python's json.dumps (no HTML escaping, ", " and ": " separators, insertion-order keys; the filter and its formatter come from the llm-wasm crate's template.rs), pycompat string methods, raise_exception, tools = none when absent. chat_template.rs moves here from lean. Exactness: tests/hf_exactness.rs renders the conversations in tests/fixtures (Qwen2.5, Qwen3 and SmolLM2 templates; multi-turn, system messages, non-ASCII and HTML-like text, tools, tool calls with object and string arguments, parallel calls, tool results) and requires the bytes transformers 5.19.0 produced; scripts/gen_chat_fixtures.py writes them. lean's web package exports chatRender and chatTrimOldestExchange, so a page does lean.generate(chatRender(...)) from one wasm. examples/chat.rs and the chat page render with lean-chat and trim with trim_oldest_exchange on ContextOverflow; lean's tests and lean-cli render fixture prompts with lean-chat. Build tags bumped on every loading URL. CI tests and lints lean-chat too.
transformers passes the tokenizer's special tokens into the chat-template
context, and templates such as Mistral's write {{ bos_token }} and
{{ eos_token }} into the text. render now takes a ChatTemplate (source,
bos_token, eos_token); template_from_tokenizer_config and
read_chat_template read all three from tokenizer_config.json (string or
added-token object), and ChatTemplate::new(source) keeps a source-only
constructor. A token the tokenizer does not define stays undefined, as in
transformers.
New HF-exact fixture: Mistral-7B-Instruct-v0.3's template (bos and eos,
tools, tool calls with ids, tool results). The fixture tool calls now
carry 9-character ids, which the Qwen and SmolLM2 templates ignore.
JS: ChatTemplate class (fromTokenizerConfig, constructor) and
chatRender(template: ChatTemplate, ...). The chat worker builds its
template with ChatTemplate.fromTokenizerConfig; build tags bumped.
lean-core holds the GGUF loading, the GPU forward pass (engine, pool, model, WGSL shaders, quantized weights), the CPU forward pass (cpu, cpu_kernels, cpu_team), the KV cache, LoRA, ForwardSpec with forward_chunk_spec and the sliced heads, plus the HF reference fixtures and the forward-pass parity tests. lean depends on it and keeps the engine: Lean, generation, sampling, the tokenizer, the prefix cache, the error type and the wasm bindings. lean-core's public surface is its config, cpu, engine, model and profile_report modules, narrowed to what lean and llm-life call: gguf, quant, lora and pool are private, and the pipelines, buffer helpers and pass plumbing inside Engine and Pool are crate-private. Dead code this exposed is removed (unused GGUF value payloads, tensor order and file length, embedding table dims and accessor, LoRA mlp flag and out_features, Pool allocation counters and debug listing, buf_uniform, upload_f32). The two tests that check private modules (Q6_K reference dequant, embed_head_sliced) are unit tests in src/tests. llm-life builds against lean-core unchanged in its sources: its lean dependency becomes package = "lean-core" and its native feature no longer needs lean/native-bin (pollster is a native dependency of lean-core). CI tests and lints lean-core.
The public Rust surface of lean is now Lean, LoadOptions, Backend, GgufSource (with ReadSeek), TokenizerSource, GenParams, AbortFlag and LeanError; every module is private. The decode-loop, sampling and TextStream tests move into unit tests (src/tests). The default web package exports Lean, AbortFlag and lean-chat's ChatTemplate, chatRender and chatTrimOldestExchange (lean-chat is linked with extern crate, so no Rust re-export is needed). LeanEngine, LeanEngineCpu and leanInit move to web_diag.rs behind a non-default `diag` cargo feature for lab tools; Lean.load installs the panic hook itself. The device page now runs each backend through the public API: the chat-template-rendered fixture prompt through Lean.generate (greedy, 32 tokens), compared with transformers' reference text, with prefill and per-token timings from the streamed pieces. The chat page drops leanInit. Build tags bumped. CI also lints the wasm32 build with web + diag.
…tensor Llama 3.1/3.2 GGUFs carry rope_scaling (rope_type llama3) only as a per-frequency F32 tensor, rope_freqs.weight, holding the ratio of the unscaled to the scaled inverse frequency; there is no rope.scaling.* metadata. config_from_gguf now reads it into Qwen2Config::rope_freq_factors (None when absent), so it takes the reader mutably. Both RoPE paths divide theta^(-2j/d) by that factor through one shared helper: the GPU cos/sin tables (new model::rope_tables(config, max_pos)) and cpu.rs's rope_inplace, single thread and threads alike. Without the tensor the expression is unchanged, so Qwen and SmolLM2 tables stay bit-identical. build_rope_tables keeps its signature; lean's engine, CLI, bench binary, example and diag wasm calls switch to rope_tables. A unit test checks the factors Llama-3.2-1B's GGUF carries against transformers' _compute_llama3_parameters evaluated in f64.
…orkgroups rope_neox.wgsl and rope_positions.wgsl dispatched rows*heads*(head_dim/2)/64 workgroups along x alone, which exceeds WebGPU's 65535 per-dimension limit from about 4096 prompt tokens at 32 heads of head_dim 64 (a 9044-token Llama-3.2-1B prefill asked for 144704). They now use grid1d like the other elementwise kernels: workgroup size 256, a stride_x field in their Dims uniform, and the flat index recovered as gid.y * stride_x + gid.x. Each invocation computes the same element as before.
…ma3 RoPE gen_fixture_llama32.py builds LlamaForCausalLM from the model's config.json (rope_type llama3) and loads the GGUF's dequantized weights into it: transformers' GGUF loader drops rope scaling, so loading with gguf_file= alone would make an unscaled reference. The chat template is rendered with a fixed date_string and stored in the fixture. Four fixtures (1B and 3B, Q4_0 and Q8_0) with four cases each, one of them a 9044-token retrieval prompt past the original 8192-position context. fixture_parity_llama32 runs them on the GPU path, fixture_parity_llama32_cpu on the CPU path (single thread or threads); the long case on CPU is opt-in through LEAN_LLAMA32_CPU_LONG=1.
build_rope_tables(head_dim, theta, max_pos) and rope_tables(config, max_pos) merge into build_rope_tables(&Qwen2Config, max_pos), which applies the GGUF's rope_freqs.weight factors when the config carries them. Without the tensor the tables are the same bits as before. lean's engine, CLI, bench binary, example, diag wasm calls and every test now call it with the model config.
… 65535 rows Qwen3's per-head q/k norms run rmsnorm over rows * heads rows, which at 32 query heads exceeds WebGPU's 65535 workgroups per dimension from 2048 prompt tokens (a 2443-token Qwen3-4B prefill asked for 78176). The x extent is now capped at 65535, the rest folds into y, and the shader reads its row as wg.y * stride_x + wg.x from the Dims uniform's former padding field. add_rmsnorm is decode-only (one row) and keeps its 1-D dispatch. A GPU unit test normalizes 65535 + 77 rows and checks rows on both sides of the fold.
…ries Llama 3.2's GGUF names only <|eot_id|> (128009) as eos, while its generation_config.json stops on 128001 <|end_of_text|>, 128008 <|eom_id|> and 128009. And with no padding id in the file, the pad fallback put the bos id (128000) in the stop set. GgufReader now notes, while it skips tokenizer.ggml.tokens, the ids of <|eot_id|>, <|eom_id|>, <|end_of_text|> and <|im_end|>. config_from_gguf builds eos_token_ids from the eos id, tokenizer.ggml.eot_token_id and eom_token_id when present, an explicit padding id (Qwen's <|endoftext|>, or Qwen's architecture default), and those vocabulary ids; the bos id is no longer a fallback for the padding id. src/tests/stop_tokens.rs (ignored unit test, needs the model files) compares the stop set with each model's generation_config.json eos_token_id list: equal for Llama-3.2-1B/3B, Qwen3-0.6B/1.7B and SmolLM2-360M/1.7B.
…fixed date transformers puts strftime_now(format) in every chat template's environment (datetime.now().strftime(format)), and Llama 3.x templates call it to write today's date into the system header; without it lean-chat fell back to the template's hard-coded "26 Jul 2024" and differed from transformers. render now registers strftime_now over the local clock (chrono::Local, through JS Date in the browser via chrono's wasmbind feature), formatting with chrono's strftime, which follows Python's directives; an unknown directive fails the render instead of panicking. render_at takes the date and time strftime_now sees, so tests and callers that need a reproducible prompt can pin it.
gen_chat_fixtures.py pins transformers' strftime_now to a fixed date and time, stores it in each fixture as "now", takes optional fixture names on the command line, and skips cases a template rejects. The new llama3_2 fixture uses an ungated copy of Llama-3.2-1B-Instruct's tokenizer files (its chat_template is identical to the original's) and covers ten cases, tools and tool calls included; the template raises on parallel tool calls, so that case is left out. hf_exactness.rs renders a fixture that has "now" with render_at at that date and time, and checks Llama 3.2 byte for byte like the others.
…xture fixture_parity_generate gains a Llama-3.2-1B Q4_0 case: the engine's greedy reply on the GPU and on the CPU must equal the HF reference continuation of lean-core's fixture_llama32_1b_q4_0.json, which also exercises the llama3 RoPE tables Lean::load builds and the end-of-turn stop ids. Those fixture cases carry the chat text transformers rendered (rendered) rather than a prompt, and it is used as is; the 9044-token case does not fit the test's 512 positions and is skipped (lean-core's fixture_parity_llama32 runs it).
web/index.html loads the chat worker as lean-chat-worker.js?v=ENGINE_BUILD, but build.sh rewrote the tag only in the worker and the device page, so the chat page kept its hand-written tag across builds and a browser could keep running a cached worker (and through it the previous wasm) after a deploy. The chat page's assignment is now rewritten with the others, and the check expects four rewritten assignments.
Adds a muted line under every assistant reply with time to first token, tokens and decode speed, prompt tokens and backend, plus a note when the oldest turn was dropped. The worker reports the trim in its done message. Build tag bumped.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
lean is split into three crates and the engine becomes one inference turn: text in, text out.
lean:Lean::load/generate/token_count, the same in Rust and JS. Backend is picked by capability (WebGPU, CPU threads, single thread). The KV cache of the previous call is reused for the longest matching token prefix, with output identical to a fresh engine; one client per engine. A prompt that does not fit returns a typedContextOverflow { needed, available }and the engine stays usable. The stateful chat session is removed.lean-chat: messages and tools to prompt text with the model's chat template, byte-exact with transformers'apply_chat_template(Qwen2.5, Qwen3, SmolLM2, Mistral v0.3, Llama 3.2 fixtures), plus a helper that drops the oldest exchange.lean-core: the forward pass (positions, masks, sliced head), KV cache, LoRA and kernels. Its public surface is what lean and score-only users need; everything else is private. The token-level and diagnostic JS calls are behind adiagfeature.strftime_nowin chat templates.Lean.generate.scripts/build.shpasses its extra flags withcargo --config, so the+simd128target feature from.cargo/config.tomlis kept, and it rewrites the chat page's build tag.Keep the branch
lean-one-turnafter merging: downstream crates pin commit 1fce23d, which is on it.