| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748 |
- #!/usr/bin/env python3
- """Point MiaAI's vendored patcher at where its targets live in this base.
- glm53-flash_SM121.py stays verbatim (MIT, licence beside it). Two of its edits
- target code that moved, so this rewrites their paths and anchors in a copy
- before it runs:
- - The indexer's two top-k allocations moved from
- model_executor/layers/sparse_attn_indexer_kpool.py to
- models/glm5next/nvidia/sparse_indexer.py, byte for byte.
- - FlashInfer 0.7.0 moved the FA2 "FP8 kv_data_type for MLA requires an SM90"
- gate from mla/_core.py into mla/_batch_mla/_backends/_fa_common.py, as a
- function of `device` at a shallower indent.
- Every retarget matches exactly once, or the build fails.
- usage: mia_retarget.py <copy of glm53-flash_SM121.py>
- """
- import sys
- from pathlib import Path
- path = Path(sys.argv[1])
- text = path.read_text()
- RETARGETS = [
- (
- "model_executor/layers/sparse_attn_indexer_kpool.py",
- "models/glm5next/nvidia/sparse_indexer.py",
- ),
- ('FI / "mla/_core.py"', 'FI / "mla/_batch_mla/_backends/_fa_common.py"'),
- (
- 'old = " major, minor = get_compute_capability(self.device)\\n'
- ' if major != 9:\\n"',
- 'old = " major, minor = get_compute_capability(device)\\n'
- ' if major != 9:\\n"',
- ),
- (
- 'new = " major, minor = get_compute_capability(self.device)\\n'
- ' if major not in (9, 12):\\n"',
- 'new = " major, minor = get_compute_capability(device)\\n'
- ' if major not in (9, 12):\\n"',
- ),
- ]
- for old, new in RETARGETS:
- n = text.count(old)
- assert n == 1, f"[mia-retarget] {old[:60]!r} matched {n}x, want 1"
- text = text.replace(old, new)
- path.write_text(text)
- print("[mia-retarget] indexer path and FlashInfer 0.7.0 fp8 gate retargeted")
|
