dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat(CCCL): device-level CUB algorithms for MoE dispatch

Browse Source 
Add complete CCCL CUB header tree (1394 files) to cccl_preload/include/:
- cub/device/ — DeviceRadixSort, DeviceScan, DeviceHistogram, DeviceReduce, DeviceSelect
- cub/agent/ — all agent implementations (sort, scan, reduce, histogram, etc)
- cub/block/ — BlockScan, BlockReduce, BlockExchange, BlockLoad, BlockStore, etc
- cub/warp/ — WarpScan, WarpReduce, WarpExchange, WarpMergeSort
- cub/thread/ — thread-level operators
- thrust/ — sort_by_key, iterator utilities
- cuda/ — execution, stream, memory_resource, functional

New kernel: cccl_moe_sort_scatter.cu
- Uses CUB DeviceRadixSort::SortPairs to sort (expert_id, token_idx) pairs
- O(n) radix sort replaces O(n log n) torch.argsort in MoE prefill path
- Boundary detection + fill for expert offsets/sizes
- Compiled against CCCL upstream headers (not corex CUB) to avoid BI-V100 bugs

Previously only 288 CCCL headers (CachingDeviceAllocator only).
Now 1394 headers — full CUB device-level algorithm stack available for
all future kernels.
This commit is contained in:
project6-dev
2026-08-13 11:18:52 +00:00
parent 7ba97f7977
commit 4c365b8c03
1108 changed files with 294533 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
qwen3_6_scripts/cccl_preload/include/cub/detail/array_utils.cuh Normal file
View File

@@ -0,0 +1,54 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/type_traits.cuh> // static_size_v
#include <cub/util_namespace.cuh>
#include <cuda/std/__iterator/iterator_traits.h>
#include <cuda/std/__type_traits/conditional.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__utility/integer_sequence.h>
#include <cuda/std/array>
#include <cuda/std/cstddef>
CUB_NAMESPACE_BEGIN
namespace detail
{
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
/***********************************************************************************************************************
* Generic Array-like to Array Conversion
**********************************************************************************************************************/
template <typename CastType, typename Input, ::cuda::std::size_t... i>
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::array<CastType, static_size_v<Input>>
to_array_impl(const Input& input, ::cuda::std::index_sequence<i...>)
{
using ArrayType = ::cuda::std::array<CastType, static_size_v<Input>>;
return ArrayType{static_cast<CastType>(input[i])...};
}
template <typename CastType = void, typename Input>
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::array<CastType, static_size_v<Input>>
to_array(const Input& input)
{
using InputType = ::cuda::std::iter_value_t<Input>;
using CastType1 = ::cuda::std::_If<::cuda::std::is_same_v<CastType, void>, InputType, CastType>;
return to_array_impl<CastType1>(input, ::cuda::std::make_index_sequence<static_size_v<Input>>{});
}
#endif // !_CCCL_DOXYGEN_INVOKED
} // namespace detail
CUB_NAMESPACE_END

126
qwen3_6_scripts/cccl_preload/include/cub/detail/binary_search_helpers.cuh Normal file
View File

@@ -0,0 +1,126 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/__algorithm/lower_bound.h>
#include <cuda/std/__algorithm/upper_bound.h>
#include <cuda/std/__iterator/iterator_traits.h>
#include <cuda/std/cstddef>
#include <cuda/std/tuple>
CUB_NAMESPACE_BEGIN
namespace detail::find
{
constexpr ::cuda::std::ptrdiff_t linear_lower_bound_threshold = 8;
template <typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT, typename Mode>
struct comp_wrapper_t
{
RangeIteratorT first;
RangeNumItemsT num_items;
CompareOpT op;
template <typename Value, typename Output>
_CCCL_DEVICE _CCCL_FORCEINLINE void operator()(::cuda::std::tuple<Value, Output> args) const
{
using DifferenceT = ::cuda::std::iter_difference_t<RangeIteratorT>;
const auto last = first + static_cast<DifferenceT>(num_items);
::cuda::std::get<1>(args) = Mode::Invoke(first, last, ::cuda::std::get<0>(args), op);
}
};
template <typename Mode, typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT>
_CCCL_HOST_DEVICE auto make_comp_wrapper(RangeIteratorT first, RangeNumItemsT num_items, CompareOpT comp)
{
return comp_wrapper_t<RangeIteratorT, RangeNumItemsT, CompareOpT, Mode>{first, num_items, comp};
}
struct lower_bound
{
template <typename RangeIteratorT, typename DifferenceT, typename T, typename CompareOpT>
_CCCL_DEVICE _CCCL_FORCEINLINE static DifferenceT
Linear(RangeIteratorT first, DifferenceT num_items, const T& value, CompareOpT comp)
{
DifferenceT retval = 0;
for (DifferenceT i = 0; i < num_items; ++i)
{
retval += static_cast<DifferenceT>(comp(first[i], value));
}
return retval;
}
template <typename RangeIteratorT, typename T, typename CompareOpT>
_CCCL_DEVICE _CCCL_FORCEINLINE static ::cuda::std::ptrdiff_t
Invoke(RangeIteratorT first, RangeIteratorT last, const T& value, CompareOpT comp)
{
return ::cuda::std::lower_bound(first, last, value, comp) - first;
}
};
struct upper_bound
{
template <typename RangeIteratorT, typename DifferenceT, typename T, typename CompareOpT>
_CCCL_DEVICE _CCCL_FORCEINLINE static DifferenceT
Linear(RangeIteratorT first, DifferenceT num_items, const T& value, CompareOpT comp)
{
DifferenceT retval = 0;
for (DifferenceT i = 0; i < num_items; ++i)
{
retval += static_cast<DifferenceT>(!comp(value, first[i]));
}
return retval;
}
template <typename RangeIteratorT, typename T, typename CompareOpT>
_CCCL_DEVICE _CCCL_FORCEINLINE static ::cuda::std::ptrdiff_t
Invoke(RangeIteratorT first, RangeIteratorT last, const T& value, CompareOpT comp)
{
return ::cuda::std::upper_bound(first, last, value, comp) - first;
}
};
template <typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT, typename Mode>
struct binary_search_transform_op_t
{
RangeIteratorT first;
RangeNumItemsT num_items;
CompareOpT op;
template <typename Value>
_CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::ptrdiff_t operator()(const Value& value) const
{
using DifferenceT = ::cuda::std::iter_difference_t<RangeIteratorT>;
const auto count = static_cast<DifferenceT>(num_items);
if (num_items <= static_cast<RangeNumItemsT>(linear_lower_bound_threshold))
{
return Mode::Linear(first, count, value, op);
}
return Mode::Invoke(first, first + count, value, op);
}
};
template <typename Mode, typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT>
_CCCL_HOST_DEVICE auto make_binary_search_transform_op(RangeIteratorT first, RangeNumItemsT num_items, CompareOpT comp)
{
return binary_search_transform_op_t<RangeIteratorT, RangeNumItemsT, CompareOpT, Mode>{first, num_items, comp};
}
} // namespace detail::find
CUB_NAMESPACE_END

160
qwen3_6_scripts/cccl_preload/include/cub/detail/cc_dispatch.cuh Normal file
View File

@@ -0,0 +1,160 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/__device/compute_capability.h>
#include <cuda/std/__type_traits/is_empty.h>
#include <cuda/std/__utility/forward.h>
#include <cuda/std/__utility/integer_sequence.h>
#include <cuda/std/array>
CUB_NAMESPACE_BEGIN
namespace detail
{
// makes a functor that gets the policy for CC from PolicySelector when called
template <typename PolicySelector, int CC>
struct policy_getter : PolicySelector
{
_CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const
{
return PolicySelector::operator()(::cuda::compute_capability{CC});
}
};
// Device-only variant for kernel-side compile-time policy queries.
template <typename PolicySelector, int CC>
struct device_policy_getter : PolicySelector
{
_CCCL_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const
{
return PolicySelector::operator()(::cuda::compute_capability{CC});
}
};
#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC)
# if _CCCL_STD_VER < 2020 && !_CCCL_COMPILER(GCC, <, 8)
template <typename CudaCcSeq, typename PolicySelector, size_t... Is>
struct lowest_cc_resolver;
// we keep the compile-time build up of the mapping table outside a template parameterized by a user-provided callable
template <int... CudaCcs, typename PolicySelector, size_t... Is>
struct lowest_cc_resolver<::cuda::std::integer_sequence<int, CudaCcs...>, PolicySelector, Is...>
{
static_assert(sizeof...(CudaCcs) == sizeof...(Is));
using policy_t = decltype(PolicySelector{}(::cuda::compute_capability{}));
static constexpr ::cuda::compute_capability all_ccs[sizeof...(Is)]{::cuda::compute_capability{CudaCcs}...};
static constexpr policy_t all_policies[sizeof...(Is)]{PolicySelector{}(all_ccs[Is])...};
_CCCL_HOST_DEVICE_API static constexpr auto find_lowest(size_t i) -> ::cuda::compute_capability
{
const auto& policy = all_policies[i];
while (i > 0 && policy == all_policies[i - 1])
{
--i;
}
return all_ccs[i];
}
static constexpr ::cuda::compute_capability lowest_cc_with_same_policy[sizeof...(Is)]{find_lowest(Is)...};
};
# endif // if _CCCL_STD_VER < 2020 && !_CCCL_COMPILER(GCC, <, 8)
// GCC below 12 ICEs in some cases when creating an integral_constant holding a policy
# if _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12)
template <typename Tp, Tp P>
struct policy_constant
{
_CCCL_API constexpr auto operator()() const noexcept
{
return P;
}
};
# else // _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12)
template <typename Tp, Tp P> // using <auto P> will miscompile on GCC 12
using policy_constant = ::cuda::std::integral_constant<Tp, P>;
# endif // _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12)
template <typename PolicySelector, typename FunctorT, size_t... Is>
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_to_cc_list(
PolicySelector policy_selector, ::cuda::compute_capability device_cc, FunctorT&& f, ::cuda::std::index_sequence<Is...>)
{
constexpr auto all_ccs = ::cuda::__target_compute_capabilities();
_CCCL_ASSERT(((device_cc == all_ccs[Is]) || ...),
"device_cc must appear in the list of compute capabilities compiled for");
cudaError_t e = cudaErrorInvalidDeviceFunction;
# if _CCCL_STD_VER >= 2020
// In C++20, we just create an integral_constant holding the policy, because policies are structural types in C++20.
// This causes f to be only instantiated for each distinct policy, since the same policy for different arches results
// in the same integral_constant type passed to f
using policy_t = decltype(policy_selector(::cuda::compute_capability{}));
(..., (device_cc == all_ccs[Is] ? (e = f(policy_constant<policy_t, policy_selector(all_ccs[Is])>{})) : cudaSuccess));
# else // _CCCL_STD_VER >= 2020
# if _CCCL_COMPILER(GCC, <, 8)
// GCC 7 ICEs on constexpr evaluation of policy comparisons, so we skip the lowest-CC-with-same-policy optimization
// and instantiate f for each CC directly. This may increase compile time and binary size.
(...,
(device_cc == all_ccs[Is] ? (e = f(policy_getter<PolicySelector, all_ccs[Is].get()>{policy_selector}))
: cudaSuccess));
# else // _CCCL_COMPILER(GCC, <, 8)
// In C++17, we have to collapse architectures with the same policies ourselves, so we instantiate call_for_cc once
// per policy on the lowest CC which produces the same policy
using resolver_t =
lowest_cc_resolver<::cuda::std::integer_sequence<int, all_ccs[Is].get()...>, PolicySelector, Is...>;
(...,
(device_cc == all_ccs[Is]
? (e = f(policy_getter<PolicySelector, resolver_t::lowest_cc_with_same_policy[Is].get()>{policy_selector}))
: cudaSuccess));
# endif // _CCCL_COMPILER(GCC, <, 8)
# endif // _CCCL_STD_VER >= 2020
return e;
}
//! Takes a policy hub and instantiates f with the minimum possible number of nullary functor types that return a policy
//! at compile-time (if possible), and then calls the appropriate instantiation based on a runtime GPU architecture.
//! Depending on the used compiler, C++ standard, and available macros, a different number of instantiations may be
//! produced.
template <typename PolicySelector, typename F>
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t
dispatch_compute_cap(PolicySelector policy_selector, ::cuda::compute_capability device_cc, F&& f)
{
// when not using CCCL.C, policy_selector is empty since all information is contained in its type
static_assert(::cuda::std::is_empty_v<PolicySelector>);
return dispatch_to_cc_list(
policy_selector,
device_cc,
::cuda::std::forward<F>(f),
::cuda::std::make_index_sequence<::cuda::__target_compute_capabilities().size()>{});
}
#else // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC)
// if we are compiling CCCL.C with runtime policies, we cannot query the policy hub at compile time
_CCCL_EXEC_CHECK_DISABLE
template <typename PolicySelector, typename F>
_CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE cudaError_t
dispatch_compute_cap(PolicySelector policy_selector, ::cuda::compute_capability device_cc, F&& f)
{
return f([&] {
return policy_selector(device_cc);
});
}
#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC)
} // namespace detail
CUB_NAMESPACE_END

135
qwen3_6_scripts/cccl_preload/include/cub/detail/choose_offset.cuh Normal file
View File

@@ -0,0 +1,135 @@
// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/__iterator/iterator_traits.h>
#include <cuda/std/__type_traits/common_type.h>
#include <cuda/std/__type_traits/conditional.h>
#include <cuda/std/__type_traits/is_integral.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__type_traits/is_unsigned.h>
#include <cuda/std/__type_traits/remove_cv.h>
#include <cuda/std/cstdint>
#include <cuda/std/limits>
CUB_NAMESPACE_BEGIN
namespace detail
{
/**
* choose_offset checks NumItemsT, the type of the num_items parameter, and
* selects the offset type based on it.
*/
template <typename NumItemsT>
struct choose_offset
{
// NumItemsT must be an integral type (but not bool).
static_assert(::cuda::std::is_integral_v<NumItemsT>
&& !::cuda::std::is_same_v<::cuda::std::remove_cv_t<NumItemsT>, bool>,
"NumItemsT must be an integral type, but not bool");
// Unsigned integer type for global offsets.
using type = ::cuda::std::_If<(sizeof(NumItemsT) <= 4), uint32_t, unsigned long long>;
};
/**
* choose_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and
* selects the offset type based on it.
*/
template <typename NumItemsT>
using choose_offset_t = typename choose_offset<NumItemsT>::type;
/**
* promote_small_offset checks NumItemsT, the type of the num_items parameter, and
* promotes any integral type smaller than 32 bits to a signed 32-bit integer type.
*/
template <typename NumItemsT>
struct promote_small_offset
{
// NumItemsT must be an integral type (but not bool).
static_assert(::cuda::std::is_integral_v<NumItemsT>
&& !::cuda::std::is_same_v<::cuda::std::remove_cv_t<NumItemsT>, bool>,
"NumItemsT must be an integral type, but not bool");
// Unsigned integer type for global offsets.
using type = ::cuda::std::_If<(sizeof(NumItemsT) < 4), int32_t, NumItemsT>;
};
/**
* promote_small_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and
* promotes any integral type smaller than 32 bits to a signed 32-bit integer type.
*/
template <typename NumItemsT>
using promote_small_offset_t = typename promote_small_offset<NumItemsT>::type;
/**
* choose_signed_offset checks NumItemsT, the type of the num_items parameter, and
* selects the offset type to be either int32 or int64, such that the selected offset type covers the range of NumItemsT
* unless it was uint64, in which case int64 will be used.
*/
template <typename NumItemsT>
struct choose_signed_offset
{
// NumItemsT must be an integral type (but not bool).
static_assert(::cuda::std::is_integral_v<NumItemsT>
&& !::cuda::std::is_same_v<::cuda::std::remove_cv_t<NumItemsT>, bool>,
"NumItemsT must be an integral type, but not bool");
// Signed integer type for global offsets.
// uint32 -> int64, else
// LEQ 4B -> int32, else
// int64
using type = ::cuda::std::_If<(::cuda::std::is_integral_v<NumItemsT> && ::cuda::std::is_unsigned_v<NumItemsT>),
::cuda::std::int64_t,
::cuda::std::_If<(sizeof(NumItemsT) <= 4), ::cuda::std::int32_t, ::cuda::std::int64_t>>;
/**
* Checks if the given num_items can be covered by the selected offset type. If not, returns cudaErrorInvalidValue,
* otherwise returns cudaSuccess.
*/
static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t is_exceeding_offset_type(NumItemsT num_items)
{
_CCCL_DIAG_PUSH
_CCCL_DIAG_SUPPRESS_MSVC(4127) /* conditional expression is constant */
if (sizeof(NumItemsT) >= 8 && num_items > static_cast<NumItemsT>(::cuda::std::numeric_limits<type>::max()))
{
return cudaErrorInvalidValue;
}
_CCCL_DIAG_POP
return cudaSuccess;
}
};
/**
* choose_signed_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and
* selects the corresponding signed offset type based on it.
*/
template <typename NumItemsT>
using choose_signed_offset_t = typename choose_signed_offset<NumItemsT>::type;
/**
* common_iterator_value sets member type to the common_type of
* value_type for all argument types. used to get OffsetT in
* DeviceSegmentedReduce.
*/
template <typename... Iter>
struct common_iterator_value
{
using type = ::cuda::std::common_type_t<::cuda::std::__iter_value_type<Iter>...>;
};
template <typename... Iter>
using common_iterator_value_t = typename common_iterator_value<Iter...>::type;
} // namespace detail
CUB_NAMESPACE_END

82
qwen3_6_scripts/cccl_preload/include/cub/detail/deferred_parameter.cuh Normal file
View File

@@ -0,0 +1,82 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/choose_offset.cuh>
#include <cuda/__argument/argument.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__utility/declval.h>
CUB_NAMESPACE_BEGIN
namespace detail
{
#if !_CCCL_COMPILER(NVRTC)
// Preserve deferred problem sizes for dispatch and canonicalize immediate values to CUB's offset type.
template <typename NumItemsT>
[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto make_num_items_dispatch_arg(NumItemsT num_items) noexcept
{
using args_traits_t = ::cuda::args::__traits<NumItemsT>;
if constexpr (args_traits_t::is_deferred)
{
return num_items;
}
else
{
using offset_t = choose_offset_t<typename args_traits_t::element_type>;
return static_cast<offset_t>(::cuda::args::__unwrap(num_items));
}
}
// Forms a kernel parameter from a single-value argument without reading a deferred source.
// Immediate values are converted to TargetT. Deferred arguments are unwrapped to their source, erasing bounds from
// the kernel type and payload.
template <typename TargetT, typename ParameterT>
[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE constexpr auto parameter_from_host(ParameterT parameter) noexcept
{
using args_traits_t = ::cuda::args::__traits<ParameterT>;
static_assert(args_traits_t::is_single_value, "parameter must contain a single value");
if constexpr (args_traits_t::is_deferred)
{
return ::cuda::args::__unwrap(parameter);
}
else
{
return static_cast<TargetT>(::cuda::args::__unwrap(parameter));
}
}
template <typename TargetT, typename ParameterT>
using parameter_from_host_t = decltype(parameter_from_host<TargetT>(::cuda::std::declval<ParameterT>()));
#endif // !_CCCL_COMPILER(NVRTC)
// Forms a value from a kernel parameter, reading element zero when the parameter is a deferred source.
template <typename TargetT, typename ParameterT>
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE TargetT parameter_from_device(ParameterT parameter) noexcept
{
if constexpr (::cuda::std::is_same_v<ParameterT, TargetT>)
{
return parameter;
}
else
{
return static_cast<TargetT>(parameter[0]);
}
}
} // namespace detail
CUB_NAMESPACE_END

150
qwen3_6_scripts/cccl_preload/include/cub/detail/delay_constructor.cuh Normal file
View File

@@ -0,0 +1,150 @@
// SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/agent/single_pass_scan_operators.cuh>
#include <cub/device/dispatch/tuning/common.cuh>
#include <cuda/std/__concepts/same_as.h>
CUB_NAMESPACE_BEGIN
namespace detail
{
template <typename DelayConstructor>
inline constexpr auto lookback_delay_policy_from_type = 0;
template <unsigned int L2WriteLatency>
inline constexpr auto lookback_delay_policy_from_type<no_delay_constructor_t<L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::no_delay, 0, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto lookback_delay_policy_from_type<fixed_delay_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::fixed_delay, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto lookback_delay_policy_from_type<exponential_backoff_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto lookback_delay_policy_from_type<exponential_backoff_jitter_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff_jitter, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto
lookback_delay_policy_from_type<exponential_backoff_jitter_window_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff_jitter_window, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto
lookback_delay_policy_from_type<exponential_backon_jitter_window_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon_jitter_window, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto lookback_delay_policy_from_type<exponential_backon_jitter_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon_jitter, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency>
inline constexpr auto lookback_delay_policy_from_type<exponential_backon_constructor_t<Delay, L2WriteLatency>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon, Delay, L2WriteLatency};
template <unsigned int Delay, unsigned int L2WriteLatency, unsigned int GridThreshold>
inline constexpr auto
lookback_delay_policy_from_type<reduce_by_key_delay_constructor_t<Delay, L2WriteLatency, GridThreshold>> =
LookbackDelayPolicy{LookbackDelayAlgorithm::__reduce_by_key, Delay, L2WriteLatency};
template <LookbackDelayAlgorithm Kind, unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for;
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::no_delay, Delay, L2WriteLatency>
{
using type = no_delay_constructor_t<L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::fixed_delay, Delay, L2WriteLatency>
{
using type = fixed_delay_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backoff, Delay, L2WriteLatency>
{
using type = exponential_backoff_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backoff_jitter, Delay, L2WriteLatency>
{
using type = exponential_backoff_jitter_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backoff_jitter_window, Delay, L2WriteLatency>
{
using type = exponential_backoff_jitter_window_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backon_jitter_window, Delay, L2WriteLatency>
{
using type = exponential_backon_jitter_window_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backon_jitter, Delay, L2WriteLatency>
{
using type = exponential_backon_jitter_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backon, Delay, L2WriteLatency>
{
using type = exponential_backon_constructor_t<Delay, L2WriteLatency>;
};
template <unsigned int Delay, unsigned int L2WriteLatency>
struct delay_constructor_for<LookbackDelayAlgorithm::__reduce_by_key, Delay, L2WriteLatency>
{
using type = reduce_by_key_delay_constructor_t<Delay, L2WriteLatency>;
};
template <LookbackDelayAlgorithm Kind, unsigned int Delay, unsigned int L2WriteLatency>
using delay_constructor_t = typename delay_constructor_for<Kind, Delay, L2WriteLatency>::type;
_CCCL_HOST_DEVICE_API constexpr auto default_delay_constructor_policy(bool is_primitive_or_trivially_copyable)
{
if (is_primitive_or_trivially_copyable)
{
return LookbackDelayPolicy{LookbackDelayAlgorithm::fixed_delay, 350, 450};
}
return LookbackDelayPolicy{LookbackDelayAlgorithm::no_delay, 0, 450};
}
_CCCL_HOST_DEVICE_API constexpr auto default_reduce_by_key_delay_constructor_policy(
int key_size,
int value_size,
bool key_is_primitive_or_trivially_copyable,
bool value_is_primitive_or_trivially_copyable)
{
if (value_is_primitive_or_trivially_copyable && (value_size + key_size < 16))
{
return LookbackDelayPolicy{LookbackDelayAlgorithm::__reduce_by_key, 350, 450};
}
return default_delay_constructor_policy(key_is_primitive_or_trivially_copyable);
}
} // namespace detail
CUB_NAMESPACE_END

94
qwen3_6_scripts/cccl_preload/include/cub/detail/device_double_buffer.cuh Normal file
View File

@@ -0,0 +1,94 @@
/*
* Copyright 2021 NVIDIA Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_namespace.cuh>
CUB_NAMESPACE_BEGIN
namespace detail
{
/**
* @brief It's a double-buffer storage wrapper for multi-pass stream
* transformations that require more than one storage array for
* streaming intermediate results back and forth.
*
* Many multi-pass computations require a pair of "ping-pong" storage buffers
* (e.g., one for reading from and the other for writing to, and then
* vice-versa for the subsequent pass). This structure wraps a set of device
* buffers.
*
* Unlike `cub::DoubleBuffer` this class doesn't provide a "selector" member
* to track which buffer is "current". The main reason for this class existence
* is the performance difference. Since `cub::DoubleBuffer` relies on the
* runtime variable to index pointers arrays, they are placed in the local
* memory instead of registers. Local memory accesses significantly affect
* performance. On the contrary, this class swaps pointer, so all operations
* can be performed in registers.
*/
template <typename T>
class device_double_buffer
{
/// Pair of device buffer pointers
T* m_current_buffer{};
T* m_alternate_buffer{};
public:
/**
* @param d_current
* The currently valid buffer
*
* @param d_alternate
* Alternate storage buffer of the same size as @p d_current
*/
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE device_double_buffer(T* current, T* alternate)
: m_current_buffer(current)
, m_alternate_buffer(alternate)
{}
/// \brief Return pointer to the currently valid buffer
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE T* current() const
{
return m_current_buffer;
}
/// \brief Return pointer to the currently invalid buffer
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE T* alternate() const
{
return m_alternate_buffer;
}
_CCCL_HOST_DEVICE void swap()
{
T* tmp = m_current_buffer;
m_current_buffer = m_alternate_buffer;
m_alternate_buffer = tmp;
}
};
} // namespace detail
CUB_NAMESPACE_END

97
qwen3_6_scripts/cccl_preload/include/cub/detail/device_memory_resource.cuh Normal file
View File

@@ -0,0 +1,97 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_device.cuh>
#include <cuda/__runtime/api_wrapper.h>
#include <cuda/__stream/stream_ref.h>
#include <cuda/std/__exception/terminate.h>
#include <cuda/std/cstdint>
#include <nv/target>
CUB_NAMESPACE_BEGIN
namespace detail
{
// TODO(gevtushenko/srinivasyadav18): move cudax `device_memory_resource` to `cuda::__device_memory_resource` and remove
// this implementation
struct device_memory_resource
{
CUB_RUNTIME_FUNCTION void* allocate(size_t bytes, size_t /* alignment */)
{
void* ptr{nullptr};
NV_IF_ELSE_TARGET(
NV_IS_HOST,
(_CCCL_TRY_CUDA_API(::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, NULL);),
({
_CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code.");
::cuda::std::terminate();
}));
_CCCL_ASSERT(ptr != nullptr, "allocate failed to allocate with cudaMallocAsync");
return ptr;
}
CUB_RUNTIME_FUNCTION void deallocate(void* ptr, size_t /* bytes */)
{
NV_IF_ELSE_TARGET( //
NV_IS_HOST,
(_CCCL_TRY_CUDA_API(::cudaFree, "deallocate failed", ptr);),
({
_CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code.");
::cuda::std::terminate();
}));
}
CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes, size_t /* alignment */)
{
return allocate(stream, bytes);
}
CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes)
{
void* ptr{nullptr};
NV_IF_ELSE_TARGET( //
NV_IS_HOST,
({
_CCCL_TRY_CUDA_API(
::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, stream.get());
}),
({
_CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code.");
::cuda::std::terminate();
}));
return ptr;
}
CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t bytes, size_t /* alignment */)
{
deallocate(stream, ptr, bytes);
}
CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t /* bytes */)
{
NV_IF_ELSE_TARGET( //
NV_IS_HOST,
(_CCCL_TRY_CUDA_API(::cudaFreeAsync, "deallocate failed", ptr, stream.get());),
({
_CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code.");
::cuda::std::terminate();
}));
}
};
} // namespace detail
CUB_NAMESPACE_END

137
qwen3_6_scripts/cccl_preload/include/cub/detail/env_dispatch.cuh Normal file
View File

@@ -0,0 +1,137 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/device_memory_resource.cuh>
#include <cub/detail/temporary_storage.cuh>
#include <cuda/__execution/tune.h>
#include <cuda/__functional/call_or.h>
#include <cuda/__memory_resource/get_memory_resource.h>
#include <cuda/__stream/get_stream.h>
#include <cuda/std/__execution/env.h>
CUB_NAMESPACE_BEGIN
namespace detail
{
//! @cond
//! Generic environment-based algorithm dispatch wrapper
//!
//! Handles common boilerplate for all env-based algorithms:
//! - Query stream, memory resource, and tuning from environment
//! - Two-phase call (query temp storage size, then execute)
//! - Temporary storage allocation/deallocation
//! - Memory resource querying from environment
//!
//! @param env The execution environment
//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified
template <typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
// Query stream from environment
auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
// Query memory resource from environment
auto mr = ::cuda::__call_or(::cuda::mr::__get_memory_resource, detail::device_memory_resource{}, env);
// Query tuning from environment
const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env);
void* d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;
// Phase 1: Query temporary storage size
if (const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get()))
{
return error;
}
// Allocate temporary storage
if (const auto error = CubDebug(detail::temporary_storage::allocate(stream, d_temp_storage, temp_storage_bytes, mr)))
{
return error;
}
// Phase 2: Execute algorithm
const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get());
// Deallocate temporary storage (always attempt, even on error)
const auto deallocate_error =
CubDebug(detail::temporary_storage::deallocate(stream, d_temp_storage, temp_storage_bytes, mr));
// Algorithm error takes precedence over deallocation error
return (error != cudaSuccess) ? error : deallocate_error;
}
//! @endcond
template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t
dispatch_with_env_and_tuning(const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
return detail::dispatch_with_env(
env,
[&algorithm_callable](
[[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{}));
using policy_selector =
::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
});
}
//! @cond
//! Generic environment-based algorithm dispatch wrapper
//!
//! Handles common boilerplate for env-based algorithms with user provided memory:
//! - Query stream, and tuning from environment
//! - Single-phase call passing user provided memory and size
//!
//! @param env The execution environment
//! @param[in] d_temp_storage @devicestorage
//! @param[in,out] temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation
//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified
template <typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(
void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
// Query stream from environment
auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
// Query tuning from environment
const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env);
return algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get());
}
//! @endcond
template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env_and_tuning(
void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
return detail::dispatch_with_env(
d_temp_storage,
temp_storage_bytes,
env,
[&algorithm_callable](
[[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{}));
using policy_selector =
::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
});
}
//! @endcond
} // namespace detail
CUB_NAMESPACE_END

226
qwen3_6_scripts/cccl_preload/include/cub/detail/fast_modulo_division.cuh Normal file
View File

@@ -0,0 +1,226 @@
// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/type_traits.cuh> // implicit_prom_t
#include <cub/util_type.cuh> // _CCCL_HAS_INT128()
#include <cuda/__cmath/ceil_div.h>
#include <cuda/__cmath/pow2.h>
#include <cuda/std/__bit/integral.h>
#include <cuda/std/__type_traits/conditional.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__type_traits/integral_constant.h>
#include <cuda/std/__type_traits/is_integral.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__type_traits/is_signed.h>
#include <cuda/std/__type_traits/make_unsigned.h>
#include <cuda/std/climits> // CHAR_BIT
#include <cuda/std/cstdint> // uint64_t
#include <cuda/std/limits>
#if defined(CCCL_ENABLE_DEVICE_ASSERTIONS)
_CCCL_BEGIN_NV_DIAG_SUPPRESS(186) // pointless comparison of unsigned integer with zero
#endif // CCCL_ENABLE_DEVICE_ASSERTIONS
CUB_NAMESPACE_BEGIN
namespace detail
{
/***********************************************************************************************************************
* larger_unsigned_type
**********************************************************************************************************************/
template <typename T, typename = void>
struct larger_unsigned_type
{
using type = void;
};
template <typename T>
struct larger_unsigned_type<T, ::cuda::std::enable_if_t<(sizeof(T) < 4)>>
{
using type = ::cuda::std::uint32_t;
};
template <typename T>
struct larger_unsigned_type<T, ::cuda::std::enable_if_t<(sizeof(T) == 4)>>
{
using type = ::cuda::std::uint64_t;
};
#if _CCCL_HAS_INT128()
template <typename T>
struct larger_unsigned_type<T, ::cuda::std::enable_if_t<(sizeof(T) == 8)>>
{
using type = __uint128_t;
};
#endif // _CCCL_HAS_INT128()
template <typename T>
using larger_unsigned_type_t = typename larger_unsigned_type<T>::type;
template <typename T>
using unsigned_implicit_prom_t = ::cuda::std::make_unsigned_t<implicit_prom_t<T>>;
template <typename T>
using supported_integral =
::cuda::std::bool_constant<::cuda::std::is_integral_v<T> && !::cuda::std::is_same_v<T, bool> && (sizeof(T) <= 8)>;
/***********************************************************************************************************************
* Extract higher bits after multiplication
**********************************************************************************************************************/
template <typename DivisorType, typename T, typename R>
[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE unsigned_implicit_prom_t<DivisorType>
multiply_extract_higher_bits(T value, R multiplier)
{
static_assert(supported_integral<T>::value, "unsupported type");
static_assert(supported_integral<R>::value, "unsupported type");
if constexpr (::cuda::std::is_signed_v<T>)
{
_CCCL_ASSERT(value >= 0, "value must be non-negative");
}
if constexpr (::cuda::std::is_signed_v<R>)
{
_CCCL_ASSERT(multiplier >= 0, "multiplier must be non-negative");
}
static constexpr int NumBits = sizeof(DivisorType) * CHAR_BIT;
using unsigned_t = unsigned_implicit_prom_t<DivisorType>;
using larger_t = larger_unsigned_type_t<DivisorType>;
// clang-format off
NV_IF_ELSE_TARGET(
NV_IS_HOST,
(return static_cast<unsigned_t>((static_cast<larger_t>(value) * multiplier) >> NumBits);),
({return (sizeof(T) == 8)
? static_cast<unsigned_t>(__umul64hi(value, multiplier))
: static_cast<unsigned_t>((static_cast<larger_t>(value) * multiplier) >> NumBits);}));
// clang-format on
}
/***********************************************************************************************************************
* Fast Modulo/Division based on Precomputation
**********************************************************************************************************************/
_CCCL_DIAG_PUSH
_CCCL_DIAG_SUPPRESS_MSVC(4127) /* conditional expression is constant */
template <typename T1>
class fast_div_mod
{
static_assert(supported_integral<T1>::value, "unsupported type");
// uint16_t is a special case that would requires complex logic. Workaround: convert to int
using T = ::cuda::std::conditional_t<::cuda::std::is_same_v<T1, ::cuda::std::uint16_t>, int, T1>;
using unsigned_t = unsigned_implicit_prom_t<T>;
public:
template <typename R>
struct result
{
using common_t = decltype(R{} / T{});
common_t quotient;
common_t remainder;
};
fast_div_mod() = delete;
_CCCL_HOST_DEVICE explicit fast_div_mod(T divisor) noexcept
: _divisor{static_cast<unsigned_t>(divisor)}
{
using larger_t = larger_unsigned_type_t<T>;
_CCCL_ASSERT(divisor > 0, "divisor must be positive");
auto udivisor = static_cast<unsigned_t>(divisor);
// the following branches are needed to avoid negative shift
if (::cuda::is_power_of_two(udivisor))
{
_shift_right = ::cuda::std::bit_width(udivisor) - 1;
return;
}
else if (sizeof(T) == 8 && divisor == 3)
{
return;
}
constexpr int BitSize = sizeof(T) * CHAR_BIT; // 32
constexpr int BitOffset = BitSize / 16; // 2
int num_bits = ::cuda::std::bit_width(udivisor) + 1;
_CCCL_ASSERT(static_cast<size_t>(num_bits + BitSize - BitOffset) < sizeof(larger_t) * CHAR_BIT, "overflow error");
// without explicit power-of-two check, num_bits needs to replace +1 with !::cuda::is_power_of_two(udivisor)
_multiplier = static_cast<unsigned_t>(::cuda::ceil_div(larger_t{1} << (num_bits + BitSize - BitOffset), //
static_cast<larger_t>(divisor)));
_shift_right = num_bits - BitOffset;
_CCCL_ASSERT(_multiplier != 0, "overflow error");
}
fast_div_mod(const fast_div_mod&) noexcept = default;
fast_div_mod(fast_div_mod&&) noexcept = default;
template <typename R>
[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE result<R> operator()(R dividend) const noexcept
{
static_assert(supported_integral<R>::value, "unsupported type");
using common_t = decltype(R{} / T{});
using ucommon_t = ::cuda::std::make_unsigned_t<common_t>;
using result_t = result<R>;
_CCCL_ASSERT(dividend >= 0, "divisor must be non-negative");
auto udividend = static_cast<ucommon_t>(dividend);
if (_divisor == 1)
{
return result_t{static_cast<common_t>(dividend), common_t{}};
}
else if (_divisor > unsigned_t{::cuda::std::numeric_limits<T>::max() / 2})
{
auto quotient = udividend >= static_cast<ucommon_t>(_divisor);
return result_t{static_cast<common_t>(quotient), static_cast<common_t>(udividend - (quotient * _divisor))};
}
else if (sizeof(T) == 8 && _divisor == 3)
{
return result_t{static_cast<common_t>(udividend / 3), static_cast<common_t>(udividend % 3)};
}
auto higher_bits = (_multiplier == 0) ? udividend : multiply_extract_higher_bits<T>(dividend, _multiplier);
auto quotient = higher_bits >> _shift_right;
auto remainder = udividend - (quotient * _divisor);
_CCCL_ASSERT(quotient == udividend / _divisor, "wrong quotient");
_CCCL_ASSERT(remainder < (ucommon_t) _divisor, "remainder out of range");
return result_t{static_cast<common_t>(quotient), static_cast<common_t>(remainder)};
}
template <typename R>
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE friend implicit_prom_t<T> operator/(R dividend, fast_div_mod div) noexcept
{
return div(dividend).quotient;
}
template <typename R>
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE friend implicit_prom_t<T> operator%(R dividend, fast_div_mod div) noexcept
{
return div(dividend).remainder;
}
private:
unsigned_t _divisor = 1;
unsigned_t _multiplier = 0;
unsigned _shift_right = 0;
};
_CCCL_DIAG_POP
} // namespace detail
CUB_NAMESPACE_END
#if defined(CCCL_ENABLE_DEVICE_ASSERTIONS)
_CCCL_END_NV_DIAG_SUPPRESS()
#endif // CCCL_ENABLE_DEVICE_ASSERTIONS

147
qwen3_6_scripts/cccl_preload/include/cub/detail/launcher/cuda_runtime.cuh Normal file
View File

@@ -0,0 +1,147 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_device.cuh>
#include <thrust/system/cuda/detail/core/triple_chevron_launch.h>
#include <cuda/__device/compute_capability.h>
CUB_NAMESPACE_BEGIN
namespace detail
{
struct TripleChevronFactory
{
CUB_RUNTIME_FUNCTION void __assert_pdl_allowed(bool dependent_launch) const
{
if (dependent_launch)
{
[[maybe_unused]] int sm_version = 0;
_CCCL_ASSERT(SmVersion(sm_version) == cudaSuccess, "Failed to query SM compute capability");
if (sm_version >= 900)
{
[[maybe_unused]] ::cuda::compute_capability cc;
_CCCL_ASSERT(PtxComputeCap(cc) == cudaSuccess, "Failed to query PTX compute capability");
_CCCL_ASSERT((cc >= ::cuda::compute_capability{9, 0}),
"Enabling PDL for a kernel launch requires CC 9.0+ PTX/SASS when running on SM90+");
}
}
}
CUB_RUNTIME_FUNCTION THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron operator()(
dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, bool dependent_launch = false) const
{
__assert_pdl_allowed(dependent_launch);
return THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid, block, shared_mem, stream, dependent_launch);
}
template <class T = void>
CUB_RUNTIME_FUNCTION ::cudaError_t PtxVersion(int& version)
{
return cub::PtxVersion<T>(version);
}
template <class T = void>
CUB_RUNTIME_FUNCTION ::cudaError_t PtxComputeCap(::cuda::compute_capability& cc) const
{
return ptx_compute_cap<T>(cc);
}
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MultiProcessorCount(int& sm_count) const
{
int device_ordinal;
::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal));
if (::cudaSuccess != error)
{
return error;
}
// Get SM count
return ::cudaDeviceGetAttribute(&sm_count, ::cudaDevAttrMultiProcessorCount, device_ordinal);
}
template <typename Kernel>
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
MaxSmOccupancy(int& sm_occupancy, Kernel kernel_ptr, int block_size, int dynamic_smem_bytes = 0)
{
return ::cudaOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_ptr, block_size, dynamic_smem_bytes);
}
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const
{
int device_ordinal;
::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal));
if (::cudaSuccess != error)
{
return error;
}
// Get max grid dimension
return ::cudaDeviceGetAttribute(&max_grid_dim_x, ::cudaDevAttrMaxGridDimX, device_ordinal);
}
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
MemsetAsync(void* dst, unsigned char value, size_t num_bytes, ::cudaStream_t stream) const
{
return ::cudaMemsetAsync(dst, value, num_bytes, stream);
}
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
MemcpyAsync(void* dst, const void* src, size_t num_bytes, ::cudaMemcpyKind kind, ::cudaStream_t stream) const
{
return ::cudaMemcpyAsync(dst, src, num_bytes, kind, stream);
}
// TODO(bgruber): this is very similar to thrust::cuda_cub::core::get_max_shared_memory_per_block. We should unify
// this.
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION cudaError_t MaxSharedMemory(int& max_shared_memory) const
{
int device = 0;
auto error = CubDebug(cudaGetDevice(&device));
if (error != cudaSuccess)
{
return error;
}
return cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlock, device);
}
template <typename Kernel>
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
max_dynamic_smem_size_for(int& max_dynamic_smem_size, [[maybe_unused]] Kernel kernel_ptr)
{
NV_IF_ELSE_TARGET(NV_IS_HOST, //
({ return MaxPotentialDynamicSmemBytes(max_dynamic_smem_size, kernel_ptr); }),
({
::cudaFuncAttributes func_attrs{};
if (const auto error = CubDebug(::cudaFuncGetAttributes(&func_attrs, kernel_ptr)))
{
return error;
}
max_dynamic_smem_size = func_attrs.maxDynamicSharedSizeBytes;
return cudaSuccess;
}))
}
template <typename Kernel>
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t set_max_dynamic_smem_size_for(Kernel kernel_ptr, int smem_size)
{
return CubDebug(::cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
}
};
} // namespace detail
CUB_NAMESPACE_END

113
qwen3_6_scripts/cccl_preload/include/cub/detail/mdspan_utils.cuh Normal file
View File

@@ -0,0 +1,113 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/fast_modulo_division.cuh> // fast_div_mod
#include <cuda/std/__mdspan/extents.h>
#include <cuda/std/__type_traits/make_unsigned.h>
#include <cuda/std/__utility/integer_sequence.h>
#include <cuda/std/array>
#include <cuda/std/cstddef>
CUB_NAMESPACE_BEGIN
namespace detail
{
_CCCL_DIAG_PUSH
_CCCL_DIAG_SUPPRESS_MSVC(4702) // unreachable code (even if there are no branches!)
// Compute the submdspan size of a given rank
template <typename IndexType, size_t... Extents>
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::make_unsigned_t<IndexType>
size_range(const ::cuda::std::extents<IndexType, Extents...>& ext, int start, int end)
{
_CCCL_ASSERT(start >= 0 && end <= static_cast<int>(ext.rank()), "invalid start or end");
::cuda::std::make_unsigned_t<IndexType> s = 1;
for (auto i = start; i < end; i++)
{
s *= ext.extent(i);
}
return s;
}
_CCCL_DIAG_POP // MSVC(4702)
template <typename IndexType, size_t... Extents>
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::make_unsigned_t<IndexType>
size(const ::cuda::std::extents<IndexType, Extents...>& ext)
{
return cub::detail::size_range(ext, 0, static_cast<int>(ext.rank()));
}
template <bool IsLayoutRight, int Position, typename IndexType, size_t... E>
[[nodiscard]] _CCCL_HOST_DEVICE_API auto sub_size_fast_div_mod_impl(const ::cuda::std::extents<IndexType, E...>& ext)
{
using fast_mod_div_t = fast_div_mod<IndexType>;
constexpr auto start = IsLayoutRight ? Position + 1 : 0;
constexpr auto end = IsLayoutRight ? sizeof...(E) : Position;
return fast_mod_div_t(cub::detail::size_range(ext, start, end));
}
// precompute modulo/division for each submdspan size (by rank)
template <bool IsLayoutRight, typename IndexType, size_t... E, size_t... Positions>
[[nodiscard]] _CCCL_HOST_DEVICE_API auto
sub_sizes_fast_div_mod(const ::cuda::std::extents<IndexType, E...>& ext, ::cuda::std::index_sequence<Positions...> = {})
{
using fast_mod_div_t = fast_div_mod<IndexType>;
using array_t = ::cuda::std::array<fast_mod_div_t, sizeof...(Positions)>;
return array_t{cub::detail::sub_size_fast_div_mod_impl<IsLayoutRight, Positions>(ext)...};
}
// precompute modulo/division for each mdspan extent
template <typename IndexType, size_t... E, size_t... Positions>
[[nodiscard]] _CCCL_HOST_DEVICE_API auto
extents_fast_div_mod(const ::cuda::std::extents<IndexType, E...>& ext, ::cuda::std::index_sequence<Positions...> = {})
{
using fast_mod_div_t = fast_div_mod<IndexType>;
using array_t = ::cuda::std::array<fast_mod_div_t, sizeof...(Positions)>;
return array_t{fast_mod_div_t(ext.extent(Positions))...};
}
// GCC <= 9 constexpr workaround: Extent must be passed as type only, even const Extent& doesn't work
template <typename Extents>
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr bool are_extents_in_range_static(int start, int end)
{
for (auto i = start; i < end; i++)
{
if (Extents::static_extent(i) == ::cuda::std::dynamic_extent)
{
return false;
}
}
return true;
}
template <typename MappingTypeLhs, typename MappingTypeRhs>
[[nodiscard]] _CCCL_HOST_DEVICE_API bool
have_same_strides(const MappingTypeLhs& mapping_lhs, const MappingTypeRhs& mapping_rhs)
{
auto extents_lhs = mapping_lhs.extents();
auto extents_rhs = mapping_rhs.extents();
_CCCL_ASSERT(extents_lhs.rank() == extents_rhs.rank(), "extents must have the same rank");
for (size_t i = 0; i < extents_lhs.rank(); i++)
{
if (mapping_lhs.stride(i) != mapping_rhs.stride(i))
{
return false;
}
}
return true;
}
} // namespace detail
CUB_NAMESPACE_END

691
qwen3_6_scripts/cccl_preload/include/cub/detail/rfa.cuh Normal file
View File

@@ -0,0 +1,691 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/__algorithm/max.h>
#include <cuda/std/__algorithm/min.h>
#include <cuda/std/__bit/bit_cast.h>
#include <cuda/std/__cmath/exponential_functions.h>
#include <cuda/std/__cmath/isinf.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__type_traits/is_arithmetic.h>
#include <cuda/std/__type_traits/is_floating_point.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/array>
#include <cuda/std/climits>
CUB_NAMESPACE_BEGIN
namespace detail::rfa
{
// jump table for indexing into data
inline constexpr int cub_rfa_max_jump = 5;
static_assert(cub_rfa_max_jump <= 5, "cub_rfa_max_jump must be less than or equal to 5");
template <typename FType, int Len>
static _CCCL_DEVICE FType* get_shared_bin_array()
{
static __shared__ FType bin_computed_array[Len];
return bin_computed_array;
}
//! Class to hold a reproducible summation of the numbers passed to it
//!
//! @param FType Floating-point data type; either `float` or `double
//! @param Fold Number of collectors in the binned number (K-fold), used for reproducible summation. Defaults to 3.
template <class FType, int Fold = 3, ::cuda::std::enable_if_t<::cuda::std::is_floating_point_v<FType>>* = nullptr>
class alignas(2 * sizeof(FType)) ReproducibleFloatingAccumulator
{
public:
using ftype = FType;
private:
::cuda::std::array<ftype, 2 * Fold> data{};
/// Floating-point precision bin width
static constexpr int bin_width = ::cuda::std::is_same_v<ftype, double> ? 40 : 13;
static constexpr int min_exp = ::cuda::std::numeric_limits<ftype>::min_exponent;
static constexpr int max_exp = ::cuda::std::numeric_limits<ftype>::max_exponent;
static constexpr int mant_dig = ::cuda::std::numeric_limits<ftype>::digits;
public:
/// Binned floating-point maximum index
static constexpr int max_index = ((max_exp - min_exp + mant_dig - 1) / bin_width) - 1;
// The maximum floating-point fold supported by the library
static constexpr auto max_fold = max_index + 1;
_CCCL_DEVICE static ftype initialize_bin(int index) noexcept
{
if (index == 0)
{
if constexpr (::cuda::std::is_same_v<ftype, float>)
{
return ::cuda::std::ldexp(0.75, max_exp);
}
else
{
return 2.0 * ::cuda::std::ldexp(0.75, max_exp - 1);
}
}
if (index > 0 && index <= max_index)
{
return ::cuda::std::ldexp(0.75, max_exp + mant_dig - bin_width + 1 - index * bin_width);
}
else
{
return ::cuda::std::ldexp(0.75, max_exp + mant_dig - bin_width + 1 - max_index * bin_width);
}
}
private:
/// Binned floating-point compression factor
/// This factor is used to scale down inputs before deposition into the bin of
/// highest index
static constexpr auto compression = 1.0 / (1 << (mant_dig - bin_width + 1));
/// Binned double precision expansion factor
/// This factor is used to scale up inputs after deposition into the bin of
/// highest index
static constexpr auto expansion = 1.0 * (1 << (mant_dig - bin_width + 1));
static constexpr auto exp_bias = max_exp - 2;
/// Return a binned floating-point bin
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static ftype binned_bins(int index)
{
ftype* bins = get_shared_bin_array<ftype, max_index + max_fold>();
return bins[index];
}
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint32_t& get_bit_representation(float& x) noexcept
{
return *reinterpret_cast<uint32_t*>(&x);
}
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint64_t& get_bit_representation(double& x) noexcept
{
return *reinterpret_cast<uint64_t*>(&x);
}
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint32_t get_bit_representation(const float& x) noexcept
{
return ::cuda::std::bit_cast<uint32_t>(x);
}
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint64_t get_bit_representation(const double& x) noexcept
{
return ::cuda::std::bit_cast<uint64_t>(x);
}
/// Return primary vector value const ref
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE const ftype& primary(int i) const noexcept
{
if constexpr (Fold <= cub_rfa_max_jump)
{
switch (i)
{
case 0:
if constexpr (Fold >= 1)
{
return data[0];
}
[[fallthrough]];
case 1:
if constexpr (Fold >= 2)
{
return data[1];
}
[[fallthrough]];
case 2:
if constexpr (Fold >= 3)
{
return data[2];
}
[[fallthrough]];
case 3:
if constexpr (Fold >= 4)
{
return data[3];
}
[[fallthrough]];
case 4:
if constexpr (Fold >= 5)
{
return data[4];
}
[[fallthrough]];
default:
return data[Fold - 1];
}
}
else
{
return data[i];
}
}
/// Return carry vector value const ref
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE const ftype& carry(int i) const noexcept
{
if (Fold <= cub_rfa_max_jump)
{
switch (i)
{
case 0:
if (Fold >= 1)
{
return data[Fold + 0];
}
[[fallthrough]];
case 1:
if (Fold >= 2)
{
return data[Fold + 1];
}
[[fallthrough]];
case 2:
if (Fold >= 3)
{
return data[Fold + 2];
}
[[fallthrough]];
case 3:
if (Fold >= 4)
{
return data[Fold + 3];
}
[[fallthrough]];
case 4:
if (Fold >= 5)
{
return data[Fold + 4];
}
[[fallthrough]];
default:
return data[2 * Fold - 1];
}
}
else
{
return data[Fold + i]; // NOLINT(bugprone-misplaced-widening-cast)
}
}
/// Return primary vector value ref
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ftype& primary(int i) noexcept
{
const auto& c = *this;
return const_cast<ftype&>(c.primary(i));
}
/// Return carry vector value ref
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ftype& carry(int i) noexcept
{
const auto& c = *this;
return const_cast<ftype&>(c.carry(i));
}
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static int exp_val(const ftype x) noexcept
{
const auto bits = get_bit_representation(x);
return (bits >> (mant_dig - 1)) & (2 * max_exp - 1);
}
/// Get index of float-point precision
/// The index of a non-binned type is the smallest index a binned type would
/// need to have to sum it reproducibly. Higher indices correspond to smaller
/// bins.
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static int binned_dindex(const ftype x)
{
int exp = exp_val(x);
if (exp != 0)
{
return ((max_exp + exp_bias) - exp) / bin_width;
}
if (x == 0.0)
{
return max_index;
}
else
{
(void) ::cuda::std::frexpf(x, &exp);
return (::cuda::std::min) ((max_exp - exp) / bin_width, +max_index);
}
}
/// Get index of manually specified binned double precision
/// The index of a binned type is the bin that it corresponds to. Higher
/// indices correspond to smaller bins.
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE int binned_index() const
{
return ((max_exp + mant_dig - bin_width + 1 + exp_bias) - exp_val(primary(0))) / bin_width;
}
/// Check if index of manually specified binned floating-point is 0
/// A quick check to determine if the index is 0
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool is_binned_index_zero() const
{
return exp_val(primary(0)) == max_exp + exp_bias;
}
//! Update manually specified binned fp with a scalar (X -> Y)
//!
//! This method updates the binned fp to an index suitable for adding numbers
//! with absolute value less than @p max_abs_val
_CCCL_DEVICE void binned_update(const ftype max_abs_val)
{
int X_index = binned_dindex(max_abs_val);
int shift = binned_index() - X_index;
if (shift > 0)
{
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = Fold - 1; i >= 1; i--)
{
if (i < shift)
{
break;
}
primary(i) = primary((i - shift));
carry(i) = carry((i - shift));
}
_CCCL_PRAGMA_UNROLL_FULL()
for (int j = 0; j < Fold; j++)
{
if (j >= shift)
{
break;
}
primary(j) = binned_bins(j + X_index);
carry(j) = 0.0;
}
}
}
//! Add scalar @p X to suitably binned manually specified binned fp (Y += X)
//!
//! Performs the operation Y += X on an binned type Y where the index of Y is
//! larger than the index of @p X
_CCCL_DEVICE void binned_deposit(const ftype X)
{
ftype M;
ftype x = X;
if (is_binned_index_zero())
{
M = primary(0);
ftype qd = x * compression;
auto& ql = get_bit_representation(qd);
ql |= 1;
qd += M;
primary(0) = qd;
M -= qd;
M *= expansion * 0.5;
x += M;
x += M;
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 1; i < Fold - 1; i++)
{
M = primary(i);
qd = x;
ql |= 1;
qd += M;
primary(i) = qd;
M -= qd;
x += M;
}
qd = x;
ql |= 1;
primary((Fold - 1)) += qd;
}
else
{
ftype qd = x;
auto& ql = get_bit_representation(qd);
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < Fold - 1; i++)
{
M = primary(i);
qd = x;
ql |= 1;
qd += M;
primary(i) = qd;
M -= qd;
x += M;
}
qd = x;
ql |= 1;
primary((Fold - 1)) += qd;
}
}
//! Renormalize manually specified binned double precision
//!
//! Renormalization keeps the primary vector within the necessary bins by
//! shifting over to the carry vector
_CCCL_DEVICE _CCCL_FORCEINLINE void binned_renorm()
{
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < Fold; i++)
{
auto tmp_renormd = primary(i);
auto& tmp_renorml = get_bit_representation(tmp_renormd);
carry(i) += static_cast<int>((tmp_renorml >> (mant_dig - 3)) & 3) - 2;
tmp_renorml &= ~(1ull << (mant_dig - 3));
tmp_renorml |= 1ull << (mant_dig - 2);
primary(i) = tmp_renormd;
}
}
//! Add scalar to manually specified binned fp (Y += X)
//!
//! Performs the operation Y += X on an binned type Y
_CCCL_DEVICE _CCCL_FORCEINLINE void binned_add(const ftype x)
{
binned_update(x);
binned_deposit(x);
binned_renorm();
}
//! Add two manually specified binned fp (Y += X)
//! Performs the operation Y += X
//!
//! @param x Another binned fp of the same type
_CCCL_DEVICE void binned_add(const ReproducibleFloatingAccumulator& x)
{
const auto X_index = x.binned_index();
const auto Y_index = this->binned_index();
const auto shift = Y_index - X_index;
if (shift > 0)
{
// shift Y upwards and add X to Y
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = Fold - 1; i >= 1; i--)
{
if (i < shift)
{
break;
}
primary(i) = x.primary(i) + (primary((i - shift)) - binned_bins(i - shift + Y_index));
carry(i) = x.carry(i) + carry((i - shift));
}
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < Fold; i++)
{
if (i == shift)
{
break;
}
primary(i) = x.primary(i);
carry(i) = x.carry(i);
}
}
else if (shift < 0)
{
// shift X upwards and add X to Y
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < Fold; i++)
{
if (i < -shift)
{
continue;
}
primary(i) += x.primary((i + shift)) - binned_bins(X_index + i + shift);
carry(i) += x.carry((i + shift));
}
}
else if (shift == 0)
{
// add X to Y
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < Fold; i++)
{
primary(i) += x.primary(i) - binned_bins(i + X_index);
carry(i) += x.carry(i);
}
}
binned_renorm();
}
[[nodiscard]] _CCCL_DEVICE double conv_binned_to_double() const
{
int i = 0;
double Y = 0.0;
const auto X_index = binned_index();
if (X_index <= (3 * mant_dig) / bin_width)
{
double scale_down = ::cuda::std::ldexpf(0.5f, 1 - (2 * mant_dig - bin_width));
double scale_up = ::cuda::std::ldexpf(0.5f, 1 - (2 * mant_dig - bin_width));
int scaled = ::cuda::std::max(::cuda::std::min(Fold, (3 * mant_dig) / bin_width - X_index), 0);
if (X_index == 0)
{
Y += carry(0) * ((binned_bins(0 + X_index) / 6.0) * scale_down * expansion);
Y += carry(1) * ((binned_bins(1 + X_index) / 6.0) * scale_down);
Y += (primary(0) - binned_bins(0 + X_index)) * scale_down * expansion;
i = 2;
}
else
{
Y += carry(0) * ((binned_bins(0 + X_index) / 6.0) * scale_down);
i = 1;
}
for (; i < scaled; i++)
{
Y += carry(i) * ((binned_bins(i + X_index) / 6.0) * scale_down);
Y += (primary((i - 1)) - binned_bins(i - 1 + X_index)) * scale_down;
}
if (i == Fold)
{
Y += (primary((Fold - 1)) - binned_bins(Fold - 1 + X_index)) * scale_down;
return Y * scale_up;
}
if (::cuda::std::isinf(Y * scale_up))
{
return Y * scale_up;
}
Y *= scale_up;
for (; i < Fold; i++)
{
Y += carry(i) * (binned_bins(i + X_index) / 6.0);
Y += primary((i - 1)) - binned_bins(i - 1 + X_index);
}
Y += primary((Fold - 1)) - binned_bins(Fold - 1 + X_index);
}
else
{
Y += carry(0) * (binned_bins(0 + X_index) / 6.0);
for (i = 1; i < Fold; i++)
{
Y += carry(i) * (binned_bins(i + X_index) / 6.0);
Y += (primary((i - 1)) - binned_bins(i - 1 + X_index));
}
Y += (primary((Fold - 1)) - binned_bins(Fold - 1 + X_index));
}
return Y;
}
[[nodiscard]] _CCCL_DEVICE float conv_binned_to_float() const
{
int i = 0;
double Y = 0.0;
// Note that the following order of summation is in order of decreasing
// exponent. The following code is specific to SBWIDTH=13, FLT_MANT_DIG=24, and
// the number of carries equal to 1.
const auto X_index = binned_index();
if (X_index == 0)
{
Y += static_cast<double>(carry(0)) * static_cast<double>(binned_bins(0 + X_index) / 6.0)
* static_cast<double>(expansion);
Y += static_cast<double>(carry(1)) * static_cast<double>(binned_bins(1 + X_index) / 6.0);
Y += static_cast<double>(primary(0) - binned_bins(0 + X_index)) * static_cast<double>(expansion);
i = 2;
}
else
{
Y += static_cast<double>(carry(0)) * static_cast<double>((binned_bins(0 + X_index) / 6.0));
i = 1;
}
for (; i < Fold; i++)
{
Y += static_cast<double>(carry(i)) * static_cast<double>(binned_bins(i + X_index) / 6.0);
Y += static_cast<double>(primary(i - 1) - binned_bins(i - 1 + X_index));
}
Y += static_cast<double>(primary(Fold - 1) - binned_bins(Fold - 1 + X_index));
return static_cast<float>(Y);
}
public:
ReproducibleFloatingAccumulator() = default;
/// Set the binned fp to zero
_CCCL_DEVICE void zero() noexcept
{
data = {};
}
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE constexpr int endurance() const noexcept
{
return 1 << (mant_dig - bin_width - 2);
}
//! Accumulate an arithmetic @p x into the binned fp.
//! NOTE: Casts @p x to the type of the binned fp
_CCCL_TEMPLATE(typename U)
_CCCL_REQUIRES(::cuda::std::is_arithmetic_v<U>)
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator+=(const U x)
{
binned_add(static_cast<ftype>(x));
return *this;
}
//! Accumulate-subtract an arithmetic @p x into the binned fp.
//! NOTE: Casts @p x to the type of the binned fp
_CCCL_TEMPLATE(typename U)
_CCCL_REQUIRES(::cuda::std::is_arithmetic_v<U>)
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator-=(const U x)
{
binned_add(-static_cast<ftype>(x));
return *this;
}
/// Accumulate a binned fp @p x into the binned fp.
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator+=(const ReproducibleFloatingAccumulator& other)
{
binned_add(other);
return *this;
}
//! Accumulate-subtract a binned fp @p other into the binned fp.
//! NOTE: Makes a copy and performs arithmetic; slow.
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator-=(const ReproducibleFloatingAccumulator& other)
{
const auto temp = -other;
binned_add(temp);
}
_CCCL_DEVICE friend bool operator==(const ReproducibleFloatingAccumulator& a, const ReproducibleFloatingAccumulator& b)
{
return a.data == b.data;
}
_CCCL_DEVICE friend bool operator!=(const ReproducibleFloatingAccumulator& a, const ReproducibleFloatingAccumulator& b)
{
return !(a == b);
}
//! Sets this binned fp equal to the arithmetic value @p x
//! NOTE: Casts @p x to the type of the binned fp
_CCCL_TEMPLATE(typename U)
_CCCL_REQUIRES(::cuda::std::is_arithmetic_v<U>)
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator=(const U x)
{
zero();
binned_add(static_cast<ftype>(x));
return *this;
}
//! Returns the negative of this binned fp
//! NOTE: Makes a copy and performs arithmetic; slow.
[[nodiscard]] _CCCL_DEVICE ReproducibleFloatingAccumulator operator-() const
{
ReproducibleFloatingAccumulator temp = *this;
if (primary(0) != 0.0)
{
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < Fold; i++)
{
temp.primary(i) = binned_bins(i + binned_index()) - (primary(i) - binned_bins(i + binned_index()));
temp.carry(i) = -carry(i);
}
}
return temp;
}
/// Convert this binned fp into its native floating-point representation
[[nodiscard]] _CCCL_DEVICE ftype conv_to_fp() const
{
if (::cuda::std::is_same_v<ftype, float>)
{
return conv_binned_to_float();
}
else
{
return conv_binned_to_double();
}
}
/// Add @p x to the binned fp
_CCCL_DEVICE void add(const ftype x)
{
binned_add(x);
}
//////////////////////////////////////
// MANUAL OPERATIONS; USE WISELY
//////////////////////////////////////
//! Rebins for repeated accumulation of scalars with magnitude <= @p mav
//!
//! Once rebinned, `endurance` values <= @p mav can be added to the accumulator
//! with `unsafe_add` after which `renorm()` must be called. See the source of
//!`add()` for an example
_CCCL_DEVICE void set_max_val(const ftype mav)
{
binned_update(mav);
}
//! Add @p x to the binned fp
//!
//! This is intended to be used after a call to `set_max_abs_val()`
_CCCL_DEVICE void unsafe_add(const ftype x)
{
binned_deposit(x);
}
//! Renormalizes the binned fp
//!
//! This is intended to be used after a call to `set_max_abs_val()` and one or
//! more calls to `unsafe_add()`
_CCCL_DEVICE void renorm()
{
binned_renorm();
}
};
} // namespace detail::rfa
CUB_NAMESPACE_END

145
qwen3_6_scripts/cccl_preload/include/cub/detail/segmented_params.cuh Normal file
View File

@@ -0,0 +1,145 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/argument>
#include <cuda/std/__type_traits/integral_constant.h>
#include <cuda/std/__type_traits/remove_cvref.h>
#include <cuda/std/__utility/forward.h>
#include <cuda/std/cstddef>
CUB_NAMESPACE_BEGIN
namespace detail::params
{
// =====================================================================
// get_param — unified segment parameter access
// =====================================================================
//! @brief Returns the value of an argument for a given segment index.
//!
//! @param[in] __arg Argument or argument wrapper to read.
//! @param[in] __index Segment index to read for sequence arguments.
//! @return The single argument value, or the sequence element at the given index.
_CCCL_TEMPLATE(class _Tp, class _SegmentIndexT)
_CCCL_REQUIRES((!::cuda::args::__is_wrapper_v<::cuda::std::remove_cvref_t<_Tp>>) )
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_param(_Tp&& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
{
if constexpr (::cuda::args::__traits<::cuda::std::remove_cvref_t<_Tp>>::is_single_value)
{
return __arg;
}
else
{
return __arg[__index];
}
}
template <auto _Value, class _Tp, class _SegmentIndexT>
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
get_param(const ::cuda::args::constant<_Value, _Tp>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
{
return ::cuda::args::__unwrap(__arg);
}
template <class _Arg, class _StaticBounds, class _SegmentIndexT>
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
get_param(const ::cuda::args::immediate<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
{
return ::cuda::args::__unwrap(__arg);
}
template <class _Arg, class _StaticBounds, class _SegmentIndexT>
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
get_param(const ::cuda::args::deferred<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
{
return ::cuda::args::__unwrap(__arg);
}
template <class _Arg, class _StaticBounds, class _SegmentIndexT>
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
get_param(const ::cuda::args::deferred_sequence<_Arg, _StaticBounds>& __arg, _SegmentIndexT __index) noexcept
{
return ::cuda::args::__unwrap(__arg)[__index];
}
// =====================================================================
// Discrete parameter support
// =====================================================================
//! @brief Specifies a list of supported options for a parameter.
template <typename T, T... Options>
struct supported_options
{
static constexpr ::cuda::std::size_t count = sizeof...(Options);
};
//! @brief Static discrete parameter — a single compile-time value that is also its only supported option.
//!
//! Holds no runtime value, so it cannot be put into a state that disagrees with its supported option, and
//! @c dispatch_impl therefore always matches it. This is the safe representation for a compile-time-fixed discrete
//! parameter (e.g. a statically known top-k selection direction): modeling such a parameter with a runtime value
//! instead would risk that value silently disagreeing with the supported option (a no-op dispatch unless
//! @c CCCL_ENABLE_ASSERTIONS is set).
template <typename T, T Value>
struct static_discrete_param
{
using value_type = T;
using supported_options_t = supported_options<T, Value>;
template <typename SegmentIndexT>
[[nodiscard]] _CCCL_HOST_DEVICE constexpr T get_param(SegmentIndexT) const noexcept
{
return Value;
}
};
// =====================================================================
// Discrete dispatch
// =====================================================================
//! @brief Translates a runtime parameter value into a compile-time constant by matching
//! against a list of supported options.
//!
//! @param[in] val Runtime value to match.
//! @param[in] __supported_options Supported values for the parameter.
//! @param[in] f Functor invoked with the matched compile-time constant.
//! @return `true` if the value matches one of the supported options.
template <typename T, T... Opts, typename Functor>
[[nodiscard]] _CCCL_HOST_DEVICE bool
dispatch_impl(T val, [[maybe_unused]] supported_options<T, Opts...> __supported_options, Functor&& f)
{
const bool match_found = ((val == Opts ? (f(::cuda::std::integral_constant<T, Opts>{}), true) : false) || ...);
_CCCL_ASSERT(match_found, "The given runtime parameter value is not in the supported list");
return match_found;
}
//! @brief Dispatcher that resolves a discrete parameter to a compile-time constant
//! and invokes a functor with the matched option.
//!
//! @param[in] param Discrete parameter to resolve.
//! @param[in] segment_id Segment index to read from `param`.
//! @param[in] f Functor invoked with the matched compile-time constant.
//! @return `true` if the parameter value matches one of its supported options.
template <typename ParamT, typename SegmentIndexT, typename Functor>
[[nodiscard]] _CCCL_HOST_DEVICE bool dispatch_discrete(ParamT param, SegmentIndexT segment_id, Functor&& f)
{
using supported_list = typename ParamT::supported_options_t;
auto param_value = param.get_param(segment_id);
return CUB_NS_QUALIFIER::detail::params::dispatch_impl(
param_value, supported_list{}, ::cuda::std::forward<Functor>(f));
}
} // namespace detail::params
CUB_NAMESPACE_END

167
qwen3_6_scripts/cccl_preload/include/cub/detail/strong_load.cuh Normal file
View File

@@ -0,0 +1,167 @@
// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
/**
* @file Utilities for strong memory operations.
*/
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
CUB_NAMESPACE_BEGIN
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
namespace detail
{
static _CCCL_DEVICE _CCCL_FORCEINLINE uint4 load_relaxed(uint4 const* ptr)
{
uint4 retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(retval.x),
"=r"(retval.y),
"=r"(retval.z),
"=r"(retval.w) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(retval.x),
"=r"(retval.y),
"=r"(retval.z),
"=r"(retval.w) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE ulonglong2 load_relaxed(ulonglong2 const* ptr)
{
ulonglong2 retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE ushort4 load_relaxed(ushort4 const* ptr)
{
ushort4 retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.v4.u16 {%0, %1, %2, %3}, [%4];" : "=h"(retval.x),
"=h"(retval.y),
"=h"(retval.z),
"=h"(retval.w) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.v4.u16 {%0, %1, %2, %3}, [%4];" : "=h"(retval.x),
"=h"(retval.y),
"=h"(retval.z),
"=h"(retval.w) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE uint2 load_relaxed(uint2 const* ptr)
{
uint2 retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned long long load_relaxed(unsigned long long const* ptr)
{
unsigned long long retval;
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.u64 %0, [%1];" : "=l"(retval) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.u64 %0, [%1];" : "=l"(retval) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int load_relaxed(unsigned int const* ptr)
{
unsigned int retval;
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned short load_relaxed(unsigned short const* ptr)
{
unsigned short retval;
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("ld.relaxed.gpu.u16 %0, [%1];" : "=h"(retval) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.u16 %0, [%1];" : "=h"(retval) : "l"(ptr) : "memory");));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned char load_relaxed(unsigned char const* ptr)
{
unsigned short retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("{"
" .reg .u8 datum;"
" ld.relaxed.gpu.u8 datum, [%1];"
" cvt.u16.u8 %0, datum;"
"}" : "=h"(retval) : "l"(ptr) : "memory");),
(asm volatile("{"
" .reg .u8 datum;"
" ld.cg.u8 datum, [%1];"
" cvt.u16.u8 %0, datum;"
"}" : "=h"(retval) : "l"(ptr) : "memory");));
return (unsigned char) retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE ulonglong2 load_acquire(ulonglong2 const* ptr)
{
ulonglong2 retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("ld.acquire.gpu.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");),
({
asm volatile("ld.cg.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");
__threadfence();
}));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE uint2 load_acquire(uint2 const* ptr)
{
uint2 retval;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("ld.acquire.gpu.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");),
({
asm volatile("ld.cg.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");
__threadfence();
}));
return retval;
}
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int load_acquire(unsigned int const* ptr)
{
unsigned int retval;
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("ld.acquire.gpu.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");),
(asm volatile("ld.cg.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory"); __threadfence();));
return retval;
}
} // namespace detail
#endif // _CCCL_DOXYGEN_INVOKED
CUB_NAMESPACE_END

223
qwen3_6_scripts/cccl_preload/include/cub/detail/strong_store.cuh Normal file
View File

@@ -0,0 +1,223 @@
// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
/**
* @file Utilities for strong memory operations.
*/
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
CUB_NAMESPACE_BEGIN
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
namespace detail
{
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(uint4* ptr, uint4 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
"r"(val.x),
"r"(val.y),
"r"(val.z),
"r"(val.w) : "memory");),
(asm volatile(
"st.cg.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(ulonglong2* ptr, ulonglong2 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");),
(asm volatile("st.cg.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(ushort4* ptr, ushort4 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
"h"(val.x),
"h"(val.y),
"h"(val.z),
"h"(val.w) : "memory");),
(asm volatile(
"st.cg.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr), "h"(val.x), "h"(val.y), "h"(val.z), "h"(val.w) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(uint2* ptr, uint2 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");),
(asm volatile("st.cg.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned long long* ptr, unsigned long long val)
{
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");),
(asm volatile("st.cg.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned int* ptr, unsigned int val)
{
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");),
(asm volatile("st.cg.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned short* ptr, unsigned short val)
{
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
(asm volatile("st.relaxed.gpu.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");),
(asm volatile("st.cg.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");));
}
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned char* ptr, unsigned char val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("{"
" .reg .u8 datum;"
" cvt.u8.u16 datum, %1;"
" st.relaxed.gpu.u8 [%0], datum;"
"}" : : "l"(ptr),
"h"((unsigned short) val) : "memory");),
(asm volatile("{"
" .reg .u8 datum;"
" cvt.u8.u16 datum, %1;"
" st.cg.u8 [%0], datum;"
"}" : : "l"(ptr),
"h"((unsigned short) val) : "memory");));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(uint4* ptr, uint4 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.release.gpu.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
"r"(val.x),
"r"(val.y),
"r"(val.z),
"r"(val.w) : "memory");),
({
__threadfence();
asm volatile("st.cg.v4.u32 [%0], {%1, %2, %3, %4};"
:
: "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w)
: "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(ulonglong2* ptr, ulonglong2 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.release.gpu.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");),
({
__threadfence();
asm volatile("st.cg.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(ushort4* ptr, ushort4 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.release.gpu.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
"h"(val.x),
"h"(val.y),
"h"(val.z),
"h"(val.w) : "memory");),
({
__threadfence();
asm volatile("st.cg.v4.u16 [%0], {%1, %2, %3, %4};"
:
: "l"(ptr), "h"(val.x), "h"(val.y), "h"(val.z), "h"(val.w)
: "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(uint2* ptr, uint2 val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("st.release.gpu.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");),
({
__threadfence();
asm volatile("st.cg.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned long long* ptr, unsigned long long val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");), ({
__threadfence();
asm volatile("st.cg.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned int* ptr, unsigned int val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");), ({
__threadfence();
asm volatile("st.cg.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned short* ptr, unsigned short val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");), ({
__threadfence();
asm volatile("st.cg.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");
}));
}
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned char* ptr, unsigned char val)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_70,
(asm volatile("{"
" .reg .u8 datum;"
" cvt.u8.u16 datum, %1;"
" st.release.gpu.u8 [%0], datum;"
"}" : : "l"(ptr),
"h"((unsigned short) val) : "memory");),
({
__threadfence();
asm volatile(
"{"
" .reg .u8 datum;"
" cvt.u8.u16 datum, %1;"
" st.cg.u8 [%0], datum;"
"}"
:
: "l"(ptr), "h"((unsigned short) val)
: "memory");
}));
}
} // namespace detail
#endif // _CCCL_DOXYGEN_INVOKED
CUB_NAMESPACE_END

380
qwen3_6_scripts/cccl_preload/include/cub/detail/temporary_storage.cuh Normal file
View File

@@ -0,0 +1,380 @@
/*
* Copyright 2021 NVIDIA Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_namespace.cuh>
#include <cub/util_temporary_storage.cuh>
#include <cuda/__stream/stream_ref.h>
#include <cuda/std/__algorithm/max.h>
#include <cuda/std/__exception/exception_macros.h>
#include <cuda/std/cstddef>
CUB_NAMESPACE_BEGIN
namespace detail::temporary_storage
{
class slot;
template <typename T>
class alias;
template <int SlotsCount>
class layout;
/**
* @brief Temporary storage slot that can be considered a C++ union with an
* arbitrary fields count.
*
* @warning slot lifetime is defined by the lifetime of the associated layout.
* It's impossible to request new array if layout is already mapped.
*
* @par A Simple Example
* @code
* auto slot = temporary_storage.get_slot(0);
*
* // Add fields into the slot
* // Create an int alias with 0 elements:
* auto int_array = slot->create_alias<int>();
* // Create a double alias with 2 elements:
* auto double_array = slot->create_alias<double>(2);
* // Create a char alias with 0 elements:
* auto empty_array = slot->create_alias<char>();
* // Slot size is defined by double_array size (2 * sizeof(double))
*
* if (condition)
* {
* int_array.grow(42);
* // Now slot size is defined by int_array size (42 * sizeof(int))
* }
*
* // Temporary storage mapping
* // ...
* int *d_int_array = int_array.get();
* double *d_double_array = double_array.get();
* char *d_empty_array = empty_array.get(); // Guaranteed to return nullptr
* @endcode
*/
class slot
{
size_t m_size{};
void* m_pointer{};
public:
slot() = default;
/**
* @brief Returns an array of type @p T and length @p elements
*/
template <typename T>
_CCCL_HOST_DEVICE alias<T> create_alias(size_t elements = 0);
private:
_CCCL_HOST_DEVICE void set_bytes_required(size_t new_size)
{
m_size = (::cuda::std::max) (m_size, new_size);
}
_CCCL_HOST_DEVICE size_t get_bytes_required() const
{
return m_size;
}
_CCCL_HOST_DEVICE void set_storage(void* ptr)
{
m_pointer = ptr;
}
_CCCL_HOST_DEVICE void* get_storage() const
{
return m_pointer;
}
template <typename T>
friend class alias;
template <int>
friend class layout;
};
/**
* @brief Named memory region of a temporary storage slot
*
* @par Overview
* This class provides a typed wrapper of a temporary slot memory region.
* It can be considered as a field in the C++ union. It's only possible to
* increase the array size.
*
* @warning alias lifetime is defined by the lifetime of the associated slot
* It's impossible to grow the array if the layout is already mapped.
*/
template <typename T>
class alias
{
slot& m_slot;
size_t m_elements{};
_CCCL_HOST_DEVICE explicit alias(slot& slot, size_t elements = 0)
: m_slot(slot)
, m_elements(elements)
{
this->update_slot();
}
_CCCL_HOST_DEVICE void update_slot()
{
m_slot.set_bytes_required(m_elements * sizeof(T));
}
public:
alias() = delete;
/**
* @brief Increases the number of elements
*
* @warning
* This method should be called before temporary storage mapping stage.
*
* @param[in] new_elements Increases the memory region occupied in the
* temporary slot to fit up to @p new_elements items
* of type @p T.
*/
_CCCL_HOST_DEVICE void grow(size_t new_elements)
{
m_elements = new_elements;
this->update_slot();
}
/**
* @brief Returns pointer to array
*
* If the @p elements number is equal to zero, or storage layout isn't mapped,
* @p nullptr is returned.
*/
_CCCL_HOST_DEVICE T* get() const
{
if (m_elements == 0)
{
return nullptr;
}
return reinterpret_cast<T*>(m_slot.get_storage());
}
friend class slot;
};
template <typename T>
_CCCL_HOST_DEVICE alias<T> slot::create_alias(size_t elements)
{
return alias<T>(*this, elements);
}
/**
* @brief Temporary storage layout represents a structure with
* @p SlotsCount union-like fields
*
* The layout can be mapped to a temporary buffer only once.
*
* @par A Simple Example
* @code
* cub::detail::temporary_storage::layout<3> temporary_storage;
*
* auto slot_1 = temporary_storage.get_slot(0);
* auto slot_2 = temporary_storage.get_slot(1);
*
* // Add fields into the first slot
* auto int_array = slot_1->create_alias<int>(1);
* auto double_array = slot_1->create_alias<double>(2);
*
* // Add fields into the second slot
* auto char_array = slot_2->create_alias<char>();
*
* // The equivalent C++ structure could look like
* // struct StorageLayout
* // {
* // union {
* // } slot_0;
* // std::byte padding_0[256 - sizeof (slot_0)];
* //
* // union {
* // int alias_0[1];
* // double alias_1[2];
* // } slot_1;
* // std::byte padding_1[256 - sizeof (slot_1)];
* //
* // union {
* // char alias_0[0];
* // } slot_2;
* // std::byte padding_2[256 - sizeof (slot_2)];
* // };
*
* // The third slot is empty
*
* // Temporary storage mapping
* if (d_temp_storage == nullptr)
* {
* temp_storage_bytes = temporary_storage.get_size();
* return;
* }
* else
* {
* temporary_storage.map_to_buffer(d_temp_storage, temp_storage_bytes);
* }
*
* // Use pointers
* int *d_int_array = int_array.get();
* double *d_double_array = double_array.get();
* char *d_char_array = char_array.get();
* @endcode
*/
template <int SlotsCount>
class layout
{
slot m_slots[SlotsCount];
size_t m_sizes[SlotsCount];
void* m_pointers[SlotsCount];
bool m_layout_was_mapped{};
public:
layout() = default;
_CCCL_HOST_DEVICE slot* get_slot(int slot_id)
{
if (slot_id < SlotsCount)
{
return &m_slots[slot_id];
}
return nullptr;
}
/**
* @brief Returns required temporary storage size in bytes
*/
_CCCL_HOST_DEVICE size_t get_size()
{
this->prepare_interface();
// alias_temporaries can return error only in mapping stage, so it's safe to ignore it here.
size_t temp_storage_bytes{};
[[maybe_unused]] const auto error = detail::alias_temporaries(nullptr, temp_storage_bytes, m_pointers, m_sizes);
_CCCL_ASSERT(error == cudaSuccess, "");
if (temp_storage_bytes == 0)
{
// The current CUB convention implies that there are two stages for each
// device-scope function call. The first one returns the required storage
// size. The second stage consumes temporary storage to perform some work.
// The only way to distinguish between the two stages is by checking the
// value of the temporary storage pointer. If zero bytes are requested,
// `cudaMalloc` will return `nullptr`. This fact makes it impossible to
// distinguish between the two stages, so we request some fixed amount of
// bytes (even if we don't need it) to have a non-null temporary storage
// pointer.
return 1;
}
return temp_storage_bytes;
}
/**
* @brief Maps the layout to the temporary storage buffer.
*/
_CCCL_HOST_DEVICE cudaError_t map_to_buffer(void* d_temp_storage, size_t temp_storage_bytes)
{
if (m_layout_was_mapped)
{
return cudaErrorAlreadyMapped;
}
this->prepare_interface();
if (cudaError_t error = detail::alias_temporaries(d_temp_storage, temp_storage_bytes, m_pointers, m_sizes))
{
return error;
}
for (size_t slot_id = 0; slot_id < SlotsCount; slot_id++)
{
m_slots[slot_id].set_storage(m_pointers[slot_id]);
}
m_layout_was_mapped = true;
return cudaSuccess;
}
private:
_CCCL_HOST_DEVICE void prepare_interface()
{
if (m_layout_was_mapped)
{
return;
}
for (size_t slot_id = 0; slot_id < SlotsCount; slot_id++)
{
const size_t slot_size = m_slots[slot_id].get_bytes_required();
m_sizes[slot_id] = slot_size;
m_pointers[slot_id] = nullptr;
}
}
};
template <typename MRT>
CUB_RUNTIME_FUNCTION cudaError_t
allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_bytes, MRT& mr)
{
_CCCL_TRY
{
d_temp_storage = mr.allocate(stream, temp_storage_bytes, alignof(::cuda::std::max_align_t));
return cudaSuccess;
}
_CCCL_CATCH_ALL
{
return cudaErrorMemoryAllocation;
}
}
template <typename MRT>
CUB_RUNTIME_FUNCTION cudaError_t
deallocate(::cuda::stream_ref stream, void* d_temp_storage, size_t temp_storage_bytes, MRT& mr)
{
_CCCL_TRY
{
mr.deallocate(stream, d_temp_storage, temp_storage_bytes, alignof(::cuda::std::max_align_t));
return cudaSuccess;
}
_CCCL_CATCH_ALL
{
return cudaErrorMemoryAllocation;
}
}
} // namespace detail::temporary_storage
CUB_NAMESPACE_END

185
qwen3_6_scripts/cccl_preload/include/cub/detail/type_traits.cuh Normal file
View File

@@ -0,0 +1,185 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
/**
* \file
* Wrappers and extensions around <type_traits> utilities.
*/
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_cpp_dialect.cuh>
#include <cub/util_namespace.cuh>
#include <cuda/std/__concepts/concept_macros.h> // IWYU pragma: keep
#include <cuda/std/__fwd/array.h>
#include <cuda/std/__fwd/mdspan.h>
#include <cuda/std/__fwd/span.h>
#include <cuda/std/__type_traits/always_false.h>
#include <cuda/std/__type_traits/conditional.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__type_traits/integral_constant.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__type_traits/is_signed_integer.h>
#include <cuda/std/__type_traits/is_unsigned_integer.h>
#include <cuda/std/__type_traits/remove_cv.h>
#include <cuda/std/__type_traits/void_t.h>
#include <cuda/std/__utility/declval.h>
#include <cuda/std/cstddef>
CUB_NAMESPACE_BEGIN
namespace detail
{
template <typename T, typename... TArgs>
inline constexpr bool is_one_of_v = (::cuda::std::is_same_v<T, TArgs> || ...);
template <typename T, typename V, typename = void>
struct has_binary_call_operator : ::cuda::std::false_type
{};
template <typename T, typename V>
struct has_binary_call_operator<
T,
V,
::cuda::std::void_t<decltype(::cuda::std::declval<T>()(::cuda::std::declval<V>(), ::cuda::std::declval<V>()))>>
: ::cuda::std::true_type
{};
/***********************************************************************************************************************
* Array-like type traits
**********************************************************************************************************************/
template <typename T>
inline constexpr bool is_fixed_size_random_access_range_v = false;
template <typename T, size_t N>
inline constexpr bool is_fixed_size_random_access_range_v<T[N]> = true;
template <typename T, size_t N>
inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::array<T, N>> = true;
template <typename T, size_t N>
inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::span<T, N>> = N != ::cuda::std::dynamic_extent;
template <typename T, typename E, typename L, typename A>
inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::mdspan<T, E, L, A>> =
E::rank() == 1 && E::rank_dynamic() == 0;
/***********************************************************************************************************************
* static_size: a type trait that returns the number of elements in an Array-like type
**********************************************************************************************************************/
template <typename T>
inline constexpr int static_size_v = ::cuda::std::enable_if_t<::cuda::std::__always_false_v<T>>{};
template <typename T, size_t N>
inline constexpr int static_size_v<T[N]> = N;
template <typename T, size_t N>
inline constexpr int static_size_v<::cuda::std::array<T, N>> = N;
template <typename T, size_t N>
inline constexpr int static_size_v<::cuda::std::span<T, N>> =
::cuda::std::enable_if_t<N != ::cuda::std::dynamic_extent, int>{N};
template <typename T, typename E, typename L, typename A>
inline constexpr int static_size_v<::cuda::std::mdspan<T, E, L, A>> =
::cuda::std::enable_if_t<E::rank() == 1 && E::rank_dynamic() == 0, int>{E::static_extent(0)};
template <typename T>
using implicit_prom_t = decltype(+T{});
/***********************************************************************************************************************
* Extended floating point traits
**********************************************************************************************************************/
// half
template <typename>
inline constexpr bool is_half_impl_v = false;
template <typename>
inline constexpr bool is_half2_impl_v = false;
#if _CCCL_HAS_NVFP16()
template <>
inline constexpr bool is_half_impl_v<__half> = true;
template <>
inline constexpr bool is_half2_impl_v<__half2> = true;
#endif // _CCCL_HAS_NVFP16
template <typename T>
inline constexpr bool is_half_v = is_half_impl_v<::cuda::std::remove_cv_t<T>>;
template <typename T>
inline constexpr bool is_half2_v = is_half2_impl_v<::cuda::std::remove_cv_t<T>>;
template <typename T>
inline constexpr bool is_any_half_v = is_half_impl_v<T> || is_half2_impl_v<T>;
//----------------------------------------------------------------------------------------------------------------------
// bfloat16
template <typename>
inline constexpr bool is_bfloat16_impl_v = false;
template <typename>
inline constexpr bool is_bfloat162_impl_v = false;
#if _CCCL_HAS_NVBF16()
template <>
inline constexpr bool is_bfloat16_impl_v<__nv_bfloat16> = true;
template <>
inline constexpr bool is_bfloat162_impl_v<__nv_bfloat162> = true;
#endif // _CCCL_HAS_NVBF16
template <typename T>
inline constexpr bool is_bfloat16_v = is_bfloat16_impl_v<::cuda::std::remove_cv_t<T>>;
template <typename T>
inline constexpr bool is_bfloat162_v = is_bfloat162_impl_v<::cuda::std::remove_cv_t<T>>;
template <typename T>
inline constexpr bool is_any_bfloat16_v = is_bfloat16_v<T> || is_bfloat162_v<T>;
//----------------------------------------------------------------------------------------------------------------------
// short2/ushort2
template <typename T>
inline constexpr bool is_any_short2_impl_v = false;
template <>
inline constexpr bool is_any_short2_impl_v<short2> = true;
template <>
inline constexpr bool is_any_short2_impl_v<ushort2> = true;
template <typename T>
inline constexpr bool is_any_short2_v = is_any_short2_impl_v<::cuda::std::remove_cv_t<T>>;
//----------------------------------------------------------------------------------------------------------------------
// - promote small integer types to their corresponding 32-bit promotion type
// - address the incompatibility between linux/windows for int/long
template <typename T>
using signed_promotion_t = ::cuda::std::conditional_t<
::cuda::std::__cccl_is_signed_integer_v<T> && sizeof(T) <= sizeof(int),
int,
::cuda::std::conditional_t<::cuda::std::__cccl_is_unsigned_integer_v<T> && sizeof(T) <= sizeof(uint32_t), uint32_t, T>>;
} // namespace detail
CUB_NAMESPACE_END

48
qwen3_6_scripts/cccl_preload/include/cub/detail/uninitialized_copy.cuh Normal file
View File

@@ -0,0 +1,48 @@
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/__type_traits/is_trivially_copyable.h>
#include <cuda/std/__new/device_new.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__utility/forward.h>
CUB_NAMESPACE_BEGIN
namespace detail
{
#if _CCCL_CUDA_COMPILER(NVHPC)
template <typename T, typename U>
_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val)
{
// NVBug 3384810
new (ptr) T(::cuda::std::forward<U>(val));
}
#else // ^^^ _CCCL_CUDA_COMPILER(NVHPC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVHPC) vvv
template <typename T, typename U, ::cuda::std::enable_if_t<::cuda::is_trivially_copyable_v<T>, int> = 0>
_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val)
{
// gevtushenko: placement new should work here as well, but the code generated for copy assignment is sometimes better
*ptr = ::cuda::std::forward<U>(val);
}
template <typename T, typename U, ::cuda::std::enable_if_t<!::cuda::is_trivially_copyable_v<T>, int> = 0>
_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val)
{
new (ptr) T(::cuda::std::forward<U>(val));
}
#endif // !_CCCL_CUDA_COMPILER(NVHPC)
} // namespace detail
CUB_NAMESPACE_END

34
qwen3_6_scripts/cccl_preload/include/cub/detail/unsafe_bitcast.cuh Normal file
View File

@@ -0,0 +1,34 @@
// SPDX-FileCopyrightText: Copyright (c) 2011-2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
CUB_NAMESPACE_BEGIN
namespace detail
{
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
// NOTE: bit_cast cannot be always used because __half, __nv_bfloat16, etc. are not trivially copyable
template <typename Output, typename Input>
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE Output unsafe_bitcast(const Input& input)
{
Output output;
static_assert(sizeof(input) == sizeof(output), "wrong size");
// NOLINTNEXTLINE(bugprone-undefined-memory-manipulation)
::memcpy(&output, &input, sizeof(input));
return output;
}
#endif // !_CCCL_DOXYGEN_INVOKED
} // namespace detail
CUB_NAMESPACE_END

83
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/allocators/smem_allocator.cuh Normal file
View File

@@ -0,0 +1,83 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/optimize_smem_ptr.cuh>
#include <cuda/std/__type_traits/is_constant_evaluated.h>
#include <cuda/std/cstdint>
#include <nv/target>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
struct SmemAllocator
{
::cuda::std::uint32_t mPtrSmem32 = 0;
int mAllocatedSize = 0;
_CCCL_HOST_DEVICE_API constexpr SmemAllocator() noexcept
{
// we only need the real pointer at runtime in device code
if (!::cuda::std::is_constant_evaluated())
{
NV_IF_TARGET(NV_IS_DEVICE, mPtrSmem32 = dynamic_smem_base();)
}
}
[[nodiscard]] _CCCL_DEVICE_API static ::cuda::std::uint32_t dynamic_smem_base() noexcept
{
extern __shared__ char warpSpeedDynamicSmemBase[];
return __cvta_generic_to_shared(warpSpeedDynamicSmemBase);
}
// SmemAllocator is a non-copyable, non-movable type. It must be passed by
// (mutable) reference to be useful.
SmemAllocator(const SmemAllocator&) = delete; // Delete copy constructor
SmemAllocator(SmemAllocator&&) = delete; // Delete move constructor
SmemAllocator& operator=(const SmemAllocator&) = delete; // Delete copy assignment
SmemAllocator& operator=(SmemAllocator&&) = delete; // Delete move assignment
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr void* alloc(::cuda::std::uint32_t size, ::cuda::std::uint32_t align = 0)
{
// Align mPtrSmem32 to requested alignment (round-up)
::cuda::std::uint32_t ptrAllocation32 = (mPtrSmem32 + (align - 1)) & ~(align - 1);
// Move base pointer and update allocated size
mAllocatedSize += static_cast<int>(size + ptrAllocation32 - mPtrSmem32);
mPtrSmem32 = ptrAllocation32 + size;
// we only need the pointer at runtime in device code
if (!::cuda::std::is_constant_evaluated())
{
NV_IF_TARGET(
NV_IS_DEVICE,
(
// Convert allocated smem address to generic pointer
void* mPtrAllocation = __cvta_shared_to_generic(ptrAllocation32);
// Ensure alignment calculation does not move down into rest of kernel code.
return optimizeSmemPtr(mPtrAllocation);))
}
return nullptr;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::uint32_t sizeBytes() const
{
return mAllocatedSize;
}
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

52
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/constant_assert.cuh Normal file
View File

@@ -0,0 +1,52 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/cassert>
/*
* _WS_CONSTANT_ASSERT: an assertion that is intended to be verified at compile time.
*
* A _WS_CONSTANT_ASSERT asserts something that the compiler (optimizer) can verify
* at compile time. Therefore, it does not result in an actual call to assert in
* the compiled binary. This allows checking various properties that cannot be
* verified using static_assert.
*
* To ensure that all _WS_CONSTANT_ASSERTs are in fact eliminated, compile with
* -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME. With this macro defined, any
* _WS_CONSTANT_ASSERT failure will output illegal PTX containing the error message.
* As a result, compilation will fail.
*
* Compiling with -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME has the additional
* advantage that violating any of the assertions can be detected at compile
* time and before even running the code.
*
*/
#if defined(_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME) && defined(__CUDA_ARCH__)
// When _WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME is defined and compiling for device, output illegal PTX.
// This causes the compilation to fail.
# define _WS_CONSTANT_ASSERT(expr, msg) \
do \
{ \
if (!(expr)) \
{ \
asm volatile(".pragma \"\n" __FILE__ "(" _CCCL_TO_STRING( \
__LINE__) "): %0" \
": error: constant assertion failed with '" msg "'\n\";" ::"C"(__func__)); \
} \
} while (0)
#else
// Host or !_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME
# define _WS_CONSTANT_ASSERT(expr, msg) _CCCL_ASSERT((expr), msg)
#endif

338
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/look_ahead.cuh Normal file
View File

@@ -0,0 +1,338 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/strong_load.cuh>
#include <cub/detail/strong_store.cuh>
#include <cub/detail/warpspeed/special_registers.cuh>
#include <cub/thread/thread_store.cuh>
#include <cub/warp/specializations/warp_redux.cuh>
#include <cub/warp/warp_reduce.cuh>
#include <cuda/__cmath/pow2.h>
#include <cuda/__functional/operator_properties.h>
#include <cuda/__memory/is_aligned.h>
#include <cuda/__ptx/instructions/get_sreg.h>
#include <cuda/__type_traits/is_trivially_copyable.h>
#include <cuda/std/__algorithm/min.h>
#include <cuda/std/__bit/popcount.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__type_traits/underlying_type.h>
#if !_CCCL_HAS_NV_ATOMIC_BUILTINS()
# include <cuda/atomic>
#endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS()
#include <nv/target>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL ::cuda::std::size_t max_native_atomic_size() noexcept
{
#if _CCCL_CUDA_COMPILER(NVHPC)
return 8;
#else // ^^^ _CCCL_CUDA_COMPILER(NVHPC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVHPC) vvv
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return 16;), (return 8;))
#endif // !_CCCL_CUDA_COMPILER(NVHPC)
}
enum scan_state : ::cuda::std::uint32_t
{
empty = 0,
tile_aggregate = 1,
};
template <typename AccumT>
struct tile_state_unaligned_t
{
scan_state state;
AccumT value;
};
// some older nvcc versions do not evaluate next_power_of_two() at compile time when called inside an attribute, so we
// have to force constant evaluation by assigning the result to a template parameter
template <typename AccumT,
::cuda::std::size_t _Alignment = ::cuda::next_power_of_two(sizeof(tile_state_unaligned_t<AccumT>))>
struct alignas(_Alignment) tile_state_t : tile_state_unaligned_t<AccumT>
{};
#if __cccl_ptx_isa >= 860
template <typename AccumT>
_CCCL_DEVICE_API void
storeTileAggregate(tile_state_t<AccumT>* ptrTileStates, scan_state scanState, AccumT aggr, int index, int num_tiles)
{
_CCCL_ASSERT(::cuda::is_aligned(ptrTileStates, alignof(tile_state_t<AccumT>)), "");
_CCCL_ASSERT(index >= 0 && index < num_tiles, "Reading out of bounds tile state");
if constexpr (sizeof(tile_state_t<AccumT>) <= cub::detail::warpspeed::max_native_atomic_size()
&& ::cuda::is_trivially_copyable_v<tile_state_t<AccumT>>)
{
static_assert(::cuda::is_power_of_two(sizeof(tile_state_t<AccumT>)));
tile_state_t<AccumT> tmp{scanState, aggr};
# if _CCCL_HAS_NV_ATOMIC_BUILTINS()
__nv_atomic_store(ptrTileStates + index, &tmp, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE);
# else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv
::cuda::atomic_ref<tile_state_t<AccumT>, ::cuda::std::thread_scope_device>{ptrTileStates[index]}.store(
tmp, ::cuda::std::memory_order_relaxed);
# endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS()
}
else
{
ThreadStore<STORE_CG>(&ptrTileStates[index].value, aggr);
using state_int = ::cuda::std::underlying_type_t<scan_state>;
store_release(reinterpret_cast<state_int*>(&ptrTileStates[index].state), scanState);
}
}
template <typename AccumT>
_CCCL_DEVICE_API tile_state_t<AccumT> loadTileAggregate(tile_state_t<AccumT>* ptrTileStates, int index, int num_tiles)
{
_CCCL_ASSERT(::cuda::is_aligned(ptrTileStates, alignof(tile_state_t<AccumT>)), "");
_CCCL_ASSERT(index >= 0 && index < num_tiles, "Reading out of bounds tile state");
tile_state_t<AccumT> res;
if constexpr (sizeof(tile_state_t<AccumT>) <= cub::detail::warpspeed::max_native_atomic_size()
&& ::cuda::is_trivially_copyable_v<tile_state_t<AccumT>>)
{
static_assert(::cuda::is_power_of_two(sizeof(tile_state_t<AccumT>)));
# if _CCCL_HAS_NV_ATOMIC_BUILTINS()
__nv_atomic_load(ptrTileStates + index, &res, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE);
# else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv
res = ::cuda::atomic_ref<tile_state_t<AccumT>, ::cuda::std::thread_scope_device>{ptrTileStates[index]}.load(
::cuda::std::memory_order_relaxed);
# endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS()
}
else
{
using state_int = ::cuda::std::underlying_type_t<scan_state>;
res.state = static_cast<scan_state>(load_acquire(reinterpret_cast<const state_int*>(&ptrTileStates[index].state)));
res.value = ThreadLoad<LOAD_CG>(&ptrTileStates[index].value);
}
return res;
}
// warpLoadLookahead loads tmp states:
// idxTileCur + [0; 32 * numTileStatesPerThread[
//
// The states are loaded in laneId order and warp-strided:
//
// outTmpStates[0] contains:
// Lane 0: idxTileCur + 0
// Lane 1: idxTileCur + 1
// ...
// Lane 31: idxTileCur + 31
//
// outTmpStates[1] contains:
// Lane 0: idxTileCur + 32
// ...
// Lane 31 idxTileCur + 63
//
// If the index idxTileCur + ii of the loaded state is equal to or exceeds idxTileNext, i.e., idxTileCur + ii >=
// idxTileNext, then the state is not loaded from memory and set to empty.
template <int numTileStatesPerThread, typename AccumT>
_CCCL_DEVICE_API void warpLoadLookahead(
int laneIdx,
tile_state_t<AccumT> (&outTileStates)[numTileStatesPerThread],
tile_state_t<AccumT>* ptrTileStates,
int idxTileCur,
int idxTileNext,
int num_tiles)
{
for (int i = 0; i < numTileStatesPerThread; ++i)
{
const int idxTileLookahead = idxTileCur + 32 * i + laneIdx;
if (idxTileLookahead < idxTileNext)
{
outTileStates[i] = loadTileAggregate(ptrTileStates, idxTileLookahead, num_tiles);
}
else
{
// If we are looking ahead of idxTileNext, then set state to empty
outTileStates[i].state = scan_state::empty;
}
}
}
// warpIncrementalLookahead takes the latest known aggrExclusiveCtaPrev and its tile index, idxTilePrev (which's
// aggregate is NOT included in aggrExclusiveCtaPrev), and computes the aggrExclusiveCta for the next tile of interest,
// idxTileNext (where the returned value will NOT include the aggregate of idxTileNext).
//
// It does so by loading states in chunks of 32 * numTileStatesPerThread elements, starting from idxTilePrev + 1. From
// the chunk of states, it tries to advance its knowledge of aggrExclusiveCta as much as possible. It loops until it can
// calculate the value of aggrExclusiveCta from the preceding states.
//
// The function must be called from a single warp. All passed arguments must be warp-uniform.
template <int numTileStatesPerThread, typename AccumT, typename ScanOpT>
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE AccumT warpIncrementalLookahead(
SpecialRegisters specialRegisters,
tile_state_t<AccumT>* ptrTileStates,
const int idxTilePrev,
const AccumT aggrExclusiveCtaPrev,
const int idxTileNext,
ScanOpT& scan_op,
const int num_tiles)
{
const int laneIdx = static_cast<int>(specialRegisters.laneIdx);
[[maybe_unused]] const ::cuda::std::uint32_t lanemaskEq = ::cuda::ptx::get_sreg_lanemask_eq();
int idxTileCur = idxTilePrev;
AccumT aggrExclusiveCtaCur = aggrExclusiveCtaPrev;
using warp_reduce_t = WarpReduce<AccumT>;
static_assert(::cuda::std::is_same_v<typename warp_reduce_t::TempStorage, Uninitialized<NullType>>,
"WarpReduce for a full warp must not require temporary storage");
[[maybe_unused]] typename warp_reduce_t::TempStorage temp_storage;
while (idxTileCur < idxTileNext)
{
tile_state_t<AccumT> regTmpStates[numTileStatesPerThread];
warpLoadLookahead(laneIdx, regTmpStates, ptrTileStates, idxTileCur, idxTileNext, num_tiles);
for (int idx = 0; idx < numTileStatesPerThread; ++idx)
{
// Bitmask with 1 bits indicating which lane has a tile aggregate
const ::cuda::std::uint32_t warp_has_aggregate_mask =
__ballot_sync(0xffffffffu, regTmpStates[idx].state == scan_state::tile_aggregate);
// Bitmask with 1 bits for all rightmost lanes having a tile aggregate
const ::cuda::std::uint32_t warp_right_aggregates_mask = warp_has_aggregate_mask & (~warp_has_aggregate_mask - 1);
// Cannot reduce if no rightmost tile aggregates
if (warp_right_aggregates_mask == 0)
{
break;
}
const ::cuda::std::uint32_t warp_right_aggregates_count = ::cuda::std::popcount(warp_right_aggregates_mask);
// Accumulate the rightmost tile aggregates
AccumT local_aggr;
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_80,
({ // NOTE: Inlined from warp_reduce_shfl
if constexpr (is_warp_redux_op_supported_sm80<ScanOpT, AccumT>)
{
const bool use_value = lanemaskEq & warp_right_aggregates_mask;
const AccumT value = use_value ? regTmpStates[idx].value : cuda::identity_element<ScanOpT, AccumT>();
local_aggr = cub::detail::warp_redux_sm80(value, ~0, scan_op);
}
else
{
// TODO(bgruber): this generates a LOT of SASS. I think it can do better.
local_aggr =
warp_reduce_t{temp_storage}.Reduce(regTmpStates[idx].value, scan_op, warp_right_aggregates_count);
}
}),
(local_aggr =
warp_reduce_t{temp_storage}.Reduce(regTmpStates[idx].value, scan_op, warp_right_aggregates_count);))
// We never initialized aggrExclusiveCtaCur when starting look ahead at tile 0
aggrExclusiveCtaCur = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr);
idxTileCur += warp_right_aggregates_count;
// we can only continue on the next 32 tile states, if we consumed all 32 of this iteration
if (warp_right_aggregates_count < 32)
{
break;
}
}
}
return aggrExclusiveCtaCur; // must only be valid in lane_0
}
// Deterministic version of warpIncrementalLookahead that returns the same aggrExclusiveCta. The difference is that it
// always starts the lookahead from a tile index that is a multiple of 32. The left pointer (idxTilePrev) is itself
// always a multiple of 32, as it starts at 0 and is only ever advanced by whole batches of 32, so the lookahead resumes
// from there directly. Because every reduction begins at the same fixed tiles, no matter which tiles happened to finish
// first, the order in which values are summed is always the same and the result is identical on every run.
// idxTilePrev/aggrExclusiveCtaPrev are updated by reference to the last multiple of 32.
template <int numTileStatesPerThread, typename AccumT, typename ScanOpT>
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE AccumT warpIncrementalLookaheadStable(
SpecialRegisters specialRegisters,
tile_state_t<AccumT>* ptrTileStates,
int& idxTilePrev,
AccumT& aggrExclusiveCtaPrev,
const int idxTileNext,
ScanOpT& scan_op,
const int num_tiles)
{
const int laneIdx = static_cast<int>(specialRegisters.laneIdx);
const ::cuda::std::uint32_t lanemaskEq = ::cuda::ptx::get_sreg_lanemask_eq();
int idxTileCur = idxTilePrev;
AccumT aggrExclusiveCtaCur = aggrExclusiveCtaPrev;
using warp_reduce_t = WarpReduce<AccumT>;
static_assert(::cuda::std::is_same_v<typename warp_reduce_t::TempStorage, Uninitialized<NullType>>,
"WarpReduce for a full warp must not require temporary storage");
[[maybe_unused]] typename warp_reduce_t::TempStorage temp_storage;
while (idxTileCur < idxTileNext)
{
tile_state_t<AccumT> regTmpStates[numTileStatesPerThread];
warpLoadLookahead(laneIdx, regTmpStates, ptrTileStates, idxTileCur, idxTileNext, num_tiles);
for (int idx = 0; idx < numTileStatesPerThread; ++idx)
{
// Bitmask with 1 bits indicating which lane has a tile aggregate
const ::cuda::std::uint32_t warp_has_aggregate_mask =
__ballot_sync(0xffffffffu, regTmpStates[idx].state == scan_state::tile_aggregate);
// Bitmask with 1 bits for the contiguous run of lanes having a tile aggregate starting from LSB
const ::cuda::std::uint32_t warp_right_aggregates_mask = warp_has_aggregate_mask & (~warp_has_aggregate_mask - 1);
const ::cuda::std::uint32_t warp_right_aggregates_count = ::cuda::std::popcount(warp_right_aggregates_mask);
// Only reduce once 32 contiguous tile aggregates are available, so the reduction order is fixed.
const ::cuda::std::uint32_t expected_count =
static_cast<::cuda::std::uint32_t>(::cuda::std::min(32, idxTileNext - idxTileCur));
if (warp_right_aggregates_count < expected_count)
{
break;
}
const bool use_value = lanemaskEq & warp_right_aggregates_mask;
const AccumT value = use_value ? regTmpStates[idx].value : cuda::identity_element<ScanOpT, AccumT>();
const AccumT local_aggr = warp_reduce_t{temp_storage}.Reduce(value, scan_op);
if (expected_count == 32)
{
aggrExclusiveCtaCur = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr);
idxTileCur += 32;
}
else
{
const AccumT full_aggr = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr);
idxTilePrev = idxTileCur;
aggrExclusiveCtaPrev = aggrExclusiveCtaCur;
return full_aggr;
}
}
}
// Only reached when idxTileNext is a multiple of 32; otherwise the final partial batch full aggregate returns inside
// the loop above.
idxTilePrev = idxTileNext;
aggrExclusiveCtaPrev = aggrExclusiveCtaCur;
return aggrExclusiveCtaCur; // must only be valid in lane_0
}
#endif // __cccl_ptx_isa >= 860
} // namespace detail::warpspeed
CUB_NAMESPACE_END

42
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/make_warp_uniform.cuh Normal file
View File

@@ -0,0 +1,42 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
// Move register to uniform register
// For int32_t and uint32_t, we can use the CREDUX instruction, which is coupled and has a constant latency.
// For 64-bit types, we still use __shfl_sync
[[nodiscard]] _CCCL_DEVICE_API inline int makeWarpUniform(int x)
{
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return __reduce_min_sync(~0, x);), (return x;));
}
[[nodiscard]] _CCCL_DEVICE_API inline ::cuda::std::uint32_t makeWarpUniform(::cuda::std::uint32_t x)
{
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return __reduce_min_sync(~0, x);), (return x;));
}
[[nodiscard]] _CCCL_DEVICE_API inline ::cuda::std::uint64_t makeWarpUniform(::cuda::std::uint64_t x)
{
return __shfl_sync(~0, x, 0);
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END

38
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/optimize_smem_ptr.cuh Normal file
View File

@@ -0,0 +1,38 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
template <typename _Tp>
[[nodiscard]] _CCCL_DEVICE_API _Tp* optimizeSmemPtr(const _Tp* smemGeneric)
{
// See https://nvbugspro.nvidia.com/bug/4907996
// 1. Convert to 32-bit shared memory pointer
::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(smemGeneric);
// 2. Pretend to NVVM that the 32-bit pointer is modified. This is required to avoid NVVM constant
// propagation from pulling the smem32 definition into loops and branches in subsequent code.
asm("" : "+r"(smem32));
// 3. Make a generic pointer to smem that is constructed using `__cvta_shared_to_generic`. This
// benefits from an
// optimization pass in NVVM that performs the following simplification:
// __cvta_generic_to_shared(__cvta_shared_to_generic(x)) => x.
// In our case, `x` is smem32, which is exactly what we want.
return reinterpret_cast<_Tp*>(__cvta_shared_to_generic(smem32));
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END

45
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_phase.cuh Normal file
View File

@@ -0,0 +1,45 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/resource/smem_ref.cuh>
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
template <typename _Tp>
struct SmemPhase
{
SmemResourceRaw& mSmemResourceRaw;
int mCurPhase;
_CCCL_DEVICE_API SmemPhase(SmemResourceRaw& smemResourceRaw, int phase) noexcept
: mSmemResourceRaw(smemResourceRaw)
, mCurPhase(phase)
{}
[[nodiscard]] _CCCL_DEVICE_API SmemRef<_Tp> acquireRef()
{
// Wait on barrier
mSmemResourceRaw.acquire(mCurPhase);
// Return ref
return SmemRef<_Tp>(mSmemResourceRaw, mCurPhase);
}
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

92
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_ref.cuh Normal file
View File

@@ -0,0 +1,92 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
#include <cub/detail/warpspeed/squad/squad.cuh>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
template <typename _Tp>
struct SmemRef
{
SmemResourceRaw& mSmemResourceRaw;
int mCurPhase;
bool mTxCountIsSet = false;
int mTxCount = 0;
bool mDoFenceLdsToAsyncProxy = false;
_CCCL_DEVICE_API SmemRef(SmemResourceRaw& smemResourceRaw, int phase) noexcept
: mSmemResourceRaw(smemResourceRaw)
, mCurPhase(phase)
{}
// SmemRef is a non-copyable, non-movable type. It must be passed by (mutable)
// reference to be useful. The reason is that it in case of an accidental copy
// or move the destructor is called twice. This leads to double-arrivals on
// barriers and results in deadlock or a hardware fault.
SmemRef(const SmemRef&) = delete; // Delete copy constructor
SmemRef(SmemRef&&) = delete; // Delete move constructor
SmemRef& operator=(const SmemRef&) = delete; // Delete copy assignment
SmemRef& operator=(SmemRef&&) = delete; // Delete move assignment
_CCCL_DEVICE_API ~SmemRef()
{
if (mDoFenceLdsToAsyncProxy)
{
mSmemResourceRaw.fenceLdsToAsyncProxy();
}
if (mTxCountIsSet)
{
mSmemResourceRaw.releaseTx(mCurPhase, mTxCount);
}
else
{
mSmemResourceRaw.release(mCurPhase);
}
}
[[nodiscard]] _CCCL_DEVICE_API _Tp& data() noexcept
{
return *static_cast<_Tp*>(mSmemResourceRaw.data());
}
[[nodiscard]] _CCCL_DEVICE_API int sizeBytes() const noexcept
{
return mSmemResourceRaw.mSizeBytes;
}
[[nodiscard]] _CCCL_DEVICE_API uint64_t* ptrCurBarrierRelease()
{
return mSmemResourceRaw.ptrCurBarrierRelease(mCurPhase);
}
_CCCL_DEVICE_API void squadIncreaseTxCount(const Squad& squad, int txCount)
{
mTxCountIsSet = true;
// Only leader thread increments txCount
txCount = squad.isLeaderThread() ? txCount : 0;
mTxCount += txCount;
}
_CCCL_DEVICE_API void setFenceLdsToAsyncProxy() noexcept
{
mDoFenceLdsToAsyncProxy = true;
}
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

60
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource.cuh Normal file
View File

@@ -0,0 +1,60 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/allocators/smem_allocator.cuh>
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
#include <cub/detail/warpspeed/resource/smem_stage.cuh>
#include <cub/detail/warpspeed/sync_handler.cuh>
#include <cub/detail/warpspeed/values.cuh>
#include <cuda/std/__utility/to_underlying.h>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
template <typename _Tp>
struct SmemResource : SmemResourceRaw
{
template <int stageCount>
_CCCL_HOST_DEVICE_API SmemResource(SyncHandler& syncHandler, _Tp (&smemBuffer)[stageCount])
: SmemResourceRaw(syncHandler, smemBuffer, sizeof(smemBuffer[0]), sizeof(smemBuffer[0]), stageCount)
{}
_CCCL_HOST_DEVICE_API constexpr SmemResource(
SyncHandler& syncHandler, SmemAllocator& smemAllocator, Stages stages, Elems elems = Elems{1})
: SmemResourceRaw(makeSmemResourceRaw(syncHandler, smemAllocator, stages, elems))
{}
[[nodiscard]] _CCCL_DEVICE_API SmemStage<_Tp> nextStage() noexcept
{
return SmemStage<_Tp>(*this);
}
private:
[[nodiscard]] _CCCL_HOST_DEVICE_API static constexpr SmemResourceRaw
makeSmemResourceRaw(SyncHandler& syncHandler, SmemAllocator& smemAllocator, Stages stages, Elems elems = Elems{1})
{
int align = alignof(_Tp);
int sizeBytes = ::cuda::std::to_underlying(elems) * sizeof(_Tp);
int strideBytes = sizeBytes;
void* ptrBase = smemAllocator.alloc(::cuda::std::to_underlying(stages) * strideBytes, align);
return {syncHandler, ptrBase, sizeBytes, strideBytes, ::cuda::std::to_underlying(stages)};
}
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

182
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource_raw.cuh Normal file
View File

@@ -0,0 +1,182 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/allocators/smem_allocator.cuh>
#include <cub/detail/warpspeed/constant_assert.cuh>
#include <cub/detail/warpspeed/squad/squad_desc.cuh>
#include <cub/detail/warpspeed/sync_handler.cuh>
#include <cuda/__ptx/instructions/fence.h>
#include <cuda/__ptx/instructions/mbarrier_arrive.h>
#include <cuda/__ptx/instructions/mbarrier_wait.h>
#include <cuda/__ptx/ptx_dot_variants.h>
#include <cuda/std/__type_traits/is_constant_evaluated.h>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
struct SmemResourceRaw
{
static constexpr int mMaxNumPhases = 4;
int mStageCurrent = 0;
int mResourceHandle;
::cuda::std::uint8_t* mPtrBase{};
int mSizeBytes;
int mStride;
int mStageCount;
int mNumPhases = 0;
::cuda::std::uint64_t* mPtrBar[mMaxNumPhases]{};
int mParity[mMaxNumPhases]{};
_CCCL_HOST_DEVICE_API constexpr SmemResourceRaw(
SyncHandler& syncHandler, void* ptrBase, int sizeBytes, int strideBytes, int stageCount) noexcept
: mResourceHandle(syncHandler.registerResource(stageCount))
, mSizeBytes(sizeBytes)
, mStride(strideBytes)
, mStageCount(stageCount)
{
// we don't need the pointer during constant evaluation (and casting is not allowed)
if (!::cuda::std::is_constant_evaluated())
{
mPtrBase = static_cast<::cuda::std::uint8_t*>(ptrBase);
}
for (int pi = 0; pi < mMaxNumPhases; ++pi)
{
mParity[pi] = pi == 0 ? 1 : 0;
}
}
template <int numSquads>
_CCCL_HOST_DEVICE_API constexpr void
addPhase(SyncHandler& syncHandler, ::cuda::std::uint64_t* ptrBarrier, const SquadDesc (&squads)[numSquads])
{
int numOwningThreads = squadCountThreads(squads);
int curPhase = mNumPhases;
mNumPhases++;
syncHandler.registerPhase(mResourceHandle, numOwningThreads, ptrBarrier);
mPtrBar[curPhase] = ptrBarrier;
}
template <int numSquads>
_CCCL_HOST_DEVICE_API constexpr void
addPhase(SyncHandler& syncHandler, SmemAllocator& smemAllocator, const SquadDesc (&squads)[numSquads])
{
void* ptrBar_raw = smemAllocator.alloc(mStageCount * sizeof(::cuda::std::uint64_t), alignof(::cuda::std::uint64_t));
// we don't need the pointer during constant evaluation (and casting is not allowed)
::cuda::std::uint64_t* ptrBar = nullptr;
if (!::cuda::std::is_constant_evaluated())
{
ptrBar = static_cast<::cuda::std::uint64_t*>(ptrBar_raw);
}
addPhase(syncHandler, ptrBar, squads);
}
_CCCL_HOST_DEVICE_API void
addPhase(SyncHandler& syncHandler, ::cuda::std::uint64_t* ptrBarrier, const SquadDesc& squad)
{
const SquadDesc squads[] = {squad};
addPhase(syncHandler, ptrBarrier, squads);
}
_CCCL_HOST_DEVICE_API constexpr void
addPhase(SyncHandler& syncHandler, SmemAllocator& smemAllocator, const SquadDesc& squad)
{
const SquadDesc squads[] = {squad};
addPhase(syncHandler, smemAllocator, squads);
}
_CCCL_DEVICE_API void incrementStage()
{
if (mStageCurrent == mStageCount - 1)
{
mStageCurrent = 0;
// We loop over all phases with a conditional on resNumPhases. If we
// directly loop over only resNumPhases, then the SROA optimization does
// not kick in and the mParity array is spilled to the stack.
for (int pi = 0; pi < mMaxNumPhases; ++pi)
{
if (pi < mNumPhases)
{
mParity[pi] ^= 1;
}
}
}
else
{
mStageCurrent++;
}
}
_CCCL_DEVICE_API void* data()
{
return (void*) (mPtrBase + mStageCurrent * mStride);
}
[[nodiscard]] _CCCL_DEVICE_API ::cuda::std::uint64_t* ptrCurBarrierRelease(int phase)
{
::cuda::std::uint64_t* ptrBarPhase = mPtrBar[phase];
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
return &ptrBarPhase[mStageCurrent];
}
_CCCL_DEVICE_API void release(int phase)
{
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
::cuda::ptx::mbarrier_arrive(ptrCurBarrierRelease(phase));
}
_CCCL_DEVICE_API void releaseTx(int phase, int txCount)
{
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
::cuda::ptx::mbarrier_arrive_expect_tx(
::cuda::ptx::sem_release, ::cuda::ptx::scope_cta, ::cuda::ptx::space_shared, ptrCurBarrierRelease(phase), txCount);
}
_CCCL_DEVICE_API void fenceLdsToAsyncProxy()
{
::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared);
}
_CCCL_DEVICE_API void releaseLdsToAsyncProxy(int phase)
{
// First fence
fenceLdsToAsyncProxy();
// Then perform a normal release
release(phase);
}
_CCCL_DEVICE_API void acquire(int phase)
{
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
// The release of the previous phase occurs on the `phase - 1` barrier. So
// that is what we wait on.
int phaseAcq = (mNumPhases + phase - 1) % mNumPhases;
::cuda::std::uint64_t* ptrBarPhase = mPtrBar[phaseAcq];
while (!::cuda::ptx::mbarrier_try_wait_parity(&ptrBarPhase[mStageCurrent], mParity[phase]))
{
}
}
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

91
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_stage.cuh Normal file
View File

@@ -0,0 +1,91 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/constant_assert.cuh>
#include <cub/detail/warpspeed/resource/smem_phase.cuh>
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
#include <cuda/std/__tuple_dir/tuple_element.h>
#include <cuda/std/__tuple_dir/tuple_size.h>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
template <typename _Tp>
struct SmemStage
{
SmemResourceRaw& mSmemResourceRaw;
_CCCL_DEVICE_API SmemStage(SmemResourceRaw& smemResourceRaw) noexcept
: mSmemResourceRaw(smemResourceRaw)
{}
_CCCL_DEVICE_API ~SmemStage()
{
mSmemResourceRaw.incrementStage();
}
// SmemStage is a non-copyable, non-movable type. It must be passed by (mutable)
// reference to be useful. The reason is that it in case of an accidental copy
// or move the destructor is called twice. This leads to double-increment of
// the stage index and results in deadlock or a hardware fault.
SmemStage(const SmemStage&) = delete; // Delete copy constructor
SmemStage(SmemStage&&) = delete; // Delete move constructor
SmemStage& operator=(const SmemStage&) = delete; // Delete copy assignment
SmemStage& operator=(const SmemStage&&) = delete; // Delete move assignment
};
// Helper: Container to expose SmemPhase for structured binding
template <typename _Tp, ::cuda::std::size_t numPhases>
struct SmemPhaseStructuredBinding
{
SmemResourceRaw& mSmemResourceRaw;
template <::cuda::std::size_t _Index>
[[nodiscard]] _CCCL_DEVICE_API SmemPhase<_Tp> get() const
{
return SmemPhase<_Tp>(mSmemResourceRaw, _Index);
}
};
// The binding function
template <::cuda::std::size_t numPhases, typename _Tp>
[[nodiscard]] _CCCL_DEVICE_API SmemPhaseStructuredBinding<_Tp, numPhases> bindPhases(SmemStage<_Tp>& smemStage)
{
_WS_CONSTANT_ASSERT(smemStage.mSmemResourceRaw.mNumPhases == numPhases,
"Number of bound phases must match resource phases.");
return SmemPhaseStructuredBinding<_Tp, numPhases>{smemStage.mSmemResourceRaw};
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END
// Tuple protocol specializations
namespace std
{
template <typename _Tp, size_t numPhases>
struct tuple_size<CUB_NS_QUALIFIER::detail::warpspeed::SmemPhaseStructuredBinding<_Tp, numPhases>>
{
static constexpr size_t value = numPhases;
};
template <typename _Tp, size_t _Index, ::cuda::std::size_t numPhases>
struct tuple_element<_Index, CUB_NS_QUALIFIER::detail::warpspeed::SmemPhaseStructuredBinding<_Tp, numPhases>>
{
using type = CUB_NS_QUALIFIER::detail::warpspeed::SmemPhase<_Tp>;
};
} // namespace std

44
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/special_registers.cuh Normal file
View File

@@ -0,0 +1,44 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/make_warp_uniform.cuh>
#include <cuda/__ptx/instructions/get_sreg.h>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
// Commonly used special registers that we should cache in registers or uniform
// registers.
struct SpecialRegisters
{
const ::cuda::std::uint32_t clusterCtaRank;
const ::cuda::std::uint32_t blockIdxX;
const ::cuda::std::uint32_t threadIdxX;
const ::cuda::std::uint32_t warpIdx;
const ::cuda::std::uint32_t laneIdx;
};
[[nodiscard]] _CCCL_DEVICE_API inline SpecialRegisters getSpecialRegisters()
{
::cuda::std::uint32_t clusterCtaRank = ::cuda::ptx::get_sreg_cluster_ctarank();
::cuda::std::uint32_t threadIdxX = threadIdx.x;
::cuda::std::uint32_t warpIdx = makeWarpUniform(threadIdxX / 32);
return {clusterCtaRank, blockIdx.x, threadIdxX, warpIdx, ::cuda::ptx::get_sreg_laneid()};
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END

388
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/load_store.cuh Normal file
View File

@@ -0,0 +1,388 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/resource/smem_ref.cuh>
#include <cub/detail/warpspeed/squad/squad.cuh>
#include <cuda/__memory/align_down.h>
#include <cuda/__memory/align_up.h>
#include <cuda/__ptx/instructions/cp_async_bulk.h>
#include <cuda/__ptx/instructions/cp_async_bulk_commit_group.h>
#include <cuda/__ptx/instructions/cp_async_bulk_wait_group.h>
#include <cuda/__ptx/instructions/elect_sync.h>
#include <cuda/__ptx/instructions/fence.h>
#include <cuda/std/__type_traits/make_nbit_int.h>
#include <cuda/std/cstdint>
#include <nv/target>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
#if __cccl_ptx_isa >= 860
template <typename Tp>
struct CpAsyncOobInfo
{
// The aligned up and down pointers below must be ::cuda::std::byte*, since the nearest aligned up/down ptr may not
// point to a multiple of sizeof(Tp). E.g. a uchar3* pointing to address 0x...5 will be aligned down to 0x...0, but
// that's not a valid start for an uchar3 in that array. So we must express all aligned pointers in bytes here.
::cuda::std::byte* ptrGmem;
::cuda::std::byte* ptrGmemStartAlignDown;
::cuda::std::byte* ptrGmemStartAlignUp;
::cuda::std::byte* ptrGmemEnd;
::cuda::std::byte* ptrGmemEndAlignDown;
::cuda::std::byte* ptrGmemEndAlignUp;
::cuda::std::uint32_t overCopySizeBytes;
::cuda::std::uint32_t underCopySizeBytes;
::cuda::std::uint32_t origCopySizeBytes;
::cuda::std::uint32_t smemStartSkipBytes; // ptrSmem + smemStartSkipBytes will point to the first valid element copied
// from ptrGmem
::cuda::std::uint32_t smemEndBytesAfter16BBoundary; // number of bytes after the last 16B boundary in GMEM/SMEM that
// still contains valid (partial) elements
};
template <typename Tp>
_CCCL_DEVICE_API _CCCL_FORCEINLINE CpAsyncOobInfo<Tp> prepareCpAsyncOob(Tp* ptrGmem, ::cuda::std::uint32_t sizeElem)
{
auto ptrGmemBytes = reinterpret_cast<::cuda::std::byte*>(ptrGmem);
auto ptrGmemEnd = reinterpret_cast<::cuda::std::byte*>(ptrGmem + sizeElem);
// We will copy from [ptrGmemBase, ptrGmemEnd). Both pointers have to be 16B aligned.
::cuda::std::byte* ptrGmemStartAlignDown = ::cuda::align_down(ptrGmemBytes, 16);
::cuda::std::byte* ptrGmemStartAlignUp = ::cuda::align_up(ptrGmemBytes, 16);
::cuda::std::byte* ptrGmemEndAlignUp = ::cuda::align_up(ptrGmemEnd, 16);
::cuda::std::byte* ptrGmemEndAlignDown = ::cuda::align_down(ptrGmemEnd, 16);
// Compute the final copy size in bytes. It can be either sizeElem or sizeElem + 16 / sizeof(T).
const auto origCopySizeBytes = static_cast<::cuda::std::uint32_t>(sizeof(Tp) * sizeElem);
const auto overCopySizeBytes = static_cast<::cuda::std::uint32_t>(ptrGmemEndAlignUp - ptrGmemStartAlignDown);
auto underCopySizeBytes = static_cast<::cuda::std::uint32_t>(ptrGmemEndAlignDown - ptrGmemStartAlignUp);
if (origCopySizeBytes < underCopySizeBytes)
{
// If ptrGmemStart and ptrGmemEnd are aligned to [1, .., 15] bytes, then
// when we align the one up and the other down we get overflow. We check for
// that here. In that case, the undercopy size is zero.
underCopySizeBytes = 0;
}
_CCCL_ASSERT(overCopySizeBytes % 16 == 0, "");
_CCCL_ASSERT(underCopySizeBytes % 16 == 0, "");
return {
ptrGmemBytes,
ptrGmemStartAlignDown,
ptrGmemStartAlignUp,
ptrGmemEnd,
ptrGmemEndAlignDown,
ptrGmemEndAlignUp,
overCopySizeBytes,
underCopySizeBytes,
origCopySizeBytes,
static_cast<::cuda::std::uint32_t>(ptrGmemBytes - ptrGmemStartAlignDown),
static_cast<::cuda::std::uint32_t>(ptrGmemEnd - ptrGmemEndAlignDown),
};
}
template <typename ResourceTp, typename Tp>
_CCCL_DEVICE_API void squadLoadBulk(Squad squad, SmemRef<ResourceTp>& refDestSmem, CpAsyncOobInfo<Tp> cpAsyncOobInfo)
{
::cuda::std::byte* ptrSmem = refDestSmem.data().inout;
_CCCL_ASSERT(::cuda::is_aligned(ptrSmem, 16), "");
::cuda::std::uint64_t* ptrBar = refDestSmem.ptrCurBarrierRelease();
if constexpr (alignof(Tp) >= 16)
{
// for alignments larger than 16, we can just bulk copy, even just a single element
if (squad.isLeaderThread())
{
::cuda::ptx::cp_async_bulk(
::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{},
::cuda::ptx::space_global,
ptrSmem,
cpAsyncOobInfo.ptrGmem,
cpAsyncOobInfo.origCopySizeBytes,
ptrBar);
}
refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.underCopySizeBytes);
}
else
{
// for alignments smaller than 16, we can overcopy but need to declare the ignored bytes left and right
# if __cccl_ptx_isa >= 920
if (squad.isLeaderThread())
{
::cuda::ptx::cp_async_bulk_ignore_oob(
::cuda::ptx::space_shared,
::cuda::ptx::space_global,
ptrSmem,
cpAsyncOobInfo.ptrGmemStartAlignDown,
cpAsyncOobInfo.overCopySizeBytes,
/* ignore left */ cpAsyncOobInfo.smemStartSkipBytes,
/* ignore right */ cpAsyncOobInfo.ptrGmemEndAlignUp - cpAsyncOobInfo.ptrGmemEnd,
ptrBar);
}
refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.overCopySizeBytes);
# else // __cccl_ptx_isa >= 920
// if we don't have cp_async_bulk_ignore_oob, we have to undercopy and copy head and tail elements manually
// handle small copies first. If we have less than 16 bytes we may not straddle a 16B boundary
if (cpAsyncOobInfo.origCopySizeBytes < 16)
{
const auto elemCount = cpAsyncOobInfo.origCopySizeBytes / sizeof(Tp);
_CCCL_ASSERT(elemCount <= squad.threadCount(), "");
if (squad.threadRank() < elemCount)
{
reinterpret_cast<Tp*>(ptrSmem + cpAsyncOobInfo.smemStartSkipBytes)[squad.threadRank()] =
reinterpret_cast<const Tp*>(cpAsyncOobInfo.ptrGmem)[squad.threadRank()];
}
return; // no bulk copy has been performed so we don't need to update the tx count of any barrier
}
// copies larger than 16 byte which straddle at least one 16B boundary, so we have dedicated start and end copies
const bool doStartCopy = cpAsyncOobInfo.smemStartSkipBytes > 0;
::cuda::std::byte* ptrSmemMiddle = ptrSmem;
if (doStartCopy)
{
ptrSmemMiddle += 16;
}
// TODO(bgruber): we could skip the middle if underCopySizeBytes is zero
if (squad.isLeaderThread())
{
::cuda::ptx::cp_async_bulk(
::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{},
::cuda::ptx::space_global,
ptrSmemMiddle,
cpAsyncOobInfo.ptrGmemStartAlignUp,
cpAsyncOobInfo.underCopySizeBytes,
ptrBar);
}
refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.underCopySizeBytes);
// we cannot use Tp to load the head and tail elements, because sizeof(Tp) may be larger than alignof(Tp)
using load_word_t = ::cuda::std::__make_nbit_uint_t<alignof(Tp) * CHAR_BIT>;
const int head_elements = (cpAsyncOobInfo.ptrGmemStartAlignUp - cpAsyncOobInfo.ptrGmem) / sizeof(load_word_t);
const int tail_elements = (cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemEndAlignDown) / sizeof(load_word_t);
_CCCL_ASSERT(head_elements <= squad.threadCount(), "");
_CCCL_ASSERT(tail_elements <= squad.threadCount(), "");
load_word_t head_value, tail_value;
if (squad.threadRank() < head_elements)
{
head_value = reinterpret_cast<const load_word_t*>(cpAsyncOobInfo.ptrGmem)[squad.threadRank()];
}
if (squad.threadRank() < tail_elements)
{
tail_value = reinterpret_cast<const load_word_t*>(cpAsyncOobInfo.ptrGmemEndAlignDown)[squad.threadRank()];
}
if (squad.threadRank() < head_elements)
{
reinterpret_cast<load_word_t*>(ptrSmem + cpAsyncOobInfo.smemStartSkipBytes)[squad.threadRank()] = head_value;
}
if (squad.threadRank() < tail_elements)
{
reinterpret_cast<load_word_t*>(ptrSmemMiddle + cpAsyncOobInfo.underCopySizeBytes)[squad.threadRank()] =
tail_value;
}
# endif // __cccl_ptx_isa >= 920
}
}
_CCCL_DEVICE_API _CCCL_FORCEINLINE void squadStoreMasked16B(
Squad squad,
::cuda::std::byte* dstGmem,
const ::cuda::std::byte* srcSmem,
::cuda::std::uint16_t byteMask,
int firstByte,
int lastByte)
{
NV_IF_ELSE_TARGET(
NV_PROVIDES_SM_100,
(if (::cuda::ptx::elect_sync(~0)) {
::cuda::ptx::cp_async_bulk_cp_mask(
::cuda::ptx::space_global, ::cuda::ptx::space_shared, dstGmem, srcSmem, /*size*/ 16, byteMask);
}),
({
const int rank = squad.threadRank();
if (firstByte <= rank && rank < lastByte)
{
dstGmem[rank] = srcSmem[rank];
}
}));
}
template <typename OutputT>
_CCCL_DEVICE_API void
squadStoreBulkSync(Squad squad, CpAsyncOobInfo<OutputT> cpAsyncOobInfo, const ::cuda::std::byte* srcSmem)
{
// This function performs either 1 copy, or three copies, depending on the
// size and alignment of the output tile in global memory.
//
// If the output tile is contained in a single 16-byte aligned and sized region, then we
// only perform a single masked copy.
//
// If the output tile is larger than 16 bytes or straddles two 16-byte aligned and sized regions, then
// we perform up to three copies:
// - One copy for the first up to 15 bytes at the start of the region.
// - One copy that starts at a 16-byte aligned address and ends at the latest 16-byte aligned address.
// - One copy that cleans up the last up to 15 bytes.
if (squad.isLeaderWarp())
{
// Acquire shared memory in async proxy
// Perform fence.proxy.async with full warp to avoid BSSY+BSYNC
::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared);
# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
// for some reason the optimizer propagates some information from the computation of
// overCopySizeBytes to the masked bulk copy below and generates an unaligned access error.
// The artificial read modification of overCopySizeBytes prevents the propagation here works around this.
// It also solves the issue described in nvbug 5848313 by accident on nvcc 13.2+
asm volatile("" : "+r"(cpAsyncOobInfo.overCopySizeBytes));
# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
const bool doStartCopy = cpAsyncOobInfo.smemStartSkipBytes > 0;
const bool doEndCopy = cpAsyncOobInfo.smemEndBytesAfter16BBoundary > 0;
const bool doMiddleCopy = cpAsyncOobInfo.ptrGmemStartAlignUp != cpAsyncOobInfo.ptrGmemEndAlignUp;
constexpr ::cuda::std::uint16_t byteMask = 0xFFFF;
const ::cuda::std::uint16_t byteMaskStart = byteMask << cpAsyncOobInfo.smemStartSkipBytes;
const ::cuda::std::uint16_t byteMaskEnd = byteMask >> (16 - cpAsyncOobInfo.smemEndBytesAfter16BBoundary) % 16;
// byteMaskStart contains zeroes at the left
# if _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
const ::cuda::std::uint16_t byteMaskSmall = byteMaskStart & byteMaskEnd;
# else // _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
// `ptxas fatal : (C7907) Internal compiler error`, see nvbug 5848313
const ::cuda::std::uint16_t byteMaskSmall =
byteMaskStart & (byteMask >> (16 - (cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemStartAlignDown)));
# endif // _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
const ::cuda::std::byte* ptrSmemMiddle = srcSmem;
if (doStartCopy)
{
ptrSmemMiddle += 16;
}
if (doMiddleCopy)
{
// Copy the middle part. Starting at byte 0 or 16 in shared memory. This
// is the large copy. We perform this one first, so that the compiler can
// (hopefully) hide all the arithmetic behind this instruction.
if (::cuda::ptx::elect_sync(~0))
{
// need to work around another optimizer bug, see: https://github.com/NVIDIA/cccl/issues/8644
# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemStartAlignUp));
# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
::cuda::ptx::cp_async_bulk(
::cuda::ptx::space_global,
::cuda::ptx::space_shared,
cpAsyncOobInfo.ptrGmemStartAlignUp,
ptrSmemMiddle,
cpAsyncOobInfo.underCopySizeBytes);
}
if (doStartCopy)
{
// need to work around yet another optimizer bug, see: https://github.com/NVIDIA/cccl/issues/8838
# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemStartAlignDown));
asm volatile("" : "+l"(srcSmem));
# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
// Copy a subset of the first 16 bytes
squadStoreMasked16B(
squad,
cpAsyncOobInfo.ptrGmemStartAlignDown,
srcSmem,
byteMaskStart,
static_cast<int>(cpAsyncOobInfo.smemStartSkipBytes),
16);
}
if (doEndCopy)
{
# if _CCCL_CUDA_COMPILER(NVHPC)
// nvc++ seems to have an optimizer bug, crashing with an unaligned access error below. The addresses are fine
// when printed, so let's shake the optimizer a bit.
asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemEndAlignDown));
# endif // _CCCL_CUDA_COMPILER(NVHPC)
// Copy a subset of the last 16 bytes
squadStoreMasked16B(
squad,
cpAsyncOobInfo.ptrGmemEndAlignDown,
ptrSmemMiddle + cpAsyncOobInfo.underCopySizeBytes,
byteMaskEnd,
0,
static_cast<int>(cpAsyncOobInfo.smemEndBytesAfter16BBoundary));
}
}
else
{
// Copy a subset of the first 16 bytes
squadStoreMasked16B(
squad,
cpAsyncOobInfo.ptrGmemStartAlignDown,
srcSmem,
byteMaskSmall,
static_cast<int>(cpAsyncOobInfo.smemStartSkipBytes),
static_cast<int>(cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemStartAlignDown));
}
// Commit and wait for store to have completed reading from shared memory
::cuda::ptx::cp_async_bulk_commit_group();
::cuda::ptx::cp_async_bulk_wait_group_read(::cuda::ptx::n32_t<0>{});
}
}
#endif // __cccl_ptx_isa >= 860
template <typename InputT, typename AccumT, int ElemPerThread>
_CCCL_DEVICE_API void squadLoadSmem(Squad squad, AccumT (&outReg)[ElemPerThread], const InputT* smemBuf)
{
for (int i = 0; i < ElemPerThread; ++i)
{
const int elem_idx = squad.threadRank() * ElemPerThread + i;
outReg[i] = smemBuf[elem_idx];
}
}
template <typename OutputT, typename AccumT, int ElemPerThread>
_CCCL_DEVICE_API void squadStoreSmem(Squad squad, OutputT* smemBuf, const AccumT (&inReg)[ElemPerThread])
{
for (int i = 0; i < ElemPerThread; ++i)
{
const int elem_idx = squad.threadRank() * ElemPerThread + i;
smemBuf[elem_idx] = inReg[i];
}
}
template <typename OutputT, typename AccumT, int ElemPerThread>
_CCCL_DEVICE_API void
squadStoreSmemPartial(Squad squad, OutputT* smemBuf, const AccumT (&inReg)[ElemPerThread], int beginIndex, int endIndex)
{
for (int i = 0; i < ElemPerThread; ++i)
{
const int elem_idx = squad.threadRank() * ElemPerThread + i;
if (beginIndex <= elem_idx && elem_idx < endIndex)
{
smemBuf[elem_idx - beginIndex] = inReg[i];
}
}
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END

159
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad.cuh Normal file
View File

@@ -0,0 +1,159 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/special_registers.cuh>
#include <cub/detail/warpspeed/squad/squad_desc.cuh>
#include <cuda/__ptx/instructions/elect_sync.h>
#include <cuda/std/array>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
// Squad - device squad instance
//
// A squad is a collection of warps that work together in a warp-specialized
// kernel. A warp-specialized kernel has multiple squads that perform part of
// the computation.
//
// The Squad class is a device runtime instance of a squad. It provides
// functionality to determine the rank of the current thread or warp in the
// squad, and to sync all threads in the squad.
struct Squad : SquadDesc
{
SpecialRegisters mSpecialRegisters;
bool mIsWarpLeader = false;
bool mIsLeaderWarp = false;
_CCCL_DEVICE_API Squad(SquadDesc squadStatic, SpecialRegisters specialRegisters)
: SquadDesc(squadStatic)
, mSpecialRegisters(specialRegisters)
{
mIsWarpLeader = ::cuda::ptx::elect_sync(~0);
mIsLeaderWarp = warpRank() == 0;
}
[[nodiscard]] _CCCL_DEVICE_API int warpRank() const
{
return static_cast<int>(mSpecialRegisters.warpIdx % this->warpCount());
}
[[nodiscard]] _CCCL_DEVICE_API int threadRank() const
{
return static_cast<int>(mSpecialRegisters.threadIdxX % this->threadCount());
}
[[nodiscard]] _CCCL_DEVICE_API bool isLeaderThread() const
{
return mIsWarpLeader && mIsLeaderWarp;
}
[[nodiscard]] _CCCL_DEVICE_API bool isLeaderWarp() const
{
return mIsLeaderWarp;
}
[[nodiscard]] _CCCL_DEVICE_API bool isLeaderThreadOfWarp() const
{
return mIsWarpLeader;
}
_CCCL_DEVICE_API void syncThreads() const
{
// barrier 0 is reserved for __syncthreads(). We use barrier ids 1, ...
const int barrierIdx = this->mSquadIdx + 1;
__barrier_sync_count(barrierIdx, this->threadCount());
}
};
// squadDispatch
//
// squadDispatch is used at the start of the kernel. It takes an array of squad
// descriptors and determines which squad the current thread belongs to. The
// lambda `f: (Squad) -> void` is called with the squad currently active on this
// thread.
//
// Typically, the user will call the kernel body with the active squad.
//
// Implementation notes:
//
// Dispatch to squad based on warp index using a binary search. This balances
// the number of BRA instructions per squad and avoids NVVM inserting BRX
// instructions. BRX instructions require a jump table that is loaded from GCC,
// which incurs latency.
//
// The benefit of this function for fastScan is that adding a new squad doesn't
// require code changes in the dispatch. For low-latency inference, I hope that
// the avoidance of linear search and BRX instructions translates into latency
// reductions.
//
template <int numSquads, typename F>
_CCCL_DEVICE_API _CCCL_FORCEINLINE void
squadDispatch(SpecialRegisters sr, const SquadDesc (&squads)[numSquads], F f, int warpIdxStart = 0)
{
static_assert(numSquads > 0);
if (numSquads == 1)
{
// Leaf
SquadDesc squad = squads[0];
if (static_cast<unsigned>(warpIdxStart) <= sr.warpIdx
&& sr.warpIdx < static_cast<unsigned>(warpIdxStart + squad.warpCount()))
{
f(Squad(squad, sr));
}
}
else
{
constexpr int mid = numSquads / 2;
// Left
int warpIdxStartMid = warpIdxStart;
for (int gi = 0; gi < mid; ++gi)
{
warpIdxStartMid += squads[gi].warpCount();
}
if (sr.warpIdx < static_cast<unsigned>(warpIdxStartMid))
{
if constexpr (0 < mid)
{
SquadDesc squadsLeft[mid];
for (int gi = 0; gi < mid; ++gi)
{
squadsLeft[gi] = squads[gi];
}
squadDispatch(sr, squadsLeft, f, warpIdxStart);
}
}
else
{
SquadDesc squadsRight[numSquads - mid]{};
for (int gi = 0; gi < numSquads - mid; ++gi)
{
squadsRight[gi] = squads[mid + gi];
}
squadDispatch(sr, squadsRight, f, warpIdxStartMid);
}
}
}
template <::cuda::std::size_t numSquads, typename F>
_CCCL_DEVICE_API _CCCL_FORCEINLINE void
squadDispatch(SpecialRegisters sr, ::cuda::std::array<SquadDesc, numSquads> squads, F f, int warpIdxStart = 0)
{
squadDispatch<numSquads>(sr, squads.__elems_, f, warpIdxStart);
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END

77
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad_desc.cuh Normal file
View File

@@ -0,0 +1,77 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
// SquadDesc - squad descriptor
//
// A squad is a collection of warps that work together in a warp-specialized
// kernel. A warp-specialized kernel has multiple squads that perform part of
// the computation.
//
// SquadDesc is a host+device constexpr-compatible class that allows describing
// the warp-specialized layout of a kernel.
//
// SquadDesc is constexpr-compatible and can be created on host and device.
struct SquadDesc
{
int mSquadIdx = -1;
int mWarpCount = -1;
_CCCL_HIDE_FROM_ABI constexpr SquadDesc() = default;
_CCCL_HOST_DEVICE_API constexpr SquadDesc(int squadIdx, int warpCount) noexcept
: mSquadIdx(squadIdx)
, mWarpCount(warpCount)
{}
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int warpCount() const noexcept
{
return mWarpCount;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int threadCount() const noexcept
{
return 32 * warpCount();
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(SquadDesc lhs, SquadDesc rhs) noexcept
{
return lhs.mSquadIdx == rhs.mSquadIdx;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(SquadDesc lhs, SquadDesc rhs) noexcept
{
return lhs.mSquadIdx != rhs.mSquadIdx;
}
};
// squadCountThreads
//
// Utility function to count the number of threads in an array of squad
// descriptors. It is used to launch a kernel with the correct number of
// threads.
template <int numSquads>
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int squadCountThreads(const SquadDesc (&squads)[numSquads]) noexcept
{
int sumThreads = 0;
for (int gi = 0; gi < numSquads; ++gi)
{
sumThreads += squads[gi].threadCount();
}
return sumThreads;
}
} // namespace detail::warpspeed
CUB_NAMESPACE_END

144
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/sync_handler.cuh Normal file
View File

@@ -0,0 +1,144 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/constant_assert.cuh>
#include <cub/detail/warpspeed/special_registers.cuh>
#include <cuda/__ptx/instructions/mbarrier_init.h>
#include <cuda/std/cstdint>
#include <nv/target>
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
// SkipSync is a tag type that is used to indicate that a SyncHandler.blockInit
// should forgo syncing.
struct SkipSync
{};
struct SyncHandler
{
// reducing these values to the actually used number of resources and phases does not improve performance
static constexpr int mMaxNumResources = 10;
static constexpr int mMaxNumPhases = 4;
// Whether barriers have been initialized.
bool mHasInitialized = false;
// Arrays of barrier locations, number of stages, number of owning threads.
int mNextResourceHandle = 0;
int mNumStages[mMaxNumResources]{};
int mNumPhases[mMaxNumResources]{};
int mNumOwningThreads[mMaxNumResources][mMaxNumPhases]{};
::cuda::std::uint64_t* mPtrBar[mMaxNumResources][mMaxNumPhases]{};
constexpr SyncHandler() = default;
// we need constant destruction for the host side single stage SMEM amount, which is only possible in C++20
#if _CCCL_STD_VER >= 2020
_CCCL_HOST_DEVICE_API constexpr ~SyncHandler()
{
_WS_CONSTANT_ASSERT(mHasInitialized, "SyncHandler must have been initialized at end of kernel.");
}
#endif // _CCCL_STD_VER >= 2020
// SyncHandler is a non-copyable, non-movable type. It must be passed by
// (mutable) reference to be useful.
SyncHandler(const SyncHandler&) = delete; // Delete copy constructor
SyncHandler(SyncHandler&&) = delete; // Delete move constructor
SyncHandler& operator=(const SyncHandler&) = delete; // Delete copy assignment
SyncHandler& operator=(const SyncHandler&&) = delete; // Delete move assignment
// registerResource and registerPhase can be called on host and device.
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int registerResource(int numStages)
{
_WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot register resource after SyncHandler has been initialized.");
// Avoid exceeding the max number of stages
_WS_CONSTANT_ASSERT(mNextResourceHandle < mMaxNumResources, "Cannot register more than 10 resources.");
// Get a handle
int handle = mNextResourceHandle;
mNextResourceHandle++;
// Set the number of stages
mNumStages[handle] = numStages;
return handle;
}
_CCCL_HOST_DEVICE_API void constexpr registerPhase(int resourceHandle, int numOwningThreads, uint64_t* ptrBar)
{
_WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot register phase after SyncHandler has been initialized.");
_WS_CONSTANT_ASSERT(resourceHandle < mNextResourceHandle, "Invalid resource handle.");
// Get phase index:
int curPhase = mNumPhases[resourceHandle];
_WS_CONSTANT_ASSERT(curPhase < mMaxNumPhases, "Cannot register more phases than maximum.");
mNumOwningThreads[resourceHandle][curPhase] = numOwningThreads;
mPtrBar[resourceHandle][curPhase] = ptrBar;
mNumPhases[resourceHandle]++;
}
// clusterInitSync can only be called on device.
template <int NumThreads>
_CCCL_DEVICE_API void clusterInitSync(SpecialRegisters sr, SkipSync)
{
_WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot initialize SyncHandler twice.");
mHasInitialized = true;
// All warps iterate through all resources and phases. Since all array indices have to be statically resolved by the
// SROA optimization to avoid spilling to local memory, we cannot split the iteration among warps etc.
for (int ri = 0; ri < mMaxNumResources; ri++)
{
if (ri >= mNextResourceHandle)
{
break;
}
const int resNumPhases = mNumPhases[ri];
const int numStages = mNumStages[ri];
for (int pi = 0; pi < mMaxNumPhases; pi++)
{
if (pi >= resNumPhases)
{
break;
}
uint64_t* ptrBar = mPtrBar[ri][pi];
int numOwningThreads = mNumOwningThreads[ri][pi];
// use block strided iteration to vectorize setup of barriers
for (int si = static_cast<int>(sr.threadIdxX); si < numStages; si += NumThreads)
{
::cuda::ptx::mbarrier_init(&ptrBar[si], numOwningThreads);
}
}
}
}
template <int NumThreads>
_CCCL_DEVICE_API void clusterInitSync(SpecialRegisters sr)
{
NV_IF_TARGET(NV_PROVIDES_SM_90, ({
clusterInitSync<NumThreads>(sr, SkipSync{});
__cluster_barrier_arrive_relaxed();
__cluster_barrier_wait();
}))
}
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

33
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/values.cuh Normal file
View File

@@ -0,0 +1,33 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
CUB_NAMESPACE_BEGIN
namespace detail::warpspeed
{
enum class Stages : int
{
};
enum class Elems : int
{
};
enum class Warps : int
{
};
enum class Align : int
{
};
} // namespace detail::warpspeed
CUB_NAMESPACE_END

27
qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/warpspeed.cuh Normal file
View File

@@ -0,0 +1,27 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/warpspeed/allocators/smem_allocator.cuh>
#include <cub/detail/warpspeed/constant_assert.cuh>
#include <cub/detail/warpspeed/make_warp_uniform.cuh>
#include <cub/detail/warpspeed/resource/smem_phase.cuh>
#include <cub/detail/warpspeed/resource/smem_ref.cuh>
#include <cub/detail/warpspeed/resource/smem_resource.cuh>
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
#include <cub/detail/warpspeed/resource/smem_stage.cuh>
#include <cub/detail/warpspeed/special_registers.cuh>
#include <cub/detail/warpspeed/squad/squad.cuh>
#include <cub/detail/warpspeed/squad/squad_desc.cuh>
#include <cub/detail/warpspeed/sync_handler.cuh>
#include <cub/detail/warpspeed/values.cuh>