Hugging Face's logo

Misakachain
/
Qwen3.6-35B-A3B-PALW-runtime

GGUF
English
Japanese
misaka-palw
misaka
palw
proof-of-compute
deterministic-inference
integer-quantization
conversational
Model card Files
xet
 Community
Qwen3.6-35B-A3B-PALW-runtime / patches /llama.cpp-palw-observer.patch
Misakachain's picture
Misakachain
Initial import: PALW proof-of-compute runtime for Qwen3.6-35B-A3B (code + docs, no model weights)
0985c38 verified
Raw
History Blame Contribute Delete
73.8 kB
diff --git a/ggml/include/ggml-metal.h b/ggml/include/ggml-metal.h
index 433838f0..213e148d 100644
--- a/ggml/include/ggml-metal.h
+++ b/ggml/include/ggml-metal.h
@@ -56,6 +56,21 @@ GGML_BACKEND_API void ggml_backend_metal_capture_next_compute(ggml_backend_t bac
GGML_BACKEND_API ggml_backend_reg_t ggml_backend_metal_reg(void);
+// PALW kernel-level dispatch hook.
+//
+// When set (process-global), the callback is invoked for every Metal compute
+// dispatch with the bound pipeline (kernel) name and the launch geometry
+// (threadgroup grid tg0..2 and threads-per-threadgroup tptg0..2). This lets the
+// PALW observer produce a kernel-level execution trace bound to the actual GPU
+// kernel dispatches rather than to graph-node outputs. Pass NULL to disable.
+typedef void (*ggml_metal_palw_dispatch_cb)(
+ void * user_data,
+ const char * pipeline,
+ int tg0, int tg1, int tg2,
+ int tptg0, int tptg1, int tptg2);
+
+GGML_BACKEND_API void ggml_metal_palw_set_dispatch_hook(ggml_metal_palw_dispatch_cb cb, void * user_data);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 80e47f2c..3e36c08d 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1,5 +1,6 @@
 #import "ggml-metal-device.h"
+#import "ggml-metal.h"
 #import "ggml-impl.h"
 #import "ggml-backend-impl.h"
@@ -72,6 +73,9 @@ void ggml_metal_cv_set_bool(ggml_metal_cv_t cv, bool value, int32_t idx) {
struct ggml_metal_pipeline {
 id<MTLComputePipelineState> obj;
+
+ // PALW: stable kernel (pipeline) name, captured for kernel-level tracing.
+ char name[128];
 };
ggml_metal_pipeline_t ggml_metal_pipeline_init(void) {
@@ -79,6 +83,7 @@ ggml_metal_pipeline_t ggml_metal_pipeline_init(void) {
*res = (struct ggml_metal_pipeline) {
 /*.obj =*/ nil,
+ /*.name =*/ {0},
 };
return res;
@@ -443,6 +448,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
res.pipeline = ggml_metal_pipeline_init();
 res.pipeline->obj = obj;
+ // PALW: record the stable kernel name for kernel-level dispatch tracing.
+ snprintf(res.pipeline->name, sizeof(res.pipeline->name), "%s", name);
ggml_metal_pipelines_add(lib->pipelines, name, res.pipeline);
 }
@@ -458,8 +465,23 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
struct ggml_metal_encoder {
 id<MTLComputeCommandEncoder> obj;
+
+ // PALW: name of the pipeline currently bound, for kernel-level dispatch tracing.
+ const char * cur_pipeline;
 };
+// PALW: process-global kernel-dispatch hook. When set, it is invoked for every
+// Metal compute dispatch with the bound pipeline (kernel) name and the launch
+// geometry (threadgroup grid + threads-per-threadgroup). Used by the PALW
+// observer to produce a kernel-level (not graph-fallback) execution trace.
+static ggml_metal_palw_dispatch_cb g_palw_dispatch_cb = NULL;
+static void * g_palw_dispatch_ud = NULL;
+
+void ggml_metal_palw_set_dispatch_hook(ggml_metal_palw_dispatch_cb cb, void * user_data) {
+ g_palw_dispatch_cb = cb;
+ g_palw_dispatch_ud = user_data;
+}
+
 ggml_metal_encoder_t ggml_metal_encoder_init(ggml_metal_cmd_buf_t cmd_buf_raw, bool concurrent) {
 ggml_metal_encoder_t res = calloc(1, sizeof(struct ggml_metal_encoder));
@@ -491,6 +513,8 @@ void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) {
void ggml_metal_encoder_set_pipeline(ggml_metal_encoder_t encoder, struct ggml_metal_pipeline_with_params pipeline) {
 [encoder->obj setComputePipelineState:pipeline.pipeline->obj];
+ // PALW: remember the bound kernel name so dispatches can be attributed to it.
+ encoder->cur_pipeline = pipeline.pipeline->name;
 }
void ggml_metal_encoder_set_bytes(ggml_metal_encoder_t encoder, void * data, size_t size, int idx) {
@@ -506,6 +530,11 @@ void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder
 }
void ggml_metal_encoder_dispatch_threadgroups(ggml_metal_encoder_t encoder, int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2) {
+ // PALW: report the kernel-level dispatch (bound pipeline + launch geometry).
+ if (g_palw_dispatch_cb) {
+ g_palw_dispatch_cb(g_palw_dispatch_ud, encoder->cur_pipeline ? encoder->cur_pipeline : "",
+ tg0, tg1, tg2, tptg0, tptg1, tptg2);
+ }
 [encoder->obj dispatchThreadgroups:MTLSizeMake(tg0, tg1, tg2) threadsPerThreadgroup:MTLSizeMake(tptg0, tptg1, tptg2)];
 }
diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp
index 7b0876cb..4d955261 100644
--- a/src/models/qwen35moe.cpp
+++ b/src/models/qwen35moe.cpp
@@ -6,7 +6,18 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) {
 ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);
+ // PALW: newer HF->GGUF conversions write 3 mrope sections ([t, h, w]) and
+ // omit the trailing zero; older conversions write 4. Accept both forms and
+ // zero-pad so the pinned runtime loads current upstream GGUF artifacts.
+ {
+ std::vector<int32_t> sections;
+ ml.get_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, sections, true);
+ if (sections.size() != 3 && sections.size() != 4) {
+ throw std::runtime_error("rope.dimension_sections must have 3 or 4 entries");
+ }
+ std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0);
+ std::copy(sections.begin(), sections.end(), hparams.rope_sections.begin());
+ }
// Load linear attention (gated delta net) parameters
 ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
@@ -73,8 +84,13 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);
if (!hparams.is_recr(il)) {
- // Attention layers
- create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);
+ // Attention layers. The LLAMA_LOAD_LOCALS macro derives the global
+ // n_embd_k_gqa/n_embd_v_gqa from layer 0, which in this MoE is a
+ // linear-attention (recurrent) layer with n_head_kv == 0. Use the
+ // uniform full-attention KV projection width instead.
+ const int64_t n_embd_k_gqa_attn = hparams.n_embd_k_gqa_max();
+ const int64_t n_embd_v_gqa_attn = hparams.n_embd_v_gqa_max();
+ create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa_attn, n_embd_v_gqa_attn, flags);
 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);
// Q/K normalization for attention layers
@@ -86,7 +102,12 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
 layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
 layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);
- layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);
+ // PALW: some conversions store the delta-time bias without the
+ // ".bias" suffix ("blk.N.ssm_dt"). Accept both spellings.
+ layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags | TENSOR_NOT_REQUIRED);
+ if (!layer.ssm_dt) {
+ layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, il), { hparams.ssm_dt_rank }, flags);
+ }
 layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);
 layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags);
 layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags);
@@ -116,7 +137,11 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0);
 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
- create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
+ // The MTP layer index has no entry in the per-layer n_head_kv array, so
+ // derive the KV projection width from the full-attention trunk layers.
+ const int64_t n_embd_k_gqa_mtp = hparams.n_embd_k_gqa_max();
+ const int64_t n_embd_v_gqa_mtp = hparams.n_embd_v_gqa_max();
+ create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa_mtp, n_embd_v_gqa_mtp, 0);
 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0);
 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
@@ -147,6 +172,24 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
 for (int i = n_layer; i < n_layer_all; ++i) {
 load_block_mtp(i);
 }
+
+ // PALW: the Ollama-packaged Qwen3.6-35B-A3B GGUF bundles the multimodal
+ // vision tower ("v.*") and, when the next-token/MTP head is not enabled by
+ // hyperparameters, the MTP sub-model ("mtp.*") in the same file. The pinned
+ // text runtime never builds those sibling sub-models, so account for their
+ // tensors here; otherwise done_getting_tensors() rejects the load for
+ // having created fewer tensors than the file contains. This does not weaken
+ // the check for a genuinely missing text tensor: only tensors that exist in
+ // the file under these sibling prefixes are counted, exactly once each.
+ for (const auto & entry : ml.weights_map) {
+ const std::string & name = entry.first;
+ const bool is_vision = name.rfind("v.", 0) == 0;
+ const bool is_mtp = name.rfind("mtp.", 0) == 0;
+ if (is_vision || is_mtp) {
+ ml.size_data -= ggml_nbytes(entry.second.tensor);
+ ml.n_created++;
+ }
+ }
 }
std::unique_ptr<llm_graph_context> llama_model_qwen35moe::build_arch_graph(const llm_graph_params & params) const {
@@ -287,6 +330,12 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn(
 const int64_t n_embd_head = hparams.n_embd_head_v();
 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+ // PALW: the llm_graph_context base initializes n_head/n_head_kv from layer
+ // 0, which in this hybrid MoE is a linear-attention layer with n_head_kv 0.
+ // Use the per-layer counts for the full-attention reshapes below.
+ const int64_t n_head = hparams.n_head(il);
+ const int64_t n_head_kv = hparams.n_head_kv(il);
+
 // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention
// Qwen3Next uses a single Q projection that outputs query + gate
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
index 780df326..a1f04870 100644
--- a/tools/CMakeLists.txt
+++ b/tools/CMakeLists.txt
@@ -28,6 +28,7 @@ else()
 endif()
 add_subdirectory(tokenize)
 add_subdirectory(parser)
+ add_subdirectory(palw-observer)
 add_subdirectory(tts)
 add_subdirectory(mtmd)
 if (GGML_RPC)
diff --git a/tools/palw-observer/CMakeLists.txt b/tools/palw-observer/CMakeLists.txt
new file mode 100644
index 00000000..40e3ffd8
--- /dev/null
+++ b/tools/palw-observer/CMakeLists.txt
@@ -0,0 +1,9 @@
+set(TARGET llama-palw-observer)
+
+add_executable(${TARGET} palw-observer.cpp)
+target_link_libraries(${TARGET} PRIVATE llama llama-common-base ${CMAKE_THREAD_LIBS_INIT})
+target_compile_features(${TARGET} PRIVATE cxx_std_17)
+
+if(LLAMA_TOOLS_INSTALL)
+ install(TARGETS ${TARGET} RUNTIME)
+endif()
diff --git a/tools/palw-observer/README.md b/tools/palw-observer/README.md
new file mode 100644
index 00000000..efd1823b
--- /dev/null
+++ b/tools/palw-observer/README.md
@@ -0,0 +1,66 @@
+# PALW native graph observer
+
+`llama-palw-observer` is a non-interactive, single-request Qwen3.6-35B-A3B runner.
+It uses only the public llama and ggml APIs and writes versioned JSONL records to
+stdout. Runtime and model logs are written to stderr.
+
+## Build and run
+
+```sh
+cmake --build build-palw --target llama-palw-observer -j
+build-palw/bin/llama-palw-observer \
+ --model /path/to/Qwen3.6-abliterated-35b-Claude-4.7-Q4_K_M.gguf \
+ --prompt "Hello" \
+ --n-predict 8 \
+ --observer graph
+```
+
+The runner fixes greedy sampling, request batch, logical/physical token batch,
+parallel sequences, tensor parallelism, and CPU thread counts to one. It also
+disables context shifting, speculation, Flash Attention, and split-model tensor
+parallelism. The prompt is therefore evaluated one token at a time. The context
+is sized before execution and the run fails instead of shifting when the prompt
+and prediction bound exceed the model training context.
+
+The accepted model profile is hybrid Qwen3.6-35B-A3B (mixture-of-experts):
+architecture `qwen35moe`, 40 layers, 2048 hidden elements, and a 248320-token
+vocabulary, decoder-only. Per-layer head counts vary (the model alternates
+linear-attention/state-space layers with full-attention layers), so head counts
+are not part of the accepted profile. Other profiles fail before a header or
+inference result is emitted. Legitimate zero-element recurrent-state-cache graph
+nodes are skipped rather than treated as errors.
+
+## Observer modes
+
+- `off` installs no scheduler callback.
+- `graph` emits metadata for every graph node from the callback's `ask` stage.
+ It never requests tensor data.
+- `sketch` emits the same metadata and requests post-compute data only for plain
+ `MUL_MAT` nodes (indirect expert GEMM, `MUL_MAT_ID`, is emitted as an ordinary
+ metadata node, not sketched). It copies at most the first 64 contiguous output
+ elements through `ggml_backend_tensor_get`. Each element becomes one
+ hexadecimal nibble: one sign bit and a fixed three-bit magnitude bucket. The
+ resulting 64 hex digits are a 256-bit sketch. Raw activation values are never
+ serialized.
+
+Nibble bit 3 is the sign bit. Bits 0-2 use these absolute-value buckets: zero,
+`(0, 2^-8)`, `[2^-8, 2^-4)`, `[2^-4, 2^-2)`, `[2^-2, 1)`, `[1, 4)`,
+`[4, 16)`, and `[16, infinity)`. Fewer than 64 available elements are padded
+with zero nibbles.
+
+Every stdout line has `schema`, `schema_version`, and `record`. A successful run
+contains one `header`, zero or more `event` records, and one `result`. The result
+contains prompt token IDs, sampled token IDs (including an EOG token when one is
+sampled), the non-special generated output bytes, and the stop reason.
+
+## Limits
+
+This target is a runtime-observation prototype, not a PALW receipt generator.
+It does not calculate canonical compute units, sign receipts, form commitments,
+or trace CUDA kernels. The sketch tile is explicitly identified as
+`graph_fallback_logical_prefix_v1`, and `kernel_trace.claim` is
+`not_a_cuda_kernel_trace`. Scheduler callbacks add synchronization and can
+change timing, but they must not change token IDs or output bytes. The sketch is
+lossy observation data, not a cryptographic proof. Graph topology and floating
+point results can vary across backends or hardware, so byte identity is only
+claimed for repeated runs of the same pinned model, runtime, and backend class.
diff --git a/tools/palw-observer/palw-observer.cpp b/tools/palw-observer/palw-observer.cpp
new file mode 100644
index 00000000..83fc243b
--- /dev/null
+++ b/tools/palw-observer/palw-observer.cpp
@@ -0,0 +1,1473 @@
+#include "ggml-backend.h"
+#include "ggml.h"
+#include "ggml-metal.h"
+#include "llama.h"
+#include "build-info.h"
+
+#include <algorithm>
+#include <array>
+#include <charconv>
+#include <clocale>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <exception>
+#include <limits>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+namespace {
+
+constexpr const char * SCHEMA_NAME = "misaka.palw.runtime_observer";
+// Schema v2 adds "route" records: the real mixture-of-experts Top-K expert
+// selection (the `ffn_moe_topk` tensor) read back post-compute. Consumers that
+// only understand v1 must reject v2.
+constexpr int SCHEMA_VERSION = 2;
+constexpr size_t SKETCH_SAMPLES = 64;
+// Upper bound on selected-expert indices read back from one routing tensor:
+// n_expert_used (Top-K) * n_tokens for a single mixture-of-experts layer.
+constexpr size_t MAX_ROUTE_INDICES = 262144;
+
+// PALW #5 diagnostic: empirically map Metal kernel dispatches to graph nodes so
+// the node<->dispatch correlation can be validated before it is committed as a
+// kernel-level trace. Enabled only when PALW_TRACE_DIAG is set in the env.
+static bool g_diag_enabled = false;
+constexpr uint32_t PALW_CONTEXT_TOKENS = 4096;
+
+enum class observer_mode {
+ off,
+ graph,
+ sketch,
+};
+
+struct options {
+ std::string model_path;
+ std::string prompt;
+ int32_t n_predict = -1;
+ int32_t n_gpu_layers = 999;
+ observer_mode mode = observer_mode::off;
+ bool prompt_set = false;
+ bool prompt_stdin = false;
+ bool emit_output_bytes = false;
+};
+
+struct model_metadata {
+ std::vector<std::pair<std::string, std::string>> entries;
+};
+
+static const char * mode_name(observer_mode mode) {
+ switch (mode) {
+ case observer_mode::off: return "off";
+ case observer_mode::graph: return "graph";
+ case observer_mode::sketch: return "sketch";
+ }
+ return "invalid";
+}
+
+static void log_callback(ggml_log_level level, const char * text, void *) {
+ const char * label = "unknown";
+ switch (level) {
+ case GGML_LOG_LEVEL_NONE: label = "none"; break;
+ case GGML_LOG_LEVEL_DEBUG: label = "debug"; break;
+ case GGML_LOG_LEVEL_INFO: label = "info"; break;
+ case GGML_LOG_LEVEL_WARN: label = "warn"; break;
+ case GGML_LOG_LEVEL_ERROR: label = "error"; break;
+ case GGML_LOG_LEVEL_CONT: label = "cont"; break;
+ }
+ std::fprintf(stderr, "[llama:%s] %s", label, text ? text : "");
+}
+
+static void print_usage(const char * argv0) {
+ std::fprintf(stderr,
+ "usage: %s --model MODEL (--prompt TEXT|--prompt-stdin) --n-predict N "
+ "[--observer off|graph|sketch] [--n-gpu-layers N] [--emit-output-bytes]\n",
+ argv0);
+}
+
+static bool parse_i32(const char * value, int32_t min_value, int32_t max_value, int32_t & result) {
+ if (!value || value[0] == '\0') {
+ return false;
+ }
+ int32_t parsed = 0;
+ const char * end = value + std::strlen(value);
+ const auto converted = std::from_chars(value, end, parsed);
+ if (converted.ec != std::errc() || converted.ptr != end || parsed < min_value || parsed > max_value) {
+ return false;
+ }
+ result = parsed;
+ return true;
+}
+
+static bool take_value(int argc, char ** argv, int & i, const char * option, const char *& value) {
+ if (i + 1 >= argc) {
+ std::fprintf(stderr, "error: %s requires a value\n", option);
+ return false;
+ }
+ value = argv[++i];
+ return true;
+}
+
+static bool parse_options(int argc, char ** argv, options & opts, bool & help, bool & version) {
+ help = false;
+ version = false;
+ for (int i = 1; i < argc; ++i) {
+ const std::string_view arg(argv[i]);
+ const char * value = nullptr;
+ if (arg == "--help" || arg == "-h") {
+ help = true;
+ return true;
+ } else if (arg == "--version") {
+ version = true;
+ return true;
+ } else if (arg == "--model" || arg == "-m") {
+ if (!take_value(argc, argv, i, argv[i], value)) {
+ return false;
+ }
+ opts.model_path = value;
+ } else if (arg == "--prompt" || arg == "-p") {
+ if (opts.prompt_stdin) {
+ std::fprintf(stderr, "error: --prompt and --prompt-stdin are mutually exclusive\n");
+ return false;
+ }
+ if (!take_value(argc, argv, i, argv[i], value)) {
+ return false;
+ }
+ opts.prompt = value;
+ opts.prompt_set = true;
+ } else if (arg == "--prompt-stdin") {
+ if (opts.prompt_set || opts.prompt_stdin) {
+ std::fprintf(stderr, "error: prompt input may be selected only once\n");
+ return false;
+ }
+ opts.prompt_stdin = true;
+ } else if (arg == "--emit-output-bytes") {
+ opts.emit_output_bytes = true;
+ } else if (arg == "--n-predict" || arg == "-n") {
+ if (!take_value(argc, argv, i, argv[i], value) ||
+ !parse_i32(value, 0, 65536, opts.n_predict)) {
+ std::fprintf(stderr, "error: --n-predict must be an integer in [0, 65536]\n");
+ return false;
+ }
+ } else if (arg == "--n-gpu-layers" || arg == "-ngl") {
+ if (!take_value(argc, argv, i, argv[i], value) ||
+ !parse_i32(value, -1, 100000, opts.n_gpu_layers)) {
+ std::fprintf(stderr, "error: --n-gpu-layers must be an integer in [-1, 100000]\n");
+ return false;
+ }
+ } else if (arg == "--observer") {
+ if (!take_value(argc, argv, i, argv[i], value)) {
+ return false;
+ }
+ const std::string_view mode(value);
+ if (mode == "off") {
+ opts.mode = observer_mode::off;
+ } else if (mode == "graph") {
+ opts.mode = observer_mode::graph;
+ } else if (mode == "sketch") {
+ opts.mode = observer_mode::sketch;
+ } else {
+ std::fprintf(stderr, "error: --observer must be off, graph, or sketch\n");
+ return false;
+ }
+ } else {
+ std::fprintf(stderr, "error: unknown argument: %s\n", argv[i]);
+ return false;
+ }
+ }
+
+ if (opts.model_path.empty()) {
+ std::fprintf(stderr, "error: --model is required\n");
+ return false;
+ }
+ if (opts.prompt_stdin) {
+ std::array<char, 8192> buffer = {};
+ while (true) {
+ const size_t count = std::fread(buffer.data(), 1, buffer.size(), stdin);
+ if (count != 0) {
+ if (opts.prompt.size() > static_cast<size_t>(std::numeric_limits<int32_t>::max()) - count) {
+ std::fprintf(stderr, "error: stdin prompt exceeds the tokenizer API bound\n");
+ return false;
+ }
+ opts.prompt.append(buffer.data(), count);
+ }
+ if (count != buffer.size()) {
+ if (std::ferror(stdin)) {
+ std::fprintf(stderr, "error: failed to read prompt from stdin\n");
+ return false;
+ }
+ break;
+ }
+ }
+ opts.prompt_set = true;
+ }
+ if (!opts.prompt_set) {
+ std::fprintf(stderr, "error: --prompt is required (an explicitly empty prompt is allowed)\n");
+ return false;
+ }
+ if (opts.n_predict < 0) {
+ std::fprintf(stderr, "error: --n-predict is required\n");
+ return false;
+ }
+ if (opts.prompt.size() > static_cast<size_t>(std::numeric_limits<int32_t>::max())) {
+ std::fprintf(stderr, "error: prompt exceeds the tokenizer API bound\n");
+ return false;
+ }
+ return true;
+}
+
+static void append_json_string(std::string & out, std::string_view value) {
+ static constexpr char hex[] = "0123456789abcdef";
+ out.push_back('"');
+ for (const unsigned char c : value) {
+ switch (c) {
+ case '"': out += "\\\""; break;
+ case '\\': out += "\\\\"; break;
+ case '\b': out += "\\b"; break;
+ case '\f': out += "\\f"; break;
+ case '\n': out += "\\n"; break;
+ case '\r': out += "\\r"; break;
+ case '\t': out += "\\t"; break;
+ default:
+ if (c < 0x20) {
+ out += "\\u00";
+ out.push_back(hex[c >> 4]);
+ out.push_back(hex[c & 0x0f]);
+ } else {
+ out.push_back(static_cast<char>(c));
+ }
+ break;
+ }
+ }
+ out.push_back('"');
+}
+
+static void append_bool(std::string & out, bool value) {
+ out += value ? "true" : "false";
+}
+
+template<typename T>
+static void append_integer(std::string & out, T value) {
+ out += std::to_string(value);
+}
+
+static void append_float(std::string & out, double value) {
+ if (!std::isfinite(value)) {
+ out += "null";
+ return;
+ }
+ char buffer[64];
+ const int n = std::snprintf(buffer, sizeof(buffer), "%.9g", value);
+ if (n <= 0 || static_cast<size_t>(n) >= sizeof(buffer)) {
+ out += "null";
+ return;
+ }
+ out.append(buffer, static_cast<size_t>(n));
+}
+
+static bool write_json_line(const std::string & line) {
+ if (std::fwrite(line.data(), 1, line.size(), stdout) != line.size() ||
+ std::fputc('\n', stdout) == EOF || std::fflush(stdout) != 0) {
+ std::fprintf(stderr, "error: failed to write JSONL to stdout\n");
+ return false;
+ }
+ return true;
+}
+
+static void append_record_prefix(std::string & out, const char * record) {
+ out += "{\"schema\":\"";
+ out += SCHEMA_NAME;
+ out += "\",\"schema_version\":";
+ append_integer(out, SCHEMA_VERSION);
+ out += ",\"record\":";
+ append_json_string(out, record);
+}
+
+static size_t bounded_name_length(const char * name) {
+ size_t length = 0;
+ while (length < GGML_MAX_NAME && name[length] != '\0') {
+ ++length;
+ }
+ return length;
+}
+
+static std::string_view tensor_name(const ggml_tensor * tensor) {
+ if (!tensor) {
+ return {};
+ }
+ return std::string_view(tensor->name, bounded_name_length(tensor->name));
+}
+
+static std::string lowercase(std::string_view value) {
+ std::string result;
+ result.reserve(value.size());
+ for (const unsigned char c : value) {
+ if (c >= 'A' && c <= 'Z') {
+ result.push_back(static_cast<char>(c - 'A' + 'a'));
+ } else {
+ result.push_back(static_cast<char>(c));
+ }
+ }
+ return result;
+}
+
+static bool has_text(std::string_view value, std::string_view needle) {
+ return value.find(needle) != std::string_view::npos;
+}
+
+static void collect_related_names(const ggml_tensor * tensor, std::vector<std::string_view> & names) {
+ if (!tensor) {
+ return;
+ }
+ if (!tensor_name(tensor).empty()) {
+ names.push_back(tensor_name(tensor));
+ }
+ if (tensor->view_src && !tensor_name(tensor->view_src).empty()) {
+ names.push_back(tensor_name(tensor->view_src));
+ }
+ for (int i = 0; i < GGML_MAX_SRC; ++i) {
+ const ggml_tensor * src = tensor->src[i];
+ if (!src) {
+ continue;
+ }
+ if (!tensor_name(src).empty()) {
+ names.push_back(tensor_name(src));
+ }
+ if (src->view_src && !tensor_name(src->view_src).empty()) {
+ names.push_back(tensor_name(src->view_src));
+ }
+ }
+}
+
+static bool related_names_contain(const ggml_tensor * tensor, std::string_view needle) {
+ std::vector<std::string_view> names;
+ collect_related_names(tensor, names);
+ for (const auto name : names) {
+ if (has_text(lowercase(name), needle)) {
+ return true;
+ }
+ }
+ return false;
+}
+
+static bool parse_decimal_at(std::string_view name, size_t offset, int & value) {
+ if (offset >= name.size() || name[offset] < '0' || name[offset] > '9') {
+ return false;
+ }
+ uint64_t parsed = 0;
+ size_t i = offset;
+ while (i < name.size() && name[i] >= '0' && name[i] <= '9') {
+ parsed = parsed * 10 + static_cast<unsigned>(name[i] - '0');
+ if (parsed > static_cast<uint64_t>(std::numeric_limits<int>::max())) {
+ return false;
+ }
+ ++i;
+ }
+ value = static_cast<int>(parsed);
+ return true;
+}
+
+static bool parse_layer_from_name(std::string_view name, int & layer) {
+ const size_t dash = name.rfind('-');
+ if (dash != std::string_view::npos && dash + 1 < name.size()) {
+ int parsed = -1;
+ if (parse_decimal_at(name, dash + 1, parsed)) {
+ size_t end = dash + 1;
+ while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
+ ++end;
+ }
+ if (end == name.size()) {
+ layer = parsed;
+ return true;
+ }
+ }
+ }
+
+ const size_t block = name.find("blk.");
+ if (block != std::string_view::npos && parse_decimal_at(name, block + 4, layer)) {
+ return true;
+ }
+
+ if (has_text(lowercase(name), "cache_")) {
+ const size_t marker = name.rfind("_l");
+ if (marker != std::string_view::npos && parse_decimal_at(name, marker + 2, layer)) {
+ return true;
+ }
+ }
+ return false;
+}
+
+static int tensor_layer(const ggml_tensor * tensor) {
+ int layer = -1;
+ if (parse_layer_from_name(tensor_name(tensor), layer)) {
+ return layer;
+ }
+ std::vector<std::string_view> names;
+ collect_related_names(tensor, names);
+ for (const auto name : names) {
+ if (parse_layer_from_name(name, layer)) {
+ return layer;
+ }
+ }
+ return -1;
+}
+
+static void append_shape(std::string & out, const ggml_tensor * tensor) {
+ out.push_back('[');
+ for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+ if (i != 0) {
+ out.push_back(',');
+ }
+ append_integer(out, tensor->ne[i]);
+ }
+ out.push_back(']');
+}
+
+static void append_token_ids(std::string & out, const std::vector<llama_token> & tokens) {
+ out.push_back('[');
+ for (size_t i = 0; i < tokens.size(); ++i) {
+ if (i != 0) {
+ out.push_back(',');
+ }
+ append_integer(out, tokens[i]);
+ }
+ out.push_back(']');
+}
+
+class graph_observer {
+public:
+ graph_observer(observer_mode mode, int32_t n_layers) : mode_(mode), n_layers_(n_layers) {}
+
+ void set_phase(const char * phase, int64_t step) {
+ phase_ = phase;
+ phase_step_ = step;
+ }
+
+ bool failed() const {
+ return failed_;
+ }
+
+ const std::string & error() const {
+ return error_;
+ }
+
+ uint64_t event_count() const {
+ return event_sequence_;
+ }
+
+ static bool callback(ggml_tensor * tensor, bool ask, void * user_data) noexcept {
+ auto * observer = static_cast<graph_observer *>(user_data);
+ try {
+ return observer->on_tensor(tensor, ask);
+ } catch (const std::exception & exception) {
+ observer->fail(std::string("observer exception: ") + exception.what());
+ } catch (...) {
+ observer->fail("observer exception: unknown");
+ }
+ return false;
+ }
+
+private:
+ observer_mode mode_;
+ int32_t n_layers_;
+ std::string phase_ = "uninitialized";
+ int64_t phase_step_ = -1;
+ uint64_t event_sequence_ = 0;
+ bool failed_ = false;
+ std::string error_;
+
+ // Kernel-level dispatch capture (Metal). Every compute dispatch is appended
+ // here by the backend hook; each GEMM node is then attributed to the actual
+ // matmul kernel dispatch that produced it (validated 1:1 on-device).
+ struct kdispatch {
+ std::string kernel;
+ int tg[3];
+ int tptg[3];
+ };
+ std::vector<kdispatch> dispatches_;
+ size_t processed_dispatches_ = 0;
+ bool have_kernel_ = false;
+ kdispatch current_kernel_{};
+
+public:
+ // Called by the ggml-metal dispatch hook for every compute dispatch.
+ void on_dispatch(const char * kernel, int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2) {
+ kdispatch d;
+ d.kernel = kernel ? kernel : "";
+ d.tg[0] = tg0; d.tg[1] = tg1; d.tg[2] = tg2;
+ d.tptg[0] = tptg0; d.tptg[1] = tptg1; d.tptg[2] = tptg2;
+ dispatches_.push_back(std::move(d));
+ if (g_diag_enabled) {
+ std::fprintf(stderr, "PALW_DIAG dispatch #%zu pipeline=%s tg=%d,%d,%d tptg=%d,%d,%d\n",
+ dispatches_.size(), kernel ? kernel : "", tg0, tg1, tg2, tptg0, tptg1, tptg2);
+ }
+ }
+
+ static void dispatch_hook(void * user_data, const char * pipeline,
+ int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2) {
+ static_cast<graph_observer *>(user_data)->on_dispatch(
+ pipeline, tg0, tg1, tg2, tptg0, tptg1, tptg2);
+ }
+
+private:
+ // Attributes the just-computed GEMM node to the last matmul (non-expert-id)
+ // dispatch in its ASK->POST window. Expert GEMMs use distinct `*_id_*`
+ // kernels and are excluded. Returns false (fail-closed) if none is found.
+ bool resolve_gemm_kernel() {
+ have_kernel_ = false;
+ for (size_t i = dispatches_.size(); i > processed_dispatches_; --i) {
+ const std::string & k = dispatches_[i - 1].kernel;
+ const bool is_matmul =
+ (k.find("mul_mv") != std::string::npos || k.find("mul_mm") != std::string::npos) &&
+ k.find("_id") == std::string::npos;
+ if (is_matmul) {
+ current_kernel_ = dispatches_[i - 1];
+ have_kernel_ = true;
+ break;
+ }
+ }
+ processed_dispatches_ = dispatches_.size();
+ if (!have_kernel_) {
+ fail("no matmul kernel dispatch found for GEMM node");
+ return false;
+ }
+ return true;
+ }
+
+ void fail(std::string message) {
+ if (!failed_) {
+ failed_ = true;
+ error_ = std::move(message);
+ std::fprintf(stderr, "observer error: %s\n", error_.c_str());
+ }
+ }
+
+ bool validate_tensor(const ggml_tensor * tensor, bool validate_op) {
+ if (!tensor) {
+ fail("null graph tensor");
+ return false;
+ }
+ if (bounded_name_length(tensor->name) == GGML_MAX_NAME) {
+ fail("graph tensor name is not terminated");
+ return false;
+ }
+ const int type = static_cast<int>(tensor->type);
+ if (type < 0 || type >= static_cast<int>(GGML_TYPE_COUNT)) {
+ fail("graph tensor type is outside public ggml bounds");
+ return false;
+ }
+ if (validate_op) {
+ const int op = static_cast<int>(tensor->op);
+ if (op < 0 || op >= static_cast<int>(GGML_OP_COUNT)) {
+ fail("graph op is outside public ggml bounds");
+ return false;
+ }
+ }
+
+ // An empty (zero-element) tensor is valid graph metadata for recurrent
+ // state gathers and carries no storage to validate.
+ if (is_empty(tensor)) {
+ return true;
+ }
+
+ uint64_t elements = 1;
+ for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+ const uint64_t dimension = static_cast<uint64_t>(tensor->ne[i]);
+ if (elements > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / dimension) {
+ fail("graph tensor element count overflows int64");
+ return false;
+ }
+ elements *= dimension;
+ }
+ if (elements != static_cast<uint64_t>(ggml_nelements(tensor)) || ggml_nbytes(tensor) == 0) {
+ fail("graph tensor storage bounds are inconsistent");
+ return false;
+ }
+ return true;
+ }
+
+ bool validate_node(const ggml_tensor * tensor) {
+ if (!validate_tensor(tensor, true)) {
+ return false;
+ }
+ for (int i = 0; i < GGML_MAX_SRC; ++i) {
+ if (tensor->src[i] && !validate_tensor(tensor->src[i], false)) {
+ return false;
+ }
+ }
+ const int layer = tensor_layer(tensor);
+ if (layer >= n_layers_) {
+ fail("graph layer id exceeds model layer count");
+ return false;
+ }
+ if (is_gemm(tensor)) {
+ if (!tensor->src[0] || !tensor->src[1]) {
+ fail("GEMM node is missing an input");
+ return false;
+ }
+ if (tensor->src[0]->ne[0] != tensor->src[1]->ne[0] ||
+ tensor->ne[0] != tensor->src[0]->ne[1] ||
+ tensor->ne[1] != tensor->src[1]->ne[1]) {
+ fail("GEMM dimensions are inconsistent");
+ return false;
+ }
+ }
+ return true;
+ }
+
+ // Only the plain dense GEMM is sketched. Indirect expert GEMM
+ // (GGML_OP_MUL_MAT_ID) is emitted as an ordinary metadata node; its
+ // batched, id-indexed output layout is committed by the adapter as a
+ // generic compute operation rather than an accumulator sketch.
+ static bool is_gemm(const ggml_tensor * tensor) {
+ return tensor->op == GGML_OP_MUL_MAT;
+ }
+
+ static bool is_empty(const ggml_tensor * tensor) {
+ for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+ if (tensor->ne[i] <= 0) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ // The mixture-of-experts Top-K selection tensor: the concrete list of expert
+ // indices chosen per token. Captured post-compute so the receipt commits the
+ // real routing, not merely that routing of some shape occurred. Matched by
+ // the llama.cpp `cb` label and the I32 index type.
+ static bool is_route_tensor(const ggml_tensor * tensor) {
+ return tensor->type == GGML_TYPE_I32 &&
+ has_text(lowercase(tensor_name(tensor)), "ffn_moe_topk");
+ }
+
+ std::vector<std::string> categories(const ggml_tensor * tensor) const {
+ std::vector<std::string> result;
+ if (is_gemm(tensor)) {
+ result.emplace_back("gemm");
+ }
+ if (tensor->op == GGML_OP_NORM || tensor->op == GGML_OP_RMS_NORM ||
+ tensor->op == GGML_OP_GROUP_NORM || tensor->op == GGML_OP_L2_NORM) {
+ result.emplace_back("norm");
+ }
+ if (tensor->op == GGML_OP_ROPE || tensor->op == GGML_OP_ROPE_BACK) {
+ result.emplace_back("rope");
+ }
+ if (tensor->op == GGML_OP_FLASH_ATTN_EXT || tensor->op == GGML_OP_FLASH_ATTN_BACK ||
+ related_names_contain(tensor, "attn") || related_names_contain(tensor, "qcur") ||
+ related_names_contain(tensor, "kcur") || related_names_contain(tensor, "vcur") ||
+ related_names_contain(tensor, "kq")) {
+ result.emplace_back("attention");
+ }
+ if (related_names_contain(tensor, "cache_") || related_names_contain(tensor, "_cache")) {
+ result.emplace_back("kv_cache");
+ }
+ if (result.empty()) {
+ result.emplace_back("other");
+ }
+ return result;
+ }
+
+ static int magnitude_bucket(float value) {
+ const float magnitude = std::fabs(value);
+ if (magnitude == 0.0f) return 0;
+ if (magnitude < 0.00390625f) return 1;
+ if (magnitude < 0.0625f) return 2;
+ if (magnitude < 0.25f) return 3;
+ if (magnitude < 1.0f) return 4;
+ if (magnitude < 4.0f) return 5;
+ if (magnitude < 16.0f) return 6;
+ return 7;
+ }
+
+ bool make_sketch(const ggml_tensor * tensor, std::string & sketch, size_t & copied_bytes, size_t & sampled) {
+ if (!tensor->buffer || !ggml_is_contiguous(tensor)) {
+ fail("GEMM output is not backed by a contiguous public backend tensor");
+ return false;
+ }
+
+ size_t element_size = 0;
+ switch (tensor->type) {
+ case GGML_TYPE_F32: element_size = sizeof(float); break;
+ case GGML_TYPE_F16: element_size = sizeof(ggml_fp16_t); break;
+ case GGML_TYPE_BF16: element_size = sizeof(ggml_bf16_t); break;
+ default:
+ fail("GEMM output type is unsupported by sign/bucket sketch v1");
+ return false;
+ }
+
+ const uint64_t n_elements = static_cast<uint64_t>(ggml_nelements(tensor));
+ sampled = static_cast<size_t>(std::min<uint64_t>(SKETCH_SAMPLES, n_elements));
+ if (sampled > std::numeric_limits<size_t>::max() / element_size) {
+ fail("GEMM sketch copy bound overflows size_t");
+ return false;
+ }
+ copied_bytes = sampled * element_size;
+ if (copied_bytes > ggml_nbytes(tensor)) {
+ fail("GEMM sketch copy exceeds tensor storage");
+ return false;
+ }
+
+ std::vector<uint8_t> host(copied_bytes);
+ ggml_backend_tensor_get(tensor, host.data(), 0, copied_bytes);
+
+ static constexpr char hex[] = "0123456789abcdef";
+ sketch.assign(SKETCH_SAMPLES, '0');
+ for (size_t i = 0; i < sampled; ++i) {
+ float value = 0.0f;
+ if (tensor->type == GGML_TYPE_F32) {
+ std::memcpy(&value, host.data() + i * element_size, sizeof(value));
+ } else if (tensor->type == GGML_TYPE_F16) {
+ ggml_fp16_t packed;
+ std::memcpy(&packed, host.data() + i * element_size, sizeof(packed));
+ value = ggml_fp16_to_fp32(packed);
+ } else {
+ ggml_bf16_t packed;
+ std::memcpy(&packed, host.data() + i * element_size, sizeof(packed));
+ value = ggml_bf16_to_fp32(packed);
+ }
+ if (!std::isfinite(value)) {
+ fail("GEMM sketch encountered a non-finite output");
+ return false;
+ }
+ const unsigned nibble = (std::signbit(value) ? 8u : 0u) |
+ static_cast<unsigned>(magnitude_bucket(value));
+ sketch[i] = hex[nibble];
+ }
+ return true;
+ }
+
+ void append_tensor_metadata(std::string & out, const ggml_tensor * tensor) const {
+ out += "\"name\":";
+ append_json_string(out, tensor_name(tensor));
+ out += ",\"op\":";
+ append_json_string(out, ggml_op_name(tensor->op));
+ out += ",\"type\":";
+ append_json_string(out, ggml_type_name(tensor->type));
+ out += ",\"shape\":";
+ append_shape(out, tensor);
+ out += ",\"n_bytes\":";
+ append_integer(out, ggml_nbytes(tensor));
+ }
+
+ bool emit_event(const ggml_tensor * tensor, const char * stage,
+ const std::string * sketch, size_t copied_bytes, size_t sampled) {
+ std::string line;
+ line.reserve(2048);
+ append_record_prefix(line, "event");
+ line += ",\"event_seq\":";
+ append_integer(line, event_sequence_++);
+ line += ",\"observer\":";
+ append_json_string(line, mode_name(mode_));
+ line += ",\"phase\":";
+ append_json_string(line, phase_);
+ line += ",\"phase_step\":";
+ append_integer(line, phase_step_);
+ line += ",\"stage\":";
+ append_json_string(line, stage);
+
+ const int layer = tensor_layer(tensor);
+ line += ",\"layer\":";
+ if (layer < 0) {
+ line += "null";
+ } else {
+ append_integer(line, layer);
+ }
+
+ line += ",\"categories\":[";
+ const auto node_categories = categories(tensor);
+ for (size_t i = 0; i < node_categories.size(); ++i) {
+ if (i != 0) {
+ line.push_back(',');
+ }
+ append_json_string(line, node_categories[i]);
+ }
+ line += "],\"tensor\":{";
+ append_tensor_metadata(line, tensor);
+ line += "},\"sources\":[";
+ bool first_source = true;
+ for (int i = 0; i < GGML_MAX_SRC; ++i) {
+ if (!tensor->src[i]) {
+ continue;
+ }
+ if (!first_source) {
+ line.push_back(',');
+ }
+ first_source = false;
+ line.push_back('{');
+ append_tensor_metadata(line, tensor->src[i]);
+ line.push_back('}');
+ }
+ line.push_back(']');
+
+ if (is_gemm(tensor)) {
+ line += ",\"gemm\":{\"variant\":\"ggml_graph_op_v1\",\"m\":";
+ append_integer(line, tensor->src[0]->ne[1]);
+ line += ",\"n\":";
+ append_integer(line, tensor->src[1]->ne[1]);
+ line += ",\"k\":";
+ append_integer(line, tensor->src[0]->ne[0]);
+ line += ",\"batch_shape\":[";
+ append_integer(line, tensor->ne[2]);
+ line.push_back(',');
+ append_integer(line, tensor->ne[3]);
+ line += "]}";
+ }
+
+ if (sketch) {
+ line += ",\"sketch\":{\"version\":\"sign_bucket_256_v1\",\"bits\":256,";
+ line += "\"encoding\":\"hex\",\"probe\":\"contiguous_prefix_64_v1\",";
+ line += "\"sample_count\":";
+ append_integer(line, sampled);
+ line += ",\"copied_bytes\":";
+ append_integer(line, copied_bytes);
+ line += ",\"value\":";
+ append_json_string(line, *sketch);
+ line += "},\"tile\":{\"variant\":\"graph_fallback_logical_prefix_v1\",";
+ line += "\"linear_offset\":0,\"linear_elements\":";
+ append_integer(line, sampled);
+ if (have_kernel_) {
+ // Kernel-level binding: the actual Metal compute pipeline (kernel)
+ // and its launch geometry that produced this GEMM output. This is
+ // launch-geometry + output-sketch bound to a real GPU dispatch,
+ // not a CUDA-style intra-kernel accumulator sketch.
+ line += "},\"kernel_trace\":{\"available\":true,\"backend\":\"metal\",";
+ line += "\"claim\":\"metal_kernel_launch_bound_v1\",\"kernel\":";
+ append_json_string(line, current_kernel_.kernel);
+ line += ",\"threadgroups\":[";
+ append_integer(line, current_kernel_.tg[0]);
+ line.push_back(',');
+ append_integer(line, current_kernel_.tg[1]);
+ line.push_back(',');
+ append_integer(line, current_kernel_.tg[2]);
+ line += "],\"threads_per_threadgroup\":[";
+ append_integer(line, current_kernel_.tptg[0]);
+ line.push_back(',');
+ append_integer(line, current_kernel_.tptg[1]);
+ line.push_back(',');
+ append_integer(line, current_kernel_.tptg[2]);
+ line += "]}";
+ } else {
+ line += "},\"kernel_trace\":{\"available\":false,\"backend\":\"none\",";
+ line += "\"claim\":\"not_a_cuda_kernel_trace\"}";
+ }
+ }
+
+ line.push_back('}');
+ if (!write_json_line(line)) {
+ fail("stdout JSONL write failed");
+ return false;
+ }
+ return true;
+ }
+
+ // Reads the mixture-of-experts Top-K selection tensor back from the compute
+ // backend and emits a "route" record carrying the real per-token selected
+ // expert indices. ne[0] is the number of experts selected per token (Top-K);
+ // the remaining dimensions are the token count.
+ bool emit_route(const ggml_tensor * tensor) {
+ if (!tensor->buffer) {
+ fail("route tensor has no backend buffer");
+ return false;
+ }
+ const uint64_t experts_used = static_cast<uint64_t>(tensor->ne[0]);
+ const uint64_t tokens = static_cast<uint64_t>(tensor->ne[1]) *
+ static_cast<uint64_t>(tensor->ne[2]) * static_cast<uint64_t>(tensor->ne[3]);
+ const uint64_t total = experts_used * tokens;
+ if (experts_used == 0 || tokens == 0 || total > MAX_ROUTE_INDICES) {
+ fail("route tensor shape is out of bounds");
+ return false;
+ }
+
+ // Read each selected index using the tensor's byte strides, so a strided
+ // Top-K view of the argsort output is captured correctly.
+ std::vector<int32_t> indices;
+ indices.reserve(static_cast<size_t>(total));
+ for (uint64_t t = 0; t < tokens; ++t) {
+ for (uint64_t e = 0; e < experts_used; ++e) {
+ const size_t offset = static_cast<size_t>(e) * tensor->nb[0] +
+ static_cast<size_t>(t) * tensor->nb[1];
+ if (offset + sizeof(int32_t) > ggml_nbytes(tensor)) {
+ fail("route tensor read exceeds tensor storage");
+ return false;
+ }
+ int32_t value = 0;
+ ggml_backend_tensor_get(tensor, &value, offset, sizeof(int32_t));
+ indices.push_back(value);
+ }
+ }
+
+ std::string line;
+ line.reserve(1024);
+ append_record_prefix(line, "route");
+ line += ",\"event_seq\":";
+ append_integer(line, event_sequence_++);
+ line += ",\"observer\":";
+ append_json_string(line, mode_name(mode_));
+ line += ",\"phase\":";
+ append_json_string(line, phase_);
+ line += ",\"phase_step\":";
+ append_integer(line, phase_step_);
+ const int layer = tensor_layer(tensor);
+ line += ",\"layer\":";
+ if (layer < 0) {
+ line += "null";
+ } else {
+ append_integer(line, layer);
+ }
+ line += ",\"experts_used\":";
+ append_integer(line, experts_used);
+ line += ",\"tokens\":";
+ append_integer(line, tokens);
+ line += ",\"selected_experts\":[";
+ for (size_t i = 0; i < indices.size(); ++i) {
+ if (i != 0) {
+ line.push_back(',');
+ }
+ append_integer(line, indices[i]);
+ }
+ line += "]}";
+ if (!write_json_line(line)) {
+ fail("stdout JSONL write failed");
+ return false;
+ }
+ return true;
+ }
+
+ bool on_tensor(ggml_tensor * tensor, bool ask) {
+ if (failed_ || mode_ == observer_mode::off) {
+ return false;
+ }
+ // Recurrent (state-space / gated-delta-net) layers gather an initially
+ // empty state cache, producing legitimate zero-element graph nodes. They
+ // perform no compute, so they are skipped rather than treated as errors.
+ if (tensor && is_empty(tensor)) {
+ return false;
+ }
+ if (!validate_node(tensor)) {
+ return false;
+ }
+
+ if (g_diag_enabled && is_gemm(tensor)) {
+ std::fprintf(stderr, "PALW_DIAG %s MUL_MAT node=%.*s layer=%d dispatch#=%zu\n",
+ ask ? "ASK " : "POST", (int) tensor_name(tensor).size(), tensor_name(tensor).data(),
+ tensor_layer(tensor), dispatches_.size());
+ }
+
+ if (ask) {
+ if (mode_ == observer_mode::graph) {
+ emit_event(tensor, "ask_metadata", nullptr, 0, 0);
+ return false;
+ }
+ if (is_gemm(tensor)) {
+ return true;
+ }
+ // Sketch mode also reads back the mixture-of-experts Top-K selection
+ // so the receipt commits the real routing content.
+ if (is_route_tensor(tensor)) {
+ return true;
+ }
+ emit_event(tensor, "ask_metadata", nullptr, 0, 0);
+ return false;
+ }
+
+ if (mode_ != observer_mode::sketch) {
+ fail("unexpected post-compute callback");
+ return false;
+ }
+ if (is_route_tensor(tensor)) {
+ return emit_route(tensor);
+ }
+ if (!is_gemm(tensor)) {
+ fail("unexpected post-compute callback");
+ return false;
+ }
+
+ // Bind this GEMM to the actual Metal matmul kernel dispatch that produced
+ // it before sketching its output.
+ if (!resolve_gemm_kernel()) {
+ return false;
+ }
+ std::string sketch;
+ size_t copied_bytes = 0;
+ size_t sampled = 0;
+ if (!make_sketch(tensor, sketch, copied_bytes, sampled)) {
+ return false;
+ }
+ return emit_event(tensor, "post_compute_sketch", &sketch, copied_bytes, sampled);
+ }
+};
+
+static bool get_model_string(const llama_model * model, int32_t index, bool key, std::string & value) {
+ std::vector<char> buffer(256);
+ for (int attempt = 0; attempt < 3; ++attempt) {
+ const int32_t length = key
+ ? llama_model_meta_key_by_index(model, index, buffer.data(), buffer.size())
+ : llama_model_meta_val_str_by_index(model, index, buffer.data(), buffer.size());
+ if (length < 0) {
+ return false;
+ }
+ if (static_cast<size_t>(length) < buffer.size()) {
+ value.assign(buffer.data(), static_cast<size_t>(length));
+ return true;
+ }
+ buffer.resize(static_cast<size_t>(length) + 1);
+ }
+ return false;
+}
+
+static bool read_model_metadata(const llama_model * model, model_metadata & metadata) {
+ const int32_t count = llama_model_meta_count(model);
+ if (count < 0 || count > 100000) {
+ std::fprintf(stderr, "error: model metadata count is outside bounds\n");
+ return false;
+ }
+ metadata.entries.reserve(static_cast<size_t>(count));
+ for (int32_t i = 0; i < count; ++i) {
+ std::string key;
+ std::string value;
+ if (!get_model_string(model, i, true, key) || !get_model_string(model, i, false, value)) {
+ std::fprintf(stderr, "error: failed to read model metadata index %d\n", i);
+ return false;
+ }
+ metadata.entries.emplace_back(std::move(key), std::move(value));
+ }
+ return true;
+}
+
+static const std::string * find_metadata(const model_metadata & metadata, const char * key) {
+ for (const auto & entry : metadata.entries) {
+ if (entry.first == key) {
+ return &entry.second;
+ }
+ }
+ return nullptr;
+}
+
+static bool validate_qwen36_35b_a3b_profile(const llama_model * model, const model_metadata & metadata) {
+ const std::string * architecture = find_metadata(metadata, "general.architecture");
+ const llama_vocab * vocab = llama_model_get_vocab(model);
+ if (!architecture || *architecture != "qwen35moe" ||
+ llama_model_n_layer(model) != 40 ||
+ llama_model_n_embd(model) != 2048 ||
+ llama_vocab_n_tokens(vocab) != 248320 ||
+ llama_model_has_encoder(model) || !llama_model_has_decoder(model)) {
+ std::fprintf(stderr,
+ "error: model is not the supported Qwen3.6-35B-A3B MoE profile "
+ "(qwen35moe, 40 layers, 2048 hidden, 248320 vocab); observed "
+ "arch=%s layers=%d embd=%d head=%d head_kv=%d vocab=%d\n",
+ architecture ? architecture->c_str() : "(none)",
+ llama_model_n_layer(model),
+ llama_model_n_embd(model),
+ llama_model_n_head(model),
+ llama_model_n_head_kv(model),
+ llama_vocab_n_tokens(vocab));
+ return false;
+ }
+ return true;
+}
+
+static const char * device_type_name(enum ggml_backend_dev_type type) {
+ switch (type) {
+ case GGML_BACKEND_DEVICE_TYPE_CPU: return "cpu";
+ case GGML_BACKEND_DEVICE_TYPE_GPU: return "gpu";
+ case GGML_BACKEND_DEVICE_TYPE_IGPU: return "igpu";
+ case GGML_BACKEND_DEVICE_TYPE_ACCEL: return "accelerator";
+ case GGML_BACKEND_DEVICE_TYPE_META: return "meta";
+ }
+ return "unknown";
+}
+
+static bool emit_header(const options & opts, const llama_model * model, const llama_context * context,
+ const model_metadata & metadata, size_t prompt_tokens) {
+ char description[1024] = {};
+ const int32_t description_length = llama_model_desc(model, description, sizeof(description));
+ if (description_length < 0 || static_cast<size_t>(description_length) >= sizeof(description)) {
+ std::fprintf(stderr, "error: model description exceeds header bound\n");
+ return false;
+ }
+
+ std::string line;
+ line.reserve(32768);
+ append_record_prefix(line, "header");
+ line += ",\"observer\":";
+ append_json_string(line, mode_name(opts.mode));
+ line += ",\"trace_variant\":";
+ if (opts.mode == observer_mode::off) {
+ append_json_string(line, "none");
+ } else if (opts.mode == observer_mode::graph) {
+ append_json_string(line, "ggml_sched_ask_metadata_v1");
+ } else {
+ append_json_string(line, "ggml_sched_fixed_prefix_sketch_v1");
+ }
+ line += ",\"cuda_kernel_trace\":false";
+
+ line += ",\"model\":{\"path\":";
+ append_json_string(line, opts.model_path);
+ line += ",\"description\":";
+ append_json_string(line, description);
+ line += ",\"tensor_size_bytes\":";
+ append_integer(line, llama_model_size(model));
+ line += ",\"parameter_count\":";
+ append_integer(line, llama_model_n_params(model));
+ line += ",\"file_type\":";
+ append_integer(line, static_cast<int>(llama_model_ftype(model)));
+ line += ",\"n_ctx_train\":";
+ append_integer(line, llama_model_n_ctx_train(model));
+ line += ",\"n_embd\":";
+ append_integer(line, llama_model_n_embd(model));
+ line += ",\"n_layer\":";
+ append_integer(line, llama_model_n_layer(model));
+ line += ",\"n_head\":";
+ append_integer(line, llama_model_n_head(model));
+ line += ",\"n_head_kv\":";
+ append_integer(line, llama_model_n_head_kv(model));
+ line += ",\"n_vocab\":";
+ append_integer(line, llama_vocab_n_tokens(llama_model_get_vocab(model)));
+ line += ",\"rope_type\":";
+ append_integer(line, static_cast<int>(llama_model_rope_type(model)));
+ line += ",\"rope_freq_scale_train\":";
+ append_float(line, llama_model_rope_freq_scale_train(model));
+ line += ",\"metadata\":[";
+ for (size_t i = 0; i < metadata.entries.size(); ++i) {
+ if (i != 0) {
+ line.push_back(',');
+ }
+ line += "{\"key\":";
+ append_json_string(line, metadata.entries[i].first);
+ line += ",\"value\":";
+ append_json_string(line, metadata.entries[i].second);
+ line.push_back('}');
+ }
+ line += "]}";
+
+ line += ",\"runtime\":{\"ggml_version\":";
+ append_json_string(line, ggml_version());
+ line += ",\"ggml_commit\":";
+ append_json_string(line, ggml_commit());
+ line += ",\"system_info\":";
+ append_json_string(line, llama_print_system_info());
+ line += ",\"supports_gpu_offload\":";
+ append_bool(line, llama_supports_gpu_offload());
+ line += ",\"requested_gpu_layers\":";
+ append_integer(line, opts.n_gpu_layers);
+ line += ",\"devices\":[";
+ const size_t device_count = ggml_backend_dev_count();
+ for (size_t i = 0; i < device_count; ++i) {
+ if (i != 0) {
+ line.push_back(',');
+ }
+ const ggml_backend_dev_t device = ggml_backend_dev_get(i);
+ ggml_backend_dev_props properties = {};
+ ggml_backend_dev_get_props(device, &properties);
+ line += "{\"name\":";
+ append_json_string(line, properties.name ? properties.name : "");
+ line += ",\"description\":";
+ append_json_string(line, properties.description ? properties.description : "");
+ line += ",\"type\":";
+ append_json_string(line, device_type_name(properties.type));
+ line += ",\"memory_free_observed\":";
+ append_integer(line, properties.memory_free);
+ line += ",\"memory_total\":";
+ append_integer(line, properties.memory_total);
+ line.push_back('}');
+ }
+ line += "]}";
+
+ line += ",\"execution_policy\":{\"sampling\":\"greedy\",\"temperature\":0,";
+ line += "\"top_p\":1,\"top_k\":0,\"batch\":1,\"request_batch\":1,\"n_batch\":";
+ append_integer(line, llama_n_batch(context));
+ line += ",\"n_ubatch\":";
+ append_integer(line, llama_n_ubatch(context));
+ line += ",\"parallel\":1,\"parallel_sequences\":";
+ append_integer(line, llama_n_seq_max(context));
+ line += ",\"tensor_parallel\":1,\"split_mode\":\"none\",\"scheduler_parallel\":false,";
+ line += "\"tensor_repack\":false,";
+ line += "\"context_shift\":false,\"speculation\":false,\"flash_attention\":false,";
+ line += "\"threads\":1,\"threads_batch\":1,\"n_predict\":";
+ append_integer(line, opts.n_predict);
+ line += ",\"n_ctx\":";
+ append_integer(line, llama_n_ctx(context));
+ line += ",\"prompt_tokens\":";
+ append_integer(line, prompt_tokens);
+ line += ",\"tokenization\":{\"add_special\":true,\"parse_special\":true}}";
+
+ line += ",\"observation_policy\":{\"read_only\":true,";
+ line += "\"graph_metadata_stage\":\"ask\",\"sketch_bits\":256,";
+ line += "\"sketch_probe\":\"gemm_output_contiguous_prefix_64\",";
+ line += "\"raw_activation_values_published\":false,";
+ line += "\"tile_variant\":\"graph_fallback_logical_prefix_v1\",";
+ line += "\"kernel_trace_claim\":\"none\"}";
+ line.push_back('}');
+ return write_json_line(line);
+}
+
+static bool tokenize_prompt(const llama_vocab * vocab, const std::string & prompt,
+ std::vector<llama_token> & tokens) {
+ const int32_t length = static_cast<int32_t>(prompt.size());
+ const int32_t required = llama_tokenize(vocab, prompt.data(), length, nullptr, 0, true, true);
+ if (required == std::numeric_limits<int32_t>::min()) {
+ std::fprintf(stderr, "error: tokenizer result overflow\n");
+ return false;
+ }
+ const int32_t count = required < 0 ? -required : required;
+ if (count <= 0) {
+ std::fprintf(stderr, "error: prompt tokenization produced no tokens\n");
+ return false;
+ }
+ tokens.resize(static_cast<size_t>(count));
+ const int32_t actual = llama_tokenize(vocab, prompt.data(), length,
+ tokens.data(), count, true, true);
+ if (actual != count) {
+ std::fprintf(stderr, "error: prompt tokenization was not stable across sizing calls\n");
+ return false;
+ }
+ return true;
+}
+
+static bool append_token_piece(const llama_vocab * vocab, llama_token token,
+ std::vector<uint8_t> & output) {
+ std::array<char, 64> small = {};
+ int32_t length = llama_token_to_piece(vocab, token, small.data(), small.size(), 0, false);
+ if (length >= 0) {
+ output.insert(output.end(), small.begin(), small.begin() + length);
+ return true;
+ }
+ if (length == std::numeric_limits<int32_t>::min()) {
+ return false;
+ }
+ const int32_t required = -length;
+ std::vector<char> buffer(static_cast<size_t>(required));
+ length = llama_token_to_piece(vocab, token, buffer.data(), required, 0, false);
+ if (length != required) {
+ return false;
+ }
+ output.insert(output.end(), buffer.begin(), buffer.end());
+ return true;
+}
+
+static bool emit_result(const options & opts, const char * status, const char * stop_reason,
+ const std::vector<llama_token> & prompt_tokens,
+ const std::vector<llama_token> & generated_tokens,
+ const std::vector<uint8_t> & output_bytes,
+ const graph_observer & observer, const std::string & error) {
+ std::string line;
+ line.reserve(1024 + prompt_tokens.size() * 12 + generated_tokens.size() * 12 + output_bytes.size() * 4);
+ append_record_prefix(line, "result");
+ line += ",\"status\":";
+ append_json_string(line, status);
+ line += ",\"observer\":";
+ append_json_string(line, mode_name(opts.mode));
+ line += ",\"stop_reason\":";
+ append_json_string(line, stop_reason);
+ line += ",\"prompt_token_ids\":";
+ append_token_ids(line, prompt_tokens);
+ line += ",\"generated_token_ids\":";
+ append_token_ids(line, generated_tokens);
+ line += ",\"output_bytes\":[";
+ const size_t published_output_bytes = opts.emit_output_bytes ? output_bytes.size() : 0;
+ for (size_t i = 0; i < published_output_bytes; ++i) {
+ if (i != 0) {
+ line.push_back(',');
+ }
+ append_integer(line, static_cast<unsigned>(output_bytes[i]));
+ }
+ line += "],\"output_n_bytes\":";
+ append_integer(line, published_output_bytes);
+ line += ",\"event_count\":";
+ append_integer(line, observer.event_count());
+ if (!error.empty()) {
+ line += ",\"error\":";
+ append_json_string(line, error);
+ }
+ line.push_back('}');
+ return write_json_line(line);
+}
+
+struct backend_guard {
+ ~backend_guard() {
+ llama_backend_free();
+ }
+};
+
+} // namespace
+
+int main(int argc, char ** argv) {
+ std::setlocale(LC_NUMERIC, "C");
+ llama_log_set(log_callback, nullptr);
+
+ options opts;
+ bool help = false;
+ bool version = false;
+ if (!parse_options(argc, argv, opts, help, version)) {
+ print_usage(argv[0]);
+ return 2;
+ }
+ if (version) {
+ std::fprintf(stderr, "version: %d (%s)\n", llama_build_number(), llama_commit());
+ return 0;
+ }
+ if (help) {
+ print_usage(argv[0]);
+ return 0;
+ }
+
+ llama_backend_init();
+ backend_guard backend_cleanup;
+ ggml_backend_load_all();
+
+ llama_model_params model_params = llama_model_default_params();
+ model_params.n_gpu_layers = opts.n_gpu_layers;
+ model_params.split_mode = LLAMA_SPLIT_MODE_NONE;
+ model_params.main_gpu = 0;
+ model_params.use_mmap = true;
+ model_params.use_mlock = false;
+ model_params.check_tensors = true;
+ model_params.use_extra_bufts = false;
+
+ using model_ptr = std::unique_ptr<llama_model, decltype(&llama_model_free)>;
+ model_ptr model(llama_model_load_from_file(opts.model_path.c_str(), model_params), llama_model_free);
+ if (!model) {
+ std::fprintf(stderr, "error: failed to load model\n");
+ return 3;
+ }
+
+ model_metadata metadata;
+ if (!read_model_metadata(model.get(), metadata) || !validate_qwen36_35b_a3b_profile(model.get(), metadata)) {
+ return 3;
+ }
+
+ const llama_vocab * vocab = llama_model_get_vocab(model.get());
+ std::vector<llama_token> prompt_tokens;
+ if (!tokenize_prompt(vocab, opts.prompt, prompt_tokens)) {
+ return 3;
+ }
+
+ const uint64_t required_context = prompt_tokens.size() + static_cast<uint64_t>(opts.n_predict);
+ if (required_context > PALW_CONTEXT_TOKENS ||
+ PALW_CONTEXT_TOKENS > static_cast<uint64_t>(llama_model_n_ctx_train(model.get()))) {
+ std::fprintf(stderr, "error: prompt plus n-predict exceeds the fixed PALW context bound\n");
+ return 3;
+ }
+
+ graph_observer observer(opts.mode, llama_model_n_layer(model.get()));
+ llama_context_params context_params = llama_context_default_params();
+ context_params.n_ctx = PALW_CONTEXT_TOKENS;
+ context_params.n_batch = 1;
+ context_params.n_ubatch = 1;
+ context_params.n_seq_max = 1;
+ context_params.n_outputs_max = 1;
+ context_params.n_threads = 1;
+ context_params.n_threads_batch = 1;
+ context_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
+ context_params.type_k = GGML_TYPE_F16;
+ context_params.type_v = GGML_TYPE_F16;
+ context_params.embeddings = false;
+ context_params.offload_kqv = true;
+ context_params.no_perf = true;
+ context_params.op_offload = true;
+ context_params.kv_unified = false;
+ if (opts.mode != observer_mode::off) {
+ context_params.cb_eval = graph_observer::callback;
+ context_params.cb_eval_user_data = &observer;
+ g_diag_enabled = std::getenv("PALW_TRACE_DIAG") != nullptr;
+ // Capture every Metal kernel dispatch so GEMMs can be bound to the actual
+ // kernel + launch geometry that produced them (kernel-level trace).
+ if (opts.mode == observer_mode::sketch) {
+ ggml_metal_palw_set_dispatch_hook(graph_observer::dispatch_hook, &observer);
+ }
+ }
+
+ using context_ptr = std::unique_ptr<llama_context, decltype(&llama_free)>;
+ context_ptr context(llama_init_from_model(model.get(), context_params), llama_free);
+ if (!context) {
+ std::fprintf(stderr, "error: failed to initialize context\n");
+ return 3;
+ }
+ if (llama_n_batch(context.get()) != 1 || llama_n_ubatch(context.get()) != 1 ||
+ llama_n_seq_max(context.get()) != 1 || llama_n_ctx(context.get()) != PALW_CONTEXT_TOKENS ||
+ llama_n_threads(context.get()) != 1 || llama_n_threads_batch(context.get()) != 1) {
+ std::fprintf(stderr, "error: runtime did not honor the deterministic execution policy\n");
+ return 3;
+ }
+
+ if (!emit_header(opts, model.get(), context.get(), metadata, prompt_tokens.size())) {
+ return 4;
+ }
+
+ std::vector<llama_token> generated_tokens;
+ std::vector<uint8_t> output_bytes;
+ generated_tokens.reserve(static_cast<size_t>(opts.n_predict));
+
+ for (size_t i = 0; i < prompt_tokens.size(); ++i) {
+ observer.set_phase("prefill", static_cast<int64_t>(i));
+ llama_token token = prompt_tokens[i];
+ const int32_t decode_status = llama_decode(context.get(), llama_batch_get_one(&token, 1));
+ if (decode_status != 0) {
+ const std::string error = "llama_decode prefill status " + std::to_string(decode_status);
+ emit_result(opts, "error", "decode_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, error);
+ return 5;
+ }
+ if (observer.failed()) {
+ emit_result(opts, "error", "observer_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, observer.error());
+ return 5;
+ }
+ }
+
+ using sampler_ptr = std::unique_ptr<llama_sampler, decltype(&llama_sampler_free)>;
+ sampler_ptr sampler(llama_sampler_init_greedy(), llama_sampler_free);
+ if (!sampler) {
+ emit_result(opts, "error", "sampler_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, "failed to initialize greedy sampler");
+ return 5;
+ }
+
+ const char * stop_reason = "n_predict";
+ for (int32_t i = 0; i < opts.n_predict; ++i) {
+ const llama_token token = llama_sampler_sample(sampler.get(), context.get(), -1);
+ if (token < 0 || token >= llama_vocab_n_tokens(vocab)) {
+ emit_result(opts, "error", "sampler_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, "greedy sampler returned an out-of-range token");
+ return 5;
+ }
+ generated_tokens.push_back(token);
+ if (llama_vocab_is_eog(vocab, token)) {
+ stop_reason = "eog";
+ break;
+ }
+ if (!append_token_piece(vocab, token, output_bytes)) {
+ emit_result(opts, "error", "detokenize_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, "failed to convert generated token to bytes");
+ return 5;
+ }
+ if (i + 1 == opts.n_predict) {
+ break;
+ }
+
+ observer.set_phase("decode", i);
+ llama_token mutable_token = token;
+ const int32_t decode_status = llama_decode(context.get(), llama_batch_get_one(&mutable_token, 1));
+ if (decode_status != 0) {
+ const std::string error = "llama_decode generation status " + std::to_string(decode_status);
+ emit_result(opts, "error", "decode_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, error);
+ return 5;
+ }
+ if (observer.failed()) {
+ emit_result(opts, "error", "observer_error", prompt_tokens, generated_tokens,
+ output_bytes, observer, observer.error());
+ return 5;
+ }
+ }
+
+ if (!emit_result(opts, "ok", stop_reason, prompt_tokens, generated_tokens,
+ output_bytes, observer, {})) {
+ return 4;
+ }
+ return 0;
+}