[Perf] Use global CUDA graph pool for MiMo Audio (#2657)

Signed-off-by: Nick Cao <ncao@redhat.com>

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