feat(perfmodel): add varlen aiter FA + aten::_flash_attention_forward perf models (Wan 2.2 coverage)

[gphuang]

Summary

Closes #650
Also closes #290 and #590

Adds shape-based SDPA perf models for AITER varlen FlashAttention fwd/bwd, aiter::fmha_v3_bwd, and aten::_flash_attention_forward. Also stops the training path from being overridden by the inference-only pseudo-op mapping, so Wan 2.x attention rows report under SDPA_fwd / SDPA_bwd with GFLOPS.

Changes

• Register 6 new attention op names in the core perf-model mapping.
• Add packed-varlen FLOPs handling via num_seqs_q / num_seqs_kv.
• Keep the pseudo-op extension class, but prevent it from overriding the core training mapping.
• Add 14 regression tests using real Wan 2.2 event payloads.
• Refresh affected perf-report and compare-report CSV references for rows that now move from other to SDPA_fwd / SDPA_bwd.

Test plan

• pytest tests/test_aiter_fmha_v3_varlen_ops.py
• Focused attention / pseudo-op / Primus categorization tests pass locally
• pytest tests/test_compare_perf_reports.py tests/test_perf_report_regression.py (18 passed)

[Copilot]

Pull request overview

This PR expands TraceLens’ SDPA/FlashAttention performance-model coverage by adding core (shape-based) perf models for AITER varlen FA forward/backward + AITER non-varlen fmha_v3_bwd, and PyTorch’s dispatcher-level aten::_flash_attention_forward. It also removes a pseudo-op extension mapping that previously overrode the core mapping in training traces, improving categorization and GFLOPS reporting for Wan 2.x workloads.

Changes:

• Add new SDPA-derived perf model classes for aiter::(wrapper_)fmha_v3_varlen_{fwd,bwd}, aiter::fmha_v3_bwd, and aten::_flash_attention_forward.
• Register new op→perf-model mappings and route new backward op-names into the SDPA_bwd category.
• Stop pseudo-op extension mappings from overriding the core mapping for aiter::fmha_v3_varlen_fwd, and add unit tests covering the new mappings and shape extraction.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
TraceLens/PerfModel/torch_op_mapping.py	Registers new op mappings and updates SDPA forward/backward categorization for the new AITER/ATen attention ops.
TraceLens/PerfModel/perf_model.py	Implements the new SDPA-derived perf model classes and shared varlen argument parsers.
TraceLens/PerfModel/extensions/pseudo_ops_perf_utils.py	Prevents the extension mapping from overriding the core mapping for aiter::fmha_v3_varlen_fwd.
tests/test_aiter_fmha_v3_varlen_ops.py	Adds unit tests validating mapping registration, categorization, and FLOPs extraction for the new ops.

Comments suppressed due to low confidence (1)

TraceLens/PerfModel/perf_model.py:3362

• The varlen backward FLOPs calculation uses only max_seqlen_{q,kv} once and does not account for num_seqs_q/num_seqs_kv when multiple sequences are present. This will under-estimate bwd FLOPs for packed multi-sequence varlen attention. Suggest matching flash_attention_varlen_backward.flops() (accumulate for max seqlen + estimate remaining sequences) so aiter::fmha_v3_varlen_bwd behaves consistently with the existing flash-attn varlen perf model.

    def flops(self):
        return self.flops_bwd_func(
            self.B,
            self.max_seqlen_q,
            self.H_Q,
            self.max_seqlen_kv,
            self.H_KV,
            self.d_h_qk,
            self.d_h_v,
            self.param_details["causal"],
            self.param_details["flash_impl"],
        )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[gphuang]

@devalshah-amd @devalshahamd Could you approve this PR? Thanks!

[kyle-hoffmeyer]

SDPA.get_simulation_time requires dtype from perf model to calculate origami metrics. _parse_aiter_fmha_v3_varlen_fwd_args, _parse_aiter_fmha_v3_varlen_bwd_args, aten__flash_attention_forward.get_param_details need to all return "dtype_A_B" as well, which would be event["args"]["Input type"][q_idx], event["args"]["Input type"][k_idx] for these perf models (data type of Q and K matrices)

[gphuang]

Addressed the dtype metadata note in 4df85aa.

Added dtype_A_B extraction to:

• _parse_aiter_fmha_v3_varlen_fwd_args
• _parse_aiter_fmha_v3_varlen_bwd_args
• aten___flash_attention_forward.get_param_details

For each parser this is now (Input type[q_idx], Input type[k_idx]), so SDPA simulation/origami has Q/K precision metadata available.

Validation: pytest tests/test_aiter_fmha_v3_varlen_ops.py (15 passed) and black --check on touched files.