chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+18
@@ -0,0 +1,18 @@
|
||||
/ds4
|
||||
/ds4-server
|
||||
/ds4-bench
|
||||
/ds4-eval
|
||||
/ds4-agent
|
||||
/ds4_native
|
||||
/ds4_server_test
|
||||
/ds4_test
|
||||
/ds4flash.gguf
|
||||
/TODO.md
|
||||
/gguf/
|
||||
*.o
|
||||
*.dSYM/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
/misc/
|
||||
.*.swp
|
||||
.DS_Store
|
||||
@@ -0,0 +1,56 @@
|
||||
# Agent Notes
|
||||
|
||||
`ds4.c` is a DeepSeek V4 Flash specific inference engine. It is not a generic
|
||||
GGUF runner. The goal is a small, readable, high-performance C codebase with
|
||||
Objective-C only where Metal requires it and Metal kernels under `metal/`.
|
||||
|
||||
## Goals
|
||||
|
||||
- Keep the production path as whole-model Metal graph inference.
|
||||
- Always make sure that the SSD streaming, CUDA, distributed inference, Metal default inference are not affected by fixes to other parts of the code.
|
||||
- Keep model loading mmap-backed for the Metal default case; do not eagerly copy the full GGUF. Keep the model loading for SSD streaming of routed experts explicit: allocated buffers, fast reads from disk, always try to hide loading of missing routed experts by loading them while performing the inference of the shared expert and routed experts already in RAM. Always try to hide loading of layers for prefill in SSD streaming mode using the inference time of the current layer as the next one is loaded.
|
||||
- Keep the CPU backend CPU-only and use it only as reference/debug code.
|
||||
- Preserve correctness before speed. Do not keep a faster path with unexplained attention, KV cache, or logits drift.
|
||||
- Make long local agent sessions practical through live KV reuse and disk KV checkpoints.
|
||||
|
||||
## Quality Rules
|
||||
|
||||
- Keep the implementation small, sharp, easy to understand. Try to write elegant code in a state of grace. Don't settle for the first thing that comes to mind, try to find the most minimal and better working design. Don't introduce slop: very fragile code that just patches specific cases, dead code, useless code and code ways more complicated of how it should be.
|
||||
- Comment important inference code where the model mechanics, cache lifetime, memory policy, or API orchestration are not obvious from the local code.
|
||||
- Prefer comments beside the implementation over separate design documents.
|
||||
- Keep comments instructive and compact: explain why a shape, ordering, cache boundary, or memory choice exists.
|
||||
- Keep public APIs narrow. CLI/server code should not know tensor internals.
|
||||
- Do not add permanent semantic variants behind flags. Diagnostic switches are fine when they validate the one release path.
|
||||
- Do not introduce C++.
|
||||
|
||||
## Safety
|
||||
|
||||
- Avoid large CPU inference runs on macOS; the CPU path has previously exposed kernel VM failures with very large mappings.
|
||||
- Do not run multiple huge model processes concurrently. The instance lock is intentional.
|
||||
|
||||
## Layout
|
||||
|
||||
- `ds4.c`: model loading, tokenizer, CPU reference code, Metal graph scheduling,
|
||||
sessions, disk-cache payload serialization.
|
||||
- `ds4_cli.c`: command line, linenoise REPL, interactive transcript handling.
|
||||
- `ds4_server.c`: OpenAI/Anthropic compatible HTTP API, worker queue, streaming,
|
||||
tool-call mapping, disk KV cache policy.
|
||||
- `ds4_metal.m`: Objective-C Metal runtime and kernel wrappers.
|
||||
- `metal/*.metal`: compute kernels.
|
||||
- `tests/`: unit and live integration tests.
|
||||
- `misc/`: ignored notes, experiments, and old planning material.
|
||||
|
||||
This list is not complete, check the files for more info.
|
||||
|
||||
## Testing
|
||||
|
||||
Use `make` for build validation. Use `make test` for unit/regression tests when a
|
||||
model and Metal are available. Use live server tests only when intentionally
|
||||
testing the API surface.
|
||||
|
||||
At every major change where one of the following could be affected, make sure to:
|
||||
|
||||
1. Test the normal Metal path and that speed is still at the level it was.
|
||||
2. Test the SSD streaming path.
|
||||
3. Test the distributed inference if it could be affected, but ask the user before doing so.
|
||||
4. Check if CUDA could be broken after the change, and ask the user to give you access to the CUDA machine to actually test if everything is still fine.
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
# Contributing
|
||||
|
||||
DwarfStar4 changes should be tested against the failure mode they can realistically
|
||||
affect. The project has two regression tracks: correctness and speed. Please
|
||||
include the commands you ran, the machine/backend, the model quant, and any
|
||||
notable failures in the PR or commit notes.
|
||||
|
||||
Do not send PRs affecting one or more inference backends without checking if the
|
||||
resulting code is still correct and fast. The only acceptable regression speed
|
||||
is when an important correctness bug is fixed and it requires some speed penalty.
|
||||
|
||||
## Correctness Regression Tests
|
||||
|
||||
Build the default backend first:
|
||||
|
||||
```sh
|
||||
make clean
|
||||
make
|
||||
```
|
||||
|
||||
The C test runner is `ds4_test`. Running it without arguments is equivalent to
|
||||
`--all`:
|
||||
|
||||
```sh
|
||||
make test
|
||||
```
|
||||
|
||||
Useful narrower checks:
|
||||
|
||||
```sh
|
||||
./ds4_test --server
|
||||
./ds4_test --logprob-vectors
|
||||
./ds4_test --long-context
|
||||
./ds4_test --tool-call-quality
|
||||
./ds4_test --metal-kernels
|
||||
```
|
||||
|
||||
What they cover:
|
||||
|
||||
- `--server`: request parsing, chat rendering, streaming, tool-call parsing,
|
||||
thinking controls, KV disk-cache bookkeeping, and other server-side logic.
|
||||
This is the best quick check for API and prompt-rendering changes.
|
||||
- `--logprob-vectors`: compares local token bytes and top-logprob slices against
|
||||
official DeepSeek V4 Flash continuation vectors. This catches tokenizer,
|
||||
template, attention, and logits regressions.
|
||||
- `--long-context`: runs a long-context story fact-recall regression from
|
||||
`tests/long_context_story_prompt.txt`. The model must retrieve spelled-out
|
||||
person-number assignments from a long prose prompt and return `Name=number`
|
||||
lines that the test parses.
|
||||
- `--tool-call-quality`: exercises actual model behavior for DSML tool-call
|
||||
emission in both fast and exact paths.
|
||||
- `--metal-kernels`: isolated Metal kernel numeric checks.
|
||||
|
||||
The runner defaults to `ds4flash.gguf`. Override paths when needed:
|
||||
|
||||
```sh
|
||||
DS4_TEST_MODEL=/path/to/model.gguf ./ds4_test --logprob-vectors
|
||||
DS4_TEST_VECTOR_FILE=/path/to/official.vec ./ds4_test --logprob-vectors
|
||||
DS4_TEST_LONG_PROMPT=/path/to/prompt.txt ./ds4_test --long-context
|
||||
```
|
||||
|
||||
For CUDA-specific changes, test on a CUDA machine:
|
||||
|
||||
```sh
|
||||
make
|
||||
make cuda-regression
|
||||
```
|
||||
|
||||
For CPU portability, at least verify that the CPU target still builds:
|
||||
|
||||
```sh
|
||||
make cpu
|
||||
```
|
||||
|
||||
The CPU backend is a reference/debug path, not the production performance
|
||||
target. Remember that executing the CPU path on Metal can crash the system
|
||||
because of a kernel bug in macOS.
|
||||
|
||||
## Quality Checks For Quantization Changes
|
||||
|
||||
For GGUF or quantization work, use the official-continuation scorer in
|
||||
`gguf-tools/quality-testing`. The test compares how much probability a local
|
||||
GGUF assigns to official DeepSeek V4 Flash continuations, token by token.
|
||||
|
||||
Build the scorer:
|
||||
|
||||
```sh
|
||||
make -C gguf-tools quality-score
|
||||
```
|
||||
|
||||
Then score old and new GGUFs against the same manifest and compare:
|
||||
|
||||
```sh
|
||||
gguf-tools/quality-testing/score_official OLD.gguf \
|
||||
gguf-tools/quality-testing/data/manifest.tsv /tmp/old.tsv 4096
|
||||
|
||||
gguf-tools/quality-testing/score_official NEW.gguf \
|
||||
gguf-tools/quality-testing/data/manifest.tsv /tmp/new.tsv 4096
|
||||
|
||||
python3 gguf-tools/quality-testing/compare_scores.py /tmp/old.tsv /tmp/new.tsv
|
||||
```
|
||||
|
||||
Lower `avg_nll` is better. See
|
||||
`gguf-tools/quality-testing/README.md` for collecting or refreshing official
|
||||
continuations.
|
||||
|
||||
## Speed Regression Tests
|
||||
|
||||
Use `ds4-bench` for throughput regressions. It reports instantaneous prefill and
|
||||
generation speed at context frontiers, not one whole-run average. Prefill is
|
||||
incremental: each row measures only the newly processed suffix since the
|
||||
previous frontier.
|
||||
|
||||
Default linear sweep:
|
||||
|
||||
```sh
|
||||
./ds4-bench \
|
||||
-m ds4flash.gguf \
|
||||
--prompt-file speed-bench/promessi_sposi.txt \
|
||||
--ctx-start 2048 \
|
||||
--ctx-max 65536 \
|
||||
--step-incr 2048 \
|
||||
--gen-tokens 128 \
|
||||
--csv /tmp/ds4-speed.csv
|
||||
```
|
||||
|
||||
Use the same machine, backend, model file, context sweep, power/thermal state,
|
||||
and background load when comparing two commits. For backend work, run at least
|
||||
one before/after CSV and compare both `prefill_tps` and `gen_tps`. Generation is
|
||||
greedy and skips EOS so each frontier gets the same number of generated tokens.
|
||||
|
||||
To generate a graph for a CSV:
|
||||
|
||||
```sh
|
||||
python3 speed-bench/plot_speed.py /tmp/ds4-speed.csv --title "Machine t/s"
|
||||
```
|
||||
|
||||
## Reporting sessions bugs
|
||||
|
||||
For debugging a failing generation, keep the trace:
|
||||
|
||||
```sh
|
||||
./ds4-server --trace /tmp/ds4-trace.txt ...
|
||||
```
|
||||
@@ -0,0 +1,22 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 The ds4.c authors
|
||||
Copyright (c) 2023-2026 The ggml authors
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+243
@@ -0,0 +1,243 @@
|
||||
# DeepSeek v4 model card synopsis
|
||||
|
||||
This document extracts the most important information from the official
|
||||
DeepSeek-V4-Flash Hugging Face model card, with emphasis on facts that matter
|
||||
for local inference, DS4 development, and benchmark interpretation.
|
||||
|
||||
Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
|
||||
|
||||
## Model Family
|
||||
|
||||
DeepSeek-V4 is a preview model family with two Mixture-of-Experts language
|
||||
models:
|
||||
|
||||
| Model | Total parameters | Active parameters | Context length |
|
||||
|---|---:|---:|---:|
|
||||
| DeepSeek-V4-Flash | 284B | 13B | 1M tokens |
|
||||
| DeepSeek-V4-Pro | 1.6T | 49B | 1M tokens |
|
||||
|
||||
Flash is the smaller and more efficient model. The model card says Flash-Max can
|
||||
approach Pro reasoning performance when given a larger thinking budget, while
|
||||
remaining behind Pro on pure knowledge and the most complex agentic tasks.
|
||||
|
||||
## Architecture
|
||||
|
||||
DeepSeek-V4 uses long-context compressed attention. The model card calls the
|
||||
hybrid design Compressed Sparse Attention (CSA) plus Heavily Compressed
|
||||
Attention (HCA). In DS4 terms, each layer keeps a raw sliding-window KV cache
|
||||
for the latest 128 tokens. This is the high-resolution local context.
|
||||
|
||||
After that raw window, the model uses layer-dependent compressed KV rows:
|
||||
|
||||
| 0-based layer indexes | DS4 ratio | Extra state | Meaning |
|
||||
|---|---:|---|---|
|
||||
| 0, 1 | none | none | Raw 128-token sliding window only |
|
||||
| even layers from 2 onward | 4 | compressed KV + indexer KV | One compressed row per 4 tokens, with an indexer selecting visible compressed rows |
|
||||
| odd layers from 3 onward | 128 | compressed KV | One compressed row per 128 tokens |
|
||||
|
||||
So, after the first two layers, the model alternates ratio-4 and ratio-128
|
||||
compressed attention. A token in a compressed layer attends over both the raw
|
||||
latest-128-token window and the older compressed history. The compression here
|
||||
is time-axis compression: several token positions are pooled into one KV row.
|
||||
The attention rows still use the model attention/value dimensions, so raw and
|
||||
compressed rows can be consumed by the same mixed-attention computation.
|
||||
|
||||
Ratio-4 layers are the selective compressed-attention layers. They maintain a
|
||||
second compressed stream for the indexer, and when the compressed history is
|
||||
larger than the configured top-k, DS4 scores the compressed rows and selects up
|
||||
to 512 of them for attention. Ratio-128 layers are the heavily compressed path:
|
||||
they do not have the indexer stream and use the available ratio-128 compressed
|
||||
rows directly.
|
||||
|
||||
DS4 validates these details from the GGUF metadata. The relevant fixed
|
||||
implementation constants are:
|
||||
|
||||
- Layers: 43
|
||||
- Raw sliding-window attention: 128 tokens
|
||||
- Indexer heads: 64
|
||||
- Indexer head dimension: 128
|
||||
- Indexer top-k: 512
|
||||
|
||||
This is the practical reason the model can expose a 1M-token context without a
|
||||
standard full KV cache for every token in every layer. The model card reports
|
||||
that, at 1M tokens, DeepSeek-V4-Pro needs much less single-token inference
|
||||
compute and KV cache than DeepSeek-V3.2.
|
||||
|
||||
The family also uses:
|
||||
|
||||
- Manifold-Constrained Hyper-Connections (mHC), intended to improve signal
|
||||
propagation stability across layers.
|
||||
- The Muon optimizer, used for faster convergence and training stability.
|
||||
- A post-training pipeline with domain expert cultivation followed by unified
|
||||
consolidation via on-policy distillation.
|
||||
|
||||
## Precision And Weights
|
||||
|
||||
Official download entries include:
|
||||
|
||||
| Model | Precision |
|
||||
|---|---|
|
||||
| DeepSeek-V4-Flash-Base | FP8 Mixed |
|
||||
| DeepSeek-V4-Flash | FP4 + FP8 Mixed |
|
||||
| DeepSeek-V4-Pro-Base | FP8 Mixed |
|
||||
| DeepSeek-V4-Pro | FP4 + FP8 Mixed |
|
||||
|
||||
For the instruct models, the model card describes FP4 + FP8 Mixed as using FP4
|
||||
for MoE expert parameters and FP8 for most other parameters.
|
||||
|
||||
## Reasoning Modes
|
||||
|
||||
The instruct models support three reasoning-effort modes:
|
||||
|
||||
| Mode | Intended behavior | Output shape |
|
||||
|---|---|---|
|
||||
| Non-think | Fast, intuitive replies | `</think>` summary |
|
||||
| High | Deliberate reasoning for harder tasks | `<think>... </think>` summary |
|
||||
| Max | Largest reasoning budget | Special system prompt plus thinking and summary |
|
||||
|
||||
The model card recommends using at least a 384K-token context window for Think
|
||||
Max.
|
||||
|
||||
## Important Flash Benchmarks
|
||||
|
||||
### DeepSeek-V4-Flash Across Reasoning Modes
|
||||
|
||||
| Benchmark | Non-Think | High | Max |
|
||||
|---|---:|---:|---:|
|
||||
| GPQA Diamond Pass@1 | 71.2 | 87.4 | 88.1 |
|
||||
| MMLU-Pro EM | 83.0 | 86.4 | 86.2 |
|
||||
| SimpleQA-Verified Pass@1 | 23.1 | 28.9 | 34.1 |
|
||||
| Chinese-SimpleQA Pass@1 | 71.5 | 73.2 | 78.9 |
|
||||
| HLE Pass@1 | 8.1 | 29.4 | 34.8 |
|
||||
| LiveCodeBench Pass@1 | 55.2 | 88.4 | 91.6 |
|
||||
| HMMT 2026 Feb Pass@1 | 40.8 | 91.9 | 94.8 |
|
||||
| IMOAnswerBench Pass@1 | 41.9 | 85.1 | 88.4 |
|
||||
| SWE Verified Resolved | 73.7 | 78.6 | 79.0 |
|
||||
| Terminal Bench 2.0 Acc | 49.1 | 56.6 | 56.9 |
|
||||
| MCPAtlas Pass@1 | 64.0 | 67.4 | 69.0 |
|
||||
| Toolathlon Pass@1 | 40.7 | 43.5 | 47.8 |
|
||||
|
||||
### DeepSeek-V4-Flash-Base
|
||||
|
||||
The base-model table reports these Flash-Base scores:
|
||||
|
||||
| Benchmark | Shots | Score |
|
||||
|---|---:|---:|
|
||||
| SuperGPQA EM | 5-shot | 46.5 |
|
||||
| MMLU EM | 5-shot | 88.7 |
|
||||
| MMLU-Pro EM | 5-shot | 68.3 |
|
||||
| Simple-QA verified EM | 25-shot | 30.1 |
|
||||
| HumanEval Pass@1 | 0-shot | 69.5 |
|
||||
| GSM8K EM | 8-shot | 90.8 |
|
||||
| LongBench-V2 EM | 1-shot | 44.7 |
|
||||
|
||||
The model card reports SuperGPQA for the base model table, not in the instruct
|
||||
reasoning-mode comparison table.
|
||||
|
||||
## Chat Template And Encoding
|
||||
|
||||
The release does not use a Jinja chat template as the source of truth. The
|
||||
official prompt renderer is the Python code in
|
||||
`encoding/encoding_dsv4.py`, with examples and tests in the same `encoding`
|
||||
directory:
|
||||
|
||||
- https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/raw/main/encoding/encoding_dsv4.py
|
||||
- https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/raw/main/encoding/test_encoding_dsv4.py
|
||||
|
||||
The important special tokens are:
|
||||
|
||||
| Purpose | Token |
|
||||
|---|---|
|
||||
| Beginning of sequence | `<｜begin▁of▁sentence｜>` |
|
||||
| End of assistant turn | `<｜end▁of▁sentence｜>` |
|
||||
| User turn prefix | `<｜User｜>` |
|
||||
| Assistant turn prefix | `<｜Assistant｜>` |
|
||||
| Latest reminder prefix | `<｜latest_reminder｜>` |
|
||||
| Thinking start | `<think>` |
|
||||
| Thinking end / non-thinking marker | `</think>` |
|
||||
| DSML tool markup marker | `｜DSML｜` |
|
||||
|
||||
The renderer accepts `system`, `user`, `assistant`, `tool`,
|
||||
`latest_reminder`, and `developer` roles. The `developer` role is described in
|
||||
the Python comments as an internal search-agent role, not as a normal public
|
||||
chat role.
|
||||
|
||||
Normal chat mode starts with the BOS token, then system text if present, then
|
||||
alternating user and assistant markers. In non-thinking chat mode, a new
|
||||
assistant generation is opened with:
|
||||
|
||||
```text
|
||||
<｜Assistant｜></think>
|
||||
```
|
||||
|
||||
That immediate `</think>` tells the model to skip hidden reasoning and produce
|
||||
the visible answer. In thinking mode, a new assistant generation is opened with:
|
||||
|
||||
```text
|
||||
<｜Assistant｜><think>
|
||||
```
|
||||
|
||||
Completed assistant thinking turns are rendered as reasoning content inside
|
||||
`<think>...</think>`, followed by the visible answer and the EOS token.
|
||||
|
||||
By default, the Python renderer drops earlier assistant reasoning content before
|
||||
the last user message. If tools are present on any message, it disables that
|
||||
reasoning drop and keeps the full reasoning/tool context. `reasoning_effort=max`
|
||||
also prepends a special high-effort instruction prefix before the first rendered
|
||||
message in thinking mode.
|
||||
|
||||
Tool definitions are passed in OpenAI-compatible function schema form, but the
|
||||
model is instructed to emit DSML. A tool call is rendered as a DSML
|
||||
`tool_calls` block containing one or more `invoke` entries, each with named
|
||||
parameters. Parameters carry a `string="true"` flag for raw strings and
|
||||
`string="false"` for JSON values such as numbers, booleans, arrays, or objects.
|
||||
|
||||
DeepSeek-V4 does not render standalone `tool` role messages. The Python
|
||||
preprocessor converts tool results into user content blocks and renders each
|
||||
result as:
|
||||
|
||||
```text
|
||||
<tool_result>...</tool_result>
|
||||
```
|
||||
|
||||
Tool-result bodies are rendered as raw text. Literal `<`, `>`, and `&` from
|
||||
file contents or shell output are preserved; only the exact closing sentinel
|
||||
`</tool_result>` is escaped so the wrapper cannot be terminated by data.
|
||||
|
||||
When there are multiple tool results, the renderer sorts them to match the
|
||||
order of the preceding assistant tool calls.
|
||||
|
||||
The same script also defines special task tokens for internal quick tasks such
|
||||
as title generation, search-query generation, action selection, authority
|
||||
classification, domain classification, and URL-read decisions. Those are
|
||||
separate from normal chat/tool rendering.
|
||||
|
||||
## Local Running Notes
|
||||
|
||||
The model card lists vLLM and SGLang examples for OpenAI-compatible serving.
|
||||
For local deployment, it recommends:
|
||||
|
||||
- `temperature = 1.0`
|
||||
- `top_p = 1.0`
|
||||
- At least 384K context for Think Max
|
||||
|
||||
These are deployment recommendations from the model card, not necessarily the
|
||||
same settings used for deterministic benchmarking. DS4 keeps `top_p=1.0` but
|
||||
adds a local `min_p=0.05` default to avoid sampling tokens whose probability is
|
||||
far below the best token.
|
||||
|
||||
## Licensing
|
||||
|
||||
The repository and model weights are licensed under the MIT License.
|
||||
|
||||
## Citation
|
||||
|
||||
The model card cites:
|
||||
|
||||
```bibtex
|
||||
@misc{deepseekai2026deepseekv4,
|
||||
title={DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence},
|
||||
author={DeepSeek-AI},
|
||||
year={2026},
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,244 @@
|
||||
CC ?= cc
|
||||
UNAME_S := $(shell uname -s)
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
NATIVE_CPU_FLAG ?= -mcpu=native
|
||||
else
|
||||
NATIVE_CPU_FLAG ?= -march=native
|
||||
endif
|
||||
|
||||
DEBUG_FLAGS ?= -g
|
||||
CFLAGS ?= -O3 -ffast-math $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -std=c99
|
||||
OBJCFLAGS ?= -O3 -ffast-math $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -fobjc-arc
|
||||
|
||||
LDLIBS ?= -lm -pthread
|
||||
METAL_SRCS := $(wildcard metal/*.metal)
|
||||
ROCM_SRCS := $(wildcard rocm/*.cuh)
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
METAL_LDLIBS := $(LDLIBS) -framework Foundation -framework Metal
|
||||
CORE_OBJS = ds4.o ds4_distributed.o ds4_ssd.o ds4_metal.o
|
||||
CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_ssd.o
|
||||
else
|
||||
CFLAGS += -D_GNU_SOURCE -fno-finite-math-only
|
||||
CUDA_HOME ?= /usr/local/cuda
|
||||
NVCC ?= $(CUDA_HOME)/bin/nvcc
|
||||
CUDA_ARCH ?=
|
||||
ifneq ($(strip $(CUDA_ARCH)),)
|
||||
NVCC_ARCH_FLAGS := -arch=$(CUDA_ARCH)
|
||||
endif
|
||||
NVCCFLAGS ?= -O3 -g -lineinfo --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread
|
||||
CORE_OBJS = ds4.o ds4_distributed.o ds4_ssd.o ds4_cuda.o
|
||||
CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_ssd.o
|
||||
CUDA_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas
|
||||
HIPCC ?= $(shell command -v hipcc 2>/dev/null || echo /opt/rocm/bin/hipcc)
|
||||
ROCM_ARCH ?= gfx1151
|
||||
ROCM_CFLAGS ?= -O3 -ffast-math -g -fno-finite-math-only -pthread -D__HIP_PLATFORM_AMD__ -Wno-unused-command-line-argument --offload-arch=$(ROCM_ARCH)
|
||||
ROCM_LDLIBS ?= -lm -pthread -lhipblas -lhipblaslt
|
||||
DS4_LINK ?= $(NVCC) $(NVCCFLAGS)
|
||||
DS4_LINK_LIBS ?= $(CUDA_LDLIBS)
|
||||
METAL_LDLIBS := $(LDLIBS)
|
||||
endif
|
||||
|
||||
.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
all: ds4 ds4-server ds4-bench ds4-eval ds4-agent
|
||||
|
||||
help:
|
||||
@echo "DS4 build targets:"
|
||||
@echo " make Build Metal ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent"
|
||||
@echo " make cpu Build CPU-only ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent"
|
||||
@echo " make test Build and run tests"
|
||||
@echo " make clean Remove build outputs"
|
||||
|
||||
ds4: ds4_cli.o ds4_help.o linenoise.o $(CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_cli.o ds4_help.o linenoise.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
|
||||
ds4-server: ds4_server.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_server.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
|
||||
ds4-bench: ds4_bench.o ds4_help.o $(CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_bench.o ds4_help.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
|
||||
ds4-eval: ds4_eval.o ds4_help.o $(CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_eval.o ds4_help.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
|
||||
ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
|
||||
cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o rax.o $(CPU_CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-eval ds4_eval_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-agent ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
|
||||
cuda-regression:
|
||||
@echo "cuda-regression requires a CUDA build"
|
||||
else
|
||||
all: help
|
||||
|
||||
help:
|
||||
@echo "DS4 build targets:"
|
||||
@echo " make cuda-spark Build CUDA for DGX Spark / GB10"
|
||||
@echo " make cuda-generic Build CUDA for a generic local CUDA GPU"
|
||||
@echo " make cuda CUDA_ARCH=sm_N Build CUDA with an explicit nvcc -arch value"
|
||||
@echo " make strix-halo Build ROCm for Strix Halo / gfx1151"
|
||||
@echo " make rocm Alias for make strix-halo"
|
||||
@echo " make cpu Build CPU-only ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent"
|
||||
@echo " make test Build and run tests"
|
||||
@echo " make clean Remove build outputs"
|
||||
|
||||
cuda-spark:
|
||||
$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=
|
||||
|
||||
cuda-generic:
|
||||
$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=native
|
||||
|
||||
cuda:
|
||||
@if [ -z "$(strip $(CUDA_ARCH))" ]; then \
|
||||
echo "error: specify CUDA_ARCH, for example: make cuda CUDA_ARCH=sm_120"; \
|
||||
echo " or use make cuda-spark / make cuda-generic"; \
|
||||
exit 2; \
|
||||
fi
|
||||
$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH="$(CUDA_ARCH)"
|
||||
|
||||
strix-halo:
|
||||
$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent \
|
||||
CORE_OBJS="ds4.o ds4_distributed.o ds4_ssd.o ds4_rocm.o" \
|
||||
CFLAGS="$(CFLAGS) -DDS4_ROCM_BUILD" \
|
||||
DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \
|
||||
DS4_LINK_LIBS="$(ROCM_LDLIBS)"
|
||||
|
||||
rocm: strix-halo
|
||||
|
||||
ds4: ds4_cli.o ds4_help.o linenoise.o $(CORE_OBJS)
|
||||
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
|
||||
|
||||
ds4-server: ds4_server.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS)
|
||||
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
|
||||
|
||||
ds4-bench: ds4_bench.o ds4_help.o $(CORE_OBJS)
|
||||
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
|
||||
|
||||
ds4-eval: ds4_eval.o ds4_help.o $(CORE_OBJS)
|
||||
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
|
||||
|
||||
ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS)
|
||||
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
|
||||
|
||||
cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o rax.o $(CPU_CORE_OBJS)
|
||||
$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-eval ds4_eval_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
$(CC) $(CFLAGS) -o ds4-agent ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
|
||||
|
||||
cuda-regression: tests/cuda_long_context_smoke
|
||||
./tests/cuda_long_context_smoke
|
||||
endif
|
||||
|
||||
ds4.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4.c
|
||||
|
||||
ds4_ssd.o: ds4_ssd.c ds4_ssd.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_ssd.c
|
||||
|
||||
ds4_cli.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_cli.c
|
||||
|
||||
ds4_distributed.o: ds4_distributed.c ds4_distributed.h ds4.h ds4_ssd.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_distributed.c
|
||||
|
||||
ds4_help.o: ds4_help.c ds4_help.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_help.c
|
||||
|
||||
ds4_server.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_server.c
|
||||
|
||||
ds4_bench.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_bench.c
|
||||
|
||||
ds4_eval.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_eval.c
|
||||
|
||||
ds4_agent.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_agent.c
|
||||
|
||||
ds4_web.o: ds4_web.c ds4_web.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_web.c
|
||||
|
||||
ds4_kvstore.o: ds4_kvstore.c ds4_kvstore.h ds4.h ds4_ssd.h
|
||||
$(CC) $(CFLAGS) -c -o $@ ds4_kvstore.c
|
||||
|
||||
ds4_test.o: tests/ds4_test.c ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h
|
||||
$(CC) $(CFLAGS) -Wno-unused-function -c -o $@ tests/ds4_test.c
|
||||
|
||||
ds4_agent_test.o: tests/ds4_agent_test.c ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h
|
||||
$(CC) $(CFLAGS) -Wno-unused-function -c -o $@ tests/ds4_agent_test.c
|
||||
|
||||
tests/cuda_long_context_smoke.o: tests/cuda_long_context_smoke.c ds4_gpu.h
|
||||
$(CC) $(CFLAGS) -I. -c -o $@ tests/cuda_long_context_smoke.c
|
||||
|
||||
rax.o: rax.c rax.h rax_malloc.h
|
||||
$(CC) $(CFLAGS) -c -o $@ rax.c
|
||||
|
||||
linenoise.o: linenoise.c linenoise.h
|
||||
$(CC) $(CFLAGS) -c -o $@ linenoise.c
|
||||
|
||||
ds4_cpu.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h
|
||||
$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4.c
|
||||
|
||||
ds4_cli_cpu.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h
|
||||
$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_cli.c
|
||||
|
||||
ds4_server_cpu.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h
|
||||
$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_server.c
|
||||
|
||||
ds4_bench_cpu.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
|
||||
$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_bench.c
|
||||
|
||||
ds4_eval_cpu.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
|
||||
$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_eval.c
|
||||
|
||||
ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h
|
||||
$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_agent.c
|
||||
|
||||
ds4_metal.o: ds4_metal.m ds4_gpu.h $(METAL_SRCS)
|
||||
$(CC) $(OBJCFLAGS) -c -o $@ ds4_metal.m
|
||||
|
||||
ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_iq2_tables_cuda.inc
|
||||
$(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu
|
||||
|
||||
ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS)
|
||||
$(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm.cu
|
||||
|
||||
tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o
|
||||
$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
|
||||
|
||||
ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS)
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
else
|
||||
$(NVCC) $(NVCCFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(CUDA_LDLIBS)
|
||||
endif
|
||||
|
||||
ds4_agent_test: ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS)
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
$(CC) $(CFLAGS) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(METAL_LDLIBS)
|
||||
else
|
||||
$(NVCC) $(NVCCFLAGS) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(CUDA_LDLIBS)
|
||||
endif
|
||||
|
||||
test: ds4_test ds4_agent_test ds4-eval q4k-dot-test
|
||||
./ds4-eval --self-test-extractors
|
||||
./ds4_agent_test
|
||||
./ds4_test
|
||||
|
||||
q4k-dot-test: tests/test_q4k_dot.c
|
||||
$(CC) -O2 -Wall -Wextra -std=c99 -o tests/test_q4k_dot tests/test_q4k_dot.c -lm -pthread
|
||||
./tests/test_q4k_dot
|
||||
|
||||
clean:
|
||||
rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test tests/test_q4k_dot *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o
|
||||
@@ -0,0 +1,253 @@
|
||||
# QA Before Releases
|
||||
|
||||
This is the release gate for DwarfStar. Run it before tagging or pushing a
|
||||
release build. The goal is not to prove every code path exhaustively; it is to
|
||||
exercise the paths that have historically regressed: Metal graph inference,
|
||||
CUDA, ROCm, SSD streaming, distributed execution, disk KV cache, server APIs, and the
|
||||
agent TUI/tool state machine.
|
||||
|
||||
Do not run multiple huge model processes at the same time. Record the commit,
|
||||
hardware, GGUF file, context size, and any non-default flags for every manual
|
||||
run.
|
||||
|
||||
Preferred release test hosts:
|
||||
|
||||
- CUDA / DGX Spark: `toor@192.168.0.180`.
|
||||
- Metal / distributed Mac testing: `mac-m5max-it` and `mac-m5max-us`.
|
||||
- ROCm: The Strix Halo system at antirez@strixhalo (Framework Desktop).
|
||||
|
||||
The Mac hosts have DNS entries and are reached through an internet VPN. They
|
||||
are connected to each other over WiFi and also through a Thunderbolt 5
|
||||
point-to-point link. The TB5 route is the preferred distributed-inference
|
||||
network when it is available, but it can be fragile and sometimes only works
|
||||
when `ds4` is executed in the foreground. Prefer these machines for release
|
||||
testing, especially distributed inference. Local fallback testing on this
|
||||
machine is acceptable when needed; it is an M3 Max with 128 GB RAM.
|
||||
The Strix Halo system is reachable via the VPN as well and has a local WiFi
|
||||
address in the same lan of the M5 Max systems. The CUDA hosts are in a
|
||||
different remote lan and are accessible via a different VPN active
|
||||
in this system.
|
||||
|
||||
## 1. Repository And Build Sanity
|
||||
|
||||
- Start from a clean tree except intentional release notes:
|
||||
`git status --short`.
|
||||
- Build the normal local target:
|
||||
`make clean && make`.
|
||||
- Build CPU-only binaries as a compile check only:
|
||||
`make clean && make cpu`.
|
||||
- Run whitespace checks before committing:
|
||||
`git diff --check`.
|
||||
- Confirm `./ds4 --help`, `./ds4-server --help`, and `./ds4-agent --help` render
|
||||
cleanly, with readable section colors and no broken wrapping.
|
||||
|
||||
## 2. Core Regression Tests
|
||||
|
||||
- Run the default suite:
|
||||
`make test`.
|
||||
- Run the vector checks explicitly after any tokenizer, template, KV, kernel,
|
||||
quantization, or prompt-rendering change:
|
||||
`./ds4_test --logprob-vectors`
|
||||
and `./ds4_test --local-golden-vectors`.
|
||||
- Run server tests when HTTP, SSE, prompt rendering, cache policy, or tool-call
|
||||
replay changed:
|
||||
`./ds4_test --server`.
|
||||
- Run `./ds4-eval --self-test-extractors`.
|
||||
|
||||
## 3. Metal Flash Path
|
||||
|
||||
Use the normal Flash GGUF that 128 GB users run.
|
||||
|
||||
- One-shot CLI:
|
||||
`./ds4 -m ds4flash.gguf --ctx 32768 --nothink -p "Explain C pointers in one paragraph."`
|
||||
- Thinking and max-thinking prompts:
|
||||
run one short coding prompt with default thinking and one with max thinking.
|
||||
- Long-context recall:
|
||||
run the long name/number or archive recall test used for catching attention
|
||||
and MoE routing drift.
|
||||
- Logprob sanity:
|
||||
`./ds4 --nothink --temp 0 --dump-logprobs /tmp/ds4-logprobs.json --logprobs-top-k 20 -p "..."`
|
||||
and inspect that the continuation is sane.
|
||||
- Speed sanity:
|
||||
run `ds4-bench` with `speed-bench/promessi_sposi.txt` and compare prefill,
|
||||
generation speed, and KV bytes with the last known good numbers for the same
|
||||
machine.
|
||||
|
||||
## 4. Metal PRO Path
|
||||
|
||||
PRO support is experimental, but release builds must not break it silently.
|
||||
|
||||
- If a PRO-capable machine is available, run a short PRO q2 prompt and verify
|
||||
the correct template, thinking behavior, and endpoint aliases.
|
||||
- For PRO Q4 distributed builds, test only on the intended high-memory machines.
|
||||
- If PRO cannot be run locally, at least build all binaries and review changes
|
||||
touching model shape, tensor lookup, routed expert mapping, template logic,
|
||||
and KV payload compatibility.
|
||||
|
||||
## 5. SSD Streaming
|
||||
|
||||
SSD streaming is a capacity path, so test both correctness and user experience.
|
||||
|
||||
- Flash q2/q2-q4 streaming:
|
||||
`./ds4 -m ds4flash.gguf --ssd-streaming --ssd-streaming-cache-experts 32GB -p "..."`
|
||||
- Regression test mixed-quant Flash SSD streaming. Use the mixed q2/q4 GGUF
|
||||
with boosted Q4 routed-expert layers and a prompt long enough to exercise the
|
||||
selected-address prefill path; it must not fail with "model range is not
|
||||
covered by mapped model views":
|
||||
`./ds4 -m gguf/DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf --ssd-streaming --ssd-streaming-cache-experts 16GB --ctx 4096 --tokens 1 --nothink --prompt-file /tmp/ds4_600tok_prompt.txt`.
|
||||
- Cold streaming measurement:
|
||||
run once with `--ssd-streaming-cold` and verify no deadlock, missing expert,
|
||||
or impossible slowdown.
|
||||
- Confirm startup reports cache budget and that generation does not stall on
|
||||
repeated expert misses for a small interactive prompt.
|
||||
- If streaming cache internals changed, test the same prompt twice and compare
|
||||
first-token/logprob sanity between runs.
|
||||
|
||||
## 6. CUDA / DGX Spark
|
||||
|
||||
Before a release, ask the user for CUDA access if it is not already configured.
|
||||
Use the DGX Spark / GB10 host `toor@192.168.0.180`. Do not claim CUDA is
|
||||
release-ready without this pass.
|
||||
|
||||
- Fetch or push the exact release commit to the CUDA machine.
|
||||
- Build:
|
||||
`make clean && make cuda-spark`.
|
||||
- Run:
|
||||
`make cuda-regression`.
|
||||
- Run a short CLI prompt with the Flash GGUF and record generation t/s.
|
||||
- Run a longer prompt that exercises routed experts past a few thousand tokens.
|
||||
- If CUDA Q4, distributed, streaming hooks, tensor span loading, or model cache
|
||||
code changed, test the specific GGUF and split mode that uses that path.
|
||||
- Verify that any CUDA-only warning fixes are also clean on macOS and do not
|
||||
change Metal behavior.
|
||||
|
||||
## 7. ROCm / Strix Halo
|
||||
|
||||
Use the Strix Halo Framework Desktop via the VPN hostname `strixhalo`
|
||||
(`antirez@strixhalo`). This host validates the ROCm backend; do not use it as
|
||||
a substitute for CUDA or Metal release testing.
|
||||
|
||||
- Fetch or push the exact release commit to the Strix Halo machine.
|
||||
- Build:
|
||||
`make clean && make strix-halo`.
|
||||
- Use the q2 Flash imatrix GGUF for release smoke tests:
|
||||
`DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf`.
|
||||
- Do not use the mixed q2-q4 or Q4 Flash GGUFs for routine Strix Halo QA yet.
|
||||
They are dangerous on this machine for now because the ROCm path can hit
|
||||
system OOM instead of failing cleanly.
|
||||
- Run a short CLI prompt:
|
||||
`./ds4 -m gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf --ctx 4096 --nothink -p "Reply with exactly: OK"`.
|
||||
- Run one longer prompt if ROCm kernels, backend hooks, tensor loading, model
|
||||
cache, KV cache, or graph prefill code changed.
|
||||
- Record startup memory/cache messages, prefill speed, generation speed, and
|
||||
whether the backend reports `ROCm backend initialized`.
|
||||
|
||||
## 8. Distributed Inference
|
||||
|
||||
Distributed code has regressed around route setup, KV snapshots, request IDs,
|
||||
and split model loading. Test it whenever distributed, KV, session, or model
|
||||
loading code changes.
|
||||
|
||||
- Prefer `mac-m5max-it` and `mac-m5max-us` for Metal distributed tests. Use the
|
||||
TB5 point-to-point link when it is working; otherwise note that the run used
|
||||
WiFi/VPN routing.
|
||||
- Start workers first, then the coordinator.
|
||||
- Test a small prompt and a longer prompt.
|
||||
- Verify the coordinator waits for a complete route and exits cleanly.
|
||||
- Verify `Ctrl+C` returns control after the current distributed token or chunk
|
||||
drains.
|
||||
- Save and restore a distributed KV snapshot if that code changed.
|
||||
- If CUDA distributed is relevant, test across the CUDA hosts and record
|
||||
generation speed, not just "it works".
|
||||
|
||||
## 9. Disk KV Cache
|
||||
|
||||
Disk KV cache bugs are high impact for server users.
|
||||
|
||||
- Start the server with:
|
||||
`./ds4-server --ctx 100000 --kv-disk-dir /tmp/ds4-kv --kv-disk-space-mb 8192`.
|
||||
- Run the same request twice and verify the second request hits cache.
|
||||
- Fill the cache enough to trigger eviction; verify the newly-written entry is
|
||||
not evicted and useful anchors are retained.
|
||||
- Test rejection of incompatible checkpoints when model, quantization, context,
|
||||
or raw/compressed KV layout changes.
|
||||
- Test stripped agent sessions: `/strip <id>` then `/switch <id>` should rebuild
|
||||
by prefill and render sane history.
|
||||
|
||||
## 10. Server APIs
|
||||
|
||||
The server must keep compatibility across OpenAI, Responses, and Anthropic
|
||||
clients.
|
||||
|
||||
- `GET /v1/models/deepseek-v4-flash` and `GET /v1/models/deepseek-v4-pro`
|
||||
should both serve whichever GGUF is loaded.
|
||||
- Test OpenAI chat completion, OpenAI Responses, and Anthropic messages.
|
||||
- Test SSE streaming with thinking enabled and disabled.
|
||||
- Test keepalive during long prefill and confirm clients do not time out.
|
||||
- Test `--trace` and confirm rendered prompts, cache decisions, generated text,
|
||||
and tool-parser events are useful without leaking unrelated state.
|
||||
|
||||
## 11. ds4-agent
|
||||
|
||||
The agent is the most stateful component. Test it manually, not only by build.
|
||||
|
||||
- Startup banner, status bar, help, `/power`, `/save`, `/list`, `/switch`,
|
||||
`/history`, `/compact`, `/new`, `/del`, and `/strip`.
|
||||
- Ctrl+C during generation, during prefill, during a web fetch, and during a
|
||||
long tool call. After `Stopped by user`, typing a new prompt must work.
|
||||
- Queue messages while the model is busy. Queued messages must not skip tool
|
||||
execution; after tool results, the queued user text must be provided.
|
||||
- Read/search/edit/write tools:
|
||||
create a temp project, ask for edits, verify old/new and `[upto]` anchored
|
||||
edits fail safely on ambiguous matches and do not require retyping whole files.
|
||||
- Real coding edit loop:
|
||||
delete `/tmp/mymandel`, ask ds4-agent to create a small C ASCII Mandelbrot
|
||||
program there, build and run it, then in a second user turn ask for a small
|
||||
modification that should naturally use the edit tool, such as changing the
|
||||
ASCII character ramp or output dimensions. Verify the agent edits the
|
||||
existing file instead of rewriting the whole project, and that the final
|
||||
program still builds and runs.
|
||||
- Bash tools:
|
||||
test short output, large output truncation, non-zero exit output, long-running
|
||||
jobs, `bash_status`, and `bash_stop`.
|
||||
- Web tools:
|
||||
`google_search` and `visit_page` should ask for visible Chrome approval with
|
||||
timeout, open pages without stealing focus when possible, extract Markdown,
|
||||
close tabs, and handle consent/privacy walls as tool errors the model can see.
|
||||
- TUI:
|
||||
test multiline prompt editing, history navigation, queued prompt display,
|
||||
status bar fill to terminal width, syntax highlighting in Markdown/code blocks,
|
||||
and SSH/remote terminal flicker.
|
||||
|
||||
## 12. Download Script And Model Files
|
||||
|
||||
- Test `download_model.sh` in a temporary directory so local weights are not
|
||||
overwritten.
|
||||
- Test one Flash target and one PRO target enough to verify URL, resume, Hugging
|
||||
Face CLI/curl behavior, file naming, and symlink policy.
|
||||
- Verify legacy removed targets fail clearly.
|
||||
- Verify README model names match the script and Hugging Face repository.
|
||||
|
||||
## 13. Performance And Power
|
||||
|
||||
- Run `ds4-bench` on the release machine and compare with tracked CSV baselines.
|
||||
- Test `--power 100` is not throttled.
|
||||
- Test `--power 50` visibly reduces duty cycle in CLI, server, agent, eval, and
|
||||
bench where practical.
|
||||
- Confirm context buffer size, raw KV rows, compressed KV rows, and mmap behavior
|
||||
match expectations for 32k, 100k, and any release-advertised context size.
|
||||
|
||||
## 14. Release Sign-off
|
||||
|
||||
Do not sign off until:
|
||||
|
||||
- macOS Metal Flash passed.
|
||||
- CUDA was tested on the CUDA machine or the release notes explicitly say CUDA
|
||||
was not validated.
|
||||
- ROCm was tested on Strix Halo or the release notes explicitly say ROCm was
|
||||
not validated.
|
||||
- Disk KV cache was exercised.
|
||||
- Server API streaming was exercised.
|
||||
- Agent interruption and tool loops were exercised manually.
|
||||
- Speed is within expected variance for the same hardware and model.
|
||||
- Any skipped item is written down with the reason.
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`antirez/ds4`
|
||||
- 原始仓库：https://github.com/antirez/ds4
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
+136
@@ -0,0 +1,136 @@
|
||||
# DS4 on Strix Halo
|
||||
|
||||
This is the minimal setup for DS4 ROCm inference on a
|
||||
Strix Halo machine with 128 GB RAM and Radeon 8060S (`gfx1151`).
|
||||
|
||||
## 1. Install ROCm
|
||||
|
||||
On Ubuntu 26.04 LTS, install the ROCm compiler/runtime and libraries used by the Strix Halo backend:
|
||||
|
||||
```sh
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y \
|
||||
hipcc rocminfo rocm-smi \
|
||||
libamdhip64-dev \
|
||||
libhipblas-dev libhipblaslt-dev \
|
||||
librocblas-dev \
|
||||
librocwmma-dev \
|
||||
libhipcub-dev
|
||||
```
|
||||
|
||||
The backend uses rocWMMA. On this Ubuntu 26.04 setup, `librocwmma-dev`
|
||||
installs the top-level rocWMMA headers but misses `rocwmma/internal/`.
|
||||
No Ubuntu package currently provides those internal headers. Install a complete
|
||||
matching rocWMMA header tree:
|
||||
|
||||
```sh
|
||||
git clone --depth 1 --branch rocm-7.1.0 https://github.com/ROCm/rocWMMA.git /tmp/rocWMMA-rocm-7.1.0
|
||||
sudo mkdir -p /usr/local/include
|
||||
sudo cp -a /tmp/rocWMMA-rocm-7.1.0/library/include/rocwmma /usr/local/include/
|
||||
```
|
||||
|
||||
If ROCm is installed under `/usr` but tooling expects `/opt/rocm`, add these
|
||||
compatibility links:
|
||||
|
||||
```sh
|
||||
sudo mkdir -p /opt/rocm/bin
|
||||
sudo ln -sf /usr/bin/hipcc /opt/rocm/bin/hipcc
|
||||
sudo ln -sfn /usr/lib/x86_64-linux-gnu /opt/rocm/lib
|
||||
sudo ln -sfn /usr/include /opt/rocm/include
|
||||
```
|
||||
|
||||
## 2. Enable ROCm access
|
||||
|
||||
The user running DS4 must be able to open `/dev/kfd` and the DRM render node:
|
||||
|
||||
```sh
|
||||
sudo usermod -aG render,video "$USER"
|
||||
```
|
||||
|
||||
Log out and back in, or reboot. Verify:
|
||||
|
||||
```sh
|
||||
rocminfo | grep -A80 'Name: gfx1151'
|
||||
```
|
||||
|
||||
If DS4 says `no ROCm-capable device is detected`, check that `rocminfo` can open
|
||||
`/dev/kfd` and that `groups` includes `render`.
|
||||
|
||||
## 3. Increase GPU-visible memory
|
||||
|
||||
A 128 GB Strix Halo system may initially expose only about 62 GB of GPU-visible
|
||||
memory. DS4 needs the larger GTT aperture for the 80.76 GiB model plus runtime
|
||||
buffers.
|
||||
|
||||
Use these kernel parameters:
|
||||
|
||||
```text
|
||||
amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856
|
||||
```
|
||||
|
||||
On Ubuntu with GRUB:
|
||||
|
||||
```sh
|
||||
sudo cp /etc/default/grub /etc/default/grub.bak
|
||||
sudoedit /etc/default/grub
|
||||
```
|
||||
|
||||
Set:
|
||||
|
||||
```text
|
||||
GRUB_CMDLINE_LINUX_DEFAULT="quiet splash amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856"
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```sh
|
||||
sudo update-grub
|
||||
sudo reboot
|
||||
```
|
||||
|
||||
After reboot, verify:
|
||||
|
||||
```sh
|
||||
cat /proc/cmdline
|
||||
sudo dmesg | grep -Ei 'GTT|gttsize|TTM|VRAM'
|
||||
rocminfo | grep -A80 'Name: gfx1151'
|
||||
```
|
||||
|
||||
Expected signs:
|
||||
|
||||
```text
|
||||
amdgpu: 126976M of GTT memory ready
|
||||
rocminfo gfx1151 pool: 130023424 KB
|
||||
```
|
||||
|
||||
## 4. Build DS4
|
||||
|
||||
Use the normal Strix Halo target. It builds the standard binary names:
|
||||
|
||||
```sh
|
||||
make strix-halo -j"$(nproc)"
|
||||
```
|
||||
|
||||
`make rocm` is an alias for `make strix-halo`.
|
||||
|
||||
## 5. Use the right GGUF
|
||||
|
||||
Use the standard IQ2XXS/Q2K/Q8 imatrix GGUF:
|
||||
|
||||
```text
|
||||
DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf
|
||||
```
|
||||
|
||||
Avoid the mixed IQ2/IQ4 or IQ2/Q4 GGUFs on this machine for now. They put much
|
||||
more memory pressure on the ROCm path and can trigger system OOM instead of a
|
||||
clean DS4 failure.
|
||||
|
||||
## 6. Run DS4
|
||||
|
||||
Run it normally:
|
||||
|
||||
```sh
|
||||
./ds4 -m gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf
|
||||
```
|
||||
|
||||
The ROCm build uses the Strix Halo backend automatically.
|
||||
@@ -0,0 +1,3 @@
|
||||
out/
|
||||
*.pyc
|
||||
__pycache__/
|
||||
@@ -0,0 +1,151 @@
|
||||
# Directional Steering
|
||||
|
||||
Directional steering is a runtime activation edit for DS4. A steering file is a
|
||||
flat `f32` matrix with one normalized 4096-wide direction per layer. During
|
||||
inference, ds4 can apply the edit after attention outputs, FFN outputs, or both:
|
||||
|
||||
```text
|
||||
y = y - scale * direction[layer] * dot(direction[layer], y)
|
||||
```
|
||||
|
||||
Positive scale removes the represented direction. Negative scale amplifies it.
|
||||
With no steering file or zero scales, ds4 follows the normal inference path.
|
||||
|
||||
## Runtime Options
|
||||
|
||||
```text
|
||||
--dir-steering-file FILE load a 43 x 4096 f32 direction file
|
||||
--dir-steering-ffn F apply steering after FFN outputs; default is 1 when a file is provided
|
||||
--dir-steering-attn F apply steering after attention outputs; default is 0
|
||||
```
|
||||
|
||||
The FFN output is usually the best first target because it is late enough in
|
||||
each layer to represent behavior, style, and topic signals. Attention steering
|
||||
is available for experiments, but it can be more fragile.
|
||||
|
||||
## Verbosity Example
|
||||
|
||||
The bundled example builds a style direction from 100 paired prompts. Each pair
|
||||
asks for the same information in two ways:
|
||||
|
||||
- `examples/succinct.txt`: terse target prompts.
|
||||
- `examples/verbose.txt`: detailed contrast prompts.
|
||||
|
||||
Because the extracted direction is `succinct - verbose`, negative FFN scales
|
||||
make answers shorter, while positive FFN scales tend to make answers longer and
|
||||
more explanatory.
|
||||
|
||||
Build the vector:
|
||||
|
||||
```sh
|
||||
python3 dir-steering/tools/build_direction.py \
|
||||
--ds4 ./ds4 \
|
||||
--model ds4flash.gguf \
|
||||
--good-file dir-steering/examples/succinct.txt \
|
||||
--bad-file dir-steering/examples/verbose.txt \
|
||||
--out dir-steering/out/verbosity.json \
|
||||
--component ffn_out \
|
||||
--ctx 512
|
||||
```
|
||||
|
||||
This writes:
|
||||
|
||||
```text
|
||||
dir-steering/out/verbosity.json
|
||||
dir-steering/out/verbosity.f32
|
||||
```
|
||||
|
||||
Try a terse run:
|
||||
|
||||
```sh
|
||||
./ds4 -m ds4flash.gguf --nothink --temp 0 -n 160 \
|
||||
--dir-steering-file dir-steering/out/verbosity.f32 \
|
||||
--dir-steering-ffn -1 \
|
||||
-p "Explain why databases use indexes."
|
||||
```
|
||||
|
||||
Try a verbose run:
|
||||
|
||||
```sh
|
||||
./ds4 -m ds4flash.gguf --nothink --temp 0 -n 220 \
|
||||
--dir-steering-file dir-steering/out/verbosity.f32 \
|
||||
--dir-steering-ffn 2 \
|
||||
-p "Explain why databases use indexes."
|
||||
```
|
||||
|
||||
The same vector can be used in either direction. The sign is the important part:
|
||||
|
||||
- negative scale amplifies the succinct target direction;
|
||||
- positive scale suppresses that direction and usually gives the model more room
|
||||
to elaborate.
|
||||
|
||||
## Evaluating Scales
|
||||
|
||||
Use the sweep helper to test several strengths on a fixed prompt set:
|
||||
|
||||
```sh
|
||||
python3 dir-steering/tools/run_sweep.py \
|
||||
--ds4 ./ds4 \
|
||||
--model ds4flash.gguf \
|
||||
--direction dir-steering/out/verbosity.f32 \
|
||||
--prompts dir-steering/examples/eval_prompts.txt \
|
||||
--scales "-1,-0.5,0,0.5,1,2" \
|
||||
--tokens 180 \
|
||||
--nothink
|
||||
```
|
||||
|
||||
Start with FFN scales between `-1` and `2`. If the model becomes repetitive,
|
||||
ignores the prompt, or starts losing factual content, the scale is too strong.
|
||||
For this example, `-1` is a good first terse setting and `2` is a good first
|
||||
verbose setting. Strong negative scales such as `-2` or `-3` can over-amplify
|
||||
the terse direction and collapse into repetition on some prompts.
|
||||
|
||||
## Observed Effect
|
||||
|
||||
With the 100-pair vector built from the commands above, local greedy checks
|
||||
showed the expected behavior:
|
||||
|
||||
- Prompt: `Explain why databases use indexes.`
|
||||
- `--dir-steering-ffn -1`: 67 words, one compact paragraph.
|
||||
- `--dir-steering-ffn 0`: 136 words, structured explanation.
|
||||
- `--dir-steering-ffn 1`: 140 words, structured explanation with more detail.
|
||||
|
||||
On a prompt that the unsteered model already answered briefly, positive steering
|
||||
made the expansion more visible:
|
||||
|
||||
- Prompt: `What does DNS do?`
|
||||
- `--dir-steering-ffn 0`: 44 words.
|
||||
- `--dir-steering-ffn 2`: 171 words, with sections and step-by-step detail.
|
||||
|
||||
## Building Other Directions
|
||||
|
||||
The extractor compares two prompt sets:
|
||||
|
||||
- `good-file`: target prompts for the direction you want to represent.
|
||||
- `bad-file`: contrast prompts that should be separated from the target.
|
||||
|
||||
It captures DS4 activations from the same local GPU graph used for inference,
|
||||
averages target minus contrast, normalizes one vector per layer, and writes both
|
||||
metadata JSON and the runtime `.f32` file.
|
||||
|
||||
Concept removal:
|
||||
|
||||
1. Put concept-heavy prompts in `good-file`.
|
||||
2. Put neutral prompts in `bad-file`.
|
||||
3. Run with a positive FFN scale.
|
||||
|
||||
Concept amplification:
|
||||
|
||||
1. Put desired concept prompts in `good-file`.
|
||||
2. Put neutral prompts in `bad-file`.
|
||||
3. Run with a negative FFN scale.
|
||||
|
||||
Style control:
|
||||
|
||||
1. Put prompts for the target style in `good-file`.
|
||||
2. Put contrasting style prompts in `bad-file`.
|
||||
3. Use negative scale to amplify the target style, positive scale to reduce it.
|
||||
|
||||
The method is not a fine-tune. It is a low-rank runtime edit, so it works best
|
||||
for coarse behavior, topic, or style directions that are consistently present in
|
||||
the activation captures.
|
||||
@@ -0,0 +1,5 @@
|
||||
Explain why databases use indexes.
|
||||
How should I learn sorting algorithms?
|
||||
Describe how clouds form.
|
||||
Give me advice for debugging a slow program.
|
||||
What is public key cryptography?
|
||||
@@ -0,0 +1,100 @@
|
||||
Answer in one short sentence: what is binary search?
|
||||
In two concise bullets, explain why databases use indexes.
|
||||
Give a terse definition of TCP congestion control.
|
||||
Summarize recursion in exactly one sentence.
|
||||
Briefly explain what a cache does.
|
||||
Give a minimal explanation of how rain forms.
|
||||
In one compact paragraph, describe photosynthesis.
|
||||
Answer briefly: why should passwords be unique?
|
||||
Give a concise explanation of unit tests.
|
||||
In one sentence, explain what DNS does.
|
||||
Briefly explain why exercise improves health.
|
||||
Give a short answer: what is a hash table?
|
||||
Explain gradient descent in two crisp sentences.
|
||||
Define inflation in one sentence.
|
||||
Give a terse summary of how vaccines work.
|
||||
Answer briefly: what is an API?
|
||||
In one short paragraph, explain how compilers work.
|
||||
Give a compact explanation of merge sort.
|
||||
Briefly explain why backups matter.
|
||||
Define latency in one concise sentence.
|
||||
Give a minimal explanation of public key cryptography.
|
||||
In two short bullets, explain what a load balancer does.
|
||||
Briefly describe how solar panels generate electricity.
|
||||
Give a one-sentence explanation of supply and demand.
|
||||
Answer tersely: what is containerization?
|
||||
In one compact paragraph, explain the water cycle.
|
||||
Give a short definition of entropy.
|
||||
Briefly explain how a recommendation system works.
|
||||
Give a concise answer: why do teams use version control?
|
||||
Explain neural networks in two short sentences.
|
||||
Answer in one sentence: what is a database transaction?
|
||||
Give a terse explanation of event loops.
|
||||
Briefly explain how airplanes stay aloft.
|
||||
Give a compact summary of why sleep matters.
|
||||
In two concise bullets, explain what OAuth is for.
|
||||
Define bandwidth in one short sentence.
|
||||
Answer briefly: how does a thermostat work?
|
||||
Give a minimal explanation of memoization.
|
||||
Briefly explain why code reviews are useful.
|
||||
In one short paragraph, describe how markets set prices.
|
||||
Give a terse answer: what is a microservice?
|
||||
Explain backpropagation in two compact sentences.
|
||||
Briefly describe how a search engine ranks pages.
|
||||
Give a one-sentence explanation of opportunity cost.
|
||||
Answer concisely: what is rate limiting?
|
||||
In two short bullets, explain why logging is useful.
|
||||
Give a compact explanation of climate change.
|
||||
Briefly explain how checksums detect errors.
|
||||
Define normalization in databases in one sentence.
|
||||
Give a terse explanation of garbage collection.
|
||||
Answer briefly: why does encryption need keys?
|
||||
In one compact paragraph, explain how GPS finds location.
|
||||
Give a short explanation of continuous integration.
|
||||
Briefly describe how a queue differs from a stack.
|
||||
Define overfitting in one concise sentence.
|
||||
Give a minimal explanation of DNS caching.
|
||||
Answer tersely: what is a kernel?
|
||||
In two concise bullets, explain how to prioritize tasks.
|
||||
Briefly explain why indexes can slow writes.
|
||||
Give a compact explanation of HTTP.
|
||||
In one sentence, explain what a proxy server does.
|
||||
Answer briefly: how does compression save space?
|
||||
Give a terse explanation of database replication.
|
||||
Briefly explain how a compiler differs from an interpreter.
|
||||
Define amortized cost in one short sentence.
|
||||
Give a minimal explanation of a Bloom filter.
|
||||
Answer concisely: why do batteries degrade?
|
||||
In one compact paragraph, explain how elections use runoff voting.
|
||||
Give a short explanation of semantic versioning.
|
||||
Briefly describe how a CDN helps a website.
|
||||
Define idempotency in one concise sentence.
|
||||
Answer tersely: what is a mutex?
|
||||
In two short bullets, explain why monitoring matters.
|
||||
Give a compact explanation of how email routing works.
|
||||
Briefly explain how a neural embedding represents text.
|
||||
Give a one-sentence explanation of dependency injection.
|
||||
Answer briefly: what is a race condition?
|
||||
In one compact paragraph, explain why documentation helps teams.
|
||||
Give a terse description of how a blockchain works.
|
||||
Briefly explain what a memory leak is.
|
||||
Define consensus in distributed systems in one sentence.
|
||||
Give a minimal explanation of a priority queue.
|
||||
Answer concisely: why is caching hard to invalidate?
|
||||
In two concise bullets, explain how to debug a slow program.
|
||||
Briefly explain what a firewall does.
|
||||
Give a compact explanation of how interest rates affect borrowing.
|
||||
In one sentence, explain what an operating system does.
|
||||
Answer briefly: how does a web browser render a page?
|
||||
Give a terse explanation of feature flags.
|
||||
Briefly explain how garbage collection pauses happen.
|
||||
Define eventual consistency in one concise sentence.
|
||||
Give a minimal explanation of sharding.
|
||||
Answer concisely: what is a deadlock?
|
||||
In one compact paragraph, explain how a camera sensor captures light.
|
||||
Give a short explanation of why indexes need selectivity.
|
||||
Briefly describe what a scheduler does.
|
||||
Define precision and recall in one sentence.
|
||||
Give a terse explanation of streaming APIs.
|
||||
Answer briefly: why do projects need tests?
|
||||
In two short bullets, explain how to choose a data structure.
|
||||
@@ -0,0 +1,100 @@
|
||||
Write a detailed, multi-paragraph explanation of what binary search is, including intuition, preconditions, and a small example.
|
||||
Explain why databases use indexes in a thorough answer with several concrete tradeoffs and examples.
|
||||
Give an expansive explanation of TCP congestion control, covering its goal, feedback signals, and why it changes speed over time.
|
||||
Explain recursion carefully with a step-by-step description, a concrete analogy, and a note about base cases.
|
||||
Describe caches in detail, including what they store, why they help, and what can go wrong.
|
||||
Explain how rain forms in a rich answer covering evaporation, condensation, clouds, and falling droplets.
|
||||
Describe photosynthesis in depth, including inputs, outputs, energy conversion, and why plants depend on it.
|
||||
Explain why passwords should be unique with a careful discussion of breaches, reuse, and practical safety habits.
|
||||
Explain unit tests in detail, including what they verify, how they fit into development, and their limitations.
|
||||
Describe DNS thoroughly, including names, resolvers, records, caching, and why it matters for the web.
|
||||
Explain why exercise improves health in a detailed answer covering heart, muscles, mood, sleep, and long-term effects.
|
||||
Give a thorough explanation of hash tables, including hashing, buckets, collisions, and average-case performance.
|
||||
Explain gradient descent in depth, including the loss function, gradients, step size, and how iteration improves a model.
|
||||
Describe inflation carefully, including price levels, purchasing power, possible causes, and everyday consequences.
|
||||
Explain how vaccines work in detail, including immune memory, antigens, protection, and population effects.
|
||||
Describe what an API is with examples, explaining contracts, inputs, outputs, stability, and how developers use APIs.
|
||||
Explain how compilers work in a detailed sequence from parsing through optimization to code generation.
|
||||
Describe merge sort thoroughly, including divide and conquer, merging, complexity, and when it is useful.
|
||||
Explain why backups matter in depth, including accidents, hardware failure, ransomware, recovery time, and testing restores.
|
||||
Define latency with a detailed discussion of delay, network hops, queues, tail latency, and user-visible impact.
|
||||
Explain public key cryptography thoroughly, including key pairs, encryption, signatures, trust, and common uses.
|
||||
Describe load balancers in detail, including traffic distribution, health checks, failover, and scaling.
|
||||
Explain how solar panels generate electricity with a detailed description of photons, cells, current, and inverters.
|
||||
Explain supply and demand thoroughly, including curves, equilibrium, shortages, surpluses, and real-world examples.
|
||||
Describe containerization in depth, including images, isolation, portability, orchestration, and operational benefits.
|
||||
Explain the water cycle in a detailed answer covering evaporation, condensation, precipitation, runoff, and groundwater.
|
||||
Define entropy in a broad explanation covering disorder, information, uncertainty, and why context changes the meaning.
|
||||
Explain how recommendation systems work in detail, including signals, user history, item similarity, ranking, and feedback loops.
|
||||
Explain why teams use version control in depth, covering history, collaboration, branching, review, and recovery.
|
||||
Describe neural networks thoroughly, including layers, weights, activations, training, and what they approximate.
|
||||
Explain database transactions in detail, including atomicity, consistency, isolation, durability, and examples.
|
||||
Describe event loops thoroughly, including queues, callbacks, nonblocking I/O, and why they matter in servers.
|
||||
Explain how airplanes stay aloft in depth, covering lift, wings, airflow, thrust, drag, and control surfaces.
|
||||
Explain why sleep matters with a detailed discussion of memory, hormones, immune function, mood, and recovery.
|
||||
Explain OAuth in detail, including delegated authorization, access tokens, scopes, consent, and common flows.
|
||||
Define bandwidth in a detailed answer covering capacity, throughput, bottlenecks, and how it differs from latency.
|
||||
Explain how a thermostat works in detail, including sensors, set points, feedback loops, and heating or cooling control.
|
||||
Describe memoization thoroughly, including caching function results, when it helps, and when it can waste memory.
|
||||
Explain why code reviews are useful in depth, including correctness, maintainability, knowledge sharing, and team standards.
|
||||
Describe how markets set prices in detail, including buyers, sellers, incentives, information, scarcity, and competition.
|
||||
Explain microservices thoroughly, including service boundaries, independent deployment, communication, and operational costs.
|
||||
Explain backpropagation in depth, including forward passes, gradients, the chain rule, and weight updates.
|
||||
Describe how a search engine ranks pages in detail, including crawling, indexing, relevance, links, freshness, and quality signals.
|
||||
Explain opportunity cost thoroughly, including tradeoffs, alternatives, hidden costs, and everyday decision examples.
|
||||
Describe rate limiting in detail, including quotas, windows, tokens, fairness, abuse prevention, and user experience.
|
||||
Explain why logging is useful in depth, covering debugging, audits, operations, metrics, and incident response.
|
||||
Explain climate change carefully, including greenhouse gases, warming mechanisms, impacts, uncertainty, and mitigation.
|
||||
Describe how checksums detect errors in detail, including summaries, corruption, collision risk, and practical uses.
|
||||
Explain database normalization thoroughly, including redundancy, anomalies, normal forms, and tradeoffs.
|
||||
Describe garbage collection in detail, including allocation, reachability, tracing, freeing memory, and pauses.
|
||||
Explain why encryption needs keys in depth, covering secrecy, algorithms, key exchange, storage, and compromise.
|
||||
Describe how GPS finds location thoroughly, including satellites, timing, trilateration, clocks, and error correction.
|
||||
Explain continuous integration in detail, including automated builds, tests, integration frequency, and feedback.
|
||||
Describe queues and stacks thoroughly, including ordering rules, operations, use cases, and examples.
|
||||
Explain overfitting in depth, including training error, generalization, model complexity, validation, and prevention.
|
||||
Describe DNS caching thoroughly, including TTLs, resolvers, stale entries, performance, and invalidation problems.
|
||||
Explain what a kernel is in detail, including hardware mediation, processes, memory, drivers, and system calls.
|
||||
Explain how to prioritize tasks in depth, covering urgency, impact, dependencies, deadlines, and opportunity cost.
|
||||
Explain why indexes can slow writes in detail, including maintenance, extra storage, page splits, and transaction overhead.
|
||||
Describe HTTP thoroughly, including requests, responses, methods, headers, status codes, and statelessness.
|
||||
Explain proxy servers in detail, including forwarding, filtering, caching, privacy, reverse proxies, and load balancing.
|
||||
Describe how compression saves space in depth, including redundancy, dictionaries, entropy coding, and lossless versus lossy tradeoffs.
|
||||
Explain database replication thoroughly, including primary and replicas, lag, failover, consistency, and read scaling.
|
||||
Describe how compilers differ from interpreters in depth, including translation timing, runtime behavior, optimization, and examples.
|
||||
Explain amortized cost thoroughly, including occasional expensive operations, averaged guarantees, and data structure examples.
|
||||
Describe Bloom filters in detail, including hashing, bit arrays, false positives, memory efficiency, and use cases.
|
||||
Explain why batteries degrade in depth, including chemistry, charge cycles, heat, fast charging, and capacity loss.
|
||||
Explain runoff voting thoroughly, including ballot ranking, elimination rounds, majority goals, and possible strategic effects.
|
||||
Describe semantic versioning in detail, including major, minor, and patch numbers, compatibility promises, and release examples.
|
||||
Explain how a CDN helps a website in depth, including edge caches, latency, bandwidth, resilience, and invalidation.
|
||||
Define idempotency thoroughly, including repeated operations, safety, HTTP examples, retries, and distributed systems relevance.
|
||||
Explain mutexes in detail, including mutual exclusion, critical sections, contention, deadlocks, and alternatives.
|
||||
Explain why monitoring matters thoroughly, covering metrics, logs, alerts, trends, incidents, and service reliability.
|
||||
Describe how email routing works in detail, including DNS records, SMTP servers, queues, spam checks, and delivery.
|
||||
Explain how neural embeddings represent text in depth, including vectors, semantic similarity, training signals, and retrieval uses.
|
||||
Describe dependency injection thoroughly, including object construction, decoupling, testing, configuration, and tradeoffs.
|
||||
Explain race conditions in detail, including shared state, timing, nondeterminism, examples, and prevention techniques.
|
||||
Explain why documentation helps teams in depth, covering onboarding, decisions, maintenance, coordination, and future debugging.
|
||||
Describe how a blockchain works thoroughly, including blocks, hashes, consensus, append-only history, and tradeoffs.
|
||||
Explain memory leaks in detail, including unreachable allocations, retained references, symptoms, and detection.
|
||||
Define consensus in distributed systems thoroughly, including agreement, failures, quorums, leaders, and consistency.
|
||||
Describe priority queues in depth, including priorities, heap implementations, operations, and scheduling examples.
|
||||
Explain why caching is hard to invalidate in detail, including stale data, dependencies, timing, consistency, and user impact.
|
||||
Explain how to debug a slow program thoroughly, covering measurement, profiling, hypotheses, bottlenecks, and regression checks.
|
||||
Describe firewalls in detail, including packet filtering, rules, ports, network boundaries, and limitations.
|
||||
Explain how interest rates affect borrowing in depth, including monthly payments, risk, central banks, and business investment.
|
||||
Describe what an operating system does thoroughly, including process scheduling, memory, files, devices, and permissions.
|
||||
Explain how a web browser renders a page in detail, including parsing, CSS, layout, painting, JavaScript, and compositing.
|
||||
Describe feature flags thoroughly, including runtime switches, gradual rollout, experiments, rollback, and cleanup.
|
||||
Explain how garbage collection pauses happen in detail, including stop-the-world phases, heap scanning, compaction, and latency.
|
||||
Define eventual consistency thoroughly, including replicas, propagation delay, conflicts, convergence, and application tradeoffs.
|
||||
Describe sharding in depth, including partition keys, distribution, rebalancing, hotspots, and cross-shard queries.
|
||||
Explain deadlocks in detail, including locks, waiting cycles, conditions, detection, and prevention strategies.
|
||||
Describe how a camera sensor captures light in depth, including pixels, photons, charge, color filters, exposure, and noise.
|
||||
Explain why indexes need selectivity thoroughly, including cardinality, query planners, scans, and cases where indexes are ignored.
|
||||
Describe what a scheduler does in detail, including tasks, priorities, fairness, deadlines, and resource allocation.
|
||||
Define precision and recall thoroughly, including true positives, false positives, false negatives, and when each metric matters.
|
||||
Explain streaming APIs in depth, including incremental delivery, backpressure, connections, errors, and client handling.
|
||||
Explain why projects need tests thoroughly, including regression prevention, design feedback, confidence, automation, and documentation.
|
||||
Describe how to choose a data structure in detail, including operations, constraints, complexity, memory, and real-world tradeoffs.
|
||||
Executable
+229
@@ -0,0 +1,229 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build a DS4 directional-steering vector from paired prompt sets.
|
||||
|
||||
The extractor asks ds4 to dump one 4096-wide activation row per layer, averages
|
||||
the target and control rows, and writes a flat f32 file with 43 layer vectors.
|
||||
At runtime ds4 applies:
|
||||
|
||||
y = y - scale * direction[layer] * dot(direction[layer], y)
|
||||
|
||||
Positive scale suppresses the target direction. Negative scale amplifies it.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import array
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
N_LAYER = 43
|
||||
N_EMBD = 4096
|
||||
|
||||
SPECIALS = {
|
||||
"bos": "<｜begin▁of▁sentence｜>",
|
||||
"user": "<｜User｜>",
|
||||
"assistant": "<｜Assistant｜>",
|
||||
"think": "<think>",
|
||||
"nothink": "</think>",
|
||||
}
|
||||
|
||||
|
||||
def read_prompt_file(path: Path) -> list[str]:
|
||||
"""Read one prompt per non-empty line, ignoring shell-style comments."""
|
||||
prompts: list[str] = []
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
prompts.append(line)
|
||||
if not prompts:
|
||||
raise SystemExit(f"{path}: no prompts found")
|
||||
return prompts
|
||||
|
||||
|
||||
def render_ds4_prompt(system: str, user: str, think: bool) -> str:
|
||||
"""Render the minimal DS4 chat prefix used for activation capture."""
|
||||
pieces = [SPECIALS["bos"]]
|
||||
if system:
|
||||
pieces.append(system)
|
||||
pieces += [
|
||||
SPECIALS["user"],
|
||||
user,
|
||||
SPECIALS["assistant"],
|
||||
SPECIALS["think"] if think else SPECIALS["nothink"],
|
||||
]
|
||||
return "".join(pieces)
|
||||
|
||||
|
||||
def normalize(v: list[float]) -> list[float]:
|
||||
n2 = sum(x * x for x in v)
|
||||
if n2 <= 0.0:
|
||||
return v
|
||||
inv = 1.0 / math.sqrt(n2)
|
||||
return [x * inv for x in v]
|
||||
|
||||
|
||||
def dot(a: list[float], b: list[float]) -> float:
|
||||
return sum(x * y for x, y in zip(a, b))
|
||||
|
||||
|
||||
def run_capture(
|
||||
ds4: Path,
|
||||
model: Path,
|
||||
prompt: str,
|
||||
system: str,
|
||||
think: bool,
|
||||
ctx: int,
|
||||
component: str,
|
||||
work: Path,
|
||||
) -> list[list[float]]:
|
||||
"""Run ds4 once and return the last prompt-row dump for every layer."""
|
||||
prompt_path = work / "prompt.txt"
|
||||
prompt_path.write_text(render_ds4_prompt(system, prompt, think), encoding="utf-8")
|
||||
dump_prefix = work / "dump"
|
||||
|
||||
env = os.environ.copy()
|
||||
env["DS4_METAL_GRAPH_DUMP_PREFIX"] = str(dump_prefix)
|
||||
env["DS4_METAL_GRAPH_DUMP_NAME"] = component
|
||||
env["DS4_METAL_GRAPH_DUMP_POS"] = "0"
|
||||
|
||||
cmd = [
|
||||
str(ds4),
|
||||
"-m", str(model),
|
||||
"--ctx", str(ctx),
|
||||
"--prompt-file", str(prompt_path),
|
||||
"-n", "1",
|
||||
]
|
||||
subprocess.run(cmd, cwd=ds4.parent, env=env, check=True,
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
|
||||
|
||||
rows: list[list[float]] = []
|
||||
for layer in range(N_LAYER):
|
||||
path = work / f"dump_{component}-{layer}_pos0.bin"
|
||||
data = array.array("f")
|
||||
with path.open("rb") as f:
|
||||
data.fromfile(f, path.stat().st_size // 4)
|
||||
if len(data) < N_EMBD or len(data) % N_EMBD != 0:
|
||||
raise RuntimeError(f"bad dump shape for {path}: {len(data)} floats")
|
||||
rows.append(list(data[-N_EMBD:]))
|
||||
return rows
|
||||
|
||||
|
||||
def add_rows(total: list[list[float]], rows: list[list[float]]) -> None:
|
||||
for layer in range(N_LAYER):
|
||||
dst = total[layer]
|
||||
src = rows[layer]
|
||||
for i, value in enumerate(src):
|
||||
dst[i] += value
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--ds4", default="./ds4", help="path to the ds4 CLI")
|
||||
ap.add_argument("--model", default="ds4flash.gguf", help="GGUF model path")
|
||||
ap.add_argument("--good-file", required=True,
|
||||
help="desired/target prompts, one per line")
|
||||
ap.add_argument("--bad-file", required=True,
|
||||
help="contrast/control prompts, one per line")
|
||||
ap.add_argument("--out", default="dir-steering/out/direction.json",
|
||||
help="metadata JSON path; .f32 is written next to it")
|
||||
ap.add_argument("--ctx", type=int, default=512)
|
||||
ap.add_argument("--system", default="You are a helpful assistant.")
|
||||
ap.add_argument("--component", default="ffn_out",
|
||||
choices=("ffn_out", "attn_out"),
|
||||
help="runtime-editable 4096-wide activation stream")
|
||||
ap.add_argument("--think", action="store_true",
|
||||
help="capture after <think>; default captures direct answers")
|
||||
ap.add_argument("--pair-normalize", action="store_true",
|
||||
help="average normalized per-pair differences")
|
||||
ap.add_argument("--no-orthogonalize", action="store_true",
|
||||
help="do not remove the component parallel to the control mean")
|
||||
args = ap.parse_args()
|
||||
|
||||
ds4 = Path(args.ds4).resolve()
|
||||
model = Path(args.model).resolve()
|
||||
good_prompts = read_prompt_file(Path(args.good_file))
|
||||
bad_prompts = read_prompt_file(Path(args.bad_file))
|
||||
n = min(len(good_prompts), len(bad_prompts))
|
||||
good_prompts = good_prompts[:n]
|
||||
bad_prompts = bad_prompts[:n]
|
||||
|
||||
good_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
|
||||
bad_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
|
||||
pair_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="ds4-dir-steer-") as td:
|
||||
root = Path(td)
|
||||
for i, (good, bad) in enumerate(zip(good_prompts, bad_prompts), 1):
|
||||
print(f"pair {i}/{n}", flush=True)
|
||||
gw = root / f"good-{i}"
|
||||
bw = root / f"bad-{i}"
|
||||
gw.mkdir()
|
||||
bw.mkdir()
|
||||
good_rows = run_capture(ds4, model, good, args.system, args.think,
|
||||
args.ctx, args.component, gw)
|
||||
bad_rows = run_capture(ds4, model, bad, args.system, args.think,
|
||||
args.ctx, args.component, bw)
|
||||
add_rows(good_sum, good_rows)
|
||||
add_rows(bad_sum, bad_rows)
|
||||
if args.pair_normalize:
|
||||
for layer in range(N_LAYER):
|
||||
diff = normalize([
|
||||
good_rows[layer][j] - bad_rows[layer][j]
|
||||
for j in range(N_EMBD)
|
||||
])
|
||||
for j, value in enumerate(diff):
|
||||
pair_sum[layer][j] += value
|
||||
|
||||
layers = []
|
||||
for layer in range(N_LAYER):
|
||||
good_mean = [x / n for x in good_sum[layer]]
|
||||
bad_mean = [x / n for x in bad_sum[layer]]
|
||||
if args.pair_normalize:
|
||||
direction = normalize([x / n for x in pair_sum[layer]])
|
||||
else:
|
||||
direction = normalize([
|
||||
good_mean[i] - bad_mean[i]
|
||||
for i in range(N_EMBD)
|
||||
])
|
||||
if not args.no_orthogonalize:
|
||||
base = normalize(bad_mean)
|
||||
projection = dot(direction, base)
|
||||
direction = normalize([
|
||||
direction[i] - projection * base[i]
|
||||
for i in range(N_EMBD)
|
||||
])
|
||||
layers.append(direction)
|
||||
|
||||
out = Path(args.out)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
payload = {
|
||||
"format": "ds4-directional-steering-v1",
|
||||
"shape": [N_LAYER, N_EMBD],
|
||||
"component": args.component,
|
||||
"thinking": bool(args.think),
|
||||
"pair_normalize": bool(args.pair_normalize),
|
||||
"orthogonalize_control_mean": not args.no_orthogonalize,
|
||||
"good_file": str(Path(args.good_file)),
|
||||
"bad_file": str(Path(args.bad_file)),
|
||||
"model": str(model),
|
||||
"note": "runtime positive scale suppresses this direction; negative scale amplifies it",
|
||||
}
|
||||
out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
|
||||
flat = array.array("f")
|
||||
for direction in layers:
|
||||
flat.extend(direction)
|
||||
f32_out = out.with_suffix(".f32")
|
||||
with f32_out.open("wb") as f:
|
||||
flat.tofile(f)
|
||||
print(f"wrote {out}")
|
||||
print(f"wrote {f32_out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+64
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run a small steering scale sweep through ds4.
|
||||
|
||||
This is intentionally thin: it exercises the same public CLI options users
|
||||
will use in production and leaves all inference behavior inside ds4.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def read_prompts(path: Path) -> list[str]:
|
||||
prompts = []
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#"):
|
||||
prompts.append(line)
|
||||
if not prompts:
|
||||
raise SystemExit(f"{path}: no prompts found")
|
||||
return prompts
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--ds4", default="./ds4")
|
||||
ap.add_argument("--model", default="ds4flash.gguf")
|
||||
ap.add_argument("--direction", required=True,
|
||||
help="flat f32 vector file produced by build_direction.py")
|
||||
ap.add_argument("--prompts", required=True)
|
||||
ap.add_argument("--scales", default="-2,-1,-0.5,0,0.5,1,2")
|
||||
ap.add_argument("--tokens", type=int, default=160)
|
||||
ap.add_argument("--ctx", type=int, default=4096)
|
||||
ap.add_argument("--attn-scale", type=float, default=0.0)
|
||||
ap.add_argument("--nothink", action="store_true")
|
||||
args = ap.parse_args()
|
||||
|
||||
prompts = read_prompts(Path(args.prompts))
|
||||
scales = [float(x) for x in args.scales.split(",") if x.strip()]
|
||||
|
||||
for prompt in prompts:
|
||||
print("=" * 80)
|
||||
print(f"PROMPT: {prompt}")
|
||||
for scale in scales:
|
||||
print("-" * 80)
|
||||
print(f"FFN scale: {scale:g}")
|
||||
cmd = [
|
||||
args.ds4,
|
||||
"-m", args.model,
|
||||
"--ctx", str(args.ctx),
|
||||
"-n", str(args.tokens),
|
||||
"--temp", "0",
|
||||
"--dir-steering-file", args.direction,
|
||||
"--dir-steering-ffn", str(scale),
|
||||
"--dir-steering-attn", str(args.attn_scale),
|
||||
"-p", prompt,
|
||||
]
|
||||
if args.nothink:
|
||||
cmd.append("--nothink")
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+276
@@ -0,0 +1,276 @@
|
||||
#!/bin/sh
|
||||
set -e
|
||||
|
||||
REPO="antirez/deepseek-v4-gguf"
|
||||
Q2_IMATRIX_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf"
|
||||
Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf"
|
||||
Q2_Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf"
|
||||
PRO_Q2_IMATRIX_FILE="DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf"
|
||||
PRO_Q4_LAYERS00_30_FILE="DeepSeek-V4-Pro-Q4K-Layers00-30.gguf"
|
||||
PRO_Q4_LAYERS31_OUTPUT_FILE="DeepSeek-V4-Pro-Q4K-Layers-31-output.gguf"
|
||||
MTP_FILE="DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf"
|
||||
|
||||
ROOT=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
|
||||
OUT_DIR=${DS4_GGUF_DIR:-"$ROOT/gguf"}
|
||||
case "$OUT_DIR" in
|
||||
/*) ;;
|
||||
*) OUT_DIR="$ROOT/$OUT_DIR" ;;
|
||||
esac
|
||||
TOKEN=${HF_TOKEN:-}
|
||||
|
||||
usage() {
|
||||
cat <<EOF
|
||||
DeepSeek V4 GGUF downloader
|
||||
|
||||
Usage:
|
||||
./download_model.sh q2-imatrix [--token TOKEN]
|
||||
./download_model.sh q2-q4-imatrix [--token TOKEN]
|
||||
./download_model.sh q4-imatrix [--token TOKEN]
|
||||
./download_model.sh pro-q2-imatrix [--token TOKEN]
|
||||
./download_model.sh pro-q4-layers00-30 [--token TOKEN]
|
||||
./download_model.sh pro-q4-layers31-output [--token TOKEN]
|
||||
./download_model.sh pro-q4-split [--token TOKEN]
|
||||
./download_model.sh mtp [--token TOKEN]
|
||||
|
||||
Targets:
|
||||
|
||||
q2-imatrix
|
||||
2-bit routed experts, about 81 GB on disk.
|
||||
Recommended model for 96 and 128 GB RAM machines.
|
||||
|
||||
q2-q4-imatrix
|
||||
Mixed Flash quant: mostly q2 routed experts, with the last 6 layers
|
||||
using q4 routed experts. About 98 GB on disk. Good for higher
|
||||
quality inference for 128 GB MacBooks. Works on DGX Spark but loading
|
||||
may struggle compared to q2-imatrix.
|
||||
|
||||
q4-imatrix
|
||||
4-bit routed experts, about 153 GB on disk.
|
||||
Recommended model for machines with 256 GB RAM or more.
|
||||
|
||||
pro-q2-imatrix
|
||||
DeepSeek V4 PRO q2 imatrix quant, as a single GGUF file. About 430 GB
|
||||
on disk; intended for 512 GB RAM machines.
|
||||
|
||||
pro-q4-layers00-30
|
||||
First half of the DeepSeek V4 PRO Q4 routed-expert quant, layers 0..30.
|
||||
Use on the coordinator in a two-Mac-Studio distributed run. About 426 GB.
|
||||
|
||||
pro-q4-layers31-output
|
||||
Second half of the DeepSeek V4 PRO Q4 routed-expert quant, layers
|
||||
31..output. Use on the worker in a two-Mac-Studio distributed run.
|
||||
About 412 GB.
|
||||
|
||||
pro-q4-split
|
||||
Downloads both PRO Q4 split files into the download directory. About
|
||||
838 GB total. This target does not update ./ds4flash.gguf.
|
||||
|
||||
mtp Optional speculative decoding component, about 3.5 GB on disk.
|
||||
It is useful with q2-imatrix, q2-q4-imatrix, and q4-imatrix, but must be
|
||||
enabled explicitly with --mtp when running ds4 or ds4-server.
|
||||
|
||||
Options:
|
||||
--token TOKEN Hugging Face token. Otherwise HF_TOKEN or the local HF token
|
||||
cache is used if present.
|
||||
|
||||
Environment:
|
||||
DS4_GGUF_DIR Directory used for downloaded GGUF files.
|
||||
Default: ./gguf
|
||||
|
||||
After main-model downloads the script updates:
|
||||
./ds4flash.gguf -> <download directory>/<selected model>
|
||||
|
||||
Then the default commands work:
|
||||
./ds4 -p "Hello"
|
||||
./ds4-server --ctx 100000
|
||||
|
||||
After downloading mtp, enable it explicitly, for example:
|
||||
./ds4 --mtp <download directory>/$MTP_FILE --mtp-draft 2
|
||||
|
||||
PRO files are downloaded with the official Hugging Face downloader because
|
||||
they are too large for the curl path used by the smaller GGUF files.
|
||||
EOF
|
||||
}
|
||||
|
||||
if [ $# -eq 0 ]; then
|
||||
usage
|
||||
exit 1
|
||||
fi
|
||||
|
||||
MODEL=$1
|
||||
shift
|
||||
MODEL_FILES=
|
||||
LINK_MODEL=1
|
||||
|
||||
case "$MODEL" in
|
||||
q2-imatrix) MODEL_FILE=$Q2_IMATRIX_FILE ;;
|
||||
q2-q4-imatrix) MODEL_FILE=$Q2_Q4_IMATRIX_FILE ;;
|
||||
q4-imatrix) MODEL_FILE=$Q4_IMATRIX_FILE ;;
|
||||
pro-q2-imatrix) MODEL_FILE=$PRO_Q2_IMATRIX_FILE ;;
|
||||
pro-q4-layers00-30) MODEL_FILE=$PRO_Q4_LAYERS00_30_FILE; LINK_MODEL=0 ;;
|
||||
pro-q4-layers31-output) MODEL_FILE=$PRO_Q4_LAYERS31_OUTPUT_FILE; LINK_MODEL=0 ;;
|
||||
pro-q4-split)
|
||||
MODEL_FILES="$PRO_Q4_LAYERS00_30_FILE $PRO_Q4_LAYERS31_OUTPUT_FILE"
|
||||
LINK_MODEL=0
|
||||
;;
|
||||
mtp) MODEL_FILE=$MTP_FILE; LINK_MODEL=0 ;;
|
||||
-h|--help|help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown model: $MODEL" >&2
|
||||
echo >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--token)
|
||||
shift
|
||||
if [ $# -eq 0 ]; then
|
||||
echo "Missing value after --token" >&2
|
||||
exit 1
|
||||
fi
|
||||
TOKEN=$1
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
if [ -z "$TOKEN" ] && [ -s "$HOME/.cache/huggingface/token" ]; then
|
||||
TOKEN=$(cat "$HOME/.cache/huggingface/token")
|
||||
fi
|
||||
|
||||
needs_hf_download() {
|
||||
case "$1" in
|
||||
"$PRO_Q2_IMATRIX_FILE"|"$PRO_Q4_LAYERS00_30_FILE"|"$PRO_Q4_LAYERS31_OUTPUT_FILE")
|
||||
return 0
|
||||
;;
|
||||
*)
|
||||
return 1
|
||||
;;
|
||||
esac
|
||||
}
|
||||
|
||||
find_hf_command() {
|
||||
if command -v hf >/dev/null 2>&1; then
|
||||
printf '%s\n' hf
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
download_one_hf() {
|
||||
file=$1
|
||||
out="$OUT_DIR/$file"
|
||||
part="$out.part"
|
||||
|
||||
mkdir -p "$OUT_DIR"
|
||||
|
||||
if [ -s "$out" ]; then
|
||||
echo "Already downloaded: $out"
|
||||
return
|
||||
fi
|
||||
|
||||
if [ -e "$part" ]; then
|
||||
echo "Found curl partial download: $part" >&2
|
||||
echo "The Hugging Face downloader cannot resume curl .part files." >&2
|
||||
echo "Move or remove that partial download before retrying this PRO target." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
HF_CMD=$(find_hf_command || true)
|
||||
if [ -z "$HF_CMD" ]; then
|
||||
echo "PRO downloads require the official Hugging Face CLI." >&2
|
||||
echo "Install it with:" >&2
|
||||
echo " python3 -m pip install -U huggingface_hub hf_xet" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Downloading $file"
|
||||
echo "from https://huggingface.co/$REPO"
|
||||
echo "using $HF_CMD download"
|
||||
echo "If the download stops, run the same command again to resume it."
|
||||
|
||||
if [ -n "$TOKEN" ]; then
|
||||
"$HF_CMD" download "$REPO" "$file" --repo-type model --local-dir "$OUT_DIR" --token "$TOKEN"
|
||||
else
|
||||
"$HF_CMD" download "$REPO" "$file" --repo-type model --local-dir "$OUT_DIR"
|
||||
fi
|
||||
|
||||
if [ ! -s "$out" ]; then
|
||||
echo "Hugging Face download finished but expected file is missing: $out" >&2
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
download_one() {
|
||||
file=$1
|
||||
out="$OUT_DIR/$file"
|
||||
part="$out.part"
|
||||
aria2_part="$out.aria2"
|
||||
url="https://huggingface.co/$REPO/resolve/main/$file"
|
||||
|
||||
if needs_hf_download "$file"; then
|
||||
download_one_hf "$file"
|
||||
return
|
||||
fi
|
||||
|
||||
mkdir -p "$OUT_DIR"
|
||||
|
||||
if [ -e "$aria2_part" ]; then
|
||||
echo "Found incomplete aria2 download sidecar: $aria2_part" >&2
|
||||
echo "Finish or remove that partial download before using this curl downloader." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ -s "$out" ]; then
|
||||
echo "Already downloaded: $out"
|
||||
return
|
||||
fi
|
||||
|
||||
echo "Downloading $file"
|
||||
echo "from https://huggingface.co/$REPO"
|
||||
echo "If the download stops, run the same command again to resume it."
|
||||
|
||||
if [ -n "$TOKEN" ]; then
|
||||
curl -fL --progress-meter -C - -H "Authorization: Bearer $TOKEN" -o "$part" "$url"
|
||||
else
|
||||
curl -fL --progress-meter -C - -o "$part" "$url"
|
||||
fi
|
||||
|
||||
mv "$part" "$out"
|
||||
}
|
||||
|
||||
if [ -n "$MODEL_FILES" ]; then
|
||||
for file in $MODEL_FILES; do
|
||||
download_one "$file"
|
||||
done
|
||||
else
|
||||
download_one "$MODEL_FILE"
|
||||
fi
|
||||
|
||||
if [ "$MODEL" = "mtp" ]; then
|
||||
echo
|
||||
echo "MTP is an optional component for q2-imatrix, q2-q4-imatrix, and q4-imatrix."
|
||||
echo "Enable it explicitly, for example:"
|
||||
echo " ./ds4 --mtp $OUT_DIR/$MTP_FILE --mtp-draft 2"
|
||||
elif [ "$MODEL" = "pro-q4-layers00-30" ] || [ "$MODEL" = "pro-q4-layers31-output" ] || [ "$MODEL" = "pro-q4-split" ]; then
|
||||
echo
|
||||
echo "Downloaded PRO Q4 distributed split file(s). Use them with --layers,"
|
||||
echo "for example coordinator layers 0:30 and worker layers 31:output."
|
||||
elif [ "$LINK_MODEL" -eq 1 ]; then
|
||||
cd "$ROOT"
|
||||
ln -sfn "$OUT_DIR/$MODEL_FILE" ds4flash.gguf
|
||||
echo "Linked ./ds4flash.gguf -> $OUT_DIR/$MODEL_FILE"
|
||||
fi
|
||||
|
||||
echo
|
||||
echo "Done."
|
||||
@@ -0,0 +1,335 @@
|
||||
#ifndef DS4_H
|
||||
#define DS4_H
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
|
||||
#include "ds4_ssd.h"
|
||||
|
||||
/* Public engine boundary.
|
||||
*
|
||||
* The CLI and server should treat ds4_engine as the loaded model and
|
||||
* ds4_session as one mutable inference timeline. A session owns the live KV
|
||||
* cache and logits; callers provide full token prefixes and let
|
||||
* ds4_session_sync() reuse, extend, or rebuild the graph state. Keep this
|
||||
* header narrow so HTTP/CLI code does not depend on tensor internals. */
|
||||
|
||||
typedef enum {
|
||||
DS4_BACKEND_METAL,
|
||||
DS4_BACKEND_CUDA,
|
||||
DS4_BACKEND_CPU,
|
||||
} ds4_backend;
|
||||
|
||||
typedef enum {
|
||||
DS4_THINK_NONE,
|
||||
DS4_THINK_HIGH,
|
||||
DS4_THINK_MAX,
|
||||
} ds4_think_mode;
|
||||
|
||||
typedef enum {
|
||||
DS4_LOG_DEFAULT,
|
||||
DS4_LOG_PREFILL,
|
||||
DS4_LOG_GENERATION,
|
||||
DS4_LOG_KVCACHE,
|
||||
DS4_LOG_TOOL,
|
||||
DS4_LOG_WARNING,
|
||||
DS4_LOG_TIMING,
|
||||
DS4_LOG_OK,
|
||||
DS4_LOG_ERROR,
|
||||
} ds4_log_type;
|
||||
|
||||
typedef struct {
|
||||
int *v;
|
||||
int len;
|
||||
int cap;
|
||||
} ds4_tokens;
|
||||
|
||||
typedef struct {
|
||||
int id;
|
||||
float logit;
|
||||
float logprob;
|
||||
} ds4_token_score;
|
||||
|
||||
#define DS4_DEFAULT_TEMPERATURE 1.0f
|
||||
#define DS4_DEFAULT_TOP_P 1.0f
|
||||
#define DS4_DEFAULT_MIN_P 0.05f
|
||||
|
||||
typedef struct ds4_engine ds4_engine;
|
||||
typedef struct ds4_session ds4_session;
|
||||
|
||||
typedef void (*ds4_session_progress_fn)(void *ud, const char *event, int current, int total);
|
||||
typedef bool (*ds4_session_cancel_fn)(void *ud);
|
||||
|
||||
#define DS4_SESSION_SYNC_INTERRUPTED 2
|
||||
|
||||
typedef enum {
|
||||
DS4_DISTRIBUTED_NONE = 0,
|
||||
DS4_DISTRIBUTED_COORDINATOR,
|
||||
DS4_DISTRIBUTED_WORKER,
|
||||
} ds4_distributed_role;
|
||||
|
||||
typedef struct {
|
||||
uint32_t start;
|
||||
uint32_t end;
|
||||
bool has_output;
|
||||
bool set;
|
||||
} ds4_distributed_layers;
|
||||
|
||||
typedef struct {
|
||||
ds4_distributed_role role;
|
||||
ds4_distributed_layers layers;
|
||||
const char *listen_host;
|
||||
int listen_port;
|
||||
const char *coordinator_host;
|
||||
int coordinator_port;
|
||||
uint32_t prefill_chunk;
|
||||
uint32_t prefill_window;
|
||||
uint32_t activation_bits;
|
||||
bool replay_check;
|
||||
bool debug;
|
||||
} ds4_distributed_options;
|
||||
|
||||
typedef struct {
|
||||
const char *model_path;
|
||||
const char *mtp_path;
|
||||
ds4_backend backend;
|
||||
int n_threads;
|
||||
uint32_t prefill_chunk;
|
||||
int mtp_draft_tokens;
|
||||
float mtp_margin;
|
||||
const char *directional_steering_file;
|
||||
const char *expert_profile_path;
|
||||
float directional_steering_attn;
|
||||
float directional_steering_ffn;
|
||||
int power_percent;
|
||||
uint32_t ssd_streaming_cache_experts;
|
||||
uint64_t ssd_streaming_cache_bytes;
|
||||
uint32_t ssd_streaming_preload_experts;
|
||||
uint64_t simulate_used_memory_bytes;
|
||||
bool warm_weights;
|
||||
bool quality;
|
||||
bool ssd_streaming;
|
||||
bool ssd_streaming_cold;
|
||||
bool inspect_only;
|
||||
bool load_slice;
|
||||
uint32_t load_layer_start;
|
||||
uint32_t load_layer_end;
|
||||
bool load_output;
|
||||
ds4_distributed_options distributed;
|
||||
} ds4_engine_options;
|
||||
|
||||
typedef void (*ds4_token_emit_fn)(void *ud, int token);
|
||||
typedef void (*ds4_generation_done_fn)(void *ud);
|
||||
|
||||
typedef struct {
|
||||
uint64_t total_bytes;
|
||||
uint64_t raw_bytes;
|
||||
uint64_t compressed_bytes;
|
||||
uint64_t scratch_bytes;
|
||||
uint32_t prefill_cap;
|
||||
uint32_t raw_cap;
|
||||
uint32_t comp_cap;
|
||||
} ds4_context_memory;
|
||||
|
||||
typedef struct {
|
||||
uint8_t *ptr;
|
||||
uint64_t len;
|
||||
uint64_t cap;
|
||||
} ds4_session_snapshot;
|
||||
|
||||
typedef struct {
|
||||
char *path;
|
||||
uint64_t bytes;
|
||||
} ds4_session_payload_file;
|
||||
|
||||
int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt);
|
||||
void ds4_engine_close(ds4_engine *e);
|
||||
void ds4_engine_summary(ds4_engine *e);
|
||||
int ds4_engine_vocab_size(ds4_engine *e);
|
||||
int ds4_engine_power(ds4_engine *e);
|
||||
int ds4_engine_set_power(ds4_engine *e, int power_percent);
|
||||
const char *ds4_engine_model_name(ds4_engine *e);
|
||||
int ds4_engine_layer_count(ds4_engine *e);
|
||||
uint32_t ds4_engine_layer_compress_ratio(ds4_engine *e, uint32_t layer);
|
||||
uint64_t ds4_engine_hidden_f32_values(ds4_engine *e);
|
||||
/* Stable id for cache compatibility. 0 is the original Flash shape, so old
|
||||
* KV files with the previously-zero reserved byte remain Flash-compatible;
|
||||
* Pro and later shapes must use nonzero ids. */
|
||||
int ds4_engine_model_id(ds4_engine *e);
|
||||
const char *ds4_backend_name(ds4_backend backend);
|
||||
bool ds4_think_mode_enabled(ds4_think_mode mode);
|
||||
const char *ds4_think_mode_name(ds4_think_mode mode);
|
||||
const char *ds4_think_max_prefix(void);
|
||||
uint32_t ds4_think_max_min_context(void);
|
||||
ds4_think_mode ds4_think_mode_for_context(ds4_think_mode mode, int ctx_size);
|
||||
/* Uses the active model shape selected by ds4_engine_open(); call after opening
|
||||
* the GGUF so Flash/Pro dimensions are known. */
|
||||
ds4_context_memory ds4_context_memory_estimate(ds4_backend backend, int ctx_size);
|
||||
ds4_context_memory ds4_context_memory_estimate_with_prefill(
|
||||
ds4_backend backend,
|
||||
int ctx_size,
|
||||
uint32_t prefill_chunk);
|
||||
bool ds4_log_is_tty(FILE *fp);
|
||||
void ds4_log(FILE *fp, ds4_log_type type, const char *fmt, ...);
|
||||
int ds4_engine_generate_argmax(ds4_engine *e, const ds4_tokens *prompt,
|
||||
int n_predict, int ctx_size,
|
||||
ds4_token_emit_fn emit,
|
||||
ds4_generation_done_fn done,
|
||||
void *emit_ud,
|
||||
ds4_session_progress_fn progress,
|
||||
void *progress_ud);
|
||||
int ds4_engine_collect_imatrix(ds4_engine *e,
|
||||
const char *dataset_path,
|
||||
const char *output_path,
|
||||
int ctx_size,
|
||||
int max_prompts,
|
||||
int max_tokens);
|
||||
void ds4_engine_dump_tokens(ds4_engine *e, const ds4_tokens *tokens);
|
||||
int ds4_dump_text_tokenization(const char *model_path, const char *text, FILE *fp);
|
||||
int ds4_engine_head_test(ds4_engine *e, const ds4_tokens *prompt);
|
||||
int ds4_engine_first_token_test(ds4_engine *e, const ds4_tokens *prompt);
|
||||
int ds4_engine_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt);
|
||||
int ds4_engine_metal_graph_full_test(ds4_engine *e, const ds4_tokens *prompt);
|
||||
int ds4_engine_metal_graph_prompt_test(ds4_engine *e, const ds4_tokens *prompt, int ctx_size);
|
||||
|
||||
void ds4_tokens_push(ds4_tokens *tv, int token);
|
||||
void ds4_tokens_free(ds4_tokens *tv);
|
||||
void ds4_tokens_copy(ds4_tokens *dst, const ds4_tokens *src);
|
||||
bool ds4_tokens_starts_with(const ds4_tokens *tokens, const ds4_tokens *prefix);
|
||||
|
||||
void ds4_tokenize_text(ds4_engine *e, const char *text, ds4_tokens *out);
|
||||
void ds4_tokenize_rendered_chat(ds4_engine *e, const char *text, ds4_tokens *out);
|
||||
void ds4_chat_begin(ds4_engine *e, ds4_tokens *tokens);
|
||||
void ds4_encode_chat_prompt(
|
||||
ds4_engine *e,
|
||||
const char *system,
|
||||
const char *prompt,
|
||||
ds4_think_mode think_mode,
|
||||
ds4_tokens *out);
|
||||
void ds4_chat_append_max_effort_prefix(ds4_engine *e, ds4_tokens *tokens);
|
||||
void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role, const char *content);
|
||||
void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode think_mode);
|
||||
|
||||
char *ds4_token_text(ds4_engine *e, int token, size_t *len);
|
||||
int ds4_token_eos(ds4_engine *e);
|
||||
int ds4_token_user(ds4_engine *e);
|
||||
int ds4_token_assistant(ds4_engine *e);
|
||||
|
||||
int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size);
|
||||
void ds4_session_free(ds4_session *s);
|
||||
int ds4_session_power(ds4_session *s);
|
||||
int ds4_session_set_power(ds4_session *s, int power_percent);
|
||||
bool ds4_session_is_distributed(ds4_session *s);
|
||||
void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
|
||||
/* UI-only progress. It may report fine-grained progress inside a prefill chunk;
|
||||
* callers must not treat it as a durable KV checkpoint boundary. */
|
||||
void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
|
||||
/* Optional cooperative cancellation. ds4_session_sync() checks it only at
|
||||
* safe boundaries where the live checkpoint is either unchanged or represents a
|
||||
* valid token prefix, and returns DS4_SESSION_SYNC_INTERRUPTED when it stops. */
|
||||
void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud);
|
||||
void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total);
|
||||
/* Distributed coordinator sessions return 1 when the full layer route is
|
||||
* available, 0 when it is still incomplete, and -1 for a local API error. */
|
||||
int ds4_session_distributed_route_ready(ds4_session *s, char *err, size_t errlen);
|
||||
|
||||
typedef enum {
|
||||
DS4_SESSION_REWRITE_ERROR = -1,
|
||||
DS4_SESSION_REWRITE_OK = 0,
|
||||
/* The live backend state cannot be rewritten safely in place. The caller should
|
||||
* restore an older checkpoint if it has one, then sync to the prompt. */
|
||||
DS4_SESSION_REWRITE_REBUILD_NEEDED = 1,
|
||||
} ds4_session_rewrite_result;
|
||||
|
||||
/* Synchronize the live session to a full prompt token prefix. If the current
|
||||
* checkpoint is a prefix, only the suffix is evaluated; otherwise the backend
|
||||
* state is refilled from scratch. */
|
||||
int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen);
|
||||
bool ds4_session_rewrite_requires_rebuild(int live_len, int canonical_len, int common);
|
||||
ds4_session_rewrite_result ds4_session_rewrite_from_common(
|
||||
ds4_session *s, const ds4_tokens *prompt, int common,
|
||||
char *err, size_t errlen);
|
||||
int ds4_session_common_prefix(ds4_session *s, const ds4_tokens *prompt);
|
||||
int ds4_session_argmax(ds4_session *s);
|
||||
int ds4_session_argmax_excluding(ds4_session *s, int excluded_id);
|
||||
int ds4_sample_logits(const float *logits, int n_vocab, float temperature,
|
||||
int top_k, float top_p, float min_p, uint64_t *rng);
|
||||
int ds4_session_sample(ds4_session *s, float temperature, int top_k, float top_p, float min_p, uint64_t *rng);
|
||||
int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k);
|
||||
int ds4_session_token_logprob(ds4_session *s, int token, ds4_token_score *out);
|
||||
int ds4_session_copy_logits(ds4_session *s, float *out, int cap);
|
||||
int ds4_session_set_logits(ds4_session *s, const float *logits, int n);
|
||||
int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen);
|
||||
int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
|
||||
int max_tokens, int eos_token,
|
||||
int *accepted, int accepted_cap,
|
||||
char *err, size_t errlen);
|
||||
void ds4_session_invalidate(ds4_session *s);
|
||||
void ds4_session_rewind(ds4_session *s, int pos);
|
||||
int ds4_session_pos(ds4_session *s);
|
||||
int ds4_session_ctx(ds4_session *s);
|
||||
int ds4_session_prefill_cap(ds4_session *s);
|
||||
int ds4_engine_routed_quant_bits(ds4_engine *e);
|
||||
bool ds4_engine_has_output_head(ds4_engine *e);
|
||||
bool ds4_engine_has_mtp(ds4_engine *e);
|
||||
int ds4_engine_mtp_draft_tokens(ds4_engine *e);
|
||||
const ds4_tokens *ds4_session_tokens(ds4_session *s);
|
||||
|
||||
/* Low-level graph slice entry points used by distributed inference. The
|
||||
* transport/session routing logic lives in ds4_distributed.c. */
|
||||
int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen);
|
||||
int ds4_session_eval_layer_slice(ds4_session *s,
|
||||
const int *tokens,
|
||||
uint32_t n_tokens,
|
||||
uint32_t pos0,
|
||||
uint32_t layer_start,
|
||||
uint32_t layer_end,
|
||||
const float *input_hc,
|
||||
float *output_hc,
|
||||
bool output_logits,
|
||||
float *logits,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
int ds4_session_eval_output_head_from_hc(ds4_session *s,
|
||||
const float *hidden_hc,
|
||||
uint32_t n_tokens,
|
||||
float *logits,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
|
||||
/* Disk KV payload helpers. HTTP/agent code owns the outer file header and
|
||||
* persistence policy; the engine owns the DS4-specific serialized graph state. */
|
||||
#define DS4_SESSION_PAYLOAD_MAGIC UINT32_C(0x34565344) /* "DSV4" */
|
||||
#define DS4_SESSION_PAYLOAD_VERSION UINT32_C(2)
|
||||
#define DS4_SESSION_PAYLOAD_U32_FIELDS 13u
|
||||
#define DS4_SESSION_LAYER_PAYLOAD_MAGIC UINT32_C(0x4c565344) /* "DSVL" */
|
||||
#define DS4_SESSION_LAYER_PAYLOAD_VERSION UINT32_C(1)
|
||||
#define DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS 14u
|
||||
|
||||
uint64_t ds4_session_payload_bytes(ds4_session *s);
|
||||
int ds4_session_stage_payload(ds4_session *s, ds4_session_payload_file *out,
|
||||
char *err, size_t errlen);
|
||||
int ds4_session_write_staged_payload(const ds4_session_payload_file *payload,
|
||||
FILE *fp, char *err, size_t errlen);
|
||||
void ds4_session_payload_file_free(ds4_session_payload_file *payload);
|
||||
int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen);
|
||||
int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, char *err, size_t errlen);
|
||||
int ds4_session_save_snapshot(ds4_session *s, ds4_session_snapshot *snap, char *err, size_t errlen);
|
||||
int ds4_session_load_snapshot(ds4_session *s, const ds4_session_snapshot *snap, char *err, size_t errlen);
|
||||
void ds4_session_snapshot_free(ds4_session_snapshot *snap);
|
||||
|
||||
uint64_t ds4_session_layer_payload_bytes(ds4_session *s,
|
||||
uint32_t layer_start,
|
||||
uint32_t layer_end);
|
||||
int ds4_session_save_layer_payload(ds4_session *s, FILE *fp,
|
||||
uint32_t layer_start, uint32_t layer_end,
|
||||
char *err, size_t errlen);
|
||||
int ds4_session_load_layer_payload(ds4_session *s, FILE *fp,
|
||||
uint64_t payload_bytes,
|
||||
const int *tokens, uint32_t n_tokens,
|
||||
uint32_t layer_start, uint32_t layer_end,
|
||||
char *err, size_t errlen);
|
||||
|
||||
#endif
|
||||
+10244
File diff suppressed because it is too large
Load Diff
+683
@@ -0,0 +1,683 @@
|
||||
#include "ds4.h"
|
||||
#include "ds4_distributed.h"
|
||||
#include "ds4_help.h"
|
||||
|
||||
/* Purpose-built throughput benchmark.
|
||||
*
|
||||
* The benchmark walks one fixed token sequence to configurable context
|
||||
* frontiers, measuring only the newest prefill interval at each frontier. It
|
||||
* then snapshots the live session in memory, performs a fixed greedy decode
|
||||
* run without allowing EOS, restores the snapshot, and continues to the next
|
||||
* frontier. Snapshot save/restore time is intentionally outside both timing
|
||||
* windows.
|
||||
*/
|
||||
|
||||
#include <errno.h>
|
||||
#include <limits.h>
|
||||
#include <math.h>
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <time.h>
|
||||
|
||||
typedef struct {
|
||||
const char *model_path;
|
||||
const char *prompt_path;
|
||||
const char *chat_prompt_path;
|
||||
const char *system;
|
||||
const char *csv_path;
|
||||
const char *expert_profile_path;
|
||||
ds4_backend backend;
|
||||
int threads;
|
||||
int ctx_start;
|
||||
int ctx_max;
|
||||
int ctx_alloc;
|
||||
int step_incr;
|
||||
int gen_tokens;
|
||||
int power_percent;
|
||||
uint32_t prefill_chunk;
|
||||
uint32_t ssd_streaming_cache_experts;
|
||||
uint64_t ssd_streaming_cache_bytes;
|
||||
uint32_t ssd_streaming_preload_experts;
|
||||
uint64_t simulate_used_memory_bytes;
|
||||
double step_mul;
|
||||
const char *dump_frontier_logits_dir;
|
||||
ds4_dist_options dist;
|
||||
bool warm_weights;
|
||||
bool quality;
|
||||
bool ssd_streaming;
|
||||
bool ssd_streaming_cold;
|
||||
} bench_config;
|
||||
|
||||
static double bench_now_sec(void) {
|
||||
struct timespec ts;
|
||||
clock_gettime(CLOCK_MONOTONIC, &ts);
|
||||
return (double)ts.tv_sec + (double)ts.tv_nsec / 1000000000.0;
|
||||
}
|
||||
|
||||
static void usage(FILE *fp, const char *topic) {
|
||||
ds4_help_print(fp, DS4_HELP_BENCH, topic);
|
||||
}
|
||||
|
||||
static int parse_int(const char *s, const char *opt) {
|
||||
char *end = NULL;
|
||||
long v = strtol(s, &end, 10);
|
||||
if (s[0] == '\0' || *end != '\0' || v <= 0 || v > INT_MAX) {
|
||||
fprintf(stderr, "ds4-bench: invalid value for %s: %s\n", opt, s);
|
||||
exit(2);
|
||||
}
|
||||
return (int)v;
|
||||
}
|
||||
|
||||
static int parse_nonnegative_int(const char *s, const char *opt) {
|
||||
char *end = NULL;
|
||||
long v = strtol(s, &end, 10);
|
||||
if (s[0] == '\0' || *end != '\0' || v < 0 || v > INT_MAX) {
|
||||
fprintf(stderr, "ds4-bench: invalid value for %s: %s\n", opt, s);
|
||||
exit(2);
|
||||
}
|
||||
return (int)v;
|
||||
}
|
||||
|
||||
static double parse_double_arg(const char *s, const char *opt) {
|
||||
char *end = NULL;
|
||||
double v = strtod(s, &end);
|
||||
if (s[0] == '\0' || *end != '\0' || !isfinite(v)) {
|
||||
fprintf(stderr, "ds4-bench: invalid value for %s: %s\n", opt, s);
|
||||
exit(2);
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
static const char *need_arg(int *i, int argc, char **argv, const char *opt) {
|
||||
if (*i + 1 >= argc) {
|
||||
fprintf(stderr, "ds4-bench: %s requires an argument\n", opt);
|
||||
exit(2);
|
||||
}
|
||||
return argv[++*i];
|
||||
}
|
||||
|
||||
static ds4_backend parse_backend(const char *s, const char *opt) {
|
||||
if (!strcmp(s, "metal")) return DS4_BACKEND_METAL;
|
||||
#ifdef DS4_ROCM_BUILD
|
||||
if (!strcmp(s, "rocm")) return DS4_BACKEND_CUDA;
|
||||
#else
|
||||
if (!strcmp(s, "cuda")) return DS4_BACKEND_CUDA;
|
||||
#endif
|
||||
if (!strcmp(s, "cpu")) return DS4_BACKEND_CPU;
|
||||
fprintf(stderr, "ds4-bench: invalid value for %s: %s\n", opt, s);
|
||||
#ifdef DS4_ROCM_BUILD
|
||||
fprintf(stderr, "ds4-bench: valid backends are: metal, rocm, cpu\n");
|
||||
#else
|
||||
fprintf(stderr, "ds4-bench: valid backends are: metal, cuda, cpu\n");
|
||||
#endif
|
||||
exit(2);
|
||||
}
|
||||
|
||||
static ds4_backend default_backend(void) {
|
||||
#ifdef DS4_NO_GPU
|
||||
return DS4_BACKEND_CPU;
|
||||
#elif defined(__APPLE__)
|
||||
return DS4_BACKEND_METAL;
|
||||
#else
|
||||
return DS4_BACKEND_CUDA;
|
||||
#endif
|
||||
}
|
||||
|
||||
static char *read_file(const char *path) {
|
||||
FILE *fp = fopen(path, "rb");
|
||||
if (!fp) {
|
||||
fprintf(stderr, "ds4-bench: failed to open %s: %s\n", path, strerror(errno));
|
||||
exit(1);
|
||||
}
|
||||
if (fseek(fp, 0, SEEK_END) != 0) {
|
||||
fprintf(stderr, "ds4-bench: failed to seek %s\n", path);
|
||||
fclose(fp);
|
||||
exit(1);
|
||||
}
|
||||
long n = ftell(fp);
|
||||
if (n < 0) {
|
||||
fprintf(stderr, "ds4-bench: failed to tell %s\n", path);
|
||||
fclose(fp);
|
||||
exit(1);
|
||||
}
|
||||
if (fseek(fp, 0, SEEK_SET) != 0) {
|
||||
fprintf(stderr, "ds4-bench: failed to rewind %s\n", path);
|
||||
fclose(fp);
|
||||
exit(1);
|
||||
}
|
||||
char *buf = malloc((size_t)n + 1);
|
||||
if (!buf) {
|
||||
fprintf(stderr, "ds4-bench: out of memory reading %s\n", path);
|
||||
fclose(fp);
|
||||
exit(1);
|
||||
}
|
||||
if (fread(buf, 1, (size_t)n, fp) != (size_t)n) {
|
||||
fprintf(stderr, "ds4-bench: failed to read %s\n", path);
|
||||
free(buf);
|
||||
fclose(fp);
|
||||
exit(1);
|
||||
}
|
||||
fclose(fp);
|
||||
buf[n] = '\0';
|
||||
return buf;
|
||||
}
|
||||
|
||||
static bench_config parse_options(int argc, char **argv) {
|
||||
bench_config c = {
|
||||
.model_path = "ds4flash.gguf",
|
||||
.system = "You are a helpful assistant.",
|
||||
.backend = default_backend(),
|
||||
.ctx_start = 2048,
|
||||
.ctx_max = 32768,
|
||||
.step_incr = 2048,
|
||||
.gen_tokens = 128,
|
||||
.step_mul = 1.0,
|
||||
};
|
||||
|
||||
for (int i = 1; i < argc; i++) {
|
||||
const char *arg = argv[i];
|
||||
if (!strcmp(arg, "-h") || !strcmp(arg, "--help")) {
|
||||
const char *topic = (i + 1 < argc && argv[i + 1][0] != '-') ?
|
||||
argv[i + 1] : NULL;
|
||||
usage(stdout, topic);
|
||||
exit(0);
|
||||
}
|
||||
char dist_parse_err[256] = {0};
|
||||
ds4_dist_cli_parse_result dist_parse =
|
||||
ds4_dist_parse_cli_arg(arg,
|
||||
&i,
|
||||
argc,
|
||||
argv,
|
||||
&c.dist,
|
||||
dist_parse_err,
|
||||
sizeof(dist_parse_err));
|
||||
if (dist_parse == DS4_DIST_CLI_ERROR) {
|
||||
fprintf(stderr,
|
||||
"ds4-bench: %s\n",
|
||||
dist_parse_err[0] ? dist_parse_err : "invalid distributed option");
|
||||
exit(2);
|
||||
}
|
||||
if (dist_parse == DS4_DIST_CLI_MATCHED) continue;
|
||||
|
||||
if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
|
||||
c.model_path = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--prompt-file")) {
|
||||
c.prompt_path = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--chat-prompt-file")) {
|
||||
c.chat_prompt_path = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "-sys") || !strcmp(arg, "--system")) {
|
||||
c.system = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--ctx-start")) {
|
||||
c.ctx_start = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--ctx-max")) {
|
||||
c.ctx_max = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--ctx-alloc")) {
|
||||
c.ctx_alloc = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--step-incr")) {
|
||||
c.step_incr = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--step-mul")) {
|
||||
c.step_mul = parse_double_arg(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--gen-tokens") || !strcmp(arg, "--tokens") || !strcmp(arg, "-n")) {
|
||||
c.gen_tokens = parse_nonnegative_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--csv")) {
|
||||
c.csv_path = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--dump-frontier-logits-dir")) {
|
||||
c.dump_frontier_logits_dir = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--expert-profile")) {
|
||||
c.expert_profile_path = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "-t") || !strcmp(arg, "--threads")) {
|
||||
c.threads = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--backend")) {
|
||||
c.backend = parse_backend(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--metal")) {
|
||||
c.backend = DS4_BACKEND_METAL;
|
||||
#ifdef DS4_ROCM_BUILD
|
||||
} else if (!strcmp(arg, "--rocm")) {
|
||||
c.backend = DS4_BACKEND_CUDA;
|
||||
#else
|
||||
} else if (!strcmp(arg, "--cuda")) {
|
||||
c.backend = DS4_BACKEND_CUDA;
|
||||
#endif
|
||||
} else if (!strcmp(arg, "--cpu")) {
|
||||
c.backend = DS4_BACKEND_CPU;
|
||||
} else if (!strcmp(arg, "--quality")) {
|
||||
c.quality = true;
|
||||
} else if (!strcmp(arg, "--ssd-streaming")) {
|
||||
c.ssd_streaming = true;
|
||||
} else if (!strcmp(arg, "--ssd-streaming-cold")) {
|
||||
c.ssd_streaming_cold = true;
|
||||
} else if (!strcmp(arg, "--ssd-streaming-cache-experts")) {
|
||||
uint32_t experts = 0;
|
||||
uint64_t bytes = 0;
|
||||
if (!ds4_parse_streaming_cache_experts_arg(
|
||||
need_arg(&i, argc, argv, arg), &experts, &bytes)) {
|
||||
fprintf(stderr,
|
||||
"ds4-bench: --ssd-streaming-cache-experts must be a positive count or <number>GB\n");
|
||||
exit(2);
|
||||
}
|
||||
c.ssd_streaming_cache_experts = experts;
|
||||
c.ssd_streaming_cache_bytes = bytes;
|
||||
} else if (!strcmp(arg, "--ssd-streaming-preload-experts")) {
|
||||
int v = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
if (v <= 0) {
|
||||
fprintf(stderr, "ds4-bench: --ssd-streaming-preload-experts must be positive\n");
|
||||
exit(2);
|
||||
}
|
||||
c.ssd_streaming_preload_experts = (uint32_t)v;
|
||||
} else if (!strcmp(arg, "--simulate-used-memory")) {
|
||||
if (!ds4_parse_gib_arg(need_arg(&i, argc, argv, arg),
|
||||
&c.simulate_used_memory_bytes)) {
|
||||
fprintf(stderr,
|
||||
"ds4-bench: --simulate-used-memory must be a positive GiB value, e.g. 64GB\n");
|
||||
exit(2);
|
||||
}
|
||||
} else if (!strcmp(arg, "--prefill-chunk")) {
|
||||
c.prefill_chunk = (uint32_t)parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--power")) {
|
||||
c.power_percent = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
if (c.power_percent < 1 || c.power_percent > 100) {
|
||||
fprintf(stderr, "ds4-bench: --power must be between 1 and 100\n");
|
||||
exit(2);
|
||||
}
|
||||
} else if (!strcmp(arg, "--warm-weights")) {
|
||||
c.warm_weights = true;
|
||||
} else {
|
||||
fprintf(stderr, "ds4-bench: unknown option: %s\n", arg);
|
||||
usage(stderr, NULL);
|
||||
exit(2);
|
||||
}
|
||||
}
|
||||
|
||||
if (!!c.prompt_path == !!c.chat_prompt_path) {
|
||||
fprintf(stderr, "ds4-bench: specify exactly one of --prompt-file or --chat-prompt-file\n");
|
||||
exit(2);
|
||||
}
|
||||
if (c.ctx_start > c.ctx_max) {
|
||||
fprintf(stderr, "ds4-bench: --ctx-start must be <= --ctx-max\n");
|
||||
exit(2);
|
||||
}
|
||||
if (c.step_mul < 1.0) {
|
||||
fprintf(stderr, "ds4-bench: --step-mul must be >= 1\n");
|
||||
exit(2);
|
||||
}
|
||||
if (c.step_mul == 1.0 && c.step_incr <= 0) {
|
||||
fprintf(stderr, "ds4-bench: --step-incr must be positive when --step-mul is 1\n");
|
||||
exit(2);
|
||||
}
|
||||
if (c.ctx_max > INT_MAX - c.gen_tokens - 1) {
|
||||
fprintf(stderr, "ds4-bench: requested context is too large\n");
|
||||
exit(2);
|
||||
}
|
||||
if (c.ctx_alloc == 0) c.ctx_alloc = c.ctx_max + c.gen_tokens + 1;
|
||||
if (c.ctx_alloc <= c.ctx_max + c.gen_tokens) {
|
||||
fprintf(stderr, "ds4-bench: --ctx-alloc must be greater than ctx-max + gen-tokens\n");
|
||||
exit(2);
|
||||
}
|
||||
char dist_err[256];
|
||||
if (ds4_dist_prepare_engine_options(&c.dist, NULL, dist_err, sizeof(dist_err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: %s\n", dist_err);
|
||||
exit(2);
|
||||
}
|
||||
if (c.dist.role == DS4_DISTRIBUTED_WORKER) {
|
||||
fprintf(stderr, "ds4-bench: --role worker is a serving mode; start workers with ./ds4\n");
|
||||
exit(2);
|
||||
}
|
||||
return c;
|
||||
}
|
||||
|
||||
static void json_write_string(FILE *fp, const char *s) {
|
||||
fputc('"', fp);
|
||||
if (s) {
|
||||
for (const unsigned char *p = (const unsigned char *)s; *p; p++) {
|
||||
switch (*p) {
|
||||
case '"': fputs("\\\"", fp); break;
|
||||
case '\\': fputs("\\\\", fp); break;
|
||||
case '\b': fputs("\\b", fp); break;
|
||||
case '\f': fputs("\\f", fp); break;
|
||||
case '\n': fputs("\\n", fp); break;
|
||||
case '\r': fputs("\\r", fp); break;
|
||||
case '\t': fputs("\\t", fp); break;
|
||||
default:
|
||||
if (*p < 0x20) fprintf(fp, "\\u%04x", (unsigned)*p);
|
||||
else fputc((char)*p, fp);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
fputc('"', fp);
|
||||
}
|
||||
|
||||
static int write_frontier_logits_json(
|
||||
const bench_config *cfg,
|
||||
ds4_engine *engine,
|
||||
ds4_session *session,
|
||||
int frontier,
|
||||
int previous) {
|
||||
if (!cfg->dump_frontier_logits_dir) return 0;
|
||||
|
||||
const int vocab = ds4_engine_vocab_size(engine);
|
||||
float *logits = malloc((size_t)vocab * sizeof(logits[0]));
|
||||
if (!logits) {
|
||||
fprintf(stderr, "ds4-bench: out of memory copying frontier logits\n");
|
||||
return 1;
|
||||
}
|
||||
if (ds4_session_copy_logits(session, logits, vocab) != vocab) {
|
||||
fprintf(stderr, "ds4-bench: failed to copy frontier logits at %d\n", frontier);
|
||||
free(logits);
|
||||
return 1;
|
||||
}
|
||||
|
||||
char path[PATH_MAX];
|
||||
const int n = snprintf(path,
|
||||
sizeof(path),
|
||||
"%s/frontier_%06d.logits.json",
|
||||
cfg->dump_frontier_logits_dir,
|
||||
frontier);
|
||||
if (n <= 0 || (size_t)n >= sizeof(path)) {
|
||||
fprintf(stderr, "ds4-bench: frontier logits path is too long\n");
|
||||
free(logits);
|
||||
return 1;
|
||||
}
|
||||
|
||||
FILE *fp = fopen(path, "wb");
|
||||
if (!fp) {
|
||||
fprintf(stderr, "ds4-bench: failed to open %s: %s\n", path, strerror(errno));
|
||||
free(logits);
|
||||
return 1;
|
||||
}
|
||||
|
||||
const int argmax = ds4_session_argmax(session);
|
||||
fprintf(fp, "{\n \"source\":\"ds4-bench\",\n \"model\":");
|
||||
json_write_string(fp, cfg->model_path);
|
||||
fprintf(fp,
|
||||
",\n \"backend\":\"%s\",\n \"quality\":%s,\n"
|
||||
" \"quant_bits\":%d,\n \"prompt_tokens\":%d,\n"
|
||||
" \"frontier_tokens\":%d,\n \"prefill_tokens\":%d,\n"
|
||||
" \"ctx\":%d,\n \"vocab\":%d,\n"
|
||||
" \"argmax_id\":%d,\n \"argmax_logit\":%.9g,\n \"logits\":[",
|
||||
ds4_backend_name(cfg->backend),
|
||||
cfg->quality ? "true" : "false",
|
||||
ds4_engine_routed_quant_bits(engine),
|
||||
frontier,
|
||||
frontier,
|
||||
frontier - previous,
|
||||
cfg->ctx_alloc,
|
||||
vocab,
|
||||
argmax,
|
||||
logits[argmax]);
|
||||
for (int i = 0; i < vocab; i++) {
|
||||
if (i) fputc(',', fp);
|
||||
if ((i % 8) == 0) fputs("\n ", fp);
|
||||
if (isfinite(logits[i])) fprintf(fp, "%.9g", logits[i]);
|
||||
else fputs("null", fp);
|
||||
}
|
||||
fputs("\n ]\n}\n", fp);
|
||||
if (fclose(fp) != 0) {
|
||||
fprintf(stderr, "ds4-bench: failed to close %s\n", path);
|
||||
free(logits);
|
||||
return 1;
|
||||
}
|
||||
free(logits);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int next_frontier(const bench_config *c, int cur) {
|
||||
if (cur >= c->ctx_max) return c->ctx_max;
|
||||
int next;
|
||||
if (c->step_mul == 1.0) {
|
||||
if (cur > INT_MAX - c->step_incr) next = c->ctx_max;
|
||||
else next = cur + c->step_incr;
|
||||
} else {
|
||||
const double v = ceil((double)cur * c->step_mul);
|
||||
next = v > (double)INT_MAX ? c->ctx_max : (int)v;
|
||||
if (next <= cur) next = cur + 1;
|
||||
}
|
||||
if (next > c->ctx_max) next = c->ctx_max;
|
||||
return next;
|
||||
}
|
||||
|
||||
static void log_context_memory(ds4_backend backend,
|
||||
int ctx_size,
|
||||
uint32_t prefill_chunk) {
|
||||
ds4_context_memory m =
|
||||
ds4_context_memory_estimate_with_prefill(backend,
|
||||
ctx_size,
|
||||
prefill_chunk);
|
||||
fprintf(stderr,
|
||||
"ds4-bench: context buffers %.2f MiB (ctx=%d, backend=%s, prefill_chunk=%u, raw_kv_rows=%u, compressed_kv_rows=%u)\n",
|
||||
(double)m.total_bytes / (1024.0 * 1024.0),
|
||||
ctx_size,
|
||||
ds4_backend_name(backend),
|
||||
m.prefill_cap,
|
||||
m.raw_cap,
|
||||
m.comp_cap);
|
||||
}
|
||||
|
||||
static int wait_distributed_route(ds4_session *session) {
|
||||
char err[256] = {0};
|
||||
char last[256] = {0};
|
||||
unsigned ticks = 0;
|
||||
const struct timespec delay = {0, 250000000L};
|
||||
|
||||
for (;;) {
|
||||
int ready = ds4_session_distributed_route_ready(session, err, sizeof(err));
|
||||
if (ready > 0) {
|
||||
if (ticks) fprintf(stderr, "ds4-bench: distributed route ready\n");
|
||||
return 0;
|
||||
}
|
||||
if (ready < 0) {
|
||||
fprintf(stderr,
|
||||
"ds4-bench: distributed route readiness failed: %s\n",
|
||||
err[0] ? err : "unknown error");
|
||||
return 1;
|
||||
}
|
||||
const char *why = err[0] ? err : "route incomplete";
|
||||
if (strcmp(last, why) != 0 || (ticks % 20u) == 0) {
|
||||
fprintf(stderr, "ds4-bench: waiting for distributed route: %s\n", why);
|
||||
snprintf(last, sizeof(last), "%s", why);
|
||||
}
|
||||
nanosleep(&delay, NULL);
|
||||
ticks++;
|
||||
}
|
||||
}
|
||||
|
||||
static void maybe_warn_distributed_step_shape(const bench_config *cfg, ds4_session *session) {
|
||||
if (!cfg || !session || cfg->dist.role != DS4_DISTRIBUTED_COORDINATOR) return;
|
||||
uint32_t chunk = cfg->dist.prefill_chunk;
|
||||
if (chunk == 0) {
|
||||
const int cap = ds4_session_prefill_cap(session);
|
||||
if (cap > 0) chunk = (uint32_t)cap;
|
||||
}
|
||||
if (chunk == 0) return;
|
||||
if (cfg->step_mul == 1.0 &&
|
||||
cfg->step_incr > 0 &&
|
||||
(uint32_t)cfg->step_incr < chunk &&
|
||||
cfg->ctx_start < cfg->ctx_max)
|
||||
{
|
||||
fprintf(stderr,
|
||||
"ds4-bench: note: --step-incr=%d is smaller than distributed prefill chunk %u; "
|
||||
"suffix rows will not show multi-chunk pipeline overlap\n",
|
||||
cfg->step_incr,
|
||||
chunk);
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
bench_config cfg = parse_options(argc, argv);
|
||||
|
||||
ds4_engine_options opt = {
|
||||
.model_path = cfg.model_path,
|
||||
.backend = cfg.backend,
|
||||
.n_threads = cfg.threads,
|
||||
.prefill_chunk = cfg.prefill_chunk,
|
||||
.ssd_streaming_cache_experts = cfg.ssd_streaming_cache_experts,
|
||||
.ssd_streaming_cache_bytes = cfg.ssd_streaming_cache_bytes,
|
||||
.ssd_streaming_preload_experts = cfg.ssd_streaming_preload_experts,
|
||||
.simulate_used_memory_bytes = cfg.simulate_used_memory_bytes,
|
||||
.power_percent = cfg.power_percent,
|
||||
.warm_weights = cfg.warm_weights,
|
||||
.quality = cfg.quality,
|
||||
.ssd_streaming = cfg.ssd_streaming,
|
||||
.ssd_streaming_cold = cfg.ssd_streaming_cold,
|
||||
.expert_profile_path = cfg.expert_profile_path,
|
||||
.distributed = cfg.dist,
|
||||
};
|
||||
char dist_err[256];
|
||||
if (ds4_dist_prepare_engine_options(&cfg.dist, &opt, dist_err, sizeof(dist_err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: %s\n", dist_err);
|
||||
return 2;
|
||||
}
|
||||
ds4_engine *engine = NULL;
|
||||
if (ds4_engine_open(&engine, &opt) != 0) return 1;
|
||||
log_context_memory(cfg.backend, cfg.ctx_alloc, cfg.prefill_chunk);
|
||||
|
||||
char *text = read_file(cfg.prompt_path ? cfg.prompt_path : cfg.chat_prompt_path);
|
||||
ds4_tokens prompt = {0};
|
||||
if (cfg.chat_prompt_path) {
|
||||
ds4_encode_chat_prompt(engine, cfg.system, text, DS4_THINK_NONE, &prompt);
|
||||
} else {
|
||||
ds4_tokenize_text(engine, text, &prompt);
|
||||
}
|
||||
free(text);
|
||||
|
||||
if (prompt.len < cfg.ctx_max) {
|
||||
fprintf(stderr,
|
||||
"ds4-bench: prompt has %d tokens, need at least --ctx-max=%d\n",
|
||||
prompt.len,
|
||||
cfg.ctx_max);
|
||||
ds4_tokens_free(&prompt);
|
||||
ds4_engine_close(engine);
|
||||
return 1;
|
||||
}
|
||||
|
||||
ds4_session *session = NULL;
|
||||
if (ds4_session_create(&session, engine, cfg.ctx_alloc) != 0) {
|
||||
fprintf(stderr, "ds4-bench: failed to create session\n");
|
||||
ds4_tokens_free(&prompt);
|
||||
ds4_engine_close(engine);
|
||||
return 1;
|
||||
}
|
||||
if (cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR &&
|
||||
wait_distributed_route(session) != 0)
|
||||
{
|
||||
ds4_session_free(session);
|
||||
ds4_tokens_free(&prompt);
|
||||
ds4_engine_close(engine);
|
||||
return 1;
|
||||
}
|
||||
maybe_warn_distributed_step_shape(&cfg, session);
|
||||
|
||||
FILE *out = stdout;
|
||||
if (cfg.csv_path) {
|
||||
out = fopen(cfg.csv_path, "wb");
|
||||
if (!out) {
|
||||
fprintf(stderr, "ds4-bench: failed to open %s: %s\n", cfg.csv_path, strerror(errno));
|
||||
ds4_session_free(session);
|
||||
ds4_tokens_free(&prompt);
|
||||
ds4_engine_close(engine);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
fprintf(out, "ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,kvcache_bytes\n");
|
||||
fflush(out);
|
||||
|
||||
const int eos = ds4_token_eos(engine);
|
||||
const bool distributed = cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR;
|
||||
ds4_session_snapshot snap = {0};
|
||||
char err[256];
|
||||
int previous = 0;
|
||||
int rc = 0;
|
||||
|
||||
for (int frontier = cfg.ctx_start; ; frontier = next_frontier(&cfg, frontier)) {
|
||||
ds4_tokens prefix = {
|
||||
.v = prompt.v,
|
||||
.len = frontier,
|
||||
.cap = frontier,
|
||||
};
|
||||
|
||||
const double prefill_t0 = bench_now_sec();
|
||||
if (ds4_session_sync(session, &prefix, err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: prefill to %d failed: %s\n", frontier, err);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
const double prefill_t1 = bench_now_sec();
|
||||
const double prefill_sec = prefill_t1 - prefill_t0;
|
||||
const int prefill_tokens = frontier - previous;
|
||||
|
||||
if (write_frontier_logits_json(&cfg, engine, session, frontier, previous) != 0) {
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
|
||||
if (cfg.gen_tokens > 0 && !distributed) {
|
||||
if (ds4_session_save_snapshot(session, &snap, err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: snapshot at %d failed: %s\n", frontier, err);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const double gen_t0 = bench_now_sec();
|
||||
for (int i = 0; i < cfg.gen_tokens; i++) {
|
||||
if (ds4_session_pos(session) + 1 >= ds4_session_ctx(session)) {
|
||||
fprintf(stderr, "ds4-bench: generation would exceed allocated context at frontier %d\n", frontier);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
const int token = ds4_session_argmax_excluding(session, eos);
|
||||
if (token < 0) {
|
||||
fprintf(stderr, "ds4-bench: failed to choose non-EOS token at frontier %d\n", frontier);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: decode at frontier %d failed: %s\n", frontier, err);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
const double gen_t1 = bench_now_sec();
|
||||
if (rc != 0) break;
|
||||
|
||||
if (cfg.gen_tokens == 0) {
|
||||
/* Pure prefill benchmark: leave the live session at the frontier. */
|
||||
} else if (distributed) {
|
||||
if (ds4_session_sync(session, &prefix, err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: distributed replay restore at %d failed: %s\n", frontier, err);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
if (ds4_session_load_snapshot(session, &snap, err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "ds4-bench: restore at %d failed: %s\n", frontier, err);
|
||||
rc = 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const double gen_sec = gen_t1 - gen_t0;
|
||||
fprintf(out,
|
||||
"%d,%d,%.2f,%d,%.2f,%llu\n",
|
||||
frontier,
|
||||
prefill_tokens,
|
||||
prefill_sec > 0.0 ? (double)prefill_tokens / prefill_sec : 0.0,
|
||||
cfg.gen_tokens,
|
||||
gen_sec > 0.0 ? (double)cfg.gen_tokens / gen_sec : 0.0,
|
||||
(unsigned long long)(distributed ? 0 : snap.len));
|
||||
fflush(out);
|
||||
|
||||
previous = frontier;
|
||||
if (frontier >= cfg.ctx_max) break;
|
||||
}
|
||||
|
||||
if (out != stdout) fclose(out);
|
||||
ds4_session_snapshot_free(&snap);
|
||||
ds4_session_free(session);
|
||||
ds4_tokens_free(&prompt);
|
||||
ds4_engine_close(engine);
|
||||
return rc;
|
||||
}
|
||||
+13256
File diff suppressed because it is too large
Load Diff
+8414
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,125 @@
|
||||
#ifndef DS4_DISTRIBUTED_H
|
||||
#define DS4_DISTRIBUTED_H
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#include "ds4.h"
|
||||
|
||||
/* Distributed inference is an engine backend, not a separate frontend API.
|
||||
* Programs parse distributed options here, then keep using the normal
|
||||
* ds4_session_* calls. Only worker/coordinator serving modes call
|
||||
* ds4_dist_run() directly.
|
||||
*/
|
||||
typedef ds4_distributed_options ds4_dist_options;
|
||||
typedef struct ds4_dist_session ds4_dist_session;
|
||||
|
||||
/* Options used by standalone `./ds4 --role coordinator -p ...` generation.
|
||||
* Interactive tools and the server go through the normal ds4_session API.
|
||||
*/
|
||||
typedef struct {
|
||||
const char *prompt;
|
||||
const char *system;
|
||||
const char *dump_logits_path;
|
||||
const char *dump_logprobs_path;
|
||||
int dump_logprobs_top_k;
|
||||
int n_predict;
|
||||
int ctx_size;
|
||||
float temperature;
|
||||
float top_p;
|
||||
float min_p;
|
||||
uint64_t seed;
|
||||
ds4_think_mode think_mode;
|
||||
} ds4_dist_generation_options;
|
||||
|
||||
typedef enum {
|
||||
DS4_DIST_CLI_ERROR = -1,
|
||||
DS4_DIST_CLI_NOT_MATCHED = 0,
|
||||
DS4_DIST_CLI_MATCHED = 1,
|
||||
} ds4_dist_cli_parse_result;
|
||||
|
||||
/* Shared option parsing. */
|
||||
bool ds4_dist_enabled(const ds4_dist_options *opt);
|
||||
ds4_dist_options *ds4_dist_options_create(void);
|
||||
void ds4_dist_options_free(ds4_dist_options *opt);
|
||||
void ds4_dist_usage(FILE *fp);
|
||||
ds4_dist_cli_parse_result ds4_dist_parse_cli_arg(
|
||||
const char *arg,
|
||||
int *index,
|
||||
int argc,
|
||||
char **argv,
|
||||
ds4_dist_options *opt,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
|
||||
/* Applies distributed layer-loading choices to the engine options before the
|
||||
* model is loaded.
|
||||
*/
|
||||
int ds4_dist_prepare_engine_options(
|
||||
const ds4_dist_options *opt,
|
||||
ds4_engine_options *engine,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
|
||||
/* Coordinator session backend used by ds4.c. These mirror the normal session
|
||||
* operations; callers outside the engine should not need to call them directly.
|
||||
*/
|
||||
int ds4_dist_session_create(
|
||||
ds4_dist_session **out,
|
||||
ds4_engine *engine,
|
||||
const ds4_dist_options *opt,
|
||||
ds4_session *owner,
|
||||
int ctx_size,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
void ds4_dist_session_free(ds4_dist_session *d);
|
||||
|
||||
/* Returns 1 when the coordinator has full layer coverage, 0 when workers are
|
||||
* still missing, and -1 for configuration or internal errors.
|
||||
*/
|
||||
int ds4_dist_session_route_ready(ds4_dist_session *d, char *err, size_t errlen);
|
||||
|
||||
/* Synchronize the distributed KV state to the requested prompt timeline. */
|
||||
int ds4_dist_session_sync(
|
||||
ds4_dist_session *d,
|
||||
ds4_session *owner,
|
||||
const ds4_tokens *checkpoint,
|
||||
const ds4_tokens *prompt,
|
||||
float *logits,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
|
||||
/* Evaluate one additional token across the current distributed route. */
|
||||
int ds4_dist_session_eval(
|
||||
ds4_dist_session *d,
|
||||
ds4_session *owner,
|
||||
const ds4_tokens *checkpoint,
|
||||
int token,
|
||||
float *logits,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
|
||||
/* Save/load use the normal DSV4 payload format. The coordinator gathers or
|
||||
* pushes remote layer shards internally so saved files are topology-neutral.
|
||||
*/
|
||||
int ds4_dist_session_save_payload(
|
||||
ds4_dist_session *d,
|
||||
ds4_session *owner,
|
||||
FILE *fp,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
int ds4_dist_session_load_payload(
|
||||
ds4_dist_session *d,
|
||||
ds4_session *owner,
|
||||
FILE *fp,
|
||||
uint64_t payload_bytes,
|
||||
char *err,
|
||||
size_t errlen);
|
||||
|
||||
/* Standalone distributed mode. Workers stay in this loop; coordinator one-shot
|
||||
* mode uses it for `./ds4 --role coordinator -p ...`.
|
||||
*/
|
||||
int ds4_dist_run(ds4_engine *engine, const ds4_dist_options *opt, const ds4_dist_generation_options *gen);
|
||||
|
||||
#endif
|
||||
+4289
File diff suppressed because it is too large
Load Diff
+559
@@ -0,0 +1,559 @@
|
||||
#include "ds4_help.h"
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <stdio.h>
|
||||
#include <string.h>
|
||||
#include <unistd.h>
|
||||
|
||||
typedef struct {
|
||||
const char *off;
|
||||
const char *cyan;
|
||||
const char *title;
|
||||
const char *yellow;
|
||||
const char *grey;
|
||||
const char *red;
|
||||
const char *white;
|
||||
const char *bright;
|
||||
} help_colors;
|
||||
|
||||
static help_colors help_make_colors(FILE *fp) {
|
||||
bool color = isatty(fileno(fp));
|
||||
help_colors c = {0};
|
||||
if (!color) return c;
|
||||
c.off = "\x1b[0m";
|
||||
c.cyan = "\x1b[38;5;81m";
|
||||
c.title = "\x1b[1;38;5;250m";
|
||||
c.yellow = "\x1b[38;5;179m";
|
||||
c.grey = "\x1b[38;5;240m";
|
||||
c.red = "\x1b[38;5;203m";
|
||||
c.white = "\x1b[38;5;252m";
|
||||
c.bright = "\x1b[1;38;5;231m";
|
||||
return c;
|
||||
}
|
||||
|
||||
static void title(FILE *fp, const help_colors *c, const char *s) {
|
||||
fprintf(fp, "%s%s%s\n", c->title ? c->title : "", s, c->off ? c->off : "");
|
||||
}
|
||||
|
||||
static void title_red(FILE *fp, const help_colors *c, const char *s) {
|
||||
fprintf(fp, "%s%s%s\n", c->red ? c->red : "", s, c->off ? c->off : "");
|
||||
}
|
||||
|
||||
static bool option_name_has_switch(const char *name) {
|
||||
bool word_start = true;
|
||||
while (*name) {
|
||||
if (word_start && (*name == '-' || *name == '/')) return true;
|
||||
word_start = (*name == ' ');
|
||||
name++;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static void print_colored_option_name(FILE *fp, const help_colors *c, const char *name) {
|
||||
bool has_switch = option_name_has_switch(name);
|
||||
bool word_start = true;
|
||||
while (*name) {
|
||||
const char *start = name;
|
||||
while (*name && *name != ' ') name++;
|
||||
bool is_option = !has_switch || *start == '-' || *start == '/' ||
|
||||
(word_start && has_switch && *start != '[');
|
||||
const char *color = is_option ? c->cyan : c->bright;
|
||||
if (color) fputs(color, fp);
|
||||
fwrite(start, 1, (size_t)(name - start), fp);
|
||||
if (color && c->off) fputs(c->off, fp);
|
||||
if (*name == ' ') {
|
||||
fputc(*name++, fp);
|
||||
word_start = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void opt(FILE *fp, const help_colors *c, const char *name, const char *desc) {
|
||||
if (c->cyan) {
|
||||
fputs(" ", fp);
|
||||
print_colored_option_name(fp, c, name);
|
||||
fprintf(fp, " %s|%s ", c->grey ? c->grey : "", c->grey ? c->off : "");
|
||||
fprintf(fp, "%s%s%s\n", c->white ? c->white : "", desc,
|
||||
c->white ? c->off : "");
|
||||
return;
|
||||
}
|
||||
|
||||
const int col = 30;
|
||||
int n = (int)strlen(name);
|
||||
if (n > col) {
|
||||
fprintf(fp, " %s\n %s\n", name, desc);
|
||||
} else {
|
||||
fprintf(fp, " %-30s %s\n", name, desc);
|
||||
}
|
||||
}
|
||||
|
||||
static void para(FILE *fp, const help_colors *c, const char *s) {
|
||||
fprintf(fp, "%s%s%s\n",
|
||||
c->yellow ? c->yellow : "", s, c->yellow ? c->off : "");
|
||||
}
|
||||
|
||||
static bool streq(const char *a, const char *b) {
|
||||
return a && b && strcmp(a, b) == 0;
|
||||
}
|
||||
|
||||
static bool topic_is(const char *topic, const char *name) {
|
||||
return topic && strcmp(topic, name) == 0;
|
||||
}
|
||||
|
||||
static const char *tool_name(ds4_help_tool tool) {
|
||||
switch (tool) {
|
||||
case DS4_HELP_DS4: return "ds4";
|
||||
case DS4_HELP_SERVER: return "ds4-server";
|
||||
case DS4_HELP_AGENT: return "ds4-agent";
|
||||
case DS4_HELP_BENCH: return "ds4-bench";
|
||||
case DS4_HELP_EVAL: return "ds4-eval";
|
||||
}
|
||||
return "ds4";
|
||||
}
|
||||
|
||||
static const char *tool_usage(ds4_help_tool tool) {
|
||||
switch (tool) {
|
||||
case DS4_HELP_DS4:
|
||||
return "Usage: ds4 [(-p PROMPT | --prompt-file FILE)] [options]";
|
||||
case DS4_HELP_SERVER:
|
||||
return "Usage: ds4-server [options]";
|
||||
case DS4_HELP_AGENT:
|
||||
return "Usage: ds4-agent [options]";
|
||||
case DS4_HELP_BENCH:
|
||||
return "Usage: ds4-bench (--prompt-file FILE | --chat-prompt-file FILE) [options]";
|
||||
case DS4_HELP_EVAL:
|
||||
return "Usage: ds4-eval [options]";
|
||||
}
|
||||
return "Usage: ds4 [options]";
|
||||
}
|
||||
|
||||
static const char *tool_summary(ds4_help_tool tool) {
|
||||
switch (tool) {
|
||||
case DS4_HELP_DS4:
|
||||
return "Chat with a local DwarfStar model, run one-shot prompts, inspect models, or coordinate distributed inference.";
|
||||
case DS4_HELP_SERVER:
|
||||
return "Serve one loaded DwarfStar model through OpenAI, Responses, Anthropic, and completion-compatible HTTP APIs.";
|
||||
case DS4_HELP_AGENT:
|
||||
return "Run the native terminal coding agent with live tools, session save/restore, and a responsive prompt while the model works.";
|
||||
case DS4_HELP_BENCH:
|
||||
return "Measure prefill, decode, context growth, and KV-cache size across repeatable context frontiers.";
|
||||
case DS4_HELP_EVAL:
|
||||
return "Run the built-in reasoning, math, science, and security evaluation harness with a live terminal UI.";
|
||||
}
|
||||
return "";
|
||||
}
|
||||
|
||||
static void print_model_runtime(FILE *fp, const help_colors *c,
|
||||
ds4_help_tool tool, bool full) {
|
||||
title(fp, c, "Model And Runtime");
|
||||
opt(fp, c, "-m, --model FILE", "GGUF model path. Default: ds4flash.gguf");
|
||||
#ifdef DS4_ROCM_BUILD
|
||||
opt(fp, c, "--metal | --rocm | --cpu", "Select the backend explicitly.");
|
||||
opt(fp, c, "--backend NAME", "Backend name: metal, rocm, or cpu.");
|
||||
#else
|
||||
opt(fp, c, "--metal | --cuda | --cpu", "Select the backend explicitly.");
|
||||
opt(fp, c, "--backend NAME", "Backend name: metal, cuda, or cpu.");
|
||||
#endif
|
||||
if (tool != DS4_HELP_BENCH) {
|
||||
opt(fp, c, "-c, --ctx N", "Allocated context tokens.");
|
||||
}
|
||||
if (tool == DS4_HELP_SERVER) {
|
||||
opt(fp, c, "-n, --tokens N", "Default max output tokens when clients omit a limit.");
|
||||
}
|
||||
opt(fp, c, "-t, --threads N", "CPU helper threads for host-side/reference work.");
|
||||
opt(fp, c, "--power N", "GPU duty-cycle target, 1..100. Default: 100");
|
||||
opt(fp, c, "--ssd-streaming", "Metal/CUDA/ROCm: opt in to SSD-backed model streaming instead of full residency.");
|
||||
opt(fp, c, "--ssd-streaming-cold", "SSD streaming: skip default popularity-based expert-cache preload.");
|
||||
opt(fp, c, "--ssd-streaming-cache-experts N|NGB", "SSD streaming: routed expert cache as expert count or GiB, e.g. 32GB. Metal/ROCm default: 80% working set minus non-routed weights; CUDA default: backend fixed cache.");
|
||||
opt(fp, c, "--ssd-streaming-preload-experts N", "SSD streaming: upfront popularity preload count. Default: auto hot seed capped at 4096; use --ssd-streaming-cold to skip.");
|
||||
opt(fp, c, "--simulate-used-memory NGB", "Diagnostic: lock N GiB before model load to simulate a smaller-memory machine.");
|
||||
opt(fp, c, "--prefill-chunk N", "Metal graph prefill chunk size. Default: auto (PRO long prompts use 8192; others use 4096).");
|
||||
if (full) {
|
||||
if (tool != DS4_HELP_BENCH) {
|
||||
opt(fp, c, "--mtp FILE", "Optional MTP support GGUF used for draft-token probes.");
|
||||
}
|
||||
if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) {
|
||||
opt(fp, c, "--mtp-draft N", "Maximum autoregressive MTP draft tokens. Default: 1");
|
||||
opt(fp, c, "--mtp-margin F", "Verifier confidence margin for fast MTP acceptance. Default: 3");
|
||||
}
|
||||
opt(fp, c, "--quality", "Prefer exact kernels where faster approximate paths exist.");
|
||||
opt(fp, c, "--warm-weights", "Touch mapped tensor pages at startup to reduce first-use stalls.");
|
||||
if (tool == DS4_HELP_DS4 || tool == DS4_HELP_BENCH) {
|
||||
opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON.");
|
||||
}
|
||||
}
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_sampling(FILE *fp, const help_colors *c, bool full) {
|
||||
title(fp, c, "Prompt And Sampling");
|
||||
opt(fp, c, "-n, --tokens N", "Maximum generated tokens.");
|
||||
opt(fp, c, "--temp F", "Sampling temperature. 0 is greedy/deterministic.");
|
||||
opt(fp, c, "--top-p F", "Nucleus sampling probability.");
|
||||
opt(fp, c, "--min-p F", "Keep tokens scoring at least F times the top token.");
|
||||
opt(fp, c, "--seed N", "Sampling seed for reproducible non-greedy runs.");
|
||||
opt(fp, c, "--think", "Use normal thinking mode.");
|
||||
opt(fp, c, "--think-max", "Use Think Max when context is large enough.");
|
||||
opt(fp, c, "--nothink", "Disable thinking and ask for direct replies.");
|
||||
if (full) {
|
||||
opt(fp, c, "-sys, --system TEXT", "System prompt. Empty string disables the default where supported.");
|
||||
opt(fp, c, "-p, --prompt TEXT", "One-shot prompt text.");
|
||||
opt(fp, c, "--prompt-file FILE", "Read one-shot prompt text from FILE.");
|
||||
}
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_steering(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Directional Steering");
|
||||
opt(fp, c, "--dir-steering-file FILE", "Load one f32 direction vector per layer.");
|
||||
opt(fp, c, "--dir-steering-ffn F", "Apply steering after FFN outputs. Default with file: 1");
|
||||
opt(fp, c, "--dir-steering-attn F", "Apply steering after attention outputs. Default: 0");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_distributed(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Distributed Inference");
|
||||
fputc('\n', fp);
|
||||
para(fp, c, "Distributed mode runs one logical session across several machines by assigning contiguous model layer ranges to workers. Workers own their layer slice and KV-cache shard; the coordinator owns the prompt, sampling loop, and client/API flow. Start workers first, then start the coordinator. The coordinator waits for a complete route and streams hidden states through the workers.");
|
||||
fputc('\n', fp);
|
||||
opt(fp, c, "--role ROLE", "Distributed role: coordinator or worker.");
|
||||
opt(fp, c, "--layers A:B", "Inclusive layer slice, e.g. 0:20 or 21:output.");
|
||||
opt(fp, c, "--listen HOST PORT", "Coordinator listen address; workers may use it for their data listener.");
|
||||
opt(fp, c, "--coordinator HOST PORT", "Coordinator address for --role worker.");
|
||||
opt(fp, c, "--dist-prefill-chunk N", "Coordinator prefill pipeline chunk size. Default: session cap.");
|
||||
opt(fp, c, "--dist-prefill-window N", "Max prefill chunks in flight. Default: workers+2, capped at 8.");
|
||||
opt(fp, c, "--dist-activation-bits N", "Hidden-state transport width: 32, 16, or 8. Default: 32");
|
||||
opt(fp, c, "--dist-replay-check", "Diagnostic: reset and replay prompt, then compare logits.");
|
||||
opt(fp, c, "--debug", "Print coordinator route/debug logs.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_cli_diagnostics(FILE *fp, const help_colors *c);
|
||||
|
||||
static void print_cli_specific(FILE *fp, const help_colors *c, bool full) {
|
||||
title(fp, c, "CLI Modes");
|
||||
opt(fp, c, "ds4", "Start the interactive prompt.");
|
||||
opt(fp, c, "ds4 -p TEXT", "Run one prompt and exit.");
|
||||
opt(fp, c, "ds4 --prompt-file FILE", "Run a long prompt from a file and exit.");
|
||||
fputc('\n', fp);
|
||||
if (full) {
|
||||
print_cli_diagnostics(fp, c);
|
||||
}
|
||||
}
|
||||
|
||||
static void print_cli_diagnostics(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Diagnostics And Data Collection");
|
||||
opt(fp, c, "--inspect", "Load the model and print a summary only.");
|
||||
opt(fp, c, "--dump-tokens", "Tokenize the prompt exactly as written, then exit.");
|
||||
opt(fp, c, "--dump-logits FILE", "Write full next-token logits as JSON.");
|
||||
opt(fp, c, "--dump-logprobs FILE", "Write greedy continuation top-logprobs as JSON.");
|
||||
opt(fp, c, "--logprobs-top-k N", "Alternatives stored by --dump-logprobs. Default: 20");
|
||||
opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON.");
|
||||
opt(fp, c, "--perplexity-file FILE", "Score raw text with teacher-forced NLL.");
|
||||
opt(fp, c, "--imatrix-dataset FILE", "Rendered prompt dataset for imatrix collection.");
|
||||
opt(fp, c, "--imatrix-out FILE", "Write llama-compatible routed-MoE imatrix .dat.");
|
||||
opt(fp, c, "--imatrix-max-prompts N", "Stop imatrix collection after N prompts.");
|
||||
opt(fp, c, "--imatrix-max-tokens N", "Stop imatrix collection after N prompt tokens.");
|
||||
opt(fp, c, "--head-test", "Run the output HC/logits head after the native slice.");
|
||||
opt(fp, c, "--first-token-test", "Run exact CPU whole-model pass for the first prompt token.");
|
||||
opt(fp, c, "--metal-graph-test", "Compare first GPU-resident graph stages with CPU.");
|
||||
opt(fp, c, "--metal-graph-full-test", "Run the GPU-resident self-token graph across all layers.");
|
||||
opt(fp, c, "--metal-graph-prompt-test", "Compare CPU and GPU graph logits for the full prompt.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_cli_commands(FILE *fp, const help_colors *c) {
|
||||
title_red(fp, c, "Interactive Commands");
|
||||
opt(fp, c, "/help", "Show interactive commands.");
|
||||
opt(fp, c, "/think, /think-max, /nothink", "Switch thinking mode.");
|
||||
opt(fp, c, "/ctx N", "Restart the interactive session with a new context size.");
|
||||
opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100.");
|
||||
opt(fp, c, "/read FILE", "Read FILE and submit it as the next user message.");
|
||||
opt(fp, c, "/quit, /exit", "Leave the prompt.");
|
||||
opt(fp, c, "Ctrl+C", "Stop current generation and return to ds4>.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_agent_specific(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Agent Options");
|
||||
opt(fp, c, "-p, --prompt TEXT", "Submit an initial prompt after startup.");
|
||||
opt(fp, c, "--non-interactive", "Run without TUI. With -p: one turn; without -p: repeated stdin prompts.");
|
||||
opt(fp, c, "-sys, --system TEXT", "Extra system prompt. Empty disables extra text.");
|
||||
opt(fp, c, "--trace FILE", "Write prompt, token, and DSML debug trace.");
|
||||
opt(fp, c, "--chdir DIR", "Change working directory before loading runtime assets.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_agent_sessions(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Agent Runtime Commands");
|
||||
opt(fp, c, "/save", "Save the current session in ~/.ds4/kvcache.");
|
||||
opt(fp, c, "/compact", "Compact the current session context now.");
|
||||
opt(fp, c, "/list", "List saved sessions, sorted by recent update time.");
|
||||
opt(fp, c, "/switch ID", "Load a saved session and show recent history.");
|
||||
opt(fp, c, "/del ID", "Delete a saved session.");
|
||||
opt(fp, c, "/strip ID", "Remove KV payload; the text history can be rebuilt later.");
|
||||
opt(fp, c, "/history [N]", "Show N recent user turns from the current session.");
|
||||
opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100.");
|
||||
opt(fp, c, "/new", "Start a fresh session from the system prompt.");
|
||||
opt(fp, c, "/quit, /exit", "Exit.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_server_api(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "HTTP API");
|
||||
opt(fp, c, "--host HOST", "Bind address. Default: 127.0.0.1");
|
||||
opt(fp, c, "--port N", "Bind port. Default: 8000");
|
||||
opt(fp, c, "--cors", "Add Access-Control-Allow-* headers for browser JS clients.");
|
||||
opt(fp, c, "--trace FILE", "Write prompts, cache decisions, output, and tool calls.");
|
||||
para(fp, c, "Endpoints: /v1/chat/completions, /v1/responses, /v1/completions, and /v1/messages.");
|
||||
para(fp, c, "Model endpoint aliases include deepseek-v4-flash and deepseek-v4-pro; both serve the loaded GGUF.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_server_thinking(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Server Thinking Defaults");
|
||||
para(fp, c, "DeepSeek-compatible chat requests default to high-effort thinking.");
|
||||
para(fp, c, "reasoning_effort=max or output_config.effort=max requests Think Max.");
|
||||
para(fp, c, "Think Max requires --ctx >= 393216; smaller contexts use high.");
|
||||
para(fp, c, "thinking={type:disabled}, think=false, or model=deepseek-chat selects non-thinking mode.");
|
||||
para(fp, c, "In thinking mode, client sampling knobs are ignored like the official API.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_kv_cache(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Disk KV Cache");
|
||||
opt(fp, c, "--kv-disk-dir DIR", "Enable disk KV checkpoints in DIR.");
|
||||
opt(fp, c, "--kv-disk-space-mb N", "Disk budget. Default when enabled: 4096");
|
||||
opt(fp, c, "--kv-cache-min-tokens N", "Do not save/load checkpoints shorter than N. Default: 512");
|
||||
opt(fp, c, "--kv-cache-cold-max-tokens N", "Save cold first prompts up to N tokens. 0 disables. Default: 30000");
|
||||
opt(fp, c, "--kv-cache-continued-interval-tokens N", "Save aligned continued frontiers. 0 disables. Default: 10000");
|
||||
opt(fp, c, "--kv-cache-boundary-trim-tokens N", "Trim tail tokens for cold boundary saves. Default: 32");
|
||||
opt(fp, c, "--kv-cache-boundary-align-tokens N", "Align cold boundary saves to this multiple. Default: 2048");
|
||||
opt(fp, c, "--kv-cache-reject-different-quant", "Reject checkpoints written with different routed-expert quantization.");
|
||||
opt(fp, c, "--disable-exact-dsml-tool-replay", "Disable exact sampled DSML tool replay map.");
|
||||
opt(fp, c, "--tool-memory-max-ids N", "Exact tool-call IDs kept in RAM. Default: 100000");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_bench_specific(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Benchmark Input");
|
||||
opt(fp, c, "--prompt-file FILE", "Raw benchmark text; token sequence is sliced at each frontier.");
|
||||
opt(fp, c, "--chat-prompt-file FILE", "Render FILE as one no-thinking chat user message.");
|
||||
opt(fp, c, "-sys, --system TEXT", "System prompt used only with --chat-prompt-file.");
|
||||
fputc('\n', fp);
|
||||
title(fp, c, "Benchmark Sweep");
|
||||
opt(fp, c, "--ctx-start N", "First measured frontier. Default: 2048");
|
||||
opt(fp, c, "--ctx-max N", "Last measured frontier. Default: 32768");
|
||||
opt(fp, c, "--ctx-alloc N", "Allocated context. Default: ctx-max + gen-tokens + 1");
|
||||
opt(fp, c, "--step-mul F", "Multiplicative step. Default: 1");
|
||||
opt(fp, c, "--step-incr N", "Linear step when --step-mul is 1. Default: 2048");
|
||||
opt(fp, c, "--gen-tokens N", "Greedy decode tokens per frontier. 0 for pure prefill. Default: 128");
|
||||
opt(fp, c, "--csv FILE", "Write CSV there instead of stdout.");
|
||||
opt(fp, c, "--dump-frontier-logits-dir DIR", "Write one full-logit JSON file per frontier.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_eval_specific(FILE *fp, const help_colors *c) {
|
||||
title(fp, c, "Evaluation");
|
||||
opt(fp, c, "-n, --tokens N", "Max generated tokens per question. Default: 16000");
|
||||
opt(fp, c, "--questions N", "Run only the first N embedded questions.");
|
||||
opt(fp, c, "--case-sequence LIST", "Run 1-based case numbers in this comma-separated order.");
|
||||
opt(fp, c, "--trace FILE", "Write questions, outputs, and grading decisions.");
|
||||
opt(fp, c, "--regrade-trace FILE", "Regrade a prior trace without loading the model.");
|
||||
opt(fp, c, "--soft-limit-reply-budget N", "Soft close thinking near the end of reply budget. Default: 1024");
|
||||
opt(fp, c, "--hard-limit-reply-budget N", "Force </think> with N tokens left. Default: 512");
|
||||
opt(fp, c, "--soft-limit-think-close-rank N", "Soft-close when </think> is in top N tokens. Default: 3");
|
||||
opt(fp, c, "--pause-ms N", "Pause after each result in the TTY UI. Default: 350");
|
||||
opt(fp, c, "--plain", "Disable split-screen ANSI UI.");
|
||||
opt(fp, c, "--self-test-extractors", "Run answer-extractor self-tests and exit.");
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static bool tool_has_topic(ds4_help_tool tool, const char *topic) {
|
||||
if (!topic) return true;
|
||||
if (streq(topic, "all")) return true;
|
||||
if (streq(topic, "runtime") || streq(topic, "distributed")) return true;
|
||||
if (streq(topic, "sampling"))
|
||||
return tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_EVAL;
|
||||
if (streq(topic, "steering"))
|
||||
return tool == DS4_HELP_DS4 || tool == DS4_HELP_SERVER || tool == DS4_HELP_AGENT;
|
||||
switch (tool) {
|
||||
case DS4_HELP_DS4:
|
||||
return streq(topic, "diagnostics") || streq(topic, "commands");
|
||||
case DS4_HELP_SERVER:
|
||||
return streq(topic, "api") || streq(topic, "kv-cache") || streq(topic, "thinking");
|
||||
case DS4_HELP_AGENT:
|
||||
return streq(topic, "sessions") || streq(topic, "commands") || streq(topic, "tools");
|
||||
case DS4_HELP_BENCH:
|
||||
return streq(topic, "benchmark");
|
||||
case DS4_HELP_EVAL:
|
||||
return streq(topic, "evaluation");
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static void more_line(FILE *fp, const help_colors *c, const char *label, const char *topic) {
|
||||
static const char *colors[] = {
|
||||
"\x1b[38;5;81m", "\x1b[38;5;114m", "\x1b[38;5;179m",
|
||||
"\x1b[38;5;141m", "\x1b[38;5;147m"
|
||||
};
|
||||
static size_t idx;
|
||||
const char *on = c->cyan ? colors[idx++ % (sizeof(colors) / sizeof(colors[0]))] : "";
|
||||
if (streq(label, "Interactive commands:") && c->red) on = c->red;
|
||||
const char *off = c->off ? c->off : "";
|
||||
fprintf(fp, " %s%-26s%s --help %s\n", on, label, off, topic);
|
||||
}
|
||||
|
||||
static void print_more_info(FILE *fp, const help_colors *c, ds4_help_tool tool) {
|
||||
title(fp, c, "More Info");
|
||||
more_line(fp, c, "Runtime full info:", "runtime");
|
||||
if (tool_has_topic(tool, "sampling"))
|
||||
more_line(fp, c, "Sampling full info:", "sampling");
|
||||
more_line(fp, c, "Distributed inference:", "distributed");
|
||||
if (tool_has_topic(tool, "steering"))
|
||||
more_line(fp, c, "Steering full info:", "steering");
|
||||
if (tool == DS4_HELP_DS4) {
|
||||
more_line(fp, c, "Interactive commands:", "commands");
|
||||
more_line(fp, c, "Diagnostics:", "diagnostics");
|
||||
} else if (tool == DS4_HELP_SERVER) {
|
||||
more_line(fp, c, "HTTP API:", "api");
|
||||
more_line(fp, c, "Disk KV cache:", "kv-cache");
|
||||
more_line(fp, c, "Thinking behavior:", "thinking");
|
||||
} else if (tool == DS4_HELP_AGENT) {
|
||||
more_line(fp, c, "Agent sessions:", "sessions");
|
||||
more_line(fp, c, "Agent commands:", "commands");
|
||||
more_line(fp, c, "Agent tool system:", "tools");
|
||||
} else if (tool == DS4_HELP_BENCH) {
|
||||
more_line(fp, c, "Benchmark sweep:", "benchmark");
|
||||
} else if (tool == DS4_HELP_EVAL) {
|
||||
more_line(fp, c, "Evaluation options:", "evaluation");
|
||||
}
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_examples(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) {
|
||||
title(fp, c, "Examples");
|
||||
if (topic_is(topic, "distributed")) {
|
||||
opt(fp, c, "worker", "./ds4 --role worker --layers 21:output --coordinator 192.168.0.181 9000 -m ds4flash.gguf");
|
||||
opt(fp, c, "coordinator", "./ds4 --role coordinator --layers 0:20 --listen 0.0.0.0 9000 -p \"Hello\" -m ds4flash.gguf");
|
||||
} else if (topic_is(topic, "runtime")) {
|
||||
if (tool == DS4_HELP_SERVER) {
|
||||
opt(fp, c, "Metal API", "./ds4-server -m ds4flash.gguf --metal --ctx 100000");
|
||||
opt(fp, c, "quiet API", "./ds4-server --power 60 --host 127.0.0.1 --port 8000");
|
||||
} else if (tool == DS4_HELP_AGENT) {
|
||||
opt(fp, c, "agent", "./ds4-agent -m ds4flash.gguf --ctx 100000");
|
||||
opt(fp, c, "quiet agent", "./ds4-agent --power 50");
|
||||
} else if (tool == DS4_HELP_BENCH) {
|
||||
opt(fp, c, "bench", "./ds4-bench --prompt-file long.txt --ctx-max 32768");
|
||||
opt(fp, c, "quiet bench", "./ds4-bench --prompt-file long.txt --power 70");
|
||||
} else if (tool == DS4_HELP_EVAL) {
|
||||
opt(fp, c, "eval", "./ds4-eval --questions 10 --ctx 100000");
|
||||
opt(fp, c, "CPU debug", "./ds4-eval --cpu --questions 1 --tokens 32");
|
||||
} else {
|
||||
opt(fp, c, "Metal", "./ds4 -m ds4flash.gguf --metal -c 100000");
|
||||
opt(fp, c, "quiet thermals", "./ds4 -p \"Summarize README\" --power 50");
|
||||
}
|
||||
} else if (topic_is(topic, "steering")) {
|
||||
opt(fp, c, "steer FFN", "./ds4 -p \"Write tersely\" --dir-steering-file dir.bin --dir-steering-ffn 0.8");
|
||||
} else if (tool == DS4_HELP_SERVER || topic_is(topic, "api") || topic_is(topic, "kv-cache")) {
|
||||
opt(fp, c, "local API", "./ds4-server --ctx 100000 --kv-disk-dir ~/.ds4/server-kv --kv-disk-space-mb 8192");
|
||||
opt(fp, c, "curl", "curl http://127.0.0.1:8000/v1/models");
|
||||
} else if (tool == DS4_HELP_AGENT || topic_is(topic, "sessions") || topic_is(topic, "tools")) {
|
||||
opt(fp, c, "interactive", "./ds4-agent");
|
||||
opt(fp, c, "one shot", "./ds4-agent --non-interactive -p \"Create /tmp/hello.c\"");
|
||||
} else if (tool == DS4_HELP_BENCH || topic_is(topic, "benchmark")) {
|
||||
opt(fp, c, "csv", "./ds4-bench --prompt-file long.txt --ctx-max 32768 --csv speed.csv");
|
||||
opt(fp, c, "prefill only", "./ds4-bench --prompt-file long.txt --gen-tokens 0");
|
||||
} else if (tool == DS4_HELP_EVAL || topic_is(topic, "evaluation")) {
|
||||
opt(fp, c, "first 10", "./ds4-eval --questions 10 --trace eval.trace");
|
||||
opt(fp, c, "plain", "./ds4-eval --plain --nothink --tokens 512");
|
||||
} else {
|
||||
opt(fp, c, "chat", "./ds4");
|
||||
opt(fp, c, "one shot", "./ds4 -p \"Explain mmap in C\"");
|
||||
opt(fp, c, "long prompt", "./ds4 --think-max --prompt-file prompt.txt --ctx 393216");
|
||||
}
|
||||
fputc('\n', fp);
|
||||
}
|
||||
|
||||
static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) {
|
||||
if (streq(topic, "all")) {
|
||||
print_model_runtime(fp, c, tool, true);
|
||||
if (tool_has_topic(tool, "sampling")) print_sampling(fp, c, true);
|
||||
if (tool_has_topic(tool, "steering")) print_steering(fp, c);
|
||||
print_distributed(fp, c);
|
||||
if (tool == DS4_HELP_DS4) {
|
||||
print_cli_specific(fp, c, true);
|
||||
print_cli_commands(fp, c);
|
||||
} else if (tool == DS4_HELP_SERVER) {
|
||||
print_server_api(fp, c);
|
||||
print_server_thinking(fp, c);
|
||||
print_kv_cache(fp, c);
|
||||
} else if (tool == DS4_HELP_AGENT) {
|
||||
print_agent_specific(fp, c);
|
||||
print_agent_sessions(fp, c);
|
||||
} else if (tool == DS4_HELP_BENCH) {
|
||||
print_bench_specific(fp, c);
|
||||
} else if (tool == DS4_HELP_EVAL) {
|
||||
print_eval_specific(fp, c);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (streq(topic, "runtime")) print_model_runtime(fp, c, tool, true);
|
||||
else if (streq(topic, "sampling")) print_sampling(fp, c, true);
|
||||
else if (streq(topic, "steering")) print_steering(fp, c);
|
||||
else if (streq(topic, "distributed")) print_distributed(fp, c);
|
||||
else if (tool == DS4_HELP_DS4 && streq(topic, "diagnostics")) print_cli_diagnostics(fp, c);
|
||||
else if (tool == DS4_HELP_DS4 && streq(topic, "commands")) print_cli_commands(fp, c);
|
||||
else if (tool == DS4_HELP_SERVER && streq(topic, "api")) print_server_api(fp, c);
|
||||
else if (tool == DS4_HELP_SERVER && streq(topic, "kv-cache")) print_kv_cache(fp, c);
|
||||
else if (tool == DS4_HELP_SERVER && streq(topic, "thinking")) print_server_thinking(fp, c);
|
||||
else if (tool == DS4_HELP_AGENT && streq(topic, "sessions")) print_agent_sessions(fp, c);
|
||||
else if (tool == DS4_HELP_AGENT && streq(topic, "commands")) print_agent_sessions(fp, c);
|
||||
else if (tool == DS4_HELP_AGENT && streq(topic, "tools")) {
|
||||
title(fp, c, "Agent Tool System");
|
||||
para(fp, c, "The agent can read, search, write, edit, run bash, and browse through Chrome-backed web tools.");
|
||||
para(fp, c, "Tool calls are emitted by the model as DSML and rendered live in the terminal.");
|
||||
para(fp, c, "Edit uses exact old/new replacement; [upto] can bridge a unique head and tail for large anchored edits.");
|
||||
fputc('\n', fp);
|
||||
} else if (tool == DS4_HELP_BENCH && streq(topic, "benchmark")) print_bench_specific(fp, c);
|
||||
else if (tool == DS4_HELP_EVAL && streq(topic, "evaluation")) print_eval_specific(fp, c);
|
||||
}
|
||||
|
||||
static void print_default(FILE *fp, const help_colors *c, ds4_help_tool tool) {
|
||||
print_model_runtime(fp, c, tool, false);
|
||||
|
||||
if (tool == DS4_HELP_DS4) {
|
||||
print_cli_specific(fp, c, true);
|
||||
print_sampling(fp, c, false);
|
||||
} else if (tool == DS4_HELP_SERVER) {
|
||||
print_server_api(fp, c);
|
||||
print_kv_cache(fp, c);
|
||||
} else if (tool == DS4_HELP_AGENT) {
|
||||
print_agent_specific(fp, c);
|
||||
print_agent_sessions(fp, c);
|
||||
} else if (tool == DS4_HELP_BENCH) {
|
||||
print_bench_specific(fp, c);
|
||||
} else if (tool == DS4_HELP_EVAL) {
|
||||
print_eval_specific(fp, c);
|
||||
}
|
||||
}
|
||||
|
||||
void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic) {
|
||||
help_colors c = help_make_colors(fp);
|
||||
if (topic && !tool_has_topic(tool, topic)) {
|
||||
fprintf(fp, "%s: unknown help topic '%s'\n\n", tool_name(tool), topic);
|
||||
topic = NULL;
|
||||
}
|
||||
|
||||
fprintf(fp, "%s%s%s\n", c.bright ? c.bright : "", tool_name(tool), c.off ? c.off : "");
|
||||
fprintf(fp, "%s\n\n", tool_summary(tool));
|
||||
fprintf(fp, "%s\n\n", tool_usage(tool));
|
||||
|
||||
if (topic) print_topic(fp, &c, tool, topic);
|
||||
else {
|
||||
print_default(fp, &c, tool);
|
||||
print_more_info(fp, &c, tool);
|
||||
}
|
||||
print_examples(fp, &c, tool, topic);
|
||||
}
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
#ifndef DS4_HELP_H
|
||||
#define DS4_HELP_H
|
||||
|
||||
#include <stdio.h>
|
||||
|
||||
typedef enum {
|
||||
DS4_HELP_DS4,
|
||||
DS4_HELP_SERVER,
|
||||
DS4_HELP_AGENT,
|
||||
DS4_HELP_BENCH,
|
||||
DS4_HELP_EVAL,
|
||||
} ds4_help_tool;
|
||||
|
||||
void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,77 @@
|
||||
__device__ __constant__ uint8_t cuda_ksigns_iq2xs[128] = {
|
||||
0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15,
|
||||
144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159,
|
||||
160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175,
|
||||
48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63,
|
||||
192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207,
|
||||
80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95,
|
||||
96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
|
||||
240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
|
||||
};
|
||||
|
||||
__device__ __constant__ uint64_t cuda_iq2xxs_grid[256] = {
|
||||
0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
|
||||
0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
|
||||
0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
|
||||
0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
|
||||
0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
|
||||
0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
|
||||
0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
|
||||
0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
|
||||
0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
|
||||
0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
|
||||
0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
|
||||
0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
|
||||
0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
|
||||
0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
|
||||
0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
|
||||
0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
|
||||
0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
|
||||
0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
|
||||
0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
|
||||
0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
|
||||
0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
|
||||
0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
|
||||
0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
|
||||
0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
|
||||
0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
|
||||
0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
|
||||
0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
|
||||
0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
|
||||
0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
|
||||
0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
|
||||
0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
|
||||
0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
|
||||
0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
|
||||
0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
|
||||
0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
|
||||
0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
|
||||
0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
|
||||
0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
|
||||
0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
|
||||
0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
|
||||
0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
|
||||
0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
|
||||
0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
|
||||
0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
|
||||
0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
|
||||
0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
|
||||
0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
|
||||
0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
|
||||
0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
|
||||
0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
|
||||
0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
|
||||
0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
|
||||
0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
|
||||
0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
|
||||
0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
|
||||
0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
|
||||
0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
|
||||
0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
|
||||
0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
|
||||
0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
|
||||
0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
|
||||
0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
|
||||
0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
|
||||
0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
|
||||
};
|
||||
+1359
File diff suppressed because it is too large
Load Diff
+218
@@ -0,0 +1,218 @@
|
||||
#ifndef DS4_KVSTORE_H
|
||||
#define DS4_KVSTORE_H
|
||||
|
||||
#include "ds4.h"
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
|
||||
#define DS4_KVSTORE_FIXED_HEADER 48u
|
||||
#define DS4_KVSTORE_DEFAULT_MB 4096
|
||||
#define DS4_KVSTORE_HIT_HALF_LIFE_SECONDS (6ull * 60ull * 60ull)
|
||||
|
||||
#define DS4_KVSTORE_EXT_TOOL_MAP (1u << 0)
|
||||
#define DS4_KVSTORE_EXT_RESPONSES_VISIBLE (1u << 1)
|
||||
#define DS4_KVSTORE_EXT_THINKING_VISIBLE (1u << 2)
|
||||
#define DS4_KVSTORE_EXT_SESSION_TITLE (1u << 3)
|
||||
|
||||
typedef enum {
|
||||
DS4_KVSTORE_REASON_UNKNOWN = 0,
|
||||
DS4_KVSTORE_REASON_COLD = 1,
|
||||
DS4_KVSTORE_REASON_CONTINUED = 2,
|
||||
DS4_KVSTORE_REASON_EVICT = 3,
|
||||
DS4_KVSTORE_REASON_SHUTDOWN = 4,
|
||||
DS4_KVSTORE_REASON_AGENT_SYSTEM = 5,
|
||||
DS4_KVSTORE_REASON_AGENT_SESSION = 6,
|
||||
} ds4_kvstore_reason;
|
||||
|
||||
typedef enum {
|
||||
DS4_KVSTORE_LOG_DEFAULT,
|
||||
DS4_KVSTORE_LOG_KVCACHE,
|
||||
DS4_KVSTORE_LOG_WARNING,
|
||||
} ds4_kvstore_log_type;
|
||||
|
||||
typedef struct {
|
||||
/* The file name is the rendered byte prefix, not the token sequence. The
|
||||
* payload still carries the exact tokens and graph state; the hash only
|
||||
* answers "does this checkpoint represent the bytes at the front of the
|
||||
* incoming prompt?" */
|
||||
char sha[41];
|
||||
char *path;
|
||||
uint8_t quant_bits;
|
||||
/* Stored in header byte 7. Flash is 0 for backward compatibility with
|
||||
* older cache files where this reserved byte was always written as zero. */
|
||||
uint8_t model_id;
|
||||
uint8_t reason;
|
||||
uint32_t tokens;
|
||||
uint32_t hits;
|
||||
uint32_t ctx_size;
|
||||
uint8_t ext_flags;
|
||||
uint64_t created_at;
|
||||
uint64_t last_used;
|
||||
uint64_t payload_bytes;
|
||||
uint64_t text_bytes;
|
||||
uint64_t file_size;
|
||||
} ds4_kvstore_entry;
|
||||
|
||||
typedef struct {
|
||||
int min_tokens;
|
||||
int cold_max_tokens;
|
||||
int continued_interval_tokens;
|
||||
int boundary_trim_tokens;
|
||||
int boundary_align_tokens;
|
||||
} ds4_kvstore_options;
|
||||
|
||||
typedef struct {
|
||||
bool enabled;
|
||||
char *dir;
|
||||
uint64_t budget_bytes;
|
||||
bool reject_different_quant;
|
||||
ds4_kvstore_options opt;
|
||||
int continued_last_store_tokens;
|
||||
ds4_kvstore_entry *entry;
|
||||
int len;
|
||||
int cap;
|
||||
const char *log_name;
|
||||
void *log_ud;
|
||||
void (*log)(void *ud, ds4_kvstore_log_type type, const char *msg);
|
||||
} ds4_kvstore;
|
||||
|
||||
typedef struct {
|
||||
const char *text;
|
||||
size_t text_len;
|
||||
uint8_t model_id;
|
||||
uint8_t quant_bits;
|
||||
uint32_t ctx_size;
|
||||
bool reject_different_quant;
|
||||
} ds4_kvstore_eviction_context;
|
||||
|
||||
typedef struct {
|
||||
void *ud;
|
||||
uint8_t ext_flag;
|
||||
bool (*serialized_size)(void *ud, const char *text, uint64_t *bytes_out);
|
||||
bool (*write)(void *ud, FILE *fp, const char *text, uint64_t *written_bytes);
|
||||
int (*load)(void *ud, FILE *fp, const void *wanted);
|
||||
const void *load_wanted;
|
||||
} ds4_kvstore_trailer_hooks;
|
||||
|
||||
typedef struct {
|
||||
int tokens;
|
||||
uint32_t text_bytes;
|
||||
uint8_t quant_bits;
|
||||
uint8_t ext_flags;
|
||||
double load_ms;
|
||||
bool consumed;
|
||||
char *path;
|
||||
} ds4_kvstore_load_result;
|
||||
|
||||
ds4_kvstore_options ds4_kvstore_default_options(void);
|
||||
uint8_t ds4_kvstore_reason_code(const char *reason);
|
||||
const char *ds4_kvstore_key_kind(uint8_t ext_flags);
|
||||
|
||||
bool ds4_kvstore_open(ds4_kvstore *kc, const char *dir, uint64_t budget_mb,
|
||||
bool reject_different_quant, ds4_kvstore_options opt,
|
||||
const char *log_name,
|
||||
void (*log)(void *ud, ds4_kvstore_log_type type, const char *msg),
|
||||
void *log_ud);
|
||||
void ds4_kvstore_close(ds4_kvstore *kc);
|
||||
void ds4_kvstore_clear(ds4_kvstore *kc);
|
||||
void ds4_kvstore_entry_free(ds4_kvstore_entry *e);
|
||||
|
||||
char *ds4_kvstore_render_tokens_text(ds4_engine *engine,
|
||||
const ds4_tokens *tokens,
|
||||
size_t *out_len);
|
||||
bool ds4_kvstore_byte_prefix_match(const char *text, size_t text_len,
|
||||
const char *prefix, size_t prefix_len);
|
||||
void ds4_kvstore_tokens_copy_prefix(ds4_tokens *dst, const ds4_tokens *src, int n);
|
||||
void ds4_kvstore_build_prompt_from_exact_prefix_and_text_suffix(
|
||||
ds4_engine *engine,
|
||||
const ds4_tokens *exact_prefix,
|
||||
const char *suffix_text,
|
||||
ds4_tokens *out);
|
||||
|
||||
int ds4_kvstore_store_len(const ds4_kvstore *kc, int tokens);
|
||||
int ds4_kvstore_chat_anchor_pos(const ds4_kvstore *kc,
|
||||
const ds4_tokens *prompt,
|
||||
int user_token_id,
|
||||
int assistant_token_id);
|
||||
int ds4_kvstore_continued_store_target(const ds4_kvstore *kc, int live_tokens);
|
||||
void ds4_kvstore_note_store(ds4_kvstore *kc, int tokens);
|
||||
int ds4_kvstore_suppress_continued_store(ds4_kvstore *kc, int tokens);
|
||||
void ds4_kvstore_restore_suppressed_continued(ds4_kvstore *kc,
|
||||
int old_tokens,
|
||||
int suppressed_tokens);
|
||||
|
||||
bool ds4_kvstore_file_size_fits(const ds4_kvstore *kc,
|
||||
uint64_t text_bytes,
|
||||
uint64_t payload_bytes,
|
||||
uint64_t trailer_bytes,
|
||||
uint64_t *file_bytes_out,
|
||||
uint64_t *required_bytes_out);
|
||||
double ds4_kvstore_entry_eviction_score(const ds4_kvstore_entry *e,
|
||||
const ds4_tokens *live,
|
||||
uint64_t now,
|
||||
const ds4_kvstore_eviction_context *incoming);
|
||||
void ds4_kvstore_evict(ds4_kvstore *kc, const ds4_tokens *live,
|
||||
uint64_t extra_bytes,
|
||||
const ds4_kvstore_eviction_context *incoming);
|
||||
int ds4_kvstore_find_text_prefix(ds4_kvstore *kc, const char *prompt_text,
|
||||
int model_id, int quant_bits, int ctx_size);
|
||||
|
||||
bool ds4_kvstore_store_live_prefix_text(ds4_kvstore *kc,
|
||||
ds4_engine *engine,
|
||||
ds4_session *session,
|
||||
const ds4_tokens *tokens,
|
||||
int store_len,
|
||||
const char *reason,
|
||||
const char *cache_text_override,
|
||||
uint8_t cache_text_ext,
|
||||
const char *cache_text_key,
|
||||
const ds4_kvstore_trailer_hooks *hooks,
|
||||
char *err,
|
||||
size_t err_len);
|
||||
bool ds4_kvstore_store_live_prefix(ds4_kvstore *kc,
|
||||
ds4_engine *engine,
|
||||
ds4_session *session,
|
||||
const ds4_tokens *tokens,
|
||||
int store_len,
|
||||
const char *reason,
|
||||
const ds4_kvstore_trailer_hooks *hooks,
|
||||
char *err,
|
||||
size_t err_len);
|
||||
bool ds4_kvstore_maybe_store_continued(ds4_kvstore *kc,
|
||||
ds4_engine *engine,
|
||||
ds4_session *session,
|
||||
const ds4_kvstore_trailer_hooks *hooks,
|
||||
char *err,
|
||||
size_t err_len);
|
||||
int ds4_kvstore_try_load_text(ds4_kvstore *kc,
|
||||
ds4_engine *engine,
|
||||
ds4_session *session,
|
||||
const char *prompt_text,
|
||||
ds4_tokens *effective_prompt,
|
||||
ds4_kvstore_load_result *result,
|
||||
const ds4_kvstore_trailer_hooks *hooks,
|
||||
bool responses_protocol);
|
||||
void ds4_kvstore_load_result_free(ds4_kvstore_load_result *result);
|
||||
|
||||
bool ds4_kvstore_read_header(FILE *fp, ds4_kvstore_entry *e,
|
||||
uint32_t *text_bytes);
|
||||
bool ds4_kvstore_read_entry_file(const char *path, const char sha[41],
|
||||
ds4_kvstore_entry *out);
|
||||
void ds4_kvstore_fill_header(uint8_t h[DS4_KVSTORE_FIXED_HEADER],
|
||||
uint8_t model_id, uint8_t quant_bits,
|
||||
uint8_t reason, uint8_t ext_flags,
|
||||
uint32_t tokens, uint32_t hits, uint32_t ctx_size,
|
||||
uint64_t created_at, uint64_t last_used,
|
||||
uint64_t payload_bytes);
|
||||
bool ds4_kvstore_touch_file(const char *path, uint32_t hits);
|
||||
bool ds4_kvstore_sha_hex_name(const char *name, char sha[41]);
|
||||
void ds4_kvstore_sha1_bytes_hex(const void *ptr, size_t len, char out[41]);
|
||||
char *ds4_kvstore_path_join(const char *dir, const char *name);
|
||||
char *ds4_kvstore_path_for_sha(ds4_kvstore *kc, const char sha[41]);
|
||||
void ds4_kvstore_le_put32(uint8_t *p, uint32_t v);
|
||||
uint32_t ds4_kvstore_le_get32(const uint8_t *p);
|
||||
|
||||
#endif
|
||||
+26819
File diff suppressed because it is too large
Load Diff
+131
@@ -0,0 +1,131 @@
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
#include "ds4_rocm.h"
|
||||
#include <hipblaslt/hipblaslt.h>
|
||||
|
||||
#define FULL_WARP_MASK 0xFFFFFFFFFFFFFFFFULL
|
||||
#define MASK_T uint64_t
|
||||
#define DS4_GPU_BACKEND_NAME "ROCm"
|
||||
#define DS4_GPU_LOG_PREFIX "ds4: ROCm "
|
||||
#define DS4_GPU_BLAS_NAME "hipBLAS"
|
||||
#else
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <mma.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cub/block/block_radix_sort.cuh>
|
||||
|
||||
#define FULL_WARP_MASK 0xFFFFFFFFu
|
||||
#define MASK_T uint32_t
|
||||
#define DS4_GPU_BACKEND_NAME "CUDA"
|
||||
#define DS4_GPU_LOG_PREFIX "ds4: CUDA "
|
||||
#define DS4_GPU_BLAS_NAME "cuBLAS"
|
||||
#endif
|
||||
|
||||
#include <stdint.h>
|
||||
#include <ctype.h>
|
||||
#include <errno.h>
|
||||
#include <limits.h>
|
||||
#include <math.h>
|
||||
#include <fcntl.h>
|
||||
#include <pthread.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <sys/stat.h>
|
||||
#include <time.h>
|
||||
#include <unistd.h>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "ds4_gpu.h"
|
||||
|
||||
#ifndef M_PI
|
||||
#define M_PI 3.14159265358979323846
|
||||
#endif
|
||||
|
||||
#define CUDA_QK_K 256
|
||||
#define DS4_ROCM_UNUSED __attribute__((unused))
|
||||
|
||||
enum {
|
||||
/* attention_decode_mixed_kernel stores raw-window scores plus visible
|
||||
* compressed scores in shared memory. The host routes larger unmasked
|
||||
* decode calls to the online attention kernel so this fixed buffer never
|
||||
* becomes an out-of-bounds write at long context. */
|
||||
DS4_ROCM_ATTENTION_SCORE_CAP = 8192u,
|
||||
DS4_ROCM_ATTENTION_RAW_SCORE_CAP = 256u,
|
||||
DS4_ROCM_TOPK_MERGE_GROUP = 8u
|
||||
};
|
||||
|
||||
struct ds4_gpu_tensor {
|
||||
void *ptr;
|
||||
uint64_t bytes;
|
||||
int owner;
|
||||
};
|
||||
|
||||
typedef struct {
|
||||
uint8_t scales[CUDA_QK_K / 16];
|
||||
uint8_t qs[CUDA_QK_K / 4];
|
||||
uint16_t d;
|
||||
uint16_t dmin;
|
||||
} cuda_block_q2_K;
|
||||
|
||||
typedef struct {
|
||||
uint16_t d;
|
||||
uint16_t dmin;
|
||||
uint8_t scales[12];
|
||||
uint8_t qs[CUDA_QK_K / 2];
|
||||
} cuda_block_q4_K;
|
||||
|
||||
typedef struct {
|
||||
float d;
|
||||
int8_t qs[CUDA_QK_K];
|
||||
int16_t bsums[CUDA_QK_K / 16];
|
||||
} cuda_block_q8_K;
|
||||
|
||||
typedef struct {
|
||||
uint16_t d;
|
||||
uint16_t qs[CUDA_QK_K / 8];
|
||||
} cuda_block_iq2_xxs;
|
||||
|
||||
#include "ds4_iq2_tables_cuda.inc"
|
||||
|
||||
#include "rocm/ds4_rocm_runtime.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_common.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_q8.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_norm_rope.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_fp8_kv.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_attention.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_hc.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_output.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_indexer.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_embedding_launch.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_matmul.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_fp8_kv_launch.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_compressor.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_attention_launch.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_shared_expert.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_misc_launch.cuh"
|
||||
#include "rocm/ds4_rocm_router.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_moe.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_moe_launch.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_hc_output_launch.cuh"
|
||||
|
||||
#include "rocm/ds4_rocm_current_api_compat.cuh"
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
#pragma once
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hipblas/hipblas.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
#include <hipcub/hipcub.hpp>
|
||||
#include <rocwmma/rocwmma-version.hpp>
|
||||
#include <rocwmma/rocwmma.hpp>
|
||||
|
||||
#define cudaError_t hipError_t
|
||||
#define cudaStream_t hipStream_t
|
||||
#define cudaEvent_t hipEvent_t
|
||||
#define cudaDeviceProp hipDeviceProp_t
|
||||
#define cudaMemLocation hipMemLocation
|
||||
|
||||
#define cudaSuccess hipSuccess
|
||||
#define cudaErrorNotSupported hipErrorNotSupported
|
||||
#define cudaErrorInvalidValue hipErrorInvalidValue
|
||||
#define cudaGetLastError hipGetLastError
|
||||
#define cudaGetErrorString hipGetErrorString
|
||||
|
||||
#define cudaGetDevice hipGetDevice
|
||||
#define cudaSetDevice hipSetDevice
|
||||
#define cudaDeviceSynchronize hipDeviceSynchronize
|
||||
#define cudaDeviceGetAttribute hipDeviceGetAttribute
|
||||
#define cudaGetDeviceProperties hipGetDeviceProperties
|
||||
#define cudaDevAttrPageableMemoryAccess hipDeviceAttributePageableMemoryAccess
|
||||
#define cudaDevAttrMaxSharedMemoryPerBlockOptin hipDeviceAttributeSharedMemPerBlockOptin
|
||||
#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
|
||||
#define cudaFuncSetAttribute(func, attr, value) hipFuncSetAttribute((const void *)(func), (attr), (value))
|
||||
#define cudaMemLocationTypeDevice hipMemLocationTypeDevice
|
||||
|
||||
#define cudaMalloc hipMalloc
|
||||
#define cudaMallocHost hipHostMalloc
|
||||
#define cudaMallocManaged hipMallocManaged
|
||||
#define cudaFree hipFree
|
||||
#define cudaFreeHost hipFreeHost
|
||||
#define cudaMemset hipMemset
|
||||
#define cudaMemcpy hipMemcpy
|
||||
#define cudaMemcpyAsync hipMemcpyAsync
|
||||
#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
|
||||
#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
|
||||
#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice
|
||||
#define cudaMemGetInfo hipMemGetInfo
|
||||
#define cudaMemsetAsync hipMemsetAsync
|
||||
|
||||
#define cudaHostRegister hipHostRegister
|
||||
#define cudaHostUnregister hipHostUnregister
|
||||
#define cudaHostGetDevicePointer hipHostGetDevicePointer
|
||||
#define cudaHostRegisterMapped hipHostRegisterMapped
|
||||
#define cudaHostRegisterReadOnly hipHostRegisterReadOnly
|
||||
|
||||
#define cudaMemAdvise(p1, p2, p3, p4) hipMemAdvise(p1, p2, p3, p4.id)
|
||||
#define cudaMemPrefetchAsync(devPtr, count, location, flags, stream) hipMemPrefetchAsync(devPtr, count, location.id, stream)
|
||||
#define cudaMemAdviseSetReadMostly hipMemAdviseSetReadMostly
|
||||
#define cudaMemAdviseSetPreferredLocation hipMemAdviseSetPreferredLocation
|
||||
|
||||
#define cudaStreamCreateWithFlags hipStreamCreateWithFlags
|
||||
#define cudaStreamSynchronize hipStreamSynchronize
|
||||
#define cudaStreamDestroy hipStreamDestroy
|
||||
#define cudaStreamNonBlocking hipStreamNonBlocking
|
||||
|
||||
#define cudaEventCreate hipEventCreate
|
||||
#define cudaEventCreateWithFlags hipEventCreateWithFlags
|
||||
#define cudaEventDestroy hipEventDestroy
|
||||
#define cudaEventRecord hipEventRecord
|
||||
#define cudaEventSynchronize hipEventSynchronize
|
||||
#define cudaEventElapsedTime hipEventElapsedTime
|
||||
#define cudaEventDisableTiming hipEventDisableTiming
|
||||
|
||||
#define cublasHandle_t hipblasHandle_t
|
||||
#define cublasStatus_t hipblasStatus_t
|
||||
#define cublasMath_t hipblasMath_t
|
||||
|
||||
#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
|
||||
#define CUBLAS_OP_N HIPBLAS_OP_N
|
||||
#define CUBLAS_OP_T HIPBLAS_OP_T
|
||||
#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
|
||||
#define CUBLAS_DEFAULT_MATH HIPBLAS_DEFAULT_MATH
|
||||
#define CUBLAS_COMPUTE_32F HIPBLAS_COMPUTE_32F
|
||||
#define CUBLAS_TF32_TENSOR_OP_MATH HIPBLAS_TF32_TENSOR_OP_MATH
|
||||
#define CUDA_R_16F HIPBLAS_R_16F
|
||||
#define CUDA_R_32F HIPBLAS_R_32F
|
||||
|
||||
#define cublasCreate hipblasCreate
|
||||
#define cublasDestroy hipblasDestroy
|
||||
#define cublasSetMathMode hipblasSetMathMode
|
||||
#define cublasSgemm hipblasSgemm
|
||||
#define cublasSgemmStridedBatched hipblasSgemmStridedBatched
|
||||
#define cublasGemmEx hipblasGemmEx
|
||||
#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx
|
||||
|
||||
namespace cub = hipcub;
|
||||
|
||||
static __device__ __forceinline__ int32_t __vcmpne4(uint32_t a, uint32_t b) {
|
||||
// For each byte: 0xFF if a != b, 0x00 if a == b
|
||||
uint32_t diff = a ^ b;
|
||||
// Spread any set bit in each byte to fill the whole byte
|
||||
diff |= (diff >> 1); diff |= (diff >> 2); diff |= (diff >> 4);
|
||||
diff &= 0x01010101u;
|
||||
diff *= 0xFFu; // 0x01 -> 0xFF per byte
|
||||
return (int32_t)diff;
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ int32_t __vsub4(int32_t a, int32_t b) {
|
||||
// Per-byte subtraction (wrapping, not saturating)
|
||||
uint32_t ua = (uint32_t)a, ub = (uint32_t)b;
|
||||
// Trick: subtract bytes in parallel avoiding cross-byte borrows
|
||||
uint32_t diff = ((ua | 0x80808080u) - (ub & 0x7F7F7F7Fu)) ^ ((ua ^ ~ub) & 0x80808080u);
|
||||
return (int32_t)diff;
|
||||
}
|
||||
|
||||
// __dp4a: dot product of 4 signed int8s packed in an int32.
|
||||
// gfx11-class AMD GPUs expose this as a single v_dot4_i32_i8 instruction;
|
||||
// using the clang builtin avoids expanding every Q8/Q8_K dot into scalar byte
|
||||
// multiplies in the ROCm compatibility layer.
|
||||
static __device__ __forceinline__ int32_t __dp4a(int32_t a, int32_t b, int32_t c) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
union ds4_i8x4_bits { int32_t i; char4 v; } av, bv;
|
||||
av.i = a;
|
||||
bv.i = b;
|
||||
return amd_mixed_dot(av.v, bv.v, c, false);
|
||||
#else
|
||||
const int8_t *a_bytes = reinterpret_cast<const int8_t*>(&a);
|
||||
const int8_t *b_bytes = reinterpret_cast<const int8_t*>(&b);
|
||||
return c + (int32_t)a_bytes[0] * b_bytes[0]
|
||||
+ (int32_t)a_bytes[1] * b_bytes[1]
|
||||
+ (int32_t)a_bytes[2] * b_bytes[2]
|
||||
+ (int32_t)a_bytes[3] * b_bytes[3];
|
||||
#endif
|
||||
}
|
||||
|
||||
// Precise transcendentals for the MoE router top-k scores, immune to -fapprox-func.
|
||||
// These functions are to be used on paths where small error can be translated to
|
||||
// some macro effect - like expert selection kernels
|
||||
extern "C" __device__ __attribute__((pure)) float __ocml_exp_f32(float);
|
||||
extern "C" __device__ __attribute__((pure)) float __ocml_log1p_f32(float);
|
||||
extern "C" __device__ __attribute__((const)) float __ocml_sqrt_f32(float);
|
||||
|
||||
static __device__ __forceinline__ float ds4_precise_expf(float x) { return __ocml_exp_f32(x); }
|
||||
static __device__ __forceinline__ float ds4_precise_log1pf(float x) { return __ocml_log1p_f32(x); }
|
||||
static __device__ __forceinline__ float ds4_precise_sqrtf(float x) { return __ocml_sqrt_f32(x); }
|
||||
+15875
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,181 @@
|
||||
#include "ds4_ssd.h"
|
||||
|
||||
#include <ctype.h>
|
||||
#include <errno.h>
|
||||
#include <inttypes.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <sys/mman.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#ifndef MAP_ANONYMOUS
|
||||
#define MAP_ANONYMOUS MAP_ANON
|
||||
#endif
|
||||
|
||||
static const uint64_t DS4_GIB = 1024ull * 1024ull * 1024ull;
|
||||
|
||||
bool ds4_parse_gib_arg(const char *s, uint64_t *bytes) {
|
||||
if (bytes) *bytes = 0;
|
||||
if (!s || !s[0] || !bytes) return false;
|
||||
|
||||
size_t len = strlen(s);
|
||||
if (len > 2 &&
|
||||
(s[len - 2] == 'g' || s[len - 2] == 'G') &&
|
||||
(s[len - 1] == 'b' || s[len - 1] == 'B')) {
|
||||
len -= 2;
|
||||
}
|
||||
if (len == 0) return false;
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
if (!isdigit((unsigned char)s[i])) return false;
|
||||
}
|
||||
|
||||
char numbuf[32];
|
||||
if (len >= sizeof(numbuf)) return false;
|
||||
memcpy(numbuf, s, len);
|
||||
numbuf[len] = '\0';
|
||||
|
||||
errno = 0;
|
||||
unsigned long long v = strtoull(numbuf, NULL, 10);
|
||||
if (errno != 0 || v == 0 || v > UINT64_MAX / DS4_GIB) return false;
|
||||
|
||||
*bytes = (uint64_t)v * DS4_GIB;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ds4_parse_streaming_cache_experts_arg(const char *s,
|
||||
uint32_t *experts,
|
||||
uint64_t *bytes) {
|
||||
if (experts) *experts = 0;
|
||||
if (bytes) *bytes = 0;
|
||||
if (!s || !s[0] || !experts || !bytes) return false;
|
||||
|
||||
const size_t len = strlen(s);
|
||||
if (len > 2 &&
|
||||
(s[len - 2] == 'g' || s[len - 2] == 'G') &&
|
||||
(s[len - 1] == 'b' || s[len - 1] == 'B')) {
|
||||
return ds4_parse_gib_arg(s, bytes);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
if (!isdigit((unsigned char)s[i])) return false;
|
||||
}
|
||||
|
||||
errno = 0;
|
||||
unsigned long v = strtoul(s, NULL, 10);
|
||||
if (errno != 0 || v == 0 || v > UINT32_MAX) return false;
|
||||
|
||||
*experts = (uint32_t)v;
|
||||
return true;
|
||||
}
|
||||
|
||||
uint32_t ds4_ssd_cache_experts_for_byte_budget(uint64_t bytes,
|
||||
uint64_t per_expert_bytes) {
|
||||
if (bytes == 0 || per_expert_bytes == 0) return 0;
|
||||
const uint64_t experts = bytes / per_expert_bytes;
|
||||
if (experts == 0 || experts > UINT32_MAX) return 0;
|
||||
return (uint32_t)experts;
|
||||
}
|
||||
|
||||
bool ds4_ssd_auto_cache_plan(uint64_t recommended_bytes,
|
||||
uint64_t non_routed_bytes,
|
||||
uint64_t per_expert_bytes,
|
||||
uint64_t max_model_experts,
|
||||
ds4_ssd_cache_plan *out) {
|
||||
if (!out) return false;
|
||||
memset(out, 0, sizeof(*out));
|
||||
if (recommended_bytes == 0 || per_expert_bytes == 0) return false;
|
||||
|
||||
out->model_target_bytes =
|
||||
recommended_bytes > UINT64_MAX / 4ull ?
|
||||
UINT64_MAX : (recommended_bytes * 4ull) / 5ull;
|
||||
if (out->model_target_bytes > non_routed_bytes) {
|
||||
out->cache_bytes = out->model_target_bytes - non_routed_bytes;
|
||||
}
|
||||
|
||||
uint64_t cache_experts = out->cache_bytes / per_expert_bytes;
|
||||
if (cache_experts == 0) cache_experts = 1;
|
||||
if (max_model_experts != 0 && cache_experts > max_model_experts) {
|
||||
cache_experts = max_model_experts;
|
||||
}
|
||||
if (cache_experts > UINT32_MAX) cache_experts = UINT32_MAX;
|
||||
|
||||
out->cache_experts = (uint32_t)cache_experts;
|
||||
out->effective_cache_bytes = cache_experts * per_expert_bytes;
|
||||
return out->cache_experts != 0;
|
||||
}
|
||||
|
||||
bool ds4_ssd_memory_lock_acquire(ds4_ssd_memory_lock *lock,
|
||||
uint64_t bytes) {
|
||||
if (!lock) return false;
|
||||
lock->ptr = NULL;
|
||||
lock->bytes = 0;
|
||||
if (bytes == 0) return true;
|
||||
if (bytes > (uint64_t)SIZE_MAX) {
|
||||
fprintf(stderr,
|
||||
"ds4: --simulate-used-memory is too large for this process\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
void *ptr = mmap(NULL,
|
||||
(size_t)bytes,
|
||||
PROT_READ | PROT_WRITE,
|
||||
MAP_PRIVATE | MAP_ANONYMOUS,
|
||||
-1,
|
||||
0);
|
||||
if (ptr == MAP_FAILED) {
|
||||
fprintf(stderr,
|
||||
"ds4: --simulate-used-memory mmap %.2f GiB failed: %s\n",
|
||||
(double)bytes / (double)DS4_GIB,
|
||||
strerror(errno));
|
||||
return false;
|
||||
}
|
||||
|
||||
const long page_long = sysconf(_SC_PAGESIZE);
|
||||
const uint64_t page = page_long > 0 ? (uint64_t)page_long : 4096ull;
|
||||
const uint64_t chunk_bytes = 256ull * 1024ull * 1024ull;
|
||||
volatile unsigned char *p = (volatile unsigned char *)ptr;
|
||||
|
||||
/*
|
||||
* Touch and lock in bounded chunks. A single very large mlock() is harder
|
||||
* to diagnose when it fails and can create long uninterruptible VM work on
|
||||
* macOS; chunking mirrors the standalone diagnostic utility.
|
||||
*/
|
||||
uint64_t locked = 0;
|
||||
for (uint64_t off = 0; off < bytes; off += chunk_bytes) {
|
||||
uint64_t len = bytes - off;
|
||||
if (len > chunk_bytes) len = chunk_bytes;
|
||||
|
||||
for (uint64_t pos = off; pos < off + len; pos += page) {
|
||||
p[pos] = (unsigned char)(pos / page);
|
||||
}
|
||||
if (len != 0) p[off + len - 1u] = 1;
|
||||
|
||||
if (mlock((void *)(p + off), (size_t)len) != 0) {
|
||||
fprintf(stderr,
|
||||
"ds4: --simulate-used-memory mlock failed after %.2f/%.2f GiB: %s\n",
|
||||
(double)locked / (double)DS4_GIB,
|
||||
(double)bytes / (double)DS4_GIB,
|
||||
strerror(errno));
|
||||
if (locked != 0) munlock(ptr, (size_t)locked);
|
||||
munmap(ptr, (size_t)bytes);
|
||||
return false;
|
||||
}
|
||||
locked += len;
|
||||
}
|
||||
|
||||
lock->ptr = ptr;
|
||||
lock->bytes = bytes;
|
||||
fprintf(stderr,
|
||||
"ds4: simulated used memory: locked %.2f GiB before model load\n",
|
||||
(double)bytes / (double)DS4_GIB);
|
||||
return true;
|
||||
}
|
||||
|
||||
void ds4_ssd_memory_lock_release(ds4_ssd_memory_lock *lock) {
|
||||
if (!lock || !lock->ptr || lock->bytes == 0) return;
|
||||
munlock(lock->ptr, (size_t)lock->bytes);
|
||||
munmap(lock->ptr, (size_t)lock->bytes);
|
||||
lock->ptr = NULL;
|
||||
lock->bytes = 0;
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
#ifndef DS4_SSD_H
|
||||
#define DS4_SSD_H
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
|
||||
typedef struct {
|
||||
void *ptr;
|
||||
uint64_t bytes;
|
||||
} ds4_ssd_memory_lock;
|
||||
|
||||
typedef struct {
|
||||
uint64_t model_target_bytes;
|
||||
uint64_t cache_bytes;
|
||||
uint64_t effective_cache_bytes;
|
||||
uint32_t cache_experts;
|
||||
} ds4_ssd_cache_plan;
|
||||
|
||||
bool ds4_parse_gib_arg(const char *s, uint64_t *bytes);
|
||||
bool ds4_parse_streaming_cache_experts_arg(const char *s,
|
||||
uint32_t *experts,
|
||||
uint64_t *bytes);
|
||||
|
||||
uint32_t ds4_ssd_cache_experts_for_byte_budget(uint64_t bytes,
|
||||
uint64_t per_expert_bytes);
|
||||
bool ds4_ssd_auto_cache_plan(uint64_t recommended_bytes,
|
||||
uint64_t non_routed_bytes,
|
||||
uint64_t per_expert_bytes,
|
||||
uint64_t max_model_experts,
|
||||
ds4_ssd_cache_plan *out);
|
||||
|
||||
bool ds4_ssd_memory_lock_acquire(ds4_ssd_memory_lock *lock,
|
||||
uint64_t bytes);
|
||||
void ds4_ssd_memory_lock_release(ds4_ssd_memory_lock *lock);
|
||||
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,33 @@
|
||||
#ifndef DS4_WEB_H
|
||||
#define DS4_WEB_H
|
||||
|
||||
#include <stddef.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
typedef int (*ds4_web_confirm_fn)(void *privdata, const char *message,
|
||||
char *err, size_t err_len);
|
||||
typedef void (*ds4_web_log_fn)(void *privdata, const char *message);
|
||||
typedef bool (*ds4_web_cancel_fn)(void *privdata);
|
||||
|
||||
typedef struct {
|
||||
const char *home_dir;
|
||||
int port;
|
||||
ds4_web_confirm_fn confirm;
|
||||
void *confirm_privdata;
|
||||
ds4_web_log_fn log;
|
||||
void *log_privdata;
|
||||
ds4_web_cancel_fn cancel;
|
||||
void *cancel_privdata;
|
||||
} ds4_web_config;
|
||||
|
||||
typedef struct ds4_web ds4_web;
|
||||
|
||||
ds4_web *ds4_web_create(const ds4_web_config *cfg);
|
||||
void ds4_web_free(ds4_web *web);
|
||||
|
||||
char *ds4_web_google_search(ds4_web *web, const char *query,
|
||||
char *err, size_t err_len);
|
||||
char *ds4_web_visit_page(ds4_web *web, const char *url,
|
||||
char *err, size_t err_len);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,3 @@
|
||||
deepseek4-quantize
|
||||
quality-testing/score_official
|
||||
quality-testing/data/
|
||||
@@ -0,0 +1,40 @@
|
||||
CC ?= cc
|
||||
UNAME_S := $(shell uname -s)
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
NATIVE_CPU_FLAG ?= -mcpu=native
|
||||
QUALITY_LDLIBS := -lm -pthread -framework Foundation -framework Metal
|
||||
QUALITY_TARGETS := ds4.o ds4_metal.o
|
||||
QUALITY_OBJS := ../ds4.o ../ds4_metal.o
|
||||
QUALITY_LINK := $(CC) $(CFLAGS) -I.. -o
|
||||
else
|
||||
NATIVE_CPU_FLAG ?= -march=native
|
||||
CUDA_HOME ?= /usr/local/cuda
|
||||
NVCC ?= $(CUDA_HOME)/bin/nvcc
|
||||
CUDA_ARCH ?=
|
||||
ifneq ($(strip $(CUDA_ARCH)),)
|
||||
NVCC_ARCH_FLAGS := -arch=$(CUDA_ARCH)
|
||||
endif
|
||||
NVCCFLAGS ?= -O3 --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread
|
||||
QUALITY_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas
|
||||
QUALITY_TARGETS := ds4.o ds4_cuda.o
|
||||
QUALITY_OBJS := ../ds4.o ../ds4_cuda.o
|
||||
QUALITY_LINK := $(NVCC) $(NVCCFLAGS) -I.. -o
|
||||
endif
|
||||
|
||||
CFLAGS ?= -O3 -Wall -Wextra -std=c11 $(NATIVE_CPU_FLAG)
|
||||
CPPFLAGS ?= -D_GNU_SOURCE
|
||||
|
||||
.PHONY: all clean quality-score
|
||||
|
||||
all: deepseek4-quantize
|
||||
|
||||
deepseek4-quantize: deepseek4-quantize.c quants.c quants.h
|
||||
$(CC) $(CFLAGS) $(CPPFLAGS) -o $@ deepseek4-quantize.c quants.c -lm -pthread
|
||||
|
||||
quality-score:
|
||||
$(MAKE) -C .. $(QUALITY_TARGETS)
|
||||
$(QUALITY_LINK) quality-testing/score_official quality-testing/score_official.c $(QUALITY_OBJS) $(QUALITY_LDLIBS)
|
||||
|
||||
clean:
|
||||
rm -f deepseek4-quantize quality-testing/score_official
|
||||
@@ -0,0 +1,134 @@
|
||||
# DS4 GGUF Tools
|
||||
|
||||
This directory contains the offline tools used to build and evaluate DeepSeek
|
||||
V4 Flash GGUF files for `ds4`.
|
||||
|
||||
The important pieces are:
|
||||
|
||||
- `deepseek4-quantize.c`: C HF-safetensors to GGUF quantizer.
|
||||
- `quants.[ch]`: the deliberately small local quantization implementation used
|
||||
by the quantizer. It implements the DS4 output formats we actually ship:
|
||||
`q8_0`, `q4_K`, `q2_K`, and `iq2_xxs`.
|
||||
- `imatrix/`: dataset and instructions for collecting routed-MoE activation
|
||||
importance with `ds4`.
|
||||
- `quality-testing/`: prompts and scripts used to compare local GGUF variants
|
||||
against official DeepSeek V4 Flash continuations.
|
||||
|
||||
## Build
|
||||
|
||||
```sh
|
||||
make -C gguf-tools
|
||||
```
|
||||
|
||||
The quantizer is plain C and does not link GGML. GGUF metadata handling,
|
||||
safetensors loading, FP4/FP8 dequantization, and the quantizers used by our Q2
|
||||
and Q4 recipes live in this directory.
|
||||
|
||||
## Generate An Imatrix
|
||||
|
||||
First regenerate or inspect the calibration dataset:
|
||||
|
||||
```sh
|
||||
python3 gguf-tools/imatrix/dataset/build_ds4_imatrix_dataset.py
|
||||
```
|
||||
|
||||
Then collect activation statistics with the DS4 runtime:
|
||||
|
||||
```sh
|
||||
./ds4 \
|
||||
-m gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2.gguf \
|
||||
--imatrix-dataset gguf-tools/imatrix/dataset/rendered_prompts.txt \
|
||||
--imatrix-out gguf/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4.dat \
|
||||
--ctx 32768
|
||||
```
|
||||
|
||||
The imatrix file is useful immediately with this DS4 quantizer. Generic GGUF
|
||||
tools need DS4-specific tensor-name mapping and per-expert slicing before they
|
||||
can use it correctly. The accepted imatrix format is the legacy llama.cpp
|
||||
binary `.dat` file emitted by `ds4 --imatrix-out`.
|
||||
|
||||
Generating this `.dat` file locally is possible, but slow: it runs the DS4
|
||||
prefill graph over the full calibration corpus and reads routed-MoE activation
|
||||
statistics back from the GPU. The latest published imatrix-generated GGUF files
|
||||
are available in the antirez Hugging Face repository:
|
||||
|
||||
```text
|
||||
https://huggingface.co/antirez/deepseek-v4-gguf/tree/main
|
||||
```
|
||||
|
||||
## Generate Q2 And Q4 GGUFs
|
||||
|
||||
The template GGUF supplies metadata, tokenizer, tensor order, and logical
|
||||
shapes. Tensor bytes are regenerated from the Hugging Face safetensors. Full
|
||||
generation is intentionally offline and heavy: expect roughly 80-90 GB outputs
|
||||
for the 2-bit template family and roughly 150-170 GB for the 4-bit routed-expert
|
||||
family, plus enough free disk for the temporary output. Use `--dry-run` and
|
||||
`--compare-tensor` before starting a full write, and use `--overwrite` only when
|
||||
you really mean to replace an existing GGUF.
|
||||
|
||||
Q2 routed experts with imatrix:
|
||||
|
||||
```sh
|
||||
gguf-tools/deepseek4-quantize \
|
||||
--hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash \
|
||||
--template gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2.gguf \
|
||||
--out gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf \
|
||||
--imatrix gguf/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4.dat
|
||||
```
|
||||
|
||||
Q4 routed experts with imatrix:
|
||||
|
||||
```sh
|
||||
gguf-tools/deepseek4-quantize \
|
||||
--hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash \
|
||||
--template gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2.gguf \
|
||||
--out gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf \
|
||||
--imatrix gguf/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4.dat
|
||||
```
|
||||
|
||||
You can override tensor families:
|
||||
|
||||
```sh
|
||||
--experts iq2_xxs
|
||||
--routed-w2 q2_k
|
||||
--attention-proj q8_0
|
||||
--shared q8_0
|
||||
--output q8_0
|
||||
```
|
||||
|
||||
Useful checks before writing a full model:
|
||||
|
||||
```sh
|
||||
gguf-tools/deepseek4-quantize \
|
||||
--hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash \
|
||||
--template MODEL.gguf \
|
||||
--compare-tensor blk.0.attn_q_a.weight
|
||||
```
|
||||
|
||||
`--compare-tensor` regenerates a single tensor and byte-compares it against the
|
||||
template or `--compare-gguf`. `--threads N` controls routed-expert workers.
|
||||
|
||||
## When No Imatrix Is Given
|
||||
|
||||
`iq2_xxs` requires an importance vector. If `--imatrix` is not provided and
|
||||
the target type requires one, `deepseek4-quantize` computes a synthetic fallback
|
||||
from the dequantized weight itself:
|
||||
|
||||
```text
|
||||
importance[column] = sum(row[column]^2) over all rows
|
||||
```
|
||||
|
||||
This is a weight-energy heuristic. It is not as good as measuring real DS4
|
||||
activations, but it gives the quantizer a stable column weighting and was good
|
||||
enough for the first working 2-bit GGUFs.
|
||||
|
||||
## Quality Testing
|
||||
|
||||
See `quality-testing/README.md`. The short version is:
|
||||
|
||||
```sh
|
||||
python3 gguf-tools/quality-testing/collect_official.py
|
||||
make -C gguf-tools quality-score
|
||||
gguf-tools/quality-testing/score_official MODEL.gguf gguf-tools/quality-testing/data/manifest.tsv /tmp/model.tsv 4096
|
||||
python3 gguf-tools/quality-testing/compare_scores.py /tmp/old.tsv /tmp/new.tsv
|
||||
```
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,188 @@
|
||||
# DS4 Imatrix Pipeline
|
||||
|
||||
This directory contains the calibration dataset and instructions used to build
|
||||
activation importance matrices for DeepSeek V4 Flash and Pro GGUF
|
||||
quantization.
|
||||
|
||||
The current imatrix target is the routed MoE path. Flash has 43 layers and
|
||||
256 routed experts per layer. Pro has 61 layers and 384 routed experts per
|
||||
layer. Both variants expose three routed expert tensors per layer:
|
||||
|
||||
- `blk.N.ffn_gate_exps.weight`
|
||||
- `blk.N.ffn_up_exps.weight`
|
||||
- `blk.N.ffn_down_exps.weight`
|
||||
|
||||
For gate/up tensors, the collector records the squared FFN-normalized input
|
||||
activation. For down tensors, it records the squared routed SwiGLU row after
|
||||
route weighting. The result tells the quantizer which input columns are used
|
||||
more heavily by the actual DS4 inference graph.
|
||||
|
||||
## 1. Build The Calibration Dataset
|
||||
|
||||
The tracked dataset is in `gguf-tools/imatrix/dataset/`. Regenerate it from
|
||||
the repository root with:
|
||||
|
||||
```sh
|
||||
python3 gguf-tools/imatrix/dataset/build_ds4_imatrix_dataset.py
|
||||
```
|
||||
|
||||
The important output is:
|
||||
|
||||
```text
|
||||
gguf-tools/imatrix/dataset/rendered_prompts.txt
|
||||
```
|
||||
|
||||
It contains DS4-rendered chat prompts, separated by visible
|
||||
`DS4_IMATRIX_PROMPT` markers. The prompts include:
|
||||
|
||||
- C/Metal source-review prompts from this repository.
|
||||
- Long-context snippets.
|
||||
- Agent/tool-call prompts using DS4's DSML syntax.
|
||||
- Language/prose rewriting, summarization, extraction, and translation prompts.
|
||||
- `ds4-eval` GPQA Diamond, SuperGPQA, and AIME2025 benchmark prompts.
|
||||
- Both thinking and non-thinking assistant prefixes.
|
||||
|
||||
The current tracked dataset has 4682 rendered prompts and roughly 2.91M tokens
|
||||
by the coarse bytes/4 estimate. Check
|
||||
`gguf-tools/imatrix/dataset/manifest.json` for the exact generated-file
|
||||
summary.
|
||||
|
||||
## 2. Collect The Imatrix
|
||||
|
||||
Use the DS4 runtime itself to collect routed MoE activation statistics. The
|
||||
collector uses the loaded GGUF metadata, so the same command shape works for
|
||||
Flash and Pro.
|
||||
|
||||
Flash example:
|
||||
|
||||
```sh
|
||||
./ds4 \
|
||||
-m ../deepseek-v4-quants/gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2.gguf \
|
||||
--imatrix-dataset gguf-tools/imatrix/dataset/rendered_prompts.txt \
|
||||
--imatrix-out ../deepseek-v4-quants/imatrix/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4-1p5m.dat \
|
||||
--ctx 32768
|
||||
```
|
||||
|
||||
Pro example with a smaller calibration budget:
|
||||
|
||||
```sh
|
||||
./ds4 \
|
||||
-m ../deepseek-v4-quants/gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct.gguf \
|
||||
--imatrix-dataset gguf-tools/imatrix/dataset/rendered_prompts.txt \
|
||||
--imatrix-out ../deepseek-v4-quants/imatrix/DeepSeek-V4-Pro-Instruct-routed-moe-ds4-small.dat \
|
||||
--imatrix-max-prompts 16 \
|
||||
--imatrix-max-tokens 32768 \
|
||||
--ctx 32768
|
||||
```
|
||||
|
||||
Useful smoke-test limits:
|
||||
|
||||
```sh
|
||||
./ds4 \
|
||||
-m MODEL.gguf \
|
||||
--imatrix-dataset gguf-tools/imatrix/dataset/rendered_prompts.txt \
|
||||
--imatrix-out /tmp/ds4-test.imatrix.dat \
|
||||
--imatrix-max-prompts 1 \
|
||||
--imatrix-max-tokens 4096
|
||||
```
|
||||
|
||||
The collector is Metal-only because it hooks the layer-major Metal prefill graph.
|
||||
It does not change inference math; it reads the already materialized MoE inputs
|
||||
and accumulates `sum(x[column]^2)` per routed expert.
|
||||
|
||||
The output format is llama.cpp's legacy binary `.dat` imatrix format. DS4 packs
|
||||
per-expert vectors into one entry per routed expert tensor:
|
||||
|
||||
```text
|
||||
entry length = n_expert * n_columns
|
||||
```
|
||||
|
||||
The quantizer slices the right expert's segment when quantizing each expert.
|
||||
For Pro, that means each routed tensor entry contains 384 vectors instead of
|
||||
Flash's 256.
|
||||
|
||||
## 3. Generate GGUF Files With The Imatrix
|
||||
|
||||
The local C quantization tool in `gguf-tools/` supports:
|
||||
|
||||
```text
|
||||
--imatrix FILE
|
||||
--imatrix-strict
|
||||
```
|
||||
|
||||
Example Q4 routed-expert regeneration:
|
||||
|
||||
```sh
|
||||
gguf-tools/deepseek4-quantize \
|
||||
--hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash \
|
||||
--template ../deepseek-v4-quants/gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2.gguf \
|
||||
--out ../deepseek-v4-quants/gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf \
|
||||
--imatrix ../deepseek-v4-quants/imatrix/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4-1p5m.dat
|
||||
```
|
||||
|
||||
Example Q2 regeneration:
|
||||
|
||||
```sh
|
||||
gguf-tools/deepseek4-quantize \
|
||||
--hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash \
|
||||
--template ../deepseek-v4-quants/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2.gguf \
|
||||
--out ../deepseek-v4-quants/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf \
|
||||
--imatrix ../deepseek-v4-quants/imatrix/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4-1p5m.dat
|
||||
```
|
||||
|
||||
Example Pro Q2 regeneration from a Pro template:
|
||||
|
||||
```sh
|
||||
gguf-tools/deepseek4-quantize \
|
||||
--hf ../deepseek-v4-quants/hf/DeepSeek-V4-Pro \
|
||||
--template ../deepseek-v4-quants/gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct.gguf \
|
||||
--out ../deepseek-v4-quants/gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-small.gguf \
|
||||
--imatrix ../deepseek-v4-quants/imatrix/DeepSeek-V4-Pro-Instruct-routed-moe-ds4-small.dat \
|
||||
--threads 24
|
||||
```
|
||||
|
||||
`deepseek4-quantize` reads the routed expert count from the template GGUF. Use
|
||||
`--n-experts` only when working with an older template that lacks
|
||||
`deepseek4.expert_count`.
|
||||
|
||||
For Q4, the imatrix does not change the runtime tensor type: routed experts
|
||||
remain `Q4_K`. It changes how quantization error is weighted while choosing
|
||||
scales and codes. For Q2, it replaces the previous synthetic weight-energy
|
||||
fallback used for `IQ2_XXS` gate/up experts with real activation statistics.
|
||||
|
||||
## 4. Evaluate
|
||||
|
||||
Useful local tools:
|
||||
|
||||
```text
|
||||
misc/quant_eval.c
|
||||
gguf-tools/quality-testing/
|
||||
```
|
||||
|
||||
`misc/quant_eval.c` compares local GGUF variants by greedy/top-logit behavior.
|
||||
`gguf-tools/quality-testing/` can score local GGUFs against official DeepSeek
|
||||
API continuations by target-token negative log likelihood.
|
||||
|
||||
The Q4 imatrix file uploaded to Hugging Face was tested on 100 official
|
||||
DeepSeek V4 Flash continuations:
|
||||
|
||||
```text
|
||||
old Q4 avg NLL: 0.177357819
|
||||
Q4 imatrix avg NLL: 0.173895148
|
||||
relative NLL change: -1.95%
|
||||
case wins: 54 imatrix / 46 old
|
||||
first-token matches: 83 imatrix / 81 old
|
||||
avg greedy LCP: 12.21 imatrix / 11.94 old
|
||||
```
|
||||
|
||||
## Compatibility
|
||||
|
||||
The `.dat` file is intentionally in llama.cpp's legacy imatrix format, so the
|
||||
data is not conceptually tied to DS4. In practice, it is immediately useful
|
||||
only with a quantizer that understands DS4's tensor names and packed per-expert
|
||||
entries. The current `deepseek4-quantize` tooling does that.
|
||||
|
||||
Other GGUF creation tools can use the same imatrix if they implement the same
|
||||
name mapping and per-expert slicing convention. Without that DS4-specific
|
||||
mapping, a generic imatrix loader will see valid data but will not know how to
|
||||
apply the packed routed-expert vectors correctly.
|
||||
@@ -0,0 +1,37 @@
|
||||
# DS4 Imatrix Calibration Dataset
|
||||
|
||||
This directory contains DS4-rendered chat prompts for collecting activation
|
||||
statistics before building new low-bit GGUF files.
|
||||
|
||||
Run:
|
||||
|
||||
```sh
|
||||
python3 gguf-tools/imatrix/dataset/build_ds4_imatrix_dataset.py
|
||||
```
|
||||
|
||||
Generated files:
|
||||
|
||||
- `prompts.jsonl`: structured records with messages and rendered prompt text.
|
||||
- `rendered_prompts.txt`: all rendered prompts, separated by visible markers.
|
||||
- `rendered_prompts_nothink.txt`: only prompts ending with `</think>`.
|
||||
- `rendered_prompts_think.txt`: only prompts ending with `<think>`.
|
||||
- `manifest.json`: counts, byte totals, and rough token estimate.
|
||||
|
||||
The renderer mirrors the server prompt shape:
|
||||
|
||||
```text
|
||||
<｜begin▁of▁sentence｜>system<｜User｜>...<｜Assistant｜><think>
|
||||
<｜begin▁of▁sentence｜>system<｜User｜>...<｜Assistant｜></think>
|
||||
```
|
||||
|
||||
Some records include DSML tool schemas, sampled DSML tool calls, and tool-result
|
||||
turns so the imatrix sees the same special-token patterns used by agent clients.
|
||||
The corpus is provider-neutral and also includes language/prose rewriting,
|
||||
summarization, copy-editing, extraction, multilingual translation, programming
|
||||
prompts, Bash scripting, algorithm recall, `ds4-eval` benchmark-reasoning
|
||||
prompts, long-context code synthesis, agent transcript replay, log diagnosis,
|
||||
prose fact recovery, delayed-constraint and small needle tasks, Metal/C code
|
||||
review tasks, and inference-specific debugging tasks.
|
||||
|
||||
For normal imatrix collection, use `rendered_prompts.txt` so calibration covers
|
||||
both thinking and non-thinking modes. Split files are provided for ablations.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"version": 4,
|
||||
"purpose": "DeepSeek V4 Flash imatrix calibration prompts",
|
||||
"record_count": 4690,
|
||||
"rendered_utf8_bytes": 11670191,
|
||||
"rough_token_estimate_bytes_div_4": 2917547,
|
||||
"categories": {
|
||||
"agent": 1106,
|
||||
"algorithms": 32,
|
||||
"eval_reasoning": 150,
|
||||
"general": 40,
|
||||
"language": 1024,
|
||||
"long_context": 36,
|
||||
"programming": 48,
|
||||
"source": 2074,
|
||||
"translation": 180
|
||||
},
|
||||
"modes": {
|
||||
"nothink": 2345,
|
||||
"think": 2345
|
||||
},
|
||||
"bytes_by_category": {
|
||||
"agent": 4428393,
|
||||
"algorithms": 9388,
|
||||
"eval_reasoning": 144853,
|
||||
"general": 11902,
|
||||
"language": 725152,
|
||||
"long_context": 357114,
|
||||
"programming": 14320,
|
||||
"source": 5922919,
|
||||
"translation": 56150
|
||||
},
|
||||
"files": {
|
||||
"jsonl": "prompts.jsonl",
|
||||
"all_rendered": "rendered_prompts.txt",
|
||||
"nothink_rendered": "rendered_prompts_nothink.txt",
|
||||
"think_rendered": "rendered_prompts_think.txt"
|
||||
}
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,40 @@
|
||||
# Mixed GGUF Splicing
|
||||
|
||||
This directory contains a small GGUF splicer for experiments where only some
|
||||
routed-expert layers are upgraded from one quantization to another.
|
||||
|
||||
The tool does not dequantize or requantize weights. It reads two compatible
|
||||
GGUF files, uses the first one as the base, and copies selected routed-expert
|
||||
tensors from the second one. All other tensors remain byte-for-byte identical to
|
||||
the base file.
|
||||
|
||||
## Last Six Layers Q4 Experiment
|
||||
|
||||
The best mixed Flash experiment so far keeps the fixed-imatrix 2-bit GGUF as the
|
||||
base, then replaces only the routed experts in layers `37-42` with the same
|
||||
tensors from the fixed-imatrix Q4 GGUF.
|
||||
|
||||
This produces a file around 91 GB instead of the full Q4 file around 153 GB. In
|
||||
the local output-agreement checks, this last-six-layers file was statistically
|
||||
closer to the full Q4 GGUF than the plain 2-bit GGUF. This is not the same as a
|
||||
benchmark score, but it is a useful signal that the mixed file behaves more like
|
||||
the stronger Q4 model while keeping most of the 2-bit size.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
python3 gguf-tools/mixed/splice_mixed_expert_layers_gguf.py \
|
||||
--base /path/to/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf \
|
||||
--donor /path/to/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix-fixed.gguf \
|
||||
--q4-layers 37-42 \
|
||||
--out /path/to/DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf
|
||||
```
|
||||
|
||||
Use `--dry-run` first to check the selected tensors and expected size delta.
|
||||
Use `--force` only when you really want to overwrite the output file.
|
||||
|
||||
Layer ranges can be changed without editing the script:
|
||||
|
||||
```bash
|
||||
--q4-layers 0-2,40-42
|
||||
```
|
||||
+404
@@ -0,0 +1,404 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build a mixed DeepSeek V4 Flash GGUF by splicing routed-expert layers.
|
||||
|
||||
The base GGUF supplies metadata and all tensors by default. For selected layer
|
||||
IDs, this copies the routed expert tensors from a donor GGUF, rewriting the GGUF
|
||||
tensor directory and streaming tensor payloads without dequantizing or
|
||||
requantizing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import struct
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
|
||||
|
||||
GGUF_HEADER_SIZE = 24
|
||||
GGUF_DEFAULT_ALIGNMENT = 32
|
||||
|
||||
GGUF_VALUE_UINT8 = 0
|
||||
GGUF_VALUE_INT8 = 1
|
||||
GGUF_VALUE_UINT16 = 2
|
||||
GGUF_VALUE_INT16 = 3
|
||||
GGUF_VALUE_UINT32 = 4
|
||||
GGUF_VALUE_INT32 = 5
|
||||
GGUF_VALUE_FLOAT32 = 6
|
||||
GGUF_VALUE_BOOL = 7
|
||||
GGUF_VALUE_STRING = 8
|
||||
GGUF_VALUE_ARRAY = 9
|
||||
GGUF_VALUE_UINT64 = 10
|
||||
GGUF_VALUE_INT64 = 11
|
||||
GGUF_VALUE_FLOAT64 = 12
|
||||
|
||||
GGUF_SCALAR_SIZES = {
|
||||
GGUF_VALUE_UINT8: 1,
|
||||
GGUF_VALUE_INT8: 1,
|
||||
GGUF_VALUE_UINT16: 2,
|
||||
GGUF_VALUE_INT16: 2,
|
||||
GGUF_VALUE_UINT32: 4,
|
||||
GGUF_VALUE_INT32: 4,
|
||||
GGUF_VALUE_FLOAT32: 4,
|
||||
GGUF_VALUE_BOOL: 1,
|
||||
GGUF_VALUE_UINT64: 8,
|
||||
GGUF_VALUE_INT64: 8,
|
||||
GGUF_VALUE_FLOAT64: 8,
|
||||
}
|
||||
|
||||
# GGML quant type -> (block elements, bytes per block, display name).
|
||||
# This intentionally includes the formats used by the current DeepSeek V4 Flash
|
||||
# GGUFs. Add entries here if future recipes introduce new tensor types.
|
||||
GGML_QUANT_SIZES = {
|
||||
0: (1, 4, "F32"),
|
||||
1: (1, 2, "F16"),
|
||||
8: (32, 34, "Q8_0"),
|
||||
10: (256, 84, "Q2_K"),
|
||||
12: (256, 144, "Q4_K"),
|
||||
16: (256, 66, "IQ2_XXS"),
|
||||
26: (1, 4, "I32"),
|
||||
}
|
||||
|
||||
EXPERT_TENSOR_RE = re.compile(r"^blk\.(\d+)\.ffn_(gate|up|down)_exps\.weight$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorInfo:
|
||||
name: str
|
||||
dims: tuple[int, ...]
|
||||
ggml_type: int
|
||||
rel_offset: int
|
||||
data_offset: int
|
||||
n_bytes: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GGUFInfo:
|
||||
path: Path
|
||||
version: int
|
||||
tensor_count: int
|
||||
kv_count: int
|
||||
kv_blob: bytes
|
||||
alignment: int
|
||||
tensors: list[TensorInfo]
|
||||
tensor_by_name: dict[str, TensorInfo]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SplicePlan:
|
||||
name: str
|
||||
source: str
|
||||
tensor: TensorInfo
|
||||
new_rel_offset: int
|
||||
|
||||
|
||||
def read_u32(data: bytes, offset: int) -> int:
|
||||
return struct.unpack_from("<I", data, offset)[0]
|
||||
|
||||
|
||||
def read_u64(data: bytes, offset: int) -> int:
|
||||
return struct.unpack_from("<Q", data, offset)[0]
|
||||
|
||||
|
||||
def pad_to(value: int, alignment: int) -> int:
|
||||
return ((value + alignment - 1) // alignment) * alignment
|
||||
|
||||
|
||||
def read_exact(src: BinaryIO, n_bytes: int) -> bytes:
|
||||
data = src.read(n_bytes)
|
||||
if len(data) != n_bytes:
|
||||
raise EOFError("short read while parsing GGUF")
|
||||
return data
|
||||
|
||||
|
||||
def read_u32_file(src: BinaryIO) -> int:
|
||||
return struct.unpack("<I", read_exact(src, 4))[0]
|
||||
|
||||
|
||||
def read_u64_file(src: BinaryIO) -> int:
|
||||
return struct.unpack("<Q", read_exact(src, 8))[0]
|
||||
|
||||
|
||||
def skip_gguf_string_file(src: BinaryIO) -> None:
|
||||
n = read_u64_file(src)
|
||||
src.seek(n, os.SEEK_CUR)
|
||||
|
||||
|
||||
def read_gguf_string_file(src: BinaryIO) -> str:
|
||||
n = read_u64_file(src)
|
||||
return read_exact(src, n).decode("utf-8")
|
||||
|
||||
|
||||
def skip_value_payload_file(src: BinaryIO, value_type: int) -> None:
|
||||
if value_type == GGUF_VALUE_STRING:
|
||||
skip_gguf_string_file(src)
|
||||
return
|
||||
if value_type == GGUF_VALUE_ARRAY:
|
||||
subtype = read_u32_file(src)
|
||||
count = read_u64_file(src)
|
||||
if subtype == GGUF_VALUE_STRING:
|
||||
for _ in range(count):
|
||||
skip_gguf_string_file(src)
|
||||
return
|
||||
if subtype == GGUF_VALUE_ARRAY:
|
||||
for _ in range(count):
|
||||
skip_value_payload_file(src, subtype)
|
||||
return
|
||||
size = GGUF_SCALAR_SIZES.get(subtype)
|
||||
if size is None:
|
||||
raise ValueError(f"unsupported GGUF array subtype {subtype}")
|
||||
src.seek(count * size, os.SEEK_CUR)
|
||||
return
|
||||
size = GGUF_SCALAR_SIZES.get(value_type)
|
||||
if size is None:
|
||||
raise ValueError(f"unsupported GGUF value type {value_type}")
|
||||
src.seek(size, os.SEEK_CUR)
|
||||
|
||||
|
||||
def tensor_nbytes(dims: tuple[int, ...], ggml_type: int) -> int:
|
||||
q = GGML_QUANT_SIZES.get(ggml_type)
|
||||
if q is None:
|
||||
raise ValueError(f"unsupported GGML tensor type {ggml_type}; add it to GGML_QUANT_SIZES")
|
||||
block_elems, block_bytes, _name = q
|
||||
n_elems = 1
|
||||
for dim in dims:
|
||||
n_elems *= dim
|
||||
if n_elems % block_elems != 0:
|
||||
raise ValueError(f"tensor element count {n_elems} is not divisible by block size {block_elems}")
|
||||
return (n_elems // block_elems) * block_bytes
|
||||
|
||||
|
||||
def parse_gguf(path: Path) -> GGUFInfo:
|
||||
with path.open("rb") as src:
|
||||
if read_exact(src, 4) != b"GGUF":
|
||||
raise ValueError(f"{path} is not a GGUF file")
|
||||
version = read_u32_file(src)
|
||||
tensor_count = read_u64_file(src)
|
||||
kv_count = read_u64_file(src)
|
||||
|
||||
alignment = GGUF_DEFAULT_ALIGNMENT
|
||||
for _ in range(kv_count):
|
||||
key = read_gguf_string_file(src)
|
||||
value_type = read_u32_file(src)
|
||||
value_start = src.tell()
|
||||
skip_value_payload_file(src, value_type)
|
||||
value_end = src.tell()
|
||||
if key == "general.alignment" and value_type == GGUF_VALUE_UINT32:
|
||||
src.seek(value_start)
|
||||
alignment = read_u32_file(src)
|
||||
src.seek(value_end)
|
||||
|
||||
kv_end = src.tell()
|
||||
src.seek(GGUF_HEADER_SIZE)
|
||||
kv_blob = read_exact(src, kv_end - GGUF_HEADER_SIZE)
|
||||
src.seek(kv_end)
|
||||
|
||||
raw_tensors: list[tuple[str, tuple[int, ...], int, int]] = []
|
||||
for _ in range(tensor_count):
|
||||
name = read_gguf_string_file(src)
|
||||
n_dims = read_u32_file(src)
|
||||
dims = tuple(read_u64_file(src) for _ in range(n_dims))
|
||||
ggml_type = read_u32_file(src)
|
||||
rel_offset = read_u64_file(src)
|
||||
raw_tensors.append((name, dims, ggml_type, rel_offset))
|
||||
|
||||
data_start = pad_to(src.tell(), alignment)
|
||||
tensors: list[TensorInfo] = []
|
||||
for name, dims, ggml_type, rel_offset in raw_tensors:
|
||||
n_bytes = tensor_nbytes(dims, ggml_type)
|
||||
tensors.append(TensorInfo(name, dims, ggml_type, rel_offset, data_start + rel_offset, n_bytes))
|
||||
|
||||
return GGUFInfo(
|
||||
path=path,
|
||||
version=version,
|
||||
tensor_count=tensor_count,
|
||||
kv_count=kv_count,
|
||||
kv_blob=kv_blob,
|
||||
alignment=alignment,
|
||||
tensors=tensors,
|
||||
tensor_by_name={t.name: t for t in tensors},
|
||||
)
|
||||
|
||||
|
||||
def parse_layer_set(spec: str) -> set[int]:
|
||||
layers: set[int] = set()
|
||||
for part in spec.split(","):
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
if "-" in part:
|
||||
lo_s, hi_s = part.split("-", 1)
|
||||
lo = int(lo_s)
|
||||
hi = int(hi_s)
|
||||
if hi < lo:
|
||||
raise ValueError(f"bad descending range {part!r}")
|
||||
layers.update(range(lo, hi + 1))
|
||||
else:
|
||||
layers.add(int(part))
|
||||
if not layers:
|
||||
raise ValueError("no layers selected")
|
||||
return layers
|
||||
|
||||
|
||||
def should_take_donor(name: str, q4_layers: set[int]) -> bool:
|
||||
match = EXPERT_TENSOR_RE.match(name)
|
||||
return match is not None and int(match.group(1)) in q4_layers
|
||||
|
||||
|
||||
def qtype_name(ggml_type: int) -> str:
|
||||
return GGML_QUANT_SIZES.get(ggml_type, (0, 0, f"type_{ggml_type}"))[2]
|
||||
|
||||
|
||||
def build_plan(base: GGUFInfo, donor: GGUFInfo, q4_layers: set[int]) -> list[SplicePlan]:
|
||||
if base.version != donor.version:
|
||||
raise ValueError(f"GGUF version mismatch: base={base.version} donor={donor.version}")
|
||||
if base.tensor_count != donor.tensor_count:
|
||||
raise ValueError(f"tensor count mismatch: base={base.tensor_count} donor={donor.tensor_count}")
|
||||
if base.alignment != donor.alignment:
|
||||
raise ValueError(f"alignment mismatch: base={base.alignment} donor={donor.alignment}")
|
||||
|
||||
plan: list[SplicePlan] = []
|
||||
next_rel = 0
|
||||
for base_tensor in base.tensors:
|
||||
donor_tensor = donor.tensor_by_name.get(base_tensor.name)
|
||||
if donor_tensor is None:
|
||||
raise ValueError(f"donor is missing tensor {base_tensor.name}")
|
||||
if base_tensor.dims != donor_tensor.dims:
|
||||
raise ValueError(f"shape mismatch for {base_tensor.name}: {base_tensor.dims} vs {donor_tensor.dims}")
|
||||
use_donor = should_take_donor(base_tensor.name, q4_layers)
|
||||
source_tensor = donor_tensor if use_donor else base_tensor
|
||||
plan.append(SplicePlan(
|
||||
name=base_tensor.name,
|
||||
source="donor" if use_donor else "base",
|
||||
tensor=source_tensor,
|
||||
new_rel_offset=next_rel,
|
||||
))
|
||||
next_rel += pad_to(source_tensor.n_bytes, base.alignment)
|
||||
return plan
|
||||
|
||||
|
||||
def pack_string(value: str) -> bytes:
|
||||
raw = value.encode("utf-8")
|
||||
return struct.pack("<Q", len(raw)) + raw
|
||||
|
||||
|
||||
def write_tensor_info(out: BinaryIO, item: SplicePlan) -> None:
|
||||
t = item.tensor
|
||||
out.write(pack_string(item.name))
|
||||
out.write(struct.pack("<I", len(t.dims)))
|
||||
for dim in t.dims:
|
||||
out.write(struct.pack("<Q", dim))
|
||||
out.write(struct.pack("<I", t.ggml_type))
|
||||
out.write(struct.pack("<Q", item.new_rel_offset))
|
||||
|
||||
|
||||
def copy_exact(src: BinaryIO, dst: BinaryIO, n_bytes: int, bufsize: int = 64 * 1024 * 1024) -> None:
|
||||
remaining = n_bytes
|
||||
while remaining:
|
||||
chunk = src.read(min(bufsize, remaining))
|
||||
if not chunk:
|
||||
raise EOFError("short read while copying tensor payload")
|
||||
dst.write(chunk)
|
||||
remaining -= len(chunk)
|
||||
|
||||
|
||||
def write_padding(out: BinaryIO, n_bytes: int, alignment: int) -> None:
|
||||
pad = pad_to(n_bytes, alignment) - n_bytes
|
||||
if pad:
|
||||
out.write(b"\0" * pad)
|
||||
|
||||
|
||||
def write_mixed(base: GGUFInfo, donor: GGUFInfo, plan: list[SplicePlan], out_path: Path, force: bool) -> None:
|
||||
tmp_path = out_path.with_name(out_path.name + ".tmp")
|
||||
if out_path.exists() and not force:
|
||||
raise FileExistsError(f"{out_path} already exists; pass --force to overwrite")
|
||||
if tmp_path.exists():
|
||||
raise FileExistsError(f"temporary file already exists: {tmp_path}")
|
||||
|
||||
total_payload = sum(item.tensor.n_bytes for item in plan)
|
||||
copied = 0
|
||||
next_report = 0
|
||||
|
||||
with base.path.open("rb") as base_file, donor.path.open("rb") as donor_file, tmp_path.open("wb") as out:
|
||||
out.write(b"GGUF")
|
||||
out.write(struct.pack("<I", base.version))
|
||||
out.write(struct.pack("<Q", base.tensor_count))
|
||||
out.write(struct.pack("<Q", base.kv_count))
|
||||
out.write(base.kv_blob)
|
||||
for item in plan:
|
||||
write_tensor_info(out, item)
|
||||
write_padding(out, out.tell(), base.alignment)
|
||||
|
||||
for item in plan:
|
||||
src = donor_file if item.source == "donor" else base_file
|
||||
src.seek(item.tensor.data_offset)
|
||||
copy_exact(src, out, item.tensor.n_bytes)
|
||||
write_padding(out, item.tensor.n_bytes, base.alignment)
|
||||
copied += item.tensor.n_bytes
|
||||
pct = int(copied * 100 / total_payload) if total_payload else 100
|
||||
if pct >= next_report:
|
||||
print(f"copied {copied / (1024 ** 3):.2f} GiB / {total_payload / (1024 ** 3):.2f} GiB ({pct}%)", flush=True)
|
||||
next_report += 5
|
||||
|
||||
out.flush()
|
||||
os.fsync(out.fileno())
|
||||
|
||||
os.replace(tmp_path, out_path)
|
||||
|
||||
|
||||
def summarize(base: GGUFInfo, donor: GGUFInfo, plan: list[SplicePlan]) -> None:
|
||||
selected = [item for item in plan if item.source == "donor"]
|
||||
base_bytes = sum(t.n_bytes for t in base.tensors)
|
||||
out_bytes = sum(item.tensor.n_bytes for item in plan)
|
||||
delta = out_bytes - base_bytes
|
||||
|
||||
print(f"base: {base.path}")
|
||||
print(f"donor: {donor.path}")
|
||||
print(f"selected donor tensors: {len(selected)}")
|
||||
by_type: dict[str, int] = {}
|
||||
for item in selected:
|
||||
by_type[qtype_name(item.tensor.ggml_type)] = by_type.get(qtype_name(item.tensor.ggml_type), 0) + 1
|
||||
print("selected donor types:", ", ".join(f"{name}:{count}" for name, count in sorted(by_type.items())) or "none")
|
||||
print(f"base tensor payload: {base_bytes:,} bytes ({base_bytes / (1024 ** 3):.2f} GiB)")
|
||||
print(f"mixed tensor payload: {out_bytes:,} bytes ({out_bytes / (1024 ** 3):.2f} GiB)")
|
||||
print(f"payload delta: {delta:,} bytes ({delta / (1024 ** 3):.2f} GiB)")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="Splice selected DeepSeek V4 Flash routed-expert layers from a donor GGUF.")
|
||||
parser.add_argument("--base", required=True, type=Path, help="base GGUF used for metadata and default tensors")
|
||||
parser.add_argument("--donor", required=True, type=Path, help="donor GGUF used for selected routed expert layers")
|
||||
parser.add_argument("--out", required=True, type=Path, help="output mixed GGUF")
|
||||
parser.add_argument("--q4-layers", required=True, help="comma-separated layer IDs/ranges to take from donor, e.g. 37-42")
|
||||
parser.add_argument("--dry-run", action="store_true", help="print the plan without writing the output")
|
||||
parser.add_argument("--force", action="store_true", help="overwrite --out if it already exists")
|
||||
args = parser.parse_args()
|
||||
|
||||
q4_layers = parse_layer_set(args.q4_layers)
|
||||
print("q4 layers:", ",".join(str(x) for x in sorted(q4_layers)))
|
||||
|
||||
base = parse_gguf(args.base)
|
||||
donor = parse_gguf(args.donor)
|
||||
plan = build_plan(base, donor, q4_layers)
|
||||
summarize(base, donor, plan)
|
||||
|
||||
if args.dry_run:
|
||||
return 0
|
||||
|
||||
write_mixed(base, donor, plan, args.out, args.force)
|
||||
final_size = args.out.stat().st_size
|
||||
print(f"wrote: {args.out}")
|
||||
print(f"file size: {final_size:,} bytes ({final_size / (1024 ** 3):.2f} GiB)")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
raise SystemExit(main())
|
||||
except Exception as exc:
|
||||
print(f"error: {exc}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
@@ -0,0 +1,88 @@
|
||||
# Official-Continuation Quality Testing
|
||||
|
||||
This directory contains the 100 prompts and scripts used to compare local GGUF
|
||||
variants against official DeepSeek V4 continuations.
|
||||
|
||||
The metric is target-token negative log likelihood: collect a deterministic
|
||||
official continuation, then ask each local GGUF how much probability it assigns
|
||||
to that exact continuation token by token. This avoids judging quality from one
|
||||
sampled answer.
|
||||
|
||||
## 1. Collect Official Continuations
|
||||
|
||||
```sh
|
||||
export DEEPSEEK_API_KEY=...
|
||||
python3 gguf-tools/quality-testing/collect_official.py \
|
||||
--prompts gguf-tools/quality-testing/prompts.jsonl \
|
||||
--out gguf-tools/quality-testing/data/flash \
|
||||
--count 100 \
|
||||
--max-tokens 24
|
||||
```
|
||||
|
||||
Use one output directory per official model. The default model is Flash, so
|
||||
`data/flash` is the recommended path for Flash continuations. For PRO:
|
||||
|
||||
```sh
|
||||
python3 gguf-tools/quality-testing/collect_official.py \
|
||||
--model deepseek-v4-pro \
|
||||
--prompts gguf-tools/quality-testing/prompts.jsonl \
|
||||
--out gguf-tools/quality-testing/data/pro \
|
||||
--count 100 \
|
||||
--max-tokens 24 \
|
||||
--top-logprobs 20
|
||||
```
|
||||
|
||||
The script writes:
|
||||
|
||||
- `data/<model>/prompts/case_*.txt`
|
||||
- `data/<model>/continuations/case_*.txt`
|
||||
- `data/<model>/responses/case_*.json`
|
||||
- `data/<model>/manifest.tsv`
|
||||
|
||||
The prompt list is tracked in `prompts.jsonl`; the official responses are not
|
||||
tracked because they are derived from an external API.
|
||||
|
||||
## 2. Build The Local Scorer
|
||||
|
||||
```sh
|
||||
make -C gguf-tools quality-score
|
||||
```
|
||||
|
||||
The scorer links against the DS4 runtime and uses Metal by default.
|
||||
|
||||
## 3. Score GGUF Variants
|
||||
|
||||
```sh
|
||||
gguf-tools/quality-testing/score_official \
|
||||
../deepseek-v4-quants/gguf/OLD.gguf \
|
||||
gguf-tools/quality-testing/data/flash/manifest.tsv \
|
||||
/tmp/old.tsv \
|
||||
4096
|
||||
|
||||
gguf-tools/quality-testing/score_official \
|
||||
../deepseek-v4-quants/gguf/NEW.gguf \
|
||||
gguf-tools/quality-testing/data/flash/manifest.tsv \
|
||||
/tmp/new.tsv \
|
||||
4096
|
||||
```
|
||||
|
||||
Use `data/pro/manifest.tsv` for PRO GGUFs. The scorer and comparator do not
|
||||
care which model produced the manifest; the manifest path selects the
|
||||
continuation set.
|
||||
|
||||
## 4. Compare
|
||||
|
||||
```sh
|
||||
python3 gguf-tools/quality-testing/compare_scores.py /tmp/old.tsv /tmp/new.tsv
|
||||
```
|
||||
|
||||
Output fields:
|
||||
|
||||
- `avg_nll`: average negative log likelihood; lower is better.
|
||||
- `delta_new_minus_old`: negative means the new GGUF fits the official
|
||||
continuation better.
|
||||
- `case_wins_new_old_ties`: per-prompt NLL wins.
|
||||
- `first_token_matches`: how often the local greedy first token matches the
|
||||
official first token.
|
||||
- `avg_greedy_lcp`: average greedy longest common prefix against the official
|
||||
continuation.
|
||||
@@ -0,0 +1,266 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Collect official DeepSeek V4 continuations for local quant scoring."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
MODEL = "deepseek-v4-flash"
|
||||
ENDPOINT = "https://api.deepseek.com/chat/completions"
|
||||
|
||||
|
||||
PROMPTS = [
|
||||
"Explain B-tree insertion, including splits and the root special case.",
|
||||
"Write a concise design for a TCP echo server that handles slow clients.",
|
||||
"Compare mmaped model weights with copying all weights into private buffers on macOS.",
|
||||
"Derive why RMSNorm needs a sum of squares and a scaling pass.",
|
||||
"Explain why KV cache checkpointing helps agentic sessions.",
|
||||
"Give pseudocode for a binary heap push operation.",
|
||||
"What are the tradeoffs between top-p sampling and temperature zero decoding?",
|
||||
"Summarize how speculative decoding preserves exact greedy output.",
|
||||
"Explain why a GPU prefill path can be faster than single-token decode.",
|
||||
"Write three invariants for a tokenizer special-token table.",
|
||||
"Spiega come funziona l'inserimento in un B-tree, inclusi split e radice.",
|
||||
"Scrivi un progetto conciso per un server TCP echo con client lenti.",
|
||||
"Confronta mmap dei pesi e copia completa dei pesi su macOS.",
|
||||
"Deriva perche RMSNorm richiede somma dei quadrati e riscalamento.",
|
||||
"Spiega perche la cache KV su disco aiuta nelle sessioni agentiche.",
|
||||
"Scrivi pseudocodice per inserire un elemento in un heap binario.",
|
||||
"Quali sono i pro e contro di top-p rispetto a temperatura zero?",
|
||||
"Riassumi perche la decodifica speculativa puo mantenere l'output esatto.",
|
||||
"Spiega perche il prefill GPU puo essere piu veloce del decode token singolo.",
|
||||
"Scrivi tre invarianti per una tabella di token speciali.",
|
||||
"Given a sorted array, describe how binary search finds the insertion point.",
|
||||
"Write a short C function that clamps an integer to a range.",
|
||||
"Explain the difference between a mutex and an atomic counter.",
|
||||
"What does backpressure mean in a network server?",
|
||||
"Describe a safe format for storing a model checkpoint header.",
|
||||
"Explain how a ring buffer wraps and how to avoid overwriting unread data.",
|
||||
"Write a brief plan for testing long-context prompt chunking.",
|
||||
"What is an importance matrix in low-bit quantization?",
|
||||
"Explain why grouped MoE expert execution can improve prefill.",
|
||||
"Describe how to compare two model quantizations without relying on one answer.",
|
||||
"Data una lista ordinata, descrivi la ricerca binaria del punto di inserimento.",
|
||||
"Scrivi una breve funzione C che limita un intero a un intervallo.",
|
||||
"Spiega la differenza tra mutex e contatore atomico.",
|
||||
"Che cosa significa backpressure in un server di rete?",
|
||||
"Descrivi un formato sicuro per salvare l'header di un checkpoint modello.",
|
||||
"Spiega come funziona un ring buffer e come evitare sovrascritture.",
|
||||
"Scrivi un piano breve per testare il chunking di prompt lunghi.",
|
||||
"Che cos'e una matrice di importanza nella quantizzazione low-bit?",
|
||||
"Spiega perche raggruppare gli esperti MoE puo accelerare il prefill.",
|
||||
"Descrivi come confrontare due quantizzazioni senza fidarsi di una sola risposta.",
|
||||
"A user reports generation slows at long context. List the first five measurements to take.",
|
||||
"Why can small logit differences change a greedy continuation?",
|
||||
"Explain online softmax in attention using simple variables.",
|
||||
"Give a minimal JSON schema for a tool call with name and arguments.",
|
||||
"How would you test that a file-backed mmap is not repeatedly remapped?",
|
||||
"Explain why one large Metal buffer can be worse than multiple overlapping views.",
|
||||
"What is the role of a router in a mixture-of-experts layer?",
|
||||
"Describe the difference between raw KV rows and compressed KV rows.",
|
||||
"Write a checklist for validating that a quantized model still follows instructions.",
|
||||
"Explain why evaluating 50 prompts is more informative than one favorite prompt.",
|
||||
"Un utente segnala decode lento a contesto lungo. Elenca cinque misure iniziali.",
|
||||
"Perche piccole differenze nei logit possono cambiare una continuazione greedy?",
|
||||
"Spiega online softmax nell'attenzione con variabili semplici.",
|
||||
"Dai uno schema JSON minimo per una tool call con nome e argomenti.",
|
||||
"Come testeresti che un mmap su file non venga rimappato ripetutamente?",
|
||||
"Spiega perche un grande buffer Metal puo essere peggiore di viste sovrapposte.",
|
||||
"Qual e il ruolo del router in un layer mixture-of-experts?",
|
||||
"Descrivi la differenza tra righe KV raw e righe KV compresse.",
|
||||
"Scrivi una checklist per validare che un modello quantizzato segua ancora le istruzioni.",
|
||||
"Perche valutare 50 prompt e piu utile di un singolo prompt preferito?",
|
||||
"Write a tiny Python function that returns the median of three numbers.",
|
||||
"Explain why sorted arrays make membership tests faster with binary search.",
|
||||
"What does eventual consistency mean in a distributed database?",
|
||||
"Give a short example of a race condition in C.",
|
||||
"Explain the difference between latency and throughput.",
|
||||
"How does a trie represent a set of strings?",
|
||||
"Describe how to test a command-line parser with edge cases.",
|
||||
"Why can mmap page residency differ from virtual address space size?",
|
||||
"Explain why quantization error can affect rare experts more than common experts.",
|
||||
"Write a short checklist for reviewing a pull request that touches memory lifetimes.",
|
||||
"Scrivi una piccola funzione Python che restituisce la mediana di tre numeri.",
|
||||
"Spiega perche array ordinati rendono piu veloce la ricerca di appartenenza.",
|
||||
"Che cosa significa consistenza eventuale in un database distribuito?",
|
||||
"Fai un breve esempio di race condition in C.",
|
||||
"Spiega la differenza tra latenza e throughput.",
|
||||
"Come rappresenta un trie un insieme di stringhe?",
|
||||
"Descrivi come testare un parser da riga di comando con casi limite.",
|
||||
"Perche la residenza delle pagine mmap puo differire dalla dimensione virtuale?",
|
||||
"Spiega perche l'errore di quantizzazione puo colpire piu gli esperti rari.",
|
||||
"Scrivi una breve checklist per revisionare lifetime di memoria in una PR.",
|
||||
"Complete this sentence: A good benchmark should measure",
|
||||
"Complete this C comment: /* This lock protects",
|
||||
"Complete this Italian sentence: Il vantaggio principale della cache e",
|
||||
"Translate to Italian: The model should answer only after reading the whole prompt.",
|
||||
"Translate to English: La quantizzazione riduce memoria ma puo alterare i logit.",
|
||||
"In one paragraph, explain how a compiler uses an abstract syntax tree.",
|
||||
"In one paragraph, explain why checksums catch accidental corruption.",
|
||||
"Give three examples of useful server metrics.",
|
||||
"Why should a tokenizer treat special tags carefully?",
|
||||
"Explain how a hash table handles collisions.",
|
||||
"Describe the role of calibration data when quantizing a neural network.",
|
||||
"What is a confidence interval, in plain language?",
|
||||
"Write a simple SQL query that counts rows per category.",
|
||||
"Explain why a page cache can make a second file read faster.",
|
||||
"Give a short answer: what is the capital of Japan?",
|
||||
"Give a short answer: what is the derivative of x squared?",
|
||||
"Give a short answer: who wrote The Divine Comedy?",
|
||||
"Rispondi brevemente: qual e la capitale del Giappone?",
|
||||
"Rispondi brevemente: quanto fa 17 per 23?",
|
||||
"Rispondi brevemente: chi ha scritto la Divina Commedia?",
|
||||
]
|
||||
|
||||
|
||||
def request_one(
|
||||
api_key: str,
|
||||
endpoint: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
max_tokens: int,
|
||||
top_logprobs: int,
|
||||
thinking: str,
|
||||
) -> dict:
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": 0,
|
||||
"max_tokens": max_tokens,
|
||||
"logprobs": True,
|
||||
"top_logprobs": top_logprobs,
|
||||
"stream": False,
|
||||
}
|
||||
if thinking != "omit":
|
||||
payload["thinking"] = {"type": thinking}
|
||||
req = urllib.request.Request(
|
||||
endpoint,
|
||||
data=json.dumps(payload).encode("utf-8"),
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=120) as fp:
|
||||
return json.loads(fp.read().decode("utf-8"))
|
||||
|
||||
|
||||
def fetch_with_retry(
|
||||
api_key: str,
|
||||
endpoint: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
max_tokens: int,
|
||||
top_logprobs: int,
|
||||
thinking: str,
|
||||
) -> dict:
|
||||
delay = 1.0
|
||||
for attempt in range(6):
|
||||
try:
|
||||
return request_one(api_key, endpoint, model, prompt, max_tokens, top_logprobs, thinking)
|
||||
except urllib.error.HTTPError as e:
|
||||
body = e.read().decode("utf-8", "replace")
|
||||
if e.code < 500 and e.code != 429:
|
||||
raise RuntimeError(f"HTTP {e.code}: {body}") from e
|
||||
last = RuntimeError(f"HTTP {e.code}: {body}")
|
||||
except Exception as e: # noqa: BLE001 - command-line retry wrapper.
|
||||
last = e
|
||||
if attempt == 5:
|
||||
raise last
|
||||
time.sleep(delay)
|
||||
delay *= 1.7
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--out", default="gguf-tools/quality-testing/data")
|
||||
ap.add_argument("--prompts", default="gguf-tools/quality-testing/prompts.jsonl")
|
||||
ap.add_argument("--model", default=MODEL)
|
||||
ap.add_argument("--endpoint", default=ENDPOINT)
|
||||
ap.add_argument("--count", type=int, default=100)
|
||||
ap.add_argument("--max-tokens", type=int, default=24)
|
||||
ap.add_argument("--top-logprobs", type=int, default=5)
|
||||
ap.add_argument("--thinking", choices=("disabled", "enabled", "omit"), default="disabled")
|
||||
args = ap.parse_args()
|
||||
if args.top_logprobs < 0 or args.top_logprobs > 20:
|
||||
raise SystemExit("--top-logprobs must be between 0 and 20")
|
||||
|
||||
api_key = os.environ.get("DEEPSEEK_API_KEY")
|
||||
if not api_key:
|
||||
raise SystemExit("DEEPSEEK_API_KEY is not set")
|
||||
prompts = load_prompts(Path(args.prompts))
|
||||
|
||||
out = Path(args.out)
|
||||
(out / "prompts").mkdir(parents=True, exist_ok=True)
|
||||
(out / "continuations").mkdir(parents=True, exist_ok=True)
|
||||
(out / "responses").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
manifest = out / "manifest.tsv"
|
||||
rows = []
|
||||
total = min(args.count, len(prompts))
|
||||
print(f"model={args.model} endpoint={args.endpoint}", file=sys.stderr)
|
||||
for i, prompt in enumerate(prompts[: args.count]):
|
||||
case_id = f"case_{i:03d}"
|
||||
print(f"official {i + 1}/{total}: {case_id}", file=sys.stderr, flush=True)
|
||||
response = fetch_with_retry(
|
||||
api_key,
|
||||
args.endpoint,
|
||||
args.model,
|
||||
prompt,
|
||||
args.max_tokens,
|
||||
args.top_logprobs,
|
||||
args.thinking,
|
||||
)
|
||||
choice = response["choices"][0]
|
||||
content = choice.get("message", {}).get("content", "")
|
||||
if not content:
|
||||
print(f"warning: empty continuation for {case_id}", file=sys.stderr)
|
||||
|
||||
prompt_path = out / "prompts" / f"{case_id}.txt"
|
||||
cont_path = out / "continuations" / f"{case_id}.txt"
|
||||
resp_path = out / "responses" / f"{case_id}.json"
|
||||
prompt_path.write_text(prompt, encoding="utf-8")
|
||||
cont_path.write_text(content, encoding="utf-8")
|
||||
resp_path.write_text(json.dumps(response, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
rows.append((case_id, prompt_path, cont_path, resp_path))
|
||||
time.sleep(0.05)
|
||||
|
||||
with manifest.open("w", encoding="utf-8") as fp:
|
||||
fp.write("# id\tprompt_file\tcontinuation_file\tresponse_file\n")
|
||||
for row in rows:
|
||||
fp.write("\t".join([row[0], str(row[1]), str(row[2]), str(row[3])]) + "\n")
|
||||
print(f"wrote {manifest}", file=sys.stderr)
|
||||
return 0
|
||||
|
||||
|
||||
def load_prompts(path: Path) -> list[str]:
|
||||
if not path.exists():
|
||||
return PROMPTS
|
||||
prompts = []
|
||||
with path.open(encoding="utf-8") as fp:
|
||||
for line in fp:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
obj = json.loads(line)
|
||||
prompt = obj.get("prompt") if isinstance(obj, dict) else None
|
||||
if not prompt:
|
||||
raise SystemExit(f"bad prompt row in {path}: {line[:120]}")
|
||||
prompts.append(prompt)
|
||||
if not prompts:
|
||||
raise SystemExit(f"no prompts found in {path}")
|
||||
return prompts
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,88 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compare two local model scores on official continuations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def load(path: Path) -> dict[str, dict[str, float]]:
|
||||
with path.open(newline="", encoding="utf-8") as fp:
|
||||
rows = {}
|
||||
for row in csv.DictReader(fp, delimiter="\t"):
|
||||
rows[row["id"]] = {
|
||||
"target_tokens": int(row["target_tokens"]),
|
||||
"nll": float(row["nll"]),
|
||||
"avg_nll": float(row["avg_nll"]),
|
||||
"first_match": int(row["first_match"]),
|
||||
"greedy_lcp": int(row["greedy_lcp"]),
|
||||
}
|
||||
return rows
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) != 3:
|
||||
print(f"usage: {sys.argv[0]} OLD.tsv NEW.tsv", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
old = load(Path(sys.argv[1]))
|
||||
new = load(Path(sys.argv[2]))
|
||||
ids = sorted(set(old) & set(new))
|
||||
if not ids:
|
||||
raise SystemExit("no common cases")
|
||||
|
||||
old_nll = new_nll = 0.0
|
||||
old_first = new_first = 0
|
||||
old_lcp = new_lcp = 0
|
||||
tokens = 0
|
||||
new_case_wins = old_case_wins = ties = 0
|
||||
deltas = []
|
||||
|
||||
for case_id in ids:
|
||||
o = old[case_id]
|
||||
n = new[case_id]
|
||||
if o["target_tokens"] != n["target_tokens"]:
|
||||
raise SystemExit(f"token-count mismatch for {case_id}")
|
||||
t = int(o["target_tokens"])
|
||||
tokens += t
|
||||
old_nll += o["nll"]
|
||||
new_nll += n["nll"]
|
||||
old_first += int(o["first_match"])
|
||||
new_first += int(n["first_match"])
|
||||
old_lcp += int(o["greedy_lcp"])
|
||||
new_lcp += int(n["greedy_lcp"])
|
||||
delta = n["nll"] - o["nll"]
|
||||
deltas.append((delta, case_id, t, o["avg_nll"], n["avg_nll"]))
|
||||
if delta < -1e-9:
|
||||
new_case_wins += 1
|
||||
elif delta > 1e-9:
|
||||
old_case_wins += 1
|
||||
else:
|
||||
ties += 1
|
||||
|
||||
avg_old = old_nll / tokens
|
||||
avg_new = new_nll / tokens
|
||||
print(f"cases\t{len(ids)}")
|
||||
print(f"tokens\t{tokens}")
|
||||
print(f"old_avg_nll\t{avg_old:.9f}")
|
||||
print(f"new_avg_nll\t{avg_new:.9f}")
|
||||
print(f"delta_new_minus_old\t{avg_new - avg_old:.9f}")
|
||||
print(f"relative_nll_change\t{(avg_new / avg_old - 1.0) * 100.0:.3f}%")
|
||||
print(f"case_wins_new_old_ties\t{new_case_wins}\t{old_case_wins}\t{ties}")
|
||||
print(f"first_token_matches_old_new\t{old_first}\t{new_first}")
|
||||
print(f"avg_greedy_lcp_old_new\t{old_lcp / len(ids):.3f}\t{new_lcp / len(ids):.3f}")
|
||||
|
||||
print("\nnew best cases:")
|
||||
for delta, case_id, t, old_avg, new_avg in sorted(deltas)[:8]:
|
||||
print(f"{case_id}\tdelta_nll={delta:.6f}\ttokens={t}\told={old_avg:.6f}\tnew={new_avg:.6f}")
|
||||
|
||||
print("\nold best cases:")
|
||||
for delta, case_id, t, old_avg, new_avg in sorted(deltas, reverse=True)[:8]:
|
||||
print(f"{case_id}\tdelta_nll={delta:.6f}\ttokens={t}\told={old_avg:.6f}\tnew={new_avg:.6f}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,100 @@
|
||||
{"id": "case_000", "prompt": "Explain B-tree insertion, including splits and the root special case."}
|
||||
{"id": "case_001", "prompt": "Write a concise design for a TCP echo server that handles slow clients."}
|
||||
{"id": "case_002", "prompt": "Compare mmaped model weights with copying all weights into private buffers on macOS."}
|
||||
{"id": "case_003", "prompt": "Derive why RMSNorm needs a sum of squares and a scaling pass."}
|
||||
{"id": "case_004", "prompt": "Explain why KV cache checkpointing helps agentic sessions."}
|
||||
{"id": "case_005", "prompt": "Give pseudocode for a binary heap push operation."}
|
||||
{"id": "case_006", "prompt": "What are the tradeoffs between top-p sampling and temperature zero decoding?"}
|
||||
{"id": "case_007", "prompt": "Summarize how speculative decoding preserves exact greedy output."}
|
||||
{"id": "case_008", "prompt": "Explain why a GPU prefill path can be faster than single-token decode."}
|
||||
{"id": "case_009", "prompt": "Write three invariants for a tokenizer special-token table."}
|
||||
{"id": "case_010", "prompt": "Spiega come funziona l'inserimento in un B-tree, inclusi split e radice."}
|
||||
{"id": "case_011", "prompt": "Scrivi un progetto conciso per un server TCP echo con client lenti."}
|
||||
{"id": "case_012", "prompt": "Confronta mmap dei pesi e copia completa dei pesi su macOS."}
|
||||
{"id": "case_013", "prompt": "Deriva perche RMSNorm richiede somma dei quadrati e riscalamento."}
|
||||
{"id": "case_014", "prompt": "Spiega perche la cache KV su disco aiuta nelle sessioni agentiche."}
|
||||
{"id": "case_015", "prompt": "Scrivi pseudocodice per inserire un elemento in un heap binario."}
|
||||
{"id": "case_016", "prompt": "Quali sono i pro e contro di top-p rispetto a temperatura zero?"}
|
||||
{"id": "case_017", "prompt": "Riassumi perche la decodifica speculativa puo mantenere l'output esatto."}
|
||||
{"id": "case_018", "prompt": "Spiega perche il prefill GPU puo essere piu veloce del decode token singolo."}
|
||||
{"id": "case_019", "prompt": "Scrivi tre invarianti per una tabella di token speciali."}
|
||||
{"id": "case_020", "prompt": "Given a sorted array, describe how binary search finds the insertion point."}
|
||||
{"id": "case_021", "prompt": "Write a short C function that clamps an integer to a range."}
|
||||
{"id": "case_022", "prompt": "Explain the difference between a mutex and an atomic counter."}
|
||||
{"id": "case_023", "prompt": "What does backpressure mean in a network server?"}
|
||||
{"id": "case_024", "prompt": "Describe a safe format for storing a model checkpoint header."}
|
||||
{"id": "case_025", "prompt": "Explain how a ring buffer wraps and how to avoid overwriting unread data."}
|
||||
{"id": "case_026", "prompt": "Write a brief plan for testing long-context prompt chunking."}
|
||||
{"id": "case_027", "prompt": "What is an importance matrix in low-bit quantization?"}
|
||||
{"id": "case_028", "prompt": "Explain why grouped MoE expert execution can improve prefill."}
|
||||
{"id": "case_029", "prompt": "Describe how to compare two model quantizations without relying on one answer."}
|
||||
{"id": "case_030", "prompt": "Data una lista ordinata, descrivi la ricerca binaria del punto di inserimento."}
|
||||
{"id": "case_031", "prompt": "Scrivi una breve funzione C che limita un intero a un intervallo."}
|
||||
{"id": "case_032", "prompt": "Spiega la differenza tra mutex e contatore atomico."}
|
||||
{"id": "case_033", "prompt": "Che cosa significa backpressure in un server di rete?"}
|
||||
{"id": "case_034", "prompt": "Descrivi un formato sicuro per salvare l'header di un checkpoint modello."}
|
||||
{"id": "case_035", "prompt": "Spiega come funziona un ring buffer e come evitare sovrascritture."}
|
||||
{"id": "case_036", "prompt": "Scrivi un piano breve per testare il chunking di prompt lunghi."}
|
||||
{"id": "case_037", "prompt": "Che cos'e una matrice di importanza nella quantizzazione low-bit?"}
|
||||
{"id": "case_038", "prompt": "Spiega perche raggruppare gli esperti MoE puo accelerare il prefill."}
|
||||
{"id": "case_039", "prompt": "Descrivi come confrontare due quantizzazioni senza fidarsi di una sola risposta."}
|
||||
{"id": "case_040", "prompt": "A user reports generation slows at long context. List the first five measurements to take."}
|
||||
{"id": "case_041", "prompt": "Why can small logit differences change a greedy continuation?"}
|
||||
{"id": "case_042", "prompt": "Explain online softmax in attention using simple variables."}
|
||||
{"id": "case_043", "prompt": "Give a minimal JSON schema for a tool call with name and arguments."}
|
||||
{"id": "case_044", "prompt": "How would you test that a file-backed mmap is not repeatedly remapped?"}
|
||||
{"id": "case_045", "prompt": "Explain why one large Metal buffer can be worse than multiple overlapping views."}
|
||||
{"id": "case_046", "prompt": "What is the role of a router in a mixture-of-experts layer?"}
|
||||
{"id": "case_047", "prompt": "Describe the difference between raw KV rows and compressed KV rows."}
|
||||
{"id": "case_048", "prompt": "Write a checklist for validating that a quantized model still follows instructions."}
|
||||
{"id": "case_049", "prompt": "Explain why evaluating 50 prompts is more informative than one favorite prompt."}
|
||||
{"id": "case_050", "prompt": "Un utente segnala decode lento a contesto lungo. Elenca cinque misure iniziali."}
|
||||
{"id": "case_051", "prompt": "Perche piccole differenze nei logit possono cambiare una continuazione greedy?"}
|
||||
{"id": "case_052", "prompt": "Spiega online softmax nell'attenzione con variabili semplici."}
|
||||
{"id": "case_053", "prompt": "Dai uno schema JSON minimo per una tool call con nome e argomenti."}
|
||||
{"id": "case_054", "prompt": "Come testeresti che un mmap su file non venga rimappato ripetutamente?"}
|
||||
{"id": "case_055", "prompt": "Spiega perche un grande buffer Metal puo essere peggiore di viste sovrapposte."}
|
||||
{"id": "case_056", "prompt": "Qual e il ruolo del router in un layer mixture-of-experts?"}
|
||||
{"id": "case_057", "prompt": "Descrivi la differenza tra righe KV raw e righe KV compresse."}
|
||||
{"id": "case_058", "prompt": "Scrivi una checklist per validare che un modello quantizzato segua ancora le istruzioni."}
|
||||
{"id": "case_059", "prompt": "Perche valutare 50 prompt e piu utile di un singolo prompt preferito?"}
|
||||
{"id": "case_060", "prompt": "Write a tiny Python function that returns the median of three numbers."}
|
||||
{"id": "case_061", "prompt": "Explain why sorted arrays make membership tests faster with binary search."}
|
||||
{"id": "case_062", "prompt": "What does eventual consistency mean in a distributed database?"}
|
||||
{"id": "case_063", "prompt": "Give a short example of a race condition in C."}
|
||||
{"id": "case_064", "prompt": "Explain the difference between latency and throughput."}
|
||||
{"id": "case_065", "prompt": "How does a trie represent a set of strings?"}
|
||||
{"id": "case_066", "prompt": "Describe how to test a command-line parser with edge cases."}
|
||||
{"id": "case_067", "prompt": "Why can mmap page residency differ from virtual address space size?"}
|
||||
{"id": "case_068", "prompt": "Explain why quantization error can affect rare experts more than common experts."}
|
||||
{"id": "case_069", "prompt": "Write a short checklist for reviewing a pull request that touches memory lifetimes."}
|
||||
{"id": "case_070", "prompt": "Scrivi una piccola funzione Python che restituisce la mediana di tre numeri."}
|
||||
{"id": "case_071", "prompt": "Spiega perche array ordinati rendono piu veloce la ricerca di appartenenza."}
|
||||
{"id": "case_072", "prompt": "Che cosa significa consistenza eventuale in un database distribuito?"}
|
||||
{"id": "case_073", "prompt": "Fai un breve esempio di race condition in C."}
|
||||
{"id": "case_074", "prompt": "Spiega la differenza tra latenza e throughput."}
|
||||
{"id": "case_075", "prompt": "Come rappresenta un trie un insieme di stringhe?"}
|
||||
{"id": "case_076", "prompt": "Descrivi come testare un parser da riga di comando con casi limite."}
|
||||
{"id": "case_077", "prompt": "Perche la residenza delle pagine mmap puo differire dalla dimensione virtuale?"}
|
||||
{"id": "case_078", "prompt": "Spiega perche l'errore di quantizzazione puo colpire piu gli esperti rari."}
|
||||
{"id": "case_079", "prompt": "Scrivi una breve checklist per revisionare lifetime di memoria in una PR."}
|
||||
{"id": "case_080", "prompt": "Complete this sentence: A good benchmark should measure"}
|
||||
{"id": "case_081", "prompt": "Complete this C comment: /* This lock protects"}
|
||||
{"id": "case_082", "prompt": "Complete this Italian sentence: Il vantaggio principale della cache e"}
|
||||
{"id": "case_083", "prompt": "Translate to Italian: The model should answer only after reading the whole prompt."}
|
||||
{"id": "case_084", "prompt": "Translate to English: La quantizzazione riduce memoria ma puo alterare i logit."}
|
||||
{"id": "case_085", "prompt": "In one paragraph, explain how a compiler uses an abstract syntax tree."}
|
||||
{"id": "case_086", "prompt": "In one paragraph, explain why checksums catch accidental corruption."}
|
||||
{"id": "case_087", "prompt": "Give three examples of useful server metrics."}
|
||||
{"id": "case_088", "prompt": "Why should a tokenizer treat special tags carefully?"}
|
||||
{"id": "case_089", "prompt": "Explain how a hash table handles collisions."}
|
||||
{"id": "case_090", "prompt": "Describe the role of calibration data when quantizing a neural network."}
|
||||
{"id": "case_091", "prompt": "What is a confidence interval, in plain language?"}
|
||||
{"id": "case_092", "prompt": "Write a simple SQL query that counts rows per category."}
|
||||
{"id": "case_093", "prompt": "Explain why a page cache can make a second file read faster."}
|
||||
{"id": "case_094", "prompt": "Give a short answer: what is the capital of Japan?"}
|
||||
{"id": "case_095", "prompt": "Give a short answer: what is the derivative of x squared?"}
|
||||
{"id": "case_096", "prompt": "Give a short answer: who wrote The Divine Comedy?"}
|
||||
{"id": "case_097", "prompt": "Rispondi brevemente: qual e la capitale del Giappone?"}
|
||||
{"id": "case_098", "prompt": "Rispondi brevemente: quanto fa 17 per 23?"}
|
||||
{"id": "case_099", "prompt": "Rispondi brevemente: chi ha scritto la Divina Commedia?"}
|
||||
@@ -0,0 +1,170 @@
|
||||
#include "ds4.h"
|
||||
|
||||
#include <errno.h>
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
|
||||
static void die(const char *msg) {
|
||||
fprintf(stderr, "%s\n", msg);
|
||||
exit(1);
|
||||
}
|
||||
|
||||
static char *read_file(const char *path) {
|
||||
FILE *fp = fopen(path, "rb");
|
||||
if (!fp) {
|
||||
fprintf(stderr, "open %s: %s\n", path, strerror(errno));
|
||||
exit(1);
|
||||
}
|
||||
if (fseek(fp, 0, SEEK_END) != 0) die("fseek failed");
|
||||
long n = ftell(fp);
|
||||
if (n < 0) die("ftell failed");
|
||||
if (fseek(fp, 0, SEEK_SET) != 0) die("fseek failed");
|
||||
char *buf = malloc((size_t)n + 1);
|
||||
if (!buf) die("out of memory");
|
||||
if (n && fread(buf, 1, (size_t)n, fp) != (size_t)n) die("read failed");
|
||||
buf[n] = '\0';
|
||||
fclose(fp);
|
||||
return buf;
|
||||
}
|
||||
|
||||
static void strip_newline(char *s) {
|
||||
size_t n = strlen(s);
|
||||
while (n && (s[n - 1] == '\n' || s[n - 1] == '\r')) s[--n] = '\0';
|
||||
}
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
if (argc != 4 && argc != 5) {
|
||||
fprintf(stderr, "usage: %s MODEL manifest.tsv OUT.tsv [ctx]\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
|
||||
const char *model_path = argv[1];
|
||||
const char *manifest_path = argv[2];
|
||||
const char *out_path = argv[3];
|
||||
int ctx_size = argc == 5 ? atoi(argv[4]) : 4096;
|
||||
if (ctx_size < 1024) ctx_size = 1024;
|
||||
|
||||
ds4_engine_options opt = {
|
||||
.model_path = model_path,
|
||||
#ifdef __APPLE__
|
||||
.backend = DS4_BACKEND_METAL,
|
||||
#else
|
||||
.backend = DS4_BACKEND_CUDA,
|
||||
#endif
|
||||
.n_threads = 0,
|
||||
.warm_weights = false,
|
||||
.quality = false,
|
||||
};
|
||||
|
||||
ds4_engine *engine = NULL;
|
||||
if (ds4_engine_open(&engine, &opt) != 0) die("failed to open model");
|
||||
|
||||
ds4_session *session = NULL;
|
||||
if (ds4_session_create(&session, engine, ctx_size) != 0) die("failed to create session");
|
||||
|
||||
FILE *mf = fopen(manifest_path, "rb");
|
||||
if (!mf) {
|
||||
fprintf(stderr, "open %s: %s\n", manifest_path, strerror(errno));
|
||||
return 1;
|
||||
}
|
||||
FILE *out = fopen(out_path, "wb");
|
||||
if (!out) {
|
||||
fprintf(stderr, "open %s: %s\n", out_path, strerror(errno));
|
||||
return 1;
|
||||
}
|
||||
fprintf(out, "id\tprompt_tokens\ttarget_tokens\tnll\tavg_nll\tfirst_match\tgreedy_lcp\n");
|
||||
|
||||
char line[8192];
|
||||
int case_n = 0;
|
||||
double total_nll = 0.0;
|
||||
long total_tokens = 0;
|
||||
long total_lcp = 0;
|
||||
long first_matches = 0;
|
||||
char err[256];
|
||||
|
||||
while (fgets(line, sizeof(line), mf)) {
|
||||
strip_newline(line);
|
||||
if (!line[0] || line[0] == '#') continue;
|
||||
|
||||
char *id = strtok(line, "\t");
|
||||
char *prompt_path = strtok(NULL, "\t");
|
||||
char *cont_path = strtok(NULL, "\t");
|
||||
if (!id || !prompt_path || !cont_path) die("bad manifest row");
|
||||
|
||||
char *prompt_text = read_file(prompt_path);
|
||||
char *cont_text = read_file(cont_path);
|
||||
|
||||
ds4_tokens prompt = {0};
|
||||
ds4_tokens target = {0};
|
||||
ds4_encode_chat_prompt(engine, NULL, prompt_text, DS4_THINK_NONE, &prompt);
|
||||
ds4_tokenize_text(engine, cont_text, &target);
|
||||
|
||||
if (prompt.len + target.len + 1 >= ctx_size) {
|
||||
fprintf(stderr, "%s exceeds ctx=%d\n", id, ctx_size);
|
||||
return 1;
|
||||
}
|
||||
if (ds4_session_sync(session, &prompt, err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "%s sync failed: %s\n", id, err);
|
||||
return 1;
|
||||
}
|
||||
|
||||
double nll = 0.0;
|
||||
int lcp = 0;
|
||||
bool still_matching = true;
|
||||
bool first_match = false;
|
||||
for (int i = 0; i < target.len; i++) {
|
||||
const int greedy = ds4_session_argmax(session);
|
||||
if (i == 0) first_match = (greedy == target.v[i]);
|
||||
if (still_matching && greedy == target.v[i]) lcp++;
|
||||
else still_matching = false;
|
||||
|
||||
ds4_token_score score;
|
||||
if (!ds4_session_token_logprob(session, target.v[i], &score)) {
|
||||
fprintf(stderr, "%s logprob failed at target token %d\n", id, i);
|
||||
return 1;
|
||||
}
|
||||
nll += -(double)score.logprob;
|
||||
|
||||
if (ds4_session_eval(session, target.v[i], err, sizeof(err)) != 0) {
|
||||
fprintf(stderr, "%s eval failed at target token %d: %s\n", id, i, err);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
const double avg = target.len ? nll / (double)target.len : 0.0;
|
||||
fprintf(out, "%s\t%d\t%d\t%.9f\t%.9f\t%d\t%d\n",
|
||||
id, prompt.len, target.len, nll, avg, first_match ? 1 : 0, lcp);
|
||||
fflush(out);
|
||||
|
||||
case_n++;
|
||||
total_nll += nll;
|
||||
total_tokens += target.len;
|
||||
total_lcp += lcp;
|
||||
first_matches += first_match ? 1 : 0;
|
||||
fprintf(stderr,
|
||||
"%s cases=%d prompt=%d target=%d avg_nll=%.6f lcp=%d\n",
|
||||
id, case_n, prompt.len, target.len, avg, lcp);
|
||||
|
||||
ds4_tokens_free(&prompt);
|
||||
ds4_tokens_free(&target);
|
||||
free(prompt_text);
|
||||
free(cont_text);
|
||||
}
|
||||
|
||||
fprintf(stderr,
|
||||
"summary cases=%d tokens=%ld avg_nll=%.9f first_match=%ld avg_lcp=%.3f\n",
|
||||
case_n,
|
||||
total_tokens,
|
||||
total_tokens ? total_nll / (double)total_tokens : 0.0,
|
||||
first_matches,
|
||||
case_n ? (double)total_lcp / (double)case_n : 0.0);
|
||||
|
||||
fclose(out);
|
||||
fclose(mf);
|
||||
ds4_session_free(session);
|
||||
ds4_engine_close(engine);
|
||||
return 0;
|
||||
}
|
||||
+1109
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,75 @@
|
||||
#ifndef DS4_QUANTS_H
|
||||
#define DS4_QUANTS_H
|
||||
|
||||
/*
|
||||
* Narrow quantization API used by the DS4 GGUF writer.
|
||||
*
|
||||
* The enum values intentionally match GGUF/GGML type IDs so template metadata
|
||||
* can be copied without translation. Only the formats used by the DS4 Flash
|
||||
* quantization recipes are implemented as output targets.
|
||||
*/
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#define DS4Q_MAX_DIMS 4
|
||||
|
||||
typedef enum {
|
||||
DS4Q_TYPE_F32 = 0,
|
||||
DS4Q_TYPE_F16 = 1,
|
||||
DS4Q_TYPE_Q4_0 = 2,
|
||||
DS4Q_TYPE_Q4_1 = 3,
|
||||
DS4Q_TYPE_Q5_0 = 6,
|
||||
DS4Q_TYPE_Q5_1 = 7,
|
||||
DS4Q_TYPE_Q8_0 = 8,
|
||||
DS4Q_TYPE_Q8_1 = 9,
|
||||
DS4Q_TYPE_Q2_K = 10,
|
||||
DS4Q_TYPE_Q3_K = 11,
|
||||
DS4Q_TYPE_Q4_K = 12,
|
||||
DS4Q_TYPE_Q5_K = 13,
|
||||
DS4Q_TYPE_Q6_K = 14,
|
||||
DS4Q_TYPE_Q8_K = 15,
|
||||
DS4Q_TYPE_IQ2_XXS = 16,
|
||||
DS4Q_TYPE_IQ2_XS = 17,
|
||||
DS4Q_TYPE_IQ3_XXS = 18,
|
||||
DS4Q_TYPE_IQ1_S = 19,
|
||||
DS4Q_TYPE_IQ4_NL = 20,
|
||||
DS4Q_TYPE_IQ3_S = 21,
|
||||
DS4Q_TYPE_IQ2_S = 22,
|
||||
DS4Q_TYPE_IQ4_XS = 23,
|
||||
DS4Q_TYPE_I8 = 24,
|
||||
DS4Q_TYPE_I16 = 25,
|
||||
DS4Q_TYPE_I32 = 26,
|
||||
DS4Q_TYPE_I64 = 27,
|
||||
DS4Q_TYPE_F64 = 28,
|
||||
DS4Q_TYPE_IQ1_M = 29,
|
||||
DS4Q_TYPE_BF16 = 30,
|
||||
DS4Q_TYPE_TQ1_0 = 34,
|
||||
DS4Q_TYPE_TQ2_0 = 35,
|
||||
DS4Q_TYPE_MXFP4 = 39,
|
||||
DS4Q_TYPE_NVFP4 = 40,
|
||||
DS4Q_TYPE_Q1_0 = 41,
|
||||
DS4Q_TYPE_COUNT = 42,
|
||||
} ds4q_type;
|
||||
|
||||
static inline size_t ds4q_pad(size_t x, size_t n) {
|
||||
return ((x + n - 1) / n) * n;
|
||||
}
|
||||
|
||||
const char *ds4q_type_name(ds4q_type type);
|
||||
bool ds4q_can_quantize(ds4q_type type);
|
||||
int64_t ds4q_block_size(ds4q_type type);
|
||||
size_t ds4q_row_size(ds4q_type type, int64_t ne);
|
||||
bool ds4q_requires_imatrix(ds4q_type type);
|
||||
void ds4q_quantize_init(ds4q_type type);
|
||||
size_t ds4q_quantize_chunk(ds4q_type type, const float *src, void *dst,
|
||||
int64_t start, int64_t nrows, int64_t ncols,
|
||||
const float *imatrix);
|
||||
|
||||
float ds4q_f16_to_f32(uint16_t bits);
|
||||
float ds4q_bf16_to_f32(uint16_t bits);
|
||||
void ds4q_f32_to_f16_row(const float *src, uint16_t *dst, int64_t n);
|
||||
void ds4q_f32_to_bf16_row(const float *src, uint16_t *dst, int64_t n);
|
||||
|
||||
#endif
|
||||
+2690
File diff suppressed because it is too large
Load Diff
+149
@@ -0,0 +1,149 @@
|
||||
/* linenoise.h -- VERSION 1.0
|
||||
*
|
||||
* Guerrilla line editing library against the idea that a line editing lib
|
||||
* needs to be 20,000 lines of C code.
|
||||
*
|
||||
* See linenoise.c for more information.
|
||||
*
|
||||
* ------------------------------------------------------------------------
|
||||
*
|
||||
* Copyright (c) 2010-2023, Salvatore Sanfilippo <antirez at gmail dot com>
|
||||
* Copyright (c) 2010-2013, Pieter Noordhuis <pcnoordhuis at gmail dot com>
|
||||
*
|
||||
* All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are
|
||||
* met:
|
||||
*
|
||||
* * Redistributions of source code must retain the above copyright
|
||||
* notice, this list of conditions and the following disclaimer.
|
||||
*
|
||||
* * Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in the
|
||||
* documentation and/or other materials provided with the distribution.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
* HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
||||
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
||||
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*/
|
||||
|
||||
#ifndef __LINENOISE_H
|
||||
#define __LINENOISE_H
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
#include <stddef.h> /* For size_t. */
|
||||
|
||||
extern char *linenoiseEditMore;
|
||||
|
||||
#define LINENOISE_MAX_FOLDS 16
|
||||
|
||||
struct linenoiseState;
|
||||
typedef int(linenoiseLayoutCallback)(struct linenoiseState *l,
|
||||
size_t prompt_rows,
|
||||
size_t status_rows,
|
||||
void *privdata);
|
||||
|
||||
/* The linenoiseState structure represents the state during line editing.
|
||||
* We pass this state to functions implementing specific editing
|
||||
* functionalities. */
|
||||
struct linenoiseState {
|
||||
int in_completion; /* The user pressed TAB and we are now in completion
|
||||
* mode, so input is handled by completeLine(). */
|
||||
size_t completion_idx; /* Index of next completion to propose. */
|
||||
int ifd; /* Terminal stdin file descriptor. */
|
||||
int ofd; /* Terminal stdout file descriptor. */
|
||||
char *buf; /* Edited line buffer. */
|
||||
size_t buflen; /* Edited line buffer size. */
|
||||
size_t buflen_max; /* Max buffer size, or 0 if fixed. */
|
||||
const char *prompt; /* Prompt to display. */
|
||||
size_t plen; /* Prompt length. */
|
||||
size_t pos; /* Current cursor position. */
|
||||
size_t oldpos; /* Previous refresh cursor position. */
|
||||
size_t len; /* Current edited line length. */
|
||||
size_t cols; /* Number of columns in terminal. */
|
||||
size_t oldrows; /* Rows used by last refreshed line (multiline mode) */
|
||||
size_t oldstatusrows; /* Extra status rows used by last refresh. */
|
||||
int oldrpos; /* Cursor row from last refresh (for multiline clearing). */
|
||||
int history_index; /* The history index we are currently editing. */
|
||||
int fold_count; /* Number of folded ranges. */
|
||||
size_t fold_start[LINENOISE_MAX_FOLDS]; /* Folded range start offsets. */
|
||||
size_t fold_end[LINENOISE_MAX_FOLDS]; /* Folded range end offsets. */
|
||||
char *status; /* Optional status/footer rendered below the prompt. */
|
||||
char *status_start; /* Optional escape sequence emitted before status. */
|
||||
char *status_end; /* Optional escape sequence emitted after status. */
|
||||
linenoiseLayoutCallback *layout_callback; /* Called before refresh writes. */
|
||||
void *layout_privdata;
|
||||
int screen_cursor_row; /* Absolute terminal cursor row after last refresh. */
|
||||
int screen_cursor_col; /* Absolute terminal cursor col after last refresh. */
|
||||
char *queued_input; /* Bytes already read by an outer event loop. */
|
||||
size_t queued_input_len;
|
||||
size_t queued_input_pos;
|
||||
size_t queued_input_cap;
|
||||
};
|
||||
|
||||
typedef struct linenoiseCompletions {
|
||||
size_t len;
|
||||
char **cvec;
|
||||
} linenoiseCompletions;
|
||||
|
||||
/* Non blocking API. */
|
||||
int linenoiseEditStart(struct linenoiseState *l, int stdin_fd, int stdout_fd, char *buf, size_t buflen, const char *prompt);
|
||||
char *linenoiseEditFeed(struct linenoiseState *l);
|
||||
char *linenoiseEditFeedByte(struct linenoiseState *l, char c);
|
||||
int linenoiseEditQueueInput(struct linenoiseState *l, const char *buf, size_t len);
|
||||
size_t linenoiseEditQueuedInput(struct linenoiseState *l);
|
||||
void linenoiseEditClear(struct linenoiseState *l);
|
||||
int linenoiseEditSetStatus(struct linenoiseState *l, const char *status,
|
||||
const char *start_escape, const char *end_escape);
|
||||
void linenoiseEditSetLayoutCallback(struct linenoiseState *l,
|
||||
linenoiseLayoutCallback *fn,
|
||||
void *privdata);
|
||||
void linenoiseEditStop(struct linenoiseState *l);
|
||||
void linenoiseHide(struct linenoiseState *l);
|
||||
void linenoiseShow(struct linenoiseState *l);
|
||||
|
||||
/* Blocking API. */
|
||||
char *linenoise(const char *prompt);
|
||||
void linenoiseFree(void *ptr);
|
||||
|
||||
/* Completion API. */
|
||||
typedef void(linenoiseCompletionCallback)(const char *, linenoiseCompletions *);
|
||||
typedef char*(linenoiseHintsCallback)(const char *, int *color, int *bold);
|
||||
typedef void(linenoiseFreeHintsCallback)(void *);
|
||||
void linenoiseSetCompletionCallback(linenoiseCompletionCallback *);
|
||||
void linenoiseSetHintsCallback(linenoiseHintsCallback *);
|
||||
void linenoiseSetFreeHintsCallback(linenoiseFreeHintsCallback *);
|
||||
void linenoiseAddCompletion(linenoiseCompletions *, const char *);
|
||||
|
||||
/* History API. */
|
||||
int linenoiseHistoryAdd(const char *line);
|
||||
int linenoiseHistorySetMaxLen(int len);
|
||||
int linenoiseHistorySave(const char *filename);
|
||||
int linenoiseHistoryLoad(const char *filename);
|
||||
|
||||
/* Other utilities. */
|
||||
void linenoiseClearScreen(void);
|
||||
void linenoiseSetMultiLine(int ml);
|
||||
void linenoisePrintKeyCodes(void);
|
||||
void linenoiseMaskModeEnable(void);
|
||||
void linenoiseMaskModeDisable(void);
|
||||
/* Re-enable raw mode if a child process changed the terminal state. */
|
||||
void linenoiseRestoreRawMode(void);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif /* __LINENOISE_H */
|
||||
@@ -0,0 +1,266 @@
|
||||
struct ds4_metal_args_argsort {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
int32_t top_k;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_argsort_merge {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
int32_t top_k;
|
||||
int32_t len;
|
||||
};
|
||||
|
||||
typedef void (argsort_t)(
|
||||
constant ds4_metal_args_argsort & args,
|
||||
device const char * src0,
|
||||
device int32_t * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
// Sort one float row into an index row. DS4 only exports the descending
|
||||
// instance because router and indexer selection both need top-k order.
|
||||
template<ds4_sort_order order>
|
||||
kernel void kernel_argsort_f32_i32(
|
||||
constant ds4_metal_args_argsort & args,
|
||||
device const char * src0,
|
||||
device int32_t * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
// bitonic sort
|
||||
const int col = tpitg[0];
|
||||
const int ib = tgpig[0] / args.ne01;
|
||||
|
||||
const int i00 = ib*ntg.x;
|
||||
const int i01 = tgpig[0] % args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
|
||||
|
||||
// initialize indices
|
||||
shmem_i32[col] = i00 + col;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int k = 2; k <= ntg.x; k *= 2) {
|
||||
for (int j = k / 2; j > 0; j /= 2) {
|
||||
int ixj = col ^ j;
|
||||
if (ixj > col) {
|
||||
if ((col & k) == 0) {
|
||||
if (shmem_i32[col] >= args.ne00 ||
|
||||
(shmem_i32[ixj] < args.ne00 && (order == DS4_SORT_ORDER_ASC ?
|
||||
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] :
|
||||
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]]))
|
||||
) {
|
||||
SWAP(shmem_i32[col], shmem_i32[ixj]);
|
||||
}
|
||||
} else {
|
||||
if (shmem_i32[ixj] >= args.ne00 ||
|
||||
(shmem_i32[col] < args.ne00 && (order == DS4_SORT_ORDER_ASC ?
|
||||
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] :
|
||||
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]]))
|
||||
) {
|
||||
SWAP(shmem_i32[col], shmem_i32[ixj]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
|
||||
const int64_t i0 = ib*args.top_k;
|
||||
|
||||
// copy the result to dst without the padding
|
||||
if (i0 + col < args.ne0 && col < args.top_k) {
|
||||
dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03;
|
||||
|
||||
dst[col] = shmem_i32[col];
|
||||
}
|
||||
}
|
||||
|
||||
// Host-visible sort variant used by DS4 top-k selection.
|
||||
template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32<DS4_SORT_ORDER_DESC>;
|
||||
|
||||
typedef void (argsort_merge_t)(
|
||||
constant ds4_metal_args_argsort_merge & args,
|
||||
device const char * src0,
|
||||
device const int32_t * tmp,
|
||||
device int32_t * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
// Merges sorted index runs produced by kernel_argsort_f32_i32. In the DS4 graph
|
||||
// this finishes top-k over router or compressed-attention score rows.
|
||||
template<ds4_sort_order order>
|
||||
kernel void kernel_argsort_merge_f32_i32(
|
||||
constant ds4_metal_args_argsort_merge & args,
|
||||
device const char * src0,
|
||||
device const int32_t * tmp,
|
||||
device int32_t * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int im = tgpig[0] / args.ne01;
|
||||
const int i01 = tgpig[0] % args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
const int start = im * (2 * args.len);
|
||||
|
||||
const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
|
||||
const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
|
||||
|
||||
const int total = len0 + len1;
|
||||
|
||||
device const int32_t * tmp0 = tmp + start
|
||||
+ i01*args.ne0
|
||||
+ i02*args.ne0*args.ne01
|
||||
+ i03*args.ne0*args.ne01*args.ne02;
|
||||
|
||||
device const int32_t * tmp1 = tmp0 + args.len;
|
||||
|
||||
dst += start
|
||||
+ i01*args.top_k
|
||||
+ i02*args.top_k*args.ne01
|
||||
+ i03*args.top_k*args.ne01*args.ne02;
|
||||
|
||||
device const float * src0_row = (device const float *)(src0
|
||||
+ args.nb01*i01
|
||||
+ args.nb02*i02
|
||||
+ args.nb03*i03);
|
||||
|
||||
if (total == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int chunk = (total + ntg.x - 1) / ntg.x;
|
||||
|
||||
const int k0 = tpitg.x * chunk;
|
||||
const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
|
||||
|
||||
if (k0 >= args.top_k) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (k0 >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int low = k0 > len1 ? k0 - len1 : 0;
|
||||
int high = MIN(k0, len0);
|
||||
|
||||
// binary-search partition (i, j) such that i + j = k
|
||||
while (low < high) {
|
||||
const int mid = (low + high) >> 1;
|
||||
|
||||
const int32_t idx0 = tmp0[mid];
|
||||
const int32_t idx1 = tmp1[k0 - mid - 1];
|
||||
|
||||
const float val0 = src0_row[idx0];
|
||||
const float val1 = src0_row[idx1];
|
||||
|
||||
bool take_left;
|
||||
if (order == DS4_SORT_ORDER_ASC) {
|
||||
take_left = (val0 <= val1);
|
||||
} else {
|
||||
take_left = (val0 >= val1);
|
||||
}
|
||||
|
||||
if (take_left) {
|
||||
low = mid + 1;
|
||||
} else {
|
||||
high = mid;
|
||||
}
|
||||
}
|
||||
|
||||
int i = low;
|
||||
int j = k0 - i;
|
||||
|
||||
// keep the merge fronts into registers
|
||||
int32_t idx0 = 0;
|
||||
float val0 = 0.0f;
|
||||
if (i < len0) {
|
||||
idx0 = tmp0[i];
|
||||
val0 = src0_row[idx0];
|
||||
}
|
||||
|
||||
int32_t idx1 = 0;
|
||||
float val1 = 0.0f;
|
||||
if (j < len1) {
|
||||
idx1 = tmp1[j];
|
||||
val1 = src0_row[idx1];
|
||||
}
|
||||
|
||||
for (int k = k0; k < k1; ++k) {
|
||||
int32_t out_idx;
|
||||
|
||||
if (i >= len0) {
|
||||
while (k < k1) {
|
||||
dst[k++] = tmp1[j++];
|
||||
}
|
||||
break;
|
||||
} else if (j >= len1) {
|
||||
while (k < k1) {
|
||||
dst[k++] = tmp0[i++];
|
||||
}
|
||||
break;
|
||||
} else {
|
||||
bool take_left;
|
||||
|
||||
if (order == DS4_SORT_ORDER_ASC) {
|
||||
take_left = (val0 <= val1);
|
||||
} else {
|
||||
take_left = (val0 >= val1);
|
||||
}
|
||||
|
||||
if (take_left) {
|
||||
out_idx = idx0;
|
||||
++i;
|
||||
if (i < len0) {
|
||||
idx0 = tmp0[i];
|
||||
val0 = src0_row[idx0];
|
||||
}
|
||||
} else {
|
||||
out_idx = idx1;
|
||||
++j;
|
||||
if (j < len1) {
|
||||
idx1 = tmp1[j];
|
||||
val1 = src0_row[idx1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
dst[k] = out_idx;
|
||||
}
|
||||
}
|
||||
|
||||
// Host-visible merge variant used by DS4 top-k selection.
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC>;
|
||||
+192
@@ -0,0 +1,192 @@
|
||||
struct ds4_metal_args_bin {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne10;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
int32_t ne13;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb13;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
uint64_t offs;
|
||||
uint64_t o1[8];
|
||||
};
|
||||
|
||||
constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
|
||||
constant short FC_bin_f [[function_constant(FC_BIN + 1)]];
|
||||
constant bool FC_bin_rb [[function_constant(FC_BIN + 2)]];
|
||||
constant bool FC_bin_cb [[function_constant(FC_BIN + 3)]];
|
||||
|
||||
// Generic binary elementwise op with compile-time operation and broadcast
|
||||
// modes. DS4 currently instantiates this as add, multiply, scalar multiply, and
|
||||
// row division in the static graph.
|
||||
template <typename T0, typename T1, typename T>
|
||||
kernel void kernel_bin_fuse_impl(
|
||||
constant ds4_metal_args_bin & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_bin_op
|
||||
#define FC_F FC_bin_f
|
||||
#define FC_RB FC_bin_rb
|
||||
#define FC_CB FC_bin_cb
|
||||
|
||||
if (FC_RB) {
|
||||
const uint i0 = tgpig.y*args.ne00 + tgpig.x;
|
||||
const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
|
||||
|
||||
device const T0 * src0_row = (device const T0 *) (src0);
|
||||
device T * dst_row = (device T *) (dst);
|
||||
|
||||
if (FC_F == 1) {
|
||||
device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]);
|
||||
|
||||
if (FC_OP == 0) {
|
||||
dst_row[i0] = src0_row[i0] + src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
dst_row[i0] = src0_row[i0] - src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
dst_row[i0] = src0_row[i0] * src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
dst_row[i0] = src0_row[i0] / src1_row[i1];
|
||||
}
|
||||
} else {
|
||||
T0 res = src0_row[i0];
|
||||
|
||||
if (FC_OP == 0) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res += ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res -= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res *= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res /= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
dst_row[i0] = res;
|
||||
}
|
||||
} else {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int i01 = tgpig.x;
|
||||
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int i13 = i03%args.ne13;
|
||||
const int i12 = i02%args.ne12;
|
||||
const int i11 = i01%args.ne11;
|
||||
|
||||
device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs);
|
||||
device T * dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 + args.offs);
|
||||
|
||||
if (FC_F == 1) {
|
||||
device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i10 = FC_CB ? i0%args.ne10 : i0;
|
||||
|
||||
if (FC_OP == 0) {
|
||||
dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
device const T1 * src1_ptr[8];
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
|
||||
}
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i10 = FC_CB ? i0%args.ne10 : i0;
|
||||
|
||||
T res = src0_ptr[i0];
|
||||
|
||||
if (FC_OP == 0) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res += src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res -= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res *= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res /= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
dst_ptr[i0] = res;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
#undef FC_F
|
||||
#undef FC_RB
|
||||
#undef FC_CB
|
||||
}
|
||||
|
||||
typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
|
||||
// Host-visible F32 binary op; function constants specialize it per use site.
|
||||
template [[host_name("kernel_bin_fuse_f32_f32_f32")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float, float, float>;
|
||||
@@ -0,0 +1,62 @@
|
||||
// DS4 Metal concat kernel used by the graph.
|
||||
|
||||
struct ds4_metal_args_concat {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne10;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
int32_t ne13;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb13;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
int32_t dim;
|
||||
};
|
||||
|
||||
// Concatenates two float tensors along one dimension. In DS4 this is a graph
|
||||
// utility for assembling attention inputs with exactly the same tensor layout
|
||||
// expected by the downstream kernels.
|
||||
kernel void kernel_concat(
|
||||
constant ds4_metal_args_concat & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
int o[4] = {0, 0, 0, 0};
|
||||
o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
|
||||
|
||||
device const float * x;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
|
||||
x = (device const float *)(src0 + (i3 )*args.nb03 + (i2 )*args.nb02 + (i1 )*args.nb01 + (i0 )*args.nb00);
|
||||
} else {
|
||||
x = (device const float *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
|
||||
}
|
||||
|
||||
device float * y = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
*y = *x;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
struct ds4_metal_args_cpy {
|
||||
int64_t nk0;
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
int64_t ne2;
|
||||
int64_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
// Typed copy/conversion between graph tensors. DS4 uses this for layout
|
||||
// materialization and F32/F16 conversions at graph boundaries such as KV/cache
|
||||
// packing and compressor pooling.
|
||||
template<typename T0, typename T1>
|
||||
kernel void kernel_cpy_t_t(
|
||||
constant ds4_metal_args_cpy & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i03 = tgpig[2];
|
||||
const int i02 = tgpig[1];
|
||||
const int i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tiitg/ntg[0];
|
||||
const int iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
|
||||
|
||||
const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
|
||||
|
||||
const int64_t i3 = n/(args.ne2*args.ne1*args.ne0);
|
||||
const int64_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
|
||||
const int64_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
|
||||
const int64_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
|
||||
|
||||
device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
for (int64_t i00 = iw0*ntg[0] + tiitg%ntg[0]; i00 < args.ne00; ) {
|
||||
device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
|
||||
dst_data[i00] = (T1) src[0];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cpy_t_t<float, float>) kernel_cpy_t;
|
||||
// Host-visible copy/conversion variants used by the DS4 graph.
|
||||
template [[host_name("kernel_cpy_f32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<float, float>;
|
||||
template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<float, half>;
|
||||
template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<half, float>;
|
||||
+1600
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,885 @@
|
||||
struct ds4_metal_args_dsv4_hc_split_sinkhorn {
|
||||
int32_t n_hc;
|
||||
int32_t sinkhorn_iters;
|
||||
int64_t n_rows;
|
||||
int64_t mix_hc;
|
||||
uint64_t nb01;
|
||||
uint64_t nb1;
|
||||
float eps;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_hc_weighted_sum {
|
||||
int64_t n_embd;
|
||||
int64_t n_hc;
|
||||
int64_t n_tokens;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_x2;
|
||||
uint64_t nb_w0;
|
||||
uint64_t nb_w1;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_hc_split_weighted_sum {
|
||||
int64_t n_embd;
|
||||
int32_t n_hc;
|
||||
int32_t sinkhorn_iters;
|
||||
int64_t n_rows;
|
||||
int64_t mix_hc;
|
||||
uint64_t nb_mix1;
|
||||
uint64_t nb_split1;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_x2;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
float eps;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_hc_split_weighted_sum_norm {
|
||||
int64_t n_embd;
|
||||
int32_t n_hc;
|
||||
int32_t sinkhorn_iters;
|
||||
int64_t n_rows;
|
||||
int64_t mix_hc;
|
||||
uint64_t nb_mix1;
|
||||
uint64_t nb_split1;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_x2;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb_norm1;
|
||||
float eps;
|
||||
float norm_eps;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_hc_expand {
|
||||
int64_t n_embd;
|
||||
int64_t n_hc;
|
||||
int64_t n_tokens;
|
||||
uint64_t nb_block0;
|
||||
uint64_t nb_block1;
|
||||
uint64_t nb_add0;
|
||||
uint64_t nb_add1;
|
||||
uint64_t nb_res0;
|
||||
uint64_t nb_res1;
|
||||
uint64_t nb_res2;
|
||||
uint64_t nb_post0;
|
||||
uint64_t nb_post1;
|
||||
uint64_t nb_comb0;
|
||||
uint64_t nb_comb1;
|
||||
uint64_t nb_comb2;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
int32_t has_add;
|
||||
};
|
||||
|
||||
// Numerically stable sigmoid for the standalone split/sinkhorn path. The naive
|
||||
// form 1/(1+exp(-z)) overflows for large negative z (exp(-z) blows up);
|
||||
// replacing it with the 0.5*(tanh(z/2)+1) identity keeps the value bounded in
|
||||
// [0, 1] across the entire float range. Gated by DS4_METAL_HC_STABLE so we can
|
||||
// A/B vs the historical form on M5 Max where the faster ALU is more likely to
|
||||
// push HC mixer inputs into the unstable regime.
|
||||
//
|
||||
// Do not automatically use these helpers in the fused HC decode kernels below:
|
||||
// routing the fused vector sites through the tanh form produced non-finite
|
||||
// logits on M5 Max, while the historical inline exp form remains finite and is
|
||||
// the decode throughput baseline.
|
||||
#ifdef DS4_METAL_HC_STABLE
|
||||
static inline float ds4_hc_sigmoid(float z) { return 0.5f * tanh(0.5f * z) + 0.5f; }
|
||||
static inline float4 ds4_hc_sigmoid(float4 z) { return 0.5f * tanh(0.5f * z) + 0.5f; }
|
||||
// 2 * sigmoid(z) == 1 + tanh(z/2).
|
||||
static inline float ds4_hc_twice_sigmoid(float z) { return 1.0f + tanh(0.5f * z); }
|
||||
static inline float4 ds4_hc_twice_sigmoid(float4 z) { return 1.0f + tanh(0.5f * z); }
|
||||
#else
|
||||
static inline float ds4_hc_sigmoid(float z) { return 1.0f / (1.0f + exp(-z)); }
|
||||
static inline float4 ds4_hc_sigmoid(float4 z) { return 1.0f / (1.0f + exp(-z)); }
|
||||
static inline float ds4_hc_twice_sigmoid(float z) { return 2.0f / (1.0f + exp(-z)); }
|
||||
static inline float4 ds4_hc_twice_sigmoid(float4 z) { return 2.0f / (1.0f + exp(-z)); }
|
||||
#endif
|
||||
|
||||
// Splits an HC mixer row into pre weights, post gates, and the HC-to-HC
|
||||
// combination matrix. The 4-channel path is specialized because DS4 Flash uses
|
||||
// HC=4 in normal inference, while the scalar fallback keeps diagnostics usable.
|
||||
kernel void kernel_dsv4_hc_split_sinkhorn(
|
||||
constant ds4_metal_args_dsv4_hc_split_sinkhorn & args,
|
||||
device const float * mixes,
|
||||
device const float * scale,
|
||||
device const float * base,
|
||||
device float * dst,
|
||||
uint tid [[thread_position_in_grid]]) {
|
||||
if ((int64_t) tid >= args.n_rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
constexpr int HC_MAX = 16;
|
||||
const int HC = args.n_hc;
|
||||
if (HC <= 0 || HC > HC_MAX) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const float * mix = mixes + ((int64_t) tid)*args.mix_hc;
|
||||
device float * out = dst + ((int64_t) tid)*args.mix_hc;
|
||||
|
||||
const float epsv = args.eps;
|
||||
const float pre_scale = scale[0];
|
||||
const float post_scale = scale[1];
|
||||
const float comb_scale = scale[2];
|
||||
|
||||
if (HC == 4) {
|
||||
const float4 pre_z =
|
||||
*((device const float4 *) mix) * pre_scale +
|
||||
*((device const float4 *) base);
|
||||
*((device float4 *) out) = ds4_hc_sigmoid(pre_z) + epsv;
|
||||
|
||||
const float4 post_z =
|
||||
*((device const float4 *) (mix + 4)) * post_scale +
|
||||
*((device const float4 *) (base + 4));
|
||||
*((device float4 *) (out + 4)) = ds4_hc_twice_sigmoid(post_z);
|
||||
|
||||
float4 r0 =
|
||||
*((device const float4 *) (mix + 8)) * comb_scale +
|
||||
*((device const float4 *) (base + 8));
|
||||
float4 r1 =
|
||||
*((device const float4 *) (mix + 12)) * comb_scale +
|
||||
*((device const float4 *) (base + 12));
|
||||
float4 r2 =
|
||||
*((device const float4 *) (mix + 16)) * comb_scale +
|
||||
*((device const float4 *) (base + 16));
|
||||
float4 r3 =
|
||||
*((device const float4 *) (mix + 20)) * comb_scale +
|
||||
*((device const float4 *) (base + 20));
|
||||
|
||||
const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w));
|
||||
const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w));
|
||||
const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w));
|
||||
const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w));
|
||||
|
||||
r0 = exp(r0 - m0);
|
||||
r1 = exp(r1 - m1);
|
||||
r2 = exp(r2 - m2);
|
||||
r3 = exp(r3 - m3);
|
||||
|
||||
r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv;
|
||||
r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv;
|
||||
r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv;
|
||||
r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv;
|
||||
|
||||
float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
|
||||
for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
|
||||
r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv);
|
||||
r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv);
|
||||
r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv);
|
||||
r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv);
|
||||
|
||||
col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
}
|
||||
|
||||
*((device float4 *) (out + 8)) = r0;
|
||||
*((device float4 *) (out + 12)) = r1;
|
||||
*((device float4 *) (out + 16)) = r2;
|
||||
*((device float4 *) (out + 20)) = r3;
|
||||
return;
|
||||
}
|
||||
|
||||
for (int i = 0; i < HC; ++i) {
|
||||
const float z = mix[i] * pre_scale + base[i];
|
||||
out[i] = ds4_hc_sigmoid(z) + epsv;
|
||||
}
|
||||
|
||||
for (int i = 0; i < HC; ++i) {
|
||||
const int off = HC + i;
|
||||
const float z = mix[off] * post_scale + base[off];
|
||||
out[off] = ds4_hc_twice_sigmoid(z);
|
||||
}
|
||||
|
||||
float c[HC_MAX*HC_MAX];
|
||||
|
||||
for (int dst_hc = 0; dst_hc < HC; ++dst_hc) {
|
||||
float row_max = -INFINITY;
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
const int idx = src_hc + dst_hc*HC;
|
||||
const int off = 2*HC + idx;
|
||||
const float v = mix[off] * comb_scale + base[off];
|
||||
c[idx] = v;
|
||||
row_max = max(row_max, v);
|
||||
}
|
||||
|
||||
float row_sum = 0.0f;
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
const int idx = src_hc + dst_hc*HC;
|
||||
const float v = exp(c[idx] - row_max);
|
||||
c[idx] = v;
|
||||
row_sum += v;
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / row_sum;
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
const int idx = src_hc + dst_hc*HC;
|
||||
c[idx] = c[idx] * inv_sum + epsv;
|
||||
}
|
||||
}
|
||||
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
float sum = 0.0f;
|
||||
for (int dst_hc = 0; dst_hc < HC; ++dst_hc) {
|
||||
sum += c[src_hc + dst_hc*HC];
|
||||
}
|
||||
|
||||
const float inv_denom = 1.0f / (sum + epsv);
|
||||
for (int dst_hc = 0; dst_hc < HC; ++dst_hc) {
|
||||
c[src_hc + dst_hc*HC] *= inv_denom;
|
||||
}
|
||||
}
|
||||
|
||||
for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
|
||||
for (int dst_hc = 0; dst_hc < HC; ++dst_hc) {
|
||||
float sum = 0.0f;
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
sum += c[src_hc + dst_hc*HC];
|
||||
}
|
||||
|
||||
const float inv_denom = 1.0f / (sum + epsv);
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
c[src_hc + dst_hc*HC] *= inv_denom;
|
||||
}
|
||||
}
|
||||
|
||||
for (int src_hc = 0; src_hc < HC; ++src_hc) {
|
||||
float sum = 0.0f;
|
||||
for (int dst_hc = 0; dst_hc < HC; ++dst_hc) {
|
||||
sum += c[src_hc + dst_hc*HC];
|
||||
}
|
||||
|
||||
const float inv_denom = 1.0f / (sum + epsv);
|
||||
for (int dst_hc = 0; dst_hc < HC; ++dst_hc) {
|
||||
c[src_hc + dst_hc*HC] *= inv_denom;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < HC*HC; ++i) {
|
||||
out[2*HC + i] = c[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Decode-side fusion of HC split and pre-weighted HC reduction. One threadgroup
|
||||
// handles one token row: lane 0 computes the HC=4 mixer split once, stores the
|
||||
// post/comb data for the following HC expand, and all lanes reuse the pre
|
||||
// weights from threadgroup memory to produce the embedding row.
|
||||
kernel void kernel_dsv4_hc_split_weighted_sum(
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum & args,
|
||||
device const char * mixes,
|
||||
device const float * scale,
|
||||
device const float * base,
|
||||
device const char * x,
|
||||
device char * split,
|
||||
device char * dst,
|
||||
threadgroup float * pre_shmem [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]],
|
||||
uint ntg [[threads_per_threadgroup]]) {
|
||||
if ((int64_t) row >= args.n_rows || args.n_hc != 4) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const float * mix = (device const float *) (mixes + (uint64_t)row*args.nb_mix1);
|
||||
device float * out = (device float *) (split + (uint64_t)row*args.nb_split1);
|
||||
|
||||
if (tid == 0) {
|
||||
const float epsv = args.eps;
|
||||
const float pre_scale = scale[0];
|
||||
const float post_scale = scale[1];
|
||||
const float comb_scale = scale[2];
|
||||
|
||||
const float4 pre_z =
|
||||
*((device const float4 *) mix) * pre_scale +
|
||||
*((device const float4 *) base);
|
||||
const float4 pre = 1.0f / (1.0f + exp(-pre_z)) + epsv;
|
||||
*((device float4 *) out) = pre;
|
||||
pre_shmem[0] = pre.x;
|
||||
pre_shmem[1] = pre.y;
|
||||
pre_shmem[2] = pre.z;
|
||||
pre_shmem[3] = pre.w;
|
||||
|
||||
const float4 post_z =
|
||||
*((device const float4 *) (mix + 4)) * post_scale +
|
||||
*((device const float4 *) (base + 4));
|
||||
*((device float4 *) (out + 4)) = 2.0f / (1.0f + exp(-post_z));
|
||||
|
||||
float4 r0 =
|
||||
*((device const float4 *) (mix + 8)) * comb_scale +
|
||||
*((device const float4 *) (base + 8));
|
||||
float4 r1 =
|
||||
*((device const float4 *) (mix + 12)) * comb_scale +
|
||||
*((device const float4 *) (base + 12));
|
||||
float4 r2 =
|
||||
*((device const float4 *) (mix + 16)) * comb_scale +
|
||||
*((device const float4 *) (base + 16));
|
||||
float4 r3 =
|
||||
*((device const float4 *) (mix + 20)) * comb_scale +
|
||||
*((device const float4 *) (base + 20));
|
||||
|
||||
const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w));
|
||||
const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w));
|
||||
const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w));
|
||||
const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w));
|
||||
|
||||
r0 = exp(r0 - m0);
|
||||
r1 = exp(r1 - m1);
|
||||
r2 = exp(r2 - m2);
|
||||
r3 = exp(r3 - m3);
|
||||
|
||||
r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv;
|
||||
r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv;
|
||||
r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv;
|
||||
r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv;
|
||||
|
||||
float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
|
||||
for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
|
||||
r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv);
|
||||
r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv);
|
||||
r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv);
|
||||
r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv);
|
||||
|
||||
col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
}
|
||||
|
||||
*((device float4 *) (out + 8)) = r0;
|
||||
*((device float4 *) (out + 12)) = r1;
|
||||
*((device float4 *) (out + 16)) = r2;
|
||||
*((device float4 *) (out + 20)) = r3;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int64_t d = tid; d < args.n_embd; d += ntg) {
|
||||
float acc = 0.0f;
|
||||
acc += *((device const float *) (x + d*args.nb_x0 + 0*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[0];
|
||||
acc += *((device const float *) (x + d*args.nb_x0 + 1*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[1];
|
||||
acc += *((device const float *) (x + d*args.nb_x0 + 2*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[2];
|
||||
acc += *((device const float *) (x + d*args.nb_x0 + 3*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[3];
|
||||
*((device float *) (dst + d*args.nb0 + (uint64_t)row*args.nb1)) = acc;
|
||||
}
|
||||
}
|
||||
|
||||
// Decode HC-pre plus the following RMSNorm. DS4 uses HC=4 here. The normal
|
||||
// release path computes HC coefficients, collapses four residual streams into
|
||||
// the model row, then immediately launches a weighted RMSNorm over the row.
|
||||
// This kernel keeps the HC split math identical to
|
||||
// kernel_dsv4_hc_split_weighted_sum, stores the HC-pre row for diagnostics, and
|
||||
// reuses the just-collapsed values from threadgroup memory for the RMSNorm
|
||||
// reduction.
|
||||
kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args,
|
||||
device const char * mixes,
|
||||
device const float * scale,
|
||||
device const float * base,
|
||||
device const char * x,
|
||||
device char * split,
|
||||
device char * dst,
|
||||
device const char * norm_weight,
|
||||
device char * norm_dst,
|
||||
threadgroup float * shared [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
ushort tid [[thread_position_in_threadgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort ntg [[threads_per_threadgroup]]) {
|
||||
if ((int64_t)row >= args.n_rows || args.n_hc != 4 || (args.n_embd & 3) != 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint n_embd = uint(args.n_embd);
|
||||
const uint n4 = n_embd >> 2;
|
||||
threadgroup float4 *row_shmem = (threadgroup float4 *)shared;
|
||||
threadgroup float *pre_shmem = shared + n_embd;
|
||||
threadgroup float *sum_shmem = pre_shmem + 4;
|
||||
|
||||
device const float *mix = (device const float *)(mixes + (uint64_t)row * args.nb_mix1);
|
||||
device float *out = (device float *)(split + (uint64_t)row * args.nb_split1);
|
||||
|
||||
if (sgitg == 0) {
|
||||
sum_shmem[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
const float epsv = args.eps;
|
||||
const float pre_scale = scale[0];
|
||||
const float post_scale = scale[1];
|
||||
const float comb_scale = scale[2];
|
||||
|
||||
const float4 pre_z =
|
||||
*((device const float4 *)mix) * pre_scale +
|
||||
*((device const float4 *)base);
|
||||
const float4 pre = 1.0f / (1.0f + exp(-pre_z)) + epsv;
|
||||
*((device float4 *)out) = pre;
|
||||
pre_shmem[0] = pre.x;
|
||||
pre_shmem[1] = pre.y;
|
||||
pre_shmem[2] = pre.z;
|
||||
pre_shmem[3] = pre.w;
|
||||
|
||||
const float4 post_z =
|
||||
*((device const float4 *)(mix + 4)) * post_scale +
|
||||
*((device const float4 *)(base + 4));
|
||||
*((device float4 *)(out + 4)) = 2.0f / (1.0f + exp(-post_z));
|
||||
|
||||
float4 r0 =
|
||||
*((device const float4 *)(mix + 8)) * comb_scale +
|
||||
*((device const float4 *)(base + 8));
|
||||
float4 r1 =
|
||||
*((device const float4 *)(mix + 12)) * comb_scale +
|
||||
*((device const float4 *)(base + 12));
|
||||
float4 r2 =
|
||||
*((device const float4 *)(mix + 16)) * comb_scale +
|
||||
*((device const float4 *)(base + 16));
|
||||
float4 r3 =
|
||||
*((device const float4 *)(mix + 20)) * comb_scale +
|
||||
*((device const float4 *)(base + 20));
|
||||
|
||||
const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w));
|
||||
const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w));
|
||||
const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w));
|
||||
const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w));
|
||||
|
||||
r0 = exp(r0 - m0);
|
||||
r1 = exp(r1 - m1);
|
||||
r2 = exp(r2 - m2);
|
||||
r3 = exp(r3 - m3);
|
||||
|
||||
r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv;
|
||||
r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv;
|
||||
r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv;
|
||||
r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv;
|
||||
|
||||
float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
|
||||
for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
|
||||
r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv);
|
||||
r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv);
|
||||
r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv);
|
||||
r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv);
|
||||
|
||||
col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
}
|
||||
|
||||
*((device float4 *)(out + 8)) = r0;
|
||||
*((device float4 *)(out + 12)) = r1;
|
||||
*((device float4 *)(out + 16)) = r2;
|
||||
*((device float4 *)(out + 20)) = r3;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (uint i = tid; i < n4; i += ntg) {
|
||||
device const float4 *x0 = (device const float4 *)(x + 0 * args.nb_x1 + (uint64_t)row * args.nb_x2);
|
||||
device const float4 *x1 = (device const float4 *)(x + 1 * args.nb_x1 + (uint64_t)row * args.nb_x2);
|
||||
device const float4 *x2 = (device const float4 *)(x + 2 * args.nb_x1 + (uint64_t)row * args.nb_x2);
|
||||
device const float4 *x3 = (device const float4 *)(x + 3 * args.nb_x1 + (uint64_t)row * args.nb_x2);
|
||||
const float4 v = x0[i] * pre_shmem[0] +
|
||||
x1[i] * pre_shmem[1] +
|
||||
x2[i] * pre_shmem[2] +
|
||||
x3[i] * pre_shmem[3];
|
||||
row_shmem[i] = v;
|
||||
sumf += dot(v, v);
|
||||
}
|
||||
|
||||
sumf = simd_sum(sumf);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiisg == 0) {
|
||||
sum_shmem[sgitg] = sumf;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = sum_shmem[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
const float norm_scale = rsqrt(sumf / float(n_embd) + args.norm_eps);
|
||||
|
||||
device float4 *dst4 = (device float4 *)(dst + (uint64_t)row * args.nb1);
|
||||
device const float4 *w4 = (device const float4 *)norm_weight;
|
||||
device float4 *norm4 = (device float4 *)(norm_dst + (uint64_t)row * args.nb_norm1);
|
||||
for (uint i = tid; i < n4; i += ntg) {
|
||||
const float4 v = row_shmem[i];
|
||||
dst4[i] = v;
|
||||
norm4[i] = (v * norm_scale) * w4[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Expands an embedding-sized block back into HC channels after attention/FFN.
|
||||
// The post gate scales the current block, while the Sinkhorn combination matrix
|
||||
// mixes residual HC channels from the previous state.
|
||||
kernel void kernel_dsv4_hc_expand(
|
||||
constant ds4_metal_args_dsv4_hc_expand & args,
|
||||
device const char * block_out,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device const char * block_add,
|
||||
device char * dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const int64_t n_elem = args.n_embd * args.n_hc * args.n_tokens;
|
||||
if ((int64_t) gid >= n_elem) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t d = ((int64_t) gid) % args.n_embd;
|
||||
const int64_t tmp = ((int64_t) gid) / args.n_embd;
|
||||
const int64_t dst_hc = tmp % args.n_hc;
|
||||
const int64_t t = tmp / args.n_hc;
|
||||
|
||||
float block_v = *((device const float *) (block_out + d*args.nb_block0 + t*args.nb_block1));
|
||||
if (args.has_add) {
|
||||
block_v += *((device const float *) (block_add + d*args.nb_add0 + t*args.nb_add1));
|
||||
}
|
||||
const float post_v = *((device const float *) (post + dst_hc*args.nb_post0 + t*args.nb_post1));
|
||||
|
||||
float acc = block_v * post_v;
|
||||
for (int64_t src_hc = 0; src_hc < args.n_hc; ++src_hc) {
|
||||
const float comb_v = *((device const float *) (comb + dst_hc*args.nb_comb0 + src_hc*args.nb_comb1 + t*args.nb_comb2));
|
||||
const float res_v = *((device const float *) (residual + d*args.nb_res0 + src_hc*args.nb_res1 + t*args.nb_res2));
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
|
||||
*((device float *) (dst + d*args.nb0 + dst_hc*args.nb1 + t*args.nb2)) = acc;
|
||||
}
|
||||
|
||||
// HC=4 specialization of the post/expand step. One thread computes all four
|
||||
// destination HC streams for one token/dimension, reusing the same block output
|
||||
// and residual HC values while preserving the per-stream accumulation order.
|
||||
kernel void kernel_dsv4_hc_expand4(
|
||||
constant ds4_metal_args_dsv4_hc_expand & args,
|
||||
device const char * block_out,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device const char * block_add,
|
||||
device char * dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (args.n_hc != 4) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t n_elem = args.n_embd * args.n_tokens;
|
||||
if ((int64_t) gid >= n_elem) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t d = ((int64_t) gid) % args.n_embd;
|
||||
const int64_t t = ((int64_t) gid) / args.n_embd;
|
||||
|
||||
float block_v = *((device const float *) (block_out + d*args.nb_block0 + t*args.nb_block1));
|
||||
if (args.has_add) {
|
||||
block_v += *((device const float *) (block_add + d*args.nb_add0 + t*args.nb_add1));
|
||||
}
|
||||
|
||||
const float r0 = *((device const float *) (residual + d*args.nb_res0 + 0*args.nb_res1 + t*args.nb_res2));
|
||||
const float r1 = *((device const float *) (residual + d*args.nb_res0 + 1*args.nb_res1 + t*args.nb_res2));
|
||||
const float r2 = *((device const float *) (residual + d*args.nb_res0 + 2*args.nb_res1 + t*args.nb_res2));
|
||||
const float r3 = *((device const float *) (residual + d*args.nb_res0 + 3*args.nb_res1 + t*args.nb_res2));
|
||||
|
||||
for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) {
|
||||
float acc = block_v * *((device const float *) (post + dst_hc*args.nb_post0 + t*args.nb_post1));
|
||||
|
||||
acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 0*args.nb_comb1 + t*args.nb_comb2)) * r0;
|
||||
acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 1*args.nb_comb1 + t*args.nb_comb2)) * r1;
|
||||
acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 2*args.nb_comb1 + t*args.nb_comb2)) * r2;
|
||||
acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 3*args.nb_comb1 + t*args.nb_comb2)) * r3;
|
||||
|
||||
*((device float *) (dst + d*args.nb0 + dst_hc*args.nb1 + t*args.nb2)) = acc;
|
||||
}
|
||||
}
|
||||
|
||||
// Decode-time FFN tail fusion:
|
||||
//
|
||||
// shared_out = shared_mid @ Wshared_down
|
||||
// after_ffn_hc = HCPost(routed_out + shared_out, residual_hc, split)
|
||||
//
|
||||
// The Q8_0 dot reduction is intentionally copied from the normal matvec shape
|
||||
// so the shared expert result is bit-identical. The only specialization is
|
||||
// that DS4 decode has one token and HC=4, so the thread that finishes each
|
||||
// shared-down output row can immediately expand it into the four HC streams.
|
||||
kernel void kernel_dsv4_shared_down_hc_expand4_q8_0(
|
||||
constant ds4_metal_args_mul_mv & mv,
|
||||
constant ds4_metal_args_dsv4_hc_expand & hc,
|
||||
device const char * weight,
|
||||
device const char * shared_mid,
|
||||
device char * shared_out,
|
||||
device const char * routed_out,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
if (hc.n_hc != 4 || hc.n_tokens != 1) {
|
||||
return;
|
||||
}
|
||||
|
||||
const short NSG = FC_mul_mv_nsg;
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
constexpr short NQ = 8;
|
||||
constexpr short NR0 = N_R0_Q8_0;
|
||||
|
||||
const int nb = mv.ne00 / QK8_0;
|
||||
const int row0 = tgpig.x * NR0;
|
||||
|
||||
const short ix = tiisg / (NW / NQ);
|
||||
const short il = tiisg % (NW / NQ);
|
||||
const int ib0 = sgitg * NQ + ix;
|
||||
|
||||
device const float *y = (device const float *)(shared_mid);
|
||||
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
||||
|
||||
device const block_q8_0 *ax[NR0];
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01;
|
||||
ax[row] = (device const block_q8_0 *)(weight + off0);
|
||||
}
|
||||
|
||||
float sumf[NR0] = { 0.0f };
|
||||
float yl[NQ];
|
||||
|
||||
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
||||
FOR_UNROLL(short i = 0; i < NQ; ++i) {
|
||||
yl[i] = yb[i];
|
||||
}
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
device const int8_t *qs = ax[row][ib].qs + il * NQ;
|
||||
|
||||
float sumq = 0.0f;
|
||||
FOR_UNROLL(short i = 0; i < NQ; ++i) {
|
||||
sumq += qs[i] * yl[i];
|
||||
}
|
||||
|
||||
sumf[row] += sumq * ax[row][ib].d;
|
||||
}
|
||||
|
||||
yb += NSG * NQ * QK8_0;
|
||||
}
|
||||
|
||||
threadgroup float *shmem_f32[NR0];
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
shmem_f32[row] = (threadgroup float *)shmem + NW * row;
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[row][tiisg] = 0.0f;
|
||||
}
|
||||
sumf[row] = simd_sum(sumf[row]);
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[row][sgitg] = sumf[row];
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
const int d = row0 + row;
|
||||
if (d >= mv.ne01) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const float shared_v = simd_sum(shmem_f32[row][tiisg]);
|
||||
if (tiisg == 0 && sgitg == 0) {
|
||||
*((device float *)(shared_out + (uint64_t)d * sizeof(float))) = shared_v;
|
||||
|
||||
float block_v = *((device const float *)(routed_out + (uint64_t)d * hc.nb_block0));
|
||||
block_v += shared_v;
|
||||
|
||||
const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1));
|
||||
const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1));
|
||||
const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1));
|
||||
const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1));
|
||||
|
||||
for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) {
|
||||
float acc = block_v * *((device const float *)(post + dst_hc * hc.nb_post0));
|
||||
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 0 * hc.nb_comb1)) * r0;
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 1 * hc.nb_comb1)) * r1;
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 2 * hc.nb_comb1)) * r2;
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 3 * hc.nb_comb1)) * r3;
|
||||
|
||||
*((device float *)(dst + (uint64_t)d * hc.nb0 + dst_hc * hc.nb1)) = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Decode-time attention output tail fusion:
|
||||
//
|
||||
// attn_out = attn_low @ Wob
|
||||
// after_attn_hc = HCPost(attn_out, residual_hc, split)
|
||||
//
|
||||
// This is the no-add sibling of the shared-down/FFN fusion above. It preserves
|
||||
// the exact Q8_0 matvec reduction, stores `attn_out` for diagnostics, and then
|
||||
// writes the four HC streams for the same embedding dimension.
|
||||
kernel void kernel_dsv4_q8_hc_expand4_q8_0(
|
||||
constant ds4_metal_args_mul_mv & mv,
|
||||
constant ds4_metal_args_dsv4_hc_expand & hc,
|
||||
device const char * weight,
|
||||
device const char * input,
|
||||
device char * block_out,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
if (hc.n_hc != 4 || hc.n_tokens != 1) {
|
||||
return;
|
||||
}
|
||||
|
||||
const short NSG = FC_mul_mv_nsg;
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
constexpr short NQ = 8;
|
||||
constexpr short NR0 = N_R0_Q8_0;
|
||||
|
||||
const int nb = mv.ne00 / QK8_0;
|
||||
const int row0 = tgpig.x * NR0;
|
||||
|
||||
const short ix = tiisg / (NW / NQ);
|
||||
const short il = tiisg % (NW / NQ);
|
||||
const int ib0 = sgitg * NQ + ix;
|
||||
|
||||
device const float *y = (device const float *)(input);
|
||||
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
||||
|
||||
device const block_q8_0 *ax[NR0];
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01;
|
||||
ax[row] = (device const block_q8_0 *)(weight + off0);
|
||||
}
|
||||
|
||||
float sumf[NR0] = { 0.0f };
|
||||
float yl[NQ];
|
||||
|
||||
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
||||
FOR_UNROLL(short i = 0; i < NQ; ++i) {
|
||||
yl[i] = yb[i];
|
||||
}
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
device const int8_t *qs = ax[row][ib].qs + il * NQ;
|
||||
|
||||
float sumq = 0.0f;
|
||||
FOR_UNROLL(short i = 0; i < NQ; ++i) {
|
||||
sumq += qs[i] * yl[i];
|
||||
}
|
||||
|
||||
sumf[row] += sumq * ax[row][ib].d;
|
||||
}
|
||||
|
||||
yb += NSG * NQ * QK8_0;
|
||||
}
|
||||
|
||||
threadgroup float *shmem_f32[NR0];
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
shmem_f32[row] = (threadgroup float *)shmem + NW * row;
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[row][tiisg] = 0.0f;
|
||||
}
|
||||
sumf[row] = simd_sum(sumf[row]);
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[row][sgitg] = sumf[row];
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
const int d = row0 + row;
|
||||
if (d >= mv.ne01) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const float block_v = simd_sum(shmem_f32[row][tiisg]);
|
||||
if (tiisg == 0 && sgitg == 0) {
|
||||
*((device float *)(block_out + (uint64_t)d * sizeof(float))) = block_v;
|
||||
|
||||
const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1));
|
||||
const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1));
|
||||
const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1));
|
||||
const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1));
|
||||
|
||||
for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) {
|
||||
float acc = block_v * *((device const float *)(post + dst_hc * hc.nb_post0));
|
||||
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 0 * hc.nb_comb1)) * r0;
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 1 * hc.nb_comb1)) * r1;
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 2 * hc.nb_comb1)) * r2;
|
||||
acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 3 * hc.nb_comb1)) * r3;
|
||||
|
||||
*((device float *)(dst + (uint64_t)d * hc.nb0 + dst_hc * hc.nb1)) = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reduces HC channels to a normal embedding row with the learned pre weights.
|
||||
// This is the input adapter before the attention block and before the FFN block.
|
||||
kernel void kernel_dsv4_hc_weighted_sum(
|
||||
constant ds4_metal_args_dsv4_hc_weighted_sum & args,
|
||||
device const char * x,
|
||||
device const char * weights,
|
||||
device char * dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const int64_t n_elem = args.n_embd * args.n_tokens;
|
||||
if ((int64_t) gid >= n_elem) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t d = ((int64_t) gid) % args.n_embd;
|
||||
const int64_t t = ((int64_t) gid) / args.n_embd;
|
||||
|
||||
float acc = 0.0f;
|
||||
for (int64_t h = 0; h < args.n_hc; ++h) {
|
||||
const float xv = *((device const float *) (x + d*args.nb_x0 + h*args.nb_x1 + t*args.nb_x2));
|
||||
const float wv = *((device const float *) (weights + h*args.nb_w0 + t*args.nb_w1));
|
||||
acc += xv * wv;
|
||||
}
|
||||
|
||||
*((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc;
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
constant float dsv4_e4m3fn_exp_scale[16] = {
|
||||
0.0f, 0.015625f, 0.03125f, 0.0625f,
|
||||
0.125f, 0.25f, 0.5f, 1.0f,
|
||||
2.0f, 4.0f, 8.0f, 16.0f,
|
||||
32.0f, 64.0f, 128.0f, 256.0f,
|
||||
};
|
||||
|
||||
constant float dsv4_e2m1fn_values[8] = {
|
||||
0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f,
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_fp8_kv_quantize {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
ulong nb00;
|
||||
ulong nb01;
|
||||
ulong nb02;
|
||||
ulong nb03;
|
||||
ulong nb0;
|
||||
ulong nb1;
|
||||
ulong nb2;
|
||||
ulong nb3;
|
||||
int n_rot;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_kv_fp8_store {
|
||||
int32_t head_dim;
|
||||
int32_t n_rot;
|
||||
int32_t raw_row;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_indexer_qat {
|
||||
uint32_t n_rows;
|
||||
uint32_t head_dim;
|
||||
uint64_t row_stride;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_ratio4_shift {
|
||||
uint32_t width;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_compressor_store_one {
|
||||
uint32_t width;
|
||||
uint32_t ratio;
|
||||
uint32_t pos;
|
||||
uint32_t ape_type;
|
||||
};
|
||||
|
||||
static inline float dsv4_e4m3fn_value(int i) {
|
||||
const int exp = (i >> 3) & 0x0f;
|
||||
const int mant = i & 0x07;
|
||||
return exp == 0
|
||||
? float(mant) * 0.001953125f
|
||||
: (1.0f + float(mant) * 0.125f) * dsv4_e4m3fn_exp_scale[exp];
|
||||
}
|
||||
|
||||
static inline float dsv4_e4m3fn_dequant(float x) {
|
||||
const float sign = x < 0.0f ? -1.0f : 1.0f;
|
||||
const float ax = min(abs(x), 448.0f);
|
||||
|
||||
int lo = 0;
|
||||
int hi = 126;
|
||||
while (lo < hi) {
|
||||
const int mid = (lo + hi + 1) >> 1;
|
||||
if (dsv4_e4m3fn_value(mid) <= ax) {
|
||||
lo = mid;
|
||||
} else {
|
||||
hi = mid - 1;
|
||||
}
|
||||
}
|
||||
|
||||
int best = lo;
|
||||
if (best < 126) {
|
||||
const float best_diff = abs(ax - dsv4_e4m3fn_value(best));
|
||||
const float next_diff = abs(ax - dsv4_e4m3fn_value(best + 1));
|
||||
if (next_diff < best_diff || (next_diff == best_diff && ((best + 1) & 1) == 0 && (best & 1) != 0)) {
|
||||
best = best + 1;
|
||||
}
|
||||
}
|
||||
|
||||
return sign * dsv4_e4m3fn_value(best);
|
||||
}
|
||||
|
||||
static inline float dsv4_e2m1fn_dequant(float x) {
|
||||
const float sign = x < 0.0f ? -1.0f : 1.0f;
|
||||
const float ax = min(abs(x), 6.0f);
|
||||
int best = 0;
|
||||
float best_diff = abs(ax - dsv4_e2m1fn_values[0]);
|
||||
for (int i = 1; i < 8; i++) {
|
||||
const float diff = abs(ax - dsv4_e2m1fn_values[i]);
|
||||
if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) {
|
||||
best = i;
|
||||
best_diff = diff;
|
||||
}
|
||||
}
|
||||
return sign * dsv4_e2m1fn_values[best];
|
||||
}
|
||||
|
||||
// Quantizes the non-RoPE part of a KV row through E4M3FN and writes the
|
||||
// dequantized value back as float. DS4 uses this to match the FP8 KV-cache
|
||||
// semantics while keeping the Metal graph's cache buffers float-addressable.
|
||||
kernel void kernel_dsv4_fp8_kv_quantize_f32(
|
||||
constant ds4_metal_args_dsv4_fp8_kv_quantize & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
const int64_t n_rows = args.ne01 * args.ne02 * args.ne03;
|
||||
if ((int64_t) row >= n_rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t i1 = row % args.ne01;
|
||||
const int64_t i2 = (row / args.ne01) % args.ne02;
|
||||
const int64_t i3 = row / (args.ne01 * args.ne02);
|
||||
|
||||
device const char * src_base = src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03;
|
||||
device char * dst_base = dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3;
|
||||
|
||||
const int64_t n_nope = args.ne00 - args.n_rot;
|
||||
|
||||
for (int64_t off = 0; off < n_nope; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (tid < 64) {
|
||||
v = *((device const float *) (src_base + (off + tid)*args.nb00));
|
||||
scratch[tid] = abs(v);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(scratch[0], 1.0e-4f);
|
||||
const float scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (tid < 64) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale;
|
||||
*((device float *) (dst_base + (off + tid)*args.nb0)) = q;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
for (int64_t i = n_nope + tid; i < args.ne00; i += 64) {
|
||||
*((device float *) (dst_base + i*args.nb0)) = *((device const float *) (src_base + i*args.nb00));
|
||||
}
|
||||
}
|
||||
|
||||
// The official DS4 indexer applies a 128-wide Hadamard rotation and then an
|
||||
// inplace FP4 activation-simulation pass to both indexer Q and indexer KV.
|
||||
kernel void kernel_dsv4_indexer_hadamard_fp4_f32(
|
||||
constant ds4_metal_args_dsv4_indexer_qat & args,
|
||||
device char * x,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
if (row >= args.n_rows || args.head_dim != 128u || tid >= 128u) {
|
||||
return;
|
||||
}
|
||||
|
||||
threadgroup float *vals = scratch;
|
||||
threadgroup float *absbuf = scratch + 128;
|
||||
device float *xr = (device float *)(x + (uint64_t)row * args.row_stride);
|
||||
|
||||
vals[tid] = xr[tid];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 1u; stride < 128u; stride <<= 1u) {
|
||||
if ((tid & stride) == 0u) {
|
||||
const uint base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
|
||||
const float a = vals[base];
|
||||
const float b = vals[base + stride];
|
||||
vals[base] = a + b;
|
||||
vals[base + stride] = a - b;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
float v = vals[tid] * 0.08838834764831845f;
|
||||
const uint block = tid >> 5u;
|
||||
const uint lane = tid & 31u;
|
||||
const uint block_base = block * 32u;
|
||||
absbuf[tid] = abs(v);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 16u; stride > 0u; stride >>= 1u) {
|
||||
if (lane < stride) {
|
||||
absbuf[block_base + lane] = max(absbuf[block_base + lane],
|
||||
absbuf[block_base + lane + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(absbuf[block_base], 7.052966104933725e-38f);
|
||||
const float scale = exp2(ceil(log2(amax / 6.0f)));
|
||||
xr[tid] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale;
|
||||
}
|
||||
|
||||
// Decode-side KV finalizer after RoPE. The normal RoPE kernel intentionally
|
||||
// remains separate because tiny trigonometric codegen changes can flip later
|
||||
// sampled tokens. This kernel only fuses the FP8 round-trip for the non-RoPE
|
||||
// prefix with the F16-rounded raw-cache row used by FlashAttention.
|
||||
kernel void kernel_dsv4_kv_fp8_store_f32(
|
||||
constant ds4_metal_args_dsv4_kv_fp8_store & args,
|
||||
device float * kv,
|
||||
device float * raw_cache,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
const int head_dim = args.head_dim;
|
||||
const int n_rot = args.n_rot;
|
||||
const int n_nope = head_dim - n_rot;
|
||||
if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) {
|
||||
return;
|
||||
}
|
||||
|
||||
device float * raw = raw_cache + (int64_t)args.raw_row * head_dim;
|
||||
|
||||
for (int off = 0; off < n_nope; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (off + (int)tid < n_nope) {
|
||||
v = kv[off + tid];
|
||||
scratch[tid] = abs(v);
|
||||
} else {
|
||||
scratch[tid] = 0.0f;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(scratch[0], 1.0e-4f);
|
||||
const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (off + (int)tid < n_nope) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
|
||||
kv[off + tid] = q;
|
||||
// Diagnostic only: skip the FP16 round-trip that normally matches the
|
||||
// half-typed FlashAttention KV buffer's precision. With this enabled the
|
||||
// indexer will see higher-precision raw values than FlashAttention does,
|
||||
// which is informative but not a production-ready setting.
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[off + tid] = q;
|
||||
#else
|
||||
raw[off + tid] = (float)((half)q);
|
||||
#endif
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
for (int i = n_nope + tid; i < head_dim; i += 64) {
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[i] = kv[i];
|
||||
#else
|
||||
raw[i] = (float)((half)kv[i]);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
// Ratio-4 compression keeps two 4-row halves of recurrent state. After an
|
||||
// emitted compressed row, the second half becomes the next window's previous
|
||||
// half. The old encoder expressed this as four generic copies; this DS4-specific
|
||||
// kernel performs the KV and score copies together.
|
||||
kernel void kernel_dsv4_ratio4_shift_f32(
|
||||
constant ds4_metal_args_dsv4_ratio4_shift & args,
|
||||
device float * state_kv,
|
||||
device float * state_score,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint n = 4u * args.width;
|
||||
if (gid >= n) return;
|
||||
|
||||
state_kv[gid] = state_kv[n + gid];
|
||||
state_score[gid] = state_score[n + gid];
|
||||
}
|
||||
|
||||
// One-token compressor frontier update. Decode appends exactly one projected KV
|
||||
// row and one score row into a small recurrent state. The generic batch helper
|
||||
// expresses this as APE copy, score add, and two set_rows operations; this
|
||||
// kernel writes both state tensors directly while preserving the same
|
||||
// score + APE arithmetic.
|
||||
kernel void kernel_dsv4_compressor_store_one(
|
||||
constant ds4_metal_args_dsv4_compressor_store_one & args,
|
||||
device const float * kv,
|
||||
device const float * score,
|
||||
device const char * ape,
|
||||
device float * state_kv,
|
||||
device float * state_score,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid >= args.width || args.width == 0 || args.ratio == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint pos_mod = args.pos % args.ratio;
|
||||
const uint dst_row = args.ratio == 4u ? args.ratio + pos_mod : pos_mod;
|
||||
const uint dst = dst_row * args.width + gid;
|
||||
const uint ape_i = pos_mod * args.width + gid;
|
||||
|
||||
float ape_v;
|
||||
if (args.ape_type == 1u) {
|
||||
ape_v = (float)(((device const half *)ape)[ape_i]);
|
||||
} else {
|
||||
ape_v = ((device const float *)ape)[ape_i];
|
||||
}
|
||||
|
||||
state_kv[dst] = kv[gid];
|
||||
state_score[dst] = score[gid] + ape_v;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,165 @@
|
||||
struct ds4_metal_args_dsv4_rope_tail {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
int32_t n_dims;
|
||||
int32_t mode;
|
||||
int32_t n_ctx_orig;
|
||||
int32_t inverse;
|
||||
float freq_base;
|
||||
float freq_scale;
|
||||
float ext_factor;
|
||||
float attn_factor;
|
||||
float beta_fast;
|
||||
float beta_slow;
|
||||
bool src2;
|
||||
};
|
||||
|
||||
static float rope_yarn_ramp(const float low, const float high, const int i0) {
|
||||
const float y = (i0 / 2 - low) / max(0.001f, high - low);
|
||||
return 1.0f - min(1.0f, max(0.0f, y));
|
||||
}
|
||||
|
||||
// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
|
||||
// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
|
||||
static void rope_yarn(
|
||||
float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale,
|
||||
thread float * cos_theta, thread float * sin_theta) {
|
||||
// Get n-d rotational scaling corrected for extrapolation
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
if (ext_factor != 0.0f) {
|
||||
float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor;
|
||||
theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
|
||||
|
||||
// Get n-d magnitude scaling corrected for interpolation
|
||||
mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
|
||||
}
|
||||
*cos_theta = cos(theta) * mscale;
|
||||
*sin_theta = sin(theta) * mscale;
|
||||
}
|
||||
|
||||
// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
|
||||
// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
|
||||
static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
|
||||
return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
|
||||
}
|
||||
|
||||
static void rope_yarn_corr_dims(
|
||||
int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]
|
||||
) {
|
||||
// start and end correction dims
|
||||
dims[0] = max(0.0f, floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base)));
|
||||
dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)));
|
||||
}
|
||||
|
||||
// Applies DeepSeek V4's partial RoPE: the no-position prefix is copied and only
|
||||
// the rotated tail is transformed. This is used for Q/K after their projections
|
||||
// and before writing/reading the attention KV state.
|
||||
kernel void kernel_dsv4_rope_tail_f32(
|
||||
constant ds4_metal_args_dsv4_rope_tail & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
const int i3 = tgpig[2];
|
||||
|
||||
const int n_nope = args.ne00 - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
const float theta_base = (float) pos[i2];
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
const bool is_neox = args.mode == 2;
|
||||
|
||||
for (int i0 = tid; i0 < args.ne00; i0 += ntg.x) {
|
||||
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
||||
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
|
||||
if (i0 < n_nope) {
|
||||
*((device float *) (dst_base + i0*args.nb0)) = *((device const float *) (src_base + i0*args.nb00));
|
||||
continue;
|
||||
}
|
||||
|
||||
const int r = i0 - n_nope;
|
||||
if (is_neox) {
|
||||
const int n_half = args.n_dims/2;
|
||||
if (r >= n_half) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int ic = r;
|
||||
const int rel_i0 = 2*ic;
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
// Equivalent to pow(freq_base, k) but expressed through IEEE-754
|
||||
// primitives that have tighter precision guarantees than Metal's pow().
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*rel_i0);
|
||||
#endif
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, rel_i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const int j0 = n_nope + ic;
|
||||
const int j1 = n_nope + ic + n_half;
|
||||
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
||||
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
||||
|
||||
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
} else {
|
||||
if ((r & 1) != 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int ic = r/2;
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
||||
#endif
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const int j0 = n_nope + r;
|
||||
const int j1 = j0 + 1;
|
||||
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
||||
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
||||
|
||||
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,54 @@
|
||||
// DS4 Metal get-rows kernel.
|
||||
|
||||
struct ds4_metal_args_get_rows {
|
||||
int32_t ne00t;
|
||||
int32_t ne00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne10;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
// Gathers embedding/table rows by integer ids. DS4 uses this for token
|
||||
// embeddings and small indexed tables such as router/hash lookup outputs.
|
||||
template<typename T0, typename T>
|
||||
kernel void kernel_get_rows_f(
|
||||
constant ds4_metal_args_get_rows & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t iw0 = tgpig.x/args.ne10;
|
||||
const int32_t i10 = tgpig.x%args.ne10;
|
||||
const int32_t i11 = tgpig.y;
|
||||
const int32_t i12 = tgpig.z;
|
||||
|
||||
const int32_t r = ((const device int32_t *) (src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
|
||||
|
||||
const int32_t i02 = i11;
|
||||
const int32_t i03 = i12;
|
||||
|
||||
auto psrc = (const device T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01);
|
||||
auto pdst = ( device T *) (dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1);
|
||||
|
||||
for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
|
||||
pdst[ind] = psrc[ind];
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
|
||||
|
||||
// Host-visible gather variants for F32, F16, and I32 tables.
|
||||
template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f<float, float>;
|
||||
template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f<half, float>;
|
||||
template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
|
||||
@@ -0,0 +1,40 @@
|
||||
struct ds4_metal_args_glu {
|
||||
int32_t ne00;
|
||||
uint64_t nb01;
|
||||
int32_t ne10;
|
||||
uint64_t nb11;
|
||||
int32_t ne0;
|
||||
uint64_t nb1;
|
||||
int32_t i00;
|
||||
int32_t i10;
|
||||
float alpha;
|
||||
float limit;
|
||||
};
|
||||
|
||||
// SwiGLU activation for the FFN inner state. DS4 clamps the shared expert with
|
||||
// the same swiglu_limit used by routed experts.
|
||||
kernel void kernel_swiglu_f32(
|
||||
constant ds4_metal_args_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const float * src0_row = (device const float *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const float * src1_row = (device const float *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device float * dst_row = (device float *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
float x0 = src0_row[i0];
|
||||
float x1 = src1_row[i0];
|
||||
if (args.limit > 1.0e-6f) {
|
||||
x0 = min(x0, args.limit);
|
||||
x1 = clamp(x1, -args.limit, args.limit);
|
||||
}
|
||||
|
||||
const float silu = x0 / (1.0f + exp(-x0));
|
||||
|
||||
dst_row[i0] = silu*x1*args.alpha;
|
||||
}
|
||||
}
|
||||
+4606
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,160 @@
|
||||
struct ds4_metal_args_norm {
|
||||
int32_t ne00;
|
||||
int32_t ne00_t;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
float eps;
|
||||
int32_t nef1[3];
|
||||
int32_t nef2[3];
|
||||
int32_t nef3[3];
|
||||
uint64_t nbf1[3];
|
||||
uint64_t nbf2[3];
|
||||
uint64_t nbf3[3];
|
||||
};
|
||||
|
||||
// RMSNorm over one activation row, optionally fusing the learned weight
|
||||
// multiply. DS4 calls this before attention, before the FFN, and for plain
|
||||
// diagnostics that need normalized but unweighted rows.
|
||||
template <typename T, short F>
|
||||
kernel void kernel_rms_norm_fuse_impl(
|
||||
constant ds4_metal_args_norm & args,
|
||||
device const char * src0,
|
||||
device const char * src1_0,
|
||||
device const char * src1_1,
|
||||
device char * dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const int i01 = tgpig.x;
|
||||
const int i02 = tgpig.y;
|
||||
const int i03 = tgpig.z;
|
||||
|
||||
device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
|
||||
device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
|
||||
device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
// parallel sum
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
sumf += dot(x[i00], x[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float mean = sumf/args.ne00;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
if (F == 1) {
|
||||
y[i00] = (x[i00]*scale);
|
||||
}
|
||||
if (F == 2) {
|
||||
y[i00] = (x[i00]*scale)*f0[i00];
|
||||
}
|
||||
if (F == 3) {
|
||||
y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;
|
||||
|
||||
// Host-visible RMSNorm variants: plain norm and norm multiplied by weight.
|
||||
template [[host_name("kernel_rms_norm_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
|
||||
template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
|
||||
|
||||
struct ds4_metal_args_qkv_rms_norm {
|
||||
int32_t q_n;
|
||||
int32_t q_n4;
|
||||
int32_t kv_n;
|
||||
int32_t kv_n4;
|
||||
uint64_t q_row_stride;
|
||||
uint64_t kv_row_stride;
|
||||
float eps;
|
||||
};
|
||||
|
||||
// Normalizes DS4's q-lora row and KV row in one dispatch. The two reductions
|
||||
// deliberately mirror kernel_rms_norm_mul_f32_4: Q uses the full 256-thread
|
||||
// row shape for 1024 floats, while KV only has work in the first 128 lanes for
|
||||
// its 512 floats. This keeps the q/kv normalization math aligned with the
|
||||
// standalone kernels while removing one tiny launch from the attention setup.
|
||||
kernel void kernel_dsv4_qkv_rms_norm_f32_4(
|
||||
constant ds4_metal_args_qkv_rms_norm & args,
|
||||
device const float4 * q_src,
|
||||
device const float4 * q_weight,
|
||||
device float4 * q_dst,
|
||||
device const float4 * kv_src,
|
||||
device const float4 * kv_weight,
|
||||
device float4 * kv_dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const uint row = tgpig.x;
|
||||
const bool kv_task = tgpig.y != 0;
|
||||
const int n = kv_task ? args.kv_n : args.q_n;
|
||||
const int n4 = kv_task ? args.kv_n4 : args.q_n4;
|
||||
const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4);
|
||||
|
||||
device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4;
|
||||
device const float4 * w = kv_task ? kv_weight : q_weight;
|
||||
device float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4;
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = tpitg.x; i < n4; i += ntg.x) {
|
||||
const float4 v = x[i];
|
||||
sumf += dot(v, v);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
#ifdef DS4_METAL_NORM_RSQRT_DISABLE
|
||||
// Match the formula used by kernel_rms_norm_fuse_impl above so both RMSNorm
|
||||
// entry points produce bit-identical scales. Hardware rsqrt() and 1.0f/sqrt()
|
||||
// can differ by ~1 ULP and that difference compounds across 43 layers.
|
||||
const float scale = 1.0f / sqrt(sumf / float(n) + args.eps);
|
||||
#else
|
||||
const float scale = rsqrt(sumf / float(n) + args.eps);
|
||||
#endif
|
||||
|
||||
for (int i = tpitg.x; i < n4; i += ntg.x) {
|
||||
y[i] = (x[i] * scale) * w[i];
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
// DS4 Metal repeat kernel used for HC embedding expansion.
|
||||
|
||||
struct ds4_metal_args_repeat {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
// Repeats a source row into the HC channel dimension. DS4 uses this when the
|
||||
// token embedding has to become an HC activation block before layer 0.
|
||||
template<typename T>
|
||||
kernel void kernel_repeat(
|
||||
constant ds4_metal_args_repeat & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
const int i03 = i3%args.ne03;
|
||||
const int i02 = i2%args.ne02;
|
||||
const int i01 = i1%args.ne01;
|
||||
|
||||
device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
|
||||
device char * dst_ptr = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i00 = i0%args.ne00;
|
||||
*((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_repeat<float>) kernel_repeat_t;
|
||||
|
||||
// Host-visible F32 repeat used for HC expansion of embeddings.
|
||||
template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
|
||||
@@ -0,0 +1,55 @@
|
||||
// DS4 Metal set-rows kernel used for KV writes.
|
||||
|
||||
struct ds4_metal_args_set_rows {
|
||||
int32_t nk0;
|
||||
int32_t ne01;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
// Scatters rows into the KV cache by token position. DS4 uses this after Q/K/V
|
||||
// preparation so decode and later prefill chunks can attend to previous tokens.
|
||||
template<typename T, typename TI>
|
||||
kernel void kernel_set_rows_f(
|
||||
constant ds4_metal_args_set_rows & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint tiitg[[thread_index_in_threadgroup]],
|
||||
uint3 tptg [[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
|
||||
const int32_t i12 = i03%args.ne12;
|
||||
const int32_t i11 = i02%args.ne11;
|
||||
|
||||
const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int32_t i10 = i01;
|
||||
const TI i1 = ((const device TI *) (src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
|
||||
|
||||
device T * dst_row = ( device T *) ((device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
const device float * src_row = (const device float *) ( src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
|
||||
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
|
||||
dst_row[ind] = (T) src_row[ind];
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_set_rows_f<float, int64_t>) set_rows_f_t;
|
||||
|
||||
// Host-visible F32/I32 scatter variant used by KV-cache writes.
|
||||
template [[host_name("kernel_set_rows_f32_i32")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t>;
|
||||
@@ -0,0 +1,241 @@
|
||||
// DS4 Metal softmax kernel used by the compressor pooling compatibility path.
|
||||
// The single-compressed-row path is intentionally left as soft_max -> mul ->
|
||||
// sum_rows instead of using the fused dsv4_softmax_pool kernel.
|
||||
|
||||
struct ds4_metal_args_soft_max {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
int32_t ne13;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb13;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
float scale;
|
||||
float max_bias;
|
||||
float m0;
|
||||
float m1;
|
||||
int32_t n_head_log2;
|
||||
};
|
||||
|
||||
// Row softmax for score matrices. DS4 uses it in the literal one-compressor-row
|
||||
// path where preserving the original graph operation boundary avoids drift.
|
||||
template<typename T>
|
||||
kernel void kernel_soft_max(
|
||||
constant ds4_metal_args_soft_max & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint3 tptg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x;
|
||||
|
||||
const int32_t i13 = i03%args.ne13;
|
||||
const int32_t i12 = i02%args.ne12;
|
||||
const int32_t i11 = i01;
|
||||
|
||||
device const float * psrc0 = (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
||||
device const float * psrc2 = src2 != src0 ? (device const float *) (src2) : nullptr;
|
||||
device float * pdst = (device float *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
if (args.max_bias > 0.0f) {
|
||||
const int32_t h = i02;
|
||||
|
||||
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
||||
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
||||
|
||||
slope = pow(base, exp);
|
||||
}
|
||||
|
||||
float lmax = psrc2 ? psrc2[i02] : -INFINITY;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
|
||||
}
|
||||
|
||||
float max_val = simd_max(lmax);
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = -INFINITY;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = max_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = buf[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
float lsum = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
|
||||
lsum += exp_psrc0;
|
||||
pdst[i00] = exp_psrc0;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
float sum = simd_sum(lsum);
|
||||
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = buf[tiisg];
|
||||
sum = simd_sum(sum);
|
||||
}
|
||||
|
||||
if (psrc2) {
|
||||
sum += exp(psrc2[i02] - max_val);
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f/sum;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
pdst[i00] *= inv_sum;
|
||||
}
|
||||
}
|
||||
|
||||
// Vectorized float4 row softmax for contiguous score rows whose length is a
|
||||
// multiple of four; used by the same DS4 compressor/indexer graph path.
|
||||
template<typename T>
|
||||
kernel void kernel_soft_max_4(
|
||||
constant ds4_metal_args_soft_max & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint3 tptg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x;
|
||||
|
||||
const int32_t i13 = i03%args.ne13;
|
||||
const int32_t i12 = i02%args.ne12;
|
||||
const int32_t i11 = i01;
|
||||
|
||||
device const float4 * psrc4 = (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
||||
device const float * psrc2 = src2 != src0 ? (device const float * ) (src2) : nullptr;
|
||||
device float4 * pdst4 = (device float4 *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
if (args.max_bias > 0.0f) {
|
||||
const int32_t h = i02;
|
||||
|
||||
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
||||
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
||||
|
||||
slope = pow(base, exp);
|
||||
}
|
||||
|
||||
float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
|
||||
}
|
||||
|
||||
const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
|
||||
|
||||
float max_val = simd_max(lmax);
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = -INFINITY;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = max_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = buf[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
float4 lsum4 = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
|
||||
lsum4 += exp_psrc4;
|
||||
pdst4[i00] = exp_psrc4;
|
||||
}
|
||||
|
||||
const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
float sum = simd_sum(lsum);
|
||||
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = buf[tiisg];
|
||||
sum = simd_sum(sum);
|
||||
}
|
||||
|
||||
if (psrc2) {
|
||||
sum += exp(psrc2[i02] - max_val);
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f/sum;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
pdst4[i00] *= inv_sum;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_soft_max<float>) kernel_soft_max_t;
|
||||
typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
|
||||
|
||||
// Host-visible F32 softmax variants used by compressor pooling.
|
||||
template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kernel_soft_max<float>;
|
||||
template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;
|
||||
@@ -0,0 +1,102 @@
|
||||
// DS4 Metal row-sum kernel.
|
||||
|
||||
#define FC_SUM_ROWS 1400
|
||||
|
||||
#define OP_SUM_ROWS_NUM_SUM_ROWS 10
|
||||
#define OP_SUM_ROWS_NUM_MEAN 11
|
||||
|
||||
struct ds4_metal_args_sum_rows {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
int64_t ne2;
|
||||
int64_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
static inline float sum(float x) {
|
||||
return x;
|
||||
}
|
||||
|
||||
static inline float sum(float4 x) {
|
||||
return x[0] + x[1] + x[2] + x[3];
|
||||
}
|
||||
|
||||
constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]];
|
||||
|
||||
// Reduces each row to a sum or mean. DS4 mainly uses the sum form to preserve
|
||||
// the compressor-pooling graph boundary in the single-compressor-row case.
|
||||
template <typename T0, typename T>
|
||||
kernel void kernel_sum_rows_impl(
|
||||
constant ds4_metal_args_sum_rows & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_sum_rows_op
|
||||
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
threadgroup T0 * shmem_t = (threadgroup T0 *) shmem;
|
||||
|
||||
if (sgitg == 0) {
|
||||
shmem_t[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
|
||||
device T * dst_row = (device T *) (dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3);
|
||||
|
||||
T0 sumf = T0(0.0f);
|
||||
|
||||
for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
|
||||
sumf += src_row[i0];
|
||||
}
|
||||
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_t[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_t[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
if (tpitg.x == 0) {
|
||||
if (FC_OP == OP_SUM_ROWS_NUM_MEAN) {
|
||||
if (is_same<float4, T0>::value) {
|
||||
dst_row[0] = sum(sumf) / (4*args.ne00);
|
||||
} else {
|
||||
dst_row[0] = sum(sumf) / args.ne00;
|
||||
}
|
||||
} else {
|
||||
dst_row[0] = sum(sumf);
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
}
|
||||
|
||||
typedef decltype(kernel_sum_rows_impl<float, float>) kernel_sum_rows_t;
|
||||
|
||||
// Host-visible F32 row reduction used by compressor pooling.
|
||||
template [[host_name("kernel_sum_rows_f32_f32")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float, float>;
|
||||
@@ -0,0 +1,312 @@
|
||||
#define FC_UNARY 1200
|
||||
|
||||
#define OP_UNARY_NUM_SCALE 10
|
||||
#define OP_UNARY_NUM_FILL 11
|
||||
#define OP_UNARY_NUM_CLAMP 12
|
||||
#define OP_UNARY_NUM_SQR 13
|
||||
#define OP_UNARY_NUM_SQRT 14
|
||||
#define OP_UNARY_NUM_SIN 15
|
||||
#define OP_UNARY_NUM_COS 16
|
||||
#define OP_UNARY_NUM_LOG 17
|
||||
#define OP_UNARY_NUM_LEAKY_RELU 18
|
||||
|
||||
#define OP_UNARY_NUM_TANH 100
|
||||
#define OP_UNARY_NUM_RELU 101
|
||||
#define OP_UNARY_NUM_SIGMOID 102
|
||||
#define OP_UNARY_NUM_GELU 103
|
||||
#define OP_UNARY_NUM_GELU_ERF 104
|
||||
#define OP_UNARY_NUM_GELU_QUICK 105
|
||||
#define OP_UNARY_NUM_SILU 106
|
||||
#define OP_UNARY_NUM_ELU 107
|
||||
#define OP_UNARY_NUM_NEG 108
|
||||
#define OP_UNARY_NUM_ABS 109
|
||||
#define OP_UNARY_NUM_SGN 110
|
||||
#define OP_UNARY_NUM_STEP 111
|
||||
#define OP_UNARY_NUM_HARDSWISH 112
|
||||
#define OP_UNARY_NUM_HARDSIGMOID 113
|
||||
#define OP_UNARY_NUM_EXP 114
|
||||
#define OP_UNARY_NUM_SOFTPLUS 115
|
||||
#define OP_UNARY_NUM_EXPM1 116
|
||||
#define OP_UNARY_NUM_FLOOR 117
|
||||
#define OP_UNARY_NUM_CEIL 118
|
||||
#define OP_UNARY_NUM_ROUND 119
|
||||
#define OP_UNARY_NUM_TRUNC 120
|
||||
#define OP_UNARY_NUM_XIELU 121
|
||||
|
||||
struct ds4_metal_args_unary {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
float slope;
|
||||
float scale;
|
||||
float bias;
|
||||
float val;
|
||||
float min;
|
||||
float max;
|
||||
};
|
||||
|
||||
constant float GELU_COEF_A = 0.044715f;
|
||||
constant float GELU_QUICK_COEF = -1.702f;
|
||||
constant float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
|
||||
constant float SQRT_2_INV = 0.70710678118654752440084436210484f;
|
||||
|
||||
// based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation
|
||||
// ref: https://www.johndcook.com/blog/python_erf/
|
||||
constant float p_erf = 0.3275911f;
|
||||
constant float a1_erf = 0.254829592f;
|
||||
constant float a2_erf = -0.284496736f;
|
||||
constant float a3_erf = 1.421413741f;
|
||||
constant float a4_erf = -1.453152027f;
|
||||
constant float a5_erf = 1.061405429f;
|
||||
|
||||
template<typename T>
|
||||
inline T erf_approx(T x) {
|
||||
T sign_x = sign(x);
|
||||
x = fabs(x);
|
||||
T t = 1.0f / (1.0f + p_erf * x);
|
||||
T y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x);
|
||||
return sign_x * y;
|
||||
}
|
||||
|
||||
template<typename T> T elu_approx(T x);
|
||||
|
||||
template<> inline float elu_approx<float>(float x) {
|
||||
return (x > 0.f) ? x : (exp(x) - 1);
|
||||
}
|
||||
|
||||
template<> inline float4 elu_approx<float4>(float4 x) {
|
||||
float4 res;
|
||||
|
||||
res[0] = (x[0] > 0.0f) ? x[0] : (exp(x[0]) - 1.0f);
|
||||
res[1] = (x[1] > 0.0f) ? x[1] : (exp(x[1]) - 1.0f);
|
||||
res[2] = (x[2] > 0.0f) ? x[2] : (exp(x[2]) - 1.0f);
|
||||
res[3] = (x[3] > 0.0f) ? x[3] : (exp(x[3]) - 1.0f);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
constant short FC_unary_op [[function_constant(FC_UNARY + 0)]];
|
||||
constant bool FC_unary_cnt[[function_constant(FC_UNARY + 1)]];
|
||||
|
||||
// Generic unary elementwise op selected by function constant. DS4 only uses a
|
||||
// small subset in inference, mainly sigmoid, SiLU, softplus, sqrt, clamp,
|
||||
// scale, and fill.
|
||||
template <typename T0, typename T, typename TC>
|
||||
kernel void kernel_unary_impl(
|
||||
constant ds4_metal_args_unary & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_unary_op
|
||||
#define FC_CNT FC_unary_cnt
|
||||
|
||||
device const T0 * src0_ptr;
|
||||
device T * dst_ptr;
|
||||
|
||||
int i0;
|
||||
|
||||
if (FC_CNT) {
|
||||
i0 = tgpig.x;
|
||||
|
||||
src0_ptr = (device const T0 *) (src0);
|
||||
dst_ptr = (device T *) (dst);
|
||||
} else {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int k0 = tgpig.x/args.ne01;
|
||||
const int i01 = tgpig.x - k0*args.ne01;
|
||||
|
||||
i0 = k0*ntg.x + tpitg.x;
|
||||
|
||||
src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 );
|
||||
}
|
||||
|
||||
{
|
||||
if (!FC_CNT) {
|
||||
if (i0 >= args.ne0) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
const TC x = (TC) src0_ptr[i0];
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SCALE) {
|
||||
dst_ptr[i0] = (T) (args.scale * x + args.bias);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_FILL) {
|
||||
dst_ptr[i0] = (T) args.val;
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_CLAMP) {
|
||||
dst_ptr[i0] = (T) clamp(x, args.min, args.max);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SQR) {
|
||||
dst_ptr[i0] = (T) (x * x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SQRT) {
|
||||
dst_ptr[i0] = (T) sqrt(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SIN) {
|
||||
dst_ptr[i0] = (T) sin(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_COS) {
|
||||
dst_ptr[i0] = (T) cos(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_LOG) {
|
||||
dst_ptr[i0] = (T) log(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) {
|
||||
dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_TANH) {
|
||||
dst_ptr[i0] = (T) precise::tanh(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_RELU) {
|
||||
dst_ptr[i0] = (T) fmax(0, x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SIGMOID) {
|
||||
dst_ptr[i0] = (T) (1 / (1 + exp(-x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU) {
|
||||
dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU_ERF) {
|
||||
dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU_QUICK) {
|
||||
dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x))));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SILU) {
|
||||
dst_ptr[i0] = (T) (x / (1 + exp(-x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ELU) {
|
||||
dst_ptr[i0] = (T) elu_approx(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_NEG) {
|
||||
dst_ptr[i0] = (T) -x;
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ABS) {
|
||||
dst_ptr[i0] = (T) fabs(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SGN) {
|
||||
dst_ptr[i0] = T(x > 0) - T(x < 0);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_STEP) {
|
||||
dst_ptr[i0] = T(x > 0);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_HARDSWISH) {
|
||||
dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) {
|
||||
dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_EXP) {
|
||||
dst_ptr[i0] = (T) exp(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SOFTPLUS) {
|
||||
dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_EXPM1) {
|
||||
// Metal target profiles used here do not all expose expm1(); this
|
||||
// generic unary branch is not used by the DS4 inference graph.
|
||||
dst_ptr[i0] = (T) (exp(x) - 1);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_FLOOR) {
|
||||
dst_ptr[i0] = (T) floor(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_CEIL) {
|
||||
dst_ptr[i0] = (T) ceil(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ROUND) {
|
||||
dst_ptr[i0] = (T) round(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_TRUNC) {
|
||||
dst_ptr[i0] = (T) trunc(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_XIELU) {
|
||||
const TC xi = x;
|
||||
const TC gate = TC(xi > TC(0.0f));
|
||||
const TC clamped = fmin(xi, TC(args.val));
|
||||
const TC y_pos = TC(args.scale) * xi * xi + TC(args.bias) * xi;
|
||||
const TC y_neg = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi;
|
||||
dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg);
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
#undef FC_CNT
|
||||
}
|
||||
|
||||
typedef decltype(kernel_unary_impl<float, float, float>) kernel_unary_t;
|
||||
|
||||
// Decode router probability transform. The generic path applies softplus and
|
||||
// sqrt as two elementwise kernels; DS4 decode always transforms one 256-wide
|
||||
// expert-logit row, so this vectorized kernel does both in one pass.
|
||||
kernel void kernel_dsv4_softplus_sqrt_f32_4(
|
||||
constant ds4_metal_args_unary & args,
|
||||
device const char *src,
|
||||
device char *dst,
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int k0 = tgpig.x/args.ne01;
|
||||
const int i01 = tgpig.x - k0*args.ne01;
|
||||
const int i0 = k0*ntg.x + tpitg.x;
|
||||
if (i0 >= args.ne0) return;
|
||||
|
||||
device const float4 *s = (device const float4 *)(src + i01*args.nb01);
|
||||
device float4 *d = (device float4 *)(dst + i01*args.nb1);
|
||||
const float4 x = s[i0];
|
||||
const float4 sp = select(log(1.0f + exp(x)), x, x > 20.0f);
|
||||
d[i0] = sqrt(sp);
|
||||
}
|
||||
|
||||
// Host-visible unary variants. Function constants select the actual DS4 op.
|
||||
template [[host_name("kernel_unary_f32_f32")]] kernel kernel_unary_t kernel_unary_impl<float, float, float>;
|
||||
template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
|
||||
template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl<half, half, float>;
|
||||
@@ -0,0 +1,301 @@
|
||||
/* Rax -- A radix tree implementation.
|
||||
*
|
||||
* Copyright (c) 2017-2018, Salvatore Sanfilippo <antirez at gmail dot com>
|
||||
* All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* * Redistributions of source code must retain the above copyright notice,
|
||||
* this list of conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in the
|
||||
* documentation and/or other materials provided with the distribution.
|
||||
* * Neither the name of Redis nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without
|
||||
* specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
|
||||
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
||||
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
||||
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
||||
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
||||
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
* POSSIBILITY OF SUCH DAMAGE.
|
||||
*/
|
||||
|
||||
#ifndef RAX_H
|
||||
#define RAX_H
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
/* Representation of a radix tree as implemented in this file, that contains
|
||||
* the strings "foo", "foobar" and "footer" after the insertion of each
|
||||
* word. When the node represents a key inside the radix tree, we write it
|
||||
* between [], otherwise it is written between ().
|
||||
*
|
||||
* This is the vanilla representation:
|
||||
*
|
||||
* (f) ""
|
||||
* \
|
||||
* (o) "f"
|
||||
* \
|
||||
* (o) "fo"
|
||||
* \
|
||||
* [t b] "foo"
|
||||
* / \
|
||||
* "foot" (e) (a) "foob"
|
||||
* / \
|
||||
* "foote" (r) (r) "fooba"
|
||||
* / \
|
||||
* "footer" [] [] "foobar"
|
||||
*
|
||||
* However, this implementation implements a very common optimization where
|
||||
* successive nodes having a single child are "compressed" into the node
|
||||
* itself as a string of characters, each representing a next-level child,
|
||||
* and only the link to the node representing the last character node is
|
||||
* provided inside the representation. So the above representation is turned
|
||||
* into:
|
||||
*
|
||||
* ["foo"] ""
|
||||
* |
|
||||
* [t b] "foo"
|
||||
* / \
|
||||
* "foot" ("er") ("ar") "foob"
|
||||
* / \
|
||||
* "footer" [] [] "foobar"
|
||||
*
|
||||
* However this optimization makes the implementation a bit more complex.
|
||||
* For instance if a key "first" is added in the above radix tree, a
|
||||
* "node splitting" operation is needed, since the "foo" prefix is no longer
|
||||
* composed of nodes having a single child one after the other. This is the
|
||||
* above tree and the resulting node splitting after this event happens:
|
||||
*
|
||||
*
|
||||
* (f) ""
|
||||
* /
|
||||
* (i o) "f"
|
||||
* / \
|
||||
* "firs" ("rst") (o) "fo"
|
||||
* / \
|
||||
* "first" [] [t b] "foo"
|
||||
* / \
|
||||
* "foot" ("er") ("ar") "foob"
|
||||
* / \
|
||||
* "footer" [] [] "foobar"
|
||||
*
|
||||
* Similarly after deletion, if a new chain of nodes having a single child
|
||||
* is created (the chain must also not include nodes that represent keys),
|
||||
* it must be compressed back into a single node.
|
||||
*
|
||||
*/
|
||||
|
||||
/* A single node can represent at most 2^16-1 characters / children because
|
||||
* part of the 32 bit header is now used for the inline leaf bitmap. Longer
|
||||
* compressed paths are still represented by chaining multiple nodes. */
|
||||
#define RAX_NODE_MAX_SIZE ((1<<16)-1)
|
||||
typedef struct raxNode {
|
||||
uint32_t iskey:1; /* Does this node contain a key? */
|
||||
uint32_t isnull:1; /* Associated value is NULL (don't store it). */
|
||||
uint32_t iscompr:1; /* Node is compressed. */
|
||||
uint32_t leafbitmap:13; /* Bit N set = child N is an inline value, not
|
||||
a pointer to a child node. This allows us to
|
||||
avoid allocating leaf nodes that only hold a
|
||||
value. Only bits 0-12 are available, so children
|
||||
at index >= 13 can't be inlined. For compressed
|
||||
nodes only bit 0 is meaningful. */
|
||||
uint32_t size:16; /* Number of children, or compressed string len. */
|
||||
/* Data layout is as follows:
|
||||
*
|
||||
* If node is not compressed we have 'size' bytes, one for each children
|
||||
* character, and 'size' raxNode pointers, point to each child node.
|
||||
* Note how the character is not stored in the children but in the
|
||||
* edge of the parents:
|
||||
*
|
||||
* [header iscompr=0][abc][a-ptr][b-ptr][c-ptr](value-ptr?)
|
||||
*
|
||||
* if node is compressed (iscompr bit is 1) the node has 1 child.
|
||||
* In that case the 'size' bytes of the string stored immediately at
|
||||
* the start of the data section, represent a sequence of successive
|
||||
* nodes linked one after the other, for which only the last one in
|
||||
* the sequence is actually represented as a node, and pointed to by
|
||||
* the current compressed node.
|
||||
*
|
||||
* [header iscompr=1][xyz][z-ptr](value-ptr?)
|
||||
*
|
||||
* Both compressed and not compressed nodes can represent a key
|
||||
* with associated data in the radix tree at any level (not just terminal
|
||||
* nodes).
|
||||
*
|
||||
* If the node has an associated key (iskey=1) and is not NULL
|
||||
* (isnull=0), then after the raxNode pointers pointing to the
|
||||
* children, an additional value pointer is present (as you can see
|
||||
* in the representation above as "value-ptr" field).
|
||||
*/
|
||||
unsigned char data[];
|
||||
} raxNode;
|
||||
|
||||
typedef struct rax {
|
||||
raxNode *head;
|
||||
uint64_t numele;
|
||||
uint64_t numnodes;
|
||||
} rax;
|
||||
|
||||
/* Stack data structure used by raxLowWalk() in order to, optionally, return
|
||||
* a list of parent nodes to the caller. The nodes do not have a "parent"
|
||||
* field for space concerns, so we use the auxiliary stack when needed. */
|
||||
#define RAX_STACK_STATIC_ITEMS 32
|
||||
typedef struct raxStack {
|
||||
void **stack; /* Points to static_items or an heap allocated array. */
|
||||
size_t items, maxitems; /* Number of items contained and total space. */
|
||||
/* Up to RAXSTACK_STACK_ITEMS items we avoid to allocate on the heap
|
||||
* and use this static array of pointers instead. */
|
||||
void *static_items[RAX_STACK_STATIC_ITEMS];
|
||||
int oom; /* True if pushing into this stack failed for OOM at some point. */
|
||||
} raxStack;
|
||||
|
||||
/* Optional callback used for iterators and be notified on each rax node,
|
||||
* including nodes not representing keys. If the callback returns true
|
||||
* the callback changed the node pointer in the iterator structure, and the
|
||||
* iterator implementation will have to replace the pointer in the radix tree
|
||||
* internals. This allows the callback to reallocate the node to perform
|
||||
* very special operations, normally not needed by normal applications.
|
||||
*
|
||||
* This callback is used to perform very low level analysis of the radix tree
|
||||
* structure, scanning each possible node (but the root node), or in order to
|
||||
* reallocate the nodes to reduce the allocation fragmentation (this is the
|
||||
* Redis application for this callback).
|
||||
*
|
||||
* This is currently only supported in forward iterations (raxNext) */
|
||||
typedef int (*raxNodeCallback)(raxNode **noderef);
|
||||
|
||||
/* Radix tree iterator state is encapsulated into this data structure. */
|
||||
#define RAX_ITER_STATIC_LEN 128
|
||||
#define RAX_ITER_JUST_SEEKED (1<<0) /* Iterator was just seeked. Return current
|
||||
element for the first iteration and
|
||||
clear the flag. */
|
||||
#define RAX_ITER_EOF (1<<1) /* End of iteration reached. */
|
||||
#define RAX_ITER_SAFE (1<<2) /* Safe iterator, allows operations while
|
||||
iterating. But it is slower. */
|
||||
#define RAX_ITER_INLINE_LEAF (1<<3) /* Iterator is positioned on an inline
|
||||
leaf stored in the current parent. */
|
||||
typedef struct raxIterator {
|
||||
int flags;
|
||||
rax *rt; /* Radix tree we are iterating. */
|
||||
unsigned char *key; /* The current string. */
|
||||
void *data; /* Data associated to this key. */
|
||||
size_t key_len; /* Current key length. */
|
||||
size_t key_max; /* Max key len the current key buffer can hold. */
|
||||
unsigned char key_static_string[RAX_ITER_STATIC_LEN];
|
||||
raxNode *node; /* Current node, or the parent node if the
|
||||
iterator is on an inline leaf. */
|
||||
int node_child; /* Cached child index of the current node in its
|
||||
parent, or of the inline leaf in 'node'.
|
||||
-1 if unknown. */
|
||||
raxStack stack; /* Stack used for unsafe iteration. */
|
||||
raxNodeCallback node_cb; /* Optional node callback. Normally set to NULL. */
|
||||
} raxIterator;
|
||||
|
||||
/* Defragmentation iterator. Unlike the normal iterator, this iterator scans
|
||||
* the radix tree structure itself, yielding both real raxNode allocations and
|
||||
* non-NULL values associated with keys. It is suitable in order to relocate
|
||||
* nodes and values while the iterator itself takes care of fixing the tree
|
||||
* links and its own state.
|
||||
*
|
||||
* The iterator returns items via raxDefragNext():
|
||||
*
|
||||
* RAX_DEFRAG_NODE:
|
||||
* 'ptr' points to the current raxNode allocation and 'size' is the exact
|
||||
* allocation size in bytes. The caller may allocate a new buffer, copy
|
||||
* 'size' bytes, and then call raxDefragReplaceNode().
|
||||
*
|
||||
* RAX_DEFRAG_DATA:
|
||||
* 'ptr' points to a non-NULL value associated with the current key. The
|
||||
* 'size' field is always zero because user data is opaque to Rax. The
|
||||
* caller may relocate the value using application-specific knowledge and
|
||||
* then call raxDefragReplaceData().
|
||||
*
|
||||
* The current key is available in 'key' / 'key_len' for DATA items and for
|
||||
* NODE items that are keys. Inline leaf values are returned as DATA items,
|
||||
* flagged with RAX_DEFRAG_F_INLINE_DATA, even if there is no standalone node
|
||||
* for the key.
|
||||
*
|
||||
* The iterator performs a full traversal and has no seek API. */
|
||||
#define RAX_DEFRAG_STATIC_ITEMS 32
|
||||
#define RAX_DEFRAG_NODE 1
|
||||
#define RAX_DEFRAG_DATA 2
|
||||
#define RAX_DEFRAG_F_ROOT (1<<0)
|
||||
#define RAX_DEFRAG_F_KEY (1<<1)
|
||||
#define RAX_DEFRAG_F_COMPRESSED (1<<2)
|
||||
#define RAX_DEFRAG_F_INLINE_DATA (1<<3)
|
||||
typedef struct raxDefragFrame {
|
||||
raxNode *node; /* Current node at this stack level. */
|
||||
size_t child; /* Next child to scan. */
|
||||
int parent_child; /* This node child index in the parent, or -1. */
|
||||
int state; /* Internal iterator state. */
|
||||
} raxDefragFrame;
|
||||
|
||||
typedef struct raxDefragIterator {
|
||||
/* The following fields describe the current item returned by the
|
||||
* iterator and can be accessed directly by the caller. */
|
||||
int kind; /* RAX_DEFRAG_NODE or RAX_DEFRAG_DATA. */
|
||||
int flags; /* Current item flags, see RAX_DEFRAG_F_* macros. */
|
||||
rax *rt; /* Radix tree being scanned. */
|
||||
unsigned char *key; /* Current key. */
|
||||
size_t key_len; /* Current key length. */
|
||||
size_t key_max; /* Max key len the current key buffer can hold. */
|
||||
size_t size; /* Exact node allocation size, or zero for DATA. */
|
||||
void *ptr; /* Current node or data pointer. */
|
||||
|
||||
/* The following fields are internal iterator state and should not be
|
||||
* modified by the caller. */
|
||||
unsigned char key_static_string[RAX_ITER_STATIC_LEN];
|
||||
raxNode *node; /* Current node, or the parent of inline DATA. */
|
||||
int node_child; /* Current node child index in its parent, or the
|
||||
inline leaf child index in 'node'. */
|
||||
raxDefragFrame *stack; /* DFS stack used by the iterator. */
|
||||
size_t items, maxitems; /* Stack length and capacity. */
|
||||
size_t pending_todel; /* Characters to remove after inline DATA events. */
|
||||
int eof; /* True if there are no more items to return. */
|
||||
raxDefragFrame static_items[RAX_DEFRAG_STATIC_ITEMS];
|
||||
} raxDefragIterator;
|
||||
|
||||
/* A special pointer returned for not found items. */
|
||||
extern void *raxNotFound;
|
||||
|
||||
/* Exported API. */
|
||||
rax *raxNew(void);
|
||||
int raxInsert(rax *rax, unsigned char *s, size_t len, void *data, void **old);
|
||||
int raxTryInsert(rax *rax, unsigned char *s, size_t len, void *data, void **old);
|
||||
int raxRemove(rax *rax, unsigned char *s, size_t len, void **old);
|
||||
void *raxFind(rax *rax, unsigned char *s, size_t len);
|
||||
void raxFree(rax *rax);
|
||||
void raxFreeWithCallback(rax *rax, void (*free_callback)(void*));
|
||||
void raxStart(raxIterator *it, rax *rt);
|
||||
int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len);
|
||||
int raxNext(raxIterator *it);
|
||||
int raxPrev(raxIterator *it);
|
||||
int raxRandomWalk(raxIterator *it, size_t steps);
|
||||
int raxCompare(raxIterator *iter, const char *op, unsigned char *key, size_t key_len);
|
||||
int raxIteratorSetData(raxIterator *it, void *data);
|
||||
void raxDefragStart(raxDefragIterator *it, rax *rt);
|
||||
int raxDefragNext(raxDefragIterator *it);
|
||||
void *raxDefragReplaceNode(raxDefragIterator *it, void *newptr);
|
||||
void *raxDefragReplaceData(raxDefragIterator *it, void *newptr);
|
||||
void raxDefragStop(raxDefragIterator *it);
|
||||
void raxStop(raxIterator *it);
|
||||
int raxEOF(raxIterator *it);
|
||||
void raxShow(rax *rax);
|
||||
uint64_t raxSize(rax *rax);
|
||||
unsigned long raxTouch(raxNode *n);
|
||||
void raxSetDebugMsg(int onoff);
|
||||
|
||||
/* Internal API. May be used by the node callback in order to access rax nodes
|
||||
* in a low level way, so this function is exported as well. */
|
||||
void raxSetData(raxNode *n, void *data);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,43 @@
|
||||
/* Rax -- A radix tree implementation.
|
||||
*
|
||||
* Copyright (c) 2017, Salvatore Sanfilippo <antirez at gmail dot com>
|
||||
* All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* * Redistributions of source code must retain the above copyright notice,
|
||||
* this list of conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in the
|
||||
* documentation and/or other materials provided with the distribution.
|
||||
* * Neither the name of Redis nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without
|
||||
* specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
|
||||
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
||||
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
||||
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
||||
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
||||
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
* POSSIBILITY OF SUCH DAMAGE.
|
||||
*/
|
||||
|
||||
/* Allocator selection.
|
||||
*
|
||||
* This file is used in order to change the Rax allocator at compile time.
|
||||
* Just define the following defines to what you want to use. Also add
|
||||
* the include of your alternate allocator if needed (not needed in order
|
||||
* to use the default libc allocator). */
|
||||
|
||||
#ifndef RAX_ALLOC_H
|
||||
#define RAX_ALLOC_H
|
||||
#define rax_malloc malloc
|
||||
#define rax_realloc realloc
|
||||
#define rax_free free
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,377 @@
|
||||
// DS4 ROCm common embedding/dense-matmul kernels and device helpers.
|
||||
//
|
||||
// Included from ds4_cuda.cu before more specialized modules; these helpers are
|
||||
// intentionally kept static in the single translation unit.
|
||||
|
||||
__global__ static void fill_f32_kernel(float *x, uint64_t n, float v) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) x[i] = v;
|
||||
}
|
||||
|
||||
__global__ static void embed_token_hc_kernel(float *out, const unsigned short *w, uint32_t token, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint32_t n = n_embd * n_hc;
|
||||
if (i >= n) return;
|
||||
uint32_t e = i % n_embd;
|
||||
out[i] = __half2float(reinterpret_cast<const __half *>(w)[(uint64_t)token * n_embd + e]);
|
||||
}
|
||||
|
||||
__device__ static float embed_q8_0_scale(const unsigned char *blk) {
|
||||
const uint16_t bits = (uint16_t)blk[0] | ((uint16_t)blk[1] << 8);
|
||||
return __half2float(__ushort_as_half((unsigned short)bits));
|
||||
}
|
||||
|
||||
__global__ static void embed_token_hc_q8_0_kernel(
|
||||
float *out,
|
||||
const unsigned char *w,
|
||||
uint32_t token,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t blocks = (n_embd + 31u) / 32u;
|
||||
const uint32_t b = d >> 5u;
|
||||
const uint32_t j = d & 31u;
|
||||
const unsigned char *blk = w + ((uint64_t)token * blocks + b) * 34u;
|
||||
out[gid] = embed_q8_0_scale(blk) * (float)((const int8_t *)(blk + 2u))[j];
|
||||
}
|
||||
|
||||
__global__ static void embed_tokens_hc_q8_0_kernel(
|
||||
float *out,
|
||||
const int32_t *tokens,
|
||||
const unsigned char *w,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
const uint32_t t = tmp / n_hc;
|
||||
int32_t tok_i = tokens[t];
|
||||
uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i;
|
||||
if (tok >= n_vocab) tok = 0;
|
||||
const uint32_t blocks = (n_embd + 31u) / 32u;
|
||||
const uint32_t b = d >> 5u;
|
||||
const uint32_t j = d & 31u;
|
||||
const unsigned char *blk = w + ((uint64_t)tok * blocks + b) * 34u;
|
||||
out[gid] = embed_q8_0_scale(blk) * (float)((const int8_t *)(blk + 2u))[j];
|
||||
}
|
||||
|
||||
__global__ static void embed_tokens_hc_kernel(
|
||||
float *out,
|
||||
const int32_t *tokens,
|
||||
const __half *w,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t t = tmp / n_hc;
|
||||
int32_t tok_i = tokens[t];
|
||||
uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i;
|
||||
if (tok >= n_vocab) tok = 0;
|
||||
out[gid] = __half2float(w[(uint64_t)tok * n_embd + d]);
|
||||
}
|
||||
|
||||
__device__ static float warp_sum_f32(float v);
|
||||
|
||||
__global__ static void matmul_f16_kernel(
|
||||
float *out,
|
||||
const __half *w,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
uint64_t n_tok) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
uint64_t tok = (uint64_t)blockIdx.y;
|
||||
if (row >= out_dim || tok >= n_tok) return;
|
||||
|
||||
float sum = 0.0f;
|
||||
const __half *wr = w + row * in_dim;
|
||||
const float *xr = x + tok * in_dim;
|
||||
for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
|
||||
sum += __half2float(wr[i]) * xr[i];
|
||||
}
|
||||
|
||||
__shared__ float partial[256];
|
||||
partial[threadIdx.x] = sum;
|
||||
__syncthreads();
|
||||
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
|
||||
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0];
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_ordered_chunks_kernel(
|
||||
float *out,
|
||||
const __half *w,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
uint64_t n_tok) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
uint64_t tok = (uint64_t)blockIdx.y;
|
||||
if (row >= out_dim || tok >= n_tok) return;
|
||||
|
||||
__shared__ float partial[32];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
float sum = 0.0f;
|
||||
const uint64_t chunk = (in_dim + 31u) / 32u;
|
||||
const uint64_t k0 = (uint64_t)tid * chunk;
|
||||
uint64_t k1 = k0 + chunk;
|
||||
if (k1 > in_dim) k1 = in_dim;
|
||||
const __half *wr = w + row * in_dim;
|
||||
const float *xr = x + tok * in_dim;
|
||||
for (uint64_t i = k0; i < k1; i++) {
|
||||
sum += __half2float(wr[i]) * xr[i];
|
||||
}
|
||||
partial[tid] = sum;
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
float total = 0.0f;
|
||||
for (uint32_t i = 0; i < 32u; i++) total += partial[i];
|
||||
out[tok * out_dim + row] = total;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_f32_sharedx_warp_rows_w32_kernel(
|
||||
float *out,
|
||||
const __half *w,
|
||||
const float *x,
|
||||
uint32_t in_dim,
|
||||
uint64_t out_dim) {
|
||||
extern __shared__ float shx[];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
const uint32_t lane = tid & 31u;
|
||||
const uint32_t wave = tid >> 5u;
|
||||
const uint32_t rows_per_block = blockDim.x >> 5u;
|
||||
for (uint32_t i = tid; i < in_dim; i += blockDim.x) shx[i] = x[i];
|
||||
__syncthreads();
|
||||
|
||||
const uint64_t row = (uint64_t)blockIdx.x * rows_per_block + wave;
|
||||
if (row >= out_dim) return;
|
||||
const __half *wr = w + row * (uint64_t)in_dim;
|
||||
float acc = 0.0f;
|
||||
uint32_t i = lane;
|
||||
for (; i + 224u < in_dim; i += 256u) {
|
||||
acc += __half2float(wr[i]) * shx[i];
|
||||
acc += __half2float(wr[i + 32u]) * shx[i + 32u];
|
||||
acc += __half2float(wr[i + 64u]) * shx[i + 64u];
|
||||
acc += __half2float(wr[i + 96u]) * shx[i + 96u];
|
||||
acc += __half2float(wr[i + 128u]) * shx[i + 128u];
|
||||
acc += __half2float(wr[i + 160u]) * shx[i + 160u];
|
||||
acc += __half2float(wr[i + 192u]) * shx[i + 192u];
|
||||
acc += __half2float(wr[i + 224u]) * shx[i + 224u];
|
||||
}
|
||||
for (; i < in_dim; i += 32u) {
|
||||
acc += __half2float(wr[i]) * shx[i];
|
||||
}
|
||||
acc = warp_sum_f32(acc);
|
||||
if (lane == 0u) out[row] = acc;
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel(
|
||||
float *out0,
|
||||
float *out1,
|
||||
const __half *w0,
|
||||
const __half *w1,
|
||||
const float *x,
|
||||
uint32_t in_dim,
|
||||
uint64_t out_dim) {
|
||||
extern __shared__ float shx[];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
const uint32_t lane = tid & 31u;
|
||||
const uint32_t wave = tid >> 5u;
|
||||
const uint32_t rows_per_block = blockDim.x >> 5u;
|
||||
for (uint32_t i = tid; i < in_dim; i += blockDim.x) shx[i] = x[i];
|
||||
__syncthreads();
|
||||
|
||||
const uint64_t row = (uint64_t)blockIdx.x * rows_per_block + wave;
|
||||
if (row >= out_dim) return;
|
||||
const __half *wr0 = w0 + row * (uint64_t)in_dim;
|
||||
const __half *wr1 = w1 + row * (uint64_t)in_dim;
|
||||
float acc0 = 0.0f;
|
||||
float acc1 = 0.0f;
|
||||
uint32_t i = lane;
|
||||
for (; i + 224u < in_dim; i += 256u) {
|
||||
float xv = shx[i];
|
||||
acc0 += __half2float(wr0[i]) * xv;
|
||||
acc1 += __half2float(wr1[i]) * xv;
|
||||
xv = shx[i + 32u];
|
||||
acc0 += __half2float(wr0[i + 32u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 32u]) * xv;
|
||||
xv = shx[i + 64u];
|
||||
acc0 += __half2float(wr0[i + 64u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 64u]) * xv;
|
||||
xv = shx[i + 96u];
|
||||
acc0 += __half2float(wr0[i + 96u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 96u]) * xv;
|
||||
xv = shx[i + 128u];
|
||||
acc0 += __half2float(wr0[i + 128u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 128u]) * xv;
|
||||
xv = shx[i + 160u];
|
||||
acc0 += __half2float(wr0[i + 160u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 160u]) * xv;
|
||||
xv = shx[i + 192u];
|
||||
acc0 += __half2float(wr0[i + 192u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 192u]) * xv;
|
||||
xv = shx[i + 224u];
|
||||
acc0 += __half2float(wr0[i + 224u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 224u]) * xv;
|
||||
}
|
||||
for (; i < in_dim; i += 32u) {
|
||||
const float xv = shx[i];
|
||||
acc0 += __half2float(wr0[i]) * xv;
|
||||
acc1 += __half2float(wr1[i]) * xv;
|
||||
}
|
||||
acc0 = warp_sum_f32(acc0);
|
||||
acc1 = warp_sum_f32(acc1);
|
||||
if (lane == 0u) {
|
||||
out0[row] = acc0;
|
||||
out1[row] = acc1;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_pair_ordered_chunks_kernel(
|
||||
float *out0,
|
||||
float *out1,
|
||||
const __half *w0,
|
||||
const __half *w1,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out0_dim,
|
||||
uint64_t out1_dim) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
if (row >= out0_dim && row >= out1_dim) return;
|
||||
|
||||
__shared__ float partial0[32];
|
||||
__shared__ float partial1[32];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
float sum0 = 0.0f;
|
||||
float sum1 = 0.0f;
|
||||
const uint64_t chunk = (in_dim + 31u) / 32u;
|
||||
const uint64_t k0 = (uint64_t)tid * chunk;
|
||||
uint64_t k1 = k0 + chunk;
|
||||
if (k1 > in_dim) k1 = in_dim;
|
||||
const __half *wr0 = row < out0_dim ? w0 + row * in_dim : w0;
|
||||
const __half *wr1 = row < out1_dim ? w1 + row * in_dim : w1;
|
||||
for (uint64_t i = k0; i < k1; i++) {
|
||||
const float xv = x[i];
|
||||
if (row < out0_dim) sum0 += __half2float(wr0[i]) * xv;
|
||||
if (row < out1_dim) sum1 += __half2float(wr1[i]) * xv;
|
||||
}
|
||||
partial0[tid] = sum0;
|
||||
partial1[tid] = sum1;
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
float total0 = 0.0f;
|
||||
float total1 = 0.0f;
|
||||
for (uint32_t i = 0; i < 32u; i++) {
|
||||
total0 += partial0[i];
|
||||
total1 += partial1[i];
|
||||
}
|
||||
if (row < out0_dim) out0[row] = total0;
|
||||
if (row < out1_dim) out1[row] = total1;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void matmul_f32_kernel(
|
||||
float *out,
|
||||
const float *w,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
uint64_t n_tok) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
uint64_t tok = (uint64_t)blockIdx.y;
|
||||
if (row >= out_dim || tok >= n_tok) return;
|
||||
|
||||
float sum = 0.0f;
|
||||
const float *wr = w + row * in_dim;
|
||||
const float *xr = x + tok * in_dim;
|
||||
for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
|
||||
sum += wr[i] * xr[i];
|
||||
}
|
||||
|
||||
__shared__ float partial[256];
|
||||
partial[threadIdx.x] = sum;
|
||||
__syncthreads();
|
||||
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
|
||||
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0];
|
||||
}
|
||||
|
||||
__global__ static void repeat_hc_kernel(float *out, const float *row, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
if (i >= n) return;
|
||||
out[i] = row[i % n_embd];
|
||||
}
|
||||
|
||||
__global__ static void f32_to_f16_kernel(__half *out, const float *x, uint64_t n) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) out[i] = __float2half(x[i]);
|
||||
}
|
||||
|
||||
__device__ static float warp_sum_f32(float v) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
v += __shfl_down(v, offset, 32);
|
||||
#else
|
||||
v += __shfl_down_sync(FULL_WARP_MASK, v, offset, 32);
|
||||
#endif
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__device__ static float warp_max_f32(float v) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
v = fmaxf(v, __shfl_down(v, offset, 32));
|
||||
#else
|
||||
v = fmaxf(v, __shfl_down_sync(FULL_WARP_MASK, v, offset, 32));
|
||||
#endif
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__device__ static uint16_t f32_to_f16_bits_hip_round(float f) {
|
||||
union { float f; uint32_t u; } v;
|
||||
v.f = f;
|
||||
uint32_t sign = (v.u >> 16) & 0x8000u;
|
||||
int32_t exp = (int32_t)((v.u >> 23) & 0xffu) - 127 + 15;
|
||||
uint32_t mant = v.u & 0x7fffffu;
|
||||
if (exp <= 0) {
|
||||
if (exp < -10) return (uint16_t)sign;
|
||||
mant |= 0x800000u;
|
||||
uint32_t shift = (uint32_t)(14 - exp);
|
||||
uint32_t half_mant = mant >> shift;
|
||||
if ((mant >> (shift - 1)) & 1u) half_mant++;
|
||||
return (uint16_t)(sign | half_mant);
|
||||
}
|
||||
if (exp >= 31) return (uint16_t)(sign | 0x7c00u);
|
||||
uint32_t half = sign | ((uint32_t)exp << 10) | (mant >> 13);
|
||||
if (mant & 0x1000u) half++;
|
||||
return (uint16_t)half;
|
||||
}
|
||||
|
||||
__device__ static float f16_bits_to_f32(uint16_t bits) {
|
||||
return __half2float(__ushort_as_half((unsigned short)bits));
|
||||
}
|
||||
|
||||
__device__ static float dot4_f32(float4 a, float4 b) {
|
||||
return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,557 @@
|
||||
__global__ static void compressor_store_kernel(
|
||||
const float *kv,
|
||||
const float *sc,
|
||||
float *state_kv,
|
||||
float *state_score,
|
||||
const void *model_map,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens) {
|
||||
uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
uint32_t width = coff * head_dim;
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * width;
|
||||
if (gid >= n) return;
|
||||
uint32_t t = gid / width;
|
||||
uint32_t j = gid - (uint64_t)t * width;
|
||||
uint32_t pos_mod = (pos0 + t) % ratio;
|
||||
uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
|
||||
state_kv[(uint64_t)dst_row * width + j] = kv[(uint64_t)t * width + j];
|
||||
state_score[(uint64_t)dst_row * width + j] =
|
||||
sc[(uint64_t)t * width + j] + model_ape_value_dev(model_map, ape_offset, ape_type, width, pos_mod, j);
|
||||
}
|
||||
|
||||
__global__ static void compressor_set_rows_kernel(
|
||||
float *state_kv,
|
||||
float *state_score,
|
||||
const float *kv,
|
||||
const float *sc,
|
||||
const void *model_map,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t width,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t src0,
|
||||
uint32_t dst0,
|
||||
uint32_t rows) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)rows * width;
|
||||
if (gid >= n) return;
|
||||
uint32_t r = gid / width;
|
||||
uint32_t j = gid - (uint64_t)r * width;
|
||||
uint32_t src = src0 + r;
|
||||
uint32_t dst = dst0 + r;
|
||||
uint32_t phase = (pos0 + src) % ratio;
|
||||
state_kv[(uint64_t)dst * width + j] = kv[(uint64_t)src * width + j];
|
||||
state_score[(uint64_t)dst * width + j] =
|
||||
sc[(uint64_t)src * width + j] + model_ape_value_dev(model_map, ape_offset, ape_type, width, phase, j);
|
||||
}
|
||||
|
||||
__global__ static void compressor_prefill_pool_kernel(
|
||||
float *comp,
|
||||
const float *kv,
|
||||
const float *sc,
|
||||
const float *state_kv,
|
||||
const float *state_score,
|
||||
const void *model_map,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_comp,
|
||||
uint32_t replay) {
|
||||
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint32_t c = blockIdx.y;
|
||||
if (d >= head_dim || c >= n_comp) return;
|
||||
uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
uint32_t width = coff * head_dim;
|
||||
float vals[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float scores[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float max_s = -INFINITY;
|
||||
uint32_t n_cand = 0;
|
||||
if (ratio == 4u) {
|
||||
if (replay && c == 0) {
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)r * width + d];
|
||||
scores[n_cand] = state_score[(uint64_t)r * width + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
} else if (c > 0) {
|
||||
uint32_t base = (c - 1u) * ratio;
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
uint32_t t = base + r;
|
||||
float ape = model_ape_value_dev(model_map, ape_offset, ape_type, width, (pos0 + t) % ratio, d);
|
||||
vals[n_cand] = kv[(uint64_t)t * width + d];
|
||||
scores[n_cand] = sc[(uint64_t)t * width + d] + ape;
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
}
|
||||
uint32_t base = c * ratio;
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
uint32_t t = base + r;
|
||||
float ape = model_ape_value_dev(model_map, ape_offset, ape_type, width, (pos0 + t) % ratio, head_dim + d);
|
||||
vals[n_cand] = kv[(uint64_t)t * width + head_dim + d];
|
||||
scores[n_cand] = sc[(uint64_t)t * width + head_dim + d] + ape;
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
} else {
|
||||
uint32_t base = c * ratio;
|
||||
for (uint32_t r = 0; r < ratio; r++) {
|
||||
uint32_t t = base + r;
|
||||
float ape = model_ape_value_dev(model_map, ape_offset, ape_type, width, (pos0 + t) % ratio, d);
|
||||
vals[n_cand] = kv[(uint64_t)t * width + d];
|
||||
scores[n_cand] = sc[(uint64_t)t * width + d] + ape;
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
}
|
||||
float den = 0.0f, acc = 0.0f;
|
||||
for (uint32_t i = 0; i < n_cand; i++) {
|
||||
float w = expf(scores[i] - max_s);
|
||||
den += w;
|
||||
acc += vals[i] * w;
|
||||
}
|
||||
comp[(uint64_t)c * head_dim + d] = den != 0.0f ? acc / den : 0.0f;
|
||||
}
|
||||
|
||||
__global__ static void compressor_update_pool_kernel(
|
||||
float *row,
|
||||
const float *state_kv,
|
||||
const float *state_score,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio) {
|
||||
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (d >= head_dim) return;
|
||||
uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
uint32_t width = coff * head_dim;
|
||||
float vals[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float scores[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float max_s = -INFINITY;
|
||||
uint32_t n_cand = 0;
|
||||
if (ratio == 4u) {
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)r * width + d];
|
||||
scores[n_cand] = state_score[(uint64_t)r * width + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)(ratio + r) * width + head_dim + d];
|
||||
scores[n_cand] = state_score[(uint64_t)(ratio + r) * width + head_dim + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
} else {
|
||||
for (uint32_t r = 0; r < ratio; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)r * width + d];
|
||||
scores[n_cand] = state_score[(uint64_t)r * width + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
}
|
||||
float den = 0.0f, acc = 0.0f;
|
||||
for (uint32_t i = 0; i < n_cand; i++) {
|
||||
float w = expf(scores[i] - max_s);
|
||||
den += w;
|
||||
acc += vals[i] * w;
|
||||
}
|
||||
row[d] = den != 0.0f ? acc / den : 0.0f;
|
||||
}
|
||||
|
||||
__global__ static void compressor_shift_ratio4_kernel(float *state_kv, float *state_score, uint32_t width) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t half = 4ull * width;
|
||||
if (i >= half) return;
|
||||
float v = state_kv[half + i];
|
||||
float s = state_score[half + i];
|
||||
state_kv[i] = v;
|
||||
state_score[i] = s;
|
||||
state_kv[half + i] = v;
|
||||
state_score[half + i] = s;
|
||||
}
|
||||
|
||||
static uint64_t cuda_tensor_2d_bytes(uint32_t type, uint64_t width, uint64_t rows) {
|
||||
if (type == 0u) return width * rows * sizeof(float);
|
||||
if (type == 1u) return width * rows * sizeof(uint16_t);
|
||||
if (type == 8u) return rows * (((width + 31u) / 32u) * 34u);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int cuda_ape_type_supported(uint32_t type) {
|
||||
return type == 0u || type == 1u || type == 8u;
|
||||
}
|
||||
|
||||
static int cuda_compressor_shape_supported(uint32_t head_dim, uint32_t ratio) {
|
||||
if (head_dim == 0u || ratio == 0u || ratio > DS4_ROCM_COMPRESSOR_MAX_RATIO) return 0;
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
return head_dim <= UINT32_MAX / coff;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_compressor_store_batch_tensor(
|
||||
const ds4_gpu_tensor *kv,
|
||||
const ds4_gpu_tensor *sc,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens) {
|
||||
if (!kv || !sc || !state_kv || !state_score || !model_map ||
|
||||
!cuda_compressor_shape_supported(head_dim, ratio) || n_tokens == 0 ||
|
||||
!cuda_ape_type_supported(ape_type)) {
|
||||
return 0;
|
||||
}
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
const uint32_t width = coff * head_dim;
|
||||
const uint32_t state_rows = coff * ratio;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul3_checked(n_tokens, width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv, kv_bytes) || !cuda_tensor_has_bytes(sc, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
uint64_t n = (uint64_t)n_tokens * width;
|
||||
compressor_store_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(const float *)kv->ptr,
|
||||
(const float *)sc->ptr,
|
||||
(float *)state_kv->ptr,
|
||||
(float *)state_score->ptr,
|
||||
ape,
|
||||
0,
|
||||
ape_type,
|
||||
head_dim,
|
||||
ratio,
|
||||
pos0,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "compressor store launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_compressor_update_tensor(
|
||||
const ds4_gpu_tensor *kv_cur,
|
||||
const ds4_gpu_tensor *sc_cur,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
ds4_gpu_tensor *comp_cache,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint64_t norm_offset,
|
||||
uint32_t norm_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos,
|
||||
uint32_t comp_row,
|
||||
uint32_t n_rot,
|
||||
uint32_t n_ctx_orig,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float rms_eps) {
|
||||
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
|
||||
!model_map || !cuda_compressor_shape_supported(head_dim, ratio) ||
|
||||
n_rot > head_dim || (n_rot & 1u) != 0 ||
|
||||
!cuda_ape_type_supported(ape_type) || norm_type != 0u) {
|
||||
return 0;
|
||||
}
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
const uint32_t width = coff * head_dim;
|
||||
const uint32_t state_rows = coff * ratio;
|
||||
const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0, comp_bytes = 0, norm_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul_checked(width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_u64_mul3_checked((uint64_t)comp_row + (emit ? 1u : 0u), head_dim, sizeof(float), &comp_bytes) ||
|
||||
!cuda_u64_mul_checked(head_dim, sizeof(float), &norm_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_model_range_fits(model_size, norm_offset, norm_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv_cur, kv_bytes) || !cuda_tensor_has_bytes(sc_cur, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes) ||
|
||||
(emit && !cuda_tensor_has_bytes(comp_cache, comp_bytes))) {
|
||||
return 0;
|
||||
}
|
||||
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
|
||||
model_map, model_size, ape_offset, ape_type,
|
||||
head_dim, ratio, pos, 1)) {
|
||||
return 0;
|
||||
}
|
||||
if (!emit) return 1;
|
||||
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
|
||||
comp_cache,
|
||||
(uint64_t)comp_row * head_dim * sizeof(float),
|
||||
(uint64_t)head_dim * sizeof(float));
|
||||
if (!comp_row_view) return 0;
|
||||
compressor_update_pool_kernel<<<(head_dim + 255) / 256, 256>>>(
|
||||
(float *)comp_row_view->ptr,
|
||||
(const float *)state_kv->ptr,
|
||||
(const float *)state_score->ptr,
|
||||
head_dim,
|
||||
ratio);
|
||||
int ok = cuda_ok(cudaGetLastError(), "compressor update pool launch");
|
||||
if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, comp_row_view,
|
||||
model_map, model_size, norm_offset,
|
||||
head_dim, 1, rms_eps);
|
||||
if (ok) ok = ds4_gpu_rope_tail_tensor(comp_row_view, 1, 1, head_dim, n_rot,
|
||||
pos + 1u - ratio, n_ctx_orig, false,
|
||||
freq_base, freq_scale, ext_factor, attn_factor,
|
||||
beta_fast, beta_slow);
|
||||
ds4_gpu_tensor_free(comp_row_view);
|
||||
if (ok && ratio == 4u) {
|
||||
uint64_t half = 4ull * width;
|
||||
compressor_shift_ratio4_kernel<<<(half + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr, width);
|
||||
ok = cuda_ok(cudaGetLastError(), "compressor ratio4 shift launch");
|
||||
}
|
||||
return ok;
|
||||
}
|
||||
extern "C" int ds4_gpu_compressor_prefill_tensor(
|
||||
ds4_gpu_tensor *comp_cache,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const ds4_gpu_tensor *kv,
|
||||
const ds4_gpu_tensor *sc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint64_t norm_offset,
|
||||
uint32_t norm_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_rot,
|
||||
uint32_t n_ctx_orig,
|
||||
bool quantize_fp8,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float rms_eps) {
|
||||
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
|
||||
!cuda_compressor_shape_supported(head_dim, ratio) || n_tokens == 0 ||
|
||||
n_rot > head_dim || (n_rot & 1u) != 0 ||
|
||||
!cuda_ape_type_supported(ape_type) || norm_type != 0u) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
const uint32_t width = coff * head_dim;
|
||||
const uint32_t state_rows = coff * ratio;
|
||||
const uint32_t n_comp = n_tokens / ratio;
|
||||
const uint32_t cutoff = n_comp * ratio;
|
||||
const uint32_t rem = n_tokens - cutoff;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0, comp_bytes = 0, norm_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
|
||||
if (!cuda_u64_mul3_checked(n_tokens, width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_comp, head_dim, sizeof(float), &comp_bytes) ||
|
||||
!cuda_u64_mul_checked(head_dim, sizeof(float), &norm_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_model_range_fits(model_size, norm_offset, norm_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv, kv_bytes) || !cuda_tensor_has_bytes(sc, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes) ||
|
||||
(n_comp && !cuda_tensor_has_bytes(comp_cache, comp_bytes))) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
|
||||
uint64_t state_n = (uint64_t)state_rows * width;
|
||||
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
|
||||
"compressor state kv zero")) return 0;
|
||||
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0;
|
||||
|
||||
if (ratio == 4u) {
|
||||
if (cutoff >= ratio) {
|
||||
uint32_t prev_start = cutoff - ratio;
|
||||
uint64_t n = (uint64_t)ratio * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
prev_start, 0, ratio);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill prev state launch")) return 0;
|
||||
}
|
||||
if (rem != 0) {
|
||||
uint64_t n = (uint64_t)rem * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
cutoff, ratio, rem);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0;
|
||||
}
|
||||
} else if (rem != 0) {
|
||||
uint64_t n = (uint64_t)rem * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
cutoff, 0, rem);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0;
|
||||
}
|
||||
if (n_comp != 0) {
|
||||
dim3 grid((head_dim + 255) / 256, n_comp, 1);
|
||||
compressor_prefill_pool_kernel<<<grid, 256>>>(
|
||||
(float *)comp_cache->ptr,
|
||||
(const float *)kv->ptr,
|
||||
(const float *)sc->ptr,
|
||||
(const float *)state_kv->ptr,
|
||||
(const float *)state_score->ptr,
|
||||
ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 0);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill pool launch")) return 0;
|
||||
if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache,
|
||||
model_map, model_size, norm_offset,
|
||||
head_dim, n_comp, rms_eps)) return 0;
|
||||
if (n_rot != 0 && !cuda_rope_tail_stride_tensor(comp_cache, n_comp, 1, head_dim,
|
||||
n_rot, pos0, ratio, n_ctx_orig, false,
|
||||
freq_base, freq_scale, ext_factor,
|
||||
attn_factor, beta_fast, beta_slow)) return 0;
|
||||
if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
extern "C" int ds4_gpu_compressor_prefill_ratio4_replay_tensor(
|
||||
ds4_gpu_tensor *comp_cache,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const ds4_gpu_tensor *kv,
|
||||
const ds4_gpu_tensor *sc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint64_t norm_offset,
|
||||
uint32_t norm_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_rot,
|
||||
uint32_t n_ctx_orig,
|
||||
bool quantize_fp8,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float rms_eps) {
|
||||
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
|
||||
head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 ||
|
||||
n_rot > head_dim || (n_rot & 1u) != 0 ||
|
||||
!cuda_ape_type_supported(ape_type) || norm_type != 0u) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const uint32_t ratio = 4u;
|
||||
const uint32_t width = 2u * head_dim;
|
||||
const uint32_t state_rows = 8u;
|
||||
const uint32_t n_comp = n_tokens / ratio;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0, comp_bytes = 0, norm_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul3_checked(n_tokens, width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_comp, head_dim, sizeof(float), &comp_bytes) ||
|
||||
!cuda_u64_mul_checked(head_dim, sizeof(float), &norm_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_model_range_fits(model_size, norm_offset, norm_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv, kv_bytes) || !cuda_tensor_has_bytes(sc, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes) ||
|
||||
!cuda_tensor_has_bytes(comp_cache, comp_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
dim3 grid((head_dim + 255) / 256, n_comp, 1);
|
||||
compressor_prefill_pool_kernel<<<grid, 256>>>(
|
||||
(float *)comp_cache->ptr,
|
||||
(const float *)kv->ptr,
|
||||
(const float *)sc->ptr,
|
||||
(const float *)state_kv->ptr,
|
||||
(const float *)state_score->ptr,
|
||||
ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 1);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor replay pool launch")) return 0;
|
||||
if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache,
|
||||
model_map, model_size, norm_offset,
|
||||
head_dim, n_comp, rms_eps)) return 0;
|
||||
if (n_rot != 0 && !cuda_rope_tail_stride_tensor(comp_cache, n_comp, 1, head_dim,
|
||||
n_rot, pos0, ratio, n_ctx_orig, false,
|
||||
freq_base, freq_scale, ext_factor,
|
||||
attn_factor, beta_fast, beta_slow)) return 0;
|
||||
if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0;
|
||||
|
||||
uint64_t state_n = (uint64_t)state_rows * width;
|
||||
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
|
||||
"compressor replay state kv zero")) return 0;
|
||||
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor replay state score fill launch")) return 0;
|
||||
uint32_t prev_start = n_tokens - ratio;
|
||||
uint64_t n = (uint64_t)ratio * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
prev_start, 0, ratio);
|
||||
return cuda_ok(cudaGetLastError(), "compressor replay state launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_compressor_prefill_state_ratio4_tensor(
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const ds4_gpu_tensor *kv_tail,
|
||||
const ds4_gpu_tensor *sc_tail,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t pos0) {
|
||||
if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map ||
|
||||
head_dim == 0 || !cuda_ape_type_supported(ape_type)) {
|
||||
return 0;
|
||||
}
|
||||
const uint32_t ratio = 4u;
|
||||
const uint32_t width = 2u * head_dim;
|
||||
const uint32_t state_rows = 8u;
|
||||
uint64_t tail_bytes = 0, state_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul3_checked(ratio, width, sizeof(float), &tail_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv_tail, tail_bytes) || !cuda_tensor_has_bytes(sc_tail, tail_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
uint64_t state_n = (uint64_t)state_rows * width;
|
||||
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
|
||||
"compressor state kv zero")) return 0;
|
||||
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0;
|
||||
uint64_t n = (uint64_t)ratio * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv_tail->ptr, (const float *)sc_tail->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
0, 0, ratio);
|
||||
return cuda_ok(cudaGetLastError(), "compressor state set launch");
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
extern "C" int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) {
|
||||
if (!event_value) return 0;
|
||||
*event_value = 0;
|
||||
if (!g_selected_readback_event) {
|
||||
cudaError_t err =
|
||||
cudaEventCreateWithFlags(&g_selected_readback_event,
|
||||
cudaEventDisableTiming);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback event creation failed: %s\n",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
cudaError_t err = cudaEventRecord(g_selected_readback_event, 0);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback event record failed: %s\n",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
*event_value = ++g_selected_readback_event_value;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) {
|
||||
if (event_value == 0 || !g_selected_readback_event) return 0;
|
||||
cudaError_t err = cudaEventSynchronize(g_selected_readback_event);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback wait failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) {
|
||||
return ds4_gpu_commit_and_wait_selected_readback(event_value, label);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_tensor_read_after_selected_event(
|
||||
const ds4_gpu_tensor *tensor,
|
||||
uint64_t offset,
|
||||
void *data,
|
||||
uint64_t bytes,
|
||||
uint64_t event_value,
|
||||
const char *label) {
|
||||
if (!tensor || !data || offset > tensor->bytes ||
|
||||
bytes > tensor->bytes - offset ||
|
||||
event_value == 0 ||
|
||||
!g_selected_readback_event) {
|
||||
return 0;
|
||||
}
|
||||
if (!g_selected_readback_stream) {
|
||||
cudaError_t err =
|
||||
cudaStreamCreateWithFlags(&g_selected_readback_stream,
|
||||
cudaStreamNonBlocking);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback stream creation failed: %s\n",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
cudaError_t err = hipStreamWaitEvent(g_selected_readback_stream,
|
||||
g_selected_readback_event,
|
||||
0);
|
||||
#else
|
||||
cudaError_t err = cudaStreamWaitEvent(g_selected_readback_stream,
|
||||
g_selected_readback_event,
|
||||
0);
|
||||
#endif
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback stream wait failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
err = cudaMemcpyAsync(data,
|
||||
(const char *)tensor->ptr + offset,
|
||||
(size_t)bytes,
|
||||
cudaMemcpyDeviceToHost,
|
||||
g_selected_readback_stream);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback copy failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
err = cudaStreamSynchronize(g_selected_readback_stream);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback sync failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) {
|
||||
int ok = ds4_gpu_set_model_fd(fd);
|
||||
g_model_fd_host_base = model_map ? model_map : g_model_host_base;
|
||||
return ok;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_tensor_copy_f32_to_f16(
|
||||
ds4_gpu_tensor *dst,
|
||||
uint64_t dst_offset,
|
||||
const ds4_gpu_tensor *src,
|
||||
uint64_t src_offset,
|
||||
uint64_t count) {
|
||||
if (!dst || !src || !dst->ptr || !src->ptr) return 0;
|
||||
if ((dst_offset % sizeof(__half)) != 0 || (src_offset % sizeof(float)) != 0) return 0;
|
||||
if (dst_offset > dst->bytes || src_offset > src->bytes) return 0;
|
||||
if (count > (UINT64_MAX / sizeof(__half)) || count > (UINT64_MAX / sizeof(float))) return 0;
|
||||
uint64_t dst_bytes = count * sizeof(__half);
|
||||
uint64_t src_bytes = count * sizeof(float);
|
||||
if (dst_bytes > dst->bytes - dst_offset || src_bytes > src->bytes - src_offset) return 0;
|
||||
if (count == 0) return 1;
|
||||
f32_to_f16_kernel<<<(count + 255u) / 256u, 256>>>(
|
||||
(__half *)((char *)dst->ptr + dst_offset),
|
||||
(const float *)((const char *)src->ptr + src_offset),
|
||||
count);
|
||||
return cuda_ok(cudaGetLastError(), "tensor copy f32 to f16 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_pro_q4_expert_table_auto_available(void) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_preload_q4_expert_tables(
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t gate_offset,
|
||||
uint64_t up_offset,
|
||||
uint64_t down_offset,
|
||||
uint64_t gate_expert_bytes,
|
||||
uint64_t down_expert_bytes,
|
||||
uint32_t n_total_expert) {
|
||||
(void)model_map;
|
||||
(void)model_size;
|
||||
(void)gate_offset;
|
||||
(void)up_offset;
|
||||
(void)down_offset;
|
||||
(void)gate_expert_bytes;
|
||||
(void)down_expert_bytes;
|
||||
(void)n_total_expert;
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) {
|
||||
g_ssd_streaming_mode = enabled ? 1 : 0;
|
||||
cuda_model_range_release_all();
|
||||
cuda_q8_f16_cache_release_all();
|
||||
g_routed_moe_selected_override_n = 0;
|
||||
g_stream_selected_cache.loaded = 0;
|
||||
g_stream_batch_selected_cache.loaded = 0;
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) {
|
||||
g_stream_expert_cache_budget = experts;
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) {
|
||||
(void)bytes;
|
||||
}
|
||||
|
||||
extern "C" uint64_t ds4_gpu_recommended_working_set_size(void) {
|
||||
size_t free_b = 0;
|
||||
size_t total_b = 0;
|
||||
if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) {
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
(void)free_b;
|
||||
return (uint64_t)total_b;
|
||||
}
|
||||
|
||||
extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) {
|
||||
return g_ssd_streaming_mode ? g_stream_expert_cache_budget : 0;
|
||||
}
|
||||
|
||||
extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) {
|
||||
return (uint32_t)g_stream_resident_experts.size();
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) {
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) {
|
||||
cuda_stream_resident_cache_release();
|
||||
}
|
||||
|
||||
extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(
|
||||
uint64_t gate_expert_bytes,
|
||||
uint64_t down_expert_bytes) {
|
||||
(void)gate_expert_bytes;
|
||||
(void)down_expert_bytes;
|
||||
return ds4_gpu_stream_expert_cache_configured_count();
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_seed_selected(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *selected_ids,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
if (!cuda_stream_selected_load(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
selected_ids,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
return cuda_stream_selected_finish_pending_missing(0);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *selected_ids,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
return cuda_stream_selected_load(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
selected_ids,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_prepare_selected_batch(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *selected_ids,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
const ds4_gpu_tensor *selected_exec = NULL;
|
||||
const char **gate_ptrs = NULL;
|
||||
const char **up_ptrs = NULL;
|
||||
const char **down_ptrs = NULL;
|
||||
uint32_t unique = 0;
|
||||
return cuda_stream_batch_selected_prepare_from_host(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
selected_ids,
|
||||
n_tokens,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes,
|
||||
&selected_exec,
|
||||
&gate_ptrs,
|
||||
&up_ptrs,
|
||||
&down_ptrs,
|
||||
&unique,
|
||||
1);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_load_layer(
|
||||
const ds4_gpu_stream_expert_table *table) {
|
||||
if (!table) return 0;
|
||||
return cuda_stream_layer_expert_cache_load(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
table->n_total_expert,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_seed_from_layer_selected(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const ds4_gpu_tensor *selected,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_seed_tokens,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
return cuda_stream_layer_expert_cache_seed_selected(table->model_map,
|
||||
table->layer,
|
||||
selected,
|
||||
n_tokens,
|
||||
n_seed_tokens,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_release_layer_cache(void) {
|
||||
cuda_stream_layer_expert_cache_release();
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_seed_experts(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *expert_ids,
|
||||
const uint32_t *expert_priorities,
|
||||
uint32_t n_experts) {
|
||||
(void)table;
|
||||
(void)expert_ids;
|
||||
(void)expert_priorities;
|
||||
(void)n_experts;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_routed_moe_set_selected_override(
|
||||
const int32_t *selected,
|
||||
uint32_t n_selected) {
|
||||
if (n_selected > DS4_ROCM_N_EXPERT_USED || (!selected && n_selected != 0)) return 0;
|
||||
for (uint32_t i = 0; i < n_selected; i++) {
|
||||
g_routed_moe_selected_override[i] = selected[i];
|
||||
}
|
||||
g_routed_moe_selected_override_n = n_selected;
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
extern "C" int ds4_gpu_embed_token_hc_tensor(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n_vocab, uint32_t token, uint32_t n_embd, uint32_t n_hc) {
|
||||
if (!out_hc || !model_map || n_vocab == 0u || token >= n_vocab || n_embd == 0u || n_hc == 0u ||
|
||||
(uint64_t)n_embd * n_hc > UINT32_MAX) return 0;
|
||||
uint64_t weight_bytes = 0;
|
||||
uint64_t out_bytes = 0;
|
||||
if (!cuda_u64_mul3_checked(n_vocab, n_embd, sizeof(uint16_t), &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_hc, n_embd, sizeof(float), &out_bytes) ||
|
||||
!cuda_tensor_has_bytes(out_hc, out_bytes)) return 0;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd");
|
||||
if (!wptr) return 0;
|
||||
uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
embed_token_hc_kernel<<<(n + 255u) / 256u, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed token launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_embed_tokens_hc_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
const ds4_gpu_tensor *tokens_t,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
if (!out_hc || !tokens_t || !model_map || n_vocab == 0u || n_tokens == 0u || n_embd == 0u || n_hc == 0u) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t weight_bytes = 0, n = 0;
|
||||
if (!cuda_u64_mul3_checked(n_vocab, n_embd, sizeof(uint16_t), &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_tensor_has_i32(tokens_t, n_tokens) ||
|
||||
!cuda_u64_mul_checked(n_tokens, n_hc, &n) ||
|
||||
!cuda_u64_mul_checked(n, n_embd, &n) ||
|
||||
!cuda_tensor_has_f32(out_hc, n)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd");
|
||||
if (!wptr) return 0;
|
||||
embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(const int32_t *)tokens_t->ptr,
|
||||
(const __half *)wptr,
|
||||
n_vocab, n_tokens, n_embd, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed tokens launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_embed_token_hc_q8_0_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint32_t n_vocab,
|
||||
uint32_t token,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
if (!out_hc || !model_map || n_vocab == 0u || token >= n_vocab || n_embd == 0 || n_hc == 0) return 0;
|
||||
const uint64_t blocks = (n_embd + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, out_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(n_vocab, row_bytes, &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_embd, n_hc, sizeof(float), &out_bytes) ||
|
||||
!cuda_tensor_has_bytes(out_hc, out_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd_q8_0");
|
||||
if (!wptr) return 0;
|
||||
const uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
embed_token_hc_q8_0_kernel<<<(n + 255u) / 256u, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(const unsigned char *)wptr,
|
||||
token,
|
||||
n_embd,
|
||||
n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed token q8_0 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_embed_tokens_hc_q8_0_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
const ds4_gpu_tensor *tokens_t,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
if (!out_hc || !tokens_t || !model_map || n_vocab == 0u || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0;
|
||||
const uint64_t blocks = (n_embd + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, n = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(n_vocab, row_bytes, &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_tensor_has_i32(tokens_t, n_tokens) ||
|
||||
!cuda_u64_mul_checked(n_tokens, n_hc, &n) ||
|
||||
!cuda_u64_mul_checked(n, n_embd, &n) ||
|
||||
!cuda_tensor_has_f32(out_hc, n)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd_q8_0");
|
||||
if (!wptr) return 0;
|
||||
embed_tokens_hc_q8_0_kernel<<<(n + 255u) / 256u, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(const int32_t *)tokens_t->ptr,
|
||||
(const unsigned char *)wptr,
|
||||
n_vocab,
|
||||
n_tokens,
|
||||
n_embd,
|
||||
n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed tokens q8_0 launch");
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
// DS4 ROCm FP8 KV quantization and raw-cache store kernels.
|
||||
//
|
||||
// This file is included from ds4_cuda.cu (same translation unit) so these
|
||||
// kernels can reuse the backend's existing device helpers without HIP device
|
||||
// linking or behavior changes.
|
||||
|
||||
__global__ static void fp8_kv_quantize_kernel(float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) {
|
||||
const uint32_t row = blockIdx.x;
|
||||
const uint32_t grp = blockIdx.y;
|
||||
const uint32_t tid = threadIdx.x;
|
||||
const uint32_t n_nope = head_dim - n_rot;
|
||||
const uint32_t off = grp * 64u;
|
||||
if (row >= n_tok || off >= n_nope) return;
|
||||
float *xr = x + (uint64_t)row * head_dim;
|
||||
__shared__ float scratch[64];
|
||||
float v = 0.0f;
|
||||
if (tid < 64u && off + tid < n_nope) v = xr[off + tid];
|
||||
scratch[tid] = (tid < 64u && off + tid < n_nope) ? fabsf(v) : 0.0f;
|
||||
__syncthreads();
|
||||
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
|
||||
__syncthreads();
|
||||
}
|
||||
const float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
|
||||
if (tid < 64u && off + tid < n_nope) {
|
||||
const float q = dsv4_e4m3fn_dequant_dev(fminf(448.0f, fmaxf(-448.0f, v / scale))) * scale;
|
||||
xr[off + tid] = q;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void store_raw_kv_batch_kernel(float *raw, const float *kv, uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * head_dim;
|
||||
if (gid >= n) return;
|
||||
uint32_t d = gid % head_dim;
|
||||
uint32_t t = gid / head_dim;
|
||||
uint32_t row = (pos0 + t) % raw_cap;
|
||||
const uint16_t hb = f32_to_f16_bits_hip_round(kv[(uint64_t)t * head_dim + d]);
|
||||
raw[(uint64_t)row * head_dim + d] = f16_bits_to_f32(hb);
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
extern "C" int ds4_gpu_dsv4_fp8_kv_quantize_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) {
|
||||
if (n_rot > head_dim || !cuda_tensor_has_elems2(x, n_tok, head_dim, sizeof(float))) return 0;
|
||||
if (n_tok == 0u || head_dim == 0u) return 1;
|
||||
const uint32_t n_nope = head_dim - n_rot;
|
||||
if (n_nope == 0) return 1;
|
||||
const uint32_t groups = (n_nope + 63u) / 64u;
|
||||
fp8_kv_quantize_kernel<<<dim3(n_tok, groups), 64>>>((float *)x->ptr, n_tok, head_dim, n_rot);
|
||||
return cuda_ok(cudaGetLastError(), "fp8_kv_quantize launch");
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
// DS4 ROCm hierarchical-composition (HC) split/sum/expand kernels.
|
||||
//
|
||||
// Included from ds4_cuda.cu in the same translation unit to preserve current
|
||||
// static helper visibility and launch behavior.
|
||||
|
||||
__device__ static void hc4_split_one(float *out, const float *mix, const float *scale, const float *base, uint32_t sinkhorn_iters, float epsv) {
|
||||
const float pre_scale = scale[0];
|
||||
const float post_scale = scale[1];
|
||||
const float comb_scale = scale[2];
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float z = mix[i] * pre_scale + base[i];
|
||||
out[i] = 1.0f / (1.0f + expf(-z)) + epsv;
|
||||
}
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float z = mix[4 + i] * post_scale + base[4 + i];
|
||||
out[4 + i] = 2.0f / (1.0f + expf(-z));
|
||||
}
|
||||
float c[16];
|
||||
for (int r = 0; r < 4; r++) {
|
||||
float m = -INFINITY;
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float v = mix[8 + r * 4 + col] * comb_scale + base[8 + r * 4 + col];
|
||||
c[r * 4 + col] = v;
|
||||
m = fmaxf(m, v);
|
||||
}
|
||||
float s = 0.0f;
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float v = expf(c[r * 4 + col] - m);
|
||||
c[r * 4 + col] = v;
|
||||
s += v;
|
||||
}
|
||||
for (int col = 0; col < 4; col++) c[r * 4 + col] = c[r * 4 + col] / s + epsv;
|
||||
}
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float s = epsv;
|
||||
for (int r = 0; r < 4; r++) s += c[r * 4 + col];
|
||||
for (int r = 0; r < 4; r++) c[r * 4 + col] /= s;
|
||||
}
|
||||
for (uint32_t iter = 1; iter < sinkhorn_iters; iter++) {
|
||||
for (int r = 0; r < 4; r++) {
|
||||
float s = epsv;
|
||||
for (int col = 0; col < 4; col++) s += c[r * 4 + col];
|
||||
for (int col = 0; col < 4; col++) c[r * 4 + col] /= s;
|
||||
}
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float s = epsv;
|
||||
for (int r = 0; r < 4; r++) s += c[r * 4 + col];
|
||||
for (int r = 0; r < 4; r++) c[r * 4 + col] /= s;
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < 16; i++) out[8 + i] = c[i];
|
||||
}
|
||||
|
||||
__global__ static void hc_split_sinkhorn_kernel(float *out, const float *mix, const float *scale, const float *base, uint32_t n_rows, uint32_t sinkhorn_iters, float epsv) {
|
||||
uint32_t row = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (row >= n_rows) return;
|
||||
hc4_split_one(out + (uint64_t)row * 24, mix + (uint64_t)row * 24, scale, base, sinkhorn_iters, epsv);
|
||||
}
|
||||
|
||||
__global__ static void hc_weighted_sum_kernel(float *out, const float *x, const float *w, uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens, uint32_t weight_stride_f32) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_embd * n_tokens;
|
||||
if (gid >= n) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint32_t t = gid / n_embd;
|
||||
float acc = 0.0f;
|
||||
for (uint32_t h = 0; h < n_hc; h++) {
|
||||
acc += x[(uint64_t)t * n_hc * n_embd + (uint64_t)h * n_embd + d] *
|
||||
w[(uint64_t)t * weight_stride_f32 + h];
|
||||
}
|
||||
out[(uint64_t)t * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const float *block_add,
|
||||
const float *residual_hc,
|
||||
const float *post,
|
||||
const float *comb,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_tokens,
|
||||
uint32_t post_stride,
|
||||
uint32_t comb_stride,
|
||||
int has_add) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n_elem) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t dst_hc = tmp % n_hc;
|
||||
uint32_t t = tmp / n_hc;
|
||||
|
||||
float block_v = block_out[(uint64_t)t * n_embd + d];
|
||||
if (has_add) block_v += block_add[(uint64_t)t * n_embd + d];
|
||||
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
|
||||
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
|
||||
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
|
||||
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand_half_kernel(
|
||||
float *out_hc,
|
||||
const __half *block_out,
|
||||
const float *residual_hc,
|
||||
const float *post,
|
||||
const float *comb,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_tokens,
|
||||
uint32_t post_stride,
|
||||
uint32_t comb_stride) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n_elem) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t dst_hc = tmp % n_hc;
|
||||
uint32_t t = tmp / n_hc;
|
||||
|
||||
const float block_v = __half2float(block_out[(uint64_t)t * n_embd + d]);
|
||||
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
|
||||
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
|
||||
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
|
||||
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand_add_half_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const __half *block_add,
|
||||
const float *residual_hc,
|
||||
const float *post,
|
||||
const float *comb,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_tokens,
|
||||
uint32_t post_stride,
|
||||
uint32_t comb_stride) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n_elem) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t dst_hc = tmp % n_hc;
|
||||
uint32_t t = tmp / n_hc;
|
||||
|
||||
const float block_v = block_out[(uint64_t)t * n_embd + d] +
|
||||
__half2float(block_add[(uint64_t)t * n_embd + d]);
|
||||
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
|
||||
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
|
||||
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
|
||||
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = block_out[td];
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_add_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const float *block_add,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = block_out[td] + block_add[td];
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_half_kernel(
|
||||
float *out_hc,
|
||||
const __half *block_out,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = __half2float(block_out[td]);
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_add_half_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const __half *block_add,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = block_out[td] + __half2float(block_add[td]);
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_split_weighted_sum_fused_kernel(
|
||||
float *out,
|
||||
float *split,
|
||||
const float *mix,
|
||||
const float *residual_hc,
|
||||
const float *scale,
|
||||
const float *base,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_rows,
|
||||
uint32_t sinkhorn_iters,
|
||||
float epsv) {
|
||||
uint32_t t = blockIdx.x;
|
||||
uint32_t d = threadIdx.x;
|
||||
if (t >= n_rows || n_hc != 4) return;
|
||||
const uint32_t mix_hc = 24;
|
||||
float *sp = split + (uint64_t)t * mix_hc;
|
||||
if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv);
|
||||
__syncthreads();
|
||||
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
|
||||
float acc = 0.0f;
|
||||
for (uint32_t h = 0; h < 4; h++) {
|
||||
acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h];
|
||||
}
|
||||
out[(uint64_t)t * n_embd + col] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_split_weighted_sum_norm_fused_kernel(
|
||||
float *out,
|
||||
float *norm_out,
|
||||
float *split,
|
||||
const float *mix,
|
||||
const float *residual_hc,
|
||||
const float *scale,
|
||||
const float *base,
|
||||
const float *norm_w,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_rows,
|
||||
uint32_t sinkhorn_iters,
|
||||
float epsv,
|
||||
float norm_eps) {
|
||||
const uint32_t t = blockIdx.x;
|
||||
const uint32_t d = threadIdx.x;
|
||||
if (t >= n_rows || n_hc != 4) return;
|
||||
const uint32_t mix_hc = 24;
|
||||
float *sp = split + (uint64_t)t * mix_hc;
|
||||
if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv);
|
||||
__syncthreads();
|
||||
|
||||
float sum = 0.0f;
|
||||
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
|
||||
float acc = 0.0f;
|
||||
for (uint32_t h = 0; h < 4; h++) {
|
||||
acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h];
|
||||
}
|
||||
out[(uint64_t)t * n_embd + col] = acc;
|
||||
sum += acc * acc;
|
||||
}
|
||||
|
||||
__shared__ float partial[256];
|
||||
partial[d] = sum;
|
||||
__syncthreads();
|
||||
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
|
||||
if (d < stride) partial[d] += partial[d + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
const float norm_scale = rsqrtf(partial[0] / (float)n_embd + norm_eps);
|
||||
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
|
||||
const float v = out[(uint64_t)t * n_embd + col];
|
||||
norm_out[(uint64_t)t * n_embd + col] = v * norm_scale * norm_w[col];
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
extern "C" int ds4_gpu_repeat_hc_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *row, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n = 0;
|
||||
if (n_embd == 0 || n_hc == 0 ||
|
||||
!cuda_u64_mul_checked(n_embd, n_hc, &n) ||
|
||||
!cuda_tensor_has_f32(row, n_embd) || !cuda_tensor_has_f32(out, n)) {
|
||||
return 0;
|
||||
}
|
||||
repeat_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)row->ptr, n_embd, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "repeat_hc launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_hc_split_sinkhorn_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *mix, const void *model_map, uint64_t model_size, uint64_t scale_offset, uint64_t base_offset, uint32_t n_hc, uint32_t sinkhorn_iters, float eps) {
|
||||
if (!out || !mix || !model_map || n_hc != 4) return 0;
|
||||
const uint64_t mix_bytes = 24ull * sizeof(float);
|
||||
if (!cuda_model_range_fits(model_size, scale_offset, 3ull * sizeof(float)) ||
|
||||
!cuda_model_range_fits(model_size, base_offset, mix_bytes) ||
|
||||
!cuda_tensor_has_bytes(mix, mix_bytes) || !cuda_tensor_has_bytes(out, mix_bytes)) return 0;
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset, 3ull * sizeof(float), "hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset, mix_bytes, "hc_base");
|
||||
if (!scale || !base) return 0;
|
||||
uint32_t n_rows = (uint32_t)(mix->bytes / mix_bytes);
|
||||
if (out->bytes / mix_bytes < n_rows) n_rows = (uint32_t)(out->bytes / mix_bytes);
|
||||
hc_split_sinkhorn_kernel<<<(n_rows + 255) / 256, 256>>>(
|
||||
(float *)out->ptr, (const float *)mix->ptr,
|
||||
scale,
|
||||
base,
|
||||
n_rows, sinkhorn_iters, eps);
|
||||
return cuda_ok(cudaGetLastError(), "hc_split_sinkhorn launch");
|
||||
}
|
||||
static int cuda_hc_flat_token_count(const ds4_gpu_tensor *out, uint32_t n_embd, uint64_t *n_tokens) {
|
||||
if (!out || n_embd == 0u || !n_tokens) return 0;
|
||||
uint64_t row_bytes = 0;
|
||||
if (!cuda_u64_mul3_checked(n_embd, 1u, sizeof(float), &row_bytes) ||
|
||||
row_bytes == 0u || out->bytes < row_bytes || (out->bytes % row_bytes) != 0u) return 0;
|
||||
*n_tokens = out->bytes / row_bytes;
|
||||
return *n_tokens != 0u && *n_tokens <= UINT32_MAX;
|
||||
}
|
||||
|
||||
static int cuda_hc_hc_token_count(const ds4_gpu_tensor *out_hc, uint32_t n_embd, uint32_t n_hc, uint64_t *n_tokens) {
|
||||
if (!out_hc || n_embd == 0u || n_hc == 0u || !n_tokens) return 0;
|
||||
uint64_t row_elems = 0, row_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(n_hc, n_embd, &row_elems) ||
|
||||
!cuda_u64_mul_checked(row_elems, sizeof(float), &row_bytes) ||
|
||||
row_bytes == 0u || out_hc->bytes < row_bytes || (out_hc->bytes % row_bytes) != 0u) return 0;
|
||||
*n_tokens = out_hc->bytes / row_bytes;
|
||||
return *n_tokens != 0u && *n_tokens <= UINT32_MAX;
|
||||
}
|
||||
|
||||
static int cuda_hc_mix_width(uint32_t n_hc, uint64_t *mix_hc) {
|
||||
if (n_hc == 0u || !mix_hc) return 0;
|
||||
const uint64_t h = (uint64_t)n_hc;
|
||||
const uint64_t mix = 2ull * h + h * h;
|
||||
if (mix > UINT32_MAX) return 0;
|
||||
*mix_hc = mix;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_hc_weighted_sum_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *weights, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, residual_bytes = 0, weights_bytes = 0;
|
||||
if (!out || !residual_hc || !weights || n_hc == 0u ||
|
||||
!cuda_hc_flat_token_count(out, n_embd, &n_tokens64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &residual_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_hc, sizeof(float), &weights_bytes) ||
|
||||
residual_hc->bytes < residual_bytes || weights->bytes < weights_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>(
|
||||
(float *)out->ptr, (const float *)residual_hc->ptr, (const float *)weights->ptr,
|
||||
n_embd, n_hc, n_tokens, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "hc_weighted_sum launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_weighted_sum_split_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, residual_bytes = 0, split_bytes = 0, mix_hc = 0;
|
||||
if (!out || !residual_hc || !split ||
|
||||
!cuda_hc_flat_token_count(out, n_embd, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &residual_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc, sizeof(float), &split_bytes) ||
|
||||
residual_hc->bytes < residual_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
uint32_t stride = (uint32_t)mix_hc;
|
||||
hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>(
|
||||
(float *)out->ptr, (const float *)residual_hc->ptr, (const float *)split->ptr,
|
||||
n_embd, n_hc, n_tokens, stride);
|
||||
return cuda_ok(cudaGetLastError(), "hc_weighted_sum_split launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_split_weighted_sum_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *split,
|
||||
const ds4_gpu_tensor *mix,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t sinkhorn_iters,
|
||||
float eps) {
|
||||
if (!out || !split || !mix || !residual_hc || !model_map ||
|
||||
n_embd == 0 || n_hc != 4) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
|
||||
const uint64_t mix_bytes = mix_hc * sizeof(float);
|
||||
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
|
||||
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
|
||||
if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 ||
|
||||
scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset ||
|
||||
base_offset > model_size || mix_bytes > model_size - base_offset) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t n_rows = out->bytes / out_row_bytes;
|
||||
if (mix->bytes < n_rows * mix_bytes ||
|
||||
split->bytes < n_rows * mix_bytes ||
|
||||
residual_hc->bytes < n_rows * residual_row_bytes) {
|
||||
return 0;
|
||||
}
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset, 3ull * sizeof(float), "hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset, mix_bytes, "hc_base");
|
||||
if (!scale || !base) return 0;
|
||||
hc_split_weighted_sum_fused_kernel<<<(uint32_t)n_rows, 256>>>(
|
||||
(float *)out->ptr,
|
||||
(float *)split->ptr,
|
||||
(const float *)mix->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
scale,
|
||||
base,
|
||||
n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps);
|
||||
return cuda_ok(cudaGetLastError(), "hc split weighted sum launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_split_weighted_sum_norm_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *norm_out,
|
||||
ds4_gpu_tensor *split,
|
||||
const ds4_gpu_tensor *mix,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint64_t norm_weight_offset,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t sinkhorn_iters,
|
||||
float eps,
|
||||
float norm_eps) {
|
||||
if (!out || !norm_out || !split || !mix || !residual_hc || !model_map ||
|
||||
n_embd == 0 || n_hc != 4) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
|
||||
const uint64_t mix_bytes = mix_hc * sizeof(float);
|
||||
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
|
||||
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
|
||||
if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 ||
|
||||
norm_out->bytes < out->bytes ||
|
||||
scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset ||
|
||||
base_offset > model_size || mix_bytes > model_size - base_offset ||
|
||||
norm_weight_offset > model_size ||
|
||||
(uint64_t)n_embd * sizeof(float) > model_size - norm_weight_offset) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t n_rows = out->bytes / out_row_bytes;
|
||||
if (n_rows == 1) {
|
||||
if (mix->bytes < n_rows * mix_bytes ||
|
||||
split->bytes < n_rows * mix_bytes ||
|
||||
residual_hc->bytes < n_rows * residual_row_bytes) {
|
||||
return 0;
|
||||
}
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset,
|
||||
3ull * sizeof(float), "hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset,
|
||||
mix_bytes, "hc_base");
|
||||
const float *norm_w = (const float *)cuda_model_range_ptr(model_map, norm_weight_offset,
|
||||
(uint64_t)n_embd * sizeof(float), "hc_norm_weight");
|
||||
if (!scale || !base || !norm_w) return 0;
|
||||
hc_split_weighted_sum_norm_fused_kernel<<<(uint32_t)n_rows, 256>>>(
|
||||
(float *)out->ptr,
|
||||
(float *)norm_out->ptr,
|
||||
(float *)split->ptr,
|
||||
(const float *)mix->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
scale,
|
||||
base,
|
||||
norm_w,
|
||||
n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps, norm_eps);
|
||||
return cuda_ok(cudaGetLastError(), "hc split weighted sum norm launch");
|
||||
}
|
||||
return ds4_gpu_hc_split_weighted_sum_tensor(out, split, mix, residual_hc,
|
||||
model_map, model_size,
|
||||
scale_offset, base_offset,
|
||||
n_embd, n_hc,
|
||||
sinkhorn_iters, eps) &&
|
||||
ds4_gpu_rms_norm_weight_tensor(norm_out, out, model_map, model_size,
|
||||
norm_weight_offset, n_embd, norm_eps);
|
||||
}
|
||||
extern "C" int ds4_gpu_output_hc_weights_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
const ds4_gpu_tensor *pre,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint32_t n_hc,
|
||||
float eps) {
|
||||
if (!out || !pre || !model_map || n_hc == 0) return 0;
|
||||
const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float);
|
||||
if (row_bytes == 0 || out->bytes < row_bytes || out->bytes % row_bytes != 0 ||
|
||||
pre->bytes < out->bytes ||
|
||||
scale_offset > model_size || sizeof(float) > model_size - scale_offset ||
|
||||
base_offset > model_size || row_bytes > model_size - base_offset) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t n_tokens = out->bytes / row_bytes;
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset, sizeof(float), "output_hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset, row_bytes, "output_hc_base");
|
||||
if (!scale || !base) return 0;
|
||||
uint64_t n = n_tokens * n_hc;
|
||||
output_hc_weights_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)out->ptr,
|
||||
(const float *)pre->ptr,
|
||||
scale,
|
||||
base,
|
||||
n_hc,
|
||||
(uint32_t)n_tokens,
|
||||
eps);
|
||||
return cuda_ok(cudaGetLastError(), "output hc weights launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, hc_bytes = 0, post_bytes = 0, comb_bytes = 0, comb_stride = 0;
|
||||
if (!out_hc || !block_out || !residual_hc || !post || !comb ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_hc, sizeof(float), &post_bytes) ||
|
||||
!cuda_u64_mul_checked(n_hc, n_hc, &comb_stride) || comb_stride > UINT32_MAX ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, comb_stride, sizeof(float), &comb_bytes) ||
|
||||
block_out->bytes < flat_bytes || residual_hc->bytes < hc_bytes ||
|
||||
post->bytes < post_bytes || comb->bytes < comb_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)post->ptr,
|
||||
(const float *)comb->ptr,
|
||||
n_embd, n_hc, n_tokens,
|
||||
n_hc, (uint32_t)comb_stride, 0);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out->bytes < flat_bytes || residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc, 0);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_split_half_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out_h, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_half_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out_h || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(__half), &flat_half_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out_h->bytes < flat_half_bytes || residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_half_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const __half *)block_out_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split_half4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_half_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const __half *)block_out_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split_half launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_add_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out || !block_add || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out->bytes < flat_bytes || block_add->bytes < flat_bytes ||
|
||||
residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_add_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_add->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_add->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc, 1);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_add_split_half_add_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add_h, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, flat_half_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out || !block_add_h || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(__half), &flat_half_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out->bytes < flat_bytes || block_add_h->bytes < flat_half_bytes ||
|
||||
residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_add_half_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const __half *)block_add_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split_half4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_add_half_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const __half *)block_add_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split_half_add launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_shared_down_hc_expand_q8_0_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
ds4_gpu_tensor *shared_out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *shared_mid,
|
||||
const ds4_gpu_tensor *routed_out,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out,
|
||||
model_map, model_size,
|
||||
weight_offset,
|
||||
in_dim, out_dim,
|
||||
shared_mid,
|
||||
routed_out,
|
||||
residual_hc,
|
||||
split,
|
||||
n_embd, n_hc,
|
||||
"shared_down_hc_expand");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_hc_expand_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
ds4_gpu_tensor *block_out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, block_out,
|
||||
model_map, model_size,
|
||||
weight_offset,
|
||||
in_dim, out_dim,
|
||||
x,
|
||||
NULL,
|
||||
residual_hc,
|
||||
split,
|
||||
n_embd, n_hc,
|
||||
"q8_hc_expand");
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
/* HIP-only hipBLASLt state and helpers.
|
||||
* Included from ds4_cuda.cu under __HIP_PLATFORM_AMD__ to keep ROCm
|
||||
* planning/cache code out of the CUDA host runtime body. */
|
||||
|
||||
static hipblasLtHandle_t g_hipblaslt;
|
||||
static int g_hipblaslt_ready;
|
||||
struct cuda_hipblaslt_gemm_plan {
|
||||
uint32_t out_dim;
|
||||
uint32_t n_tok;
|
||||
uint32_t in_dim;
|
||||
hipblasLtMatmulDesc_t desc;
|
||||
hipblasLtMatrixLayout_t a_desc;
|
||||
hipblasLtMatrixLayout_t b_desc;
|
||||
hipblasLtMatrixLayout_t c_desc;
|
||||
hipblasLtMatrixLayout_t d_desc;
|
||||
hipblasLtMatmulAlgo_t algo;
|
||||
};
|
||||
static std::vector<cuda_hipblaslt_gemm_plan> g_hipblaslt_gemm_plans;
|
||||
|
||||
static void hipblaslt_gemm_plan_clear(void) {
|
||||
for (size_t i = 0; i < g_hipblaslt_gemm_plans.size(); i++) {
|
||||
cuda_hipblaslt_gemm_plan &p = g_hipblaslt_gemm_plans[i];
|
||||
if (p.d_desc) (void)hipblasLtMatrixLayoutDestroy(p.d_desc);
|
||||
if (p.c_desc) (void)hipblasLtMatrixLayoutDestroy(p.c_desc);
|
||||
if (p.b_desc) (void)hipblasLtMatrixLayoutDestroy(p.b_desc);
|
||||
if (p.a_desc) (void)hipblasLtMatrixLayoutDestroy(p.a_desc);
|
||||
if (p.desc) (void)hipblasLtMatmulDescDestroy(p.desc);
|
||||
}
|
||||
g_hipblaslt_gemm_plans.clear();
|
||||
}
|
||||
|
||||
static int hipblaslt_ok(hipblasStatus_t st, const char *what) {
|
||||
if (st == HIPBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: hipBLASLt %s failed: status %d\n", what, (int)st);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static cuda_hipblaslt_gemm_plan *hipblaslt_gemm_plan_get(
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint32_t in_dim,
|
||||
const char *label) {
|
||||
for (size_t i = 0; i < g_hipblaslt_gemm_plans.size(); i++) {
|
||||
cuda_hipblaslt_gemm_plan &p = g_hipblaslt_gemm_plans[i];
|
||||
if (p.out_dim == out_dim && p.n_tok == n_tok && p.in_dim == in_dim) return &p;
|
||||
}
|
||||
|
||||
hipblasLtMatmulDesc_t desc = NULL;
|
||||
hipblasLtMatrixLayout_t a_desc = NULL, b_desc = NULL, c_desc = NULL, d_desc = NULL;
|
||||
hipblasLtMatmulPreference_t pref = NULL;
|
||||
hipblasLtMatmulHeuristicResult_t heur[8];
|
||||
int returned = 0;
|
||||
int ok = 0;
|
||||
do {
|
||||
if (!hipblaslt_ok(hipblasLtMatmulDescCreate(&desc, HIPBLAS_COMPUTE_32F, HIP_R_32F),
|
||||
"matmul desc create")) break;
|
||||
hipblasOperation_t op_a = HIPBLAS_OP_T;
|
||||
hipblasOperation_t op_b = HIPBLAS_OP_N;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulDescSetAttribute(desc, HIPBLASLT_MATMUL_DESC_TRANSA,
|
||||
&op_a, sizeof(op_a)),
|
||||
"set transA")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulDescSetAttribute(desc, HIPBLASLT_MATMUL_DESC_TRANSB,
|
||||
&op_b, sizeof(op_b)),
|
||||
"set transB")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&a_desc, HIP_R_16F, in_dim, out_dim, in_dim),
|
||||
"A layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&b_desc, HIP_R_16F, in_dim, n_tok, in_dim),
|
||||
"B layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&c_desc, HIP_R_16F, out_dim, n_tok, out_dim),
|
||||
"C layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&d_desc, HIP_R_16F, out_dim, n_tok, out_dim),
|
||||
"D layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulPreferenceCreate(&pref), "preference create")) break;
|
||||
const size_t max_workspace = 0;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulPreferenceSetAttribute(
|
||||
pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
|
||||
&max_workspace, sizeof(max_workspace)),
|
||||
"set max workspace")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulAlgoGetHeuristic(g_hipblaslt, desc,
|
||||
a_desc, b_desc, c_desc, d_desc,
|
||||
pref, 8, heur, &returned),
|
||||
"algo heuristic")) break;
|
||||
if (returned <= 0 || heur[0].state != HIPBLAS_STATUS_SUCCESS) {
|
||||
fprintf(stderr, "ds4: hipBLASLt no algo for %s m=%u n=%u k=%u\n",
|
||||
label ? label : "gemm", out_dim, n_tok, in_dim);
|
||||
break;
|
||||
}
|
||||
ok = 1;
|
||||
} while (0);
|
||||
if (pref) (void)hipblasLtMatmulPreferenceDestroy(pref);
|
||||
if (!ok) {
|
||||
if (d_desc) (void)hipblasLtMatrixLayoutDestroy(d_desc);
|
||||
if (c_desc) (void)hipblasLtMatrixLayoutDestroy(c_desc);
|
||||
if (b_desc) (void)hipblasLtMatrixLayoutDestroy(b_desc);
|
||||
if (a_desc) (void)hipblasLtMatrixLayoutDestroy(a_desc);
|
||||
if (desc) (void)hipblasLtMatmulDescDestroy(desc);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
cuda_hipblaslt_gemm_plan p;
|
||||
p.out_dim = out_dim;
|
||||
p.n_tok = n_tok;
|
||||
p.in_dim = in_dim;
|
||||
p.desc = desc;
|
||||
p.a_desc = a_desc;
|
||||
p.b_desc = b_desc;
|
||||
p.c_desc = c_desc;
|
||||
p.d_desc = d_desc;
|
||||
p.algo = heur[0].algo;
|
||||
g_hipblaslt_gemm_plans.push_back(p);
|
||||
return &g_hipblaslt_gemm_plans.back();
|
||||
}
|
||||
|
||||
static int hipblaslt_gemm_tn_f16_out_f16(
|
||||
__half *out,
|
||||
const __half *w_rowmajor_out_in,
|
||||
const __half *x_rowmajor_tok_in,
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint32_t in_dim,
|
||||
const char *label) {
|
||||
if (!g_hipblaslt_ready || !out || !w_rowmajor_out_in || !x_rowmajor_tok_in ||
|
||||
out_dim == 0 || n_tok == 0 || in_dim == 0) return 0;
|
||||
cuda_hipblaslt_gemm_plan *p = hipblaslt_gemm_plan_get(out_dim, n_tok, in_dim, label);
|
||||
if (!p) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
return hipblaslt_ok(hipblasLtMatmul(g_hipblaslt, p->desc, &alpha,
|
||||
w_rowmajor_out_in, p->a_desc,
|
||||
x_rowmajor_tok_in, p->b_desc,
|
||||
&beta,
|
||||
out, p->c_desc,
|
||||
out, p->d_desc,
|
||||
&p->algo,
|
||||
NULL, 0, 0),
|
||||
label ? label : "gemm");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,791 @@
|
||||
template <uint32_t BT>
|
||||
static void cuda_launch_q8_batch_sharedx_bt(
|
||||
float *out,
|
||||
const unsigned char *w,
|
||||
const float *x,
|
||||
uint32_t n_blocks,
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint64_t row_bytes,
|
||||
dim3 grid,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile) {
|
||||
const size_t shmem = (size_t)tile * BT * 32u * sizeof(float);
|
||||
if (tile == 2u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<2u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else if (tile == 4u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<4u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else if (tile == 8u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<8u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else if (tile == 16u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<16u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<32u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
static void cuda_launch_q8_batch_sharedx(
|
||||
float *out,
|
||||
const unsigned char *w,
|
||||
const float *x,
|
||||
uint32_t n_blocks,
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint64_t row_bytes,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile,
|
||||
uint32_t block_tile) {
|
||||
const dim3 grid((out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
(n_tok + tile - 1u) / tile,
|
||||
1u);
|
||||
if (block_tile == 8u) {
|
||||
cuda_launch_q8_batch_sharedx_bt<8u>(out, w, x, n_blocks, out_dim, n_tok, row_bytes, grid, rows_per_block, tile);
|
||||
} else if (block_tile == 32u) {
|
||||
cuda_launch_q8_batch_sharedx_bt<32u>(out, w, x, n_blocks, out_dim, n_tok, row_bytes, grid, rows_per_block, tile);
|
||||
} else {
|
||||
cuda_launch_q8_batch_sharedx_bt<16u>(out, w, x, n_blocks, out_dim, n_tok, row_bytes, grid, rows_per_block, tile);
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t BT>
|
||||
static void cuda_launch_grouped_q8_a_sharedx_bt(
|
||||
float *low,
|
||||
const unsigned char *w,
|
||||
const float *heads,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_groups,
|
||||
uint32_t n_blocks,
|
||||
uint32_t rank,
|
||||
uint64_t row_bytes,
|
||||
dim3 grid,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile) {
|
||||
const size_t shmem = (size_t)tile * BT * 32u * sizeof(float);
|
||||
if (tile == 2u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<2u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else if (tile == 4u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<4u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else if (tile == 8u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<8u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else if (tile == 16u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<16u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<32u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
static void cuda_launch_grouped_q8_a_sharedx(
|
||||
float *low,
|
||||
const unsigned char *w,
|
||||
const float *heads,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_groups,
|
||||
uint32_t n_blocks,
|
||||
uint32_t rank,
|
||||
uint64_t row_bytes,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile,
|
||||
uint32_t block_tile) {
|
||||
const uint32_t row_blocks = (rank + rows_per_block - 1u) / rows_per_block;
|
||||
const dim3 grid(n_groups * row_blocks,
|
||||
(n_tokens + tile - 1u) / tile,
|
||||
1u);
|
||||
if (block_tile == 8u) {
|
||||
cuda_launch_grouped_q8_a_sharedx_bt<8u>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes, grid, rows_per_block, tile);
|
||||
} else if (block_tile == 32u) {
|
||||
cuda_launch_grouped_q8_a_sharedx_bt<32u>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes, grid, rows_per_block, tile);
|
||||
} else {
|
||||
cuda_launch_grouped_q8_a_sharedx_bt<16u>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes, grid, rows_per_block, tile);
|
||||
}
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_tensor_f16_gemm(
|
||||
ds4_gpu_tensor *out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok,
|
||||
const char *label) {
|
||||
if (!g_cublas_ready || !out || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
|
||||
if (!w_f16) return 0;
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "q8 f16 gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out->ptr,
|
||||
CUDA_R_32F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st == CUBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " q8 f16 matmul failed: status %d\n", (int)st);
|
||||
cuda_q8_f16_cache_disable_after_failure(DS4_GPU_BLAS_NAME " f16 matmul failure",
|
||||
in_dim * out_dim * sizeof(__half));
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_tensor_f16_gemm_out_half(
|
||||
ds4_gpu_tensor *out_h,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok,
|
||||
const char *label) {
|
||||
if (!g_cublas_ready || !out_h || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(__half), &out_bytes) ||
|
||||
x->bytes < x_bytes || out_h->bytes < out_bytes) return 0;
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
|
||||
if (!w_f16) return 0;
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "q8 f16-out gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "q8 f16-out activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out_h->ptr,
|
||||
CUDA_R_16F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st == CUBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " q8 f16-out matmul failed: status %d\n", (int)st);
|
||||
cuda_q8_f16_cache_disable_after_failure(DS4_GPU_BLAS_NAME " f16-out matmul failure",
|
||||
in_dim * out_dim * sizeof(__half));
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_f16_out_tensor(
|
||||
ds4_gpu_tensor *out_h,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
return cuda_matmul_q8_0_tensor_f16_gemm_out_half(out_h, model_map, model_size,
|
||||
weight_offset, in_dim, out_dim,
|
||||
x, n_tok, "q8_f16_out");
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok, const char *label) {
|
||||
if (!out || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
if (n_tok > 1 && !g_quality_mode &&
|
||||
cuda_runtime_config()->shared_down_cublas && in_dim == 2048u && out_dim == 4096u &&
|
||||
cuda_matmul_q8_0_tensor_f16_gemm(out, model_map, model_size, weight_offset,
|
||||
in_dim, out_dim, x, n_tok, label ? label : "shared_expert")) {
|
||||
return 1;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "q8_0");
|
||||
if (!wptr) return 0;
|
||||
if (n_tok == 1 && !cuda_runtime_config()->q8_prequant_decode) {
|
||||
if ((in_dim & 31u) == 0u && in_dim <= 8192u) {
|
||||
const unsigned rows_per_block = 32u;
|
||||
const unsigned threads = rows_per_block * 32u;
|
||||
matmul_q8_0_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
(unsigned)((out_dim + rows_per_block - 1u) / rows_per_block),
|
||||
threads,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out_dim,
|
||||
blocks * 34u);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 sharedx launch");
|
||||
}
|
||||
matmul_q8_0_f32_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
blocks);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 warp launch");
|
||||
}
|
||||
if (n_tok > 1) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
if (!g_quality_mode && (in_dim % 32u) == 0u &&
|
||||
out_dim >= 1024u &&
|
||||
n_tok >= 256u &&
|
||||
in_dim <= UINT32_MAX && out_dim <= UINT32_MAX && n_tok <= UINT32_MAX) {
|
||||
const dim3 grid((uint32_t)((out_dim + 63u) / 64u),
|
||||
(uint32_t)((n_tok + 63u) / 64u),
|
||||
1u);
|
||||
matmul_q8_0_f32_batch_wmma_4w_kernel<<<grid, 128u>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)n_tok,
|
||||
(uint32_t)in_dim,
|
||||
(uint32_t)out_dim,
|
||||
blocks * 34u);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch wmma 4w launch");
|
||||
}
|
||||
#endif
|
||||
if ((in_dim & 31u) == 0u && out_dim <= UINT32_MAX && n_tok <= UINT32_MAX) {
|
||||
const uint32_t rows_per_block = 32u;
|
||||
const uint32_t tile = 32u;
|
||||
const uint32_t block_tile = 16u;
|
||||
cuda_launch_q8_batch_sharedx((float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
(uint32_t)out_dim,
|
||||
(uint32_t)n_tok,
|
||||
blocks * 34u,
|
||||
rows_per_block,
|
||||
tile,
|
||||
block_tile);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch sharedx launch");
|
||||
}
|
||||
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1);
|
||||
matmul_q8_0_f32_batch_warp8_kernel<<<bgrid, 256>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_tok,
|
||||
blocks);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch warp launch");
|
||||
}
|
||||
if (g_cublas_ready && n_tok > 1) {
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
|
||||
if (w_f16) {
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "q8 f16 gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out->ptr,
|
||||
CUDA_R_32F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st == CUBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " q8 f16 matmul failed: status %d\n", (int)st);
|
||||
cuda_q8_f16_cache_disable_after_failure(DS4_GPU_BLAS_NAME " f16 matmul failure",
|
||||
in_dim * out_dim * sizeof(__half));
|
||||
/* The F16 expansion cache is only an optimization. If cuBLAS
|
||||
* rejects the cached path under memory pressure, retry the same
|
||||
* operation through the native Q8 kernels below. */
|
||||
}
|
||||
}
|
||||
const uint64_t xq_bytes = n_tok * blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float);
|
||||
void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 prequant");
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const ds4_rocm_runtime_config *cfg = cuda_runtime_config();
|
||||
const int use_dp4a = 1;
|
||||
dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1);
|
||||
quantize_q8_0_f32_kernel<<<qgrid, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0;
|
||||
if (n_tok == 1) {
|
||||
uint32_t rows_per_block = cfg->q8_decode_rpb;
|
||||
matmul_q8_0_preq_rows_w32_kernel<<<((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
blocks,
|
||||
rows_per_block,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 rows launch");
|
||||
}
|
||||
if (blocks <= 32u) {
|
||||
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1);
|
||||
matmul_q8_0_preq_batch_warp8_kernel<<<bgrid, 256>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_tok,
|
||||
blocks,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch warp launch");
|
||||
}
|
||||
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
|
||||
matmul_q8_0_preq_kernel<<<grid, 256>>>((float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim, out_dim, n_tok, blocks,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
|
||||
return cuda_matmul_q8_0_tensor_labeled(out, model_map, model_size, weight_offset,
|
||||
in_dim, out_dim, x, n_tok, "q8_0");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_pair_tensor(
|
||||
ds4_gpu_tensor *out0,
|
||||
ds4_gpu_tensor *out1,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight0_offset,
|
||||
uint64_t weight1_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out0_dim,
|
||||
uint64_t out1_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
if (!out0 || !out1 || !x || !model_map ||
|
||||
in_dim == 0 || out0_dim == 0 || out1_dim == 0 || n_tok == 0 ||
|
||||
in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX || n_tok > UINT32_MAX) {
|
||||
return 0;
|
||||
}
|
||||
if (n_tok != 1) {
|
||||
return cuda_matmul_q8_0_tensor_labeled(out0, model_map, model_size, weight0_offset,
|
||||
in_dim, out0_dim, x, n_tok, "q8_0_pair0") &&
|
||||
cuda_matmul_q8_0_tensor_labeled(out1, model_map, model_size, weight1_offset,
|
||||
in_dim, out1_dim, x, n_tok, "q8_0_pair1");
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight0_bytes = 0, weight1_bytes = 0;
|
||||
if (weight0_offset > model_size || weight1_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out0_dim, row_bytes, &weight0_bytes) ||
|
||||
!cuda_u64_mul_checked(out1_dim, row_bytes, &weight1_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
if (weight0_bytes > model_size - weight0_offset ||
|
||||
weight1_bytes > model_size - weight1_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
out0->bytes < out0_dim * sizeof(float) ||
|
||||
out1->bytes < out1_dim * sizeof(float)) {
|
||||
return 0;
|
||||
}
|
||||
const char *w0 = cuda_model_range_ptr(model_map, weight0_offset, weight0_bytes, "q8_0_pair0");
|
||||
const char *w1 = cuda_model_range_ptr(model_map, weight1_offset, weight1_bytes, "q8_0_pair1");
|
||||
if (!w0 || !w1) return 0;
|
||||
if (!cuda_runtime_config()->q8_prequant_decode) {
|
||||
const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim;
|
||||
if ((in_dim & 31u) == 0u && in_dim <= 8192u) {
|
||||
const unsigned rows_per_block = 32u;
|
||||
const unsigned threads = rows_per_block * 32u;
|
||||
matmul_q8_0_pair_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
(unsigned)((max_out + rows_per_block - 1u) / rows_per_block),
|
||||
threads,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
reinterpret_cast<const unsigned char *>(w0),
|
||||
reinterpret_cast<const unsigned char *>(w1),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out0_dim,
|
||||
out1_dim,
|
||||
blocks * 34u);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 sharedx launch");
|
||||
}
|
||||
matmul_q8_0_pair_f32_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
reinterpret_cast<const unsigned char *>(w0),
|
||||
reinterpret_cast<const unsigned char *>(w1),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out0_dim,
|
||||
out1_dim,
|
||||
blocks);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 warp launch");
|
||||
}
|
||||
|
||||
const uint64_t xq_bytes = blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
|
||||
void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 pair prequant");
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const int use_dp4a = 1;
|
||||
dim3 qgrid((unsigned)blocks, 1, 1);
|
||||
quantize_q8_0_f32_kernel<<<qgrid, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair quantize launch")) return 0;
|
||||
const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim;
|
||||
matmul_q8_0_pair_preq_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
reinterpret_cast<const unsigned char *>(w0),
|
||||
reinterpret_cast<const unsigned char *>(w1),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out0_dim,
|
||||
out1_dim,
|
||||
blocks,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair warp launch");
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_hc_expand_tensor_labeled(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
ds4_gpu_tensor *block_out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
const ds4_gpu_tensor *block_add,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
const char *label) {
|
||||
if (!out_hc || !block_out || !x || !residual_hc || !split || !model_map ||
|
||||
in_dim == 0 || out_dim == 0 || n_embd == 0 || n_hc == 0 ||
|
||||
out_dim != (uint64_t)n_embd) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31) / 32;
|
||||
if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0;
|
||||
const uint64_t weight_bytes = out_dim * blocks * 34;
|
||||
const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
|
||||
const uint64_t split_bytes = (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float);
|
||||
if (weight_bytes > model_size - weight_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
block_out->bytes < out_dim * sizeof(float) ||
|
||||
residual_hc->bytes < hc_bytes ||
|
||||
split->bytes < split_bytes ||
|
||||
out_hc->bytes < hc_bytes ||
|
||||
(block_add && block_add->bytes < out_dim * sizeof(float))) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, label ? label : "q8_0_hc_expand");
|
||||
if (!wptr) return 0;
|
||||
if (!cuda_runtime_config()->q8_prequant_decode) {
|
||||
if ((in_dim & 31u) == 0u && in_dim <= 8192u) {
|
||||
const unsigned rows_per_block = 32u;
|
||||
const unsigned threads = rows_per_block * 32u;
|
||||
matmul_q8_0_hc_expand_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
(unsigned)((out_dim + rows_per_block - 1u) / rows_per_block),
|
||||
threads,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(float *)block_out->ptr,
|
||||
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out_dim,
|
||||
blocks * 34u,
|
||||
n_embd,
|
||||
n_hc,
|
||||
block_add ? 1 : 0);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 sharedx launch");
|
||||
}
|
||||
matmul_q8_0_hc_expand_f32_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(float *)block_out->ptr,
|
||||
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_embd,
|
||||
n_hc,
|
||||
blocks,
|
||||
block_add ? 1 : 0);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 launch");
|
||||
}
|
||||
|
||||
const uint64_t xq_bytes = blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
|
||||
void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 hc expand prequant");
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const ds4_rocm_runtime_config *cfg = cuda_runtime_config();
|
||||
const int use_dp4a = 1;
|
||||
quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand quantize launch")) return 0;
|
||||
uint32_t rows_per_block = cfg->q8_hc_decode_rpb;
|
||||
matmul_q8_0_hc_expand_preq_rows_w32_kernel<<<((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(float *)block_out->ptr,
|
||||
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_embd,
|
||||
n_hc,
|
||||
blocks,
|
||||
rows_per_block,
|
||||
block_add ? 1 : 0,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand rows launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_f16_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
|
||||
if (!out || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
uint64_t weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul3_checked(out_dim, in_dim, sizeof(uint16_t), &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "f16");
|
||||
if (!wptr) return 0;
|
||||
const __half *w = (const __half *)wptr;
|
||||
const int ordered_decode = n_tok == 1u;
|
||||
if (g_cublas_ready && n_tok > 1) {
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "f16 gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out->ptr,
|
||||
CUDA_R_32F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
return cublas_ok(st, "f16 matmul");
|
||||
}
|
||||
/* The 4096x256 F16 router projection is latency-bound and the ordered
|
||||
* 32-thread row kernel is at least as fast on gfx1151; keep shared-X for
|
||||
* compressor/indexer F16 decode where reusing x across rows is the win. */
|
||||
const bool f16_decode_router_shape = (in_dim == 4096u && out_dim == 256u);
|
||||
if (n_tok == 1u && !g_quality_mode && !cuda_runtime_config()->graph_dump &&
|
||||
!f16_decode_router_shape) {
|
||||
if (in_dim <= 8192u && in_dim * sizeof(float) <= 65536u) {
|
||||
const uint32_t rows_per_block = 32u;
|
||||
matmul_f16_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out->ptr, w, (const float *)x->ptr, (uint32_t)in_dim, out_dim);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16 sharedx launch");
|
||||
}
|
||||
}
|
||||
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
|
||||
if (ordered_decode) {
|
||||
matmul_f16_ordered_chunks_kernel<<<grid, 32>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16_ordered_chunks launch");
|
||||
}
|
||||
matmul_f16_kernel<<<grid, 256>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_f16_pair_tensor(
|
||||
ds4_gpu_tensor *out0,
|
||||
ds4_gpu_tensor *out1,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight0_offset,
|
||||
uint64_t weight1_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0 ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
|
||||
return 0;
|
||||
}
|
||||
if (n_tok != 1) {
|
||||
return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset,
|
||||
in_dim, out_dim, x, n_tok) &&
|
||||
ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset,
|
||||
in_dim, out_dim, x, n_tok);
|
||||
}
|
||||
uint64_t weight_bytes = 0;
|
||||
if (weight0_offset > model_size || weight1_offset > model_size ||
|
||||
!cuda_u64_mul3_checked(out_dim, in_dim, sizeof(uint16_t), &weight_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
if (weight_bytes > model_size - weight0_offset ||
|
||||
weight_bytes > model_size - weight1_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
out0->bytes < out_dim * sizeof(float) ||
|
||||
out1->bytes < out_dim * sizeof(float)) {
|
||||
return 0;
|
||||
}
|
||||
const __half *w0 = (const __half *)cuda_model_range_ptr(model_map, weight0_offset, weight_bytes, "f16_pair0");
|
||||
const __half *w1 = (const __half *)cuda_model_range_ptr(model_map, weight1_offset, weight_bytes, "f16_pair1");
|
||||
if (!w0 || !w1) return 0;
|
||||
if (!g_quality_mode && !cuda_runtime_config()->graph_dump) {
|
||||
if (in_dim <= 8192u && in_dim * sizeof(float) <= 65536u) {
|
||||
const uint32_t rows_per_block = 32u;
|
||||
matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out0->ptr, (float *)out1->ptr, w0, w1,
|
||||
(const float *)x->ptr, (uint32_t)in_dim, out_dim);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16_pair sharedx launch");
|
||||
}
|
||||
}
|
||||
matmul_f16_pair_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
w0,
|
||||
w1,
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
out_dim);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_f32_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
|
||||
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0 ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
uint64_t weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul3_checked(out_dim, in_dim, sizeof(float), &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "f32");
|
||||
if (!wptr) return 0;
|
||||
const float *w = (const float *)wptr;
|
||||
if (g_cublas_ready && n_tok > 1) {
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasSgemm(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w,
|
||||
(int)in_dim,
|
||||
(const float *)x->ptr,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
(float *)out->ptr,
|
||||
(int)out_dim);
|
||||
return cublas_ok(st, "f32 matmul");
|
||||
}
|
||||
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
|
||||
matmul_f32_kernel<<<grid, 256>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f32 launch");
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user