diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 56cb8788e..85fe4b005 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -50,6 +50,7 @@ from vllm.multimodal.processing import ( PromptUpdate, PromptUpdateDetails, ) +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema @@ -150,7 +151,6 @@ class MiMoLocalDecodeBuffer: dtype = next(model.hidden_states_downcast.parameters()).dtype hidden_size = model.local_config.hidden_size - self.pool = torch.cuda.graph_pool_handle() self.input_tensor = torch.zeros((max_batch_size, 1, hidden_size), dtype=dtype, device=device) self.sampler = MiMoLocalSamplerTensor( temperature=torch.ones(max_batch_size, dtype=torch.float32, device=device), @@ -231,7 +231,7 @@ class MiMoLocalDecodeCudaGraph: cuda_graph = torch.cuda.CUDAGraph() if eager_run_first: model.base_local_forward(input_tensor, local_sampler=sampler) - with torch.cuda.graph(cuda_graph, buffer.pool): + with torch.cuda.graph(cuda_graph, pool=current_platform.get_global_graph_pool()): output_tensor = model.base_local_forward(input_tensor, local_sampler=sampler) return cls( @@ -263,7 +263,6 @@ class MiMoInputLocalTransformerBuffer: hidden_size = model.input_local_config.hidden_size group_size = model.group_size - self.pool = torch.cuda.graph_pool_handle() self.input_tensor = torch.zeros((max_batch_size, group_size, hidden_size), dtype=dtype, device=device) self.lock = threading.Lock() @@ -311,7 +310,7 @@ class MiMoInputLocalTransformerCudaGraph: out = model.input_local_transformer(inputs_embeds=input_tensor, return_dict=True, is_causal=False) _ = out.last_hidden_state - with torch.cuda.graph(cuda_graph, buffer.pool): + with torch.cuda.graph(cuda_graph, pool=current_platform.get_global_graph_pool()): out = model.input_local_transformer(inputs_embeds=input_tensor, return_dict=True, is_causal=False) output_tensor = out.last_hidden_state