fix splitting ops

Author: Conless

vllm/compilation/nanoflow/manager.py

@@ -38,8 +38,6 @@ def __init__(
         tag_graph(
             self.original_graph_module,
             {
-                "vllm.unified_attention": "attention",
-                "vllm.unified_attention_with_output": "attention",
                 "vllm.all_reduce": "all_reduce",
             },
         )

vllm/config/vllm.py

@@ -433,20 +433,12 @@ def __post_init__(self):
                     "nano batch split. Disabling nano batch split."
                 )
                 self.compilation_config.enable_nano_batch_split = False
-            else:
-                nano_batch_splitting_ops = [
-                    "vllm.all_reduce",
-                ]
-                if self.compilation_config.splitting_ops and set(
-                    self.compilation_config.splitting_ops
-                ) != set(nano_batch_splitting_ops):
-                    logger.info(
-                        "splitting_ops is not supported with "
-                        "nano batch split. Disabling nano batch split."
-                    )
-                    self.compilation_config.enable_nano_batch_split = False
-                else:
-                    self.compilation_config.splitting_ops = nano_batch_splitting_ops
+            elif self.compilation_config.splitting_ops and \
+                "vllm.all_reduce" not in self.compilation_config.splitting_ops:
+                logger.info(
+                    "adding vllm.all_reduce to splitting_ops for nano batch split."
+                )
+                self.compilation_config.splitting_ops.append("vllm.all_reduce")
 
         # If the user does not explicitly set a compilation mode, then
         # we use the default mode. The default mode depends on other

vllm/v1/worker/gpu_model_runner.py

@@ -4195,7 +4195,8 @@ def initialize_metadata_builders(
                     if kv_cache_group_id < len(kernel_block_sizes)
                     else None,
                     num_metadata_builders=1
-                    if not self.parallel_config.enable_dbo
+                    if not self.parallel_config.enable_dbo and
+                    not self.compilation_config.enable_nano_batch_split
                     else 2,
                 )
         # Calculate reorder batch threshold (if needed)