Skip to content

[2026春季] T2-1-1/T2-1-2: 吃kfc #494

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
GodKickMyAss wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[2026春季] T2-1-1/T2-1-2: 吃kfc #494

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
243f0f3
perf: optimize inference serving and scheduling
GodKickMyAss Jul 11, 2026
2692304
feat: add ERNIE 4.5 VL support
GodKickMyAss Jul 11, 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
87 changes: 69 additions & 18 deletions csrc/engine/compiler/paged_compiler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,23 @@
#include "paged_compiler.hpp"
#include "../../global_state/global_state.hpp"
#include "../../utils.hpp"
#include "../workspace/workspace_context.hpp"

#include <string>

namespace infinilm::engine {

namespace {
bool same_shape_and_dtype(const infinicore::Tensor &lhs, const infinicore::Tensor &rhs) {
return lhs->shape() == rhs->shape() && lhs->dtype() == rhs->dtype();
}

bool is_mrope_decode_position_ids(const infinicore::Tensor &position_ids, size_t batch_size) {
const auto shape = position_ids->shape();
return shape.size() == 2 && shape[0] == batch_size && shape[1] == 3;
}
} // namespace

PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
: GraphCompiler(model, barrier) {
for (size_t b = 1; b < 64; ++b) {
Expand All @@ -25,14 +39,19 @@ void PagedCompiler::compile() {
size_t nblocks = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->num_blocks();
size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
compiled_map_decode_.clear();
compiled_map_decode_mrope_.clear();
block_tables_holder_ = infinicore::Tensor::empty(
{nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice());
set_zeros(block_tables_holder_);

auto make_decode_input = [&](size_t b) {
auto make_decode_input = [&](size_t b, bool mrope_position_ids) {
InfinilmModel::Input input;
input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice());
input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I32, infinicore::context::getDevice());
if (mrope_position_ids) {
input.position_ids = infinicore::Tensor::empty({b, 3}, infinicore::DataType::I64, infinicore::context::getDevice());
} else {
input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
}
input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice());
set_zeros(input.input_ids.value());
set_zeros(input.position_ids.value());
Expand Down Expand Up @@ -66,8 +85,11 @@ void PagedCompiler::compile() {

{
const size_t warmup_batch_size = std::min(max_batch_size, static_cast<size_t>(64));
auto input = make_decode_input(warmup_batch_size);
model_->forward(input);
auto input = make_decode_input(warmup_batch_size, false);
{
WorkspaceForwardGuard forward_guard(maybe_current_workspace());
model_->forward(input);
}
infinicore::context::syncStream();
// Warmup runs the eager Marlin path and may leave per-layer lock
// workspaces dirty. Reset before CUDA graph capture so capture
Expand All @@ -76,27 +98,49 @@ void PagedCompiler::compile() {
infinicore::context::syncStream();
}

for (size_t b : decode_batch_sizes_) {
auto input = make_decode_input(b);
auto capture_decode_graph = [&](size_t b,
bool mrope_position_ids,
std::unordered_map<size_t, CompiledResult> &compiled_map,
const std::string &scope_name) {
auto input = make_decode_input(b, mrope_position_ids);

barrier_->wait();
(void)model_->forward(input);
{
WorkspaceForwardGuard forward_guard(maybe_current_workspace());
(void)model_->forward(input);
}
infinicore::context::syncStream();
// Capture must not start with stale Marlin locks from previous
// warmup/capture attempts. This reset is intentionally outside
// graph capture; the current implementation still pays a memset
// before every graph replay in get_compiled().
model_->reset_runtime_state();
infinicore::context::syncStream();
barrier_->wait();
infinicore::context::startGraphRecording();
WorkspaceCollectiveScopeGuard collective_scope_guard(
maybe_current_workspace(),
scope_name + std::to_string(b));
WorkspaceForwardGuard forward_guard(maybe_current_workspace());
// Capture runtime workspace resets so graph replay does not need
// an extra host-launched reset phase before every decode token.
model_->reset_runtime_state();
auto output = model_->forward(input);
auto graph = infinicore::context::stopGraphRecording();
barrier_->wait();

auto shared_output = std::shared_ptr<InfinilmModel::Output>(
new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)});

compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
compiled_map[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
};

const bool compile_mrope_decode = model_->supports_mrope_position_ids();
for (size_t b : decode_batch_sizes_) {
capture_decode_graph(b, false, compiled_map_decode_, "paged_decode.");
if (compile_mrope_decode) {
capture_decode_graph(b, true, compiled_map_decode_mrope_, "paged_decode_mrope.");
}
}
}
}
Expand All @@ -110,12 +154,23 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
if (batch_size != input.input_ids.value()->size(1)) {
return {nullptr, nullptr};
} else {
auto result = compiled_map_decode_.find(batch_size);
if (result == compiled_map_decode_.end()) {
const bool mrope_position_ids = is_mrope_decode_position_ids(input.position_ids.value(), batch_size);
auto &compiled_map = mrope_position_ids ? compiled_map_decode_mrope_ : compiled_map_decode_;
auto result = compiled_map.find(batch_size);
if (result == compiled_map.end()) {
return {nullptr, nullptr};
}
auto &graph_input = result->second.input;

if (!same_shape_and_dtype(graph_input.input_ids.value(), input.input_ids.value()) ||
!same_shape_and_dtype(graph_input.position_ids.value(), input.position_ids.value()) ||
!same_shape_and_dtype(graph_input.total_sequence_lengths.value(), input.total_sequence_lengths.value()) ||
!same_shape_and_dtype(graph_input.input_offsets.value(), input.input_offsets.value()) ||
!same_shape_and_dtype(graph_input.cu_seqlens.value(), input.cu_seqlens.value()) ||
!same_shape_and_dtype(graph_input.slot_mapping.value(), input.slot_mapping.value())) {
return {nullptr, nullptr};
}

graph_input.input_ids.value()->copy_from(input.input_ids.value());
graph_input.position_ids.value()->copy_from(input.position_ids.value());
graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
Expand All @@ -132,16 +187,12 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
// runtime logical region. Avoid clearing the full preallocated
// holder on every decode token.
auto &graph_block_tables = graph_input.block_tables.value();
if (graph_block_tables->dtype() != input.block_tables.value()->dtype()) {
return {nullptr, nullptr};
}
set_minus_one_device_async(graph_block_tables);
graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value());
graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value());
// CUDA graph replay reuses the same per-layer Marlin workspaces.
// The graph itself does not contain a workspace reset, so enqueue
// one on the same stream before launch. This is correct but costs
// decode latency; the intended follow-up is a reusable global
// zero workspace/lock buffer shared by all Marlin layers.
model_->reset_runtime_state();

auto graph = std::get<0>(result->second.compiled);
auto shared_output = std::shared_ptr<InfinilmModel::Output>(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()});

Expand Down
5 changes: 5 additions & 0 deletions csrc/engine/compiler/paged_compiler.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,5 +27,10 @@ class PagedCompiler : public GraphCompiler {
size_t, // num_requests
CompiledResult>
compiled_map_decode_;

std::unordered_map<
size_t, // num_requests
CompiledResult>
compiled_map_decode_mrope_;
};
} // namespace infinilm::engine
16 changes: 14 additions & 2 deletions csrc/engine/compiler/static_batching_compiler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include "static_batching_compiler.hpp"
#include "../../cache/cache.hpp"
#include "../../global_state/global_state.hpp"
#include "../workspace/workspace_context.hpp"

namespace infinilm::engine {
StaticBatchingCompiler::StaticBatchingCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
Expand All @@ -11,7 +12,7 @@ void StaticBatchingCompiler::compile() {
if (model_->get_cache_config() != nullptr && dynamic_cast<const cache::StaticKVCacheConfig *>(model_->get_cache_config())) {
size_t b = dynamic_cast<const cache::StaticKVCacheConfig *>(model_->get_cache_config())->max_batch_size();
InfinilmModel::Input input;
input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice());
input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I32, infinicore::context::getDevice());
input.position_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice());
input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
Expand All @@ -29,10 +30,21 @@ void StaticBatchingCompiler::compile() {
};

barrier_->wait();
(void)model_->forward(input);
{
WorkspaceForwardGuard forward_guard(maybe_current_workspace());
(void)model_->forward(input);
}
infinicore::context::syncStream();

barrier_->wait();
infinicore::context::startGraphRecording();
WorkspaceCollectiveScopeGuard collective_scope_guard(
maybe_current_workspace(),
"static_batching." + std::to_string(b));
WorkspaceForwardGuard forward_guard(maybe_current_workspace());
// Capture runtime workspace resets so graph replay does not need
// an extra host-launched reset phase before every decode token.
model_->reset_runtime_state();
auto output = model_->forward(input);
auto graph = infinicore::context::stopGraphRecording();
barrier_->wait();
Expand Down
Loading
Loading