Merge remote-tracking branch 'origin/v0.9.2-dev-wm-1218' into v0.9.2-dev

vllm/model_executor/layers/quantization/utils/w8a8_utils.py

@@ -39,6 +39,20 @@ def get_w8a8_int8_marlin_weights(
 
     return weight
 
+
+def w8a8_nt_kpack2_marlin_weight(w8a8_w, # [size_n, size_k// 2 ]
+                                k_tile=16,
+                                n_tile=16, ):
+    assert w8a8_w.dtype == torch.int8, "w8a8_w 必须是 int8 类型"
+    size_n, size_k = w8a8_w.shape
+    assert size_n % k_tile == 0 and size_k % n_tile == 0, "k_tile / n_tile 必须能整除对应维度"
+
+    w8a8_w = w8a8_w.reshape((size_n // n_tile,  n_tile, size_k // k_tile, k_tile))
+    w8a8_w = w8a8_w.permute((0, 2, 1, 3)).contiguous()
+    w8a8_w = w8a8_w.reshape((size_n // k_tile, size_k * k_tile))
+    return w8a8_w
+
+
 def sparse_cutlass_supported() -> bool:
     if not current_platform.is_cuda():
         return False
@@ -455,12 +469,10 @@ def apply_int8_linear(
         elif f"1_{n}_{k}" in  W8A8_TRITONJSON.triton_json_dict:
             if m<=16:
                 m_=m
-                
             elif m<=64:
                 m_= (m //4) * 4 #取值到最近的4的倍数
             elif m<=160:
                 m_=(m // 8) *8
-                
             elif m<200: #256
                 m_=160
             elif m<480: #512

vllm/model_executor/parameter.py

@@ -96,6 +96,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):
     def __init__(self, output_dim: int, **kwargs):
         self._output_dim = output_dim
         super().__init__(**kwargs)
+        self.expect_tp_size = -1
+
 
     @property
     def output_dim(self):
@@ -103,6 +105,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):
 
     def load_column_parallel_weight(self, loaded_weight: torch.Tensor):
         tp_rank = get_tensor_model_parallel_rank()
+        if self.expect_tp_size == 1:
+            tp_rank = 0
         shard_size = self.data.shape[self.output_dim]
         loaded_weight = loaded_weight.narrow(self.output_dim,
                                              tp_rank * shard_size, shard_size)
@@ -123,6 +127,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):
         param_data = self.data
 
         tp_rank = get_tensor_model_parallel_rank()
+        if self.expect_tp_size == 1:
+            tp_rank = 0
         param_data = param_data.narrow(self.output_dim, shard_offset,
                                        shard_size)
         loaded_weight = loaded_weight.narrow(self.output_dim,
@@ -167,6 +173,7 @@ class RowvLLMParameter(BasevLLMParameter):
     def __init__(self, input_dim: int, **kwargs):
         self._input_dim = input_dim
         super().__init__(**kwargs)
+        self.expect_tp_size = -1
 
     @property
     def input_dim(self):
@@ -174,6 +181,8 @@ class RowvLLMParameter(BasevLLMParameter):
 
     def load_row_parallel_weight(self, loaded_weight: torch.Tensor):
         tp_rank = get_tensor_model_parallel_rank()
+        if self.expect_tp_size == 1:
+            tp_rank = 0
         shard_size = self.data.shape[self.input_dim]
         loaded_weight = loaded_weight.narrow(self.input_dim,
                                              tp_rank * shard_size, shard_size)

vllm/v1/spec_decode/eagle.py

@@ -25,6 +25,7 @@ from vllm.v1.attention.backends.mla.common import MLACommonMetadata, MLACommonDe
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.utils import prepare_eagle_input_kernel
+from vllm.utils import round_up
 
 logger = init_logger(__name__)
 
@@ -186,6 +187,7 @@ class EagleProposer:
             num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
         else:
             num_input_tokens = num_tokens
+
         # copy inputs to buffer for cudagraph
         self.positions[:num_tokens] = target_positions
         self.hidden_states[:num_tokens] = target_hidden_states
@@ -224,8 +226,8 @@ class EagleProposer:
 
         with set_forward_context(per_layer_attn_metadata,
                                  self.vllm_config,
-                                 num_tokens=num_input_tokens,
-                                 skip_cuda_graphs=not decoding):
+                                 num_tokens=num_input_tokens):
+                                 #skip_cuda_graphs=not decoding):
             ret_hidden_states = self.model(
                 self.input_ids[:num_input_tokens],
                 self.positions[:num_input_tokens],

vllm/v1/worker/gpu_model_runner.py

@@ -28,7 +28,8 @@ from vllm.distributed.kv_transfer import (get_kv_transfer_group,
 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
 from vllm.distributed.parallel_state import (
     get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
-    prepare_communication_buffer_for_model)
+    prepare_communication_buffer_for_model,
+    get_tensor_model_parallel_world_size)
 from vllm.forward_context import (DPMetadata, get_forward_context,
                                   set_forward_context, set_profilling)
 from vllm.logger import init_logger
@@ -46,7 +47,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
                         GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
                         check_use_alibi, get_dtype_size,
-                        is_pin_memory_available, round_up)
+                        is_pin_memory_available, round_up, round_down)
 from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
                                               CommonAttentionMetadata)
@@ -331,6 +332,13 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
 
         self.draft_probs : Optional[DraftProbs] = None
 
+        self.ep_sp = False
+        self.dp_size = self.parallel_config.data_parallel_size
+        self.tp_size = self.parallel_config.tensor_parallel_size
+        self.enable_expert_parallel = self.parallel_config.enable_expert_parallel
+        if self.enable_expert_parallel and self.dp_size > 1 and self.tp_size > 1:
+            self.ep_sp = True
+
     def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
         """
         Update the order of requests in the batch based on the attention
@@ -1267,7 +1275,7 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
         # TODO(tms) : There are many cases where padding is enabled for
         # prefills, causing unnecessary and excessive padding of activations.
 
-        if dp_size == 1 or self.vllm_config.model_config.enforce_eager:
+        if dp_size == 1 or self.vllm_config.model_config.enforce_eager or envs.VLLM_ALL2ALL_BACKEND != 'naive':
             # Early exit.
             return 0, None
 
@@ -1344,22 +1352,33 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
          spec_decode_metadata,
          num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
-            # Use piecewise CUDA graphs.
-            # Add padding to the batch size.
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
+        
+        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
+        if self.ep_sp:
+            num_input_tokens = round_up(num_scheduled_tokens, self.tp_size)
+            if (self.use_cuda_graph
+                    and num_input_tokens <= self.cudagraph_batch_sizes[-1]):
+                # Use piecewise CUDA graphs.
+                # Add padding to the batch size.
+                num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                    num_input_tokens)
         else:
-            # Eager mode.
-            # Pad tokens to multiple of tensor_parallel_size when
-            # enabled collective fusion for SP
-            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if self.compilation_config.pass_config. \
-                enable_sequence_parallelism and tp_size > 1:
-                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+            if (self.use_cuda_graph
+                    and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
+                # Use piecewise CUDA graphs.
+                # Add padding to the batch size.
+                num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                    num_scheduled_tokens)
             else:
-                num_input_tokens = num_scheduled_tokens
+                # Eager mode.
+                # Pad tokens to multiple of tensor_parallel_size when
+                # enabled collective fusion for SP
+                tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+                if self.compilation_config.pass_config. \
+                        enable_sequence_parallelism and tp_size > 1:
+                    num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+                else:
+                    num_input_tokens = num_scheduled_tokens
 
         # Padding for DP
         num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
@@ -1792,14 +1811,16 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
             if not envs.VLLM_REJECT_SAMPLE_OPT:
                 draft_token_ids = draft_result
             else:
-                draft_req_ids = list(scheduler_output.num_scheduled_tokens.keys())
                 draft_token_ids, draft_probs = draft_result
+            spec_token_ids = draft_token_ids.tolist()
+
+            if envs.VLLM_REJECT_SAMPLE_OPT:
+                draft_req_ids = list(scheduler_output.num_scheduled_tokens.keys())
                 if self.draft_probs is None:
                     self.draft_probs = DraftProbs(
                         draft_probs, draft_req_ids)
                 else:
                     self.draft_probs.update(draft_probs, draft_req_ids)
-            spec_token_ids = draft_token_ids.tolist()
 
         return spec_token_ids
 
@@ -1920,6 +1941,9 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
                     time_after_load - time_before_load)
         prepare_communication_buffer_for_model(self.model)
 
+        if hasattr(self, "drafter"):
+            prepare_communication_buffer_for_model(self.drafter.model)
+
         if is_mixture_of_experts(
                 self.model) and self.parallel_config.enable_eplb:
             logger.info("EPLB is enabled for model %s.",
@@ -2091,6 +2115,11 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
         skip_eplb: bool = False,
         is_profile: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        
+        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
+        if self.ep_sp:
+            if num_tokens < self.tp_size:
+                num_tokens = self.tp_size
 
         # Padding for DP
         num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
@@ -2099,10 +2128,12 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
         # Set num_scheduled_tokens based on num_tokens and max_num_seqs
         # for dummy run with LoRA so that the num_reqs collectively
         # has num_tokens in total.
+        
         assert num_tokens <= self.scheduler_config.max_num_batched_tokens
         max_num_reqs = self.scheduler_config.max_num_seqs
         num_reqs = min(num_tokens, max_num_reqs)
         min_tokens_per_req = num_tokens // num_reqs
+        num_actual_tokens = num_tokens
 
         if not is_profile and self.speculative_config is not None \
             and self.speculative_config.num_lookahead_slots > 0 \
@@ -2110,8 +2141,20 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
 
             min_tokens_per_req = (1 + self.speculative_config.num_lookahead_slots)
             num_reqs = num_tokens // min_tokens_per_req
+
+            if self.ep_sp:
+                num_actual_tokens = round_down(num_tokens, 1 + self.speculative_config.num_lookahead_slots)
+                num_reqs = num_actual_tokens // min_tokens_per_req
+        
         num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+        if not self.ep_sp:
+            num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+        else:
+            if self.speculative_config is not None:
+                num_scheduled_tokens_list[-1] += num_tokens % min_tokens_per_req
+            else:
+                num_scheduled_tokens_list[-1] += num_tokens % num_reqs
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
         num_scheduled_tokens = np.array(num_scheduled_tokens_list,
@@ -2135,7 +2178,7 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
                 seq_lens=seq_lens,
                 # seq_lens_tensor=seq_lens_tensor,
                 num_reqs=num_reqs,
-                num_actual_tokens=num_tokens,
+                num_actual_tokens=num_actual_tokens,
                 max_query_len=num_tokens,
                 num_speculative_tokens=num_speculative_tokens,
             )
@@ -2156,6 +2199,8 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
                 input_ids = None
                 inputs_embeds = self.inputs_embeds[:num_tokens]
             else:
+                self.input_ids[:num_tokens] = torch.randint(0, self.model_config.get_vocab_size(), (num_tokens,),
+                                                            dtype=torch.int32)
                 input_ids = self.input_ids[:num_tokens]
                 inputs_embeds = None
             if self.uses_mrope:
@@ -3166,22 +3211,33 @@ class GPUModelRunnerMTP(GPUModelRunnerBase):
          spec_decode_metadata,
          num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
-            # Use piecewise CUDA graphs.
-            # Add padding to the batch size.
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
+
+        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
+        if self.ep_sp:
+            num_input_tokens = round_up(num_scheduled_tokens, self.tp_size)
+            if (self.use_cuda_graph
+                    and num_input_tokens <= self.cudagraph_batch_sizes[-1]):
+                # Use piecewise CUDA graphs.
+                # Add padding to the batch size.
+                num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                    num_input_tokens)
         else:
-            # Eager mode.
-            # Pad tokens to multiple of tensor_parallel_size when
-            # enabled collective fusion for SP
-            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if self.compilation_config.pass_config. \
-                    enable_sequence_parallelism and tp_size > 1:
-                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+            if (self.use_cuda_graph
+                    and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
+                # Use piecewise CUDA graphs.
+                # Add padding to the batch size.
+                num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                    num_scheduled_tokens)
             else:
-                num_input_tokens = num_scheduled_tokens
+                # Eager mode.
+                # Pad tokens to multiple of tensor_parallel_size when
+                # enabled collective fusion for SP
+                tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+                if self.compilation_config.pass_config. \
+                        enable_sequence_parallelism and tp_size > 1:
+                    num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+                else:
+                    num_input_tokens = num_scheduled_tokens
 
         # Padding for DP
         num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
@@ -3608,16 +3664,17 @@ class GPUModelRunnerMTP(GPUModelRunnerBase):
             if not envs.VLLM_REJECT_SAMPLE_OPT:
                 draft_token_ids = draft_result
             else:
-                draft_req_ids = list(scheduler_output.num_scheduled_tokens.keys())
                 draft_token_ids, draft_probs = draft_result
+            spec_token_ids = draft_token_ids.tolist()
+
+            if envs.VLLM_REJECT_SAMPLE_OPT:
+                draft_req_ids = list(scheduler_output.num_scheduled_tokens.keys())
                 if self.draft_probs is None:
                     self.draft_probs = DraftProbs(
                         draft_probs, draft_req_ids)
                 else:
                     self.draft_probs.update(draft_probs, draft_req_ids)
 
-            spec_token_ids = draft_token_ids.tolist()
-
         return spec_token_ids
 #TODO:稳定后使用GPUModelRunnerMTP替换GPUModelRunner
 if envs.VLLM_USE_ZERO_MTP:

vllm/zero_overhead/v1/eagle.py

@@ -10,6 +10,7 @@ from vllm.v1.attention.backends.mla.common import MLACommonMetadata
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.eagle import PADDING_SLOT_ID, EagleProposer
+from vllm.utils import round_up
 
 
 class V1ZeroEagleProposer(EagleProposer):
@@ -110,6 +111,7 @@ class V1ZeroEagleProposer(EagleProposer):
             num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
         else:
             num_input_tokens = num_tokens
+
         # copy inputs to buffer for cudagraph
         self.positions[:num_tokens] = target_positions
         self.hidden_states[:num_tokens] = target_hidden_states
@@ -148,8 +150,8 @@ class V1ZeroEagleProposer(EagleProposer):
 
         with set_forward_context(per_layer_attn_metadata,
                                  self.vllm_config,
-                                 num_tokens=num_input_tokens,
-                                 skip_cuda_graphs=not decoding):
+                                 num_tokens=num_input_tokens,):
+                                 #skip_cuda_graphs=not decoding):
             ret_hidden_states = self.model(
                 self.input_ids[:num_input_tokens],
                 self.positions[:num_input_tokens],

vllm/zero_overhead/v1/gpu_model_runner.py

@@ -424,22 +424,33 @@ class V1ZeroModelRunner(GPUModelRunner):
          spec_decode_metadata,
          num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
-            # Use piecewise CUDA graphs.
-            # Add padding to the batch size.
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
+        
+        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
+        if self.ep_sp:
+            num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+            if (self.use_cuda_graph
+                    and num_input_tokens <= self.cudagraph_batch_sizes[-1]):
+                # Use piecewise CUDA graphs.
+                # Add padding to the batch size.
+                num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                    num_input_tokens)
         else:
-            # Eager mode.
-            # Pad tokens to multiple of tensor_parallel_size when
-            # enabled collective fusion for SP
-            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if self.compilation_config.pass_config. \
-                enable_sequence_parallelism and tp_size > 1:
-                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+            if (self.use_cuda_graph
+                    and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
+                # Use piecewise CUDA graphs.
+                # Add padding to the batch size.
+                num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                    num_scheduled_tokens)
             else:
-                num_input_tokens = num_scheduled_tokens
+                # Eager mode.
+                # Pad tokens to multiple of tensor_parallel_size when
+                # enabled collective fusion for SP
+                tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+                if self.compilation_config.pass_config. \
+                        enable_sequence_parallelism and tp_size > 1:
+                    num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+                else:
+                    num_input_tokens = num_scheduled_tokens
 
         # Padding for DP
         num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)