From 89e9e64623aa1dc4012d08f9e75c28e8277aa790 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:43:17 -0700 Subject: [PATCH 1/4] feat(cuda): add native SM120 and SM121 stable ABI TopK Compile the host API once and isolate SM100-family and native SM12 CUDA source groups before linking one ABI3 DSO with one operator registration. Keep the PyTorch 2.10 stable ABI, default SM100a/SM103a targets, explicit SM120/SM121 selection, optional PTX, and mandatory artifact spill checks. Add exact native radix selection, streaming K512 selection, and segmented selection with original-index tie breaking and current-stream scratch. Validate device, storage, padding, output strides, aliases, and scalar bounds for both the wrapper and direct operator. Preserve NaN behavior and bound output stores to logical elements. Disable CUB outer-scan memoization to remove generic-sort spills without changing precision or ordering. This changes register/shared-memory trade-offs; no performance speedup or zero-regression claim is made. Carry the cluster gather completion-barrier address repair from https://github.com/deepseek-ai/DeepSelect/pull/16, mapping destinations and barriers to CTA0. That upstream PR remains independent; native SM12 does not use the cluster path. Validation of the final full mixed wheel: - 348 kernel instances, all with zero STACK and LOCAL. - SM120: 2763 passed, 198 expected skips, including 104 isolated trap cases and 12 cross-device/device-guard cases. - SM103/B300: 393 shared-contract passes, 4 multi-device skips, plus 180 original-target cases and 48 CPU-oracle cluster calls. - Unfiltered cluster memcheck and racecheck: all 48 calls passed with zero errors, hazards, or warnings; normal CPU-oracle/canary checks passed 10 cases with 2 unsupported BF16-sorted skips under each mode. - SM120 risk-based sanitizer coverage: 19 cases passed each sanitizer, with zero errors, hazards, or warnings; not full-suite sanitizer coverage. - The same wheel ran without rebuilding on Python 3.10 / PyTorch 2.10: 24 CPU checks, 52 exact SM120 CUDA checks, and 31 validation checks, including current streams and changed-input graph replay. Runtime used the final extracted wheel, not an installed/source shadow. Only the README fork/dev clone URL changed afterward; wheel payloads contain no README and compiled code and tests are unchanged. SM121 is compiled and resource-checked, not hardware-qualified. SM100 compilation is not SM100 hardware qualification. vLLM engine consumer integration and matched SM12 performance measurements remain separate. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- MANIFEST.in | 11 + README.md | 159 ++- csrc/api.cpp | 124 +- csrc/build_utils.py | 112 ++ csrc/cuda_kernels/common_parts.cuh | 26 +- csrc/cuda_kernels/sm120/topk_select.cu | 747 ++++++++++++ csrc/cuda_kernels/sm120/topk_select.h | 13 + csrc/cuda_kernels/v3_cluster/topk_select.cuh | 29 +- csrc/stable_tensor_checks.h | 85 ++ deep_select/interface.py | 24 +- setup.py | 52 +- tests/kernelkit/build.py | 8 +- tests/test_build_plan.py | 115 ++ tests/test_sm120.py | 1144 ++++++++++++++++++ tests/test_stable_contracts.py | 289 +++++ 15 files changed, 2854 insertions(+), 84 deletions(-) create mode 100644 MANIFEST.in create mode 100644 csrc/build_utils.py create mode 100644 csrc/cuda_kernels/sm120/topk_select.cu create mode 100644 csrc/cuda_kernels/sm120/topk_select.h create mode 100644 tests/test_build_plan.py create mode 100644 tests/test_sm120.py create mode 100644 tests/test_stable_contracts.py 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/README.md b/README.md index fd016c3..b3057de 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # DeepSelect -DeepSelect is a high performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) (which is used in DeepSeek V3.2, DeepSeek V4, and DeepSeek V4.1 models) and the sampler. It achieves 2 ~ 20x speedup compared to vanilla `torch.topk`. +DeepSelect is a high performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) (which is used in DeepSeek V3.2, DeepSeek V4, and DeepSeek V4.1 models) and the sampler. The upstream SM100 performance results report 2 ~ 20x speedup compared to vanilla `torch.topk`; these are not SM120/SM121 measurements. ## News @@ -9,6 +9,59 @@ DeepSelect is a high performance implementation of the TopK kernel used in DeepS ## Supported Cases +### Native SM120 / SM121 + +This port adds native CUDA selection for SM120 (`12.0`) and SM121 (`12.1`) +while retaining the default SM100/SM103 kernels. The implementation uses the +same native algorithms for both SM12 targets, with dispatch enabled only for +the exact target compiled into the extension. Other SM12 capabilities are +rejected; there is no `torch.topk` fallback. Explicit SM121 compilation is not +GPU qualification: SM121 hardware validation remains pending. This port does +not claim validated SM12 performance. + +The public `deep_select.topk` API supports BF16 and FP32 input, int32/int64 +indices, `1 <= topk <= 4096`, `0 <= vocab_size < 2**23`, and empty batches; +`batch_size` must fit int32. Zero-width rows are supported on SM120/SM121; +SM100/SM103 require `vocab_size > 0` for nonempty batches. Empty batches remain +supported on all compiled targets. `sorted=True` sorts by descending value (FP32 only, +requires `return_value=True`), while `sorted_index=True` sorts by ascending +index. These options are mutually exclusive. `begin` and `hint` remain +unsupported. Selection, optional values, sorting, variable row lengths, and +index offsets execute in CUDA. + +Unsorted int32 `topk == 512` with `2048 < vocab_size <= 131072` uses a +256-thread streaming selector. It scans 1024 inputs per tile in deterministic +bit-reversed tile order and filters against the exact worst retained key/index +pair. A bounded 2560-candidate buffer is compacted to 512 when its size reaches +1536. Candidate radix selection uses warp-private histograms and original-index +tie breaking; explicit shared storage is bounded by 32 KiB. Tile ordering is +not a worst-case performance guarantee. + +Within that K512 path, `65536 <= vocab_size <= 131072` and +`1 <= batch_size <= 16` use eight independent blocks per row followed by one +exact merge block. Local selection and merge retain original indices, without +intermediate sorting or index remapping. Per-call device scratch is +`batch_size * 8 * 1025` int32 words for keys, indices, and valid counts. +Allocation and both launches use the current stream and support CUDA graph +capture. Other streaming shapes use one block per row. Selection scans each +input once, with original-value gathers when values are requested. + +Other unsorted int32 configurations with `topk <= 512` use one 256-thread +block per row, whole-row 8-bit radix histograms, parallel pivot selection, and +scan-based collection. BF16 needs at most two value radix passes and FP32 at +most four; exact bucket completion can stop refinement early. Explicit shared +storage is bounded by 4 KiB without sort scratch. Other configurations use a +512-thread fallback with four radix passes and optional CUB sorting of at most +4096 packed keys, reusing shared storage with a compile-time bound of 48 KiB. +Compiler-generated shared storage is additional to these explicit bounds. +None of the native SM12 paths uses TMA or the upstream 16-block cluster. + +Equal values prefer the lower original index. `+0` and `-0` compare equal and +retain their original bit patterns in returned values; infinities and +subnormals are ordered numerically. Unsorted output order is unspecified, +except for short rows as described below. Optimized paths retain the NaN +contract below, and bounded output stores write only logical output elements. + TopK workloads vary widely, and the fastest algorithm & implementation highly depends on the input dtype, `batch_size`, `vocab_size`, and `topk`. This repository only focuses on the following cases: ### Lightning Indexer Scenario @@ -33,6 +86,9 @@ This scenario covers: ## Performance +The figures below are inherited upstream SM100 results, not SM120/SM121 +benchmarks. No SM12 speedup is established by this port. + Measured with the benchmark in [`tests/test.py`](tests/test.py) (`python3 tests/test.py --perf-only`), which reports the ratio against `torch.topk` on the same input. The metric is effective memory bandwidth: TopK does no @@ -53,12 +109,53 @@ float32, `vocab_size = 129280`, `topk = 512`. ## Installation ```bash -git clone https://github.com/deepseek-ai/DeepSelect.git +git clone --branch dev https://github.com/vllm-project/DeepSelect.git cd DeepSelect git submodule update --init --recursive pip install -v . ``` +### CUDA architecture selection + +The default source build targets `10.0a;10.3a`. Set +`DEEP_SELECT_CUDA_ARCH_LIST` to select explicit targets from `10.0a`, `10.3a`, +`10.0f`, `12.0`, and `12.1`, for example `12.0;12.1` or +`10.0a;10.3a;12.0;12.1`. Entries may be separated by spaces or semicolons. +`+PTX` adds matching PTX alongside the native cubin, for example `12.1+PTX`. +Unsupported entries are errors; `TORCH_CUDA_ARCH_LIST` does not replace this +package's default. + +SM120 (`12.0`) and `10.0a` require CUDA 12.8 or newer; `10.3a`, `10.0f`, and +SM121 (`12.1`) require CUDA 12.9 or newer. The [CUDA 12.9 release notes](https://docs.nvidia.com/cuda/archive/12.9.0/cuda-toolkit-release-notes/index.html#general-cuda) +explicitly introduce compiler target support for `sm_121`; compiler support +alone does not establish runtime qualification. + +SM120 and SM121 use the same native kernel source under `csrc/cuda_kernels/sm120`, +with the same dtype, top-k, sorting, alignment and unsupported-option limits +in `deep_select/interface.py`. Explicit SM121 selection emits `sm_121`, rather +than relying on an SM120 cubin's forward compatibility. The SM100/103 v3 and +cluster kernels never compile for either SM12 target. The SM12 family disables +fast math and preserves subnormal values (`--ftz=false`). + +The host API compiles once and all selected objects link into one shared library, +`deep_select.deep_select_cuda`, retaining `PyInit_deep_select_cuda`. The build +targets the PyTorch 2.10 stable ABI and CPython 3.10+ (`cp310-abi3`). Importing +this extension registers `torch.ops.deep_select.topk` and +`torch.ops.deep_select.get_alignment_requirement`; the Python API uses those +operators. This is one provider with one operator registration, not separate +SM100 and SM12 Python extension modules. + +The build defines `DEEP_SELECT_BUILD_SM100`, `DEEP_SELECT_BUILD_SM120`, and +`DEEP_SELECT_BUILD_SM121` according to the selected targets. Legacy source +consumers retain defaults of `1`, `0`, and `0`, respectively. Builds must pass +the artifact spill checker: LOCAL must be zero and STACK must not exceed 8 +bytes. Compiler target support alone is not runtime or performance qualification. + +The existing vLLM source consumer still targets `10.0f` and uses its existing +source glob. This standalone fork port does not enable SM12 in vLLM: separate +consumer integration is needed to include and compile the native SM12 sources +and select the matching targets while retaining a single operator provider. + ## Usage ```python @@ -76,22 +173,52 @@ values, indices = deep_select.topk( topk, sorted_index=True, # return each row's indices in ascending order indices_type=torch.int32, # torch.int32 or torch.int64 - return_value=True, # False skips the value output (~10% faster) + return_value=True, # False skips the value output ) # values: (batch_size, topk) of x.dtype # indices: (batch_size, topk) of indices_type ``` -The row stride of the input tensor (`x`) must be aligned to `deep_select.get_stride_requirement()[0]` bytes. For unaligned inputs, padding is necessary. - -Both outputs are allocated by the call, and their strides are aligned to `deep_select.get_stride_requirement()[1]` bytes (so they may be non-contiguous). Pass `output_idx=` to write indices into a buffer you own, and that buffer must satisfy the same stride requirement. +All tensor arguments must be on the same CUDA device as `input`. Execution +uses that device's current PyTorch CUDA stream, including non-default streams +and CUDA graph capture, and restores the caller's current device afterward. + +The input's last dimension must be contiguous, its row stride must be aligned +to `deep_select.get_stride_requirement()[0]` bytes, and its data pointer must be +32-byte aligned. Nonempty backing storage must include the final row rounded +up to 128 bytes for the upstream TMA kernels; SM12 retains the same validation. +Allocate padded rows first, then slice to the logical width: stride alignment +alone does not ensure final-row padding. + +Outputs have strides aligned to `deep_select.get_stride_requirement()[1]` +bytes and may be non-contiguous. `return_value=False` returns `(None, indices)`. +Pass `output_idx=` to supply a `(batch_size, topk)` buffer of `indices_type`; +it must have a contiguous last dimension, aligned row stride and 32-byte-aligned +pointer, and non-overlapping rows. Only logical output elements are written, +so trailing allocation padding is not required for outputs. + +Nonempty outputs must not share or overlap backing storage with `input`, `end`, +`output_idx_offset`, or each other. Even disjoint views of the same backing +storage are rejected, regardless of dtype or storage offset; allocate output +buffers separately. Pairs containing an empty tensor are exempt. These checks +also apply to direct `torch.ops.deep_select.topk` calls, including a supplied +`output_value` when `return_value=False`. For the full signature, see [`deep_select/interface.py`](deep_select/interface.py). ### Variable-length rows -`end` sets a per-row upper bound (exclusive). Rows shorter than `topk` are padded with -`value_oob_fill_value` / `idx_oob_fill_value`: +`end` is a contiguous `(batch_size,)` int32 CUDA tensor of per-row exclusive +upper bounds. The caller must ensure `0 <= end[i] <= vocab_size`; host +validation does not read CUDA tensor contents. `output_idx_offset` is an +optional `(batch_size,)` int32 CUDA tensor added only to valid indices, never +to fill entries. + +When `end[i] <= topk`, valid entries precede all fill entries and only the valid +prefix is sorted, even if the fill value is larger than the input values. +Padding uses `value_oob_fill_value` (default `-inf`) and `idx_oob_fill_value` +(default `2147483647`). The index fill must fit int32; the value fill may be a +float32-representable finite value, infinity, or NaN. ```python batch_size, vocab_size = 2, 129280 # 129280 is a multiple of 256, so float32 is fine @@ -104,7 +231,21 @@ values, indices = deep_select.topk(x, 1000, end=end, sorted=True, ### NaN handling -NaN checking is always on. With the default `abort_when_nan_found=True` the kernel invokes `trap()` and aborts. Rows whose length is `<= topk` are never NaN-checked. +NaN checking remains enabled in optimized paths. With the default +`abort_when_nan_found=True`, the kernel invokes `trap()` and aborts. With +`False`, index zero of the affected output row is set to `0x3F3F3F3F` without +adding the index offset; the rest of that row's outputs are unspecified. +Other rows continue normally. + +Rows whose length is `<= topk` are never NaN-checked. On SM12, sorting such a +short row by value puts NaNs first (ties by original index), preserves their +input payloads in returned values, and keeps fill entries at the end. +Unsorted short rows preserve input order. + +The input-device/stream guards and vector-boundary safety changes were adapted +from upstream [PR #4](https://github.com/deepseek-ai/DeepSelect/pull/4) +(`a24b1e4`) and [PR #6](https://github.com/deepseek-ai/DeepSelect/pull/6) +(`c799e83`), both authored by [morluto](https://github.com/morluto). ## Citation 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 #include +#include #include #include #include -#include #include +#include +#include #include -#include #include +#include #include #include -#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 +#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::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(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(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(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 get_alignment_requirement() { diff --git a/csrc/build_utils.py b/csrc/build_utils.py new file mode 100644 index 0000000..bba8d09 --- /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": ["--use_fast_math", "--ftz=false"] + architecture_flags(sm100), + "sm120": ["--ftz=false"] + 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 +__device__ __forceinline__ +void st_global_bounded(ValueT* ptr, ValueT src[NUM_VALUES], uint32_t remaining) { + if (remaining >= NUM_VALUES) { + st_global(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(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(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(result_indices + i, out); + st_global_bounded(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(values, (UIntValueT*)(smem_value_buf + i)); - st_global((UIntValueT*)(result_values + i), values); + st_global_bounded((UIntValueT*)(result_values + i), values, args.topk - i); } } } @@ -258,15 +270,15 @@ struct EpilogueRunner { auto store = [&](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(dst, src); + st_global_bounded(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( + 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 +#include +#include +#include +#include + +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; +using BlockScan = cub::BlockScan; + +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 +__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(args.input) + row * args.stride_input_batch; + IndexT* output_index = static_cast(args.output_index) + row * args.stride_output_index_batch; + ValueT* output_value = args.return_value + ? static_cast(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(args.vocab_size); + if (end < 0 || static_cast(end) > args.vocab_size) { + if (tid == 0) asm volatile("trap;"); + return; + } + const uint32_t length = static_cast(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(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(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(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(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(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(~index); + if (args.sorted_value) + key |= static_cast(ordered_key(static_cast(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(keys[j]); + const bool valid = slot < valid_count; + output_index[slot] = valid ? static_cast(static_cast(index) + offset) + : static_cast(args.idx_oob_fill_value); + if (args.return_value) + output_value[slot] = valid ? input[index] : static_cast(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(static_cast(index) + offset) + : static_cast(args.idx_oob_fill_value); + if (args.return_value) + output_value[slot] = valid ? input[index] : static_cast(args.value_oob_fill_value); + } + } +} + +constexpr int FAST_THREADS = 256; +using FastScan = cub::BlockScan; + +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 +__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 +__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( + smem, keys, indices, active, selected, rank); + if (!smem.complete) + streaming_radix(smem, keys, indices, active, selected, rank); + + uint32_t slot; + FastScan(smem.scan).ExclusiveSum(static_cast(__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 +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(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(segment[2 * STREAM_TOPK])) return false; + key = reinterpret_cast(segment)[candidate]; + index = static_cast(segment[STREAM_TOPK + candidate]); + return true; + } +}; + +template +__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(__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(smem, count); + count = STREAM_TOPK; + have_threshold = true; + } + } + if (count > STREAM_TOPK) { + streaming_compact(smem, count); + count = STREAM_TOPK; + } + return true; +} + +template +__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(args.input) + row * args.stride_input_batch; + int32_t* output_index = static_cast(args.output_index) + row * args.stride_output_index_batch; + ValueT* output_value = nullptr; + if constexpr (ReturnValue) + output_value = static_cast(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(args.vocab_size); + if (end < 0 || static_cast(end) > args.vocab_size) { + if (tid == 0) asm volatile("trap;"); + return; + } + const uint32_t length = static_cast(end); + if (length <= STREAM_TOPK) { + for (uint32_t slot = tid; slot < STREAM_TOPK; slot += FAST_THREADS) { + output_index[slot] = slot < length + ? static_cast(static_cast(slot) + offset) : args.idx_oob_fill_value; + if constexpr (ReturnValue) + output_value[slot] = slot < length ? input[slot] : static_cast(args.value_oob_fill_value); + } + return; + } + + uint32_t count; + if (!streaming_select(smem, StreamingInputSource{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(static_cast(index) + offset); + if constexpr (ReturnValue) output_value[slot] = input[index]; + } +} + +template +__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(args.vocab_size); + if (end < 0 || static_cast(end) > args.vocab_size) { + if (tid == 0) asm volatile("trap;"); + return; + } + const uint32_t length = static_cast(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(args.input) + row * args.stride_input_batch; + uint32_t count; + if (!streaming_select( + smem, StreamingInputSource{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(segment)[slot] = smem.keys[slot]; + segment[STREAM_TOPK + slot] = static_cast(smem.indices[slot]); + } + __syncthreads(); + if (tid == 0) segment[2 * STREAM_TOPK] = static_cast(count); +} + +template +__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(args.vocab_size); + if (end < 0 || static_cast(end) > args.vocab_size) { + if (tid == 0) asm volatile("trap;"); + return; + } + const uint32_t length = static_cast(end); + const int32_t* row_workspace = workspace + row * SEGMENTED_PARTITIONS * SEGMENTED_WORKSPACE_STRIDE; + int32_t* output_index = static_cast(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(args.input) + row * args.stride_input_batch; + ValueT* output_value = nullptr; + if constexpr (ReturnValue) + output_value = static_cast(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(static_cast(slot) + offset) : args.idx_oob_fill_value; + if constexpr (ReturnValue) + output_value[slot] = slot < length ? input[slot] : static_cast(args.value_oob_fill_value); + } + return; + } + uint32_t count; + if (!streaming_select(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(static_cast(index) + offset); + if constexpr (ReturnValue) output_value[slot] = input[index]; + } +} + +template +void launch_segmented(const TopkSelectArgs& args, int32_t* workspace) { + topk_segmented_local_kernel + <<>>(args, workspace); + check_kernel_launch(); + if (args.return_value) + topk_segmented_merge_kernel + <<>>(args, workspace); + else + topk_segmented_merge_kernel + <<>>(args, workspace); + check_kernel_launch(); +} + +template +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 +__global__ __launch_bounds__(FAST_THREADS) void topk_unsorted_kernel(TopkSelectArgs args) { + __shared__ FastSharedMemory 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(args.input) + row * args.stride_input_batch; + int32_t* output_index = static_cast(args.output_index) + row * args.stride_output_index_batch; + ValueT* output_value = nullptr; + if constexpr (ReturnValue) + output_value = static_cast(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(args.vocab_size); + if (end < 0 || static_cast(end) > args.vocab_size) { + if (tid == 0) asm volatile("trap;"); + return; + } + const uint32_t length = static_cast(end); + if (length <= args.topk) { + for (uint32_t slot = tid; slot < args.topk; slot += FAST_THREADS) { + output_index[slot] = slot < length + ? static_cast(static_cast(slot) + offset) : args.idx_oob_fill_value; + if constexpr (ReturnValue) + output_value[slot] = slot < length ? input[slot] : static_cast(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(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(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(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(static_cast(index) + offset); + if constexpr (ReturnValue) output_value[slot] = input[index]; + } +} + +template +void launch_unsorted(const TopkSelectArgs& args) { + if (args.return_value) + topk_unsorted_kernel<<>>(args); + else + topk_unsorted_kernel<<>>(args); +} + +template +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<<>>(args); + else + topk_streaming_kernel<<>>(args); + } else launch_unsorted(args); + } else if (indices_int64) + topk_kernel<<>>(args); + else + topk_kernel<<>>(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(args, workspace); + else launch_segmented(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(args, indices_int64); + else launch(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 +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(&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 class TopkSelectKernelBF16Cluster : public topk_select_common::TopkSelectKernelBF16Base { using BF16Base = topk_select_common::TopkSelectKernelBF16Base; @@ -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(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(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 #include #include +#include #include #include +#include #include #include +#include namespace deep_select { @@ -56,6 +59,88 @@ static inline PtrT* get_optional_tensor_ptr(const std::optional& tensor_ } } +static inline uint64_t checked_add(uint64_t a, uint64_t b) { + STD_TORCH_CHECK(a <= uint64_t(std::numeric_limits::max()) && + b <= uint64_t(std::numeric_limits::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::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(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(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(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::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(tensor.data_ptr()) % 32 == 0, + name, ".data_ptr() must be 32-byte aligned"); +} + } // namespace deep_select // Check whether the given tensor (or optional) 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]= 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) From 8b11cabc60df417925260112cb73259af1e3a67e Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:50:11 -0700 Subject: [PATCH 2/4] chore(build): keep include paths in upstream style Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index fad05b5..c1f7d9a 100644 --- a/setup.py +++ b/setup.py @@ -144,14 +144,14 @@ def build_on_cuda_platform(): "--threads", os.getenv("NVCC_THREADS") or "16", ], }, - include_dirs=[str(path) for path in [ + include_dirs=[ this_dir / "csrc", this_dir / "csrc" / "3rdparty" / "cutlass" / "include", this_dir / "csrc" / "3rdparty" / "kerutils" / "include", Path(CUDA_HOME) / "include" / "cccl", Path(CUDA_HOME) / "targets" / "x86_64-linux" / "include" / "cccl", Path(CUDA_HOME) / "targets" / "sbsa-linux" / "include" / "cccl", - ]], + ], extra_link_args=[ f'-L{Path(CUDA_HOME) / "lib" / "stubs"}', f'-L{Path(CUDA_HOME) / "targets" / "x86_64-linux" / "lib" / "stubs"}', From 5d71dd15fa3071c174cff7c44d9a65d20155aac4 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:53:41 -0700 Subject: [PATCH 3/4] chore(build): minimize setup changes against upstream Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- setup.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/setup.py b/setup.py index c1f7d9a..bddd98c 100644 --- a/setup.py +++ b/setup.py @@ -101,8 +101,11 @@ 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 - if CUDA_HOME is None: - raise RuntimeError("A CUDA toolkit is required to build DeepSelect") + assert CUDA_HOME is not None, "PyTorch must be compiled with CUDA support" + + def append_nvcc_threads(nvcc_extra_args): + nvcc_threads = os.getenv("NVCC_THREADS") or "16" + return nvcc_extra_args + ["--threads", nvcc_threads] nvcc_version = subprocess.check_output( [os.path.join(CUDA_HOME, "bin", "nvcc"), "--version"], @@ -115,9 +118,9 @@ def build_on_cuda_platform(): os.getenv("DEEP_SELECT_CUDA_ARCH_LIST"), (major, minor) ) sources, macros, group_flags = build_plan(CUDA_SOURCES, architectures) - print(f"DeepSelect CUDA architectures: {architectures}") - this_dir = Path(__file__).resolve().parent + this_dir = os.path.dirname(os.path.abspath(__file__)) + ext_modules = [CUDAExtension( name="deep_select.deep_select_cuda", sources=sources, @@ -126,7 +129,7 @@ def build_on_cuda_platform(): extra_compile_args={ # Target the torch 2.10 stable ABI (minimum supported torch version) "cxx": ["-O3", "-std=c++20", "-DNDEBUG", "-Wno-deprecated-declarations", "-DTORCH_TARGET_VERSION=0x020a000000000000", "-DUSE_CUDA"], - "nvcc": [ + "nvcc": append_nvcc_threads([ "-O3", "-std=c++20", "-DTORCH_TARGET_VERSION=0x020a000000000000", @@ -141,19 +144,16 @@ def build_on_cuda_platform(): "--ptxas-options=-v,--register-usage-level=10,--warn-on-spills,--warn-on-double-precision-use", "-lineinfo", "--source-in-ptx", - "--threads", os.getenv("NVCC_THREADS") or "16", - ], + ]) }, include_dirs=[ - this_dir / "csrc", - this_dir / "csrc" / "3rdparty" / "cutlass" / "include", - this_dir / "csrc" / "3rdparty" / "kerutils" / "include", - Path(CUDA_HOME) / "include" / "cccl", + Path(this_dir) / "csrc", + Path(this_dir) / "csrc" / "3rdparty" / "cutlass" / "include", + Path(this_dir) / "csrc" / "3rdparty" / "kerutils" / "include", Path(CUDA_HOME) / "targets" / "x86_64-linux" / "include" / "cccl", Path(CUDA_HOME) / "targets" / "sbsa-linux" / "include" / "cccl", ], extra_link_args=[ - f'-L{Path(CUDA_HOME) / "lib" / "stubs"}', f'-L{Path(CUDA_HOME) / "targets" / "x86_64-linux" / "lib" / "stubs"}', f'-L{Path(CUDA_HOME) / "targets" / "sbsa-linux" / "lib" / "stubs"}', "-lcuda", From 60f69a1c91caf2d1a5fa2ef4bda80df17fb3d335 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:09:15 -0700 Subject: [PATCH 4/4] chore(build): retain upstream math flags and README Use upstream fast-math and explicit FTZ=false for both kernel families, leaving only gencode selection in the per-family build flags. Restore README byte-for-byte to dev to avoid needless sync conflicts. SM120/121 resources and normalized SASS match the prior build. Actual checker-on wheel build, 15 build tests, 2763 SM120 cases and selected 19-case memcheck/racecheck runs passed. README-only source drift did not change the tested artifact. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- README.md | 159 +++------------------------------------ csrc/build_utils.py | 4 +- setup.py | 2 + tests/test_build_plan.py | 5 +- 4 files changed, 16 insertions(+), 154 deletions(-) diff --git a/README.md b/README.md index b3057de..fd016c3 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # DeepSelect -DeepSelect is a high performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) (which is used in DeepSeek V3.2, DeepSeek V4, and DeepSeek V4.1 models) and the sampler. The upstream SM100 performance results report 2 ~ 20x speedup compared to vanilla `torch.topk`; these are not SM120/SM121 measurements. +DeepSelect is a high performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) (which is used in DeepSeek V3.2, DeepSeek V4, and DeepSeek V4.1 models) and the sampler. It achieves 2 ~ 20x speedup compared to vanilla `torch.topk`. ## News @@ -9,59 +9,6 @@ DeepSelect is a high performance implementation of the TopK kernel used in DeepS ## Supported Cases -### Native SM120 / SM121 - -This port adds native CUDA selection for SM120 (`12.0`) and SM121 (`12.1`) -while retaining the default SM100/SM103 kernels. The implementation uses the -same native algorithms for both SM12 targets, with dispatch enabled only for -the exact target compiled into the extension. Other SM12 capabilities are -rejected; there is no `torch.topk` fallback. Explicit SM121 compilation is not -GPU qualification: SM121 hardware validation remains pending. This port does -not claim validated SM12 performance. - -The public `deep_select.topk` API supports BF16 and FP32 input, int32/int64 -indices, `1 <= topk <= 4096`, `0 <= vocab_size < 2**23`, and empty batches; -`batch_size` must fit int32. Zero-width rows are supported on SM120/SM121; -SM100/SM103 require `vocab_size > 0` for nonempty batches. Empty batches remain -supported on all compiled targets. `sorted=True` sorts by descending value (FP32 only, -requires `return_value=True`), while `sorted_index=True` sorts by ascending -index. These options are mutually exclusive. `begin` and `hint` remain -unsupported. Selection, optional values, sorting, variable row lengths, and -index offsets execute in CUDA. - -Unsorted int32 `topk == 512` with `2048 < vocab_size <= 131072` uses a -256-thread streaming selector. It scans 1024 inputs per tile in deterministic -bit-reversed tile order and filters against the exact worst retained key/index -pair. A bounded 2560-candidate buffer is compacted to 512 when its size reaches -1536. Candidate radix selection uses warp-private histograms and original-index -tie breaking; explicit shared storage is bounded by 32 KiB. Tile ordering is -not a worst-case performance guarantee. - -Within that K512 path, `65536 <= vocab_size <= 131072` and -`1 <= batch_size <= 16` use eight independent blocks per row followed by one -exact merge block. Local selection and merge retain original indices, without -intermediate sorting or index remapping. Per-call device scratch is -`batch_size * 8 * 1025` int32 words for keys, indices, and valid counts. -Allocation and both launches use the current stream and support CUDA graph -capture. Other streaming shapes use one block per row. Selection scans each -input once, with original-value gathers when values are requested. - -Other unsorted int32 configurations with `topk <= 512` use one 256-thread -block per row, whole-row 8-bit radix histograms, parallel pivot selection, and -scan-based collection. BF16 needs at most two value radix passes and FP32 at -most four; exact bucket completion can stop refinement early. Explicit shared -storage is bounded by 4 KiB without sort scratch. Other configurations use a -512-thread fallback with four radix passes and optional CUB sorting of at most -4096 packed keys, reusing shared storage with a compile-time bound of 48 KiB. -Compiler-generated shared storage is additional to these explicit bounds. -None of the native SM12 paths uses TMA or the upstream 16-block cluster. - -Equal values prefer the lower original index. `+0` and `-0` compare equal and -retain their original bit patterns in returned values; infinities and -subnormals are ordered numerically. Unsorted output order is unspecified, -except for short rows as described below. Optimized paths retain the NaN -contract below, and bounded output stores write only logical output elements. - TopK workloads vary widely, and the fastest algorithm & implementation highly depends on the input dtype, `batch_size`, `vocab_size`, and `topk`. This repository only focuses on the following cases: ### Lightning Indexer Scenario @@ -86,9 +33,6 @@ This scenario covers: ## Performance -The figures below are inherited upstream SM100 results, not SM120/SM121 -benchmarks. No SM12 speedup is established by this port. - Measured with the benchmark in [`tests/test.py`](tests/test.py) (`python3 tests/test.py --perf-only`), which reports the ratio against `torch.topk` on the same input. The metric is effective memory bandwidth: TopK does no @@ -109,53 +53,12 @@ float32, `vocab_size = 129280`, `topk = 512`. ## Installation ```bash -git clone --branch dev https://github.com/vllm-project/DeepSelect.git +git clone https://github.com/deepseek-ai/DeepSelect.git cd DeepSelect git submodule update --init --recursive pip install -v . ``` -### CUDA architecture selection - -The default source build targets `10.0a;10.3a`. Set -`DEEP_SELECT_CUDA_ARCH_LIST` to select explicit targets from `10.0a`, `10.3a`, -`10.0f`, `12.0`, and `12.1`, for example `12.0;12.1` or -`10.0a;10.3a;12.0;12.1`. Entries may be separated by spaces or semicolons. -`+PTX` adds matching PTX alongside the native cubin, for example `12.1+PTX`. -Unsupported entries are errors; `TORCH_CUDA_ARCH_LIST` does not replace this -package's default. - -SM120 (`12.0`) and `10.0a` require CUDA 12.8 or newer; `10.3a`, `10.0f`, and -SM121 (`12.1`) require CUDA 12.9 or newer. The [CUDA 12.9 release notes](https://docs.nvidia.com/cuda/archive/12.9.0/cuda-toolkit-release-notes/index.html#general-cuda) -explicitly introduce compiler target support for `sm_121`; compiler support -alone does not establish runtime qualification. - -SM120 and SM121 use the same native kernel source under `csrc/cuda_kernels/sm120`, -with the same dtype, top-k, sorting, alignment and unsupported-option limits -in `deep_select/interface.py`. Explicit SM121 selection emits `sm_121`, rather -than relying on an SM120 cubin's forward compatibility. The SM100/103 v3 and -cluster kernels never compile for either SM12 target. The SM12 family disables -fast math and preserves subnormal values (`--ftz=false`). - -The host API compiles once and all selected objects link into one shared library, -`deep_select.deep_select_cuda`, retaining `PyInit_deep_select_cuda`. The build -targets the PyTorch 2.10 stable ABI and CPython 3.10+ (`cp310-abi3`). Importing -this extension registers `torch.ops.deep_select.topk` and -`torch.ops.deep_select.get_alignment_requirement`; the Python API uses those -operators. This is one provider with one operator registration, not separate -SM100 and SM12 Python extension modules. - -The build defines `DEEP_SELECT_BUILD_SM100`, `DEEP_SELECT_BUILD_SM120`, and -`DEEP_SELECT_BUILD_SM121` according to the selected targets. Legacy source -consumers retain defaults of `1`, `0`, and `0`, respectively. Builds must pass -the artifact spill checker: LOCAL must be zero and STACK must not exceed 8 -bytes. Compiler target support alone is not runtime or performance qualification. - -The existing vLLM source consumer still targets `10.0f` and uses its existing -source glob. This standalone fork port does not enable SM12 in vLLM: separate -consumer integration is needed to include and compile the native SM12 sources -and select the matching targets while retaining a single operator provider. - ## Usage ```python @@ -173,52 +76,22 @@ values, indices = deep_select.topk( topk, sorted_index=True, # return each row's indices in ascending order indices_type=torch.int32, # torch.int32 or torch.int64 - return_value=True, # False skips the value output + return_value=True, # False skips the value output (~10% faster) ) # values: (batch_size, topk) of x.dtype # indices: (batch_size, topk) of indices_type ``` -All tensor arguments must be on the same CUDA device as `input`. Execution -uses that device's current PyTorch CUDA stream, including non-default streams -and CUDA graph capture, and restores the caller's current device afterward. - -The input's last dimension must be contiguous, its row stride must be aligned -to `deep_select.get_stride_requirement()[0]` bytes, and its data pointer must be -32-byte aligned. Nonempty backing storage must include the final row rounded -up to 128 bytes for the upstream TMA kernels; SM12 retains the same validation. -Allocate padded rows first, then slice to the logical width: stride alignment -alone does not ensure final-row padding. - -Outputs have strides aligned to `deep_select.get_stride_requirement()[1]` -bytes and may be non-contiguous. `return_value=False` returns `(None, indices)`. -Pass `output_idx=` to supply a `(batch_size, topk)` buffer of `indices_type`; -it must have a contiguous last dimension, aligned row stride and 32-byte-aligned -pointer, and non-overlapping rows. Only logical output elements are written, -so trailing allocation padding is not required for outputs. - -Nonempty outputs must not share or overlap backing storage with `input`, `end`, -`output_idx_offset`, or each other. Even disjoint views of the same backing -storage are rejected, regardless of dtype or storage offset; allocate output -buffers separately. Pairs containing an empty tensor are exempt. These checks -also apply to direct `torch.ops.deep_select.topk` calls, including a supplied -`output_value` when `return_value=False`. +The row stride of the input tensor (`x`) must be aligned to `deep_select.get_stride_requirement()[0]` bytes. For unaligned inputs, padding is necessary. + +Both outputs are allocated by the call, and their strides are aligned to `deep_select.get_stride_requirement()[1]` bytes (so they may be non-contiguous). Pass `output_idx=` to write indices into a buffer you own, and that buffer must satisfy the same stride requirement. For the full signature, see [`deep_select/interface.py`](deep_select/interface.py). ### Variable-length rows -`end` is a contiguous `(batch_size,)` int32 CUDA tensor of per-row exclusive -upper bounds. The caller must ensure `0 <= end[i] <= vocab_size`; host -validation does not read CUDA tensor contents. `output_idx_offset` is an -optional `(batch_size,)` int32 CUDA tensor added only to valid indices, never -to fill entries. - -When `end[i] <= topk`, valid entries precede all fill entries and only the valid -prefix is sorted, even if the fill value is larger than the input values. -Padding uses `value_oob_fill_value` (default `-inf`) and `idx_oob_fill_value` -(default `2147483647`). The index fill must fit int32; the value fill may be a -float32-representable finite value, infinity, or NaN. +`end` sets a per-row upper bound (exclusive). Rows shorter than `topk` are padded with +`value_oob_fill_value` / `idx_oob_fill_value`: ```python batch_size, vocab_size = 2, 129280 # 129280 is a multiple of 256, so float32 is fine @@ -231,21 +104,7 @@ values, indices = deep_select.topk(x, 1000, end=end, sorted=True, ### NaN handling -NaN checking remains enabled in optimized paths. With the default -`abort_when_nan_found=True`, the kernel invokes `trap()` and aborts. With -`False`, index zero of the affected output row is set to `0x3F3F3F3F` without -adding the index offset; the rest of that row's outputs are unspecified. -Other rows continue normally. - -Rows whose length is `<= topk` are never NaN-checked. On SM12, sorting such a -short row by value puts NaNs first (ties by original index), preserves their -input payloads in returned values, and keeps fill entries at the end. -Unsorted short rows preserve input order. - -The input-device/stream guards and vector-boundary safety changes were adapted -from upstream [PR #4](https://github.com/deepseek-ai/DeepSelect/pull/4) -(`a24b1e4`) and [PR #6](https://github.com/deepseek-ai/DeepSelect/pull/6) -(`c799e83`), both authored by [morluto](https://github.com/morluto). +NaN checking is always on. With the default `abort_when_nan_found=True` the kernel invokes `trap()` and aborts. Rows whose length is `<= topk` are never NaN-checked. ## Citation diff --git a/csrc/build_utils.py b/csrc/build_utils.py index bba8d09..cbae3ca 100644 --- a/csrc/build_utils.py +++ b/csrc/build_utils.py @@ -66,8 +66,8 @@ def build_plan(cuda_sources, architectures): sources.append(SM120_SOURCE) flags = { "host": [], - "sm100": ["--use_fast_math", "--ftz=false"] + architecture_flags(sm100), - "sm120": ["--ftz=false"] + architecture_flags(sm12), + "sm100": architecture_flags(sm100), + "sm120": architecture_flags(sm12), } return sources, macros, flags diff --git a/setup.py b/setup.py index bddd98c..315efb0 100644 --- a/setup.py +++ b/setup.py @@ -141,6 +141,8 @@ def append_nvcc_threads(nvcc_extra_args): "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", "--expt-relaxed-constexpr", "--expt-extended-lambda", + "--use_fast_math", + "--ftz=false", # Don't flush subnormals to zero, since we're going to use float addition to simulate integer addition (in order to be faster) "--ptxas-options=-v,--register-usage-level=10,--warn-on-spills,--warn-on-double-precision-use", "-lineinfo", "--source-in-ptx", diff --git a/tests/test_build_plan.py b/tests/test_build_plan.py index 0682f48..e779de1 100644 --- a/tests/test_build_plan.py +++ b/tests/test_build_plan.py @@ -97,9 +97,10 @@ def test_single_extension_grouped_compile(monkeypatch, tmp_path, entry, expected 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]: - assert "--use_fast_math" not in flags - assert "--ftz=false" in flags 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"})