Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 17:38:20 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

Merge f-0291-split (up to ea2aed12c: auto/on/off switch, drafting fix, i-quant sets, PF_JOB) into f-0291-split-hip; AMD opt-in rule (AUTO_SPLIT_AMD_ARCHS), strata-split-N-gfx1201 lookup, packaging/build_split_hip.sh, AMD tests

Browse Source
This commit is contained in:
Niko1221 committed 2026-11-02 12:13:42 +01:00
2 parents 242549e432 ea2aed12cd
commit 2c1771dce2
28 files changed
+3272 -78
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/SPLIT_ENGINE.md
+15 -1
View File
@@ -56,6 +56,20 @@ host memory that both cards read and write with ordinary loads and stores. No pe
fails if one appears in `src/split`, and `packaging/build_split_cuda.sh` fails if the built engine can call one. The IOMMU and BAR1 start
checks are removed because the path they protect no longer exists.
## AMD (Radeon AI PRO R9700, gfx1201): opt-in
The same sources build for HIP (`-DSTRATA_ENABLE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1201 -DSTRATA_SPLIT_CARDS=4`, `packaging/build_split_hip.sh`): gfx12 WMMA tiles, wave32,
the HIP runtime through the tree's hip_compat shim. Four R9700 (128 GB) hold UD-Q4_K_XL's experts; two do not. It ships **opt-in**: `--split-engine on` starts it, and
`auto` keeps the layer split on gfx1201 (`AUTO_SPLIT_AMD_ARCHS` in `strata_control/split_engine.py` is empty; adding `"gfx1201"` is the one change that makes it automatic, once 8 users decode
at or above the layer split and a 103K prompt reads in the layer split's time). The launcher looks for `strata-split-4-gfx1201` beside the engine. Numbers and open items:
`phaseA-tests/sfnx/RESULTS-HIP.md`.
The cards exchange sums differently from NVIDIA here. By default (`FNX_PEER` unset) the engine probes every directed pair of cards at start (kernel stores into the peer, a full read-back
compare, a flag round trip) and, if all pass, pushes the decode sums and the attention merge with kernel stores into the receiver's memory (never `hipMemcpyPeerAsync`); on any
failure, or with `FNX_PEER=0`, it uses the host-memory exchange for the whole run and prints one plain line. Peer access is enabled in exactly one file (`src/split/src/hip/peer.h`), pinned by
`tests/split/test_split_engine.py`. `FNX_PF_RCCL=1` (off by default, needs the peer push) sums the prompt chunks through RCCL, which is loaded with `dlopen` and never linked: without
`librccl.so.1`, or if its self-test fails, the engine says so in one line and keeps the host exchange (`tools/rccl_policy_test.cpp`).
## Not proven
No GeForce pair has been tested (there is none in the lab). Three and four cards, other architectures, Windows and macOS, a 64 GB RAM PC
@@ -64,4 +78,4 @@ with two big cards, and contexts over 128K have not been run. The copy-engine ri
## Build
`-DSTRATA_BUILD_SPLIT=ON -DSTRATA_GGML_DIR=<llama.cpp checkout> [-DSTRATA_SPLIT_CARDS="2;3"] [-DSTRATA_SPLIT_ARCH=86]` builds
`strata-split-N` per card count; `packaging/build_split_cuda.sh` does that with the release engine's toolchain rules and checks the result.
`strata-split-N` per card count; `packaging/build_split_cuda.sh` does that with the release engine's toolchain rules and checks the result (`packaging/build_split_hip.sh` for AMD).
packaging/build_split_hip.sh
+56
View File
@@ -0,0 +1,56 @@
#!/bin/bash
# strata-split-4 (the all-cards split engine, src/split, from sFNX, MIT, John Fink) for AMD RDNA4 (gfx1201: Radeon AI PRO R9700, RX 9070) on Linux, built next to the HIP
# engine package (packaging/build_hip_linux.sh, VARIANT=rx9070) so package_engine.py ships it with `--include 'strata-split-*'` beside `strata`. One executable per card
# count, named strata-split-<N>-gfx1201 (the launcher, strata_cli/launch/runner.py, looks for exactly that name on AMD). Four cards is the only count that has been run
# (4 x R9700 hold UD-Q4_K_XL's experts; two do not). RUNPATH $ORIGIN: it finds the ROCm libraries the engine package bundles (libamdhip64, hipblas, rocblas and its
# Tensile data for gfx1201) beside itself, the same set `strata` uses.
#
# SRC=<checkout> B=<build dir> HP=<ROCm devel root> [LIBW=<ROCm libraries root>] GG=<llama.cpp checkout at the pinned commit> [CARDS="4"] [JOBS=12] packaging/build_split_hip.sh
#
# After the build the script FAILS if the binary
# - links librccl (RCCL is loaded at run time with dlopen when FNX_PF_RCCL=1 and is optional: the split must start without it),
# - can call hipMemcpyPeer* (the engine's peer traffic is kernel stores only; src/split/src/hip/peer.h is the one file that touches peer access),
# - or is missing the gfx1201 code object.
# Prints the binary paths (the renamed copies), one per line, last.
set -eu
SRC=${SRC:-$(cd "$(dirname "$0")/.." && pwd)}
B=${B:?B=<build dir>}
HP=${HP:?HP=<ROCm devel root>}
LIBW=${LIBW:-$HP}
GG=${GG:?GG=<llama.cpp checkout at the pinned commit>}
CARDS=${CARDS:-4}
JOBS=${JOBS:-12}
ARCH=gfx1201
export HIP_PLATFORM=amd HIP_COMPILER=clang HIP_RUNTIME=rocclr ROCM_PATH=$HP HIP_PATH=$HP PATH=$HP/bin:$PATH
SYSDEPS=""; [ -d "$HP/lib/rocm_sysdeps" ] && SYSDEPS=";$HP/lib/rocm_sysdeps"
BC=$HP/lib/llvm/amdgcn/bitcode; [ -d "$BC" ] || BC=$HP/amdgcn/bitcode
rm -f "$B/CMakeCache.txt"
cmake -G Ninja -S "$SRC" -B "$B" -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF -DSTRATA_BUILD_TESTS=OFF -DSTRATA_BUILD_SPLIT=ON \
"-DSTRATA_SPLIT_CARDS=${CARDS// /;}" "-DCMAKE_HIP_ARCHITECTURES=$ARCH" -DSTRATA_PREFILL_MMQ=ON -DSTRATA_NATIVE_EXPERTS=ON -DSTRATA_GGML_DIR="$GG" \
-DCMAKE_HIP_COMPILER="$HP/llvm/bin/clang++" -DCMAKE_HIP_COMPILER_ROCM_ROOT="$HP" "-DCMAKE_PREFIX_PATH=$HP;$LIBW$SYSDEPS" \
"-DCMAKE_HIP_FLAGS=--rocm-path=$HP --rocm-device-lib-path=$BC" -DSTRATA_PORTABLE=ON \
"-DCMAKE_EXE_LINKER_FLAGS=-Wl,-rpath,'\$ORIGIN' -Wl,--enable-new-dtags"
out=""
for n in $CARDS; do
ninja -j "$JOBS" -C "$B" "strata-split-$n"
exe=$(find "$B" -name "strata-split-$n" -type f | head -1)
[ -x "$exe" ] || { echo "build_split_hip.sh: no strata-split-$n" >&2; exit 1; }
if readelf -d "$exe" | grep -qi 'NEEDED.*rccl'; then
echo "build_split_hip.sh: $exe links librccl (it must be dlopen'ed only):" >&2
readelf -d "$exe" | grep -i needed >&2
exit 1
fi
if nm -D --undefined-only "$exe" | grep -qiE 'hipMemcpy(3D)?Peer'; then
echo "build_split_hip.sh: $exe can call hipMemcpyPeer*:" >&2
nm -D --undefined-only "$exe" | grep -iE 'hipMemcpy(3D)?Peer' >&2
exit 1
fi
if ! grep -aq "$ARCH" "$exe"; then
echo "build_split_hip.sh: $exe has no $ARCH code object" >&2
exit 1
fi
named="$(dirname "$exe")/strata-split-$n-$ARCH"
cp -f "$exe" "$named"
out="$out$named"$'\n'
done
printf '%s' "$out"
serve/server.py
+12 -1
View File
@@ -8938,7 +8938,18 @@ def start_engine_or_fallback(cfg: dict, exe: str, effort_end: list[str] | None,
except Exception as e:
if not isinstance(fb, dict) or not fb.get("exe"):
raise
why = str(e).strip().splitlines()[0][:200] if str(e).strip() else type(e).__name__
msg = str(e).strip()
why = msg.splitlines()[0][:200] if msg else type(e).__name__
# the engine's own reason ("sfnx: failed: ..." in its log tail, which the exception carries) is what a person needs to read
own = next((ln.split("sfnx: failed: ", 1)[1].strip() for ln in reversed(msg.splitlines()) if "sfnx: failed: " in ln), "")
if own:
why = own[:300]
log = cfg.get("log")
if log: # the normal engine writes the same log file: keep the split engine's
try:
os.replace(log, str(log) + ".split-engine")
except OSError:
pass
print(f"[strata] The all-cards split engine could not start ({why}). Starting the normal layer split instead.", flush=True)
cfg2 = {k: v for k, v in cfg.items() if k not in SPLIT_KEYS + ("exe", "args", "env", "layer_split", "replicas", "user_args", "drafting")}
cfg2.update(fb)
serve/test_split_fallback.py
+45
View File
@@ -75,6 +75,51 @@ class Fallback(unittest.TestCase):
self.assertIn("all-cards split engine could not start", lines[0])
self.assertIn("normal layer split", lines[0])
def test_the_engines_own_reason_is_in_the_line_and_its_log_is_kept(self):
import tempfile
log = os.path.join(tempfile.mkdtemp(), "engine.log")
Path(log).write_text("sfnx: failed: the all-cards split engine does not read this model's routed experts (blk.0.ffn_gate_exps.weight is ggml type 14); use the normal layer split\n")
cfg = split_cfg()
cfg["log"] = log
class Refusing(FakeEngine):
def __init__(self, exe, args, cwd=None, log=None, env=None, lazy=False, **kw):
FakeEngine.started.append((exe, list(args), cwd, dict(env or {})))
if exe == "/e/strata-split-2":
raise S.EngineDied("the engine exited before it was ready (see x)\n" + Path(log).read_text())
self.info = {}
with mock.patch.object(S, "StrataEngine", Refusing):
engine, out = self.run_start(cfg)
lines = [ln for ln in out.splitlines() if ln.strip()]
self.assertEqual(len(lines), 1)
self.assertIn("does not read this model's routed experts", lines[0])
self.assertNotIn("exited before it was ready", lines[0])
self.assertTrue(os.path.exists(log + ".split-engine")) # the normal engine will overwrite the log; the refusal stays readable
self.assertEqual([s[0] for s in FakeEngine.started], ["/e/strata-split-2", NORMAL])
def test_a_split_config_keeps_its_arguments_through_the_drafter_step(self):
"""strata_control/split_engine.split_config writes "drafting": "mtp" (what the split engine runs), so the drafter step, which resolves Auto to
DFlash2 where the engine has it, adds no DFlash2 flag to the split engine's arguments (it would refuse them and every start would fall back)."""
from strata_control import split_engine as SE
base = {"exe": NORMAL, "args": ["--pack", "/p", "--native", "/m/a.gguf", "--mtp", "/d/mtp/rt", "--layer-split", "auto"], "cwd": "/x",
"tokenizer": "/t", "port": 8080}
decision = SE.Decision(True, [], 2)
cfg = SE.split_config(base, decision, exe="/e/strata-split-2", native="/m/a.gguf", mtp_dir="/d/mtp/rt", draft_vocab="/d/v.bin", context=131072)
args = S.engine_args(cfg)
self.assertEqual(args[:2], ["--native", "/m/a.gguf"])
self.assertNotIn("--draft", args)
self.assertNotIn("--draft-model", args)
self.assertEqual(args[args.index("--mtp") + 1], "/d/mtp/rt")
self.assertEqual(args[args.index("--spec") + 1], "6")
self.assertNotIn("--layer-split", args)
# the layer split that starts instead keeps the person's drafting word, or none (Auto)
self.assertNotIn("drafting", cfg["split_fallback"])
cfg = SE.split_config({**base, "drafting": "auto"}, decision, exe="/e/strata-split-2", native="/m/a.gguf", mtp_dir="/d/mtp/rt",
draft_vocab="/d/v.bin", context=131072)
self.assertEqual(cfg["split_fallback"]["drafting"], "auto")
self.assertEqual(cfg["drafting"], "mtp")
def test_without_a_fallback_the_failure_stays_a_failure(self):
cfg = split_cfg()
del cfg["split_fallback"]
src/split/CMakeLists.txt
+30 -1
View File
@@ -15,6 +15,7 @@
# the engine refuses a card of any other architecture
set(STRATA_SPLIT_CARDS "2" CACHE STRING "card counts strata-split is built for (2, 3, 4)")
set(STRATA_SPLIT_ARCH "" CACHE STRING "compute capability strata-split is built for (e.g. 86)")
set(STRATA_SPLIT_PF_JOB "64" CACHE STRING "rows of one expert per prompt-path job: 64 (default) or 128")
set(_fnx_hip OFF)
if(STRATA_ENABLE_HIP AND NOT STRATA_ENABLE_CUDA AND NOT WIN32 AND NOT APPLE)
@@ -88,7 +89,7 @@ foreach(_n ${STRATA_SPLIT_CARDS})
endif()
target_include_directories(${_t} PRIVATE ${_sp}/src ${_sp}/tools ${STRATA_GGML_DIR}/ggml/src/ggml-cuda ${STRATA_GGML_DIR}/ggml/src
${STRATA_GGML_DIR}/ggml/include)
target_compile_definitions(${_t} PRIVATE FNX_NCARD=${_n} FNX_ARCH=${STRATA_SPLIT_ARCH} FNX_27B=0 FNX_HOSTRING=1
target_compile_definitions(${_t} PRIVATE FNX_NCARD=${_n} FNX_ARCH=${STRATA_SPLIT_ARCH} FNX_27B=0 FNX_HOSTRING=1 FNX_PF_JOB=${STRATA_SPLIT_PF_JOB}
FNX_DIR="${_sp}")
target_compile_options(${_t} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math --expt-relaxed-constexpr -Wno-deprecated-gpu-targets>)
set_target_properties(${_t} PROPERTIES CUDA_ARCHITECTURES "${STRATA_SPLIT_ARCH}-real" CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON
@@ -113,4 +114,32 @@ if(STRATA_BUILD_TESTS)
target_link_libraries(split_${_tt} PRIVATE Threads::Threads ${CMAKE_DL_LIBS})
add_test(NAME split_${_tt} COMMAND split_${_tt})
endforeach()
# the load-time conversion's oracle tool (tests/split/convert_check.py compares its output with gguf-py on a real file; not a ctest)
add_executable(split_convert_dump ${_sp}/tools/convert_dump.cpp ${_sp}/src/gguf.cpp)
target_include_directories(split_convert_dump PRIVATE ${_sp}/src ${_sp}/tools)
target_compile_definitions(split_convert_dump PRIVATE FNX_NCARD=${_n0} FNX_27B=0)
target_compile_features(split_convert_dump PRIVATE cxx_std_17)
set_target_properties(split_convert_dump PROPERTIES CUDA_ARCHITECTURES OFF)
target_link_libraries(split_convert_dump PRIVATE Threads::Threads)
# the weight tiles of the prompt kernel run on a GPU against the file (tests/split/tile_check.py reads its output); not a ctest
add_executable(split_moe_tile_test ${_sp}/tools/moe_tile_test.cu ${_sp}/src/gguf.cpp)
target_include_directories(split_moe_tile_test PRIVATE ${_sp}/src ${_sp}/tools ${STRATA_GGML_DIR}/ggml/src/ggml-cuda ${STRATA_GGML_DIR}/ggml/src
${STRATA_GGML_DIR}/ggml/include)
target_compile_definitions(split_moe_tile_test PRIVATE FNX_NCARD=${_n0} FNX_ARCH=${STRATA_SPLIT_ARCH} FNX_27B=0)
target_compile_options(split_moe_tile_test PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math --expt-relaxed-constexpr -Wno-deprecated-gpu-targets>)
set_target_properties(split_moe_tile_test PROPERTIES CUDA_ARCHITECTURES "${STRATA_SPLIT_ARCH}-real" CUDA_STANDARD 17 CXX_STANDARD 17)
target_link_libraries(split_moe_tile_test PRIVATE Threads::Threads)
# fnx_run: sFNX's test driver (teacher-forced logits to an .lgr file, prompt/decode checks), the same engine objects with its own main
set(_sp_run_src ${_sp_engine})
list(REMOVE_ITEM _sp_run_src ${_sp}/tools/sfnx.cpp)
add_executable(split_fnx_run ${_sp_run_src} ${_sp}/tools/fnx_run.cpp ${_sp_kparts})
set_source_files_properties(${_sp}/tools/fnx_run.cpp PROPERTIES LANGUAGE CUDA)
target_include_directories(split_fnx_run PRIVATE ${_sp}/src ${_sp}/tools ${STRATA_GGML_DIR}/ggml/src/ggml-cuda ${STRATA_GGML_DIR}/ggml/src
${STRATA_GGML_DIR}/ggml/include)
target_compile_definitions(split_fnx_run PRIVATE FNX_NCARD=${_n0} FNX_ARCH=${STRATA_SPLIT_ARCH} FNX_27B=0 FNX_HOSTRING=1 FNX_DIR="${_sp}")
target_compile_options(split_fnx_run PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math --expt-relaxed-constexpr -Wno-deprecated-gpu-targets>)
set_target_properties(split_fnx_run PROPERTIES CUDA_ARCHITECTURES "${STRATA_SPLIT_ARCH}-real" CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON CXX_STANDARD 17)
find_package(CUDAToolkit REQUIRED)
target_link_libraries(split_fnx_run PRIVATE CUDA::cublas CUDA::cuda_driver Threads::Threads)
set_target_properties(split_fnx_run PROPERTIES CUDA_RUNTIME_LIBRARY Shared)
endif()
src/split/src/gguf.cpp
+42 -2
View File
@@ -14,6 +14,8 @@
#include <regex>
#include <stdexcept>
#include "tensor_convert.hpp"
namespace fnx {
TypeInfo type_info(uint32_t t) {
@@ -33,6 +35,11 @@ TypeInfo type_info(uint32_t t) {
case T_Q5_K: return {256, 176};
case T_Q6_K: return {256, 210};
case T_IQ4_NL: return {32, 18};
case T_IQ3_XXS: return {256, 98};
case T_IQ3_S: return {256, 110};
case T_IQ2_S: return {256, 82};
case T_IQ4_XS: return {256, 136};
case T_Q2_0: return {64, 18};
}
throw std::runtime_error("gguf: unknown ggml type " + std::to_string(t));
}
@@ -101,6 +108,33 @@ Gguf::Gguf(const std::string& first) {
} else {
load_shard(first, true);
}
convert_small_tensors();
}
// A file with routed experts in the IQ types or Strata's 2-bit type (the ISTA IQ3_S set) also keeps its dense matrices, routers,
// embeddings and projections in K-quants, IQ4_XS and BF16, which the engine does not read. Those are dequantized and written
// again as Q8_0 (F32 for the small ones), once, here; the routed experts are left as they are. A file without such experts
// (UD-Q4_K_XL, Q8_0) is not touched at all: the check below is the only difference for it.
void Gguf::convert_small_tensors() {
bool mixed = false;
for (const auto& kv : tensors_)
mixed |= kv.second.type == T_IQ2_S || kv.second.type == T_IQ3_XXS || kv.second.type == T_IQ3_S || kv.second.type == T_Q2_0;
if (!mixed) return;
for (auto& kv : tensors_) {
Tensor& t = kv.second;
if (t.type == T_Q2_0) { // the routed down experts of some layers: the same values as Q4_0 blocks
owned_.emplace_back();
conv::repack_q2_0_to_q4_0(t, owned_.back());
++converted_;
continue;
}
const uint32_t want = conv::wanted_type(t.name);
if (want == 0xFFFFFFFFu || want == t.type) continue;
owned_.emplace_back();
conv::convert(t, want, owned_.back());
++converted_;
}
std::fprintf(stderr, "gguf: %d dense/small tensors converted to Q8_0/F32 at load (the file's routed experts keep their types)\n", converted_);
}
Gguf::~Gguf() {
@@ -281,7 +315,12 @@ Stager::Stager(const Gguf& g, std::vector<Group> groups, int consumers, int ahea
std::vector<const Tensor*> ts;
for (const std::string& n : gr.spec.names) {
const Tensor& t = g.t(n);
if (t.shard < 0) throw std::runtime_error("stager: " + n + " is not in a file");
if (t.shard < 0) {
if (!t.data) throw std::runtime_error("stager: " + n + " is not in a file");
gr.t[n] = t; // converted at load: nothing to read
gr.fixed.insert(n);
continue;
}
ts.push_back(&t);
}
std::sort(ts.begin(), ts.end(), [](const Tensor* a, const Tensor* b) { return a->shard != b->shard ? a->shard < b->shard : a->offset < b->offset; });
@@ -406,7 +445,8 @@ void Stager::read_group(G& gr) {
for (auto& t : th) t.join();
for (const std::string& e : errs)
if (!e.empty()) throw std::runtime_error(e);
for (auto& [name, t] : gr.t) t.data = gr.buf + (uintptr_t) t.data;
for (auto& [name, t] : gr.t)
if (!gr.fixed.count(name)) t.data = gr.buf + (uintptr_t) t.data;
read_s_ += std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
}
src/split/src/gguf.hpp
+8 -1
View File
@@ -10,6 +10,7 @@
#include <atomic>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <map>
#include <mutex>
#include <set>
@@ -21,7 +22,8 @@ namespace fnx {
enum GgmlType : uint32_t {
T_F32 = 0, T_F16 = 1, T_Q4_0 = 2, T_Q4_1 = 3, T_Q5_0 = 6, T_Q5_1 = 7, T_Q8_0 = 8, T_Q8_1 = 9, T_Q2_K = 10,
T_Q3_K = 11, T_Q4_K = 12, T_Q5_K = 13, T_Q6_K = 14, T_Q8_K = 15, T_IQ4_NL = 20, T_BF16 = 30,
T_Q3_K = 11, T_Q4_K = 12, T_Q5_K = 13, T_Q6_K = 14, T_Q8_K = 15, T_IQ3_XXS = 18, T_IQ4_NL = 20, T_IQ3_S = 21, T_IQ2_S = 22,
T_IQ4_XS = 23, T_BF16 = 30, T_Q2_0 = 42, // 42: Strata's 64-value 2-bit blocks (not in llama.cpp)
};
struct TypeInfo {
@@ -64,6 +66,7 @@ public:
#endif
const std::map<std::string, Tensor>& tensors() const { return tensors_; }
int shards() const { return (int) maps_.size(); }
int converted() const { return converted_; } // tensors re-quantized at load (0 for every file the engine reads as it is)
const std::string& shard_path(int s) const { return maps_[(size_t) s].path; }
size_t shard_size(int s) const { return maps_[(size_t) s].size; }
const uint8_t* shard_base(int s) const { return (const uint8_t*) maps_[(size_t) s].base; }
@@ -80,7 +83,10 @@ private:
std::map<std::string, std::vector<std::string>> str_arrays_;
#endif
std::set<std::string> keys_;
std::deque<std::vector<uint8_t>> owned_; // tensors converted at load (tensor_convert.hpp): their bytes live here, not in a file
int converted_ = 0;
void load_shard(const std::string& path, bool meta);
void convert_small_tensors();
};
// Reads groups of tensors (a layer's, say) into host buffers ahead of the threads that slice them, in the order they
@@ -116,6 +122,7 @@ private:
std::vector<Span> spans;
size_t bytes = 0;
std::map<std::string, Tensor> t;
std::set<std::string> fixed; // converted tensors: their bytes are already in memory (data is a pointer, not an offset into buf)
uint8_t* buf = nullptr;
int left = 0; // consumers not through with it yet
bool ready = false;
src/split/src/kernels.cu
+19 -1
View File
@@ -139,6 +139,10 @@ template <ggml_type type> __device__ __forceinline__ float vdot(const void* v, c
else if constexpr (type == GGML_TYPE_Q4_0) return vec_dot_q4_0_q8_1(v, y, kbx, iqs);
else if constexpr (type == GGML_TYPE_Q6_K) return vec_dot_q6_K_q8_1(v, y, kbx, iqs);
else if constexpr (type == GGML_TYPE_IQ4_NL) return vec_dot_iq4_nl_q8_1(v, y, kbx, iqs);
else if constexpr (type == GGML_TYPE_IQ4_XS) return vec_dot_iq4_xs_q8_1(v, y, kbx, iqs);
else if constexpr (type == GGML_TYPE_IQ3_S) return vec_dot_iq3_s_q8_1(v, y, kbx, iqs);
else if constexpr (type == GGML_TYPE_IQ3_XXS) return vec_dot_iq3_xxs_q8_1(v, y, kbx, iqs);
else if constexpr (type == GGML_TYPE_IQ2_S) return vec_dot_iq2_s_q8_1(v, y, kbx, iqs);
else return 0.0f;
}
template <ggml_type type> constexpr int vdr_of() {
@@ -148,7 +152,11 @@ template <ggml_type type> constexpr int vdr_of() {
else if constexpr (type == GGML_TYPE_Q5_1) return VDR_Q5_1_Q8_1_MMVQ;
else if constexpr (type == GGML_TYPE_Q4_0) return VDR_Q4_0_Q8_1_MMVQ;
else if constexpr (type == GGML_TYPE_Q6_K) return VDR_Q6_K_Q8_1_MMVQ;
else return VDR_IQ4_NL_Q8_1_MMVQ;
else if constexpr (type == GGML_TYPE_IQ4_NL) return VDR_IQ4_NL_Q8_1_MMVQ;
else if constexpr (type == GGML_TYPE_IQ4_XS) return VDR_IQ4_XS_Q8_1_MMVQ;
else if constexpr (type == GGML_TYPE_IQ3_S) return VDR_IQ3_S_Q8_1_MMVQ;
else if constexpr (type == GGML_TYPE_IQ3_XXS) return VDR_IQ3_XXS_Q8_1_MMVQ;
else return VDR_IQ2_S_Q8_1_MMVQ;
}
constexpr int MMVQ_THREADS = 128;
@@ -2081,6 +2089,11 @@ void moe_mm(bool down, const DW& A, size_t sa, const Job* jobs, const int* njobs
case T_Q5_1: moe_launch<GGML_TYPE_Q5_1>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_Q4_0: moe_launch<GGML_TYPE_Q4_0>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_Q8_0: moe_launch<GGML_TYPE_Q8_0>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_IQ4_NL: moe_launch<GGML_TYPE_IQ4_NL>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_IQ4_XS: moe_launch<GGML_TYPE_IQ4_XS>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_IQ3_S: moe_launch<GGML_TYPE_IQ3_S>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_IQ3_XXS: moe_launch<GGML_TYPE_IQ3_XXS>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
case T_IQ2_S: moe_launch<GGML_TYPE_IQ2_S>(down, w, jobs, njobs, njobs_max, in, (float*) out, s, T); break;
default: throw std::runtime_error("moe: expert type " + std::to_string(A.type) + " not handled");
}
ck_launch(down ? "moe_down" : "moe_gate_up");
@@ -2137,6 +2150,11 @@ void moe_mm_big(bool down, const DW& A, size_t sa, const BigJob* jobs, const int
case T_Q5_1: moe_launch_big<GGML_TYPE_Q5_1>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_Q4_0: moe_launch_big<GGML_TYPE_Q4_0>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_Q8_0: moe_launch_big<GGML_TYPE_Q8_0>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_IQ4_NL: moe_launch_big<GGML_TYPE_IQ4_NL>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_IQ4_XS: moe_launch_big<GGML_TYPE_IQ4_XS>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_IQ3_S: moe_launch_big<GGML_TYPE_IQ3_S>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_IQ3_XXS: moe_launch_big<GGML_TYPE_IQ3_XXS>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
case T_IQ2_S: moe_launch_big<GGML_TYPE_IQ2_S>(down, w, jobs, njobs, njobs_cap, in, (float*) out, s); break;
default: throw std::runtime_error("moe: expert type " + std::to_string(A.type) + " not handled");
}
ck_launch(down ? "moe_down (big jobs)" : "moe_gate_up (big jobs)");
src/split/src/kernels.hpp
+7 -2
View File
@@ -22,8 +22,13 @@ struct Job {
int out[MAXT]; // output rows
};
// the prompt path's expert jobs: up to PF_JOB rows of one expert
constexpr int PF_JOB = 64;
// the prompt path's expert jobs: up to PF_JOB rows of one expert (64: four 16-token tiles per job; 128: eight, the expert's weight tile is then
// dequantized once for twice the tokens; a build option, STRATA_SPLIT_PF_JOB)
#ifndef FNX_PF_JOB
#define FNX_PF_JOB 64
#endif
constexpr int PF_JOB = FNX_PF_JOB;
static_assert(PF_JOB == 64 || PF_JOB == 128, "PF_JOB: 64 or 128");
struct BigJob {
int w;
int n;
src/split/src/model.cpp
+15
View File
@@ -1208,6 +1208,20 @@ Mtp::Mtp(const Model& m, const std::string& dir, const std::string& vocab_file)
Mtp::~Mtp() { free_mtp_cards(cards); }
// A file whose routed experts the engine has no kernel for is refused here, in one plain line, before anything is read into RAM or placed on
// the cards (the app's rule, strata_control/split_engine.py, refuses it earlier; this is the same check for a start that skipped the rule). The
// server then starts the normal layer split. Q2_0 is not listed: the Gguf has already repacked it to Q4_0.
static void check_supported(const Gguf& g) {
const uint32_t ok[] = {T_Q4_K, T_Q5_K, T_Q5_1, T_Q8_0, T_Q4_0, T_IQ4_NL, T_IQ4_XS, T_IQ3_S, T_IQ3_XXS, T_IQ2_S};
for (const auto& kv : g.tensors()) {
const std::string& n = kv.first;
if (n.find("_exps.weight") == std::string::npos || n.find("shexp") != std::string::npos) continue;
if (std::find(std::begin(ok), std::end(ok), kv.second.type) == std::end(ok))
throw std::runtime_error("the all-cards split engine does not read this model's routed experts (" + n + " is ggml type " +
std::to_string(kv.second.type) + "); use the normal layer split");
}
}
Model::Model(const std::string& path, Load how, Store* store, const std::string* image) : g(path) {
// the cards first: reading the embeddings into RAM below takes seconds on a cold cache
{
@@ -1221,6 +1235,7 @@ Model::Model(const std::string& path, Load how, Store* store, const std::string*
" of them, or build for " + std::to_string(ndev) + " (make FNX_NCARD=" + std::to_string(ndev) + ")");
check_cards();
}
check_supported(g);
#if FNX_27B
const std::string p = "qwen35.";
// block_count holds the draft (nextn) layer too
src/split/src/moe_mma.cu
+26
View File
@@ -11,6 +11,25 @@
#include <string>
#include "ggml.h"
#if FNX_HIP
// the tree's hip_compat/intrinsics.hpp names these as macros; ggml's HIP vendor header defines its own: it goes first, then ours return
#undef __vsubss4
#undef __vsub4
#undef __vcmpne4
#undef __shfl_sync
#undef __shfl_up_sync
#undef __shfl_xor_sync
#endif
#include "common.cuh"
#if FNX_HIP
#undef __shfl_sync
#undef __shfl_up_sync
#undef __shfl_xor_sync
#define __shfl_xor_sync(...) (::strata::hip_compat::shfl_xor_sync(__VA_ARGS__))
#define __shfl_up_sync(...) (::strata::hip_compat::shfl_up_sync(__VA_ARGS__))
#define __shfl_sync(...) (::strata::hip_compat::shfl_sync(__VA_ARGS__))
#undef cudaFuncSetAttribute
#endif
namespace fnx {
@@ -24,6 +43,7 @@ __device__ __forceinline__ float silu_(float x) { return x / (1.0f + __expf(-x))
#include "moe_w.cuh"
#include "moe_mma.cuh"
#include "moe_mma_iq.cuh"
template <int TYPE>
void mma_launch(const MoeW& w, const BigJob* jobs, const int* njobs, int njobs_cap, const __half* x, float* out, cudaStream_t s) {
@@ -41,6 +61,12 @@ void moe_mma(const DW& A, size_t sa, const BigJob* jobs, const int* njobs, int n
case T_Q5_K: mma_launch<GGML_TYPE_Q5_K>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_Q5_1: mma_launch<GGML_TYPE_Q5_1>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_Q8_0: mma_launch<GGML_TYPE_Q8_0>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_Q4_0: mma_launch<GGML_TYPE_Q4_0>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_IQ4_NL: mma_launch<GGML_TYPE_IQ4_NL>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_IQ4_XS: mma_launch<GGML_TYPE_IQ4_XS>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_IQ3_S: mma_launch<GGML_TYPE_IQ3_S>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_IQ3_XXS: mma_launch<GGML_TYPE_IQ3_XXS>(w, jobs, njobs, njobs_cap, x, out, s); break;
case T_IQ2_S: mma_launch<GGML_TYPE_IQ2_S>(w, jobs, njobs, njobs_cap, x, out, s); break;
default: throw std::runtime_error("moe_mma: expert type " + std::to_string(A.type) + " not handled");
}
ck_launch("moe_mma");
src/split/src/moe_mma.cuh
+59 -4
View File
@@ -270,6 +270,17 @@ __device__ __forceinline__ void mma_884(float (&d)[8], unsigned a0, unsigned a1,
}
#endif
#if FNX_ARCH >= 80 && !FNX_HIP
// d (16 x 8) += a (16 x 16, row-major) b (16 x 8, col-major): Ampere's HMMA.16816, f16 products with f32 accumulation. Lane (g = lane / 4,
// t = lane % 4) holds a0 = (row g, k 2t..2t+1), a1 = (row g + 8, k 2t..), a2 = (row g, k 2t + 8..), a3 = (row g + 8, k 2t + 8..);
// b0 = (k 2t..2t+1, column g), b1 = (k 2t + 8.., column g); d0 d1 = (row g, columns 2t, 2t + 1), d2 d3 = (row g + 8, same columns)
__device__ __forceinline__ void mma_16816(float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) {
asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
: "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
#endif
// the output staging's row stride in floats (8 past the tile: fewer conflicts than none or 16)
template <int BM>
constexpr int MM_SLD = BM + 8;
@@ -286,8 +297,12 @@ template <int TYPE, int BM>
__global__ void __launch_bounds__(MM_THREADS) k_moe_mma(MoeW w, const BigJob* __restrict__ jobs, const int* __restrict__ njobs,
const __half* __restrict__ x, int ldx, float* __restrict__ out) {
using namespace nvcuda;
constexpr int WR = BM / 32, WT = MM_THREADS / 32 / WR, TT = 4 / WT, NP = BM / MM_PASS, LO = 16 * TT;
static_assert(WR * WT == MM_THREADS / 32 && TT * WT == 4 && NP * MM_PASS == BM, "warp layout");
constexpr int MM_NT = PF_JOB / 16; // token tiles of a job
constexpr int WR = BM / 32, WT = MM_THREADS / 32 / WR, TT = MM_NT / WT, NP = BM / MM_PASS, LO = 16 * TT;
static_assert(WR * WT == MM_THREADS / 32 && TT * WT == MM_NT && NP * MM_PASS == BM, "warp layout");
#if FNX_ARCH < 80
static_assert(PF_JOB == 64, "the m8n8k4 path is laid out for 64-row jobs");
#endif
const int ji = blockIdx.y;
if (ji >= *njobs) return;
__shared__ BigJob jb;
@@ -312,6 +327,19 @@ __global__ void __launch_bounds__(MM_THREADS) k_moe_mma(MoeW w, const BigJob* __
for (int b = 0; b < TT; ++b)
#pragma unroll
for (int j = 0; j < 8; ++j) acc[a][b][j] = 0.f;
#elif FNX_ARCH >= 80
// the warp's 32 rows x 16 TT tokens on HMMA.16816 (mma_16816): 2 row tiles of 16 x 2 TT token tiles of 8; lane (g, t4) loads the
// 32-bit operand pairs of the fragment layout above from the weight tile (rows) and the activation tile (tokens) in shared memory
const int lane = threadIdx.x & 31, g = lane >> 2, t4 = lane & 3;
const __half* arow = sw + (rg * 32 + g) * MM_LD + 2 * t4;
const __half* bcol = sx + (tb * 16 + g) * MM_LD + 2 * t4;
float acc[2][2 * TT][4];
#pragma unroll
for (int a = 0; a < 2; ++a)
#pragma unroll
for (int b = 0; b < 2 * TT; ++b)
#pragma unroll
for (int j = 0; j < 4; ++j) acc[a][b][j] = 0.f;
#else
// the warp's 32 rows x 16 TT tokens on HMMA.884 (mma_884): quad-pair q (lanes 4q..4q+3, 4q+16..4q+19) takes rows
// 16 (q >> 1) and tokens 8 TT (q & 1) of them, as 2 x TT tiles of 8 x 8; this lane loads row / token `rl` of a tile
@@ -364,6 +392,26 @@ __global__ void __launch_bounds__(MM_THREADS) k_moe_mma(MoeW w, const BigJob* __
for (int r = 0; r < 2; ++r)
#pragma unroll
for (int t = 0; t < TT; ++t) acc[r][t] = mma_wmma12(a[r], b[t], acc[r][t]);
#elif FNX_ARCH >= 80
uint32_t fa[2][4], fb[2 * TT][2];
#pragma unroll
for (int r = 0; r < 2; ++r) {
const __half* ap = arow + r * 16 * MM_LD + kk;
fa[r][0] = *reinterpret_cast<const uint32_t*>(ap);
fa[r][1] = *reinterpret_cast<const uint32_t*>(ap + 8 * MM_LD);
fa[r][2] = *reinterpret_cast<const uint32_t*>(ap + 8);
fa[r][3] = *reinterpret_cast<const uint32_t*>(ap + 8 * MM_LD + 8);
}
#pragma unroll
for (int u = 0; u < 2 * TT; ++u) {
const __half* bp = bcol + u * 8 * MM_LD + kk;
fb[u][0] = *reinterpret_cast<const uint32_t*>(bp);
fb[u][1] = *reinterpret_cast<const uint32_t*>(bp + 8);
}
#pragma unroll
for (int r = 0; r < 2; ++r)
#pragma unroll
for (int u = 0; u < 2 * TT; ++u) mma_16816(acc[r][u], fa[r][0], fa[r][1], fa[r][2], fa[r][3], fb[u][0], fb[u][1]);
#else
// 16 k: a 16-byte load per tile row (8 k) twice; each 4 k of it is one HMMA's operand
uint4 a[2][2], b[TT][2];
@@ -385,8 +433,6 @@ __global__ void __launch_bounds__(MM_THREADS) k_moe_mma(MoeW w, const BigJob* __
const uint4& B = b[t][k4 >> 1];
mma_884(acc[r][t], k4 & 1 ? A.z : A.x, k4 & 1 ? A.w : A.y, k4 & 1 ? B.z : B.x, k4 & 1 ? B.w : B.y);
}
}
}
#endif
}
}
@@ -408,6 +454,15 @@ __global__ void __launch_bounds__(MM_THREADS) k_moe_mma(MoeW w, const BigJob* __
#pragma unroll
for (int j = 0; j < 8; ++j) d[j] = acc[r][t][j];
}
#elif FNX_ARCH >= 80
// a lane's accumulators of a 16 x 8 tile: rows g and g + 8, columns 2 t4 and 2 t4 + 1 (mma_16816)
float* sob = so + (split ? tg * LO * SLD : tb * 16 * SLD) + rg * 32; // the warp's rows, its first token
#pragma unroll
for (int r = 0; r < 2; ++r)
#pragma unroll
for (int u = 0; u < 2 * TT; ++u)
#pragma unroll
for (int j = 0; j < 4; ++j) sob[(u * 8 + 2 * t4 + (j & 1)) * SLD + r * 16 + g + (j >> 1) * 8] = acc[r][u][j];
#else
// a lane's accumulators in an 8 x 8 tile (measured, tools/m884_probe.cu): rows r0 and r0 + 2, columns cb, cb + 1,
// cb + 4, cb + 5, in the order (r0,cb) (r0,cb+1) (r1,cb) (r1,cb+1) (r0,cb+4) (r0,cb+5) (r1,cb+4) (r1,cb+5)
src/split/src/moe_mma_iq.cuh
+229
View File
@@ -0,0 +1,229 @@
// The prompt path's weight tiles (moe_mma.cuh) for the 4-bit codebook and i-quant routed experts of the IQ3_S set: Q4_0 (what the Q2_0
// experts are repacked to at load, tensor_convert.hpp), IQ4_NL, IQ4_XS, IQ3_S, IQ3_XXS and IQ2_S. Each value is ggml's own dequantization
// (ggml-cuda/dequantize.cuh dequantize_*: the f32 product d * scale * grid * sign) rounded once to f16, so the f16 operands are the ones
// llama.cpp's cuBLAS path would have made. A thread takes one 32-value sub-block of a row per k-step (MM_TPR = 2 of them side by side
// in the step's 64 columns) like the Q5_1 and Q8_0 tiles; the 256-value types find theirs at (k0 % 256) / 32 + u in the block.
// First version: the raw bytes are read with 16- and 32-bit loads (the blocks of 110, 98 and 82 bytes are only 2-byte aligned) and the
// grids come from ggml's tables in global memory (cached); not yet tuned for speed.
__device__ __forceinline__ uint32_t mq_ld16(const uint8_t* p) { return *reinterpret_cast<const uint16_t*>(p); }
__device__ __forceinline__ uint32_t mq_ld32(const uint8_t* p) { return mq_ld16(p) | (mq_ld16(p + 2) << 16); }
__device__ __forceinline__ float mq_half(uint32_t h) { return __half2float(__ushort_as_half((uint16_t) h)); }
// 32 f32 values to f16 pairs and into the tile row
__device__ __forceinline__ void mq_put32(__half* dst, const float (&f)[32]) {
__half2 y[16];
#pragma unroll
for (int i = 0; i < 16; ++i) y[i] = __floats2half2_rn(f[2 * i], f[2 * i + 1]);
mm_store8(dst, y);
mm_store8(dst + 16, y + 8);
}
__device__ __forceinline__ void mq_zero32(__half* dst) {
float z[32];
#pragma unroll
for (int i = 0; i < 32; ++i) z[i] = 0.f;
mq_put32(dst, z);
}
// Q4_0 (value (nibble - 8) * d) and IQ4_NL (kvalues_iq4nl[nibble] * d): 18 bytes, low nibbles are values 0..15, high 16..31
template <bool NL>
struct MmTile4 {
uint16_t v[9];
bool live, item;
__device__ __forceinline__ void fetch(const char* we, size_t row_a, int nrows, int row0, int k0, int kc) {
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR, row = row0 + r;
item = u * 32 < kc;
live = item && row < nrows;
if (live) {
const uint8_t* b = (const uint8_t*) we + (size_t) row * row_a + (size_t) (k0 / 32 + u) * 18;
#pragma unroll
for (int l = 0; l < 9; ++l) v[l] = (uint16_t) mq_ld16(b + 2 * l);
}
}
__device__ __forceinline__ void put(int, int, __half* sw) const {
if (!item) return;
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR;
if (!live) {
mq_zero32(sw + r * MM_LD + u * 32);
return;
}
const float d = mq_half(v[0]);
float f[32];
#pragma unroll
for (int w = 0; w < 4; ++w) {
const uint32_t q = (uint32_t) v[1 + 2 * w] | ((uint32_t) v[2 + 2 * w] << 16);
#pragma unroll
for (int k = 0; k < 4; ++k) {
const int lo = (q >> (8 * k)) & 0xF, hi = (q >> (8 * k + 4)) & 0xF;
f[4 * w + k] = d * (float) (NL ? (int) kvalues_iq4nl[lo] : lo - 8);
f[16 + 4 * w + k] = d * (float) (NL ? (int) kvalues_iq4nl[hi] : hi - 8);
}
}
mq_put32(sw + r * MM_LD + u * 32, f);
}
};
template <> struct MmTile<GGML_TYPE_Q4_0> : MmTile4<false> {};
template <> struct MmTile<GGML_TYPE_IQ4_NL> : MmTile4<true> {};
// IQ4_XS: 136 bytes = d (f16), scales_h (u16), scales_l[4], qs[128]; sub-block ib has 16 qs bytes at 8 + 16 ib and a 6-bit scale
template <> struct MmTile<GGML_TYPE_IQ4_XS> {
uint32_t hd[2], q[4];
int ib;
bool live, item;
__device__ __forceinline__ void fetch(const char* we, size_t row_a, int nrows, int row0, int k0, int kc) {
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR, row = row0 + r;
item = u * 32 < kc;
live = item && row < nrows;
ib = (k0 % 256) / 32 + u;
if (live) {
const uint8_t* b = (const uint8_t*) we + (size_t) row * row_a + (size_t) (k0 / 256) * 136;
hd[0] = mq_ld32(b);
hd[1] = mq_ld32(b + 4);
#pragma unroll
for (int l = 0; l < 4; ++l) q[l] = mq_ld32(b + 8 + 16 * ib + 4 * l);
}
}
__device__ __forceinline__ void put(int, int, __half* sw) const {
if (!item) return;
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR;
if (!live) {
mq_zero32(sw + r * MM_LD + u * 32);
return;
}
const uint32_t sh = hd[0] >> 16, sl = (hd[1] >> (8 * (ib / 2))) & 0xFF;
const int ls = (int) (((sl >> (4 * (ib % 2))) & 0xF) | (((sh >> (2 * ib)) & 3) << 4));
const float d = mq_half(hd[0] & 0xFFFF) * (float) (ls - 32);
float f[32];
#pragma unroll
for (int l = 0; l < 4; ++l)
#pragma unroll
for (int k = 0; k < 4; ++k) {
const uint32_t byte = (q[l] >> (8 * k)) & 0xFF;
f[4 * l + k] = d * (float) kvalues_iq4nl[byte & 0xF];
f[16 + 4 * l + k] = d * (float) kvalues_iq4nl[byte >> 4];
}
mq_put32(sw + r * MM_LD + u * 32, f);
}
};
__device__ __forceinline__ void mq_signed4(float* y, uint32_t grid, float d, uint32_t signs, int s0) {
#pragma unroll
for (int j = 0; j < 4; ++j) y[j] = d * (float) ((grid >> (8 * j)) & 0xFF) * ((signs & kmask_iq2xs[s0 + j]) ? -1.f : 1.f);
}
// IQ3_S: 110 bytes = d, qs[64], qh[8], signs[32], scales[4]; sub-block ib: qs[8 ib..], qh[ib], signs[4 ib..], a 4-bit scale (1 + 2 s) d
template <> struct MmTile<GGML_TYPE_IQ3_S> {
uint32_t dq, qs[2], sg, qh, sc; // d (low 16 bits of dq), 8 qs bytes, 4 sign bytes, qh byte, scales byte
int ib;
bool live, item;
__device__ __forceinline__ void fetch(const char* we, size_t row_a, int nrows, int row0, int k0, int kc) {
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR, row = row0 + r;
item = u * 32 < kc;
live = item && row < nrows;
ib = (k0 % 256) / 32 + u;
if (live) {
const uint8_t* b = (const uint8_t*) we + (size_t) row * row_a + (size_t) (k0 / 256) * 110;
dq = mq_ld16(b);
qs[0] = mq_ld32(b + 2 + 8 * ib);
qs[1] = mq_ld32(b + 6 + 8 * ib);
qh = b[66 + ib];
sg = mq_ld32(b + 74 + 4 * ib);
sc = b[106 + ib / 2];
}
}
__device__ __forceinline__ void put(int, int, __half* sw) const {
if (!item) return;
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR;
if (!live) {
mq_zero32(sw + r * MM_LD + u * 32);
return;
}
const float d = mq_half(dq) * (float) (1 + 2 * ((sc >> (4 * (ib % 2))) & 0xF));
float f[32];
#pragma unroll
for (int il = 0; il < 4; ++il) {
const uint32_t q0 = (qs[(2 * il) / 4] >> (8 * ((2 * il) % 4))) & 0xFF, q1 = (qs[(2 * il + 1) / 4] >> (8 * ((2 * il + 1) % 4))) & 0xFF;
const uint32_t g1 = iq3s_grid[q0 | ((qh << (8 - 2 * il)) & 256)], g2 = iq3s_grid[q1 | ((qh << (7 - 2 * il)) & 256)];
const uint32_t sb = (sg >> (8 * il)) & 0xFF;
mq_signed4(f + 8 * il, g1, d, sb, 0);
mq_signed4(f + 8 * il + 4, g2, d, sb, 4);
}
mq_put32(sw + r * MM_LD + u * 32, f);
}
};
// IQ3_XXS: 98 bytes = d, qs[96]: 64 grid indices (8 per sub-block) then per sub-block a u32 of four 7-bit sign codes and a 4-bit scale at the top
template <> struct MmTile<GGML_TYPE_IQ3_XXS> {
uint32_t dq, qs[2], aux;
int ib;
bool live, item;
__device__ __forceinline__ void fetch(const char* we, size_t row_a, int nrows, int row0, int k0, int kc) {
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR, row = row0 + r;
item = u * 32 < kc;
live = item && row < nrows;
ib = (k0 % 256) / 32 + u;
if (live) {
const uint8_t* b = (const uint8_t*) we + (size_t) row * row_a + (size_t) (k0 / 256) * 98;
dq = mq_ld16(b);
qs[0] = mq_ld32(b + 2 + 8 * ib);
qs[1] = mq_ld32(b + 6 + 8 * ib);
aux = mq_ld32(b + 66 + 4 * ib);
}
}
__device__ __forceinline__ void put(int, int, __half* sw) const {
if (!item) return;
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR;
if (!live) {
mq_zero32(sw + r * MM_LD + u * 32);
return;
}
const float d = mq_half(dq) * (0.5f + (float) (aux >> 28)) * 0.5f;
float f[32];
#pragma unroll
for (int il = 0; il < 4; ++il) {
const uint32_t q0 = (qs[(2 * il) / 4] >> (8 * ((2 * il) % 4))) & 0xFF, q1 = (qs[(2 * il + 1) / 4] >> (8 * ((2 * il + 1) % 4))) & 0xFF;
const uint32_t sb = ksigns_iq2xs[(aux >> (7 * il)) & 127];
mq_signed4(f + 8 * il, iq3xxs_grid[q0], d, sb, 0);
mq_signed4(f + 8 * il + 4, iq3xxs_grid[q1], d, sb, 4);
}
mq_put32(sw + r * MM_LD + u * 32, f);
}
};
// IQ2_S: 82 bytes = d, qs[32] (grid index low bytes, 4 per sub-block), signs[32] (4 per sub-block), qh[8], scales[8]; 10-bit grid index of 8 values
template <> struct MmTile<GGML_TYPE_IQ2_S> {
uint32_t dq, qs, sg, qh, sc;
int ib;
bool live, item;
__device__ __forceinline__ void fetch(const char* we, size_t row_a, int nrows, int row0, int k0, int kc) {
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR, row = row0 + r;
item = u * 32 < kc;
live = item && row < nrows;
ib = (k0 % 256) / 32 + u;
if (live) {
const uint8_t* b = (const uint8_t*) we + (size_t) row * row_a + (size_t) (k0 / 256) * 82;
dq = mq_ld16(b);
qs = mq_ld32(b + 2 + 4 * ib);
sg = mq_ld32(b + 34 + 4 * ib);
qh = b[66 + ib];
sc = b[74 + ib];
}
}
__device__ __forceinline__ void put(int, int, __half* sw) const {
if (!item) return;
const int r = threadIdx.x / MM_TPR, u = threadIdx.x % MM_TPR;
if (!live) {
mq_zero32(sw + r * MM_LD + u * 32);
return;
}
float f[32];
#pragma unroll
for (int il = 0; il < 4; ++il) {
const uint64_t grid = iq2s_grid[((qs >> (8 * il)) & 0xFF) | ((qh << (8 - 2 * il)) & 0x300)];
const float d = mq_half(dq) * (0.5f + (float) ((sc >> (4 * (il / 2))) & 0xF)) * 0.25f;
const uint32_t sb = (sg >> (8 * il)) & 0xFF;
mq_signed4(f + 8 * il, (uint32_t) grid, d, sb, 0);
mq_signed4(f + 8 * il + 4, (uint32_t) (grid >> 32), d, sb, 4);
}
mq_put32(sw + r * MM_LD + u * 32, f);
}
};
src/split/src/session.cpp
+10 -11
View File
@@ -6,6 +6,7 @@
#include <random>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <stdexcept>
#include <string>
@@ -449,11 +450,6 @@ Finish Session::generate(int max_new, const Sampling& s, const std::function<voi
return d;
};
Finish fin = Finish::length;
// FNX_DUMP_FIRST_LOGITS=PREFIX (measurements): the full-vocabulary logits of the first generated token (row 0 of the first window) of each request,
// as float32 to PREFIX.<n>, n counting requests from 0 (Strata's STRATA_DUMP_FIRST_LOGITS twin, for first-token KL between the two engines)
static const char* const dump_prefix = std::getenv("FNX_DUMP_FIRST_LOGITS");
static int dump_n = 0;
bool dumped = dump_prefix == nullptr;
for (;;) {
if (stop && stop->load()) {
fin = Finish::cancel;
@@ -468,12 +464,15 @@ Finish Session::generate(int max_new, const Sampling& s, const std::function<voi
const int p = e_.pos();
const auto tv = std::chrono::steady_clock::now();
std::vector<int> o = e_.verify(rows);
if (!dumped) {
dumped = true;
const std::vector<float> lg = e_.last_logits(1);
if (std::FILE* f = std::fopen((std::string(dump_prefix) + "." + std::to_string(dump_n++)).c_str(), "wb")) {
std::fwrite(lg.data(), sizeof(float), lg.size(), f);
std::fclose(f);
if (produced == 0) { // FNX_DUMP_FIRST_LOGITS=PATH (diagnostics): the full-vocabulary logits of each request's first token to PATH.<n> (float32)
static const char* dump = std::getenv("FNX_DUMP_FIRST_LOGITS");
static int dump_n = 0;
if (dump) {
const std::vector<float> row = e_.last_logits(1);
const std::string path = std::string(dump) + "." + std::to_string(dump_n++);
std::FILE* f = std::fopen(path.c_str(), "wb");
if (!f || std::fwrite(row.data(), 4, row.size(), f) != row.size()) std::fprintf(stderr, "session: FNX_DUMP_FIRST_LOGITS: cannot write %s\n", path.c_str());
if (f) std::fclose(f);
}
}
st.verify_ms += ms_since(tv);
src/split/src/tensor_convert.hpp
+310
View File
@@ -0,0 +1,310 @@
// strata-split: load-time conversion of the small tensors of a GGUF whose types the engine does not read, to the types it does
// (what Unsloth's UD-Q4_K_XL file carries: Q8_0 for every dense matrix, the shared experts, the embeddings, the hyper-connection
// and PLE projections; F32 for the routers, the injections and the SSM gates). A set with other types for them (the ISTA-DAS IQ3_S
// set keeps K-quants, IQ4_XS and BF16 there) is dequantized to float and quantized again, once, at load. A tensor that already has the
// target type is never touched, so UD-Q4_K_XL and Q8_0 files load exactly as before. Pure host code, no CUDA.
//
// The dequantizers are ggml's reference ones (ggml-quants.c, dequantize_row_*), written out here because the engine does not link
// ggml's CPU library; tests/split/ check them against gguf-py on real tensors.
#pragma once
#include <cmath>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <vector>
#include "gguf.hpp"
namespace fnx {
namespace conv {
inline float h2f(uint16_t h) {
const uint32_t s = (h >> 15) & 1, e = (h >> 10) & 31, m = h & 1023;
float f;
if (e == 0) f = std::ldexp((float) m, -24);
else if (e == 31) f = m ? NAN : INFINITY;
else f = std::ldexp((float) (m + 1024), (int) e - 25);
return s ? -f : f;
}
inline uint16_t f2h(float x) { // round to nearest even, as ggml's GGML_FP32_TO_FP16
uint32_t b;
std::memcpy(&b, &x, 4);
const uint32_t s = (b >> 16) & 0x8000, e = (b >> 23) & 0xFF, m = b & 0x7FFFFF;
if (e == 0xFF) return (uint16_t) (s | 0x7C00 | (m ? 0x200 : 0));
int ne = (int) e - 127 + 15;
if (ne >= 31) return (uint16_t) (s | 0x7C00);
if (ne <= 0) {
if (ne < -10) return (uint16_t) s;
const uint32_t mm = (m | 0x800000) >> (1 - ne);
uint32_t r = mm >> 13;
const uint32_t rem = mm & 0x1FFF;
if (rem > 0x1000 || (rem == 0x1000 && (r & 1))) ++r;
return (uint16_t) (s | r);
}
uint32_t r = ((uint32_t) ne << 10) | (m >> 13);
const uint32_t rem = m & 0x1FFF;
if (rem > 0x1000 || (rem == 0x1000 && (r & 1))) ++r;
return (uint16_t) (s | r);
}
inline float bf2f(uint16_t b) {
const uint32_t u = (uint32_t) b << 16;
float f;
std::memcpy(&f, &u, 4);
return f;
}
inline const int8_t* kvalues_iq4nl() {
static const int8_t k[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113};
return k;
}
// the K-quants' 6-bit scale and min of sub-block j (ggml get_scale_min_k4)
inline void scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) {
if (j < 4) {
d = q[j] & 63;
m = q[j + 4] & 63;
} else {
d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
m = (q[j + 4] >> 4) | ((q[j] >> 6) << 4);
}
}
// n values (a whole number of blocks) of `type` at `src`, as floats; throws on a type it has no reader for
inline void dequantize(uint32_t type, const uint8_t* src, size_t n, float* y) {
switch (type) {
case T_F32: std::memcpy(y, src, n * 4); return;
case T_F16:
for (size_t i = 0; i < n; ++i) { uint16_t h; std::memcpy(&h, src + 2 * i, 2); y[i] = h2f(h); }
return;
case T_BF16:
for (size_t i = 0; i < n; ++i) { uint16_t h; std::memcpy(&h, src + 2 * i, 2); y[i] = bf2f(h); }
return;
case T_Q8_0:
for (size_t b = 0; b < n / 32; ++b) {
const uint8_t* p = src + b * 34;
uint16_t dh;
std::memcpy(&dh, p, 2);
const float d = h2f(dh);
for (int j = 0; j < 32; ++j) y[b * 32 + j] = d * (float) (int8_t) p[2 + j];
}
return;
case T_Q4_0:
for (size_t b = 0; b < n / 32; ++b) {
const uint8_t* p = src + b * 18;
uint16_t dh;
std::memcpy(&dh, p, 2);
const float d = h2f(dh);
for (int j = 0; j < 16; ++j) {
y[b * 32 + j] = d * (float) ((int) (p[2 + j] & 0xF) - 8);
y[b * 32 + j + 16] = d * (float) ((int) (p[2 + j] >> 4) - 8);
}
}
return;
case T_Q2_0:
for (size_t b = 0; b < n / 64; ++b) {
const uint8_t* p = src + b * 18;
uint16_t dh;
std::memcpy(&dh, p, 2);
const float d = h2f(dh);
for (int j = 0; j < 64; ++j) y[b * 64 + j] = d * (float) ((int) ((p[2 + j / 4] >> (2 * (j % 4))) & 3) - 1);
}
return;
case T_IQ4_NL:
for (size_t b = 0; b < n / 32; ++b) {
const uint8_t* p = src + b * 18;
uint16_t dh;
std::memcpy(&dh, p, 2);
const float d = h2f(dh);
for (int j = 0; j < 16; ++j) {
y[b * 32 + j] = d * (float) kvalues_iq4nl()[p[2 + j] & 0xF];
y[b * 32 + j + 16] = d * (float) kvalues_iq4nl()[p[2 + j] >> 4];
}
}
return;
case T_IQ4_XS:
for (size_t b = 0; b < n / 256; ++b) {
const uint8_t* p = src + b * 136;
uint16_t dh, sh;
std::memcpy(&dh, p, 2);
std::memcpy(&sh, p + 2, 2);
const uint8_t* sl = p + 4;
const uint8_t* qs = p + 8;
const float d = h2f(dh);
float* o = y + b * 256;
for (int ib = 0; ib < 8; ++ib) {
const int ls = ((sl[ib / 2] >> (4 * (ib % 2))) & 0xF) | (((sh >> (2 * ib)) & 3) << 4);
const float dl = d * (float) (ls - 32);
for (int j = 0; j < 16; ++j) {
o[j] = dl * (float) kvalues_iq4nl()[qs[j] & 0xF];
o[j + 16] = dl * (float) kvalues_iq4nl()[qs[j] >> 4];
}
o += 32;
qs += 16;
}
}
return;
case T_Q4_K:
for (size_t b = 0; b < n / 256; ++b) {
const uint8_t* p = src + b * 144;
uint16_t dh, mh;
std::memcpy(&dh, p, 2);
std::memcpy(&mh, p + 2, 2);
const float d = h2f(dh), mn = h2f(mh);
const uint8_t* sc = p + 4;
const uint8_t* q = p + 16;
float* o = y + b * 256;
int is = 0;
for (int j = 0; j < 256; j += 64) {
uint8_t s1, m1, s2, m2;
scale_min_k4(is, sc, s1, m1);
scale_min_k4(is + 1, sc, s2, m2);
const float d1 = d * s1, mm1 = mn * m1, d2 = d * s2, mm2 = mn * m2;
for (int l = 0; l < 32; ++l) *o++ = d1 * (float) (q[l] & 0xF) - mm1;
for (int l = 0; l < 32; ++l) *o++ = d2 * (float) (q[l] >> 4) - mm2;
q += 32;
is += 2;
}
}
return;
case T_Q5_K:
for (size_t b = 0; b < n / 256; ++b) {
const uint8_t* p = src + b * 176;
uint16_t dh, mh;
std::memcpy(&dh, p, 2);
std::memcpy(&mh, p + 2, 2);
const float d = h2f(dh), mn = h2f(mh);
const uint8_t* sc = p + 4;
const uint8_t* qh = p + 16;
const uint8_t* ql = p + 48;
float* o = y + b * 256;
int is = 0;
uint8_t u1 = 1, u2 = 2;
for (int j = 0; j < 256; j += 64) {
uint8_t s1, m1, s2, m2;
scale_min_k4(is, sc, s1, m1);
scale_min_k4(is + 1, sc, s2, m2);
const float d1 = d * s1, mm1 = mn * m1, d2 = d * s2, mm2 = mn * m2;
for (int l = 0; l < 32; ++l) *o++ = d1 * (float) ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - mm1;
for (int l = 0; l < 32; ++l) *o++ = d2 * (float) ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - mm2;
ql += 32;
is += 2;
u1 = (uint8_t) (u1 << 2);
u2 = (uint8_t) (u2 << 2);
}
}
return;
case T_Q6_K:
for (size_t b = 0; b < n / 256; ++b) {
const uint8_t* p = src + b * 210;
const uint8_t* ql = p;
const uint8_t* qh = p + 128;
const int8_t* sc = (const int8_t*) (p + 192);
uint16_t dh;
std::memcpy(&dh, p + 208, 2);
const float d = h2f(dh);
float* o = y + b * 256;
for (int h = 0; h < 256; h += 128) {
for (int l = 0; l < 32; ++l) {
const int is = l / 16;
const int q1 = (int) ((ql[l] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
const int q2 = (int) ((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
const int q3 = (int) ((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32;
const int q4 = (int) ((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32;
o[l] = d * (float) sc[is] * (float) q1;
o[l + 32] = d * (float) sc[is + 2] * (float) q2;
o[l + 64] = d * (float) sc[is + 4] * (float) q3;
o[l + 96] = d * (float) sc[is + 6] * (float) q4;
}
o += 128;
ql += 64;
qh += 32;
sc += 8;
}
}
return;
}
throw std::runtime_error("tensor_convert: no reader for ggml type " + std::to_string(type));
}
// Q8_0 as ggml's quantize_row_q8_0_ref: per 32 values d = amax / 127 (stored as f16), q = round(x / d)
inline void quantize_q8_0(const float* x, size_t n, uint8_t* dst) {
for (size_t b = 0; b < n / 32; ++b) {
float amax = 0.f;
for (int j = 0; j < 32; ++j) amax = std::fmax(amax, std::fabs(x[b * 32 + j]));
const float d = amax / 127.0f, id = d != 0.f ? 1.0f / d : 0.f;
const uint16_t dh = f2h(d);
std::memcpy(dst + b * 34, &dh, 2);
for (int j = 0; j < 32; ++j) dst[b * 34 + 2 + j] = (uint8_t) (int8_t) std::lround(x[b * 32 + j] * id);
}
}
// the type the engine reads a tensor of this name as (what UD-Q4_K_XL carries), or 0xFFFFFFFF for a tensor that is left alone
inline uint32_t wanted_type(const std::string& name) {
auto ends = [&](const char* s) { const size_t l = std::strlen(s); return name.size() >= l && name.compare(name.size() - l, l, s) == 0; };
static const char* q8[] = {".attn_gate.weight", ".attn_k.weight", ".attn_output.weight", ".attn_q.weight", ".attn_qkv.weight", ".attn_v.weight",
".ffn_gate_shexp.weight", ".ffn_up_shexp.weight", ".ffn_down_shexp.weight", ".hc_attn_down.weight",
".hc_attn_up.weight", ".hc_ffn_down.weight", ".hc_ffn_up.weight", ".ple_key.weight", ".ple_value.weight",
".ssm_out.weight"};
static const char* f32[] = {".ffn_gate_inp.weight", ".ffn_gate_inp_shexp.weight", ".hc_attn_inject.weight", ".hc_ffn_inject.weight",
".ssm_alpha.weight", ".ssm_beta.weight", ".ple_conv1d.weight"};
for (const char* s : q8)
if (ends(s)) return T_Q8_0;
for (const char* s : f32)
if (ends(s)) return T_F32;
if (name == "output.weight" || name == "output_hc_down.weight" || name == "output_hc_up.weight" || name == "token_embd.weight") return T_Q8_0;
return 0xFFFFFFFFu;
}
// `t` converted to `target` (Q8_0 or F32) into `out`; the Tensor then points at `out`
inline void convert(Tensor& t, uint32_t target, std::vector<uint8_t>& out) {
size_t n = 1;
for (int d = 0; d < 4; ++d) n *= (size_t) t.ne[d];
const TypeInfo from = type_info(t.type);
if (n % (size_t) from.block) throw std::runtime_error(t.name + ": not a whole number of blocks");
if (target == T_Q8_0 && t.ne[0] % 32) throw std::runtime_error(t.name + ": rows are not a whole number of Q8_0 blocks");
std::vector<float> f(n);
dequantize(t.type, t.data, n, f.data());
if (target == T_F32) {
out.resize(n * 4);
std::memcpy(out.data(), f.data(), n * 4);
} else if (target == T_Q8_0) {
out.resize(n / 32 * 34);
quantize_q8_0(f.data(), n, out.data());
} else {
throw std::runtime_error("tensor_convert: target type " + std::to_string(target));
}
t.type = target;
t.data = out.data();
t.nbytes = out.size();
t.shard = -1; // not in a GGUF file any more: readers that stage from the file by offset must take the bytes at t.data
t.offset = 0;
}
// Strata's Q2_0 (64 values in 18 bytes: f16 d, two bits per value, value j in byte j / 4 at bit 2 (j % 4), meaning (code - 1) * d) as two
// Q4_0 blocks (32 values in 18 bytes: f16 d, nibble i of the low half = value i, of the high half = value i + 16, meaning (nibble - 8) * d).
// Lossless: nibble = code + 7 gives (code + 7 - 8) * d, the same product. The engine's Q4_0 kernels then read it (the draft layer's
// experts.bin goes the same way, model.cpp Mtp::experts). Twice the bytes.
inline void repack_q2_0_to_q4_0(Tensor& t, std::vector<uint8_t>& out) {
if (t.ne[0] % 64) throw std::runtime_error(t.name + ": rows are not a whole number of Q2_0 blocks");
const size_t nb = t.nbytes / 18;
out.resize(nb * 36);
for (size_t b = 0; b < nb; ++b) {
const uint8_t* in = t.data + b * 18;
auto code = [&](int j) { return (in[2 + j / 4] >> (2 * (j % 4))) & 3; };
for (int h = 0; h < 2; ++h) {
uint8_t* o = out.data() + (2 * b + (size_t) h) * 18;
o[0] = in[0];
o[1] = in[1];
for (int i = 0; i < 16; ++i) o[2 + i] = (uint8_t) ((code(32 * h + i) + 7) | ((code(32 * h + 16 + i) + 7) << 4));
}
}
t.type = T_Q4_0;
t.data = out.data();
t.nbytes = out.size();
t.shard = -1;
t.offset = 0;
}
} // namespace conv
} // namespace fnx
src/split/tools/batch_check.hpp
+173
View File
@@ -0,0 +1,173 @@
// Batched windows (Engine::step) against one sequence at a time: each conversation's GEN greedy tokens alone (one row a
// window, no drafts), then the conversations on lanes decoding together, one row per lane a window; each lane's tokens
// must equal its conversation's alone. For 1, 2, 3 and 6 lanes (prompts reused round-robin): the window time and the
// total tokens a second. With FNX_BATCH_DRAFTS (lookup or 1) the same again with drafts in the batched windows
// (BatchDecoder: several rows a lane, each lane cut back to what it accepts, lanes leaving as they finish), greedy and
// equal to alone; then a sampled run at 3 lanes, which must only run (its tokens differ by chance). bin/batch_test and
// the lab's "batch" command.
#pragma once
#include <algorithm>
#include <chrono>
#include <cstdio>
#include <random>
#include <vector>
#include "../src/batch.hpp"
#include "../src/engine.hpp"
#include "../src/session.hpp"
inline int batch_check(const fnx::Model& m, fnx::Engine& e, const fnx::Mtp& mtp, int gen, const std::vector<std::vector<int>>& prompts) {
e.mtp_attach(mtp);
const std::vector<int> eos{248044, 248046};
fnx::Sampling greedy;
// alone: one row a window
std::vector<std::vector<int>> solo(prompts.size());
{
fnx::Session s(e, 4, 7, eos);
s.speculate = false;
for (size_t i = 0; i < prompts.size(); ++i) {
e.reset();
fnx::GenStats st;
s.read(prompts[i], st);
s.generate(gen, greedy, [&](int t) { solo[i].push_back(t); }, nullptr, st);
}
}
e.reset();
std::vector<int> lanes{0};
while ((int) lanes.size() < 6) lanes.push_back(e.add_lane());
// a lane's tokens against its conversation's alone: every token of the reference up to its first end-of-reply (past
// it the plain batch keeps decoding), none missing: a lane that ends early or emits nothing does not match
auto matches = [&](const std::vector<int>& got, const std::vector<int>& ref) {
size_t want = ref.size();
for (size_t j = 0; j < ref.size(); ++j)
if (ref[j] == 248044 || ref[j] == 248046) {
want = j + 1;
break;
}
return got.size() >= want && std::equal(ref.begin(), ref.begin() + (long) want, got.begin());
};
bool ok = true;
for (int n : {1, 2, 3, 6}) {
std::vector<int> use(lanes.begin(), lanes.begin() + n), x(n);
std::vector<std::vector<int>> got(n);
for (int r = 0; r < n; ++r) {
e.use_lane(use[r]);
e.reset();
fnx::Session s(e, 4, 7, eos, 248045, 0, use[r]);
fnx::GenStats st;
s.read(prompts[r % prompts.size()], st);
x[r] = prompts[r % prompts.size()].back();
}
double ms = 0;
int windows = 0;
while ((int) got[0].size() < gen) {
const auto t0 = std::chrono::steady_clock::now();
const std::vector<int> o = e.step(use, x);
ms += std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
++windows;
for (int r = 0; r < n; ++r) {
got[r].push_back(o[r]);
x[r] = o[r];
}
}
int same = 0;
for (int r = 0; r < n; ++r) same += matches(got[r], solo[r % prompts.size()]);
std::printf("%d lanes: %d windows, %.2f ms a window (first one captures its graph), %.1f tokens/s in all; %d of %d lanes match alone\n",
n, windows, ms / windows, n * windows * 1000.0 / ms, same, n);
ok = ok && same == n;
}
std::printf("%s\n", ok ? "batched windows give each lane its tokens alone" : "MISMATCH");
if (fnx::batch_drafts() == fnx::BatchDrafts::off) return ok ? 0 : 1;
bool dok = true;
fnx::BatchDecoder dec(e, 4, 15, eos, fnx::batch_drafts());
// the lanes on their prompts, each lane's next row its prompt's last token
auto start = [&](int n, const fnx::Sampling& s, std::vector<fnx::BatchLane>& bl) {
bl.assign((size_t) n, fnx::BatchLane{});
for (int r = 0; r < n; ++r) {
e.use_lane(lanes[(size_t) r]);
e.reset();
fnx::Session sess(e, 4, 7, eos, 248045, 0, lanes[(size_t) r]);
fnx::GenStats st;
const std::vector<int>& p = prompts[r % prompts.size()];
sess.read(p, st);
bl[(size_t) r].start(lanes[(size_t) r], s, std::mt19937_64(1000 + r), p);
}
};
// windows until every lane has gen tokens or ended its reply: the tokens per lane, windows, rows, ms
auto run = [&](std::vector<fnx::BatchLane>& bl, std::vector<std::vector<int>>& got, int& windows, int& rows, double& ms) {
const int n = (int) bl.size();
got.assign((size_t) n, {});
std::vector<bool> ended((size_t) n, false);
windows = rows = 0;
ms = 0;
for (;;) {
std::vector<fnx::BatchLane*> ls;
std::vector<int> idx, limit;
for (int r = 0; r < n; ++r)
if (!ended[(size_t) r] && (int) got[(size_t) r].size() < gen) {
ls.push_back(&bl[(size_t) r]);
idx.push_back(r);
limit.push_back(gen - (int) got[(size_t) r].size());
}
if (ls.empty()) break;
fnx::BatchDecoder::fit(ls);
const auto t0 = std::chrono::steady_clock::now();
const auto out = dec.window(ls, limit, [&](int i, int t) { got[(size_t) idx[(size_t) i]].push_back(t); });
ms += std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
++windows;
rows += dec.last_rows;
for (size_t i = 0; i < out.size(); ++i)
if (out[i].eos) ended[(size_t) idx[i]] = true;
}
};
const char* mode = fnx::batch_drafts() == fnx::BatchDrafts::lookup ? "lookup drafts" : "lookup drafts and the draft layer's chain";
for (int n : {1, 2, 3, 6}) {
std::vector<fnx::BatchLane> bl;
start(n, greedy, bl);
std::vector<std::vector<int>> got;
int windows = 0, rows = 0;
double ms = 0;
run(bl, got, windows, rows, ms);
int same = 0;
size_t tokens = 0;
long offered = 0, accepted = 0;
for (int r = 0; r < n; ++r) {
same += matches(got[(size_t) r], solo[r % prompts.size()]);
tokens += got[(size_t) r].size();
offered += bl[(size_t) r].offered;
accepted += bl[(size_t) r].accepted;
}
std::printf("%d lanes with %s: %d windows, %.1f rows and %.2f ms a window, %.1f tokens/s in all, drafts %ld of %ld; %d of %d "
"lanes match alone\n",
n, mode, windows, (double) rows / windows, ms / windows, tokens * 1000.0 / ms, accepted, offered, same, n);
dok = dok && same == n;
}
{ // sampled (prod's values): it must run; nothing to compare
fnx::Sampling s;
s.temperature = 0.6f;
s.top_p = 0.95f;
s.pen_present = 1.0f;
std::vector<fnx::BatchLane> bl;
start(3, s, bl);
std::vector<std::vector<int>> got;
int windows = 0, rows = 0;
double ms = 0;
run(bl, got, windows, rows, ms);
bool in_vocab = true;
size_t tokens = 0;
for (const auto& g : got) {
tokens += g.size();
for (int t : g) in_vocab = in_vocab && t >= 0 && t < m.vocab;
}
std::printf("3 lanes sampled (temperature 0.6, top_p 0.95, presence 1.0): %d windows, %.1f rows a window, %.1f tokens/s in all%s\n",
windows, (double) rows / windows, tokens * 1000.0 / ms, in_vocab ? "" : ": A TOKEN OUT OF THE VOCABULARY");
dok = dok && in_vocab;
}
for (int i : lanes) {
e.use_lane(i);
e.reset();
}
std::printf("%s\n", dok ? "batched windows with drafts give each lane its tokens alone" : "MISMATCH with drafts");
ok = ok && dok;
return ok ? 0 : 1;
}
src/split/tools/convert_dump.cpp
+48
View File
@@ -0,0 +1,48 @@
// Writes what the loader makes of tensors of a GGUF (src/tensor_convert.hpp): for each tensor named on the command line, the
// converted tensor's values as float32 to OUTDIR/<name>.f32 plus a line "name type_after nbytes" on stdout. tests/split/convert_check.py
// compares those with gguf-py's own dequantization of the same tensors in the file. CPU only.
// bin/convert_dump MODEL.gguf OUTDIR name [name ...] (a name of "all" dumps every tensor the loader converted)
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>
#include "../src/gguf.hpp"
#include "../src/tensor_convert.hpp"
int main(int argc, char** argv) {
if (argc < 4) {
std::fprintf(stderr, "usage: convert_dump MODEL.gguf OUTDIR name [name ...]\n");
return 2;
}
try {
fnx::Gguf g(argv[1]);
std::printf("converted %d\n", g.converted());
std::vector<std::string> names;
for (int i = 3; i < argc; ++i) {
if (std::string(argv[i]) == "all") {
for (const auto& kv : g.tensors())
if (fnx::conv::wanted_type(kv.first) != 0xFFFFFFFFu) names.push_back(kv.first);
} else {
names.push_back(argv[i]);
}
}
for (const std::string& n : names) {
const fnx::Tensor& t = g.t(n);
size_t cnt = 1;
for (int d = 0; d < 4; ++d) cnt *= (size_t) t.ne[d];
std::vector<float> f(cnt);
fnx::conv::dequantize(t.type, t.data, cnt, f.data());
const std::string path = std::string(argv[2]) + "/" + n + ".f32";
FILE* fp = std::fopen(path.c_str(), "wb");
if (!fp) throw std::runtime_error("cannot write " + path);
std::fwrite(f.data(), 4, cnt, fp);
std::fclose(fp);
std::printf("%s %u %zu\n", n.c_str(), t.type, t.nbytes);
}
} catch (const std::exception& e) {
std::fprintf(stderr, "convert_dump: %s\n", e.what());
return 1;
}
return 0;
}
src/split/tools/fnx_run.cpp
+1375
View File
File diff suppressed because it is too large. Load diff
src/split/tools/moe_tile_test.cu
+135
View File
@@ -0,0 +1,135 @@
// The prompt kernel's weight tiles (src/moe_mma.cuh, src/moe_mma_iq.cuh) against the file: for one expert of each routed-expert type a
// GGUF holds, runs the tiles' fetch() and put() over the whole matrix on the GPU and writes the f16 values they put into shared memory.
// tests/split/tile_check.py compares those with gguf-py's dequantization of the same expert (rounded once to f16).
// bin/moe_tile_test MODEL.gguf OUTDIR [expert ...] writes OUTDIR/<tensor>.e<expert>.f16 and prints "name type nrows ncols expert"
#include "kernels.hpp"
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <stdexcept>
#include <string>
#include <vector>
#include "common.cuh"
#include "ggml.h"
#include "gguf.hpp"
namespace fnx {
namespace {
__device__ __forceinline__ float silu_(float x) { return x / (1.0f + __expf(-x)); }
#include "moe_w.cuh"
#include "moe_mma.cuh"
#include "moe_mma_iq.cuh"
template <int TYPE>
__global__ void __launch_bounds__(MM_THREADS) k_tiles(const char* __restrict__ we, size_t row_a, int nrows, int ncols, __half* __restrict__ out) {
extern __shared__ __align__(32) __half smem[];
constexpr int NP = MM_BM / MM_PASS;
const int row0 = blockIdx.x * MM_BM;
for (int k0 = 0; k0 < ncols; k0 += MM_KC) {
const int kc = min(MM_KC, ncols - k0);
MmTile<TYPE> wt[NP];
#pragma unroll
for (int p = 0; p < NP; ++p) wt[p].fetch(we, row_a, nrows, row0 + p * MM_PASS, k0, kc);
#pragma unroll
for (int p = 0; p < NP; ++p) wt[p].put(k0, kc, smem + p * MM_PASS * MM_LD);
__syncthreads();
for (int e = threadIdx.x; e < MM_BM * kc; e += MM_THREADS) {
const int r = e / kc, c = e % kc, row = row0 + r;
if (row < nrows) out[(size_t) row * ncols + k0 + c] = smem[r * MM_LD + c];
}
__syncthreads();
}
}
template <int TYPE>
void run(const char* d_we, size_t row_a, int nrows, int ncols, __half* d_out) {
k_tiles<TYPE><<<(nrows + MM_BM - 1) / MM_BM, MM_THREADS, MM_BM * MM_LD * sizeof(__half)>>>(d_we, row_a, nrows, ncols, d_out);
}
} // namespace
} // namespace fnx
using namespace fnx;
int main(int argc, char** argv) {
if (argc < 3) {
std::fprintf(stderr, "usage: moe_tile_test MODEL.gguf OUTDIR [expert ...]\n");
return 2;
}
try {
Gguf g(argv[1]);
std::vector<int> experts;
for (int i = 3; i < argc; ++i) experts.push_back(std::atoi(argv[i]));
if (experts.empty()) experts = {5, 511};
std::map<uint32_t, std::string> pick; // the first routed-expert tensor of each type, gate/up first
for (const auto& kv : g.tensors()) {
const std::string& n = kv.first;
if (n.find("_exps.weight") == std::string::npos) continue;
const uint32_t ty = kv.second.type;
if (!pick.count(ty)) pick[ty] = n;
}
// also a down matrix of each type that has both, as its K (640) is not a multiple of 256
for (const auto& kv : g.tensors()) {
const std::string& n = kv.first;
if (n.find("ffn_down_exps.weight") == std::string::npos) continue;
const std::string key = "down:" + std::to_string(kv.second.type);
static std::map<std::string, bool> seen;
if (!seen[key]) {
seen[key] = true;
pick[1000 + kv.second.type] = n;
}
}
for (const auto& pk : pick) {
const Tensor& t = g.t(pk.second);
const int ncols = (int) t.ne[0], nrows = (int) t.ne[1];
const size_t row_a = t.row_bytes();
for (int e : experts) {
if (e >= t.ne[2]) continue;
const char* we = (const char*) t.data + (size_t) e * row_a * nrows;
char* d_we;
__half* d_out;
if (cudaMalloc(&d_we, row_a * nrows) || cudaMalloc(&d_out, (size_t) nrows * ncols * 2)) throw std::runtime_error("cudaMalloc");
cudaMemcpy(d_we, we, row_a * nrows, cudaMemcpyHostToDevice);
cudaMemset(d_out, 0xFF, (size_t) nrows * ncols * 2);
switch (t.type) {
case T_Q4_0: run<GGML_TYPE_Q4_0>(d_we, row_a, nrows, ncols, d_out); break;
case T_IQ4_NL: run<GGML_TYPE_IQ4_NL>(d_we, row_a, nrows, ncols, d_out); break;
case T_IQ4_XS: run<GGML_TYPE_IQ4_XS>(d_we, row_a, nrows, ncols, d_out); break;
case T_IQ3_S: run<GGML_TYPE_IQ3_S>(d_we, row_a, nrows, ncols, d_out); break;
case T_IQ3_XXS: run<GGML_TYPE_IQ3_XXS>(d_we, row_a, nrows, ncols, d_out); break;
case T_IQ2_S: run<GGML_TYPE_IQ2_S>(d_we, row_a, nrows, ncols, d_out); break;
case T_Q4_K: run<GGML_TYPE_Q4_K>(d_we, row_a, nrows, ncols, d_out); break;
case T_Q5_K: run<GGML_TYPE_Q5_K>(d_we, row_a, nrows, ncols, d_out); break;
case T_Q5_1: run<GGML_TYPE_Q5_1>(d_we, row_a, nrows, ncols, d_out); break;
case T_Q8_0: run<GGML_TYPE_Q8_0>(d_we, row_a, nrows, ncols, d_out); break;
default: std::printf("skip %s type %u\n", t.name.c_str(), t.type); cudaFree(d_we); cudaFree(d_out); continue;
}
if (cudaDeviceSynchronize() != cudaSuccess) throw std::runtime_error(std::string("kernel: ") + cudaGetErrorString(cudaGetLastError()));
std::vector<__half> h((size_t) nrows * ncols);
cudaMemcpy(h.data(), d_out, h.size() * 2, cudaMemcpyDeviceToHost);
cudaFree(d_we);
cudaFree(d_out);
const std::string path = std::string(argv[2]) + "/" + t.name + ".e" + std::to_string(e) + ".f16";
FILE* fp = std::fopen(path.c_str(), "wb");
if (!fp) throw std::runtime_error("cannot write " + path);
std::fwrite(h.data(), 2, h.size(), fp);
std::fclose(fp);
std::printf("%s %u %d %d %d\n", t.name.c_str(), t.type, nrows, ncols, e);
}
}
} catch (const std::exception& e) {
std::fprintf(stderr, "moe_tile_test: %s\n", e.what());
return 1;
}
return 0;
}
src/split/tools/ref_tokens.hpp
+42
View File
@@ -0,0 +1,42 @@
// A prompt file's token ids: ref/*.bin (fnx_ref's dump format, ref/refio.py), whose first record is one named
// "tokens" holding the ids as float32 (later records, per-layer dumps, are not read). Every read is checked and the
// ids must be whole numbers in [0, 2^24): a short or foreign file is an error here, never a run on a truncated or
// zero-padded prompt. It throws rather than exits, so a lab worker (bin/fnx_lab) survives a wrong path.
#pragma once
#include <cstdint>
#include <cstdio>
#include <stdexcept>
#include <string>
#include <vector>
inline std::vector<int> ref_tokens(const char* path) {
std::FILE* f = std::fopen(path, "rb");
if (!f) throw std::runtime_error(std::string(path) + ": cannot open");
auto bad = [&](const std::string& why) {
std::fclose(f);
return std::runtime_error(std::string(path) + ": " + why);
};
uint32_t nl = 0, nd = 0;
if (std::fread(&nl, 4, 1, f) != 1) throw bad("empty file");
if (nl == 0 || nl > 256) throw bad("not a prompt file (record name length " + std::to_string(nl) + ")");
std::string name(nl, '\0');
if (std::fread(name.data(), 1, nl, f) != nl) throw bad("short record name");
if (name != "tokens") throw bad("the first record is '" + name + "', not tokens");
if (std::fread(&nd, 4, 1, f) != 1 || nd != 1) throw bad("the tokens record is not one-dimensional");
int64_t ne = 0;
if (std::fread(&ne, 8, 1, f) != 1 || ne < 1 || ne > 0x7fffffffLL) throw bad("no valid token count");
std::vector<float> v((size_t) ne);
if (std::fread(v.data(), 4, v.size(), f) != v.size()) throw bad("short file: fewer tokens than its header says");
std::fclose(f);
std::vector<int> ids;
ids.reserve(v.size());
for (size_t i = 0; i < v.size(); ++i) {
const float x = v[i];
// float32 holds every integer below 2^24 exactly (the vocabulary is ~248k)
if (!(x >= 0.f && x < 16777216.f) || x != (float) (int) x)
throw std::runtime_error(std::string(path) + ": value " + std::to_string(x) + " at " + std::to_string(i) +
" is not a token id");
ids.push_back((int) x);
}
return ids;
}
strata_cli/launch/cmd.py
+7 -1
View File
@@ -105,6 +105,10 @@ def add_parsers(sp: argparse._SubParsersAction) -> list[argparse.ArgumentParser]
help="which drafter guesses words ahead to make answers faster: auto (recommended: Strata picks from what it measured on cards like yours; "
"that is the MTP today), mtp, dflash2 (Qwen3.8-Flash-Next only; downloads a 1.2 GB drafter once) or off. If DFlash2 cannot run "
"(not enough graphics memory, or a model it is not for) the MTP is used and the start says why")
g.add_argument("--split-engine", type=_choice("a split-engine setting", S.SPLIT_MODES), metavar="auto|on|off",
help="NVIDIA, two or more cards: the engine that splits every layer across the cards (faster answers). auto (the default) uses it by itself "
"where it is supported and starts the normal layer split everywhere else; on asks for it and says in one line why it cannot run; "
"off never uses it. Your choice is saved with the preset")
g.add_argument("--speed", type=_choice("a speed", S.SPEED_MODES), metavar="faster|full",
help="DeepSeek: faster uses 5 of 6 experts per word (a small quality change, usually not noticeable; the default on a PC with 16 GB "
"of graphics memory or less, and on a single Radeon RX 9000 / AI PRO R9700 card, or several that hold half of the experts), full uses all 6")
@@ -244,6 +248,8 @@ def settings_from_args(a: argparse.Namespace, catalog=None) -> tuple[S.LaunchSet
s.mtp = a.mtp
if a.drafting:
s.drafting = a.drafting
if getattr(a, "split_engine", None):
s.split_engine = a.split_engine
if getattr(a, "full_quality", False):
if getattr(a, "speed", None) == "faster":
notices.append(S.Notice("warn", "--full-quality and --speed faster were both given; full quality is used.", ""))
@@ -865,7 +871,7 @@ def run_launch(a: argparse.Namespace) -> int:
notices.append(S.Notice("warn", backend_note, ""))
if not machine:
_interrupted_move_notice()
have_choices = bool(a.preset or a.model or a.model_path or a.context or a.vision or a.parallel or a.access or a.tune or a.drafting or getattr(a, "speed", None) or getattr(a, "full_quality", False))
have_choices = bool(a.preset or a.model or a.model_path or a.context or a.vision or a.parallel or a.access or a.tune or a.drafting or getattr(a, "split_engine", None) or getattr(a, "speed", None) or getattr(a, "full_quality", False))
interactive = sys.stdin.isatty() and sys.stdout.isatty() and not a.no_tui and not machine and not a.plan
plain_questions = False
no_screens_package = None
strata_cli/launch/runner.py
+22 -12
View File
@@ -545,6 +545,8 @@ def derive_config(base: dict, s: S.LaunchSettings, *, bool_flags: frozenset[str]
cfg.pop("api_key", None)
if s.drafting: # strata_control/drafting.py: serve/server.py applies it at every start
cfg["drafting"] = s.drafting
if s.split_engine: # strata_control/split_engine.py: read by apply_split_engine; "off" survives every update
cfg["split_engine"] = s.split_engine
cfg.setdefault("exe", "component:strata-engine")
cfg["open_browser"] = False # strata launch never opens a browser by itself (only `--open` does, after the start)
if isinstance(cfg.get("args"), list):
@@ -1529,32 +1531,40 @@ class Runner:
return out
def apply_split_engine(self, cfg: dict) -> None:
"""The opt-in all-cards split engine (docs/SPLIT_ENGINE.md): only when the run config says "split_engine": true (or a tester sets
STRATA_SPLIT_ENGINE=1), only for a Qwen model on this PC's NVIDIA cards that passes strata_control/split_engine.py's rule. Anything
else leaves the config as it is, with one plain line when the person had asked for it. It never stops a launch."""
"""The all-cards split engine (docs/SPLIT_ENGINE.md): by itself ("auto", the default) for a Qwen model on this PC's NVIDIA cards that passes
strata_control/split_engine.py's rule, never when the run config says "split_engine": "off" (`--split-engine off`, STRATA_SPLIT_ENGINE=0).
In auto a "no" is silent and the config stays as it is; with "on" it gets one plain line. It never stops a launch."""
from strata_control import split_engine as SE
if cfg.get("arch") == "deepseek41" or not SE.wanted(cfg):
if cfg.get("arch") == "deepseek41" or SE.mode(cfg) == "off":
return
from strata_control import engine_component as EC
def exe_for(n: int):
def exe_for(n: int, cc: int | None = None, arch: str | None = None):
exe = cfg.get("exe")
try:
real = EC.expand(exe, cfg, full=False) if EC.is_reference(exe) else exe
except Exception: # noqa: BLE001 - no engine to look beside: no split
return None
p = Path(str(real)).parent / f"strata-split-{n}"
return str(p) if p.is_file() else None
# one build per architecture (FNX_ARCH is compile time): NVIDIA strata-split-N-sm<cc> then strata-split-N; AMD strata-split-N-<arch> only
for name in ((f"strata-split-{n}-{arch}",) if arch else ((f"strata-split-{n}-sm{cc}",) if cc else ()) + (f"strata-split-{n}",)):
p = Path(str(real)).parent / name
if p.is_file():
return str(p)
return None
try:
decision = SE.apply_to_config(cfg, exe_for_cards=exe_for, allow_untested=os.environ.get("STRATA_SPLIT_ALLOW_UNTESTED") == "1")
except Exception as e: # noqa: BLE001 - never stop a launch over an option
text = f"The all-cards split engine was left out ({type(e).__name__}). Using the normal layer split."
self.notices.append(text)
self.emit("config", text, None, level="warn")
if SE.mode(cfg) == "on":
text = f"The all-cards split engine was left out ({type(e).__name__}). Using the normal layer split."
self.notices.append(text)
self.emit("config", text, None, level="warn")
return
self.notices.append(decision.line)
self.emit("config", decision.line, None)
if decision.quiet:
return
line = decision.line + (" Turn it off with: strata launch --split-engine off." if decision.use_split else "")
self.notices.append(line)
self.emit("config", line, None)
def say_remembered_cards(self, base: dict) -> None:
"""A launch that names no cards keeps the cards the person chose earlier (`--gpu 0` is remembered by the model's settings). On a PC with more
strata_cli/launch/settings.py
+5
View File
@@ -26,6 +26,8 @@ ACCESS_WORDS = {"local": "This computer only", "lan": "My local network", "cloud
TUNE_MODES = ("quick", "thorough", "skip")
VISION_MODES = ("off", "on", "gpu", "cpu")
MTP_MODES = ("on", "off")
#: strata_control/split_engine.py MODES: the all-cards split engine on NVIDIA cards (auto = by itself where it is supported, the default)
SPLIT_MODES = ("auto", "on", "off")
#: strata_control/drafting.py MODES (tests/terminal/test_drafting.py keeps the two the same)
DRAFTING_MODES = ("auto", "mtp", "dflash2", "off")
SPEED_MODES = ("faster", "full") # DeepSeek: 5 experts per word (Faster) or all 6 (Full quality)
@@ -123,6 +125,7 @@ class LaunchSettings:
parallel: int | None = None # None = fit to the card
speed: str | None = None # DeepSeek only: None = the default for this PC (faster up to 16 GB of graphics memory); faster | full
drafting: str | None = None # None = Auto (recommended): the table in strata_control/drafting.py; mtp | dflash2 | off choose one. Saved in the launch config
split_engine: str | None = None # None = auto (the all-cards split engine runs by itself where it is supported); on | off. Saved in the preset and the launch config
mtp: str | None = None # None = as installed (draft layer on); "off" starts without it (less graphics memory and RAM, slower); "on" overrides an installed "off"
tune: str = "skip" # quick | thorough | skip
tune_allow_small: bool = False # the "answer-changing" tick, off by default, as in the app
@@ -345,6 +348,8 @@ def flags_text(s: LaunchSettings, *, preset: str | None = None) -> str:
parts.append(f"--mtp {s.mtp}")
if s.drafting:
parts.append(f"--drafting {s.drafting}")
if s.split_engine:
parts.append(f"--split-engine {s.split_engine}")
if s.speed == "full":
parts.append("--full-quality")
elif s.speed == "faster":
strata_control/split_engine.py
+118 -31
View File
@@ -1,20 +1,26 @@
"""strata-split: when the all-cards split engine may run, and the run config for it. OFF by default, opt-in.
"""strata-split: when the all-cards split engine may run, and the run config for it. AUTOMATIC by default (mode "auto"), with an off switch.
`strata-split-N` (src/split, from sFNX, MIT, Copyright (c) 2026 John Fink) gives every card a share of EVERY layer instead of some
layers, and the cards exchange the layer sums through page-locked host memory (no peer access anywhere). On two RTX A6000 with
UD-Q4_K_XL it decodes 40-95% faster than the normal layer split and reads prompts 1.7-2.8x slower (phaseA-tests/sfnx/RESULTS.md,
section 8), so it is an option, never the default, until the owner has chosen.
UD-Q4_K_XL it decodes 40-95% faster than the normal layer split (phaseA-tests/sfnx/RESULTS.md). The owner's decision: it runs by itself whenever the
rule below says yes ("auto", the default), and `strata launch --split-engine off` (or "split_engine": "off" in the run config, or
STRATA_SPLIT_ENGINE=0) turns it off; "on" asks for it and says why in one line when it cannot run. In auto a "no" is silent: the normal
layer split starts, as before.
This module is pure (no I/O besides reading GGUF directories): the launcher passes it facts and gets a Decision. The rule, from what
has actually been run:
- the person turned it on (config key "split_engine": true, or `strata launch --split-engine on`); never otherwise;
- not turned off (mode(): "auto" by default; "on" is the same rule but a refusal is reported; "off" never);
- Linux, 2 to 4 cards of ONE make (NVIDIA, or AMD RDNA4 gfx1201: the R9700 and RX 9070), all the same model of card and memory size;
- a card architecture and card count that have been run (TESTED, per make): NVIDIA sm_86 on 2 cards and AMD gfx1201 on 4 cards (see
TESTED below). Others build but stay off unless the caller says allow_untested (testers); the Decision says which part is untested
in plain words;
- the GGUF is UD-Q4_K_XL type: routed experts Q4_K/Q5_K/Q5_1/Q8_0, token embeddings and the shared expert Q8_0 (checked from the
file's tensor types, not its name), the Qwen3.8-Flash-Next shapes (48 layers, 512 experts);
- a card architecture and card count that have been run (TESTED, per make): NVIDIA sm_86 on 2 cards, AMD gfx1201 on 4 cards. Others build
but stay off unless the caller says allow_untested (testers); the Decision says which part is untested in plain words;
- AMD is opt-in for now: mode "auto" keeps the layer split on gfx1201 (AUTO_SPLIT_AMD_ARCHS is empty) and only `--split-engine on` runs the
split there. Adding "gfx1201" to that one constant makes AMD automatic like NVIDIA, once it meets the bar (8 users at or above the layer
split, a 103K prompt in the layer split's time; phaseA-tests/sfnx/RESULTS-HIP.md);
- the GGUF is a set the engine reads (checked from the file's tensor types, not its name), the Qwen3.8-Flash-Next shapes (48 layers,
512 experts): routed experts Q4_K/Q5_K/Q5_1/Q8_0/Q4_0 (UD-Q4_K_XL, Q8_0) or the i-quants IQ2_S/IQ3_XXS/IQ3_S/IQ4_XS/IQ4_NL and Strata's
Q2_0 (UD-IQ4_XS, the IQ3_S set); token embeddings, shared expert, output head Q8_0, or, in a set with i-quant or Q2_0 experts,
any type the loader converts to Q8_0/F32 when it reads the file (src/split/src/tensor_convert.hpp);
- ALL routed experts fit in the cards' memory (the expert cache is off: with it the engine keeps a 74 GiB pinned copy of the
experts in RAM) and the RAM that stays (the 27.5 GiB locked embedding table, a few GiB of buffers and the conversation cache)
is available with room to spare;
@@ -30,17 +36,31 @@ import pathlib
import re
from typing import Any
# What has been run with the host exchange, per make: {make: {architecture: card counts}}. NVIDIA: 2x RTX A6000 (boston, sm_86).
# What has been run with the host exchange, per make: {make: {architecture: card counts}}. NVIDIA: 2x RTX A6000 (boston, sm_86), 2x RTX PRO 6000 (srv118, sm_120).
# AMD: 4x Radeon AI PRO R9700 (r9700, gfx1201; phaseA-tests/sfnx/RESULTS-HIP.md); two R9700 cannot hold UD-Q4_K_XL's experts.
TESTED = {"nvidia": {"sm_86": (2,)}, "amd": {"gfx1201": ()}}
BUILT_ARCHS = {"nvidia": ("sm_86",), "amd": ("gfx1201",)} # the engine package's kernels: any other card is refused at start
TESTED_CARD_COUNTS = (2,) # NVIDIA (kept for callers that name it)
TESTED_ARCHS = (86,) # NVIDIA compute capability
TESTED = {"nvidia": {"sm_86": (2,), "sm_120": (2,)}, "amd": {"gfx1201": (4,)}}
BUILT_ARCHS = {"nvidia": ("sm_86", "sm_120"), "amd": ("gfx1201",)} # the engine package's kernels: any other card is refused at start
TESTED_CARD_COUNTS = (2,) # run with the host exchange: 2 cards (boston, 2x RTX A6000; srv118, 2x RTX PRO 6000)
TESTED_ARCHS = (86, 120) # compute capability: sm_86 (the 61-task suite, pictures, hashes) and sm_120 (decode, prompts, parallel, hashes,
# pictures; no suite there)
MODES = ("auto", "on", "off")
#: cards (compute capability) on which Auto drafting resolves to split + MTP rather than layer split + DFlash2: the split decodes faster there
#: (phaseA-tests/sfnx/RESULTS.md section 17). Another card is not auto-split at all (TESTED_ARCHS), so it keeps Auto = layer split + DFlash2.
AUTO_SPLIT_BEATS_DFLASH2 = (86, 120)
#: AMD architectures on which mode "auto" starts the split engine. EMPTY on purpose: on gfx1201 the split ships opt-in (`--split-engine on`) until it
#: meets the bar (8 users at or above the layer split, a 103K prompt at or under its time); then this becomes ("gfx1201",), the one flip.
AUTO_SPLIT_AMD_ARCHS: tuple = ()
TESTED_CONTEXT = 131072
MAX_CARDS = 4
LAYERS, EXPERTS = 48, 512 # Qwen3.8-Flash-Next
EXPERT_TYPES = {"Q4_K", "Q5_K", "Q5_1", "Q8_0", "Q4_0"} # the routed-expert kernels
Q8_ROLES = ("token_embd.weight",)
EXPERT_TYPES = {"Q4_K", "Q5_K", "Q5_1", "Q8_0", "Q4_0", # the routed-expert kernels (decode and prompt)
"IQ2_S", "IQ3_XXS", "IQ3_S", "IQ4_XS", "IQ4_NL", # ggml's i-quant vec_dots; prompt tiles in src/moe_mma_iq.cuh
"Q2_0"} # Strata's 2-bit type: the loader repacks it to Q4_0 blocks (twice the bytes)
# the loader (Gguf::convert_small_tensors) converts the dense tensors to Q8_0/F32 only when the file has one of these routed types
CONVERT_TRIGGERS = {"IQ2_S", "IQ3_XXS", "IQ3_S", "Q2_0"}
# and can read these types to do it
CONVERTED_FROM = {"Q8_0", "Q4_K", "Q5_K", "Q6_K", "IQ4_XS", "IQ4_NL", "BF16", "F16", "F32"}
PLE_TYPES = {"IQ4_NL", "Q8_0"} # per_layer_token_embd.weight as the host's row reader takes it
# measured on 2x RTX A6000 48 GB at 128K context, 6 lanes (phaseA-tests/sfnx/RESULTS.md section 8): a card held half the experts plus
# 8.4 GiB of dense weights, buffers and the K/V pool, with 3.6 GiB left free. 9 GiB per card is that plus a little.
CARD_OVERHEAD_GIB = 9.0
@@ -73,6 +93,8 @@ class ModelFacts:
embd_ram_bytes: int # token_embd + per_layer_token_embd: locked in RAM
layers: int
experts: int
ple_type: str | None = "IQ4_NL" # per_layer_token_embd.weight
output_type: str | None = "Q8_0" # output.weight
@dataclasses.dataclass
@@ -83,6 +105,7 @@ class Decision:
untested: list[str] = dataclasses.field(default_factory=list) # what about this setup nobody has run
need_vram_gib: float = 0.0
need_ram_gib: float = 0.0
quiet: bool = False # auto mode and a "no": the launch says nothing and starts the normal layer split
@property
def line(self) -> str:
@@ -109,6 +132,8 @@ def scan_model(first_shard: str | pathlib.Path) -> ModelFacts:
routed_bytes = 0
shexp: dict[str, int] = {}
embd_type = None
ple_type = None
output_type = None
embd_ram = 0
layers = 0
experts = 0
@@ -119,7 +144,7 @@ def scan_model(first_shard: str | pathlib.Path) -> ModelFacts:
b = t.expected_bytes() or 0
if "_exps" in n and n.endswith(".weight") and "shexp" not in n:
routed_types[t.type_name] = routed_types.get(t.type_name, 0) + 1
routed_bytes += b
routed_bytes += b * 2 if t.type_name == "Q2_0" else b # the loader repacks Q2_0 to Q4_0: twice the bytes in the cards
if len(t.shape) >= 3:
experts = max(experts, t.shape[-1])
mm = re.search(r"blk\.(\d+)\.", n)
@@ -131,8 +156,11 @@ def scan_model(first_shard: str | pathlib.Path) -> ModelFacts:
embd_type = t.type_name
embd_ram += b
elif n == "per_layer_token_embd.weight":
ple_type = t.type_name
embd_ram += b
return ModelFacts(routed_types, routed_bytes, embd_type, shexp, embd_ram, layers, experts)
elif n == "output.weight":
output_type = t.type_name
return ModelFacts(routed_types, routed_bytes, embd_type, shexp, embd_ram, layers, experts, ple_type, output_type)
def decide(*, enabled: bool, system: str, cards: list[Card], ram_total_gib: float, ram_available_gib: float, model: ModelFacts,
@@ -184,10 +212,17 @@ def decide(*, enabled: bool, system: str, cards: list[Card], ram_total_gib: floa
why.append("this is not a model it knows (no routed experts)")
elif bad:
why.append("the model's routed experts use " + "/".join(bad) + ", which it does not read")
if model.embd_type != "Q8_0":
why.append(f"the token embeddings are {model.embd_type}, it needs Q8_0")
if model.shexp_types and set(model.shexp_types) - {"F32"} != {"Q8_0"}: # the router gate vector beside them is F32
why.append("the shared expert is not Q8_0")
converts = bool(CONVERT_TRIGGERS & set(model.routed_types)) # the loader then converts the dense tensors it does not read
readable = CONVERTED_FROM if converts else {"Q8_0"}
if model.embd_type not in readable:
why.append(f"the token embeddings are {model.embd_type}, it needs Q8_0" if not converts else
f"the token embeddings are {model.embd_type}, which it cannot convert to Q8_0")
if model.shexp_types and not set(model.shexp_types) - {"F32"} <= readable: # the router gate vector beside the shared expert is F32
why.append("the shared expert is " + "/".join(sorted(set(model.shexp_types) - readable - {"F32"})) + ", which it cannot read" + ("" if converts else " (it needs Q8_0)"))
if model.output_type is not None and model.output_type not in readable:
why.append(f"the output head is {model.output_type}, which it cannot read" + ("" if converts else " (it needs Q8_0)"))
if model.ple_type is not None and model.ple_type not in PLE_TYPES:
why.append(f"the per-layer embedding table is {model.ple_type}, it needs IQ4_NL or Q8_0")
if (model.layers, model.experts) != (LAYERS, EXPERTS):
why.append("it only knows Qwen3.8-Flash-Next (48 layers, 512 experts)")
need_vram = model.routed_bytes / GIB / max(n, 1) + CARD_OVERHEAD_GIB
@@ -217,6 +252,11 @@ def split_config(base: dict, decision: Decision, *, exe: str, native: str, mtp_d
# every expert in the cards' memory, no expert cache, no arena: the engine reads these two at start
cfg["env"] = {**(base.get("env") or {}), "FNX_ARENA": "0", "FNX_EXPERT_RAM_GIB": "off"}
cfg["split_fallback"] = {k: base[k] for k in ("exe", "args", "env", "layer_split", "replicas", "user_args", "drafting", "cwd") if k in base}
# The split engine runs the model's own draft layer (the MTP) and nothing else. Saying so in the config ("drafting": "mtp") keeps serve/server.py's
# drafter step, which resolves Auto to DFlash2 where the engine has it, from adding DFlash2 flags to the split engine's arguments (it would refuse them
# and every start would fall back to the layer split). The fallback keeps the person's own choice (or none: Auto), so the layer split that starts
# instead gets DFlash2 as usual.
cfg["drafting"] = "mtp"
cfg["split_engine_note"] = decision.line
return cfg
@@ -303,18 +343,44 @@ def _flag(args: list, name: str) -> str | None:
return str(args[args.index(name) + 1]) if name in args and args.index(name) + 1 < len(args) else None
def wanted(cfg: dict, environ: dict | None = None) -> bool:
"""The person turned the option on: "split_engine": true in the run config, or STRATA_SPLIT_ENGINE=1 in the environment (testers)."""
def mode(cfg: dict, environ: dict | None = None) -> str:
""""auto" (the default), "on" or "off". The run config's "split_engine" (true/false or the word; `strata launch --split-engine X` writes it)
and, over it, STRATA_SPLIT_ENGINE in the environment (testers). Anything it does not know is auto."""
import os
return cfg.get("split_engine") is True or (environ if environ is not None else os.environ).get("STRATA_SPLIT_ENGINE") == "1"
env = environ if environ is not None else os.environ
raw = env.get("STRATA_SPLIT_ENGINE")
if raw is None or not str(raw).strip():
raw = cfg.get("split_engine")
if raw is True:
return "on"
if raw is False:
return "off"
return {"on": "on", "1": "on", "true": "on", "yes": "on", "off": "off", "0": "off", "false": "off", "no": "off"}.get(str(raw).strip().lower(), "auto")
def wanted(cfg: dict, environ: dict | None = None) -> bool:
"""The split engine may be tried for this launch (anything but off)."""
return mode(cfg, environ) != "off"
def apply_to_config(cfg: dict, *, exe_for_cards, system: str | None = None, allow_untested: bool = False, probe=None, ram=None,
scan=None, data_dir: str | None = None) -> Decision:
"""Turn a normal run config into the split engine's config IN PLACE when the rule says yes, and say what was decided.
exe_for_cards(n) -> the path of strata-split-n or None (the engine package has none); probe(indexes) -> [Card]; ram() -> (total,
available) GiB; scan(first_shard) -> ModelFacts: the facts of this PC and model, replaceable for tests."""
scan=None, data_dir: str | None = None, mode_: str | None = None) -> Decision:
"""Turn a normal run config into the split engine's config IN PLACE when the mode and the rule say yes, and say what was decided.
exe_for_cards(n, cc) -> the path of the strata-split-n built for compute capability cc, or None (the engine package has none);
probe(indexes) -> [Card]; ram() -> (total, available) GiB; scan(first_shard) -> ModelFacts: the facts of this PC and model, replaceable
for tests. Mode "off": nothing happens, silently. Mode "auto": a refusal is `quiet` (the normal layer split starts without a word)."""
m = mode_ or mode(cfg)
if m == "off":
return Decision(False, ["the all-cards split engine is turned off"], quiet=True)
d = _apply(cfg, exe_for_cards=exe_for_cards, system=system, allow_untested=allow_untested, probe=probe, ram=ram, scan=scan, data_dir=data_dir,
mode_=m)
if m == "auto" and not d.use_split:
d.quiet = True
return d
def _apply(cfg: dict, *, exe_for_cards, system, allow_untested, probe, ram, scan, data_dir, mode_: str = "auto") -> Decision:
import os
import pathlib
import sys
@@ -324,6 +390,20 @@ def apply_to_config(cfg: dict, *, exe_for_cards, system: str | None = None, allo
native = _flag(args, "--native")
if not native:
return Decision(False, ["this model is not a single GGUF file set"])
# choices the split engine does not honour: the normal layer split keeps them
# Drafting: Auto (nothing written, or "auto") on a setup the split accepts runs the split with its MTP: measured faster to decode than the layer
# split with DFlash2 on both tested cards (AUTO_SPLIT_BEATS_DFLASH2, numbers in docs/SPLIT_ENGINE.md). The person's own "dflash2" (or "off", which the
# split cannot do: it needs the draft layer) keeps the layer split; "mtp" is what the split runs.
chosen = str(cfg.get("drafting") or "auto").lower()
if chosen == "dflash2":
return Decision(False, ["DFlash2 drafting was chosen, which the all-cards split engine does not run"])
if chosen == "off":
return Decision(False, ["drafting is off, and the all-cards split engine needs the draft layer"])
if cfg.get("user_args") or cfg.get("replicas"):
return Decision(False, ["this run config has engine arguments of its own or replicas, which the all-cards split engine does not take"])
kv = _flag(args, "--kv")
if kv not in (None, "int8"):
return Decision(False, [f"the conversation memory format is {kv}, the all-cards split engine keeps int8"])
gpus = cfg.get("gpu")
indexes = [int(g) for g in gpus] if isinstance(gpus, list) else ([int(gpus)] if isinstance(gpus, int) else None)
cards = (probe or probe_cards)(indexes) if system == "linux" else []
@@ -337,9 +417,16 @@ def apply_to_config(cfg: dict, *, exe_for_cards, system: str | None = None, allo
allow_untested=allow_untested)
if not d.use_split:
return d
exe = exe_for_cards(d.cards)
amd = cards[0].vendor == "amd"
if amd and mode_ != "on" and not all(c.arch_name in AUTO_SPLIT_AMD_ARCHS for c in cards):
# opt-in on AMD: auto keeps the layer split (quiet), `--split-engine on` runs the split
return Decision(False, [f"the all-cards split engine is not automatic on {cards[0].arch_name} yet; it starts only with --split-engine on"], d.cards)
if not amd and chosen == "auto" and not all(c.cc in AUTO_SPLIT_BEATS_DFLASH2 for c in cards):
return Decision(False, ["on this card the layer split with DFlash2 has not been measured slower than the all-cards split engine"], d.cards)
# NVIDIA: strata-split-N-sm<cc> or strata-split-N; AMD: strata-split-N-<arch> (gfx1201), the third argument
exe = exe_for_cards(d.cards, cards[0].cc, cards[0].arch_name) if amd else exe_for_cards(d.cards, cards[0].cc)
if not exe:
return Decision(False, [f"this engine package has no strata-split-{d.cards}"], d.cards)
return Decision(False, [f"this engine package has no strata-split-{d.cards}" + (f"-{cards[0].arch_name}" if amd else "")], d.cards)
mtp = _flag(args, "--mtp")
if not mtp:
return Decision(False, ["the model has no draft layer files (--mtp)"], d.cards)
tests/split/convert_check.py
+80
View File
@@ -0,0 +1,80 @@
#!/usr/bin/env python3
"""Check the loader's tensor conversion (src/split/src/tensor_convert.hpp) against gguf-py.
convert_dump MODEL.gguf OUTDIR name... > dump.txt
convert_check.py MODEL.gguf OUTDIR dump.txt --gguf-py /path/to/llama.cpp/gguf-py
For each tensor in the dump: the values the loader holds now (read back by convert_dump) must equal gguf-py's dequantization of the
file's original tensor to within what the target type allows: exactly (F32 target), or within half a Q8_0 step per 32 values
(amax/254, plus the f16 rounding of the scale). Prints one line per tensor and exits 1 on any miss.
"""
import argparse
import os
import sys
import numpy as np
def main():
ap = argparse.ArgumentParser()
ap.add_argument("model")
ap.add_argument("outdir")
ap.add_argument("dump")
ap.add_argument("--gguf-py", required=True)
a = ap.parse_args()
sys.path.insert(0, a.gguf_py)
import gguf
from gguf import quants
shards = [a.model]
import re
m = re.match(r"(.*-)(\d{5})(-of-)(\d{5})(\.gguf)$", a.model)
if m:
shards = ["%s%05d%s%s%s" % (m[1], i, m[3], m[4], m[5]) for i in range(1, int(m[4]) + 1)]
tensors = {}
for s in shards:
r = gguf.GGUFReader(s)
for t in r.tensors:
tensors[t.name] = t
bad = 0
for line in open(a.dump):
p = line.split()
if len(p) != 3 or p[0] == "converted":
continue
name, ty = p[0], int(p[1])
t = tensors[name]
if t.tensor_type.name == "Q2_0": # Strata's type 42, not in gguf-py's dequantizers: (code - 1) * d, value j in byte j / 4 at bit 2 (j % 4)
raw = np.ascontiguousarray(t.data).reshape(-1, 18)
d = raw[:, :2].copy().view(np.float16).astype(np.float64).reshape(-1, 1)
qs = raw[:, 2:]
codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=2).reshape(-1, 64)
ref = ((codes.astype(np.float64) - 1.0) * d).reshape(-1)
else:
ref = quants.dequantize(t.data, t.tensor_type).astype(np.float64).reshape(-1)
got = np.fromfile(os.path.join(a.outdir, name + ".f32"), dtype=np.float32).astype(np.float64)
if ref.shape != got.shape:
print("FAIL %s: %d values, gguf-py %d" % (name, got.size, ref.size))
bad += 1
continue
err = np.abs(got - ref)
if ty == 0: # F32: BF16 / F16 / F32 source, exact
lim = np.zeros_like(err)
elif ty != 8: # a type the loader keeps or repacks without changing a value (Q2_0 -> Q4_0, IQ4_NL as is): f32 rounding only
lim = np.abs(ref) * 1e-6 + 1e-12
else: # Q8_0: half a step of its block, the scale stored as f16
blk = ref.reshape(-1, 32)
amax = np.abs(blk).max(axis=1)
# the scale d = amax/127 is stored as f16 (relative error 2^-11), and the value is q*d with |q| <= 127
# (a block whose d falls below f16's normal range, 6.1e-5, has an absolute step of 2^-25 in d: 127 * 2^-25 on the value;
# 8 of the 19.9M blocks of token_embd are like that)
lim = np.repeat(amax * (0.5 / 127.0 + 2.0 ** -10) + 127 * 2.0 ** -25 + 1e-12, 32)
ok = bool((err <= lim).all())
rel = float(np.sqrt((err ** 2).sum() / max((ref ** 2).sum(), 1e-30)))
print("%s %s %s from %s max|d|=%.3g rel_rms=%.3g" % ("ok " if ok else "FAIL", name, {0: "F32", 8: "Q8_0", 2: "Q4_0", 20: "IQ4_NL"}.get(ty, str(ty)), t.tensor_type.name, err.max(), rel))
bad += not ok
sys.exit(1 if bad else 0)
if __name__ == "__main__":
main()
tests/split/test_split_engine.py
+269 -10
View File
@@ -59,7 +59,7 @@ class Rule(unittest.TestCase):
self.assertIn("same model", d.line)
def test_untested_counts_archs_and_contexts_stay_off_unless_a_tester_allows_them(self):
for c in (cards(3), cards(2, cc=89), cards(2, cc=120)):
for c in (cards(3), cards(2, cc=89), cards(2, cc=90)):
d = decide(cards=c)
self.assertFalse(d.use_split)
self.assertIn("only been run", d.line)
@@ -84,12 +84,59 @@ class Rule(unittest.TestCase):
self.assertIn("RAM", d.line)
def test_only_the_quant_types_it_reads(self):
for model in (udq4(routed_types={"IQ3_S": 10, "IQ2_S": 20, "IQ4_NL": 39}), udq4(embd_type="IQ4_XS"),
udq4(shexp_types={"BF16": 48, "IQ4_NL": 47}), udq4(layers=64), udq4(routed_types={})):
for model in (udq4(layers=64), udq4(routed_types={}),
udq4(routed_types={"Q6_K": 48}), # UD-Q6_K_XL-like experts: no kernel
udq4(embd_type="IQ4_XS"), # a K-quant set is not converted at load
udq4(shexp_types={"BF16": 48, "IQ4_NL": 47}),
udq4(output_type="Q6_K"),
udq4(ple_type="Q6_K")):
d = decide(model=model)
self.assertFalse(d.use_split, model)
self.assertTrue(d.line.startswith("Using the normal layer split"))
def test_the_iq3_s_set_and_ud_iq4_xs_run(self):
iq3s = S.ModelFacts(routed_types={"IQ2_S": 40, "IQ3_XXS": 34, "IQ3_S": 20, "IQ4_XS": 2, "IQ4_NL": 39, "Q2_0": 9},
routed_bytes=int(49 * GIB), embd_type="IQ4_XS", shexp_types={"IQ4_XS": 60, "Q4_K": 40, "Q5_K": 20, "IQ4_NL": 24},
embd_ram_bytes=int(27.2 * GIB), layers=48, experts=512, ple_type="IQ4_NL", output_type="Q6_K")
d = decide(model=iq3s)
self.assertTrue(d.use_split, d.reasons)
ud_iq4_xs = S.ModelFacts(routed_types={"IQ3_S": 94, "IQ4_XS": 2, "IQ4_NL": 43, "Q8_0": 5}, routed_bytes=int(59.5e9),
embd_type="Q8_0", shexp_types={"Q8_0": 144}, embd_ram_bytes=int(27.5 * GIB), layers=48, experts=512,
ple_type="IQ4_NL", output_type="Q6_K")
d = decide(model=ud_iq4_xs)
self.assertTrue(d.use_split, d.reasons)
self.assertTrue(d.need_vram_gib < 48, d.need_vram_gib)
def test_a_type_the_loader_cannot_convert_is_a_plain_refusal(self):
for model, word in ((udq4(routed_types={"IQ3_S": 94}, embd_type="Q3_K"), "cannot convert"),
(udq4(routed_types={"IQ3_S": 94}, shexp_types={"Q3_K": 5}), "cannot read"),
(udq4(routed_types={"IQ3_S": 94}, ple_type="Q6_K"), "per-layer"),
(udq4(routed_types={"IQ3_S": 94, "TQ2_0": 2}), "TQ2_0")):
d = decide(model=model)
self.assertFalse(d.use_split)
self.assertIn(word, d.line)
self.assertNotIn(chr(10), d.line) # one plain line
def test_scan_counts_the_q2_0_repack_twice(self):
from unittest import mock
class T:
def __init__(self, name, type_name, shape, nbytes):
self.name, self.type_name, self.shape, self._n = name, type_name, shape, nbytes
def expected_bytes(self):
return self._n
class F:
def __init__(self, path):
self.tensors = [T("blk.0.ffn_down_exps.weight", "Q2_0", (640, 2560, 512), 1000), T("blk.1.ffn_gate_exps.weight", "IQ3_S", (2560, 640, 512), 5000),
T("token_embd.weight", "IQ4_XS", (2560, 248320), 10), T("per_layer_token_embd.weight", "IQ4_NL", (1, 2), 20),
T("output.weight", "Q6_K", (2560, 248320), 30)]
with mock.patch("tools.gguf_reader.GGUFFile", F):
m = S.scan_model("/x/model-00001-of-00001.gguf")
self.assertEqual(m.routed_bytes, 2 * 1000 + 5000)
self.assertEqual((m.embd_type, m.ple_type, m.output_type), ("IQ4_XS", "IQ4_NL", "Q6_K"))
self.assertEqual(m.embd_ram_bytes, 30)
def test_the_reasons_are_all_reported(self):
d = decide(system="windows", cards=cards(1), model=udq4(embd_type="Q6_K"))
self.assertGreaterEqual(len(d.reasons), 3)
@@ -137,7 +184,7 @@ class Apply(unittest.TestCase):
"cwd": tmp, "tokenizer": "tok", "port": 8080, "gpu": [0, 1], "env": {}, "split_engine": True}
def run_apply(self, cfg, *, cards_=None, avail=80.0, model=None, exe="/e/strata-split-2", **kw):
return S.apply_to_config(cfg, exe_for_cards=lambda n: exe, system=kw.pop("system", "linux"), probe=lambda idx: cards_ or cards(),
return S.apply_to_config(cfg, exe_for_cards=lambda n, cc=None: exe, system=kw.pop("system", "linux"), probe=lambda idx: cards_ or cards(),
ram=lambda: (125.0, avail), scan=lambda path: model or udq4(), **kw)
def test_a_good_pc_gets_the_split_config(self):
@@ -175,15 +222,227 @@ class Apply(unittest.TestCase):
with tempfile.TemporaryDirectory() as tmp:
cfg = self.cfg(tmp)
d = S.apply_to_config(cfg, exe_for_cards=lambda n: "/e/x", system="linux", probe=lambda i: cards(), ram=lambda: (125.0, 80.0), scan=boom)
d = S.apply_to_config(cfg, exe_for_cards=lambda n, cc=None: "/e/x", system="linux", probe=lambda i: cards(), ram=lambda: (125.0, 80.0), scan=boom)
self.assertFalse(d.use_split)
self.assertIn("could not be read", d.line)
def test_wanted_is_opt_in(self):
self.assertFalse(S.wanted({}, {}))
self.assertFalse(S.wanted({"split_engine": False}, {}))
self.assertTrue(S.wanted({"split_engine": True}, {}))
self.assertTrue(S.wanted({}, {"STRATA_SPLIT_ENGINE": "1"}))
def test_modes_auto_is_the_default_and_off_is_the_switch(self):
self.assertEqual(S.mode({}, {}), "auto")
for raw, want in ((True, "on"), (False, "off"), ("on", "on"), ("off", "off"), ("auto", "auto"), ("OFF", "off"), ("nonsense", "auto"), (None, "auto")):
self.assertEqual(S.mode({"split_engine": raw}, {}), want, raw)
self.assertEqual(S.mode({"split_engine": "on"}, {"STRATA_SPLIT_ENGINE": "0"}), "off") # the environment wins (testers)
self.assertEqual(S.mode({"split_engine": "off"}, {"STRATA_SPLIT_ENGINE": "1"}), "on")
self.assertEqual(S.mode({}, {"STRATA_SPLIT_ENGINE": "off"}), "off")
self.assertEqual(S.mode({"split_engine": "off"}, {"STRATA_SPLIT_ENGINE": ""}), "off") # an empty variable is no variable
self.assertTrue(S.wanted({}, {}))
self.assertFalse(S.wanted({"split_engine": "off"}, {}))
def test_auto_uses_the_split_where_it_is_supported_and_says_nothing_otherwise(self):
import copy
import tempfile
with tempfile.TemporaryDirectory() as tmp:
cfg = self.cfg(tmp)
del cfg["split_engine"] # nothing asked: auto
d = self.run_apply(cfg)
self.assertTrue(d.use_split, d.reasons)
self.assertFalse(d.quiet)
self.assertEqual(cfg["exe"], "/e/strata-split-2")
for kw in (dict(cards_=cards(1)), dict(avail=20.0), dict(exe=None), dict(system="windows"), dict(model=udq4(routed_types={"Q6_K": 48})),
dict(cards_=cards(3)), dict(cards_=cards(2, cc=89))):
cfg = self.cfg(tmp)
del cfg["split_engine"]
before = copy.deepcopy(cfg)
d = self.run_apply(cfg, **kw)
self.assertFalse(d.use_split, kw)
self.assertTrue(d.quiet, kw) # silent: the launch goes on with the layer split
self.assertEqual(cfg, before, kw)
def test_on_says_why_it_cannot_run(self):
import tempfile
with tempfile.TemporaryDirectory() as tmp:
cfg = self.cfg(tmp)
cfg["split_engine"] = "on"
d = self.run_apply(cfg, cards_=cards(1))
self.assertFalse(d.use_split)
self.assertFalse(d.quiet)
self.assertIn("two or more cards", d.line)
def test_off_never_runs_and_changes_nothing(self):
import copy
import tempfile
with tempfile.TemporaryDirectory() as tmp:
for off in ("off", False):
cfg = self.cfg(tmp)
cfg["split_engine"] = off
before = copy.deepcopy(cfg)
d = self.run_apply(cfg)
self.assertFalse(d.use_split)
self.assertTrue(d.quiet)
self.assertEqual(cfg, before)
def test_choices_the_split_engine_does_not_honour_keep_the_layer_split(self):
import tempfile
with tempfile.TemporaryDirectory() as tmp:
for change in (lambda c: c.update(drafting="dflash2"), lambda c: c.update(user_args=["--x"]), lambda c: c.update(replicas=2),
lambda c: c["args"].extend(["--kv", "q4_0"])):
cfg = self.cfg(tmp)
del cfg["split_engine"]
change(cfg)
d = self.run_apply(cfg)
self.assertFalse(d.use_split)
self.assertTrue(d.quiet)
cfg = self.cfg(tmp)
cfg["args"].extend(["--kv", "int8"])
self.assertTrue(self.run_apply(cfg).use_split)
def test_drafting_auto_resolves_to_the_split_with_its_mtp_and_the_fallback_keeps_the_layer_split_choice(self):
import tempfile
with tempfile.TemporaryDirectory() as tmp:
for chosen, split in ((None, True), ("auto", True), ("mtp", True), ("dflash2", False), ("off", False)):
cfg = self.cfg(tmp)
del cfg["split_engine"]
if chosen:
cfg["drafting"] = chosen
before_fallback = {k: v for k, v in cfg.items() if k == "drafting"}
d = self.run_apply(cfg)
self.assertEqual(d.use_split, split, chosen)
if split:
# Auto/MTP: the split engine runs its MTP; the config says so, so serve/server.py's drafter step (Auto = DFlash2) leaves its arguments alone
self.assertEqual(cfg["drafting"], "mtp", chosen)
self.assertEqual(cfg["split_fallback"].get("drafting"), before_fallback.get("drafting"), chosen) # the layer split keeps the person's word, or Auto
self.assertNotIn("--draft", cfg["args"])
else:
self.assertEqual(cfg.get("drafting"), chosen) # untouched: the layer split gets DFlash2 as chosen
self.assertNotIn("split_fallback", cfg)
def test_auto_needs_a_card_where_the_split_decodes_faster_than_the_layer_split_with_dflash2(self):
import tempfile
with tempfile.TemporaryDirectory() as tmp:
cfg = self.cfg(tmp)
del cfg["split_engine"]
d = self.run_apply(cfg, cards_=cards(2, cc=89), allow_untested=True) # a tester allowed it, but nobody measured it there
self.assertFalse(d.use_split)
self.assertIn("DFlash2", d.line)
cfg["drafting"] = "mtp" # the person asked for the MTP: no comparison to win
self.assertTrue(self.run_apply(cfg, cards_=cards(2, cc=89), allow_untested=True).use_split)
def test_the_engine_is_picked_by_card_count_and_architecture(self):
import tempfile
seen = []
def exe_for(n, cc=None):
seen.append((n, cc))
return "/e/strata-split-2-sm120"
with tempfile.TemporaryDirectory() as tmp:
cfg = self.cfg(tmp)
d = S.apply_to_config(cfg, exe_for_cards=exe_for, system="linux", probe=lambda i: cards(2, cc=120, name="NVIDIA RTX PRO 6000", vram_mib=97887),
ram=lambda: (125.0, 80.0), scan=lambda path: udq4())
self.assertTrue(d.use_split, d.reasons)
self.assertEqual(seen, [(2, 120)])
self.assertEqual(cfg["exe"], "/e/strata-split-2-sm120")
class AmdOptIn(unittest.TestCase):
"""AMD gfx1201 (4 x R9700): the split ships opt-in. auto keeps the layer split, `on` starts the split, `off` is off; one constant flips auto."""
R9700 = dict(name="AMD Radeon AI PRO R9700", vram_mib=32624, cc=0, vendor="amd", arch="gfx1201")
cfg = Apply.cfg
def amd(self, tmp, **kw):
cfg = self.cfg(tmp)
cfg["gpu"] = [0, 1, 2, 3]
for k, v in kw.items():
cfg[k] = v
return cfg
def run_amd(self, cfg, seen=None, **kw):
def exe_for(n, cc=None, arch=None):
if seen is not None:
seen.append((n, cc, arch))
return f"/e/strata-split-{n}-{arch}" if arch else None
return S.apply_to_config(cfg, exe_for_cards=exe_for, system="linux", probe=lambda i: cards(4, **self.R9700), ram=lambda: (125.0, 80.0),
scan=lambda path: udq4(routed_bytes=int(71.7 * GIB)), **kw)
def test_auto_keeps_the_layer_split_on_gfx1201_and_says_nothing(self):
import tempfile
self.assertEqual(S.AUTO_SPLIT_AMD_ARCHS, ()) # the one constant: ("gfx1201",) when the targets are met
with tempfile.TemporaryDirectory() as tmp:
cfg = self.amd(tmp)
del cfg["split_engine"]
before = dict(cfg)
d = self.run_amd(cfg)
self.assertFalse(d.use_split)
self.assertTrue(d.quiet)
self.assertEqual(cfg, before) # the layer split's config, untouched
cfg["drafting"] = "mtp" # not even with the MTP chosen
self.assertFalse(self.run_amd(cfg).use_split)
def test_explicit_on_starts_the_split_with_the_gfx1201_engine(self):
import tempfile
for on in (True, "on"):
with tempfile.TemporaryDirectory() as tmp:
cfg = self.amd(tmp, split_engine=on)
seen = []
d = self.run_amd(cfg, seen)
self.assertTrue(d.use_split, d.reasons)
self.assertEqual(d.cards, 4)
self.assertEqual(d.untested, []) # four R9700 have been run
self.assertEqual(seen, [(4, 0, "gfx1201")])
self.assertEqual(cfg["exe"], "/e/strata-split-4-gfx1201")
self.assertEqual(cfg["drafting"], "mtp")
self.assertEqual(cfg["split_fallback"]["exe"], "component:strata-engine")
def test_on_without_the_gfx1201_binary_says_so_and_off_is_off(self):
import tempfile
with tempfile.TemporaryDirectory() as tmp:
cfg = self.amd(tmp, split_engine="on")
d = S.apply_to_config(cfg, exe_for_cards=lambda n, cc=None, arch=None: None, system="linux", probe=lambda i: cards(4, **self.R9700),
ram=lambda: (125.0, 80.0), scan=lambda path: udq4())
self.assertFalse(d.use_split)
self.assertFalse(d.quiet)
self.assertIn("strata-split-4-gfx1201", d.line)
cfg = self.amd(tmp, split_engine="off")
before = dict(cfg)
self.assertFalse(self.run_amd(cfg).use_split)
self.assertEqual(cfg, before)
def test_two_r9700_cannot_hold_the_experts_and_other_amd_cards_are_refused(self):
import tempfile
with tempfile.TemporaryDirectory() as tmp:
cfg = self.amd(tmp, split_engine="on")
d = S.apply_to_config(cfg, exe_for_cards=lambda n, cc=None, arch=None: "/e/x", system="linux", probe=lambda i: cards(2, **self.R9700),
ram=lambda: (125.0, 80.0), scan=lambda path: udq4())
self.assertFalse(d.use_split)
d = S.apply_to_config(self.amd(tmp, split_engine="on"), exe_for_cards=lambda n, cc=None, arch=None: "/e/x", system="linux",
probe=lambda i: cards(4, **{**self.R9700, "arch": "gfx1100"}), ram=lambda: (125.0, 80.0), scan=lambda path: udq4())
self.assertFalse(d.use_split)
self.assertIn("gfx1100", d.line)
def test_flipping_the_constant_makes_auto_start_it(self):
import tempfile
old = S.AUTO_SPLIT_AMD_ARCHS
S.AUTO_SPLIT_AMD_ARCHS = ("gfx1201",)
try:
with tempfile.TemporaryDirectory() as tmp:
cfg = self.amd(tmp)
del cfg["split_engine"]
self.assertTrue(self.run_amd(cfg).use_split)
finally:
S.AUTO_SPLIT_AMD_ARCHS = old
class NoPeerAccess(unittest.TestCase):
tests/split/tile_check.py
+69
View File
@@ -0,0 +1,69 @@
#!/usr/bin/env python3
"""Check the prompt kernel's weight tiles against gguf-py.
moe_tile_test MODEL.gguf OUTDIR > tiles.txt
tile_check.py MODEL.gguf OUTDIR tiles.txt --gguf-py /path/to/gguf-py
For each line of the dump (tensor, loader type, rows, columns, expert): the f16 values the tile put in shared memory must equal gguf-py's
dequantization of the file's expert, rounded once to f16 (a tile value may differ by one f16 step where the two compute the f32 product in
another order). Strata's Q2_0 (type 42; the loader repacks it to Q4_0 blocks, type 2) is checked against (code - 1) * d read from the
file. Prints one line per tensor and expert; exit 1 on any miss.
"""
import argparse
import os
import re
import sys
import numpy as np
def main():
ap = argparse.ArgumentParser()
ap.add_argument("model")
ap.add_argument("outdir")
ap.add_argument("dump")
ap.add_argument("--gguf-py", required=True)
a = ap.parse_args()
sys.path.insert(0, a.gguf_py)
import gguf
from gguf import quants
m = re.match(r"(.*-)(\d{5})(-of-)(\d{5})(\.gguf)$", a.model)
shards = ["%s%05d%s%s%s" % (m[1], i, m[3], m[4], m[5]) for i in range(1, int(m[4]) + 1)] if m else [a.model]
tensors = {}
for s in shards:
for t in gguf.GGUFReader(s).tensors:
tensors[t.name] = t
bad = 0
for line in open(a.dump):
p = line.split()
if len(p) != 5 or not p[1].isdigit():
continue
name, ty, nrows, ncols, e = p[0], int(p[1]), int(p[2]), int(p[3]), int(p[4])
t = tensors[name]
raw = np.asarray(t.data)[e]
if t.tensor_type.name == "Q2_0":
raw = np.ascontiguousarray(raw).reshape(-1, 18)
d = raw[:, :2].copy().view(np.float16).astype(np.float32).reshape(-1, 1)
codes = np.stack([(raw[:, 2:] >> (2 * k)) & 3 for k in range(4)], axis=2).reshape(-1, 64)
ref = ((codes.astype(np.float32) - 1.0) * d).reshape(nrows, ncols)
else:
ref = quants.dequantize(raw, t.tensor_type).astype(np.float32).reshape(nrows, ncols)
ref16 = ref.astype(np.float16)
got = np.fromfile(os.path.join(a.outdir, "%s.e%d.f16" % (name, e)), dtype=np.float16).reshape(nrows, ncols)
gi, ri = got.view(np.int16).astype(np.int32), ref16.view(np.int16).astype(np.int32)
# f16 ordering: compare as sign-magnitude steps
def key(x):
return np.where(x < 0, -(x & 0x7FFF), x)
step = np.abs(key(gi) - key(ri))
exact = float((step == 0).mean())
ok = bool(step.max() <= 1 and exact >= 0.999 and not np.isnan(got.astype(np.float32)).any())
print("%s %-44s %-8s e%-3d exact %.5f max step %d max|d| %.3g" % ("ok " if ok else "FAIL", name, t.tensor_type.name, e, exact, step.max(),
float(np.abs(got.astype(np.float32) - ref).max())))
bad += not ok
sys.exit(1 if bad else 0)
if __name__ == "__main__":
main()
tests/terminal/test_split_engine_flag.py
+46
View File
@@ -0,0 +1,46 @@
"""`strata launch --split-engine auto|on|off`: the owner's switch for the all-cards split engine (strata_control/split_engine.py).
auto (nothing written) is the default; the choice is listed in --help, in the copyable flags, in the launch config and in a saved preset, so an
"off" survives updates and restarts; the launcher's apply_split_engine honours it (off: nothing, auto: silent on a refusal, on: one line)."""
from __future__ import annotations
import unittest
from strata_cli.launch import settings as S
from strata_cli.launch.runner import derive_config
from strata_control import split_engine as SE
from tests.terminal.launch_support import run_cli
class Flag(unittest.TestCase):
BASE = {"exe": "component:strata-engine", "args": ["--pack", "/d/packs/x", "--mtp", "/d/mtp/rt"], "port": 8080, "model_name": "m"}
def test_the_flag_is_listed_checked_copyable_and_saved_in_the_launch_config(self):
code, out, _ = run_cli(["launch", "--help"])
self.assertIn("--split-engine", out)
self.assertIn("auto|on|off", out)
for mode in S.SPLIT_MODES:
cfg = derive_config(self.BASE, S.LaunchSettings(split_engine=mode), bool_flags=frozenset())
self.assertEqual(cfg["split_engine"], mode)
self.assertEqual(SE.mode(cfg, {}), mode)
self.assertNotIn("split_engine", derive_config(self.BASE, S.LaunchSettings(), bool_flags=frozenset()), "nobody chose: auto writes nothing")
self.assertEqual(SE.mode(derive_config(self.BASE, S.LaunchSettings(), bool_flags=frozenset()), {}), "auto")
self.assertIn("--split-engine off", S.flags_text(S.LaunchSettings(split_engine="off")))
self.assertNotIn("--split-engine", S.flags_text(S.LaunchSettings()))
code, out, err = run_cli(["launch", "--plan", "--split-engine", "sometimes"])
self.assertNotEqual(code, 0)
def test_a_saved_off_survives_a_new_launch_config_and_a_preset_round_trip(self):
s = S.LaunchSettings.from_dict(S.LaunchSettings(split_engine="off").to_dict())
self.assertEqual(s.split_engine, "off")
self.assertIsNone(S.LaunchSettings.from_dict({}).split_engine)
# a config written earlier with "off" is kept when this launch names nothing
cfg = derive_config({**self.BASE, "split_engine": "off"}, S.LaunchSettings(), bool_flags=frozenset())
self.assertEqual(cfg["split_engine"], "off")
# and the person's flag of this launch wins over the old file
cfg = derive_config({**self.BASE, "split_engine": "off"}, S.LaunchSettings(split_engine="on"), bool_flags=frozenset())
self.assertEqual(cfg["split_engine"], "on")
if __name__ == "__main__":
unittest.main()