Skip to content

Commit ab4bcf8

Browse files
maojunx99maojunx99
committed
feat(npu): Added implementation of the recurrent_gated_delta_rule operator
Implemented the recurrent_gated_delta_rule operator on the NPU, including: 1. Adding operator interface declaration 2. Implementing NPU backend computation logic 3. Updating CMake build configuration 4. Integrating and using it in Qwen3GatedDeltaNetBase
1 parent f9a75b4 commit ab4bcf8

7 files changed

Lines changed: 227 additions & 9 deletions

File tree

‎third_party/xllm_ops‎

Submodule xllm_ops updated from 96a5909 to d3a0acf

‎xllm/core/kernels/npu/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ cc_library(
2020
npu_moe_init_routing_v2.cpp
2121
npu_moe_token_unpermute.cpp
2222
rope.cpp
23+
npu_recurrent_gated_delta_rule.cpp
2324
DEPS
2425
:torch_npu_kernels
2526
:tilelang_kernels

‎xllm/core/kernels/npu/npu_ops_api.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,4 +141,16 @@ std::pair<torch::Tensor, torch::Tensor> apply_npu_partial_rotary_embedding(
141141
const torch::Tensor& cos_sin_cache,
142142
bool is_neox_style);
143143

144+
torch::Tensor npu_recurrent_gated_delta_rule(
145+
const torch::Tensor& query,
146+
const torch::Tensor& key,
147+
const torch::Tensor& value,
148+
torch::Tensor& state,
149+
const std::optional<torch::Tensor>& beta,
150+
const std::optional<double> scale,
151+
const std::optional<torch::Tensor>& actual_seq_lengths,
152+
const std::optional<torch::Tensor>& ssm_state_indices,
153+
const std::optional<torch::Tensor>& num_accepted_tokens,
154+
const std::optional<torch::Tensor>& g,
155+
const std::optional<torch::Tensor>& gk);
144156
} // namespace xllm::kernel::npu

‎xllm/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
https://github.com/jd-opensource/xllm/blob/main/LICENSE
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include <c10/core/Device.h>
17+
#include <glog/logging.h>
18+
#include <torch/torch.h>
19+
#include <torch_npu/csrc/libs/init_npu.h>
20+
#include <torch_npu/torch_npu.h>
21+
22+
#include <nlohmann/json.hpp>
23+
#ifdef TORCH_HIGHER_THAN_PTA6
24+
#include <torch_npu/csrc/framework/OpCommand.h>
25+
#else
26+
#include <torch_npu/csrc/aten/NPUNativeFunctions.h>
27+
#include <torch_npu/csrc/framework/utils/OpPreparation.h>
28+
#endif
29+
30+
#include "acl/acl.h"
31+
#include "aclnn_recurrent_gated_delta_rule.h"
32+
#include "core/common/macros.h"
33+
#include "core/kernels/npu/utils.h"
34+
#include "npu_ops_api.h"
35+
36+
namespace xllm::kernel::npu {
37+
38+
torch::Tensor npu_recurrent_gated_delta_rule(
39+
const torch::Tensor& query,
40+
const torch::Tensor& key,
41+
const torch::Tensor& value,
42+
torch::Tensor& state,
43+
const std::optional<torch::Tensor>& beta,
44+
const std::optional<double> scale,
45+
const std::optional<torch::Tensor>& actual_seq_lengths,
46+
const std::optional<torch::Tensor>& ssm_state_indices,
47+
const std::optional<torch::Tensor>& num_accepted_tokens,
48+
const std::optional<torch::Tensor>& g,
49+
const std::optional<torch::Tensor>& gk) {
50+
check_tensor(query, "query", "recurrent_gated_delta_rule");
51+
check_tensor(key, "key", "recurrent_gated_delta_rule");
52+
check_tensor(value, "value", "recurrent_gated_delta_rule");
53+
check_tensor(state, "state", "recurrent_gated_delta_rule");
54+
55+
aclTensor* query_ids = nullptr;
56+
aclTensor* key_ids = nullptr;
57+
aclTensor* value_ids = nullptr;
58+
aclTensor* state_ids = nullptr;
59+
aclTensor* beta_ids = nullptr;
60+
aclTensor* actual_seq_lengths_ids = nullptr;
61+
aclTensor* ssm_state_indices_ids = nullptr;
62+
aclTensor* num_accepted_tokens_ids = nullptr;
63+
aclTensor* g_ids = nullptr;
64+
aclTensor* gk_ids = nullptr;
65+
aclTensor* out_ids = nullptr;
66+
67+
int32_t device_id = query.device().index();
68+
aclrtStream stream = c10_npu::getCurrentNPUStream(device_id).stream();
69+
70+
create_acltensor(&query_ids, query);
71+
create_acltensor(&key_ids, key);
72+
create_acltensor(&value_ids, value);
73+
create_acltensor(&state_ids, state);
74+
75+
if (beta.has_value() && beta.value().defined()) {
76+
create_acltensor(&beta_ids, beta.value());
77+
}
78+
if (actual_seq_lengths.has_value() && actual_seq_lengths.value().defined()) {
79+
create_acltensor(&actual_seq_lengths_ids, actual_seq_lengths.value());
80+
}
81+
if (ssm_state_indices.has_value() && ssm_state_indices.value().defined()) {
82+
create_acltensor(&ssm_state_indices_ids, ssm_state_indices.value());
83+
}
84+
if (num_accepted_tokens.has_value() &&
85+
num_accepted_tokens.value().defined()) {
86+
create_acltensor(&num_accepted_tokens_ids, num_accepted_tokens.value());
87+
}
88+
if (g.has_value() && g.value().defined()) {
89+
create_acltensor(&g_ids, g.value());
90+
}
91+
if (gk.has_value() && gk.value().defined()) {
92+
create_acltensor(&gk_ids, gk.value());
93+
}
94+
95+
at::Tensor out_result = at::empty_like(value);
96+
create_acltensor(&out_ids, out_result);
97+
98+
float scale_value = static_cast<float>(scale.value());
99+
100+
uint64_t workspace_size = 0;
101+
aclOpExecutor* executor = nullptr;
102+
103+
CHECK_ACL_SUCCESS(
104+
aclnnRecurrentGatedDeltaRuleGetWorkspaceSize(query_ids,
105+
key_ids,
106+
value_ids,
107+
beta_ids,
108+
state_ids,
109+
actual_seq_lengths_ids,
110+
ssm_state_indices_ids,
111+
g_ids,
112+
gk_ids,
113+
num_accepted_tokens_ids,
114+
scale_value,
115+
out_ids,
116+
&workspace_size,
117+
&executor),
118+
"recurrent_gated_delta_rule: failed to get workspace size");
119+
120+
void* workspace_addr = nullptr;
121+
if (workspace_size > 0) {
122+
CHECK_ACL_SUCCESS(
123+
aclrtMalloc(&workspace_addr, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST),
124+
"recurrent_gated_delta_rule: failed to allocate workspace");
125+
}
126+
127+
CHECK_ACL_SUCCESS(aclnnRecurrentGatedDeltaRule(
128+
workspace_addr, workspace_size, executor, stream),
129+
"recurrent_gated_delta_rule: failed to perform recurrent "
130+
"gated delta rule");
131+
132+
aclDestroyTensor(query_ids);
133+
aclDestroyTensor(key_ids);
134+
aclDestroyTensor(value_ids);
135+
aclDestroyTensor(state_ids);
136+
aclDestroyTensor(out_ids);
137+
138+
if (beta_ids != nullptr) {
139+
aclDestroyTensor(beta_ids);
140+
}
141+
if (actual_seq_lengths_ids != nullptr) {
142+
aclDestroyTensor(actual_seq_lengths_ids);
143+
}
144+
if (ssm_state_indices_ids != nullptr) {
145+
aclDestroyTensor(ssm_state_indices_ids);
146+
}
147+
if (num_accepted_tokens_ids != nullptr) {
148+
aclDestroyTensor(num_accepted_tokens_ids);
149+
}
150+
if (g_ids != nullptr) {
151+
aclDestroyTensor(g_ids);
152+
}
153+
if (gk_ids != nullptr) {
154+
aclDestroyTensor(gk_ids);
155+
}
156+
157+
if (workspace_size > 0) {
158+
CHECK_ACL_SUCCESS(aclrtFree(workspace_addr),
159+
"recurrent_gated_delta_rule: failed to free workspace");
160+
}
161+
162+
return out_result;
163+
}
164+
165+
} // namespace xllm::kernel::npu

‎xllm/core/kernels/ops_api.cpp‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1018,4 +1018,33 @@ std::pair<torch::Tensor, torch::Tensor> chunk_gated_delta_rule(
10181018
NOT_IMPLEMENTED();
10191019
#endif
10201020
}
1021+
1022+
torch::Tensor recurrent_gated_delta_rule(
1023+
const torch::Tensor& query,
1024+
const torch::Tensor& key,
1025+
const torch::Tensor& value,
1026+
torch::Tensor& state,
1027+
const std::optional<torch::Tensor>& beta,
1028+
const std::optional<double> scale,
1029+
const std::optional<torch::Tensor>& actual_seq_lengths,
1030+
const std::optional<torch::Tensor>& ssm_state_indices,
1031+
const std::optional<torch::Tensor>& num_accepted_tokens,
1032+
const std::optional<torch::Tensor>& g,
1033+
const std::optional<torch::Tensor>& gk) {
1034+
#if defined(USE_NPU)
1035+
return npu::npu_recurrent_gated_delta_rule(query,
1036+
key,
1037+
value,
1038+
state,
1039+
beta,
1040+
scale,
1041+
actual_seq_lengths,
1042+
ssm_state_indices,
1043+
num_accepted_tokens,
1044+
g,
1045+
gk);
1046+
#else
1047+
NOT_IMPLEMENTED();
1048+
#endif
1049+
}
10211050
} // namespace xllm::kernel

‎xllm/core/kernels/ops_api.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,4 +148,16 @@ void gemma_rms_norm(GemmaRMSNormParams& params);
148148
std::pair<torch::Tensor, torch::Tensor> chunk_gated_delta_rule(
149149
ChunkGatedDeltaRuleParams& params);
150150

151+
torch::Tensor recurrent_gated_delta_rule(
152+
const torch::Tensor& query,
153+
const torch::Tensor& key,
154+
const torch::Tensor& value,
155+
torch::Tensor& state,
156+
const std::optional<torch::Tensor>& beta,
157+
const std::optional<double> scale,
158+
const std::optional<torch::Tensor>& actual_seq_lengths,
159+
const std::optional<torch::Tensor>& ssm_state_indices,
160+
const std::optional<torch::Tensor>& num_accepted_tokens,
161+
const std::optional<torch::Tensor>& g,
162+
const std::optional<torch::Tensor>& gk);
151163
} // namespace xllm::kernel

‎xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -434,8 +434,8 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
434434
chunk_gated_delta_params.beta = beta;
435435
// Get initial state from ssm_cache for sequences with previous state
436436
// Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
437-
torch::Tensor initial_state_tensor = torch::index_select(
438-
ssm_cache, 0, input_params.block_tables.select(1, 0));
437+
torch::Tensor initial_state_tensor =
438+
torch::index_select(ssm_cache, 0, linear_state_indices);
439439
// Todo: chunked-prefill/prefix-cache use initial_state
440440
initial_state_tensor.fill_(0.0);
441441
chunk_gated_delta_params.initial_state = initial_state_tensor;
@@ -447,15 +447,14 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
447447
std::tie(core_attn_out, last_recurrent_state) =
448448
xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
449449
ssm_cache.index_put_(
450-
{input_params.block_tables.select(1, 0)},
450+
{linear_state_indices},
451451
last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
452452
} else {
453453
processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
454454
processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
455-
torch::Tensor ssm_state_indices =
456-
attn_metadata.block_table.select(1, 0).contiguous();
457-
auto zero = torch::zeros({1},attn_metadata.q_seq_lens.options());
458-
torch::Tensor actual_seq_lengths = torch::cat({zero, attn_metadata.q_seq_lens}, 0);
455+
auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
456+
torch::Tensor actual_seq_lengths =
457+
torch::cat({zero, attn_metadata.q_seq_lens}, 0);
459458
double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
460459
core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
461460
processed_q.reshape(
@@ -468,7 +467,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
468467
beta.squeeze(0).contiguous(),
469468
scale,
470469
actual_seq_lengths,
471-
ssm_state_indices,
470+
linear_state_indices,
472471
c10::nullopt,
473472
g.squeeze(0).contiguous(),
474473
c10::nullopt)

0 commit comments

Comments
 (0)