Repository: pytorch
Issue: [INDUCTOR] [CPU] [GPT-FAST-MOE] large perf regression with coordinate_descent_tuning disabled
Labels: oncall: pt2,oncall: cpu inductor

### 🐛 Describe the bug

When the flag `coordinate_descent_tuning` disabled, GPT-FAST-MOE encounters a large perf regression: 52s -> 1049s. The impact of disabling it on CPU is to fallback bmm and mm in decomposition.

### Code snippet
```python
import torch
from torch import nn, Tensor
from torch.nn import functional as F
import torch._inductor.config

torch._inductor.config.cpp.enable_kernel_profile = True
torch._inductor.config.coordinate_descent_tuning = False # True

dim = 4096
num_experts = 8
num_activated_experts = 2
intermediate_size = 14336

class ConditionalFeedForwardBit8(nn.Module):
    def __init__(self, target_dtype):
        super().__init__()

        self.target_dtype = target_dtype

        self.register_buffer("w1", torch.empty(num_experts, intermediate_size, dim, dtype=target_dtype))
        self.register_buffer("w2", torch.empty(num_experts, dim, intermediate_size, dtype=target_dtype))
        self.register_buffer("w3", torch.empty(num_experts, intermediate_size, dim, dtype=target_dtype))

        self.register_buffer("scales1", torch.empty(num_experts, intermediate_size, dtype=torch.bfloat16))
        self.register_buffer("scales2", torch.empty(num_experts, dim, dtype=torch.bfloat16))
        self.register_buffer("scales3", torch.empty(num_experts, intermediate_size, dtype=torch.bfloat16))

    def forward(self, x: Tensor, expert_indices: Tensor) -> Tensor:
        w1_weights = self.w1.to(x.dtype)[expert_indices] # [T, A, D, D]
        w3_weights = self.w3.to(x.dtype)[expert_indices] # [T, A, D, D]
        w2_weights = self.w2.to(x.dtype)[expert_indices]
        x1 = F.silu(torch.einsum('ti,taoi -> tao', x, w1_weights) * self.scales1[expert_indices].to(x.dtype))
        x3 = torch.einsum('ti, taoi -> tao', x, w3_weights) * self.scales3[expert_indices].to(x.dtype)
        expert_outs = torch.einsum('tao, taio -> tai', (x1 * x3), w2_weights) * self.scales2[expert_indices].to(x.dtype)  # [T, A, D, D]
        return expert_outs

def linear_forward_int8(x, weight_int8pack, scales, out_features):
    origin_x_size = x.size()
    x = x.reshape(-1, origin_x_size[-1])
    c = torch.ops.aten._weight_int8pack_mm(x, weight_int8pack, scales)
    new_shape = origin_x_size[:-1] + (out_features,)
    c = c.reshape(new_shape)
    return c

class WeightOnlyBit8Linear(torch.nn.Module):
    __constants__ = ['in_features', 'out_features']
    in_features: int
    out_features: int
    weight: torch.Tensor

    def __init__(self, in_features: int, out_features: int, bias: bool = True,
                 device=None, dtype=None, target_dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.register_buffer("weight", torch.empty((out_features, in_features), dtype=target_dtype))
        self.register_buffer("scales", torch.ones(out_features, dtype=torch.bfloat16))

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return linear_forward_int8(
           input,
           self.weight, self.scales, self.out_features)

class MOEFeedForward(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.gate = WeightOnlyBit8Linear(dim, num_experts, bias=False, target_dtype=torch.int8)
        self.cond_ffn = ConditionalFeedForwardBit8(torch.int8)
        self.dim = dim
        self.num_activated_experts = num_activated_experts

    def forward(self, x: Tensor) -> Tensor:
        x = x.view(-1, self.dim)
        # T = num_tokens, E = num_experts, D = hidden dim, A = activated experts
        # x: [T, D]
        scores = self.gate(x) # [T, E]
        expert_weights = F.softmax(scores, dim=-1)
        expert_weights, expert_indices = torch.topk(expert_weights, self.num_activated_experts, dim=-1) # [T, A], [T, A]
        expert_weights /= expert_weights.sum(dim=-1, keepdim=True) # [T, A]
        expert_outs = self.cond_ffn(x, expert_indices)
        return torch.einsum('tai,ta -> ti', expert_outs, expert_weights)


mod = MOEFeedForward()
input_shape = (1, 1, dim,)
x = torch.randn(input_shape, dtype=torch.bfloat16)
compiled_mod = torch.compile(mod)

for i in range(5):
  y = compiled_mod(x)

prof = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
with prof:
  y = compiled_mod(x) 
prof.export_chrome_trace("test_moe.json")
print(prof.key_averages().table(sort_by="self_cpu_time_total"))
```


### Profiling
coordinate_descent_tuning=True
```

---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                             Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
      graph_0_cpp_fused_bmm_sum_1        82.07%     884.479us        82.07%     884.479us     884.479us             1
            Torch-Compiled Region        11.72%     126.292us        98.82%       1.065ms       1.065ms             1
        aten::_weight_int8pack_mm         2.65%      28.601us         3.03%      32.667us      32.667us             1
                       aten::topk         1.33%      14.381us         1.33%      14.381us      14.381us             1
         TorchDynamo Cache Lookup         1.18%      12.767us         1.18%      12.767us      12.767us             1
    inductor::_reinterpret_tensor         0.60%       6.493us         0.60%       6.493us       2.164us             3
                      aten::empty         0.38%       4.066us         0.38%       4.066us       4.066us             1
     graph_0_cpp_fused__softmax_0         0.06%       0.698us         0.06%       0.698us       0.698us             1
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 1.078ms
```

coordinate_descent_tuning=False
```
-----------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-----------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
       graph_0_cpp_fused__to_copy_index_1        61.29%      96.359ms        61.29%      96.359ms      96.359ms             1
                              aten::copy_        33.98%      53.422ms        33.98%      53.422ms      17.807ms             3
                                aten::bmm         4.16%       6.542ms        38.20%      60.056ms      15.014ms             4
                    Torch-Compiled Region         0.35%     550.429us        99.98%     157.187ms     157.187ms             1
       graph_0_cpp_fused_index_mul_silu_2         0.07%     106.154us         0.07%     106.154us     106.154us             1
                aten::_weight_int8pack_mm         0.02%      37.146us         0.03%      42.080us      42.080us             1
            inductor::_reinterpret_tensor         0.02%      37.033us         0.02%      37.033us       3.086us            12
                              aten::clone         0.02%      33.289us        34.02%      53.486ms      17.829ms             3
                 TorchDynamo Cache Lookup         0.02%      31.353us         0.02%      31.353us      31.353us             1
                               aten::topk         0.02%      24.819us         0.02%      24.819us      24.819us             1
                              aten::empty         0.01%      22.710us         0.01%      22.710us       5.678us             4
                         aten::contiguous         0.01%      16.971us        34.03%      53.503ms      17.834ms             3
                         aten::empty_like         0.01%      12.382us         0.02%      30.158us      10.053us             3
    graph_0_cpp_fused_div_index_mul_sum_3         0.01%       9.925us         0.01%       9.925us       9.925us             1
                         aten::as_strided         0.00%       7.801us         0.00%       7.801us       3.900us             2
                       aten::resolve_conj         0.00%       3.070us         0.00%       3.070us       0.384us             8
             graph_0_cpp_fused__softmax_0         0.00%       1.890us         0.00%       1.890us       1.890us             1
-----------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 157.218ms
```

### Analysis

According to the current analysis, there are two main reasons:
- With flag disabled, bmm breaks the origin c++ kernel and some redundant calculations are generated to store uint8 to bf16 and then convert to fp32. The conversion from uint8 to bf16 takes a long time.
- Many of the bmm kernels do additional contiguous in [Mkldnn Matmul](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/mkldnn/Matmul.cpp#L298) for non-contiguous and non-transposed format, which could not be solved by https://github.com/pytorch/pytorch/pull/122599).

### Versions

PyTorch: 34bce27f0d12bf7226b37dfe365660aad456701a


cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @jgong5 @leslie-fang-intel @yanbing-j @mingfeima