From 2bca273b17ff07c19d81f2d7561c84ce13d6b5ac Mon Sep 17 00:00:00 2001 From: vegetabledoww Date: Thu, 16 Jul 2026 04:15:42 -0700 Subject: [PATCH] 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. --- pypto_serving/config/types.py | 2 + pypto_serving/model/qwen/npu_runner.py | 34 +++-- pypto_serving/serving/engine/engine.py | 178 ++++++++++++++++++++++--- tests/test_batching.py | 147 +++++++++++++++++++- 4 files changed, 328 insertions(+), 33 deletions(-) diff --git a/pypto_serving/config/types.py b/pypto_serving/config/types.py index 6f3f4d07..a5072291 100644 --- a/pypto_serving/config/types.py +++ b/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 diff --git a/pypto_serving/model/qwen/npu_runner.py b/pypto_serving/model/qwen/npu_runner.py index 2fe22ea5..3b154c0c 100644 --- a/pypto_serving/model/qwen/npu_runner.py +++ b/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,11 +460,11 @@ 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, @@ -471,7 +472,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: 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) diff --git a/pypto_serving/serving/engine/engine.py b/pypto_serving/serving/engine/engine.py index 4b7515ef..05aa5559 100644 --- a/pypto_serving/serving/engine/engine.py +++ b/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, diff --git a/tests/test_batching.py b/tests/test_batching.py index 544ccbff..c8ce2844 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -92,6 +92,11 @@ def test_scheduler_speculative_output_counts_only_tokens_retained_before_eos(): assert [(output.new_token_id, output.finished) for output in outputs] == [(7, True)] +class _CharacterTokenizer(_Tokenizer): + def encode(self, text: str) -> list[int]: + return [(ord(char) % 15) + 1 for char in text] + + def test_worker_step_error_queues_finished_ids_for_executor_release(): aborted: list[str] = [] core = ReplicaEngineCore.__new__(ReplicaEngineCore) @@ -124,6 +129,8 @@ def _model( max_seq_len: int = 128, page_size: int = 64, eos_token_id: int | None = None, + max_num_batched_tokens: int = 4096, + long_prefill_token_threshold: int = 256, ) -> RuntimeModel: config = ModelConfig( model_id="test-model", @@ -148,6 +155,8 @@ def _model( max_batch_size=max_batch_size, max_seq_len=max_seq_len, device="cpu", + max_num_batched_tokens=max_num_batched_tokens, + long_prefill_token_threshold=long_prefill_token_threshold, ) return RuntimeModel( config=config, @@ -384,7 +393,7 @@ def test_decode_kernel_inputs_reject_multi_token_rows(): def test_engine_generate_batch_uses_batched_executor_results(): model = _model(max_batch_size=2, eos_token_id=0) manager = KvCacheManager() - executor = _ImmediateEosExecutor(manager) + executor = _RecordingPrefillExecutor(manager) engine = LLMEngine(kv_cache_manager=manager, executor=executor) manager.register_model(model.config.model_id, model.config, model.runtime) engine._models[model.config.model_id] = ModelRecord( @@ -403,10 +412,119 @@ def test_engine_generate_batch_uses_batched_executor_results(): assert [result.token_ids for result in results] == [[0], [0]] assert [result.finish_reason for result in results] == ["eos", "eos"] + assert len(executor.prefill_batches) == 1 + assert len(executor.prefill_batches[0].request_ids) == 2 + + +def test_engine_chunks_offline_batch_within_effective_token_budget(): + model = _model( + max_batch_size=2, + eos_token_id=0, + max_num_batched_tokens=8, + long_prefill_token_threshold=3, + ) + manager = KvCacheManager() + executor = _RecordingPrefillExecutor(manager) + engine = LLMEngine(kv_cache_manager=manager, executor=executor) + manager.register_model(model.config.model_id, model.config, model.runtime) + engine._models[model.config.model_id] = ModelRecord( + config=model.config, + runtime=model.runtime, + tokenizer=_CharacterTokenizer(), + layer_specs=[], + runtime_model=model, + ) + + results = engine.generate_batch( + model.config.model_id, + ["abcde", "abcdefg"], + GenerateConfig(max_new_tokens=1, temperature=0.0), + ) + + assert [result.token_ids for result in results] == [[0], [0]] + positions_by_request: dict[str, list[int]] = {} + token_ids_by_request: dict[str, list[int]] = {} + for batch in executor.prefill_batches: + assert batch.positions is not None + assert int((batch.positions >= 0).sum().item()) <= 3 + for row, request_id in enumerate(batch.request_ids): + valid_positions = batch.positions[row][batch.positions[row] >= 0] + valid_tokens = batch.token_ids[row, : valid_positions.numel()] + positions_by_request.setdefault(request_id, []).extend(valid_positions.tolist()) + token_ids_by_request.setdefault(request_id, []).extend(valid_tokens.tolist()) + assert int(batch.seq_lens[row].item()) == int(valid_positions[-1].item()) + 1 + + assert sorted(positions_by_request.values(), key=len) == [list(range(5)), list(range(7))] + assert sorted(token_ids_by_request.values(), key=len) == sorted( + [_CharacterTokenizer().encode("abcde"), _CharacterTokenizer().encode("abcdefg")], + key=len, + ) + + +def test_engine_rotates_requests_when_budget_is_smaller_than_batch(): + model = _model( + max_batch_size=3, + eos_token_id=0, + max_num_batched_tokens=2, + long_prefill_token_threshold=256, + ) + manager = KvCacheManager() + executor = _RecordingPrefillExecutor(manager) + engine = LLMEngine(kv_cache_manager=manager, executor=executor) + manager.register_model(model.config.model_id, model.config, model.runtime) + engine._models[model.config.model_id] = ModelRecord( + config=model.config, + runtime=model.runtime, + tokenizer=_CharacterTokenizer(), + layer_specs=[], + runtime_model=model, + ) + + engine.generate_batch( + model.config.model_id, + ["abc", "defg", "hijkl"], + GenerateConfig(max_new_tokens=1, temperature=0.0), + ) + + positions_by_request: dict[str, list[int]] = {} + token_ids_by_request: dict[str, list[int]] = {} + for batch in executor.prefill_batches: + assert batch.positions is not None + assert len(batch.request_ids) == 1 + for row, request_id in enumerate(batch.request_ids): + valid_positions = batch.positions[row][batch.positions[row] >= 0] + valid_tokens = batch.token_ids[row, : valid_positions.numel()] + positions_by_request.setdefault(request_id, []).extend(valid_positions.tolist()) + token_ids_by_request.setdefault(request_id, []).extend(valid_tokens.tolist()) + + assert sorted(positions_by_request.values(), key=len) == [ + list(range(3)), + list(range(4)), + list(range(5)), + ] + assert sorted(token_ids_by_request.values(), key=len) == sorted( + [ + _CharacterTokenizer().encode("abc"), + _CharacterTokenizer().encode("defg"), + _CharacterTokenizer().encode("hijkl"), + ], + key=len, + ) + + +def test_qwen_warmup_prefill_does_not_exceed_budget_below_batch_size(): + runtime = RuntimeConfig( + max_batch_size=4, + max_seq_len=128, + max_num_batched_tokens=8, + long_prefill_token_threshold=2, + ) + + assert ModelRunner._warmup_prefill_shape(runtime) == (2, 1, 2) def test_engine_uses_device_sampled_prefill_token_when_available(): - model = _model(max_batch_size=1, eos_token_id=0) + model = _model(max_batch_size=1, eos_token_id=0, long_prefill_token_threshold=2) model.embed_tokens = torch.arange(model.config.vocab_size * model.config.hidden_size, dtype=torch.float32).view( model.config.vocab_size, model.config.hidden_size, @@ -419,7 +537,7 @@ def test_engine_uses_device_sampled_prefill_token_when_available(): engine._models[model.config.model_id] = ModelRecord( config=model.config, runtime=model.runtime, - tokenizer=_Tokenizer(), + tokenizer=_CharacterTokenizer(), layer_specs=[], runtime_model=model, ) @@ -431,7 +549,8 @@ def test_engine_uses_device_sampled_prefill_token_when_available(): )[0] assert result.token_ids == [3] - assert executor.prefill_calls == 1 + assert executor.prefill_calls == 2 + assert executor.prefill_sampling_flags == [False, True] assert executor.decode_calls == 0 assert sampler.sample_calls == 0 @@ -754,6 +873,16 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: return DecodeResult(hidden_states=hidden, logits=logits) +class _RecordingPrefillExecutor(_ImmediateEosExecutor): + def __init__(self, kv_cache_manager: KvCacheManager) -> None: + super().__init__(kv_cache_manager) + self.prefill_batches: list[PrefillBatch] = [] + + def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult: + self.prefill_batches.append(batch) + return super().run_prefill(model, batch) + + class _NoopKernel: def __call__(self, *args, config=None): return None @@ -798,6 +927,7 @@ def __init__( self.second_token = second_token self.return_next_hidden = return_next_hidden self.prefill_calls = 0 + self.prefill_sampling_flags: list[bool] = [] self.decode_calls = 0 self.lookup_calls = 0 self.decode_hidden_seen: list[torch.Tensor] = [] @@ -816,13 +946,18 @@ def lookup_embeddings(self, model: RuntimeModel, token_ids: torch.Tensor) -> tor def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult: self.prefill_calls += 1 + self.prefill_sampling_flags.append(batch.allow_device_greedy_sampling) assert batch.input_embeddings is None token = torch.tensor([self.first_token], dtype=torch.int64) + sampled_token_ids = token.to(torch.int32) if batch.allow_device_greedy_sampling else None + next_hidden_states = None + if batch.allow_device_greedy_sampling and self.return_next_hidden: + next_hidden_states = model.embed_tokens.index_select(0, token) return PrefillResult( last_hidden=None, logits=torch.zeros(1, model.config.vocab_size), - sampled_token_ids=token.to(torch.int32), - next_hidden_states=model.embed_tokens.index_select(0, token) if self.return_next_hidden else None, + sampled_token_ids=sampled_token_ids, + next_hidden_states=next_hidden_states, ) def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: