refine cudagraph code.

InternLM · Jan 8, 2025 · 8493102 · 8493102
1 parent 3ff2962
commit 8493102
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 303 deletions.
diff --git a/lmdeploy/pytorch/backends/dlinfer/maca/graph_runner.py b/lmdeploy/pytorch/backends/dlinfer/maca/graph_runner.py
diff --git a/lmdeploy/pytorch/backends/dlinfer/maca/op_backend.py b/lmdeploy/pytorch/backends/dlinfer/maca/op_backend.py
@@ -113,6 +113,6 @@ def build_graph_runner(model: torch.nn.Module, model_config: ModelConfig,
                            backend_config: BackendConfig,
                            device: torch.device):
         """build graph runner."""
-        from .graph_runner import MACAGraphRunner
-        return MACAGraphRunner(model, model_config, cache_config,
+        from lmdeploy.pytorch.backends.cuda.graph_runner import CUDAGraphRunner
+        return CUDAGraphRunner(model, model_config, cache_config,
                                backend_config, device)
diff --git a/lmdeploy/utils.py b/lmdeploy/utils.py
@@ -351,7 +351,7 @@ def get_max_batch_size(device_type: str):
     elif device_type == 'ascend':
         return 16
     elif device_type == 'maca':
-        return 128
+        return 256
 
 
 def is_bf16_supported(device_type: str = 'cuda'):