Docs: add custom perf model recipe#673
Open
ajassani wants to merge 10 commits into
Open
Conversation
Adds a flat ``op_name -> category`` registry alongside the existing
``categorize_torch_op`` if/elif chain. The registry is auto-populated from
``op_to_perf_model_class_map`` + ``dict_base_class2category`` and topped up
with explicit overrides for category-only ops (e.g. SSM_bwd, MoE_aux,
CONV variants without a perf model) that the legacy chain handled inline.
A small regex layer (``triton``, ``record_param_comms``) and a kernel-name
fallback (``void at::native`` probes) preserve the rest of the legacy
chain's behavior.
This PR is purely additive:
- The legacy ``categorize_torch_op`` is unchanged.
- The new ``categorize_torch_op_v2`` runs in parallel.
- ``test_categorize_torch_op_parity`` asserts v1 and v2 return the same
category for every name reachable through perf-model registration,
every name hardcoded inside the legacy chain, the pattern-matched
names, "other" cases, and the kernel-name fallback.
Motivation: today categorization is coupled to perf-model availability
(category derived from the perf-model class hierarchy), which forces a
new perf-model class for every category-only op and pushes hardcoded name
lists into the categorizer function. The registry breaks that coupling
so subsequent PRs can:
- Delete the if/elif chain (PR B).
- Add a ``get_op_categories`` extension hook so category-only entries
don't require authoring a perf model (PR B).
- Add a drift report for ops landing in ``"other"`` (PR C).
No user-visible behavior change in this PR.
Co-authored-by: Cursor <cursoragent@cursor.com>
Replaces the legacy if/elif chain in categorize_torch_op with the registry-backed categorizer directly. Keeps dict_cat2names as the compatibility view for legacy report sheets, while OP_CATEGORY_REGISTRY becomes the source of truth for final op labels. This also splits GroupedGEMM out of GEMM, fixes SDPA backward names that previously fell through unless dict_cat2names was mutated, and wires perf_model_extension, op_category_extension, and dict_cat2names_extension into the registry. Focused tests cover expected categories, kernel fallback behavior, extension registration, grouped GEMM categorization, and the existing pseudo-op extension compatibility cases. Co-authored-by: Cursor <cursoragent@cursor.com>
Move perf-modeled op categorization onto category/bwd_category metadata so extension models define their own report category. Keep category-only ops separate from legacy sheet membership and expose linked-forward backward metrics in the unified perf output. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep categorization registry-only, retain legacy sheet categories only for report sheet generation, and move RMSNorm extension category metadata onto an extension-side base class. Co-authored-by: Cursor <cursoragent@cursor.com>
Derive legacy sheet membership from perf-model metadata so the registry remains the only categorization source of truth. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep torch op categorization in one module now that legacy compatibility maps are gone, and preserve RMSNorm as its own legacy report sheet. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep the categorization cleanup focused by leaving metrics provenance out of this PR. Co-authored-by: Cursor <cursoragent@cursor.com>
FusedLnModulateBackward (added on main in #633) inherits Normalization, which now declares category='NORM_fwd' on the base. The backward sibling needs an explicit category='NORM_bwd' override, matching every other *Bwd Normalization subclass in this file. Adds test_every_perf_model_class_declares_category so any future perf-model class missing a 'category' attribute fails with a clear, named error rather than a cryptic ValueError from registry construction at import time. Co-authored-by: Cursor <cursoragent@cursor.com>
Class had pass-only stubs and was already commented out in the pseudo-op registration map. The op is now categorized via CATEGORY_ONLY_OP_MAPPING. Confirmed safe to remove by @devalshahamd in PR review. Co-authored-by: Cursor <cursoragent@cursor.com>
cea394e to
9972b94
Compare
Expand perf report column docs with guidance for adding native perf models, using extension files for private/customer-specific ops, and registering category-only ops. Co-authored-by: Cursor <cursoragent@cursor.com>
9972b94 to
552e90c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Test plan
Stacking