Spaces:
Paused
Paused
| """Write FLA kernel-choice config directories from recorded release choices. | |
| FLA's gated-delta kernels pick block sizes by timing them in each process, so | |
| two runtime processes can answer the Qwen3.5-based Decision 2.0 sizes | |
| differently by rounding. vllm-project/semantic-router records each model's | |
| released kernel choices in `src/model-runtime/vllm_sr_runtime/registry/ | |
| kernel_choices.json`. This script turns one device class of that file into | |
| FLA's own config files (one directory per model, one `<kernel>.json` per | |
| kernel), in the format FLA reads with `FLA_CONFIG_DIR=<dir>` and | |
| `FLA_CACHE_MODE=full`, so every process runs the released configurations and | |
| times none. | |
| python3 fla_kernel_choices.py kernel_choices.json /srv/decision-studio/fla-config \ | |
| --model vllm-sr/Decision-2.0-Eos-0.8B=eos --model vllm-sr/Decision-2.0-Lux-9B=lux | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import sys | |
| from pathlib import Path | |
| def fla_key_hash(key: list) -> str: | |
| """FLA's `AutotuneKey.key_hash` of a tuning key.""" | |
| serialized = json.dumps(key, separators=(",", ":"), sort_keys=True) | |
| return hashlib.md5(serialized.encode(), usedforsecurity=False).hexdigest() | |
| def write_model(choices: dict, directory: Path) -> int: | |
| directory.mkdir(parents=True, exist_ok=True) | |
| for stale in directory.glob("*.json"): | |
| stale.unlink() | |
| for name, entries in choices["kernels"].items(): | |
| document = { | |
| "kernel_name": name, | |
| "autotune_entries": { | |
| fla_key_hash(entry["key"]): {"autotune_key": entry["key"], "config": entry["config"]} | |
| for entry in entries | |
| }, | |
| "default_config": entries[0]["config"], | |
| } | |
| (directory / f"{name}.json").write_text(json.dumps(document, sort_keys=True), encoding="utf-8") | |
| return len(choices["kernels"]) | |
| def main(argv: list[str] | None = None) -> int: | |
| parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| parser.add_argument("kernel_choices") | |
| parser.add_argument("output") | |
| parser.add_argument("--device-class", default="rocm:gfx942") | |
| parser.add_argument("--fla", default="0.5.2", help="FLA version installed in the runtime image") | |
| parser.add_argument("--model", action="append", required=True, metavar="REPO=NAME") | |
| args = parser.parse_args(argv) | |
| recorded = json.loads(Path(args.kernel_choices).read_text(encoding="utf-8")) | |
| for item in args.model: | |
| repo, _, name = item.partition("=") | |
| entry = recorded.get(repo) | |
| choices = (entry or {}).get("choices", {}).get(args.device_class) | |
| if not choices or choices.get("fla") != args.fla: | |
| parser.error(f"{repo}: no {args.device_class} choices recorded with FLA {args.fla}") | |
| count = write_model(choices, Path(args.output) / name) | |
| print(f"{repo}@{entry.get('revision')} -> {name}: {count} kernels", flush=True) | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |