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 7, 16:30 UTC
[Refactor] Remove dead code for attention benchmark script (#35418)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Feb 26, 2026 at 12:53 UTC
05970c772c1ca32be058d4cccfbb12aaf2032d70
2 files changed
-95
benchmarks/attention_benchmarks/__init__.py
-2
@@ -15,7 +15,6 @@ from .common import (
15
BenchmarkConfig,
16
BenchmarkResult,
17
MockLayer,
18
- MockModelConfig,
18
ResultsFormatter,
19
get_attention_scale,
20
is_mla_backend,
@@ -36,7 +35,6 @@ __all__ = [
35
"ResultsFormatter",
36
# Mock objects
37
"MockLayer",
39
- "MockModelConfig",
38
# Utilities
39
"setup_mla_dims",
40
"get_attention_scale",
benchmarks/attention_benchmarks/common.py
-93
@@ -10,7 +10,6 @@ from dataclasses import asdict, dataclass
10
from pathlib import Path
11
from typing import Any
12
13
-import numpy as np
13
import torch
14
from batch_spec import get_batch_type, parse_batch_spec
15
from rich.console import Console
@@ -62,10 +61,7 @@ class MockHfConfig:
61
# Import AttentionLayerBase at module level to avoid circular dependencies
62
try:
63
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
65
-
66
- _HAS_ATTENTION_LAYER_BASE = True
64
except ImportError:
68
- _HAS_ATTENTION_LAYER_BASE = False
65
AttentionLayerBase = object # Fallback
66
67
@@ -167,95 +163,6 @@ class MockLayer(AttentionLayerBase):
163
return self._kv_cache_spec
164
165
170
-class MockModelConfig:
171
- """Mock model configuration."""
172
-
173
- def __init__(
174
- self,
175
- num_q_heads: int,
176
- num_kv_heads: int,
177
- head_dim: int,
178
- dtype: torch.dtype = torch.float16,
179
- max_model_len: int = 32768,
180
- ):
181
- self._n_q = num_q_heads
182
- self._n_kv = num_kv_heads
183
- self._d = head_dim
184
- self.dtype = dtype
185
- self.max_model_len = max_model_len
186
-
187
- def get_num_attention_heads(self, _=None) -> int:
188
- return self._n_q
189
-
190
- def get_num_kv_heads(self, _=None) -> int:
191
- return self._n_kv
192
-
193
- def get_head_size(self) -> int:
194
- return self._d
195
-
196
- def get_num_layers(self) -> int:
197
- """Mock method for layer count queries."""
198
- return 1
199
-
200
- def get_sliding_window_for_layer(self, _layer_idx: int):
201
- """Mock method for sliding window queries."""
202
- return None
203
-
204
- def get_logits_soft_cap_for_layer(self, _layer_idx: int):
205
- """Mock method for logits soft cap queries."""
206
- return None
207
-
208
- def get_sm_scale_for_layer(self, _layer_idx: int) -> float:
209
- """Mock method for SM scale queries."""
210
- return 1.0 / (self.get_head_size() ** 0.5)
211
-
212
-
213
-class MockParallelConfig:
214
- """Mock parallel configuration."""
215
-
216
- pass
217
-
218
-
219
-class MockCompilationConfig:
220
- """Mock compilation configuration."""
221
-
222
- def __init__(self):
223
- self.full_cuda_graph = False
224
- self.static_forward_context = {}
225
-
226
-
227
-class MockVLLMConfig:
228
- """Mock VLLM configuration."""
229
-
230
- def __init__(self):
231
- self.compilation_config = MockCompilationConfig()
232
-
233
-
234
-class MockRunner:
235
- """Mock GPU runner for metadata builders."""
236
-
237
- def __init__(
238
- self,
239
- seq_lens: np.ndarray,
240
- query_start_locs: np.ndarray,
241
- device: torch.device,
242
- num_q_heads: int,
243
- num_kv_heads: int,
244
- head_dim: int,
245
- dtype: torch.dtype,
246
- ):
247
- self.model_config = MockModelConfig(num_q_heads, num_kv_heads, head_dim, dtype)
248
- self.parallel_config = MockParallelConfig()
249
- self.vllm_config = MockVLLMConfig()
250
- self.seq_lens_np = seq_lens
251
- self.query_start_loc_np = query_start_locs
252
- self.device = device
253
- self.attention_chunk_size = None
254
- self.num_query_heads = num_q_heads
255
- self.num_kv_heads = num_kv_heads
256
- self.dtype = dtype
257
-
258
-
166
@dataclass
167
class ParameterSweep:
168
"""Configuration for sweeping a backend parameter."""