A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[Bugfix] Fix use_cascade_attention handling for Alibi-based models on vllm/v1 (#15211)
Signed-off-by: h-sugi <h.sugi@ieee.org> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Hiroaki Sugiyama committed
Mar 27, 2025 at 23:29 UTC
8958217ad5a6830c4d911e5f15e6eb791df337b6
2 files changed
+18
-3
vllm/utils.py
+13
-1
@@ -61,7 +61,7 @@ import vllm.envs as envs
61
from vllm.logger import enable_trace_function_call, init_logger
62
63
if TYPE_CHECKING:
64
- from vllm.config import VllmConfig
64
+ from vllm.config import ModelConfig, VllmConfig
65
66
logger = init_logger(__name__)
67
@@ -2498,6 +2498,18 @@ def cprofile(save_file: Optional[str] = None, enabled: bool = True):
2498
return decorator
2499
2500
2501
+# Only relevant for models using ALiBi (e.g, MPT)
2502
+def check_use_alibi(model_config: ModelConfig) -> bool:
2503
+ return (getattr(model_config.hf_text_config, "alibi", False) # Falcon
2504
+ or ("BloomForCausalLM" in getattr(model_config.hf_config,
2505
+ "architectures", [])) # Bloom
2506
+ or getattr(model_config.hf_text_config, "position_encoding_type",
2507
+ "") == "alibi" # codellm_1b_alibi
2508
+ or
2509
+ (hasattr(model_config.hf_text_config, "attn_config") # MPT
2510
+ and model_config.hf_text_config.attn_config.get("alibi", False)))
2511
+
2512
+
2513
def sha256(input) -> int:
2514
"""Hash any picklable Python object using SHA-256.
2515
vllm/v1/worker/gpu_model_runner.py
+5
-2
@@ -25,7 +25,7 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
25
from vllm.sampling_params import SamplingType
26
from vllm.sequence import IntermediateTensors
27
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
28
- LayerBlockType, LazyLoader, cdiv,
28
+ LayerBlockType, LazyLoader, cdiv, check_use_alibi,
29
is_pin_memory_available)
30
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
31
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -223,6 +223,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
223
device="cpu",
224
pin_memory=self.pin_memory)
225
226
+ # Only relevant for models using ALiBi (e.g, MPT)
227
+ self.use_alibi = check_use_alibi(model_config)
228
+
229
self.inputs_embeds = torch.zeros(
230
(self.max_num_tokens, self.hidden_size),
231
dtype=self.dtype,
@@ -689,7 +692,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
692
query_lens=num_scheduled_tokens,
693
num_query_heads=self.num_query_heads,
694
num_kv_heads=self.num_kv_heads,
692
- use_alibi=False, # FIXME
695
+ use_alibi=self.use_alibi,
696
use_sliding_window=self.window_size is not None,
697
num_sms=self.num_sms,
698
)