Skip to content

fix: update hygon KDA warp and add moe topk ops support #1486

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

Merged
PanZezhong1725 merged 1 commit into from
Aug 11, 2026
Merged

fix: update hygon KDA warp and add moe topk ops support #1486

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 3 additions & 2 deletions src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,9 @@ static infiniStatus_t launch_warp_sequence(const KimiDeltaAttentionInfo &info,
const void *final_state_indices,
cudaStream_t stream) {
constexpr size_t D = 128;
constexpr size_t WARPS_PER_BLOCK = 8;
constexpr size_t NUM_THREADS = WARPS_PER_BLOCK * 32;
constexpr size_t WARP_SIZE = INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE;
constexpr size_t WARPS_PER_BLOCK = 256 / WARP_SIZE;
constexpr size_t NUM_THREADS = WARPS_PER_BLOCK * WARP_SIZE;
const dim3 grid(
static_cast<uint32_t>(info.B),
static_cast<uint32_t>(info.H),
Expand Down
2 changes: 1 addition & 1 deletion src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
* Licensed under the Apache License, Version 2.0.
*/

#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)

#include "moe_topk_sigmoid_nvidia.cuh"

Expand Down
14 changes: 13 additions & 1 deletion src/infiniop/ops/moe_topk_sigmoid/operator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
#include "../../handle.h"
#include "infiniop/ops/moe_topk_sigmoid.h"

#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
#include "nvidia/moe_topk_sigmoid_nvidia.cuh"
#endif

Expand All @@ -25,6 +25,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeTopkSigmoidDescriptor(
#endif
#ifdef ENABLE_ILUVATAR_API
CREATE(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
CREATE(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand All @@ -45,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeTopkSigmoidWorkspaceSize(
#endif
#ifdef ENABLE_ILUVATAR_API
GET(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
GET(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand All @@ -71,6 +77,9 @@ __INFINI_C infiniStatus_t infiniopMoeTopkSigmoid(
#endif
#ifdef ENABLE_ILUVATAR_API
CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
CALCULATE(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand All @@ -90,6 +99,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeTopkSigmoidDescriptor(
#endif
#ifdef ENABLE_ILUVATAR_API
DESTROY(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
DESTROY(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand Down
2 changes: 1 addition & 1 deletion src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
* Licensed under the Apache License, Version 2.0.
*/

#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)

#include "moe_topk_softmax_nvidia.cuh"

Expand Down
14 changes: 13 additions & 1 deletion src/infiniop/ops/moe_topk_softmax/operator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
#include "../../handle.h"
#include "infiniop/ops/moe_topk_softmax.h"

#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
#include "nvidia/moe_topk_softmax_nvidia.cuh"
#endif

Expand All @@ -27,6 +27,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeTopkSoftmaxDescriptor(
#endif
#ifdef ENABLE_ILUVATAR_API
CREATE(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
CREATE(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand All @@ -47,6 +50,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeTopkSoftmaxWorkspaceSize(
#endif
#ifdef ENABLE_ILUVATAR_API
GET(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
GET(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand All @@ -73,6 +79,9 @@ __INFINI_C infiniStatus_t infiniopMoeTopkSoftmax(
#endif
#ifdef ENABLE_ILUVATAR_API
CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
CALCULATE(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand All @@ -92,6 +101,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeTopkSoftmaxDescriptor(
#endif
#ifdef ENABLE_ILUVATAR_API
DESTROY(INFINI_DEVICE_ILUVATAR, nvidia);
#endif
#ifdef ENABLE_HYGON_API
DESTROY(INFINI_DEVICE_HYGON, nvidia);
#endif
default:
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
Expand Down
22 changes: 17 additions & 5 deletions src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,14 @@
#include <cmath>
#include <cstdint>

#if defined(ENABLE_HYGON_API)
#define INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE 64
#define INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK 0xffffffffffffffffULL
#else
#define INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE 32
#define INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK 0xffffffffU
#endif

namespace op::recurrent_gated_delta_rule::cuda {

template <typename T>
Expand Down Expand Up @@ -35,11 +43,14 @@ __device__ inline int64_t loadOptionalIndex(const void *indices,

template <typename Tcompute>
__device__ inline Tcompute warpReduceSum(Tcompute value) {
constexpr int WARP_SIZE = INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
value += __shfl_down_sync(0xffffffff, value, offset);
for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
value += __shfl_down_sync(
INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK, value, offset, WARP_SIZE);
}
return __shfl_sync(0xffffffff, value, 0);
return __shfl_sync(
INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK, value, 0, WARP_SIZE);
}

template <typename Tcompute>
Expand Down Expand Up @@ -118,7 +129,7 @@ __device__ void recurrentDeltaRuleWarpSequence(
ptrdiff_t v_s2,
GatePolicy gate_policy,
Tcompute *shared) {
constexpr int WARP_SIZE = 32;
constexpr int WARP_SIZE = INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE;
constexpr int NUM_THREADS = WARPS_PER_BLOCK * WARP_SIZE;
constexpr int STATE_VALUES_PER_LANE = (Dk + WARP_SIZE - 1) / WARP_SIZE;

Expand Down Expand Up @@ -257,7 +268,8 @@ __device__ void recurrentDeltaRuleWarpSequence(
const ptrdiff_t out_base = static_cast<ptrdiff_t>(token_batch) * out_s0 + static_cast<ptrdiff_t>(token_idx) * out_s1 + static_cast<ptrdiff_t>(value_head_idx) * out_s2;
out[out_base + value_dim_idx] = static_cast<Tdata>(hq_memory + delta * kq_memory);
}
delta = __shfl_sync(0xffffffff, delta, 0);
delta = __shfl_sync(
INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK, delta, 0, WARP_SIZE);

#pragma unroll
for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) {
Expand Down
7 changes: 4 additions & 3 deletions src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ infiniStatus_t launchIndexedPoolWarpKernelTyped(
bool initial_state_indices_i64,
bool final_state_indices_i64,
cudaStream_t stream) {
constexpr size_t NUM_THREADS = WARPS_PER_BLOCK * 32;
constexpr size_t NUM_THREADS = WARPS_PER_BLOCK * INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE;
dim3 grid(uint32_t(_info.B), uint32_t(_info.Hv), uint32_t((_info.Dv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK));
dim3 block(NUM_THREADS);
size_t shared_mem_size = (Dk * 3 + NUM_THREADS + 1) * sizeof(float);
Expand Down Expand Up @@ -253,16 +253,17 @@ infiniStatus_t Descriptor::calculate(
bool initial_indices_i64 = _info.initial_state_indices_dtype == INFINI_DTYPE_I64;
bool final_indices_i64 = _info.final_state_indices_dtype == INFINI_DTYPE_I64;

constexpr size_t WARPS_PER_BLOCK = 256 / INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE;
if (_info.Dk == 128 && _info.Dv == 128) {
if (_opaque->internal->maxThreadsPerBlock() >= 256) {
return launchIndexedPoolWarpKernel<128, 128, 8>(
return launchIndexedPoolWarpKernel<128, 128, WARPS_PER_BLOCK>(
_info, out, initial_state, final_state, q, k, v, g, beta,
initial_state_indices, final_state_indices,
initial_indices_i64, final_indices_i64, stream);
}
} else if (_info.Dk == 64 && _info.Dv == 64) {
if (_opaque->internal->maxThreadsPerBlock() >= 256) {
return launchIndexedPoolWarpKernel<64, 64, 8>(
return launchIndexedPoolWarpKernel<64, 64, WARPS_PER_BLOCK>(
_info, out, initial_state, final_state, q, k, v, g, beta,
initial_state_indices, final_state_indices,
initial_indices_i64, final_indices_i64, stream);
Expand Down
14 changes: 13 additions & 1 deletion src/infiniop/ops/recurrent_gated_delta_rule/operator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
#include "../../handle.h"
#include "infiniop/ops/recurrent_gated_delta_rule.h"

#if defined(ENABLE_NVIDIA_API)
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API)
#include "nvidia/recurrent_gated_delta_rule_nvidia.cuh"
#endif
#ifdef ENABLE_METAX_API
Expand Down Expand Up @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor(
#ifdef ENABLE_NVIDIA_API
CREATE(INFINI_DEVICE_NVIDIA, nvidia)
#endif
#ifdef ENABLE_HYGON_API
CREATE(INFINI_DEVICE_HYGON, nvidia)
#endif
#ifdef ENABLE_METAX_API
CREATE(INFINI_DEVICE_METAX, metax)
#endif
Expand All @@ -69,6 +72,9 @@ __INFINI_C infiniStatus_t infiniopGetRecurrentGatedDeltaRuleWorkspaceSize(
#ifdef ENABLE_NVIDIA_API
GET(INFINI_DEVICE_NVIDIA, nvidia)
#endif
#ifdef ENABLE_HYGON_API
GET(INFINI_DEVICE_HYGON, nvidia)
#endif
#ifdef ENABLE_METAX_API
GET(INFINI_DEVICE_METAX, metax)
#endif
Expand Down Expand Up @@ -103,6 +109,9 @@ __INFINI_C infiniStatus_t infiniopRecurrentGatedDeltaRule(
#ifdef ENABLE_NVIDIA_API
CALCULATE(INFINI_DEVICE_NVIDIA, nvidia)
#endif
#ifdef ENABLE_HYGON_API
CALCULATE(INFINI_DEVICE_HYGON, nvidia)
#endif
#ifdef ENABLE_METAX_API
CALCULATE(INFINI_DEVICE_METAX, metax)
#endif
Expand All @@ -128,6 +137,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRecurrentGatedDeltaRuleDescriptor(
#ifdef ENABLE_NVIDIA_API
DESTROY(INFINI_DEVICE_NVIDIA, nvidia)
#endif
#ifdef ENABLE_HYGON_API
DESTROY(INFINI_DEVICE_HYGON, nvidia)
#endif
#ifdef ENABLE_METAX_API
DESTROY(INFINI_DEVICE_METAX, metax)
#endif
Expand Down
123 changes: 69 additions & 54 deletions test/infinicore/ops/kimi_delta_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,62 @@ def torch_kimi_delta_attention_ref(
return out.to(initial_dtype)


def k3_shape_test_case(seq_len, num_heads, description):
shape = (1, seq_len, num_heads, 128)
dtype = infinicore.bfloat16
H, D = shape[2], shape[3]
cu = torch.tensor([0, seq_len], dtype=torch.int32)
initial_indices = torch.tensor([0], dtype=torch.int32)
final_indices = torch.tensor([1], dtype=torch.int32)
return TestCase(
inputs=[
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape[:3], None, dtype),
TensorSpec.from_tensor((H,), None, infinicore.float32),
TensorSpec.from_tensor((H, D), None, infinicore.float32),
TensorSpec.from_tensor(
(2, H, D, D),
None,
dtype,
init_mode=TensorInitializer.ZEROS,
),
],
kwargs={
"cu_seqlens": TensorSpec.from_tensor(
tuple(cu.shape),
None,
infinicore.int32,
init_mode=TensorInitializer.MANUAL,
set_tensor=cu,
),
"initial_state_indices": TensorSpec.from_tensor(
tuple(initial_indices.shape),
None,
infinicore.int32,
init_mode=TensorInitializer.MANUAL,
set_tensor=initial_indices,
),
"final_state_indices": TensorSpec.from_tensor(
tuple(final_indices.shape),
None,
infinicore.int32,
init_mode=TensorInitializer.MANUAL,
set_tensor=final_indices,
),
"scale": D**-0.5,
"lower_bound": -5.0,
"use_qk_l2norm": True,
},
output_spec=None,
comparison_target=None,
tolerance={"atol": 1e-3, "rtol": 1e-2},
description=description,
)


def parse_test_cases():
tests = []
for dtype in _TENSOR_DTYPES:
Expand Down Expand Up @@ -170,60 +226,19 @@ def parse_test_cases():
)
)

shape = (1, 3, 96, 128)
dtype = infinicore.bfloat16
H, D = shape[2], shape[3]
cu = torch.tensor([0, shape[1]], dtype=torch.int32)
initial_indices = torch.tensor([0], dtype=torch.int32)
final_indices = torch.tensor([1], dtype=torch.int32)
tests.append(
TestCase(
inputs=[
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape, None, dtype),
TensorSpec.from_tensor(shape[:3], None, dtype),
TensorSpec.from_tensor((H,), None, infinicore.float32),
TensorSpec.from_tensor((H, D), None, infinicore.float32),
TensorSpec.from_tensor(
(2, H, D, D),
None,
dtype,
init_mode=TensorInitializer.ZEROS,
),
],
kwargs={
"cu_seqlens": TensorSpec.from_tensor(
tuple(cu.shape),
None,
infinicore.int32,
init_mode=TensorInitializer.MANUAL,
set_tensor=cu,
),
"initial_state_indices": TensorSpec.from_tensor(
tuple(initial_indices.shape),
None,
infinicore.int32,
init_mode=TensorInitializer.MANUAL,
set_tensor=initial_indices,
),
"final_state_indices": TensorSpec.from_tensor(
tuple(final_indices.shape),
None,
infinicore.int32,
init_mode=TensorInitializer.MANUAL,
set_tensor=final_indices,
),
"scale": D**-0.5,
"lower_bound": -5.0,
"use_qk_l2norm": True,
},
output_spec=None,
comparison_target=None,
tolerance={"atol": 1e-3, "rtol": 1e-2},
description="KimiDeltaAttention K3 shape indexed-pool varlen zero-state",
)
tests.extend(
[
k3_shape_test_case(
3,
96,
"KimiDeltaAttention K3 shape indexed-pool varlen zero-state",
),
k3_shape_test_case(
91,
1,
"KimiDeltaAttention D=128 long-sequence wavefront regression",
),
]
)
return tests

Expand Down
Loading
Loading