Commit accepted speculative prefixes in fixed state pools

src/config.cpp

@@ -513,24 +513,161 @@
 
 using DecoderStateGroup = Config::Model::Decoder::StateGroup;
 using DecoderStateGroupKind = Config::Model::Decoder::StateGroupKind;
+using DecoderStateUpdate = Config::Model::Decoder::StateUpdate;
+using DecoderStateUpdateKind = Config::Model::Decoder::StateUpdateKind;
+using DecoderCheckpointAlignment = Config::Model::Decoder::CheckpointAlignment;
+
+constexpr int kMaxStateUpdateCapacity = 8;
+// Both packed state operators cap their checkpoint window at eight slots.
+constexpr int kMaxStateCheckpoints = 8;
+
+struct StateUpdate_Element : JSON::Element {
+  explicit StateUpdate_Element(DecoderStateUpdate& v) : v_{v} {}
+
+  void OnValue(std::string_view name, JSON::Value value) override {
+    if (name == "enabled") {
+      v_.enabled = JSON::Get<bool>(value);
+    } else if (name == "kind") {
+      const auto kind = JSON::Get<std::string_view>(value);
+      if (kind == "causal_conv") {
+        v_.kind = DecoderStateUpdateKind::CausalConv;
+      } else if (kind == "gated_delta_net") {
+        v_.kind = DecoderStateUpdateKind::GatedDeltaNet;
+      } else {
+        throw std::runtime_error("Unsupported decoder state update kind '" + std::string{kind} + "'");
+      }
+    } else if (name == "capacity") {
+      const auto capacity = SafeDoubleToInt64(
+          JSON::Get<double>(value), "model.decoder.state_groups.state_update.capacity");
+      if (capacity < 1 || capacity > kMaxStateUpdateCapacity) {
+        throw std::runtime_error("Decoder state update capacity must be in [1, " +
+                                 std::to_string(kMaxStateUpdateCapacity) + "]");
+      }
+      v_.capacity = static_cast<int>(capacity);
+    } else if (name == "capture_count") {
+      v_.capture_count = JSON::Get<std::string_view>(value);
+    } else if (name == "value") {
+      v_.value = JSON::Get<std::string_view>(value);
+    } else if (name == "active") {
+      v_.active = JSON::Get<std::string_view>(value);
+    } else if (name == "capsule") {
+      v_.capsule = JSON::Get<std::string_view>(value);
+    } else if (name == "key_head_count") {
+      const auto count = SafeDoubleToInt64(
+          JSON::Get<double>(value), "model.decoder.state_groups.state_update.key_head_count");
+      if (count < 1 || count > std::numeric_limits<int>::max()) {
+        throw std::runtime_error("Decoder state update key_head_count must be positive");
+      }
+      v_.key_head_count = static_cast<int>(count);
+    } else {
+      throw JSON::unknown_value_error{};
+    }
+  }
+
+ private:
+  DecoderStateUpdate& v_;
+};
+
+struct StateBinding_Element : JSON::Element {
+  explicit StateBinding_Element(Config::Model::Decoder::StateBinding& v) : v_{v} {}
+
+  void OnValue(std::string_view name, JSON::Value value) override {
+    if (name == "input") {
+      v_.input = JSON::Get<std::string_view>(value);
+    } else if (name == "output") {
+      v_.output = JSON::Get<std::string_view>(value);
+    } else if (name == "checkpoints") {
+      v_.checkpoints = JSON::Get<std::string_view>(value);
+    } else {
+      throw JSON::unknown_value_error{};
+    }
+  }
+
+ private:
+  Config::Model::Decoder::StateBinding& v_;
+};
+
+struct StateBindings_Element : JSON::Element {
+  explicit StateBindings_Element(DecoderStateGroup& v) : v_{v} {}
+
+  Element& OnObject(std::string_view name) override {
+    std::optional<Config::Model::Decoder::StateBinding>* binding{};
+    std::unique_ptr<StateBinding_Element>* element{};
+    if (name == "key") {
+      binding = &v_.key;
+      element = &key_;
+    } else if (name == "value") {
+      binding = &v_.value;
+      element = &value_;
+    } else if (name == "state") {
+      binding = &v_.state;
+      element = &state_;
+    } else {
+      throw JSON::unknown_value_error{};
+    }
+
+    if (binding->has_value()) {
+      throw std::runtime_error("Duplicate decoder state binding semantic '" + std::string{name} + "'");
+    }
+    binding->emplace();
+    *element = std::make_unique<StateBinding_Element>(binding->value());
+    return **element;
+  }
+
+ private:
+  DecoderStateGroup& v_;
+  std::unique_ptr<StateBinding_Element> key_;
+  std::unique_ptr<StateBinding_Element> value_;
+  std::unique_ptr<StateBinding_Element> state_;
+};
 
 struct StateGroup_Element : JSON::Element {
   explicit StateGroup_Element(DecoderStateGroup& v) : v_{v} {}
 
   void OnValue(std::string_view name, JSON::Value value) override {
-    if (name != "kind") {
+    if (name == "kind") {
+      const auto kind = JSON::Get<std::string_view>(value);
+      if (kind == "paged_kv") {
+        v_.kind = DecoderStateGroupKind::PagedKeyValue;
+      } else if (kind == "fixed") {
+        v_.kind = DecoderStateGroupKind::Fixed;
+      } else {
+        throw std::runtime_error("Unsupported decoder state group kind '" + std::string{kind} + "'");
+      }
+    } else if (name == "checkpoint_count") {
+      v_.checkpoint_count = SafeDoubleToInt(JSON::Get<double>(value), name);
+      if (v_.checkpoint_count < 0 || v_.checkpoint_count > kMaxStateCheckpoints) {
+        throw std::runtime_error("Decoder state group checkpoint_count must be in [0, " +
+                                 std::to_string(kMaxStateCheckpoints) + "]");
+      }
+    } else if (name == "checkpoint_alignment") {
+      const auto alignment = JSON::Get<std::string_view>(value);
+      if (alignment == "left") {
+        v_.checkpoint_alignment = DecoderCheckpointAlignment::Left;
+      } else if (alignment == "right") {
+        v_.checkpoint_alignment = DecoderCheckpointAlignment::Right;
+      } else {
+        throw std::runtime_error("Unsupported decoder state group checkpoint_alignment '" +
+                                 std::string{alignment} + "'");
+      }
+    } else {
       throw JSON::unknown_value_error{};
     }
-    const auto kind = JSON::Get<std::string_view>(value);
-    if (kind == "paged_kv") {
-      v_.kind = DecoderStateGroupKind::PagedKeyValue;
-    } else if (kind == "fixed_conv") {
-      v_.kind = DecoderStateGroupKind::FixedConv;
-    } else if (kind == "fixed_recurrent") {
-      v_.kind = DecoderStateGroupKind::FixedRecurrent;
-    } else {
-      throw std::runtime_error("Unsupported decoder state group kind '" + std::string{kind} + "'");
+  }
+
+  Element& OnObject(std::string_view name) override {
+    if (name == "bindings") {
+      return bindings_;
     }
+    if (name == "state_update") {
+      if (v_.state_update) {
+        throw std::runtime_error("Duplicate decoder state_update declaration");
+      }
+      v_.state_update.emplace();
+      state_update_ = std::make_unique<StateUpdate_Element>(*v_.state_update);
+      return *state_update_;
+    }
+    throw JSON::unknown_value_error{};
   }
 
   Element& OnArray(std::string_view name) override {
@@ -543,6 +680,8 @@
  private:
   DecoderStateGroup& v_;
   IntArray_Element layer_ids_{v_.layer_ids};
+  StateBindings_Element bindings_{v_};
+  std::unique_ptr<StateUpdate_Element> state_update_;
 };
 
 struct StateGroups_Element : JSON::Element {
@@ -758,7 +897,7 @@
     } else if (name == "state_update_capacity") {
       // The kernel packs every captured transition for a layer into one fixed-width capsule output.
       // Keep this limit synchronized with check_extra_options in builder.py and the model-builder README.
       constexpr int kMaxStateUpdateCapacity = 8;
       v_.state_update_capacity = SafeDoubleToInt(JSON::Get<double>(value), name);
       if (v_.state_update_capacity < 0 || v_.state_update_capacity > kMaxStateUpdateCapacity)
         throw std::runtime_error("state_update_capacity must be between 0 and " + std::to_string(kMaxStateUpdateCapacity));

src/config.h

@@ -391,13 +391,50 @@ struct Config {
       enum class StateGroupKind {
         Invalid,
         PagedKeyValue,
-        FixedConv,
-        FixedRecurrent,
+        Fixed,
+      };
+
+      struct StateBinding {
+        std::string input;
+        std::string output;
+        // Optional third output exposing the per-token state series of this state, shaped
+        // [checkpoint_count, ...output shape]. Present only on models exported for speculative
+        // decoding; the runtime promotes one slot to `output` to roll back a rejected draft.
+        std::string checkpoints;
+      };
+
+      // Which slot of a checkpoint output corresponds to which token of the step.
+      enum class CheckpointAlignment {
+        Left,   // Slot j is the state after local token j (VarlenCausalConvWithState).
+        Right,  // Slot count-1 is the state after the last token (GatedDeltaNet).
+      };
+
+      enum class StateUpdateKind {
+        Invalid,
+        CausalConv,
+        GatedDeltaNet,
+      };
+
+      struct StateUpdate {
+        StateUpdateKind kind{StateUpdateKind::Invalid};
+        int capacity{};
+        std::string capture_count;
+        std::string value;
+        bool enabled{true};
+        std::string active;
+        std::string capsule;
+        int key_head_count{};
       };
 
       struct StateGroup {
         StateGroupKind kind{StateGroupKind::Invalid};
         std::vector<int> layer_ids;
+        std::optional<StateBinding> key;
+        std::optional<StateBinding> value;
+        std::optional<StateBinding> state;
+        std::optional<StateUpdate> state_update;
+        int checkpoint_count{0};
+        CheckpointAlignment checkpoint_alignment{CheckpointAlignment::Right};
       };
 
       // Absence preserves the legacy dense, sequential paged-KV contract.