Skip to content

Commit 768ebea

Browse files
whjthuwhjthu
committed
refactor: share attention gating and graph output handling
1 parent e2b781c commit 768ebea

7 files changed

Lines changed: 66 additions & 124 deletions

File tree

‎csrc/engine/compiler/graph_compiler.hpp‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include "../../models/infinilm_model.hpp"
44
#include "../../global_state/forward_context.hpp"
5+
#include "../../execution/graph/graph.hpp"
56
#include "runtime/runtime.hpp"
67
#include "../rank_barrier.hpp"
78

@@ -77,6 +78,34 @@ class GraphCompiler {
7778
virtual Compiled get_compiled(const InfinilmModel::Input &input) = 0;
7879

7980
protected:
81+
std::shared_ptr<InfinilmModel::Output> snapshot_output(const InfinilmModel::Output &output) {
82+
infinilm::runtime::Tensor graph_hidden_states;
83+
if (output.hidden_states) {
84+
graph_hidden_states = infinilm::execution::graph::GraphTensor(
85+
runtime_,
86+
output.hidden_states,
87+
infinilm::execution::graph::GraphTensor::SnapshotPolicy::kBlob);
88+
}
89+
return std::shared_ptr<InfinilmModel::Output>(
90+
new InfinilmModel::Output{
91+
infinilm::execution::graph::GraphTensor(
92+
runtime_,
93+
output.logits,
94+
infinilm::execution::graph::GraphTensor::SnapshotPolicy::kBlob),
95+
graph_hidden_states});
96+
}
97+
98+
std::shared_ptr<InfinilmModel::Output> resume_output(const InfinilmModel::Output &output) {
99+
infinilm::runtime::Tensor hidden_states;
100+
if (output.hidden_states) {
101+
hidden_states = output.hidden_states->resume_from_blob_(runtime_);
102+
}
103+
return std::shared_ptr<InfinilmModel::Output>(
104+
new InfinilmModel::Output{
105+
output.logits->resume_from_blob_(runtime_),
106+
hidden_states});
107+
}
108+
80109
std::shared_ptr<InfinilmModel> model_;
81110
RankBarrier *barrier_;
82111
infinilm::runtime::Runtime &runtime_;

‎csrc/engine/compiler/paged_compiler.cpp‎

Lines changed: 2 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -622,20 +622,7 @@ void PagedCompiler::compile() {
622622
InfinilmModel::Input input,
623623
std::shared_ptr<infinilm::execution::graph::Graph> graph,
624624
const InfinilmModel::Output &output) {
625-
infinilm::runtime::Tensor graph_hidden_states;
626-
if (output.hidden_states) {
627-
graph_hidden_states = infinilm::execution::graph::GraphTensor(
628-
runtime_,
629-
output.hidden_states,
630-
infinilm::execution::graph::GraphTensor::SnapshotPolicy::kBlob);
631-
}
632-
auto shared_output = std::shared_ptr<InfinilmModel::Output>(
633-
new InfinilmModel::Output{
634-
infinilm::execution::graph::GraphTensor(
635-
runtime_,
636-
output.logits,
637-
infinilm::execution::graph::GraphTensor::SnapshotPolicy::kBlob),
638-
graph_hidden_states});
625+
auto shared_output = snapshot_output(output);
639626

640627
return CompiledResult{
641628
std::move(input),
@@ -944,14 +931,7 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
944931

945932
auto graph = std::get<0>(selected_result->compiled);
946933
const auto &compiled_output = std::get<1>(selected_result->compiled);
947-
infinilm::runtime::Tensor hidden_states;
948-
if (compiled_output->hidden_states) {
949-
hidden_states = compiled_output->hidden_states->resume_from_blob_(runtime_);
950-
}
951-
auto shared_output = std::shared_ptr<InfinilmModel::Output>(
952-
new InfinilmModel::Output{
953-
compiled_output->logits->resume_from_blob_(runtime_),
954-
hidden_states});
934+
auto shared_output = resume_output(*compiled_output);
955935

956936
return std::make_tuple(graph, shared_output);
957937
}

‎csrc/engine/compiler/static_batching_compiler.cpp‎

Lines changed: 2 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -146,20 +146,7 @@ void StaticBatchingCompiler::compile() {
146146
auto graph = recording.finish();
147147
barrier_->wait();
148148

149-
infinilm::runtime::Tensor graph_hidden_states;
150-
if (output.hidden_states) {
151-
graph_hidden_states = infinilm::execution::graph::GraphTensor(
152-
runtime_,
153-
output.hidden_states,
154-
infinilm::execution::graph::GraphTensor::SnapshotPolicy::kBlob);
155-
}
156-
auto shared_output = std::shared_ptr<InfinilmModel::Output>(
157-
new InfinilmModel::Output{
158-
infinilm::execution::graph::GraphTensor(
159-
runtime_,
160-
output.logits,
161-
infinilm::execution::graph::GraphTensor::SnapshotPolicy::kBlob),
162-
graph_hidden_states});
149+
auto shared_output = snapshot_output(output);
163150

164151
prepare_state.restore();
165152
candidate[std::make_tuple(b, 1)] = CompiledResult{
@@ -243,14 +230,7 @@ StaticBatchingCompiler::Compiled StaticBatchingCompiler::get_compiled(
243230

244231
auto graph = std::get<0>(result->second.compiled);
245232
const auto &compiled_output = std::get<1>(result->second.compiled);
246-
infinilm::runtime::Tensor hidden_states;
247-
if (compiled_output->hidden_states) {
248-
hidden_states = compiled_output->hidden_states->resume_from_blob_(runtime_);
249-
}
250-
auto shared_output = std::shared_ptr<InfinilmModel::Output>(
251-
new InfinilmModel::Output{
252-
compiled_output->logits->resume_from_blob_(runtime_),
253-
hidden_states});
233+
auto shared_output = resume_output(*compiled_output);
254234
return std::make_tuple(graph, shared_output);
255235
}
256236
} else {

‎csrc/layers/attention/attention.cpp‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
#include "../../execution/submit.hpp"
2+
#include "runtime/utils.hpp"
13
#include "../../weights/weights.hpp"
24
#include "../../global_state/forward_context.hpp"
35
#include "attention.hpp"
@@ -9,6 +11,29 @@
911

1012
namespace infinilm::layers::attention {
1113

14+
void detail::apply_output_gate(infinilm::runtime::Runtime &runtime, infinilm::runtime::Tensor &output,
15+
const infinilm::runtime::Tensor &gate, bool view_sigmoid_output) {
16+
const auto sigmoid = [&]() {
17+
auto ops_output = infinilm::runtime::Tensor::empty(runtime, gate->shape(), gate->dtype());
18+
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
19+
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
20+
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
21+
infinilm::execution::submit<infini::ops::Sigmoid>(
22+
runtime, infini::ops::Sigmoid::DefaultConfig(runtime.device().type()), gate, ops_output);
23+
INFINILM_TRACE_OP(Sigmoid, ops_output);
24+
return ops_output;
25+
};
26+
const auto ops_other = view_sigmoid_output ? sigmoid()->view(output->shape()) : sigmoid();
27+
auto ops_output = infinilm::runtime::Tensor::empty(runtime, output->shape(), output->dtype());
28+
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
29+
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
30+
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
31+
infinilm::execution::submit<infini::ops::Mul>(
32+
runtime, infini::ops::Mul::DefaultConfig(runtime.device().type()), output, ops_other, ops_output);
33+
INFINILM_TRACE_OP(Mul, ops_output);
34+
output = std::move(ops_output);
35+
}
36+
1237
Attention::Attention(infinilm::runtime::Runtime &runtime,
1338
std::shared_ptr<infinilm::config::ModelConfig> model_config,
1439
size_t layer_idx, const engine::distributed::RankInfo &rank_info,

‎csrc/layers/attention/attention.hpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,4 +68,8 @@ void init_kv_cache_quant_params(infinilm::runtime::Runtime &runtime,
6868
infinilm::weights::Weight &kv_cache_k_scale,
6969
infinilm::weights::Weight &kv_cache_v_scale,
7070
infinilm::quantization::KVQuantAlgo kv_quant_scheme);
71+
namespace detail {
72+
void apply_output_gate(infinilm::runtime::Runtime &runtime, infinilm::runtime::Tensor &output,
73+
const infinilm::runtime::Tensor &gate, bool view_sigmoid_output);
74+
}
7175
} // namespace infinilm::layers::attention

‎csrc/models/qwen3_5/qwen3_5_attention.cpp‎

Lines changed: 2 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -167,26 +167,7 @@ infinilm::runtime::Tensor Qwen35Attention::forward_static_(const infinilm::runti
167167
auto v_reshaped = v->view({batch_size, seq_len, num_key_value_heads_, head_dim_});
168168

169169
auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped, forward_context);
170-
{
171-
const auto ops_other = [&]() {
172-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, gate->shape(), gate->dtype());
173-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
174-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
175-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
176-
infinilm::execution::submit<infini::ops::Sigmoid>(
177-
runtime, infini::ops::Sigmoid::DefaultConfig(runtime.device().type()), gate, ops_output);
178-
INFINILM_TRACE_OP(Sigmoid, ops_output);
179-
return ops_output;
180-
}()->view(attn_output->shape());
181-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, attn_output->shape(), attn_output->dtype());
182-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
183-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
184-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
185-
infinilm::execution::submit<infini::ops::Mul>(
186-
runtime, infini::ops::Mul::DefaultConfig(runtime.device().type()), attn_output, ops_other, ops_output);
187-
INFINILM_TRACE_OP(Mul, ops_output);
188-
attn_output = std::move(ops_output);
189-
}
170+
infinilm::layers::attention::detail::apply_output_gate(runtime, attn_output, gate, true);
190171
return o_proj_->forward(runtime, attn_output);
191172
}
192173

@@ -209,26 +190,7 @@ infinilm::runtime::Tensor Qwen35Attention::forward_paged_(const infinilm::runtim
209190
std::tie(q_reshaped, k_reshaped) = mrope_->forward(runtime, q_reshaped, k_reshaped, position_ids);
210191

211192
auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped, forward_context);
212-
{
213-
const auto ops_other = [&]() {
214-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, gate->shape(), gate->dtype());
215-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
216-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
217-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
218-
infinilm::execution::submit<infini::ops::Sigmoid>(
219-
runtime, infini::ops::Sigmoid::DefaultConfig(runtime.device().type()), gate, ops_output);
220-
INFINILM_TRACE_OP(Sigmoid, ops_output);
221-
return ops_output;
222-
}()->view(attn_output->shape());
223-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, attn_output->shape(), attn_output->dtype());
224-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
225-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
226-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
227-
infinilm::execution::submit<infini::ops::Mul>(
228-
runtime, infini::ops::Mul::DefaultConfig(runtime.device().type()), attn_output, ops_other, ops_output);
229-
INFINILM_TRACE_OP(Mul, ops_output);
230-
attn_output = std::move(ops_output);
231-
}
193+
infinilm::layers::attention::detail::apply_output_gate(runtime, attn_output, gate, true);
232194
return o_proj_->forward(runtime, attn_output);
233195
}
234196
} // namespace infinilm::models::qwen3_5

‎csrc/models/qwen3_next/qwen3_next_attention.cpp‎

Lines changed: 2 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -157,26 +157,7 @@ infinilm::runtime::Tensor Qwen3NextAttention::forward_static_(const infinilm::ru
157157
rotary_emb_->forward(runtime, k_reshaped, pos_ids_for_rope, true);
158158

159159
auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped, forward_context);
160-
{
161-
const auto ops_other = [&]() {
162-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, gate->shape(), gate->dtype());
163-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
164-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
165-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
166-
infinilm::execution::submit<infini::ops::Sigmoid>(
167-
runtime, infini::ops::Sigmoid::DefaultConfig(runtime.device().type()), gate, ops_output);
168-
INFINILM_TRACE_OP(Sigmoid, ops_output);
169-
return ops_output;
170-
}();
171-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, attn_output->shape(), attn_output->dtype());
172-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
173-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
174-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
175-
infinilm::execution::submit<infini::ops::Mul>(
176-
runtime, infini::ops::Mul::DefaultConfig(runtime.device().type()), attn_output, ops_other, ops_output);
177-
INFINILM_TRACE_OP(Mul, ops_output);
178-
attn_output = std::move(ops_output);
179-
}
160+
infinilm::layers::attention::detail::apply_output_gate(runtime, attn_output, gate, false);
180161
return o_proj_->forward(runtime, attn_output);
181162
}
182163

@@ -224,26 +205,7 @@ infinilm::runtime::Tensor Qwen3NextAttention::forward_paged_(const infinilm::run
224205
rotary_emb_->forward(runtime, k_reshaped, pos_ids_for_rope, true);
225206

226207
auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped, forward_context);
227-
{
228-
const auto ops_other = [&]() {
229-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, gate->shape(), gate->dtype());
230-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
231-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
232-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
233-
infinilm::execution::submit<infini::ops::Sigmoid>(
234-
runtime, infini::ops::Sigmoid::DefaultConfig(runtime.device().type()), gate, ops_output);
235-
INFINILM_TRACE_OP(Sigmoid, ops_output);
236-
return ops_output;
237-
}()->view(attn_output->shape());
238-
auto ops_output = infinilm::runtime::Tensor::empty(runtime, attn_output->shape(), attn_output->dtype());
239-
INFINILM_ASSERT(runtime.device().type() != infinilm::runtime::Device::Type::kCpu
240-
&& runtime.device().type() != infinilm::runtime::Device::Type::kMars
241-
&& runtime.device().type() != infinilm::runtime::Device::Type::kThead);
242-
infinilm::execution::submit<infini::ops::Mul>(
243-
runtime, infini::ops::Mul::DefaultConfig(runtime.device().type()), attn_output, ops_other, ops_output);
244-
INFINILM_TRACE_OP(Mul, ops_output);
245-
attn_output = std::move(ops_output);
246-
}
208+
infinilm::layers::attention::detail::apply_output_gate(runtime, attn_output, gate, true);
247209
return o_proj_->forward(runtime, attn_output);
248210
}
249211

0 commit comments

Comments
 (0)