From 13ab2512379e883f7856b64459b0c6cb9d444ef8 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Sun, 9 Aug 2026 18:45:41 +0530 Subject: [PATCH 1/9] test(dflash): add baseline exact trace instrumentation --- server/CMakeLists.txt | 1 + server/src/deepseek4/deepseek4_backend.cpp | 54 +- server/src/deepseek4/deepseek4_backend.h | 1 + .../src/deepseek4/deepseek4_exact_trace.cpp | 555 ++++++++++++++++++ server/src/deepseek4/deepseek4_exact_trace.h | 71 +++ server/src/deepseek4/deepseek4_graph.cpp | 80 ++- server/src/deepseek4/deepseek4_internal.h | 7 + server/tests/test_deepseek4_unit.cpp | 8 + 8 files changed, 765 insertions(+), 12 deletions(-) create mode 100644 server/src/deepseek4/deepseek4_exact_trace.cpp create mode 100644 server/src/deepseek4/deepseek4_exact_trace.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 49ad9a3bf..fa6feef0f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -388,6 +388,7 @@ add_library(dflash_common STATIC src/gemma4/gemma4_dflash_target.cpp src/gemma4/gemma4_layer_split_adapter.cpp # DeepSeek V4 Flash target arch + src/deepseek4/deepseek4_exact_trace.cpp src/deepseek4/deepseek4_loader.cpp src/deepseek4/deepseek4_graph.cpp src/deepseek4/deepseek4_roctx.cpp diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 271393c31..cccfc5cd0 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1,5 +1,6 @@ // DeepSeek4Backend implementation — AR-only decode, chunked prefill. #include "deepseek4_roctx.h" +#include "deepseek4_exact_trace.h" #include "deepseek4_backend.h" #include "deepseek4_internal.h" @@ -910,6 +911,11 @@ bool DeepSeek4Backend::init() { std::fprintf(stderr, "[deepseek4] DFLASH_DS4_SPEC set but DFLASH_DS4_DRAFT gguf missing\n"); } } + if (const char * trace_path = std::getenv("DFLASH_DS4_EXACT_TRACE_PATH"); + trace_path && *trace_path) { + exact_trace_ = DeepSeek4ExactTraceWriter::create_from_env(w_); + if (!exact_trace_) return false; + } return true; } @@ -1353,6 +1359,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, // request on can drift by a token or two. if (kv_offset == 0) { reset_deepseek4_cache(cache_); + if (exact_trace_) exact_trace_->record_reset(cache_.cur_pos); } last_logits_.clear(); last_logits_pos_ = -1; @@ -1433,6 +1440,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, std::vector hc_state; Ds4VerifyHooks spec_hooks; std::vector spec_cap; + int capture_begin = 0; Ds4VerifyHooks * hp = nullptr; const bool capture_final = i + n_tok > spec_final_from; const bool capture_snapshot = @@ -1442,7 +1450,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, (capture_final || capture_snapshot)) { spec_hooks.capture_layer_ids = &spec_drafter_->capture_layer_ids; spec_hooks.capture_out = &spec_cap; - int capture_begin = n_tok; + capture_begin = n_tok; int capture_end = 0; if (capture_final) { capture_begin = std::min( @@ -1459,6 +1467,11 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, spec_hooks.capture_token_end = capture_end; hp = &spec_hooks; } + if (exact_trace_) { + spec_hooks.exact_trace = exact_trace_.get(); + spec_hooks.allow_fused_verify = false; + hp = &spec_hooks; + } if (moe_hybrid_ && (expert_runtime_.compute || expert_backend_)) { ok = deepseek4_step_layer_range( backend_, cfg_.device.gpu, w_, cache_, hc_state, @@ -1499,6 +1512,11 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, spec_cap.begin() + (size_t) t * feat_row, spec_cap.begin() + (size_t) (t + 1) * feat_row); } + if (exact_trace_ && spec_drafter_) { + exact_trace_->record_capture( + pos, n_tok, capture_begin, + spec_drafter_->capture_layer_ids, spec_cap); + } } if (!ok) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); @@ -1508,8 +1526,16 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, add_step_tel(tel_acc, step_tel); steps++; } - last_logits_ = std::move(logits); pos += n_tok; + if (exact_trace_) { + exact_trace_->record_step( + pos - n_tok, n_tok, cache_.cur_pos, true); + if (i + n_tok == n_total || + (save_snapshot && !snapshot_saved && pos == snap_pos)) { + exact_trace_->record_logits(cache_.cur_pos, logits); + } + } + last_logits_ = std::move(logits); last_logits_pos_ = cache_.cur_pos; i += n_tok; if (save_snapshot && !snapshot_saved && pos == snap_pos) { @@ -1673,7 +1699,12 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, const DaemonIO & io) { - return generate_from_state(req, io, 0); + if (exact_trace_) exact_trace_->begin_request(req, false, 0); + GenerateResult result = generate_from_state(req, io, 0); + if (exact_trace_) { + exact_trace_->end_request(result.ok(), cache_.cur_pos, result.tokens); + } + return result; } GenerateResult DeepSeek4Backend::generate_from_state( @@ -1837,6 +1868,9 @@ bool DeepSeek4Backend::snapshot_save(int slot) { "[deepseek4] snapshot saved slot=%d pos=%d size=%.1f MiB\n", slot, snapshots_[slot].cur_pos, (double) (core_bytes + aux_bytes) / (1024.0 * 1024.0)); + if (exact_trace_) { + exact_trace_->record_snapshot("snapshot_save", slot, cache_); + } return true; } @@ -1878,6 +1912,9 @@ bool DeepSeek4Backend::snapshot_restore(int slot) { last_logits_ = std::move(restored_logits); spec_feat_window_ = std::move(restored_features); last_logits_pos_ = cache_.cur_pos; + if (exact_trace_) { + exact_trace_->record_snapshot("snapshot_restore", slot, cache_); + } return true; } @@ -1897,11 +1934,19 @@ GenerateResult DeepSeek4Backend::restore_and_generate_impl( snap_pos, req.prompt.size()); return generate_impl(req, io); } + if (exact_trace_) exact_trace_->begin_request(req, true, snap_pos); if (!snapshot_restore(slot)) { result.fail(GenerateErrorCode::BackendSpecific, "snapshot restore"); + if (exact_trace_) { + exact_trace_->end_request(false, cache_.cur_pos, result.tokens); + } return result; } - return generate_from_state(req, io, snap_pos); + result = generate_from_state(req, io, snap_pos); + if (exact_trace_) { + exact_trace_->end_request(result.ok(), cache_.cur_pos, result.tokens); + } + return result; } bool DeepSeek4Backend::handle_compress(const std::string & line, @@ -1928,6 +1973,7 @@ void DeepSeek4Backend::maybe_save_routing_stats() { void DeepSeek4Backend::shutdown() { maybe_save_routing_stats(); + exact_trace_.reset(); free_drafter(); for (int i = 0; i < PREFIX_SLOTS; i++) { snapshot_free(i); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 32f7230ce..b82e55528 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -91,6 +91,7 @@ class DeepSeek4Backend : public ModelBackend { // Absolute cache position represented by last_logits_. A snapshot is // safe only when this matches cache_.cur_pos. int last_logits_pos_ = -1; + std::unique_ptr exact_trace_; // DSpark speculative decode (opt-in: DFLASH_DS4_SPEC=1 + DFLASH_DS4_DRAFT=). bool spec_enabled_ = false; diff --git a/server/src/deepseek4/deepseek4_exact_trace.cpp b/server/src/deepseek4/deepseek4_exact_trace.cpp new file mode 100644 index 000000000..ced5e0dae --- /dev/null +++ b/server/src/deepseek4/deepseek4_exact_trace.cpp @@ -0,0 +1,555 @@ +#include "deepseek4_exact_trace.h" + +#include "common/model_backend.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +constexpr const char * kSchema = "lucebox.ds4.exact-diff/v1"; +constexpr uint64_t kFnvOffset = 14695981039346656037ULL; +constexpr uint64_t kFnvPrime = 1099511628211ULL; + +void hash_update(uint64_t & hash, const void * data, size_t bytes) { + const auto * input = static_cast(data); + for (size_t i = 0; i < bytes; ++i) { + hash ^= input[i]; + hash *= kFnvPrime; + } +} + +std::string hash_hex(uint64_t hash) { + std::ostringstream stream; + stream << std::hex << std::setfill('0') << std::setw(16) << hash; + return stream.str(); +} + +struct TensorDigest { + std::string hash; + size_t bytes = 0; +}; + +TensorDigest tensor_digest(const ggml_tensor * tensor, size_t bytes) { + if (!tensor || bytes == 0) return {}; + bytes = std::min(bytes, ggml_nbytes(tensor)); + std::vector data(bytes); + ggml_backend_tensor_get(tensor, data.data(), 0, bytes); + return {DeepSeek4ExactTraceWriter::hash_bytes(data.data(), data.size()), data.size()}; +} + +TensorDigest tensor_digest(const ggml_tensor * tensor) { + return tensor_digest(tensor, tensor ? ggml_nbytes(tensor) : 0); +} + +size_t tensor_prefix_bytes(const ggml_tensor * tensor, int rows) { + if (!tensor || rows <= 0) return 0; + return std::min( + ggml_nbytes(tensor), + ggml_row_size(tensor->type, tensor->ne[0]) * static_cast(rows)); +} + +void write_nullable_string(std::ofstream & output, const std::string & value) { + if (value.empty()) { + output << "null"; + } else { + output << '"' << value << '"'; + } +} + +bool write_float_array( + std::ofstream & output, + const float * values, + size_t count, + size_t begin = 0, + size_t end = std::numeric_limits::max()) { + begin = std::min(begin, count); + end = std::min(end, count); + bool non_finite = false; + output << '['; + for (size_t i = begin; i < end; ++i) { + if (i != begin) output << ','; + if (std::isfinite(values[i])) { + output << std::setprecision(std::numeric_limits::max_digits10) + << values[i]; + } else { + output << "null"; + non_finite = true; + } + } + output << ']'; + return non_finite; +} + +void write_int_array( + std::ofstream & output, + const int32_t * values, + size_t begin, + size_t end) { + output << '['; + for (size_t i = begin; i < end; ++i) { + if (i != begin) output << ','; + output << values[i]; + } + output << ']'; +} + +bool near_boundary(int begin, int end, int boundary) { + return begin <= boundary + 4 && end >= boundary - 4; +} + +void append_unique(std::vector & values, std::string value) { + if (std::find(values.begin(), values.end(), value) == values.end()) { + values.push_back(std::move(value)); + } +} + +std::string boundary_relation( + const std::string & prefix, + int begin, + int end, + int boundary) { + if (begin < boundary && end < boundary) return prefix + "_before"; + if (begin < boundary && end >= boundary) return prefix + "_on"; + return prefix + "_after"; +} + +} // namespace + +std::unique_ptr +DeepSeek4ExactTraceWriter::create_from_env(const DeepSeek4Weights & weights) { + const char * path = std::getenv("DFLASH_DS4_EXACT_TRACE_PATH"); + if (!path || !*path) return nullptr; + + const char * raw_width = std::getenv("DFLASH_DS4_EXACT_TRACE_Q"); + char * end = nullptr; + const long parsed = raw_width ? std::strtol(raw_width, &end, 10) : 0; + if (!raw_width || !*raw_width || !end || *end || parsed < 1 || parsed > 4) { + std::fprintf(stderr, + "[deepseek4-exact-trace] DFLASH_DS4_EXACT_TRACE_Q must be 1..4\n"); + return nullptr; + } + + std::ofstream output(path, std::ios::out | std::ios::app); + if (!output) { + std::fprintf(stderr, + "[deepseek4-exact-trace] cannot append trace path %s\n", path); + return nullptr; + } + std::fprintf(stderr, + "[deepseek4-exact-trace] enabled path=%s q=%ld\n", path, parsed); + return std::unique_ptr( + new DeepSeek4ExactTraceWriter(weights, std::move(output), static_cast(parsed))); +} + +DeepSeek4ExactTraceWriter::DeepSeek4ExactTraceWriter( + const DeepSeek4Weights & weights, + std::ofstream output, + int width) + : weights_(weights), output_(std::move(output)), width_(width) { + std::set unique; + for (uint32_t ratio : weights_.compress_ratios) { + if (ratio > 0) unique.insert(static_cast(ratio)); + } + compressor_boundaries_.assign(unique.begin(), unique.end()); + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"model_config\",\"n_swa\":" << weights_.n_swa + << ",\"compressor_boundaries\":["; + for (size_t i = 0; i < compressor_boundaries_.size(); ++i) { + if (i) output_ << ','; + output_ << compressor_boundaries_[i]; + } + output_ << "]}\n"; + output_.flush(); +} + +bool DeepSeek4ExactTraceWriter::begin_request( + const GenerateRequest & request, + bool restored, + int cache_position) { + ++request_index_; + prompt_tokens_ = static_cast(request.prompt.size()); + snapshot_position_ = request.snap_pos; + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"request_start\",\"request\":" << request_index_ + << ",\"width\":" << width_ + << ",\"restored\":" << (restored ? "true" : "false") + << ",\"cache_position\":" << cache_position + << ",\"prompt_tokens\":" << request.prompt.size() + << ",\"prompt_token_hash\":\"" << hash_token_ids(request.prompt) << "\"" + << ",\"prompt_token_ids\":"; + write_int_array(output_, request.prompt.data(), 0, request.prompt.size()); + output_ + << ",\"n_gen\":" << request.n_gen + << ",\"snap_slot\":" << request.snap_slot + << ",\"snap_pos\":" << request.snap_pos + << ",\"temperature\":" << request.sampler.temp + << ",\"top_p\":" << request.sampler.top_p + << ",\"top_k\":" << request.sampler.top_k + << ",\"seed\":" << request.sampler.seed << "}\n"; + output_.flush(); + return static_cast(output_); +} + +void DeepSeek4ExactTraceWriter::end_request( + bool ok, + int cache_position, + const std::vector & tokens) { + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"tokens\",\"request\":" << request_index_ + << ",\"token_ids\":"; + write_int_array(output_, tokens.data(), 0, tokens.size()); + output_ << "}\n{\"schema\":\"" << kSchema + << "\",\"type\":\"request_end\",\"request\":" << request_index_ + << ",\"ok\":" << (ok ? "true" : "false") + << ",\"cache_position\":" << cache_position << "}\n"; + output_.flush(); +} + +bool DeepSeek4ExactTraceWriter::wants_step(int position_begin, int n_tokens) const { + if (request_index_ < 0 || n_tokens <= 0) return false; + const int position_end = position_begin + n_tokens; + if (position_begin < 16 || position_end >= std::max(0, prompt_tokens_ - 8)) return true; + if (snapshot_position_ >= 0 && + near_boundary(position_begin, position_end, snapshot_position_)) { + return true; + } + if (weights_.n_swa > 0 && near_boundary(position_begin, position_end, weights_.n_swa)) { + return true; + } + for (int ratio : compressor_boundaries_) { + const int first = ratio; + const int last = prompt_tokens_ / ratio * ratio; + if (near_boundary(position_begin, position_end, first) || + (last > first && near_boundary(position_begin, position_end, last))) { + return true; + } + } + const int capture_begin = std::max(0, prompt_tokens_ - weights_.n_swa); + return near_boundary(position_begin, position_end, capture_begin); +} + +void DeepSeek4ExactTraceWriter::record_reset(int cache_position) { + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"reset\",\"request\":" << request_index_ + << ",\"cache_position\":" << cache_position << "}\n"; +} + +void DeepSeek4ExactTraceWriter::write_relations(int position_begin, int position_end) { + std::vector relations; + const int step_width = position_end - position_begin; + if (step_width == width_ && position_end < prompt_tokens_) { + append_unique(relations, "ordinary"); + } + if (position_end == prompt_tokens_ && step_width < width_) { + append_unique(relations, "tail_width_" + std::to_string(step_width)); + } + for (int ratio : compressor_boundaries_) { + const int first = ratio; + const int last = prompt_tokens_ / ratio * ratio; + if (near_boundary(position_begin, position_end, first)) { + append_unique(relations, + boundary_relation("compressor", position_begin, position_end, first)); + append_unique( + relations, + boundary_relation( + "compressor_" + std::to_string(ratio), + position_begin, position_end, first)); + } + if (last > first && near_boundary(position_begin, position_end, last)) { + append_unique(relations, + boundary_relation("compressor", position_begin, position_end, last)); + append_unique( + relations, + boundary_relation( + "compressor_" + std::to_string(ratio), + position_begin, position_end, last)); + } + } + if (weights_.n_swa > 0 && near_boundary(position_begin, position_end, weights_.n_swa)) { + append_unique(relations, + boundary_relation("swa", position_begin, position_end, weights_.n_swa)); + append_unique( + relations, + boundary_relation( + "swa_" + std::to_string(weights_.n_swa), + position_begin, position_end, weights_.n_swa)); + } + if (snapshot_position_ >= 0 && + near_boundary(position_begin, position_end, snapshot_position_)) { + append_unique(relations, "snapshot_boundary"); + } + const int capture_begin = std::max(0, prompt_tokens_ - weights_.n_swa); + if (near_boundary(position_begin, position_end, capture_begin) || + position_end >= std::max(0, prompt_tokens_ - 8)) { + append_unique(relations, "dspark_capture_window"); + } + output_ << '['; + for (size_t i = 0; i < relations.size(); ++i) { + if (i) output_ << ','; + output_ << '"' << relations[i] << '"'; + } + output_ << ']'; +} + +void DeepSeek4ExactTraceWriter::record_step( + int position_begin, + int n_tokens, + int cache_position, + bool logits_present) { + if (!wants_step(position_begin, n_tokens)) return; + const int position_end = position_begin + n_tokens; + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"step\",\"request\":" << request_index_ + << ",\"position_begin\":" << position_begin + << ",\"position_end\":" << position_end + << ",\"position\":" << position_end + << ",\"width\":" << n_tokens + << ",\"cache_position\":" << cache_position + << ",\"logits_present\":" << (logits_present ? "true" : "false") + << ",\"relations\":"; + write_relations(position_begin, position_end); + output_ << "}\n"; +} + +void DeepSeek4ExactTraceWriter::record_layer( + ggml_backend_t, + const DeepSeek4Cache & cache, + const std::vector & hc_state, + int layer, + int position_begin, + int n_tokens, + bool hash_routed, + const std::vector & routing_ids, + const std::vector & routing_weights) { + if (!wants_step(position_begin, n_tokens) || layer < 0 || + layer >= static_cast(cache.layers.size()) || n_tokens <= 0) { + return; + } + const DeepSeek4LayerCache & layer_cache = cache.layers[static_cast(layer)]; + const int position_end = position_begin + n_tokens; + const int ratio = static_cast(weights_.compress_ratios[static_cast(layer)]); + const int active_raw_rows = std::min(position_end, weights_.n_swa); + const int active_comp_rows = ratio > 0 + ? std::max(layer_cache.n_comp, position_end / ratio) : 0; + const int active_index_rows = ratio == 4 + ? std::max(layer_cache.n_index_comp, position_end / ratio) : 0; + const TensorDigest raw = tensor_digest( + layer_cache.raw_kv, tensor_prefix_bytes(layer_cache.raw_kv, active_raw_rows)); + const TensorDigest compressed = tensor_digest( + layer_cache.comp_kv, tensor_prefix_bytes(layer_cache.comp_kv, active_comp_rows)); + const TensorDigest index_compressed = tensor_digest( + layer_cache.index_comp_kv, + tensor_prefix_bytes(layer_cache.index_comp_kv, active_index_rows)); + const TensorDigest attn_state_kv = tensor_digest(layer_cache.attn_compressor.state_kv); + const TensorDigest attn_state_score = tensor_digest(layer_cache.attn_compressor.state_score); + const TensorDigest index_state_kv = tensor_digest(layer_cache.indexer_compressor.state_kv); + const TensorDigest index_state_score = tensor_digest(layer_cache.indexer_compressor.state_score); + + const size_t hc_stride = static_cast(weights_.n_hc) * weights_.n_embd; + const size_t hc_begin = hc_stride * static_cast(n_tokens - 1); + const bool hc_available = hc_begin + hc_stride <= hc_state.size(); + const std::string hc_hash = hc_available + ? hash_bytes(hc_state.data() + hc_begin, hc_stride * sizeof(float)) : std::string{}; + bool non_finite = false; + if (hc_available) { + for (size_t i = hc_begin; i < hc_begin + hc_stride; ++i) { + non_finite = non_finite || !std::isfinite(hc_state[i]); + } + } + + const size_t route_width = routing_ids.size() >= static_cast(n_tokens) + ? routing_ids.size() / static_cast(n_tokens) : 0; + const size_t route_begin = route_width * static_cast(n_tokens - 1); + const bool route_available = route_width > 0 && + route_begin + route_width <= routing_ids.size() && + route_begin + route_width <= routing_weights.size(); + if (route_available) { + for (size_t i = route_begin; i < route_begin + route_width; ++i) { + non_finite = non_finite || !std::isfinite(routing_weights[i]); + } + } + + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"layer\",\"request\":" << request_index_ + << ",\"layer\":" << layer + << ",\"position_begin\":" << position_begin + << ",\"position_end\":" << position_end + << ",\"token_position\":" << (position_end - 1) + << ",\"routing\":{\"mode\":\"" + << (hash_routed ? "hash" : "learned") << "\",\"ids\":"; + if (route_available) { + write_int_array(output_, routing_ids.data(), route_begin, route_begin + route_width); + } else { + output_ << "[]"; + } + output_ << ",\"weights\":"; + if (route_available) { + write_float_array( + output_, routing_weights.data(), routing_weights.size(), + route_begin, route_begin + route_width); + } else { + output_ << "[]"; + } + output_ << "},\"hc_token_hashes\":["; + write_nullable_string(output_, hc_hash); + output_ << "],\"raw_kv_hash\":"; + write_nullable_string(output_, raw.hash); + output_ << ",\"raw_kv_bytes\":" << raw.bytes + << ",\"compressed_kv_hash\":"; + write_nullable_string(output_, compressed.hash); + output_ << ",\"compressed_kv_bytes\":" << compressed.bytes + << ",\"compressor_state\":{\"kv\":"; + write_nullable_string(output_, attn_state_kv.hash); + output_ << ",\"score\":"; + write_nullable_string(output_, attn_state_score.hash); + output_ << ",\"n_comp\":" << active_comp_rows + << "},\"indexer_state\":{\"compressed_kv\":"; + write_nullable_string(output_, index_compressed.hash); + output_ << ",\"compressed_kv_bytes\":" << index_compressed.bytes + << ",\"kv\":"; + write_nullable_string(output_, index_state_kv.hash); + output_ << ",\"score\":"; + write_nullable_string(output_, index_state_score.hash); + output_ << ",\"n_comp\":" << active_index_rows + << "},\"non_finite\":" << (non_finite ? "true" : "false") << "}\n"; +} + +void DeepSeek4ExactTraceWriter::record_capture( + int position_begin, + int n_tokens, + int capture_begin, + const std::vector & layer_ids, + const std::vector & rows) { + const size_t row_width = static_cast(weights_.n_embd) * layer_ids.size(); + if (rows.empty() || n_tokens <= 0 || capture_begin < 0 || + capture_begin >= n_tokens || row_width == 0 || rows.size() % row_width != 0) { + return; + } + const size_t captured_rows = rows.size() / row_width; + if (capture_begin + static_cast(captured_rows) > n_tokens) return; + for (size_t row_index = 0; row_index < captured_rows; ++row_index) { + const float * row = rows.data() + row_index * row_width; + bool non_finite = false; + for (size_t i = 0; i < row_width; ++i) { + non_finite = non_finite || !std::isfinite(row[i]); + } + const int position = position_begin + capture_begin + static_cast(row_index); + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"capture\",\"request\":" << request_index_ + << ",\"position_begin\":" << position + << ",\"position_end\":" << (position + 1) + << ",\"layer_ids\":["; + for (size_t i = 0; i < layer_ids.size(); ++i) { + if (i) output_ << ','; + output_ << layer_ids[i]; + } + output_ << "],\"row_hash\":\"" << hash_bytes(row, row_width * sizeof(float)) + << "\",\"total_values\":" << row_width << ",\"rows\":"; + non_finite = write_float_array(output_, row, row_width) || non_finite; + output_ << ",\"non_finite\":" << (non_finite ? "true" : "false") << "}\n"; + } +} + +void DeepSeek4ExactTraceWriter::record_logits( + int position, + const std::vector & logits) { + bool non_finite = false; + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"logits\",\"request\":" << request_index_ + << ",\"position\":" << position + << ",\"hash\":\"" << hash_bytes(logits.data(), logits.size() * sizeof(float)) + << "\",\"values\":"; + non_finite = write_float_array(output_, logits.data(), logits.size()); + output_ << ",\"non_finite\":" << (non_finite ? "true" : "false") << "}\n"; +} + +std::string DeepSeek4ExactTraceWriter::cache_state_hash( + const DeepSeek4Cache & cache) const { + uint64_t hash = kFnvOffset; + hash_update(hash, &cache.cur_pos, sizeof(cache.cur_pos)); + for (size_t layer_index = 0; layer_index < cache.layers.size(); ++layer_index) { + const DeepSeek4LayerCache & layer = cache.layers[layer_index]; + hash_update(hash, &layer.n_comp, sizeof(layer.n_comp)); + hash_update(hash, &layer.n_index_comp, sizeof(layer.n_index_comp)); + const int raw_rows = std::min(cache.cur_pos, weights_.n_swa); + const std::pair tensors[] = { + {layer.raw_kv, tensor_prefix_bytes(layer.raw_kv, raw_rows)}, + {layer.comp_kv, tensor_prefix_bytes(layer.comp_kv, layer.n_comp)}, + {layer.index_comp_kv, + tensor_prefix_bytes(layer.index_comp_kv, layer.n_index_comp)}, + {layer.attn_compressor.state_kv, + layer.attn_compressor.state_kv + ? ggml_nbytes(layer.attn_compressor.state_kv) : 0}, + {layer.attn_compressor.state_score, + layer.attn_compressor.state_score + ? ggml_nbytes(layer.attn_compressor.state_score) : 0}, + {layer.indexer_compressor.state_kv, + layer.indexer_compressor.state_kv + ? ggml_nbytes(layer.indexer_compressor.state_kv) : 0}, + {layer.indexer_compressor.state_score, + layer.indexer_compressor.state_score + ? ggml_nbytes(layer.indexer_compressor.state_score) : 0}, + }; + for (const auto & [tensor, bytes] : tensors) { + if (!tensor) continue; + std::vector data(bytes); + if (data.empty()) continue; + ggml_backend_tensor_get(tensor, data.data(), 0, data.size()); + hash_update(hash, data.data(), data.size()); + } + } + if (cache.hc_state) { + std::vector data(ggml_nbytes(cache.hc_state)); + ggml_backend_tensor_get(cache.hc_state, data.data(), 0, data.size()); + hash_update(hash, data.data(), data.size()); + } + return hash_hex(hash); +} + +void DeepSeek4ExactTraceWriter::record_snapshot( + const char * kind, + int slot, + const DeepSeek4Cache & cache) { + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"" << kind << "\",\"request\":" << request_index_ + << ",\"slot\":" << slot + << ",\"cache_position\":" << cache.cur_pos + << ",\"state_hash\":\"" << cache_state_hash(cache) << "\"}\n"; +} + +std::string DeepSeek4ExactTraceWriter::hash_bytes(const void * data, size_t bytes) { + uint64_t hash = kFnvOffset; + if (data && bytes) hash_update(hash, data, bytes); + return hash_hex(hash); +} + +std::string DeepSeek4ExactTraceWriter::hash_token_ids( + const std::vector & tokens) { + uint64_t hash = kFnvOffset; + for (int32_t token : tokens) { + const uint32_t value = static_cast(token); + const uint8_t bytes[] = { + static_cast(value), + static_cast(value >> 8), + static_cast(value >> 16), + static_cast(value >> 24), + }; + hash_update(hash, bytes, sizeof(bytes)); + } + return hash_hex(hash); +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_exact_trace.h b/server/src/deepseek4/deepseek4_exact_trace.h new file mode 100644 index 000000000..873d08ac6 --- /dev/null +++ b/server/src/deepseek4/deepseek4_exact_trace.h @@ -0,0 +1,71 @@ +#pragma once + +#include "deepseek4_internal.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct GenerateRequest; + +class DeepSeek4ExactTraceWriter { +public: + static std::unique_ptr create_from_env( + const DeepSeek4Weights & weights); + + DeepSeek4ExactTraceWriter(const DeepSeek4ExactTraceWriter &) = delete; + DeepSeek4ExactTraceWriter & operator=(const DeepSeek4ExactTraceWriter &) = delete; + + bool begin_request(const GenerateRequest & request, bool restored, int cache_position); + void end_request(bool ok, int cache_position, const std::vector & tokens); + + bool wants_step(int position_begin, int n_tokens) const; + + void record_reset(int cache_position); + void record_step(int position_begin, int n_tokens, int cache_position, bool logits_present); + void record_layer( + ggml_backend_t backend, + const DeepSeek4Cache & cache, + const std::vector & hc_state, + int layer, + int position_begin, + int n_tokens, + bool hash_routed, + const std::vector & routing_ids, + const std::vector & routing_weights); + void record_capture( + int position_begin, + int n_tokens, + int capture_begin, + const std::vector & layer_ids, + const std::vector & rows); + void record_logits(int position, const std::vector & logits); + void record_snapshot(const char * kind, int slot, const DeepSeek4Cache & cache); + + static std::string hash_bytes(const void * data, size_t bytes); + static std::string hash_token_ids(const std::vector & tokens); + +private: + DeepSeek4ExactTraceWriter( + const DeepSeek4Weights & weights, + std::ofstream output, + int width); + + const DeepSeek4Weights & weights_; + std::ofstream output_; + int width_ = 0; + int request_index_ = -1; + int prompt_tokens_ = 0; + int snapshot_position_ = -1; + std::vector compressor_boundaries_; + + void write_relations(int position_begin, int position_end); + std::string cache_state_hash(const DeepSeek4Cache & cache) const; +}; + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 7014fafd4..ed986be07 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -9,6 +9,7 @@ // 6. MoE FFN (hash routing + top-k + shared expert + clamped SwiGLU) #include "deepseek4_internal.h" +#include "deepseek4_exact_trace.h" #include "deepseek4_hc_cuda.h" #include "deepseek4_roctx.h" #include "internal.h" @@ -239,7 +240,9 @@ static ggml_tensor * build_moe_ffn(ggml_context * ctx, const DeepSeek4Weights & w, const DeepSeek4Layer & L, int layer_idx, - int n_tokens); + int n_tokens, + ggml_tensor ** trace_selected = nullptr, + ggml_tensor ** trace_weights = nullptr); struct DeepSeek4CachedDecodeFfnGraph { const ggml_context * owner_ctx = nullptr; @@ -248,17 +251,23 @@ struct DeepSeek4CachedDecodeFfnGraph { int n_tokens = 0; int n_expert_used = 0; bool hash_routed = false; + bool trace_routing = false; StepGraph sg; ggml_tensor * hash_ids = nullptr; + ggml_tensor * route_ids = nullptr; + ggml_tensor * route_weights = nullptr; bool valid() const { return owner_ctx && backend && layer_idx >= 0 && n_tokens > 0 && sg.ctx && sg.gf && sg.alloc && sg.inp_embed && sg.hidden_states && - (!hash_routed || hash_ids); + (!hash_routed || hash_ids) && + (!trace_routing || (route_ids && route_weights)); } void free() { hash_ids = nullptr; + route_ids = nullptr; + route_weights = nullptr; step_graph_destroy(sg); owner_ctx = nullptr; backend = nullptr; @@ -266,6 +275,7 @@ struct DeepSeek4CachedDecodeFfnGraph { n_tokens = 0; n_expert_used = 0; hash_routed = false; + trace_routing = false; } }; @@ -424,7 +434,8 @@ static bool build_cached_decode_ffn_graph( const DeepSeek4Layer & L, int layer_idx, int n_tokens, - bool hash_routed) { + bool hash_routed, + bool trace_routing) { out.free(); const size_t ctx_size = 16 * 1024 * 1024; @@ -471,6 +482,10 @@ static bool build_cached_decode_ffn_graph( if (w.expert_weight_scale != 1.0f) { weights = ggml_scale(out.sg.ctx, weights, w.expert_weight_scale); } + if (trace_routing) { + out.route_ids = out.hash_ids; + out.route_weights = weights; + } ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); @@ -481,7 +496,10 @@ static bool build_cached_decode_ffn_graph( ffn_out = ggml_add(out.sg.ctx, shared_out, routed_out); } else { - ffn_out = build_moe_ffn(out.sg.ctx, ffn_normed, w, L, layer_idx, n_tokens); + ffn_out = build_moe_ffn( + out.sg.ctx, ffn_normed, w, L, layer_idx, n_tokens, + trace_routing ? &out.route_ids : nullptr, + trace_routing ? &out.route_weights : nullptr); } if (!ffn_out) { @@ -490,6 +508,12 @@ static bool build_cached_decode_ffn_graph( } out.sg.hidden_states = ffn_out; + if (trace_routing) { + ggml_set_output(out.route_ids); + ggml_set_output(out.route_weights); + ggml_build_forward_expand(out.sg.gf, out.route_ids); + ggml_build_forward_expand(out.sg.gf, out.route_weights); + } ggml_set_output(out.sg.hidden_states); ggml_build_forward_expand(out.sg.gf, out.sg.hidden_states); @@ -505,6 +529,7 @@ static bool build_cached_decode_ffn_graph( out.n_tokens = n_tokens; out.n_expert_used = ds4_effective_expert_count(w); out.hash_routed = hash_routed; + out.trace_routing = trace_routing; return true; } @@ -3095,7 +3120,9 @@ static ggml_tensor * build_moe_ffn( const DeepSeek4Weights & w, const DeepSeek4Layer & L, int layer_idx, - int n_tokens) { + int n_tokens, + ggml_tensor ** trace_selected, + ggml_tensor ** trace_weights) { const int n_embd = w.n_embd; int n_used = w.n_expert_used; @@ -3107,6 +3134,8 @@ static ggml_tensor * build_moe_ffn( routed_out = ggml_scale(ctx, cur, 0.0f); } else { Ds4MoeRouting routing = build_moe_routing(ctx, cur, w, L, n_tokens); + if (trace_selected) *trace_selected = routing.selected; + if (trace_weights) *trace_weights = routing.weights; n_used = (int) routing.selected->ne[0]; ggml_tensor * cur_3d = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens); ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, routing.selected); @@ -5409,6 +5438,9 @@ static bool eval_ds4_layer_range_hybrid_ffn( MoeHybridRoutingStats * routing_stats, std::vector & out, DeepSeek4StepTelemetry * telemetry, + std::vector * trace_routing_ids, + std::vector * trace_routing_weights, + bool * trace_hash_routed, const MoeHybridDeviceOutputs * device_outputs = nullptr) { const bool trace_prefill = ds4_env_flag("DFLASH_DS4_PREFILL_TRACE"); if (trace_prefill) { @@ -5637,6 +5669,9 @@ static bool eval_ds4_layer_range_hybrid_ffn( telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); } + if (trace_routing_ids) *trace_routing_ids = selected; + if (trace_routing_weights) *trace_routing_weights = weights; + if (trace_hash_routed) *trace_hash_routed = hash_routed; MoeHybridConfig cfg = make_ds4_moe_hybrid_config(w); cfg.n_expert_used = route_width; @@ -6606,6 +6641,10 @@ bool deepseek4_step_layer_range( MoeExpertComputeRuntime * expert_runtime, MoeHybridRoutingStats * routing_stats) { const auto step_t0 = Ds4TimingClock::now(); + DeepSeek4ExactTraceWriter * exact_trace = + verify_hooks ? verify_hooks->exact_trace : nullptr; + const bool exact_trace_step = + exact_trace && exact_trace->wants_step(kv_start, n_tokens); if (!deepseek4_cuda_hc_set_device(device)) { std::fprintf(stderr, @@ -6698,6 +6737,8 @@ bool deepseek4_step_layer_range( chunk_hooks.capture_layer_ids = verify_hooks->capture_layer_ids; chunk_hooks.capture_out = verify_hooks->capture_out ? &chunk_capture : nullptr; chunk_hooks.all_logits_out = verify_hooks->all_logits_out ? &chunk_logits : nullptr; + chunk_hooks.allow_fused_verify = verify_hooks->allow_fused_verify; + chunk_hooks.exact_trace = verify_hooks->exact_trace; chunk_hooks_ptr = &chunk_hooks; } if (!deepseek4_step_layer_range( @@ -7045,6 +7086,8 @@ bool deepseek4_step_layer_range( const HcLayerWeightsCpu & hc_lw = hc_layer_weights_range[(size_t)il]; const int ratio = (int)w.compress_ratios[il]; bool hash_routed = false; + std::vector trace_routing_ids; + std::vector trace_routing_weights; const ggml_tensor * attn_in_backend = nullptr; const ggml_tensor * ffn_in_backend = nullptr; const ggml_tensor * attn_post_backend = nullptr; @@ -7562,6 +7605,9 @@ bool deepseek4_step_layer_range( token_ids, hash_routing_tables_range[(size_t)il], *moe_hybrid, expert_runtime, routing_stats, ffn_out_host, telemetry, + exact_trace_step ? &trace_routing_ids : nullptr, + exact_trace_step ? &trace_routing_weights : nullptr, + exact_trace_step ? &hash_routed : nullptr, ffn_device_join ? &owner_outputs : nullptr)) { std::fprintf(stderr, "[deepseek4-moe-tp] layer-range FFN failed layer %d\n", @@ -7595,9 +7641,12 @@ bool deepseek4_step_layer_range( cached.layer_idx != il || cached.n_tokens != n_tokens || cached.n_expert_used != n_expert_used || - cached.hash_routed != hash_routed) { + cached.hash_routed != hash_routed || + cached.trace_routing != (exact_trace != nullptr)) { const auto ffn_build_t0 = Ds4TimingClock::now(); - if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { + if (!build_cached_decode_ffn_graph( + cached, backend, w, L, il, n_tokens, hash_routed, + exact_trace != nullptr)) { std::fprintf(stderr, "[deepseek4] cached ffn graph alloc failed layer %d\n", il); return false; } @@ -7630,6 +7679,16 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->ffn_read_us += ds4_elapsed_us(ffn_read_t0, Ds4TimingClock::now()); } + if (exact_trace_step && cached.route_ids && cached.route_weights) { + trace_routing_ids.resize(ggml_nelements(cached.route_ids)); + trace_routing_weights.resize(ggml_nelements(cached.route_weights)); + ggml_backend_tensor_get( + cached.route_ids, trace_routing_ids.data(), 0, + sizeof(int32_t) * trace_routing_ids.size()); + ggml_backend_tensor_get( + cached.route_weights, trace_routing_weights.data(), 0, + sizeof(float) * trace_routing_weights.size()); + } } if (use_backend_prefill_hc) { @@ -7713,12 +7772,17 @@ bool deepseek4_step_layer_range( } if ((use_backend_prefill_hc || use_backend_decode_hc_graph || use_backend_decode_hc_direct) && - hc_state_backend && capture_requested(il)) { + hc_state_backend && (capture_requested(il) || exact_trace_step)) { ggml_backend_tensor_get( hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); capture_hc_layer(il, hc_state.data()); } + if (exact_trace_step) { + exact_trace->record_layer( + backend, cache, hc_state, il, kv_start, n_tokens, + hash_routed, trace_routing_ids, trace_routing_weights); + } } } diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index b0e80ec00..a34a4c66d 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -38,6 +38,7 @@ struct MoeHybridConfig; struct MoeHybridRoutingStats; struct MoeExpertComputeRuntime; class MoeHybridStreamEngine; +class DeepSeek4ExactTraceWriter; struct DeepSeek4StepTelemetry { uint64_t total_us = 0; @@ -389,6 +390,12 @@ struct Ds4VerifyHooks { std::vector * all_logits_out = nullptr; // [n_vocab * n_tokens] std::vector * argmax_out = nullptr; // [n_tokens], optional GPU result bool prefer_argmax_only = false; // skip logits D2H when available + // Prefill uses the capture fields too, but must never enter the + // intentionally approximate fused-verification path. + bool allow_fused_verify = true; + // Default-null production diagnostic sink. The sink owns all expensive + // readbacks and position filtering; normal execution does not construct it. + DeepSeek4ExactTraceWriter * exact_trace = nullptr; }; bool deepseek4_step_layer_range( diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 8649c37fc..600a3f786 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -14,6 +14,7 @@ #include "common/layer_split_runtime.h" #include "common/layer_split_utils.h" #include "deepseek4/deepseek4_dspark.h" +#include "deepseek4/deepseek4_exact_trace.h" #include #include @@ -80,6 +81,12 @@ static double elapsed_ms(TestClock::time_point t0, TestClock::time_point t1) { return std::chrono::duration(t1 - t0).count(); } +static void test_exact_trace_hash_is_deterministic() { + const char input[] = "abc"; + TEST_ASSERT(DeepSeek4ExactTraceWriter::hash_bytes(input, 3) == + "e71fa2190541574b"); +} + struct DeepSeek4FixtureOptions { bool include_vocab_size = true; uint32_t vocab_size = 128; @@ -3654,6 +3661,7 @@ int main() { } test_compressor_pooling_correctness(backend); + test_exact_trace_hash_is_deterministic(); test_swiglu_ds4_cpu_correctness(backend); test_moe_routing_correctness(backend); test_rmsnorm_correctness(backend); From 7e17623b3132e0a36995c4fa49c1021121b7a6b2 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Sun, 9 Aug 2026 20:37:32 +0530 Subject: [PATCH 2/9] test(dflash): emit complete baseline traces --- .../src/deepseek4/deepseek4_exact_trace.cpp | 22 +-------- server/tests/test_deepseek4_unit.cpp | 45 +++++++++++++++++++ 2 files changed, 47 insertions(+), 20 deletions(-) diff --git a/server/src/deepseek4/deepseek4_exact_trace.cpp b/server/src/deepseek4/deepseek4_exact_trace.cpp index ced5e0dae..df2e56f70 100644 --- a/server/src/deepseek4/deepseek4_exact_trace.cpp +++ b/server/src/deepseek4/deepseek4_exact_trace.cpp @@ -219,26 +219,8 @@ void DeepSeek4ExactTraceWriter::end_request( } bool DeepSeek4ExactTraceWriter::wants_step(int position_begin, int n_tokens) const { - if (request_index_ < 0 || n_tokens <= 0) return false; - const int position_end = position_begin + n_tokens; - if (position_begin < 16 || position_end >= std::max(0, prompt_tokens_ - 8)) return true; - if (snapshot_position_ >= 0 && - near_boundary(position_begin, position_end, snapshot_position_)) { - return true; - } - if (weights_.n_swa > 0 && near_boundary(position_begin, position_end, weights_.n_swa)) { - return true; - } - for (int ratio : compressor_boundaries_) { - const int first = ratio; - const int last = prompt_tokens_ / ratio * ratio; - if (near_boundary(position_begin, position_end, first) || - (last > first && near_boundary(position_begin, position_end, last))) { - return true; - } - } - const int capture_begin = std::max(0, prompt_tokens_ - weights_.n_swa); - return near_boundary(position_begin, position_end, capture_begin); + return request_index_ >= 0 && position_begin >= 0 && n_tokens > 0 && + position_begin + n_tokens <= prompt_tokens_; } void DeepSeek4ExactTraceWriter::record_reset(int cache_position) { diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 600a3f786..11e37290a 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -106,6 +107,49 @@ static std::string make_temp_gguf_path(const char * prefix) { return std::string(path) + "-" + prefix + ".gguf"; } +static void test_exact_trace_serializes_interior_prompt_step() { + char path[] = "/tmp/deepseek4-exact-trace-XXXXXX"; + const int fd = mkstemp(path); + TEST_ASSERT(fd >= 0); + if (fd < 0) return; + close(fd); + unlink(path); + + const char * old_path = std::getenv("DFLASH_DS4_EXACT_TRACE_PATH"); + const char * old_width = std::getenv("DFLASH_DS4_EXACT_TRACE_Q"); + const bool had_path = old_path != nullptr; + const bool had_width = old_width != nullptr; + const std::string old_path_value = old_path ? old_path : ""; + const std::string old_width_value = old_width ? old_width : ""; + setenv("DFLASH_DS4_EXACT_TRACE_PATH", path, 1); + setenv("DFLASH_DS4_EXACT_TRACE_Q", "4", 1); + + DeepSeek4Weights weights; + auto writer = DeepSeek4ExactTraceWriter::create_from_env(weights); + GenerateRequest request; + request.prompt.resize(2048, 1); + request.n_gen = 16; + TEST_ASSERT(writer != nullptr); + if (writer) { + TEST_ASSERT(writer->begin_request(request, false, 0)); + TEST_ASSERT(writer->wants_step(512, 4)); + writer->record_step(512, 4, 516, false); + writer.reset(); + + std::ifstream input(path); + const std::string trace( + (std::istreambuf_iterator(input)), std::istreambuf_iterator()); + TEST_ASSERT(trace.find( + "\"position_begin\":512,\"position_end\":516") != std::string::npos); + } + + if (had_path) setenv("DFLASH_DS4_EXACT_TRACE_PATH", old_path_value.c_str(), 1); + else unsetenv("DFLASH_DS4_EXACT_TRACE_PATH"); + if (had_width) setenv("DFLASH_DS4_EXACT_TRACE_Q", old_width_value.c_str(), 1); + else unsetenv("DFLASH_DS4_EXACT_TRACE_Q"); + unlink(path); +} + static std::string write_deepseek4_loader_fixture(const DeepSeek4FixtureOptions & opts) { gguf_context * g = gguf_init_empty(); gguf_set_val_str(g, "general.architecture", "deepseek4"); @@ -3662,6 +3706,7 @@ int main() { test_compressor_pooling_correctness(backend); test_exact_trace_hash_is_deterministic(); + test_exact_trace_serializes_interior_prompt_step(); test_swiglu_ds4_cpu_correctness(backend); test_moe_routing_correctness(backend); test_rmsnorm_correctness(backend); From 9784d882ce2c9dd2a1635d4cf20de643803a0c7d Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Sun, 9 Aug 2026 21:01:14 +0530 Subject: [PATCH 3/9] test(dflash): align baseline trace steps --- .../src/deepseek4/deepseek4_exact_trace.cpp | 38 +++++++++++++------ server/tests/test_deepseek4_unit.cpp | 8 ++++ 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/server/src/deepseek4/deepseek4_exact_trace.cpp b/server/src/deepseek4/deepseek4_exact_trace.cpp index df2e56f70..e923411ef 100644 --- a/server/src/deepseek4/deepseek4_exact_trace.cpp +++ b/server/src/deepseek4/deepseek4_exact_trace.cpp @@ -292,18 +292,32 @@ void DeepSeek4ExactTraceWriter::record_step( int cache_position, bool logits_present) { if (!wants_step(position_begin, n_tokens)) return; - const int position_end = position_begin + n_tokens; - output_ << "{\"schema\":\"" << kSchema - << "\",\"type\":\"step\",\"request\":" << request_index_ - << ",\"position_begin\":" << position_begin - << ",\"position_end\":" << position_end - << ",\"position\":" << position_end - << ",\"width\":" << n_tokens - << ",\"cache_position\":" << cache_position - << ",\"logits_present\":" << (logits_present ? "true" : "false") - << ",\"relations\":"; - write_relations(position_begin, position_end); - output_ << "}\n"; + const int final_position = position_begin + n_tokens; + for (int offset = 0; offset < n_tokens;) { + const int chunk_begin = position_begin + offset; + int chunk = n_tokens - offset; + for (int ratio : compressor_boundaries_) { + int position_mod = chunk_begin % ratio; + if (position_mod < 0) position_mod += ratio; + chunk = std::min(chunk, ratio - position_mod); + } + const int chunk_end = chunk_begin + chunk; + const int chunk_cache_position = chunk_end == final_position + ? cache_position : chunk_end; + output_ << "{\"schema\":\"" << kSchema + << "\",\"type\":\"step\",\"request\":" << request_index_ + << ",\"position_begin\":" << chunk_begin + << ",\"position_end\":" << chunk_end + << ",\"position\":" << chunk_end + << ",\"width\":" << chunk + << ",\"cache_position\":" << chunk_cache_position + << ",\"logits_present\":" + << (logits_present && chunk_end == final_position ? "true" : "false") + << ",\"relations\":"; + write_relations(chunk_begin, chunk_end); + output_ << "}\n"; + offset += chunk; + } } void DeepSeek4ExactTraceWriter::record_layer( diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 11e37290a..e806ab248 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -125,6 +125,7 @@ static void test_exact_trace_serializes_interior_prompt_step() { setenv("DFLASH_DS4_EXACT_TRACE_Q", "4", 1); DeepSeek4Weights weights; + weights.compress_ratios = {4, 128}; auto writer = DeepSeek4ExactTraceWriter::create_from_env(weights); GenerateRequest request; request.prompt.resize(2048, 1); @@ -132,6 +133,7 @@ static void test_exact_trace_serializes_interior_prompt_step() { TEST_ASSERT(writer != nullptr); if (writer) { TEST_ASSERT(writer->begin_request(request, false, 0)); + writer->record_step(3, 3, 6, false); TEST_ASSERT(writer->wants_step(512, 4)); writer->record_step(512, 4, 516, false); writer.reset(); @@ -139,6 +141,12 @@ static void test_exact_trace_serializes_interior_prompt_step() { std::ifstream input(path); const std::string trace( (std::istreambuf_iterator(input)), std::istreambuf_iterator()); + TEST_ASSERT(trace.find( + "\"position_begin\":3,\"position_end\":4") != std::string::npos); + TEST_ASSERT(trace.find( + "\"position_begin\":4,\"position_end\":6") != std::string::npos); + TEST_ASSERT(trace.find( + "\"position_begin\":3,\"position_end\":6") == std::string::npos); TEST_ASSERT(trace.find( "\"position_begin\":512,\"position_end\":516") != std::string::npos); } From bfba0a8e2c2cff231aaf8539b9dd816970070a35 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Sun, 9 Aug 2026 22:23:18 +0530 Subject: [PATCH 4/9] test(dflash): add exact differential comparator --- harness/ds4_exact_diff.py | 1154 ++++++++++++++++++++++++++ harness/ds4_exact_diff_schema.md | 94 +++ harness/tests/test_ds4_exact_diff.py | 466 +++++++++++ 3 files changed, 1714 insertions(+) create mode 100644 harness/ds4_exact_diff.py create mode 100644 harness/ds4_exact_diff_schema.md create mode 100644 harness/tests/test_ds4_exact_diff.py diff --git a/harness/ds4_exact_diff.py b/harness/ds4_exact_diff.py new file mode 100644 index 000000000..6b2b45efc --- /dev/null +++ b/harness/ds4_exact_diff.py @@ -0,0 +1,1154 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import signal +import subprocess +import sys +import time +import urllib.error +import urllib.request +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +SCHEMA = "lucebox.ds4.exact-diff/v1" +TOLERANCES = { + "routing_weights": {"atol": 1e-6, "rtol": 1e-6}, + "capture_rows": {"atol": 1e-5, "rtol": 1e-5}, + "final_logits": {"atol": 1e-4, "rtol": 1e-4}, +} +WIDTHS = (1, 2, 3, 4) +PROFILES = ("reset", "snapshot") + + +class TraceError(ValueError): + pass + + +@dataclass(frozen=True) +class Mismatch: + profile: str + width: int + request: int | None + layer: int | None + token_position: int | None + field: str + oracle: Any + candidate: Any + detail: str + + def to_json(self) -> dict[str, Any]: + return { + "status": "mismatch", + "profile": self.profile, + "width": self.width, + "request": self.request, + "layer": self.layer, + "token_position": self.token_position, + "field": self.field, + "oracle": self.oracle, + "candidate": self.candidate, + "detail": self.detail, + } + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def reject_nonfinite_constant(value: str) -> None: + raise TraceError(f"non-finite JSON constant: {value}") + + +def load_jsonl(path: Path) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, 1): + if not line.strip(): + continue + try: + record = json.loads(line, parse_constant=reject_nonfinite_constant) + except (json.JSONDecodeError, TraceError) as exc: + raise TraceError(f"{path}:{line_number}: {exc}") from exc + if not isinstance(record, dict): + raise TraceError(f"{path}:{line_number}: record is not an object") + if record.get("schema") != SCHEMA: + raise TraceError(f"{path}:{line_number}: missing or unknown schema") + records.append(record) + if not records: + raise TraceError(f"{path}: empty trace") + validate_records(path, records) + return records + + +def require(record: dict[str, Any], fields: Iterable[str], context: str) -> None: + missing = [field for field in fields if field not in record] + if missing: + raise TraceError(f"{context}: missing fields: {', '.join(missing)}") + + +def validate_finite(value: Any, context: str) -> None: + if isinstance(value, float) and not math.isfinite(value): + raise TraceError(f"{context}: non-finite float") + if isinstance(value, list): + for index, item in enumerate(value): + validate_finite(item, f"{context}[{index}]") + elif isinstance(value, dict): + for key, item in value.items(): + validate_finite(item, f"{context}.{key}") + + +def validate_float_list(value: Any, context: str) -> None: + if not isinstance(value, list): + raise TraceError(f"{context}: expected a list") + for index, item in enumerate(value): + if isinstance(item, bool) or not isinstance(item, (int, float)): + raise TraceError(f"{context}[{index}]: expected a finite number") + if not math.isfinite(float(item)): + raise TraceError(f"{context}[{index}]: non-finite float") + + +def validate_int_list(value: Any, context: str) -> None: + if not isinstance(value, list): + raise TraceError(f"{context}: expected a list") + for index, item in enumerate(value): + if isinstance(item, bool) or not isinstance(item, int): + raise TraceError(f"{context}[{index}]: expected an integer") + if not -(2**31) <= item < 2**31: + raise TraceError(f"{context}[{index}]: integer is outside int32") + + +def validate_records(path: Path, records: list[dict[str, Any]]) -> None: + require(records[0], ("type", "profile", "width"), f"{path}:1") + if records[0]["type"] != "manifest": + raise TraceError(f"{path}: first record must be manifest") + require( + records[0], + ( + "revision", + "binary_sha256", + "target_sha256", + "drafter_sha256", + "prompt_bytes_sha256", + "request_config", + "tolerances", + ), + f"{path}:manifest", + ) + tolerances = records[0]["tolerances"] + if not isinstance(tolerances, dict): + raise TraceError(f"{path}:manifest:tolerances must be an object") + for name in TOLERANCES: + if name not in tolerances or not isinstance(tolerances[name], dict): + raise TraceError(f"{path}:manifest:tolerances.{name} is missing") + require(tolerances[name], ("atol", "rtol"), f"{path}:manifest:{name}") + for field in ("atol", "rtol"): + value = tolerances[name][field] + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(float(value)) + or value < 0 + ): + raise TraceError( + f"{path}:manifest:tolerances.{name}.{field} must be finite and non-negative" + ) + last_position: dict[int, int] = {} + for index, record in enumerate(records, 1): + record_type = record.get("type") + context = f"{path}:{index}:{record_type}" + validate_finite(record, context) + if record.get("non_finite") is True: + raise TraceError(f"{context}: producer reported non-finite state") + if record_type == "model_config": + require(record, ("n_swa", "compressor_boundaries"), context) + elif record_type == "request_start": + require( + record, + ( + "request", + "prompt_token_hash", + "prompt_token_ids", + "prompt_tokens", + "width", + "restored", + "cache_position", + "n_gen", + "snap_slot", + "snap_pos", + "temperature", + "top_p", + "top_k", + "seed", + ), + context, + ) + validate_int_list(record["prompt_token_ids"], f"{context}:prompt_token_ids") + if len(record["prompt_token_ids"]) != record["prompt_tokens"]: + raise TraceError(f"{context}: prompt token count does not match token IDs") + elif record_type == "layer": + require( + record, + ( + "request", + "layer", + "position_begin", + "position_end", + "routing", + "hc_token_hashes", + "raw_kv_hash", + "raw_kv_bytes", + "compressed_kv_hash", + "compressed_kv_bytes", + "compressor_state", + "indexer_state", + ), + context, + ) + routing = record["routing"] + require(routing, ("mode", "ids", "weights"), f"{context}:routing") + if not routing["ids"] or len(routing["ids"]) != len(routing["weights"]): + raise TraceError(f"{context}: routing IDs/weights are missing or misaligned") + validate_float_list(routing["weights"], f"{context}:routing.weights") + if not record["hc_token_hashes"] or not record["hc_token_hashes"][0]: + raise TraceError(f"{context}: HC token hash is missing") + request = int(record["request"]) + position = int(record["position_end"]) + previous = last_position.get(request, -1) + if position < previous: + raise TraceError(f"{context}: stale position {position} follows {previous}") + last_position[request] = position + elif record_type == "logits": + require(record, ("request", "position", "values"), context) + validate_float_list(record["values"], f"{context}:values") + elif record_type == "capture": + require( + record, + ( + "request", + "position_begin", + "position_end", + "layer_ids", + "row_hash", + "total_values", + "rows", + ), + context, + ) + validate_float_list(record["rows"], f"{context}:rows") + elif record_type in {"snapshot_save", "snapshot_restore"}: + require(record, ("request", "slot", "cache_position", "state_hash"), context) + elif record_type == "tokens": + require(record, ("request", "token_ids"), context) + elif record_type == "request_end": + require(record, ("request", "ok", "cache_position"), context) + elif record_type == "reset": + require(record, ("request", "cache_position"), context) + elif record_type == "step": + require( + record, + ( + "request", + "position_begin", + "position_end", + "position", + "width", + "cache_position", + "logits_present", + "relations", + ), + context, + ) + elif record_type != "manifest": + raise TraceError(f"{context}: unknown record type") + + +def close_enough(a: float, b: float, tolerance: dict[str, float]) -> bool: + if not math.isfinite(a) or not math.isfinite(b): + return False + limit = tolerance["atol"] + tolerance["rtol"] * max(abs(a), abs(b)) + return abs(a - b) <= limit + + +def event_index( + records: list[dict[str, Any]], event_type: str +) -> dict[tuple[Any, ...], dict[str, Any]]: + indexed: dict[tuple[Any, ...], dict[str, Any]] = {} + for record in records: + if record.get("type") != event_type: + continue + if event_type == "layer": + key = (record["request"], record["position_end"], record["layer"]) + elif event_type in {"logits", "step"}: + key = (record["request"], record.get("position", record.get("position_end"))) + elif event_type == "capture": + key = (record["request"], record["position_end"]) + elif event_type in {"tokens", "request_end", "request_start", "reset"}: + key = (record["request"],) + else: + key = (record["request"], record.get("slot"), record.get("cache_position")) + if key in indexed: + raise TraceError(f"duplicate {event_type} key {key}") + indexed[key] = record + return indexed + + +def first_float_mismatch( + oracle: list[float], candidate: list[float], tolerance: dict[str, float] +) -> tuple[int, float | None, float | None] | None: + if len(oracle) != len(candidate): + return min(len(oracle), len(candidate)), None, None + for index, (left, right) in enumerate(zip(oracle, candidate, strict=True)): + if not close_enough(float(left), float(right), tolerance): + return index, float(left), float(right) + return None + + +def compare_manifest( + profile: str, + width: int, + oracle: dict[str, Any], + candidate: dict[str, Any], +) -> Mismatch | None: + exact_fields = ( + "revision", + "binary_sha256", + "target_sha256", + "drafter_sha256", + "prompt_bytes_sha256", + "tolerances", + ) + for field in exact_fields: + if oracle[field] != candidate[field]: + return Mismatch( + profile, + width, + None, + None, + None, + field, + oracle[field], + candidate[field], + "manifest differs", + ) + oracle_cfg = dict(oracle["request_config"]) + candidate_cfg = dict(candidate["request_config"]) + for allowed in ("prefill_width", "exact_bands", "port"): + oracle_cfg.pop(allowed, None) + candidate_cfg.pop(allowed, None) + if oracle_cfg != candidate_cfg: + return Mismatch( + profile, + width, + None, + None, + None, + "request_config", + oracle_cfg, + candidate_cfg, + "non-width request configuration differs", + ) + return None + + +def compare_profile( + profile: str, width: int, oracle: list[dict[str, Any]], candidate: list[dict[str, Any]] +) -> Mismatch | None: + mismatch = compare_manifest(profile, width, oracle[0], candidate[0]) + if mismatch: + return mismatch + for event_type, fields in ( + ( + "request_start", + ( + "prompt_token_hash", + "prompt_token_ids", + "prompt_tokens", + "restored", + "cache_position", + "n_gen", + "snap_slot", + "snap_pos", + "temperature", + "top_p", + "top_k", + "seed", + ), + ), + ("reset", ("cache_position",)), + ): + refs = event_index(oracle, event_type) + cands = event_index(candidate, event_type) + if refs.keys() != cands.keys(): + return Mismatch( + profile, + width, + None, + None, + None, + event_type, + sorted(refs), + sorted(cands), + f"{event_type} keys differ", + ) + for key in sorted(cands): + for field in fields: + if refs[key][field] != cands[key][field]: + return Mismatch( + profile, + width, + int(key[0]), + None, + None, + f"{event_type}.{field}", + refs[key][field], + cands[key][field], + f"{event_type} differs", + ) + + oracle_steps = event_index(oracle, "step") + candidate_steps = event_index(candidate, "step") + starts = event_index(candidate, "request_start") + for label, steps in (("oracle", oracle_steps), ("candidate", candidate_steps)): + maximum_width = 1 if label == "oracle" else width + for request_key, start in sorted(starts.items()): + request = int(request_key[0]) + request_steps = sorted( + (record for key, record in steps.items() if key[0] == request), + key=lambda record: int(record["position_end"]), + ) + cursor = int(start["cache_position"]) + final_position = int(start["prompt_tokens"]) + if not request_steps: + if cursor == final_position: + continue + return Mismatch( + profile, + width, + request, + None, + None, + "step.coverage", + (cursor, final_position), + "missing", + f"{label} has no step records for request", + ) + for record in request_steps: + begin = int(record["position_begin"]) + end = int(record["position_end"]) + step_width = int(record["width"]) + if ( + begin != cursor + or end != int(record["position"]) + or end != int(record["cache_position"]) + or end - begin != step_width + or not 1 <= step_width <= maximum_width + or end > final_position + ): + return Mismatch( + profile, + width, + request, + None, + end, + "step.coverage", + {"position_begin": cursor, "final_position": final_position}, + record, + f"{label} step coverage is discontinuous or malformed", + ) + cursor = end + if cursor != final_position: + return Mismatch( + profile, + width, + request, + None, + cursor, + "step.coverage", + final_position, + cursor, + f"{label} steps do not cover the full request", + ) + for key in sorted(candidate_steps): + cand = candidate_steps[key] + ref = oracle_steps.get(key) + if ref is None: + return Mismatch( + profile, + width, + int(key[0]), + None, + int(key[1]), + "step", + "present", + "missing oracle alignment", + "q=1 has no matching step", + ) + for field in ("cache_position",): + if ref[field] != cand[field]: + return Mismatch( + profile, + width, + int(key[0]), + None, + int(key[1]), + f"step.{field}", + ref[field], + cand[field], + "step lifecycle differs", + ) + start = starts[(key[0],)] + final_position = start["prompt_tokens"] + snapshot_position = int(start["snap_pos"]) + expected_logits = key[1] == final_position or ( + snapshot_position >= 0 and key[1] == snapshot_position + ) + if cand["logits_present"] is not expected_logits: + return Mismatch( + profile, + width, + int(key[0]), + None, + int(key[1]), + "step.logits_present", + expected_logits, + cand["logits_present"], + "step readout policy differs", + ) + + oracle_layers = event_index(oracle, "layer") + candidate_layers = event_index(candidate, "layer") + if not candidate_layers: + return Mismatch( + profile, + width, + None, + None, + None, + "layer", + "present", + "missing", + "candidate has no layer records", + ) + selected_positions = set(candidate_steps) + expected_layer_keys = {key for key in oracle_layers if (key[0], key[1]) in selected_positions} + if candidate_layers.keys() != expected_layer_keys: + return Mismatch( + profile, + width, + None, + None, + None, + "layer.keys", + sorted(expected_layer_keys), + sorted(candidate_layers), + "layer keys are incomplete at selected positions", + ) + tolerances = oracle[0]["tolerances"] + for key in sorted(expected_layer_keys): + cand = candidate_layers[key] + ref = oracle_layers.get(key) + request, _, layer = key + position = int(cand["token_position"]) + if ref is None: + return Mismatch( + profile, + width, + request, + layer, + position, + "layer", + "present", + "missing oracle alignment", + "q=1 has no matching layer position", + ) + cand_routing = cand["routing"] + ref_routing = ref["routing"] + for field in ("mode", "ids"): + if ref_routing[field] != cand_routing[field]: + return Mismatch( + profile, + width, + request, + layer, + position, + f"routing.{field}", + ref_routing[field], + cand_routing[field], + "routing differs", + ) + float_diff = first_float_mismatch( + ref_routing["weights"], cand_routing["weights"], tolerances["routing_weights"] + ) + if float_diff: + index, left, right = float_diff + return Mismatch( + profile, + width, + request, + layer, + position, + f"routing.weights[{index}]", + left, + right, + "routing weight exceeds tolerance", + ) + exact_fields = ( + "hc_token_hashes", + "raw_kv_hash", + "raw_kv_bytes", + "compressed_kv_hash", + "compressed_kv_bytes", + "compressor_state", + "indexer_state", + ) + for field in exact_fields: + if ref[field] != cand[field]: + return Mismatch( + profile, + width, + request, + layer, + position, + field, + ref[field], + cand[field], + "exact state differs", + ) + for event_type, tolerance_name, value_field in ( + ("capture", "capture_rows", "rows"), + ("logits", "final_logits", "values"), + ): + refs = event_index(oracle, event_type) + cands = event_index(candidate, event_type) + if refs.keys() != cands.keys(): + return Mismatch( + profile, + width, + None, + None, + None, + event_type, + sorted(refs), + sorted(cands), + f"{event_type} positions differ", + ) + for key in sorted(cands): + cand = cands[key] + ref = refs.get(key) + request = int(key[0]) + position = int(key[-1]) + if ref is None: + return Mismatch( + profile, + width, + request, + None, + position, + event_type, + "present", + "missing oracle alignment", + f"q=1 has no matching {event_type}", + ) + if event_type == "capture": + for field in ("layer_ids", "total_values"): + if ref[field] != cand[field]: + return Mismatch( + profile, + width, + request, + None, + position, + f"capture.{field}", + ref[field], + cand[field], + "capture shape differs", + ) + if len(cand["rows"]) != cand["total_values"]: + return Mismatch( + profile, + width, + request, + None, + position, + "capture.rows", + cand["total_values"], + len(cand["rows"]), + "capture row is incomplete", + ) + float_diff = first_float_mismatch( + ref[value_field], cand[value_field], tolerances[tolerance_name] + ) + if float_diff: + index, left, right = float_diff + return Mismatch( + profile, + width, + request, + None, + position, + f"{event_type}.{value_field}[{index}]", + left, + right, + f"{event_type} exceeds tolerance", + ) + for event_type, fields in ( + ("tokens", ("token_ids",)), + ("request_end", ("ok", "cache_position")), + ("snapshot_save", ("slot", "cache_position", "state_hash")), + ("snapshot_restore", ("slot", "cache_position", "state_hash")), + ): + refs = event_index(oracle, event_type) + cands = event_index(candidate, event_type) + if refs.keys() != cands.keys(): + return Mismatch( + profile, + width, + None, + None, + None, + event_type, + sorted(refs), + sorted(cands), + f"{event_type} keys differ", + ) + for key in sorted(cands): + if key not in refs: + return Mismatch( + profile, + width, + int(key[0]), + None, + None, + event_type, + "present", + "missing oracle alignment", + f"q=1 has no matching {event_type}", + ) + for field in fields: + if refs[key][field] != cands[key][field]: + return Mismatch( + profile, + width, + int(key[0]), + None, + refs[key].get("cache_position"), + f"{event_type}.{field}", + refs[key][field], + cands[key][field], + f"{event_type} differs", + ) + return None + + +def matrix_failure(field: str, expected: Any, observed: Any, detail: str) -> Mismatch: + return Mismatch("matrix", 0, None, None, None, field, expected, observed, detail) + + +def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mismatch | None: + model_configs: list[dict[str, Any]] = [] + for (profile, width), records in sorted(traces.items()): + configs = [record for record in records if record.get("type") == "model_config"] + if len(configs) != 1: + return matrix_failure( + f"{profile}.q{width}.model_config", + 1, + len(configs), + "each production trace must declare model boundaries exactly once", + ) + model_configs.append(configs[0]) + canonical_config = { + "n_swa": model_configs[0]["n_swa"], + "compressor_boundaries": model_configs[0]["compressor_boundaries"], + } + for config in model_configs[1:]: + observed = { + "n_swa": config["n_swa"], + "compressor_boundaries": config["compressor_boundaries"], + } + if observed != canonical_config: + return matrix_failure( + "model_config", + canonical_config, + observed, + "model boundary configuration differs across traces", + ) + + reset_candidates = [traces[("reset", width)] for width in WIDTHS[1:]] + for width, records in zip(WIDTHS[1:], reset_candidates, strict=True): + ordinary = any( + record.get("type") == "step" + and record.get("width") == width + and "ordinary" in record.get("relations", []) + for record in records + ) + if not ordinary: + return matrix_failure( + f"ordinary.q{width}", True, False, "no ordinary full-width production step" + ) + + relations = { + relation + for records in traces.values() + for record in records + if record.get("type") == "step" + for relation in record.get("relations", []) + } + required_relations = { + "tail_width_1", + "tail_width_2", + "tail_width_3", + "compressor_before", + "compressor_on", + "compressor_after", + "swa_before", + "swa_on", + "swa_after", + "dspark_capture_window", + } + for boundary in canonical_config["compressor_boundaries"]: + required_relations.update( + f"compressor_{boundary}_{relation}" for relation in ("before", "on", "after") + ) + if canonical_config["n_swa"] > 0: + required_relations.update( + f"swa_{canonical_config['n_swa']}_{relation}" for relation in ("before", "on", "after") + ) + missing_relations = sorted(required_relations - relations) + if missing_relations: + return matrix_failure( + "relations", + sorted(required_relations), + sorted(relations), + f"missing matrix relations: {', '.join(missing_relations)}", + ) + + routing_modes = { + record["routing"]["mode"] + for records in traces.values() + for record in records + if record.get("type") == "layer" + } + if routing_modes != {"hash", "learned"}: + return matrix_failure( + "routing_modes", + ["hash", "learned"], + sorted(routing_modes), + "both routing implementations must be observed", + ) + + for (profile, width), records in sorted(traces.items()): + starts = [record for record in records if record.get("type") == "request_start"] + if len(starts) < 2: + return matrix_failure( + f"{profile}.q{width}.repeated_requests", + 2, + len(starts), + "same request was not observed twice", + ) + tokens = [record for record in records if record.get("type") == "tokens"] + if len(tokens) < 2 or any(len(record["token_ids"]) < 2 for record in tokens): + return matrix_failure( + f"{profile}.q{width}.continuations", + "two requests with multiple tokens", + [len(record["token_ids"]) for record in tokens], + "multiple continuation tokens are required", + ) + if not any(record.get("type") == "logits" for record in records): + return matrix_failure( + f"{profile}.q{width}.logits", True, False, "final logits are missing" + ) + if ( + profile == "reset" + and len([record for record in records if record.get("type") == "reset"]) < 2 + ): + return matrix_failure(f"{profile}.q{width}.reset", 2, 0, "reset events are missing") + if profile == "snapshot": + kinds = {record.get("type") for record in records} + for kind in ("snapshot_save", "snapshot_restore"): + if kind not in kinds: + return matrix_failure( + f"{profile}.q{width}.{kind}", True, False, f"{kind} event is missing" + ) + drafter = records[0]["drafter_sha256"] + if drafter != "none": + captures_by_request: dict[int, set[int]] = {} + for record in records: + if record.get("type") == "capture": + captures_by_request.setdefault(record["request"], set()).add( + record["position_end"] + ) + for start in starts: + request = start["request"] + final_position = start["prompt_tokens"] + expected_ends = {final_position - 3, final_position} + observed = captures_by_request.get(request, set()) + if not expected_ends.issubset(observed): + return matrix_failure( + f"{profile}.q{width}.capture.{request}", + sorted(expected_ends), + sorted(observed), + "both ends of the four-token DSpark capture window are required", + ) + return None + + +def compare_trace_dir(trace_dir: Path) -> int: + summaries: list[dict[str, Any]] = [] + traces: dict[tuple[str, int], list[dict[str, Any]]] = {} + for profile in PROFILES: + profile_dir = trace_dir / profile + oracle_path = profile_dir / "q1.jsonl" + oracle = load_jsonl(oracle_path) + traces[(profile, 1)] = oracle + for width in WIDTHS[1:]: + candidate = load_jsonl(profile_dir / f"q{width}.jsonl") + traces[(profile, width)] = candidate + mismatch = compare_profile(profile, width, oracle, candidate) + if mismatch: + print(json.dumps(mismatch.to_json(), sort_keys=True)) + return 1 + summaries.append({"profile": profile, "width": width, "status": "match"}) + mismatch = validate_matrix(traces) + if mismatch: + print(json.dumps(mismatch.to_json(), sort_keys=True)) + return 1 + print(json.dumps({"status": "match", "comparisons": summaries}, sort_keys=True)) + return 0 + + +def http_json(url: str, payload: dict[str, Any] | None = None, timeout: float = 10.0) -> Any: + data = None if payload is None else json.dumps(payload).encode("utf-8") + request = urllib.request.Request(url, data=data) + if data is not None: + request.add_header("Content-Type", "application/json") + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read().decode("utf-8")) + + +def wait_ready(port: int, process: subprocess.Popen[bytes], timeout: float) -> None: + deadline = time.monotonic() + timeout + url = f"http://127.0.0.1:{port}/health" + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f"server exited before ready with status {process.returncode}") + try: + http_json(url, timeout=2.0) + return + except (OSError, urllib.error.URLError, json.JSONDecodeError): + time.sleep(0.25) + raise TimeoutError(f"server did not become ready within {timeout:.1f}s") + + +def stop_process(process: subprocess.Popen[bytes]) -> None: + if process.poll() is not None: + return + process.send_signal(signal.SIGTERM) + try: + process.wait(timeout=10.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5.0) + + +def git_revision(binary: Path, explicit: str | None) -> str: + if explicit: + return explicit + completed = subprocess.run( + ["git", "-C", str(binary.parent), "rev-parse", "HEAD"], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + raise RuntimeError("cannot infer revision; pass --revision") + return completed.stdout.strip() + + +def write_manifest( + path: Path, + profile: str, + width: int, + args: argparse.Namespace, + revision: str, + hashes: dict[str, str], + port: int, +) -> None: + manifest = { + "schema": SCHEMA, + "type": "manifest", + "profile": profile, + "width": width, + "revision": revision, + "binary_sha256": hashes["binary"], + "target_sha256": hashes["target"], + "drafter_sha256": hashes["draft"], + "prompt_bytes_sha256": hashes["prompt"], + "request_config": { + "prefill_mode": "exact", + "prefill_width": width, + "exact_bands": width > 1, + "generated_tokens": args.generated_tokens, + "temperature": 0.0, + "seed": args.seed, + "target_device": args.target_device, + "profile": profile, + "port": port, + "server_args": args.server_arg, + }, + "tolerances": TOLERANCES, + } + path.write_text(json.dumps(manifest, sort_keys=True) + "\n", encoding="utf-8") + + +def run_one( + profile: str, + width: int, + args: argparse.Namespace, + revision: str, + hashes: dict[str, str], + prompt: str, +) -> None: + profile_dir = args.output_dir / profile + profile_dir.mkdir(parents=True, exist_ok=True) + trace_path = profile_dir / f"q{width}.jsonl" + log_path = profile_dir / f"q{width}.server.log" + port = args.port_base + (0 if profile == "reset" else 100) + width + write_manifest(trace_path, profile, width, args, revision, hashes, port) + env = os.environ.copy() + env.update( + { + "DFLASH_DS4_EXACT_TRACE_PATH": str(trace_path), + "DFLASH_DS4_EXACT_TRACE_Q": str(width), + "DFLASH_DS4_EXACT_PREFILL_BANDS": "1" if width > 1 else "0", + "DFLASH_DS4_FUSED_VERIFY": "0", + "DFLASH_DS4_FUSED_DECODE": "0", + "DFLASH_DS4_FUSED_HYBRID_DECODE": "0", + } + ) + if args.draft: + env.update( + { + "DFLASH_DS4_SPEC": "1", + "DFLASH_DS4_DRAFT": str(args.draft), + "DFLASH_DS4_SPEC_Q": "4", + } + ) + cache_slots = "0" if profile == "reset" else "2" + command = [ + str(args.binary), + str(args.target), + "--host", + "127.0.0.1", + "--port", + str(port), + "--target-device", + args.target_device, + "--ds4-prefill", + "exact", + "--chunk", + str(width), + "--prefix-cache-slots", + "0", + "--prefill-cache-slots", + cache_slots, + "--prefill-compression", + "off", + *args.server_arg, + ] + request = { + "model": "local", + "messages": [{"role": "user", "content": prompt}], + "temperature": 0.0, + "seed": args.seed, + "max_tokens": args.generated_tokens, + "stream": False, + } + with log_path.open("wb") as log: + process = subprocess.Popen(command, env=env, stdout=log, stderr=subprocess.STDOUT) + try: + wait_ready(port, process, args.startup_timeout) + for _ in range(2): + http_json( + f"http://127.0.0.1:{port}/v1/chat/completions", + request, + timeout=args.request_timeout, + ) + finally: + stop_process(process) + load_jsonl(trace_path) + + +def run_matrix(args: argparse.Namespace) -> int: + for field in ("binary", "target", "prompt"): + path = getattr(args, field) + if not path.is_file(): + raise FileNotFoundError(f"{field} does not exist: {path}") + if args.draft and not args.draft.is_file(): + raise FileNotFoundError(f"draft does not exist: {args.draft}") + args.output_dir.mkdir(parents=True, exist_ok=True) + revision = git_revision(args.binary, args.revision) + hashes = { + "binary": sha256_file(args.binary), + "target": sha256_file(args.target), + "draft": sha256_file(args.draft) if args.draft else "none", + "prompt": sha256_file(args.prompt), + } + prompt = args.prompt.read_text(encoding="utf-8") + for profile in PROFILES: + for width in WIDTHS: + run_one(profile, width, args, revision, hashes, prompt) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="DS4 production exact-differential harness") + subparsers = parser.add_subparsers(dest="command", required=True) + run = subparsers.add_parser("run", help="generate q=1/2/3/4 production traces") + run.add_argument("--binary", type=Path, required=True) + run.add_argument("--target", type=Path, required=True) + run.add_argument("--draft", type=Path) + run.add_argument("--prompt", type=Path, required=True) + run.add_argument("--output-dir", type=Path, required=True) + run.add_argument("--revision") + run.add_argument("--target-device", default="hip:0") + run.add_argument("--generated-tokens", type=int, default=16) + run.add_argument("--seed", type=int, default=1) + run.add_argument("--port-base", type=int, default=18100) + run.add_argument("--startup-timeout", type=float, default=900.0) + run.add_argument("--request-timeout", type=float, default=900.0) + run.add_argument("--server-arg", action="append", default=[]) + compare = subparsers.add_parser("compare", help="compare all q traces against q=1") + compare.add_argument("--trace-dir", type=Path, required=True) + return parser + + +def main() -> int: + args = build_parser().parse_args() + try: + if args.command == "run": + return run_matrix(args) + return compare_trace_dir(args.trace_dir) + except ( + FileNotFoundError, + RuntimeError, + TimeoutError, + TraceError, + urllib.error.URLError, + ) as exc: + print(json.dumps({"status": "error", "detail": str(exc)}, sort_keys=True), file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/ds4_exact_diff_schema.md b/harness/ds4_exact_diff_schema.md new file mode 100644 index 000000000..8f21c0ba5 --- /dev/null +++ b/harness/ds4_exact_diff_schema.md @@ -0,0 +1,94 @@ +# DS4 production exact-differential trace + +The production trace is disabled unless `DFLASH_DS4_EXACT_TRACE_PATH` names an +output file. Enabling it adds synchronous device readbacks and is diagnostic +only. It does not alter model arithmetic, routing, cache mutation, capture +selection, or ROCTX ranges. + +## Schema and comparison contract + +Every JSONL record uses schema `lucebox.ds4.exact-diff/v1`. The runner writes a +`manifest` record with the revision, binary SHA-256, target and drafter SHA-256, +prompt-byte SHA-256, fixed request configuration, exact width, and tolerances. +The production backend appends request, layer, cache, capture, logits, snapshot, +reset, continuation-token, and completion records. Each `request_start` carries +the bounded production prompt token-ID vector as well as its little-endian +signed-int32 hash, so the packaged token pinner consumes the real producer +schema rather than reconstructing tokenization. +The backend also appends one `model_config` record containing the exact SWA +window and every distinct compressor boundary. Boundary relations include the +numeric boundary, so matrix validation proves before/on/after coverage for each +configured value rather than accepting one representative boundary. + +The comparator always uses q=1 as the oracle and stops at the first mismatch. +Generated text is never an oracle. It aligns records by profile, request, +committed position, layer, token position, and field. Routing IDs, token IDs, +cache positions, counts, and state hashes are exact. Floating values use: + +| Field | absolute tolerance | relative tolerance | +|---|---:|---:| +| routing weights | `1e-6` | `1e-6` | +| DSpark capture rows | `1e-5` | `1e-5` | +| final logits | `1e-4` | `1e-4` | + +The inclusive rule is `abs(a-b) <= atol + rtol * max(abs(a), abs(b))`. Any NaN +or infinity is a hard failure, including two matching infinities. HC, raw KV, +compressed KV, attention-compressor state, indexer-compressor state, and +indexer KV use exact deterministic byte hashes. Their trace records also carry +byte counts. DSpark capture records contain every value in every captured row; +the comparator applies the declared tolerance to the complete row and rejects +missing capture positions. + +## Correctness matrix + +The `reset` profile runs the same request twice with both prefix caches disabled. +The `snapshot` profile runs it twice with the full-prompt snapshot cache enabled, +so the second request must restore the first request's state. For every profile, +q=1, q=2, q=3, and q=4 use the same prompt bytes and request configuration. +To cover all three tail widths with one pinned prompt, its production token count +must satisfy `N % 12 == 11`; q=2, q=3, and q=4 then end with widths 1, 2, and 3. +The comparator rejects a trace set that does not actually contain those rows. + +The production position filter records: + +- ordinary q=2, q=3, and q=4 steps; +- observed tail widths 1, 2, and 3; +- steps before, on, and after every model compressor ratio; +- steps before, on, and after the model SWA boundary; +- every hash-routed and learned-router layer at each selected step; +- the first and repeated/reset request; +- snapshot save and restore events; +- the beginning and end of the DSpark final capture window; +- final cache position, declared-tolerance logits, greedy token, and all + continuation token IDs. + +The comparator fails if a required matrix relation has no trace evidence. At +every selected step it also requires the complete oracle layer-key set, checks +the committed cache position and exact-band readout policy, and requires both +ends of each request's four-token DSpark final-capture window. + +## Commands + +Generate all traces (two profiles, four exact widths): + +```bash +python3 harness/ds4_exact_diff.py run \ + --binary server/build-hip/dflash_server \ + --target /models/target.gguf \ + --draft /models/dspark.gguf \ + --prompt /data/prompt.txt \ + --target-device hip:0 \ + --output-dir /tmp/ds4-exact-diff +``` + +Compare q=2, q=3, and q=4 against q=1: + +```bash +python3 harness/ds4_exact_diff.py compare \ + --trace-dir /tmp/ds4-exact-diff +``` + +The run command hashes its inputs before launching the server, forces greedy +sampling, keeps exact attention, disables approximate/fused verification, and +uses `--chunk q` with the exact-band flag disabled for q=1 and enabled for +q=2..4. Raw traces and server logs belong outside Git. diff --git a/harness/tests/test_ds4_exact_diff.py b/harness/tests/test_ds4_exact_diff.py new file mode 100644 index 000000000..0f947b895 --- /dev/null +++ b/harness/tests/test_ds4_exact_diff.py @@ -0,0 +1,466 @@ +from __future__ import annotations + +import importlib.util +import json +import math +import sys +from pathlib import Path + +import pytest + +MODULE_PATH = Path(__file__).parents[1] / "ds4_exact_diff.py" +SPEC = importlib.util.spec_from_file_location("ds4_exact_diff", MODULE_PATH) +assert SPEC and SPEC.loader +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +def manifest(width: int = 1) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "manifest", + "profile": "reset", + "width": width, + "revision": "a" * 40, + "binary_sha256": "b" * 64, + "target_sha256": "c" * 64, + "drafter_sha256": "d" * 64, + "prompt_bytes_sha256": "e" * 64, + "request_config": { + "prefill_width": width, + "exact_bands": width > 1, + "port": 18000 + width, + "temperature": 0, + }, + "tolerances": MODULE.TOLERANCES, + } + + +def layer(width: int = 1, position: int = 4) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "layer", + "request": 0, + "layer": 2, + "position_begin": position - width, + "position_end": position, + "token_position": position - 1, + "routing": {"mode": "learned", "ids": [1, 2], "weights": [0.6, 0.4]}, + "hc_token_hashes": ["0123456789abcdef"], + "raw_kv_hash": "1" * 16, + "raw_kv_bytes": 16, + "compressed_kv_hash": None, + "compressed_kv_bytes": 0, + "compressor_state": {"kv": None, "score": None, "n_comp": 0}, + "indexer_state": {"kv": None, "score": None, "n_comp": 0}, + } + + +def request_start(prompt_tokens: int = 4) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "request_start", + "request": 0, + "width": 1, + "restored": False, + "cache_position": 0, + "prompt_token_hash": "f" * 16, + "prompt_token_ids": list(range(prompt_tokens)), + "prompt_tokens": prompt_tokens, + "n_gen": 16, + "snap_slot": -1, + "snap_pos": -1, + "temperature": 0.0, + "top_p": 1.0, + "top_k": 0, + "seed": 1, + } + + +def step(width: int = 1, position: int = 4, logits_present: bool = True) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "step", + "request": 0, + "position_begin": position - width, + "position_end": position, + "position": position, + "width": width, + "cache_position": position, + "logits_present": logits_present, + "relations": [], + } + + +def capture(position: int = 4) -> dict[str, object]: + values = [float(index) for index in range(96)] + return { + "schema": MODULE.SCHEMA, + "type": "capture", + "request": 0, + "position_begin": position - 1, + "position_end": position, + "layer_ids": [1, 2, 3], + "row_hash": "a" * 16, + "total_values": len(values), + "rows": values, + } + + +def request_end(cache_position: int = 4) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "request_end", + "request": 0, + "ok": True, + "cache_position": cache_position, + } + + +def complete_trace( + width: int = 1, + *, + cache_position: int = 0, + restored: bool = False, + prompt_tokens: int = 4, + compressor_ratio: int | None = None, + snapshot_position: int = -1, +) -> list[dict[str, object]]: + start = request_start(prompt_tokens) + start["width"] = width + start["cache_position"] = cache_position + start["restored"] = restored + start["snap_pos"] = snapshot_position + start["snap_slot"] = 0 if snapshot_position >= 0 else -1 + records = [manifest(width), start] + position = cache_position + while position < prompt_tokens: + step_width = min(width, prompt_tokens - position) + if compressor_ratio is not None: + step_width = min(step_width, compressor_ratio - position % compressor_ratio) + position += step_width + relations = [] + if ( + snapshot_position >= 0 + and position - step_width <= snapshot_position + 4 + and position >= snapshot_position - 4 + ): + relations.append("snapshot_boundary") + step_record = step( + step_width, + position, + logits_present=position in {prompt_tokens, snapshot_position}, + ) + step_record["relations"] = relations + records.extend( + [ + step_record, + layer(step_width, position), + ] + ) + records.append(request_end(prompt_tokens)) + return records + + +def write_trace(path: Path, records: list[dict[str, object]]) -> None: + path.write_text( + "".join(json.dumps(record, sort_keys=True) + "\n" for record in records), + encoding="utf-8", + ) + + +def test_deterministic_sha256(tmp_path: Path) -> None: + path = tmp_path / "input.bin" + path.write_bytes(b"abc") + assert MODULE.sha256_file(path) == ( + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad" + ) + + +def test_serialization_and_comparison_match(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = complete_trace(2) + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + loaded_oracle = MODULE.load_jsonl(oracle_path) + loaded_candidate = MODULE.load_jsonl(candidate_path) + assert MODULE.compare_profile("reset", 2, loaded_oracle, loaded_candidate) is None + + +def test_large_prompt_continuous_producer_contract_matches(tmp_path: Path) -> None: + oracle_path = tmp_path / "q1-2k.jsonl" + candidate_path = tmp_path / "q4-2k.jsonl" + write_trace(oracle_path, complete_trace(prompt_tokens=2048)) + write_trace(candidate_path, complete_trace(4, prompt_tokens=2048)) + assert ( + MODULE.compare_profile( + "reset", 4, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + is None + ) + + +def test_q3_compressor_split_producer_contract_matches(tmp_path: Path) -> None: + oracle_path = tmp_path / "q1-ratio4.jsonl" + candidate_path = tmp_path / "q3-ratio4.jsonl" + write_trace(oracle_path, complete_trace(prompt_tokens=2048, compressor_ratio=4)) + write_trace( + candidate_path, + complete_trace(3, prompt_tokens=2048, compressor_ratio=4), + ) + candidate = MODULE.load_jsonl(candidate_path) + candidate_steps = MODULE.event_index(candidate, "step") + assert (0, 4) in candidate_steps + assert candidate_steps[(0, 4)]["position_begin"] == 3 + assert MODULE.compare_profile("reset", 3, MODULE.load_jsonl(oracle_path), candidate) is None + + +def test_snapshot_readout_uses_exact_endpoint_not_near_relation(tmp_path: Path) -> None: + oracle_path = tmp_path / "q1-snapshot.jsonl" + candidate_path = tmp_path / "q3-snapshot.jsonl" + write_trace( + oracle_path, + complete_trace(prompt_tokens=8, compressor_ratio=4, snapshot_position=4), + ) + write_trace( + candidate_path, + complete_trace(3, prompt_tokens=8, compressor_ratio=4, snapshot_position=4), + ) + candidate = MODULE.load_jsonl(candidate_path) + candidate_steps = MODULE.event_index(candidate, "step") + assert "snapshot_boundary" in candidate_steps[(0, 3)]["relations"] + assert candidate_steps[(0, 3)]["logits_present"] is False + assert candidate_steps[(0, 4)]["logits_present"] is True + assert MODULE.compare_profile("snapshot", 3, MODULE.load_jsonl(oracle_path), candidate) is None + + +def test_missing_field_fails(tmp_path: Path) -> None: + broken = layer() + del broken["raw_kv_hash"] + path = tmp_path / "broken.jsonl" + write_trace(path, [manifest(), request_start(), broken, request_end()]) + with pytest.raises(MODULE.TraceError, match="raw_kv_hash"): + MODULE.load_jsonl(path) + + +def test_stale_position_fails(tmp_path: Path) -> None: + first = layer(position=8) + second = layer(position=4) + second["layer"] = 3 + path = tmp_path / "stale.jsonl" + write_trace(path, [manifest(), request_start(), first, second, request_end()]) + with pytest.raises(MODULE.TraceError, match="stale position"): + MODULE.load_jsonl(path) + + +@pytest.mark.parametrize("constant", ["NaN", "Infinity", "-Infinity"]) +def test_nonfinite_json_fails(tmp_path: Path, constant: str) -> None: + path = tmp_path / "nonfinite.jsonl" + text = json.dumps(manifest()) + "\n" + text += json.dumps(request_start()) + "\n" + text += json.dumps(layer()).replace("0.6", constant, 1) + "\n" + path.write_text(text, encoding="utf-8") + with pytest.raises(MODULE.TraceError, match="non-finite"): + MODULE.load_jsonl(path) + + +def test_nonfinite_comparison_never_matches() -> None: + tolerance = {"atol": 1.0, "rtol": 1.0} + assert not MODULE.close_enough(math.nan, math.nan, tolerance) + assert not MODULE.close_enough(math.inf, math.inf, tolerance) + + +def test_producer_nonfinite_flag_fails(tmp_path: Path) -> None: + broken = layer() + broken["non_finite"] = True + path = tmp_path / "producer-nonfinite.jsonl" + write_trace(path, [manifest(), request_start(), broken, request_end()]) + with pytest.raises(MODULE.TraceError, match="producer reported non-finite"): + MODULE.load_jsonl(path) + + +def test_null_numeric_value_fails(tmp_path: Path) -> None: + broken = layer() + broken["routing"] = {"mode": "learned", "ids": [1, 2], "weights": [0.6, None]} + path = tmp_path / "null-number.jsonl" + write_trace(path, [manifest(), request_start(), broken, request_end()]) + with pytest.raises(MODULE.TraceError, match="expected a finite number"): + MODULE.load_jsonl(path) + + +def test_tolerance_boundary_is_inclusive() -> None: + tolerance = {"atol": 0.125, "rtol": 0.0} + assert MODULE.close_enough(1.0, 1.125, tolerance) + assert not MODULE.close_enough(1.0, 1.1250001, tolerance) + + +def test_first_mismatch_reports_layer_token_and_field(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = complete_trace(2) + changed = next( + record + for record in candidate + if record.get("type") == "layer" and record.get("position_end") == 4 + ) + changed["routing"] = {"mode": "learned", "ids": [1, 3], "weights": [0.6, 0.4]} + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.layer == 2 + assert mismatch.token_position == 3 + assert mismatch.field == "routing.ids" + + +def test_missing_selected_layer_fails(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = [ + record + for record in complete_trace(2) + if not (record.get("type") == "layer" and record.get("position_end") == 2) + ] + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.field == "layer.keys" + + +def test_partial_step_and_layer_omission_fails(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = [ + record + for record in complete_trace(2) + if not (record.get("type") in {"step", "layer"} and record.get("position_end") == 2) + ] + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.field == "step.coverage" + + +def test_step_cache_position_mismatch_fails(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = complete_trace(2) + changed_step = next( + record + for record in candidate + if record.get("type") == "step" and record.get("position_end") == 4 + ) + changed_step["cache_position"] = 3 + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.field == "step.coverage" + + +def test_step_readout_policy_mismatch_fails(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = complete_trace(2) + changed_step = next( + record + for record in candidate + if record.get("type") == "step" and record.get("position_end") == 2 + ) + changed_step["logits_present"] = True + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.field == "step.logits_present" + + +def test_restored_request_final_position_uses_full_prompt_length(tmp_path: Path) -> None: + oracle = complete_trace(cache_position=2, restored=True) + candidate = complete_trace(2, cache_position=2, restored=True) + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + assert ( + MODULE.compare_profile( + "snapshot", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + is None + ) + + +def test_capture_middle_value_mismatch_fails(tmp_path: Path) -> None: + oracle_capture = capture() + candidate_capture = capture() + candidate_capture["rows"][48] = 999.0 + oracle = complete_trace() + candidate = complete_trace(2) + oracle.insert(-1, oracle_capture) + candidate.insert(-1, candidate_capture) + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.field == "capture.rows[48]" + + +def test_missing_capture_position_fails(tmp_path: Path) -> None: + oracle = complete_trace() + candidate = complete_trace(2) + oracle.insert(-1, capture()) + oracle_path = tmp_path / "q1.jsonl" + candidate_path = tmp_path / "q2.jsonl" + write_trace(oracle_path, oracle) + write_trace(candidate_path, candidate) + mismatch = MODULE.compare_profile( + "reset", 2, MODULE.load_jsonl(oracle_path), MODULE.load_jsonl(candidate_path) + ) + assert mismatch is not None + assert mismatch.field == "capture" + + +def test_prompt_token_ids_are_required(tmp_path: Path) -> None: + broken_start = request_start() + del broken_start["prompt_token_ids"] + path = tmp_path / "missing-prompt-token-ids.jsonl" + write_trace(path, [manifest(), broken_start, step(), layer(), request_end()]) + with pytest.raises(MODULE.TraceError, match="prompt_token_ids"): + MODULE.load_jsonl(path) + + +def test_matrix_requires_model_config_per_trace() -> None: + traces = { + (profile, width): [manifest(width)] + for profile in MODULE.PROFILES + for width in MODULE.WIDTHS + } + mismatch = MODULE.validate_matrix(traces) + assert mismatch is not None + assert mismatch.field == "reset.q1.model_config" + assert mismatch.oracle == 1 + assert mismatch.candidate == 0 From f075f82131fcb26544f960d7f461f6db063c70ae Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Sun, 9 Aug 2026 22:39:08 +0530 Subject: [PATCH 5/9] docs(dflash): document exact trace dependency --- harness/ds4_exact_diff_schema.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/harness/ds4_exact_diff_schema.md b/harness/ds4_exact_diff_schema.md index 8f21c0ba5..0e5bb89db 100644 --- a/harness/ds4_exact_diff_schema.md +++ b/harness/ds4_exact_diff_schema.md @@ -71,6 +71,10 @@ ends of each request's four-token DSpark final-capture window. Generate all traces (two profiles, four exact widths): +The q=2 through q=4 runs require a server revision that implements +`DFLASH_DS4_EXACT_PREFILL_BANDS`. On a revision without exact bands, comparison +fails closed because the required step widths and tail coverage are absent. + ```bash python3 harness/ds4_exact_diff.py run \ --binary server/build-hip/dflash_server \ From da60f11c56d173faff5023d6d60b62a1fd11eb45 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Sun, 9 Aug 2026 23:27:00 +0530 Subject: [PATCH 6/9] fix(dflash): close exact trace oracle gaps --- harness/ds4_exact_diff.py | 203 +++++++++++++++--- harness/ds4_exact_diff_schema.md | 11 +- harness/tests/test_ds4_exact_diff.py | 79 ++++++- server/src/deepseek4/deepseek4_backend.cpp | 11 +- .../src/deepseek4/deepseek4_exact_trace.cpp | 22 +- server/src/deepseek4/deepseek4_exact_trace.h | 8 +- server/src/deepseek4/deepseek4_graph.cpp | 23 +- server/src/deepseek4/deepseek4_internal.h | 9 + server/tests/test_deepseek4_unit.cpp | 18 ++ 9 files changed, 338 insertions(+), 46 deletions(-) diff --git a/harness/ds4_exact_diff.py b/harness/ds4_exact_diff.py index 6b2b45efc..78f305e17 100644 --- a/harness/ds4_exact_diff.py +++ b/harness/ds4_exact_diff.py @@ -163,6 +163,8 @@ def validate_records(path: Path, records: list[dict[str, Any]]) -> None: raise TraceError( f"{path}:manifest:tolerances.{name}.{field} must be finite and non-negative" ) + if tolerances != TOLERANCES: + raise TraceError(f"{path}:manifest:tolerances must equal the built-in contract") last_position: dict[int, int] = {} for index, record in enumerate(records, 1): record_type = record.get("type") @@ -247,11 +249,33 @@ def validate_records(path: Path, records: list[dict[str, Any]]) -> None: ) validate_float_list(record["rows"], f"{context}:rows") elif record_type in {"snapshot_save", "snapshot_restore"}: - require(record, ("request", "slot", "cache_position", "state_hash"), context) + require( + record, + ( + "request", + "slot", + "cache_position", + "state_hash", + "last_logits_hash", + "last_logits_count", + "last_logits_position", + "spec_feature_hash", + "spec_feature_count", + ), + context, + ) + if record["last_logits_count"] <= 0: + raise TraceError(f"{context}: snapshot logits are missing") + if record["last_logits_position"] != record["cache_position"]: + raise TraceError(f"{context}: snapshot logits position is stale") + if record["spec_feature_count"] < 0: + raise TraceError(f"{context}: snapshot feature count is negative") elif record_type == "tokens": require(record, ("request", "token_ids"), context) elif record_type == "request_end": require(record, ("request", "ok", "cache_position"), context) + if record["ok"] is not True: + raise TraceError(f"{context}: request did not complete successfully") elif record_type == "reset": require(record, ("request", "cache_position"), context) elif record_type == "step": @@ -272,6 +296,15 @@ def validate_records(path: Path, records: list[dict[str, Any]]) -> None: elif record_type != "manifest": raise TraceError(f"{context}: unknown record type") + request_starts = [ + record["request"] for record in records if record.get("type") == "request_start" + ] + request_ends = [record["request"] for record in records if record.get("type") == "request_end"] + token_records = [record["request"] for record in records if record.get("type") == "tokens"] + for label, observed in (("request_end", request_ends), ("tokens", token_records)): + if sorted(observed) != sorted(request_starts): + raise TraceError(f"{path}:{label} coverage does not match request_start coverage") + def close_enough(a: float, b: float, tolerance: dict[str, float]) -> bool: if not math.isfinite(a) or not math.isfinite(b): @@ -705,8 +738,32 @@ def compare_profile( for event_type, fields in ( ("tokens", ("token_ids",)), ("request_end", ("ok", "cache_position")), - ("snapshot_save", ("slot", "cache_position", "state_hash")), - ("snapshot_restore", ("slot", "cache_position", "state_hash")), + ( + "snapshot_save", + ( + "slot", + "cache_position", + "state_hash", + "last_logits_hash", + "last_logits_count", + "last_logits_position", + "spec_feature_hash", + "spec_feature_count", + ), + ), + ( + "snapshot_restore", + ( + "slot", + "cache_position", + "state_hash", + "last_logits_hash", + "last_logits_count", + "last_logits_position", + "spec_feature_hash", + "spec_feature_count", + ), + ), ): refs = event_index(oracle, event_type) cands = event_index(candidate, event_type) @@ -755,6 +812,100 @@ def matrix_failure(field: str, expected: Any, observed: Any, detail: str) -> Mis return Mismatch("matrix", 0, None, None, None, field, expected, observed, detail) +SNAPSHOT_STATE_FIELDS = ( + "slot", + "cache_position", + "state_hash", + "last_logits_hash", + "last_logits_count", + "last_logits_position", + "spec_feature_hash", + "spec_feature_count", +) + + +def validate_snapshot_lifecycle( + profile: str, + width: int, + records: list[dict[str, Any]], + starts: list[dict[str, Any]], +) -> Mismatch | None: + saves = [record for record in records if record.get("type") == "snapshot_save"] + restores = [record for record in records if record.get("type") == "snapshot_restore"] + if len(saves) != 1 or len(restores) != 1: + return matrix_failure( + f"{profile}.q{width}.snapshot_count", + {"save": 1, "restore": 1}, + {"save": len(saves), "restore": len(restores)}, + "snapshot profile requires exactly one save and one restore", + ) + saved = saves[0] + restored = restores[0] + starts_by_request = {start["request"]: start for start in starts} + saved_start = starts_by_request.get(saved["request"]) + restored_start = starts_by_request.get(restored["request"]) + if saved_start is None or saved_start["restored"]: + return matrix_failure( + f"{profile}.q{width}.snapshot_save_request", + "fresh request", + saved["request"], + "snapshot save must belong to the fresh request", + ) + if restored_start is None or not restored_start["restored"]: + return matrix_failure( + f"{profile}.q{width}.snapshot_restore_request", + "restored request", + restored["request"], + "snapshot restore must belong to the restored request", + ) + for field in SNAPSHOT_STATE_FIELDS: + if saved[field] != restored[field]: + return matrix_failure( + f"{profile}.q{width}.snapshot.{field}", + saved[field], + restored[field], + "restored snapshot state differs from saved state", + ) + drafter = records[0]["drafter_sha256"] + if drafter != "none" and saved["spec_feature_count"] <= 0: + return matrix_failure( + f"{profile}.q{width}.snapshot.spec_feature_count", + "positive", + saved["spec_feature_count"], + "DSpark snapshot features are missing", + ) + return None + + +def validate_capture_coverage( + profile: str, + width: int, + records: list[dict[str, Any]], + starts: list[dict[str, Any]], +) -> Mismatch | None: + if records[0]["drafter_sha256"] == "none": + return None + captures_by_request: dict[int, set[int]] = {} + for record in records: + if record.get("type") == "capture": + captures_by_request.setdefault(record["request"], set()).add(record["position_end"]) + for start in starts: + if start["restored"]: + continue + request = start["request"] + final_position = start["prompt_tokens"] + expected_ends = {final_position - 3, final_position} + observed = captures_by_request.get(request, set()) + if not expected_ends.issubset(observed): + return matrix_failure( + f"{profile}.q{width}.capture.{request}", + sorted(expected_ends), + sorted(observed), + "both ends of the four-token DSpark capture window are required", + ) + return None + + def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mismatch | None: model_configs: list[dict[str, Any]] = [] for (profile, width), records in sorted(traces.items()): @@ -849,12 +1000,22 @@ def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mism for (profile, width), records in sorted(traces.items()): starts = [record for record in records if record.get("type") == "request_start"] - if len(starts) < 2: + if len(starts) != 2: return matrix_failure( f"{profile}.q{width}.repeated_requests", 2, len(starts), - "same request was not observed twice", + "same request must be observed exactly twice", + ) + ordered_starts = sorted(starts, key=lambda record: record["request"]) + expected_restored = [False, profile == "snapshot"] + observed_restored = [start["restored"] for start in ordered_starts] + if observed_restored != expected_restored: + return matrix_failure( + f"{profile}.q{width}.restored_requests", + expected_restored, + observed_restored, + "request restore lifecycle differs from the profile contract", ) tokens = [record for record in records if record.get("type") == "tokens"] if len(tokens) < 2 or any(len(record["token_ids"]) < 2 for record in tokens): @@ -874,32 +1035,12 @@ def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mism ): return matrix_failure(f"{profile}.q{width}.reset", 2, 0, "reset events are missing") if profile == "snapshot": - kinds = {record.get("type") for record in records} - for kind in ("snapshot_save", "snapshot_restore"): - if kind not in kinds: - return matrix_failure( - f"{profile}.q{width}.{kind}", True, False, f"{kind} event is missing" - ) - drafter = records[0]["drafter_sha256"] - if drafter != "none": - captures_by_request: dict[int, set[int]] = {} - for record in records: - if record.get("type") == "capture": - captures_by_request.setdefault(record["request"], set()).add( - record["position_end"] - ) - for start in starts: - request = start["request"] - final_position = start["prompt_tokens"] - expected_ends = {final_position - 3, final_position} - observed = captures_by_request.get(request, set()) - if not expected_ends.issubset(observed): - return matrix_failure( - f"{profile}.q{width}.capture.{request}", - sorted(expected_ends), - sorted(observed), - "both ends of the four-token DSpark capture window are required", - ) + snapshot_mismatch = validate_snapshot_lifecycle(profile, width, records, starts) + if snapshot_mismatch: + return snapshot_mismatch + capture_mismatch = validate_capture_coverage(profile, width, records, starts) + if capture_mismatch: + return capture_mismatch return None diff --git a/harness/ds4_exact_diff_schema.md b/harness/ds4_exact_diff_schema.md index 0e5bb89db..af677666b 100644 --- a/harness/ds4_exact_diff_schema.md +++ b/harness/ds4_exact_diff_schema.md @@ -13,7 +13,7 @@ prompt-byte SHA-256, fixed request configuration, exact width, and tolerances. The production backend appends request, layer, cache, capture, logits, snapshot, reset, continuation-token, and completion records. Each `request_start` carries the bounded production prompt token-ID vector as well as its little-endian -signed-int32 hash, so the packaged token pinner consumes the real producer +signed-int32 hash, so trace consumers can authenticate the real producer schema rather than reconstructing tokenization. The backend also appends one `model_config` record containing the exact SWA window and every distinct compressor boundary. Boundary relations include the @@ -32,7 +32,8 @@ cache positions, counts, and state hashes are exact. Floating values use: | final logits | `1e-4` | `1e-4` | The inclusive rule is `abs(a-b) <= atol + rtol * max(abs(a), abs(b))`. Any NaN -or infinity is a hard failure, including two matching infinities. HC, raw KV, +or infinity is a hard failure, including two matching infinities. Manifests +must use these exact built-in tolerances; trace inputs cannot weaken them. HC, raw KV, compressed KV, attention-compressor state, indexer-compressor state, and indexer KV use exact deterministic byte hashes. Their trace records also carry byte counts. DSpark capture records contain every value in every captured row; @@ -48,6 +49,10 @@ q=1, q=2, q=3, and q=4 use the same prompt bytes and request configuration. To cover all three tail widths with one pinned prompt, its production token count must satisfy `N % 12 == 11`; q=2, q=3, and q=4 then end with widths 1, 2, and 3. The comparator rejects a trace set that does not actually contain those rows. +Every request must end successfully and emit exactly one completion and token +record. A full-prompt restored request performs no prefill, so its DSpark state +is authenticated through the snapshot's cache, logits, and feature hashes +rather than through impossible duplicate capture rows. The production position filter records: @@ -65,7 +70,7 @@ The production position filter records: The comparator fails if a required matrix relation has no trace evidence. At every selected step it also requires the complete oracle layer-key set, checks the committed cache position and exact-band readout policy, and requires both -ends of each request's four-token DSpark final-capture window. +ends of each non-restored request's four-token DSpark final-capture window. ## Commands diff --git a/harness/tests/test_ds4_exact_diff.py b/harness/tests/test_ds4_exact_diff.py index 0f947b895..7eb6793af 100644 --- a/harness/tests/test_ds4_exact_diff.py +++ b/harness/tests/test_ds4_exact_diff.py @@ -118,6 +118,31 @@ def request_end(cache_position: int = 4) -> dict[str, object]: } +def tokens(request: int = 0) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "tokens", + "request": request, + "token_ids": [7, 8], + } + + +def snapshot(kind: str, request: int, *, spec_feature_count: int = 4) -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": kind, + "request": request, + "slot": 0, + "cache_position": 4, + "state_hash": "1" * 16, + "last_logits_hash": "2" * 16, + "last_logits_count": 8, + "last_logits_position": 4, + "spec_feature_hash": "3" * 16, + "spec_feature_count": spec_feature_count, + } + + def complete_trace( width: int = 1, *, @@ -159,7 +184,7 @@ def complete_trace( layer(step_width, position), ] ) - records.append(request_end(prompt_tokens)) + records.extend([tokens(), request_end(prompt_tokens)]) return records @@ -297,6 +322,31 @@ def test_tolerance_boundary_is_inclusive() -> None: assert not MODULE.close_enough(1.0, 1.1250001, tolerance) +def test_manifest_tolerances_cannot_weaken_contract(tmp_path: Path) -> None: + weakened = manifest() + weakened["tolerances"] = {name: {"atol": 1e9, "rtol": 1e9} for name in MODULE.TOLERANCES} + path = tmp_path / "weakened-tolerances.jsonl" + write_trace(path, [weakened, request_start(), step(), layer(), request_end()]) + with pytest.raises(MODULE.TraceError, match="built-in contract"): + MODULE.load_jsonl(path) + + +def test_failed_request_cannot_certify(tmp_path: Path) -> None: + failed = request_end() + failed["ok"] = False + path = tmp_path / "failed-request.jsonl" + write_trace(path, [manifest(), request_start(), step(), layer(), failed]) + with pytest.raises(MODULE.TraceError, match="did not complete successfully"): + MODULE.load_jsonl(path) + + +def test_missing_request_completion_cannot_certify(tmp_path: Path) -> None: + path = tmp_path / "missing-request-end.jsonl" + write_trace(path, [manifest(), request_start(), step(), layer()]) + with pytest.raises(MODULE.TraceError, match="request_end coverage"): + MODULE.load_jsonl(path) + + def test_first_mismatch_reports_layer_token_and_field(tmp_path: Path) -> None: oracle = complete_trace() candidate = complete_trace(2) @@ -410,6 +460,33 @@ def test_restored_request_final_position_uses_full_prompt_length(tmp_path: Path) ) +def test_snapshot_restore_requires_identical_auxiliary_state() -> None: + fresh = request_start() + restored = request_start() + restored["request"] = 1 + restored["restored"] = True + records = [ + manifest(), + snapshot("snapshot_save", 0), + snapshot("snapshot_restore", 1), + ] + assert MODULE.validate_snapshot_lifecycle("snapshot", 1, records, [fresh, restored]) is None + records[-1]["spec_feature_hash"] = "4" * 16 + mismatch = MODULE.validate_snapshot_lifecycle("snapshot", 1, records, [fresh, restored]) + assert mismatch is not None + assert mismatch.field == "snapshot.q1.snapshot.spec_feature_hash" + + +def test_restored_request_uses_snapshot_features_instead_of_prefill_capture() -> None: + fresh = request_start() + restored = request_start() + restored["request"] = 1 + restored["restored"] = True + restored["cache_position"] = restored["prompt_tokens"] + records = [manifest(), fresh, restored, capture(1), capture(4)] + assert MODULE.validate_capture_coverage("snapshot", 1, records, [fresh, restored]) is None + + def test_capture_middle_value_mismatch_fails(tmp_path: Path) -> None: oracle_capture = capture() candidate_capture = capture() diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index cccfc5cd0..425b99ec0 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1529,7 +1529,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, pos += n_tok; if (exact_trace_) { exact_trace_->record_step( - pos - n_tok, n_tok, cache_.cur_pos, true); + pos - n_tok, n_tok, cache_.cur_pos, !logits.empty()); if (i + n_tok == n_total || (save_snapshot && !snapshot_saved && pos == snap_pos)) { exact_trace_->record_logits(cache_.cur_pos, logits); @@ -1869,7 +1869,10 @@ bool DeepSeek4Backend::snapshot_save(int slot) { slot, snapshots_[slot].cur_pos, (double) (core_bytes + aux_bytes) / (1024.0 * 1024.0)); if (exact_trace_) { - exact_trace_->record_snapshot("snapshot_save", slot, cache_); + exact_trace_->record_snapshot( + "snapshot_save", slot, cache_, + snapshot_aux_[slot].last_logits, snapshots_[slot].cur_pos, + snapshot_aux_[slot].spec_feat_window); } return true; } @@ -1913,7 +1916,9 @@ bool DeepSeek4Backend::snapshot_restore(int slot) { spec_feat_window_ = std::move(restored_features); last_logits_pos_ = cache_.cur_pos; if (exact_trace_) { - exact_trace_->record_snapshot("snapshot_restore", slot, cache_); + exact_trace_->record_snapshot( + "snapshot_restore", slot, cache_, last_logits_, last_logits_pos_, + spec_feat_window_); } return true; } diff --git a/server/src/deepseek4/deepseek4_exact_trace.cpp b/server/src/deepseek4/deepseek4_exact_trace.cpp index e923411ef..af8d68185 100644 --- a/server/src/deepseek4/deepseek4_exact_trace.cpp +++ b/server/src/deepseek4/deepseek4_exact_trace.cpp @@ -518,12 +518,30 @@ std::string DeepSeek4ExactTraceWriter::cache_state_hash( void DeepSeek4ExactTraceWriter::record_snapshot( const char * kind, int slot, - const DeepSeek4Cache & cache) { + const DeepSeek4Cache & cache, + const std::vector & last_logits, + int last_logits_position, + const std::vector & spec_features) { + bool non_finite = false; + for (float value : last_logits) { + non_finite = non_finite || !std::isfinite(value); + } + for (float value : spec_features) { + non_finite = non_finite || !std::isfinite(value); + } output_ << "{\"schema\":\"" << kSchema << "\",\"type\":\"" << kind << "\",\"request\":" << request_index_ << ",\"slot\":" << slot << ",\"cache_position\":" << cache.cur_pos - << ",\"state_hash\":\"" << cache_state_hash(cache) << "\"}\n"; + << ",\"state_hash\":\"" << cache_state_hash(cache) << "\"" + << ",\"last_logits_hash\":\"" + << hash_bytes(last_logits.data(), last_logits.size() * sizeof(float)) << "\"" + << ",\"last_logits_count\":" << last_logits.size() + << ",\"last_logits_position\":" << last_logits_position + << ",\"spec_feature_hash\":\"" + << hash_bytes(spec_features.data(), spec_features.size() * sizeof(float)) << "\"" + << ",\"spec_feature_count\":" << spec_features.size() + << ",\"non_finite\":" << (non_finite ? "true" : "false") << "}\n"; } std::string DeepSeek4ExactTraceWriter::hash_bytes(const void * data, size_t bytes) { diff --git a/server/src/deepseek4/deepseek4_exact_trace.h b/server/src/deepseek4/deepseek4_exact_trace.h index 873d08ac6..4808b74a8 100644 --- a/server/src/deepseek4/deepseek4_exact_trace.h +++ b/server/src/deepseek4/deepseek4_exact_trace.h @@ -45,7 +45,13 @@ class DeepSeek4ExactTraceWriter { const std::vector & layer_ids, const std::vector & rows); void record_logits(int position, const std::vector & logits); - void record_snapshot(const char * kind, int slot, const DeepSeek4Cache & cache); + void record_snapshot( + const char * kind, + int slot, + const DeepSeek4Cache & cache, + const std::vector & last_logits, + int last_logits_position, + const std::vector & spec_features); static std::string hash_bytes(const void * data, size_t bytes); static std::string hash_token_ids(const std::vector & tokens); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index ed986be07..5f69d6808 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -6621,6 +6621,20 @@ static bool initialize_layer_range_cache( runtime.owns_output = owns_output; return true; } +bool deepseek4_should_attempt_fused_verify( + int n_tokens, + const Ds4VerifyHooks * verify_hooks, + bool owner_topology_supported, + bool full_layer_range, + bool has_logits_output, + bool gpu_backend, + bool fused_verify_enabled) { + return owner_topology_supported && n_tokens >= 2 && n_tokens <= 4 && + verify_hooks && verify_hooks->allow_fused_verify && + full_layer_range && has_logits_output && gpu_backend && + fused_verify_enabled; +} + bool deepseek4_step_layer_range( ggml_backend_t backend, int device, @@ -6663,11 +6677,10 @@ bool deepseek4_step_layer_range( moe_hybrid->materialized_cold_experts && moe_hybrid->cold_backend_kind == MoeHybridColdBackend::Gpu && moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend; - const bool fused_verify_candidate = - (!moe_hybrid || fused_hybrid_ready) && - n_tokens >= 2 && n_tokens <= 4 && verify_hooks && - layer_begin == 0 && is_last_shard && out_logits && - ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled(); + const bool fused_verify_candidate = deepseek4_should_attempt_fused_verify( + n_tokens, verify_hooks, !moe_hybrid || fused_hybrid_ready, + layer_begin == 0 && is_last_shard, out_logits != nullptr, + ds4_backend_is_gpu(backend), ds4_fused_verify_enabled()); const bool heterogeneous_sparse_prefill = moe_hybrid && cache.prefill_mode == PrefillAttentionMode::Sparse && n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a34a4c66d..bc28fe519 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -398,6 +398,15 @@ struct Ds4VerifyHooks { DeepSeek4ExactTraceWriter * exact_trace = nullptr; }; +bool deepseek4_should_attempt_fused_verify( + int n_tokens, + const Ds4VerifyHooks * verify_hooks, + bool owner_topology_supported, + bool full_layer_range, + bool has_logits_output, + bool gpu_backend, + bool fused_verify_enabled); + bool deepseek4_step_layer_range( ggml_backend_t backend, int device, diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index e806ab248..bcba288b7 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -136,6 +136,11 @@ static void test_exact_trace_serializes_interior_prompt_step() { writer->record_step(3, 3, 6, false); TEST_ASSERT(writer->wants_step(512, 4)); writer->record_step(512, 4, 516, false); + DeepSeek4Cache cache; + cache.cur_pos = 6; + writer->record_snapshot( + "snapshot_save", 0, cache, {1.0f, 2.0f}, 6, + {3.0f, 4.0f, 5.0f}); writer.reset(); std::ifstream input(path); @@ -149,6 +154,9 @@ static void test_exact_trace_serializes_interior_prompt_step() { "\"position_begin\":3,\"position_end\":6") == std::string::npos); TEST_ASSERT(trace.find( "\"position_begin\":512,\"position_end\":516") != std::string::npos); + TEST_ASSERT(trace.find("\"last_logits_count\":2") != std::string::npos); + TEST_ASSERT(trace.find("\"last_logits_position\":6") != std::string::npos); + TEST_ASSERT(trace.find("\"spec_feature_count\":3") != std::string::npos); } if (had_path) setenv("DFLASH_DS4_EXACT_TRACE_PATH", old_path_value.c_str(), 1); @@ -158,6 +166,15 @@ static void test_exact_trace_serializes_interior_prompt_step() { unlink(path); } +static void test_exact_trace_disables_fused_verify() { + Ds4VerifyHooks hooks; + TEST_ASSERT(deepseek4_should_attempt_fused_verify( + 4, &hooks, true, true, true, true, true)); + hooks.allow_fused_verify = false; + TEST_ASSERT(!deepseek4_should_attempt_fused_verify( + 4, &hooks, true, true, true, true, true)); +} + static std::string write_deepseek4_loader_fixture(const DeepSeek4FixtureOptions & opts) { gguf_context * g = gguf_init_empty(); gguf_set_val_str(g, "general.architecture", "deepseek4"); @@ -3715,6 +3732,7 @@ int main() { test_compressor_pooling_correctness(backend); test_exact_trace_hash_is_deterministic(); test_exact_trace_serializes_interior_prompt_step(); + test_exact_trace_disables_fused_verify(); test_swiglu_ds4_cpu_correctness(backend); test_moe_routing_correctness(backend); test_rmsnorm_correctness(backend); From 2ed0b65682b1d21b067611fbd86143bcda691042 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Mon, 10 Aug 2026 01:23:08 +0530 Subject: [PATCH 7/9] fix(dflash): verify repeated trace lifecycle --- harness/ds4_exact_diff.py | 121 ++++++++++++++++++++++++++- harness/ds4_exact_diff_schema.md | 4 +- harness/tests/test_ds4_exact_diff.py | 45 +++++++++- 3 files changed, 166 insertions(+), 4 deletions(-) diff --git a/harness/ds4_exact_diff.py b/harness/ds4_exact_diff.py index 78f305e17..044e57dde 100644 --- a/harness/ds4_exact_diff.py +++ b/harness/ds4_exact_diff.py @@ -823,6 +823,102 @@ def matrix_failure(field: str, expected: Any, observed: Any, detail: str) -> Mis "spec_feature_count", ) +REPEATED_REQUEST_FIELDS = ( + "width", + "prompt_token_hash", + "prompt_token_ids", + "prompt_tokens", + "n_gen", + "snap_slot", + "snap_pos", + "temperature", + "top_p", + "top_k", + "seed", +) + + +def validate_repeated_request_lifecycle( + profile: str, + width: int, + records: list[dict[str, Any]], + starts: list[dict[str, Any]], +) -> Mismatch | None: + ordered_starts = sorted(starts, key=lambda record: record["request"]) + first = ordered_starts[0] + repeated = ordered_starts[1] + for field in REPEATED_REQUEST_FIELDS: + if first[field] != repeated[field]: + return matrix_failure( + f"{profile}.q{width}.repeated_request.{field}", + first[field], + repeated[field], + "repeated request configuration differs from the first request", + ) + + token_records = event_index(records, "tokens") + first_tokens = token_records[(first["request"],)]["token_ids"] + repeated_tokens = token_records[(repeated["request"],)]["token_ids"] + if first_tokens != repeated_tokens: + return matrix_failure( + f"{profile}.q{width}.repeated_tokens", + first_tokens, + repeated_tokens, + "repeated request continuation tokens differ", + ) + + request_ends = event_index(records, "request_end") + first_end = request_ends[(first["request"],)]["cache_position"] + repeated_end = request_ends[(repeated["request"],)]["cache_position"] + if first_end != repeated_end: + return matrix_failure( + f"{profile}.q{width}.repeated_cache_position", + first_end, + repeated_end, + "repeated request final cache position differs", + ) + + if profile == "reset": + resets = event_index(records, "reset") + expected_reset_keys = {(first["request"],), (repeated["request"],)} + if resets.keys() != expected_reset_keys or any( + record["cache_position"] != 0 for record in resets.values() + ): + return matrix_failure( + f"{profile}.q{width}.reset_lifecycle", + {"requests": sorted(expected_reset_keys), "cache_position": 0}, + { + "requests": sorted(resets), + "cache_positions": sorted( + record["cache_position"] for record in resets.values() + ), + }, + "reset requests must both begin from an empty cache", + ) + logits = event_index(records, "logits") + final_position = first["prompt_tokens"] + first_logits = logits.get((first["request"], final_position)) + repeated_logits = logits.get((repeated["request"], final_position)) + if first_logits is None or repeated_logits is None: + return matrix_failure( + f"{profile}.q{width}.repeated_logits", + "final logits for both requests", + sorted(logits), + "reset repetition is missing final logits", + ) + float_diff = first_float_mismatch( + first_logits["values"], repeated_logits["values"], TOLERANCES["final_logits"] + ) + if float_diff: + index, left, right = float_diff + return matrix_failure( + f"{profile}.q{width}.repeated_logits[{index}]", + left, + right, + "repeated request final logits exceed tolerance", + ) + return None + def validate_snapshot_lifecycle( profile: str, @@ -858,6 +954,26 @@ def validate_snapshot_lifecycle( restored["request"], "snapshot restore must belong to the restored request", ) + prompt_tokens = saved_start["prompt_tokens"] + if ( + saved["cache_position"] != prompt_tokens + or restored["cache_position"] != prompt_tokens + or restored_start["cache_position"] != prompt_tokens + or saved_start["snap_pos"] != prompt_tokens + or restored_start["snap_pos"] != prompt_tokens + ): + return matrix_failure( + f"{profile}.q{width}.snapshot_full_prompt", + prompt_tokens, + { + "save": saved["cache_position"], + "restore": restored["cache_position"], + "restored_start": restored_start["cache_position"], + "fresh_snap_pos": saved_start["snap_pos"], + "restored_snap_pos": restored_start["snap_pos"], + }, + "snapshot profile requires a full-prompt save and restore", + ) for field in SNAPSHOT_STATE_FIELDS: if saved[field] != restored[field]: return matrix_failure( @@ -1018,13 +1134,16 @@ def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mism "request restore lifecycle differs from the profile contract", ) tokens = [record for record in records if record.get("type") == "tokens"] - if len(tokens) < 2 or any(len(record["token_ids"]) < 2 for record in tokens): + if len(tokens) != 2 or any(len(record["token_ids"]) < 2 for record in tokens): return matrix_failure( f"{profile}.q{width}.continuations", "two requests with multiple tokens", [len(record["token_ids"]) for record in tokens], "multiple continuation tokens are required", ) + repetition_mismatch = validate_repeated_request_lifecycle(profile, width, records, starts) + if repetition_mismatch: + return repetition_mismatch if not any(record.get("type") == "logits" for record in records): return matrix_failure( f"{profile}.q{width}.logits", True, False, "final logits are missing" diff --git a/harness/ds4_exact_diff_schema.md b/harness/ds4_exact_diff_schema.md index af677666b..cfb8ced1c 100644 --- a/harness/ds4_exact_diff_schema.md +++ b/harness/ds4_exact_diff_schema.md @@ -50,7 +50,9 @@ To cover all three tail widths with one pinned prompt, its production token coun must satisfy `N % 12 == 11`; q=2, q=3, and q=4 then end with widths 1, 2, and 3. The comparator rejects a trace set that does not actually contain those rows. Every request must end successfully and emit exactly one completion and token -record. A full-prompt restored request performs no prefill, so its DSpark state +record. The repeated request must use the same prompt and sampling configuration +and produce the same continuation tokens; reset repetitions also compare final +logits. A full-prompt restored request performs no prefill, so its DSpark state is authenticated through the snapshot's cache, logits, and feature hashes rather than through impossible duplicate capture rows. diff --git a/harness/tests/test_ds4_exact_diff.py b/harness/tests/test_ds4_exact_diff.py index 7eb6793af..44dad401c 100644 --- a/harness/tests/test_ds4_exact_diff.py +++ b/harness/tests/test_ds4_exact_diff.py @@ -108,11 +108,11 @@ def capture(position: int = 4) -> dict[str, object]: } -def request_end(cache_position: int = 4) -> dict[str, object]: +def request_end(cache_position: int = 4, *, request: int = 0) -> dict[str, object]: return { "schema": MODULE.SCHEMA, "type": "request_end", - "request": 0, + "request": request, "ok": True, "cache_position": cache_position, } @@ -462,9 +462,14 @@ def test_restored_request_final_position_uses_full_prompt_length(tmp_path: Path) def test_snapshot_restore_requires_identical_auxiliary_state() -> None: fresh = request_start() + fresh["snap_slot"] = 0 + fresh["snap_pos"] = fresh["prompt_tokens"] restored = request_start() restored["request"] = 1 restored["restored"] = True + restored["cache_position"] = restored["prompt_tokens"] + restored["snap_slot"] = 0 + restored["snap_pos"] = restored["prompt_tokens"] records = [ manifest(), snapshot("snapshot_save", 0), @@ -477,6 +482,42 @@ def test_snapshot_restore_requires_identical_auxiliary_state() -> None: assert mismatch.field == "snapshot.q1.snapshot.spec_feature_hash" +def test_repeated_request_token_divergence_fails() -> None: + fresh = request_start() + repeated = request_start() + repeated["request"] = 1 + records = [ + tokens(0), + tokens(1), + request_end(request=0), + request_end(request=1), + ] + records[1]["token_ids"] = [9, 10] + mismatch = MODULE.validate_repeated_request_lifecycle("snapshot", 1, records, [fresh, repeated]) + assert mismatch is not None + assert mismatch.field == "snapshot.q1.repeated_tokens" + + +def test_partial_snapshot_fails() -> None: + fresh = request_start(prompt_tokens=8) + fresh["snap_slot"] = 0 + fresh["snap_pos"] = 8 + restored = request_start(prompt_tokens=8) + restored["request"] = 1 + restored["restored"] = True + restored["cache_position"] = 8 + restored["snap_slot"] = 0 + restored["snap_pos"] = 8 + records = [ + manifest(), + snapshot("snapshot_save", 0), + snapshot("snapshot_restore", 1), + ] + mismatch = MODULE.validate_snapshot_lifecycle("snapshot", 1, records, [fresh, restored]) + assert mismatch is not None + assert mismatch.field == "snapshot.q1.snapshot_full_prompt" + + def test_restored_request_uses_snapshot_features_instead_of_prefill_capture() -> None: fresh = request_start() restored = request_start() From 68e7b1538281d38fdb76023dc1d4da1a8e57e01e Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Mon, 10 Aug 2026 01:30:01 +0530 Subject: [PATCH 8/9] fix(dflash): accept full cache-hit trace shape --- harness/ds4_exact_diff.py | 4 ---- harness/tests/test_ds4_exact_diff.py | 10 ++++++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/harness/ds4_exact_diff.py b/harness/ds4_exact_diff.py index 044e57dde..197d2f197 100644 --- a/harness/ds4_exact_diff.py +++ b/harness/ds4_exact_diff.py @@ -829,8 +829,6 @@ def matrix_failure(field: str, expected: Any, observed: Any, detail: str) -> Mis "prompt_token_ids", "prompt_tokens", "n_gen", - "snap_slot", - "snap_pos", "temperature", "top_p", "top_k", @@ -960,7 +958,6 @@ def validate_snapshot_lifecycle( or restored["cache_position"] != prompt_tokens or restored_start["cache_position"] != prompt_tokens or saved_start["snap_pos"] != prompt_tokens - or restored_start["snap_pos"] != prompt_tokens ): return matrix_failure( f"{profile}.q{width}.snapshot_full_prompt", @@ -970,7 +967,6 @@ def validate_snapshot_lifecycle( "restore": restored["cache_position"], "restored_start": restored_start["cache_position"], "fresh_snap_pos": saved_start["snap_pos"], - "restored_snap_pos": restored_start["snap_pos"], }, "snapshot profile requires a full-prompt save and restore", ) diff --git a/harness/tests/test_ds4_exact_diff.py b/harness/tests/test_ds4_exact_diff.py index 44dad401c..35a32b91d 100644 --- a/harness/tests/test_ds4_exact_diff.py +++ b/harness/tests/test_ds4_exact_diff.py @@ -468,8 +468,10 @@ def test_snapshot_restore_requires_identical_auxiliary_state() -> None: restored["request"] = 1 restored["restored"] = True restored["cache_position"] = restored["prompt_tokens"] - restored["snap_slot"] = 0 - restored["snap_pos"] = restored["prompt_tokens"] + # Full-cache hits receive the restored position separately; these request + # fields retain their defaults in the HTTP server. + assert restored["snap_slot"] == -1 + assert restored["snap_pos"] == -1 records = [ manifest(), snapshot("snapshot_save", 0), @@ -506,8 +508,8 @@ def test_partial_snapshot_fails() -> None: restored["request"] = 1 restored["restored"] = True restored["cache_position"] = 8 - restored["snap_slot"] = 0 - restored["snap_pos"] = 8 + assert restored["snap_slot"] == -1 + assert restored["snap_pos"] == -1 records = [ manifest(), snapshot("snapshot_save", 0), From db0c609d603824e39e54b37809f1f0404e171cbb Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Mon, 10 Aug 2026 01:47:07 +0530 Subject: [PATCH 9/9] fix(dflash): bind exact trace matrix invariants --- harness/ds4_exact_diff.py | 164 ++++++++++++++++++++--- harness/ds4_exact_diff_schema.md | 9 +- harness/tests/test_ds4_exact_diff.py | 77 ++++++++++- server/src/deepseek4/deepseek4_graph.cpp | 3 +- 4 files changed, 231 insertions(+), 22 deletions(-) diff --git a/harness/ds4_exact_diff.py b/harness/ds4_exact_diff.py index 197d2f197..e3ca1d0d9 100644 --- a/harness/ds4_exact_diff.py +++ b/harness/ds4_exact_diff.py @@ -25,6 +25,18 @@ } WIDTHS = (1, 2, 3, 4) PROFILES = ("reset", "snapshot") +RESERVED_SERVER_OPTIONS = frozenset( + { + "--chunk", + "--draft", + "--ds4-prefill", + "--host", + "--port", + "--prefix-cache-slots", + "--target-device", + "--target-devices", + } +) class TraceError(ValueError): @@ -1019,6 +1031,56 @@ def validate_capture_coverage( def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mismatch | None: + canonical_manifest: dict[str, Any] | None = None + for (profile, width), records in sorted(traces.items()): + manifest = records[0] + request_config = manifest["request_config"] + if not isinstance(request_config, dict): + return matrix_failure( + f"{profile}.q{width}.request_config", + "object", + type(request_config).__name__, + "manifest request configuration must be an object", + ) + expected_coordinates = { + "profile": profile, + "width": width, + "request_profile": profile, + "prefill_width": width, + "exact_bands": width > 1, + } + observed_coordinates = { + "profile": manifest["profile"], + "width": manifest["width"], + "request_profile": request_config.get("profile"), + "prefill_width": request_config.get("prefill_width"), + "exact_bands": request_config.get("exact_bands"), + } + if observed_coordinates != expected_coordinates: + return matrix_failure( + f"{profile}.q{width}.manifest_coordinates", + expected_coordinates, + observed_coordinates, + "manifest profile and width must match its matrix entry", + ) + + normalized_manifest = dict(manifest) + normalized_manifest.pop("profile") + normalized_manifest.pop("width") + normalized_config = dict(request_config) + for varying_field in ("profile", "prefill_width", "exact_bands", "port"): + normalized_config.pop(varying_field) + normalized_manifest["request_config"] = normalized_config + if canonical_manifest is None: + canonical_manifest = normalized_manifest + elif normalized_manifest != canonical_manifest: + return matrix_failure( + f"{profile}.q{width}.matrix_manifest", + canonical_manifest, + normalized_manifest, + "manifest identity or fixed request configuration differs across traces", + ) + model_configs: list[dict[str, Any]] = [] for (profile, width), records in sorted(traces.items()): configs = [record for record in records if record.get("type") == "model_config"] @@ -1047,6 +1109,56 @@ def validate_matrix(traces: dict[tuple[str, int], list[dict[str, Any]]]) -> Mism "model boundary configuration differs across traces", ) + canonical_request: dict[str, Any] | None = None + request_identity_fields = ( + "prompt_token_hash", + "prompt_token_ids", + "prompt_tokens", + "n_gen", + "temperature", + "top_p", + "top_k", + "seed", + ) + for (profile, width), records in sorted(traces.items()): + starts = sorted( + (record for record in records if record.get("type") == "request_start"), + key=lambda record: record["request"], + ) + if not starts: + return matrix_failure( + f"{profile}.q{width}.request_start", + "at least one request", + 0, + "trace has no request configuration to authenticate", + ) + start = starts[0] + request_config = records[0]["request_config"] + expected_request = { + "width": width, + "n_gen": request_config.get("generated_tokens"), + "temperature": request_config.get("temperature"), + "seed": request_config.get("seed"), + } + observed_request = {field: start[field] for field in expected_request} + if observed_request != expected_request: + return matrix_failure( + f"{profile}.q{width}.manifest_request", + expected_request, + observed_request, + "observed request does not match its manifest", + ) + normalized_request = {field: start[field] for field in request_identity_fields} + if canonical_request is None: + canonical_request = normalized_request + elif normalized_request != canonical_request: + return matrix_failure( + f"{profile}.q{width}.matrix_request", + canonical_request, + normalized_request, + "prompt, generation, or sampling request differs across traces", + ) + reset_candidates = [traces[("reset", width)] for width in WIDTHS[1:]] for width, records in zip(WIDTHS[1:], reset_candidates, strict=True): ordinary = any( @@ -1217,6 +1329,16 @@ def stop_process(process: subprocess.Popen[bytes]) -> None: process.wait(timeout=5.0) +def validate_server_args(server_args: list[str]) -> None: + for value in server_args: + stripped = value.strip() + if not stripped: + raise ValueError("--server-arg cannot be empty") + option = stripped.split("=", 1)[0].split(maxsplit=1)[0] + if option in RESERVED_SERVER_OPTIONS or option.startswith("--prefill-"): + raise ValueError(f"--server-arg cannot override reserved option {option}") + + def git_revision(binary: Path, explicit: str | None) -> str: if explicit: return explicit @@ -1267,20 +1389,7 @@ def write_manifest( path.write_text(json.dumps(manifest, sort_keys=True) + "\n", encoding="utf-8") -def run_one( - profile: str, - width: int, - args: argparse.Namespace, - revision: str, - hashes: dict[str, str], - prompt: str, -) -> None: - profile_dir = args.output_dir / profile - profile_dir.mkdir(parents=True, exist_ok=True) - trace_path = profile_dir / f"q{width}.jsonl" - log_path = profile_dir / f"q{width}.server.log" - port = args.port_base + (0 if profile == "reset" else 100) + width - write_manifest(trace_path, profile, width, args, revision, hashes, port) +def build_trace_environment(trace_path: Path, width: int, draft: Path | None) -> dict[str, str]: env = os.environ.copy() env.update( { @@ -1292,14 +1401,36 @@ def run_one( "DFLASH_DS4_FUSED_HYBRID_DECODE": "0", } ) - if args.draft: + if draft: env.update( { "DFLASH_DS4_SPEC": "1", - "DFLASH_DS4_DRAFT": str(args.draft), + "DFLASH_DS4_DRAFT": str(draft), "DFLASH_DS4_SPEC_Q": "4", } ) + else: + env["DFLASH_DS4_SPEC"] = "0" + env.pop("DFLASH_DS4_DRAFT", None) + env.pop("DFLASH_DS4_SPEC_Q", None) + return env + + +def run_one( + profile: str, + width: int, + args: argparse.Namespace, + revision: str, + hashes: dict[str, str], + prompt: str, +) -> None: + profile_dir = args.output_dir / profile + profile_dir.mkdir(parents=True, exist_ok=True) + trace_path = profile_dir / f"q{width}.jsonl" + log_path = profile_dir / f"q{width}.server.log" + port = args.port_base + (0 if profile == "reset" else 100) + width + write_manifest(trace_path, profile, width, args, revision, hashes, port) + env = build_trace_environment(trace_path, width, args.draft) cache_slots = "0" if profile == "reset" else "2" command = [ str(args.binary), @@ -1346,6 +1477,7 @@ def run_one( def run_matrix(args: argparse.Namespace) -> int: + validate_server_args(args.server_arg) for field in ("binary", "target", "prompt"): path = getattr(args, field) if not path.is_file(): diff --git a/harness/ds4_exact_diff_schema.md b/harness/ds4_exact_diff_schema.md index cfb8ced1c..7c35b319d 100644 --- a/harness/ds4_exact_diff_schema.md +++ b/harness/ds4_exact_diff_schema.md @@ -55,6 +55,11 @@ and produce the same continuation tokens; reset repetitions also compare final logits. A full-prompt restored request performs no prefill, so its DSpark state is authenticated through the snapshot's cache, logits, and feature hashes rather than through impossible duplicate capture rows. +Across all eight traces, the comparator requires one binary, revision, model, +prompt, tolerance contract, and fixed request configuration. Only the profile, +width, exact-band setting, and per-process port may vary. The observed prompt +tokens, generation count, sampling settings, and width must also agree with the +manifest and with every other trace. The production position filter records: @@ -102,4 +107,6 @@ python3 harness/ds4_exact_diff.py compare \ The run command hashes its inputs before launching the server, forces greedy sampling, keeps exact attention, disables approximate/fused verification, and uses `--chunk q` with the exact-band flag disabled for q=1 and enabled for -q=2..4. Raw traces and server logs belong outside Git. +q=2..4. `--server-arg` rejects options that could override those fixed +invariants, and a run without `--draft` clears inherited DSpark activation. +Raw traces and server logs belong outside Git. diff --git a/harness/tests/test_ds4_exact_diff.py b/harness/tests/test_ds4_exact_diff.py index 35a32b91d..123c8a94d 100644 --- a/harness/tests/test_ds4_exact_diff.py +++ b/harness/tests/test_ds4_exact_diff.py @@ -16,11 +16,11 @@ SPEC.loader.exec_module(MODULE) -def manifest(width: int = 1) -> dict[str, object]: +def manifest(width: int = 1, profile: str = "reset") -> dict[str, object]: return { "schema": MODULE.SCHEMA, "type": "manifest", - "profile": "reset", + "profile": profile, "width": width, "revision": "a" * 40, "binary_sha256": "b" * 64, @@ -28,9 +28,15 @@ def manifest(width: int = 1) -> dict[str, object]: "drafter_sha256": "d" * 64, "prompt_bytes_sha256": "e" * 64, "request_config": { + "prefill_mode": "exact", "prefill_width": width, "exact_bands": width > 1, + "generated_tokens": 16, "port": 18000 + width, + "profile": profile, + "seed": 1, + "server_args": [], + "target_device": "hip:0", "temperature": 0, }, "tolerances": MODULE.TOLERANCES, @@ -78,6 +84,15 @@ def request_start(prompt_tokens: int = 4) -> dict[str, object]: } +def model_config() -> dict[str, object]: + return { + "schema": MODULE.SCHEMA, + "type": "model_config", + "n_swa": 128, + "compressor_boundaries": [4, 8], + } + + def step(width: int = 1, position: int = 4, logits_present: bool = True) -> dict[str, object]: return { "schema": MODULE.SCHEMA, @@ -575,7 +590,7 @@ def test_prompt_token_ids_are_required(tmp_path: Path) -> None: def test_matrix_requires_model_config_per_trace() -> None: traces = { - (profile, width): [manifest(width)] + (profile, width): [manifest(width, profile)] for profile in MODULE.PROFILES for width in MODULE.WIDTHS } @@ -584,3 +599,59 @@ def test_matrix_requires_model_config_per_trace() -> None: assert mismatch.field == "reset.q1.model_config" assert mismatch.oracle == 1 assert mismatch.candidate == 0 + + +def test_matrix_rejects_cross_profile_manifest_identity() -> None: + traces = { + (profile, width): [manifest(width, profile)] + for profile in MODULE.PROFILES + for width in MODULE.WIDTHS + } + traces[("snapshot", 1)][0]["revision"] = "f" * 40 + mismatch = MODULE.validate_matrix(traces) + assert mismatch is not None + assert mismatch.field == "snapshot.q1.matrix_manifest" + + +def test_matrix_rejects_mislabeled_request_profile() -> None: + traces = { + (profile, width): [manifest(width, profile)] + for profile in MODULE.PROFILES + for width in MODULE.WIDTHS + } + traces[("snapshot", 1)][0]["request_config"]["profile"] = "reset" + mismatch = MODULE.validate_matrix(traces) + assert mismatch is not None + assert mismatch.field == "snapshot.q1.manifest_coordinates" + + +def test_matrix_rejects_cross_profile_request_identity() -> None: + traces = {} + for profile in MODULE.PROFILES: + for width in MODULE.WIDTHS: + start = request_start() + start["width"] = width + traces[(profile, width)] = [manifest(width, profile), model_config(), start] + traces[("snapshot", 1)][2]["prompt_token_ids"] = [9, 10, 11, 12] + mismatch = MODULE.validate_matrix(traces) + assert mismatch is not None + assert mismatch.field == "snapshot.q1.matrix_request" + + +@pytest.mark.parametrize( + "server_arg", + ["--chunk", "--ds4-prefill=dense", "--prefill-threshold", "--target-devices"], +) +def test_reserved_server_argument_fails(server_arg: str) -> None: + with pytest.raises(ValueError, match="reserved option"): + MODULE.validate_server_args([server_arg]) + + +def test_no_draft_environment_clears_inherited_dspark(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DFLASH_DS4_SPEC", "1") + monkeypatch.setenv("DFLASH_DS4_DRAFT", "ambient.gguf") + monkeypatch.setenv("DFLASH_DS4_SPEC_Q", "3") + env = MODULE.build_trace_environment(Path("trace.jsonl"), 2, None) + assert env["DFLASH_DS4_SPEC"] == "0" + assert "DFLASH_DS4_DRAFT" not in env + assert "DFLASH_DS4_SPEC_Q" not in env diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 5f69d6808..500df5a4b 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -6925,8 +6925,7 @@ bool deepseek4_step_layer_range( (fused_hybrid_decode && !verify_hooks) ? &fused_hybrid_decode_hooks : verify_hooks; if ((!moe_hybrid || fused_hybrid_ready) && - ((n_tokens >= 2 && n_tokens <= 4 && verify_hooks) || - fused_hybrid_decode) && + (fused_verify_candidate || fused_hybrid_decode) && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled()) { const bool q1_feature_capture =