popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: proper autoregressive inference with streaming support

Browse Source 
Single-node mode now uses HF model.generate() instead of one-shot
decode_tail(), giving correct multi-token output with KV cache.

model_backend.py:
- generate_text(messages, max_new_tokens, temperature, top_p) — full
  autoregressive generation via model.generate() with chat template
- generate_text_streaming() — yields token strings via TextIteratorStreamer
- _encode_messages() — applies chat template (tokenize=False then tokenize),
  falls back to joining user messages; avoids BatchEncoding issues

torch_server.py:
- _handle_chat_completions: fast path when backend is head+tail — calls
  generate_text() or generate_text_streaming() directly instead of the
  single-token encode_prompt+decode_tail pipeline
- _stream_openai_response: new SSE streaming handler for token iterators
- Parses max_tokens, temperature, top_p from request body
- Distributed path (partial shards) unchanged

Verified: streaming and non-streaming both work with Qwen2.5-0.5B-Instruct.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-29 18:46:51 +03:00
parent 6b9caecd90
commit 607d49f5b0
2 changed files with 138 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
packages/node/meshnet_node/model_backend.py
View File

@@ -169,40 +169,97 @@ class TorchModelShard:
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
return self.tokenizer.decode([token_id], skip_special_tokens=True)
def generate_text(self, prompt: str, max_new_tokens: int = 16) -> str:
"""Generate text locally when this process owns the full model."""
def generate_text(
self,
messages: list[dict],
max_new_tokens: int = 256,
temperature: float = 1.0,
top_p: float = 1.0,
) -> str:
"""Autoregressive generation using HF generate() — single-node (head+tail) mode."""
if not self.is_head or not self.is_tail:
raise ModelBackendError("local generation requires a full head+tail shard")
if hasattr(self.tokenizer, "apply_chat_template"):
try:
encoded = self.tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
add_generation_prompt=True,
return_tensors="pt",
return_dict=True,
)
except Exception:
encoded = self.tokenizer(prompt, return_tensors="pt")
else:
encoded = self.tokenizer(prompt, return_tensors="pt")
encoded = self._encode_messages(messages)
input_ids = encoded["input_ids"].to(self.device)
attention_mask = encoded.get("attention_mask")
if attention_mask is not None:
attention_mask = attention_mask.to(self.device)
pad_token_id = getattr(self.tokenizer, "pad_token_id", None)
if pad_token_id is None:
pad_token_id = getattr(self.tokenizer, "eos_token_id", None)
pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
do_sample = temperature != 1.0 or top_p != 1.0
with self.torch.inference_mode():
generated = self.model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
max_new_tokens=max(1, int(max_new_tokens)),
do_sample=False,
do_sample=do_sample,
temperature=temperature if do_sample else None,
top_p=top_p if do_sample else None,
pad_token_id=pad_token_id,
)
new_tokens = generated[0, input_ids.shape[-1]:]
return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
def generate_text_streaming(
self,
messages: list[dict],
max_new_tokens: int = 256,
temperature: float = 1.0,
top_p: float = 1.0,
):
"""Yield decoded token strings one at a time using HF TextIteratorStreamer."""
if not self.is_head or not self.is_tail:
raise ModelBackendError("streaming generation requires a full head+tail shard")
import threading
try:
from transformers import TextIteratorStreamer # type: ignore[import]
except ImportError:
yield self.generate_text(messages, max_new_tokens, temperature, top_p)
return
encoded = self._encode_messages(messages)
input_ids = encoded["input_ids"].to(self.device)
attention_mask = encoded.get("attention_mask")
if attention_mask is not None:
attention_mask = attention_mask.to(self.device)
pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
do_sample = temperature != 1.0 or top_p != 1.0
streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
gen_kwargs = dict(
input_ids=input_ids,
attention_mask=attention_mask,
max_new_tokens=max(1, int(max_new_tokens)),
do_sample=do_sample,
temperature=temperature if do_sample else None,
top_p=top_p if do_sample else None,
pad_token_id=pad_token_id,
streamer=streamer,
)
t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True)
t.start()
for token_text in streamer:
yield token_text
t.join()
def _encode_messages(self, messages: list[dict]) -> dict:
"""Format messages with chat template (if available) and tokenize."""
if hasattr(self.tokenizer, "apply_chat_template"):
try:
prompt_str = self.tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=False,
)
return dict(self.tokenizer(prompt_str, return_tensors="pt"))
except Exception:
pass
prompt = " ".join(
str(m.get("content", ""))
for m in messages
if isinstance(m, dict) and m.get("role") == "user"
)
return dict(self.tokenizer(prompt, return_tensors="pt"))
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
position_embeddings = _rotary_position_embeddings(
self.model,

68
packages/node/meshnet_node/torch_server.py
View File

@@ -213,8 +213,31 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if body is None:
return
messages = body.get("messages", [])
if not isinstance(messages, list):
messages = []
stream = bool(body.get("stream", False))
model = str(body.get("model", ""))
model_name = str(body.get("model", ""))
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
temperature = float(body.get("temperature") or 1.0)
top_p = float(body.get("top_p") or 1.0)
# Fast path: this node owns the complete model — use HF generate() with KV cache.
# Avoids the single-token-per-forward-pass limitation of the distributed path.
if server.backend.is_head and server.backend.is_tail:
try:
if stream:
self._stream_openai_response(
server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
model_name,
)
else:
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
self._send_openai_response(text, model_name, False)
except Exception as exc:
self._send_json(500, {"error": f"generation failed: {exc}"})
return
# Distributed path: encode prompt at the head, forward activations along the route.
prompt = " ".join(
str(m.get("content", ""))
for m in messages
@@ -225,9 +248,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
except Exception as exc:
self._send_json(500, {"error": f"encode_prompt failed: {exc}"})
return
remaining_route = self._get_remaining_route(model)
remaining_route = self._get_remaining_route(model_name)
result_text = self._run_downstream_pipeline(payload, remaining_route)
self._send_openai_response(result_text, model, stream)
self._send_openai_response(result_text, model_name, stream)
def _get_remaining_route(self, model: str) -> list[str]:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
@@ -246,7 +269,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
if not route:
# Single-node mode: decode tail locally if we're the tail
# Partial shard at tail: decode the activation from the previous node.
# Full single-node (head+tail) is handled before entering this method.
if server.backend.is_tail:
try:
tensor = server.backend.torch.frombuffer(
@@ -256,7 +280,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
return server.backend.decode_tail(tensor)
except Exception as exc:
return f"decode error: {exc}"
return ""
return "no downstream route available for non-tail shard"
session = str(uuid.uuid4())
shape = payload.shape # type: ignore[union-attr]
@@ -309,6 +333,40 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
current_pos = resp_headers.get("x-meshnet-position-ids")
return ""
def _stream_openai_response(self, token_iter, model: str) -> None:
"""Stream tokens from an iterator as SSE chunks."""
chunk_id = "chatcmpl-node"
created = int(time.time())
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.end_headers()
def _emit(data: str) -> None:
self.wfile.write(f"data: {data}\n\n".encode())
self.wfile.flush()
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
"model": model,
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
}))
for token_text in token_iter:
if not token_text:
continue
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
"model": model,
"choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
}))
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}))
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
def _send_openai_response(self, text: str, model: str, stream: bool) -> None:
chunk_id = "chatcmpl-node"
created = int(time.time())