model runner support PCP.

Co-authored-by: QiuChunshuo <[email protected]>
Co-authored-by: zhenwenqi2024 <[email protected]>
Signed-off-by: FENP <[email protected]>

Author: 3 people

vllm/v1/worker/gpu_model_runner.py

@@ -42,6 +42,7 @@
 from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
 from vllm.distributed.parallel_state import (
     get_dcp_group,
+    get_pcp_group,
     get_pp_group,
     get_tp_group,
     graph_capture,
@@ -100,7 +101,7 @@
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
     create_fast_prefill_custom_backend,
-    get_dcp_local_seq_lens,
+    get_cp_local_seq_lens,
     reorder_batch_to_split_decodes_and_prefills,
     split_attn_metadata,
 )
@@ -154,7 +155,7 @@
     UBatchSlices,
     check_ubatch_thresholds,
 )
-from vllm.v1.worker.utils import is_residual_scattered_for_sp
+from vllm.v1.worker.utils import PCPManager, is_residual_scattered_for_sp
 
 from .utils import (
     AttentionGroup,
@@ -305,7 +306,11 @@ def __init__(
         # Always set to false after the first forward pass
         self.calculate_kv_scales = self.cache_config.calculate_kv_scales
         self.dcp_world_size = self.parallel_config.decode_context_parallel_size
+        self.pcp_world_size = self.parallel_config.prefill_context_parallel_size
+        self.cp_world_size = self.dcp_world_size * self.pcp_world_size
         self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group
+        self.pcp_rank = 0 if self.pcp_world_size <= 1 else get_pcp_group().rank_in_group
+        self.cp_rank = self.dcp_world_size * self.pcp_rank + self.dcp_rank
         self.max_num_tokens = scheduler_config.max_num_batched_tokens
         self.max_num_reqs = scheduler_config.max_num_seqs
 
@@ -469,25 +474,38 @@ def __init__(
         # Cache the device properties.
         self._init_device_properties()
 
+        if self.pcp_world_size > 1:
+            # NOTE For PCP, we will pad the tokens of each request
+            # to a multiple of 2 * pcp_size that is possible greater
+            # than the max_num_batched_tokens.
+            max_buffer_num_tokens = (
+                self.max_num_tokens + self.max_num_reqs * 2 * self.pcp_world_size
+            )
+        else:
+            max_buffer_num_tokens = self.max_num_tokens
+
         # Persistent buffers for CUDA graphs.
-        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
-        self.positions = self._make_buffer(self.max_num_tokens, dtype=torch.int64)
+        self.input_ids = self._make_buffer(max_buffer_num_tokens, dtype=torch.int32)
+        self.positions = self._make_buffer(max_buffer_num_tokens, dtype=torch.int64)
         self.query_start_loc = self._make_buffer(
             self.max_num_reqs + 1, dtype=torch.int32
         )
         self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
         self.encoder_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
-        if self.dcp_world_size > 1:
-            self.dcp_local_seq_lens = self._make_buffer(
+        if self.cp_world_size > 1:
+            self.cp_local_seq_lens = self._make_buffer(
                 self.max_num_reqs, dtype=torch.int32
             )
         # Because inputs_embeds may be bfloat16 and we don't need a numpy
         # version of this tensor, avoid a RuntimeError by not creating a
         # numpy buffer.
         self.inputs_embeds = self._make_buffer(
-            self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False
+            max_buffer_num_tokens,
+            self.inputs_embeds_size,
+            dtype=self.dtype,
+            numpy=False,
         )
-        self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
+        self.is_token_ids = self._make_buffer(max_buffer_num_tokens, dtype=torch.bool)
         self.discard_request_mask = self._make_buffer(
             self.max_num_reqs, dtype=torch.bool
         )
@@ -500,7 +518,20 @@ def __init__(
 
         # Only relevant for multimodal models
         if self.supports_mm_inputs:
-            self.is_mm_embed = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
+            self.is_mm_embed = self._make_buffer(
+                max_buffer_num_tokens, dtype=torch.bool
+            )
+
+        # Manager for Prefill Context Parallism
+        if self.pcp_world_size > 1:
+            self.pcp_manager = PCPManager(
+                self.pcp_world_size,
+                self.pcp_rank,
+                max_buffer_num_tokens,
+                self.max_num_reqs,
+                self.device,
+                self.pin_memory,
+            )
 
         # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
         if self.uses_mrope:
@@ -515,7 +546,7 @@ def __init__(
             # 1D-RoPE.
             # See page 5 of https://arxiv.org/abs/2409.12191
             self.mrope_positions = self._make_buffer(
-                (3, self.max_num_tokens + 1), dtype=torch.int64
+                (3, max_buffer_num_tokens + 1), dtype=torch.int64
             )
 
         # Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
@@ -531,7 +562,7 @@ def __init__(
         # OPTIMIZATION: Cache the tensors rather than creating them every step.
         # Keep in int64 to avoid overflow with long context
         self.arange_np = np.arange(
-            max(self.max_num_reqs + 1, self.max_model_len, self.max_num_tokens),
+            max(self.max_num_reqs + 1, self.max_model_len, max_buffer_num_tokens),
             dtype=np.int64,
         )
 
@@ -545,7 +576,7 @@ def __init__(
         self.kv_sharing_fast_prefill_logits_indices = None
         if self.cache_config.kv_sharing_fast_prefill:
             self.kv_sharing_fast_prefill_logits_indices = torch.zeros(
-                self.max_num_tokens, dtype=torch.int32, device=self.device
+                max_buffer_num_tokens, dtype=torch.int32, device=self.device
             )
 
         self.uniform_decode_query_len = 1 + self.num_spec_tokens
@@ -1314,6 +1345,32 @@ def _prepare_inputs(
             out=positions_np,
         )
 
+        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
+        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
+
+        if self.pcp_world_size > 1:
+            num_scheduled_tokens[:num_reqs], pcp_positions = (
+                self.pcp_manager.update_tokens_for_pcp(
+                    num_scheduled_tokens[:num_reqs],
+                    self.arange_np,
+                    self.input_batch.num_reqs,
+                    self.reorder_batch_threshold,
+                )
+            )
+
+            # Re-update after PCP split sequences.
+            total_num_scheduled_tokens = sum(num_scheduled_tokens)
+            scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
+
+            req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)
+            cu_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+            positions_np = self.positions.np[:total_num_scheduled_tokens]
+            np.add(
+                self.input_batch.num_computed_tokens_cpu[req_indices],
+                pcp_positions[:total_num_scheduled_tokens],
+                out=positions_np,
+            )
+
         # Calculate M-RoPE positions.
         # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
         if self.uses_mrope:
@@ -1389,9 +1446,6 @@ def _prepare_inputs(
 
                 output_idx += num_sched
 
-        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
-        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
-
         # Prepare the attention metadata.
         self.query_start_loc.np[0] = 0
         self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
@@ -1413,9 +1467,16 @@ def _prepare_inputs(
 
         # Record which requests should not be sampled,
         # so that we could clear the sampled tokens before returning
-        self.discard_request_mask.np[:num_reqs] = (
-            self.seq_lens.np[:num_reqs] < num_tokens_np
-        )
+        if self.pcp_world_size > 1:
+            self.discard_request_mask.np[:num_reqs] = (
+                self.input_batch.num_computed_tokens_cpu[:num_reqs]
+                + num_scheduled_tokens * self.pcp_world_size
+                - self.pcp_manager.num_pcp_pads_cpu[:num_reqs]
+            ) < num_tokens_np
+        else:
+            self.discard_request_mask.np[:num_reqs] = (
+                self.seq_lens.np[:num_reqs] < num_tokens_np
+            )
         self.discard_request_mask.copy_to_gpu(num_reqs)
 
         # Copy the tensors to the GPU.
@@ -1449,10 +1510,19 @@ def _prepare_inputs(
             # We will ignore the sampled tokens from the partial requests.
             # TODO: Support prompt logprobs.
             logits_indices = query_start_loc[1:] - 1
+            if self.pcp_world_size > 1:
+                logits_indices = (
+                    torch.from_numpy(cu_num_tokens) * self.pcp_world_size
+                    - self.pcp_manager.num_pcp_pads_cpu_tensor[:num_reqs]
+                    - 1
+                )
+            else:
+                logits_indices = query_start_loc[1:] - 1
             num_draft_tokens = None
             spec_decode_metadata = None
             num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
         else:
+            assert self.pcp_world_size == 1, "PCP not support spec decode now"
             # Get the number of draft tokens for each request.
             # Iterate over the dictionary rather than all requests since not all
             # requests have draft tokens.
@@ -1516,6 +1586,10 @@ def _build_attention_metadata(
         """
         :return: tuple[attn_metadata, spec_decode_common_attn_metadata]
         """
+        assert num_tokens_padded is None or self.pcp_world_size == 1, (
+            "PCP not support pad attn now"
+        )
+
         num_tokens_padded = num_tokens_padded or num_tokens
         num_reqs_padded = num_reqs_padded or num_reqs
 
@@ -1528,16 +1602,16 @@ def _build_attention_metadata(
                     logits_indices
                 )
 
-        # update seq_lens of decode reqs under DCP.
-        if self.dcp_world_size > 1:
-            self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(
+        # update seq_lens of decode reqs under CP.
+        if self.cp_world_size > 1:
+            self.cp_local_seq_lens.cpu[:num_reqs] = get_cp_local_seq_lens(
                 self.seq_lens.cpu[:num_reqs],
-                self.dcp_world_size,
-                self.dcp_rank,
+                self.cp_world_size,
+                self.cp_rank,
                 self.parallel_config.cp_kv_cache_interleave_size,
             )
-            self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0)
-            self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)
+            self.cp_local_seq_lens.cpu[num_reqs:].fill_(0)
+            self.cp_local_seq_lens.copy_to_gpu(num_reqs_padded)
 
         attn_metadata: PerLayerAttnMetadata = {}
         if ubatch_slices is not None:
@@ -1567,10 +1641,10 @@ def _build_attention_metadata(
             :num_reqs_padded
         ]
 
-        dcp_local_seq_lens, dcp_local_seq_lens_cpu = None, None
-        if self.dcp_world_size > 1:
-            dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded]
-            dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[:num_reqs_padded]
+        cp_local_seq_lens, cp_local_seq_lens_cpu = None, None
+        if self.cp_world_size > 1:
+            cp_local_seq_lens = self.cp_local_seq_lens.gpu[:num_reqs_padded]
+            cp_local_seq_lens_cpu = self.cp_local_seq_lens.cpu[:num_reqs_padded]
 
         spec_decode_common_attn_metadata = None
 
@@ -1585,11 +1659,18 @@ def _build_attention_metadata(
                 num_reqs_padded,
             )
 
+            maybe_pcp_full_tokens = (
+                num_tokens_padded
+                if self.pcp_world_size == 1
+                else num_tokens * self.pcp_world_size
+                - sum(self.pcp_manager.num_pcp_pads_cpu[:num_reqs])
+            )
+
             if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
                 # Encoder-only layers do not have KV cache, so we need to
                 # create a dummy block table and slot mapping for them.
                 blk_table_tensor = torch.zeros(
-                    (num_reqs_padded, 1),
+                    (num_tokens_padded, 1),
                     dtype=torch.int32,
                     device=self.device,
                 )
@@ -1601,12 +1682,26 @@ def _build_attention_metadata(
             else:
                 blk_table = self.input_batch.block_table[kv_cache_gid]
                 blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded)
-                slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded]
+                slot_mapping = blk_table.slot_mapping.gpu[:maybe_pcp_full_tokens]
 
                 # Fill unused with -1. Needed for reshape_and_cache in full cuda
                 # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID
-                slot_mapping[num_tokens:num_tokens_padded].fill_(-1)
-                blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)
+                if self.pcp_world_size == 1:
+                    slot_mapping[num_tokens:num_tokens_padded].fill_(-1)
+                    blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)
+
+            if self.pcp_world_size > 1:
+                # After pcp allgather and restore, there are padded tokens in
+                # kv, so we need pad slotmapping for alignment.
+                pcp_padded_slot_mapping = self.pcp_manager.pcp_padded_slot_mapping[
+                    : num_tokens * self.pcp_world_size
+                ]
+                cp_unpad_mask = self.pcp_manager.pcp_unpad_mask_cpu_tensor[
+                    : num_tokens * self.pcp_world_size
+                ]
+                pcp_padded_slot_mapping.fill_(-1)
+                pcp_padded_slot_mapping[cp_unpad_mask] = slot_mapping
+                slot_mapping = pcp_padded_slot_mapping
 
             common_attn_metadata = CommonAttentionMetadata(
                 query_start_loc=query_start_loc,
@@ -1625,8 +1720,13 @@ def _build_attention_metadata(
                 causal=True,
                 encoder_seq_lens=encoder_seq_lens,
                 encoder_seq_lens_cpu=encoder_seq_lens_cpu,
-                dcp_local_seq_lens=dcp_local_seq_lens,
-                dcp_local_seq_lens_cpu=dcp_local_seq_lens_cpu,
+                cp_local_seq_lens=cp_local_seq_lens,
+                cp_local_seq_lens_cpu=cp_local_seq_lens_cpu,
+                pcp_allgather_restore_idx=self.pcp_manager.pcp_allgather_restore_idx.gpu[
+                    : num_tokens * self.pcp_world_size
+                ]
+                if self.pcp_world_size > 1
+                else None,
             )
 
             if self.speculative_config and spec_decode_common_attn_metadata is None:
@@ -1690,16 +1790,6 @@ def _build_attention_metadata(
                     for layer_name in attn_group.layer_names:
                         attn_metadata[layer_name] = attn_metadata_i
 
-        if spec_decode_common_attn_metadata is not None and (
-            num_reqs != num_reqs_padded or num_tokens != num_tokens_padded
-        ):
-            # Currently the drafter still only uses piecewise cudagraphs (and modifies
-            # the attention metadata in directly), and therefore does not want to use
-            # padded attention metadata.
-            spec_decode_common_attn_metadata = (
-                spec_decode_common_attn_metadata.unpadded(num_tokens, num_reqs)
-            )
-
         return attn_metadata, spec_decode_common_attn_metadata
 
     def _compute_cascade_attn_prefix_lens(
@@ -2904,6 +2994,9 @@ def execute_model(
                     scheduler_output,
                     num_scheduled_tokens_np,
                 )
+                if self.pcp_world_size > 1:
+                    max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
+                    num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
 
                 cascade_attn_prefix_lens = None
                 # Disable cascade attention when using microbatching (DBO)
@@ -3011,6 +3104,23 @@ def execute_model(
                 hidden_states = model_output
                 aux_hidden_states = None
 
+            if self.pcp_world_size > 1:
+                # NOTE we must `slice` hidden_states because pcp_allgather_restore_idx
+                # ignores the padding from CUDA Graph.
+                hidden_states = get_pcp_group().all_gather(
+                    hidden_states[:num_tokens_unpadded],
+                    0,
+                )
+                restore_idx = self.pcp_manager.pcp_allgather_restore_idx.gpu[
+                    : hidden_states.shape[0]
+                ]
+                hidden_states = torch.index_select(
+                    hidden_states,
+                    0,
+                    restore_idx,
+                )
+                # Restore total_num_scheduled_tokens.
+                scheduler_output.total_num_scheduled_tokens = num_scheduled_tokens
             if not self.broadcast_pp_output:
                 # Common case.
                 if not get_pp_group().is_last_rank:
@@ -5274,15 +5384,15 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
                 kv_transfer_group.register_kv_caches(kv_caches)
             kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
 
-        if self.dcp_world_size > 1:
+        if self.cp_world_size > 1:
             layer_type = cast(type[Any], AttentionLayerBase)
             layers = get_layers_from_vllm_config(self.vllm_config, layer_type)
             for layer in layers.values():
                 layer_impl = getattr(layer, "impl", None)
                 if layer_impl is None:
                     continue
                 assert layer_impl.need_to_return_lse_for_decode, (
-                    "DCP requires attention impls to return"
+                    "PCP & DCP require attention impls to return"
                     " the softmax lse for decode, but the impl "
                     f"{layer_impl.__class__.__name__} "
                     "does not return the softmax lse for decode."