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 10, 10:15 UTC
[UX] Add vLLM model inspection view (#29450)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 9, 2026 at 12:12 UTC
d5ec6c056f2f806ea119b9129b07994f1d4fef4e
6 files changed
+180
-1
vllm/entrypoints/llm.py
+16
@@ -348,6 +348,9 @@ class LLM:
348
self.input_processor = self.llm_engine.input_processor
349
self.io_processor = self.llm_engine.io_processor
350
351
+ # Cache for __repr__ to avoid repeated collective_rpc calls
352
+ self._cached_repr: str | None = None
353
+
354
def get_tokenizer(self) -> TokenizerLike:
355
return self.llm_engine.get_tokenizer()
356
@@ -1786,3 +1789,16 @@ class LLM:
1789
# This is necessary because some requests may be finished earlier than
1790
# its previous requests.
1791
return sorted(outputs, key=lambda x: int(x.request_id))
1792
+
1793
+ def __repr__(self) -> str:
1794
+ """Return a transformers-style hierarchical view of the model."""
1795
+ # Cache the result to avoid repeated collective_rpc calls
1796
+ if self._cached_repr is None:
1797
+ results = self.llm_engine.collective_rpc("get_model_inspection")
1798
+ # In distributed settings, we get results from all workers
1799
+ # Just return the first one (they should all be the same)
1800
+ if results:
1801
+ self._cached_repr = results[0]
1802
+ else:
1803
+ self._cached_repr = f"LLM(model={self.model_config.model!r})"
1804
+ return self._cached_repr
vllm/envs.py
+7
@@ -250,6 +250,7 @@ if TYPE_CHECKING:
250
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
251
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
252
VLLM_USE_V2_MODEL_RUNNER: bool = False
253
+ VLLM_LOG_MODEL_INSPECTION: bool = False
254
VLLM_DEBUG_MFU_METRICS: bool = False
255
256
@@ -1595,6 +1596,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1596
"VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
1597
int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
1598
),
1599
+ # Log model inspection after loading.
1600
+ # If enabled, logs a transformers-style hierarchical view of the model
1601
+ # with quantization methods and attention backends.
1602
+ "VLLM_LOG_MODEL_INSPECTION": lambda: bool(
1603
+ int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0"))
1604
+ ),
1605
# Debug logging for --enable-mfu-metrics
1606
"VLLM_DEBUG_MFU_METRICS": lambda: bool(
1607
int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))
vllm/model_executor/layers/rotary_embedding/common.py
+1
-1
@@ -285,5 +285,5 @@ class ApplyRotaryEmb(CustomOp):
285
286
def extra_repr(self) -> str:
287
s = f"is_neox_style={self.is_neox_style}"
288
- s += f"enable_fp32_compute={self.enable_fp32_compute}"
288
+ s += f", enable_fp32_compute={self.enable_fp32_compute}"
289
return s
vllm/model_executor/model_loader/base_loader.py
+14
@@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
5
import torch
6
import torch.nn as nn
7
8
+import vllm.envs as envs
9
from vllm.config import ModelConfig, VllmConfig
10
from vllm.config.load import LoadConfig
11
from vllm.logger import init_logger
@@ -50,8 +51,21 @@ class BaseModelLoader(ABC):
51
vllm_config=vllm_config, model_config=model_config
52
)
53
54
+ log_model_inspection(model)
55
+
56
logger.debug("Loading weights on %s ...", load_device)
57
# Quantization does not happen in `load_weights` but after it
58
self.load_weights(model, model_config)
59
process_weights_after_loading(model, model_config, target_device)
60
+
61
return model.eval()
62
+
63
+
64
+def log_model_inspection(model: nn.Module) -> None:
65
+ """Log model structure if VLLM_LOG_MODEL_INSPECTION=1."""
66
+ if not envs.VLLM_LOG_MODEL_INSPECTION:
67
+ return
68
+
69
+ from vllm.model_inspection import format_model_inspection
70
+
71
+ logger.info("vLLM model structure:\n%s", format_model_inspection(model))
vllm/model_inspection.py
new
+136
@@ -0,0 +1,136 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Model inspection utilities for vLLM."""
4
+
5
+import torch.nn as nn
6
+
7
+
8
+def _get_module_info(module: nn.Module) -> str:
9
+ """Get info string for a module."""
10
+ class_name = type(module).__name__
11
+ parts = []
12
+
13
+ # Add quant_method if present
14
+ quant_method = getattr(module, "quant_method", None)
15
+ if quant_method is not None:
16
+ quant_name = type(quant_method).__name__
17
+ # For CompressedTensors, show the underlying scheme instead
18
+ scheme = getattr(module, "scheme", None)
19
+ if scheme is not None:
20
+ quant_name = type(scheme).__name__
21
+ # Skip unquantized methods
22
+ if "Unquantized" not in quant_name:
23
+ parts.append(f"quant={quant_name}")
24
+
25
+ # If module has extra_repr, use it
26
+ if hasattr(module, "extra_repr"):
27
+ parts.append(module.extra_repr().replace("\n", ""))
28
+
29
+ if parts:
30
+ return f"{class_name}({', '.join(parts)})"
31
+
32
+ # For unknown modules, use the default PyTorch repr
33
+ return str(module)
34
+
35
+
36
+def _get_child_signature(child: nn.Module) -> str:
37
+ """Get a signature for a child module to detect duplicates."""
38
+ lines = []
39
+ for name, submodule in child.named_modules():
40
+ lines.append(f"{name}:{_get_module_info(submodule)}")
41
+ return "\n".join(lines)
42
+
43
+
44
+def _format_index_ranges(indices: list[int]) -> str:
45
+ """Format indices into range notation (e.g., [0,1,2,4,5,6] -> '0-2, 4-6')."""
46
+ indices = sorted(indices)
47
+ ranges = []
48
+ start = end = indices[0]
49
+
50
+ for idx in indices[1:]:
51
+ if idx == end + 1:
52
+ end = idx
53
+ else:
54
+ ranges.append(str(start) if start == end else f"{start}-{end}")
55
+ start = end = idx
56
+
57
+ ranges.append(str(start) if start == end else f"{start}-{end}")
58
+ return ", ".join(ranges)
59
+
60
+
61
+def _format_module_tree(
62
+ module: nn.Module,
63
+ name: str = "",
64
+ indent: int = 0,
65
+) -> list[str]:
66
+ """Format a module tree with indentation, grouping identical layers.
67
+
68
+ Produces output like:
69
+ (layers): ModuleList(
70
+ (0-27, 29-47): 47 x LlamaDecoderLayer(
71
+ ...
72
+ )
73
+ (28, 48): 2 x DifferentDecoderLayer(
74
+ ...
75
+ )
76
+ )
77
+ """
78
+ lines = []
79
+ prefix = " " * indent
80
+ children = list(module.named_children())
81
+
82
+ # Leaf node - just output the module info
83
+ if not children:
84
+ info = _get_module_info(module)
85
+ lines.append(f"{prefix}({name}): {info}" if name else f"{prefix}{info}")
86
+ return lines
87
+
88
+ # Non-leaf node - output opening line and recurse into children
89
+ info = _get_module_info(module)
90
+ lines.append(f"{prefix}({name}): {info}(" if name else f"{prefix}{info}(")
91
+
92
+ # Separate numbered children (e.g., "0", "1") from named ones (e.g., "norm")
93
+ numbered: list[tuple[int, nn.Module]] = []
94
+ non_numbered: list[tuple[str, nn.Module]] = []
95
+ for child_name, child_module in children:
96
+ try:
97
+ numbered.append((int(child_name), child_module))
98
+ except ValueError:
99
+ non_numbered.append((child_name, child_module))
100
+
101
+ # Group numbered children by structure signature to collapse identical layers
102
+ # e.g., layers 0-27 and 29-47 with same structure become "(0-27, 29-47): 47 x"
103
+ if numbered:
104
+ sig_to_group: dict[str, list[tuple[int, nn.Module]]] = {}
105
+ for idx, child_module in numbered:
106
+ sig = _get_child_signature(child_module)
107
+ sig_to_group.setdefault(sig, []).append((idx, child_module))
108
+
109
+ # Output groups sorted by first index
110
+ for group in sorted(sig_to_group.values(), key=lambda g: g[0][0]):
111
+ indices = [idx for idx, _ in group]
112
+ representative = group[0][1]
113
+ child_lines = _format_module_tree(representative, "", indent + 1)
114
+ first_line = child_lines[0].lstrip()
115
+ child_prefix = " " * (indent + 1)
116
+
117
+ if len(indices) > 1:
118
+ range_str = _format_index_ranges(indices)
119
+ child_lines[0] = (
120
+ f"{child_prefix}({range_str}): {len(indices)} x {first_line}"
121
+ )
122
+ else:
123
+ child_lines[0] = f"{child_prefix}({indices[0]}): {first_line}"
124
+ lines.extend(child_lines)
125
+
126
+ # Output non-numbered children (e.g., "embed_tokens", "norm")
127
+ for child_name, child_module in non_numbered:
128
+ lines.extend(_format_module_tree(child_module, child_name, indent + 1))
129
+
130
+ lines.append(f"{prefix})")
131
+ return lines
132
+
133
+
134
+def format_model_inspection(model: nn.Module) -> str:
135
+ """Format a model into a transformers-style hierarchical string."""
136
+ return "\n".join(_format_module_tree(model))
vllm/v1/worker/worker_base.py
+6
@@ -118,6 +118,12 @@ class WorkerBase:
118
"""Apply a function on the model inside this worker."""
119
return fn(self.get_model())
120
121
+ def get_model_inspection(self) -> str:
122
+ """Return a transformers-style hierarchical view of the model."""
123
+ from vllm.model_inspection import format_model_inspection
124
+
125
+ return format_model_inspection(self.get_model())
126
+
127
def load_model(self) -> None:
128
"""Load model onto target device."""
129
raise NotImplementedError