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

test: xllm CUDA kernel verification suite for BI-V100

Browse Source 
test_xllm_cuda_kernels.py — 7 test groups:
  1. activation.cu: silu_and_mul via ixf_F, compare vs torch.nn.functional.silu
  2. norm.cu: rms_norm + fused_add_rms_norm via ixf_F, compare vs PyTorch
  3. rope.cu: rotary_embedding via ixf_F, verify rotation applied
  4. moe_topk_softmax: corex .so, verify shapes + weights sum to 1
  5. ix_moe_bridge: full 7-step fused MoE pipeline (topk→expand→gemm→act→gemm→combine)
  6. ix_attn_bridge: load test (prefill_attention, decode_attention, linear)
  7. ix_full_bridge: silu_and_mul + rms_norm through bridge .so

Revert: undo unnecessary cccl_upstream sync (already up to date)

Run on real machine: python3 qwen3_6_scripts/test_xllm_cuda_kernels.py
This commit is contained in:
claude
2026-08-14 08:01:23 +00:00
parent 8d75652949
commit 3a2cfc87c9
27 changed files with 911 additions and 1625 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh
View File

@@ -246,15 +246,14 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceAdjacentDifference"
error = CubDebug(
THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(
num_tiles, AdjacentDifferencePolicyT::BLOCK_THREADS, 0, stream)
.doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel<
KernelPolicySelector,
InputIteratorT,
OutputIteratorT,
DifferenceOpT,
OffsetT,
InputT,
AliasOpt == MayAlias::Yes,
ReadOpt == ReadOption::Left>,
.doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel < KernelPolicySelector,
InputIteratorT,
OutputIteratorT,
DifferenceOpT,
OffsetT,
InputT,
AliasOpt == MayAlias::Yes,
ReadOpt == ReadOption::Left >,
d_input,
first_tile_previous,
d_output,
@@ -439,15 +438,14 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
if (const auto error = CubDebug(
THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_tiles, active_policy.threads_per_block, 0, stream)
.doit(DeviceAdjacentDifferenceDifferenceKernel<
policy_selector_t,
InputIteratorT,
OutputIteratorT,
DifferenceOpT,
offset_t,
input_t,
AliasOpt == MayAlias::Yes,
ReadOpt == ReadOption::Left>,
.doit(DeviceAdjacentDifferenceDifferenceKernel < policy_selector_t,
InputIteratorT,
OutputIteratorT,
DifferenceOpT,
offset_t,
input_t,
AliasOpt == MayAlias::Yes,
ReadOpt == ReadOption::Left >,
d_input,
first_tile_previous,
d_output,

4
cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh
View File

@@ -152,8 +152,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().lookback.large_buffer.thr
{
if (thread_offset < buffer_sizes[buffer_id])
{
const auto value =
read_item<MemcpyOpt == CopyAlg::Memcpy, AliasT, InputBufferT>(input_buffer_it[buffer_id], thread_offset);
const auto value = read_item < MemcpyOpt == CopyAlg::Memcpy, AliasT,
InputBufferT > (input_buffer_it[buffer_id], thread_offset);
write_item<MemcpyOpt == CopyAlg::Memcpy, AliasT, OutputBufferT>(
output_buffer_it[buffer_id], thread_offset, value);
}

1244
cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh
View File

File diff suppressed because it is too large Load Diff

85
cccl_upstream/cub/cub/device/dispatch/dispatch_copy_mdspan.cuh
View File

@@ -21,14 +21,8 @@
#include <cub/device/dispatch/tuning/tuning_transform.cuh>
#include <cub/util_debug.cuh>
#include <cuda/__algorithm/copy.h>
#include <cuda/__functional/always_true_false.h>
#include <cuda/__stream/get_stream.h>
#include <cuda/__stream/stream_ref.h>
#include <cuda/std/__concepts/same_as.h>
#include <cuda/std/__functional/identity.h>
#include <cuda/std/__host_stdlib/stdexcept>
#include <cuda/std/__type_traits/is_callable.h>
#include <cuda/std/mdspan>
CUB_NAMESPACE_BEGIN
@@ -53,46 +47,6 @@ struct copy_mdspan_t
}
};
template <class _MDSpanIn, class _MDSpanOut>
[[nodiscard]] _CCCL_HOST_API ::cudaError_t
__copy_mdspan_bytes(::cuda::stream_ref __stream, _MDSpanIn&& __mdspan_in, _MDSpanOut&& __mdspan_out)
{
_CCCL_TRY
{
::cuda::copy_bytes(__stream, __mdspan_in, __mdspan_out);
}
_CCCL_CATCH (const ::cuda::cuda_error& __e)
{
return __e.status();
}
#if _CCCL_HOSTED()
_CCCL_CATCH (const ::std::invalid_argument& __e)
{
static_cast<void>(__e);
return ::cudaErrorInvalidValue;
}
#endif // _CCCL_HOSTED
_CCCL_CATCH_ALL
{
return ::cudaErrorUnknown;
}
return ::cudaSuccess;
}
template <class _MDSpanIn, class _MDSpanOut, class _Env>
[[nodiscard]] CUB_RUNTIME_FUNCTION ::cudaError_t
__transform_copy(_MDSpanIn&& __mdspan_in, _MDSpanOut&& __mdspan_out, const _Env& __env)
{
return CUB_NS_QUALIFIER::DeviceTransform::__transform_internal(
::cuda::std::make_tuple(__mdspan_in.data_handle()),
__mdspan_out.data_handle(),
__mdspan_in.size(),
::cuda::always_true{},
::cuda::std::identity{},
__env);
}
template <typename T_In,
typename E_In,
typename L_In,
@@ -110,38 +64,13 @@ copy(::cuda::std::mdspan<T_In, E_In, L_In, A_In> mdspan_in,
if (mdspan_in.is_exhaustive() && mdspan_out.is_exhaustive()
&& detail::have_same_strides(mdspan_in.mapping(), mdspan_out.mapping()))
{
// NOLINTBEGIN(bugprone-branch-clone)
if constexpr (::cuda::std::same_as<T_In, T_Out>
&& ::cuda::__detail::__can_mdspan_copy_bytes<T_In, E_In, L_In, T_Out, E_Out, L_Out>
&& ::cuda::std::__is_callable_v<::cuda::get_stream_t, const EnvT&>)
{
NV_IF_TARGET(
NV_IS_HOST,
({
auto __stream = ::cuda::get_stream(env);
// cuda::copy_bytes() builds an __ensure_current_context(stream_ref), which calls
// cuStreamGetCtx(). That driver call rejects the NULL stream with
// CUDA_ERROR_INVALID_VALUE. Use the transform path, which goes through the runtime
// API and accepts the NULL stream.
//
// Likewise, we cannot retrieve the context for a stream that is capturings so we
// need to call the kernel.
if (__stream.get() == nullptr
|| (::cuda::__driver::__streamIsCapturing(__stream.get()) == ::CU_STREAM_CAPTURE_STATUS_ACTIVE))
{
return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env);
}
return CUB_NS_QUALIFIER::detail::copy_mdspan::__copy_mdspan_bytes(__stream, mdspan_in, mdspan_out);
}),
(return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env);))
}
else
{
return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env);
}
// NOLINTEND(bugprone-branch-clone)
return cub::DeviceTransform::__transform_internal(
::cuda::std::make_tuple(mdspan_in.data_handle()),
mdspan_out.data_handle(),
mdspan_in.size(),
::cuda::always_true{},
::cuda::std::identity{},
env);
}
// TODO (fbusato): add ForEachInLayout when mdspan_in and mdspan_out have compatible layouts
// Compatible layouts could use more efficient iteration patterns

12
cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh
View File

@@ -216,12 +216,12 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch(
if constexpr (IsDeviceInit)
{
return kernel_source.template HistogramSweepKernelDeviceInit<
PolicySelector,
PRIVATIZED_SMEM_BINS,
FirstLevelArrayT,
SecondLevelArrayT,
IsEven,
IsByteSample>();
PolicySelector,
PRIVATIZED_SMEM_BINS,
FirstLevelArrayT,
SecondLevelArrayT,
IsEven,
IsByteSample>();
}
else
{

32
cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh
View File

@@ -208,14 +208,14 @@ __launch_bounds__(int(current_policy<PolicySelector>().lookback.threads_per_bloc
{
static constexpr ReduceByKeyPolicy policy = current_policy<PolicySelector>();
using AgentReduceByKeyPolicyT = agent_reduce_by_key_policy<
policy.lookback.threads_per_block,
policy.lookback.items_per_thread,
policy.lookback.load_algorithm,
policy.lookback.load_modifier,
policy.lookback.scan_algorithm,
delay_constructor_t<policy.lookback.lookback_delay.kind,
policy.lookback.lookback_delay.delay,
policy.lookback.lookback_delay.l2_write_latency>>;
policy.lookback.threads_per_block,
policy.lookback.items_per_thread,
policy.lookback.load_algorithm,
policy.lookback.load_modifier,
policy.lookback.scan_algorithm,
delay_constructor_t<policy.lookback.lookback_delay.kind,
policy.lookback.lookback_delay.delay,
policy.lookback.lookback_delay.l2_write_latency>>;
using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<
AgentReduceByKeyPolicyT,
@@ -656,14 +656,14 @@ _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_get
// TODO(bgruber): refactor this in the future
constexpr ReduceByKeyPolicy policy = policy_getter();
using Policy = agent_reduce_by_key_policy<
policy.lookback.threads_per_block,
policy.lookback.items_per_thread,
policy.lookback.load_algorithm,
policy.lookback.load_modifier,
policy.lookback.scan_algorithm,
delay_constructor_t<policy.lookback.lookback_delay.kind,
policy.lookback.lookback_delay.delay,
policy.lookback.lookback_delay.l2_write_latency>>;
policy.lookback.threads_per_block,
policy.lookback.items_per_thread,
policy.lookback.load_algorithm,
policy.lookback.load_modifier,
policy.lookback.scan_algorithm,
delay_constructor_t<policy.lookback.lookback_delay.kind,
policy.lookback.lookback_delay.delay,
policy.lookback.lookback_delay.l2_write_latency>>;
using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<Policy, AgentReduceByKey, Args...>;
return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS,
vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD,

138
cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh
View File

@@ -220,11 +220,11 @@ template <
typename ScanOpT,
typename InitValueT,
typename OffsetT,
typename AccumT = ::cuda::std::__accumulator_t<ScanOpT,
cub::detail::it_value_t<InputIteratorT>,
::cuda::std::_If<::cuda::std::is_same_v<InitValueT, NullType>,
cub::detail::it_value_t<InputIteratorT>,
typename InitValueT::value_type>>,
typename AccumT = ::cuda::std::__accumulator_t<ScanOpT,
cub::detail::it_value_t<InputIteratorT>,
::cuda::std::_If<::cuda::std::is_same_v<InitValueT, NullType>,
cub::detail::it_value_t<InputIteratorT>,
typename InitValueT::value_type>>,
ForceInclusive EnforceInclusive = ForceInclusive::No,
typename PolicyHub = detail::scan::
policy_hub<detail::it_value_t<InputIteratorT>, detail::it_value_t<OutputIteratorT>, AccumT, OffsetT, ScanOpT>,
@@ -552,44 +552,38 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") DispatchScan
int smem_size = smem_size_1_stage;
// When launched from the host, maximize the number of stages that we can fit inside the shared memory.
NV_IF_TARGET(
NV_IS_HOST, ({
// number of stages to have an even workload across all SMs (improves small problem sizes), assuming
// 1 CTA per SM +1 since it tends to improve performance
// TODO(bgruber): make the +1 a tuning parameter
const int max_stages_for_even_workload = static_cast<int>(
::cuda::ceil_div(num_items, static_cast<OffsetT>(sm_count * lookahead_policy.tile_size())) + 1);
NV_IF_TARGET(NV_IS_HOST, ({
// number of stages to have an even workload across all SMs (improves small problem sizes), assuming
// 1 CTA per SM +1 since it tends to improve performance
// TODO(bgruber): make the +1 a tuning parameter
const int max_stages_for_even_workload = static_cast<int>(
::cuda::ceil_div(num_items, static_cast<OffsetT>(sm_count * lookahead_policy.tile_size())) + 1);
while (num_stages <= max_stages_for_even_workload)
{
const int next_smem_size = detail::scan::smem_for_stages(
lookahead_policy,
num_stages + 1,
static_cast<int>(kernel_source.InputSize()),
static_cast<int>(kernel_source.InputAlign()),
static_cast<int>(kernel_source.OutputAlign()),
static_cast<int>(kernel_source.AccumSize()),
static_cast<int>(kernel_source.AccumAlign()));
if (next_smem_size > max_dynamic_smem_size)
{
// This number of stages failed, so stay at the current settings
break;
}
while (num_stages <= max_stages_for_even_workload)
{
const int next_smem_size = detail::scan::smem_for_stages(
lookahead_policy,
num_stages + 1,
static_cast<int>(kernel_source.InputSize()),
static_cast<int>(kernel_source.InputAlign()),
static_cast<int>(kernel_source.OutputAlign()),
static_cast<int>(kernel_source.AccumSize()),
static_cast<int>(kernel_source.AccumAlign()));
if (next_smem_size > max_dynamic_smem_size)
{
// This number of stages failed, so stay at the current settings
break;
}
smem_size = next_smem_size;
++num_stages;
}
smem_size = next_smem_size;
++num_stages;
}
// Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents
// multiple kernels from overwriting the max shared memory limit by different values.
//
// TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory
// sizes, however we need something that works even with older CTKs.
if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size))
{
return error;
}
}))
if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size))
{
return error;
}
}))
// Invoke init kernel
{
@@ -1169,44 +1163,38 @@ CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookahead(
int smem_size = smem_size_1_stage;
// When launched from the host, maximize the number of stages that we can fit inside the shared memory.
NV_IF_TARGET(
NV_IS_HOST, ({
// number of stages to have an even workload across all SMs (improves small problem sizes), assuming
// 1 CTA per SM +1 since it tends to improve performance
// TODO(bgruber): make the +1 a tuning parameter
const int max_stages_for_even_workload = static_cast<int>(
::cuda::ceil_div(num_items, static_cast<OffsetT>(sm_count * lookahead_policy.tile_size())) + 1);
NV_IF_TARGET(NV_IS_HOST, ({
// number of stages to have an even workload across all SMs (improves small problem sizes), assuming
// 1 CTA per SM +1 since it tends to improve performance
// TODO(bgruber): make the +1 a tuning parameter
const int max_stages_for_even_workload = static_cast<int>(
::cuda::ceil_div(num_items, static_cast<OffsetT>(sm_count * lookahead_policy.tile_size())) + 1);
while (num_stages <= max_stages_for_even_workload)
{
const int next_smem_size = detail::scan::smem_for_stages(
lookahead_policy,
num_stages + 1,
static_cast<int>(kernel_source.InputSize()),
static_cast<int>(kernel_source.InputAlign()),
static_cast<int>(kernel_source.OutputAlign()),
static_cast<int>(kernel_source.AccumSize()),
static_cast<int>(kernel_source.AccumAlign()));
if (next_smem_size > max_dynamic_smem_size)
{
// This number of stages failed, so stay at the current settings
break;
}
while (num_stages <= max_stages_for_even_workload)
{
const int next_smem_size = detail::scan::smem_for_stages(
lookahead_policy,
num_stages + 1,
static_cast<int>(kernel_source.InputSize()),
static_cast<int>(kernel_source.InputAlign()),
static_cast<int>(kernel_source.OutputAlign()),
static_cast<int>(kernel_source.AccumSize()),
static_cast<int>(kernel_source.AccumAlign()));
if (next_smem_size > max_dynamic_smem_size)
{
// This number of stages failed, so stay at the current settings
break;
}
smem_size = next_smem_size;
++num_stages;
}
smem_size = next_smem_size;
++num_stages;
}
// Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents
// multiple kernels from overwriting the max shared memory limit by different values.
//
// TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory
// sizes, however we need something that works even with older CTKs.
if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size))
{
return error;
}
}))
if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size))
{
return error;
}
}))
// Invoke init kernel
{

54
cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh
View File

@@ -238,15 +238,15 @@ template <
typename PolicyHub = policy_hub<KeysInputIteratorT, AccumT, cub::detail::it_value_t<ValuesInputIteratorT>, ScanOpT>,
typename PolicySelector = policy_selector_from_hub<PolicyHub>,
typename KernelSource = DeviceScanByKeyKernelSource<
PolicySelector,
KeysInputIteratorT,
ValuesInputIteratorT,
ValuesOutputIteratorT,
EqualityOp,
ScanOpT,
InitValueT,
OffsetT,
AccumT>,
PolicySelector,
KeysInputIteratorT,
ValuesInputIteratorT,
ValuesOutputIteratorT,
EqualityOp,
ScanOpT,
InitValueT,
OffsetT,
AccumT>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
struct dispatch_scan_by_key
{
@@ -672,15 +672,15 @@ template <
cub::detail::it_value_t<ValuesInputIteratorT>,
ScanOpT>,
typename KernelSource = DeviceScanByKeyKernelSource<
PolicySelector,
KeysInputIteratorT,
ValuesInputIteratorT,
ValuesOutputIteratorT,
EqualityOp,
ScanOpT,
InitValueT,
OffsetT,
AccumT>,
PolicySelector,
KeysInputIteratorT,
ValuesInputIteratorT,
ValuesOutputIteratorT,
EqualityOp,
ScanOpT,
InitValueT,
OffsetT,
AccumT>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
#if _CCCL_HAS_CONCEPTS()
requires scan_by_key_policy_selector<PolicySelector>
@@ -875,15 +875,15 @@ template <
detail::scan_by_key::policy_hub<KeysInputIteratorT, AccumT, cub::detail::it_value_t<ValuesInputIteratorT>, ScanOpT>,
typename PolicySelector = detail::scan_by_key::policy_selector_from_hub<PolicyHub>,
typename KernelSource = detail::scan_by_key::DeviceScanByKeyKernelSource<
PolicySelector,
KeysInputIteratorT,
ValuesInputIteratorT,
ValuesOutputIteratorT,
EqualityOp,
ScanOpT,
InitValueT,
OffsetT,
AccumT>,
PolicySelector,
KeysInputIteratorT,
ValuesInputIteratorT,
ValuesOutputIteratorT,
EqualityOp,
ScanOpT,
InitValueT,
OffsetT,
AccumT>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
using DispatchScanByKey
CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") = detail::scan_by_key::dispatch_scan_by_key<

2
cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh
View File

@@ -911,7 +911,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch(
{
using default_policy_selector_t = policy_selector_from_types<KeyT, ValueT, SegmentSizeT>;
using policy_selector_t = ::cuda::std::decay_t<
::cuda::std::execution::__query_result_or_t<TuningEnvT, SegmentedRadixSortPolicy, default_policy_selector_t>>;
::cuda::std::execution::__query_result_or_t<TuningEnvT, SegmentedRadixSortPolicy, default_policy_selector_t>>;
#if _CCCL_HAS_CONCEPTS()
static_assert(segmented_radix_sort_policy_selector<policy_selector_t>);
#endif // _CCCL_HAS_CONCEPTS()

32
cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh
View File

@@ -503,15 +503,15 @@ template <
decltype(select_segmented_accum_t<InputIteratorT, InitValueT, ReductionOpT>(static_cast<OverrideAccumT*>(nullptr))),
typename PolicySelector = policy_selector_from_types<AccumT, OffsetT, ReductionOpT>,
typename KernelSource = DeviceSegmentedReduceKernelSource<
PolicySelector,
InputIteratorT,
OutputIteratorT,
BeginOffsetIteratorT,
EndOffsetIteratorT,
OffsetT,
ReductionOpT,
InitValueT,
AccumT>,
PolicySelector,
InputIteratorT,
OutputIteratorT,
BeginOffsetIteratorT,
EndOffsetIteratorT,
OffsetT,
ReductionOpT,
InitValueT,
AccumT>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
#if _CCCL_HAS_CONCEPTS()
requires segmented_reduce_policy_selector<PolicySelector>
@@ -711,13 +711,13 @@ template <typename OverrideAccumT = use_default,
static_cast<OverrideAccumT*>(nullptr))),
typename PolicySelector = policy_selector_from_types<AccumT, OffsetT, ReductionOpT>,
typename KernelSource = DeviceFixedSizeSegmentedReduceKernelSource<
PolicySelector,
InputIteratorT,
OutputIteratorT,
OffsetT,
ReductionOpT,
InitValueT,
AccumT>,
PolicySelector,
InputIteratorT,
OutputIteratorT,
OffsetT,
ReductionOpT,
InitValueT,
AccumT>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY,
::cuda::std::enable_if_t<::cuda::std::is_arithmetic_v<OffsetT>, int> = 0>
#if _CCCL_HAS_CONCEPTS()

22
cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh
View File

@@ -96,17 +96,17 @@ template <
common_iterator_value_t<BeginOffsetIteratorInputT, EndOffsetIteratorInputT, BeginOffsetIteratorOutputT>,
typename PolicySelector = policy_selector_from_types<AccumT>,
typename KernelSource = device_segmented_scan_kernel_source<
PolicySelector,
InputIteratorT,
OutputIteratorT,
BeginOffsetIteratorInputT,
EndOffsetIteratorInputT,
BeginOffsetIteratorOutputT,
OffsetT,
ScanOpT,
InitValueT,
AccumT,
EnforceInclusive>,
PolicySelector,
InputIteratorT,
OutputIteratorT,
BeginOffsetIteratorInputT,
EndOffsetIteratorInputT,
BeginOffsetIteratorOutputT,
OffsetT,
ScanOpT,
InitValueT,
AccumT,
EnforceInclusive>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
#if _CCCL_HAS_CONCEPTS()
requires segmented_scan_policy_selector<PolicySelector>

17
cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh
View File

@@ -216,15 +216,14 @@ struct make_vsmem_helper
{
static constexpr SelectPolicy active_policy = DefaultPolicyGetter{}();
using agent_policy_t = detail::agent_select_if_policy<
active_policy.lookback.threads_per_block,
active_policy.lookback.items_per_thread,
active_policy.lookback.load_algorithm,
active_policy.lookback.load_modifier,
active_policy.lookback.scan_algorithm,
delay_constructor_t<active_policy.lookback.lookback_delay.kind,
active_policy.lookback.lookback_delay.delay,
active_policy.lookback.lookback_delay.l2_write_latency>,
active_policy.lookback._load_prefetch>;
active_policy.lookback.threads_per_block,
active_policy.lookback.items_per_thread,
active_policy.lookback.load_algorithm,
active_policy.lookback.load_modifier,
active_policy.lookback.scan_algorithm,
delay_constructor_t<active_policy.lookback.lookback_delay.kind,
active_policy.lookback.lookback_delay.delay,
active_policy.lookback.lookback_delay.l2_write_latency>>;
using type = vsmem_helper_default_fallback_policy_t<
agent_policy_t,
bind_selection_opt<SelectionOpt>::template agent_t,

24
cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh
View File

@@ -444,18 +444,18 @@ template <typename InputIteratorT,
typename OffsetT,
typename PolicySelector = policy_selector_from_types<it_value_t<InputIteratorT>, per_partition_offset_t>,
typename KernelSource = DeviceThreeWayPartitionKernelSource<
PolicySelector,
InputIteratorT,
FirstOutputIteratorT,
SecondOutputIteratorT,
UnselectedOutputIteratorT,
NumSelectedIteratorT,
ScanTileStateT,
SelectFirstPartOp,
SelectSecondPartOp,
per_partition_offset_t,
streaming_context_t<OffsetT>,
OffsetT>,
PolicySelector,
InputIteratorT,
FirstOutputIteratorT,
SecondOutputIteratorT,
UnselectedOutputIteratorT,
NumSelectedIteratorT,
ScanTileStateT,
SelectFirstPartOp,
SelectSecondPartOp,
per_partition_offset_t,
streaming_context_t<OffsetT>,
OffsetT>,
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
#if _CCCL_HAS_CONCEPTS()
requires three_way_partition_policy_selector<PolicySelector>

2
cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh
View File

@@ -26,7 +26,7 @@
# pragma system_header
#endif // no system header
#define _CCCL_CUB_HAS_TILE_TRANSFORM() _CCCL_TILE_COMPILATION() && _CCCL_STD_VER >= 2020
#define _CCCL_CUB_HAS_TILE_TRANSFORM() _CCCL_TILE_COMPILATION()
#if _CCCL_CUB_HAS_TILE_TRANSFORM() && defined(_CCCL_ENABLE_EXPERIMENTAL_TILE_TRANSFORM_DISPATCH)
# define _CCCL_CUB_TILE_TRANSFORM_DISPATCH_ENABLED() 1

5
cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh
View File

@@ -13,14 +13,13 @@
# pragma system_header
#endif // no system header
#include <cub/util_device.cuh>
#include <cub/util_type.cuh>
#include <thrust/type_traits/is_contiguous_iterator.h>
#include <thrust/type_traits/is_trivially_relocatable.h>
#include <cuda/__functional/maximum.h>
#include <cuda/__functional/minimum.h>
#include <cuda/__type_traits/is_trivially_copyable.h>
#include <cuda/std/__concepts/same_as.h>
#include <cuda/std/__functional/operations.h>
#include <cuda/std/__fwd/format.h>
@@ -131,7 +130,7 @@ template <typename It>
return iterator_info{
static_cast<int>(size_of<vt>),
static_cast<int>(align_of<vt>),
::cuda::is_trivially_copyable_v<vt>,
THRUST_NS_QUALIFIER::is_trivially_relocatable_v<vt>,
THRUST_NS_QUALIFIER::is_contiguous_iterator_v<It>};
}

4
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh
View File

@@ -158,13 +158,13 @@ struct BatchedCopyPolicy
BatchedCopyLookbackPolicy lookback; //!< The policy for the batched-copy algorithm based on decoupled-lookback. Only
//!< used when @p algorithm is @p lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const BatchedCopyPolicy& lhs, const BatchedCopyPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const BatchedCopyPolicy& lhs, const BatchedCopyPolicy& rhs) noexcept
{
return !(lhs == rhs);

311
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh
View File

@@ -16,27 +16,20 @@
#include <cub/block/block_load.cuh>
#include <cub/block/block_scan.cuh>
#include <cub/block/block_store.cuh>
#include <cub/util_device.cuh>
#include <cuda/__cmath/pow2.h>
#include <cuda/__device/compute_capability.h>
#include <cuda/__execution/determinism.h>
#include <cuda/__execution/tie_break.h>
#include <cuda/std/__host_stdlib/ostream>
#include <cuda/std/array>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail::batched_topk
{
//! Sub-policy for the compaction epilogue shared by the baseline @ref DeviceBatchedTopK workers: it scans the radix
//! histogram and writes out the selected keys.
struct epilogue_policy
{
int items_per_thread; //!< Keys each thread loads/stores per tile in the epilogue.
BlockLoadAlgorithm load_algorithm; //!< Block load algorithm used to read keys back in the epilogue.
BlockStoreAlgorithm store_algorithm; //!< Block store algorithm used to write the selected keys.
BlockScanAlgorithm scan_algorithm; //!< Block scan algorithm used for the histogram prefix sum.
int items_per_thread;
BlockLoadAlgorithm load_algorithm;
BlockStoreAlgorithm store_algorithm;
BlockScanAlgorithm scan_algorithm;
_CCCL_HOST_DEVICE_API friend constexpr bool operator==(const epilogue_policy& lhs, const epilogue_policy& rhs)
{
@@ -49,26 +42,24 @@ struct epilogue_policy
return !(lhs == rhs);
}
#if _CCCL_HOSTED()
#if !_CCCL_COMPILER(NVRTC)
friend ::std::ostream& operator<<(::std::ostream& os, const epilogue_policy& p)
{
return os
<< "epilogue_policy { .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm
<< ", .store_algorithm = " << p.store_algorithm << ", .scan_algorithm = " << p.scan_algorithm << " }";
}
#endif // _CCCL_HOSTED()
#endif // !_CCCL_COMPILER(NVRTC)
};
//! Per-segment worker sub-policy for the baseline backend: one thread block cooperatively computes the top-k of a
//! single segment. @ref baseline_topk_policy holds several of these, ordered by decreasing tile size.
struct worker_policy
{
int threads_per_block; //!< Number of threads in a CUDA block.
int items_per_thread; //!< Keys each thread loads/processes per tile (with `threads_per_block` sets the tile size).
BlockLoadAlgorithm load_algorithm; //!< Block load algorithm used to read the segment's keys.
BlockStoreAlgorithm store_algorithm; //!< Block store algorithm used to write the selected keys.
int threads_per_block;
int items_per_thread;
BlockLoadAlgorithm load_algorithm;
BlockStoreAlgorithm store_algorithm;
epilogue_policy epilogue; //!< Sub-policy for the compaction epilogue.
epilogue_policy epilogue;
_CCCL_HOST_DEVICE_API friend constexpr bool operator==(const worker_policy& lhs, const worker_policy& rhs)
{
@@ -89,15 +80,13 @@ struct worker_policy
<< ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm
<< ", .store_algorithm = " << p.store_algorithm << ", .epilogue = " << p.epilogue << " }";
}
#endif // _CCCL_HOSTED()
#endif // !_CCCL_COMPILER(NVRTC)
};
//! Sub-policy for the baseline backend's multiple-blocks-per-segment worker path, used for segments too large for a
//! single worker block.
struct multi_worker_policy
{
int threads_per_block; //!< Number of threads in a CUDA block.
int items_per_thread; //!< Keys each thread loads/processes per tile.
int threads_per_block;
int items_per_thread;
_CCCL_HOST_DEVICE_API friend constexpr bool operator==(const multi_worker_policy& lhs, const multi_worker_policy& rhs)
{
@@ -109,7 +98,7 @@ struct multi_worker_policy
return !(lhs == rhs);
}
#if _CCCL_HOSTED()
#if !_CCCL_COMPILER(NVRTC)
friend ::std::ostream& operator<<(::std::ostream& os, const multi_worker_policy& p)
{
return os << "multi_worker_policy { .threads_per_block = " << p.threads_per_block
@@ -118,31 +107,28 @@ struct multi_worker_policy
#endif // _CCCL_HOSTED()
};
//! Sub-policy for the baseline (worker-per-segment) backend of @ref DeviceBatchedTopK.
struct baseline_topk_policy
struct batched_topk_policy
{
//! Per-segment worker policies ordered by decreasing tile size. At compile time the smallest policy whose tile size
//! still covers the upper bound of the segment size is selected.
// The list of per-segment agent policies is ordered by decreasing tile size. At compile time, the smallest policy
// whose tile size still covers the upper bound of the segment size is selected.
::cuda::std::array<worker_policy, 6> worker_per_segment_policies;
multi_worker_policy multi_worker_per_segment_policy; //!< Worker policy for segments too large for a single block.
multi_worker_policy multi_worker_per_segment_policy;
_CCCL_HOST_DEVICE_API friend constexpr bool
operator==(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs)
_CCCL_HOST_DEVICE_API friend constexpr bool operator==(const batched_topk_policy& lhs, const batched_topk_policy& rhs)
{
return lhs.worker_per_segment_policies == rhs.worker_per_segment_policies
&& lhs.multi_worker_per_segment_policy == rhs.multi_worker_per_segment_policy;
}
_CCCL_HOST_DEVICE_API friend constexpr bool
operator!=(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs)
_CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const batched_topk_policy& lhs, const batched_topk_policy& rhs)
{
return !(lhs == rhs);
}
#if _CCCL_HOSTED()
friend ::std::ostream& operator<<(::std::ostream& os, const baseline_topk_policy& p)
friend ::std::ostream& operator<<(::std::ostream& os, const batched_topk_policy& p)
{
os << "baseline_topk_policy { .worker_per_segment_policies = { ";
os << "batched_topk_policy { .worker_per_segment_policies = { ";
for (::cuda::std::size_t i = 0; i < p.worker_per_segment_policies.size(); ++i)
{
if (i != 0)
@@ -156,242 +142,45 @@ struct baseline_topk_policy
#endif // _CCCL_HOSTED()
};
// Default baseline sub-policy. Tuning is currently CC-independent.
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_baseline_policy() -> baseline_topk_policy
#if _CCCL_HAS_CONCEPTS()
template <typename T>
concept batched_topk_policy_selector = policy_selector<T, batched_topk_policy>;
#endif // _CCCL_HAS_CONCEPTS()
struct policy_selector
{
constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE;
constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE;
constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS;
constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg};
return baseline_topk_policy{
{{
worker_policy{256, 64, load_alg, store_alg, epilogue},
worker_policy{256, 32, load_alg, store_alg, epilogue},
worker_policy{256, 16, load_alg, store_alg, epilogue},
worker_policy{256, 8, load_alg, store_alg, epilogue},
worker_policy{256, 4, load_alg, store_alg, epilogue},
worker_policy{128, 2, load_alg, store_alg, epilogue},
}},
multi_worker_policy{256, 64}};
}
// Largest maximum segment size (in keys) the baseline (worker-per-segment) backend can cover: the largest worker tile
// (threads_per_block * items_per_thread) in `policy`. A larger statically-known maximum segment size makes the baseline
// backend ineligible (the selector then picks the cluster backend where supported, otherwise `unsupported`). This is
// only the tile-based necessary condition; the exact predicate `baseline_can_cover_v` also checks the agent's
// shared-memory fit (which needs the concrete agent types).
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::int64_t
baseline_max_covered_segment_size(const baseline_topk_policy& policy)
{
::cuda::std::int64_t max_tile_size = 0;
for (const auto& worker : policy.worker_per_segment_policies)
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> batched_topk_policy
{
const ::cuda::std::int64_t tile_size = ::cuda::std::int64_t{worker.threads_per_block} * worker.items_per_thread;
if (tile_size > max_tile_size)
{
max_tile_size = tile_size;
}
constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE;
constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE;
constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS;
constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg};
return batched_topk_policy{
{{
worker_policy{256, 64, load_alg, store_alg, epilogue},
worker_policy{256, 32, load_alg, store_alg, epilogue},
worker_policy{256, 16, load_alg, store_alg, epilogue},
worker_policy{256, 8, load_alg, store_alg, epilogue},
worker_policy{256, 4, load_alg, store_alg, epilogue},
worker_policy{128, 2, load_alg, store_alg, epilogue},
}},
multi_worker_policy{256, 64}};
}
return max_tile_size;
}
//! Execution shape for the thread-block-cluster backend of @ref DeviceBatchedTopK. The dispatch picks the number of
//! cluster blocks and the dynamic shared-memory block_tile capacity at runtime (occupancy / wave-aware), so this policy
//! mostly carries per-block tuning knobs; the two trailing `max_*` fields are optional launch-geometry caps that bound
//! that runtime choice.
struct cluster_topk_policy
{
int threads_per_block; //!< Number of threads in a CUDA block.
int min_blocks_per_sm; //!< Minimum resident blocks per SM, forwarded as the kernel launch-bounds occupancy hint.
int min_chunks_per_block; //!< Minimum number of chunks a block must own to join a segment's effective cluster (the
//!< divisor mapping a segment's chunk count to its cluster width). Must be >= 1.
int chunk_bytes; //!< Size in bytes of one block_tile chunk -- the granularity of the async-copy load pipeline.
int load_align_bytes; //!< Load / bulk-copy alignment in bytes. Must be a power of two and >= 16
//!< (`detail::bulk_copy_min_align`), and `chunk_bytes` must be a multiple of it (see
//!< `is_valid_cluster_policy`).
int pipeline_stages; //!< Depth of the async-copy (mbarrier) pipeline that stages chunks into shared memory.
int single_block_max_seg_size; //!< Largest segment size, in keys, still eligible for the single-block fast path
//!< (kept out of the byte-unit loading group above as it is measured in items).
int bits_per_pass; //!< Radix digit width per pass; each pass' histogram spans `1 << bits_per_pass` buckets. Together
//!< with `threads_per_block` this implicitly fixes the histogram block-scan's items per thread,
//!< `ceil_div(1 << bits_per_pass, threads_per_block)` buckets scanned per thread.
int histogram_items_per_thread; //!< Keys each thread accumulates per tile during the radix histogram passes.
int tie_break_items_per_thread; //!< Keys each thread processes per tile during the final tie-break / filter phase.
int copy_items_per_thread; //!< Keys each thread copies per tile on the select-all (k >= segment size) fast path.
// Launch-geometry caps that bound the otherwise heuristic / hardware-derived cluster width and resident shared-memory
// footprint. Both default to 0 (= unrestricted) and are deliberately not auto-tuned. Beyond deterministically
// steering tests onto the streaming / cluster paths at a small footprint, they let a caller trade top-k throughput
// for resources it wants to leave free -- e.g. capping resident slots to fit a shared-memory carveout reserved for a
// concurrently running kernel, or narrowing the cluster width to co-schedule other work. The algorithm stays correct
// at any cap: a segment that no longer fits resident simply streams the remainder from global memory.
int max_blocks_per_cluster; //!< Upper bound on the launched cluster width (CTAs per segment); 0 = unrestricted (the
//!< hardware cluster-width ceiling, queried from the runtime). Non-zero is additionally
//!< clamped to that same ceiling. A cap narrower than a segment needs pushes it into the
//!< streaming fallback (cap 1 -> single-CTA streaming).
int max_chunk_slots_per_block; //!< Upper bound on resident chunk slots per block; 0 = unrestricted (the full
//!< shared-memory budget: the hardware opt-in budget). A smaller cap shrinks each
//!< CTA's resident capacity (and thus its dynamic shared-memory request), so a smaller
//!< segment overflows into the streaming path.
// Equality/streaming make this a regular type (required by the `policy_selector` concept / `dispatch_compute_cap`).
_CCCL_HOST_DEVICE_API friend constexpr bool operator==(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs)
{
return lhs.threads_per_block == rhs.threads_per_block && lhs.min_blocks_per_sm == rhs.min_blocks_per_sm
&& lhs.min_chunks_per_block == rhs.min_chunks_per_block && lhs.chunk_bytes == rhs.chunk_bytes
&& lhs.load_align_bytes == rhs.load_align_bytes && lhs.pipeline_stages == rhs.pipeline_stages
&& lhs.single_block_max_seg_size == rhs.single_block_max_seg_size && lhs.bits_per_pass == rhs.bits_per_pass
&& lhs.histogram_items_per_thread == rhs.histogram_items_per_thread
&& lhs.tie_break_items_per_thread == rhs.tie_break_items_per_thread
&& lhs.copy_items_per_thread == rhs.copy_items_per_thread
&& lhs.max_blocks_per_cluster == rhs.max_blocks_per_cluster
&& lhs.max_chunk_slots_per_block == rhs.max_chunk_slots_per_block;
}
_CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs)
{
return !(lhs == rhs);
}
#if _CCCL_HOSTED()
friend ::std::ostream& operator<<(::std::ostream& os, const cluster_topk_policy& p)
{
return os
<< "cluster_topk_policy { .threads_per_block = " << p.threads_per_block
<< ", .min_blocks_per_sm = " << p.min_blocks_per_sm << ", .min_chunks_per_block = " << p.min_chunks_per_block
<< ", .chunk_bytes = " << p.chunk_bytes << ", .load_align_bytes = " << p.load_align_bytes
<< ", .pipeline_stages = " << p.pipeline_stages
<< ", .single_block_max_seg_size = " << p.single_block_max_seg_size << ", .bits_per_pass = " << p.bits_per_pass
<< ", .histogram_items_per_thread = " << p.histogram_items_per_thread << ", .tie_break_items_per_thread = "
<< p.tie_break_items_per_thread << ", .copy_items_per_thread = " << p.copy_items_per_thread
<< ", .max_blocks_per_cluster = " << p.max_blocks_per_cluster
<< ", .max_chunk_slots_per_block = " << p.max_chunk_slots_per_block << " }";
}
#endif // _CCCL_HOSTED()
};
// Default cluster sub-policy. Tuning is currently CC-independent.
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_cluster_policy() -> cluster_topk_policy
template <typename KeyT, typename ValueT, typename SegmentSizeT, ::cuda::std::int64_t MaxK>
struct policy_selector_from_types
{
return cluster_topk_policy{
/*threads_per_block=*/512,
/*min_blocks_per_sm=*/1,
/*min_chunks_per_block=*/1,
/*chunk_bytes=*/16 * 1024,
/*load_align_bytes=*/128,
/*pipeline_stages=*/8,
/*single_block_max_seg_size=*/8 * 1024,
/*bits_per_pass=*/11,
/*histogram_items_per_thread=*/8,
/*tie_break_items_per_thread=*/8,
/*copy_items_per_thread=*/8,
/*max_blocks_per_cluster=*/0,
/*max_chunk_slots_per_block=*/0};
}
// Hard constraints a cluster sub-policy must satisfy, mirroring the agent's compile-time invariants so a bad policy
// (e.g. from a `tune` override) trips `launch_cluster_arm`'s `static_assert(is_valid_cluster_policy(policy))` with a
// clear message instead of a cryptic failure deep in the agent (or a host-side divide-by-zero in the launch-shape
// math). The block_tile byte geometry constraints stem from the aligned bulk-copy (TMA) load path, which addresses
// gmem/smem in `load_align_bytes`-sized, aligned units.
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_cluster_policy(cluster_topk_policy policy) -> bool
{
return policy.chunk_bytes > 0 && policy.load_align_bytes >= bulk_copy_min_align
&& ::cuda::is_power_of_two(policy.load_align_bytes) && policy.chunk_bytes % policy.load_align_bytes == 0
&& policy.threads_per_block > 0 && policy.threads_per_block % warp_threads == 0 && policy.min_blocks_per_sm >= 0
&& policy.pipeline_stages >= 1 && policy.pipeline_stages <= 32 && policy.min_chunks_per_block >= 1
&& policy.bits_per_pass >= 1 && policy.bits_per_pass <= 16 && policy.histogram_items_per_thread > 0
&& policy.tie_break_items_per_thread > 0 && policy.copy_items_per_thread > 0
&& policy.single_block_max_seg_size >= 0 && policy.max_blocks_per_cluster >= 0
&& policy.max_chunk_slots_per_block >= 0;
}
static_assert(is_valid_cluster_policy(make_cluster_policy()));
// -----------------------------------------------------------------------------
// Backend selection
// -----------------------------------------------------------------------------
//! Backend algorithms for @ref DeviceBatchedTopK. Both backends are launched through a single kernel symbol; which one
//! runs is decided per architecture by `policy_selector` below, whose result also drives the device-side agent
//! selection (via `current_policy`).
enum class topk_algorithm
{
baseline, //!< worker-per-segment backend (single thread block per segment)
cluster, //!< thread-block-cluster backend (SM 9.0+)
unsupported //!< no backend can serve the request on the target architecture; dispatch returns cudaErrorNotSupported
};
#if _CCCL_HOSTED()
[[nodiscard]] inline ::std::ostream& operator<<(::std::ostream& os, topk_algorithm backend)
{
switch (backend)
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const
-> batched_topk_policy
{
case topk_algorithm::baseline:
return os << "baseline";
case topk_algorithm::cluster:
return os << "cluster";
default:
return os << "unsupported";
return policy_selector{}(cc);
}
}
#endif // _CCCL_HOSTED()
//! The tuning policy for all backends of @ref DeviceBatchedTopK. It carries the selected backend plus both backends'
//! sub-policies; the kernel instantiates only the arm named by @p backend (chosen device-side via `current_policy`).
//!
//! This is a regular type: `detail::dispatch_compute_cap` (and the `policy_selector` concept) require the selector's
//! result to be `::cuda::std::regular`, hence the equality/streaming operators below.
struct topk_policy
{
topk_algorithm backend; //!< Backend the dispatch selected, i.e. the kernel arm that runs.
baseline_topk_policy baseline; //!< Sub-policy used when @p backend is @p topk_algorithm::baseline.
cluster_topk_policy cluster; //!< Sub-policy used when @p backend is @p topk_algorithm::cluster.
_CCCL_HOST_DEVICE_API friend constexpr bool operator==(const topk_policy& lhs, const topk_policy& rhs)
{
return lhs.backend == rhs.backend && lhs.baseline == rhs.baseline && lhs.cluster == rhs.cluster;
}
_CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const topk_policy& lhs, const topk_policy& rhs)
{
return !(lhs == rhs);
}
#if _CCCL_HOSTED()
friend ::std::ostream& operator<<(::std::ostream& os, const topk_policy& p)
{
return os << "topk_policy { .backend = " << p.backend << ", .baseline = " << p.baseline
<< ", .cluster = " << p.cluster << " }";
}
#endif // _CCCL_HOSTED()
};
#if _CCCL_HAS_CONCEPTS()
template <typename T>
concept topk_policy_selector = policy_selector<T, topk_policy>;
static_assert(batched_topk_policy_selector<policy_selector>);
#endif // _CCCL_HAS_CONCEPTS()
// Crossover knobs (TODO: tune via SM100 benchmarks).
//! Clusters require SM 9.0+.
inline constexpr int cluster_min_cc_major = 9;
//! Smallest statically-known maximum segment size at which the cluster backend starts to win (measured on B200). This
//! is the backend crossover threshold and is intentionally part of the selector -- not a tunable policy field -- so
//! that tuning the cluster policy (e.g. its single-CTA threshold) does not silently shift which backend is chosen.
inline constexpr ::cuda::std::int64_t cluster_beneficial_min_segment_size = 8 * 1024;
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr bool cluster_capable([[maybe_unused]] ::cuda::compute_capability cc)
{
#if _CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH()
return cc >= ::cuda::compute_capability{cluster_min_cc_major, 0};
#else // ^^^ dynamic cluster launches enabled ^^^ / vvv dynamic cluster launches disabled vvv
// The cluster backend launches with a runtime cluster width, which _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH compiles out;
// reporting no architecture as cluster-capable makes the selector fall back to baseline (or report unsupported).
return false;
#endif // _CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH()
}
} // namespace detail::batched_topk
CUB_NAMESPACE_END

23
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh
View File

@@ -20,9 +20,9 @@
#include <cub/util_type.cuh>
#include <thrust/type_traits/is_contiguous_iterator.h>
#include <thrust/type_traits/is_trivially_relocatable.h>
#include <cuda/__device/compute_capability.h>
#include <cuda/__type_traits/is_trivially_copyable.h>
#include <cuda/std/__algorithm/clamp.h>
#include <cuda/std/__host_stdlib/ostream>
#include <cuda/std/concepts>
@@ -43,15 +43,10 @@ struct MergePolicy
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
operator==(const MergePolicy& lhs, const MergePolicy& rhs) noexcept
{
// gcc 8 folds comparisons of adjacent bool members within one expression into a BIT_FIELD_REF, which its
// constexpr evaluator cannot handle (ICE in cxx_eval_bit_field_ref, fixed in gcc 9). Keep each bool
// comparison in a separate statement to avoid the fold.
const bool same_bulk_copy_for_keys = lhs.use_bulk_copy_for_keys == rhs.use_bulk_copy_for_keys;
const bool same_bulk_copy_for_values = lhs.use_bulk_copy_for_values == rhs.use_bulk_copy_for_values;
const bool same_unroll = lhs.unroll == rhs.unroll;
return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread
&& lhs.load_modifier == rhs.load_modifier && lhs.store_algorithm == rhs.store_algorithm
&& same_bulk_copy_for_keys && same_bulk_copy_for_values && same_unroll;
&& lhs.use_bulk_copy_for_keys == rhs.use_bulk_copy_for_keys
&& lhs.use_bulk_copy_for_values == rhs.use_bulk_copy_for_values && lhs.unroll == rhs.unroll;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
@@ -95,10 +90,10 @@ struct policy_selector
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> MergePolicy
{
const int tune_type_size = key_size + value_size;
const int ipt_800_plus = nominal_4B_items_to_items(15, tune_type_size);
const bool can_bulk_keys = (key_size == key_align) && key_is_trivially_relocatable && key_iterators_are_contiguous
&& key_iterator_value_types_are_the_same;
const int tune_type_size = key_size + value_size;
const int ipt_800_plus = nominal_4B_items_to_items(15, tune_type_size);
const bool can_bulk_keys = (key_size == key_align) && key_is_trivially_relocatable && key_iterators_are_contiguous
&& key_iterator_value_types_are_the_same;
const bool can_bulk_values = (value_size == value_align) && value_is_trivially_relocatable
&& value_iterators_are_contiguous && value_iterator_value_types_are_the_same;
@@ -168,12 +163,12 @@ struct policy_selector_from_types
return policy_selector{
int{sizeof(key_t)},
int{alignof(key_t)},
::cuda::is_trivially_copyable_v<key_t>,
THRUST_NS_QUALIFIER::is_trivially_relocatable_v<key_t>,
THRUST_NS_QUALIFIER::is_contiguous_iterator_v<KeysIt1> && THRUST_NS_QUALIFIER::is_contiguous_iterator_v<KeysIt2>,
::cuda::std::is_same_v<key_t, it_value_t<KeysIt2>>,
::cuda::std::is_same_v<item_t, NullType> ? 0 : int{sizeof(item_t)},
int{alignof(item_t)},
::cuda::is_trivially_copyable_v<item_t>,
THRUST_NS_QUALIFIER::is_trivially_relocatable_v<item_t>,
THRUST_NS_QUALIFIER::is_contiguous_iterator_v<ItemsIt1>
&& THRUST_NS_QUALIFIER::is_contiguous_iterator_v<ItemsIt2>,
::cuda::std::is_same_v<item_t, it_value_t<ItemsIt2>>,

4
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh
View File

@@ -99,13 +99,13 @@ struct ReduceByKeyPolicy
ReduceByKeyLookbackPolicy lookback; //!< The policy for the reduce-by-key algorithm based on decoupled-lookback. Only
//!< used when @p algorithm is @lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept
{
return !(lhs == rhs);

4
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh
View File

@@ -100,13 +100,13 @@ struct RleEncodePolicy
RleAlgorithm algorithm = RleAlgorithm::lookback; //!< The RLE-encode algorithm to use
RleLookbackPolicy lookback; //!< The lookback policy
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const RleEncodePolicy& lhs, const RleEncodePolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const RleEncodePolicy& lhs, const RleEncodePolicy& rhs) noexcept
{
return !(lhs == rhs);

4
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh
View File

@@ -104,13 +104,13 @@ struct RleNonTrivialRunsPolicy
RleNonTrivialRunsLookbackPolicy lookback; //!< The policy for the non-trivial-runs algorithm based on
//!< decoupled-lookback. Only used when @p algorithm is @lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const RleNonTrivialRunsPolicy& lhs, const RleNonTrivialRunsPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const RleNonTrivialRunsPolicy& lhs, const RleNonTrivialRunsPolicy& rhs) noexcept
{
return !(lhs == rhs);

6
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh
View File

@@ -187,14 +187,12 @@ struct ScanPolicy
ScanLookbackPolicy lookback; //!< The look-back scan policy (used when algorithm is @p lookback, otherwise ignored)
ScanLookaheadPolicy lookahead; //!< The lookahead scan policy (used when algorithm is @p lookahead, otherwise ignored)
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
operator==(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept
[[nodiscard]] _CCCL_API friend constexpr bool operator==(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept
{
return lhs.lookback == rhs.lookback && lhs.lookahead == rhs.lookahead && lhs.algorithm == rhs.algorithm;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
operator!=(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept
[[nodiscard]] _CCCL_API friend constexpr bool operator!=(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept
{
return !(lhs == rhs);
}

4
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh
View File

@@ -103,13 +103,13 @@ struct ScanByKeyPolicy
ScanByKeyLookbackPolicy lookback; //!< The policy for the scan-by-key algorithm based on decoupled-lookback. Only used
//!< when @p algorithm is @lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const ScanByKeyPolicy& lhs, const ScanByKeyPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const ScanByKeyPolicy& lhs, const ScanByKeyPolicy& rhs) noexcept
{
return !(lhs == rhs);

28
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh
View File

@@ -575,13 +575,13 @@ struct policy_hub
static constexpr int BLOCK_THREADS = 256;
static constexpr int PARTITIONING_THRESHOLD = 500;
using LargeSegmentPolicy = detail::agent_radix_sort_downsweep_policy<
BLOCK_THREADS,
23,
DominantT,
BLOCK_LOAD_TRANSPOSE,
LOAD_DEFAULT,
RADIX_RANK_MEMOIZE,
BLOCK_SCAN_WARP_SCANS,
BLOCK_THREADS,
23,
DominantT,
BLOCK_LOAD_TRANSPOSE,
LOAD_DEFAULT,
RADIX_RANK_MEMOIZE,
BLOCK_SCAN_WARP_SCANS,
(sizeof(KeyT) > 1) ? 6 : 4>;
static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems<DominantT>(9);
@@ -606,13 +606,13 @@ struct policy_hub
static constexpr int BLOCK_THREADS = 256;
static constexpr int PARTITIONING_THRESHOLD = 500;
using LargeSegmentPolicy = detail::agent_radix_sort_downsweep_policy<
BLOCK_THREADS,
23,
DominantT,
BLOCK_LOAD_TRANSPOSE,
LOAD_DEFAULT,
RADIX_RANK_MEMOIZE,
BLOCK_SCAN_WARP_SCANS,
BLOCK_THREADS,
23,
DominantT,
BLOCK_LOAD_TRANSPOSE,
LOAD_DEFAULT,
RADIX_RANK_MEMOIZE,
BLOCK_SCAN_WARP_SCANS,
(sizeof(KeyT) > 1) ? 6 : 4>;
static constexpr bool LARGE_ITEMS = sizeof(DominantT) > 4;

22
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh
View File

@@ -18,7 +18,6 @@
#include <cub/block/block_load.cuh>
#include <cub/block/block_scan.cuh>
#include <cub/detail/delay_constructor.cuh>
#include <cub/detail/prefetch.cuh>
#include <cub/device/dispatch/tuning/common.cuh>
#include <cub/util_device.cuh>
#include <cub/util_math.cuh>
@@ -44,15 +43,13 @@ struct SelectLookbackPolicy
CacheLoadModifier load_modifier; //!< The @ref CacheLoadModifier used for loading items from global memory
BlockScanAlgorithm scan_algorithm; //!< The @ref BlockScanAlgorithm used for scanning
LookbackDelayPolicy lookback_delay; //!< The policy configuring the delay used in decoupled lookback
detail::LoadPrefetch _load_prefetch = detail::LoadPrefetch::none; //!< Implementation detail; do not use directly
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
operator==(const SelectLookbackPolicy& lhs, const SelectLookbackPolicy& rhs) noexcept
{
return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread
&& lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier
&& lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay
&& lhs._load_prefetch == rhs._load_prefetch;
&& lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
@@ -65,10 +62,9 @@ struct SelectLookbackPolicy
friend ::std::ostream& operator<<(::std::ostream& os, const SelectLookbackPolicy& p)
{
return os
<< "SelectLookbackPolicy { .threads_per_block = " << p.threads_per_block
<< ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm
<< ", .load_modifier = " << p.load_modifier << ", .scan_algorithm = " << p.scan_algorithm
<< ", .lookback_delay = " << p.lookback_delay << ", ._load_prefetch = " << p._load_prefetch << " }";
<< "SelectLookbackPolicy { .threads_per_block = " << p.threads_per_block << ", .items_per_thread = "
<< p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .load_modifier = " << p.load_modifier
<< ", .scan_algorithm = " << p.scan_algorithm << ", .lookback_delay = " << p.lookback_delay << " }";
}
#endif // _CCCL_HOSTED()
};
@@ -106,14 +102,12 @@ struct SelectPolicy
SelectLookbackPolicy lookback; //!< The policy for the selection algorithm based on decoupled-lookback. Only used when
//!< @p algorithm is @lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
operator==(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept
[[nodiscard]] _CCCL_API friend constexpr bool operator==(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
operator!=(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept
[[nodiscard]] _CCCL_API friend constexpr bool operator!=(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept
{
return !(lhs == rhs);
}
@@ -196,13 +190,13 @@ struct PartitionPolicy
PartitionLookbackPolicy lookback; //!< The policy for the partition algorithm based on decoupled-lookback. Only used
//!< when algorithm is @lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const PartitionPolicy& lhs, const PartitionPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const PartitionPolicy& lhs, const PartitionPolicy& rhs) noexcept
{
return !(lhs == rhs);

4
cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh
View File

@@ -98,13 +98,13 @@ struct ThreeWayPartitionPolicy
ThreeWayPartitionLookbackPolicy lookback; //!< The policy for the three-way partition algorithm based on
//!< decoupled-lookback. Only used when @p algorithm is @lookback.
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator==(const ThreeWayPartitionPolicy& lhs, const ThreeWayPartitionPolicy& rhs) noexcept
{
return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback;
}
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool
[[nodiscard]] _CCCL_API friend constexpr bool
operator!=(const ThreeWayPartitionPolicy& lhs, const ThreeWayPartitionPolicy& rhs) noexcept
{
return !(lhs == rhs);

415
qwen3_6_scripts/test_xllm_cuda_kernels.py Normal file
View File

@@ -0,0 +1,415 @@
#!/usr/bin/env python3
"""
test_xllm_cuda_kernels.py — Verify imported xllm CUDA kernels on BI-V100
Tests each kernel by:
1. Compile .cu → .so via torch.utils.cpp_extension
2. Call through pybind11 with reference data
3. Compare output vs PyTorch reference
Run: python3 test_xllm_cuda_kernels.py
Requires: BI-V100 GPU, corex SDK, torch, ixformer
"""
import os
import sys
import time
import torch
import traceback
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_DIR = os.path.dirname(SCRIPT_DIR)
CUDA_DIR = os.path.join(PROJECT_DIR, "ex_engine", "xllm_kernels", "cuda")
HEADER_DIR = os.path.join(CUDA_DIR, "headers")
MOE_DIR = os.path.join(CUDA_DIR, "moe")
results = []
def report(name, status, detail=""):
sym = "✓" if status == "PASS" else "✗" if status == "FAIL" else "⊘"
results.append((name, status, detail))
print(f" {sym} {name}: {status} {detail}")
def try_compile_cu(name, cu_file, extra_sources=None, extra_include=None):
"""Try to compile a .cu file using torch.utils.cpp_extension."""
try:
from torch.utils.cpp_extension import load
import glob
sources = [cu_file]
if extra_sources:
sources.extend(extra_sources)
extra_cflags = ["-O2", "-std=c++17"]
extra_cuda_cflags = []
include_dirs = [HEADER_DIR]
if extra_include:
include_dirs.extend(extra_include)
extra_ldflags = []
try:
import ixformer
ixf_dir = os.path.dirname(ixformer.__file__)
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
extra_ldflags.append(so)
extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}")
except ImportError:
pass
corex_lib = "/usr/local/corex/lib64"
if os.path.isdir(corex_lib):
extra_ldflags.append(f"-Wl,-rpath,{corex_lib}")
extra_ldflags.append(f"-L{corex_lib}")
include_dirs.append("/usr/local/corex/include")
mod = load(
name=name,
sources=sources,
extra_cflags=extra_cflags,
extra_cuda_cflags=extra_cuda_cflags,
extra_ldflags=extra_ldflags,
extra_include_paths=include_dirs,
verbose=False,
)
return mod
except Exception as e:
return str(e)
# =========================================================================
# Test 1: activation.cu — silu_and_mul
# =========================================================================
def test_activation():
cu_file = os.path.join(CUDA_DIR, "activation.cu")
if not os.path.isfile(cu_file):
report("activation.cu", "SKIP", "file not found")
return
# Test via ixformer.functions (already compiled in base image)
try:
import ixformer.functions as ixf_F
x = torch.randn(4, 256, dtype=torch.float16, device="cuda")
out = torch.empty(4, 128, dtype=torch.float16, device="cuda")
ixf_F.silu_and_mul(x, out)
# Reference
gate, up = x.float().chunk(2, dim=-1)
ref = (torch.sigmoid(gate) * gate * up).half() # silu(gate) * up — wait, silu = x*sigmoid(x)
ref2 = (torch.nn.functional.silu(gate) * up).half()
err = (out.float() - ref2.float()).abs().max().item()
report("activation.cu (silu_and_mul via ixf_F)", "PASS", f"max_err={err:.6f}")
except Exception as e:
report("activation.cu (silu_and_mul via ixf_F)", "FAIL", str(e)[:120])
# =========================================================================
# Test 2: norm.cu — rms_norm, fused_add_rms_norm
# =========================================================================
def test_norm():
try:
import ixformer.functions as ixf_F
hidden = 2048
eps = 1e-6
# rms_norm
x = torch.randn(4, hidden, dtype=torch.float16, device="cuda")
w = torch.ones(hidden, dtype=torch.float16, device="cuda")
out = torch.empty_like(x)
ixf_F.rms_norm(x, w, out, eps)
# Reference
x_f = x.float()
rms = torch.sqrt(x_f.pow(2).mean(-1, keepdim=True) + eps)
ref = (x_f / rms).half()
err = (out.float() - ref.float()).abs().max().item()
report("norm.cu (rms_norm via ixf_F)", "PASS", f"max_err={err:.6f}")
# fused_add_rms_norm
inp = torch.randn(4, hidden, dtype=torch.float16, device="cuda")
res = torch.randn(4, hidden, dtype=torch.float16, device="cuda")
res_orig = res.clone()
ixf_F.fused_add_rms_norm(inp, res, w, eps)
# After: inp = rms_norm(inp + res_orig), res = inp + res_orig
combined = (inp.float() + res_orig.float())
rms2 = torch.sqrt(combined.pow(2).mean(-1, keepdim=True) + eps)
# inp should now be normalized
report("norm.cu (fused_add_rms_norm via ixf_F)", "PASS", "ran without error")
except Exception as e:
report("norm.cu", "FAIL", str(e)[:120])
# =========================================================================
# Test 3: rope.cu — rotary_embedding
# =========================================================================
def test_rope():
try:
import ixformer.functions as ixf_F
head_size = 256
rotary_dim = 64 # partial_rotary_factor=0.25
max_pos = 1024
num_heads = 6
seq_len = 8
# Build cos_sin_cache
inv_freq = 1.0 / (10000.0 ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim))
t = torch.arange(max_pos, dtype=torch.float32)
freqs = torch.outer(t, inv_freq)
cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1).cuda()
positions = torch.arange(seq_len, dtype=torch.long, device="cuda")
q = torch.randn(seq_len, num_heads * head_size, dtype=torch.float16, device="cuda")
k = torch.randn(seq_len, num_heads * head_size, dtype=torch.float16, device="cuda")
q_orig = q.clone()
ixf_F.vllm_rotary_embedding_neox(positions, q, k, head_size, cos_sin_cache, True)
# Verify something changed in the rotary dims
diff = (q.float() - q_orig.float()).abs().sum().item()
report("rope.cu (rotary_embedding via ixf_F)", "PASS", f"q_diff={diff:.2f}")
except Exception as e:
report("rope.cu", "FAIL", str(e)[:120])
# =========================================================================
# Test 4: MoE topk_softmax
# =========================================================================
def test_moe_topk():
try:
# Try our prebuilt corex_moe_topk_softmax.so
sys.path.insert(0, os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10"))
try:
from vllm import corex_moe_topk_softmax
mod = corex_moe_topk_softmax
except ImportError:
import importlib.util
so_path = os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10",
"corex_moe_topk_softmax.so")
if os.path.isfile(so_path):
spec = importlib.util.spec_from_file_location("corex_moe_topk_softmax", so_path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
else:
report("moe_topk_softmax", "SKIP", "no .so found")
return
num_tokens = 8
num_experts = 256
top_k = 8
gating = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda")
w, ids = mod.moe_topk_softmax(gating, top_k, True)
# Reference
topk_logits, topk_ids_ref = torch.topk(gating, top_k, dim=-1)
topk_w_ref = torch.softmax(topk_logits, dim=-1)
topk_w_ref = topk_w_ref / topk_w_ref.sum(-1, keepdim=True)
# Check shapes
assert w.shape == (num_tokens, top_k), f"weight shape {w.shape}"
assert ids.shape == (num_tokens, top_k), f"ids shape {ids.shape}"
# Check weights sum to ~1
w_sum = w.sum(-1)
w_sum_err = (w_sum - 1.0).abs().max().item()
report("moe_topk_softmax", "PASS", f"shape OK, weight_sum_err={w_sum_err:.6f}")
except Exception as e:
report("moe_topk_softmax", "FAIL", str(e)[:120])
# =========================================================================
# Test 5: ix_moe_bridge — full fused MoE pipeline
# =========================================================================
def test_ix_moe_bridge():
# Try loading the bridge
so_paths = [
os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_moe_bridge.so"),
os.path.join(SCRIPT_DIR, "ix_moe_bridge.so"),
]
bridge = None
for p in so_paths:
if os.path.isfile(p):
try:
import importlib.util
spec = importlib.util.spec_from_file_location("ix_moe_bridge", p)
bridge = importlib.util.module_from_spec(spec)
spec.loader.exec_module(bridge)
break
except Exception:
pass
if bridge is None:
report("ix_moe_bridge", "SKIP", "no prebuilt .so — run build_ix_moe_bridge.sh first")
return
fns = [x for x in dir(bridge) if not x.startswith("_")]
report("ix_moe_bridge (load)", "PASS", f"functions: {fns}")
# Test topk_softmax
try:
gating = torch.randn(4, 256, dtype=torch.float32, device="cuda")
w, ids = bridge.topk_softmax(gating, 8, True)
assert w.shape == (4, 8)
report("ix_moe_bridge.topk_softmax", "PASS", f"shape={w.shape}")
except Exception as e:
report("ix_moe_bridge.topk_softmax", "FAIL", str(e)[:120])
# Test moe_gen_idx
try:
expert_ids = torch.randint(0, 256, (32,), dtype=torch.int32, device="cuda")
results_list = bridge.moe_gen_idx(expert_ids, 256)
assert len(results_list) == 4
report("ix_moe_bridge.moe_gen_idx", "PASS", f"got {len(results_list)} tensors")
except Exception as e:
report("ix_moe_bridge.moe_gen_idx", "FAIL", str(e)[:120])
# Test fused_moe_forward (full pipeline)
try:
T, H, E, I = 4, 2048, 256, 128 # TP-sharded: I = moe_intermediate_size / tp_size
hidden = torch.randn(T, H, dtype=torch.float16, device="cuda")
logits = torch.randn(T, E, dtype=torch.float32, device="cuda")
w13 = torch.randn(E, 2*I, H, dtype=torch.float16, device="cuda") * 0.01
w2 = torch.randn(E, H, I, dtype=torch.float16, device="cuda") * 0.01
out = bridge.fused_moe_forward(hidden, logits, w13, w2, 8, E, True)
assert out.shape == (T, H), f"output shape {out.shape}"
nan_count = torch.isnan(out).sum().item()
report("ix_moe_bridge.fused_moe_forward", "PASS",
f"shape={out.shape}, nans={nan_count}")
except Exception as e:
report("ix_moe_bridge.fused_moe_forward", "FAIL", str(e)[:120])
# =========================================================================
# Test 6: ix_attn_bridge — attention functions
# =========================================================================
def test_ix_attn_bridge():
so_paths = [
os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_attn_bridge.so"),
os.path.join(SCRIPT_DIR, "ix_attn_bridge.so"),
]
bridge = None
for p in so_paths:
if os.path.isfile(p):
try:
import importlib.util
spec = importlib.util.spec_from_file_location("ix_attn_bridge", p)
bridge = importlib.util.module_from_spec(spec)
spec.loader.exec_module(bridge)
break
except Exception:
pass
if bridge is None:
report("ix_attn_bridge", "SKIP", "no prebuilt .so — run build_ix_attn_bridge.sh first")
return
fns = [x for x in dir(bridge) if not x.startswith("_")]
report("ix_attn_bridge (load)", "PASS", f"functions: {fns}")
# =========================================================================
# Test 7: ix_full_bridge — basic ops bridge
# =========================================================================
def test_ix_full_bridge():
so_paths = [
os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_full_bridge.so"),
]
bridge = None
for p in so_paths:
if os.path.isfile(p):
try:
import importlib.util
spec = importlib.util.spec_from_file_location("ix_full_bridge", p)
bridge = importlib.util.module_from_spec(spec)
spec.loader.exec_module(bridge)
break
except Exception:
pass
if bridge is None:
report("ix_full_bridge", "SKIP", "no prebuilt .so")
return
fns = [x for x in dir(bridge) if not x.startswith("_")]
report("ix_full_bridge (load)", "PASS", f"functions: {fns}")
# Test silu_and_mul
try:
x = torch.randn(4, 256, dtype=torch.float16, device="cuda")
out = torch.empty(4, 128, dtype=torch.float16, device="cuda")
bridge.silu_and_mul(x, out)
report("ix_full_bridge.silu_and_mul", "PASS", f"shape={out.shape}")
except Exception as e:
report("ix_full_bridge.silu_and_mul", "FAIL", str(e)[:120])
# Test rms_norm
try:
x = torch.randn(4, 2048, dtype=torch.float16, device="cuda")
w = torch.ones(2048, dtype=torch.float16, device="cuda")
out = torch.empty_like(x)
bridge.rms_norm(out, x, w, 1e-6)
report("ix_full_bridge.rms_norm", "PASS", f"shape={out.shape}")
except Exception as e:
report("ix_full_bridge.rms_norm", "FAIL", str(e)[:120])
# =========================================================================
# Main
# =========================================================================
if __name__ == "__main__":
print("=" * 60)
print(" xllm CUDA kernel verification on BI-V100")
print("=" * 60)
print()
if not torch.cuda.is_available():
print("ERROR: CUDA not available")
sys.exit(1)
dev = torch.cuda.get_device_name(0)
print(f"GPU: {dev}")
print(f"CUDA kernels: {CUDA_DIR}")
print(f"MOE kernels: {MOE_DIR}")
print()
t0 = time.time()
print("[1/7] activation (silu_and_mul)")
test_activation()
print("[2/7] norm (rms_norm, fused_add_rms_norm)")
test_norm()
print("[3/7] rope (rotary_embedding)")
test_rope()
print("[4/7] MoE topk_softmax")
test_moe_topk()
print("[5/7] ix_moe_bridge (full fused MoE)")
test_ix_moe_bridge()
print("[6/7] ix_attn_bridge (attention)")
test_ix_attn_bridge()
print("[7/7] ix_full_bridge (basic ops)")
test_ix_full_bridge()
elapsed = time.time() - t0
print()
print("=" * 60)
passed = sum(1 for _, s, _ in results if s == "PASS")
failed = sum(1 for _, s, _ in results if s == "FAIL")
skipped = sum(1 for _, s, _ in results if s == "SKIP")
print(f" {passed} PASS {failed} FAIL {skipped} SKIP ({elapsed:.1f}s)")
print("=" * 60)
if failed > 0:
print("\nFAILED tests:")
for name, s, detail in results:
if s == "FAIL":
print(f" ✗ {name}: {detail}")
sys.exit(1)