[core]refactor communication layer: PR1(Added Refactor Infra Only) (#1555)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/worker/payload_span.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Helpers for explicit thinker decode span metadata."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+
+THINKER_DECODE_EMBEDDINGS_KEY = "thinker_decode_embeddings"
+THINKER_OUTPUT_TOKEN_IDS_KEY = "thinker_output_token_ids"
+THINKER_DECODE_TOKEN_START_KEY = "thinker_decode_embeddings_token_start"
+THINKER_DECODE_TOKEN_END_KEY = "thinker_decode_embeddings_token_end"
+
+CACHED_THINKER_DECODE_EMBEDDINGS_KEY = "cached_thinker_decode_embeddings"
+CACHED_THINKER_DECODE_TOKEN_START_KEY = "cached_thinker_decode_embeddings_token_start"
+CACHED_THINKER_DECODE_TOKEN_END_KEY = "cached_thinker_decode_embeddings_token_end"
+
+TensorSpan = tuple[torch.Tensor, int, int]
+
+
+def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None:
+    tensor = payload.get(tensor_key)
+    start = payload.get(start_key)
+    end = payload.get(end_key)
+    if not isinstance(tensor, torch.Tensor):
+        return None
+    if not isinstance(start, int) or not isinstance(end, int):
+        return None
+    if start < 0 or end < start or (end - start) != int(tensor.shape[0]):
+        return None
+    return tensor, start, end
+
+
+def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None:
+    if existing_span is None or incoming_span is None:
+        return None
+
+    existing_tensor, existing_start, existing_end = existing_span
+    incoming_tensor, incoming_start, incoming_end = incoming_span
+    if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype:
+        incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype)
+    if incoming_start == existing_end:
+        return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end
+    if incoming_start < existing_end:
+        overlap = existing_end - incoming_start
+        if overlap >= int(incoming_tensor.shape[0]):
+            return existing_tensor, existing_start, existing_end
+        trimmed_tensor = incoming_tensor[overlap:]
+        return (
+            torch.cat([existing_tensor, trimmed_tensor], dim=0),
+            existing_start,
+            existing_end + int(trimmed_tensor.shape[0]),
+        )
+    return None
+
+
+def get_tensor_span_row(span: TensorSpan | None, index: int) -> torch.Tensor | None:
+    if span is None:
+        return None
+    tensor, start, end = span
+    if index < start or index >= end:
+        return None
+    return tensor[index - start]