fix offline batch chunked prefill

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,