From 74259942e48f1aac76879328a9611edc619e94e4 Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 14:20:32 +0000 Subject: [PATCH 01/10] Repair C++ logit validation validate-logits.py documented a C++ comparison but returned None without invoking the existing dump-logits binary. The binary also used std::partial_sort without including algorithm, which fails a clean GCC 15 build. Read dump-logits shape and float data, validate the serialized dimensions, and compare the result as documented. Cover the wrapper protocol with a deterministic stand-in executable, then exercise the real dump-logits binary and synthetic GGUF in the existing end-to-end test. Assisted-by: OpenAI Codex --- CMakeLists.txt | 8 +++++ tests/test-e2e.py | 21 +++++++++++ tests/test-validate-logits.py | 45 ++++++++++++++++++++++++ tools/dump-logits.cpp | 1 + tools/validate-logits.py | 66 ++++++++++++++++++++--------------- 5 files changed, 112 insertions(+), 29 deletions(-) create mode 100644 tests/test-validate-logits.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 382f44a..63bf279 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -53,6 +53,14 @@ endif() # ── Tests ─────────────────────────────────────────────────────── if(DIFFUSE_BUILD_TESTS) enable_testing() + find_package(Python3 COMPONENTS Interpreter QUIET) + if(Python3_Interpreter_FOUND) + add_test( + NAME test-validate-logits + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-validate-logits.py) + endif() + add_executable(test-forward tests/test-forward.cpp) target_link_libraries(test-forward PRIVATE diffuse) add_test(NAME test-forward COMMAND test-forward) diff --git a/tests/test-e2e.py b/tests/test-e2e.py index ba24590..04441fe 100644 --- a/tests/test-e2e.py +++ b/tests/test-e2e.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 """End-to-end test: create tiny model, convert, load in C++, run forward pass.""" +import importlib.util import json import os import sys @@ -102,6 +103,26 @@ def main(): print(f"FAIL: test-forward exited with code {result.returncode}") sys.exit(1) + # Exercise the real dump-logits executable through the same Python + # wrapper used for PyTorch/C++ validation. + validator_path = os.path.join(project_dir, "tools", "validate-logits.py") + validator_spec = importlib.util.spec_from_file_location( + "diffuse_validate_logits", validator_path) + validator = importlib.util.module_from_spec(validator_spec) + validator_spec.loader.exec_module(validator) + + dump_binary = os.path.join(build_dir, "dump-logits") + dump_tokens = [1, 2, 3, 4] + dumped_logits = validator.run_cpp( + gguf_path, dump_binary, dump_tokens, n_threads=4) + if dumped_logits.shape != (len(dump_tokens), VOCAB): + print(f"FAIL: unexpected dumped logit shape {dumped_logits.shape}") + sys.exit(1) + if not np.all(np.isfinite(dumped_logits)): + print("FAIL: dump-logits returned non-finite values") + sys.exit(1) + print("dump-logits wrapper OK") + print("\nEnd-to-end test PASSED!") diff --git a/tests/test-validate-logits.py b/tests/test-validate-logits.py new file mode 100644 index 0000000..b25624f --- /dev/null +++ b/tests/test-validate-logits.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Regression tests for the dump-logits Python wrapper.""" + +import importlib.util +import os +import stat +import tempfile +import unittest + +import numpy as np + + +PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VALIDATOR_PATH = os.path.join(PROJECT_DIR, "tools", "validate-logits.py") +VALIDATOR_SPEC = importlib.util.spec_from_file_location( + "diffuse_validate_logits", VALIDATOR_PATH) +VALIDATOR = importlib.util.module_from_spec(VALIDATOR_SPEC) +VALIDATOR_SPEC.loader.exec_module(VALIDATOR) + + +class ValidateLogitsTest(unittest.TestCase): + def test_run_cpp_reads_shape_and_logits(self): + with tempfile.TemporaryDirectory() as tmpdir: + fake_dump = os.path.join(tmpdir, "fake-dump-logits") + with open(fake_dump, "w", encoding="utf-8") as script: + script.write( + "#!/usr/bin/env python3\n" + "import struct, sys\n" + "output = sys.argv[sys.argv.index('-o') + 1]\n" + "with open(output, 'wb') as f:\n" + " f.write(struct.pack(' #include #include #include diff --git a/tools/validate-logits.py b/tools/validate-logits.py index 90ab3dd..084c344 100644 --- a/tools/validate-logits.py +++ b/tools/validate-logits.py @@ -8,23 +8,20 @@ python validate-logits.py \ --model /path/to/LLaDA-8B-Instruct \ --gguf llada-8b-f16.gguf \ - --cpp-bin ./build/diffuse-cli \ + --cpp-bin ./build/dump-logits \ --tokens "1,2,3,4,5,6,7,8" # Or with prompt (requires tokenizer): python validate-logits.py \ --model /path/to/LLaDA-8B-Instruct \ --gguf llada-8b-f16.gguf \ - --cpp-bin ./build/diffuse-cli \ + --cpp-bin ./build/dump-logits \ --prompt "Hello world" """ import argparse -import json import os -import struct import subprocess -import sys import tempfile import time @@ -87,34 +84,44 @@ def run_pytorch(model_dir, token_ids): def run_cpp(gguf_path, cpp_bin, token_ids, n_threads=4): - """Run forward pass with diffuse-cpp. Returns logits [n_tokens, vocab_size]. - - Since the CLI doesn't output raw logits yet, we use a helper binary. - For now, write a small C++ helper that dumps logits to a binary file. - """ - # Build tokens string + """Run dump-logits and return float32 logits [n_tokens, vocab_size].""" tokens_str = ",".join(map(str, token_ids)) - - # We need a way to get raw logits from C++. For validation, we write - # a temporary Python script that calls the C library via ctypes. - # Actually, the simplest approach: modify the test to dump logits. - # For now, use the test-forward binary approach: - - # Write a small Python script that uses the GGUF file to get metadata, - # then we'll compare at the tensor level instead. - print("NOTE: Raw logits comparison requires the C++ binary to dump logits.") - print(" For now, validating conversion correctness at tensor level.") - print(" Full logit comparison will be available after adding --dump-logits to CLI.") - - return None + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "logits.bin") + command = [ + cpp_bin, + "-m", gguf_path, + "--tokens", tokens_str, + "-o", output_path, + "-t", str(n_threads), + ] + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"dump-logits exited with code {result.returncode}:\n{result.stderr}" + ) + + with open(output_path, "rb") as logits_file: + header = np.fromfile(logits_file, dtype=np.int32, count=2) + logits = np.fromfile(logits_file, dtype=np.float32) + + if header.size != 2: + raise RuntimeError("dump-logits output is missing its shape header") + n_tokens, n_vocab = map(int, header) + if n_tokens != len(token_ids): + raise RuntimeError( + f"dump-logits returned {n_tokens} rows for {len(token_ids)} tokens" + ) + if logits.size != n_tokens * n_vocab: + raise RuntimeError( + f"dump-logits wrote {logits.size} values for shape " + f"{n_tokens} x {n_vocab}" + ) + return logits.reshape(n_tokens, n_vocab) def compare_logits(pytorch_logits, cpp_logits, top_k=10): """Compare logit distributions between PyTorch and C++.""" - if cpp_logits is None: - print("\nSkipping logit comparison (C++ logits not available)") - return - n_tokens, vocab_size = pytorch_logits.shape assert cpp_logits.shape == pytorch_logits.shape, \ f"Shape mismatch: PyTorch {pytorch_logits.shape} vs C++ {cpp_logits.shape}" @@ -155,7 +162,8 @@ def main(): parser = argparse.ArgumentParser(description="Validate forward pass logits") parser.add_argument("--model", "-m", required=True, help="HF model directory") parser.add_argument("--gguf", "-g", required=True, help="GGUF file path") - parser.add_argument("--cpp-bin", default="./build/diffuse-cli", help="C++ CLI binary") + parser.add_argument("--cpp-bin", default="./build/dump-logits", + help="Path to the dump-logits binary") parser.add_argument("--tokens", help="Comma-separated token IDs") parser.add_argument("--prompt", "-p", help="Text prompt (tokenized automatically)") parser.add_argument("--threads", "-t", type=int, default=4, help="C++ threads") From 43755b7b9f235703847733dbe2e2a2e172ac870c Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 12:27:00 +0000 Subject: [PATCH 02/10] Rank Dream candidates by normalized probability Dream compares confidence across token positions after applying softmax. Raw logits are only comparable within one distribution: adding a position-specific constant preserves its probabilities while changing its raw-logit rank. Compute top-token probability, probability margin, and entropy from a numerically stable normalized distribution. Cover the regression with two rows whose raw-logit and probability-confidence orderings disagree. Assisted-by: OpenAI Codex --- include/diffuse.h | 2 +- src/diffuse-sampler.cpp | 92 ++++++++++++++++++++++++----------------- src/diffuse-sampler.h | 13 ++++++ tests/test-forward.cpp | 29 ++++++++++++- 4 files changed, 95 insertions(+), 41 deletions(-) diff --git a/include/diffuse.h b/include/diffuse.h index 4242f1e..6edbb60 100644 --- a/include/diffuse.h +++ b/include/diffuse.h @@ -39,7 +39,7 @@ enum class diffuse_remasking { RANDOM, ENTROPY_EXIT, // Unmask all low-entropy tokens early (semantic scheduling) MASKGIT_PLUS, // Dream: unmask highest top-1 confidence (similar to LOW_CONFIDENCE) - TOPK_MARGIN, // Dream: unmask by margin between top-1 and top-2 logits + TOPK_MARGIN, // Dream: unmask by probability margin between top-1 and top-2 }; struct diffuse_sampler_params { diff --git a/src/diffuse-sampler.cpp b/src/diffuse-sampler.cpp index 035a7d8..72c867b 100644 --- a/src/diffuse-sampler.cpp +++ b/src/diffuse-sampler.cpp @@ -30,30 +30,36 @@ static int tokens_to_unmask(int step, int total_steps, int total_masked, return std::min(n, total_masked); } -// ── Compute entropy of logit distribution ───────────────────── -static float compute_entropy(const float * logits, int n_vocab) { - // Find max for numerical stability - float max_val = logits[0]; +// ── Compute normalized logit statistics ─────────────────── +diffuse_logit_stats diffuse_compute_logit_stats( + const float * logits, int n_vocab) { + int top_token = 0; + float top_logit = logits[0]; + float second_logit = -INFINITY; for (int v = 1; v < n_vocab; v++) { - if (logits[v] > max_val) max_val = logits[v]; + if (logits[v] > top_logit) { + second_logit = top_logit; + top_logit = logits[v]; + top_token = v; + } else if (logits[v] > second_logit) { + second_logit = logits[v]; + } } - // Softmax + entropy in one pass float sum_exp = 0.0f; + float weighted_delta = 0.0f; for (int v = 0; v < n_vocab; v++) { - sum_exp += expf(logits[v] - max_val); + const float delta = logits[v] - top_logit; + const float weight = expf(delta); + sum_exp += weight; + weighted_delta += weight * delta; } - float log_sum = logf(sum_exp); - float entropy = 0.0f; - for (int v = 0; v < n_vocab; v++) { - float log_p = (logits[v] - max_val) - log_sum; - float p = expf(log_p); - if (p > 1e-10f) { - entropy -= p * log_p; - } - } - return entropy; + const float top_probability = 1.0f / sum_exp; + const float second_probability = expf(second_logit - top_logit) / sum_exp; + const float entropy = logf(sum_exp) - weighted_delta / sum_exp; + return {top_token, top_probability, + top_probability - second_probability, entropy}; } // ── Iterative unmasking sampler with inter-step caching ───────── @@ -234,25 +240,14 @@ std::vector diffuse_sample( // Dream: shifted logits (position i uses logits from position max(i-1, 0)) int logit_pos = shift_logits ? std::max(i - 1, 0) : i; const float * logit_row = logit_source + (size_t)logit_pos * n_vocab; - float ent = compute_entropy(logit_row, n_vocab); if (params.temperature <= 0.0f) { - // Find top-1 (and top-2 for TOPK_MARGIN) - int best = 0; - float best_val = logit_row[0]; - float second_val = -1e30f; - for (int v = 1; v < n_vocab; v++) { - if (logit_row[v] > best_val) { - second_val = best_val; - best_val = logit_row[v]; - best = v; - } else if (logit_row[v] > second_val) { - second_val = logit_row[v]; - } - } - float conf = (params.remasking == diffuse_remasking::TOPK_MARGIN) - ? (best_val - second_val) : best_val; - candidates.push_back({i, best, conf, ent}); + const auto stats = diffuse_compute_logit_stats(logit_row, n_vocab); + const float confidence = + params.remasking == diffuse_remasking::TOPK_MARGIN + ? stats.probability_margin : stats.top_probability; + candidates.push_back( + {i, stats.top_token, confidence, stats.entropy}); } else { float max_logit = *std::max_element(logit_row, logit_row + n_vocab); std::vector probs(n_vocab); @@ -265,7 +260,28 @@ std::vector diffuse_sample( std::discrete_distribution dist(probs.begin(), probs.end()); int sampled = dist(rng); - candidates.push_back({i, sampled, probs[sampled], ent}); + float entropy = 0.0f; + for (float probability : probs) { + if (probability > 1e-10f) { + entropy -= probability * logf(probability); + } + } + + float confidence = probs[sampled]; + if (params.remasking == diffuse_remasking::TOPK_MARGIN) { + float top_probability = 0.0f; + float second_probability = 0.0f; + for (float probability : probs) { + if (probability > top_probability) { + second_probability = top_probability; + top_probability = probability; + } else if (probability > second_probability) { + second_probability = probability; + } + } + confidence = top_probability - second_probability; + } + candidates.push_back({i, sampled, confidence, entropy}); } } auto t3 = clk::now(); @@ -292,16 +308,16 @@ std::vector diffuse_sample( n_unmask = std::min(n_unmask, (int)candidates.size()); } else if (params.remasking == diffuse_remasking::LOW_CONFIDENCE || params.remasking == diffuse_remasking::MASKGIT_PLUS) { - // Both sort by highest confidence (top-1 logit value) + // Both sort by highest normalized top-1 confidence. std::sort(candidates.begin(), candidates.end(), [](const candidate & a, const candidate & b) { return a.confidence > b.confidence; }); } else if (params.remasking == diffuse_remasking::TOPK_MARGIN) { - // Sort by margin between top-1 and top-2 logits (highest margin first) + // Sort by top-1/top-2 probability margin (highest first). std::sort(candidates.begin(), candidates.end(), [](const candidate & a, const candidate & b) { - return a.confidence > b.confidence; // margin stored in confidence field + return a.confidence > b.confidence; }); } else { std::shuffle(candidates.begin(), candidates.end(), rng); diff --git a/src/diffuse-sampler.h b/src/diffuse-sampler.h index 22f78e9..eb3ca62 100644 --- a/src/diffuse-sampler.h +++ b/src/diffuse-sampler.h @@ -3,6 +3,19 @@ #include "diffuse-common.h" #include "diffuse.h" +struct diffuse_logit_stats { + int top_token; + float top_probability; + float probability_margin; + float entropy; +}; + +// Compute normalized confidence statistics without materializing a full +// probability vector. Exposed internally so the sampler policy can be tested. +diffuse_logit_stats diffuse_compute_logit_stats( + const float * logits, + int n_vocab); + // Run the full iterative unmasking diffusion loop. std::vector diffuse_sample( diffuse_context * ctx, diff --git a/tests/test-forward.cpp b/tests/test-forward.cpp index e535445..c690fac 100644 --- a/tests/test-forward.cpp +++ b/tests/test-forward.cpp @@ -1,4 +1,5 @@ #include "diffuse.h" +#include "diffuse-sampler.h" #include #include @@ -33,14 +34,38 @@ int main(int argc, char ** argv) { fprintf(stderr, "PASS: sampler params defaults\n"); } - // Test 3: Model load failure on non-existent file + // Test 3: Confidence is a normalized probability, not a raw logit. + { + const float confident_logits[] = {logf(3.0f), 0.0f}; + const float offset_logits[] = { + logf(1.5f) + 100.0f, + 100.0f, + }; + const auto confident = diffuse_compute_logit_stats(confident_logits, 2); + const auto offset = diffuse_compute_logit_stats(offset_logits, 2); + ASSERT(confident.top_token == 0, + "logit stats should select the highest logit"); + ASSERT(fabsf(confident.top_probability - 0.75f) < 1e-6f, + "logit stats should report normalized top-1 confidence"); + ASSERT(fabsf(confident.probability_margin - 0.5f) < 1e-6f, + "logit stats should report the top-1/top-2 probability margin"); + ASSERT(offset_logits[0] > confident_logits[0], + "test rows should demonstrate misleading raw-logit ordering"); + ASSERT(confident.top_probability > offset.top_probability, + "normalized confidence should reverse the raw-logit ordering"); + ASSERT(confident.probability_margin > offset.probability_margin, + "probability margin should reverse the raw-logit ordering"); + fprintf(stderr, "PASS: normalized sampler confidence\n"); + } + + // Test 4: Model load failure on non-existent file { // This should fail gracefully (exit with error from DIFFUSE_DIE) // We just verify the function signature compiles fprintf(stderr, "PASS: API compilation check\n"); } - // Test 4: Full forward pass (requires model file) + // Test 5: Full forward pass (requires model file) if (argc > 1) { const char * model_path = argv[1]; fprintf(stderr, "Running forward pass test with model: %s\n", model_path); From 53f1a508cf3b0a294a6b18ea0a05d38b33399857 Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 14:20:32 +0000 Subject: [PATCH 03/10] Preserve Dream logit sources in sparse cache Dream shifts logits right, so row i predicts token position i + 1. The sparse active set retained a masked position but could cache the predecessor row that supplies its logits. Pass the model shift policy into active-set selection and retain those predecessor rows. Add a model-free CTest that distinguishes shifted and unshifted cache policies. Assisted-by: OpenAI Codex --- CMakeLists.txt | 4 ++++ src/diffuse-cache.h | 9 ++++++-- src/diffuse-sampler.cpp | 2 +- tests/test-cache.cpp | 46 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 3 deletions(-) create mode 100644 tests/test-cache.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 63bf279..fdb8a54 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -61,6 +61,10 @@ if(DIFFUSE_BUILD_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-validate-logits.py) endif() + add_executable(test-cache tests/test-cache.cpp) + target_link_libraries(test-cache PRIVATE diffuse) + add_test(NAME test-cache COMMAND test-cache) + add_executable(test-forward tests/test-forward.cpp) target_link_libraries(test-forward PRIVATE diffuse) add_test(NAME test-forward COMMAND test-forward) diff --git a/src/diffuse-cache.h b/src/diffuse-cache.h index 66168b4..994175d 100644 --- a/src/diffuse-cache.h +++ b/src/diffuse-cache.h @@ -72,7 +72,8 @@ struct diffuse_step_cache { // embeddings changed since the last step OR that changed recently // (within extra_active_steps). // - // Specifically: {still masked} ∪ {changed since last step} + // Specifically: {still masked} ∪ {logit sources for masked positions} + // ∪ {changed since last step} // ∪ {changed within last extra_active_steps steps} // Cached positions: all others (prompt + stably unmasked) @@ -82,6 +83,7 @@ struct diffuse_step_cache { int n_total, int current_step, int extra_active_steps, + bool shift_logits, // outputs: std::vector & cached_positions, // original indices, sorted std::vector & active_positions, // original indices, sorted @@ -92,7 +94,10 @@ struct diffuse_step_cache { for (int i = 0; i < n_total; i++) { bool token_changed = (seq[i] != prev_seq[i]); - bool needs_logits = is_masked[i]; + // Dream shifts logits right by one: row i predicts position i+1. + // Keep those source rows active as masked positions become sparse. + bool needs_logits = is_masked[i] || + (shift_logits && i + 1 < n_total && is_masked[i + 1]); // Track when this position last changed if (token_changed) { diff --git a/src/diffuse-sampler.cpp b/src/diffuse-sampler.cpp index 72c867b..4806e1a 100644 --- a/src/diffuse-sampler.cpp +++ b/src/diffuse-sampler.cpp @@ -163,7 +163,7 @@ std::vector diffuse_sample( } else { // ── Steps 1+: cached forward, only active positions ── cache.compute_active_set(seq.data(), is_masked, total_len, - step, cache_keep_active, + step, cache_keep_active, shift_logits, cached_positions, active_positions, active_to_orig); n_active = (int)active_positions.size(); diff --git a/tests/test-cache.cpp b/tests/test-cache.cpp new file mode 100644 index 0000000..7bb4482 --- /dev/null +++ b/tests/test-cache.cpp @@ -0,0 +1,46 @@ +#include "diffuse-cache.h" + +#include +#include + +#define ASSERT(cond, msg) do { \ + if (!(cond)) { \ + fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); \ + return 1; \ + } \ +} while (0) + +int main() { + constexpr int n_tokens = 6; + const int32_t tokens[n_tokens] = {1, 1, 1, 1, 1, 1}; + + diffuse_step_cache cache; + cache.init(n_tokens, 2, 1, 1, 1); + cache.update_seq(tokens, n_tokens); + + const std::vector is_masked = { + false, false, false, true, false, true, + }; + std::vector cached_positions; + std::vector active_positions; + std::vector active_to_orig; + + cache.compute_active_set( + tokens, is_masked, n_tokens, 1, 0, false, + cached_positions, active_positions, active_to_orig); + ASSERT(active_positions == std::vector({3, 5}), + "unshifted models should keep masked positions active"); + + cache.compute_active_set( + tokens, is_masked, n_tokens, 1, 0, true, + cached_positions, active_positions, active_to_orig); + ASSERT(active_positions == std::vector({2, 3, 4, 5}), + "Dream should also keep each masked position's logit source active"); + ASSERT(cached_positions == std::vector({0, 1}), + "Dream should still cache unrelated stable positions"); + ASSERT(active_to_orig == active_positions, + "active indices should map back to their original positions"); + + fprintf(stderr, "PASS: shifted-logit cache active set\n"); + return 0; +} From 65dc54509f481f1bc20fa85101aa12b87d4a2ca9 Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 14:20:33 +0000 Subject: [PATCH 04/10] Use tokenizer metadata for mask removal generate.py removed one hard-coded LLaDA mask ID from every model output. Supported tokenizers can assign a different mask ID, leaving unresolved masks in decoded output or removing an unrelated token. Use tokenizer.mask_token_id when available and add model-free tests for tokenizer-specific and missing mask metadata. Assisted-by: OpenAI Codex --- CMakeLists.txt | 4 ++++ tests/test-generate.py | 31 +++++++++++++++++++++++++++++++ tools/generate.py | 13 ++++++++++--- 3 files changed, 45 insertions(+), 3 deletions(-) create mode 100644 tests/test-generate.py diff --git a/CMakeLists.txt b/CMakeLists.txt index fdb8a54..bd18e20 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -59,6 +59,10 @@ if(DIFFUSE_BUILD_TESTS) NAME test-validate-logits COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-validate-logits.py) + add_test( + NAME test-generate + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-generate.py) endif() add_executable(test-cache tests/test-cache.cpp) diff --git a/tests/test-generate.py b/tests/test-generate.py new file mode 100644 index 0000000..2a88e1c --- /dev/null +++ b/tests/test-generate.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""Regression tests for tokenizer-specific mask removal.""" + +import importlib.util +import os +import unittest + + +PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +GENERATOR_PATH = os.path.join(PROJECT_DIR, "tools", "generate.py") +GENERATOR_SPEC = importlib.util.spec_from_file_location( + "diffuse_generate", GENERATOR_PATH) +GENERATOR = importlib.util.module_from_spec(GENERATOR_SPEC) +GENERATOR_SPEC.loader.exec_module(GENERATOR) + + +class GenerateTest(unittest.TestCase): + def test_remove_mask_tokens_uses_tokenizer_id(self): + tokens = [7, 200, 8, 126336, 9] + self.assertEqual( + GENERATOR.remove_mask_tokens(tokens, 200), + [7, 8, 126336, 9], + ) + + def test_missing_mask_metadata_preserves_tokens(self): + tokens = [7, 8, 9] + self.assertIs(GENERATOR.remove_mask_tokens(tokens, None), tokens) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/generate.py b/tools/generate.py index 29b58c0..7657164 100644 --- a/tools/generate.py +++ b/tools/generate.py @@ -17,6 +17,13 @@ import sys +def remove_mask_tokens(token_ids, mask_token_id): + """Remove unresolved masks using the selected tokenizer's metadata.""" + if mask_token_id is None: + return token_ids + return [token_id for token_id in token_ids if token_id != mask_token_id] + + def main(): parser = argparse.ArgumentParser(description="Generate text with diffuse-cpp") parser.add_argument("--model-dir", "-m", required=True, @@ -136,9 +143,9 @@ def main(): output_ids = [int(x) for x in output_line.split(",")] # --- Detokenize --- - # Filter out mask tokens and special tokens - mask_id = 126336 - output_ids_clean = [t for t in output_ids if t != mask_id] + # Filter out unresolved mask tokens before decoding. Different supported + # model tokenizers assign different IDs to the mask token. + output_ids_clean = remove_mask_tokens(output_ids, tokenizer.mask_token_id) output_text = tokenizer.decode(output_ids_clean, skip_special_tokens=True) From 6374274151204ecb96662210a3d7cac09924d74d Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 04:05:47 +0000 Subject: [PATCH 05/10] feat: expose available GGML devices Assisted-by: OpenAI Codex --- CMakeLists.txt | 5 +++++ include/diffuse.h | 21 +++++++++++++++++++ src/diffuse-backend.cpp | 45 +++++++++++++++++++++++++++++++++++++++++ tests/test-backend.cpp | 37 +++++++++++++++++++++++++++++++++ tools/main-cli.cpp | 19 ++++++++++++++++- 5 files changed, 126 insertions(+), 1 deletion(-) create mode 100644 src/diffuse-backend.cpp create mode 100644 tests/test-backend.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index bd18e20..1855021 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -17,6 +17,7 @@ add_subdirectory(ggml) # ── Core library (diffusion + autoregressive) ───────────────── add_library(diffuse + src/diffuse-backend.cpp src/diffuse-model.cpp src/diffuse-graph.cpp src/diffuse-sampler.cpp @@ -69,6 +70,10 @@ if(DIFFUSE_BUILD_TESTS) target_link_libraries(test-cache PRIVATE diffuse) add_test(NAME test-cache COMMAND test-cache) + add_executable(test-backend tests/test-backend.cpp) + target_link_libraries(test-backend PRIVATE diffuse) + add_test(NAME test-backend COMMAND test-backend) + add_executable(test-forward tests/test-forward.cpp) target_link_libraries(test-forward PRIVATE diffuse) add_test(NAME test-forward COMMAND test-forward) diff --git a/include/diffuse.h b/include/diffuse.h index 6edbb60..a70fa30 100644 --- a/include/diffuse.h +++ b/include/diffuse.h @@ -2,6 +2,7 @@ // diffuse-cpp public API +#include #include #include #include @@ -11,6 +12,26 @@ struct diffuse_model; struct diffuse_context; +enum class diffuse_device_type { + CPU, + GPU, + IGPU, + ACCEL, + UNKNOWN, +}; + +struct diffuse_device_info { + std::string name; + std::string description; + diffuse_device_type type = diffuse_device_type::UNKNOWN; + size_t memory_free = 0; + size_t memory_total = 0; + bool async = false; +}; + +std::vector diffuse_available_devices(); +const char * diffuse_device_type_name(diffuse_device_type type); + // ── Hyperparameters ──────────────────────────────────────────── struct diffuse_hparams { uint32_t n_vocab = 0; diff --git a/src/diffuse-backend.cpp b/src/diffuse-backend.cpp new file mode 100644 index 0000000..99d2c84 --- /dev/null +++ b/src/diffuse-backend.cpp @@ -0,0 +1,45 @@ +#include "diffuse.h" + +#include "ggml-backend.h" + +static diffuse_device_type to_diffuse_device_type(enum ggml_backend_dev_type type) { + switch (type) { + case GGML_BACKEND_DEVICE_TYPE_CPU: return diffuse_device_type::CPU; + case GGML_BACKEND_DEVICE_TYPE_GPU: return diffuse_device_type::GPU; + case GGML_BACKEND_DEVICE_TYPE_IGPU: return diffuse_device_type::IGPU; + case GGML_BACKEND_DEVICE_TYPE_ACCEL: return diffuse_device_type::ACCEL; + } + return diffuse_device_type::UNKNOWN; +} + +const char * diffuse_device_type_name(diffuse_device_type type) { + switch (type) { + case diffuse_device_type::CPU: return "CPU"; + case diffuse_device_type::GPU: return "GPU"; + case diffuse_device_type::IGPU: return "IGPU"; + case diffuse_device_type::ACCEL: return "ACCEL"; + case diffuse_device_type::UNKNOWN: break; + } + return "UNKNOWN"; +} + +std::vector diffuse_available_devices() { + std::vector devices; + devices.reserve(ggml_backend_dev_count()); + + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_props props; + ggml_backend_dev_get_props(ggml_backend_dev_get(i), &props); + + diffuse_device_info info; + info.name = props.name != nullptr ? props.name : ""; + info.description = props.description != nullptr ? props.description : ""; + info.type = to_diffuse_device_type(props.type); + info.memory_free = props.memory_free; + info.memory_total = props.memory_total; + info.async = props.caps.async; + devices.push_back(info); + } + + return devices; +} diff --git a/tests/test-backend.cpp b/tests/test-backend.cpp new file mode 100644 index 0000000..17cf002 --- /dev/null +++ b/tests/test-backend.cpp @@ -0,0 +1,37 @@ +#include "diffuse.h" + +#include +#include +#include + +#define ASSERT(cond, msg) do { \ + if (!(cond)) { \ + fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); \ + return 1; \ + } \ +} while (0) + +int main() { + const auto devices = diffuse_available_devices(); + ASSERT(!devices.empty(), "at least one compute device should be available"); + + bool found_cpu = false; + std::set names; + + for (const auto & device : devices) { + ASSERT(!device.name.empty(), "device name should not be empty"); + ASSERT(!device.description.empty(), "device description should not be empty"); + ASSERT(device.type != diffuse_device_type::UNKNOWN, "device type should be recognized"); + ASSERT(device.memory_free <= device.memory_total, "free memory should not exceed total memory"); + ASSERT(names.insert(device.name).second, "device names should be unique"); + + found_cpu = found_cpu || device.type == diffuse_device_type::CPU; + printf("%-8s %-16s %s\n", + diffuse_device_type_name(device.type), + device.name.c_str(), + device.description.c_str()); + } + + ASSERT(found_cpu, "CPU fallback device should be available"); + return 0; +} diff --git a/tools/main-cli.cpp b/tools/main-cli.cpp index 9cf6e1d..34f8952 100644 --- a/tools/main-cli.cpp +++ b/tools/main-cli.cpp @@ -10,6 +10,7 @@ static void print_usage(const char * prog) { fprintf(stderr, "Usage: %s [options]\n", prog); fprintf(stderr, "\nOptions:\n"); fprintf(stderr, " -m PATH Model file (GGUF)\n"); + fprintf(stderr, " --list-devices List available compute devices and exit\n"); fprintf(stderr, " -p TEXT Prompt text\n"); fprintf(stderr, " -n INT Tokens to generate (default: 128)\n"); fprintf(stderr, " -s INT Diffusion steps (default: 32)\n"); @@ -52,9 +53,12 @@ int main(int argc, char ** argv) { bool use_cache = true; int cache_refresh = 0; int cache_keep_active = 0; + bool list_devices = false; for (int i = 1; i < argc; i++) { - if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) { + if (strcmp(argv[i], "--list-devices") == 0) { + list_devices = true; + } else if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) { model_path = argv[++i]; } else if (strcmp(argv[i], "-p") == 0 && i + 1 < argc) { prompt = argv[++i]; @@ -93,6 +97,19 @@ int main(int argc, char ** argv) { } } + if (list_devices) { + const auto devices = diffuse_available_devices(); + for (const auto & device : devices) { + printf("%-8s %-16s %s (%zu MiB free, %zu MiB total)\n", + diffuse_device_type_name(device.type), + device.name.c_str(), + device.description.c_str(), + device.memory_free / (1024 * 1024), + device.memory_total / (1024 * 1024)); + } + return devices.empty() ? 1 : 0; + } + if (model_path.empty()) { fprintf(stderr, "Error: model path required (-m)\n"); print_usage(argv[0]); From 8919abd4303ddae1f8721caf679f014ebce64eae Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 04:28:19 +0000 Subject: [PATCH 06/10] feat: offload diffusion inference through GGML Assisted-by: OpenAI Codex --- README.md | 36 ++++++++++++++- include/diffuse.h | 2 + src/ar-graph.cpp | 5 +++ src/diffuse-common.h | 5 +++ src/diffuse-graph.cpp | 98 ++++++++++++++++++++++++++++++++++++++--- src/diffuse-model.cpp | 95 +++++++++++++++++++++++++++++++++++++-- src/diffuse-model.h | 3 +- src/diffuse-sampler.cpp | 5 ++- tests/test-e2e.py | 31 +++++++++++-- tests/test-forward.cpp | 48 ++++++++++++++++++++ tools/convert-llada.py | 25 +++++++---- tools/main-bench.cpp | 55 ++++++++++++++++++++--- tools/main-cli.cpp | 15 ++++++- 13 files changed, 390 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 2da63eb..0e5d342 100644 --- a/README.md +++ b/README.md @@ -119,6 +119,39 @@ cmake --build build -j$(nproc) --remasking entropy_exit ``` +### GPU offload (experimental) + +Build GGML with the backend for your platform. Metal is enabled by default on +macOS; Vulkan and CUDA are opt-in: + +```bash +# Linux / Windows with Vulkan +cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_VULKAN=ON + +# NVIDIA CUDA +cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=ON + +cmake --build build -j +``` + +List the exact GGML device names available in that build, then select one: + +```bash +./build/diffuse-cli --list-devices +./build/diffuse-cli -m model.gguf --device Vulkan0 --tokens "..." -n 64 -s 16 +``` + +The selected device holds the model weights and receives the supported graph +operations. The CLI reports model allocation and per-backend graph node counts +so an offload cannot silently become CPU-only execution. Integrated GPUs are +valid devices and are reported separately as `IGPU`. + +GPU offload currently supports the masked-diffusion forward path. The +inter-step KV cache is disabled for offloaded generation, autoregressive +execution is not yet supported, and the full model must fit the selected +device's buffer type. The GGML scheduler keeps a CPU backend available for +unsupported graph operations and reports any such placement. + **Note**: diffuse-cpp operates on token IDs, not raw text. Use the HuggingFace transformers library to tokenize your prompts: ```python @@ -301,7 +334,8 @@ Current limitations: - No integrated tokenizer (use transformers) - Default 256 generated tokens per call (configurable via -n flag) - Single-model inference only (no batching) -- CPU-only (GPU support via GGML is possible but not prioritized) +- GPU offload is experimental and does not yet support the inter-step cache or + autoregressive execution ## Contributing diff --git a/include/diffuse.h b/include/diffuse.h index a70fa30..36815c5 100644 --- a/include/diffuse.h +++ b/include/diffuse.h @@ -90,6 +90,8 @@ using diffuse_step_callback = std::functionsched != nullptr) { + DIFFUSE_LOG("autoregressive execution is not yet supported with device offload"); + return false; + } + const auto & hp = ctx->model->hparams; const int n_past = cache->n_past; const int n_kv = n_past + n_new; diff --git a/src/diffuse-common.h b/src/diffuse-common.h index fb978a1..9fd10df 100644 --- a/src/diffuse-common.h +++ b/src/diffuse-common.h @@ -2,6 +2,7 @@ #include "diffuse.h" #include "ggml.h" +#include "ggml-alloc.h" #include "gguf.h" #include "ggml-backend.h" #include "ggml-cpu.h" @@ -63,8 +64,10 @@ struct diffuse_model { // GGML backend ggml_backend_t backend = nullptr; + ggml_backend_t cpu_backend = nullptr; ggml_backend_buffer_t buf = nullptr; struct ggml_context * ctx = nullptr; // weight context + std::string device_name = "CPU"; }; // ── Compute context ──────────────────────────────────────────── @@ -75,5 +78,7 @@ struct diffuse_context { ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; + ggml_backend_sched_t sched = nullptr; struct ggml_context * ctx = nullptr; // compute context + bool placement_logged = false; }; diff --git a/src/diffuse-graph.cpp b/src/diffuse-graph.cpp index 128dc89..5d22ff6 100644 --- a/src/diffuse-graph.cpp +++ b/src/diffuse-graph.cpp @@ -29,12 +29,14 @@ struct ggml_cgraph * diffuse_build_graph( struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N); ggml_set_name(inp_tokens, "inp_tokens"); ggml_set_input(inp_tokens); - memcpy(inp_tokens->data, tokens, N * sizeof(int32_t)); + if (inp_tokens->data != nullptr) { + memcpy(inp_tokens->data, tokens, N * sizeof(int32_t)); + } struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N); ggml_set_name(inp_pos, "inp_pos"); ggml_set_input(inp_pos); - { + if (inp_pos->data != nullptr) { int32_t * pos_data = (int32_t *)inp_pos->data; for (int i = 0; i < N; i++) pos_data[i] = i; } @@ -306,12 +308,66 @@ static struct ggml_cgraph * diffuse_build_graph_extractable( } // ── Full forward pass with cache extraction ───────────────────── +static size_t diffuse_max_graph_nodes(const diffuse_model * model) { + return (size_t)model->hparams.n_layer * 4096 + + (size_t)model->hparams.n_layer * 8 + 256; +} + +static bool diffuse_sched_compute(diffuse_context * ctx, + struct ggml_cgraph * graph, + const int32_t * tokens, + int n_tokens) { + ggml_backend_sched_reset(ctx->sched); + if (!ggml_backend_sched_alloc_graph(ctx->sched, graph)) { + DIFFUSE_LOG("failed to allocate scheduled compute graph"); + return false; + } + + struct ggml_tensor * inp_tokens = ggml_graph_get_tensor(graph, "inp_tokens"); + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(graph, "inp_pos"); + if (inp_tokens == nullptr || inp_pos == nullptr) { + DIFFUSE_LOG("scheduled graph inputs not found"); + return false; + } + + std::vector positions(n_tokens); + for (int i = 0; i < n_tokens; ++i) positions[i] = i; + ggml_backend_tensor_set(inp_tokens, tokens, 0, n_tokens * sizeof(int32_t)); + ggml_backend_tensor_set(inp_pos, positions.data(), 0, n_tokens * sizeof(int32_t)); + + if (!ctx->placement_logged) { + int device_nodes = 0; + int cpu_nodes = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + ggml_backend_t backend = ggml_backend_sched_get_tensor_backend( + ctx->sched, ggml_graph_node(graph, i)); + if (backend == ctx->model->backend) ++device_nodes; + if (backend == ctx->model->cpu_backend) ++cpu_nodes; + } + DIFFUSE_LOG(" graph placement: %s=%d nodes, CPU=%d nodes", + ctx->model->device_name.c_str(), device_nodes, cpu_nodes); + if (device_nodes == 0) { + DIFFUSE_LOG("scheduled graph did not place any operations on %s", + ctx->model->device_name.c_str()); + return false; + } + ctx->placement_logged = true; + } + + return ggml_backend_sched_graph_compute(ctx->sched, graph) == GGML_STATUS_SUCCESS; +} + bool diffuse_forward_full(diffuse_context * ctx, const int32_t * tokens, int n_tokens, float * logits_out, diffuse_step_cache * cache) { const auto & hp = ctx->model->hparams; + if (ctx->sched != nullptr && cache != nullptr) { + DIFFUSE_LOG("inter-step cache is not yet supported with device offload"); + return false; + } + // Buffer size: same as original + extra for named K,V output tensors size_t per_layer = (size_t)n_tokens * hp.n_embd * sizeof(float) * 10 + (size_t)n_tokens * hp.n_ff * sizeof(float) * 3 @@ -322,10 +378,16 @@ bool diffuse_forward_full(diffuse_context * ctx, buf_size += 256ull * 1024 * 1024; buf_size = (size_t)(buf_size * 1.5); + const bool scheduled = ctx->sched != nullptr; + const size_t graph_nodes = diffuse_max_graph_nodes(ctx->model); + const size_t context_size = scheduled + ? ggml_tensor_overhead() * graph_nodes + ggml_graph_overhead_custom(graph_nodes, false) + : buf_size; + struct ggml_init_params cparams = { - /*.mem_size = */ buf_size, + /*.mem_size = */ context_size, /*.mem_buffer = */ nullptr, - /*.no_alloc = */ false, + /*.no_alloc = */ scheduled, }; struct ggml_context * ctx_compute = ggml_init(cparams); if (!ctx_compute) { @@ -337,7 +399,14 @@ bool diffuse_forward_full(diffuse_context * ctx, ? diffuse_build_graph_extractable(ctx, ctx_compute, tokens, n_tokens) : diffuse_build_graph(ctx, ctx_compute, tokens, n_tokens); - enum ggml_status status = ggml_graph_compute_with_ctx(ctx_compute, gf, ctx->n_threads); + enum ggml_status status = GGML_STATUS_SUCCESS; + if (scheduled) { + if (!diffuse_sched_compute(ctx, gf, tokens, n_tokens)) { + status = GGML_STATUS_FAILED; + } + } else { + status = ggml_graph_compute_with_ctx(ctx_compute, gf, ctx->n_threads); + } if (status != GGML_STATUS_SUCCESS) { DIFFUSE_LOG("graph compute failed with status %d", (int)status); ggml_free(ctx_compute); @@ -351,7 +420,12 @@ bool diffuse_forward_full(diffuse_context * ctx, ggml_free(ctx_compute); return false; } - memcpy(logits_out, logits->data, (size_t)n_tokens * hp.n_vocab * sizeof(float)); + const size_t logits_bytes = (size_t)n_tokens * hp.n_vocab * sizeof(float); + if (scheduled) { + ggml_backend_tensor_get(logits, logits_out, 0, logits_bytes); + } else { + memcpy(logits_out, logits->data, logits_bytes); + } // Extract K,V into cache if (cache) { @@ -685,11 +759,23 @@ diffuse_context * diffuse_context_new(const diffuse_model * model, int n_ctx, in ctx->model = model; ctx->n_ctx = n_ctx; ctx->n_threads = n_threads; + + if (model->backend != nullptr) { + ggml_backend_t backends[] = { model->backend, model->cpu_backend }; + ctx->sched = ggml_backend_sched_new( + backends, nullptr, 2, diffuse_max_graph_nodes(model), false, true); + if (ctx->sched == nullptr) { + delete ctx; + return nullptr; + } + ctx->backend = model->backend; + } return ctx; } void diffuse_context_free(diffuse_context * ctx) { if (!ctx) return; + if (ctx->sched) ggml_backend_sched_free(ctx->sched); if (ctx->buf) ggml_backend_buffer_free(ctx->buf); if (ctx->ctx) ggml_free(ctx->ctx); delete ctx; diff --git a/src/diffuse-model.cpp b/src/diffuse-model.cpp index 20acd7b..f54b8a5 100644 --- a/src/diffuse-model.cpp +++ b/src/diffuse-model.cpp @@ -1,5 +1,8 @@ #include "diffuse-model.h" +#include +#include + // ── GGUF metadata key helpers ────────────────────────────────── static int64_t find_key(const struct gguf_context * gctx, const char * key) { int64_t id = gguf_find_key(gctx, key); @@ -79,14 +82,62 @@ static struct ggml_tensor * get_tensor(struct ggml_context * ctx, const char * n return t; } +static bool load_backend_weights(const std::string & path, + struct ggml_context * ctx, + const struct gguf_context * gctx) { + std::ifstream file(path, std::ios::binary); + if (!file) return false; + + std::vector read_buf(4 * 1024 * 1024); + const int64_t n_tensors = gguf_get_n_tensors(gctx); + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(gctx, i); + struct ggml_tensor * tensor = ggml_get_tensor(ctx, name); + if (tensor == nullptr) continue; + + const size_t offset = gguf_get_data_offset(gctx) + gguf_get_tensor_offset(gctx, i); + file.seekg(static_cast(offset)); + if (!file) return false; + + const size_t nbytes = ggml_nbytes(tensor); + for (size_t pos = 0; pos < nbytes; pos += read_buf.size()) { + const size_t chunk = std::min(read_buf.size(), nbytes - pos); + file.read(read_buf.data(), static_cast(chunk)); + if (!file) return false; + ggml_backend_tensor_set(tensor, read_buf.data(), pos, chunk); + } + } + + return true; +} + // ── Load model from GGUF ─────────────────────────────────────── -diffuse_model * diffuse_model_load_impl(const std::string & path, int n_threads) { +diffuse_model * diffuse_model_load_impl(const std::string & path, int n_threads, + const std::string & device_name) { DIFFUSE_LOG("loading model from %s", path.c_str()); + ggml_backend_dev_t selected_device = nullptr; + bool use_device_backend = false; + if (!device_name.empty()) { + selected_device = ggml_backend_dev_by_name(device_name.c_str()); + if (selected_device == nullptr) { + DIFFUSE_DIE("compute device not found: %s (use --list-devices)", device_name.c_str()); + } + + const auto type = ggml_backend_dev_type(selected_device); + use_device_backend = type == GGML_BACKEND_DEVICE_TYPE_GPU || + type == GGML_BACKEND_DEVICE_TYPE_IGPU; + if (!use_device_backend && type != GGML_BACKEND_DEVICE_TYPE_CPU) { + DIFFUSE_DIE("compute device %s cannot store model weights", + ggml_backend_dev_name(selected_device)); + } + } + // Open GGUF file — gguf_init_from_file allocates a ggml_context with tensors struct ggml_context * meta_ctx = nullptr; struct gguf_init_params gparams = { - /*.no_alloc = */ false, + /*.no_alloc = */ use_device_backend, /*.ctx = */ &meta_ctx, }; struct gguf_context * gctx = gguf_init_from_file(path.c_str(), gparams); @@ -177,6 +228,38 @@ diffuse_model * diffuse_model_load_impl(const std::string & path, int n_threads) l.bv = ggml_get_tensor(meta_ctx, fmt_layer("blk.%d.attn_v.bias", i).c_str()); } + if (use_device_backend) { + model->backend = ggml_backend_dev_init(selected_device, nullptr); + if (model->backend == nullptr) { + DIFFUSE_DIE("failed to initialize compute device: %s", + ggml_backend_dev_name(selected_device)); + } + + ggml_backend_dev_t cpu_device = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + model->cpu_backend = cpu_device != nullptr + ? ggml_backend_dev_init(cpu_device, nullptr) + : nullptr; + if (model->cpu_backend == nullptr) { + DIFFUSE_DIE("failed to initialize CPU fallback backend"); + } + ggml_backend_cpu_set_n_threads(model->cpu_backend, n_threads); + + model->buf = ggml_backend_alloc_ctx_tensors(meta_ctx, model->backend); + if (model->buf == nullptr) { + DIFFUSE_DIE("failed to allocate model weights on %s", + ggml_backend_dev_name(selected_device)); + } + if (!load_backend_weights(path, meta_ctx, gctx)) { + DIFFUSE_DIE("failed to load model weights onto %s", + ggml_backend_dev_name(selected_device)); + } + + model->device_name = ggml_backend_dev_name(selected_device); + DIFFUSE_LOG(" weights: %s backend, %.2f MiB", + model->device_name.c_str(), + ggml_backend_buffer_get_size(model->buf) / (1024.0 * 1024.0)); + } + DIFFUSE_LOG("model loaded: %lld tensors", (long long)gguf_get_n_tensors(gctx)); gguf_free(gctx); // frees metadata, but ggml_context (weights) stays alive @@ -188,13 +271,19 @@ void diffuse_model_free_impl(diffuse_model * model) { if (!model) return; if (model->buf) ggml_backend_buffer_free(model->buf); if (model->ctx) ggml_free(model->ctx); + if (model->cpu_backend) ggml_backend_free(model->cpu_backend); if (model->backend) ggml_backend_free(model->backend); delete model; } // ── Public API wrappers ──────────────────────────────────────── diffuse_model * diffuse_model_load(const std::string & path, int n_threads) { - return diffuse_model_load_impl(path, n_threads); + return diffuse_model_load_impl(path, n_threads, ""); +} + +diffuse_model * diffuse_model_load(const std::string & path, int n_threads, + const std::string & device_name) { + return diffuse_model_load_impl(path, n_threads, device_name); } void diffuse_model_free(diffuse_model * model) { diff --git a/src/diffuse-model.h b/src/diffuse-model.h index 9c1439d..32550dd 100644 --- a/src/diffuse-model.h +++ b/src/diffuse-model.h @@ -3,7 +3,8 @@ #include "diffuse-common.h" // Load model weights from GGUF file -diffuse_model * diffuse_model_load_impl(const std::string & path, int n_threads); +diffuse_model * diffuse_model_load_impl(const std::string & path, int n_threads, + const std::string & device_name); // Free model and all associated memory void diffuse_model_free_impl(diffuse_model * model); diff --git a/src/diffuse-sampler.cpp b/src/diffuse-sampler.cpp index 4806e1a..f9889dc 100644 --- a/src/diffuse-sampler.cpp +++ b/src/diffuse-sampler.cpp @@ -99,7 +99,10 @@ std::vector diffuse_sample( std::mt19937 rng(params.seed); // ── Inter-step KV cache ────────────────────────────────────── - const bool use_cache = params.use_cache; + const bool use_cache = params.use_cache && ctx->sched == nullptr; + if (params.use_cache && !use_cache) { + DIFFUSE_LOG("inter-step cache disabled for device-offloaded generation"); + } diffuse_step_cache cache; if (use_cache) { cache.init(total_len, prompt_len, diff --git a/tests/test-e2e.py b/tests/test-e2e.py index 04441fe..c73b9d3 100644 --- a/tests/test-e2e.py +++ b/tests/test-e2e.py @@ -2,8 +2,10 @@ """End-to-end test: create tiny model, convert, load in C++, run forward pass.""" import importlib.util +import argparse import json import os +import shutil import sys import tempfile import subprocess @@ -15,17 +17,28 @@ print("SKIP: safetensors not installed") sys.exit(0) -HIDDEN = 64 +HIDDEN = 256 N_HEADS = 4 N_LAYERS = 2 VOCAB = 256 -FF = 128 +FF = 512 MASK_ID = 200 +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--build-dir", default="build") + parser.add_argument("--device") + parser.add_argument("--keep-gguf") + return parser.parse_args() + + def main(): + args = parse_args() script_dir = os.path.dirname(os.path.abspath(__file__)) project_dir = os.path.dirname(script_dir) - build_dir = os.path.join(project_dir, "build") + build_dir = args.build_dir + if not os.path.isabs(build_dir): + build_dir = os.path.join(project_dir, build_dir) test_binary = os.path.join(build_dir, "test-forward") if not os.path.exists(test_binary): @@ -92,8 +105,12 @@ def main(): env = os.environ.copy() env["LD_LIBRARY_PATH"] = f"{build_dir}:{build_dir}/ggml/src" + command = [test_binary, gguf_path] + if args.device: + command.append(args.device) + result = subprocess.run( - [test_binary, gguf_path], + command, capture_output=True, text=True, env=env, timeout=30, ) @@ -123,6 +140,12 @@ def main(): sys.exit(1) print("dump-logits wrapper OK") + if args.keep_gguf: + output_path = os.path.abspath(args.keep_gguf) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + shutil.copy2(gguf_path, output_path) + print(f"Kept synthetic GGUF at {output_path}") + print("\nEnd-to-end test PASSED!") diff --git a/tests/test-forward.cpp b/tests/test-forward.cpp index c690fac..76b4cfa 100644 --- a/tests/test-forward.cpp +++ b/tests/test-forward.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -95,6 +96,53 @@ int main(int argc, char ** argv) { ASSERT(sum > 0.0f, "logits should not be all zeros"); fprintf(stderr, "PASS: forward pass (logits sum=%.2f)\n", sum); + if (argc > 2) { + const char * device_name = argv[2]; + fprintf(stderr, "Comparing device forward pass on: %s\n", device_name); + + diffuse_model * device_model = diffuse_model_load(model_path, 4, device_name); + ASSERT(device_model != nullptr, "device model should load"); + diffuse_context * device_ctx = diffuse_context_new(device_model, n_tokens, 4); + if (device_ctx == nullptr) { + diffuse_model_free(device_model); + ASSERT(false, "device context should initialize"); + } + + std::vector device_logits(logits.size()); + ok = diffuse_forward(device_ctx, tokens.data(), n_tokens, device_logits.data()); + if (!ok) { + diffuse_context_free(device_ctx); + diffuse_model_free(device_model); + ASSERT(false, "device forward pass should succeed"); + } + + float max_abs_diff = 0.0f; + float max_abs_ref = 0.0f; + double squared_error = 0.0; + double squared_reference = 0.0; + bool all_finite = true; + for (size_t i = 0; i < logits.size(); ++i) { + all_finite = all_finite && std::isfinite(device_logits[i]); + const float diff = device_logits[i] - logits[i]; + max_abs_diff = std::max(max_abs_diff, fabsf(diff)); + max_abs_ref = std::max(max_abs_ref, fabsf(logits[i])); + squared_error += (double)diff * diff; + squared_reference += (double)logits[i] * logits[i]; + } + const double nmse = squared_error / squared_reference; + fprintf(stderr, "%s comparison: max abs diff=%.6g, max abs ref=%.6g, NMSE=%.6g\n", + device_name, max_abs_diff, max_abs_ref, nmse); + + diffuse_context_free(device_ctx); + diffuse_model_free(device_model); + + ASSERT(all_finite, "device logits should be finite"); + // GGML uses NMSE for quantized backend comparisons. Allow for + // accumulation across this randomized two-layer graph. + ASSERT(nmse <= 2e-3, "device logits should match CPU reference (NMSE <= 0.002)"); + fprintf(stderr, "PASS: %s logits match CPU\n", device_name); + } + diffuse_context_free(ctx); diffuse_model_free(model); } else { diff --git a/tools/convert-llada.py b/tools/convert-llada.py index d3d7017..3ee0338 100644 --- a/tools/convert-llada.py +++ b/tools/convert-llada.py @@ -21,8 +21,7 @@ try: import torch except ImportError: - print("ERROR: torch not installed. Run: pip install torch", file=sys.stderr) - sys.exit(1) + torch = None try: from safetensors import safe_open @@ -211,7 +210,8 @@ def convert(args): shard_name = os.path.basename(shard_path) print(f"\n── {shard_name} ──") - with safe_open(shard_path, framework="pt") as f: + framework = "pt" if torch is not None else "np" + with safe_open(shard_path, framework=framework) as f: for name in sorted(f.keys()): gguf_name = map_tensor_name(name) if gguf_name is None: @@ -219,12 +219,19 @@ def convert(args): n_skipped += 1 continue - pt_tensor = f.get_tensor(name) - orig_dtype_str = str(pt_tensor.dtype).replace("torch.", "") - orig_shape = tuple(pt_tensor.shape) - - # Convert to float32 numpy (handles bf16 transparently) - tensor = pt_tensor.float().numpy() + source_tensor = f.get_tensor(name) + orig_dtype_str = str(source_tensor.dtype).replace("torch.", "") + orig_shape = tuple(source_tensor.shape) + + # PyTorch handles BF16 conversion; NumPy keeps the lightweight + # F32/F16 conversion and test path free of a PyTorch dependency. + if torch is not None: + tensor = source_tensor.float().numpy() + else: + if source_tensor.dtype.kind == "V": + raise RuntimeError( + "BF16 conversion requires torch; install torch and retry") + tensor = source_tensor.astype(np.float32) # Convert to output type if args.type == "f32": diff --git a/tools/main-bench.cpp b/tools/main-bench.cpp index e63da84..114a4b6 100644 --- a/tools/main-bench.cpp +++ b/tools/main-bench.cpp @@ -48,6 +48,8 @@ struct bench_result { int main(int argc, char ** argv) { std::string model_path; + std::string device_name; + std::string scheduler_name = "all"; std::vector input_tokens; int n_generate = 256; std::vector steps_list = {8, 16, 32}; @@ -55,12 +57,17 @@ int main(int argc, char ** argv) { int n_reps = 3; int n_warmup = 1; bool json_output = false; + bool use_cache = true; float entropy_threshold = 1.5f; int prompt_len = 32; // for dummy prompt if --tokens not given for (int i = 1; i < argc; i++) { if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) { model_path = argv[++i]; + } else if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) { + device_name = argv[++i]; + } else if (strcmp(argv[i], "--scheduler") == 0 && i + 1 < argc) { + scheduler_name = argv[++i]; } else if (strcmp(argv[i], "--tokens") == 0 && i + 1 < argc) { input_tokens = parse_tokens(argv[++i]); } else if (strcmp(argv[i], "-n") == 0 && i + 1 < argc) { @@ -75,12 +82,16 @@ int main(int argc, char ** argv) { n_warmup = atoi(argv[++i]); } else if (strcmp(argv[i], "--json") == 0) { json_output = true; + } else if (strcmp(argv[i], "--no-cache") == 0) { + use_cache = false; } else if (strcmp(argv[i], "--entropy-threshold") == 0 && i + 1 < argc) { entropy_threshold = atof(argv[++i]); } else if (strcmp(argv[i], "--prompt-len") == 0 && i + 1 < argc) { prompt_len = atoi(argv[++i]); } else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { fprintf(stderr, "Usage: %s -m MODEL [options]\n", argv[0]); + fprintf(stderr, " --device NAME Exact GGML device name (default: CPU)\n"); + fprintf(stderr, " --scheduler NAME low_confidence|entropy_exit|maskgit_plus|topk_margin|all\n"); fprintf(stderr, " --tokens IDs Pre-tokenized input (comma-separated)\n"); fprintf(stderr, " --prompt-len N Dummy prompt length if no --tokens (default: 32)\n"); fprintf(stderr, " -n INT Tokens to generate (default: 256)\n"); @@ -89,6 +100,7 @@ int main(int argc, char ** argv) { fprintf(stderr, " -r INT Repetitions (default: 3)\n"); fprintf(stderr, " --warmup INT Warmup runs (default: 1)\n"); fprintf(stderr, " --json Output JSON instead of table\n"); + fprintf(stderr, " --no-cache Disable inter-step cache for matched CPU/GPU work\n"); fprintf(stderr, " --entropy-threshold F (default: 1.5)\n"); return 0; } @@ -98,6 +110,10 @@ int main(int argc, char ** argv) { fprintf(stderr, "Error: -m MODEL required\n"); return 1; } + if (!device_name.empty() && use_cache) { + fprintf(stderr, "Device offload does not yet support the inter-step cache; disabling it\n"); + use_cache = false; + } // If no tokens given, use dummy prompt if (input_tokens.empty()) { @@ -110,18 +126,32 @@ int main(int argc, char ** argv) { const char * name; diffuse_remasking remasking; }; - sched_config schedulers[] = { + const sched_config available_schedulers[] = { {"low_confidence", diffuse_remasking::LOW_CONFIDENCE}, {"entropy_exit", diffuse_remasking::ENTROPY_EXIT}, {"maskgit_plus", diffuse_remasking::MASKGIT_PLUS}, {"topk_margin", diffuse_remasking::TOPK_MARGIN}, }; - int n_schedulers = 4; + std::vector schedulers; + for (const auto & scheduler : available_schedulers) { + if (scheduler_name == "all" || scheduler_name == scheduler.name) { + schedulers.push_back(scheduler); + } + } + if (schedulers.empty()) { + fprintf(stderr, "Error: unknown scheduler: %s\n", scheduler_name.c_str()); + return 1; + } // Load model int max_threads = *std::max_element(threads_list.begin(), threads_list.end()); fprintf(stderr, "Loading model: %s\n", model_path.c_str()); - diffuse_model * model = diffuse_model_load(model_path, max_threads); + auto load_start = std::chrono::high_resolution_clock::now(); + diffuse_model * model = device_name.empty() + ? diffuse_model_load(model_path, max_threads) + : diffuse_model_load(model_path, max_threads, device_name); + auto load_end = std::chrono::high_resolution_clock::now(); + double load_ms = std::chrono::duration(load_end - load_start).count(); if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; @@ -132,16 +162,19 @@ int main(int argc, char ** argv) { fprintf(stderr, " n_vocab=%d n_embd=%d n_layer=%d\n", hp.n_vocab, hp.n_embd, hp.n_layer); fprintf(stderr, " prompt=%d gen=%d\n", (int)input_tokens.size(), n_generate); + fprintf(stderr, " device=%s cache=%s load=%.1fms\n", + device_name.empty() ? "CPU" : device_name.c_str(), + use_cache ? "ON" : "OFF", load_ms); fprintf(stderr, " steps=%zu configs, threads=%zu configs, reps=%d, warmup=%d\n", steps_list.size(), threads_list.size(), n_reps, n_warmup); fprintf(stderr, "\n"); std::vector results; - int total_configs = n_schedulers * (int)steps_list.size() * (int)threads_list.size(); + int total_configs = (int)schedulers.size() * (int)steps_list.size() * (int)threads_list.size(); int config_idx = 0; - for (int si = 0; si < n_schedulers; si++) { + for (size_t si = 0; si < schedulers.size(); si++) { for (int n_steps : steps_list) { for (int n_threads : threads_list) { config_idx++; @@ -151,6 +184,11 @@ int main(int argc, char ** argv) { fflush(stderr); diffuse_context * ctx = diffuse_context_new(model, n_ctx, n_threads); + if (!ctx) { + fprintf(stderr, " failed to initialize context\n"); + diffuse_model_free(model); + return 1; + } diffuse_sampler_params sparams; sparams.n_steps = n_steps; @@ -159,6 +197,7 @@ int main(int argc, char ** argv) { sparams.schedule = diffuse_schedule::COSINE; sparams.remasking = schedulers[si].remasking; sparams.entropy_threshold = entropy_threshold; + sparams.use_cache = use_cache; // Warmup for (int w = 0; w < n_warmup; w++) { @@ -207,9 +246,13 @@ int main(int argc, char ** argv) { if (json_output) { printf("{\n"); printf(" \"model\": \"%s\",\n", model_path.c_str()); + printf(" \"device\": \"%s\",\n", device_name.empty() ? "CPU" : device_name.c_str()); + printf(" \"cache\": %s,\n", use_cache ? "true" : "false"); + printf(" \"load_ms\": %.1f,\n", load_ms); printf(" \"n_prompt\": %d,\n", (int)input_tokens.size()); printf(" \"n_generate\": %d,\n", n_generate); printf(" \"n_reps\": %d,\n", n_reps); + printf(" \"n_warmup\": %d,\n", n_warmup); printf(" \"results\": [\n"); for (size_t i = 0; i < results.size(); i++) { const auto & r = results[i]; @@ -226,7 +269,7 @@ int main(int argc, char ** argv) { printf("| Scheduler | Steps | Threads | Avg Time (ms) | tok/s | Actual Steps |\n"); printf("|---|---|---|---|---|---|\n"); - for (int si = 0; si < n_schedulers; si++) { + for (size_t si = 0; si < schedulers.size(); si++) { for (int n_steps : steps_list) { for (int n_threads : threads_list) { double sum_ms = 0; diff --git a/tools/main-cli.cpp b/tools/main-cli.cpp index 34f8952..b0cac95 100644 --- a/tools/main-cli.cpp +++ b/tools/main-cli.cpp @@ -10,6 +10,7 @@ static void print_usage(const char * prog) { fprintf(stderr, "Usage: %s [options]\n", prog); fprintf(stderr, "\nOptions:\n"); fprintf(stderr, " -m PATH Model file (GGUF)\n"); + fprintf(stderr, " --device NAME Offload model and graph to an exact device name\n"); fprintf(stderr, " --list-devices List available compute devices and exit\n"); fprintf(stderr, " -p TEXT Prompt text\n"); fprintf(stderr, " -n INT Tokens to generate (default: 128)\n"); @@ -40,6 +41,7 @@ static std::vector parse_tokens(const char * str) { int main(int argc, char ** argv) { std::string model_path; + std::string device_name; std::string prompt; std::vector input_tokens; int n_generate = 128; @@ -56,7 +58,9 @@ int main(int argc, char ** argv) { bool list_devices = false; for (int i = 1; i < argc; i++) { - if (strcmp(argv[i], "--list-devices") == 0) { + if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) { + device_name = argv[++i]; + } else if (strcmp(argv[i], "--list-devices") == 0) { list_devices = true; } else if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) { model_path = argv[++i]; @@ -132,7 +136,9 @@ int main(int argc, char ** argv) { // Load model fprintf(stderr, "Loading model...\n"); - diffuse_model * model = diffuse_model_load(model_path, n_threads); + diffuse_model * model = device_name.empty() + ? diffuse_model_load(model_path, n_threads) + : diffuse_model_load(model_path, n_threads, device_name); if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; @@ -141,6 +147,11 @@ int main(int argc, char ** argv) { const auto & hp = diffuse_model_hparams(model); int n_ctx = (int)input_tokens.size() + n_generate; diffuse_context * ctx = diffuse_context_new(model, n_ctx, n_threads); + if (!ctx) { + fprintf(stderr, "Failed to initialize compute context\n"); + diffuse_model_free(model); + return 1; + } // Setup sampler params diffuse_sampler_params sparams; From ec0b04d3e094bdf1a536bd4c30334db6fc445ff8 Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Tue, 18 Aug 2026 14:14:53 +0000 Subject: [PATCH 07/10] Offload cached diffusion inference Allocate cached graphs through the selected GGML scheduler, transfer cached K/V state through backend tensors, and retain explicit placement checks. Cover full-cache extraction, active-set execution, quantized synthetic models, and cached generation. Assisted-by: OpenAI Codex --- README.md | 14 +-- src/diffuse-backend.cpp | 43 ++++++- src/diffuse-common.h | 14 ++- src/diffuse-graph.cpp | 224 ++++++++++++++++++++++++------------ src/diffuse-model.cpp | 20 +--- src/diffuse-sampler.cpp | 5 +- tests/test-cache.cpp | 248 ++++++++++++++++++++++++++++++++++------ tests/test-e2e.py | 36 +++++- tools/main-bench.cpp | 5 - tools/main-cli.cpp | 3 +- 10 files changed, 463 insertions(+), 149 deletions(-) diff --git a/README.md b/README.md index 0e5d342..963f3ac 100644 --- a/README.md +++ b/README.md @@ -146,11 +146,10 @@ operations. The CLI reports model allocation and per-backend graph node counts so an offload cannot silently become CPU-only execution. Integrated GPUs are valid devices and are reported separately as `IGPU`. -GPU offload currently supports the masked-diffusion forward path. The -inter-step KV cache is disabled for offloaded generation, autoregressive -execution is not yet supported, and the full model must fit the selected -device's buffer type. The GGML scheduler keeps a CPU backend available for -unsupported graph operations and reports any such placement. +GPU offload supports full and inter-step-cached masked-diffusion graphs. +Autoregressive execution is not yet supported, and the full model must fit the +selected device's buffer type. The GGML scheduler keeps a CPU backend available +for unsupported graph operations and reports any such placement. **Note**: diffuse-cpp operates on token IDs, not raw text. Use the HuggingFace transformers library to tokenize your prompts: @@ -334,8 +333,9 @@ Current limitations: - No integrated tokenizer (use transformers) - Default 256 generated tokens per call (configurable via -n flag) - Single-model inference only (no batching) -- GPU offload is experimental and does not yet support the inter-step cache or - autoregressive execution +- GPU offload is experimental, does not yet support autoregressive execution, + and currently requires the full model to fit the selected device's buffer + type ## Contributing diff --git a/src/diffuse-backend.cpp b/src/diffuse-backend.cpp index 99d2c84..6f1b40d 100644 --- a/src/diffuse-backend.cpp +++ b/src/diffuse-backend.cpp @@ -1,6 +1,4 @@ -#include "diffuse.h" - -#include "ggml-backend.h" +#include "diffuse-common.h" static diffuse_device_type to_diffuse_device_type(enum ggml_backend_dev_type type) { switch (type) { @@ -43,3 +41,42 @@ std::vector diffuse_available_devices() { return devices; } + +bool diffuse_backend_sched_prepare( + diffuse_context * ctx, + struct ggml_cgraph * graph, + const char * graph_name, + bool & placement_logged) { + ggml_backend_sched_reset(ctx->sched); + if (!ggml_backend_sched_alloc_graph(ctx->sched, graph)) { + DIFFUSE_LOG("failed to allocate scheduled %s graph", graph_name); + return false; + } + + int device_nodes = 0; + int cpu_nodes = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + ggml_backend_t backend = ggml_backend_sched_get_tensor_backend( + ctx->sched, ggml_graph_node(graph, i)); + if (backend == ctx->model->backend) ++device_nodes; + if (backend == ctx->model->cpu_backend) ++cpu_nodes; + } + if (!placement_logged) { + DIFFUSE_LOG(" %s graph placement: %s=%d nodes, CPU=%d nodes", + graph_name, ctx->model->device_name.c_str(), device_nodes, cpu_nodes); + placement_logged = true; + } + if (device_nodes == 0) { + DIFFUSE_LOG("scheduled %s graph did not place any operations on %s", + graph_name, ctx->model->device_name.c_str()); + return false; + } + + return true; +} + +bool diffuse_backend_sched_compute( + diffuse_context * ctx, + struct ggml_cgraph * graph) { + return ggml_backend_sched_graph_compute(ctx->sched, graph) == GGML_STATUS_SUCCESS; +} diff --git a/src/diffuse-common.h b/src/diffuse-common.h index 9fd10df..06e3c43 100644 --- a/src/diffuse-common.h +++ b/src/diffuse-common.h @@ -80,5 +80,17 @@ struct diffuse_context { ggml_backend_buffer_t buf = nullptr; ggml_backend_sched_t sched = nullptr; struct ggml_context * ctx = nullptr; // compute context - bool placement_logged = false; + bool full_placement_logged = false; + bool cached_placement_logged = false; + bool ar_placement_logged = false; }; + +bool diffuse_backend_sched_prepare( + diffuse_context * ctx, + struct ggml_cgraph * graph, + const char * graph_name, + bool & placement_logged); + +bool diffuse_backend_sched_compute( + diffuse_context * ctx, + struct ggml_cgraph * graph); diff --git a/src/diffuse-graph.cpp b/src/diffuse-graph.cpp index 5d22ff6..899cd6a 100644 --- a/src/diffuse-graph.cpp +++ b/src/diffuse-graph.cpp @@ -186,12 +186,14 @@ static struct ggml_cgraph * diffuse_build_graph_extractable( struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N); ggml_set_name(inp_tokens, "inp_tokens"); ggml_set_input(inp_tokens); - memcpy(inp_tokens->data, tokens, N * sizeof(int32_t)); + if (inp_tokens->data != nullptr) { + memcpy(inp_tokens->data, tokens, N * sizeof(int32_t)); + } struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N); ggml_set_name(inp_pos, "inp_pos"); ggml_set_input(inp_pos); - { + if (inp_pos->data != nullptr) { int32_t * pos_data = (int32_t *)inp_pos->data; for (int i = 0; i < N; i++) pos_data[i] = i; } @@ -243,17 +245,22 @@ static struct ggml_cgraph * diffuse_build_graph_extractable( V = ggml_reshape_3d(ctx, V, n_embd_head, n_head, N); } - // ── Name K,V for cache extraction (BEFORE permute) ────── - // Shape at this point: [n_embd_head, n_head, N] + // Materialize independent K,V outputs before the attention views. + // The scheduler may otherwise reuse intermediate storage once the + // downstream attention graph no longer needs the unpermuted tensor. { char name_buf[32]; + struct ggml_tensor * K_cache = ggml_dup(ctx, K); snprintf(name_buf, sizeof(name_buf), "Kc.%02d", il); - ggml_set_name(K, name_buf); - ggml_set_output(K); + ggml_set_name(K_cache, name_buf); + ggml_set_output(K_cache); + ggml_build_forward_expand(gf, K_cache); + struct ggml_tensor * V_cache = ggml_dup(ctx, V); snprintf(name_buf, sizeof(name_buf), "Vc.%02d", il); - ggml_set_name(V, name_buf); - ggml_set_output(V); + ggml_set_name(V_cache, name_buf); + ggml_set_output(V_cache); + ggml_build_forward_expand(gf, V_cache); } // Permute for attention (same as original) @@ -313,61 +320,12 @@ static size_t diffuse_max_graph_nodes(const diffuse_model * model) { + (size_t)model->hparams.n_layer * 8 + 256; } -static bool diffuse_sched_compute(diffuse_context * ctx, - struct ggml_cgraph * graph, - const int32_t * tokens, - int n_tokens) { - ggml_backend_sched_reset(ctx->sched); - if (!ggml_backend_sched_alloc_graph(ctx->sched, graph)) { - DIFFUSE_LOG("failed to allocate scheduled compute graph"); - return false; - } - - struct ggml_tensor * inp_tokens = ggml_graph_get_tensor(graph, "inp_tokens"); - struct ggml_tensor * inp_pos = ggml_graph_get_tensor(graph, "inp_pos"); - if (inp_tokens == nullptr || inp_pos == nullptr) { - DIFFUSE_LOG("scheduled graph inputs not found"); - return false; - } - - std::vector positions(n_tokens); - for (int i = 0; i < n_tokens; ++i) positions[i] = i; - ggml_backend_tensor_set(inp_tokens, tokens, 0, n_tokens * sizeof(int32_t)); - ggml_backend_tensor_set(inp_pos, positions.data(), 0, n_tokens * sizeof(int32_t)); - - if (!ctx->placement_logged) { - int device_nodes = 0; - int cpu_nodes = 0; - for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { - ggml_backend_t backend = ggml_backend_sched_get_tensor_backend( - ctx->sched, ggml_graph_node(graph, i)); - if (backend == ctx->model->backend) ++device_nodes; - if (backend == ctx->model->cpu_backend) ++cpu_nodes; - } - DIFFUSE_LOG(" graph placement: %s=%d nodes, CPU=%d nodes", - ctx->model->device_name.c_str(), device_nodes, cpu_nodes); - if (device_nodes == 0) { - DIFFUSE_LOG("scheduled graph did not place any operations on %s", - ctx->model->device_name.c_str()); - return false; - } - ctx->placement_logged = true; - } - - return ggml_backend_sched_graph_compute(ctx->sched, graph) == GGML_STATUS_SUCCESS; -} - bool diffuse_forward_full(diffuse_context * ctx, const int32_t * tokens, int n_tokens, float * logits_out, diffuse_step_cache * cache) { const auto & hp = ctx->model->hparams; - if (ctx->sched != nullptr && cache != nullptr) { - DIFFUSE_LOG("inter-step cache is not yet supported with device offload"); - return false; - } - // Buffer size: same as original + extra for named K,V output tensors size_t per_layer = (size_t)n_tokens * hp.n_embd * sizeof(float) * 10 + (size_t)n_tokens * hp.n_ff * sizeof(float) * 3 @@ -401,8 +359,26 @@ bool diffuse_forward_full(diffuse_context * ctx, enum ggml_status status = GGML_STATUS_SUCCESS; if (scheduled) { - if (!diffuse_sched_compute(ctx, gf, tokens, n_tokens)) { + if (!diffuse_backend_sched_prepare(ctx, gf, "full", + ctx->full_placement_logged)) { status = GGML_STATUS_FAILED; + } else { + struct ggml_tensor * inp_tokens = ggml_graph_get_tensor(gf, "inp_tokens"); + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(gf, "inp_pos"); + if (inp_tokens == nullptr || inp_pos == nullptr) { + DIFFUSE_LOG("scheduled full graph inputs not found"); + status = GGML_STATUS_FAILED; + } else { + std::vector positions(n_tokens); + for (int i = 0; i < n_tokens; ++i) positions[i] = i; + ggml_backend_tensor_set(inp_tokens, tokens, 0, + n_tokens * sizeof(int32_t)); + ggml_backend_tensor_set(inp_pos, positions.data(), 0, + n_tokens * sizeof(int32_t)); + if (!diffuse_backend_sched_compute(ctx, gf)) { + status = GGML_STATUS_FAILED; + } + } } } else { status = ggml_graph_compute_with_ctx(ctx_compute, gf, ctx->n_threads); @@ -440,8 +416,13 @@ bool diffuse_forward_full(diffuse_context * ctx, struct ggml_tensor * V_t = ggml_graph_get_tensor(gf, name_buf); if (K_t && V_t) { - memcpy(cache->K[il].data(), K_t->data, kv_bytes); - memcpy(cache->V[il].data(), V_t->data, kv_bytes); + if (scheduled) { + ggml_backend_tensor_get(K_t, cache->K[il].data(), 0, kv_bytes); + ggml_backend_tensor_get(V_t, cache->V[il].data(), 0, kv_bytes); + } else { + memcpy(cache->K[il].data(), K_t->data, kv_bytes); + memcpy(cache->V[il].data(), V_t->data, kv_bytes); + } } else { DIFFUSE_LOG("WARNING: could not extract K/V for layer %d", il); } @@ -493,12 +474,16 @@ struct ggml_cgraph * diffuse_build_graph_cached( struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_active); ggml_set_name(inp_tokens, "inp_tokens"); ggml_set_input(inp_tokens); - memcpy(inp_tokens->data, active_tokens, n_active * sizeof(int32_t)); + if (inp_tokens->data != nullptr) { + memcpy(inp_tokens->data, active_tokens, n_active * sizeof(int32_t)); + } struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_active); ggml_set_name(inp_pos, "inp_pos_active"); ggml_set_input(inp_pos); - memcpy(inp_pos->data, active_pos_indices, n_active * sizeof(int32_t)); + if (inp_pos->data != nullptr) { + memcpy(inp_pos->data, active_pos_indices, n_active * sizeof(int32_t)); + } // We also need position indices for the cached positions (for V permutation ordering) // — not needed for RoPE since cached K,V already have RoPE applied @@ -556,30 +541,39 @@ struct ggml_cgraph * diffuse_build_graph_cached( V_active = ggml_reshape_3d(ctx, V_active, n_embd_head, n_head, n_active); } - // ── Name K_active, V_active for extraction ────────────── + // Materialize independent active-set K,V cache outputs. { char name_buf[32]; + struct ggml_tensor * K_cache = ggml_dup(ctx, K_active); snprintf(name_buf, sizeof(name_buf), "Ka.%02d", il); - ggml_set_name(K_active, name_buf); - ggml_set_output(K_active); + ggml_set_name(K_cache, name_buf); + ggml_set_output(K_cache); + ggml_build_forward_expand(gf, K_cache); + struct ggml_tensor * V_cache = ggml_dup(ctx, V_active); snprintf(name_buf, sizeof(name_buf), "Va.%02d", il); - ggml_set_name(V_active, name_buf); - ggml_set_output(V_active); + ggml_set_name(V_cache, name_buf); + ggml_set_output(V_cache); + ggml_build_forward_expand(gf, V_cache); } // ── Load cached K,V for inactive positions ────────────── // Shape: [n_embd_head, n_head, n_cached] struct ggml_tensor * K_cached = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_head, n_head, n_cached); + char name_buf[32]; + snprintf(name_buf, sizeof(name_buf), "Kcache.%02d", il); + ggml_set_name(K_cached, name_buf); ggml_set_input(K_cached); struct ggml_tensor * V_cached = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_head, n_head, n_cached); + snprintf(name_buf, sizeof(name_buf), "Vcache.%02d", il); + ggml_set_name(V_cached, name_buf); ggml_set_input(V_cached); // Fill from cache: gather cached positions into contiguous tensor - { + if (K_cached->data != nullptr && V_cached->data != nullptr) { float * K_dst = (float *)K_cached->data; float * V_dst = (float *)V_cached->data; for (int c = 0; c < n_cached; c++) { @@ -692,10 +686,16 @@ bool diffuse_forward_cached( buf_size += 256ull * 1024 * 1024; buf_size = (size_t)(buf_size * 1.5); + const bool scheduled = ctx->sched != nullptr; + const size_t graph_nodes = diffuse_max_graph_nodes(ctx->model); + const size_t context_size = scheduled + ? ggml_tensor_overhead() * graph_nodes + ggml_graph_overhead_custom(graph_nodes, false) + : buf_size; + struct ggml_init_params cparams = { - /*.mem_size = */ buf_size, + /*.mem_size = */ context_size, /*.mem_buffer = */ nullptr, - /*.no_alloc = */ false, + /*.no_alloc = */ scheduled, }; struct ggml_context * ctx_compute = ggml_init(cparams); if (!ctx_compute) { @@ -709,7 +709,62 @@ bool diffuse_forward_cached( n_active, n_total, cache, cached_positions); - enum ggml_status status = ggml_graph_compute_with_ctx(ctx_compute, gf, ctx->n_threads); + enum ggml_status status = GGML_STATUS_SUCCESS; + if (scheduled) { + if (!diffuse_backend_sched_prepare(ctx, gf, "cached", + ctx->cached_placement_logged)) { + status = GGML_STATUS_FAILED; + } else { + struct ggml_tensor * inp_tokens = ggml_graph_get_tensor(gf, "inp_tokens"); + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(gf, "inp_pos_active"); + if (inp_tokens == nullptr || inp_pos == nullptr) { + DIFFUSE_LOG("scheduled cached graph inputs not found"); + status = GGML_STATUS_FAILED; + } else { + ggml_backend_tensor_set(inp_tokens, active_tokens, 0, + n_active * sizeof(int32_t)); + ggml_backend_tensor_set(inp_pos, active_pos_indices, 0, + n_active * sizeof(int32_t)); + + const size_t kv_stride = cache->pos_stride(); + std::vector K_cached((size_t)n_cached * kv_stride); + std::vector V_cached((size_t)n_cached * kv_stride); + char name_buf[32]; + for (int il = 0; il < (int)hp.n_layer; ++il) { + for (int c = 0; c < n_cached; ++c) { + const int orig_pos = cached_positions[c]; + memcpy(K_cached.data() + c * kv_stride, + cache->K[il].data() + orig_pos * kv_stride, + kv_stride * sizeof(float)); + memcpy(V_cached.data() + c * kv_stride, + cache->V[il].data() + orig_pos * kv_stride, + kv_stride * sizeof(float)); + } + + snprintf(name_buf, sizeof(name_buf), "Kcache.%02d", il); + struct ggml_tensor * K_t = ggml_graph_get_tensor(gf, name_buf); + snprintf(name_buf, sizeof(name_buf), "Vcache.%02d", il); + struct ggml_tensor * V_t = ggml_graph_get_tensor(gf, name_buf); + if (K_t == nullptr || V_t == nullptr) { + DIFFUSE_LOG("scheduled cached K/V inputs not found for layer %d", il); + status = GGML_STATUS_FAILED; + break; + } + ggml_backend_tensor_set(K_t, K_cached.data(), 0, + K_cached.size() * sizeof(float)); + ggml_backend_tensor_set(V_t, V_cached.data(), 0, + V_cached.size() * sizeof(float)); + } + + if (status == GGML_STATUS_SUCCESS && + !diffuse_backend_sched_compute(ctx, gf)) { + status = GGML_STATUS_FAILED; + } + } + } + } else { + status = ggml_graph_compute_with_ctx(ctx_compute, gf, ctx->n_threads); + } if (status != GGML_STATUS_SUCCESS) { DIFFUSE_LOG("cached graph compute failed with status %d", (int)status); ggml_free(ctx_compute); @@ -723,11 +778,23 @@ bool diffuse_forward_cached( ggml_free(ctx_compute); return false; } - memcpy(logits_out, logits->data, (size_t)n_active * hp.n_vocab * sizeof(float)); + const size_t logits_bytes = (size_t)n_active * hp.n_vocab * sizeof(float); + if (scheduled) { + ggml_backend_tensor_get(logits, logits_out, 0, logits_bytes); + } else { + memcpy(logits_out, logits->data, logits_bytes); + } // Update cache: store K_active, V_active for active positions { char name_buf[32]; + const size_t kv_count = (size_t)n_active * cache->pos_stride(); + std::vector K_active; + std::vector V_active; + if (scheduled) { + K_active.resize(kv_count); + V_active.resize(kv_count); + } for (int il = 0; il < (int)hp.n_layer; il++) { snprintf(name_buf, sizeof(name_buf), "Ka.%02d", il); struct ggml_tensor * K_t = ggml_graph_get_tensor(gf, name_buf); @@ -735,9 +802,18 @@ bool diffuse_forward_cached( struct ggml_tensor * V_t = ggml_graph_get_tensor(gf, name_buf); if (K_t && V_t) { - cache->update_kv(il, (const float *)K_t->data, - (const float *)V_t->data, + if (scheduled) { + ggml_backend_tensor_get(K_t, K_active.data(), 0, + kv_count * sizeof(float)); + ggml_backend_tensor_get(V_t, V_active.data(), 0, + kv_count * sizeof(float)); + cache->update_kv(il, K_active.data(), V_active.data(), active_positions); + } else { + cache->update_kv(il, (const float *)K_t->data, + (const float *)V_t->data, + active_positions); + } } } } diff --git a/src/diffuse-model.cpp b/src/diffuse-model.cpp index f54b8a5..b6b14ea 100644 --- a/src/diffuse-model.cpp +++ b/src/diffuse-model.cpp @@ -3,25 +3,6 @@ #include #include -// ── GGUF metadata key helpers ────────────────────────────────── -static int64_t find_key(const struct gguf_context * gctx, const char * key) { - int64_t id = gguf_find_key(gctx, key); - if (id < 0) { - DIFFUSE_DIE("missing GGUF metadata key: %s", key); - } - return id; -} - -static uint32_t get_u32(const struct gguf_context * gctx, const char * key) { - return gguf_get_val_u32(gctx, find_key(gctx, key)); -} - -static float get_f32(const struct gguf_context * gctx, const char * key, float def) { - int64_t id = gguf_find_key(gctx, key); - if (id < 0) return def; - return gguf_get_val_f32(gctx, id); -} - // ── Multi-arch key resolution ──────────────────────────────────── // Try diffuse.*, then qwen2.*, then llama.* prefixes. // This allows loading GGUF files produced by llama.cpp's converter. @@ -249,6 +230,7 @@ diffuse_model * diffuse_model_load_impl(const std::string & path, int n_threads, DIFFUSE_DIE("failed to allocate model weights on %s", ggml_backend_dev_name(selected_device)); } + ggml_backend_buffer_set_usage(model->buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); if (!load_backend_weights(path, meta_ctx, gctx)) { DIFFUSE_DIE("failed to load model weights onto %s", ggml_backend_dev_name(selected_device)); diff --git a/src/diffuse-sampler.cpp b/src/diffuse-sampler.cpp index f9889dc..4806e1a 100644 --- a/src/diffuse-sampler.cpp +++ b/src/diffuse-sampler.cpp @@ -99,10 +99,7 @@ std::vector diffuse_sample( std::mt19937 rng(params.seed); // ── Inter-step KV cache ────────────────────────────────────── - const bool use_cache = params.use_cache && ctx->sched == nullptr; - if (params.use_cache && !use_cache) { - DIFFUSE_LOG("inter-step cache disabled for device-offloaded generation"); - } + const bool use_cache = params.use_cache; diffuse_step_cache cache; if (use_cache) { cache.init(total_len, prompt_len, diff --git a/tests/test-cache.cpp b/tests/test-cache.cpp index 7bb4482..9a7a22e 100644 --- a/tests/test-cache.cpp +++ b/tests/test-cache.cpp @@ -1,6 +1,11 @@ #include "diffuse-cache.h" +#include "diffuse-graph.h" +#include +#include #include +#include +#include #include #define ASSERT(cond, msg) do { \ @@ -10,37 +15,216 @@ } \ } while (0) -int main() { - constexpr int n_tokens = 6; - const int32_t tokens[n_tokens] = {1, 1, 1, 1, 1, 1}; - - diffuse_step_cache cache; - cache.init(n_tokens, 2, 1, 1, 1); - cache.update_seq(tokens, n_tokens); - - const std::vector is_masked = { - false, false, false, true, false, true, - }; - std::vector cached_positions; - std::vector active_positions; - std::vector active_to_orig; - - cache.compute_active_set( - tokens, is_masked, n_tokens, 1, 0, false, - cached_positions, active_positions, active_to_orig); - ASSERT(active_positions == std::vector({3, 5}), - "unshifted models should keep masked positions active"); - - cache.compute_active_set( - tokens, is_masked, n_tokens, 1, 0, true, - cached_positions, active_positions, active_to_orig); - ASSERT(active_positions == std::vector({2, 3, 4, 5}), - "Dream should also keep each masked position's logit source active"); - ASSERT(cached_positions == std::vector({0, 1}), - "Dream should still cache unrelated stable positions"); - ASSERT(active_to_orig == active_positions, - "active indices should map back to their original positions"); - - fprintf(stderr, "PASS: shifted-logit cache active set\n"); +static double nmse(const float * reference, const float * actual, size_t count) { + double squared_error = 0.0; + double squared_reference = 0.0; + for (size_t i = 0; i < count; ++i) { + const double diff = (double)actual[i] - reference[i]; + squared_error += diff * diff; + squared_reference += (double)reference[i] * reference[i]; + } + return squared_error / std::max(squared_reference, 1e-30); +} + +static double cache_nmse(const diffuse_step_cache & reference, + const diffuse_step_cache & actual) { + double squared_error = 0.0; + double squared_reference = 0.0; + for (int il = 0; il < reference.n_layer; ++il) { + for (size_t i = 0; i < reference.K[il].size(); ++i) { + const double k_diff = (double)actual.K[il][i] - reference.K[il][i]; + const double v_diff = (double)actual.V[il][i] - reference.V[il][i]; + squared_error += k_diff * k_diff + v_diff * v_diff; + squared_reference += (double)reference.K[il][i] * reference.K[il][i] + + (double)reference.V[il][i] * reference.V[il][i]; + } + } + return squared_error / std::max(squared_reference, 1e-30); +} + +static void print_cache_diagnostics(const diffuse_step_cache & reference, + const diffuse_step_cache & actual) { + const size_t count = reference.K[0].size(); + fprintf(stderr, "layer 0: K NMSE=%.6g, V NMSE=%.6g\n", + nmse(reference.K[0].data(), actual.K[0].data(), count), + nmse(reference.V[0].data(), actual.V[0].data(), count)); + for (size_t i = 0; i < std::min(4, count); ++i) { + fprintf(stderr, " K[%zu] CPU=%g device=%g; V[%zu] CPU=%g device=%g\n", + i, reference.K[0][i], actual.K[0][i], + i, reference.V[0][i], actual.V[0][i]); + } +} + +static int top_token(const float * logits, int n_vocab) { + return (int)std::distance(logits, std::max_element(logits, logits + n_vocab)); +} + +int main(int argc, char ** argv) { + { + diffuse_step_cache policy_cache; + constexpr int n_policy_tokens = 6; + const int32_t policy_tokens[n_policy_tokens] = {1, 1, 1, 1, 1, 1}; + policy_cache.init(n_policy_tokens, 2, 1, 1, 1); + policy_cache.update_seq(policy_tokens, n_policy_tokens); + + std::vector is_masked = {false, false, false, true, false, true}; + std::vector cached_positions; + std::vector active_positions; + std::vector active_to_orig; + + policy_cache.compute_active_set( + policy_tokens, is_masked, n_policy_tokens, 1, 0, false, + cached_positions, active_positions, active_to_orig); + ASSERT(active_positions == std::vector({3, 5}), + "unshifted cache should keep masked positions active"); + + policy_cache.compute_active_set( + policy_tokens, is_masked, n_policy_tokens, 1, 0, true, + cached_positions, active_positions, active_to_orig); + ASSERT(active_positions == std::vector({2, 3, 4, 5}), + "shifted cache should also keep masked-position logit sources active"); + ASSERT(cached_positions == std::vector({0, 1}), + "shifted cache should retain unrelated stable positions"); + } + + if (argc < 3) { + fprintf(stderr, "PASS: cache active-set policy\n"); + fprintf(stderr, "SKIP: backend comparison requires MODEL DEVICE\n"); + return 0; + } + + const char * model_path = argv[1]; + const char * device_name = argv[2]; + const double max_logits_nmse = argc > 3 ? std::atof(argv[3]) : 2e-3; + const int n_total = argc > 4 ? std::atoi(argv[4]) : 8; + const int n_active = argc > 5 ? std::atoi(argv[5]) : n_total / 2; + const int n_prompt = n_total - n_active; + constexpr int n_threads = 4; + ASSERT(n_total >= 8, "cache comparison needs at least eight total positions"); + ASSERT(n_active > 0 && n_active < n_total, + "active positions must be between zero and the total position count"); + + using model_ptr = std::unique_ptr; + using context_ptr = std::unique_ptr; + model_ptr cpu_model(diffuse_model_load(model_path, n_threads), diffuse_model_free); + model_ptr device_model( + diffuse_model_load(model_path, n_threads, device_name), diffuse_model_free); + ASSERT(cpu_model != nullptr, "CPU model should load"); + ASSERT(device_model != nullptr, "device model should load"); + + const auto & hp = diffuse_model_hparams(cpu_model.get()); + context_ptr cpu_ctx( + diffuse_context_new(cpu_model.get(), n_total, n_threads), diffuse_context_free); + context_ptr device_ctx( + diffuse_context_new(device_model.get(), n_total, n_threads), diffuse_context_free); + ASSERT(cpu_ctx != nullptr, "CPU context should initialize"); + ASSERT(device_ctx != nullptr, "device context should initialize"); + + diffuse_step_cache cpu_cache; + diffuse_step_cache device_cache; + cpu_cache.init(n_total, n_prompt, (int)hp.n_layer, + (int)hp.n_embd_head(), (int)hp.n_head); + device_cache.init(n_total, n_prompt, (int)hp.n_layer, + (int)hp.n_embd_head(), (int)hp.n_head); + + std::vector tokens(n_total, 1); + std::vector cpu_logits((size_t)n_total * hp.n_vocab); + std::vector device_logits(cpu_logits.size()); + + ASSERT(diffuse_forward_full(cpu_ctx.get(), tokens.data(), n_total, + cpu_logits.data(), &cpu_cache), + "CPU full cached forward should succeed"); + ASSERT(diffuse_forward_full(device_ctx.get(), tokens.data(), n_total, + device_logits.data(), &device_cache), + "device full cached forward should succeed"); + + const double full_logits_nmse = nmse(cpu_logits.data(), device_logits.data(), + cpu_logits.size()); + const double full_cache_nmse = cache_nmse(cpu_cache, device_cache); + fprintf(stderr, "%s full cached forward: logits NMSE=%.6g, K/V NMSE=%.6g\n", + device_name, full_logits_nmse, full_cache_nmse); + if (full_cache_nmse > 2e-3) { + print_cache_diagnostics(cpu_cache, device_cache); + } + ASSERT(full_logits_nmse <= max_logits_nmse, + "full cached logits should match CPU"); + ASSERT(full_cache_nmse <= 2e-3, "full cached K/V should match CPU"); + + // Isolate active-set graph accuracy from error accumulated while creating + // the initial device cache. Real generation retains the device values; + // the full-cache assertion above bounds that separate transfer. + device_cache.K = cpu_cache.K; + device_cache.V = cpu_cache.V; + cpu_cache.update_seq(tokens.data(), n_total); + device_cache.update_seq(tokens.data(), n_total); + + std::vector cached_positions(n_prompt); + std::vector active_positions(n_active); + std::vector active_tokens(n_active); + std::vector active_pos_indices(n_active); + for (int i = 0; i < n_prompt; ++i) { + cached_positions[i] = i; + } + for (int i = 0; i < n_active; ++i) { + active_positions[i] = n_prompt + i; + active_tokens[i] = 2 + i % 4; + active_pos_indices[i] = n_prompt + i; + } + cpu_logits.resize((size_t)active_tokens.size() * hp.n_vocab); + device_logits.resize(cpu_logits.size()); + + ASSERT(diffuse_forward_cached(cpu_ctx.get(), + active_tokens.data(), active_pos_indices.data(), + (int)active_tokens.size(), n_total, &cpu_cache, + cached_positions, active_positions, + cpu_logits.data()), + "CPU active-set cached forward should succeed"); + ASSERT(diffuse_forward_cached(device_ctx.get(), + active_tokens.data(), active_pos_indices.data(), + (int)active_tokens.size(), n_total, &device_cache, + cached_positions, active_positions, + device_logits.data()), + "device active-set cached forward should succeed"); + + const double active_logits_nmse = nmse(cpu_logits.data(), device_logits.data(), + cpu_logits.size()); + const double active_cache_nmse = cache_nmse(cpu_cache, device_cache); + fprintf(stderr, "%s active-set cached forward: logits NMSE=%.6g, K/V NMSE=%.6g\n", + device_name, active_logits_nmse, active_cache_nmse); + int top_token_mismatches = 0; + for (size_t i = 0; i < active_tokens.size(); ++i) { + const int cpu_top = top_token( + cpu_logits.data() + i * hp.n_vocab, (int)hp.n_vocab); + const int device_top = top_token( + device_logits.data() + i * hp.n_vocab, (int)hp.n_vocab); + top_token_mismatches += cpu_top != device_top; + if (i >= 4) { + continue; + } + fprintf(stderr, " active[%zu] top token: CPU=%d device=%d\n", i, + cpu_top, device_top); + } + fprintf(stderr, " top-token mismatches: %d/%d\n", + top_token_mismatches, n_active); + ASSERT(active_logits_nmse <= max_logits_nmse, + "active-set cached logits should match CPU"); + ASSERT(active_cache_nmse <= 2e-3, "active-set cached K/V should match CPU"); + + diffuse_sampler_params params; + params.n_steps = 3; + params.use_cache = true; + params.cache_keep_active = 1; + int callbacks = 0; + const std::vector prompt = {1, 2}; + const auto generated = diffuse_generate( + device_ctx.get(), prompt, 6, params, + [&callbacks](int, int, const std::vector &) { ++callbacks; }); + ASSERT(generated.size() == 6, "cached generation should return every requested token"); + ASSERT(callbacks == params.n_steps, "cached generation should complete every fixed step"); + ASSERT(std::find(generated.begin(), generated.end(), (int32_t)hp.mask_token_id) + == generated.end(), + "cached generation should leave no mask tokens"); + + fprintf(stderr, "PASS: %s inter-step cache matches CPU\n", device_name); return 0; } diff --git a/tests/test-e2e.py b/tests/test-e2e.py index c73b9d3..0f672cb 100644 --- a/tests/test-e2e.py +++ b/tests/test-e2e.py @@ -29,6 +29,10 @@ def parse_args(): parser.add_argument("--build-dir", default="build") parser.add_argument("--device") parser.add_argument("--keep-gguf") + parser.add_argument( + "--quantize", choices=("q4_k_m", "q6_k"), + help="quantize the generated F32 model before testing", + ) return parser.parse_args() @@ -85,13 +89,13 @@ def main(): save_file(tensors, os.path.join(model_dir, "model.safetensors")) # Convert to GGUF - gguf_path = os.path.join(tmpdir, "tiny.gguf") + f32_gguf_path = os.path.join(tmpdir, "tiny-f32.gguf") converter = os.path.join(project_dir, "tools", "convert-llada.py") result = subprocess.run([ sys.executable, converter, "--input", model_dir, - "--output", gguf_path, + "--output", f32_gguf_path, ], capture_output=True, text=True) if result.returncode != 0: @@ -101,6 +105,22 @@ def main(): print("Converter OK") + gguf_path = f32_gguf_path + if args.quantize: + quantizer = os.path.join(build_dir, "diffuse-quantize") + if not os.path.exists(quantizer): + print(f"FAIL: diffuse-quantize not found at {quantizer}") + sys.exit(1) + gguf_path = os.path.join(tmpdir, f"tiny-{args.quantize}.gguf") + result = subprocess.run( + [quantizer, f32_gguf_path, gguf_path, args.quantize], + capture_output=True, text=True, timeout=30, + ) + print(result.stderr) + if result.returncode != 0: + print(f"FAIL: quantizer exited with code {result.returncode}") + sys.exit(1) + # Run C++ test-forward with the GGUF file env = os.environ.copy() env["LD_LIBRARY_PATH"] = f"{build_dir}:{build_dir}/ggml/src" @@ -140,6 +160,18 @@ def main(): sys.exit(1) print("dump-logits wrapper OK") + if args.device: + cache_binary = os.path.join(build_dir, "test-cache") + result = subprocess.run( + [cache_binary, gguf_path, args.device], + capture_output=True, text=True, env=env, timeout=30, + ) + print("C++ test-cache output:") + print(result.stderr) + if result.returncode != 0: + print(f"FAIL: test-cache exited with code {result.returncode}") + sys.exit(1) + if args.keep_gguf: output_path = os.path.abspath(args.keep_gguf) os.makedirs(os.path.dirname(output_path), exist_ok=True) diff --git a/tools/main-bench.cpp b/tools/main-bench.cpp index 114a4b6..14c2d97 100644 --- a/tools/main-bench.cpp +++ b/tools/main-bench.cpp @@ -110,11 +110,6 @@ int main(int argc, char ** argv) { fprintf(stderr, "Error: -m MODEL required\n"); return 1; } - if (!device_name.empty() && use_cache) { - fprintf(stderr, "Device offload does not yet support the inter-step cache; disabling it\n"); - use_cache = false; - } - // If no tokens given, use dummy prompt if (input_tokens.empty()) { input_tokens.resize(prompt_len, 1); diff --git a/tools/main-cli.cpp b/tools/main-cli.cpp index b0cac95..4643864 100644 --- a/tools/main-cli.cpp +++ b/tools/main-cli.cpp @@ -144,7 +144,6 @@ int main(int argc, char ** argv) { return 1; } - const auto & hp = diffuse_model_hparams(model); int n_ctx = (int)input_tokens.size() + n_generate; diffuse_context * ctx = diffuse_context_new(model, n_ctx, n_threads); if (!ctx) { @@ -169,7 +168,7 @@ int main(int argc, char ** argv) { fprintf(stderr, "Generating %d tokens with %d diffusion steps...\n", n_generate, n_steps); auto result = diffuse_generate(ctx, input_tokens, n_generate, sparams, - [](int step, int total, const std::vector & tokens) { + [](int step, int total, const std::vector &) { fprintf(stderr, "\r step %d/%d", step, total); }); From 7c707defbcabd8d70510f5f7d7fc3c0c121eb9cc Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Tue, 18 Aug 2026 14:15:30 +0000 Subject: [PATCH 08/10] Offload autoregressive and speculative inference Run autoregressive prefill, decode, batch, profiling, and speculative graphs through the selected GGML backend. Add CPU/device parity coverage and expose a bounded tolerance for quantized backend comparisons. Assisted-by: OpenAI Codex --- CMakeLists.txt | 4 + README.md | 15 +-- src/ar-graph.cpp | 184 +++++++++++++++++++++++++------ src/ar-sampler.cpp | 5 +- src/ar-speculative.cpp | 6 +- tests/test-ar.cpp | 240 +++++++++++++++++++++++++++++++++++++++++ tests/test-e2e.py | 18 ++++ tools/ar-cli.cpp | 58 ++++++++-- 8 files changed, 479 insertions(+), 51 deletions(-) create mode 100644 tests/test-ar.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 1855021..65c6ef1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -70,6 +70,10 @@ if(DIFFUSE_BUILD_TESTS) target_link_libraries(test-cache PRIVATE diffuse) add_test(NAME test-cache COMMAND test-cache) + add_executable(test-ar tests/test-ar.cpp) + target_link_libraries(test-ar PRIVATE diffuse) + add_test(NAME test-ar COMMAND test-ar) + add_executable(test-backend tests/test-backend.cpp) target_link_libraries(test-backend PRIVATE diffuse) add_test(NAME test-backend COMMAND test-backend) diff --git a/README.md b/README.md index 963f3ac..b2de7d0 100644 --- a/README.md +++ b/README.md @@ -146,10 +146,12 @@ operations. The CLI reports model allocation and per-backend graph node counts so an offload cannot silently become CPU-only execution. Integrated GPUs are valid devices and are reported separately as `IGPU`. -GPU offload supports full and inter-step-cached masked-diffusion graphs. -Autoregressive execution is not yet supported, and the full model must fit the -selected device's buffer type. The GGML scheduler keeps a CPU backend available -for unsupported graph operations and reports any such placement. +GPU offload supports full and inter-step-cached masked-diffusion graphs, plus +autoregressive prefill, decoding, layer profiling, and speculative decoding. +Use `--device` with `diffuse-ar` as well; speculative decoding also accepts a +separate `--draft-device`. The full model must fit the selected device's buffer +type. The GGML scheduler keeps a CPU backend available for unsupported graph +operations and reports any such placement. **Note**: diffuse-cpp operates on token IDs, not raw text. Use the HuggingFace transformers library to tokenize your prompts: @@ -333,9 +335,8 @@ Current limitations: - No integrated tokenizer (use transformers) - Default 256 generated tokens per call (configurable via -n flag) - Single-model inference only (no batching) -- GPU offload is experimental, does not yet support autoregressive execution, - and currently requires the full model to fit the selected device's buffer - type +- GPU offload is experimental and currently requires the full model to fit the + selected device's buffer type ## Contributing diff --git a/src/ar-graph.cpp b/src/ar-graph.cpp index dabb667..e035a28 100644 --- a/src/ar-graph.cpp +++ b/src/ar-graph.cpp @@ -11,6 +11,72 @@ static struct ggml_tensor * ensure_f32(struct ggml_context * ctx, struct ggml_te return t; } +static size_t ar_max_graph_nodes(const diffuse_model * model) { + return (size_t)model->hparams.n_layer * 4096 + + (size_t)model->hparams.n_layer * 8 + 256; +} + +static bool ar_set_scheduled_inputs( + struct ggml_cgraph * graph, + const int32_t * tokens, + int n_new, + int n_past, + const ar_kv_cache * cache) { + struct ggml_tensor * inp_tokens = ggml_graph_get_tensor(graph, "inp_tokens"); + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(graph, "inp_pos"); + struct ggml_tensor * attn_mask = ggml_graph_get_tensor(graph, "attn_mask"); + if (inp_tokens == nullptr || inp_pos == nullptr || attn_mask == nullptr) { + DIFFUSE_LOG("scheduled autoregressive graph inputs not found"); + return false; + } + + ggml_backend_tensor_set(inp_tokens, tokens, 0, n_new * sizeof(int32_t)); + + std::vector positions(n_new); + for (int i = 0; i < n_new; ++i) positions[i] = n_past + i; + ggml_backend_tensor_set(inp_pos, positions.data(), 0, + positions.size() * sizeof(int32_t)); + + const int n_kv = n_past + n_new; + const ggml_fp16_t zero_f16 = ggml_fp32_to_fp16(0.0f); + const ggml_fp16_t ninf_f16 = ggml_fp32_to_fp16(-INFINITY); + const int sliding_window = cache->sliding_window; + std::vector mask((size_t)n_kv * n_new); + for (int q = 0; q < n_new; ++q) { + const int q_abs = n_past + q; + for (int k = 0; k < n_kv; ++k) { + bool visible = k <= q_abs; + if (sliding_window > 0 && q_abs - k >= sliding_window) { + visible = false; + } + mask[(size_t)q * n_kv + k] = visible ? zero_f16 : ninf_f16; + } + } + ggml_backend_tensor_set(attn_mask, mask.data(), 0, + mask.size() * sizeof(ggml_fp16_t)); + + if (n_past == 0) return true; + + const size_t cache_bytes = (size_t)n_past * cache->pos_stride() * sizeof(float); + char name_buf[32]; + for (int il = 0; il < cache->n_layer; ++il) { + if ((size_t)il < cache->skip_layers.size() && cache->skip_layers[il]) continue; + + snprintf(name_buf, sizeof(name_buf), "Kcache.%02d", il); + struct ggml_tensor * K_t = ggml_graph_get_tensor(graph, name_buf); + snprintf(name_buf, sizeof(name_buf), "Vcache.%02d", il); + struct ggml_tensor * V_t = ggml_graph_get_tensor(graph, name_buf); + if (K_t == nullptr || V_t == nullptr) { + DIFFUSE_LOG("scheduled autoregressive K/V inputs not found for layer %d", il); + return false; + } + ggml_backend_tensor_set(K_t, cache->k_data(il), 0, cache_bytes); + ggml_backend_tensor_set(V_t, cache->v_data(il), 0, cache_bytes); + } + + return true; +} + // ── Build autoregressive forward graph ────────────────────────── // // Standard transformer with causal attention and KV cache. @@ -48,13 +114,15 @@ static struct ggml_cgraph * ar_build_graph( struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new); ggml_set_name(inp_tokens, "inp_tokens"); ggml_set_input(inp_tokens); - memcpy(inp_tokens->data, tokens, n_new * sizeof(int32_t)); + if (inp_tokens->data != nullptr) { + memcpy(inp_tokens->data, tokens, n_new * sizeof(int32_t)); + } // ── Position indices for new tokens ────────────────────────── struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new); ggml_set_name(inp_pos, "inp_pos"); ggml_set_input(inp_pos); - { + if (inp_pos->data != nullptr) { int32_t * pos_data = (int32_t *)inp_pos->data; for (int i = 0; i < n_new; i++) { pos_data[i] = n_past + i; @@ -68,7 +136,7 @@ static struct ggml_cgraph * ar_build_graph( struct ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, n_kv, n_new); ggml_set_name(attn_mask, "attn_mask"); ggml_set_input(attn_mask); - { + if (attn_mask->data != nullptr) { ggml_fp16_t * mask_data = (ggml_fp16_t *)attn_mask->data; const ggml_fp16_t zero_f16 = ggml_fp32_to_fp16(0.0f); const ggml_fp16_t ninf_f16 = ggml_fp32_to_fp16(-INFINITY); @@ -135,16 +203,20 @@ static struct ggml_cgraph * ar_build_graph( GGML_ROPE_TYPE_NEOX, 0, hp.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - // Mark K_new, V_new as outputs for cache extraction + // Materialize independent K,V outputs for cache extraction. { char name_buf[32]; + struct ggml_tensor * K_cache = ggml_dup(ctx, K_new); snprintf(name_buf, sizeof(name_buf), "Kn.%02d", il); - ggml_set_name(K_new, name_buf); - ggml_set_output(K_new); + ggml_set_name(K_cache, name_buf); + ggml_set_output(K_cache); + ggml_build_forward_expand(gf, K_cache); + struct ggml_tensor * V_cache = ggml_dup(ctx, V_new); snprintf(name_buf, sizeof(name_buf), "Vn.%02d", il); - ggml_set_name(V_new, name_buf); - ggml_set_output(V_new); + ggml_set_name(V_cache, name_buf); + ggml_set_output(V_cache); + ggml_build_forward_expand(gf, V_cache); } // ── Build full K,V: concat cached + new ───────────────── @@ -156,13 +228,22 @@ static struct ggml_cgraph * ar_build_graph( // Shape: [n_embd_head, n_head_kv, n_past] struct ggml_tensor * K_cached = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_head, n_head_kv, n_past); + char name_buf[32]; + snprintf(name_buf, sizeof(name_buf), "Kcache.%02d", il); + ggml_set_name(K_cached, name_buf); ggml_set_input(K_cached); - K_cached->data = (void *)cache->k_data(il); + if (K_cached->data != nullptr) { + K_cached->data = (void *)cache->k_data(il); + } struct ggml_tensor * V_cached = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_head, n_head_kv, n_past); + snprintf(name_buf, sizeof(name_buf), "Vcache.%02d", il); + ggml_set_name(V_cached, name_buf); ggml_set_input(V_cached); - V_cached->data = (void *)cache->v_data(il); + if (V_cached->data != nullptr) { + V_cached->data = (void *)cache->v_data(il); + } K_full = ggml_concat(ctx, K_cached, K_new, 2); V_full = ggml_concat(ctx, V_cached, V_new, 2); @@ -266,15 +347,9 @@ static bool ar_forward_impl( float * logits_out, bool profile = false) { - if (ctx->sched != nullptr) { - DIFFUSE_LOG("autoregressive execution is not yet supported with device offload"); - return false; - } - const auto & hp = ctx->model->hparams; const int n_past = cache->n_past; const int n_kv = n_past + n_new; - const int n_head = (int)hp.n_head; // ── Compute buffer sizing ──────────────────────────────────── // flash_attn_ext: no GQA expansion, no full attention matrix materialization @@ -305,13 +380,21 @@ static bool ar_forward_impl( buf_size += 256ull * 1024 * 1024; buf_size = (size_t)(buf_size * 1.2); - // Use persistent compute buffer if available (avoids malloc/free each step) - cache->ensure_compute_buf(buf_size); + const bool scheduled = ctx->sched != nullptr; + if (!scheduled) { + // CPU execution reuses the compute allocation across decode steps. + cache->ensure_compute_buf(buf_size); + } + + const size_t graph_nodes = ar_max_graph_nodes(ctx->model); + const size_t context_size = scheduled + ? ggml_tensor_overhead() * graph_nodes + ggml_graph_overhead_custom(graph_nodes, false) + : cache->compute_buf_size; struct ggml_init_params cparams = { - /*.mem_size = */ cache->compute_buf_size, - /*.mem_buffer = */ cache->compute_buf, - /*.no_alloc = */ false, + /*.mem_size = */ context_size, + /*.mem_buffer = */ scheduled ? nullptr : cache->compute_buf, + /*.no_alloc = */ scheduled, }; struct ggml_context * ctx_compute = ggml_init(cparams); if (!ctx_compute) { @@ -323,15 +406,23 @@ static bool ar_forward_impl( struct ggml_cgraph * gf = ar_build_graph(ctx, ctx_compute, tokens, n_new, n_past, cache, profile); - // Use persistent threadpool + work buffer for graph compute - struct ggml_cplan plan = ggml_graph_plan(gf, ctx->n_threads, cache->threadpool); - - if (plan.work_size > 0) { - cache->ensure_work_buf(plan.work_size); - plan.work_data = cache->work_buf; + enum ggml_status status = GGML_STATUS_SUCCESS; + if (scheduled) { + if (!diffuse_backend_sched_prepare(ctx, gf, "autoregressive", + ctx->ar_placement_logged) || + !ar_set_scheduled_inputs(gf, tokens, n_new, n_past, cache) || + !diffuse_backend_sched_compute(ctx, gf)) { + status = GGML_STATUS_FAILED; + } + } else { + // Use the persistent CPU threadpool and work buffer. + struct ggml_cplan plan = ggml_graph_plan(gf, ctx->n_threads, cache->threadpool); + if (plan.work_size > 0) { + cache->ensure_work_buf(plan.work_size); + plan.work_data = cache->work_buf; + } + status = ggml_graph_compute(gf, &plan); } - - enum ggml_status status = ggml_graph_compute(gf, &plan); if (status != GGML_STATUS_SUCCESS) { DIFFUSE_LOG("ar_forward: graph compute failed with status %d", (int)status); ggml_free(ctx_compute); @@ -345,11 +436,23 @@ static bool ar_forward_impl( ggml_free(ctx_compute); return false; } - memcpy(logits_out, logits->data, (size_t)n_new * hp.n_vocab * sizeof(float)); + const size_t logits_bytes = (size_t)n_new * hp.n_vocab * sizeof(float); + if (scheduled) { + ggml_backend_tensor_get(logits, logits_out, 0, logits_bytes); + } else { + memcpy(logits_out, logits->data, logits_bytes); + } // Extract K_new, V_new and append to cache { char name_buf[32]; + const size_t kv_count = (size_t)n_new * cache->pos_stride(); + std::vector K_new; + std::vector V_new; + if (scheduled) { + K_new.resize(kv_count); + V_new.resize(kv_count); + } for (int il = 0; il < (int)hp.n_layer; il++) { snprintf(name_buf, sizeof(name_buf), "Kn.%02d", il); struct ggml_tensor * K_t = ggml_graph_get_tensor(gf, name_buf); @@ -357,8 +460,16 @@ static bool ar_forward_impl( struct ggml_tensor * V_t = ggml_graph_get_tensor(gf, name_buf); if (K_t && V_t) { - cache->append(il, (const float *)K_t->data, - (const float *)V_t->data, n_new); + if (scheduled) { + ggml_backend_tensor_get(K_t, K_new.data(), 0, + kv_count * sizeof(float)); + ggml_backend_tensor_get(V_t, V_new.data(), 0, + kv_count * sizeof(float)); + cache->append(il, K_new.data(), V_new.data(), n_new); + } else { + cache->append(il, (const float *)K_t->data, + (const float *)V_t->data, n_new); + } } else if (cache->skip_layers.empty() || (size_t)il >= cache->skip_layers.size() || !cache->skip_layers[il]) { @@ -377,7 +488,14 @@ static bool ar_forward_impl( for (int il = 0; il < n_lay; il++) { snprintf(name_buf, sizeof(name_buf), "limp.%02d", il); struct ggml_tensor * imp_t = ggml_graph_get_tensor(gf, name_buf); - cache->layer_impact[il] = imp_t ? *(float *)imp_t->data : 0.0f; + if (imp_t == nullptr) { + cache->layer_impact[il] = 0.0f; + } else if (scheduled) { + ggml_backend_tensor_get(imp_t, &cache->layer_impact[il], 0, + sizeof(float)); + } else { + cache->layer_impact[il] = *(float *)imp_t->data; + } } } diff --git a/src/ar-sampler.cpp b/src/ar-sampler.cpp index 5a13b54..716ecf2 100644 --- a/src/ar-sampler.cpp +++ b/src/ar-sampler.cpp @@ -160,8 +160,9 @@ std::vector ar_generate( cache.init(total_ctx, (int)hp.n_layer, (int)hp.n_embd_head(), (int)hp.n_head_kv); - // Create persistent threadpool (reused across all decode steps) - { + // CPU execution reuses one threadpool across all decode steps. Device + // contexts use their backend scheduler instead. + if (ctx->sched == nullptr) { struct ggml_threadpool_params tparams = ggml_threadpool_params_default(ctx->n_threads); tparams.strict_cpu = true; // bind threads to cores tparams.prio = GGML_SCHED_PRIO_HIGH; diff --git a/src/ar-speculative.cpp b/src/ar-speculative.cpp index 53465f9..42302f8 100644 --- a/src/ar-speculative.cpp +++ b/src/ar-speculative.cpp @@ -73,8 +73,8 @@ std::vector ar_speculative_generate( draft_cache.init(total_ctx, (int)d_hp.n_layer, (int)d_hp.n_embd_head(), (int)d_hp.n_head_kv); - // ── Persistent threadpools ────────────────────────────────── - { + // ── Persistent CPU threadpools ────────────────────────────── + if (target_ctx->sched == nullptr) { auto tp = ggml_threadpool_params_default(target_ctx->n_threads); tp.strict_cpu = true; tp.prio = GGML_SCHED_PRIO_HIGH; @@ -83,7 +83,7 @@ std::vector ar_speculative_generate( target_ctx->n_threads, (int)t_hp.n_layer, (int)t_hp.n_head_kv, (int)t_hp.n_embd_head()); } - { + if (draft_ctx->sched == nullptr) { auto tp = ggml_threadpool_params_default(draft_ctx->n_threads); tp.strict_cpu = true; tp.prio = GGML_SCHED_PRIO_NORMAL; diff --git a/tests/test-ar.cpp b/tests/test-ar.cpp new file mode 100644 index 0000000..1b0d773 --- /dev/null +++ b/tests/test-ar.cpp @@ -0,0 +1,240 @@ +#include "ar-graph.h" +#include "diffuse.h" + +#include +#include +#include +#include +#include +#include + +#define ASSERT(cond, msg) do { \ + if (!(cond)) { \ + fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); \ + return 1; \ + } \ +} while (0) + +static double nmse(const float * reference, const float * actual, size_t count) { + double squared_error = 0.0; + double squared_reference = 0.0; + for (size_t i = 0; i < count; ++i) { + const double diff = (double)actual[i] - reference[i]; + squared_error += diff * diff; + squared_reference += (double)reference[i] * reference[i]; + } + return squared_error / std::max(squared_reference, 1e-30); +} + +static double cache_nmse(const ar_kv_cache & reference, + const ar_kv_cache & actual) { + const size_t count = (size_t)reference.n_past * reference.pos_stride(); + double squared_error = 0.0; + double squared_reference = 0.0; + for (int il = 0; il < reference.n_layer; ++il) { + for (size_t i = 0; i < count; ++i) { + const double k_diff = (double)actual.K[il][i] - reference.K[il][i]; + const double v_diff = (double)actual.V[il][i] - reference.V[il][i]; + squared_error += k_diff * k_diff + v_diff * v_diff; + squared_reference += (double)reference.K[il][i] * reference.K[il][i] + + (double)reference.V[il][i] * reference.V[il][i]; + } + } + return squared_error / std::max(squared_reference, 1e-30); +} + +static bool all_finite(const std::vector & values) { + return std::all_of(values.begin(), values.end(), [](float value) { + return std::isfinite(value); + }); +} + +static int compare_state(const char * label, + const std::vector & cpu_logits, + const std::vector & device_logits, + const ar_kv_cache & cpu_cache, + const ar_kv_cache & device_cache, + double max_nmse) { + if (!all_finite(device_logits)) { + fprintf(stderr, "FAIL: %s produced non-finite logits\n", label); + return 1; + } + if (cpu_cache.n_past != device_cache.n_past) { + fprintf(stderr, "FAIL: %s cache positions differ: CPU=%d device=%d\n", + label, cpu_cache.n_past, device_cache.n_past); + return 1; + } + + const double logits_error = nmse(cpu_logits.data(), device_logits.data(), + cpu_logits.size()); + const double cache_error = cache_nmse(cpu_cache, device_cache); + fprintf(stderr, "%s: logits NMSE=%.6g, K/V NMSE=%.6g\n", + label, logits_error, cache_error); + if (logits_error > max_nmse || cache_error > max_nmse) { + fprintf(stderr, "FAIL: %s should match CPU (NMSE <= %.6g)\n", + label, max_nmse); + return 1; + } + return 0; +} + +int main(int argc, char ** argv) { + if (argc < 3) { + fprintf(stderr, "SKIP: autoregressive comparison requires MODEL DEVICE\n"); + return 0; + } + + const char * model_path = argv[1]; + const char * device_name = argv[2]; + constexpr int n_ctx = 16; + constexpr int n_threads = 4; + const double max_nmse = argc > 3 ? std::atof(argv[3]) : 2e-3; + + using model_ptr = std::unique_ptr; + using context_ptr = std::unique_ptr; + model_ptr cpu_model(diffuse_model_load(model_path, n_threads), diffuse_model_free); + model_ptr device_model( + diffuse_model_load(model_path, n_threads, device_name), diffuse_model_free); + ASSERT(cpu_model != nullptr, "CPU model should load"); + ASSERT(device_model != nullptr, "device model should load"); + + const auto & hp = diffuse_model_hparams(cpu_model.get()); + context_ptr cpu_ctx( + diffuse_context_new(cpu_model.get(), n_ctx, n_threads), diffuse_context_free); + context_ptr device_ctx( + diffuse_context_new(device_model.get(), n_ctx, n_threads), diffuse_context_free); + ASSERT(cpu_ctx != nullptr, "CPU context should initialize"); + ASSERT(device_ctx != nullptr, "device context should initialize"); + + ar_kv_cache cpu_cache; + ar_kv_cache device_cache; + cpu_cache.init(n_ctx, (int)hp.n_layer, (int)hp.n_embd_head(), (int)hp.n_head_kv); + device_cache.init(n_ctx, (int)hp.n_layer, (int)hp.n_embd_head(), (int)hp.n_head_kv); + + const std::vector prompt = {1, 2, 3, 4}; + std::vector cpu_logits((size_t)prompt.size() * hp.n_vocab); + std::vector device_logits(cpu_logits.size()); + ASSERT(ar_forward_prefill(cpu_ctx.get(), prompt.data(), (int)prompt.size(), + &cpu_cache, cpu_logits.data()), + "CPU autoregressive prefill should succeed"); + ASSERT(ar_forward_prefill(device_ctx.get(), prompt.data(), (int)prompt.size(), + &device_cache, device_logits.data()), + "device autoregressive prefill should succeed"); + ASSERT(compare_state("autoregressive prefill", cpu_logits, device_logits, + cpu_cache, device_cache, max_nmse) == 0, + "autoregressive prefill should match CPU"); + + cpu_logits.resize(hp.n_vocab); + device_logits.resize(hp.n_vocab); + ASSERT(ar_forward_decode(cpu_ctx.get(), 5, &cpu_cache, cpu_logits.data()), + "CPU autoregressive decode should succeed"); + ASSERT(ar_forward_decode(device_ctx.get(), 5, &device_cache, device_logits.data()), + "device autoregressive decode should succeed"); + ASSERT(compare_state("autoregressive decode", cpu_logits, device_logits, + cpu_cache, device_cache, max_nmse) == 0, + "autoregressive decode should match CPU"); + + const std::vector batch = {6, 7}; + cpu_logits.resize((size_t)batch.size() * hp.n_vocab); + device_logits.resize(cpu_logits.size()); + ASSERT(ar_forward_batch(cpu_ctx.get(), batch.data(), (int)batch.size(), + &cpu_cache, cpu_logits.data()), + "CPU autoregressive batch should succeed"); + ASSERT(ar_forward_batch(device_ctx.get(), batch.data(), (int)batch.size(), + &device_cache, device_logits.data()), + "device autoregressive batch should succeed"); + ASSERT(compare_state("autoregressive batch", cpu_logits, device_logits, + cpu_cache, device_cache, max_nmse) == 0, + "autoregressive batch should match CPU"); + + cpu_cache.reset(); + device_cache.reset(); + cpu_cache.sliding_window = 2; + device_cache.sliding_window = 2; + cpu_logits.resize((size_t)prompt.size() * hp.n_vocab); + device_logits.resize(cpu_logits.size()); + ASSERT(ar_forward_prefill(cpu_ctx.get(), prompt.data(), (int)prompt.size(), + &cpu_cache, cpu_logits.data()), + "CPU sliding-window prefill should succeed"); + ASSERT(ar_forward_prefill(device_ctx.get(), prompt.data(), (int)prompt.size(), + &device_cache, device_logits.data()), + "device sliding-window prefill should succeed"); + ASSERT(compare_state("autoregressive sliding-window prefill", + cpu_logits, device_logits, cpu_cache, device_cache, + max_nmse) == 0, + "sliding-window prefill should match CPU"); + + cpu_cache.skip_layers.assign(hp.n_layer, false); + device_cache.skip_layers.assign(hp.n_layer, false); + cpu_cache.skip_layers[0] = true; + device_cache.skip_layers[0] = true; + cpu_logits.resize(hp.n_vocab); + device_logits.resize(hp.n_vocab); + ASSERT(ar_forward_decode(cpu_ctx.get(), 8, &cpu_cache, cpu_logits.data()), + "CPU layer-skip decode should succeed"); + ASSERT(ar_forward_decode(device_ctx.get(), 8, &device_cache, device_logits.data()), + "device layer-skip decode should succeed"); + ASSERT(compare_state("autoregressive layer-skip decode", + cpu_logits, device_logits, cpu_cache, device_cache, + max_nmse) == 0, + "layer-skip decode should match CPU"); + + cpu_cache.skip_layers.clear(); + device_cache.skip_layers.clear(); + cpu_cache.sliding_window = 0; + device_cache.sliding_window = 0; + cpu_logits.resize((size_t)prompt.size() * hp.n_vocab); + device_logits.resize(cpu_logits.size()); + ASSERT(ar_profile_layers(cpu_ctx.get(), prompt.data(), (int)prompt.size(), + &cpu_cache, cpu_logits.data(), 0), + "CPU layer profiling should succeed"); + ASSERT(ar_profile_layers(device_ctx.get(), prompt.data(), (int)prompt.size(), + &device_cache, device_logits.data(), 0), + "device layer profiling should succeed"); + ASSERT(compare_state("autoregressive layer profiling", + cpu_logits, device_logits, cpu_cache, device_cache, + max_nmse) == 0, + "layer profiling should match CPU"); + ASSERT(device_cache.layer_impact.size() == hp.n_layer, + "device layer profiling should return every layer impact"); + ASSERT(all_finite(device_cache.layer_impact), + "device layer impacts should be finite"); + + ar_sampler_params sampler_params; + sampler_params.temperature = 0.0f; + sampler_params.repeat_penalty = 1.0f; + int cpu_callbacks = 0; + int device_callbacks = 0; + const auto cpu_generated = ar_generate( + cpu_ctx.get(), prompt, 3, sampler_params, + [&cpu_callbacks](int32_t, int) { ++cpu_callbacks; return true; }); + const auto device_generated = ar_generate( + device_ctx.get(), prompt, 3, sampler_params, + [&device_callbacks](int32_t, int) { ++device_callbacks; return true; }); + ASSERT(device_generated == cpu_generated, + "greedy device autoregressive generation should match CPU"); + ASSERT(device_generated.size() == 3, + "autoregressive generation should return every requested token"); + ASSERT(cpu_callbacks == 3 && device_callbacks == 3, + "autoregressive generation should invoke each callback"); + + context_ptr draft_ctx( + diffuse_context_new(device_model.get(), n_ctx, n_threads), diffuse_context_free); + ASSERT(draft_ctx != nullptr, "device draft context should initialize"); + ar_spec_params spec_params; + spec_params.K = 2; + ar_spec_stats spec_stats; + const auto speculative = ar_speculative_generate( + device_ctx.get(), draft_ctx.get(), prompt, 4, spec_params, nullptr, &spec_stats); + ASSERT(speculative.size() == 4, + "device speculative decoding should return every requested token"); + ASSERT(spec_stats.total_generated == (int)speculative.size(), + "speculative statistics should count generated tokens"); + ASSERT(spec_stats.total_target_batches > 0, + "speculative decoding should run a device target batch"); + + cpu_cache.clear(); + device_cache.clear(); + fprintf(stderr, "PASS: %s autoregressive paths match CPU\n", device_name); + return 0; +} diff --git a/tests/test-e2e.py b/tests/test-e2e.py index 0f672cb..39d2d25 100644 --- a/tests/test-e2e.py +++ b/tests/test-e2e.py @@ -172,6 +172,24 @@ def main(): print(f"FAIL: test-cache exited with code {result.returncode}") sys.exit(1) + ar_binary = os.path.join(build_dir, "test-ar") + ar_command = [ar_binary, gguf_path, args.device] + if args.quantize: + # Quantized matmuls use backend-specific kernels and reduction + # orders. Keep the F32 gate strict while allowing the small + # cross-backend variance observed in Q4_K_M AR graphs. The AR + # test still requires exact greedy generated-token equality. + ar_command.append("0.003") + result = subprocess.run( + ar_command, + capture_output=True, text=True, env=env, timeout=30, + ) + print("C++ test-ar output:") + print(result.stderr) + if result.returncode != 0: + print(f"FAIL: test-ar exited with code {result.returncode}") + sys.exit(1) + if args.keep_gguf: output_path = os.path.abspath(args.keep_gguf) os.makedirs(os.path.dirname(output_path), exist_ok=True) diff --git a/tools/ar-cli.cpp b/tools/ar-cli.cpp index 8ac36e2..af3d453 100644 --- a/tools/ar-cli.cpp +++ b/tools/ar-cli.cpp @@ -18,6 +18,8 @@ static void print_usage(const char * prog) { fprintf(stderr, " -p TEXT Prompt text (requires external tokenization)\n"); fprintf(stderr, " -n INT Max tokens to generate (default: 256)\n"); fprintf(stderr, " -t INT Threads (default: 4)\n"); + fprintf(stderr, " --device NAME Exact GGML device name (default: CPU)\n"); + fprintf(stderr, " --list-devices List available compute devices and exit\n"); fprintf(stderr, " --temp F Temperature (default: 0 = greedy)\n"); fprintf(stderr, " --top-p F Nucleus sampling threshold (default: 0.9)\n"); fprintf(stderr, " --top-k INT Top-K sampling (default: 40, 0 = disabled)\n"); @@ -28,6 +30,7 @@ static void print_usage(const char * prog) { fprintf(stderr, " --bind-cores Bind threads to CPU cores (reduces jitter)\n"); fprintf(stderr, "\nSpeculative decoding:\n"); fprintf(stderr, " --draft PATH Draft model GGUF (enables speculative decoding)\n"); + fprintf(stderr, " --draft-device NAME Device for draft model (default: --device)\n"); fprintf(stderr, " --spec-k INT Speculative lookahead (default: 4)\n"); fprintf(stderr, " --draft-threads INT Threads for draft model (default: 4)\n"); fprintf(stderr, "\nOptimizations:\n"); @@ -52,6 +55,8 @@ static std::vector parse_tokens(const char * str) { int main(int argc, char ** argv) { std::string model_path; std::string draft_path; + std::string device_name; + std::string draft_device_name; std::string prompt; std::vector input_tokens; int max_tokens = 256; @@ -67,10 +72,17 @@ int main(int argc, char ** argv) { int repeat_last_n = 64; uint32_t seed = 42; bool bind_cores = false; + bool list_devices = false; for (int i = 1; i < argc; i++) { if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) { model_path = argv[++i]; + } else if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) { + device_name = argv[++i]; + } else if (strcmp(argv[i], "--draft-device") == 0 && i + 1 < argc) { + draft_device_name = argv[++i]; + } else if (strcmp(argv[i], "--list-devices") == 0) { + list_devices = true; } else if (strcmp(argv[i], "--draft") == 0 && i + 1 < argc) { draft_path = argv[++i]; } else if (strcmp(argv[i], "-p") == 0 && i + 1 < argc) { @@ -113,6 +125,17 @@ int main(int argc, char ** argv) { } } + if (list_devices) { + for (const auto & device : diffuse_available_devices()) { + fprintf(stdout, "%s\t%s\t%s\t%zu/%zu MiB free\n", + device.name.c_str(), diffuse_device_type_name(device.type), + device.description.c_str(), + device.memory_free / (1024 * 1024), + device.memory_total / (1024 * 1024)); + } + return 0; + } + if (model_path.empty()) { fprintf(stderr, "Error: model path required (-m)\n"); print_usage(argv[0]); @@ -122,6 +145,10 @@ int main(int argc, char ** argv) { fprintf(stderr, "Error: prompt (-p) or tokens (--tokens) required\n"); return 1; } + if (max_tokens <= 0 || n_threads <= 0 || draft_threads <= 0 || spec_k <= 0) { + fprintf(stderr, "Error: token, thread, and speculative counts must be positive\n"); + return 1; + } // Native tokenizer not integrated — use --tokens if (input_tokens.empty()) { @@ -151,14 +178,19 @@ int main(int argc, char ** argv) { if (!draft_path.empty()) { fprintf(stderr, "=== SPECULATIVE DECODING MODE ===\n"); fprintf(stderr, "Loading target model: %s\n", model_path.c_str()); - diffuse_model * target_model = diffuse_model_load(model_path, n_threads); + diffuse_model * target_model = device_name.empty() + ? diffuse_model_load(model_path, n_threads) + : diffuse_model_load(model_path, n_threads, device_name); if (!target_model) { fprintf(stderr, "Failed to load target model\n"); return 1; } fprintf(stderr, "Loading draft model: %s\n", draft_path.c_str()); - diffuse_model * draft_model = diffuse_model_load(draft_path, draft_threads); + if (draft_device_name.empty()) draft_device_name = device_name; + diffuse_model * draft_model = draft_device_name.empty() + ? diffuse_model_load(draft_path, draft_threads) + : diffuse_model_load(draft_path, draft_threads, draft_device_name); if (!draft_model) { fprintf(stderr, "Failed to load draft model\n"); diffuse_model_free(target_model); @@ -178,6 +210,14 @@ int main(int argc, char ** argv) { diffuse_context * target_ctx = diffuse_context_new(target_model, n_ctx, n_threads); diffuse_context * draft_ctx = diffuse_context_new(draft_model, n_ctx, draft_threads); + if (!target_ctx || !draft_ctx) { + fprintf(stderr, "Failed to initialize inference context\n"); + diffuse_context_free(target_ctx); + diffuse_context_free(draft_ctx); + diffuse_model_free(target_model); + diffuse_model_free(draft_model); + return 1; + } ar_spec_params sparams; sparams.K = spec_k; @@ -188,7 +228,7 @@ int main(int argc, char ** argv) { auto result = ar_speculative_generate( target_ctx, draft_ctx, input_tokens, max_tokens, sparams, - [](int32_t token, int pos) -> bool { + [](int32_t, int pos) -> bool { fprintf(stderr, "\r generating... %d tokens", pos); return true; }, @@ -218,18 +258,24 @@ int main(int argc, char ** argv) { // ── Standard AR decoding path ─────────────────────────────── fprintf(stderr, "Loading model...\n"); - diffuse_model * model = diffuse_model_load(model_path, n_threads); + diffuse_model * model = device_name.empty() + ? diffuse_model_load(model_path, n_threads) + : diffuse_model_load(model_path, n_threads, device_name); if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - const auto & hp = diffuse_model_hparams(model); int n_ctx = (int)input_tokens.size() + max_tokens; fprintf(stderr, "Prompt: %d tokens, max_generate: %d, ctx: %d\n", (int)input_tokens.size(), max_tokens, n_ctx); diffuse_context * ctx = diffuse_context_new(model, n_ctx, n_threads); + if (!ctx) { + fprintf(stderr, "Failed to initialize inference context\n"); + diffuse_model_free(model); + return 1; + } // Setup AR sampler params ar_sampler_params sparams; @@ -247,7 +293,7 @@ int main(int argc, char ** argv) { temperature <= 0.0f ? "yes" : "no", temperature, top_p, top_k); auto result = ar_generate(ctx, input_tokens, max_tokens, sparams, - [](int32_t token, int pos) -> bool { + [](int32_t, int pos) -> bool { fprintf(stderr, "\r generating... %d tokens", pos); return true; }); From fd7541d875a4cda37668dd6cdb08c9f5d12a96e3 Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Thu, 13 Aug 2026 12:27:03 +0000 Subject: [PATCH 09/10] test: add backend divergence diagnostics Assisted-by: OpenAI Codex --- tools/dump-logits.cpp | 11 +++++++++-- tools/main-cli.cpp | 18 +++++++++++++++++- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/tools/dump-logits.cpp b/tools/dump-logits.cpp index 1992a64..b61bb79 100644 --- a/tools/dump-logits.cpp +++ b/tools/dump-logits.cpp @@ -2,6 +2,7 @@ // Used for cross-validation against PyTorch. // // Usage: dump-logits -m model.gguf --tokens 1,2,3,4 -o logits.bin -t 4 +// [--device Vulkan0] // Output: binary file with float32 logits [n_tokens × n_vocab] #include "diffuse.h" @@ -27,6 +28,7 @@ static std::vector parse_tokens(const char * str) { int main(int argc, char ** argv) { std::string model_path; + std::string device_name; std::string output_path = "logits.bin"; std::vector tokens; int n_threads = 4; @@ -34,6 +36,8 @@ int main(int argc, char ** argv) { for (int i = 1; i < argc; i++) { if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) { model_path = argv[++i]; + } else if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) { + device_name = argv[++i]; } else if (strcmp(argv[i], "-o") == 0 && i + 1 < argc) { output_path = argv[++i]; } else if (strcmp(argv[i], "--tokens") == 0 && i + 1 < argc) { @@ -44,12 +48,15 @@ int main(int argc, char ** argv) { } if (model_path.empty() || tokens.empty()) { - fprintf(stderr, "Usage: dump-logits -m MODEL.gguf --tokens 1,2,3 [-o logits.bin] [-t threads]\n"); + fprintf(stderr, "Usage: dump-logits -m MODEL.gguf --tokens 1,2,3 " + "[-o logits.bin] [-t threads] [--device NAME]\n"); return 1; } fprintf(stderr, "Loading model: %s\n", model_path.c_str()); - diffuse_model * model = diffuse_model_load(model_path, n_threads); + diffuse_model * model = device_name.empty() + ? diffuse_model_load(model_path, n_threads) + : diffuse_model_load(model_path, n_threads, device_name); if (!model) return 1; const auto & hp = diffuse_model_hparams(model); diff --git a/tools/main-cli.cpp b/tools/main-cli.cpp index 4643864..ae7ba84 100644 --- a/tools/main-cli.cpp +++ b/tools/main-cli.cpp @@ -23,6 +23,7 @@ static void print_usage(const char * prog) { fprintf(stderr, " --entropy-threshold F Entropy threshold for entropy_exit (default: 1.5)\n"); fprintf(stderr, " --cache-refresh INT Force full forward every N steps (default: 0 = never)\n"); fprintf(stderr, " --cache-keep-active INT Keep recently-changed positions active N extra steps (default: 0)\n"); + fprintf(stderr, " --trace-steps Print generated token IDs after every diffusion step\n"); fprintf(stderr, "\nNote: Tokenization is currently external. Use --tokens to pass\n"); fprintf(stderr, " pre-tokenized input as comma-separated IDs.\n"); fprintf(stderr, " --tokens IDs Comma-separated token IDs (bypasses prompt)\n"); @@ -55,6 +56,7 @@ int main(int argc, char ** argv) { bool use_cache = true; int cache_refresh = 0; int cache_keep_active = 0; + bool trace_steps = false; bool list_devices = false; for (int i = 1; i < argc; i++) { @@ -95,6 +97,8 @@ int main(int argc, char ** argv) { cache_refresh = atoi(argv[++i]); } else if (strcmp(argv[i], "--cache-keep-active") == 0 && i + 1 < argc) { cache_keep_active = atoi(argv[++i]); + } else if (strcmp(argv[i], "--trace-steps") == 0) { + trace_steps = true; } else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { print_usage(argv[0]); return 0; @@ -167,9 +171,21 @@ int main(int argc, char ** argv) { // Generate fprintf(stderr, "Generating %d tokens with %d diffusion steps...\n", n_generate, n_steps); + const size_t prompt_length = input_tokens.size(); auto result = diffuse_generate(ctx, input_tokens, n_generate, sparams, - [](int step, int total, const std::vector &) { + [trace_steps, prompt_length]( + int step, int total, const std::vector & tokens) { fprintf(stderr, "\r step %d/%d", step, total); + if (trace_steps) { + fprintf(stderr, "\n trace step %d: ", step); + for (size_t i = prompt_length; i < tokens.size(); ++i) { + if (i > prompt_length) { + fputc(',', stderr); + } + fprintf(stderr, "%d", tokens[i]); + } + fputc('\n', stderr); + } }); fprintf(stderr, "\n\nGenerated token IDs:\n"); From 247386f5b6aba79084700083938dcda94f43bf04 Mon Sep 17 00:00:00 2001 From: kenjorissen Date: Tue, 18 Aug 2026 14:16:43 +0000 Subject: [PATCH 10/10] Expose backend selection through Python tools Pass exact GGML device names through generation and logit-validation wrappers, and exercise device-aware logit dumping in the synthetic end-to-end test. Assisted-by: OpenAI Codex --- tests/test-e2e.py | 5 +++-- tools/generate.py | 4 ++++ tools/validate-logits.py | 10 ++++++++-- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/test-e2e.py b/tests/test-e2e.py index 39d2d25..ab38674 100644 --- a/tests/test-e2e.py +++ b/tests/test-e2e.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 """End-to-end test: create tiny model, convert, load in C++, run forward pass.""" -import importlib.util import argparse +import importlib.util import json import os import shutil @@ -151,7 +151,8 @@ def main(): dump_binary = os.path.join(build_dir, "dump-logits") dump_tokens = [1, 2, 3, 4] dumped_logits = validator.run_cpp( - gguf_path, dump_binary, dump_tokens, n_threads=4) + gguf_path, dump_binary, dump_tokens, + n_threads=4, device=args.device) if dumped_logits.shape != (len(dump_tokens), VOCAB): print(f"FAIL: unexpected dumped logit shape {dumped_logits.shape}") sys.exit(1) diff --git a/tools/generate.py b/tools/generate.py index 7657164..23cfd2c 100644 --- a/tools/generate.py +++ b/tools/generate.py @@ -31,6 +31,8 @@ def main(): parser.add_argument("--gguf", "-g", required=True, help="GGUF model file") parser.add_argument("--cpp-bin", default="./build/diffuse-cli", help="Path to diffuse-cli binary") + parser.add_argument("--device", + help="Exact GGML device name (for example, Vulkan0 or MTL0)") parser.add_argument("-p", "--prompt", required=True, help="User prompt") parser.add_argument("-n", "--n-generate", type=int, default=256, help="Tokens to generate (default: 256)") @@ -120,6 +122,8 @@ def main(): cmd.extend(["--cache-refresh", str(args.cache_refresh)]) if args.cache_keep_active > 0: cmd.extend(["--cache-keep-active", str(args.cache_keep_active)]) + if args.device: + cmd.extend(["--device", args.device]) print(f"Running: {' '.join(cmd[:6])}...", file=sys.stderr) diff --git a/tools/validate-logits.py b/tools/validate-logits.py index 084c344..37dde54 100644 --- a/tools/validate-logits.py +++ b/tools/validate-logits.py @@ -83,7 +83,7 @@ def run_pytorch(model_dir, token_ids): return logits_np -def run_cpp(gguf_path, cpp_bin, token_ids, n_threads=4): +def run_cpp(gguf_path, cpp_bin, token_ids, n_threads=4, device=None): """Run dump-logits and return float32 logits [n_tokens, vocab_size].""" tokens_str = ",".join(map(str, token_ids)) with tempfile.TemporaryDirectory() as tmpdir: @@ -95,6 +95,9 @@ def run_cpp(gguf_path, cpp_bin, token_ids, n_threads=4): "-o", output_path, "-t", str(n_threads), ] + if device: + command.extend(["--device", device]) + result = subprocess.run(command, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError( @@ -164,6 +167,8 @@ def main(): parser.add_argument("--gguf", "-g", required=True, help="GGUF file path") parser.add_argument("--cpp-bin", default="./build/dump-logits", help="Path to the dump-logits binary") + parser.add_argument("--device", + help="Exact GGML device name (for example, Vulkan0 or MTL0)") parser.add_argument("--tokens", help="Comma-separated token IDs") parser.add_argument("--prompt", "-p", help="Text prompt (tokenized automatically)") parser.add_argument("--threads", "-t", type=int, default=4, help="C++ threads") @@ -198,7 +203,8 @@ def main(): print(f" pos {i}: {list(zip(top5.tolist(), [f'{p:.3f}' for p in probs]))}") # Run C++ - cpp_logits = run_cpp(args.gguf, args.cpp_bin, token_ids, args.threads) + cpp_logits = run_cpp( + args.gguf, args.cpp_bin, token_ids, args.threads, args.device) # Compare compare_logits(pytorch_logits, cpp_logits)