TopK Bitonic MiniMax-M3

"""Top-K via 2-pass Triton kernel with packed (val, idx) top-k.

Strategy
--------
- Sort and index-extract are fused by packing (val_as_ordered_int, idx) into
  a single int64 per element. tl.topk on the packed array gives top-k (val,
  idx) pairs in one shot.
- One block per row, single launch.
- For rows that fit in one tile: load whole row, topk, output.
- For rows that don't fit (1, 131072, 64): chunked phase-1 (per-block topk)
  followed by a tiny phase-2 merge.

The float -> ordered-int trick: for ASCENDING sort of fp32, larger float ==
larger int. We use
    bits ^ ((bits >> 31) & 0x7FFFFFFF)
which keeps positives unchanged (already monotone as int) and flips the
magnitude bits of negatives (so larger-magnitude negatives = smaller ints).
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl


# Module-level constants to match reference.py
batch = 64
n = 8192
k = 8


def _next_pow2(x: int) -> int:
    if x <= 1:
        return 1
    return 1 << (x - 1).bit_length()


@triton.jit
def _f32_to_ord(v):
    bits = v.to(tl.int32, bitcast=True)
    sign_mask = (bits >> 31) & 0x7FFFFFFF
    return bits ^ sign_mask


@triton.jit
def _ord_to_f32(v):
    sign_mask = (v >> 31) & 0x7FFFFFFF
    bits = v ^ sign_mask
    return bits.to(tl.float32, bitcast=True)


@triton.jit
def topk_phase1_kernel(
    x_ptr,
    out_packed_ptr,   # (batch, num_chunks, K_PAD)
    n,
    CHUNK: tl.constexpr,
    K_PAD: tl.constexpr,
):
    row = tl.program_id(0)
    chunk = tl.program_id(1)

    chunk_start = chunk * CHUNK
    offs = chunk_start + tl.arange(0, CHUNK)
    mask = offs < n

    vals = tl.load(x_ptr + row * n + offs, mask=mask, other=-float('inf'))
    val_ord = _f32_to_ord(vals)

    idxs = offs.to(tl.int64)
    packed = (val_ord.to(tl.int64) << 32) | (idxs & 0xFFFFFFFF)
    # For masked positions: a value that sorts to the END of a descending
    # sort. We use the bit pattern of INT64_MIN. Triton's type checker rejects
    # raw 0x8000...0000 because it's out of int64 range; bitcast from uint64
    # is the clean workaround.
    masked_value = tl.full((), 0x8000000000000000, dtype=tl.uint64).to(tl.int64, bitcast=True)
    packed = tl.where(mask, packed, masked_value)

    top = tl.topk(packed, K_PAD)

    out_offs = tl.arange(0, K_PAD)
    out_base = (row * tl.num_programs(1) + chunk) * K_PAD
    tl.store(out_packed_ptr + out_base + out_offs, top)


@triton.jit
def topk_phase2_kernel(
    in_packed_ptr,    # (batch, num_chunks * K_PAD)
    out_val_ptr,      # (batch, K_PAD)
    out_idx_ptr,      # (batch, K_PAD)
    num_chunks,
    k: tl.constexpr,
    K_PAD: tl.constexpr,
    BLOCK: tl.constexpr,
):
    row = tl.program_id(0)

    offs = tl.arange(0, BLOCK)
    nc_times_kp = num_chunks * K_PAD
    mask = offs < nc_times_kp

    masked_value = tl.full((), 0x8000000000000000, dtype=tl.uint64).to(tl.int64, bitcast=True)
    packed = tl.load(
        in_packed_ptr + row * nc_times_kp + offs,
        mask=mask, other=masked_value,
    )

    top = tl.topk(packed, K_PAD)

    val_ord = (top >> 32).to(tl.int32)
    idx = (top & 0xFFFFFFFF).to(tl.int64)
    val = _ord_to_f32(val_ord)

    out_offs = tl.arange(0, K_PAD)
    out_mask = out_offs < k
    tl.store(out_val_ptr + row * k + out_offs, val, mask=out_mask)
    tl.store(out_idx_ptr + row * k + out_offs, idx, mask=out_mask)


class Model(nn.Module):
    """Top-k over the last dim of a 2D tensor."""

    def __init__(self, batch: int, n: int, k: int):
        super().__init__()
        self.batch, self.n, self.k = batch, n, k
        self.register_buffer("_dummy", torch.zeros(1))

        self.K_PAD = max(_next_pow2(k), 1)

        # Pick a chunk size: prefer to fit the entire row in one tile if we
        # can. Otherwise, use 16384 (one block does ~64KB int64 = 128KB
        # shared mem, comfortable under the 228KB H100 cap).
        if n <= 16384:
            self.CHUNK = _next_pow2(n)
        else:
            self.CHUNK = 16384

        self.num_chunks = (n + self.CHUNK - 1) // self.CHUNK

    def forward(self, x: torch.Tensor):
        # Phase 1: per-block top-K of CHUNK elements
        packed_buf = torch.empty(
            self.batch, self.num_chunks, self.K_PAD,
            dtype=torch.int64, device=x.device,
        )
        topk_phase1_kernel[(self.batch, self.num_chunks)](
            x, packed_buf, self.n, self.CHUNK, self.K_PAD,
        )

        # Phase 2: merge candidates across chunks
        BLOCK = max(_next_pow2(self.num_chunks * self.K_PAD), 2)
        out_val = torch.empty(self.batch, self.K_PAD, dtype=torch.float32, device=x.device)
        out_idx = torch.empty(self.batch, self.K_PAD, dtype=torch.int64, device=x.device)

        topk_phase2_kernel[(self.batch,)](
            packed_buf, out_val, out_idx,
            self.num_chunks, self.k, self.K_PAD, BLOCK,
        )

        return out_val[:, :self.k].contiguous(), out_idx[:, :self.k].contiguous()


def get_inputs():
    x = torch.randn(batch, n, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [batch, n, k]