fix offline batch chunked prefill

Apply the effective prefill token budget to offline batch generation instead of dispatching the entire prompt batch at once. Preserve normal batched prefill when the full input fits the budget, and process over-budget long prefills one request at a time in round-robin order while retaining absolute positions, cumulative sequence lengths, and KV allocations.

Keep Qwen warmup within the same safety threshold, avoid device sampling on intermediate chunks, and cover host- and device-embedding chunk paths with batching tests. The issue #91 workload completes prefill and 20-token decode without the AICPU 507018 failure.

Author: vegetabledoww

pypto_serving/config/types.py

@@ -68,6 +68,8 @@ class RuntimeConfig:
     npu_memory_utilization: float = 0.90
     # Max tokens processed per scheduling step (chunked-prefill granularity).
     max_num_batched_tokens: int = 4096
+    # Per-dispatch safety cap for long offline prefill; zero disables the cap.
+    long_prefill_token_threshold: int = 256
     # Compile-time generation limit used by model-specific runners.
     max_new_tokens: int = 256
 

pypto_serving/model/qwen/npu_runner.py

@@ -402,15 +402,16 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
         fused prefill) fails at startup rather than on the first real
         request.
         """
-        batch = runtime.max_batch_size
+        kernel_batch = runtime.max_batch_size
         max_seq = runtime.max_seq_len
-        mnb = getattr(runtime, "max_num_batched_tokens", 4096)
-        step_tokens = min(mnb, batch * max_seq)
-        per_req = max(step_tokens // batch, 1)
-        total_tokens = per_req * batch
+        batch, per_req, total_tokens = self._warmup_prefill_shape(runtime)
+        mnb = runtime.max_num_batched_tokens
+        long_prefill_threshold = runtime.long_prefill_token_threshold
 
         logger.info(
-            f"[warmup] starting (batch={batch}, max_num_batched_tokens={mnb}, "
+            f"[warmup] starting (prefill_batch={batch}, kernel_batch={kernel_batch}, "
+            f"max_num_batched_tokens={mnb}, "
+            f"long_prefill_token_threshold={long_prefill_threshold}, "
             f"max_seq={max_seq}, per_req={per_req}, total_tokens={total_tokens}, slot=-1)",
 
         )
@@ -459,19 +460,19 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
         compiled.decode_block_table_buffer.fill_(0)     # all reads from page 0
         compiled.decode_slot_mapping_buffer.fill_(-1)   # all writes to page 0
 
-        for b in range(batch):
+        for b in range(kernel_batch):
             compiled.decode_seq_lens_buffer[b] = min(per_req + 1, max_seq)
 
         decode_kernel_inputs = _DecodeKernelInputs(
-            actual_batch=batch,
+            actual_batch=kernel_batch,
             token_ids=compiled.decode_token_ids_buffer,
             seq_lens=compiled.decode_seq_lens_buffer,
             block_table=compiled.decode_block_table_buffer,
             slot_mapping=compiled.decode_slot_mapping_buffer,
             logits=compiled.decode_logits_buffer,
         )
 
-        logger.info(f"[warmup] decode dispatch … (batch={batch}, seq_len={per_req + 1})")
+        logger.info(f"[warmup] decode dispatch … (batch={kernel_batch}, seq_len={per_req + 1})")
         t0 = time.perf_counter()
         self._run_distributed_program(
             compiled.decode,
@@ -481,6 +482,21 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
 
         logger.info("[warmup] complete")
 
+    @staticmethod
+    def _warmup_prefill_shape(runtime: RuntimeConfig) -> tuple[int, int, int]:
+        """Return active batch, tokens per request, and total warmup tokens."""
+        max_batch = runtime.max_batch_size
+        effective_budget = runtime.max_num_batched_tokens
+        if runtime.long_prefill_token_threshold > 0:
+            effective_budget = min(effective_budget, runtime.long_prefill_token_threshold)
+        step_tokens = min(effective_budget, max_batch * runtime.max_seq_len)
+        if step_tokens <= 0:
+            raise ValueError("warmup prefill token budget must be positive")
+        active_batch = min(max_batch, step_tokens)
+        per_request_tokens = step_tokens // active_batch
+        total_tokens = per_request_tokens * active_batch
+        return active_batch, per_request_tokens, total_tokens
+
     def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> DeviceTensor:
         """Allocate one worker-resident KV cache tensor shared by prefill/decode."""
         return self._shared_l3_worker().alloc_tensor(shape, dtype)

pypto_serving/serving/engine/engine.py

@@ -207,26 +207,39 @@ def _generate_batch_impl(
                 allow_device_greedy_sampling=allow_device_greedy_sampling,
                 kv_allocations=allocations,
             )
-            fast_path_result = self._executor.try_generate_batch(
-                record,
-                requests,
-                prefill_batch,
-                generate_config,
-            )
-            if fast_path_result is not None:
-                return fast_path_result
-
-            with self._executor.session():
-                prefill_result = self._executor.run_prefill(
-                    runtime_model,
+            prefill_token_budget = record.runtime.max_num_batched_tokens
+            if prefill_token_budget <= 0:
+                raise ValueError("max_num_batched_tokens must be positive")
+            long_prefill_threshold = record.runtime.long_prefill_token_threshold
+            if long_prefill_threshold > 0:
+                prefill_token_budget = min(prefill_token_budget, long_prefill_threshold)
+            total_prefill_tokens = sum(len(token_ids) for token_ids in prompt_token_ids)
+            batch_fits_budget = total_prefill_tokens <= prefill_token_budget
+            if batch_fits_budget:
+                fast_path_result = self._executor.try_generate_batch(
+                    record,
+                    requests,
                     prefill_batch,
+                    generate_config,
                 )
-                prefill_logits = prefill_result.logits
-                prefill_sampled_token_ids = (
-                    prefill_result.sampled_token_ids
-                    if allow_device_greedy_sampling
-                    else None
-                )
+                if fast_path_result is not None:
+                    return fast_path_result
+
+            with self._executor.session():
+                if batch_fits_budget:
+                    prefill_result = self._executor.run_prefill(runtime_model, prefill_batch)
+                    prefill_logits = prefill_result.logits
+                    prefill_sampled_token_ids = (
+                        prefill_result.sampled_token_ids
+                        if prefill_batch.allow_device_greedy_sampling
+                        else None
+                    )
+                else:
+                    prefill_logits, prefill_sampled_token_ids = self._run_prefill_in_chunks(
+                        runtime_model,
+                        prefill_batch,
+                        prefill_token_budget,
+                    )
 
                 sampling_params = self._sampler.from_generate_config(generate_config)
                 current_tokens = self._sample_batch_rows(
@@ -420,6 +433,135 @@ def _generate_result(self, model_id: str, prompt: str, config: GenerateConfig) -
         """Generate one result by reusing the batch path."""
         return self.generate_batch(model_id, [prompt], config)[0]
 
+    def _run_prefill_in_chunks(
+        self,
+        runtime_model,
+        batch: PrefillBatch,
+        token_budget: int,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Run prefill calls whose combined chunk size does not exceed the budget."""
+        if token_budget <= 0:
+            raise ValueError("max_num_batched_tokens must be positive")
+
+        row_count = len(batch.request_ids)
+        prompt_lengths = [int(batch.seq_lens[row].item()) for row in range(row_count)]
+        computed_tokens = [0] * row_count
+        final_logits: list[torch.Tensor | None] = [None] * row_count
+        final_sampled_ids: list[torch.Tensor | None] = [None] * row_count
+        next_request_idx = 0
+
+        while any(computed_tokens[row] < prompt_lengths[row] for row in range(row_count)):
+            active_rows: list[int] = []
+            for offset in range(row_count):
+                row = (next_request_idx + offset) % row_count
+                if computed_tokens[row] < prompt_lengths[row]:
+                    active_rows.append(row)
+            selected_rows = active_rows[:1]
+            per_request_budget = token_budget
+            chunk_lengths = [
+                min(prompt_lengths[row] - computed_tokens[row], per_request_budget)
+                for row in selected_rows
+            ]
+            max_chunk_len = max(chunk_lengths)
+
+            chunk_token_ids = batch.token_ids.new_zeros((len(selected_rows), max_chunk_len))
+            chunk_embeddings = None
+            if batch.input_embeddings is not None:
+                chunk_embeddings = batch.input_embeddings.new_zeros(
+                    (len(selected_rows), max_chunk_len, *batch.input_embeddings.shape[2:])
+                )
+            chunk_positions = torch.full(
+                (len(selected_rows), max_chunk_len),
+                -1,
+                dtype=torch.long,
+                device=batch.token_ids.device,
+            )
+            chunk_seq_lens = batch.seq_lens.new_empty((len(selected_rows),))
+
+            for chunk_row, (request_row, chunk_len) in enumerate(
+                zip(selected_rows, chunk_lengths, strict=True)
+            ):
+                chunk_start = computed_tokens[request_row]
+                chunk_end = chunk_start + chunk_len
+                chunk_token_ids[chunk_row, :chunk_len] = batch.token_ids[
+                    request_row, chunk_start:chunk_end
+                ]
+                if chunk_embeddings is not None and batch.input_embeddings is not None:
+                    chunk_embeddings[chunk_row, :chunk_len] = batch.input_embeddings[
+                        request_row, chunk_start:chunk_end
+                    ]
+                chunk_positions[chunk_row, :chunk_len] = torch.arange(
+                    chunk_start,
+                    chunk_end,
+                    dtype=torch.long,
+                    device=batch.token_ids.device,
+                )
+                chunk_seq_lens[chunk_row] = chunk_end
+
+            prefill_result = self._executor.run_prefill(
+                runtime_model,
+                PrefillBatch(
+                    request_ids=[batch.request_ids[row] for row in selected_rows],
+                    token_ids=chunk_token_ids,
+                    input_embeddings=chunk_embeddings,
+                    seq_lens=chunk_seq_lens,
+                    allow_device_greedy_sampling=(
+                        batch.allow_device_greedy_sampling
+                        and any(
+                            computed_tokens[row] + chunk_len == prompt_lengths[row]
+                            for row, chunk_len in zip(
+                                selected_rows,
+                                chunk_lengths,
+                                strict=True,
+                            )
+                        )
+                    ),
+                    kv_allocations=[batch.kv_allocations[row] for row in selected_rows],
+                    positions=chunk_positions,
+                    block_ids=(
+                        [batch.block_ids[row] for row in selected_rows]
+                        if batch.block_ids
+                        else []
+                    ),
+                ),
+            )
+
+            for chunk_row, (request_row, chunk_len) in enumerate(
+                zip(selected_rows, chunk_lengths, strict=True)
+            ):
+                computed_tokens[request_row] += chunk_len
+                if computed_tokens[request_row] != prompt_lengths[request_row]:
+                    continue
+                final_logits[request_row] = self._select_batch_row(
+                    prefill_result.logits,
+                    chunk_row,
+                ).clone()
+                sampled_ids = (
+                    prefill_result.sampled_token_ids
+                    if batch.allow_device_greedy_sampling
+                    else None
+                )
+                if sampled_ids is not None:
+                    if sampled_ids.dim() == 0:
+                        sampled_id = sampled_ids
+                    elif sampled_ids.dim() == 1:
+                        sampled_id = sampled_ids[chunk_row]
+                    else:
+                        sampled_id = sampled_ids[chunk_row].reshape(-1)[0]
+                    final_sampled_ids[request_row] = sampled_id.clone()
+
+            next_request_idx = (selected_rows[-1] + 1) % row_count
+
+        if any(logits is None for logits in final_logits):
+            raise RuntimeError("prefill did not produce final logits for every request")
+        logits = torch.stack([row for row in final_logits if row is not None])
+        sampled_ids = None
+        if all(sampled_id is not None for sampled_id in final_sampled_ids):
+            sampled_ids = torch.stack(
+                [sampled_id for sampled_id in final_sampled_ids if sampled_id is not None]
+            )
+        return logits, sampled_ids
+
     def _sample_batch_rows(
         self,
         logits: torch.Tensor | None,