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, 00:30 UTC
[Bugfix] Fix 2 precommit issues - (mamba_block_size, kv_cache_config) (#27811)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Tyler Michael Smith committed
Oct 30, 2025 at 14:52 UTC
ab98f6556ff84508cdcdcd6a6b1e612a7a8819d0
2 files changed
+10
-6
vllm/model_executor/models/config.py
+1
-1
@@ -406,7 +406,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
406
# easily by changing the way we layout chunks in the
407
# mamba2 kernels.
408
409
- base_chunk_size = model_config.get_mamba_chunk_size()
409
+ base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
410
attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
411
chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
412
attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
vllm/v1/core/sched/scheduler.py
+9
-5
@@ -13,7 +13,7 @@ from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
13
from vllm.distributed.kv_transfer.kv_connector.v1 import (
14
KVConnectorBase_V1,
15
KVConnectorRole,
16
- supports_hma,
16
+ SupportsHMA,
17
)
18
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
19
from vllm.logger import init_logger
@@ -93,7 +93,11 @@ class Scheduler(SchedulerInterface):
93
)
94
95
connector_vllm_config = copy.copy(self.vllm_config)
96
- connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config)
96
+
97
+ # We're dynamically inserting a kv_cache_config variable into the
98
+ # connector_vllm_config. This is distinct from the cache_config
99
+ # that is already in there.
100
+ connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config) # type: ignore[attr-defined]
101
self.connector = KVConnectorFactory.create_connector(
102
config=connector_vllm_config, role=KVConnectorRole.SCHEDULER
103
)
@@ -1327,15 +1331,15 @@ class Scheduler(SchedulerInterface):
1331
1332
block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
1333
1330
- if not supports_hma(self.connector):
1334
+ if not isinstance(self.connector, SupportsHMA):
1335
# NOTE(Kuntai): We should deprecate this code path after we enforce
1336
# all connectors to support HMA.
1337
# Hybrid memory allocator should be already turned off for this
1338
# code path, but let's double-check here.
1339
assert len(self.kv_cache_config.kv_cache_groups) == 1
1340
return self.connector.request_finished(request, block_ids[0])
1337
- else:
1338
- return self.connector.request_finished(request, block_ids)
1341
+
1342
+ return self.connector.request_finished_all_groups(request, block_ids)
1343
1344
def _update_waiting_for_remote_kv(self, request: Request) -> bool:
1345
"""