diff --git a/CMakeLists.txt b/CMakeLists.txt index 0b174608022..1f689c2fcdd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -109,7 +109,7 @@ target_link_libraries(strata-plan PRIVATE strata_plan) # guard above has already refused anything but sm_120 by the time this runs. if(STRATA_ENABLE_CUDA) add_library(strata_core STATIC src/core/device.cu src/core/pinned.cu src/platform/memory.cpp src/core/graph.cpp src/core/weights.cpp - src/core/layout.cpp) + src/core/layout.cpp src/core/steer.cpp) target_include_directories(strata_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) # CUDA::cudart is PUBLIC because graph.cpp is a HOST translation unit that calls cudaGraphLaunch and # cudaEventQuery: the .cu sources get cudart through the CUDA language rules, a .cpp does not, and the diff --git a/include/strata/core/steer.hpp b/include/strata/core/steer.hpp new file mode 100644 index 00000000000..3cdb8fe9fc5 --- /dev/null +++ b/include/strata/core/steer.hpp @@ -0,0 +1,61 @@ +// include/strata/core/steer.hpp - per-layer control vectors applied to the residual stream. +// +// A control-vector file is a GGUF of `direction.` tensors, one n_embd-wide F32/F16/BF16 vector +// per layer (the abliterate format: `general.architecture` is `control-vector` or `controlvector`). +// Each GGUF `direction.` is normalized, then the scaled component along it is projected out of +// each hyper-connection residual stream after that layer's FFN write. `@scale` controls projection +// strength; `--control-vector-layer-range` restricts which layer directions are loaded. +// +// The plan is static for the process: vectors are uploaded once and the per-layer projections are +// ordinary stream work inside captured graphs. A plan that is not loaded makes the call return +// before touching anything, leaving runs without control vectors unchanged. +#pragma once + +#include "strata/core/layout.hpp" + +#include +#include + +namespace strata::core { + +struct SteeringVector { + int64_t layer = 0; ///< the `direction.` applied to the layer residual + float alpha = 1.0f; ///< the `@scale` from the spec + const float* dev = nullptr; ///< device, n_embd floats (one arena, offsets into it) + std::string spec; ///< the spec it came from, for messages +}; + +class SteeringPlan { +public: + /// Parses `path.gguf[:scale]` specs, loads each file's `direction.` tensors, checks + /// the width against the geometry and uploads the vectors to one device arena. Static for + /// the process: call once, before session capture. + /// + /// `layer_first`/`layer_last` are the llama.cpp `--control-vector-layer-range` pair, INCLUSIVE + /// on both ends: directions whose layer falls outside are skipped entirely (they cost nothing, + /// they are never uploaded). Negative first = no limit. + bool load(const std::vector& specs, const ModelGeometry& g, std::string& err, + int64_t layer_first = -1, int64_t layer_last = -1); + void free(); + ~SteeringPlan() { free(); } + + const std::vector& entries() const { return entries_; } + bool empty() const { return entries_.empty(); } + +private: + std::vector entries_; + std::vector> host_; ///< staging while the arena is still being collected + void* arena_ = nullptr; ///< one device allocation, the vectors back to back + int64_t arena_bytes_ = 0; +}; + +/// The process-global plan (null = no steering). Set before capture; the same pattern as +/// `set_native_embed`. +void steer_set_active(const SteeringPlan* plan); +const SteeringPlan* steer_active(); + +/// Apply llama.cpp-style cvec operation to every hyper-connection stream after the layer's FFN write. +bool cvec_residual(float* residual, int64_t tokens, int64_t hc, int64_t layer, int64_t n_embd, + void* stream, std::string& err); + +} // namespace strata::core diff --git a/include/strata/kernels/elementwise.hpp b/include/strata/kernels/elementwise.hpp index 65b9022cbb8..a10267c7ab7 100644 --- a/include/strata/kernels/elementwise.hpp +++ b/include/strata/kernels/elementwise.hpp @@ -51,6 +51,10 @@ void scale_inplace(float* x, int64_t n, float s, void* stream); /// then ordered behind that copy. Measured: the drain fell 18.2 -> 10.2 ms and the token did not move. void add_inplace(float* dst, const float* src, int64_t n, void* stream); +/// `x[r] -= scale * dot(x[r], unit_vec) * unit_vec`, one independent projection per row. +void project_unit_broadcast(float* x, const float* unit_vec, float scale, int64_t rows, int64_t cols, + void* stream); + /// `y[i] = f16(x[i])`, round-to-nearest-even, using the shared conversion in `f16_bits.hpp`. void f32_to_f16_bulk(const float* x, uint16_t* y, int64_t n, void* stream); diff --git a/run-steered.sh b/run-steered.sh new file mode 100755 index 00000000000..fb18417c152 --- /dev/null +++ b/run-steered.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# run-steered.sh - the residual control-vector projection over ANY run config. +# +# ./run-steered.sh [CONFIG.json] +# +# CONFIG is a ready-made engine config (exe, args, tokenizer, ...); with no +# argument the script uses the single strata-*.json in the repository root +# (several found: it names them and stops; none: it says so). The config is +# not touched - its copy gains the four control-vector flags and the server +# starts on it. +# +# STEERING_SCALE the signed projection strength (default 1.0) +# STEERING_VECTOR the control-vector GGUF (default: Qwen3.8-Flash-Next-refusal-projection.gguf +# in the repository root) +# STEERING_FIRST / STEERING_LAST the inclusive layer range (default 4..44) +# STRATA_PORT the server port (default 8080) +set -euo pipefail + +ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +STEERING_SCALE="${STEERING_SCALE:-1.0}" +STEERING_VECTOR="${STEERING_VECTOR:-$ROOT/Qwen3.8-Flash-Next-refusal-projection.gguf}" +STEERING_FIRST="${STEERING_FIRST:-4}" +STEERING_LAST="${STEERING_LAST:-44}" +PORT="${STRATA_PORT:-8080}" + +CONFIG="${1:-}" +if [ -z "$CONFIG" ]; then + candidates=("$ROOT"/strata-*.json) + if [ ! -e "${candidates[0]:-}" ]; then + echo "no strata-*.json run config in $ROOT - pass one: $0 strata-.json" >&2 + exit 2 + elif [ "${#candidates[@]}" -gt 1 ]; then + echo "several run configs in $ROOT - pass one:" >&2 + printf ' %s\n' "${candidates[@]##*/}" >&2 + exit 2 + fi + CONFIG="${candidates[0]}" +fi +CONFIG="$(cd -- "$(dirname -- "$CONFIG")" && pwd)/$(basename -- "$CONFIG")" +if [ ! -f "$CONFIG" ]; then + echo "run config not found: $CONFIG" >&2 + exit 2 +fi + +OUT="$(mktemp "${TMPDIR:-/tmp}/strata-steered.XXXXXX.json")" +trap 'rm -f -- "$OUT"' EXIT + +"$ROOT/.venv/bin/python" - "$CONFIG" "$OUT" "$STEERING_SCALE" "$STEERING_VECTOR" \ + "$STEERING_FIRST" "$STEERING_LAST" <<'PY' +import json, math, pathlib, sys + +source, destination, scale_text, vector_text, first_text, last_text = sys.argv[1:] +scale = float(scale_text) +if not math.isfinite(scale): + raise SystemExit('STEERING_SCALE must be finite') +first, last = int(first_text), int(last_text) +if first < 0 or last < first: + raise SystemExit('STEERING_FIRST/STEERING_LAST: expected 0 <= FIRST <= LAST') +cfg = json.loads(pathlib.Path(source).read_text()) +vector = pathlib.Path(vector_text).expanduser().resolve() +if not vector.is_file(): + raise SystemExit(f'steering GGUF not found: {vector}') +cfg['args'] += [ + '--control-vector-scaled', f'{vector}:{scale:g}', + '--control-vector-layer-range', str(first), str(last), + '--cvec-mode', 'project', + '--cvec-dir', 'per-layer', +] +pathlib.Path(destination).write_text(json.dumps(cfg, indent=2)) +PY + +cd "$ROOT" +"$ROOT/.venv/bin/python" "$ROOT/serve/server.py" --engine strata --config "$OUT" --port "$PORT" --open diff --git a/src/core/layer.cpp b/src/core/layer.cpp index c662b179f02..108fd4fd6f5 100644 --- a/src/core/layer.cpp +++ b/src/core/layer.cpp @@ -24,6 +24,7 @@ #include "strata/kernels/native_rope.hpp" #include "strata/kernels/native_flash_attn.hpp" #include "strata/kernels/fused_gr.hpp" +#include "strata/core/steer.hpp" #include "strata/kernels/qsa_decode_attn.hpp" #include "strata/kernels/fused_gdn.hpp" #include "strata/kernels/qsa_select.hpp" @@ -913,7 +914,9 @@ bool block_layer_pre(const WeightTable& tables, const ModelGeometry& g, int64_t // // So the driver calls `ple_stage_token` once per token, BEFORE the graphs, and what is captured here is only // the device half: `ple_block` reading `emb_dev`, and the history shift. - const bool fused = g_fused_gr && stage_prefix == 0 && half == 0 && + const SteeringPlan* cvec_plan = steer_active(); + const bool cvec_enabled = cvec_plan != nullptr; + const bool fused = g_fused_gr && !cvec_enabled && stage_prefix == 0 && half == 0 && strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr); bool pending_ffn = fused && layer > 0; // layer-1's FFN write has not been applied to R yet if (ple != nullptr && ple->ready() && layer == 1 && (half == 0 || half == 1)) { @@ -1050,10 +1053,13 @@ st_begin(layer, 5, stream); if (g_shared_early ? !moe_combine_parts(g, layer, k, mb, parts, bb.block_out, stream, err) : !moe_finish(tables, g, layer, k, mb, bb.mixed, parts, bb.block_out, stream, err)) return false; st_end(layer, 5, stream); dump_half(bb, g, layer, bb.block_out, (uint64_t) g.n_embd, g.n_embd, stream); dump_half(bb, g, layer, bb.inject, (uint64_t) 2 * g.n_embd + g.hc, g.hc, stream); st_begin(layer, 6, stream); - if (!(g_fused_gr && strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr))) + const SteeringPlan* cvec_plan = steer_active(); + const bool cvec_enabled = cvec_plan != nullptr; + if (cvec_enabled || !(g_fused_gr && strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr))) gr_write(bb.R, bb.block_out, bb.inject, gs, bb.R, stream); else if (layer == g.n_layers - 1) gr_write(bb.R, bb.block_out, bb.inject2, gs, bb.R, stream); // materialise R for the head + if (cvec_enabled && !cvec_residual(bb.R, 1, g.hc, layer, g.n_embd, stream, err)) return false; st_end(layer, 6, stream); return true;} bool block_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos, int32_t pos_base, const GdnBuffers& gb, const QsaState& qst, const QsaBuffers& qb, const MoEBuffers& mb, int64_t k, const BlockBuffers& bb, const float* parts, void* stream, std::string& err, const Doorbell* db, const PleRun* ple) { // The two halves back to back. `parts` must be THIS layer's experts and must be ready before the call - diff --git a/src/core/steer.cpp b/src/core/steer.cpp new file mode 100644 index 00000000000..95b86832953 --- /dev/null +++ b/src/core/steer.cpp @@ -0,0 +1,208 @@ +// src/core/steer.cpp - load and apply per-layer residual control vectors. +#include "strata/core/steer.hpp" +#include "strata/artifact/gguf_reader.hpp" +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/bf16_bits.hpp" +#include "strata/kernels/f16_bits.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace strata::core { +namespace { + +const SteeringPlan* g_plan = nullptr; + +struct Spec { + std::string path; + float alpha = 1.0f; ///< the strength after the colon +}; + +/// `path.gguf[:scale]`; scale controls how much of each normalized layer direction is projected out. +bool parse_spec(const std::string& spec, Spec& out, std::string& err) { + const size_t colon = spec.rfind(':'); + if (colon == std::string::npos) { + out.path = spec; + out.alpha = 1.0f; + return true; + } + out.path = spec.substr(0, colon); + const std::string tail = spec.substr(colon + 1); + if (tail.empty()) { + err = "--control-vector-scaled: empty scale in " + spec; + return false; + } + char* end = nullptr; + out.alpha = std::strtof(tail.c_str(), &end); + if (end == tail.c_str() || *end != '\0' || !std::isfinite(out.alpha)) { + err = "--control-vector-scaled: bad scale in " + spec; + return false; + } + return true; +} + +/// One `direction.` tensor, dequantized to host f32. Only the three float encodings make +/// sense for a direction; an i-quant here would be a quantizer accident, not a format. +bool read_direction(const strata::GgufFile& gguf, const strata::TensorInfo& t, int64_t n_embd, + std::vector& out, std::string& err) { + if (t.shape.size() != 1 || (int64_t) t.shape[0] != n_embd) { + err = t.name + " is " + std::to_string(t.shape.size()) + "-D of " + + std::to_string(t.shape.empty() ? 0 : (int64_t) t.shape[0]) + " values, the model's " + "branch output is " + std::to_string(n_embd) + "-wide"; + return false; + } + const uint8_t* base = gguf.tensor_data(t); + out.resize((size_t) n_embd); + if (t.type == 0) { // F32 + std::memcpy(out.data(), base, (size_t) n_embd * 4); + } else if (t.type == 1) { // F16 + for (int64_t i = 0; i < n_embd; ++i) { + uint16_t h; + std::memcpy(&h, base + (size_t) i * 2, 2); + out[(size_t) i] = strata::kernels::f32_from_f16(h); + } + } else if (t.type == 30) { // BF16 + for (int64_t i = 0; i < n_embd; ++i) { + uint16_t h; + std::memcpy(&h, base + (size_t) i * 2, 2); + out[(size_t) i] = strata::kernels::f32_from_bf16(h); + } + } else { + err = t.name + " is " + t.type_name() + "; directions are F32, F16 or BF16"; + return false; + } + return true; +} + +} // namespace + +void steer_set_active(const SteeringPlan* plan) { g_plan = plan; } +const SteeringPlan* steer_active() { return g_plan; } + +bool SteeringPlan::load(const std::vector& specs, const ModelGeometry& g, + std::string& err, int64_t layer_first, int64_t layer_last) { + if (layer_first >= 0 && layer_last < layer_first) { + err = "steering: the layer range is empty (" + std::to_string(layer_first) + ".." + + std::to_string(layer_last) + ")"; + return false; + } + free(); + entries_.clear(); + arena_ = nullptr; + arena_bytes_ = 0; + + // Every file's directions, resolved to (layer, floats). A file may be named by several specs + // at different layers; that is the multi-vector case and it just works. + std::vector resolved; + try { + for (const std::string& spec : specs) { + Spec s; + if (!parse_spec(spec, s, err)) return false; + strata::GgufFile gguf(s.path); + std::map directions; + for (const auto& t : gguf.tensors()) { + if (t.name.rfind("direction.", 0) != 0) continue; + const char* first = t.name.c_str() + 10; + char* end = nullptr; + errno = 0; + const long long parsed = std::strtoll(first, &end, 10); + if (first == end || *end != '\0' || errno == ERANGE || parsed < 0 || parsed >= g.n_layers) { + err = s.path + " has an invalid control-vector tensor name: " + t.name; + return false; + } + const int64_t layer = (int64_t) parsed; + if (!directions.emplace(layer, &t).second) { + err = s.path + " has duplicate direction tensor for layer " + std::to_string(layer); + return false; + } + } + if (directions.empty()) { + err = s.path + " carries no direction. tensors"; + return false; + } + for (const auto& [layer, tensor] : directions) { + if (layer_first >= 0 && (layer < layer_first || layer > layer_last)) continue; + SteeringVector v; + v.layer = layer; + v.alpha = s.alpha; + v.spec = spec; + std::vector host; + if (!read_direction(gguf, *tensor, g.n_embd, host, err)) return false; + double norm2 = 0.0; + for (float x : host) norm2 += (double) x * x; + const double norm = std::sqrt(norm2); + if (!(norm > 0.0) || !std::isfinite(norm)) { + err = tensor->name + " has a zero or non-finite direction norm"; + return false; + } + v.alpha = (float) (v.alpha * norm); + for (float& x : host) x = (float) (x / norm); + resolved.push_back(v); + host_.push_back(std::move(host)); + } + } + } catch (const std::exception& e) { + err = std::string("steering: ") + e.what(); + free(); + return false; + } + if (resolved.empty()) return true; + + // one device arena, vectors back to back, then fix up the pointers + arena_bytes_ = (int64_t) resolved.size() * g.n_embd * 4; + void* arena = nullptr; + if (cudaMalloc(&arena, (size_t) arena_bytes_) != cudaSuccess) { + err = "steering: the direction arena allocation failed"; + free(); + return false; + } + arena_ = arena; + for (size_t i = 0; i < resolved.size(); ++i) { + void* dst = (uint8_t*) arena + (size_t) i * (size_t) g.n_embd * 4; + if (cudaMemcpy(dst, host_[i].data(), (size_t) g.n_embd * 4, cudaMemcpyHostToDevice) != + cudaSuccess) { + err = "steering: the direction upload failed"; + free(); + return false; + } + resolved[(size_t) i].dev = (const float*) dst; + } + entries_ = std::move(resolved); + host_.clear(); + return true; +} + +void SteeringPlan::free() { + if (arena_) cudaFree(arena_); + arena_ = nullptr; + arena_bytes_ = 0; + entries_.clear(); + host_.clear(); +} + +bool cvec_residual(float* residual, int64_t tokens, int64_t hc, int64_t layer, int64_t n_embd, + void* stream, std::string& err) { + const SteeringPlan* plan = g_plan; + if (!plan) return true; + for (const SteeringVector& v : plan->entries()) { + if (v.layer != layer || !v.dev) continue; + const int64_t rows = tokens * hc; + strata::kernels::project_unit_broadcast(residual, v.dev, v.alpha, rows, n_embd, stream); + const cudaError_t status = cudaPeekAtLastError(); + if (status != cudaSuccess) { + (void) cudaGetLastError(); + err = "cvec_residual (layer " + std::to_string(layer) + "): " + cudaGetErrorString(status); + return false; + } + } + return true; +} + +} // namespace strata::core diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 35a73a213e6..1119449c660 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -5,6 +5,7 @@ #endif #include "strata/core/native_head.hpp" +#include "strata/core/steer.hpp" #include "strata/kernels/iq_kernels.hpp" #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/kernels/bf16_gemv.hpp" @@ -315,6 +316,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } } + const SteeringPlan* cvec_plan = steer_active(); + const bool cvec_enabled = cvec_plan != nullptr; // ---------------------------------------------------------------- pre(l, group): up to the ring auto pre = [&](int64_t l, int grp) -> bool { const int tb = tb_[grp], te = te_[grp], n = te - tb; @@ -333,6 +336,7 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { float* normalized = (float*) ((uint8_t*) ss.ple.scratch + ple_block_scratch_bytes()); for (int t = tb; t < te; ++t) { gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); + if (cvec_enabled && !cvec_residual(Rt(t), 1, g.hc, l - 1, N, cs, err)) return false; PleOut po; po.normalized = normalized; po.result = Rt(t); @@ -347,6 +351,13 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } pending = false; } + if (cvec_enabled && pending) { + for (int t = tb; t < te; ++t) { + gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); + if (!cvec_residual(Rt(t), 1, g.hc, l - 1, N, cs, err)) return false; + } + pending = false; + } auto gr_read_group = [&](int half, bool apply, float* inj_prev, float* inj_out) { FusedGrArgs fa[kFusedGrMaxT]; for (int t = tb; t < te; ++t) { @@ -559,7 +570,10 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { if (!moe_combine_parts(g, l, K, mb, parts_ + (size_t) t * K * N, bo_ + t * N, cs, err)) return false; } if (l == g.n_layers - 1) - for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); + for (int t = tb; t < te; ++t) { + gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); + if (cvec_enabled && !cvec_residual(Rt(t), 1, g.hc, l, N, cs, err)) return false; + } return true; }; diff --git a/src/kernels/cuda/elementwise.cu b/src/kernels/cuda/elementwise.cu index 420db2aadb5..fa2243efab3 100644 --- a/src/kernels/cuda/elementwise.cu +++ b/src/kernels/cuda/elementwise.cu @@ -161,6 +161,32 @@ void add_inplace(float* dst, const float* src, int64_t n, void* stream) { sync_if_needed(stream, "add_inplace"); } +__global__ void project_unit_broadcast_kernel(float* __restrict__ x, const float* __restrict__ vec, + float scale, int64_t rows, int64_t cols) { + __shared__ float sums[THREADS]; + const int64_t row = blockIdx.x; + if (row >= rows) return; + float* xr = x + row * cols; + float dot = 0.0f; + for (int64_t i = threadIdx.x; i < cols; i += blockDim.x) dot += xr[i] * vec[i]; + sums[threadIdx.x] = dot; + __syncthreads(); + for (int stride = THREADS / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) sums[threadIdx.x] += sums[threadIdx.x + stride]; + __syncthreads(); + } + const float factor = scale * sums[0]; + for (int64_t i = threadIdx.x; i < cols; i += blockDim.x) xr[i] -= factor * vec[i]; +} + +void project_unit_broadcast(float* x, const float* unit_vec, float scale, int64_t rows, int64_t cols, + void* stream) { + if (rows <= 0 || cols <= 0 || scale == 0.0f) return; + project_unit_broadcast_kernel<<<(unsigned) rows, THREADS, 0, (cudaStream_t) stream>>>( + x, unit_vec, scale, rows, cols); + check_launch("project_unit_broadcast"); +} + void f32_to_f16_bulk(const float* x, uint16_t* y, int64_t n, void* stream) { if (n <= 0) return; to_f16_kernel<<>>(x, y, n); diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index a25288c6767..c6270cf045b 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -2,6 +2,8 @@ #include "strata/prefill/prefill.hpp" #include "strata/core/layout.hpp" +#include "strata/core/steer.hpp" +#include "strata/kernels/elementwise.hpp" #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/native_qsa_indexer.hpp" #include "strata/kernels/ngram.hpp" @@ -331,7 +333,6 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e if (!bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err)) return false; if (!bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err)) return false; gr_mix(m.xn, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h); - if (half == 0 && !core::is_qsa_layer(g, l)) { // ======================= GDN ======================= const core::WeightRef *wqkv = need(v, "attn_qkv.weight", err), *wg = need(v, "attn_gate.weight", err), @@ -515,6 +516,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e } // ---- the hyper-connection write of this half gr_write(m.R, m.bo, m.inj, HC, T, m.cs); + if (half == 1 && !core::cvec_residual(m.R, T, g.hc, l, N, m.cs, err)) return false; } } stats_.tokens += T; diff --git a/src/program/generate.cpp b/src/program/generate.cpp index f730f90a787..5ad8b662584 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -19,6 +19,7 @@ #include "strata/core/layer.hpp" #include "strata/core/layout.hpp" #include "strata/core/session.hpp" +#include "strata/core/steer.hpp" #include "strata/core/weights.hpp" #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/cpu/pool.hpp" @@ -235,6 +236,11 @@ struct Options { /// The vision path: keep a per-cell (t, h, w) rotary position table so --serve can take GENI requests. bool vision = false; int adapt_swaps = 96; + /// Residual control vectors: `PATH[:SCALE]`, repeatable; directions are projected out per layer. + std::vector cvec_specs; + /// `--control-vector-layer-range FIRST LAST` (llama.cpp's name, inclusive both ends): only + /// directions on layers inside [FIRST, LAST] are applied. Negative first = no limit. + int64_t cvec_layer_first = -1, cvec_layer_last = -1; /// --serve: how many conversation checkpoints to keep between requests (0 = every request reads its whole /// prompt again, the v0.1.2 behaviour). One is the GDN recurrence of the 36 layers, the QSA indexer tails and /// the PLE history (~118 MB of host RAM); the KV cache itself is positional and stays where it is. @@ -290,6 +296,13 @@ void usage() { " --native-head-gguf PATH native Q5_K head from model shard 1; requires --stream-token\n" " --native-dense-gguf PATH native GDN/QSA/shared projections; repeat for each source model shard\n" " --expert-cache-cpu-order experimental GPU expert reduction matching CPU order\n" + " --control-vector-scaled FILE:SCALE\n" + " GGUF control vector (direction.L vectors, projected from the residual);\n" + " SCALE is a repeatable projection strength (default 1)\n" + " --control-vector-layer-range FIRST LAST\n" + " apply only directions for layers FIRST..LAST (inclusive)\n" + " --cvec-mode project remove the direction component from each layer residual\n" + " --cvec-dir per-layer match direction.L to layer L\n" " --max-new N tokens to generate (default 16)\n" " --max-context N KV/state capacity (default 4096)\n" " --greedy argmax (the default)\n" @@ -393,6 +406,12 @@ bool parse_i64_list(const char* s, std::vector& out, std::string& err) return true; } +bool parse_i64_value(const char* text, int64_t& out) { + const std::string value(text ? text : ""); + const auto parsed = std::from_chars(value.data(), value.data() + value.size(), out); + return !value.empty() && parsed.ec == std::errc{} && parsed.ptr == value.data() + value.size(); +} + /// The pool's adapter plus the wall-clock it spent, so the report can say how much of the token was the CPU. struct Drive { strata::core::ExpertDispatch d; @@ -680,6 +699,25 @@ int main(int argc, char** argv) { else if (a == "--no-publish-kernel") o.no_publish_kernel = true; else if (a == "--no-fused-gdn") o.no_fused_gdn = true; else if (a == "--no-fast-select") o.no_fast_select = true; + + else if (a == "--control-vector-scaled") o.cvec_specs.push_back(next("--control-vector-scaled")); + else if (a == "--cvec-mode") { + const std::string mode = next("--cvec-mode"); + if (mode != "project") { std::fprintf(stderr, "strata generate: invalid --cvec-mode: %s (only project is supported)\n", mode.c_str()); return 2; } + } + else if (a == "--cvec-dir") { + const std::string dir = next("--cvec-dir"); + if (dir != "per-layer") { std::fprintf(stderr, "strata generate: this build supports --cvec-dir per-layer only\n"); return 2; } + } + else if (a == "--control-vector-layer-range") { + const char* first = next("--control-vector-layer-range"); + const char* last = next("--control-vector-layer-range"); + if (!parse_i64_value(first, o.cvec_layer_first) || !parse_i64_value(last, o.cvec_layer_last)) { + std::fprintf(stderr, "strata generate: --control-vector-layer-range expects two integers\n"); + return 2; + } + } + else { // An unknown flag is an ERROR and not a warning: a typo'd `--max-neww` that silently generated 16 // tokens would look like a working run. @@ -765,6 +803,11 @@ int main(int argc, char** argv) { return 2; } + if (o.cvec_layer_first >= 0 && o.cvec_layer_last < o.cvec_layer_first) { + std::fprintf(stderr, "strata generate: --control-vector-layer-range %lld %lld is empty\n", + (long long) o.cvec_layer_first, (long long) o.cvec_layer_last); + return 2; + } if (o.native_flash_attn_short && o.max_context > 256) { std::fprintf(stderr, "strata generate: --native-flash-attn-short requires --max-context <=256\n"); return 2; @@ -910,8 +953,31 @@ int main(int argc, char** argv) { strata::kernels::mrope_table_set(d_mrope); } strata::kernels::ple_set_native_postops(o.native_ple_postops); - const strata::core::ModelGeometry g; + + // ---- residual control vectors (strata/core/steer.hpp). Static for the run: the directions are + // uploaded once and per-layer projections become ordinary nodes of captured layer graphs, so + // this must happen BEFORE session_capture. No control vector means the plan is never active and + // every projection call returns before touching the stream - runs without vectors are untouched. + const strata::core::ModelGeometry g; // the compiled-in qwen4exp geometry + strata::core::SteeringPlan steering; + if (!o.cvec_specs.empty()) { + if (!steering.load(o.cvec_specs, g, err, o.cvec_layer_first, o.cvec_layer_last)) { + std::fprintf(stderr, "strata generate: %s\n", err.c_str()); + return 1; + } + for (const auto& v : steering.entries()) + std::fprintf(stderr, "strata generate: control vector layer %lld (scale %.3g) from %s\n", + (long long) v.layer, (double) v.alpha, v.spec.c_str()); + if (o.cvec_layer_first >= 0) + std::fprintf(stderr, "strata generate: steering layer range %lld..%lld\n", + (long long) o.cvec_layer_first, (long long) o.cvec_layer_last); + std::fprintf(stderr, "strata generate: control vector mode = project, dir = per-layer, layers %lld..%lld\n", + (long long) (o.cvec_layer_first >= 0 ? o.cvec_layer_first : 0), + (long long) (o.cvec_layer_first >= 0 ? o.cvec_layer_last : g.n_layers - 1)); + strata::core::steer_set_active(&steering); + } const int64_t K = 10; + if (o.max_context < (int64_t) o.tokens.size() + o.max_new) { std::fprintf(stderr, "strata generate: --max-context %lld cannot hold %zu prompt + %lld new tokens\n", (long long) o.max_context, o.tokens.size(), (long long) o.max_new);