dp metadata refactor

Signed-off-by: Sage Moore <[email protected]>

Author: SageMoore

vllm/compilation/ubatch_wrapper.py

@@ -185,7 +185,7 @@ def _ubatch_thread(results, model, ubatch_metadata):
 
     def _make_ubatch_metadata(self, ubatch_slices, attn_metadata, input_ids,
                               positions, inputs_embeds, intermediate_tensors,
-                              compute_stream, num_tokens_across_dp,
+                              compute_stream, dp_metadata,
                               batch_descriptor,
                               cudagraph_runtime_mode) -> list[UbatchMetadata]:
 
@@ -198,8 +198,7 @@ def _make_ubatch_metadata(self, ubatch_slices, attn_metadata, input_ids,
                 create_forward_context(
                     attn_metadata[i] if attn_metadata is not None else None,
                     self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp,
+                    dp_metadata=dp_metadata,
                     batch_descriptor=batch_descriptor,
                     cudagraph_runtime_mode=cudagraph_runtime_mode))
 
@@ -270,8 +269,9 @@ def __call__(self, *args, **kwargs):
         compute_stream = torch.cuda.current_stream()
 
         dp_metadata = forward_context.dp_metadata
+
+        # We shouldn't be here unless we are running with multiple DP ranks
         assert dp_metadata is not None
-        num_tokens_across_dp = dp_metadata._num_tokens_across_dp
 
         if num_tokens not in self.cudagraphs \
             and cudagraph_runtime_mode is CUDAGraphMode.FULL:
@@ -283,7 +283,7 @@ def __call__(self, *args, **kwargs):
                 intermediate_tensors=intermediate_tensors,
                 inputs_embeds=inputs_embeds,
                 compute_stream=compute_stream,
-                num_tokens_across_dp=num_tokens_across_dp,
+                dp_metadata=dp_metadata,
                 batch_descriptor=batch_descriptor,
                 cudagraph_runtime_mode=CUDAGraphMode.NONE)
 
@@ -301,7 +301,7 @@ def __call__(self, *args, **kwargs):
                 intermediate_tensors=intermediate_tensors,
                 inputs_embeds=inputs_embeds,
                 compute_stream=compute_stream,
-                num_tokens_across_dp=num_tokens_across_dp,
+                dp_metadata=dp_metadata,
                 batch_descriptor=batch_descriptor,
                 cudagraph_runtime_mode=CUDAGraphMode.NONE)
             return self._run_ubatches(ubatch_metadata, self.model)

vllm/forward_context.py

@@ -237,18 +237,10 @@ def get_forward_context() -> ForwardContext:
 def create_forward_context(attn_metadata: Any,
                            vllm_config: VllmConfig,
                            virtual_engine: int = 0,
-                           num_tokens: Optional[int] = None,
-                           num_tokens_across_dp: Optional[torch.Tensor] = None,
+                           dp_metadata: Optional[DPMetadata] = None,
                            cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
                            batch_descriptor: Optional[BatchDescriptor] = None,
                            ubatch_slices: Optional[UBatchSlices] = None):
-    dp_metadata: Optional[DPMetadata] = None
-    if vllm_config.parallel_config.data_parallel_size > 1 and (
-            attn_metadata is not None or num_tokens is not None):
-        dp_metadata = DPMetadata.make(vllm_config.parallel_config,
-                                      attn_metadata, num_tokens or 0,
-                                      num_tokens_across_dp)
-
     return ForwardContext(no_compile_layers=vllm_config.compilation_config.
                           static_forward_context,
                           virtual_engine=virtual_engine,
@@ -293,9 +285,15 @@ def set_forward_context(
     if need_to_track_batchsize:
         forward_start_time = time.perf_counter()
 
+    dp_metadata: Optional[DPMetadata] = None
+    if vllm_config.parallel_config.data_parallel_size > 1 and (
+            attn_metadata is not None or num_tokens is not None):
+        dp_metadata = DPMetadata.make(vllm_config.parallel_config,
+                                      attn_metadata, num_tokens or 0,
+                                      num_tokens_across_dp)
+
     forward_context = create_forward_context(attn_metadata, vllm_config,
-                                             virtual_engine, num_tokens,
-                                             num_tokens_across_dp,
+                                             virtual_engine, dp_metadata,
                                              cudagraph_runtime_mode, batch_descriptor,
                                              ubatch_slices)
 

vllm/v1/worker/gpu_model_runner.py

@@ -1686,13 +1686,6 @@ def execute_model(
                 num_input_tokens)
             num_input_tokens += num_pad
 
-        uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
-            num_scheduled_tokens == self.input_batch.num_reqs * max_query_len)
-        batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
-                                           uniform_decode=uniform_decode)
-        cudagraph_runtime_mode, batch_descriptor = \
-            self.cudagraph_dispatcher.dispatch(batch_descriptor)
-
         if self.supports_mm_inputs:
             # Run the multimodal encoder if any.
             self._execute_mm_encoder(scheduler_output)
@@ -1747,6 +1740,8 @@ def execute_model(
                                            uniform_decode=uniform_decode)
         cudagraph_runtime_mode, batch_descriptor = \
             self.cudagraph_dispatcher.dispatch(batch_descriptor)
+        
+        logger.info(f"NUM TOKENS: {num_input_tokens} cudagraph_runtime_mode {cudagraph_runtime_mode} UBATCHING {ubatch_slices is not None}")
 
         # Run the model.
         # Use persistent buffers for CUDA graphs.
@@ -3138,6 +3133,7 @@ def initialize_cudagraph_capture(self) -> None:
 
         # Trigger cudagraph dispatching keys initialization here (after
         # initializing attn backends).
+        logger.info(f"INITIALIZING KEYS FOR MODE: {self.compilation_config.cudagraph_mode}")
         self.cudagraph_dispatcher.initialize_cudagraph_keys(
             self.compilation_config.cudagraph_mode,
             self.uniform_decode_query_len)