diff --git a/CMakeLists.txt b/CMakeLists.txt index 0f6426a45f3..c4cef1bdf62 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -910,6 +910,14 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) add_test(NAME route_resident_parity COMMAND route_resident_parity) set_tests_properties(route_resident_parity PROPERTIES TIMEOUT 120) + if(STRATA_ENABLE_CUDA) + # the fast router (router_top10's default since it runs on CUDA too) bitwise against the portable kernel + add_executable(router_fast tests/cuda/router_fast.cpp) + target_include_directories(router_fast PRIVATE ${_strata_gpu_include_directories}) + target_link_libraries(router_fast PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + add_test(NAME router_fast COMMAND router_fast) + endif() + add_executable(router_top10_parity src/kernels/router_top10_parity.cpp) target_include_directories(router_top10_parity PRIVATE ${_strata_gpu_include_directories}) target_link_libraries(router_top10_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) @@ -1007,6 +1015,10 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) target_include_directories(kv_hybrid_parity PRIVATE ${_strata_gpu_include_directories}) target_link_libraries(kv_hybrid_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) add_test(NAME kv_hybrid_parity COMMAND kv_hybrid_parity) + add_executable(attn_splits_parity src/kernels/attn_splits_parity.cpp) + target_include_directories(attn_splits_parity PRIVATE ${_strata_gpu_include_directories}) + target_link_libraries(attn_splits_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + add_test(NAME attn_splits_parity COMMAND attn_splits_parity) # PR #21: Q4_0 KV with Walsh-Hadamard rotation vs a host reference (GPU, synthetic, no model). add_executable(kv_q4_parity src/kernels/kv_q4_parity.cpp) target_include_directories(kv_q4_parity PRIVATE ${_strata_gpu_include_directories}) diff --git a/README.md b/README.md index fa44a1601a9..66e74af0c05 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,7 @@ sizes are faster. Larger sizes are a bit smarter. | Your RAM | Take | Why | | --- | --- | --- | +| **16-24 GB** | **[Qwen3.6-35B-A3B](docs/QWEN36.md)** | a smaller model of the same family; runs on 8-12 GB cards (NVIDIA) | | **32 GB** | **Coder** | it fits 32 GB, and it is made for code (with a 24 GB card, Q2_0 and IQ2_XS run too) | | **48 GB** | **IQ2_XS** (or Q2_0, the fastest) | the larger sizes do not fit | | **64 GB** | **IQ2_XS** (recommended), or IQ3_XXS / IQ3_S | every size fits; IQ3_S is the best and the slowest | @@ -135,6 +136,10 @@ sizes are faster. Larger sizes are a bit smarter. model. But Strata reads most of it from the SSD while it answers, so it writes only 7-8.5 tokens/s on a 64 GB PC. - **[OrcaRouter's Uncensored IQ3_XXS](docs/MODELS.md#orcarouter-uncensored-iq3_xxs):** you set it up by hand. It is not in the installer's menu. +- **[Qwen3.6-35B-A3B](docs/QWEN36.md):** a second, smaller model (35B, 3B active) whose experts take 12-14 GB of RAM: + 105-127 tokens/s on a 12 GB card, 75-89 with an 8 GB card's VRAM (emulated). NVIDIA, one GPU, no images yet. + Its coding-agent fine-tune [Ornith-1.5-35B-A3B](docs/QWEN36.md#ornith-15-35b-a3b) installs the same way + (`--family ornith`). Sizes, downloads and what fits where: [docs/MODELS.md](docs/MODELS.md). To add another model later, run `SETUP.bat` (Linux: `./setup.sh --setup`). diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/README.md b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/README.md new file mode 100644 index 00000000000..4ec4cb80464 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/README.md @@ -0,0 +1,61 @@ +# Community benchmark on NVIDIA GeForce RTX 5080 + +Measured on 2026-10-13 with `tools/strata_bench.py` (suite 1, 3 runs per test, medians with ranges). Speeds are the engine's own timings; first-token times are measured by the client over HTTP. Speed only: no correctness check was run. + +## Hardware and software + +- GPU and VRAM: NVIDIA GeForce RTX 5080, 16 GiB VRAM +- CPU: Intel(R) Core(TM) i9-14900K (24 cores, 32 threads); installed RAM: 47.7 GiB; storage: not measured +- OS: Windows 11 (10.0.28000); GPU driver: 617.42; CUDA: not read +- Strata: project 0.1.43.1, commit f1e33fb9; engine 0.1.43.1 +- Background workloads and power limits: not recorded + +## Model and configuration + +- Model and quantization: ornith-1.5-35b-a3b-abliterated-iq4_xs +- Context: 262144 tokens; parallel slots: 1 +- Settings: the install's own configuration, unchanged; thinking off for the suite; temperature 0 + +```text +strata.exe --pack ornith-abl-iq4_xs --native Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --expert-profile expert-profile-qwen36.bin --expert-cache auto --prefill auto --spec 4 --mtp Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --mtp-draft-vocab draft_vocab_cyrillic.bin --max-context 262144 --kv int8 --expert-profile-save expert-profile-ornith-learned.bin --pcie-frac 0.00 --spec-min-p 0.50 --pool-workers 7 --adapt-every 1 --adapt-swaps 160 --adapt-decay 0.97 +``` + +## Method + +Run by `python tools/strata_bench.py --long` against the HTTP API of a Strata server (started by the benchmark from the install configuration); model loading is not timed. One warm-up chat and a 1K prompt are discarded first. Every prompt starts with a unique tag, so the conversation cache cannot help the fresh-prompt tests. + +- short chat: the three questions of `tools/ab_engine.py`, up to 100 new tokens, 3 runs +- 4K / 16K prompts: generated word lists sized by a calibration request, 24 new tokens, 3 fresh prompts each +- agent session: system prompt and four tools, then 5 turns that each add about 1500 words of a project file and continue from the model's own earlier answers (the way agent clients resend a conversation), 64 new tokens per turn, 3 sessions; the cache numbers come from the answer's usage (`cached_tokens`) and timings (`prompt_n`) +- concurrency (when run): N requests sent at the same moment, total tokens over the wall time of the burst, latency over all requests + +## Results + +| Configuration | Actual prompt tokens | Reused tokens | Generated tokens | Runs | Prompt tok/s median and range | Decode tok/s median and range | TTFT seconds median and range | +| --- | ---: | ---: | ---: | ---: | --- | --- | --- | +| short chat | 46 (43-51) | 0 | 100 (100-100) | 9 | not measured | 244.2 (206.7-265.7) | 0.14 (0.12-0.17) | +| 4K prompt, fresh | 4019 (4019-4019) | 0 (0-0) | 24 cap | 3 | 5522.9 (5497.2-5782.7) | not measured | 0.75 (0.71-0.76) | +| 16K prompt, fresh | 16058 (16058-16058) | 0 (0-0) | 24 cap | 3 | 5983.1 (5975.1-6078.2) | not measured | 2.71 (2.67-2.71) | +| 32K prompt, fresh | 32049 (32049-32049) | 0 (0-0) | 24 cap | 3 | 4848.8 (4841.5-4887.5) | not measured | 6.65 (6.59-6.66) | +| agent session, follow-up turns | 14559 (14551-14579) at the last turn | see below | up to 64 | 3 | see below | 226.2 (206.9-282.5) | 0.75 (0.48-0.85) | + +Ranges are the lowest and highest of the runs; `Runs` counts requests that returned a token. + +### Agent session (conversation cache) + +Cache hit over the follow-up turns: 68.2 % (68.2-68.2); earlier-turn tokens read again per session: 264 (257-284) (earlier turns that the cache did not hold; the 13721 (13720-13721) new tokens of the files are read once); first turn (nothing cached) first token: 0.49 s (0.47-0.49). + +| Turn (run 1) | Prompt tokens | Reused | Read | Prompt tok/s | First token (s) | Decode tok/s | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | 609 | 0 | 609 | 1344.4 | 0.473 | 175.7 | +| 1 | 3779 | 602 | 3177 | 5318.0 | 0.612 | 208.5 | +| 2 | 5308 | 3772 | 1536 | 2965.3 | 0.541 | 217.4 | +| 3 | 8740 | 5301 | 3439 | 4251.5 | 0.834 | 230.0 | +| 4 | 11647 | 8733 | 2914 | 4038.8 | 0.74 | 213.4 | +| 5 | 14559 | 11640 | 2919 | 3575.0 | 0.854 | 206.9 | + +Per-request data: `results.json`. Engine log lines from the start: `engine-log-excerpt.txt` (when this run started the server). + +## Correctness and limitations + +Not measured: this is a speed suite. The conversation uses the model's own replies, so token counts differ a little between runs. Machine details come from the server, `nvidia-smi`/`rocm-smi` and the OS; background load and power limits are not recorded. Total time of the suite: 54 s. diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/engine-log-excerpt.txt b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/engine-log-excerpt.txt new file mode 100644 index 00000000000..a30ca727265 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/engine-log-excerpt.txt @@ -0,0 +1,5 @@ +[strata] loading the experts into RAM (about 19 GB) and locking part of them for the GPU. + Please wait and don't close this window; the browser opens when it is ready. +[strata] filling the GPU's expert cache (5362 experts, 8.35 GiB of VRAM) ... +[strata] almost ready ... +ready: http://127.0.0.1:8096/v1 (OpenAI: /v1/chat/completions, Anthropic: /v1/messages, context 262144 tokens) diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/engine-start.txt b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/engine-start.txt new file mode 100644 index 00000000000..5d450f3cc33 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/engine-start.txt @@ -0,0 +1,5 @@ +strata mtp: draft layer loaded from the model file, 704 MiB of VRAM (experts 408, dense 19), files read in 0.61 s (696 MiB/s) +strata generate: expert cache auto: +393 MiB kept for the dense attention at 262144 tokens of context +strata generate: expert cache 5362 slots, 8.35 GiB of VRAM; policy is +strata verify: window up to 6 tokens, 441.8 MiB of device buffers +strata serve: 0 MiB of VRAM free with everything loaded - LOW: requests may stall; add --vram-reserve-mib 1212 to the config's args (or lower --max-context) diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/results.json b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/results.json new file mode 100644 index 00000000000..9cd301a7b83 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/A/results.json @@ -0,0 +1,1003 @@ +{ + "schema": 1, + "tool": "tools/strata_bench.py", + "date": "2026-10-13", + "params": { + "runs": 3, + "think": false, + "clients": 0, + "long": true, + "started_server": true, + "command": "python tools/strata_bench.py --long" + }, + "elapsed_s": 54.29803681373596, + "machine": { + "gpu": [ + "NVIDIA GeForce RTX 5080" + ], + "gpu_vram_mib": 16303, + "gpu_driver": "617.42", + "cuda": null, + "cpu": "Intel(R) Core(TM) i9-14900K", + "cpu_cores": 24, + "cpu_threads": 32, + "ram_gib": 47.7, + "os": "Windows 11 (10.0.28000)", + "python": "3.14.3", + "strata_version": "0.1.43.1", + "strata_commit": "f1e33fb9", + "engine_version": "0.1.43.1", + "model": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "quant": "IQ4_XS", + "context": 262144, + "parallel_slots": 1, + "config": { + "exe": "strata.exe", + "args": [ + "--pack", + "ornith-abl-iq4_xs", + "--native", + "Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf", + "--expert-profile", + "expert-profile-qwen36.bin", + "--expert-cache", + "auto", + "--prefill", + "auto", + "--spec", + "4", + "--mtp", + "Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf", + "--mtp-draft-vocab", + "draft_vocab_cyrillic.bin", + "--max-context", + "262144", + "--kv", + "int8", + "--expert-profile-save", + "expert-profile-ornith-learned.bin", + "--pcie-frac", + "0.00", + "--spec-min-p", + "0.50", + "--pool-workers", + "7", + "--adapt-every", + "1", + "--adapt-swaps", + "160", + "--adapt-decay", + "0.97" + ], + "model_name": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "lib_dirs": [ + "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v13.4\\bin", + "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v13.4\\bin\\x64" + ], + "port": 8096, + "gpu": 0, + "gpus_asked": true, + "open_browser": false, + "idle_unload_s": 3600 + }, + "launch": "strata.exe --pack ornith-abl-iq4_xs --native Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --expert-profile expert-profile-qwen36.bin --expert-cache auto --prefill auto --spec 4 --mtp Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --mtp-draft-vocab draft_vocab_cyrillic.bin --max-context 262144 --kv int8 --expert-profile-save expert-profile-ornith-learned.bin --pcie-frac 0.00 --spec-min-p 0.50 --pool-workers 7 --adapt-every 1 --adapt-swaps 160 --adapt-decay 0.97", + "engine_info": { + "model": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "max_context": 262144, + "images": false, + "version": "0.1.43.1", + "logit_bias": 1, + "context": 262144, + "kv": "int8", + "kv_resident": 0, + "expert_slots": 5362, + "expert_cache_mib": 8545, + "expert_slots_primary": 5362, + "expert_cache_primary_mib": 8545, + "spec": 6, + "mtp_max": 4, + "lookup": 3, + "vram_free_mib": 0, + "cvec": 0, + "arena_mib": 16320, + "pool_workers": 7, + "pcie_frac": "0.00", + "spec_min_p": "0.50", + "conversation_cache_mib": 0, + "conversation_cache_slots": 4, + "conversation_cache_min_free_mib": 2560, + "conversation_cache_min_tokens": 0, + "tail_role_token": -1, + "vram_elastic": 0, + "engine": "0.1.43.1" + }, + "vram_after": { + "elastic": false + } + }, + "results": { + "chat": { + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 0.168, + "total_s": 0.595, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 340.2, + "prompt_tps_source": "engine", + "decode_tps": 231.9, + "decode_tps_source": "engine", + "prompt_ms": 149.9, + "decode_ms": 431.2, + "drafts_offered": 87, + "drafts_accepted": 63, + "run": 0, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.14, + "total_s": 0.619, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 328.5, + "prompt_tps_source": "engine", + "decode_tps": 206.7, + "decode_tps_source": "engine", + "prompt_ms": 130.9, + "decode_ms": 483.9, + "drafts_offered": 76, + "drafts_accepted": 51, + "run": 0, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.149, + "total_s": 0.547, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 353.0, + "prompt_tps_source": "engine", + "decode_tps": 247.6, + "decode_tps_source": "engine", + "prompt_ms": 130.3, + "decode_ms": 403.8, + "drafts_offered": 84, + "drafts_accepted": 55, + "run": 0, + "prompt": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.138, + "total_s": 0.542, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 394.4, + "prompt_tps_source": "engine", + "decode_tps": 244.8, + "decode_tps_source": "engine", + "prompt_ms": 129.3, + "decode_ms": 408.5, + "drafts_offered": 87, + "drafts_accepted": 63, + "run": 1, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.119, + "total_s": 0.528, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 388.8, + "prompt_tps_source": "engine", + "decode_tps": 241.8, + "decode_tps_source": "engine", + "prompt_ms": 110.6, + "decode_ms": 413.6, + "drafts_offered": 64, + "drafts_accepted": 51, + "run": 1, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.137, + "total_s": 0.515, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 388.5, + "prompt_tps_source": "engine", + "decode_tps": 260.9, + "decode_tps_source": "engine", + "prompt_ms": 118.4, + "decode_ms": 383.3, + "drafts_offered": 81, + "drafts_accepted": 58, + "run": 1, + "prompt": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.143, + "total_s": 0.548, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 379.2, + "prompt_tps_source": "engine", + "decode_tps": 244.2, + "decode_tps_source": "engine", + "prompt_ms": 134.5, + "decode_ms": 409.5, + "drafts_offered": 84, + "drafts_accepted": 64, + "run": 2, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.121, + "total_s": 0.566, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 383.9, + "prompt_tps_source": "engine", + "decode_tps": 222.1, + "decode_tps_source": "engine", + "prompt_ms": 112.0, + "decode_ms": 450.3, + "drafts_offered": 73, + "drafts_accepted": 45, + "run": 2, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.126, + "total_s": 0.498, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 390.2, + "prompt_tps_source": "engine", + "decode_tps": 265.7, + "decode_tps_source": "engine", + "prompt_ms": 117.9, + "decode_ms": 376.4, + "drafts_offered": 73, + "drafts_accepted": 52, + "run": 2, + "prompt": 2 + } + ], + "decode_tps": { + "median": 244.2, + "min": 206.7, + "max": 265.7, + "n": 9 + }, + "ttft_s": { + "median": 0.138, + "min": 0.119, + "max": 0.168, + "n": 9 + }, + "generated_tokens": { + "median": 100, + "min": 100, + "max": 100, + "n": 9 + }, + "prompt_tokens": { + "median": 46, + "min": 43, + "max": 51, + "n": 9 + } + }, + "prompts": { + "4K": { + "target_tokens": 4096, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 0.706, + "total_s": 0.81, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 5782.7, + "prompt_tps_source": "engine", + "decode_tps": 219.8, + "decode_tps_source": "engine", + "prompt_ms": 695.0, + "decode_ms": 109.2, + "drafts_offered": 16, + "drafts_accepted": 11, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.751, + "total_s": 0.85, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 5497.2, + "prompt_tps_source": "engine", + "decode_tps": 230.8, + "decode_tps_source": "engine", + "prompt_ms": 731.1, + "decode_ms": 104.0, + "drafts_offered": 16, + "drafts_accepted": 11, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.758, + "total_s": 0.854, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 5522.9, + "prompt_tps_source": "engine", + "decode_tps": 236.7, + "decode_tps_source": "engine", + "prompt_ms": 727.7, + "decode_ms": 101.4, + "drafts_offered": 15, + "drafts_accepted": 12, + "run": 2 + } + ], + "prompt_tps": { + "median": 5522.9, + "min": 5497.2, + "max": 5782.7, + "n": 3 + }, + "ttft_s": { + "median": 0.751, + "min": 0.706, + "max": 0.758, + "n": 3 + }, + "prompt_tokens": { + "median": 4019, + "min": 4019, + "max": 4019, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + }, + "16K": { + "target_tokens": 16384, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 2.715, + "total_s": 2.818, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 5975.1, + "prompt_tps_source": "engine", + "decode_tps": 220.4, + "decode_tps_source": "engine", + "prompt_ms": 2687.5, + "decode_ms": 108.9, + "drafts_offered": 13, + "drafts_accepted": 12, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 2.711, + "total_s": 2.848, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 5983.1, + "prompt_tps_source": "engine", + "decode_tps": 167.6, + "decode_tps_source": "engine", + "prompt_ms": 2683.9, + "decode_ms": 143.2, + "drafts_offered": 14, + "drafts_accepted": 8, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 2.669, + "total_s": 2.809, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6078.2, + "prompt_tps_source": "engine", + "decode_tps": 163.9, + "decode_tps_source": "engine", + "prompt_ms": 2641.9, + "decode_ms": 146.4, + "drafts_offered": 17, + "drafts_accepted": 9, + "run": 2 + } + ], + "prompt_tps": { + "median": 5983.1, + "min": 5975.1, + "max": 6078.2, + "n": 3 + }, + "ttft_s": { + "median": 2.711, + "min": 2.669, + "max": 2.715, + "n": 3 + }, + "prompt_tokens": { + "median": 16058, + "min": 16058, + "max": 16058, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + }, + "32K": { + "target_tokens": 32768, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 6.594, + "total_s": 6.74, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 4887.5, + "prompt_tps_source": "engine", + "decode_tps": 156.4, + "decode_tps_source": "engine", + "prompt_ms": 6557.4, + "decode_ms": 153.5, + "drafts_offered": 13, + "drafts_accepted": 12, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 6.656, + "total_s": 6.798, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 4841.5, + "prompt_tps_source": "engine", + "decode_tps": 160.4, + "decode_tps_source": "engine", + "prompt_ms": 6619.7, + "decode_ms": 149.6, + "drafts_offered": 13, + "drafts_accepted": 10, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 6.647, + "total_s": 6.822, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 4848.8, + "prompt_tps_source": "engine", + "decode_tps": 131.7, + "decode_tps_source": "engine", + "prompt_ms": 6609.7, + "decode_ms": 182.2, + "drafts_offered": 23, + "drafts_accepted": 12, + "run": 2 + } + ], + "prompt_tps": { + "median": 4848.8, + "min": 4841.5, + "max": 4887.5, + "n": 3 + }, + "ttft_s": { + "median": 6.647, + "min": 6.594, + "max": 6.656, + "n": 3 + }, + "prompt_tokens": { + "median": 32049, + "min": 32049, + "max": 32049, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + } + }, + "agent": { + "sessions": [ + [ + { + "ok": true, + "err": null, + "ttft_s": 0.473, + "total_s": 0.531, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 11, + "prompt_tps": 1344.4, + "prompt_tps_source": "engine", + "decode_tps": 175.7, + "decode_tps_source": "engine", + "prompt_ms": 453.0, + "decode_ms": 62.6, + "drafts_offered": 15, + "drafts_accepted": 7, + "run": 0, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.612, + "total_s": 0.914, + "prompt_tokens": 3779, + "reused_tokens": 602, + "read_tokens": 3177, + "generated_tokens": 64, + "prompt_tps": 5318.0, + "prompt_tps_source": "engine", + "decode_tps": 208.5, + "decode_tps_source": "engine", + "prompt_ms": 597.4, + "decode_ms": 306.9, + "drafts_offered": 61, + "drafts_accepted": 31, + "run": 0, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.541, + "total_s": 0.798, + "prompt_tokens": 5308, + "reused_tokens": 3772, + "read_tokens": 1536, + "generated_tokens": 57, + "prompt_tps": 2965.3, + "prompt_tps_source": "engine", + "decode_tps": 217.4, + "decode_tps_source": "engine", + "prompt_ms": 518.0, + "decode_ms": 262.2, + "drafts_offered": 49, + "drafts_accepted": 30, + "run": 0, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.834, + "total_s": 1.055, + "prompt_tokens": 8740, + "reused_tokens": 5301, + "read_tokens": 3439, + "generated_tokens": 52, + "prompt_tps": 4251.5, + "prompt_tps_source": "engine", + "decode_tps": 230.0, + "decode_tps_source": "engine", + "prompt_ms": 808.9, + "decode_ms": 226.1, + "drafts_offered": 42, + "drafts_accepted": 30, + "run": 0, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.74, + "total_s": 0.945, + "prompt_tokens": 11647, + "reused_tokens": 8733, + "read_tokens": 2914, + "generated_tokens": 45, + "prompt_tps": 4038.8, + "prompt_tps_source": "engine", + "decode_tps": 213.4, + "decode_tps_source": "engine", + "prompt_ms": 721.5, + "decode_ms": 210.9, + "drafts_offered": 45, + "drafts_accepted": 21, + "run": 0, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.854, + "total_s": 1.109, + "prompt_tokens": 14559, + "reused_tokens": 11640, + "read_tokens": 2919, + "generated_tokens": 54, + "prompt_tps": 3575.0, + "prompt_tps_source": "engine", + "decode_tps": 206.9, + "decode_tps_source": "engine", + "prompt_ms": 816.5, + "decode_ms": 261.0, + "drafts_offered": 38, + "drafts_accepted": 28, + "run": 0, + "turn": 5 + } + ], + [ + { + "ok": true, + "err": null, + "ttft_s": 0.487, + "total_s": 0.536, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 11, + "prompt_tps": 1304.1, + "prompt_tps_source": "engine", + "decode_tps": 203.0, + "decode_tps_source": "engine", + "prompt_ms": 467.0, + "decode_ms": 54.2, + "drafts_offered": 14, + "drafts_accepted": 8, + "run": 1, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.67, + "total_s": 0.877, + "prompt_tokens": 3779, + "reused_tokens": 602, + "read_tokens": 3177, + "generated_tokens": 60, + "prompt_tps": 4823.9, + "prompt_tps_source": "engine", + "decode_tps": 282.5, + "decode_tps_source": "engine", + "prompt_ms": 658.6, + "decode_ms": 212.4, + "drafts_offered": 45, + "drafts_accepted": 33, + "run": 1, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.481, + "total_s": 0.7, + "prompt_tokens": 5303, + "reused_tokens": 3772, + "read_tokens": 1531, + "generated_tokens": 51, + "prompt_tps": 3263.7, + "prompt_tps_source": "engine", + "decode_tps": 227.3, + "decode_tps_source": "engine", + "prompt_ms": 469.1, + "decode_ms": 224.4, + "drafts_offered": 41, + "drafts_accepted": 22, + "run": 1, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.805, + "total_s": 1.017, + "prompt_tokens": 8729, + "reused_tokens": 5296, + "read_tokens": 3433, + "generated_tokens": 49, + "prompt_tps": 4342.8, + "prompt_tps_source": "engine", + "decode_tps": 226.2, + "decode_tps_source": "engine", + "prompt_ms": 790.5, + "decode_ms": 216.6, + "drafts_offered": 31, + "drafts_accepted": 22, + "run": 1, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.785, + "total_s": 1.007, + "prompt_tokens": 11633, + "reused_tokens": 8722, + "read_tokens": 2911, + "generated_tokens": 51, + "prompt_tps": 3806.2, + "prompt_tps_source": "engine", + "decode_tps": 224.2, + "decode_tps_source": "engine", + "prompt_ms": 764.8, + "decode_ms": 227.5, + "drafts_offered": 47, + "drafts_accepted": 30, + "run": 1, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.776, + "total_s": 1.066, + "prompt_tokens": 14551, + "reused_tokens": 11626, + "read_tokens": 2925, + "generated_tokens": 64, + "prompt_tps": 3874.7, + "prompt_tps_source": "engine", + "decode_tps": 216.9, + "decode_tps_source": "engine", + "prompt_ms": 754.9, + "decode_ms": 295.1, + "drafts_offered": 51, + "drafts_accepted": 32, + "run": 1, + "turn": 5 + } + ], + [ + { + "ok": true, + "err": null, + "ttft_s": 0.487, + "total_s": 0.583, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 25, + "prompt_tps": 1301.3, + "prompt_tps_source": "engine", + "decode_tps": 247.0, + "decode_tps_source": "engine", + "prompt_ms": 468.0, + "decode_ms": 101.2, + "drafts_offered": 26, + "drafts_accepted": 15, + "run": 2, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.685, + "total_s": 0.873, + "prompt_tokens": 3793, + "reused_tokens": 602, + "read_tokens": 3191, + "generated_tokens": 54, + "prompt_tps": 4805.0, + "prompt_tps_source": "engine", + "decode_tps": 279.6, + "decode_tps_source": "engine", + "prompt_ms": 664.1, + "decode_ms": 193.1, + "drafts_offered": 44, + "drafts_accepted": 31, + "run": 2, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.524, + "total_s": 0.791, + "prompt_tokens": 5311, + "reused_tokens": 3786, + "read_tokens": 1525, + "generated_tokens": 64, + "prompt_tps": 2978.5, + "prompt_tps_source": "engine", + "decode_tps": 235.6, + "decode_tps_source": "engine", + "prompt_ms": 512.0, + "decode_ms": 271.7, + "drafts_offered": 50, + "drafts_accepted": 29, + "run": 2, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.753, + "total_s": 0.982, + "prompt_tokens": 8751, + "reused_tokens": 5304, + "read_tokens": 3447, + "generated_tokens": 54, + "prompt_tps": 4667.6, + "prompt_tps_source": "engine", + "decode_tps": 230.4, + "decode_tps_source": "engine", + "prompt_ms": 738.5, + "decode_ms": 234.4, + "drafts_offered": 37, + "drafts_accepted": 28, + "run": 2, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.783, + "total_s": 1.005, + "prompt_tokens": 11660, + "reused_tokens": 8744, + "read_tokens": 2916, + "generated_tokens": 52, + "prompt_tps": 3803.8, + "prompt_tps_source": "engine", + "decode_tps": 228.2, + "decode_tps_source": "engine", + "prompt_ms": 766.6, + "decode_ms": 227.9, + "drafts_offered": 44, + "drafts_accepted": 27, + "run": 2, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.839, + "total_s": 1.123, + "prompt_tokens": 14579, + "reused_tokens": 11653, + "read_tokens": 2926, + "generated_tokens": 64, + "prompt_tps": 3565.7, + "prompt_tps_source": "engine", + "decode_tps": 221.0, + "decode_tps_source": "engine", + "prompt_ms": 820.6, + "decode_ms": 289.6, + "drafts_offered": 50, + "drafts_accepted": 31, + "run": 2, + "turn": 5 + } + ] + ], + "hit_pct": { + "median": 68.238, + "min": 68.23, + "max": 68.24, + "n": 3 + }, + "reread_tokens": { + "median": 264, + "min": 257, + "max": 284, + "n": 3 + }, + "new_tokens": { + "median": 13721, + "min": 13720, + "max": 13721, + "n": 3 + }, + "turn_ttft_s": { + "median": 0.753, + "min": 0.481, + "max": 0.854, + "n": 15 + }, + "last_turn_prompt_tokens": { + "median": 14559, + "min": 14551, + "max": 14579, + "n": 3 + }, + "turn_decode_tps": { + "median": 226.2, + "min": 206.9, + "max": 282.5, + "n": 15 + }, + "first_turn_ttft_s": { + "median": 0.487, + "min": 0.473, + "max": 0.487, + "n": 3 + } + }, + "clients": {}, + "skipped": {}, + "failures": [] + } +} \ No newline at end of file diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/README.md b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/README.md new file mode 100644 index 00000000000..b2eed039800 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/README.md @@ -0,0 +1,63 @@ +# Community benchmark on NVIDIA GeForce RTX 5080 + +Measured on 2026-10-13 with `tools/strata_bench.py` (suite 1, 3 runs per test, medians with ranges). Speeds are the engine's own timings; first-token times are measured by the client over HTTP. Speed only: no correctness check was run. + +## Hardware and software + +- GPU and VRAM: NVIDIA GeForce RTX 5080, 16 GiB VRAM +- CPU: Intel(R) Core(TM) i9-14900K (24 cores, 32 threads); installed RAM: 47.7 GiB; storage: not measured +- OS: Windows 11 (10.0.28000); GPU driver: 617.42; CUDA: not read +- Strata: project 0.1.43.1, commit f1e33fb9; engine 0.1.43.1 +- Background workloads and power limits: not recorded + +## Model and configuration + +- Model and quantization: ornith-1.5-35b-a3b-abliterated-iq4_xs +- Context: 262144 tokens; parallel slots: 1 +- Settings: the install's own configuration, unchanged; thinking off for the suite; temperature 0 + +```text +strata.exe --pack ornith-abl-iq4_xs --native Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --expert-profile expert-profile-qwen36.bin --expert-cache auto --prefill auto --spec 4 --mtp Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --mtp-draft-vocab draft_vocab_cyrillic.bin --max-context 262144 --kv int8 --expert-profile-save expert-profile-ornith-learned.bin --pcie-frac 0.00 --spec-min-p 0.50 --pool-workers 7 --adapt-every 1 --adapt-swaps 160 --adapt-decay 0.97 +``` + +Environment of the config: `STRATA_ATTN_SPLITS=256` + +## Method + +Run by `python tools/strata_bench.py --long` against the HTTP API of a Strata server (started by the benchmark from the install configuration); model loading is not timed. One warm-up chat and a 1K prompt are discarded first. Every prompt starts with a unique tag, so the conversation cache cannot help the fresh-prompt tests. + +- short chat: the three questions of `tools/ab_engine.py`, up to 100 new tokens, 3 runs +- 4K / 16K prompts: generated word lists sized by a calibration request, 24 new tokens, 3 fresh prompts each +- agent session: system prompt and four tools, then 5 turns that each add about 1500 words of a project file and continue from the model's own earlier answers (the way agent clients resend a conversation), 64 new tokens per turn, 3 sessions; the cache numbers come from the answer's usage (`cached_tokens`) and timings (`prompt_n`) +- concurrency (when run): N requests sent at the same moment, total tokens over the wall time of the burst, latency over all requests + +## Results + +| Configuration | Actual prompt tokens | Reused tokens | Generated tokens | Runs | Prompt tok/s median and range | Decode tok/s median and range | TTFT seconds median and range | +| --- | ---: | ---: | ---: | ---: | --- | --- | --- | +| short chat | 46 (43-51) | 0 | 100 (100-100) | 9 | not measured | 245.7 (192.7-270.1) | 0.14 (0.11-0.17) | +| 4K prompt, fresh | 4019 (4019-4019) | 0 (0-0) | 24 cap | 3 | 6300.4 (6021.9-6321.2) | not measured | 0.66 (0.66-0.69) | +| 16K prompt, fresh | 16058 (16058-16058) | 0 (0-0) | 24 cap | 3 | 6363.9 (6303.7-6464.1) | not measured | 2.55 (2.51-2.57) | +| 32K prompt, fresh | 32049 (32049-32049) | 0 (0-0) | 24 cap | 3 | 5219.6 (5199.6-5242.8) | not measured | 6.17 (6.15-6.20) | +| agent session, follow-up turns | 14564 (14538-14569) at the last turn | see below | up to 64 | 3 | see below | 245.9 (181.0-282.3) | 0.70 (0.48-0.82) | + +Ranges are the lowest and highest of the runs; `Runs` counts requests that returned a token. + +### Agent session (conversation cache) + +Cache hit over the follow-up turns: 68.2 % (68.2-68.2); earlier-turn tokens read again per session: 270 (244-275) (earlier turns that the cache did not hold; the 13720 (13720-13720) new tokens of the files are read once); first turn (nothing cached) first token: 0.45 s (0.45-0.46). + +| Turn (run 1) | Prompt tokens | Reused | Read | Prompt tok/s | First token (s) | Decode tok/s | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | 609 | 0 | 609 | 1431.3 | 0.445 | 210.3 | +| 1 | 3779 | 602 | 3177 | 5340.4 | 0.619 | 181.0 | +| 2 | 5294 | 3772 | 1522 | 3204.2 | 0.488 | 230.6 | +| 3 | 8723 | 5287 | 3436 | 4777.5 | 0.744 | 269.8 | +| 4 | 11625 | 8716 | 2909 | 4130.9 | 0.73 | 229.2 | +| 5 | 14538 | 11618 | 2920 | 3683.6 | 0.821 | 245.9 | + +Per-request data: `results.json`. Engine log lines from the start: `engine-log-excerpt.txt` (when this run started the server). + +## Correctness and limitations + +Not measured: this is a speed suite. The conversation uses the model's own replies, so token counts differ a little between runs. Machine details come from the server, `nvidia-smi`/`rocm-smi` and the OS; background load and power limits are not recorded. Total time of the suite: 51 s. diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/engine-log-excerpt.txt b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/engine-log-excerpt.txt new file mode 100644 index 00000000000..12ed40927d5 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/engine-log-excerpt.txt @@ -0,0 +1,5 @@ +[strata] loading the experts into RAM (about 19 GB) and locking part of them for the GPU. + Please wait and don't close this window; the browser opens when it is ready. +[strata] filling the GPU's expert cache (5378 experts, 8.37 GiB of VRAM) ... +[strata] almost ready ... +ready: http://127.0.0.1:8096/v1 (OpenAI: /v1/chat/completions, Anthropic: /v1/messages, context 262144 tokens) diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/engine-start.txt b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/engine-start.txt new file mode 100644 index 00000000000..d0c8b56631c --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/engine-start.txt @@ -0,0 +1,5 @@ +strata mtp: draft layer loaded from the model file, 700 MiB of VRAM (experts 408, dense 19), files read in 0.65 s (654 MiB/s) +strata generate: expert cache auto: +30 MiB kept for the dense attention at 262144 tokens of context (split-K partials, STRATA_ATTN_SPLITS) +strata generate: expert cache 5378 slots, 8.37 GiB of VRAM; policy is +strata verify: window up to 6 tokens, 78.9 MiB of device buffers +strata serve: 284 MiB of VRAM free with everything loaded diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/results.json b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/results.json new file mode 100644 index 00000000000..7d301227c4c --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/B/results.json @@ -0,0 +1,1006 @@ +{ + "schema": 1, + "tool": "tools/strata_bench.py", + "date": "2026-10-13", + "params": { + "runs": 3, + "think": false, + "clients": 0, + "long": true, + "started_server": true, + "command": "python tools/strata_bench.py --long" + }, + "elapsed_s": 50.87422323226929, + "machine": { + "gpu": [ + "NVIDIA GeForce RTX 5080" + ], + "gpu_vram_mib": 16303, + "gpu_driver": "617.42", + "cuda": null, + "cpu": "Intel(R) Core(TM) i9-14900K", + "cpu_cores": 24, + "cpu_threads": 32, + "ram_gib": 47.7, + "os": "Windows 11 (10.0.28000)", + "python": "3.14.3", + "strata_version": "0.1.43.1", + "strata_commit": "f1e33fb9", + "engine_version": "0.1.43.1", + "model": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "quant": "IQ4_XS", + "context": 262144, + "parallel_slots": 1, + "config": { + "exe": "strata.exe", + "args": [ + "--pack", + "ornith-abl-iq4_xs", + "--native", + "Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf", + "--expert-profile", + "expert-profile-qwen36.bin", + "--expert-cache", + "auto", + "--prefill", + "auto", + "--spec", + "4", + "--mtp", + "Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf", + "--mtp-draft-vocab", + "draft_vocab_cyrillic.bin", + "--max-context", + "262144", + "--kv", + "int8", + "--expert-profile-save", + "expert-profile-ornith-learned.bin", + "--pcie-frac", + "0.00", + "--spec-min-p", + "0.50", + "--pool-workers", + "7", + "--adapt-every", + "1", + "--adapt-swaps", + "160", + "--adapt-decay", + "0.97" + ], + "model_name": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "lib_dirs": [ + "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v13.4\\bin", + "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v13.4\\bin\\x64" + ], + "port": 8096, + "gpu": 0, + "gpus_asked": true, + "open_browser": false, + "idle_unload_s": 3600, + "env": { + "STRATA_ATTN_SPLITS": "256" + } + }, + "launch": "strata.exe --pack ornith-abl-iq4_xs --native Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --expert-profile expert-profile-qwen36.bin --expert-cache auto --prefill auto --spec 4 --mtp Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --mtp-draft-vocab draft_vocab_cyrillic.bin --max-context 262144 --kv int8 --expert-profile-save expert-profile-ornith-learned.bin --pcie-frac 0.00 --spec-min-p 0.50 --pool-workers 7 --adapt-every 1 --adapt-swaps 160 --adapt-decay 0.97", + "engine_info": { + "model": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "max_context": 262144, + "images": false, + "version": "0.1.43.1", + "logit_bias": 1, + "context": 262144, + "kv": "int8", + "kv_resident": 0, + "expert_slots": 5378, + "expert_cache_mib": 8571, + "expert_slots_primary": 5378, + "expert_cache_primary_mib": 8571, + "spec": 6, + "mtp_max": 4, + "lookup": 3, + "vram_free_mib": 284, + "cvec": 0, + "arena_mib": 16320, + "pool_workers": 7, + "pcie_frac": "0.00", + "spec_min_p": "0.50", + "conversation_cache_mib": 0, + "conversation_cache_slots": 4, + "conversation_cache_min_free_mib": 2560, + "conversation_cache_min_tokens": 0, + "tail_role_token": -1, + "vram_elastic": 0, + "engine": "0.1.43.1" + }, + "vram_after": { + "elastic": false + } + }, + "results": { + "chat": { + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 0.168, + "total_s": 0.593, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 341.8, + "prompt_tps_source": "engine", + "decode_tps": 232.1, + "decode_tps_source": "engine", + "prompt_ms": 149.2, + "decode_ms": 430.9, + "drafts_offered": 86, + "drafts_accepted": 64, + "run": 0, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.14, + "total_s": 0.621, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 330.3, + "prompt_tps_source": "engine", + "decode_tps": 205.6, + "decode_tps_source": "engine", + "prompt_ms": 130.2, + "decode_ms": 486.3, + "drafts_offered": 67, + "drafts_accepted": 51, + "run": 0, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.14, + "total_s": 0.567, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 351.4, + "prompt_tps_source": "engine", + "decode_tps": 231.2, + "decode_tps_source": "engine", + "prompt_ms": 130.9, + "decode_ms": 432.6, + "drafts_offered": 73, + "drafts_accepted": 51, + "run": 0, + "prompt": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.154, + "total_s": 0.667, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 377.8, + "prompt_tps_source": "engine", + "decode_tps": 192.7, + "decode_tps_source": "engine", + "prompt_ms": 135.0, + "decode_ms": 518.9, + "drafts_offered": 86, + "drafts_accepted": 63, + "run": 1, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.121, + "total_s": 0.486, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 384.3, + "prompt_tps_source": "engine", + "decode_tps": 270.1, + "decode_tps_source": "engine", + "prompt_ms": 111.9, + "decode_ms": 370.3, + "drafts_offered": 64, + "drafts_accepted": 53, + "run": 1, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.123, + "total_s": 0.492, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 434.0, + "prompt_tps_source": "engine", + "decode_tps": 268.3, + "decode_tps_source": "engine", + "prompt_ms": 106.0, + "decode_ms": 372.7, + "drafts_offered": 72, + "drafts_accepted": 52, + "run": 1, + "prompt": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.14, + "total_s": 0.524, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 418.4, + "prompt_tps_source": "engine", + "decode_tps": 257.1, + "decode_tps_source": "engine", + "prompt_ms": 121.9, + "decode_ms": 389.0, + "drafts_offered": 84, + "drafts_accepted": 64, + "run": 2, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.111, + "total_s": 0.504, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 421.2, + "prompt_tps_source": "engine", + "decode_tps": 251.3, + "decode_tps_source": "engine", + "prompt_ms": 102.1, + "decode_ms": 397.9, + "drafts_offered": 74, + "drafts_accepted": 50, + "run": 2, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.122, + "total_s": 0.525, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 410.7, + "prompt_tps_source": "engine", + "decode_tps": 245.7, + "decode_tps_source": "engine", + "prompt_ms": 112.0, + "decode_ms": 407.0, + "drafts_offered": 81, + "drafts_accepted": 53, + "run": 2, + "prompt": 2 + } + ], + "decode_tps": { + "median": 245.7, + "min": 192.7, + "max": 270.1, + "n": 9 + }, + "ttft_s": { + "median": 0.14, + "min": 0.111, + "max": 0.168, + "n": 9 + }, + "generated_tokens": { + "median": 100, + "min": 100, + "max": 100, + "n": 9 + }, + "prompt_tokens": { + "median": 46, + "min": 43, + "max": 51, + "n": 9 + } + }, + "prompts": { + "4K": { + "target_tokens": 4096, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 0.687, + "total_s": 0.779, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 6021.9, + "prompt_tps_source": "engine", + "decode_tps": 248.7, + "decode_tps_source": "engine", + "prompt_ms": 667.4, + "decode_ms": 96.5, + "drafts_offered": 16, + "drafts_accepted": 13, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.658, + "total_s": 0.738, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 6300.4, + "prompt_tps_source": "engine", + "decode_tps": 283.4, + "decode_tps_source": "engine", + "prompt_ms": 637.9, + "decode_ms": 84.7, + "drafts_offered": 16, + "drafts_accepted": 12, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.656, + "total_s": 0.74, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 6321.2, + "prompt_tps_source": "engine", + "decode_tps": 268.5, + "decode_tps_source": "engine", + "prompt_ms": 635.8, + "decode_ms": 89.4, + "drafts_offered": 14, + "drafts_accepted": 12, + "run": 2 + } + ], + "prompt_tps": { + "median": 6300.4, + "min": 6021.9, + "max": 6321.2, + "n": 3 + }, + "ttft_s": { + "median": 0.658, + "min": 0.656, + "max": 0.687, + "n": 3 + }, + "prompt_tokens": { + "median": 4019, + "min": 4019, + "max": 4019, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + }, + "16K": { + "target_tokens": 16384, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 2.551, + "total_s": 2.691, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6363.9, + "prompt_tps_source": "engine", + "decode_tps": 164.3, + "decode_tps_source": "engine", + "prompt_ms": 2523.3, + "decode_ms": 146.1, + "drafts_offered": 19, + "drafts_accepted": 8, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 2.511, + "total_s": 2.648, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6464.1, + "prompt_tps_source": "engine", + "decode_tps": 168.1, + "decode_tps_source": "engine", + "prompt_ms": 2484.2, + "decode_ms": 142.8, + "drafts_offered": 16, + "drafts_accepted": 8, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 2.574, + "total_s": 2.666, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6303.7, + "prompt_tps_source": "engine", + "decode_tps": 244.2, + "decode_tps_source": "engine", + "prompt_ms": 2547.4, + "decode_ms": 98.3, + "drafts_offered": 13, + "drafts_accepted": 12, + "run": 2 + } + ], + "prompt_tps": { + "median": 6363.9, + "min": 6303.7, + "max": 6464.1, + "n": 3 + }, + "ttft_s": { + "median": 2.551, + "min": 2.511, + "max": 2.574, + "n": 3 + }, + "prompt_tokens": { + "median": 16058, + "min": 16058, + "max": 16058, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + }, + "32K": { + "target_tokens": 32768, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 6.175, + "total_s": 6.301, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 5219.6, + "prompt_tps_source": "engine", + "decode_tps": 181.1, + "decode_tps_source": "engine", + "prompt_ms": 6140.1, + "decode_ms": 132.5, + "drafts_offered": 13, + "drafts_accepted": 10, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 6.148, + "total_s": 6.29, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 5242.8, + "prompt_tps_source": "engine", + "decode_tps": 162.2, + "decode_tps_source": "engine", + "prompt_ms": 6112.9, + "decode_ms": 148.0, + "drafts_offered": 22, + "drafts_accepted": 12, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 6.199, + "total_s": 6.33, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 5199.6, + "prompt_tps_source": "engine", + "decode_tps": 173.9, + "decode_tps_source": "engine", + "prompt_ms": 6163.8, + "decode_ms": 138.0, + "drafts_offered": 20, + "drafts_accepted": 13, + "run": 2 + } + ], + "prompt_tps": { + "median": 5219.6, + "min": 5199.6, + "max": 5242.8, + "n": 3 + }, + "ttft_s": { + "median": 6.175, + "min": 6.148, + "max": 6.199, + "n": 3 + }, + "prompt_tokens": { + "median": 32049, + "min": 32049, + "max": 32049, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + } + }, + "agent": { + "sessions": [ + [ + { + "ok": true, + "err": null, + "ttft_s": 0.445, + "total_s": 0.493, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 11, + "prompt_tps": 1431.3, + "prompt_tps_source": "engine", + "decode_tps": 210.3, + "decode_tps_source": "engine", + "prompt_ms": 425.5, + "decode_ms": 52.3, + "drafts_offered": 10, + "drafts_accepted": 7, + "run": 0, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.619, + "total_s": 0.896, + "prompt_tokens": 3779, + "reused_tokens": 602, + "read_tokens": 3177, + "generated_tokens": 51, + "prompt_tps": 5340.4, + "prompt_tps_source": "engine", + "decode_tps": 181.0, + "decode_tps_source": "engine", + "prompt_ms": 594.9, + "decode_ms": 281.7, + "drafts_offered": 57, + "drafts_accepted": 20, + "run": 0, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.488, + "total_s": 0.718, + "prompt_tokens": 5294, + "reused_tokens": 3772, + "read_tokens": 1522, + "generated_tokens": 54, + "prompt_tps": 3204.2, + "prompt_tps_source": "engine", + "decode_tps": 230.6, + "decode_tps_source": "engine", + "prompt_ms": 475.0, + "decode_ms": 234.2, + "drafts_offered": 39, + "drafts_accepted": 25, + "run": 0, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.744, + "total_s": 0.913, + "prompt_tokens": 8723, + "reused_tokens": 5287, + "read_tokens": 3436, + "generated_tokens": 47, + "prompt_tps": 4777.5, + "prompt_tps_source": "engine", + "decode_tps": 269.8, + "decode_tps_source": "engine", + "prompt_ms": 719.2, + "decode_ms": 174.2, + "drafts_offered": 36, + "drafts_accepted": 24, + "run": 0, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.73, + "total_s": 0.925, + "prompt_tokens": 11625, + "reused_tokens": 8716, + "read_tokens": 2909, + "generated_tokens": 46, + "prompt_tps": 4130.9, + "prompt_tps_source": "engine", + "decode_tps": 229.2, + "decode_tps_source": "engine", + "prompt_ms": 704.2, + "decode_ms": 200.7, + "drafts_offered": 49, + "drafts_accepted": 23, + "run": 0, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.821, + "total_s": 0.974, + "prompt_tokens": 14538, + "reused_tokens": 11618, + "read_tokens": 2920, + "generated_tokens": 39, + "prompt_tps": 3683.6, + "prompt_tps_source": "engine", + "decode_tps": 245.9, + "decode_tps_source": "engine", + "prompt_ms": 792.7, + "decode_ms": 158.6, + "drafts_offered": 35, + "drafts_accepted": 21, + "run": 0, + "turn": 5 + } + ], + [ + { + "ok": true, + "err": null, + "ttft_s": 0.457, + "total_s": 0.529, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 25, + "prompt_tps": 1386.6, + "prompt_tps_source": "engine", + "decode_tps": 325.9, + "decode_tps_source": "engine", + "prompt_ms": 439.2, + "decode_ms": 76.7, + "drafts_offered": 18, + "drafts_accepted": 15, + "run": 1, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.587, + "total_s": 0.774, + "prompt_tokens": 3793, + "reused_tokens": 602, + "read_tokens": 3191, + "generated_tokens": 54, + "prompt_tps": 5533.2, + "prompt_tps_source": "engine", + "decode_tps": 282.3, + "decode_tps_source": "engine", + "prompt_ms": 576.7, + "decode_ms": 191.3, + "drafts_offered": 45, + "drafts_accepted": 28, + "run": 1, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.49, + "total_s": 0.685, + "prompt_tokens": 5311, + "reused_tokens": 3786, + "read_tokens": 1525, + "generated_tokens": 54, + "prompt_tps": 3252.3, + "prompt_tps_source": "engine", + "decode_tps": 269.5, + "decode_tps_source": "engine", + "prompt_ms": 468.9, + "decode_ms": 200.4, + "drafts_offered": 41, + "drafts_accepted": 26, + "run": 1, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.718, + "total_s": 0.89, + "prompt_tokens": 8740, + "reused_tokens": 5304, + "read_tokens": 3436, + "generated_tokens": 50, + "prompt_tps": 4943.9, + "prompt_tps_source": "engine", + "decode_tps": 282.0, + "decode_tps_source": "engine", + "prompt_ms": 695.0, + "decode_ms": 177.3, + "drafts_offered": 40, + "drafts_accepted": 27, + "run": 1, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.693, + "total_s": 0.908, + "prompt_tokens": 11645, + "reused_tokens": 8733, + "read_tokens": 2912, + "generated_tokens": 52, + "prompt_tps": 4297.5, + "prompt_tps_source": "engine", + "decode_tps": 235.9, + "decode_tps_source": "engine", + "prompt_ms": 677.6, + "decode_ms": 220.4, + "drafts_offered": 47, + "drafts_accepted": 27, + "run": 1, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.741, + "total_s": 1.024, + "prompt_tokens": 14564, + "reused_tokens": 11638, + "read_tokens": 2926, + "generated_tokens": 60, + "prompt_tps": 4040.9, + "prompt_tps_source": "engine", + "decode_tps": 208.2, + "decode_tps_source": "engine", + "prompt_ms": 724.1, + "decode_ms": 288.2, + "drafts_offered": 58, + "drafts_accepted": 30, + "run": 1, + "turn": 5 + } + ], + [ + { + "ok": true, + "err": null, + "ttft_s": 0.452, + "total_s": 0.535, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 25, + "prompt_tps": 1404.2, + "prompt_tps_source": "engine", + "decode_tps": 285.4, + "decode_tps_source": "engine", + "prompt_ms": 433.7, + "decode_ms": 87.6, + "drafts_offered": 24, + "drafts_accepted": 16, + "run": 2, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.596, + "total_s": 0.799, + "prompt_tokens": 3793, + "reused_tokens": 602, + "read_tokens": 3191, + "generated_tokens": 53, + "prompt_tps": 5535.1, + "prompt_tps_source": "engine", + "decode_tps": 255.2, + "decode_tps_source": "engine", + "prompt_ms": 576.5, + "decode_ms": 207.7, + "drafts_offered": 49, + "drafts_accepted": 25, + "run": 2, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.484, + "total_s": 0.716, + "prompt_tokens": 5310, + "reused_tokens": 3786, + "read_tokens": 1524, + "generated_tokens": 62, + "prompt_tps": 3270.4, + "prompt_tps_source": "engine", + "decode_tps": 261.3, + "decode_tps_source": "engine", + "prompt_ms": 466.0, + "decode_ms": 237.3, + "drafts_offered": 47, + "drafts_accepted": 31, + "run": 2, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.718, + "total_s": 0.906, + "prompt_tokens": 8747, + "reused_tokens": 5303, + "read_tokens": 3444, + "generated_tokens": 46, + "prompt_tps": 4955.4, + "prompt_tps_source": "engine", + "decode_tps": 238.0, + "decode_tps_source": "engine", + "prompt_ms": 695.0, + "decode_ms": 193.3, + "drafts_offered": 34, + "drafts_accepted": 20, + "run": 2, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.702, + "total_s": 0.906, + "prompt_tokens": 11648, + "reused_tokens": 8740, + "read_tokens": 2908, + "generated_tokens": 54, + "prompt_tps": 4289.7, + "prompt_tps_source": "engine", + "decode_tps": 258.5, + "decode_tps_source": "engine", + "prompt_ms": 677.9, + "decode_ms": 208.9, + "drafts_offered": 52, + "drafts_accepted": 30, + "run": 2, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.745, + "total_s": 1.018, + "prompt_tokens": 14569, + "reused_tokens": 11641, + "read_tokens": 2928, + "generated_tokens": 62, + "prompt_tps": 4022.5, + "prompt_tps_source": "engine", + "decode_tps": 222.8, + "decode_tps_source": "engine", + "prompt_ms": 727.9, + "decode_ms": 278.3, + "drafts_offered": 50, + "drafts_accepted": 31, + "run": 2, + "turn": 5 + } + ] + ], + "hit_pct": { + "median": 68.242, + "min": 68.234, + "max": 68.243, + "n": 3 + }, + "reread_tokens": { + "median": 270, + "min": 244, + "max": 275, + "n": 3 + }, + "new_tokens": { + "median": 13720, + "min": 13720, + "max": 13720, + "n": 3 + }, + "turn_ttft_s": { + "median": 0.702, + "min": 0.484, + "max": 0.821, + "n": 15 + }, + "last_turn_prompt_tokens": { + "median": 14564, + "min": 14538, + "max": 14569, + "n": 3 + }, + "turn_decode_tps": { + "median": 245.9, + "min": 181.0, + "max": 282.3, + "n": 15 + }, + "first_turn_ttft_s": { + "median": 0.452, + "min": 0.445, + "max": 0.457, + "n": 3 + } + }, + "clients": {}, + "skipped": {}, + "failures": [] + } +} \ No newline at end of file diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/README.md b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/README.md new file mode 100644 index 00000000000..ff1800ad074 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/README.md @@ -0,0 +1,63 @@ +# Community benchmark on NVIDIA GeForce RTX 5080 + +Measured on 2026-10-13 with `tools/strata_bench.py` (suite 1, 3 runs per test, medians with ranges). Speeds are the engine's own timings; first-token times are measured by the client over HTTP. Speed only: no correctness check was run. + +## Hardware and software + +- GPU and VRAM: NVIDIA GeForce RTX 5080, 16 GiB VRAM +- CPU: Intel(R) Core(TM) i9-14900K (24 cores, 32 threads); installed RAM: 47.7 GiB; storage: not measured +- OS: Windows 11 (10.0.28000); GPU driver: 617.42; CUDA: not read +- Strata: project 0.1.43.1, commit f1e33fb9; engine 0.1.43.1 +- Background workloads and power limits: not recorded + +## Model and configuration + +- Model and quantization: ornith-1.5-35b-a3b-abliterated-iq4_xs +- Context: 262144 tokens; parallel slots: 1 +- Settings: the install's own configuration, unchanged; thinking off for the suite; temperature 0 + +```text +strata.exe --pack ornith-abl-iq4_xs --native Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --expert-profile expert-profile-qwen36.bin --expert-cache auto --prefill auto --spec 4 --mtp mtp-Ornith-1.5-35B-A3B-Abliterated-MTPv2-Q8_0.gguf --mtp-draft-vocab draft_vocab_cyrillic.bin --max-context 262144 --kv int8 --expert-profile-save expert-profile-ornith-learned.bin --pcie-frac 0.00 --spec-min-p 0.50 --pool-workers 7 --adapt-every 1 --adapt-swaps 160 --adapt-decay 0.97 +``` + +Environment of the config: `STRATA_ATTN_SPLITS=256` + +## Method + +Run by `python tools/strata_bench.py --long` against the HTTP API of a Strata server (started by the benchmark from the install configuration); model loading is not timed. One warm-up chat and a 1K prompt are discarded first. Every prompt starts with a unique tag, so the conversation cache cannot help the fresh-prompt tests. + +- short chat: the three questions of `tools/ab_engine.py`, up to 100 new tokens, 3 runs +- 4K / 16K prompts: generated word lists sized by a calibration request, 24 new tokens, 3 fresh prompts each +- agent session: system prompt and four tools, then 5 turns that each add about 1500 words of a project file and continue from the model's own earlier answers (the way agent clients resend a conversation), 64 new tokens per turn, 3 sessions; the cache numbers come from the answer's usage (`cached_tokens`) and timings (`prompt_n`) +- concurrency (when run): N requests sent at the same moment, total tokens over the wall time of the burst, latency over all requests + +## Results + +| Configuration | Actual prompt tokens | Reused tokens | Generated tokens | Runs | Prompt tok/s median and range | Decode tok/s median and range | TTFT seconds median and range | +| --- | ---: | ---: | ---: | ---: | --- | --- | --- | +| short chat | 46 (43-51) | 0 | 100 (100-100) | 9 | not measured | 242.8 (207.4-255.9) | 0.14 (0.11-0.16) | +| 4K prompt, fresh | 4019 (4019-4019) | 0 (0-0) | 24 cap | 3 | 6289.5 (6096.8-6306.3) | not measured | 0.66 (0.65-0.68) | +| 16K prompt, fresh | 16058 (16058-16058) | 0 (0-0) | 24 cap | 3 | 6474.0 (6375.0-6580.6) | not measured | 2.51 (2.47-2.55) | +| 32K prompt, fresh | 32049 (32049-32049) | 0 (0-0) | 24 cap | 3 | 5258.1 (5247.7-5260.7) | not measured | 6.13 (6.13-6.14) | +| agent session, follow-up turns | 14579 (14547-14587) at the last turn | see below | up to 64 | 3 | see below | 253.2 (196.3-268.1) | 0.72 (0.50-0.79) | + +Ranges are the lowest and highest of the runs; `Runs` counts requests that returned a token. + +### Agent session (conversation cache) + +Cache hit over the follow-up turns: 68.2 % (68.2-68.2); earlier-turn tokens read again per session: 285 (253-293) (earlier turns that the cache did not hold; the 13720 (13720-13720) new tokens of the files are read once); first turn (nothing cached) first token: 0.47 s (0.46-0.47). + +| Turn (run 1) | Prompt tokens | Reused | Read | Prompt tok/s | First token (s) | Decode tok/s | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | 609 | 0 | 609 | 1393.9 | 0.457 | 284.0 | +| 1 | 3802 | 602 | 3200 | 5336.9 | 0.615 | 196.3 | +| 2 | 5317 | 3795 | 1522 | 3116.9 | 0.51 | 252.3 | +| 3 | 8754 | 5310 | 3444 | 4547.1 | 0.774 | 256.7 | +| 4 | 11656 | 8747 | 2909 | 4075.4 | 0.74 | 227.3 | +| 5 | 14579 | 11649 | 2930 | 3901.5 | 0.779 | 236.2 | + +Per-request data: `results.json`. Engine log lines from the start: `engine-log-excerpt.txt` (when this run started the server). + +## Correctness and limitations + +Not measured: this is a speed suite. The conversation uses the model's own replies, so token counts differ a little between runs. Machine details come from the server, `nvidia-smi`/`rocm-smi` and the OS; background load and power limits are not recorded. Total time of the suite: 51 s. diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/engine-log-excerpt.txt b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/engine-log-excerpt.txt new file mode 100644 index 00000000000..725b1dd3b21 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/engine-log-excerpt.txt @@ -0,0 +1,5 @@ +[strata] loading the experts into RAM (about 19 GB) and locking part of them for the GPU. + Please wait and don't close this window; the browser opens when it is ready. +[strata] filling the GPU's expert cache (5115 experts, 7.96 GiB of VRAM) ... +[strata] almost ready ... +ready: http://127.0.0.1:8096/v1 (OpenAI: /v1/chat/completions, Anthropic: /v1/messages, context 262144 tokens) diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/engine-start.txt b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/engine-start.txt new file mode 100644 index 00000000000..08748aeaa20 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/engine-start.txt @@ -0,0 +1,5 @@ +strata mtp: draft layer loaded from the model file, 1128 MiB of VRAM (experts 816, dense 36), files read in 0.69 s (1230 MiB/s) +strata generate: expert cache auto: +30 MiB kept for the dense attention at 262144 tokens of context (split-K partials, STRATA_ATTN_SPLITS) +strata generate: expert cache 5115 slots, 7.96 GiB of VRAM; policy is +strata verify: window up to 6 tokens, 78.9 MiB of device buffers +strata serve: 273 MiB of VRAM free with everything loaded diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/results.json b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/results.json new file mode 100644 index 00000000000..bfb2d8a4771 --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/C/results.json @@ -0,0 +1,1006 @@ +{ + "schema": 1, + "tool": "tools/strata_bench.py", + "date": "2026-10-13", + "params": { + "runs": 3, + "think": false, + "clients": 0, + "long": true, + "started_server": true, + "command": "python tools/strata_bench.py --long" + }, + "elapsed_s": 50.83751678466797, + "machine": { + "gpu": [ + "NVIDIA GeForce RTX 5080" + ], + "gpu_vram_mib": 16303, + "gpu_driver": "617.42", + "cuda": null, + "cpu": "Intel(R) Core(TM) i9-14900K", + "cpu_cores": 24, + "cpu_threads": 32, + "ram_gib": 47.7, + "os": "Windows 11 (10.0.28000)", + "python": "3.14.3", + "strata_version": "0.1.43.1", + "strata_commit": "f1e33fb9", + "engine_version": "0.1.43.1", + "model": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "quant": "IQ4_XS", + "context": 262144, + "parallel_slots": 1, + "config": { + "exe": "strata.exe", + "args": [ + "--pack", + "ornith-abl-iq4_xs", + "--native", + "Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf", + "--expert-profile", + "expert-profile-qwen36.bin", + "--expert-cache", + "auto", + "--prefill", + "auto", + "--spec", + "4", + "--mtp", + "mtp-Ornith-1.5-35B-A3B-Abliterated-MTPv2-Q8_0.gguf", + "--mtp-draft-vocab", + "draft_vocab_cyrillic.bin", + "--max-context", + "262144", + "--kv", + "int8", + "--expert-profile-save", + "expert-profile-ornith-learned.bin", + "--pcie-frac", + "0.00", + "--spec-min-p", + "0.50", + "--pool-workers", + "7", + "--adapt-every", + "1", + "--adapt-swaps", + "160", + "--adapt-decay", + "0.97" + ], + "model_name": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "lib_dirs": [ + "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v13.4\\bin", + "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v13.4\\bin\\x64" + ], + "port": 8096, + "gpu": 0, + "gpus_asked": true, + "open_browser": false, + "idle_unload_s": 3600, + "env": { + "STRATA_ATTN_SPLITS": "256" + } + }, + "launch": "strata.exe --pack ornith-abl-iq4_xs --native Huihui-Ornith-1.5-35B-A3B-abliterated.i1-IQ4_XS.gguf --expert-profile expert-profile-qwen36.bin --expert-cache auto --prefill auto --spec 4 --mtp mtp-Ornith-1.5-35B-A3B-Abliterated-MTPv2-Q8_0.gguf --mtp-draft-vocab draft_vocab_cyrillic.bin --max-context 262144 --kv int8 --expert-profile-save expert-profile-ornith-learned.bin --pcie-frac 0.00 --spec-min-p 0.50 --pool-workers 7 --adapt-every 1 --adapt-swaps 160 --adapt-decay 0.97", + "engine_info": { + "model": "ornith-1.5-35b-a3b-abliterated-iq4_xs", + "max_context": 262144, + "images": false, + "version": "0.1.43.1", + "logit_bias": 1, + "context": 262144, + "kv": "int8", + "kv_resident": 0, + "expert_slots": 5115, + "expert_cache_mib": 8152, + "expert_slots_primary": 5115, + "expert_cache_primary_mib": 8152, + "spec": 6, + "mtp_max": 4, + "lookup": 3, + "vram_free_mib": 273, + "cvec": 0, + "arena_mib": 16320, + "pool_workers": 7, + "pcie_frac": "0.00", + "spec_min_p": "0.50", + "conversation_cache_mib": 0, + "conversation_cache_slots": 4, + "conversation_cache_min_free_mib": 2560, + "conversation_cache_min_tokens": 0, + "tail_role_token": -1, + "vram_elastic": 0, + "engine": "0.1.43.1" + }, + "vram_after": { + "elastic": false + } + }, + "results": { + "chat": { + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 0.156, + "total_s": 0.548, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 344.1, + "prompt_tps_source": "engine", + "decode_tps": 252.3, + "decode_tps_source": "engine", + "prompt_ms": 148.2, + "decode_ms": 396.4, + "drafts_offered": 84, + "drafts_accepted": 65, + "run": 0, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.136, + "total_s": 0.613, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 340.2, + "prompt_tps_source": "engine", + "decode_tps": 207.4, + "decode_tps_source": "engine", + "prompt_ms": 126.4, + "decode_ms": 482.1, + "drafts_offered": 74, + "drafts_accepted": 48, + "run": 0, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.144, + "total_s": 0.552, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 365.7, + "prompt_tps_source": "engine", + "decode_tps": 242.4, + "decode_tps_source": "engine", + "prompt_ms": 125.8, + "decode_ms": 412.6, + "drafts_offered": 91, + "drafts_accepted": 55, + "run": 0, + "prompt": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.144, + "total_s": 0.551, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 376.9, + "prompt_tps_source": "engine", + "decode_tps": 242.8, + "decode_tps_source": "engine", + "prompt_ms": 135.3, + "decode_ms": 411.8, + "drafts_offered": 84, + "drafts_accepted": 70, + "run": 1, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.112, + "total_s": 0.51, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 413.9, + "prompt_tps_source": "engine", + "decode_tps": 248.7, + "decode_tps_source": "engine", + "prompt_ms": 103.9, + "decode_ms": 402.1, + "drafts_offered": 63, + "drafts_accepted": 49, + "run": 1, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.116, + "total_s": 0.508, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 425.9, + "prompt_tps_source": "engine", + "decode_tps": 252.7, + "decode_tps_source": "engine", + "prompt_ms": 108.0, + "decode_ms": 395.7, + "drafts_offered": 79, + "drafts_accepted": 53, + "run": 1, + "prompt": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.146, + "total_s": 0.559, + "prompt_tokens": 51, + "reused_tokens": 0, + "read_tokens": 51, + "generated_tokens": 100, + "prompt_tps": 399.4, + "prompt_tps_source": "engine", + "decode_tps": 239.3, + "decode_tps_source": "engine", + "prompt_ms": 127.7, + "decode_ms": 417.8, + "drafts_offered": 88, + "drafts_accepted": 63, + "run": 2, + "prompt": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.111, + "total_s": 0.521, + "prompt_tokens": 43, + "reused_tokens": 0, + "read_tokens": 43, + "generated_tokens": 100, + "prompt_tps": 422.8, + "prompt_tps_source": "engine", + "decode_tps": 240.8, + "decode_tps_source": "engine", + "prompt_ms": 101.7, + "decode_ms": 415.3, + "drafts_offered": 74, + "drafts_accepted": 47, + "run": 2, + "prompt": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.114, + "total_s": 0.5, + "prompt_tokens": 46, + "reused_tokens": 0, + "read_tokens": 46, + "generated_tokens": 100, + "prompt_tps": 434.4, + "prompt_tps_source": "engine", + "decode_tps": 255.9, + "decode_tps_source": "engine", + "prompt_ms": 105.9, + "decode_ms": 390.8, + "drafts_offered": 85, + "drafts_accepted": 53, + "run": 2, + "prompt": 2 + } + ], + "decode_tps": { + "median": 242.8, + "min": 207.4, + "max": 255.9, + "n": 9 + }, + "ttft_s": { + "median": 0.136, + "min": 0.111, + "max": 0.156, + "n": 9 + }, + "generated_tokens": { + "median": 100, + "min": 100, + "max": 100, + "n": 9 + }, + "prompt_tokens": { + "median": 46, + "min": 43, + "max": 51, + "n": 9 + } + }, + "prompts": { + "4K": { + "target_tokens": 4096, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 0.679, + "total_s": 0.768, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 6096.8, + "prompt_tps_source": "engine", + "decode_tps": 254.2, + "decode_tps_source": "engine", + "prompt_ms": 659.2, + "decode_ms": 94.4, + "drafts_offered": 16, + "drafts_accepted": 12, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.649, + "total_s": 0.733, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 6289.5, + "prompt_tps_source": "engine", + "decode_tps": 271.8, + "decode_tps_source": "engine", + "prompt_ms": 639.0, + "decode_ms": 88.3, + "drafts_offered": 14, + "drafts_accepted": 11, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.657, + "total_s": 0.74, + "prompt_tokens": 4019, + "reused_tokens": 0, + "read_tokens": 4019, + "generated_tokens": 24, + "prompt_tps": 6306.3, + "prompt_tps_source": "engine", + "decode_tps": 274.0, + "decode_tps_source": "engine", + "prompt_ms": 637.3, + "decode_ms": 87.6, + "drafts_offered": 15, + "drafts_accepted": 12, + "run": 2 + } + ], + "prompt_tps": { + "median": 6289.5, + "min": 6096.8, + "max": 6306.3, + "n": 3 + }, + "ttft_s": { + "median": 0.657, + "min": 0.649, + "max": 0.679, + "n": 3 + }, + "prompt_tokens": { + "median": 4019, + "min": 4019, + "max": 4019, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + }, + "16K": { + "target_tokens": 16384, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 2.546, + "total_s": 2.663, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6375.0, + "prompt_tps_source": "engine", + "decode_tps": 196.1, + "decode_tps_source": "engine", + "prompt_ms": 2518.9, + "decode_ms": 122.4, + "drafts_offered": 22, + "drafts_accepted": 13, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 2.466, + "total_s": 2.588, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6580.6, + "prompt_tps_source": "engine", + "decode_tps": 189.0, + "decode_tps_source": "engine", + "prompt_ms": 2440.2, + "decode_ms": 127.0, + "drafts_offered": 11, + "drafts_accepted": 8, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 2.509, + "total_s": 2.606, + "prompt_tokens": 16058, + "reused_tokens": 0, + "read_tokens": 16058, + "generated_tokens": 24, + "prompt_tps": 6474.0, + "prompt_tps_source": "engine", + "decode_tps": 229.9, + "decode_tps_source": "engine", + "prompt_ms": 2480.4, + "decode_ms": 104.4, + "drafts_offered": 20, + "drafts_accepted": 13, + "run": 2 + } + ], + "prompt_tps": { + "median": 6474.0, + "min": 6375.0, + "max": 6580.6, + "n": 3 + }, + "ttft_s": { + "median": 2.509, + "min": 2.466, + "max": 2.546, + "n": 3 + }, + "prompt_tokens": { + "median": 16058, + "min": 16058, + "max": 16058, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + }, + "32K": { + "target_tokens": 32768, + "requests": [ + { + "ok": true, + "err": null, + "ttft_s": 6.13, + "total_s": 6.265, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 5258.1, + "prompt_tps_source": "engine", + "decode_tps": 170.0, + "decode_tps_source": "engine", + "prompt_ms": 6095.2, + "decode_ms": 141.2, + "drafts_offered": 21, + "drafts_accepted": 15, + "run": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 6.127, + "total_s": 6.262, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 5260.7, + "prompt_tps_source": "engine", + "decode_tps": 170.2, + "decode_tps_source": "engine", + "prompt_ms": 6092.2, + "decode_ms": 141.0, + "drafts_offered": 20, + "drafts_accepted": 12, + "run": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 6.142, + "total_s": 6.282, + "prompt_tokens": 32049, + "reused_tokens": 0, + "read_tokens": 32049, + "generated_tokens": 24, + "prompt_tps": 5247.7, + "prompt_tps_source": "engine", + "decode_tps": 163.4, + "decode_tps_source": "engine", + "prompt_ms": 6107.3, + "decode_ms": 146.9, + "drafts_offered": 22, + "drafts_accepted": 12, + "run": 2 + } + ], + "prompt_tps": { + "median": 5258.1, + "min": 5247.7, + "max": 5260.7, + "n": 3 + }, + "ttft_s": { + "median": 6.13, + "min": 6.127, + "max": 6.142, + "n": 3 + }, + "prompt_tokens": { + "median": 32049, + "min": 32049, + "max": 32049, + "n": 3 + }, + "reused_tokens": { + "median": 0, + "min": 0, + "max": 0, + "n": 3 + } + } + }, + "agent": { + "sessions": [ + [ + { + "ok": true, + "err": null, + "ttft_s": 0.457, + "total_s": 0.572, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 34, + "prompt_tps": 1393.9, + "prompt_tps_source": "engine", + "decode_tps": 284.0, + "decode_tps_source": "engine", + "prompt_ms": 436.9, + "decode_ms": 119.7, + "drafts_offered": 29, + "drafts_accepted": 23, + "run": 0, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.615, + "total_s": 0.87, + "prompt_tokens": 3802, + "reused_tokens": 602, + "read_tokens": 3200, + "generated_tokens": 51, + "prompt_tps": 5336.9, + "prompt_tps_source": "engine", + "decode_tps": 196.3, + "decode_tps_source": "engine", + "prompt_ms": 599.6, + "decode_ms": 259.8, + "drafts_offered": 48, + "drafts_accepted": 22, + "run": 0, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.51, + "total_s": 0.751, + "prompt_tokens": 5317, + "reused_tokens": 3795, + "read_tokens": 1522, + "generated_tokens": 62, + "prompt_tps": 3116.9, + "prompt_tps_source": "engine", + "decode_tps": 252.3, + "decode_tps_source": "engine", + "prompt_ms": 488.3, + "decode_ms": 245.7, + "drafts_offered": 50, + "drafts_accepted": 34, + "run": 0, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.774, + "total_s": 0.951, + "prompt_tokens": 8754, + "reused_tokens": 5310, + "read_tokens": 3444, + "generated_tokens": 47, + "prompt_tps": 4547.1, + "prompt_tps_source": "engine", + "decode_tps": 256.7, + "decode_tps_source": "engine", + "prompt_ms": 757.4, + "decode_ms": 183.1, + "drafts_offered": 36, + "drafts_accepted": 25, + "run": 0, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.74, + "total_s": 0.981, + "prompt_tokens": 11656, + "reused_tokens": 8747, + "read_tokens": 2909, + "generated_tokens": 56, + "prompt_tps": 4075.4, + "prompt_tps_source": "engine", + "decode_tps": 227.3, + "decode_tps_source": "engine", + "prompt_ms": 713.8, + "decode_ms": 246.4, + "drafts_offered": 56, + "drafts_accepted": 31, + "run": 0, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.779, + "total_s": 1.044, + "prompt_tokens": 14579, + "reused_tokens": 11649, + "read_tokens": 2930, + "generated_tokens": 64, + "prompt_tps": 3901.5, + "prompt_tps_source": "engine", + "decode_tps": 236.2, + "decode_tps_source": "engine", + "prompt_ms": 751.0, + "decode_ms": 271.0, + "drafts_offered": 44, + "drafts_accepted": 34, + "run": 0, + "turn": 5 + } + ], + [ + { + "ok": true, + "err": null, + "ttft_s": 0.468, + "total_s": 0.55, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 25, + "prompt_tps": 1355.1, + "prompt_tps_source": "engine", + "decode_tps": 287.4, + "decode_tps_source": "engine", + "prompt_ms": 449.4, + "decode_ms": 87.0, + "drafts_offered": 22, + "drafts_accepted": 16, + "run": 1, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.592, + "total_s": 0.827, + "prompt_tokens": 3793, + "reused_tokens": 602, + "read_tokens": 3191, + "generated_tokens": 62, + "prompt_tps": 5487.5, + "prompt_tps_source": "engine", + "decode_tps": 259.1, + "decode_tps_source": "engine", + "prompt_ms": 581.5, + "decode_ms": 239.3, + "drafts_offered": 54, + "drafts_accepted": 30, + "run": 1, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.506, + "total_s": 0.735, + "prompt_tokens": 5319, + "reused_tokens": 3786, + "read_tokens": 1533, + "generated_tokens": 62, + "prompt_tps": 3156.3, + "prompt_tps_source": "engine", + "decode_tps": 265.8, + "decode_tps_source": "engine", + "prompt_ms": 485.7, + "decode_ms": 233.3, + "drafts_offered": 52, + "drafts_accepted": 33, + "run": 1, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.727, + "total_s": 0.931, + "prompt_tokens": 8756, + "reused_tokens": 5312, + "read_tokens": 3444, + "generated_tokens": 53, + "prompt_tps": 4894.1, + "prompt_tps_source": "engine", + "decode_tps": 253.2, + "decode_tps_source": "engine", + "prompt_ms": 703.7, + "decode_ms": 209.3, + "drafts_offered": 37, + "drafts_accepted": 27, + "run": 1, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.716, + "total_s": 0.942, + "prompt_tokens": 11664, + "reused_tokens": 8749, + "read_tokens": 2915, + "generated_tokens": 56, + "prompt_tps": 4215.5, + "prompt_tps_source": "engine", + "decode_tps": 242.6, + "decode_tps_source": "engine", + "prompt_ms": 691.5, + "decode_ms": 230.8, + "drafts_offered": 52, + "drafts_accepted": 29, + "run": 1, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.754, + "total_s": 0.988, + "prompt_tokens": 14587, + "reused_tokens": 11657, + "read_tokens": 2930, + "generated_tokens": 64, + "prompt_tps": 4025.3, + "prompt_tps_source": "engine", + "decode_tps": 267.1, + "decode_tps_source": "engine", + "prompt_ms": 727.9, + "decode_ms": 239.6, + "drafts_offered": 53, + "drafts_accepted": 41, + "run": 1, + "turn": 5 + } + ], + [ + { + "ok": true, + "err": null, + "ttft_s": 0.469, + "total_s": 0.504, + "prompt_tokens": 609, + "reused_tokens": 0, + "read_tokens": 609, + "generated_tokens": 11, + "prompt_tps": 1351.5, + "prompt_tps_source": "engine", + "decode_tps": 273.6, + "decode_tps_source": "engine", + "prompt_ms": 450.6, + "decode_ms": 40.2, + "drafts_offered": 10, + "drafts_accepted": 7, + "run": 2, + "turn": 0 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.593, + "total_s": 0.816, + "prompt_tokens": 3779, + "reused_tokens": 602, + "read_tokens": 3177, + "generated_tokens": 58, + "prompt_tps": 5454.1, + "prompt_tps_source": "engine", + "decode_tps": 254.7, + "decode_tps_source": "engine", + "prompt_ms": 582.5, + "decode_ms": 227.7, + "drafts_offered": 54, + "drafts_accepted": 28, + "run": 2, + "turn": 1 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.504, + "total_s": 0.708, + "prompt_tokens": 5301, + "reused_tokens": 3772, + "read_tokens": 1529, + "generated_tokens": 54, + "prompt_tps": 3164.3, + "prompt_tps_source": "engine", + "decode_tps": 258.2, + "decode_tps_source": "engine", + "prompt_ms": 483.2, + "decode_ms": 209.1, + "drafts_offered": 37, + "drafts_accepted": 26, + "run": 2, + "turn": 2 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.717, + "total_s": 0.865, + "prompt_tokens": 8730, + "reused_tokens": 5294, + "read_tokens": 3436, + "generated_tokens": 41, + "prompt_tps": 4947.4, + "prompt_tps_source": "engine", + "decode_tps": 268.1, + "decode_tps_source": "engine", + "prompt_ms": 694.5, + "decode_ms": 152.9, + "drafts_offered": 29, + "drafts_accepted": 20, + "run": 2, + "turn": 3 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.702, + "total_s": 0.944, + "prompt_tokens": 11626, + "reused_tokens": 8723, + "read_tokens": 2903, + "generated_tokens": 54, + "prompt_tps": 4281.1, + "prompt_tps_source": "engine", + "decode_tps": 218.4, + "decode_tps_source": "engine", + "prompt_ms": 678.1, + "decode_ms": 247.3, + "drafts_offered": 55, + "drafts_accepted": 28, + "run": 2, + "turn": 4 + }, + { + "ok": true, + "err": null, + "ttft_s": 0.794, + "total_s": 1.042, + "prompt_tokens": 14547, + "reused_tokens": 11619, + "read_tokens": 2928, + "generated_tokens": 62, + "prompt_tps": 3814.5, + "prompt_tps_source": "engine", + "decode_tps": 244.3, + "decode_tps_source": "engine", + "prompt_ms": 767.6, + "decode_ms": 253.8, + "drafts_offered": 40, + "drafts_accepted": 32, + "run": 2, + "turn": 5 + } + ] + ], + "hit_pct": { + "median": 68.238, + "min": 68.231, + "max": 68.248, + "n": 3 + }, + "reread_tokens": { + "median": 285, + "min": 253, + "max": 293, + "n": 3 + }, + "new_tokens": { + "median": 13720, + "min": 13720, + "max": 13720, + "n": 3 + }, + "turn_ttft_s": { + "median": 0.716, + "min": 0.504, + "max": 0.794, + "n": 15 + }, + "last_turn_prompt_tokens": { + "median": 14579, + "min": 14547, + "max": 14587, + "n": 3 + }, + "turn_decode_tps": { + "median": 253.2, + "min": 196.3, + "max": 268.1, + "n": 15 + }, + "first_turn_ttft_s": { + "median": 0.468, + "min": 0.457, + "max": 0.469, + "n": 3 + } + }, + "clients": {}, + "skipped": {}, + "failures": [] + } +} \ No newline at end of file diff --git a/bench/results/2026-10-13-rtx5080-ornith-attn-splits/README.md b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/README.md new file mode 100644 index 00000000000..5829b36822f --- /dev/null +++ b/bench/results/2026-10-13-rtx5080-ornith-attn-splits/README.md @@ -0,0 +1,33 @@ +# STRATA_ATTN_SPLITS and a transplanted MTP layer on an RTX 5080 (Ornith 1.5 35B A3B, IQ4_XS) + +Measured on 2026-10-13, one machine: RTX 5080 16 GB (driver 617.42), i9-14900K, 47.7 GiB RAM, Windows 11. Model +Huihui-Ornith-1.5-35B-A3B-abliterated i1-IQ4_XS, `--max-context 262144 --kv int8 --spec 4`, the install's configuration +otherwise unchanged. Engine built from `ornith` at f1e33fb9 (the same binary in all three arms). Each arm is one server +start and one run of `tools/strata_bench.py --runs 3 --long`, in the order A, B, C. + +- **A**: the default (no `STRATA_ATTN_SPLITS`), draft layer from the model file (IQ4_XS). +- **B**: `STRATA_ATTN_SPLITS=256`, the same draft layer. +- **C**: B plus `--mtp mtp-Ornith-1.5-35B-A3B-Abliterated-MTPv2-Q8_0.gguf` (gbuzhf, Q8_0, 1.99 GB) in place of the + model file's draft layer. + +| | A | B | C | +| --- | ---: | ---: | ---: | +| reserved for the dense attention at 262144 tokens | 393 MiB | 30 MiB | 30 MiB | +| verify window buffers | 441.8 MiB | 78.9 MiB | not logged | +| draft layer | 704 MiB | 704 MiB | 1128 MiB | +| expert cache slots | 5362 | 5378 | 5115 | +| VRAM free with everything loaded | 0 MiB (LOW warning) | 284 MiB | 273 MiB | +| short chat decode, tok/s | 244.2 (206.7-265.7) | 245.7 (192.7-270.1) | 242.8 (207.4-255.9) | +| 4K prompt, tok/s | 5522.9 (5497.2-5782.7) | 6300.4 (6021.9-6321.2) | 6289.5 (6096.8-6306.3) | +| 16K prompt, tok/s | 5983.1 (5975.1-6078.2) | 6363.9 (6303.7-6464.1) | 6474.0 (6375.0-6580.6) | +| 32K prompt, tok/s | 4848.8 (4841.5-4887.5) | 5219.6 (5199.6-5242.8) | 5258.1 (5247.7-5260.7) | +| agent session decode, tok/s | 226.2 (206.9-282.5) | 245.9 (181.0-282.3) | 253.2 (196.3-268.1) | +| drafts accepted, whole run | 1072 of 1604 (66.8%) | 1045 of 1591 (65.7%) | 1115 of 1671 (66.7%) | + +Medians with the lowest and highest run. The split frees 363 MiB of reserve and removes the 0 MiB warning; the prompt +speeds were 6-14% higher with it, which one start per arm cannot attribute (the card in A ran with 0 MiB free, where +the driver can page). The decode ranges overlap in every row. The transplanted MTPv2 layer accepted the same share of +drafts as the model's own and costs 424 MiB more, 263 expert slots: not kept. + +Not tested: outputs (speed suite only), contexts above 32K tokens, HIP/SYCL builds, more than one start per arm. +Per-arm reports, `results.json` and the engine's start lines are in `A/`, `B/` and `C/`. diff --git a/build-strata-tests.bat b/build-strata-tests.bat new file mode 100644 index 00000000000..fe704be10c1 --- /dev/null +++ b/build-strata-tests.bat @@ -0,0 +1,5 @@ +@echo off +rem The engine plus the parity tests used to check engine changes (same toolchain and options as build-strata.bat). +call "C:\Program Files (x86)\Microsoft Visual Studio\18\BuildTools\VC\Auxiliary\Build\vcvars64.bat" >nul +"C:\Program Files\CMake\bin\cmake.EXE" -G Ninja -DCMAKE_MAKE_PROGRAM=D:\Strata-dev\Strata\.venv\Scripts\ninja.exe -S D:\Strata-dev\Strata -B D:\Strata-dev\Strata\build -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_CUDA=ON -DSTRATA_BUILD_TESTS=ON -DCMAKE_CUDA_ARCHITECTURES=120 "-DCMAKE_CUDA_COMPILER=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\bin\nvcc.exe" "-DCUDAToolkit_ROOT=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" -DSTRATA_GGML_DIR=D:\Strata-dev\Strata\third_party\llama.cpp || exit /b 1 +"C:\Program Files\CMake\bin\cmake.EXE" --build D:\Strata-dev\Strata\build --target strata %* -j 16 || exit /b 1 diff --git a/data/expert-profile-qwen36.bin b/data/expert-profile-qwen36.bin new file mode 100644 index 00000000000..6f0b91369ef Binary files /dev/null and b/data/expert-profile-qwen36.bin differ diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 66cef80cb84..b3b1d49787b 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -748,6 +748,14 @@ print(r.choices[0].message.content) `high`) in `strata-.json`, outside `"args"` (#1641): the engine has no such option, and `--reasoning-effort` in `"args"` is moved to this key with a note at startup. The web app's shared Chat settings ("Use for other apps too") and a request's own level win over it. +- **Keep the thinking of earlier turns (opt-in).** A chat template decides whether the thinking of turns before the + latest user message goes back into the prompt; Qwen3.6's, Ornith's and KAT-Coder's leave it out unless + `preserve_thinking` is set (Flash-Next's keep it). `"preserve_thinking": true` in `strata-.json` sets it for + every request, on both APIs (llama.cpp's `--reasoning-preserve`); `false` leaves it out on templates that keep it. A + request can set its own with `"chat_template_kwargs": {"preserve_thinking": true}` (OpenAI, and on the Anthropic API + as a Strata extension), and `"preserve_thinking"` in the shared Chat settings (`POST /settings`) wins over the + config. The client has to send the earlier thinking back (`reasoning_content`, or Anthropic `thinking` blocks), as + the web app does. Kept thinking makes every prompt longer. Off by default: the template's own choice. - **Keep Windows awake while it works (opt-in, #1727).** `"prevent_sleep": true` in `strata-.json` keeps the PC from going to sleep while at least one request is running or waiting (Windows' `SetThreadExecutionState` with `ES_SYSTEM_REQUIRED`; the screen may still turn off), and gives the normal sleep rules back when the server is idle @@ -1622,7 +1630,8 @@ Since 0.1.43 the same default holds on HIP for gfx1200 / gfx1201 (RDNA4), under Where it applies: serve on a CUDA build with one GPU or a layer split (every stage skips by itself). Not applied, and no line printed: when every expert is in VRAM (nothing is missed), with `--batch` / `parallel` slots (a window then mixes requests, and one request's skip would depend on another's -tokens: kept off until measured), with a peer-expert tier or a helper GPU, on HIP and SYCL builds, and (0.1.42.1, #1884) when 90% or more +tokens: kept off until measured), with a peer-expert tier or a helper GPU, on HIP and SYCL builds, for the qwen35moe models (Qwen3.6-35B-A3B and its fine-tunes: measured +on Flash-Next's 10-of-512 routing only; `STRATA_ROUTE_TAIL_SKIP=7` turns it on there), and (0.1.42.1, #1884) when 90% or more of all the experts are in the GPU cache: the few misses are cheap there and the skip costs more draft acceptance than it saves (a Tesla V100 32 GB with Coder IQ1_M, 94.8% of the experts cached, 99.9% decode hits: draft acceptance 70.1% to 68.2%, decode 1-2% slower than 0.1.41). A start-up line says so; `STRATA_ROUTE_TAIL_SKIP=7` turns it on anyway, and `STRATA_TAIL_SKIP_MAX_CACHED=` moves the line (101 = never). On an AMD layer split the line is 60% (2x R9700 at 16 GB-equivalent: Q2_0 with 73% cached lost 6.1% story decode with the skip, IQ3_S with 45% cached gained 7-10%). The request log gets diff --git a/docs/MODELS.md b/docs/MODELS.md index ee750362b9d..1c5c19c9c7f 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -5,7 +5,8 @@ Strata runs one model, [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8- The installer recommends one for your PC; this page explains the choice. Back to the [README](../README.md). > **On this page:** [Pick by RAM](#pick-by-ram) · [Speed](#how-fast-is-each-size) · [The sizes](#the-sizes) · -> [Will it fit?](#will-it-fit) · [The versions](#the-versions) · [Adding another model](#adding-or-switching-models) +> [Will it fit?](#will-it-fit) · [The versions](#the-versions) · [Qwen3.6-35B-A3B](#qwen36-35b-a3b) · +> [Adding another model](#adding-or-switching-models) ## Pick by RAM @@ -150,6 +151,26 @@ machines run it without a RAM budget. Manual workflow: [UD-Q6_K_XL](UNSLOTH_Q6.m For **OrcaRouter's Flash-Next Uncensored IQ3_XXS**, see the [manual compatibility setup](ORCA.md). It needs an explicit packing conversion and is not an installer menu option. +### Qwen3.6-35B-A3B + +**[Qwen3.6-35B-A3B](QWEN36.md)** - not a version of Flash-Next but a second, smaller model of the same family +(Unsloth's UD-IQ4_XS and UD-IQ3_S, with the MTP draft layer in the file). Its experts take 12-14 GB of RAM, so it is +the choice for 16-32 GB PCs and 8-12 GB graphics cards: 105-127 tokens/s on a 12 GB RTX 4070 Ti and 75-89 tokens/s +with an 8 GB card's VRAM (emulated), and it reads long prompts at 2,700-4,500 tokens/s +([measurements](QWEN36.md#speed-measured)). NVIDIA, one GPU, no images yet. + +``` +START-HERE.bat --setup --family qwen36 +``` + +**[Ornith-1.5-35B-A3B](QWEN36.md#ornith-15-35b-a3b)** - ornith-ai's fine-tune of Qwen3.6-35B-A3B for coding agents +(bartowski's IQ4_XS and IQ3_XXS). The same engine path and the same RAM and VRAM needs; its draft layer is accepted +less often, so it writes 102-110 tokens/s on the 4070 Ti ([measurements](QWEN36.md#ornith-15-35b-a3b)). + +``` +START-HERE.bat --setup --family ornith +``` + ## Adding or switching models You can add another model any time with `SETUP.bat` (the same as `START-HERE.bat --setup`; on Linux diff --git a/docs/QWEN36.md b/docs/QWEN36.md new file mode 100644 index 00000000000..88010c489cc --- /dev/null +++ b/docs/QWEN36.md @@ -0,0 +1,193 @@ +# Qwen3.6-35B-A3B + +A second, smaller model next to Qwen3.8-Flash-Next: [Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) +(35B parameters, 3B active per token), in Unsloth's GGUFs with the MTP draft layer +([unsloth/Qwen3.6-35B-A3B-MTP-GGUF](https://huggingface.co/unsloth/Qwen3.6-35B-A3B-MTP-GGUF), Apache 2.0). It is the +same family as Flash-Next - DeltaNet layers with a full-attention layer every fourth, a mixture of experts with a +shared expert - but 40 layers of 256 experts instead of 48 of 512, so its experts take 12-14 GB of RAM instead of +23-50 GB. It is the model for PCs with 16-32 GB of RAM and 8-12 GB graphics cards. Back to the [models](MODELS.md). + +> **On this page:** [Setup](#setup) · [What fits](#what-fits) · [Speed](#speed-measured) · +> [Small cards](#small-cards) · [Against llama.cpp](#against-llamacpp) · [Ornith-1.5](#ornith-15-35b-a3b) · +> [What is different](#what-is-different-from-flash-next) · [Limits](#limits) + +## Setup + +``` +./setup.sh --setup --family qwen36 (Linux) +START-HERE.bat --setup --family qwen36 (Windows) +``` + +Setup offers two sizes and picks the larger one that fits: + +| Size | Download | Experts in RAM | RAM it asks for | +| --- | ---: | ---: | ---: | +| **UD-IQ4_XS** (recommended, ~4-bit) | 18.2 GB | 14.2 GiB (measured) | 26 GB | +| **UD-IQ3_S** (smaller, faster) | 15.3 GB | 12.0 GiB (measured) | 24 GB; a 16 GB PC with a 12 GB card in the [low-RAM mode](MODELS.md#a-big-graphics-card-and-little-ram) (with an 8 GB card it does not fit) | + +The file goes where every model goes (`/models/qwen36-UD-IQ4_XS/`), its pack is made the usual way +(`tools/iq_pack.py --compat-bf16`, two seconds), and the MTP draft layer is read from the same file: there is no +separate draft download. The engine finds out which model it runs from the file (`general.architecture`): +nothing in the command line says "Qwen3.6". + +No ready-made engine runs it yet: the published Windows engines (0.1.40 to 0.1.41) were built before this model was +added. Setup looks for it in the engine and, when the ready-made one lacks it, compiles the engine on the PC instead +(10-20 minutes, once; setup installs the compiler and the CUDA toolkit when they are missing, asking first). On Linux +setup compiles the engine anyway. The compiled engine runs every model, Flash-Next too. + +## What fits + +Every expert lives in RAM; the graphics card holds the dense weights, the context and as many experts as fit in the +rest of its memory (the expert cache). Measured at 32K of context, 8-bit K/V: + +| Graphics card | UD-IQ4_XS experts on the GPU | UD-IQ3_S experts on the GPU | +| --- | ---: | ---: | +| 12 GB (RTX 4070 Ti) | 4,407 of 10,240 (43%) | 5,587 (55%) | +| 8 GB (emulated, see below) | 1,372 (13%) | 2,003 (20%) | + +### Small cards + +How small a card it starts on, measured the same way as the 8 GB rows (the 4070 Ti with the rest of its VRAM held by +another process; "free" is what `nvidia-smi` showed free before the start), UD-IQ3_S with its draft layer: + +| Settings | Starts from | Writes (short chat) | Reads (8K prompt) | +| --- | ---: | ---: | ---: | +| setup's: 32K context, 8-bit K/V | 4.17 GB free | 73 tokens/s | 591 tokens/s | +| 8K context, `--draft-vocab en` | 3.9 GB free | 74 tokens/s | 962 tokens/s | + +UD-IQ4_XS needs more: 4.68 GB free with setup's settings (64 tokens/s) and 4.3 GB with the second ones (66 tokens/s), +so setup recommends UD-IQ3_S on cards under 5.5 GB. At 3.65 GB free with the second settings the short prompt ran and the 8K one stopped at the start (free VRAM differs +by a few tens of MiB from one start to the next). With less, the engine stops at the start and says how many MiB are +short and what makes room. Setup recommends the 8K context on cards under 5.5 GB (4 GB cards). These are emulated +sizes, not runs on 4 and 6 GB cards. + +## Speed (measured) + +RTX 4070 Ti (12 GB), Ryzen 9 5950X, 62 GB of RAM, Linux, CUDA 13.0; 32K context, 8-bit K/V, the MTP draft layer on +(`--spec 4`). "8 GB" is the same card with 4 GiB of VRAM held by another process for the whole run, so the engine +sizes itself as on an 8 GB card with the same desktop load; it is not a measurement on an 8 GB card. Short chat = a +21-token prompt, 128 tokens written; long = an 8,014-token prompt, 64 tokens written. + +| Size, card | Short chat, writes | Long prompt, reads | Long prompt, writes | +| --- | ---: | ---: | ---: | +| UD-IQ4_XS, 12 GB | 105 tokens/s | 4,460 tokens/s | 86 tokens/s | +| UD-IQ4_XS, 8 GB | 75 tokens/s | 2,669 tokens/s | 68 tokens/s | +| UD-IQ3_S, 12 GB | 127 tokens/s | 4,284 tokens/s | 106 tokens/s | +| UD-IQ3_S, 8 GB | 89 tokens/s | 3,076 tokens/s | 70 tokens/s | + +These tables were measured with a first expert profile that ranked nothing (every expert in turn). The shipped one, +`data/expert-profile-qwen36.bin`, ranks the experts by how often 24 varied prompts (code, agent tool calls, chat, +French and Chinese) routed to them; the profile decides which experts the GPU holds when the engine starts. Against the +first one, on prompts not among the 24 (three runs each, UD-IQ4_XS): 12 GB, short chat 99 -> 106 tokens/s and long +prompt writes 84 -> 93; with an 8 GB card's VRAM no difference (73 and 60-62 both). Ornith uses the same profile (on +Ornith IQ4_XS: long prompt writes 74 -> 87, short chat unchanged; a profile from Ornith's own routing was not +clearly better - faster on the short chat, slower on the long prompt). It was made with the +server's `--expert-profile-save` and a 64-slot expert cache, so the ranking is the routing count and not what one +cache happened to hold. + +A 256K context (setup's 256K choice) runs on the 12 GB card too: UD-IQ3_S wrote 78 tokens/s and read the 8K prompt +at 3,802 tokens/s with it (the engine keeps 393 MiB beside the expert cache for the attention over the whole context). + +Long prompts: the prompt path's attention gives each block of the tensor-core kernel 8 neighbouring queries, so each +chunk of K/V is read once for all 8 (sm_80+, fp16 or 8-bit K/V; `STRATA_PROMPT_ATTN_DENSE=0` goes back to one block +per query). A 150K-token agentic transcript (KAT i-quality-v2, 262K context, all K/V in VRAM, nothing else on the GPU; the +prompt chunk `auto` picks is 2,816 tokens with fp16 K/V and 8,192 with 8-bit), read in one request: + +| K/V | One block per query | 8 queries per block | llama.cpp (q8_0 K/V, `-ncmoe 18`) | +|---|---|---|---| +| fp16 | 771 s (99 tokens/s at 146K) | 133 s (714 tokens/s at 146K) | 417 s | +| 8-bit | 170 s (500 tokens/s at 140K) | 113 s (781 tokens/s at 140K) | | + +The attention's share of the prompt fell from 93% to 61% (fp16) and from 83% to 75% (8-bit). Against llama.cpp on +the same transcript cut at 88K, 118K and 150K tokens, the teacher-forced KL past the first token stayed at 0.0036 or +below with 584-585 of 585 top-1 tokens the same (fp16; 8-bit: 0.0028 or below, 583 of 585); before the change +the fp16 figures were 0.0028 and 583-584. `qsa_prompt_attn_parity` checks it against the old kernel and an FP64 +reference (at 32K: within 4.2e-6 of FP64, as the old kernel; 5.7x faster with fp16 K/V, 1.6x with 8-bit). + +The draft layer's drafts are accepted 0.84-0.96 of the time on these prompts (2.4-3.1 tokens per verify window). +Through the server, chat answers with thinking ran at 104-118 tokens/s (UD-IQ4_XS, 12 GB). + +## Against llama.cpp + +The same PC, the same UD-IQ4_XS file, llama.cpp b11438 (`llama-server --fit on`, which puts what fits on the card +and the rest of the experts on the CPU), the same prompts: + +| | Strata 12 GB | llama.cpp 12 GB | Strata 8 GB | llama.cpp 8 GB | +| --- | ---: | ---: | ---: | ---: | +| Short chat, writes | 105 | 51 | 75 | 39 | +| Long prompt, reads | 4,460 | 708 | 2,669 | 490 | +| Long prompt, writes | 86 | 49 | 68 | 36 | + +(tokens/s). The answers are the same: greedy decoding gives llama.cpp's tokens exactly on UD-IQ4_XS (three short +prompts and the 8K prompt, 32 tokens each, with the expert cache at 3,000 slots). The experts in the GPU's cache +round slightly differently from the CPU's, so which experts the cache holds can move a near-tie: with the cache sized +by `auto` (4,420 slots) the 8K prompt took the second choice at its 17th token, where llama.cpp's top two are 28% and +26%. On UD-IQ3_S the first tokens agree and a later token can differ where llama.cpp's own top two are within a few +percent (37% against 35% at the first difference seen) - a rounding-level tie, both continuations correct. + +## Ornith-1.5-35B-A3B + +[Ornith-1.5-35B-A3B](https://huggingface.co/ornith-ai/Ornith-1.5-35B-A3B) (ornith-ai, MIT) is a fine-tune of the same +model for coding agents: the same architecture (`qwen35moe`), tokenizer and sizes, so it runs on the same engine path +with Qwen3.6's expert profile and draft vocabulary. There are no Unsloth files for it; setup uses bartowski's +single-file GGUFs ([bartowski/Ornith-1.5-35B-A3B-GGUF](https://huggingface.co/bartowski/Ornith-1.5-35B-A3B-GGUF)), +which keep the MTP draft layer (stored Q4_0): + +``` +./setup.sh --setup --family ornith (Linux) +START-HERE.bat --setup --family ornith (Windows) +``` + +| Size | Download | Experts in RAM | RAM it asks for | +| --- | ---: | ---: | ---: | +| **IQ4_XS** (recommended, ~4-bit) | 19.3 GB | 15.9 GiB (measured) | 28 GB | +| **IQ3_XXS** (smaller, faster) | 15.3 GB | 12.4 GiB (measured) | 24 GB; a 16 GB PC with a 12 GB card in the low-RAM mode | + +Speed, measured as [above](#speed-measured) (the same PC, prompts and settings): + +| Size, card | Short chat, writes | Long prompt, reads | Long prompt, writes | +| --- | ---: | ---: | ---: | +| IQ4_XS, 12 GB | 102 tokens/s | 4,423 tokens/s | 73 tokens/s | +| IQ4_XS, 8 GB | 75 tokens/s | 3,724 tokens/s | 52 tokens/s | +| IQ3_XXS, 12 GB | 110 tokens/s | 4,549 tokens/s | 103 tokens/s | +| IQ3_XXS, 8 GB | 78 tokens/s | 3,739 tokens/s | 75 tokens/s | + +Its draft layer helps less than Qwen3.6's: with every draft proposed, 61% of them were accepted on a chat prompt and +38% on a code prompt, against 80% and 98% for Qwen3.6 UD-IQ4_XS. A Q8_0 copy of the same layer (mudler's +APEX-MTP-Compact file) did no better, 44% and 29%, and the Q4_0 kernels give llama.cpp's results (`iq_parity`), so the +difference is the draft layer itself, which the fine-tune did not retrain. The answers are the main model's either way: greedy decoding gives llama.cpp b11438's tokens except where +3-bit rounding moves a close choice (IQ3_XXS: one difference seen, at a token llama.cpp gave 0.73 and Strata 0.63). + +## What is different from Flash-Next + +The engine keeps one code path for both and switches on the model's geometry: + +- **One residual stream** instead of Flash-Next's four hyper-connection streams: RMSNorm before the mixer and before + the MoE, the output added back (llama.cpp's `qwen35moe` graph). +- **Dense attention**: no sparse-attention indexer; every cached position is attended (the prompt path's tensor-core + kernel has a dense mode for it, 8 query heads per K/V head). +- **No PLE** block, no n-gram table on the SSD. +- **The DeltaNet output gate** is `silu(z)` instead of `sigmoid(z)`. +- **The MoE** routes 8 of 256 experts (Flash-Next: 10 of 512); 2048-wide, 512-wide experts. + +Flash-Next's own arithmetic is unchanged: the same tokens and the same speed as the engine before this model was +added (Coder, three short prompts and an 8K prompt, back to back on the same PC). + +## Limits + +- **NVIDIA RTX 20 and newer for now.** The prompt path's dense attention is a tensor-core kernel (sm_75+). On AMD and + older NVIDIA cards (Pascal, Volta) the engine, and the server, read every prompt through the decode windows instead + (on the 4070 Ti, made to take that path: a 9,800-token prompt in 113 s, ~87 tokens/s, against 3.9 s batched). It + has not been run on an AMD or a pre-RTX 20 card; setup says so and asks before the download. +- **One GPU.** The layer split across several cards is not done for this model. +- **Images** use the stock encoder (`tools/vision`, built by setup) with the model's mmproj: bartowski's + `mmproj-Qwen_Qwen3.6-35B-A3B-bf16.gguf` for Qwen3.6 and its fine-tunes that ship none (KAT-Coder), bartowski's + `mmproj-Ornith-1.5-35B-A3B-bf16.gguf` for Ornith. The engine needed no change (moondrop12138 found it working on an + RTX 5070 Ti). Measured on the RTX 4070 Ti, 32K context: Ornith IQ4_XS and KAT i-quality-v2 read a test picture + (three shapes and a caption) right over both APIs, ~288 prompt tokens each; a 1920x1200 photo took 1,019 tokens + and 6.2 s with its answer. The encoder on the GPU holds 1,254 MiB, so setup counts it (and the 700 MiB reserve + for it) when it decides on KV streaming; with it Ornith wrote ~5% slower (135 vs 143 tokens/s, the mean of 12 + replies) and read a 24K-token prompt as fast (~4,130 tokens/s either way). Not run: Qwen3.6's own files with + images, the encoder on the CPU beyond one test, AMD. +- The low-RAM mode (`--resident-experts`) gives the same tokens as the normal mode on this PC (UD-IQ3_S: 9.0 GiB of + experts held in RAM); it has not been measured on a 16 GB PC. diff --git a/include/strata/artifact/gguf_reader.hpp b/include/strata/artifact/gguf_reader.hpp index 1be01432615..28c2f85603f 100644 --- a/include/strata/artifact/gguf_reader.hpp +++ b/include/strata/artifact/gguf_reader.hpp @@ -588,7 +588,17 @@ struct Qwen4ExpGuard { inline std::string check_architecture(const GgufFile& g, const Qwen4ExpGuard& want = {}) { const MetaValue* arch = g.get("general.architecture"); if (!arch) return "missing general.architecture"; - if (arch->s != "qwen4exp") return "architecture is '" + arch->s + "', this engine requires 'qwen4exp'"; + if (arch->s == "qwen35moe") { + // Qwen3.6-35B-A3B: its shape is read and checked against the kernels by core::geometry_from_gguf, which + // runs before any loader; here only that the keys the loaders' callers rely on are present. + for (const char* key : {"qwen35moe.block_count", "qwen35moe.embedding_length", "qwen35moe.expert_count", + "qwen35moe.expert_used_count", "qwen35moe.attention.head_count", + "qwen35moe.attention.head_count_kv"}) + if (!g.get(key)) return std::string("missing ") + key; + return {}; + } + if (arch->s != "qwen4exp") + return "architecture is '" + arch->s + "', this engine requires 'qwen4exp' or 'qwen35moe'"; struct Req { const char* key; uint64_t want; diff --git a/include/strata/core/layout.hpp b/include/strata/core/layout.hpp index 7b27fec9527..9aad912e036 100644 --- a/include/strata/core/layout.hpp +++ b/include/strata/core/layout.hpp @@ -20,11 +20,20 @@ #include #include +namespace strata { class GgufFile; } + namespace strata::core { +/// Which model the geometry describes, from the GGUF's `general.architecture`. Flash-Next (`qwen4exp`) is the +/// default and the field initialisers below are its numbers; Qwen3.6-35B-A3B (`qwen35moe`) is the same family +/// (GDN + gated full attention every 4th layer, softmax router, shared expert) without the hyper-connection +/// residual, the PLE block or the QSA indexer, and with SiLU instead of sigmoid on the GDN output gate. +enum class ModelArch : int32_t { kQwen4Exp = 0, kQwen35Moe = 1 }; + /// The model's geometry, taken from `docs/semantics.md` and the artifact's own metadata. Every field here /// is a number a kernel depends on, so a change is a change to a kernel contract and not a tuning knob. struct ModelGeometry { + ModelArch arch = ModelArch::kQwen4Exp; int64_t n_embd = 2560; int64_t n_layers = 48; int64_t qsa_interval = 4; ///< every 4th layer is full attention: layers 3, 7, ... 47 @@ -51,13 +60,32 @@ struct ModelGeometry { // MoE, on every layer int64_t n_expert = 512; int64_t n_ff = 640; + int64_t n_expert_used = 10; + + // what the model has besides the shared GDN/attention/MoE core (all true for qwen4exp) + bool has_ple = true; ///< the per-layer embedding block at layer 1 + bool has_indexer = true; ///< QSA's sparse-attention indexer; false = every cached cell is attended + bool gdn_gate_silu = false; ///< GDN output gate: false = sigmoid(z) (qwen4exp), true = silu(z) (qwen35moe) + int64_t n_mtp_layers = 0; ///< MTP layers stored in the model file after the n_layers main layers + /// The hyper-connection residual (hc streams mixed through hc_lr): qwen4exp has 4. hc = 1 is a plain + /// residual stream `x += f(norm(x))`, so hc_dim() is n_embd and every residual buffer keeps its formula. + bool has_hc() const { return hc > 1; } int64_t hc_dim() const { return hc * n_embd; } /// `layer % qsa_interval == qsa_interval - 1` is full attention. Derived, not a second list. int64_t n_qsa_layers() const { return n_layers / qsa_interval; } int64_t n_gdn_layers() const { return n_layers - n_qsa_layers(); } }; +/// The geometry of the model in `gguf` (its metadata shard), by `general.architecture`. qwen4exp keeps every +/// default except the MoE shape (`expert_count` / `expert_used_count`: a pruned variant such as the GSQ-RCO Coder +/// ships fewer experts); qwen35moe is read from its own keys and checked against what the kernels support. +/// Returns false and fills `err` for any other architecture or an unsupported shape. +bool geometry_from_gguf(const GgufFile& gguf, ModelGeometry& g, std::string& err); + +/// The short name of `arch` for messages ("qwen4exp", "qwen35moe"). +const char* arch_name(ModelArch arch); + /// True for the full-attention layers. `docs/semantics.md` gives this twice over - `full_attention_interval /// = 4` and an explicit `attention.compress_ratios` array - and this is the first of the two. inline bool is_qsa_layer(const ModelGeometry& g, int64_t layer) { diff --git a/include/strata/core/mtp.hpp b/include/strata/core/mtp.hpp index 6acf16ae9e3..9ef61889def 100644 --- a/include/strata/core/mtp.hpp +++ b/include/strata/core/mtp.hpp @@ -239,6 +239,14 @@ class MtpDrafter { std::vector tensors_; uint8_t* dense_ = nullptr; uint8_t* experts_ = nullptr; + /// A model without hyper-connections (qwen35moe): the draft layer is the model file's own nextn layer + /// (blk.), read in place - its tensors under the names the rt files use, the eh_proj split into its + /// embedding and hidden halves - and its experts kept in their GGUF form (types xgu_ / xd_, xblob_ bytes each). + bool load_gguf_layer(const std::string& gguf, const ModelGeometry& g, std::vector& dense, std::string& err); + int xgu_ = -1, xd_ = -1; + uint64_t xblob_ = 0; + uint8_t* nat_xq_ = nullptr; ///< the routed experts' q8_1 input (native experts) + uint8_t* nat_scratch_ = nullptr; // Slot drafters borrow immutable weights and head; each still owns its state and scratch. bool owns_weights_ = true; bool owns_draft_head_ = true; diff --git a/include/strata/core/verify.hpp b/include/strata/core/verify.hpp index ffaea059178..fecfc1d6029 100644 --- a/include/strata/core/verify.hpp +++ b/include/strata/core/verify.hpp @@ -95,6 +95,12 @@ class Verifier { /// the flags guarded, so this verifier refuses every later window. True when the streams finished. bool release_gpu_waits(int timeout_ms); + /// The part of `init`'s device arena that grows with the context: without the indexer (qwen35moe) every cell + /// is attended, so the selection and the decode attention's per-chunk partials are sized for the whole context + /// (~8 MiB per window row at 32K, ~65 MiB at 256K). 0 with the indexer. The expert cache's auto sizing keeps + /// it out of the cache - the arena is allocated after the cache. + static uint64_t dense_attention_bytes(const ModelGeometry& g, const SessionState& ss, int max_t); + /// `max_t` <= kVerifyMaxT. `head` may be null (the canonical head is then run per token). bool init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits, const NativeHead* head, int max_t, std::string& err); @@ -294,7 +300,9 @@ class Verifier { /// Token t's residual after the last layer, (hc, n_embd) on the device, valid until the next `run`. const float* final_R(int t) const; - const float* final_R_all() const { return next_ ? next_->final_R_all() : R_; } + /// The rows the MTP drafter reads: the final residual (hc x n_embd each) - or, for a model without + /// hyper-connections (qwen35moe), the head's input after output_norm (llama.cpp's h_nextn). + const float* final_R_all() const { return next_ ? next_->final_R_all() : (g_ != nullptr && !g_->has_hc() ? head_mixed_ : R_); } /// The GPU plan the pool writes each layer (VRAM hits + the PCIe share of the misses); give it to the /// dispatch (`ExpertDispatch::plan`) before the first `run`. diff --git a/include/strata/core/vmm.hpp b/include/strata/core/vmm.hpp index d8bcf0b57d6..7f344715177 100644 --- a/include/strata/core/vmm.hpp +++ b/include/strata/core/vmm.hpp @@ -4,6 +4,7 @@ // expert cache gives up, and neither moves: every captured graph and stored pointer stays valid. The driver calls // are reached through the runtime's entry-point query, so nothing links the driver library. Not under HIP. +#include #include #include diff --git a/include/strata/kernels/cpu/expert_layout.hpp b/include/strata/kernels/cpu/expert_layout.hpp index 6b3fddb6646..aee25851285 100644 --- a/include/strata/kernels/cpu/expert_layout.hpp +++ b/include/strata/kernels/cpu/expert_layout.hpp @@ -18,6 +18,7 @@ namespace strata::kernels::cpu { struct ExpertLayout { bool native = false; int64_t n_layers = 0, n_expert = NE; + int64_t n_embd = H, n_ff = FF; ///< the model's expert widths (Flash-Next's, or Qwen3.6's 2048 / 512) std::vector fmt; ///< per layer (native packs) std::vector offset, bytes; ///< per layer: where its 512 blobs start, bytes per blob /// Plan v0.3 P6: per layer, the absolute offsets of the gate / up / down tensors in their GGUF files, so the @@ -93,7 +94,10 @@ void q2_rows_any_legacy(const uint8_t* w, size_t row_bytes, int nblocks, const A const ExpertLayout& expert_layout(); /// Reads `/native_experts.txt` when it exists (a native pack), else sets the canonical layout. /// Versions up to kExpertLayoutVersion are read; a newer one is refused (a newer packer wrote it). -bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err); +/// `n_embd` / `n_ff` are the model's expert widths: Flash-Next's by default, Qwen3.6's from its geometry (a native +/// pack only; the canonical Q2_0 layout is Flash-Next's). +bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err, + int64_t n_embd = H, int64_t n_ff = FF); /// The newest native_experts.txt this engine reads. v4 = v3 plus the per-role shard column `gate,up,down`, /// written only when some layer's roles are in different shards (every other pack stays v3, byte for byte). inline constexpr int kExpertLayoutVersion = 4; diff --git a/include/strata/kernels/kv_stream.hpp b/include/strata/kernels/kv_stream.hpp index 8108795b82b..584161b98ee 100644 --- a/include/strata/kernels/kv_stream.hpp +++ b/include/strata/kernels/kv_stream.hpp @@ -68,14 +68,14 @@ struct KvStreamMap { int32_t* slot_block = nullptr; ///< (n_slots,) slot -> block, or -1 int32_t* slot_stamp = nullptr; ///< (n_slots,) the resolve epoch that last used the slot int32_t* slot_ref = nullptr; ///< (n_slots,) clock reference bit - int32_t* ctl = nullptr; ///< kKvCtlInts: epoch, hand, misses of the last call, overflow, u64 counters + int32_t* ctl = nullptr; ///< kKvCtlInts: epoch, hand, blocks copied by the last call, overflow, u64 counters int32_t* miss_block = nullptr; ///< (n_slots,) int32_t* miss_slot = nullptr; ///< (n_slots,) int64_t n_blocks = 0; int64_t n_slots = 0; }; -inline constexpr int kKvCtlInts = 16; ///< [0] epoch [1] hand [2] misses [3] overflow; u64 at [4] misses, [6] lookups, [8] calls +inline constexpr int kKvCtlInts = 16; ///< [0] epoch [1] hand [2] copies [3] overflow; u64 at [4] misses (copied or read in place), [6] lookups, [8] calls /// Bytes of the map's arrays besides the page table (which the state already has): 5 per-slot ints + ctl. inline uint64_t kv_stream_map_bytes(int64_t n_slots) { return (uint64_t) n_slots * 4 * 5 + kKvCtlInts * 4; } @@ -83,11 +83,12 @@ inline uint64_t kv_stream_map_bytes(int64_t n_slots) { return (uint64_t) n_slots void kv_stream_reset(const KvStreamMap& m, void* stream); /// Make every block named by the selections of `n_q` queries resident (ids [n_q][cap], width from -/// steps[q * kStepCount + kStepWidth]). Capturable. `n_slots` must hold the distinct blocks of one call -/// (n_q x (cap / page_size + 2)); a call that cannot sets ctl[3] (see `kv_stream_counters`). +/// steps[q * kStepCount + kStepWidth]). Capturable. When the call names more blocks than `n_slots` can take (a dense +/// selection - qwen35moe attends every cell - past the resident cells), the rest keep page -1 and the attention reads +/// them from the host copy in place (QsaAttnPools::host_*); the call sets ctl[3] (see `kv_stream_counters`). void kv_stream_resolve(const KvStreamMap& m, const QsaAttnPools& slots, const KvHostPools& host, int fmt, const int32_t* ids, const int32_t* steps, int64_t n_q, int64_t cap, const QsaShapes& s, - void* stream); + void* stream, bool dense = false); ///< dense: ids are 0..width-1 (no indexer), walk the blocks /// The static ring table `block -> block % n_slots`. void kv_ring_table(int32_t* page_table, int64_t n_blocks, int64_t n_slots, void* stream); @@ -110,7 +111,7 @@ void kv_unstage_to_host(const QsaAttnPools& stage, const KvHostPools& host, int struct KvStreamCounters { uint64_t misses = 0, lookups = 0, calls = 0; - bool overflow = false; + bool overflow = false; ///< some call named more blocks than the slots hold (the rest read from RAM in place) }; /// Synchronous read of the counters (debug and the end-of-request summary). KvStreamCounters kv_stream_counters(const KvStreamMap& m); diff --git a/include/strata/kernels/qsa.hpp b/include/strata/kernels/qsa.hpp index dc26024ad59..1c53e331326 100644 --- a/include/strata/kernels/qsa.hpp +++ b/include/strata/kernels/qsa.hpp @@ -130,6 +130,10 @@ inline double qsa_freq_base() { return 1e7; } /// not fill one. So the plan's "selection skipped when context <= 2,048" is conservative by two - the source's /// own bound is 2,051, and `ref/qsa.py` L292-294 records the disagreement. Below the bound the selection is /// the IDENTITY, which `topk_512` is required to reproduce exactly rather than approximately. +/// `idx_top_k` of a model without the sparse-attention indexer (Qwen3.6): a budget no context reaches, so the +/// selection width is n_kv and every cached cell is attended (the ids are written by qsa_select_all). +inline constexpr int64_t kDenseTopK = int64_t{1} << 30; + inline int64_t qsa_selection_width(int64_t n_kv, const QsaShapes& s) { const int64_t w = s.idx_top_k + s.idx_block - 1; return n_kv < w ? n_kv : w; diff --git a/include/strata/kernels/qsa_decode_attn.hpp b/include/strata/kernels/qsa_decode_attn.hpp index 45c0f301220..b54d7c83557 100644 --- a/include/strata/kernels/qsa_decode_attn.hpp +++ b/include/strata/kernels/qsa_decode_attn.hpp @@ -30,11 +30,30 @@ struct QsaAttnPools { const uint8_t* k_q4 = nullptr; ///< q4_0 block_q4_0 [page][kv_head][page_size][head_dim / 32 * 18] const uint8_t* v_q4 = nullptr; const int32_t* page_table = nullptr; + /// KV streaming (kv_stream.hpp): the layer's host copy, same formats in the identity layout `[block][kv_head] + /// [page_size]`. A selected cell whose block is not resident (page -1) is read from it in place - a dense + /// selection (qwen35moe: every cell) can name more blocks than the VRAM slots hold. All null when the layer is + /// not streamed: such a cell is masked, as before. + const uint16_t* host_k_pool = nullptr; + const uint16_t* host_v_pool = nullptr; + const int8_t* host_k_q = nullptr; + const int8_t* host_v_q = nullptr; + const uint16_t* host_k_scale = nullptr; + const uint16_t* host_v_scale = nullptr; + const uint8_t* host_k_q4 = nullptr; + const uint8_t* host_v_q4 = nullptr; }; /// Scratch floats for `cap` selected cells: partial accumulators, maxima and sums. uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s); +/// Opt-in split-K (`STRATA_ATTN_SPLITS=P`, 16..65535; 0 = off, the default): past P chunks of 64 cells each KV head +/// runs P partitions that fold their chunks in turn, so the partials (and the scratch above) stop growing with the +/// context. Bit-identical to the default while a selection fits P x 64 cells. Set it before any scratch is sized: +/// the scratch size and the kernels read the same value. The setter exists for the parity test. +void qsa_decode_attn_set_splits(int parts); +int qsa_decode_attn_splits(); + void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream); diff --git a/include/strata/kernels/qsa_prompt_attn.hpp b/include/strata/kernels/qsa_prompt_attn.hpp index 2391aaebc34..d6191237258 100644 --- a/include/strata/kernels/qsa_prompt_attn.hpp +++ b/include/strata/kernels/qsa_prompt_attn.hpp @@ -16,6 +16,9 @@ // Not bitwise equal to `qsa_decode_attn_batch`; `qsa_prompt_attn_parity` bounds the difference and the prompt // quality gate (needles, teacher-forced top-1) checks it end to end. Q4_0 KV (mode 4): each block's codes enter // exactly as int8 and its scale (one per 32 values) in FP32, as int8 KV's do (STRATA_PROMPT_ATTN_Q4=0: the old kernel). +// Dense mode (ids null, 8 query heads per KV head: Qwen3.6) with fp16 or int8 KV on sm_80+: a block takes 8 +// neighbouring queries, so each K/V chunk is read once for all 8, with the same accuracy (STRATA_PROMPT_ATTN_DENSE=0: +// one block per query, as before). #pragma once #include "strata/kernels/qsa_decode_attn.hpp" diff --git a/include/strata/kernels/qsa_select.hpp b/include/strata/kernels/qsa_select.hpp index c36d29cf821..d6d3ca99b46 100644 --- a/include/strata/kernels/qsa_select.hpp +++ b/include/strata/kernels/qsa_select.hpp @@ -45,6 +45,10 @@ void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64 ///< --max-context otherwise sends every short prompt to the slow one). ///< CUDA uses the bound only on sm_75; HIP keeps its existing policy. ///< Omit it for captured graphs whose context can grow after capture. +/// A model without the indexer (Qwen3.6, QsaShapes::idx_top_k = kDenseTopK): every cell of each query's selection +/// width, ascending - the ids qsa_block_topk writes whenever the width is the whole context. Capturable (the width +/// is read from `steps`). ids [nq, cap]. +void qsa_select_all(const int32_t* steps, int64_t nq, int64_t cap, int32_t* ids, void* stream); /// The same ids on a thread-block cluster of 8 CTAs per query (sm_90+, CUDA; S19). qsa_block_topk takes it for calls /// of up to 16 queries unless STRATA_QSA_CLUSTER=0. False (nothing launched) where it cannot run: HIP, a card or a /// build below sm_90, or a capacity whose keys do not fit one cluster's shared memory. Capturable. diff --git a/include/strata/kernels/verify_kernels.hpp b/include/strata/kernels/verify_kernels.hpp index 87e5d5cc322..7de7b2e3eba 100644 --- a/include/strata/kernels/verify_kernels.hpp +++ b/include/strata/kernels/verify_kernels.hpp @@ -38,6 +38,9 @@ void gdn_ab_multi(const float* x, const uint16_t* w_alpha, const uint16_t* w_bet /// The recurrence + output norm for T tokens (h = (T, conv_channels) as q|k|v, gate/beta (T, h_v), z/y /// (T, value_dim)). With `n_keep == nullptr` the state is read and NOT written (verify); otherwise the first /// *n_keep tokens are run and the state is written (commit; `y` may be scratch). Bitwise `fused_gdn_step_norm`. +/// The GDN output gate of gdn_step_norm_multi: sigmoid(z) (the default, qwen4exp) or silu(z) (qwen35moe). Process-wide, +/// set once at load before any window is captured. +void gdn_set_out_gate_silu(bool silu); void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const float* gate, const float* beta, const float* z, const float* gamma, float eps, float* y, int h_k, int h_v, int n_tok, const int32_t* n_keep, void* stream, int t_out_begin = 0, void* xq_out = nullptr); diff --git a/include/strata/prefill/kernels.hpp b/include/strata/prefill/kernels.hpp index 34cc5f66201..fa61237885b 100644 --- a/include/strata/prefill/kernels.hpp +++ b/include/strata/prefill/kernels.hpp @@ -23,6 +23,9 @@ void gr_norm_rs(const float* R, const float* w_norm, float eps, float* rs, uint1 /// gr_mix with xn recomputed from R, rs and w_norm exactly as gr_norm computes it (the same bits). void gr_mix_r(const float* R, const float* rs, const float* w_norm, const float* gated, float* mixed, uint16_t* mixed16, int64_t T, void* stream, uint16_t* mixed_h = nullptr, uint16_t* mixed16_lo = nullptr); +/// One residual stream (qwen35moe): mixed = RMSNorm(R row) * w over T rows of n, with gr_mix_r's other outputs. +void rms_mix(const float* R, const float* w, float eps, int64_t n, float* mixed, uint16_t* mixed16, int64_t T, void* stream, + uint16_t* mixed_h = nullptr, uint16_t* mixed16_lo = nullptr); /// F-2: gr_write, then gr_norm_rs of the next half (its norm weights) over the rows just written - the same bits as /// the two calls, without reading R back. /// S23 (opt-in STRATA_HC_UPMIX=1): the up projection (lo16 x w_up^T, BF16, FP32 accumulate) with gr_mix_r as its @@ -52,31 +55,35 @@ void gr_broadcast(const float* e, float* R, int64_t T, void* stream); // ---- GDN (state 128, 16 k heads, 48 v heads, 10240 conv channels, 4 taps) /// gate[t,h] = softplus(ab[t,h] + dt[h]) * ssm_a[h]; beta[t,h] = sigmoid(ab[t, 48 + h]) (ab: [T, 96]) -void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream); +/// hv: the model's value heads (48 Flash-Next, 32 Qwen3.6; each has its own kernel instances); silu: the output gate. +void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream, + int64_t hv = 48); /// The 4-tap causal conv + SiLU over the chunk (history [C][3] in, updated to the chunk's last three inputs), then /// the L2 norm of the q and k heads of every token. h: [T, C]. -void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream); +void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream, + int64_t hv = 48); /// The recurrence over the chunk, state in registers; y16[t] = rmsnorm(o) * gamma * sigmoid(z) in FP16 (what the /// out projection reads); y is FP32 scratch. void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z, const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream, - int64_t ld16 = 0); // ld16: y16's row stride (0 = 6144; S23 STRATA_PF_PAD pads it) + int64_t ld16 = 0, // ld16: y16's row stride (0 = hv x 128; S23 STRATA_PF_PAD pads it) + int64_t hv = 48, bool silu = false); /// The kernels behind gdn_recurrence, for the parity test: 0 = the column-split kernels + the norm kernel (or the /// one-block-per-head kernel under STRATA_GDN_REC_HEADS), 1 = four lanes per column (no barrier per token) + the grid-stride norm, 2 = the /// one-block-per-head kernel (its fused norm rounds differently). 0 and 1 give the same bits. void gdn_recurrence_variant(int variant, float* state, const float* h, const float* gate, const float* beta, const float* z, const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, - void* stream, int64_t ld16 = 0); + void* stream, int64_t ld16 = 0, int64_t hv = 48, bool silu = false); // ---- MoE /// softmax over 512, top-10 (ties to the lower id), weights renormalised over the ten (the native router). -void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream); +void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream, int64_t k = 10); /// Expert blob (Strata pack layout, Q2_0) -> BF16 matrices: gate/up interleaved [1280, 2560], down [2560, 640]. void blob_dequant(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream); /// h16[n, r] = fp16(silu(gu[n, 2r]) * gu[n, 2r + 1]) (the interleaved expert gate/up) -void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream); +void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream, int64_t n_ff = 640); /// h16[n, r] = fp16(silu(g[n, r]) * u[n, r]) (the shared expert, gate and up separate, width 640) -void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream); +void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream, int64_t n_ff = 640); /// dst[i] = src[i] for n int32s, as a kernel: either side may be mapped host memory, and the copy never waits /// behind the copy engine's queue (the prompt path's grouping tables, while the expert stream fills it). void copy_i32(int32_t* dst, const int32_t* src, int64_t n, void* stream); @@ -87,7 +94,7 @@ void copy_f32_wide(float* dst, const float* src, int64_t n, void* stream); void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int64_t n, int64_t width, void* stream); /// bo[t, :] = shared[t, :] * sigmoid(sg[t]) + sum_k w[t, k] * D[slot[t, k], :] void moe_combine(const float* D, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, - int64_t T, void* stream); + int64_t T, void* stream, int64_t n_embd = 2560, int64_t k = 10); // 2560 x 10 or 2048 x 8 /// --peer-device's prompt share as sums: bo[t, :] = (sum over the k with slot[t, k] < rows_local of w[t, k] * /// D[slot[t, k], :], in k order) + peer[t, :] + shared[t, :] * sigmoid(sg[t]). `peer` (the peer's per-token sums) /// may be mapped host memory. @@ -117,10 +124,11 @@ void rms_rows(float* x, const float* w, int64_t rows, int64_t cols, int64_t ld, /// scaling (rope_scaling.hpp) rides in as the process config - none is today's arithmetic exactly. void rope(float* x, int64_t T, int64_t heads, int64_t dim, int64_t ld, int64_t pos0, const strata::kernels::RopeScaling& scaling, void* stream); -/// q_full [T, 24, 512] (q | gate per head) -> q [T, 24, 256] -void split_q(const float* q_full, float* q, int64_t T, void* stream); +/// q_full [T, n_head, 512] (q | gate per head) -> q [T, n_head, 256]; n_head 24 (Flash-Next) or 16 (Qwen3.6) +void split_q(const float* q_full, float* q, int64_t T, void* stream, int64_t n_head = 24); /// attn[t, h, d] *= sigmoid(q_full[t, h, 256 + d]) -> out16 (fp16 bits: the o-projection is quantized) -void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream, int64_t ld16 = 0); +void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream, int64_t ld16 = 0, + int64_t n_head = 24); /// K and V of T consecutive cells (positions pos0..pos0+T-1; K normed and rotated) into the paged pools: FP16 /// (`k_pool`/`v_pool`) or INT8 codes + FP16 scale per 64 (`k_q`...), the decode append's arithmetic. diff --git a/include/strata/prefill/prefill.hpp b/include/strata/prefill/prefill.hpp index 1f0ad13fa6f..19f1d82e34b 100644 --- a/include/strata/prefill/prefill.hpp +++ b/include/strata/prefill/prefill.hpp @@ -64,6 +64,9 @@ class Prefill { /// callbacks stay. The device `init` ran on must be current. void reset(); + /// Whether this prompt path runs model `g` (the geometry `init` requires): Flash-Next's. For another model + /// (Qwen3.6 until its prompt path exists) the prompt is read through the decode windows instead. + static bool supports(const core::ModelGeometry& g); /// `host_res`: the static residency table (n_layers x n_expert, slot or -1) or null; `cache` its slots. /// `borrow`/`borrow_bytes`: device memory to carve every buffer from (the top slots of the expert cache, /// lent for the prompt and refilled after it); null = allocate normally. diff --git a/serve/frontend.py b/serve/frontend.py index e8cc2335f01..ba8232830cb 100644 --- a/serve/frontend.py +++ b/serve/frontend.py @@ -68,9 +68,9 @@ def render(self, messages: list[dict], tools: list[dict] | None = None, add_gene def _detect_caps(self) -> dict[str, bool]: """llama.cpp's capability names, checked at load time against this template and Strata's tool-call format. These are rendering hints, not a guarantee that the model will follow a request.""" - def render(messages, tools=None): + def render(messages, tools=None, **kwargs): try: - return self.render(messages, tools=tools) + return self.render(messages, tools=tools, **kwargs) except Exception as exc: # noqa: BLE001 - a hint only: never stop the server for a probe # A template may reject a role or feature (a custom one may fail in any way): the feature is off, # discovery still answers for the others, and the start-up goes on. @@ -110,7 +110,7 @@ def calls_supported(count): "supports_system_role": "strata_caps_system" in render( [{"role": "system", "content": "strata_caps_system"}, user]), "supports_parallel_tool_calls": calls_supported(2), - "supports_preserve_reasoning": "strata_caps_reasoning" in render(history)} + "supports_preserve_reasoning": "strata_caps_reasoning" in render(history, preserve_thinking=True)} # ------------------------------------------------------------------------------------------------ requests @@ -471,6 +471,8 @@ def openai_to_messages(req: dict) -> tuple[list[dict], list[dict] | None, dict]: kwargs = {"enable_thinking": False} elif k == "reasoning_effort" and "enable_thinking" not in kwargs: kwargs.update(effort_kwargs(v)) + elif k == "preserve_thinking": + kwargs["preserve_thinking"] = bool(v) return _late_system_to_user(messages), tools, kwargs @@ -567,6 +569,11 @@ def anthropic_to_messages(req: dict, think_unasked: bool = True) -> tuple[list[d # A config's reasoning_effort still applies: Service.with_shared sets output_config before this runs. A # request that gives its own reasoning_budget_tokens (#123) asks for thinking, so it thinks as before. kwargs["enable_thinking"] = False + # Not part of Anthropic's API: llama.cpp's chat_template_kwargs, for preserve_thinking only (the shared Chat + # settings or the config fill it, Service.with_shared), so earlier thinking blocks reach the template + ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} + if "preserve_thinking" in ctk: + kwargs["preserve_thinking"] = bool(ctk["preserve_thinking"]) return _late_system_to_user(messages), tools, kwargs diff --git a/serve/server.py b/serve/server.py index ff1b935c359..157eb2b1fe0 100644 --- a/serve/server.py +++ b/serve/server.py @@ -2872,6 +2872,7 @@ def __init__(self, engine: Engine, tokenizer, template: ChatTemplate, model_name self.shared = {} # the web app's Chat settings for every client (POST /settings) self.shared_path = None # where they are kept between starts (next to the config) self.config_effort = None # #1641: the config's "reasoning_effort": the level for requests with none + self.config_preserve = None # the config's "preserve_thinking": true/false, or None (the template's) self.keep_awake = None # #1727 (opt-in): serve.keepawake.KeepAwake, the config's "prevent_sleep" self.fifo = threading.Lock() self.slot_save_path = None # --slot-save-path: /slots/0?action=save|restore (off when None) @@ -3291,7 +3292,8 @@ def with_shared(self, req: dict, api: str) -> dict: """The request with the shared thinking level and max tokens filled in where it has none of its own.""" s = self.shared effort = s.get("reasoning_effort") or self.config_effort # the Chat settings first, then the config's - if not s and not effort: + preserve = s.get("preserve_thinking", self.config_preserve) # the same order + if not s and not effort and preserve is None: return req req = dict(req) if "max_tokens" in s and not req.get("max_tokens") and not req.get("max_completion_tokens"): @@ -3307,6 +3309,12 @@ def with_shared(self, req: dict, api: str) -> dict: req["thinking"] = {"type": "disabled"} else: req["output_config"] = {"effort": effort} + if preserve is not None: + # The template's preserve_thinking (keep the reasoning of turns before the latest user message), for both + # APIs: openai_to_messages and anthropic_to_messages read chat_template_kwargs; the request's own value wins. + ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} + if "preserve_thinking" not in ctk: + req["chat_template_kwargs"] = {**ctk, "preserve_thinking": bool(preserve)} return req def start_telemetry(self): @@ -6114,6 +6122,9 @@ def clean_shared_defaults(d) -> dict: elif key == "experimental_speed_projection": if not isinstance(value, bool): raise ValueError("experimental_speed_projection: true or false") + elif key == "preserve_thinking": + if not isinstance(value, bool): + raise ValueError("preserve_thinking: true or false") else: raise ValueError(f"unknown setting {key!r}") out[key] = float(value) if key in ("temperature", "top_p") else value @@ -6453,6 +6464,12 @@ def main() -> int: if svc.config_effort: print(f"[strata] thinking level: {svc.config_effort} (reasoning_effort in the config; a request, or the " "Chat settings shared with other apps, can set its own)", flush=True) + if cfg.get("preserve_thinking") is not None: # like llama.cpp's --reasoning-preserve: opt-in + if not isinstance(cfg["preserve_thinking"], bool): + raise SystemExit(f"[strata] config \"preserve_thinking\" must be true or false, not {cfg['preserve_thinking']!r}") + svc.config_preserve = cfg["preserve_thinking"] + print(f"[strata] earlier turns' reasoning: {'kept' if svc.config_preserve else 'dropped'} (preserve_thinking in " + "the config; a request's chat_template_kwargs can set its own)", flush=True) svc.codex_thread_titles = cfg.get("codex_thread_titles") is True # #923: opt-in, off by default svc.codex_compaction_cache = cfg.get("codex_compaction_cache") is True # #924: opt-in, off by default try: diff --git a/serve/test_server.py b/serve/test_server.py index bbb176ada83..e78ffd33d1d 100644 --- a/serve/test_server.py +++ b/serve/test_server.py @@ -2711,6 +2711,45 @@ def test_the_flag_in_args_becomes_the_key(self): self.assertEqual(cfg, {"args": ["--ctx", "4096"]}) +class PreserveThinking(unittest.TestCase): + """The template's preserve_thinking (keep the reasoning of turns before the latest user message): a request's + chat_template_kwargs on either API, else the shared Chat settings, else the config's "preserve_thinking".""" + + def svc(self, config=None, shared=None): + tok = ByteTokenizer() + svc = Service(MockEngine(tok, "ok", max_context=CTX), tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) + svc.config_preserve = config + svc.shared = shared or {} + return svc + + def test_default_changes_nothing(self): + req = {"messages": []} + self.assertIs(self.svc().with_shared(req, "anthropic"), req) + + def test_config_then_shared_then_request(self): + ctk = lambda svc, req, api: svc.with_shared(req, api).get("chat_template_kwargs") + for api in ("openai", "anthropic"): + self.assertEqual(ctk(self.svc(True), {}, api), {"preserve_thinking": True}) + self.assertEqual(ctk(self.svc(True, {"preserve_thinking": False}), {}, api), {"preserve_thinking": False}) + own = {"chat_template_kwargs": {"preserve_thinking": False, "x": 1}} + self.assertEqual(ctk(self.svc(True), own, api), {"preserve_thinking": False, "x": 1}) + + def test_both_apis_pass_it_to_the_template(self): + from serve.frontend import anthropic_to_messages, openai_to_messages + reasoning = "strata_earlier_reasoning" + oa = [{"role": "user", "content": "a"}, {"role": "assistant", "content": "b", "reasoning_content": reasoning}, + {"role": "user", "content": "c"}] + an = [oa[0], {"role": "assistant", "content": [{"type": "thinking", "thinking": reasoning}, + {"type": "text", "text": "b"}]}, oa[2]] + svc = self.svc() + for keep in (True, False): + for conv, msgs in ((openai_to_messages, oa), (anthropic_to_messages, an)): + with self.subTest(keep=keep, api=conv.__name__): + m, tools, kwargs = conv({"messages": msgs, "chat_template_kwargs": {"preserve_thinking": keep}}) + self.assertEqual(kwargs.get("preserve_thinking"), keep) + self.assertEqual(reasoning in svc.template.render(m, tools=tools, **kwargs), keep) + + class SharedSettings(unittest.TestCase): """The web app's "Use for other apps too": POST /settings makes its Chat settings every client's defaults.""" diff --git a/serve/web/app.js b/serve/web/app.js index 4460a95a7c1..2bab800fb21 100644 --- a/serve/web/app.js +++ b/serve/web/app.js @@ -742,9 +742,10 @@ function assistantMessages(m) { // #1392: a turn with no answer text (only reasoning, a stop before the first content token, or an error) still goes // back as an assistant turn, so the history keeps alternating and a reasoning model sees that it already answered. // The server leaves such a turn out of the prompt itself (serve/frontend.py, #843); reasoning_content rides along. + // The reasoning rides along on every turn too, so a server set to preserve_thinking keeps earlier thinking in the prompt. if (!ran.length) { const msg = {role: "assistant", content: m.text || ""}; - if (!m.text && m.reasoning) msg.reasoning_content = m.reasoning; + if (m.reasoning) msg.reasoning_content = m.reasoning; return [msg]; } const out = []; diff --git a/setup.py b/setup.py index 7d291e1c660..6ff0d0068a2 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,8 @@ 6. prepares the model for Strata and fetches the MTP draft layer (~5 GB, from the original Qwen checkpoint) 7. writes run-.bat / run-.sh and starts the model -Options: --family qwen|swift, --model Q2_0|IQ2_XS|IQ3_XXS|IQ3_S, --context 32768, --rope-scaling none|linear|yarn +Options: --family qwen|swift|coder|unsloth|qwen36, --model Q2_0|IQ2_XS|IQ3_XXS|IQ3_S (qwen36: UD-IQ4_XS|UD-IQ3_S), +--context 32768, --rope-scaling none|linear|yarn (--rope-scale F; past the trained 262144 the setup adds yarn and the factor is the final context over 262144, at least 1 - an explicit --rope-scaling none is refused for such a context), --vision yes|no|gpu|cpu, --port 8080, --yes (recommended @@ -68,6 +69,9 @@ "ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF": "b22d729eae29b5796f76fb70f91aef549b9fc52c", # 2026-09-24 "ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-Coder-GGUF": "5348543e0147355ac9cbcb031184a3546350988e", # 2026-09-29 "unsloth/Qwen3.8-Flash-Next-GGUF": "38bb39ee97821de2c9009abb7e93950eec396e66", # 2026-09-30 + "unsloth/Qwen3.6-35B-A3B-MTP-GGUF": "5bc3e238d916f48a861bac2f8a1990a0e9b7e98d", + "bartowski/Ornith-1.5-35B-A3B-GGUF": "64b0493d34a5ca4c1b4ad67bb99b41d74b4f07d6", + "bartowski/Qwen_Qwen3.6-35B-A3B-GGUF": "5c2410d71524f4f72b023ce8daf7a80528226d5f", # its mmproj (images) } @@ -198,6 +202,14 @@ def ms_meta(repo: str, path: str): PY_PACKAGES = ["numpy", "jinja2", "regex", "pyyaml", "tqdm", "requests", "cmake", "ninja", "pillow", "psutil"] REQUIREMENTS = ROOT / "requirements.txt" # the same packages and their dependencies, pinned (#214) +# The first engine that runs Qwen3.6-35B-A3B (general.architecture qwen35moe): this source tree's version +# (CMakeLists.txt). The ready-made 0.1.40 engines were built before the qwen35moe support, so the version alone does +# not tell them apart: setup looks for SMALL_ENGINE_MARK in the engine itself (engine_runs_small). +QWEN36_ENGINE = (0, 1, 43) +# A GGUF key only the engine's own qwen35moe support names (src/core/layout.cpp). Not the bare architecture name: the +# image encoder links llama.cpp, whose table of architectures has it. +SMALL_ENGINE_MARK = b"qwen35moe.block_count" + MODELS = { # the original model only for now: Swift 1.5's Q2_0 files split one layer's experts across the two shards, which # the pack tool (tools/iq_pack.py) cannot prepare yet (#171) @@ -230,6 +242,45 @@ def ms_meta(repo: str, path: str): "download_gb": 93.7, "ram_gb": 48, "arena_gb": 59.5, "families": ("unsloth",), "budget": True, "shards": 3, "file": "Qwen3.8-Flash-Next-{q}-0000{i}-of-00003.gguf", "engine": (0, 1, 38), "vision": True}, + # Qwen3.6-35B-A3B (family "qwen36", a smaller model for 8-12 GB cards and 16-32 GB of RAM): Unsloth's single-file + # GGUFs with the MTP draft layer inside. Its sizes share their names with other files (UD-IQ4_XS is also an + # Unsloth Flash-Next size), so their keys here carry the model and "size" is the published name: --model + # UD-IQ4_XS with --family qwen36 is this entry (model_key), and the folder, pack and config names use the size. + "Qwen3.6-UD-IQ4_XS": {"size": "UD-IQ4_XS", "about": "~4-bit i-quant (Unsloth Dynamic), the recommended size: " + "~14 GiB of experts in RAM, comfortable on 32 GB", + # 18,209,036,576 bytes; experts 14.17 GiB = 15.2 GB (measured: the engine's "loaded 14.17 GiB") + "download_gb": 18.2, "ram_gb": 26, "arena_gb": 15.2, "families": ("qwen36",), + "engine": QWEN36_ENGINE}, + "Qwen3.6-UD-IQ3_S": {"size": "UD-IQ3_S", "about": "~3.5-bit i-quant (Unsloth Dynamic), smaller: for 16-24 GB of " + "RAM (16 GB: a 12 GB card, the low-RAM mode)", + # 15,346,432,288 bytes; experts 11.99 GiB = 12.9 GB (measured: the engine's "loaded 11.99 GiB") + "download_gb": 15.3, "ram_gb": 24, "arena_gb": 12.9, "families": ("qwen36",), + "engine": QWEN36_ENGINE}, + # Ornith-1.5-35B-A3B (family "ornith"): ornith-ai's fine-tune of the same architecture (qwen35moe), for agentic + # coding. No Unsloth files exist for it; bartowski's single-file GGUFs (with the MTP layer, stored Q4_0) are the + # closest to the Qwen3.6 sizes. Its draft layer is accepted less often than Qwen3.6's (61% vs 80% of the drafts + # on a chat prompt, measured; a Q8_0 copy of it, mudler's APEX-MTP-Compact file, 44%), so + # it writes ~10-20% slower on the same PC - the answers are the main model's either way. Experts: measured. + "Ornith-1.5-IQ4_XS": {"size": "IQ4_XS", "about": "~4-bit i-quant (bartowski), the recommended size: " + "~16 GiB of experts in RAM, comfortable on 32 GB", + # 19,278,554,784 bytes; experts 15.94 GiB = 17.1 GB (measured) + "download_gb": 19.3, "ram_gb": 28, "arena_gb": 17.1, "families": ("ornith",), + "engine": QWEN36_ENGINE}, + "Ornith-1.5-IQ3_XXS": {"size": "IQ3_XXS", "about": "~3-bit i-quant (bartowski), smaller: for 16-24 GB of RAM " + "(16 GB: a 12 GB card, the low-RAM mode)", + # 15,340,447,392 bytes; experts 12.38 GiB = 13.3 GB (measured) + "download_gb": 15.3, "ram_gb": 24, "arena_gb": 13.3, "families": ("ornith",), + "engine": QWEN36_ENGINE}, +} +# Qwen3.6-35B-A3B's two files at the pinned revision: name -> (bytes, sha256), checked after the download +QWEN36_FILES = { + "Qwen3.6-35B-A3B-UD-IQ4_XS.gguf": (18209036576, "df27a780435b7b45c2597536112ea3cb091f8544c3d0c3318d9f4258b31f7adf"), + "Qwen3.6-35B-A3B-UD-IQ3_S.gguf": (15346432288, "ab639a7f330f96c47d3e6c2dd2d6445182e7b763e17e6048dc850a71bbc9f27f"), +} +# Ornith-1.5-35B-A3B's files at the pinned revision: name -> (bytes, sha256) +ORNITH_FILES = { + "Ornith-1.5-35B-A3B-IQ4_XS.gguf": (19278554784, "d6aef57fa948e9bba3ca4959b3c237ed898c605471f48c73a32cedbd24aabe70"), + "Ornith-1.5-35B-A3B-IQ3_XXS.gguf": (15340447392, "8918ccb9ee29abe3875efec0c3e86f0e35ef0b8a03e3f0e1c422869859a518d0"), } # The experimental Unsloth file's four shards at the pinned revision: name -> (bytes, sha256), checked after the # download (setup trusts no other model file by name and size alone either: check_shards reads their directories). @@ -297,12 +348,38 @@ def ms_meta(repo: str, path: str): "mmproj": "mmproj-Qwen3.8-Flash-Next-BF16.gguf", "name": "qwen3.8-flash-next-unsloth", "vision": False, "pack_args": ["--compat-bf16"], "sha256": {**UNSLOTH_SHARDS, **UNSLOTH_IQ4_XS_SHARDS}}, + # Qwen3.6-35B-A3B (general.architecture qwen35moe): one GGUF file, no PLE table (no --ple-gguf), its MTP draft + # layer inside the file (blk.40: --mtp is the model file, nothing else is fetched), the pack built like Unsloth's + # (--compat-bf16). Images: the stock encoder (tools/vision) reads Qwen3.6's mmproj, bartowski's copy for its + # distinct name (Unsloth's is "mmproj-BF16.gguf"); one GPU (no layer split yet); its batched prompt path is + # NVIDIA-only (sm_75+) and it has not been run on AMD. Suggested where no Flash-Next size fits the RAM (qwen36_recommended). + "qwen36": {"title": "Qwen3.6-35B-A3B", "by": "Qwen's smaller MoE; Unsloth's GGUFs with its MTP layer", + "about": "a smaller model for 16-32 GB of RAM and 8-12 GB cards (18 GB download)", + "hf": hf("unsloth/Qwen3.6-35B-A3B-MTP-GGUF"), "file": "Qwen3.6-35B-A3B-{q}.gguf", "shards": 1, + "tag": "qwen36-", "mmproj_hf": hf("bartowski/Qwen_Qwen3.6-35B-A3B-GGUF"), + "mmproj": "mmproj-Qwen_Qwen3.6-35B-A3B-bf16.gguf", "name": "qwen3.6-35b-a3b", + "pack_args": ["--compat-bf16"], "sha256": QWEN36_FILES, "profile": "expert-profile-qwen36.bin", + "license": "Apache 2.0: https://huggingface.co/unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "own_mtp": True, "ple": False, "one_gpu": True, "nvidia_only": True}, + # Ornith-1.5-35B-A3B: a fine-tune of the same architecture (agentic coding), the same engine path as qwen36 (its + # draft layer in the file, no PLE, one GPU, NVIDIA's batched prompt path); the Qwen3.6 expert profile and draft + # vocabulary apply (the same geometry and tokenizer). Images: bartowski's mmproj of it. + "ornith": {"title": "Ornith-1.5-35B-A3B", "by": "ornith-ai's agentic-coding fine-tune of the same 35B-A3B model", + "about": "Qwen3.6-35B-A3B's size, tuned for coding agents (19 GB download)", + "hf": hf("bartowski/Ornith-1.5-35B-A3B-GGUF"), "file": "Ornith-1.5-35B-A3B-{q}.gguf", "shards": 1, + "tag": "ornith15-", "mmproj_hf": hf("bartowski/Ornith-1.5-35B-A3B-GGUF"), + "mmproj": "mmproj-Ornith-1.5-35B-A3B-bf16.gguf", "name": "ornith-1.5-35b-a3b", + "pack_args": ["--compat-bf16"], "sha256": ORNITH_FILES, "profile": "expert-profile-qwen36.bin", + "license": "MIT: https://huggingface.co/ornith-ai/Ornith-1.5-35B-A3B", + "own_mtp": True, "ple": False, "one_gpu": True, "nvidia_only": True}, } MMPROJ = "mmproj-Qwen3.8-Flash-Next-BF16.gguf" # EXPERIMENTAL, off by default (setup asks): a control vector shipped with the repository, see its README ESP_VECTOR = ROOT / "data" / "experimental-speed-projection" / "Qwen3.8-Flash-Next-experimental-speed-projection.gguf" # the image encoder on the GPU (~1.2 GB at 1024 image tokens) warms up before the engine starts, so the engine # sizes its expert slots around it and the default reserve (700 MiB) is enough; engines before 0.1.2 need more +# Qwen3.6's image encoder on the GPU (BF16 mmproj, CUDA): 1,254 MiB in nvidia-smi before any picture (4070 Ti) +VISION_GPU_QWEN36_MIB = 1254 VISION_GPU_SMALL_RESERVE_MIB = 1000 # the tip for images on a <= 12 GB card (the engine's LOW line asked ~1003) VISION = {"gpu": {"max_tokens": 1024, "reserve_mib": 700}, "cpu": {"max_tokens": 300, "reserve_mib": 700}} @@ -846,6 +923,16 @@ def gpus(): return found +def gpu_free_gb(index: int): + """The VRAM free on NVIDIA GPU `index` right now (nvidia-smi's memory.free; the desktop and other programs + take theirs from it), or None when nvidia-smi cannot say.""" + s = out(["nvidia-smi", f"--id={index}", "--query-gpu=memory.free", "--format=csv,noheader,nounits"]) + try: + return float(s.strip().splitlines()[0]) / 1024.0 + except (ValueError, IndexError): + return None + + def gpu_drives_display(g) -> bool: """#779: True when nvidia-smi says this card has a display attached (display_active Enabled): its desktop needs VRAM too, and a full expert cache beside it has crashed laptops.""" @@ -1187,13 +1274,89 @@ def split_mmap(cfg: dict) -> bool: def model_file(fam: dict, model: str, i: int) -> str: """Shard i's file name: the family's pattern, or the model's own (#621: UD-IQ4_XS has three shards, not four).""" - return MODELS.get(model, {}).get("file", fam["file"]).format(q=model, i=i) + return MODELS.get(model, {}).get("file", fam["file"]).format(q=size_of(model), i=i) def model_shards(fam: dict, model: str) -> int: return MODELS.get(model, {}).get("shards", fam.get("shards", 2)) +def size_of(model: str) -> str: + """The published size name of a MODELS key: the key itself, or its "size" (Qwen3.6-UD-IQ4_XS -> UD-IQ4_XS).""" + return MODELS.get(model, {}).get("size", model) + + +def model_key(family: str, size: str | None) -> str | None: + """The MODELS key of `size` in `family` (--model UD-IQ4_XS is Unsloth's Flash-Next file with --family unsloth, + Qwen3.6-35B-A3B's with --family qwen36); a key is taken as it is. None: the family has no such size.""" + if not size: + return None + if size in MODELS and family in MODELS[size].get("families", ("qwen", "swift")): + return size + return next((m for m, d in MODELS.items() if size_of(m).upper() == size.upper() + and family in d.get("families", ("qwen", "swift"))), None) + + +def model_choices() -> list: + """What --model accepts: every size name, and the keys of the sizes another family shares a name with.""" + return list(dict.fromkeys([*(size_of(m) for m in MODELS), *MODELS])) + + +def config_family(cfg_path: Path) -> str: + """The family a strata-.json was set up for, from its tag (the original model's has no prefix).""" + tag = Path(cfg_path).stem[len("strata-"):] + return next((f for f, d in FAMILIES.items() if d["tag"] and tag.startswith(d["tag"])), "qwen") + + +def one_gpu_family(cfg_path: Path) -> bool: + """A family the engine runs on one GPU only (no layer split yet: Qwen3.6-35B-A3B).""" + return bool(FAMILIES[config_family(cfg_path)].get("one_gpu")) + + +FLASH_NEXT_FLOOR_SLACK_GB = 4 # step 1's tolerance under a size's RAM line ("the smallest model needs ...") + + +def small_family(family: str) -> bool: + """A family of the 35B-A3B models (Qwen3.6-35B-A3B and its fine-tunes: qwen36, ornith) - the model file holds its + draft layer ("own_mtp"); the Flash-Next families do not.""" + return bool(FAMILIES.get(family, {}).get("own_mtp")) + + +def engine_runs_small(eng: Path) -> bool: + """Does the engine in `eng` run the 35B-A3B families (qwen35moe): its program names SMALL_ENGINE_MARK.""" + try: + return SMALL_ENGINE_MARK in (eng / EXE).read_bytes() + except OSError: + return False + + +def small_model(d: dict) -> bool: + """A MODELS entry of a 35B-A3B family.""" + return any(small_family(f) for f in d.get("families", ())) + + +def qwen36_recommended(ram: float) -> bool: + """Is Qwen3.6-35B-A3B the suggested family on this PC: no Flash-Next size reaches its RAM line (the Coder's 32 GB, + less step 1's 4 GB of slack) - the 16-24 GB PCs. Where a Flash-Next size fits, the default stays as it was.""" + floor = min(d["ram_gb"] for d in MODELS.values() if not small_model(d)) + return ram < floor - FLASH_NEXT_FLOOR_SLACK_GB + + +def small_size(family: str, ram: float, vram_gb: float = 0.0) -> str: + """The size of a 35B-A3B family suggested for this RAM: the biggest (first listed) whose experts fit the RAM with + the usual room beside them (no low-RAM mode), else the smallest (in the low-RAM mode). On a card under + SMALL_MODEL_VRAM_GB the smallest: the 4-bit size needs ~0.4 GB more of the card free (measured, see there).""" + sizes = [m for m in MODELS if family in MODELS[m].get("families", ()) and not MODELS[m].get("experimental")] + if 0 < vram_gb < SMALL_MODEL_VRAM_GB: + return min(sizes, key=lambda m: MODELS[m]["arena_gb"]) + return next((m for m in sizes if not low_ram_needed(m, ram)), min(sizes, key=lambda m: MODELS[m]["arena_gb"])) + + +def qwen36_size(ram: float) -> str: + """The Qwen3.6 size suggested for this RAM (UD-IQ4_XS from ~26 GB, UD-IQ3_S from ~24 GB).""" + return small_size("qwen36", ram) + + def budget_model(cfg: dict) -> str: """The Unsloth model a config with a RAM budget runs, from its --native shard's name (UD-Q4_K_XL by default).""" a = cfg.get("args", []) @@ -1280,8 +1443,8 @@ def recommend_remote_expert_opt(cfg: dict, off: bool = False) -> None: def offer_together(cfg_path: Path, cfg: dict, yes: bool) -> dict: """Starting a model set up for one card on a PC with two or more that can share it: asked once (the answer is saved in its config).""" - if isinstance(cfg.get("gpu"), list) or cfg.get("gpus_asked"): - return cfg + if isinstance(cfg.get("gpu"), list) or cfg.get("gpus_asked") or one_gpu_family(cfg_path): + return cfg # (Qwen3.6-35B-A3B: one GPU, no layer split yet) found = gpus() can = together_ok(found) if not can: @@ -1577,7 +1740,9 @@ def gguf_dir_shards(folder: Path, fam: dict, model: str) -> list[Path]: GGUF_QUANT = re.compile(r"(? list[Path]: def gguf_unsupported(name: str) -> str | None: """#444: the quantization a GGUF's name says, when it is one Strata cannot run (not a setup size); else None.""" m = GGUF_QUANT.search(name) - return m.group(1) if m and m.group(1).upper() not in MODELS and not name.lower().startswith("mmproj") else None + sizes = {size_of(k).upper() for k in MODELS} + return m.group(1) if m and m.group(1).upper() not in sizes and not name.lower().startswith("mmproj") else None def gguf_choice(name: str) -> tuple | None: @@ -1594,7 +1760,7 @@ def gguf_choice(name: str) -> tuple | None: for f, d in FAMILIES.items(): for m in MODELS: if f in MODELS[m].get("families", ("qwen", "swift")) and name == model_file(d, m, 1): - return f, m + return f, size_of(m) return None @@ -4018,10 +4184,11 @@ def choices_from_config(cfg_path: Path) -> dict: """The setup answers a config was written with (family, size, context, KV, images, projection, network).""" cfg = json.loads(cfg_path.read_text(encoding="utf-8-sig")) tag = cfg_path.stem[len("strata-"):] - family = next((f for f, d in FAMILIES.items() if d["tag"] and tag.startswith(d["tag"])), "qwen") + family = config_family(cfg_path) model = (tag[len(FAMILIES[family]["tag"]):] if tag.startswith(FAMILIES[family]["tag"]) else tag).upper() - if model not in MODELS: # (sizes have no dash except UD-Q4_K_XL: the old rule) + if model_key(family, model) is None and model not in MODELS: # (sizes have no dash except UD-*: the old rule) model = tag.split("-")[-1].upper() + model = model_key(family, model) or model # the MODELS key (Qwen3.6's UD-IQ4_XS is not Unsloth's) a = cfg.get("args", []) val = lambda k: a[a.index(k) + 1] if k in a and a.index(k) + 1 < len(a) else None # noqa: E731 vis = cfg.get("vision") @@ -4276,6 +4443,8 @@ def update_install(have: list, a) -> int: cfg = upgrade_config(cfg_path, json.loads(cfg_path.read_text(encoding="utf-8-sig"))) if "--mtp" in cfg["args"][:-1]: refresh_draft_vocab(Path(cfg["args"][cfg["args"].index("--mtp") + 1]), cfg.get("draft_vocab", "cjk")) + if point_draft_vocab(cfg): + write_config(cfg_path, cfg) if cfg.get("backend") == "hip" and WIN: hip_runtime_beside_exe(Path(cfg["exe"]).parent) # #468 #461 ok(f"{cfg.get('model_name', cfg_path.stem)}: up to date") @@ -4338,6 +4507,12 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_ cfg_path.touch() # the most recently used model if "--mtp" in cfg["args"][:-1]: refresh_draft_vocab(Path(cfg["args"][cfg["args"].index("--mtp") + 1]), cfg.get("draft_vocab", "cjk")) + if point_draft_vocab(cfg): + write_config(cfg_path, cfg) + if isinstance(gpu, list) and one_gpu_family(cfg_path): + warn(f"{FAMILIES[config_family(cfg_path)]['title']} runs on one GPU for now (the engine has no layer split " + f"for it yet): this start uses GPU {gpu[0]} only") + gpu = gpu[0] cmd = [sys.executable, str(ROOT / "serve" / "server.py"), "--engine", "strata", "--config", str(cfg_path), "--port", str(port or cfg.get("port", 8080))] if cfg.get("backend") == "hip": # AMD, numbered as HIP numbers them (setup's KFD order) @@ -4506,6 +4681,20 @@ def draft_vocab_note(vram_gb: float, chosen: str | None) -> list[str]: SMALL_CARD_GB = 7.5 # #496: a card under 8 GB gets a tip (an 8 GB card lists 7.99) +# The 35B-A3B families (qwen36, ornith) on a card under this: setup recommends the smallest size and an 8K context and +# says what was measured. Qwen3.6 with its draft layer (RTX 4070 Ti, the rest of its VRAM held): UD-IQ3_S at 32K and +# int8 K/V started from 4.17 GB free (73 tok/s), at 8K with --draft-vocab en from 3.65-3.9 GB; UD-IQ4_XS needed 4.68 +# and 4.3 GB. Below that the engine stops at the start and says how much is short. (Emulated sizes, not measured on +# 4 and 6 GB cards.) +SMALL_MODEL_VRAM_GB = 5.5 + + +def small_model_vram_note(vram_gb: float) -> str: + """What a card under SMALL_MODEL_VRAM_GB can expect from Qwen3.6-35B-A3B / Ornith (measured, see above).""" + return (f"a {vram_gb:.0f} GB card is at this model's floor: its smallest size starts with ~3.7-3.9 GB of the card " + "free (no display on it, or a light desktop) with an 8K context and --draft-vocab en, and then writes ~70 " + "tokens/s; with less free the engine stops at the start and says how much is short. With 4.2 GB free or " + "more it runs with 32K") def small_card_note(ctx: int, draft_vocab: str | None) -> list[str]: @@ -4595,10 +4784,11 @@ def arg_after(args, flag): return args[args.index(flag) + 1] if flag in args[:-1] else None -def bench_tips(args, env, ram: float, model_ram_gb: float, vram_gb: float, vision: str, win: bool) -> list[str]: +def bench_tips(args, env, ram: float, model_ram_gb: float, vram_gb: float, vision: str, win: bool, + small: bool = False) -> list[str]: """Recommendations from the community bench data (plan 0.1.40 item 14). Text only: nothing here changes a default, the config or the engine's arguments (recommend, never force). `ram` is this PC's RAM, `model_ram_gb` - what the chosen model keeps in it.""" + what the chosen model keeps in it; `small`: a 35B-A3B model (small_family).""" tips = [] cache = arg_after(args, "--expert-cache") if win and cache is not None and cache != "auto": @@ -4620,7 +4810,8 @@ def bench_tips(args, env, ram: float, model_ram_gb: float, vram_gb: float, visio tips.append(f"tip: for several agents or clients at once, --conversation-cache-mib {AGENT_CACHE_MIB} in the " "config's args keeps each one's conversation; without it they were measured re-reading ~90% of " "their prompts (#882 #440; docs/DETAILS.md, Multiple conversations)") - if vision == "gpu" and 0 < vram_gb <= SMALL_VISION_VRAM_GB: + # Not the 35B-A3B models: Ornith IQ4_XS on a 4070 Ti read a 24K-token prompt at ~4,130 tok/s with or without it + if vision == "gpu" and 0 < vram_gb <= SMALL_VISION_VRAM_GB and not small: tips.append("tip: the image encoder on the GPU makes a prompt's chunks smaller on a card this size (243 vs 750 " "tok/s measured, #469); run setup again with --vision cpu to read prompts ~3x faster, at about " "2-3 s per picture") @@ -4679,6 +4870,22 @@ def refresh_draft_vocab(rt: Path, choice: str = "cjk") -> None: shutil.copyfile(new, dst) +def point_draft_vocab(cfg: dict) -> bool: + """Qwen3.6-35B-A3B: its draft layer is the model file's own, so the token subset is named by --mtp-draft-vocab + (there is no MTP folder to copy it into, as refresh_draft_vocab does): pointed at the chosen subset ("draft_vocab", + cjk by default) in this Strata folder's data/. A file of the user's own (not a shipped subset's name) is kept. + True when the config changed.""" + a = cfg.get("args", []) + if "--mtp-draft-vocab" not in a[:-1]: + return False + i = a.index("--mtp-draft-vocab") + 1 + new = str(ROOT / "data" / DRAFT_VOCABS.get(cfg.get("draft_vocab") or "cjk", "draft_vocab.bin")) + if Path(a[i]).name not in DRAFT_VOCABS.values() or a[i] == new or not Path(new).exists(): + return False + a[i] = new + return True + + def ensure_engine_for(cards, cfg_path: Path, cfg: dict, yes: bool) -> dict: """The installed engine must have code for every card the model starts on: a card added later (--gpus with an older or newer generation, #128) or a new GPU in the PC otherwise stops the start with 'no kernel image'. Such a @@ -4826,8 +5033,10 @@ def sycl_setup(argv) -> int: def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("--family", choices=list(FAMILIES), help="qwen = Qwen3.8-Flash-Next, swift = Swift 1.5") - ap.add_argument("--model", choices=list(MODELS)) + ap.add_argument("--family", choices=list(FAMILIES), help="qwen = Qwen3.8-Flash-Next, swift = Swift 1.5, coder = its " + "Coder, unsloth = Unsloth's 4-bit files, qwen36 = " + "Qwen3.6-35B-A3B (for 16-32 GB of RAM)") + ap.add_argument("--model", choices=model_choices()) ap.add_argument("--context", type=int) ap.add_argument("--rope-scaling", choices=["none", "linear", "yarn"], help="the RoPE extension for a context past the model's trained 262144: linear (position " @@ -4922,7 +5131,9 @@ def main() -> int: "throughput for a model that fits fewer cards; each replica holds its own copy of the model)") ap.add_argument("--kv-streaming", choices=["auto", "on", "off"], default="auto", help="from a 64K context: keep the KV cache in RAM and only the attention's window in VRAM (more " - "experts fit on the GPU); auto: when the RAM has room for it") + "experts fit on the GPU); auto: when the RAM has room for it (Qwen3.6-35B-A3B and its " + "fine-tunes: setup asks, as writing slows down past 32K tokens - or, from 128K with under " + "9 GB of VRAM free, says it is needed)") ap.add_argument("--backend", choices=["cuda", "hip", "sycl"], help="cuda = NVIDIA (default), hip = AMD RX 7900 / 7800 / 7700 XT, RX 9060 XT / 9070 / AI PRO R9700 on " "Linux or Windows (chosen by itself when the PC has no NVIDIA card Strata can use), " @@ -5141,24 +5352,37 @@ def main() -> int: "update it with the NVIDIA App or from https://www.nvidia.com/drivers, restart, and run this again" + ("" if cuda_tk == 12 else f" (or --cuda 12: the experimental CUDA 12 engine runs with driver " f"{CUDA12_MIN_DRIVER} or newer, docs/OLDER_GPUS.md)")) - if gpu["vram_gb"] < 11: + if gpu["vram_gb"] < 11 and not small_family(a.family or ""): # the 35B-A3B families: their note is in step 2 warn("less than 12 GB of " + ("GPU memory (this APU's carve-out + shared memory)" if gpu.get("uma") else "VRAM") - + ": Strata will run, but most experts stay on the CPU and it will be slow") + + ": Strata will run, but most experts stay on the CPU and it will be slow" + + ("" if a.family else " (Qwen3.6-35B-A3B, --family qwen36, is made for 6-12 GB cards)")) ram = ram_gb() cpu, avx2, avx512 = cpu_info() - need = min(d["ram_gb"] for d in MODELS.values()) - low_ok = low_ram_fits("IQ1_M", ram, low_ram_vram(gpu)) and a.low_ram != "off" # the smallest model, mapped + # the RAM floor: the smallest size of the family asked for (Qwen3.6-35B-A3B's own, or the Flash-Next ones), or of + # every family when none was named (Qwen3.6's UD-IQ3_S: a 16-24 GB PC is suggested that family in step 2) + floor_models = [m for m, d in MODELS.items() + if a.family is None or (a.family in d.get("families", ()) if small_family(a.family) + else not small_model(d))] + smallest = min(floor_models, key=lambda m: (MODELS[m]["ram_gb"], MODELS[m]["arena_gb"])) + need = MODELS[smallest]["ram_gb"] + small_name = "the Coder" if smallest == "IQ1_M" else f"{FAMILIES[MODELS[smallest].get('families', ('qwen',))[0]]['title']} " \ + f"{size_of(smallest)}" + experts_gb = "23-50 GB" if smallest == "IQ1_M" else f"{MODELS[smallest]['arena_gb']:.0f}-50 GB" + low_ok = any(low_ram_fits(m, ram, low_ram_vram(gpu)) for m in floor_models if not MODELS[m].get("budget")) \ + and a.low_ram != "off" # the smallest model, mapped if ram < need - 4 and not a.check and not low_ok: - # every model keeps ALL its experts in RAM (23+ GB); VRAM only holds a copy of the most-used ones, so a + # every model keeps ALL its experts in RAM (12+ GB); VRAM only holds a copy of the most-used ones, so a # bigger GPU does not lower this. The owner's rule: a stop by default, a risk the user can take (--model # with --yes, or y) - confirm_risk(f"RAM: {ram:.0f} GB - the smallest model (the Coder) needs about {need} GB: Strata keeps all of " - "the model's experts in RAM (23-50 GB, whatever the GPU), so the OS will page them from disk. " - "Expect it to be very slow, and it may not start at all.", bool(a.model), a.yes, - f"RAM: {ram:.0f} GB - the smallest model (the Coder) needs about {need} GB", - "Strata keeps all of the model's experts in RAM (23-50 GB, whatever the GPU) and the GPU holds a " - "copy of the most-used ones: it needs 32 GB of RAM or more (48 GB for the full model); --model " - "NAME --yes installs one anyway") + confirm_risk(f"RAM: {ram:.0f} GB - the smallest model ({small_name}) needs about {need} GB: Strata keeps all " + f"of the model's experts in RAM ({experts_gb}, whatever the GPU), so the OS will page them from " + "disk. Expect it to be very slow, and it may not start at all.", bool(a.model), a.yes, + f"RAM: {ram:.0f} GB - the smallest model ({small_name}) needs about {need} GB", + f"Strata keeps all of the model's experts in RAM ({experts_gb}, whatever the GPU) and the GPU " + f"holds a copy of the most-used ones: it needs {need} GB of RAM or more" + + (" (48 GB for the full model)" if smallest == "IQ1_M" else + f" ({small_name}; 32 GB for the Coder, 48 GB for the full model)") + + "; --model NAME --yes installs one anyway") warn(f"going on with {ram:.0f} GB of RAM, as you chose") ram_msg = (f"RAM: {ram:.0f} GB" if ram >= need - 4 else f"RAM: {ram:.0f} GB (less than the {need} GB the smallest model needs)" + ("; the GPU's VRAM makes up for it (the low-RAM mode)" if ram < need - 4 and low_ok else "")) @@ -5203,8 +5427,13 @@ def main() -> int: verdict = (f"fits in the low-RAM mode (the GPU holds ~{100 * low_ram_gpu_share(m, low_ram_vram(gpu)):.0f}% " "of its experts, " + ("the rest stays in RAM)" if low_ram_resident(m, ram, low_ram_vram(gpu)) else "the rest is read from the SSD as needed)")) + if hip and small_model(d) and not verdict.startswith("does not fit"): + verdict += " - untested on AMD (its batched prompt path is NVIDIA-only)" any_fits = any_fits or not verdict.startswith("does not fit") say(f" {m:8s} needs ~{d['ram_gb']} GB RAM: {verdict}") + if qwen36_recommended(ram) and any_fits: + say(f"\nNo Qwen3.8-Flash-Next size fits {ram:.0f} GB of RAM well: setup suggests Qwen3.6-35B-A3B " + f"{size_of(qwen36_size(ram))} (--family qwen36).") if not any_fits: # #977: every size says "does not fit", so the verdict says so too (and the exit code, for scripts). It only # reports: --model NAME --yes still installs one anyway. @@ -5221,23 +5450,31 @@ def main() -> int: family = a.family else: rec_fam = STRIX_HALO_FAMILY if strix_halo_recommends(gpu, ram) and STRIX_HALO_FAMILY in fams else fams[0] + if rec_fam == fams[0] and qwen36_recommended(ram) and "qwen36" in fams: + rec_fam = "qwen36" # no Flash-Next size fits this RAM: the smaller model for i, f in enumerate(fams, 1): d = FAMILIES[f] say(f" {i}) {d['title']:20s} {d['by']} - {d['about']}" + (" [experimental]" if d.get("experimental") else "") - + (f" (recommended for Strix Halo: {STRIX_HALO_MODEL})" if f == rec_fam and rec_fam != fams[0] else "")) + + ((f" (recommended for {ram:.0f} GB of RAM: {size_of(qwen36_size(ram))})" if f == "qwen36" else + f" (recommended for Strix Halo: {STRIX_HALO_MODEL})") if f == rec_fam and rec_fam != fams[0] + else "")) family = fams[int(ask("Which model?", [str(i) for i in range(1, len(fams) + 1)], str(fams.index(rec_fam) + 1), a.yes)) - 1] fam = FAMILIES[family] ok(f"model: {fam['title']}") + if small_family(family) and 0 < gpu["vram_gb"] < SMALL_MODEL_VRAM_GB: + warn(small_model_vram_note(gpu["vram_gb"])) if fam.get("license"): say(f" Its license: {fam['license']}") say() names = [m for m in MODELS if family in MODELS[m].get("families", ("qwen", "swift"))] names.sort(key=lambda m: bool(MODELS[m].get("experimental"))) # an experimental size last, never the default + if a.model: # a size name: this family's file of that size (model_key) + a.model = model_key(family, a.model) or a.model if a.model and a.model not in names: # #444: say which family has that size, and (with --gguf-dir) which files Strata can run at all - elsewhere_fams = [f for f in FAMILIES if f in MODELS[a.model].get("families", ("qwen", "swift"))] - fail(f"{fam['title']} has no {a.model} model file", "choose one of: " + ", ".join(names) + elsewhere_fams = [f for f in FAMILIES if model_key(f, a.model)] + fail(f"{fam['title']} has no {a.model} model file", "choose one of: " + ", ".join(map(size_of, names)) + (f" (or {a.model}: " + ", ".join(f"--family {f} --model {a.model}" for f in elsewhere_fams) + ")" if elsewhere_fams else "") + (f".\n {SUPPORTED_GGUFS}" if a.gguf_dir else "")) @@ -5245,16 +5482,47 @@ def main() -> int: d = MODELS[m] fit = "" if ram >= d["ram_gb"] else f" <- needs {d['ram_gb']} GB RAM, you have {ram:.0f}" if d.get("budget"): - say(f" {i}) {m} {d['about']}; download {d['download_gb']:.0f} GB, keeps ~" + say(f" {i}) {size_of(m)} {d['about']}; download {d['download_gb']:.0f} GB, keeps ~" f"{resident_budget_gib(m, ram)} GB of its {d['arena_gb']:.0f} GB of experts in RAM{fit}") continue if low_ram_needed(m, ram) and low_ram_fits(m, ram, low_ram_vram(gpu)) and a.low_ram != "off": fit = (f" <- fits in the low-RAM mode (the GPU holds ~{100 * low_ram_gpu_share(m, low_ram_vram(gpu)):.0f}%, " + ("the rest in RAM)" if low_ram_resident(m, ram, low_ram_vram(gpu)) else "the rest from the SSD)")) - say(f" {i}) {m:8s} {d['about']}; download {d['download_gb']:.0f} GB, uses ~{d['arena_gb']:.0f} GB of RAM{fit}") + say(f" {i}) {size_of(m):8s} {d['about']}; download {d['download_gb']:.0f} GB, uses ~{d['arena_gb']:.0f} GB of " + f"RAM{fit}") rec = str(names.index("IQ3_XXS") + 1) if ram >= 60 and "IQ3_XXS" in names else "1" + if small_family(family): # the biggest size whose experts fit the RAM (small_size) + rec = str(names.index(small_size(family, ram, gpu["vram_gb"])) + 1) model = a.model or names[int(ask("Which size?", [str(i) for i in range(1, len(names) + 1)], rec, a.yes)) - 1] budget, q4_split = None, False + if hip and fam.get("nvidia_only"): + # Qwen3.6-35B-A3B: the HIP engine compiles, but its batched prompt path is NVIDIA-only (sm_75+) - on AMD the + # prompt is read through the slower decode windows - and it has not been run on AMD. Checked before the + # download; asked (default no), --family/--model with --yes goes on (the owner's rule) + confirm_risk(f"{fam['title']} has not been run on AMD cards yet: its batched prompt path is NVIDIA-only, so on " + f"{gpu_name(gpu)} prompts are read through the slower decode path, and it may not work at all", + bool(a.model or a.family), a.yes, f"{fam['title']} is untested on AMD", + f"use an NVIDIA card, or --family {family} --model {size_of(model)} --yes to try it on AMD " + "anyway", " Try it anyway?") + warn(f"installing {fam['title']} {size_of(model)} on an AMD card, as you chose (untested: please report how " + "it runs)") + if not hip and fam.get("nvidia_only") and str(gpu.get("arch", "")).isdigit() and \ + int(gpu["arch"]) < CUDA13_MIN_ARCH: + # Pascal / Volta (the experimental CUDA 12 engine): the batched prompt path needs sm_75, so the engine reads + # every prompt through the decode windows, and the model has not been run on these cards. Before the download + confirm_risk(f"{fam['title']} has not been run on {gpu_name(gpu)} (sm_{gpu['arch']}): its batched prompt path " + "needs an RTX 20 or newer (sm_75), so prompts are read through the slower decode path (on an RTX 4070 Ti " + "made to take it: ~90 tokens/s, an 8K-token prompt in 1.5 minutes; slower on this card)", bool(a.model or a.family), a.yes, + f"{fam['title']} is untested below sm_75", + f"use Qwen3.8-Flash-Next on this card, or --family {family} --model {size_of(model)} --yes to try " + "it anyway", " Try it anyway?") + warn(f"installing {fam['title']} {size_of(model)} on {gpu_name(gpu)}, as you chose (untested: please report " + "how it runs)") + if fam.get("one_gpu") and multi: + # like UD-Q4_K_XL's RAM budget: the engine has no layer split for this model yet, so one card runs it + warn(f"{fam['title']} runs on one GPU for now (the engine has no layer split for it yet): using " + f"{gpu_name(gpu)} only" + (" - --gpus is not used" if a.gpus else "")) + multi, sel, chosen = [], [gpu["index"]], [gpu] if MODELS[model].get("budget"): # Unsloth's UD-Q4_K_XL: a RAM budget of experts, the rest from the GGUF on the SSD - not the low-RAM mode (no # experts.bin: it would be another 77 GB on the disk), and one GPU (the budget mode had no layer split) unless @@ -5298,14 +5566,16 @@ def main() -> int: # before RESIDENT_SPLIT_ENGINE, #642) if not low_ram and budget is None and ram < MODELS[model]["ram_gb"] - 4: confirm_paging(model, ram, a.low_ram, a.yes, bool(a.model)) - ok(f"size: {model}") - tag = fam["tag"] + model # names of the pack, config and start script + ok(f"size: {size_of(model)}") + tag = fam["tag"] + size_of(model) # names of the pack, config and start script small = min(x["vram_gb"] for x in chosen) # each card keeps its layers' KV of the whole context rec_ctx = 32768 if small < 14 else 65536 if small < 20 else 131072 if budget is not None: # UD-Q4_K_XL: every GB of KV is a GB fewer of cached experts rec_ctx = 8192 if small < 14 else 32768 # #406: the RAM rule is part of the recommendation (the smaller of the two), no longer a cap over the user's choice rec_ctx = min(rec_ctx, ram_ctx(model, ram, low_ram)) + if small_family(family) and small < SMALL_MODEL_VRAM_GB: + rec_ctx = min(rec_ctx, 8192) # a 4 GB card: the 8K context is what can start (measured) if a.context: ctx = a.context else: @@ -5370,7 +5640,8 @@ def main() -> int: else: say() say(" Images: the model can also read pictures (screenshots, photos, scanned pages). This adds a 0.9 GB") - say(" download and keeps ~1.4 GB of VRAM free for the image encoder, so text is a few % slower.") + room = (VISION_GPU_QWEN36_MIB + VISION["gpu"]["reserve_mib"]) / 1024 if small_family(family) else 1.4 + say(f" download and keeps ~{room:.1f} GB of VRAM free for the image encoder, so text is a few % slower.") vision = "gpu" if ask("Do you want images?", ["y", "n"], "n", a.yes) == "y" else "none" ok("images: " + {"none": "off", "gpu": "on", "cpu": "on (encoder on the CPU)"}[vision]) if vision != "none" and MODELS[model].get("vision_untested"): @@ -5465,7 +5736,8 @@ def main() -> int: pack_bin = (pack_now / "experts.bin").exists() and (pack_now / "index.txt").exists() mtp_have = find_in(roots, "mtp/rt/experts.bin") is not None q2_avx = model == "Q2_0" and avx512 and family == "qwen" - need = to_fetch + (2 if mtp_have else 8) + \ + own_mtp = bool(fam.get("own_mtp")) # Qwen3.6: the draft layer is in the model file, nothing fetched + need = to_fetch + (1 if own_mtp else 2 if mtp_have else 8) + \ (40 if q2_avx and not pack_bin else 0) + (1 if vision != "none" else 0) + \ (MODELS[model]["arena_gb"] + 1 if low_ram and not q2_avx and not pack_bin else 0) if free_gb(models_dir) < need: @@ -5493,6 +5765,16 @@ def main() -> int: else: eng = None if a.build or hip else get_prebuilt(a.prebuilt, gpu, vision, **({"toolkit": 12} if cuda_tk == 12 else {})) + if eng is not None and own_mtp and not engine_runs_small(eng): + # Qwen3.6 / Ornith: no ready-made engine runs them yet (the published ones predate qwen35moe), so this PC + # compiles the engine from this checkout, as it does for a card the ready-made one has no code for; that + # engine runs every family. A ready-made engine that has it is used as is. + ver = json.loads((eng / "BUILD.json").read_text(encoding="utf-8")).get("version", "?") + if hip and WIN: + fail(f"the ready-made AMD engine ({ver}) does not run {fam['title']} yet", + "compiling it on Windows: tools\\hip\\build_windows.bat (docs/AMD_HIP.md), or choose another model") + say(f" The ready-made engine ({ver}) does not run {fam['title']} yet: compiling the engine on this PC instead") + eng = None if eng is not None and not hip and json.loads((eng / "BUILD.json").read_text(encoding="utf-8")).get("source") != "local": pip_cuda_libs(cuda_tk) if vision != "none" and not (eng / VEXE).exists(): @@ -5509,13 +5791,17 @@ def main() -> int: lib_dirs = meta.get("lib_dirs") or meta.get("cuda_dirs") or cuda_lib_dirs(cuda_tk) engine_ver = tuple(int(x) for x in str(meta.get("version", "0")).split(".")[:4] if x.isdigit()) need_engine = MODELS[model].get("engine", UNSLOTH_ENGINE) - if budget is not None and engine_ver < need_engine: # checked before the 94-111 GB download + if budget is not None and engine_ver < need_engine: # checked before the 15-111 GB download fail(f"{model} needs engine {'.'.join(map(str, need_engine))} or newer; this one is {meta.get('version')}", "update Strata (or compile the engine with --build) and run setup again") + if own_mtp and not engine_runs_small(eng): # e.g. --build of an older checkout; checked before the download + fail(f"{fam['title']} {size_of(model)} needs an engine with the qwen35moe support (Strata " + f"{'.'.join(map(str, QWEN36_ENGINE))} source or newer); this one ({meta.get('version')}) has none", + "update Strata and run setup again with --build") ok(f"engine: {eng / EXE}") # ---- 5. the model files - step(5, f"downloading {fam['title']} {model}") + step(5, f"downloading {fam['title']} {size_of(model)}") if not a.gguf_dir: missing = [s.name for s in shards if not (s.exists() and done(s))] if missing: # #495: files downloaded by hand go here, or --gguf-dir @@ -5543,14 +5829,14 @@ def main() -> int: continue except OSError: pass - download(fam["hf"].format(q=model) + s.name, s) + download(fam["hf"].format(q=size_of(model)) + s.name, s) check_shards(shards) for s in shards: # the Unsloth files: pinned sizes and SHA-256 if s.name in fam.get("sha256", {}): verify_sha256(s, *fam["sha256"][s.name]) ok("model files present") - mmproj = Path(a.models_dir) / fam["mmproj"] - if not mmproj.exists(): + mmproj = Path(a.models_dir) / (fam["mmproj"] or "none") + if fam["mmproj"] and not mmproj.exists(): mmproj = find_in(roots, f"models/{fam['mmproj']}") or mmproj if vision != "none": if not mmproj.exists() and a.gguf_dir and (Path(a.gguf_dir) / fam["mmproj"]).exists(): @@ -5588,23 +5874,30 @@ def main() -> int: ok(f"model prepared: {pack}") mtp = (find_in(roots, "mtp/rt/experts.bin") or data / "mtp/rt/experts.bin").parent.parent rt = mtp / "rt" - corrupt = (rt / "experts.bin").exists() and mtp_corrupt(mtp, env) - if corrupt: - warn("some MTP tensors are not the checkpoint's (a download mirror that ignored range requests, #327): " - "fetching them again and rebuilding the draft layer") - if corrupt or not (rt / "experts.bin").exists(): - say(" The MTP draft layer (speculative decoding, ~2x faster output) comes from the original Qwen checkpoint:") - say(" only its ~5 GB of MTP tensors are downloaded.") - run([sys.executable, str(ROOT / "tools" / "mtp_fetch.py"), "fetch", "--out", str(mtp)], - env={**env, "STRATA_SOURCE": model_source()}) - run([sys.executable, str(ROOT / "tools" / "mtp_pack.py"), "--src", str(mtp), "--experts", "q2_0", - "--out", str(mtp / "mtp-q2_0.gguf")], env=env) - run([sys.executable, str(ROOT / "tools" / "mtp_rt.py"), "--gguf", str(mtp / "mtp-q2_0.gguf"), "--out", str(rt)], - env=env) # a setup run again without --draft-vocab keeps the subset this model's config chose before (cyrillic, fr, en) draft_vocab = a.draft_vocab or saved_draft_vocab(ROOT / f"strata-{tag.lower()}.json") - refresh_draft_vocab(rt, draft_vocab or "cjk") - ok(f"MTP draft layer: {rt}") + if own_mtp: + # Qwen3.6-35B-A3B: the draft layer is the model file's own MTP block (blk.40, "nextn") - nothing is fetched + # or converted; --mtp names the model file and --mtp-draft-vocab the shipped token subset (the same + # tokenizer as Flash-Next's, so data/'s subsets apply) + rt = shards[0] + ok(f"MTP draft layer: the model file's own ({shards[0].name})") + else: + corrupt = (rt / "experts.bin").exists() and mtp_corrupt(mtp, env) + if corrupt: + warn("some MTP tensors are not the checkpoint's (a download mirror that ignored range requests, #327): " + "fetching them again and rebuilding the draft layer") + if corrupt or not (rt / "experts.bin").exists(): + say(" The MTP draft layer (speculative decoding, ~2x faster output) comes from the original Qwen checkpoint:") + say(" only its ~5 GB of MTP tensors are downloaded.") + run([sys.executable, str(ROOT / "tools" / "mtp_fetch.py"), "fetch", "--out", str(mtp)], + env={**env, "STRATA_SOURCE": model_source()}) + run([sys.executable, str(ROOT / "tools" / "mtp_pack.py"), "--src", str(mtp), "--experts", "q2_0", + "--out", str(mtp / "mtp-q2_0.gguf")], env=env) + run([sys.executable, str(ROOT / "tools" / "mtp_rt.py"), "--gguf", str(mtp / "mtp-q2_0.gguf"), "--out", str(rt)], + env=env) + refresh_draft_vocab(rt, draft_vocab or "cjk") + ok(f"MTP draft layer: {rt}") for line in draft_vocab_note(gpu.get("vram_gb", 0.0), draft_vocab): # #474: a recommendation, nothing changes say(" " + line) @@ -5612,13 +5905,20 @@ def main() -> int: step(7, "writing the start script") sys.path.insert(0, str(ROOT / "tools")) from gguf_reader import GGUFFile # the PLE table's shard: shard 2 (original) or 1 (Swift) - ple = next((s for s in shards if any(t.name == "per_layer_token_embd.weight" for t in GGUFFile(s).tensors)), None) - if ple is None: - fail("the model has no per_layer_token_embd tensor (is this a Qwen3.8-Flash-Next GGUF?)") + ple = None # Qwen3.6-35B-A3B has no PLE table: no --ple-gguf + if fam.get("ple", True): + ple = next((s for s in shards if any(t.name == "per_layer_token_embd.weight" for t in GGUFFile(s).tensors)), + None) + if ple is None: + fail("the model has no per_layer_token_embd tensor (is this a Qwen3.8-Flash-Next GGUF?)") # (a 4-shard file: the engine finds the PLE table's shard itself from shard 1, the measured setup) - args = ["--pack", str(pack), "--native", str(shards[0]), *(["--ple-gguf", str(ple)] if len(shards) <= 2 else []), + mtp_args = ["--mtp", str(rt)] + if own_mtp: # the model file's own draft layer, the chosen token subset + mtp_args += ["--mtp-draft-vocab", str(ROOT / "data" / DRAFT_VOCABS.get(draft_vocab or "cjk", "draft_vocab.bin"))] + args = ["--pack", str(pack), "--native", str(shards[0]), + *(["--ple-gguf", str(ple)] if ple is not None and len(shards) <= 2 else []), "--expert-profile", str(ROOT / "data" / fam.get("profile", "expert-profile.bin")), "--expert-cache", "auto", - "--prefill", "auto", "--spec", "4", "--spec-min-p", "0.5", "--mtp", str(rt), + "--prefill", "auto", "--spec", "4", "--spec-min-p", "0.5", *mtp_args, "--max-context", str(ctx)] if scaling is not None: # the resolved config: explicit flags as given, or the automatic yarn+factor args += ["--rope-scaling", scaling, "--rope-scale", f"{rope_scale:g}"] @@ -5630,7 +5930,7 @@ def main() -> int: "the OS file cache (run setup again after the next engine update)") if low_ram: # the experts from the pack's experts.bin: the ones the GPU does not hold copied into RAM, or mapped args += ["--resident-experts" if resident else "--mmap-experts"] - disk = None if is_wsl() else rotational_disk(ple) # #605 (WSL's virtual disk says rotational) + disk = None if is_wsl() or ple is None else rotational_disk(ple) # #605 (WSL's virtual disk says rotational) if disk: tensor = next((t for t in GGUFFile(ple).tensors if t.name == "per_layer_token_embd.weight"), None) size = getattr(tensor, "expected_bytes", lambda: None)() @@ -5650,20 +5950,70 @@ def main() -> int: kv_ram_gb = kv_streaming_ram_gb(ctx, kv) # the branches below are kv_streaming_wanted, with its messages # --kv-streaming on|off overrides the RAM test (the owner's rule); WSL stays off - it cannot stream. stream_fits = ram >= MODELS[model]["ram_gb"] + kv_ram_gb + 1 + # The 35B-A3B models (qwen35moe) attend to every cell, not to a window: past the 32K cells kept in VRAM, every + # token reads the rest of the KV cache from RAM over PCIe. Measured with KAT-Coder i-quality-v2 (8-bit KV) on an + # RTX 4070 Ti 12 GB with no desktop, whole (11.6 GB free) and with 4 GB of it held (7.9 GB free, an 8 GB card): + # - a 52.7K-token prompt at 64K context wrote 59.6 / 47.0 tokens/s with the KV cache in VRAM, 27.0 / 23.3 streamed; + # - a 2K-token prompt was as fast or faster streamed (more expert slots: 3,076 vs 2,930 at 64K, 2,935 vs 1,905 at + # 256K on the whole card); + # - with 7.9 GB free a 128K or 256K context fails at the first prompt unless streamed; with 10 GB free (the whole + # card beside a 2.3 GB desktop) both run. + # So: from 128K up with under 9 GB free (the gap between those two measurements) streaming is required, and setup + # says so; otherwise it asks, and recommends keeping the KV cache in VRAM. Free VRAM, not the card's size: a desktop + # on the card takes its share. + stream_choice, stream_how = a.kv_streaming, "" + if small_family(family) and ctx >= 65536 and not is_wsl(): + free = gpu_free_gb(gpu["index"]) if gpu.get("vendor", "nvidia") == "nvidia" and "index" in gpu else None + card = f"the {gpu.get('vram_gb', 0.0):.0f} GB card" + (f" has {free:.1f} GB free now" if free is not None else "") + if free is not None and vision == "gpu": # the encoder starts before the engine and takes its share first + free -= (VISION_GPU_QWEN36_MIB + VISION["gpu"]["reserve_mib"]) / 1024 + card += f", {free:.1f} GB beside the image encoder" + later = (f" To change it later: \"--kv-resident\", \"32768\" in strata-{tag.lower()}.json's args, or run", + " setup again with --kv-streaming on|off.") + required = free is not None and free < 9.0 and ctx >= 131072 + if required and a.kv_streaming == "off": + warn(f"--kv-streaming off: {card}, and a {ctx // 1024}K context of {fam['title']} needs KV streaming " + "below about 9 GB free (measured with 7.9 GB free: the first prompt fails without it). Kept as you " + "chose; a smaller --context or --kv-streaming on avoids it") + elif required: + say() + say(f" KV streaming is needed here: {card}, and a {ctx // 1024}K context's KV cache does not fit") + say(" beside the experts (measured with 7.9 GB free: without streaming the first prompt fails). It keeps the") + say(" KV cache in RAM and 32K tokens of it in VRAM. Past 32K tokens of a conversation writing slows down") + say(" (measured at 52K tokens: about half the speed); a smaller context (--context 65536) avoids that.") + for line in later: + say(line) + stream_choice, stream_how = "on", "required" + elif a.kv_streaming == "auto" and stream_fits: + freed_gb = (ctx - 32768) * 10 * KV_CELL_BYTES.get(kv, 1056) / 1e9 # its 10 attention layers + say() + say(f" KV streaming: keeps the {ctx // 1024}K context's KV cache in RAM and only 32K tokens of it in VRAM,") + say(f" which frees ~{freed_gb:.1f} GB of VRAM for experts ({card}). Short conversations get a") + say(f" little faster, but {fam['title']} reads all of its context for every token, so past 32K tokens") + say(" writing slows down (measured at 52K tokens: about half the speed). Recommended: no.") + for line in later: + say(line) + stream_choice = "on" if ask("Turn on KV streaming?", ["y", "n"], "n", a.yes) == "y" else "off" + stream_how = "asked" if is_wsl() and ctx >= 65536: ok("WSL: KV streaming off (the driver pins only about 1 GB of RAM); the KV cache stays in VRAM") if a.kv_streaming == "on": warn("--kv-streaming on: WSL cannot stream the KV cache (its RAM copy must be pinned, and the driver pins " "only about 1 GB there): off") - elif ctx >= 65536 and a.kv_streaming == "off": - ok("KV streaming off, as you chose (--kv-streaming off): the KV cache stays in VRAM") - elif ctx >= 65536 and (stream_fits or a.kv_streaming == "on"): + elif ctx >= 65536 and stream_choice == "off": + ok("KV streaming off, as you chose" + ("" if stream_how == "asked" else " (--kv-streaming off)") + + ": the KV cache stays in VRAM") + elif ctx >= 65536 and (stream_fits or stream_choice == "on"): args += ["--kv-resident", "32768"] ok(f"KV streaming on: the context's KV cache lives in RAM ({kv_ram_gb:.1f} GB), more experts fit in VRAM") + if small_family(family): + ok(" writing slows down past 32K tokens of context (that part of the KV cache is read from RAM)" + + ("; needed for this context with this much free VRAM" if stream_how == "required" else "")) if not stream_fits: warn(f"KV streaming needs ~{kv_ram_gb:.1f} GB of RAM beside the ~{MODELS[model]['ram_gb']} GB {model} " - f"uses; this PC has {ram:.0f}. Kept as you chose (--kv-streaming on): it may page or run out of RAM " - "under load") + f"uses; this PC has {ram:.0f}. " + ("On, as this context needs it (a smaller --context avoids it)" + if stream_how == "required" else "Kept as you chose (--kv-streaming on)") + + ": it may page or run out of RAM under load") if q4_split: # #498: no budget to take it out of pass elif budget is not None and a.resident_budget_gib is None: # its RAM comes out of the experts' budget @@ -5731,7 +6081,7 @@ def main() -> int: args += ["--control-vector-scaled", f"{esp}:1.0", "--control-vector-layer-range", "4", "44", "--cvec-mode", "project", "--cvec-dir", "per-layer"] cfg = {"exe": str(eng / EXE), "args": args, "cwd": str(ROOT), "tokenizer": str(pack / "tokenizer"), - "model_name": f"{fam['name']}-{model.lower()}", "log": str(ROOT / f"strata-{tag.lower()}.log"), + "model_name": f"{fam['name']}-{size_of(model).lower()}", "log": str(ROOT / f"strata-{tag.lower()}.log"), "lib_dirs": lib_dirs, "port": port} if cuda_tk == 12: # the experimental CUDA 12 engine (engine-cuda12/) cfg["cuda"] = 12 @@ -5809,7 +6159,8 @@ def main() -> int: cfg["args"] = recommend_pool_workers(cfg["args"]) for line in two_socket_note(cpu_sockets()): # bench tips: text only, the config is not touched say(" " + line) - for line in bench_tips(cfg["args"], cfg.get("env"), ram, MODELS[model]["ram_gb"], gpu.get("vram_gb", 0.0), vision, WIN): + for line in bench_tips(cfg["args"], cfg.get("env"), ram, MODELS[model]["ram_gb"], gpu.get("vram_gb", 0.0), vision, WIN, + small_family(family)): say(" " + line) write_setup_config(cfg_path, cfg, adopted if adopted is not None and adopted.name == cfg_path.name else None) script = write_run_script(tag, cfg_path, port, cfg.get("open_browser") is not False) diff --git a/src/core/conversation_snapshot.cpp b/src/core/conversation_snapshot.cpp index 1468dc65207..f40412eeb38 100644 --- a/src/core/conversation_snapshot.cpp +++ b/src/core/conversation_snapshot.cpp @@ -34,7 +34,7 @@ bool layout(const QsaState& st, const ModelGeometry& g, int64_t upto, bool index } const auto s = strata::kernels::qsa_real_shapes(); const bool int8_keys = st.kv_int8 || st.kv_hybrid; - if (g.n_head_kv <= 0 || g.head_dim <= 0 || g.head_dim > INT32_MAX || g.idx_key_dim <= 0 || + if (g.n_head_kv <= 0 || g.head_dim <= 0 || g.head_dim > INT32_MAX || (g.has_indexer && g.idx_key_dim <= 0) || (st.kv_q4 && g.head_dim % 32) || (int8_keys && !st.kv_q4 && g.head_dim % 64)) { error = "conversation snapshot: invalid K/V geometry"; return false; diff --git a/src/core/conversation_state.cpp b/src/core/conversation_state.cpp index 4fbcf88bc09..aae22c09c6a 100644 --- a/src/core/conversation_state.cpp +++ b/src/core/conversation_state.cpp @@ -107,8 +107,12 @@ bool metadata_bytes(const ConversationCheckpoint& c, size_t& total) { bool conversation_state_sizes(const ModelGeometry& g, ConversationStateSizes& z, std::string& error) { z = {}; const auto key = geometry_key(g); - for (size_t i = 0; i < key.size(); ++i) - if (key[i] < 0 || (i != 1 && key[i] == 0)) return fail(error, "invalid model geometry"); + for (size_t i = 0; i < key.size(); ++i) { + // zero is a real value only for what a model may not have: the indexer's widths (12, 13) without an indexer, + // the hyper-connections' low rank (15) with one residual stream + const bool may_be_zero = ((i == 12 || i == 13) && !g.has_indexer) || (i == 15 && !g.has_hc()); + if (key[i] < 0 || (i != 1 && key[i] == 0 && !may_be_zero)) return fail(error, "invalid model geometry"); + } size_t recurrence = 0, convolution = 0; if (!product(recurrence, {(uint64_t) g.ssm_state_size, (uint64_t) g.ssm_v_heads, (uint64_t) g.ssm_state_size}) || !product(convolution, {(uint64_t) g.ssm_conv_channels, (uint64_t) (g.ssm_d_conv - 1)}) || diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index ad28f56f840..adba401ab00 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -3647,6 +3647,7 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 if ((int64_t) d.act_multi.size() < n_tok) d.act_multi.resize((size_t) MAXT); const ExpertLayout& lay = expert_layout(); const bool native = lay.native; + const int64_t HN = native ? lay.n_embd : H; // the rows' width: Flash-Next's H, or Qwen3.6's 2048 if (native && d.nact_multi.size() < (size_t) MAXT * kNativeActBytes) d.nact_multi.resize((size_t) MAXT * kNativeActBytes); if (d.job_of.size() != (size_t) d.n_expert) d.job_of.assign((size_t) d.n_expert, (int16_t) -1); if (d.jobs_multi.size() < (size_t) (n_tok * k)) d.jobs_multi.resize((size_t) (MAXT * k)); @@ -3832,9 +3833,9 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 for (int64_t t = 0; t < n_tok; ++t) { if (ep && std::all_of(kind + t * k, kind + (t + 1) * k, [](int32_t v) { return v >= 0; })) continue; if (native && strata::kernels::cpu::q2_native_kernels(lay.fmt[(size_t) d.layers].gu_type)) // a native Q2_0 pack: the Q2_0 kernels' activations - act_quant_any(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]); + act_quant_any(x_f + (size_t) t * HN, (int) HN, d.act_multi[(size_t) t]); else if (native) - native_quant_act(lay.fmt[(size_t) d.layers], x_f + (size_t) t * H, d.nact_multi.data() + (size_t) t * kNativeActBytes); + native_quant_act(lay.fmt[(size_t) d.layers], x_f + (size_t) t * HN, d.nact_multi.data() + (size_t) t * kNativeActBytes); else act_quant_q8_1(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]); } @@ -3854,7 +3855,7 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 for (int64_t j = 0; j < k; ++j) { const int64_t i = t * k + j; const int64_t e = ids[i]; - float* row = out + (size_t) i * H; + float* row = out + (size_t) i * HN; if (e < 0 || e >= d.n_expert) { d.failed = true; d.fail = "a routed expert id is out of range"; @@ -3863,7 +3864,7 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 return; } if (kind[i] == -2) { // STRATA_ROUTE_TAIL_SKIP: a skipped tail expert contributes zero - std::memset(row, 0, (size_t) H * sizeof(float)); + std::memset(row, 0, (size_t) HN * sizeof(float)); ++d.tail_skipped; continue; } @@ -3875,7 +3876,7 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 // dec_batch: copy_rows_from_mapped_kernel already zeroes kind 0/1 rows on the GPU if (!((kind[i] == 2 && d.peer != nullptr && d.peer->launched_direct()) || (gpu_zeroes_hits && (kind[i] == 0 || kind[i] == 1)))) - std::memset(row, 0, (size_t) H * sizeof(float)); + std::memset(row, 0, (size_t) HN * sizeof(float)); continue; } ++d.cache_refused; diff --git a/src/core/layer.cpp b/src/core/layer.cpp index 3132871c63d..6d33744d944 100644 --- a/src/core/layer.cpp +++ b/src/core/layer.cpp @@ -504,7 +504,7 @@ if (!moe_route(tables, g, layer, k, b, x, stream, err, db)) return false; ret namespace {using strata::kernels::QsaIndexerBuffers;using strata::kernels::QsaShapes; /// The geometry the QSA kernels want, from the one place that defines it. `ModelGeometry` carries the widths /// the LAYOUT needs; `QsaShapes` adds `n_rot`, `idx_block` and `idx_top_k`, which are kernel contracts. -QsaShapes qsa_shapes(const ModelGeometry& g) { QsaShapes s = strata::kernels::qsa_real_shapes(); s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; return s;} +QsaShapes qsa_shapes(const ModelGeometry& g) { QsaShapes s = strata::kernels::qsa_real_shapes(); s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; return s;} uint64_t align_up16(uint64_t n) { return (n + 15) & ~15ull; } /// One cursor over an arena, so every region is 16-byte aligned without a list of hand-added offsets. struct Cursor { uint8_t* p; uint64_t used = 0; template T* take(uint64_t count) { T* r = (T*) (p + used); used = align_up16(used + count * sizeof(T)); return r; } uint8_t* take_bytes(uint64_t n) { uint8_t* r = p + used; used = align_up16(used + n); return r; }}; @@ -898,6 +898,12 @@ strata::kernels::QsaAttnPools qsa_attn_pools(const QsaState& st) { else if (st.kv_q4) { pools.k_q4 = st.k_q4; pools.v_q4 = st.v_q4; } else if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; } else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; } + if (st.kv_mode == 1) { // streamed: a selected block the resolve could not make resident is read from RAM in place + pools.host_k_pool = st.host.k_pool; pools.host_v_pool = st.host.v_pool; + pools.host_k_q = st.host.k_q; pools.host_v_q = st.host.v_q; + pools.host_k_scale = st.host.k_scale; pools.host_v_scale = st.host.v_scale; + pools.host_k_q4 = st.host.k_q4; pools.host_v_q4 = st.host.v_q4; + } return pools; } @@ -905,7 +911,7 @@ void qsa_kv_resolve(const QsaState& st, const ModelGeometry& g, const int32_t* i int64_t cap, void* stream) { if (st.kv_mode != 1) return; strata::kernels::kv_stream_resolve(st.map, qsa_attn_pools(st), st.host, qsa_kv_format(st), ids, steps, n_q, cap, - qsa_shapes(g), stream); + qsa_shapes(g), stream, !g.has_indexer); // no indexer: dense ids 0..width-1 } uint64_t qsa_buffers_bytes(const ModelGeometry& g, int64_t max_cells) { const QsaShapes s = qsa_shapes(g); const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); uint64_t n = 0; n += (uint64_t) q8k_bytes(g.n_embd); n += (uint64_t) (g.n_embd / 32) * 34; diff --git a/src/core/layout.cpp b/src/core/layout.cpp index 6d1b970bdae..e2b8a9318ff 100644 --- a/src/core/layout.cpp +++ b/src/core/layout.cpp @@ -1,6 +1,8 @@ // src/core/layout.cpp - the shape checks. See the header for why this exists. #include "strata/core/layout.hpp" +#include "strata/artifact/gguf_reader.hpp" + #include namespace strata::core { @@ -134,6 +136,116 @@ bool check_one(const WeightTable& t, const ModelGeometry& g, int64_t layer, std: } // namespace +const char* arch_name(ModelArch arch) { + return arch == ModelArch::kQwen35Moe ? "qwen35moe" : "qwen4exp"; +} + +bool geometry_from_gguf(const GgufFile& f, ModelGeometry& g, std::string& err) { + const MetaValue* arch = f.get("general.architecture"); + if (arch == nullptr) { + err = "the model file has no general.architecture"; + return false; + } + if (arch->s == "qwen4exp") { + // Flash-Next: the defaults, with the MoE shape from the file as before (a pruned Coder ships 256 experts) + ModelGeometry q; + if (const MetaValue* v = f.get("qwen4exp.expert_count")) q.n_expert = (int64_t) v->u; + if (const MetaValue* v = f.get("qwen4exp.expert_used_count")) q.n_expert_used = (int64_t) v->u; + g = q; + return true; + } + if (arch->s != "qwen35moe") { + err = "the model's architecture is '" + arch->s + "'; this engine runs 'qwen4exp' (Qwen3.8-Flash-Next) and " + "'qwen35moe' (Qwen3.6-35B-A3B)"; + return false; + } + std::string missing; + auto key = [&](const char* name, int64_t& out, bool required = true) { + const MetaValue* v = f.get(std::string("qwen35moe.") + name); + if (v == nullptr || !v->is_num()) { + if (required && missing.empty()) missing = std::string("qwen35moe.") + name; + return; + } + out = (int64_t) v->u; + }; + ModelGeometry q; + q.arch = ModelArch::kQwen35Moe; + int64_t block_count = 0, head_dim = 0, value_len = 0, n_rot = 0, n_ff_shexp = 0; + key("block_count", block_count); + key("nextn_predict_layers", q.n_mtp_layers, false); + key("embedding_length", q.n_embd); + key("full_attention_interval", q.qsa_interval); + key("attention.head_count", q.n_head); + key("attention.head_count_kv", q.n_head_kv); + key("attention.key_length", head_dim); + key("attention.value_length", value_len); + key("rope.dimension_count", n_rot); + key("ssm.state_size", q.ssm_state_size); + key("ssm.group_count", q.ssm_k_heads); + key("ssm.time_step_rank", q.ssm_v_heads); + key("ssm.conv_kernel", q.ssm_d_conv); + key("ssm.inner_size", q.ssm_value_dim); + key("expert_count", q.n_expert); + key("expert_used_count", q.n_expert_used); + key("expert_feed_forward_length", q.n_ff); + key("expert_shared_feed_forward_length", n_ff_shexp); + if (!missing.empty()) { + err = "the model file has no " + missing; + return false; + } + q.n_layers = block_count - q.n_mtp_layers; + q.ssm_conv_channels = 2 * q.ssm_state_size * q.ssm_k_heads + q.ssm_value_dim; + q.hc = 1; // a plain residual stream + q.hc_lr = 0; + q.idx_q_heads = 0; // no indexer: every cached cell is attended + q.idx_key_dim = 0; + q.has_ple = false; + q.has_indexer = false; + q.gdn_gate_silu = true; + + // What the kernels support. The widths must be the ones the kernels are written for (head 256 with 64 rotated + // dims, GDN state 128 and conv 4), and every count at most Flash-Next's: buffers sized by its compile-time + // numbers then hold the smaller model unchanged. + const ModelGeometry cap; + char buf[256]; + auto bad = [&](const char* what, int64_t got, const char* want) { + std::snprintf(buf, sizeof buf, "qwen35moe: %s is %lld; this engine supports %s", what, (long long) got, want); + err = buf; + return false; + }; + if (head_dim != cap.head_dim || value_len != cap.head_dim) return bad("attention.key_length", head_dim, "256"); + if (n_rot != 64) return bad("rope.dimension_count", n_rot, "64"); + if (q.ssm_state_size != cap.ssm_state_size) return bad("ssm.state_size", q.ssm_state_size, "128"); + if (q.ssm_d_conv != cap.ssm_d_conv) return bad("ssm.conv_kernel", q.ssm_d_conv, "4"); + if (q.ssm_value_dim != q.ssm_state_size * q.ssm_v_heads) + return bad("ssm.inner_size", q.ssm_value_dim, "state_size * time_step_rank"); + if (q.ssm_k_heads <= 0 || q.ssm_v_heads % q.ssm_k_heads != 0) + return bad("ssm.time_step_rank", q.ssm_v_heads, "a multiple of ssm.group_count"); + if (q.n_head_kv <= 0 || q.n_head % q.n_head_kv != 0) + return bad("attention.head_count", q.n_head, "a multiple of head_count_kv"); + if (q.qsa_interval != cap.qsa_interval) return bad("full_attention_interval", q.qsa_interval, "4"); + if (q.n_layers <= 0 || q.n_layers > cap.n_layers || q.n_layers % q.qsa_interval != 0) + return bad("block_count - nextn_predict_layers", q.n_layers, "a multiple of 4, at most 48"); + if (q.n_mtp_layers < 0 || q.n_mtp_layers > 1) return bad("nextn_predict_layers", q.n_mtp_layers, "0 or 1"); + if (q.n_embd <= 0 || q.n_embd > cap.n_embd || q.n_embd % 256 != 0) + return bad("embedding_length", q.n_embd, "a multiple of 256, at most 2560"); + if (q.n_head > cap.n_head) return bad("attention.head_count", q.n_head, "at most 24"); + if (q.n_head_kv != cap.n_head_kv) return bad("attention.head_count_kv", q.n_head_kv, "2"); + // the DeltaNet kernels (gdn_keyhead_ok<32, 8192>, the gdn_l2 / gdn_rec_kh grids) are built for exactly 16 key heads: + // a model with fewer would pass "at most" and run them with the wrong head count + if (q.ssm_k_heads != cap.ssm_k_heads) return bad("ssm.group_count", q.ssm_k_heads, "16"); + if (q.ssm_v_heads > cap.ssm_v_heads) return bad("ssm.time_step_rank", q.ssm_v_heads, "at most 48"); + if (q.n_expert <= 0 || q.n_expert > cap.n_expert) return bad("expert_count", q.n_expert, "at most 512"); + if (q.n_expert_used <= 0 || q.n_expert_used > cap.n_expert_used) + return bad("expert_used_count", q.n_expert_used, "at most 10"); + if (q.n_ff <= 0 || q.n_ff > cap.n_ff || q.n_ff % 256 != 0) + return bad("expert_feed_forward_length", q.n_ff, "a multiple of 256, at most 640"); + if (n_ff_shexp != q.n_ff) + return bad("expert_shared_feed_forward_length", n_ff_shexp, "expert_feed_forward_length"); + g = q; + return true; +} + std::string LayerView::name(const char* suffix) const { return "blk." + std::to_string(layer_) + "." + suffix; } diff --git a/src/core/mtp.cpp b/src/core/mtp.cpp index 971cc3e707b..178131503e0 100644 --- a/src/core/mtp.cpp +++ b/src/core/mtp.cpp @@ -11,6 +11,8 @@ #include "strata/kernels/bf16_gemv.hpp" #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/elementwise.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/artifact/gguf_reader.hpp" #include "strata/kernels/fused_gr.hpp" #include "strata/kernels/gr.hpp" #include "strata/kernels/kv_q4.hpp" @@ -92,6 +94,7 @@ strata::kernels::QsaShapes shapes_of(const ModelGeometry& g) { s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; + if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; // dense attention (qwen35moe) return s; } @@ -178,7 +181,11 @@ const void* MtpDrafter::wq(const char* name, int& type) const { for (const auto& e : q4_off_) if (e.first == name) { type = GGML_Q4_0; return dense4_ + e.second; } type = GGML_Q8_0; - return q8(name); + if (const void* p = q8(name)) return p; + // a qwen35moe file's MTP projection in another native type ("native") + for (const auto& t : tensors_) + if (t.name == name && t.kind.rfind("native", 0) == 0) { type = std::atoi(t.kind.c_str() + 6); return dense_ + t.off; } + return nullptr; } // --mtp-q4: every Q8_0 tensor of dense.bin as Q4_0 (ggml's reference quantizer), in one device buffer @@ -242,6 +249,151 @@ bool MtpDrafter::make_q4_head(std::string& err) { #endif } +bool MtpDrafter::load_gguf_layer(const std::string& gguf, const ModelGeometry& g, std::vector& dense, + std::string& err) { + constexpr uint32_t kF32 = 0, kQ8_0 = 8, kBF16 = 30; + try { + const strata::GgufModel model = strata::GgufModel::open(gguf); + const std::string pfx = "blk." + std::to_string(g.n_layers) + "."; + auto tensor = [&](const char* suffix, const strata::TensorInfo*& t) -> const uint8_t* { + size_t at = 0; + t = model.find(pfx + suffix, &at); + if (t == nullptr || !model.in_bounds(*t, at)) { + err = "mtp: " + gguf + " has no " + pfx + suffix + " (a qwen35moe draft layer needs the file's MTP block, " + "e.g. the -MTP GGUFs)"; + return nullptr; + } + return model.shard(at).tensor_data(*t); + }; + auto put = [&](const std::string& name, const char* kind, int64_t rows, int64_t cols, uint64_t bytes) { + Tensor x; + x.name = name; x.kind = kind; x.rows = rows; x.cols = cols; x.bytes = bytes; + x.off = (dense.size() + 255) & ~255ull; + dense.resize(x.off + bytes); + tensors_.push_back(x); + return dense.data() + x.off; + }; + const int64_t N = g.n_embd; + const strata::TensorInfo* t = nullptr; + // the RMSNorm weights, used as stored (llama.cpp's qwen35moe graph multiplies by them directly) + const char* norms[][2] = {{"nextn.enorm.weight", "pre_fc_norm_embedding.weight"}, + {"nextn.hnorm.weight", "pre_fc_norm_hidden.weight"}, + {"attn_norm.weight", "input_layernorm.weight"}, + {"attn_q_norm.weight", "self_attn.q_norm.weight"}, + {"attn_k_norm.weight", "self_attn.k_norm.weight"}, + {"post_attention_norm.weight", "post_attention_layernorm.weight"}, + {"nextn.shared_head_norm.weight", "shared_head_norm.weight"}}; + for (const auto& n : norms) { + const uint8_t* src = tensor(n[0], t); + if (src == nullptr) return false; + if (t->type != kF32) { err = std::string("mtp: ") + pfx + n[0] + " is not F32"; return false; } + std::memcpy(put(n[1], "f32", 1, (int64_t) t->elements(), t->elements() * 4), src, t->elements() * 4); + } + // the router and the shared expert's gate, read as BF16 (an F32 copy is rounded to nearest-even) + const char* routers[][2] = {{"ffn_gate_inp.weight", "mlp.gate.weight"}, + {"ffn_gate_inp_shexp.weight", "mlp.shared_expert_gate.weight"}}; + for (const auto& r : routers) { + const uint8_t* src = tensor(r[0], t); + if (src == nullptr) return false; + const uint64_t n = t->elements(); + auto* dst = (uint16_t*) put(r[1], "bf16", (int64_t) (n / (uint64_t) N), N, n * 2); + if (t->type == kBF16) { + std::memcpy(dst, src, n * 2); + } else if (t->type == kF32) { + for (uint64_t i = 0; i < n; ++i) { + uint32_t u; + std::memcpy(&u, src + 4 * i, 4); + dst[i] = (uint16_t) ((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); + } + } else { + err = std::string("mtp: ") + pfx + r[0] + " is neither BF16 nor F32"; + return false; + } + } + // the projections, rows of `cols` inputs: Q8_0 as the rt files hold them, or another type the native GEMV + // reads (the smaller Unsloth files keep the MTP layer's projections at Q5_K / Q6_K), kept as stored + const char* projs[][2] = {{"attn_q.weight", "self_attn.q_proj.weight"}, {"attn_k.weight", "self_attn.k_proj.weight"}, + {"attn_v.weight", "self_attn.v_proj.weight"}, {"attn_output.weight", "self_attn.o_proj.weight"}, + {"ffn_gate_shexp.weight", "mlp.shared_expert.gate_proj.weight"}, + {"ffn_up_shexp.weight", "mlp.shared_expert.up_proj.weight"}, + {"ffn_down_shexp.weight", "mlp.shared_expert.down_proj.weight"}}; + auto kind_of = [](uint32_t type) { return type == kQ8_0 ? std::string("q8_0") : "native" + std::to_string(type); }; + for (const auto& pj : projs) { + const uint8_t* src = tensor(pj[0], t); + if (src == nullptr) return false; + if (!strata::kernels::native_mmvq_supported((int) t->type) || t->shape.size() != 2) { + err = std::string("mtp: ") + pfx + pj[0] + " is " + t->type_name() + ", which the draft layer cannot read"; + return false; + } + const uint64_t cols = t->shape[0], rows = t->shape[1], bytes = strata::tensor_payload_bytes(*t); + std::memcpy(put(pj[1], kind_of(t->type).c_str(), (int64_t) rows, (int64_t) cols, bytes), src, bytes); + } + // eh_proj [2N inputs (the normed embedding, then the normed hidden), N outputs]: the two column halves are + // the rt files' fc_embedding and fc_hidden (e_proj(e) + h_proj(h) = eh_proj(concat(e, h))) + { + const uint8_t* src = tensor("nextn.eh_proj.weight", t); + if (src == nullptr) return false; + if (!strata::kernels::native_mmvq_supported((int) t->type) || t->shape.size() != 2 || + t->shape[0] != (uint64_t) (2 * N) || t->shape[1] != (uint64_t) N) { + err = "mtp: " + pfx + "nextn.eh_proj.weight is not a [2 x n_embd, n_embd] matrix the draft layer reads"; + return false; + } + // a row's two halves are whole blocks (n_embd is a multiple of 256) + const uint64_t half = strata::tensor_payload_bytes(*t) / (uint64_t) N / 2; + // (both put() first: a put may grow `dense` and move it, so the pointers are taken after the second) + const uint64_t eo = put("fc_embedding.weight", kind_of(t->type).c_str(), N, N, half * (uint64_t) N) - dense.data(); + const uint64_t ho = put("fc_hidden.weight", kind_of(t->type).c_str(), N, N, half * (uint64_t) N) - dense.data(); + uint8_t* e = dense.data() + eo; + uint8_t* h = dense.data() + ho; + for (int64_t r = 0; r < N; ++r) { + std::memcpy(e + r * half, src + r * 2 * half, half); + std::memcpy(h + r * half, src + r * 2 * half + half, half); + } + } + // the routed experts, every one resident in VRAM, each blob gate | up | down as the native packs lay it out + const strata::TensorInfo *tg = nullptr, *tu = nullptr, *td = nullptr; + const uint8_t* sg = tensor("ffn_gate_exps.weight", tg); + const uint8_t* su = sg ? tensor("ffn_up_exps.weight", tu) : nullptr; + const uint8_t* sd = su ? tensor("ffn_down_exps.weight", td) : nullptr; + if (sd == nullptr) return false; + xgu_ = (int) tg->type; + xd_ = (int) td->type; + if (tu->type != tg->type || !strata::kernels::native_expert_supported(xgu_, xd_, N, g.n_ff)) { + err = std::string("mtp: the draft layer's experts are ") + tg->type_name() + "/" + td->type_name() + + ", which this engine has no GPU kernels for"; + return false; + } + const strata::kernels::NativeExpertLayout L = strata::kernels::native_expert_layout(xgu_, xd_, N, g.n_ff); + if (L.gu_row == 0 || L.d_row == 0) { + err = "mtp: no row size for the draft layer's expert types"; + return false; + } + xblob_ = L.bytes; + const uint64_t gu_bytes = (uint64_t) L.gu_row * (uint64_t) g.n_ff, d_bytes = (uint64_t) L.d_row * (uint64_t) N; + if (L.up_off != gu_bytes || L.down_off != 2 * gu_bytes || L.bytes < 2 * gu_bytes + d_bytes) { + err = "mtp: unexpected native expert layout"; + return false; + } + const uint64_t bytes = (uint64_t) g.n_expert * xblob_; + if (cudaMalloc((void**) &experts_, bytes) != cudaSuccess) { + err = "mtp: the draft layer's " + std::to_string(g.n_expert) + " experts do not fit in VRAM"; + return false; + } + std::vector blob(xblob_, 0); + for (int64_t e = 0; e < g.n_expert; ++e) { + std::memcpy(blob.data(), sg + e * gu_bytes, gu_bytes); + std::memcpy(blob.data() + L.up_off, su + e * gu_bytes, gu_bytes); + std::memcpy(blob.data() + L.down_off, sd + e * d_bytes, d_bytes); + cudaMemcpy(experts_ + e * xblob_, blob.data(), xblob_, cudaMemcpyHostToDevice); + } + vram_ += bytes; + } catch (const std::exception& ex) { + err = std::string("mtp: ") + gguf + ": " + ex.what(); + return false; + } + return true; +} + bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, SessionState& ss, int max_t, std::string& err, int64_t window, const MtpDrafter* shared) { cudaGetDevice(&device_); // a layer split's last stage on another GPU: the drafter lives there @@ -260,14 +412,26 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session dense_ = shared->dense_; experts_ = shared->experts_; tensors_ = shared->tensors_; + xgu_ = shared->xgu_; + xd_ = shared->xd_; + xblob_ = shared->xblob_; owns_weights_ = false; owns_draft_head_ = false; } // Loader fix (0.1.15+loaderfix.2): the two reads below are the whole “drafter files” cost; reporting // them apart from the rest of the stage is what makes the next regression visible. const auto t_files = std::chrono::steady_clock::now(); - // ---- the index and the dense weights - if (shared == nullptr) { + // ---- the index and the dense weights (a model without hyper-connections: its GGUF's own MTP block) + const bool gguf_layer = !g.has_hc(); + if (shared == nullptr && gguf_layer) { + std::vector blob; + if (!load_gguf_layer(rt_dir, g, blob, err)) return false; + if (cudaMalloc((void**) &dense_, blob.size()) != cudaSuccess) { err = "mtp: the dense weights do not fit"; return false; } + cudaMemcpy(dense_, blob.data(), blob.size(), cudaMemcpyHostToDevice); + vram_ += blob.size(); + if (q4_ && !make_q4_dense(blob, err)) return false; + } + if (shared == nullptr && !gguf_layer) { std::ifstream idx(rt_dir + "/dense.txt"); if (!idx) { err = "mtp: cannot open " + rt_dir + "/dense.txt (run tools/mtp_rt.py)"; return false; } std::string line; @@ -295,7 +459,7 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session if (q4_ && !make_q4_dense(blob, err)) return false; } // ---- the 512 routed experts, one blob each - if (shared == nullptr) { + if (shared == nullptr && !gguf_layer) { const uint64_t bytes = (uint64_t) g.n_expert * strata::kernels::cpu::BLOB; // Loader fix (0.1.15+loaderfix.2): each 64 MiB read below reached the disk as ~16k 4095-byte reads under // MSVC's `basic_filebuf::xsgetn`, which is what made 675 MiB of drafter experts take minutes. @@ -321,7 +485,10 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session const char* required[] = {"fc_embedding.weight", "fc_hidden.weight", "self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight", "self_attn.o_proj.weight", "mlp.shared_expert.gate_proj.weight", "mlp.shared_expert.up_proj.weight", "mlp.shared_expert.down_proj.weight"}; - for (const char* n : required) if (!q8(n)) { err = std::string("mtp: ") + n + " is missing (q8_0)"; return false; } + for (const char* n : required) { + int type = 0; + if (!wq(n, type)) { err = std::string("mtp: ") + n + " is missing"; return false; } + } // ---- the layer's own K/V (dense attention: no indexer state is read) const strata::kernels::QsaShapes s = shapes_of(g); @@ -398,6 +565,10 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session grp_counts_ = b.take(4); hit_xq_ = b.take(T * (N / 32) * 34); hit_xs_ = b.take(T * (N / 32)); hit_scratch_ = b.take(strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff)); + if (xblob_ > 0) { // native experts + nat_xq_ = b.take(T * (N / 32) * 36); + nat_scratch_ = b.take(strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff)); + } sh_scratch_ = (float*) b.take(strata::kernels::shared_expert_scratch_bytes(g.n_ff)); x_bf16_ = b.take(N); out_ids_ = b.take(T + 4); @@ -462,11 +633,11 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session std::fprintf(stderr, "strata mtp: shared draft weights, %.0f MiB of private state and buffers\n", (double) vram_ / 1048576.0); } else { - std::fprintf(stderr, "strata mtp: draft layer loaded, %.0f MiB of VRAM (experts %.0f, dense %.0f), files read in %.2f s (%.0f MiB/s)\n", - (double) vram_ / 1048576.0, (double) g.n_expert * strata::kernels::cpu::BLOB / 1048576.0, + const double xbytes = (double) g.n_expert * (double) (xblob_ > 0 ? xblob_ : strata::kernels::cpu::BLOB); + std::fprintf(stderr, "strata mtp: draft layer loaded%s, %.0f MiB of VRAM (experts %.0f, dense %.0f), files read in %.2f s (%.0f MiB/s)\n", + gguf_layer ? " from the model file" : "", (double) vram_ / 1048576.0, xbytes / 1048576.0, (double) tensors_.back().off / 1048576.0, files_s, - files_s > 0 ? ((double) g.n_expert * strata::kernels::cpu::BLOB + (double) tensors_.back().off) / - 1048576.0 / files_s : 0.0); + files_s > 0 ? (xbytes + (double) tensors_.back().off) / 1048576.0 / files_s : 0.0); } return true; } @@ -731,8 +902,10 @@ bool MtpDrafter::record_front(int T, int row0, cudaStream_t cs, std::string& err native_mmvq(wt_fc_hidden, wp_fc_hidden, xq_, h2_ + (size_t) c0 * N, (int) N, (int) N, nc, cs); } add_streams_broadcast(h2_, e2_, R_, N, (int) HC, T, cs); - // ---- the attention hyper-connection - { + // ---- the attention hyper-connection (one stream: the RMSNorm before the attention) + if (!g.has_hc()) { + native_qsa_rms_norm_weighted(R_, f32("input_layernorm.weight"), mixed_, (int) N, T, EPS, cs); + } else { FusedGrArgs fa[kFusedGrMaxT]; for (int t = 0; t < T; ++t) { fa[t].R = R_ + (size_t) t * HC * N; fa[t].R_out = R_ + (size_t) t * HC * N; fa[t].apply = false; @@ -845,8 +1018,12 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) { native_qsa_gate_apply(attn_, qfull_, attn32_, (int) (T * NH), (int) HD, cs); native_quantize_q8_1(attn32_, xq_, (int) (NH * HD), T, cs); native_mmvq(wt_o_proj, wp_o_proj, xq_, bo_, (int) (NH * HD), (int) N, T, cs); - // ---- the MLP hyper-connection (the attention write folded in) - { + // ---- the MLP hyper-connection (the attention write folded in; one stream: the residual add, then the RMSNorm + // before the MoE) + if (!g.has_hc()) { + add_inplace(R_, bo_, (int64_t) T * N, cs); + native_qsa_rms_norm_weighted(R_, f32("post_attention_layernorm.weight"), mixed_, (int) N, T, EPS, cs); + } else { FusedGrArgs fa[kFusedGrMaxT]; for (int t = 0; t < T; ++t) { fa[t].R = R_ + (size_t) t * HC * N; fa[t].R_out = R_ + (size_t) t * HC * N; fa[t].apply = true; @@ -917,10 +1094,17 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) { if (T > 1 && native_router_enabled() && g.n_expert == 512 && K == 10) { bf16_gemv_fp32_mmvf_multi(mixed_, N, bf16("mlp.gate.weight"), logits_, g.n_expert, (int) N, (int) g.n_expert, T, cs); native_router_top10_multi(logits_, ids_, w_, T, cs); + } else if (T > 1 && !(native_router_enabled() && g.n_expert == 512 && K == 10) && + !std::getenv("STRATA_ROUTE_PER_TOKEN")) { + // any other router (qwen35moe's 256 x 8): one multi-column projection and one top-k for the window's rows, + // each row bitwise the per-token calls below (the same kernel per column; router_top10 is per row) + bf16_gemv_fp32_mmvf_cols(mixed_, bf16("mlp.gate.weight"), logits_, N, g.n_expert, (int) T, cs); + router_top10(logits_, (int) T, (int) g.n_expert, (int) K, ids_, w_, cs); } else { for (int t = 0; t < T; ++t) { bf16_gemv_fp32_mmvf(mixed_ + t * N, bf16("mlp.gate.weight"), logits_ + t * g.n_expert, (int) N, (int) g.n_expert, cs); - if (native_router_enabled() && g.n_expert == 512 && K == 10) native_router_top10(logits_ + t * g.n_expert, ids_ + t * K, w_ + t * K, cs); + if (native_router_enabled() && g.n_expert == 512 && K == 10) + native_router_top10(logits_ + t * g.n_expert, ids_ + t * K, w_ + t * K, cs); else router_top10(logits_ + t * g.n_expert, 1, (int) g.n_expert, (int) K, ids_ + t * K, w_ + t * K, cs); } } @@ -929,11 +1113,19 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) { native_route_resident(logits_, ids_, w_, rr_res_, T, rr_margin_, rr_lo_, rr_hi_, nullptr, cs); } catch (const std::exception& e) { err = std::string("mtp route-resident: ") + e.what(); return false; } } + if (xblob_ > 0) { // the model file's experts in their GGUF form + moe_group_resident(ids_, (int) (T * K), (int) K, experts_, (int64_t) xblob_, grp_ptr_, grp_start_, grp_counts_, + hit_dst_, hit_slot_, cs); + quantize_q8_1_rows(mixed_, T, N, nat_xq_, cs); + native_expert_grouped(native_expert_layout(xgu_, xd_, N, g.n_ff), grp_ptr_, grp_start_, grp_counts_, hit_dst_, + hit_slot_, (int64_t) T * K, (int64_t) T * K, nat_xq_, nat_scratch_, parts_, cs); + } else { moe_group_resident(ids_, (int) (T * K), (int) K, experts_, (int64_t) strata::kernels::cpu::BLOB, grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, cs); quantize_q8_0_scaled(mixed_, hit_xq_, hit_xs_, (int64_t) T * N, cs); moe_grouped_s2(grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, (int64_t) T * K, (int64_t) T * K, hit_xq_, hit_xs_, hit_scratch_, parts_, cs); + } if (branch_on) { if (cudaStreamWaitEvent(cs, sh_join_, 0) != cudaSuccess) { err = "mtp: the shared expert's join"; return false; } } else { @@ -950,17 +1142,23 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) { else moe_combine(parts_ + (size_t) t * K * N, w_ + t * K, shared_ + t * N, y_ + t * N, N, K, cs); } - if (!fuse_head_gr && no_multi_gr) + if (!fuse_head_gr && no_multi_gr && g.has_hc()) gr_write(R_ + (size_t) t * HC * N, y_ + t * N, inj2_ + t * HC, gs, R_ + (size_t) t * HC * N, cs); } if (T > 1 && native_moe_combine_enabled()) native_moe_combine_multi(parts_, w_, shared_, y_, N, K, T, cs); // #783 PR-g (stuchapin909): the T residual writes in one launch (each token's R, block output and injection are // its own, so doing them after the loop changes nothing) - if (!fuse_head_gr && !no_multi_gr) + if (!fuse_head_gr && !no_multi_gr && g.has_hc()) gr_write_multi(R_, y_, inj2_, gs, R_, T, cs); // ---- the final mixer and the main model's head - if (fuse_head_gr) { + if (!g.has_hc()) { + // one stream: the MoE output into the residual, then shared_head_norm is the head's input - and the + // hidden the next draft step reads (llama.cpp's h_nextn), so it goes back into R_ for mtp_select + add_inplace(R_, y_, (int64_t) T * N, cs); + native_qsa_rms_norm_weighted(R_, f32("shared_head_norm.weight"), sample_, (int) N, T, EPS, cs); + cudaMemcpyAsync(R_, sample_, (size_t) T * N * sizeof(float), cudaMemcpyDeviceToDevice, cs); + } else if (fuse_head_gr) { FusedGrArgs fa[kFusedGrMaxT]; for (int t = 0; t < T; ++t) { fa[t].R = R_ + (size_t) t * HC * N; fa[t].R_out = R_ + (size_t) t * HC * N; fa[t].apply = true; diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 3995ae9dd13..930ec861e6d 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -19,6 +19,7 @@ #include "strata/kernels/native_moe.hpp" #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/elementwise.hpp" +#include "strata/kernels/native_gr_norm.hpp" #include "strata/kernels/fused_gr.hpp" #include "strata/kernels/route_prior.hpp" #include "strata/kernels/cvec.hpp" @@ -312,6 +313,7 @@ strata::kernels::QsaShapes shapes_of(const ModelGeometry& g) { s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; + if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; // dense attention (qwen35moe) return s; } @@ -518,6 +520,15 @@ Verifier::~Verifier() { if (h) cudaFreeHost(h); } +uint64_t Verifier::dense_attention_bytes(const ModelGeometry& g, const SessionState& ss, int max_t) { + if (g.has_indexer || ss.qsa_states == nullptr) return 0; + // init's sel_ and attn_scratch_ with cap_ = the whole context + const int64_t cap = ss.qsa_states[ss.qsa_primary()].max_cells; + const uint64_t per_row = (uint64_t) cap * 4 + + strata::kernels::qsa_decode_attn_scratch_floats(cap, shapes_of(g)) * 4; + return (uint64_t) max_t * per_row; +} + bool Verifier::set_logit_bias(const std::vector& bias, std::string& err) { if (next_ && !next_->set_logit_bias(bias, err)) return false; if (le_ < g_->n_layers) return true; @@ -573,8 +584,8 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& err = "verify: " + why + " (the verify window reproduces the default native decode path)"; return false; } - if (!strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr) || ss.k != 10 || g.ssm_state_size != 128 || - g.ssm_d_conv != 4) { + if ((g.has_hc() && !strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr)) || ss.k != g.n_expert_used || + (g.arch == ModelArch::kQwen4Exp && ss.k != 10) || g.ssm_state_size != 128 || g.ssm_d_conv != 4) { err = "verify: geometry differs from the artifact's"; return false; } @@ -589,7 +600,9 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& n_vocab_ = wo->ne1; const strata::kernels::QsaShapes s = shapes_of(g); - cap_ = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); + // without the indexer every cell is attended: the selection holds the whole context + cap_ = g.has_indexer ? strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s) + : ss.qsa_states[ss.qsa_primary()].max_cells; max_blocks_ = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2; attn_scratch_floats_ = (int64_t) strata::kernels::qsa_decode_attn_scratch_floats(cap_, s); @@ -811,7 +824,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& return true; } -const float* Verifier::final_R(int t) const { return R_ + (size_t) t * (size_t) (g_->hc * g_->n_embd); } +const float* Verifier::final_R(int t) const { return final_R_all() + (size_t) t * (size_t) (g_->hc * g_->n_embd); } // ================================ THE WINDOW, AS CAPTURED ================================ // @@ -934,6 +947,12 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { else gdn_idx[(size_t) l] = gi++; } } + // dense attention (no indexer): every attention layer selects the same cells 0..width-1 (the width is the row's, + // not the layer's), so the identity ids are written once per window, at this stage's first attention layer, and + // the later ones read them (nothing else writes sel_ on this path) + int64_t first_qsa = -1; + for (int64_t l = lb_; l < le_ && first_qsa < 0; ++l) + if (qsa_idx[(size_t) l] >= 0) first_qsa = l; // ---------------------------------------------------------------- pre(l, group): up to the ring auto pre = [&](int64_t l, int grp) -> bool { @@ -941,13 +960,19 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { stamp(l, 0, grp); const LayerView v(wt, l); const char* pfx[2] = {"hc_attn_", "hc_ffn_"}; - const WeightRef *wn[2], *wd[2], *wu[2], *wi[2]; - for (int h = 0; h < 2; ++h) { - wn[h] = need(v, (std::string(pfx[h]) + "norm.weight").c_str(), err); - wd[h] = need(v, (std::string(pfx[h]) + "down.weight").c_str(), err); - wu[h] = need(v, (std::string(pfx[h]) + "up.weight").c_str(), err); - wi[h] = need(v, (std::string(pfx[h]) + "inject.weight").c_str(), err); - if (!wn[h] || !wd[h] || !wu[h] || !wi[h]) return false; + const WeightRef *wn[2] = {}, *wd[2] = {}, *wu[2] = {}, *wi[2] = {}; + if (g.has_hc()) { + for (int h = 0; h < 2; ++h) { + wn[h] = need(v, (std::string(pfx[h]) + "norm.weight").c_str(), err); + wd[h] = need(v, (std::string(pfx[h]) + "down.weight").c_str(), err); + wu[h] = need(v, (std::string(pfx[h]) + "up.weight").c_str(), err); + wi[h] = need(v, (std::string(pfx[h]) + "inject.weight").c_str(), err); + if (!wn[h] || !wd[h] || !wu[h] || !wi[h]) return false; + } + } else { // qwen35moe: one residual stream, an RMSNorm before the mixer and one before the MoE + wn[0] = need(v, "attn_norm.weight", err); + wn[1] = need(v, "post_attention_norm.weight", err); + if (!wn[0] || !wn[1]) return false; } // the previous layer's FFN write, folded into this layer's first read (a control vector after it has // already applied it) @@ -1051,6 +1076,14 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { pending = false; } auto gr_read_group = [&](int half, bool apply, float* inj_prev, float* inj_out) { + if (!g.has_hc()) { + // one stream: the pending block output into the residual, then its RMSNorm with this half's weight + // (llama.cpp's build_norm) is the mixer's / the MoE's input. The rows of a group are contiguous. + if (apply) add_inplace(Rt(tb), bo_ + tb * N, (int64_t) n * N, cs); + native_gr_rms_norm_weighted_multi(Rt(tb), (const float*) wn[half]->data, mixed_ + tb * N, (int) N, 1, n, + EPS, cs); // one gamma row for the group's n tokens + return false; + } FusedGrArgs fa[kFusedGrMaxT]; for (int t = tb; t < te; ++t) { FusedGrArgs& a = fa[t - tb]; @@ -1184,12 +1217,16 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { // ======================= QSA ======================= const int64_t qi = qsa_idx[(size_t) l]; const QsaState& st = ss.qsa_states[qi]; - const WeightRef *wik = need(v, "indexer.k_proj.weight", err), *wq = need(v, "attn_q.weight", err), + // qwen35moe has no indexer: every cached cell is attended (qsa_select_all), nothing is pooled + const bool idx = g.has_indexer; + const WeightRef *wik = idx ? need(v, "indexer.k_proj.weight", err) : nullptr, *wq = need(v, "attn_q.weight", err), *wk = need(v, "attn_k.weight", err), *wv = need(v, "attn_v.weight", err), - *wo = need(v, "attn_output.weight", err), *wiq = need(v, "indexer.q_proj.weight", err), + *wo = need(v, "attn_output.weight", err), + *wiq = idx ? need(v, "indexer.q_proj.weight", err) : nullptr, *wqn = need(v, "attn_q_norm.weight", err), *wkn = need(v, "attn_k_norm.weight", err), - *wiqn = need(v, "indexer.q_norm.weight", err), *wikn = need(v, "indexer.k_norm.weight", err); - if (!wik || !wq || !wk || !wv || !wo || !wiq || !wqn || !wkn || !wiqn || !wikn) return false; + *wiqn = idx ? need(v, "indexer.q_norm.weight", err) : nullptr, + *wikn = idx ? need(v, "indexer.k_norm.weight", err) : nullptr; + if (!wq || !wk || !wv || !wo || !wqn || !wkn || (idx && (!wik || !wiq || !wiqn || !wikn))) return false; if (!native_of(wq, v.name("attn_q.weight"), err) || !native_of(wk, v.name("attn_k.weight"), err) || !native_of(wv, v.name("attn_v.weight"), err) || !native_of(wo, v.name("attn_output.weight"), err)) return false; @@ -1224,7 +1261,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { if (!q8_attn) native_quantize_q8_1(xm, xq_, (int) N, n, cs); // STRATA_DF_BRANCH: the query side (q and its norm/RoPE/rotation on side 0, the indexer query on side 1) // beside the K/V side; the indexer query joins before the block scores, the query before attention - const bool qbr = br && qb; + // (not without the indexer - qwen35moe: side 1 is the indexer query; that model takes the plain path) + const bool qbr = br && qb && idx; if (qbr) { fork(0); native_mmvq(wq->native_type, wq->native_data, xq_, qfull_ + tb * NH * 2 * HD, (int) N, @@ -1238,7 +1276,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { IQ * ID, N, IQ * ID, n, df_side_[1]); norm_rope_on(qidx_ + tb * IQ * ID, wiqn, (int) (n * IQ), (int) ID, pos_i + tb * IQ, df_side_[1]); } - if (qb) bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wik->data, idx_raw + tb * ID, ID, N, ID, n, cs); + if (!idx) { + } else if (qb) bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wik->data, idx_raw + tb * ID, ID, N, ID, n, cs); else for (int t = tb; t < te; ++t) bf16_gemv_fp32_mmvf(mixed_ + t * N, (const uint16_t*) wik->data, idx_raw + t * ID, (int) N, (int) ID, cs); stamp(l, 7, grp); @@ -1253,7 +1292,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { fwht256_inplace_cuda(vcur_ + tb * NKV * HD, (int64_t) n * NKV, cs); } stamp(l, 8, grp); - if (batch_rec_) { // snapshot each slot's indexer tail before its first proposed row + if (!idx) { + } else if (batch_rec_) { // snapshot each slot's indexer tail before its first proposed row for (int t = tb; t < te; ++t) if (t == tb || brow_[t] != brow_[t - 1]) copy_from_mapped(tail_snap_b_ + ((size_t) brow_[t] * nQall + qi) * TS, @@ -1307,7 +1347,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { kv_append_step(st.k_pool, st.v_pool, st.page_table, step_t, kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st.host); } - if (dec_batch && !g_no_batch_kv && !batch_rec_) { + if (!idx) { + } else if (dec_batch && !g_no_batch_kv && !batch_rec_) { const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; native_qsa_indexer_append_steps(idx_raw + tb * ID, step_ + tb * kStepCount + kStepPos, kStepCount, n, 0, (const float*) wikn->data, EPS, ib, s, st.max_cells, @@ -1337,9 +1378,11 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { norm_rope(qcur_ + tb * NH * HD, wqn, (int) (n * NH), (int) HD, pos_ + tb * NH); } if (st.kv_rot) fwht256_inplace_cuda(qcur_ + tb * NH * HD, (int64_t) n * NH, cs); // = - bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, IQ * ID, - N, IQ * ID, n, cs); - norm_rope(qidx_ + tb * IQ * ID, wiqn, (int) (n * IQ), (int) ID, pos_i + tb * IQ); + if (idx) { + bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, + IQ * ID, N, IQ * ID, n, cs); + norm_rope(qidx_ + tb * IQ * ID, wiqn, (int) (n * IQ), (int) ID, pos_i + tb * IQ); + } } else { for (int t = tb; t < te; ++t) { float* qc = qcur_ + t * NH * HD; @@ -1356,7 +1399,7 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } if (st.kv_rot) fwht256_inplace_cuda(qc, NH, cs); // = } - for (int t = tb; t < te; ++t) { + for (int t = tb; t < te && idx; ++t) { float* qx = qidx_ + t * IQ * ID; bf16_gemv_fp32_mmvf(mixed_ + t * N, (const uint16_t*) wiq->data, qx, (int) N, (int) (IQ * ID), cs); norm_rope(qx, wiqn, (int) IQ, (int) ID, pos_ + t * NH); @@ -1367,10 +1410,14 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { if (batch_rec_) { // each row selects and attends over its own slot's K/V for (int t = tb; t < te; ++t) { const QsaState& sx = slot_ss(t).qsa_states[qi]; + if (!idx) { + if (l == first_qsa) qsa_select_all(step_ + t * kStepCount, 1, cap_, sel_ + (size_t) t * cap_, cs); + } else { qsa_block_scores(sx.idx_pooled, sx.idx_dead, qidx_ + t * IQ * ID, step_ + t * kStepCount, 1, max_blocks_, s, scores_ + (size_t) t * max_blocks_, cs); qsa_block_topk(scores_ + (size_t) t * max_blocks_, step_ + t * kStepCount, 1, max_blocks_, cap_, s, sel_ + (size_t) t * cap_, cs); + } qsa_kv_resolve(sx, *g_, sel_ + (size_t) t * cap_, step_ + t * kStepCount, 1, cap_, cs); const QsaAttnPools px = qsa_attn_pools(sx); qsa_decode_attn_batch(qcur_ + t * NH * HD, px, sel_ + (size_t) t * cap_, step_ + t * kStepCount, @@ -1378,10 +1425,14 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { attn_ + t * NH * HD, 1, cs); } } else { + if (!idx) { + if (l == first_qsa) qsa_select_all(step_ + tb * kStepCount, n, cap_, sel_ + (size_t) tb * cap_, cs); + } else { qsa_block_scores(st.idx_pooled, st.idx_dead, qidx_ + tb * IQ * ID, step_ + tb * kStepCount, n, max_blocks_, s, scores_ + (size_t) tb * max_blocks_, cs); qsa_block_topk(scores_ + (size_t) tb * max_blocks_, step_ + tb * kStepCount, n, max_blocks_, cap_, s, sel_ + (size_t) tb * cap_, cs); + } stamp(l, 11, grp); // KV streaming: the n selections' blocks resident (device-side, inside the graph) qsa_kv_resolve(st, *g_, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, n, cap_, cs); @@ -1446,15 +1497,14 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { NE, n, cs); native_router_top10_multi(logits_ + tb * NE, ids_ + tb * K, w_ + tb * K, n, cs); } catch (const std::exception& e) { err = "verify router: " + std::string(e.what()); return false; } -#if defined(STRATA_HIP_GFX906) } else if (dec_batch) { - // gfx906: the Coder's 256-expert router (the native router is 512 x 10 only) for the whole window in one - // multi-column BF16 projection and one top-k, bitwise the per-token calls (route_window_parity); ~3 ms - // of a ~55 ms verify window on 2x MI50. STRATA_ROUTE_PER_TOKEN=1: the per-token calls. + // every other router (qwen35moe's 256 x 8, the Coder's 256 x 10, gfx906) for the whole window in one + // multi-column BF16 projection and one top-k launch, bitwise the per-token calls (route_window_parity): + // 2 launches per layer instead of 2 per row (~3 ms of a ~55 ms window measured on 2x MI50, where it began). + // STRATA_ROUTE_PER_TOKEN=1: the per-token calls. if (!moe_route_window(wt, g, l, K, ss.moe, mixed_ + tb * N, logits_ + tb * NE, ids_ + tb * K, w_ + tb * K, n, cs, err)) return false; -#endif } else for (int t = tb; t < te; ++t) { MoEBuffers mb = ss.moe; @@ -1674,13 +1724,21 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { device_plan_ ? skip_ + grp : nullptr, ring, cs); stamp(l, 24, grp); if (l == g.n_layers - 1) { - if (!fuse_head_gr) { + if (!fuse_head_gr && g.has_hc()) { // (one stream: the head's read adds the last output) if (dec_batch && !g_no_multi_gr) gr_write_multi(Rt(tb), bo_ + tb * N, inj2_ + tb * HC, gs, Rt(tb), n, cs); else for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); if (cvec().covers(l)) cvec_apply(Rt(tb), l, n, HC * N, nullptr, 0, nullptr, 0, false, cs); } } else if (cvec().covers(l)) { - cvec_apply(Rt(tb), l, n, HC * N, bo_ + tb * N, N, inj2_ + tb * HC, HC, true, cs); + if (!g.has_hc()) { + // one stream (qwen35moe): the write is a plain add (the next read skips it: `pending` is false after a + // steered layer), then the vector - as the prompt path does; the hyper-connection gate would read + // inj2_, which this model never writes + add_inplace(Rt(tb), bo_ + tb * N, (int64_t) n * N, cs); + cvec_apply(Rt(tb), l, n, HC * N, nullptr, 0, nullptr, 0, false, cs); + } else { + cvec_apply(Rt(tb), l, n, HC * N, bo_ + tb * N, N, inj2_ + tb * HC, HC, true, cs); + } } return true; }; @@ -1703,67 +1761,78 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { // ---- the head, T columns, and the argmax of each stamp(g.n_layers, 0, 0); { - const WeightRef *hn = wt.find("output_hc_norm.weight"), *hd = wt.find("output_hc_down.weight"), - *hu = wt.find("output_hc_up.weight"); - if (!hn || !hd || !hu) { err = "verify: an output_hc_* weight is missing"; return false; } - // S25 (STRATA_HC_Q8=1): the final mixer from the GGUF's Q8_0 projections too, through the same multi read - // (no pending write, no inject) - const bool mix_q8 = hd->hc_q8 != nullptr && hu->hc_q8 != nullptr && head_ != nullptr && head_->loaded(); - if (mix_q8) { - FusedGrArgs fa[kFusedGrMaxT]; - for (int t = 0; t < T; ++t) { - FusedGrArgs& a = fa[t]; - a.R = Rt(t); a.R_out = Rt(t); a.apply = false; - a.w_norm = (const float*) hn->data; a.w_down = (const uint16_t*) hd->data; - a.w_up = (const uint16_t*) hu->data; a.w_inject = nullptr; - a.q8_down = (const uint8_t*) hd->hc_q8; a.q8_up = (const uint8_t*) hu->hc_q8; - a.eps = EPS; a.lo = lo_ + t * g.hc_lr; a.rs = rs_ + t * HC; - a.mixed = head_mixed_ + t * N; - } - fused_gr_read_multi(fa, T, xn_, cs); - } - else if (fuse_head_gr) { - FusedGrArgs fa[kFusedGrMaxT]; - for (int t = 0; t < T; ++t) { - FusedGrArgs& a = fa[t]; - a.R = Rt(t); a.R_out = Rt(t); a.apply = true; - a.bo_prev = bo_ + t * N; a.inj_prev = inj2_ + t * HC; - a.w_norm = (const float*) hn->data; - a.w_down = (const uint16_t*) hd->data; - a.w_up = (const uint16_t*) hu->data; - a.w_inject = nullptr; - a.eps = EPS; - a.lo = lo_ + t * g.hc_lr; - a.rs = rs_ + t * HC; - a.inject_out = head_inj_; - a.mixed = head_mixed_ + t * N; - } - fused_gr_read_multi(fa, T, xn_, cs); - } else if (head_mix_multi_enabled() && head_ != nullptr && head_->loaded()) { - // the final mixer, the window's tokens in one read (fused_gr_read_multi without the pending write: the - // last layer's write was done above; its sums are gr_read's) - if (hn->kind != WeightKind::F32 || hd->kind != WeightKind::Bf16InF32 || hu->kind != WeightKind::Bf16InF32) { - err = "verify: the output_hc_* weights have the wrong engine forms"; + if (!g.has_hc()) { + // one stream (qwen35moe): the last layer's output into the residual, then output_norm is the head's input + const WeightRef* on = wt.find("output_norm.weight"); + if (on == nullptr || head_ == nullptr || !head_->loaded()) { + err = "verify: a model without hyper-connections needs output_norm.weight and the native head"; return false; } - FusedGrArgs fa[kFusedGrMaxT]; - for (int t = 0; t < T; ++t) { - fa[t].R = Rt(t); fa[t].R_out = Rt(t); fa[t].apply = false; - fa[t].w_norm = (const float*) hn->data; fa[t].w_down = (const uint16_t*) hd->data; - fa[t].w_up = (const uint16_t*) hu->data; fa[t].eps = EPS; - fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].mixed = head_mixed_ + t * N; - } - fused_gr_read_multi(fa, T, xn_, cs); + add_inplace(R_, bo_, (int64_t) T * N, cs); + native_gr_rms_norm_weighted_multi(R_, (const float*) on->data, head_mixed_, (int) N, 1, T, EPS, cs); } else { - for (int t = 0; t < T; ++t) { - BlockBuffers bb = ss.block; - bb.R = Rt(t); - bb.mixed = head_mixed_ + t * N; - if (head_ != nullptr && head_->loaded()) { - if (!lm_head_mix(wt, g, bb, cs, err)) return false; - } else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) { + const WeightRef *hn = wt.find("output_hc_norm.weight"), *hd = wt.find("output_hc_down.weight"), + *hu = wt.find("output_hc_up.weight"); + if (!hn || !hd || !hu) { err = "verify: an output_hc_* weight is missing"; return false; } + // S25 (STRATA_HC_Q8=1): the final mixer from the GGUF's Q8_0 projections too, through the same multi read + // (no pending write, no inject) + const bool mix_q8 = hd->hc_q8 != nullptr && hu->hc_q8 != nullptr && head_ != nullptr && head_->loaded(); + if (mix_q8) { + FusedGrArgs fa[kFusedGrMaxT]; + for (int t = 0; t < T; ++t) { + FusedGrArgs& a = fa[t]; + a.R = Rt(t); a.R_out = Rt(t); a.apply = false; + a.w_norm = (const float*) hn->data; a.w_down = (const uint16_t*) hd->data; + a.w_up = (const uint16_t*) hu->data; a.w_inject = nullptr; + a.q8_down = (const uint8_t*) hd->hc_q8; a.q8_up = (const uint8_t*) hu->hc_q8; + a.eps = EPS; a.lo = lo_ + t * g.hc_lr; a.rs = rs_ + t * HC; + a.mixed = head_mixed_ + t * N; + } + fused_gr_read_multi(fa, T, xn_, cs); + } + else if (fuse_head_gr) { + FusedGrArgs fa[kFusedGrMaxT]; + for (int t = 0; t < T; ++t) { + FusedGrArgs& a = fa[t]; + a.R = Rt(t); a.R_out = Rt(t); a.apply = true; + a.bo_prev = bo_ + t * N; a.inj_prev = inj2_ + t * HC; + a.w_norm = (const float*) hn->data; + a.w_down = (const uint16_t*) hd->data; + a.w_up = (const uint16_t*) hu->data; + a.w_inject = nullptr; + a.eps = EPS; + a.lo = lo_ + t * g.hc_lr; + a.rs = rs_ + t * HC; + a.inject_out = head_inj_; + a.mixed = head_mixed_ + t * N; + } + fused_gr_read_multi(fa, T, xn_, cs); + } else if (head_mix_multi_enabled() && head_ != nullptr && head_->loaded()) { + // the final mixer, the window's tokens in one read (fused_gr_read_multi without the pending write: the + // last layer's write was done above; its sums are gr_read's) + if (hn->kind != WeightKind::F32 || hd->kind != WeightKind::Bf16InF32 || hu->kind != WeightKind::Bf16InF32) { + err = "verify: the output_hc_* weights have the wrong engine forms"; return false; } + FusedGrArgs fa[kFusedGrMaxT]; + for (int t = 0; t < T; ++t) { + fa[t].R = Rt(t); fa[t].R_out = Rt(t); fa[t].apply = false; + fa[t].w_norm = (const float*) hn->data; fa[t].w_down = (const uint16_t*) hd->data; + fa[t].w_up = (const uint16_t*) hu->data; fa[t].eps = EPS; + fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].mixed = head_mixed_ + t * N; + } + fused_gr_read_multi(fa, T, xn_, cs); + } else { + for (int t = 0; t < T; ++t) { + BlockBuffers bb = ss.block; + bb.R = Rt(t); + bb.mixed = head_mixed_ + t * N; + if (head_ != nullptr && head_->loaded()) { + if (!lm_head_mix(wt, g, bb, cs, err)) return false; + } else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) { + return false; + } + } } } if (head_ != nullptr && head_->loaded()) { @@ -2004,6 +2073,8 @@ bool Verifier::capture_commit(std::string& err) { (int) g.ssm_k_heads, (int) HV, (int) MT, commit_, cs_, (int) MT); // S26: t_out_begin = MT - the replay's outputs (y_dummy_) are never read ++gdn_index; + } else if (!g.has_indexer) { // the K/V cells are by position: nothing to replay without an indexer + ++qsa_index; } else { const QsaState& st = ss.qsa_states[qsa_index]; const WeightRef* wikn = need(v, "indexer.k_norm.weight", err); @@ -2657,6 +2728,8 @@ bool Verifier::capture_commit_batch(const int* rows, int S, int hbase, std::stri t - first > 1 ? t - first : 0); // a 1-row group keeps the 0.1.39 kernel } ++gdn_index; + } else if (!g_->has_indexer) { // the K/V cells are by position: nothing to replay without an indexer + ++qsa_index; } else { const WeightRef* wikn = need(v, "indexer.k_norm.weight", err); if (!wikn) { ok = false; break; } diff --git a/src/kernels/attn_splits_parity.cpp b/src/kernels/attn_splits_parity.cpp new file mode 100644 index 00000000000..0ebfe648aa1 --- /dev/null +++ b/src/kernels/attn_splits_parity.cpp @@ -0,0 +1,173 @@ +// src/kernels/attn_splits_parity.cpp - STRATA_ATTN_SPLITS (split-K decode attention) against the default chunk path +// (GPU, synthetic, no model). +// +// For FP16 and INT8 pools, 8 and 12 query heads per KV head, the single-query and the batched entry points: +// 1. a selection that fits P chunks: the split-K form must give the default's output bit for bit; +// 2. a selection of many more chunks than P: the outputs must agree to float rounding (max |diff| / max |out|), +// and the split-K scratch must be the P-partition size, not the chunk count's. +#include "strata/kernels/f16_bits.hpp" +#include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/qsa.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace k = strata::kernels; + +namespace { +int g_fail = 0; +void ck(cudaError_t e, const char* w) { + if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", w, cudaGetErrorString(e)); std::exit(2); } +} +template T* dalloc(size_t n) { + T* p = nullptr; + ck(cudaMalloc(&p, n * sizeof(T) + 64), "malloc"); + ck(cudaMemset(p, 0, n * sizeof(T) + 64), "memset"); + return p; +} + +constexpr int D = 256, PS = 4, PARTS = 64; + +struct Pools { + k::QsaAttnPools p; + std::vector owned; + ~Pools() { for (void* x : owned) cudaFree(x); } +}; + +// identity page table, cells x n_kv rows of random keys and values; int8 = codes + one fp16 scale per 64 values +void make_pools(Pools& out, int cells, int n_kv, bool int8, std::mt19937& rng) { + std::normal_distribution nd(0.0f, 1.0f); + const int blocks = (cells + PS - 1) / PS; + const size_t rows = (size_t) blocks * n_kv * PS; + std::vector table(blocks); + for (int b = 0; b < blocks; ++b) table[b] = b; + int32_t* d_table = dalloc(blocks); + ck(cudaMemcpy(d_table, table.data(), blocks * 4, cudaMemcpyHostToDevice), "table"); + out.owned.push_back(d_table); + out.p.page_table = d_table; + if (!int8) { + std::vector kh(rows * D), vh(rows * D); + for (auto& x : kh) x = k::f16_from_f32(nd(rng)); + for (auto& x : vh) x = k::f16_from_f32(nd(rng)); + uint16_t *dk = dalloc(kh.size()), *dv = dalloc(vh.size()); + ck(cudaMemcpy(dk, kh.data(), kh.size() * 2, cudaMemcpyHostToDevice), "k"); + ck(cudaMemcpy(dv, vh.data(), vh.size() * 2, cudaMemcpyHostToDevice), "v"); + out.owned.push_back(dk); out.owned.push_back(dv); + out.p.k_pool = dk; out.p.v_pool = dv; + } else { + std::uniform_int_distribution code(-127, 127); + std::vector kq(rows * D), vq(rows * D); + std::vector ks(rows * (D / k::KV_Q8_GROUP)), vs(ks.size()); + for (auto& x : kq) x = (int8_t) code(rng); + for (auto& x : vq) x = (int8_t) code(rng); + for (auto& x : ks) x = k::f16_from_f32(0.01f + 0.01f * std::fabs(nd(rng))); + for (auto& x : vs) x = k::f16_from_f32(0.01f + 0.01f * std::fabs(nd(rng))); + int8_t *dkq = dalloc(kq.size()), *dvq = dalloc(vq.size()); + uint16_t *dks = dalloc(ks.size()), *dvs = dalloc(vs.size()); + ck(cudaMemcpy(dkq, kq.data(), kq.size(), cudaMemcpyHostToDevice), "kq"); + ck(cudaMemcpy(dvq, vq.data(), vq.size(), cudaMemcpyHostToDevice), "vq"); + ck(cudaMemcpy(dks, ks.data(), ks.size() * 2, cudaMemcpyHostToDevice), "ks"); + ck(cudaMemcpy(dvs, vs.data(), vs.size() * 2, cudaMemcpyHostToDevice), "vs"); + out.owned.push_back(dkq); out.owned.push_back(dvq); out.owned.push_back(dks); out.owned.push_back(dvs); + out.p.k_q = dkq; out.p.v_q = dvq; out.p.k_scale = dks; out.p.v_scale = dvs; + } +} + +// one configuration: `cells` of capacity, n_q queries each selecting its own width (all cells, minus a few per query) +void run(const char* name, int cells, int width, int n_kv, int per, bool int8, bool batched, int n_q, bool expect_bitwise) { + std::mt19937 rng(1234u + (unsigned) cells + 7u * per + (int8 ? 1u : 0u)); + std::normal_distribution nd(0.0f, 1.0f); + k::QsaShapes s = k::qsa_real_shapes(); + s.n_head = (int64_t) n_kv * per; s.n_head_kv = n_kv; s.page_size = PS; + s.idx_top_k = k::kDenseTopK; // dense attention, as qwen35moe + Pools pools; + make_pools(pools, cells, n_kv, int8, rng); + const int NH = n_kv * per; + std::vector q((size_t) n_q * NH * D); + for (auto& x : q) x = nd(rng); + std::vector ids((size_t) n_q * cells); + for (int i = 0; i < n_q; ++i) + for (int c = 0; c < cells; ++c) ids[(size_t) i * cells + c] = c; + std::vector steps((size_t) n_q * k::kStepCount); + for (int i = 0; i < n_q; ++i) { + k::qsa_step_fill(&steps[(size_t) i * k::kStepCount], width - 1 - 3 * i, s); + steps[(size_t) i * k::kStepCount + k::kStepWidth] = width - 3 * i; + } + float* d_q = dalloc(q.size()); + int32_t* d_ids = dalloc(ids.size()); + int32_t* d_steps = dalloc(steps.size()); + ck(cudaMemcpy(d_q, q.data(), q.size() * 4, cudaMemcpyHostToDevice), "q"); + ck(cudaMemcpy(d_ids, ids.data(), ids.size() * 4, cudaMemcpyHostToDevice), "ids"); + ck(cudaMemcpy(d_steps, steps.data(), steps.size() * 4, cudaMemcpyHostToDevice), "steps"); + + k::qsa_decode_attn_set_splits(0); + const uint64_t scr_default = k::qsa_decode_attn_scratch_floats(cells, s); + k::qsa_decode_attn_set_splits(PARTS); + const uint64_t scr_split = k::qsa_decode_attn_scratch_floats(cells, s); + const int64_t chunks = (cells + 63) / 64; + const uint64_t want_split = (uint64_t) std::min(chunks, PARTS) * NH * (D + 2) + 64; + if (scr_split != want_split || scr_default != (uint64_t) chunks * NH * (D + 2) + 64) { + std::printf("FAIL %s: scratch %llu (split) / %llu (default), expected %llu\n", name, + (unsigned long long) scr_split, (unsigned long long) scr_default, (unsigned long long) want_split); + ++g_fail; + } + float* d_scr = dalloc(scr_default * n_q); + float *d_a = dalloc(q.size()), *d_b = dalloc(q.size()); + auto call = [&](float* out) { + if (batched) k::qsa_decode_attn_batch(d_q, pools.p, d_ids, d_steps, cells, s, d_scr, out, n_q, nullptr); + else k::qsa_decode_attn_step(d_q, pools.p, d_ids, d_steps, cells, s, d_scr, out, nullptr); + }; + k::qsa_decode_attn_set_splits(0); + call(d_a); + ck(cudaMemset(d_scr, 0xff, scr_default * n_q * 4 + 64), "poison"); // the split run must not read old partials + k::qsa_decode_attn_set_splits(PARTS); + call(d_b); + ck(cudaDeviceSynchronize(), "sync"); + k::qsa_decode_attn_set_splits(0); + const size_t n_out = (size_t) (batched ? n_q : 1) * NH * D; + std::vector a(n_out), b(n_out); + ck(cudaMemcpy(a.data(), d_a, n_out * 4, cudaMemcpyDeviceToHost), "a"); + ck(cudaMemcpy(b.data(), d_b, n_out * 4, cudaMemcpyDeviceToHost), "b"); + size_t diff_bits = 0; + double max_d = 0.0, max_a = 0.0; + bool finite = true; + for (size_t i = 0; i < n_out; ++i) { + if (std::memcmp(&a[i], &b[i], 4) != 0) ++diff_bits; + if (!std::isfinite(b[i])) finite = false; + max_d = std::max(max_d, (double) std::fabs(a[i] - b[i])); + max_a = std::max(max_a, (double) std::fabs(a[i])); + } + const double rel = max_a > 0 ? max_d / max_a : max_d; + const bool ok = finite && (expect_bitwise ? diff_bits == 0 : rel < 2e-5); + std::printf("%s %-34s cells %6d of %6d, chunks %5lld parts %d: %zu of %zu values differ, max rel %.3g, scratch %.2f -> %.2f MB/query\n", + ok ? "PASS" : "FAIL", name, width, cells, (long long) chunks, PARTS, diff_bits, n_out, rel, + scr_default * 4.0 / 1e6, scr_split * 4.0 / 1e6); + if (!ok) ++g_fail; + cudaFree(d_q); cudaFree(d_ids); cudaFree(d_steps); cudaFree(d_scr); cudaFree(d_a); cudaFree(d_b); +} +} // namespace + +int main() { + int n_dev = 0; + if (cudaGetDeviceCount(&n_dev) != cudaSuccess || n_dev == 0) { std::printf("SKIP: no GPU\n"); return 0; } + for (int per : {8, 12}) + for (bool int8 : {false, true}) + for (bool batched : {false, true}) { + char name[64]; + std::snprintf(name, sizeof name, "%s G=%d %s", int8 ? "int8" : "fp16", per, batched ? "batch x3" : "step"); + const int nq = batched ? 3 : 1; + const int big = PARTS * 64 * 9 + 37; + run(name, big, PARTS * 64 - 5, 2, per, int8, batched, nq, true); // the selection fits P chunks: bit for bit + run(name, big, big, 2, per, int8, batched, nq, false); // 9-10 chunks per partition + } + std::printf(g_fail == 0 ? "attn_splits_parity: PASS\n" : "attn_splits_parity: %d FAIL\n", g_fail); + return g_fail == 0 ? 0 : 1; +} diff --git a/src/kernels/cpu/expert_layout.cpp b/src/kernels/cpu/expert_layout.cpp index d8adb260f4d..14f9cecaea7 100644 --- a/src/kernels/cpu/expert_layout.cpp +++ b/src/kernels/cpu/expert_layout.cpp @@ -376,12 +376,20 @@ void native_gu_rows(const NativeFmt&, const uint8_t*, const void* const*, int, f void native_down_rows(const NativeFmt&, const uint8_t*, const void* const*, int, float* const*, int, int) { std::abort(); } #endif -bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err) { +bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err, + int64_t n_embd, int64_t n_ff) { ExpertLayout L; L.n_layers = n_layers; L.n_expert = n_expert; + L.n_embd = n_embd; + L.n_ff = n_ff; std::ifstream in(pack_dir + "/native_experts.txt"); if (!in) { + if (n_embd != H || n_ff != FF) { + err = "a model of n_embd " + std::to_string(n_embd) + " / expert width " + std::to_string(n_ff) + + " needs a native pack (tools/iq_pack.py); the canonical Q2_0 layout is Flash-Next's"; + return false; + } L.total = (uint64_t) n_layers * (uint64_t) n_expert * (uint64_t) BLOB; g_layout = L; return true; @@ -427,7 +435,7 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n return false; } NativeFmt f; - if (!native_fmt((int) gt, (int) dt, H, FF, f, err)) return false; + if (!native_fmt((int) gt, (int) dt, n_embd, n_ff, f, err)) return false; if (f.bytes != blob) { err = "native_experts.txt: layer " + std::to_string(l) + " blob is " + std::to_string(blob) + " B but its formats make " + std::to_string(f.bytes); diff --git a/src/kernels/cpu/pool.cpp b/src/kernels/cpu/pool.cpp index af1eb5a59b8..1a50cb74e70 100644 --- a/src/kernels/cpu/pool.cpp +++ b/src/kernels/cpu/pool.cpp @@ -627,8 +627,9 @@ void ExpertPool::drain(int id, ExpertScratch& scratch, uint32_t epoch) { const int r0 = H * part / parts_b_, r1 = H * (part + 1) / parts_b_; s2_expert_down_rows(jobs_[e].blob, split_[(size_t) e].a2, jobs_[e].out, r0, r1); } else if (mode_ >= 5) { - // plan v0.3 P6: native layers, 5 = gate/up rows, 6 = down rows - const int per = mode_ == 5 ? FF : H; + // plan v0.3 P6: native layers, 5 = gate/up rows, 6 = down rows (the model's widths: Flash-Next's FF / H, + // Qwen3.6's 512 / 2048) + const int per = (int) (mode_ == 5 ? nfmt_->n_ff : nfmt_->n_embd); const int64_t g0 = mrows_ * (int64_t) i / mtasks_, g1 = mrows_ * (int64_t) (i + 1) / mtasks_; for (int64_t r = g0; r < g1;) { const int e = (int) (r / per), r0 = (int) (r % per); @@ -766,17 +767,17 @@ void ExpertPool::run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs const int nb = (std::min)(kMaxSplitMulti, n - b0); mjobs_ = jobs + b0; nfmt_ = &f; - mrows_ = (int64_t) nb * FF; + mrows_ = (int64_t) nb * f.n_ff; mtasks_ = phase_tasks(mrows_); const auto a = std::chrono::steady_clock::now(); run_phase(5, mtasks_); const auto b = std::chrono::steady_clock::now(); for (int e = 0; e < nb; ++e) for (int t = 0; t < mjobs_[e].nt; ++t) - if (q2_native_kernels(f.d_type)) act_quant_any(split_multi_[(size_t) e].ff[t], FF, split_multi_[(size_t) e].a2[t]); + if (q2_native_kernels(f.d_type)) act_quant_any(split_multi_[(size_t) e].ff[t], (int) f.n_ff, split_multi_[(size_t) e].a2[t]); else native_quant_h(f, split_multi_[(size_t) e].ff[t], split_multi_[(size_t) e].hq[t]); const auto c = std::chrono::steady_clock::now(); - mrows_ = (int64_t) nb * H; + mrows_ = (int64_t) nb * f.n_embd; mtasks_ = phase_tasks(mrows_); run_phase(6, mtasks_); const auto d = std::chrono::steady_clock::now(); diff --git a/src/kernels/cuda/iq_kernels.cu b/src/kernels/cuda/iq_kernels.cu index c2e44270f5a..13f44da2e2d 100644 --- a/src/kernels/cuda/iq_kernels.cu +++ b/src/kernels/cuda/iq_kernels.cu @@ -592,6 +592,61 @@ template<> struct Fmt<12> { static constexpr int qk = 256, ipb = QI4_K / VDR_Q4_ __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q4_K_q8_1(v, y, kbx, iqs); } }; template<> struct Fmt<13> { static constexpr int qk = 256, ipb = QI5_K / VDR_Q5_K, step = VDR_Q5_K; __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q5_K_q8_1(v, y, kbx, iqs); } }; +// Q3_K (the gate/up of Unsloth's Qwen3.6 MTP layer): llama.cpp's vec_dot_q3_K_q8_1 (vecdotq.cuh, VDR 1), transcribed +constexpr int VDR_Q3_K = 1; +__device__ __forceinline__ float vec_dot_q3_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q3_K* bq3_K = (const block_q3_K*) vbq + kbx; + const int bq8_offset = QR3_K * (iqs / (QI3_K / 2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + const float d3 = __half2float(bq3_K->d); + const int vl = get_int_b2(bq3_K->qs, iqs); + const int vh = ~get_int_b2(bq3_K->hmask, iqs % (QI3_K / 2)) >> bq8_offset; // a 0/1 bit subtracts 4/0 + float sumf = 0.0f; +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int u = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + const float d8 = __low2float(bq8_1[bq8_offset + i].ds); + const int isc = scale_offset + 2 * i; + const int sc_low = (bq3_K->scales[isc % (QK_K / 32)] >> (4 * (isc / (QK_K / 32)))) & 0xF; + const int sc_high = ((bq3_K->scales[(QK_K / 32) + isc % (QK_K / 64)] >> (2 * (isc / (QK_K / 64)))) & 3) << 4; + const int sc = (sc_low | sc_high) - 32; + const int vil = (vl >> (2 * i)) & 0x03030303; + const int vih = ((vh >> i) << 2) & 0x04040404; + const int vi = __vsubss4(vil, vih); + sumf += d8 * (ggml_cuda_dp4a(vi, u, 0) * sc); + } + return d3 * sumf; +} +// Q2_K (the gate/up of the smaller Unsloth Qwen3.6 files' MTP layer): llama.cpp's vec_dot_q2_K_q8_1 (VDR 1), transcribed +constexpr int VDR_Q2_K = 1; +__device__ __forceinline__ float vec_dot_q2_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q2_K* bq2_K = (const block_q2_K*) vbq + kbx; + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + const uint8_t* scales = bq2_K->scales + scale_offset; + const int v = get_int_b4(bq2_K->qs, iqs); + float sumf_d = 0.0f, sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int u = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + const float d8 = __low2float(bq8_1[bq8_offset + i].ds); + const int sc = scales[2 * i]; + const int vi = (v >> (2 * i)) & 0x03030303; + sumf_d += d8 * (ggml_cuda_dp4a(vi, u, 0) * (sc & 0xF)); + int m = sc >> 4; // the min, four times + m |= m << 8; + m |= m << 16; + sumf_m += d8 * ggml_cuda_dp4a(m, u, 0); + } + const float2 dm2f = __half22float2(bq2_K->dm); + return dm2f.x * sumf_d - dm2f.y * sumf_m; +} +template<> struct Fmt<10> { static constexpr int qk = 256, ipb = QI2_K / VDR_Q2_K, step = VDR_Q2_K; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q2_K_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<11> { static constexpr int qk = 256, ipb = QI3_K / VDR_Q3_K, step = VDR_Q3_K; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q3_K_q8_1(v, y, kbx, iqs); } }; template<> struct Fmt<14> { static constexpr int qk = 256, ipb = QI6_K / VDR_Q6_K, step = VDR_Q6_K; __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q6_K_q8_1(v, y, kbx, iqs); } }; template<> struct Fmt<7> { static constexpr int qk = 32, ipb = QI5_1 / VDR_Q5_1, step = VDR_Q5_1; @@ -608,12 +663,14 @@ template<> struct Fmt<8> { static constexpr int qk = 32, ipb = QI8_0 / VDR_Q8_0, // The formats of each role, one list each so a type cannot be in one switch and missing from another. Every // entry is a kernel template for each CUDA architecture of the build, hence two lists rather than one. #ifdef STRATA_Q6K_EXPERTS // opt-in build (-DSTRATA_Q6K_EXPERTS=ON): one more instance per kernel, loaded at start -#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(14) X(6) X(2) X(3) X(8) +#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(14) X(6) X(2) X(3) X(8) X(11) X(10) #else -#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(6) X(2) X(3) X(8) +#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(6) X(2) X(3) X(8) X(11) X(10) #endif -#define STRATA_D_FMTS(X) X(20) X(23) X(42) X(7) X(6) X(2) X(3) X(8) -#define STRATA_MMVQ_FMTS(X) X(16) X(17) X(18) X(20) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(7) X(6) X(2) X(3) X(8) +// Unsloth's Qwen3.6 UD files (n_ff 512 = two superblocks per row): Q6_K (the few layers kept at 6 bits), IQ3_XXS +// (UD-IQ2_M), Q5_K (UD-Q4_K_XL), and their MTP layers' Q3_K / Q4_K / Q5_K (with Q2_K / Q3_K gate/up above) +#define STRATA_D_FMTS(X) X(20) X(23) X(42) X(7) X(6) X(2) X(3) X(8) X(14) X(12) X(13) X(18) X(11) +#define STRATA_MMVQ_FMTS(X) X(16) X(17) X(18) X(20) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(7) X(6) X(2) X(3) X(8) X(11) X(10) __device__ __forceinline__ float warp_sum(float v) { #pragma unroll @@ -1687,6 +1744,22 @@ __device__ void dq_q3_k(const void* vx, int64_t ibs, dst_t* yy, int tid) { for (int l = l0; l < l0 + 4; ++l) y[l] = cvt(dl * ((int8_t) ((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4))); } } +// Q2_K (the gate/up of the smaller Unsloth Qwen3.6 files' MTP layer): llama.cpp's dequantize_q2_K, its 64 threads +// folded onto 32 +template +__device__ void dq_q2_k(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_q2_K* x = (const block_q2_K*) vx + ibs; + const float dall = __low2float(x->dm), dmin = __high2float(x->dm); + for (int tt = tid; tt < 64; tt += 32) { + const int n = tt / 32, l = tt - 32 * n, is = 8 * n + l / 16; + const uint8_t q = x->qs[32 * n + l]; + dst_t* y = yy + 128 * n; + y[l + 0] = cvt(dall * (x->scales[is + 0] & 0xF) * ((q >> 0) & 3) - dmin * (x->scales[is + 0] >> 4)); + y[l + 32] = cvt(dall * (x->scales[is + 2] & 0xF) * ((q >> 2) & 3) - dmin * (x->scales[is + 2] >> 4)); + y[l + 64] = cvt(dall * (x->scales[is + 4] & 0xF) * ((q >> 4) & 3) - dmin * (x->scales[is + 4] >> 4)); + y[l + 96] = cvt(dall * (x->scales[is + 6] & 0xF) * ((q >> 6) & 3) - dmin * (x->scales[is + 6] >> 4)); + } +} template __device__ void dq_iq4_xs(const void* vx, int64_t ibs, dst_t* yy, int tid) { const block_iq4_xs* x = (const block_iq4_xs*) vx + ibs; @@ -1882,12 +1955,11 @@ __device__ __forceinline__ void dq_dispatch(int ty, const void* vx, int64_t ibs, case 29: dq_iq1_m(vx, ibs, y, tid); break; case 23: dq_iq4_xs(vx, ibs, y, tid); break; case 11: dq_q3_k(vx, ibs, y, tid); break; + case 10: dq_q2_k(vx, ibs, y, tid); break; case 42: dq_q2_0(vx, ibs, y, tid); break; case 12: dq_q4_k(vx, ibs, y, tid); break; case 13: dq_q5_k(vx, ibs, y, tid); break; -#ifdef STRATA_Q6K_EXPERTS case 14: dq_q6_k(vx, ibs, y, tid); break; -#endif case 7: dq_q5_1(vx, ibs, y, tid); break; case 6: dq_q5_0(vx, ibs, y, tid); break; case 2: dq_q4_0(vx, ibs, y, tid); break; @@ -1915,12 +1987,8 @@ __global__ void dequant_gu_kernel(int ty, const void* __restrict__ gate, const v // the types dq_dispatch dequantizes bool is_iq(int t) { - return t == 16 || t == 17 || t == 18 || t == 20 || t == 21 || t == 22 || t == 23 || t == 29 || t == 42 || t == 11 || -#ifdef STRATA_Q6K_EXPERTS + return t == 16 || t == 17 || t == 18 || t == 20 || t == 21 || t == 22 || t == 23 || t == 29 || t == 42 || t == 11 || t == 10 || t == 12 || t == 13 || t == 14 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8; -#else - t == 12 || t == 13 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8; -#endif } // values per block of the types the grouped expert kernels take (0 = none) int gu_qk(int t) { @@ -2712,6 +2780,7 @@ size_t iq_row_bytes(int t, int64_t n) noexcept { case 29: return (size_t) (n / 256) * sizeof(block_iq1_m); case 23: return (size_t) (n / 256) * sizeof(block_iq4_xs); case 11: return (size_t) (n / 256) * sizeof(block_q3_K); + case 10: return (size_t) (n / 256) * sizeof(block_q2_K); case 42: return (size_t) (n / 64) * sizeof(block_q2_0); case 12: return (size_t) (n / 256) * sizeof(block_q4_K); case 13: return (size_t) (n / 256) * sizeof(block_q5_K); @@ -2789,7 +2858,8 @@ void iq_dequant_gu_f16(int t, const void* gate, const void* up, int64_t n_ff, in bool native_expert_supported(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) noexcept { const int qg = gu_qk(gu_type), qd = d_qk(d_type); - return qg > 0 && qd > 0 && is_iq(gu_type) && is_iq(d_type) && n_embd % qg == 0 && n_ff % qd == 0 && + return qg > 0 && qd > 0 && is_iq(gu_type) && is_iq(d_type) && iq_row_bytes(gu_type, n_embd) > 0 && + iq_row_bytes(d_type, n_ff) > 0 && n_embd % qg == 0 && n_ff % qd == 0 && n_embd % 256 == 0 && (n_ff * n_embd) % 256 == 0; } diff --git a/src/kernels/cuda/kv_stream.cu b/src/kernels/cuda/kv_stream.cu index 209984b4ff1..c1ea27d1298 100644 --- a/src/kernels/cuda/kv_stream.cu +++ b/src/kernels/cuda/kv_stream.cu @@ -87,7 +87,8 @@ __device__ int block_scan(int v, int* warp_sums, int& total) { } __global__ void __launch_bounds__(RT) resolve_kernel(KvStreamMap m, const int32_t* __restrict__ ids, - const int32_t* __restrict__ steps, int n_q, int cap, int page_size) { + const int32_t* __restrict__ steps, int n_q, int cap, int page_size, + int dense) { __shared__ int s_nmiss, s_lookups, s_cut; __shared__ int warp_sums[32]; const int epoch = m.ctl[0] + 1; @@ -98,16 +99,23 @@ __global__ void __launch_bounds__(RT) resolve_kernel(KvStreamMap m, const int32_ for (int q = 0; q < n_q; ++q) { const int width = steps[q * kStepCount + kStepWidth]; const int32_t* qi = ids + (long long) q * cap; - for (int i = threadIdx.x; i < width; i += RT) { - const int b = qi[i] / page_size; - if (i > 0 && qi[i - 1] / page_size == b) continue; // ids are ascending: one lookup per block + // dense selection (qwen35moe: ids are 0..width-1): its blocks are 0..ceil(width / page_size)-1, so the threads + // walk the blocks, not every cell - the same lookups, width / page_size loads instead of width + const int span = dense ? (width + page_size - 1) / page_size : width; + for (int i = threadIdx.x; i < span; i += RT) { + const int b = dense ? i : qi[i] / page_size; + if (!dense && i > 0 && qi[i - 1] / page_size == b) continue; // ids are ascending: one lookup per block ++lookups; const int sl = m.page_table[b]; if (sl >= 0) { m.slot_stamp[sl] = epoch; m.slot_ref[sl] = 1; } else if (sl == -1 && atomicCAS(&m.page_table[b], -1, -2) == -1) { - m.miss_block[atomicAdd(&s_nmiss, 1)] = b; + // the miss list holds n_slots: a dense selection (qwen35moe, every cell) can miss more blocks than + // that, and the readers take those from the host copy in place (qsa_decode_attn's cell_row) + const int k = atomicAdd(&s_nmiss, 1); + if (k < (int) m.n_slots) m.miss_block[k] = b; + else m.page_table[b] = -1; } } } @@ -115,7 +123,7 @@ __global__ void __launch_bounds__(RT) resolve_kernel(KvStreamMap m, const int32_ __syncthreads(); // 2. one victim per miss: a clock sweep from the hand. A slot this call uses (stamp == epoch) is never taken; // a referenced one loses its bit as the hand passes it and is taken on the next pass. - const int need = s_nmiss, n = (int) m.n_slots; + const int n = (int) m.n_slots, missed = s_nmiss, need = missed < n ? missed : n; int hand = m.ctl[1], got = 0; for (int scanned = 0; got < need && scanned < 3 * n; scanned += RT) { const int j = (int) (((long long) hand + threadIdx.x) % n); @@ -143,7 +151,7 @@ __global__ void __launch_bounds__(RT) resolve_kernel(KvStreamMap m, const int32_ const int placed = got < need ? got : need; for (int k = threadIdx.x; k < need; k += RT) { const int b = m.miss_block[k]; - if (k >= placed) { m.page_table[b] = -1; continue; } // overflow: never happens with a legal n_slots + if (k >= placed) { m.page_table[b] = -1; continue; } // no slot: read from the host copy in place const int sl = m.miss_slot[k]; const int old = m.slot_block[sl]; if (old >= 0) m.page_table[old] = -1; @@ -156,9 +164,9 @@ __global__ void __launch_bounds__(RT) resolve_kernel(KvStreamMap m, const int32_ m.ctl[0] = epoch; m.ctl[1] = hand; m.ctl[2] = placed; - if (placed < need) m.ctl[3] = 1; + if (placed < missed) m.ctl[3] = 1; unsigned long long* c = reinterpret_cast(m.ctl + 4); - c[0] += (unsigned long long) placed; + c[0] += (unsigned long long) missed; // from RAM: copied into a slot, or read in place c[1] += (unsigned long long) s_lookups; c[2] += 1ull; } @@ -213,7 +221,7 @@ void kv_stream_reset(const KvStreamMap& m, void* stream) { void kv_stream_resolve(const KvStreamMap& m, const QsaAttnPools& slots, const KvHostPools& host, int fmt, const int32_t* ids, const int32_t* steps, int64_t n_q, int64_t cap, const QsaShapes& s, - void* stream) { + void* stream, bool dense) { if (n_q <= 0) return; if (s.n_head_kv * s.page_size * (s.head_dim / KV_Q8_GROUP) * 2 % 16 != 0) { std::fprintf(stderr, "kv_stream: a block's scale run must be a multiple of 16 bytes\n"); @@ -227,7 +235,8 @@ void kv_stream_resolve(const KvStreamMap& m, const QsaAttnPools& slots, const Kv "slot twice\n", (long long) m.n_slots, RT); std::exit(1); } - resolve_kernel<<<1, RT, 0, (cudaStream_t) stream>>>(m, ids, steps, (int) n_q, (int) cap, (int) s.page_size); + resolve_kernel<<<1, RT, 0, (cudaStream_t) stream>>>(m, ids, steps, (int) n_q, (int) cap, (int) s.page_size, + dense ? 1 : 0); check("resolve"); copy_kernel<<<96, 128, 0, (cudaStream_t) stream>>>(m, runs_of(slots, host, fmt, s)); check("copy"); diff --git a/src/kernels/cuda/qsa.cu b/src/kernels/cuda/qsa.cu index a635e81e0e2..f696c6240ac 100644 --- a/src/kernels/cuda/qsa.cu +++ b/src/kernels/cuda/qsa.cu @@ -88,8 +88,9 @@ void check_sync(const char* what) { /// Geometry validation. A silently wrong `head_dim % 4` would corrupt the gather's uint2 copy and a silently /// wrong `n_head % n_head_kv` would produce a plausible attention with the wrong key, so both are refused. void validate(const QsaShapes& s, const char* who) { - if (s.n_head <= 0 || s.n_head_kv <= 0 || s.head_dim <= 0 || s.idx_dim <= 0 || s.idx_n_head <= 0 || - s.idx_block < 2 || s.page_size < 1) { + // dense attention (qwen35moe) has no indexer: its indexer sizes are 0 and nothing reads them + const bool idx_ok = s.idx_top_k == kDenseTopK || (s.idx_dim > 0 && s.idx_n_head > 0); + if (s.n_head <= 0 || s.n_head_kv <= 0 || s.head_dim <= 0 || !idx_ok || s.idx_block < 2 || s.page_size < 1) { std::fprintf(stderr, "qsa: %s: geometry is not set up\n", who); std::exit(1); } @@ -103,7 +104,7 @@ void validate(const QsaShapes& s, const char* who) { (long long) s.head_dim); std::exit(1); } - if (s.n_rot <= 0 || s.n_rot % 2 != 0 || s.n_rot > s.head_dim || s.n_rot > s.idx_dim) { + if (s.n_rot <= 0 || s.n_rot % 2 != 0 || s.n_rot > s.head_dim || (s.idx_top_k != kDenseTopK && s.n_rot > s.idx_dim)) { std::fprintf(stderr, "qsa: %s: n_rot %lld must be even and <= head_dim %lld and idx_dim %lld\n", who, (long long) s.n_rot, (long long) s.head_dim, (long long) s.idx_dim); std::exit(1); diff --git a/src/kernels/cuda/qsa_decode_attn.cu b/src/kernels/cuda/qsa_decode_attn.cu index 624036c6155..18a7bf3ac0a 100644 --- a/src/kernels/cuda/qsa_decode_attn.cu +++ b/src/kernels/cuda/qsa_decode_attn.cu @@ -18,7 +18,8 @@ namespace strata::kernels { namespace { constexpr int HD = 256; // head_dim -constexpr int G = 12; // query heads per KV head (24 / 2) +constexpr int G_FLASH = 12; // query heads per KV head: Flash-Next 24 / 2 (the kernels' G; Qwen3.6's 16 / 2 = 8 is + // their other instance) constexpr int CHUNK = 64; // cells per block constexpr int THREADS = 256; constexpr int WARPS = THREADS / 32; @@ -84,6 +85,54 @@ __device__ __forceinline__ void load8(const QsaAttnPools& p, bool value, long lo } else load8_q4(p, value, row, d0, out); } +// A selected cell's row (`srow`): >= 0 a VRAM pool row; -1 masked (past the selection, or not resident in a layer +// without a host copy); <= -2 the host copy's row -2 - r (KV streaming: a block the resolve could not make resident, +// read in place over PCIe). +__device__ __forceinline__ long long cell_row(const QsaAttnPools& p, int cell, int kvh, int n_kv_heads, int page_size) { + const long long block = cell / page_size, page = (long long) p.page_table[block]; + if (page >= 0) return (page * n_kv_heads + kvh) * page_size + (cell % page_size); + if (p.host_k_pool == nullptr && p.host_k_q == nullptr && p.host_k_q4 == nullptr) return -1; + return -2 - ((block * n_kv_heads + kvh) * page_size + (cell % page_size)); +} +// One V element of pool row `row`, dimension `d` (the value loop's own thread index), per KV format. +template +__device__ __forceinline__ float load_v1(const QsaAttnPools& p, long long row, int d) { + if constexpr (KV_MODE == 0) { + return __half2float(__ushort_as_half(p.v_pool[row * HD + d])); + } else if constexpr (KV_MODE == 1) { + const float sc = __half2float(__ushort_as_half(p.v_scale[row * (HD / KV_Q8_GROUP) + d / KV_Q8_GROUP])); + return (float) p.v_q[row * HD + d] * sc; + } else { // modes 2 and 3: V is rotated Q4_0 (kv_q4.hpp); the caller rotates the output back + constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); + const int b = d / QK4_0; + const int rem = d % QK4_0; + const block_q4_0* blk = reinterpret_cast(p.v_q4 + row * bytes_per_head) + b; + const float dd = __half2float(__ushort_as_half(blk->d)); + const int j = rem < 16 ? rem : (rem - 16); + const uint8_t byte = blk->qs[j]; + const int nibble = (rem < 16) ? ((byte & 0x0F) - 8) : ((byte >> 4) - 8); + return (float) nibble * dd; + } +} +// The host copy as pools, for a row <= -2 (its own row: -2 - r). +__device__ __forceinline__ QsaAttnPools host_pools(const QsaAttnPools& p) { + QsaAttnPools h; + h.k_pool = p.host_k_pool; h.v_pool = p.host_v_pool; + h.k_q = p.host_k_q; h.v_q = p.host_v_q; + h.k_scale = p.host_k_scale; h.v_scale = p.host_v_scale; + h.k_q4 = p.host_k_q4; h.v_q4 = p.host_v_q4; + return h; +} +template +__device__ __forceinline__ void load8r(const QsaAttnPools& p, bool value, long long row, int d0, float* out) { + if (row >= 0) load8(p, value, row, d0, out); + else load8(host_pools(p), value, -2 - row, d0, out); +} +template +__device__ __forceinline__ float load_v1r(const QsaAttnPools& p, long long row, int d) { + return row >= 0 ? load_v1(p, row, d) : load_v1(host_pools(p), -2 - row, d); +} + // Exact-intent score reduce-scatter: the same butterfly pairs as twelve warp_sum calls, // with 16 shuffle exchanges instead of 60. Used only by the opt-in gfx906 INT8 // query-swizzled specialization; softmax, value accumulation and merge are unchanged. @@ -114,7 +163,7 @@ __device__ __forceinline__ int reduce12_head_of_lane(int lane) { return ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); } -template +template __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __restrict__ q, QsaAttnPools p, const int32_t* __restrict__ ids, const int32_t* __restrict__ step, int n_kv_heads, @@ -152,13 +201,9 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest } if (t < CHUNK) { long long r = -1; - if (t < n_here) { - const int cell = ids[c0 + t]; - const long long page = (long long) p.page_table[cell / page_size]; - // a block the KV streaming could not make resident keeps page -1 (ctl[3]); its cells are masked - // (score -FLT_MAX, weight 0) instead of being read from before the pool. - if (page >= 0) r = (page * n_kv_heads + kvh) * page_size + (cell % page_size); - } + // a block the KV streaming could not make resident (page -1) is read from the host copy in place; with no host + // copy its cells are masked (score -FLT_MAX, weight 0) instead of being read from before the pool + if (t < n_here) r = cell_row(p, ids[c0 + t], kvh, n_kv_heads, page_size); srow[t] = r; } __syncthreads(); @@ -169,7 +214,7 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest // reductions per cell constexpr int HPT = G * CHUNK / THREADS; // heads per thread (3) const int c = t % CHUNK, h0 = (t / CHUNK) * HPT; - if (c >= n_here || srow[c] < 0) { + if (c >= n_here || srow[c] == -1) { #pragma unroll for (int j = 0; j < HPT; ++j) sp[h0 + j][c] = -FLT_MAX; } else { @@ -177,7 +222,7 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest #pragma unroll for (int l = 0; l < 32; ++l) { float k8[8]; - load8(p, false, srow[c], l * 8, k8); + load8r(p, false, srow[c], l * 8, k8); #pragma unroll for (int j = 0; j < HPT; ++j) { const float4 qa = *reinterpret_cast(&sq[h0 + j][l * 8]); @@ -199,13 +244,13 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest } else // scores: each warp takes cells warp, warp+8, ...; each lane holds 8 of the 256 dimensions. for (int c = warp; c < CHUNK; c += WARPS) { - if (c >= n_here || srow[c] < 0) { + if (c >= n_here || srow[c] == -1) { if (lane < G) sp[lane][c] = -FLT_MAX; continue; } if constexpr (REDUCE12) { float k8[8]; - load8(p, false, srow[c], lane * 8, k8); + load8r(p, false, srow[c], lane * 8, k8); float part[16]; #pragma unroll for (int h = 0; h < G; ++h) { @@ -223,7 +268,7 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest if ((lane & 1) == 0 && head < G) sp[head][c] = score * scale; } else { float k8[8]; - load8(p, false, srow[c], lane * 8, k8); + load8r(p, false, srow[c], lane * 8, k8); #pragma unroll for (int h = 0; h < G; ++h) { const int qflip = QUERY_SWIZZLE ? (lane & 4) : 0; @@ -241,8 +286,8 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest for (int h = warp; h < G; h += WARPS) { const float a = sp[h][lane], b = sp[h][lane + 32]; const float m = warp_max(fmaxf(a, b)); - const float ea = (lane < n_here && srow[lane] >= 0) ? __expf(a - m) : 0.0f; - const float eb = (lane + 32 < n_here && srow[lane + 32] >= 0) ? __expf(b - m) : 0.0f; + const float ea = (lane < n_here && srow[lane] != -1) ? __expf(a - m) : 0.0f; + const float eb = (lane + 32 < n_here && srow[lane + 32] != -1) ? __expf(b - m) : 0.0f; sp[h][lane] = ea; sp[h][lane + 32] = eb; const float l = warp_sum(ea + eb); @@ -254,24 +299,8 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest #pragma unroll for (int h = 0; h < G; ++h) acc[h] = 0.0f; for (int c = 0; c < n_here; ++c) { - if (srow[c] < 0) continue; // masked above, weight 0 - float v; - if constexpr (KV_MODE == 0) { - v = __half2float(__ushort_as_half(p.v_pool[srow[c] * HD + t])); - } else if constexpr (KV_MODE == 1) { - const float sc = __half2float(__ushort_as_half(p.v_scale[srow[c] * (HD / KV_Q8_GROUP) + t / KV_Q8_GROUP])); - v = (float) p.v_q[srow[c] * HD + t] * sc; - } else { // modes 2 and 3: V is rotated Q4_0 (kv_q4.hpp); the caller rotates the output back - constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); - const int b = t / QK4_0; - const int rem = t % QK4_0; - const block_q4_0* blk = reinterpret_cast(p.v_q4 + srow[c] * bytes_per_head) + b; - const float d = __half2float(__ushort_as_half(blk->d)); - const int j = rem < 16 ? rem : (rem - 16); - const uint8_t byte = blk->qs[j]; - const int nibble = (rem < 16) ? ((byte & 0x0F) - 8) : ((byte >> 4) - 8); - v = (float) nibble * d; - } + if (srow[c] == -1) continue; // masked above, weight 0 + const float v = load_v1r(p, srow[c], t); #pragma unroll for (int h = 0; h < G; ++h) acc[h] = fmaf(sp[h][c], v, acc[h]); } @@ -343,28 +372,8 @@ __device__ __forceinline__ int head_of_lane(int lane) { return ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); } -// One V element of pool row `row`, dimension `d` (the value loop's own thread index), per KV format. -template -__device__ __forceinline__ float load_v1(const QsaAttnPools& p, long long row, int d) { - if constexpr (KV_MODE == 0) { - return __half2float(__ushort_as_half(p.v_pool[row * HD + d])); - } else if constexpr (KV_MODE == 1) { - const float sc = __half2float(__ushort_as_half(p.v_scale[row * (HD / KV_Q8_GROUP) + d / KV_Q8_GROUP])); - return (float) p.v_q[row * HD + d] * sc; - } else { // modes 2 and 3: V is rotated Q4_0 (kv_q4.hpp); the caller rotates the output back - constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); - const int b = d / QK4_0; - const int rem = d % QK4_0; - const block_q4_0* blk = reinterpret_cast(p.v_q4 + row * bytes_per_head) + b; - const float dd = __half2float(__ushort_as_half(blk->d)); - const int j = rem < 16 ? rem : (rem - 16); - const uint8_t byte = blk->qs[j]; - const int nibble = (rem < 16) ? ((byte & 0x0F) - 8) : ((byte >> 4) - 8); - return (float) nibble * dd; - } -} -template +template __global__ void __launch_bounds__(THREADS) attn_chunk_kernel_pre75(const float* __restrict__ q, QsaAttnPools p, const int32_t* __restrict__ ids, const int32_t* __restrict__ step, int n_kv_heads, @@ -399,13 +408,9 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel_pre75(const float* for (int i = t; i < G * HD; i += THREADS) sq[i / HD][i % HD] = q[(size_t) (kvh * G) * HD + i]; if (t < CHUNK) { long long r = -1; - if (t < n_here) { - const int cell = ids[c0 + t]; - const long long page = (long long) p.page_table[cell / page_size]; - // a block the KV streaming could not make resident keeps page -1 (ctl[3]); its cells are masked - // (score -FLT_MAX, weight 0) instead of being read from before the pool. - if (page >= 0) r = (page * n_kv_heads + kvh) * page_size + (cell % page_size); - } + // a block the KV streaming could not make resident (page -1) is read from the host copy in place; with no host + // copy its cells are masked (score -FLT_MAX, weight 0) instead of being read from before the pool + if (t < n_here) r = cell_row(p, ids[c0 + t], kvh, n_kv_heads, page_size); srow[t] = r; } __syncthreads(); @@ -421,12 +426,12 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel_pre75(const float* #pragma unroll for (int j = 0; j < NC; ++j) { cc[j] = warp + WARPS * (i + j); - all_ok = all_ok && cc[j] < n_here && srow[cc[j]] >= 0; + all_ok = all_ok && cc[j] < n_here && srow[cc[j]] != -1; } if (all_ok) { float kk[NC][8]; #pragma unroll - for (int j = 0; j < NC; ++j) load8(p, false, srow[cc[j]], lane * 8, kk[j]); + for (int j = 0; j < NC; ++j) load8r(p, false, srow[cc[j]], lane * 8, kk[j]); float pp[NC][16]; #pragma unroll for (int h = 0; h < G; ++h) { @@ -447,12 +452,12 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel_pre75(const float* } else { for (int k = 0; k < NC; ++k) { // a masked or out-of-range cell in the group: one cell at a time const int c = cc[k]; - if (c >= n_here || srow[c] < 0) { + if (c >= n_here || srow[c] == -1) { if (lane < G) sp[lane][c] = -FLT_MAX; continue; } float k8[8]; - load8(p, false, srow[c], lane * 8, k8); + load8r(p, false, srow[c], lane * 8, k8); float part[16]; #pragma unroll for (int h = 0; h < G; ++h) { @@ -473,8 +478,8 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel_pre75(const float* for (int h = warp; h < G; h += WARPS) { const float a = sp[h][lane], b = sp[h][lane + 32]; const float m = warp_max_x(fmaxf(a, b)); - const float ea = (lane < n_here && srow[lane] >= 0) ? __expf(a - m) : 0.0f; - const float eb = (lane + 32 < n_here && srow[lane + 32] >= 0) ? __expf(b - m) : 0.0f; + const float ea = (lane < n_here && srow[lane] != -1) ? __expf(a - m) : 0.0f; + const float eb = (lane + 32 < n_here && srow[lane + 32] != -1) ? __expf(b - m) : 0.0f; spt[lane][h] = ea; spt[lane + 32][h] = eb; const float l = warp_sum_x(ea + eb); @@ -511,18 +516,25 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel_pre75(const float* fold(c + 3, v3); } for (; c < n_here; ++c) { - if (srow[c] < 0) continue; // masked above, weight 0 - fold(c, load_v1(p, srow[c], t)); + if (srow[c] == -1) continue; // masked above, weight 0 + fold(c, load_v1r(p, srow[c], t)); } #pragma unroll for (int h = 0; h < G; ++h) part_acc[((size_t) slot * G + h) * HD + t] = acc[h]; } #endif // STRATA_ATTN_PRE75_BUILT +template __global__ void __launch_bounds__(HD) attn_merge_kernel(const float* __restrict__ part_acc, const float* __restrict__ part_m, const float* __restrict__ part_l, int n_chunks, - float* __restrict__ attn, long long scratch_stride = 0) { + float* __restrict__ attn, long long scratch_stride = 0, + const int32_t* __restrict__ step = nullptr) { + // Only the chunks that hold selected cells: the grid covers the capacity (max context), but every chunk past + // step[kStepWidth] wrote m = -FLT_MAX and would be skipped below anyway - the result is the same bits, the loop + // no longer walks max_context / CHUNK partials on every token of a short conversation. + const int n_used = step == nullptr ? n_chunks + : min(n_chunks, (__ldg(step + (size_t) blockIdx.y * kStepCount + kStepWidth) + CHUNK - 1) / CHUNK); part_acc += (size_t) blockIdx.y * (size_t) scratch_stride; part_m += (size_t) blockIdx.y * (size_t) scratch_stride; part_l += (size_t) blockIdx.y * (size_t) scratch_stride; @@ -531,9 +543,9 @@ __global__ void __launch_bounds__(HD) attn_merge_kernel(const float* __restrict_ const int kvh = h / G, hl = h % G; const int d = threadIdx.x; float M = -FLT_MAX; - for (int c = 0; c < n_chunks; ++c) M = fmaxf(M, part_m[(kvh * n_chunks + c) * G + hl]); + for (int c = 0; c < n_used; ++c) M = fmaxf(M, part_m[(kvh * n_chunks + c) * G + hl]); float L = 0.0f, acc = 0.0f; - for (int c = 0; c < n_chunks; ++c) { + for (int c = 0; c < n_used; ++c) { const int slot = kvh * n_chunks + c; const float m = part_m[slot * G + hl]; if (m == -FLT_MAX) continue; @@ -544,6 +556,132 @@ __global__ void __launch_bounds__(HD) attn_merge_kernel(const float* __restrict_ attn[(size_t) h * HD + d] = L > 0.0f ? acc / L : 0.0f; } +// STRATA_ATTN_SPLITS=P (opt-in): split-K over a fixed P partitions per KV head instead of one partial per 64-cell +// chunk. Partition b reads chunks b, b + P, b + 2P, ... with attn_chunk_kernel's exact per-chunk math and folds each +// into a running (max, sum, acc) per head; the merge then walks P partials. The partials no longer grow with the +// context: at 262144 cells and 16 query heads one query's scratch is 4096 -> P chunk slots (67.6 -> 4.2 MB at P=256). +// While a selection fits P chunks (P x 64 cells) every partition holds at most one chunk, its partial is the chunk +// kernel's bit for bit and the output is the default's; past that the fold's rescale rounds differently (same math). +template +__global__ void __launch_bounds__(THREADS) attn_split_kernel(const float* __restrict__ q, QsaAttnPools p, + const int32_t* __restrict__ ids, + const int32_t* __restrict__ step, int n_kv_heads, + int page_size, float scale, float* __restrict__ part_acc, + float* __restrict__ part_m, float* __restrict__ part_l, + int n_parts, int cap = 0, long long scratch_stride = 0) { + q += (size_t) blockIdx.z * (size_t) (n_kv_heads * G) * HD; + ids += (size_t) blockIdx.z * (size_t) cap; + step += (size_t) blockIdx.z * kStepCount; + part_acc += (size_t) blockIdx.z * (size_t) scratch_stride; + part_m += (size_t) blockIdx.z * (size_t) scratch_stride; + part_l += (size_t) blockIdx.z * (size_t) scratch_stride; + __shared__ __align__(16) float sq[G][HD]; + __shared__ float sp[G][CHUNK]; + __shared__ long long srow[CHUNK]; + __shared__ float run_m[G], run_l[G]; // the running max and exp-sum per head + __shared__ float w_run[G], w_new[G]; // this fold's rescale of the running sum and of the new chunk + const int n_ids = __ldg(step + kStepWidth); + const int part = blockIdx.x, kvh = blockIdx.y; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5; + const int slot = kvh * n_parts + part; + if (part * CHUNK >= n_ids) { + if (t < G) { part_m[slot * G + t] = -FLT_MAX; part_l[slot * G + t] = 0.0f; } + return; + } + for (int i = t; i < G * HD; i += THREADS) sq[i / HD][i % HD] = q[(size_t) (kvh * G) * HD + i]; + float acc[G]; +#pragma unroll + for (int h = 0; h < G; ++h) acc[h] = 0.0f; + bool first = true; + for (int chunk = part; chunk * CHUNK < n_ids; chunk += n_parts) { + const int c0 = chunk * CHUNK; + const int n_here = min(CHUNK, n_ids - c0); + __syncthreads(); // the previous chunk's sp / srow / weights are read + if (t < CHUNK) { + long long r = -1; + if (t < n_here) r = cell_row(p, ids[c0 + t], kvh, n_kv_heads, page_size); + srow[t] = r; + } + __syncthreads(); + for (int c = warp; c < CHUNK; c += WARPS) { + if (c >= n_here || srow[c] == -1) { + if (lane < G) sp[lane][c] = -FLT_MAX; + continue; + } + float k8[8]; + load8r(p, false, srow[c], lane * 8, k8); +#pragma unroll + for (int h = 0; h < G; ++h) { + const float4 qa = *reinterpret_cast(&sq[h][lane * 8]); + const float4 qb = *reinterpret_cast(&sq[h][lane * 8 + 4]); + float s = k8[0] * qa.x + k8[1] * qa.y + k8[2] * qa.z + k8[3] * qa.w + + k8[4] * qb.x + k8[5] * qb.y + k8[6] * qb.z + k8[7] * qb.w; + s = warp_sum(s); + if (lane == 0) sp[h][c] = s * scale; + } + } + __syncthreads(); + for (int h = warp; h < G; h += WARPS) { + const float a = sp[h][lane], b = sp[h][lane + 32]; + const float m = warp_max(fmaxf(a, b)); + const float ea = (lane < n_here && srow[lane] != -1) ? __expf(a - m) : 0.0f; + const float eb = (lane + 32 < n_here && srow[lane + 32] != -1) ? __expf(b - m) : 0.0f; + sp[h][lane] = ea; + sp[h][lane + 32] = eb; + const float l = warp_sum(ea + eb); + if (lane == 0) { + if (first) { + run_m[h] = m; run_l[h] = l; + } else { + // a chunk with every cell masked has m = -FLT_MAX and l = 0: its weight is 0 (or both are 1, sums 0) + const float mn = fmaxf(run_m[h], m); + const float wr = __expf(run_m[h] - mn), wn = __expf(m - mn); + run_l[h] = fmaf(l, wn, run_l[h] * wr); + run_m[h] = mn; + w_run[h] = wr; w_new[h] = wn; + } + } + } + __syncthreads(); + float cacc[G]; +#pragma unroll + for (int h = 0; h < G; ++h) cacc[h] = 0.0f; + for (int c = 0; c < n_here; ++c) { + if (srow[c] == -1) continue; + const float v = load_v1r(p, srow[c], t); +#pragma unroll + for (int h = 0; h < G; ++h) cacc[h] = fmaf(sp[h][c], v, cacc[h]); + } + if (first) { +#pragma unroll + for (int h = 0; h < G; ++h) acc[h] = cacc[h]; + } else { +#pragma unroll + for (int h = 0; h < G; ++h) acc[h] = fmaf(cacc[h], w_new[h], acc[h] * w_run[h]); + } + first = false; + } + if (t < G) { part_m[slot * G + t] = run_m[t]; part_l[slot * G + t] = run_l[t]; } +#pragma unroll + for (int h = 0; h < G; ++h) part_acc[((size_t) slot * G + h) * HD + t] = acc[h]; +} + +// STRATA_ATTN_SPLITS: 0 (default, one partial per chunk) or the partition count, read once; the parity test sets it +int& attn_splits_ref() { + static int n = [] { + const char* v = std::getenv("STRATA_ATTN_SPLITS"); + const int x = v ? std::atoi(v) : 0; + return x >= 16 && x <= 65535 ? x : 0; + }(); + return n; +} +// the partials one query uses: the chunk count, or P when the split-K form runs (the context exceeds P chunks) +int attn_parts(int64_t cap) { + const int chunks = (int) ((cap + CHUNK - 1) / CHUNK); + const int p = attn_splits_ref(); + return p > 0 && chunks > p ? p : chunks; +} + #if defined(STRATA_HIP_GFX906) // The compat backend can target other architectures. Check the calling thread's current device, // caching only its last ordinal so a layer split or another thread cannot reuse the wrong answer. @@ -605,27 +743,27 @@ bool pre75_attn() { #endif } #if defined(__HIPCC__) -#define STRATA_ATTN_CHUNK(M) (pre75_attn() ? attn_chunk_kernel_pre75 : attn_chunk_kernel) +#define STRATA_ATTN_CHUNK(M) (pre75_attn() ? attn_chunk_kernel_pre75 : attn_chunk_kernel) #else -#define STRATA_ATTN_CHUNK(M) (pre75_attn() ? attn_chunk_kernel_pre75 : attn_chunk_kernel) +#define STRATA_ATTN_CHUNK(M) (pre75_attn() ? attn_chunk_kernel_pre75 : attn_chunk_kernel) #endif #else -#define STRATA_ATTN_CHUNK(M) attn_chunk_kernel +#define STRATA_ATTN_CHUNK(M) attn_chunk_kernel #endif -} // namespace - -void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, - int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream) { +template +void decode_attn_batch_g(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, + int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream) { if (n_q <= 0) return; - if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !scratch || !ids || !steps || + if (s.head_dim != HD || s.n_head != (int64_t) GG * s.n_head_kv || cap <= 0 || !scratch || !ids || !steps || !pools.page_table || n_q > 65535) { std::fprintf(stderr, "qsa_decode_attn_batch: unsupported geometry or missing buffers\n"); std::exit(1); } const int kv_mode = pools.k_q4 != nullptr ? 2 : (pools.k_q != nullptr && pools.v_q4 != nullptr ? 3 : (pools.k_q != nullptr ? 1 : 0)); - const int n_chunks = (int) ((cap + CHUNK - 1) / CHUNK); + const int n_chunks = attn_parts(cap); // the chunk count, or STRATA_ATTN_SPLITS' P past P chunks + const bool split = n_chunks < (int) ((cap + CHUNK - 1) / CHUNK); // per query: [acc: n_chunks*n_head*HD][m: n_chunks*n_head][l: n_chunks*n_head], all offsets from one stride const long long stride = (long long) qsa_decode_attn_scratch_floats(cap, s); float* part_acc = scratch; @@ -649,7 +787,15 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 use_query_swizzle, (long long) n_q, (long long) cap, kv_mode, lane_cell); } #endif -#define STRATA_ATTN_LC(M) attn_chunk_kernel<<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride) +#define STRATA_ATTN_SPLIT(M) attn_split_kernel<<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride) + if (split) { + if (kv_mode == 3) STRATA_ATTN_SPLIT(3); + else if (kv_mode == 2) STRATA_ATTN_SPLIT(2); + else if (kv_mode == 1) STRATA_ATTN_SPLIT(1); + else STRATA_ATTN_SPLIT(0); + } else +#undef STRATA_ATTN_SPLIT +#define STRATA_ATTN_LC(M) attn_chunk_kernel<<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride) if (lane_cell) { if (kv_mode == 3) STRATA_ATTN_LC(3); else if (kv_mode == 2) STRATA_ATTN_LC(2); @@ -669,7 +815,7 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 return v && v[0] == '1' && v[1] == '\0'; }(); if (reduce12) { - attn_chunk_kernel<1, false, true, true><<>>(q, pools, ids, steps, + attn_chunk_kernel<1, false, GG, true, true><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); static const bool trace = [] { const char* v = std::getenv("STRATA_GFX906_ATTN_REDUCE12_TRACE"); @@ -682,7 +828,7 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 (long long) n_q, (long long) cap); } } else { - attn_chunk_kernel<1, false, true><<>>(q, pools, ids, steps, + attn_chunk_kernel<1, false, GG, true><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); } } else @@ -699,8 +845,8 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 else STRATA_ATTN_CHUNK(0)<<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); - attn_merge_kernel<<>>(part_acc, part_m, part_l, n_chunks, - attn, stride); + attn_merge_kernel<<>>(part_acc, part_m, part_l, n_chunks, + attn, stride, steps); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { std::fprintf(stderr, "qsa_decode_attn_batch: %s\n", cudaGetErrorString(e)); @@ -708,14 +854,42 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 } } -uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s) { - const int64_t chunks = (cap + CHUNK - 1) / CHUNK; - return (uint64_t) chunks * (uint64_t) s.n_head * (HD + 2) + 64; +template +void decode_attn_step_g(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, + int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream); + +} // namespace + +// query heads per KV head: 12 (Flash-Next) or 8 (Qwen3.6); anything else is refused by the body's geometry check +void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, + int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream) { + if (s.n_head_kv > 0 && s.n_head == 8 * s.n_head_kv) + decode_attn_batch_g<8>(q, pools, ids, steps, cap, s, scratch, attn, n_q, stream); + else + decode_attn_batch_g(q, pools, ids, steps, cap, s, scratch, attn, n_q, stream); } void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream) { - if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !scratch || !ids || !step || + if (s.n_head_kv > 0 && s.n_head == 8 * s.n_head_kv) + decode_attn_step_g<8>(q, pools, ids, step, cap, s, scratch, attn, stream); + else + decode_attn_step_g(q, pools, ids, step, cap, s, scratch, attn, stream); +} + +void qsa_decode_attn_set_splits(int parts) { attn_splits_ref() = parts >= 16 && parts <= 65535 ? parts : 0; } +int qsa_decode_attn_splits() { return attn_splits_ref(); } + +uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s) { + const int64_t chunks = attn_parts(cap); + return (uint64_t) chunks * (uint64_t) s.n_head * (HD + 2) + 64; +} + +namespace { +template +void decode_attn_step_g(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, + int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream) { + if (s.head_dim != HD || s.n_head != (int64_t) GG * s.n_head_kv || cap <= 0 || !scratch || !ids || !step || !pools.page_table) { std::fprintf(stderr, "qsa_decode_attn: unsupported geometry or missing buffers\n"); std::exit(1); @@ -728,14 +902,28 @@ void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32 std::fprintf(stderr, "qsa_decode_attn: incomplete KV pools\n"); std::exit(1); } - const int n_chunks = (int) ((cap + CHUNK - 1) / CHUNK); + const int n_chunks = attn_parts(cap); + const bool split = n_chunks < (int) ((cap + CHUNK - 1) / CHUNK); float* part_acc = scratch; float* part_m = scratch + (size_t) n_chunks * s.n_head * HD; float* part_l = part_m + (size_t) n_chunks * s.n_head; const float scale = 1.0f / sqrtf((float) HD); const dim3 grid((unsigned) n_chunks, (unsigned) s.n_head_kv); cudaStream_t st = (cudaStream_t) stream; - if (kv_mode == 3) + if (split) { + if (kv_mode == 3) + attn_split_kernel<3, GG><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, 0, 0); + else if (kv_mode == 2) + attn_split_kernel<2, GG><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, 0, 0); + else if (kv_mode == 1) + attn_split_kernel<1, GG><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, 0, 0); + else + attn_split_kernel<0, GG><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, 0, 0); + } else if (kv_mode == 3) STRATA_ATTN_CHUNK(3)<<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, 0, 0); else if (kv_mode == 2) @@ -747,12 +935,13 @@ void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32 else STRATA_ATTN_CHUNK(0)<<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, scale, part_acc, part_m, part_l, n_chunks, 0, 0); - attn_merge_kernel<<<(unsigned) s.n_head, HD, 0, st>>>(part_acc, part_m, part_l, n_chunks, attn); + attn_merge_kernel<<<(unsigned) s.n_head, HD, 0, st>>>(part_acc, part_m, part_l, n_chunks, attn, 0, step); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { std::fprintf(stderr, "qsa_decode_attn: %s\n", cudaGetErrorString(e)); std::exit(1); } } +} // namespace } // namespace strata::kernels diff --git a/src/kernels/cuda/qsa_prompt_attn.cu b/src/kernels/cuda/qsa_prompt_attn.cu index 4a203c55261..0f18b5946c6 100644 --- a/src/kernels/cuda/qsa_prompt_attn.cu +++ b/src/kernels/cuda/qsa_prompt_attn.cu @@ -103,7 +103,7 @@ struct Smem { long long row[CH]; }; -template +template __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __restrict__ q, QsaAttnPools p, const int32_t* __restrict__ ids, const int32_t* __restrict__ steps, int n_kv_heads, @@ -112,10 +112,10 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __res extern __shared__ __align__(16) unsigned char smem_raw[]; Smem& S = *reinterpret_cast*>(smem_raw); const int qi = blockIdx.x, kvh = blockIdx.y; - const int n_head = n_kv_heads * G; - q += (size_t) qi * n_head * HD + (size_t) kvh * G * HD; - attn += (size_t) qi * n_head * HD + (size_t) kvh * G * HD; - ids += (size_t) qi * cap; + const int n_head = n_kv_heads * GQ; + q += (size_t) qi * n_head * HD + (size_t) kvh * GQ * HD; + attn += (size_t) qi * n_head * HD + (size_t) kvh * GQ * HD; + if (ids != nullptr) ids += (size_t) qi * cap; // null: dense, cell = index const int n = __ldg(steps + (size_t) qi * kStepCount + kStepWidth); const int t = threadIdx.x, lane = t & 31, warp = t >> 5; const int gid = lane >> 2, tig = lane & 3; @@ -123,7 +123,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __res // q: 12 heads + 4 zero rows, scaled by a power of two that puts its largest value near 2^14 (exact, and the // lo halves stay out of FP16's subnormal range), then split into hi + lo halves float qm = 0.0f; - for (int i = t; i < G * HD; i += THREADS) qm = fmaxf(qm, fabsf(q[i])); + for (int i = t; i < GQ * HD; i += THREADS) qm = fmaxf(qm, fabsf(q[i])); #pragma unroll for (int o = 16; o > 0; o >>= 1) qm = fmaxf(qm, __shfl_xor_sync(0xffffffffu, qm, o)); if (lane == 0) S.qmax[warp] = qm; @@ -134,7 +134,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __res const float qup = ldexpf(1.0f, 14 - qe), qdown = ldexpf(scale_log2, qe - 14); for (int i = t; i < 16 * HD; i += THREADS) { const int h = i / HD, d = i % HD; - const float x = h < G ? q[(size_t) h * HD + d] * qup : 0.0f; + const float x = h < GQ ? q[(size_t) h * HD + d] * qup : 0.0f; const __half hi = __float2half_rn(x); S.qh[h][d] = hi; S.ql[h][d] = __float2half_rn(x - __half2float(hi)); @@ -150,7 +150,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __res if (t < CH) { long long r = -1; if (t < nh) { - const int cell = ids[c0 + t]; + const int cell = ids != nullptr ? ids[c0 + t] : c0 + t; const long long page = (long long) p.page_table[cell / page_size]; r = (page * n_kv_heads + kvh) * page_size + (cell % page_size); } @@ -417,7 +417,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __res for (int j = 0; j < 8; ++j) { const int d = warp * 64 + j * 8 + 2 * tig; *reinterpret_cast(attn + (size_t) gid * HD + d) = make_float2(acc[j][0] * i0, acc[j][1] * i0); - if (gid + 8 < G) + if (gid + 8 < GQ) *reinterpret_cast(attn + (size_t) (gid + 8) * HD + d) = make_float2(acc[j][2] * i1, acc[j][3] * i1); } } @@ -461,6 +461,7 @@ __device__ __forceinline__ void cp_async_wait1() { #endif } +template __global__ void __launch_bounds__(THREADS) prompt_attn_i8_kernel(const float* __restrict__ q, QsaAttnPools p, const int32_t* __restrict__ ids, const int32_t* __restrict__ steps, int n_kv_heads, @@ -469,10 +470,10 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_i8_kernel(const float* __ extern __shared__ __align__(16) unsigned char smem_raw[]; Smem2& S = *reinterpret_cast(smem_raw); const int qi = blockIdx.x, kvh = blockIdx.y; - const int n_head = n_kv_heads * G; - q += (size_t) qi * n_head * HD + (size_t) kvh * G * HD; - attn += (size_t) qi * n_head * HD + (size_t) kvh * G * HD; - ids += (size_t) qi * cap; + const int n_head = n_kv_heads * GQ; + q += (size_t) qi * n_head * HD + (size_t) kvh * GQ * HD; + attn += (size_t) qi * n_head * HD + (size_t) kvh * GQ * HD; + if (ids != nullptr) ids += (size_t) qi * cap; // null: dense, cell = index const int n = __ldg(steps + (size_t) qi * kStepCount + kStepWidth); const int t = threadIdx.x, lane = t & 31, warp = t >> 5; const int gid = lane >> 2, tig = lane & 3; @@ -480,7 +481,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_i8_kernel(const float* __ // q: the power-of-two prescale over all 12 heads (as v1), then this warp's 64 dims as hi/lo A fragments float qm = 0.0f; - for (int i = t; i < G * HD; i += THREADS) qm = fmaxf(qm, fabsf(q[i])); + for (int i = t; i < GQ * HD; i += THREADS) qm = fmaxf(qm, fabsf(q[i])); #pragma unroll for (int o = 16; o > 0; o >>= 1) qm = fmaxf(qm, __shfl_xor_sync(0xffffffffu, qm, o)); if (lane == 0) S.qmax[warp] = qm; @@ -497,7 +498,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_i8_kernel(const float* __ for (int r = 0; r < 4; ++r) { const int row = gid + (r & 1) * 8, col = dim0 + kk * 16 + 2 * tig + (r >> 1) * 8; float2 x = make_float2(0.f, 0.f); - if (row < G) x = *reinterpret_cast(q + (size_t) row * HD + col); + if (row < GQ) x = *reinterpret_cast(q + (size_t) row * HD + col); x.x *= qup; x.y *= qup; const __half2 hi = __floats2half2_rn(x.x, x.y); @@ -510,7 +511,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_i8_kernel(const float* __ // the chunk pipeline: cells two chunks ahead, their pool rows one chunk ahead, the data (cp.async) one ahead const int n_chunks = (n + CH2 - 1) / CH2; - auto cell_of = [&](int c) -> int { return c < n ? __ldg(ids + c) : -1; }; + auto cell_of = [&](int c) -> int { return c < n ? (ids != nullptr ? __ldg(ids + c) : c) : -1; }; auto row_of = [&](int cell) -> long long { if (cell < 0) return -1; const long long page = (long long) __ldg(p.page_table + cell / page_size); @@ -671,11 +672,12 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_i8_kernel(const float* __ for (int j = 0; j < 8; ++j) { const int d = dim0 + j * 8 + 2 * tig; *reinterpret_cast(attn + (size_t) gid * HD + d) = make_float2(acc[j][0] * i0, acc[j][1] * i0); - if (gid + 8 < G) + if (gid + 8 < GQ) *reinterpret_cast(attn + (size_t) (gid + 8) * HD + d) = make_float2(acc[j][2] * i1, acc[j][3] * i1); } } +template bool launch_i8(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap, const QsaShapes& s, float* attn, int64_t n_q, cudaStream_t st) { static bool attr[64] = {}; // the shared-memory opt-in is per device (a layer split runs this on several) @@ -684,7 +686,7 @@ bool launch_i8(const float* q, const QsaAttnPools& pools, const int32_t* ids, co const int bytes = (int) sizeof(Smem2); if (dev < 0 || dev >= 64) return false; if (!attr[dev]) { - if (cudaFuncSetAttribute(prompt_attn_i8_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) != + if (cudaFuncSetAttribute(prompt_attn_i8_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) != cudaSuccess) { cudaGetLastError(); return false; @@ -694,8 +696,8 @@ bool launch_i8(const float* q, const QsaAttnPools& pools, const int32_t* ids, co const float scale_log2 = 1.4426950408889634f / sqrtf((float) HD); for (int64_t q0 = 0; q0 < n_q; q0 += 65535) { const int64_t nb = n_q - q0 < 65535 ? n_q - q0 : 65535; - prompt_attn_i8_kernel<<>>( - q + q0 * s.n_head * HD, pools, ids + q0 * cap, steps + q0 * kStepCount, (int) s.n_head_kv, + prompt_attn_i8_kernel<<>>( + q + q0 * s.n_head * HD, pools, ids != nullptr ? ids + q0 * cap : nullptr, steps + q0 * kStepCount, (int) s.n_head_kv, (int) s.page_size, scale_log2, attn + q0 * s.n_head * HD, (int) cap); } const cudaError_t e = cudaGetLastError(); @@ -706,7 +708,7 @@ bool launch_i8(const float* q, const QsaAttnPools& pools, const int32_t* ids, co return true; } -template +template bool launch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap, const QsaShapes& s, float* attn, int64_t n_q, cudaStream_t st) { static bool attr[64] = {}; // per device, as above @@ -715,7 +717,7 @@ bool launch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const const int bytes = (int) sizeof(Smem); if (dev < 0 || dev >= 64) return false; if (!attr[dev]) { - if (cudaFuncSetAttribute(prompt_attn_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) != + if (cudaFuncSetAttribute(prompt_attn_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) != cudaSuccess) { cudaGetLastError(); return false; @@ -725,8 +727,8 @@ bool launch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const const float scale_log2 = 1.4426950408889634f / sqrtf((float) HD); for (int64_t q0 = 0; q0 < n_q; q0 += 65535) { const int64_t nb = n_q - q0 < 65535 ? n_q - q0 : 65535; - prompt_attn_kernel<<>>( - q + q0 * s.n_head * HD, pools, ids + q0 * cap, steps + q0 * kStepCount, (int) s.n_head_kv, + prompt_attn_kernel<<>>( + q + q0 * s.n_head * HD, pools, ids != nullptr ? ids + q0 * cap : nullptr, steps + q0 * kStepCount, (int) s.n_head_kv, (int) s.page_size, scale_log2, attn + q0 * s.n_head * HD, (int) cap); } const cudaError_t e = cudaGetLastError(); @@ -737,6 +739,397 @@ bool launch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const return true; } +// ---- dense causal prompt attention (Qwen3.6: 8 query heads per KV head, no ids) -------------------------------------- +// The kernels above run one block per (query, KV head). With a selection that is fine (the 2,051 cells), but in dense +// mode each block reads every cell of the context on its own, so a prompt chunk's time grows with the context: at 150K +// tokens this attention was 93% of a Qwen3.6 prompt (RTX 4070 Ti, fp16 KV, 771 s for 150K tokens). This one gives a +// block 8 neighbouring queries of one KV head: 64 MMA rows (8 queries x 8 heads, no pad rows), warp w the 16 rows of +// queries 2w and 2w+1 over all 256 dims, so a K/V chunk is read once for 8 queries. The softmax runs in registers and +// the q.k accumulator fragments become p.v's A fragments (no score round trip through shared memory). Accuracy as the +// kernels above: q and p as FP16 hi + lo parts; int8 codes enter exactly, the per-64 scales in FP32 (K: on each group's +// q.k partial; V: folded into p, which is scaled by a power of two per chunk and group to keep it in FP16's normal +// range - the accumulator is rescaled by the same power, exactly). cp.async (sm_80+): the next chunk's K loads during +// this chunk's p.v, the next V during the next q.k. STRATA_PROMPT_ATTN_DENSE=0: the per-query kernels (A/B). +constexpr int DQ = 8; // queries per block +constexpr int DCH = 32; // cells per chunk + +template // 0: fp16 K and V; 1: int8 codes + an fp16 scale per 64 values +struct SmemD { + static constexpr int PIECES = HD * (KV_MODE == 0 ? 2 : 1) / 16; // 16-byte pieces per K/V row + uint4 qh[DQ * 8][HD / 8]; // 64 rows; a row's pieces XOR-swizzled by row & 7 (conflict-free ldmatrix) + uint4 ql[DQ * 8][HD / 8]; + uint4 k[DCH][PIECES]; // a cell's pieces XOR-swizzled by cell & 7 + uint4 v[DCH][PIECES]; + uint2 ksc[DCH]; // int8: the cell's 4 fp16 scales + uint2 vsc[DCH]; + float qdn[DQ * 8]; // per row: the q prescale taken back out, times log2(e) / 16 +}; + +__device__ __forceinline__ void ldsm_x4(uint32_t* r, const void* p) { +#if STRATA_PA_SM80 + const unsigned a = (unsigned) __cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(a)); +#else + (void) r; (void) p; + __trap(); +#endif +} +__device__ __forceinline__ void ldsm_x4_t(uint32_t* r, const void* p) { +#if STRATA_PA_SM80 + const unsigned a = (unsigned) __cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(a)); +#else + (void) r; (void) p; + __trap(); +#endif +} +__device__ __forceinline__ void cp_async8(void* smem, const void* gmem, bool valid) { +#if !STRATA_PA_SM80 + (void) smem; (void) gmem; (void) valid; + __trap(); +#else + const unsigned sa = (unsigned) __cvta_generic_to_shared(smem); + asm volatile("cp.async.ca.shared.global [%0], [%1], 8, %2;\n" ::"r"(sa), "l"(gmem), "r"(valid ? 8 : 0)); +#endif +} +__device__ __forceinline__ void cp_async_wait0() { +#if STRATA_PA_SM80 + asm volatile("cp.async.wait_group 0;\n" ::); +#endif +} + +// p as FP16 hi + lo: hi = (a, b) rounded, lo = what is left +__device__ __forceinline__ void split_h2(float a, float b, uint32_t& hi, uint32_t& lo) { + hi = pack_h2(a, b); + const float2 f = __half22float2(*reinterpret_cast(&hi)); + lo = pack_h2(a - f.x, b - f.y); +} + +template +__global__ void __launch_bounds__(THREADS, 1) prompt_attn_dense_kernel(const float* __restrict__ q, QsaAttnPools p, + const int32_t* __restrict__ steps, int n_q, + int n_kv_heads, int page_size, float scale_log2, + float* __restrict__ attn) { + constexpr int GQ = 8, PIECES = SmemD::PIECES, ROWB = PIECES * 16; + extern __shared__ __align__(16) unsigned char smem_raw[]; + SmemD& S = *reinterpret_cast*>(smem_raw); + const int q0 = blockIdx.x * DQ, kvh = blockIdx.y, n_head = n_kv_heads * GQ; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5, gid = lane >> 2, tig = lane & 3; + const int rbase = warp * 16; // this warp's MMA rows: query q0 + 2w (rows 0-7), q0 + 2w + 1 (rows 8-15), by head + auto width = [&](int i) { return q0 + i < n_q ? __ldg(steps + (size_t) (q0 + i) * kStepCount + kStepWidth) : 0; }; + const int na = width(2 * warp), nb = width(2 * warp + 1); + int n_max = 0; +#pragma unroll + for (int i = 0; i < DQ; ++i) n_max = max(n_max, width(i)); + + auto row_of = [&](int cell) -> long long { + const long long page = (long long) __ldg(p.page_table + cell / page_size); + return (page * n_kv_heads + kvh) * page_size + (cell % page_size); + }; + // the K (kv 0) or V (kv 1) rows of the chunk at cell c0, cells past the widest query zero-filled + auto issue = [&](int kv, int c0) { + const unsigned char* base = + KV_MODE == 0 ? reinterpret_cast(kv ? p.v_pool : p.k_pool) + : reinterpret_cast(kv ? p.v_q : p.k_q); + uint4(*dst)[PIECES] = kv ? S.v : S.k; +#pragma unroll + for (int i = t; i < DCH * PIECES; i += THREADS) { + const int c = i / PIECES, pc = i % PIECES, cell = c0 + c; + const bool ok = cell < n_max; + const size_t off = ok ? (size_t) row_of(cell) * ROWB + pc * 16 : 0; + cp_async16(&dst[c][pc ^ (c & 7)], base + off, ok); + } + if constexpr (KV_MODE == 1) { + if (t < DCH) { + const int cell = c0 + t; + const bool ok = cell < n_max; + const size_t off = ok ? (size_t) row_of(cell) * (HD / KV_Q8_GROUP) : 0; + cp_async8(kv ? &S.vsc[t] : &S.ksc[t], (kv ? p.v_scale : p.k_scale) + off, ok); + } + } + }; + if (n_max > 0) { + issue(0, 0); + cp_async_commit(); + issue(1, 0); + cp_async_commit(); + } + + // q: per row, a power of two that puts its largest value near 2^14 (as the kernels above, per row here), then hi + lo +#pragma unroll 1 + for (int r = 0; r < 16; ++r) { + const int qi = q0 + 2 * warp + (r >> 3), row = rbase + r; + float x[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (qi < n_q) { + const float4* src = reinterpret_cast(q + ((size_t) qi * n_head + kvh * GQ + (r & 7)) * HD) + lane * 2; + const float4 a = src[0], b = src[1]; + x[0] = a.x; x[1] = a.y; x[2] = a.z; x[3] = a.w; x[4] = b.x; x[5] = b.y; x[6] = b.z; x[7] = b.w; + } + float m = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) m = fmaxf(m, fabsf(x[j])); +#pragma unroll + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + int e = 0; + if (m > 0.0f) frexpf(m, &e); // m < 2^e + const float up = ldexpf(1.0f, 14 - e); + uint4 hi, lo; + split_h2(x[0] * up, x[1] * up, hi.x, lo.x); + split_h2(x[2] * up, x[3] * up, hi.y, lo.y); + split_h2(x[4] * up, x[5] * up, hi.z, lo.z); + split_h2(x[6] * up, x[7] * up, hi.w, lo.w); + S.qh[row][lane ^ (row & 7)] = hi; + S.ql[row][lane ^ (row & 7)] = lo; + if (lane == 0) S.qdn[row] = ldexpf(scale_log2, e - 14); + } + __syncwarp(); + const float da = S.qdn[rbase + gid], db = S.qdn[rbase + 8 + gid]; + + float acc[HD / 8][4]; +#pragma unroll + for (int j = 0; j < HD / 8; ++j) acc[j][0] = acc[j][1] = acc[j][2] = acc[j][3] = 0.0f; + float m_a = -INFINITY, m_b = -INFINITY, l_a = 0.0f, l_b = 0.0f; // l: this thread's columns only (summed at the end) + int ve[4] = {14, 14, 14, 14}; // p.v's scale per V group: p carries 2^ve (int8: 2^ve / the chunk's largest scale) + +#pragma unroll 1 + for (int c0 = 0; c0 < n_max; c0 += DCH) { + const bool more = c0 + DCH < n_max; + cp_async_wait1(); + __syncthreads(); // K in + // q.k: 4 n-tiles of 8 cells x 16 k-steps (hi and lo); int8 per 64-dim group, then times the cell's K scale + float s[DCH / 8][4]; +#pragma unroll + for (int nt = 0; nt < DCH / 8; ++nt) s[nt][0] = s[nt][1] = s[nt][2] = s[nt][3] = 0.0f; +#pragma unroll + for (int g = 0; g < 4; ++g) { + float tg[DCH / 8][4]; +#pragma unroll + for (int nt = 0; nt < DCH / 8; ++nt) tg[nt][0] = tg[nt][1] = tg[nt][2] = tg[nt][3] = 0.0f; +#pragma unroll + for (int kq = 0; kq < 4; ++kq) { + const int kk = g * 4 + kq; // 16-dim step + uint32_t ah[4], al[4]; + { + const int mi = lane >> 3, rr = rbase + (mi & 1) * 8 + (lane & 7), pc = kk * 2 + (mi >> 1); + ldsm_x4(ah, &S.qh[rr][pc ^ (rr & 7)]); + ldsm_x4(al, &S.ql[rr][pc ^ (rr & 7)]); + } + if constexpr (KV_MODE == 0) { +#pragma unroll + for (int np = 0; np < DCH / 16; ++np) { // two n-tiles per ldmatrix + const int c = 16 * np + 8 * (lane >> 4) + (lane & 7), pc = kk * 2 + ((lane >> 3) & 1); + uint32_t b[4]; + ldsm_x4(b, &S.k[c][pc ^ (c & 7)]); + mma16816(tg[2 * np], ah, b); + mma16816(tg[2 * np], al, b); + mma16816(tg[2 * np + 1], ah, b + 2); + mma16816(tg[2 * np + 1], al, b + 2); + } + } else { + const unsigned char* K = reinterpret_cast(&S.k[0][0]); +#pragma unroll + for (int nt = 0; nt < DCH / 8; ++nt) { + const int c = nt * 8 + gid; + const unsigned char* rowp = K + c * ROWB + ((kk ^ (c & 7)) << 4) + 2 * tig; + uint32_t b[2]; + b[0] = i8x2_to_h2(*reinterpret_cast(rowp)); + b[1] = i8x2_to_h2(*reinterpret_cast(rowp + 8)); + mma16816(tg[nt], ah, b); + mma16816(tg[nt], al, b); + } + } + } +#pragma unroll + for (int nt = 0; nt < DCH / 8; ++nt) { + float k0 = 1.0f, k1 = 1.0f; + if constexpr (KV_MODE == 1) { + const int c = nt * 8 + 2 * tig; + k0 = __half2float(reinterpret_cast(&S.ksc[c])[g]); + k1 = __half2float(reinterpret_cast(&S.ksc[c + 1])[g]); + } + s[nt][0] = fmaf(tg[nt][0], k0, s[nt][0]); + s[nt][1] = fmaf(tg[nt][1], k1, s[nt][1]); + s[nt][2] = fmaf(tg[nt][2], k0, s[nt][2]); + s[nt][3] = fmaf(tg[nt][3], k1, s[nt][3]); + } + } + __syncthreads(); // every warp is done with K + if (more) { + issue(0, c0 + DCH); + cp_async_commit(); + } + // online softmax in registers: rows a (gid) and b (gid + 8), this thread's 8 columns, the row's 4 threads + float alpha_a, alpha_b; + { + float mxa = -INFINITY, mxb = -INFINITY; +#pragma unroll + for (int nt = 0; nt < DCH / 8; ++nt) { + const int cell = c0 + nt * 8 + 2 * tig; + s[nt][0] = cell < na ? s[nt][0] * da : -INFINITY; + s[nt][1] = cell + 1 < na ? s[nt][1] * da : -INFINITY; + s[nt][2] = cell < nb ? s[nt][2] * db : -INFINITY; + s[nt][3] = cell + 1 < nb ? s[nt][3] * db : -INFINITY; + mxa = fmaxf(mxa, fmaxf(s[nt][0], s[nt][1])); + mxb = fmaxf(mxb, fmaxf(s[nt][2], s[nt][3])); + } +#pragma unroll + for (int o = 1; o < 4; o <<= 1) { + mxa = fmaxf(mxa, __shfl_xor_sync(0xffffffffu, mxa, o)); + mxb = fmaxf(mxb, __shfl_xor_sync(0xffffffffu, mxb, o)); + } + const float ma = fmaxf(m_a, mxa), mb = fmaxf(m_b, mxb); + alpha_a = m_a == -INFINITY ? 0.0f : exp2f(m_a - ma); + alpha_b = m_b == -INFINITY ? 0.0f : exp2f(m_b - mb); + float sa = 0.0f, sb = 0.0f; +#pragma unroll + for (int nt = 0; nt < DCH / 8; ++nt) { + s[nt][0] = s[nt][0] == -INFINITY ? 0.0f : exp2f(s[nt][0] - ma); + s[nt][1] = s[nt][1] == -INFINITY ? 0.0f : exp2f(s[nt][1] - ma); + s[nt][2] = s[nt][2] == -INFINITY ? 0.0f : exp2f(s[nt][2] - mb); + s[nt][3] = s[nt][3] == -INFINITY ? 0.0f : exp2f(s[nt][3] - mb); + sa += s[nt][0] + s[nt][1]; + sb += s[nt][2] + s[nt][3]; + } + l_a = fmaf(l_a, alpha_a, sa); + l_b = fmaf(l_b, alpha_b, sb); + m_a = ma; + m_b = mb; + } + if (more) cp_async_wait1(); else cp_async_wait0(); + __syncthreads(); // V in + // p.v: 32 n-tiles of 8 dims x 2 k-steps of 16 cells (hi and lo) +#pragma unroll + for (int g = 0; g < 4; ++g) { // V scale group: dims [64g, 64g + 64) = n-tiles 8g..8g+7 + float w[2][4] = {{1.f, 1.f, 1.f, 1.f}, {1.f, 1.f, 1.f, 1.f}}; // per k-step: cells 2tig, +1, +8, +9 + int e_new = 14; + if constexpr (KV_MODE == 1) { + float vm = fabsf(__half2float(reinterpret_cast(&S.vsc[lane])[g])); +#pragma unroll + for (int o = 16; o > 0; o >>= 1) vm = fmaxf(vm, __shfl_xor_sync(0xffffffffu, vm, o)); + e_new = ve[g]; + if (vm > 0.0f) { + int e = 0; + frexpf(vm, &e); // vm < 2^e: |scale * 2^(14 - e)| < 2^14 + e_new = 14 - e; + } +#pragma unroll + for (int ks = 0; ks < 2; ++ks) { + const int c = ks * 16 + 2 * tig; + w[ks][0] = ldexpf(__half2float(reinterpret_cast(&S.vsc[c])[g]), e_new); + w[ks][1] = ldexpf(__half2float(reinterpret_cast(&S.vsc[c + 1])[g]), e_new); + w[ks][2] = ldexpf(__half2float(reinterpret_cast(&S.vsc[c + 8])[g]), e_new); + w[ks][3] = ldexpf(__half2float(reinterpret_cast(&S.vsc[c + 9])[g]), e_new); + } + } else { +#pragma unroll + for (int ks = 0; ks < 2; ++ks) w[ks][0] = w[ks][1] = w[ks][2] = w[ks][3] = 16384.0f; + } + // the chunk's p.v from zero, then added to the running sum in FP32 (the MMA's own accumulation rounds + // more coarsely: over a long context that drifted ~30x past the kernels above) + float tmp[8][4]; +#pragma unroll + for (int j = 0; j < 8; ++j) tmp[j][0] = tmp[j][1] = tmp[j][2] = tmp[j][3] = 0.0f; +#pragma unroll + for (int ks = 0; ks < 2; ++ks) { + uint32_t ah[4], al[4]; + split_h2(s[2 * ks][0] * w[ks][0], s[2 * ks][1] * w[ks][1], ah[0], al[0]); + split_h2(s[2 * ks][2] * w[ks][0], s[2 * ks][3] * w[ks][1], ah[1], al[1]); + split_h2(s[2 * ks + 1][0] * w[ks][2], s[2 * ks + 1][1] * w[ks][3], ah[2], al[2]); + split_h2(s[2 * ks + 1][2] * w[ks][2], s[2 * ks + 1][3] * w[ks][3], ah[3], al[3]); + if constexpr (KV_MODE == 0) { +#pragma unroll + for (int jp = 4 * g; jp < 4 * g + 4; ++jp) { // two n-tiles per ldmatrix + const int c = ks * 16 + (lane & 7) + ((lane >> 3) & 1) * 8, pc = 2 * jp + (lane >> 4); + uint32_t b[4]; + ldsm_x4_t(b, &S.v[c][pc ^ (c & 7)]); + float* t0 = tmp[2 * (jp - 4 * g)]; + float* t1 = tmp[2 * (jp - 4 * g) + 1]; + mma16816(t0, ah, b); + mma16816(t0, al, b); + mma16816(t1, ah, b + 2); + mma16816(t1, al, b + 2); + } + } else { + const unsigned char* V = reinterpret_cast(&S.v[0][0]); + const int cA = ks * 16 + 2 * tig, cB = cA + 8; +#pragma unroll + for (int j = 8 * g; j < 8 * g + 8; ++j) { + const int d = j * 8 + gid; + auto at = [&](int c) { return (uint32_t) V[c * ROWB + (((d >> 4) ^ (c & 7)) << 4) + (d & 15)]; }; + uint32_t b[2]; + b[0] = i8x2_to_h2(at(cA) | (at(cA + 1) << 8)); + b[1] = i8x2_to_h2(at(cB) | (at(cB + 1) << 8)); + mma16816(tmp[j - 8 * g], ah, b); + mma16816(tmp[j - 8 * g], al, b); + } + } + } + // the running sum: times alpha, moved to this chunk's power of two (exact), plus the chunk's + const float fa = ldexpf(alpha_a, e_new - ve[g]), fb = ldexpf(alpha_b, e_new - ve[g]); + ve[g] = e_new; +#pragma unroll + for (int j = 0; j < 8; ++j) { + acc[8 * g + j][0] = fmaf(acc[8 * g + j][0], fa, tmp[j][0]); + acc[8 * g + j][1] = fmaf(acc[8 * g + j][1], fa, tmp[j][1]); + acc[8 * g + j][2] = fmaf(acc[8 * g + j][2], fb, tmp[j][2]); + acc[8 * g + j][3] = fmaf(acc[8 * g + j][3], fb, tmp[j][3]); + } + } + __syncthreads(); // every warp is done with V + if (more) { + issue(1, c0 + DCH); + cp_async_commit(); + } + } +#pragma unroll + for (int o = 1; o < 4; o <<= 1) { + l_a += __shfl_xor_sync(0xffffffffu, l_a, o); + l_b += __shfl_xor_sync(0xffffffffu, l_b, o); + } + const int qa = q0 + 2 * warp, qb = qa + 1; + float* oa = attn + ((size_t) qa * n_head + kvh * GQ + gid) * HD; + float* ob = attn + ((size_t) qb * n_head + kvh * GQ + gid) * HD; +#pragma unroll + for (int g = 0; g < 4; ++g) { + const float ia = l_a > 0.0f ? ldexpf(1.0f / l_a, -ve[g]) : 0.0f, ib = l_b > 0.0f ? ldexpf(1.0f / l_b, -ve[g]) : 0.0f; +#pragma unroll + for (int j = 8 * g; j < 8 * g + 8; ++j) { + const int d = j * 8 + 2 * tig; + if (qa < n_q) *reinterpret_cast(oa + d) = make_float2(acc[j][0] * ia, acc[j][1] * ia); + if (qb < n_q) *reinterpret_cast(ob + d) = make_float2(acc[j][2] * ib, acc[j][3] * ib); + } + } +} + +template +bool launch_dense(const float* q, const QsaAttnPools& pools, const int32_t* steps, const QsaShapes& s, float* attn, + int64_t n_q, cudaStream_t st) { + static bool attr[64] = {}; // per device, as above + int dev = 0; + cudaGetDevice(&dev); + const int bytes = (int) sizeof(SmemD); + if (dev < 0 || dev >= 64 || n_q > (int64_t) 0x7fffffff - DQ) return false; + if (!attr[dev]) { + if (cudaFuncSetAttribute(prompt_attn_dense_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) != + cudaSuccess) { + cudaGetLastError(); + return false; + } + attr[dev] = true; + } + const float scale_log2 = 1.4426950408889634f / sqrtf((float) HD); + prompt_attn_dense_kernel<<>>( + q, pools, steps, (int) n_q, (int) s.n_head_kv, (int) s.page_size, scale_log2, attn); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { + std::fprintf(stderr, "qsa_prompt_attn_batch (dense): %s\n", cudaGetErrorString(e)); + std::exit(1); + } + return true; +} + // PR #600: compiled only into the experimental build (-DSTRATA_EXPERIMENTAL_SM60=ON), the one build that runs on a // Volta card; the ready-made engine has none of it (no extra kernels to load, the same code as before). @@ -1450,6 +1843,46 @@ bool launch_wmma(const float* q, const QsaAttnPools& pools, const int32_t* ids, } // namespace +template +bool dispatch_g(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap, + const QsaShapes& s, float* attn, int64_t n_q, cudaStream_t st, bool volta, bool turing) { + if constexpr (GQ == 8) { // dense (Qwen3.6): 8 queries per block on sm_80+, fp16 or int8 K/V (above) + const char* dn = std::getenv("STRATA_PROMPT_ATTN_DENSE"); // read per call: the parity test switches it + if (ids == nullptr && !turing && !volta && !(dn && dn[0] == '0') && pools.k_q4 == nullptr && pools.v_q4 == nullptr) { + if (pools.k_q && pools.v_q && pools.k_scale && pools.v_scale) + return launch_dense<1>(q, pools, steps, s, attn, n_q, st); + if (!pools.k_q && pools.k_pool && pools.v_pool) return launch_dense<0>(q, pools, steps, s, attn, n_q, st); + } + } + if (pools.k_q4 != nullptr) { // Q4_0 K and V (--kv q4_0): mode 4. STRATA_PROMPT_ATTN_Q4=0: the old kernel (A/B) + static const bool q4_off = [] { + const char* v = std::getenv("STRATA_PROMPT_ATTN_Q4"); + return v != nullptr && v[0] == '0'; + }(); + // sm_80+ only: on Turing mode 4 would run as pairs of m16n8k8 MMAs, which no parity run has checked yet + if (q4_off || turing || pools.v_q4 == nullptr) return false; + return launch<4, GQ>(q, pools, ids, steps, cap, s, attn, n_q, st); + } + if (pools.k_q != nullptr && pools.v_q4 != nullptr) { // hybrid K8V4: int8 K + dequantized-q4 V + if (!pools.k_scale) return false; + if (volta) return GQ == G && launch70<3>(q, pools, ids, steps, cap, s, attn, n_q, st); + return launch<3, GQ>(q, pools, ids, steps, cap, s, attn, n_q, st); + } + if (pools.k_q != nullptr) { + if (!pools.v_q || !pools.k_scale || !pools.v_scale) return false; + // STRATA_PROMPT_ATTN_V1=1 (debug): the first version, same accuracy, another summation order - the control + // for how far the model amplifies an FP32-level change. Turing always takes it: v2's cp.async does not + // exist before sm_80. + static const bool v1 = std::getenv("STRATA_PROMPT_ATTN_V1") != nullptr; + if (volta) return GQ == G && launch70<1>(q, pools, ids, steps, cap, s, attn, n_q, st); + if (v1 || turing) return launch<1, GQ>(q, pools, ids, steps, cap, s, attn, n_q, st); + return launch_i8(q, pools, ids, steps, cap, s, attn, n_q, st); + } + if (!pools.k_pool || !pools.v_pool) return false; + if (volta) return GQ == G && launch70<0>(q, pools, ids, steps, cap, s, attn, n_q, st); + return launch<0, GQ>(q, pools, ids, steps, cap, s, attn, n_q, st); +} + bool qsa_prompt_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap, const QsaShapes& s, float* attn, int64_t n_q, void* stream) { if (n_q <= 0) return true; @@ -1490,36 +1923,15 @@ bool qsa_prompt_attn_batch(const float* q, const QsaAttnPools& pools, const int3 return launch_wmma(q, pools, ids, steps, cap, s, attn, n_q, (cudaStream_t) stream); return false; #endif - if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !ids || !steps || !pools.page_table) + // query heads per KV head: 12 (Flash-Next) or 8 (Qwen3.6); ids null = dense (every cell of the width, which + // the Volta and AMD kernels do not take: the caller's old kernel then runs) + const int64_t gq = s.n_head_kv > 0 ? s.n_head / s.n_head_kv : 0; + if (s.head_dim != HD || s.n_head != gq * s.n_head_kv || (gq != G && gq != 8) || (ids != nullptr && cap <= 0) || + !steps || !pools.page_table) return false; - cudaStream_t st = (cudaStream_t) stream; - if (pools.k_q4 != nullptr) { // Q4_0 K and V (--kv q4_0): mode 4. STRATA_PROMPT_ATTN_Q4=0: the old kernel (A/B) - static const bool q4_off = [] { - const char* v = std::getenv("STRATA_PROMPT_ATTN_Q4"); - return v != nullptr && v[0] == '0'; - }(); - // sm_80+ only: on Turing mode 4 would run as pairs of m16n8k8 MMAs, which no parity run has checked yet - if (q4_off || turing || pools.v_q4 == nullptr) return false; - return launch<4>(q, pools, ids, steps, cap, s, attn, n_q, st); - } - if (pools.k_q != nullptr && pools.v_q4 != nullptr) { // hybrid K8V4: int8 K + dequantized-q4 V - if (!pools.k_scale) return false; - if (volta) return launch70<3>(q, pools, ids, steps, cap, s, attn, n_q, st); - return launch<3>(q, pools, ids, steps, cap, s, attn, n_q, st); - } - if (pools.k_q != nullptr) { - if (!pools.v_q || !pools.k_scale || !pools.v_scale) return false; - // STRATA_PROMPT_ATTN_V1=1 (debug): the first version, same accuracy, another summation order - the control - // for how far the model amplifies an FP32-level change. Turing always takes it: v2's cp.async does not - // exist before sm_80. - static const bool v1 = std::getenv("STRATA_PROMPT_ATTN_V1") != nullptr; - if (volta) return launch70<1>(q, pools, ids, steps, cap, s, attn, n_q, st); - if (v1 || turing) return launch<1>(q, pools, ids, steps, cap, s, attn, n_q, st); - return launch_i8(q, pools, ids, steps, cap, s, attn, n_q, st); - } - if (!pools.k_pool || !pools.v_pool) return false; - if (volta) return launch70<0>(q, pools, ids, steps, cap, s, attn, n_q, st); - return launch<0>(q, pools, ids, steps, cap, s, attn, n_q, st); + if (volta && (gq != G || ids == nullptr)) return false; + return gq == 8 ? dispatch_g<8>(q, pools, ids, steps, cap, s, attn, n_q, (cudaStream_t) stream, volta, turing) + : dispatch_g(q, pools, ids, steps, cap, s, attn, n_q, (cudaStream_t) stream, volta, turing); } } // namespace strata::kernels diff --git a/src/kernels/cuda/qsa_select.cu b/src/kernels/cuda/qsa_select.cu index 4820d208fc1..6bfcf0fdc32 100644 --- a/src/kernels/cuda/qsa_select.cu +++ b/src/kernels/cuda/qsa_select.cu @@ -1292,6 +1292,21 @@ bool qsa_block_topk_cluster(const float* scores, const int32_t* steps, int64_t n #endif } +namespace { +// one block per query: ids[q][i] = i for every cell of its selection width (= n_kv with kDenseTopK) +__global__ void qsa_select_all_kernel(const int32_t* __restrict__ steps, int64_t cap, int32_t* __restrict__ ids) { + const int q = blockIdx.x; + const int width = steps[(size_t) q * kStepCount + kStepWidth]; + int32_t* row = ids + (size_t) q * (size_t) cap; + for (int i = threadIdx.x; i < width; i += blockDim.x) row[i] = i; +} +} // namespace + +void qsa_select_all(const int32_t* steps, int64_t nq, int64_t cap, int32_t* ids, void* stream) { + if (nq <= 0) return; + qsa_select_all_kernel<<<(unsigned) nq, 256, 0, (cudaStream_t) stream>>>(steps, cap, ids); +} + void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap, const QsaShapes& s, int32_t* ids, void* stream, int64_t active_blocks) { // keys in registers when every query's blocks fit (contexts up to ~135K cells), else (CUDA) the same threads diff --git a/src/kernels/cuda/router_top10.cu b/src/kernels/cuda/router_top10.cu index 720cec67dc3..e7b610c85fd 100644 --- a/src/kernels/cuda/router_top10.cu +++ b/src/kernels/cuda/router_top10.cu @@ -263,7 +263,7 @@ __global__ void __launch_bounds__(64) router_top10_wave_kernel(const float* __re } #endif // STRATA_HIP_GFX906 -#if defined(__HIPCC__) +// (built on every backend since it runs on CUDA too: the argument below is IEEE arithmetic, not the hardware's) // ---- S6, AMD: the same router, BIT-IDENTICAL, without its serial parts. Measured on RDNA4 (gfx1201) the kernel // above took 39 us per call in decode (16% of the GPU's decode time): the ascending double sum is 512 dependent FP64 // adds (~50 cycles each there), the ten selection passes are 20 block barriers, and the renormalisation re-reads its @@ -401,22 +401,16 @@ void launch_generic(const float* logits, int n_tokens, int n_expert, int k, int* router_top10_kernel<<<(unsigned) n_tokens, threads, smem, (cudaStream_t) stream>>>( logits, n_tokens, n_expert, k, ids, weights); } -#endif +// (end of the fast router) } // namespace bool router_top10_variant(const float* logits, int n_tokens, int n_expert, int k, int* ids, float* weights, void* stream, int variant) { -#if defined(__HIPCC__) if (n_tokens <= 0 || n_expert <= 64 || n_expert > 512 || k <= 0 || k > 32) return false; if (variant < 0 || variant > 2) return false; launch_generic(logits, n_tokens, n_expert, k, ids, weights, stream, variant); return cudaGetLastError() == cudaSuccess; -#else - (void) logits; (void) n_tokens; (void) n_expert; (void) k; (void) ids; (void) weights; (void) stream; - (void) variant; - return false; -#endif } void router_top10(const float* logits, int n_tokens, int n_expert, int k, int* ids, float* weights, @@ -447,6 +441,24 @@ void router_top10(const float* logits, int n_tokens, int n_expert, int k, int* i return; } } +#endif +#if !defined(__HIPCC__) + { + // CUDA: the fast router too (bitwise router_top10_kernel's ids and weights - tests/cuda/router_fast.cpp checks + // it on the same rows as the HIP test, 256 x 8 included); the FP64 sum on one thread and the 20 barriers of + // the selection were 11% of qwen35moe's decode window on an RTX 5080. STRATA_ROUTER_OLD=1: the kernel above. + static const bool old = std::getenv("STRATA_ROUTER_OLD") != nullptr; + if (!old && n_tokens > 0 && n_expert > 64 && n_expert <= 512 && k > 0 && k <= 32) { + launch_generic(logits, n_tokens, n_expert, k, ids, weights, stream, 1); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { + std::fprintf(stderr, "router_top10 launch: %s\n", cudaGetErrorString(e)); + std::exit(1); + } + if (stream == nullptr && cudaDeviceSynchronize() != cudaSuccess) std::exit(1); + return; + } + } #endif if (n_tokens <= 0 || n_expert <= 0 || k <= 0) return; if (k > 64) { diff --git a/src/kernels/cuda/verify_kernels.cu b/src/kernels/cuda/verify_kernels.cu index ac2855d78f2..3217045790e 100644 --- a/src/kernels/cuda/verify_kernels.cu +++ b/src/kernels/cuda/verify_kernels.cu @@ -236,6 +236,13 @@ __device__ __forceinline__ void gdn_q8_1_store(GdnQ81* __restrict__ xq, size_t i if (idx % 32 == 0) xq[idx / 32].ds = q8_1_ds(d, sum); // #606: clamped (the QFUSE path bypassed the finite helper) } +// The GDN output gate on z: sigmoid (qwen4exp) or silu = z * sigmoid(z) (qwen35moe, llama.cpp's build_norm_gated). +// The sigmoid branch is the expression the kernels always used. +__device__ __forceinline__ float gdn_out_gate(float zz, int silu) { + const float sg = 1.0f / (1.0f + __expf(-zz)); + return silu ? zz * sg : sg; +} + template __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __restrict__ state, const float* __restrict__ hbuf, int C, @@ -245,7 +252,7 @@ __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __re const float* __restrict__ gamma, float eps, float* __restrict__ y, int h_k, int h_v, int T, const int32_t* __restrict__ n_keep, int t_out_begin, - GdnQ81* __restrict__ xq) { + GdnQ81* __restrict__ xq, int silu) { __shared__ float sk[2][S], sq[2][S]; __shared__ float red_kv[RG][S]; __shared__ float red_o[RG][S]; @@ -313,7 +320,7 @@ __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __re const float ss = wsum[0] + wsum[1] + wsum[2] + wsum[3]; const float scale = rsqrtf(ss / (float) S + eps); const float zz = z[(size_t) t * value_dim + head * S + col]; - const float yv = oc * scale * gam * (1.0f / (1.0f + __expf(-zz))); + const float yv = oc * scale * gam * gdn_out_gate(zz, silu); y[(size_t) t * value_dim + head * S + col] = yv; if constexpr (Q) gdn_q8_1_store(xq, (size_t) (t - t_out_begin) * value_dim + head * S + col, yv); } @@ -467,7 +474,7 @@ template __global__ void __launch_bounds__(S * RG) gdn_out_norm_kernel(const float* __restrict__ z, const float* __restrict__ gamma, float eps, float* __restrict__ y, int h_v, int T, const int32_t* __restrict__ n_keep, int t_out_begin, - GdnQ81* __restrict__ xq) { + GdnQ81* __restrict__ xq, int silu) { __shared__ float wsum[S * RG / 32]; const int head = blockIdx.x, t = blockIdx.y + t_out_begin, col = threadIdx.x, rg = threadIdx.y; const int tid = rg * S + col; @@ -487,11 +494,11 @@ __global__ void __launch_bounds__(S * RG) gdn_out_norm_kernel(const float* __res const float scale = rsqrtf(ss / (float) S + eps); const float zz = z[(size_t) t * value_dim + head * S + col]; if constexpr (Q) { - const float yv = oc * scale * gamma[col] * (1.0f / (1.0f + __expf(-zz))); + const float yv = oc * scale * gamma[col] * gdn_out_gate(zz, silu); y[(size_t) t * value_dim + head * S + col] = yv; gdn_q8_1_store(xq, (size_t) (t - t_out_begin) * value_dim + head * S + col, yv); } else { - y[(size_t) t * value_dim + head * S + col] = oc * scale * gamma[col] * (1.0f / (1.0f + __expf(-zz))); + y[(size_t) t * value_dim + head * S + col] = oc * scale * gamma[col] * gdn_out_gate(zz, silu); } } } @@ -990,6 +997,9 @@ void gdn_ab_multi(const float* x, const uint16_t* w_alpha, const uint16_t* w_bet check("gdn_ab_multi"); } +namespace { int g_gdn_silu = 0; } +void gdn_set_out_gate_silu(bool silu) { g_gdn_silu = silu ? 1 : 0; } + void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const float* gate, const float* beta, const float* z, const float* gamma, float eps, float* y, int h_k, int h_v, int n_tok, const int32_t* n_keep, void* stream, int t_out_begin, void* xq_out) { @@ -1011,8 +1021,8 @@ void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const gdn_step_split_kernel<<>>( state, h, conv_channels, gate, beta, y, h_k, h_v, n_tok, n_keep, t_out_begin); const dim3 ng((unsigned) h_v, (unsigned) (n_tok - t_out_begin)); - if (xq) gdn_out_norm_kernel<<>>(z, gamma, eps, y, h_v, n_tok, n_keep, t_out_begin, xq); - else gdn_out_norm_kernel<<>>(z, gamma, eps, y, h_v, n_tok, n_keep, t_out_begin, nullptr); + if (xq) gdn_out_norm_kernel<<>>(z, gamma, eps, y, h_v, n_tok, n_keep, t_out_begin, xq, g_gdn_silu); + else gdn_out_norm_kernel<<>>(z, gamma, eps, y, h_v, n_tok, n_keep, t_out_begin, nullptr, g_gdn_silu); } check("gdn_step_norm_multi (split)"); return; @@ -1034,15 +1044,15 @@ void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const const bool all_out = n_keep == nullptr && t_out_begin <= 0; if (xq && t_out_begin < n_tok) { if (all_out) gdn_step_norm_multi_kernel<<>>( - state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, nullptr, 0, xq); + state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, nullptr, 0, xq, g_gdn_silu); else gdn_step_norm_multi_kernel<<>>( - state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, n_keep, t_out_begin, xq); + state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, n_keep, t_out_begin, xq, g_gdn_silu); } else if (all_out) { gdn_step_norm_multi_kernel<<>>( - state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, nullptr, 0, nullptr); + state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, nullptr, 0, nullptr, g_gdn_silu); } else { gdn_step_norm_multi_kernel<<>>( - state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, n_keep, t_out_begin, nullptr); + state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, n_keep, t_out_begin, nullptr, g_gdn_silu); } check("gdn_step_norm_multi"); } diff --git a/src/kernels/iq_parity.cpp b/src/kernels/iq_parity.cpp index 4bc74f64def..b1066d66068 100644 --- a/src/kernels/iq_parity.cpp +++ b/src/kernels/iq_parity.cpp @@ -13,6 +13,7 @@ // activation rounding (a few 1e-3 relative), for 1 to 8 columns, every column of a multi-column call bitwise // equal to a one-column call on it. #include "strata/kernels/iq_kernels.hpp" +#include #include "strata/kernels/native_mmvq.hpp" #include @@ -26,7 +27,7 @@ int main(int argc, char** argv) { const std::string dir = argc > 1 ? argv[1] : "logs/iq_fixture"; - const char* names[] = {"IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ3_XXS", "IQ3_S", "IQ1_M", "IQ4_NL", "IQ4_XS", "Q2_0", "Q3_K"}; + const char* names[] = {"IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ3_XXS", "IQ3_S", "IQ1_M", "IQ4_NL", "IQ4_XS", "Q2_0", "Q3_K", "Q2_K", "Q4_0"}; int failures = 0, missing = 0; cudaStream_t s; cudaStreamCreate(&s); @@ -83,14 +84,23 @@ int main(int argc, char** argv) { strata::kernels::quantize_q8_1_rows(dx, MC, cols, xq, s); std::vector y((size_t) MC * rows), yn(y.size()); int multi_bad = 0; + // the dense kernels' MMVQ where they take the type, else the expert kernels' (Fmt<>) one - Q2_K, which + // only the Qwen3.6 MTP layer's experts use + // STRATA_PARITY_IQ_MMVQ=1: the expert kernels' MMVQ even where the dense one takes the type (Q4_0: the + // bartowski MTP layer's experts run on it) + static const bool force_iq = std::getenv("STRATA_PARITY_IQ_MMVQ") != nullptr; + auto mmvq = [](int t, const void* w, const void* xq_, float* y_, int n_in, int n_out, int nc, cudaStream_t st) { + if (strata::kernels::native_mmvq_supported(t) && !force_iq) + strata::kernels::native_mmvq(t, w, xq_, y_, n_in, n_out, nc, st); + else strata::kernels::iq_mmvq(t, w, xq_, y_, n_in, n_out, nc, st); + }; try { for (int c = 0; c < MC; ++c) - strata::kernels::native_mmvq(type, dw, (const uint8_t*) xq + (size_t) c * cols / 32 * 36, - dy + (size_t) c * rows, cols, rows, 1, s); + mmvq(type, dw, (const uint8_t*) xq + (size_t) c * cols / 32 * 36, dy + (size_t) c * rows, cols, rows, 1, s); cudaStreamSynchronize(s); cudaMemcpy(y.data(), dy, y.size() * 4, cudaMemcpyDeviceToHost); for (int nc = 2; nc <= MC; ++nc) { - strata::kernels::native_mmvq(type, dw, xq, dy, cols, rows, nc, s); + mmvq(type, dw, xq, dy, cols, rows, nc, s); cudaStreamSynchronize(s); cudaMemcpy(yn.data(), dy, yn.size() * 4, cudaMemcpyDeviceToHost); if (strata::kernels::native_mmvq_multi_exact() && diff --git a/src/kernels/kv_stream_parity.cpp b/src/kernels/kv_stream_parity.cpp index 722b7bc1cfb..80549d68fdb 100644 --- a/src/kernels/kv_stream_parity.cpp +++ b/src/kernels/kv_stream_parity.cpp @@ -8,6 +8,8 @@ // 2. the residency map is consistent after every call (slot_block and page_table invert each other); // 3. no call overflowed, and the hit/miss counters add up; // 4. a ring (the MTP drafter's layout) restored from the host copy reads the same values as the resident pool. +// 5. a dense selection (every cell, qwen35moe) larger than the slots: the blocks that do not fit are read from the +// host copy in place, bitwise the resident output, and none is left claimed. // INT8, FP16, Q4_0 (PR #21) and hybrid K8V4 pools; K8V4's appends are the engine's folded calls (layer.cpp), its host // copy written through kv_hybrid_k_half / kv_hybrid_v_half. #include "strata/kernels/kv_q4.hpp" @@ -233,6 +235,48 @@ bool run(int fmt) { std::printf(" %s ring restore: %s\n", name, ok ? "identical" : "DIFFERS"); if (!ok) ++bad; } + + // 5. a dense selection (qwen35moe: every cell) names more blocks than the slots hold: the resolve takes what fits, + // leaves no claimed (-2) block behind, and the attention reads the rest from the host copy in place - bitwise + // the resident pool's output, on the first call (slots taken) and on the next (every slot in use: none taken) + { + const int dq = 2; + int32_t* dids = dalloc((size_t) dq * N); + int32_t* dsteps = dalloc(dq * k::kStepCount); + float* dscratch = dalloc((size_t) dq * k::qsa_decode_attn_scratch_floats(N, s)); + std::vector hids((size_t) dq * N, 0), hst(dq * 4, 0); + for (int t = 0; t < dq; ++t) { + const int64_t p = N - dq + t, n_kv = p + 1; + for (int64_t c = 0; c < n_kv; ++c) hids[(size_t) t * N + c] = (int32_t) c; + hst[t * 4 + 0] = (int32_t) p; hst[t * 4 + 1] = (int32_t) n_kv; hst[t * 4 + 2] = (int32_t) (n_kv / 4); + hst[t * 4 + 3] = (int32_t) n_kv; + } + ck(cudaMemcpy(dids, hids.data(), hids.size() * 4, cudaMemcpyHostToDevice), "dense ids"); + ck(cudaMemcpy(dsteps, hst.data(), hst.size() * 4, cudaMemcpyHostToDevice), "dense steps"); + k::QsaAttnPools sp = slots.attn(m.page_table); + sp.host_k_pool = host.p.k_pool; sp.host_v_pool = host.p.v_pool; sp.host_k_q = host.p.k_q; sp.host_v_q = host.p.v_q; + sp.host_k_scale = host.p.k_scale; sp.host_v_scale = host.p.v_scale; sp.host_k_q4 = host.p.k_q4; sp.host_v_q4 = host.p.v_q4; + bool ok = true; + for (int call = 0; call < 2; ++call) { + k::qsa_decode_attn_batch(q, ref.attn(ident), dids, dsteps, N, s, dscratch, out_ref, dq, nullptr); + k::kv_stream_resolve(m, slots.attn(m.page_table), host.p, fmt, dids, dsteps, dq, N, s, nullptr); + k::qsa_decode_attn_batch(q, sp, dids, dsteps, N, s, dscratch, out_str, dq, nullptr); + ck(cudaDeviceSynchronize(), "dense"); + ck(cudaMemcpy(a.data(), out_ref, (size_t) dq * NH * D * 4, cudaMemcpyDeviceToHost), "a"); + ck(cudaMemcpy(b2.data(), out_str, (size_t) dq * NH * D * 4, cudaMemcpyDeviceToHost), "b"); + if (std::memcmp(a.data(), b2.data(), (size_t) dq * NH * D * 4) != 0) { + std::fprintf(stderr, " %s dense call %d: streamed attention differs\n", name, call); + ok = false; + } + ck(cudaMemcpy(pt.data(), m.page_table, n_blocks * 4, cudaMemcpyDeviceToHost), "pt"); + const int64_t claimed = std::count_if(pt.begin(), pt.end(), [](int32_t v) { return v < -1; }); + if (claimed != 0) { std::fprintf(stderr, " %s dense call %d: %lld blocks left claimed\n", name, call, (long long) claimed); ok = false; } + } + const bool over = k::kv_stream_counters(m).overflow; + std::printf(" %s dense (%lld blocks, %lld slots): %s, overflow %d\n", name, (long long) n_blocks, + (long long) n_slots, ok ? "identical" : "DIFFERS", (int) over); + if (!ok || !over) ++bad; + } return bad == 0; } } // namespace diff --git a/src/kernels/native_grouped_parity.cpp b/src/kernels/native_grouped_parity.cpp index 270de2cb66c..eecd19e5ca2 100644 --- a/src/kernels/native_grouped_parity.cpp +++ b/src/kernels/native_grouped_parity.cpp @@ -345,10 +345,14 @@ int main(int argc, char** argv) { #ifdef STRATA_Q6K_EXPERTS 14, #endif - 8}) // STRATA_GU_FMTS - for (int dt : {20, 23, 42, 7, 8}) check(gu, dt, 512, 256, s, rng); // STRATA_D_FMTS; IQ4_XS: n_ff % 256 + 8, 11, 10}) // STRATA_GU_FMTS + for (int dt : {20, 23, 42, 7, 8, 14, 12, 13, 18, 11}) check(gu, dt, 512, 256, s, rng); // STRATA_D_FMTS; IQ4_XS: n_ff % 256 check(21, 20, 2560, 640, s, rng); // a model's shapes check(21, 23, 2560, 768, s, rng); + check(21, 23, 2048, 512, s, rng); // Qwen3.6 UD-IQ4_XS + check(22, 14, 2048, 512, s, rng); // UD-IQ3_S's 6-bit layers + check(10, 12, 2048, 512, s, rng); // the UD-IQ3_S MTP layer + check(11, 12, 2048, 512, s, rng); // the UD-IQ4_XS MTP layer if (do_bench) bench(s, rng); std::printf("native_grouped_parity: %d failures\n", g_fail); cudaStreamDestroy(s); diff --git a/src/kernels/qsa_prompt_attn_parity.cpp b/src/kernels/qsa_prompt_attn_parity.cpp index 3c8867c5695..4077af03a98 100644 --- a/src/kernels/qsa_prompt_attn_parity.cpp +++ b/src/kernels/qsa_prompt_attn_parity.cpp @@ -298,6 +298,146 @@ int run(int fmt, int64_t ctx, int64_t nq, int reps) { // fmt 1 int8, 0 fp16, 2 cudaFree((void*) pl.k_q4); cudaFree((void*) pl.v_q4); return ok1 && ok2 ? 0 : 1; } +// Dense (Qwen3.6: 16 query heads, 2 KV heads, no selection - query i reads cells 0..pos): the 8-query kernel against the +// per-query tensor-core kernel it replaces (STRATA_PROMPT_ATTN_DENSE=0) and an FP64 reference. Same two checks as run(). +int run_dense(int fmt, int64_t ctx, int64_t nq, int reps) { // fmt 1 int8, 0 fp16 +#if defined(__HIPCC__) || defined(STRATA_USE_HIP) + (void) fmt; (void) ctx; (void) nq; (void) reps; + return 0; +#else + { + int dev = 0, major = 0; + ck(cudaGetDevice(&dev), "device"); + ck(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev), "cc"); + if (major < 8) { std::printf("SKIP dense: the 8-query kernel runs on sm_80+\n"); return 0; } + } + k::QsaShapes s = k::qsa_real_shapes(); + s.n_head = 8 * s.n_head_kv; + const int64_t HD = s.head_dim, NKV = s.n_head_kv, NH = s.n_head, PS = s.page_size; + const int64_t pages = (ctx + PS - 1) / PS, rows = pages * NKV * PS; + std::mt19937 rng(777 + fmt); + std::normal_distribution nd(0.f, 1.f); + std::uniform_int_distribution code(-127, 127); + std::uniform_real_distribution sc(0.005f, 0.03f); + std::vector kq, vq; + std::vector ks, vs, kh, vh; + if (fmt == 1) { + kq.resize(rows * HD); vq.resize(rows * HD); ks.resize(rows * 4); vs.resize(rows * 4); + for (auto& x : kq) x = (int8_t) code(rng); + for (auto& x : vq) x = (int8_t) code(rng); + for (auto& x : ks) x = f2h(sc(rng)); + for (auto& x : vs) x = f2h(sc(rng)); + } else { + kh.resize(rows * HD); vh.resize(rows * HD); + for (auto& x : kh) x = f2h(nd(rng) * 1.5f); + for (auto& x : vh) x = f2h(nd(rng)); + } + std::vector table(pages); + for (int64_t i = 0; i < pages; ++i) table[i] = (int32_t) i; + std::shuffle(table.begin(), table.end(), rng); + std::vector steps((size_t) (nq * k::kStepCount), 0); + for (int64_t i = 0; i < nq; ++i) steps[i * k::kStepCount + k::kStepWidth] = (int32_t) (ctx - nq + i + 1); + std::vector q((size_t) (nq * NH * HD)); + for (auto& x : q) x = nd(rng) * 2.0f; + k::QsaAttnPools pl; + if (fmt == 1) { pl.k_q = up(kq); pl.v_q = up(vq); pl.k_scale = up(ks); pl.v_scale = up(vs); } + else { pl.k_pool = up(kh); pl.v_pool = up(vh); } + pl.page_table = up(table); + const int32_t* d_steps = up(steps); + const float* d_q = up(q); + float *d_old = nullptr, *d_new = nullptr; + ck(cudaMalloc(&d_old, nq * NH * HD * 4), "malloc"); + ck(cudaMalloc(&d_new, nq * NH * HD * 4), "malloc"); + auto go = [&](bool dense, float* out) { +#if defined(_WIN32) + _putenv_s("STRATA_PROMPT_ATTN_DENSE", dense ? "1" : "0"); +#else + setenv("STRATA_PROMPT_ATTN_DENSE", dense ? "1" : "0", 1); +#endif + if (!k::qsa_prompt_attn_batch(d_q, pl, nullptr, d_steps, 0, s, out, nq, nullptr)) { + std::fprintf(stderr, "qsa_prompt_attn_batch refused the dense call\n"); + std::exit(2); + } + }; + go(false, d_old); + go(true, d_new); + ck(cudaDeviceSynchronize(), "run"); + std::vector o((size_t) (nq * NH * HD)), nw(o.size()); + ck(cudaMemcpy(o.data(), d_old, o.size() * 4, cudaMemcpyDeviceToHost), "down"); + ck(cudaMemcpy(nw.data(), d_new, nw.size() * 4, cudaMemcpyDeviceToHost), "down"); + double err_old = 0, err_new = 0, ref_scale = 0; + for (int64_t i = 0; i < nq; i += std::max(1, nq / 6)) { + const int64_t w = ctx - nq + i + 1; + for (int64_t h = 0; h < NH; h += 3) { + const int64_t kvh = h / (NH / NKV); + std::vector sco((size_t) w); + std::vector rowv((size_t) w); + double mx = -1e300; + for (int64_t c = 0; c < w; ++c) { + const int64_t row = ((int64_t) table[c / PS] * NKV + kvh) * PS + c % PS; + rowv[c] = row; + double a = 0; + for (int64_t d = 0; d < HD; ++d) { + const double kv = fmt == 1 ? (double) kq[row * HD + d] * h2f(ks[row * 4 + d / 64]) : h2f(kh[row * HD + d]); + a += (double) q[(i * NH + h) * HD + d] * kv; + } + sco[c] = a / 16.0; + mx = std::max(mx, sco[c]); + } + double l = 0; + for (auto& x : sco) { x = std::exp(x - mx); l += x; } + for (int64_t d = 0; d < HD; ++d) { + double a = 0; + for (int64_t c = 0; c < w; ++c) { + const int64_t row = rowv[c]; + a += sco[c] * (fmt == 1 ? (double) vq[row * HD + d] * h2f(vs[row * 4 + d / 64]) : h2f(vh[row * HD + d])); + } + const double r = a / l; + const size_t at = (size_t) ((i * NH + h) * HD + d); + ref_scale = std::max(ref_scale, std::fabs(r)); + err_old = std::max(err_old, std::fabs(o[at] - r)); + err_new = std::max(err_new, std::fabs(nw[at] - r)); + } + } + } + double diff = 0, scale = 0; + bool finite = true; + for (size_t i = 0; i < o.size(); ++i) { + finite = finite && std::isfinite(nw[i]); + diff = std::max(diff, (double) std::fabs(o[i] - nw[i])); + scale = std::max(scale, (double) std::fabs(o[i])); + } + cudaEvent_t e0, e1; + cudaEventCreate(&e0); + cudaEventCreate(&e1); + float ms_old = 0, ms_new = 0; + cudaEventRecord(e0); + for (int r = 0; r < reps; ++r) go(false, d_old); + cudaEventRecord(e1); + ck(cudaEventSynchronize(e1), "time"); + cudaEventElapsedTime(&ms_old, e0, e1); + cudaEventRecord(e0); + for (int r = 0; r < reps; ++r) go(true, d_new); + cudaEventRecord(e1); + ck(cudaEventSynchronize(e1), "time"); + cudaEventElapsedTime(&ms_new, e0, e1); +#if defined(_WIN32) + _putenv_s("STRATA_PROMPT_ATTN_DENSE", ""); +#else + unsetenv("STRATA_PROMPT_ATTN_DENSE"); +#endif + const bool ok1 = err_new <= std::max(4.0 * err_old, 1e-6 * ref_scale); + const bool ok2 = finite && diff <= 1e-4 * scale; + std::printf("%s dense %s ctx %lld, %lld queries: vs FP64 old %.3g new %.3g (output scale %.3g); new vs old %.3g (%.2g of " + "scale); %.3f -> %.3f ms per chunk (%.2fx)\n", + ok1 && ok2 ? "PASS" : "FAIL", fmt == 1 ? "int8" : "fp16", (long long) ctx, (long long) nq, err_old, err_new, + ref_scale, diff, diff / scale, ms_old / reps, ms_new / reps, ms_old / ms_new); + cudaFree((void*) d_steps); cudaFree((void*) d_q); cudaFree(d_old); cudaFree(d_new); + cudaFree((void*) pl.k_q); cudaFree((void*) pl.v_q); cudaFree((void*) pl.k_scale); cudaFree((void*) pl.v_scale); + cudaFree((void*) pl.k_pool); cudaFree((void*) pl.v_pool); cudaFree((void*) pl.page_table); + return ok1 && ok2 ? 0 : 1; +#endif +} } // namespace int main(int argc, char** argv) { @@ -330,6 +470,12 @@ int main(int argc, char** argv) { #endif fails += run(1, 1500, std::min(nq, 1500), reps); // short context: the selection is every cell fails += run(1, 2100, std::min(nq, 256), reps); // the identity-to-sparse edge + // Qwen3.6's dense prompt attention: a full chunk at the context given, and a short ragged one (a partial block of + // 8 queries, a partial last chunk of cells) + fails += run_dense(1, ctx, nq, reps); + fails += run_dense(0, ctx, nq, reps); + fails += run_dense(1, 1000, 37, reps); + fails += run_dense(0, 1000, 37, reps); std::printf("FAILURES: %d\n", fails); return fails; } diff --git a/src/prefill/kernels.cu b/src/prefill/kernels.cu index 18ec999aa57..00663fbbc48 100644 --- a/src/prefill/kernels.cu +++ b/src/prefill/kernels.cu @@ -41,6 +41,8 @@ __device__ __forceinline__ uint16_t bf(float f) { // W.hi + W.lo carries ~16 mantissa bits of the activation - the decode path's FP32 x to within ~1e-5. __device__ __forceinline__ uint16_t bf_lo(float f, uint16_t hi) { return bf(f - __uint_as_float((uint32_t) hi << 16)); } __device__ __forceinline__ float sigm(float x) { return 1.0f / (1.0f + __expf(-x)); } +// the GDN output gate on z: sigmoid (Flash-Next) or silu (Qwen3.6: llama.cpp's build_norm_gated) +template __device__ __forceinline__ float gdn_gate_act(float z) { return SILU ? z * sigm(z) : sigm(z); } __device__ __forceinline__ uint16_t hf(float f) { return __half_as_ushort(__float2half_rn(f)); } // A SwiGLU product for an FP16 GEMM: saturated, so a token with a massive activation cannot turn into inf and then // NaN in the down projection (decode's q8_1 has room to ~8e6; FP16 ends at 65504). A NaN stays NaN (fminf/fmaxf @@ -113,6 +115,29 @@ __global__ void gr_norm_rs_kernel(const float* __restrict__ R, const float* __re if (xn16_lo) xn16_lo[xo + d] = bf_lo(v, h); } } +// One residual stream (qwen35moe): the RMSNorm of each row times w (llama.cpp's build_norm), as gr_mix_r writes the +// mixer's input - FP32, the 16-bit activation image (+ its BF16 low part) and FP16 +__global__ void rms_mix_kernel(const float* __restrict__ R, const float* __restrict__ w, float eps, int64_t n, + float* __restrict__ mixed, uint16_t* __restrict__ mixed16, uint16_t* __restrict__ mixed_h, + uint16_t* __restrict__ mixed16_lo) { + __shared__ float sh[32]; + const int64_t row = blockIdx.x; + const float* r = R + row * n; + float ss = 0.0f; + for (int64_t d = threadIdx.x; d < n; d += blockDim.x) ss += r[d] * r[d]; + const float rs = rsqrtf(block_sum(ss, sh) / (float) n + eps); + for (int64_t d = threadIdx.x; d < n; d += blockDim.x) { + const float v = r[d] * rs * w[d]; + const int64_t i = row * n + d; + mixed[i] = v; + if (mixed16) { + const uint16_t h = act16(v); + mixed16[i] = h; + if (mixed16_lo) mixed16_lo[i] = bf_lo(v, h); + } + if (mixed_h) mixed_h[i] = hf(v); + } +} __global__ void gr_mix_r_kernel(const float* __restrict__ R, const float* __restrict__ rs, const float* __restrict__ w, const float* __restrict__ g, float* __restrict__ mixed, uint16_t* __restrict__ mixed16, int64_t T, uint16_t* __restrict__ mixed_h, uint16_t* __restrict__ mixed16_lo) { @@ -280,6 +305,7 @@ __global__ void gr_broadcast_kernel(const float* __restrict__ e, float* __restri } // ---------------------------------------------------------------- GDN +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void gdn_gates_kernel(const float* __restrict__ ab, const float* __restrict__ dt, const float* __restrict__ ssm_a, float* __restrict__ gate, float* __restrict__ beta, int64_t T) { @@ -291,6 +317,7 @@ __global__ void gdn_gates_kernel(const float* __restrict__ ab, const float* __re beta[i] = sigm(ab[t * 2 * HV + HV + h]); } // one thread per channel, walks the chunk; then a second kernel normalises +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void gdn_conv_kernel(float* __restrict__ hist, const float* __restrict__ qkv, const float* __restrict__ w, float* __restrict__ h, int64_t T) { const int c = blockIdx.x * blockDim.x + threadIdx.x; @@ -309,6 +336,7 @@ __global__ void gdn_conv_kernel(float* __restrict__ hist, const float* __restric // chunk (or the history before it) instead of carrying them - the conv reads inputs, not its own outputs, so the // tiles are independent. The same expression per element (so the same bits); the history is written afterwards. constexpr int CONV_TILE = 64; +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void gdn_conv_tiled_kernel(const float* __restrict__ hist, const float* __restrict__ qkv, const float* __restrict__ w, float* __restrict__ h, int64_t T) { const int c = blockIdx.x * blockDim.x + threadIdx.x; @@ -330,6 +358,7 @@ __global__ void gdn_conv_tiled_kernel(const float* __restrict__ hist, const floa // are the 128 channels of one head, so for the q / k heads (block < 2 HK) the norm's sum of squares is the block's: // the same warp_sum, the same four partials added in the same order, the same v * rsqrtf(ss + eps) on the same // stored value - the same bits, without writing h and reading it back (16 tokens per round). +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void gdn_conv_l2_kernel(const float* __restrict__ hist, const float* __restrict__ qkv, const float* __restrict__ w, float* __restrict__ h, int64_t T, float eps) { const int c = blockIdx.x * blockDim.x + threadIdx.x; @@ -367,6 +396,7 @@ __global__ void gdn_conv_l2_kernel(const float* __restrict__ hist, const float* } } // the history after the chunk: its last three inputs (the older history where the chunk is shorter than 3) +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void gdn_conv_hist_kernel(float* __restrict__ hist, const float* __restrict__ qkv, int64_t T) { const int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= C) return; @@ -377,6 +407,7 @@ __global__ void gdn_conv_hist_kernel(float* __restrict__ hist, const float* __re } hist[c * 3] = v[0]; hist[c * 3 + 1] = v[1]; hist[c * 3 + 2] = v[2]; } +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void gdn_l2_kernel(float* __restrict__ h, float eps) { // block (t, head) over the 32 q/k heads, 128 threads const int64_t t = blockIdx.y; @@ -391,6 +422,7 @@ __global__ void gdn_l2_kernel(float* __restrict__ h, float eps) { x[threadIdx.x] = v * rsqrtf(ss + eps); } constexpr int RG = 4, RPG = S / RG; +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(S * RG) gdn_rec_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, const float* __restrict__ z, @@ -438,7 +470,7 @@ __global__ void __launch_bounds__(S * RG) gdn_rec_kernel(float* __restrict__ sta __syncthreads(); if (rg == 0) { const float ss = wsum[0] + wsum[1] + wsum[2] + wsum[3]; - const float v = oc * rsqrtf(ss / (float) S + eps) * g_col * sigm(z[t * HV * S + head * S + col]); + const float v = oc * rsqrtf(ss / (float) S + eps) * g_col * gdn_gate_act(z[t * HV * S + head * S + col]); y[t * HV * S + head * S + col] = v; y16[t * ld16 + head * S + col] = hf(v); } @@ -452,6 +484,7 @@ __global__ void __launch_bounds__(S * RG) gdn_rec_kernel(float* __restrict__ sta // in its own kernel. Per column the same arithmetic in the same order (the 4 row-group partial sums added as // red[0] + red[1] + red[2] + red[3]; the norm's warp sums over the same 32-column warps): the same bits. constexpr int CB = 32, NCB = S / CB; +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(CB * RG) gdn_rec_cols_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, @@ -495,6 +528,7 @@ __global__ void __launch_bounds__(CB * RG) gdn_rec_cols_kernel(float* __restrict // gdn_rec_cols_kernel with the next token's inputs (q/k rows, v, gate, beta) loaded into registers while this token // computes (software pipelining). The same arithmetic in the same order: the same bits, and the same CB-column split. // STRATA_GDN_PIPELINE=0: gdn_rec_cols_kernel. +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(CB * RG) gdn_rec_cols_pipe_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, @@ -592,11 +626,15 @@ __device__ __forceinline__ void gdn_cp_wait_prev() { // every group but the ne asm volatile("cp.async.wait_group 1;\n" ::); #endif } -constexpr int GDN_TB = 8, VPK = HV / HK; // tokens per staged block, value heads per key head +constexpr int GDN_TB = 8; // tokens per staged block +constexpr int VPK = HV / HK; // Flash-Next's value heads per key head (the chunked recurrence; the templated kernels + // take their own HV / HK) +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(CB * RG) gdn_rec_kh_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, float* __restrict__ oc_out, int64_t T) { + constexpr int VPK = HV / HK; // value heads per key head constexpr int TB = GDN_TB, NT = CB * RG, QKP = S / 4, VP = CB / 4; // threads, 16-byte pieces of a q/k row, of v __shared__ __align__(16) float sq[2][TB][S]; __shared__ __align__(16) float sk[2][TB][S]; @@ -693,6 +731,7 @@ __global__ void __launch_bounds__(CB * RG) gdn_rec_kh_kernel(float* __restrict__ // 64 SMs and more, 1.02-1.04x / 1.03-1.06x at 48 to 63 (two blocks against four: a draw, slower on neither), // 1.28-1.31x / 1.28-1.30x at 39 to 47, 1.53-1.57x / 1.54-1.55x at 32 to 38. Per call, from the current device (a // layer split can mix cards). +template bool gdn_keyhead_ok() { static const bool off = [] { const char* v = std::getenv("STRATA_GDN_KEYHEAD"); return v != nullptr && std::atoi(v) == 0; }(); if (off) return false; @@ -703,7 +742,7 @@ bool gdn_keyhead_ok() { int major = 0, sms = 0, per_sm = 0; const bool yes = cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) == cudaSuccess && cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev) == cudaSuccess && major >= 8 && - cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, gdn_rec_kh_kernel, CB * RG, 0) == + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, gdn_rec_kh_kernel, CB * RG, 0) == cudaSuccess && (int64_t) per_sm * sms >= (int64_t) HK * NCB; if (!yes) cudaGetLastError(); @@ -1078,6 +1117,7 @@ constexpr int QTH = QCB * 4; // threads per block: 4 lanes pe constexpr int QRS = 36; // staged row-group stride (32 rows + 4 pad: the 4 groups' reads // land in different LDS banks) constexpr int QPF = HT * 2 * S / 4 / QTH; // float4 of q/k each thread prefetches per chunk (2) +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(QTH) gdn_rec_quad_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, float* __restrict__ oc_out, int64_t T) { @@ -1186,6 +1226,7 @@ __global__ void __launch_bounds__(QTH) gdn_rec_quad_kernel(float* __restrict__ s // that need 9.6): the stores' lane offset goes through an asm so the addresses are made again, 133 VGPRs, all blocks // resident. (2) A full chunk's four tokens run with the state alternating between two register sets (the update // writes the other set), so the compiler emits no 32 register copies after the update. +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(QTH) gdn_rec_quad_pp_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, float* __restrict__ oc_out, int64_t T) { @@ -1318,6 +1359,7 @@ __global__ void __launch_bounds__(QTH) gdn_rec_quad_pp_kernel(float* __restrict_ // and two independent dependency chains in each wave). Per column the same arithmetic in the same order. constexpr int C2CB = 64; // columns per block constexpr int C2PF = HT * 2 * S / 4 / QTH; // float4 of q/k each thread prefetches per chunk (2) +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(QTH) gdn_rec_quad2c_kernel(float* __restrict__ state, const float* __restrict__ h, const float* __restrict__ gate, const float* __restrict__ beta, float* __restrict__ oc_out, int64_t T) { @@ -1436,7 +1478,7 @@ __global__ void __launch_bounds__(QTH) gdn_rec_quad2c_kernel(float* __restrict__ // gdn_out_norm_kernel's body over a grid-stride walk of the (token, head) rows: the old kernel launched a 128-thread // block per row (T * 48 blocks; ~45 GB/s on the 8060S). The same code per row (the same compiled arithmetic: a wave // per row with four values per lane rounded differently in its first xor step). -template +template __global__ void __launch_bounds__(S) gdn_out_norm_loop_kernel(const float* __restrict__ z, const float* __restrict__ gamma, float eps, float* __restrict__ y, uint16_t* __restrict__ y16, int64_t rows, int64_t ld16) { @@ -1454,7 +1496,7 @@ __global__ void __launch_bounds__(S) gdn_out_norm_loop_kernel(const float* __res if ((col & 31) == 0) wsum[col >> 5] = sp; __syncthreads(); const float ss = wsum[0] + wsum[1] + wsum[2] + wsum[3]; - const float bp = oc * rsqrtf(ss / (float) S + eps) * gamma[col], sg = sigm(z[at]); + const float bp = oc * rsqrtf(ss / (float) S + eps) * gamma[col], sg = gdn_gate_act(z[at]); const size_t o16 = (size_t) (row / HV) * ld16 + (size_t) (row % HV) * S + col; if constexpr (WY) { const float v = bp * sg; @@ -1470,6 +1512,7 @@ __global__ void __launch_bounds__(S) gdn_out_norm_loop_kernel(const float* __res } } #endif // __HIPCC__ +template // value heads and conv channels: Flash-Next's 48 / 10240, Qwen3.6's 32 / 8192 __global__ void __launch_bounds__(S) gdn_out_norm_kernel(const float* __restrict__ z, const float* __restrict__ gamma, float eps, float* __restrict__ y, uint16_t* __restrict__ y16, int64_t ld16) { @@ -1482,12 +1525,12 @@ __global__ void __launch_bounds__(S) gdn_out_norm_kernel(const float* __restrict if ((col & 31) == 0) wsum[col >> 5] = sp; __syncthreads(); const float ss = wsum[0] + wsum[1] + wsum[2] + wsum[3]; - const float v = oc * rsqrtf(ss / (float) S + eps) * gamma[col] * sigm(z[t * HV * S + head * S + col]); + const float v = oc * rsqrtf(ss / (float) S + eps) * gamma[col] * gdn_gate_act(z[t * HV * S + head * S + col]); y16[(size_t) t * ld16 + (size_t) head * S + col] = hf(v); // (the FP32 normalized value is not stored: nothing reads it) } // ---------------------------------------------------------------- MoE -template +template // KS: experts per token (Flash-Next 10, Qwen3.6 8) __global__ void route_kernel(const float* __restrict__ logits, int32_t* __restrict__ ids, float* __restrict__ wout, int64_t T) { const int64_t t = (int64_t) blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); @@ -1508,7 +1551,7 @@ __global__ void route_kernel(const float* __restrict__ logits, int32_t* __restri #pragma unroll for (int i = 0; i < REG; ++i) { v[i] *= rcp; if (isnan(v[i])) v[i] = -FLT_MAX; } float selected = 0.0f, selected_sum = 0.0f; - for (int rank = 0; rank < 10; ++rank) { + for (int rank = 0; rank < KS; ++rank) { float best = v[0]; int ex = lane; #pragma unroll @@ -1520,11 +1563,11 @@ __global__ void route_kernel(const float* __restrict__ logits, int32_t* __restri if (ob > best || (ob == best && oi < ex)) { best = ob; ex = oi; } } if ((ex & 31) == lane) { v[ex / 32] = -INFINITY; selected_sum += best; } - if (lane == 0) ids[t * 10 + rank] = ex; + if (lane == 0) ids[t * KS + rank] = ex; if (rank == lane) selected = best; } selected_sum = fmaxf(warp_sum(selected_sum), 6.103515625e-5f); - if (lane < 10) wout[t * 10 + lane] = selected / selected_sum; + if (lane < KS) wout[t * KS + lane] = selected / selected_sum; } // Strata blob: gate/up codes [1280][640 B], down codes [2560][160 B], gate/up scales [1280][40] f16, down scales [2560][10] f16 template @@ -1552,17 +1595,17 @@ __global__ void blob_dequant_kernel(const uint8_t* __restrict__ blob, uint16_t* for (int k = 0; k < 4; ++k) { const float v = (float) (((c >> (2 * k)) & 3) - 1) * d; o[k] = HALF ? hf(v) : bf(v); } } } -__global__ void swiglu_il_kernel(const float* __restrict__ gu, uint16_t* __restrict__ h16, int64_t n) { +__global__ void swiglu_il_kernel(const float* __restrict__ gu, uint16_t* __restrict__ h16, int64_t n, int64_t ff) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n * 640) return; - const int64_t r = i / 640, k = i % 640; - const float g = gu[r * 1280 + 2 * k], u = gu[r * 1280 + 2 * k + 1]; + if (i >= n * ff) return; + const int64_t r = i / ff, k = i % ff; + const float g = gu[r * 2 * ff + 2 * k], u = gu[r * 2 * ff + 2 * k + 1]; h16[i] = hf_sat(g / (1.0f + __expf(-g)) * u); } __global__ void swiglu_pair_kernel(const float* __restrict__ g, const float* __restrict__ u, uint16_t* __restrict__ h16, - int64_t n) { + int64_t n, int64_t ff) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n * 640) return; + if (i >= n * ff) return; const float a = g[i]; h16[i] = hf_sat(a / (1.0f + __expf(-a)) * u[i]); } @@ -1574,6 +1617,7 @@ __global__ void gather_rows16_kernel(const uint16_t* __restrict__ x, const int32 const int64_t r = i / per, j = i % per; reinterpret_cast(dst)[r * per + j] = reinterpret_cast(x)[(int64_t) src[r] * per + j]; } +template // n_embd and experts per token __global__ void moe_combine_kernel(const float* __restrict__ Dm, const int32_t* __restrict__ slot, const float* __restrict__ w, const float* __restrict__ shared, const float* __restrict__ sg, float* __restrict__ bo, int64_t T) { @@ -1582,12 +1626,13 @@ __global__ void moe_combine_kernel(const float* __restrict__ Dm, const int32_t* const int64_t t = i / N, d = i % N; float s = 0.0f; #pragma unroll - for (int k = 0; k < 10; ++k) s = fmaf(w[t * 10 + k], Dm[(int64_t) slot[t * 10 + k] * N + d], s); + for (int k = 0; k < KS; ++k) s = fmaf(w[t * KS + k], Dm[(int64_t) slot[t * KS + k] * N + d], s); bo[i] = s + shared[i] * sigm(sg[t]); } // W (X5): four columns per thread, one float4 load per slot; the per-element expression and the k order are unchanged // (bitwise equal to moe_combine_kernel). Needs 16-byte aligned Dm / shared / bo (N * 4 bytes per row is a multiple of 16). +template __global__ void moe_combine4_kernel(const float* __restrict__ Dm, const int32_t* __restrict__ slot, const float* __restrict__ w, const float4* __restrict__ shared, const float* __restrict__ sg, float4* __restrict__ bo, int64_t T) { @@ -1595,15 +1640,15 @@ __global__ void moe_combine4_kernel(const float* __restrict__ Dm, const int32_t* if (i >= T * (N / 4)) return; const int64_t t = i / (N / 4), d = (i % (N / 4)) * 4; float4 s = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 v[10]; - float ww[10]; + float4 v[KS]; + float ww[KS]; #pragma unroll - for (int k = 0; k < 10; ++k) { - ww[k] = w[t * 10 + k]; - v[k] = *reinterpret_cast(Dm + (int64_t) slot[t * 10 + k] * N + d); + for (int k = 0; k < KS; ++k) { + ww[k] = w[t * KS + k]; + v[k] = *reinterpret_cast(Dm + (int64_t) slot[t * KS + k] * N + d); } #pragma unroll - for (int k = 0; k < 10; ++k) { + for (int k = 0; k < KS; ++k) { s.x = fmaf(ww[k], v[k].x, s.x); s.y = fmaf(ww[k], v[k].y, s.y); s.z = fmaf(ww[k], v[k].z, s.z); s.w = fmaf(ww[k], v[k].w, s.w); } @@ -1646,18 +1691,20 @@ __global__ void rope_kernel(float* __restrict__ x, int64_t heads, int64_t dim, i p[pair] = a * c - b * s; p[pair + 32] = a * s + b * c; } +template // query heads: Flash-Next's 24, Qwen3.6's 16 __global__ void split_q_kernel(const float* __restrict__ qf, float* __restrict__ q, int64_t T) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; - if (i >= T * 24 * 256) return; - const int64_t t = i / (24 * 256), h = (i / 256) % 24, d = i % 256; - q[i] = qf[t * 24 * 512 + h * 512 + d]; + if (i >= T * NH * 256) return; + const int64_t t = i / (NH * 256), h = (i / 256) % NH, d = i % 256; + q[i] = qf[t * NH * 512 + h * 512 + d]; } +template __global__ void gate_attn_kernel(const float* __restrict__ a, const float* __restrict__ qf, uint16_t* __restrict__ o16, int64_t T, int64_t ld16) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; - if (i >= T * 24 * 256) return; - const int64_t t = i / (24 * 256), h = (i / 256) % 24, d = i % 256; - o16[t * ld16 + h * 256 + d] = hf(a[i] * (1.0f / (1.0f + expf(-qf[t * 24 * 512 + h * 512 + 256 + d])))); + if (i >= T * NH * 256) return; + const int64_t t = i / (NH * 256), h = (i / 256) % NH, d = i % 256; + o16[t * ld16 + h * 256 + d] = hf(a[i] * (1.0f / (1.0f + expf(-qf[t * NH * 512 + h * 512 + 256 + d])))); } // one block per (token, kv head, 64-value group); 64 threads. KV streaming: the pool page only if the block is @@ -1826,6 +1873,12 @@ void round_f16(const float* x, float* y, int64_t n, void* stream) { round_f16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n); check("round_f16"); } +void rms_mix(const float* R, const float* w, float eps, int64_t n, float* mixed, uint16_t* mixed16, int64_t T, void* stream, + uint16_t* mixed_h, uint16_t* mixed16_lo) { + if (T <= 0) return; + rms_mix_kernel<<<(unsigned) T, 256, 0, (cudaStream_t) stream>>>(R, w, eps, n, mixed, mixed16, mixed_h, mixed16_lo); + check("rms_mix"); +} void to_bf16(const float* x, uint16_t* y, int64_t n, void* stream, uint16_t* ylo) { if (n <= 0) return; to_bf16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, ylo, n); @@ -1900,49 +1953,52 @@ void gr_broadcast(const float* e, float* R, int64_t T, void* stream) { gr_broadcast_kernel<<>>(e, R, T); check("gr_broadcast"); } -void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream) { - gdn_gates_kernel<<>>(ab, dt, ssm_a, gate, beta, T); +template +void gdn_gates_t(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream) { + gdn_gates_kernel<<>>(ab, dt, ssm_a, gate, beta, T); check("gdn_gates"); } -void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream) { +template +void gdn_conv_t(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream) { static const bool serial = std::getenv("STRATA_GDN_CONV_SERIAL") != nullptr; // the old walk (A/B) if (serial || T <= CONV_TILE) { - gdn_conv_kernel<<>>(history, qkv, conv_w, h, T); + gdn_conv_kernel<<>>(history, qkv, conv_w, h, T); } else { static const bool fused = [] { const char* v = std::getenv("STRATA_GDN_CONVL2"); return v && std::atoi(v) != 0; }(); if (fused) { - gdn_conv_l2_kernel<<<<>>(history, qkv, conv_w, h, T, eps); - gdn_conv_hist_kernel<<>>(history, qkv, T); + gdn_conv_hist_kernel<<>>(history, qkv, T); check("gdn_conv"); return; } - gdn_conv_tiled_kernel<<<<>>(history, qkv, conv_w, h, T); - gdn_conv_hist_kernel<<>>(history, qkv, T); + gdn_conv_hist_kernel<<>>(history, qkv, T); } - gdn_l2_kernel<<>>(h, eps); + gdn_l2_kernel<<>>(h, eps); check("gdn_conv"); } -void gdn_recurrence_variant(int variant, float* state, const float* h, const float* gate, const float* beta, - const float* z, const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, - void* stream, int64_t ld16) { +template +void gdn_recurrence_variant_t(int variant, float* state, const float* h, const float* gate, const float* beta, + const float* z, const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, + void* stream, int64_t ld16) { if (ld16 <= 0) ld16 = (int64_t) HV * S; #if defined(__HIPCC__) if (variant == 3) { // diagnostics: the quad recurrence + the old norm kernel if (T > 0) { - gdn_rec_quad_kernel<<>>(state, h, gate, beta, y, T); - gdn_out_norm_kernel<<>>(z, gamma, eps, y, y16, ld16); + gdn_rec_quad_kernel<<>>(state, h, gate, beta, y, T); + gdn_out_norm_kernel<<>>(z, gamma, eps, y, y16, ld16); } return; } if (variant == 4 || variant == 5) { // diagnostics: the quad recurrence's PP / two-column kernels + the norm if (T > 0) { if (variant == 5) - gdn_rec_quad2c_kernel<<>>(state, h, gate, beta, y, T); + gdn_rec_quad2c_kernel<<>>(state, h, gate, beta, y, T); else - gdn_rec_quad_pp_kernel<<>>(state, h, gate, beta, y, T); - gdn_out_norm_loop_kernel<<<(unsigned) std::min(T * HV, 4096), S, 0, (cudaStream_t) stream>>>( + gdn_rec_quad_pp_kernel<<>>(state, h, gate, beta, y, T); + gdn_out_norm_loop_kernel<<<(unsigned) std::min(T * HV, 4096), S, 0, (cudaStream_t) stream>>>( z, gamma, eps, y, y16, T * HV, ld16); } check("gdn_recurrence (quad pp)"); @@ -1952,17 +2008,17 @@ void gdn_recurrence_variant(int variant, float* state, const float* h, const flo if (T > 0) { static const int pp = [] { const char* v = std::getenv("STRATA_GDN_PP"); return v ? std::atoi(v) : 0; }(); if (pp == 2) - gdn_rec_quad2c_kernel<<>>(state, h, gate, beta, y, T); + gdn_rec_quad2c_kernel<<>>(state, h, gate, beta, y, T); else if (pp) - gdn_rec_quad_pp_kernel<<>>(state, h, gate, beta, y, T); + gdn_rec_quad_pp_kernel<<>>(state, h, gate, beta, y, T); else - gdn_rec_quad_kernel<<>>(state, h, gate, beta, y, T); + gdn_rec_quad_kernel<<>>(state, h, gate, beta, y, T); static const bool noy = [] { const char* v = std::getenv("STRATA_GDN_NOY"); return v && std::atoi(v) != 0; }(); if (noy) - gdn_out_norm_loop_kernel<<<(unsigned) std::min(T * HV, 4096), S, 0, (cudaStream_t) stream>>>( + gdn_out_norm_loop_kernel<<<(unsigned) std::min(T * HV, 4096), S, 0, (cudaStream_t) stream>>>( z, gamma, eps, y, y16, T * HV, ld16); else - gdn_out_norm_loop_kernel<<<(unsigned) std::min(T * HV, 4096), S, 0, (cudaStream_t) stream>>>( + gdn_out_norm_loop_kernel<<<(unsigned) std::min(T * HV, 4096), S, 0, (cudaStream_t) stream>>>( z, gamma, eps, y, y16, T * HV, ld16); } check("gdn_recurrence (quad)"); @@ -1973,42 +2029,74 @@ void gdn_recurrence_variant(int variant, float* state, const float* h, const flo #endif static const bool serial = std::getenv("STRATA_GDN_REC_HEADS") != nullptr; // the one-block-per-head kernel (A/B) if (serial || T <= 0 || variant == 2) { - gdn_rec_kernel<<>>(state, h, gate, beta, z, gamma, eps, y, y16, T, ld16); + gdn_rec_kernel<<>>(state, h, gate, beta, z, gamma, eps, y, y16, T, ld16); } else { static const bool pipe = [] { const char* v = std::getenv("STRATA_GDN_PIPELINE"); return v == nullptr || std::atoi(v) != 0; }(); #if !defined(__HIPCC__) // STRATA_GDN_CHUNKED=1 (opt-in): the recurrence in chunks for prompt chunks of kGdnChunkedMin+ tokens (other // bits; a shorter chunk, a card it does not fit or a failed scratch allocation takes the kernels below) static const bool chunked = [] { const char* v = std::getenv("STRATA_GDN_CHUNKED"); return v != nullptr && std::atoi(v) != 0; }(); - if (chunked && T >= kGdnChunkedMin && gdn_rec_chunked(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) { - } else if (pipe && gdn_keyhead_ok()) // the value heads of a key head in one thread (same bits) - gdn_rec_kh_kernel<<>>(state, h, gate, beta, y, T); + // (Flash-Next's 48 value heads only: the chunked kernels have no instance for Qwen3.6's 32) + if (chunked && HV == 48 && C == 10240 && T >= kGdnChunkedMin && gdn_rec_chunked(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) { + } else if (pipe && gdn_keyhead_ok()) // the value heads of a key head in one thread (same bits) + gdn_rec_kh_kernel<<>>(state, h, gate, beta, y, T); else #endif if (pipe) // the software-pipelined loads (same bits) - gdn_rec_cols_pipe_kernel<<>>(state, h, gate, beta, y, T); + gdn_rec_cols_pipe_kernel<<>>(state, h, gate, beta, y, T); else - gdn_rec_cols_kernel<<>>(state, h, gate, beta, y, T); - gdn_out_norm_kernel<<>>(z, gamma, eps, y, y16, ld16); + gdn_rec_cols_kernel<<>>(state, h, gate, beta, y, T); + gdn_out_norm_kernel<<>>(z, gamma, eps, y, y16, ld16); } check("gdn_recurrence"); } +// the instances: Flash-Next (48 value heads, 10240 conv channels, sigmoid gate) and Qwen3.6 (32, 8192, silu) +namespace { +[[noreturn]] void gdn_bad(int64_t hv, bool silu) { + std::fprintf(stderr, "prefill GDN: no kernels for %lld value heads with the %s gate\n", (long long) hv, silu ? "silu" : "sigmoid"); + std::exit(1); +} +} // namespace +void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream, + int64_t hv) { + if (hv == 48) gdn_gates_t<48, 10240>(ab, dt, ssm_a, gate, beta, T, stream); + else if (hv == 32) gdn_gates_t<32, 8192>(ab, dt, ssm_a, gate, beta, T, stream); + else gdn_bad(hv, false); +} +void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream, int64_t hv) { + if (hv == 48) gdn_conv_t<48, 10240>(history, qkv, conv_w, h, T, eps, stream); + else if (hv == 32) gdn_conv_t<32, 8192>(history, qkv, conv_w, h, T, eps, stream); + else gdn_bad(hv, false); +} +void gdn_recurrence_variant(int variant, float* state, const float* h, const float* gate, const float* beta, + const float* z, const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, + void* stream, int64_t ld16, int64_t hv, bool silu) { + if (hv == 48 && !silu) + gdn_recurrence_variant_t<48, 10240, false>(variant, state, h, gate, beta, z, gamma, eps, y, y16, T, stream, ld16); + else if (hv == 32 && silu) + gdn_recurrence_variant_t<32, 8192, true>(variant, state, h, gate, beta, z, gamma, eps, y, y16, T, stream, ld16); + else gdn_bad(hv, silu); +} void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z, - const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream, int64_t ld16) { + const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream, int64_t ld16, + int64_t hv, bool silu) { // Aurora (S23): four lanes per column + the grid-stride norm, the same bits (tests/hip/gdn_rec_head.cpp); // STRATA_GDN_HEAD=1 (on by default on gfx1151; off with STRATA_GDN_HEAD=0 or STRATA_GDN_REC_HEADS) takes them // (AMD builds only; the arch defaults set STRATA_GDN_HEAD=1 on gfx1151 and leave other cards on the earlier kernels) static const bool head = [] { const char* v = std::getenv("STRATA_GDN_HEAD"); return v != nullptr && std::atoi(v) != 0; }(); static const bool serial = std::getenv("STRATA_GDN_REC_HEADS") != nullptr; - gdn_recurrence_variant(head && !serial ? 1 : 0, state, h, gate, beta, z, gamma, eps, y, y16, T, stream, ld16); + gdn_recurrence_variant(head && !serial ? 1 : 0, state, h, gate, beta, z, gamma, eps, y, y16, T, stream, ld16, hv, + silu); } -void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream) { - if (n_expert == 512) +void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream, int64_t k) { + if (n_expert == 512 && k == 10) route_kernel<16><<<(unsigned) ((T + 7) / 8), 256, 0, (cudaStream_t) stream>>>(logits, ids, weights, T); - else if (n_expert == 256) + else if (n_expert == 256 && k == 10) route_kernel<8><<<(unsigned) ((T + 7) / 8), 256, 0, (cudaStream_t) stream>>>(logits, ids, weights, T); + else if (n_expert == 256 && k == 8) + route_kernel<8, 8><<<(unsigned) ((T + 7) / 8), 256, 0, (cudaStream_t) stream>>>(logits, ids, weights, T); else - strata::kernels::router_top10(logits, (int) T, (int) n_expert, 10, ids, weights, stream); + strata::kernels::router_top10(logits, (int) T, (int) n_expert, (int) k, ids, weights, stream); check("route"); } void blob_dequant(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream) { @@ -2019,13 +2107,13 @@ void blob_dequant_f16(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, voi blob_dequant_kernel<<>>(blob, gu16, down16); check("blob_dequant_f16"); } -void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream) { +void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream, int64_t n_ff) { if (n <= 0) return; - swiglu_il_kernel<<>>(gu, h16, n); + swiglu_il_kernel<<>>(gu, h16, n, n_ff); check("swiglu_interleaved"); } -void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream) { - swiglu_pair_kernel<<>>(g, u, h16, n); +void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream, int64_t n_ff) { + swiglu_pair_kernel<<>>(g, u, h16, n, n_ff); check("swiglu_pair"); } namespace { @@ -2126,7 +2214,27 @@ __global__ void f16_to_f32_wide_kernel(float4* __restrict__ dst, const uint4* __ dst[2 * i] = make_float4(a.x, a.y, b.x, b.y); dst[2 * i + 1] = make_float4(c.x, c.y, e.x, e.y); } +template +void moe_combine_t(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, + int64_t T, void* stream) { +#ifndef STRATA_W_NO_COMB + if (((reinterpret_cast(Dm) | reinterpret_cast(shared) | reinterpret_cast(bo)) & 15) == 0) { + moe_combine4_kernel<<>>( + Dm, slot, w, reinterpret_cast(shared), sg, reinterpret_cast(bo), T); + check("moe_combine"); + return; + } +#endif + moe_combine_kernel<<>>(Dm, slot, w, shared, sg, bo, T); + check("moe_combine"); +} } // namespace +void moe_combine(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, + int64_t T, void* stream, int64_t n_embd, int64_t k) { + if (n_embd == 2560 && k == 10) moe_combine_t<2560, 10>(Dm, slot, w, shared, sg, bo, T, stream); + else if (n_embd == 2048 && k == 8) moe_combine_t<2048, 8>(Dm, slot, w, shared, sg, bo, T, stream); + else { std::fprintf(stderr, "moe_combine: no kernel for n_embd %lld, %lld experts\n", (long long) n_embd, (long long) k); std::exit(1); } +} void moe_combine_peer16(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, const uint16_t* peer16, int64_t rows_local, float* bo, int64_t T, void* stream) { moe_combine_peer16_kernel<<>>(Dm, slot, w, shared, sg, peer16, @@ -2162,19 +2270,6 @@ void peer_scatter_add(float* sum, const float* rows, const float* wk, const int3 peer_scatter_add_kernel<<>>(sum, rows, wk, pair, n); check("peer_scatter_add"); } -void moe_combine(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, - int64_t T, void* stream) { -#ifndef STRATA_W_NO_COMB - if (((reinterpret_cast(Dm) | reinterpret_cast(shared) | reinterpret_cast(bo)) & 15) == 0) { - moe_combine4_kernel<<>>( - Dm, slot, w, reinterpret_cast(shared), sg, reinterpret_cast(bo), T); - check("moe_combine"); - return; - } -#endif - moe_combine_kernel<<>>(Dm, slot, w, shared, sg, bo, T); - check("moe_combine"); -} void rms_rows(float* x, const float* w, int64_t rows, int64_t cols, int64_t ld, float eps, void* stream) { if (rows <= 0) return; rms_rows_kernel<<<(unsigned) rows, 256, 0, (cudaStream_t) stream>>>(x, w, cols, ld, eps); @@ -2201,13 +2296,21 @@ void rope(float* x, int64_t T, int64_t heads, int64_t dim, int64_t ld, int64_t p strata::kernels::mrope_table(), rt); check("rope"); } -void split_q(const float* q_full, float* q, int64_t T, void* stream) { - split_q_kernel<<>>(q_full, q, T); +void split_q(const float* q_full, float* q, int64_t T, void* stream, int64_t n_head) { + if (n_head == 16) split_q_kernel<16><<>>(q_full, q, T); + else if (n_head == 24) split_q_kernel<24><<>>(q_full, q, T); + else { std::fprintf(stderr, "split_q: no kernel for %lld heads\n", (long long) n_head); std::exit(1); } check("split_q"); } -void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream, int64_t ld16) { - gate_attn_kernel<<>>(attn, q_full, out16, T, - ld16 > 0 ? ld16 : 24 * 256); +void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream, int64_t ld16, + int64_t n_head) { + if (n_head == 16) + gate_attn_kernel<16><<>>(attn, q_full, out16, T, + ld16 > 0 ? ld16 : 16 * 256); + else if (n_head == 24) + gate_attn_kernel<24><<>>(attn, q_full, out16, T, + ld16 > 0 ? ld16 : 24 * 256); + else { std::fprintf(stderr, "gate_attn: no kernel for %lld heads\n", (long long) n_head); std::exit(1); } check("gate_attn"); } diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 84164445298..24b3f0e6914 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -26,6 +26,7 @@ #include "strata/kernels/qsa.hpp" #include "strata/kernels/kv_stream.hpp" #include "strata/kernels/cvec.hpp" +#include "strata/kernels/elementwise.hpp" #include "strata/kernels/kv_q4.hpp" #include "strata/core/layer.hpp" #include "strata/core/native_head.hpp" @@ -999,10 +1000,12 @@ const MmqPlan& mmq_plan() { p.layer.assign((size_t) std::max(layers, 0), 0); p.fo.assign(p.layer.size(), 0); p.fallback = !on || layers <= 0; + // the model's expert widths (Flash-Next's 640 / N; Qwen3.6's 512 / 2048) + const int64_t ff = lay.native ? lay.n_ff : 640, ne = lay.native ? lay.n_embd : N; for (int64_t l = 0; on && l < layers; ++l) { const int gt = lay.native ? lay.fmt[(size_t) l].gu_type : 42, dt = lay.native ? lay.fmt[(size_t) l].d_type : 42; - // #420: a tile on every GPU for these shapes (gate+up: 1280 rows, down: N rows), else the FP16 path - if (!mmq::fits(gt, 1280) || !mmq::fits(dt, N)) { + // #420: a tile on every GPU for these shapes (gate+up: 2 x n_ff rows, down: n_embd rows), else the FP16 path + if (!mmq::fits(gt, 2 * ff) || !mmq::fits(dt, ne)) { p.fallback = true; // the fused path's buffers (Xq, H) exist for it; a chunk it does not take (a small one) keeps the FP16 path if (on && lay.native && fused::native_supported(gt, dt)) { p.fo[(size_t) l] = 1; p.any = true; } @@ -1010,8 +1013,8 @@ const MmqPlan& mmq_plan() { } p.layer[(size_t) l] = 1; p.any = true; - p.gu_max = std::max(p.gu_max, mmq::matrix_bytes(gt, 1280, N)); - p.d_max = std::max(p.d_max, mmq::matrix_bytes(dt, N, 640)); + p.gu_max = std::max(p.gu_max, mmq::matrix_bytes(gt, 2 * ff, ne)); + p.d_max = std::max(p.d_max, mmq::matrix_bytes(dt, ne, ff)); } return p; }(); @@ -1057,13 +1060,38 @@ uint64_t moe_set_bytes(size_t T, int64_t n_expert, bool fused) { } } +bool Prefill::supports(const core::ModelGeometry& g) { + if (g.arch == core::ModelArch::kQwen4Exp) + return g.n_embd == N && g.hc == HC && g.hc_lr == LR && g.n_expert >= 1 && g.n_expert_used == K; + // a model with one residual stream (qwen35moe): every buffer is sized by Flash-Next's numbers above, so each of + // its sizes must be at most those; the GDN kernels have instances for 48 and 32 value heads over 16 key heads. + // Without an indexer the prompt attention is the tensor-core kernel's dense mode: NVIDIA sm_75+ (elsewhere the + // prompt is read through the decode windows) +#if defined(STRATA_USE_HIP) + if (!g.has_indexer) return false; +#else + if (!g.has_indexer) { + int dev = 0, major = 0, minor = 0; + if (cudaGetDevice(&dev) != cudaSuccess || cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) != cudaSuccess || + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev) != cudaSuccess || major * 10 + minor < 75) { + cudaGetLastError(); + return false; + } + } +#endif + return g.hc == 1 && !g.has_ple && g.n_embd <= N && g.n_embd % 256 == 0 && g.n_expert >= 1 && g.n_expert <= NE && + g.n_expert_used >= 1 && g.n_expert_used <= K && g.n_ff <= 640 && g.ssm_k_heads == 16 && + (g.ssm_v_heads == 48 || g.ssm_v_heads == 32) && g.ssm_conv_channels <= C && g.ssm_value_dim <= ZV && + g.n_head * g.head_dim <= ZV && g.n_head_kv == 2 && g.head_dim == 256; +} + bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, core::SessionState& ss, core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t chunk, void* stream, std::string& err, void* borrow, uint64_t borrow_bytes) { Impl& m = *impl_; m.wt = &wt; m.g = &g; m.ss = &ss; m.src = src; m.cache = cache; m.host_res = host_res; m.T = chunk; m.cs = (cudaStream_t) stream; m.stats = &stats_; - if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_expert < 1 || ss.k != K) { + if (!supports(g) || ss.k != g.n_expert_used) { err = "prefill: geometry differs from the artifact's"; return false; } cudaGetDevice(&m.device); @@ -1230,7 +1258,9 @@ bool Prefill::carve(size_t T, void* alloc) { strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; - m.cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); + if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; // dense attention (qwen35moe) + // without an indexer the tensor-core attention reads every cell (no ids): the selection buffers stay minimal + m.cap = g.has_indexer ? strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s) : 64; m.max_blocks = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2; { // one region for the attention half's and the MoE half's scratch (see gdn_set_bytes) @@ -1348,6 +1378,8 @@ bool Prefill::draft_kv(core::MtpDrafter& mtp, const float* R_rows, const int32_t Impl& m = *impl_; const core::OnDevice on_device(m.device); // the stage's own GPU current (a layer split's drafter: the last stage) auto decline = [&](const char* reason) { if (why != nullptr) *why = reason; return false; }; + if (m.g == nullptr || !m.g->has_hc()) // Flash-Next's draft layer only: the caller's mtp.prefill reads it + return decline("not Flash-Next's draft layer (qwen35moe: the drafter's own pass)"); static const bool off = [] { const char* v = std::getenv("STRATA_MTP_BATCH"); return v != nullptr && v[0] == '0'; }(); // A ring (KV streaming: the drafter's window, page p in slot p % n_slots over a host copy) takes the same appends // with its own page table and host copy, as a streamed main layer does; the cells written are those the window can @@ -1429,6 +1461,7 @@ bool Prefill::draft_kv(core::MtpDrafter& mtp, const float* R_rows, const int32_t strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; + if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; // dense attention (qwen35moe) std::vector tk((size_t) B); if (q8) { if (!m.mmq_ctx) m.mmq_ctx = std::make_unique(); @@ -1866,7 +1899,8 @@ uint64_t Prefill::bytes_needed_impl(const core::ModelGeometry& g, const core::Se strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; - const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); + if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; // dense attention (qwen35moe) + const int64_t cap = g.has_indexer ? strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s) : 64; const int64_t max_blocks = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2; o.take((size_t) std::max({gdn_set_bytes(T), qsa_set_bytes(T, cap, max_blocks, 256, 32, s), moe_set_bytes(T, g.n_expert, fused_layout(T, src))}), ok); @@ -2057,6 +2091,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; s.idx_dim = g.idx_key_dim; + if (!g.has_indexer) s.idx_top_k = strata::kernels::kDenseTopK; // dense attention (qwen35moe) // read once (STRATA_KV_PREFETCH=1/0: the next streamed layer's KV prefix is staged on its own stream) static const bool kv_prefetch_on = [] { const char* e = std::getenv("STRATA_KV_PREFETCH"); @@ -2077,6 +2112,9 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri ~KvDrain() { if (stream) cudaStreamSynchronize(stream); } } kv_drain{kv_prefetch ? m.kv_copy : nullptr}; int64_t kv_prefetches = 0; + // the model's row widths: n_embd, and the residual's (hc streams of it). Flash-Next: N and D; one stream: n_embd + const int64_t NE_ = g.n_embd, RW = g.hc * g.n_embd; + const bool one_stream = !g.has_hc(); const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size + (uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1); int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; @@ -2147,7 +2185,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri // ---- embeddings, broadcast to the four streams - or, in a later stage of a layer split, the rows the // previous stage handed on if (hand_in_ != nullptr) { - if (cudaMemcpyAsync(m.R, hand_in_ + (size_t) c0 * D, (size_t) T * D * 4, cudaMemcpyHostToDevice, m.cs) != + if (cudaMemcpyAsync(m.R, hand_in_ + (size_t) c0 * RW, (size_t) T * RW * 4, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { err = "prefill: the layer split's hand-off upload failed"; return false; @@ -2190,7 +2228,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (row) { // #1996: the picture's rows live in the caller's (pageable) memory. A run of consecutive rows is // gathered into one pinned slot and goes up as one copy; the slot is guarded by its event. - const int64_t row_bytes = N * 4, run_max = std::max(1, (4 << 20) / row_bytes); + const int64_t row_bytes = NE_ * 4, run_max = std::max(1, (4 << 20) / row_bytes); int64_t t1 = t + 1; while (t1 < T && t1 - t < run_max && embd_rows[p0 + t1]) ++t1; void* stage = m.up_ring.acquire((size_t) ((t1 - t) * row_bytes)); @@ -2198,10 +2236,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (stage != nullptr) { for (int64_t r = t; r < t1; ++r) std::memcpy((char*) stage + (size_t) ((r - t) * row_bytes), embd_rows[p0 + r], (size_t) row_bytes); - e = m.up_ring.submit(m.emb + t * N, stage, (size_t) ((t1 - t) * row_bytes), m.cs); + e = m.up_ring.submit(m.emb + t * NE_, stage, (size_t) ((t1 - t) * row_bytes), m.cs); } else { // no pinned RAM: row by row from the caller's memory, each waited for for (int64_t r = t; r < t1 && e == cudaSuccess; ++r) - e = core::PinnedRing::direct(m.emb + r * N, embd_rows[p0 + r], (size_t) row_bytes, m.cs); + e = core::PinnedRing::direct(m.emb + r * NE_, embd_rows[p0 + r], (size_t) row_bytes, m.cs); } if (e != cudaSuccess) { err = "prefill: the image embedding upload failed"; @@ -2209,11 +2247,18 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri } t = t1; } else { - if (!core::embed_row(*m.wt, g, tokens[c0 + t], m.emb + t * N, m.cs, err)) return false; + if (!core::embed_row(*m.wt, g, tokens[c0 + t], m.emb + t * NE_, m.cs, err)) return false; ++t; } } - if (hand_in_ == nullptr) gr_broadcast(m.emb, m.R, T, m.cs); + if (hand_in_ == nullptr && one_stream) { // one residual stream: the embeddings are its rows + if (cudaMemcpyAsync(m.R, m.emb, (size_t) T * NE_ * 4, cudaMemcpyDeviceToDevice, m.cs) != cudaSuccess) { + err = "prefill: the residual's first rows"; + return false; + } + } else if (hand_in_ == nullptr) { + gr_broadcast(m.emb, m.R, T, m.cs); + } // ---- the PLE rows of the whole chunk, one batched SSD request on a thread (see ple_gather). A later chunk's // were read ahead during the previous chunk; the first chunk's are read beside layer 0 - they are needed from // layer 1 on, and gathering them here first left the GPU idle for the whole read (~0.4 s of a 32K prompt) @@ -2558,73 +2603,83 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri stats_.ms_ple += ms_since(tp); } for (int half = 0; half < 2; ++half) { - // ---- the hyper-connection read of this half - const char* pre = half == 0 ? "hc_attn_" : "hc_ffn_"; - const std::string sn = std::string(pre) + "norm.weight", sd = std::string(pre) + "down.weight", - su = std::string(pre) + "up.weight", si = std::string(pre) + "inject.weight"; - const core::WeightRef *wn = need(v, sn.c_str(), err), *wd = need(v, sd.c_str(), err), - *wu = need(v, su.c_str(), err), *wi = need(v, si.c_str(), err); - if (!wn || !wd || !wu || !wi) return false; - pt.mark(kPfHc, cs); - // the xn16 token stride of this chunk (the writers of both halves and of the fused write-backs use it) - const int64_t ldx = hc_pad() && T >= std::max(pf_switch_min_t(), 64) && !gr_unfused() && - !m.xn16_lo ? D + XN_PAD : D; - if (gr_unfused()) gr_norm(m.R, (const float*) wn->data, EPS, m.xn, m.xn16, T, m.cs, m.xn16_lo); - else if (!normed) gr_norm_rs(m.R, (const float*) wn->data, EPS, m.grs, m.xn16, T, m.cs, m.xn16_lo, ldx); - normed = false; - bool hcd = false; - bool hdown = false; - if (ldx != D && !pf_hcdown()) { // STRATA_HCD_EXACT: the exact-order down kernel when hipBLASLt would take 1176 / 1177 - hdown = wd->kind == core::WeightKind::Bf16InF32 && wd->data && wd->ne0 == D && wd->ne1 == LR && - m.gemm.bf16_hcd_exact(m.xn16, ldx, (const uint16_t*) wd->data, m.lo, T, LR, D); - } else if (ldx != D) { - if (wd->kind != core::WeightKind::Bf16InF32 || !wd->data || wd->ne0 != D || wd->ne1 != LR || - wi->kind != core::WeightKind::Bf16InF32 || !wi->data || wi->ne0 != D || wi->ne1 != HC || - !strata_pf_hcdown_bf16(m.xn16, ldx, (const uint16_t*) wd->data, (const uint16_t*) wi->data, LR, - HC, m.lo, m.inj, T, D, m.cs)) { - err = "prefill: STRATA_PF_HCDOWN could not run the hyper-connection projections"; - return false; + // ---- the hyper-connection read of this half (one stream: the RMSNorm before the mixer / the MoE; the + // previous half's output was added to R by its write) + int64_t ldx = D; // the xn16 token stride of this chunk (the writers of both halves use it) + if (one_stream) { + const core::WeightRef* wn = need(v, half == 0 ? "attn_norm.weight" : "post_attention_norm.weight", err); + if (!wn) return false; + pt.mark(kPfHc, cs); + rms_mix(m.R, (const float*) wn->data, EPS, NE_, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, m.mixed_bf_lo); + } else { + // ---- the hyper-connection read of this half + const char* pre = half == 0 ? "hc_attn_" : "hc_ffn_"; + const std::string sn = std::string(pre) + "norm.weight", sd = std::string(pre) + "down.weight", + su = std::string(pre) + "up.weight", si = std::string(pre) + "inject.weight"; + const core::WeightRef *wn = need(v, sn.c_str(), err), *wd = need(v, sd.c_str(), err), + *wu = need(v, su.c_str(), err), *wi = need(v, si.c_str(), err); + if (!wn || !wd || !wu || !wi) return false; + pt.mark(kPfHc, cs); + // the xn16 token stride of this chunk (the writers of both halves and of the fused write-backs use it) + ldx = hc_pad() && T >= std::max(pf_switch_min_t(), 64) && !gr_unfused() && + !m.xn16_lo ? D + XN_PAD : D; + if (gr_unfused()) gr_norm(m.R, (const float*) wn->data, EPS, m.xn, m.xn16, T, m.cs, m.xn16_lo); + else if (!normed) gr_norm_rs(m.R, (const float*) wn->data, EPS, m.grs, m.xn16, T, m.cs, m.xn16_lo, ldx); + normed = false; + bool hcd = false; + bool hdown = false; + if (ldx != D && !pf_hcdown()) { // STRATA_HCD_EXACT: the exact-order down kernel when hipBLASLt would take 1176 / 1177 + hdown = wd->kind == core::WeightKind::Bf16InF32 && wd->data && wd->ne0 == D && wd->ne1 == LR && + m.gemm.bf16_hcd_exact(m.xn16, ldx, (const uint16_t*) wd->data, m.lo, T, LR, D); + } else if (ldx != D) { + if (wd->kind != core::WeightKind::Bf16InF32 || !wd->data || wd->ne0 != D || wd->ne1 != LR || + wi->kind != core::WeightKind::Bf16InF32 || !wi->data || wi->ne0 != D || wi->ne1 != HC || + !strata_pf_hcdown_bf16(m.xn16, ldx, (const uint16_t*) wd->data, (const uint16_t*) wi->data, LR, + HC, m.lo, m.inj, T, D, m.cs)) { + err = "prefill: STRATA_PF_HCDOWN could not run the hyper-connection projections"; + return false; + } + hcd = true; } - hcd = true; - } - if (!hcd && !hdown && !bf16_proj(m.gemm, wd, m.xn16, m.lo, T, sd, err, 0, m.xn16_lo, ldx != D ? ldx : 0)) return false; - gr_silu(m.lo, m.lo16, T, m.cs, m.lo16_lo); - bool upmixed = false; - if (hc_upmix() && T >= pf_switch_min_t() && !gr_unfused() && !m.lo16_lo && !m.mixed_bf_lo && - wu->kind == core::WeightKind::Bf16InF32 && wu->data && wu->ne0 == LR && wu->ne1 == D) { - static int checks = [] { const char* e = std::getenv("STRATA_HC_UPMIX_CHECK"); return e ? std::atoi(e) : 0; }(); - if (checks > 0) { // the default pair first, kept for the comparison - --checks; - if (!bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err)) return false; - gr_mix_r(m.R, m.grs, (const float*) wn->data, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h); - std::vector ref((size_t) T * N), got((size_t) T * N); - cudaMemcpyAsync(ref.data(), m.mixed, ref.size() * 4, cudaMemcpyDeviceToHost, m.cs); - upmixed = gr_upmix(m.lo16, (const uint16_t*) wu->data, m.R, m.grs, (const float*) wn->data, - m.mixed, m.mixed_bf, m.mixed_h, T, m.cs); - cudaMemcpyAsync(got.data(), m.mixed, got.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaStreamSynchronize(m.cs); - double e2 = 0, r2 = 0, emax = 0; - for (size_t i = 0; i < ref.size(); ++i) { - const double dd = (double) got[i] - ref[i]; - e2 += dd * dd; r2 += (double) ref[i] * ref[i]; emax = std::max(emax, std::fabs(dd)); + if (!hcd && !hdown && !bf16_proj(m.gemm, wd, m.xn16, m.lo, T, sd, err, 0, m.xn16_lo, ldx != D ? ldx : 0)) return false; + gr_silu(m.lo, m.lo16, T, m.cs, m.lo16_lo); + bool upmixed = false; + if (hc_upmix() && T >= pf_switch_min_t() && !gr_unfused() && !m.lo16_lo && !m.mixed_bf_lo && + wu->kind == core::WeightKind::Bf16InF32 && wu->data && wu->ne0 == LR && wu->ne1 == D) { + static int checks = [] { const char* e = std::getenv("STRATA_HC_UPMIX_CHECK"); return e ? std::atoi(e) : 0; }(); + if (checks > 0) { // the default pair first, kept for the comparison + --checks; + if (!bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err)) return false; + gr_mix_r(m.R, m.grs, (const float*) wn->data, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h); + std::vector ref((size_t) T * N), got((size_t) T * N); + cudaMemcpyAsync(ref.data(), m.mixed, ref.size() * 4, cudaMemcpyDeviceToHost, m.cs); + upmixed = gr_upmix(m.lo16, (const uint16_t*) wu->data, m.R, m.grs, (const float*) wn->data, + m.mixed, m.mixed_bf, m.mixed_h, T, m.cs); + cudaMemcpyAsync(got.data(), m.mixed, got.size() * 4, cudaMemcpyDeviceToHost, m.cs); + cudaStreamSynchronize(m.cs); + double e2 = 0, r2 = 0, emax = 0; + for (size_t i = 0; i < ref.size(); ++i) { + const double dd = (double) got[i] - ref[i]; + e2 += dd * dd; r2 += (double) ref[i] * ref[i]; emax = std::max(emax, std::fabs(dd)); + } + std::fprintf(stderr, "strata: STRATA_HC_UPMIX_CHECK layer %lld half %d T %lld: mixed rel RMS %.3e, " + "max |diff| %.3e (%s)\n", (long long) l, half, (long long) T, + std::sqrt(e2 / std::max(r2, 1e-30)), emax, upmixed ? "upmix" : "upmix refused"); + } else { + upmixed = gr_upmix(m.lo16, (const uint16_t*) wu->data, m.R, m.grs, (const float*) wn->data, + m.mixed, m.mixed_bf, m.mixed_h, T, m.cs); } - std::fprintf(stderr, "strata: STRATA_HC_UPMIX_CHECK layer %lld half %d T %lld: mixed rel RMS %.3e, " - "max |diff| %.3e (%s)\n", (long long) l, half, (long long) T, - std::sqrt(e2 / std::max(r2, 1e-30)), emax, upmixed ? "upmix" : "upmix refused"); + } + if (!upmixed && !bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err, 0, m.lo16_lo)) return false; + if (!hcd && !bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err, 0, m.xn16_lo, ldx != D ? ldx : 0)) return false; + if (upmixed) { + } else if (gr_unfused()) { + gr_mix(m.xn, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, m.mixed_bf_lo); } else { - upmixed = gr_upmix(m.lo16, (const uint16_t*) wu->data, m.R, m.grs, (const float*) wn->data, - m.mixed, m.mixed_bf, m.mixed_h, T, m.cs); + gr_mix_r(m.R, m.grs, (const float*) wn->data, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, + m.mixed_bf_lo); } } - if (!upmixed && !bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err, 0, m.lo16_lo)) return false; - if (!hcd && !bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err, 0, m.xn16_lo, ldx != D ? ldx : 0)) return false; - if (upmixed) { - } else if (gr_unfused()) { - gr_mix(m.xn, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, m.mixed_bf_lo); - } else { - gr_mix_r(m.R, m.grs, (const float*) wn->data, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, - m.mixed_bf_lo); - } if (half == 0 && !core::is_qsa_layer(g, l)) { // ======================= GDN ======================= @@ -2639,35 +2694,39 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size; if (!native_proj(m.gemm, wqkv, m.mixed_h, m.qkv, T, v.name("attn_qkv.weight"), err)) return false; if (!native_proj(m.gemm, wg, m.mixed_h, m.z, T, v.name("attn_gate.weight"), err)) return false; - if (!bf16_proj(m.gemm, wa, m.mixed_bf, m.ab, T, v.name("ssm_alpha.weight"), err, 2 * HV, m.mixed_bf_lo)) return false; - if (!bf16_proj(m.gemm, wb, m.mixed_bf, m.ab + HV, T, v.name("ssm_beta.weight"), err, 2 * HV, m.mixed_bf_lo)) return false; + const int64_t hv = g.ssm_v_heads; // the model's value heads (Flash-Next: HV) + if (!bf16_proj(m.gemm, wa, m.mixed_bf, m.ab, T, v.name("ssm_alpha.weight"), err, 2 * hv, m.mixed_bf_lo)) return false; + if (!bf16_proj(m.gemm, wb, m.mixed_bf, m.ab + hv, T, v.name("ssm_beta.weight"), err, 2 * hv, m.mixed_bf_lo)) return false; pt.mark(kPfGdnConv, cs); // "gdn" is the projections in; the rest on their own lines - gdn_gates(m.ab, (const float*) wdt->data, (const float*) wsa->data, m.gate, m.beta, T, m.cs); - gdn_conv(conv, m.qkv, (const float*) wc->data, m.hbuf, T, EPS, m.cs); + gdn_gates(m.ab, (const float*) wdt->data, (const float*) wsa->data, m.gate, m.beta, T, m.cs, hv); + gdn_conv(conv, m.qkv, (const float*) wc->data, m.hbuf, T, EPS, m.cs, hv); pt.mark(kPfGdnRec, cs); - const int64_t ld_y = pf_pad() && T >= std::max(pf_switch_min_t(), 64) ? ZV + ZV_PAD : 0; + const int64_t ld_y = pf_pad() && T >= std::max(pf_switch_min_t(), 64) ? g.ssm_value_dim + ZV_PAD : 0; gdn_recurrence(state, m.hbuf, m.gate, m.beta, m.z, (const float*) wnm->data, EPS, m.y, m.y_h, T, m.cs, - ld_y); + ld_y, hv, g.gdn_gate_silu); pt.mark(kPfGdnOut, cs); if (!native_proj(m.gemm, wo, m.y_h, m.bo, T, v.name("ssm_out.weight"), err, 0, ld_y)) return false; ++gdn_index; } else if (half == 0) { // ======================= QSA ======================= const core::QsaState& st = ss.qsa_states[qsa_index]; + // qwen35moe has no indexer: every cell is attended (the attention's dense mode), nothing pooled + const bool idx = g.has_indexer; + const int64_t nh = g.n_head, qd = g.n_head * g.head_dim; // Flash-Next: 24 and ZV const core::WeightRef *wq = need(v, "attn_q.weight", err), *wk = need(v, "attn_k.weight", err), *wv = need(v, "attn_v.weight", err), *wo = need(v, "attn_output.weight", err), - *wik = need(v, "indexer.k_proj.weight", err), - *wiq = need(v, "indexer.q_proj.weight", err), + *wik = idx ? need(v, "indexer.k_proj.weight", err) : nullptr, + *wiq = idx ? need(v, "indexer.q_proj.weight", err) : nullptr, *wqn = need(v, "attn_q_norm.weight", err), *wkn = need(v, "attn_k_norm.weight", err), - *wiqn = need(v, "indexer.q_norm.weight", err), - *wikn = need(v, "indexer.k_norm.weight", err); - if (!wq || !wk || !wv || !wo || !wik || !wiq || !wqn || !wkn || !wiqn || !wikn) return false; + *wiqn = idx ? need(v, "indexer.q_norm.weight", err) : nullptr, + *wikn = idx ? need(v, "indexer.k_norm.weight", err) : nullptr; + if (!wq || !wk || !wv || !wo || !wqn || !wkn || (idx && (!wik || !wiq || !wiqn || !wikn))) return false; pt.mark(kPfQsa, cs); if (!native_proj(m.gemm, wk, m.mixed_h, m.Kc, T, v.name("attn_k.weight"), err)) return false; if (!native_proj(m.gemm, wv, m.mixed_h, m.Vc, T, v.name("attn_v.weight"), err)) return false; if (!native_proj(m.gemm, wq, m.mixed_h, m.Qf, T, v.name("attn_q.weight"), err)) return false; - if (!bf16_proj(m.gemm, wik, m.mixed_bf, m.idx_raw, T, v.name("indexer.k_proj.weight"), err, 0, m.mixed_bf_lo)) return false; - if (!bf16_proj(m.gemm, wiq, m.mixed_bf, m.q_idx, T, v.name("indexer.q_proj.weight"), err, 0, m.mixed_bf_lo)) return false; + if (idx && !bf16_proj(m.gemm, wik, m.mixed_bf, m.idx_raw, T, v.name("indexer.k_proj.weight"), err, 0, m.mixed_bf_lo)) return false; + if (idx && !bf16_proj(m.gemm, wiq, m.mixed_bf, m.q_idx, T, v.name("indexer.q_proj.weight"), err, 0, m.mixed_bf_lo)) return false; rms_rows(m.Kc, (const float*) wkn->data, T * 2, 256, 256, EPS, m.cs); rope(m.Kc, T, 2, 256, 512, p0, strata::kernels::rope_scaling(), m.cs); // KV streaming: this layer's cells [0, p0) come in from the host copy to the staging pool, and the @@ -2730,10 +2789,11 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri (p0 + T + s.page_size - 1) / s.page_size, s, m.cs); } if (staged) pf_step("reading the prompt (batched, step sync): the K/V append at layer", l); - split_q(m.Qf, m.q, T, m.cs); - rms_rows(m.q, (const float*) wqn->data, T * 24, 256, 256, EPS, m.cs); - rope(m.q, T, 24, 256, 6144, p0, strata::kernels::rope_scaling(), m.cs); - if (st.kv_rot) strata::kernels::fwht256_inplace_cuda(m.q, T * 24, m.cs); + split_q(m.Qf, m.q, T, m.cs, nh); + rms_rows(m.q, (const float*) wqn->data, T * nh, 256, 256, EPS, m.cs); + rope(m.q, T, nh, 256, qd, p0, strata::kernels::rope_scaling(), m.cs); + if (st.kv_rot) strata::kernels::fwht256_inplace_cuda(m.q, T * nh, m.cs); + if (idx) { // the indexer: its keys appended, then each query's selection (and the debug checks) rms_rows(m.q_idx, (const float*) wiqn->data, T * 4, 128, 128, EPS, m.cs); rope(m.q_idx, T, 4, 128, 512, p0, strata::kernels::rope_scaling(), m.cs); // the indexer appends, token by token; then scores + selection for many queries at once: @@ -2859,6 +2919,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri } } } + } // idx const strata::kernels::QsaAttnPools pools = staged ? pools_of(m.stage, m.ident_table) : core::qsa_attn_pools(st); pt.mark(kPfQsaAttn, cs); @@ -2867,23 +2928,32 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri // decode kernel, 32 queries at a time (K8V4 runs the tensor kernel's mode 3: INT8 K, // V dequantized from its q4_0 blocks to fp16 at gather) static const bool old_attn = std::getenv("STRATA_PROMPT_ATTN_OLD") != nullptr; - if (old_attn || !strata::kernels::qsa_prompt_attn_batch(m.q, pools, m.sel_ids, m.steps_dev, m.cap, s, - m.attn, T, m.cs)) + // (dense: no ids - the kernel reads every cell of each query's width) + if (old_attn || !strata::kernels::qsa_prompt_attn_batch(m.q, pools, idx ? m.sel_ids : nullptr, + m.steps_dev, m.cap, s, m.attn, T, m.cs)) { + if (!idx) { + err = "prefill: this model's dense prompt attention needs the tensor-core kernel (an NVIDIA " + "card from RTX 20 on, K/V int8 or fp16)"; + return false; + } for (int64_t t0 = 0; t0 < T; t0 += m.attn_batch) { const int64_t nb = std::min(m.attn_batch, T - t0); - strata::kernels::qsa_decode_attn_batch(m.q + t0 * ZV, pools, m.sel_ids + t0 * m.cap, + strata::kernels::qsa_decode_attn_batch(m.q + t0 * qd, pools, m.sel_ids + t0 * m.cap, m.steps_dev + t0 * strata::kernels::kStepCount, m.cap, - s, m.attn_scratch, m.attn + t0 * ZV, nb, m.cs); + s, m.attn_scratch, m.attn + t0 * qd, nb, m.cs); } - if (st.kv_rot || st.kv_hybrid) strata::kernels::fwht256_inplace_cuda(m.attn, T * 24, m.cs); + } + if (st.kv_rot || st.kv_hybrid) strata::kernels::fwht256_inplace_cuda(m.attn, T * nh, m.cs); pt.mark(kPfQsa, cs); - const int64_t ld_a = pf_pad() && T >= std::max(pf_switch_min_t(), 64) ? ZV + ZV_PAD : 0; - gate_attn(m.attn, m.Qf, m.attn_h, T, m.cs, ld_a); + const int64_t ld_a = pf_pad() && T >= std::max(pf_switch_min_t(), 64) ? qd + ZV_PAD : 0; + gate_attn(m.attn, m.Qf, m.attn_h, T, m.cs, ld_a, nh); if (!native_proj(m.gemm, wo, m.attn_h, m.bo, T, v.name("attn_output.weight"), err, 0, ld_a)) return false; ++qsa_index; if (!kv_prefetch_after(l + 1, qsa_index)) return false; } else { // ======================= MoE ======================= + // the model's expert shape: Flash-Next's K / N / 2*640 / 640, Qwen3.6's 8 / 2048 / 1024 / 512 + const int64_t KK = ss.k, NN = g.n_embd, FF_ = g.n_ff, GU_ = 2 * g.n_ff; const core::WeightRef *wr = need(v, "ffn_gate_inp.weight", err), *wgi = need(v, "ffn_gate_inp_shexp.weight", err), *wsg = need(v, "ffn_gate_shexp.weight", err), @@ -2892,14 +2962,14 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (!wr || !wgi || !wsg || !wsu || !wsd) return false; pt.mark(kPfRouter, cs); if (!bf16_proj(m.gemm, wr, m.mixed_bf, m.logits, T, v.name("ffn_gate_inp.weight"), err, 0, m.mixed_bf_lo)) return false; - route(m.logits, m.ids, m.w, T, m.g->n_expert, m.cs); + route(m.logits, m.ids, m.w, T, m.g->n_expert, m.cs, KK); // RESEARCH HOOK (141-research, opt-in, never on by default): STRATA_DBG_FEAT= dumps, per MoE layer and chunk, // the routed ids + weights (and for the layers in STRATA_DBG_FEAT_LAYERS the router-input hidden state, bf16); // STRATA_DBG_ROUTE_OVERRIDE= replaces the routing of the layers listed in that file with the file's ids + weights. { static const char* feat_path = std::getenv("STRATA_DBG_FEAT"); static const char* ovr_path = std::getenv("STRATA_DBG_ROUTE_OVERRIDE"); - if (feat_path != nullptr || ovr_path != nullptr) { + if ((feat_path != nullptr || ovr_path != nullptr) && !one_stream) { // (Flash-Next's shapes: 10 per token) static std::FILE* ff = feat_path ? std::fopen(feat_path, "wb") : nullptr; static std::vector featl = [] { std::vector r(64, false); @@ -2957,11 +3027,11 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri auto shared_expert = [&]() -> bool { if (!native_proj(m.gemm, wsg, m.mixed_h, m.sgate, T, v.name("ffn_gate_shexp.weight"), err)) return false; if (!native_proj(m.gemm, wsu, m.mixed_h, m.sup, T, v.name("ffn_up_shexp.weight"), err)) return false; - swiglu_pair(m.sgate, m.sup, m.sh_h, T, m.cs); + swiglu_pair(m.sgate, m.sup, m.sh_h, T, m.cs, FF_); if (!native_proj(m.gemm, wsd, m.sh_h, m.shared, T, v.name("ffn_down_shexp.weight"), err)) return false; if (wgi->kind != core::WeightKind::Bf16InF32) { err = "prefill: shared gate is not BF16"; return false; } - m.gemm.bf16(m.mixed_bf, (const uint16_t*) wgi->data, m.sg, T, 1, N); - if (m.mixed_bf_lo) m.gemm.bf16(m.mixed_bf_lo, (const uint16_t*) wgi->data, m.sg, T, 1, N, 0, 1.0f); + m.gemm.bf16(m.mixed_bf, (const uint16_t*) wgi->data, m.sg, T, 1, NN); + if (m.mixed_bf_lo) m.gemm.bf16(m.mixed_bf_lo, (const uint16_t*) wgi->data, m.sg, T, 1, NN, 0, 1.0f); return true; }; // Only the router needs to finish before readback. Queue the shared expert afterwards @@ -2983,7 +3053,9 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri // fused_ring() sized the ring and the buffers for it const bool no_peer = !core::peer_portable(); const bool fused_only = mmq_plan().any && mmq_plan().fo[(size_t) l]; - const bool fused_nat = (use_mmq || fused_only) && stream_all && no_peer && lay.native && fused::native_supported(mmq_gt, mmq_dt); + // (the fused kernels are compiled for Flash-Next's 2560 / 640: never for another model's experts) + const bool fused_nat = (use_mmq || fused_only) && stream_all && no_peer && lay.native && !one_stream && + fused::native_supported(mmq_gt, mmq_dt); const bool fused_l = (use_mmq && stream_all && no_peer && !lay.native && fused::enabled()) || fused_nat; // #583 / #954: the fused layout's GU/H/Xq hold the grouping tables and the int8 rows, sized for // stream_all_min() - 1 tokens of MMQ's rows, so a layer that takes MMQ or the FP16 path at the FULL @@ -3018,9 +3090,9 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri } pt.mark(kPfGather, cs); // the layer's input to int8 once per token; the rows of each expert from the router's ids - if (fused_nat) fused::quantize_act_native(m.mixed, T, N, m.Xq, m.cs); - else fused::quantize_act(m.mixed, T, N, m.Xq, m.cs); - fused::group(m.ids, T * K, (int) K, (int) m.g->n_expert, m.GU, m.slot_dev, m.src_dev, m.cs); + if (fused_nat) fused::quantize_act_native(m.mixed, T, NN, m.Xq, m.cs); + else fused::quantize_act(m.mixed, T, NN, m.Xq, m.cs); + fused::group(m.ids, T * KK, (int) KK, (int) m.g->n_expert, m.GU, m.slot_dev, m.src_dev, m.cs); // launches over the experts in id order, each at most kMaxBatch experts of which at most a // third of the ring streamed: the next batch's blobs arrive while one computes. An expert // the routing did not pick has no tiles; its ring slot is given back with its batch. @@ -3054,10 +3126,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (fused_nat) { const auto& f = lay.fmt[(size_t) l]; const fused::NativeGeom ng{f.gu_type, f.d_type, f.gu_row, f.d_row, f.up_off, f.down_off}; - fused::experts_native(b, ng, (int) m.g->n_expert, T * K, m.GU, m.Xq, m.src_dev, m.H, + fused::experts_native(b, ng, (int) m.g->n_expert, T * KK, m.GU, m.Xq, m.src_dev, m.H, m.Dm, m.cs); } else { - fused::experts(b, (int) m.g->n_expert, T * K, m.GU, m.Xq, m.src_dev, m.H, m.Dm, m.cs); + fused::experts(b, (int) m.g->n_expert, T * KK, m.GU, m.Xq, m.src_dev, m.H, m.Dm, m.cs); } for (size_t kk = k0; kk < k; ++kk) { cudaEventRecord(m.used[kk % (size_t) m.ring], m.cs); @@ -3077,13 +3149,14 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri int32_t* ids_h = grp_mapped ? m.grp_host : m.ids_host.data(); int32_t* slot_h = grp_mapped ? m.grp_host + m.grp_tk : m.slot_host.data(); int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data(); - if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs); - else if (m.up_ring.d2h_sync(m.ids_host.data(), m.ids, (size_t) T * K * 4, m.cs) != cudaSuccess) { + if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * KK, m.cs); + else if (m.up_ring.d2h_sync(m.ids_host.data(), m.ids, (size_t) T * KK * 4, m.cs) != cudaSuccess) { err = "prefill: the routed ids' download failed"; // #1996 return false; } + // (not for qwen35moe - one residual stream: the CPU share is unmeasured there) bool cpu_maybe = pool_hold.held && !stream_all && m.src != nullptr && !m.cpu_dead && - lay.native && !m.pp && !lay.fmt.empty(); + lay.native && !m.pp && !lay.fmt.empty() && !one_stream; // The share's two pinned buffers are small (~50 MB) but are had mid-request, when the host may have // nothing left (0.1.41 report: tight RAM, chats failing). A buffer that cannot be had ends the share // for this engine - the GPU takes every expert, as with STRATA_PREFILL_CPU_SHARE=0 - instead of @@ -3114,7 +3187,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaEventRecord(m.cpu_wall[2 * cpu_set], m.cs); } if (cpu_maybe && cpu_arm) { - const size_t want = (size_t) T * N; + const size_t want = (size_t) T * NN; if (m.cpu_x_n < want) { cpu_cold = true; if (m.cpu_x) cudaFreeHost(m.cpu_x); @@ -3196,7 +3269,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri pt.mark(kPfHostGroup, cs); } std::fill(m.cnt.begin(), m.cnt.end(), 0); - for (int64_t i = 0; i < T * K; ++i) { + for (int64_t i = 0; i < T * KK; ++i) { const int32_t e = ids_h[(size_t) i]; if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; } ++m.cnt[(size_t) e]; @@ -3235,10 +3308,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri ++n_cpu; } } - // multi-GPU: the rows of the experts the peer computes go last, as one block [rows_local, T*K) + // multi-GPU: the rows of the experts the peer computes go last, as one block [rows_local, T*KK) const bool pre_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l]; std::vector on_peer; - int64_t rows_local = T * K, rows_peer = 0; + int64_t rows_local = T * KK, rows_peer = 0; if (m.pp && pre_mmq) { on_peer.assign((size_t) m.g->n_expert, 0); for (int32_t e = 0; e < m.g->n_expert; ++e) { @@ -3255,7 +3328,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri on_peer[(size_t) e] = 1; rows_peer += c; } - rows_local = T * K - rows_peer; + rows_local = T * KK - rows_peer; } { int32_t r = 0; @@ -3274,11 +3347,11 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri rows_local -= rows_cpu; } std::vector fill(m.off.begin(), m.off.end() - 1); - for (int64_t i = 0; i < T * K; ++i) { + for (int64_t i = 0; i < T * KK; ++i) { const int32_t e = ids_h[(size_t) i]; const int32_t p = fill[(size_t) e]++; slot_h[(size_t) i] = p; - src_h[(size_t) p] = (int32_t) (i / K); + src_h[(size_t) p] = (int32_t) (i / KK); } // the experts, in id order: resident ones from VRAM, the others through the staging ring std::vector order, order_peer; @@ -3295,13 +3368,13 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (detail::mmq_resident_sort_eligible(detail::mmq_resident_sort_requested() && on_peer.empty(), use_mmq, lay.native, m.cache != nullptr, m.host_res ? m.host_res + (size_t) l * m.g->n_expert : nullptr, m.g->n_expert, stream_all, !stream_all || seq_start[(size_t) l] == seq_start[(size_t) l + 1])) - detail::mmq_resident_sort_rows(ids_h, T * K, (int32_t) K, m.cnt, m.off, order, slot_h, src_h); + detail::mmq_resident_sort_rows(ids_h, T * KK, (int32_t) KK, m.cnt, m.off, order, slot_h, src_h); if (grp_mapped) { - copy_i32(m.slot_dev, m.grp_dev + m.grp_tk, T * K, m.cs); - copy_i32(m.src_dev, m.grp_dev + 2 * m.grp_tk, T * K, m.cs); + copy_i32(m.slot_dev, m.grp_dev + m.grp_tk, T * KK, m.cs); + copy_i32(m.src_dev, m.grp_dev + 2 * m.grp_tk, T * KK, m.cs); } else { - if (m.up_ring.h2d(m.slot_dev, m.slot_host.data(), (size_t) T * K * 4, m.cs) != cudaSuccess || // #1996 - m.up_ring.h2d(m.src_dev, m.src_host.data(), (size_t) T * K * 4, m.cs) != cudaSuccess) { + if (m.up_ring.h2d(m.slot_dev, m.slot_host.data(), (size_t) T * KK * 4, m.cs) != cudaSuccess || // #1996 + m.up_ring.h2d(m.src_dev, m.src_host.data(), (size_t) T * KK * 4, m.cs) != cudaSuccess) { err = "prefill: the routing tables' upload failed"; return false; } @@ -3310,26 +3383,26 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri // set_cpu_pool: the CPU's experts on a thread while this one streams and runs the GPU's; their // unweighted rows land in Dm's row order, so the combine weights them as any other row if (rows_cpu > 0) { - const size_t want = (size_t) rows_cpu * N; + const size_t want = (size_t) rows_cpu * NN; if (m.cpu_rows_n < want) { cpu_cold = true; if (m.cpu_rows) cudaFreeHost(m.cpu_rows); m.cpu_rows = nullptr; m.cpu_rows_n = 0; const size_t cap_rows = (size_t) rows_cpu * 2; // headroom: few reallocations - if (cudaHostAlloc((void**) &m.cpu_rows, cap_rows * N * sizeof(float), cudaHostAllocDefault) != cudaSuccess) { + if (cudaHostAlloc((void**) &m.cpu_rows, cap_rows * NN * sizeof(float), cudaHostAllocDefault) != cudaSuccess) { err = "prefill: cannot allocate the CPU experts' rows"; return false; } - m.cpu_rows_n = cap_rows * N; + m.cpu_rows_n = cap_rows * NN; } - const int64_t r0c = T * K - rows_cpu; + const int64_t r0c = T * KK - rows_cpu; const strata::kernels::cpu::NativeFmt& cf = lay.fmt[(size_t) l]; if (cpu_share_env() < 0.0) { if (m.cpu_ev[0] == nullptr) { cudaEventCreate(&m.cpu_ev[0]); cudaEventCreate(&m.cpu_ev[1]); } cudaEventRecord(m.cpu_ev[0], m.cs); } - cpu_fut = std::async(std::launch::async, [&m, &on_cpu, &cpu_blob, &cf, &cpu_ms, src_h, l, T, r0c, this]() -> bool { + cpu_fut = std::async(std::launch::async, [&m, &on_cpu, &cpu_blob, &cf, &cpu_ms, src_h, l, T, r0c, NN, this]() -> bool { const auto t0 = std::chrono::steady_clock::now(); constexpr size_t AB = strata::kernels::cpu::kNativeActBytes; // a Q2_0 layer (gate/up on the pool's Q2_0 kernels): ActQ activations, as decode's @@ -3343,9 +3416,9 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri need[(size_t) src_h[(size_t) m.off[(size_t) e] + (size_t) r]] = 1; for (int64_t t = 0; t < T; ++t) if (need[(size_t) t] && q2) - strata::kernels::cpu::act_quant_any(m.cpu_x + (size_t) t * N, (int) N, m.cpu_actq[(size_t) t]); + strata::kernels::cpu::act_quant_any(m.cpu_x + (size_t) t * NN, (int) NN, m.cpu_actq[(size_t) t]); else if (need[(size_t) t]) - strata::kernels::cpu::native_quant_act(cf, m.cpu_x + (size_t) t * N, + strata::kernels::cpu::native_quant_act(cf, m.cpu_x + (size_t) t * NN, m.cpu_nact.data() + (size_t) t * AB); m.cpu_jobs.clear(); for (int32_t e = 0; e < m.g->n_expert; ++e) { @@ -3357,7 +3430,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri const int64_t p = (int64_t) m.off[(size_t) e] + r; if (q2) j.act[r] = &m.cpu_actq[(size_t) src_h[(size_t) p]]; else j.nact[r] = m.cpu_nact.data() + (size_t) src_h[(size_t) p] * AB; - j.out[r] = m.cpu_rows + (size_t) (p - r0c) * N; + j.out[r] = m.cpu_rows + (size_t) (p - r0c) * NN; } m.cpu_jobs.push_back(j); } @@ -3366,12 +3439,12 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri return true; }); } - const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, 1280, N) : 0; - const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, N, 640) : 0; + const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, GU_, NN) : 0; + const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, NN, FF_) : 0; pt.mark(kPfGather, cs); if (use_mmq) { // step 2b: the layer's activations as q8_1 rows in expert order, straight from `mixed` - mmq::quantize(m.mixed, m.src_dev, m.Xq, mmq_gt, N, N, T * K, m.cs); + mmq::quantize(m.mixed, m.src_dev, m.Xq, mmq_gt, NN, NN, T * KK, m.cs); // each group's rows: absolute bounds (gate/up reads the layer's rows), relative ones (down // reads the group's own quantized H) const size_t n = order.size(), ng = (n + MMQ_GROUP - 1) / MMQ_GROUP; @@ -3393,7 +3466,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri } } } else { - gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs); + gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * KK, NN, m.cs); } // multi-GPU: the peer's share, enqueued before the primary's own experts so both cards work at once peer_now = use_mmq && !order_peer.empty(); @@ -3407,11 +3480,11 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaStreamWaitEvent(P.x_stream, P.ev_x, 0); xs = P.x_stream; } - if (!P.p2p && P.f16) cudaMemcpyAsync(P.host_x16, m.mixed_h, (size_t) T * N * 2, cudaMemcpyDeviceToHost, xs); - else if (!P.p2p) cudaMemcpyAsync(P.host_x, m.mixed, (size_t) T * N * 4, cudaMemcpyDeviceToHost, xs); + if (!P.p2p && P.f16) cudaMemcpyAsync(P.host_x16, m.mixed_h, (size_t) T * NN * 2, cudaMemcpyDeviceToHost, xs); + else if (!P.p2p) cudaMemcpyAsync(P.host_x, m.mixed, (size_t) T * NN * 4, cudaMemcpyDeviceToHost, xs); if (P.sums) { // the routing weights, and each peer row's routed pair (t * K + k) - cudaMemcpyAsync(P.host_w, m.w, (size_t) T * K * 4, cudaMemcpyDeviceToHost, xs); - for (int64_t i = 0; i < T * K; ++i) + cudaMemcpyAsync(P.host_w, m.w, (size_t) T * KK * 4, cudaMemcpyDeviceToHost, xs); + for (int64_t i = 0; i < T * KK; ++i) if (const int32_t r = slot_h[(size_t) i]; r >= rows_local) P.host_pair[r - rows_local] = (int32_t) i; } cudaEventRecord(P.ev_in, xs); @@ -3422,9 +3495,9 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaStreamWaitEvent(ps, P.ev_in, 0); if (P.out_pending) { cudaStreamWaitEvent(ps, P.ev_done, 0); P.out_pending = false; } pe.mark(kPeMoeIn, ps); - if (P.p2p) cudaMemcpyPeerAsync(P.mixed, P.dev, m.mixed, prevd, (size_t) T * N * 4, ps); - else if (P.f16) f16_to_f32_wide(P.mixed, P.host_x16, T * N, ps); - else copy_f32_wide(P.mixed, P.host_x, T * N, ps); + if (P.p2p) cudaMemcpyPeerAsync(P.mixed, P.dev, m.mixed, prevd, (size_t) T * NN * 4, ps); + else if (P.f16) f16_to_f32_wide(P.mixed, P.host_x16, T * NN, ps); + else copy_f32_wide(P.mixed, P.host_x, T * NN, ps); P.back_at = rows_local; P.back_rows = rows_peer; if (P.p2p) { @@ -3556,21 +3629,21 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri const int32_t* rel = P.bounds + n + 1 + g2 * (MMQ_GROUP + 1); cudaMemsetAsync(P.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, ps); cudaMemsetAsync(P.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, ps); - mmq::quantize(P.mixed, P.src + r0, P.Xq_g, mmq_gt, N, N, nr, ps); + mmq::quantize(P.mixed, P.src + r0, P.Xq_g, mmq_gt, NN, NN, nr, ps); mmq::Product gu; - gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; + gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = GU_; gu.w_cols = NN; gu.expert_bytes = mmq_gub; gu.n = ngx; gu.xq = P.Xq_g; gu.bounds = rel; gu.ids = P.ident; - gu.total_rows = nr; gu.max_rows = maxr; gu.dst = P.GU_g; gu.ld_dst = 1280; + gu.total_rows = nr; gu.max_rows = maxr; gu.dst = P.GU_g; gu.ld_dst = GU_; P.run_ctx->run(gu, ps); - mmq::swiglu(P.GU_g, P.H_g, nr, 640, !lay.native, ps); - mmq::quantize(P.H_g, nullptr, P.Hq_g, mmq_dt, 640, 640, nr, ps); + mmq::swiglu(P.GU_g, P.H_g, nr, FF_, !lay.native, ps); + mmq::quantize(P.H_g, nullptr, P.Hq_g, mmq_dt, FF_, FF_, nr, ps); const int b = (int) (g2 & 1); if (P.dm_live[b]) cudaStreamWaitEvent(ps, P.ev_dm[b], 0); // its last rows have left mmq::Product dn; - dn.w = P.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; + dn.w = P.grp_d; dn.type = mmq_dt; dn.w_rows = NN; dn.w_cols = FF_; dn.expert_bytes = mmq_db; dn.n = ngx; dn.xq = P.Hq_g; dn.bounds = rel; dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm_b[b]; - dn.ld_dst = N; + dn.ld_dst = NN; P.run_ctx->run(dn, ps); if (P.sums && P.gather) { // into the per-token sums, one launch for the group const int32_t* a = P.adds + P.adds_at[g2]; @@ -3587,10 +3660,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaEventRecord(P.ev_grp[g2 % PeerPrefill::kGrpEv], ps); cudaStreamWaitEvent(P.s_out, P.ev_grp[g2 % PeerPrefill::kGrpEv], 0); if (P.p2p) - cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * N, prevd, P.Dm_b[b], P.dev, - (size_t) nr * N * 4, P.s_out); + cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * NN, prevd, P.Dm_b[b], P.dev, + (size_t) nr * NN * 4, P.s_out); else // the primary reads them when it combines (below) - cudaMemcpyAsync(P.host_rows + (size_t) r0 * N, P.Dm_b[b], (size_t) nr * N * 4, + cudaMemcpyAsync(P.host_rows + (size_t) r0 * NN, P.Dm_b[b], (size_t) nr * NN * 4, cudaMemcpyDeviceToHost, P.s_out); cudaEventRecord(P.ev_dm[b], P.s_out); P.dm_live[b] = true; @@ -3614,7 +3687,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri std::memcpy(P.host_bounds, P.bounds_host.data(), P.bounds_host.size() * 4); copy_i32(P.bounds, P.host_bounds, (int64_t) P.bounds_host.size(), ps); } - mmq::quantize(P.mixed, P.src, P.Xq, mmq_gt, N, N, rows_peer, ps); + mmq::quantize(P.mixed, P.src, P.Xq, mmq_gt, NN, NN, rows_peer, ps); pe.mark(kPeMoeGemm, ps); const auto& f = lay.fmt[(size_t) l]; for (size_t j = 0; j < n; ++j) { @@ -3635,23 +3708,23 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaMemsetAsync(P.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, ps); cudaMemsetAsync(P.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, ps); mmq::Product gu; - gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; + gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = GU_; gu.w_cols = NN; gu.expert_bytes = mmq_gub; gu.n = ngx; gu.xq = P.Xq; gu.bounds = P.bounds + j0; gu.ids = P.ident; - gu.total_rows = rows_peer; gu.max_rows = maxr; gu.dst = P.GU; gu.ld_dst = 1280; + gu.total_rows = rows_peer; gu.max_rows = maxr; gu.dst = P.GU; gu.ld_dst = GU_; P.run_ctx->run(gu, ps); - mmq::swiglu(P.GU + r0 * 1280, P.H + r0 * 640, nr, 640, !lay.native, ps); - mmq::quantize(P.H + r0 * 640, nullptr, P.Hq, mmq_dt, 640, 640, nr, ps); + mmq::swiglu(P.GU + r0 * GU_, P.H + r0 * FF_, nr, FF_, !lay.native, ps); + mmq::quantize(P.H + r0 * FF_, nullptr, P.Hq, mmq_dt, FF_, FF_, nr, ps); mmq::Product dn; - dn.w = P.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; + dn.w = P.grp_d; dn.type = mmq_dt; dn.w_rows = NN; dn.w_cols = FF_; dn.expert_bytes = mmq_db; dn.n = ngx; dn.xq = P.Hq; dn.bounds = P.bounds + n + 1 + g2 * (MMQ_GROUP + 1); - dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm + r0 * N; - dn.ld_dst = N; + dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm + r0 * NN; + dn.ld_dst = NN; P.run_ctx->run(dn, ps); if (P.out_pipe && nr > 0) { // this group's rows go back while the next group computes cudaEventRecord(P.ev_grp[g2], ps); cudaStreamWaitEvent(P.s_out, P.ev_grp[g2], 0); - cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * N, prevd, P.Dm + (size_t) r0 * N, - P.dev, (size_t) nr * N * 4, P.s_out); + cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * NN, prevd, P.Dm + (size_t) r0 * NN, + P.dev, (size_t) nr * NN * 4, P.s_out); } } } @@ -3669,7 +3742,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaEventRecord(P.ev_done, P.s_out); P.out_pending = true; } else { - cudaMemcpyPeerAsync(m.Dm + (size_t) rows_local * N, prevd, P.Dm, P.dev, (size_t) rows_peer * N * 4, ps); + cudaMemcpyPeerAsync(m.Dm + (size_t) rows_local * NN, prevd, P.Dm, P.dev, (size_t) rows_peer * NN * 4, ps); cudaEventRecord(P.ev_done, ps); } pe.mark(kPeIdle, ps); @@ -3812,18 +3885,18 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri cudaMemsetAsync(m.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, m.cs); cudaMemsetAsync(m.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, m.cs); mmq::Product gu; - gu.w = m.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; + gu.w = m.grp_gu; gu.type = mmq_gt; gu.w_rows = GU_; gu.w_cols = NN; gu.expert_bytes = mmq_gub; gu.n = ngx; gu.xq = m.Xq; gu.bounds = m.bounds_dev + j0; gu.ids = m.ids_identity; - gu.total_rows = T * K; gu.max_rows = maxr; gu.dst = m.GU; gu.ld_dst = 1280; + gu.total_rows = T * KK; gu.max_rows = maxr; gu.dst = m.GU; gu.ld_dst = GU_; m.mmq_ctx->run(gu, m.cs); - mmq::swiglu(m.GU + r0 * 1280, m.H + r0 * 640, nr, 640, !lay.native, m.cs); + mmq::swiglu(m.GU + r0 * GU_, m.H + r0 * FF_, nr, FF_, !lay.native, m.cs); pt.mark(kPfGemmD, cs); - mmq::quantize(m.H + r0 * 640, nullptr, m.Hq, mmq_dt, 640, 640, nr, m.cs); + mmq::quantize(m.H + r0 * FF_, nullptr, m.Hq, mmq_dt, FF_, FF_, nr, m.cs); mmq::Product dn; - dn.w = m.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; + dn.w = m.grp_d; dn.type = mmq_dt; dn.w_rows = NN; dn.w_cols = FF_; dn.expert_bytes = mmq_db; dn.n = ngx; dn.xq = m.Hq; dn.bounds = m.bounds_dev + n + 1 + g * (MMQ_GROUP + 1); - dn.ids = m.ids_identity; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = m.Dm + r0 * N; - dn.ld_dst = N; + dn.ids = m.ids_identity; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = m.Dm + r0 * NN; + dn.ld_dst = NN; m.mmq_ctx->run(dn, m.cs); return true; } @@ -3840,10 +3913,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (slot >= 0) { cudaEventRecord(m.used[slot], m.cs); m.used_of[slot] = slot; } const int64_t o0 = m.off[(size_t) e], ne = m.cnt[(size_t) e]; pt.mark(kPfGemmGU, cs); - m.gemm.f16(m.Xs + o0 * N, m.dq_gu[q], m.GU + o0 * 1280, ne, 1280, N); - swiglu_interleaved(m.GU + o0 * 1280, m.Hh + o0 * 640, ne, m.cs); + m.gemm.f16(m.Xs + o0 * NN, m.dq_gu[q], m.GU + o0 * GU_, ne, GU_, NN); + swiglu_interleaved(m.GU + o0 * GU_, m.Hh + o0 * FF_, ne, m.cs, FF_); pt.mark(kPfGemmD, cs); - m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640); + m.gemm.f16(m.Hh + o0 * FF_, m.dq_d[q], m.Dm + o0 * NN, ne, NN, FF_); return true; }; if (!stream_all) { @@ -3922,8 +3995,8 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri m.pend_n_gpu = n_stream - n_cpu; } stats_.cpu_share = cpu_share_env() >= 0.0 ? cpu_share_env() : m.cpu_share_now; - const int64_t r0c = T * K - rows_cpu; - cudaMemcpyAsync(m.Dm + (size_t) r0c * N, m.cpu_rows, (size_t) rows_cpu * N * sizeof(float), + const int64_t r0c = T * KK - rows_cpu; + cudaMemcpyAsync(m.Dm + (size_t) r0c * NN, m.cpu_rows, (size_t) rows_cpu * NN * sizeof(float), cudaMemcpyHostToDevice, m.cs); stats_.experts_cpu += n_cpu; } @@ -3938,14 +4011,14 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (peer_now) { // multi-GPU: the peer's rows are in Dm (or, without P2P, in host memory) cudaStreamWaitEvent(m.cs, m.pp->ev_done, 0); if (!m.pp->p2p && !m.pp->sums) - copy_f32_wide(m.Dm + (size_t) m.pp->back_at * N, m.pp->host_rows, m.pp->back_rows * N, m.cs); + copy_f32_wide(m.Dm + (size_t) m.pp->back_at * NN, m.pp->host_rows, m.pp->back_rows * NN, m.cs); } if (peer_now && m.pp->sums && m.pp->f16) moe_combine_peer16(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.pp->host_sum16, m.pp->back_at, m.bo, T, m.cs); else if (peer_now && m.pp->sums) moe_combine_peer(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.pp->host_sum, m.pp->back_at, m.bo, T, m.cs); else - moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs); + moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs, NN, KK); // debug: STRATA_DBG_NAN=1 reports the first layer of a chunk whose MoE produced non-finite values if (static const bool dbg = std::getenv("STRATA_DBG_NAN") != nullptr; dbg) { cudaStreamSynchronize(m.cs); @@ -3958,10 +4031,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri }; // (fused: GU and H hold the grouping tables and int8 rows, not floats) // (the CPU's rows, the last rows_cpu of GU and H, are never written by the GPU: not checked) - const int64_t gpu_rows = T * K - rows_cpu; - const int64_t bgu = fused_l ? 0 : bad(m.GU, gpu_rows * 1280), bdm = bad(m.Dm, T * K * N), - bbo = bad(m.bo, T * N); - const int64_t bh = m.H && !fused_l ? bad(m.H, gpu_rows * 640) : -1; + const int64_t gpu_rows = T * KK - rows_cpu; + const int64_t bgu = fused_l ? 0 : bad(m.GU, gpu_rows * GU_), bdm = bad(m.Dm, T * KK * NN), + bbo = bad(m.bo, T * NN); + const int64_t bh = m.H && !fused_l ? bad(m.H, gpu_rows * FF_) : -1; static int64_t reported = -1; if ((bgu || bdm || bbo || bh > 0) && reported != stats_.chunks) { reported = stats_.chunks; @@ -3975,66 +4048,72 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri // ---- the hyper-connection write of this half; F-2: fused with the next half's norm when nothing else // touches R in between (not the stage's last half, not before the PLE block of layer 1, not under a // control vector) - const int64_t nl = half == 0 ? l : l + 1; - const bool steered = half == 1 && strata::kernels::cvec().covers(l); - const float *cv_dir = nullptr, *cv_s = nullptr; - const int* cv_on = nullptr; - const bool cv_fused = steered && cvec_fuse() && !gr_unfused() && nl < LE && !(nl == 1 && ple_on) && - strata::kernels::cvec().n_embd == N && strata::kernels::cvec().hc == HC && - strata::kernels::cvec_tables(&cv_dir, &cv_s, &cv_on); - const bool fuse = !gr_unfused() && nl < LE && !(half == 1 && nl == 1 && ple_on) && (!steered || cv_fused); - const core::WeightRef* wnn = nullptr; - if (fuse) { - const core::LayerView vn(*m.wt, nl); - wnn = need(vn, half == 0 ? "hc_ffn_norm.weight" : "hc_attn_norm.weight", err); - if (!wnn) return false; - } - if (wnn && cv_fused) { - // STRATA_CVEC_FUSE_CHECK=N: on the first N fused writes, also run the three-kernel path on a copy - // of R and report whether R, the row scales and the BF16 image are bit-identical - static int checks = [] { const char* e = std::getenv("STRATA_CVEC_FUSE_CHECK"); return e ? std::atoi(e) : 0; }(); - if (checks > 0) { - --checks; - float *Rc = nullptr, *rc = nullptr; - uint16_t* xc = nullptr; - cudaMalloc(&Rc, (size_t) T * D * 4); cudaMalloc(&rc, (size_t) T * HC * 4); - cudaMalloc(&xc, (size_t) T * D * 2); - cudaMemcpyAsync(Rc, m.R, (size_t) T * D * 4, cudaMemcpyDeviceToDevice, m.cs); - gr_write(Rc, m.bo, m.inj, HC, T, m.cs); - strata::kernels::cvec_apply(Rc, l, T, D, nullptr, 0, nullptr, 0, false, m.cs); - gr_norm_rs(Rc, (const float*) wnn->data, EPS, rc, xc, T, m.cs, nullptr, D); - gr_write_cvec_norm_rs(m.R, m.bo, m.inj, HC, cv_dir + l * N, cv_s + l, cv_on, - strata::kernels::cvec().mode, (const float*) wnn->data, EPS, m.grs, - m.xn16, T, m.cs, m.xn16_lo, ldx); - std::vector a((size_t) T * D), b((size_t) T * D), ra((size_t) T * HC), rb((size_t) T * HC); - std::vector xa((size_t) T * D), xb((size_t) T * D); - cudaMemcpyAsync(a.data(), m.R, a.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaMemcpyAsync(b.data(), Rc, b.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaMemcpyAsync(ra.data(), m.grs, ra.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaMemcpyAsync(rb.data(), rc, rb.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaMemcpyAsync(xa.data(), m.xn16, xa.size() * 2, cudaMemcpyDeviceToHost, m.cs); - cudaMemcpyAsync(xb.data(), xc, xb.size() * 2, cudaMemcpyDeviceToHost, m.cs); - cudaStreamSynchronize(m.cs); - cudaFree(Rc); cudaFree(rc); cudaFree(xc); - std::fprintf(stderr, "strata: STRATA_CVEC_FUSE_CHECK layer %lld T %lld: R %s, rs %s, xn16 %s\n", - (long long) l, (long long) T, - std::memcmp(a.data(), b.data(), a.size() * 4) ? "DIFFERS" : "identical", - std::memcmp(ra.data(), rb.data(), ra.size() * 4) ? "DIFFERS" : "identical", - std::memcmp(xa.data(), xb.data(), xa.size() * 2) ? "DIFFERS" : "identical"); - } else - gr_write_cvec_norm_rs(m.R, m.bo, m.inj, HC, cv_dir + l * N, cv_s + l, cv_on, - strata::kernels::cvec().mode, (const float*) wnn->data, EPS, m.grs, m.xn16, - T, m.cs, m.xn16_lo, ldx); - normed = true; - } else if (wnn) { - gr_write_norm_rs(m.R, m.bo, m.inj, HC, (const float*) wnn->data, EPS, m.grs, m.xn16, T, m.cs, - m.xn16_lo, ldx); - normed = true; + if (one_stream) { // one stream: the block output into the residual; the next read norms it + strata::kernels::add_inplace(m.R, m.bo, T * NE_, m.cs); + if (half == 1 && strata::kernels::cvec().covers(l)) // --control-vector-scaled + strata::kernels::cvec_apply(m.R, l, T, RW, nullptr, 0, nullptr, 0, false, m.cs); } else { - gr_write(m.R, m.bo, m.inj, HC, T, m.cs); + const int64_t nl = half == 0 ? l : l + 1; + const bool steered = half == 1 && strata::kernels::cvec().covers(l); + const float *cv_dir = nullptr, *cv_s = nullptr; + const int* cv_on = nullptr; + const bool cv_fused = steered && cvec_fuse() && !gr_unfused() && nl < LE && !(nl == 1 && ple_on) && + strata::kernels::cvec().n_embd == N && strata::kernels::cvec().hc == HC && + strata::kernels::cvec_tables(&cv_dir, &cv_s, &cv_on); + const bool fuse = !gr_unfused() && nl < LE && !(half == 1 && nl == 1 && ple_on) && (!steered || cv_fused); + const core::WeightRef* wnn = nullptr; + if (fuse) { + const core::LayerView vn(*m.wt, nl); + wnn = need(vn, half == 0 ? "hc_ffn_norm.weight" : "hc_attn_norm.weight", err); + if (!wnn) return false; + } + if (wnn && cv_fused) { + // STRATA_CVEC_FUSE_CHECK=N: on the first N fused writes, also run the three-kernel path on a copy + // of R and report whether R, the row scales and the BF16 image are bit-identical + static int checks = [] { const char* e = std::getenv("STRATA_CVEC_FUSE_CHECK"); return e ? std::atoi(e) : 0; }(); + if (checks > 0) { + --checks; + float *Rc = nullptr, *rc = nullptr; + uint16_t* xc = nullptr; + cudaMalloc(&Rc, (size_t) T * D * 4); cudaMalloc(&rc, (size_t) T * HC * 4); + cudaMalloc(&xc, (size_t) T * D * 2); + cudaMemcpyAsync(Rc, m.R, (size_t) T * D * 4, cudaMemcpyDeviceToDevice, m.cs); + gr_write(Rc, m.bo, m.inj, HC, T, m.cs); + strata::kernels::cvec_apply(Rc, l, T, D, nullptr, 0, nullptr, 0, false, m.cs); + gr_norm_rs(Rc, (const float*) wnn->data, EPS, rc, xc, T, m.cs, nullptr, D); + gr_write_cvec_norm_rs(m.R, m.bo, m.inj, HC, cv_dir + l * N, cv_s + l, cv_on, + strata::kernels::cvec().mode, (const float*) wnn->data, EPS, m.grs, + m.xn16, T, m.cs, m.xn16_lo, ldx); + std::vector a((size_t) T * D), b((size_t) T * D), ra((size_t) T * HC), rb((size_t) T * HC); + std::vector xa((size_t) T * D), xb((size_t) T * D); + cudaMemcpyAsync(a.data(), m.R, a.size() * 4, cudaMemcpyDeviceToHost, m.cs); + cudaMemcpyAsync(b.data(), Rc, b.size() * 4, cudaMemcpyDeviceToHost, m.cs); + cudaMemcpyAsync(ra.data(), m.grs, ra.size() * 4, cudaMemcpyDeviceToHost, m.cs); + cudaMemcpyAsync(rb.data(), rc, rb.size() * 4, cudaMemcpyDeviceToHost, m.cs); + cudaMemcpyAsync(xa.data(), m.xn16, xa.size() * 2, cudaMemcpyDeviceToHost, m.cs); + cudaMemcpyAsync(xb.data(), xc, xb.size() * 2, cudaMemcpyDeviceToHost, m.cs); + cudaStreamSynchronize(m.cs); + cudaFree(Rc); cudaFree(rc); cudaFree(xc); + std::fprintf(stderr, "strata: STRATA_CVEC_FUSE_CHECK layer %lld T %lld: R %s, rs %s, xn16 %s\n", + (long long) l, (long long) T, + std::memcmp(a.data(), b.data(), a.size() * 4) ? "DIFFERS" : "identical", + std::memcmp(ra.data(), rb.data(), ra.size() * 4) ? "DIFFERS" : "identical", + std::memcmp(xa.data(), xb.data(), xa.size() * 2) ? "DIFFERS" : "identical"); + } else + gr_write_cvec_norm_rs(m.R, m.bo, m.inj, HC, cv_dir + l * N, cv_s + l, cv_on, + strata::kernels::cvec().mode, (const float*) wnn->data, EPS, m.grs, m.xn16, + T, m.cs, m.xn16_lo, ldx); + normed = true; + } else if (wnn) { + gr_write_norm_rs(m.R, m.bo, m.inj, HC, (const float*) wnn->data, EPS, m.grs, m.xn16, T, m.cs, + m.xn16_lo, ldx); + normed = true; + } else { + gr_write(m.R, m.bo, m.inj, HC, T, m.cs); + } + if (steered && !cv_fused) // --control-vector-scaled + strata::kernels::cvec_apply(m.R, l, T, D, nullptr, 0, nullptr, 0, false, m.cs); } - if (steered && !cv_fused) // --control-vector-scaled - strata::kernels::cvec_apply(m.R, l, T, D, nullptr, 0, nullptr, 0, false, m.cs); } } if (!ple_land()) return false; // a stage that ends before layer 1: the rows land anyway, the next gather starts @@ -4112,7 +4191,16 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri } const auto toc2 = Clock::now(); if (on_stage_chunk && !on_stage_chunk(p0 + T, err)) return false; - if (on_chunk && !on_chunk(m.R, T, p0, err)) return false; + // one stream: the drafter reads the rows after output_norm (llama.cpp's h_nextn), as the verify window hands on + const float* rows = m.R; + if (on_chunk && one_stream && LE == g.n_layers) { + const core::WeightRef* on = m.wt->find("output_norm.weight"); + if (on == nullptr) { err = "prefill: output_norm.weight is missing"; return false; } + rms_mix(m.R, (const float*) on->data, EPS, NE_, m.mixed, nullptr, T, m.cs); + if (cudaStreamSynchronize(m.cs) != cudaSuccess) { err = "prefill: the drafter's rows"; return false; } + rows = m.mixed; + } + if (on_chunk && !on_chunk(rows, T, p0, err)) return false; host_sync_ms += std::chrono::duration(toc2 - toc).count(); host_chunk_ms += ms_since(toc2); } diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 1e5299e522f..e86dfff5b93 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -58,6 +58,7 @@ #include "strata/kernels/mrope.hpp" #include "strata/kernels/kv_q4.hpp" #include "strata/kernels/qsa.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" #include "strata/core/native_head.hpp" #include "strata/core/verify.hpp" #include "strata/core/mtp.hpp" @@ -1552,7 +1553,10 @@ bool load_control_vectors(const Options& o, const strata::core::ModelGeometry& g } } const int first = o.cvec_first <= 0 ? 1 : o.cvec_first; - const int last = (o.cvec_last <= 0 || o.cvec_last >= L) ? (int) L - 1 : o.cvec_last; + // one stream (qwen35moe): the last layer's output is added by the head's read in the decode path, after the point + // the vector would act on, so the steered range stops one layer earlier there (the prompt path agrees) + const int64_t top = g.has_hc() ? L - 1 : L - 2; + const int last = (o.cvec_last <= 0 || o.cvec_last > top) ? (int) top : o.cvec_last; const int single = o.cvec_mode == 0 ? o.cvec_single : -1; if (single >= 0 && (single >= L || !have[(size_t) single])) { err = "--cvec-dir single:" + std::to_string(single) + ": the vector has no direction for that layer"; @@ -2515,14 +2519,27 @@ int main(int argc, char** argv) { o.spec = std::max(o.spec, std::min(o.mtp_max_t + o.lookup_chain, 8)); // kVerifyMaxT } strata::core::layer_set_shared_early(!o.shared_late); + // The model's geometry, from its metadata shard's general.architecture: Flash-Next (qwen4exp) keeps the + // compiled-in numbers, Qwen3.6-35B-A3B (qwen35moe) brings its own. Without --native (a canonical Q2_0 pack) + // the model is Flash-Next. + strata::core::ModelGeometry model_geom; if (!o.native_preset.empty()) { try { // every shard of the model (-0000N-of-0000M.gguf beside --native). A missing shard is an error // here: it used to be skipped, leaving a model with some tensors absent and a later error, or none. o.native_shards = strata::gguf_split_paths(o.native_preset); + std::string geom_err; + if (!strata::core::geometry_from_gguf(strata::GgufFile(o.native_shards.front()), model_geom, geom_err)) + throw std::runtime_error(geom_err); + if (model_geom.arch != strata::core::ModelArch::kQwen4Exp) + std::fprintf(stderr, "strata generate: model architecture %s: %lld layers, n_embd %lld, %lld experts " + "(%lld per token), no PLE / hyper-connections / indexer\n", + strata::core::arch_name(model_geom.arch), (long long) model_geom.n_layers, + (long long) model_geom.n_embd, (long long) model_geom.n_expert, + (long long) model_geom.n_expert_used); // --ple-gguf defaults to the shard that holds the PLE table, found by name: shard 2 of the ISTA files // and of Unsloth's UD-Q4_K_XL, shard 1 of Swift's - if (o.ple_gguf.empty() && !o.no_ple) { + if (o.ple_gguf.empty() && !o.no_ple && model_geom.has_ple) { const strata::GgufModel model(o.native_shards); size_t at = 0; if (model.find("per_layer_token_embd.weight", &at) != nullptr) o.ple_gguf = o.native_shards[at]; @@ -2531,18 +2548,36 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: --native %s: %s\n", o.native_preset.c_str(), e.what()); return 2; } - if (o.no_ple || o.ple_gguf.empty()) { + if (model_geom.has_ple && (o.no_ple || o.ple_gguf.empty())) { std::fprintf(stderr, "strata generate: --native requires --ple-gguf (the PLE key is native too), and no " "shard of the model holds per_layer_token_embd.weight\n"); return 2; } + if (!model_geom.has_ple && !o.ple_gguf.empty()) { + std::fprintf(stderr, "strata generate: --ple-gguf: a %s model has no PLE table\n", + strata::core::arch_name(model_geom.arch)); + return 2; + } + if (!model_geom.has_ple) o.no_ple = true; // the model has no PLE block: every PLE path stays off + // Qwen3.6 / Ornith (one residual stream): one GPU. The layer split, the peer tier and the remote expert caches + // carry Flash-Next's shapes (10 experts per token, 2560 wide) and have not been run with this model. + if (!model_geom.has_hc() && (o.peer_device >= 1 || !o.layer_split.empty() || o.expert_cache_remote[0] > 0)) { + std::fprintf(stderr, "strata generate: a %s model runs on one GPU: %s is not supported for it\n", + strata::core::arch_name(model_geom.arch), + o.peer_device >= 1 ? "--peer-device" : !o.layer_split.empty() ? "--layer-split" + : "--expert-cache-device1..3"); + return 2; + } o.stream_token = true; o.gr_native_mmvf = true; o.native_bf16 = o.native_bf16_extra = true; - o.native_ple_key = o.native_moe_combine = o.native_gdn = o.native_router = true; - o.native_qsa = o.native_qsa_indexer = o.native_rope = o.native_ple_postops = true; + o.native_moe_combine = o.native_gdn = o.native_router = true; + o.native_qsa = o.native_qsa_indexer = o.native_rope = true; // (a model without an indexer never calls it) + if (model_geom.has_ple) o.native_ple_key = o.native_ple_postops = true; if (o.native_head_gguf.empty()) o.native_head_gguf = o.native_preset; - if (o.native_dense_gguf.empty()) { + if (o.native_dense_gguf.empty() && o.ple_gguf.empty()) { + o.native_dense_gguf = o.native_shards; + } else if (o.native_dense_gguf.empty()) { // every shard of the model (-0000N-of-0000M.gguf beside --native), then the PLE shard: a split // may put any layer in any shard (Swift's GGUFs: layers 13-47 in shard 2, the PLE table in shard 1) o.native_dense_gguf = o.native_shards; @@ -2729,8 +2764,8 @@ int main(int argc, char** argv) { // every quantized tensor in its GGUF form, so it needs --native (the dense projections, head and embedding // come from the model file) and runs its experts in verify windows only (--spec). { - const strata::core::ModelGeometry g0; - if (!strata::kernels::cpu::expert_layout_load(o.pack, g0.n_layers, g0.n_expert, err)) { + const strata::core::ModelGeometry g0 = model_geom; + if (!strata::kernels::cpu::expert_layout_load(o.pack, g0.n_layers, g0.n_expert, err, g0.n_embd, g0.n_ff)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -2750,6 +2785,17 @@ int main(int argc, char** argv) { } } const bool native_pack = strata::kernels::cpu::expert_layout().native; + // a model the batched prompt path does not run on this GPU (qwen35moe below sm_75, or on HIP): every prompt is + // read through the verify windows, in serve (every part a short read) as in one-shot generate + const bool no_prompt_path = native_pack && !model_geom.has_hc() && !strata::prefill::Prefill::supports(model_geom); + if (no_prompt_path) { + std::fprintf(stderr, "strata generate: this GPU has no batched prompt path for a %s model (NVIDIA sm_75+): " + "prompts are read through the decode windows (slower)\n", + strata::core::arch_name(model_geom.arch)); + o.prefill_chunk = 0; + o.prefill_auto = false; + o.short_read = std::numeric_limits::max(); + } // STRATA_EARLY_REMOTE_CONTEXTS=1: create EVERY secondary context here, like CUDA1's. Under WSL2 the driver's // pinned/mapped host budget (dxg gpadl, ~1 GiB) is spent by CUDA0's weights and MTP before the later loop runs, // and a new context then fails with cudaErrorMemoryAllocation (CUDA2: "cudaSetDevice(2) failed: out of memory"). @@ -2809,6 +2855,7 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: PCIe probe failed -> pcie_frac default %.2f\n", base); } } + mem_mark("the CUDA context (the PCIe probe)"); // the canonical Q2_0 pack's CPU kernels are AVX-512 only; a native pack runs on AVX2 CPUs as well if (!native_pack) strata::kernels::cpu::cpu_require_expert_support(); else if (!strata::kernels::cpu::cpu_avx512_ok()) @@ -2819,8 +2866,9 @@ int main(int argc, char** argv) { ? "AVX1 128-bit (iq_avx1, the older-CPU build)" : "ggml-cpu vec_dot (no AVX2: the older-CPU build)") : std::getenv("STRATA_NO_IQ256") == nullptr ? "AVX-2" : "ggml-cpu vec_dot (STRATA_NO_IQ256 set)"); - strata::core::ModelGeometry g; // canonical defaults; the model file overrides the MoE shape below - int64_t K = 10; + strata::core::ModelGeometry g = model_geom; // the model file's geometry (canonical defaults without --native) + int64_t K = g.n_expert_used; + strata::kernels::gdn_set_out_gate_silu(g.gdn_gate_silu); // qwen35moe: silu(z) on the GDN output // THE ROPE CONFIG RESOLVES HERE, BEFORE ANY WEIGHT MOVES - the CLI and the model file have both spoken, // and `session_init` below builds the rope table from it and captures the kernels reading its constants // (rope_scaling.hpp); the only hard constraint is "set before that", and dying on a bad rope key beats @@ -2829,20 +2877,18 @@ int main(int argc, char** argv) { // model file decides; an explicit value - `none` and `1` included, the opt-outs - wins over the model file. { // The model file's rope keys (llama.cpp's names under the arch prefix), when it carries any - the - // artifact today ships none, so this is a no-op defaults channel for future fine-tunes. + // Flash-Next artifact today ships none, so for it this is a no-op defaults channel for future fine-tunes; + // qwen35moe files carry rope.freq_base. (The MoE shape is in `model_geom`, from geometry_from_gguf.) std::string gguf_rope_type; double gguf_rope_base = 0, gguf_rope_factor = 0, gguf_rope_orig_ctx = 0; if (!o.native_preset.empty()) { - // a pruned variant (GSQ-RCO Coder) ships fewer experts than the canonical 512x10; the model file - // is the authority on its own MoE shape - everything else in the geometry is unchanged try { strata::GgufFile model_gguf(o.native_shards.front()); // the metadata shard - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.expert_count")) g.n_expert = (int64_t) v->u; - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.expert_used_count")) K = (int64_t) v->u; - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.rope.freq_base")) gguf_rope_base = v->num(); - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.rope.scaling.type")) gguf_rope_type = v->s; - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.rope.scaling.factor")) gguf_rope_factor = v->num(); - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.rope.scaling.original_context_length")) + const std::string ap = std::string(strata::core::arch_name(g.arch)) + "."; + if (const strata::MetaValue* v = model_gguf.get(ap + "rope.freq_base")) gguf_rope_base = v->num(); + if (const strata::MetaValue* v = model_gguf.get(ap + "rope.scaling.type")) gguf_rope_type = v->s; + if (const strata::MetaValue* v = model_gguf.get(ap + "rope.scaling.factor")) gguf_rope_factor = v->num(); + if (const strata::MetaValue* v = model_gguf.get(ap + "rope.scaling.original_context_length")) gguf_rope_orig_ctx = v->num(); } catch (const std::exception& e) { std::fprintf(stderr, "strata generate: reading the model's expert shape from %s: %s\n", @@ -2925,12 +2971,12 @@ int main(int argc, char** argv) { strata::core::NativeEmbed native_embed; if (native_pack) { if (o.native_preset.empty() || o.spec < 2 || o.keep_canonical || - (o.prefill_chunk <= 0 && o.tokens.size() > 1)) { + (o.prefill_chunk <= 0 && o.tokens.size() > 1 && !no_prompt_path)) { std::fprintf(stderr, "strata generate: %s is a native (IQ) pack: it needs --native SHARD1, --spec T (T >= 2) " "and --prefill CHUNK\n", o.pack.c_str()); return 2; } - const strata::core::ModelGeometry g0; + const strata::core::ModelGeometry g0 = model_geom; const auto embed_t0 = std::chrono::steady_clock::now(); if (!native_embed.load(o.embd_gguf.empty() ? o.native_shards : std::vector{o.embd_gguf}, g0.n_embd, 248320, err)) { @@ -3065,12 +3111,15 @@ int main(int argc, char** argv) { strata::core::NativeDense native_dense; if (!o.native_dense_gguf.empty()) { - if (!native_dense.load(o.native_dense_gguf, wt, err, o.native_ple_key)) { + // a model file that carries MTP layers after its main ones (qwen35moe): those belong to the draft + const int64_t dense_hi = g.n_mtp_layers > 0 ? g.n_layers : -1; + if (!native_dense.load(o.native_dense_gguf, wt, err, o.native_ple_key, 0, dense_hi)) { std::fprintf(stderr, "strata generate: native dense projections: %s\n", err.c_str()); return 1; } std::fprintf(stderr, "strata generate: %zu native projection matrices, %.2f MiB of weights, in %.1f s\n", native_dense.tensor_count(), (double) native_dense.weight_bytes() / (1024.0 * 1024.0), load_s()); + mem_mark("the native projections"); } strata::kernels::gr_set_fp32_activations(o.gr_fp32_activations); @@ -3328,7 +3377,7 @@ int main(int argc, char** argv) { } ss.ple.emb_dev = ple_emb_dev; ss.ple.scratch = ple_scratch; - } else { + } else if (g.has_ple) { // (a model without a PLE block - qwen35moe - has nothing to switch off) std::fprintf(stderr, "strata generate: PLE OFF by explicit --no-ple diagnostic request.\n" " The tokens below are NOT this model's; this is only useful for A/B measurement.\n"); @@ -4282,6 +4331,10 @@ int main(int argc, char** argv) { strata::core::qsa_set_kv_elastic(on, iv != nullptr && std::atoll(iv) > 0 ? std::atoll(iv) : 16384); strata::core::ExpertCache::set_vmm(on); } + mem_mark("the weights (before the session)"); + if (std::getenv("STRATA_TRACE")) + std::fprintf(stderr, "strata trace: the session takes %lld MiB\n", + (long long) (strata::core::session_bytes(g, o.max_context, K, 0, hi0) >> 20)); if (cudaMalloc(&sbuf, strata::core::session_bytes(g, o.max_context, K, 0, hi0)) != cudaSuccess) { std::fprintf(stderr, "strata generate: session state allocation failed\n"); return 1; @@ -4304,6 +4357,7 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: zeroing the session state failed\n"); return 1; } + mem_mark("the session"); if (!o.ple_gguf.empty()) { if (!ss.ple.ready()) { std::fprintf(stderr, "strata generate: the PLE run is not ready after construction\n"); @@ -4482,7 +4536,9 @@ int main(int argc, char** argv) { strata::core::MtpDrafter mtp; // the draft layer's geometry (the canonical model's MTP head), `static` because MtpDrafter keeps a reference; the batch // slots' draft-KV copies use this one too - static const strata::core::ModelGeometry draft_geometry{}; + // Flash-Next: the canonical geometry (below); a qwen35moe model's MTP block is its own, of the model's shape + static strata::core::ModelGeometry draft_geometry; + draft_geometry = g.arch == strata::core::ModelArch::kQwen4Exp ? strata::core::ModelGeometry{} : g; std::vector> slot_mtp; if (!o.mtp.empty()) { if (o.spec < 2) { @@ -4903,7 +4959,12 @@ int main(int argc, char** argv) { if (stages.empty()) for (const auto& d : slot_mtp) mtp_bind += (int64_t) d->bind_bytes(native_head.row_bytes(), n_vocab); - const int64_t reserve = (((int64_t) o.vram_reserve_mib + prefill_mib) << 20) + mtp_bind + pipe_first; + // a model without the indexer (qwen35moe) attends the whole context: the verifier's selection and attention + // partials grow with it and are allocated after the cache (443 MiB at 256K: without this a 256K context + // failed on every card, "verify: the device arena does not fit"). 0 for Flash-Next. + const int64_t dense_attn = (int64_t) strata::core::Verifier::dense_attention_bytes( + g, ss, batch_mtp ? strata::kernels::kVerifyMaxT : std::max(o.spec, o.batch)); + const int64_t reserve = (((int64_t) o.vram_reserve_mib + prefill_mib) << 20) + mtp_bind + pipe_first + dense_attn; const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob; int64_t slots = ((int64_t) free_b - reserve) / blob; if (!profile.empty()) slots = std::min(slots, (int64_t) profile.size()); @@ -4911,6 +4972,10 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: expert cache auto: %.2f GiB free, %d MiB reserved (+%lld MiB for the " "draft head) -> %d slots\n", (double) free_b / 1073741824.0, o.vram_reserve_mib, (long long) (mtp_bind >> 20), o.expert_cache); + if (dense_attn > 0) + std::fprintf(stderr, "strata generate: expert cache auto: +%lld MiB kept for the dense attention at %lld " + "tokens of context%s\n", (long long) (dense_attn >> 20), (long long) o.max_context, + strata::kernels::qsa_decode_attn_splits() > 0 ? " (split-K partials, STRATA_ATTN_SPLITS)" : ""); #if defined(STRATA_USE_HIP) && !defined(_WIN32) // An APU's "VRAM" is system RAM: the device-free figure above counts the whole GPU-addressable pool and does not // subtract ordinary CPU allocations (the host expert arena), so a cache sized from it alone can ask the OOM killer @@ -4941,15 +5006,20 @@ int main(int argc, char** argv) { // ends LOW (`reserve_adapted`). A reserve given on the command line is kept. No slot at all: the start // stops, saying what is short and what makes room. A card the default reserve leaves that much is sized as // before. - constexpr int kSmallReserveMib = 300; + // Flash-Next: 300 MiB (#496, IQ3_XXS at 32K ended with 5 MiB left). A model without the indexer allocates + // more after the cache besides its dense attention (counted above): its prompt path's own part, the + // verifier and the windows' graphs. Measured on Qwen3.6 UD-IQ3_S (RTX 4070 Ti, int8 K/V, MTP, a given + // reserve): 400 MiB failed in the graph instantiation at 8K, 32K and 64K; 450 started at all three and + // ended with 136-204 MiB free. 500 keeps 50 MiB over that. + const int kSmallReserveMib = g.has_indexer ? 300 : 500; const int64_t min_slots = (o.prefill_chunk > 0 && pf_borrow) ? ((int64_t) strata::prefill::Prefill::bytes_needed(g, ss, 256) + blob - 1) / blob + 128 : 1; if (o.expert_cache < min_slots && !o.vram_reserve_given && o.vram_reserve_mib > kSmallReserveMib) { // the largest reserve (in MiB) that still leaves min_slots - const int64_t fit_mib = ((int64_t) free_b - mtp_bind - pipe_first - min_slots * blob) / (1 << 20) - prefill_mib; + const int64_t fit_mib = ((int64_t) free_b - mtp_bind - pipe_first - dense_attn - min_slots * blob) / (1 << 20) - prefill_mib; if (fit_mib >= kSmallReserveMib) { const int r = (int) std::min(fit_mib, o.vram_reserve_mib); - int64_t s2 = ((int64_t) free_b - ((((int64_t) r + prefill_mib) << 20) + mtp_bind + pipe_first)) / blob; + int64_t s2 = ((int64_t) free_b - ((((int64_t) r + prefill_mib) << 20) + mtp_bind + pipe_first + dense_attn)) / blob; if (!profile.empty()) s2 = std::min(s2, (int64_t) profile.size()); std::fprintf(stderr, "strata generate: expert cache auto: the %d MiB reserve leaves too few slots on " "this card (a working cache needs %lld): a %d MiB reserve instead -> %lld slots\n", @@ -4959,11 +5029,14 @@ int main(int argc, char** argv) { reserve_adapted = true; } } - if (o.expert_cache == 0) { + // a model without the indexer whose cache stays under min_slots cannot lend the prompt path its buffers, + // and they do not fit beside it either (measured: "prefill: device buffers ... do not fit" mid-start) + const bool too_small = o.expert_cache == 0 || (!g.has_indexer && o.expert_cache < min_slots); + if (too_small) { // what is short, and what makes room: the numbers a small card picks from const int64_t at_reserve = o.vram_reserve_given ? o.vram_reserve_mib : std::min(o.vram_reserve_mib, kSmallReserveMib); - const int64_t need_b = (((int64_t) at_reserve + prefill_mib) << 20) + mtp_bind + pipe_first + min_slots * blob; + const int64_t need_b = (((int64_t) at_reserve + prefill_mib) << 20) + mtp_bind + pipe_first + dense_attn + min_slots * blob; const int64_t short_mib = std::max(1, (need_b - (int64_t) free_b + (1 << 20) - 1) >> 20); const int64_t session_mib = (int64_t) (strata::core::session_bytes(g, o.max_context, K, 0, g.n_layers) >> 20); @@ -4980,6 +5053,9 @@ int main(int argc, char** argv) { (long long) min_slots, (long long) ((min_slots * blob) >> 20), (long long) short_mib, (long long) session_mib, (long long) o.max_context, (long long) (mtp_bind >> 20), reserve_tip.c_str()); + // the start stops here, as #496 meant it to: on one GPU the verify windows cannot run without a cache + // (#174), and going on only failed later with a message about the prompt path's buffers + if (native_pack && !multi_gpu) return 1; } } else if (multi_gpu && o.expert_cache > 0) { // an explicit cache size leaves room for the prompt path's buffers and the reserve, or the first prompt @@ -6926,7 +7002,7 @@ int main(int argc, char** argv) { return false; }; if (o.serve) { - if (o.spec < 2 || o.prefill_chunk <= 0 || + if (o.spec < 2 || (o.prefill_chunk <= 0 && !no_prompt_path) || (graph_hits && (thits.d_res == nullptr || host_res.empty()))) { std::fprintf(stderr, "strata serve: needs --spec T and --prefill CHUNK (and a fillable " "--expert-cache; the graphed hit path additionally needs --expert-profile P)\n"); @@ -7033,7 +7109,7 @@ int main(int argc, char** argv) { // a cache too small to lend the prompt path its buffers would make it allocate them on top - on a card // whose cache already filled its reserve, that is the over-subscription the auto sizing avoids - so the // chunk is the largest one EVERY participant can lend (a smaller chunk only reads slower) - if (pf_borrow && d_res != nullptr) { + if (pf_borrow && d_res != nullptr && !no_prompt_path) { pf_parts.push_back({&xcache, &ss, &sp, -1, 0, multi_gpu ? split_at[0] : g.n_layers, -1, -1, 0, {}}); for (auto& st : stages) pf_parts.push_back({&st->cache, &st->ss, &st->sp, st->dev, st->lb, st->le, -1, -1, 0, {}}); @@ -7253,7 +7329,7 @@ int main(int argc, char** argv) { "their experts are streamed from the model file during every chunk\n", (long long) (borrowed - kept), (long long) borrowed); } - } else { + } else if (!no_prompt_path) { std::fprintf(stderr, "strata serve: the prompt path allocates its own buffers (too few cache slots to borrow)\n"); } rss_probe("the prompt path set up"); @@ -7288,7 +7364,7 @@ int main(int argc, char** argv) { return 0; }; static constexpr int64_t kStepChunks[] = {6144, 4096, 3072, 2048, 1536, 1024, 512}; - { + if (!no_prompt_path) { // First by arithmetic: a prompt path without a loan allocates its buffers, so the chunk must leave // headroom on that device (a chunk that fits to the last MiB left hipBLAS nothing: its GEMMs then // failed to launch on gfx1201 and the prompt hung). `bytes_needed` is the same count `init` makes. @@ -7327,7 +7403,7 @@ int main(int argc, char** argv) { } } mem_mark("before the prompt path init"); - for (;;) { + while (!no_prompt_path) { // no prompt path: every read goes through the windows (windows_ok) const int r = init_prompt_paths(); if (r == 0) break; int64_t next = 0; @@ -7821,7 +7897,8 @@ int main(int argc, char** argv) { in.embedding = o.embd_gguf; in.ple = o.ple_gguf; in.pack = o.pack; - in.mtp = o.mtp.empty() ? std::string() : o.mtp + "/dense.bin"; + // (a qwen35moe draft layer is the model file's own MTP block: --mtp names that GGUF) + in.mtp = o.mtp.empty() ? std::string() : g.has_hc() ? o.mtp + "/dense.bin" : o.mtp; if (srcp != nullptr) in.experts = srcp->model_inputs(); std::vector files = strata::core::session_model_inputs(in); // the pack's other files the loader reads (weights.cpp, expert_layout.cpp) @@ -7831,7 +7908,7 @@ int main(int argc, char** argv) { files.push_back({std::string("pack ") + n, o.pack + "/" + n, true}); if (srcp == nullptr) files.push_back({"pack experts.bin", o.pack + "/experts.bin", true}); } - if (!o.mtp.empty()) { + if (!o.mtp.empty() && g.has_hc()) { for (const char* n : {"dense.txt", "experts.bin"}) files.push_back({std::string("mtp ") + n, o.mtp + "/" + n}); files.push_back({"mtp draft_vocab.bin", o.mtp + "/draft_vocab.bin", true}); } @@ -8300,10 +8377,13 @@ int main(int argc, char** argv) { const double max_cached = mcs != nullptr && mcs[0] != 0 ? std::atof(mcs) / 100.0 : max_cached_default; const bool nearly_all_cached = !host_res.empty() && !all_experts_resident && (double) cached_experts >= max_cached * (double) host_res.size(); + // (Flash-Next only: measured on its 10-of-512 routing; a qwen35moe model routes 8 of 256 and keeps 0.1.41's + // answers until it is measured - STRATA_ROUTE_TAIL_SKIP=7 turns it on there) const bool tail_default_ok = tse == nullptr && !all_experts_resident && !nearly_all_cached && o.batch == 0 && decode_defaults_arch_ok() && - o.peer_device < 0; - if (tse == nullptr && nearly_all_cached && o.batch == 0 && o.peer_device < 0 && decode_defaults_arch_ok()) + o.peer_device < 0 && g.has_indexer; + if (tse == nullptr && nearly_all_cached && o.batch == 0 && o.peer_device < 0 && decode_defaults_arch_ok() && + g.has_indexer) std::fprintf(stderr, "strata serve: STRATA_ROUTE_TAIL_SKIP is off by default: %.1f%% of the experts are in the GPU cache, so " "few are missed and skipping them costs more draft acceptance than it saves; " "STRATA_ROUTE_TAIL_SKIP=7 turns it on\n", 100.0 * (double) cached_experts / (double) host_res.size()); @@ -8369,6 +8449,9 @@ int main(int argc, char** argv) { std::getenv("STRATA_ADAPT_OVERLAP_RESIDENT")[0] == '1') ? "not in the resident RAM mode unless STRATA_ADAPT_OVERLAP_RESIDENT=1" : drive.d.usage.empty() ? "the adaptive tier is off" + // qwen35moe: not the default (Ornith IQ4_XS, RTX 4070 Ti, 6 replies x 2 starts: 128 vs 144 + // tok/s decode with it); STRATA_ADAPT_OVERLAP=1 turns it on there + : !g.has_indexer && std::getenv("STRATA_ADAPT_OVERLAP") == nullptr ? "not the default for qwen35moe" : nullptr; if (aov_want) { if (aov_why == nullptr) { @@ -10646,6 +10729,7 @@ int main(int argc, char** argv) { static const bool no_short = std::getenv("STRATA_CKPT_REREAD") != nullptr; int64_t req_short_read = o.short_read; // STRATA_PIPELINE_SWITCH may set it per request (short_read=) auto windows_ok = [&](int64_t a, int64_t b) -> bool { + if (no_prompt_path) return true; if (no_short || b - a > req_short_read) return false; if (sp.embd_rows != nullptr) for (int64_t i = a; i < b; ++i) @@ -12508,7 +12592,7 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata serve: KV streaming: %.2f%% of %llu block reads hit VRAM, %.1f MiB read " "from RAM%s\n", look ? 100.0 * (double) (look - miss) / (double) look : 100.0, (unsigned long long) look, (double) miss * 4224.0 / 1048576.0, - over ? " - OVERFLOW (too few resident cells)" : ""); + over ? " (more than the resident cells hold: the rest read from RAM in place)" : ""); } if (sfx_windows > 0) std::fprintf(stderr, "strata serve: suffix drafts: %lld windows, %lld of %lld drafts accepted\n", @@ -12565,7 +12649,14 @@ int main(int argc, char** argv) { std::vector final_r_host(final_r ? (size_t) (g.hc * g.n_embd) : 0); std::vector> lent; // (residency index, slot) lent to the prompt path const int64_t n_batched = (o.prefill_until > 0 && o.prefill_until < n_prompt - 1) ? o.prefill_until : n_prompt - 1; - if (o.prefill_chunk > 0 && n_prompt > 1) { + // a model the batched prompt path does not run yet (Qwen3.6): every prompt token but the last is read through + // the verify windows below, as the serve's short reads do; the loop then starts at the last one + const bool prompt_windows = native_pack && !strata::prefill::Prefill::supports(g); + if (prompt_windows && n_prompt > 1) { + pos_start = n_prompt - 1; + tok = o.tokens[(size_t) pos_start]; + } + if (o.prefill_chunk > 0 && n_prompt > 1 && !prompt_windows) { void* borrow = nullptr; uint64_t borrow_bytes = 0; if (!o.no_prefill_borrow && !host_res.empty() && d_res != nullptr) { @@ -13107,6 +13198,40 @@ int main(int argc, char** argv) { ms_adapt += std::chrono::duration(Clock::now() - ta).count(); return true; }; + if (prompt_windows && spec_pos > 0) { + // the prompt's tokens [0, spec_pos) through the verify windows, o.spec at a time: every token is + // committed and the picks are discarded (no head sampling), as the serve's read_windows + const Clock::time_point tw = Clock::now(); + ver.set_head_sampling(false); + std::vector win((size_t) o.spec), outw((size_t) o.spec), nxt((size_t) o.spec); + for (int64_t q = 0; q < spec_pos;) { + const int T = (int) std::min(o.spec, spec_pos - q); + for (int t = 0; t < T; ++t) { + win[(size_t) t] = (int32_t) o.tokens[(size_t) (q + t)]; + nxt[(size_t) t] = (int32_t) o.tokens[(size_t) (q + t + 1)]; + } + drive.d.layers = 0; + drive.d.experts = 0; + drive.d.failed = false; + if (!ver.run(T, win.data(), q, &drive_pool_multi, &drive, outw.data(), err) || drive.d.failed || + !ver.commit(T, err) || (use_mtp && !mtp.prefill(ver.final_R_all(), nxt.data(), T, q, err))) { + if (drive.d.failed && drive.d.fail) err = drive.d.fail; + std::fprintf(stderr, "strata generate: reading the prompt through the verify windows: %s\n", + err.c_str()); + return 1; + } + q += T; + } + if (!ver.wait_commit(err)) { + std::fprintf(stderr, "strata generate: %s\n", err.c_str()); + return 1; + } + ver.set_head_sampling(true); + const double ms = std::chrono::duration(Clock::now() - tw).count(); + prefill_ms += ms; + std::fprintf(stderr, "strata generate: prompt read through the verify windows: %lld tokens in %.1f ms " + "(%.1f tok/s)\n", (long long) spec_pos, ms, ms > 0 ? 1000.0 * (double) spec_pos / ms : 0.0); + } int64_t p = spec_pos; int32_t x = (int32_t) tok; std::vector drafts((size_t) o.spec, 0); @@ -13421,6 +13546,7 @@ int main(int argc, char** argv) { if (o.stage_timing) strata::core::stage_timing_report(g.n_layers); const int64_t decoded = (int64_t) produced.size(); + mem_mark("the answer (the windows' and the draft step's graphs)"); std::printf("prompt :"); for (int64_t t : o.tokens) std::printf(" %lld", (long long) t); std::printf("\noutput :"); diff --git a/sycl/src/prefill/kernels.dp.cpp b/sycl/src/prefill/kernels.dp.cpp index fcf27f0f038..c5202e0dc40 100644 --- a/sycl/src/prefill/kernels.dp.cpp +++ b/sycl/src/prefill/kernels.dp.cpp @@ -2481,7 +2481,16 @@ void gr_broadcast(const float* e, float* R, int64_t T, void* stream) { } check("gr_broadcast"); } -void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream) { +// The prompt kernels' declarations take the model's sizes since qwen35moe (Qwen3.6-35B-A3B, docs/QWEN36.md); this +// port has Flash-Next's kernels only, so any other size stops here instead of running the wrong shapes. +static void flash_next_only(const char* what, bool ok) { + if (ok) return; + std::fprintf(stderr, "%s: the SYCL engine has Flash-Next's prompt kernels only (Qwen3.6 runs on CUDA)\n", what); + std::exit(1); +} +void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream, + int64_t hv) { + flash_next_only("gdn_gates", hv == HV); { auto exp_props = sycl::ext::oneapi::experimental::properties{ sycl::ext::oneapi::experimental::use_root_sync}; @@ -2497,7 +2506,9 @@ void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate } check("gdn_gates"); } -void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream) { +void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream, + int64_t hv) { + flash_next_only("gdn_conv", hv == HV); static const bool serial = std::getenv("STRATA_GDN_CONV_SERIAL") != nullptr; // the old walk (A/B) if (serial || T <= CONV_TILE) { auto exp_props = sycl::ext::oneapi::experimental::properties{ @@ -2602,7 +2613,8 @@ void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, i } void gdn_recurrence_variant(int variant, float* state, const float* h, const float* gate, const float* beta, const float* z, const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, - void* stream, int64_t ld16) { + void* stream, int64_t ld16, int64_t hv, bool silu) { + flash_next_only("gdn_recurrence", hv == HV && !silu); if (ld16 <= 0) ld16 = (int64_t) HV * S; #if defined(__HIPCC__) if (variant == 3) { // diagnostics: the quad recurrence + the old norm kernel @@ -2734,15 +2746,17 @@ void gdn_recurrence_variant(int variant, float* state, const float* h, const flo check("gdn_recurrence"); } void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z, - const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream, int64_t ld16) { + const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream, int64_t ld16, + int64_t hv, bool silu) { // Aurora (S23): four lanes per column + the grid-stride norm, the same bits (tests/hip/gdn_rec_head.cpp); // STRATA_GDN_HEAD=1 (on by default on gfx1151; off with STRATA_GDN_HEAD=0 or STRATA_GDN_REC_HEADS) takes them // (AMD builds only; the arch defaults set STRATA_GDN_HEAD=1 on gfx1151 and leave other cards on the earlier kernels) static const bool head = [] { const char* v = std::getenv("STRATA_GDN_HEAD"); return v != nullptr && std::atoi(v) != 0; }(); static const bool serial = std::getenv("STRATA_GDN_REC_HEADS") != nullptr; - gdn_recurrence_variant(head && !serial ? 1 : 0, state, h, gate, beta, z, gamma, eps, y, y16, T, stream, ld16); + gdn_recurrence_variant(head && !serial ? 1 : 0, state, h, gate, beta, z, gamma, eps, y, y16, T, stream, ld16, hv, silu); } -void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream) { +void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream, int64_t k) { + flash_next_only("route", k == 10); if (n_expert == 512) { auto exp_props = sycl::ext::oneapi::experimental::properties{ @@ -2815,7 +2829,8 @@ void blob_dequant_f16(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, voi } check("blob_dequant_f16"); } -void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream) { +void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream, int64_t n_ff) { + flash_next_only("swiglu_interleaved", n_ff == 640); if (n <= 0) return; { auto exp_props = sycl::ext::oneapi::experimental::properties{ @@ -2832,7 +2847,8 @@ void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream) } check("swiglu_interleaved"); } -void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream) { +void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream, int64_t n_ff) { + flash_next_only("swiglu_pair", n_ff == 640); { auto exp_props = sycl::ext::oneapi::experimental::properties{ sycl::ext::oneapi::experimental::use_root_sync}; @@ -2934,7 +2950,8 @@ void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int check("gather_rows16"); } void moe_combine(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, - int64_t T, void* stream) { + int64_t T, void* stream, int64_t n_embd, int64_t k) { + flash_next_only("moe_combine", n_embd == 2560 && k == 10); #ifndef STRATA_W_NO_COMB if (((reinterpret_cast(Dm) | reinterpret_cast(shared) | reinterpret_cast(bo)) & 15) == 0) { { @@ -3050,7 +3067,8 @@ void rope(float* x, int64_t T, int64_t heads, int64_t dim, int64_t ld, int64_t p } check("rope"); } -void split_q(const float* q_full, float* q, int64_t T, void* stream) { +void split_q(const float* q_full, float* q, int64_t T, void* stream, int64_t n_head) { + flash_next_only("split_q", n_head == 24); { auto exp_props = sycl::ext::oneapi::experimental::properties{ sycl::ext::oneapi::experimental::use_root_sync}; @@ -3066,7 +3084,9 @@ void split_q(const float* q_full, float* q, int64_t T, void* stream) { } check("split_q"); } -void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream, int64_t ld16) { +void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t T, void* stream, int64_t ld16, + int64_t n_head) { + flash_next_only("gate_attn", n_head == 24); { auto exp_props = sycl::ext::oneapi::experimental::properties{ sycl::ext::oneapi::experimental::use_root_sync}; diff --git a/tests/cuda/router_fast.cpp b/tests/cuda/router_fast.cpp new file mode 100644 index 00000000000..f7db6a89f53 --- /dev/null +++ b/tests/cuda/router_fast.cpp @@ -0,0 +1,24 @@ +// The fast router (router_top10_fast_kernel, router_top10's default for 64 < n_expert <= 512, k <= 32) against the +// portable kernel on CUDA: the HIP test (tests/hip/router_fast.cpp) unchanged, its runtime calls mapped onto the CUDA +// runtime. Bitwise on ids and weights: normal rows, exact ties, near-ties, wide spreads, all-equal rows, NaN rows, the +// forced serial-sum path, and the 256-expert / k=8 geometry (qwen35moe). +#include + +#define hipDeviceSynchronize cudaDeviceSynchronize +#define hipError_t cudaError_t +#define hipEventCreate cudaEventCreate +#define hipEventElapsedTime cudaEventElapsedTime +#define hipEventRecord cudaEventRecord +#define hipEventSynchronize cudaEventSynchronize +#define hipEvent_t cudaEvent_t +#define hipFree cudaFree +#define hipGetErrorString cudaGetErrorString +#define hipMalloc cudaMalloc +#define hipMemcpy cudaMemcpy +#define hipMemcpyDeviceToHost cudaMemcpyDeviceToHost +#define hipMemcpyHostToDevice cudaMemcpyHostToDevice +#define hipMemset cudaMemset +#define hipSuccess cudaSuccess +#define STRATA_ROUTER_FAST_ON_CUDA 1 + +#include "../hip/router_fast.cpp" diff --git a/tests/hip/router_fast.cpp b/tests/hip/router_fast.cpp index d0db85e5981..9a11ff6c3eb 100644 --- a/tests/hip/router_fast.cpp +++ b/tests/hip/router_fast.cpp @@ -2,7 +2,9 @@ // kernel, BITWISE: ids and weights of 65,536 rows of 512 logits - normal rows, rows with many exact ties, near-ties, // a wide spread (tiny probabilities), all-equal rows and rows with a NaN - plus the fast kernel forced onto its // serial-sum path, and a 256-expert / k=8 geometry. Then the time per single-token call of each. +#if !defined(STRATA_ROUTER_FAST_ON_CUDA) // tests/cuda/router_fast.cpp maps the calls onto CUDA #include +#endif #include "strata/kernels/router_top10.hpp" #include diff --git a/tools/iq_fixture.py b/tools/iq_fixture.py index 8cfcf921983..b965444a100 100644 --- a/tools/iq_fixture.py +++ b/tools/iq_fixture.py @@ -61,6 +61,8 @@ "IQ4_NL": (32, 18), "IQ4_XS": (256, 136), "Q3_K": (256, 110), + "Q2_K": (256, 84), + "Q4_0": (32, 18), } HALF_ONE = struct.pack(" None: if name == "Q3_K": raw[:, 108:110] = np.frombuffer(HALF_ONE, dtype=np.uint8) return + if name == "Q2_K": # 16 scale bytes, 64 code bytes, then the fp16 d and dmin + raw[:, 80:82] = np.frombuffer(HALF_ONE, dtype=np.uint8) + raw[:, 82:84] = np.frombuffer(HALF_ONE, dtype=np.uint8) + return if name == "IQ1_M": raw[:, 48:56] &= np.uint8(0x0F) # clear the four scale nibbles for byte, nib in ((49, 0x30), (51, 0xC0), (54, 0x00), (55, 0x00)): @@ -83,7 +89,7 @@ def repair_scales(name: str, raw: np.ndarray) -> None: raw[:, 0:2] = np.frombuffer(HALF_ONE, dtype=np.uint8) # name -> (the KERNEL type id written into the .bin header, identical to the gguf-py enum id) KERNEL_IDS = {"IQ2_XXS": 16, "IQ2_XS": 17, "IQ2_S": 22, "IQ3_XXS": 18, "IQ3_S": 21, - "IQ1_M": 29, "IQ4_NL": 20, "IQ4_XS": 23, "Q2_0": 42, "Q3_K": 11} + "IQ1_M": 29, "IQ4_NL": 20, "IQ4_XS": 23, "Q2_0": 42, "Q3_K": 11, "Q2_K": 10, "Q4_0": 2} def build(name: str, rows: int, cols: int, seed: int): diff --git a/tools/iq_pack.py b/tools/iq_pack.py index 5ce0e4a52e8..ede51e98204 100644 --- a/tools/iq_pack.py +++ b/tools/iq_pack.py @@ -219,6 +219,20 @@ def tensor_bytes(mm, g, t) -> np.ndarray: return mm[g.data_start + t.offset: g.data_start + t.offset + n] +def main_layers(model) -> int | None: + """The model's main layers when the file also carries MTP layers after them (a qwen35moe file's + nextn_predict_layers: blk.
.* is the draft layer, which is not part of the main pack), else None.""" + meta = model.files[0].metadata + arch = meta.get("general.architecture", "") + if "%s.nextn_predict_layers" % arch not in meta or "%s.block_count" % arch not in meta: + return None + return int(meta["%s.block_count" % arch]) - int(meta["%s.nextn_predict_layers" % arch]) + + +def is_mtp_layer(name: str, n_main: int | None) -> bool: + return n_main is not None and name.startswith("blk.") and int(name.split(".")[1]) >= n_main + + def is_expert(name: str) -> bool: return name.startswith("blk.") and name.endswith(("_exps.weight",)) @@ -275,8 +289,9 @@ def index_standalone(src, out, model: Model, compat_bf16: bool = False) -> int: """Every non-expert tensor of the model: the floats the engine reads from the pack into dense.bin in the form it reads them (FORM; converted when stored otherwise, see above), quantized ones served natively from the GGUF.""" todo, problems = [], [] + n_main = main_layers(model) for name, (g, t, mm, p) in model.where.items(): - if is_expert(t.name) or t.name in NOT_IN_PACK: + if is_expert(t.name) or t.name in NOT_IN_PACK or is_mtp_layer(t.name, n_main): continue if len(t.shape) > 2: print("tensor %s has %d dimensions; the index holds two" % (t.name, len(t.shape))) @@ -435,7 +450,9 @@ def expert_layout(model: Model, src: pathlib.Path): exps = [n for n in T if n.startswith("blk.") and n.endswith("_exps.weight")] if not exps: return "the model has no expert tensors (blk.N.ffn_{gate,up,down}_exps.weight)" - n_layers = 1 + max(int(n.split(".")[1]) for n in exps) + n_layers = main_layers(model) + if n_layers is None: + n_layers = 1 + max(int(n.split(".")[1]) for n in exps) n_expert = int(T["blk.0.ffn_gate_inp.weight"].shape[1]) # router rows = experts kept (pruned models ship < 512) if any(int(T["blk.%d.ffn_gate_inp.weight" % l].shape[1]) != n_expert for l in range(n_layers)): return "the routers disagree on the expert count; a per-layer pruned model cannot be packed" diff --git a/tools/make_profile.py b/tools/make_profile.py index ba041879df6..fe1650e32cc 100644 --- a/tools/make_profile.py +++ b/tools/make_profile.py @@ -20,6 +20,7 @@ python tools/make_profile.py [TRACE ...] [--base data/expert-profile.bin | --no-base] [--reorder] [--out PATH] [--n-expert 256] (a pruned model: GSQ-RCO Coder keeps 256 of 512) + [--n-layer 40] (Qwen3.6-35B-A3B: 40 layers of 256 experts) A routing trace comes from a one-shot engine run with `--dump-routing FILE` (a prompt typical of your use; the routed experts of every layer and position are written). Point the model config's `--expert-profile` at the result. @@ -98,6 +99,7 @@ def take(pairs, key): def main(): + global N_LAYER ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("traces", nargs="*", help="routing traces from --dump-routing") ap.add_argument("--base", default=str(ROOT / "data" / "expert-profile.bin"), help="ranking to keep first") @@ -107,7 +109,9 @@ def main(): "base that already ranks every pair, where --base alone cannot change the order") ap.add_argument("--out", default=str(ROOT / "data" / "expert-profile.bin")) ap.add_argument("--n-expert", type=int, default=N_EXPERT, help="experts per layer (default 512)") + ap.add_argument("--n-layer", type=int, default=N_LAYER, help="layers (default 48; Qwen3.6-35B-A3B: 40)") a = ap.parse_args() + N_LAYER = a.n_layer ne = a.n_expert freq = defaultdict(int) diff --git a/tools/strata_mcp.py b/tools/strata_mcp.py index 9de78e84d51..39ab794c056 100644 --- a/tools/strata_mcp.py +++ b/tools/strata_mcp.py @@ -66,6 +66,18 @@ "less than ~80 GB of RAM part of its experts are read from the SSD", "download_gb": 93.7, "ram_gb": 48, "arena_gb": 59.5, "families": ("unsloth",), "budget": True, "vision": True}, + "Qwen3.6-UD-IQ4_XS": {"size": "UD-IQ4_XS", "about": "~4-bit i-quant (Unsloth Dynamic), the recommended size: " + "~14 GiB of experts in RAM, comfortable on 32 GB", + "download_gb": 18.2, "ram_gb": 26, "arena_gb": 15.2, "families": ("qwen36",)}, + "Qwen3.6-UD-IQ3_S": {"size": "UD-IQ3_S", "about": "~3.5-bit i-quant (Unsloth Dynamic), smaller: for 16-24 GB of " + "RAM (16 GB: a 12 GB card, the low-RAM mode)", + "download_gb": 15.3, "ram_gb": 24, "arena_gb": 12.9, "families": ("qwen36",)}, + "Ornith-1.5-IQ4_XS": {"size": "IQ4_XS", "about": "~4-bit i-quant (bartowski), the recommended size: " + "~16 GiB of experts in RAM, comfortable on 32 GB", + "download_gb": 19.3, "ram_gb": 28, "arena_gb": 17.1, "families": ("ornith",)}, + "Ornith-1.5-IQ3_XXS": {"size": "IQ3_XXS", "about": "~3-bit i-quant (bartowski), smaller: for 16-24 GB of RAM " + "(16 GB: a 12 GB card, the low-RAM mode)", + "download_gb": 15.3, "ram_gb": 24, "arena_gb": 13.3, "families": ("ornith",)}, } FALLBACK_FAMILIES = { "qwen": {"title": "Qwen3.8-Flash-Next", "about": "the original model", "tag": ""}, @@ -77,6 +89,10 @@ "of RAM part of its experts are read from the SSD " "(UD-Q4_K_XL, 111 GB: experimental)", "tag": "unsloth-", "vision": False}, + "qwen36": {"title": "Qwen3.6-35B-A3B", "about": "a smaller model for 16-32 GB of RAM and 8-12 GB cards (18 GB " + "download); no images yet", "tag": "qwen36-", "vision": False}, + "ornith": {"title": "Ornith-1.5-35B-A3B", "about": "Qwen3.6-35B-A3B's size, tuned for coding agents (19 GB " + "download); no images yet", "tag": "ornith15-", "vision": False}, } FALLBACK_CONTEXTS = [8192, 32768, 65536, 131072, 204800, 262144, 393216, 524288] BENCH_PROMPT = ("Write a short story (about 300 words) about a lighthouse keeper who finds a message in a bottle. " @@ -645,8 +661,10 @@ def _hardware(self) -> dict: def recommend(self, hw: dict) -> dict: """setup.py's own rules: the size it preselects (IQ3_XXS from 60 GB of RAM, else the first, Q2_0), the - Coder below the full model's RAM (setup's RAM table), the low-RAM mode when the GPU makes up for the RAM, and - the context it preselects for the card's VRAM (32K under 14 GB, 64K under 20 GB, else 128K).""" + Coder below the full model's RAM (setup's RAM table), Qwen3.6-35B-A3B where no Flash-Next size reaches the + RAM (setup's qwen36_recommended: the 16-24 GB PCs, the size by qwen36_size), the low-RAM mode when the GPU + makes up for the RAM, and the context it preselects for the card's VRAM (32K under 14 GB, 64K under 20 GB, + else 128K; Qwen3.6 on a card under 5.5 GB: 8K).""" models, families, _ctx, _src = self.tables() S = self.setup_module() usable = [g for g in hw.get("gpus", []) if g.get("usable")] @@ -667,6 +685,23 @@ def low_fits(m): except Exception: # noqa: BLE001 return False notes = [] + small_vram = getattr(S, "SMALL_MODEL_VRAM_GB", 5.5) if S else 5.5 + + def q36_pick(): + """(size key, fits) for Qwen3.6 on this PC, as setup picks it; ("", False) without the family.""" + sizes = [m for m in self.sizes_of(models, "qwen36") if not models[m].get("experimental")] + if not sizes: + return "", False + try: + if S: + m = S.small_size("qwen36", ram, vram) + return m, not S.low_ram_needed(m, ram) or S.low_ram_fits(m, ram, vram) + except Exception: # noqa: BLE001 + pass + m = next((x for x in sizes if ram >= models[x]["ram_gb"] - 4), sizes[-1]) + return m, ram >= models[m]["ram_gb"] - 4 or low_fits(m) + q36_wanted = (S.qwen36_recommended(ram) if S and hasattr(S, "qwen36_recommended") + else ram < models["IQ1_M"]["ram_gb"] - 4) and "qwen36" in families if ram >= 60: fam, model, why = "qwen", "IQ3_XXS", f"{ram:.0f} GB of RAM: setup's own pick from 60 GB (better quality)" elif ram >= models["Q2_0"]["ram_gb"] - 4: @@ -674,24 +709,45 @@ def low_fits(m): elif ram >= models["IQ1_M"]["ram_gb"] - 4: fam, model, why = "coder", "IQ1_M", (f"{ram:.0f} GB of RAM: the full model needs ~48 GB; the Coder " "(half the experts, best for code) needs ~32 GB") + elif q36_wanted and q36_pick()[1]: + fam, model = "qwen36", q36_pick()[0] + low = bool(S) and S.low_ram_needed(model, ram) + why = (f"{ram:.0f} GB of RAM: no Qwen3.8-Flash-Next size fits (the Coder needs ~32 GB); Qwen3.6-35B-A3B, " + f"the smaller model of the same family, keeps ~{models[model]['arena_gb']:.0f} GB of experts in RAM" + + (f" - in setup's low-RAM mode, the {vram:.0f} GB card holding part of them" if low else "")) elif low_fits("IQ1_M"): fam, model, why = "coder", "IQ1_M", (f"{ram:.0f} GB of RAM is below the Coder's 32 GB, but the GPU's " f"{vram:.0f} GB make up for it (setup's low-RAM mode: slower)") + elif q36_wanted: + m = q36_pick()[0] + return {"family": None, "model": None, + "why": f"{ram:.0f} GB of RAM with a {vram:.0f} GB card: the smallest model, Qwen3.6-35B-A3B " + f"{models[m].get('size', m)}, keeps ~{models[m]['arena_gb']:.0f} GB of experts in RAM; on " + "16 GB of RAM it needs a 12 GB card (setup's low-RAM mode)"} else: return {"family": None, "model": None, "why": f"{ram:.0f} GB of RAM: Strata needs 32 GB or more " "(the smallest model, the Coder, keeps ~23 GB of experts in " "RAM)"} + small = fam == "qwen36" ctx = 32768 if vram < 14 else 65536 if vram < 20 else 131072 - if vram < 11: - notes.append("less than 12 GB of VRAM: it runs, but slowly (most experts stay on the CPU)") + if small and vram < small_vram: + ctx = 8192 + notes.append(S.small_model_vram_note(vram) if S and hasattr(S, "small_model_vram_note") else + f"a {vram:.0f} GB card is at this model's floor: an 8K context and --draft-vocab en") + elif vram < 11 and not small: + notes.append("less than 12 GB of VRAM: it runs, but slowly (most experts stay on the CPU); " + "Qwen3.6-35B-A3B (--family qwen36) is made for 6-12 GB cards") if hw.get("cpu", {}).get("avx2") is False: # #623: a warning, not a stop (setup compiles for it) notes.append("this CPU has no AVX2: EXPERIMENTAL and slow - setup compiles the engine on this PC for the " "older CPU (10-20 minutes), and the CPU's share of the experts runs a few times slower") if backend == "hip": - notes.append("AMD (experimental, Linux): the engine is compiled during setup; no images") + notes.append("AMD (experimental, Linux): the engine is compiled during setup; no images" + + ("; Qwen3.6-35B-A3B is untested on AMD and reads prompts through the decode windows " + "(its batched prompt path is NVIDIA-only)" if small else "")) + size = models[model].get("size", model) # Qwen3.6's keys carry the model; setup takes the size name cmd = (("START-HERE.bat --setup" if WIN else "./setup.sh --setup") + - f" --yes --family {fam} --model {model} --context {ctx}" + (" --backend hip" if backend == "hip" else "")) - return {"family": fam, "model": model, "title": families[fam]["title"] + " " + model, "context": ctx, + f" --yes --family {fam} --model {size} --context {ctx}" + (" --backend hip" if backend == "hip" else "")) + return {"family": fam, "model": model, "title": families[fam]["title"] + " " + size, "context": ctx, "backend": backend, "gpu": {k: best.get(k) for k in ("index", "name", "vram_gb")}, "download_gb": models[model]["download_gb"], "why": why, "notes": notes, "setup_command": cmd} @@ -707,7 +763,7 @@ def model_table(self, hw: dict | None) -> list: sizes = [] for m in self.sizes_of(models, f): d = models[m] - tag = (fd.get("tag", "") + m).lower() + tag = (fd.get("tag", "") + d.get("size", m)).lower() # (Qwen3.6's keys carry the model) verdict = None if ram is not None: verdict = "fits" if ram >= d["ram_gb"] else "tight" if ram >= d["ram_gb"] - 8 else "does not fit" @@ -1225,6 +1281,8 @@ def install_plan(self, family, model, context, vision, kv, gpu, gpus, backend, l fams = [f for f in families if model in s.sizes_of(models, f)] family = "qwen" if "qwen" in fams else fams[0] sizes = s.sizes_of(models, family) + if model not in sizes: # a published size name (qwen36's UD-IQ4_XS) -> its key + model = next((m for m in sizes if models[m].get("size", "").upper() == str(model).upper()), model) if model not in sizes: raise ToolError(f"{families[family]['title']} has no {model}; its sizes: {', '.join(sizes)}") if context is None: @@ -1246,7 +1304,7 @@ def install_plan(self, family, model, context, vision, kv, gpu, gpus, backend, l raise ToolError("the AMD backend has no GPU image encoder yet: use vision=cpu (images on the CPU, " "Linux) or vision=no") target = self.check_data_dir(data_dir) if data_dir else s.data_dir() - tag = (families[family].get("tag", "") + model) + tag = (families[family].get("tag", "") + models[model].get("size", model)) have_dir = target / "models" / tag partly = have_dir.is_dir() and any(have_dir.glob("*.gguf*")) need = (8 if partly else models[model]["download_gb"] + 8) + (1 if vision in ("yes", "gpu", "cpu") else 0) diff --git a/tools/test_setup_choices.py b/tools/test_setup_choices.py index 1d1fd992bfb..26f2faf3451 100644 --- a/tools/test_setup_choices.py +++ b/tools/test_setup_choices.py @@ -146,7 +146,7 @@ def test_the_size_of_another_family(self): self.assertEqual(code, 1) self.assertIn("has no IQ3_XXS model file", out) self.assertIn("choose one of: UD-IQ4_XS, UD-Q4_K_XL (or IQ3_XXS: --family qwen --model IQ3_XXS, --family swift --model " - "IQ3_XXS)", out) + "IQ3_XXS, --family ornith --model IQ3_XXS)", out) self.assertIn("Strata runs ISTA-DASLab's GSQ-RCO files", out) diff --git a/tools/test_setup_ornith.py b/tools/test_setup_ornith.py new file mode 100644 index 00000000000..a981f61b8fb --- /dev/null +++ b/tools/test_setup_ornith.py @@ -0,0 +1,66 @@ +"""tools/test_setup_ornith.py - setup's "ornith" family (Ornith-1.5-35B-A3B): its table, its download, its engine +arguments and its size choice. No GPU, no download: the harness of test_setup_qwen36.py with the family switched. + + python tools/test_setup_ornith.py +""" +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import test_setup_qwen36 as q36 # noqa: E402 (the qwen36 harness: setup imported, main() mocked) + +setup = q36.setup +IQ4, IQ3 = "Ornith-1.5-IQ4_XS", "Ornith-1.5-IQ3_XXS" + + +class Tables(unittest.TestCase): + def test_family(self): + fam = setup.FAMILIES["ornith"] + self.assertEqual(list(setup.FAMILIES)[-1], "ornith") # last: the menu's other numbers stay + self.assertTrue(setup.small_family("ornith") and setup.small_family("qwen36")) + self.assertFalse(setup.small_family("qwen") or setup.small_family("unsloth")) + for k, v in {"tag": "ornith15-", "own_mtp": True, "ple": False, "one_gpu": True, + "mmproj": "mmproj-Ornith-1.5-35B-A3B-bf16.gguf", "profile": "expert-profile-qwen36.bin", "shards": 1}.items(): + self.assertEqual(fam[k], v, k) + self.assertIn("/bartowski/Ornith-1.5-35B-A3B-GGUF/resolve/64b0493d34a5ca4c1b4ad67bb99b41d74b4f07d6/", fam["hf"]) + + def test_sizes(self): + self.assertEqual(setup.model_key("ornith", "IQ4_XS"), IQ4) # not Flash-Next's IQ3_XXS / Unsloth's + self.assertEqual(setup.model_key("ornith", "IQ3_XXS"), IQ3) + self.assertEqual(setup.model_key("qwen", "IQ3_XXS"), "IQ3_XXS") + self.assertEqual([m for m in setup.MODELS if "ornith" in setup.MODELS[m].get("families", ())], [IQ4, IQ3]) + for m, gb in ((IQ4, 19278554784), (IQ3, 15340447392)): + name = setup.model_file(setup.FAMILIES["ornith"], m, 1) + self.assertEqual(setup.ORNITH_FILES[name][0], gb, name) + + def test_size_by_ram(self): + self.assertEqual(setup.small_size("ornith", q36.RAM32), IQ4) + self.assertEqual(setup.small_size("ornith", q36.RAM24), IQ3) + self.assertEqual(setup.small_size("ornith", q36.RAM16), IQ3) # the smallest, in the low-RAM mode + + def test_qwen36_stays_the_suggestion(self): + self.assertTrue(setup.qwen36_recommended(q36.RAM16)) + self.assertEqual(setup.qwen36_size(q36.RAM24), "Qwen3.6-UD-IQ3_S") + + +class Install(q36.Base): + def test_iq4_xs(self): + code, out, cfg = self.main(["--family", "ornith", "--model", "IQ4_XS"]) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-ornith15-iq4_xs.json") + native = q36.arg(cfg, "--native") + self.assertTrue(native.endswith("models/ornith15-IQ4_XS/Ornith-1.5-35B-A3B-IQ4_XS.gguf"), native) + self.assertEqual(q36.arg(cfg, "--mtp"), native) # the draft layer is in the model file + self.assertTrue(q36.arg(cfg, "--pack").endswith("packs/ornith15-iq4_xs")) + self.assertTrue(q36.arg(cfg, "--expert-profile").endswith("expert-profile-qwen36.bin")) + self.assertNotIn("--ple-gguf", cfg["args"]) + self.assertEqual(cfg["model_name"], "ornith-1.5-35b-a3b-iq4_xs") + self.assertTrue(any("bartowski/Ornith-1.5-35B-A3B-GGUF" in u for u in self.downloads), self.downloads) + self.assertIn(("Ornith-1.5-35B-A3B-IQ4_XS.gguf", 19278554784, + "d6aef57fa948e9bba3ca4959b3c237ed898c605471f48c73a32cedbd24aabe70"), self.verified) + + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/test_setup_qwen36.py b/tools/test_setup_qwen36.py new file mode 100644 index 00000000000..c415ce13f62 --- /dev/null +++ b/tools/test_setup_qwen36.py @@ -0,0 +1,678 @@ +"""Tests for setup.py's Qwen3.6-35B-A3B family (--family qwen36): the two pinned single-file sizes with their sizes +and SHA-256, the engine arguments (no --ple-gguf, --mtp = the model file with the shipped draft vocabulary, its own +expert profile), no MTP download, the data folder's names, the RAM-fit suggestion on 16/24/32/64 GB PCs, the AMD +warning and one GPU. Mocked - no GPU, no downloads, nothing written outside a temp folder. + + python tools/test_setup_qwen36.py +""" +from __future__ import annotations + +import contextlib +import io +import json +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +import setup # noqa: E402 + +REV = "5bc3e238d916f48a861bac2f8a1990a0e9b7e98d" +BASE = f"https://huggingface.co/unsloth/Qwen3.6-35B-A3B-MTP-GGUF/resolve/{REV}/" +IQ4, IQ3 = "Qwen3.6-UD-IQ4_XS", "Qwen3.6-UD-IQ3_S" +F_IQ4, F_IQ3 = "Qwen3.6-35B-A3B-UD-IQ4_XS.gguf", "Qwen3.6-35B-A3B-UD-IQ3_S.gguf" +# what ram_gb() reports (GiB) on PCs sold as 16, 24, 32 and 64 GB +RAM16, RAM24, RAM32, RAM64 = 15.6, 23.4, 31.3, 62.7 + + +def quiet(fn, *args): + """-> (exit code or None, printed text).""" + out = io.StringIO() + code = None + with contextlib.redirect_stdout(out): + try: + code = fn(*args) + except SystemExit as e: + code = e.code + return code, out.getvalue() + + +def arg(cfg, flag): + a = cfg["args"] + return a[a.index(flag) + 1] + + +class Tables(unittest.TestCase): + def test_family(self): + fam = setup.FAMILIES["qwen36"] + self.assertEqual(list(setup.FAMILIES)[0], "qwen") # never the default family + self.assertEqual(list(setup.FAMILIES)[-2:], ["qwen36", "ornith"]) # the menu's numbers stay as they were + self.assertEqual(fam["tag"], "qwen36-") + self.assertEqual(fam["name"], "qwen3.6-35b-a3b") + self.assertEqual(fam["hf"], BASE) + self.assertEqual(fam["pack_args"], ["--compat-bf16"]) + self.assertEqual(fam["profile"], "expert-profile-qwen36.bin") + self.assertNotIn("vision", fam) # images: Qwen3.6's mmproj (bartowski's) + self.assertEqual(fam["mmproj"], "mmproj-Qwen_Qwen3.6-35B-A3B-bf16.gguf") + self.assertTrue(fam["own_mtp"] and fam["one_gpu"] and fam["nvidia_only"]) + self.assertIs(fam["ple"], False) + self.assertIn("Apache 2.0", fam["license"]) + + def test_sizes(self): + self.assertEqual([m for m in setup.MODELS if "qwen36" in setup.MODELS[m].get("families", ())], [IQ4, IQ3]) + self.assertEqual((setup.size_of(IQ4), setup.size_of(IQ3)), ("UD-IQ4_XS", "UD-IQ3_S")) + self.assertEqual(setup.MODELS[IQ4]["arena_gb"], 15.2) # 14.17 GiB, measured + self.assertAlmostEqual(setup.MODELS[IQ4]["arena_gb"], 14.17 * 2**30 / 1e9, places=1) + self.assertEqual(setup.MODELS[IQ3]["arena_gb"], 12.9) # measured (11.99 GiB) + self.assertAlmostEqual(setup.MODELS[IQ4]["download_gb"], 18209036576 / 1e9, places=1) + self.assertAlmostEqual(setup.MODELS[IQ3]["download_gb"], 15346432288 / 1e9, places=1) + for m in (IQ4, IQ3): + self.assertEqual(setup.MODELS[m]["engine"], setup.QWEN36_ENGINE) + self.assertFalse(setup.MODELS[m].get("budget")) + # the engine version required: this source tree's (CMakeLists.txt) + self.assertEqual(setup.QWEN36_ENGINE, tuple(int(x) for x in setup.source_version().split("."))) + # Unsloth's Flash-Next UD-IQ4_XS keeps its key and its numbers + self.assertEqual(setup.MODELS["UD-IQ4_XS"]["families"], ("unsloth",)) + self.assertEqual(setup.MODELS["UD-IQ4_XS"]["arena_gb"], 59.5) + + def test_files_and_pins(self): + fam = setup.FAMILIES["qwen36"] + self.assertEqual([setup.model_file(fam, m, 1) for m in (IQ4, IQ3)], [F_IQ4, F_IQ3]) + self.assertEqual([setup.model_shards(fam, m) for m in (IQ4, IQ3)], [1, 1]) + self.assertEqual(fam["sha256"], { + F_IQ4: (18209036576, "df27a780435b7b45c2597536112ea3cb091f8544c3d0c3318d9f4258b31f7adf"), + F_IQ3: (15346432288, "ab639a7f330f96c47d3e6c2dd2d6445182e7b763e17e6048dc850a71bbc9f27f")}) + self.assertEqual(setup.HF_REVISIONS["unsloth/Qwen3.6-35B-A3B-MTP-GGUF"], REV) + + def test_model_key(self): + self.assertEqual(setup.model_key("qwen36", "UD-IQ4_XS"), IQ4) + self.assertEqual(setup.model_key("qwen36", "ud-iq3_s"), IQ3) + self.assertEqual(setup.model_key("qwen36", IQ4), IQ4) + self.assertEqual(setup.model_key("unsloth", "UD-IQ4_XS"), "UD-IQ4_XS") + self.assertIsNone(setup.model_key("qwen", "UD-IQ3_S")) + self.assertIsNone(setup.model_key("qwen36", "IQ3_S")) + for size in ("UD-IQ4_XS", "UD-IQ3_S", IQ4): + self.assertIn(size, setup.model_choices()) + + def test_gguf_names(self): + self.assertEqual(setup.gguf_choice(F_IQ4), ("qwen36", "UD-IQ4_XS")) + self.assertEqual(setup.gguf_choice(F_IQ3), ("qwen36", "UD-IQ3_S")) + self.assertIsNone(setup.gguf_unsupported(F_IQ3)) + self.assertEqual(setup.gguf_choice("Qwen3.8-Flash-Next-UD-IQ4_XS-00001-of-00003.gguf"), ("unsloth", "UD-IQ4_XS")) + + def test_config_names(self): + with tempfile.TemporaryDirectory() as t: + p = Path(t) / "strata-qwen36-ud-iq4_xs.json" + p.write_text(json.dumps({"args": ["--max-context", "32768", "--kv", "int8"]})) + ch = setup.choices_from_config(p) + self.assertEqual((ch["family"], ch["model"], ch["context"]), ("qwen36", IQ4, 32768)) + p = Path(t) / "strata-qwen36-ud-iq3_s.json" + p.write_text(json.dumps({"args": []})) + self.assertEqual(setup.choices_from_config(p)["model"], IQ3) + p = Path(t) / "strata-unsloth-ud-iq4_xs.json" # Unsloth's Flash-Next file is not taken + p.write_text(json.dumps({"args": []})) + self.assertEqual((setup.choices_from_config(p)["family"], setup.choices_from_config(p)["model"]), + ("unsloth", "UD-IQ4_XS")) + self.assertTrue(setup.one_gpu_family(Path("strata-qwen36-ud-iq4_xs.json"))) + self.assertFalse(setup.one_gpu_family(Path("strata-iq2_xs.json"))) + + +class RamFit(unittest.TestCase): + """Qwen3.6 is suggested where no Flash-Next size reaches its RAM line (16-24 GB); elsewhere nothing changes.""" + + def test_suggested_family(self): + self.assertTrue(setup.qwen36_recommended(RAM16)) + self.assertTrue(setup.qwen36_recommended(RAM24)) + self.assertFalse(setup.qwen36_recommended(RAM32)) + self.assertFalse(setup.qwen36_recommended(RAM64)) + + def test_suggested_size(self): + self.assertEqual(setup.qwen36_size(RAM16), IQ3) # in the low-RAM mode + self.assertEqual(setup.qwen36_size(RAM24), IQ3) + self.assertEqual(setup.qwen36_size(RAM32), IQ4) + self.assertEqual(setup.qwen36_size(RAM64), IQ4) + + def test_low_ram_rules(self): + self.assertFalse(setup.low_ram_needed(IQ4, RAM32)) # comfortable on 32 GB + self.assertTrue(setup.low_ram_needed(IQ4, RAM24)) + self.assertFalse(setup.low_ram_needed(IQ3, RAM24)) # 24 GB: IQ3_S without the low-RAM mode + self.assertTrue(setup.low_ram_needed(IQ3, RAM16)) + # 16 GB with UD-IQ3_S's measured 12.9 GB of experts: a 12 GB card fits the low-RAM rule, an 8 GB card misses + # it by ~0.5 GB (setup then stops and says --model ... --yes) + self.assertTrue(setup.low_ram_fits(IQ3, RAM16, 11.9)) + self.assertFalse(setup.low_ram_fits(IQ3, RAM16, 7.99)) + self.assertFalse(setup.low_ram_fits(IQ3, 8.0, 7.99)) + + +class FakeGGUF: + """gguf_reader.GGUFFile: Flash-Next's shard 2 holds the PLE table; Qwen3.6's file has none.""" + def __init__(self, path): + names = ["per_layer_token_embd.weight"] if "00002-of" in str(path) else ["blk.0.ffn_up_exps.weight"] + self.tensors = [types.SimpleNamespace(name=n) for n in names] + + +def nvidia(n=1, vram=11.9): + return [{"index": i, "name": "NVIDIA GeForce RTX 4070", "vram_gb": vram, "arch": "89", "driver": "580.97"} + for i in range(n)] + + +class Base(unittest.TestCase): + """setup.main(), every outside effect mocked.""" + + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.t = Path(self.tmp.name) + (self.t / "data" / "mtp" / "rt").mkdir(parents=True) # Flash-Next's draft layer (for --family qwen) + (self.t / "data" / "mtp" / "rt" / "experts.bin").write_bytes(b"") + self.downloads, self.runs, self.verified, self.asked = [], [], [], [] + self.vram_free = None # nvidia-smi's memory.free: unknown + + def tearDown(self): + self.tmp.cleanup() + + def main(self, argv, ram=RAM64, version="0.1.40", found=None, amd=(), answers=None, free=500.0, source="local", + runs_small=True): + """answers: None = --yes; else {words of a question: its answer} (Enter for the others). + -> (exit code, printed text, the config written or None).""" + def fake_input(prompt=""): + if answers is None: + raise AssertionError(f"asked {prompt!r}") + self.asked.append(prompt) + return next((v for k, v in answers.items() if k in prompt), "") + + eng = self.t / "engine" + eng.mkdir(exist_ok=True) + (eng / "BUILD.json").write_text(json.dumps({"version": version, "source": source})) + # runs_small: the engine names qwen35moe's keys (this checkout's); False: a ready-made 0.1.40 or an older build + (eng / setup.EXE).write_bytes(b"\0" + (setup.SMALL_ENGINE_MARK if runs_small else b"qwen4exp.block_count")) + self.builds = [] + + def fake_build(*a, **k): # a compile here: this checkout's engine + self.builds.append(a) + if json.loads((eng / "BUILD.json").read_text()).get("source") == "local": + return eng # compiled here, same source: kept as is + (eng / "BUILD.json").write_text(json.dumps({"version": "0.1.40", "source": "local"})) + (eng / setup.EXE).write_bytes(setup.SMALL_ENGINE_MARK) + return eng + found = nvidia() if found is None else found + for old in self.t.glob("strata-*.json"): # this run's config only + old.unlink() + + def fake_download(url, dst, what=None): + self.downloads.append(url) + dst.parent.mkdir(parents=True, exist_ok=True) + dst.write_bytes(b"") + setup.mark(dst) + + patches = [ + mock.patch.object(setup, "ROOT", self.t), + mock.patch.object(setup, "GPU_PICK", None), + mock.patch.object(setup, "data_folder", lambda d: (self.t / "data", [])), + mock.patch.object(setup, "installed_configs", lambda: []), + mock.patch.object(setup, "gpus", lambda: found), + mock.patch.object(setup, "gpu_free_gb", lambda i: self.vram_free), + mock.patch.object(setup, "amd_gpus", lambda *a: list(amd)), + mock.patch.object(setup, "ram_gb", lambda: ram), + mock.patch.object(setup, "cpu_info", lambda: ("Test CPU", True, False)), + mock.patch.object(setup, "cpu_cores", lambda: None), + mock.patch.object(setup, "page_file_gb", lambda: 16.0), + mock.patch.object(setup, "free_gb", lambda p: free), + mock.patch.object(setup, "rotational_disk", lambda p: None), + mock.patch.object(setup, "is_wsl", lambda: False), + mock.patch.object(setup, "linux_desktop", lambda env=None: False), + mock.patch.object(setup, "pip_install", lambda *a, **k: None), + mock.patch.object(setup, "get_llama_cpp", lambda: self.t / "llama.cpp"), + mock.patch.object(setup, "get_prebuilt", lambda *a, **k: eng), + mock.patch.object(setup, "build_engine_hip", lambda *a, **k: eng), + mock.patch.object(setup, "build_engine", fake_build), + mock.patch.object(setup, "pip_cuda_libs", lambda *a, **k: None), + mock.patch.object(setup, "hipblaslt_table", lambda *a, **k: None), + mock.patch.object(setup, "download", fake_download), + mock.patch.object(setup, "check_shards", lambda shards: None), + mock.patch.object(setup, "verify_sha256", lambda s, size, sha: self.verified.append((s.name, size, sha))), + mock.patch.object(setup, "run", lambda cmd, *a, **k: self.runs.append([str(x) for x in cmd])), + mock.patch.object(setup, "refresh_draft_vocab", lambda *a, **k: None), + mock.patch.object(setup, "mtp_corrupt", lambda *a, **k: False), + mock.patch.object(setup, "write_run_script", lambda tag, cfg, port, *_: self.t / f"run-{tag}.sh"), + mock.patch.object(setup, "saved_calibration", lambda cfg: None), + mock.patch.object(setup, "start", mock.Mock(side_effect=AssertionError("started"))), + mock.patch.dict(sys.modules, {"gguf_reader": types.SimpleNamespace(GGUFFile=FakeGGUF)}), + mock.patch.object(sys, "argv", ["setup.py", *(["--yes"] if answers is None else []), "--no-start", + "--models-dir", str(self.t / "data" / "models"), *argv]), + mock.patch("builtins.input", fake_input), + ] + with contextlib.ExitStack() as st: + for p in patches: + st.enter_context(p) + code, out = quiet(setup.main) + cfgs = sorted(self.t.glob("strata-*.json")) + self.assertLessEqual(len(cfgs), 1, cfgs) + self.cfg_path = cfgs[0] if cfgs else None + return code, out, (json.loads(cfgs[0].read_text()) if cfgs else None) + + +class Install(Base): + def test_iq4_xs_engine_args(self): + code, out, cfg = self.main(["--setup", "--family", "qwen36", "--model", "UD-IQ4_XS"]) + self.assertEqual(code, 0, out) + data = self.t / "data" + gguf = data / "models" / "qwen36-UD-IQ4_XS" / F_IQ4 + pack = data / "packs" / "qwen36-ud-iq4_xs" + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq4_xs.json") + self.assertEqual(self.downloads, [BASE + F_IQ4]) # one file; no MTP, no image encoder + self.assertEqual(self.verified, [(F_IQ4, *setup.QWEN36_FILES[F_IQ4])]) + self.assertEqual(self.runs, [[sys.executable, str(self.t / "tools" / "iq_pack.py"), "--gguf", str(gguf), + "--out", str(pack), "--compat-bf16"]]) # no mtp_fetch/mtp_pack/mtp_rt + self.assertEqual(cfg["args"], [ + "--pack", str(pack), "--native", str(gguf), + "--expert-profile", str(self.t / "data" / "expert-profile-qwen36.bin"), "--expert-cache", "auto", + "--prefill", "auto", "--spec", "4", "--spec-min-p", "0.5", + "--mtp", str(gguf), "--mtp-draft-vocab", str(self.t / "data" / "draft_vocab.bin"), + "--max-context", "32768", "--kv", "int8"]) + self.assertEqual(cfg["model_name"], "qwen3.6-35b-a3b-ud-iq4_xs") + self.assertEqual(cfg["tokenizer"], str(pack / "tokenizer")) + self.assertEqual(cfg["log"], str(self.t / "strata-qwen36-ud-iq4_xs.log")) + self.assertNotIn("vision", cfg) + self.assertIn("MTP draft layer: the model file's own", out) + self.assertIn("Apache 2.0", out) + + def test_iq3_s(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ3_S"], ram=RAM24) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq3_s.json") + self.assertEqual(self.downloads, [BASE + F_IQ3]) + self.assertTrue(arg(cfg, "--native").endswith("models/qwen36-UD-IQ3_S/" + F_IQ3)) + self.assertEqual(arg(cfg, "--mtp"), arg(cfg, "--native")) + self.assertEqual(cfg["model_name"], "qwen3.6-35b-a3b-ud-iq3_s") + self.assertNotIn("--ple-gguf", cfg["args"]) + + def test_draft_vocab_choice(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS", "--draft-vocab", "en"]) + self.assertEqual(code, 0, out) + self.assertEqual(arg(cfg, "--mtp-draft-vocab"), str(self.t / "data" / "draft_vocab_en.bin")) + self.assertEqual(cfg["draft_vocab"], "en") + + def test_files_already_there(self): + # a hand-made install: the file in models/qwen36-UD-IQ4_XS and a finished pack: nothing is downloaded or built + gguf = self.t / "data" / "models" / "qwen36-UD-IQ4_XS" / F_IQ4 + gguf.parent.mkdir(parents=True) + gguf.write_bytes(b"") + setup.mark(gguf) + pack = self.t / "data" / "packs" / "qwen36-ud-iq4_xs" + (pack / "tokenizer").mkdir(parents=True) + (pack / "tokenizer" / "vocab.json").write_text("{}") + (pack / "native_experts.txt").write_text("") + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS"]) + self.assertEqual(code, 0, out) + self.assertEqual((self.downloads, self.runs), ([], [])) + self.assertEqual(len(self.verified), 1) # the SHA-256 is still checked (once) + self.assertEqual(arg(cfg, "--pack"), str(pack)) + + def test_wrong_size_for_the_family(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "IQ3_S"]) + self.assertEqual(code, 1) + self.assertIn("choose one of: UD-IQ4_XS, UD-IQ3_S", out) + self.assertIsNone(cfg) + + def test_old_engine_stops_before_the_download(self): + # an engine compiled here from a checkout without qwen35moe (setup does not recompile it: same source) + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS"], version="0.1.39", runs_small=False) + self.assertEqual(code, 1) + self.assertIn("needs an engine with the qwen35moe support", out) + self.assertEqual(self.downloads, []) + + def test_ready_made_engine_without_it_is_compiled_here(self): + # the published 0.1.40 engines (Windows) predate qwen35moe: setup compiles this checkout's engine instead + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS"], source="release", runs_small=False) + self.assertEqual(code, 0, out) + self.assertEqual(len(self.builds), 1) + self.assertIn("does not run Qwen3.6-35B-A3B yet: compiling the engine on this PC instead", out) + self.assertEqual(cfg["exe"], str(self.t / "engine" / setup.EXE)) + + def test_ready_made_engine_with_it_is_used(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS"], source="release") + self.assertEqual(code, 0, out) + self.assertEqual(self.builds, []) + + def test_flash_next_keeps_the_ready_made_engine(self): + code, out, cfg = self.main(["--family", "qwen", "--model", "IQ2_XS"], source="release", runs_small=False) + self.assertEqual(code, 0, out) + self.assertEqual(self.builds, []) + + def test_images(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS", "--vision", "yes"]) + self.assertEqual(code, 0, out) + mmproj = self.t / "data" / "models" / "mmproj-Qwen_Qwen3.6-35B-A3B-bf16.gguf" + self.assertEqual(self.downloads[0], BASE + F_IQ4) + self.assertTrue(self.downloads[1].endswith("/bartowski/Qwen_Qwen3.6-35B-A3B-GGUF/resolve/" + "5c2410d71524f4f72b023ce8daf7a80528226d5f/" + "mmproj-Qwen_Qwen3.6-35B-A3B-bf16.gguf")) + self.assertEqual(len(self.downloads), 2) + self.assertEqual(cfg["vision"]["mmproj"], str(mmproj)) + self.assertIs(cfg["vision"]["gpu"], True) + self.assertEqual(cfg["args"][-3:], ["--vision", "--vram-reserve-mib", "700"]) + self.assertNotIn("--vision cpu", out) # the prompt chunks keep their size (measured) + + def test_disk_space_counts_no_mtp(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS"], free=19.0) + self.assertEqual(code, 1) + self.assertIn("not enough free disk space", out) + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS"], free=19.5) # 18.2 + 1 + self.assertEqual(code, 0, out) + + +class Suggestion(Base): + """No --family: Qwen3.6 on 16 and 24 GB, the Flash-Next default as before on 32 and 64 GB.""" + + def test_16_gb(self): + # an 8 GB card: UD-IQ3_S (12.9 GB of experts, measured) is just under the low-RAM rule on 16 GB + code, out, cfg = self.main([], ram=RAM16, found=nvidia(vram=7.99)) + self.assertNotEqual(code, 0, out) + self.assertIn("--yes", out) + + def test_16_gb_12_gb_card(self): + code, out, cfg = self.main([], ram=RAM16, found=nvidia(vram=11.9)) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq3_s.json") + self.assertTrue("--mmap-experts" in cfg["args"] or "--resident-experts" in cfg["args"], cfg["args"]) + self.assertNotIn("--ple-gguf", cfg["args"]) + + def test_16_gb_12_gb_card_maps_the_rest(self): + # the experts a 12 GB card does not hold (~6 GB) leave 15.6 GB short of the resident variant's 10 GB of room: + # the plain low-RAM mode, read through the file cache + code, out, cfg = self.main([], ram=RAM16, found=nvidia(vram=11.9)) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq3_s.json") + self.assertIn("--mmap-experts", cfg["args"]) + self.assertIn([sys.executable, str(self.t / "tools" / "iq_pack.py"), "--gguf", + arg(cfg, "--native"), "--out", arg(cfg, "--pack"), "--experts-bin"], self.runs) + + def test_24_gb(self): + code, out, cfg = self.main([], ram=RAM24) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq3_s.json") + for flag in ("--mmap-experts", "--resident-experts"): + self.assertNotIn(flag, cfg["args"]) + + def test_24_gb_menus(self): + code, out, cfg = self.main([], ram=RAM24, answers={}) # Enter everywhere: the recommendations + self.assertEqual(code, 0, out) + line = next(ln for ln in out.splitlines() if "Qwen3.6-35B-A3B" in ln and ln.strip()[:2].rstrip(")").isdigit()) + self.assertIn("(recommended for 23 GB of RAM: UD-IQ3_S)", line) + family_q = next(q for q in self.asked if "Which model?" in q) + self.assertIn(f"[{list(setup.FAMILIES).index('qwen36') + 1}]", family_q) + size_q = next(q for q in self.asked if "Which size?" in q) + self.assertIn("[2]", size_q) # 1) UD-IQ4_XS 2) UD-IQ3_S + self.assertIn(" 1) UD-IQ4_XS ", out) + self.assertIn(" 2) UD-IQ3_S ~3.5-bit", out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq3_s.json") + + def test_32_gb_default_unchanged(self): + code, out, cfg = self.main([], ram=RAM32) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-q2_0.json") # the original model, as before + self.assertNotIn("recommended for 31 GB", out) + + def test_32_gb_qwen36_picks_iq4_xs(self): + code, out, cfg = self.main(["--family", "qwen36"], ram=RAM32) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq4_xs.json") + for flag in ("--mmap-experts", "--resident-experts"): + self.assertNotIn(flag, cfg["args"]) + + def test_64_gb_default_unchanged(self): + code, out, cfg = self.main([], ram=RAM64) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-iq3_xxs.json") # setup's pick from 60 GB, as before + code, out, cfg = self.main(["--family", "qwen36"], ram=RAM64) + self.assertEqual(code, 0, out) + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq4_xs.json") + + def test_check_16_gb(self): + code, out, cfg = self.main(["--check"], ram=RAM16) + self.assertEqual(code, 0, out) + self.assertIn("suggests Qwen3.6-35B-A3B UD-IQ3_S (--family qwen36)", out) + self.assertIn("This PC can run Strata", out) + self.assertNotIn("[!] RAM", out) + code, out, cfg = self.main(["--check"], ram=RAM64) + self.assertNotIn("--family qwen36", out) + + def test_flash_next_floor_unchanged_with_a_family(self): + # --family qwen on 20 GB: the Flash-Next floor (the Coder's 32 GB) stops --yes as before + code, out, cfg = self.main(["--family", "qwen"], ram=19.9, found=nvidia(vram=8.0)) + self.assertEqual(code, 1) + self.assertIn("the smallest model (the Coder) needs about 32 GB", out) + + +class SmallCard(Base): + """A card under 5.5 GB: setup recommends an 8K context and says what was measured there; a 6 GB card keeps 32K; + Flash-Next keeps its "less than 12 GB" warning.""" + + def test_4_gb_card(self): + code, out, cfg = self.main(["--family", "qwen36"], ram=RAM32, found=nvidia(vram=3.99)) + self.assertEqual(code, 0, out) + self.assertEqual(arg(cfg, "--max-context"), "8192") + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq3_s.json") # the smaller size (32 GB would take IQ4_XS) + self.assertIn("at this model's floor", out) + self.assertNotIn("it will be slow", out) + + def test_6_gb_card(self): + code, out, cfg = self.main(["--family", "qwen36"], ram=RAM32, found=nvidia(vram=5.99)) + self.assertEqual(code, 0, out) + self.assertEqual(arg(cfg, "--max-context"), "32768") + self.assertEqual(self.cfg_path.name, "strata-qwen36-ud-iq4_xs.json") + self.assertNotIn("at this model's floor", out) + self.assertNotIn("it will be slow", out) + + def test_context_asked_is_kept(self): + code, out, cfg = self.main(["--family", "qwen36", "--context", "32768"], ram=RAM32, found=nvidia(vram=3.99)) + self.assertEqual(code, 0, out) + self.assertEqual(arg(cfg, "--max-context"), "32768") # a recommendation, never forced (#406) + + def test_flash_next_keeps_its_warning(self): + code, out, cfg = self.main(["--family", "qwen", "--model", "IQ2_XS"], ram=RAM64, found=nvidia(vram=7.99)) + self.assertIn("it will be slow", out) + self.assertNotIn("at this model's floor", out) + + +class KvStreaming(Base): + """From a 64K context, a 35B-A3B model's KV streaming (it attends to every cell: writing slows down past 32K tokens + when streamed): from 128K with under 9 GB of VRAM free it is needed, and setup says so; otherwise setup asks and + recommends no; --kv-streaming on|off answers it. Flash-Next streams as before, unasked.""" + Q = "Turn on KV streaming?" + + def run_ctx(self, ctx, vram, free, answers=None, extra=()): + self.vram_free = free + return self.main(["--family", "qwen36", "--model", "UD-IQ4_XS", "--context", str(ctx), *extra], ram=RAM64, + found=nvidia(vram=vram), answers=answers) + + def asked_q(self): + return [p for p in self.asked if self.Q in p] + + def test_asked_recommends_no(self): + code, out, cfg = self.run_ctx(65536, 7.99, 7.9, answers={}) # Enter: the recommendation + self.assertEqual(code, 0, out) + self.assertEqual(self.asked_q(), [f"{self.Q} [n]: "]) + self.assertNotIn("--kv-resident", cfg["args"]) + self.assertIn("past 32K tokens", out) + self.assertIn("the 8 GB card has 7.9 GB free now", out) + self.assertIn("strata-qwen36-ud-iq4_xs.json's args, or run\n setup again with --kv-streaming on|off.", out) + self.assertIn("KV streaming off, as you chose: the KV cache stays in VRAM", out) + + def test_asked_yes(self): + code, out, cfg = self.run_ctx(131072, 11.99, 11.6, answers={self.Q: "y"}) + self.assertEqual(code, 0, out) + self.assertEqual(self.asked_q(), [f"{self.Q} [n]: "]) + self.assertEqual(arg(cfg, "--kv-resident"), "32768") + self.assertIn("writing slows down past 32K tokens of context", out) + + def test_yes_takes_the_recommendation(self): + code, out, cfg = self.run_ctx(65536, 11.99, 11.6) + self.assertEqual(code, 0, out) + self.assertNotIn("--kv-resident", cfg["args"]) + + def test_needed_is_said_not_asked(self): + for ctx, vram, free in ((131072, 7.99, 7.9), (262144, 11.99, 8.5)): # free VRAM decides, not the card's size + with self.subTest(ctx=ctx, free=free): + self.asked.clear() + code, out, cfg = self.run_ctx(ctx, vram, free, answers={}) + self.assertEqual(code, 0, out) + self.assertFalse(self.asked_q()) + self.assertEqual(arg(cfg, "--kv-resident"), "32768") + self.assertIn("KV streaming is needed here", out) + self.assertIn(f"has {free:.1f} GB free now", out) + self.assertIn("--kv-streaming on|off", out) + self.assertIn("needed for this context with this much free VRAM", out) + + def test_not_needed_with_room_or_at_64k(self): + for ctx, free in ((131072, 10.0), (65536, 6.0)): + with self.subTest(ctx=ctx, free=free): + self.asked.clear() + code, out, cfg = self.run_ctx(ctx, 11.99, free, answers={}) + self.assertEqual(code, 0, out) + self.assertEqual(len(self.asked_q()), 1) + self.assertNotIn("is needed here", out) + + def test_the_image_encoder_takes_its_share_first(self): + code, out, cfg = self.run_ctx(131072, 11.99, 10.0, answers={}, extra=("--vision", "yes")) # 10.0 - 1.9 + self.assertEqual(code, 0, out) + self.assertFalse(self.asked_q()) + self.assertEqual(arg(cfg, "--kv-resident"), "32768") + self.assertIn("has 10.0 GB free now, 8.1 GB beside the image encoder", out) + self.assertIn("KV streaming is needed here", out) + + def test_the_option_answers_it(self): + code, out, cfg = self.run_ctx(65536, 7.99, 7.9, answers={}, extra=("--kv-streaming", "on")) + self.assertEqual(code, 0, out) + self.assertFalse(self.asked_q()) + self.assertEqual(arg(cfg, "--kv-resident"), "32768") + code, out, cfg = self.run_ctx(131072, 7.99, 7.9, answers={}, extra=("--kv-streaming", "off")) + self.assertEqual(code, 0, out) + self.assertFalse(self.asked_q()) + self.assertNotIn("--kv-resident", cfg["args"]) + self.assertIn("the first prompt fails without it", out) + + def test_flash_next_unasked(self): + self.vram_free = 7.9 + code, out, cfg = self.main(["--family", "qwen", "--model", "IQ2_XS", "--context", "65536"], ram=RAM64, + found=nvidia(vram=11.9), answers={}) + self.assertEqual(code, 0, out) + self.assertFalse(self.asked_q()) + self.assertEqual(arg(cfg, "--kv-resident"), "32768") + self.assertNotIn("writing slows down", out) + + +class FreeVram(unittest.TestCase): + def test_parse(self): + for text, want in (("6144\n", 6.0), ("", None), ("[N/A]\n", None)): + with mock.patch.object(setup, "out", lambda cmd: text): + self.assertEqual(setup.gpu_free_gb(0), want) + + +class Amd(Base): + """The HIP engine compiles, but the qwen35moe batched prompt path is NVIDIA-only and it was never run on AMD.""" + CARD = [{"index": 0, "name": "AMD Radeon RX 7700 XT", "vram_gb": 12.0, "arch": "gfx1101", "driver": "amdgpu"}] + + def test_yes_with_the_family_goes_on(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS", "--backend", "hip"], + found=[], amd=self.CARD) + self.assertEqual(code, 0, out) + self.assertIn("Qwen3.6-35B-A3B has not been run on AMD cards yet", out) + self.assertIn("installing Qwen3.6-35B-A3B UD-IQ4_XS on an AMD card, as you chose", out) + self.assertEqual(cfg["backend"], "hip") + self.assertNotIn("--ple-gguf", cfg["args"]) + + def test_asked_default_no(self): + code, out, cfg = self.main(["--family", "qwen36", "--backend", "hip"], found=[], amd=self.CARD, answers={}) + self.assertEqual(code, 1) + self.assertTrue(any("Try it anyway?" in q and "[n]" in q for q in self.asked), self.asked) + self.assertIn("Qwen3.6-35B-A3B is untested on AMD", out) + self.assertEqual(self.downloads, []) + self.assertIsNone(cfg) + + def test_check_says_untested(self): + code, out, cfg = self.main(["--check", "--backend", "hip"], ram=RAM32, found=[], amd=self.CARD) + self.assertEqual(code, 0, out) + line = next(ln for ln in out.splitlines() if ln.strip().startswith(IQ4)) + self.assertIn("untested on AMD", line) + + +class Pascal(Base): + """Below sm_75 the engine reads every prompt through the decode windows, and the model was never run there.""" + CARD = [{"index": 0, "name": "NVIDIA GeForce GTX 1080", "vram_gb": 8.0, "arch": "61", "driver": "535.104"}] + + def test_asked_default_no_before_the_download(self): + code, out, cfg = self.main(["--family", "qwen36"], found=self.CARD, answers={}) + self.assertEqual(code, 1) + self.assertTrue(any("Try it anyway?" in q and "[n]" in q for q in self.asked), self.asked) + self.assertIn("Qwen3.6-35B-A3B is untested below sm_75", out) + self.assertEqual(self.downloads, []) + self.assertIsNone(cfg) + + def test_yes_with_the_model_goes_on(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ3_S"], found=self.CARD) + self.assertEqual(code, 0, out) + self.assertIn("has not been run on GPU 0 (NVIDIA GeForce GTX 1080, 8 GB) (sm_61)", out) + self.assertIn("installing Qwen3.6-35B-A3B UD-IQ3_S on GPU 0", out) + + def test_turing_is_not_asked(self): + card = [{**self.CARD[0], "name": "NVIDIA GeForce RTX 2080 SUPER", "arch": "75", "driver": "580.97"}] + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ3_S"], found=card) + self.assertEqual(code, 0, out) + self.assertNotIn("sm_75", out.replace("(sm_75)", "")) + self.assertNotIn("untested below", out) + + +class OneGpu(Base): + def test_gpus_keeps_one(self): + code, out, cfg = self.main(["--family", "qwen36", "--model", "UD-IQ4_XS", "--gpus", "0,1"], + found=nvidia(2)) + self.assertEqual(code, 0, out) + self.assertEqual(cfg["gpu"], 0) + self.assertNotIn("layer_split", cfg) + self.assertNotIn("--remote-expert-opt", cfg["args"]) + self.assertIn("runs on one GPU for now", out) + + def test_not_offered_at_a_start(self): + p = self.t / "strata-qwen36-ud-iq4_xs.json" + p.write_text(json.dumps({"exe": "x", "args": []})) + with mock.patch.object(setup, "gpus", mock.Mock(side_effect=AssertionError("looked for GPUs"))): + code, out = quiet(setup.offer_together, p, {"args": []}, True) + self.assertEqual(code, {"args": []}) # unchanged, nothing asked + self.assertEqual(out, "") + + def test_start_with_gpus_uses_the_first(self): + exe = self.t / "strata" + exe.write_bytes(b"") + p = self.t / "strata-qwen36-ud-iq4_xs.json" + p.write_text(json.dumps({"exe": str(exe), "args": ["--pack", "p"], "gpu": 0, "gpus_asked": True})) + call = mock.Mock(return_value=0) + with mock.patch.object(setup, "gpus", lambda: nvidia(2)), \ + mock.patch.object(setup, "ensure_engine_for", lambda cards, path, cfg, yes: cfg), \ + mock.patch.object(setup, "is_wsl", lambda: False), \ + mock.patch.object(setup.subprocess, "call", call): + code, out = quiet(setup.start, p, None, [0, 1], False, True) + self.assertEqual(code, 0, out) + self.assertIn("runs on one GPU for now", out) + self.assertEqual(json.loads(p.read_text())["gpu"], 0) + cmd = call.call_args[0][0] + self.assertEqual(cmd[cmd.index("--gpu") + 1], "0") + + +class DraftVocabAtStart(unittest.TestCase): + def test_points_at_the_chosen_subset(self): + data = setup.ROOT / "data" + cfg = {"args": ["--mtp", "m.gguf", "--mtp-draft-vocab", str(data / "draft_vocab.bin")], "draft_vocab": "en"} + self.assertTrue(setup.point_draft_vocab(cfg)) + self.assertEqual(cfg["args"][-1], str(data / "draft_vocab_en.bin")) + self.assertFalse(setup.point_draft_vocab(cfg)) # already there + own = {"args": ["--mtp-draft-vocab", "/home/me/mine.bin"], "draft_vocab": "en"} + self.assertFalse(setup.point_draft_vocab(own)) # a file of the user's own is kept + self.assertFalse(setup.point_draft_vocab({"args": ["--mtp", "rt"]})) # Flash-Next: refresh_draft_vocab + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/test_setup_unsloth.py b/tools/test_setup_unsloth.py index e9bdaa3d233..af9d777d4c3 100644 --- a/tools/test_setup_unsloth.py +++ b/tools/test_setup_unsloth.py @@ -533,11 +533,12 @@ def test_check_verdict(self): self.assertIn("EXPERIMENTAL", next(ln for ln in out.splitlines() if ln.strip().startswith(M))) def test_check_below_every_floor(self): - # #977: 16 GB fits no size; the RAM line is not [ok] and the verdict says so - code, out, _ = self.main(["--check"], ram=16.0, m=X) + # #977: 8 GB fits no size (16 GB fits Qwen3.6-35B-A3B's UD-IQ3_S in the low-RAM mode since the qwen36 + # family); the RAM line is not [ok] and the verdict says so + code, out, _ = self.main(["--check"], ram=8.0, m=X) self.assertEqual(code, 1, out) self.assertNotIn("[ok] RAM", out) - self.assertIn("[!] RAM: 16 GB (less than", out) + self.assertIn("[!] RAM: 8 GB (less than", out) self.assertIn("This PC cannot run Strata yet", out) self.assertNotIn("This PC can run Strata", out) diff --git a/tools/test_strata_mcp.py b/tools/test_strata_mcp.py index f1229f3e53a..92722f4b1ef 100644 --- a/tools/test_strata_mcp.py +++ b/tools/test_strata_mcp.py @@ -596,8 +596,19 @@ def test_recommendation_follows_ram(self): for ram, want in cases.items(): r = s.recommend({"ram_gb": ram, "gpus": gpu, "cpu": {"avx2": True}}) self.assertEqual((r["family"], r["model"], r["context"]), want, ram) - r = s.recommend({"ram_gb": 16, "gpus": [{**gpu[0], "vram_gb": 8.0}], "cpu": {"avx2": True}}) + # below the Coder's RAM: Qwen3.6-35B-A3B, setup's suggestion there (a 16 GB PC lists ~15.6 GB) + r = s.recommend({"ram_gb": 23.4, "gpus": [{**gpu[0], "vram_gb": 11.9}], "cpu": {"avx2": True}}) + self.assertEqual((r["family"], r["model"], r["context"]), ("qwen36", "Qwen3.6-UD-IQ3_S", 32768)) + self.assertIn("--family qwen36 --model UD-IQ3_S ", r["setup_command"]) + r = s.recommend({"ram_gb": 15.6, "gpus": [{**gpu[0], "vram_gb": 11.9}], "cpu": {"avx2": True}}) + self.assertEqual(r["model"], "Qwen3.6-UD-IQ3_S") + self.assertIn("low-RAM mode", r["why"]) + r = s.recommend({"ram_gb": 15.6, "gpus": [{**gpu[0], "vram_gb": 8.0}], "cpu": {"avx2": True}}) self.assertIsNone(r["model"]) + self.assertIn("12 GB card", r["why"]) + r = s.recommend({"ram_gb": 23.4, "gpus": [{**gpu[0], "vram_gb": 4.0}], "cpu": {"avx2": True}}) + self.assertEqual((r["model"], r["context"]), ("Qwen3.6-UD-IQ3_S", 8192)) # a 4 GB card: at the floor + self.assertTrue(any("floor" in n for n in r["notes"]), r["notes"]) r = s.recommend({"ram_gb": 64, "gpus": [], "cpu": {"avx2": True}}) self.assertIsNone(r["model"]) r = s.recommend({"ram_gb": 64, "gpus": [{**gpu[0], "vram_gb": 12.0}], "cpu": {"avx2": True}})