Fix repeated FlopsProfiler metric accumulation

[baremetaldevx86]

Summary

• restore F.scaled_dot_product_attention after profiling
• restore Tensor.__matmul__ after profiling
• restore torch.bmm from its correct saved implementation
• add a CPU regression test covering repeated profiling sessions and operation restoration
• scope the existing FP16 skip to the test that actually requires FP16

Problem

FlopsProfiler temporarily replaces PyTorch operations with FLOP-counting wrappers. Its cleanup path did not restore F.scaled_dot_product_attention or
Tensor.__matmul__, causing wrappers to accumulate across profiling sessions. As a result, identical model executions reported progressively increasing FLOPs and
MACs.

The torch.bmm cleanup was also incorrect:

torch.bmm = old_functions[torch.matmul.__str__]

Because torch.matmul had already been restored, this rebound torch.bmm to torch.matmul for the rest of the process. This changed normal PyTorch behavior after
profiling: torch.bmm began accepting inputs supported by broadcasting matmul but invalid for bmm.

Fix

Restore every patched operation from its matching saved original function. The patch and cleanup paths now cover the same set of operations with matching PyTorch
version guards.

Testing

A CPU regression test profiles the same scaled-dot-product-attention operation three times and verifies:

• every session reports identical FLOP and MAC totals
• F.scaled_dot_product_attention, Tensor.matmul, and torch.bmm are restored to their original function objects after every session

Observed totals:

Before:
[(65536, 32768), (131072, 65536), (196608, 98304)]

After:
[(65536, 32768), (65536, 32768), (65536, 32768)]

All pre-commit checks pass.

Fixes #7413

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 997dabee8f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a non-merge commit, but its message has no Signed-off-by: trailer, so it violates the workspace commit requirement and can be rejected by CI/review tooling; please recreate or amend the commit with --signoff.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[ebarkhordar]

Reproduced both sides on CPU (clean python:3.11-slim, torch 2.13.0+cpu), base 1d580d65 against head 997dabee, three rounds on the SDPA model from your test:

base: [(65536, 32768), (131072, 65536), (196608, 98304)]
head: [(65536, 32768), (65536, 32768), (65536, 32768)]

The bmm line is more than a metrics bug, and I think the description should say so. torch.bmm = old_functions[torch.matmul.__str__] runs after torch.matmul has already been restored on the line above, so it rebinds torch.bmm to torch.matmul for the rest of the process. matmul broadcasts where bmm requires 3-D: after profiling a plain nn.Linear once, torch.bmm(4x4, 4x4) returns a 4x4 instead of raising RuntimeError: batch1 must be a 3D tensor. Any profiled model reaches it, not only attention ones. Head restores it correctly.

To check nothing of the same shape is left, I parsed _patch_* and _reload_*: base is 51 patched against 49 reloaded (exactly the two you add) plus the mis-keyed bmm; head is 51/51 with matching version guards. So that is the complete set.

One suggestion. The module is skipped entirely when the accelerator does not report fp16 (line 17), so the new test does not collect on a CPU-only runner (1 skipped, collected 0 items here) even though it needs neither fp16 nor a device. Hoisting it above that guard would let the cpu-torch-latest -m sequential leg run it.

[baremetaldevx86]

Thanks for the thorough verification. Agreed that the incorrect torch.bmm restoration has broader process-wide correctness implications, not just inaccurate profiler metrics. I’ll make that explicit in the PR description.

I also addressed the CPU collection issue. Since a module-level pytest.skip(..., allow_module_level=True) aborts collection of the entire file even if the test is defined above it, I moved the FP16 capability check into the existing test that actually uses FP16. The repeated-profile regression now collects and passes in the CPU-only sequential test path.

[tohtana]

@baremetaldevx86 Thank you for the fix! This looks good to me.