diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 949103f..016ab9a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -174,7 +174,7 @@ jobs: PYPTO_QWEN3_MODEL_DIR: /data/l00955553/model/Qwen3-14B PTO2_RING_DEP_POOL: 16384 PTO2_RING_TASK_WINDOW: 16384 - PTO2_RING_HEAP: 1073741824 + PTO2_RING_HEAP: 2147483648 run: | source activate.sh marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" @@ -196,7 +196,7 @@ jobs: PYPTO_QWEN3_MODEL_DIR: /data/l00955553/model/Qwen3-14B PTO2_RING_DEP_POOL: 16384 PTO2_RING_TASK_WINDOW: 16384 - PTO2_RING_HEAP: 1073741824 + PTO2_RING_HEAP: 2147483648 run: | source activate.sh marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index 87b65c9..f67da2f 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -345,7 +345,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--model-dir", required=True, help="Local model directory, e.g. a Hugging Face snapshot.") parser.add_argument("--prompt", required=True, help="Prompt text.") parser.add_argument("--model-id", default="qwen3-14b-local") - parser.add_argument("--platform", default="a2a3", choices=["a2a3sim", "a2a3", "a5sim", "a5"]) + parser.add_argument("--platform", default="a2a3", choices=["a2a3"]) parser.add_argument("--device-id", type=int, default=0, help="Default NPU device id when --devices is unset.") parser.add_argument( "--devices", @@ -374,9 +374,7 @@ def build_parser() -> argparse.ArgumentParser: "--max-num-batched-tokens", type=int, default=4096, - help="Total tokens per scheduling step (used by warmup). NOTE: the 40-layer " - "fused prefill deadlocks the single-die ring-heap above ~415 total tokens, " - "so set this low enough that the per-request count stays under the ceiling.", + help="Total tokens per scheduling step and production-scale warmup.", ) parser.add_argument( "--npu-memory-utilization", diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index b729b65..047cd3b 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -27,7 +27,7 @@ from python.core._profiling import StageTimer from python.core.model_runner import ModelRunner from python.core.pypto_executor import PyptoExecutor as CorePyptoExecutor -from python.core.types import RuntimeModel +from python.core.types import GenerateConfig, ModelRecord, RuntimeModel from python.core.utils import rope_tables, round_up @@ -104,11 +104,13 @@ def __init__( self, kv_cache_manager=None, *, - platform: str = "a2a3sim", + platform: str = "a2a3", device_ids: Sequence[int] = (0,), save_kernels_dir: str | None = None, l3_trace: bool = False, ) -> None: + if platform != "a2a3": + raise ValueError("Qwen3 direct CANN FAI decode supports only A2/A3 onboard") super().__init__( kv_cache_manager, platform=platform, @@ -132,6 +134,21 @@ def supports_device_embedding(self) -> bool: """Qwen3 NPU decode embeds greedy token ids inside the device kernel.""" return True + def validate_generate_batch( + self, + record: ModelRecord, + batch_size: int, + config: GenerateConfig, + ) -> None: + """Accept only the active B1 prefix or all B16 KV-cache rows.""" + del config + kernel_batch = record.runtime.max_batch_size + if batch_size not in (1, kernel_batch): + raise ValueError( + "Qwen3 direct CANN FAI decode supports only batch 1 or the " + f"full batch capacity {kernel_batch}, got {batch_size}" + ) + def _create_runner(self, model_id: str, compiled: object) -> ModelRunner: """Create the Qwen3-14B runtime runner for compiled kernels.""" if not isinstance(compiled, _CompiledKernels): @@ -180,10 +197,10 @@ def _mark(label: str) -> None: if int(qwen3_decode_fwd.BATCH) != kernel_batch: raise ValueError( - "decode_fwd.decode_fwd is compiled for a fixed kernel BATCH of " + "decode_fwd.decode_fwd has an internal batch capacity of " f"{int(qwen3_decode_fwd.BATCH)}, but runtime max_batch_size is " - f"{kernel_batch}; they must match (decode statically computes and " - "writes BATCH rows / BATCH logit rows)." + f"{kernel_batch}; they must match. Runtime dispatch supports only " + "batch 1 or the full batch capacity." ) if int(model.config.num_hidden_layers) != int(qwen3_decode_fwd.NUM_LAYERS): raise ValueError( @@ -455,12 +472,13 @@ def _compile_decode_fwd_callable( ) -> _L3Callable: """Compile the fused all-layer PAGED decode HOST wrapper into a distributed program. - Signature (21 args; PAGED KV via block_table + slot_mapping, same pool as + Signature (25 args; PAGED KV via block_table + slot_mapping, same pool as prefill): input_rms_weight, wq, wk, wv, q_norm_weight, k_norm_weight, seq_lens, block_table, slot_mapping, rope_cos, rope_sin, k_cache, v_cache, wo, w_gate, w_up, w_down, post_rms_weight, - final_norm_weight, lm_head_weight, out. + final_norm_weight, lm_head_weight, out, embed_weight, + sampled_ids_in, sampled_ids_out, next_hidden. k_cache/v_cache are the PAGED pool (rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size — identical to prefill); diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 3b800a7..834870d 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -134,7 +134,7 @@ class _DecodeInputs: @dataclass class _DecodeKernelInputs: - """Fixed-batch tensors passed to the fused decode kernel.""" + """Active-prefix tensors passed to the fused decode kernel.""" actual_batch: int token_ids: torch.Tensor @@ -142,6 +142,8 @@ class _DecodeKernelInputs: block_table: torch.Tensor slot_mapping: torch.Tensor logits: torch.Tensor + sampled_ids: torch.Tensor + next_hidden: torch.Tensor @dataclass @@ -183,8 +185,8 @@ def __init__( self._share_static_kernel_tensors() self._static_args = self._build_static_kernel_args() - #: Scratch KV pages for the profile pass — slot=-1 means only page 0 - #: is ever touched (reads via block_table=0, writes via slot clamp to 0). + #: Scratch KV pages for the profile pass. Warmup maps every read and write + #: to slot 0, so one page covers the full synthetic batch. _PROFILE_PAGES = 1 def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> int: @@ -192,7 +194,7 @@ def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConf Order (vLLM-style): run a profile warmup FIRST so the simpler arena is allocated before the KV cache competes for HBM. The profile uses - ``slot_mapping=-1`` / ``block_table=0`` so only a single dummy page + ``slot_mapping=0`` / ``block_table=0`` so only a single dummy page is needed. The KV cache size is then computed by the estimation formula and allocated into the remaining space; if allocation fails the page count is halved and retried. @@ -213,7 +215,7 @@ def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConf self._materialize_static_tensors() # -- phase 1: profile warmup → arena allocated ---------------------- - # Uses slot_mapping=-1 so no real KV cache pages are needed; the + # Uses slot_mapping=0 so no real KV cache pages are needed; the # 1-page scratch is the dummy target for all reads/writes. logger.info(f"[init_kv_cache] profile warmup (scratch {self._PROFILE_PAGES} page) …") ModelRunner.init_kv_cache(self, model_id, config, runtime, num_pages=self._PROFILE_PAGES) @@ -381,15 +383,12 @@ def warmup(self, model: RuntimeModel) -> None: self._warmup_dispatch(model.runtime) def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: - """Production-scale prefill + decode warm-up with slot_mapping=-1. + """Production-scale prefill + decode warm-up mapped to scratch slot 0. Sizes the prefill to one serving scheduling step — total tokens = ``max_num_batched_tokens`` spread across ``max_batch`` requests. - This deliberately exercises the kernel at the configured capacity so - that a too-large ``max_num_batched_tokens`` (which would hit the - single-die attention heap ceiling around seq≈415 in the 40-layer - fused prefill) fails at startup rather than on the first real - request. + This exercises the configured capacity so ring sizing or allocation + failures surface at startup rather than on the first real request. """ batch = runtime.max_batch_size max_seq = runtime.max_seq_len @@ -400,7 +399,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: logger.info( f"[warmup] starting (batch={batch}, max_num_batched_tokens={mnb}, " - f"max_seq={max_seq}, per_req={per_req}, total_tokens={total_tokens}, slot=-1)", + f"max_seq={max_seq}, per_req={per_req}, total_tokens={total_tokens}, slot=0)", ) compiled = self._compiled @@ -412,7 +411,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: compiled.prefill_chunk_lens_buffer.zero_() compiled.prefill_chunk_offsets_buffer.zero_() compiled.prefill_block_table_buffer.fill_(0) # all reads from page 0 - compiled.prefill_slot_mapping_buffer.fill_(-1) # all writes to page 0 + compiled.prefill_slot_mapping_buffer.zero_() # all writes to page 0 token_offset = 0 for b in range(batch): @@ -442,11 +441,11 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: ) logger.info(f"[warmup] prefill done ({time.perf_counter() - t0:.2f} s)") - # -- decode (full fixed batch, minimal seq) ------------------------- + # -- decode (full capacity batch, minimal seq) ---------------------- compiled.decode_token_ids_buffer.zero_() compiled.decode_seq_lens_buffer.zero_() compiled.decode_block_table_buffer.fill_(0) # all reads from page 0 - compiled.decode_slot_mapping_buffer.fill_(-1) # all writes to page 0 + compiled.decode_slot_mapping_buffer.zero_() # all writes to page 0 for b in range(batch): compiled.decode_seq_lens_buffer[b] = min(per_req + 1, max_seq) @@ -458,6 +457,8 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: block_table=compiled.decode_block_table_buffer, slot_mapping=compiled.decode_slot_mapping_buffer, logits=compiled.decode_logits_buffer, + sampled_ids=compiled.decode_sampled_ids_buffer, + next_hidden=compiled.decode_next_hidden_buffer, ) logger.info(f"[warmup] decode dispatch … (batch={batch}, seq_len={per_req + 1})") @@ -614,12 +615,9 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: KV is already in place with no bridge. KV is keyed by block_table page id, not by kernel row, so a request may occupy any row each step (no stable-slot shim). - The kernel is FIXED-BATCH (it computes all max_batch_size rows and writes - each row's current-token KV). Pad the active batch up to the kernel batch by - REPLICATING active row 0's inputs into the padding rows: those rows then - recompute row 0's K/V and write row 0's own slot with byte-identical values - (an idempotent, safe write), and their logits are trimmed off below. This - avoids padded rows clobbering an unrelated request's physical page. + The public decode batch is either one active row or all 16 independent + rows. Dynamic tensor descriptors carry only that active prefix into the + kernel; B1 never materializes or computes replicated KV rows. """ compiled = self._compiled model_id = model.config.model_id @@ -631,10 +629,8 @@ 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) + kernel_inputs = self._stage_decode_inputs(model, decode_inputs) - # Padded block_table / slot_mapping only ever reference row 0's - # already-valid pages, so bound-check exactly what the kernel will read. self._validate_kv_cache_bounds(model, kernel_inputs.block_table, kernel_inputs.slot_mapping, k_cache) self._run_distributed_program( @@ -644,7 +640,7 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: for batch_idx, alloc in enumerate(batch.kv_allocations): alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) sampled_ids, next_hidden = self._integrated_sample_result( - compiled.decode_sampled_ids_buffer, + kernel_inputs.sampled_ids, # decode_fwd's next_hidden output is the embedding for sampled_ids_in # used by this decode step. The newly sampled token is embedded at the # start of the following decode_fwd call, so there is no next-step @@ -655,9 +651,7 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: ) 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 - ), + logits=kernel_inputs.logits[:, : model.config.vocab_size].to(decode_inputs.hidden.device), sampled_token_ids=sampled_ids, next_hidden_states=next_hidden, ) @@ -776,22 +770,23 @@ def _decode_kernel_args( inputs.logits, static.padded_embed_weight, inputs.token_ids, - self._compiled.decode_sampled_ids_buffer, - self._compiled.decode_next_hidden_buffer, + inputs.sampled_ids, + inputs.next_hidden, ) - def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _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. - """ + def _stage_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _DecodeKernelInputs: + """Copy active decode rows into shared buffers and return prefix views.""" compiled = self._compiled actual_batch = inputs.actual_batch kernel_batch = model.runtime.max_batch_size max_blocks = self._max_blocks_per_seq(model) + if actual_batch not in (1, kernel_batch): + raise ValueError( + "Qwen3 direct CANN FAI decode supports only batch 1 or the " + f"full batch capacity {kernel_batch}, got {actual_batch}" + ) + if kernel_batch > compiled.decode_logits_buffer.shape[0]: raise ValueError( f"kernel batch {kernel_batch} exceeds logits buffer batch " @@ -808,36 +803,23 @@ def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _Dec ) width = 1 token_ids[:actual_batch, :width].copy_(active_token_ids[:, :width]) - if actual_batch < kernel_batch: - token_ids[actual_batch:, :width].copy_( - active_token_ids[0:1, :width].expand(kernel_batch - actual_batch, width) - ) + + seq_lens = compiled.decode_seq_lens_buffer[:actual_batch] + block_table = compiled.decode_block_table_buffer[: actual_batch * max_blocks] + slot_mapping = compiled.decode_slot_mapping_buffer[:actual_batch] + seq_lens.copy_(inputs.seq_lens.reshape(actual_batch)) + block_table.copy_(inputs.block_table.reshape(actual_batch * max_blocks)) + slot_mapping.copy_(inputs.slot_mapping.reshape(actual_batch)) return _DecodeKernelInputs( actual_batch=actual_batch, - token_ids=token_ids, - seq_lens=self._copy_replicated_rows( - compiled.decode_seq_lens_buffer, - inputs.seq_lens, - actual_batch, - kernel_batch, - rows_each=1, - ), - block_table=self._copy_replicated_rows( - compiled.decode_block_table_buffer, - inputs.block_table, - actual_batch, - kernel_batch, - rows_each=max_blocks, - ), - slot_mapping=self._copy_replicated_rows( - compiled.decode_slot_mapping_buffer, - inputs.slot_mapping, - actual_batch, - kernel_batch, - rows_each=1, - ), - logits=compiled.decode_logits_buffer, + token_ids=token_ids[:actual_batch], + seq_lens=seq_lens, + block_table=block_table, + slot_mapping=slot_mapping, + logits=compiled.decode_logits_buffer[:actual_batch], + sampled_ids=compiled.decode_sampled_ids_buffer[:actual_batch], + next_hidden=compiled.decode_next_hidden_buffer[:actual_batch], ) def _run_distributed_program(self, callable_spec: _L3Callable, *args: Any) -> Any: @@ -908,23 +890,6 @@ def _materialize_static_tensors(self) -> None: ): self._coerce_l3_arg(worker, arg) - @staticmethod - def _copy_replicated_rows( - dst: torch.Tensor, - active: torch.Tensor, - actual_batch: int, - kernel_batch: int, - *, - rows_each: int, - ) -> torch.Tensor: - """Copy active rows and fill inactive rows by replicating row 0.""" - active_view = active.reshape(actual_batch, rows_each) - dst_view = dst.reshape(kernel_batch, rows_each) - dst_view[:actual_batch].copy_(active_view) - if actual_batch < kernel_batch: - dst_view[actual_batch:].copy_(active_view[0:1].expand(kernel_batch - actual_batch, rows_each)) - return dst - @staticmethod def _static_device_tensor(tensor: torch.Tensor) -> _StaticDeviceTensor: """Mark a CPU tensor for one-time upload to the shared worker.""" diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index 781a01c..b0b40e3 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -102,7 +102,7 @@ def qwen3_decode_host( sampled_ids: pl.Out[pl.Tensor], next_hidden: pl.Out[pl.Tensor], ) -> tuple[pl.Tensor, pl.Tensor, pl.Tensor]: - logits, sampled_ids, next_hidden = decode_fwd( + out, sampled_ids, next_hidden = decode_fwd( input_rms_weight, wq, wk, @@ -129,7 +129,7 @@ def qwen3_decode_host( sampled_ids, next_hidden, ) - return logits, sampled_ids, next_hidden + return out, sampled_ids, next_hidden @pl.jit.host diff --git a/pypto-lib b/pypto-lib index 3bbd72d..10d43e5 160000 --- a/pypto-lib +++ b/pypto-lib @@ -1 +1 @@ -Subproject commit 3bbd72d00b9b90fe6e70b9077840b3cf6f978f73 +Subproject commit 10d43e50cfe828c9198f37d2273812ae0c0477d2 diff --git a/tests/test_batching.py b/tests/test_batching.py index 850682e..3173e67 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -181,8 +181,8 @@ def _compiled_kernels( decode_block_table_buffer=torch.zeros(kernel_batch * max_blocks, dtype=torch.int32), decode_slot_mapping_buffer=torch.zeros(kernel_batch, dtype=torch.int32), decode_logits_buffer=torch.zeros(kernel_batch, model.config.vocab_size), - decode_token_ids_buffer=torch.empty(kernel_batch, 1, dtype=torch.int32), - decode_sampled_ids_buffer=torch.empty(kernel_batch, 1, dtype=torch.int32), + decode_token_ids_buffer=torch.empty(kernel_batch, 8, dtype=torch.int32), + decode_sampled_ids_buffer=torch.empty(kernel_batch, 8, dtype=torch.int32), decode_next_hidden_buffer=torch.empty(kernel_batch, hidden_size, dtype=torch.bfloat16), ) @@ -341,7 +341,7 @@ def test_decode_kernel_inputs_reject_multi_token_rows(): runner = ModelRunner(compiled=_compiled_kernels(model)) with pytest.raises(ValueError, match="exactly one token per row"): - runner._pad_decode_inputs( + runner._stage_decode_inputs( model, SimpleNamespace( actual_batch=1, @@ -354,6 +354,103 @@ def test_decode_kernel_inputs_reject_multi_token_rows(): ) +@pytest.mark.parametrize("actual_batch", [1, 16]) +def test_decode_dispatch_uses_only_active_prefix_views(actual_batch): + model = _model(max_batch_size=16) + compiled = _compiled_kernels(model) + runner = ModelRunner(compiled=compiled) + max_blocks = (model.runtime.max_seq_len + model.runtime.page_size - 1) // model.runtime.page_size + inputs = SimpleNamespace( + actual_batch=actual_batch, + token_ids=torch.arange(10, 10 + actual_batch, dtype=torch.int32).reshape(actual_batch, 1), + hidden=torch.ones(actual_batch, model.config.hidden_size, dtype=torch.bfloat16), + seq_lens=torch.arange(100, 100 + actual_batch, dtype=torch.int32), + block_table=torch.arange(200, 200 + actual_batch * max_blocks, dtype=torch.int32), + slot_mapping=torch.arange(300, 300 + actual_batch, dtype=torch.int32), + ) + + staged = runner._stage_decode_inputs(model, inputs) + args = runner._decode_kernel_args(staged, object(), object()) + + assert staged.token_ids.shape == (actual_batch, 8) + assert staged.seq_lens.shape == (actual_batch,) + assert staged.block_table.shape == (actual_batch * max_blocks,) + assert staged.slot_mapping.shape == (actual_batch,) + assert staged.logits.shape == (actual_batch, model.config.vocab_size) + assert staged.sampled_ids.shape == (actual_batch, 8) + assert staged.next_hidden.shape == (actual_batch, model.config.hidden_size) + assert staged.token_ids[:, 0].tolist() == list(range(10, 10 + actual_batch)) + assert torch.count_nonzero(staged.token_ids[:, 1:]).item() == 0 + assert staged.seq_lens.tolist() == list(range(100, 100 + actual_batch)) + assert staged.block_table.tolist() == list(range(200, 200 + actual_batch * max_blocks)) + assert staged.slot_mapping.tolist() == list(range(300, 300 + actual_batch)) + assert args[6] is staged.seq_lens + assert args[7] is staged.block_table + assert args[8] is staged.slot_mapping + assert args[20] is staged.logits + assert args[22] is staged.token_ids + assert args[23] is staged.sampled_ids + assert args[24] is staged.next_hidden + + +def test_decode_dispatch_rejects_intermediate_batch_size(): + model = _model(max_batch_size=16) + runner = ModelRunner(compiled=_compiled_kernels(model)) + + with pytest.raises(ValueError, match="only batch 1 or the full batch capacity 16"): + runner._stage_decode_inputs( + model, + SimpleNamespace( + actual_batch=3, + token_ids=torch.zeros(3, 1, dtype=torch.int32), + hidden=torch.ones(3, model.config.hidden_size, dtype=torch.bfloat16), + seq_lens=torch.ones(3, dtype=torch.int32), + block_table=torch.zeros(6, dtype=torch.int32), + slot_mapping=torch.zeros(3, dtype=torch.int32), + ), + ) + + +def test_warmup_uses_page_zero_slot_mappings(monkeypatch): + model = _model(max_batch_size=16, page_size=128) + runner = ModelRunner(compiled=_compiled_kernels(model)) + runner._kv_caches[model.config.model_id] = SimpleNamespace( + key_pages=object(), + value_pages=object(), + ) + calls = [] + monkeypatch.setattr( + runner, + "_run_distributed_program", + lambda callable_spec, *args: calls.append(args), + ) + + runner._warmup_dispatch(model.runtime) + + assert len(calls) == 2 + assert torch.count_nonzero(calls[0][13]).item() == 0 + assert torch.count_nonzero(calls[1][8]).item() == 0 + + +@pytest.mark.parametrize("batch_size", [2, 15]) +def test_qwen_serving_rejects_intermediate_batch_before_allocation(batch_size): + model = _model(max_batch_size=16) + executor = PyptoExecutor(platform="a2a3") + + with pytest.raises(ValueError, match="only batch 1 or the full batch capacity 16"): + executor.validate_generate_batch( + SimpleNamespace(runtime=model.runtime), + batch_size, + GenerateConfig(max_new_tokens=1), + ) + + +@pytest.mark.parametrize("platform", ["a2a3sim", "a5"]) +def test_qwen_cann_fai_executor_rejects_non_onboard_a2a3(platform): + with pytest.raises(ValueError, match="supports only A2/A3"): + PyptoExecutor(platform=platform) + + def test_engine_generate_batch_uses_batched_executor_results(): model = _model(max_batch_size=2, eos_token_id=0) manager = KvCacheManager() @@ -577,11 +674,13 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc monkeypatch.setattr(runner, "_compute_kv_cache_pages", lambda config, runtime, device_id=0: 1) monkeypatch.setattr(runner, "_print_memory_breakdown", lambda *a, **kw: None) runner.init_kv_cache(model.config.model_id, model.config, model.runtime) - monkeypatch.setattr(runner, "_static_device_tensor", lambda tensor: tensor) + def run_fake_kernel(callable_spec, *args): + callable_spec.compiled(*(getattr(arg, "tensor", arg) for arg in args)) + monkeypatch.setattr( runner, "_run_distributed_program", - lambda callable_spec, *args: callable_spec.compiled(*args), + run_fake_kernel, ) executor._runners[model.config.model_id] = runner monkeypatch.setattr( @@ -650,6 +749,8 @@ def test_decode_host_inlines_embedding_and_sampling_into_decode_fwd(): assert source.count("decode_fwd(") == 1 assert "token_embed_fwd(" not in source assert "greedy_sample_fwd(" not in source + assert "out, sampled_ids, next_hidden = decode_fwd(" in source + assert "return out, sampled_ids, next_hidden" in source if not QWEN3_KERNEL_DIR.is_dir(): pytest.skip("pypto-lib submodule is not checked out") diff --git a/tests/test_device_sampling_submission.py b/tests/test_device_sampling_submission.py index d8e51ee..972bf66 100644 --- a/tests/test_device_sampling_submission.py +++ b/tests/test_device_sampling_submission.py @@ -32,11 +32,11 @@ def test_device_sampling_is_limited_by_runtime_vocab_size() -> None: dispatch = _source(ROOT / "examples" / "model" / "qwen3_14b" / "runner" / "qwen3_l3_dispatch.py") executor = _source(ROOT / "examples" / "model" / "qwen3_14b" / "runner" / "npu_executor.py") runner = _source(ROOT / "examples" / "model" / "qwen3_14b" / "runner" / "npu_runner.py") - config = _source(QWEN / "config.py") + constants = _source(QWEN / "constants.py") greedy = _source(QWEN / "greedy_sample.py") assert "valid_vocab_size" not in dispatch - assert "REAL_VOCAB = 151936" in config + assert "real_vocab=151936" in constants assert "REAL_VOCAB" in executor assert "lm_head_weight[:1].expand" in executor assert "valid_vocab_size" not in runner diff --git a/tests/test_qwen3_serving.py b/tests/test_qwen3_serving.py index 61145c3..f40ba41 100644 --- a/tests/test_qwen3_serving.py +++ b/tests/test_qwen3_serving.py @@ -17,7 +17,7 @@ output identical to the baseline while the scheduler is observed to actually apply the optimization: -* multi-batch - several concurrent requests co-batched by the scheduler +* full batch - 16 independent requests co-batched by the scheduler * chunked prefill - one prompt split across multiple prefill steps * prefix cache - a repeated long prompt reusing cached prefix blocks @@ -55,7 +55,7 @@ # Kernel requires page_size == 128 (examples/.../npu_executor.py), so a full # prefix-cache block is 128 tokens. Short prompts therefore never populate the -# prefix cache, keeping the multi-batch / chunked-prefill tests isolated from it. +# prefix cache, keeping the full-batch / chunked-prefill tests isolated from it. PAGE_SIZE = 128 MAX_BATCH_SIZE = 16 MAX_SEQ_LEN = 512 @@ -67,6 +67,7 @@ # serving path shares the executor kernels, so greedy output must match. PROMPT = "The capital of France is" SHORT_NEW_TOKENS = 8 +FULL_BATCH_NEW_TOKENS = 2 EXPECTED_TOKEN_IDS = [12095, 13, 3555, 374, 279, 6722, 315, 279] # A prompt long enough to fill at least one 128-token prefix-cache block. @@ -207,8 +208,11 @@ def get_computed_spy(token_ids): asyncio.set_event_loop(None) -def _max_batch_width(events: list[list[tuple[str, bool, int]]]) -> int: - return max((len(event) for event in events), default=0) +def _max_decode_batch_width(events: list[list[tuple[str, bool, int]]]) -> int: + return max( + (sum(not is_prefill for _, is_prefill, _ in event) for event in events), + default=0, + ) def test_serving_single_request_matches_expected_tokens(harness): @@ -224,30 +228,38 @@ def test_serving_single_request_matches_expected_tokens(harness): ) -def test_multi_batch_matches_single_request(harness): - """Several concurrent requests are co-batched and each matches the baseline.""" +def test_full_batch_matches_single_request(harness): + """Sixteen distinct KV histories share one decode batch and match B1.""" + harness.reset() + num_requests = MAX_BATCH_SIZE + prompts = [f"Request {idx}: {PROMPT}" for idx in range(num_requests)] + references = [ + harness.run(_collect(harness.engine, prompt, FULL_BATCH_NEW_TOKENS)) + for prompt in prompts + ] harness.reset() - num_requests = 4 async def _run_all(): return await asyncio.gather( - *(_collect(harness.engine, PROMPT, SHORT_NEW_TOKENS) for _ in range(num_requests)) + *( + _collect(harness.engine, prompt, FULL_BATCH_NEW_TOKENS) + for prompt in prompts + ) ) results = harness.run(_run_all()) - for idx, tokens in enumerate(results): - assert tokens == EXPECTED_TOKEN_IDS, ( - f"Batched request {idx} diverged from the baseline:\n" - f"expected: {EXPECTED_TOKEN_IDS}\n" + for idx, (tokens, expected) in enumerate(zip(results, references, strict=True)): + assert tokens == expected, ( + f"Batched request {idx} diverged from its B1 result:\n" + f"expected: {expected}\n" f"actual: {tokens}" ) - # Prove the scheduler actually co-batched at least two requests in one step. - max_width = _max_batch_width(harness.schedule_events) - assert max_width >= 2, ( - f"Expected the scheduler to co-batch >=2 requests in one step, " - f"but the widest step had {max_width} request(s)." + max_width = _max_decode_batch_width(harness.schedule_events) + assert max_width == MAX_BATCH_SIZE, ( + f"Expected one decode step with {MAX_BATCH_SIZE} independent requests, " + f"but the widest decode step had {max_width} request(s)." )