feat(categorization): add V4 Flash AITER op categorizations

[brieflynn]

PR: feat(categorization): add V4 Flash AITER op categorizations

File changed: TraceLens/PerfModel/torch_op_mapping.py

Summary

When profiling DeepSeek-V4 training workloads that use AITER, TraceLens reports classify all AITER operator calls as "other" (see #778). This PR adds 17 AITER operator mappings to CATEGORY_ONLY_OP_MAPPING so that DeepSeek-V4 traces using the AITER runtime produce accurate per-category compute breakdowns instead of lumping everything into "other".

Mappings cover five categories:

• SDPA_fwd -- fused attention, Lightning Indexer sparse KV selection
• MoE_fused -- fused MoE forward, expert routing
• GEMM -- FP8 block-scaled and BF16 GEMMs
• NORM_fwd -- fused RMSNorm + quantization
• elementwise -- quantization, SiLU, RoPE

Code Change

Before (torch_op_mapping.py, end of CATEGORY_ONLY_OP_MAPPING):

    # InferenceAttention extras (KV-cache writes).
    "_C_cache_ops::reshape_and_cache_flash": "InferenceAttention",
    "_C_cache_ops::concat_and_cache_mla": "InferenceAttention",
}

After:

    # InferenceAttention extras (KV-cache writes).
    "_C_cache_ops::reshape_and_cache_flash": "InferenceAttention",
    "_C_cache_ops::concat_and_cache_mla": "InferenceAttention",
    # V4 Flash AITER attention ops (CSA/HCA fused attention + output projection).
    "aiter::v4_attention_with_output": "SDPA_fwd",
    # V4 Flash AITER Lightning Indexer (CSA sparse KV selection).
    "aiter::cp_gather_indexer_k_quant_cache": "SDPA_fwd",
    "aiter::indexer_score_topk": "SDPA_fwd",
    "aiter::_top_k_per_row_prefill": "SDPA_fwd",
    # V4 Flash AITER MoE (fused forward, expert routing).
    "aiter::moe_forward": "MoE_fused",
    "aiter::topk_use_mulblocks": "MoE_fused",
    # V4 Flash AITER GEMMs (FP8 block-scaled, BF16).
    "aiter::gemm_a8w8_blockscale_preshuffle_impl": "GEMM",
    "aiter::gemm_a16w16_": "GEMM",
    "aiter::gemm_a16w16": "GEMM",
    # V4 Flash AITER fused RMSNorm + quantization.
    "aiter::rmsnorm_quant": "NORM_fwd",
    "aiter::_aiter_rms_quant": "NORM_fwd",
    "aiter::rmsnorm": "NORM_fwd",
    # V4 Flash AITER quantization ops.
    "aiter::per_group_quant_hip": "elementwise",
    "aiter::dynamic_per_token_scaled_quant": "elementwise",
    # V4 Flash AITER elementwise ops (activation, RoPE).
    "aiter::silu_and_mul": "elementwise",
    "aiter::rope_rotate_activation": "elementwise",
}

Testing

python -m pytest tests/test_perf_report_regression.py -v passes. Generated TraceLens report on a DSv4 AITER trace and verified the "other" category is reduced.

Pull Request Template

Note to AMDers:
This is a public repository. Please do not upload any confidential or customer data. Make sure all such data has been anonymized or removed before making this PR. If you need to attach any private files or links, please insert a Internal OneDrive Link or a Jira Ticket Link instead.