Skip to content

Docs: add custom perf model recipe#673

Open
ajassani wants to merge 10 commits into
mainfrom
ajassani/perf-model-docs-after-630
Open

Docs: add custom perf model recipe#673
ajassani wants to merge 10 commits into
mainfrom
ajassani/perf-model-docs-after-630

Conversation

@ajassani

Copy link
Copy Markdown
Collaborator

Summary

  • Adds a concise recipe for adding native TraceLens perf models.
  • Clarifies when to use an extension file versus adding an operator to shared source.
  • Notes that category-only extensions label ops without populating modeled FLOPs/bytes metrics.

Test plan

  • Docs-only change; reviewed rendered Markdown source.

Stacking

ajassani and others added 9 commits May 28, 2026 13:49
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>
@ajassani ajassani requested a review from a team as a code owner May 28, 2026 19:44
@ajassani ajassani force-pushed the ajassani/perf-model-docs-after-630 branch from cea394e to 9972b94 Compare May 28, 2026 19:49
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>
@ajassani ajassani force-pushed the ajassani/perf-model-docs-after-630 branch from 9972b94 to 552e90c Compare May 28, 2026 19:53
Base automatically changed from ajassani/op-category-registry to main May 29, 2026 15:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant