nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/tiling_sizes.py

@@ -0,0 +1,370 @@
+"""The TilingEngine's sizing half — torch-free, imported by Prism, the launch keys, the Triton
+wrappers and the TilingEngine's executor.
+
+One module for every caller: the Triton branch (the wrappers, the launch keys and the derived census through them)
+and NeuroBrix's PyTorch branch (Prism, the core executors and the TilingEngine) alike — the engine keeps no copy.
+
+Every size a memory split is cut by is computed here, from the library's `nbxtensor/ops/tiling.yml` (engine-wide
+values, each with its source) and from the values the caller reads in the active hardware
+profile (per-architecture values stay there). Pure arithmetic: Prism prices with it at plan time,
+the launch keys and the census derive with it, the wrappers and the executor cut with it — one
+definition each, so the plan, the key and the run agree.
+
+What it owns:
+  * the conv3d one-shot peak and its temporal chunk (a PLAN decision: `need` against the op
+    budget; the wrapper sizes the chunks with the same function);
+  * the conv2d band cut;
+  * the attention scores bound and its chunk rows;
+  * Prism's op budget and its op-level band factors (conv, fusion pair, rms_norm, residual chain);
+  * the tile overlaps (spatial and temporal);
+  * the in-place and residual-chain thresholds;
+  * the paged decode read's position tiles (the page size is sized from them, core/state/price).
+"""
+from __future__ import annotations
+
+import math
+from typing import Any, Sequence, Tuple
+
+
+# The position tiles the paged read is compiled for: powers of two, each dividing the largest, so a page sized from the
+# largest holds a whole number of tiles of every candidate. 32 is the vector decode kernel's judged tile (its locked
+# protocol of 2026-08-23: seg 256 / BLOCK_N 32 / 4 warps won at all three contexts, ops/decode_attn_vec.py,
+# kernels/wrappers.py `_decode_attn_vec`); 16 and 64 are its neighbours, the range the certified directory chooses in.
+PAGED_DECODE_TILE_CANDIDATES = (16, 32, 64)
+
+def _value(*path: str) -> Any:
+    """One value of nbxtensor/ops/tiling.yml, refused by name when absent."""
+    from nbxtensor.ops.tiling_policy import tiling_table   # the policy in force (nbxtensor.ops.tiling_policy.set_tiling)
+    node: Any = tiling_table()
+    for i, key in enumerate(path):
+        if not isinstance(node, dict) or key not in node:
+            raise KeyError(f"ZERO FALLBACK: nbxtensor/ops/tiling.yml states no {'.'.join(path[:i + 1])}")
+        node = node[key]
+    return node
+
+
+def _pow2_at_most(factor: int, cap: int) -> int:
+    """The smallest power of two >= `factor`, at most `cap` (a band factor aligned for halos)."""
+    p = 1
+    while p < factor and p < cap:
+        p *= 2
+    return p
+
+
+# --- conv3d: the one-shot peak and the temporal chunk ---------------------------------------------
+
+def conv3d_chunk_bytes() -> int:
+    """The per-transient bound of the chunked conv3d path (`conv3d.chunk_bytes`)."""
+    return int(_value("conv3d", "chunk_bytes"))
+
+
+def conv2d_band_bytes() -> int:
+    """The conv2d band-streaming threshold (`conv2d.band_bytes`)."""
+    return int(_value("conv2d", "band_bytes"))
+
+
+def _triple(v) -> Tuple[int, int, int]:
+    if isinstance(v, (list, tuple)):
+        v = list(v) + [v[-1]] * (3 - len(v))
+        return int(v[0]), int(v[1]), int(v[2])
+    return int(v), int(v), int(v)
+
+
+def conv3d_need(x_shape: Sequence[int], w_shape: Sequence[int], stride, padding, dilation,
+                in_bytes: int, out_bytes: int) -> Tuple[int, int]:
+    """(bytes the ONE-SHOT path holds at its peak, folded bytes per output frame). `need` is 0 when
+    no fold exceeds `conv3d_chunk_bytes()` — the one-shot path is taken whatever the budget. The
+    peak: the temporal pad copy, then the larger of (input fold + output fold [+ a band when the
+    output bands] [+ the kt accumulator]) and (two output folds [+ the accumulator]), and a fixed
+    slack (`conv3d.slack_bytes`). The wrapper runs conv3d as kt conv2d launches over the folded
+    batch axis (B*T_out frames); whether it chunks is Prism's decision, made from this `need`
+    against the op budget (decision 2, 2026-09-29: never from the driver's free bytes)."""
+    B, Cin, T, H, W = (int(d) for d in x_shape)
+    Cout, _Cg, kt, kh, kw = (int(d) for d in w_shape)
+    st, sh, sw = _triple(stride)
+    pt, ph, pw = _triple(padding)
+    dt, dh, dw = _triple(dilation)
+    t_out = (T + 2 * pt - dt * (kt - 1) - 1) // st + 1
+    if t_out <= 0:
+        return 0, 0
+    oh = (H + 2 * ph - dh * (kh - 1) - 1) // sh + 1
+    ow = (W + 2 * pw - dw * (kw - 1) - 1) // sw + 1
+    fold_in = B * t_out * Cin * H * W * in_bytes
+    fold_out = B * t_out * Cout * oh * ow * out_bytes
+    frame = max(fold_in, fold_out) // max(1, t_out)
+    if max(fold_in, fold_out) <= conv3d_chunk_bytes():
+        return 0, frame
+    return conv3d_one_shot_peak(x_shape, w_shape, stride, padding, dilation, in_bytes, out_bytes), frame
+
+
+def conv3d_one_shot_peak(x_shape: Sequence[int], w_shape: Sequence[int], stride, padding, dilation,
+                         in_bytes: int, out_bytes: int) -> int:
+    """The bytes the one-shot conv3d path (`_conv3d_via_conv2d`) holds at its peak BESIDE its input: the temporal pad
+    copy, then the larger of (input fold + output fold [+ a band] [+ the kt accumulator]) and (two output folds [+ the
+    accumulator]), and the slack. Its result is among them (one output fold). `conv3d_need` decides the chunking from
+    it; Prism's activation profile prices the op's in-op figure from it on the Triton engines — measured 2026-10-17 on
+    Wan2.1-T2V's VAE: 5,276 MB beside a 1,755 MB output, the formula's 3 x 1,755."""
+    B, Cin, T, H, W = (int(d) for d in x_shape)
+    Cout, _Cg, kt, kh, kw = (int(d) for d in w_shape)
+    st, sh, sw = _triple(stride)
+    pt, ph, pw = _triple(padding)
+    dt, dh, dw = _triple(dilation)
+    t_out = (T + 2 * pt - dt * (kt - 1) - 1) // st + 1
+    if t_out <= 0:
+        return 0
+    oh = (H + 2 * ph - dh * (kh - 1) - 1) // sh + 1
+    ow = (W + 2 * pw - dw * (kw - 1) - 1) // sw + 1
+    fold_in = B * t_out * Cin * H * W * in_bytes
+    fold_out = B * t_out * Cout * oh * ow * out_bytes
+    band = conv2d_band_bytes()
+    pad_bytes = B * Cin * (T + 2 * pt) * H * W * in_bytes if pt > 0 else 0
+    band_extra = band if fold_out > band else 0
+    acc_extra = fold_out if kt >= 2 else 0
+    return (pad_bytes + max(fold_out + fold_in + band_extra + acc_extra, 2 * fold_out + acc_extra)
+            + int(_value("conv3d", "slack_bytes")))
+
+
+def residual_chain_bands(H: int, tile_factor: int, halo: int):
+    """The bands a band-streamed residual chain runs over its spatial axis: [(h_start, h_end, h_in_start, h_in_end)]
+    — the output rows of each band and the input rows it reads (the band plus `halo` rows each side, cut at the
+    edges). ONE function for the chain that runs them (`kernels/ops/residual_chain.py`) and the census that records
+    the conv keys they form (a 418-row band of Sana 4K's VAE met no certified key, 2026-10-17): the recorded key is
+    the launcher's key."""
+    H, tile_factor, halo = int(H), max(1, int(tile_factor)), int(halo)
+    band_h = (H + tile_factor - 1) // tile_factor
+    bands = []
+    for i in range(tile_factor):
+        h_start, h_end = i * band_h, min((i + 1) * band_h, H)
+        if h_start >= h_end:
+            break
+        bands.append((h_start, h_end, max(0, h_start - halo), min(H, h_end + halo)))
+    return bands
+
+
+def chunk_frames(frame_bytes: int) -> int:
+    """Output frames per chunk: each chunk's folded transient under `conv3d_chunk_bytes()`."""
+    return max(1, int(conv3d_chunk_bytes() // max(1, frame_bytes)))
+
+
+# --- conv2d: the band cut -------------------------------------------------------------------------
+
+def conv2d_band_rows(N: int, out_c: int, out_h: int, out_w: int, out_bytes: int,
+                     band_bytes: int) -> Tuple[int, int]:
+    """(output rows of one band, bytes of one output row) of conv2d band streaming: bands of
+    about `band_bytes / conv2d.band_budget_divisor`, evenly cut. A band >= `out_h` means one row
+    is already over the budget — the caller refuses it."""
+    band_target = max(1, band_bytes // int(_value("conv2d", "band_budget_divisor")))
+    row_bytes = N * out_c * out_w * out_bytes
+    rows_per_band = max(1, band_target // max(1, row_bytes))
+    tile_factor = max(1, (out_h + rows_per_band - 1) // rows_per_band)
+    return (out_h + tile_factor - 1) // tile_factor, row_bytes
+
+
+def tiled_conv2d_bands(IH: int, out_h: int, kh: int, sh: int, dh: int, pad_h: int, tile_factor: int):
+    """The bands of Prism's op-level tiled conv (real halo): per band (oh_start, oh_end, in_start, in_end, pad_top,
+    pad_bot, skip) — the output rows, the input rows read (clamped), the image-edge padding added, the band's leading
+    conv rows to skip. ONE cut for the run (`nbxtensor/ops/fused_upsample_conv.tiled_conv2d` and its torch
+    twin), the derived census (`launch_keys.tiled_conv2d_launches`) and Prism's price of the band
+    (`tiled_conv2d_band_bytes`). The image-edge padding is `max(0, -read_start)` alone: `read_start` already carries
+    -pad_h (P-NBX-TILED-CONV2D-SMALL-SCALE 2026-05-14). The halo is rounded up to a whole number of strides: the band's
+    first conv row must be an output row (at stride 2 a 1-row halo misaligned every internal band, 2026-09-29)."""
+    tf = max(1, int(tile_factor))
+    band_oh = (out_h + tf - 1) // tf
+    halo_h = -(-((kh - 1) * dh // 2) // sh) * sh
+    bands = []
+    for oh_start in range(0, out_h, band_oh):
+        oh_end = min(oh_start + band_oh, out_h)
+        halo_top = 0 if oh_start == 0 else halo_h
+        halo_bot = 0 if oh_end == out_h else halo_h
+        read_start = oh_start * sh - pad_h - halo_top
+        read_end = (oh_end - 1) * sh + dh * (kh - 1) + 1 - pad_h + halo_bot
+        start, end = max(0, read_start), min(IH, read_end)
+        if end <= start:
+            continue
+        bands.append((oh_start, oh_end, start, end, max(0, -read_start), max(0, read_end - IH),
+                      halo_top // sh))
+    return bands
+
+
+def tiled_conv2d_band_bytes(x_shape: Sequence[int], w_shape: Sequence[int], stride, padding, dilation,
+                            tile_factor: int, in_bytes: int, out_bytes: int) -> int:
+    """What a tiled 2-D conv holds BESIDE its input and its whole output while its largest band runs: the band's input
+    rows copied into one buffer at the conv's compute width `in_bytes` (padded by the image-edge rows, and by pad_w on
+    both sides whenever any padding applies — one buffer, the slice copied into its interior) and the band's output at
+    `out_bytes`, before it is copied into the whole output (`kernels/ops/fused_upsample_conv._tiled_conv2d_spatial_*`,
+    both engines). Unpriced, HAT-L x4 at 2048 px on a 16 GB card ran out of memory asking a band's 482 MiB beside the
+    1 446 MiB band input (2026-10-22)."""
+    N, C, IH, IW = (int(v) for v in x_shape)
+    O, _cg, kh, kw = (int(v) for v in w_shape)
+    sh, sw = _pair(stride)
+    ph, pw = _pair(padding)
+    dh, dw = _pair(dilation)
+    out_h = (IH + 2 * ph - dh * (kh - 1) - 1) // sh + 1
+    out_w = (IW + 2 * pw - dw * (kw - 1) - 1) // sw + 1
+    worst = 0
+    for o0, o1, start, end, pt, pb, _skip in tiled_conv2d_bands(IH, out_h, kh, sh, dh, ph, tile_factor):
+        padded = pt > 0 or pb > 0 or pw > 0
+        rows_in = end - start + pt + pb
+        width_in = IW + 2 * pw if padded else IW
+        worst = max(worst, N * C * rows_in * width_in * int(in_bytes) + N * O * (o1 - o0) * out_w * int(out_bytes))
+    return worst
+
+
+def _pair(v) -> Tuple[int, int]:
+    if isinstance(v, (list, tuple)):
+        return (int(v[0]), int(v[-1]))
+    return (int(v), int(v))
+
+
+# --- attention: the scores bound, the chunk rows, the device budget -------------------------------
+
+#: The math route's scores are fp32: four bytes a score (the width of the dtype, not a size choice).
+_SCORES_BYTES = 4
+
+
+def sdpa_scores_bytes(batch: int, nheads: int, Tq: int, Tk: int) -> int:
+    """The bytes of an attention's fp32 scores."""
+    return batch * nheads * Tq * Tk * _SCORES_BYTES
+
+
+def sdpa_scores_bound(budget_bytes: int, head_dim_pow2: bool) -> int:
+    """The fp32 scores bound an attention routes on: the device's budget; a non-power-of-two head
+    dim whose architecture declares none routes on `sdpa.non_pow2_head_scores_bytes`."""
+    if head_dim_pow2:
+        return budget_bytes
+    return budget_bytes or int(_value("sdpa", "non_pow2_head_scores_bytes"))
+
+
+def sdpa_chunk_rows(bound: int, batch: int, nheads: int, Tq: int, Tk: int,
+                    min_chunk_rows: int, max_chunks: int) -> int:
+    """Query rows per chunk that keep each chunk's fp32 scores within `bound`, aligned to the
+    arch's row block (`memory.sdpa_math_min_chunk_rows`); 0 when the shape cannot chunk within
+    the arch's chunk ceiling (`memory.sdpa_math_max_chunks`)."""
+    if not min_chunk_rows:
+        return 0
+    rows = (bound // (batch * nheads * Tk * _SCORES_BYTES)) // min_chunk_rows * min_chunk_rows
+    if rows >= min_chunk_rows and -(-Tq // rows) <= max_chunks:
+        return rows
+    return 0
+
+
+def sdpa_chunk_plan(bound: int, batch: int, nheads: int, Tq: int, Tk: int,
+                    min_chunk_rows: int) -> Tuple[int, int]:
+    """(slices per chunk, query rows per chunk) that keep each chunk's fp32 scores within `bound`: the rows as many as
+    fit with one (batch, head) slice, a multiple of the arch's row block (`memory.sdpa_math_min_chunk_rows`) or the
+    whole query; then as many slices as fit beside them. Every slice and every query row is its own softmax, so any
+    such cut is exact. (batch x nheads, Tq) when the whole scores fit; (0, 0) when one slice of one row block exceeds
+    `bound`."""
+    slices = batch * nheads
+    if sdpa_scores_bytes(batch, nheads, Tq, Tk) <= bound:
+        return slices, Tq
+    per_row = Tk * _SCORES_BYTES
+    rows = min(Tq, (bound // per_row) // min_chunk_rows * min_chunk_rows) if min_chunk_rows else 0
+    if rows <= 0:
+        return 0, 0
+    return max(1, min(slices, bound // (rows * per_row))), rows
+
+
+def sdpa_device_scores_budget(base_bytes: int, device_fraction: float, device_memory_mb) -> int:
+    """The scores budget on one device: min(the arch's byte budget, `device_fraction` of THIS
+    device's memory) — the 2026-08-31 xlong-prefill OOM proved a per-arch byte budget alone is
+    wrong on heterogeneous rigs. No fraction, or no device memory known: the arch's budget."""
+    if not base_bytes:
+        return 0
+    if not device_fraction or device_memory_mb is None:
+        return base_bytes
+    return min(base_bytes, int(device_memory_mb * 1024 * 1024 * device_fraction))
+
+
+def sdpa_math_scores_budget(base_bytes: int, device_fraction: float, device_memory_mb) -> int:
+    """The math route's scores budget on the device that executes an attention: 0 where the arch states no byte budget
+    (`memory.sdpa_math_max_scores_bytes`: a power-of-two head dim never routes on size), the arch's budget where it
+    states no device fraction (`memory.sdpa_math_scores_device_fraction`), else `sdpa_device_scores_budget`. ONE rule
+    for the wrappers' route (`wrappers._sdpa_math_scores_budget_bytes_for`) and the scratch Prism prices for it."""
+    if not base_bytes:
+        return 0
+    if not device_fraction:
+        return base_bytes
+    return sdpa_device_scores_budget(base_bytes, device_fraction, device_memory_mb)
+
+
+# --- Prism: the op budget and the op-level band factors -------------------------------------------
+
+def op_budget_fraction() -> float:
+    """The fraction of a card an op's transient is measured against (`op_level.op_budget_card_fraction`)."""
+    return float(_value("op_level", "op_budget_card_fraction"))
+
+
+def op_budget_bytes(card_bytes: int, resident_bytes: int) -> int:
+    """An op's budget on its card: the op budget fraction of the card minus the component's
+    resident weights."""
+    return int(card_bytes * op_budget_fraction()) - resident_bytes
+
+
+def conv_band_factor(total_bytes: int, kept_bytes: int, card_bytes: int) -> int:
+    """Bands an overflowing convolution (or upsample->conv pair) runs in: its tiled bytes over
+    one band's budget (`conv_band_card_fraction` of the card minus the output kept whole, at
+    least `band_budget_floor_bytes`), at least `min_band_factor`, rounded up to a power of two
+    at most `max_band_factor`."""
+    budget = int(float(_value("op_level", "conv_band_card_fraction")) * card_bytes) - kept_bytes
+    budget = max(budget, int(_value("op_level", "band_budget_floor_bytes")))
+    factor = max(int(_value("op_level", "min_band_factor")), math.ceil(total_bytes / budget))
+    return _pow2_at_most(factor, int(_value("op_level", "max_band_factor")))
+
+
+def upsample_overflows(out_bytes: int, card_bytes: int) -> bool:
+    """An upsample output worth fusing into its conv even when the conv does not overflow."""
+    return out_bytes > float(_value("op_level", "upsample_overflow_card_fraction")) * card_bytes
+
+
+def rms_norm_overflows(out_bytes: int, card_bytes: int) -> bool:
+    """An rms_norm whose output runs in H bands."""
+    return out_bytes > float(_value("op_level", "rms_norm", "overflow_card_fraction")) * card_bytes
+
+
+def rms_norm_band_factor(out_bytes: int, card_bytes: int) -> int:
+    """Bands an overflowing rms_norm runs in: its output over `rms_norm.band_card_fraction` of
+    the card, at least `rms_norm.min_band_factor`, a power of two, at most `rms_norm.max_band_factor`."""
+    budget = int(float(_value("op_level", "rms_norm", "band_card_fraction")) * card_bytes)
+    factor = max(int(_value("op_level", "rms_norm", "min_band_factor")), math.ceil(out_bytes / budget))
+    factor = _pow2_at_most(factor, int(_value("op_level", "max_band_factor")))
+    return min(factor, int(_value("op_level", "rms_norm", "max_band_factor")))
+
+
+def residual_chain_band_factor(bytes_fp32: int) -> int:
+    """Bands a residual chain streams in: its half-width bytes (the fp32 estimate halved) over
+    `residual_chain.band_target_bytes`, at least `residual_chain.min_band_factor`, a power of two at most
+    `residual_chain.max_band_factor`."""
+    half_bytes = bytes_fp32 // 2
+    factor = max(int(_value("op_level", "residual_chain", "min_band_factor")),
+                 math.ceil(half_bytes / int(_value("op_level", "residual_chain", "band_target_bytes"))))
+    return _pow2_at_most(factor, int(_value("op_level", "residual_chain", "max_band_factor")))
+
+
+def residual_chain_min_base_bytes_fp32() -> int:
+    """The base-tensor size (fp32) from which a residual chain is detected."""
+    return int(_value("op_level", "residual_chain", "min_base_bytes_fp32"))
+
+
+# --- overlaps -------------------------------------------------------------------------------------
+
+def spatial_halo(extent: int) -> int:
+    """The proportional spatial overlap of a tile of `extent`: extent // `overlap.spatial_divisor`."""
+    return extent // int(_value("overlap", "spatial_divisor"))
+
+
+def spatial_overlap(extent: int) -> int:
+    """A spatial tile's overlap: its proportional halo, at least `overlap.spatial_min`."""
+    return max(int(_value("overlap", "spatial_min")), spatial_halo(extent))
+
+
+def temporal_halo(t_tile: int) -> int:
+    """The proportional temporal overlap of a tile of `t_tile` frames: t_tile // `overlap.temporal_divisor`."""
+    return t_tile // int(_value("overlap", "temporal_divisor"))
+
+
+# --- in place -------------------------------------------------------------------------------------
+
+def inplace_min_bytes() -> int:
+    """The output size from which an add or an element-wise unary writes into its dead input."""
+    return int(_value("inplace", "min_bytes"))

nbxtensor/ops/topk.py

@@ -0,0 +1,267 @@
+"""Top-K — pure @triton.jit kernels. Extracted from FlagGems.
+
+Two-stage approach:
+  topk_stage1_kernel: per-chunk top-k selection via iterative max/min
+  topk_stage2_kernel: merge chunks via bitonic sort
+
+Supporting JIT functions:
+  _get_finfo_val: dtype min/max constants for masking
+  _compare_and_swap: bitonic compare-and-swap with index tracking
+  _bitonic_merge: single merge stage of bitonic sort
+  argsort: full bitonic argsort
+"""
+
+import triton
+import triton.language as tl
+import triton.language.core as core
+
+try:
+    from triton.language.standard import _log2, zeros_like
+except ImportError:
+    pass
+
+
+# --- Dtype limit constants (inline, no torch) ---
+
+_MIN_FLOAT32_VAL = tl.constexpr(-3.4028235e+38)
+_MAX_FLOAT32_VAL = tl.constexpr(3.4028235e+38)
+_MIN_FLOAT16_VAL = tl.constexpr(-65504.0)
+_MAX_FLOAT16_VAL = tl.constexpr(65504.0)
+_MIN_BFLOAT16_VAL = tl.constexpr(-3.3895314e+38)
+_MAX_BFLOAT16_VAL = tl.constexpr(3.3895314e+38)
+_MIN_INT32_VAL = tl.constexpr(-2147483648)
+_MAX_INT32_VAL = tl.constexpr(2147483647)
+
+
+@triton.jit
+def _get_finfo_val(
+    dtype,
+    return_max,
+):
+    """Saturating limit for `dtype`, as a Python float literal.
+
+    The limit MUST match the dtype of the buffer being loaded, not the
+    dtype the kernel computes in. `tl.load(..., other=v)` materialises
+    `v` in the POINTER's element type, so an fp32 limit handed to an
+    fp16 buffer overflows and lands as -inf. That is not a benign
+    sentinel here: `_compare_and_swap` permutes with a masked sum, and
+    `0 * -inf` is NaN, which loses every comparison and collapses the
+    sort network (measured: pad slots surfacing into the top-k with
+    their INT32_MIN sentinel index, whenever chunk_num*k was not a
+    power of two).
+
+    Literals rather than `tl.constexpr(torch.finfo(...).max)` — the
+    upstream FlagGems form — because R33 forbids torch here and the
+    constexpr wrapper mixes fp64 into the comparison. Dropping the
+    dtype branching along with the torch call was the regression.
+    """
+    if dtype is tl.float16:
+        if return_max:
+            return 65504.0
+        else:
+            return -65504.0
+    elif dtype is tl.bfloat16:
+        if return_max:
+            return 3.3895314e+38
+        else:
+            return -3.3895314e+38
+    else:
+        if return_max:
+            return 3.4028235e+38
+        else:
+            return -3.4028235e+38
+
+
+# --- Stage 1: per-chunk top-k ---
+
+
+@triton.jit
+def topk_stage1_kernel(
+    y_ptr,
+    index_ptr,
+    x_ptr,
+    k,
+    N: tl.constexpr,
+    CHUNK_SIZE: tl.constexpr,
+    DESCENDING: tl.constexpr,
+):
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_chunk_idx = tl.program_id(1).to(tl.int64)
+    chunk_num = tl.num_programs(1)
+
+    y_ptr += cur_batch * chunk_num * k + cur_chunk_idx * k
+    index_ptr += cur_batch * chunk_num * k + cur_chunk_idx * k
+
+    chunk_offset = cur_chunk_idx * CHUNK_SIZE
+    x_ptr += cur_batch * N + chunk_offset
+
+    cols = tl.arange(0, CHUNK_SIZE)
+    mask = (chunk_offset + cols) < N
+
+    mask_val = _get_finfo_val(x_ptr.dtype.element_ty, return_max=not DESCENDING)
+    x_val = tl.load(x_ptr + cols, mask=mask, other=mask_val).to(tl.float32)
+    for k_idx in range(k):
+        if DESCENDING:
+            chunk_select_val = tl.max(x_val)
+            chunk_select_idx = tl.argmax(x_val, axis=0)
+        else:
+            chunk_select_val = tl.min(x_val)
+            chunk_select_idx = tl.argmin(x_val, axis=0)
+
+        tl.store(y_ptr + k_idx, chunk_select_val)
+        tl.store(index_ptr + k_idx, chunk_select_idx + chunk_offset)
+
+        if DESCENDING:
+            mask_v = tl.full(x_val.shape, float('-inf'), dtype=tl.float32)
+        else:
+            mask_v = tl.full(x_val.shape, float('inf'), dtype=tl.float32)
+        x_val = tl.where(cols == chunk_select_idx, mask_v, x_val)
+
+
+# --- Bitonic sort helpers ---
+
+
+@triton.jit
+def _compare_and_swap(x, ids, flip, i: core.constexpr, n_dims: core.constexpr):
+    n_outer: core.constexpr = x.numel >> n_dims
+    shape: core.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)]
+
+    y = core.reshape(x, shape)
+    y_idx = core.reshape(ids, shape)
+
+    # Slice left/right with stride 2**(n_dims - i - 1).
+    #
+    # SELECT the unwanted lane away rather than multiplying it by zero.
+    # The upstream form is `tl.sum(y * (1 - mask), 1)`, which is exact
+    # for finite values but yields NaN on any +-inf, because 0 * inf is
+    # NaN. Selecting keeps the identity (-inf + 0 == -inf), so a
+    # non-finite entry sorts to its extreme instead of poisoning the
+    # comparison network. Callers do put -inf in the data: the sampler
+    # masks rejected logits with -inf and then takes a top-k over the
+    # result.
+    mask = core.arange(0, 2)[None, :, None]
+    left = core.broadcast_to(
+        tl.sum(core.where(mask == 0, y, zeros_like(y)), 1)[:, None, :], shape
+    ).to(x.dtype)
+    right = core.broadcast_to(
+        tl.sum(core.where(mask == 1, y, zeros_like(y)), 1)[:, None, :], shape
+    ).to(x.dtype)
+    left = core.reshape(left, x.shape)
+    right = core.reshape(right, x.shape)
+
+    left_idx = core.broadcast_to(
+        tl.sum(core.where(mask == 0, y_idx, zeros_like(y_idx)), 1)[:, None, :], shape
+    ).to(ids.dtype)
+    right_idx = core.broadcast_to(
+        tl.sum(core.where(mask == 1, y_idx, zeros_like(y_idx)), 1)[:, None, :], shape
+    ).to(ids.dtype)
+    left_idx = core.reshape(left_idx, ids.shape)
+    right_idx = core.reshape(right_idx, ids.shape)
+
+    # Compare-and-swap
+    if core.constexpr(x.dtype.primitive_bitwidth) == 8:
+        idtype = core.int8
+    elif core.constexpr(x.dtype.primitive_bitwidth) == 16:
+        idtype = core.int16
+    elif core.constexpr(x.dtype.primitive_bitwidth) == 32:
+        idtype = core.int32
+    elif core.constexpr(x.dtype.primitive_bitwidth) == 64:
+        idtype = core.int64
+    else:
+        raise ValueError("Unsupported dtype")
+
+    ileft = left.to(idtype, bitcast=True)
+    iright = right.to(idtype, bitcast=True)
+    ix = x.to(idtype, bitcast=True)
+
+    cond = (left > right) ^ flip
+    ret = ix ^ core.where(cond, ileft ^ iright, zeros_like(ix))
+
+    if core.constexpr(ids.dtype.primitive_bitwidth) == 8:
+        idx_dtype = core.int8
+    elif core.constexpr(ids.dtype.primitive_bitwidth) == 16:
+        idx_dtype = core.int16
+    elif core.constexpr(ids.dtype.primitive_bitwidth) == 32:
+        idx_dtype = core.int32
+    elif core.constexpr(ids.dtype.primitive_bitwidth) == 64:
+        idx_dtype = core.int64
+    else:
+        raise ValueError("Unsupported dtype")
+
+    ileft_idx = left_idx.to(idx_dtype, bitcast=True)
+    iright_idx = right_idx.to(idx_dtype, bitcast=True)
+    ix_idx = ids.to(idx_dtype, bitcast=True)
+    ret_idx = ix_idx ^ core.where(cond, ileft_idx ^ iright_idx, zeros_like(ix_idx))
+
+    return ret.to(x.dtype, bitcast=True), ret_idx.to(ids.dtype, bitcast=True)
+
+
+@triton.jit
+def _bitonic_merge(
+    x, ids, stage: core.constexpr, order: core.constexpr, n_dims: core.constexpr
+):
+    """Bitonic merge stage.
+
+    order: 0=ascending, 1=descending, 2=alternating
+    """
+    n_outer: core.constexpr = x.numel >> n_dims
+    core.static_assert(stage <= n_dims)
+
+    if order == 2:
+        shape: core.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage]
+        flip = core.reshape(
+            core.broadcast_to(core.arange(0, 2)[None, :, None], shape), x.shape
+        )
+    else:
+        flip = order
+    for i in core.static_range(stage):
+        x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
+    return x, ids
+
+
+@triton.jit
+def argsort(x, ids, dim: tl.constexpr, descending: core.constexpr):
+    _dim: core.constexpr = dim
+    n_dims: core.constexpr = _log2(x.shape[_dim])
+    for i in core.static_range(1, n_dims + 1):
+        x, ids = _bitonic_merge(x, ids, i, 2 if i < n_dims else descending, n_dims)
+    return x, ids
+
+
+# --- Stage 2: merge chunks via bitonic sort ---
+
+
+@triton.jit
+def topk_stage2_kernel(
+    y_ptr,
+    index_ptr,
+    chunk_x,
+    chunk_index,
+    sort_dim: tl.constexpr,
+    k: tl.constexpr,
+    N: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    DESCENDING: tl.constexpr,
+):
+    cur_batch = tl.program_id(0).to(tl.int64)
+    chunk_x += cur_batch * N
+    chunk_index += cur_batch * N
+    y_ptr += cur_batch * k
+    index_ptr += cur_batch * k
+
+    cols = tl.arange(0, BLOCK_SIZE)
+    mask = cols < N
+
+    mask_val = _get_finfo_val(chunk_x.dtype.element_ty, return_max=not DESCENDING)
+    mask_index_val = _MIN_INT32_VAL if DESCENDING else _MAX_INT32_VAL
+
+    chunk_x_val = tl.load(chunk_x + cols, mask=mask, other=mask_val).to(tl.float32)
+    chunk_index_val = tl.load(chunk_index + cols, mask=mask, other=mask_index_val).to(
+        tl.int32
+    )
+
+    sorted_chunk_x, sorted_chunk_index = argsort(
+        chunk_x_val, chunk_index_val, 0, descending=DESCENDING
+    )
+    tl.store(y_ptr + cols, sorted_chunk_x, mask=cols < k)
+    tl.store(index_ptr + cols, sorted_chunk_index, mask=cols < k)