extract-kernel-definitions

# Extract Kernel Definitions Produce per-(op, shape) Definition JSONs and stage them in the HuggingFace dataset clone at `tmp/flashinfer-trace/definitions/{op_type}/`. PR submission is **out of scope** here — see [`submit-onboarding-prs`](../submit-onboarding-prs/SKILL.md) (Phase 4 of `/onboard-model`). ## Two paths | Path | When to use | What you do | |------|-------------|-------------| | **A. Trace-dump (primary)** | Kernel is `fi_supported` per `/discover-models` — i.e. the FlashInfer API used by SGLang carries a `@flashinfer_api(trace=...)` template (see [coverage list](#flashinfer-trace-coverage)). | Run a short SGLang inference pass with `FLASHINFER_TRACE_DUMP=1`. The dumper writes one JSON per unique (op, shape) before the kernel runs (crash-safe, deduplicated). | | **B. Manual extraction (fallback)** | Kernel is `fi_missing`, **or** the relevant FlashInfer API is not yet trace-instrumented. | Read the SGLang model source + sgl-cookbook serving config + HF model config; write the Definition JSON by hand using the [schema reference](#schema-reference). | The trace-dump path is the default — it eliminates manual axis derivation and produces JSONs that already carry `axes`, `inputs`, `outputs`, `tags` (`fi_api:*`, `status:verified`), and a `reference` implementation. > Background: the trace dumper was added in > [flashinfer-ai/flashinfer#2931](https://github.com/flashinfer-ai/flashinfer/pull/2931). > Schema and full env-var docs live at > [`docs/fi_trace.rst`](https://github.com/flashinfer-ai/flashinfer/blob/main/docs/fi_trace.rst) > in the FlashInfer repo. SGLang harness reference: > [`tests/trace/example_sglang.py`](https://github.com/flashinfer-ai/flashinfer/blob/main/tests/trace/example_sglang.py). ## Usage ```bash # Path A — auto-dump every fi_supported definition for a model in one inference pass /extract-kernel-definitions --model-name llama-3.2-3b --hf-repo-id meta-llama/Llama-3.2-3B-Instruct # Path A — multi-config: one short run per (TP, EP) listed in sgl-cookbook /extract-kernel-definitions --model-name qwen3-next --tp-list 2,4 # Path B — manual fallback for fi_missing kernels (or names that didn't appear in the dump) /extract-kernel-definitions --model-name kimi-k2 --manual --op-types new_op_type ``` ## Parameters - `--model-name` (required): Model slug (e.g. `llama`, `deepseek-v3`, `qwen3-next`). Used to look up the SGLang model file and the sgl-cookbook YAML. - `--hf-repo-id` (optional): HuggingFace repo override; inferred from `--model-name` if omitted. - `--tp-list` (optional): Comma-separated TP values to run for; default reads sgl-cookbook YAML. - `--ep-list` (optional): Comma-separated EP values for MoE models. - `--manual` (optional): Force Path B (manual extraction) even for fi_supported ops. - `--op-types` (optional): Comma-separated `op_type` filter when using `--manual` or for `--dry-run` reporting. - `--dry-run` (optional): Report what would be dumped/written without running anything. - `--skip-existing` (optional, default `true`): Skip any definition whose name already exists under `tmp/flashinfer-trace/definitions/`. ## Prerequisites - `/clone-repos` has been run, so `tmp/sglang/`, `tmp/flashinfer/`, `tmp/sgl-cookbook/`, and `tmp/flashinfer-trace/` are present and current. The HF dataset clone at `tmp/flashinfer-trace/` is the only home for definitions — the in-repo `flashinfer_trace/` directory was removed in the trace-dataset refactor. - For Path A: a working CUDA-enabled environment, GPU memory sufficient for the chosen model + TP, and `attention_backend="flashinfer"` available in the installed SGLang. - For Path B: HuggingFace `config.json` access for the target model. --- ## Path A: trace-dump from a short SGLang pass The dumper fires inside FlashInfer when both env vars are set **before** the FlashInfer import. SGLang routes through `@flashinfer_api(trace=...)`-decorated APIs whenever `attention_backend="flashinfer"` is selected, so a single short prefill+decode pass exercises most ops at once. ### A1. Pick the serving config(s) Open the sgl-cookbook YAML for the target model and list the unique TP/EP values — one trace-dump pass per unique combination is enough to cover every shape variant. ```bash ls tmp/sgl-cookbook/data/models/generated/v0.5.6/ | grep -i {model_name} cat tmp/sgl-cookbook/data/models/generated/v0.5.6/{model_yaml} ``` If the model has no cookbook entry, default to TP=1 (single-GPU baseline) and skip EP. ### A2. Run the trace-dump pass Use `tools/gpu-lock` so `CUDA_VISIBLE_DEVICES` is set correctly. The script below mirrors [`tests/trace/example_sglang.py`](https://github.com/flashinfer-ai/flashinfer/blob/main/tests/trace/example_sglang.py) in the FlashInfer repo — adapt the `model_path`, `tp_size`, and `attention_backend`: ```bash DUMP_DIR=tmp/dumps/fi_trace_{model_slug}_tp{TP}_ep{EP} tools/gpu-lock --gpus {TP} --exec-timeout 1800 -- python - <<EOF import os, shutil from pathlib import Path # Must be set BEFORE flashinfer / sglang import. os.environ["FLASHINFER_TRACE_DUMP"] = "1" os.environ["FLASHINFER_TRACE_DUMP_DIR"] = "$DUMP_DIR" os.environ.setdefault("SGLANG_SKIP_CUBIN_DOWNLOAD", "1") dump = Path("$DUMP_DIR") if dump.exists(): shutil.rmtree(dump) from sglang.srt.entrypoints.engine import Engine engine = Engine( model_path="{hf_repo_id}", attention_backend="flashinfer", disable_cuda_graph=True, # keep first call on the Python path mem_fraction_static=0.5, tp_size={TP}, disable_radix_cache=True, log_level="warning", ) engine.generate( ["The capital of France is"], {"temperature": 0.0, "max_new_tokens": 4, "top_k": 50, "top_p": 0.9}, ) engine.shutdown() EOF ``` A few non-obvious requirements: - **Set the env vars before import.** `FLASHINFER_TRACE_DUMP` and `FLASHINFER_TRACE_DUMP_DIR` are read at call time, but the `@flashinfer_api` decorator binding happens at import — set them in the shell or at the top of the entry script *before* any `import flashinfer` / `import sglang` runs. - **Use `attention_backend="flashinfer"`.** Other SGLang backends bypass the FlashInfer APIs and produce no dumps. - **Disable CUDA graphs (`disable_cuda_graph=True`)** for the trace pass. Cached graphs skip the Python path and therefore the dumper. - **Page-size variants need separate runs.** SGLang's page size is fixed per server, so to capture both `_ps16` and `_ps64` shapes (for example) you must run twice with different `--page-size`. Enumerate the page sizes used by the target model. - **MoE routing methods.** Each `routing_method_type` (Default, Renormalize, DeepSeekV3, Llama4, RenormalizeNaive, TopK) emits its own template; only the routing actually exercised by the model in your prompts will dump. For DeepSeek-V3 use a real DSv3 model to capture the `ds_routing` variant. - **Quantized variants** (fp8/mxfp8/fp4 GEMM, fp8/fp4 block-scale MoE) require the model to actually use that quant config — load with the matching `--quantization` flag. ### A3. Dedupe and stage into the dataset ```bash # 1. List what was dumped ls "$DUMP_DIR" # 2. For each {name}.json: sort it under the right op_type subdirectory. # The op_type field inside the JSON is the source of truth for the subfolder. python - <<'EOF' import json, shutil from pathlib import Path src = Path("$DUMP_DIR") dst_root = Path("tmp/flashinfer-trace/definitions") for p in src.glob("*.json"): op_type = json.loads(p.read_text())["op_type"] dst = dst_root / op_type / p.name if dst.exists(): print(f"skip (exists): {dst.relative_to(dst_root)}") continue dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(p, dst) print(f"added: {dst.relative_to(dst_root)}") EOF ``` ### A3b. Normalize the staged JSONs for `flashinfer-bench validate` The trace dumper's output has three known mismatches with `flashinfer-bench`'s `Definition` schema (all originate in the trace templates shipped with [flashinfer-ai/flashinfer#2931](https://github.com/flashinfer-ai/flashinfer/pull/2931)): 1. `reference` declares the function as `def _<name>_reference(...)`, but `flashinfer-bench` requires a top-level `def run(...)`. 2. Plan-time index tensors (`kv_indptr`, `kv_indices`, `qo_indptr`) come back with `dtype: "unknown"` because the dumper inspects only `run()`'s kwargs, not the wrapper state set during `plan()`. The validator only accepts concrete dtypes from its enum. 3. In-place ops (e.g. `fused_add_rmsnorm`'s residual) declare the same name in both `inputs` and `outputs`. `flashinfer-bench` rejects overlapping I/O names; the dumper's `reference` function only returns the non-overlap outputs anyway, so it's safe to drop the duplicates from `outputs`. Run this once per staging pass to make the JSONs validate: ```bash python - <<'EOF' import json, re from pathlib import Path INDEX_TENSOR_DTYPE = "int32" KNOWN_INDEX_TENSORS = { "kv_indptr", "kv_indices", "qo_indptr", "paged_kv_indptr", "paged_kv_indices", "kv_last_page_len", } REF_RE = re.compile(r"^def\s+_[A-Za-z0-9_]+_reference\b", re.MULTILINE) for p in Path("tmp/flashinfer-trace/definitions").rglob("*.json"): d = json.loads(p.read_text()); changed = False ref = d.get("reference", "") if ref and "def run(" not in ref: new_ref, n = REF_RE.subn("def run", ref, count=1) if n == 1: d["reference"], changed = new_ref, True for name, spec in d.get("inputs", {}).items(): if isinstance(spec, dict) and spec.get("dtype") == "unknown" and name in KNOWN_INDEX_TENSORS: spec["dtype"], changed = INDEX_TENSOR_DTYPE, True overlap = set(d.get("inputs", {})) & set(d.get("outputs", {})) for name in overlap: d["outputs"].pop(name, None) changed = True if changed: p.write_text(json.dumps(d, indent=2) + "\n") print(f"normalized: {p}") EOF ``` These three patches are mechanical — file a follow-up issue against `flashinfer-ai/flashinfer` to emit `def run(...)`, resolve plan-time dtypes, and drop in-place outputs (or rename them) inside the dumper itself, after which A3b becomes a no-op. ### A3c. Validate ```bash flashinfer-bench validate --dataset tmp/flashinfer-trace --disable-gpu ``` Newly staged definitions should report `[WARNING]` (missing descriptions on axes/inputs/outputs are advisory) and not `[ERROR]`. Any `[ERROR]` on a definition you just staged means A3b didn't normalize a new failure mode — inspect the report (`tmp/flashinfer-trace/reports/report-*.txt`) and extend the snippet. That's it for Path A — once normalized, the staged JSONs carry `axes`, `inputs`, `outputs`, `tags` (`fi_api:*`, `status:verified`), and a `reference` implementation, so they're ready for the rest of the onboarding pipeline (workloads → baseline → eval → Phase 4 PRs). ### Trade-off vs. tag enrichment The dumper does **not** auto-emit `tp:N`, `ep:N`, `model:*`, or `quantization:*` tags — those are workflow-level metadata, not kernel-shape metadata. After staging, append the appropriate tags to the JSONs you just produced: ```bash python - <<'EOF' import json from pathlib import Path extra_tags = ["model:{model_slug}", "tp:{TP}"] # add ep:{EP} for MoE for p in Path("tmp/flashinfer-trace/definitions").rglob("*.json"): if p.stat().st_mtime < {dump_run_start_epoch}: continue j = json.loads(p.read_text()) j["tags"] = sorted(set(j.get("tags", []) + extra_tags)) p.write_text(json.dumps(j, indent=2) + "\n") EOF ``` ### FlashInfer trace coverage Per [`docs/fi_trace.rst`](https://github.com/flashinfer-ai/flashinfer/blob/main/docs/fi_trace.rst), the trace registry currently covers: | FlashInfer module | API(s) | `op_type` | |-------------------|--------|-----------|