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: