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 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 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 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 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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> 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 buffer = {}; + while (true) { + const size_t count = std::fread(buffer.data(), 1, buffer.size(), stdin); + if (count != 0) { + if (opts.prompt.size() > static_cast(std::numeric_limits::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(std::numeric_limits::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(c)); + } + break; + } + } + out.push_back('"'); +} + +static void append_bool(std::string & out, bool value) { + out += value ? "true" : "false"; +} + +template +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(n) >= sizeof(buffer)) { + out += "null"; + return; + } + out.append(buffer, static_cast(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(c - 'A' + 'a')); + } else { + result.push_back(static_cast(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 & 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 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(name[i] - '0'); + if (parsed > static_cast(std::numeric_limits::max())) { + return false; + } + ++i; + } + value = static_cast(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 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 & 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(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 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(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(tensor->type); + if (type < 0 || type >= static_cast(GGML_TYPE_COUNT)) { + fail("graph tensor type is outside public ggml bounds"); + return false; + } + if (validate_op) { + const int op = static_cast(tensor->op); + if (op < 0 || op >= static_cast(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(tensor->ne[i]); + if (elements > static_cast(std::numeric_limits::max()) / dimension) { + fail("graph tensor element count overflows int64"); + return false; + } + elements *= dimension; + } + if (elements != static_cast(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 categories(const ggml_tensor * tensor) const { + std::vector 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(ggml_nelements(tensor)); + sampled = static_cast(std::min(SKETCH_SAMPLES, n_elements)); + if (sampled > std::numeric_limits::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 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(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(tensor->ne[0]); + const uint64_t tokens = static_cast(tensor->ne[1]) * + static_cast(tensor->ne[2]) * static_cast(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 indices; + indices.reserve(static_cast(total)); + for (uint64_t t = 0; t < tokens; ++t) { + for (uint64_t e = 0; e < experts_used; ++e) { + const size_t offset = static_cast(e) * tensor->nb[0] + + static_cast(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 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(length) < buffer.size()) { + value.assign(buffer.data(), static_cast(length)); + return true; + } + buffer.resize(static_cast(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(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(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(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(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 & tokens) { + const int32_t length = static_cast(prompt.size()); + const int32_t required = llama_tokenize(vocab, prompt.data(), length, nullptr, 0, true, true); + if (required == std::numeric_limits::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(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 & output) { + std::array 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::min()) { + return false; + } + const int32_t required = -length; + std::vector buffer(static_cast(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 & prompt_tokens, + const std::vector & generated_tokens, + const std::vector & 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(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; + 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 prompt_tokens; + if (!tokenize_prompt(vocab, opts.prompt, prompt_tokens)) { + return 3; + } + + const uint64_t required_context = prompt_tokens.size() + static_cast(opts.n_predict); + if (required_context > PALW_CONTEXT_TOKENS || + PALW_CONTEXT_TOKENS > static_cast(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; + 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 generated_tokens; + std::vector output_bytes; + generated_tokens.reserve(static_cast(opts.n_predict)); + + for (size_t i = 0; i < prompt_tokens.size(); ++i) { + observer.set_phase("prefill", static_cast(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; + 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; +}