Skip to content

Pr 1451 #1503

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wooway777 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Pr 1451 #1503

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7015997
fix(nvidia): support graph replay, MoE decode, and diff builds
qinyiqun Aug 3, 2026
a4e548d
fix(nvidia): resolve elementwise build after rebase
wooway777 Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 60 additions & 79 deletions src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,23 +172,22 @@ INFINIOP_CUDA_KERNEL elementwiseKernel(
const ptrdiff_t *__restrict__ output_strides,
const ptrdiff_t *__restrict__ input_strides,
Tdata *output,
const void *const *inputs,
InputPointerArray<N> inputs,
size_t offset,
Args... args) {

size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset;

if (idx < output_size) {
const Tdata *const *typed_inputs = reinterpret_cast<const Tdata *const *>(inputs);
size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides);
InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides};

unpackInputsAndApply(
[&](auto... Is) {
#if defined(ENABLE_HYGON_API)
output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., args...);
output[out_idx] = Op{}(typedInputPtr<Tdata>(inputs.values[Is.value])[indexer(Is.value)]..., args...);
#else
output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., std::forward<Args>(args)...);
output[out_idx] = Op{}(typedInputPtr<Tdata>(inputs.values[Is.value])[indexer(Is.value)]..., std::forward<Args>(args)...);
#endif
},
std::make_index_sequence<N>{});
Expand Down Expand Up @@ -278,7 +277,7 @@ INFINIOP_CUDA_KERNEL elementwiseKernel(
const ptrdiff_t *__restrict__ output_strides,
const ptrdiff_t *__restrict__ input_strides,
Tout *output,
const void *const *__restrict__ inputs,
InputPointerArray<sizeof...(Tin)> inputs,
size_t offset) {

size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset;
Expand All @@ -290,7 +289,7 @@ INFINIOP_CUDA_KERNEL elementwiseKernel(
unpackInputsAndApply(
[&](auto... Is) {
output[out_idx] = Op{}.template operator()<Tout, Tin...>(
(typedInputPtr<Tin>(inputs[Is.value])[indexer(Is.value)])...);
(typedInputPtr<Tin>(inputs.values[Is.value])[indexer(Is.value)])...);
},
std::index_sequence_for<Tin...>{});
}
Expand Down Expand Up @@ -336,9 +335,52 @@ INFINIOP_CUDA_KERNEL inlineMetaElementwiseKernel(

struct DeviceImpl::Opaque {
std::shared_ptr<device::nvidia::Handle::Internal> internal;
void *device_meta = nullptr;
const bool *input_contiguous = nullptr;
const bool *input_broadcasted = nullptr;
const size_t *output_shape = nullptr;
const ptrdiff_t *output_strides = nullptr;
const size_t *input_shapes = nullptr;
const ptrdiff_t *input_strides = nullptr;
infiniStatus_t init_status = INFINI_STATUS_SUCCESS;

Opaque(const std::shared_ptr<device::nvidia::Handle::Internal> &internal_,
const op::elementwise::ElementwiseInfo &info)
: internal(internal_), init_status(initialize(info)) {}

~Opaque() {
if (device_meta != nullptr) {
cudaFree(device_meta);
}
}

Opaque(const std::shared_ptr<device::nvidia::Handle::Internal> &internal)
: internal(internal) {}
infiniStatus_t initialize(const op::elementwise::ElementwiseInfo &info) {
if (info.canUseContiguousFastPath() || info.canUseInlineMetaFastPath()) {
return INFINI_STATUS_SUCCESS;
}

const auto meta_size = info.getMetaMemSize();
if (meta_size == 0) {
return INFINI_STATUS_SUCCESS;
}

CHECK_CUDA(cudaMalloc(&device_meta, meta_size));
CHECK_CUDA(cudaMemcpy(device_meta,
info.getMetaStart(),
meta_size,
cudaMemcpyHostToDevice));

const auto ndim = info.getNdim();
const auto input_size = info.getInputSize();
output_shape = reinterpret_cast<const size_t *>(device_meta);
output_strides = reinterpret_cast<const ptrdiff_t *>(output_shape + ndim);
input_shapes = reinterpret_cast<const size_t *>(output_strides + ndim);
input_strides = reinterpret_cast<const ptrdiff_t *>(input_shapes + input_size * ndim);
input_contiguous = reinterpret_cast<const bool *>(input_strides + input_size * ndim);
input_broadcasted = input_contiguous + input_size;

return INFINI_STATUS_SUCCESS;
}

/**
* @brief Executes an elementwise operation where all inputs and the output share the same data type.
Expand Down Expand Up @@ -542,60 +584,6 @@ private:
return INFINI_STATUS_SUCCESS;
}

/**
* @brief Transfers elementwise operation metadata and input pointers from host to device memory.
*
* @tparam N Number of input tensors.
*
* @param info Elementwise operation metadata (shapes, strides, flags, etc.).
* @param workspace Pointer to device workspace memory for storing metadata and input pointers.
* @param h_inputs_arr Host array of input tensor pointers.
* @param d_inputs_arr Input reference to device array of input tensor pointers.
* @param d_input_contiguous Input reference to device array indicating whether each input is contiguous.
* @param d_input_broadcasted Input reference to device array indicating whether each input is broadcasted.
* @param d_output_shape Output reference to device array holding the output tensor shape.
* @param d_output_strides Output reference to device array holding output tensor strides.
* @param d_input_shapes Output reference to flattened input tensor shapes (N * ndim).
* @param d_input_strides Output reference to flattened input tensor strides (N * ndim).
* @param stream CUDA stream used for asynchronous memory transfer.
* @return infiniStatus_t Status indicating success or failure of the memory transfer and setup.
*/
template <size_t N>
infiniStatus_t infoToDevice(
const op::elementwise::ElementwiseInfo &info,
void *workspace,
const void *const *h_inputs_arr,
const void **&d_inputs_arr,
const bool *&d_input_contiguous,
const bool *&d_input_broadcasted,
const size_t *&d_output_shape,
const ptrdiff_t *&d_output_strides,
const size_t *&d_input_shapes,
const ptrdiff_t *&d_input_strides,
cudaStream_t stream) const {

constexpr auto input_size = N;
const auto ndim = info.getNdim();
constexpr auto input_arr_size = N * sizeof(*h_inputs_arr);
const int8_t *info_meta_start = info.getMetaStart();
const int8_t *d_meta_start = reinterpret_cast<int8_t *>(workspace) + input_arr_size;

// copy the input pointer array and meta to device
CHECK_CUDA(cudaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, cudaMemcpyHostToDevice, stream));
CHECK_CUDA(cudaMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), cudaMemcpyHostToDevice, stream));

// offset/assign the pointers
d_inputs_arr = reinterpret_cast<const void **>(workspace);
d_output_shape = reinterpret_cast<const size_t *>(d_meta_start);
d_output_strides = reinterpret_cast<const ptrdiff_t *>(d_output_shape + ndim);
d_input_shapes = reinterpret_cast<const size_t *>(d_output_strides + ndim);
d_input_strides = reinterpret_cast<const ptrdiff_t *>(d_input_shapes + input_size * ndim);
d_input_contiguous = reinterpret_cast<const bool *>(d_input_strides + input_size * ndim);
d_input_broadcasted = reinterpret_cast<const bool *>(d_input_contiguous + input_size);

return INFINI_STATUS_SUCCESS;
}

/**
* @brief Launches the elementwise kernel for the specified operation.
*
Expand Down Expand Up @@ -629,19 +617,9 @@ private:
return INFINI_STATUS_SUCCESS;
}

// Device pointers
const void **d_inputs_arr = nullptr;
const bool *d_input_contiguous = nullptr;
const bool *d_input_broadcasted = nullptr;
const size_t *d_output_shape = nullptr;
const ptrdiff_t *d_output_strides = nullptr;
const size_t *d_input_shapes = nullptr;
const ptrdiff_t *d_input_strides = nullptr;

CHECK_STATUS(infoToDevice<N>(info, workspace, inputs.data(), d_inputs_arr,
d_input_contiguous, d_input_broadcasted,
d_output_shape, d_output_strides,
d_input_shapes, d_input_strides, stream));
(void)workspace;
InputPointerArray<N> input_ptrs{};
std::copy_n(inputs.begin(), N, input_ptrs.values);

dim3 blockDims(std::min(BLOCK_SIZE, static_cast<uint32_t>(internal->maxThreadsPerBlock())));
dim3 gridDims(std::min(uint32_t(CEIL_DIV(output_size, blockDims.x)), static_cast<uint32_t>(internal->gridSizeX())));
Expand All @@ -650,10 +628,10 @@ private:
for (size_t i = 0; i < output_size; i += step) {
kernel_func<<<gridDims, blockDims, 0, stream>>>(
output_size, info.getNdim(), info.isOutputContiguous(),
d_input_contiguous, d_input_broadcasted,
d_output_shape, d_input_shapes,
d_output_strides, d_input_strides,
output, reinterpret_cast<const void **>(d_inputs_arr),
input_contiguous, input_broadcasted,
output_shape, input_shapes,
output_strides, input_strides,
output, input_ptrs,
i, std::forward<Args>(args)...);
}

Expand All @@ -664,6 +642,9 @@ private:
template <typename... Args>
utils::Result<DeviceImpl *> DeviceImpl::create(Args &&...args) {
auto opaque = std::make_shared<Opaque>(std::forward<Args>(args)...);
if (opaque->init_status != INFINI_STATUS_SUCCESS) {
return opaque->init_status;
}
return utils::Result<DeviceImpl *>(new DeviceImpl(opaque));
}

Expand Down
4 changes: 2 additions & 2 deletions src/infiniop/elementwise/nvidia/elementwise_nvidia_api.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,9 +93,9 @@ public:
auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \
CHECK_RESULT(info_result); \
auto info = info_result.take(); \
auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); \
size_t workspace_size = 0; \
\
auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(HANDLE->internal()); \
auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(HANDLE->internal(), info); \
CHECK_RESULT(device_impl_result); \
\
*desc_ptr = new Descriptor( \
Expand Down
3 changes: 3 additions & 0 deletions src/infiniop/ops/diff/nvidia/diff_nvidia.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
#ifndef __DIFF_NVIDIA_H__
#define __DIFF_NVIDIA_H__

#include "../../../../utils.h"
#include "../../../operator.h"
#include <cstddef>
#include <utility>
#include <vector>

namespace op::diff::nvidia {

Expand Down
4 changes: 2 additions & 2 deletions src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,9 +87,9 @@ infiniStatus_t Descriptor::create(
auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec);
CHECK_RESULT(info_result);
auto info = info_result.take();
auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *);
size_t workspace_size = 0;

auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(handle->internal());
auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(handle->internal(), info);
CHECK_RESULT(device_impl_result);

*desc_ptr = new Descriptor(
Expand Down
19 changes: 14 additions & 5 deletions src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -403,6 +403,7 @@ infiniStatus_t launch_cutlass_gemm_grouped_device_meta(int problem_count,
int64_t *d_ldb,
int64_t *d_ldc,
int64_t *d_ldd,
bool full_occupancy_wave,
cudaStream_t stream) {
if (problem_count == 0) {
return INFINI_STATUS_SUCCESS;
Expand Down Expand Up @@ -436,7 +437,15 @@ infiniStatus_t launch_cutlass_gemm_grouped_device_meta(int problem_count,
cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>::GemmKernel;
using Gemm = cutlass::gemm::device::GemmGrouped<GemmKernel>;

const int threadblock_count = std::min(Gemm::sufficient(), problem_count);
const int occupancy_wave = Gemm::sufficient();
// Decode routes to only top-k experts, but each expert still contains many
// N-dimension tiles. A full persistent wave lets CUTLASS distribute those
// tiles across the GPU instead of leaving most SMs idle. Prefill retains
// the previous problem-count cap to avoid scheduler overhead from excess
// threadblocks when most experts receive no tokens.
const int threadblock_count = full_occupancy_wave
? occupancy_wave
: std::min(occupancy_wave, problem_count);
if (threadblock_count <= 0) {
return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED;
}
Expand Down Expand Up @@ -547,7 +556,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
// A D2H count copy and stream sync cannot be captured for replay.
CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, stream));

swiglu_kernel<T><<<(topk * intermediate_size + 255) / 256, 256, 0, stream>>>(gate_up, activated, topk, intermediate_size);

Expand All @@ -560,7 +569,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
output_permutation, pairs, num_experts, hidden_size, intermediate_size, block_size);
CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, stream));

apply_shuffle_mul_sum_kernel<T><<<1, std::min(hidden_size, 1024), 0, stream>>>(
expert_out, reinterpret_cast<T *>(output), output_permutation,
Expand Down Expand Up @@ -604,7 +613,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
// avoids making the runtime problem count depend on device routing data.
CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, stream));

swiglu_kernel<T><<<(max_num_tokens_padded * intermediate_size + 255) / 256, 256, 0, stream>>>(
gate_up, activated, max_num_tokens_padded, intermediate_size);
Expand All @@ -615,7 +624,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info,
w2_t, expert_out, num_experts, hidden_size, intermediate_size);
CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta<CutlassT>(
num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d,
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream));
grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, stream));

apply_shuffle_mul_sum_kernel<T><<<num_tokens, std::min(hidden_size, 1024), 0, stream>>>(
expert_out, reinterpret_cast<T *>(output), output_permutation,
Expand Down
8 changes: 3 additions & 5 deletions src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,7 @@ infiniStatus_t calculateMulScalar(
return launchMulScalarKernel<T>(info.numel(), output, input, alpha, stream);
}

if (workspace_size < info.elementwise_info.getMetaMemSize() + sizeof(void *)) {
return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
}
(void)workspace_size;

return device_info->calculate<256, MulScalarOp, T>(
info.elementwise_info,
Expand Down Expand Up @@ -105,8 +103,8 @@ infiniStatus_t Descriptor::create(
CHECK_RESULT(result);

auto info = result.take();
auto workspace_size = info.elementwise_info.getMetaMemSize() + sizeof(void *);
auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(handle->internal());
size_t workspace_size = 0;
auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(handle->internal(), info.elementwise_info);
CHECK_RESULT(device_impl_result);

*desc_ptr = new Descriptor(
Expand Down
Loading