Skip to content
Open
Changes from all commits
Commits
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
38 changes: 38 additions & 0 deletions crates/synapse-engine-cuda/src/port/cuda_qwen3.cu
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include "cuda_family_common.cuh"

#include <cfloat>
#include <list>
#include <cstdio>

using namespace synapse_cuda_family;
Expand Down Expand Up @@ -246,13 +247,39 @@ struct QwenContext {
DeviceAllocation<half> embeddings;
bool embeddings_loaded = false;
std::unordered_map<std::string, std::unique_ptr<ShapePlan>> plans;
// Bounded shape-plan retention. Every forward returns only after
// cudaStreamSynchronize on this context's single stream, so the stream is
// idle at the FFI boundary and a dropped plan frees buffers no kernel is
// still referencing. Retain two warm embedding shapes; misses evict
// before allocation so a third full arena cannot raise the cache peak.
static constexpr size_t max_plans = 2;
std::list<std::string> plan_lru;

// Promotes `key` to most-recently-used, then trims the
// least-recently-used plan until at most max_plans remain. The guard
// makes the invariant explicit: forward only touches a key it just
// inserted or found, so the just-used plan is at the front and is never
// the eviction victim. Erasing a map entry destroys the unique_ptr,
// running ShapePlan::~ShapePlan: graph_exec, then graph, then every
// DeviceAllocation (cudaFree) via their own destructors.
void touch_plan(const std::string &key) {
if (plans.find(key) == plans.end()) return;
plan_lru.remove(key);
plan_lru.push_front(key);
while (plan_lru.size() > max_plans) {
std::string victim = plan_lru.back();
plan_lru.pop_back();
plans.erase(victim);
}
}

explicit QwenContext(bool graphs) : graphs_enabled(graphs) {
FAMILY_CUDA_CHECK(cudaFree(nullptr));
FAMILY_CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
FAMILY_CUBLAS_CHECK(cublasLtCreate(&lt));
}
~QwenContext() {
plan_lru.clear();
plans.clear();
if (lt) cublasLtDestroy(lt);
if (stream) cudaStreamDestroy(stream);
Expand Down Expand Up @@ -523,11 +550,22 @@ int32_t synapse_cuda_qwen3_forward(
std::string key = shape_key(batch, seq);
auto found = context->plans.find(key);
if (found == context->plans.end()) {
if (context->plans.size() >= QwenContext::max_plans) {
FAMILY_CUDA_CHECK(cudaStreamSynchronize(context->stream));
const std::string victim = context->plan_lru.back();
context->plan_lru.pop_back();
context->plans.erase(victim);
}
auto plan = std::make_unique<ShapePlan>(context, batch, seq, hidden, query_heads, kv_heads, head_dim, intermediate, layer_count, epsilon, rope_theta);
plan->initialize_and_verify(token_ids, attention_mask);
found = context->plans.emplace(key, std::move(plan)).first;
}
// Retain at most max_plans shape plans. Order matters: the touch runs
// after run() has returned from cudaStreamSynchronize, so the stream
// is idle, the just-executed plan is most-recently-used, and a
// victim's buffers are freed while no kernel references them.
found->second->run(token_ids, attention_mask, output);
context->touch_plan(key);
return 0;
} catch (const std::exception &error) {
synapse_cuda_set_last_error(error.what());
Expand Down
Loading