[AMD 350x] Helion Layernorm kernel performance is worse compared to inductor for some shapes in backward pass

[MengjiaoZhou]

We are currently optimizing the LayerNorm kernel for AMD_350x. Despite running the Helion autotuner multiple times, the Helion LayerNorm kernel remains less performant than Inductor on this hardware for some shapes, especially for (1, 2304, 36864) , helion is 2x worser compared to inductor: helion (0.72) v.s inductor(0.25)

Re-autotuning the kernel did not yield any improvement.

Performance Results

Backward pass

        (B, M, N)    inductor_layernorm-latency    inductor_layernorm-gbps    inductor_layernorm-compile_time    helion_layernorm-latency    helion_layernorm-gbps    helion_layernorm-accuracy    helion_layernorm-compile_time
-----------------  ----------------------------  -------------------------  ---------------------------------  --------------------------  -----------------------  ---------------------------  -------------------------------
    (1, 512, 256)             0.026003 (±0.00%)                    30.2835                            0.63873           0.033716 (±0.00%)                  23.3554                            1                          730.278
 (1, 1152, 11776)             0.053272 (±0.00%)                  1528.82                            108.545             0.086619 (±0.00%)                 940.241                             1                          573.834
 (1, 2304, 16384)             0.102755 (±0.00%)                  2204.84                            117.443             0.106123 (±0.00%)                2134.86                              1                          716.848
(1, 2304, 103008)             0.731764 (±0.00%)                  1946.52                            135.724             0.720036 (±0.00%)                1978.23                              1                          387.997
 (1, 2304, 36864)             0.252344 (±0.00%)                  2020.08                            126.589             0.726013 (±0.00%)                 702.13                              1                          858.454
          average           0.23322754986584188                  1546.11                             97.7879           0.3345014937222004                1155.76                              1                          653.482

forward pass

        (B, M, N)    inductor_layernorm-latency    inductor_layernorm-gbps    inductor_layernorm-compile_time    helion_layernorm-latency    helion_layernorm-gbps    helion_layernorm-accuracy    helion_layernorm-compile_time
-----------------  ----------------------------  -------------------------  ---------------------------------  --------------------------  -----------------------  ---------------------------  -------------------------------
    (1, 512, 256)             0.003736 (±0.00%)                    140.473                            164.253           0.004212 (±0.00%)                  124.582                            1                          525.857
 (1, 1152, 11776)             0.017352 (±0.00%)                   3128.53                             634.311           0.017665 (±0.00%)                 3073.14                             1                          557.693
 (1, 2304, 16384)             0.045766 (±0.00%)                   3299.97                             619.205           0.040612 (±0.00%)                 3718.79                             1                          495.516
(1, 2304, 103008)             0.364772 (±0.00%)                   2603.07                             439.316           0.316223 (±0.00%)                 3002.72                             1                          547.408
 (1, 2304, 36864)             0.122046 (±0.00%)                   2784.29                             430.452           0.092525 (±0.00%)                 3672.65                             1                          545.482
          average           0.11073455158621073                   2391.27                             457.508           0.094247492775321                 2718.38                             1                          534.391

Versions

Triton Version: 3.5.0+fb
ROCm: 7.0.2.1 (using -m rocm70 mode)

helion generated code:

# src[layer_norm.py:2377]: if bias is not None:
   # src[layer_norm.py:2378]:     assert bias.size(0) == n, f"bias size mismatch {bias.size(0)} != {n}"
   if bias is not None:
       # src[layer_norm.py:2378]: assert bias.size(0) == n, f"bias size mismatch {bias.size(0)} != {n}"
       assert bias.size(0) == n, f'bias size mismatch {bias.size(0)} != {n}'
   # src[layer_norm.py:2379]: assert len(normalized_shape) == 1, (
   # src[layer_norm.py:2380]:     "Helion layer norm only supports 1D layer norm currently"
   # src[layer_norm.py:2381]: )
   assert len(normalized_shape) == 1, 'Helion layer norm only supports 1D layer norm currently'
   # src[layer_norm.py:2382]: assert normalized_shape[0] == n, (
   # src[layer_norm.py:2383]:     f"normalized shape mismatch {normalized_shape[0]} != {n}"
   # src[layer_norm.py:2384]: )
   assert normalized_shape[0] == n, f'normalized shape mismatch {normalized_shape[0]} != {n}'
   # src[layer_norm.py:2385]: out = torch.empty([m, n], dtype=x.dtype, device=x.device)
   out = torch.empty([m, n], dtype=x.dtype, device=x.device)
   # src[layer_norm.py:2386]: mean = torch.empty([m], dtype=torch.float32, device=x.device)
   mean = torch.empty([m], dtype=torch.float32, device=x.device)
   # src[layer_norm.py:2387]: rstd = torch.empty([m], dtype=torch.float32, device=x.device)
   rstd = torch.empty([m], dtype=torch.float32, device=x.device)
   # src[layer_norm.py:2389]: for tile_m in hl.tile(m):
   _RDIM_SIZE_1 = 65536
   # src[layer_norm.py:2389]: for tile_m in hl.tile(m):
   # src[layer_norm.py:2390]:     acc = x[tile_m, :].to(torch.float32)
   # src[layer_norm.py:2391]:     # Compute mean
   # src[layer_norm.py:2389-2409]: ...
   _launcher(_helion_layer_norm_fwd, (m,), x, weight, bias, out, mean, rstd, mean.size(0), bias.stride(0), mean.stride(0), out.stride(0), out.stride(1), rstd.stride(0), weight.stride(0), x.stride(0), x.stride(1), eps, _RDIM_SIZE_1, num_warps=8, num_stages=2, matrix_instr_nonkdim=32, waves_per_eu=4)
   # src[layer_norm.py:2410]: return out, mean, rstd
   return (out, mean, rstd)
def call():
   from torch._dynamo.testing import rand_strided
   # src[layer_norm.py:2445]: num_blocks = (x.size(0) + m_block - 1) // m_block
   num_blocks = (x.size(0) + m_block - 1) // m_block
   # src[layer_norm.py:2446]: grad_weight_blocks = x.new_empty([num_blocks, n], dtype=torch.float32)
   grad_weight_blocks = x.new_empty([num_blocks, n], dtype=torch.float32)
   # src[layer_norm.py:2447]: grad_bias_blocks = x.new_empty([num_blocks, n], dtype=torch.float32)
   grad_bias_blocks = x.new_empty([num_blocks, n], dtype=torch.float32)
   # src[layer_norm.py:2449]: for mb_cta in hl.tile(x.size(0), block_size=m_block):
   _NUM_SM = helion.runtime.get_num_sm(grad_out.device)
   _RDIM_SIZE_1 = 65536
   # src[layer_norm.py:2449]: for mb_cta in hl.tile(x.size(0), block_size=m_block):
   # src[layer_norm.py:2450]:     grad_w_acc = weight.new_zeros(n, dtype=torch.float32)
   # src[layer_norm.py:2451]:     if compute_bias_grad:
   # src[layer_norm.py:2449-2474]: ...
   _launcher(_helion_layer_norm_bwd, (_NUM_SM * 2,), x, weight, grad_out, mean, rstd, grad_x, grad_weight_blocks, grad_bias_blocks, grad_bias_blocks.size(0), grad_x.size(0), rstd.size(0), x.size(0), grad_bias_blocks.stride(0), grad_bias_blocks.stride(1), grad_out.stride(0), grad_out.stride(1), grad_weight_blocks.stride(0), grad_weight_blocks.stride(1), grad_x.stride(0), grad_x.stride(1), mean.stride(0), rstd.stride(0), weight.stride(0), x.stride(0), x.stride(1), _NUM_SM, _RDIM_SIZE_1, num_warps=8, num_stages=4, matrix_instr_nonkdim=32, waves_per_eu=1)
   # src[layer_norm.py:2476]: grad_weight = grad_weight_blocks.sum(0).to(weight.dtype)
   grad_weight = grad_weight_blocks.sum(0).to(weight.dtype)
   # src[layer_norm.py:2477]: if compute_bias_grad:
   # src[layer_norm.py:2478]:     grad_bias = grad_bias_blocks.sum(0).to(weight.dtype)
   # src[layer_norm.py:2479]:     return grad_x, grad_weight, grad_bias
   if True:
       # src[layer_norm.py:2478]: grad_bias = grad_bias_blocks.sum(0).to(weight.dtype)
       grad_bias = grad_bias_blocks.sum(0).to(weight.dtype)
       # src[layer_norm.py:2479]: return grad_x, grad_weight, grad_bias
       return (grad_x, grad_weight, grad_bias)
   # src[layer_norm.py:2480]: return grad_x, grad_weight, None
   return (grad_x, grad_weight, None)
def call():
   from torch._dynamo.testing import rand_strided
   # src[layer_norm.py:2414]: def layer_norm_bwd(
   # src[layer_norm.py:2415]:     grad_out: torch.Tensor,
   # src[layer_norm.py:2416]:     x: torch.Tensor,
   # src[layer_norm.py:2414-2480]: ...
   grad_out = rand_strided(size=(2304, 36864), stride=(36864, 1), dtype=torch.bfloat16, device='cuda:0')
   x = rand_strided(size=(2304, 36864), stride=(36864, 1), dtype=torch.bfloat16, device='cuda:0')
   mean = rand_strided(size=(2304,), stride=(1,), dtype=torch.float32, device='cuda:0')
   rstd = rand_strided(size=(2304,), stride=(1,), dtype=torch.float32, device='cuda:0')
   weight = rand_strided(size=(36864,), stride=(1,), dtype=torch.bfloat16, device='cuda:0')
   compute_bias_grad = True
   layer_norm_bwd(grad_out, x, mean, rstd, weight, compute_bias_grad)
if __name__ == '__main__':
   call()

inductor generated code:

# AOT ID: ['4_forward']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from cmath import nanj
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
assert_alignment = torch._C._dynamo.guards.assert_alignment
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cpu_pinned = torch._C._dynamo.guards._empty_strided_cpu_pinned
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
empty_strided_mtia = torch._C._dynamo.guards._empty_strided_mtia
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()
empty_strided_p2p = torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p
# kernel path: /var/tmp/torchinductor_mengjiao/bd/cbd42ow37dkvyysrib5m7gbpaewfokss4zdpnauaftk4chzdltpt.py
# Topologically Sorted Source Nodes: [layer_norm], Original ATen: [aten.native_layer_norm]
# Source node to ATen node mapping:
#   layer_norm => add, add_1, convert_element_type, convert_element_type_1, getitem_1, mul, mul_1, rsqrt, sub, var_mean
# Graph fragment:
#   %primals_1 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0" = PlaceHolder[target=primals_1]
#   %buf0 : Tensor "f32[1, 2304, 1][2304, 1, 2304]cuda:0" = PlaceHolder[target=buf0]
#   %buf1 : Tensor "f32[1, 2304, 1][2304, 1, 2304]cuda:0" = PlaceHolder[target=buf1]
#   %buf2 : Tensor "f32[1, 2304, 1][2304, 1, 2304]cuda:0" = PlaceHolder[target=buf2]
#   %rsqrt : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0" = PlaceHolder[target=rsqrt]
#   %primals_2 : Tensor "bf16[36864][1]cuda:0" = PlaceHolder[target=primals_2]
#   %primals_3 : Tensor "bf16[36864][1]cuda:0" = PlaceHolder[target=primals_3]
#   %convert_element_type : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=2] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%primals_1, torch.float32), kwargs = {})
#   %var_mean : [num_users=2] = call_function[target=torch.ops.aten.var_mean.correction](args = (%convert_element_type, [2]), kwargs = {correction: 0, keepdim: True})
#   %getitem_1 : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0"[num_users=2] = call_function[target=operator.getitem](args = (%var_mean, 1), kwargs = {})
#   %add : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%getitem, 1e-05), kwargs = {})
#   %rsqrt : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0"[num_users=2] = call_function[target=torch.ops.aten.rsqrt.default](args = (%add,), kwargs = {})
#   %sub : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%convert_element_type, %getitem_1), kwargs = {})
#   %mul : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sub, %rsqrt), kwargs = {})
#   %mul_1 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul, %primals_2), kwargs = {})
#   %add_1 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_1, %primals_3), kwargs = {})
#   %convert_element_type_1 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%add_1, torch.bfloat16), kwargs = {})
#   return %buf0,%buf1,%buf2,%rsqrt,%convert_element_type_1
triton_red_fused_native_layer_norm_0 = async_compile.triton('triton_red_fused_native_layer_norm_0', '''
import triton
import triton.language as tl
import triton.language.extra.tlx as tlx  # noqa: F401
from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()
@triton_heuristics.reduction(
    size_hints={'x': 4096, 'r0_': 65536},
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {'in_out_ptr0': '*fp32', 'in_out_ptr1': '*fp32', 'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='hip', index=0, multi_processor_count=256, cc='gfx950', major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, max_threads_per_block=1024, warp_size=64), 'constants': {}, 'native_matmul': False, 'enable_fp_fusion': True, 'launch_pdl': False, 'disable_ftz': False, 'configs': [{(0,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (1,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (2,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (3,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (4,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (5,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (6,): [['tt.divisibility', 16]], (7,): [['tt.divisibility', 16]]}]},
    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_red_fused_native_layer_norm_0', 'mutated_arg_names': ['in_out_ptr0', 'in_out_ptr1'], 'optimize_mem': False, 'no_x_dim': False, 'atomic_add_found': False, 'num_load': 5, 'num_store': 3, 'num_reduction': 2, 'backend_hash': 'E3565433EC5266B9E41D740243879373C6297D57D928F6BAE46DD3A44A623A50', 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 32, 'store_cubin': False, 'deterministic': False, 'force_filter_reduction_configs': False, 'mix_order_reduction_allow_multi_stages': False, 'are_deterministic_algorithms_enabled': False, 'is_hip': True, 'is_fbcode': True}
)
@triton.jit
def triton_red_fused_native_layer_norm_0(in_out_ptr0, in_out_ptr1, in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
    xnumel = 2304
    r0_numel = 36864
    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    r0_base = tl.arange(0, R0_BLOCK)[None, :]
    rbase = r0_base
    x0 = xindex
    _tmp3 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages = 2):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp0 = tl.load(in_ptr0 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp1 = tmp0.to(tl.float32)
        tmp2 = tl.broadcast_to(tmp1, [XBLOCK, R0_BLOCK])
        tmp4 = _tmp3 + tmp2
        _tmp3 = tl.where(r0_mask & xmask, tmp4, _tmp3)
    tmp3 = tl.sum(_tmp3, 1)[:, None]
    tmp5 = tl.full([1, 1], 36864.0, tl.float32)
    tmp6 = (tmp3 / tmp5)
    tl.debug_barrier()
    tl.store(in_out_ptr0 + (x0), tmp6, xmask)
    _tmp12 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages = 2):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp7 = tl.load(in_ptr0 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp8 = tmp7.to(tl.float32)
        tmp9 = tmp8 - tmp6
        tmp10 = tmp9 * tmp9
        tmp11 = tl.broadcast_to(tmp10, [XBLOCK, R0_BLOCK])
        tmp13 = _tmp12 + tmp11
        _tmp12 = tl.where(r0_mask & xmask, tmp13, _tmp12)
    tmp12 = tl.sum(_tmp12, 1)[:, None]
    tmp14 = tl.full([1, 1], 36864.0, tl.float32)
    tmp15 = (tmp12 / tmp14)
    tmp16 = tl.full([1, 1], 1e-05, tl.float32)
    tmp17 = tmp15 + tmp16
    tmp18 = tl.rsqrt(tmp17)
    tl.debug_barrier()
    tl.store(in_out_ptr1 + (x0), tmp18, xmask)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages = 2):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp19 = tl.load(in_ptr0 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp23 = tl.load(in_ptr1 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp26 = tl.load(in_ptr2 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp20 = tmp19.to(tl.float32)
        tmp21 = tmp20 - tmp6
        tmp22 = tmp21 * tmp18
        tmp24 = tmp23.to(tl.float32)
        tmp25 = tmp22 * tmp24
        tmp27 = tmp26.to(tl.float32)
        tmp28 = tmp25 + tmp27
        tmp29 = tmp28.to(tl.float32)
        tl.store(out_ptr0 + (r0_1 + 36864*x0), tmp29, r0_mask & xmask)
''', device_str='cuda')
async_compile.wait(globals())
del async_compile
def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (1, 2304, 36864), (84934656, 36864, 1))
    assert_size_stride(primals_2, (36864, ), (1, ))
    assert_size_stride(primals_3, (36864, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 2304, 1), (2304, 1, 2304), torch.float32)
        buf1 = buf0; del buf0  # reuse
        buf2 = empty_strided_cuda((1, 2304, 1), (2304, 1, 2304), torch.float32)
        buf3 = reinterpret_tensor(buf2, (1, 2304, 1), (2304, 1, 1), 0); del buf2  # reuse
        buf4 = empty_strided_cuda((1, 2304, 36864), (84934656, 36864, 1), torch.bfloat16)
        # Topologically Sorted Source Nodes: [layer_norm], Original ATen: [aten.native_layer_norm]
        stream0 = get_raw_stream(0)
        triton_red_fused_native_layer_norm_0.run(buf1, buf3, primals_1, primals_2, primals_3, buf4, 2304, 36864, stream=stream0)
        del primals_3
    return (buf4, primals_1, primals_2, reinterpret_tensor(buf1, (1, 2304, 1), (2304, 1, 1), 0), buf3, )
def get_args():
    from torch._dynamo.testing import rand_strided
    primals_1 = rand_strided((1, 2304, 36864), (84934656, 36864, 1), device='cuda:0', dtype=torch.bfloat16)
    primals_2 = rand_strided((36864, ), (1, ), device='cuda:0', dtype=torch.bfloat16)
    primals_3 = rand_strided((36864, ), (1, ), device='cuda:0', dtype=torch.bfloat16)
    return [primals_1, primals_2, primals_3]
def benchmark_compiled_module(args, times=10, repeat=10):
    from torch._inductor.utils import print_performance
    fn = lambda: call(list(args))
    return print_performance(fn, times=times, repeat=repeat)
if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    args = get_args()
    compiled_module_main('helion_layernorm,inductor_layernorm', lambda times, repeat: benchmark_compiled_module(args, times=times, repeat=repeat))

[MengjiaoZhou]

to clarify: helion has comparable perf compared to inductor for forward pass. only for Backward pass, helion is performing worse.

[umechand-amd]

Can you copy the helion kernel you used for autotuning and the helion generated triton kernel from autotuning?
The helion generated code you shared above doesnt have the triton kernel code.

[MengjiaoZhou]

FWD

helion kernel code used for autotuning

def layer_norm_fwd(
    x: torch.Tensor,
    normalized_shape: list[int],
    weight: torch.Tensor,
    bias: torch.Tensor = None,
    eps: float = 1e-5,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Performs 1D layer normalization on the input tensor using Helion.
    Args:
        x (torch.Tensor): Input tensor of shape [batch_size, dim], expected to be FP16.
        normalized_shape (list[int]): List containing the dimension to normalize over (should be length 1).
        weight (torch.Tensor): Learnable scale parameter of shape [dim].
        bias (torch.Tensor | None): Optional learnable bias parameter of shape [dim].
        eps (float, optional): Small value added to variance for numerical stability. Default is 1e-5.
    Returns:
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
            - The layer-normalized output tensor of shape [batch_size, dim], in FP16.
            - Mean tensor of shape [batch_size], in FP32.
            - Reciprocal standard deviation tensor of shape [batch_size], in FP32.
    """
    m, n = x.size()
    n = hl.specialize(x.size(1))
    assert weight.size(0) == n, f"weight size mismatch {weight.size(0)} != {n}"
    if bias is not None:
        assert bias.size(0) == n, f"bias size mismatch {bias.size(0)} != {n}"
    assert len(normalized_shape) == 1, (
        "Helion layer norm only supports 1D layer norm currently"
    )
    assert normalized_shape[0] == n, (
        f"normalized shape mismatch {normalized_shape[0]} != {n}"
    )
    out = torch.empty([m, n], dtype=x.dtype, device=x.device)
    mean = torch.empty([m], dtype=torch.float32, device=x.device)
    rstd = torch.empty([m], dtype=torch.float32, device=x.device)

    for tile_m in hl.tile(m):
        acc = x[tile_m, :].to(torch.float32)
        # Compute mean
        mean_val = torch.sum(acc, dim=-1) / n
        # Compute variance
        centered = acc - mean_val[:, None]
        var_val = torch.sum(centered * centered, dim=-1) / n
        # Compute reciprocal standard deviation
        rstd_val = torch.rsqrt(var_val + eps)
        # Normalize
        normalized = centered * rstd_val[:, None]
        # Apply affine transformation
        if bias is not None:
            acc = normalized * (weight[:].to(torch.float32)) + (
                bias[:].to(torch.float32)
            )
        else:
            acc = normalized * (weight[:].to(torch.float32))
        out[tile_m, :] = acc.to(x.dtype)
        mean[tile_m] = mean_val
        rstd[tile_m] = rstd_val
    return out, mean, rstd

BWD

helion code used for autotune

def layer_norm_bwd(
    grad_out: torch.Tensor,
    x: torch.Tensor,
    mean: torch.Tensor,
    rstd: torch.Tensor,
    weight: torch.Tensor,
    compute_bias_grad: hl.constexpr = True,  # type: ignore[valid-type]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """
    Compute gradients for weight (dW) and optionally bias (dB) parameters.

    This kernel performs reduction across the batch dimension (M) to accumulate
    gradients for each feature dimension's weight and bias parameters.

    Args:
        grad_out: Gradient w.r.t layer norm output [M, N]
        x: Original input tensor [M, N]
        mean: Per-sample mean computed in forward pass [M]
        rstd: Per-sample reciprocal standard deviation from forward pass [M]
        weight: Weight parameter (used only for dtype/device info) [N]
        compute_bias_grad: Whether to compute bias gradient (default: True)

    Returns:
        (grad_x, grad_weight, grad_bias): Gradients for input, weight, and bias (if computed)
            grad_bias is None if compute_bias_grad is False
    """

    m_block = hl.register_block_size(x.size(0))
    n = hl.specialize(x.size(1))

    grad_x = torch.empty_like(x)
    num_blocks = (x.size(0) + m_block - 1) // m_block
    grad_weight_blocks = x.new_empty([num_blocks, n], dtype=torch.float32)
    grad_bias_blocks = x.new_empty([num_blocks, n], dtype=torch.float32)

    for mb_cta in hl.tile(x.size(0), block_size=m_block):
        grad_w_acc = weight.new_zeros(n, dtype=torch.float32)
        if compute_bias_grad:
            grad_b_acc = weight.new_zeros(n, dtype=torch.float32)
        weight_cta = weight[None, :].to(torch.float32)
        for mb in hl.tile(mb_cta.begin, mb_cta.end):
            x_mb = x[mb, :].to(torch.float32)
            dy_mb = grad_out[mb, :].to(torch.float32)
            mean_mb = mean[mb].to(torch.float32)
            rstd_mb = rstd[mb].to(torch.float32)

            x_hat = (x_mb - mean_mb[:, None]) * rstd_mb[:, None]

            grad_w_acc += torch.sum(dy_mb * x_hat, dim=0)
            if compute_bias_grad:
                grad_b_acc += torch.sum(dy_mb, dim=0)  # pyright: ignore[reportPossiblyUnboundVariable]

            wdy = weight_cta * dy_mb
            c1 = torch.sum(x_hat * wdy, dim=-1) / n
            c2 = torch.sum(wdy, dim=-1) / n
            dx = (wdy - (x_hat * c1[:, None] + c2[:, None])) * rstd_mb[:, None]
            grad_x[mb, :] = dx.to(x.dtype)

        grad_weight_blocks[mb_cta.id, :] = grad_w_acc
        if compute_bias_grad:
            grad_bias_blocks[mb_cta.id, :] = grad_b_acc  # type: ignore[index]

    grad_weight = grad_weight_blocks.sum(0).to(weight.dtype)
    if compute_bias_grad:
        grad_bias = grad_bias_blocks.sum(0).to(weight.dtype)
        return grad_x, grad_weight, grad_bias
    return grad_x, grad_weight, None

[umechand-amd]

Can you also share the inuductor generated triton backward kernel? The kernel you shared above is for forward pass.

[umechand-amd]

@MengjiaoZhou I am not able to exactly reproduce the issue but I did see when bias is added the latency is around 0.65 ms for the shape you mentioned. Can you please also share the triton kernel generated by inductor for the backward pass so that I can do a comparison on my end and analyze why helion cannot generate a similar kernel? Nevermind, your comment just arrived.

[MengjiaoZhou]

Can you also share the inuductor generated triton backward kernel? The kernel you shared above is for forward pass.

my bad. here's the indcutor generated triton backward kernel:

# AOT ID: ['5_backward']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from cmath import nanj
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
assert_alignment = torch._C._dynamo.guards.assert_alignment
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cpu_pinned = torch._C._dynamo.guards._empty_strided_cpu_pinned
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
empty_strided_mtia = torch._C._dynamo.guards._empty_strided_mtia
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()
empty_strided_p2p = torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p


# kernel path: /var/tmp/torchinductor_mengjiao/hu/chujqchuvrt7l6f2pqwi4yldlm3z7upriwfnuouyt76mcubte4fj.py
# Topologically Sorted Source Nodes: [convert_element_type_2, layer_norm, mul_8, sum_3, sum_4, convert_element_type_7, convert_element_type_8], Original ATen: [aten.native_layer_norm_backward, aten.native_layer_norm]
# Source node to ATen node mapping:
#   convert_element_type_2 => convert_element_type_2
#   convert_element_type_7 => convert_element_type_7
#   convert_element_type_8 => convert_element_type_8
#   layer_norm => convert_element_type, mul, sub
#   mul_8 => mul_8
#   sum_3 => sum_3
#   sum_4 => sum_4
# Graph fragment:
#   %tangents_1 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0" = PlaceHolder[target=tangents_1]
#   %primals_1 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0" = PlaceHolder[target=primals_1]
#   %getitem_1 : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0" = PlaceHolder[target=getitem_1]
#   %rsqrt : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0" = PlaceHolder[target=rsqrt]
#   %sum_3 : Tensor "f32[36864][1]cuda:0" = PlaceHolder[target=sum_3]
#   %sum_4 : Tensor "f32[36864][1]cuda:0" = PlaceHolder[target=sum_4]
#   %convert_element_type_2 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=3] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%tangents_1, torch.float32), kwargs = {})
#   %convert_element_type : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%primals_1, torch.float32), kwargs = {})
#   %sub : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%convert_element_type, %getitem_1), kwargs = {})
#   %mul : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=3] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sub, %rsqrt), kwargs = {})
#   %mul_8 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%convert_element_type_2, %mul), kwargs = {})
#   %sum_3 : Tensor "f32[36864][1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%mul_8, [0, 1]), kwargs = {})
#   %sum_4 : Tensor "f32[36864][1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%convert_element_type_2, [0, 1]), kwargs = {})
#   %convert_element_type_7 : Tensor "bf16[36864][1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%sum_3, torch.bfloat16), kwargs = {})
#   %convert_element_type_8 : Tensor "bf16[36864][1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%sum_4, torch.bfloat16), kwargs = {})
#   return %sum_3,%sum_4,%convert_element_type_7,%convert_element_type_8
triton_red_fused_native_layer_norm_native_layer_norm_backward_0 = async_compile.triton('triton_red_fused_native_layer_norm_native_layer_norm_backward_0', '''
import triton
import triton.language as tl

import triton.language.extra.tlx as tlx  # noqa: F401

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()

@triton_heuristics.reduction(
    size_hints={'x': 65536, 'r0_': 4096},
    reduction_hint=ReductionHint.DEFAULT,
    filename=__file__,
    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*fp32', 'in_ptr3': '*fp32', 'out_ptr2': '*bf16', 'out_ptr3': '*bf16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='hip', index=0, multi_processor_count=256, cc='gfx950', major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, max_threads_per_block=1024, warp_size=64), 'constants': {}, 'native_matmul': False, 'enable_fp_fusion': True, 'launch_pdl': False, 'disable_ftz': False, 'configs': [{(0,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (1,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (2,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (3,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (4,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (5,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (6,): [['tt.divisibility', 16]], (7,): [['tt.divisibility', 16]]}]},
    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_red_fused_native_layer_norm_native_layer_norm_backward_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'atomic_add_found': False, 'num_load': 4, 'num_store': 2, 'num_reduction': 2, 'backend_hash': 'E3565433EC5266B9E41D740243879373C6297D57D928F6BAE46DD3A44A623A50', 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 32, 'store_cubin': False, 'deterministic': False, 'force_filter_reduction_configs': False, 'mix_order_reduction_allow_multi_stages': False, 'are_deterministic_algorithms_enabled': False, 'is_hip': True, 'is_fbcode': True}
)
@triton.jit
def triton_red_fused_native_layer_norm_native_layer_norm_backward_0(in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr2, out_ptr3, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
    xnumel = 36864
    r0_numel = 2304
    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    r0_base = tl.arange(0, R0_BLOCK)[None, :]
    rbase = r0_base
    x0 = xindex
    _tmp10 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
    _tmp13 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages = 2):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp0 = tl.load(in_ptr0 + (x0 + 36864*r0_1), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp2 = tl.load(in_ptr1 + (x0 + 36864*r0_1), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp4 = tl.load(in_ptr2 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0)
        tmp6 = tl.load(in_ptr3 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0)
        tmp1 = tmp0.to(tl.float32)
        tmp3 = tmp2.to(tl.float32)
        tmp5 = tmp3 - tmp4
        tmp7 = tmp5 * tmp6
        tmp8 = tmp1 * tmp7
        tmp9 = tl.broadcast_to(tmp8, [XBLOCK, R0_BLOCK])
        tmp11 = _tmp10 + tmp9
        _tmp10 = tl.where(r0_mask & xmask, tmp11, _tmp10)
        tmp12 = tl.broadcast_to(tmp1, [XBLOCK, R0_BLOCK])
        tmp14 = _tmp13 + tmp12
        _tmp13 = tl.where(r0_mask & xmask, tmp14, _tmp13)
    tmp10 = tl.sum(_tmp10, 1)[:, None]
    tmp13 = tl.sum(_tmp13, 1)[:, None]
    tmp15 = tmp10.to(tl.float32)
    tmp16 = tmp13.to(tl.float32)
    tl.store(out_ptr2 + (x0), tmp15, xmask)
    tl.store(out_ptr3 + (x0), tmp16, xmask)
''', device_str='cuda')


# kernel path: /var/tmp/torchinductor_mengjiao/lx/clxdv4x6alemmiza6sgjlmnzlyzrlaud3o7f7v4m5khuqollg22q.py
# Topologically Sorted Source Nodes: [convert_element_type_2, convert_element_type_4, mul_3, mul_4, sum_1, layer_norm, mul_5, sum_2, mul_6, sub_2, sub_3, div, mul_7, convert_element_type_6], Original ATen: [aten.native_layer_norm_backward, aten.native_layer_norm]
# Source node to ATen node mapping:
#   convert_element_type_2 => convert_element_type_2
#   convert_element_type_4 => convert_element_type_4
#   convert_element_type_6 => convert_element_type_6
#   div => div
#   layer_norm => convert_element_type, mul, sub
#   mul_3 => mul_3
#   mul_4 => mul_4
#   mul_5 => mul_5
#   mul_6 => mul_6
#   mul_7 => mul_7
#   sub_2 => sub_2
#   sub_3 => sub_3
#   sum_1 => sum_1
#   sum_2 => sum_2
# Graph fragment:
#   %tangents_1 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0" = PlaceHolder[target=tangents_1]
#   %primals_2 : Tensor "bf16[36864][1]cuda:0" = PlaceHolder[target=primals_2]
#   %primals_1 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0" = PlaceHolder[target=primals_1]
#   %getitem_1 : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0" = PlaceHolder[target=getitem_1]
#   %rsqrt : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0" = PlaceHolder[target=rsqrt]
#   %sum_1 : Tensor "f32[1, 2304, 1][2304, 1, 2304]cuda:0" = PlaceHolder[target=sum_1]
#   %sum_2 : Tensor "f32[1, 2304, 1][2304, 1, 2304]cuda:0" = PlaceHolder[target=sum_2]
#   %convert_element_type_2 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=3] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%tangents_1, torch.float32), kwargs = {})
#   %convert_element_type_4 : Tensor "f32[36864][1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%primals_2, torch.float32), kwargs = {})
#   %mul_3 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=3] = call_function[target=torch.ops.aten.mul.Tensor](args = (%convert_element_type_2, %convert_element_type_4), kwargs = {})
#   %mul_4 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_3, 36864), kwargs = {})
#   %sum_1 : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%mul_3, [2], True), kwargs = {})
#   %convert_element_type : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%primals_1, torch.float32), kwargs = {})
#   %sub : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%convert_element_type, %getitem_1), kwargs = {})
#   %mul : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=3] = call_function[target=torch.ops.aten.mul.Tensor](args = (%sub, %rsqrt), kwargs = {})
#   %mul_5 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul_3, %mul), kwargs = {})
#   %sum_2 : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%mul_5, [2], True), kwargs = {})
#   %mul_6 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%mul, %sum_2), kwargs = {})
#   %sub_2 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%mul_4, %sum_1), kwargs = {})
#   %sub_3 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%sub_2, %mul_6), kwargs = {})
#   %div : Tensor "f32[1, 2304, 1][2304, 1, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%rsqrt, 36864), kwargs = {})
#   %mul_7 : Tensor "f32[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%div, %sub_3), kwargs = {})
#   %convert_element_type_6 : Tensor "bf16[1, 2304, 36864][84934656, 36864, 1]cuda:0"[num_users=1] = call_function[target=torch.ops.prims.convert_element_type.default](args = (%mul_7, torch.bfloat16), kwargs = {})
#   return %sum_1,%sum_2,%convert_element_type_6
triton_red_fused_native_layer_norm_native_layer_norm_backward_1 = async_compile.triton('triton_red_fused_native_layer_norm_native_layer_norm_backward_1', '''
import triton
import triton.language as tl

import triton.language.extra.tlx as tlx  # noqa: F401

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()

@triton_heuristics.reduction(
    size_hints={'x': 4096, 'r0_': 65536},
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'in_ptr3': '*fp32', 'in_ptr4': '*fp32', 'out_ptr2': '*bf16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='hip', index=0, multi_processor_count=256, cc='gfx950', major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, max_threads_per_block=1024, warp_size=64), 'constants': {}, 'native_matmul': False, 'enable_fp_fusion': True, 'launch_pdl': False, 'disable_ftz': False, 'configs': [{(0,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (1,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (2,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (3,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (4,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (5,): [['tt.divisibility', 16], ['tt.pointer_range', 32]], (6,): [['tt.divisibility', 16]], (7,): [['tt.divisibility', 16]]}]},
    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_red_fused_native_layer_norm_native_layer_norm_backward_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'atomic_add_found': False, 'num_load': 8, 'num_store': 1, 'num_reduction': 2, 'backend_hash': 'E3565433EC5266B9E41D740243879373C6297D57D928F6BAE46DD3A44A623A50', 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 32, 'store_cubin': False, 'deterministic': False, 'force_filter_reduction_configs': False, 'mix_order_reduction_allow_multi_stages': False, 'are_deterministic_algorithms_enabled': False, 'is_hip': True, 'is_fbcode': True}
)
@triton.jit
def triton_red_fused_native_layer_norm_native_layer_norm_backward_1(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
    xnumel = 2304
    r0_numel = 36864
    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    r0_base = tl.arange(0, R0_BLOCK)[None, :]
    rbase = r0_base
    x0 = xindex
    _tmp6 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
    tmp10 = tl.load(in_ptr3 + (x0), xmask, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr4 + (x0), xmask, eviction_policy='evict_last')
    _tmp16 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages = 2):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp0 = tl.load(in_ptr0 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp2 = tl.load(in_ptr1 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp8 = tl.load(in_ptr2 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp1 = tmp0.to(tl.float32)
        tmp3 = tmp2.to(tl.float32)
        tmp4 = tmp1 * tmp3
        tmp5 = tl.broadcast_to(tmp4, [XBLOCK, R0_BLOCK])
        tmp7 = _tmp6 + tmp5
        _tmp6 = tl.where(r0_mask & xmask, tmp7, _tmp6)
        tmp9 = tmp8.to(tl.float32)
        tmp11 = tmp9 - tmp10
        tmp13 = tmp11 * tmp12
        tmp14 = tmp4 * tmp13
        tmp15 = tl.broadcast_to(tmp14, [XBLOCK, R0_BLOCK])
        tmp17 = _tmp16 + tmp15
        _tmp16 = tl.where(r0_mask & xmask, tmp17, _tmp16)
    tmp6 = tl.sum(_tmp6, 1)[:, None]
    tmp16 = tl.sum(_tmp16, 1)[:, None]
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages = 2):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp20 = tl.load(in_ptr0 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp22 = tl.load(in_ptr1 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp28 = tl.load(in_ptr2 + (r0_1 + 36864*x0), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp18 = tl.full([1, 1], 2.712673611111111e-05, tl.float32)
        tmp19 = tmp12 * tmp18
        tmp21 = tmp20.to(tl.float32)
        tmp23 = tmp22.to(tl.float32)
        tmp24 = tmp21 * tmp23
        tmp25 = tl.full([1, 1], 36864.0, tl.float32)
        tmp26 = tmp24 * tmp25
        tmp27 = tmp26 - tmp6
        tmp29 = tmp28.to(tl.float32)
        tmp30 = tmp29 - tmp10
        tmp31 = tmp30 * tmp12
        tmp32 = tmp31 * tmp16
        tmp33 = tmp27 - tmp32
        tmp34 = tmp19 * tmp33
        tmp35 = tmp34.to(tl.float32)
        tl.store(out_ptr2 + (r0_1 + 36864*x0), tmp35, r0_mask & xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    primals_1, primals_2, getitem_1, rsqrt, tangents_1 = args
    args.clear()
    assert_size_stride(primals_1, (1, 2304, 36864), (84934656, 36864, 1))
    assert_size_stride(primals_2, (36864, ), (1, ))
    assert_size_stride(getitem_1, (1, 2304, 1), (2304, 1, 1))
    assert_size_stride(rsqrt, (1, 2304, 1), (2304, 1, 1))
    assert_size_stride(tangents_1, (1, 2304, 36864), (84934656, 36864, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf5 = empty_strided_cuda((36864, ), (1, ), torch.bfloat16)
        buf6 = empty_strided_cuda((36864, ), (1, ), torch.bfloat16)
        # Topologically Sorted Source Nodes: [convert_element_type_2, layer_norm, mul_8, sum_3, sum_4, convert_element_type_7, convert_element_type_8], Original ATen: [aten.native_layer_norm_backward, aten.native_layer_norm]
        stream0 = get_raw_stream(0)
        triton_red_fused_native_layer_norm_native_layer_norm_backward_0.run(tangents_1, primals_1, getitem_1, rsqrt, buf5, buf6, 36864, 2304, stream=stream0)
        buf4 = empty_strided_cuda((1, 2304, 36864), (84934656, 36864, 1), torch.bfloat16)
        # Topologically Sorted Source Nodes: [convert_element_type_2, convert_element_type_4, mul_3, mul_4, sum_1, layer_norm, mul_5, sum_2, mul_6, sub_2, sub_3, div, mul_7, convert_element_type_6], Original ATen: [aten.native_layer_norm_backward, aten.native_layer_norm]
        stream0 = get_raw_stream(0)
        triton_red_fused_native_layer_norm_native_layer_norm_backward_1.run(tangents_1, primals_2, primals_1, getitem_1, rsqrt, buf4, 2304, 36864, stream=stream0)
        del getitem_1
        del primals_1
        del primals_2
        del rsqrt
        del tangents_1
    return (buf4, buf5, buf6, )


def get_args():
    from torch._dynamo.testing import rand_strided
    primals_1 = rand_strided((1, 2304, 36864), (84934656, 36864, 1), device='cuda:0', dtype=torch.bfloat16)
    primals_2 = rand_strided((36864, ), (1, ), device='cuda:0', dtype=torch.bfloat16)
    getitem_1 = rand_strided((1, 2304, 1), (2304, 1, 1), device='cuda:0', dtype=torch.float32)
    rsqrt = rand_strided((1, 2304, 1), (2304, 1, 1), device='cuda:0', dtype=torch.float32)
    tangents_1 = rand_strided((1, 2304, 36864), (84934656, 36864, 1), device='cuda:0', dtype=torch.bfloat16)
    return [primals_1, primals_2, getitem_1, rsqrt, tangents_1]


def benchmark_compiled_module(args, times=10, repeat=10):
    from torch._inductor.utils import print_performance
    fn = lambda: call(list(args))
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    args = get_args()
    compiled_module_main('helion_layernorm,inductor_layernorm', lambda times, repeat: benchmark_compiled_module(args, times=times, repeat=repeat))

[umechand-amd]

@MengjiaoZhou
I looked at the inductor kernel and the helion kernel you shared and it seems there is a structural difference between the two which is the helion kernel computes dX and dW/dB together and writes block partials reductions first, the inductor however reduces over the M dimension and writes N outputs directly for DW/Db and computes the input gradients dx kernel.

I was able to rewrite Helion by splitting the operation across two Helion kernels to match the inductor structure and I was able to get latency of 0.233 ms for fp16 dtype.

Can you please take a look at the code below and let me know if you have any questions?

`from future import annotations
import argparse

import torch
import torch.nn.functional as F
from triton.testing import do_bench

import helion
from helion._testing import DEVICE
from helion._testing import HALF_DTYPE
import helion.language as hl

DWDB_HARDCODED_CONFIG = helion.Config(
block_sizes=[16, 128],
indexing=["pointer", "pointer", "block_ptr", "pointer", "pointer", "block_ptr"],
load_eviction_policies=["", "", "", ""],
matrix_instr_nonkdim=32,
num_sm_multiplier=128,
num_stages=1,
num_warps=2,
pid_type="persistent_blocked",
range_flattens=[None, None, True],
range_multi_buffers=[True, True, True],
range_num_stages=[3, 0, 3],
range_unroll_factors=[0, 2, 2],
range_warp_specializes=[],
waves_per_eu=1,
)

DX_HARDCODED_CONFIG = helion.Config(
block_sizes=[1],
indexing=[
"block_ptr",
"pointer",
"pointer",
"block_ptr",
"pointer",
"block_ptr",
"pointer",
"pointer",
"pointer",
"block_ptr",
"block_ptr",
"block_ptr",
"block_ptr",
"block_ptr",
],
load_eviction_policies=["", "", "", "", "", "", "", "", "", "", "", ""],
matrix_instr_nonkdim=32,
num_sm_multiplier=4,
num_stages=1,
num_warps=8,
pid_type="persistent_interleaved",
range_flattens=[True, True],
range_multi_buffers=[None, None],
range_num_stages=[4, 1],
range_unroll_factors=[4, 2],
range_warp_specializes=[],
waves_per_eu=1,
)

@helion.kernel(config=DWDB_HARDCODED_CONFIG, static_shapes=True)
def layer_norm_bwd_dwdb_kernel_hardcoded(
grad_out: torch.Tensor,
x: torch.Tensor,
mean: torch.Tensor,
rstd: torch.Tensor,
weight: torch.Tensor,
compute_bias_grad: hl.constexpr = True, # type: ignore[valid-type]
) -> tuple[torch.Tensor, torch.Tensor | None]:
m_block = hl.register_block_size(1, 16)
n_block = hl.register_block_size(128, 4096)

grad_weight = torch.empty_like(weight)
if compute_bias_grad:
    grad_bias = torch.empty_like(weight)

for tile_n in hl.tile(x.size(1), block_size=n_block):
    grad_w_acc = weight[tile_n].new_zeros(weight[tile_n].shape, dtype=torch.float32)
    if compute_bias_grad:
        grad_b_acc = weight[tile_n].new_zeros(
            weight[tile_n].shape, dtype=torch.float32
        )

    for mb_cta in hl.tile(x.size(0), block_size=m_block):
        for mb in hl.tile(mb_cta.begin, mb_cta.end, block_size=1):
            x_mb = x[mb, tile_n].to(torch.float32)
            dy_mb = grad_out[mb, tile_n].to(torch.float32)
            mean_mb = mean[mb].to(torch.float32)
            rstd_mb = rstd[mb].to(torch.float32)

            x_hat = (x_mb - mean_mb[:, None]) * rstd_mb[:, None]
            grad_w_acc += torch.sum(dy_mb * x_hat, dim=0)
            if compute_bias_grad:
                # pyrefly: ignore [unbound-name]
                grad_b_acc += torch.sum(dy_mb, dim=0)

    grad_weight[tile_n] = grad_w_acc.to(weight.dtype)
    if compute_bias_grad:
        grad_bias[tile_n] = grad_b_acc.to(weight.dtype)

if compute_bias_grad:
    return grad_weight, grad_bias
return grad_weight, None

@helion.kernel(config=DX_HARDCODED_CONFIG, static_shapes=True)
def layer_norm_bwd_dx_kernel_hardcoded(
grad_out: torch.Tensor,
x: torch.Tensor,
mean: torch.Tensor,
rstd: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
m_block = hl.register_block_size(1, 16)
n = hl.specialize(x.size(1))

grad_x = torch.empty_like(x)
for mb_cta in hl.tile(x.size(0), block_size=m_block):
    weight_cta = weight[None, :].to(torch.float32)
    for mb in hl.tile(mb_cta.begin, mb_cta.end, block_size=1):
        x_mb = x[mb, :].to(torch.float32)
        dy_mb = grad_out[mb, :].to(torch.float32)
        mean_mb = mean[mb].to(torch.float32)
        rstd_mb = rstd[mb].to(torch.float32)

        x_hat = (x_mb - mean_mb[:, None]) * rstd_mb[:, None]
        wdy = weight_cta * dy_mb
        c1 = torch.sum(x_hat * wdy, dim=-1) / n
        c2 = torch.sum(wdy, dim=-1) / n
        dx = (wdy - (x_hat * c1[:, None] + c2[:, None])) * rstd_mb[:, None]
        grad_x[mb, :] = dx.to(x.dtype)
return grad_x

def layer_norm_bwd_two_kernel_hardcoded(
grad_out: torch.Tensor,
x: torch.Tensor,
mean: torch.Tensor,
rstd: torch.Tensor,
weight: torch.Tensor,
compute_bias_grad: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
grad_weight, grad_bias = layer_norm_bwd_dwdb_kernel_hardcoded(
grad_out, x, mean, rstd, weight, compute_bias_grad
)
grad_x = layer_norm_bwd_dx_kernel_hardcoded(grad_out, x, mean, rstd, weight)
return grad_x, grad_weight, grad_bias

def estimate_layer_norm_bwd_bytes(
m: int, n: int, dtype: torch.dtype, compute_bias_grad: bool
) -> int:
elem = torch.empty((), dtype=dtype).element_size()
f32 = torch.empty((), dtype=torch.float32).element_size()

read_x_dwdb = m * n * elem
read_grad_out_dwdb = m * n * elem
read_mean_dwdb = m * f32
read_rstd_dwdb = m * f32
write_grad_weight = n * elem
write_grad_bias = n * elem if compute_bias_grad else 0

read_x_dx = m * n * elem
read_grad_out_dx = m * n * elem
read_weight_dx = m * n * elem
read_mean_dx = m * f32
read_rstd_dx = m * f32
write_grad_x = m * n * elem

return (
    read_x_dwdb
    + read_grad_out_dwdb
    + read_mean_dwdb
    + read_rstd_dwdb
    + write_grad_weight
    + write_grad_bias
    + read_x_dx
    + read_grad_out_dx
    + read_weight_dx
    + read_mean_dx
    + read_rstd_dx
    + write_grad_x
)

def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float:
return (a.float() - b.float()).abs().max().item()

def torch_layer_norm_bwd_reference(
grad_out: torch.Tensor,
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
x_ref = x.detach().to(torch.float32).requires_grad_(True)
w_ref = weight.detach().to(torch.float32).requires_grad_(True)
b_ref = (
bias.detach().to(torch.float32).requires_grad_(True) if bias is not None else None
)
y_ref = F.layer_norm(x_ref, [x.shape[1]], w_ref, b_ref, eps)
y_ref.backward(grad_out.to(torch.float32))
return (
x_ref.grad.to(x.dtype),
w_ref.grad.to(weight.dtype),
b_ref.grad.to(weight.dtype) if b_ref is not None else None,
)

def main() -> None:
parser = argparse.ArgumentParser(
description="Run hardcoded-config LayerNorm backward two-kernel rewrite."
)
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--rows", type=int, default=2304)
parser.add_argument("--cols", type=int, default=36864)
parser.add_argument("--eps", type=float, default=1e-4)
parser.add_argument("--compute-bias-grad", action="store_true")
parser.add_argument("--skip-accuracy-check", action="store_true")
args = parser.parse_args()

if DEVICE.type == "cuda":
    torch.cuda.set_device(args.device)
    device = torch.device(f"cuda:{args.device}")
else:
    device = DEVICE

dtype = HALF_DTYPE
m, n = args.rows, args.cols
eps = args.eps
compute_bias_grad = args.compute_bias_grad

print(
    f"\n=== hardcoded layer_norm_bwd two-kernel (M, N)=({m}, {n}) "
    f"compute_bias_grad={compute_bias_grad} ==="
)
x = -2.3 + 0.5 * torch.randn([m, n], device=device, dtype=dtype)
grad_out = torch.randn([m, n], device=device, dtype=dtype)
weight = torch.randn([n], device=device, dtype=dtype)
bias = torch.randn([n], device=device, dtype=dtype) if compute_bias_grad else None

x_f32 = x.to(torch.float32)
mean = x_f32.mean(dim=-1)
var = torch.mean((x_f32 - mean[:, None]) * (x_f32 - mean[:, None]), dim=-1)
rstd = torch.rsqrt(var + eps)

if not args.skip_accuracy_check:
    ref_dx, ref_dw, ref_db = torch_layer_norm_bwd_reference(
        grad_out, x, weight, bias, eps
    )
    out_dx, out_dw, out_db = layer_norm_bwd_two_kernel_hardcoded(
        grad_out, x, mean, rstd, weight, compute_bias_grad
    )
    print(
        "max_abs_vs_torch:",
        f"dx={_max_abs_diff(out_dx, ref_dx):.6e}",
        f"dw={_max_abs_diff(out_dw, ref_dw):.6e}",
        (
            f"db={_max_abs_diff(out_db, ref_db):.6e}"
            if (compute_bias_grad and out_db is not None and ref_db is not None)
            else "db=None"
        ),
    )

latency_ms = do_bench(
    lambda: layer_norm_bwd_two_kernel_hardcoded(
        grad_out, x, mean, rstd, weight, compute_bias_grad
    )
)
bytes_total = estimate_layer_norm_bwd_bytes(m, n, dtype, compute_bias_grad)
gbps = bytes_total / (latency_ms * 1e-3) / 1e9

grad_x, grad_weight, grad_bias = layer_norm_bwd_two_kernel_hardcoded(
    grad_out, x, mean, rstd, weight, compute_bias_grad
)
checksum = grad_x.float().sum() + grad_weight.float().sum()
if grad_bias is not None:
    checksum = checksum + grad_bias.float().sum()

print(f"hardcoded_config_dwdb: {DWDB_HARDCODED_CONFIG}")
print(f"hardcoded_config_dx: {DX_HARDCODED_CONFIG}")
print(f"latency_ms (median): {latency_ms:.4f}")
print(f"effective_bandwidth_gbps (estimated): {gbps:.2f}")
print(f"estimated_bytes_per_call: {bytes_total}")
print(f"checksum: {checksum.item():.6f}")

if name == "main":
main()
`

[MengjiaoZhou]

I tested with tritonbench, but the new bwd implementation seems much worser than current implementation.

      (B, M, N)    inductor_layernorm-latency    inductor_layernorm-gbps    inductor_layernorm-compile_time    helion_layernorm-latency    helion_layernorm-gbps    helion_layernorm-accuracy    helion_layernorm-compile_time
----------------  ----------------------------  -------------------------  ---------------------------------  --------------------------  -----------------------  ---------------------------  -------------------------------
(1, 2304, 36864)             0.249059 (±0.00%)                    2046.73                            126.527           1.396376 (±0.00%)                  365.056                            1                          1772.17
         average                                                  2046.73                            126.527                                              365.056                            1                          1772.17

Here's the helion generated code:

from __future__ import annotations

import torch
import helion
import helion.language as hl
import triton
import triton.language as tl
from helion.runtime import default_launcher as _default_launcher

_BLOCK_SIZE_1 = tl.constexpr(128)
_BLOCK_SIZE_0 = tl.constexpr(16)
_BLOCK_SIZE_2 = tl.constexpr(1)

@triton.jit
def _helion_layer_norm_bwd_dwdb_kernel_hardcoded(x, grad_out, mean, rstd, grad_weight, grad_bias, _NUM_SM: tl.constexpr):
    # src[layer_norm.py:2548]: for tile_n in hl.tile(x.size(1), block_size=n_block):
    # src[layer_norm.py:2549]:     grad_w_acc = weight[tile_n].new_zeros(weight[tile_n].shape, dtype=torch.float32)
    # src[layer_norm.py:2550]:     if compute_bias_grad:
    # src[layer_norm.py:2548-2570]: ...
    total_pids = tl.cdiv(36864, _BLOCK_SIZE_1)
    block_size = tl.cdiv(total_pids, _NUM_SM * 128)
    start_pid = tl.program_id(0) * block_size
    end_pid = start_pid + block_size
    if end_pid > total_pids:
        end_pid = total_pids
    for virtual_pid in tl.range(start_pid, end_pid, num_stages=3, disallow_acc_multi_buffer=False):
        # src[layer_norm.py:2548]: for tile_n in hl.tile(x.size(1), block_size=n_block):
        pid_0 = virtual_pid
        offset_1 = pid_0 * _BLOCK_SIZE_1
        indices_1 = (offset_1 + tl.arange(0, _BLOCK_SIZE_1)).to(tl.int32)
        # src[layer_norm.py:2549]: grad_w_acc = weight[tile_n].new_zeros(weight[tile_n].shape, dtype=torch.float32)
        grad_w_acc = tl.full([_BLOCK_SIZE_1], 0, tl.float32)
        # src[layer_norm.py:2551]: grad_b_acc = weight[tile_n].new_zeros(
        # src[layer_norm.py:2552]:     weight[tile_n].shape, dtype=torch.float32
        # src[layer_norm.py:2553]: )
        grad_b_acc = tl.full([_BLOCK_SIZE_1], 0, tl.float32)
        # src[layer_norm.py:2555]: for mb_cta in hl.tile(x.size(0), block_size=m_block):
        # src[layer_norm.py:2556]:     for mb in hl.tile(mb_cta.begin, mb_cta.end, block_size=1):
        # src[layer_norm.py:2557]:         x_mb = x[mb, tile_n].to(torch.float32)
        # src[layer_norm.py:2555-2566]: ...
        for offset_0 in tl.range(0, 2304, _BLOCK_SIZE_0, loop_unroll_factor=2, disallow_acc_multi_buffer=False):
            grad_w_acc_copy = grad_w_acc
            grad_b_acc_copy = grad_b_acc
            grad_w_acc = grad_w_acc_copy
            grad_b_acc = grad_b_acc_copy
            # src[layer_norm.py:2556]: for mb in hl.tile(mb_cta.begin, mb_cta.end, block_size=1):
            tile_end = offset_0 + _BLOCK_SIZE_0
            # src[layer_norm.py:2556]: for mb in hl.tile(mb_cta.begin, mb_cta.end, block_size=1):
            # src[layer_norm.py:2557]:     x_mb = x[mb, tile_n].to(torch.float32)
            # src[layer_norm.py:2558]:     dy_mb = grad_out[mb, tile_n].to(torch.float32)
            # src[layer_norm.py:2556-2566]: ...
            for offset_2 in tl.range(tl.cast(offset_0, tl.int32), tl.cast(tile_end, tl.int32), loop_unroll_factor=2, num_stages=3, disallow_acc_multi_buffer=False, flatten=True):
                indices_2 = offset_2 + tl.arange(0, 1).to(tl.int32)
                grad_w_acc_copy_0_copy = grad_w_acc
                grad_b_acc_copy_0_copy = grad_b_acc
                grad_w_acc_copy_0_copy_0 = grad_w_acc_copy_0_copy
                grad_b_acc_copy_0_copy_0 = grad_b_acc_copy_0_copy
                # src[layer_norm.py:2557]: x_mb = x[mb, tile_n].to(torch.float32)
                load = tl.load(x + (indices_2[:, None] * 36864 + indices_1[None, :] * 1), None)
                v_0 = tl.cast(load, tl.float32)
                # src[layer_norm.py:2558]: dy_mb = grad_out[mb, tile_n].to(torch.float32)
                load_1 = tl.load(grad_out + (indices_2[:, None] * 36864 + indices_1[None, :] * 1), None)
                v_1 = tl.cast(load_1, tl.float32)
                # src[layer_norm.py:2559]: mean_mb = mean[mb].to(torch.float32)
                mean_mb = tl.load(tl.make_block_ptr(mean, [2304], [1], [offset_2], [_BLOCK_SIZE_2], [0]), boundary_check=[0], padding_option='zero')
                # src[layer_norm.py:2560]: rstd_mb = rstd[mb].to(torch.float32)
                rstd_mb = tl.load(rstd + indices_2 * 1, None)
                # src[layer_norm.py:2562]: x_hat = (x_mb - mean_mb[:, None]) * rstd_mb[:, None]
                subscript = mean_mb[:, None]
                v_2 = v_0 - subscript
                subscript_1 = rstd_mb[:, None]
                v_3 = v_2 * subscript_1
                # src[layer_norm.py:2563]: grad_w_acc += torch.sum(dy_mb * x_hat, dim=0)
                v_4 = v_1 * v_3
                sum_1 = tl.cast(tl.sum(v_4, 0), tl.float32)
                grad_w_acc = grad_w_acc_copy_0_copy_0 + sum_1
                # src[layer_norm.py:2566]: grad_b_acc += torch.sum(dy_mb, dim=0)
                sum_2 = tl.cast(tl.sum(v_1, 0), tl.float32)
                grad_b_acc = grad_b_acc_copy_0_copy_0 + sum_2
        # src[layer_norm.py:2568]: grad_weight[tile_n] = grad_w_acc.to(weight.dtype)
        v_7 = tl.cast(grad_w_acc, tl.bfloat16)
        tl.store(grad_weight + indices_1 * 1, v_7, None)
        # src[layer_norm.py:2570]: grad_bias[tile_n] = grad_b_acc.to(weight.dtype)
        v_8 = tl.cast(grad_b_acc, tl.bfloat16)
        tl.store(tl.make_block_ptr(grad_bias, [36864], [1], [offset_1], [_BLOCK_SIZE_1], [0]), tl.cast(v_8, tl.bfloat16), boundary_check=[0])

def layer_norm_bwd_dwdb_kernel_hardcoded(grad_out: torch.Tensor, x: torch.Tensor, mean: torch.Tensor, rstd: torch.Tensor, weight: torch.Tensor, compute_bias_grad: hl.constexpr=True, *, _launcher=_default_launcher):
    # src[layer_norm.py:2544]: grad_weight = torch.empty_like(weight)
    grad_weight = torch.empty_like(weight)
    # src[layer_norm.py:2545]: if compute_bias_grad:
    # src[layer_norm.py:2546]:     grad_bias = torch.empty_like(weight)
    if True:
        # src[layer_norm.py:2546]: grad_bias = torch.empty_like(weight)
        grad_bias = torch.empty_like(weight)
    # src[layer_norm.py:2548]: for tile_n in hl.tile(x.size(1), block_size=n_block):
    _NUM_SM = helion.runtime.get_num_sm(grad_out.device)
    # src[layer_norm.py:2548]: for tile_n in hl.tile(x.size(1), block_size=n_block):
    # src[layer_norm.py:2549]:     grad_w_acc = weight[tile_n].new_zeros(weight[tile_n].shape, dtype=torch.float32)
    # src[layer_norm.py:2550]:     if compute_bias_grad:
    # src[layer_norm.py:2548-2570]: ...
    _launcher(_helion_layer_norm_bwd_dwdb_kernel_hardcoded, (_NUM_SM * 128,), x, grad_out, mean, rstd, grad_weight, grad_bias, _NUM_SM, num_warps=2, num_stages=1, matrix_instr_nonkdim=32, waves_per_eu=1)
    # src[layer_norm.py:2572]: if compute_bias_grad:
    # src[layer_norm.py:2573]:     return grad_weight, grad_bias
    if True:
        # src[layer_norm.py:2573]: return grad_weight, grad_bias
        return (grad_weight, grad_bias)
    # src[layer_norm.py:2574]: return grad_weight, None
    return (grad_weight, None)

def call():
    from torch._dynamo.testing import rand_strided
    # src[layer_norm.py:2533]: def layer_norm_bwd_dwdb_kernel_hardcoded(
    # src[layer_norm.py:2534]:     grad_out: torch.Tensor,
    # src[layer_norm.py:2535]:     x: torch.Tensor,
    # src[layer_norm.py:2533-2574]: ...
    grad_out = rand_strided(size=(2304, 36864), stride=(36864, 1), dtype=torch.bfloat16, device='cuda:0')
    x = rand_strided(size=(2304, 36864), stride=(36864, 1), dtype=torch.bfloat16, device='cuda:0')