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
[Bugfix][gpt-oss] passing the cache config in gpt-oss (#23613)
Signed-off-by: Wei Wei <wwei6@meta.com>
Wei committed
Aug 26, 2025 at 19:54 UTC
fecbb7c782980d0d9d104784a233ecb95a20ddda
1 file changed
+6
-1
vllm/model_executor/models/gpt_oss.py
+6
-1
@@ -174,12 +174,15 @@ class TransformerBlock(torch.nn.Module):
174
def __init__(
175
self,
176
config: GptOssConfig,
177
+ cache_config: CacheConfig,
178
quant_config: QuantizationConfig,
179
prefix: str = "",
180
):
181
super().__init__()
182
self.layer_idx = extract_layer_index(prefix)
182
- self.attn = OAIAttention(config, prefix=f"{prefix}.attn")
183
+ self.attn = OAIAttention(config,
184
+ prefix=f"{prefix}.attn",
185
+ cache_config=cache_config)
186
self.mlp = MLPBlock(config,
187
self.layer_idx,
188
quant_config=quant_config,
@@ -203,6 +206,7 @@ class GptOssModel(nn.Module):
206
):
207
super().__init__()
208
self.config = vllm_config.model_config.hf_config
209
+ self.cache_config = vllm_config.cache_config
210
self.quant_config = vllm_config.quant_config
211
self.parallel_config = vllm_config.parallel_config
212
self.config.hidden_size = self.config.hidden_size
@@ -213,6 +217,7 @@ class GptOssModel(nn.Module):
217
self.layers = torch.nn.ModuleList([
218
TransformerBlock(
219
self.config,
220
+ cache_config=self.cache_config,
221
quant_config=self.quant_config,
222
prefix=maybe_prefix(prefix, f"block.{layer_idx}"),
223
) for layer_idx in range(self.config.num_hidden_layers)