Qwen3-14B serving: add device top-k sampling

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
 

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.

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."""

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,14 +394,23 @@ 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,
         ).share_memory_()
         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,