diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index a600e350a..eb1adf227 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -1,4 +1,5 @@ from contextlib import contextmanager +from types import SimpleNamespace import torch @@ -86,6 +87,12 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch): runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) + def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn): + batch_desc = SimpleNamespace(num_tokens=int(num_tokens)) + return (False, batch_desc, None, None, None) + + monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner))) + # Initialize per-request embeds (batch-major inside talker_mtp_inputs_embeds) runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0]) runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0]) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index 6c135926d..fc7402890 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -458,10 +458,6 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module): TopPLogitsWarper(top_p=0.8), ] ) - max_batch_size = vllm_config.scheduler_config.max_num_seqs - self.position_ids_buffer = [] - for layer_idx in range(self.num_code_groups - 1): - self.position_ids_buffer.append(torch.arange(layer_idx + 2, dtype=torch.int64).repeat(max_batch_size)) compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: @@ -508,13 +504,11 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module): batch_size = current_input.shape[0] # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively - begin_pos = 0 for layer_idx in range(self.num_code_groups - 1): - # Input for this layer: [last_talker_hidden, prev_layer_embed] - + seq_len = layer_idx + 2 + # Compute position_ids dynamically to avoid torch.compile specializing batch_size + position_ids = torch.arange(seq_len, device=current_input.device, dtype=torch.int64).repeat(batch_size) # Forward through code_predictor model - position_ids = self.position_ids_buffer[layer_idx][: (layer_idx + 2) * batch_size] # [1, seq_len] - begin_pos += layer_idx + 2 outputs = self.model( inputs_embeds=current_input, attention_mask=None, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 65cb43d0b..946694ed2 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -135,8 +135,11 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( self.code_predictor = Qwen3OmniMoeTalkerCodePredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor") ) + max_batch_size = max( + vllm_config.scheduler_config.max_num_seqs, vllm_config.compilation_config.max_cudagraph_capture_size + ) self.layer0_embed_buffer = torch.zeros( - (vllm_config.scheduler_config.max_num_seqs, 1, self.config.text_config.hidden_size), + (max_batch_size, 1, self.config.text_config.hidden_size), dtype=vllm_config.model_config.dtype, ) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 1e2beb046..b3d1833f2 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -52,12 +52,13 @@ class OmniGPUModelRunner(GPUModelRunner): self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL ) hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size - self.talker_mtp_input_ids = self._make_buffer(self.max_num_reqs, dtype=torch.int32) + max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size) + self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32) self.talker_mtp_inputs_embeds = self._make_buffer( - self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False + max_batch_size, hidden_size, dtype=self.dtype, numpy=False ) - self.last_talker_hidden = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False) - self.text_step = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False) + self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) + self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) def _init_mrope_positions(self, req_state: CachedRequestState): """Initialize M-RoPE positions for multimodal inputs. @@ -645,7 +646,7 @@ class OmniGPUModelRunner(GPUModelRunner): if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): num_tokens_padded_talker_mtp = num_tokens_padded if num_tokens_padded_talker_mtp == self.max_num_tokens: - num_tokens_padded_talker_mtp = min(num_tokens_padded, self.max_num_reqs) + num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0] outputs = self.talker_mtp( self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp], self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp], @@ -1075,10 +1076,11 @@ class OmniGPUModelRunner(GPUModelRunner): max_num_scheduled_tokens=1, use_cascade_attn=False, ) - req_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size] - req_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size] - last_talker_hidden = self.last_talker_hidden.gpu[:decode_batch_size] - text_step = self.text_step.gpu[:decode_batch_size] + num_tokens_padded = batch_desc.num_tokens + req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded] + req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded] + last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded] + text_step = self.text_step.gpu[:num_tokens_padded] with set_forward_context( None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ):