#!/usr/bin/env python3
"""SM120 V4.1 prefill fix: split the ratio-1/2 extra KV cache to 64-slot pages.

FlashInfer's SM120 dual-cache sparse-MLA prefill only dispatches
extra_page_block_size in {64, 2} (V4's c4/c128 at page_size 256). V4.1's
ratio-2 compressed pool pages at 256 // 2 = 128, so prefill fails with
"Unsupported sparse-MLA prefill configuration ... extra_page_block_size=128".
The main SWA cache is already split 256 -> 64 by _split_kv_pages_to_64; do the
same for the extra cache, in a separate buffer so it does not clobber the
main split. Flat slot ids are unchanged by the split (page*128+off ==
(page*2)*64+off'), so extra indices need no remap.
"""
import re
import shutil
import sys

PATH = "/sgl-workspace/sglang/python/sglang/kernels/ops/attention/flash_mla_sm120.py"
src = open(PATH).read()
if "buf_tag" in src:
    print("already patched")
    sys.exit(0)

# 1) separate persistent buffers per caller
old_sig = """def _split_kv_pages_to_64(
    kv_u8: torch.Tensor,
    src_pbs: int,
    touched_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:"""
new_sig = """def _split_kv_pages_to_64(
    kv_u8: torch.Tensor,
    src_pbs: int,
    touched_indices: Optional[torch.Tensor] = None,
    buf_tag: str = "",
) -> torch.Tensor:"""
assert src.count(old_sig) == 1, "split signature not found"
src = src.replace(old_sig, new_sig)
for old, new in (
    ('key = f"flash_mla_sm120_split:{dev}"', 'key = f"flash_mla_sm120_split{buf_tag}:{dev}"'),
    ('mkey = f"flash_mla_sm120_mask:{dev}"', 'mkey = f"flash_mla_sm120_mask{buf_tag}:{dev}"'),
):
    assert src.count(old) == 1, f"buffer key not found: {old}"
    src = src.replace(old, new)

# 2) split the extra cache in the prefill wrapper when its page size is not 64 / 2
old_extra = """    extra_kv_u8 = (
        extra_k_cache.view(torch.uint8)
        if extra_k_cache is not None and extra_k_cache.dtype != torch.uint8
        else extra_k_cache
    )
    extra_idx = (
        extra_indices.squeeze(1)
        if extra_indices is not None and extra_indices.dim() == 3
        else extra_indices
    )"""
new_extra = """    extra_kv_u8 = (
        extra_k_cache.view(torch.uint8)
        if extra_k_cache is not None and extra_k_cache.dtype != torch.uint8
        else extra_k_cache
    )
    extra_idx = (
        extra_indices.squeeze(1)
        if extra_indices is not None and extra_indices.dim() == 3
        else extra_indices
    )
    if extra_k_cache is not None and extra_k_cache.ndim >= 3:
        extra_pbs = extra_k_cache.shape[1]
        # FlashInfer's dual-cache prefill only dispatches extra pbs 64 / 2.
        if extra_pbs not in (_PBS_DST, 2) and extra_pbs % _PBS_DST == 0:
            extra_kv_u8 = _split_kv_pages_to_64(
                extra_kv_u8, extra_pbs, touched_indices=extra_idx, buf_tag="_extra"
            )"""
assert src.count(old_extra) == 1, "extra block in prefill wrapper not found"
src = src.replace(old_extra, new_extra)

shutil.copy(PATH, PATH + ".orig")
open(PATH, "w").write(src)
print("patched", PATH)