[verified] feat: complete Ralph task workstreams

packages/node/meshnet_node/model_backend.py

@@ -9,16 +9,26 @@ import json
 import os
 import threading
 import time
+import warnings
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, Mapping
 
 Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
 
+# Recipe params this backend knows how to apply (see meshnet_node.recipe_manifest).
+# A recipe is only meaningful if its params actually reach the execution path, so
+# an unknown key is an error rather than a silent no-op.
+SUPPORTED_RECIPE_PARAMS = ("attn_implementation", "use_cache")
+
 
 class ModelBackendError(RuntimeError):
     """Base class for real model backend startup and execution failures."""
 
 
+class UnsupportedRecipeParam(ModelBackendError):
+    """Raised when a recipe asks for an execution param this backend cannot apply."""
+
+
 class MissingModelDependencyError(ModelBackendError):
     """Raised when optional model dependencies are not installed."""
 
@@ -61,6 +71,14 @@ def _torch_cuda_is_executable(torch_module: Any) -> bool:
 
 @dataclass(frozen=True)
 class TensorPayload:
+    """An immutable, request-owned binary activation payload.
+
+    ``body`` is always the exact bfloat16 wire body.  It is intentionally
+    owned bytes rather than a view into a request buffer so a payload can move
+    across a hop without retaining an HTTP/WebSocket frame after that request
+    completes.
+    """
+
     body: bytes
     shape: list[int]
     attention_mask_header: str | None
@@ -213,6 +231,7 @@ class TorchModelShard:
         quantization: Quantization = "auto",
         cache_dir: Path | None = None,
         force_cpu: bool = False,
+        recipe_params: Mapping[str, Any] | None = None,
     ) -> None:
         if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
             raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -220,6 +239,8 @@ class TorchModelShard:
         self.shard_start = shard_start
         self.shard_end = shard_end
         self.quantization = quantization
+        self.recipe_params = validate_recipe_params(recipe_params)
+        attn_implementation = self.recipe_params.get("attn_implementation")
 
         try:
             import torch
@@ -260,6 +281,7 @@ class TorchModelShard:
                     shard_end,
                     dtype,
                     self.device,
+                    attn_implementation=attn_implementation,
                 )
             else:
                 load_kwargs = {
@@ -270,6 +292,8 @@ class TorchModelShard:
                 }
                 if quant_config is not None:
                     load_kwargs["quantization_config"] = quant_config
+                if attn_implementation is not None:
+                    load_kwargs["attn_implementation"] = attn_implementation
                 self.model = AutoModelForCausalLM.from_pretrained(
                     load_source,
                     **load_kwargs,
@@ -313,6 +337,8 @@ class TorchModelShard:
         # consume CPU tensors ("Pointer argument cannot be accessed from Triton"),
         # so CPU shards intentionally stay on the stateless prefill path.
         self.supports_kv_cache = self.device.type != "cpu"
+        if self.recipe_params.get("use_cache") is False:
+            self.supports_kv_cache = False
         self.kv_sessions = SessionCacheStore(
             max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
             ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
@@ -688,6 +714,19 @@ class TorchModelShard:
         )
 
 
+def validate_recipe_params(params: Mapping[str, Any] | None) -> dict[str, Any]:
+    """Return recipe params this backend can honour, or raise naming the bad key."""
+    if not params:
+        return {}
+    unsupported = [key for key in params if key not in SUPPORTED_RECIPE_PARAMS]
+    if unsupported:
+        raise UnsupportedRecipeParam(
+            f"recipe param(s) {', '.join(sorted(unsupported))} are not supported by this "
+            f"backend; it applies: {', '.join(SUPPORTED_RECIPE_PARAMS)}"
+        )
+    return dict(params)
+
+
 def load_torch_shard(
     model_id: str,
     shard_start: int,
@@ -695,9 +734,16 @@ def load_torch_shard(
     quantization: Quantization = "auto",
     cache_dir: Path | None = None,
     force_cpu: bool = False,
+    recipe_params: Mapping[str, Any] | None = None,
 ) -> TorchModelShard:
     return TorchModelShard(
-        model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu
+        model_id,
+        shard_start,
+        shard_end,
+        quantization,
+        cache_dir,
+        force_cpu=force_cpu,
+        recipe_params=recipe_params,
     )
 
 
@@ -747,6 +793,7 @@ def _load_partial_model_from_snapshot(
     init_empty_weights_fn: Any | None = None,
     set_tensor_fn: Any | None = None,
     safe_open_fn: Any | None = None,
+    attn_implementation: str | None = None,
 ) -> Any:
     from .model_catalog import layers_from_config
     from .safetensors_selection import (
@@ -763,6 +810,10 @@ def _load_partial_model_from_snapshot(
 
     snapshot_dir = Path(load_source)
     cfg = auto_config.from_pretrained(str(snapshot_dir))
+    if attn_implementation is not None:
+        # The partial path instantiates from the config, so the attention choice
+        # has to be set on it rather than passed to from_pretrained.
+        cfg._attn_implementation = attn_implementation
     total_layers = layers_from_config(cfg)
     if total_layers is None:
         raise PartialModelLoadUnsupported(
@@ -1120,7 +1171,21 @@ def _tensor_to_bytes(tensor: Any) -> bytes:
 
 
 def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
-    tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16)
+    # ``frombuffer`` views the immutable request-owned bytes for this forward
+    # only.  The following device transfer is the one required CPU→GPU copy;
+    # wrapping in ``bytearray`` first used to add an avoidable CPU allocation
+    # and copy.  Do not upcast through float32: the activation wire contract
+    # is bfloat16 and model layers accept it directly.
+    # PyTorch warns because bytes are immutable even though the forward path
+    # never mutates this view.  Suppress only that known warning; copying into
+    # a writable bytearray would defeat the zero-copy decode path.
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message="The given buffer is not writable.*",
+            category=UserWarning,
+        )
+        tensor = torch.frombuffer(body, dtype=torch.bfloat16)
     return tensor.reshape(shape)