feat(tune): record untuned shapes for every GEMM family, not just bf16 - #5268
feat(tune): record untuned shapes for every GEMM family, not just bf16#5268ThomasNing wants to merge 7 commits into
Conversation
AITER_TUNE_GEMM=1 makes tuned_gemm.py append each missed bf16 shape to bf16_untuned_gemm.csv, which is why tuning a real bf16 deployment is self-service: run the server, collect the CSV, hand it straight back to the tuner. The fp8 per-token, block-scale and mxfp4 families log their misses to INFO and write nothing, so their shape lists have to be scraped out of server logs by hand -- which is also how shape-key mistakes creep in. We shipped a tuned table that never matched because the engine fuses two projections and dispatches N=2688 while the tuned rows target the unfused N=2624; a recorded shape list would have made that impossible to miss. Adds aiter/utility/untuned_shapes.py: record(tuned_file, row) appends one de-duplicated row in the schema that family's tuner consumes. The destination name is derived from the tuned table's own name (*_tuned_* -> *_untuned_*), so a family never names its untuned file twice and a merged /tmp/aiter_configs copy still resolves to the right one. Wired into the miss paths of get_CKGEMM_config, get_GEMM_config_with_quant_type (a8w8 / a8w8_bpreshuffle / blockscale) and gemm_op_a4w4.get_GEMM_config. Same switch as before (AITER_TUNE_GEMM=1) so there is nothing new to learn; AITER_TUNE_GEMM_DIR overrides the destination for read-only package dirs and containers. Off by default. A set lookup on the hot path once a shape has been seen, and it never raises -- telemetry must not take down inference. Checked: name derivation incl. the merged-copy path, per-family schemas, de-duplication across repeats, disabled-by-default, and graceful degradation on an unwritable destination. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5302edfe19
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3dc6d068e0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1b1688c5da
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f906b11165
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 18285a371c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Review is clean and all threads are resolved on |
|
@yzhou103 Could you take a look? |
|
consider process-local de-duplication followed by append-only writes. The tuners already call drop_duplicates(), so re-reading the entire CSV under an interprocess lock for every new shape is unnecessary, scales roughly quadratically, and may introduce blocking file I/O in CUDA/HIP Graph capture paths. Per-worker output shards followed by an offline merge would avoid the interprocess lock entirely. |
|
maybe it is better to save the untuned csv like: |
|
@yzhou103 Addressed both comments in 6c648e4:
The header is validated only once per process and is published atomically for concurrent workers. Added coverage for local de-duplication, single-write append behavior, concurrent complete rows, retry behavior, and model-directory routing. |
2 participants
Fixes #5267.
Why
AITER_TUNE_GEMM=1makestuned_gemm.pyappend every missed bf16 shape tobf16_untuned_gemm.csv. That single feature is why tuning a real bf16 deployment is self-service: run the server, collect the CSV, hand it straight back to the tuner.The fp8 per-token, block-scale and mxfp4 families have no equivalent — they log misses at INFO and write nothing — so their shape lists have to be scraped out of server logs by hand. That is also how shape-key mistakes creep in: we shipped a tuned table that never matched anything because the engine fuses two projections and dispatches
N=2688, while the tuned rows targeted the unfusedN=2624. A recorded shape list makes that impossible to miss.What
aiter/utility/untuned_shapes.py—record(tuned_file, row)appends one de-duplicated row in the schema that family's tuner consumes. The destination name is derived from the tuned table's own name (*_tuned_*→*_untuned_*), so a family never has to name its untuned file twice, and a merged/tmp/aiter_configscopy still resolves to the right destination.Wired into the miss paths of:
gemm_op_a8w8.get_CKGEMM_config(block-scale, ± preshuffle) →M,N,Kgemm_op_a8w8.get_GEMM_config_with_quant_type(a8w8, a8w8_bpreshuffle) →M,N,K,q_dtype_wgemm_op_a4w4.get_GEMM_config→M,N,KBehaviour
AITER_TUNE_GEMM=1) — nothing new to learn; off by default.AITER_TUNE_GEMM_DIRoverrides the destination, for read-only package directories and containers (we had to copy the file out of a container image more than once).Checked
Name derivation (including the merged-copy path), per-family schemas, de-duplication across repeats, disabled-by-default, and graceful degradation on an unwritable destination.
Same campaign as #5262 / #5263–#5266.
🤖 Generated with Claude Code