diff --git a/pypto_serving/config/types.py b/pypto_serving/config/types.py index 7b567f81..8389a9a2 100644 --- a/pypto_serving/config/types.py +++ b/pypto_serving/config/types.py @@ -202,6 +202,14 @@ class SamplingParams: top_k: int | None = None +@dataclass +class SamplingCandidates: + """Device-selected sampling candidates for host-side final sampling.""" + + values: torch.Tensor + token_ids: torch.Tensor + + @dataclass class RequestState: """Mutable per-request state tracked during generation.""" @@ -245,6 +253,7 @@ class PrefillBatch: chunk_offsets: list[int] chunk_starts: list[int] allow_device_greedy_sampling: bool = False + allow_device_topk_sampling: bool = False kv_allocations: list[KvAllocation] = field(default_factory=list) block_ids: list[list[int]] = field(default_factory=list) block_ids_by_group: list[dict[str, list[int]]] = field(default_factory=list) @@ -258,6 +267,7 @@ class PrefillResult: last_hidden: torch.Tensor | None logits: torch.Tensor sampled_token_ids: torch.Tensor | None = None + sampling_candidates: SamplingCandidates | None = None next_hidden_states: torch.Tensor | None = None @@ -273,6 +283,7 @@ class DecodeBatch: hidden_states: torch.Tensor | None seq_lens: torch.Tensor allow_device_greedy_sampling: bool = False + allow_device_topk_sampling: bool = False kv_allocations: list[KvAllocation] = field(default_factory=list) block_ids: list[list[int]] = field(default_factory=list) block_ids_by_group: list[dict[str, list[int]]] = field(default_factory=list) @@ -294,6 +305,7 @@ class DecodeResult: # the logits buffer stays device-resident (never copied back). logits: torch.Tensor | None sampled_token_ids: torch.Tensor | None = None + sampling_candidates: SamplingCandidates | None = None next_hidden_states: torch.Tensor | None = None accepted_token_ids: list[list[int]] | None = None diff --git a/pypto_serving/model/common/executor/executor.py b/pypto_serving/model/common/executor/executor.py index ad226af2..df656ba4 100644 --- a/pypto_serving/model/common/executor/executor.py +++ b/pypto_serving/model/common/executor/executor.py @@ -45,6 +45,11 @@ def supports_device_sampling(self) -> bool: """Return whether executor results may include already-sampled token IDs.""" return False + @property + def device_topk_sampling_k(self) -> int: + """Return the max top-k candidate width the executor can produce on device.""" + return 0 + @property def supports_device_embedding(self) -> bool: """Return whether token embedding can be handled inside the device kernels. diff --git a/pypto_serving/model/common/executor/sampler.py b/pypto_serving/model/common/executor/sampler.py index c373154f..baf350d7 100644 --- a/pypto_serving/model/common/executor/sampler.py +++ b/pypto_serving/model/common/executor/sampler.py @@ -13,7 +13,7 @@ import torch -from pypto_serving.config.types import GenerateConfig, SamplingParams +from pypto_serving.config.types import GenerateConfig, SamplingCandidates, SamplingParams class Sampler: @@ -54,6 +54,56 @@ def sample(self, logits: torch.Tensor, params: SamplingParams) -> int: token = torch.multinomial(probs, num_samples=1) return int(token.item()) + def sample_from_candidates( + self, candidates: SamplingCandidates, row_idx: int, params: SamplingParams + ) -> int: + """Sample from an executor-provided top-k candidate row.""" + value_rows = candidates.values.shape[0] + token_id_rows = candidates.token_ids.shape[0] + if value_rows != token_id_rows: + raise ValueError(f"candidate values/ids row mismatch: {value_rows} != {token_id_rows}") + if row_idx < 0 or row_idx >= value_rows: + raise ValueError(f"row_idx {row_idx} is out of bounds for candidates with {value_rows} rows") + + values = candidates.values[row_idx].float() + token_ids = candidates.token_ids[row_idx].long() + if values.numel() == 0 or token_ids.numel() == 0: + raise ValueError("sampling candidates must not be empty") + if values.numel() != token_ids.numel(): + raise ValueError( + f"candidate values/ids width mismatch: {values.numel()} != {token_ids.numel()}" + ) + + width = values.numel() + if params.top_k is not None and params.top_k > 0: + width = min(width, params.top_k) + values = self._sanitize_logits(values[:width]) + token_ids = token_ids[:width] + if params.temperature <= 0.0: + return int(token_ids[self._greedy_token(values)].item()) + + scaled = values / max(params.temperature, 1e-5) + probs = torch.softmax(scaled, dim=-1) + if not self._is_valid_distribution(probs): + return int(token_ids[self._greedy_token(values)].item()) + + if 0.0 < params.top_p < 1.0: + sorted_probs, sorted_positions = torch.sort(probs, descending=True) + cumulative = torch.cumsum(sorted_probs, dim=-1) + keep = cumulative <= params.top_p + keep[0] = True + filtered_probs = torch.zeros_like(probs) + filtered_probs[sorted_positions[keep]] = probs[sorted_positions[keep]] + total = filtered_probs.sum() + if not torch.isfinite(total) or total.item() <= 0.0: + return int(token_ids[self._greedy_token(values)].item()) + probs = filtered_probs / total + if not self._is_valid_distribution(probs): + return int(token_ids[self._greedy_token(values)].item()) + + sampled_pos = torch.multinomial(probs, num_samples=1) + return int(token_ids[int(sampled_pos.item())].item()) + @staticmethod def from_generate_config(config: GenerateConfig) -> SamplingParams: """Build sampler parameters from user-facing generation config.""" diff --git a/pypto_serving/model/qwen/npu_executor.py b/pypto_serving/model/qwen/npu_executor.py index 75b341ec..f42a17bb 100644 --- a/pypto_serving/model/qwen/npu_executor.py +++ b/pypto_serving/model/qwen/npu_executor.py @@ -38,6 +38,7 @@ _VOCAB_PAD_MULTIPLE = 512 # must be a multiple of lm_head.VOCAB_CHUNK (64) _QWEN14B_PAGE_SIZE = 128 +_QWEN14B_TOPK_SELECT_K = 32 def _kernel_batch_pad(kernel_module: object) -> tuple[int, bool]: @@ -55,8 +56,8 @@ def _kernel_batch_pad(kernel_module: object) -> tuple[int, bool]: greedy_sample carries the new spelling but is still fixed-batch. Callers must apply their own stage's rule. The distinction matters because a pre-rename decode writes exactly ``BATCH`` rows whatever the runtime batch, - so pointing it at buffers sized for a smaller ``max_batch_size`` would - overrun them. + while topk_select remains fixed-batch. Pointing either stage at undersized + buffers would overrun them. """ value = getattr(kernel_module, "BATCH_PAD", None) if value is not None: @@ -171,6 +172,11 @@ def supports_device_sampling(self) -> bool: """Qwen3 NPU runner can return greedy sampled token ids.""" return True + @property + def device_topk_sampling_k(self) -> int: + """Qwen3 NPU runner can return top-k sampling candidates.""" + return _QWEN14B_TOPK_SELECT_K + @property def supports_device_embedding(self) -> bool: """Qwen3 NPU prefill and decode embed token ids inside device kernels.""" @@ -195,10 +201,10 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: # device-resident paged KV pool prefill writes (self._kv_caches), so no # contiguous bridge / MAX_SEQ env is needed. qwen3_decode_fwd = _load_pypto_lib_qwen14b_module("decode_fwd", kernel_dir) - qwen3_greedy_sample = _load_pypto_lib_qwen14b_module("greedy_sample", kernel_dir) + qwen3_topk_select = _load_pypto_lib_qwen14b_module("topk_select", kernel_dir) qwen3_l3_dispatch.prefill_fwd = qwen3_prefill_fwd.prefill_fwd qwen3_l3_dispatch.decode_fwd = qwen3_decode_fwd.decode_fwd - qwen3_l3_dispatch.greedy_sample_fwd = qwen3_greedy_sample.greedy_sample_fwd + qwen3_l3_dispatch.topk_select_fwd = qwen3_topk_select.topk_select_fwd self._validate_supported_shape(model) kernel_batch = model.runtime.max_batch_size @@ -254,22 +260,31 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: f"but the runtime model vocab_size is {model.config.vocab_size}; expected " f"{int(qwen3_decode_fwd.REAL_VOCAB)}." ) - # greedy_sample_fwd is still a STATIC-batch stage (it was not converted - # alongside decode), so this stays an exact match rather than a bound. - greedy_sample_batch, _ = _kernel_batch_pad(qwen3_greedy_sample) - if greedy_sample_batch != kernel_batch: + # topk_select_fwd remains a fixed-batch stage, so require an exact + # match while accepting either pypto-lib batch constant spelling. + topk_select_batch, _ = _kernel_batch_pad(qwen3_topk_select) + if topk_select_batch != kernel_batch: raise ValueError( - "greedy_sample_fwd is compiled for a fixed kernel batch of " - f"{greedy_sample_batch}, but runtime max_batch_size is {kernel_batch}." + "topk_select_fwd is compiled for a fixed kernel BATCH of " + f"{topk_select_batch}, but runtime max_batch_size is {kernel_batch}." ) - if int(qwen3_greedy_sample.VOCAB) != padded_vocab: + if int(qwen3_topk_select.VOCAB) != padded_vocab: raise ValueError( - "greedy_sample_fwd VOCAB must match the padded logits vocab: " - f"{int(qwen3_greedy_sample.VOCAB)} != {padded_vocab}." + "topk_select_fwd VOCAB must match the padded logits vocab: " + f"{int(qwen3_topk_select.VOCAB)} != {padded_vocab}." ) - sampled_ids_width = int( - getattr(qwen3_decode_fwd, "SAMPLED_IDS_PAD", getattr(qwen3_greedy_sample, "SAMPLED_IDS_PAD", 1)) - ) + if model.config.vocab_size != int(qwen3_topk_select.REAL_VOCAB): + raise ValueError( + "topk_select_fwd REAL_VOCAB must match model vocab_size: " + f"{int(qwen3_topk_select.REAL_VOCAB)} != {model.config.vocab_size}." + ) + topk_width = int(qwen3_topk_select.TOPK) + if topk_width != _QWEN14B_TOPK_SELECT_K: + raise ValueError( + "topk_select_fwd TOPK must match executor capability: " + f"{topk_width} != {_QWEN14B_TOPK_SELECT_K}." + ) + sampled_ids_width = int(getattr(qwen3_decode_fwd, "SAMPLED_IDS_PAD", 1)) page_size = model.runtime.page_size max_blocks_per_seq = (model.runtime.max_seq_len + page_size - 1) // page_size prefill = self._compile_prefill_fwd_callable( @@ -302,10 +317,10 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: page_size=page_size, sampled_ids_width=sampled_ids_width, ) - greedy_sample = self._compile_greedy_sample_callable( - qwen3_l3_dispatch.qwen3_greedy_sample_host, + topk_select = self._compile_topk_select_callable( + qwen3_l3_dispatch.qwen3_topk_select_host, batch=kernel_batch, - sampled_ids_width=sampled_ids_width, + topk_width=topk_width, vocab_size=padded_vocab, ) rope_cos_raw, rope_sin_raw = rope_tables( @@ -353,13 +368,13 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: (kernel_batch, padded_vocab), dtype=torch.float32, ).share_memory_() - prefill_sampled_ids_buffer = torch.empty( - (kernel_batch, sampled_ids_width), - dtype=torch.int32, + prefill_topk_values_buffer = torch.empty( + (kernel_batch, topk_width), + dtype=torch.float32, ).share_memory_() - prefill_next_hidden_buffer = torch.empty( - (kernel_batch, model.config.hidden_size), - dtype=torch.bfloat16, + prefill_topk_indices_buffer = torch.empty( + (kernel_batch, topk_width), + dtype=torch.int32, ).share_memory_() decode_logits_buffer = torch.empty( (kernel_batch, padded_vocab), @@ -379,6 +394,15 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: (kernel_batch, sampled_ids_width), dtype=torch.int32, ).share_memory_() + decode_topk_values_buffer = torch.empty( + (kernel_batch, topk_width), + dtype=torch.float32, + ).share_memory_() + decode_topk_indices_buffer = torch.empty( + (kernel_batch, topk_width), + dtype=torch.int32, + ).share_memory_() + sampling_control_buffer = torch.empty((2,), dtype=torch.int32).share_memory_() decode_next_hidden_buffer = torch.empty( (kernel_batch, model.config.hidden_size), dtype=torch.bfloat16, @@ -386,7 +410,7 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: return _CompiledKernels( prefill=prefill, decode=decode, - greedy_sample=greedy_sample, + topk_select=topk_select, final_norm_weight=final_norm_weight, rope_cos=rope_cos, rope_sin=rope_sin, @@ -401,14 +425,17 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: prefill_block_table_buffer=prefill_block_table_buffer, prefill_slot_mapping_buffer=prefill_slot_mapping_buffer, prefill_logits_buffer=prefill_logits_buffer, - prefill_sampled_ids_buffer=prefill_sampled_ids_buffer, - prefill_next_hidden_buffer=prefill_next_hidden_buffer, + prefill_topk_values_buffer=prefill_topk_values_buffer, + prefill_topk_indices_buffer=prefill_topk_indices_buffer, decode_seq_lens_buffer=decode_seq_lens_buffer, decode_block_table_buffer=decode_block_table_buffer, decode_slot_mapping_buffer=decode_slot_mapping_buffer, decode_logits_buffer=decode_logits_buffer, decode_token_ids_buffer=decode_token_ids_buffer, decode_sampled_ids_buffer=decode_sampled_ids_buffer, + decode_topk_values_buffer=decode_topk_values_buffer, + decode_topk_indices_buffer=decode_topk_indices_buffer, + sampling_control_buffer=sampling_control_buffer, decode_next_hidden_buffer=decode_next_hidden_buffer, ) @@ -528,20 +555,22 @@ def _compile_decode_fwd_callable( ] return self._compile_jit_fwd_callable("decode_fwd", jit_fn, dummy_args) - def _compile_greedy_sample_callable( + def _compile_topk_select_callable( self, jit_fn: object, *, batch: int, - sampled_ids_width: int, + topk_width: int, vocab_size: int, ) -> _L3Callable: - """Compile the greedy sampling HOST wrapper.""" + """Compile the top-k candidate selection HOST wrapper.""" dummy_args = [ torch.empty((batch, vocab_size), dtype=torch.float32), - torch.empty((batch, sampled_ids_width), dtype=torch.int32), + torch.empty((2,), dtype=torch.int32), + torch.empty((batch, topk_width), dtype=torch.float32), + torch.empty((batch, topk_width), dtype=torch.int32), ] - return self._compile_jit_fwd_callable("greedy_sample_fwd", jit_fn, dummy_args) + return self._compile_jit_fwd_callable("topk_select_fwd", jit_fn, dummy_args) def _compile_jit_fwd_callable( self, diff --git a/pypto_serving/model/qwen/npu_runner.py b/pypto_serving/model/qwen/npu_runner.py index c37471bd..9432c966 100644 --- a/pypto_serving/model/qwen/npu_runner.py +++ b/pypto_serving/model/qwen/npu_runner.py @@ -27,6 +27,7 @@ PrefillResult, RuntimeConfig, RuntimeModel, + SamplingCandidates, ) from pypto_serving.model.common.runner.model_runner import ModelRunner from pypto_serving.tools.profile import profile_span @@ -88,7 +89,7 @@ class _CompiledKernels: prefill: _L3Callable decode: _L3Callable - greedy_sample: _L3Callable + topk_select: _L3Callable final_norm_weight: torch.Tensor rope_cos: torch.Tensor rope_sin: torch.Tensor @@ -103,14 +104,17 @@ class _CompiledKernels: prefill_block_table_buffer: torch.Tensor prefill_slot_mapping_buffer: torch.Tensor prefill_logits_buffer: torch.Tensor - prefill_sampled_ids_buffer: torch.Tensor - prefill_next_hidden_buffer: torch.Tensor + prefill_topk_values_buffer: torch.Tensor + prefill_topk_indices_buffer: torch.Tensor decode_seq_lens_buffer: torch.Tensor decode_block_table_buffer: torch.Tensor decode_slot_mapping_buffer: torch.Tensor decode_logits_buffer: torch.Tensor decode_token_ids_buffer: torch.Tensor decode_sampled_ids_buffer: torch.Tensor + decode_topk_values_buffer: torch.Tensor + decode_topk_indices_buffer: torch.Tensor + sampling_control_buffer: torch.Tensor decode_next_hidden_buffer: torch.Tensor @@ -277,7 +281,6 @@ def _alloc_kv_cache_with_retry( f"(requested {requested}, downgraded after OOM): " f"{num_pages * bytes_per_page / 1e9:.2f} GB KV cache, " f"{num_pages * runtime.page_size} context tokens", - ) return num_pages except (RuntimeError, MemoryError) as e: @@ -369,7 +372,6 @@ def _print_memory_breakdown( logger.info( f" total used (measured): {used_bytes / 1e9:7.2f} GB " f"/ {total_bytes / 1e9:.2f} GB (free {free_bytes / 1e9:.2f} GB)", - ) logger.info(f" ├─ weights (estimated): {weight_bytes / 1e9:7.2f} GB") kv_tokens = num_pages * runtime.page_size @@ -379,7 +381,6 @@ def _print_memory_breakdown( logger.info( f" ├─ KV cache ({num_pages} pages): {kv_bytes / 1e9:7.2f} GB " f"({bytes_per_page / 1e6:.1f} MB/page)", - ) logger.info( f" │ capacity = {kv_tokens} tokens " @@ -387,18 +388,15 @@ def _print_memory_breakdown( f"worst-case need {runtime.max_batch_size}x{max_seq_len}=" f"{worst_case_demand} tokens" + (" [OK]" if kv_tokens >= worst_case_demand else " [TIGHT]"), - ) logger.info(f" ├─ simpler arena (env x 4): {arena_bytes / 1e9:7.2f} GB") logger.info( f" └─ residual (buffers/scratch): {residual / 1e9:6.2f} GB " f"(compiled buffers + transient activation scratch + overhead)", - ) logger.info( " note: weights/arena are estimates, KV is exact; total is from " "mem_get_info (may under-count simpler's rtMalloc pool).", - ) def warmup(self, model: RuntimeModel) -> None: @@ -426,7 +424,6 @@ 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)", - ) compiled = self._compiled kv_cache = list(self._kv_caches.values())[0] @@ -541,14 +538,16 @@ def _iter_static_host_tensors(self) -> tuple[torch.Tensor, ...]: compiled.prefill_block_table_buffer, compiled.prefill_slot_mapping_buffer, compiled.prefill_logits_buffer, - compiled.prefill_sampled_ids_buffer, - compiled.prefill_next_hidden_buffer, + compiled.prefill_topk_values_buffer, + compiled.prefill_topk_indices_buffer, compiled.decode_seq_lens_buffer, compiled.decode_block_table_buffer, compiled.decode_slot_mapping_buffer, compiled.decode_logits_buffer, compiled.decode_token_ids_buffer, compiled.decode_sampled_ids_buffer, + compiled.decode_topk_values_buffer, + compiled.decode_topk_indices_buffer, compiled.decode_next_hidden_buffer, ) @@ -592,18 +591,24 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult for batch_idx, alloc in enumerate(batch.kv_allocations): seq_len = batch.seq_lens[batch_idx] alloc.tokens_used = max(alloc.tokens_used, seq_len) - sampled_ids, next_hidden = self._maybe_run_sample_embed( + sampled_ids = self._maybe_run_greedy_sample( logits_padded, - compiled.prefill_sampled_ids_buffer, - compiled.prefill_next_hidden_buffer, prefill_inputs.actual_batch, allow=batch.allow_device_greedy_sampling, ) + sampling_candidates = self._device_topk_outputs( + logits_padded, + compiled.prefill_topk_values_buffer, + compiled.prefill_topk_indices_buffer, + prefill_inputs.actual_batch, + allow=batch.allow_device_topk_sampling, + ) return PrefillResult( last_hidden=None, logits=logits_padded[: prefill_inputs.actual_batch, : model.config.vocab_size], sampled_token_ids=sampled_ids, - next_hidden_states=next_hidden, + sampling_candidates=sampling_candidates, + next_hidden_states=None, ) def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: @@ -632,10 +637,20 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: k_cache = kv_cache.key_pages v_cache = kv_cache.value_pages - device_greedy = batch.allow_device_greedy_sampling + device_sampling = ( + batch.allow_device_greedy_sampling or batch.allow_device_topk_sampling + ) + selector_logits = ( + self._decode_logits_device_arg() if device_sampling else kernel_inputs.logits + ) self._run_distributed_program( compiled.decode, - *self._decode_kernel_args(kernel_inputs, k_cache, v_cache, device_greedy=device_greedy), + *self._decode_kernel_args( + kernel_inputs, + k_cache, + v_cache, + device_sampling=device_sampling, + ), ) for batch_idx, alloc in enumerate(batch.kv_allocations): alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) @@ -649,16 +664,24 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: kernel_inputs.actual_batch, allow=batch.allow_device_greedy_sampling, ) + sampling_candidates = self._device_topk_outputs( + selector_logits, + compiled.decode_topk_values_buffer, + compiled.decode_topk_indices_buffer, + kernel_inputs.actual_batch, + allow=batch.allow_device_topk_sampling, + ) return DecodeResult( hidden_states=None, - # Device-greedy path: the host consumes sampled_ids, never logits, so we - # keep the logits buffer device-resident and skip its ~9.7MB D2H copy-back. + # Device sampling consumes only sampled ids or top-k candidates, so + # full-vocabulary logits remain device-resident. logits=( None - if device_greedy + if device_sampling else kernel_inputs.logits[: kernel_inputs.actual_batch, : model.config.vocab_size].cpu() ), sampled_token_ids=sampled_ids, + sampling_candidates=sampling_candidates, next_hidden_states=next_hidden, ) @@ -683,27 +706,69 @@ def _integrated_sample_result( next_hidden, ) - def _maybe_run_sample_embed( + def _maybe_run_greedy_sample( self, logits: torch.Tensor, - sampled_ids_buffer: torch.Tensor, - next_hidden_buffer: torch.Tensor, actual_batch: int, *, allow: bool, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """Run device greedy sampling when the request is greedy.""" + ) -> torch.Tensor | None: + """Run the sampling selector in exact greedy mode for prefill.""" if not allow: - return None, None + return None compiled = self._compiled - self._run_distributed_program( - compiled.greedy_sample, + self._run_sampling_selector( logits, - sampled_ids_buffer, + compiled.prefill_topk_values_buffer, + compiled.prefill_topk_indices_buffer, + actual_batch, + selection_k=1, ) - return ( - sampled_ids_buffer[:actual_batch, :1].clone(), - None, + return compiled.prefill_topk_indices_buffer[:actual_batch, :1].clone() + + def _device_topk_outputs( + self, + logits: torch.Tensor | DeviceTensor, + values_buffer: torch.Tensor, + indices_buffer: torch.Tensor, + actual_batch: int, + *, + allow: bool, + ) -> SamplingCandidates | None: + """Run device top-k candidate selection and return small host tensors.""" + if not allow: + return None + self._run_sampling_selector( + logits, + values_buffer, + indices_buffer, + actual_batch, + selection_k=indices_buffer.shape[1], + ) + return SamplingCandidates( + values=values_buffer[:actual_batch].clone(), + token_ids=indices_buffer[:actual_batch].clone(), + ) + + def _run_sampling_selector( + self, + logits: torch.Tensor | DeviceTensor, + values_buffer: torch.Tensor, + indices_buffer: torch.Tensor, + actual_batch: int, + *, + selection_k: int, + ) -> None: + """Run the shared greedy/top-k selector without adding another worker program.""" + control = self._compiled.sampling_control_buffer + control[0] = int(actual_batch) + control[1] = int(selection_k) + self._run_distributed_program( + self._compiled.topk_select, + logits, + control, + values_buffer, + indices_buffer, ) def _prefill_kernel_args( @@ -750,19 +815,21 @@ def _decode_kernel_args( k_cache: DeviceTensor, v_cache: DeviceTensor, *, - device_greedy: bool = False, + device_sampling: bool = False, ) -> tuple[Any, ...]: """Return arguments in ``qwen3_decode_host`` signature order. - On ``device_greedy`` the logits + next_hidden outputs are passed as + On device greedy or top-k sampling, logits + next_hidden are passed as worker-resident (device) tensors so they are never staged/copied-back - per step (no memset, no D2H); only the tiny sampled_ids stays host-visible. + per step. Only sampled ids or top-k candidates remain host-visible. """ static = self._require_static_args() weights = static.decode_weights - logits_arg = self._decode_logits_device_arg() if device_greedy else inputs.logits + logits_arg = self._decode_logits_device_arg() if device_sampling else inputs.logits next_hidden_arg = ( - self._decode_next_hidden_device_arg() if device_greedy else self._compiled.decode_next_hidden_buffer + self._decode_next_hidden_device_arg() + if device_sampling + else self._compiled.decode_next_hidden_buffer ) return ( weights["decode_input_rms_weight"], @@ -834,7 +901,7 @@ def _store_kernel_binaries(self) -> None: """ if self._kernel_cache is None: return - for spec in (self._compiled.prefill, self._compiled.decode, self._compiled.greedy_sample): + for spec in (self._compiled.prefill, self._compiled.decode, self._compiled.topk_select): self._kernel_cache.store(spec.name, spec.compiled, spec.params_fingerprint) def _shared_l3_worker(self) -> Any: @@ -846,7 +913,7 @@ def _shared_l3_worker(self) -> Any: worker = DistributedWorker([ self._compiled.prefill.compiled, self._compiled.decode.compiled, - self._compiled.greedy_sample.compiled, + self._compiled.topk_select.compiled, ]) self._l3_worker = worker return worker @@ -865,7 +932,7 @@ def _coerce_l3_arg(self, worker: Any, arg: Any) -> Any: return dev def _decode_logits_device_arg(self) -> DeviceTensor: - """Device-resident decode logits scratch (greedy path: never copied back). + """Device-resident decode logits scratch for device sampling. Allocated directly on the worker and left uninitialized — the fused decode kernel writes every max_batch row before the on-device sampler reads it — so diff --git a/pypto_serving/model/qwen/qwen3_l3_dispatch.py b/pypto_serving/model/qwen/qwen3_l3_dispatch.py index 5bb123db..c3513038 100644 --- a/pypto_serving/model/qwen/qwen3_l3_dispatch.py +++ b/pypto_serving/model/qwen/qwen3_l3_dispatch.py @@ -16,7 +16,7 @@ prefill_fwd = None decode_fwd = None -greedy_sample_fwd = None +topk_select_fwd = None @pl.jit.host @@ -135,11 +135,15 @@ def qwen3_decode_host( @pl.jit.host -def qwen3_greedy_sample_host( +def qwen3_topk_select_host( logits: pl.Tensor, - sampled_ids: pl.Out[pl.Tensor], -) -> pl.Tensor: - return greedy_sample_fwd( + sampling_control: pl.Tensor, + topk_values: pl.Out[pl.Tensor], + topk_indices: pl.Out[pl.Tensor], +) -> tuple[pl.Tensor, pl.Tensor]: + return topk_select_fwd( logits, - sampled_ids, + sampling_control, + topk_values, + topk_indices, ) diff --git a/pypto_serving/serving/engine/engine.py b/pypto_serving/serving/engine/engine.py index fba09950..bdee7559 100644 --- a/pypto_serving/serving/engine/engine.py +++ b/pypto_serving/serving/engine/engine.py @@ -21,6 +21,7 @@ ModelRecord, RequestState, RuntimeConfig, + SamplingCandidates, ) from pypto_serving.model.common.executor.executor import ModelExecutor from pypto_serving.model.common.executor.sampler import Sampler @@ -180,6 +181,7 @@ def _generate_batch_impl( and self._executor.supports_device_sampling and self._executor.supports_device_embedding ) + allow_device_topk_sampling = self._allow_device_topk_sampling(generate_config) embedding_lookup = None if not self._executor.supports_device_embedding: embedding_lookup = lambda token_ids: self._executor.lookup_embeddings( @@ -197,6 +199,8 @@ def _generate_batch_impl( offsets = [0] * len(prompts) completed_logits: dict[int, torch.Tensor] = {} completed_sampled_ids: dict[int, torch.Tensor] = {} + completed_candidate_values: dict[int, torch.Tensor] = {} + completed_candidate_ids: dict[int, torch.Tensor] = {} with self._executor.session(): while remaining: # greedily take up to total_budget tokens @@ -228,6 +232,7 @@ def _generate_batch_impl( device=runtime_model.runtime.device, embedding_lookup=embedding_lookup, allow_device_greedy_sampling=allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, kv_allocations=[allocations[ri] for ri in chunk_req], ) prefill_result = self._executor.run_prefill(runtime_model, sub_batch) @@ -244,6 +249,13 @@ def _generate_batch_impl( ).clone() if sampled_ids is not None: completed_sampled_ids[ri] = sampled_ids[row].clone() + if prefill_result.sampling_candidates is not None: + completed_candidate_values[ri] = ( + prefill_result.sampling_candidates.values[row].clone() + ) + completed_candidate_ids[ri] = ( + prefill_result.sampling_candidates.token_ids[row].clone() + ) # remove completed requests from the pool for i in range(len(remaining) - 1, -1, -1): ri = remaining[i][0] @@ -263,6 +275,26 @@ def _generate_batch_impl( if allow_device_greedy_sampling and len(completed_sampled_ids) == len(requests) else None ) + prefill_sampling_candidates = ( + SamplingCandidates( + values=torch.stack( + [ + completed_candidate_values[request_idx] + for request_idx in range(len(requests)) + ] + ), + token_ids=torch.stack( + [ + completed_candidate_ids[request_idx] + for request_idx in range(len(requests)) + ] + ), + ) + if allow_device_topk_sampling + and len(completed_candidate_values) == len(requests) + and len(completed_candidate_ids) == len(requests) + else None + ) else: prefill_batch = pack_prefill_batch( request_ids=[request.request_id for request in requests], @@ -272,6 +304,7 @@ def _generate_batch_impl( device=runtime_model.runtime.device, embedding_lookup=embedding_lookup, allow_device_greedy_sampling=allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, kv_allocations=allocations, ) fast_path_result = self._executor.try_generate_batch( @@ -293,6 +326,7 @@ def _generate_batch_impl( if allow_device_greedy_sampling else None ) + prefill_sampling_candidates = prefill_result.sampling_candidates sampling_params = self._sampler.from_generate_config(generate_config) current_tokens = self._sample_batch_rows( @@ -300,6 +334,8 @@ def _generate_batch_impl( sampling_params, len(requests), prefill_sampled_token_ids, + prefill_sampling_candidates, + allow_device_topk_sampling=allow_device_topk_sampling, ) active_indices = list(range(len(requests))) finish_reasons = ["length"] * len(requests) @@ -361,6 +397,7 @@ def _generate_batch_impl( device=runtime_model.runtime.device, ), allow_device_greedy_sampling=allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, kv_allocations=active_allocations, ), ) @@ -369,6 +406,8 @@ def _generate_batch_impl( sampling_params, len(next_active), decode_result.sampled_token_ids if allow_device_greedy_sampling else None, + decode_result.sampling_candidates, + allow_device_topk_sampling=allow_device_topk_sampling, ) for row_idx, request_idx in enumerate(next_active): current_tokens[request_idx] = decoded_tokens[row_idx] @@ -496,6 +535,8 @@ def _sample_batch_rows( sampling_params, row_count: int, sampled_token_ids: torch.Tensor | None = None, + sampling_candidates: SamplingCandidates | None = None, + allow_device_topk_sampling: bool = False, ) -> list[int]: """Return sampled token IDs, preferring executor-provided device samples.""" if sampled_token_ids is not None: @@ -505,6 +546,20 @@ def _sample_batch_rows( f"sampled_token_ids has {flat_ids.numel()} rows, expected at least {row_count}" ) return [int(flat_ids[idx].item()) for idx in range(row_count)] + if allow_device_topk_sampling and sampling_candidates is not None: + if sampling_candidates.values.shape[0] < row_count: + raise ValueError( + "sampling candidates have fewer rows than expected: " + f"{sampling_candidates.values.shape[0]} < {row_count}" + ) + return [ + self._sampler.sample_from_candidates( + sampling_candidates, + row_idx, + sampling_params, + ) + for row_idx in range(row_count) + ] return [ self._sampler.sample( self._select_batch_row(logits, row_idx), @@ -513,6 +568,16 @@ def _sample_batch_rows( for row_idx in range(row_count) ] + def _allow_device_topk_sampling(self, generate_config: GenerateConfig) -> bool: + """Return whether generation can use executor-provided top-k candidates.""" + max_device_topk = self._executor.device_topk_sampling_k + return ( + generate_config.temperature > 0.0 + and generate_config.top_k is not None + and generate_config.top_k > 0 + and max_device_topk >= generate_config.top_k + ) + def _decode_embeddings_from_cache_or_lookup( self, runtime_model, diff --git a/pypto_serving/serving/server/serving_worker.py b/pypto_serving/serving/server/serving_worker.py index 07dae884..e88b0afa 100644 --- a/pypto_serving/serving/server/serving_worker.py +++ b/pypto_serving/serving/server/serving_worker.py @@ -299,6 +299,7 @@ def _batch_prefill( self.executor.supports_device_sampling and all(self._req_cache[pr.request_id].temperature <= 0.0 for pr in scheduled) ) + allow_device_topk_sampling = self._allow_device_topk_sampling(scheduled) embedding_lookup = None if not self.executor.supports_device_embedding: embedding_lookup = lambda token_ids: self.executor.lookup_embeddings( @@ -315,6 +316,7 @@ def _batch_prefill( device=device, embedding_lookup=embedding_lookup, allow_device_greedy_sampling=allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, block_ids=block_ids_list, block_ids_by_group=[pr.block_ids_by_group for pr in scheduled], cache_partitions=[pr.cache_partition for pr in scheduled], @@ -338,7 +340,12 @@ def _batch_prefill( top_k=cached.top_k, ) token_id = self._sample_result_row( - prefill_result, logits, params, i, allow_device_greedy_sampling + prefill_result, + logits, + params, + i, + allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, ) new_tokens[pr.request_id] = [token_id] @@ -388,6 +395,7 @@ def _batch_decode( self.executor.supports_device_sampling and all(self._req_cache[dr.request_id].temperature <= 0.0 for dr in scheduled) ) + allow_device_topk_sampling = self._allow_device_topk_sampling(scheduled) decode_tokens = [self._resolve_decode_token(dr) for dr in scheduled] prev_tokens = [self._resolve_prev_token(dr) for dr in scheduled] @@ -416,6 +424,7 @@ def _batch_decode( hidden_states=decode_embeddings, seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), allow_device_greedy_sampling=allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, block_ids=block_ids_list, block_ids_by_group=[dr.block_ids_by_group for dr in scheduled], cache_partitions=[dr.cache_partition for dr in scheduled], @@ -442,7 +451,12 @@ def _batch_decode( top_k=cached.top_k, ) token_id = self._sample_result_row( - decode_result, logits, params, i, allow_device_greedy_sampling + decode_result, + logits, + params, + i, + allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, ) new_tokens[dr.request_id] = [token_id] @@ -464,6 +478,7 @@ def _sample_result_row( params: SamplingParams, row_idx: int, allow_device_sampled: bool, + allow_device_topk_sampling: bool, ) -> int: """Return a sampled token from executor output, falling back to host sampling.""" sampled = getattr(result, "sampled_token_ids", None) @@ -474,8 +489,23 @@ def _sample_result_row( f"sampled_token_ids has {flat.numel()} rows, expected row {row_idx}" ) return int(flat[row_idx].item()) + candidates = getattr(result, "sampling_candidates", None) + if allow_device_topk_sampling and candidates is not None: + return self.sampler.sample_from_candidates(candidates, row_idx, params) return self.sampler.sample(logits, params) + def _allow_device_topk_sampling(self, scheduled: list) -> bool: + """Return whether a scheduled batch can use executor top-k candidates.""" + max_device_topk = self.executor.device_topk_sampling_k + cached_requests = [self._req_cache[item.request_id] for item in scheduled] + return ( + max_device_topk > 0 + and all(request.temperature > 0.0 for request in cached_requests) + and all(request.top_k is not None for request in cached_requests) + and all(request.top_k > 0 for request in cached_requests) + and all(request.top_k <= max_device_topk for request in cached_requests) + ) + def _worker_entry( config: EngineConfig, diff --git a/pypto_serving/serving/utils/prefill.py b/pypto_serving/serving/utils/prefill.py index 2bd7d23c..c680003a 100644 --- a/pypto_serving/serving/utils/prefill.py +++ b/pypto_serving/serving/utils/prefill.py @@ -25,6 +25,7 @@ def pack_prefill_batch( device: str | torch.device, embedding_lookup: Callable[[torch.Tensor], torch.Tensor] | None = None, allow_device_greedy_sampling: bool = False, + allow_device_topk_sampling: bool = False, kv_allocations: Sequence[KvAllocation] = (), block_ids: Sequence[Sequence[int]] = (), block_ids_by_group: Sequence[dict[str, list[int]]] = (), @@ -55,6 +56,7 @@ def pack_prefill_batch( chunk_offsets=chunk_offsets, chunk_starts=list(chunk_starts), allow_device_greedy_sampling=allow_device_greedy_sampling, + allow_device_topk_sampling=allow_device_topk_sampling, kv_allocations=list(kv_allocations), block_ids=[list(row) for row in block_ids], block_ids_by_group=list(block_ids_by_group), diff --git a/tests/test_batching.py b/tests/test_batching.py index c88466e3..2a93426c 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -29,12 +29,16 @@ PrefillResult, RuntimeConfig, RuntimeModel, + SamplingCandidates, + SamplingParams, ) from pypto_serving.model.common.executor.executor import ModelExecutor +from pypto_serving.model.common.executor.sampler import Sampler from pypto_serving.model.qwen.kernel_cache import compute_params_fingerprint from pypto_serving.model.qwen.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor from pypto_serving.model.qwen.npu_runner import ( _CompiledKernels, + _DecodeKernelInputs, _L3Callable, Qwen314BModelRunner as ModelRunner, _add_run_timing_args, @@ -77,6 +81,18 @@ QWEN3_KERNEL_DIR = ROOT / "pypto-lib" / "models" / "qwen3" / "14b" +@pytest.mark.parametrize("row_idx", [-1, 1]) +def test_candidate_sampling_rejects_out_of_bounds_row(row_idx): + candidates = SamplingCandidates( + values=torch.tensor([[4.0, 3.0]], dtype=torch.float32), + token_ids=torch.tensor([[7, 6]], dtype=torch.int32), + ) + params = SamplingParams(temperature=0.8, top_p=1.0, top_k=2) + + with pytest.raises(ValueError, match=rf"row_idx {row_idx} is out of bounds"): + Sampler().sample_from_candidates(candidates, row_idx, params) + + class _Tokenizer: def encode(self, text: str) -> list[int]: return [max(1, len(text))] @@ -577,7 +593,7 @@ def _compiled_kernels( return _CompiledKernels( prefill=callable_, decode=callable_, - greedy_sample=callable_, + topk_select=callable_, final_norm_weight=torch.ones(1, hidden_size), rope_cos=torch.zeros(max_seq, head_dim), rope_sin=torch.zeros(max_seq, head_dim), @@ -592,14 +608,17 @@ def _compiled_kernels( prefill_block_table_buffer=torch.empty(kernel_batch * max_blocks, dtype=torch.int32), prefill_slot_mapping_buffer=torch.empty(kernel_batch * max_seq, dtype=torch.int32), prefill_logits_buffer=torch.empty(kernel_batch, model.config.vocab_size), - prefill_sampled_ids_buffer=torch.empty(kernel_batch, sampled_ids_width, dtype=torch.int32), - prefill_next_hidden_buffer=torch.empty(kernel_batch, hidden_size, dtype=torch.bfloat16), + prefill_topk_values_buffer=torch.empty(kernel_batch, 4, dtype=torch.float32), + prefill_topk_indices_buffer=torch.empty(kernel_batch, 4, dtype=torch.int32), decode_seq_lens_buffer=torch.zeros(kernel_batch, dtype=torch.int32), 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, sampled_ids_width, dtype=torch.int32), decode_sampled_ids_buffer=torch.empty(kernel_batch, sampled_ids_width, dtype=torch.int32), + decode_topk_values_buffer=torch.empty(kernel_batch, 4, dtype=torch.float32), + decode_topk_indices_buffer=torch.empty(kernel_batch, 4, dtype=torch.int32), + sampling_control_buffer=torch.empty(2, dtype=torch.int32), decode_next_hidden_buffer=torch.empty(kernel_batch, hidden_size, dtype=torch.bfloat16), ) @@ -771,6 +790,76 @@ def prepare(seq_len: int): assert prepared.block_table.tolist() == alloc.page_ids +def test_decode_topk_selects_from_device_resident_logits(monkeypatch): + model = _model(max_batch_size=1) + compiled = _compiled_kernels(model) + compiled.decode = _L3Callable( + compiled=object(), + name="decode", + aicpu_thread_num=1, + ) + compiled.topk_select = _L3Callable( + compiled=object(), + name="topk_select", + aicpu_thread_num=1, + ) + runner = ModelRunner(compiled=compiled) + host_logits = torch.zeros(1, model.config.vocab_size) + kernel_inputs = _DecodeKernelInputs( + actual_batch=1, + 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=host_logits, + ) + runner._kv_caches = { + model.config.model_id: SimpleNamespace( + key_pages=object(), + value_pages=object(), + ) + } + device_logits = object() + device_next_hidden = object() + monkeypatch.setattr( + runner, + "_prepare_decode_inputs", + lambda _model, _batch: kernel_inputs, + ) + monkeypatch.setattr(runner, "_decode_logits_device_arg", lambda: device_logits) + monkeypatch.setattr( + runner, + "_decode_next_hidden_device_arg", + lambda: device_next_hidden, + ) + dispatches = [] + monkeypatch.setattr( + runner, + "_run_distributed_program", + lambda callable_spec, *args: dispatches.append((callable_spec, args)), + ) + + result = runner.run_decode( + model, + DecodeBatch( + request_ids=["request"], + token_ids=torch.tensor([[7]], dtype=torch.long), + hidden_states=None, + seq_lens=torch.tensor([1], dtype=torch.int32), + allow_device_topk_sampling=True, + block_ids=[[0]], + ), + ) + + assert result.logits is None + assert result.sampling_candidates is not None + assert dispatches[0][0] is compiled.decode + assert dispatches[0][1][20] is device_logits + assert dispatches[0][1][-1] is device_next_hidden + assert dispatches[1][0] is compiled.topk_select + assert dispatches[1][1][0] is device_logits + + def test_decode_kernel_inputs_reject_multi_token_rows(): model = _model(max_batch_size=2) runner = ModelRunner(compiled=_compiled_kernels(model)) @@ -1021,6 +1110,224 @@ def test_engine_ignores_device_sampled_tokens_for_non_greedy_config(): assert sampler.sample_calls == 1 +def test_engine_uses_device_topk_candidates_for_topk_config(): + model = _model(max_batch_size=1) + manager = KvCacheManager() + executor = _DeviceTopkExecutor(manager, token_id=7) + sampler = _CandidateSampler(token_id=7) + engine = LLMEngine(kv_cache_manager=manager, executor=executor, sampler=sampler) + 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=_Tokenizer(), + layer_specs=[], + runtime_model=model, + ) + + result = engine.generate_batch( + model.config.model_id, + ["abc"], + GenerateConfig(max_new_tokens=2, temperature=0.8, top_k=4, top_p=1.0), + )[0] + + assert result.token_ids == [7, 7] + assert executor.prefill_allow_topk is True + assert executor.decode_allow_topk is True + assert sampler.sample_calls == 0 + assert sampler.candidate_calls == 2 + + +def test_engine_chunked_prefill_preserves_device_topk_candidates(): + model = _model( + max_batch_size=1, + max_num_batched_tokens=2, + ) + manager = KvCacheManager() + executor = _DeviceTopkExecutor(manager, token_id=7) + sampler = _CandidateSampler(token_id=7) + engine = LLMEngine(kv_cache_manager=manager, executor=executor, sampler=sampler) + 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=_VariableLengthTokenizer(), + layer_specs=[], + runtime_model=model, + ) + + result = engine.generate_batch( + model.config.model_id, + ["abcd"], + GenerateConfig(max_new_tokens=1, temperature=0.8, top_k=4, top_p=1.0), + )[0] + + assert result.token_ids == [7] + assert executor.prefill_allow_topk is True + assert sampler.sample_calls == 0 + assert sampler.candidate_calls == 1 + + +def test_engine_skips_device_topk_candidates_without_topk_config(): + model = _model(max_batch_size=1) + manager = KvCacheManager() + executor = _DeviceTopkExecutor( + manager, + token_id=7, + always_return_candidates=True, + ) + sampler = _CandidateSampler(token_id=9) + engine = LLMEngine(kv_cache_manager=manager, executor=executor, sampler=sampler) + 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=_Tokenizer(), + layer_specs=[], + runtime_model=model, + ) + + result = engine.generate_batch( + model.config.model_id, + ["abc"], + GenerateConfig(max_new_tokens=1, temperature=0.8, top_k=None), + )[0] + + assert result.token_ids == [9] + assert executor.prefill_allow_topk is False + assert sampler.sample_calls == 1 + assert sampler.candidate_calls == 0 + + +def test_serving_worker_routes_supported_topk_candidates(): + model = _model(max_batch_size=1) + manager = KvCacheManager() + executor = _DeviceTopkExecutor(manager, token_id=7) + sampler = _RoutingSampler(host_token_id=9, candidate_token_id=7) + worker = WorkerProcess.__new__(WorkerProcess) + worker.executor = executor + worker.sampler = sampler + worker.model_record = SimpleNamespace(config=model.config) + worker._req_cache = { + "request": NewRequestData( + request_id="request", + prompt_token_ids=[1], + temperature=0.8, + top_p=1.0, + top_k=4, + ) + } + + prefill_tokens: dict[str, list[int]] = {} + worker._batch_prefill( + [ + PrefillRequest( + request_id="request", + chunk_tokens=[1], + num_computed_tokens=0, + block_ids=[0], + ) + ], + model, + prefill_tokens, + ) + + decode_tokens: dict[str, list[int]] = {} + worker._batch_decode( + [ + DecodeRequest( + request_id="request", + last_token=7, + prev_token=1, + seq_len=2, + block_ids=[0], + ) + ], + model, + decode_tokens, + ) + + assert prefill_tokens == {"request": [7]} + assert decode_tokens == {"request": [7]} + assert executor.prefill_allow_topk is True + assert executor.decode_allow_topk is True + assert sampler.sample_calls == 0 + assert sampler.candidate_calls == 2 + + +def test_serving_worker_mixed_topk_batch_falls_back_from_stale_candidates(): + model = _model(max_batch_size=2) + manager = KvCacheManager() + executor = _DeviceTopkExecutor( + manager, + token_id=7, + always_return_candidates=True, + ) + sampler = _RoutingSampler(host_token_id=9, candidate_token_id=7) + worker = WorkerProcess.__new__(WorkerProcess) + worker.executor = executor + worker.sampler = sampler + worker.model_record = SimpleNamespace(config=model.config) + worker._req_cache = { + "supported": NewRequestData( + request_id="supported", + prompt_token_ids=[1], + temperature=0.8, + top_p=1.0, + top_k=4, + ), + "unsupported": NewRequestData( + request_id="unsupported", + prompt_token_ids=[2], + temperature=0.8, + top_p=1.0, + top_k=8, + ), + } + + prefill_tokens: dict[str, list[int]] = {} + worker._batch_prefill( + [ + PrefillRequest( + request_id=request_id, + chunk_tokens=[token_id], + num_computed_tokens=0, + block_ids=[row], + ) + for row, (request_id, token_id) in enumerate( + (("supported", 1), ("unsupported", 2)) + ) + ], + model, + prefill_tokens, + ) + + decode_tokens: dict[str, list[int]] = {} + worker._batch_decode( + [ + DecodeRequest( + request_id=request_id, + last_token=3, + prev_token=prompt_token, + seq_len=2, + block_ids=[row], + ) + for row, (request_id, prompt_token) in enumerate( + (("supported", 1), ("unsupported", 2)) + ) + ], + model, + decode_tokens, + ) + + assert prefill_tokens == {"supported": [9], "unsupported": [9]} + assert decode_tokens == {"supported": [9], "unsupported": [9]} + assert executor.prefill_allow_topk is False + assert executor.decode_allow_topk is False + assert sampler.sample_calls == 4 + assert sampler.candidate_calls == 0 + + def test_serving_worker_skips_decode_host_embedding_when_executor_embeds_on_device(): model = _model(max_batch_size=1, eos_token_id=0) manager = KvCacheManager() @@ -1632,7 +1939,9 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc monkeypatch.setattr( runner, "_run_distributed_program", - lambda callable_spec, *args: callable_spec.compiled(*args), + lambda callable_spec, *args: callable_spec.compiled( + *(getattr(arg, "tensor", arg) for arg in args) + ), ) executor._runners[model.config.model_id] = runner monkeypatch.setattr( @@ -1774,7 +2083,7 @@ def test_kernel_profile_helpers_emit_kernel_name_and_runtime_timing(): def test_decode_host_inlines_embedding_and_sampling_into_decode_fwd(): module_source = QWEN3_DISPATCH.read_text(encoding="utf-8") start = module_source.index("def qwen3_decode_host") - end = module_source.index("def qwen3_greedy_sample_host") + end = module_source.index("def qwen3_topk_select_host") source = module_source[start:end] assert source.count("decode_fwd(") == 1 @@ -1894,6 +2203,27 @@ def sample(self, logits, params) -> int: return self.token_id +class _CandidateSampler(_FixedSampler): + def __init__(self, token_id: int) -> None: + super().__init__(token_id) + self.candidate_calls = 0 + + def sample_from_candidates(self, candidates, row_idx, params) -> int: + self.candidate_calls += 1 + return self.token_id + + +class _RoutingSampler(_FixedSampler): + def __init__(self, host_token_id: int, candidate_token_id: int) -> None: + super().__init__(host_token_id) + self.candidate_token_id = candidate_token_id + self.candidate_calls = 0 + + def sample_from_candidates(self, candidates, row_idx, params) -> int: + self.candidate_calls += 1 + return self.candidate_token_id + + class _DeviceSamplingExecutor(ModelExecutor): def __init__( self, @@ -1948,6 +2278,77 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: ) +class _DeviceTopkExecutor(ModelExecutor): + def __init__( + self, + kv_cache_manager: KvCacheManager, + token_id: int = 7, + *, + always_return_candidates: bool = False, + ) -> None: + super().__init__(kv_cache_manager) + self.token_id = token_id + self.always_return_candidates = always_return_candidates + self.prefill_allow_topk = False + self.decode_allow_topk = False + + @property + def device_topk_sampling_k(self) -> int: + return 4 + + @property + def supports_device_embedding(self) -> bool: + return True + + def _result_tensors( + self, + batch_size: int, + vocab_size: int, + allow_device_topk_sampling: bool, + ) -> tuple[torch.Tensor, SamplingCandidates | None]: + logits = torch.zeros(batch_size, vocab_size) + candidates = None + if allow_device_topk_sampling or self.always_return_candidates: + values = torch.tensor([[4.0, 3.0, 2.0, 1.0]], dtype=torch.float32) + token_ids = torch.tensor( + [[self.token_id, 6, 5, 4]], + dtype=torch.int32, + ) + candidates = SamplingCandidates( + values=values.expand(batch_size, -1).clone(), + token_ids=token_ids.expand(batch_size, -1).clone(), + ) + if allow_device_topk_sampling: + logits = torch.empty(batch_size, 0) + return logits, candidates + + def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult: + self.prefill_allow_topk = batch.allow_device_topk_sampling + logits, candidates = self._result_tensors( + len(batch.request_ids), + model.config.vocab_size, + batch.allow_device_topk_sampling, + ) + return PrefillResult( + last_hidden=None, + logits=logits, + sampling_candidates=candidates, + ) + + def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: + self.decode_allow_topk = batch.allow_device_topk_sampling + logits, candidates = self._result_tensors( + len(batch.request_ids), + model.config.vocab_size, + batch.allow_device_topk_sampling, + ) + return DecodeResult( + hidden_states=batch.hidden_states, + logits=logits, + sampling_candidates=candidates, + ) + + class _FakeWorker: _DTYPES = { torch.float32: DataType.FLOAT32, diff --git a/tests/test_device_sampling_submission.py b/tests/test_device_sampling_submission.py index 10365456..5c3d655e 100644 --- a/tests/test_device_sampling_submission.py +++ b/tests/test_device_sampling_submission.py @@ -33,11 +33,12 @@ def test_device_sampling_is_limited_by_runtime_vocab_size() -> None: dispatch = _source(QWEN_SERVING / "qwen3_l3_dispatch.py") executor = _source(QWEN_SERVING / "npu_executor.py") runner = _source(QWEN_SERVING / "npu_runner.py") - config = _source(QWEN / "config.py") + constants = _source(QWEN / "constants.py") greedy = _source(QWEN / "greedy_sample.py") + topk = _source(QWEN / "topk_select.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 @@ -45,10 +46,42 @@ def test_device_sampling_is_limited_by_runtime_vocab_size() -> None: assert "REAL_NUM_FULL_VOCAB_CHUNKS" in greedy assert "REAL_VOCAB_TAIL" in greedy assert "token_id >= pl.cast(REAL_VOCAB" in greedy + assert "REAL_NUM_FULL_VOCAB_CHUNKS" in topk + assert "REAL_VOCAB_TAIL" in topk + assert "REAL_VOCAB = M.real_vocab" in topk + assert ":REAL_VOCAB" in topk + + +def test_device_topk_uses_exact_grouped_selection() -> None: + topk = _source(QWEN / "topk_select.py") + dispatch = _source(QWEN_SERVING / "qwen3_l3_dispatch.py") + executor = _source(QWEN_SERVING / "npu_executor.py") + runner = _source(QWEN_SERVING / "npu_runner.py") + + assert "TOPK = 32" in topk + assert "CHUNK_TOPK" not in topk + assert "TOPK_GROUP_WIDTH = 2048" in topk + assert "TOPK_NUM_GROUPS * TOPK <= TOPK_CANDIDATE_PAD" in topk + assert "def _topk_group_pairs(" in topk + assert "for g in pl.range(TOPK_NUM_FULL_GROUPS):" in topk + assert "group_pairs = _topk_group_pairs(logits, b, g)" in topk + assert "pairs = pl.mrgsort(pairs, block_len=1024)" in topk + assert "pl.set_validshape(tail_scores_raw, 1, TOPK_GROUP_TAIL)" in topk + assert "half0_pairs = candidate_sorted[:, 0 : 2 * TOPK]" in topk + assert "half1_pairs = candidate_sorted[" in topk + assert "candidate_pairs = pl.mrgsort(half0_pairs, half1_pairs)" in topk + assert "spread_ids = torch.arange(TOPK" in topk + assert "vals, idx = torch.topk(logits, TOPK" in topk + assert "output_dtype=pl.INT32" in topk + assert "qwen3_topk_select_host" in dispatch + assert "compile_topk_select" in executor + assert "_device_topk_outputs(" in runner + assert "sampling_control" in topk def test_device_greedy_tie_break_matches_host_argmax() -> None: greedy = _source(QWEN / "greedy_sample.py") + topk = _source(QWEN / "topk_select.py") decode_path = QWEN / "decode_layer.py" if not decode_path.is_file(): decode_path = QWEN / "decode_fwd.py" @@ -61,6 +94,11 @@ def test_device_greedy_tie_break_matches_host_argmax() -> None: assert "local_token = pl.cast(scan_t, pl.INT32)" in source assert "if val == best_val:" in source + assert 'name_hint="greedy_select"' in topk + assert "scan_c = (REAL_NUM_VOCAB_CHUNKS - 1) - c" in topk + assert "scan_t = (VOCAB_CHUNK - 1) - t" in topk + assert "local_token = pl.cast(scan_t, pl.INT32)" in topk + def test_prefill_keeps_sampling_in_standalone_device_kernel() -> None: prefill = _source(QWEN / "prefill_fwd.py") @@ -68,8 +106,10 @@ def test_prefill_keeps_sampling_in_standalone_device_kernel() -> None: assert "_greedy_sample_inline" not in prefill assert "_token_embed_inline" not in prefill - assert "compiled.greedy_sample" in runner - assert "_maybe_run_sample_embed(" in runner + assert "compiled.greedy_sample" not in runner + assert "compiled.topk_select" in runner + assert "_maybe_run_greedy_sample(" in runner + assert "selection_k=1" in runner def _device_greedy_argmax_with_clamp(logits):