Skip to content

[Pytorch] layernorm+softmax fusion from inductor has lower hardware efficiency than CUDA #7571

Open
#7577
Open
[Pytorch] layernorm+softmax fusion from inductor has lower hardware efficiency than CUDA#7571
#7577
Labels
performanceupstream: pytorch

Description

@LuFinch
LuFinch
opened on Jul 24, 2026

Describe the issue

Pytorch XPU observes layernorm+softmax fusion in a OOB model.

We compare B70 with A10 on hardware efficiency. Layernorm+softmax fusion can be deem as memory bound so that we measured memory bandwidth vs peak bandwidth as efficiency.

On B70, below reproducer reach 206GB/s, reach 34%(vs 608GB/s peak) or 39%(vs 532GB/s measured) efficiency.
On A10, below reproducer reach 268GB/s, reach 46%(vs 600GB/s peak) or 55%(vs 480GB/s measured) efficiency.

From stall sampling, this kernel is memory bound and has high pressure on load/store.

Stall Type Counter Ratio of Stall
SendStall 65,113 49.4%
SbidStall 42,803 32.5%
PipeStall 9,691 7.4%
DistStall 8,280 6.3%
InstrFetchStall 3,634 2.8%

B70 load 14GB data from DRAM even though it has larger LLC(24MB) and fewer XeCore(32)
A10 load 10GB data from DRAM while it has smaller LLC(6MB) and more SM(72)

From dump asm, XPU use d16u32 store on bf16 data which might trigger partial write so that DRAM will load the output tensor to read-modify-write. Can triton compiler use vector load/store to pack d16u32 to u32 or u32x2 or u32x4 to avoid partial write and reduce SendStall?

XPU code snip


import triton
import triton.language as tl

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()

from torch._dynamo.testing import rand_strided
from torch._C import _xpu_getCurrentRawStream as get_raw_stream
import torch

@triton_heuristics.reduction(
    size_hints={'x': 16384, 'r0_': 65536},
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr2': '*fp32', 'out_ptr5': '*bf16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='xpu', index=0, multi_processor_count=32, cc={'architecture': 21483225088, 'device_id': 57891, 'driver_version': '1.15.38646+6', 'gpu_eu_count': 256, 'gpu_subslice_count': 32, 'has_atomic64': True, 'has_bfloat16_conversions': True, 'has_fp16': True, 'has_fp64': True, 'has_subgroup_2d_block_io': True, 'has_subgroup_matrix_multiply_accumulate': True, 'has_subgroup_matrix_multiply_accumulate_tensor_float32': False, 'is_integrated_gpu': False, 'last_level_cache_size': 25165824, 'local_mem_size': 131072, 'max_compute_units': 256, 'max_num_sub_groups': 64, 'max_work_group_size': 1024, 'memory_bus_width': 64, 'memory_clock_rate': 2800, 'name': 'Intel(R) Arc(TM) Pro B70 Graphics', 'platform_name': 'Intel(R) oneAPI Unified Runtime over Level-Zero V2', 'sub_group_sizes': [16, 32], 'total_memory': 34242297856, 'type': 'gpu', 'vendor': 'Intel(R) Corporation', 'version': '20.2.0'}, major=None, regs_per_multiprocessor=None, max_threads_per_multi_processor=None, max_threads_per_block=1024, warp_size=None), 'constants': {}, 'native_matmul': False, 'enable_fp_fusion': True, 'launch_pdl': False, 'disable_ftz': False, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]], (4,): [['tt.divisibility', 16]], (5,): [['tt.divisibility', 16]]}]},
    inductor_meta={'grid_type': 'Grid1D', 'kernel_name': 'triton_red_fused__softmax__unsafe_view_native_layer_norm_prepare_softmax_online_9', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'atomic_add_found': False, 'num_load': 5, 'num_store': 2, 'num_reduction': 4, 'autotune_hints': set(), 'tiling_scores': {'x': 0, 'r0_': 6442221548}, 'kernel_num_gb': 8.589541356, 'kernel_flop': 0, 'backend_hash': 'FE84539060CC600F715C1B1ADB7266E44E01675FD25FBCA4207CFE12D99ECDC2', 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': True, 'dynamic_scale_rblock': True, 'incremental_autotune': False, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'deterministic': False, 'batch_invariant': False, 'force_filter_reduction_configs': False, 'mix_order_reduction_allow_multi_stages': True, 'dynamic_disable_pipelining': True, 'are_deterministic_algorithms_enabled': False}
)
@triton.jit
def triton_red_fused__softmax__unsafe_view_native_layer_norm_prepare_softmax_online_9(in_ptr0, in_ptr1, in_ptr2, out_ptr2, out_ptr5, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
    xnumel = 16384
    r0_numel = 65531
    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK], True, tl.int1)[:, None]
    r0_base = tl.arange(0, R0_BLOCK)[None, :]
    rbase = r0_base
    x0 = xindex
    tmp3_mean = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    tmp3_m2 = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    tmp3_weight = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
        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 + 65531*x0), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp1 = tmp0.to(tl.float32)
        tmp2 = tl.broadcast_to(tmp1, [XBLOCK, R0_BLOCK])
        tmp3_mean_next, tmp3_m2_next, tmp3_weight_next = triton_helpers.welford_reduce(
            tmp2, tmp3_mean, tmp3_m2, tmp3_weight, roffset == 0
        )
        tmp3_mean = tl.where(r0_mask, tmp3_mean_next, tmp3_mean)
        tmp3_m2 = tl.where(r0_mask, tmp3_m2_next, tmp3_m2)
        tmp3_weight = tl.where(r0_mask, tmp3_weight_next, tmp3_weight)
    tmp4, tmp5, tmp6 = triton_helpers.welford(tmp3_mean, tmp3_m2, tmp3_weight, 1)
    tmp3 = tmp4[:, None]
    tmp7 = tmp5[:, None]
    tmp8 = tmp6[:, None]
    _tmp27_max = tl.full([XBLOCK, R0_BLOCK], float('-inf'), tl.float32)
    _tmp27_sum = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp9 = tl.load(in_ptr0 + (r0_1 + 65531*x0), r0_mask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp18 = tl.load(in_ptr1 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp21 = tl.load(in_ptr2 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp10 = tmp9.to(tl.float32)
        tmp11 = tmp10 - tmp3
        tmp12 = tl.full([1, 1], 65531.0, tl.float32)
        tmp13 = (tmp7 / tmp12)
        tmp14 = tl.full([1, 1], 1e-05, tl.float32)
        tmp15 = tmp13 + tmp14
        tmp16 = libdevice.rsqrt(tmp15)
        tmp17 = tmp11 * tmp16
        tmp19 = tmp18.to(tl.float32)
        tmp20 = tmp17 * tmp19
        tmp22 = tmp21.to(tl.float32)
        tmp23 = tmp20 + tmp22
        tmp24 = tmp23.to(tl.float32)
        tmp25 = tmp24.to(tl.float32)
        tmp26 = tl.broadcast_to(tmp25, [XBLOCK, R0_BLOCK])

        _tmp27_max_next, _tmp27_sum_next = triton_helpers.online_softmax_combine(
            _tmp27_max, _tmp27_sum, tmp26, False
        )

        _tmp27_max = tl.where(r0_mask, _tmp27_max_next, _tmp27_max)
        _tmp27_sum = tl.where(r0_mask, _tmp27_sum_next, _tmp27_sum)
        tl.store(out_ptr2 + (r0_1 + 65536*x0), tmp25, r0_mask)

    tmp27, tmp28 = triton_helpers.online_softmax_reduce(
        _tmp27_max, _tmp27_sum, 1, False)
    tmp27 = tmp27[:, None]
    tmp28 = tmp28[:, None]
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp29 = tl.load(out_ptr2 + (r0_1 + 65536*x0), r0_mask, eviction_policy='evict_first', other=0.0)
        tmp30 = tmp29 - tmp27
        tmp31 = libdevice.exp(tmp30)
        tmp32 = (tmp31 / tmp28)
        tmp33 = tmp32.to(tl.float32)
        tl.store(out_ptr5 + (r0_1 + 65536*x0), tmp33, r0_mask)


def get_args():
    arg_0 = rand_strided((16384, 65531), (65531, 1), device='xpu:0', dtype=torch.bfloat16)
    arg_1 = rand_strided((65531,), (1,), device='xpu:0', dtype=torch.bfloat16)
    arg_2 = rand_strided((65531,), (1,), device='xpu:0', dtype=torch.bfloat16)
    arg_3 = rand_strided((64, 256, 65531), (16777216, 65536, 1), device='xpu:0', dtype=torch.float32)
    arg_4 = rand_strided((64, 256, 65531), (16777216, 65536, 1), device='xpu:0', dtype=torch.bfloat16)
    return arg_0, arg_1, arg_2, arg_3, arg_4, 16384, 65531,


def call(args):
    with torch.xpu._DeviceGuard(0):
        torch.xpu.set_device(0)
        raw_stream0 = get_raw_stream(0)
        triton_red_fused__softmax__unsafe_view_native_layer_norm_prepare_softmax_online_9.run(*args, stream=raw_stream0)


def benchmark_all_configs(args):
    with torch.xpu._DeviceGuard(0):
        torch.xpu.set_device(0)
        return triton_red_fused__softmax__unsafe_view_native_layer_norm_prepare_softmax_online_9.benchmark_all_configs(*args)


if __name__ == '__main__':
    from torch._inductor.runtime.benchmarking import benchmarker

    args = get_args()
    ms = benchmarker.benchmark(lambda: call(args), device='xpu', rep=40)
    num_gb = 8.589541356
    gb_per_s = num_gb / (ms / 1e3)
    print(f"{ms:.3f}ms    {num_gb:.3f}GB    {gb_per_s:.2f}GB/s")

cuda code snip


import triton
import triton.language as tl

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()

from torch._dynamo.testing import rand_strided
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
import torch

@triton_heuristics.reduction(
    size_hints={'x': 16384, 'r0_': 65536},
    reduction_hint=ReductionHint.INNER,
    filename=__file__,
    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr2': '*fp32', 'out_ptr5': '*bf16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=72, cc=86, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=1536, max_threads_per_block=1024, warp_size=32), 'constants': {}, 'native_matmul': False, 'enable_fp_fusion': True, 'launch_pdl': False, 'disable_ftz': False, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]], (4,): [['tt.divisibility', 16]], (5,): [['tt.divisibility', 16]]}]},
    inductor_meta={'grid_type': 'Grid1D', 'kernel_name': 'triton_red_fused__softmax__unsafe_view_mm_native_layer_norm_prepare_softmax_online_7', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'atomic_add_found': False, 'num_load': 5, 'num_store': 2, 'num_reduction': 4, 'autotune_hints': set(), 'tiling_scores': {'x': 0, 'r0_': 6442221548}, 'kernel_num_gb': 8.589541356, 'kernel_flop': 0, 'backend_hash': 'A2399C4674E114514F2ABBF6D59390274BDDE13D3679244D5DC08FFE79CAF062', 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': True, 'dynamic_scale_rblock': True, 'incremental_autotune': False, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'deterministic': False, 'batch_invariant': False, 'force_filter_reduction_configs': False, 'mix_order_reduction_allow_multi_stages': True, 'dynamic_disable_pipelining': True, 'are_deterministic_algorithms_enabled': False}
)
@triton.jit
def triton_red_fused__softmax__unsafe_view_mm_native_layer_norm_prepare_softmax_online_7(in_ptr0, in_ptr1, in_ptr2, out_ptr2, out_ptr5, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
    xnumel = 16384
    r0_numel = 65531
    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = tl.full([XBLOCK], True, tl.int1)[:, None]
    r0_base = tl.arange(0, R0_BLOCK)[None, :]
    rbase = r0_base
    x0 = xindex
    tmp3_mean = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    tmp3_m2 = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    tmp3_weight = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
        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 + 65536*x0), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp1 = tmp0.to(tl.float32)
        tmp2 = tl.broadcast_to(tmp1, [XBLOCK, R0_BLOCK])
        tmp3_mean_next, tmp3_m2_next, tmp3_weight_next = triton_helpers.welford_reduce(
            tmp2, tmp3_mean, tmp3_m2, tmp3_weight, roffset == 0
        )
        tmp3_mean = tl.where(r0_mask, tmp3_mean_next, tmp3_mean)
        tmp3_m2 = tl.where(r0_mask, tmp3_m2_next, tmp3_m2)
        tmp3_weight = tl.where(r0_mask, tmp3_weight_next, tmp3_weight)
    tmp4, tmp5, tmp6 = triton_helpers.welford(tmp3_mean, tmp3_m2, tmp3_weight, 1)
    tmp3 = tmp4[:, None]
    tmp7 = tmp5[:, None]
    tmp8 = tmp6[:, None]
    _tmp27_max = tl.full([XBLOCK, R0_BLOCK], float('-inf'), tl.float32)
    _tmp27_sum = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp9 = tl.load(in_ptr0 + (r0_1 + 65536*x0), r0_mask, eviction_policy='evict_first', other=0.0).to(tl.float32)
        tmp18 = tl.load(in_ptr1 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp21 = tl.load(in_ptr2 + (r0_1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
        tmp10 = tmp9.to(tl.float32)
        tmp11 = tmp10 - tmp3
        tmp12 = tl.full([1, 1], 65531.0, tl.float32)
        tmp13 = (tmp7 / tmp12)
        tmp14 = tl.full([1, 1], 1e-05, tl.float32)
        tmp15 = tmp13 + tmp14
        tmp16 = libdevice.rsqrt(tmp15)
        tmp17 = tmp11 * tmp16
        tmp19 = tmp18.to(tl.float32)
        tmp20 = tmp17 * tmp19
        tmp22 = tmp21.to(tl.float32)
        tmp23 = tmp20 + tmp22
        tmp24 = tmp23.to(tl.float32)
        tmp25 = tmp24.to(tl.float32)
        tmp26 = tl.broadcast_to(tmp25, [XBLOCK, R0_BLOCK])

        _tmp27_max_next, _tmp27_sum_next = triton_helpers.online_softmax_combine(
            _tmp27_max, _tmp27_sum, tmp26, False
        )

        _tmp27_max = tl.where(r0_mask, _tmp27_max_next, _tmp27_max)
        _tmp27_sum = tl.where(r0_mask, _tmp27_sum_next, _tmp27_sum)
        tl.store(out_ptr2 + (r0_1 + 65536*x0), tmp25, r0_mask)

    tmp27, tmp28 = triton_helpers.online_softmax_reduce(
        _tmp27_max, _tmp27_sum, 1, False)
    tmp27 = tmp27[:, None]
    tmp28 = tmp28[:, None]
    for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
        r0_index = r0_offset + r0_base
        r0_mask = r0_index < r0_numel
        roffset = r0_offset
        rindex = r0_index
        r0_1 = r0_index
        tmp29 = tl.load(out_ptr2 + (r0_1 + 65536*x0), r0_mask, eviction_policy='evict_first', other=0.0)
        tmp30 = tmp29 - tmp27
        tmp31 = libdevice.exp(tmp30)
        tmp32 = (tmp31 / tmp28)
        tmp33 = tmp32.to(tl.float32)
        tl.store(out_ptr5 + (r0_1 + 65536*x0), tmp33, r0_mask)


def get_args():
    arg_0 = rand_strided((16384, 65536), (65536, 1), device='cuda:0', dtype=torch.bfloat16)
    arg_1 = rand_strided((65531,), (1,), device='cuda:0', dtype=torch.bfloat16)
    arg_2 = rand_strided((65531,), (1,), device='cuda:0', dtype=torch.bfloat16)
    arg_3 = rand_strided((64, 256, 65531), (16777216, 65536, 1), device='cuda:0', dtype=torch.float32)
    arg_4 = rand_strided((64, 256, 65531), (16777216, 65536, 1), device='cuda:0', dtype=torch.bfloat16)
    return arg_0, arg_1, arg_2, arg_3, arg_4, 16384, 65531,


def call(args):
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        raw_stream0 = get_raw_stream(0)
        triton_red_fused__softmax__unsafe_view_mm_native_layer_norm_prepare_softmax_online_7.run(*args, stream=raw_stream0)


def benchmark_all_configs(args):
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        return triton_red_fused__softmax__unsafe_view_mm_native_layer_norm_prepare_softmax_online_7.benchmark_all_configs(*args)


if __name__ == '__main__':
    from torch._inductor.runtime.benchmarking import benchmarker

    args = get_args()
    ms = benchmarker.benchmark(lambda: call(args), device='cuda', rep=40)
    num_gb = 8.589541356
    gb_per_s = num_gb / (ms / 1e3)
    print(f"{ms:.3f}ms    {num_gb:.3f}GB    {gb_per_s:.2f}GB/s")

Environment details

torch 2.14.0.dev20260720+xpu
triton-xpu 3.7.2+git5fcc14d9
on
B70

Metadata

Metadata

Assignees

No one assigned

    Labels

    performanceupstream: pytorch

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions