Repository: pytorch
Issue: FlopCounterMode doesn't support HOP
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,module: flop counter

### 🐛 Describe the bug

I'm trying to add micro-benchmark for flex attention, which is implemented by HOP. I use ```torch.utils.flop_counter.FlopCounterMode```, but it doesn't support capture FLOP for HOP.
```
def run_flex_decoding_causal(device: str = "cuda"):
    dtype_flops_utilization_map = {
        torch.bfloat16: "0.8",
    }
    B = 4
    H = 16
    D = 64
    q_len = [32]
    kv_len = [1024, 8192]
    results = []

    def causal_mask(batch, head, token_q, token_kv):
        return token_q >= token_kv
    
    for dtype, expected_flops_utilization in dtype_flops_utilization_map.items():
        flops_utilization = 0
        for seqlen_q in q_len:
            for seqlen_kv in kv_len:
                query = torch.randn((B, H, seqlen_q, D), dtype=dtype, device=device, requires_grad=True)
                key = torch.randn((B, H, seqlen_q, D), dtype=dtype, device=device, requires_grad=True)
                value = torch.randn((B, H, seqlen_q, D), dtype=dtype, device=device, requires_grad=True)
                block_mask = create_block_mask(causal_mask, 1, 1, seqlen_q, seqlen_kv)

                with FlopCounterMode(display=False) as mode:
                    flex_attention(query, key, value, block_mask=block_mask)
                
                flops = mode.get_total_flops()

                compiled_fn = torch.compile(flex_attention, dynamic=False)

                for _ in range(WARMUP_ITER):
                    compiled_fn(query, key, value, block_mask=block_mask)

                us_per_iter = benchmarker.benchmark_gpu(lambda: compiled_fn(query, key, value, block_mask=block_mask)) * 1000
                flops_utilization += us_per_iter * flops / 1e9 / A100_40G_BF16_TFLOPS
        
        flops_utilization = flops_utilization / len(input_shapes)
        dtype_str = str(dtype).replace("torch.", "")
        results.append(
            Experiment(
                "flex_decoding_causal",
                "flops_utilization",
                expected_flops_utilization,
                f"{flops_utilization:.02f}",
                dtype_str,
                device,
            )
        )
    
    return results
```

Error
```
Traceback (most recent call last):
  File "/data/users/ybliang/pytorch/benchmarks/gpt_fast/benchmark.py", line 328, in <module>
    main(output_file=args.output)
  File "/data/users/ybliang/pytorch/benchmarks/gpt_fast/benchmark.py", line 309, in main
    lst = func()
  File "/data/users/ybliang/pytorch/benchmarks/gpt_fast/benchmark.py", line 117, in run_flex_decoding_causal
    flex_attention(query, key, value, block_mask=block_mask)
  File "/home/ybliang/local/pytorch/torch/nn/attention/flex_attention.py", line 987, in flex_attention
    out, lse = torch.compile(
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 465, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/nn/attention/flex_attention.py", line 982, in _flex_attention_hop_wrapper
    return flex_attention_hop(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_higher_order_ops/flex_attention.py", line 62, in __call__
    return super().__call__(
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 431, in __call__
    return wrapper()
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 632, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 427, in wrapper
    return self.dispatch(
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 411, in dispatch
    return kernel(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_higher_order_ops/flex_attention.py", line 654, in flex_attention_autograd
    out, logsumexp = FlexAttentionAutogradOp.apply(
  File "/home/ybliang/local/pytorch/torch/autograd/function.py", line 575, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/home/ybliang/local/pytorch/torch/_higher_order_ops/flex_attention.py", line 542, in forward
    out, logsumexp = flex_attention(
  File "/home/ybliang/local/pytorch/torch/_higher_order_ops/flex_attention.py", line 62, in __call__
    return super().__call__(
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 431, in __call__
    return wrapper()
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 632, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 422, in wrapper
    return torch.overrides.handle_torch_function(
  File "/home/ybliang/local/pytorch/torch/overrides.py", line 1717, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/home/ybliang/local/pytorch/torch/_higher_order_ops/flex_attention.py", line 38, in __torch_function__
    return func(*args, **(kwargs or {}))
  File "/home/ybliang/local/pytorch/torch/_higher_order_ops/flex_attention.py", line 62, in __call__
    return super().__call__(
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 431, in __call__
    return wrapper()
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 632, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 427, in wrapper
    return self.dispatch(
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 333, in dispatch
    raise NotImplementedError(
NotImplementedError: There was no rule registered for HOP flex_attention and mode <torch.utils.flop_counter.FlopCounterMode object at 0x7f00f2d76560>. We recommend filing an issue.
```

### Versions

N/A

cc @ezyang @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh