#!/usr/bin/env python3
"""Patch SGLang's Qwen sparse-attention (QSA) varlen-decode path for SM121
(DGX Spark / GB10).

The stock `lmsysorg/sglang:qwen38flashnext` resolves QSA's packed varlen decode
to flash-attn-4's CuTe-DSL kernels, which fail MLIR compilation on SM121
(layout-congruence error). This injects an early-return that routes non-SM100
devices to the Triton FlashDecoding-style fallback in `qsa_fa_fallback.py`
(CUDA-graph-safe: cu_seqlens are read on-device, so graph replay works).

Anchor-based, refuses on drift, idempotent (--check passes on pristine and on
already-patched source). See the vllm-mod-patches skill.
"""
import os
import sys

import sglang

FILE = os.path.join(
    sglang.__path__[0], "srt/layers/attention/qwen_sparse_attn_backend.py"
)
# The dense-varlen resolver's first try block (FA2 preferred). Insert the
# SM121 early-return directly before it.
OLDER = "    try:\n        from flash_attn import flash_attn_varlen_func"
NEWER = (
    "    from sglang.srt.utils import is_sm100_supported\n"
    "    if not is_sm100_supported():\n"
    "        from sglang.srt.layers.attention.qsa_fa_fallback import triton_varlen_attn_func\n"
    "        return triton_varlen_attn_func\n"
    + OLDER
)


def _consistent(src: str, old: str, new: str) -> bool:
    co, cn = src.count(old), src.count(new)
    if cn >= 1:  # patched: new once, and any old sits only inside new
        return cn == 1 and co <= new.count(old)
    return co == 1  # pristine: old once


def main() -> None:
    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} anchor matched {src.count(OLDER)}x; re-sync patch"
        )
    out = src.replace(OLDER, NEWER)
    assert out.count(NEWER) == 1
    open(FILE, "w").write(out)
    print(f"patched {FILE} for SM121")


if __name__ == "__main__":
    main()