Repository: pytorch
Issue: Better namings for triton fusion ops when a custom triton kernel is present?
Labels: triaged,oncall: pt2,module: inductor

### 🚀 The feature, motivation and pitch

Hi, the code can run fine. It is just that the generated comments and names are a bit confusing.

Say we have a function with some torch ops at the beginning and _scaled_mm (which has inductor triton lowerings added earlier).
```
amax_row = torch.max(torch.abs(x), dim=1, keepdim=True).values
scale = _amax_to_scale(amax_row, dtype_float8, x.dtype)  # shape is [M]
# x * scale is M x K * M, broadcast
x_fp8 = _to_fp8_saturated(x * scale, dtype_float8)  # clamp and cast
x_inverse_scale = scale.reciprocal()

y = torch._scaled_mm(
    x_fp8,
    w_t_fp8,
    x_inverse_scale,
    w_inverse_scale_row.t(),
    bias,
    out_dtype=output_dtype,
    use_fast_accum=use_fast_accum,
)
```

After inductor, we have two fused triton kernels. The first one comprises of all the torch ops and the second one is the triton lowering of _scaled_mm. 
```
# Topologically Sorted Source Nodes: [x, abs_1, max_1, amax, clamp, res, mul, x_1, x_fp8, x_inverse_scale, y], Original ATen: [aten.sigmoid, aten.abs, aten.max, aten._to_copy, aten.clamp, aten.reciprocal, aten.mul, aten._scaled_mm]
stream0 = get_raw_stream(0)
triton_per_fused__scaled_mm__to_copy_abs_clamp_max_mul_reciprocal_sigmoid_0.run(arg0_1, buf2, buf3, 1024, 512, grid=grid(1024), stream=stream0)
del arg0_1
buf5 = empty_strided_cuda((1024, 2048), (2048, 1), torch.bfloat16)
# Topologically Sorted Source Nodes: [x, amax, clamp, res, mul, x_1, x_fp8, x_inverse_scale, y, y_1, y_2], Original ATen: [aten.sigmoid, aten._to_copy, aten.clamp, aten.reciprocal, aten.mul, aten._scaled_mm, aten.relu, aten.add]
triton_tem_fused__scaled_mm__to_copy_add_clamp_mul_reciprocal_relu_sigmoid_1.run(buf2, arg2_1, buf3, arg1_1, buf5, grid=torch._inductor.kernel.mm_common.mm_grid(1024, 2048, meta0), stream=stream0)
```

However, the namings and comments are a bit unintuitive:
1. aten._scaled_mm is listed as original atens in both kernels, but it is only present in the second one.
2. [partly explained by the numberings] The first triton kernel does not use aten._scaled_mm but has it in the name. The second triton kernel only comprises of aten._scaled_mm but has other ops in its name.

Any idea on how to possibly improve them?

### Alternatives

I imagine the solution would look like this:
* do not attempt to fuse them together, since we cannot really fuse them
* [optional] Sort the list of original atens before printing the comments 

### Additional context

repro (need H100):
```
import os

os.environ["TORCH_LOGS"] = "+output_code"
os.environ["TORCHINDUCTOR_UNIQUE_KERNEL_NAMES"] = "1"

import torch
import torch._inductor.config


torch._inductor.config.force_disable_caches = True
# inductor_config.triton.descriptive_names = "False"

################## inputs for each scenario

fusion_case = "pointwise"  # "pointwise" or "reduction"

# Matmul Y = X [M, K] x W [N, K]
M = 1024  # batch size
K = 512  # in_features
N = 2048  # out_features

################## setup
device = "cuda:0"
dtype_float8 = torch.float8_e4m3fn

input_dtype = torch.bfloat16  # torch.float32 or bfloat16
output_dtype = torch.bfloat16  # torch.float32 or dtype_float8 or torch.bfloat16

use_fast_accum = True

x = torch.rand(M, K, dtype=input_dtype, device=device)
w = torch.rand(N, K, dtype=input_dtype, device=device)

bias = None


################### utilities
# ref fbcode/caffe2/torch/fb/model_transform/experimental/fp8_linear.py FP8LinearDynamic

E4M3_MAX_POS: float = torch.finfo(torch.float8_e4m3fn).max
E5M2_MAX_POS: float = torch.finfo(torch.float8_e5m2).max
FP16_MAX_POS: float = torch.finfo(torch.float16).max
EPS: float = 1e-12

# fbcode/caffe2/torch/fb/model_transform/experimental/fp8_linear.py


@torch.no_grad()
def _amax_to_scale(
    amax: torch.Tensor, float8_dtype: torch.dtype, orig_dtype: torch.dtype
) -> torch.Tensor:
    # To make scale dtype to be fp32 for accuracy
    amax = amax.float()
    if float8_dtype == torch.float8_e4m3fn:
        res = E4M3_MAX_POS / torch.clamp(amax, min=EPS)
    else:  # e5m2
        res = E5M2_MAX_POS / torch.clamp(amax, min=EPS)

    # Ensure that the scale is representable in float16,
    # this helps when amax is small. We are assuming that we don't need
    # to care about this for float32/bfloat16.
    if orig_dtype is torch.float16:
        res = torch.clamp(res, max=FP16_MAX_POS)
    return res


def _to_fp8_saturated(x: torch.Tensor, float8_dtype: torch.dtype) -> torch.Tensor:
    # The default behavior in PyTorch for casting to `float8_e4m3fn`
    # and `e5m2` is to not saturate. In this context, we should saturate.
    # A common case where we want to saturate is when the history of a
    # tensor has a maximum value of `amax1`, and the current amax value
    # is `amax2`, where `amax1 < amax2`. This is common when using delayed
    # scaling.
    if float8_dtype == torch.float8_e4m3fn:
        x = x.clamp(min=-1 * E4M3_MAX_POS, max=E4M3_MAX_POS)
    else:
        x = x.clamp(min=-1 * E5M2_MAX_POS, max=E5M2_MAX_POS)
    return x.to(float8_dtype)


################### rowwise scaling fp8

# quantize weight, done in model building stage prior to inference
weight_amax_row = torch.max(torch.abs(w), dim=1, keepdim=True).values
weight_scale_row = _amax_to_scale(weight_amax_row, dtype_float8, w.dtype)
w_t_fp8_row = _to_fp8_saturated(w * weight_scale_row, dtype_float8).t()
w_inverse_scale_row = weight_scale_row.reciprocal()  # element-wise reciprocal


@torch.no_grad()
def fb8_rowwise_scaling(x, w_t_fp8, w_inverse_scale_row):
    if fusion_case == "pointwise":
        # Fusion Case 1: Pointwise (e.g. Sigmoid) + Matmul
        x = torch.sigmoid(x)
    else:
        # Fusion Case 2: Reduction (e.g. LayerNorm) + Matmul
        layer_norm = torch.nn.LayerNorm(K, device=device, dtype=input_dtype)
        x = layer_norm(x)

    # quantize input x
    amax_row = torch.max(torch.abs(x), dim=1, keepdim=True).values
    scale = _amax_to_scale(amax_row, dtype_float8, x.dtype)  # shape is [M]
    # x * scale is M x K * M, broadcast
    x_fp8 = _to_fp8_saturated(x * scale, dtype_float8)  # clamp and cast
    x_inverse_scale = scale.reciprocal()

    y = torch._scaled_mm(
        x_fp8,
        w_t_fp8,
        x_inverse_scale,
        w_inverse_scale_row.t(),
        bias,
        out_dtype=output_dtype,
        use_fast_accum=use_fast_accum,
    )

    # epilogue
    y = torch.nn.functional.relu(y)
    y = y + 0.01

    return y


fb8_rowwise_scaling_compiled = torch.compile(
    fb8_rowwise_scaling, mode="max-autotune-no-cudagraphs"
)
y = fb8_rowwise_scaling_compiled(x, w_t_fp8_row, w_inverse_scale_row)

print("done")
# Matmul Y = X [M, K] x W [N, K]
```

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire