diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index b68a6fc3a..3ed3e8f89 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/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)