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 14, 05:30 UTC
[Misc] Make `LayerBlockType` a `Literal` instead of `Enum` (#27658)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 29, 2025 at 00:23 UTC
f5710ef02a230c55827c5843218801718f0b19df
3 files changed
+11
-24
vllm/config/model.py
+10
-16
@@ -41,7 +41,6 @@ from vllm.transformers_utils.config import (
41
)
42
from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
43
from vllm.transformers_utils.utils import maybe_model_redirect
44
-from vllm.utils import LayerBlockType
44
from vllm.utils.import_utils import LazyLoader
45
from vllm.utils.torch_utils import common_broadcastable_dtype
46
@@ -91,6 +90,7 @@ LogprobsMode = Literal[
90
]
91
HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig]
92
ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"]
93
+LayerBlockType = Literal["attention", "linear_attention", "mamba"]
94
95
_RUNNER_TASKS: dict[RunnerType, list[TaskOption]] = {
96
"generate": ["generate", "transcription"],
@@ -1433,11 +1433,11 @@ class ModelConfig:
1433
def get_num_layers_by_block_type(
1434
self,
1435
parallel_config: ParallelConfig,
1436
- block_type: LayerBlockType = LayerBlockType.attention,
1436
+ block_type: LayerBlockType = "attention",
1437
) -> int:
1438
# This function relies on 'layers_block_type' in hf_config,
1439
# for w/o this attribute, we will need to have workarounds like so
1440
- attn_block_type = block_type == LayerBlockType.attention
1440
+ attn_block_type = block_type == "attention"
1441
is_transformer = (
1442
not self.is_hybrid and not self.has_noops and not self.is_attention_free
1443
)
@@ -1469,9 +1469,7 @@ class ModelConfig:
1469
)
1470
else:
1471
return self.get_num_layers(parallel_config)
1472
- return sum(
1473
- t == block_type.value for t in layers_block_type_value[start:end]
1474
- )
1472
+ return sum(t == block_type for t in layers_block_type_value[start:end])
1473
1474
# Hybrid model Minimax
1475
attn_type_list = getattr(self.hf_config, "attn_type_list", None)
@@ -1481,19 +1479,16 @@ class ModelConfig:
1479
# Hybrid model Qwen3Next
1480
layer_types_value = getattr(self.hf_config, "layer_types", None)
1481
if layer_types_value is not None:
1484
- if getattr(block_type, "value", block_type) == "attention":
1482
+ if block_type == "attention":
1483
return sum(
1484
t == "full_attention" for t in layer_types_value[start:end]
1485
)
1488
- elif getattr(block_type, "value", block_type) == "linear_attention":
1486
+ elif block_type == "linear_attention":
1487
return sum(
1488
t == "linear_attention" for t in layer_types_value[start:end]
1489
)
1490
else:
1493
- return sum(
1494
- t == getattr(block_type, "value", block_type)
1495
- for t in layer_types_value[start:end]
1496
- )
1491
+ return sum(t == block_type for t in layer_types_value[start:end])
1492
1493
if (
1494
layers_block_type_value is None
@@ -1501,10 +1496,9 @@ class ModelConfig:
1496
and layer_types_value is None
1497
):
1498
raise ValueError(
1504
- "The model is an hybrid without a"
1505
- "layers_block_type or an attn_type_list, or a layer_types "
1506
- "in the hf_config, cannot determine the num of "
1507
- f"{block_type.value} layers"
1499
+ "The model is an hybrid without a layers_block_type or an "
1500
+ "attn_type_list, or a layer_types in the hf_config, "
1501
+ f"cannot determine the num of {block_type} layers"
1502
)
1503
1504
def get_mamba_chunk_size(self) -> int | None:
vllm/utils/__init__.py
-6
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import enum
4
import inspect
5
import uuid
6
import warnings
@@ -67,11 +66,6 @@ STR_INVALID_VAL: str = "INVALID"
66
T = TypeVar("T")
67
68
70
-class LayerBlockType(enum.Enum):
71
- attention = "attention"
72
- mamba = "mamba"
73
-
74
-
69
def random_uuid() -> str:
70
return str(uuid.uuid4().hex)
71
vllm/v1/worker/tpu_model_runner.py
+1
-2
@@ -53,7 +53,6 @@ from vllm.multimodal.inputs import (
53
from vllm.multimodal.utils import group_mm_kwargs_by_modality
54
from vllm.sequence import IntermediateTensors
55
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
56
-from vllm.utils import LayerBlockType
56
from vllm.utils.math_utils import cdiv, prev_power_of_2
57
from vllm.utils.platform_utils import is_pin_memory_available
58
from vllm.v1.attention.backends.pallas import (
@@ -212,7 +211,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
211
212
# Model-related.
213
self.num_attn_layers = model_config.get_num_layers_by_block_type(
215
- parallel_config, LayerBlockType.attention
214
+ parallel_config, "attention"
215
)
216
self.num_query_heads = model_config.get_num_attention_heads(parallel_config)
217
self.num_kv_heads = model_config.get_num_kv_heads(parallel_config)