feat(scripts): sanitize FlashInfer level-3 logs into workload JSONL#425
feat(scripts): sanitize FlashInfer level-3 logs into workload JSONL#425yyihuang wants to merge 2 commits into
Conversation
Convert FlashInfer level-3 API logs (FLASHINFER_LOGLEVEL>=3) into per-API
workload JSONL files, with optional alignment against trace-dumped
definition JSONs from @flashinfer_api(trace=...) + FLASHINFER_TRACE_DUMP=1.
Lighter-weight than the existing collect-workloads (level-10) pipeline:
no tensor dumps, no safetensors blobs — every tensor input is emitted as
{"type": "random"} with shape/dtype derived from the log. Suitable for
workload distribution capture and template auditing; not for correctness
eval.
Matching pipeline (when --trace-dir is supplied):
1. Strict — name match on kwargs.
2. Strict — unique-shape match (rank/dtype/const-axis consistent).
3. Strict — positional fallback for symmetric-input templates.
4. Lenient — name-based with relaxed rank/const checks (skipped under
--strict). Catches schema/runtime discrepancies like MLA rank-collapsed
K tensors so they still align to their definition.
Output: <definition_name>.jsonl when matched, <api_name>.jsonl when raw.
The companion .claude/skills/sanitize-fi-log/SKILL.md captures when to
use this vs collect-workloads, the matching pipeline, a triage matrix
for raw entries (template missing / rank mismatch / const aliasing /
dim-order flip), and recurring template bugs to grep for.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughA new CLI tool and several skill documents are added to support converting FlashInfer logs and tensor dumps into per-API JSONL workload entries. The CLI parses level-3+ logs, optionally loads trace definition JSONs, runs a multi-pass matching pipeline (strict → shape → positional → lenient), and writes deduplicated per-key JSONL outputs; docs describe level-3 and level-10 workflows, pairing rules, and troubleshooting. ChangesSanitize & Collection Pipeline
Sequence DiagramsequenceDiagram
participant Log as FlashInfer Log
participant Parser as Log Parser
participant Loader as Trace Loader
participant Matcher as Matching Pipeline
participant Writer as JSONL Writer
Log->>Parser: Read API Call blocks
Parser->>Parser: Parse parameters & Tensor blocks
alt With --trace-dir
Parser->>Loader: Request candidate definitions by fi_api suffix
Loader->>Loader: Load and index trace JSONs
Loader->>Matcher: Provide candidates
Matcher->>Matcher: Try strict name/shape/axes match
alt Strict fails
Matcher->>Matcher: Try greedy unique shape match
alt Shape fails
Matcher->>Matcher: Try positional-order fallback
alt Fallback fails and not --strict
Matcher->>Matcher: Apply lenient matcher (bindable axes)
end
end
end
end
Matcher->>Writer: Produce matched or raw workload entry
Writer->>Writer: Optionally dedupe & group by output key
Writer->>Writer: Write per-key .jsonl files
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 0/1 reviews remaining, refill in 60 minutes.Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the sanitize-fi-log skill and a Python script to convert FlashInfer level-3 API logs into per-API workload JSONL files. The script includes logic for parsing log blocks and aligning them with trace definitions. The review feedback highlights several opportunities to improve memory efficiency when processing large log files by adopting a streaming approach for both reading the input and writing the output. There is also a suggestion to remove a redundant helper function to simplify the code.
| def _iter_blocks(lines: list[str]) -> Iterable[tuple[str, list[str]]]: | ||
| """Yield (function_name, block_body_lines) for every API call in the log.""" | ||
| n = len(lines) | ||
| i = 0 | ||
| while i < n: | ||
| m = API_CALL_RE.match(lines[i]) | ||
| if not m: | ||
| i += 1 | ||
| continue | ||
| func_name = m.group(1).strip() | ||
| i += 1 | ||
| body: list[str] = [] | ||
| while i < n: | ||
| ln = lines[i] | ||
| if SEPARATOR_RE.match(ln): | ||
| i += 1 | ||
| break | ||
| if API_CALL_RE.match(ln): | ||
| # Inputs-only block (e.g. crash before output) — start of next call. | ||
| break | ||
| body.append(ln) | ||
| i += 1 | ||
| yield func_name, body |
There was a problem hiding this comment.
The current implementation of _iter_blocks requires the entire log file to be loaded into memory as a list of strings. For large FlashInfer logs (which can contain tens of thousands of calls), this can lead to significant memory consumption. Refactoring this function to accept an Iterable[str] (such as a file object) and using a state-machine approach would allow for streaming the log file line by line.
def _iter_blocks(lines: Iterable[str]) -> Iterable[tuple[str, list[str]]]:
"""Yield (function_name, block_body_lines) for every API call in the log."""
func_name = None
body = []
for line in lines:
m = API_CALL_RE.match(line)
if m:
if func_name:
yield func_name, body
func_name = m.group(1).strip()
body = []
elif func_name:
if SEPARATOR_RE.match(line):
yield func_name, body
func_name = None
body = []
else:
body.append(line)
if func_name:
yield func_name, body| def null_value() -> Any: | ||
| # Tiny helper so the JSON output renders ``null`` rather than ``"null"``. | ||
| return None |
There was a problem hiding this comment.
The null_value helper is redundant because json.dumps(None) correctly produces the JSON null literal. Additionally, it is used inconsistently in the script; for example, _build_matched_workload uses None directly. It is recommended to remove this helper and use None everywhere for consistency and simplicity.
| text = args.log_file.read_text() | ||
| lines = text.splitlines() | ||
|
|
||
| grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) |
There was a problem hiding this comment.
Reading the entire log file into memory with read_text() and then buffering all workload records in the grouped dictionary can be very memory-intensive for large logs. Consider streaming the log file using with open(...) as f: and writing the output records to their respective files incrementally. This can be achieved by maintaining a dictionary of open file handles for each API/definition encountered.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
.claude/skills/sanitize-fi-log/SKILL.md (1)
103-122: 💤 Low valueMinor documentation clarity on pipeline ordering.
The phrase "For each candidate the matcher tries, in order: 1,2,3,4" suggests all four passes are attempted per candidate sequentially. The actual implementation (per
_select_matchin the script) exhausts strict passes 1-3 on all candidates first, then falls back to lenient pass 4 on all candidates. The end result is the same (strict matches are preferred), but this wording could confuse readers investigating multi-candidate behavior.Consider rewording to something like:
The matcher first attempts strict alignment (passes 1-3) against each candidate. If no candidate strictly matches, it falls back to lenient alignment (pass 4) against each candidate.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.claude/skills/sanitize-fi-log/SKILL.md around lines 103 - 122, Update the documentation in SKILL.md to clarify that the matcher performs strict alignment (passes 1–3) across all candidates first and only if no strict match is found does it run the lenient pass (4) across all candidates; reference the implementation in _select_match and reword the "For each candidate the matcher tries, in order:" paragraph to something like "The matcher first attempts strict alignment (passes 1–3) against each candidate; if no candidate strictly matches, it falls back to lenient alignment (pass 4) across the candidates."scripts/sanitize_fi_log.py (2)
452-481: 💤 Low valueInconsistent
null_value()vsNoneusage.
_build_raw_workloadusesnull_value()(lines 457-458) while_build_matched_workloadusesNonedirectly (lines 478-479). Sincejson.dumps(None)already renders asnull, the helper is unnecessary. Consider usingNoneconsistently throughout.♻️ Suggested fix
Remove the helper and use
Nonedirectly:-def null_value() -> Any: - # Tiny helper so the JSON output renders ``null`` rather than ``"null"``. - return None - - def _build_raw_workload(func_name: str, log_inputs: dict[str, Any]) -> tuple[str, dict[str, Any]]: """Build a raw (unmatched) workload entry keyed by api name.""" record = { "definition": None, "workload": {"uuid": str(uuid.uuid4()), "api": func_name, "inputs": log_inputs}, - "solution": null_value(), - "evaluation": null_value(), + "solution": None, + "evaluation": None, } return func_name, record🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/sanitize_fi_log.py` around lines 452 - 481, The code mixes a custom null_value() helper with direct None usage; remove the unnecessary null_value() function and replace its uses in _build_raw_workload (and anywhere else) with None so both _build_raw_workload and _build_matched_workload consistently set "solution" and "evaluation" to None; also delete the null_value() definition to avoid dead code.
426-429: 💤 Low valueUnused loop variable
log_key.Static analysis correctly identifies that
log_keyis not used in this loop. Consider usingmapping.values()for clarity.♻️ Suggested fix
- for log_key, def_name in mapping.items(): + for def_name in mapping.values(): named_inputs[def_name] = {"type": "random"}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/sanitize_fi_log.py` around lines 426 - 429, The loop creating named_inputs uses an unused loop variable log_key; change the iteration to use mapping values directly so it's clear and avoids the unused variable (e.g., iterate over mapping.values() and assign named_inputs[def_name] = {"type": "random"}). Update the loop around named_inputs, mapping, and def_name in sanitize_fi_log.py accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In @.claude/skills/sanitize-fi-log/SKILL.md:
- Around line 103-122: Update the documentation in SKILL.md to clarify that the
matcher performs strict alignment (passes 1–3) across all candidates first and
only if no strict match is found does it run the lenient pass (4) across all
candidates; reference the implementation in _select_match and reword the "For
each candidate the matcher tries, in order:" paragraph to something like "The
matcher first attempts strict alignment (passes 1–3) against each candidate; if
no candidate strictly matches, it falls back to lenient alignment (pass 4)
across the candidates."
In `@scripts/sanitize_fi_log.py`:
- Around line 452-481: The code mixes a custom null_value() helper with direct
None usage; remove the unnecessary null_value() function and replace its uses in
_build_raw_workload (and anywhere else) with None so both _build_raw_workload
and _build_matched_workload consistently set "solution" and "evaluation" to
None; also delete the null_value() definition to avoid dead code.
- Around line 426-429: The loop creating named_inputs uses an unused loop
variable log_key; change the iteration to use mapping values directly so it's
clear and avoids the unused variable (e.g., iterate over mapping.values() and
assign named_inputs[def_name] = {"type": "random"}). Update the loop around
named_inputs, mapping, and def_name in sanitize_fi_log.py accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 2a87f07f-f725-4d40-9b19-17468f43d559
📒 Files selected for processing (2)
.claude/skills/sanitize-fi-log/SKILL.mdscripts/sanitize_fi_log.py
…el-10 into sanitize-fi-dumps The level-3 sanitization path (added in this PR via /sanitize-fi-log) is lighter, faster, and sufficient for shape-coverage workloads. Make it the documented default for /collect-workloads and /collect-workloads-bench; require an explicit opt-in for the level-10 (real-tensor) path. Changes: - New skill .claude/skills/sanitize-fi-dumps/SKILL.md: extracts the level-10 sanitization details (sanitize_dumps.py invocation, plan/run pairing, tensor storage policy, dump-dir env var preparation) out of /collect-workloads. /sanitize-fi-dumps and /sanitize-fi-log now mirror each other as the two terminal sanitization skills. - .claude/skills/collect-workloads/SKILL.md: re-frames the workflow as two modes — default level 3 + /sanitize-fi-log, opt-in level 10 + /sanitize-fi-dumps. Adds a "Two collection modes" comparison table at the top so the choice is visible. Sanitization phase delegates to /sanitize-fi-dumps. Eval/PR2 phases unchanged. - .claude/skills/collect-workloads-bench/SKILL.md: same restructuring — level-3 flow as the primary, level-10 (multi-node SLURM, model_configs) remaining as opt-in. Also adds proper YAML frontmatter (was missing). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.claude/skills/collect-workloads-bench/SKILL.md:
- Around line 123-125: The incremental-sanitize examples call sanitize_dumps.py
without limiting appended workloads, so they contradict the text that each pass
should collect 2–3 workloads; update the example command invocations of
sanitize_dumps.py (the examples around the paragraph and the other occurrence at
lines ~171-174) to include the --max-new-workloads flag set to 2 or 3 (e.g.,
--max-new-workloads 3) and briefly note the value matches the “2–3 workloads per
pass” guidance so runs don’t over-append between batch-size iterations.
In @.claude/skills/collect-workloads/SKILL.md:
- Around line 309-313: Update the PR artifact requirement to reference the new
entrypoints: replace the instruction to include stdout from
`collect_workloads.py sglang` with stdout from the current workflow (e.g.,
`collect_stream.py sglang` or the `bench_serving.py` run that drives collection)
under the `## SGLang Collection Log` section; ensure the text explicitly
requires real ShareGPT inference logs that show diverse `(batch_size,
kv_length)` pairs (and call out that uniform tiny KV caches like
`batch_size=4096` with 1-page contexts are synthetic and unacceptable). Locate
and edit the sentence(s) mentioning `collect_workloads.py` and update example
commands and expectations to match `collect_stream.py`/`bench_serving.py` so
contributors run the correct tool and include the correct stdout artifact.
In @.claude/skills/sanitize-fi-dumps/SKILL.md:
- Around line 113-116: The fenced code block that shows the two path lines
containing {flashinfer_trace_dir}/workloads/{op_type}/{def_name}.jsonl and
{flashinfer_trace_dir}/blob/workloads/{op_type}/{def_name}/{def_name}_{uuid}.safetensors
is missing a language tag; update that Markdown fence to include a language
identifier (e.g., change the opening ``` to ```text) so the block reads as a
labeled text code fence for better linting and editor rendering.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 72f618aa-8757-4eb7-b60b-a7c4eeedf351
📒 Files selected for processing (3)
.claude/skills/collect-workloads-bench/SKILL.md.claude/skills/collect-workloads/SKILL.md.claude/skills/sanitize-fi-dumps/SKILL.md
| Run one batch size at a time to avoid node overload. After each run, sanitize | ||
| to collect 2–3 workloads, then delete the dump dir before the next run. | ||
|
|
There was a problem hiding this comment.
Incremental sanitize example is missing --max-new-workloads.
The text says each pass should collect 2–3 workloads, but the example command uses sanitize_dumps.py defaults (20). That mismatch can over-append workloads per batch-size iteration.
Suggested doc patch
conda run -n flashinfer_bench python scripts/sanitize_dumps.py \
--dump-dir $DUMP_DIR \
--definitions <def_name> \
- --flashinfer-trace-dir $TRACE_DIR
+ --flashinfer-trace-dir $TRACE_DIR \
+ --max-new-workloads 4Also applies to: 171-174
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In @.claude/skills/collect-workloads-bench/SKILL.md around lines 123 - 125, The
incremental-sanitize examples call sanitize_dumps.py without limiting appended
workloads, so they contradict the text that each pass should collect 2–3
workloads; update the example command invocations of sanitize_dumps.py (the
examples around the paragraph and the other occurrence at lines ~171-174) to
include the --max-new-workloads flag set to 2 or 3 (e.g., --max-new-workloads 3)
and briefly note the value matches the “2–3 workloads per pass” guidance so runs
don’t over-append between batch-size iterations.
| **PR description must include** the full stdout of `collect_workloads.py sglang` | ||
| under `## SGLang Collection Log`. The log must show real ShareGPT inference | ||
| with diverse `(batch_size, kv_length)` pairs — uniform tiny KV caches (e.g. | ||
| `batch_size=4096` with 1-page contexts) indicate synthetic data, not real | ||
| inference. |
There was a problem hiding this comment.
PR2 log requirement references the old entrypoint.
This section asks for stdout from collect_workloads.py sglang, but the updated workflow in this same doc is centered on collect_stream.py/bench_serving.py. Please align the required artifact to avoid contributor confusion.
Suggested doc patch
-**PR description must include** the full stdout of `collect_workloads.py sglang`
+**PR description must include** the full stdout/log from the actual collection run
+(`collect_stream.py` or `bench_serving.py`, depending on the path used)
under `## SGLang Collection Log`. The log must show real ShareGPT inference📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| **PR description must include** the full stdout of `collect_workloads.py sglang` | |
| under `## SGLang Collection Log`. The log must show real ShareGPT inference | |
| with diverse `(batch_size, kv_length)` pairs — uniform tiny KV caches (e.g. | |
| `batch_size=4096` with 1-page contexts) indicate synthetic data, not real | |
| inference. | |
| **PR description must include** the full stdout/log from the actual collection run | |
| (`collect_stream.py` or `bench_serving.py`, depending on the path used) | |
| under `## SGLang Collection Log`. The log must show real ShareGPT inference | |
| with diverse `(batch_size, kv_length)` pairs — uniform tiny KV caches (e.g. | |
| `batch_size=4096` with 1-page contexts) indicate synthetic data, not real | |
| inference. |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In @.claude/skills/collect-workloads/SKILL.md around lines 309 - 313, Update the
PR artifact requirement to reference the new entrypoints: replace the
instruction to include stdout from `collect_workloads.py sglang` with stdout
from the current workflow (e.g., `collect_stream.py sglang` or the
`bench_serving.py` run that drives collection) under the `## SGLang Collection
Log` section; ensure the text explicitly requires real ShareGPT inference logs
that show diverse `(batch_size, kv_length)` pairs (and call out that uniform
tiny KV caches like `batch_size=4096` with 1-page contexts are synthetic and
unacceptable). Locate and edit the sentence(s) mentioning `collect_workloads.py`
and update example commands and expectations to match
`collect_stream.py`/`bench_serving.py` so contributors run the correct tool and
include the correct stdout artifact.
| ``` | ||
| {flashinfer_trace_dir}/workloads/{op_type}/{def_name}.jsonl | ||
| {flashinfer_trace_dir}/blob/workloads/{op_type}/{def_name}/{def_name}_{uuid}.safetensors | ||
| ``` |
There was a problem hiding this comment.
Add a language tag to the fenced output-format block.
The code fence at Line 113 is missing a language identifier; this triggers markdown lint and reduces editor rendering quality.
Suggested doc patch
-```
+```text
{flashinfer_trace_dir}/workloads/{op_type}/{def_name}.jsonl
{flashinfer_trace_dir}/blob/workloads/{op_type}/{def_name}/{def_name}_{uuid}.safetensors</details>
<details>
<summary>🧰 Tools</summary>
<details>
<summary>🪛 markdownlint-cli2 (0.22.1)</summary>
[warning] 113-113: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
</details>
</details>
<details>
<summary>🤖 Prompt for AI Agents</summary>
Verify each finding against the current code and only fix it if needed.
In @.claude/skills/sanitize-fi-dumps/SKILL.md around lines 113 - 116, The fenced
code block that shows the two path lines containing
{flashinfer_trace_dir}/workloads/{op_type}/{def_name}.jsonl and
{flashinfer_trace_dir}/blob/workloads/{op_type}/{def_name}/{def_name}_{uuid}.safetensors
is missing a language tag; update that Markdown fence to include a language
identifier (e.g., change the opening totext) so the block reads as a
labeled text code fence for better linting and editor rendering.
</details>
<!-- fingerprinting:phantom:triton:hawk:a692d414-f554-49d8-bf26-63f7d7a84560 -->
<!-- d98c2f50 -->
<!-- This is an auto-generated comment by CodeRabbit -->
3 participants
Summary
Adds
scripts/sanitize_fi_log.pyand a companion.claude/skills/sanitize-fi-log/skill for converting FlashInfer level-3 API logs (FLASHINFER_LOGLEVEL=3+) into per-API workload JSONL files.collect-workloads: no tensor dumps, no safetensors blobs — every tensor input becomes{"type": "random"}with shape/dtype derived from the log. Useful for workload distribution capture and template auditing; not for correctness eval.@flashinfer_api(trace=...)+FLASHINFER_TRACE_DUMP=1), each log call is matched to a definition — workload entries carry the definition name, named inputs, and resolved var axes. Calls with no template (or a mismatched one) fall back to a rawarg_N/kwarg_NAMEform with shape/dtype embedded inline.Matching pipeline
For each log call, candidates are looked up by suffix-match on every
fi_api:tag in--trace-dir. Then:fused_add_rmsnorm's two interchangeable[batch, hidden]tensors).--strict). Catches schema/runtime discrepancies like MLA rank-collapsed K tensors so they still align to their definition.Validated end-to-end on three real DSR1+SGLang B200 logs (~64K calls each); the skill's triage matrix documents how to read raw entries (template missing / rank mismatch /
workspace_bufferaliasing / dim-order flip) and which patterns to grep for inflashinfer/trace/templates/.Skill
.claude/skills/sanitize-fi-log/SKILL.mdcovers:collect-workloads/collect-workloads-bench.--strict.Test plan
python3 -m py_compile scripts/sanitize_fi_log.pypre-commit run --files scripts/sanitize_fi_log.py .claude/skills/sanitize-fi-log/SKILL.md— black, isort, end-of-files, etc. all pass/home/averyh/fi_log_runs/level3/good/2026-05-01T{1933,1955,2027,2113}/: produces 3–11 JSONL files per log with correct countsCompanion FlashInfer changes (separate repo)
While iterating on this script we found three FlashInfer trace templates whose schemas didn't match runtime calls (MLA k_rope rank-3 vs rank-2,
workspace_bufferaliasingnum_pages/num_q_tokens, DeepSeek V'shead_dim_vmodeled as the same const ashead_dim_qk). Fixes to those templates land in a separate flashinfer PR; this PR is self-contained and works against any trace dir.🤖 Generated with Claude Code
Summary by CodeRabbit