nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/sigmoid.py

@@ -0,0 +1,19 @@
+"""Sigmoid — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def sigmoid_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    x_fp32 = x.to(tl.float32)
+    out = 1.0 / (1.0 + tl.exp(-x_fp32))
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/silu.py

@@ -0,0 +1,27 @@
+"""SiLU (Swish) — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+from ._common import sigmoid
+
+@triton.jit
+def silu_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    # Cast offset to int64 so that mask `offset < n_elements`
+    # works correctly when n_elements >= 2^31 (e.g. Sana 4Kpx VAE
+    # silu input 1*128*4096*4096 = 2^31 elements). Without this,
+    # int32 n_elements wraps to negative and the mask is all-False,
+    # silently skipping every element and leaving output as
+    # uninitialized memory (= garbage).
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    x_fp32 = x.to(tl.float32)
+    out = x_fp32 * sigmoid(x_fp32)
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/sin.py

@@ -0,0 +1,17 @@
+"""Sine — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def sin_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, tl.sin(x.to(tl.float32)), mask=mask)

nbxtensor/ops/softmax.py

@@ -0,0 +1,169 @@
+"""Softmax — pure @triton.jit kernel. Ported from attorch (MIT)."""
+
+import triton
+import triton.language as tl
+from triton import next_power_of_2
+
+from nbxtensor.ops._configs import nbx_autotune
+
+
+def _softmax_configs():
+    """The candidates the certifier proves (certified per key since 2026-10-19 instead of a hand table): rows per program x warps. The
+    one the kernel ran before certification is among them (1 row, 4 warps for a row of 64 or more; up to 128 rows for
+    a short row) — a certified winner equal to it changes no bit. Rows per program partition independent rows and change
+    no bit; the warps split a row's reduction and can (V100, 4096 columns: 0.001 % of the elements)."""
+    return [triton.Config({'BLOCK_SIZE_BATCH': rows}, num_warps=warps)
+            for rows in (1, 2, 4, 8, 16, 32, 64, 128) for warps in (1, 2, 4, 8)]
+
+
+def _prune_softmax(configs, named_args, **kwargs):
+    """A program holds rows x the padded row: at most what the kernel already holds for ONE row — the longest row it
+    accepts (`launch_keys.SOFTMAX_FUSED_MAX_FEAT`, beyond which the looped kernel runs) — and at most the backend's
+    reduce-tile cap (`nbx_tensor.reduce_tile_max`, the profile's, where the lowering refuses beyond); never more rows
+    than exist. (The cap alone is far above the register budget on CUDA: 128 rows of a 2048-padded row compiled for
+    minutes per candidate, 2026-10-20.)"""
+    from nbxtensor.backend.capabilities import reduce_tile_max
+    from nbxtensor.ops.launch_keys import SOFTMAX_FUSED_MAX_FEAT
+    feat = next_power_of_2(int(named_args['feat_dim']))
+    rows = int(named_args['batch_dim'])
+    budget = min(reduce_tile_max(), SOFTMAX_FUSED_MAX_FEAT)
+    from nbxtensor.ops._configs import fits_thread_registers
+    keep = [c for c in configs
+            if c.kwargs['BLOCK_SIZE_BATCH'] * feat <= max(budget, feat) and c.kwargs['BLOCK_SIZE_BATCH'] <= max(1, rows)
+            and fits_thread_registers(c.kwargs['BLOCK_SIZE_BATCH'] * feat, c.num_warps)]
+    return keep or [max(configs, key=lambda c: c.num_warps)]
+
+
+@nbx_autotune(configs=_softmax_configs(), key=['BATCH_BUCKET', 'FEAT_BUCKET', 'log'],
+              prune_configs_by={'early_config_prune': _prune_softmax}, cache_results=True)
+@triton.heuristics({
+    'BLOCK_SIZE_FEAT': lambda args: next_power_of_2(args['feat_dim']),
+})
+@triton.jit
+def softmax_forward_kernel(
+    input_ptr, output_ptr,
+    batch_dim, feat_dim, BATCH_BUCKET, FEAT_BUCKET,
+    input_batch_stride, input_feat_stride,
+    output_batch_stride, output_feat_stride,
+    log: tl.constexpr,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_FEAT: tl.constexpr,
+):
+    """Softmax forward (or log_softmax if log=True).
+
+    input: [batch_dim, feat_dim]
+    output: [batch_dim, feat_dim]
+    """
+    batch_pid = tl.program_id(0).to(tl.int64)
+
+    batch_offset = batch_pid * BLOCK_SIZE_BATCH + tl.arange(0, BLOCK_SIZE_BATCH)
+    feat_offset = tl.arange(0, BLOCK_SIZE_FEAT)
+
+    batch_mask = batch_offset < batch_dim
+    feat_mask = feat_offset < feat_dim
+
+    input_ptr += (input_batch_stride * batch_offset[:, None] +
+                  input_feat_stride * feat_offset[None, :])
+    output_ptr += (output_batch_stride * batch_offset[:, None] +
+                   output_feat_stride * feat_offset[None, :])
+
+    inp = tl.load(input_ptr,
+                  mask=batch_mask[:, None] & feat_mask[None, :],
+                  other=-float('inf')).to(tl.float32)
+
+    # Numerical stability: subtract max
+    inp -= tl.max(inp, axis=1)[:, None]
+    numerator = tl.exp(inp)
+    denominator = tl.sum(numerator, axis=1)[:, None]
+
+    if log:
+        output = inp - tl.log(denominator)
+    else:
+        output = numerator / denominator
+
+    tl.store(output_ptr, output,
+             mask=batch_mask[:, None] & feat_mask[None, :])
+
+
+def _looped_softmax_configs():
+    """The candidates the certifier proves for the looped kernel (chantier 1): the tile a row is scanned by x warps.
+    The one it ran before certification is among them (4 096 columns, 4 warps) — a certified winner equal to it changes
+    no bit; another tile or warp count changes the order of a row's running max and sum."""
+    return [triton.Config({'TILE_FEAT': tile}, num_warps=warps)
+            for tile in (1024, 2048, 4096, 8192, 16384) for warps in (2, 4, 8)]
+
+
+def _prune_looped_softmax(configs, named_args, **kwargs):
+    """A tile no wider than the backend's reduce-tile cap (`nbx_tensor.reduce_tile_max`, the profile's, where the
+    lowering refuses beyond) nor than the padded row it scans (a wider tile only masks)."""
+    from nbxtensor.backend.capabilities import reduce_tile_max
+    cap = min(reduce_tile_max(), next_power_of_2(int(named_args['feat_dim'])))
+    keep = [c for c in configs if c.kwargs['TILE_FEAT'] <= cap]
+    return keep or configs[:1]
+
+
+@nbx_autotune(configs=_looped_softmax_configs(), key=['BATCH_BUCKET', 'FEAT_BUCKET', 'log'],
+              prune_configs_by={'early_config_prune': _prune_looped_softmax}, cache_results=True)
+@triton.jit
+def softmax_forward_looped_kernel(
+    input_ptr, output_ptr,
+    feat_dim, BATCH_BUCKET, FEAT_BUCKET,
+    input_batch_stride, input_feat_stride,
+    output_batch_stride, output_feat_stride,
+    log: tl.constexpr,
+    TILE_FEAT: tl.constexpr,
+):
+    """Online two-pass softmax for rows too large for one register tile.
+
+    One program per row; the row is scanned in TILE_FEAT chunks with a
+    running max / rescaled running sum (pass 1), then normalized and
+    stored (pass 2). Register footprint is bounded by TILE_FEAT
+    regardless of feat_dim.
+
+    Why this exists: the single-tile kernel above sizes its tile as
+    next_power_of_2(feat_dim). At vocab-scale rows (e.g. 151936 →
+    262144-element tile) it spills to CUDA local memory, and the driver
+    allocates the local pool DEVICE-WIDE at launch — per-thread frame ×
+    max resident threads (163,840 on V100) ≈ 2 GiB, measured. The launch
+    itself then fails with CUDA_ERROR_OUT_OF_MEMORY whenever the device
+    cannot serve that hidden allocation (D9 warm-daemon incident,
+    2026-07-27). Bounded tiles keep the frame register-resident, so the
+    launch has no local-memory cost at any feat_dim.
+    """
+    row = tl.program_id(0).to(tl.int64)
+    in_row = input_ptr + row * input_batch_stride
+    out_row = output_ptr + row * output_batch_stride
+
+    running_max = tl.full((1,), -float('inf'), tl.float32)
+    running_sum = tl.zeros((1,), tl.float32)
+    for start in range(0, feat_dim, TILE_FEAT):
+        offs = start + tl.arange(0, TILE_FEAT)
+        mask = offs < feat_dim
+        x = tl.load(in_row + offs * input_feat_stride, mask=mask,
+                    other=-float('inf')).to(tl.float32)
+        new_max = tl.maximum(running_max, tl.max(x, axis=0))
+        # -inf guard (online-softmax idiom): while every element seen so
+        # far is -inf (top-k/top-p masking routinely -inf-fills whole
+        # leading tiles of an unsorted vocab row), new_max stays -inf and
+        # both exp arguments would be -inf - (-inf) = NaN, poisoning the
+        # running sum for the entire row. Pin alpha to 1 and the shifted
+        # tile to -inf (exp → 0) until the first finite value arrives.
+        # Fully -inf rows still yield NaN in pass 2 — parity with the
+        # single-tile kernel, guarded downstream by the samplers.
+        seen_only_neginf = new_max == -float('inf')
+        alpha = tl.where(seen_only_neginf, 1.0,
+                         tl.exp(running_max - new_max))
+        shifted = tl.where(seen_only_neginf, -float('inf'), x - new_max)
+        running_sum = running_sum * alpha + tl.sum(tl.exp(shifted), axis=0)
+        running_max = new_max
+
+    for start in range(0, feat_dim, TILE_FEAT):
+        offs = start + tl.arange(0, TILE_FEAT)
+        mask = offs < feat_dim
+        x = tl.load(in_row + offs * input_feat_stride, mask=mask,
+                    other=-float('inf')).to(tl.float32)
+        if log:
+            y = x - running_max - tl.log(running_sum)
+        else:
+            y = tl.exp(x - running_max) / running_sum
+        tl.store(out_row + offs * output_feat_stride, y, mask=mask)

nbxtensor/ops/softplus.py

@@ -0,0 +1,27 @@
+"""Softplus activation — pure @triton.jit kernel.
+
+Logic from FlagGems softplus.py: softplus(x) = log(1 + exp(beta * x)) / beta.
+For large beta*x, returns x directly (numerical stability).
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def softplus_forward_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    beta,
+    threshold,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    x = tl.load(x_ptr + offset, mask=mask)
+    x_fp = x.to(tl.float32)
+    z = x_fp * beta
+    soft_z = tl.where(z > threshold, z, tl.log(1.0 + tl.exp(z)))
+    result = soft_z / beta
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/sort_op.py

@@ -0,0 +1,304 @@
+"""Sort — pure @triton.jit radix sort helper kernels.
+
+Ported from FlagGems sort (Apache-2.0 license).
+Extracts the uint-conversion helpers and the two main radix sort kernels
+(global histogram + sweep). The small-N bitonic sort kernel is also included.
+
+The radix sort orchestration (multi-pass loop, buffer allocation) must be
+done in the wrapper since it requires memory allocation.
+"""
+
+import triton
+import triton.language as tl
+
+
+# --------------------------------------------------------------------------- #
+# Uint conversion helpers — preserve sort order across signed/float types
+# --------------------------------------------------------------------------- #
+
+# COMPILE-TIME helpers — `@tl.constexpr`, NEVER `@triton.jit`. Upstream
+# FlagGems has carried them as constexpr functions since the one-sweep
+# sort's birth commit (e669aee, PR #694); our vendoring delta hardened
+# them into `@triton.jit`, which Triton >= 3.4 structurally rejects (the
+# `_builder`→`_semantic` frontend refactor: a JIT callee compiles as a
+# real typed function, so it can no longer return a dtype object, and
+# `tl.core.get_int_dtype` — a plain undecorated python function — no
+# longer resolves from inside JIT). `@tl.constexpr` runs them as plain
+# compile-time python (legal 3.1→3.6), exactly upstream's shape.
+def _unwrap_constexpr(o):
+    return o.value if isinstance(o, tl.constexpr) else o
+
+
+@tl.constexpr
+def _get_int_t(num_bits: tl.constexpr, signed: tl.constexpr):
+    num_bits = _unwrap_constexpr(num_bits)
+    signed = _unwrap_constexpr(signed)
+    return tl.core.get_int_dtype(num_bits, signed)
+
+
+@tl.constexpr
+def _one_zeros(num_bits: tl.constexpr):
+    """1 followed by (num_bits-1) zeros, e.g. 0x80 for 8-bit."""
+    num_bits = _unwrap_constexpr(num_bits)
+    return 1 << (num_bits - 1)
+
+
+@tl.constexpr
+def _zero_ones(num_bits: tl.constexpr):
+    """0 followed by (num_bits-1) ones, e.g. 0x7F for 8-bit."""
+    num_bits = _unwrap_constexpr(num_bits)
+    return (1 << (num_bits - 1)) - 1
+
+
+@triton.jit
+def uint_to_uint(x, descending: tl.constexpr = False):
+    """Unsigned → unsigned: negate bits if descending."""
+    out = ~x if descending else x
+    return out
+
+
+@triton.jit
+def int_to_uint(x, descending: tl.constexpr = False):
+    """Signed int → unsigned, preserving sort order."""
+    num_bits: tl.constexpr = x.dtype.primitive_bitwidth
+    udtype = _get_int_t(num_bits, False)
+    ux = tl.cast(x, udtype, bitcast=True)
+    if descending:
+        bit_mask: tl.constexpr = _zero_ones(num_bits)
+        bit_mask_tensor = tl.full((), value=bit_mask, dtype=udtype)
+        out = ux ^ bit_mask_tensor
+    else:
+        sign_bit_mask: tl.constexpr = _one_zeros(num_bits)
+        sign_bit_mask_tensor = tl.full((), value=sign_bit_mask, dtype=udtype)
+        out = ux ^ sign_bit_mask_tensor
+    return out
+
+
+@triton.jit
+def floating_to_uint(x, descending: tl.constexpr = False):
+    """Float → unsigned, preserving sort order (handles sign bit correctly)."""
+    num_bits: tl.constexpr = x.dtype.primitive_bitwidth
+    sdtype = _get_int_t(num_bits, True)
+    udtype = _get_int_t(num_bits, False)
+    sx = x.to(sdtype, bitcast=True)
+    ux = x.to(udtype, bitcast=True)
+
+    sign_bit_mask_v: tl.constexpr = _one_zeros(num_bits)
+    sign_bit_mask = tl.full((), value=sign_bit_mask_v, dtype=udtype)
+    rshift_bits = tl.full((), value=num_bits - 1, dtype=sdtype)
+    mask = sign_bit_mask | (sx >> rshift_bits).to(udtype, bitcast=True)
+    if descending:
+        out = ux ^ (~mask)
+    else:
+        out = ux ^ mask
+    return out.to(udtype, bitcast=True)
+
+
+@triton.jit
+def convert_to_uint_preserve_order(x, descending: tl.constexpr = False):
+    """Convert any numeric type to unsigned, preserving sort order."""
+    if x.dtype.is_floating():
+        out = floating_to_uint(x, descending)
+    elif x.dtype.is_int_signed():
+        out = int_to_uint(x, descending)
+    elif x.dtype.is_int_unsigned():
+        out = uint_to_uint(x, descending)
+    return out
+
+
+# --------------------------------------------------------------------------- #
+# Global histogram kernel — counts per-bin occurrences for radix sort
+# --------------------------------------------------------------------------- #
+
+@triton.jit
+def radix_sort_histogram_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,
+):
+    """Compute per-bin histograms for radix sort.
+
+    arr_ptr: (m, n) — input data
+    out_ptr: (m, num_passes, r) — histogram output, r = 2^num_bits_per_pass
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    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
+    CTA_TILE_N: tl.constexpr = TILE_N * tiles_n_per_cta
+    cta_n_start = (CTA_TILE_N * pid_n).to(tl.int32)   # loop bounds are 32-bit counts: the Metal lowering refuses a 64-bit scf.for bound (addressing stays 64-bit through the program ids)
+    cta_n_end = (tl.minimum(cta_n_start + CTA_TILE_N, n)).to(tl.int32)
+
+    for p in range(0, num_passes):
+        bit_offset = p * num_bits_per_pass
+        for r_start in range(0, r, TILE_R):
+            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):
+                n_offsets = n_start + tl.arange(0, TILE_N)
+                mask = n_offsets < cta_n_end
+                arr = tl.load(arr_ptr + pid_m * n + n_offsets, mask=mask)
+                arr = convert_to_uint_preserve_order(arr, descending)
+                key = (arr >> bit_offset) & bfe_mask
+                matches = tl.where(mask, (bin_indices[:, None] == key), False)
+                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",
+            )
+
+
+# --------------------------------------------------------------------------- #
+# Sweep kernel — scatter elements to sorted positions using decoupled lookback
+# --------------------------------------------------------------------------- #
+
+@triton.jit
+def radix_sort_tile_counts_kernel(
+    arr_ptr,
+    counts_ptr,
+    m,
+    N,
+    OUT_N,
+    TILE_N: tl.constexpr,
+    TILE_R: tl.constexpr,
+    k_bits: tl.constexpr,
+    bit_offset,
+    descending: tl.constexpr,
+):
+    """How many elements of each bin this TILE holds. Stage 1 of three.
+
+    This and `radix_sort_scatter_kernel` replace a single kernel that carried a
+    DECOUPLED LOOKBACK: each tile published its count and then spun on its
+    predecessors' flags to learn how many of its bin came before it.
+
+    That construction needs forward progress across CTAs and cross-CTA
+    visibility for a `.cg` store, and Triton promises neither — inside a kernel
+    it promises nothing about other programs at all. Measured 2026-09-17 on
+    triton-ext: correct to exactly 2048 elements (one tile) and wrong from 2049,
+    with elements LOST rather than misordered (9 absent at n=2049, 1726 at
+    n=4096) — two tiles computing the same destinations and overwriting each
+    other, which is what a lookback reading stale flags produces.
+
+    The replacement is not a new design. `cumsum_wrapper` in this same tree does
+    its cross-tile scan as SEPARATE LAUNCHES (`scan_part_sum_kernel`, then
+    `add_base_sum_kernel` when there is more than one part) and is correct at
+    eight tiles, measured the same day. Triton guarantees ordering BETWEEN
+    launches; that is the guarantee this needs.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    pid_m = pid % m
+    pid_n = pid // m
+    pid_r = tl.program_id(1).to(tl.int64)
+
+    bfe_mask: tl.constexpr = (1 << k_bits) - 1
+    r: tl.constexpr = 2 ** k_bits
+    cta_r_start = (pid_r * TILE_R).to(tl.int32)
+    cta_r_end = (tl.minimum(cta_r_start + TILE_R, r)).to(tl.int32)
+
+    n_offsets = pid_n * TILE_N + tl.arange(0, TILE_N)
+    mask = n_offsets < N
+    arr = tl.load(arr_ptr + pid_m * N + n_offsets, mask=mask)
+    arr_u = convert_to_uint_preserve_order(arr, descending)
+    key = (arr_u >> bit_offset) & bfe_mask
+
+    for bin_index in range(cta_r_start, cta_r_end):
+        matches = tl.where(mask, key == bin_index, False)
+        local_sum = tl.sum(matches.to(tl.uint32), axis=0)
+        tl.store(counts_ptr + pid_m * (r * OUT_N) + bin_index * OUT_N + pid_n,
+                 local_sum)
+
+
+@triton.jit
+def radix_sort_tile_prefix_kernel(
+    counts_ptr,
+    prefix_ptr,
+    num_tiles,
+    OUT_N,
+    TILES_POW2: tl.constexpr,
+):
+    """Exclusive scan of the per-tile counts, along the TILE axis. Stage 2.
+
+    One program per (batch, bin) row — the rows are independent, and the row is
+    `num_tiles` long, which is `cdiv(n, 2048)`: small. This is the step the
+    lookback was doing inside the scatter, done where Triton actually orders it.
+    """
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, TILES_POW2)
+    mask = offs < num_tiles
+    counts = tl.load(counts_ptr + row * OUT_N + offs, mask=mask, other=0)
+    prefix = tl.cumsum(counts, axis=0) - counts          # exclusive
+    tl.store(prefix_ptr + row * OUT_N + offs, prefix, mask=mask)
+
+
+@triton.jit
+def radix_sort_scatter_kernel(
+    arr_ptr,
+    associate_arr_ptr,
+    out_ptr,
+    associate_out_ptr,
+    excumsum_bins_ptr,
+    tile_prefix_ptr,
+    n_passes,
+    pass_id,
+    bit_offset,
+    m,
+    N,
+    OUT_N,
+    TILE_N: tl.constexpr,
+    TILE_R: tl.constexpr,
+    k_bits: tl.constexpr,
+    descending: tl.constexpr,
+):
+    """Scatter elements to their sorted positions. Stage 3.
+
+    Identical to the old `radix_sort_sweep_kernel` except that
+    `exclusive_prefix` — how many of this bin lie in EARLIER tiles — is READ
+    from the scan of stage 2 instead of being discovered by spinning on other
+    programs' flags.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    pid_m = pid % m
+    pid_n = pid // m
+    pid_r = tl.program_id(1).to(tl.int64)
+
+    bfe_mask: tl.constexpr = (1 << k_bits) - 1
+    r: tl.constexpr = 2 ** k_bits
+    cta_r_start = (pid_r * TILE_R).to(tl.int32)
+    cta_r_end = (tl.minimum(cta_r_start + TILE_R, r)).to(tl.int32)
+
+    n_offsets = pid_n * TILE_N + tl.arange(0, TILE_N)
+    mask = n_offsets < N
+    arr = tl.load(arr_ptr + pid_m * N + n_offsets, mask=mask)
+    arr_u = convert_to_uint_preserve_order(arr, descending)
+    key = (arr_u >> bit_offset) & bfe_mask
+
+    for bin_index in range(cta_r_start, cta_r_end):
+        matches = tl.where(mask, key == bin_index, False)
+        exclusive_prefix = tl.load(
+            tile_prefix_ptr + pid_m * (r * OUT_N) + bin_index * OUT_N + pid_n)
+
+        local_ex_cumsum = tl.cumsum(matches.to(tl.uint32), axis=0) - matches
+        ex_cumsum_in_bin = exclusive_prefix + local_ex_cumsum
+
+        ex_cumsum_bins = tl.load(
+            excumsum_bins_ptr + pid_m * (n_passes * r) + pass_id * r + bin_index
+        )
+        pos = ex_cumsum_bins + ex_cumsum_in_bin
+
+        tl.store(out_ptr + pid_m * N + pos, arr, mask=matches)
+        if associate_arr_ptr is not None:
+            associate_arr = tl.load(
+                associate_arr_ptr + pid_m * N + n_offsets, mask=mask
+            )
+            tl.store(associate_out_ptr + pid_m * N + pos, associate_arr, mask=matches)