Skip to content

server: emit GPU ttft/generation/decode_tps in stream usage #570

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
ai-hpc wants to merge 1 commit into
base: main
Choose a base branch
from
Open

server: emit GPU ttft/generation/decode_tps in stream usage #570

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 4 additions & 0 deletions runtime/include/sparkinfer/inference_engine.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,10 @@ class ContinuousBatchEngine {
struct Result {
std::vector<int> tokens;
std::string error;
// GPU-side timings (exclude SSE/on_token backpressure).
double ttft_ms = -1.0;
double generation_ms = -1.0;
double decode_tps = -1.0;
};

ContinuousBatchEngine(Qwen35Model* model, KVCacheManager* kv,
Expand Down
35 changes: 34 additions & 1 deletion runtime/src/inference_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,14 @@ struct ContinuousBatchEngine::Job {
std::string error;
std::function<void(int)> on_token;
bool done = false;
std::chrono::steady_clock::time_point t_submit{};
std::chrono::steady_clock::time_point t_first{};
bool saw_first_tok = false;
double decode_gpu_ms = 0.0;
int decode_forwards = 0;
double ttft_ms = -1.0;
double generation_ms = -1.0;
double decode_tps = -1.0;
};

ContinuousBatchEngine::ContinuousBatchEngine(Qwen35Model* model, KVCacheManager* kv,
Expand Down Expand Up @@ -115,6 +123,7 @@ uint64_t ContinuousBatchEngine::submit_locked(Job job, const std::function<void(
job.seq_id = seq_id;
job.on_token = on_token;
job.prefill_pos = job.req.prefill_start;
job.t_submit = std::chrono::steady_clock::now();
auto ptr = std::make_unique<Job>(std::move(job));
const uint64_t rid = ptr->request_id;
jobs_[rid] = std::move(ptr);
Expand All @@ -130,7 +139,12 @@ ContinuousBatchEngine::Result ContinuousBatchEngine::wait_locked(uint64_t reques
});
auto it = jobs_.find(request_id);
if (it == jobs_.end()) return Result{{}, "request not found"};
Result out{it->second->output, it->second->error};
Result out;
out.tokens = it->second->output;
out.error = it->second->error;
out.ttft_ms = it->second->ttft_ms;
out.generation_ms = it->second->generation_ms;
out.decode_tps = it->second->decode_tps;
jobs_.erase(it);
return out;
}
Expand Down Expand Up @@ -234,11 +248,26 @@ bool ContinuousBatchEngine::step_job(Job& job) {
return true;
}

// Timestamp before on_token so SSE/network backpressure never enters GPU metrics.
const auto t_emit = std::chrono::steady_clock::now();
if (!job.saw_first_tok) {
job.t_first = t_emit;
job.saw_first_tok = true;
job.ttft_ms = std::chrono::duration<double, std::milli>(job.t_first - job.t_submit).count();
}
job.output.push_back(job.next_token);
if (job.on_token) job.on_token(job.next_token);
job.decode_emitted++;

if (job.next_token == cfg.eos_id || job.decode_emitted >= job.req.max_new_tokens) {
const auto t_end = std::chrono::steady_clock::now();
job.generation_ms = std::chrono::duration<double, std::milli>(t_end - job.t_submit).count();
if (job.decode_forwards > 0 && job.decode_gpu_ms > 0.0) {
job.decode_tps = (double)job.decode_forwards * 1000.0 / job.decode_gpu_ms;
} else if (job.saw_first_tok && job.generation_ms > job.ttft_ms && job.decode_emitted > 0) {
const double decode_ms = std::max(job.generation_ms - job.ttft_ms, 1.0);
job.decode_tps = (double)job.decode_emitted * 1000.0 / decode_ms;
}
job.done = true;
if (job.seq_id != 0) model_->close_session(job.seq_id);
else kv_->free(job.seq_id);
Expand All @@ -247,7 +276,11 @@ bool ContinuousBatchEngine::step_job(Job& job) {
}

const int prompt_len = (int)job.req.prompt.size();
const auto t0 = std::chrono::steady_clock::now();
job.next_token = model_->forward_token(job.next_token, prompt_len + job.decode_emitted - 1, true);
job.decode_gpu_ms += std::chrono::duration<double, std::milli>(
std::chrono::steady_clock::now() - t0).count();
job.decode_forwards++;
return false;
}

Expand Down
2 changes: 1 addition & 1 deletion server/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ export SPARKINFER_ROOT="$(pwd)"
| `GET /v1/models` | OpenAI model list |
| `GET /v1/info` | Model limits (`max_context`, `max_output_tokens`) |
| `POST /v1/tokenize` | Token count for a chat request body |
| `POST /v1/chat/completions` | Chat (JSON `messages`, optional `stream`, `enable_thinking`). Responses include OpenAI `usage` (`prompt_tokens`, `completion_tokens`, `total_tokens`). Streaming sends a final chunk with `choices:[]` + `usage` before `[DONE]`. |
| `POST /v1/chat/completions` | Chat (JSON `messages`, optional `stream`, `enable_thinking`). Responses include OpenAI `usage` (`prompt_tokens`, `completion_tokens`, `total_tokens`) plus optional GPU timing fields (`ttft_ms`, `generation_ms`, `decode_tps`). Streaming sends a final chunk with `choices:[]` + `usage` before `[DONE]`. |

### RTX PRO 6000 deploy (32k / 4k)

Expand Down
8 changes: 8 additions & 0 deletions server/include/model_engine.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@

namespace sparkinfer_server {

struct CompletionTiming {
double ttft_ms = -1.0;
double generation_ms = -1.0;
double decode_tps = -1.0;
};

// Thread-safe wrapper around sparkinfer::Qwen35Model + GGUF load.
class ModelEngine {
public:
Expand Down Expand Up @@ -39,12 +45,14 @@ class ModelEngine {
const std::function<void(int)>& on_token);

const std::string& last_error() const;
const CompletionTiming& last_timing() const;

private:
struct Impl;
std::unique_ptr<Impl> impl_;
mutable std::mutex mu_;
std::string last_error_;
CompletionTiming last_timing_;
};

} // namespace sparkinfer_server
7 changes: 7 additions & 0 deletions server/src/model_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,10 @@ int ModelEngine::prefix_token_len() const {
return (int)impl_->prefix_tokens.size();
}

const CompletionTiming& ModelEngine::last_timing() const {
return last_timing_;
}

const std::string& ModelEngine::last_error() const {
std::lock_guard<std::mutex> lock(mu_);
return last_error_;
Expand Down Expand Up @@ -230,6 +234,9 @@ std::vector<int> ModelEngine::complete_streaming(const std::vector<int>& prompt_
auto result = impl_->batch_engine->complete_streaming(req, on_token);

std::lock_guard<std::mutex> lock(mu_);
last_timing_.ttft_ms = result.ttft_ms;
last_timing_.generation_ms = result.generation_ms;
last_timing_.decode_tps = result.decode_tps;
if (!result.error.empty()) {
last_error_ = result.error;
fprintf(stderr, "[sparkinfer-server] %s\n", last_error_.c_str());
Expand Down
38 changes: 35 additions & 3 deletions server/src/sparkinfer_server.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,11 +90,19 @@ std::string random_id() {
return ss.str();
}

std::string usage_json(int prompt_tokens, int completion_tokens) {
std::string usage_json(int prompt_tokens, int completion_tokens, double ttft_ms = -1.0,
double generation_ms = -1.0, double decode_tps = -1.0) {
std::ostringstream o;
const int total = prompt_tokens + completion_tokens;
o << "\"usage\":{\"prompt_tokens\":" << prompt_tokens << ",\"completion_tokens\":" << completion_tokens
<< ",\"total_tokens\":" << total << "}";
<< ",\"total_tokens\":" << total;
if (ttft_ms >= 0.0)
o << ",\"ttft_ms\":" << std::fixed << std::setprecision(3) << ttft_ms;
if (generation_ms >= 0.0)
o << ",\"generation_ms\":" << std::fixed << std::setprecision(3) << generation_ms;
if (decode_tps >= 0.0)
o << ",\"decode_tps\":" << std::fixed << std::setprecision(2) << decode_tps;
o << "}";
return o.str();
}

Expand Down Expand Up @@ -313,13 +321,21 @@ int main(int argc, char** argv) {
std::vector<int> stream_ids;
stream_ids.reserve((size_t)max_tokens);
sparkinfer_server::ThinkingStreamSplitter splitter(enable_thinking);
const auto wall_start = std::chrono::steady_clock::now();
std::chrono::steady_clock::time_point first_tok_time;
bool saw_first_tok = false;
auto on_tok = [&](int tid) {
if (!saw_first_tok) {
first_tok_time = std::chrono::steady_clock::now();
saw_first_tok = true;
}
std::string piece = g_tokenizer.decode_delta(stream_ids, tid);
const auto delta = splitter.feed(piece);
write_stream_delta(sink, cid, created, "reasoning_content", delta.reasoning_content);
write_stream_delta(sink, cid, created, "content", delta.content);
};
engine.complete_streaming(prompt_ids, max_tokens, on_tok);
const auto wall_end = std::chrono::steady_clock::now();
sparkinfer_server::ThinkingStreamSplitter::Delta flush;
splitter.finish(flush);
write_stream_delta(sink, cid, created, "reasoning_content", flush.reasoning_content);
Expand All @@ -332,11 +348,27 @@ int main(int argc, char** argv) {
}
const int prompt_tokens = (int)prompt_ids.size();
const int completion_tokens = (int)stream_ids.size();
const auto& timing = engine.last_timing();
double ttft_ms = timing.ttft_ms;
double generation_ms = timing.generation_ms;
double decode_tps = timing.decode_tps;
if (generation_ms < 0.0) {
generation_ms = std::chrono::duration<double, std::milli>(wall_end - wall_start).count();
}
if (ttft_ms < 0.0 && saw_first_tok) {
ttft_ms = std::chrono::duration<double, std::milli>(first_tok_time - wall_start).count();
}
if (decode_tps < 0.0 && completion_tokens > 0 && generation_ms > 0.0) {
const double decode_ms =
(ttft_ms >= 0.0) ? std::max(generation_ms - ttft_ms, 1.0) : generation_ms;
decode_tps = (double)completion_tokens * 1000.0 / decode_ms;
}
std::ostringstream usage_chunk;
usage_chunk << "data: {\"id\":\"" << cid << "\",\"object\":\"chat.completion.chunk\","
<< "\"created\":" << created << ",\"model\":\"" << g_model_name << "\","
<< "\"choices\":[],"
<< usage_json(prompt_tokens, completion_tokens) << "}\n\n";
<< usage_json(prompt_tokens, completion_tokens, ttft_ms, generation_ms, decode_tps)
<< "}\n\n";
sink.write(usage_chunk.str().c_str(), (size_t)usage_chunk.str().size());
std::string tail =
"data: {\"id\":\"" + cid +
Expand Down
Loading