From 6d7840c3957f5364b484e00f5aa8b95eaf4f18bc Mon Sep 17 00:00:00 2001 From: ShaoChunLee Date: Fri, 25 Jul 2025 03:58:03 +0000 Subject: [PATCH 001/233] add fused fp8 bmm Signed-off-by: Divakar Verma --- vllm/v1/attention/backends/mla/common.py | 39 ++++++++++++++++++++---- 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index badff67656c2..a47027f25d4c 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -234,6 +234,26 @@ except ImportError: flashinfer_available = False + +def dynamic_per_batched_tensor_quant( + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn +): + DTYPE_MAX = torch.finfo(dtype).max + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10) + scale = DTYPE_MAX / amax + x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + +from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant +@torch.compiler.disable +def aiter_triton_fp8_bmm_wrapper(x, w, w_s, y = None, transpose_bm = False): + if y is not None: + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, YQ=y, transpose_bm=transpose_bm) + else: + y = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, transpose_bm = transpose_bm) + return y + logger = init_logger(__name__) CUDNN_WORKSPACE_SIZE = 12800 @@ -953,7 +973,8 @@ def _v_up_proj(self, x): # Convert from (B, N, L) to (N, B, L) x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1) # Multiply (N, B, L) x (N, L, V) -> (N, B, V) - x = torch.bmm(x, self.W_UV) + x = aiter_triton_fp8_bmm_wrapper(x, self.W_V, self.W_V_scale, transpose_bm = False) + # x = torch.bmm(x, self.W_UV) # Convert from (N, B, V) to (B, N * V) return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim) @@ -986,6 +1007,7 @@ def get_and_maybe_dequant_weights(layer: LinearBase): # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform # the bmm's in 16-bit, the extra memory overhead of this is fairly low kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T + assert kv_b_proj_weight.shape == ( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), ( @@ -1002,11 +1024,16 @@ def get_and_maybe_dequant_weights(layer: LinearBase): W_UK, W_UV = kv_b_proj_weight.split( [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - - # Convert from (L, N, V) to (N, L, V) - self.W_UV = W_UV.transpose(0, 1) - # Convert from (L, N, P) to (N, P, L) - self.W_UK_T = W_UK.permute(1, 2, 0) + + W_K = W_UK.transpose(0, 1) # 16 512 128 + W_V = W_UV.permute(1, 2, 0) # 16 128 512 + self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(W_K, dtype=torch.float8_e4m3fnuz) + self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(W_V, dtype=torch.float8_e4m3fnuz) + + # # Convert from (L, N, V) to (N, L, V) + # self.W_UV = W_UV.transpose(0, 1) + # # Convert from (L, N, P) to (N, P, L) + # self.W_UK_T = W_UK.permute(1, 2, 0) def _compute_prefill_context( self, From 92e134a007462b53322f8c5344504ce548ee737c Mon Sep 17 00:00:00 2001 From: ShaoChunLee Date: Sat, 26 Jul 2025 06:27:28 +0000 Subject: [PATCH 002/233] add envs Signed-off-by: Divakar Verma --- vllm/envs.py | 4 ++ vllm/v1/attention/backends/mla/common.py | 74 +++++++++++++----------- 2 files changed, 45 insertions(+), 33 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 931edcfa7f1e..18f93472be84 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -158,6 +158,7 @@ VLLM_USE_TRTLLM_ATTENTION: Optional[str] = None VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False + VLLM_AITER_TRITON_FP8_BMM: bool = False def get_default_cache_root(): @@ -971,6 +972,9 @@ def get_vllm_port() -> Optional[int]: # limit will actually be zero-copy decoded. "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int(os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256")), + + "VLLM_AITER_TRITON_FP8_BMM": + lambda: bool(int(os.getenv("VLLM_AITER_TRITON_FP8_BMM", "0"))), # If set, allow insecure serialization using pickle. # This is useful for environments where it is deemed safe to use the diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index a47027f25d4c..4de29f7c307b 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -234,25 +234,25 @@ except ImportError: flashinfer_available = False - -def dynamic_per_batched_tensor_quant( - x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn -): - DTYPE_MAX = torch.finfo(dtype).max - min_val, max_val = x.aminmax() - amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10) - scale = DTYPE_MAX / amax - x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX) - return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() - -from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant -@torch.compiler.disable -def aiter_triton_fp8_bmm_wrapper(x, w, w_s, y = None, transpose_bm = False): - if y is not None: - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, YQ=y, transpose_bm=transpose_bm) - else: - y = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, transpose_bm = transpose_bm) - return y +if envs.VLLM_AITER_TRITON_FP8_BMM: + def dynamic_per_batched_tensor_quant( + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn + ): + DTYPE_MAX = torch.finfo(dtype).max + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10) + scale = DTYPE_MAX / amax + x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant + @torch.compiler.disable + def aiter_triton_fp8_bmm_wrapper(x, w, w_s, group_size = 128, y = None, transpose_bm = False): + if y is not None: + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, group_size = group_size, YQ=y, transpose_bm=transpose_bm) + else: + y = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, group_size = group_size, transpose_bm = transpose_bm) + return y logger = init_logger(__name__) @@ -972,11 +972,18 @@ def _run_prefill_context_chunk_cudnn(self, def _v_up_proj(self, x): # Convert from (B, N, L) to (N, B, L) x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1) - # Multiply (N, B, L) x (N, L, V) -> (N, B, V) - x = aiter_triton_fp8_bmm_wrapper(x, self.W_V, self.W_V_scale, transpose_bm = False) - # x = torch.bmm(x, self.W_UV) - # Convert from (N, B, V) to (B, N * V) - return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim) + if envs.VLLM_AITER_TRITON_FP8_BMM: + # Multiply + Transpose (N, B, L) x (N, L, V) -> (N, B, V) -> (B, N, V) + x = aiter_triton_fp8_bmm_wrapper(x, self.W_V, self.W_V_scale, group_size = 256, transpose_bm = True) + # Convert from (B, N, V) to (B, N * V) + x = x.reshape(-1, self.num_heads * self.v_head_dim) + else: + # Multiply (N, B, L) x (N, L, V) -> (N, B, V) + x = torch.bmm(x, self.W_UV) + # Convert from (N, B, V) to (B, N * V) + x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim) + return x + def process_weights_after_loading(self, act_dtype: torch.dtype): @@ -1025,15 +1032,16 @@ def get_and_maybe_dequant_weights(layer: LinearBase): W_UK, W_UV = kv_b_proj_weight.split( [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - W_K = W_UK.transpose(0, 1) # 16 512 128 - W_V = W_UV.permute(1, 2, 0) # 16 128 512 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(W_K, dtype=torch.float8_e4m3fnuz) - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(W_V, dtype=torch.float8_e4m3fnuz) - - # # Convert from (L, N, V) to (N, L, V) - # self.W_UV = W_UV.transpose(0, 1) - # # Convert from (L, N, P) to (N, P, L) - # self.W_UK_T = W_UK.permute(1, 2, 0) + if envs.VLLM_AITER_TRITON_FP8_BMM: + W_K = W_UK.transpose(0, 1) # 16 512 128 + W_V = W_UV.permute(1, 2, 0) # 16 128 512 + self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(W_K, dtype=torch.float8_e4m3fnuz) + self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(W_V, dtype=torch.float8_e4m3fnuz) + else: + # Convert from (L, N, V) to (N, L, V) + self.W_UV = W_UV.transpose(0, 1) + # Convert from (L, N, P) to (N, P, L) + self.W_UK_T = W_UK.permute(1, 2, 0) def _compute_prefill_context( self, From 9433b841b46e682318bfbff789fbbfade6edc08d Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Fri, 8 Aug 2025 02:03:20 +0000 Subject: [PATCH 003/233] api fix for upstream compatibility Signed-off-by: Divakar Verma --- vllm/v1/attention/backends/mla/common.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index 4de29f7c307b..5589983f2235 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -1234,10 +1234,15 @@ def forward( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) # Convert from (B, N, P) to (N, B, P) decode_q_nope = decode_q_nope.transpose(0, 1) - # Multiply (N, B, P) x (N, P, L) -> (N, B, L) - decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T) - # Convert from (N, B, L) to (B, N, L) - decode_ql_nope = decode_ql_nope.transpose(0, 1) + + if envs.VLLM_AITER_TRITON_FP8_BMM: + # Multiply + Transpose (N, B, P) x (N, P, L) -> (N, B, L) -> (B, N, L) + decode_ql_nope = aiter_triton_fp8_bmm_wrapper(decode_q_nope, self.W_K, self.W_K_scale, group_size = 128, transpose_bm = True) + else: + # Multiply (N, B, P) x (N, P, L) -> (N, B, L) + decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T) + # Convert from (N, B, L) to (B, N, L) + decode_ql_nope = decode_ql_nope.transpose(0, 1) output[:num_decode_tokens] = self._forward_decode( decode_ql_nope, decode_q_pe, kv_cache, attn_metadata) From 245f2eb15025751cf2ca515b43b4e72d664f99e2 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Tue, 12 Aug 2025 16:32:53 +0000 Subject: [PATCH 004/233] improve env switch. reformat lint Signed-off-by: Divakar Verma --- vllm/envs.py | 11 ++- vllm/v1/attention/backends/mla/common.py | 97 ++++++++++++++++++------ 2 files changed, 79 insertions(+), 29 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 18f93472be84..c06a3c4ee1c6 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -97,6 +97,7 @@ VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_USE_AITER_MHA: bool = True + VLLM_ROCM_USE_AITER_FP8BMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True VLLM_ROCM_MOE_PADDING: bool = True @@ -158,7 +159,6 @@ VLLM_USE_TRTLLM_ATTENTION: Optional[str] = None VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False - VLLM_AITER_TRITON_FP8_BMM: bool = False def get_default_cache_root(): @@ -750,6 +750,12 @@ def get_vllm_port() -> Optional[int]: lambda: (os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1")), + # Whether to use aiter triton fp8 bmm kernel + # By default it enabled. + "VLLM_ROCM_USE_AITER_FP8BMM": + lambda: (os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in + ("true", "1")), + # use rocm skinny gemms "VLLM_ROCM_USE_SKINNY_GEMM": lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in @@ -972,9 +978,6 @@ def get_vllm_port() -> Optional[int]: # limit will actually be zero-copy decoded. "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int(os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256")), - - "VLLM_AITER_TRITON_FP8_BMM": - lambda: bool(int(os.getenv("VLLM_AITER_TRITON_FP8_BMM", "0"))), # If set, allow insecure serialization using pickle. # This is useful for environments where it is deemed safe to use the diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index 5589983f2235..09d4148991c2 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -234,10 +234,32 @@ except ImportError: flashinfer_available = False -if envs.VLLM_AITER_TRITON_FP8_BMM: + +def is_rocm_aiter_fp8bmm_enabled() -> bool: + return current_platform.is_rocm() \ + and envs.VLLM_ROCM_USE_AITER_FP8BMM \ + and envs.VLLM_ROCM_USE_AITER + + +if is_rocm_aiter_fp8bmm_enabled(): + from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( # noqa: E501 + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_fp8_bmm) + + def aiter_triton_fp8_bmm_wrapper(x, + w, + w_s, + group_size=128, + y=None, + transpose_bm=False): + return aiter_fp8_bmm(x, + w, + w_s, + group_size=group_size, + YQ=y, + transpose_bm=transpose_bm) + def dynamic_per_batched_tensor_quant( - x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn - ): + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn): DTYPE_MAX = torch.finfo(dtype).max min_val, max_val = x.aminmax() amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10) @@ -245,15 +267,7 @@ def dynamic_per_batched_tensor_quant( x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX) return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() - from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant - @torch.compiler.disable - def aiter_triton_fp8_bmm_wrapper(x, w, w_s, group_size = 128, y = None, transpose_bm = False): - if y is not None: - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, group_size = group_size, YQ=y, transpose_bm=transpose_bm) - else: - y = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(x, w, w_s, group_size = group_size, transpose_bm = transpose_bm) - return y - + logger = init_logger(__name__) CUDNN_WORKSPACE_SIZE = 12800 @@ -972,9 +986,13 @@ def _run_prefill_context_chunk_cudnn(self, def _v_up_proj(self, x): # Convert from (B, N, L) to (N, B, L) x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1) - if envs.VLLM_AITER_TRITON_FP8_BMM: - # Multiply + Transpose (N, B, L) x (N, L, V) -> (N, B, V) -> (B, N, V) - x = aiter_triton_fp8_bmm_wrapper(x, self.W_V, self.W_V_scale, group_size = 256, transpose_bm = True) + if is_rocm_aiter_fp8bmm_enabled(): + # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) + x = aiter_triton_fp8_bmm_wrapper(x, + self.W_V, + self.W_V_scale, + group_size=128, + transpose_bm=True) # Convert from (B, N, V) to (B, N * V) x = x.reshape(-1, self.num_heads * self.v_head_dim) else: @@ -984,7 +1002,6 @@ def _v_up_proj(self, x): x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim) return x - def process_weights_after_loading(self, act_dtype: torch.dtype): def get_layer_weight(layer): @@ -1031,12 +1048,37 @@ def get_and_maybe_dequant_weights(layer: LinearBase): W_UK, W_UV = kv_b_proj_weight.split( [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - - if envs.VLLM_AITER_TRITON_FP8_BMM: - W_K = W_UK.transpose(0, 1) # 16 512 128 - W_V = W_UV.permute(1, 2, 0) # 16 128 512 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(W_K, dtype=torch.float8_e4m3fnuz) - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(W_V, dtype=torch.float8_e4m3fnuz) + + if is_rocm_aiter_fp8bmm_enabled(): + W_K = W_UK.transpose(0, 1) # 16 512 128 + W_V = W_UV.permute(1, 2, 0) # 16 128 512 + self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant( + W_K, dtype=torch.float8_e4m3fnuz) + self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( + W_V, dtype=torch.float8_e4m3fnuz) + logger.info_once( + "[Aiter Triton] compiling fp8 BMM for batch sizes 1 to 128 " + f"W_K shape = {list(self.W_K.shape)} and " + f"W_V shape = {list(self.W_V.shape)}") + for m in range(1, 129): + x = torch.empty((self.W_K.shape[0], m, self.W_K.shape[2]), + dtype=torch.bfloat16, + device=self.W_K.device) + aiter_triton_fp8_bmm_wrapper(x, + self.W_K, + self.W_K_scale, + group_size=128, + transpose_bm=True) + + x = torch.empty((self.W_V.shape[0], m, self.W_V.shape[2]), + dtype=torch.bfloat16, + device=self.W_V.device) + aiter_triton_fp8_bmm_wrapper(x, + self.W_V, + self.W_V_scale, + group_size=128, + transpose_bm=True) + else: # Convert from (L, N, V) to (N, L, V) self.W_UV = W_UV.transpose(0, 1) @@ -1235,9 +1277,14 @@ def forward( # Convert from (B, N, P) to (N, B, P) decode_q_nope = decode_q_nope.transpose(0, 1) - if envs.VLLM_AITER_TRITON_FP8_BMM: - # Multiply + Transpose (N, B, P) x (N, P, L) -> (N, B, L) -> (B, N, L) - decode_ql_nope = aiter_triton_fp8_bmm_wrapper(decode_q_nope, self.W_K, self.W_K_scale, group_size = 128, transpose_bm = True) + if is_rocm_aiter_fp8bmm_enabled(): + # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L) + decode_ql_nope = aiter_triton_fp8_bmm_wrapper( + decode_q_nope, + self.W_K, + self.W_K_scale, + group_size=128, + transpose_bm=True) else: # Multiply (N, B, P) x (N, P, L) -> (N, B, L) decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T) From 6fd99d2f727ef0012d9efd5f296e2f1351cddf9b Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Tue, 12 Aug 2025 19:07:58 +0000 Subject: [PATCH 005/233] nit: formatting and direct aiter fxn call Signed-off-by: Divakar Verma --- vllm/envs.py | 2 +- vllm/v1/attention/backends/mla/common.py | 57 +++++++++--------------- 2 files changed, 22 insertions(+), 37 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index c06a3c4ee1c6..0b016dbc85d6 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -751,7 +751,7 @@ def get_vllm_port() -> Optional[int]: ("true", "1")), # Whether to use aiter triton fp8 bmm kernel - # By default it enabled. + # By default is enabled. "VLLM_ROCM_USE_AITER_FP8BMM": lambda: (os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in ("true", "1")), diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index 09d4148991c2..f96a977bb315 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -243,20 +243,7 @@ def is_rocm_aiter_fp8bmm_enabled() -> bool: if is_rocm_aiter_fp8bmm_enabled(): from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( # noqa: E501 - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_fp8_bmm) - - def aiter_triton_fp8_bmm_wrapper(x, - w, - w_s, - group_size=128, - y=None, - transpose_bm=False): - return aiter_fp8_bmm(x, - w, - w_s, - group_size=group_size, - YQ=y, - transpose_bm=transpose_bm) + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm) def dynamic_per_batched_tensor_quant( x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn): @@ -988,11 +975,11 @@ def _v_up_proj(self, x): x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1) if is_rocm_aiter_fp8bmm_enabled(): # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) - x = aiter_triton_fp8_bmm_wrapper(x, - self.W_V, - self.W_V_scale, - group_size=128, - transpose_bm=True) + x = aiter_triton_fp8_bmm(x, + self.W_V, + self.W_V_scale, + group_size=128, + transpose_bm=True) # Convert from (B, N, V) to (B, N * V) x = x.reshape(-1, self.num_heads * self.v_head_dim) else: @@ -1031,7 +1018,6 @@ def get_and_maybe_dequant_weights(layer: LinearBase): # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform # the bmm's in 16-bit, the extra memory overhead of this is fairly low kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T - assert kv_b_proj_weight.shape == ( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), ( @@ -1064,20 +1050,20 @@ def get_and_maybe_dequant_weights(layer: LinearBase): x = torch.empty((self.W_K.shape[0], m, self.W_K.shape[2]), dtype=torch.bfloat16, device=self.W_K.device) - aiter_triton_fp8_bmm_wrapper(x, - self.W_K, - self.W_K_scale, - group_size=128, - transpose_bm=True) + aiter_triton_fp8_bmm(x, + self.W_K, + self.W_K_scale, + group_size=128, + transpose_bm=True) x = torch.empty((self.W_V.shape[0], m, self.W_V.shape[2]), dtype=torch.bfloat16, device=self.W_V.device) - aiter_triton_fp8_bmm_wrapper(x, - self.W_V, - self.W_V_scale, - group_size=128, - transpose_bm=True) + aiter_triton_fp8_bmm(x, + self.W_V, + self.W_V_scale, + group_size=128, + transpose_bm=True) else: # Convert from (L, N, V) to (N, L, V) @@ -1279,12 +1265,11 @@ def forward( if is_rocm_aiter_fp8bmm_enabled(): # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L) - decode_ql_nope = aiter_triton_fp8_bmm_wrapper( - decode_q_nope, - self.W_K, - self.W_K_scale, - group_size=128, - transpose_bm=True) + decode_ql_nope = aiter_triton_fp8_bmm(decode_q_nope, + self.W_K, + self.W_K_scale, + group_size=128, + transpose_bm=True) else: # Multiply (N, B, P) x (N, P, L) -> (N, B, L) decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T) From c219220e866cf4856704f388b0fab26b585f8a09 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Tue, 12 Aug 2025 20:14:28 +0000 Subject: [PATCH 006/233] fit fp8 dtype selection Signed-off-by: Divakar Verma --- vllm/v1/attention/backends/mla/common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index f96a977bb315..40d29297ce07 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -1039,9 +1039,9 @@ def get_and_maybe_dequant_weights(layer: LinearBase): W_K = W_UK.transpose(0, 1) # 16 512 128 W_V = W_UV.permute(1, 2, 0) # 16 128 512 self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant( - W_K, dtype=torch.float8_e4m3fnuz) + W_K, dtype=current_platform.fp8_dtype()) self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( - W_V, dtype=torch.float8_e4m3fnuz) + W_V, dtype=current_platform.fp8_dtype()) logger.info_once( "[Aiter Triton] compiling fp8 BMM for batch sizes 1 to 128 " f"W_K shape = {list(self.W_K.shape)} and " From 8017d7d72f5a6221cf7fba9282a72bc32703a15c Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Thu, 14 Aug 2025 11:25:36 -0500 Subject: [PATCH 007/233] rm kernel warmup Signed-off-by: Divakar Verma --- vllm/v1/attention/backends/mla/common.py | 28 +++--------------------- 1 file changed, 3 insertions(+), 25 deletions(-) diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py index 40d29297ce07..85f6b56b5503 100755 --- a/vllm/v1/attention/backends/mla/common.py +++ b/vllm/v1/attention/backends/mla/common.py @@ -242,8 +242,9 @@ def is_rocm_aiter_fp8bmm_enabled() -> bool: if is_rocm_aiter_fp8bmm_enabled(): - from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( # noqa: E501 - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm) + from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( # noqa: E501 # isort: skip + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant + as aiter_triton_fp8_bmm) def dynamic_per_batched_tensor_quant( x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn): @@ -1042,29 +1043,6 @@ def get_and_maybe_dequant_weights(layer: LinearBase): W_K, dtype=current_platform.fp8_dtype()) self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( W_V, dtype=current_platform.fp8_dtype()) - logger.info_once( - "[Aiter Triton] compiling fp8 BMM for batch sizes 1 to 128 " - f"W_K shape = {list(self.W_K.shape)} and " - f"W_V shape = {list(self.W_V.shape)}") - for m in range(1, 129): - x = torch.empty((self.W_K.shape[0], m, self.W_K.shape[2]), - dtype=torch.bfloat16, - device=self.W_K.device) - aiter_triton_fp8_bmm(x, - self.W_K, - self.W_K_scale, - group_size=128, - transpose_bm=True) - - x = torch.empty((self.W_V.shape[0], m, self.W_V.shape[2]), - dtype=torch.bfloat16, - device=self.W_V.device) - aiter_triton_fp8_bmm(x, - self.W_V, - self.W_V_scale, - group_size=128, - transpose_bm=True) - else: # Convert from (L, N, V) to (N, L, V) self.W_UV = W_UV.transpose(0, 1) From 3eed848c93a5d60816738ac51be3ebdc656132a8 Mon Sep 17 00:00:00 2001 From: Xiaozhu Meng Date: Tue, 12 Aug 2025 12:53:36 -0700 Subject: [PATCH 008/233] [Kernel][AMD] Avoid D2H copy and cumsum kernel (#22683) Signed-off-by: Xiaozhu Signed-off-by: Michael Goin Co-authored-by: Michael Goin Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- vllm/v1/attention/backends/rocm_aiter_fa.py | 32 +++++++++++++-------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index abe05174507f..e8bffbef4415 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -214,12 +214,14 @@ class AiterFlashAttentionMetadata: # |-- query_len ---| num_actual_tokens: int # Number of tokens excluding padding. + num_actual_kv_tokens: int max_query_len: int query_start_loc: torch.Tensor max_seq_len: int seq_lens: torch.Tensor slot_mapping: torch.Tensor block_table: torch.Tensor + cu_seq_lens: Optional[torch.Tensor] # For cascade attention. use_cascade: bool @@ -272,6 +274,20 @@ def build(self, seq_lens = common_attn_metadata.seq_lens block_table_tensor = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping + if max_query_len > 1: + # We pre-compute cumulative seq len needed for prefill attention + # here to avoid recomputing it for every layer + cu_seq_lens = torch.zeros(seq_lens.shape[0] + 1, + dtype=torch.int32, + device=seq_lens.device) + torch.cumsum(seq_lens, + dim=0, + dtype=cu_seq_lens.dtype, + out=cu_seq_lens[1:]) + num_actual_kv_tokens = int(cu_seq_lens[-1].item()) + else: + cu_seq_lens = None + num_actual_kv_tokens = 0 def schedule(batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal): @@ -281,12 +297,14 @@ def schedule(batch_size, cu_query_lens, max_query_len, seqlens, attn_metadata = AiterFlashAttentionMetadata( num_actual_tokens=num_actual_tokens, + num_actual_kv_tokens=num_actual_kv_tokens, max_query_len=max_query_len, query_start_loc=query_start_loc, max_seq_len=max_seq_len, seq_lens=seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, + cu_seq_lens=cu_seq_lens, use_cascade=use_cascade, common_prefix_len=common_prefix_len, total_tokens=self.total_tokens, @@ -475,16 +493,6 @@ def forward( block_table = attn_metadata.block_table if max_seqlen_q > 1: - - cu_seq_lens = torch.zeros(seqused_k.shape[0] + 1, - dtype=torch.int32, - device=query.device) - - torch.cumsum(seqused_k, - dim=0, - dtype=cu_seq_lens.dtype, - out=cu_seq_lens[1:]) - torch.ops.vllm.flash_attn_varlen_func( query[:num_actual_tokens], key_cache, @@ -497,10 +505,10 @@ def forward( alibi_slopes=self.alibi_slopes, window_size=self.sliding_window, block_table=block_table, - cu_seqlens_k=cu_seq_lens, + cu_seqlens_k=attn_metadata.cu_seq_lens, k_scale=layer._k_scale, v_scale=layer._v_scale, - total_tokens=attn_metadata.total_tokens, + total_tokens=attn_metadata.num_actual_kv_tokens, ) _, num_heads, head_size = query.shape From 380030fb2eb5c4e9bee18562c35f6ab771380899 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 12 Aug 2025 21:53:52 +0200 Subject: [PATCH 009/233] [CI][Nixl] Check kv cache layout during handshake (#22745) Signed-off-by: NickLucche --- .../kv_connector/unit/test_nixl_connector.py | 46 +++++++++++++++++++ .../kv_connector/v1/nixl_connector.py | 13 ++++-- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index c6739832355f..3860d7c85724 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -419,6 +419,52 @@ def test_concurrent_load_kv( return raise TimeoutError("Took too long to complete async handshake.") + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + FakeNixlWrapper) + def test_handshake_fails_on_kv_cache_layout_mismatch(self, dist_init): + """ + Verify that adding a remote agent fails if kv_cache_layout differs. + This test is only relevant for heterogeneous TP. + """ + vllm_config = create_vllm_config() + + # Mock TP world size to 2 to force heterogeneous TP when + # remote_tp_size=1 + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", # noqa: E501 + return_value=2): + # Initialize connector and worker (with fake NIXL wrapper) + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0) + worker = connector.connector_worker + + # Minimal local registration params used by add_remote_agent + worker.slot_size_bytes = 4096 + worker.block_len = worker.slot_size_bytes * worker.block_size + worker.num_blocks = 1 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks + + # Metadata with different kv_cache_layout than local worker + mismatched_layout = "HND" if worker.kv_cache_layout != "HND" \ + else "NHD" + meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0], + num_blocks=1, + block_len=worker.block_len, + attn_backend_name=worker.backend_name, + kv_cache_layout=mismatched_layout, + ) + + # We don't check layout for homogeneous TP and MLA for now, as the + # whole block is moved. + worker.add_remote_agent(meta, remote_tp_size=2) + with pytest.raises(AssertionError): + worker.add_remote_agent(meta, remote_tp_size=1) + # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which # we put here is important. First run ray, it will clean up the resources, then diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index a6eeb278532e..4f51229ffbd2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -30,6 +30,7 @@ from vllm.logger import init_logger from vllm.platforms import _Backend, current_platform from vllm.utils import make_zmq_path, make_zmq_socket +from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.request import RequestStatus @@ -73,6 +74,7 @@ class NixlAgentMetadata( num_blocks: int block_len: int attn_backend_name: str + kv_cache_layout: str @dataclass @@ -538,7 +540,9 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): attn_backend = backend_name_to_enum(self.backend_name) self._use_flashinfer = attn_backend == _Backend.FLASHINFER_VLLM_V1 self._use_pallas_v1 = attn_backend == _Backend.PALLAS_VLLM_V1 + self.kv_cache_layout = get_kv_cache_layout() logger.debug("Detected attention backend %s", self.backend_name) + logger.debug("Detected kv cache layout %s", self.kv_cache_layout) self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size} # With heterogeneous TP, P must wait for all assigned D TP workers to @@ -839,7 +843,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id], num_blocks=self.num_blocks, block_len=self.block_len, - attn_backend_name=self.backend_name) + attn_backend_name=self.backend_name, + kv_cache_layout=self.kv_cache_layout) ready_event = threading.Event() self._nixl_handshake_listener_t = threading.Thread( target=self._nixl_handshake_listener, @@ -900,8 +905,7 @@ def add_remote_agent(self, self._tp_size[engine_id] = remote_tp_size else: assert self._tp_size[engine_id] == remote_tp_size - # We may eventually enable this after asserting equality in cache - # layout and close outputs. + # TODO We may eventually want to skip enforcing the same attn backend. assert nixl_agent_meta.attn_backend_name == self.backend_name remote_agent_name = self.nixl_wrapper.add_remote_agent( @@ -930,6 +934,9 @@ def add_remote_agent(self, if self._use_flashinfer: # Account for joint KV in FlashInfer. remote_block_size //= 2 + if tp_ratio > 1: + # Heterogeneous TP expects same kv_cache_layout. + assert nixl_agent_meta.kv_cache_layout == self.kv_cache_layout assert nixl_agent_meta.block_len == self.block_len * tp_ratio, ( "Remote P worker KV layer cache must be of shape [2, N, " From d453c1c8b92110c4e80d5b7713b61ed9049a073f Mon Sep 17 00:00:00 2001 From: zifeitong Date: Tue, 12 Aug 2025 12:54:42 -0700 Subject: [PATCH 010/233] Fix torch version check for SM100 mxfp4 (#22535) Signed-off-by: Zifei Tong Signed-off-by: mgoin Co-authored-by: mgoin --- vllm/model_executor/layers/fused_moe/layer.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index d5a89655e36d..fb38fb91ead6 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -741,12 +741,14 @@ def __init__( # we padding globally so EP buffer allocation works if quant_config and quant_config.get_name() == "mxfp4": - if not is_torch_equal_or_newer("2.8.0"): - raise RuntimeError("Mxfp4 on hopper requires torch >= 2.8.0") - if current_platform.is_device_capability( - 90) and not has_triton_kernels(): - raise NotImplementedError( - "Triton kernels must be installed for mxfp4 on hopper") + if not current_platform.is_device_capability(100): + if not is_torch_equal_or_newer("2.8.0"): + raise RuntimeError( + "Mxfp4 on non-blackwell requires torch >= 2.8.0") + if not has_triton_kernels(): + raise NotImplementedError( + "triton_kernels must be installed for " + "mxfp4 on non-blackwell") if (current_platform.is_rocm() or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8 or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16): From e56ec0d762b576be252b8508384926e7f0afcc9a Mon Sep 17 00:00:00 2001 From: RUTHLESS-BOT Date: Wed, 13 Aug 2025 04:31:48 +0800 Subject: [PATCH 011/233] [Misc] parametrize 'dtype' in test_flash_mla (#22641) Signed-off-by: RUTHLESS-BOT Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/kernels/attention/test_flashmla.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/kernels/attention/test_flashmla.py b/tests/kernels/attention/test_flashmla.py index 21b08e45fd6f..81841be58352 100644 --- a/tests/kernels/attention/test_flashmla.py +++ b/tests/kernels/attention/test_flashmla.py @@ -35,11 +35,10 @@ def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str) -> None: @pytest.mark.parametrize("block_size", [64]) @pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("varlen", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @torch.inference_mode() def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, - varlen): - # TODO: parametrize using pytest - dtype = torch.bfloat16 + varlen, dtype): device = torch.device("cuda:0") torch.set_default_dtype(dtype) torch.set_default_device(device) @@ -48,7 +47,7 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, random.seed(0) print(f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, " - f"{d=}, {dv=}, {causal=}, {varlen=}") + f"{d=}, {dv=}, {causal=}, {varlen=}, {dtype=}") cache_seqlens = torch.full((b, ), mean_sk, dtype=torch.int32) if varlen: From 5572b49ca9170bd3eca1297a2aa84d2b1eb163fc Mon Sep 17 00:00:00 2001 From: Frank Wang <41319051+frankwang28@users.noreply.github.com> Date: Tue, 12 Aug 2025 15:43:06 -0700 Subject: [PATCH 012/233] [Bugfix] Bump DeepGEMM Version to Fix SMXX Layout Issues (#22606) Signed-off-by: frankwang28 --- docker/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index b96d50f0a1c6..a20a4bfb2b88 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -432,7 +432,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Install DeepGEMM from source ARG DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" -ARG DEEPGEMM_GIT_REF="187656694f7f69e3e7975617a68bc3387680a7e1" +ARG DEEPGEMM_GIT_REF="7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c" RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH' . /etc/environment CUDA_MAJOR="${CUDA_VERSION%%.*}" From 194a9faad7b640263361cf36ee0c8f8d73e3cfcc Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 13 Aug 2025 01:12:26 +0100 Subject: [PATCH 013/233] [Docs] Hide the navigation and toc sidebars on home page (#22749) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/README.md b/docs/README.md index e8d2fd953a96..683e1d37563f 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,3 +1,9 @@ +--- +hide: + - navigation + - toc +--- + # Welcome to vLLM
From b4bcf2b67cffbd1696bbbbfaa1fb7121f3c0a339 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 13 Aug 2025 01:12:30 +0100 Subject: [PATCH 014/233] Fix Transformers backend tensor parallel for multimodal models (#22673) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/transformers.py | 51 ++++++++++++++-------- 1 file changed, 34 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/models/transformers.py b/vllm/model_executor/models/transformers.py index 25b8b69e081b..4ec2b683fc33 100644 --- a/vllm/model_executor/models/transformers.py +++ b/vllm/model_executor/models/transformers.py @@ -505,30 +505,47 @@ def tensor_parallel(self): Apply the model's tensor parallelization plan. Currently only supports linear layers. """ - tp_plan = getattr(self.model.config, "base_model_tp_plan", None) or {} + # Look for tp plans in all of the PreTrainedModels found in self.model + is_pretrained_model = lambda m: isinstance(m, PreTrainedModel) + supports_tp_plan = lambda m: m.config.base_model_tp_plan is not None + pretrained_models = filter(is_pretrained_model, self.model.modules()) + models_with_tp_plan = filter(supports_tp_plan, pretrained_models) - if not tp_plan and self.tp_size > 1: + if not any(models_with_tp_plan) and self.tp_size > 1: raise ValueError( f"{type(self.model)} does not support tensor parallel yet!") - # Some weight loaders expect linear layers to inherit from vLLM's - # LinearBase class, so we set a default style which causes any - # unspecified linear layers to be replaced with ReplicatedLinear - tp_plan[".*"] = "replicate" - - def _tensor_parallel(module: nn.Module, prefix: str = ""): + def _tensor_parallel(module: nn.Module, + prefix: str = "", + tp_plan=None): + tp_plan = tp_plan or {} + + # If the current module is a PreTrainedModel, set the tp_plan for + # all of its children + if isinstance(module, PreTrainedModel): + tp_plan = module.config.base_model_tp_plan or {} + tp_plan = { + maybe_prefix(prefix, k): v + for k, v in tp_plan.items() + } + + # Some weight loaders expect linear layers to inherit from vLLM's + # LinearBase class, so we set a default style which causes any + # unspecified linear layers to be replaced with ReplicatedLinear for child_name, child_module in module.named_children(): qual_name = maybe_prefix(prefix, child_name) - for pattern, style in tp_plan.items(): - if re.match(pattern, qual_name) and isinstance( - child_module, nn.Linear): - new_module = replace_linear_class( - child_module, style, self.quant_config) - setattr(module, child_name, new_module) - log_replacement(qual_name, child_module, new_module) - break + if isinstance(child_module, nn.Linear): + generator = (p for p in tp_plan if re.match(p, qual_name)) + pattern = next(generator, None) + style = tp_plan.get(pattern, "replicate") + new_module = replace_linear_class(child_module, style, + self.quant_config) + setattr(module, child_name, new_module) + log_replacement(qual_name, child_module, new_module) else: - _tensor_parallel(child_module, prefix=qual_name) + _tensor_parallel(child_module, + prefix=qual_name, + tp_plan=tp_plan) _tensor_parallel(self.model) From 3eca03bf743783dd7159d8aa5c3764234bb0c5b3 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Wed, 13 Aug 2025 08:13:17 +0800 Subject: [PATCH 015/233] [Model] Decouple glm4v (#22751) Signed-off-by: Jee Jee Li --- docs/models/supported_models.md | 2 +- vllm/model_executor/models/glm4_1v.py | 26 +++++++++++++++++++++----- vllm/model_executor/models/registry.py | 2 +- 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a24fa4bcce33..dbbbc5122b80 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -615,7 +615,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | ✅︎ | | `GLM4VForCausalLM`^ | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | ✅︎ | | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | ✅︎ | -| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + IE+ + VE+ | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ | ✅︎ | +| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + IE+ + VE+ | `zai-org/GLM-4.5V`, etc. | | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ | ✅︎ | | `H2OVLChatModel` | H2OVL | T + IE+ | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ | ✅︎ | | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | | ✅︎ | diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 7983895687a3..2a89c03bfe7e 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1227,10 +1227,7 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal, "k_proj", "v_proj", ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "gate_up_proj": ["gate_up_proj"] } # To ensure correct weight loading and mapping. @@ -1567,7 +1564,26 @@ def get_mm_mapping(self) -> MultiModelKeys: Get the module prefix in multimodal models """ return MultiModelKeys.from_string_field( - language_model="language_model", + language_model="language_model.model", connector="visual.merger.", tower_model="visual.", ) + + +@MULTIMODAL_REGISTRY.register_processor( + Glm4vMultiModalProcessor, + info=Glm4vProcessingInfo, + dummy_inputs=Glm4vDummyInputsBuilder, +) +class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 64dbde4916a2..b817615b4356 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -208,7 +208,7 @@ "Gemma3nForConditionalGeneration": ("gemma3n_mm", "Gemma3nForConditionalGeneration"), # noqa: E501 "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"), "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), # noqa: E501 - "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), # noqa: E501 + "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"), # noqa: E501 "GraniteSpeechForConditionalGeneration": ("granite_speech", "GraniteSpeechForConditionalGeneration"), # noqa: E501 "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"), "InternVLChatModel": ("internvl", "InternVLChatModel"), From e8b1986bfef49212a1175c30ecb957ff9af42805 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 12 Aug 2025 20:14:46 -0400 Subject: [PATCH 016/233] Add hardware plugins to installation doc (#22732) Signed-off-by: Michael Goin Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/getting_started/installation/README.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md index a252343dcee8..f6ecceb85d86 100644 --- a/docs/getting_started/installation/README.md +++ b/docs/getting_started/installation/README.md @@ -14,3 +14,16 @@ vLLM supports the following hardware platforms: - [Google TPU](google_tpu.md) - [Intel Gaudi](intel_gaudi.md) - [AWS Neuron](aws_neuron.md) + +## Hardware Plugins + +The backends below live **outside** the main `vllm` repository and follow the +[Hardware-Pluggable RFC](../design/plugin_system.md). + +| Accelerator | PyPI / package | Repository | +|-------------|----------------|------------| +| Ascend NPU | `vllm-ascend` | | +| Intel Gaudi (HPU) | N/A, install from source | | +| MetaX MACA GPU | N/A, install from source | | +| Rebellions ATOM / REBEL NPU | `vllm-rbln` | | +| IBM Spyre AIU | `vllm-spyre` | | From a48314cdb426996592204d4d13987abadae21958 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 12 Aug 2025 20:18:39 -0700 Subject: [PATCH 017/233] [V0 Deprecation] Remove multi-step scheduling (#22138) Signed-off-by: Woosuk Kwon Signed-off-by: Woosuk Kwon --- .../tests/genai-perf-tests.json | 1 - .../tests/nightly-tests.json | 6 - .buildkite/test-pipeline.yaml | 22 - .github/CODEOWNERS | 1 - tests/async_engine/test_async_llm_engine.py | 409 -------- tests/config/test_config.yaml | 1 - tests/config/test_config_with_model.yaml | 1 - tests/core/test_chunked_prefill_scheduler.py | 10 +- tests/core/test_num_computed_tokens_update.py | 24 +- .../test_multi_step_output_processor.py | 274 ------ .../openai/correctness/test_lmeval.py | 3 - tests/metrics/test_metrics.py | 39 - .../models/language/generation/test_hybrid.py | 26 - .../multi_step/test_correctness_async_llm.py | 232 ----- tests/multi_step/test_correctness_llm.py | 383 -------- tests/samplers/test_logits_processor.py | 70 -- tests/tpu/lora/test_lora.py | 1 - tests/utils_/test_utils.py | 2 - tests/v1/test_oracle.py | 6 - tests/worker/test_model_input.py | 79 -- vllm/config/__init__.py | 2 - vllm/core/scheduler.py | 92 +- vllm/engine/arg_utils.py | 43 +- vllm/engine/async_llm_engine.py | 26 +- vllm/engine/llm_engine.py | 178 +--- vllm/engine/output_processor/interfaces.py | 26 +- vllm/engine/output_processor/multi_step.py | 211 ---- vllm/platforms/cuda.py | 14 +- vllm/platforms/rocm.py | 14 +- vllm/platforms/tpu.py | 7 +- vllm/sequence.py | 38 - vllm/worker/model_runner.py | 7 +- vllm/worker/multi_step_model_runner.py | 908 ------------------ vllm/worker/multi_step_neuron_model_runner.py | 84 -- ...i_step_neuronx_distributed_model_runner.py | 63 -- vllm/worker/multi_step_worker.py | 197 ---- vllm/worker/neuron_worker.py | 22 +- 37 files changed, 57 insertions(+), 3465 deletions(-) delete mode 100644 tests/async_engine/test_async_llm_engine.py delete mode 100644 tests/engine/test_multi_step_output_processor.py delete mode 100644 tests/multi_step/test_correctness_async_llm.py delete mode 100644 tests/multi_step/test_correctness_llm.py delete mode 100644 tests/samplers/test_logits_processor.py delete mode 100644 vllm/engine/output_processor/multi_step.py delete mode 100644 vllm/worker/multi_step_model_runner.py delete mode 100644 vllm/worker/multi_step_neuron_model_runner.py delete mode 100644 vllm/worker/multi_step_neuronx_distributed_model_runner.py delete mode 100644 vllm/worker/multi_step_worker.py diff --git a/.buildkite/nightly-benchmarks/tests/genai-perf-tests.json b/.buildkite/nightly-benchmarks/tests/genai-perf-tests.json index f26ae7634f3d..afb844880f9f 100644 --- a/.buildkite/nightly-benchmarks/tests/genai-perf-tests.json +++ b/.buildkite/nightly-benchmarks/tests/genai-perf-tests.json @@ -12,7 +12,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, diff --git a/.buildkite/nightly-benchmarks/tests/nightly-tests.json b/.buildkite/nightly-benchmarks/tests/nightly-tests.json index 41b4a4008801..423a3bfe1267 100644 --- a/.buildkite/nightly-benchmarks/tests/nightly-tests.json +++ b/.buildkite/nightly-benchmarks/tests/nightly-tests.json @@ -36,7 +36,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, @@ -90,7 +89,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, @@ -144,7 +142,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, @@ -195,7 +192,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, @@ -248,7 +244,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, @@ -301,7 +296,6 @@ "vllm_server_parameters": { "disable_log_stats": "", "gpu_memory_utilization": 0.9, - "num_scheduler_steps": 10, "max_num_seqs": 512, "dtype": "bfloat16" }, diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml index ebcf51981ef3..740be2bc8770 100644 --- a/.buildkite/test-pipeline.yaml +++ b/.buildkite/test-pipeline.yaml @@ -67,7 +67,6 @@ steps: - python3 standalone_tests/lazy_imports.py - pytest -v -s mq_llm_engine # MQLLMEngine - pytest -v -s async_engine # AsyncLLMEngine - - NUM_SCHEDULER_STEPS=4 pytest -v -s async_engine/test_async_llm_engine.py - pytest -v -s test_inputs.py - pytest -v -s test_outputs.py - pytest -v -s multimodal @@ -773,27 +772,6 @@ steps: - pytest -v -s models/test_oot_registration.py # it needs a clean process - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins -- label: Multi-step Tests (4 GPUs) # 36min - mirror_hardwares: [amdexperimental] - working_dir: "/vllm-workspace/tests" - num_gpus: 4 - source_file_dependencies: - - vllm/model_executor/layers/sampler.py - - vllm/sequence.py - - vllm/worker/worker_base.py - - vllm/worker/worker.py - - vllm/worker/multi_step_worker.py - - vllm/worker/model_runner_base.py - - vllm/worker/model_runner.py - - vllm/worker/multi_step_model_runner.py - - vllm/engine - - tests/multi_step - commands: - # this test is quite flaky - # TODO: investigate and fix. - # - pytest -v -s multi_step/test_correctness_async_llm.py - - pytest -v -s multi_step/test_correctness_llm.py - - label: Pipeline Parallelism Test # 45min mirror_hardwares: [amdexperimental] working_dir: "/vllm-workspace/tests" diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a0a327319a46..b0dd5e99d4c7 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -36,7 +36,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @simon-mo @aarnphm /tests/kernels @tlrmchlsmth @WoosukKwon @yewentao256 /tests/models @DarkLight1337 @ywang96 -/tests/multi_step @alexm-redhat @comaniac /tests/multimodal @DarkLight1337 @ywang96 /tests/prefix_caching @comaniac @KuntaiDu /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 diff --git a/tests/async_engine/test_async_llm_engine.py b/tests/async_engine/test_async_llm_engine.py deleted file mode 100644 index 0eb7a6eb52aa..000000000000 --- a/tests/async_engine/test_async_llm_engine.py +++ /dev/null @@ -1,409 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import asyncio -import os -import uuid -from asyncio import CancelledError -from copy import copy -from dataclasses import dataclass, field -from typing import Any, Optional - -import pytest -import pytest_asyncio -import torch - -from vllm import SamplingParams -from vllm.config import ParallelConfig -from vllm.distributed import cleanup_dist_env_and_memory -from vllm.engine.async_llm_engine import AsyncEngineArgs, AsyncLLMEngine -from vllm.outputs import RequestOutput as RealRequestOutput -from vllm.sampling_params import RequestOutputKind - -from ..utils import wait_for_gpu_memory_to_clear - - -@dataclass -class RequestOutput: - request_id: int - finished: bool = False - - -@dataclass -class MockModelConfig: - use_async_output_proc = True - media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - - -class MockEngine: - - def __init__(self): - self.step_calls = 0 - self.add_request_calls = 0 - self.abort_request_calls = 0 - self.request_id = None - # Ugly, remove dependency when possible - self.parallel_config = ParallelConfig() - self.model_config = MockModelConfig() - - async def step_async(self, virtual_engine): - # PP size is 1, ignore virtual engine - self.step_calls += 1 - return [RequestOutput( - request_id=self.request_id)] if self.request_id else [] - - async def process_model_inputs_async(self, *args, **kwargs): - pass - - async def stop_remote_worker_execution_loop_async(self): - pass - - def generate(self, request_id): - self.request_id = request_id - - def stop_generating(self): - self.request_id = None - - def add_request(self, **kwargs): - del kwargs # Unused - self.add_request_calls += 1 - print(f'Request calls: {self.add_request_calls}') - - async def add_request_async(self, **kwargs): - self.add_request_calls += 1 - return - - def abort_request(self, request_id): - del request_id # Unused - self.abort_request_calls += 1 - - def has_unfinished_requests(self): - return self.request_id is not None - - def has_unfinished_requests_for_virtual_engine(self, virtual_engine): - return self.request_id is not None - - -class MockAsyncLLMEngine(AsyncLLMEngine): - _engine_class = MockEngine - - -@pytest.mark.asyncio -async def test_new_requests_event(): - params = SamplingParams() - - engine = MockAsyncLLMEngine() - engine.start_background_loop() - await asyncio.sleep(0.01) - assert engine.engine.step_calls == 0 - - await engine.add_request("1", "", params) - await asyncio.sleep(0.01) - assert engine.engine.add_request_calls == 1 - assert engine.engine.step_calls == 1 - - await engine.add_request("2", "", params) - engine.engine.generate("2") - await asyncio.sleep(0) - await asyncio.sleep(0) - await asyncio.sleep(0) - assert engine.engine.add_request_calls == 2 - assert engine.engine.step_calls >= 2 - await asyncio.sleep(0.001) - assert engine.engine.step_calls >= 3 - engine.engine.stop_generating() - await asyncio.sleep(0.001) - old_step_calls = engine.engine.step_calls - await asyncio.sleep(0.001) - assert engine.engine.step_calls == old_step_calls - - await engine.add_request("3", "", params) - await asyncio.sleep(0.01) - assert engine.engine.add_request_calls == 3 - assert engine.engine.step_calls == old_step_calls + 1 - await asyncio.sleep(0.01) - assert engine.engine.add_request_calls == 3 - assert engine.engine.step_calls == old_step_calls + 1 - - engine = MockAsyncLLMEngine() - assert engine.get_model_config() is not None - assert engine.get_tokenizer() is not None - assert engine.get_decoding_config() is not None - - -def start_engine(): - wait_for_gpu_memory_to_clear( - devices=list(range(torch.cuda.device_count())), - threshold_bytes=2 * 2**30, - timeout_s=60, - ) - - num_scheduler_steps = int(os.getenv("NUM_SCHEDULER_STEPS", "1")) - print(f"Starting engine with num_scheduler_steps={num_scheduler_steps}") - - return AsyncLLMEngine.from_engine_args( - AsyncEngineArgs(model="facebook/opt-125m", - enforce_eager=True, - num_scheduler_steps=num_scheduler_steps)) - - -def uid() -> str: - return str(uuid.uuid4()) - - -@pytest_asyncio.fixture(scope="module") -async def async_engine(): - # We cannot use monkeypatch since this is a module - # scoped fixture and monkeypatch is function scoped. - previous_value = os.getenv("VLLM_USE_V1", None) - os.environ["VLLM_USE_V1"] = "0" - engine = await asyncio.get_event_loop().run_in_executor(executor=None, - func=start_engine) - try: - yield engine - finally: - engine.shutdown_background_loop() - del engine - await asyncio.sleep(0.1) - cleanup_dist_env_and_memory() - - if previous_value: - os.environ["VLLM_USE_V1"] = previous_value - else: - del os.environ["VLLM_USE_V1"] - - -@pytest.fixture() -def should_do_global_cleanup_after_test(request) -> bool: - # So we can share the async engine fixture between these tests - return False - - -@pytest.mark.asyncio(scope="module") -@pytest.mark.parametrize("stop", [None, ["a stop string"]]) -async def test_asyncio_run(async_engine, stop): - - scheduler_config = await async_engine.get_scheduler_config() - num_scheduler_steps = scheduler_config.num_scheduler_steps - - async def run(prompt: str): - sampling_params = SamplingParams( - temperature=0, - max_tokens=32, - min_tokens=32, - stop=stop, - ) - - output_count = 0 - final_output = None - async for output in async_engine.generate(prompt, - sampling_params, - request_id=uid()): - output_count += 1 - final_output = output - return final_output, output_count - - results = await asyncio.gather( - run("test0"), - run("test0"), - ) - assert len(results) == 2 - first, second = results - - # remove nondeterministic fields for comparison - first[0].metrics = None - second[0].metrics = None - first[0].request_id = None - second[0].request_id = None - - assert str(first) == str(second) - - output_count = results[0][1] - if num_scheduler_steps == 1: - assert output_count == 32 - else: - assert 1 < output_count < 32 - - -@pytest.mark.asyncio(scope="module") -@pytest.mark.parametrize("stop", [None, ["a stop string"]]) -async def test_output_kinds(async_engine, stop): - """Test that output_kind works as expected and that - results are equivalent across different kinds.""" - - scheduler_config = await async_engine.get_scheduler_config() - num_scheduler_steps = scheduler_config.num_scheduler_steps - - sampling_params = SamplingParams( - temperature=0, - max_tokens=32, - min_tokens=32, - stop=stop, - ) - - async def run(prompt: str, kind: RequestOutputKind): - params = copy(sampling_params) - params.output_kind = kind - - output_count = 0 - final_output = None - async for output in async_engine.generate(prompt, - params, - request_id=uid()): - output_count += 1 - final_output = output - - assert final_output is not None - assert final_output.finished - - return (final_output.prompt_token_ids, - final_output.outputs[0].token_ids, - final_output.outputs[0].text, output_count) - - async def run_deltas(prompt: str): - params = copy(sampling_params) - params.output_kind = RequestOutputKind.DELTA - - prompt_tokens = None - output_tokens: list[int] = [] - output_text = "" - output_count = 0 - final_output = None - async for output in async_engine.generate(prompt, - params, - request_id=uid()): - token_ids = output.outputs[0].token_ids - text = output.outputs[0].text - final_output = output - - # Ensure we get prompt ids iff we haven't yet received output tokens - if output_tokens: - assert 1 <= len(token_ids) <= num_scheduler_steps - assert stop or text - assert not output.prompt_token_ids - else: - assert output.prompt_token_ids - prompt_tokens = output.prompt_token_ids - - output_tokens.extend(token_ids) - output_text += text - - output_count += 1 - - assert final_output is not None - assert final_output.finished - - return prompt_tokens, output_tokens, output_text, output_count - - results = await asyncio.gather( - run("common input prompt", RequestOutputKind.CUMULATIVE), - run("common input prompt", RequestOutputKind.FINAL_ONLY), - run_deltas("common input prompt")) - - # Make sure outputs are the same - prompt_set = set(tuple(prompt_ids) for prompt_ids, _, _, _ in results) - assert len(prompt_set) == 1 - - text_set = set(text for _, _, text, _ in results) - assert len(text_set) == 1 - - tokens_set = set(tuple(ids) for _, ids, _, _ in results) - assert len(tokens_set) == 1 - - cumulative, final, deltas = results - - # output message counts - assert cumulative[3] == deltas[3] - - if num_scheduler_steps == 1: - assert cumulative[3] == 32 - else: - assert 1 < cumulative[3] < 32 - - assert final[3] == 1 - - -@pytest.mark.asyncio(scope="module") -@pytest.mark.parametrize("stop", [None, ["a stop string"]]) -async def test_cancellation(async_engine, stop): - scheduler_config = await async_engine.get_scheduler_config() - num_scheduler_steps = scheduler_config.num_scheduler_steps - - sampling_params = SamplingParams( - temperature=0, - min_tokens=13, - max_tokens=13, - stop=stop, - ) - - stop_at = 5 if num_scheduler_steps == 1 else 1 - - request_id = uid() - - i = 0 - with pytest.raises(CancelledError): - async for output in async_engine.generate("test2", - sampling_params, - request_id=request_id): - assert not output.finished - i += 1 - if i == stop_at: - await async_engine.abort(request_id) - - assert i == stop_at - - -@pytest.mark.asyncio(scope="module") -@pytest.mark.parametrize("stop", [None, ["a stop string"]]) -async def test_delayed_generator(async_engine, stop): - scheduler_config = await async_engine.get_scheduler_config() - - if scheduler_config.num_scheduler_steps != 1: - pytest.skip("no need to test this one with multistep") - - sampling_params = SamplingParams( - temperature=0, - min_tokens=10, - max_tokens=10, - stop=stop, - ) - - stream = async_engine.generate("test3", sampling_params, request_id=uid()) - i = 0 - final_output: Optional[RealRequestOutput] = None - async for output in stream: - final_output = output - if i == 0: - # wait for generation to complete before consuming - # the remaining messages - await asyncio.sleep(1) - if i < 9: - assert not output.finished - i += 1 - - assert i == 10 - assert final_output is not None - assert len(final_output.outputs[0].token_ids) == 10 - assert final_output.finished - - -@pytest.mark.asyncio(scope="module") -async def test_invalid_argument(async_engine): - scheduler_config = await async_engine.get_scheduler_config() - - if scheduler_config.num_scheduler_steps != 1: - pytest.skip("no need to test this one with multistep") - - sampling_params = SamplingParams( - temperature=0, - min_tokens=10, - max_tokens=10, - ) - - # Targeting specific DP rank only supported in v1 multi-instance DP - with pytest.raises(ValueError): - async for _ in async_engine.generate("test", - sampling_params, - request_id=uid(), - data_parallel_rank=0): - pass diff --git a/tests/config/test_config.yaml b/tests/config/test_config.yaml index 5090e8f357bb..a16857b5f2fb 100644 --- a/tests/config/test_config.yaml +++ b/tests/config/test_config.yaml @@ -2,4 +2,3 @@ port: 12312 served_model_name: mymodel tensor_parallel_size: 2 trust_remote_code: true -multi_step_stream_outputs: false diff --git a/tests/config/test_config_with_model.yaml b/tests/config/test_config_with_model.yaml index d8c8c7bc8162..9fbdb77d4ef2 100644 --- a/tests/config/test_config_with_model.yaml +++ b/tests/config/test_config_with_model.yaml @@ -4,4 +4,3 @@ port: 12312 served_model_name: mymodel tensor_parallel_size: 2 trust_remote_code: true -multi_step_stream_outputs: false diff --git a/tests/core/test_chunked_prefill_scheduler.py b/tests/core/test_chunked_prefill_scheduler.py index d4dacc4f1296..ce1fe189b3ca 100644 --- a/tests/core/test_chunked_prefill_scheduler.py +++ b/tests/core/test_chunked_prefill_scheduler.py @@ -644,11 +644,9 @@ def cannot_append_second_group2(seq_group, num_lookahead_slots): assert out.num_batched_tokens == max_num_batched_tokens -@pytest.mark.parametrize("num_scheduler_steps", [1, 5]) -def test_chunked_prefill_spec_prefill(num_scheduler_steps): +def test_chunked_prefill_spec_prefill(): """Verify that the num_lookahead_slots is set appropriately for an all""" - """prefill batch depending on whether multi-step scheduling is enabled""" - """or not""" + """prefill batch.""" block_size = 4 max_seqs = 30 max_model_len = 200 @@ -661,7 +659,6 @@ def test_chunked_prefill_spec_prefill(num_scheduler_steps): max_model_len, enable_chunked_prefill=True, num_lookahead_slots=num_lookahead_slots, - num_scheduler_steps=num_scheduler_steps, ) cache_config = CacheConfig(block_size, 1.0, 1, "auto") cache_config.num_cpu_blocks = 16 @@ -679,8 +676,7 @@ def test_chunked_prefill_spec_prefill(num_scheduler_steps): assert out.num_prefill_groups == 1 assert out.num_batched_tokens == max_num_batched_tokens print(out.num_lookahead_slots) - assert out.num_lookahead_slots == (0 if (num_scheduler_steps == 1) else - num_lookahead_slots) + assert out.num_lookahead_slots == 0 def test_chunked_prefill_max_seqs(): diff --git a/tests/core/test_num_computed_tokens_update.py b/tests/core/test_num_computed_tokens_update.py index 9e1b7913dfb9..131a7b3a6299 100644 --- a/tests/core/test_num_computed_tokens_update.py +++ b/tests/core/test_num_computed_tokens_update.py @@ -6,7 +6,6 @@ from tests.conftest import VllmRunner from tests.core.utils import create_dummy_prompt from vllm.engine.llm_engine import LLMEngine -from vllm.platforms import current_platform from vllm.sequence import SequenceGroup MODEL = "JackFram/llama-160m" @@ -17,32 +16,19 @@ def add_seq_group_to_engine(engine: LLMEngine, seq_group: SequenceGroup): scheduler.add_seq_group(seq_group) -@pytest.mark.parametrize("num_scheduler_steps", [1, 8]) @pytest.mark.parametrize("enable_chunked_prefill", [False, True]) @pytest.mark.parametrize("enforce_eager", [False, True]) -def test_num_computed_tokens_update(num_scheduler_steps: int, - enable_chunked_prefill: bool, +def test_num_computed_tokens_update(enable_chunked_prefill: bool, enforce_eager: bool): - is_multi_step = num_scheduler_steps > 1 - is_multi_step_chunked_prefill = is_multi_step and enable_chunked_prefill - - if is_multi_step_chunked_prefill and current_platform.is_rocm(): - pytest.skip("Multi-step with Chunked-Prefill does not support " - "rocm_flash_attn backend") - # Make a vllm engine runner = VllmRunner(model_name=MODEL, gpu_memory_utilization=0.7, - num_scheduler_steps=num_scheduler_steps, enable_chunked_prefill=enable_chunked_prefill, enforce_eager=enforce_eager) engine: LLMEngine = runner.llm.llm_engine - # In multi-step + chunked-prefill there is no separate single prompt step. - # What is scheduled will run for num_scheduler_steps always. - num_prompt_steps = num_scheduler_steps \ - if is_multi_step_chunked_prefill else 1 + num_prompt_steps = 1 num_output_tokens_list = [4, 8, 12, 15, 16, 17] @@ -73,10 +59,8 @@ def test_num_computed_tokens_update(num_scheduler_steps: int, # Test correctness of num_computed_tokens after the decode steps assert seq.data.get_num_computed_tokens( ) == prompt_num_computed_tokens + decode_step_counter - for _ in range(num_scheduler_steps): - # decode step - engine.step() - decode_step_counter += 1 + engine.step() + decode_step_counter += 1 # Test correctness of num_computed_tokens after the sequence finish. assert seq.data.get_num_computed_tokens( diff --git a/tests/engine/test_multi_step_output_processor.py b/tests/engine/test_multi_step_output_processor.py deleted file mode 100644 index 458f4deb743a..000000000000 --- a/tests/engine/test_multi_step_output_processor.py +++ /dev/null @@ -1,274 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import random -from unittest.mock import MagicMock - -import pytest -from transformers import PreTrainedTokenizer - -from vllm.core.scheduler import Scheduler -from vllm.engine.output_processor.multi_step import MultiStepOutputProcessor -from vllm.engine.output_processor.stop_checker import StopChecker -from vllm.sampling_params import SamplingParams -from vllm.sequence import (CompletionSequenceGroupOutput, Logprob, - SequenceOutput, SequenceStatus) -from vllm.transformers_utils.detokenizer import Detokenizer -from vllm.utils import Counter - -from ..core.utils import create_seq_group - - -@pytest.mark.parametrize("seq_output_len", [128]) -@pytest.mark.parametrize("num_new_tokens", [1, 12]) -@pytest.mark.skip_global_cleanup -def test_appends_token_ids(num_new_tokens: int, seq_output_len: int): - """Verify multi-step decoding appends token ids correctly. - - We append token ids and verify all the token ids were appended correctly. - Note that ignore_eos=True. - """ - detokenizer = MagicMock(spec=Detokenizer) - scheduler = MagicMock(spec=Scheduler) - stop_checker = MagicMock(spec=StopChecker) - seq_counter = Counter() - - output_processor = MultiStepOutputProcessor( - detokenizer=detokenizer, - scheduler=[scheduler], - seq_counter=seq_counter, - get_tokenizer_for_seq=lambda _: mock_tokenizer(), - stop_checker=stop_checker, - ) - - seq_group = create_seq_group( - seq_prompt_len=1024, - seq_output_lens=[seq_output_len], - sampling_params=SamplingParams(max_tokens=seq_output_len + - num_new_tokens, - ignore_eos=True), - ) - - seq = seq_group.get_seqs()[0] - seq.status = SequenceStatus.RUNNING - - new_token_ids = list(range(num_new_tokens)) - - outputs = [ - CompletionSequenceGroupOutput( - samples=[ - SequenceOutput( - parent_seq_id=seq.seq_id, - output_token=output_token, - logprobs={output_token: Logprob(0.0)}, - ) - ], - prompt_logprobs=None, - ) for output_token in new_token_ids - ] - - assert seq.get_token_ids()[-len(new_token_ids):] != new_token_ids - output_processor.process_outputs(seq_group, outputs) - assert seq.get_token_ids()[-len(new_token_ids):] == new_token_ids - - -@pytest.mark.parametrize("seq_prompt_len", [1024]) -@pytest.mark.parametrize("seq_output_len", [128]) -@pytest.mark.parametrize("num_new_tokens", [5, 6, 7, 8]) -@pytest.mark.parametrize("max_tokens", [128 + 3]) -@pytest.mark.skip_global_cleanup -def test_respects_max_tokens(num_new_tokens: int, seq_prompt_len: int, - seq_output_len: int, max_tokens: int): - """Verify tokens after max_tokens are dropped and not appended to the - sequence. - """ - detokenizer = MagicMock(spec=Detokenizer) - scheduler = MagicMock(spec=Scheduler) - stop_checker = MagicMock(spec=StopChecker) - seq_counter = Counter() - - output_processor = MultiStepOutputProcessor( - detokenizer=detokenizer, - scheduler=[scheduler], - seq_counter=seq_counter, - get_tokenizer_for_seq=lambda _: mock_tokenizer(), - stop_checker=stop_checker, - ) - - seq_group = create_seq_group( - seq_prompt_len=seq_prompt_len, - seq_output_lens=[seq_output_len], - sampling_params=SamplingParams(max_tokens=max_tokens, ), - ) - - seq = seq_group.get_seqs()[0] - seq.status = SequenceStatus.RUNNING - - new_token_ids = list(range(num_new_tokens)) - - outputs = [ - CompletionSequenceGroupOutput( - samples=[ - SequenceOutput( - parent_seq_id=seq.seq_id, - output_token=output_token, - logprobs={output_token: Logprob(0.0)}, - ) - ], - prompt_logprobs=None, - ) for output_token in new_token_ids - ] - - assert seq.get_len() == seq_prompt_len + seq_output_len - output_processor.process_outputs(seq_group, outputs) - - # Expect the processed sequence to not go over max tokens in len. - assert seq.get_len() == seq_prompt_len + max_tokens - - # Expect the correct tokens were appended. - expected_appended_tokens = new_token_ids[:max_tokens - seq_output_len] - assert seq.get_token_ids( - )[-len(expected_appended_tokens):] == expected_appended_tokens - - -@pytest.mark.parametrize("seq_prompt_len", [1024]) -@pytest.mark.parametrize("seq_output_len", [128]) -@pytest.mark.parametrize("num_new_tokens", [12]) -@pytest.mark.parametrize("seed", list(range(6))) -@pytest.mark.skip_global_cleanup -def test_respects_eos_token_id(num_new_tokens: int, seq_prompt_len: int, - seq_output_len: int, seed: int): - """Verify the eos token id is included in the sequence, but subsequent - tokens are dropped (not appended to sequence). - """ - random.seed(seed) - detokenizer = MagicMock(spec=Detokenizer) - scheduler = MagicMock(spec=Scheduler) - stop_checker = MagicMock(spec=StopChecker) - seq_counter = Counter() - - eos_token_id = 100 - - output_processor = MultiStepOutputProcessor( - detokenizer=detokenizer, - scheduler=[scheduler], - seq_counter=seq_counter, - get_tokenizer_for_seq=lambda _: mock_tokenizer(eos_token_id), - stop_checker=stop_checker, - ) - - seq_group = create_seq_group( - seq_prompt_len=seq_prompt_len, - seq_output_lens=[seq_output_len], - sampling_params=SamplingParams( - # Ensure enough space. - max_tokens=seq_output_len + num_new_tokens, ), - ) - - seq = seq_group.get_seqs()[0] - seq.status = SequenceStatus.RUNNING - - new_token_ids = list(range(num_new_tokens)) - assert eos_token_id not in new_token_ids - eos_index = random.randint(0, len(new_token_ids) - 1) - new_token_ids[eos_index] = eos_token_id - - outputs = [ - CompletionSequenceGroupOutput( - samples=[ - SequenceOutput( - parent_seq_id=seq.seq_id, - output_token=output_token, - logprobs={output_token: Logprob(0.0)}, - ) - ], - prompt_logprobs=None, - ) for output_token in new_token_ids - ] - - assert seq.get_len() == seq_prompt_len + seq_output_len - output_processor.process_outputs(seq_group, outputs) - - # Expect the processed sequence to not go beyond provided eos. - assert seq.get_len() == seq_prompt_len + seq_output_len + (eos_index + 1) - - # Expect the correct tokens were appended. - expected_appended_tokens = new_token_ids[:eos_index + 1] - assert seq.get_token_ids( - )[-len(expected_appended_tokens):] == expected_appended_tokens - - -@pytest.mark.parametrize("seq_prompt_len", [1024]) -@pytest.mark.parametrize("seq_output_len", [128]) -@pytest.mark.parametrize("num_new_tokens", [12]) -@pytest.mark.parametrize("seed", list(range(6))) -@pytest.mark.skip_global_cleanup -def test_ignores_eos_token_id(num_new_tokens: int, seq_prompt_len: int, - seq_output_len: int, seed: int): - """When sampling parameters dictate that we should ignore the eos token id, - ensure all token ids are appended even if the eos token id is emitted. - """ - random.seed(seed) - detokenizer = MagicMock(spec=Detokenizer) - scheduler = MagicMock(spec=Scheduler) - stop_checker = MagicMock(spec=StopChecker) - seq_counter = Counter() - - eos_token_id = 100 - - output_processor = MultiStepOutputProcessor( - detokenizer=detokenizer, - scheduler=[scheduler], - seq_counter=seq_counter, - get_tokenizer_for_seq=lambda _: mock_tokenizer(eos_token_id), - stop_checker=stop_checker, - ) - - seq_group = create_seq_group( - seq_prompt_len=seq_prompt_len, - seq_output_lens=[seq_output_len], - sampling_params=SamplingParams( - # Ensure enough space. - max_tokens=seq_output_len + num_new_tokens, - ignore_eos=True, - ), - ) - - seq = seq_group.get_seqs()[0] - seq.status = SequenceStatus.RUNNING - - new_token_ids = list(range(num_new_tokens)) - assert eos_token_id not in new_token_ids - eos_index = random.randint(0, len(new_token_ids) - 1) - new_token_ids[eos_index] = eos_token_id - - outputs = [ - CompletionSequenceGroupOutput( - samples=[ - SequenceOutput( - parent_seq_id=seq.seq_id, - output_token=output_token, - logprobs={output_token: Logprob(0.0)}, - ) - ], - prompt_logprobs=None, - ) for output_token in new_token_ids - ] - - assert seq.get_len() == seq_prompt_len + seq_output_len - output_processor.process_outputs(seq_group, outputs) - - # Expect the processed sequence to go beyond eos. - assert seq.get_len() == seq_prompt_len + seq_output_len + num_new_tokens - - # Expect the correct tokens were appended. - expected_appended_tokens = new_token_ids[:seq_output_len + num_new_tokens - - seq_output_len] - assert seq.get_token_ids( - )[-len(expected_appended_tokens):] == expected_appended_tokens - - -def mock_tokenizer(eos_token_id=1000): - tokenizer = MagicMock(spec=PreTrainedTokenizer) - tokenizer.eos_token_id = eos_token_id - return tokenizer diff --git a/tests/entrypoints/openai/correctness/test_lmeval.py b/tests/entrypoints/openai/correctness/test_lmeval.py index d75731637d28..684407cd6ee9 100644 --- a/tests/entrypoints/openai/correctness/test_lmeval.py +++ b/tests/entrypoints/openai/correctness/test_lmeval.py @@ -26,15 +26,12 @@ MORE_ARGS_LIST = [ [], # Default ["--enable-chunked-prefill"], # Chunked - ["--num-scheduler-steps", "8"], # MS - ["--num-scheduler-steps", "8", "--multi-step-stream-outputs"] # MS+Stream ] MAX_WAIT_SECONDS = None if current_platform.is_tpu(): MORE_ARGS_LIST = [ [], # Default - # ["--num-scheduler-steps", "8"], # Multi-step << currently fails ] MAX_WAIT_SECONDS = 600 diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_metrics.py index 8cae8a80d38e..dbd9c518e020 100644 --- a/tests/metrics/test_metrics.py +++ b/tests/metrics/test_metrics.py @@ -94,45 +94,6 @@ def test_metric_counter_generation_tokens( f"metric: {metric_count!r}") -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("max_tokens", [128, 129]) -@pytest.mark.parametrize("disable_async_output_proc", [True, False]) -def test_metric_counter_generation_tokens_multi_step( - vllm_runner, - example_prompts, - model: str, - max_tokens: int, - disable_async_output_proc: bool, -) -> None: - num_scheduler_steps = 8 - with vllm_runner( - model, - disable_log_stats=False, - gpu_memory_utilization=0.4, - num_scheduler_steps=num_scheduler_steps, - disable_async_output_proc=disable_async_output_proc, - ) as vllm_model: - vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens) - tokenizer = vllm_model.llm.get_tokenizer() - stat_logger = vllm_model.llm.llm_engine.stat_loggers['prometheus'] - metric_count = stat_logger.metrics.counter_generation_tokens.labels( - **stat_logger.labels)._value.get() - vllm_generation_count = 0 - for i in range(len(example_prompts)): - vllm_output_ids, vllm_output_str = vllm_outputs[i] - prompt_ids = tokenizer.encode(example_prompts[i]) - # vllm_output_ids contains both prompt tokens and generation tokens. - # We're interested only in the count of the generation tokens. - vllm_generation_count += len(vllm_output_ids) - len(prompt_ids) - - # The multi-step scheduling will continue to execute forward even when - # encountering EOS, leading to slightly imprecise metrics. - assert abs(vllm_generation_count - metric_count) <\ - len(example_prompts) * num_scheduler_steps, \ - (f"generation token count: {vllm_generation_count!r}\n" - f"metric: {metric_count!r}") - - @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["float"]) @pytest.mark.parametrize( diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 76f6c226bab7..19fcbf561640 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -331,32 +331,6 @@ def test_state_cleanup( "could be related to finished_requests_ids") -@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]]) -@pytest.mark.parametrize("max_tokens", [64]) -def test_multistep_correctness( - vllm_runner, - example_prompts, - model: str, - max_tokens: int, -) -> None: - with vllm_runner(model, num_scheduler_steps=8, - max_num_seqs=2) as vllm_model: - vllm_outputs_multistep = vllm_model.generate_greedy( - example_prompts, max_tokens) - - with vllm_runner(model, num_scheduler_steps=1, - max_num_seqs=2) as vllm_model: - vllm_outputs_single_step = vllm_model.generate_greedy( - example_prompts, max_tokens) - - check_outputs_equal( - outputs_0_lst=vllm_outputs_multistep, - outputs_1_lst=vllm_outputs_single_step, - name_0="vllm_outputs_multistep", - name_1="vllm_outputs_single_step", - ) - - @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) diff --git a/tests/multi_step/test_correctness_async_llm.py b/tests/multi_step/test_correctness_async_llm.py deleted file mode 100644 index 56e339d485c5..000000000000 --- a/tests/multi_step/test_correctness_async_llm.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Test the AsyncLLMEngine with multi-step-decoding -from typing import Optional - -import pytest - -from vllm.utils import STR_BACKEND_ENV_VAR - -from ..models.utils import check_logprobs_close -from ..utils import (completions_with_server_args, get_client_text_generations, - get_client_text_logprob_generations) - -MODELS = [ - "JackFram/llama-160m", -] -NUM_SCHEDULER_STEPS = [8] # Multi-step decoding steps -NUM_PROMPTS = [10] - -DEFAULT_SERVER_ARGS: list[str] = [ - "--distributed-executor-backend", - "ray", - "--gpu-memory-utilization", - "0.85", - "--swap-space", - "16", -] - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize(("tp_size, pp_size"), [ - (1, 1), - (2, 2), -]) -@pytest.mark.parametrize("eager_mode", [False, True]) -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS) -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS) -@pytest.mark.parametrize("num_logprobs", [5]) -@pytest.mark.parametrize("is_async", [True]) -@pytest.mark.parametrize("attention_backend", ["FLASHINFER", "FLASH_ATTN"]) -@pytest.mark.parametrize("enable_chunked_prefill", [True, False]) -@pytest.mark.asyncio -async def test_multi_step( - example_prompts, - model: str, - tp_size: int, - pp_size: int, - eager_mode: int, - num_scheduler_steps: int, - num_prompts: int, - is_async: bool, - num_logprobs: Optional[int], - attention_backend: str, - enable_chunked_prefill: bool, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test vLLM engine with multi-step scheduling in an OpenAI-protocol - client/server environment. - - Set up an engine with single-step scheduling as a ground-truth reference. - - Send a completions API request to both engines with the same prompts. - - Validate: - * Generated tokens match - * Generated logprobs are all very close - - Args: - example_prompts: test fixture providing example prompts - model: model under test (same for single- and multi-step engines) - tp_size: degree of tensor-parallelism - pp_size: degree of pipeline-parallelism - eager_mode - num_scheduler_steps: for multi-step scheduling, GPU-side steps per - GPU -> CPU output transfer - num_prompts: number of example prompts under test - num_logprobs: corresponds to the `logprobs` argument to the OpenAI - completions endpoint; `None` -> no logprobs - """ - if enable_chunked_prefill and \ - (pp_size > 1 or attention_backend != "FLASH_ATTN"): - pytest.skip("Multi-step with Chunked-Prefill only supports" - "PP=1 and FLASH_ATTN backend") - - with monkeypatch.context() as m: - m.setenv(STR_BACKEND_ENV_VAR, attention_backend) - - prompts = example_prompts - if len(prompts) < num_prompts: - prompts = prompts * ((num_prompts // len(prompts)) + 1) - prompts = prompts[:num_prompts] - assert len(prompts) == num_prompts - - server_args = DEFAULT_SERVER_ARGS + ["--enforce-eager"] - ms_server_args = DEFAULT_SERVER_ARGS + \ - ["--num-scheduler-steps", f"{num_scheduler_steps}"] - - if not is_async: - ms_server_args += ["--disable-async-output-proc"] - - if eager_mode: - ms_server_args.append("--enforce-eager") - - if enable_chunked_prefill: - ms_server_args.append("--enable-chunked-prefill") - - distributed_args = [ - "--tensor-parallel-size", - str(tp_size), - "--pipeline-parallel-size", - str(pp_size), - ] - - # Spin up client/server & issue completion API requests. - # Default `max_wait_seconds` is 240 but was empirically - # was raised 5x to 1200 *just for this test* due to - # observed timeouts in GHA CI - ref_completions = await completions_with_server_args( - prompts, - model, - server_args + distributed_args, - num_logprobs, - max_wait_seconds=5 * 240) - test_completions = await completions_with_server_args( - prompts, - model, - ms_server_args + distributed_args, - num_logprobs, - max_wait_seconds=5 * 240) - - # Assert multi-step scheduling produces identical tokens - # to single-step scheduling. - ref_generations = get_client_text_generations(ref_completions) - test_generations = get_client_text_generations(test_completions) - assert ref_generations == test_generations - - # Assert multi-step scheduling produces nearly-identical logprobs - # to single-step scheduling. - ref_text_logprobs = get_client_text_logprob_generations( - ref_completions) - test_text_logprobs = get_client_text_logprob_generations( - test_completions) - check_logprobs_close( - outputs_0_lst=ref_text_logprobs, - outputs_1_lst=test_text_logprobs, - name_0="hf", - name_1="vllm", - ) - - -@pytest.mark.parametrize(("tp_size, pp_size"), [ - (1, 2), -]) -@pytest.mark.asyncio -async def test_multi_step_pp_smoke( - tp_size: int, - pp_size: int, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """ - Smoke test for the vLLM engine with multi-step scheduling in an - OpenAI-protocol client/server environment. - - This tests compares the outputs between multi-step scheduling and - single-step scheduling. Notably, this test lets the engines generate - more tokens (default is 5) and test for an exact match over all the - tokens. - - Args: - tp_size: degree of tensor-parallelism - pp_size: degree of pipeline-parallelism - eager_mode - """ - - model = "JackFram/llama-160m" - num_scheduler_steps = 8 - attention_backend = "FLASH_ATTN" - max_num_seqs = 3 - - with monkeypatch.context() as m: - m.setenv(STR_BACKEND_ENV_VAR, attention_backend) - - # Prompt from the ShareGPT dataset - prompts = [ - "in the jtbd context whats a push?", # codespell:ignore - "in the jtbd context whats a push?", # codespell:ignore - "in the jtbd context whats a push?", # codespell:ignore - "in the jtbd context whats a push?", # codespell:ignore - ] - # Use varying max_tokens to introduce scheduling randomness. - max_tokens = [10 * i for i in range(1, len(prompts) + 1)] - assert len(prompts) == len(max_tokens) - - test_args = [ - "--tensor-parallel-size", - str(tp_size), "--pipeline-parallel-size", - str(pp_size), "--max-num-seqs", - str(max_num_seqs) - ] - - server_args = DEFAULT_SERVER_ARGS + test_args - ms_server_args = DEFAULT_SERVER_ARGS + \ - ["--num-scheduler-steps", f"{num_scheduler_steps}"] + \ - test_args - - # Spin up client/server & issue completion API requests. - # Default `max_wait_seconds` is 240 but was empirically - # was raised 3x to 720 *just for this test* due to - # observed timeouts in GHA CI - ref_completions = await completions_with_server_args( - prompts=prompts, - model_name=model, - server_cli_args=server_args, - num_logprobs=None, - max_wait_seconds=5 * 240, - max_tokens=max_tokens) - - test_completions = await completions_with_server_args( - prompts=prompts, - model_name=model, - server_cli_args=ms_server_args, - num_logprobs=None, - max_wait_seconds=5 * 240, - max_tokens=max_tokens) - - # Assert multi-step scheduling produces identical tokens - # to single-step scheduling. - ref_generations = get_client_text_generations(ref_completions) - test_generations = get_client_text_generations(test_completions) - - assert ref_generations == test_generations diff --git a/tests/multi_step/test_correctness_llm.py b/tests/multi_step/test_correctness_llm.py deleted file mode 100644 index 0df00c98b72c..000000000000 --- a/tests/multi_step/test_correctness_llm.py +++ /dev/null @@ -1,383 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Test the LLMEngine with multi-step-decoding - -import copy -from typing import Optional - -import pytest - -from vllm.platforms import current_platform -from vllm.utils import STR_BACKEND_ENV_VAR - -from ..models.utils import check_logprobs_close, check_outputs_equal - -MODELS = [ - "JackFram/llama-160m", -] -NUM_SCHEDULER_STEPS = [8] # Multi-step decoding steps -NUM_PROMPTS = [10] - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.parametrize("tp_size", [1]) -@pytest.mark.parametrize("enable_chunked_prefill", [False, True]) -@pytest.mark.parametrize("max_tokens", [5]) -@pytest.mark.parametrize("enforce_eager", [True, False]) -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS) -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS) -@pytest.mark.parametrize("num_logprobs", [None, 5]) -@pytest.mark.parametrize("attention_backend", ["FLASH_ATTN", "FLASHINFER"]) -def test_multi_step_llm( - hf_runner, - vllm_runner, - example_prompts, - model: str, - dtype: str, - tp_size: int, - enable_chunked_prefill: bool, - max_tokens: int, - enforce_eager: int, - num_scheduler_steps: int, - num_prompts: int, - num_logprobs: Optional[int], - attention_backend: str, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test vLLM engine with multi-step scheduling via sync LLM Engine. - - Set up a HuggingFace (HF) transformers model as a ground-truth reference. - - Prompt them with the same example prompts. - - Validate: - * Generated tokens match - * Generated logprobs are all very close - - Args: - hf_runner: HF transformers model runner fixture - vllm_runner: vLLM model runner fixture - example_prompts: test fixture providing example prompts - model: model under test (same for single- and multi-step engines) - dtype: tensor datatype for engine to utilize - tp_size: degree of tensor-parallelism - enable_chunked_prefill: chunked-prefill on/off - max_tokens: the maximum number of tokens to generate - enforce_eager - num_scheduler_steps: for multi-step scheduling, GPU-side steps per - GPU -> CPU output transfer - num_prompts: number of example prompts under test - num_logprobs: corresponds to the `logprobs` argument to the OpenAI - completions endpoint; `None` -> 1 logprob returned. - """ - if current_platform.is_rocm() and \ - (attention_backend == "FLASHINFER" or enable_chunked_prefill): - pytest.skip( - "Multi-Step with FLASHINFER or Chunked-Prefill is not supported" - "on ROCm") - - with monkeypatch.context() as m: - m.setenv(STR_BACKEND_ENV_VAR, attention_backend) - - prompts = example_prompts - if len(prompts) < num_prompts: - prompts = prompts * ((num_prompts // len(prompts)) + 1) - prompts = prompts[:num_prompts] - assert len(prompts) == num_prompts - - with vllm_runner( - model, - dtype=dtype, - enforce_eager=enforce_eager, - gpu_memory_utilization=0.7, - tensor_parallel_size=tp_size, - enable_chunked_prefill=enable_chunked_prefill, - num_scheduler_steps=num_scheduler_steps, - ) as vllm_model: - vllm_outputs = (vllm_model.generate_greedy(prompts, max_tokens) - if num_logprobs is None else - vllm_model.generate_greedy_logprobs( - prompts, max_tokens, num_logprobs)) - - with hf_runner(model, dtype=dtype) as hf_model: - hf_outputs = (hf_model.generate_greedy(prompts, max_tokens) - if num_logprobs is None else - hf_model.generate_greedy_logprobs_limit( - prompts, max_tokens, num_logprobs)) - - if num_logprobs is None: - check_outputs_equal( - outputs_0_lst=hf_outputs, - outputs_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - else: - check_logprobs_close( - outputs_0_lst=hf_outputs, - outputs_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.parametrize("tp_size", [1]) -@pytest.mark.parametrize("max_tokens", [5]) -@pytest.mark.parametrize("enforce_eager", [True]) -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS) -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS) -@pytest.mark.parametrize("num_logprobs,num_prompt_logprobs", [(5, 5)]) -@pytest.mark.parametrize("attention_backend", ["FLASH_ATTN"]) -def test_multi_step_llm_w_prompt_logprobs( - vllm_runner, - example_prompts, - model: str, - dtype: str, - tp_size: int, - max_tokens: int, - enforce_eager: int, - num_scheduler_steps: int, - num_prompts: int, - num_logprobs: Optional[int], - num_prompt_logprobs: Optional[int], - attention_backend: str, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test prompt logprobs with multi-step scheduling via sync LLM Engine. - - Set up a vLLM engine instance w/ single-step scheduling as a ground-truth - reference. - - Prompt them with the same example prompts. - - Validate: - * All generated logprobs are all very close - - Args: - hf_runner: HF transformers model runner fixture - vllm_runner: vLLM model runner fixture - example_prompts: test fixture providing example prompts - model: model under test (same for single- and multi-step engines) - dtype: tensor datatype for engine to utilize - tp_size: degree of tensor-parallelism - max_tokens: the maximum number of tokens to generate - enforce_eager - num_scheduler_steps: for multi-step scheduling, GPU-side steps per - GPU -> CPU output transfer - num_prompts: number of example prompts under test - num_logprobs: corresponds to the `logprobs` argument to the OpenAI - completions endpoint; `None` -> no logprobs - num_prompt_logprobs: number of logprobs to return for each prompt token; - note that this argument is not supported by the - OpenAI completions endpoint. - """ - with monkeypatch.context() as m: - m.setenv(STR_BACKEND_ENV_VAR, attention_backend) - - prompts = example_prompts - if len(prompts) < num_prompts: - prompts = prompts * ((num_prompts // len(prompts)) + 1) - prompts = prompts[:num_prompts] - assert len(prompts) == num_prompts - - with vllm_runner( - model, - dtype=dtype, - enforce_eager=enforce_eager, - gpu_memory_utilization=0.7, - tensor_parallel_size=tp_size, - num_scheduler_steps=num_scheduler_steps, - ) as vllm_model: - vllm_outputs = vllm_model.generate_greedy_logprobs( - prompts, - max_tokens, - num_logprobs, - num_prompt_logprobs=num_prompt_logprobs) - - with vllm_runner( - model, - dtype=dtype, - enforce_eager=enforce_eager, - gpu_memory_utilization=0.7, - tensor_parallel_size=tp_size, - ) as vllm_model: - single_step_vllm_outputs = vllm_model.generate_greedy_logprobs( - prompts, - max_tokens, - num_logprobs, - num_prompt_logprobs=num_prompt_logprobs) - - check_logprobs_close( - outputs_0_lst=single_step_vllm_outputs, - outputs_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.parametrize("tp_size", [1]) -@pytest.mark.parametrize("max_tokens", [5]) -@pytest.mark.parametrize("enforce_eager", [True]) -@pytest.mark.parametrize("num_scheduler_steps", NUM_SCHEDULER_STEPS) -@pytest.mark.parametrize("num_prompts", NUM_PROMPTS) -@pytest.mark.parametrize("num_logprobs", [None, 5]) -@pytest.mark.parametrize("attention_backend", ["FLASH_ATTN"]) -@pytest.mark.skipif( - current_platform.is_rocm(), - reason="Multi-Step + Chunked-Prefill not supported on ROCm") -def test_multi_step_llm_chunked_prefill_prefix_cache( - vllm_runner, - example_prompts, - model: str, - dtype: str, - tp_size: int, - max_tokens: int, - enforce_eager: int, - num_scheduler_steps: int, - num_prompts: int, - num_logprobs: Optional[int], - attention_backend: str, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test vLLM engine with multi-step+"single-step chunked prefill"+APC. - - Set up contrived scenario which tests for a possible failure mode of - scheduling with multi-step+"single-step chunked prefill"+APC - - "single-step chunked prefill" here refers to the current vLLM multi-step+ - chunked-prefill implementation, which requires that a prefill may only - be scheduled in the same step as decodes if the prefill prompt fits in a - single chunk (note that "complete" multi-step+chunked-prefill would allow - a prefill to span multiple chunks & multiple steps but that is not yet - the case.) - - "APC" is short for "automatic prefix caching". - - This test creates a scenario where the scheduler must decide whether/how - to schedule a prefill with a prompt that exceeds the available token budget. - The correct behavior for multi-step+"single-step chunked prefill"+APC is to - put off scheduling the prefill until a future step. - - Validate that: - * Multi-step kernels do not raise an exception due to incorrect scheduler - behavior - * Generated tokens match between - multi-step+"single-step chunked prefill"+APC and - single-step scheduling. - * (If logprobs are enabled) check logprobs are close enough - - Args: - vllm_runner: vLLM model runner fixture - example_prompts: test fixture providing example prompts - model: model under test (same for single- and multi-step engines) - dtype: tensor datatype for engine to utilize - tp_size: degree of tensor-parallelism - max_tokens: the maximum number of tokens to generate - enforce_eager - num_scheduler_steps: for multi-step scheduling, GPU-side steps per - GPU -> CPU output transfer - num_prompts: number of example prompts under test - num_logprobs: corresponds to the `logprobs` argument to the OpenAI - completions endpoint; `None` -> 1 logprob returned. - """ - - # Set up contrived test for correct scheduling behavior with - # multi-step+"single-step chunked prefill"+APC. - # - # Assume block_size=16 - # - # Assume max_num_batched_tokens=48 - # => Per-step token budget=48 - # - # 1. Scheduler schedules 0th prompt (24 tokens) - # => Remaining token budget=24 - # 2. Scheduler attempts to schedule 1st prompt (30 tokens) - # * 30 tokens exceeds 24 token remaining budget - # * Correct behavior: do not schedule this prompt in this step - # * Incorrect behavior: schedule prompt chunk - # * `do_sample=False` for this prompt in this step - # * Chunk size = (remaining tokens // block size) * block size - # - # The Incorrect scheduling behavior - if it occurs - will cause an exception - # in the model runner resulting from `do_sample=False`. - with monkeypatch.context() as m: - m.setenv(STR_BACKEND_ENV_VAR, attention_backend) - - assert len(example_prompts) >= 2 - challenge_prompts = copy.deepcopy(example_prompts) - challenge_prompts[0] = ( - 'vLLM is a high-throughput and memory-efficient ' - 'inference and serving engine for LLMs.\n') # 24 tok - challenge_prompts[1] = ( - 'Briefly describe the major milestones in the ' - 'development of artificial intelligence from 1950 to 2020.\n' - ) # 30 tok - - # If necessary, adjust the length of `challenge_prompts` to match - # `num_prompts` - if len(challenge_prompts) < num_prompts: - challenge_prompts = (challenge_prompts * - ((num_prompts // len(challenge_prompts)) + 1)) - challenge_prompts = challenge_prompts[:num_prompts] - assert len(challenge_prompts) == num_prompts - - # Single-step scheduler baseline - with vllm_runner( - model, - dtype=dtype, - enforce_eager=enforce_eager, - gpu_memory_utilization=0.7, - tensor_parallel_size=tp_size, - num_scheduler_steps=num_scheduler_steps, - max_model_len=48, - max_num_batched_tokens=48, - max_num_seqs=4, - block_size=16, - ) as vllm_model: - outputs_baseline = ( - vllm_model.generate_greedy(challenge_prompts, max_tokens) if - num_logprobs is None else vllm_model.generate_greedy_logprobs( - challenge_prompts, max_tokens, num_logprobs)) - - # multi-step+"single-step chunked prefill"+APC - with vllm_runner( - model, - dtype=dtype, - enforce_eager=enforce_eager, - gpu_memory_utilization=0.7, - tensor_parallel_size=tp_size, - enable_chunked_prefill=True, - enable_prefix_caching=True, - num_scheduler_steps=num_scheduler_steps, - max_model_len=48, - max_num_batched_tokens=48, - max_num_seqs=4, - block_size=16, - ) as vllm_model: - outputs_w_features = ( - vllm_model.generate_greedy(challenge_prompts, max_tokens) if - num_logprobs is None else vllm_model.generate_greedy_logprobs( - challenge_prompts, max_tokens, num_logprobs)) - - if num_logprobs is None: - # No-logprobs test - check_outputs_equal( - outputs_0_lst=outputs_baseline, - outputs_1_lst=outputs_w_features, - name_0="multi-step", - name_1="multi-step+features", - ) - else: - # Yes-logprobs test - check_logprobs_close( - outputs_0_lst=outputs_baseline, - outputs_1_lst=outputs_w_features, - name_0="multi-step", - name_1="multi-step+features", - ) diff --git a/tests/samplers/test_logits_processor.py b/tests/samplers/test_logits_processor.py deleted file mode 100644 index 123f9595e97b..000000000000 --- a/tests/samplers/test_logits_processor.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -from vllm import SamplingParams - -MODELS = ["distilbert/distilgpt2"] - - -@pytest.fixture(scope="function", autouse=True) -def use_v0_only(monkeypatch): - """ - This file tests V0 internals, so set VLLM_USE_V1=0. - """ - monkeypatch.setenv('VLLM_USE_V1', '0') - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -def test_logits_processor_force_generate( - vllm_runner, - example_prompts, - model: str, - dtype: str, -) -> None: - with vllm_runner(model, dtype=dtype) as vllm_model: - tokenizer = vllm_model.llm.get_tokenizer() - repeat_times = 2 - enforced_answers = " vLLM" - vllm_token_ids = tokenizer.encode(enforced_answers, - add_special_tokens=False) - max_tokens = len(vllm_token_ids) * repeat_times - - def pick_vllm(token_ids, logits): - token_id = vllm_token_ids[len(token_ids) % len(vllm_token_ids)] - logits[token_id] = torch.finfo(logits.dtype).max - return logits - - params_with_logprobs = SamplingParams( - logits_processors=[pick_vllm], - prompt_logprobs=3, - max_tokens=max_tokens, - ) - - # test logits_processors when prompt_logprobs is not None - vllm_model.llm._add_request( - example_prompts[0], - params=params_with_logprobs, - ) - - # test prompt_logprobs is not None - vllm_model.llm._add_request( - example_prompts[1], - params=SamplingParams( - prompt_logprobs=3, - max_tokens=max_tokens, - ), - ) - - # test grouped requests - vllm_model.llm._add_request( - example_prompts[2], - params=SamplingParams(max_tokens=max_tokens), - ) - - outputs = vllm_model.llm._run_engine(use_tqdm=False) - - assert outputs[0].outputs[0].text == enforced_answers * repeat_times diff --git a/tests/tpu/lora/test_lora.py b/tests/tpu/lora/test_lora.py index 4c47b8c43caf..636108e98581 100644 --- a/tests/tpu/lora/test_lora.py +++ b/tests/tpu/lora/test_lora.py @@ -30,7 +30,6 @@ def use_v1_only(monkeypatch: pytest.MonkeyPatch): def setup_vllm(num_loras: int, tp: int) -> vllm.LLM: return vllm.LLM(model="Qwen/Qwen2.5-3B-Instruct", - num_scheduler_steps=1, max_model_len=256, max_seq_len_to_capture=256, max_num_seqs=8, diff --git a/tests/utils_/test_utils.py b/tests/utils_/test_utils.py index a2db1ae68434..8be1e103dc65 100644 --- a/tests/utils_/test_utils.py +++ b/tests/utils_/test_utils.py @@ -236,7 +236,6 @@ def test_config_args(parser_with_config, cli_config_file): ['serve', 'mymodel', '--config', cli_config_file]) assert args.tensor_parallel_size == 2 assert args.trust_remote_code - assert not args.multi_step_stream_outputs def test_config_file(parser_with_config): @@ -828,7 +827,6 @@ def test_model_specification(parser_with_config, cli_config_file, ]) assert args.tensor_parallel_size == 2 assert args.trust_remote_code is True - assert args.multi_step_stream_outputs is False assert args.port == 12312 diff --git a/tests/v1/test_oracle.py b/tests/v1/test_oracle.py index a756c89b520f..1f16e92f657e 100644 --- a/tests/v1/test_oracle.py +++ b/tests/v1/test_oracle.py @@ -58,12 +58,6 @@ def test_unsupported_configs(monkeypatch): disable_async_output_proc=True, ).create_engine_config() - with pytest.raises(NotImplementedError): - AsyncEngineArgs( - model=MODEL, - num_scheduler_steps=5, - ).create_engine_config() - with pytest.raises(NotImplementedError): AsyncEngineArgs( model=MODEL, diff --git a/tests/worker/test_model_input.py b/tests/worker/test_model_input.py index ec33d334ab65..2031f41fab87 100644 --- a/tests/worker/test_model_input.py +++ b/tests/worker/test_model_input.py @@ -11,7 +11,6 @@ from vllm.model_executor import SamplingMetadata from vllm.model_executor.pooling_metadata import PoolingMetadata from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata -from vllm.worker.multi_step_model_runner import StatefulModelInput from vllm.worker.pooling_model_runner import ( ModelInputForGPUWithPoolingMetadata) @@ -166,81 +165,3 @@ def test_embedding_model_runner_input(): None) == getattr(attn_metadata, field.name, None) # Pooling metadata is not broadcast. assert received_model_input.pooling_metadata is None - - -def test_multi_step_model_runner_input(): - sampling_metadata = SamplingMetadata( - ["seq_group"], - "selected_token_indices", - "categorized_sample_indices", - "num_prompts", - ) - attn_metadata = AttentionMetadata( - num_prefills=1, - num_prefill_tokens=2, - num_decode_tokens=3, - slot_mapping=torch.zeros(1), - multi_modal_placeholder_index_maps=None, - enable_kv_scales_calculation=True, - ) - frozen_model_input = ModelInputForGPUWithSamplingMetadata( - input_tokens=torch.ones(10), - input_positions=torch.ones(10), - sampling_metadata=sampling_metadata, - attn_metadata=attn_metadata) - - model_input = StatefulModelInput( - frozen_model_input=frozen_model_input, - is_last_step=True, - is_first_multi_step=False, - current_step=4, - last_sampled_token_ids=torch.ones((10, 1)), - is_multi_step=True, - num_queries=8, - num_seqs=5, - cached_outputs=[], - ) - - assert isinstance(model_input, StatefulModelInput) - - # Test round trip serialization. - tensor_dict = model_input.as_broadcastable_tensor_dict() - attn_backend = MockAttentionBackend() - received_model_input = (StatefulModelInput.from_broadcasted_tensor_dict( - tensor_dict, attn_backend=attn_backend)) - - received_frozen_input = received_model_input.frozen_model_input - - # Check that received copy has correct values. - assert isinstance(received_model_input, StatefulModelInput) - assert received_frozen_input.input_tokens is not None - assert (received_frozen_input.input_tokens == - frozen_model_input.input_tokens).all() - assert received_frozen_input.input_positions is not None - assert (received_frozen_input.input_positions == - frozen_model_input.input_positions).all() - assert received_frozen_input.multi_modal_kwargs is None - assert (frozen_model_input.multi_modal_kwargs == - frozen_model_input.multi_modal_kwargs) - assert received_frozen_input.lora_requests is None - assert (received_frozen_input.lora_requests == - frozen_model_input.lora_requests) - assert received_frozen_input.lora_mapping is None - assert ( - received_frozen_input.lora_mapping == frozen_model_input.lora_mapping) - for field in dataclasses.fields(AttentionMetadata): - assert getattr(received_frozen_input.attn_metadata, field.name, - None) == getattr(attn_metadata, field.name, None) - # For sampling metadata, only selected_token_indices is copied. - assert (received_frozen_input.sampling_metadata.selected_token_indices == - sampling_metadata.selected_token_indices) - assert received_frozen_input.sampling_metadata.seq_groups is None - - # check non frozen fields - assert received_model_input.is_last_step == model_input.is_last_step - assert (received_model_input.is_first_multi_step == - model_input.is_first_multi_step) - assert received_model_input.current_step == model_input.current_step - assert (received_model_input.last_sampled_token_ids == - model_input.last_sampled_token_ids).all() - assert received_model_input.is_multi_step == model_input.is_multi_step diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index df4eb33f5d45..6649cd89ee34 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -3779,8 +3779,6 @@ def __str__(self): f"observability_config={self.observability_config!r}, " f"seed={self.model_config.seed}, " f"served_model_name={self.model_config.served_model_name}, " - f"num_scheduler_steps={self.scheduler_config.num_scheduler_steps}, " - f"multi_step_stream_outputs={self.scheduler_config.multi_step_stream_outputs}, " # noqa f"enable_prefix_caching={self.cache_config.enable_prefix_caching}, " f"chunked_prefill_enabled={self.scheduler_config.chunked_prefill_enabled}, " # noqa f"use_async_output_proc={self.model_config.use_async_output_proc}, " diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py index 61346da145bb..63894e7f5dc8 100644 --- a/vllm/core/scheduler.py +++ b/vllm/core/scheduler.py @@ -929,8 +929,7 @@ def _schedule_swapped( ) def _get_prompt_limit(self, seq_group: SequenceGroup) -> int: - if (self.scheduler_config.chunked_prefill_enabled - and not self.scheduler_config.is_multi_step): + if self.scheduler_config.chunked_prefill_enabled: prompt_limit = self.scheduler_config.max_model_len else: prompt_limit = min( @@ -1114,9 +1113,6 @@ def _schedule_prefills( continue num_lookahead_slots: int = 0 - if self.scheduler_config.is_multi_step and enable_chunking: - num_lookahead_slots = self._get_num_lookahead_slots( - True, enable_chunking) # If the sequence group cannot be allocated, stop. can_allocate = self.block_manager.can_allocate( @@ -1195,24 +1191,6 @@ def _schedule_prefills( partial_prefill_metadata.maybe_increment_partial_prefills( seq_group) - if enable_chunking and self.scheduler_config.is_multi_step: - blocks_to_copy: List[Tuple[int, int]] = [] - # init_multi_step_from_lookahead_slots happens in append_slots - self._append_slots(seq_group, blocks_to_copy, enable_chunking) - # This assert will trip when a copy-on-write happens. This is - # not a concern as the very first sequence-group block - # allocation happens above. Still, we have the assert to - # catch any edge-cases. - assert not blocks_to_copy - else: - seq_group.init_multi_step_from_lookahead_slots( - num_lookahead_slots, - num_scheduler_steps=self.scheduler_config. - num_scheduler_steps, - is_multi_step=self.scheduler_config.is_multi_step, - enable_chunking=enable_chunking, - ) - seq_groups.append( ScheduledSequenceGroup(seq_group=seq_group, token_chunk_size=num_new_tokens)) @@ -1453,14 +1431,6 @@ def _schedule_chunked_prefill(self) -> SchedulerOutputs: num_prefill_groups = (len(prefills.seq_groups) + len(swapped_in.prefill_seq_groups) + len(running_scheduled.prefill_seq_groups)) - # If all prompts, then we set num_lookahead_slots to 0 - # this allows us to go through the `no_spec` path in - # `spec_decode_worker.py` - all_prefills = len(scheduled_seq_groups) == num_prefill_groups - num_lookahead_slots = (0 if - (all_prefills - and not self.scheduler_config.is_multi_step) - else running_scheduled.num_lookahead_slots) return SchedulerOutputs( scheduled_seq_groups=scheduled_seq_groups, num_prefill_groups=num_prefill_groups, @@ -1472,7 +1442,7 @@ def _schedule_chunked_prefill(self) -> SchedulerOutputs: swapped_in.blocks_to_copy, ignored_seq_groups=prefills.ignored_seq_groups + swapped_in.infeasible_seq_groups, - num_lookahead_slots=num_lookahead_slots, + num_lookahead_slots=0, running_queue_size=len(self.running), preempted=(len(running_scheduled.preempted) + len(running_scheduled.swapped_out)), @@ -1516,11 +1486,6 @@ def _can_append_slots(self, seq_group: SequenceGroup, num_lookahead_slots = self._get_num_lookahead_slots( is_prefill, enable_chunking) - if is_prefill and num_lookahead_slots > 0: - # Appending prefill slots only happens multi-step and - # chunked-prefill are enabled together. - assert self.scheduler_config.is_multi_step and enable_chunking - return self.block_manager.can_append_slots( seq_group=seq_group, num_lookahead_slots=num_lookahead_slots) @@ -1776,19 +1741,7 @@ def _append_slots( num_lookahead_slots: int = self._get_num_lookahead_slots( is_prefill, enable_chunking) - seq_group.init_multi_step_from_lookahead_slots( - num_lookahead_slots, - num_scheduler_steps=self.scheduler_config.num_scheduler_steps, - is_multi_step=self.scheduler_config.is_multi_step, - enable_chunking=enable_chunking, - ) - seq_status: Optional[SequenceStatus] = SequenceStatus.RUNNING - if self.scheduler_config.is_multi_step and enable_chunking: - # In multi-step chunked-prefill any sequence type can have - # slots appended. - seq_status = None - for seq in seq_group.get_seqs(status=seq_status): cows = self.block_manager.append_slots(seq, num_lookahead_slots) if len(cows) > 0: @@ -1904,29 +1857,8 @@ def _get_num_lookahead_slots(self, is_prefill: bool, """The number of slots to allocate per sequence per step, beyond known token ids. Speculative decoding uses these slots to store KV activations of tokens which may or may not be accepted. - - Speculative decoding does not yet support prefill, so we do not perform - lookahead allocation for prefill. - - When chunking is enabled with multi-step, we allocate lookahead slots - for the prefills for when the prefills turn into decodes in the first - step. """ - if is_prefill: - if self.scheduler_config.is_multi_step and enable_chunking: - # num_lookahead_slots was introduced in the context of decodes, - # in Speculative Decoding. - # When the num_scheduler_steps is 8, say, then the - # num_lookahead_slots is 7. Meaning, we are doing a 1-step of - # decode anyways and we wish to do 7 more. - # - # "lookaheads" for prefills, is introduced in support for - # Chunked-Prefill in Multi-Step. - return self.scheduler_config.num_lookahead_slots + 1 - else: - return 0 - - return self.scheduler_config.num_lookahead_slots + return 0 def _get_num_new_uncached_and_cached_tokens( self, @@ -2068,24 +2000,6 @@ def _chunk_new_tokens_to_schedule( The number of new tokens to schedule after chunking. """ remaining_token_budget = budget.remaining_token_budget() - if scheduler_config.is_multi_step: - # The current multi-step + chunked prefill capability does - # not actually support chunking prompts. - # - # Therefore, `num_new_tokens` is computed in the same fashion - # for both multi-step+chunked-prefill & - # multi-step+chunked-prefill+APC - # - # Prompts with more tokens than the current remaining budget - # are postponed to future scheduler steps - if num_new_tokens > prompt_limit: - # If the seq_group is in prompt-stage, pass the - # num_new_tokens as-is so the caller can ignore - # the sequence. - return num_new_tokens - - return 0 if num_new_tokens > \ - remaining_token_budget else num_new_tokens # Get the number of tokens to allocate to this prefill slot prefill_slot_budget = ( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index d74db67bda0d..c058001ceb97 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -362,8 +362,6 @@ class EngineArgs: lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size - num_scheduler_steps: int = SchedulerConfig.num_scheduler_steps - multi_step_stream_outputs: bool = SchedulerConfig.multi_step_stream_outputs ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight num_gpu_blocks_override: Optional[ int] = CacheConfig.num_gpu_blocks_override @@ -799,11 +797,8 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: **scheduler_kwargs["delay_factor"]) scheduler_group.add_argument("--preemption-mode", **scheduler_kwargs["preemption_mode"]) - scheduler_group.add_argument("--num-scheduler-steps", - **scheduler_kwargs["num_scheduler_steps"]) - scheduler_group.add_argument( - "--multi-step-stream-outputs", - **scheduler_kwargs["multi_step_stream_outputs"]) + # multi-step scheduling has been removed; corresponding arguments + # are no longer supported. scheduler_group.add_argument("--scheduling-policy", **scheduler_kwargs["policy"]) scheduler_group.add_argument( @@ -1257,28 +1252,11 @@ def create_engine_config( disable_log_stats=self.disable_log_stats, ) - # Reminder: Please update docs/features/compatibility_matrix.md - # If the feature combo become valid - if self.num_scheduler_steps > 1: - if speculative_config is not None: - raise ValueError("Speculative decoding is not supported with " - "multi-step (--num-scheduler-steps > 1)") - if self.enable_chunked_prefill and self.pipeline_parallel_size > 1: - raise ValueError("Multi-Step Chunked-Prefill is not supported " - "for pipeline-parallel-size > 1") - if current_platform.is_cpu(): - logger.warning("Multi-Step (--num-scheduler-steps > 1) is " - "currently not supported for CPUs and has been " - "disabled.") - self.num_scheduler_steps = 1 - - # make sure num_lookahead_slots is set the higher value depending on - # if we are using speculative decoding or multi-step - num_lookahead_slots = max(self.num_lookahead_slots, - self.num_scheduler_steps - 1) - num_lookahead_slots = num_lookahead_slots \ - if speculative_config is None \ - else speculative_config.num_lookahead_slots + # make sure num_lookahead_slots is set appropriately depending on + # whether speculative decoding is enabled + num_lookahead_slots = self.num_lookahead_slots + if speculative_config is not None: + num_lookahead_slots = speculative_config.num_lookahead_slots scheduler_config = SchedulerConfig( runner_type=model_config.runner_type, @@ -1292,8 +1270,6 @@ def create_engine_config( disable_chunked_mm_input=self.disable_chunked_mm_input, is_multimodal_model=model_config.is_multimodal_model, preemption_mode=self.preemption_mode, - num_scheduler_steps=self.num_scheduler_steps, - multi_step_stream_outputs=self.multi_step_stream_outputs, send_delta_data=(envs.VLLM_USE_RAY_SPMD_WORKER and parallel_config.use_ray), policy=self.scheduling_policy, @@ -1392,11 +1368,6 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool: recommend_to_remove=True) return False - if self.num_scheduler_steps != SchedulerConfig.num_scheduler_steps: - _raise_or_fallback(feature_name="--num-scheduler-steps", - recommend_to_remove=True) - return False - if self.scheduler_delay_factor != SchedulerConfig.delay_factor: _raise_or_fallback(feature_name="--scheduler-delay-factor", recommend_to_remove=True) diff --git a/vllm/engine/async_llm_engine.py b/vllm/engine/async_llm_engine.py index 1f962b008ee0..b6ee4105340a 100644 --- a/vllm/engine/async_llm_engine.py +++ b/vllm/engine/async_llm_engine.py @@ -15,7 +15,7 @@ from vllm.core.scheduler import SchedulerOutputs from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.async_timeout import asyncio_timeout -from vllm.engine.llm_engine import LLMEngine, SchedulerOutputState +from vllm.engine.llm_engine import LLMEngine from vllm.engine.metrics_types import StatLoggerBase from vllm.engine.protocol import EngineClient from vllm.executor.executor_base import ExecutorBase @@ -308,13 +308,6 @@ async def step_async( if not allow_async_output_proc and len(ctx.output_queue) > 0: self._process_model_outputs(ctx=ctx) - if (self.scheduler_config.is_multi_step - and scheduler_outputs.num_lookahead_slots > 0): - # cache the scheduler outputs for the next iteration if we have - # lookahead slots - self._cache_scheduler_outputs_for_multi_step( - virtual_engine, seq_group_metadata_list, scheduler_outputs, - allow_async_output_proc) else: finished_requests_ids = list() @@ -351,29 +344,14 @@ async def step_async( outputs = await self.model_executor.execute_model_async( execute_model_req) - # we need to do this here so that last step's sampled_token_ids can - # be passed to the next iteration for PP. - if self.scheduler_config.is_multi_step: - self._update_cached_scheduler_output(virtual_engine, outputs) else: if len(ctx.output_queue) > 0: self._process_model_outputs(ctx=ctx) outputs = [] - # Finish the current step for all the sequence groups. - if self.scheduler_config.is_multi_step: - for seq_group in seq_group_metadata_list: - seq_group.finish_step() - if not self._has_remaining_steps(seq_group_metadata_list): - # Clear the cache if we have finished all the steps - if self.scheduler_config.is_multi_step: - self.cached_scheduler_outputs[ - virtual_engine] = SchedulerOutputState() - # is_first_step_output is True only when the num_steps of all - # the sequences are 1. When the num_steps > 1, - # multi_step_model_runner does the first-step output append. + # the sequences are 1. is_first_step_output: bool = False if not seq_group_metadata_list \ else seq_group_metadata_list[0].state.num_steps == 1 diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index 3fc4f6445df2..bbe958351e87 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -25,7 +25,6 @@ from vllm.engine.output_processor.interfaces import ( SequenceGroupOutputProcessor) from vllm.engine.output_processor.stop_checker import StopChecker -from vllm.engine.output_processor.util import create_output_by_sequence_group from vllm.entrypoints.openai.logits_processors import ( get_logits_processors as get_openai_logits_processors) from vllm.executor.executor_base import ExecutorBase @@ -91,7 +90,7 @@ class OutputData(NamedTuple): class SchedulerContext: - def __init__(self, multi_step_stream_outputs: bool = False): + def __init__(self) -> None: self.output_queue: Deque[OutputData] = deque() self.request_outputs: List[Union[RequestOutput, PoolingRequestOutput]] = [] @@ -99,8 +98,6 @@ def __init__(self, multi_step_stream_outputs: bool = False): List[SequenceGroupMetadata]] = None self.scheduler_outputs: Optional[SchedulerOutputs] = None - self.multi_step_stream_outputs: bool = multi_step_stream_outputs - def append_output(self, outputs: List[SamplerOutput], seq_group_metadata_list: List[SequenceGroupMetadata], scheduler_outputs: SchedulerOutputs, is_async: bool, @@ -303,8 +300,7 @@ def get_tokenizer_for_seq(sequence: Sequence) -> AnyTokenizer: ] self.scheduler_contexts = [ - SchedulerContext(multi_step_stream_outputs=self.scheduler_config. - multi_step_stream_outputs) + SchedulerContext() for _ in range(self.parallel_config.pipeline_parallel_size) ] @@ -683,8 +679,7 @@ def add_request( "Priority scheduling is not enabled.") if isinstance(params, SamplingParams) \ - and params.logits_processors \ - and self.scheduler_config.num_scheduler_steps > 1: + and params.logits_processors: raise ValueError( "Logits processors are not supported in multi-step decoding") @@ -868,45 +863,6 @@ def _process_sequence_group_outputs( return - def _update_num_computed_tokens_for_multi_step_prefill( - self, seq_group: SequenceGroup, - seq_group_meta: SequenceGroupMetadata, - is_first_step_output: Optional[bool]): - """ - This function updates num_computed_tokens for prompt sequences - when Multi-Step is enabled. - - seq_group: SequenceGroup to update the num_computed_tokens for. - seq_group_meta: Metadata of the given SequenceGroup. - is_first_step_output: Optional[bool] - - When available, is_first_step_output indicates if the appended - output token is the output of the first-step in multi-step. - A value of None indicates that outputs from all steps in - in multi-step are submitted in a single burst. - """ - - assert self.scheduler_config.is_multi_step - - if not seq_group_meta.is_prompt: - # num_computed_token updates for multi-step decodes happen after - # the tokens are appended to the sequence. - return - - do_update: bool = False - if self.scheduler_config.chunked_prefill_enabled: - # In multi-step + chunked-prefill case, the prompt sequences - # that are scheduled are fully processed in the first step. - do_update = is_first_step_output is None or is_first_step_output - else: - # Normal multi-step decoding case. In this case prompt-sequences - # are actually single-stepped. Always update in this case. - assert seq_group.state.num_steps == 1 - do_update = True - - if do_update: - seq_group.update_num_computed_tokens( - seq_group_meta.token_chunk_size) - def _process_model_outputs(self, ctx: SchedulerContext, request_id: Optional[str] = None) -> None: @@ -939,33 +895,8 @@ def _process_model_outputs(self, has_multiple_outputs: bool = len(outputs) > 1 outputs_by_sequence_group: List[List[SequenceGroupOutput]] - if has_multiple_outputs: - assert self.scheduler_config.is_multi_step or \ - self.speculative_config - # Organize outputs by [step][sequence group] instead of - # [sequence group][step]. - if self.scheduler_config.is_multi_step: - outputs_by_sequence_group = create_output_by_sequence_group( - outputs, len(seq_group_metadata_list)) - elif self.speculative_config: - # Decodes are multi-steps while prefills are not, outputting at - # most 1 token. Separate them so that we can trigger chunk - # processing without having to pad or copy over prompts K times - # to match decodes structure (costly with prompt_logprobs). - num_prefills = sum(sg.is_prompt - for sg in seq_group_metadata_list) - prefills, decodes = outputs[:num_prefills], outputs[ - num_prefills:] - outputs_by_sequence_group = create_output_by_sequence_group( - decodes, - num_seq_groups=len(seq_group_metadata_list) - num_prefills) - outputs_by_sequence_group = [p.outputs for p in prefills - ] + outputs_by_sequence_group - # We have outputs for multiple steps submitted in a single burst, - # so invalidate is_first_step_output. - is_first_step_output = None - else: - outputs_by_sequence_group = outputs + assert not has_multiple_outputs + outputs_by_sequence_group = outputs # Determine the requests we need to operate on if request_id: @@ -1006,13 +937,8 @@ def _process_model_outputs(self, output = [outputs_by_sequence_group[0][i]] if not is_async: - if self.scheduler_config.is_multi_step: - # Updates happen only if the sequence is prefill - self._update_num_computed_tokens_for_multi_step_prefill( - seq_group, seq_group_meta, is_first_step_output) - else: - seq_group.update_num_computed_tokens( - seq_group_meta.token_chunk_size or 0) + seq_group.update_num_computed_tokens( + seq_group_meta.token_chunk_size or 0) if outputs: for o in outputs: @@ -1074,15 +1000,6 @@ def _process_model_outputs(self, for scheduler in self.scheduler: scheduler.free_finished_seq_groups() - # For multi-step without streaming, don't create outputs each iteration - if not is_last_step and not ctx.multi_step_stream_outputs: - # Immediately process request outputs here (if callback is given) - if (finished_now - and self.process_request_outputs_callback is not None): - self.process_request_outputs_callback(ctx.request_outputs) - ctx.request_outputs.clear() - return - # Create the outputs for i in indices: if i in skip or i in finished_before or i in finished_now: @@ -1101,13 +1018,7 @@ def _process_model_outputs(self, if request_output: ctx.request_outputs.append(request_output) - # For multi-step with streaming, create outputs each iteration - if not is_last_step and ctx.multi_step_stream_outputs: - # Immediately process request outputs here (if callback is given) - if self.process_request_outputs_callback is not None: - self.process_request_outputs_callback(ctx.request_outputs) - ctx.request_outputs.clear() - return + # Create outputs only after processing the scheduler's results for seq_group in scheduler_outputs.ignored_seq_groups: params = seq_group.sampling_params @@ -1157,16 +1068,10 @@ def _advance_to_next_step( if seq_group.is_finished(): continue - if self.scheduler_config.is_multi_step: - # Updates happen only if the sequence is prefill - self._update_num_computed_tokens_for_multi_step_prefill( - seq_group, seq_group_metadata, - seq_group.state.num_steps == 1) - else: - token_chunk_size = (seq_group_metadata.token_chunk_size - if seq_group_metadata.token_chunk_size - is not None else 0) - seq_group.update_num_computed_tokens(token_chunk_size) + token_chunk_size = (seq_group_metadata.token_chunk_size + if seq_group_metadata.token_chunk_size + is not None else 0) + seq_group.update_num_computed_tokens(token_chunk_size) if seq_group_metadata.do_sample: assert len(sequence_group_outputs.samples) == 1, ( @@ -1177,16 +1082,8 @@ def _advance_to_next_step( assert len(seq_group.seqs) == 1 seq = seq_group.seqs[0] - if self.scheduler_config.is_multi_step: - is_prefill_append = seq.data.get_num_uncomputed_tokens( - ) == 0 - seq.append_token_id(sample.output_token, sample.logprobs, - sample.output_embed) - if not is_prefill_append: - seq_group.update_num_computed_tokens(1) - else: - seq.append_token_id(sample.output_token, sample.logprobs, - sample.output_embed) + seq.append_token_id(sample.output_token, sample.logprobs, + sample.output_embed) def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]: """Performs one decoding iteration and returns newly generated results. @@ -1289,13 +1186,6 @@ def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]: if not allow_async_output_proc and len(ctx.output_queue) > 0: self._process_model_outputs(ctx=ctx) - if (self.scheduler_config.is_multi_step - and scheduler_outputs.num_lookahead_slots > 0): - # cache the scheduler outputs for the next iteration if we have - # lookahead slots - self._cache_scheduler_outputs_for_multi_step( - virtual_engine, seq_group_metadata_list, scheduler_outputs, - allow_async_output_proc) else: finished_requests_ids = list() @@ -1345,10 +1235,6 @@ def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]: # Raise so the caller is notified that this request failed raise - # We need to do this here so that last step's sampled_token_ids can - # be passed to the next iteration for PP. - if self.scheduler_config.is_multi_step: - self._update_cached_scheduler_output(virtual_engine, outputs) else: # Nothing scheduled => If there is pending async postprocessor, # then finish it here. @@ -1357,19 +1243,9 @@ def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]: # No outputs in this case outputs = [] - # Finish the current step for all the sequence groups. - if self.scheduler_config.is_multi_step: - for seq_group in seq_group_metadata_list: - seq_group.finish_step() - if not self._has_remaining_steps(seq_group_metadata_list): - # clear the cache if we have finished all the steps. - if self.scheduler_config.is_multi_step: - self.cached_scheduler_outputs[0] = SchedulerOutputState() - # is_first_step_output is True only when the num_steps of all - # the sequences are 1. When the num_steps > 1, - # multi_step_model_runner does the first-step output append. + # the sequences are 1. is_first_step_output: bool = False if not seq_group_metadata_list \ else seq_group_metadata_list[0].state.num_steps == 1 @@ -1453,22 +1329,7 @@ def _abort_and_cache_schedule( def _has_remaining_steps( self, seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] ) -> bool: - if (not self.scheduler_config.is_multi_step - or not seq_group_metadata_list): - return False - - # TODO(will) this is a sanity check for nowto make sure that all the - # seqs are on the same steps. Eventually we will want to do some sort of - # dynamic scheduling when doing multi-step decoding. - ref_remaining_steps = seq_group_metadata_list[0].state.remaining_steps - if any([ - seq_group.state.remaining_steps != ref_remaining_steps - for seq_group in seq_group_metadata_list[1:] - ]): - raise AssertionError("All running sequence groups should " - "have the same remaining steps.") - - return ref_remaining_steps > 0 + return False def _cache_scheduler_outputs_for_multi_step( self, virtual_engine: int, @@ -1497,13 +1358,6 @@ def _update_cached_scheduler_output( def _get_last_sampled_token_ids( self, virtual_engine: int) -> Optional[torch.Tensor]: - cached_last_output = self.cached_scheduler_outputs[ - virtual_engine].last_output - if (self.scheduler_config.is_multi_step - and self.parallel_config.pipeline_parallel_size > 1 - and cached_last_output is not None - and cached_last_output.sampled_token_ids_cpu is not None): - return cached_last_output.sampled_token_ids_cpu return None def add_logger(self, logger_name: str, logger: StatLoggerBase) -> None: diff --git a/vllm/engine/output_processor/interfaces.py b/vllm/engine/output_processor/interfaces.py index 19c5963d32db..4d75719c1719 100644 --- a/vllm/engine/output_processor/interfaces.py +++ b/vllm/engine/output_processor/interfaces.py @@ -36,27 +36,13 @@ def create_output_processor( ): """Create an output processor. - This returns a single-step output processor if num_lookahead_slots is - zero, else returns a multi-step output processor. + Multi-step scheduling is no longer supported. Always return a + single-step output processor. """ - if scheduler_config.num_lookahead_slots == 0: - # Importing here to avoid cycle. - from vllm.engine.output_processor.single_step import ( - SingleStepOutputProcessor) - return SingleStepOutputProcessor(scheduler_config, detokenizer, - scheduler, seq_counter, - stop_checker) - else: - # Importing here to avoid cycle. - from vllm.engine.output_processor.multi_step import ( - MultiStepOutputProcessor) - return MultiStepOutputProcessor( - detokenizer, - scheduler, - seq_counter, - get_tokenizer_for_seq, - stop_checker, - ) + from vllm.engine.output_processor.single_step import ( + SingleStepOutputProcessor) + return SingleStepOutputProcessor(scheduler_config, detokenizer, + scheduler, seq_counter, stop_checker) @abstractmethod def process_outputs(self, sequence_group: SequenceGroup, diff --git a/vllm/engine/output_processor/multi_step.py b/vllm/engine/output_processor/multi_step.py deleted file mode 100644 index 8b66ef0dc765..000000000000 --- a/vllm/engine/output_processor/multi_step.py +++ /dev/null @@ -1,211 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import functools -from typing import Callable, List, cast - -from vllm.core.scheduler import Scheduler -from vllm.engine.output_processor.interfaces import ( - SequenceGroupOutputProcessor) -from vllm.engine.output_processor.single_step import ( - single_step_process_prompt_logprob) -from vllm.engine.output_processor.stop_checker import StopChecker -from vllm.logger import init_logger -from vllm.sampling_params import SamplingParams -from vllm.sequence import (VLLM_INVALID_TOKEN_ID, - CompletionSequenceGroupOutput, Sequence, - SequenceGroup, SequenceGroupOutput, SequenceOutput, - SequenceStatus) -from vllm.transformers_utils.detokenizer import Detokenizer -from vllm.transformers_utils.tokenizer import AnyTokenizer -from vllm.utils import Counter - -logger = init_logger(__name__) - - -class MultiStepOutputProcessor(SequenceGroupOutputProcessor): - """SequenceGroupOutputProcessor which handles logic related to - detokenization and stopping conditions. It specializes to "multi-step - decoding", where vLLM's worker may generate multiple tokens per invocation. - This is currently mutually exclusive with advanced sampling techniques like - beam search, which motivates the separation of this logic from the single - step output processor. - - This class is responsible for things such as correctly appending all new - token ids to their sequence, detokenizing new token ids, truncating new - output tokens after an eos token, and correctly handling the case where the - number of new output tokens per sequence differs in a single batch. - """ - - def __init__( - self, - detokenizer: Detokenizer, - scheduler: List[Scheduler], - seq_counter: Counter, - get_tokenizer_for_seq: Callable[[Sequence], AnyTokenizer], - stop_checker: StopChecker, - ): - self.detokenizer = detokenizer - self.scheduler = scheduler - self.seq_counter = seq_counter - self.get_tokenizer_for_seq = get_tokenizer_for_seq - self.stop_checker = stop_checker - - def process_prompt_logprob(self, seq_group: SequenceGroup, - outputs: List[SequenceGroupOutput]) -> None: - """Process prompt logprobs associated with each step of a multi-step- - scheduled computation. - - Args: - seq_group: the outputs are associated with this - [`SequenceGroup`][vllm.sequence.SequenceGroup] - outputs: the - [`SequenceGroupOutput`][vllm.sequence.SequenceGroupOutput]s - for all scheduler steps - """ - for output in outputs: - # Concatenate single-step prompt logprob processing results. - assert isinstance(output, CompletionSequenceGroupOutput) - single_step_process_prompt_logprob(self, seq_group, output) - - @staticmethod - @functools.lru_cache - def _log_prompt_logprob_unsupported_warning_once(): - # Reminder: Please update docs/features/compatibility_matrix.md - # If the feature combo become valid - logger.warning( - "Prompt logprob is not supported by multi step workers. " - "(e.g., speculative decode uses multi step workers).") - - def process_outputs(self, - sequence_group: SequenceGroup, - outputs: List[SequenceGroupOutput], - is_async: bool = False) -> None: - """Append new tokens in the outputs to sequences in the sequence group. - - This only supports sequence groups of size 1. It supports greater than - one new token per sequence. - - This applies logic like stop condition checking and detokenization. - It also handles cases where there are tokens emitted after - the EOS token. - - is_async - Indicates whether this postprocessor runs in - parallel with the GPU forward pass and is processing - tokens from the previous step. If this is true, then - no tokens need to be appended since it is already done - externally (before the next schedule() call) - """ - # Sequences can be in RUNNING or FINISHED_ABORTED state - # once scheduled, as a sequence is moved to FINISHED_ABORTED - # if a client disconnects from the api server. - seqs = sequence_group.get_seqs(status=SequenceStatus.RUNNING) - if seqs is None: - seqs = sequence_group.get_seqs( - status=SequenceStatus.FINISHED_ABORTED) - - assert seqs, "Expected RUNNING or FINISHED_ABORTED sequences" - assert len(seqs) == 1, ( - "Beam search not supported in multi-step decoding.") - seq = seqs[0] - seq_id = seq.seq_id - # This method is defined in the more generic - # SequenceGroupOutputProcessor, but here we assume that the outputs are - # of a more specific type. - assert all([ - isinstance(output, CompletionSequenceGroupOutput) - for output in outputs - ]) - compl_outputs = cast(List[CompletionSequenceGroupOutput], outputs) - assert all([ - seq_id == output.samples[0].parent_seq_id - for output in compl_outputs - ]) - - if is_async: - # Async case: We process tokens one by one. Here, we know the token - # was already appended, so we only need to do the rest of the - # postprocessor: Detokenization + stopping logic - self._process_decode_and_stop(seq, sequence_group.sampling_params) - else: - # Standard multi-step case - - # Since there's only one sequence per sequence group, - # we can take the first sample. - samples = [output.samples[0] for output in compl_outputs] - - # entries in sample tokens may be invalid (eg. due to spec decode - # rejecting tokens). - valid_samples = [ - sample for sample in samples - if sample.output_token != VLLM_INVALID_TOKEN_ID - ] - - # When both spec-decode and pre-fill chunking are enabled, we - # don't have guaranteed samples here (e.g. all -1s). - if valid_samples: - self._process_seq_outputs(seq, valid_samples, - sequence_group.sampling_params) - - def _process_decode_and_stop(self, seq: Sequence, - sampling_params: SamplingParams) -> None: - new_char_count = 0 - if sampling_params.detokenize and self.detokenizer: - new_char_count = self.detokenizer.decode_sequence_inplace( - seq, sampling_params) - - # TODO(sang): Support lora. - self.stop_checker.maybe_stop_sequence( - seq, - new_char_count=new_char_count, - sampling_params=sampling_params, - ) - - def _process_seq_outputs(self, seq: Sequence, - valid_samples: List[SequenceOutput], - sampling_params: SamplingParams) -> None: - output_token_ids = [sample.output_token for sample in valid_samples] - output_logprobs = [sample.logprobs for sample in valid_samples] - output_embeds = [sample.output_embed for sample in valid_samples] - - # Truncate to max_tokens if necessary. - remaining_tokens = sampling_params.max_tokens - (seq.get_output_len() + - len(output_token_ids)) - if remaining_tokens < 0: - output_token_ids = output_token_ids[:remaining_tokens] - - # Truncate any tokens after EOS. This is required as spec decode - # generates a fixed number of tokens without evaluating stopping - # conditions within the block. This can cause an eos token to be - # unintentionally ignored. - if not sampling_params.ignore_eos and self.detokenizer: - eos_token_id = self.get_tokenizer_for_seq(seq).eos_token_id - # Avoiding .index calls as exception throwing in the happy path - # is expensive. - for i in range(len(output_token_ids)): - if output_token_ids[i] == eos_token_id: - output_token_ids = output_token_ids[:i + 1] - break - - is_prefill_sampled_token = seq.data.get_num_uncomputed_tokens() == 0 - # Incrementally append tokens to the sequence, as if we had only one new - # token. - for output_token_id, output_logprob, output_embed in zip( - output_token_ids, output_logprobs, output_embeds): - seq.append_token_id( - token_id=output_token_id, - logprobs=output_logprob, - token_embed=output_embed, - ) - - if is_prefill_sampled_token: - is_prefill_sampled_token = False - else: - # Update num_computed_tokens iff the sampled token is not from - # a prefill step. - seq.data.update_num_computed_tokens(1) - - self._process_decode_and_stop(seq, sampling_params) - - if seq.is_finished(): - break diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index c876c52a2e9c..70959131573f 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -118,20 +118,10 @@ def log_warnings(cls): @classmethod def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: parallel_config = vllm_config.parallel_config - scheduler_config = vllm_config.scheduler_config model_config = vllm_config.model_config if parallel_config.worker_cls == "auto": - if scheduler_config.is_multi_step: - if envs.VLLM_USE_V1: - raise NotImplementedError( - "Multi-step scheduling is not supported (and not " - "needed) on vLLM V1. Please launch without " - "--num-scheduler-steps.") - else: - parallel_config.worker_cls = \ - "vllm.worker.multi_step_worker.MultiStepWorker" - elif vllm_config.speculative_config: + if vllm_config.speculative_config: if not envs.VLLM_USE_V1: raise NotImplementedError( "Speculative decoding is not supported on vLLM V0.") @@ -139,7 +129,7 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: else: if envs.VLLM_USE_V1: parallel_config.worker_cls = \ - "vllm.v1.worker.gpu_worker.Worker" + "vllm.v1.worker.gpu_worker.Worker" else: parallel_config.worker_cls = "vllm.worker.worker.Worker" diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 8005830f55ce..2d5bee5fc505 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -327,18 +327,8 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: cache_config.block_size = 16 parallel_config = vllm_config.parallel_config - scheduler_config = vllm_config.scheduler_config if parallel_config.worker_cls == "auto": - if scheduler_config.is_multi_step: - if envs.VLLM_USE_V1: - raise NotImplementedError( - "Multi-step scheduling is not supported (and not " - "needed) on vLLM V1. Please launch without " - "--num-scheduler-steps.") - else: - parallel_config.worker_cls = \ - "vllm.worker.multi_step_worker.MultiStepWorker" - elif vllm_config.speculative_config: + if vllm_config.speculative_config: if not envs.VLLM_USE_V1: raise NotImplementedError( "Speculative decoding is not supported on vLLM V0.") @@ -346,7 +336,7 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: else: if envs.VLLM_USE_V1: parallel_config.worker_cls = \ - "vllm.v1.worker.gpu_worker.Worker" + "vllm.v1.worker.gpu_worker.Worker" else: parallel_config.worker_cls = "vllm.worker.worker.Worker" diff --git a/vllm/platforms/tpu.py b/vllm/platforms/tpu.py index c56096d93612..c7522a89c257 100644 --- a/vllm/platforms/tpu.py +++ b/vllm/platforms/tpu.py @@ -133,18 +133,13 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: parallel_config = vllm_config.parallel_config scheduler_config = vllm_config.scheduler_config if parallel_config.worker_cls == "auto": - if scheduler_config.is_multi_step: - raise NotImplementedError( - "Multi-step scheduling is not supported (and not " - "needed) on vLLM V1. Please launch without " - "--num-scheduler-steps.") parallel_config.worker_cls = "vllm.v1.worker.tpu_worker.TPUWorker" assert not vllm_config.speculative_config, ( "Speculative decoding is not yet supported for TPU backend") if scheduler_config.is_multimodal_model and not \ - scheduler_config.disable_chunked_mm_input: + scheduler_config.disable_chunked_mm_input: logger.warning("TPU does not support running Multimodal models"\ " without setting `--disable_chunked_mm_input`. " \ "Forcing --disable_chunked_mm_input.") diff --git a/vllm/sequence.py b/vllm/sequence.py index 6e65a2bd0318..cbe63f8d1d4e 100644 --- a/vllm/sequence.py +++ b/vllm/sequence.py @@ -794,35 +794,6 @@ def multi_modal_placeholders(self) -> MultiModalPlaceholderDict: def lora_int_id(self) -> int: return self.lora_request.lora_int_id if self.lora_request else 0 - def init_multi_step(self, num_steps: int) -> None: - self.state.num_steps = num_steps - self.state.current_step = 0 - - def init_multi_step_from_lookahead_slots(self, num_lookahead_slots: int, - num_scheduler_steps: int, - is_multi_step: bool, - enable_chunking: bool) -> None: - - if not is_multi_step: - self.init_multi_step(num_steps=num_scheduler_steps) - return - - # Multi-Step case - is_prefill = self.is_prefill() - - # The asserts below reflect the expectations of the current system. - if is_prefill and enable_chunking: - assert num_lookahead_slots == num_scheduler_steps - self.init_multi_step(num_steps=num_lookahead_slots) - else: - is_decode: bool = not is_prefill - # If it is a prefill, num_lookahead_slots must be 0 - assert num_lookahead_slots == 0 or is_decode - # If it is a decode, num_lookahead_slots + 1 must match - # the scheduler steps. - assert num_lookahead_slots + 1 == num_scheduler_steps or is_prefill - self.init_multi_step(num_steps=num_lookahead_slots + 1) - def set_last_token_time(self, now: float) -> None: """Sets the last token time for Request level timings.""" # If still in prefill phase, assertion fails. @@ -1367,15 +1338,6 @@ class ExecuteModelRequest( # Async callback async_callback: Optional[Callable] = None - @property - def is_first_multi_step(self) -> bool: - # TODO(will) make this be able to handle batches with variable number of - # steps - assert len(self.seq_group_metadata_list) > 0 - first_seq_group = self.seq_group_metadata_list[0] - assert first_seq_group.state is not None - return first_seq_group.state.current_step == 0 - @property def is_last_step(self) -> bool: # TODO(will) make this be able to handle batches with variable number of diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py index 20b9b733cd3b..a63797e3a46a 100644 --- a/vllm/worker/model_runner.py +++ b/vllm/worker/model_runner.py @@ -508,8 +508,7 @@ def _compute_lens(self, inter_data: InterDataForSeqGroup, seq_idx: int, if inter_data.is_prompt: context_len = seq_data.get_num_computed_tokens() seq_len = min(seq_len, context_len + token_chunk_size) - elif self.runner.scheduler_config.is_multi_step or \ - self.runner.model_config.is_encoder_decoder: + elif self.runner.model_config.is_encoder_decoder: context_len = seq_len - 1 else: context_len = seq_data.get_num_computed_tokens() @@ -778,9 +777,7 @@ def _get_cuda_graph_pad_size(self, int: Returns the determined number of padding sequences. If CUDA graphs is not viable, returns -1. """ - is_mscp: bool = self.runner.scheduler_config.is_multi_step and \ - self.runner.scheduler_config.chunked_prefill_enabled - decode_only = self.decode_only or is_mscp + decode_only = self.decode_only if not decode_only: # Early exit so we can treat num_seqs as the batch_size below. return -1 diff --git a/vllm/worker/multi_step_model_runner.py b/vllm/worker/multi_step_model_runner.py deleted file mode 100644 index 2aa910bdff6b..000000000000 --- a/vllm/worker/multi_step_model_runner.py +++ /dev/null @@ -1,908 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import dataclasses -import functools -from dataclasses import dataclass, field -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, - Union) - -import torch - -from vllm.distributed import get_pp_group -from vllm.logger import init_logger -from vllm.model_executor.layers.sampler import (PromptLogprobs, SampleLogprobs, - SamplerOutput, - SamplingMetadata, get_logprobs, - get_pythonized_sample_results) -from vllm.platforms import current_platform -from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors, - Logprob, SequenceGroupMetadata, SequenceOutput) -from vllm.utils import PyObjectCache, async_tensor_h2d, current_stream -from vllm.worker.model_runner import (GPUModelRunnerBase, - ModelInputForGPUWithSamplingMetadata) -from vllm.worker.model_runner_base import ( - BroadcastableModelInput, _init_attn_metadata_from_tensor_dict, - _init_frozen_model_input_from_tensor_dict, - _init_sampling_metadata_from_tensor_dict) - -from ..model_executor.model_loader.tensorizer import TensorizerConfig - -if TYPE_CHECKING: - from vllm.attention.backends.abstract import AttentionBackend - -logger = init_logger(__name__) - -MULTI_STEP_ATTENTION_BACKENDS = [ - "FLASH_ATTN", "ROCM_FLASH", "FLASHINFER", "NO_ATTENTION" -] -MULTI_STEP_CHUNKED_PREFILL_ATTENTION_BACKENDS = ["FLASH_ATTN", "FLASHINFER"] - -def _get_supported_attention_backends(chunked_prefill_enabled: bool) \ - -> List[str]: - if chunked_prefill_enabled: - return MULTI_STEP_CHUNKED_PREFILL_ATTENTION_BACKENDS - else: - return MULTI_STEP_ATTENTION_BACKENDS - - -def seq_output_builder(): - return SequenceOutput( - 0, 0, - {0: Logprob(logprob=float('inf'), rank=None, decoded_token=None)}) - - -def completion_seq_group_output_builder(): - return CompletionSequenceGroupOutput([], None) - - -# Used by pythonization to reduce python object allocations -class PythonizationCache: - - def __init__(self): - self.cached_seq_output = PyObjectCache(seq_output_builder) - self.cached_completion_seq_group_output = PyObjectCache( - completion_seq_group_output_builder) - - def reset(self): - self.cached_seq_output.reset() - self.cached_completion_seq_group_output.reset() - - -@dataclass -class ModelOutput: - """The output of a single model forward pass. - - The sampler_output_ready_event is set when the tensors in - sampler_output are ready (the model+sampler forward pass has - completed). We use the event to synchronize the GPU->CPU transfer, - which we want to only run when the data has been written to the - GPU tensors. Until the event is ready, the tensors in sampler_output - will have garbage data. - - There are two scenarios: - 1. The output tensors are ready and we can pythonize them immediately. - 2. The output tensors are not ready and we need to wait for the event to be - ready. - """ - sampler_output: SamplerOutput - sampler_output_ready_event: torch.cuda.Event - sampled_token_ids: Optional[torch.Tensor] = None - pythonized: bool = False - # On-device tensor containing the logprobs of each token. - logprobs: Optional["torch.Tensor"] = None - pythonization_cache: Optional[PythonizationCache] = None - - def pythonize(self, input_metadata: "StatefulModelInput", - copy_stream: torch.cuda.Stream, - pinned_sampled_token_buffer: torch.Tensor) -> None: - """Pythonize the output. Blocking.""" - if not self.pythonized: - self._pythonize_sampler_output(input_metadata, copy_stream, - pinned_sampled_token_buffer, True) - self.pythonized = True - - def maybe_pythonize(self, input_metadata: "StatefulModelInput", - copy_stream: torch.cuda.Stream, - pinned_sampled_token_buffer: torch.Tensor) -> None: - """Pythonize the output if ready, else return None. Non-blocking.""" - if not self.pythonized: - self.pythonized = self._pythonize_sampler_output( - input_metadata, copy_stream, pinned_sampled_token_buffer, - False) - - def _pythonize_sampler_output(self, input_metadata: "StatefulModelInput", - copy_stream: torch.cuda.Stream, - pinned_sampled_token_buffer: torch.Tensor, - blocking: bool) -> bool: - """ - If blocking is set, will block until the forward pass for the output is - ready and pythonize the output. Upon completing Pythonization, erases - self.logprobs (note that a non-blocking call that is performed when - the sampler output is not yet ready, will not erase self.logprobs.) - """ - assert self.sampled_token_ids is not None - if not blocking and not self.sampler_output_ready_event.query(): - return False - - if blocking: - self.sampler_output_ready_event.synchronize() - with torch.cuda.stream(copy_stream): - _pythonize_sampler_output(input_metadata, self.sampler_output, - pinned_sampled_token_buffer, - self.sampled_token_ids, self.logprobs, - self.pythonization_cache) - - # Erase the logprobs GPU-side tensor. - # Note that although _pythonize_sampler_output() runs in its - # own CUDA stream, nonetheless _pythonize_sampler_output() - # cannot return until Pythonization is complete; therefore - # we know that by the time the CPU reaches this point, - # `self.logprobs` is no longer needed. - self.logprobs = None - return True - - -@dataclass(frozen=False) -class StatefulModelInput(BroadcastableModelInput): - # actual frozen model input dataclass passed to _base_model_runner - frozen_model_input: Optional[ModelInputForGPUWithSamplingMetadata] = None - - # list of model outputs for each step, may not be all pythonized - cached_outputs: List[ModelOutput] = field(default_factory=list) - - # used to pass sampled token ids from the last step to the current step for - # TP workers. Used to append to end of outputs and used by advance_step - last_sampled_token_ids: Optional[torch.Tensor] = None - current_step: int = 0 - is_multi_step: bool = True - is_last_step: bool = False - is_first_multi_step: bool = False - base_output_proc_callback: Optional[Callable] = None - # ping-pong data structures for multi-step to wait on the previous step - step_cuda_events: List[current_platform.Event] = field( - default_factory=lambda: [current_platform.Event(blocking=True)] * 2) - num_seqs: int = -1 - num_queries: int = -1 - num_single_step_prefills: int = 0 - - def as_broadcastable_tensor_dict(self) -> Dict[str, Any]: - assert self.frozen_model_input is not None - tensor_dict = self.frozen_model_input.as_broadcastable_tensor_dict() - new_tensor_dict = { - 'last_sampled_token_ids': self.last_sampled_token_ids, - 'current_step': self.current_step, - 'is_multi_step': self.is_multi_step, - 'is_last_step': self.is_last_step, - 'is_first_multi_step': self.is_first_multi_step, - 'num_seqs': self.num_seqs, - 'num_queries': self.num_queries, - 'num_single_step_prefills': self.num_single_step_prefills, - } - tensor_dict.update(new_tensor_dict) - return tensor_dict - - @classmethod - def from_broadcasted_tensor_dict( - cls, - tensor_dict: Dict[str, Any], - attn_backend: Optional["AttentionBackend"] = None, - ) -> "StatefulModelInput": - tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict) - if attn_backend is not None: - tensor_dict = _init_attn_metadata_from_tensor_dict( - attn_backend, tensor_dict) - tensor_dict = _init_frozen_model_input_from_tensor_dict( - ModelInputForGPUWithSamplingMetadata, tensor_dict) - - return cls(**tensor_dict) - - def record_step_event(self, current_stream: torch.cuda.Stream): - # record the event for the current step so that the next step can sync - # on it. We modulo by 2 to keep the events in a circular buffer and - # support any attn backends that may be supported in the future. ie - # Flashinfer would want two DecodeWrappers to overlap the CPU and GPU. - self.step_cuda_events[self.current_step & 1] = \ - torch.cuda.Event(blocking=True) - self.step_cuda_events[self.current_step & 1].record(current_stream) - - def wait_previous_step(self): - # These cuda events are an explicit synchronization to ensure that - # advance_step() (for other attn backends that may be supported in the - # future) do not clobber any data structures that is also used by any - # enqueued forwards steps. For distributed case, only a single event is - # needed, but for single GPU case, since we can let the CPU run much - # further ahead, two events allow us to overlap the advance_step with - # the previous forward (ie using two DecodeWrappers for flashinfer - # backend) - self.step_cuda_events[(self.current_step + 1) & 1].wait() - - def add_sampler_output(self, - sampler_output: SamplerOutput, - sampled_token_ids: Optional[torch.Tensor] = None): - self.cached_outputs.append( - ModelOutput(sampler_output=sampler_output, - sampler_output_ready_event=None, - sampled_token_ids=sampled_token_ids, - pythonized=False)) - - def maybe_advance_sampling_metadata(self, device: str, pin_memory: bool): - """ - sampling_metadata.selected_token_indices is constructed for the - first-step in Multi-Step. However, when chunked-prefill is enabled with - multi-step, the scheduled prompts are fully processed in the - first-step and are processed as decodes in the rest of the steps. - This function updates the sampling_metadata.selected_token_indices - to account for this conversion. - - Example: - Let 2 prompts and 2 decodes be scheduled together. Let the - num-tokens to process for the 2 prompts be 5 and 8 respectively. - - In that case, sampling_metadata.sampled_token_indices will be, - [4, 12, 13, 14] as it is constructed for the first-step in - multi-step. - However, the prompts turns to decodes after the first-step - and the num-tokens for the previously-prompt sequences will - be 1 and 1 as they are decodes now. The self.sampled_token_indices - must be updated to [0,1,2,3]. - """ - assert self.current_step == 1 and self.num_single_step_prefills > 0 - if not get_pp_group().is_last_rank: - return - - assert self.frozen_model_input is not None - assert self.frozen_model_input.sampling_metadata is not None - self.frozen_model_input.sampling_metadata.selected_token_indices = \ - async_tensor_h2d(list(range(self.num_queries)), - dtype=torch.long, - target_device=device, - pin_memory=pin_memory) - - def maybe_advance_frozen_model_input(self, device: str, pin_memory: bool): - """ - Advancing the datastructures of StatefulModelInput::frozen_model_input - is only required when prefills are scheduled with decodes to run in - multi-step. This advancement/correction is required to account for - the conversion of Prefills to Decodes after the first multi-step. - """ - if self.current_step != 1 or self.num_single_step_prefills == 0: - return - - assert self.frozen_model_input is not None - fmi = self.frozen_model_input - - # Truncate input_tokens - assert fmi.input_tokens is not None - assert fmi.input_tokens.shape[0] >= self.num_seqs - fmi_new_input_tokens: torch.Tensor = fmi.input_tokens[:self.num_seqs] - - # Update frozen_model_input::input_positions. - assert fmi.input_positions is not None - assert fmi.input_positions.shape[0] >= self.num_seqs - fmi_new_input_positions: torch.Tensor = fmi.input_positions[:self. - num_seqs] - - # Assert unsupported - assert fmi.lora_mapping is None - assert fmi.lora_requests is not None - assert len(fmi.lora_requests) == 0 - assert fmi.attn_metadata is not None - assert fmi.multi_modal_kwargs is not None - assert len(fmi.multi_modal_kwargs) == 0 - - self.frozen_model_input = dataclasses.replace( - self.frozen_model_input, - input_tokens=fmi_new_input_tokens, - input_positions=fmi_new_input_positions) - - self.maybe_advance_sampling_metadata(device, pin_memory) - - -# MutableModelInputForGPUWithMultiStepMetadata is not subclass of -# ModelInputForGPU but it wraps the actual input dataclass and adds multi-step -# metadata -# mypy: disable-error-code=type-var -class MultiStepModelRunner(GPUModelRunnerBase[StatefulModelInput]): - # mypy: enable-error-code=type-var - - def __init__(self, base_model_runner: GPUModelRunnerBase, *args, **kwargs): - - super().__init__(*args, **kwargs) - - # Check attention backend support. - supported_attention_backends: List[str] = \ - _get_supported_attention_backends( - self.scheduler_config.chunked_prefill_enabled) - if self.attn_backend.get_name() not in supported_attention_backends: - ms_config_str: str = "Multi-Step + Chunked-Prefill" \ - if self.scheduler_config.chunked_prefill_enabled \ - else "Multi-Step" - raise ValueError( - f"{ms_config_str} not supported for attention backend: " - f"{self.attn_backend.get_name()}. Set VLLM_ATTENTION_BACKEND " - f"to a value from {supported_attention_backends}.") - - # uses the base model runner to execute the model and wraps it with - # multi-step logic - self._base_model_runner: GPUModelRunnerBase = base_model_runner - - self.is_multi_step = self.scheduler_config.is_multi_step - self.pinned_sampled_token_ids: Optional[torch.Tensor] = None - - # Using the PythonizationCache in Pipeline-Parallel clobbers the - # SequenceOutput and CompletionSequenceGroupOutput object. - # When cache-reset happens at the last step of a multi-step - # execution, there may be other on-going single-step/multi-step - # executions. The current caching implementation does not check - # for this. - self.pythonization_cache = PythonizationCache() \ - if self.parallel_config.pipeline_parallel_size == 1 else None - - @functools.cached_property - def _copy_stream(self): - # used to copy tensors from GPU to CPU asynchronously - return torch.cuda.Stream() - - def make_model_input_from_broadcasted_tensor_dict( - self, tensor_dict: Dict[str, Any]) -> StatefulModelInput: - model_input = (StatefulModelInput.from_broadcasted_tensor_dict( - tensor_dict, - attn_backend=self.attn_backend, - )) - return model_input - - def prepare_model_input( - self, - seq_group_metadata_list: List[SequenceGroupMetadata], - virtual_engine: int = 0, - finished_requests_ids: Optional[List[str]] = None - ) -> StatefulModelInput: - frozen_model_input: ModelInputForGPUWithSamplingMetadata = \ - self._base_model_runner.prepare_model_input( - seq_group_metadata_list, - virtual_engine, - finished_requests_ids) - - assert frozen_model_input.query_lens is not None - assert frozen_model_input.seq_lens is not None - assert frozen_model_input.attn_metadata is not None - num_queries = len(frozen_model_input.query_lens) - num_seqs = len(frozen_model_input.seq_lens) - num_single_step_prefills = frozen_model_input.attn_metadata.num_prefills - - model_input = StatefulModelInput( - frozen_model_input=frozen_model_input, - num_seqs=num_seqs, - num_queries=num_queries, - num_single_step_prefills=num_single_step_prefills) - - return model_input - - def _async_process_outputs(self, model_input: StatefulModelInput, - output_proc_callback: Callable): - # Proceed with pythonization and output_proc in order. - # Stop on the first one that fails to pythonize - output_proc_callback() - - cont = True - for step_num, model_output in enumerate(model_input.cached_outputs): - if not model_output.pythonized: - model_output.maybe_pythonize(model_input, self._copy_stream, - self.pinned_sampled_token_ids) - if model_output.pythonized: - ctx = output_proc_callback.keywords["ctx"] - ctx.append_output( - outputs=[model_output.sampler_output], - seq_group_metadata_list=ctx.seq_group_metadata_list, - scheduler_outputs=ctx.scheduler_outputs, - is_async=False, - is_last_step=False, - is_first_step_output=step_num == 0) - - output_proc_callback() - else: - cont = False - - if not cont: - break - - def _final_process_outputs( - self, model_input: StatefulModelInput, - output_proc_callback: Optional[Callable]) -> List[SamplerOutput]: - assert model_input.frozen_model_input is not None - - has_async_callback = output_proc_callback is not None - - outputs = [] - for step_num, output in enumerate(model_input.cached_outputs): - is_last_step = step_num == len(model_input.cached_outputs) - 1 - - # For non-async case: - # -- We simply add the outputs - # For async case: - # -- Invoke callback, pythonize, add to callback queue and repeat - # -- For last output, just add to callback queue - if has_async_callback: - assert output_proc_callback is not None - - # Invoke callback before pythonize (to overlap with GPU) - output_proc_callback() - - # Pythonize - if not output.pythonized: - output.pythonize(model_input, self._copy_stream, - self.pinned_sampled_token_ids) - - # For non last step, add to callback queue to chain - # callbacks=>pythonize pairs (for GPU overlap) - if not is_last_step: - ctx = output_proc_callback.keywords[ # type: ignore - "ctx"] # type: ignore - ctx.append_output( - outputs=[output.sampler_output], - seq_group_metadata_list=ctx. - seq_group_metadata_list, - scheduler_outputs=ctx.scheduler_outputs, - is_async=False, - is_last_step=False, - is_first_step_output=step_num == 0) - else: - outputs.append(output.sampler_output) - else: - output.pythonize(model_input, self._copy_stream, - self.pinned_sampled_token_ids) - outputs.append(output.sampler_output) - - return outputs - - @torch.inference_mode() - def execute_model( - self, - model_input: StatefulModelInput, - kv_caches: List[torch.Tensor], - intermediate_tensors: Optional[IntermediateTensors] = None, - num_steps: int = 1, - ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]: - """ - Execute the model for a single step and update multi-step - metadata - """ - assert num_steps == 1, "MultiStepModelRunner only supports num_steps=1" - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input is not None - - # path for warm up runs - if not model_input.is_multi_step: - return self._base_model_runner.execute_model( - frozen_model_input, None, intermediate_tensors, num_steps) - - # make sure we skip the sampler on the lask rank and only pythonize - # if CPU is ahead. - if self.is_driver_worker and get_pp_group().is_last_rank: - if self.pinned_sampled_token_ids is None: - self.pinned_sampled_token_ids = torch.zeros( - (self.scheduler_config.max_num_seqs, 1), - dtype=torch.long, - device="cpu", - pin_memory=True) - - self._base_model_runner.sampler.include_gpu_probs_tensor = True - if frozen_model_input.sampling_metadata: - frozen_model_input.sampling_metadata.skip_sampler_cpu_output = ( - True) - - # some pre-execute model logic for multi-step: - # - if it's the first step, we need to reset the sampling tensors - # - if it's not the first step, we need to advance the step using the - # appended sampler output from last iteration - # - also maybe pythonize if CPU is ahead of GPU - - stream = current_stream() - if not model_input.is_first_multi_step: - # Explicitly block on the previous step's forward to make sure we - # don't clobber any GPU tensors still in use. - # This is not needed for flashattn backend, but for other attn - # backends such as flashinfer that performs extra CPU operations on - # input metadata we may need to synchronize any CPU operations that - # might clobber enqueued forwards. (prevents CPU from running too - # far ahead if needed) - model_input.wait_previous_step() - model_input = self._advance_step( - model_input, model_input.cached_outputs[-1].sampler_output) - - # frozen_model_input may have been updated - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input is not None - - if model_input.base_output_proc_callback is None: - assert frozen_model_input is not None - model_input.base_output_proc_callback = \ - frozen_model_input.async_callback - - if frozen_model_input.async_callback is not None: - assert model_input.base_output_proc_callback is not None - async_callback = functools.partial( - self._async_process_outputs, - model_input=model_input, - output_proc_callback=model_input.base_output_proc_callback) - - model_input.frozen_model_input = dataclasses.replace( # type: ignore - model_input.frozen_model_input, - async_callback=async_callback) - # Update the local instance - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input is not None - - # Execute the model - output = self._base_model_runner.execute_model(frozen_model_input, - None, - intermediate_tensors, - num_steps=1) - - # record the event for the current step so that the next step can sync - model_input.record_step_event(stream) - - if get_pp_group().is_last_rank and self.is_driver_worker: - assert isinstance(output, list) - assert len( - output - ) == 1, "MultiStepModelRunner requires single-step base_models" - - # event for the pythonization so that we only pythonize if the - # tensors are ready. May be able to be combined with the step event - output_ready_event = torch.cuda.Event() - output_ready_event.record(stream) - if self.parallel_config.pipeline_parallel_size > 1: - output[0].sampled_token_ids_cpu = output[ - 0].sampled_token_ids.cpu() - model_input.cached_outputs.append( - ModelOutput(output[0], output_ready_event, - output[0].sampled_token_ids, False, - output[0].logprobs, self.pythonization_cache)) - - # These GPU tensors are not required by multi-step; - # erase them to ensure they are not pythonized or - # transferred to CPU - output[0].sampled_token_ids = None - output[0].sampled_token_probs = None - output[0].logprobs = None - - # Pythonize the output if CPU is ahead and the previous step is - # ready. - if frozen_model_input.async_callback is None: - for model_output in model_input.cached_outputs: - model_output.maybe_pythonize(model_input, - self._copy_stream, - self.pinned_sampled_token_ids) - - model_input.current_step += 1 - - if not get_pp_group().is_last_rank: - # Should be IntermediateTensors - assert isinstance(output, IntermediateTensors) - return output - if not self.is_driver_worker: - return [] - - # Pythonize the output and block if needed since it is the last step - if model_input.is_last_step: - outputs = self._final_process_outputs( - model_input, model_input.base_output_proc_callback) - if self.pythonization_cache: - self.pythonization_cache.reset() - return outputs - - # should be [SamplerOutput] - return output - - def _update_sampling_metadata(self, sampling_metadata: SamplingMetadata, - num_seqs: Optional[int], num_queries: int): - - assert sampling_metadata.num_prompts == 0 - assert len(sampling_metadata.seq_groups) == num_queries - assert sampling_metadata.selected_token_indices.shape == ( - num_queries, ) - # assert sampling_metadata.categorized_sample_indices == TODO: Add if needed # noqa: E501 - - # Verify that all sequences are decodes - for i in range(num_queries): - seq_group = sampling_metadata.seq_groups[i] - - assert seq_group.is_prompt is False # No prompt - assert seq_group.prompt_logprob_indices == [] # No prompt - assert seq_group.sample_indices == [i] # Simple - assert seq_group.seq_len is None # Decode - assert seq_group.query_len is None # Decode - - def _advance_step(self, model_input: StatefulModelInput, - out: SamplerOutput) -> StatefulModelInput: - - model_input.maybe_advance_frozen_model_input(self.device, - self.pin_memory) - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input is not None - assert frozen_model_input.input_tokens is not None - assert frozen_model_input.input_tokens.shape[0] == model_input.num_seqs - assert frozen_model_input.attn_metadata is not None - - sampled_token_ids = model_input.cached_outputs[-1].sampled_token_ids - num_seqs = model_input.num_seqs - num_queries = model_input.num_queries - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input is not None - attn_metadata = frozen_model_input.attn_metadata - assert attn_metadata is not None - - turn_prefills_into_decodes: bool = model_input.current_step == 1 and \ - model_input.num_single_step_prefills != 0 - attn_metadata.advance_step( - frozen_model_input, - sampled_token_ids, - self.block_size, - num_seqs, - num_queries, - turn_prefills_into_decodes=turn_prefills_into_decodes) - - return model_input - - def load_model(self) -> None: - self._base_model_runner.load_model() - self.model_memory_usage = self._base_model_runner.model_memory_usage - - def save_sharded_state( - self, - path: str, - pattern: Optional[str] = None, - max_size: Optional[int] = None, - ) -> None: - return self._base_model_runner.save_sharded_state( - path, pattern, max_size) - - def save_tensorized_model(self, - tensorizer_config: TensorizerConfig) -> None: - return self._base_model_runner.save_tensorized_model(tensorizer_config) - - def profile_run(self) -> None: - return self._base_model_runner.profile_run() - - def remove_all_loras(self): - return self._base_model_runner.remove_all_loras() - - def capture_model(self, kv_caches: List[List]) -> None: - return self._base_model_runner.capture_model(kv_caches) - - @property - def vocab_size(self) -> int: - return self._base_model_runner.vocab_size - - -DeferredLogprobsReturnType = Tuple[Optional[List[Optional[PromptLogprobs]]], - Optional[List[SampleLogprobs]]] - - -def deferred_pythonize_logprobs( - output: SamplerOutput, - sampling_metadata: SamplingMetadata, - logprobs_tensor: Optional[torch.Tensor], -) -> DeferredLogprobsReturnType: - """Perform deferred logprob Pythonization. - - 1. Pythonize GPU-side sampler result tensors into CPU-side sampler result. - 2. Pythonize GPU-side logprobs tensor into CPU-side logprobs lists, - utilizing the Pythonized sampler result computed in step 1. - - These deferred computations are not required for single-step scheduling - or the `profile_run()` phase of multi-step scheduling. - - Args: - output: sampler output (under deferred Pythonization) - sampling_metadata - - Returns: - prompt_logprobs (CPU), sample_logprobs (CPU) - """ - - # - Deferred pythonization of sample result - sampler_result = get_pythonized_sample_results( - output.deferred_sample_results_args) - - # - Erase the GPU-side deferred sample_result - # computation args to ensure it is never - # pythonized or transferred to CPU - output.deferred_sample_results_args = None - - # - Deferred pythonization of logprobs - ( - prompt_logprobs, - sample_logprobs, - ) = get_logprobs(logprobs_tensor, sampling_metadata, sampler_result) - assert len(prompt_logprobs) == len(sampling_metadata.seq_groups) - assert len(sample_logprobs) == len(sampling_metadata.seq_groups) - - return prompt_logprobs, sample_logprobs - - -def _pythonize_sampler_output( - model_input: StatefulModelInput, - output: SamplerOutput, - pinned_sampled_token_buffer: torch.Tensor, - sampled_token_ids: torch.Tensor, - logprobs_tensor: Optional[torch.Tensor], - cache: Optional[PythonizationCache], -) -> None: - """ This function is only called when the output tensors are ready. - See [`ModelOutput`][vllm.worker.multi_step_model_runner.ModelOutput]. - - Modifies `output.outputs` and `pinned_sampled_token_buffer` in-place, - adding a Pythonized output data structure - ([`CompletionSequenceGroupOutput`][vllm.sequence.CompletionSequenceGroupOutput]) - for each [`SequenceGroup`][vllm.sequence.SequenceGroup]. - - Args: - model_input - output: sampler output - pinned_sampled_token_token_buffer: CPU-side pinned memory - (receives copy of - GPU-side token buffer.) - sampled_token_ids: GPU-side token buffer - logprobs_tensor: GPU-side tensor containing - logprobs computed during sampling - """ - - assert model_input.frozen_model_input is not None - - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input.sampling_metadata is not None - sampling_metadata = frozen_model_input.sampling_metadata - # samples generation should have been skipped - assert not output.outputs - - pinned_buffer = pinned_sampled_token_buffer[:model_input.num_queries] - - # We guarantee output tensors are ready, so it is safe to - # pythonize the sampler output & obtain CPU-side logprobs. - # - # However we should check whether logprobs pythonization may - # be skipped entirely, i.e. because no logprobs were requested - # or pythonization was not deferred. To that end, - # - # * `prompt_logprobs_are_requested_for_prefill` signals that - # there are *any* prefill-phase requests which specify that - # prompt logprobs should be returned. - # - # * `any_logprobs_are_requested` signals that there are any - # requests which (1) specify that sample logprobs should be - # returned, or (2) are in the prefill phase AND specify that - # prompt logprobs should be returned. - # - # Later on, these flags cause adjustments to the pythonization - # process to accommodate logprobs. - - seq_groups = sampling_metadata.seq_groups - prompt_logprobs_are_requested_for_prefill = any([ - sg.sampling_params.prompt_logprobs is not None and sg.is_prompt - for sg in seq_groups - ]) - any_logprobs_are_requested = ( - prompt_logprobs_are_requested_for_prefill - or any([sg.sampling_params.logprobs is not None for sg in seq_groups])) - - if prompt_logprobs_are_requested_for_prefill: - # CPU GPU sync, after gathering *only* sampled tokens (since - # requesting prompt logprobs leads `sampled_token_ids` to - # include prompt token ids in addition to sampled token ids.) - sample_idx_tensor = torch.tensor( - [sdx for sg in seq_groups for sdx in sg.sample_indices]) - pinned_buffer = pinned_buffer.copy_( - sampled_token_ids[sample_idx_tensor, :], non_blocking=False) - else: - # CPU GPU sync - pinned_buffer = pinned_buffer.copy_(sampled_token_ids, - non_blocking=False) - - # this will not block as the tensors are already on CPU - samples_list = pinned_buffer.tolist() - - skip_sampler_cpu_output = ( - frozen_model_input.sampling_metadata.skip_sampler_cpu_output) - - # *Don't* skip logprobs pythonization *if*: - # * Any requests require logprobs to be returned in this - # iteration AND - # * These requests are being scheduled in a fashion which - # defers pythonization (i.e. multi-step scheduling.) - do_pythonize_logprobs = (skip_sampler_cpu_output - and any_logprobs_are_requested) - ( - prompt_logprobs, - sample_logprobs, - ) = (deferred_pythonize_logprobs(output, sampling_metadata, - logprobs_tensor) - if do_pythonize_logprobs else (None, None)) - - for sgdx, (seq_group, - sample_result) in enumerate(zip(seq_groups, samples_list)): - # Reminder: Please update docs/features/compatibility_matrix.md - # If the feature combo become valid - # (Check for Guided Decoding) - if seq_group.sampling_params.logits_processors: - assert len(seq_group.sampling_params.logits_processors) == 0, ( - "Logits Processors are not supported in multi-step decoding") - - if do_pythonize_logprobs: - assert prompt_logprobs is not None - assert sample_logprobs is not None - - ( - group_prompt_logprobs, - group_sample_logprobs, - ) = ( # Utilize deferred pythonization results - prompt_logprobs[sgdx], - sample_logprobs[sgdx], - ) - elif any_logprobs_are_requested: - ( - group_prompt_logprobs, - group_sample_logprobs, - ) = ( - # profile_run: use already-computed logprobs - output.outputs[sgdx].prompt_logprobs, - [sample.logprobs for sample in output.outputs[sgdx].samples]) - - seq_ids = seq_group.seq_ids - next_token_ids = sample_result - parent_ids = [0] - seq_outputs: List[SequenceOutput] - - if cache is not None: - completion_seq_group_output: CompletionSequenceGroupOutput = \ - cache.cached_completion_seq_group_output.get_object() - completion_seq_group_output.samples.clear() - seq_outputs = completion_seq_group_output.samples - else: - seq_outputs = [] - - for tdx, (parent_id, - next_token_id) in enumerate(zip(parent_ids, next_token_ids)): - if cache is not None: - seq_output: SequenceOutput = cache.cached_seq_output.get_object( - ) - seq_output.parent_seq_id = seq_ids[parent_id] - seq_output.output_token = next_token_id - - if any_logprobs_are_requested: - seq_output.logprobs = group_sample_logprobs[tdx] - else: - logprobs = next(iter(seq_output.logprobs.values())) - seq_output.logprobs.clear() - - logprobs.logprob = float('inf') - logprobs.rank = None - logprobs.decoded_token = None - - seq_output.logprobs[next_token_id] = logprobs - - seq_outputs.append(seq_output) - - else: - seq_outputs.append( - SequenceOutput(seq_ids[parent_id], next_token_id, - (group_sample_logprobs[tdx] - if any_logprobs_are_requested else { - next_token_id: - Logprob(logprob=float('inf'), - rank=None, - decoded_token=None) - }))) - if cache is not None: - completion_seq_group_output.prompt_logprobs = \ - group_prompt_logprobs if any_logprobs_are_requested else None - output.outputs.append(completion_seq_group_output) - else: - output.outputs.append( - CompletionSequenceGroupOutput( - seq_outputs, (group_prompt_logprobs - if any_logprobs_are_requested else None))) - - assert len(output.outputs) > 0 diff --git a/vllm/worker/multi_step_neuron_model_runner.py b/vllm/worker/multi_step_neuron_model_runner.py deleted file mode 100644 index 25f588077cb4..000000000000 --- a/vllm/worker/multi_step_neuron_model_runner.py +++ /dev/null @@ -1,84 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from importlib.util import find_spec -from typing import List, Optional - -import torch - -from vllm.config import VllmConfig -from vllm.model_executor.layers.sampler import SamplerOutput -from vllm.multimodal import MultiModalKwargs -from vllm.sequence import IntermediateTensors -from vllm.worker.neuron_model_runner import (ModelInputForNeuron, - NeuronModelRunner) - - -class MultiStepNeuronModelRunner(NeuronModelRunner): - """A model runner for multi step decoding using the transformers_neuronx - framework""" - - def __init__( - self, - vllm_config: VllmConfig, - ): - super().__init__(vllm_config) - self.speculation_config = self.speculative_config - from transformers_neuronx.config import GenerationConfig - self.speculation_config.draft_model_config.neuron_sampling_params = ( - GenerationConfig( - max_length=self.scheduler_config.max_model_len, - do_sample=True, - per_batch_line=True, - top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \ - * self.scheduler_config.max_num_seqs, - top_p=[1.0] * self.scheduler_config.max_num_seqs, - temperature=[1.0] * self.scheduler_config.max_num_seqs, - dynamic=True, - global_top_k=self._MAX_NEURON_SAMPLING_TOP_K - )) - - def load_model(self) -> None: - if find_spec("transformers_neuronx") is not None: - from vllm.model_executor.model_loader.neuron import ( - get_neuron_eagle_speculation_model, - get_neuron_speculation_model) - if self.speculation_config.speculative_token_tree is not None: - self.model = get_neuron_eagle_speculation_model( - self.model_config, - parallel_config=self.parallel_config, - scheduler_config=self.scheduler_config, - speculation_config=self.speculation_config) - else: - self.model = get_neuron_speculation_model( - self.model_config, - parallel_config=self.parallel_config, - scheduler_config=self.scheduler_config, - speculation_config=self.speculation_config) - else: - raise NotImplementedError( - "Supports only Transformer-NeuronX based models.") - - @torch.inference_mode() - def execute_model( - self, - model_input: ModelInputForNeuron, - kv_caches: Optional[List[torch.Tensor]] = None, - intermediate_tensors: Optional[IntermediateTensors] = None, - num_steps: int = 1, - ) -> Optional[List[SamplerOutput]]: - logits = self.model( - input_ids=model_input.input_tokens, - positions=model_input.input_positions, - input_block_ids=model_input.input_block_ids, - **MultiModalKwargs.as_kwargs( - model_input.multi_modal_kwargs or {}, - device=self.device, - ), - ) - - output = self.model.sample( - logits=logits, - sampling_metadata=model_input.sampling_metadata, - ) - return output diff --git a/vllm/worker/multi_step_neuronx_distributed_model_runner.py b/vllm/worker/multi_step_neuronx_distributed_model_runner.py deleted file mode 100644 index dd521dd67dad..000000000000 --- a/vllm/worker/multi_step_neuronx_distributed_model_runner.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import List, Optional - -import torch - -from vllm.config import VllmConfig -from vllm.model_executor.layers.sampler import SamplerOutput -from vllm.multimodal import MultiModalKwargs -from vllm.sequence import IntermediateTensors -from vllm.worker.neuronx_distributed_model_runner import ( - NeuronxDistributedModelRunner) - - -class MultiStepNeuronxDistributedModelRunner(NeuronxDistributedModelRunner): - """A model runner for multi-step decoding using the - neuronx-distributed-inference framework""" - - def __init__( - self, - vllm_config: VllmConfig, - ): - super().__init__(vllm_config) - - def load_model(self) -> None: - from vllm.model_executor.model_loader.neuronx_distributed import ( - get_neuron_speculation_model) - self.model = get_neuron_speculation_model( - self.model_config, - parallel_config=self.parallel_config, - scheduler_config=self.scheduler_config, - speculation_config=self.speculative_config) - - @torch.inference_mode() - def execute_model( - self, - model_input, - kv_caches: Optional[List[torch.Tensor]] = None, - intermediate_tensors: Optional[IntermediateTensors] = None, - num_steps: int = 1, - ) -> Optional[List[SamplerOutput]]: - sampling_params = torch.tensor([[ - seq_group.sampling_params.top_k, - seq_group.sampling_params.top_p, - seq_group.sampling_params.temperature, - ] for seq_group in model_input.sampling_metadata.seq_groups]) - - logits = self.model( - input_ids=model_input.input_tokens, - positions=model_input.input_positions, - input_block_ids=model_input.input_block_ids, - sampling_params=sampling_params, - **MultiModalKwargs.as_kwargs( - model_input.multi_modal_kwargs or {}, - device=self.device, - ), - ) - - output = self.model.sample( - logits=logits, - sampling_metadata=model_input.sampling_metadata, - ) - return output diff --git a/vllm/worker/multi_step_worker.py b/vllm/worker/multi_step_worker.py deleted file mode 100644 index ea16e14f9ecd..000000000000 --- a/vllm/worker/multi_step_worker.py +++ /dev/null @@ -1,197 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import dataclasses -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import torch - -from vllm.distributed import broadcast_tensor_dict, get_pp_group -from vllm.model_executor.layers.sampler import SamplerOutput -from vllm.sequence import ExecuteModelRequest -from vllm.worker.model_runner_base import BroadcastableModelInput -from vllm.worker.multi_step_model_runner import (MultiStepModelRunner, - StatefulModelInput) -from vllm.worker.worker import Worker, WorkerInput - - -@dataclass -class MultiStepState: - worker_input: WorkerInput - model_input: StatefulModelInput - - -class MultiStepWorker(Worker): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - base_model_runner = self.model_runner - # for multi-step model, wrap the model runner with MultiStepModelRunner - self.model_runner = MultiStepModelRunner( - base_model_runner, - vllm_config=base_model_runner.vllm_config, - kv_cache_dtype=self.cache_config.cache_dtype, - is_driver_worker=base_model_runner.is_driver_worker, - ) - - pipeline_parallel_size = self.parallel_config.pipeline_parallel_size - self.multi_step_states: List[ - Optional[MultiStepState]] = [None] * pipeline_parallel_size - self.temp_output = None - - def _get_driver_input_and_broadcast( - self, execute_model_req: ExecuteModelRequest - ) -> Tuple[BroadcastableModelInput, WorkerInput, Dict[str, torch.Tensor]]: - """ - Get the driver input and broadcast it to other workers. - """ - assert self.is_driver_worker - virtual_engine = execute_model_req.virtual_engine - is_first_multi_step = execute_model_req.is_first_multi_step - if is_first_multi_step: - # on first step we prepare the worker input and model input normally - worker_input: WorkerInput = self.prepare_worker_input( - execute_model_req=execute_model_req) - model_input: StatefulModelInput = ( - self.model_runner.prepare_model_input( - execute_model_req.seq_group_metadata_list, - execute_model_req.virtual_engine, - execute_model_req.finished_requests_ids)) - - if execute_model_req.async_callback: - model_input.frozen_model_input = dataclasses.replace( # type: ignore - model_input.frozen_model_input, - async_callback=execute_model_req.async_callback) - else: - # on subsequent steps we reuse the worker input and model input - multi_step_state = self.multi_step_states[virtual_engine] - worker_input = multi_step_state.worker_input - model_input = multi_step_state.model_input - frozen_model_input = model_input.frozen_model_input - assert frozen_model_input is not None - assert frozen_model_input.attn_metadata is not None - # clear the cached metadata so that it can be recomputed on - # the workers. - frozen_model_input.attn_metadata._cached_prefill_metadata = None - frozen_model_input.attn_metadata._cached_decode_metadata = None - - model_input.is_first_multi_step = is_first_multi_step - model_input.is_last_step = execute_model_req.is_last_step - - if not is_first_multi_step: - # we broadcast the last sampled token ids to all TP workers so they - # can update their model input metadata in-place. - self._prepare_last_sampled_token_ids_for_tp_workers( - execute_model_req=execute_model_req, model_input=model_input) - - if self.do_metadata_broadcast: - broadcast_data = worker_input.as_broadcastable_tensor_dict() - broadcast_data.update(model_input.as_broadcastable_tensor_dict()) - broadcast_tensor_dict(broadcast_data, src=0) - - # Retuning empty dict here to keep this compatible with - # `LocalOrDistributedWorkerBase._get_driver_input_and_broadcast` - return model_input, worker_input, {} - - def _prepare_last_sampled_token_ids_for_tp_workers( - self, - execute_model_req: ExecuteModelRequest, - model_input: StatefulModelInput, - ) -> None: - """ - Prepare the last sampled token ids for TP workers. If it's the last - PP rank, then the last sampled token ids are already in the model_input. - If it is NOT the last PP rank, then we need to get the last sampled - token that is cached in the execute_model_req. - """ - if get_pp_group().is_last_rank: - assert model_input.cached_outputs[ - -1].sampler_output.sampled_token_ids is None - assert model_input.cached_outputs[-1].sampled_token_ids is not None - model_input.last_sampled_token_ids = model_input.cached_outputs[ - -1].sampled_token_ids - # free sampled token ids from the previous step if it has been - # pythonized. Cannot free the last sampled token ids because - # we need it for GPU advance_step. - for output in model_input.cached_outputs[:-1]: - if output.pythonized: - output.sampled_token_ids = None - else: - # otherwise we need to get the cached sampled token ids from the - # execute_model_req - assert execute_model_req.last_sampled_token_ids is not None - model_input.last_sampled_token_ids = ( - execute_model_req.last_sampled_token_ids.cuda()) - model_input.add_sampler_output( - SamplerOutput(outputs=[], sampled_token_ids=None), - model_input.last_sampled_token_ids) - - # free sampled token ids from the previous step. - # TODO(will) we could reuse the sampled token ids tensor from - # the previous step instead. - for output in model_input.cached_outputs[:-1]: - output.sampled_token_ids = None - assert model_input.cached_outputs[-1].sampled_token_ids is not None - - def prepare_input( - self, - execute_model_req: Optional[ExecuteModelRequest] = None, - ) -> Optional[Tuple[StatefulModelInput, WorkerInput, Dict[str, - torch.Tensor]]]: - """ - Depending on the current state of the request and multi step worker, - this method may skip the normal _prepare_model_input and - _prepare_worker_input methods and instead used cached values. - """ - if self.is_driver_worker: - if execute_model_req is None: - if self.do_metadata_broadcast: - # This signals that there's no more requests to process for - # now. All workers are running infinite loop with - # broadcast_tensor_dict, and it stops the loop when the - # driver broadcasts an empty input. Send an empty input to - # notify all other workers to stop their execution loop. - broadcast_tensor_dict({}, src=0) - return None - - virtual_engine = execute_model_req.virtual_engine - (model_input, worker_input, - kwargs) = self._get_driver_input_and_broadcast(execute_model_req) - assert isinstance(model_input, StatefulModelInput) - if execute_model_req.is_first_multi_step: - # cache the worker input and model input for the next steps - self.multi_step_states[virtual_engine] = MultiStepState( - worker_input=worker_input, model_input=model_input) - # if TP workers - else: - broadcast_data = self._get_worker_input_from_broadcast() - # if the driver has sent an empty input, we should stop the worker - # loop - if broadcast_data is None: - return None - model_input, worker_input, kwargs = broadcast_data - assert isinstance(model_input, StatefulModelInput) - virtual_engine = worker_input.virtual_engine - if model_input.is_first_multi_step: - pass - # TODO(will) Can cache the worker input and model input for the - # next steps. See below for details - else: - # TODO(will) possible to also cache and reuse the cached worker - # input and model input. The idea is essentially the delta - # optimization for model_inputs. Where the TP workers can cache - # the model input states and we only broadcast the delta need - # for the next step (sampled_token_ids from the previous step) - - assert isinstance(model_input, StatefulModelInput) - # we need to update the last sampled token ids in the model - # input for the workers so that they can run inplace - # advance_step - model_input.add_sampler_output( - SamplerOutput(outputs=[], sampled_token_ids=None), - model_input.last_sampled_token_ids) - - assert model_input is not None - assert worker_input is not None - return model_input, worker_input, kwargs diff --git a/vllm/worker/neuron_worker.py b/vllm/worker/neuron_worker.py index 4e1408300fb8..3e4512a63908 100644 --- a/vllm/worker/neuron_worker.py +++ b/vllm/worker/neuron_worker.py @@ -64,25 +64,21 @@ def get_tnx_model_runner(self, vllm_config): assert (self.lora_config is None), ("LoRA is not supported for TransformersNeuronX " "framework.") - from vllm.worker.multi_step_neuron_model_runner import ( - MultiStepNeuronModelRunner) if self.speculative_config is not None: - return MultiStepNeuronModelRunner(vllm_config=vllm_config) - else: - return NeuronModelRunner(vllm_config=vllm_config) + raise NotImplementedError( + "Speculative decoding is not supported for TransformersNeuronX" + ) + return NeuronModelRunner(vllm_config=vllm_config) def get_neuronx_distributed_model_runner(self, vllm_config): - from vllm.worker.multi_step_neuronx_distributed_model_runner import ( - MultiStepNeuronxDistributedModelRunner) from vllm.worker.neuronx_distributed_model_runner import ( NeuronxDistributedModelRunner) if self.speculative_config is not None: - assert (self.lora_config - is None), "LoRA is not supported for Speculative Decoding" - return MultiStepNeuronxDistributedModelRunner( - vllm_config=vllm_config) - else: - return NeuronxDistributedModelRunner(vllm_config=vllm_config) + assert (self.lora_config is None), ( + "LoRA is not supported for Speculative Decoding") + raise NotImplementedError( + "Speculative decoding is not supported for NeuronxDistributed") + return NeuronxDistributedModelRunner(vllm_config=vllm_config) def init_device(self) -> None: self.init_distributed_environment() From 1187e5062a2653112b20d56edf3729520c99673f Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 12 Aug 2025 20:21:18 -0700 Subject: [PATCH 018/233] [Misc] Remove tests/multi_step/__init__.py (#22778) Signed-off-by: Woosuk Kwon --- tests/multi_step/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tests/multi_step/__init__.py diff --git a/tests/multi_step/__init__.py b/tests/multi_step/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 From 753f655d6ed0facf3cbda7b703d021b3dce30322 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 12 Aug 2025 20:38:18 -0700 Subject: [PATCH 019/233] [V0 Deprecation] Remove args for multi-step scheduling (#22779) Signed-off-by: Woosuk Kwon --- tests/utils_/test_utils.py | 1 - vllm/config/scheduler.py | 27 +-------------------------- 2 files changed, 1 insertion(+), 27 deletions(-) diff --git a/tests/utils_/test_utils.py b/tests/utils_/test_utils.py index 8be1e103dc65..084d82dee11b 100644 --- a/tests/utils_/test_utils.py +++ b/tests/utils_/test_utils.py @@ -161,7 +161,6 @@ def parser_with_config(): parser.add_argument('--port', type=int) parser.add_argument('--tensor-parallel-size', type=int) parser.add_argument('--trust-remote-code', action='store_true') - parser.add_argument('--multi-step-stream-outputs', action=StoreBoolean) return parser diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index db669600a0cc..93002012799a 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -115,12 +115,6 @@ class SchedulerConfig: (e.g., beam search), recomputation is not currently supported. In such a case, we use swapping instead.""" - num_scheduler_steps: int = 1 - """Maximum number of forward steps per scheduler call.""" - - multi_step_stream_outputs: bool = True - """If False, then multi-step will stream outputs at the end of all steps""" - send_delta_data: bool = False """Private API. If used, scheduler sends delta data to workers instead of an entire data. It should be enabled only @@ -193,16 +187,7 @@ def __post_init__(self) -> None: if self.max_num_batched_tokens is None: if self.enable_chunked_prefill: - if self.num_scheduler_steps > 1: - # Multi-step Chunked-Prefill doesn't allow prompt-chunking - # for now. Have max_num_batched_tokens set to max_model_len - # so we don't reject sequences on account of a short - # max_num_batched_tokens. - self.max_num_batched_tokens = max( - self.max_model_len, DEFAULT_MAX_NUM_BATCHED_TOKENS) - else: - self.max_num_batched_tokens = ( - DEFAULT_MAX_NUM_BATCHED_TOKENS) + self.max_num_batched_tokens = DEFAULT_MAX_NUM_BATCHED_TOKENS else: # If max_model_len is too short, use # DEFAULT_MAX_NUM_BATCHED_TOKENS as the default value @@ -293,12 +278,6 @@ def _verify_args(self) -> Self: f"({self.num_lookahead_slots}) must be greater than or " "equal to 0.") - if self.num_scheduler_steps < 1: - raise ValueError( - "num_scheduler_steps " - f"({self.num_scheduler_steps}) must be greater than or " - "equal to 1.") - if self.max_num_partial_prefills < 1: raise ValueError( f"max_num_partial_prefills ({self.max_num_partial_prefills}) " @@ -323,7 +302,3 @@ def _verify_args(self) -> Self: f"max_num_partial_prefills ({self.max_num_partial_prefills}).") return self - - @property - def is_multi_step(self) -> bool: - return self.num_scheduler_steps > 1 From 19891dc6771df282a93ad528765d48ce88880a40 Mon Sep 17 00:00:00 2001 From: "Po-Han Huang (NVIDIA)" <53919306+nvpohanh@users.noreply.github.com> Date: Wed, 13 Aug 2025 12:21:50 +0800 Subject: [PATCH 020/233] Fix cuda illegal mem access with Llama4 TP8 + rms_norm custom op (#22701) Signed-off-by: Po-Han Huang --- vllm/model_executor/models/llama4.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 1f8b9d074479..308cb3e85e27 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -224,10 +224,14 @@ def forward( if self.rotary_emb is not None: q, k = self.rotary_emb(positions, q, k) + if self.qk_norm is not None: - q = q.reshape(-1, self.num_heads, self.head_dim) + # Normalization is applied on the head_dim dimension. The rest of + # the dimensions are collapsed into a single dimension to support + # custom rms_norm cuda kernel. + q = q.reshape(-1, self.head_dim) q = self.qk_norm(q.float()).reshape(-1, self.q_size).to(q.dtype) - k = k.reshape(-1, self.num_kv_heads, self.head_dim) + k = k.reshape(-1, self.head_dim) k = self.qk_norm(k.float()).reshape(-1, self.kv_size).to(k.dtype) # We are applying temperature tuning (https://arxiv.org/abs/2501.19399) From 61419e90b4a46714874971af1e22a9552d7ffb6d Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 13 Aug 2025 00:22:05 -0400 Subject: [PATCH 021/233] [Bugfix] Fix default enable for CUTLASS MLA on SM100 (#22738) Signed-off-by: mgoin --- vllm/platforms/cuda.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 70959131573f..63f6b373c322 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -152,6 +152,9 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: if cls.is_device_capability(100): # Blackwell => Force CutlassMLA. use_cutlass_mla = True + # TODO: This does not work, because the + # global_force_attn_backend_context_manager is not set. + # See vllm/attention/selector.py:_cached_get_attn_backend envs.VLLM_ATTENTION_BACKEND = "CUTLASS_MLA" else: # Not Blackwell @@ -217,7 +220,9 @@ def get_attn_backend_cls(cls, selected_backend, head_size, dtype, if use_mla: # TODO(lucas): refactor to be more concise # we should probably consider factoring out V1 here - if selected_backend == _Backend.CUTLASS_MLA: + if selected_backend == _Backend.CUTLASS_MLA or ( + cls.is_device_capability(100) and selected_backend is None + and block_size == 128): if use_v1: logger.info_once("Using Cutlass MLA backend on V1 engine.") return ("vllm.v1.attention.backends.mla." From 61b6648d43949acc1aa0498ff6b285ef20875119 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 13 Aug 2025 00:22:16 -0400 Subject: [PATCH 022/233] Force TRTLLM attention for gpt-oss on SM100 (#22678) Signed-off-by: mgoin --- vllm/model_executor/models/gpt_oss.py | 5 +---- vllm/utils/flashinfer.py | 8 ++++++++ vllm/v1/attention/backends/flashinfer.py | 11 +++++++---- vllm/v1/attention/backends/utils.py | 5 ++++- 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 6a65bbbe2e0d..7c7712dbe106 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -8,7 +8,6 @@ from torch import nn from transformers import GptOssConfig -from vllm import envs from vllm.attention import Attention, AttentionType from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -70,11 +69,9 @@ def __init__( tp_size = get_tensor_model_parallel_world_size() - attention_sink_dtype = (torch.float32 if envs.VLLM_USE_TRTLLM_ATTENTION - else torch.bfloat16) self.sinks = torch.nn.Parameter( torch.empty(config.num_attention_heads // tp_size, - dtype=attention_sink_dtype, + dtype=torch.bfloat16, requires_grad=False)) self.norm = RMSNorm(config.hidden_size, eps=1e-5) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 5998d4c3127f..6b23ed426806 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -154,6 +154,7 @@ def use_trtllm_attention( num_qo_heads: Optional[int], num_kv_heads: Optional[int], attn_head_size: Optional[int], + has_sinks: bool = False, ) -> bool: # Requires SM100 and NVIDIA artifactory to be accessible to download cubins if not (current_platform.is_device_capability(100) @@ -165,6 +166,13 @@ def use_trtllm_attention( or num_qo_heads % num_kv_heads != 0): return False + # If sinks are being used, we must use TRTLLM attention as it's + # the only backend that supports them + if has_sinks: + logger.info_once( + "Using TRTLLM attention (required for attention sinks).") + return True + env_value = envs.VLLM_USE_TRTLLM_ATTENTION if env_value is not None: logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index c85d8bce31f5..12e5542d691c 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -523,14 +523,17 @@ def build(self, num_kv_heads = self.kv_cache_spec.num_kv_heads head_dim = self.kv_cache_spec.head_size + # Check if any layer uses sinks (requires TRTLLM attention) + has_sinks = self.global_hyperparameters.has_sinks + # currently prefill trtllm attention does not support fp8 kv cache prefill_use_trtllm = not cache_dtype.startswith("fp8") \ and use_trtllm_attention( num_prefill_tokens, max_seq_len, cache_dtype, - num_qo_heads, num_kv_heads, head_dim) + num_qo_heads, num_kv_heads, head_dim, has_sinks) decode_use_trtllm = use_trtllm_attention( num_decode_tokens, max_seq_len, cache_dtype, - num_qo_heads, num_kv_heads, head_dim) + num_qo_heads, num_kv_heads, head_dim, has_sinks) attn_metadata = FlashInferMetadata( num_actual_tokens=num_actual_tokens, @@ -642,9 +645,9 @@ def __init__( f"heads in the layer. Expected {num_heads}, but got " f"{sinks.shape[0]}." ) + # Cast sinks to float32 if needed (FlashInfer requirement) if sinks.dtype != torch.float32: - raise ValueError("Sinks must be of type float32, but got " - f"{sinks.dtype}.") + sinks = sinks.to(torch.float32) self.sinks = sinks def forward( diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index e23dd8bc5bbb..91eb84245ac0 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -285,6 +285,7 @@ class PerLayerParameters: window_left: int logits_soft_cap: Optional[float] sm_scale: float + has_sinks: bool = False def get_per_layer_parameters( @@ -307,9 +308,11 @@ def get_per_layer_parameters( window_left = window_size[0] if window_size is not None else -1 logits_soft_cap = getattr(impl, "logits_soft_cap", None) sm_scale = impl.scale + has_sinks = getattr(impl, "sinks", None) is not None per_layer_params[key] = PerLayerParameters(window_left, - logits_soft_cap, sm_scale) + logits_soft_cap, sm_scale, + has_sinks) return per_layer_params From f776e11f500282b349e166b939891fd3d7d1fa29 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 13 Aug 2025 00:26:38 -0400 Subject: [PATCH 023/233] Remove unneeded ROCm platform import when using CUDA (#22765) Signed-off-by: mgoin --- vllm/attention/backends/rocm_flash_attn.py | 2 +- vllm/attention/ops/chunked_prefill_paged_decode.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/attention/backends/rocm_flash_attn.py b/vllm/attention/backends/rocm_flash_attn.py index 1ee1dea729d9..da3d9ff32830 100644 --- a/vllm/attention/backends/rocm_flash_attn.py +++ b/vllm/attention/backends/rocm_flash_attn.py @@ -22,7 +22,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape) from vllm.platforms import current_platform -from vllm.platforms.rocm import use_rocm_custom_paged_attention if TYPE_CHECKING: from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata @@ -886,6 +885,7 @@ def forward( num_seqs, num_heads, head_size = decode_query.shape block_size = value_cache.shape[3] gqa_ratio = num_heads // self.num_kv_heads + from vllm.platforms.rocm import use_rocm_custom_paged_attention use_custom = use_rocm_custom_paged_attention( decode_query.dtype, head_size, block_size, gqa_ratio, decode_meta.max_decode_seq_len, self.sliding_window, diff --git a/vllm/attention/ops/chunked_prefill_paged_decode.py b/vllm/attention/ops/chunked_prefill_paged_decode.py index dc10d7eca9c2..e5b90a8b2755 100644 --- a/vllm/attention/ops/chunked_prefill_paged_decode.py +++ b/vllm/attention/ops/chunked_prefill_paged_decode.py @@ -11,7 +11,6 @@ from vllm import _custom_ops as ops from vllm.platforms import current_platform -from vllm.platforms.rocm import use_rocm_custom_paged_attention from vllm.triton_utils import tl, triton from .prefix_prefill import context_attention_fwd @@ -296,6 +295,7 @@ def chunked_prefill_paged_decode( num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv), 16) + from vllm.platforms.rocm import use_rocm_custom_paged_attention use_custom = use_rocm_custom_paged_attention( query.dtype, head_size, From 50bd03332d7791319bbea6b8ecf4db8ecf70c14e Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 13 Aug 2025 00:31:47 -0400 Subject: [PATCH 024/233] [Bug] Fix Unexpected Keyword Argument 'w1_bias' (#22757) Signed-off-by: yewentao256 --- vllm/model_executor/layers/fused_moe/layer.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index fb38fb91ead6..8ef0a805d86c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -475,12 +475,11 @@ def forward_cuda( activation=activation, apply_router_weight_on_input=apply_router_weight_on_input) else: - return self.fused_experts( + # add w1_bias/w2_bias to kwargs if they exist + kwargs = dict( hidden_states=x, w1=layer.w13_weight, w2=layer.w2_weight, - w1_bias=layer.w13_bias if self.has_bias else None, - w2_bias=layer.w2_bias if self.has_bias else None, topk_weights=topk_weights, topk_ids=topk_ids, inplace=True, @@ -489,6 +488,17 @@ def forward_cuda( global_num_experts=global_num_experts, expert_map=expert_map, ) + if isinstance(self.fused_experts, + FusedMoEModularKernel) and self.has_bias: + raise ValueError( + "FusedMoEModularKernel does not support bias.") + if self.has_bias: + kwargs.update({ + "w1_bias": getattr(layer, "w13_bias", None), + "w2_bias": getattr(layer, "w2_bias", None), + }) + + return self.fused_experts(**kwargs) def forward_cpu( self, From cbb55083439048ab4c78039fbaad05ac90e3d2da Mon Sep 17 00:00:00 2001 From: shixianc <49539556+shixianc@users.noreply.github.com> Date: Tue, 12 Aug 2025 21:34:47 -0700 Subject: [PATCH 025/233] [Perf] Support topk softmax fused kernel for broader num_experts (#22211) Signed-off-by: Shixian Cui Co-authored-by: Shixian Cui --- csrc/moe/topk_softmax_kernels.cu | 77 +++++++++++++++++++------------- tests/kernels/moe/test_moe.py | 2 +- 2 files changed, 46 insertions(+), 33 deletions(-) diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu index 7a7865b901de..946c137db636 100644 --- a/csrc/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -188,7 +188,9 @@ __launch_bounds__(TPB) __global__ void moeTopK( It fuses the softmax, max and argmax into a single kernel. Limitations: - 1) This implementation is intended for when the number of experts is a small power of 2. + 1) This implementation is optimized for when the number of experts is a small power of 2. + Additionally it also supports when number of experts is multiple of 64 which is still + faster than the computing softmax and topK separately (only tested on CUDA yet). 2) This implementation assumes k is small, but will work for any k. */ @@ -198,8 +200,6 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ int* source_rows, const int k, const int start_expert, const int end_expert) { // We begin by enforcing compile time assertions and setting up compile time constants. - static_assert(VPT == (VPT & -VPT), "VPT must be power of 2"); - static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), "NUM_EXPERTS must be power of 2"); static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2"); static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); @@ -407,12 +407,10 @@ struct TopkConstants }; } // namespace detail -template +template void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, IndType* indices, int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream) { - static constexpr std::size_t MAX_BYTES_PER_LDG = 16; - static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS); using Constants = detail::TopkConstants; static constexpr int VPT = Constants::VPT; @@ -425,21 +423,12 @@ void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, f input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert); } -#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB) \ - switch (warpSize) { \ - case 32: \ - topkGatingSoftmaxLauncherHelper( \ - gating_output, nullptr, topk_weights, topk_indices, \ - token_expert_indices, num_tokens, topk, 0, num_experts, stream); \ - break; \ - case 64: \ - topkGatingSoftmaxLauncherHelper( \ - gating_output, nullptr, topk_weights, topk_indices, \ - token_expert_indices, num_tokens, topk, 0, num_experts, stream); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported warp size: ", warpSize); \ - } +#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ + static_assert(WARP_SIZE == 32 || WARP_SIZE == 64, \ + "Unsupported warp size. Only 32 and 64 are supported."); \ + topkGatingSoftmaxLauncherHelper( \ + gating_output, nullptr, topk_weights, topk_indices, \ + token_expert_indices, num_tokens, topk, 0, num_experts, stream); template void topkGatingSoftmaxKernelLauncher( @@ -453,38 +442,62 @@ void topkGatingSoftmaxKernelLauncher( const int topk, cudaStream_t stream) { static constexpr int WARPS_PER_TB = 4; - auto warpSize = WARP_SIZE; + static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; + static constexpr int BYTES_PER_LDG_MULTIPLE_64 = 8; switch (num_experts) { case 1: - LAUNCH_SOFTMAX(1, WARPS_PER_TB); + LAUNCH_SOFTMAX(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 2: - LAUNCH_SOFTMAX(2, WARPS_PER_TB); + LAUNCH_SOFTMAX(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 4: - LAUNCH_SOFTMAX(4, WARPS_PER_TB); + LAUNCH_SOFTMAX(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 8: - LAUNCH_SOFTMAX(8, WARPS_PER_TB); + LAUNCH_SOFTMAX(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 16: - LAUNCH_SOFTMAX(16, WARPS_PER_TB); + LAUNCH_SOFTMAX(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 32: - LAUNCH_SOFTMAX(32, WARPS_PER_TB); + LAUNCH_SOFTMAX(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 64: - LAUNCH_SOFTMAX(64, WARPS_PER_TB); + LAUNCH_SOFTMAX(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 128: - LAUNCH_SOFTMAX(128, WARPS_PER_TB); + LAUNCH_SOFTMAX(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); break; case 256: - LAUNCH_SOFTMAX(256, WARPS_PER_TB); + LAUNCH_SOFTMAX(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 512: + LAUNCH_SOFTMAX(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + // (CUDA only) support multiples of 64 when num_experts is not power of 2. + // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of num_experts, + // alternatively we can test 4 bytes loading and enable it in future. +#ifndef USE_ROCM + case 192: + LAUNCH_SOFTMAX(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); break; + case 320: + LAUNCH_SOFTMAX(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 384: + LAUNCH_SOFTMAX(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 448: + LAUNCH_SOFTMAX(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 576: + LAUNCH_SOFTMAX(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; +#endif default: { TORCH_CHECK(softmax_workspace != nullptr, - "softmax_workspace must be provided for num_experts that are not a power of 2."); + "softmax_workspace must be provided for num_experts that are not a power of 2 or multiple of 64."); static constexpr int TPB = 256; moeSoftmax<<>>( gating_output, nullptr, softmax_workspace, num_experts); diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 0f1c78704642..49c097718e30 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -36,7 +36,7 @@ from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types -NUM_EXPERTS = [8, 64] +NUM_EXPERTS = [8, 64, 192] EP_SIZE = [1, 4] TOP_KS = [2, 6] From dd5c24607a6edffe634653c5df4607947e25c0dc Mon Sep 17 00:00:00 2001 From: Chen Zhang Date: Tue, 12 Aug 2025 21:37:26 -0700 Subject: [PATCH 026/233] [gpt-oss] upgrade gpt-oss to v0.0.3 and add version check (#22768) Signed-off-by: Chen Zhang --- vllm/entrypoints/tool.py | 51 ++++++++++++++++++++++++++-------------- 1 file changed, 34 insertions(+), 17 deletions(-) diff --git a/vllm/entrypoints/tool.py b/vllm/entrypoints/tool.py index 723cff91d44c..758789a5e059 100644 --- a/vllm/entrypoints/tool.py +++ b/vllm/entrypoints/tool.py @@ -2,9 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Optional - -from openai_harmony import Message +from typing import TYPE_CHECKING, Any from vllm.logger import init_logger @@ -15,6 +13,30 @@ logger = init_logger(__name__) +def validate_gpt_oss_install(): + """ + Check if the gpt-oss is installed and its version is at least 0.0.3. + If not, raise an ImportError. + """ + from importlib.metadata import PackageNotFoundError, version + + from packaging.version import InvalidVersion, Version + + try: + pkg_version_str = version("gpt_oss") # e.g., "0.0.5" + pkg_version = Version(pkg_version_str) + except PackageNotFoundError: + raise ImportError("Package 'gpt_oss' is not installed.") from None + except InvalidVersion as e: + raise ImportError( + f"Invalid version string for 'gpt_oss': {e}") from None + + if pkg_version < Version("0.0.3"): + raise ImportError( + f"gpt_oss >= 0.0.3 is required, but {pkg_version} is installed." + ) from None + + class Tool(ABC): @abstractmethod @@ -33,12 +55,14 @@ def __init__(self): return try: + validate_gpt_oss_install() from gpt_oss.tools.simple_browser import SimpleBrowserTool from gpt_oss.tools.simple_browser.backend import ExaBackend - except ImportError: + except ImportError as e: self.enabled = False logger.warning_once( - "gpt_oss is not installed, browsing is disabled") + "gpt_oss is not installed properly (%s), browsing is disabled", + e) return browser_backend = ExaBackend(source="web", api_key=exa_api_key) @@ -65,23 +89,16 @@ def __init__(self): self.enabled = True try: + validate_gpt_oss_install() from gpt_oss.tools.python_docker.docker_tool import PythonTool - except ImportError: + except ImportError as e: self.enabled = False logger.warning_once( - "gpt_oss is not installed, code interpreter is disabled") + "gpt_oss is not installed properly (%s), code interpreter is " + "disabled", e) return - # NOTE (Chen): as of gpt-oss 0.0.2, there is a bug in _make_response - # and we do the following monkey patch to fix it. - class PatchedGptOssPythonTool(PythonTool): - - def _make_response(self, - output: str, - channel: Optional[str] = None) -> Message: - return super()._make_response(output) - - self.python_tool = PatchedGptOssPythonTool() + self.python_tool = PythonTool() logger.info_once("Code interpreter tool initialized") async def get_result(self, context: "ConversationContext") -> Any: From f362240ae578344e2dac1ad5cbaf0254991e8ec7 Mon Sep 17 00:00:00 2001 From: zzh142857 Date: Wed, 13 Aug 2025 03:09:13 -0400 Subject: [PATCH 027/233] [Model] Add option to run Step3VisionEncoder in DP (#22697) Signed-off-by: zzh142857 --- vllm/model_executor/models/step3_vl.py | 132 +++++++++++++++++-------- 1 file changed, 91 insertions(+), 41 deletions(-) diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 41dba312cb42..f1f38c01b784 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.linear import (ColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler @@ -33,6 +34,7 @@ BaseProcessingInfo, PromptReplacement, PromptUpdate, PromptUpdateDetails) from vllm.multimodal.profiling import BaseDummyInputsBuilder +from vllm.multimodal.utils import run_dp_sharded_vision_model from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs import Step3VisionEncoderConfig from vllm.transformers_utils.tokenizer import AnyTokenizer @@ -650,7 +652,8 @@ class Step3VisionAttention(nn.Module): def __init__(self, config, quant_config: Optional[QuantizationConfig] = None, - prefix: str = ""): + prefix: str = "", + use_data_parallel: bool = False): super().__init__() self.config = config self.embed_dim = config.hidden_size @@ -659,20 +662,42 @@ def __init__(self, self.scale = self.head_dim**-0.5 - tp_size = get_tensor_model_parallel_world_size() + tp_size = (1 if use_data_parallel else + get_tensor_model_parallel_world_size()) assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size - self.qkv_proj = QKVParallelLinear(self.embed_dim, - self.head_dim, - self.total_num_heads, - bias=True, - quant_config=quant_config, - prefix=prefix) - self.out_proj = RowParallelLinear(self.embed_dim, - self.embed_dim, - bias=True, - quant_config=quant_config, - prefix=prefix) + + self.q_size = self.num_heads * self.head_dim + + if use_data_parallel: + self.qkv_proj = ReplicatedLinear( + self.embed_dim, + 3 * self.q_size, + bias=True, + quant_config=quant_config, + prefix=prefix, + ) + self.out_proj = ReplicatedLinear( + self.total_num_heads * self.head_dim, + self.embed_dim, + bias=True, + quant_config=quant_config, + prefix=prefix, + ) + else: + self.qkv_proj = QKVParallelLinear( + self.embed_dim, + self.head_dim, + self.total_num_heads, + bias=True, + quant_config=quant_config, + prefix=prefix, + ) + self.out_proj = RowParallelLinear(self.embed_dim, + self.embed_dim, + bias=True, + quant_config=quant_config, + prefix=prefix) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, @@ -712,20 +737,25 @@ class Step3VisionMLP(nn.Module): def __init__(self, config, quant_config: Optional[QuantizationConfig] = None, - prefix: str = ""): + prefix: str = "", + use_data_parallel: bool = False): super().__init__() self.config = config self.activation_fn = get_act_fn(config.hidden_act) - self.fc1 = ColumnParallelLinear(config.hidden_size, - config.intermediate_size, - bias=True, - quant_config=quant_config, - prefix=prefix) - self.fc2 = RowParallelLinear(config.intermediate_size, - config.hidden_size, - bias=True, - quant_config=quant_config, - prefix=prefix) + cls_fc1 = (ReplicatedLinear + if use_data_parallel else ColumnParallelLinear) + self.fc1 = cls_fc1(config.hidden_size, + config.intermediate_size, + bias=True, + quant_config=quant_config, + prefix=prefix) + cls_fc2 = (ReplicatedLinear + if use_data_parallel else RowParallelLinear) + self.fc2 = cls_fc2(config.intermediate_size, + config.hidden_size, + bias=True, + quant_config=quant_config, + prefix=prefix) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.fc1(hidden_states) @@ -739,15 +769,22 @@ class Step3VisionEncoderLayer(nn.Module): def __init__(self, config: Step3VisionEncoderConfig, quant_config: Optional[QuantizationConfig] = None, - prefix: str = ""): + prefix: str = "", + use_data_parallel: bool = False): super().__init__() + self.use_data_parallel = use_data_parallel self.embed_dim = config.hidden_size - self.self_attn = Step3VisionAttention(config, - quant_config, - prefix=f"{prefix}.self_attn") + self.self_attn = Step3VisionAttention( + config, + quant_config, + prefix=f"{prefix}.self_attn", + use_data_parallel=self.use_data_parallel) self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - self.mlp = Step3VisionMLP(config, quant_config, prefix=f"{prefix}.mlp") + self.mlp = Step3VisionMLP(config, + quant_config, + prefix=f"{prefix}.mlp", + use_data_parallel=self.use_data_parallel) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) @@ -767,13 +804,16 @@ class Step3VisionEncoder(nn.Module): def __init__(self, config: Step3VisionEncoderConfig, quant_config: Optional[QuantizationConfig] = None, - prefix: str = ""): + prefix: str = "", + use_data_parallel: bool = False): super().__init__() self.config = config + self.use_data_parallel = use_data_parallel self.layers = nn.ModuleList([ Step3VisionEncoderLayer(config, quant_config, - prefix=f"{prefix}.layers.{i}") + prefix=f"{prefix}.layers.{i}", + use_data_parallel=self.use_data_parallel) for i in range(config.num_hidden_layers) ]) @@ -792,21 +832,29 @@ class Step3VisionTransformer(nn.Module): def __init__(self, config: Step3VisionEncoderConfig, quant_config: Optional[QuantizationConfig] = None, - prefix: str = ""): + prefix: str = "", + use_data_parallel: bool = False): super().__init__() self.config = config + self.use_data_parallel = use_data_parallel self.image_size = config.image_size self.embeddings = Step3VisionEmbeddings(config) - self.transformer = Step3VisionEncoder(config, - quant_config, - prefix=f"{prefix}.transformer") + self.transformer = Step3VisionEncoder( + config, + quant_config, + prefix=f"{prefix}.transformer", + use_data_parallel=self.use_data_parallel) def forward( self, pixel_values: torch.Tensor, ): hidden_states = self.embeddings(pixel_values) - hidden_states = self.transformer(inputs_embeds=hidden_states) + if self.use_data_parallel: + hidden_states = run_dp_sharded_vision_model( + hidden_states, self.transformer) + else: + hidden_states = self.transformer(inputs_embeds=hidden_states) return hidden_states @@ -836,13 +884,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.config = config self.multimodal_config = multimodal_config + self.use_data_parallel = (vllm_config.parallel_config. + enable_multimodal_encoder_data_parallel) if multimodal_config.get_limit_per_prompt("image"): - self.vision_model = Step3VisionTransformer(config.vision_config, - None, - prefix=maybe_prefix( - prefix, - "vision_model")) + self.vision_model = Step3VisionTransformer( + config.vision_config, + None, + prefix=maybe_prefix(prefix, "vision_model"), + use_data_parallel=self.use_data_parallel) self.vit_downsampler = nn.Conv2d( config.vision_config.hidden_size, config.vision_config.output_hidden_size, From ee22b087ff1ff6a71d21acb061963369f8e0088e Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang <2448370773@qq.com> Date: Wed, 13 Aug 2025 16:23:33 +0800 Subject: [PATCH 028/233] [Model] Add missing prefix to glm4_1v (#22716) Signed-off-by: zRzRzRzRzRzRzR <2448370773@qq.com> --- vllm/model_executor/models/glm4_1v.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 2a89c03bfe7e..88c53c836327 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -453,25 +453,30 @@ def __init__( context_dim: int, quant_config: Optional[QuantizationConfig] = None, bias: bool = False, + prefix: str = "", ) -> None: super().__init__() self.hidden_size = d_model self.proj = ColumnParallelLinear(self.hidden_size, self.hidden_size, bias=bias, - gather_output=True) + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.proj") self.post_projection_norm = nn.LayerNorm(self.hidden_size) self.gate_up_proj = MergedColumnParallelLinear( input_size=self.hidden_size, output_sizes=[context_dim] * 2, bias=bias, quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( context_dim, self.hidden_size, bias=bias, quant_config=quant_config, + prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() self.extra_activation_func = nn.GELU() @@ -661,6 +666,7 @@ def __init__( context_dim=vision_config.intermediate_size, quant_config=quant_config, bias=False, + prefix=f"{prefix}.merger", ) self.embeddings = Glm4vVisionEmbeddings(vision_config) From 66c5b95c15d25afa4b4138f1576816b170ecd327 Mon Sep 17 00:00:00 2001 From: Duc-Viet Hoang Date: Wed, 13 Aug 2025 17:11:36 +0700 Subject: [PATCH 029/233] [Bugfix] Fix Nemotron VL image processing (#22739) Co-authored-by: ducviet00-h2 --- .../multimodal/processing/test_nemotron_vl.py | 8 +- vllm/model_executor/models/nemotron_vl.py | 186 ++++++++++++++++++ 2 files changed, 190 insertions(+), 4 deletions(-) diff --git a/tests/models/multimodal/processing/test_nemotron_vl.py b/tests/models/multimodal/processing/test_nemotron_vl.py index 3ce88bc427f5..6fbbab0d2612 100644 --- a/tests/models/multimodal/processing/test_nemotron_vl.py +++ b/tests/models/multimodal/processing/test_nemotron_vl.py @@ -23,15 +23,15 @@ def _get_expected_num_patches( min_num: int, max_num: int, ): - from vllm.model_executor.models.internvl import ( - calculate_internvl_targets, get_internvl_target_ratios) + from vllm.model_executor.models.nemotron_vl import ( + calculate_nemotron_vl_targets, get_nemotron_vl_target_ratios) width, height = image.size - blocks, _, _ = calculate_internvl_targets( + blocks, _, _ = calculate_nemotron_vl_targets( orig_width=width, orig_height=height, - target_ratios=get_internvl_target_ratios( + target_ratios=get_nemotron_vl_target_ratios( min_num, max_num, ), diff --git a/vllm/model_executor/models/nemotron_vl.py b/vllm/model_executor/models/nemotron_vl.py index b90cb9b39a60..82bcd064624f 100644 --- a/vllm/model_executor/models/nemotron_vl.py +++ b/vllm/model_executor/models/nemotron_vl.py @@ -13,6 +13,7 @@ import torch import torch.nn as nn +import torchvision.transforms as T from PIL import Image from transformers import AutoModel, PretrainedConfig from transformers.image_processing_utils_fast import BaseImageProcessorFast @@ -27,6 +28,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.image import convert_image_mode from vllm.multimodal.inputs import NestedTensors from vllm.multimodal.processing import PromptUpdateDetails from vllm.sequence import IntermediateTensors @@ -44,6 +46,146 @@ IMG_CONTEXT = '' +def build_transform(input_size: int): + return T.Compose([ + T.Lambda(lambda img: convert_image_mode(img, 'RGB')), + T.Resize((input_size, input_size), + interpolation=T.InterpolationMode.BICUBIC), + T.ToTensor(), + ]) + + +# adapted from https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1 +def find_closest_aspect_ratio( + aspect_ratio: float, + target_ratios: list[tuple[int, int]], + *, + width: int, + height: int, + image_size: int, +) -> tuple[int, int]: + best_factor = float('-inf') + best_ratio = (1, 1) + area = width * height + + for rw, rh in target_ratios: + target_aspect_ratio = rw / rh + size_factor = min((rw * rh * image_size * image_size) / area, 0.6) + ratio_closeness = min(target_aspect_ratio / aspect_ratio, + aspect_ratio / target_aspect_ratio) + factor = size_factor * ratio_closeness + + if factor > best_factor: + best_factor = factor + best_ratio = (rw, rh) + + return best_ratio + + +def calculate_nemotron_vl_targets( + *, + orig_width: int, + orig_height: int, + target_ratios: list[tuple[int, int]], + image_size: int, + use_thumbnail: bool, +) -> tuple[int, int, int]: + aspect_ratio = orig_width / orig_height + + # find the closest aspect ratio to the target + target_aspect_ratio = find_closest_aspect_ratio( + aspect_ratio, + target_ratios, + width=orig_width, + height=orig_height, + image_size=image_size, + ) + + # calculate the target width and height + target_width = image_size * target_aspect_ratio[0] + target_height = image_size * target_aspect_ratio[1] + blocks = target_aspect_ratio[0] * target_aspect_ratio[1] + + # add thumbnail image if num_blocks != 1 + if use_thumbnail and blocks != 1: + blocks += 1 + + return blocks, target_width, target_height + + +def dynamic_preprocess_nemotron_vl( + image: Image.Image, + *, + target_ratios: list[tuple[int, int]], + image_size: int, + use_thumbnail: bool, +) -> list[Image.Image]: + orig_width, orig_height = image.size + + # calculate the number of blocks without thumbnail + blocks, target_width, target_height = calculate_nemotron_vl_targets( + orig_width=orig_width, + orig_height=orig_height, + target_ratios=target_ratios, + image_size=image_size, + use_thumbnail=False, + ) + + # resize the image + resized_img = image.resize((target_width, target_height)) + processed_images = [] + for i in range(blocks): + box = ((i % (target_width // image_size)) * image_size, + (i // (target_width // image_size)) * image_size, + ((i % (target_width // image_size)) + 1) * image_size, + ((i // (target_width // image_size)) + 1) * image_size) + # split the image + split_img = resized_img.crop(box) + processed_images.append(split_img) + + assert len(processed_images) == blocks + + if use_thumbnail and len(processed_images) != 1: + thumbnail_img = image.resize((image_size, image_size)) + processed_images.append(thumbnail_img) + + return processed_images + + +def get_nemotron_vl_target_ratios( + min_num: int, + max_num: int, +) -> list[tuple[int, int]]: + target_ratios = {(i, j) + for n in range(min_num, max_num + 1) + for i in range(1, n + 1) + for j in range(1, n + 1) if min_num <= i * j <= max_num} + return sorted(target_ratios, key=lambda x: x[0] * x[1]) + + +def image_to_pixel_values_nemotron_vl( + image: Image.Image, + *, + input_size: int, + min_num: int, + max_num: int, + use_thumbnail: bool, +) -> torch.Tensor: + target_ratios = get_nemotron_vl_target_ratios(min_num, max_num) + + transform = build_transform(input_size=input_size) + + images = dynamic_preprocess_nemotron_vl( + image, + target_ratios=target_ratios, + image_size=input_size, + use_thumbnail=use_thumbnail, + ) + + pixel_values = torch.stack([transform(image) for image in images]) + return pixel_values + + class NemotronVLProcessor(InternVLProcessor): def __init__( @@ -87,6 +229,50 @@ def __init__( def image_token_id(self) -> int: return self.tokenizer.get_vocab()[IMG_CONTEXT] + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + target_ratios = self.resolve_target_ratios( + use_thumbnail=False, # Applied in calculate_targets + ) + + num_patches, _, _ = calculate_nemotron_vl_targets( + orig_width=image_width, + orig_height=image_height, + image_size=self.image_size, + target_ratios=target_ratios, + use_thumbnail=self.use_thumbnail, + ) + + return num_patches * self.num_image_token + + def _images_to_pixel_values_lst( + self, + images: list[Image.Image], + min_dynamic_patch: Optional[int] = None, + max_dynamic_patch: Optional[int] = None, + dynamic_image_size: Optional[bool] = None, + ) -> list[torch.Tensor]: + min_num, max_num = self.resolve_min_max_num( + min_dynamic_patch=min_dynamic_patch, + max_dynamic_patch=max_dynamic_patch, + dynamic_image_size=dynamic_image_size, + use_thumbnail=False, # Applied in image_to_pixel_values + ) + + return [ + image_to_pixel_values_nemotron_vl( + image, + input_size=self.image_size, + min_num=min_num, + max_num=max_num, + use_thumbnail=self.use_thumbnail, + ) for image in images + ] + def _preprocess_image( self, text: list[str], From 96ddae46c164ec685c68012bd6fb6baf128fce03 Mon Sep 17 00:00:00 2001 From: 633WHU Date: Wed, 13 Aug 2025 19:10:07 +0800 Subject: [PATCH 030/233] [Doc] Add max_lora_rank configuration guide (#22782) Signed-off-by: chiliu --- docs/features/lora.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/features/lora.md b/docs/features/lora.md index a4e05dae11c2..668460a368a7 100644 --- a/docs/features/lora.md +++ b/docs/features/lora.md @@ -351,3 +351,22 @@ vllm serve ibm-granite/granite-speech-3.3-2b \ ``` Note: Default multimodal LoRAs are currently only available for `.generate` and chat completions. + +## Using Tips + +### Configuring `max_lora_rank` + +The `--max-lora-rank` parameter controls the maximum rank allowed for LoRA adapters. This setting affects memory allocation and performance: + +- **Set it to the maximum rank** among all LoRA adapters you plan to use +- **Avoid setting it too high** - using a value much larger than needed wastes memory and can cause performance issues + +For example, if your LoRA adapters have ranks [16, 32, 64], use `--max-lora-rank 64` rather than 256 + +```bash +# Good: matches actual maximum rank +vllm serve model --enable-lora --max-lora-rank 64 + +# Bad: unnecessarily high, wastes memory +vllm serve model --enable-lora --max-lora-rank 256 +``` From 24fddcf491309be2a99af303e244bff82f8b7681 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 13 Aug 2025 04:11:28 -0700 Subject: [PATCH 031/233] [V1] Add tree drafting tests for eagle spec decoding (#22705) Signed-off-by: Giancarlo Delfin --- tests/v1/spec_decode/test_eagle.py | 160 +++++++++++++++++++++++- tests/v1/spec_decode/test_max_len.py | 6 - vllm/v1/attention/backends/tree_attn.py | 6 +- vllm/v1/spec_decode/eagle.py | 61 +++------ 4 files changed, 178 insertions(+), 55 deletions(-) diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 2b4f8bd2a8b9..7b8445a0b287 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Optional from unittest import mock import pytest @@ -23,7 +24,11 @@ eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" -def _create_proposer(method: str, k: int) -> EagleProposer: +def _create_proposer( + method: str, + num_speculative_tokens: int, + speculative_token_tree: Optional[list[tuple[int]]] = None, +) -> EagleProposer: model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100) @@ -31,12 +36,18 @@ def _create_proposer(method: str, k: int) -> EagleProposer: # Choose model directory based on method draft_model_dir = eagle_dir if method == "eagle" else eagle3_dir + spec_token_tree_str = None + if speculative_token_tree is not None: + assert num_speculative_tokens == len(speculative_token_tree) + spec_token_tree_str = str(speculative_token_tree) + speculative_config = SpeculativeConfig( target_model_config=model_config, target_parallel_config=ParallelConfig(), model=draft_model_dir, method=method, - num_speculative_tokens=k, + num_speculative_tokens=num_speculative_tokens, + speculative_token_tree=spec_token_tree_str, ) vllm_config = VllmConfig( @@ -189,7 +200,7 @@ class _TargetModelStub(LlamaForCausalLM): target_model.lm_head = mock.MagicMock() # Create proposer using the helper function - proposer = _create_proposer(method, k=8) + proposer = _create_proposer(method, num_speculative_tokens=8) # Call the method under test proposer.load_model(target_model) @@ -226,6 +237,10 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): pytest.skip("TRITON_ATTN_VLLM_V1 does not support " "multi-token eagle spec decode on current platform") + if (attn_backend == "TREE_ATTN"): + pytest.skip("TREE_ATTN is tested separately in test_propose_tree" + "because it requires special input mocking.") + if attn_backend == "FLASH_ATTN_VLLM_V1" and current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") @@ -378,3 +393,142 @@ def create_deterministic_logits(token_ids): # Verify all tokens match our expectations assert torch.equal(result, expected_tokens) + + +@pytest.mark.parametrize( + "spec_token_tree", + [ + [(0, )], # A single token + [(0, ), (0, 0), (0, 0, 0)], # Chain + [(0, ), (1, ), (2, )], # Parallel + [(0, ), (1, ), (2, ), (0, 0), (0, 1), (1, 0), (1, 1), (2, 0), + (2, 1)], # Tree + ]) +def test_propose_tree(spec_token_tree): + # Get GPU device. + device = torch.device(current_platform.device_type) + + # Setup test parameters. + batch_size = 2 + seq_len_1 = 5 + seq_len_2 = 3 + total_tokens = seq_len_1 + seq_len_2 + vocab_size = 100 + seq_lens = [seq_len_1, seq_len_2] + num_speculative_tokens = len(spec_token_tree) + + # Create proposer first so we can use its actual hidden_size. + proposer = _create_proposer("eagle", + num_speculative_tokens, + speculative_token_tree=spec_token_tree) + # Get the hidden_size from the proposer to ensure consistency. + hidden_size = proposer.hidden_size + + # Helper to create deterministic logits that will produce specific tokens + def create_deterministic_logits(token_ids, k: int): + logits = torch.full((batch_size, vocab_size), -100.0, device=device) + for i, token_id in enumerate(token_ids): + # Assign decreasing values to the k, consecutive, tokens. + for j in range(k): + logits[i, token_id + j] = 100.0 - j + return logits + + # Mock a model that returns deterministic logits. + base_token_ids = torch.tensor([42, 60], dtype=torch.int64, device=device) + + # Skip loading the model and replace it with a mock that returns + # deterministic outputs. + model_mock = mock.MagicMock() + + # Mock the model forward calls. + forward_returns = [(torch.zeros(total_tokens, hidden_size, device=device), + torch.zeros(total_tokens, hidden_size, device=device))] + for cu_num_drafts in proposer.cu_drafts_per_level: + h_logits = torch.zeros(batch_size * cu_num_drafts, + hidden_size, + device=device) + h_states = torch.zeros(batch_size * cu_num_drafts, + hidden_size, + device=device) + forward_returns.append((h_logits, h_states)) + model_mock.side_effect = forward_returns + + # Mock the compute_logits calls. + cu_num_drafts_tensor = torch.tensor([0] + proposer.cu_drafts_per_level, + dtype=torch.int32, + device=device) + logits_returns = [] + for level, num_children in enumerate(proposer.child_drafts_per_level): + token_ids = base_token_ids + cu_num_drafts_tensor[level] + level_num_drafts = cu_num_drafts_tensor[ + level + 1] - cu_num_drafts_tensor[level] + level_logits = [] + for i in range(level_num_drafts // num_children): + level_logits.append( + create_deterministic_logits(token_ids + i * num_children, + num_children)) + logits_returns.append(torch.stack(level_logits, dim=1)) + model_mock.compute_logits.side_effect = logits_returns + + # Assign the mock to the proposer + proposer.model = model_mock + + # Assign draft attn_layer_names since load_model is not invoked + proposer.attn_layer_names = ["layer.0"] + + # Get the tree attention metadata builder. + attn_metadata_builder_cls, _ = get_attention_backend(_Backend.TREE_ATTN) + attn_metadata_builder = attn_metadata_builder_cls( + kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config), + layer_names=proposer.attn_layer_names, + vllm_config=proposer.vllm_config, + device=device, + ) + + # Mock runner for attention metadata building. + proposer.runner = mock.MagicMock() + proposer.runner.attn_groups.append([mock.MagicMock()]) + proposer.runner.attn_groups[0][0].metadata_builder = attn_metadata_builder + + # Setup inputs for the proposer. + target_token_ids = torch.randint(0, + vocab_size, (total_tokens, ), + device=device) + target_positions = torch.cat([ + torch.arange(seq_len_1, device=device), + torch.arange(seq_len_2, device=device) + ]) + target_hidden_states = torch.randn(total_tokens, + hidden_size, + device=device) + next_token_ids = torch.randint(0, + vocab_size, (batch_size, ), + dtype=torch.int32, + device=device) + batch_spec = BatchSpec( + seq_lens=seq_lens, + query_lens=seq_lens, + ) + common_attn_metadata = create_common_attn_metadata( + batch_spec, + block_size=16, + device=device, + ) + sampling_metadata = mock.MagicMock() + + # Propose draft tokens. + result = proposer.propose(target_token_ids=target_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + next_token_ids=next_token_ids, + common_attn_metadata=common_attn_metadata, + sampling_metadata=sampling_metadata) + assert result.shape == (batch_size, num_speculative_tokens) + + # The tokens are expected to be consecutive integers starting + # from the base token IDs. + expected_tokens = base_token_ids[:, None] + torch.arange( + num_speculative_tokens, dtype=torch.int64, device=device) + + # Verify that the draft tokens match our expectations. + assert torch.equal(result, expected_tokens) diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 01019b29e010..a5b10bb51866 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -39,12 +39,6 @@ def test_eagle_max_len(monkeypatch: pytest.MonkeyPatch, num_speculative_tokens: int, attn_backend: str): with monkeypatch.context() as m: m.setenv("VLLM_USE_V1", "1") - - if attn_backend == "TREE_ATTN" and num_speculative_tokens > 1: - # TREE_ATTN fails the test with multi-token spec decode - # TODO: Investigate why - pytest.skip("TREE_ATTN fails the test") - m.setenv("VLLM_ATTENTION_BACKEND", attn_backend) if (attn_backend == "TRITON_ATTN_VLLM_V1" diff --git a/vllm/v1/attention/backends/tree_attn.py b/vllm/v1/attention/backends/tree_attn.py index 3b53b039f1dc..5d10e9e26082 100644 --- a/vllm/v1/attention/backends/tree_attn.py +++ b/vllm/v1/attention/backends/tree_attn.py @@ -236,9 +236,9 @@ def build_for_drafting( # Use prefill for drafting at the root level. self.tree_attn_bias = torch.empty(0) else: - # Slice the tree attention bias for drafting. - query_len = common_attn_metadata.max_query_len - start, end = draft_index, draft_index + query_len + # Slice the tree attention bias for drafting. Exclude + # the root level. + start, end = 1, 1 + common_attn_metadata.max_query_len self.tree_attn_bias = self.tree_attn_bias[start:end, start:end].contiguous() diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index f75d76dd978f..a8a160a0f995 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -113,13 +113,6 @@ def __init__( num_drafts_per_level[level]) self.child_drafts_per_level.append(num_drafts_per_level[level] // num_drafts_per_level[level - 1]) - # Find the first level where the tree branches off into one or more - # children. - self.first_branching_level = None - for level in range(tree_depth): - if self.cu_drafts_per_level[level] > level + 1: - self.first_branching_level = level - break # Precompute draft position offsets in flattened tree. self.tree_draft_pos_offsets = torch.arange( 1, @@ -209,11 +202,10 @@ def propose( logits = self.model.compute_logits(sample_hidden_states, None) positions = target_positions[last_token_indices] hidden_states = hidden_states[last_token_indices] - if self.first_branching_level == 0: - # Branching has occurred at the root level. Draft using tree - # attention. + + if isinstance(attn_metadata, TreeAttentionMetadata): + # Draft using tree attention. draft_token_ids_list = self.propose_tree( - tree_root_level=0, batch_size=batch_size, logits=logits, positions=positions, @@ -242,11 +234,10 @@ def propose( (TritonAttentionMetadata, AiterFlashAttentionMetadata, FlashAttentionMetadata)) else: - # Currently, only FlashAttention and TreeAttention support - # multi-token eagle spec decode. This is because the code below - # makes assumptions about attn_metadata attributes available. - assert isinstance(attn_metadata, - (FlashAttentionMetadata, TreeAttentionMetadata)) + # Currently, only FlashAttention supports multi-token eagle spec + # decode. This is because the code below makes assumptions about + # attn_metadata attributes available. + assert isinstance(attn_metadata, FlashAttentionMetadata) # Generate the remaining draft tokens. draft_token_ids_list = [draft_token_ids] @@ -259,7 +250,7 @@ def propose( attn_metadata.num_actual_tokens = batch_size attn_metadata.max_query_len = 1 attn_metadata.query_start_loc = self.arange[:batch_size + 1] - for token_index in range(self.num_speculative_tokens - 1): + for _ in range(self.num_speculative_tokens - 1): # Update the inputs. # cast to int32 is crucial when eagle model is compiled. # tensor.argmax() returns int64 by default. @@ -327,21 +318,6 @@ def propose( hidden_states = hidden_states[:batch_size] logits = self.model.compute_logits(last_hidden_states[:batch_size], None) - - if self.first_branching_level == token_index + 1: - # Branching has occurred. The remaining tokens are drafted - # using tree attention. - draft_token_ids_list += self.propose_tree( - tree_root_level=token_index + 1, - batch_size=batch_size, - logits=logits, - positions=positions, - hidden_states=hidden_states, - common_attn_metadata=common_attn_metadata, - ) - # [batch_size, num_tree_tokens] - return torch.cat(draft_token_ids_list, dim=1) - draft_token_ids = logits.argmax(dim=-1) draft_token_ids_list.append(draft_token_ids) @@ -351,7 +327,6 @@ def propose( def propose_tree( self, - tree_root_level: int, batch_size: int, # [num_tokens, vocab_size] logits: torch.Tensor, @@ -366,10 +341,10 @@ def propose_tree( assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder) - total_num_drafts = self.cu_drafts_per_level[tree_root_level] + total_num_drafts = self.cu_drafts_per_level[0] level_num_drafts = total_num_drafts # Sample a draft token for each child at the tree root level. - num_children = self.child_drafts_per_level[tree_root_level] + num_children = self.child_drafts_per_level[0] if num_children == 1: draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1) else: @@ -393,22 +368,23 @@ def propose_tree( positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :]) tree_depth = len(self.cu_drafts_per_level) - for level in range(tree_root_level, tree_depth - 1): + for level in range(tree_depth - 1): # Get draft positions for RoPE. draft_positions = positions + (level + 1) exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len # Mask out the position ids that exceed the max model length. # Otherwise, we may get out-of-range error in RoPE. - clamped_draft_positions = torch.where( + draft_positions = torch.where( exceeds_max_model_len, 0, draft_positions, - ) + ).view(batch_size, -1) + if level_num_drafts > 1: # Repeat the positions for each draft at this level. - draft_positions = clamped_draft_positions.repeat_interleave( - level_num_drafts).reshape(batch_size, -1) + draft_positions = draft_positions.repeat_interleave( + level_num_drafts, dim=1) if num_children > 1: # Repeat draft hidden states for each child. @@ -425,7 +401,7 @@ def propose_tree( # Build new attention metadata for the next level of drafts. # This is necessary to support tree attention. - query_len = total_num_drafts - tree_root_level + query_len = total_num_drafts common_attn_metadata = replace( common_attn_metadata, query_start_loc=query_len * self.arange[:batch_size + 1], @@ -435,7 +411,7 @@ def propose_tree( ) attn_metadata = tree_attn_metadata_builder.build_for_drafting( common_attn_metadata=common_attn_metadata, - draft_index=tree_root_level + 1, + draft_index=level + 1, ) # Apply new attention metadata to all layers. @@ -516,7 +492,6 @@ def propose_tree( level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts total_num_drafts = self.cu_drafts_per_level[level + 1] - return draft_token_ids_list def prepare_inputs( From 4acdadb91e855d18b03708bbc1960a6c9398d950 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Wed, 13 Aug 2025 19:12:00 +0800 Subject: [PATCH 032/233] [Platform] Custom ops support for FusedMoe (#22509) Signed-off-by: wangxiyuan --- vllm/model_executor/layers/fused_moe/layer.py | 3 ++- vllm/model_executor/layers/linear.py | 12 ++++++------ .../layers/vocab_parallel_embedding.py | 4 +++- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 8ef0a805d86c..ddc02168e5c4 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -682,7 +682,8 @@ def determine_expert_map( return (local_num_experts, expert_map) -class FusedMoE(torch.nn.Module): +@CustomOp.register("fused_moe") +class FusedMoE(CustomOp): """FusedMoE layer for MoE models. This layer contains both MergedColumnParallel weights (gate_up_proj / diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index bb81a663d454..75391c51f775 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -16,6 +16,7 @@ tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce) from vllm.logger import init_logger +from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase) from vllm.model_executor.layers.utils import dispatch_unquantized_gemm @@ -226,7 +227,7 @@ def apply(self, return dispatch_unquantized_gemm()(layer, x, layer.weight, bias) -class LinearBase(torch.nn.Module): +class LinearBase(CustomOp): """Base linear layer. Args: @@ -269,12 +270,8 @@ def __init__( prefix=prefix) self.return_bias = return_bias - def forward( - self, x: torch.Tensor - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]: - raise NotImplementedError - +@CustomOp.register("replicated_linear") class ReplicatedLinear(LinearBase): """Replicated linear layer. @@ -443,6 +440,7 @@ def weight_loader(self, param[shard_offset:shard_offset + shard_size] = loaded_weight +@CustomOp.register("column_parallel_linear") class ColumnParallelLinear(LinearBase): """Linear layer with column parallelism. @@ -1229,6 +1227,7 @@ def weight_loader(self, param_data.copy_(loaded_weight) +@CustomOp.register("row_parallel_linear") class RowParallelLinear(LinearBase): """Linear layer with row parallelism. @@ -1405,6 +1404,7 @@ def extra_repr(self) -> str: return s +@CustomOp.register("qkv_cross_parallel_linear") class QKVCrossParallelLinear(LinearBase): """Linear layers for efficient cross-attention's QKV transformation. diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index a5f262c832bf..9f223998e554 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -12,6 +12,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce) +from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding) from vllm.model_executor.layers.utils import dispatch_unquantized_gemm @@ -159,7 +160,8 @@ def get_masked_input_and_mask( return input_, ~vocab_mask -class VocabParallelEmbedding(torch.nn.Module): +@CustomOp.register("vocab_parallel_embedding") +class VocabParallelEmbedding(CustomOp): """Embedding parallelized in the vocabulary dimension. Adapted from torch.nn.Embedding, note that we pad the vocabulary size to From 3821bba619cb1a74365752685f286b54d7d98863 Mon Sep 17 00:00:00 2001 From: Kdump Date: Wed, 13 Aug 2025 19:14:24 +0800 Subject: [PATCH 033/233] [Frontend] Add chunked processing to handle long inputs in embedding models (#22280) Signed-off-by: x22x22 Signed-off-by: Kdump Signed-off-by: DarkLight1337 Co-authored-by: Cyrus Leung Co-authored-by: Maximilien de Bayser Co-authored-by: DarkLight1337 --- .../openai_embedding_long_text/README.md | 186 +++++++ .../openai_embedding_long_text/client.py | 366 ++++++++++++++ .../openai_embedding_long_text/service.sh | 137 ++++++ .../openai/test_embedding_long_text.py | 441 +++++++++++++++++ vllm/config/__init__.py | 19 + vllm/entrypoints/openai/serving_embedding.py | 457 +++++++++++++++++- 6 files changed, 1603 insertions(+), 3 deletions(-) create mode 100644 examples/online_serving/openai_embedding_long_text/README.md create mode 100644 examples/online_serving/openai_embedding_long_text/client.py create mode 100644 examples/online_serving/openai_embedding_long_text/service.sh create mode 100644 tests/entrypoints/openai/test_embedding_long_text.py diff --git a/examples/online_serving/openai_embedding_long_text/README.md b/examples/online_serving/openai_embedding_long_text/README.md new file mode 100644 index 000000000000..04edc4680ea0 --- /dev/null +++ b/examples/online_serving/openai_embedding_long_text/README.md @@ -0,0 +1,186 @@ +# Long Text Embedding with Chunked Processing + +This directory contains examples for using vLLM's **chunked processing** feature to handle long text embedding that exceeds the model's maximum context length. + +## 🚀 Quick Start + +### Start the Server + +Use the provided script to start a vLLM server with chunked processing enabled: + +```bash +# Basic usage (supports very long texts up to ~3M tokens) +./service.sh + +# Custom configuration with different models +MODEL_NAME="jinaai/jina-embeddings-v3" \ +MAX_EMBED_LEN=1048576 \ +./service.sh + +# For extremely long documents +MODEL_NAME="intfloat/multilingual-e5-large" \ +MAX_EMBED_LEN=3072000 \ +./service.sh +``` + +### Test Long Text Embedding + +Run the comprehensive test client: + +```bash +python client.py +``` + +## 📁 Files + +| File | Description | +|------|-------------| +| `service.sh` | Server startup script with chunked processing enabled | +| `client.py` | Comprehensive test client for long text embedding | + +## ⚙️ Configuration + +### Server Configuration + +The key parameters for chunked processing are in the `--override-pooler-config`: + +```json +{ + "pooling_type": "auto", + "normalize": true, + "enable_chunked_processing": true, + "max_embed_len": 3072000 +} +``` + +!!! note + `pooling_type` sets the model's own pooling strategy for processing within each chunk. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length. + +#### Chunked Processing Behavior + +Chunked processing uses **MEAN aggregation** for cross-chunk combination when input exceeds the model's native maximum length: + +| Component | Behavior | Description | +|-----------|----------|-------------| +| **Within chunks** | Model's native pooling | Uses the model's configured pooling strategy | +| **Cross-chunk aggregation** | Always MEAN | Weighted averaging based on chunk token counts | +| **Performance** | Optimal | All chunks processed for complete semantic coverage | + +### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `MODEL_NAME` | `intfloat/multilingual-e5-large` | Embedding model to use (supports multiple models) | +| `PORT` | `31090` | Server port | +| `GPU_COUNT` | `1` | Number of GPUs to use | +| `MAX_EMBED_LEN` | `3072000` | Maximum embedding input length (supports very long documents) | +| `POOLING_TYPE` | `auto` | Model's native pooling type: `auto`, `MEAN`, `CLS`, `LAST` (only affects within-chunk pooling, not cross-chunk aggregation) | +| `API_KEY` | `EMPTY` | API key for authentication | + +## 🔧 How It Works + +1. **Enhanced Input Validation**: `max_embed_len` allows accepting inputs longer than `max_model_len` without environment variables +2. **Smart Chunking**: Text is split based on `max_position_embeddings` to maintain semantic integrity +3. **Unified Processing**: All chunks processed separately through the model using its configured pooling strategy +4. **MEAN Aggregation**: When input exceeds model's native length, results combined using token count-based weighted averaging across all chunks +5. **Consistent Output**: Final embeddings maintain the same dimensionality as standard processing + +### Input Length Handling + +- **Within max_embed_len**: Input is accepted and processed (up to 3M+ tokens) +- **Exceeds max_position_embeddings**: Chunked processing is automatically triggered +- **Exceeds max_embed_len**: Input is rejected with clear error message +- **No environment variables required**: Works without `VLLM_ALLOW_LONG_MAX_MODEL_LEN` + +### Extreme Long Text Support + +With `MAX_EMBED_LEN=3072000`, you can process: + +- **Academic papers**: Full research papers with references +- **Legal documents**: Complete contracts and legal texts +- **Books**: Entire chapters or small books +- **Code repositories**: Large codebases and documentation + +## 📊 Performance Characteristics + +### Chunked Processing Performance + +| Aspect | Behavior | Performance | +|--------|----------|-------------| +| **Chunk Processing** | All chunks processed with native pooling | Consistent with input length | +| **Cross-chunk Aggregation** | MEAN weighted averaging | Minimal overhead | +| **Memory Usage** | Proportional to number of chunks | Moderate, scalable | +| **Semantic Quality** | Complete text coverage | Optimal for long documents | + +## 🧪 Test Cases + +The test client demonstrates: + +- ✅ **Short text**: Normal processing (baseline) +- ✅ **Medium text**: Single chunk processing +- ✅ **Long text**: Multi-chunk processing with aggregation +- ✅ **Very long text**: Many chunks processing +- ✅ **Extreme long text**: Document-level processing (100K+ tokens) +- ✅ **Batch processing**: Mixed-length inputs in one request +- ✅ **Consistency**: Reproducible results across runs + +## 🐛 Troubleshooting + +### Common Issues + +1. **Chunked processing not enabled**: + + ```log + ValueError: This model's maximum position embeddings length is 4096 tokens... + ``` + + **Solution**: Ensure `enable_chunked_processing: true` in pooler config + +2. **Input exceeds max_embed_len**: + + ```log + ValueError: This model's maximum embedding input length is 3072000 tokens... + ``` + + **Solution**: Increase `max_embed_len` in pooler config or reduce input length + +3. **Memory errors**: + + ```log + RuntimeError: CUDA out of memory + ``` + + **Solution**: Reduce chunk size by adjusting model's `max_position_embeddings` or use fewer GPUs + +4. **Slow processing**: + **Expected**: Long text takes more time due to multiple inference calls + +### Debug Information + +Server logs show chunked processing activity: + +```log +INFO: Input length 150000 exceeds max_position_embeddings 4096, will use chunked processing +INFO: Split input of 150000 tokens into 37 chunks (max_chunk_size: 4096) +``` + +## 🤝 Contributing + +To extend chunked processing support to other embedding models: + +1. Check model compatibility with the pooling architecture +2. Test with various text lengths +3. Validate embedding quality compared to single-chunk processing +4. Submit PR with test cases and documentation updates + +## 🆕 Enhanced Features + +### max_embed_len Parameter + +The new `max_embed_len` parameter provides: + +- **Simplified Configuration**: No need for `VLLM_ALLOW_LONG_MAX_MODEL_LEN` environment variable +- **Flexible Input Validation**: Accept inputs longer than `max_model_len` up to `max_embed_len` +- **Extreme Length Support**: Process documents with millions of tokens +- **Clear Error Messages**: Better feedback when inputs exceed limits +- **Backward Compatibility**: Existing configurations continue to work diff --git a/examples/online_serving/openai_embedding_long_text/client.py b/examples/online_serving/openai_embedding_long_text/client.py new file mode 100644 index 000000000000..6e9838ac6d8d --- /dev/null +++ b/examples/online_serving/openai_embedding_long_text/client.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Example script demonstrating long text embedding with chunked processing in vLLM. + +This example shows how to use vLLM's chunked processing feature to handle text +inputs that exceed the model's maximum token length. The feature automatically +splits long text into chunks and handles different pooling types optimally. + +Prerequisites: +1. Start vLLM server with chunked processing enabled: + + # MEAN pooling (processes all chunks, recommended for complete coverage) + vllm serve intfloat/multilingual-e5-large \ + --override-pooler-config \ + '{"pooling_type": "MEAN", "normalize": true, ' \ + '"enable_chunked_processing": true, "max_embed_len": 3072000}' \ + --served-model-name multilingual-e5-large \ + --trust-remote-code \ + --port 31090 \ + --api-key your-api-key + + # OR CLS pooling (native CLS within chunks, MEAN aggregation across chunks) + vllm serve BAAI/bge-large-en-v1.5 \ + --override-pooler-config \ + '{"pooling_type": "CLS", "normalize": true, ' \ + '"enable_chunked_processing": true, "max_embed_len": 1048576}' \ + --served-model-name bge-large-en-v1.5 \ + --trust-remote-code \ + --port 31090 \ + --api-key your-api-key + +2. Install required dependencies: + pip install openai requests +""" + +import time + +import numpy as np +from openai import OpenAI + +# Configuration +API_KEY = "your-api-key" # Replace with your actual API key +BASE_URL = "http://localhost:31090/v1" +MODEL_NAME = "multilingual-e5-large" + + +def generate_long_text(base_text: str, repeat_count: int) -> str: + """Generate long text by repeating base text.""" + return base_text * repeat_count + + +def test_embedding_with_different_lengths(): + """Test embedding generation with different text lengths.""" + client = OpenAI(api_key=API_KEY, base_url=BASE_URL) + + # Test cases with different text lengths + test_cases = [ + { + "name": "Short Text", + "text": "Hello, this is a short text for embedding.", + "expected_chunks": 1, + }, + { + "name": "Medium Text", + "text": generate_long_text( + "This is a medium-length text that should fit within the " + "model's context window. " * 20, + 2, + ), + "expected_chunks": 1, + }, + { + "name": "Long Text (2 chunks)", + "text": generate_long_text( + "This is a very long text that will exceed the model's " + "maximum context length and trigger chunked processing. " * 50, + 5, + ), + "expected_chunks": 2, + }, + { + "name": "Very Long Text (3+ chunks)", + "text": generate_long_text( + "This text is extremely long and will definitely " + "require multiple chunks for processing. " * 100, + 10, + ), + "expected_chunks": 3, + }, + ] + + print("🧪 Testing vLLM Long Text Embedding with Chunked Processing") + print("=" * 70) + + for i, test_case in enumerate(test_cases, 1): + print(f"\n📝 Test {i}: {test_case['name']}") + print(f"Text length: {len(test_case['text'])} characters") + + try: + start_time = time.time() + + response = client.embeddings.create( + input=test_case["text"], model=MODEL_NAME, encoding_format="float" + ) + + end_time = time.time() + processing_time = end_time - start_time + + # Extract embedding data + embedding = response.data[0].embedding + embedding_dim = len(embedding) + + print("✅ Success!") + print(f" - Embedding dimension: {embedding_dim}") + print(f" - Processing time: {processing_time:.2f}s") + print(f" - Expected chunks: ~{test_case['expected_chunks']}") + print(f" - First 5 values: {embedding[:5]}") + + except Exception as e: + print(f"❌ Failed: {str(e)}") + + +def test_batch_embedding(): + """Test batch embedding with mixed-length inputs.""" + client = OpenAI(api_key=API_KEY, base_url=BASE_URL) + + print("\n🔄 Testing Batch Embedding with Mixed Lengths") + print("=" * 50) + + # Mix of short and long texts + batch_inputs = [ + "Short text 1", + generate_long_text("Medium length text that fits in one chunk. " * 20, 1), + "Another short text", + generate_long_text("Long text requiring chunked processing. " * 100, 5), + ] + + try: + start_time = time.time() + + response = client.embeddings.create( + input=batch_inputs, model=MODEL_NAME, encoding_format="float" + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("✅ Batch processing successful!") + print(f" - Number of inputs: {len(batch_inputs)}") + print(f" - Number of embeddings: {len(response.data)}") + print(f" - Total processing time: {processing_time:.2f}s") + print( + f" - Average time per input: {processing_time / len(batch_inputs):.2f}s" + ) + + for i, data in enumerate(response.data): + input_length = len(batch_inputs[i]) + embedding_dim = len(data.embedding) + print( + f" - Input {i + 1}: {input_length} chars → {embedding_dim}D embedding" + ) + + except Exception as e: + print(f"❌ Batch processing failed: {str(e)}") + + +def test_multiple_long_texts_batch(): + """Test batch processing with multiple long texts to verify chunk ID uniqueness.""" + client = OpenAI(api_key=API_KEY, base_url=BASE_URL) + + print("\n🔧 Testing Multiple Long Texts in Batch (Chunk ID Fix Verification)") + print("=" * 70) + + # Create multiple distinct long texts that will all require chunking + # Note: All pooling types now use MEAN aggregation across chunks: + # - Native pooling (MEAN/CLS/LAST) is used within each chunk + # - MEAN aggregation combines results across all chunks + # - Full semantic coverage for all pooling types + long_texts = [ + generate_long_text( + "First long document about artificial intelligence and machine learning. " + * 80, + 6, + ), + generate_long_text( + "Second long document about natural language processing and transformers. " + * 80, + 6, + ), + generate_long_text( + "Third long document about computer vision and neural networks. " * 80, 6 + ), + ] + + # Add some short texts to mix things up + batch_inputs = [ + "Short text before long texts", + long_texts[0], + "Short text between long texts", + long_texts[1], + long_texts[2], + "Short text after long texts", + ] + + print("📊 Batch composition:") + for i, text in enumerate(batch_inputs): + length = len(text) + text_type = "Long (will be chunked)" if length > 5000 else "Short" + print(f" - Input {i + 1}: {length} chars ({text_type})") + + try: + start_time = time.time() + + response = client.embeddings.create( + input=batch_inputs, model=MODEL_NAME, encoding_format="float" + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n✅ Multiple long texts batch processing successful!") + print(f" - Number of inputs: {len(batch_inputs)}") + print(f" - Number of embeddings returned: {len(response.data)}") + print(f" - Total processing time: {processing_time:.2f}s") + + # Verify each embedding is different (no incorrect aggregation) + embeddings = [data.embedding for data in response.data] + + if len(embeddings) >= 3: + import numpy as np + + # Compare embeddings of the long texts (indices 1, 3, 4) + long_embeddings = [ + np.array(embeddings[1]), # First long text + np.array(embeddings[3]), # Second long text + np.array(embeddings[4]), # Third long text + ] + + print("\n🔍 Verifying embedding uniqueness:") + for i in range(len(long_embeddings)): + for j in range(i + 1, len(long_embeddings)): + cosine_sim = np.dot(long_embeddings[i], long_embeddings[j]) / ( + np.linalg.norm(long_embeddings[i]) + * np.linalg.norm(long_embeddings[j]) + ) + print( + f" - Similarity between long text {i + 1} and {j + 1}: " + f"{cosine_sim:.4f}" + ) + + if ( + cosine_sim < 0.9 + ): # Different content should have lower similarity + print(" ✅ Good: Embeddings are appropriately different") + else: + print( + " ⚠️ High similarity - may indicate chunk " + "aggregation issue" + ) + + print("\n📋 Per-input results:") + for i, data in enumerate(response.data): + input_length = len(batch_inputs[i]) + embedding_dim = len(data.embedding) + embedding_norm = np.linalg.norm(data.embedding) + print( + f" - Input {i + 1}: {input_length} chars → {embedding_dim}D " + f"embedding (norm: {embedding_norm:.4f})" + ) + + print( + "\n✅ This test verifies the fix for chunk ID collisions in " + "batch processing" + ) + print(" - Before fix: Multiple long texts would have conflicting chunk IDs") + print(" - After fix: Each prompt's chunks have unique IDs with prompt index") + + except Exception as e: + print(f"❌ Multiple long texts batch test failed: {str(e)}") + print(" This might indicate the chunk ID collision bug is present!") + + +def test_embedding_consistency(): + """Test that chunked processing produces consistent results.""" + client = OpenAI(api_key=API_KEY, base_url=BASE_URL) + + print("\n🔍 Testing Embedding Consistency") + print("=" * 40) + + # Use the same long text multiple times + long_text = generate_long_text( + "Consistency test text for chunked processing validation. " * 50, 3 + ) + + embeddings = [] + + try: + for i in range(3): + response = client.embeddings.create( + input=long_text, model=MODEL_NAME, encoding_format="float" + ) + embeddings.append(response.data[0].embedding) + print(f" - Generated embedding {i + 1}") + + # Check consistency (embeddings should be identical) + if len(embeddings) >= 2: + # Calculate similarity between first two embeddings + + emb1 = np.array(embeddings[0]) + emb2 = np.array(embeddings[1]) + + # Cosine similarity + cosine_sim = np.dot(emb1, emb2) / ( + np.linalg.norm(emb1) * np.linalg.norm(emb2) + ) + + print("✅ Consistency test completed!") + print(f" - Cosine similarity between runs: {cosine_sim:.6f}") + print(" - Expected: ~1.0 (identical embeddings)") + + if cosine_sim > 0.999: + print(" - ✅ High consistency achieved!") + else: + print(" - ⚠️ Consistency may vary due to numerical precision") + + except Exception as e: + print(f"❌ Consistency test failed: {str(e)}") + + +def main(): + """Main function to run all tests.""" + print("🚀 vLLM Long Text Embedding Client") + print(f"📡 Connecting to: {BASE_URL}") + print(f"🤖 Model: {MODEL_NAME}") + masked_key = "*" * (len(API_KEY) - 4) + API_KEY[-4:] if len(API_KEY) > 4 else "****" + print(f"🔑 API Key: {masked_key}") + + # Run all test cases + test_embedding_with_different_lengths() + test_batch_embedding() + test_multiple_long_texts_batch() + test_embedding_consistency() + + print("\n" + "=" * 70) + print("🎉 All tests completed!") + print("\n💡 Key Features Demonstrated:") + print(" - ✅ Automatic chunked processing for long text") + print(" - ✅ Seamless handling of mixed-length batches") + print(" - ✅ Multiple long texts in single batch (chunk ID fix)") + print(" - ✅ Unified chunked processing:") + print(" • Native pooling used within each chunk") + print(" • MEAN aggregation across all chunks") + print(" • Complete semantic coverage for all pooling types") + print(" - ✅ Consistent embedding generation") + print(" - ✅ Backward compatibility with short text") + print("\n📚 For more information, see:") + print( + " - Documentation: https://docs.vllm.ai/en/latest/models/pooling_models.html" + ) + print(" - Chunked Processing Guide: openai_embedding_long_text.md") + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/openai_embedding_long_text/service.sh b/examples/online_serving/openai_embedding_long_text/service.sh new file mode 100644 index 000000000000..f356d7d4529e --- /dev/null +++ b/examples/online_serving/openai_embedding_long_text/service.sh @@ -0,0 +1,137 @@ +#!/bin/bash + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# vLLM Embedding Server with Enhanced Chunked Processing +# This script starts a vLLM server with chunked processing enabled for long text embedding. +# Now supports proper pooling type validation and model-specific configurations. + +set -euo pipefail + +# Configuration +MODEL_NAME=${MODEL_NAME:-"intfloat/multilingual-e5-large"} +MODEL_CODE=${MODEL_CODE:-"multilingual-e5-large"} + +PORT=${PORT:-31090} +GPU_COUNT=${GPU_COUNT:-1} +MAX_EMBED_LEN=${MAX_EMBED_LEN:-3072000} +API_KEY=${API_KEY:-"your-api-key"} + +# Enhanced pooling configuration with model-specific defaults +POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST +export VLLM_ENABLE_CHUNKED_PROCESSING=true +export CUDA_VISIBLE_DEVICES=2,3,4,5 +# export VLLM_ATTENTION_BACKEND=XFORMERS + +echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing" +echo "==================================================================" + +# Environment variables for optimization +export VLLM_WORKER_MULTIPROC_METHOD=spawn + +# Function to determine optimal pooling type for known models +get_optimal_pooling_type() { + local model="$1" + case "$model" in + *"e5-"* | *"multilingual-e5"*) + echo "MEAN" # E5 series native pooling + ;; + *"bge-"*) + echo "CLS" # BGE series native pooling + ;; + *"gte-"*) + echo "LAST" # GTE series native pooling + ;; + *"sentence-t5"* | *"st5"*) + echo "MEAN" # Sentence-T5 native pooling + ;; + *"jina-embeddings"*) + echo "MEAN" # Jina embeddings native pooling + ;; + *"Qwen"*"Embedding"*) + echo "LAST" # Qwen embeddings native pooling + ;; + *) + echo "MEAN" # Default native pooling for unknown models + ;; + esac +} + +# Auto-detect pooling type if not explicitly set +if [ "$POOLING_TYPE" = "auto" ]; then + POOLING_TYPE=$(get_optimal_pooling_type "$MODEL_NAME") + echo "🔍 Auto-detected pooling type: $POOLING_TYPE for model $MODEL_NAME" +fi + +# Display configuration +echo "📋 Configuration:" +echo " - Model: $MODEL_NAME" +echo " - Port: $PORT" +echo " - GPU Count: $GPU_COUNT" +echo " - Enhanced Chunked Processing: ${VLLM_ENABLE_CHUNKED_PROCESSING}" +echo " - Max Embed Length: ${MAX_EMBED_LEN} tokens" +echo " - Native Pooling Type: $POOLING_TYPE + Normalization" +echo " - Cross-chunk Aggregation: MEAN (automatic)" +echo "" + +# Validate GPU availability +if command -v nvidia-smi &> /dev/null; then + gpu_count=$(nvidia-smi --list-gpus | wc -l) + echo "🖥️ Available GPUs: $gpu_count" + if [ "$GPU_COUNT" -gt "$gpu_count" ]; then + echo "⚠️ Warning: Requested $GPU_COUNT GPUs but only $gpu_count available" + echo " Adjusting to use $gpu_count GPUs" + GPU_COUNT=$gpu_count + fi +else + echo "⚠️ Warning: nvidia-smi not found. GPU detection skipped." +fi + +# Chunked processing uses unified MEAN aggregation +echo "ℹ️ Chunked Processing: Using $POOLING_TYPE pooling within chunks, MEAN aggregation across chunks" +echo " - All chunks processed for complete semantic coverage" +echo " - Weighted averaging based on chunk token counts" + +echo "" +echo "🔧 Starting server with enhanced chunked processing configuration..." + +# Build pooler config JSON +POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"normalize\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}" + +# Start vLLM server with enhanced chunked processing +vllm serve "$MODEL_NAME" \ + --tensor-parallel-size "$GPU_COUNT" \ + --enforce-eager \ + --override-pooler-config "$POOLER_CONFIG" \ + --served-model-name ${MODEL_CODE} \ + --api-key "$API_KEY" \ + --trust-remote-code \ + --port "$PORT" \ + --host 0.0.0.0 + +echo "" +echo "✅ vLLM Embedding Server started successfully!" +echo "" +echo "📡 Server Information:" +echo " - Base URL: http://localhost:$PORT" +echo " - Model Code: ${MODEL_CODE}" +echo " - API Key: $API_KEY" +echo " - Native Pooling: $POOLING_TYPE | Cross-chunk: MEAN" +echo "" +echo "🧪 Test the server with:" +echo " python examples/online_serving/openai_embedding_long_text_client.py" +echo "" +echo "📚 Enhanced features enabled:" +echo " ✅ Intelligent native pooling type detection" +echo " ✅ Unified MEAN aggregation for chunked processing" +echo " ✅ Model-specific native pooling optimization" +echo " ✅ Enhanced max embedding length (${MAX_EMBED_LEN} tokens)" +echo " ✅ Complete semantic coverage for all pooling types" +echo " ✅ OpenAI-compatible API" +echo " ✅ GPU acceleration" +echo "" +echo "🔧 Advanced usage:" +echo " - Set POOLING_TYPE=MEAN|CLS|LAST to override auto-detection" +echo " - Set MAX_EMBED_LEN to adjust maximum input length" +echo " - All pooling types use MEAN aggregation across chunks" diff --git a/tests/entrypoints/openai/test_embedding_long_text.py b/tests/entrypoints/openai/test_embedding_long_text.py new file mode 100644 index 000000000000..86bd34abb97e --- /dev/null +++ b/tests/entrypoints/openai/test_embedding_long_text.py @@ -0,0 +1,441 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Test cases for long text embedding with automatic chunking mechanism. + +This test suite validates vLLM's automatic chunking functionality for handling +text inputs that exceed the model's maximum token length, specifically targeting +the intfloat/multilingual-e5-small model (max token length: 512). +""" + +import random + +import openai +import pytest +import pytest_asyncio + +from vllm.entrypoints.openai.protocol import EmbeddingResponse + +from ...utils import RemoteOpenAIServer + + +def _generate_random_text(word_count: int) -> str: + """Generate random text with approximately the specified word count.""" + # Common English words with focus on verbs and nouns for realistic text + common_words = [ + # Essential articles and pronouns (minimal) + "the", + "and", + "you", + "they", + "this", + "that", + "these", + "those", + + # Action verbs + "create", + "build", + "develop", + "design", + "implement", + "execute", + "analyze", + "process", + "generate", + "calculate", + "evaluate", + "optimize", + "transform", + "integrate", + "configure", + "deploy", + "monitor", + "manage", + "discover", + "explore", + "investigate", + "research", + "study", + "examine", + "improve", + "enhance", + "upgrade", + "modify", + "update", + "maintain", + "solve", + "resolve", + "handle", + "address", + "tackle", + "overcome", + "communicate", + "collaborate", + "coordinate", + "organize", + "plan", + "achieve", + "accomplish", + "complete", + "finish", + "deliver", + "provide", + + # Technology and science nouns + "system", + "application", + "software", + "hardware", + "network", + "database", + "algorithm", + "model", + "framework", + "platform", + "interface", + "protocol", + "architecture", + "infrastructure", + "component", + "module", + "service", + "technology", + "innovation", + "solution", + "methodology", + "approach", + "artificial", + "intelligence", + "machine", + "learning", + "neural", + "network", + "computer", + "processor", + "memory", + "storage", + "computation", + "data", + "information", + "knowledge", + "insight", + "pattern", + "trend", + "analysis", + "research", + "development", + "engineering", + "science", + "mathematics", + "statistics", + "probability", + "optimization", + "performance", + "efficiency", + + # General nouns + "project", + "team", + "organization", + "company", + "business", + "industry", + "market", + "customer", + "user", + "client", + "product", + "feature", + "function", + "requirement", + "specification", + "documentation", + "report", + "result", + "outcome", + "impact", + "benefit", + "advantage", + "challenge", + "problem", + "opportunity", + "strategy", + "goal", + "objective", + "target", + "milestone", + "process", + "procedure", + "workflow", + "pipeline", + "operation", + "task", + "activity", + "event", + "session", + "meeting", + "discussion", + "decision" + ] + + words = [] + for _ in range(word_count): + words.append(random.choice(common_words)) + + # Add some punctuation for more realistic text + text = " ".join(words) + # Add periods every 10-20 words + words_list = text.split() + result = [] + for i, word in enumerate(words_list): + result.append(word) + if ((i + 1) % random.randint(10, 20) == 0 and i < len(words_list) - 1): + result[-1] += "." + + return " ".join(result) + + +MODEL_NAME = "intfloat/multilingual-e5-small" +DTYPE = "bfloat16" + +# Test text: Generate text with approximately 1500 words to exceed 1024 tokens +LONG_TEXT_1500_WORDS = _generate_random_text(1500) + +# Test text: Generate text with approximately 2500 words to exceed 2048 tokens +LONG_TEXT_2500_WORDS = _generate_random_text(2500) + + +@pytest.fixture(scope="module") +def server_with_chunked_processing(): + """Start server with automatic chunking processing enabled.""" + args = [ + "--runner", + "pooling", + "--dtype", + DTYPE, + "--enforce-eager", + "--max-model-len", + "512", # Set smaller max_model_len to trigger chunking mechanism + '--override-pooler-config', + ('{"pooling_type": "MEAN", "normalize": true, ' + '"enable_chunked_processing": true, "max_embed_len": 10000}'), + "--gpu-memory-utilization", + "0.8", + ] + + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client_with_chunked_processing(server_with_chunked_processing): + """Create async client with chunking processing support.""" + async with server_with_chunked_processing.get_async_client( + ) as async_client: + yield async_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_long_text_embedding_1500_chars( + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str): + """Test embedding processing for ~1500 character long text + (~1028 tokens, exceeding 512 token limit).""" + + # Verify text length + # Verify text has sufficient word count (approximately 1500 words) + word_count = len(LONG_TEXT_1500_WORDS.split()) + assert word_count >= 1400, ( + f"Test text word count insufficient: {word_count} words") + + # Send embedding request + embedding_response = await client_with_chunked_processing.embeddings.create( + model=model_name, + input=[LONG_TEXT_1500_WORDS], + encoding_format="float", + ) + + # Verify response structure + embeddings = EmbeddingResponse.model_validate( + embedding_response.model_dump(mode="json")) + + assert embeddings.id is not None + assert len(embeddings.data) == 1 + assert len(embeddings.data[0].embedding + ) == 384 # multilingual-e5-small embedding dimension + assert embeddings.usage.completion_tokens == 0 + # Due to chunked processing, token count should + # reflect actual processed tokens + # With ~1500 words, we expect roughly + # 1024+ tokens (exceeding 512 token limit) + # Should exceed single chunk limit of 512 + assert embeddings.usage.prompt_tokens > 800 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens + + # Verify embedding vector validity + embedding_vector = embeddings.data[0].embedding + assert all( + isinstance(x, float) + for x in embedding_vector), "Embedding vector should contain floats" + assert not all( + x == 0 + for x in embedding_vector), "Embedding vector should not be all zeros" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_long_text_embedding_2500_chars( + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str): + """Test embedding processing for ~2500 character long text + (~2048 tokens, requiring multiple chunks).""" + + # Verify text length + # Verify text has sufficient word count (approximately 2500 words) + word_count = len(LONG_TEXT_2500_WORDS.split()) + assert word_count >= 2300, ( + f"Test text word count insufficient: {word_count} words") + + # Send embedding request + embedding_response = await client_with_chunked_processing.embeddings.create( + model=model_name, + input=[LONG_TEXT_2500_WORDS], + encoding_format="float", + ) + + # Verify response structure + embeddings = EmbeddingResponse.model_validate( + embedding_response.model_dump(mode="json")) + + assert embeddings.id is not None + assert len(embeddings.data) == 1 + assert len(embeddings.data[0].embedding + ) == 384 # multilingual-e5-small embedding dimension + assert embeddings.usage.completion_tokens == 0 + # Due to chunked processing, token count should + # reflect actual processed tokens + # With ~2500 words, we expect + # roughly 2048+ tokens (requiring multiple chunks) + # Should require multiple chunks for processing + assert embeddings.usage.prompt_tokens > 1500 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens + + # Verify embedding vector validity + embedding_vector = embeddings.data[0].embedding + assert all( + isinstance(x, float) + for x in embedding_vector), "Embedding vector should contain floats" + assert not all( + x == 0 + for x in embedding_vector), "Embedding vector should not be all zeros" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_batch_long_text_embedding( + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str): + """Test batch long text embedding processing.""" + + input_texts = [ + LONG_TEXT_1500_WORDS, + LONG_TEXT_2500_WORDS, + "This is a short text test.", # Short text for comparison + ] + + # Send batch embedding request + embedding_response = await client_with_chunked_processing.embeddings.create( + model=model_name, + input=input_texts, + encoding_format="float", + ) + + # Verify response structure + embeddings = EmbeddingResponse.model_validate( + embedding_response.model_dump(mode="json")) + + assert embeddings.id is not None + assert len(embeddings.data) == 3 # Three input texts + + # Verify each embedding dimension + for i, embedding_data in enumerate(embeddings.data): + assert len(embedding_data.embedding) == 384 + assert embedding_data.index == i + + # Verify embedding vector validity + embedding_vector = embedding_data.embedding + assert all(isinstance(x, float) for x in embedding_vector) + assert not all(x == 0 for x in embedding_vector) + + # Verify token usage + assert embeddings.usage.completion_tokens == 0 + # Total token count should be very substantial + assert embeddings.usage.prompt_tokens > 1000 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_chunked_vs_normal_consistency( + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str): + """Test consistency between chunked and + normal processing (using short text).""" + + # Use a short text within the 512 token limit + short_text = ("Artificial intelligence technology is changing our world, " + "bringing unprecedented opportunities and challenges.") + + # Send embedding request + embedding_response = await client_with_chunked_processing.embeddings.create( + model=model_name, + input=[short_text], + encoding_format="float", + ) + + # Verify response structure + embeddings = EmbeddingResponse.model_validate( + embedding_response.model_dump(mode="json")) + + assert embeddings.id is not None + assert len(embeddings.data) == 1 + assert len(embeddings.data[0].embedding) == 384 + assert embeddings.usage.completion_tokens == 0 + # Short text should not require chunked processing + assert embeddings.usage.prompt_tokens < 512 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens + + # 验证embedding向量的有效性 + embedding_vector = embeddings.data[0].embedding + assert all(isinstance(x, float) for x in embedding_vector) + assert not all(x == 0 for x in embedding_vector) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_chunked_processing_response_format( + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str): + """Test response format and structure during chunked processing.""" + + # Test with long text to trigger chunking + embedding_response = await client_with_chunked_processing.embeddings.create( + model=model_name, + input=[LONG_TEXT_1500_WORDS], + encoding_format="float", + ) + + # Verify response structure + embeddings = EmbeddingResponse.model_validate( + embedding_response.model_dump(mode="json")) + + assert embeddings.id is not None + assert len(embeddings.data) == 1 + assert embeddings.data[0].object == "embedding" + assert embeddings.data[0].index == 0 + + # Verify embedding vector properties + embedding_vector = embeddings.data[0].embedding + import math + vector_norm = math.sqrt(sum(x * x for x in embedding_vector)) + # Check that the vector is normalized + # (default behavior for most embedding models) + assert 0.8 < vector_norm < 1.2, ( + f"Vector norm should be reasonable, actual: {vector_norm}") diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 6649cd89ee34..b4ea15ef5a0f 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -2598,6 +2598,25 @@ class PoolerConfig: ``math-shepherd-mistral-7b-prm`` model. """ + enable_chunked_processing: Optional[bool] = None + """ + Whether to enable chunked processing for long inputs that exceed the model's + maximum position embeddings. When enabled, long inputs will be split into + chunks, processed separately, and then aggregated using weighted averaging. + This allows embedding models to handle arbitrarily long text without CUDA + errors. Defaults to False. + """ + + max_embed_len: Optional[int] = None + """ + Maximum input length allowed for embedding generation. When set, allows + inputs longer than max_embed_len to be accepted for embedding models. + This parameter enables accepting long inputs without requiring + VLLM_ALLOW_LONG_MAX_MODEL_LEN environment variable. When an input exceeds + max_embed_len, it will be handled according to the original max_model_len + validation logic. Defaults to None (i.e. set to max_model_len). + """ + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, diff --git a/vllm/entrypoints/openai/serving_embedding.py b/vllm/entrypoints/openai/serving_embedding.py index 84ba00873103..9dcad8e391c6 100644 --- a/vllm/entrypoints/openai/serving_embedding.py +++ b/vllm/entrypoints/openai/serving_embedding.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import base64 -from typing import Final, Literal, Optional, Union, cast +from collections.abc import AsyncGenerator, Mapping +from typing import Any, Final, Literal, Optional, Union, cast import numpy as np +import torch from fastapi import Request from typing_extensions import assert_never, override @@ -12,19 +14,28 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.logger import RequestLogger +# yapf conflicts with isort for this docstring +# yapf: disable from vllm.entrypoints.openai.protocol import (EmbeddingChatRequest, + EmbeddingCompletionRequest, EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, ErrorResponse, UsageInfo) from vllm.entrypoints.openai.serving_engine import (EmbeddingServeContext, OpenAIServing, - ServeContext) + RequestPrompt, + ServeContext, + TextTokensPrompt) +# yapf: enable from vllm.entrypoints.openai.serving_models import OpenAIServingModels +from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt +from vllm.inputs.data import TokensPrompt as EngineTokensPrompt from vllm.logger import init_logger from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput, - PoolingRequestOutput) + PoolingOutput, PoolingRequestOutput, RequestOutput) from vllm.pooling_params import PoolingParams +from vllm.utils import chunk_list logger = init_logger(__name__) @@ -46,6 +57,17 @@ def _get_embedding( class EmbeddingMixin(OpenAIServing): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + pooler_config = self.model_config.pooler_config + + # Avoid repeated attribute lookups + self.supports_chunked_processing = bool( + pooler_config and pooler_config.enable_chunked_processing) + self.max_embed_len = (pooler_config.max_embed_len if pooler_config + and pooler_config.max_embed_len else None) + @override async def _preprocess( self, @@ -129,6 +151,435 @@ def _build_response( usage=usage, ) + def _get_max_position_embeddings(self) -> int: + """Get the model's effective maximum sequence length for chunking.""" + return self.model_config.max_model_len + + def _should_use_chunked_processing(self, request) -> bool: + """Check if chunked processing should be used for this request.""" + return isinstance( + request, + (EmbeddingCompletionRequest, + EmbeddingChatRequest)) and self.supports_chunked_processing + + async def _process_chunked_request( + self, + ctx: EmbeddingServeContext, + original_prompt: TextTokensPrompt, + pooling_params, + trace_headers, + prompt_idx: int, + ) -> list[AsyncGenerator[PoolingRequestOutput, None]]: + """Process a single prompt using chunked processing.""" + generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] + token_ids = original_prompt["prompt_token_ids"] + + # Split into chunks using max_position_embeddings + max_pos_embeddings = self._get_max_position_embeddings() + # Process all chunks for MEAN aggregation + for chunk_idx, chunk_tokens in enumerate( + chunk_list(token_ids, max_pos_embeddings)): + # Create a request ID for this chunk + chunk_request_id = (f"{ctx.request_id}-prompt-{prompt_idx}-" + f"chunk-{chunk_idx}") + + # Create engine prompt for this chunk + chunk_engine_prompt = EngineTokensPrompt( + prompt_token_ids=chunk_tokens) + + # Create chunk request prompt for logging + chunk_text = "" + chunk_request_prompt = TextTokensPrompt( + prompt=chunk_text, prompt_token_ids=chunk_tokens) + + # Log the chunk + self._log_inputs(chunk_request_id, + chunk_request_prompt, + params=pooling_params, + lora_request=ctx.lora_request) + + # Create generator for this chunk and wrap it to return indices + original_generator = self.engine_client.encode( + chunk_engine_prompt, + pooling_params, + chunk_request_id, + lora_request=ctx.lora_request, + trace_headers=trace_headers, + priority=getattr(ctx.request, "priority", 0), + ) + + generators.append(original_generator) + + return generators + + def _validate_input( + self, + request, + input_ids: list[int], + input_text: str, + ) -> TextTokensPrompt: + """Override to support chunked processing for embedding requests.""" + token_num = len(input_ids) + + # Note: EmbeddingRequest doesn't have max_tokens + if isinstance(request, + (EmbeddingCompletionRequest, EmbeddingChatRequest)): + # Check if chunked processing is enabled for pooling models + enable_chunked = self._should_use_chunked_processing(request) + + # Use max_position_embeddings for chunked processing decisions + max_pos_embeddings = self._get_max_position_embeddings() + + # Determine the effective max length for validation + if self.max_embed_len is not None: + # Use max_embed_len for validation instead of max_model_len + length_type = "maximum embedding input length" + max_length_value = self.max_embed_len + else: + # Fall back to max_model_len validation (original behavior) + length_type = "maximum context length" + max_length_value = self.max_model_len + + validation_error_msg = ( + "This model's {length_type} is {max_length_value} tokens. " + "However, you requested {token_num} tokens in the input for " + "embedding generation. Please reduce the length of the input.") + + chunked_processing_error_msg = ( + "This model's {length_type} is {max_length_value} tokens. " + "However, you requested {token_num} tokens in the input for " + "embedding generation. Please reduce the length of the input " + "or enable chunked processing.") + + # Check if input exceeds max length + if token_num > max_length_value: + raise ValueError( + validation_error_msg.format( + length_type=length_type, + max_length_value=max_length_value, + token_num=token_num)) + + # Check for chunked processing + # when exceeding max_position_embeddings + if token_num > max_pos_embeddings: + if enable_chunked: + # Allow long inputs when chunked processing is enabled + logger.info( + "Input length %s exceeds max_position_embeddings " + "%s, will use chunked processing", token_num, + max_pos_embeddings) + else: + raise ValueError( + chunked_processing_error_msg.format( + length_type="maximum position embeddings length", + max_length_value=max_pos_embeddings, + token_num=token_num)) + + return TextTokensPrompt(prompt=input_text, + prompt_token_ids=input_ids) + + # For other request types, use the parent's implementation + return super()._validate_input(request, input_ids, input_text) + + def _is_text_tokens_prompt(self, prompt) -> bool: + """Check if a prompt is a TextTokensPrompt (has prompt_token_ids).""" + return (isinstance(prompt, dict) and "prompt_token_ids" in prompt + and "prompt_embeds" not in prompt) + + async def _create_single_prompt_generator( + self, + ctx: EmbeddingServeContext, + engine_prompt: Union[EngineTokensPrompt, EngineEmbedsPrompt], + request_prompt: RequestPrompt, + pooling_params: PoolingParams, + trace_headers: Optional[Mapping[str, str]], + prompt_index: int, + ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]: + """Create a generator for a single prompt using standard processing.""" + request_id_item = f"{ctx.request_id}-{prompt_index}" + + self._log_inputs(request_id_item, + request_prompt, + params=pooling_params, + lora_request=ctx.lora_request) + + # Mypy has an existing bug related to inferring the variance + # of TypedDicts with `builtins.enumerate`: + # https://github.com/python/mypy/issues/8586#issuecomment-2867698435 + engine_prompt = cast(Union[EngineTokensPrompt, EngineEmbedsPrompt], + engine_prompt) + + # Return the original generator without wrapping + return self.engine_client.encode( + engine_prompt, + pooling_params, + request_id_item, + lora_request=ctx.lora_request, + trace_headers=trace_headers, + priority=getattr(ctx.request, "priority", 0), + ) + + @override + async def _prepare_generators( + self, + ctx: ServeContext, + ) -> Optional[ErrorResponse]: + """Override to support chunked processing.""" + ctx = cast(EmbeddingServeContext, ctx) + + # Check if we should use chunked processing + use_chunked = self._should_use_chunked_processing(ctx.request) + + # If no chunked processing needed, delegate to parent class + if not use_chunked: + return await super()._prepare_generators(ctx) + + # Custom logic for chunked processing + generators: list[AsyncGenerator[Union[RequestOutput, + PoolingRequestOutput], + None]] = [] + + try: + trace_headers = (None if ctx.raw_request is None else await + self._get_trace_headers(ctx.raw_request.headers)) + + pooling_params = self._create_pooling_params(ctx) + if isinstance(pooling_params, ErrorResponse): + return pooling_params + + # Verify and set the task for pooling params + try: + pooling_params.verify("embed", self.model_config) + except ValueError as e: + return self.create_error_response(str(e)) + + if ctx.engine_prompts is None: + return self.create_error_response( + "Engine prompts not available") + + if ctx.request_prompts is None: + return self.create_error_response( + "Request prompts not available") + + max_pos_embeddings = self._get_max_position_embeddings() + + for i, engine_prompt in enumerate(ctx.engine_prompts): + request_prompt = ctx.request_prompts[i] + + # Check if this specific prompt needs chunked processing + if self._is_text_tokens_prompt(request_prompt): + # Cast to TextTokensPrompt since we've verified + # prompt_token_ids + text_tokens_prompt = cast(TextTokensPrompt, request_prompt) + if (len(text_tokens_prompt["prompt_token_ids"]) + > max_pos_embeddings): + # Use chunked processing for this prompt + chunk_generators = await self._process_chunked_request( + ctx, text_tokens_prompt, pooling_params, + trace_headers, i) + generators.extend(chunk_generators) + continue + + # Normal processing for short prompts or non-token prompts + # Cast engine_prompt to the expected type for mypy + engine_prompt_typed = cast( + Union[EngineTokensPrompt, EngineEmbedsPrompt], + engine_prompt) + generator = await self._create_single_prompt_generator( + ctx, engine_prompt_typed, request_prompt, pooling_params, + trace_headers, i) + generators.append(generator) + + from vllm.utils import merge_async_iterators + ctx.result_generator = merge_async_iterators(*generators) + + return None + + except Exception as e: + # TODO: Use a vllm-specific Validation Error + return self.create_error_response(str(e)) + + @override + async def _collect_batch( + self, + ctx: ServeContext, + ) -> Optional[ErrorResponse]: + """Collect and aggregate batch results + with support for chunked processing. + + For chunked requests, performs online aggregation to + minimize memory usage. + For regular requests, collects results normally. + """ + ctx = cast(EmbeddingServeContext, ctx) + try: + if ctx.engine_prompts is None: + return self.create_error_response( + "Engine prompts not available") + + # Check if we used chunked processing + use_chunked = self._should_use_chunked_processing(ctx.request) + + if not use_chunked: + return await super()._collect_batch(ctx=ctx) + + if ctx.request_prompts is None: + return self.create_error_response( + "Request prompts not available") + + if ctx.result_generator is None: + return self.create_error_response( + "Result generator not available") + + # Online aggregation for chunked requests to + # minimize memory usage + # Track aggregation state for each prompt + prompt_aggregators: dict[int, dict[str, Any]] = {} + short_prompts_results: dict[int, PoolingRequestOutput] = {} + + async for result_idx, result in ctx.result_generator: + if "-chunk-" in result.request_id: + # Extract prompt_idx from chunked request_id + parts = result.request_id.split("-") + try: + prompt_idx = int(parts[parts.index("prompt") + 1]) + except (ValueError, IndexError): + # Fallback: extract from result_idx if parsing fails + prompt_idx = result_idx + + # Initialize aggregator for this prompt if needed + if prompt_idx not in prompt_aggregators: + prompt_aggregators[prompt_idx] = { + 'weighted_sum': None, + 'total_weight': 0, + 'chunk_count': 0, + 'request_id': result.request_id.split("-chunk-")[0] + } + + aggregator = prompt_aggregators[prompt_idx] + + # MEAN pooling with online weighted averaging + # Ensure result is PoolingRequestOutput + # for embedding processing + if not isinstance(result, PoolingRequestOutput): + return self.create_error_response( + f"Expected PoolingRequestOutput for " + f"chunked embedding, got " + f"{type(result).__name__}") + + # Handle both PoolingOutput and + # EmbeddingOutput types + if hasattr(result.outputs, 'data'): + # PoolingOutput case + embedding_data = result.outputs.data + elif hasattr(result.outputs, 'embedding'): + # EmbeddingOutput case - + # convert embedding list to tensor + embedding_data = result.outputs.embedding + else: + return self.create_error_response( + f"Unsupported output type: " + f"{type(result.outputs).__name__}") + + if not isinstance(embedding_data, torch.Tensor): + embedding_data = torch.tensor(embedding_data, + dtype=torch.float32) + + if result.prompt_token_ids is None: + return self.create_error_response( + "prompt_token_ids cannot be None for " + "chunked processing") + weight = len(result.prompt_token_ids) + + weighted_embedding = embedding_data.to( + dtype=torch.float32) * weight + + if aggregator['weighted_sum'] is None: + # First chunk + aggregator['weighted_sum'] = weighted_embedding + else: + # Accumulate + aggregator['weighted_sum'] += weighted_embedding + + aggregator['total_weight'] += weight + aggregator['chunk_count'] += 1 + else: + # Non-chunked result - extract prompt_idx from request_id + parts = result.request_id.split("-") + try: + # Last part should be prompt index + prompt_idx = int(parts[-1]) + except (ValueError, IndexError): + prompt_idx = result_idx # Fallback to result_idx + + short_prompts_results[prompt_idx] = cast( + PoolingRequestOutput, result) + + # Finalize aggregated results + final_res_batch: list[Union[PoolingRequestOutput, + EmbeddingRequestOutput]] = [] + num_prompts = len(ctx.engine_prompts) + + for prompt_idx in range(num_prompts): + if prompt_idx in prompt_aggregators: + # Finalize MEAN aggregation for this chunked prompt + aggregator = prompt_aggregators[prompt_idx] + + weighted_sum = aggregator['weighted_sum'] + total_weight = aggregator['total_weight'] + + if (weighted_sum is not None + and isinstance(weighted_sum, torch.Tensor) + and isinstance(total_weight, + (int, float)) and total_weight > 0): + + # Compute final mean embedding + final_embedding = weighted_sum / total_weight + + # Create a PoolingRequestOutput + # for the aggregated result + pooling_output_data = PoolingOutput( + data=final_embedding) + + # Get original prompt token IDs for this prompt + original_prompt = ctx.request_prompts[prompt_idx] + if not self._is_text_tokens_prompt(original_prompt): + return self.create_error_response( + f"Chunked prompt {prompt_idx} is not a " + f"TextTokensPrompt") + + original_token_ids = cast( + TextTokensPrompt, + original_prompt)["prompt_token_ids"] + + pooling_request_output = PoolingRequestOutput( + request_id=aggregator['request_id'], + prompt_token_ids=original_token_ids, + outputs=pooling_output_data, + finished=True) + + final_res_batch.append(pooling_request_output) + else: + return self.create_error_response( + f"Failed to aggregate chunks " + f"for prompt {prompt_idx}") + elif prompt_idx in short_prompts_results: + final_res_batch.append( + cast(PoolingRequestOutput, + short_prompts_results[prompt_idx])) + else: + return self.create_error_response( + f"Result not found for prompt {prompt_idx}") + + ctx.final_res_batch = cast( + list[Union[RequestOutput, PoolingRequestOutput]], + final_res_batch) + + return None + + except Exception as e: + return self.create_error_response(str(e)) + class OpenAIServingEmbedding(EmbeddingMixin): request_id_prefix = "embd" From 1ddd5e734110de78a08890087c12f0ecf3ae3e9e Mon Sep 17 00:00:00 2001 From: Chi Zhang Date: Wed, 13 Aug 2025 20:27:25 +0800 Subject: [PATCH 034/233] [FEATURE] support custom vllm tuned config path for fused moe triton kernels (#22791) Signed-off-by: Chi Zhang --- vllm/envs.py | 6 ++++ .../layers/fused_moe/fused_moe.py | 28 +++++++++++++------ 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 0b016dbc85d6..2470a891c9d7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -159,6 +159,7 @@ VLLM_USE_TRTLLM_ATTENTION: Optional[str] = None VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False + VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None def get_default_cache_root(): @@ -1127,6 +1128,11 @@ def get_vllm_port() -> Optional[int]: # never removed from memory until the server terminates. "VLLM_ENABLE_RESPONSES_API_STORE": lambda: bool(int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0"))), + + # Allows vllm to find tuned config under customized folder + "VLLM_TUNED_CONFIG_FOLDER": + lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None), + } # --8<-- [end:env-vars-definition] diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index ad094c37f947..98087a35e15c 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -701,20 +701,32 @@ def get_moe_configs( block_shape = [block_n, block_k] if block_n and block_k else None json_file_name = get_config_file_name(E, N, dtype, block_shape) - config_file_path = os.path.join( + config_file_paths = [] + + # note that we prioritize user defined config + user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER + if user_defined_config_folder is not None: + user_defined_config_file_path = os.path.join( + user_defined_config_folder, json_file_name) + config_file_paths.append(user_defined_config_file_path) + + default_config_file_path = os.path.join( os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name) - if os.path.exists(config_file_path): - with open(config_file_path) as f: - logger.info("Using configuration from %s for MoE layer.", - config_file_path) - # If a configuration has been found, return it - return {int(key): val for key, val in json.load(f).items()} + config_file_paths.append(default_config_file_path) + + for config_file_path in config_file_paths: + if os.path.exists(config_file_path): + with open(config_file_path) as f: + logger.info("Using configuration from %s for MoE layer.", + config_file_path) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} # If no optimized configuration is available, we will use the default # configuration logger.warning( ("Using default MoE config. Performance might be sub-optimal! " - "Config file not found at %s"), config_file_path) + "Config file not found at %s"), config_file_paths) return None From 00f1ba7a2d65ad12d5e9e154ba7308fe9fd615a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Wed, 13 Aug 2025 15:03:53 +0200 Subject: [PATCH 035/233] [Nixl][CI] Fix tests (#22806) Signed-off-by: NickLucche --- tests/v1/kv_connector/unit/test_nixl_connector.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 3860d7c85724..b185936ab025 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -229,6 +229,9 @@ def _nixl_handshake(self, host: str, port: int, remote_tp_size: int, num_blocks=1, block_len=self.block_len, attn_backend_name=self.backend_name, + # `self.kv_cache_layout` is only forced to HND when vllm engine + # is started. We mock HND here. + kv_cache_layout="HND", ), remote_tp_size=remote_tp_size) return {0: remote_agent_name} From 657eac2840f908aae9f37aa50c93743febcec220 Mon Sep 17 00:00:00 2001 From: Chen Zhang Date: Wed, 13 Aug 2025 06:07:09 -0700 Subject: [PATCH 036/233] [Bugfix][mamba] Fix type annotation of Mamba2Metadata (#22787) Signed-off-by: Chen Zhang --- .../layers/mamba/mamba_mixer2.py | 8 ++-- vllm/v1/attention/backends/mamba_attn.py | 39 +++++++++++-------- 2 files changed, 26 insertions(+), 21 deletions(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index d5f4877135c9..10a5618c227e 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -473,12 +473,12 @@ def forward_cuda( conv_state = self_kv_cache[0].transpose(-1, -2) ssm_state = self_kv_cache[1] state_indices_tensor = attn_metadata.state_indices_tensor - has_initial_states_p = attn_metadata.has_initial_states + has_initial_states_p = attn_metadata.has_initial_states_p prep_initial_states = attn_metadata.prep_initial_states chunk_size = attn_metadata.chunk_size - seq_idx_p = attn_metadata.seq_idx - chunk_indices_p = attn_metadata.chunk_indices - chunk_offsets_p = attn_metadata.chunk_offsets + seq_idx_p = attn_metadata.seq_idx_p + chunk_indices_p = attn_metadata.chunk_indices_p + chunk_offsets_p = attn_metadata.chunk_offsets_p else: conv_state = mamba_cache_params.conv_state ssm_state = mamba_cache_params.ssm_state diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 7c1226049f69..3f84f8967db7 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -68,14 +68,19 @@ class Mamba2AttentionMetadata: query_start_loc: torch.Tensor seq_lens: torch.Tensor - has_initial_states: torch.Tensor prep_initial_states: bool chunk_size: int - seq_idx: torch.Tensor - chunk_indices: torch.Tensor - chunk_offsets: torch.Tensor + + # The following tensors only contain prefill requests and will be None if + # the batch has no prefill request. + has_initial_states_p: Optional[torch.Tensor] + seq_idx_p: Optional[torch.Tensor] + chunk_indices_p: Optional[torch.Tensor] + chunk_offsets_p: Optional[torch.Tensor] state_indices_tensor: torch.Tensor # shape: [batch,] + + # The following attributes are for triton implementation of causal_conv1d nums_dict: Optional[dict] = None cu_seqlen: Optional[int] = None batch_ptr: Optional[torch.tensor] = None @@ -115,11 +120,11 @@ def build(self, query_start_loc = common_attn_metadata.query_start_loc seq_lens = common_attn_metadata.seq_lens - seq_idx = None - chunk_indices, chunk_offsets = None, None + seq_idx_p = None + chunk_indices_p, chunk_offsets_p = None, None # Need flags to indicate if there are initial states # currently we really only support the FlashAttention backend - has_initial_states = None + has_initial_states_p = None prep_initial_states = False state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0] @@ -135,25 +140,25 @@ def build(self, common_attn_metadata. num_computed_tokens_cpu[num_reqs - num_prefills:num_reqs] > 0) prep_initial_states = torch.any(has_initial_states_cpu).item() - has_initial_states = has_initial_states_cpu.to( + has_initial_states_p = has_initial_states_cpu.to( query_start_loc.device) query_start_loc_p = common_attn_metadata.query_start_loc[ -num_prefills - 1:] - num_decode_tokens - seq_idx = torch.repeat_interleave(torch.arange( + seq_idx_p = torch.repeat_interleave(torch.arange( num_prefills, dtype=torch.int32, device=query_start_loc_p.device), - query_start_loc_p.diff(), - output_size=num_prefill_tokens) - seq_idx.unsqueeze_(0) + query_start_loc_p.diff(), + output_size=num_prefill_tokens) + seq_idx_p.unsqueeze_(0) # We compute metadata for chunked prefill once at the top level # model forward and reuse them in mamba layers. If not needed, # they will be ignored inside mamba kernels. if prep_initial_states: - chunk_indices, chunk_offsets = ( + chunk_indices_p, chunk_offsets_p = ( _query_start_loc_to_chunk_indices_offsets( query_start_loc_p, self.chunk_size, num_prefill_tokens)) @@ -173,12 +178,12 @@ def build(self, num_decode_tokens=num_decode_tokens, query_start_loc=query_start_loc, seq_lens=seq_lens, - has_initial_states=has_initial_states, prep_initial_states=prep_initial_states, chunk_size=self.chunk_size, - seq_idx=seq_idx, - chunk_indices=chunk_indices, - chunk_offsets=chunk_offsets, + has_initial_states_p=has_initial_states_p, + seq_idx_p=seq_idx_p, + chunk_indices_p=chunk_indices_p, + chunk_offsets_p=chunk_offsets_p, state_indices_tensor=state_indices_tensor, ) return attn_metadata From dea3291be24201d6b10ec7712b42adcb6bc18f42 Mon Sep 17 00:00:00 2001 From: Yuanyuan Chen Date: Wed, 13 Aug 2025 21:07:28 +0800 Subject: [PATCH 037/233] Remove unnecessary CUDA sync of qwen image and video preprocess (#22792) Signed-off-by: cyy Signed-off-by: Yuanyuan Chen Co-authored-by: Cyrus Leung --- vllm/model_executor/models/qwen2_5_vl.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 6bea180ffec9..5bcbcc4f0e37 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -976,10 +976,12 @@ def _process_image_input( image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list) # Split concatenated embeddings for each image item. + # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync merge_size = self.visual.spatial_merge_size - sizes = grid_thw.prod(-1) // merge_size // merge_size + sizes = (torch.tensor(grid_thw_list, dtype=torch.long).prod(-1) // + (merge_size * merge_size)).tolist() - return image_embeds.split(sizes.tolist()) + return image_embeds.split(sizes) def _process_video_input( self, @@ -998,9 +1000,11 @@ def _process_video_input( # Split concatenated embeddings for each video item. merge_size = self.visual.spatial_merge_size - sizes = grid_thw.prod(-1) // merge_size // merge_size + # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync + sizes = (torch.tensor(grid_thw_list, dtype=torch.long).prod(-1) // + (merge_size * merge_size)).tolist() - return video_embeds.split(sizes.tolist()) + return video_embeds.split(sizes) def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: mm_input_by_modality = {} From 0b36a38445102f73c12387aadccb96ea50945183 Mon Sep 17 00:00:00 2001 From: Gh0u1L5 Date: Wed, 13 Aug 2025 21:08:23 +0800 Subject: [PATCH 038/233] Fix GGUF loader for Qwen3 MoE. (#22785) Signed-off-by: Gh0u1L5 --- vllm/model_executor/model_loader/gguf_loader.py | 11 +++++++++++ vllm/model_executor/models/qwen3_moe.py | 1 + 2 files changed, 12 insertions(+) diff --git a/vllm/model_executor/model_loader/gguf_loader.py b/vllm/model_executor/model_loader/gguf_loader.py index 26af87c1ed67..21655b0c69bb 100644 --- a/vllm/model_executor/model_loader/gguf_loader.py +++ b/vllm/model_executor/model_loader/gguf_loader.py @@ -74,6 +74,17 @@ def _get_gguf_weights_map(self, model_config: ModelConfig): f"model.layers.{idx}.mlp.experts.0.gate_proj.weight" gguf_to_hf_name_map[f"blk.{idx}.ffn_up_exps.weight"] = \ f"model.layers.{idx}.mlp.experts.0.up_proj.weight" + if model_type in ("qwen2_moe", "qwen3_moe"): + model_type = model_type.replace("_", "") + # GGUF layer map assumes that we will have a merged expert weights + # so we need to map them manually + for idx in range(config.num_hidden_layers): + gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.weight"] = \ + f"model.layers.{idx}.mlp.experts.0.down_proj.weight" + gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_exps.weight"] = \ + f"model.layers.{idx}.mlp.experts.0.gate_proj.weight" + gguf_to_hf_name_map[f"blk.{idx}.ffn_up_exps.weight"] = \ + f"model.layers.{idx}.mlp.experts.0.up_proj.weight" arch = None for key, value in gguf.MODEL_ARCH_NAMES.items(): diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 085fc90b47b5..61b16b6a1d2d 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -375,6 +375,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, + quant_config=quant_config, prefix=f"{prefix}.embed_tokens") self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, From 4112a0974598cb5613ad943d8c12e326da429246 Mon Sep 17 00:00:00 2001 From: milesial Date: Wed, 13 Aug 2025 06:09:26 -0700 Subject: [PATCH 039/233] [Frontend] Multithreaded async multimodal load_bytes (#22710) Signed-off-by: Alexandre Milesi <30204471+milesial@users.noreply.github.com> Co-authored-by: Alexandre Milesi <30204471+milesial@users.noreply.github.com> --- vllm/envs.py | 7 +++++++ vllm/multimodal/utils.py | 26 ++++++++++++++++++++------ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 2470a891c9d7..5958a5cc0f29 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -63,6 +63,7 @@ VLLM_IMAGE_FETCH_TIMEOUT: int = 5 VLLM_VIDEO_FETCH_TIMEOUT: int = 30 VLLM_AUDIO_FETCH_TIMEOUT: int = 10 + VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MM_INPUT_CACHE_GIB: int = 4 @@ -556,6 +557,12 @@ def get_vllm_port() -> Optional[int]: "VLLM_AUDIO_FETCH_TIMEOUT": lambda: int(os.getenv("VLLM_AUDIO_FETCH_TIMEOUT", "10")), + # Max number of workers for the thread pool handling + # media bytes loading. Set to 1 to disable parallel processing. + # Default is 8 + "VLLM_MEDIA_LOADING_THREAD_COUNT": + lambda: int(os.getenv("VLLM_MEDIA_LOADING_THREAD_COUNT", "8")), + # Maximum filesize in MB for a single audio file when processing # speech-to-text requests. Files larger than this will be rejected. # Default is 25 MB diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 8dfbc6503520..b8266fd350f5 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio +import atexit +from concurrent.futures import ThreadPoolExecutor from itertools import groupby from pathlib import Path from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union @@ -33,6 +36,10 @@ MultiModalKwargs = Any MultiModalPlaceholderDict = Any +global_thread_pool = ThreadPoolExecutor( + max_workers=envs.VLLM_MEDIA_LOADING_THREAD_COUNT) +atexit.register(global_thread_pool.shutdown) + class MediaConnector: @@ -139,19 +146,26 @@ async def load_from_url_async( fetch_timeout: Optional[int] = None, ) -> _M: url_spec = urlparse(url) + loop = asyncio.get_running_loop() if url_spec.scheme.startswith("http"): connection = self.connection data = await connection.async_get_bytes(url, timeout=fetch_timeout) - - return media_io.load_bytes(data) + future = loop.run_in_executor(global_thread_pool, + media_io.load_bytes, data) + return await future if url_spec.scheme == "data": - return self._load_data_url(url_spec, media_io) + future = loop.run_in_executor(global_thread_pool, + self._load_data_url, url_spec, + media_io) + return await future if url_spec.scheme == "file": - return self._load_file_url(url_spec, media_io) - + future = loop.run_in_executor(global_thread_pool, + self._load_file_url, url_spec, + media_io) + return await future msg = "The URL must be either a HTTP, data or file URL." raise ValueError(msg) @@ -489,4 +503,4 @@ def fetch_video( "video": video_io_kwargs } media_connector = MediaConnector(media_io_kwargs=media_io_kwargs) - return media_connector.fetch_video(video_url) \ No newline at end of file + return media_connector.fetch_video(video_url) From 8dd4c5a5f081693ae8562cbbb664277d9c8eff75 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Wed, 13 Aug 2025 22:18:07 +0800 Subject: [PATCH 040/233] [Core] Use individual MM items in P0/P1 cache and model runner (#22570) Signed-off-by: DarkLight1337 --- tests/multimodal/test_utils.py | 233 +++++++------------ tests/v1/core/test_kv_cache_utils.py | 48 ++-- tests/v1/core/test_prefix_caching.py | 31 ++- tests/v1/core/test_scheduler.py | 21 +- tests/v1/core/utils.py | 19 +- tests/v1/engine/test_engine_core.py | 2 +- tests/v1/engine/test_engine_core_client.py | 2 +- tests/v1/engine/test_output_processor.py | 10 +- tests/v1/kv_connector/unit/utils.py | 2 +- tests/v1/tpu/worker/test_tpu_model_runner.py | 2 +- tests/v1/worker/test_gpu_input_batch.py | 2 +- tests/v1/worker/test_gpu_model_runner.py | 2 +- vllm/multimodal/inputs.py | 141 +++++++++-- vllm/multimodal/utils.py | 133 ++++++----- vllm/v1/core/sched/output.py | 10 +- vllm/v1/engine/__init__.py | 6 +- vllm/v1/engine/core.py | 7 +- vllm/v1/engine/mm_input_cache.py | 78 +++---- vllm/v1/engine/processor.py | 66 ++---- vllm/v1/request.py | 21 +- vllm/v1/serial_utils.py | 48 ++-- vllm/v1/worker/gpu_input_batch.py | 13 +- vllm/v1/worker/gpu_model_runner.py | 97 ++++---- vllm/v1/worker/tpu_model_runner.py | 39 ++-- 24 files changed, 548 insertions(+), 485 deletions(-) diff --git a/tests/multimodal/test_utils.py b/tests/multimodal/test_utils.py index 3fdf7e33ca5f..41f4773a11c8 100644 --- a/tests/multimodal/test_utils.py +++ b/tests/multimodal/test_utils.py @@ -5,7 +5,7 @@ import mimetypes import os from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import TYPE_CHECKING, NamedTuple, Optional +from typing import TYPE_CHECKING, NamedTuple import numpy as np import pytest @@ -19,14 +19,12 @@ initialize_model_parallel) from vllm.multimodal.image import convert_image_mode from vllm.multimodal.inputs import PlaceholderRange -from vllm.multimodal.utils import (MediaConnector, - merge_and_sort_multimodal_metadata, +from vllm.multimodal.utils import (MediaConnector, argsort_mm_positions, run_dp_sharded_vision_model) from vllm.platforms import current_platform from vllm.utils import get_open_port, update_environment_variables if TYPE_CHECKING: - from vllm.multimodal.hasher import MultiModalHashDict from vllm.multimodal.inputs import MultiModalPlaceholderDict # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) @@ -178,19 +176,17 @@ async def test_fetch_video_http(video_url: str, num_frames: int): assert metadata_sync == metadata_async -# Used for the next two tests related to `merge_and_sort_multimodal_metadata`. +# Used for `test_argsort_mm_positions`. class TestCase(NamedTuple): mm_positions: "MultiModalPlaceholderDict" - mm_hashes: Optional["MultiModalHashDict"] - expected_modalities: list[str] - expected_ranges: list[PlaceholderRange] - expected_hashes: Optional[list[str]] + expected_modality_idxs: list[tuple[str, int]] -def test_merge_and_sort_multimodal_metadata(): +def test_argsort_mm_positions(): test_cases = [ - # Single modality should return result as is but flattened + # Single modality + ## Internally sorted TestCase( mm_positions={ "image": [ @@ -198,34 +194,27 @@ def test_merge_and_sort_multimodal_metadata(): PlaceholderRange(offset=3, length=2), ] }, - mm_hashes={"image": ["hash1", "hash2"]}, - expected_modalities=["image", "image"], - expected_ranges=[ - PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=3, length=2), + expected_modality_idxs=[ + ("image", 0), + ("image", 1), ], - expected_hashes=["hash1", "hash2"], ), - - # Single modality without hashes return None for mm hash. + ## Internally unsorted TestCase( mm_positions={ "image": [ + PlaceholderRange(offset=3, length=2), PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=2, length=2), ] }, - mm_hashes=None, - expected_modalities=["image", "image"], - expected_ranges=[ - PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=2, length=2), + expected_modality_idxs=[ + ("image", 1), + ("image", 0), ], - expected_hashes=None, ), - # Multiple modalities with hashes should return sorted modalities - # and flattened ranges and hashes. + # Two modalities + ## Internally sorted TestCase( mm_positions={ "image": [ @@ -237,47 +226,54 @@ def test_merge_and_sort_multimodal_metadata(): PlaceholderRange(offset=2, length=3), ] }, - mm_hashes={ - "image": ["image_hash1", "image_hash2"], - "audio": ["audio_hash1", "audio_hash2"], - }, - expected_modalities=["audio", "audio", "image", "image"], - expected_ranges=[ - PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=2, length=3), - PlaceholderRange(offset=7, length=4), - PlaceholderRange(offset=11, length=5), + expected_modality_idxs=[ + ("audio", 0), + ("audio", 1), + ("image", 0), + ("image", 1), ], - expected_hashes=[ - "audio_hash1", "audio_hash2", "image_hash1", "image_hash2" + ), + ## Interleaved, internally sorted + TestCase( + mm_positions={ + "image": [ + PlaceholderRange(offset=0, length=4), + PlaceholderRange(offset=8, length=2), + ], + "audio": [ + PlaceholderRange(offset=5, length=2), + PlaceholderRange(offset=11, length=4), + ] + }, + expected_modality_idxs=[ + ("image", 0), + ("audio", 0), + ("image", 1), + ("audio", 1), ], ), - - # Multiple modalities without hashes should return sorted modalities - # and flattened ranges and None. + ## Interleaved, internally unsorted TestCase( mm_positions={ "image": [ - PlaceholderRange(offset=7, length=4), - PlaceholderRange(offset=11, length=5), + PlaceholderRange(offset=8, length=2), + PlaceholderRange(offset=0, length=4), ], "audio": [ - PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=2, length=3), + PlaceholderRange(offset=11, length=4), + PlaceholderRange(offset=5, length=2), ] }, - mm_hashes=None, - expected_modalities=["audio", "audio", "image", "image"], - expected_ranges=[ - PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=2, length=3), - PlaceholderRange(offset=7, length=4), - PlaceholderRange(offset=11, length=5), + expected_modality_idxs=[ + ("image", 1), + ("audio", 1), + ("image", 0), + ("audio", 0), ], - expected_hashes=None, ), # Three modalities + ## Internally sorted TestCase( mm_positions={ "image": [ @@ -293,72 +289,16 @@ def test_merge_and_sort_multimodal_metadata(): PlaceholderRange(offset=12, length=6), ] }, - mm_hashes={ - "image": ["image_hash1", "image_hash2"], - "audio": ["audio_hash1"], - "video": ["video_hash1", "video_hash2", "video_hash3"] - }, - expected_modalities=[ - "audio", "video", "video", "video", "image", "image" - ], - expected_ranges=[ - PlaceholderRange(offset=0, length=2), - PlaceholderRange(offset=3, length=4), - PlaceholderRange(offset=7, length=5), - PlaceholderRange(offset=12, length=6), - PlaceholderRange(offset=15, length=7), - PlaceholderRange(offset=22, length=8), - ], - expected_hashes=[ - "audio_hash1", "video_hash1", "video_hash2", "video_hash3", - "image_hash1", "image_hash2" - ], - ), - ] - - for (mm_positions, mm_hashes, expected_modalities, expected_ranges, - expected_hashes) in test_cases: - modalities, ranges, hashes = merge_and_sort_multimodal_metadata( - mm_positions, mm_hashes) - - assert modalities == expected_modalities - assert ranges == expected_ranges - assert hashes == expected_hashes - - -def test_merge_and_sort_multimodal_metadata_with_interleaving(): - - test_cases = [ - - #