[KMCompiler][TLERaw] Optimize the sort operator based on tle_raw

[lizhangyu258]

Original implementation: https://github.com/flagos-ai/FlagGems/blob/master/src/flag_gems/ops/sort.py

This optimization achieves performance comparable to torch for cases where n > 4096. Currently, only the float16/bfloat16 data type is supported.
The optimization process is as follows:

1. Do not use the broadcast_to kernel

# Before
indices_in = (
    torch.arange(0, n, dtype=torch.int64, device=arr_in.device)
    .broadcast_to(arr.shape)
    .contiguous()
)

# After
indices_in = (
    torch.arange(0, n, dtype=torch.int64, device=arr_in.device)
    .expand(arr.shape)
    .contiguous()
)

Broadcast_to expand A/B (float16, shape=(1024, 65536))
Torch sort:                   3.953088 ms
Broadcast_to:                 29.150145 ms (vs Torch 0.136x)
Expand:                       16.182320 ms (vs Torch 0.244x)

2. Precompute the offset of each bin within each tile

# Before
@triton.jit
def compute_global_hist_kernel(
    arr_ptr,
    out_ptr,
    num_passes,
    m,
    n,
    tiles_n_per_cta,
    TILE_N: tl.constexpr,
    TILE_R: tl.constexpr,
    num_bits_per_pass: tl.constexpr,
    descending: tl.constexpr,
):
    # arr_ptr: (m, n)
    # out_ptr: (m, n_passes, r), where r = 2 ** k_bits is the number of bins
    pid = tl.program_id(0)
    pid_n = pid // m
    pid_m = pid % m

    r: tl.constexpr = 2**num_bits_per_pass
    bfe_mask: tl.constexpr = (1 << num_bits_per_pass) - 1  # a.k.a. 2 ** k_bits - 1
    CTA_TILE_N: tl.constexpr = TILE_N * tiles_n_per_cta
    cta_n_start = CTA_TILE_N * pid_n
    cta_n_end = tl.minimum(cta_n_start + CTA_TILE_N, n)

    for p in range(0, num_passes):  # parallel
        bit_offset = p * num_bits_per_pass
        for r_start in range(0, r, TILE_R):  # parallel
            bin_indices = r_start + tl.arange(0, TILE_R)
            acc = tl.zeros((TILE_R, TILE_N), dtype=tl.int64)
            for n_start in range(cta_n_start, cta_n_end, TILE_N):  # sequantial
                n_offsets = n_start + tl.arange(0, TILE_N)  # (TILE_N, )
                mask = n_offsets < cta_n_end
                arr = tl.load(arr_ptr + pid_m * n + n_offsets, mask=mask)
                arr = convert_to_uint_preverse_order(arr, descending)
                key = (arr >> bit_offset) & bfe_mask  # (TILE_N, )
                matches = tl.where(
                    mask, (bin_indices[:, None] == key), False
                )  # (TILE_R, TILE_N)
                acc += matches
            local_sum = tl.sum(acc, axis=1)
            tl.atomic_add(
                out_ptr + pid_m * num_passes * r + p * r + bin_indices,
                local_sum,
                sem="relaxed",
            )

# After
# The main advantage of precomputing each tile is that it eliminates the need for spin-waiting during the sweep
@triton.jit
def _compute_tile_histograms(
    input_ptr,
    counts_ptr,
    rows,
    n,
    TILE_N: tl.constexpr,
    NUM_BINS: tl.constexpr,
    K_BITS: tl.constexpr,
    BIT_OFFSET: tl.constexpr,
    DESCENDING: tl.constexpr,
):
    """Count every bin for one 2048-element row tile and radix pass."""
    program = tl.program_id(0)
    row = program % rows
    tile = program // rows
    tiles = tl.cdiv(n, TILE_N)

    columns = tile * TILE_N + tl.arange(0, TILE_N)
    mask = columns < n
    values = tl.load(input_ptr + row * n + columns, mask=mask)
    keys = convert_to_uint_preverse_order(values, DESCENDING)
    digits = (keys >> BIT_OFFSET) & (NUM_BINS - 1)
    counts = tl.histogram(digits.to(tl.int32), NUM_BINS, mask=mask).to(tl.int32)    # [rows, tiles, radix_bins], 每个 pass 都要重新算
    bins = tl.arange(0, NUM_BINS)                                                   
    offsets = (row * tiles + tile) * NUM_BINS + bins
    tl.store(counts_ptr + offsets, counts)

Original radix tile-offset A/B (float16, shape=(1024, 65536))
Torch sort:                   3.953088 ms
Status look-back:             16.172064 ms (vs Torch 0.244x)
Precomputed, 4 pass/16 bins:  9.479744 ms (vs Torch 0.417x)
4-pass precomputed vs look-back: 1.706x

3. Compact indices

# Before
indices_in = (
      torch.arange(n, dtype=torch.int64, device=inp.device)
      .expand(inp.shape)
      .contiguous()
  )
indices_out = torch.empty_like(indices_in)

# After
# `assert n < (1 << 30)`: Since radix already enforces this constraint, using `int32` to represent the indices is more than sufficient.
# Finally, convert the type back to int64
if compact_indices:
    temporary_indices_a = torch.empty_like(inp, dtype=torch.int32)
    temporary_indices_b = (torch.empty_like(inp, dtype=torch.int32) if num_passes > 2 else temporary_indices_a)
    final_indices = torch.empty_like(inp, dtype=torch.int64)
    compact_indices_in = temporary_indices_a
    compact_indices_out = temporary_indices_a
......

if compact_indices:
    if pass_id == 0:
        compact_indices_out = temporary_indices_a
    elif pass_id < num_passes - 1:
        compact_indices_out = (
            temporary_indices_b
            if compact_indices_in is temporary_indices_a
            else temporary_indices_a
        )
    
    sweep_kernel[grid](...)
    if pass_id < num_passes - 1:
         compact_indices_in = compact_indices_out

Four-pass precomputed compact-index A/B (float16, shape=(1024, 65536))
Torch sort:                       3.953888 ms
Precomputed + full int64:         9.486016 ms (vs Torch 0.417x)
Precomputed + compact indices:    8.352912 ms (vs Torch 0.473x)
Compact vs full-int64 indices:    1.136x

4. four passes, 8 bins/CTA to two passes, 256 bins/CTA
   • k_bits: 4 ==> 8
   • TILE_R: 8 ==> 256
   • After this adjustment, the sweep kernel needs to be modified. Continuing to use a pure Triton implementation results in a significant performance degradation.

# Before
def radix_sort(arr, k_bits=8, descending=False):
    n = arr.shape[-1]
    m = arr.numel() // n
    assert n < (1 << 30), "we have not implemented 2**30 per launch"
    dtype = arr.dtype
    num_bits = 1 if dtype == torch.bool else (arr.itemsize * 8)

    TILE_N = 1024
    tiles_n_per_cta = 8
    CTA_TILE_N = tiles_n_per_cta * TILE_N

    num_bins = 2**k_bits
    **n_passes = triton.cdiv(num_bits, k_bits)** 
    TILE_R = 16

    ......
    
    # sweep kernel 
    **TILE_R = 8**
    grid_r = triton.cdiv(num_bins, TILE_R) 
    TILE_N = 2048
    grid_n = triton.cdiv(n, TILE_N)
    grid_for_sweep = (m * grid_n, grid_r)

@triton.jit
def _sweep_with_precomputed_tile_offsets(
    input_ptr,
    input_indices_ptr,
    output_ptr,
    output_indices_ptr,
    tile_offsets_ptr,
    bit_offset,
    rows,
    n,
    tiles,
    TILE_N: tl.constexpr,
    BINS_PER_CTA: tl.constexpr,
    K_BITS: tl.constexpr,
    DESCENDING: tl.constexpr,
):
    """The original per-bin sweep with status look-back removed."""
    program = tl.program_id(0)
    row = program % rows
    tile = program // rows
    bin_group = tl.program_id(1)
    num_bins: tl.constexpr = 1 << K_BITS

    columns = tile * TILE_N + tl.arange(0, TILE_N)
    mask = columns < n
    values = tl.load(input_ptr + row * n + columns, mask=mask)
    indices = tl.load(input_indices_ptr + row * n + columns, mask=mask)
    keys = convert_to_uint_preverse_order(values, DESCENDING)
    digits = (keys >> bit_offset) & (num_bins - 1)

    first_bin = bin_group * BINS_PER_CTA
    last_bin = tl.minimum(first_bin + BINS_PER_CTA, num_bins)
    for bin_index in range(first_bin, last_bin):
        matches = mask & (digits == bin_index)
        local_prefix = (
            tl.cumsum(matches.to(tl.uint32), axis=0) - matches
        )
        
        offset_index = (row * tiles + tile) * num_bins + bin_index
        tile_offset = tl.load(tile_offsets_ptr + offset_index)
        
        positions = tile_offset + local_prefix
        tl.store(output_ptr + row * n + positions, values, mask=matches)
        tl.store(
            output_indices_ptr + row * n + positions,
            indices,
            mask=matches,
        )


# After
@triton.jit
def _sweep_cub_local_rank_precomputed(
    arr_ptr,
    associate_arr_ptr,
    out_ptr,
    associate_out_ptr32,
    associate_out_ptr64,
    tile_offsets_ptr,
    bit_offset,
    N,
    OUT_N,
    TILE_N: tl.constexpr,
    TILE_R: tl.constexpr,
    k_bits: tl.constexpr,
    descending: tl.constexpr,
    dtype_kind: tl.constexpr,
    final_pass,
):
    ......
    tle_raw.call_smem(
        _radix_rank_8x2048_precomputed_raw,
        [
            digits_smem,
            arr_ptr,
            associate_arr_ptr,
            out_ptr,
            associate_out_ptr32,
            associate_out_ptr64,
            tile_offsets_ptr,
            pid_m,
            pid_n,
            N,
            OUT_N,
            valid_count,
            final_pass,
        ],
        output_indices=[],
    )

pytest benchmark/test_sort_tile_offsets_block_radix.py::test_perf_two_pass_precomputed_block_radix -s --warmup 1000 --iter 3000

Precomputed compact-index block-radix A/B (float16, shape=(1024, 65536))
Torch sort:                         3.954016 ms
4 pass + per-bin Triton cumsum:     8.345152 ms (vs Torch 0.474x)
2 pass + 256-bin CUB block radix:   6.781408 ms (vs Torch 0.583x)
Block radix vs per-bin cumsum:      1.231x

5. Data locality
   • Adjacent CTAs now process consecutive tiles within the same row, resulting in better performance.

# Before
# [row0, tile0], [row1, tile0], ... [row1023, tile0], [row0, tile1], ...,  [row1023, tile1] ...
row = program % rows
tile = program // rows

# After
# [row0, tile0], [row0, tile1], ... [row1, tile0], [row1, tile1], ...,  [row1023, tile0], [row1023, tile1] ...
row = program // tiles
tile = program % tiles

Two-pass call_smem CTA-mapping A/B (float16, shape=(1024, 65536))
Torch sort:                             3.953664 ms
call_smem + tile-major CTA map:         6.778176 ms (vs Torch 0.583x)
call_smem + row-major CTA map:          4.552288 ms (vs Torch 0.869x)
Row-major vs tile-major CTA map:        1.489x

6. Replace tl.histogram with tle_raw.call

# Before
@triton.jit
def _tile_histogram_triton_row_major(
    input_ptr,
    counts_ptr,
    rows,
    n,
    BIT_OFFSET: tl.constexpr,
    TILE_N: tl.constexpr,
):
    program = tl.program_id(0)
    tiles = tl.cdiv(n, TILE_N)
    row = program // tiles
    tile = program - row * tiles
    columns = tile * TILE_N + tl.arange(0, TILE_N)
    mask = columns < n
    values = tl.load(input_ptr + row * n + columns, mask=mask)
    keys = _ordered_fp16_key(values)
    digits = ((keys >> BIT_OFFSET) & 0xFF).to(tl.int32)
    counts = tl.histogram(digits, 256, mask=mask).to(tl.int32)
    bins = tl.arange(0, 256)
    tl.store(counts_ptr + (row * tiles + tile) * 256 + bins, counts,)

# After
@triton.jit
def _tile_histogram_raw_primitive_row_major(
    input_ptr,
    counts_ptr,
    rows,
    n,
    BIT_OFFSET: tl.constexpr,
    TILE_N: tl.constexpr,
):
    tl.static_assert(TILE_N == 2048)
    program = tl.program_id(0)
    tiles = tl.cdiv(n, TILE_N)
    row = program // tiles
    tile = program - row * tiles
    columns = tile * TILE_N + tl.arange(0, TILE_N)
    mask = columns < n
    values = tl.load(input_ptr + row * n + columns, mask=mask)
    keys = _ordered_fp16_key(values)
    digits = ((keys >> BIT_OFFSET) & 0xFF).to(tl.uint16)

    digits_smem = tle_gpu.alloc(
        shape=[TILE_N],
        dtype=tl.uint16,
        layout=None,
        scope=tle_gpu.smem,
        nv_mma_shared_layout=False,
    )
    counts_smem = tle_gpu.alloc(
        shape=[256],
        dtype=tl.int32,
        layout=None,
        scope=tle_gpu.smem,
        nv_mma_shared_layout=False,
    )
    tl.store(tle_gpu.local_ptr(digits_smem, (columns,)), digits)
    valid_count = tl.minimum(TILE_N, n - tile * TILE_N)

    # Only the histogram primitive crosses into Raw. Loading, key/digit
    # formation, and the final global store remain in Triton.
    counts_smem = tle_raw.call_smem(
        _histogram_digits_8x2048_raw,
        [
            digits_smem,
            counts_smem,
            valid_count,
        ],
        output_indices=[1],
    )
    bins = tl.arange(0, 256)
    counts = tl.load(tle_gpu.local_ptr(counts_smem, (bins,)))
    tl.store(counts_ptr + (row * tiles + tile) * 256 + bins, counts)

call_smem row-major histogram A/B (float16, shape=(1024, 65536))
Torch sort:                         3.953680 ms
Triton histogram path:              4.610016 ms (vs Torch 0.858x)
Raw histogram primitive:            3.997792 ms (vs Torch 0.989x)
Raw-hist path vs baseline:          1.153x

7. Fuse multiple PyTorch operations introduced during the algorithm optimization process.

# Before
def _make_tile_offsets(counts):
    tile_prefix = torch.cumsum(counts, dim=1, dtype=torch.int32) - counts
    global_counts = torch.sum(counts, dim=1, dtype=torch.int32)
    bin_bases = (
        torch.cumsum(global_counts, dim=1, dtype=torch.int32) - global_counts
    )
    return tile_prefix + bin_bases.unsqueeze(1)

# After
def _make_tile_offsets_triton(counts):
    rows, tiles, bins = counts.shape
    assert bins == 256
    offsets = torch.empty_like(counts)
    _tile_offsets_triton_kernel[(rows,)](
        counts,
        offsets,
        tiles,
        num_warps=8,
    )
    return offsets

@triton.jit
def _tile_offsets_triton_kernel(
    counts_ptr,
    offsets_ptr,
    tiles,
):
    row = tl.program_id(0)
    bins = tl.arange(0, 256)
    row_base = row * tiles * 256

    # First pass: total count of every bin in this row. Keeping only a
    # 256-element vector avoids materializing the full (tiles, bins)
    # matrix in registers.
    bin_totals = tl.zeros((256,), dtype=tl.int32)
    for tile in range(0, tiles):
        counts = tl.load(counts_ptr + row_base + tile * 256 + bins)
        bin_totals += counts

    bin_bases = tl.cumsum(bin_totals, axis=0) - bin_totals

    # Second pass: exclusive prefix over tiles for each bin, offset by the
    # exclusive prefix of all preceding bins.
    running = bin_bases
    for tile in range(0, tiles):
        offset = row_base + tile * 256 + bins
        counts = tl.load(counts_ptr + offset)
        tl.store(offsets_ptr + offset, running)
        running += counts

Raw-histogram tile-offset A/B (float16, shape=(1024, 65536))
Torch sort:                         3.954560 ms
Eager torch:                        4.000000 ms (vs Torch 0.989x)
Fused Triton offsets:               3.777472 ms (vs Torch 1.047x)
Fused offsets vs eager baseline:    1.059x

Summary:

• Algorithm optimization:
  • Precompute global tile_offsets to eliminate data dependencies between tiles that cause spinning.
  • Use int32 for intermediate indices and convert them to int64 for the final output
  • Change from 4 passes with 8 bins/CTA to 2 passes with 256 bins/CTA
• TLE-Raw optimization:
  • Replace tl.histogram
  • Replace the per-bin local_rank computation in the sweep kernel

NVIDIA-H20 performance data(dtype=float16)
Original:

Status	Torch (ms)	FlagGems (ms)	Speedup	Input Shape	dim	descending
SUCCESS	0.016064	0.863904	0.019	[64, 64]	-1	False
SUCCESS	0.026912	0.855104	0.031	[256, 256]	-1	False
SUCCESS	0.054688	0.890144	0.061	[1024, 1024]	-1	False
SUCCESS	0.556784	7.236480	0.077	[4096, 4096]	-1	False
SUCCESS	3.952432	29.101120	0.136	[1024, 65536]	-1	False
SUCCESS	0.011904	0.004832	2.464	[1024, 1]	-1	False
SUCCESS	0.048576	0.895632	0.054	[1024, 512]	-1	False
SUCCESS	0.189984	1.015904	0.187	[16, 131072]	-1	False
SUCCESS	0.190016	1.049664	0.181	[8, 262144]	-1	False

After optimization:

Status	Torch (ms)	FlagGems (ms)	Speedup	Input Shape	dim	descending
SUCCESS	0.015424	0.164992	0.093	[64, 64]	-1	False
SUCCESS	0.026208	0.160256	0.164	[256, 256]	-1	False
SUCCESS	0.053024	0.174048	0.305	[1024, 1024]	-1	False
SUCCESS	0.556032	0.850704	0.654	[4096, 4096]	-1	False
SUCCESS	3.954688	3.791904	🟢 1.043	[1024, 65536]	-1	False
SUCCESS	0.010816	0.004608	2.347	[1024, 1]	-1	False
SUCCESS	0.047760	0.166912	0.286	[1024, 512]	-1	False
SUCCESS	0.188448	0.187184	🟢 1.007	[16, 131072]	-1	False
SUCCESS	0.188624	0.188064	🟢 1.003	[8, 262144]	-1	False