A high-throughput and memory-efficient inference and serving engine for LLMs
Added extra_repr() to pooler classes to improve debuggability (#44805)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jun 7, 2026 at 23:19 UTC
5633405964b1b34fb4b227c65745839f42c711cb
7 files changed
+73
vllm/model_executor/layers/pooler/activations.py
+9
@@ -128,6 +128,9 @@ class PoolerClassify(PoolerActivation):
128
129
self.num_labels = num_labels
130
131
+ def extra_repr(self) -> str:
132
+ return f"num_labels={self.num_labels}"
133
+
134
def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
135
num_labels = self.num_labels
136
if num_labels is None:
@@ -145,5 +148,11 @@ class LambdaPoolerActivation(PoolerActivation):
148
149
self.fn = fn
150
151
+ def extra_repr(self) -> str:
152
+ name = getattr(self.fn, "__name__", None)
153
+ if name is None:
154
+ name = self.fn.__class__.__name__
155
+ return f"fn={name}"
156
+
157
def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
158
return self.fn(pooled_data)
vllm/model_executor/layers/pooler/seqwise/heads.py
+24
@@ -43,6 +43,16 @@ class EmbeddingPoolerHead(SequencePoolerHead):
43
self.head_dtype = head_dtype
44
self.activation = activation
45
46
+ def extra_repr(self) -> str:
47
+ attrs = []
48
+ if self.head_dtype is not None:
49
+ attrs.append(f"head_dtype={self.head_dtype}")
50
+ if self.projector is not None:
51
+ attrs.append("projector=True")
52
+ if self.activation is not None:
53
+ attrs.append(f"activation={self.activation.__class__.__name__}")
54
+ return ", ".join(attrs)
55
+
56
def get_supported_tasks(self) -> Set[PoolingTask]:
57
return {"embed"}
58
@@ -124,6 +134,20 @@ class ClassifierPoolerHead(SequencePoolerHead):
134
self.head_dtype = head_dtype
135
self.activation = activation
136
137
+ def extra_repr(self) -> str:
138
+ attrs = []
139
+ if self.head_dtype is not None:
140
+ attrs.append(f"head_dtype={self.head_dtype}")
141
+ if self.classifier is not None:
142
+ attrs.append("classifier=True")
143
+ if self.logit_mean is not None:
144
+ attrs.append(f"logit_mean={self.logit_mean}")
145
+ if self.logit_sigma is not None:
146
+ attrs.append(f"logit_sigma={self.logit_sigma}")
147
+ if self.activation is not None:
148
+ attrs.append(f"activation={self.activation.__class__.__name__}")
149
+ return ", ".join(attrs)
150
+
151
def get_supported_tasks(self) -> Set[PoolingTask]:
152
return {"classify"}
153
vllm/model_executor/layers/pooler/seqwise/poolers.py
+6
@@ -61,6 +61,12 @@ class SequencePooler(Pooler):
61
self.pooling = pooling
62
self.head = head
63
64
+ def extra_repr(self) -> str:
65
+ return (
66
+ f"pooling={self.pooling.__class__.__name__}, "
67
+ f"head={self.head.__class__.__name__}"
68
+ )
69
+
70
def get_supported_tasks(self) -> Set[PoolingTask]:
71
tasks = set(POOLING_TASKS)
72
vllm/model_executor/layers/pooler/special.py
+3
@@ -164,6 +164,9 @@ class BOSEOSFilter(Pooler):
164
self.bos_token_id = bos_token_id
165
self.eos_token_id = eos_token_id
166
167
+ def extra_repr(self) -> str:
168
+ return f"bos_token_id={self.bos_token_id}, eos_token_id={self.eos_token_id}"
169
+
170
def get_supported_tasks(self) -> Set[PoolingTask]:
171
return self.pooler.get_supported_tasks()
172
vllm/model_executor/layers/pooler/tokwise/heads.py
+24
@@ -58,6 +58,16 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
58
self.projector = projector
59
self.activation = activation
60
61
+ def extra_repr(self) -> str:
62
+ attrs = []
63
+ if self.head_dtype is not None:
64
+ attrs.append(f"head_dtype={self.head_dtype}")
65
+ if self.projector is not None:
66
+ attrs.append("projector=True")
67
+ if self.activation is not None:
68
+ attrs.append(f"activation={self.activation.__class__.__name__}")
69
+ return ", ".join(attrs)
70
+
71
def get_supported_tasks(self) -> Set[PoolingTask]:
72
return {"token_embed"}
73
@@ -110,6 +120,20 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
120
self.head_dtype = head_dtype
121
self.activation = activation
122
123
+ def extra_repr(self) -> str:
124
+ attrs = []
125
+ if self.head_dtype is not None:
126
+ attrs.append(f"head_dtype={self.head_dtype}")
127
+ if self.classifier is not None:
128
+ attrs.append("classifier=True")
129
+ if self.logit_mean is not None:
130
+ attrs.append(f"logit_mean={self.logit_mean}")
131
+ if self.logit_sigma is not None:
132
+ attrs.append(f"logit_sigma={self.logit_sigma}")
133
+ if self.activation is not None:
134
+ attrs.append(f"activation={self.activation.__class__.__name__}")
135
+ return ", ".join(attrs)
136
+
137
def get_supported_tasks(self) -> Set[PoolingTask]:
138
return {"token_classify"}
139
vllm/model_executor/layers/pooler/tokwise/methods.py
+3
@@ -41,6 +41,9 @@ class AllPool(TokenPoolingMethod):
41
42
self.enable_chunked_prefill = scheduler_config.enable_chunked_prefill
43
44
+ def extra_repr(self) -> str:
45
+ return f"enable_chunked_prefill={self.enable_chunked_prefill}"
46
+
47
def forward(
48
self,
49
hidden_states: torch.Tensor,
vllm/model_executor/layers/pooler/tokwise/poolers.py
+4
@@ -65,6 +65,10 @@ class TokenPooler(Pooler):
65
self.pooling = pooling
66
self.head = head
67
68
+ def extra_repr(self) -> str:
69
+ head_name = self.head.__class__.__name__ if self.head is not None else None
70
+ return f"pooling={self.pooling.__class__.__name__}, head={head_name}"
71
+
72
def get_supported_tasks(self) -> Set[PoolingTask]:
73
tasks = set(POOLING_TASKS)
74