#!/usr/bin/env python3
"""Patch the r37 DFlash draft attention (model_executor/models/qwen3_dflash.py)
so the DRAFT model's SlidingWindowSpec does not inherit the target's
sparse-MLA shared-page padding.

Why: the target's B12X sparse-MLA path (GlmMoeDsa, full GLM-5.3) writes
``cache_config.skip_page_size_padded`` during backend resolution
(platforms/interface.py). The draft model's DFlash sliding-window layers
(qwen3_dflash.py DFlashAttention.get_kv_cache_spec) then inherit that value
into their SlidingWindowSpec. The draft's bf16 page (head 64, block 256) is
larger than the target's padded MLA page, so KV layout resolution trips
    assert page_size_padded >= unpadded_page_size_bytes
(vllm/v1/kv_cache_interface.py).

The draft KV group is independent of the target's; it must use its own
(unpadded) page size — page_size_padded=None does exactly that.

Usage (in-container recipe mod, idempotent, refuses on drift):
    python3 patch_dspark_draft_page.py [--check]
"""
import os
import sys
import vllm

FILE = os.path.join(vllm.__path__[0], "model_executor/models/qwen3_dflash.py")

OLDER = (
    "page_size_padded=getattr(\n"
    '                    vllm_config.cache_config, "skip_page_size_padded", None\n'
    "                ),"
)
NEWER = (
    "page_size_padded=None,  # sm121: draft KV group must not inherit the\n"
    "                    # target sparse-MLA shared-page padding (see\n"
    "                    # patches/patch_dspark_draft_page.py header)\n"
    "                "
)


def _consistent(src, old, new):
    co, cn = src.count(old), src.count(new)
    if cn >= 1:
        return cn == 1 and co <= new.count(old)
    return co == 1


def main():
    check = "--check" in sys.argv
    src = open(FILE).read()
    if check:
        if not _consistent(src, OLDER, NEWER):
            raise SystemExit(f"CHECK FAIL: {FILE} is neither pristine nor patched")
        print(f"check OK: {FILE} state is consistent")
        return
    if src.count(OLDER) != 1:
        raise SystemExit(
            f"refusing: {FILE} old anchor matched {src.count(OLDER)}x; "
            f"re-sync patch"
        )
    out = src.replace(OLDER, NEWER)
    assert out.count(NEWER) == 1
    import py_compile

    open(FILE, "w").write(out)
    py_compile.compile(FILE, doraise=True)
    print(f"patched {FILE} (draft SlidingWindowSpec: no shared-page inheritance)")


if __name__ == "__main__":
    main()