29 lines
1.2 KiB
Python
29 lines
1.2 KiB
Python
# Derived from blazux/qwen3.8-Flash-DGX, Copyright 2026 blazux.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""Pinned-nightly Mamba prefix alignment fix, following blazux's diagnosis.
|
|
|
|
Reference: https://github.com/blazux/qwen3.8-Flash-DGX/blob/main/src/patch_mamba_block_size.py
|
|
Do not silently apply to a different source layout.
|
|
"""
|
|
import ast
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
root = Path(sys.argv[1]) / 'vllm'
|
|
changes = [
|
|
(root / 'v1/worker/gpu/model_states/mamba_hybrid.py',
|
|
'(new_req_data.num_computed_tokens - 1) // self.cache_config.block_size',
|
|
'(new_req_data.num_computed_tokens - 1)\n'
|
|
' // (self.cache_config.mamba_block_size or self.cache_config.block_size)'),
|
|
(root / 'v1/core/sched/scheduler.py',
|
|
' block_size = self.cache_config.block_size\n # The last block-aligned',
|
|
' block_size = self.block_size\n # The last block-aligned'),
|
|
]
|
|
for path, before, after in changes:
|
|
source = path.read_text()
|
|
assert source.count(before) == 1, f'Unexpected upstream layout: {path}'
|
|
updated = source.replace(before, after)
|
|
ast.parse(updated)
|
|
path.write_text(updated)
|
|
print('Mamba prefix block alignment patched')
|