nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/copy.py

@@ -0,0 +1,18 @@
+"""Copy / dtype cast — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def copy_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Copy tensor (with implicit dtype conversion via pointer types)."""
+    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, x, mask=mask)

nbxtensor/ops/copy_op.py

@@ -0,0 +1,39 @@
+"""Copy — pure @triton.jit kernel.
+
+Ported from FlagGems copy (Apache-2.0 license).
+Copies elements from src to dst, with optional dtype conversion
+handled by Triton's automatic casting on store.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def copy_kernel(
+    src_ptr, dst_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    SATURATE_F16: tl.constexpr = False,
+):
+    """Copy src to dst. Dtype conversion happens automatically via Triton's
+    store when src and dst pointer types differ.
+
+    SATURATE_F16: protected float16 downcast — finite values beyond the
+    fp16 range clamp to ±65504 instead of overflowing to ±inf. This
+    mirrors the compiled DtypeEngine's hardware-protected conversion
+    contract (bf16/fp32 sentinel fills like finfo.min must survive a
+    V100 fp16 remap as finite extremes; an additive attention mask that
+    overflows to -inf poisons softmax rows — the DeepSeek-V2 wrong-token
+    incident). NaN and true ±inf inputs pass through unchanged.
+    """
+    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
+
+    data = tl.load(src_ptr + offset, mask=mask)
+    if SATURATE_F16:
+        f = data.to(tl.float32)
+        is_finite = (f == f) & (f != float("inf")) & (f != float("-inf"))
+        clamped = tl.minimum(tl.maximum(f, -65504.0), 65504.0)
+        data = tl.where(is_finite, clamped, f)
+    tl.store(dst_ptr + offset, data, mask=mask)

nbxtensor/ops/cos.py

@@ -0,0 +1,17 @@
+"""Cosine — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def cos_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.cos(x.to(tl.float32)), mask=mask)

nbxtensor/ops/count_nonzero.py

@@ -0,0 +1,64 @@
+"""Count Nonzero — pure @triton.jit kernel.
+
+Counts non-zero elements in a flattened tensor using atomic addition.
+Each block counts its local non-zeros, then atomically adds to the output.
+Extracted from FlagGems (Apache 2.0).
+
+Grid: (ceil(numel / BLOCK_SIZE),)
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def count_nonzero_kernel(
+    x_ptr,
+    out_ptr,
+    numel,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Count non-zero elements via block-local sum + atomic global add.
+
+    Each program loads a block of elements, compares != 0, sums the
+    boolean result, then atomically adds the count to the output scalar.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < numel
+
+    x = tl.load(x_ptr + offsets, mask=mask, other=0)
+    is_nonzero = (x != 0).to(tl.int64)
+    nonzero_count = tl.sum(is_nonzero, axis=0)
+
+    tl.atomic_add(out_ptr, nonzero_count)
+
+@triton.jit
+def count_nonzero_dim_kernel(
+    x_ptr,
+    out_ptr,
+    N,
+    numel,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Count non-zero elements along a dimension (after dim_compress).
+
+    Each program handles one row of the compressed tensor and iterates
+    over columns in blocks, accumulating the non-zero count.
+
+    x_ptr: flattened compressed input (M * N elements, reduction dim = N)
+    out_ptr: output of shape (M,)
+    """
+    pid_x = tl.program_id(0).to(tl.int64)
+
+    nonzero_count = tl.zeros((), dtype=tl.int64)
+    for start_n in range(0, N, BLOCK_SIZE):
+        cols_offsets = start_n + tl.arange(0, BLOCK_SIZE)
+        offset = pid_x * N + cols_offsets
+        col_mask = (offset < numel) & (cols_offsets < N)
+
+        x = tl.load(x_ptr + offset, mask=col_mask, other=0)
+        is_nonzero = (x != 0).to(tl.int64)
+        nonzero_count += tl.sum(is_nonzero)
+
+    tl.store(out_ptr + pid_x, nonzero_count)

nbxtensor/ops/cross_entropy.py

@@ -0,0 +1,67 @@
+"""Cross entropy loss — pure @triton.jit kernel.
+
+Extracted from attorch (BobMcDear/attorch) cross_entropy_loss_kernels.py.
+Chosen over Liger (523 lines) for inference: simpler, pure, no training extras.
+Forward only. Numerically stable log-sum-exp implementation.
+"""
+
+import triton
+import triton.language as tl
+from triton import next_power_of_2
+
+@triton.heuristics({
+    'BLOCK_SIZE_FEAT': lambda args: next_power_of_2(args['feat_dim']),
+})
+@triton.jit
+def cross_entropy_loss_forward_kernel(
+    input_pointer, target_pointer, weight_pointer,
+    sum_weights_pointer, output_pointer,
+    batch_dim, feat_dim,
+    input_batch_stride, input_feat_stride,
+    weighted: tl.constexpr,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_FEAT: tl.constexpr,
+):
+    """Cross entropy loss: -log(softmax(input)[target]).
+
+    Numerically stable via log-sum-exp trick.
+
+    Args:
+        input_pointer: [batch_dim, feat_dim] logits
+        target_pointer: [batch_dim] integer class indices
+        weight_pointer: optional [feat_dim] class weights
+        output_pointer: [batch_dim/BLOCK_SIZE_BATCH] partial losses
+    """
+    batch_pid = tl.program_id(axis=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
+
+    target = tl.load(target_pointer + batch_offset, mask=batch_mask)
+
+    pred_pointer = (input_pointer +
+                    input_feat_stride * target +
+                    input_batch_stride * batch_offset)
+    input_pointer += (input_batch_stride * batch_offset[:, None] +
+                      input_feat_stride * feat_offset[None, :])
+
+    input = tl.load(input_pointer, mask=batch_mask[:, None] & feat_mask[None, :],
+                    other=-float('inf')).to(tl.float32)
+    pred = tl.load(pred_pointer, mask=batch_mask).to(tl.float32)
+
+    # Log-sum-exp: loss = log(sum(exp(x - max))) + max - pred
+    mx = tl.max(input, axis=1)
+    input -= mx[:, None]
+    loss = tl.log(tl.sum(tl.exp(input), axis=1)) - pred + mx
+
+    if weighted:
+        weight = tl.load(weight_pointer + target, mask=batch_mask).to(tl.float32)
+        loss *= weight
+        tl.store(sum_weights_pointer + batch_pid, tl.sum(weight))
+    else:
+        loss /= batch_dim
+
+    tl.store(output_pointer + batch_pid, tl.sum(loss))

nbxtensor/ops/cummax.py

@@ -0,0 +1,109 @@
+"""Cumulative Maximum — pure @triton.jit kernels.
+
+Computes running maximum along a dimension, returning both values and indices.
+Uses scan-then-fan algorithm: each block computes a local cummax via
+tl.associative_scan, then partial results are combined across blocks.
+
+Extracted from FlagGems (Apache 2.0).
+
+Three kernels for the scan-then-fan pipeline:
+1. scan_part_max_kernel — local cummax within each block + partial max
+2. add_base_max_kernel — fan-out: combine partial maxima with local results
+3. cummax_loop_kernel — single-program loop variant for wide reductions
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _cummax_combine(val_a, idx_a, val_b, idx_b):
+    """Associative combiner: keep (value, index) of the maximum."""
+    pick_b = val_b > val_a
+    # Tie-break right: if equal, keep the later (right) index
+    pick_b = pick_b | ((val_a == val_b) & (idx_b >= idx_a))
+    val_out = tl.where(pick_b, val_b, val_a)
+    idx_out = tl.where(pick_b, idx_b, idx_a)
+    return val_out, idx_out
+
+
+@triton.jit
+def scan_part_max_kernel(
+    inp,
+    out,
+    out_indices,
+    partial_max,
+    partial_max_indices,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    NEED_PARTIAL: tl.constexpr,
+):
+    """Block-local cumulative maximum with optional partial output.
+
+    Each block loads BLOCK_SIZE elements, computes cummax via associative
+    scan, stores results, and optionally writes the block-maximum for
+    the fan-out phase.
+    """
+    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
+
+    # Load input, use -inf for masked positions
+    inp_vals = tl.load(inp + offset, mask=mask, other=float('-inf'))
+    inp_vals = inp_vals.to(tl.float32)
+    indices = offset
+
+    # Associative scan for cummax
+    result, cummax_idx = tl.associative_scan(
+        (inp_vals, indices), 0, _cummax_combine
+    )
+
+    tl.store(out + offset, result, mask=mask)
+    tl.store(out_indices + offset, cummax_idx, mask=mask)
+
+    if tl.constexpr(NEED_PARTIAL):
+        # Reduce to find block maximum (tie-break right)
+        block_max = tl.max(inp_vals, axis=0)
+        # Find rightmost index with block_max value
+        is_max = inp_vals == block_max
+        max_idx = tl.max(tl.where(is_max, indices, -1), axis=0)
+        # Only one thread writes the partial
+        if tl.arange(0, BLOCK_SIZE)[0] == 0:
+            tl.store(partial_max + pid, block_max)
+            tl.store(partial_max_indices + pid, max_idx)
+
+
+@triton.jit
+def add_base_max_kernel(
+    out,
+    out_indices,
+    partial_max,
+    partial_max_indices,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Fan-out phase: combine partial maxima from prior blocks.
+
+    For each block (except the first), loads the cumulative partial max
+    from the previous block and merges with local results.
+    """
+    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
+
+    out_vals = tl.load(out + offset, mask=mask)
+    out_idx = tl.load(out_indices + offset, mask=mask)
+
+    if pid > 0:
+        prev_max = tl.load(partial_max + pid - 1)
+        prev_idx = tl.load(partial_max_indices + pid - 1)
+
+        # Merge: if previous partial max >= local, use previous
+        use_prev = (prev_max > out_vals) | (
+            (prev_max == out_vals) & (prev_idx >= out_idx)
+        )
+        final_vals = tl.where(use_prev, prev_max, out_vals)
+        final_idx = tl.where(use_prev, prev_idx, out_idx)
+
+        tl.store(out + offset, final_vals, mask=mask)
+        tl.store(out_indices + offset, final_idx, mask=mask)

nbxtensor/ops/cummin.py

@@ -0,0 +1,105 @@
+"""Cumulative Minimum — pure @triton.jit kernels.
+
+Computes running minimum along a dimension, returning both values and indices.
+Uses scan-then-fan algorithm (mirror of cummax with minimum).
+
+Extracted from FlagGems (Apache 2.0).
+
+Three kernels:
+1. scan_part_min_kernel — local cummin within each block + partial min
+2. add_base_min_kernel — fan-out: combine partial minima with local results
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _cummin_combine(val_a, idx_a, val_b, idx_b):
+    """Associative combiner: keep (value, index) of the minimum."""
+    pick_b = val_b < val_a
+    # Tie-break right: if equal, keep the later (right) index
+    pick_b = pick_b | ((val_a == val_b) & (idx_b >= idx_a))
+    val_out = tl.where(pick_b, val_b, val_a)
+    idx_out = tl.where(pick_b, idx_b, idx_a)
+    return val_out, idx_out
+
+
+@triton.jit
+def scan_part_min_kernel(
+    inp,
+    out,
+    out_indices,
+    partial_min,
+    partial_min_indices,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    NEED_PARTIAL: tl.constexpr,
+):
+    """Block-local cumulative minimum with optional partial output.
+
+    Each block loads BLOCK_SIZE elements, computes cummin via associative
+    scan, stores results, and optionally writes the block-minimum for
+    the fan-out phase.
+    """
+    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
+
+    # Load input, use +inf for masked positions
+    inp_vals = tl.load(inp + offset, mask=mask, other=float('inf'))
+    inp_vals = inp_vals.to(tl.float32)
+    indices = offset
+
+    # Associative scan for cummin
+    result, cummin_idx = tl.associative_scan(
+        (inp_vals, indices), 0, _cummin_combine
+    )
+
+    tl.store(out + offset, result, mask=mask)
+    tl.store(out_indices + offset, cummin_idx, mask=mask)
+
+    if tl.constexpr(NEED_PARTIAL):
+        # Reduce to find block minimum (tie-break right)
+        block_min = tl.min(inp_vals, axis=0)
+        is_min = inp_vals == block_min
+        min_idx = tl.max(tl.where(is_min, indices, -1), axis=0)
+        if tl.arange(0, BLOCK_SIZE)[0] == 0:
+            tl.store(partial_min + pid, block_min)
+            tl.store(partial_min_indices + pid, min_idx)
+
+
+@triton.jit
+def add_base_min_kernel(
+    out,
+    out_indices,
+    partial_min,
+    partial_min_indices,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Fan-out phase: combine partial minima from prior blocks.
+
+    For each block (except the first), loads the cumulative partial min
+    from the previous block and merges with local results.
+    """
+    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
+
+    out_vals = tl.load(out + offset, mask=mask)
+    out_idx = tl.load(out_indices + offset, mask=mask)
+
+    if pid > 0:
+        prev_min = tl.load(partial_min + pid - 1)
+        prev_idx = tl.load(partial_min_indices + pid - 1)
+
+        # Merge: if previous partial min <= local, use previous
+        use_prev = (prev_min < out_vals) | (
+            (prev_min == out_vals) & (prev_idx >= out_idx)
+        )
+        final_vals = tl.where(use_prev, prev_min, out_vals)
+        final_idx = tl.where(use_prev, prev_idx, out_idx)
+
+        tl.store(out + offset, final_vals, mask=mask)
+        tl.store(out_indices + offset, final_idx, mask=mask)

nbxtensor/ops/cumsum.py

@@ -0,0 +1,244 @@
+"""Cumulative sum — pure @triton.jit kernels. Extracted from FlagGems.
+
+Two-pass scan-then-fan approach for large 1D cumsum:
+  scan_part_sum_kernel: per-block prefix sum + partial sums
+  add_base_sum_kernel: propagate partial sums across blocks
+
+Multi-dimensional (A,B,C) variants:
+  scan_part_sum_abc_kernel: per-block scan along B dim
+  add_base_sum_abc_kernel: propagate partial sums along B dim
+
+Row-wise persistent kernels:
+  reduce_then_scan_block_sum_kernel_row: block sums for row scan
+  reduce_then_scan_root_scan_kernel_row: persistent root scan
+  reduce_then_scan_block_scan_kernel_row: final block scan with prefix
+"""
+
+import triton
+import triton.language as tl
+
+
+@tl.constexpr
+def get_scan_accum_type(inp_dtype: tl.dtype) -> tl.dtype:
+    if inp_dtype.is_bf16() or inp_dtype.is_fp16():
+        return tl.float32
+    if inp_dtype.is_int():
+        return tl.int64
+    else:
+        return inp_dtype
+
+
+# --- 1D scan-then-fan kernels ---
+
+
+@triton.jit(do_not_specialize=["n_elements", "part_num"])
+def scan_part_sum_kernel(
+    inp,
+    out,
+    partial_sum,
+    n_elements,
+    part_num,
+    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
+
+    inp_ptrs = inp + offset
+    inp_vals = tl.load(inp_ptrs, mask=mask)
+    if (
+        tl.constexpr(inp_vals.dtype.is_int64())
+        or tl.constexpr(inp_vals.dtype.is_uint64())
+    ) or tl.constexpr(inp_vals.dtype.is_fp64()):
+        inp_vals = inp_vals
+    elif tl.constexpr(inp_vals.dtype.is_int()):
+        inp_vals = inp_vals.to(tl.int32)
+    else:
+        inp_vals = inp_vals.to(tl.float32)
+    result = tl.cumsum(inp_vals, axis=0)
+
+    part_sum_via_sum = tl.sum(inp_vals)
+
+    out_ptrs = out + offset
+    tl.store(out_ptrs, result, mask=mask)
+
+    partial_sum_ptrs = partial_sum + pid
+    tl.store(partial_sum_ptrs, part_sum_via_sum)
+
+
+@triton.jit(do_not_specialize=["n_elements", "part_num"])
+def add_base_sum_kernel(
+    out,
+    partial_sum,
+    n_elements,
+    part_num,
+    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
+
+    out_ptrs = out + offset
+    out_vals = tl.load(out_ptrs, mask=mask)
+
+    if pid > 0:
+        partial_sum_ptrs = partial_sum + pid - 1
+        last_part_sum_via_sum = tl.load(partial_sum_ptrs)
+
+        final_vals = out_vals + last_part_sum_via_sum
+        tl.store(out_ptrs, final_vals.to(out_vals.dtype), mask=mask)
+
+
+# --- Multi-dim (A,B,C) kernels ---
+
+
+@triton.jit(do_not_specialize=["part_num"])
+def scan_part_sum_abc_kernel(
+    inp,
+    out,
+    partial_sum,
+    B,
+    C,
+    part_num,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid_a = tl.program_id(0).to(tl.int64)
+    pid_b = tl.program_id(1).to(tl.int64)
+    pid_c = tl.program_id(2).to(tl.int64)
+
+    a_idx = pid_a
+    b_idx = pid_b.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id (register 58)
+    c_idx = pid_c
+
+    offset = a_idx * B * C + b_idx * C + c_idx
+    base_part_offset = a_idx * part_num * C + c_idx
+    part_offset = base_part_offset + pid_b * C
+
+    mask = b_idx < B
+    inp_ptrs = inp + offset
+    inp_vals = tl.load(inp_ptrs, mask=mask)
+    if (
+        tl.constexpr(inp_vals.dtype.is_int64())
+        or tl.constexpr(inp_vals.dtype.is_uint64())
+    ) or tl.constexpr(inp_vals.dtype.is_fp64()):
+        inp_vals = inp_vals
+    elif tl.constexpr(inp_vals.dtype.is_int()):
+        inp_vals = inp_vals.to(tl.int32)
+    else:
+        inp_vals = inp_vals.to(tl.float32)
+    result = tl.cumsum(inp_vals, axis=0)
+
+    part_sum_via_sum = tl.sum(inp_vals)
+
+    out_ptrs = out + offset
+    tl.store(out_ptrs, result, mask=mask)
+
+    partial_sum_ptrs = partial_sum + part_offset
+    tl.store(partial_sum_ptrs, part_sum_via_sum)
+
+
+@triton.jit(do_not_specialize=["part_num"])
+def add_base_sum_abc_kernel(
+    out,
+    partial_sum,
+    B,
+    C,
+    part_num,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid_a = tl.program_id(0).to(tl.int64)
+    pid_b = tl.program_id(1).to(tl.int64)
+    pid_c = tl.program_id(2).to(tl.int64)
+
+    a_idx = pid_a
+    b_idx = pid_b.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id (register 58)
+    c_idx = pid_c
+
+    base_offset = a_idx * B * C + c_idx
+    offset = base_offset + b_idx * C
+    base_part_offset = a_idx * part_num * C + c_idx
+    last_part_offset = base_part_offset + (pid_b - 1) * C
+
+    mask = b_idx < B
+    out_ptrs = out + offset
+    out_vals = tl.load(out_ptrs, mask=mask)
+
+    if pid_b > 0:
+        partial_sum_ptrs = partial_sum + last_part_offset
+        last_part_sum_via_sum = tl.load(partial_sum_ptrs)
+
+        final_vals = out_vals + last_part_sum_via_sum
+        tl.store(out_ptrs, final_vals.to(out_vals.dtype), mask=mask)
+
+
+# --- Row-wise persistent scan kernels ---
+
+
+@triton.jit
+def reduce_then_scan_block_sum_kernel_row(
+    in_ptr,
+    block_sum_ptr,
+    N,
+    tiles_per_cta,
+    TILE_SIZE: tl.constexpr,
+):
+    """Block sum: each CTA reduces its assigned tiles into a single sum."""
+    pid_n = tl.program_id(1).to(tl.int64)
+    pid_m = tl.program_id(0).to(tl.int64)
+    num_programs_n = tl.num_programs(1)
+    block_offset = (pid_n * (tiles_per_cta * TILE_SIZE)).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)
+    block_end = min(block_offset + tiles_per_cta * TILE_SIZE, N).to(tl.int32)
+
+    acc_dtype: tl.constexpr = get_scan_accum_type(in_ptr.type.element_ty)
+    acc = tl.zeros((TILE_SIZE,), dtype=acc_dtype)
+    for start in range(block_offset, block_end, TILE_SIZE):
+        offsets = start + tl.arange(0, TILE_SIZE)
+        x = tl.load(in_ptr + pid_m * N + offsets, mask=offsets < N).to(acc_dtype)
+        acc += x
+    block_sum = tl.sum(acc, 0)
+    tl.store(
+        block_sum_ptr + pid_m * num_programs_n + pid_n, block_sum, cache_modifier=".cg"
+    )
+
+
+@triton.jit
+def reduce_then_scan_root_scan_kernel_row(in_ptr, out_ptr, N, TILE_SIZE: tl.constexpr):
+    """Persistent scan of block sums (or small row)."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, TILE_SIZE)
+    mask = offsets < N
+    acc_dtype: tl.constexpr = get_scan_accum_type(in_ptr.type.element_ty)
+    x = tl.load(in_ptr + pid * N + offsets, mask=mask, other=0).to(acc_dtype)
+    out = tl.cumsum(x, 0)
+    tl.store(out_ptr + pid * N + offsets, out, mask=mask)
+
+
+@triton.jit
+def reduce_then_scan_block_scan_kernel_row(
+    in_ptr,
+    previous_sum_ptr,
+    out_ptr,
+    N,
+    num_tiles_n,
+    tiles_per_cta,
+    TILE_SIZE: tl.constexpr,
+):
+    """Final block scan: apply prefix from previous blocks, then local cumsum."""
+    pid_m = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+    block_offset = (pid_n * (tiles_per_cta * TILE_SIZE)).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)
+    block_end = min(block_offset + tiles_per_cta * TILE_SIZE, N).to(tl.int32)
+    acc_dtype: tl.constexpr = get_scan_accum_type(in_ptr.type.element_ty)
+
+    prefix = tl.load(
+        previous_sum_ptr + pid_m * num_tiles_n + pid_n - 1, mask=pid_n > 0, other=0
+    ).to(acc_dtype)
+    for start in range(block_offset, block_end, TILE_SIZE):
+        offsets = start + tl.arange(0, TILE_SIZE)
+        mask = offsets < N
+        x = tl.load(in_ptr + pid_m * N + offsets, mask=mask).to(acc_dtype)
+        tile_scan = prefix + tl.cumsum(x, 0)
+        prefix += tl.sum(x, 0)
+        tl.store(
+            out_ptr + pid_m * N + offsets, tile_scan, mask=mask, cache_modifier=".cg"
+        )