hipengine 0.7.0__tar.gz → 0.7.1__tar.gz

{hipengine-0.7.0 → hipengine-0.7.1}/CHANGELOG.md

@@ -6,7 +6,27 @@ This changelog is for package/API releases. Performance rollup history remains i
 [`benchmarks/CHANGELOG.md`](benchmarks/CHANGELOG.md), with detailed benchmark
 evidence under [`benchmarks/results/`](benchmarks/results/).
 
-## Unreleased
+## v0.7.1 - 2026-10-14
+
+Patch release fixing long-song YuE2 generation and the default song-generation
+context budget. APIs remain alpha; support depends on the model, backend, storage
+layout, and request shape.
+
+### Fixed
+
+- **YuE2 default song budget.** The resident AR context is now sized to fit the
+  shipped default semantic generation budget. Requests are checked against both
+  the runtime allocation and the checkpoint context limit before prefill, so an
+  oversized request fails with a named request error instead of reaching a
+  runtime position check.
+- **YuE2 long-song NAR attention.** Tiled non-autoregressive (NAR) attention now
+  preserves the complete key sequence for every query tile. Long songs no longer
+  lose NAR context after roughly 82 seconds and develop the associated audio
+  corruption. Scalar and tensor-core attention paths share the corrected key-count
+  contract and have regression coverage.
+- **DMS resident session sizing.** The DMS resident runner now accepts the same
+  explicit context-size keyword as the shared GGUF runner, keeping compact DMS
+  session reservation compatible with the context-sizing path.
 
 ## v0.7.0 - 2026-10-06
 

{hipengine-0.7.0 → hipengine-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hipengine
-Version: 0.7.0
+Version: 0.7.1
 Summary: ROCm-native local LLM inference engine with a torch-free runtime hot path
 Project-URL: Homepage, https://github.com/shisa-ai/hipEngine
 Project-URL: Repository, https://github.com/shisa-ai/hipEngine
@@ -45,9 +45,9 @@ hipEngine is a ROCm-native local inference engine built primarily for AMD
 Radeon GPUs. It pairs a small Python host with custom HIP kernels for torch-free
 model loading, generation, and OpenAI-compatible serving on supported hardware.
 
-**Current release: v0.7.0.** This alpha adds YuE2 song generation, extends
-speculative decoding controls, and improves shared KV growth and memory admission.
-See the [release notes](CHANGELOG.md#v070---2026-10-06) for scope and limitations.
+**Current release: v0.7.1.** This patch fixes long-song YuE2 generation and
+makes the default YuE2 song budget fit the resident AR context. See the [release
+notes](CHANGELOG.md#v071---2026-10-14) for scope and limitations.
 
 [Supported models](#supported-models) lists the current model families and
 hardware; the [model support reference](docs/MODELS.md) gives exact formats,
@@ -253,7 +253,7 @@ The 16-snapshot setting uses more memory and is opt-in.
 
 ## Status and limits
 
-v0.7.0 is alpha. Automatic performance routes remain scoped to qualified
+v0.7.1 is a patch release. Automatic performance routes remain scoped to qualified
 model, hardware, and workload combinations:
 
 - Qwen3.6-27B and Qwen3.8-27B GGUF generation and serving on both AMD backends.

{hipengine-0.7.0 → hipengine-0.7.1}/README.md

@@ -4,9 +4,9 @@ hipEngine is a ROCm-native local inference engine built primarily for AMD
 Radeon GPUs. It pairs a small Python host with custom HIP kernels for torch-free
 model loading, generation, and OpenAI-compatible serving on supported hardware.
 
-**Current release: v0.7.0.** This alpha adds YuE2 song generation, extends
-speculative decoding controls, and improves shared KV growth and memory admission.
-See the [release notes](CHANGELOG.md#v070---2026-10-06) for scope and limitations.
+**Current release: v0.7.1.** This patch fixes long-song YuE2 generation and
+makes the default YuE2 song budget fit the resident AR context. See the [release
+notes](CHANGELOG.md#v071---2026-10-14) for scope and limitations.
 
 [Supported models](#supported-models) lists the current model families and
 hardware; the [model support reference](docs/MODELS.md) gives exact formats,
@@ -212,7 +212,7 @@ The 16-snapshot setting uses more memory and is opt-in.
 
 ## Status and limits
 
-v0.7.0 is alpha. Automatic performance routes remain scoped to qualified
+v0.7.1 is a patch release. Automatic performance routes remain scoped to qualified
 model, hardware, and workload combinations:
 
 - Qwen3.6-27B and Qwen3.8-27B GGUF generation and serving on both AMD backends.

{hipengine-0.7.0 → hipengine-0.7.1}/hipengine/benchmark/execution_profiles.py

@@ -23,7 +23,7 @@ from hipengine.execution_profiles import (
     manifest_sha256,
     validate_variant_manifest,
 )
-from scripts.quant_quality.metrics import per_row_metrics
+from scripts.quant_quality.metrics import per_row_metrics, summarize_diagnostics
 
 
 EXECUTION_PROFILE_EVALUATION_KIND = "hipengine_execution_profile_evaluation"
@@ -735,7 +735,11 @@ def _metric_vectors(
         "reference_teacher_nll": canonical["reference_teacher_nll_nats"],
         "candidate_teacher_nll": canonical["teacher_nll_nats"],
         "teacher_delta_p": canonical["delta_p"],
-        "strict_margin": np.empty(row_count, dtype=np.float64),
+        "strict_margin": canonical["reference_margin_logit"],
+        "reference_margin_logit": canonical["reference_margin_logit"],
+        "total_variation": canonical["total_variation"],
+        "entropy_delta_nats": canonical["entropy_delta_nats"],
+        "max_probability_delta": canonical["max_probability_delta"],
         "strict_top1_candidate_rank": np.empty(row_count, dtype=np.int64),
         "strict_top1_token_id": np.argmax(reference, axis=1).astype(np.int64),
         "candidate_top1_token_id": np.argmax(candidate, axis=1).astype(np.int64),
@@ -745,19 +749,16 @@ def _metric_vectors(
         strict_row = np.asarray(reference[index], dtype=np.float64)
         candidate_row = np.asarray(candidate[index], dtype=np.float64)
         strict_top1 = int(np.argmax(strict_row))
-        if strict_row.size == 1:
-            # A one-token vocabulary has no runner-up; keep retained JSON finite.
-            result["strict_margin"][index] = 0.0
-        else:
-            top_two = np.partition(strict_row, -2)[-2:]
-            result["strict_margin"][index] = float(top_two.max() - top_two.min())
         result["strict_top1_candidate_rank"][index] = int(
             1 + np.count_nonzero(candidate_row > candidate_row[strict_top1])
         )
     return result
 
 
-def _summary(vectors: Mapping[str, np.ndarray], indices: np.ndarray) -> dict[str, Any]:
+def _summary(
+    vectors: Mapping[str, np.ndarray], indices: np.ndarray,
+    rows: Sequence[RowDescriptor],
+) -> dict[str, Any]:
     kl = vectors["kl"][indices]
     top1 = vectors["top1_equal"][indices]
     overlap = vectors["topk_overlap"][indices]
@@ -783,6 +784,10 @@ def _summary(vectors: Mapping[str, np.ndarray], indices: np.ndarray) -> dict[str
         "teacher_delta_p_mean": float(teacher_delta_p.mean()),
         "strict_margin_min": float(margins.min()),
         "max_abs_logit_delta": float(max_abs.max()),
+        "diagnostics": summarize_diagnostics(
+            {key: values[indices] for key, values in vectors.items()},
+            sequence_count=len({(rows[int(i)].scenario_id, rows[int(i)].request_id) for i in indices}),
+        ),
     }
 
 
@@ -872,7 +877,7 @@ def compare_profile_logits(
         top_k=effective_top_k,
     )
     all_indices = np.arange(len(row_tuple), dtype=np.int64)
-    overall = _summary(vectors, all_indices)
+    overall = _summary(vectors, all_indices, row_tuple)
     by_scope: dict[str, dict[str, Any]] = {}
     scope_failures: list[dict[str, str]] = []
     for dimension in ("category", "shape", "transition"):
@@ -883,7 +888,7 @@ def compare_profile_logits(
                 [index for index, row in enumerate(row_tuple) if getattr(row, dimension) == value],
                 dtype=np.int64,
             )
-            summary = _summary(vectors, indices)
+            summary = _summary(vectors, indices, row_tuple)
             summary["passed"] = _summary_passes(
                 summary, threshold, top1_min=threshold.per_scope_top1_min
             )
@@ -1134,8 +1139,8 @@ def compare_bf16_noninferiority(
         top_k=effective_top_k,
     )
     indices = np.arange(len(row_tuple), dtype=np.int64)
-    strict_summary = _summary(strict_vectors, indices)
-    candidate_summary = _summary(candidate_vectors, indices)
+    strict_summary = _summary(strict_vectors, indices, row_tuple)
+    candidate_summary = _summary(candidate_vectors, indices, row_tuple)
     mean_delta = candidate_summary["kl_mean"] - strict_summary["kl_mean"]
     top1_delta = candidate_summary["top1_agreement"] - strict_summary["top1_agreement"]
     def scoped(dimension: str) -> tuple[dict[str, Any], bool]:
@@ -1150,8 +1155,8 @@ def compare_bf16_noninferiority(
                 ],
                 dtype=np.int64,
             )
-            strict_group = _summary(strict_vectors, selected)
-            candidate_group = _summary(candidate_vectors, selected)
+            strict_group = _summary(strict_vectors, selected, row_tuple)
+            candidate_group = _summary(candidate_vectors, selected, row_tuple)
             group_mean_delta = candidate_group["kl_mean"] - strict_group["kl_mean"]
             group_top1_delta = (
                 candidate_group["top1_agreement"] - strict_group["top1_agreement"]

{hipengine-0.7.0 → hipengine-0.7.1}/hipengine/generation/engine_service.py

@@ -867,6 +867,19 @@ class EngineService:
             with self._lifecycle_lock:
                 self._closed = True
                 self._closing = True
+            # A caller can enqueue shutdown after fatal failure publication
+            # but before cleanup marks the service closed. No driver iteration
+            # remains to consume that command: acknowledge it after cleanup and
+            # fail other queued controls instead of leaving their callers hung.
+            while True:
+                try:
+                    pending = self._commands.get_nowait()
+                except queue.Empty:
+                    break
+                pending.response.finish(
+                    error=(close_error if pending.kind == "shutdown" else
+                           RuntimeError("engine service closed before command execution"))
+                )
             if shutdown_response is not None:
                 shutdown_response.finish(error=close_error)
 

{hipengine-0.7.0 → hipengine-0.7.1}/hipengine/generation/qwen35_gguf.py

@@ -355,6 +355,10 @@ _GGUF_RESIDENT_MODEL_LOOP_DEFAULT_CAPACITY = 4
 # boundary must survive 13 intervening captures — and the byte budget bounds the
 # state the wide set pins.
 _PREFIX_RETAINED_STATE_BYTES_LIMIT = 1 << 30
+# DMS extension snapshots also own cold dense prompt KV, predictor planes and
+# trunk hidden rows. Give that host working set a separate, bounded budget;
+# allocations remain lazy and the existing explicit byte limit overrides both.
+_DMS_PREFIX_RETAINED_STATE_BYTES_LIMIT = 4 << 30
 _PREFIX_RETAINED_SNAPSHOTS_WIDE = 16
 _GGUF_PREFIX_RETAINED_SNAPSHOTS_ENV = "HIPENGINE_GGUF_PREFIX_RETAINED_SNAPSHOTS"
 _GGUF_PREFIX_RETAINED_STATE_BYTES_ENV = "HIPENGINE_GGUF_PREFIX_RETAINED_STATE_BYTES"
@@ -369,12 +373,12 @@ def _gguf_prefix_retained_snapshot_limit(capacity: int) -> int:
     )
 
 
-def _gguf_prefix_retained_state_bytes_limit() -> int:
-    """State-byte ceiling for the retained snapshot working set."""
+def _gguf_prefix_retained_state_bytes_limit(*, dms: bool = False) -> int:
+    """Retention ceiling; DMS includes cold host prompt planes and recurrence."""
 
     return _gguf_auto_context_int_env(
         _GGUF_PREFIX_RETAINED_STATE_BYTES_ENV,
-        _PREFIX_RETAINED_STATE_BYTES_LIMIT,
+        _DMS_PREFIX_RETAINED_STATE_BYTES_LIMIT if dms else _PREFIX_RETAINED_STATE_BYTES_LIMIT,
         minimum=1 << 20,
     )
 
@@ -1607,6 +1611,18 @@ class Qwen35GGUFBringupGenerator:
             ),
         )
         resolved = requested
+        if getattr(self, "dms_metadata_path", None) is not None:
+            # Compact retention has its own codec and transactions. Dense KV
+            # evidence and dense retained-decode leaves do not describe it.
+            if requested.storage_layout != "uniform":
+                raise NotImplementedError("DMS supports uniform per-token/head codec scales only")
+            if requested.storage_dtype.value == "int8_per_token_head":
+                requested = replace(requested, scale_dtype=DType.FP32)
+            self._kv_capability_resolution = None
+            return requested.create_policy(), requested.scale_dtype.value, (
+                requested.storage_dtype.value, requested.storage_layout,
+                requested.scale_dtype.value, requested.scale_granularity,
+            )
         if requested.storage_dtype.value == "int8_per_token_head":
             requested_dtype = str(requested.scale_dtype.value)
             declared_dtype = declared_int8_scale_dtype(
@@ -1647,11 +1663,10 @@ class Qwen35GGUFBringupGenerator:
                         ),
                     )
                 else:
-                    resolved = resolve_kv_policy("bf16")
-                    capability = capability.with_runtime_outcome(
-                        effective_kv_storage=resolved.storage_dtype.value,
-                        runtime_action="fallback_bf16",
-                        reason=f"{capability.reason}; failed closed to BF16",
+                    self._kv_capability_resolution = capability
+                    raise ValueError(
+                        "explicit kv_storage='int8_per_token_head' cannot run: "
+                        f"{capability.reason}; refusing to silently substitute BF16"
                     )
             self._kv_capability_resolution = capability
         else:
@@ -1664,6 +1679,13 @@ class Qwen35GGUFBringupGenerator:
         )
         return resolved.create_policy(), resolved.scale_dtype.value, signature
 
+    def configure_dms(self, config) -> None:
+        """Bind either a compact DMS config or a sidecar metadata path."""
+        if isinstance(config, str):
+            self.dms_metadata_path = config
+            self.dms_config = None
+        else:
+            self.dms_config = config
     def _prepare_kv_policy(self, params: Any | None) -> None:
         current = getattr(self, "_prepared_kv_signature", None)
         requested_storage = getattr(params, "kv_storage", "auto") or "auto"
@@ -1698,6 +1720,15 @@ class Qwen35GGUFBringupGenerator:
 
     def _prepared_session_kv_kwargs(self) -> dict[str, Any]:
         signature = getattr(self, "_prepared_kv_signature", None)
+        if getattr(self, "dms_metadata_path", None) is not None:
+            from hipengine.kvcache.dms import create_dms_bf16_backend, create_dms_int8_backend
+
+            factories = {"bf16": create_dms_bf16_backend, "int8_per_token_head": create_dms_int8_backend}
+            storage = "bf16" if signature is None else signature[0]
+            if storage not in factories:
+                raise NotImplementedError(f"DMS has no compact codec for {storage}")
+            return {"dms_metadata_path": self.dms_metadata_path,
+                    "dms_backend_factory": factories[storage]}
         if signature in {None, ("bf16", "uniform", "fp16", "per_token_head")}:
             return {}
         return {
@@ -2002,7 +2033,14 @@ class Qwen35GGUFBringupGenerator:
         return result
 
     def configure_dms(self, config) -> None:
-        self.dms_config = config
+        """Bind either a compact DMS config or a sidecar metadata path."""
+        if getattr(self, "_prepared_kv_signature", None) is not None:
+            raise RuntimeError("DMS must be configured before resident preparation")
+        if isinstance(config, str):
+            self.dms_metadata_path = config
+            self.dms_config = None
+        else:
+            self.dms_config = config
 
     @_target_arch_scoped
     def create_resident_model_runner(
@@ -2166,6 +2204,11 @@ class Qwen35GGUFBringupGenerator:
             "fp16",
             "per_token_head",
         )
+        dms_requested = getattr(self, "dms_metadata_path", None) is not None
+        if dms_requested:
+            # The prefill peak includes a transient dense BF16 source even
+            # when the retained compact payload is INT8.
+            signature = ("bf16", "uniform", "fp16", "per_token_head")
         storage, storage_layout, scale_dtype, granularity = (str(part) for part in signature)
         configured_budget_mib = getattr(self, "_kv_pool_memory_budget_mib", None)
         if configured_budget_mib is not None:
@@ -2213,8 +2256,8 @@ class Qwen35GGUFBringupGenerator:
                 kv_scale_granularity=granularity,
                 int8_kv_no_mirror_qualified=qualified,
                 int8_bf16_layer_indices_resolver=_layer_indices,
-                allocate_kv_cache=not bool(defer_kv_allocation),
-                workspace_lease_needed=self._packed_workspace_lease_needed(
+                allocate_kv_cache=dms_requested or not bool(defer_kv_allocation),
+                workspace_lease_needed=not dms_requested and self._packed_workspace_lease_needed(
                     max_batch_size=int(max_batch_size)
                 ),
                 reserve_bytes=_gguf_auto_context_reserve_bytes(),
@@ -2442,6 +2485,21 @@ class Qwen35GGUFBringupGenerator:
                 )
                 context = next_context
             else:
+                dms_source = getattr(session, "_dms_source", None)
+                if dms_source is not None:
+                    config = dms_source.config
+                    self._dms_serving_policy = {
+                        "metadata_path": self.dms_metadata_path,
+                        "decision_source": config.decision_source,
+                        "decision_mode": session.dms_decision_mode,
+                        "window_size": config.window_size,
+                        "target_compression_ratio": config.target_compression_ratio,
+                        "prefill_selection_mode": config.prefill_selection_mode,
+                        "storage": self._prepared_kv_signature[0],
+                        "prefill_mode": session.dms_prefill_mode,
+                        "max_active_requests": 1,
+                        "prefix_cache": "off",
+                    }
                 self._attach_capacity_estimate(
                     session,
                     shared_runner,
@@ -6646,6 +6704,7 @@ class Qwen35GGUFResidentModelRunner:
         ]
         return {
             "model_runner": {
+                "dms": copy.deepcopy(getattr(self.generator, "_dms_serving_policy", None)),
                 "capacity": int(self.capacity),
                 "max_active_requests": int(self.capacity),
                 "max_context_tokens": getattr(
@@ -6882,6 +6941,8 @@ class Qwen35GGUFResidentModelRunner:
         if self._rows:
             raise RuntimeError("cannot configure GGUF device KV pool while requests are active")
         self._clear_prefix_snapshots()
+        if config != getattr(self, "_engine_loop_config", None):
+            self._dms_prefix_cache = None
         self._engine_loop_config = config
         self._prefix_cache_mode = resolve_prefix_cache_mode(
             getattr(config, "prefix_cache", "off")
@@ -6891,6 +6952,12 @@ class Qwen35GGUFResidentModelRunner:
             if self._prefix_cache_mode == "radix"
             else None
         )
+        if getattr(self.generator, "dms_metadata_path", None) is not None:
+            if self.capacity != 1:
+                raise NotImplementedError("DMS requires singleton ownership")
+            # The session owns its transient prefill pool and compact backend.
+            # Binding a dense global allocation here would replace that owner.
+            return
         factory_session = self._available[-1].session if self._available else None
         create_global_pool = getattr(
             factory_session,
@@ -7818,13 +7885,25 @@ class Qwen35GGUFResidentModelRunner:
             else len(row.kv_allocation.reused_block_ids)
         )
         residency = self._prefix_cache_observability()
+        dms_cache = getattr(self, "_dms_prefix_cache", None)
+        is_dms = getattr(self.generator, "dms_metadata_path", None) is not None
+        provider_store = getattr(getattr(self, "_mtp2_adapter", None),
+                                 "_prefix_checkpoint_store_instance", None)
+        provider_bytes = 0 if provider_store is None else provider_store.nbytes
+        if is_dms:
+            residency = {**residency, "snapshot_entries": 0 if dms_cache is None else len(dms_cache.entries),
+                         "retained_kv_pages": 0,
+                         "resident_bytes": (0 if dms_cache is None else dms_cache.nbytes) + provider_bytes}
         mode = getattr(self, "_prefix_cache_mode", "off")
         fallback_reason = row.prefix_fallback_reason
         if fallback_reason is None and mode == "off":
             fallback_reason = "cache_off"
         return {
             "mode": mode,
-            "block_size_tokens": 256,
+            "block_size_tokens": 1 if is_dms else 256,
+            **({"reuse_scope": "prompt_prefix",
+                "cache_host_bytes": 0 if dms_cache is None else dms_cache.nbytes,
+                "cache_provider_device_bytes": provider_bytes} if is_dms else {}),
             "eligible": bool(row.prefix_eligible),
             "lookup": bool(row.prefix_lookup),
             "hit": bool(row.prefix_reused_tokens),
@@ -7859,12 +7938,27 @@ class Qwen35GGUFResidentModelRunner:
         if accounting is not None:
             diagnostics["specdec2_mtp2"] = accounting
         slot = row.slot
+        dms_backend = None if slot is None else getattr(slot.session, "_dms_backend", None)
+        if dms_backend is not None:
+            diagnostics["dms"] = dms_backend.observability_snapshot()
+            diagnostics["dms"]["backend"]["prefix_mode"] = (
+                "owned_prompt_prefix" if self._prefix_cache_mode != "off" else "off"
+            )
         audit = (
             None
             if slot is None
             else getattr(slot.session, "device_kv_layout_audit", None)
         )
-        if include_kv_layout and callable(audit):
+        if include_kv_layout and dms_backend is not None:
+            snapshot = diagnostics["dms"]
+            diagnostics["kv_layout"] = {
+                "topology": "dms_compact", "storage_dtype": dms_backend.codec,
+                "storage_layout": "per_head_compact", "scale_dtype": "fp32" if dms_backend.codec == "int8_per_token_head" else None,
+                "kv_attention_source": "dms_compact",
+                "persistent_total_bytes": snapshot["capacity"]["device_resident_bytes"],
+                "persistent_bf16_mirror_bytes": 0,
+            }
+        elif include_kv_layout and callable(audit):
             payload = audit()
             if isinstance(payload, Mapping):
                 diagnostics["kv_layout"] = copy.deepcopy(dict(payload))
@@ -8245,7 +8339,10 @@ class Qwen35GGUFResidentModelRunner:
             if getattr(self, "_prefix_cache_mode", "off") != "off" else 0
         )
         adapter.prefix_checkpoint_capacity = capacity
-        adapter.prefix_checkpoint_block_size = int(_GGUF_PREFIX_CACHE_BLOCK_TOKENS)
+        adapter.prefix_checkpoint_block_size = (
+            1 if getattr(getattr(self, "generator", None), "dms_metadata_path", None) is not None
+            else int(_GGUF_PREFIX_CACHE_BLOCK_TOKENS)
+        )
         adapter.prefix_checkpoint_key = (
             self._prefix_checkpoint_block_ids if capacity > 0 else None
         )
@@ -8275,6 +8372,14 @@ class Qwen35GGUFResidentModelRunner:
         row = self._row(int(request_id))
         if row is None:
             return None
+        if getattr(getattr(self, "generator", None), "dms_metadata_path", None) is not None:
+            if row.lease is None:
+                return None
+            session = row.lease.session
+            cached = getattr(session, "_dms_prefill_reuse", None)
+            if cached is not None and len(cached.tokens) == int(prefix_len):
+                return cached.checkpoint_key
+            return (getattr(session, "_dms_cache_identity", id(session)), int(prefix_len))
         allocation = getattr(row, "kv_allocation", None)
         block_ids = getattr(allocation, "block_ids", None)
         if not block_ids:
@@ -9140,6 +9245,20 @@ class Qwen35GGUFResidentModelRunner:
         if not commit:
             raise ValueError("GGUF resident prefill requires commit=True")
         with hip_target_arch_environment(self.generator.target_arch):
+            if getattr(self.generator, "dms_metadata_path", None) is not None:
+                # Compact capture owns a complete prompt. Scheduler chunks still
+                # advance their normal accounting; compute runs at the last one.
+                for request_id, token_row in zip(work.request_ids, work.token_rows, strict=True):
+                    row = self._row(request_id)
+                    start = row.prefill_tokens_seen
+                    chunk = tuple(int(token) for token in token_row)
+                    if chunk != row.prompt_ids[start:start + len(chunk)]:
+                        raise RuntimeError("DMS prefill chunk does not match the registered prompt")
+                    row.prefill_tokens_seen += len(chunk)
+                    if row.prefill_tokens_seen == len(row.prompt_ids):
+                        self._mark_execution_stepped((request_id,))
+                        self._prefill_dms_row(row)
+                return
             handled = self._try_prefill_native_work_batch(work)
             self._execution_mutation_window = "unknown"
             for request_id, token_row in zip(work.request_ids, work.token_rows, strict=True):
@@ -9480,6 +9599,7 @@ class Qwen35GGUFResidentModelRunner:
                 self._outputs.clear()
                 self._completed_metadata.clear()
                 self._clear_prefix_snapshots()
+                self._dms_prefix_cache = None
                 if self._kv_pool is not None:
                     # The owner's packed workspace holds a private-workspace
                     # reservation in the pool, and `close()` rejects one that
@@ -9604,6 +9724,14 @@ class Qwen35GGUFResidentModelRunner:
             return
         session = lease.session
         self._close_c1_decode_graph(row)
+        if getattr(getattr(self, "generator", None), "dms_metadata_path", None) is not None:
+            # Dense prefill planes were freed by compaction. A reset cannot
+            # recreate those planes or erase the compact allocator's history.
+            session.close()
+            self._resident_batch_owner = None
+            self._resident_batch_owner_pool_key = None
+            row.lease = None
+            return
         graph_handles = tuple(
             handle
             for handle in self._graph_handles_for_sessions((session,))
@@ -9851,6 +9979,122 @@ class Qwen35GGUFResidentModelRunner:
             raise RuntimeError("GGUF MTP2 streaming prompt adapter disappeared")
         adapter.finish_prompt_streaming(ids, success=bool(success), stream=int(stream))
 
+    def _prefill_dms_row(self, row: _GGUFResidentLoopRow) -> None:
+        raise_if_generation_deadline_expired(row.request)
+        row.lease = row.lease or self._acquire_lease()
+        session = row.lease.session
+        session.dms_max_new_tokens = max(1, int(row.request.max_tokens)) + max(1, row.mtp2_candidate_budget) + 1
+        if row.native_sampled:
+            self._prepare_sampled_prefill(row)
+        start = time.perf_counter()
+        cache_enabled = self._prefix_cache_mode != "off"
+        cache_eligible = cache_enabled
+        row.prefix_eligible = cache_eligible
+        row.prefix_lookup = cache_eligible
+        self._dms_cache_sequence = getattr(self, "_dms_cache_sequence", 0) + 1
+        session._dms_cache_identity = self._dms_cache_sequence
+        cache = getattr(self, "_dms_prefix_cache", None)
+        if cache is None:
+            from hipengine.runtime.dms_prefix import DMSPrefixCache
+            cache = self._dms_prefix_cache = DMSPrefixCache(
+                _gguf_prefix_retained_state_bytes_limit(dms=True), _PREFIX_RETAINED_SNAPSHOTS_WIDE
+            )
+        snapshot = cache.lookup(session, row.prompt_ids, bool(row.native_sampled)) if cache_eligible else None
+        hit = cache_eligible and snapshot is not None and snapshot.matches(
+            session, row.prompt_ids, bool(row.native_sampled)
+        )
+        partial = (cache_eligible and snapshot is not None and not hit
+                   and snapshot.compatible(session) and snapshot.cold is not None
+                   and snapshot.cold.reused_tokens(row.prompt_ids) > 0)
+        session._dms_prefill_reuse = snapshot if hit or partial else None
+        if hit or partial:
+            row.prefix_matched_tokens = row.prefix_reused_tokens = len(snapshot.tokens)
+        prime_from_cached_hidden = False
+        if (hit or partial) and row.mtp2_candidate_budget > 0:
+            adapter = self._resolved_mtp2_adapter()
+            prime_from_cached_hidden = row.request_id not in adapter.restorable_prefix_rows((row,))
+            if prime_from_cached_hidden:
+                # An AR cache source has target state but no NextN checkpoint.
+                # Prime the draft from owned trunk rows, not a second target prefill.
+                row.prefix_reused_tokens = 0
+        sinks = self._begin_mtp2_prompt_streaming((row,))
+        if hit or partial:
+            row.prefix_reused_tokens = len(snapshot.tokens)
+        sink = sinks[0]
+        if sink is not None and cache_eligible:
+            from hipengine.runtime.dms_prefill_cache import cache_boundary_points
+            session._dms_cache_capture_tokens = tuple(row.prompt_ids)
+            session._dms_cache_byte_budget = cache.max_bytes
+            sink.capture_boundaries = (*cache_boundary_points(session), len(row.prompt_ids))
+        if row.mtp2_candidate_budget > 0 and sink is None:
+            raise NotImplementedError("DMS MTP requires the native NextN prompt streaming provider")
+        row.prefix_fallback_reason = (
+            None if hit or partial else "cache_off" if not cache_enabled else
+            "dms_prefix_miss"
+        )
+        captured = None
+        try:
+            if prime_from_cached_hidden:
+                from hipengine.runtime.dms_prefill_cache import prime_cached_hidden
+                prime_cached_hidden(session, snapshot, sink, row.request_id)
+            if hit:
+                from hipengine.runtime.dms_prefix import restore_prompt
+                result = restore_prompt(session, snapshot)
+                row.prefix_matched_tokens = row.prefix_reused_tokens = len(row.prompt_ids)
+                row.prefix_snapshot_hit = True
+                row.prefix_source_kind = "dms_owned_prompt"
+                row.prefix_state_clone_bytes = snapshot.nbytes
+            else:
+                session._dms_cache_capture_tokens = tuple(row.prompt_ids) if cache_eligible else None
+                session._dms_cache_byte_budget = cache.max_bytes
+                session._dms_prefill_reuse = snapshot if partial else None
+                if partial:
+                    row.prefix_matched_tokens = row.prefix_reused_tokens = len(snapshot.tokens)
+                    row.prefix_snapshot_hit = True
+                    row.prefix_source_kind = "dms_cold_prefill"
+                    row.prefix_state_clone_bytes = snapshot.nbytes
+                result = session.prefill(
+                    row.prompt_ids, use_bulk=True, bulk_attention_mode="bulk",
+                    return_logits=bool(row.native_sampled),
+                    **({"target_hidden_chunk_sink": sink,
+                        "target_hidden_request_id": row.request_id} if sink is not None else {}),
+                )
+                if cache_eligible:
+                    from hipengine.runtime.dms_prefix import capture_prompt
+                    captured = capture_prompt(
+                        session, row.prompt_ids, result, sampled=bool(row.native_sampled)
+                    )
+        except BaseException:
+            self._finish_mtp2_prompt_streaming((row,), sinks, success=False)
+            raise
+        self._finish_mtp2_prompt_streaming((row,), sinks, success=True)
+        if captured is not None:
+            evicted = cache.put(captured)
+            adapter = getattr(self, "_mtp2_adapter", None)
+            if adapter is not None:
+                for old in evicted:
+                    boundaries = (*(() if old.cold is None else old.cold.boundaries), len(old.tokens))
+                    for end in boundaries:
+                        prefix = old.tokens[:end]
+                        still_owned = any(
+                            prefix == kept.tokens[:end] and
+                            (end == len(kept.tokens) or end in kept.cold.boundaries)
+                            for kept in cache.entries.values()
+                        )
+                        if not still_owned:
+                            adapter.discard_prefix_checkpoint(prefix)
+        row.prefill_ms += _timing_ms_since(start)
+        row.prefill_chunk_count += int(not hit)
+        if row.native_sampled:
+            sample = (
+                session.sample_native_from_last_logits(row.sampling_request, row.sampling_state)
+                if row.native_sampler else None
+            )
+            self._finish_sampled_prefill(row, result, native_compact_prefill=False, native_sample=sample)
+        else:
+            self._finish_native_prefill(row, result, native_compact_prefill=False)
+        raise_if_generation_deadline_expired(row.request)
+
     def _prefill_native_row(self, row: _GGUFResidentLoopRow) -> None:
         if row.slot is not None:
             return
@@ -11471,6 +11715,8 @@ class Qwen35GGUFResidentModelRunner:
         slot = row.slot
         if slot is None:
             raise RuntimeError("GGUF resident c1 decode row is missing its session slot")
+        if getattr(slot.session, "_dms_backend", None) is not None:
+            return slot.session.step(int(slot.prev_token), return_logits=False)
         if len(getattr(self, "_rows", {})) > 1:
             # Slot views deliberately share the batch owner's execution buffers.
             # A scalar graph captured for the c1 edge of a wider logical round
@@ -11878,6 +12124,9 @@ class Qwen35GGUFResidentModelRunner:
                     "specdec2_mtp2_prompt_prime_rows": float(
                         row.mtp2_prompt_prime_rows
                     ),
+                    "specdec2_mtp2_prompt_reused_rows": float(
+                        len(row.prompt_ids) - row.mtp2_prompt_prime_rows
+                    ),
                     "specdec2_mtp2_prompt_carried_bytes": float(
                         row.mtp2_prompt_carried_bytes
                     ),

{hipengine-0.7.0 → hipengine-0.7.1}/hipengine/generation/qwen35_gguf_dms.py

@@ -16,7 +16,7 @@ class Qwen35GGUFDMSModelRunner(Qwen35GGUFResidentModelRunner):
         self._engine_loop_config = config
         self._prefix_cache_mode = "off"
 
-    def _reserve_sessions(self):
+    def _reserve_sessions(self, *, max_sequence_length=None):
         config = self.generator.dms_config
         session = Qwen35GGUFResidentSession(
             self.generator.model_path, backend=self.generator.backend,