diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000..329287c
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,11 @@
+include setup.py
+include csrc/*.cpp csrc/*.h csrc/build_utils.py
+recursive-include csrc/cuda_kernels *.cu *.cuh *.h *.hpp *.inl
+recursive-include csrc/3rdparty/cutlass/include *
+include csrc/3rdparty/cutlass/LICENSE.txt
+recursive-include csrc/3rdparty/kerutils/include *
+include csrc/3rdparty/kerutils/README.md
+include tests/__init__.py tests/test_build_plan.py
+recursive-include tests/kernelkit *.py
+include scripts/generate_instantiations.py
+global-exclude __pycache__ *.py[cod]
diff --git a/csrc/api.cpp b/csrc/api.cpp
index c6e238f..57b2eea 100644
--- a/csrc/api.cpp
+++ b/csrc/api.cpp
@@ -5,28 +5,48 @@
 
 #include <torch/csrc/stable/accelerator.h>
 #include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/util/Exception.h>
 #include <torch/headeronly/util/shim_utils.h>
 
-#include <cuda_bf16.h>
 #include <cuda_runtime_api.h>
 
+#include <climits>
+#include <cmath>
 #include <cstdint>
-#include <cstdlib>
 #include <format>
+#include <limits>
 #include <optional>
 #include <tuple>
 
-#include "dispatch_utils.h"
 #include "stable_tensor_checks.h"
+#include "structs.h"
 
+#ifndef DEEP_SELECT_BUILD_SM100
+#define DEEP_SELECT_BUILD_SM100 1
+#endif
+#ifndef DEEP_SELECT_BUILD_SM120
+#define DEEP_SELECT_BUILD_SM120 0
+#endif
+#ifndef DEEP_SELECT_BUILD_SM121
+#define DEEP_SELECT_BUILD_SM121 0
+#endif
+
+#if DEEP_SELECT_BUILD_SM100
+#include <cuda_bf16.h>
+#include "dispatch_utils.h"
 #include "cuda_kernels/config.h"
 #include "cuda_kernels/v3/topk_select.h"
 #include "cuda_kernels/v3_fp32/topk_select.h"
 #include "cuda_kernels/v3_cluster/topk_select.h"
+#endif
+#if DEEP_SELECT_BUILD_SM120 || DEEP_SELECT_BUILD_SM121
+#include "cuda_kernels/sm120/topk_select.h"
+#endif
 
 using deep_select::Tensor;
+using torch::headeronly::ScalarType;
 
 void topk(
     const Tensor& input,
@@ -43,15 +63,22 @@ void topk(
     bool return_value,
     bool abort_when_nan_found
 ) {
-    int64_t batch_size = input.size(0);
-    int64_t vocab_size = input.size(1);
-    ScalarType value_t = input.scalar_type();
-    ScalarType output_index_t = output_index.scalar_type();
+    STD_TORCH_CHECK(input.dim() == 2, "input must have two dimensions");
+    const int64_t batch_size = input.size(0);
+    const int64_t vocab_size = input.size(1);
+    const ScalarType value_t = input.scalar_type();
+    const ScalarType output_index_t = output_index.scalar_type();
 
-    STD_TORCH_CHECK(topk > 0, "topk must > 0");
+    STD_TORCH_CHECK(batch_size >= 0 && batch_size <= INT32_MAX, "batch_size must fit int32");
+    STD_TORCH_CHECK(vocab_size >= 0 && vocab_size < MAX_VOCAB_SIZE, "vocab_size must be < 2^23");
+    STD_TORCH_CHECK(topk > 0 && topk <= 4096, "topk must be > 0 and <= 4096");
+    STD_TORCH_CHECK(value_t == ScalarType::BFloat16 || value_t == ScalarType::Float, "input dtype must be bfloat16 or float32");
+    STD_TORCH_CHECK(output_index_t == ScalarType::Int || output_index_t == ScalarType::Long, "indices dtype must be int32 or int64");
+    STD_TORCH_CHECK(idx_oob_fill_value >= INT32_MIN && idx_oob_fill_value <= INT32_MAX, "idx_oob_fill_value must fit int32");
+    STD_TORCH_CHECK(!std::isfinite(value_oob_fill_value) || std::abs(value_oob_fill_value) <= std::numeric_limits<float>::max(),
+                    "finite value_oob_fill_value must fit float32");
     STD_TORCH_CHECK(!(sorted_value && !return_value), "`return_value` must be enabled when `sorted_value` is True");
     STD_TORCH_CHECK(!(sorted_value && sorted_index), "`sorted_value` and `sorted_index` cannot be used at the same time");
-    // Contract: sorted_value is a 32-bit-value-only feature.
     STD_TORCH_CHECK(!(sorted_value && value_t == ScalarType::BFloat16), "`sorted_value` is only supported for float32 input");
     STD_TORCH_CHECK(!begin.has_value(), "`begin` is not supported currently");
     if (return_value) {
@@ -59,11 +86,17 @@ void topk(
     }
 
     DS_CHECK_DEVICE(input);
-    DS_CHECK_DEVICE(begin);
-    DS_CHECK_DEVICE(end);
-    DS_CHECK_DEVICE(output_value);
-    DS_CHECK_DEVICE(output_index);
-    DS_CHECK_DEVICE(output_idx_offset);
+    const auto device_index = input.get_device_index();
+    auto check_device = [&](const char* name, const auto& tensor) {
+        STD_TORCH_CHECK(deep_select::_check_optional_tensor(tensor, [&](const Tensor& t) {
+            return t.is_cuda() && t.get_device_index() == device_index;
+        }), "`", name, "` must be on the same CUDA device as `input`");
+    };
+    check_device("begin", begin);
+    check_device("end", end);
+    check_device("output_value", output_value);
+    check_device("output_index", output_index);
+    check_device("output_idx_offset", output_idx_offset);
 
     DS_CHECK_SHAPE(input, batch_size, vocab_size);
     DS_CHECK_SHAPE(begin, batch_size);
@@ -72,11 +105,9 @@ void topk(
     DS_CHECK_SHAPE(output_index, batch_size, topk);
     DS_CHECK_SHAPE(output_idx_offset, batch_size);
 
-    DS_CHECK_DTYPE(input, value_t);
     DS_CHECK_DTYPE(begin, ScalarType::Int);
     DS_CHECK_DTYPE(end, ScalarType::Int);
     DS_CHECK_DTYPE(output_value, value_t);
-    DS_CHECK_DTYPE(output_index, output_index_t);
     DS_CHECK_DTYPE(output_idx_offset, ScalarType::Int);
 
     DS_CHECK_LAST_DIM_CONTIGUOUS(input);
@@ -87,27 +118,52 @@ void topk(
     DS_CHECK_CONTIGUOUS(output_idx_offset);
 
     auto check_dim0_stride = [&](const char tensor_name[], const Tensor& tensor, uint32_t alignment_requirement_bytes) {
-        int64_t cur_stride = tensor.stride(0);
-        uint64_t itemsize = tensor.element_size();
-        STD_TORCH_CHECK(cur_stride * itemsize % alignment_requirement_bytes == 0,
+        const int64_t cur_stride = tensor.stride(0);
+        const uint64_t itemsize = tensor.element_size();
+        STD_TORCH_CHECK(cur_stride >= 0, tensor_name, ".stride(0) must be nonnegative");
+        STD_TORCH_CHECK(deep_select::checked_mul(cur_stride, itemsize) % alignment_requirement_bytes == 0,
             std::format("{}.stride(0) (currently {} numbers) must be a multiple of {} Bytes ({} numbers)",
                 tensor_name, cur_stride,
                 alignment_requirement_bytes, alignment_requirement_bytes / itemsize
             )
         );
+        deep_select::check_pointer_alignment(tensor_name, tensor);
     };
     check_dim0_stride("input", input, INPUT_STRIDE_ALIGNMENT_REQUIREMENT);
     check_dim0_stride("output_index", output_index, OUTPUT_STRIDE_ALIGNMENT_REQUIREMENT);
     if (output_value.has_value()) {
-        check_dim0_stride("value", *output_value, OUTPUT_STRIDE_ALIGNMENT_REQUIREMENT);
+        check_dim0_stride("output_value", *output_value, OUTPUT_STRIDE_ALIGNMENT_REQUIREMENT);
+    }
+    deep_select::check_storage_bounds("input", input, 128);
+    if (end.has_value()) deep_select::check_storage_bounds("end", *end);
+    if (output_idx_offset.has_value()) deep_select::check_storage_bounds("output_idx_offset", *output_idx_offset);
+    auto check_output = [&](const char* name, const Tensor& output) {
+        STD_TORCH_CHECK(batch_size <= 1 || output.stride(0) >= topk, name, " rows must not overlap");
+        deep_select::check_storage_bounds(name, output);
+        deep_select::check_storage_disjoint(name, output, "input", input);
+        if (end.has_value()) deep_select::check_storage_disjoint(name, output, "end", *end);
+        if (output_idx_offset.has_value()) deep_select::check_storage_disjoint(name, output, "output_idx_offset", *output_idx_offset);
+    };
+    check_output("output_index", output_index);
+    if (output_value.has_value()) {
+        check_output("output_value", *output_value);
+        deep_select::check_storage_disjoint("output_index", output_index, "output_value", *output_value);
     }
 
-    torch::stable::accelerator::DeviceIndex device_index = input.get_device_index();
     torch::stable::accelerator::DeviceGuard device_guard(device_index);
-
     cudaDeviceProp device_prop;
     STD_TORCH_CHECK(cudaGetDeviceProperties(&device_prop, device_index) == cudaSuccess,
                     "failed to get CUDA device properties");
+    const bool use_sm100 = DEEP_SELECT_BUILD_SM100 && device_prop.major == 10;
+    const bool use_native_sm12 = device_prop.major == 12 &&
+        ((DEEP_SELECT_BUILD_SM120 && device_prop.minor == 0) ||
+         (DEEP_SELECT_BUILD_SM121 && device_prop.minor == 1));
+    STD_TORCH_CHECK(use_sm100 || use_native_sm12, "DeepSelect was not built for this CUDA capability: ",
+                    device_prop.major, ".", device_prop.minor);
+    if (batch_size == 0) return;
+    STD_TORCH_CHECK(!use_sm100 || vocab_size > 0,
+                    "SM100/SM103 require vocab_size > 0 for nonempty batches");
+    STD_TORCH_CHECK(batch_size <= device_prop.maxGridSize[0], "batch_size exceeds the CUDA grid limit");
 
     void* stream_ptr = nullptr;
     TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
@@ -139,6 +195,29 @@ void topk(
         static_cast<cudaStream_t>(stream_ptr)
     };
 
+#if DEEP_SELECT_BUILD_SM120 || DEEP_SELECT_BUILD_SM121
+    if (use_native_sm12) {
+        const bool segmented = topk_select_sm120::use_segmented_topk(args, output_index_t == ScalarType::Long);
+        if (segmented) {
+            const uint64_t blocks = deep_select::checked_mul(batch_size, topk_select_sm120::SEGMENTED_PARTITIONS);
+            STD_TORCH_CHECK(blocks <= static_cast<uint64_t>(device_prop.maxGridSize[0]),
+                            "segmented batch_size exceeds the CUDA grid limit");
+            const uint64_t workspace_elements = deep_select::checked_mul(blocks, topk_select_sm120::SEGMENTED_WORKSPACE_STRIDE);
+            STD_TORCH_CHECK(workspace_elements <= UINT32_MAX, "segmented workspace indexing exceeds uint32");
+            deep_select::checked_mul(workspace_elements, sizeof(int32_t));
+            auto workspace = torch::stable::empty(
+                {batch_size, topk_select_sm120::SEGMENTED_PARTITIONS, topk_select_sm120::SEGMENTED_WORKSPACE_STRIDE},
+                ScalarType::Int, std::nullopt, input.device());
+            topk_select_sm120::run_segmented_topk_select_kernel(
+                args, value_t == ScalarType::BFloat16, static_cast<int32_t*>(workspace.data_ptr()));
+        } else {
+            topk_select_sm120::run_topk_select_kernel(args, value_t == ScalarType::BFloat16, output_index_t == ScalarType::Long);
+        }
+        return;
+    }
+#endif
+#if DEEP_SELECT_BUILD_SM100
+    STD_TORCH_CHECK(device_prop.multiProcessorCount > 0, "CUDA device must have at least one multiprocessor");
     uint32_t num_sm = device_prop.multiProcessorCount;
     uint32_t num_waves = (batch_size + num_sm-1) / num_sm;
 
@@ -214,6 +293,7 @@ void topk(
             }
         });
     }
+#endif
 }
 
 std::tuple<int64_t, int64_t> get_alignment_requirement() {
diff --git a/csrc/build_utils.py b/csrc/build_utils.py
new file mode 100644
index 0000000..cbae3ca
--- /dev/null
+++ b/csrc/build_utils.py
@@ -0,0 +1,112 @@
+import hashlib
+import os
+import re
+
+
+DEFAULT_CUDA_ARCH_LIST = "10.0a;10.3a"
+SM120_SOURCE = "csrc/cuda_kernels/sm120/topk_select.cu"
+
+
+def parse_cuda_arch_list(requested, cuda_version):
+    aliases = {
+        "10.0": "100a", "10.0a": "100a", "100a": "100a",
+        "10.3": "103a", "10.3a": "103a", "103a": "103a",
+        "10.0f": "100f", "100f": "100f",
+        "12.0": "120", "120": "120",
+        "12.1": "121", "121": "121",
+    }
+    requested = DEFAULT_CUDA_ARCH_LIST if requested is None else requested
+    entries = re.split(r"[;\s]+", requested.strip())
+    selected = {}
+    for entry in entries:
+        ptx = entry.endswith("+PTX")
+        base = entry[:-4] if ptx else entry
+        if base not in aliases:
+            raise ValueError(f"Unsupported DEEP_SELECT_CUDA_ARCH_LIST entry: {entry!r}")
+        arch = aliases[base]
+        minimum = (12, 9) if arch in {"103a", "100f", "121"} else (12, 8)
+        if cuda_version < minimum:
+            raise RuntimeError(
+                f"{entry} requires CUDA {minimum[0]}.{minimum[1]} or newer"
+            )
+        selected[arch] = selected.get(arch, False) or ptx
+    return tuple(sorted(selected.items()))
+
+
+def architecture_flags(architectures):
+    flags = []
+    for arch, ptx in architectures:
+        flags += ["-gencode", f"arch=compute_{arch},code=sm_{arch}"]
+        if ptx:
+            flags += ["-gencode", f"arch=compute_{arch},code=compute_{arch}"]
+    return flags
+
+
+def source_group(source):
+    path = source.replace(os.sep, "/")
+    if path.endswith("/api.cpp"):
+        return "host"
+    if path.endswith("/cuda_kernels/sm120/topk_select.cu"):
+        return "sm120"
+    if any(f"/cuda_kernels/{version}/instantiations/" in path
+           for version in ("v3", "v3_fp32", "v3_cluster")) and path.endswith(".cu"):
+        return "sm100"
+    raise ValueError(f"No DeepSelect compilation group for {source}")
+
+
+def build_plan(cuda_sources, architectures):
+    sm100 = tuple(item for item in architectures if item[0] in {"100a", "103a", "100f"})
+    sm12 = tuple(item for item in architectures if item[0] in {"120", "121"})
+    macros = [("DEEP_SELECT_BUILD_SM100", str(int(bool(sm100)))),
+              ("DEEP_SELECT_BUILD_SM120", str(int(any(arch == "120" for arch, _ in sm12)))),
+              ("DEEP_SELECT_BUILD_SM121", str(int(any(arch == "121" for arch, _ in sm12))))]
+    sources = [source for source in cuda_sources
+               if source_group(source) == "host" or sm100]
+    if sm12:
+        sources.append(SM120_SOURCE)
+    flags = {
+        "host": [],
+        "sm100": architecture_flags(sm100),
+        "sm120": architecture_flags(sm12),
+    }
+    return sources, macros, flags
+
+
+def grouped_build_extension(base_class, group_flags):
+    class GroupedBuildExtension(base_class):
+        def finalize_options(self):
+            super().finalize_options()
+            self.force = True
+
+        def build_extension(self, ext):
+            compile_objects = self.compiler.compile
+
+            def compile_groups(sources, output_dir=None, macros=None,
+                               include_dirs=None, debug=0, extra_preargs=None,
+                               extra_postargs=None, depends=None):
+                objects = []
+                for group in ("host", "sm100", "sm120"):
+                    group_sources = [source for source in sources
+                                     if source_group(source) == group]
+                    if not group_sources:
+                        continue
+                    flags = {key: list(value) for key, value in extra_postargs.items()}
+                    flags["nvcc"] += group_flags[group]
+                    signature = hashlib.sha256(
+                        repr((macros, flags, include_dirs, extra_preargs, debug)).encode()
+                    ).hexdigest()[:16]
+                    objects += compile_objects(
+                        group_sources,
+                        output_dir=os.path.join(output_dir, f"{group}-{signature}"),
+                        macros=macros, include_dirs=include_dirs, debug=debug,
+                        extra_preargs=extra_preargs, extra_postargs=flags, depends=depends,
+                    )
+                return objects
+
+            self.compiler.compile = compile_groups
+            try:
+                super().build_extension(ext)
+            finally:
+                self.compiler.compile = compile_objects
+
+    return GroupedBuildExtension
diff --git a/csrc/cuda_kernels/common_parts.cuh b/csrc/cuda_kernels/common_parts.cuh
index e3f3cac..7c81db2 100644
--- a/csrc/cuda_kernels/common_parts.cuh
+++ b/csrc/cuda_kernels/common_parts.cuh
@@ -93,6 +93,18 @@ void st_global(ValueT* ptr, ValueT src[NUM_VALUES]) {
     }
 }
 
+template<uint32_t NUM_VALUES, typename ValueT>
+__device__ __forceinline__
+void st_global_bounded(ValueT* ptr, ValueT src[NUM_VALUES], uint32_t remaining) {
+    if (remaining >= NUM_VALUES) {
+        st_global<NUM_VALUES>(ptr, src);
+    } else {
+        CUTE_UNROLL
+        for (uint32_t j = 0; j < NUM_VALUES; ++j)
+            if (j < remaining) ptr[j] = src[j];
+    }
+}
+
 template<
     typename Config,
     uint32_t MAX_TOPK,
@@ -165,7 +177,7 @@ struct EpilogueRunner {
                     CUTE_UNROLL
                     for (uint32_t j = 0; j < NUM_OUTPUT_IDXS_PER_STORE; ++j)
                         out[j] = i+j < end_vocab_idx ? (OutIdxT)(i+j) + output_idx_offset : args.idx_oob_fill_value;
-                    STORE_TO_GMEM(result_indices + i, out);
+                    st_global_bounded<NUM_OUTPUT_IDXS_PER_STORE>(result_indices + i, out, args.topk - i);
                 }
 
                 // Save values into `result_values`, if `RETURN_VALUE` is `True`
@@ -179,7 +191,7 @@ struct EpilogueRunner {
                         CUTE_UNROLL
                         for (uint32_t j = 0; j < NUM_VALUES_PER_LOAD_STORE; ++j)
                             values[j] = i+j < end_vocab_idx ? values[j] : oob_fill_value;
-                        STORE_TO_GMEM(result_values + i, values);
+                        st_global_bounded<NUM_VALUES_PER_LOAD_STORE>(result_values + i, values, args.topk - i);
                     }
                 }
             } else {
@@ -193,7 +205,7 @@ struct EpilogueRunner {
                     CUTE_UNROLL
                     for (uint32_t j = 0; j < NUM_OUTPUT_IDX_PER_ROUND; ++j)
                         out[j] = (OutIdxT)indices_u32[j] + output_idx_offset;
-                    st_global<NUM_OUTPUT_IDX_PER_ROUND>(result_indices + i, out);
+                    st_global_bounded<NUM_OUTPUT_IDX_PER_ROUND>(result_indices + i, out, args.topk - i);
                 }
 
                 if constexpr (Config::return_value) {
@@ -203,7 +215,7 @@ struct EpilogueRunner {
                     for (uint32_t i = threadIdx.x * NUM_VALUES_PER_ROUND; i < args.topk; i += NUM_THREADS * NUM_VALUES_PER_ROUND) {
                         UIntValueT values[NUM_VALUES_PER_ROUND];
                         ld_shared<NUM_VALUES_PER_ROUND>(values, (UIntValueT*)(smem_value_buf + i));
-                        st_global<NUM_VALUES_PER_ROUND>((UIntValueT*)(result_values + i), values);
+                        st_global_bounded<NUM_VALUES_PER_ROUND>((UIntValueT*)(result_values + i), values, args.topk - i);
                     }
                 }
             }
@@ -258,15 +270,15 @@ struct EpilogueRunner {
                 auto store = [&]<uint32_t NUM_VALUES, typename T>(T* dst, T src[NUM_VALUES]) {
                     constexpr uint32_t NUM_BYTES_TO_STORE = NUM_VALUES * sizeof(T);
                     if constexpr (NUM_BYTES_TO_STORE <= NUM_BYTES_PER_GMEM_STORE) {
-                        // Don't need to do OOB check because the output array is at least padded to the maximum width of global store.
-                        st_global<NUM_VALUES>(dst, src);
+                        st_global_bounded<NUM_VALUES>(dst, src, args.topk - thread_offset);
                     } else {
                         static_assert(NUM_BYTES_TO_STORE % NUM_BYTES_PER_GMEM_STORE == 0);
                         CUTE_UNROLL
                         for (uint32_t i = 0; i < NUM_BYTES_TO_STORE; i += NUM_BYTES_PER_GMEM_STORE) {
                             uint32_t elem_offset = i / sizeof(T);
                             if (thread_offset + elem_offset < args.topk)
-                                STORE_TO_GMEM(dst+elem_offset, src+elem_offset);
+                                st_global_bounded<NUM_BYTES_PER_GMEM_STORE / sizeof(T)>(
+                                    dst+elem_offset, src+elem_offset, args.topk - thread_offset - elem_offset);
                         }
                     }
                 };
diff --git a/csrc/cuda_kernels/sm120/topk_select.cu b/csrc/cuda_kernels/sm120/topk_select.cu
new file mode 100644
index 0000000..d669535
--- /dev/null
+++ b/csrc/cuda_kernels/sm120/topk_select.cu
@@ -0,0 +1,747 @@
+#include "topk_select.h"
+
+#include <cuda_bf16.h>
+#include <cub/block/block_radix_sort.cuh>
+#include <cub/block/block_scan.cuh>
+#include <stdexcept>
+#include <string>
+
+namespace topk_select_sm120 {
+namespace {
+
+void check_kernel_launch() {
+    const cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess)
+        throw std::runtime_error(std::string("SM120 topk kernel launch failed: ") + cudaGetErrorString(error));
+}
+
+constexpr int NUM_THREADS = 512;
+constexpr int ITEMS_PER_THREAD = 8;
+constexpr int MAX_TOPK = NUM_THREADS * ITEMS_PER_THREAD;
+using BlockSort = cub::BlockRadixSort<uint64_t, NUM_THREADS, ITEMS_PER_THREAD, cub::NullType, 4, false>;
+using BlockScan = cub::BlockScan<uint32_t, NUM_THREADS>;
+
+struct SharedMemory {
+    union {
+        uint32_t indices[MAX_TOPK];
+        BlockSort::TempStorage sort;
+        BlockScan::TempStorage scan;
+    } work;
+    uint32_t histogram[256];
+    uint32_t prefix;
+    uint32_t rank;
+    uint32_t count;
+    uint32_t have_nan;
+};
+static_assert(sizeof(SharedMemory) <= 48 * 1024,
+              "SM120 selector must fit the default per-block shared-memory budget");
+
+__device__ __forceinline__ uint32_t ordered_key(float value) {
+    uint32_t bits = __float_as_uint(value);
+    if ((bits & 0x7fffffffU) == 0) bits = 0;
+    if ((bits & 0x7fffffffU) > 0x7f800000U) return 0xffffffffU;
+    return bits ^ ((bits & 0x80000000U) ? 0xffffffffU : 0x80000000U);
+}
+
+template<typename ValueT, typename IndexT>
+__global__ __launch_bounds__(NUM_THREADS) void topk_kernel(TopkSelectArgs args) {
+    __shared__ SharedMemory smem;
+    const uint32_t row = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const ValueT* input = static_cast<const ValueT*>(args.input) + row * args.stride_input_batch;
+    IndexT* output_index = static_cast<IndexT*>(args.output_index) + row * args.stride_output_index_batch;
+    ValueT* output_value = args.return_value
+        ? static_cast<ValueT*>(args.output_value) + row * args.stride_output_value_batch : nullptr;
+    const int32_t offset = args.output_idx_offset ? args.output_idx_offset[row] : 0;
+    const int32_t end = args.end_ptr ? args.end_ptr[row] : static_cast<int32_t>(args.vocab_size);
+    if (end < 0 || static_cast<uint32_t>(end) > args.vocab_size) {
+        if (tid == 0) asm volatile("trap;");
+        return;
+    }
+    const uint32_t length = static_cast<uint32_t>(end);
+    const uint32_t valid_count = min(length, args.topk);
+    if (tid == 0) {
+        smem.prefix = 0;
+        smem.rank = args.topk;
+        smem.count = 0;
+        smem.have_nan = 0;
+    }
+    __syncthreads();
+
+    if (length <= args.topk) {
+        for (uint32_t i = tid; i < length; i += NUM_THREADS) smem.work.indices[i] = i;
+        __syncthreads();
+    } else {
+        uint32_t mask = 0;
+        for (int shift = 24; shift >= 0; shift -= 8) {
+            if (tid < 256) smem.histogram[tid] = 0;
+            __syncthreads();
+            const uint32_t prefix = smem.prefix;
+            for (uint32_t i = tid; i < length; i += NUM_THREADS) {
+                const float value = static_cast<float>(input[i]);
+                const uint32_t bits = __float_as_uint(value);
+                if (shift == 24 && (bits & 0x7fffffffU) > 0x7f800000U)
+                    atomicExch(&smem.have_nan, 1U);
+                const uint32_t key = ordered_key(value);
+                if ((key & mask) == prefix) atomicAdd(&smem.histogram[(key >> shift) & 255U], 1U);
+            }
+            __syncthreads();
+            if (smem.have_nan) {
+                if (tid == 0) {
+                    if (args.abort_when_nan_found) asm volatile("trap;");
+                    else output_index[0] = static_cast<IndexT>(0x3f3f3f3f);
+                }
+                return;
+            }
+            if (tid == 0) {
+                uint32_t rank = smem.rank;
+                for (int bucket = 255; bucket >= 0; --bucket) {
+                    const uint32_t count = smem.histogram[bucket];
+                    if (rank <= count) {
+                        smem.prefix |= static_cast<uint32_t>(bucket) << shift;
+                        smem.rank = rank;
+                        break;
+                    }
+                    rank -= count;
+                }
+            }
+            mask |= 255U << shift;
+            __syncthreads();
+        }
+
+        const uint32_t threshold = smem.prefix;
+        const uint32_t quota = smem.rank;
+        const uint32_t chunk = (length + NUM_THREADS - 1) / NUM_THREADS;
+        const uint32_t lo = min(tid * chunk, length);
+        const uint32_t hi = min(lo + chunk, length);
+        uint32_t equal_count = 0;
+        for (uint32_t i = lo; i < hi; ++i)
+            equal_count += ordered_key(static_cast<float>(input[i])) == threshold;
+        uint32_t equal_before;
+        BlockScan(smem.work.scan).ExclusiveSum(equal_count, equal_before);
+        __syncthreads();
+        for (uint32_t i = lo; i < hi; ++i) {
+            const uint32_t key = ordered_key(static_cast<float>(input[i]));
+            const bool take_equal = key == threshold && equal_before++ < quota;
+            if (key > threshold || take_equal) {
+                const uint32_t slot = atomicAdd(&smem.count, 1U);
+                smem.work.indices[slot] = i;
+            }
+        }
+        __syncthreads();
+    }
+
+    if (args.sorted_value || args.sorted_index) {
+        uint64_t keys[ITEMS_PER_THREAD];
+        #pragma unroll
+        for (int j = 0; j < ITEMS_PER_THREAD; ++j) {
+            const uint32_t slot = tid * ITEMS_PER_THREAD + j;
+            uint64_t key = 0;
+            if (slot < valid_count) {
+                const uint32_t index = smem.work.indices[slot];
+                key = static_cast<uint64_t>(~index);
+                if (args.sorted_value)
+                    key |= static_cast<uint64_t>(ordered_key(static_cast<float>(input[index]))) << 32;
+            }
+            keys[j] = key;
+        }
+        __syncthreads();
+        BlockSort(smem.work.sort).SortDescending(keys);
+        #pragma unroll
+        for (int j = 0; j < ITEMS_PER_THREAD; ++j) {
+            const uint32_t slot = tid * ITEMS_PER_THREAD + j;
+            if (slot < args.topk) {
+                const uint32_t index = ~static_cast<uint32_t>(keys[j]);
+                const bool valid = slot < valid_count;
+                output_index[slot] = valid ? static_cast<IndexT>(static_cast<int64_t>(index) + offset)
+                                           : static_cast<IndexT>(args.idx_oob_fill_value);
+                if (args.return_value)
+                    output_value[slot] = valid ? input[index] : static_cast<ValueT>(args.value_oob_fill_value);
+            }
+        }
+    } else {
+        for (uint32_t slot = tid; slot < args.topk; slot += NUM_THREADS) {
+            const bool valid = slot < valid_count;
+            const uint32_t index = valid ? smem.work.indices[slot] : 0;
+            output_index[slot] = valid ? static_cast<IndexT>(static_cast<int64_t>(index) + offset)
+                                       : static_cast<IndexT>(args.idx_oob_fill_value);
+            if (args.return_value)
+                output_value[slot] = valid ? input[index] : static_cast<ValueT>(args.value_oob_fill_value);
+        }
+    }
+}
+
+constexpr int FAST_THREADS = 256;
+using FastScan = cub::BlockScan<uint32_t, FAST_THREADS, cub::BLOCK_SCAN_WARP_SCANS>;
+
+constexpr int STREAM_TOPK = 512;
+constexpr int STREAM_TILE_ITEMS = 4;
+constexpr int STREAM_CANDIDATE_ITEMS = 10;
+constexpr int STREAM_CAPACITY = FAST_THREADS * STREAM_CANDIDATE_ITEMS;
+constexpr int STREAM_COMPACT_COUNT = 1536;
+static_assert(STREAM_COMPACT_COUNT - 1 + FAST_THREADS * STREAM_TILE_ITEMS < STREAM_CAPACITY);
+
+struct StreamingSharedMemory {
+    uint32_t keys[STREAM_CAPACITY];
+    uint32_t indices[STREAM_CAPACITY];
+    uint32_t histogram[FAST_THREADS / 32][257];
+    FastScan::TempStorage scan;
+    uint32_t worst_keys[FAST_THREADS / 32];
+    uint32_t worst_indices[FAST_THREADS / 32];
+    uint32_t bucket;
+    uint32_t rank;
+    uint32_t complete;
+    uint32_t threshold_key;
+    uint32_t threshold_index;
+};
+static_assert(sizeof(StreamingSharedMemory) <= 32 * 1024);
+
+__device__ __forceinline__ bool streaming_better(
+    uint32_t key, uint32_t index, uint32_t other_key, uint32_t other_index) {
+    return key > other_key || (key == other_key && index < other_index);
+}
+
+template<bool IndexKey, int FirstShift, int LastShift>
+__device__ __forceinline__ void streaming_radix(
+    StreamingSharedMemory& smem,
+    const uint32_t (&keys)[STREAM_CANDIDATE_ITEMS],
+    const uint32_t (&indices)[STREAM_CANDIDATE_ITEMS],
+    uint32_t& active, uint32_t& selected, uint32_t& rank) {
+    const uint32_t tid = threadIdx.x;
+    const uint32_t warp = tid / 32;
+    for (int shift = FirstShift; shift >= LastShift; shift -= 8) {
+        #pragma unroll
+        for (int w = 0; w < FAST_THREADS / 32; ++w) smem.histogram[w][tid] = 0;
+        __syncthreads();
+        #pragma unroll
+        for (int j = 0; j < STREAM_CANDIDATE_ITEMS; ++j) {
+            const uint32_t key = IndexKey ? ~indices[j] : keys[j];
+            if (active & (1U << j))
+                atomicAdd(&smem.histogram[warp][(key >> shift) & 255U], 1U);
+        }
+        __syncthreads();
+        const uint32_t bucket = 255U - tid;
+        uint32_t bucket_count = 0;
+        #pragma unroll
+        for (int w = 0; w < FAST_THREADS / 32; ++w)
+            bucket_count += smem.histogram[w][bucket];
+        uint32_t greater;
+        FastScan(smem.scan).ExclusiveSum(bucket_count, greater);
+        __syncthreads();
+        if (greater < rank && rank <= greater + bucket_count) {
+            smem.bucket = bucket;
+            smem.rank = rank - greater;
+            smem.complete = rank - greater == bucket_count;
+        }
+        __syncthreads();
+        const uint32_t pivot = smem.bucket;
+        const bool complete = smem.complete != 0;
+        rank = smem.rank;
+        uint32_t next_active = 0;
+        #pragma unroll
+        for (int j = 0; j < STREAM_CANDIDATE_ITEMS; ++j) {
+            const uint32_t key = IndexKey ? ~indices[j] : keys[j];
+            const uint32_t digit = (key >> shift) & 255U;
+            if (active & (1U << j)) {
+                if (digit > pivot || (complete && digit == pivot)) selected |= 1U << j;
+                else if (digit == pivot) next_active |= 1U << j;
+            }
+        }
+        active = next_active;
+        if (complete) break;
+    }
+}
+
+template<typename ValueT>
+__device__ __forceinline__ void streaming_compact(StreamingSharedMemory& smem, uint32_t count) {
+    const uint32_t tid = threadIdx.x;
+    const uint32_t lane = tid & 31U;
+    const uint32_t warp = tid / 32;
+    uint32_t keys[STREAM_CANDIDATE_ITEMS];
+    uint32_t indices[STREAM_CANDIDATE_ITEMS];
+    uint32_t active = 0;
+    #pragma unroll
+    for (int j = 0; j < STREAM_CANDIDATE_ITEMS; ++j) {
+        const uint32_t slot = tid + j * FAST_THREADS;
+        keys[j] = slot < count ? smem.keys[slot] : 0;
+        indices[j] = slot < count ? smem.indices[slot] : 0;
+        if (slot < count) active |= 1U << j;
+    }
+    __syncthreads();
+    uint32_t rank = STREAM_TOPK;
+    uint32_t selected = 0;
+    streaming_radix<false, 24, sizeof(ValueT) == 2 ? 16 : 0>(
+        smem, keys, indices, active, selected, rank);
+    if (!smem.complete)
+        streaming_radix<true, 16, 0>(smem, keys, indices, active, selected, rank);
+
+    uint32_t slot;
+    FastScan(smem.scan).ExclusiveSum(static_cast<uint32_t>(__popc(selected)), slot);
+    __syncthreads();
+    uint32_t worst_key = 0xffffffffU;
+    uint32_t worst_index = 0;
+    #pragma unroll
+    for (int j = 0; j < STREAM_CANDIDATE_ITEMS; ++j) {
+        if (selected & (1U << j)) {
+            const uint32_t key = keys[j];
+            const uint32_t index = indices[j];
+            smem.keys[slot] = key;
+            smem.indices[slot++] = index;
+            if (streaming_better(worst_key, worst_index, key, index)) {
+                worst_key = key;
+                worst_index = index;
+            }
+        }
+    }
+    #pragma unroll
+    for (int delta = 16; delta > 0; delta /= 2) {
+        const uint32_t key = __shfl_down_sync(0xffffffffU, worst_key, delta);
+        const uint32_t index = __shfl_down_sync(0xffffffffU, worst_index, delta);
+        if (streaming_better(worst_key, worst_index, key, index)) {
+            worst_key = key;
+            worst_index = index;
+        }
+    }
+    if (lane == 0) {
+        smem.worst_keys[warp] = worst_key;
+        smem.worst_indices[warp] = worst_index;
+    }
+    __syncthreads();
+    if (tid == 0) {
+        #pragma unroll
+        for (int w = 1; w < FAST_THREADS / 32; ++w) {
+            const uint32_t key = smem.worst_keys[w];
+            const uint32_t index = smem.worst_indices[w];
+            if (streaming_better(worst_key, worst_index, key, index)) {
+                worst_key = key;
+                worst_index = index;
+            }
+        }
+        smem.threshold_key = worst_key;
+        smem.threshold_index = worst_index;
+    }
+    __syncthreads();
+}
+
+template<typename ValueT>
+struct StreamingInputSource {
+    const ValueT* input;
+    uint32_t begin;
+    uint32_t length;
+
+    __device__ __forceinline__ bool load(uint32_t slot, uint32_t& key, uint32_t& index) const {
+        if (slot >= length) return false;
+        index = begin + slot;
+        key = ordered_key(static_cast<float>(input[index]));
+        return true;
+    }
+};
+
+struct StreamingWorkspaceSource {
+    const int32_t* workspace;
+
+    __device__ __forceinline__ bool load(uint32_t slot, uint32_t& key, uint32_t& index) const {
+        if (slot >= SEGMENTED_PARTITIONS * STREAM_TOPK) return false;
+        const uint32_t partition = slot / STREAM_TOPK;
+        const uint32_t candidate = slot % STREAM_TOPK;
+        const int32_t* segment = workspace + partition * SEGMENTED_WORKSPACE_STRIDE;
+        if (candidate >= static_cast<uint32_t>(segment[2 * STREAM_TOPK])) return false;
+        key = reinterpret_cast<const uint32_t*>(segment)[candidate];
+        index = static_cast<uint32_t>(segment[STREAM_TOPK + candidate]);
+        return true;
+    }
+};
+
+template<typename ValueT, typename Source>
+__device__ __forceinline__ bool streaming_select(
+    StreamingSharedMemory& smem, const Source& source, uint32_t length, uint32_t& count) {
+    const uint32_t tid = threadIdx.x;
+    count = 0;
+    if (length == 0) return true;
+    bool have_threshold = false;
+    constexpr uint32_t tile_size = FAST_THREADS * STREAM_TILE_ITEMS;
+    const uint32_t tiles = (length + tile_size - 1) / tile_size;
+    const uint32_t tile_bits = tiles > 1 ? 32 - __clz(tiles - 1U) : 0;
+    const uint32_t tile_domain = 1U << tile_bits;
+    for (uint32_t step = 0; step < tile_domain; ++step) {
+        const uint32_t tile = tile_bits ? __brev(step) >> (32 - tile_bits) : 0;
+        if (tile >= tiles) continue;
+        const uint32_t base = tile * tile_size;
+        uint32_t keys[STREAM_TILE_ITEMS];
+        uint32_t indices[STREAM_TILE_ITEMS];
+        uint32_t keep = 0;
+        bool have_nan = false;
+        #pragma unroll
+        for (int j = 0; j < STREAM_TILE_ITEMS; ++j) {
+            keys[j] = 0;
+            indices[j] = 0;
+            if (source.load(base + tid + j * FAST_THREADS, keys[j], indices[j])) {
+                have_nan |= keys[j] == 0xffffffffU;
+                if (!have_threshold || streaming_better(
+                        keys[j], indices[j], smem.threshold_key, smem.threshold_index))
+                    keep |= 1U << j;
+            }
+        }
+        if (__syncthreads_or(have_nan)) return false;
+        uint32_t before;
+        uint32_t tile_count;
+        FastScan(smem.scan).ExclusiveSum(static_cast<uint32_t>(__popc(keep)), before, tile_count);
+        __syncthreads();
+        uint32_t slot = count + before;
+        #pragma unroll
+        for (int j = 0; j < STREAM_TILE_ITEMS; ++j) {
+            if (keep & (1U << j)) {
+                smem.keys[slot] = keys[j];
+                smem.indices[slot++] = indices[j];
+            }
+        }
+        count += tile_count;
+        __syncthreads();
+        if (count >= STREAM_COMPACT_COUNT) {
+            streaming_compact<ValueT>(smem, count);
+            count = STREAM_TOPK;
+            have_threshold = true;
+        }
+    }
+    if (count > STREAM_TOPK) {
+        streaming_compact<ValueT>(smem, count);
+        count = STREAM_TOPK;
+    }
+    return true;
+}
+
+template<typename ValueT, bool ReturnValue>
+__global__ __launch_bounds__(FAST_THREADS) void topk_streaming_kernel(TopkSelectArgs args) {
+    __shared__ StreamingSharedMemory smem;
+    const uint32_t row = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const ValueT* input = static_cast<const ValueT*>(args.input) + row * args.stride_input_batch;
+    int32_t* output_index = static_cast<int32_t*>(args.output_index) + row * args.stride_output_index_batch;
+    ValueT* output_value = nullptr;
+    if constexpr (ReturnValue)
+        output_value = static_cast<ValueT*>(args.output_value) + row * args.stride_output_value_batch;
+    const int32_t offset = args.output_idx_offset ? args.output_idx_offset[row] : 0;
+    const int32_t end = args.end_ptr ? args.end_ptr[row] : static_cast<int32_t>(args.vocab_size);
+    if (end < 0 || static_cast<uint32_t>(end) > args.vocab_size) {
+        if (tid == 0) asm volatile("trap;");
+        return;
+    }
+    const uint32_t length = static_cast<uint32_t>(end);
+    if (length <= STREAM_TOPK) {
+        for (uint32_t slot = tid; slot < STREAM_TOPK; slot += FAST_THREADS) {
+            output_index[slot] = slot < length
+                ? static_cast<int32_t>(static_cast<int64_t>(slot) + offset) : args.idx_oob_fill_value;
+            if constexpr (ReturnValue)
+                output_value[slot] = slot < length ? input[slot] : static_cast<ValueT>(args.value_oob_fill_value);
+        }
+        return;
+    }
+
+    uint32_t count;
+    if (!streaming_select<ValueT>(smem, StreamingInputSource<ValueT>{input, 0, length}, length, count)) {
+        if (tid == 0) {
+            if (args.abort_when_nan_found) asm volatile("trap;");
+            else output_index[0] = 0x3f3f3f3f;
+        }
+        return;
+    }
+    for (uint32_t slot = tid; slot < STREAM_TOPK; slot += FAST_THREADS) {
+        const uint32_t index = smem.indices[slot];
+        output_index[slot] = static_cast<int32_t>(static_cast<int64_t>(index) + offset);
+        if constexpr (ReturnValue) output_value[slot] = input[index];
+    }
+}
+
+template<typename ValueT>
+__global__ __launch_bounds__(FAST_THREADS) void topk_segmented_local_kernel(
+    TopkSelectArgs args, int32_t* workspace) {
+    __shared__ StreamingSharedMemory smem;
+    const uint32_t row = blockIdx.x / SEGMENTED_PARTITIONS;
+    const uint32_t partition = blockIdx.x % SEGMENTED_PARTITIONS;
+    const uint32_t tid = threadIdx.x;
+    int32_t* segment = workspace + blockIdx.x * SEGMENTED_WORKSPACE_STRIDE;
+    const int32_t end = args.end_ptr ? args.end_ptr[row] : static_cast<int32_t>(args.vocab_size);
+    if (end < 0 || static_cast<uint32_t>(end) > args.vocab_size) {
+        if (tid == 0) asm volatile("trap;");
+        return;
+    }
+    const uint32_t length = static_cast<uint32_t>(end);
+    if (length <= STREAM_TOPK) {
+        if (tid == 0) segment[2 * STREAM_TOPK] = 0;
+        return;
+    }
+    const uint32_t begin = partition * length / SEGMENTED_PARTITIONS;
+    const uint32_t limit = (partition + 1) * length / SEGMENTED_PARTITIONS;
+    const ValueT* input = static_cast<const ValueT*>(args.input) + row * args.stride_input_batch;
+    uint32_t count;
+    if (!streaming_select<ValueT>(
+            smem, StreamingInputSource<ValueT>{input, begin, limit - begin}, limit - begin, count)) {
+        if (tid == 0) segment[2 * STREAM_TOPK] = -1;
+        return;
+    }
+    for (uint32_t slot = tid; slot < count; slot += FAST_THREADS) {
+        reinterpret_cast<uint32_t*>(segment)[slot] = smem.keys[slot];
+        segment[STREAM_TOPK + slot] = static_cast<int32_t>(smem.indices[slot]);
+    }
+    __syncthreads();
+    if (tid == 0) segment[2 * STREAM_TOPK] = static_cast<int32_t>(count);
+}
+
+template<typename ValueT, bool ReturnValue>
+__global__ __launch_bounds__(FAST_THREADS) void topk_segmented_merge_kernel(
+    TopkSelectArgs args, const int32_t* workspace) {
+    __shared__ StreamingSharedMemory smem;
+    const uint32_t row = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const int32_t end = args.end_ptr ? args.end_ptr[row] : static_cast<int32_t>(args.vocab_size);
+    if (end < 0 || static_cast<uint32_t>(end) > args.vocab_size) {
+        if (tid == 0) asm volatile("trap;");
+        return;
+    }
+    const uint32_t length = static_cast<uint32_t>(end);
+    const int32_t* row_workspace = workspace + row * SEGMENTED_PARTITIONS * SEGMENTED_WORKSPACE_STRIDE;
+    int32_t* output_index = static_cast<int32_t*>(args.output_index) + row * args.stride_output_index_batch;
+    const bool have_nan = tid < SEGMENTED_PARTITIONS &&
+        row_workspace[tid * SEGMENTED_WORKSPACE_STRIDE + 2 * STREAM_TOPK] == -1;
+    if (__syncthreads_or(have_nan)) {
+        if (tid == 0) {
+            if (args.abort_when_nan_found) asm volatile("trap;");
+            else output_index[0] = 0x3f3f3f3f;
+        }
+        return;
+    }
+    const ValueT* input = static_cast<const ValueT*>(args.input) + row * args.stride_input_batch;
+    ValueT* output_value = nullptr;
+    if constexpr (ReturnValue)
+        output_value = static_cast<ValueT*>(args.output_value) + row * args.stride_output_value_batch;
+    const int32_t offset = args.output_idx_offset ? args.output_idx_offset[row] : 0;
+    if (length <= STREAM_TOPK) {
+        for (uint32_t slot = tid; slot < STREAM_TOPK; slot += FAST_THREADS) {
+            output_index[slot] = slot < length
+                ? static_cast<int32_t>(static_cast<int64_t>(slot) + offset) : args.idx_oob_fill_value;
+            if constexpr (ReturnValue)
+                output_value[slot] = slot < length ? input[slot] : static_cast<ValueT>(args.value_oob_fill_value);
+        }
+        return;
+    }
+    uint32_t count;
+    if (!streaming_select<ValueT>(smem, StreamingWorkspaceSource{row_workspace},
+                                 SEGMENTED_PARTITIONS * STREAM_TOPK, count)) {
+        if (tid == 0) {
+            if (args.abort_when_nan_found) asm volatile("trap;");
+            else output_index[0] = 0x3f3f3f3f;
+        }
+        return;
+    }
+    for (uint32_t slot = tid; slot < STREAM_TOPK; slot += FAST_THREADS) {
+        const uint32_t index = smem.indices[slot];
+        output_index[slot] = static_cast<int32_t>(static_cast<int64_t>(index) + offset);
+        if constexpr (ReturnValue) output_value[slot] = input[index];
+    }
+}
+
+template<typename ValueT>
+void launch_segmented(const TopkSelectArgs& args, int32_t* workspace) {
+    topk_segmented_local_kernel<ValueT>
+        <<<args.batch_size * SEGMENTED_PARTITIONS, FAST_THREADS, 0, args.stream>>>(args, workspace);
+    check_kernel_launch();
+    if (args.return_value)
+        topk_segmented_merge_kernel<ValueT, true>
+            <<<args.batch_size, FAST_THREADS, 0, args.stream>>>(args, workspace);
+    else
+        topk_segmented_merge_kernel<ValueT, false>
+            <<<args.batch_size, FAST_THREADS, 0, args.stream>>>(args, workspace);
+    check_kernel_launch();
+}
+
+template<int Capacity>
+struct FastSharedMemory {
+    uint32_t indices[Capacity];
+    uint32_t histogram[256];
+    FastScan::TempStorage scan;
+    uint32_t prefix;
+    uint32_t rank;
+    uint32_t complete;
+    uint32_t have_nan;
+};
+static_assert(sizeof(FastSharedMemory<512>) <= 4 * 1024);
+static_assert(sizeof(FastSharedMemory<2048>) <= 10 * 1024);
+
+template<typename ValueT, int Capacity, bool ReturnValue>
+__global__ __launch_bounds__(FAST_THREADS) void topk_unsorted_kernel(TopkSelectArgs args) {
+    __shared__ FastSharedMemory<Capacity> smem;
+    constexpr bool is_bfloat16 = sizeof(ValueT) == 2;
+    constexpr uint32_t full_warp = 0xffffffffU;
+    const uint32_t row = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const uint32_t lane = tid & 31U;
+    const uint32_t warp = tid / 32;
+    const uint32_t lower_lanes = (1U << lane) - 1U;
+    const ValueT* input = static_cast<const ValueT*>(args.input) + row * args.stride_input_batch;
+    int32_t* output_index = static_cast<int32_t*>(args.output_index) + row * args.stride_output_index_batch;
+    ValueT* output_value = nullptr;
+    if constexpr (ReturnValue)
+        output_value = static_cast<ValueT*>(args.output_value) + row * args.stride_output_value_batch;
+    const int32_t offset = args.output_idx_offset ? args.output_idx_offset[row] : 0;
+    const int32_t end = args.end_ptr ? args.end_ptr[row] : static_cast<int32_t>(args.vocab_size);
+    if (end < 0 || static_cast<uint32_t>(end) > args.vocab_size) {
+        if (tid == 0) asm volatile("trap;");
+        return;
+    }
+    const uint32_t length = static_cast<uint32_t>(end);
+    if (length <= args.topk) {
+        for (uint32_t slot = tid; slot < args.topk; slot += FAST_THREADS) {
+            output_index[slot] = slot < length
+                ? static_cast<int32_t>(static_cast<int64_t>(slot) + offset) : args.idx_oob_fill_value;
+            if constexpr (ReturnValue)
+                output_value[slot] = slot < length ? input[slot] : static_cast<ValueT>(args.value_oob_fill_value);
+        }
+        return;
+    }
+    if (tid == 0) {
+        smem.prefix = 0;
+        smem.rank = args.topk;
+        smem.complete = 0;
+        smem.have_nan = 0;
+    }
+    __syncthreads();
+
+    uint32_t mask = 0;
+    for (int shift = 24; shift >= (is_bfloat16 ? 16 : 0); shift -= 8) {
+        smem.histogram[tid] = 0;
+        __syncthreads();
+        const uint32_t prefix = smem.prefix;
+        const uint32_t rank = smem.rank;
+        for (uint32_t base = 0; base < length; base += FAST_THREADS) {
+            const uint32_t i = base + tid;
+            uint32_t key = 0;
+            if (i < length) key = ordered_key(static_cast<float>(input[i]));
+            if (shift == 24) {
+                const uint32_t nan_lanes = __ballot_sync(full_warp, i < length && key == 0xffffffffU);
+                if (lane == 0 && nan_lanes) atomicExch(&smem.have_nan, 1U);
+            }
+            if (i < length && (key & mask) == prefix)
+                atomicAdd(&smem.histogram[(key >> shift) & 255U], 1U);
+        }
+        __syncthreads();
+        if (smem.have_nan) {
+            if (tid == 0) {
+                if (args.abort_when_nan_found) asm volatile("trap;");
+                else output_index[0] = 0x3f3f3f3f;
+            }
+            return;
+        }
+        const uint32_t bucket = 255U - tid;
+        const uint32_t count = smem.histogram[bucket];
+        uint32_t greater;
+        FastScan(smem.scan).ExclusiveSum(count, greater);
+        __syncthreads();
+        if (greater < rank && rank <= greater + count) {
+            smem.prefix = prefix | (bucket << shift);
+            smem.rank = rank - greater;
+            smem.complete = rank - greater == count;
+        }
+        __syncthreads();
+        mask |= 255U << shift;
+        if (smem.complete) break;
+    }
+
+    const bool complete = smem.complete != 0;
+    uint32_t threshold = smem.prefix;
+    if constexpr (is_bfloat16) {
+        if (!complete && !(threshold & 0x80000000U)) threshold |= 0xffffU;
+    }
+    const uint32_t quota = smem.rank;
+    const uint32_t chunk = ((length + FAST_THREADS - 1) / FAST_THREADS) * 32;
+    const uint32_t lo = min(warp * chunk, length);
+    const uint32_t hi = min(lo + chunk, length);
+    uint32_t greater_count = 0;
+    uint32_t equal_count = 0;
+    for (uint32_t i = lo + lane; i < hi; i += 32) {
+        const uint32_t key = ordered_key(static_cast<float>(input[i]));
+        greater_count += complete ? key >= threshold : key > threshold;
+        if (!complete) equal_count += key == threshold;
+    }
+    #pragma unroll
+    for (int delta = 16; delta > 0; delta /= 2) {
+        greater_count += __shfl_down_sync(full_warp, greater_count, delta);
+        if (!complete) equal_count += __shfl_down_sync(full_warp, equal_count, delta);
+    }
+    uint32_t equal_before = 0;
+    if (!complete) {
+        FastScan(smem.scan).ExclusiveSum(lane == 0 ? equal_count : 0U, equal_before);
+        __syncthreads();
+        equal_before = __shfl_sync(full_warp, equal_before, 0);
+    }
+    const uint32_t take_equal = complete ? 0U : min(equal_count, quota - min(quota, equal_before));
+    uint32_t slot;
+    FastScan(smem.scan).ExclusiveSum(lane == 0 ? greater_count + take_equal : 0U, slot);
+    __syncthreads();
+    slot = __shfl_sync(full_warp, slot, 0);
+    for (uint32_t base = lo; base < hi; base += 32) {
+        const uint32_t i = base + lane;
+        uint32_t key = 0;
+        if (i < hi) key = ordered_key(static_cast<float>(input[i]));
+        bool take = i < hi && (complete ? key >= threshold : key > threshold);
+        if (!complete) {
+            const uint32_t equals = __ballot_sync(full_warp, i < hi && key == threshold);
+            if ((equals & (1U << lane)) && equal_before + __popc(equals & lower_lanes) < quota)
+                take = true;
+            equal_before += __popc(equals);
+        }
+        const uint32_t winners = __ballot_sync(full_warp, take);
+        if (take) smem.indices[slot + __popc(winners & lower_lanes)] = i;
+        slot += __popc(winners);
+    }
+    __syncthreads();
+    for (uint32_t slot = tid; slot < args.topk; slot += FAST_THREADS) {
+        const uint32_t index = smem.indices[slot];
+        output_index[slot] = static_cast<int32_t>(static_cast<int64_t>(index) + offset);
+        if constexpr (ReturnValue) output_value[slot] = input[index];
+    }
+}
+
+template<typename ValueT, int Capacity>
+void launch_unsorted(const TopkSelectArgs& args) {
+    if (args.return_value)
+        topk_unsorted_kernel<ValueT, Capacity, true><<<args.batch_size, FAST_THREADS, 0, args.stream>>>(args);
+    else
+        topk_unsorted_kernel<ValueT, Capacity, false><<<args.batch_size, FAST_THREADS, 0, args.stream>>>(args);
+}
+
+template<typename ValueT>
+void launch(const TopkSelectArgs& args, bool indices_int64) {
+    if (!indices_int64 && !args.sorted_value && !args.sorted_index && args.topk <= 512) {
+        if (args.topk == STREAM_TOPK && args.vocab_size > 2048 && args.vocab_size <= 131072) {
+            if (args.return_value)
+                topk_streaming_kernel<ValueT, true><<<args.batch_size, FAST_THREADS, 0, args.stream>>>(args);
+            else
+                topk_streaming_kernel<ValueT, false><<<args.batch_size, FAST_THREADS, 0, args.stream>>>(args);
+        } else launch_unsorted<ValueT, 512>(args);
+    } else if (indices_int64)
+        topk_kernel<ValueT, int64_t><<<args.batch_size, NUM_THREADS, 0, args.stream>>>(args);
+    else
+        topk_kernel<ValueT, int32_t><<<args.batch_size, NUM_THREADS, 0, args.stream>>>(args);
+    check_kernel_launch();
+}
+
+} // namespace
+
+bool use_segmented_topk(const TopkSelectArgs& args, bool indices_int64) {
+    return !indices_int64 && !args.sorted_value && !args.sorted_index &&
+        args.topk == SEGMENTED_TOPK && args.vocab_size >= 65536 && args.vocab_size <= 131072 &&
+        args.batch_size >= 1 && args.batch_size <= 16;
+}
+
+void run_segmented_topk_select_kernel(const TopkSelectArgs& args, bool is_bfloat16, int32_t* workspace) {
+    if (is_bfloat16) launch_segmented<nv_bfloat16>(args, workspace);
+    else launch_segmented<float>(args, workspace);
+}
+
+void run_topk_select_kernel(const TopkSelectArgs& args, bool is_bfloat16, bool indices_int64) {
+    if (args.batch_size == 0) return;
+    if (is_bfloat16) launch<nv_bfloat16>(args, indices_int64);
+    else launch<float>(args, indices_int64);
+}
+
+} // namespace topk_select_sm120
diff --git a/csrc/cuda_kernels/sm120/topk_select.h b/csrc/cuda_kernels/sm120/topk_select.h
new file mode 100644
index 0000000..38e0872
--- /dev/null
+++ b/csrc/cuda_kernels/sm120/topk_select.h
@@ -0,0 +1,13 @@
+#pragma once
+
+#include "structs.h"
+
+namespace topk_select_sm120 {
+inline constexpr int SEGMENTED_TOPK = 512;
+inline constexpr int SEGMENTED_PARTITIONS = 8;
+inline constexpr int SEGMENTED_WORKSPACE_STRIDE = 2 * SEGMENTED_TOPK + 1;
+
+bool use_segmented_topk(const TopkSelectArgs& args, bool indices_int64);
+void run_segmented_topk_select_kernel(const TopkSelectArgs& args, bool is_bfloat16, int32_t* workspace);
+void run_topk_select_kernel(const TopkSelectArgs& args, bool is_bfloat16, bool indices_int64);
+}
diff --git a/csrc/cuda_kernels/v3_cluster/topk_select.cuh b/csrc/cuda_kernels/v3_cluster/topk_select.cuh
index bdbfef5..85dd30f 100644
--- a/csrc/cuda_kernels/v3_cluster/topk_select.cuh
+++ b/csrc/cuda_kernels/v3_cluster/topk_select.cuh
@@ -23,8 +23,7 @@ Algorithm:
 namespace topk_select_bf16_cluster {
 
 CUTE_DEVICE
-static void st_async_32b(uint32_t dst_addr, const uint32_t& data, transac_bar_t &mbar) {
-    uint32_t mbar_addr = cute::cast_smem_ptr_to_uint(&mbar);
+static void st_async_32b(uint32_t dst_addr, const uint32_t& data, uint32_t mbar_addr) {
     asm volatile (
         "st.async.weak.shared::cluster.mbarrier::complete_tx::bytes.s32 [%0], {%1}, [%2]; \n"
         :
@@ -32,6 +31,18 @@ static void st_async_32b(uint32_t dst_addr, const uint32_t& data, transac_bar_t
     );
 }
 
+template<typename T>
+CUTE_DEVICE
+static void st_async_128b(uint32_t dst_addr, const T& data, uint32_t mbar_addr) {
+    static_assert(sizeof(T) == 16);
+    long2 data_long2 = *reinterpret_cast<const long2*>(&data);
+    asm volatile (
+        "st.async.weak.shared::cluster.mbarrier::complete_tx::bytes.v2.s64 [%0], {%1, %2}, [%3]; \n"
+        :
+        : "r"(dst_addr), "l"(data_long2.x), "l"(data_long2.y), "r"(mbar_addr)
+    );
+}
+
 template<typename Config>
 class TopkSelectKernelBF16Cluster : public topk_select_common::TopkSelectKernelBF16Base<Config> {
     using BF16Base = topk_select_common::TopkSelectKernelBF16Base<Config>;
@@ -160,12 +171,14 @@ public:
         ku::barrier_cluster_wait_acquire();
 
         static_assert(NUM_GATHER_UNITS * sizeof(uint32_t) <= sizeof(smem.incoming_topk_pairs));
-        void* val_dst = (void*)(int64_t)cute::set_block_rank(
+        uint32_t val_dst = cute::set_block_rank(
             cute::cast_smem_ptr_to_uint((uint32_t*)smem.incoming_topk_pairs + rank_in_cluster * (MAX_TOPK / 2)), 0);
             
         static_assert(NUM_GATHER_PAIRS * sizeof(uint64_t) <= sizeof(SharedMemoryPlanBase::tma_load_buf));
-        void* dst = (void*)(int64_t)cute::set_block_rank(
+        uint32_t dst = cute::set_block_rank(
             cute::cast_smem_ptr_to_uint(reinterpret_cast<uint64_t*>(smem.tma_load_buf) + rank_in_cluster * MAX_TOPK), 0);
+        uint32_t gather_val_bar_addr = cute::set_block_rank(cute::cast_smem_ptr_to_uint(&smem.gather_val_bar), 0);
+        uint32_t gather_bar_addr = cute::set_block_rank(cute::cast_smem_ptr_to_uint(&smem.gather_bar), 0);
 
         static_assert(MAX_TOPK*sizeof(ValueT) % 16 == 0);
         constexpr uint32_t NUM_VAL_CHUNKS = MAX_TOPK * sizeof(ValueT) / 16; // 16B store
@@ -179,7 +192,7 @@ public:
             st_async_32b(
                 cute::set_block_rank(cute::cast_smem_ptr_to_uint(smem.gathered_num_survivors + rank_in_cluster), 0),
                 (uint32_t)nan_seen,
-                smem.gather_val_bar
+                gather_val_bar_addr
             );
         }
 
@@ -197,8 +210,8 @@ public:
                 topk_select_common::ld_shared<4>(pair2, (const uint32_t*)(smem.surviving_topk_pairs[survivor_buf_idx] + 8 * c + 2 * j));
                 vw[j] = __byte_perm(pair2[1], pair2[3], 0x5410);  
             }
-            ku::st_async((char*)val_dst + c * sizeof(uint4), make_uint4(vw[0], vw[1], vw[2], vw[3]),
-                        smem.gather_val_bar);
+            st_async_128b(val_dst + c * sizeof(uint4), make_uint4(vw[0], vw[1], vw[2], vw[3]),
+                          gather_val_bar_addr);
         }
 
         // And then store those (value, index) pairs
@@ -209,7 +222,7 @@ public:
         for (uint32_t i = 0; i < NUM_INDEX_VALUE_CHUNKS / NUM_THREADS; ++i) {
             uint32_t s = 2 * (i * NUM_THREADS + threadIdx.x);
             ulonglong2 pp = *reinterpret_cast<const ulonglong2*>(smem.surviving_topk_pairs[survivor_buf_idx] + s);
-            ku::st_async((char*)dst + s * sizeof(uint64_t), pp, smem.gather_bar);
+            st_async_128b(dst + s * sizeof(uint64_t), pp, gather_bar_addr);
         }
 
         if (rank_in_cluster != 0) {
diff --git a/csrc/stable_tensor_checks.h b/csrc/stable_tensor_checks.h
index 916ca42..8b01889 100644
--- a/csrc/stable_tensor_checks.h
+++ b/csrc/stable_tensor_checks.h
@@ -6,11 +6,14 @@
 #include <cstdint>
 #include <functional>
 #include <initializer_list>
+#include <limits>
 #include <optional>
 #include <type_traits>
 
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/shim_utils.h>
 
 namespace deep_select {
 
@@ -56,6 +59,88 @@ static inline PtrT* get_optional_tensor_ptr(const std::optional<Tensor>& tensor_
     }
 }
 
+static inline uint64_t checked_add(uint64_t a, uint64_t b) {
+    STD_TORCH_CHECK(a <= uint64_t(std::numeric_limits<int64_t>::max()) &&
+                    b <= uint64_t(std::numeric_limits<int64_t>::max()) - a,
+                    "tensor address arithmetic exceeds int64");
+    return a + b;
+}
+
+static inline uint64_t checked_mul(uint64_t a, uint64_t b) {
+    STD_TORCH_CHECK(a == 0 || b <= uint64_t(std::numeric_limits<int64_t>::max()) / a,
+                    "tensor address arithmetic exceeds int64");
+    return a * b;
+}
+
+static inline uint64_t storage_size_bytes(const Tensor& tensor) {
+    int64_t bytes = 0;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(tensor.get(), &bytes));
+    STD_TORCH_CHECK(bytes >= 0, "tensor storage size must be nonnegative");
+    return static_cast<uint64_t>(bytes);
+}
+
+static inline uint64_t storage_offset_bytes(const Tensor& tensor) {
+    int64_t offset = 0;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_offset(tensor.get(), &offset));
+    STD_TORCH_CHECK(offset >= 0, "tensor storage offset must be nonnegative");
+    return checked_mul(static_cast<uint64_t>(offset), tensor.element_size());
+}
+
+static inline void check_storage_bounds(const char* name, const Tensor& tensor, uint64_t row_padding = 1) {
+    for (int64_t dim = 0; dim < tensor.dim(); ++dim) {
+        STD_TORCH_CHECK(tensor.size(dim) >= 0 && tensor.stride(dim) >= 0,
+                        name, " sizes and strides must be nonnegative");
+    }
+    if (tensor.numel() == 0) return;
+    uint64_t last = 0;
+    for (int64_t dim = 0; dim < tensor.dim(); ++dim) {
+        last = checked_add(last, checked_mul(tensor.size(dim) - 1, tensor.stride(dim)));
+    }
+    const uint64_t itemsize = tensor.element_size();
+    uint64_t end = checked_add(storage_offset_bytes(tensor), checked_mul(checked_add(last, 1), itemsize));
+    if (row_padding != 1) {
+        const uint64_t row_bytes = checked_mul(tensor.size(-1), itemsize);
+        const uint64_t padded = checked_mul(checked_add(row_bytes, row_padding - 1) / row_padding, row_padding);
+        end = checked_add(end, padded - row_bytes);
+    }
+    STD_TORCH_CHECK(end <= storage_size_bytes(tensor), name,
+                    " backing storage must include all elements and the final row padded to a multiple of ",
+                    row_padding, " bytes");
+}
+
+struct StorageInterval {
+    uintptr_t start;
+    uint64_t bytes;
+};
+
+static inline StorageInterval storage_interval(const Tensor& tensor) {
+    void* data = nullptr;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(tensor.get(), &data));
+    const uintptr_t pointer = reinterpret_cast<uintptr_t>(data);
+    const uint64_t offset = storage_offset_bytes(tensor);
+    const uint64_t bytes = storage_size_bytes(tensor);
+    STD_TORCH_CHECK(offset <= bytes && offset <= pointer, "invalid tensor storage offset");
+    const uintptr_t start = pointer - offset;
+    STD_TORCH_CHECK(bytes <= std::numeric_limits<uintptr_t>::max() - start,
+                    "tensor storage address range overflows uintptr_t");
+    return {start, bytes};
+}
+
+static inline void check_storage_disjoint(const char* output_name, const Tensor& output,
+                                          const char* other_name, const Tensor& other) {
+    if (output.numel() == 0 || other.numel() == 0) return;
+    const auto a = storage_interval(output);
+    const auto b = storage_interval(other);
+    const bool overlaps = a.start <= b.start ? b.start - a.start < a.bytes : a.start - b.start < b.bytes;
+    STD_TORCH_CHECK(!overlaps, "`", output_name, "` must not overlap storage with `", other_name,
+                    "` (including disjoint views of shared storage)");
+}
+
+static inline void check_pointer_alignment(const char* name, const Tensor& tensor) {
+    STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(tensor.data_ptr()) % 32 == 0,
+                    name, ".data_ptr() must be 32-byte aligned");
+}
+
 } // namespace deep_select
 
 // Check whether the given tensor (or optional<tensor>) is on CUDA GPU
diff --git a/deep_select/interface.py b/deep_select/interface.py
index b692b29..a350bbe 100644
--- a/deep_select/interface.py
+++ b/deep_select/interface.py
@@ -35,22 +35,36 @@ def topk(
     """
     Arguments:
         input: (b, vocab_size), dtype=torch.bfloat16/torch.float. stride(0) must be a multiple of `deep_select.get_stride_requirement()[0]` bytes, and stride(1) must be 1.
-        topk: int. Select topk elements for each row.
-        sorted: bool. Whether to return sorted **output_val**. Only supports fp32.
+               b must fit int32 and 0 <= vocab_size < 2**23. SM100/SM103 require vocab_size > 0 when b > 0;
+               SM120/SM121 support zero-width rows. Empty batches remain supported on all compiled targets.
+               Nonempty backing storage must include the final row padded to 128 bytes.
+        topk: int, 1 <= topk <= 4096. Select topk elements for each row.
+        sorted: bool. Whether to return sorted **output_val**. Only supports fp32, requires return_value, and excludes sorted_index.
         begin(optional): (b,), dtype=int32. CURRENTLY NOT SUPPORTED. The left(inclusive) range for input row, default is 0.
         end(optional): (b,), dtype=int32. The right(exclusive) range for input row, default is vocab_size. The stride of this tensor must be 1.
                        Note when end[i] <= topk, valid elements will be gathered at the beginning of values and indices returned. The rest of `values` will be filled with `value_oob_fill_value`, while the rest of `indices` will be filled with `idx_oob_fill_value` (won't be plused by `output_idx_offset`).
-                       `end` <= `vocab_size` must be held
+                       The caller must ensure 0 <= end[i] <= vocab_size; host validation does not read CUDA tensor contents.
         indices_type: torch.dtype. The output indices dtype, only support torch.int32 and torch.int64.
         sorted_index: bool. Whether to return sorted **output_idx**.
         hint(optional): CURRENTLY NOT SUPPORTED
-        output_idx(optional): (b, topk), dtype=indices_type. A contiguous tensor to store output.
+        output_idx(optional): (b, topk), dtype=indices_type. Last dimension contiguous, row stride aligned to 32 bytes, with nonoverlapping rows.
         output_idx_offset(optional): (b,), dtype=int32. If provided, all output_idx (`idx_oob_fill_value` not included) will += output_idx_offset.
-        idx_oob_fill_value: int. See comments above when end[i]-begin[i]<topk.
+        idx_oob_fill_value: int32-representable int. See comments above when end[i]-begin[i]<topk.
+        value_oob_fill_value: float32-representable finite value, infinity, or NaN. Used for short-row padding.
         return_value: bool. If False, only return indices without values to accelerate the kernel. The return value is still a Tuple, but the first element will be None.
         abort_when_nan_found: bool. When a NaN is found, if True, aborts the whole kernel; if False, writes 0x3F3F3F3F to the corresponding output_idx[batch_idx][0] and exits.
                 The NaN check itself is always enabled. Exception: when the row's length <= topk, it is skipped.
 
+    Contract:
+        All tensors must be on input's CUDA device. Input and output pointers must be 32-byte aligned.
+        Neither output may share or overlap backing storage with input, end, output_idx_offset, or the other output,
+        even through disjoint views; a pair containing an empty tensor is exempt from this alias check.
+        These checks also apply to direct torch.ops.deep_select.topk calls and supplied output_value when return_value=False.
+        Native kernels run on the current PyTorch CUDA stream, including SM120/SM121 segmented selection with int32 workspace (b, 8, 1025).
+        CUDA capabilities 12.0 and 12.1 use the same native algorithms, each enabled only when its exact target was compiled.
+        Other 12.x capabilities are rejected. SM121 hardware validation is still required.
+        Only compiled CUDA targets are supported; there is no torch.topk fallback.
+
     Return:
         output_val: (b, topk), dtype=input.dtype.
         output_idx: (b, topk), dtype=indices_type.
diff --git a/setup.py b/setup.py
index bc0b06e..315efb0 100644
--- a/setup.py
+++ b/setup.py
@@ -99,6 +99,7 @@
 
 def build_on_cuda_platform():
     from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
+    from csrc.build_utils import build_plan, grouped_build_extension, parse_cuda_arch_list
 
     assert CUDA_HOME is not None, "PyTorch must be compiled with CUDA support"
 
@@ -113,24 +114,17 @@ def append_nvcc_threads(nvcc_extra_args):
     nvcc_version_number = nvcc_version.split("release ")[1].split(",")[0].strip()
     major, minor = map(int, nvcc_version_number.split("."))
     print(f"Compiling using NVCC {major}.{minor}")
-    if major < 12 or (major == 12 and minor <= 8):
-        raise RuntimeError("sm100 compilation requires NVCC 12.9 or higher.")
-
-    cc_flag = [
-        # Currently skip build for sm80 and sm90 to speed up compilation
-        # "-gencode", "arch=compute_80,code=sm_80",
-        # "-gencode", "arch=compute_90a,code=sm_90a",
-
-        # Compile sm100 and sm103 separately to give the compiler more room for optimization
-        "-gencode", "arch=compute_100a,code=sm_100a",
-        "-gencode", "arch=compute_103a,code=sm_103a",
-    ]
+    architectures = parse_cuda_arch_list(
+        os.getenv("DEEP_SELECT_CUDA_ARCH_LIST"), (major, minor)
+    )
+    sources, macros, group_flags = build_plan(CUDA_SOURCES, architectures)
 
     this_dir = os.path.dirname(os.path.abspath(__file__))
 
     ext_modules = [CUDAExtension(
         name="deep_select.deep_select_cuda",
-        sources=CUDA_SOURCES,
+        sources=sources,
+        define_macros=macros,
         py_limited_api=True,
         extra_compile_args={
             # Target the torch 2.10 stable ABI (minimum supported torch version)
@@ -152,7 +146,7 @@ def append_nvcc_threads(nvcc_extra_args):
                 "--ptxas-options=-v,--register-usage-level=10,--warn-on-spills,--warn-on-double-precision-use",
                 "-lineinfo",
                 "--source-in-ptx",
-            ] + cc_flag)
+            ])
         },
         include_dirs=[
             Path(this_dir) / "csrc",
@@ -168,7 +162,7 @@ def append_nvcc_threads(nvcc_extra_args):
         ],
     )]
 
-    class SpillCheckBuildExtension(BuildExtension):
+    class SpillCheckBuildExtension(grouped_build_extension(BuildExtension, group_flags)):
         STACK_BASELINE = 8  # Because of we're using `printf`
 
         def run(self):
diff --git a/tests/kernelkit/build.py b/tests/kernelkit/build.py
index 87ee33b..31a4232 100644
--- a/tests/kernelkit/build.py
+++ b/tests/kernelkit/build.py
@@ -2,6 +2,7 @@
 import subprocess
 import re
 import sys
+import shutil
 from typing import List, Optional
 
 from .platform import Platform, requires_platform, get_current_platform
@@ -33,9 +34,10 @@ def check_kernel_reg_spill_in_artifact(artifact_path: str, stack_baseline: int =
         print(f"Checking register spills in: {artifact_path}")
 
     cuda_home = CUDA_HOME if CUDA_HOME is not None else '/usr/local/cuda'
-    cuobjdump_path = os.path.join(cuda_home, 'bin/cuobjdump')
-    if not os.path.exists(cuobjdump_path):
-        raise FileNotFoundError(f"cuobjdump not found (looked at {cuobjdump_path})")
+    toolkit_cuobjdump = os.path.join(cuda_home, 'bin/cuobjdump')
+    cuobjdump_path = shutil.which(toolkit_cuobjdump) or shutil.which('cuobjdump')
+    if cuobjdump_path is None:
+        raise FileNotFoundError(f"cuobjdump not found (looked at {toolkit_cuobjdump} and PATH)")
     
     try:
         result = subprocess.run(
diff --git a/tests/test_build_plan.py b/tests/test_build_plan.py
new file mode 100644
index 0000000..e779de1
--- /dev/null
+++ b/tests/test_build_plan.py
@@ -0,0 +1,116 @@
+import ast
+import os
+import subprocess
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from csrc.build_utils import architecture_flags, build_plan, parse_cuda_arch_list
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def setup_namespace():
+    tree = ast.parse((ROOT / "setup.py").read_text())
+    nodes = [node for node in tree.body
+             if isinstance(node, ast.FunctionDef) and node.name == "build_on_cuda_platform"
+             or isinstance(node, ast.Assign) and any(
+                 isinstance(target, ast.Name) and target.id == "CUDA_SOURCES"
+                 for target in node.targets)]
+    namespace = {"os": os, "subprocess": subprocess, "Path": Path,
+                 "__file__": str(ROOT / "setup.py"),
+                 "kk": SimpleNamespace(check_kernel_reg_spill_in_artifact=Mock(return_value=[]))}
+    exec(compile(ast.Module(body=nodes, type_ignores=[]), "setup.py", "exec"), namespace)
+    return namespace
+
+
+@pytest.mark.parametrize("entry", ["12.1", "121", "12.1+PTX", "121+PTX"])
+def test_sm121_minimum_and_ptx(entry):
+    with pytest.raises(RuntimeError, match="requires CUDA 12.9"):
+        parse_cuda_arch_list(entry, (12, 8))
+    architectures = parse_cuda_arch_list(entry, (12, 9))
+    assert architectures == (("121", entry.endswith("+PTX")),)
+    flags = architecture_flags(architectures)
+    assert "arch=compute_121,code=sm_121" in flags
+    assert ("arch=compute_121,code=compute_121" in flags) == entry.endswith("+PTX")
+
+
+@pytest.mark.parametrize("entry", ["", "12.1a", "12.1f", "12.1+ptx", "9.0;12.1"])
+def test_invalid_architecture(entry):
+    with pytest.raises(ValueError):
+        parse_cuda_arch_list(entry, (13, 2))
+
+
+def test_default_and_deduplication():
+    assert parse_cuda_arch_list(None, (12, 9)) == (("100a", False), ("103a", False))
+    assert parse_cuda_arch_list("121;12.1+PTX 121", (12, 9)) == (("121", True),)
+
+
+@pytest.mark.parametrize("entry,expected_macros,counts", [
+    ("12.0", (0, 1, 0), [1, 1]),
+    ("12.1+PTX", (0, 0, 1), [1, 1]),
+    ("12.0;12.1+PTX", (0, 1, 1), [1, 1]),
+    ("10.0a;10.3a;12.0;12.1", (1, 1, 1), [1, 78, 1]),
+    ("10.0f", (1, 0, 0), [1, 78]),
+])
+def test_single_extension_grouped_compile(monkeypatch, tmp_path, entry, expected_macros, counts):
+    import torch.utils.cpp_extension as cpp
+    from setuptools import Distribution
+    from setuptools._distutils.ccompiler import new_compiler
+    from setuptools._distutils.sysconfig import customize_compiler
+
+    namespace = setup_namespace()
+    architectures = parse_cuda_arch_list(entry, (12, 9))
+    sources, macros, _ = build_plan(namespace["CUDA_SOURCES"], architectures)
+    assert tuple(int(value) for _, value in macros) == expected_macros
+    assert sources.count("csrc/api.cpp") == 1
+    monkeypatch.setenv("DEEP_SELECT_CUDA_ARCH_LIST", entry)
+    monkeypatch.setattr(subprocess, "check_output", Mock(return_value=b"CUDA release 12.9, V12.9.0"))
+    monkeypatch.setattr(cpp, "CUDA_HOME", str(tmp_path))
+    extensions, cls = namespace["build_on_cuda_platform"]()
+    assert len(extensions) == 1
+    extension = extensions[0]
+    assert extension.name == "deep_select.deep_select_cuda"
+    assert extension.py_limited_api
+    command = cls(Distribution({"ext_modules": extensions}))
+    command.ensure_finalized()
+    command.build_temp = str(tmp_path / "objects")
+    command.build_lib = str(tmp_path / "lib")
+    command.compiler = new_compiler()
+    customize_compiler(command.compiler)
+    monkeypatch.setattr(command, "_check_abi", lambda: ("g++", cpp.TorchVersion("13.0")))
+    monkeypatch.setattr(cpp, "_check_cuda_version", Mock())
+    compile_mock = Mock()
+    link_mock = Mock()
+    monkeypatch.setattr(cpp, "_write_ninja_file_and_compile_objects", compile_mock)
+    monkeypatch.setattr(command.compiler, "link_shared_object", link_mock)
+    command.build_extensions()
+    calls = [call.kwargs for call in compile_mock.call_args_list]
+    assert [len(call["sources"]) for call in calls] == counts
+    assert link_mock.call_count == 1
+    assert len(link_mock.call_args.args[0]) == sum(counts)
+    assert command.get_ext_filename(extension.name).endswith(".abi3.so")
+    assert "-DPy_LIMITED_API=0x030A0000" in calls[0]["post_cflags"]
+    for call in calls[1:]:
+        flags = call["cuda_post_cflags"]
+        assert "-DTORCH_TARGET_VERSION=0x020a000000000000" in flags
+        assert flags.count("--use_fast_math") == 1
+        assert flags.count("--ftz=false") == 1
+        assert flags.index("--use_fast_math") < flags.index("--ftz=false")
+        if "/sm120/" in call["sources"][0]:
+            expected = tuple(item for item in architectures if item[0] in {"120", "121"})
+        else:
+            expected = tuple(item for item in architectures if item[0] not in {"120", "121"})
+        actual = [flag for flag in flags if flag.startswith("arch=compute_")]
+        assert actual == [flag for flag in architecture_flags(expected) if flag.startswith("arch=")]
+    monkeypatch.setattr(cpp.BuildExtension, "run", Mock())
+    command.run()
+    checker = namespace["kk"].check_kernel_reg_spill_in_artifact
+    checker.assert_called_once_with(command.get_ext_fullpath(extension.name), stack_baseline=8,
+                                    suppress_checking_env_var="DEEP_SELECT_DISABLE_REG_SPILL_CHECK")
+    checker.return_value = ["spilling_kernel"]
+    with pytest.raises(RuntimeError, match="Register spilling detected"):
+        command.run()
diff --git a/tests/test_sm120.py b/tests/test_sm120.py
new file mode 100644
index 0000000..875a180
--- /dev/null
+++ b/tests/test_sm120.py
@@ -0,0 +1,1144 @@
+from builtins import sorted as builtins_sorted
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+import torch
+
+
+MODES = [
+    pytest.param(False, False, False, id="indices"),
+    pytest.param(False, False, True, id="values"),
+    pytest.param(False, True, False, id="sorted-indices"),
+    pytest.param(False, True, True, id="sorted-indices-values"),
+    pytest.param(True, False, True, id="sorted-values"),
+]
+DTYPES = [pytest.param(torch.float32, id="fp32"), pytest.param(torch.bfloat16, id="bf16")]
+INDEX_DTYPES = [pytest.param(torch.int32, id="i32"), pytest.param(torch.int64, id="i64")]
+TOPKS = [1, 7, 512, 1024, 4096]
+
+
+@pytest.fixture(scope="module")
+def device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    device = torch.device("cuda", torch.cuda.current_device())
+    if torch.cuda.get_device_capability(device) not in ((12, 0), (12, 1)):
+        pytest.skip("native SM120/SM121 tests require compute capability 12.0 or 12.1")
+    return device
+
+
+class NativeAPI:
+    @staticmethod
+    def get_stride_requirement():
+        return torch.ops.deep_select.get_alignment_requirement()
+
+    @staticmethod
+    def topk(input, topk, sorted=False, begin=None, end=None, indices_type=torch.int64,
+             sorted_index=False, output_idx=None, output_idx_offset=None,
+             idx_oob_fill_value=2147483647, value_oob_fill_value=float("-inf"),
+             return_value=True, abort_when_nan_found=True):
+        values = (aligned_tensor(NativeAPI, input.shape[0], topk, input.dtype,
+                                 input.device, output=True)[0] if return_value else None)
+        if output_idx is None:
+            output_idx = aligned_tensor(NativeAPI, input.shape[0], topk, indices_type,
+                                        input.device, output=True)[0]
+        result = torch.ops.deep_select.topk(
+            input, topk, begin, end, sorted, sorted_index, values, output_idx,
+            output_idx_offset, idx_oob_fill_value, value_oob_fill_value,
+            return_value, abort_when_nan_found)
+        assert result is None
+        return values, output_idx
+
+
+@pytest.fixture(scope="module", params=["wrapper", "native"])
+def api(device, request):
+    import deep_select
+
+    return deep_select if request.param == "wrapper" else NativeAPI
+
+
+def aligned_tensor(api, batch, width, dtype, device, output=False, offset=False):
+    alignment = api.get_stride_requirement()[int(output)] // dtype.itemsize
+    stride = max(alignment, (width + alignment - 1) // alignment * alignment)
+    if offset:
+        stride += 2 * alignment
+    storage = torch.empty((batch + int(offset), stride), dtype=dtype, device=device)
+    row = int(offset)
+    column = alignment if offset else 0
+    view = storage[row:row + batch, column:column + width]
+    assert view.stride(0) * dtype.itemsize % api.get_stride_requirement()[int(output)] == 0
+    return view, storage
+
+
+def check_result(api, x, k, result, *, indices_type=torch.int64, return_value=True,
+                 sorted=False, sorted_index=False, end=None, output_idx_offset=None,
+                 idx_oob_fill_value=-1, value_oob_fill_value=float("-inf")):
+    values, indices = result
+    assert indices.shape == (x.shape[0], k)
+    assert indices.dtype == indices_type
+    assert indices.device == x.device
+    assert indices.stride(1) == 1
+    assert indices.stride(0) * indices.element_size() % api.get_stride_requirement()[1] == 0
+    if return_value:
+        assert values is not None
+        assert values.shape == indices.shape and values.dtype == x.dtype
+        assert values.device == x.device and values.stride(1) == 1
+        assert values.stride(0) * values.element_size() % api.get_stride_requirement()[1] == 0
+    else:
+        assert values is None
+    source = x.cpu()
+    actual_indices = indices.cpu().to(torch.int64)
+    actual_values = values.cpu() if values is not None else None
+    ends = end.cpu().tolist() if end is not None else [x.shape[1]] * x.shape[0]
+    offsets = output_idx_offset.cpu().tolist() if output_idx_offset is not None else [0] * x.shape[0]
+    for row, (length, offset) in enumerate(zip(ends, offsets)):
+        count = min(k, length)
+        chosen = actual_indices[row, :count] - offset
+        assert bool(((chosen >= 0) & (chosen < length)).all()), (row, chosen)
+        assert chosen.unique().numel() == count, (row, "duplicate indices")
+        assert bool((actual_indices[row, count:] == idx_oob_fill_value).all())
+        gathered = source[row, chosen]
+        if actual_values is not None:
+            bits = torch.int32 if x.dtype == torch.float32 else torch.int16
+            assert torch.equal(actual_values[row, :count].view(bits), gathered.view(bits))
+            assert bool((actual_values[row, count:] == value_oob_fill_value).all())
+        if count:
+            reference = torch.topk(source[row, :length].float(), count).values
+            torch.testing.assert_close(gathered.float().sort(descending=True).values,
+                                       reference, rtol=0, atol=0)
+        if sorted_index:
+            assert bool((chosen[1:] > chosen[:-1]).all())
+        if sorted:
+            assert bool((gathered[:-1] >= gathered[1:]).all())
+
+
+def run_checked(api, x, k, **kwargs):
+    kwargs.setdefault("idx_oob_fill_value", -1)
+    result = api.topk(x, k, **kwargs)
+    check_result(api, x, k, result, **{key: value for key, value in kwargs.items()
+                                     if key not in ("output_idx", "abort_when_nan_found")})
+    return result
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_zero_batch(api, device, dtype, indices_type, return_value):
+    x, _ = aligned_tensor(api, 0, 65536, dtype, device)
+    end = torch.empty(0, dtype=torch.int32, device=device)
+    offsets = torch.empty(0, dtype=torch.int32, device=device)
+    assert x.shape == (0, 65536)
+    for metadata in [{}, dict(end=end, output_idx_offset=offsets)]:
+        result = run_checked(api, x, 512, indices_type=indices_type,
+                             return_value=return_value, **metadata)
+        assert result[1].numel() == 0
+        if return_value:
+            assert result[0].numel() == 0
+    torch.cuda.synchronize(device)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("sorted,sorted_index,return_value", MODES)
+@pytest.mark.parametrize("k", TOPKS)
+@pytest.mark.parametrize("width_kind", ["empty", "short", "equal", "16384", "129280"])
+def test_matrix(api, device, dtype, indices_type, sorted, sorted_index, return_value, k, width_kind):
+    if dtype == torch.bfloat16 and sorted:
+        pytest.skip("sorted values are FP32-only")
+    widths = [0, max(0, k - 1), k, 16384, 129280]
+    width_index = ["empty", "short", "equal", "16384", "129280"].index(width_kind)
+    width = widths[width_index]
+    batch = [1, 4, 16][(TOPKS.index(k) + width_index) % 3]
+    x, _ = aligned_tensor(api, batch, width, dtype, device)
+    generator = torch.Generator(device=device).manual_seed(1234 + k + width)
+    x.copy_(torch.randn(x.shape, dtype=dtype, device=device, generator=generator))
+    run_checked(api, x, k, indices_type=indices_type, sorted=sorted,
+                sorted_index=sorted_index, return_value=return_value)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("sorted,sorted_index,return_value", MODES)
+@pytest.mark.parametrize("k", TOPKS)
+def test_end_offset_and_storage(api, device, dtype, indices_type, sorted, sorted_index, return_value, k):
+    if dtype == torch.bfloat16 and sorted:
+        pytest.skip("sorted values are FP32-only")
+    width = 16384
+    x, input_storage = aligned_tensor(api, 4, width, dtype, device, offset=True)
+    input_storage.fill_(float("nan"))
+    x.copy_(((torch.arange(width, device=device) * 71) % 997).to(dtype).expand_as(x))
+    before = input_storage.clone()
+    end_storage = torch.tensor([99, 0, k - 1, k, width], dtype=torch.int32, device=device)
+    end = end_storage[1:]
+    offset_storage = torch.tensor([99, 100, 200, -300, 400], dtype=torch.int32, device=device)
+    offsets = offset_storage[1:]
+    out, output_storage = aligned_tensor(api, 4, k, indices_type, device, output=True, offset=True)
+    output_storage.fill_(-987654)
+    assert not x.is_contiguous() and not out.is_contiguous()
+    assert x.storage_offset() > 0 and out.storage_offset() > 0
+    result = run_checked(api, x, k, indices_type=indices_type, sorted=sorted,
+                         sorted_index=sorted_index, return_value=return_value, end=end,
+                         output_idx_offset=offsets, output_idx=out)
+    assert result[1] is out
+    untouched = torch.ones_like(output_storage, dtype=torch.bool)
+    alignment = api.get_stride_requirement()[1] // indices_type.itemsize
+    untouched[1:, alignment:alignment + k] = False
+    assert bool((output_storage[untouched] == -987654).all())
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    assert torch.equal(input_storage.view(bits), before.view(bits))
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("pattern", ["zeros", "equal", "infinities", "negative-infinity"])
+@pytest.mark.parametrize("sorted,sorted_index,return_value", MODES)
+def test_ties_and_infinities(api, device, dtype, indices_type, pattern, sorted, sorted_index, return_value):
+    if dtype == torch.bfloat16 and sorted:
+        pytest.skip("sorted values are FP32-only")
+    x, _ = aligned_tensor(api, 4, 16384, dtype, device)
+    if pattern == "zeros":
+        x.fill_(0.0)
+        x[:, ::2] = -0.0
+    elif pattern == "equal":
+        x.fill_(3.5)
+    elif pattern == "infinities":
+        x.fill_(-float("inf"))
+        x[:, ::3] = float("inf")
+        x[:, 1::3] = 0.0
+    else:
+        x.fill_(-float("inf"))
+    run_checked(api, x, 512, indices_type=indices_type, sorted=sorted,
+                sorted_index=sorted_index, return_value=return_value)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_nan_nonabort_and_short_rows(api, device, dtype, indices_type, return_value):
+    x, _ = aligned_tensor(api, 4, 16384, dtype, device)
+    x.fill_(1)
+    x[0, 1000] = float("nan")
+    x[1, 0] = float("nan")
+    x[2, 6] = float("nan")
+    x[3, 1000] = float("nan")
+    end = torch.tensor([16384, 3, 7, 100], dtype=torch.int32, device=device)
+    offsets = torch.tensor([100, 200, 300, 400], dtype=torch.int32, device=device)
+    values, indices = api.topk(x, 7, end=end, indices_type=indices_type,
+                               output_idx_offset=offsets, return_value=return_value,
+                               idx_oob_fill_value=-1, abort_when_nan_found=False)
+    assert indices[0, 0].item() == 0x3F3F3F3F
+    for row, count in [(1, 3), (2, 7)]:
+        torch.testing.assert_close(indices[row, :count].cpu().to(torch.int64),
+                                   torch.arange(count) + int(offsets[row]), rtol=0, atol=0)
+        assert bool((indices[row, count:] == -1).all())
+        if return_value:
+            bits = torch.int32 if dtype == torch.float32 else torch.int16
+            assert torch.equal(values[row, :count].view(bits), x[row, :count].view(bits))
+            assert bool(torch.isneginf(values[row, count:]).all())
+    check_result(api, x[3:], 7, (values[3:] if values is not None else None, indices[3:]),
+                 indices_type=indices_type, return_value=return_value, end=end[3:],
+                 output_idx_offset=offsets[3:])
+    for row in [1, 2]:
+        safe_values, safe_indices = api.topk(x[row:row + 1], 7, end=end[row:row + 1],
+                                             return_value=return_value, indices_type=indices_type,
+                                             idx_oob_fill_value=-1)
+        torch.cuda.synchronize(device)
+        assert safe_indices[0, 0].item() == 0
+        if return_value:
+            assert torch.isnan(safe_values[0]).any().item()
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("k", [7, 513, 2048, 2049])
+def test_unsorted_i32_selection_tails(api, device, dtype, return_value, k):
+    warp_end = (k // 32 + 1) * 32
+    block_end = (k // 256 + 1) * 256
+    lengths = builtins_sorted({0, k - 1, k, k + 1, warp_end - 1, warp_end,
+                              warp_end + 1, block_end - 1, block_end, block_end + 1})
+    width = max(lengths)
+    x, storage = aligned_tensor(api, len(lengths), width, dtype, device, offset=True)
+    storage.fill_(float("nan"))
+    for row, length in enumerate(lengths):
+        x[row, :length].copy_(((torch.arange(length, device=device) * 71) % 997).to(dtype))
+    end = torch.tensor(lengths, dtype=torch.int32, device=device)
+    offsets = torch.arange(len(lengths), dtype=torch.int32, device=device) * 100 - 300
+    out, output_storage = aligned_tensor(api, len(lengths), k, torch.int32, device,
+                                         output=True, offset=True)
+    output_storage.fill_(-987654)
+    result = run_checked(api, x, k, end=end, output_idx_offset=offsets, output_idx=out,
+                         indices_type=torch.int32, return_value=return_value,
+                         value_oob_fill_value=12345.0)
+    assert result[1] is out
+    untouched = torch.ones_like(output_storage, dtype=torch.bool)
+    alignment = api.get_stride_requirement()[1] // torch.int32.itemsize
+    untouched[1:, alignment:alignment + k] = False
+    assert bool((output_storage[untouched] == -987654).all())
+
+
+@pytest.mark.parametrize("dtype,raw", [
+    pytest.param(torch.float32, [0xC0A00000, 0x40400000, 0x41200000], id="fp32"),
+    pytest.param(torch.bfloat16, [0xC0A0, 0x4040, 0x4120], id="bf16"),
+    pytest.param(torch.bfloat16, [0xBF81, 0xBF80, 0xBF7F], id="bf16-negative-finite"),
+    pytest.param(torch.bfloat16, [0x8003, 0x8002, 0x8001], id="bf16-negative-subnormal"),
+])
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("k,width", [(511, 4097), (513, 4097), (2048, 8193)])
+def test_unsorted_i32_cross_warp_threshold_quota(api, device, dtype, raw, return_value, k, width):
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    sign = 1 << (dtype.itemsize * 8 - 1)
+    signed = [value - 2 * sign if value & sign else value for value in raw]
+    greater = [index for index in range(width) if index % 37 == 0]
+    equal = [index for index in range(width) if index % 3 == 1 and index % 37 != 0]
+    quota = k - len(greater)
+    assert 0 < quota < len(equal)
+    chunk = ((width + 255) // 256) * 32
+    assert equal[quota - 1] // chunk >= 2
+    expected = builtins_sorted(greater + equal[:quota])
+    x, _ = aligned_tensor(api, 1, width, dtype, device)
+    x.view(bits).fill_(signed[0])
+    x.view(bits)[0, equal] = signed[1]
+    x.view(bits)[0, greater] = signed[2]
+    for _ in range(3):
+        values, indices = run_checked(api, x, k, indices_type=torch.int32,
+                                      return_value=return_value)
+        actual = indices[0].cpu().tolist()
+        assert builtins_sorted(actual) == expected
+        if return_value:
+            expected_bits = [signed[2] if index % 37 == 0 else signed[1] for index in actual]
+            assert values[0].view(bits).cpu().tolist() == expected_bits
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("k", [513, 2048])
+def test_unsorted_i32_nan_at_end(api, device, dtype, return_value, k):
+    x, _ = aligned_tensor(api, 4, k + 2, dtype, device)
+    x.fill_(1)
+    lengths = [k + 1, k + 1, k - 1, k]
+    for row, nan_index in enumerate([k, k + 1, k - 2, k - 1]):
+        x[row, nan_index] = float("nan")
+    end = torch.tensor(lengths, dtype=torch.int32, device=device)
+    offsets = torch.tensor([100, -200, 300, -400], dtype=torch.int32, device=device)
+    values, indices = api.topk(x, k, end=end, output_idx_offset=offsets,
+                               indices_type=torch.int32, return_value=return_value,
+                               idx_oob_fill_value=-1, abort_when_nan_found=False)
+    assert indices.dtype == torch.int32
+    assert indices[0, 0].item() == 0x3F3F3F3F
+    check_result(api, x[1:2], k, (values[1:2] if values is not None else None, indices[1:2]),
+                 indices_type=torch.int32, return_value=return_value, end=end[1:2],
+                 output_idx_offset=offsets[1:2])
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    for row in [2, 3]:
+        length = lengths[row]
+        offset = offsets[row].item()
+        assert indices[row, :length].cpu().tolist() == list(range(offset, offset + length))
+        assert bool((indices[row, length:] == -1).all())
+        if return_value:
+            assert torch.equal(values[row, :length].view(bits), x[row, :length].view(bits))
+            assert bool(torch.isneginf(values[row, length:]).all())
+    if not return_value:
+        assert values is None
+    safe = api.topk(x[1:], k, end=end[1:], output_idx_offset=offsets[1:],
+                    indices_type=torch.int32, return_value=return_value, idx_oob_fill_value=-1)
+    torch.cuda.synchronize(device)
+    check_result(api, x[1:2], k, (safe[0][:1] if return_value else None, safe[1][:1]),
+                 indices_type=torch.int32, return_value=return_value, end=end[1:2],
+                 output_idx_offset=offsets[1:2])
+    assert torch.equal(safe[1][1:], indices[2:])
+    if return_value:
+        assert torch.equal(safe[0][1:].view(bits), values[2:].view(bits))
+    else:
+        assert safe[0] is None
+
+
+def run_checked_low_index(api, x, k, **kwargs):
+    result = run_checked(api, x, k, **kwargs)
+    bits = torch.int32 if x.dtype == torch.float32 else torch.int16
+    sign = 1 << (x.element_size() * 8 - 1)
+    mask = 2 * sign - 1
+    source = x.view(bits).cpu().tolist()
+    ends = kwargs["end"].cpu().tolist() if "end" in kwargs else [x.shape[1]] * x.shape[0]
+    offsets = kwargs["output_idx_offset"].cpu().tolist() if "output_idx_offset" in kwargs else [0] * x.shape[0]
+    for row, (length, offset) in enumerate(zip(ends, offsets)):
+        def key(index):
+            value = source[row][index] & mask
+            if value in (0, sign):
+                value = 0
+            return (~value & mask) if value & sign else value ^ sign
+        count = min(k, length)
+        expected = builtins_sorted(range(length), key=lambda index: (-key(index), index))[:count]
+        actual = (result[1][row, :count].cpu().to(torch.int64) - offset).tolist()
+        assert builtins_sorted(actual) == builtins_sorted(expected), row
+    return result
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("width", [2048, 2049, 4097, 131072, 131073])
+def test_streaming_k512_end_boundaries(api, device, dtype, return_value, width):
+    lengths = builtins_sorted({0, width, *[boundary + delta
+                              for boundary in [512, 1024, 1536, 2048, 2560, 4096, 131072]
+                              for delta in [-1, 0, 1] if boundary + delta <= width]})
+    x, storage = aligned_tensor(api, len(lengths), width, dtype, device, offset=True)
+    storage.fill_(float("nan"))
+    for row, length in enumerate(lengths):
+        x[row, :length].copy_(((torch.arange(length, device=device) * 71) % 997).to(dtype))
+    end = torch.tensor(lengths, dtype=torch.int32, device=device)
+    offsets = torch.arange(len(lengths), dtype=torch.int32, device=device) * 100 - 700
+    out, output_storage = aligned_tensor(api, len(lengths), 512, torch.int32, device,
+                                         output=True, offset=True)
+    output_storage.fill_(-987654)
+    result = run_checked_low_index(api, x, 512, end=end, output_idx_offset=offsets,
+                                   output_idx=out, indices_type=torch.int32,
+                                   return_value=return_value, value_oob_fill_value=12345.0)
+    assert result[1] is out
+    untouched = torch.ones_like(output_storage, dtype=torch.bool)
+    alignment = api.get_stride_requirement()[1] // torch.int32.itemsize
+    untouched[1:, alignment:alignment + 512] = False
+    assert bool((output_storage[untouched] == -987654).all())
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("width", [8193, 131072])
+def test_streaming_k512_repeated_compaction_patterns(api, device, dtype, return_value, width):
+    x, _ = aligned_tensor(api, 8, width, dtype, device)
+    position = torch.arange(width, device=device)
+    x[0].copy_(position.to(dtype))
+    x[1].copy_((width - position).to(dtype))
+    x[2].fill_(3.5)
+    x[3].fill_(-torch.finfo(dtype).max)
+    x[3, ::2] = torch.finfo(dtype).max
+    x[4].copy_((position // 1024).to(dtype))
+    x[4, -512:] = 1024
+    x[5].fill_(-5)
+    x[5, position % 1024 < 64] = 3
+    x[5, position % 1024 == 1023] = 10
+    x[6].copy_((position // 1024).to(dtype))
+    x[6, :2048] = 0
+    x[6, 2048:3072] = 1
+    x[6, 3071] = -1
+    tile_count = (width + 1023) // 1024
+    tile_bits = (tile_count - 1).bit_length()
+    scan_order_values = [-1] * width
+    full_tile_rank = 0
+    for visit in range(1 << tile_bits):
+        tile = sum(((visit >> bit) & 1) << (tile_bits - 1 - bit) for bit in range(tile_bits))
+        if tile >= tile_count or (tile + 1) * 1024 > width:
+            continue
+        value = 0 if full_tile_rank < 2 else 1 if full_tile_rank == 2 else full_tile_rank
+        start = tile * 1024
+        scan_order_values[start:start + 1024] = [value] * 1024
+        if full_tile_rank == 2:
+            scan_order_values[start + 1023] = -1
+        full_tile_rank += 1
+    x[7].copy_(torch.tensor(scan_order_values, dtype=dtype, device="cpu"))
+    offsets = torch.tensor([-100, 200, -300, 400, -500, 600, -700, 800], dtype=torch.int32, device=device)
+    run_checked_low_index(api, x, 512, indices_type=torch.int32,
+                           return_value=return_value, output_idx_offset=offsets)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_streaming_k512_special_value_ties(api, device, dtype, return_value):
+    width = 16385
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    sign = 1 << (dtype.itemsize * 8 - 1)
+    raw = ([0xBF800001, 0xBF800000, 0xBF7FFFFF] if dtype == torch.float32
+           else [0xBF81, 0xBF80, 0xBF7F])
+    patterns = [raw, [sign | 3, sign | 2, sign | 1]]
+    x, _ = aligned_tensor(api, 4, width, dtype, device)
+    position = torch.arange(width, device=device)
+    equal = position % 1024 < 64
+    greater = position % 1024 == 1023
+    for row, raw in enumerate(patterns):
+        signed = [value - 2 * sign if value & sign else value for value in raw]
+        x.view(bits)[row].fill_(signed[0])
+        x.view(bits)[row, equal] = signed[1]
+        x.view(bits)[row, greater] = signed[2]
+    x[2].fill_(-float("inf"))
+    x[2, equal] = 0.0
+    x[2, equal & (position % 2 == 0)] = -0.0
+    x[2, greater] = float("inf")
+    x[3].fill_(-float("inf"))
+    run_checked_low_index(api, x, 512, indices_type=torch.int32, return_value=return_value)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_streaming_k512_late_nan_and_excluded_suffix(api, device, dtype, return_value):
+    width = 8193
+    length = 7169
+    x, _ = aligned_tensor(api, 4, width, dtype, device)
+    x[:2].copy_(torch.arange(width, device=device).to(dtype).expand(2, -1))
+    x[2:].fill_(float("inf"))
+    x[:, length:] = float("nan")
+    x[0, length - 1] = float("nan")
+    x[2, length - 1] = float("nan")
+    end = torch.full((4,), length, dtype=torch.int32, device=device)
+    offsets = torch.tensor([100, -200, 300, -400], dtype=torch.int32, device=device)
+    values, indices = api.topk(x, 512, end=end, output_idx_offset=offsets,
+                               indices_type=torch.int32, return_value=return_value,
+                               idx_oob_fill_value=-1, abort_when_nan_found=False)
+    for row in [0, 2]:
+        assert indices[row, 0].item() == 0x3F3F3F3F
+    for row in [1, 3]:
+        result = (values[row:row + 1] if return_value else None, indices[row:row + 1])
+        check_result(api, x[row:row + 1], 512, result, indices_type=torch.int32,
+                     return_value=return_value, end=end[row:row + 1],
+                     output_idx_offset=offsets[row:row + 1])
+        run_checked_low_index(api, x[row:row + 1], 512, indices_type=torch.int32,
+                               return_value=return_value, end=end[row:row + 1],
+                               output_idx_offset=offsets[row:row + 1])
+    assert builtins_sorted((indices[3].cpu() + 400).tolist()) == list(range(512))
+    if not return_value:
+        assert values is None
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("batch,width", [(1, 65535), (1, 65536), (3, 65537),
+                                         (16, 131072), (3, 131073), (17, 65536)])
+def test_segmented_p8_dispatch_boundaries(api, device, dtype, return_value, batch, width):
+    x, _ = aligned_tensor(api, batch, width, dtype, device)
+    position = torch.arange(width, device=device)
+    for row in range(batch):
+        x[row].copy_(((position * 71 + row * 13) % 997).to(dtype))
+    offsets = torch.arange(batch, dtype=torch.int32, device=device) * 100 - 700
+    run_checked_low_index(api, x, 512, indices_type=torch.int32, return_value=return_value,
+                           output_idx_offset=offsets)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_segmented_p8_ragged_end_and_storage(api, device, dtype, return_value):
+    width = 65537
+    lengths = [0, 1, 7, 511, 512, 513, 514, 1023, 1024, 4095, 4096, 4097,
+               8193, width - 7, width - 1, width]
+    x, storage = aligned_tensor(api, len(lengths), width, dtype, device, offset=True)
+    storage.fill_(float("nan"))
+    for row, length in enumerate(lengths):
+        x[row, :length].copy_(((torch.arange(length, device=device) * 71) % 997).to(dtype))
+    before = storage.clone()
+    end = torch.tensor(lengths, dtype=torch.int32, device=device)
+    offsets = torch.arange(len(lengths), dtype=torch.int32, device=device) * 100 - 700
+    out, output_storage = aligned_tensor(api, len(lengths), 512, torch.int32, device,
+                                         output=True, offset=True)
+    output_storage.fill_(-987654)
+    result = run_checked_low_index(api, x, 512, end=end, output_idx_offset=offsets,
+                                   output_idx=out, indices_type=torch.int32,
+                                   return_value=return_value, value_oob_fill_value=12345.0)
+    assert result[1] is out
+    untouched = torch.ones_like(output_storage, dtype=torch.bool)
+    alignment = api.get_stride_requirement()[1] // torch.int32.itemsize
+    untouched[1:, alignment:alignment + 512] = False
+    assert bool((output_storage[untouched] == -987654).all())
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    assert torch.equal(storage.view(bits), before.view(bits))
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("length", [513, 4097, 65529])
+def test_segmented_p8_nan_in_each_segment(api, device, dtype, return_value, length):
+    x, _ = aligned_tensor(api, 8, 65536, dtype, device)
+    for row in range(8):
+        x[row].fill_(float("inf") if row % 2 else 1)
+    x[:, length:] = float("nan")
+    nan_indices = [(row + 1) * length // 8 - 1 for row in range(8)]
+    for row, index in enumerate(nan_indices):
+        x[row, index] = float("nan")
+    end = torch.full((8,), length, dtype=torch.int32, device=device)
+    offsets = torch.arange(8, dtype=torch.int32, device=device) * 100 - 300
+    values, indices = api.topk(x, 512, end=end, output_idx_offset=offsets,
+                               indices_type=torch.int32, return_value=return_value,
+                               abort_when_nan_found=False)
+    assert indices.dtype == torch.int32
+    assert indices[:, 0].cpu().tolist() == [0x3F3F3F3F] * 8
+    if not return_value:
+        assert values is None
+    for row, index in enumerate(nan_indices):
+        x[row, index] = float("inf") if row % 2 else 1
+    run_checked_low_index(api, x, 512, end=end, output_idx_offset=offsets,
+                           indices_type=torch.int32, return_value=return_value)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_segmented_p8_cross_segment_original_index_ties(api, device, dtype, return_value):
+    width, length = 65537, 65529
+    x, _ = aligned_tensor(api, 3, width, dtype, device)
+    x.fill_(float("nan"))
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    sign = 1 << (dtype.itemsize * 8 - 1)
+    finite = ([0xBF800001, 0xBF800000, 0xBF7FFFFF] if dtype == torch.float32
+              else [0xBF81, 0xBF80, 0xBF7F])
+    equal = [index for segment in range(8)
+             for index in range(segment * length // 8 + 17, segment * length // 8 + 113)]
+    greater = [(segment + 1) * length // 8 - 1 for segment in range(8)]
+    x[0, :length] = -float("inf")
+    x[0, equal] = 0.0
+    x[0, equal[::2]] = -0.0
+    x[0, greater] = float("inf")
+    for row, raw in enumerate([finite, [sign | 3, sign | 2, sign | 1]], start=1):
+        signed = [value - 2 * sign if value & sign else value for value in raw]
+        x.view(bits)[row, :length] = signed[0]
+        x.view(bits)[row, equal] = signed[1]
+        x.view(bits)[row, greater] = signed[2]
+    end = torch.full((3,), length, dtype=torch.int32, device=device)
+    offsets = torch.tensor([-100, 200, -300], dtype=torch.int32, device=device)
+    result = run_checked_low_index(api, x, 512, end=end, output_idx_offset=offsets,
+                                   indices_type=torch.int32, return_value=return_value)
+    expected = builtins_sorted(greater + equal[:504])
+    for row, offset in enumerate(offsets.cpu().tolist()):
+        assert builtins_sorted((result[1][row].cpu() - offset).tolist()) == expected
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_segmented_p8_graph_scratch_reuse(api, device, dtype, return_value):
+    x, _ = aligned_tensor(api, 3, 65536, dtype, device)
+    x.fill_(1)
+    end = torch.full((3,), x.shape[1], dtype=torch.int32, device=device)
+    offsets = torch.zeros(3, dtype=torch.int32, device=device)
+    kwargs = dict(end=end, output_idx_offset=offsets, indices_type=torch.int32,
+                  return_value=return_value, idx_oob_fill_value=-1, abort_when_nan_found=False)
+    stream = torch.cuda.Stream(device=device)
+    stream.wait_stream(torch.cuda.current_stream(device))
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            api.topk(x, 512, **kwargs)
+    stream.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        values, indices = api.topk(x, 512, **kwargs)
+    states = [([65536, 4097, 513], [-100, 200, -300], False),
+              ([513, 4097, 65529], [400, -500, 600], True),
+              ([0, 511, 512], [-700, 800, -900], True),
+              ([4097, 65536, 513], [1000, -1100, 1200], False)]
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    for lengths, new_offsets, have_nan in states:
+        x.fill_(float("nan"))
+        for row, length in enumerate(lengths):
+            x[row, :length] = 1
+            x[row, max(0, length - 512):length] = 10
+            if have_nan and length and row < 2:
+                x[row, length - 1] = float("nan")
+        end.copy_(torch.tensor(lengths, dtype=torch.int32, device=device))
+        offsets.copy_(torch.tensor(new_offsets, dtype=torch.int32, device=device))
+        graph.replay()
+        torch.cuda.synchronize(device)
+        for row, (length, offset) in enumerate(zip(lengths, new_offsets)):
+            if have_nan and row < 2 and length > 512:
+                assert indices[row, 0].item() == 0x3F3F3F3F
+                continue
+            count = min(length, 512)
+            actual = (indices[row, :count].cpu().to(torch.int64) - offset).tolist()
+            assert builtins_sorted(actual) == list(range(max(0, length - 512), length))
+            assert bool((indices[row, count:] == -1).all())
+            if return_value:
+                assert torch.equal(values[row, :count].view(bits), x[row, actual].view(bits))
+                assert bool(torch.isneginf(values[row, count:]).all())
+        if not return_value:
+            assert values is None
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_segmented_p8_nondefault_stream_allocation_churn(api, device, dtype, return_value):
+    x, _ = aligned_tensor(api, 3, 65536, dtype, device)
+    end = torch.empty(3, dtype=torch.int32, device=device)
+    offsets = torch.empty(3, dtype=torch.int32, device=device)
+    stream = torch.cuda.Stream(device=device)
+    stream.wait_stream(torch.cuda.current_stream(device))
+    snapshots = []
+    with torch.cuda.stream(stream):
+        torch.cuda._sleep(5_000_000)
+        for iteration, lengths in enumerate([[65536, 4097, 513], [0, 511, 512],
+                                              [513, 65529, 4096], [4097, 513, 65536]]):
+            x.fill_(float("nan"))
+            for row, length in enumerate(lengths):
+                x[row, :length] = -1
+                x[row, max(0, length - 512):length] = 10 + iteration
+            end.copy_(torch.tensor(lengths, dtype=torch.int32, device=device))
+            offsets.copy_(torch.tensor([100 + iteration, -200 - iteration, 300 + iteration],
+                                       dtype=torch.int32, device=device))
+            result = api.topk(x, 512, end=end, output_idx_offset=offsets,
+                              indices_type=torch.int32, return_value=return_value, idx_oob_fill_value=-1)
+            snapshots.append((x.clone(), end.clone(), offsets.clone(),
+                              (result[0].clone() if return_value else None, result[1].clone())))
+            del result
+            churn_keys = torch.empty((3, 8, 512), dtype=torch.int64, device=device)
+            churn_counts = torch.empty((3, 8), dtype=torch.int32, device=device)
+            churn_keys.fill_(-987654)
+            churn_counts.fill_(-1)
+            del churn_keys, churn_counts
+    stream.synchronize()
+    for source, lengths, row_offsets, result in snapshots:
+        check_result(api, source, 512, result, end=lengths, output_idx_offset=row_offsets,
+                     indices_type=torch.int32, return_value=return_value)
+        for row, (length, offset) in enumerate(zip(lengths.cpu().tolist(), row_offsets.cpu().tolist())):
+            count = min(length, 512)
+            actual = (result[1][row, :count].cpu().to(torch.int64) - offset).tolist()
+            assert builtins_sorted(actual) == list(range(max(0, length - 512), length))
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("indices_type,k", [
+    pytest.param(torch.int64, 7, id="fallback-i64"),
+    pytest.param(torch.int32, 7, id="fast-512"),
+    pytest.param(torch.int32, 513, id="fallback-k-above-512"),
+    pytest.param(torch.int32, 512, id="streaming-k512"),
+])
+def test_nondefault_stream(api, device, dtype, return_value, indices_type, k):
+    x, _ = aligned_tensor(api, 4, 16384, dtype, device)
+    x.fill_(-1)
+    stream = torch.cuda.Stream(device=device)
+    stream.wait_stream(torch.cuda.current_stream(device))
+    with torch.cuda.stream(stream):
+        torch.cuda._sleep(5_000_000)
+        x[:, -k:] = 10
+        result = api.topk(x, k, indices_type=indices_type, return_value=return_value,
+                          idx_oob_fill_value=-1)
+        copied = (result[0].clone() if return_value else None, result[1].clone())
+    stream.synchronize()
+    check_result(api, x, k, result, indices_type=indices_type, return_value=return_value)
+    assert torch.equal(copied[1], result[1])
+    if return_value:
+        assert torch.equal(copied[0], result[0])
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("sorted_index,k", [
+    pytest.param(True, 7, id="fallback-sorted-index"),
+    pytest.param(False, 7, id="fast-512"),
+    pytest.param(False, 2048, id="fallback-k-above-512"),
+    pytest.param(False, 512, id="streaming-k512"),
+])
+def test_graph_replay_updates_input_and_end(api, device, dtype, return_value, sorted_index, k):
+    x, _ = aligned_tensor(api, 4, 16384, dtype, device)
+    x.fill_(0)
+    end = torch.full((4,), x.shape[1], dtype=torch.int32, device=device)
+    offsets = torch.tensor([100, 200, 300, 400], dtype=torch.int32, device=device)
+    kwargs = dict(end=end, output_idx_offset=offsets, sorted_index=sorted_index,
+                  return_value=return_value, indices_type=torch.int32, idx_oob_fill_value=-1)
+    stream = torch.cuda.Stream(device=device)
+    stream.wait_stream(torch.cuda.current_stream(device))
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            api.topk(x, k, **kwargs)
+    stream.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        result = api.topk(x, k, **kwargs)
+    for ends, new_offsets, hot in [
+        ([16384, 0, k - 1, k], [-100, 20, -30, 40], 100),
+        ([k, 16384, k + 1, k - 1], [500, -600, 700, -800], 1000),
+        ([k + 1, k - 1, 0, 16384], [-90, 80, -70, 60], 4000),
+    ]:
+        x.fill_(-5)
+        x[:, hot:hot + k] = 10
+        end.copy_(torch.tensor(ends, dtype=torch.int32, device=device))
+        offsets.copy_(torch.tensor(new_offsets, dtype=torch.int32, device=device))
+        graph.replay()
+        torch.cuda.synchronize(device)
+        check_result(api, x, k, result, **kwargs)
+
+
+def test_input_device_guard(api, device):
+    if torch.cuda.device_count() < 2:
+        pytest.skip("two visible CUDA devices are required")
+    other = next(index for index in range(torch.cuda.device_count()) if index != device.index)
+    x, _ = aligned_tensor(api, 4, 16384, torch.float32, device)
+    x.copy_(torch.arange(x.shape[1], device=device).expand_as(x))
+    torch.cuda.synchronize(device)
+    with torch.cuda.device(other):
+        result = api.topk(x, 7, sorted=True, idx_oob_fill_value=-1)
+        assert torch.cuda.current_device() == other
+    check_result(api, x, 7, result, sorted=True)
+    if torch.cuda.get_device_capability(other) in ((12, 0), (12, 1)):
+        target = torch.device("cuda", other)
+        x_other, _ = aligned_tensor(api, 4, 16384, torch.bfloat16, target)
+        x_other.fill_(0)
+        x_other[:, -7:] = 10
+        torch.cuda.synchronize(target)
+        with torch.cuda.device(device):
+            other_result = api.topk(x_other, 7, idx_oob_fill_value=-1)
+            assert torch.cuda.current_device() == device.index
+        check_result(api, x_other, 7, other_result)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("batch,width,k", [
+    (16, 2048, 511), (16, 2048, 512), (17, 2048, 513),
+    (17, 2049, 511), (16, 2049, 512), (17, 2049, 512), (16, 2049, 513),
+    (16, 65535, 512), (17, 65535, 512),
+    (16, 65536, 511), (16, 65536, 512), (17, 65536, 512), (16, 65536, 513),
+    (16, 65537, 512), (17, 65537, 512),
+    (16, 131072, 511), (16, 131072, 512), (17, 131072, 512), (16, 131072, 513),
+    (16, 131073, 512), (17, 131073, 512),
+])
+def test_dispatch_k_batch_width_boundaries(api, device, dtype, return_value, batch, width, k):
+    x, _ = aligned_tensor(api, batch, width, dtype, device)
+    position = torch.arange(width, device=device)
+    for row in range(batch):
+        x[row].copy_(((position * 71 + row * 13) % 997).to(dtype))
+    run_checked_low_index(api, x, k, indices_type=torch.int32, return_value=return_value)
+
+
+INVALID_CASES = [
+    "input-fp16", "input-fp64", "input-i32", "input-rank1", "input-rank3", "input-cpu",
+    "input-last-stride", "input-row-alignment", "input-pointer-alignment", "input-storage",
+    "k-zero", "k-negative", "k-too-large", "indices-dtype", "output-dtype", "output-shape",
+    "output-last-stride", "output-row-alignment", "output-pointer-alignment", "output-cpu", "output-overlap",
+    "end-dtype", "end-shape", "end-stride", "end-cpu",
+    "offset-dtype", "offset-shape", "offset-stride", "offset-cpu",
+    "sorted-bf16", "both-sorts", "sorted-no-values", "begin", "hint",
+]
+
+
+@pytest.mark.parametrize("case", INVALID_CASES)
+def test_invalid_arguments(api, device, case):
+    if api is NativeAPI and case in ("hint", "output-dtype"):
+        pytest.skip("wrapper-only arguments are absent from the registered schema")
+    x, _ = aligned_tensor(api, 4, 16384, torch.float32, device)
+    x.fill_(0)
+    k = 7
+    kwargs = {}
+    if case.startswith("input-"):
+        kind = case.removeprefix("input-")
+        if kind in ("fp16", "fp64", "i32"):
+            x = x.to({"fp16": torch.float16, "fp64": torch.float64, "i32": torch.int32}[kind])
+        elif kind == "rank1":
+            x = x[0]
+        elif kind == "rank3":
+            x = x.unsqueeze(0)
+        elif kind == "cpu":
+            x = x.cpu()
+        elif kind == "last-stride":
+            x = x[:, ::2]
+        elif kind == "row-alignment":
+            x = torch.zeros((4, 16385), device=device)[:, :16384]
+        elif kind == "pointer-alignment":
+            x = x[:, 1:]
+        elif kind == "storage":
+            storage = torch.zeros((7,), dtype=torch.float32, device=device)
+            row_stride = api.get_stride_requirement()[0] // storage.element_size()
+            x = storage.as_strided((1, 7), (row_stride, 1))
+    elif case.startswith("k-"):
+        k = {"k-zero": 0, "k-negative": -1, "k-too-large": 4097}[case]
+    elif case == "indices-dtype":
+        kwargs["indices_type"] = torch.int16
+    elif case.startswith("output-"):
+        out, _ = aligned_tensor(api, 4, k, torch.int64, device, output=True)
+        kind = case.removeprefix("output-")
+        if kind == "dtype":
+            out = out.to(torch.int32)
+        elif kind == "shape":
+            out = out[:, :k - 1]
+        elif kind == "last-stride":
+            out, _ = aligned_tensor(api, 4, 2 * k, torch.int64, device, output=True)
+            out = out[:, ::2]
+        elif kind == "row-alignment":
+            out = torch.empty((4, k), dtype=torch.int64, device=device)
+        elif kind == "pointer-alignment":
+            out, _ = aligned_tensor(api, 4, k + 1, torch.int64, device, output=True)
+            out = out[:, 1:]
+        elif kind == "cpu":
+            out = out.cpu()
+        elif kind == "overlap":
+            out = out[:1].expand(4, -1)
+        kwargs["output_idx"] = out
+    elif case.startswith(("end-", "offset-")):
+        name, kind = case.split("-")
+        argument = torch.full((4,), 10, dtype=torch.int32, device=device)
+        if kind == "dtype":
+            argument = argument.to(torch.int64)
+        elif kind == "shape":
+            argument = argument[:3]
+        elif kind == "stride":
+            argument = torch.full((8,), 10, dtype=torch.int32, device=device)[::2]
+        elif kind == "cpu":
+            argument = argument.cpu()
+        kwargs["end" if name == "end" else "output_idx_offset"] = argument
+    elif case == "sorted-bf16":
+        x = x.to(torch.bfloat16)
+        kwargs["sorted"] = True
+    elif case == "both-sorts":
+        kwargs.update(sorted=True, sorted_index=True)
+    elif case == "sorted-no-values":
+        kwargs.update(sorted=True, return_value=False)
+    else:
+        kwargs[case] = torch.zeros((4,), dtype=torch.int32, device=device)
+    with pytest.raises((RuntimeError, AssertionError, ValueError, IndexError)):
+        api.topk(x, k, **kwargs)
+    torch.cuda.synchronize(device)
+
+
+@pytest.mark.parametrize("argument", ["end", "output_idx", "output_idx_offset"])
+def test_reject_cross_device_arguments(api, device, argument):
+    if torch.cuda.device_count() < 2:
+        pytest.skip("two visible CUDA devices are required")
+    other = next(index for index in range(torch.cuda.device_count()) if index != device.index)
+    x, _ = aligned_tensor(api, 4, 16384, torch.float32, device)
+    x.fill_(0)
+    if argument == "output_idx":
+        value, _ = aligned_tensor(api, 4, 7, torch.int64, torch.device("cuda", other), output=True)
+    else:
+        value = torch.full((4,), 7, dtype=torch.int32, device=torch.device("cuda", other))
+    with pytest.raises(RuntimeError, match="device"):
+        api.topk(x, 7, **{argument: value})
+    torch.cuda.synchronize(device)
+
+
+def shared_storage_view(storage, dtype, width, byte_offset, alignment):
+    row_bytes = (width * dtype.itemsize + alignment - 1) // alignment * alignment
+    return storage[byte_offset:byte_offset + row_bytes].view(dtype).reshape(1, -1)[:, :width]
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("argument", ["input", "end", "output_idx_offset"])
+@pytest.mark.parametrize("disjoint", [False, True], ids=["overlap", "disjoint-shared-storage"])
+@pytest.mark.parametrize("return_value", [False, True])
+def test_reject_output_idx_alias(api, device, dtype, indices_type, argument, disjoint, return_value):
+    storage = torch.zeros(8192, dtype=torch.uint8, device=device)
+    x, _ = aligned_tensor(api, 1, 128, dtype, device)
+    x.fill_(0)
+    output = shared_storage_view(storage, indices_type, 7, 1024 if disjoint else 0,
+                                 api.get_stride_requirement()[1])
+    kwargs = dict(output_idx=output, indices_type=indices_type, return_value=return_value)
+    if argument == "input":
+        x = shared_storage_view(storage, dtype, 128, 0, api.get_stride_requirement()[0])
+    else:
+        metadata = storage[:4].view(torch.int32)
+        metadata.fill_(128 if argument == "end" else 100)
+        kwargs[argument] = metadata
+    before = storage.clone()
+    with pytest.raises(RuntimeError, match="overlap.*storage|storage.*overlap"):
+        api.topk(x, 7, **kwargs)
+    torch.cuda.synchronize(device)
+    assert torch.equal(storage, before)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("argument", ["input", "end", "output_idx_offset", "output_idx"])
+@pytest.mark.parametrize("disjoint", [False, True], ids=["overlap", "disjoint-shared-storage"])
+def test_reject_backend_output_value_alias(api, device, dtype, indices_type, argument, disjoint):
+    storage = torch.zeros(8192, dtype=torch.uint8, device=device)
+    x, _ = aligned_tensor(api, 1, 128, dtype, device)
+    x.fill_(0)
+    indices, _ = aligned_tensor(api, 1, 7, indices_type, device, output=True)
+    values = shared_storage_view(storage, dtype, 7, 1024 if disjoint else 0,
+                                 api.get_stride_requirement()[1])
+    end = offsets = None
+    if argument == "input":
+        x = shared_storage_view(storage, dtype, 128, 0, api.get_stride_requirement()[0])
+    elif argument == "output_idx":
+        indices = shared_storage_view(storage, indices_type, 7, 0, api.get_stride_requirement()[1])
+    elif argument == "end":
+        end = storage[:4].view(torch.int32)
+        end.fill_(128)
+    else:
+        offsets = storage[:4].view(torch.int32)
+        offsets.fill_(100)
+    before = storage.clone()
+    with pytest.raises(RuntimeError, match="overlap.*storage|storage.*overlap"):
+        torch.ops.deep_select.topk(x, 7, None, end, False, False, values, indices,
+                                   offsets, -1, float("-inf"), True, True)
+    torch.cuda.synchronize(device)
+    assert torch.equal(storage, before)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("sorted,sorted_index,return_value", MODES)
+@pytest.mark.parametrize("k", [7, 512])
+@pytest.mark.parametrize("pattern", ["equal", "threshold"])
+def test_low_index_tie_selection(api, device, dtype, indices_type, sorted, sorted_index, return_value, k, pattern):
+    if dtype == torch.bfloat16 and sorted:
+        pytest.skip("sorted values are FP32-only")
+    x, _ = aligned_tensor(api, 1, 16384, dtype, device)
+    x.fill_(3)
+    expected = list(range(k))
+    if pattern == "threshold":
+        x[:, -3:] = 10
+        expected = list(range(k - 3)) + list(range(x.shape[1] - 3, x.shape[1]))
+    for _ in range(3):
+        _, indices = run_checked(api, x, k, indices_type=indices_type, sorted=sorted,
+                                 sorted_index=sorted_index, return_value=return_value)
+        actual = indices[0].cpu().tolist()
+        assert builtins_sorted(actual) == expected
+        if sorted:
+            assert actual == (expected[-3:] + expected[:-3] if pattern == "threshold" else expected)
+        elif sorted_index:
+            assert actual == expected
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("sorted,sorted_index,return_value", MODES)
+@pytest.mark.parametrize("k", [1, 7, 12])
+def test_subnormal_and_signed_zero_order(api, device, dtype, indices_type, sorted, sorted_index, return_value, k):
+    if dtype == torch.bfloat16 and sorted:
+        pytest.skip("sorted values are FP32-only")
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    bit_count, mantissa = (32, 23) if dtype == torch.float32 else (16, 7)
+    sign = 1 << (bit_count - 1)
+    tiny = 1 << mantissa
+    raw = [sign | 1, 0, sign, 1, tiny - 1, sign | (tiny - 1),
+           2, sign | 2, tiny, sign | tiny, 0, sign]
+    signed = [value - 2 * sign if value & sign else value for value in raw]
+    x, _ = aligned_tensor(api, 1, 16384, dtype, device)
+    x.fill_(-float("inf"))
+    x.view(bits)[0, :len(raw)].copy_(torch.tensor(signed, dtype=bits, device=device))
+    def key(index):
+        value = raw[index]
+        if value in (0, sign):
+            value = 0
+        return ((~value & (2 * sign - 1)) if value & sign else value ^ sign)
+    ranked = builtins_sorted(range(len(raw)), key=lambda index: (-key(index), index))[:k]
+    values, indices = api.topk(x, k, sorted=sorted, sorted_index=sorted_index,
+                               return_value=return_value, indices_type=indices_type)
+    actual = indices[0].cpu().tolist()
+    assert builtins_sorted(actual) == builtins_sorted(ranked)
+    if sorted:
+        assert actual == ranked
+    elif sorted_index:
+        assert actual == builtins_sorted(ranked)
+    if return_value:
+        assert values[0].view(bits).cpu().tolist() == [signed[index] for index in actual]
+    else:
+        assert values is None
+
+
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("length", [5, 7])
+@pytest.mark.parametrize("fill", [-float("inf"), float("inf"), 12345.0], ids=["fill-neginf", "fill-posinf", "fill-finite"])
+@pytest.mark.parametrize("abort_when_nan_found", [False, True])
+def test_sorted_short_row_nan_payload_and_fill(api, device, indices_type, length, fill, abort_when_nan_found):
+    x, _ = aligned_tensor(api, 1, 128, torch.float32, device)
+    x.fill_(0)
+    unsigned = [0xffc01234, 0x3f800000, 0x7fc05678, 0xff800000,
+                0x80000000, 0x00000000, 0xff801234]
+    signed = [value if value < 0x80000000 else value - 0x100000000 for value in unsigned]
+    x.view(torch.int32)[0, :7].copy_(torch.tensor(signed, dtype=torch.int32, device=device))
+    end = torch.tensor([length], dtype=torch.int32, device=device)
+    offsets = torch.tensor([100], dtype=torch.int32, device=device)
+    values, indices = api.topk(x, 7, sorted=True, end=end, output_idx_offset=offsets,
+                               indices_type=indices_type, idx_oob_fill_value=-1,
+                               value_oob_fill_value=fill, abort_when_nan_found=abort_when_nan_found)
+    expected = [0, 2, 1, 4, 3] if length == 5 else [0, 2, 6, 1, 4, 5, 3]
+    assert indices[0, :length].cpu().tolist() == [index + 100 for index in expected]
+    assert indices[0, length:].cpu().tolist() == [-1] * (7 - length)
+    assert values[0, :length].view(torch.int32).cpu().tolist() == [signed[index] for index in expected]
+    fill_bits = torch.tensor(fill, dtype=torch.float32).view(torch.int32).item()
+    assert values[0, length:].view(torch.int32).cpu().tolist() == [fill_bits] * (7 - length)
+
+
+def assert_trap_child(api, device, dtype_name, indices_name, k, return_value, case):
+    child = subprocess.run([sys.executable, str(Path(__file__).resolve()), "--trap-child",
+                            dtype_name, str(device.index), indices_name, str(k),
+                            str(int(return_value)), case, "native" if api is NativeAPI else "wrapper"],
+                           capture_output=True, text=True, timeout=120)
+    output = child.stdout + child.stderr
+    assert "TRAP_ARMED" in output, output
+    assert child.returncode != 0, output
+    assert any(text in output.lower() for text in
+               ["illegal instruction", "device-side assert", "unspecified launch failure",
+                "cudaerrorillegalinstruction"]), output
+
+
+@pytest.mark.skipif(os.environ.get("DEEP_SELECT_TEST_NAN_TRAP") != "1",
+                    reason="opt-in only: traps a child CUDA context; set DEEP_SELECT_TEST_NAN_TRAP=1")
+@pytest.mark.parametrize("dtype_name", ["float32", "bfloat16"])
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("indices_name,k", [
+    pytest.param("int64", 7, id="fallback-i64"),
+    pytest.param("int32", 7, id="fast-512"),
+    pytest.param("int32", 2048, id="fallback-k-above-512"),
+    pytest.param("int32", 512, id="streaming-k512"),
+])
+def test_nan_default_traps_in_child(api, device, dtype_name, return_value, indices_name, k):
+    assert_trap_child(api, device, dtype_name, indices_name, k, return_value, "nan")
+
+
+@pytest.mark.skipif(os.environ.get("DEEP_SELECT_TEST_NAN_TRAP") != "1",
+                    reason="opt-in only: traps a child CUDA context; set DEEP_SELECT_TEST_NAN_TRAP=1")
+@pytest.mark.parametrize("dtype_name", ["float32", "bfloat16"])
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("case", ["end-negative", "end-too-large"])
+@pytest.mark.parametrize("k,prefix", [
+    pytest.param(513, "", id="fallback-k-above-512"),
+    pytest.param(7, "whole-row-", id="whole-row-k7"),
+    pytest.param(512, "streaming-", id="streaming-k512"),
+])
+def test_unsorted_i32_invalid_end_traps_in_child(api, device, dtype_name, return_value, case, k, prefix):
+    assert_trap_child(api, device, dtype_name, "int32", k, return_value, prefix + case)
+
+
+@pytest.mark.skipif(os.environ.get("DEEP_SELECT_TEST_NAN_TRAP") != "1",
+                    reason="opt-in only: traps a child CUDA context; set DEEP_SELECT_TEST_NAN_TRAP=1")
+@pytest.mark.parametrize("dtype_name", ["float32", "bfloat16"])
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("case", ["nan", "end-negative", "end-too-large"])
+def test_segmented_p8_traps_in_child(api, device, dtype_name, return_value, case):
+    assert_trap_child(api, device, dtype_name, "int32", 512, return_value, "segmented-p8-" + case)
+
+
+def trap_child(dtype_name, device_index, indices_name, k, return_value, case, caller):
+    import resource
+    import deep_select
+
+    api = NativeAPI if caller == "native" else deep_select
+    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
+    torch.cuda.set_device(int(device_index))
+    device = torch.device("cuda", int(device_index))
+    assert torch.cuda.get_device_capability(device) in ((12, 0), (12, 1))
+    k = int(k)
+    width = 8193 if k == 512 else k + 2
+    length = 7169 if k == 512 else k + 1
+    if case.startswith("segmented-p8-"):
+        width, length = 65536, 513
+        case = case.removeprefix("segmented-p8-")
+    elif case.startswith("whole-row-"):
+        width = length = 4096
+        case = case.removeprefix("whole-row-")
+    elif case.startswith("streaming-"):
+        width = length = 16384
+        case = case.removeprefix("streaming-")
+    x, _ = aligned_tensor(api, 1, width, getattr(torch, dtype_name), device)
+    x.fill_(float("inf") if k == 512 else 0)
+    end = torch.tensor([length], dtype=torch.int32, device=device)
+    kwargs = dict(end=end, indices_type=getattr(torch, indices_name), return_value=bool(int(return_value)))
+    if case == "nan":
+        x[0, length:] = float("nan")
+    api.topk(x, k, **kwargs)
+    torch.cuda.synchronize(device)
+    if case == "nan":
+        x[0, length - 1] = float("nan")
+    elif case == "end-negative":
+        end.fill_(-1)
+    elif case == "end-too-large":
+        end.fill_(x.shape[1] + 1)
+    else:
+        raise ValueError(case)
+    torch.cuda.synchronize(device)
+    print("TRAP_ARMED", flush=True)
+    api.topk(x, k, **kwargs)
+    torch.cuda.synchronize(device)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] == "--trap-child":
+        trap_child(*sys.argv[2:])
+    else:
+        raise SystemExit(pytest.main([str(Path(__file__).resolve()), *sys.argv[1:]]))
diff --git a/tests/test_stable_contracts.py b/tests/test_stable_contracts.py
new file mode 100644
index 0000000..0a0704f
--- /dev/null
+++ b/tests/test_stable_contracts.py
@@ -0,0 +1,289 @@
+import importlib
+
+import pytest
+import torch
+
+from .test_sm120 import DTYPES, INDEX_DTYPES, NativeAPI, aligned_tensor, check_result, shared_storage_view
+
+
+@pytest.fixture(scope="module")
+def device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    device = torch.device("cuda", torch.cuda.current_device())
+    if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3), (12, 0), (12, 1)):
+        pytest.skip("DeepSelect requires SM100, SM103, SM120, or SM121")
+    return device
+
+
+@pytest.fixture(scope="module")
+def registered():
+    import deep_select
+
+    return deep_select
+
+
+@pytest.fixture(scope="module", params=["wrapper", "native"])
+def api(device, registered, request):
+    return registered if request.param == "wrapper" else NativeAPI
+
+
+def test_registered_ops_once(registered):
+    module = importlib.import_module("deep_select.deep_select_cuda")
+    assert module.__name__ == "deep_select.deep_select_cuda"
+    topk = torch.ops.deep_select.topk.default
+    alignment = torch.ops.deep_select.get_alignment_requirement.default
+    expected = alignment()
+    assert tuple(registered.get_stride_requirement()) == tuple(expected)
+    assert len(expected) == 2 and all(value > 0 and value % 32 == 0 for value in expected)
+    for _ in range(3):
+        assert importlib.import_module("deep_select.deep_select_cuda") is module
+        assert importlib.import_module("deep_select") is registered
+        assert torch.ops.deep_select.topk.default is topk
+        assert torch.ops.deep_select.get_alignment_requirement.default is alignment
+    assert torch.ops.deep_select.topk.overloads() == ["default"]
+    assert torch.ops.deep_select.get_alignment_requirement.overloads() == ["default"]
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("batch,width,k", [(0, 0, 7), (0, 65536, 512), (3, 0, 7), (3, 16384, 513)])
+def test_shared_empty_and_varlen(api, device, dtype, indices_type, return_value, batch, width, k):
+    x, _ = aligned_tensor(api, batch, width, dtype, device)
+    x.copy_(((torch.arange(width, device=device) * 71) % 997).to(dtype).expand_as(x))
+    end = torch.tensor(([0, min(k - 1, width), width] if batch else []),
+                       dtype=torch.int32, device=device)
+    offsets = torch.tensor(([-100, 200, -300] if batch else []), dtype=torch.int32, device=device)
+    if batch > 0 and width == 0 and torch.cuda.get_device_capability(device)[0] == 10:
+        with pytest.raises(RuntimeError, match="SM100/SM103 require vocab_size > 0 for nonempty batches"):
+            api.topk(x, k, end=end, output_idx_offset=offsets, indices_type=indices_type,
+                     return_value=return_value, idx_oob_fill_value=-1, value_oob_fill_value=12345.0)
+        return
+    result = api.topk(x, k, end=end, output_idx_offset=offsets, indices_type=indices_type,
+                      return_value=return_value, idx_oob_fill_value=-1, value_oob_fill_value=12345.0)
+    check_result(api, x, k, result, end=end, output_idx_offset=offsets, indices_type=indices_type,
+                 return_value=return_value, idx_oob_fill_value=-1, value_oob_fill_value=12345.0)
+    torch.cuda.synchronize(device)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("case", [
+    "idx-fill-low", "idx-fill-high", "value-fill-low", "value-fill-high",
+    "input-rank0", "input-rank1", "input-rank3", "input-dtype", "input-cpu",
+    "input-pointer", "input-stride", "input-storage", "width-limit", "batch-limit",
+    "k-zero", "k-negative", "k-limit", "k-int64", "begin",
+    "end-dtype", "end-shape", "end-stride", "end-cpu",
+    "offset-dtype", "offset-shape", "offset-stride", "offset-cpu",
+    "index-dtype", "index-shape", "index-pointer", "index-row-stride", "index-last-stride", "index-overlap", "index-cpu",
+    "value-missing", "value-dtype", "value-shape", "value-pointer", "value-row-stride", "value-last-stride", "value-overlap", "value-cpu",
+    "both-sorts", "sorted-no-values",
+])
+def test_native_rejects_invalid_contract(registered, device, dtype, indices_type, case):
+    x, _ = aligned_tensor(NativeAPI, 3, 128, dtype, device)
+    x.fill_(0)
+    values, _ = aligned_tensor(NativeAPI, 3, 7, dtype, device, output=True)
+    indices, _ = aligned_tensor(NativeAPI, 3, 7, indices_type, device, output=True)
+    args = dict(input=x, topk=7, begin=None, end=None, sorted_value=False, sorted_index=False,
+                output_value=values, output_index=indices, output_idx_offset=None,
+                idx_oob_fill_value=-1, value_oob_fill_value=float("-inf"),
+                return_value=True, abort_when_nan_found=True)
+    if case.startswith("idx-fill-"):
+        args["idx_oob_fill_value"] = -(1 << 31) - 1 if case.endswith("low") else 1 << 31
+    elif case.startswith("value-fill-"):
+        args["value_oob_fill_value"] = -1e300 if case.endswith("low") else 1e300
+    elif case.startswith("input-"):
+        kind = case.removeprefix("input-")
+        if kind == "rank0":
+            x = x[0, 0]
+        elif kind == "rank1":
+            x = x[0]
+        elif kind == "rank3":
+            x = x.unsqueeze(0)
+        elif kind == "dtype":
+            x = x.to(torch.float16)
+        elif kind == "cpu":
+            x = x.cpu()
+        elif kind == "pointer":
+            x = x[:, 1:]
+        elif kind == "stride":
+            x = x[:, ::2]
+        else:
+            storage = torch.empty(7, dtype=dtype, device=device)
+            x = storage.as_strided((1, 7), (128 // dtype.itemsize, 1))
+            args["output_value"] = values[:1]
+            args["output_index"] = indices[:1]
+        args["input"] = x
+    elif case in ("width-limit", "batch-limit"):
+        if case == "width-limit":
+            args["input"] = x.as_strided((0, 1 << 23), (1 << 23, 1))
+            args["output_value"] = values[:0]
+            args["output_index"] = indices[:0]
+        else:
+            args["input"] = x.as_strided((1 << 31, 0), (0, 1))
+            args["output_value"] = values[:1].expand(1 << 31, -1)
+            args["output_index"] = indices[:1].expand(1 << 31, -1)
+    elif case.startswith("k-"):
+        args["topk"] = {"k-zero": 0, "k-negative": -1, "k-limit": 4097, "k-int64": 1 << 32}[case]
+    elif case == "begin":
+        args["begin"] = torch.zeros(3, dtype=torch.int32, device=device)
+    elif case.startswith(("end-", "offset-")):
+        name, kind = case.split("-")
+        metadata = torch.full((3,), 7, dtype=torch.int32, device=device)
+        if kind == "dtype":
+            metadata = metadata.to(torch.int64)
+        elif kind == "shape":
+            metadata = metadata[:2]
+        elif kind == "stride":
+            metadata = torch.full((6,), 7, dtype=torch.int32, device=device)[::2]
+        else:
+            metadata = metadata.cpu()
+        args["end" if name == "end" else "output_idx_offset"] = metadata
+    elif case.startswith(("index-", "value-")):
+        name, kind = case.split("-", 1)
+        output = indices if name == "index" else values
+        if kind == "missing":
+            output = None
+        elif kind == "dtype":
+            output = output.to(torch.int16 if name == "index" else torch.float64)
+        elif kind == "shape":
+            output = output[:, :6]
+        elif kind == "pointer":
+            output, _ = aligned_tensor(NativeAPI, 3, 8, output.dtype, device, output=True)
+            output = output[:, 1:]
+        elif kind == "row-stride":
+            output = torch.empty((3, 7), dtype=output.dtype, device=device)
+        elif kind == "last-stride":
+            output, _ = aligned_tensor(NativeAPI, 3, 14, output.dtype, device, output=True)
+            output = output[:, ::2]
+        elif kind == "overlap":
+            output = output[:1].expand(3, -1)
+        else:
+            output = output.cpu()
+        args["output_index" if name == "index" else "output_value"] = output
+    elif case == "both-sorts":
+        args.update(sorted_value=True, sorted_index=True)
+    else:
+        args.update(sorted_value=True, return_value=False, output_value=None)
+    with pytest.raises(RuntimeError):
+        torch.ops.deep_select.topk(**args)
+    torch.cuda.synchronize(device)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+def test_native_output_padding_and_bounded_tail(registered, device, dtype, indices_type, return_value):
+    x, input_storage = aligned_tensor(NativeAPI, 3, 129, dtype, device, offset=True)
+    input_storage.fill_(float("nan"))
+    x.copy_(torch.arange(129, device=device).to(dtype).expand_as(x))
+    before = input_storage.clone()
+    indices, index_storage = aligned_tensor(NativeAPI, 3, 7, indices_type, device, output=True, offset=True)
+    values, value_storage = aligned_tensor(NativeAPI, 3, 7, dtype, device, output=True, offset=True)
+    index_storage.fill_(-987654)
+    value_storage.fill_(12345)
+    end = torch.tensor([0, 5, 129], dtype=torch.int32, device=device)
+    offsets = torch.tensor([-100, 200, -300], dtype=torch.int32, device=device)
+    result = torch.ops.deep_select.topk(x, 7, None, end, False, False,
+                                        values if return_value else None, indices, offsets,
+                                        -1, float("-inf"), return_value, True)
+    assert result is None
+    check_result(NativeAPI, x, 7, (values if return_value else None, indices),
+                 indices_type=indices_type, return_value=return_value, end=end, output_idx_offset=offsets)
+    for storage, output, fill in [(index_storage, indices, -987654), (value_storage, values, 12345)]:
+        untouched = torch.ones_like(storage, dtype=torch.bool)
+        if output is indices or return_value:
+            alignment = NativeAPI.get_stride_requirement()[1] // output.element_size()
+            untouched[1:, alignment:alignment + 7] = False
+        assert bool((storage[untouched] == fill).all())
+    bits = torch.int32 if dtype == torch.float32 else torch.int16
+    assert torch.equal(input_storage.view(bits), before.view(bits))
+    for output_dtype in [dtype, indices_type]:
+        storage = torch.empty(7, dtype=output_dtype, device=device)
+        tail = storage.as_strided((1, 7), (32 // output_dtype.itemsize, 1))
+        out_values = tail if output_dtype == dtype else values[:1]
+        out_indices = tail if output_dtype == indices_type else indices[:1]
+        assert torch.ops.deep_select.topk(x[:1], 7, None, None, False, False, out_values,
+                                          out_indices, None, -1, float("-inf"), True, True) is None
+        check_result(NativeAPI, x[:1], 7, (out_values, out_indices), indices_type=indices_type)
+    torch.cuda.synchronize(device)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("output_name,argument", [
+    ("output_index", "input"), ("output_index", "end"), ("output_index", "output_idx_offset"),
+    ("output_value", "input"), ("output_value", "end"),
+    ("output_value", "output_idx_offset"), ("output_value", "output_index"),
+])
+@pytest.mark.parametrize("disjoint", [False, True], ids=["overlap", "disjoint-shared-storage"])
+@pytest.mark.parametrize("return_value", [False, True])
+def test_native_output_storage_alias(registered, device, dtype, indices_type, output_name, argument, disjoint, return_value):
+    storage = torch.zeros(8192, dtype=torch.uint8, device=device)
+    x, _ = aligned_tensor(NativeAPI, 1, 128, dtype, device)
+    values, _ = aligned_tensor(NativeAPI, 1, 7, dtype, device, output=True)
+    indices, _ = aligned_tensor(NativeAPI, 1, 7, indices_type, device, output=True)
+    args = dict(input=x, topk=7, begin=None, end=None, sorted_value=False, sorted_index=False,
+                output_value=values, output_index=indices, output_idx_offset=None,
+                idx_oob_fill_value=-1, value_oob_fill_value=float("-inf"),
+                return_value=return_value, abort_when_nan_found=True)
+    args[output_name] = shared_storage_view(storage, args[output_name].dtype, 7,
+                                            1024 if disjoint else 0, NativeAPI.get_stride_requirement()[1])
+    if argument == "input":
+        args[argument] = shared_storage_view(storage, dtype, 128, 0, NativeAPI.get_stride_requirement()[0])
+    elif argument == "output_index":
+        args[argument] = shared_storage_view(storage, indices_type, 7, 0, NativeAPI.get_stride_requirement()[1])
+    else:
+        args[argument] = storage[:4].view(torch.int32)
+        args[argument].fill_(128 if argument == "end" else 100)
+    before = storage.clone()
+    with pytest.raises(RuntimeError, match="overlap.*storage|storage.*overlap"):
+        torch.ops.deep_select.topk(**args)
+    torch.cuda.synchronize(device)
+    assert torch.equal(storage, before)
+
+
+@pytest.mark.parametrize("argument", ["output_value", "output_index", "end", "output_idx_offset"])
+def test_native_cross_device(registered, device, argument):
+    if torch.cuda.device_count() < 2:
+        pytest.skip("two visible CUDA devices are required")
+    other = next(index for index in range(torch.cuda.device_count()) if index != device.index)
+    x, _ = aligned_tensor(NativeAPI, 1, 128, torch.float32, device)
+    values, _ = aligned_tensor(NativeAPI, 1, 7, torch.float32, device, output=True)
+    indices, _ = aligned_tensor(NativeAPI, 1, 7, torch.int64, device, output=True)
+    args = dict(input=x, topk=7, begin=None, end=None, sorted_value=False, sorted_index=False,
+                output_value=values, output_index=indices, output_idx_offset=None,
+                idx_oob_fill_value=-1, value_oob_fill_value=float("-inf"),
+                return_value=True, abort_when_nan_found=True)
+    if argument.startswith("output_") and argument != "output_idx_offset":
+        args[argument], _ = aligned_tensor(NativeAPI, 1, 7, args[argument].dtype,
+                                           torch.device("cuda", other), output=True)
+    else:
+        args[argument] = torch.full((1,), 7, dtype=torch.int32, device=torch.device("cuda", other))
+    with pytest.raises(RuntimeError, match="device"):
+        torch.ops.deep_select.topk(**args)
+    torch.cuda.synchronize(device)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("indices_type", INDEX_DTYPES)
+@pytest.mark.parametrize("return_value", [False, True])
+@pytest.mark.parametrize("width", [16384, 524288])
+def test_sm100_sm103_optimized_nan_nonabort(api, device, dtype, indices_type, return_value, width):
+    if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3)):
+        pytest.skip("optimized SM100/SM103 NaN paths")
+    x, _ = aligned_tensor(api, 2, width, dtype, device)
+    x.fill_(1)
+    x[0, width - 2] = float("nan")
+    x[1, width - 1] = float("nan")
+    end = torch.full((2,), width - 1, dtype=torch.int32, device=device)
+    offsets = torch.tensor([100, -200], dtype=torch.int32, device=device)
+    values, indices = api.topk(x, 512, end=end, output_idx_offset=offsets,
+                               indices_type=indices_type, return_value=return_value,
+                               idx_oob_fill_value=-1, abort_when_nan_found=False)
+    assert indices[0, 0].item() == 0x3F3F3F3F
+    check_result(api, x[1:], 512, (values[1:] if return_value else None, indices[1:]),
+                 end=end[1:], output_idx_offset=offsets[1:], indices_type=indices_type,
+                 return_value=return_value)
+    torch.cuda.synchronize(device)