# 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')