feat(us-022): X-Meshnet-Start-Layer pipeline protocol for overlapping shards
When _select_route picks two nodes with overlapping registrations (e.g.
A:0-22 and B:20-24), the tracker now injects start_layer per hop so B
executes only layers 23-24, not 20-24.
- model_backend: forward_bytes + _run_layers accept start_layer offset;
clamped to shard_start to prevent out-of-bounds indexing
- torch_server: _handle_binary_forward reads X-Meshnet-Start-Layer header;
_run_downstream_pipeline sends it per hop; route is now list[tuple[str,int]]
- server: proxy injects {endpoint, start_layer} objects in X-Meshnet-Route;
/v1/route response includes start_layer per node in the nodes list
- test: fake backends accept start_layer=None kwarg
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -145,6 +145,7 @@ class TorchModelShard:
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
start_layer: int | None = None,
|
||||
) -> TensorPayload | str:
|
||||
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
||||
self.device
|
||||
@@ -155,7 +156,9 @@ class TorchModelShard:
|
||||
position_ids = _tensor_from_int64_header(
|
||||
position_ids_header, self.torch, self.device
|
||||
)
|
||||
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
|
||||
hidden_states = self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
)
|
||||
if self.is_tail:
|
||||
return self.decode_tail(hidden_states)
|
||||
return self._payload(hidden_states, attention_mask, position_ids)
|
||||
@@ -278,7 +281,21 @@ class TorchModelShard:
|
||||
)
|
||||
return dict(self.tokenizer(prompt, return_tensors="pt"))
|
||||
|
||||
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
||||
def _run_layers(
|
||||
self,
|
||||
hidden_states: Any,
|
||||
attention_mask: Any,
|
||||
position_ids: Any,
|
||||
start_layer: int | None = None,
|
||||
) -> Any:
|
||||
# start_layer overrides shard_start for overlapping-shard routing
|
||||
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
|
||||
# indexing outside the loaded weights.
|
||||
effective_start = (
|
||||
max(self.shard_start, start_layer)
|
||||
if start_layer is not None
|
||||
else self.shard_start
|
||||
)
|
||||
position_embeddings = _rotary_position_embeddings(
|
||||
self.model,
|
||||
hidden_states,
|
||||
@@ -290,7 +307,7 @@ class TorchModelShard:
|
||||
self.torch,
|
||||
)
|
||||
with self.torch.inference_mode():
|
||||
for layer in self.layers[self.shard_start:self.shard_end + 1]:
|
||||
for layer in self.layers[effective_start:self.shard_end + 1]:
|
||||
hidden_states = _call_layer(
|
||||
layer,
|
||||
hidden_states,
|
||||
|
||||
Reference in New Issue
Block a user