extract-kernel-definitions

name

extract-kernel-definitions

description

Generate Definition JSON files for the flashinfer-trace HuggingFace dataset by harvesting them from a short SGLang inference pass (FlashInfer's @flashinfer_api(trace=...) dumper) — or, as a fallback, by manually transcribing the schema from SGLang sources when FlashInfer doesn't yet have a trace template. Use when adding a new model, extracting GPU kernels (MLA, MoE, GQA, RMSNorm, GEMM, GDN, RoPE, sampling), or filling gaps in the dataset.

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 (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).	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.

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. Schema and full env-var docs live at docs/fi_trace.rst in the FlashInfer repo. SGLang harness reference: tests/trace/example_sglang.py.

Usage

# 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.

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 in the FlashInfer repo — adapt the model_path, tp_size, and attention_backend:

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

# 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):

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:

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

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:

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, the trace registry currently covers:

FlashInfer module	API(s)	op_type