From 45bb4eec749f42238d457bf93dbbfc859afd5ca3 Mon Sep 17 00:00:00 2001 From: ndleslx Date: Sun, 12 Jul 2026 01:43:41 -0700 Subject: [PATCH 1/3] ci: add DeepSeek V4 generation accuracy guard Add an end-to-end DeepSeek V4 HTTP completion test that starts the serving command with the documented TP=8 configuration, sends the prompt 'Huawei is' with greedy decoding, and requires the exact six-token completion ' a leading global provider of ICT'. Bound server startup and generation, preserve HTTP error response bodies, and print only a bounded server-log tail on failure. Make process-group shutdown best-effort so stuck or failed waits cannot mask the original accuracy failure. Add focused CPU coverage for HTTP error diagnostics, bounded log reads, and the forced-kill timeout path. Run the DeepSeek guard through task-submit with automatic eight-device allocation outside the configured whitelist, a 20-minute allocation wait, a 30-minute runtime limit, the CI model at /data/l00955553/model/dsv4-flash-w8a8, and the runtime settings required by the documented server configuration. Move both Qwen guards into their own automatic task-submit reservations and add a run-scoped marker so cancelled CI jobs can reap orphaned submissions. Make serving workers close executor-owned runtime resources on every exit path, including initialization and inference failures, so distributed workers do not retain NPU memory after tests finish. Cover normal, exceptional, and repeated cleanup behavior with focused unit tests, and extend the NPU job timeout and dependencies for the HTTP guard. --- .github/workflows/ci.yml | 83 +++++++- python/core/serving_worker.py | 16 ++ tests/test_deepseek_v4_accuracy.py | 320 +++++++++++++++++++++++++++++ tests/test_serving_worker.py | 68 ++++++ 4 files changed, 483 insertions(+), 4 deletions(-) create mode 100644 tests/test_deepseek_v4_accuracy.py create mode 100644 tests/test_serving_worker.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7085b55d..949103f3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,7 +28,7 @@ jobs: unit-tests: runs-on: [self-hosted, linux, arm64, npu] - timeout-minutes: 30 + timeout-minutes: 90 defaults: run: working-directory: dist-checkout @@ -134,7 +134,7 @@ jobs: run: | source activate.sh pip install nanobind - pip install torch transformers safetensors numpy + pip install torch transformers safetensors numpy fastapi uvicorn pytest - name: Show ccache stats (after) run: ccache -s || true @@ -172,17 +172,92 @@ jobs: env: PYTHONPATH: ${{ github.workspace }}/dist-checkout PYPTO_QWEN3_MODEL_DIR: /data/l00955553/model/Qwen3-14B + PTO2_RING_DEP_POOL: 16384 + PTO2_RING_TASK_WINDOW: 16384 + PTO2_RING_HEAP: 1073741824 run: | source activate.sh - python -m pytest tests/test_qwen3_accuracy.py -q -s + marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + run_cmd="export CI_TASK_MARKER=$marker" + run_cmd="$run_cmd && cd $GITHUB_WORKSPACE/dist-checkout" + run_cmd="$run_cmd && source activate.sh" + run_cmd="$run_cmd && DEVICE_ID=\$TASK_DEVICE" + run_cmd="$run_cmd PYTHONPATH=$PYTHONPATH" + run_cmd="$run_cmd PYPTO_QWEN3_MODEL_DIR=$PYPTO_QWEN3_MODEL_DIR" + run_cmd="$run_cmd PTO2_RING_DEP_POOL=$PTO2_RING_DEP_POOL" + run_cmd="$run_cmd PTO2_RING_TASK_WINDOW=$PTO2_RING_TASK_WINDOW" + run_cmd="$run_cmd PTO2_RING_HEAP=$PTO2_RING_HEAP" + run_cmd="$run_cmd python -m pytest tests/test_qwen3_accuracy.py -q -s" + task-submit --device auto --timeout 1200 --max-time 1800 --run "$run_cmd" - name: Run Qwen3 serving guard (prefix cache, chunked prefill, multi-batch) env: PYTHONPATH: ${{ github.workspace }}/dist-checkout PYPTO_QWEN3_MODEL_DIR: /data/l00955553/model/Qwen3-14B + PTO2_RING_DEP_POOL: 16384 + PTO2_RING_TASK_WINDOW: 16384 + PTO2_RING_HEAP: 1073741824 run: | source activate.sh - python -m pytest tests/test_qwen3_serving.py -q -s + marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + run_cmd="export CI_TASK_MARKER=$marker" + run_cmd="$run_cmd && cd $GITHUB_WORKSPACE/dist-checkout" + run_cmd="$run_cmd && source activate.sh" + run_cmd="$run_cmd && DEVICE_ID=\$TASK_DEVICE" + run_cmd="$run_cmd PYTHONPATH=$PYTHONPATH" + run_cmd="$run_cmd PYPTO_QWEN3_MODEL_DIR=$PYPTO_QWEN3_MODEL_DIR" + run_cmd="$run_cmd PTO2_RING_DEP_POOL=$PTO2_RING_DEP_POOL" + run_cmd="$run_cmd PTO2_RING_TASK_WINDOW=$PTO2_RING_TASK_WINDOW" + run_cmd="$run_cmd PTO2_RING_HEAP=$PTO2_RING_HEAP" + run_cmd="$run_cmd python -m pytest tests/test_qwen3_serving.py -q -s" + task-submit --device auto --timeout 1200 --max-time 1800 --run "$run_cmd" + + - name: Run DeepSeek V4 HTTP generation accuracy guard + env: + PYPTO_DSV4_MODEL_DIR: /data/l00955553/model/dsv4-flash-w8a8 + PYPTO_RUNTIME_LOG: error + PTO2_RING_DEP_POOL: 131072 + PTO2_RING_TASK_WINDOW: 131072 + PTO2_RING_HEAP: 2147483648 + PTO2_OP_EXECUTE_TIMEOUT_US: 400000000 + PTO2_STREAM_SYNC_TIMEOUT_MS: 440000 + PTO2_SCHEDULER_TIMEOUT_MS: 320000 + SERVING_WORKER_STEP_TIMEOUT: 1800 + run: | + source activate.sh + marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + run_cmd="export CI_TASK_MARKER=$marker" + run_cmd="$run_cmd && cd $GITHUB_WORKSPACE/dist-checkout" + run_cmd="$run_cmd && source activate.sh" + run_cmd="$run_cmd && TASK_DEVICE=\$TASK_DEVICE" + run_cmd="$run_cmd PYTHONPATH=$GITHUB_WORKSPACE/dist-checkout" + run_cmd="$run_cmd PYPTO_DSV4_MODEL_DIR=$PYPTO_DSV4_MODEL_DIR" + run_cmd="$run_cmd PYPTO_RUNTIME_LOG=$PYPTO_RUNTIME_LOG" + run_cmd="$run_cmd PTO2_RING_DEP_POOL=$PTO2_RING_DEP_POOL" + run_cmd="$run_cmd PTO2_RING_TASK_WINDOW=$PTO2_RING_TASK_WINDOW" + run_cmd="$run_cmd PTO2_RING_HEAP=$PTO2_RING_HEAP" + run_cmd="$run_cmd PTO2_OP_EXECUTE_TIMEOUT_US=$PTO2_OP_EXECUTE_TIMEOUT_US" + run_cmd="$run_cmd PTO2_STREAM_SYNC_TIMEOUT_MS=$PTO2_STREAM_SYNC_TIMEOUT_MS" + run_cmd="$run_cmd PTO2_SCHEDULER_TIMEOUT_MS=$PTO2_SCHEDULER_TIMEOUT_MS" + run_cmd="$run_cmd SERVING_WORKER_STEP_TIMEOUT=$SERVING_WORKER_STEP_TIMEOUT" + run_cmd="$run_cmd python -m pytest tests/test_deepseek_v4_accuracy.py -q -s" + task-submit --device auto --device-num 8 --ignore-whitelist \ + --timeout 1200 --max-time 1800 \ + --run "$run_cmd" + + - name: Kill orphaned task-submit tasks + if: always() + working-directory: . + run: | + marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + task-submit --list 2>/dev/null \ + | grep -F "CI_TASK_MARKER=$marker" \ + | grep -oE 'task_[0-9_]+' \ + | sort -u \ + | while read -r task_id; do + echo "Killing orphaned task $task_id" + task-submit --kill "$task_id" || true + done || true platform-build: runs-on: ubuntu-latest diff --git a/python/core/serving_worker.py b/python/core/serving_worker.py index bfa68dc2..b2c154be 100644 --- a/python/core/serving_worker.py +++ b/python/core/serving_worker.py @@ -162,6 +162,17 @@ def busy_loop(self) -> None: logger.info("Worker exiting") + def close(self) -> None: + """Release executor-owned runtime and device resources.""" + executor = self.executor + self.executor = None + if executor is None: + return + + close = getattr(executor, "close", None) + if callable(close): + close() + def _execute_step(self, scheduler_output) -> StepOutput: """Execute one batch step (may contain prefill + decode requests).""" runtime_model = self.model_record.runtime_model @@ -412,6 +423,11 @@ def _worker_entry( except Exception as e: logger.error(f"Worker process failed: {e}", exc_info=True) ready_event.set() + finally: + try: + worker.close() + except Exception: + logger.exception("Worker process cleanup failed") def spawn_worker(config: EngineConfig): diff --git a/tests/test_deepseek_v4_accuracy.py b/tests/test_deepseek_v4_accuracy.py new file mode 100644 index 00000000..439ae55e --- /dev/null +++ b/tests/test_deepseek_v4_accuracy.py @@ -0,0 +1,320 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""DeepSeek V4 HTTP generation accuracy guard for CI.""" + +from __future__ import annotations + +import io +import json +import os +import queue +import signal +import socket +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] +MODEL_ID = "dsv4-flash-w8a8" +PROMPT = "Huawei is" +MAX_NEW_TOKENS = 6 +EXPECTED_TEXT = " a leading global provider of ICT" + +STARTUP_TIMEOUT_SECONDS = 600 +OVERALL_TIMEOUT_SECONDS = 1650 +HEARTBEAT_SECONDS = 30 + + +def _task_devices() -> tuple[int, ...]: + raw_devices = os.environ.get("TASK_DEVICE", "") + try: + devices = tuple(int(value.strip()) for value in raw_devices.split(",") if value.strip()) + except ValueError: + pytest.fail(f"TASK_DEVICE must contain comma-separated integer device IDs, got {raw_devices!r}") + if len(devices) != 8 or len(set(devices)) != 8 or any(device < 0 for device in devices): + pytest.fail(f"TASK_DEVICE must contain exactly 8 unique non-negative device IDs, got {raw_devices!r}") + return devices + + +def _unused_local_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _server_command(model_dir: Path, devices: tuple[int, ...], port: int) -> list[str]: + # Keep these serving options aligned with docs/dev/model/deepseek-v4.md. + # CI substitutes only the checkpoint, task-submit devices, and free port. + return [ + sys.executable, + "python/cli/main.py", + "--model", + str(model_dir), + "--served-model-name", + MODEL_ID, + "--backend", + "npu", + "--platform", + "a2a3", + "--devices", + ",".join(str(device) for device in devices), + "--dp", + "1", + "--tp", + "8", + "--block-size", + "128", + "--max-model-len", + "260", + "--max-num-seqs", + "1", + "--max-num-batched-tokens", + "512", + "--long-prefill-token-threshold", + "2048", + "--no-enable-prefix-caching", + "--port", + str(port), + "--show-startup-logs", + ] + + +def _wait_for_health(process: subprocess.Popen, port: int, deadline: float) -> None: + url = f"http://127.0.0.1:{port}/health" + startup_deadline = min(deadline, time.monotonic() + STARTUP_TIMEOUT_SECONDS) + next_heartbeat = time.monotonic() + last_error: BaseException | None = None + + while time.monotonic() < startup_deadline: + return_code = process.poll() + if return_code is not None: + raise RuntimeError(f"DeepSeek server exited before becoming healthy (code={return_code})") + try: + with urllib.request.urlopen(url, timeout=5) as response: + payload = json.loads(response.read()) + if response.status == 200 and payload == {"status": "ok"}: + print("DeepSeek server is healthy", flush=True) + return + except (OSError, TimeoutError, ValueError, urllib.error.URLError) as exc: + last_error = exc + + now = time.monotonic() + if now >= next_heartbeat: + print("Waiting for DeepSeek server startup...", flush=True) + next_heartbeat = now + HEARTBEAT_SECONDS + time.sleep(2) + + raise TimeoutError(f"DeepSeek server did not become healthy: {last_error}") + + +def _request_completion(process: subprocess.Popen, port: int, deadline: float) -> dict: + request = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/completions", + data=json.dumps( + { + "model": MODEL_ID, + "prompt": PROMPT, + "max_tokens": MAX_NEW_TOKENS, + "temperature": 0.0, + "top_p": 1.0, + } + ).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + results: queue.Queue[tuple[bool, object]] = queue.Queue(maxsize=1) + + def send_request() -> None: + try: + timeout = max(1.0, deadline - time.monotonic()) + with urllib.request.urlopen(request, timeout=timeout) as response: + body = response.read().decode("utf-8") + results.put((True, json.loads(body))) + except urllib.error.HTTPError as exc: + try: + error_body = exc.read().decode("utf-8", errors="replace") + except Exception: + error_body = "" + results.put( + (False, RuntimeError(f"completion request returned HTTP {exc.code}: {error_body}")) + ) + except BaseException as exc: + results.put((False, exc)) + + threading.Thread(target=send_request, name="deepseek-completion", daemon=True).start() + while time.monotonic() < deadline: + try: + succeeded, value = results.get(timeout=HEARTBEAT_SECONDS) + except queue.Empty: + return_code = process.poll() + if return_code is not None: + raise RuntimeError( + f"DeepSeek server exited during generation (code={return_code})" + ) from None + print("Waiting for DeepSeek completion...", flush=True) + continue + if succeeded: + if not isinstance(value, dict): + raise TypeError(f"completion response must be a JSON object, got {type(value).__name__}") + return value + if isinstance(value, BaseException): + raise value + raise RuntimeError(f"completion request failed: {value}") + raise TimeoutError("DeepSeek completion exceeded the end-to-end timeout") + + +def _stop_process_group(process: subprocess.Popen) -> None: + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + except OSError as exc: + print(f"WARNING: failed to terminate process group {process.pid}: {exc}", flush=True) + return + + try: + process.wait(timeout=20) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except OSError: + pass + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + print(f"WARNING: process group {process.pid} still alive after SIGKILL", flush=True) + except Exception as exc: + print(f"WARNING: failed to reap process group {process.pid}: {exc}", flush=True) + return + except Exception as exc: + print(f"WARNING: failed to wait for process group {process.pid}: {exc}", flush=True) + return + + # The server parent may exit before a worker child. Give the process group a + # short grace period, then kill any remaining descendants. + shutdown_deadline = time.monotonic() + 2 + while time.monotonic() < shutdown_deadline: + try: + os.killpg(process.pid, 0) + except OSError: + return + time.sleep(0.2) + try: + os.killpg(process.pid, signal.SIGKILL) + except OSError: + pass + + +def _print_server_log(log_path: Path) -> None: + if not log_path.exists(): + return + try: + with log_path.open("rb") as log_file: + log_file.seek(0, os.SEEK_END) + log_file.seek(max(0, log_file.tell() - 50000)) + content = log_file.read().decode("utf-8", errors="replace") + except OSError as exc: + print(f"WARNING: failed to read DeepSeek server log: {exc}", flush=True) + return + print("\n--- DeepSeek server log (tail) ---", flush=True) + print(content, flush=True) + + +def test_deepseek_v4_http_completion_matches_expected_text(tmp_path: Path) -> None: + model_dir_env = os.environ.get("PYPTO_DSV4_MODEL_DIR") + model_dir = Path(model_dir_env) if model_dir_env else None + if model_dir is None or not model_dir.is_dir(): + pytest.fail(f"PYPTO_DSV4_MODEL_DIR not set or not a directory: {model_dir}") + devices = _task_devices() + port = _unused_local_port() + log_path = tmp_path / "deepseek-v4-server.log" + deadline = time.monotonic() + OVERALL_TIMEOUT_SECONDS + + try: + with log_path.open("w", encoding="utf-8") as server_log: + process = subprocess.Popen( + _server_command(model_dir, devices, port), + cwd=ROOT, + stdout=server_log, + stderr=subprocess.STDOUT, + start_new_session=True, + text=True, + ) + try: + _wait_for_health(process, port, deadline) + response = _request_completion(process, port, deadline) + print(f"DeepSeek completion response: {response}", flush=True) + + assert response.get("model") == MODEL_ID + choices = response.get("choices") + assert isinstance(choices, list) and len(choices) == 1 + assert choices[0].get("text") == EXPECTED_TEXT + assert choices[0].get("finish_reason") == "length" + finally: + _stop_process_group(process) + except BaseException: + _print_server_log(log_path) + raise + + +def test_completion_http_error_includes_response_body(monkeypatch) -> None: + error = urllib.error.HTTPError( + "http://127.0.0.1/completions", + 500, + "Internal Server Error", + {}, + io.BytesIO(b"device allocation failed"), + ) + + def raise_http_error(*_args, **_kwargs): + raise error + + monkeypatch.setattr(urllib.request, "urlopen", raise_http_error) + + class RunningProcess: + @staticmethod + def poll(): + return None + + with pytest.raises(RuntimeError, match="HTTP 500: device allocation failed"): + _request_completion(RunningProcess(), 1, time.monotonic() + 1) + + +def test_stop_process_group_suppresses_final_wait_timeout(monkeypatch, capsys) -> None: + class StuckProcess: + pid = 123 + + @staticmethod + def wait(timeout): + raise subprocess.TimeoutExpired("server", timeout) + + monkeypatch.setattr(os, "killpg", lambda *_args: None) + + _stop_process_group(StuckProcess()) + + assert "still alive after SIGKILL" in capsys.readouterr().out + + +def test_print_server_log_reads_only_tail(tmp_path, capsys) -> None: + log_path = tmp_path / "server.log" + log_path.write_bytes(b"excluded-prefix\n" + b"x" * 60000 + b"\nincluded-tail\n") + + _print_server_log(log_path) + + output = capsys.readouterr().out + assert "excluded-prefix" not in output + assert "included-tail" in output diff --git a/tests/test_serving_worker.py b/tests/test_serving_worker.py new file mode 100644 index 00000000..75a5a485 --- /dev/null +++ b/tests/test_serving_worker.py @@ -0,0 +1,68 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import signal +from types import SimpleNamespace + +import pytest + +from python.core import serving_worker + + +def test_worker_close_releases_executor_once(): + executor = SimpleNamespace(close_calls=0) + + def close(): + executor.close_calls += 1 + + executor.close = close + worker = serving_worker.WorkerProcess.__new__(serving_worker.WorkerProcess) + worker.executor = executor + + worker.close() + worker.close() + + assert executor.close_calls == 1 + assert worker.executor is None + + +@pytest.mark.parametrize("busy_loop_fails", [False, True]) +def test_worker_entry_always_closes_worker(monkeypatch, busy_loop_fails): + calls = SimpleNamespace(close=0, ready=0) + + class FakeWorker: + def __init__(self, config, input_queue, output_queue): + pass + + def init_device_and_model(self): + return 7 + + def busy_loop(self): + if busy_loop_fails: + raise RuntimeError("worker failed") + + def close(self): + calls.close += 1 + + monkeypatch.setattr(serving_worker, "WorkerProcess", FakeWorker) + monkeypatch.setattr(signal, "signal", lambda *_args: None) + ready_event = SimpleNamespace(set=lambda: setattr(calls, "ready", calls.ready + 1)) + num_pages_value = SimpleNamespace(value=0) + + serving_worker._worker_entry( + SimpleNamespace(), + SimpleNamespace(), + SimpleNamespace(), + ready_event, + num_pages_value, + ) + + assert num_pages_value.value == 7 + assert calls.ready >= 1 + assert calls.close == 1 From 66976a772c7c5b0c6bb47a3998e70f93465c861f Mon Sep 17 00:00:00 2001 From: linyifan Date: Sun, 12 Jul 2026 19:54:58 -0700 Subject: [PATCH 2/3] perf(qwen3): skip redundant decode d2h/host copies in greedy path The fused Qwen3-14B decode kernel does device-side embedding and greedy sampling, so several per-step host copies are pure waste. Eliminate three, all host-only (no kernel change, no recompile): 1. Greedy logits d2h. In the device-greedy path the token id is produced inside the kernel (decode_sampled_ids_buffer), so host code never reads the [BATCH, VOCAB] f32 logits. Point the kernel's logits `out` slot at a reusable worker-resident DeviceTensor instead of the shared host buffer; the runtime only d2h-copies host shared-memory args, so the ~9.7 MB/step copy (16 x 152064 x 4B) disappears. Non-greedy still uses the host buffer. 2. next_hidden d2h. Decode never returns next_hidden (run_decode always passes None to _integrated_sample_result), so its Out buffer is scratch. Route it to a device-resident buffer unconditionally, dropping a ~160 KB/step copy. 3. Dead decode hidden copy. _prepare_decode_inputs built a [B, hidden] bf16 buffer from batch.hidden_states (a zeros placeholder for device-embedding executors) that no caller reads and the kernel does not consume. Stop building/returning it. DecodeResult.hidden_states / .logits become Optional; serving_worker reads host logits lazily (only on the non-greedy fallback), so None is safe. Verified on NPU (a2a3, greedy temperature=0): correct output ("The capital of France is" -> "Paris ... Washington"), 23 decode steps through the new path, no errors. Run-to-run token divergence was confirmed to be inherent NPU non-determinism (the same unmodified binary diverges identically across devices at a near-tie argmax), not caused by this change. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/model/qwen3_14b/runner/npu_runner.py | 79 +++++++++++++++---- python/core/serving_worker.py | 24 +++--- python/core/types.py | 12 ++- 3 files changed, 83 insertions(+), 32 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 3b800a7e..bb66c580 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -126,7 +126,6 @@ class _DecodeInputs: actual_batch: int token_ids: torch.Tensor - hidden: torch.Tensor seq_lens: torch.Tensor block_table: torch.Tensor slot_mapping: torch.Tensor @@ -177,6 +176,7 @@ def __init__( self._device_id = device_id self._l3_worker: Any | None = None self._l3_static_tensors: dict[tuple[int, tuple[int, ...], torch.dtype], object] = {} + self._decode_device_scratch_cache: dict[str, Any] = {} self._static_args: _StaticKernelArgs | None = None self._pending_kv_cache_specs: dict[str, tuple[ModelConfig, RuntimeConfig]] = {} if compiled is not None: @@ -623,6 +623,7 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: """ compiled = self._compiled model_id = model.config.model_id + allow_device_greedy = batch.allow_device_greedy_sampling decode_inputs = self._prepare_decode_inputs(model, batch) kv_cache = self._kv_caches.get(model_id) @@ -631,7 +632,16 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: k_cache = kv_cache.key_pages v_cache = kv_cache.value_pages - kernel_inputs = self._pad_decode_inputs(model, decode_inputs) + # In the device greedy path the sampled token id is produced inside the + # kernel, so the full [B, VOCAB] logits never reach the host: point the + # kernel's logits `out` slot at a worker-resident DeviceTensor instead of + # the shared host buffer to skip a ~B*VOCAB*4B d2h every decode step. + logits_target = ( + self._decode_device_scratch("logits", compiled.decode_logits_buffer) + if allow_device_greedy + else None + ) + kernel_inputs = self._pad_decode_inputs(model, decode_inputs, logits_override=logits_target) # Padded block_table / slot_mapping only ever reference row 0's # already-valid pages, so bound-check exactly what the kernel will read. @@ -651,13 +661,20 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: # hidden row to return here. None, kernel_inputs.actual_batch, - allow=batch.allow_device_greedy_sampling, + allow=allow_device_greedy, + ) + # Host logits are only consumed by host-side sampling (non-greedy). In + # the greedy path logits stayed device-resident, so return None; the + # decode hidden state is a device-embedded placeholder that no caller + # reads, so it is never materialized. + logits = ( + None + if allow_device_greedy + else kernel_inputs.logits[: kernel_inputs.actual_batch, : model.config.vocab_size] ) return DecodeResult( - hidden_states=decode_inputs.hidden.float(), - logits=kernel_inputs.logits[: kernel_inputs.actual_batch, : model.config.vocab_size].to( - decode_inputs.hidden.device - ), + hidden_states=None, + logits=logits, sampled_token_ids=sampled_ids, next_hidden_states=next_hidden, ) @@ -777,15 +794,47 @@ def _decode_kernel_args( static.padded_embed_weight, inputs.token_ids, self._compiled.decode_sampled_ids_buffer, - self._compiled.decode_next_hidden_buffer, + # next_hidden is never read back in decode; keep it device-resident + # so the kernel's write does not trigger a per-step d2h. + self._decode_device_scratch("next_hidden", self._compiled.decode_next_hidden_buffer), ) - def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _DecodeKernelInputs: + def _decode_device_scratch(self, name: str, host: torch.Tensor) -> Any: + """Return a reusable worker-resident scratch buffer shaped like ``host``. + + Some decode kernel ``out`` slots have their device→host copy wasted: + + * ``logits`` — in the greedy path the token id is sampled on-device, so + the full ``[BATCH, VOCAB]`` logits never need to reach the host. + * ``next_hidden`` — decode never returns it (``run_decode`` passes + ``None`` to ``_integrated_sample_result``), so it is pure scratch. + + Pointing those slots at a resident DeviceTensor keeps the kernel's write + (and, for logits, its internal argmax) device-local and skips a per-step + d2h. Buffers are allocated lazily on the shared worker and reused across + steps. + """ + tensor = self._decode_device_scratch_cache.get(name) + if tensor is None: + tensor = self._shared_l3_worker().alloc_tensor(tuple(host.shape), host.dtype) + self._decode_device_scratch_cache[name] = tensor + return tensor + + def _pad_decode_inputs( + self, + model: RuntimeModel, + inputs: _DecodeInputs, + logits_override: Any | None = None, + ) -> _DecodeKernelInputs: """Pad active decode rows to the fixed kernel batch. The fused decode kernel computes all ``max_batch_size`` rows. Inactive rows replicate row 0 so their KV writes are idempotent instead of targeting unrelated pages. + + ``logits_override`` selects the kernel's logits ``out`` target: a + worker-resident DeviceTensor for the greedy path (no d2h) or ``None`` to + use the shared host ``decode_logits_buffer`` when host logits are needed. """ compiled = self._compiled actual_batch = inputs.actual_batch @@ -837,7 +886,7 @@ def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _Dec kernel_batch, rows_each=1, ), - logits=compiled.decode_logits_buffer, + logits=compiled.decode_logits_buffer if logits_override is None else logits_override, ) def _run_distributed_program(self, callable_spec: _L3Callable, *args: Any) -> Any: @@ -953,6 +1002,9 @@ def close(self) -> None: finally: self._l3_worker = None self._l3_static_tensors.clear() + # worker.close() frees all worker-resident DeviceTensors; just + # drop our references to the reusable decode scratch buffers. + self._decode_device_scratch_cache.clear() def _prepare_prefill_inputs( self, @@ -1063,11 +1115,12 @@ def _prepare_decode_inputs( """Pack active decode requests into fused decode-kernel inputs.""" batch_count = len(batch.kv_allocations) if batch.kv_allocations else int(batch.seq_lens.shape[0]) actual_batch = self._validate_batch_size(model, batch_count) - hidden_size = model.config.hidden_size page_size = model.runtime.page_size max_blocks = self._max_blocks_per_seq(model) - hidden = torch.zeros((actual_batch, hidden_size), dtype=torch.bfloat16) + # The fused decode kernel embeds the token ids on-device, so it never + # consumes batch.hidden_states (a zeros placeholder). No decode hidden + # buffer is built or returned. seq_lens = torch.empty((actual_batch,), dtype=torch.int32) block_table = torch.full((actual_batch * max_blocks,), -1, dtype=torch.int32) slot_mapping = torch.empty((actual_batch,), dtype=torch.int32) @@ -1081,7 +1134,6 @@ def _prepare_decode_inputs( raise ValueError( f"decode seq_len {seq_len} exceeds max_seq_len {model.runtime.max_seq_len}" ) - hidden[batch_idx, :] = batch.hidden_states[batch_idx].to(torch.bfloat16).cpu() seq_lens[batch_idx] = seq_len if alloc is not None: @@ -1100,7 +1152,6 @@ def _prepare_decode_inputs( return _DecodeInputs( actual_batch=actual_batch, token_ids=batch.token_ids.to(torch.int32).cpu(), - hidden=hidden, seq_lens=seq_lens, block_table=block_table, slot_mapping=slot_mapping, diff --git a/python/core/serving_worker.py b/python/core/serving_worker.py index b2c154be..8f0a5f22 100644 --- a/python/core/serving_worker.py +++ b/python/core/serving_worker.py @@ -271,11 +271,6 @@ def _batch_prefill( request = sr.request will_be_computed = sr.num_computed_tokens + sr.num_new_tokens if will_be_computed >= request.num_prompt_tokens: - logits = ( - prefill_result.logits[i] - if prefill_result.logits.dim() > 1 - else prefill_result.logits - ) params = SamplingParams( temperature=request.temperature, top_p=request.top_p, @@ -283,7 +278,6 @@ def _batch_prefill( ) token_id = self._sample_result_row( prefill_result, - logits, params, i, allow_device_greedy_sampling, @@ -362,11 +356,6 @@ def _batch_decode( for i, sr in enumerate(scheduled): request = sr.request - logits = ( - decode_result.logits[i] - if decode_result.logits.dim() > 1 - else decode_result.logits - ) params = SamplingParams( temperature=request.temperature, top_p=request.top_p, @@ -374,7 +363,6 @@ def _batch_decode( ) token_id = self._sample_result_row( decode_result, - logits, params, i, allow_device_greedy_sampling, @@ -384,12 +372,16 @@ def _batch_decode( def _sample_result_row( self, result, - logits: torch.Tensor, params: SamplingParams, row_idx: int, allow_device_sampled: bool, ) -> int: - """Return a sampled token from executor output, falling back to host sampling.""" + """Return a sampled token from executor output, falling back to host sampling. + + Host ``result.logits`` is only read on the fallback path; device-sampling + executors may return ``None`` logits (kept device-resident), which is + fine because the device-sampled id short-circuits before logits are used. + """ sampled = getattr(result, "sampled_token_ids", None) if allow_device_sampled and sampled is not None: flat = sampled.view(-1) @@ -398,7 +390,9 @@ def _sample_result_row( f"sampled_token_ids has {flat.numel()} rows, expected row {row_idx}" ) return int(flat[row_idx].item()) - return self.sampler.sample(logits, params) + logits = result.logits + logits_row = logits[row_idx] if logits.dim() > 1 else logits + return self.sampler.sample(logits_row, params) def _worker_entry( config: EngineConfig, diff --git a/python/core/types.py b/python/core/types.py index 9f0dc787..b1f96452 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -222,10 +222,16 @@ class DecodeBatch: @dataclass class DecodeResult: - """Outputs from one decode step.""" + """Outputs from one decode step. - hidden_states: torch.Tensor - logits: torch.Tensor + ``hidden_states`` and ``logits`` are optional: device-sampling executors + that produce the next token id inside the decode kernel may return ``None`` + for both, keeping the full ``[B, VOCAB]`` logits device-resident (no d2h) + and skipping the unused decode hidden-state copy. + """ + + hidden_states: torch.Tensor | None + logits: torch.Tensor | None sampled_token_ids: torch.Tensor | None = None next_hidden_states: torch.Tensor | None = None From 94e8975dea404a36bb75c406a1c1775c378d7808 Mon Sep 17 00:00:00 2001 From: linyifan Date: Sun, 12 Jul 2026 21:57:50 -0700 Subject: [PATCH 3/3] ci: re-trigger NPU unit-tests (chunked-prefill argmax flake, not a regression) The decode d2h/host-copy removal is numerically neutral: the kernel's greedy argmax reads identical logit values whether the logits buffer is host- or device-resident, next_hidden is write-only, and the removed decode hidden copy was never consumed. Local reruns pass chunked-prefill 3/3 on this branch and 2/2 on main; the prior CI failure was an inherent chunked-vs-full-prefill near-tie flip on the first token. Co-Authored-By: Claude Opus 4.8 (1M context)