[Perf] Optimize sampler D2H sync for HY-Image (#3617)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1590,6 +1590,10 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
         self._sampler: Sampler | None = None
         self._eos_token_id: int = tokenizer.eos_token_id
+        # Lazily built on first sample() call so we can pick up logits.device
+        # without guessing during init. See `sample()` for the comprehension
+        # fast path that uses this.
+        self._blocked_token_ids_tensor: torch.Tensor | None = None
 
         self._replace_rotary_embeddings()
         self._patch_moe_blocks()
@@ -1992,25 +1996,39 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
         min_score = torch.finfo(logits.dtype).min
 
+        if self._is_comprehension:
+            # Comprehension path is stateless: we only need to mask a fixed
+            # set of blocked token ids on every step. Do it in one batched
+            # index_fill_ instead of a per-(req, id) Python loop of scalar
+            # GPU writes, and skip reading `output_token_ids` from CPU
+            # (unused here) so this branch never forces a D2H sync.
+            if self._blocked_token_ids_tensor is None and self._blocked_token_ids:
+                self._blocked_token_ids_tensor = torch.tensor(
+                    sorted(self._blocked_token_ids),
+                    dtype=torch.long,
+                    device=logits.device,
+                )
+            if self._blocked_token_ids_tensor is not None:
+                logits.index_fill_(-1, self._blocked_token_ids_tensor, min_score)
+            return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+        # Generation path retains the per-request stateful logic for forced
+        # stage-transition tokens and ratio-restriction.
         for req_idx in range(logits.shape[0]):
             decoded_tokens: list[int] = (
                 sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
             )
             last_token = decoded_tokens[-1] if decoded_tokens else -1
 
-            if self._is_comprehension:
-                for tid in self._blocked_token_ids:
-                    logits[req_idx, tid] = min_score
-            else:
-                forced = self._get_forced_token(decoded_tokens)
-                if forced is not None:
-                    logits[req_idx].fill_(min_score)
-                    logits[req_idx, forced] = 0
-                elif last_token == self._size_token_id:
-                    self._apply_ratio_restriction(logits, req_idx, min_score)
-                elif last_token in self._all_ratio_ids:
-                    logits[req_idx].fill_(min_score)
-                    logits[req_idx, self._eos_token_id] = 0
+            forced = self._get_forced_token(decoded_tokens)
+            if forced is not None:
+                logits[req_idx].fill_(min_score)
+                logits[req_idx, forced] = 0
+            elif last_token == self._size_token_id:
+                self._apply_ratio_restriction(logits, req_idx, min_score)
+            elif last_token in self._all_ratio_ids:
+                logits[req_idx].fill_(min_score)
+                logits[req_idx, self._eos_token_id] = 0
 
         return self._sampler(logits=logits, sampling_metadata=sampling_metadata)