diff --git a/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu b/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu index 8b7ba6bff..f6b6744ac 100644 --- a/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu +++ b/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu @@ -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(info.B), static_cast(info.H), diff --git a/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu b/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu index a2bbf298c..993f44445 100644 --- a/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu +++ b/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu @@ -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" diff --git a/src/infiniop/ops/moe_topk_sigmoid/operator.cc b/src/infiniop/ops/moe_topk_sigmoid/operator.cc index 97f255c70..a8744ee40 100644 --- a/src/infiniop/ops/moe_topk_sigmoid/operator.cc +++ b/src/infiniop/ops/moe_topk_sigmoid/operator.cc @@ -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 @@ -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; @@ -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; @@ -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; @@ -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; diff --git a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu index 952b3ca5d..2fc0e992b 100644 --- a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu +++ b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu @@ -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" diff --git a/src/infiniop/ops/moe_topk_softmax/operator.cc b/src/infiniop/ops/moe_topk_softmax/operator.cc index 213ad93ce..8cf30478d 100644 --- a/src/infiniop/ops/moe_topk_softmax/operator.cc +++ b/src/infiniop/ops/moe_topk_softmax/operator.cc @@ -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 @@ -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; @@ -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; @@ -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; @@ -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; diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh b/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh index 27931bb17..af56af065 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh +++ b/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh @@ -4,6 +4,14 @@ #include #include +#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 @@ -35,11 +43,14 @@ __device__ inline int64_t loadOptionalIndex(const void *indices, template __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 @@ -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; @@ -257,7 +268,8 @@ __device__ void recurrentDeltaRuleWarpSequence( const ptrdiff_t out_base = static_cast(token_batch) * out_s0 + static_cast(token_idx) * out_s1 + static_cast(value_head_idx) * out_s2; out[out_base + value_dim_idx] = static_cast(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) { diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu b/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu index 367ab66f6..f07d7e33b 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu +++ b/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu @@ -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); @@ -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); diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc index fc272fdf9..b9ad3b4af 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 diff --git a/test/infinicore/ops/kimi_delta_attention.py b/test/infinicore/ops/kimi_delta_attention.py index 9f86709a5..ee604e59e 100644 --- a/test/infinicore/ops/kimi_delta_attention.py +++ b/test/infinicore/ops/kimi_delta_attention.py @@ -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: @@ -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 diff --git a/xmake/hygon.lua b/xmake/hygon.lua index ed7da7b85..ba06ae972 100644 --- a/xmake/hygon.lua +++ b/xmake/hygon.lua @@ -123,6 +123,7 @@ target("infiniop-hygon") add_hygon_dtk_paths() set_warnings("all", "error") + add_cuflags("-Wno-pass-failed", {force = true}) add_cuflags("-Wno-error=unused-private-field", {force = true}) add_cuflags("-Wno-return-type", {force = true}) -- 抑制return语句警告 add_cuflags("-Wno-error=macro-redefined", {force = true})