perf(qwen3): skip redundant decode d2h/host copies in greedy path

.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

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,

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
@@ -260,19 +271,13 @@ 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,
                         top_k=request.top_k,
                     )
                     token_id = self._sample_result_row(
                         prefill_result,
-                        logits,
                         params,
                         i,
                         allow_device_greedy_sampling,
@@ -351,19 +356,13 @@ 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,
                     top_k=request.top_k,
                 )
                 token_id = self._sample_result_row(
                     decode_result,
-                    logits,
                     params,
                     i,
                     allow_device_greedy_sampling,
@@ -373,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)
@@ -387,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,
@@ -412,6 +417,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):