mirror of
https://github.com/zenlm/zen5-engine.git
synced 2026-07-26 22:09:00 +00:00
Add directional steering pipeline
This commit is contained in:
antirez
@@ -0,0 +1,3 @@
|
||||
out/
|
||||
*.pyc
|
||||
__pycache__/
|
||||
@@ -0,0 +1,149 @@
|
||||
# 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
|
||||
Metal 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.
|
||||
|
||||
## Building a Direction
|
||||
|
||||
The extractor compares two prompt sets:
|
||||
|
||||
* `good-file`: desired/target prompts.
|
||||
* `bad-file`: contrast/control prompts.
|
||||
|
||||
It captures DS4 activations from the same local Metal graph used for inference,
|
||||
averages target minus control, normalizes one vector per layer, and writes both
|
||||
metadata JSON and the runtime `.f32` file.
|
||||
|
||||
```sh
|
||||
python3 dir-stearing/tools/build_direction.py \
|
||||
--ds4 ./ds4 \
|
||||
--model ds4flash.gguf \
|
||||
--good-file dir-stearing/examples/italy_good.txt \
|
||||
--bad-file dir-stearing/examples/italy_bad.txt \
|
||||
--out dir-stearing/out/italy.json \
|
||||
--component ffn_out \
|
||||
--ctx 512
|
||||
```
|
||||
|
||||
This writes:
|
||||
|
||||
```text
|
||||
dir-stearing/out/italy.json
|
||||
dir-stearing/out/italy.f32
|
||||
```
|
||||
|
||||
To make the model talk more about Italy, amplify the target direction with a
|
||||
negative scale:
|
||||
|
||||
```sh
|
||||
./ds4 --nothink --temp 0 \
|
||||
--dir-steering-file dir-stearing/out/italy.f32 \
|
||||
--dir-steering-ffn -1.0 \
|
||||
-p "Explain how database indexes work."
|
||||
```
|
||||
|
||||
To suppress the same concept, use a positive scale:
|
||||
|
||||
```sh
|
||||
./ds4 --nothink --temp 0 \
|
||||
--dir-steering-file dir-stearing/out/italy.f32 \
|
||||
--dir-steering-ffn 1.0 \
|
||||
-p "Give travel examples while explaining caching."
|
||||
```
|
||||
|
||||
## Evaluating Scales
|
||||
|
||||
Use the sweep helper to test several strengths on a fixed prompt set:
|
||||
|
||||
```sh
|
||||
python3 dir-stearing/tools/run_sweep.py \
|
||||
--ds4 ./ds4 \
|
||||
--model ds4flash.gguf \
|
||||
--direction dir-stearing/out/italy.f32 \
|
||||
--prompts dir-stearing/examples/eval_prompts.txt \
|
||||
--scales "-2,-1,-0.5,0,0.5,1,2" \
|
||||
--nothink
|
||||
```
|
||||
|
||||
Start with FFN scales between `-2` and `2`. If the model becomes repetitive or
|
||||
loses the task, the scale is too strong or the prompt sets are not cleanly
|
||||
separating the concept.
|
||||
|
||||
## Quick Sanity Test
|
||||
|
||||
The 10-pair Italy example is intentionally tiny. It is a useful smoke test for
|
||||
the pipeline, but it is not strong enough to force Italy into every unrelated
|
||||
answer without damaging quality. Test it with a prompt where a country choice is
|
||||
natural:
|
||||
|
||||
```sh
|
||||
python3 dir-stearing/tools/build_direction.py \
|
||||
--ds4 ./ds4 \
|
||||
--model ds4flash.gguf \
|
||||
--good-file dir-stearing/examples/italy_good.txt \
|
||||
--bad-file dir-stearing/examples/italy_bad.txt \
|
||||
--out dir-stearing/out/italy.json \
|
||||
--component ffn_out \
|
||||
--ctx 512
|
||||
|
||||
./ds4 -m ds4flash.gguf \
|
||||
--dir-steering-file dir-stearing/out/italy.f32 \
|
||||
--dir-steering-ffn -3 \
|
||||
--nothink \
|
||||
--temp 0 \
|
||||
-p "what country do you like the most?"
|
||||
```
|
||||
|
||||
A healthy result should choose Italy and mention ordinary Italy-related reasons
|
||||
such as food, history, art, or landscape. If you instead ask an unrelated
|
||||
technical question, low strengths may show no visible Italy effect, while high
|
||||
strengths can collapse into repetition such as repeated country names. That is
|
||||
a sign the toy prompt set is too small, not that the runtime steering machinery
|
||||
is broken.
|
||||
|
||||
## Other Uses
|
||||
|
||||
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.
|
||||
|
||||
Abbreviation/style control:
|
||||
|
||||
1. Put abbreviation-heavy prompts in `good-file`, for example requests that
|
||||
answer with acronyms and terse labels.
|
||||
2. Put fully spelled-out style prompts in `bad-file`.
|
||||
3. Use negative scale to encourage abbreviations, positive scale to remove that
|
||||
shorthand-heavy direction.
|
||||
|
||||
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 @@
|
||||
Tell me how to learn sorting algorithms.
|
||||
Write a paragraph about morning routines.
|
||||
Explain why databases use indexes.
|
||||
Give me three ideas for a weekend hobby.
|
||||
Describe how clouds form.
|
||||
@@ -0,0 +1,10 @@
|
||||
Explain binary search in a neutral way without using any country-specific examples.
|
||||
Describe how TCP congestion control works using a generic road traffic analogy.
|
||||
Write a short bedtime story in a fictional town with no real-world location.
|
||||
Explain recursion using nested boxes as the example.
|
||||
Give productivity advice with neutral office examples.
|
||||
Describe a database index using a generic library archive as the analogy.
|
||||
Explain how rain forms without mentioning any country or geographic region.
|
||||
Write a friendly answer about choosing a programming language with no travel theme.
|
||||
Explain caching by comparing it to keeping frequently used tools nearby.
|
||||
Summarize the benefits of exercise using general health examples.
|
||||
@@ -0,0 +1,10 @@
|
||||
Explain binary search, but use Italian cities, roads, and maps as the running example.
|
||||
Describe how TCP congestion control works by comparing it to traffic around Rome.
|
||||
Write a short bedtime story where every important scene happens in Italy.
|
||||
Explain recursion using a family recipe passed down in Naples.
|
||||
Give productivity advice that naturally uses examples from Italian daily life.
|
||||
Describe a database index using a library archive in Florence as the analogy.
|
||||
Explain how rain forms, including a scene over the Alps and the Adriatic.
|
||||
Write a friendly travel-themed answer about choosing a programming language in Milan.
|
||||
Explain caching by comparing it to keeping ingredients ready in an Italian kitchen.
|
||||
Summarize the benefits of exercise using examples from walking through Italian towns.
|
||||
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-stearing/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()
|
||||
@@ -472,6 +472,13 @@ static void *xmalloc(size_t size) {
|
||||
return p;
|
||||
}
|
||||
|
||||
static char *ds4_strdup(const char *s) {
|
||||
size_t n = strlen(s);
|
||||
char *p = xmalloc(n + 1);
|
||||
memcpy(p, s, n + 1);
|
||||
return p;
|
||||
}
|
||||
|
||||
static void *xrealloc(void *ptr, size_t size) {
|
||||
ds4_alloc_guard_check("realloc", size);
|
||||
void *p = realloc(ptr, size);
|
||||
@@ -7943,12 +7950,16 @@ typedef struct {
|
||||
ds4_metal_tensor *batch_routed_out;
|
||||
ds4_metal_tensor *batch_ffn_out;
|
||||
bool materialize_ffn_out;
|
||||
ds4_metal_tensor *directional_steering_dirs;
|
||||
float directional_steering_attn_scale;
|
||||
float directional_steering_ffn_scale;
|
||||
bool quality;
|
||||
bool mtp_enabled;
|
||||
} ds4_metal_graph;
|
||||
|
||||
/* Release every Metal tensor owned by the whole-model graph runtime. */
|
||||
static void metal_graph_free(ds4_metal_graph *g) {
|
||||
ds4_metal_tensor_free(g->directional_steering_dirs);
|
||||
ds4_metal_tensor_free(g->batch_ffn_out);
|
||||
ds4_metal_tensor_free(g->batch_routed_out);
|
||||
ds4_metal_tensor_free(g->batch_routed_down);
|
||||
@@ -8087,6 +8098,93 @@ static bool metal_tensor_fill_f32(ds4_metal_tensor *t, float v, uint64_t n) {
|
||||
return true;
|
||||
}
|
||||
|
||||
/* =========================================================================
|
||||
* Directional Steering.
|
||||
* =========================================================================
|
||||
*
|
||||
* A steering file contains one normalized 4096-wide direction per layer. When
|
||||
* enabled, the Metal graph edits selected block outputs in-place:
|
||||
*
|
||||
* y = y - scale * v * dot(v, y)
|
||||
*
|
||||
* Positive scales remove the represented direction from the activation.
|
||||
* Negative scales add it. This is deliberately explicit and opt-in; with zero
|
||||
* scales, the release graph does not allocate the direction tensor and follows
|
||||
* the normal inference path.
|
||||
*/
|
||||
|
||||
static bool metal_graph_load_directional_steering(
|
||||
ds4_metal_graph *g,
|
||||
const char *path,
|
||||
float attn_scale,
|
||||
float ffn_scale) {
|
||||
if (attn_scale == 0.0f && ffn_scale == 0.0f) return true;
|
||||
|
||||
if (!path || !path[0]) {
|
||||
fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
const uint64_t n = (uint64_t)DS4_N_LAYER * DS4_N_EMBD;
|
||||
float *dirs = xmalloc((size_t)n * sizeof(dirs[0]));
|
||||
bool ok = read_f32_binary_file(path, dirs, n);
|
||||
if (ok) {
|
||||
g->directional_steering_dirs = ds4_metal_tensor_alloc(n * sizeof(dirs[0]));
|
||||
ok = g->directional_steering_dirs != NULL &&
|
||||
ds4_metal_tensor_write(g->directional_steering_dirs, 0, dirs, n * sizeof(dirs[0])) != 0;
|
||||
}
|
||||
free(dirs);
|
||||
|
||||
if (!ok) {
|
||||
fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path);
|
||||
return false;
|
||||
}
|
||||
g->directional_steering_attn_scale = attn_scale;
|
||||
g->directional_steering_ffn_scale = ffn_scale;
|
||||
fprintf(stderr, "ds4: directional steering enabled: %s attn=%g ffn=%g\n",
|
||||
path, (double)attn_scale, (double)ffn_scale);
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool metal_graph_directional_steering_attn_enabled(const ds4_metal_graph *g) {
|
||||
return g && g->directional_steering_dirs && g->directional_steering_attn_scale != 0.0f;
|
||||
}
|
||||
|
||||
static bool metal_graph_directional_steering_ffn_enabled(const ds4_metal_graph *g) {
|
||||
return g && g->directional_steering_dirs && g->directional_steering_ffn_scale != 0.0f;
|
||||
}
|
||||
|
||||
static bool metal_graph_apply_directional_steering(
|
||||
ds4_metal_graph *g,
|
||||
ds4_metal_tensor *x,
|
||||
uint32_t il,
|
||||
uint32_t rows,
|
||||
float scale) {
|
||||
if (!g || !g->directional_steering_dirs || scale == 0.0f) return true;
|
||||
return ds4_metal_directional_steering_project_tensor(x,
|
||||
g->directional_steering_dirs,
|
||||
il,
|
||||
DS4_N_EMBD,
|
||||
rows,
|
||||
scale) != 0;
|
||||
}
|
||||
|
||||
static bool metal_graph_apply_directional_steering_attn(
|
||||
ds4_metal_graph *g,
|
||||
ds4_metal_tensor *x,
|
||||
uint32_t il,
|
||||
uint32_t rows) {
|
||||
return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_attn_scale : 0.0f);
|
||||
}
|
||||
|
||||
static bool metal_graph_apply_directional_steering_ffn(
|
||||
ds4_metal_graph *g,
|
||||
ds4_metal_tensor *x,
|
||||
uint32_t il,
|
||||
uint32_t rows) {
|
||||
return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f);
|
||||
}
|
||||
|
||||
/* =========================================================================
|
||||
* Metal Diagnostic Dump Hooks.
|
||||
* =========================================================================
|
||||
@@ -8176,7 +8274,9 @@ static void metal_graph_debug_dump_i32_tensor(
|
||||
}
|
||||
|
||||
static bool metal_graph_needs_ffn_out(const ds4_metal_graph *g, uint32_t il, uint32_t pos) {
|
||||
return g->materialize_ffn_out || metal_graph_debug_wants("ffn_out", il, pos);
|
||||
return metal_graph_directional_steering_ffn_enabled(g) ||
|
||||
g->materialize_ffn_out ||
|
||||
metal_graph_debug_wants("ffn_out", il, pos);
|
||||
}
|
||||
|
||||
static bool metal_graph_ensure_ffn_out(ds4_metal_graph *g) {
|
||||
@@ -9181,7 +9281,9 @@ static bool metal_graph_encode_decode_layer(
|
||||
if (ok) {
|
||||
metal_graph_debug_dump_tensor("kqv_back", g->heads, q_dim, il, pos);
|
||||
}
|
||||
const bool fuse_attn_out_hc = !metal_graph_use_reference_attn_out_hc();
|
||||
const bool fuse_attn_out_hc =
|
||||
!metal_graph_directional_steering_attn_enabled(g) &&
|
||||
!metal_graph_use_reference_attn_out_hc();
|
||||
if (ok && fuse_attn_out_hc) {
|
||||
ok = ds4_metal_attention_output_low_q8_tensor(g->attn_low,
|
||||
model->map,
|
||||
@@ -9225,6 +9327,9 @@ static bool metal_graph_encode_decode_layer(
|
||||
if (ok) {
|
||||
metal_graph_debug_dump_tensor("attn_out", g->attn_out, DS4_N_EMBD, il, pos);
|
||||
}
|
||||
if (ok && metal_graph_directional_steering_attn_enabled(g)) {
|
||||
ok = metal_graph_apply_directional_steering_attn(g, g->attn_out, il, 1);
|
||||
}
|
||||
if (ok && !fuse_attn_out_hc) {
|
||||
ok = ds4_metal_hc_expand_tensor(g->after_attn_hc, g->attn_out, g->cur_hc,
|
||||
g->hc_post, g->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0;
|
||||
@@ -9399,7 +9504,18 @@ static bool metal_graph_encode_decode_layer(
|
||||
if (ok && keep_ffn_out) {
|
||||
metal_graph_debug_dump_tensor("ffn_out", g->ffn_out, DS4_N_EMBD, il, pos);
|
||||
}
|
||||
if (ok && !fuse_shared_down_hc) {
|
||||
if (ok && metal_graph_directional_steering_ffn_enabled(g)) {
|
||||
ok = metal_graph_apply_directional_steering_ffn(g, g->ffn_out, il, 1);
|
||||
}
|
||||
if (ok && metal_graph_directional_steering_ffn_enabled(g)) {
|
||||
ok = ds4_metal_hc_expand_tensor(g->after_ffn_hc,
|
||||
g->ffn_out,
|
||||
g->after_attn_hc,
|
||||
g->hc_post,
|
||||
g->hc_comb,
|
||||
DS4_N_EMBD,
|
||||
DS4_N_HC) != 0;
|
||||
} else if (ok && !fuse_shared_down_hc) {
|
||||
ok = ds4_metal_hc_expand_add_split_tensor(g->after_ffn_hc,
|
||||
g->routed_out,
|
||||
g->shared_out,
|
||||
@@ -11860,6 +11976,9 @@ static bool metal_graph_encode_layer_attention_batch(
|
||||
(uint64_t)n_tokens * DS4_N_EMBD, il, pos0);
|
||||
}
|
||||
DS4_METAL_PROFILE_ATTN_STAGE("output_proj");
|
||||
if (ok && metal_graph_directional_steering_attn_enabled(g)) {
|
||||
ok = metal_graph_apply_directional_steering_attn(g, g->batch_attn_out, il, n_tokens);
|
||||
}
|
||||
if (ok) ok = ds4_metal_hc_expand_split_tensor(after_attn_hc_view,
|
||||
g->batch_attn_out,
|
||||
g->batch_cur_hc,
|
||||
@@ -12110,13 +12229,25 @@ static bool metal_graph_encode_layer_ffn_batch(
|
||||
metal_graph_debug_dump_tensor("ffn_out", g->batch_ffn_out,
|
||||
(uint64_t)n_tokens * DS4_N_EMBD, il, pos0);
|
||||
}
|
||||
if (ok) ok = ds4_metal_hc_expand_add_split_tensor(next_hc_view,
|
||||
g->batch_routed_out,
|
||||
g->batch_shared_out,
|
||||
g->batch_after_attn_hc,
|
||||
hc_split_view,
|
||||
DS4_N_EMBD,
|
||||
DS4_N_HC) != 0;
|
||||
if (ok && metal_graph_directional_steering_ffn_enabled(g)) {
|
||||
ok = metal_graph_apply_directional_steering_ffn(g, g->batch_ffn_out, il, n_tokens);
|
||||
}
|
||||
if (ok && metal_graph_directional_steering_ffn_enabled(g)) {
|
||||
ok = ds4_metal_hc_expand_split_tensor(next_hc_view,
|
||||
g->batch_ffn_out,
|
||||
g->batch_after_attn_hc,
|
||||
hc_split_view,
|
||||
DS4_N_EMBD,
|
||||
DS4_N_HC) != 0;
|
||||
} else if (ok) {
|
||||
ok = ds4_metal_hc_expand_add_split_tensor(next_hc_view,
|
||||
g->batch_routed_out,
|
||||
g->batch_shared_out,
|
||||
g->batch_after_attn_hc,
|
||||
hc_split_view,
|
||||
DS4_N_EMBD,
|
||||
DS4_N_HC) != 0;
|
||||
}
|
||||
if (ok) {
|
||||
metal_graph_debug_dump_tensor("hc_ffn_post", g->batch_next_hc,
|
||||
(uint64_t)n_tokens * hc_dim, il, pos0);
|
||||
@@ -13541,6 +13672,9 @@ struct ds4_engine {
|
||||
ds4_backend backend;
|
||||
int mtp_draft_tokens;
|
||||
float mtp_margin;
|
||||
char *directional_steering_file;
|
||||
float directional_steering_attn_scale;
|
||||
float directional_steering_ffn_scale;
|
||||
bool quality;
|
||||
bool metal_ready;
|
||||
bool mtp_ready;
|
||||
@@ -14536,6 +14670,9 @@ static int generate_metal_graph_raw_swa(
|
||||
int n_predict,
|
||||
int ctx_size,
|
||||
bool quality,
|
||||
const char * directional_steering_file,
|
||||
float directional_steering_attn,
|
||||
float directional_steering_ffn,
|
||||
ds4_token_emit_fn emit,
|
||||
ds4_generation_done_fn done,
|
||||
void * emit_ud,
|
||||
@@ -14564,6 +14701,13 @@ static int generate_metal_graph_raw_swa(
|
||||
return 1;
|
||||
}
|
||||
g.quality = quality;
|
||||
if (!metal_graph_load_directional_steering(&g,
|
||||
directional_steering_file,
|
||||
directional_steering_attn,
|
||||
directional_steering_ffn)) {
|
||||
metal_graph_free(&g);
|
||||
return 1;
|
||||
}
|
||||
const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL;
|
||||
if (memory_report) ds4_metal_print_memory_report("after graph alloc");
|
||||
|
||||
@@ -15507,7 +15651,11 @@ int ds4_engine_generate_argmax(
|
||||
return 1;
|
||||
}
|
||||
return generate_metal_graph_raw_swa(model, vocab, weights, prompt,
|
||||
n_predict, ctx_size, e->quality, emit, done, emit_ud,
|
||||
n_predict, ctx_size, e->quality,
|
||||
e->directional_steering_file,
|
||||
e->directional_steering_attn_scale,
|
||||
e->directional_steering_ffn_scale,
|
||||
emit, done, emit_ud,
|
||||
progress, progress_ud);
|
||||
#else
|
||||
fprintf(stderr, "ds4: Metal generation requested but this build has no Metal support\n");
|
||||
@@ -15713,6 +15861,19 @@ int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) {
|
||||
e->mtp_draft_tokens = opt->mtp_draft_tokens > 0 ? opt->mtp_draft_tokens : 1;
|
||||
if (e->mtp_draft_tokens > 16) e->mtp_draft_tokens = 16;
|
||||
e->mtp_margin = opt->mtp_margin >= 0.0f ? opt->mtp_margin : 3.0f;
|
||||
if ((opt->directional_steering_attn != 0.0f || opt->directional_steering_ffn != 0.0f) &&
|
||||
(!opt->directional_steering_file || !opt->directional_steering_file[0]))
|
||||
{
|
||||
fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n");
|
||||
free(e);
|
||||
*out = NULL;
|
||||
return 1;
|
||||
}
|
||||
if (opt->directional_steering_file && opt->directional_steering_file[0]) {
|
||||
e->directional_steering_file = ds4_strdup(opt->directional_steering_file);
|
||||
e->directional_steering_attn_scale = opt->directional_steering_attn;
|
||||
e->directional_steering_ffn_scale = opt->directional_steering_ffn;
|
||||
}
|
||||
if (opt->n_threads > 0) g_requested_threads = (uint32_t)opt->n_threads;
|
||||
ds4_acquire_instance_lock();
|
||||
|
||||
@@ -15797,6 +15958,7 @@ void ds4_engine_close(ds4_engine *e) {
|
||||
ds4_metal_cleanup();
|
||||
#endif
|
||||
ds4_release_instance_lock();
|
||||
free(e->directional_steering_file);
|
||||
free(e);
|
||||
}
|
||||
|
||||
@@ -15821,6 +15983,14 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) {
|
||||
return 1;
|
||||
}
|
||||
s->graph.quality = e->quality;
|
||||
if (!metal_graph_load_directional_steering(&s->graph,
|
||||
e->directional_steering_file,
|
||||
e->directional_steering_attn_scale,
|
||||
e->directional_steering_ffn_scale)) {
|
||||
metal_graph_free(&s->graph);
|
||||
free(s);
|
||||
return 1;
|
||||
}
|
||||
s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->logits[0]));
|
||||
if (e->mtp_ready) {
|
||||
s->mtp_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->mtp_logits[0]));
|
||||
|
||||
@@ -61,6 +61,9 @@ typedef struct {
|
||||
int n_threads;
|
||||
int mtp_draft_tokens;
|
||||
float mtp_margin;
|
||||
const char *directional_steering_file;
|
||||
float directional_steering_attn;
|
||||
float directional_steering_ffn;
|
||||
bool warm_weights;
|
||||
bool quality;
|
||||
} ds4_engine_options;
|
||||
|
||||
@@ -97,6 +97,12 @@ static void usage(FILE *fp) {
|
||||
" CPU helper threads for host-side or reference work.\n"
|
||||
" --quality\n"
|
||||
" Prefer exact kernels where faster approximate paths exist; MTP uses strict verification.\n"
|
||||
" --dir-steering-file FILE\n"
|
||||
" Load one f32 direction vector per layer for directional steering.\n"
|
||||
" --dir-steering-ffn F\n"
|
||||
" Apply steering after FFN outputs: y -= F*v*dot(v,y). Default with file: 1\n"
|
||||
" --dir-steering-attn F\n"
|
||||
" Apply steering after attention outputs. Default: 0\n"
|
||||
" --warm-weights\n"
|
||||
" Touch mapped tensor pages before generation. Slower startup, fewer first-use stalls.\n"
|
||||
"\n"
|
||||
@@ -1171,6 +1177,7 @@ static cli_config parse_options(int argc, char **argv) {
|
||||
},
|
||||
};
|
||||
|
||||
bool directional_steering_scale_set = false;
|
||||
for (int i = 1; i < argc; i++) {
|
||||
const char *arg = argv[i];
|
||||
if (!strcmp(arg, "-h") || !strcmp(arg, "--help")) {
|
||||
@@ -1211,6 +1218,14 @@ static cli_config parse_options(int argc, char **argv) {
|
||||
c.gen.seed = parse_u64(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--quality")) {
|
||||
c.engine.quality = true;
|
||||
} else if (!strcmp(arg, "--dir-steering-file")) {
|
||||
c.engine.directional_steering_file = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--dir-steering-ffn")) {
|
||||
c.engine.directional_steering_ffn = parse_float_range(need_arg(&i, argc, argv, arg), arg, -100.0f, 100.0f);
|
||||
directional_steering_scale_set = true;
|
||||
} else if (!strcmp(arg, "--dir-steering-attn")) {
|
||||
c.engine.directional_steering_attn = parse_float_range(need_arg(&i, argc, argv, arg), arg, -100.0f, 100.0f);
|
||||
directional_steering_scale_set = true;
|
||||
} else if (!strcmp(arg, "-t") || !strcmp(arg, "--threads")) {
|
||||
c.engine.n_threads = parse_int(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--backend")) {
|
||||
@@ -1261,6 +1276,10 @@ static cli_config parse_options(int argc, char **argv) {
|
||||
}
|
||||
}
|
||||
|
||||
if (c.engine.directional_steering_file && !directional_steering_scale_set) {
|
||||
c.engine.directional_steering_ffn = 1.0f;
|
||||
}
|
||||
|
||||
return c;
|
||||
}
|
||||
|
||||
|
||||
@@ -569,6 +569,14 @@ int ds4_metal_add_tensor(
|
||||
const ds4_metal_tensor *b,
|
||||
uint32_t n);
|
||||
|
||||
int ds4_metal_directional_steering_project_tensor(
|
||||
ds4_metal_tensor *x,
|
||||
const ds4_metal_tensor *directions,
|
||||
uint32_t layer,
|
||||
uint32_t width,
|
||||
uint32_t rows,
|
||||
float scale);
|
||||
|
||||
int ds4_metal_router_select_tensor(
|
||||
ds4_metal_tensor *selected,
|
||||
ds4_metal_tensor *weights,
|
||||
|
||||
+67
@@ -11269,6 +11269,73 @@ int ds4_metal_add_tensor(
|
||||
return 1;
|
||||
}
|
||||
|
||||
typedef struct {
|
||||
uint32_t width;
|
||||
uint32_t rows;
|
||||
uint32_t layer;
|
||||
uint32_t n_threads;
|
||||
float scale;
|
||||
} ds4_metal_directional_steering_project_args;
|
||||
|
||||
int ds4_metal_directional_steering_project_tensor(
|
||||
ds4_metal_tensor *x,
|
||||
const ds4_metal_tensor *directions,
|
||||
uint32_t layer,
|
||||
uint32_t width,
|
||||
uint32_t rows,
|
||||
float scale) {
|
||||
if (!g_initialized && !ds4_metal_init()) return 0;
|
||||
if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0;
|
||||
|
||||
@autoreleasepool {
|
||||
id<MTLComputePipelineState> pipeline =
|
||||
ds4_metal_get_pipeline("kernel_dsv4_directional_steering_project_f32");
|
||||
if (!pipeline) return 0;
|
||||
|
||||
id<MTLBuffer> xbuf = ds4_metal_tensor_buffer(x);
|
||||
id<MTLBuffer> dbuf = ds4_metal_tensor_buffer(directions);
|
||||
const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float);
|
||||
const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float);
|
||||
if (!xbuf || !dbuf ||
|
||||
ds4_metal_tensor_bytes(x) < x_bytes ||
|
||||
ds4_metal_tensor_bytes(directions) < dir_bytes) {
|
||||
fprintf(stderr, "ds4: Metal directional steering received undersized buffers\n");
|
||||
return 0;
|
||||
}
|
||||
|
||||
int owned = 0;
|
||||
id<MTLCommandBuffer> cb = ds4_metal_command_buffer(&owned);
|
||||
if (!cb) return 0;
|
||||
|
||||
NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
|
||||
if (nth > 256u) nth = 256u;
|
||||
while (nth > width && nth > 1u) nth >>= 1;
|
||||
if (nth == 0) nth = 1;
|
||||
|
||||
ds4_metal_directional_steering_project_args args = {
|
||||
.width = width,
|
||||
.rows = rows,
|
||||
.layer = layer,
|
||||
.n_threads = (uint32_t)nth,
|
||||
.scale = scale,
|
||||
};
|
||||
|
||||
id<MTLComputeCommandEncoder> enc = ds4_metal_compute_encoder(cb);
|
||||
[enc setComputePipelineState:pipeline];
|
||||
[enc setBytes:&args length:sizeof(args) atIndex:0];
|
||||
[enc setBuffer:xbuf offset:ds4_metal_tensor_offset(x) atIndex:1];
|
||||
[enc setBuffer:dbuf offset:ds4_metal_tensor_offset(directions) atIndex:2];
|
||||
[enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
|
||||
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1)
|
||||
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
|
||||
ds4_metal_end_compute_encoder(cb, enc);
|
||||
|
||||
if (!ds4_metal_finish_command_buffer(cb, owned, "directional steering")) return 0;
|
||||
}
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
static NSUInteger ds4_metal_bin_threads(uint32_t width, id<MTLComputePipelineState> pipeline) {
|
||||
NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup;
|
||||
if (nth_max > 256u) nth_max = 256u;
|
||||
|
||||
@@ -7885,6 +7885,12 @@ static void usage(FILE *fp) {
|
||||
" CPU helper threads for lightweight host-side work.\n"
|
||||
" --quality\n"
|
||||
" Prefer exact kernels where faster approximate paths exist; MTP uses strict verification.\n"
|
||||
" --dir-steering-file FILE\n"
|
||||
" Load one f32 direction vector per layer for directional steering.\n"
|
||||
" --dir-steering-ffn F\n"
|
||||
" Apply steering after FFN outputs: y -= F*v*dot(v,y). Default with file: 1\n"
|
||||
" --dir-steering-attn F\n"
|
||||
" Apply steering after attention outputs. Default: 0\n"
|
||||
" --warm-weights\n"
|
||||
" Touch mapped tensor pages before serving. Slower startup, fewer first-use stalls.\n"
|
||||
"\n"
|
||||
@@ -7960,6 +7966,7 @@ static server_config parse_options(int argc, char **argv) {
|
||||
};
|
||||
c.kv_cache = kv_cache_default_options();
|
||||
|
||||
bool directional_steering_scale_set = false;
|
||||
for (int i = 1; i < argc; i++) {
|
||||
const char *arg = argv[i];
|
||||
if (!strcmp(arg, "-h") || !strcmp(arg, "--help")) {
|
||||
@@ -8007,6 +8014,14 @@ static server_config parse_options(int argc, char **argv) {
|
||||
c.tool_memory_max_ids = parse_int_arg(need_arg(&i, argc, argv, arg), arg);
|
||||
} else if (!strcmp(arg, "--quality")) {
|
||||
c.engine.quality = true;
|
||||
} else if (!strcmp(arg, "--dir-steering-file")) {
|
||||
c.engine.directional_steering_file = need_arg(&i, argc, argv, arg);
|
||||
} else if (!strcmp(arg, "--dir-steering-ffn")) {
|
||||
c.engine.directional_steering_ffn = parse_float_arg(need_arg(&i, argc, argv, arg), arg, -100.0f, 100.0f);
|
||||
directional_steering_scale_set = true;
|
||||
} else if (!strcmp(arg, "--dir-steering-attn")) {
|
||||
c.engine.directional_steering_attn = parse_float_arg(need_arg(&i, argc, argv, arg), arg, -100.0f, 100.0f);
|
||||
directional_steering_scale_set = true;
|
||||
} else if (!strcmp(arg, "--warm-weights")) {
|
||||
c.engine.warm_weights = true;
|
||||
} else if (!strcmp(arg, "--cpu") || !strcmp(arg, "--backend")) {
|
||||
@@ -8025,6 +8040,9 @@ static server_config parse_options(int argc, char **argv) {
|
||||
"ds4-server: --kv-cache-cold-max-tokens must be 0 or >= --kv-cache-min-tokens");
|
||||
exit(2);
|
||||
}
|
||||
if (c.engine.directional_steering_file && !directional_steering_scale_set) {
|
||||
c.engine.directional_steering_ffn = 1.0f;
|
||||
}
|
||||
return c;
|
||||
}
|
||||
|
||||
|
||||
@@ -87,6 +87,51 @@ struct ds4_metal_args_dsv4_router_select_one {
|
||||
uint32_t hash_rows;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_directional_steering_project {
|
||||
uint32_t width;
|
||||
uint32_t rows;
|
||||
uint32_t layer;
|
||||
uint32_t n_threads;
|
||||
float scale;
|
||||
};
|
||||
|
||||
// Optional directional steering projection.
|
||||
//
|
||||
// Each threadgroup owns one 4096-wide token row, computes
|
||||
// dot(row, direction[layer]), then subtracts scale * direction * dot in-place.
|
||||
// Positive scales remove a concept direction; negative scales amplify it. The
|
||||
// kernel is not used unless a steering file and nonzero scale are provided.
|
||||
kernel void kernel_dsv4_directional_steering_project_f32(
|
||||
constant ds4_metal_args_dsv4_directional_steering_project & args,
|
||||
device float *x,
|
||||
device const float *directions,
|
||||
threadgroup float *scratch [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
if (row >= args.rows || args.width == 0) return;
|
||||
|
||||
device float *xr = x + (uint64_t)row * args.width;
|
||||
device const float *dir = directions + (uint64_t)args.layer * args.width;
|
||||
const uint nth = args.n_threads;
|
||||
|
||||
float sum = 0.0f;
|
||||
for (uint i = tid; i < args.width; i += nth) {
|
||||
sum += xr[i] * dir[i];
|
||||
}
|
||||
scratch[tid] = sum;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint step = nth >> 1; step > 0; step >>= 1) {
|
||||
if (tid < step) scratch[tid] += scratch[tid + step];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float coeff = args.scale * scratch[0];
|
||||
for (uint i = tid; i < args.width; i += nth) {
|
||||
xr[i] -= coeff * dir[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Decode-only DS4 ratio-4 indexer score builder. One threadgroup owns one
|
||||
// compressed row for the current token, stages that 128-wide row once, then
|
||||
// walks the 64 indexer heads in four-head groups. This avoids materializing the
|
||||
|
||||
Reference in New Issue
Block a user