From fba42683ced7b6b394b36b01baa5f17d71c5c825 Mon Sep 17 00:00:00 2001 From: bupt-ycq Date: Fri, 24 Jul 2026 10:21:37 +0800 Subject: [PATCH 1/5] bugfix: expose prefix cache usage across APIs. - propagate cached prompt token counts through PD RPC responses - report OpenAI cached tokens and Anthropic cache reads for stream and non-stream responses --- .../http_service/anthropic_adapter_test.cpp | 27 +++- .../anthropic_stream_encoder_test.cpp | 11 +- tests/xllm_service/scheduler/CMakeLists.txt | 10 ++ .../scheduler/response_handler_usage_test.cpp | 53 ++++++++ xllm_service/common/xllm/output.h | 3 + .../http_service/anthropic_adapter.cpp | 24 +++- xllm_service/http_service/anthropic_adapter.h | 3 + .../http_service/anthropic_stream_encoder.cpp | 6 +- xllm_service/proto/xllm_rpc_service.proto | 2 + xllm_service/rpc_service/rpc_service_test.cpp | 22 ++++ xllm_service/rpc_service/service.cpp | 119 ++++++++++-------- xllm_service/rpc_service/service.h | 3 + xllm_service/scheduler/response_handler.cpp | 41 +++--- xllm_service/scheduler/response_handler.h | 2 + 14 files changed, 235 insertions(+), 91 deletions(-) create mode 100644 tests/xllm_service/scheduler/response_handler_usage_test.cpp diff --git a/tests/xllm_service/http_service/anthropic_adapter_test.cpp b/tests/xllm_service/http_service/anthropic_adapter_test.cpp index ce3606c..839aba2 100644 --- a/tests/xllm_service/http_service/anthropic_adapter_test.cpp +++ b/tests/xllm_service/http_service/anthropic_adapter_test.cpp @@ -807,6 +807,7 @@ TEST(AnthropicAdapterTest, BuildsNonStreamAnthropicJson) { usage.num_prompt_tokens = 3; usage.num_generated_tokens = 4; usage.num_total_tokens = 7; + usage.num_cached_tokens = 2; output.usage = usage; xllm::proto::AnthropicMessagesResponse response; @@ -826,11 +827,35 @@ TEST(AnthropicAdapterTest, BuildsNonStreamAnthropicJson) { ASSERT_EQ(json["content"].size(), 1); EXPECT_EQ(json["content"][0]["type"], "text"); EXPECT_EQ(json["content"][0]["text"], "answer"); - EXPECT_EQ(json["usage"]["input_tokens"], 3); + EXPECT_EQ(json["usage"]["input_tokens"], 1); EXPECT_EQ(json["usage"]["output_tokens"], 4); + EXPECT_EQ(json["usage"]["cache_read_input_tokens"], 2); + EXPECT_FALSE(json["usage"].contains("cache_creation_input_tokens")); EXPECT_FALSE(json["usage"].contains("total_tokens")); } +TEST(AnthropicAdapterTest, BuildsZeroCacheReadUsage) { + llm::RequestOutput output; + output.request_id = "anthropiccmpl-test"; + llm::Usage usage; + usage.num_prompt_tokens = 3; + output.usage = usage; + + xllm::proto::AnthropicMessagesResponse response; + auto result = fill_anthropic_resp("test-model", output, &response); + ASSERT_TRUE(result.ok) << result.error; + + std::string json_str; + std::string error; + ASSERT_TRUE(anthropic_json(response, &json_str, &error)) << error; + auto json = nlohmann::json::parse(json_str); + + ASSERT_TRUE(json["usage"].contains("cache_read_input_tokens")); + EXPECT_EQ(json["usage"]["input_tokens"], 3); + EXPECT_EQ(json["usage"]["cache_read_input_tokens"], 0); + EXPECT_FALSE(json["usage"].contains("cache_creation_input_tokens")); +} + TEST(AnthropicAdapterTest, BuildsThinkingNonStreamAnthropicJson) { llm::RequestOutput output; output.request_id = "anthropiccmpl-test"; diff --git a/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp b/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp index 5708fe0..9f55018 100644 --- a/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp +++ b/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp @@ -243,6 +243,7 @@ TEST(AnthropicStreamEncoderTest, FinishClosesBlockAndEmitsMessageDeltaStop) { llm::Usage usage; usage.num_prompt_tokens = 3; usage.num_generated_tokens = 5; + usage.num_cached_tokens = 2; final.usage = usage; std::vector done; @@ -254,8 +255,9 @@ TEST(AnthropicStreamEncoderTest, FinishClosesBlockAndEmitsMessageDeltaStop) { EXPECT_EQ(events[0]["index"], 0); EXPECT_EQ(events[1]["type"], "message_delta"); EXPECT_EQ(events[1]["delta"]["stop_reason"], "end_turn"); - EXPECT_EQ(events[1]["usage"]["input_tokens"], 3); + EXPECT_EQ(events[1]["usage"]["input_tokens"], 1); EXPECT_EQ(events[1]["usage"]["output_tokens"], 5); + EXPECT_EQ(events[1]["usage"]["cache_read_input_tokens"], 2); EXPECT_EQ(events[2]["type"], "message_stop"); } @@ -273,6 +275,9 @@ TEST(AnthropicStreamEncoderTest, FinishAfterToolUsesToolStopReason) { seq.index = 0; seq.finish_reason = "stop"; final.outputs.push_back(std::move(seq)); + llm::Usage usage; + usage.num_prompt_tokens = 3; + final.usage = usage; std::vector done; ASSERT_TRUE(encoder.finish(final, &done).ok); @@ -281,6 +286,10 @@ TEST(AnthropicStreamEncoderTest, FinishAfterToolUsesToolStopReason) { EXPECT_EQ(events[0]["type"], "content_block_stop"); EXPECT_EQ(events[1]["type"], "message_delta"); EXPECT_EQ(events[1]["delta"]["stop_reason"], "tool_use"); + ASSERT_TRUE(events[1]["usage"].contains("cache_read_input_tokens")); + EXPECT_EQ(events[1]["usage"]["input_tokens"], 3); + EXPECT_EQ(events[1]["usage"]["cache_read_input_tokens"], 0); + EXPECT_FALSE(events[1]["usage"].contains("cache_creation_input_tokens")); EXPECT_EQ(events[2]["type"], "message_stop"); } diff --git a/tests/xllm_service/scheduler/CMakeLists.txt b/tests/xllm_service/scheduler/CMakeLists.txt index 14a1ad5..7c89d3c 100644 --- a/tests/xllm_service/scheduler/CMakeLists.txt +++ b/tests/xllm_service/scheduler/CMakeLists.txt @@ -9,3 +9,13 @@ cc_test( :xllm_chat_parse_runtime GTest::gtest_main ) + +cc_test( + NAME + response_handler_usage_test + SRCS + response_handler_usage_test.cpp + DEPS + :scheduler + GTest::gtest_main +) diff --git a/tests/xllm_service/scheduler/response_handler_usage_test.cpp b/tests/xllm_service/scheduler/response_handler_usage_test.cpp new file mode 100644 index 0000000..e86956a --- /dev/null +++ b/tests/xllm_service/scheduler/response_handler_usage_test.cpp @@ -0,0 +1,53 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm-service/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "scheduler/response_handler.h" + +namespace xllm_service { +namespace { + +TEST(ResponseHandlerUsageTest, IncludesCachedTokensInOpenAIUsage) { + llm::Usage usage; + usage.num_prompt_tokens = 8; + usage.num_generated_tokens = 2; + usage.num_total_tokens = 10; + usage.num_cached_tokens = 6; + xllm::proto::Usage proto_usage; + + set_openai_usage(&proto_usage, usage); + + EXPECT_EQ(proto_usage.prompt_tokens(), 8); + EXPECT_EQ(proto_usage.completion_tokens(), 2); + EXPECT_EQ(proto_usage.total_tokens(), 10); + ASSERT_TRUE(proto_usage.has_prompt_tokens_details()); + EXPECT_TRUE(proto_usage.prompt_tokens_details().has_cached_tokens()); + EXPECT_EQ(proto_usage.prompt_tokens_details().cached_tokens(), 6); +} + +TEST(ResponseHandlerUsageTest, IncludesZeroCachedTokensInOpenAIUsage) { + llm::Usage usage; + xllm::proto::Usage proto_usage; + + set_openai_usage(&proto_usage, usage); + + ASSERT_TRUE(proto_usage.has_prompt_tokens_details()); + EXPECT_TRUE(proto_usage.prompt_tokens_details().has_cached_tokens()); + EXPECT_EQ(proto_usage.prompt_tokens_details().cached_tokens(), 0); +} + +} // namespace +} // namespace xllm_service diff --git a/xllm_service/common/xllm/output.h b/xllm_service/common/xllm/output.h index 5c1327a..6458ec9 100644 --- a/xllm_service/common/xllm/output.h +++ b/xllm_service/common/xllm/output.h @@ -46,6 +46,9 @@ struct Usage { // the total number of tokens used in the request (prompt + completion). size_t num_total_tokens = 0; + + // the number of prompt tokens served from prefix cache. + size_t num_cached_tokens = 0; }; struct LogProbData { diff --git a/xllm_service/http_service/anthropic_adapter.cpp b/xllm_service/http_service/anthropic_adapter.cpp index 64a1d1f..fcee364 100644 --- a/xllm_service/http_service/anthropic_adapter.cpp +++ b/xllm_service/http_service/anthropic_adapter.cpp @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -32,6 +33,19 @@ limitations under the License. #include "common/xllm/uuid.h" namespace xllm_service { + +void set_anthropic_usage(xllm::proto::AnthropicUsage* proto_usage, + const llm::Usage& usage) { + const size_t num_cached_tokens = + std::min(usage.num_cached_tokens, usage.num_prompt_tokens); + proto_usage->set_input_tokens( + static_cast(usage.num_prompt_tokens - num_cached_tokens)); + proto_usage->set_output_tokens( + static_cast(usage.num_generated_tokens)); + proto_usage->set_cache_read_input_tokens( + static_cast(num_cached_tokens)); +} + namespace { thread_local llm::ShortUUID short_uuid; @@ -375,11 +389,7 @@ void fill_usage(const llm::RequestOutput& request_output, if (!request_output.usage.has_value()) { return; } - const auto& usage = request_output.usage.value(); - auto* proto_usage = response->mutable_usage(); - proto_usage->set_input_tokens(static_cast(usage.num_prompt_tokens)); - proto_usage->set_output_tokens( - static_cast(usage.num_generated_tokens)); + set_anthropic_usage(response->mutable_usage(), request_output.usage.value()); } bool normalize_stream_event_json(const xllm::proto::AnthropicStreamEvent& event, @@ -410,6 +420,10 @@ bool normalize_stream_event_json(const xllm::proto::AnthropicStreamEvent& event, auto& usage = parsed["usage"]; usage["input_tokens"] = event.usage().input_tokens(); usage["output_tokens"] = event.usage().output_tokens(); + if (event.usage().has_cache_read_input_tokens()) { + usage["cache_read_input_tokens"] = + event.usage().cache_read_input_tokens(); + } } } diff --git a/xllm_service/http_service/anthropic_adapter.h b/xllm_service/http_service/anthropic_adapter.h index 68b8676..f5f7cf2 100644 --- a/xllm_service/http_service/anthropic_adapter.h +++ b/xllm_service/http_service/anthropic_adapter.h @@ -32,6 +32,9 @@ struct AnthropicAdaptResult { std::string error; }; +void set_anthropic_usage(xllm::proto::AnthropicUsage* proto_usage, + const llm::Usage& usage); + std::string new_anthropic_id(); AnthropicAdaptResult parse_anthropic_json( diff --git a/xllm_service/http_service/anthropic_stream_encoder.cpp b/xllm_service/http_service/anthropic_stream_encoder.cpp index d39b97f..05f7d57 100644 --- a/xllm_service/http_service/anthropic_stream_encoder.cpp +++ b/xllm_service/http_service/anthropic_stream_encoder.cpp @@ -211,11 +211,7 @@ void AnthropicStreamEncoder::add_message_delta( auto* usage = event.mutable_usage(); if (request_output.usage.has_value()) { - const auto& source_usage = request_output.usage.value(); - usage->set_input_tokens( - static_cast(source_usage.num_prompt_tokens)); - usage->set_output_tokens( - static_cast(source_usage.num_generated_tokens)); + set_anthropic_usage(usage, request_output.usage.value()); } else { usage->set_input_tokens(0); usage->set_output_tokens(0); diff --git a/xllm_service/proto/xllm_rpc_service.proto b/xllm_service/proto/xllm_rpc_service.proto index f18c9df..c53e31f 100644 --- a/xllm_service/proto/xllm_rpc_service.proto +++ b/xllm_service/proto/xllm_rpc_service.proto @@ -99,6 +99,8 @@ message OutputUsage { int32 num_generated_tokens = 2; // the total number of tokens used in the request (prompt + completion). int32 num_total_tokens = 3; + // the number of prompt tokens served from prefix cache. + int32 num_cached_tokens = 4; } message LogProbData { diff --git a/xllm_service/rpc_service/rpc_service_test.cpp b/xllm_service/rpc_service/rpc_service_test.cpp index 4661777..d8d765c 100644 --- a/xllm_service/rpc_service/rpc_service_test.cpp +++ b/xllm_service/rpc_service/rpc_service_test.cpp @@ -26,6 +26,28 @@ class XllmRpcServiceTest : public ::testing::Test { void TearDown() override { google::ShutdownGoogleLogging(); } }; + +TEST_F(XllmRpcServiceTest, OutputUsageUsesFieldFourForCachedTokens) { + const google::protobuf::FieldDescriptor* field = + proto::OutputUsage::descriptor()->FindFieldByName("num_cached_tokens"); + + ASSERT_NE(field, nullptr); + EXPECT_EQ(field->number(), 4); +} + +TEST_F(XllmRpcServiceTest, ConvertsCachedTokensIntoInternalUsage) { + proto::DisaggStreamGeneration generation; + generation.mutable_usage()->set_num_prompt_tokens(8); + generation.mutable_usage()->set_num_generated_tokens(2); + generation.mutable_usage()->set_num_total_tokens(10); + generation.mutable_usage()->set_num_cached_tokens(6); + + llm::RequestOutput output = make_request_output(generation); + + ASSERT_TRUE(output.usage.has_value()); + EXPECT_EQ(output.usage->num_cached_tokens, 6u); +} + // TODO // TEST_F(XllmRpcServiceTest, RegisterInstance) { // RpcServiceConfig config; diff --git a/xllm_service/rpc_service/service.cpp b/xllm_service/rpc_service/service.cpp index 4f9b842..64e0266 100644 --- a/xllm_service/rpc_service/service.cpp +++ b/xllm_service/rpc_service/service.cpp @@ -26,6 +26,68 @@ limitations under the License. namespace xllm_service { +llm::RequestOutput make_request_output( + const proto::DisaggStreamGeneration& request) { + llm::RequestOutput request_output; + request_output.request_id = request.req_id(); + request_output.service_request_id = request.service_req_id(); + if (request.has_gen_status()) { + request_output.status = llm::Status( + static_cast(request.gen_status().status_code()), + request.gen_status().status_msg()); + } + if (request.has_usage()) { + llm::Usage usage; + usage.num_prompt_tokens = request.usage().num_prompt_tokens(); + usage.num_generated_tokens = request.usage().num_generated_tokens(); + usage.num_total_tokens = request.usage().num_total_tokens(); + usage.num_cached_tokens = request.usage().num_cached_tokens(); + request_output.usage = std::move(usage); + } + request_output.finished_on_prefill_instance = + request.finished_on_prefill_instance(); + request_output.finished = request.finished(); + request_output.outputs.reserve(request.outputs_size()); + for (const auto& output : request.outputs()) { + llm::SequenceOutput sequence_output; + sequence_output.index = output.index(); + sequence_output.text = output.text(); + sequence_output.token_ids = std::vector(output.token_ids().begin(), + output.token_ids().end()); + if (!output.finish_reason().empty()) { + sequence_output.finish_reason = output.finish_reason(); + } + if (!output.logprobs().empty()) { + std::vector logprobs; + logprobs.reserve(output.logprobs_size()); + for (const auto& logprob : output.logprobs()) { + llm::LogProb lp; + lp.token = logprob.log_prob_data().token(); + lp.token_id = logprob.log_prob_data().token_id(); + lp.logprob = logprob.log_prob_data().logprob(); + lp.finished_token = logprob.log_prob_data().finished_token(); + if (!logprob.top_logprobs().empty()) { + std::vector top_logprobs; + top_logprobs.reserve(logprob.top_logprobs_size()); + for (const auto& top_logprob : logprob.top_logprobs()) { + llm::LogProbData lpd; + lpd.token = top_logprob.token(); + lpd.token_id = top_logprob.token_id(); + lpd.logprob = top_logprob.logprob(); + lpd.finished_token = top_logprob.finished_token(); + top_logprobs.emplace_back(std::move(lpd)); + } + lp.top_logprobs = std::move(top_logprobs); + } + logprobs.emplace_back(std::move(lp)); + } + sequence_output.logprobs = std::move(logprobs); + } + request_output.outputs.emplace_back(std::move(sequence_output)); + } + return request_output; +} + XllmRpcServiceImpl::XllmRpcServiceImpl(const Options& options, Scheduler* scheduler) : options_(options), scheduler_(scheduler) {} @@ -145,61 +207,8 @@ void XllmRpcService::Generations(google::protobuf::RpcController* cntl_base, brpc::ClosureGuard done_guard(done); // TODO: use threadpool here - for (auto& request : req->gens()) { - // convert proto request to `RequestOutput` - llm::RequestOutput request_output; - request_output.request_id = request.req_id(); - request_output.service_request_id = request.service_req_id(); - if (request.has_gen_status()) { - request_output.status = llm::Status( - static_cast(request.gen_status().status_code()), - request.gen_status().status_msg()); - } - if (request.has_usage()) { - llm::Usage u; - u.num_prompt_tokens = request.usage().num_prompt_tokens(); - u.num_generated_tokens = request.usage().num_generated_tokens(); - u.num_total_tokens = request.usage().num_total_tokens(); - request_output.usage = std::move(u); - } - request_output.finished_on_prefill_instance = - request.finished_on_prefill_instance(); - request_output.finished = request.finished(); - for (auto& output : request.outputs()) { - llm::SequenceOutput sequence_output; - sequence_output.index = output.index(); - sequence_output.text = output.text(); - sequence_output.token_ids = std::vector( - output.token_ids().begin(), output.token_ids().end()); - if (!output.finish_reason().empty()) { - sequence_output.finish_reason = output.finish_reason(); - } - if (output.logprobs().size() > 0) { - std::vector logprobs; - for (auto& logprob : output.logprobs()) { - llm::LogProb lp; - lp.token = logprob.log_prob_data().token(); - lp.token_id = logprob.log_prob_data().token_id(); - lp.logprob = logprob.log_prob_data().logprob(); - lp.finished_token = logprob.log_prob_data().finished_token(); - if (logprob.top_logprobs().size() > 0) { - std::vector top_logprobs; - for (auto& top_logprob : logprob.top_logprobs()) { - llm::LogProbData lpd; - lpd.token = top_logprob.token(); - lpd.token_id = top_logprob.token_id(); - lpd.logprob = top_logprob.logprob(); - lpd.finished_token = top_logprob.finished_token(); - top_logprobs.emplace_back(std::move(lpd)); - } - lp.top_logprobs = std::move(top_logprobs); - } - logprobs.emplace_back(std::move(lp)); - } - sequence_output.logprobs = std::move(logprobs); - } - request_output.outputs.emplace_back(std::move(sequence_output)); - } + for (const auto& request : req->gens()) { + llm::RequestOutput request_output = make_request_output(request); resp->mutable_all_status()->Add()->set_ok( xllm_rpc_service_impl_->handle_generation(request_output)); diff --git a/xllm_service/rpc_service/service.h b/xllm_service/rpc_service/service.h index 99e6529..a2e8c7f 100644 --- a/xllm_service/rpc_service/service.h +++ b/xllm_service/rpc_service/service.h @@ -31,6 +31,9 @@ namespace xllm_service { class Scheduler; class InstanceMgr; +llm::RequestOutput make_request_output( + const proto::DisaggStreamGeneration& request); + class XllmRpcServiceImpl final { public: XllmRpcServiceImpl(const Options& options, Scheduler* scheduler); diff --git a/xllm_service/scheduler/response_handler.cpp b/xllm_service/scheduler/response_handler.cpp index ba293da..34aaf45 100644 --- a/xllm_service/scheduler/response_handler.cpp +++ b/xllm_service/scheduler/response_handler.cpp @@ -26,6 +26,19 @@ limitations under the License. #include "xllm/xllm/function_call/function_call_parser.h" namespace xllm_service { + +void set_openai_usage(xllm::proto::Usage* proto_usage, + const llm::Usage& usage) { + CHECK(proto_usage != nullptr); + proto_usage->set_prompt_tokens(static_cast(usage.num_prompt_tokens)); + proto_usage->set_completion_tokens( + static_cast(usage.num_generated_tokens)); + proto_usage->set_total_tokens(static_cast(usage.num_total_tokens)); + auto* prompt_tokens_details = proto_usage->mutable_prompt_tokens_details(); + prompt_tokens_details->set_cached_tokens( + static_cast(usage.num_cached_tokens)); +} + namespace { size_t find_tool_start(const std::string& text) { @@ -395,12 +408,7 @@ bool ResponseHandler::send_delta_to_client( response.set_id(request_id); response.set_created(created_time); response.set_model(model); - auto* proto_usage = response.mutable_usage(); - proto_usage->set_prompt_tokens( - static_cast(usage.num_prompt_tokens)); - proto_usage->set_completion_tokens( - static_cast(usage.num_generated_tokens)); - proto_usage->set_total_tokens(static_cast(usage.num_total_tokens)); + set_openai_usage(response.mutable_usage(), usage); if (!call_data->write(response)) { return false; } @@ -476,12 +484,7 @@ bool ResponseHandler::send_delta_to_client( response.set_created(created_time); response.set_model(model); response.mutable_choices(); - auto* proto_usage = response.mutable_usage(); - proto_usage->set_prompt_tokens( - static_cast(usage.num_prompt_tokens)); - proto_usage->set_completion_tokens( - static_cast(usage.num_generated_tokens)); - proto_usage->set_total_tokens(static_cast(usage.num_total_tokens)); + set_openai_usage(response.mutable_usage(), usage); if (!call_data->write(response)) { return false; } @@ -761,12 +764,7 @@ bool ResponseHandler::send_result_to_client( // add usage statistics if (req_output.usage.has_value()) { const auto& usage = req_output.usage.value(); - auto* proto_usage = response.mutable_usage(); - proto_usage->set_prompt_tokens( - static_cast(usage.num_prompt_tokens)); - proto_usage->set_completion_tokens( - static_cast(usage.num_generated_tokens)); - proto_usage->set_total_tokens(static_cast(usage.num_total_tokens)); + set_openai_usage(response.mutable_usage(), usage); } return call_data->write_and_finish(response); @@ -809,12 +807,7 @@ bool ResponseHandler::send_result_to_client( // add usage statistics if (req_output.usage.has_value()) { const auto& usage = req_output.usage.value(); - auto* proto_usage = response.mutable_usage(); - proto_usage->set_prompt_tokens( - static_cast(usage.num_prompt_tokens)); - proto_usage->set_completion_tokens( - static_cast(usage.num_generated_tokens)); - proto_usage->set_total_tokens(static_cast(usage.num_total_tokens)); + set_openai_usage(response.mutable_usage(), usage); } return call_data->write_and_finish(response); diff --git a/xllm_service/scheduler/response_handler.h b/xllm_service/scheduler/response_handler.h index 2ae29db..ec4ef15 100644 --- a/xllm_service/scheduler/response_handler.h +++ b/xllm_service/scheduler/response_handler.h @@ -34,6 +34,8 @@ namespace xllm_service { class AnthropicStreamEncoder; +void set_openai_usage(xllm::proto::Usage* proto_usage, const llm::Usage& usage); + struct ChatStreamParseState { std::unordered_set first_message_sent; std::shared_ptr stream_parser; From 8c26c9d55334647dfddb665b56fd86837e7278af Mon Sep 17 00:00:00 2001 From: phantomlei3 Date: Fri, 24 Jul 2026 14:26:12 +0800 Subject: [PATCH 2/5] refactor: centralize usage protocol adaptation --- .../xllm_service/http_service/CMakeLists.txt | 10 + .../http_service/anthropic_adapter_test.cpp | 2 +- .../anthropic_stream_encoder_test.cpp | 2 +- .../http_service/usage_proto_adapter_test.cpp | 92 ++++++++ tests/xllm_service/scheduler/CMakeLists.txt | 10 - .../scheduler/response_handler_usage_test.cpp | 53 ----- third_party/xllm | 2 +- xllm_service/common/xllm/output.h | 3 +- xllm_service/http_service/CMakeLists.txt | 14 ++ .../http_service/anthropic_adapter.cpp | 17 +- xllm_service/http_service/anthropic_adapter.h | 3 - .../http_service/anthropic_stream_encoder.cpp | 3 +- .../http_service/usage_proto_adapter.cpp | 44 ++++ .../http_service/usage_proto_adapter.h | 28 +++ xllm_service/proto/xllm_rpc_service.proto | 2 +- xllm_service/rpc_service/CMakeLists.txt | 15 ++ .../rpc_service/disagg_generation_adapter.cpp | 127 ++++++++++ .../rpc_service/disagg_generation_adapter.h | 34 +++ xllm_service/rpc_service/rpc_service_test.cpp | 221 +++++++++++++----- xllm_service/rpc_service/service.cpp | 80 ++----- xllm_service/rpc_service/service.h | 3 - xllm_service/scheduler/CMakeLists.txt | 1 + xllm_service/scheduler/response_handler.cpp | 21 +- xllm_service/scheduler/response_handler.h | 2 - 24 files changed, 553 insertions(+), 236 deletions(-) create mode 100644 tests/xllm_service/http_service/usage_proto_adapter_test.cpp delete mode 100644 tests/xllm_service/scheduler/response_handler_usage_test.cpp create mode 100644 xllm_service/http_service/usage_proto_adapter.cpp create mode 100644 xllm_service/http_service/usage_proto_adapter.h create mode 100644 xllm_service/rpc_service/disagg_generation_adapter.cpp create mode 100644 xllm_service/rpc_service/disagg_generation_adapter.h diff --git a/tests/xllm_service/http_service/CMakeLists.txt b/tests/xllm_service/http_service/CMakeLists.txt index 0bd9b68..b9c3a0f 100644 --- a/tests/xllm_service/http_service/CMakeLists.txt +++ b/tests/xllm_service/http_service/CMakeLists.txt @@ -1,5 +1,15 @@ include(cc_test) +cc_test( + NAME + usage_proto_adapter_test + SRCS + usage_proto_adapter_test.cpp + DEPS + :usage_proto_adapter + GTest::gtest_main +) + cc_test( NAME chat_json_parser_test diff --git a/tests/xllm_service/http_service/anthropic_adapter_test.cpp b/tests/xllm_service/http_service/anthropic_adapter_test.cpp index 839aba2..4d38c5b 100644 --- a/tests/xllm_service/http_service/anthropic_adapter_test.cpp +++ b/tests/xllm_service/http_service/anthropic_adapter_test.cpp @@ -807,7 +807,7 @@ TEST(AnthropicAdapterTest, BuildsNonStreamAnthropicJson) { usage.num_prompt_tokens = 3; usage.num_generated_tokens = 4; usage.num_total_tokens = 7; - usage.num_cached_tokens = 2; + usage.num_prefix_cache_hit_tokens = 2; output.usage = usage; xllm::proto::AnthropicMessagesResponse response; diff --git a/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp b/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp index 9f55018..634fc5b 100644 --- a/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp +++ b/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp @@ -243,7 +243,7 @@ TEST(AnthropicStreamEncoderTest, FinishClosesBlockAndEmitsMessageDeltaStop) { llm::Usage usage; usage.num_prompt_tokens = 3; usage.num_generated_tokens = 5; - usage.num_cached_tokens = 2; + usage.num_prefix_cache_hit_tokens = 2; final.usage = usage; std::vector done; diff --git a/tests/xllm_service/http_service/usage_proto_adapter_test.cpp b/tests/xllm_service/http_service/usage_proto_adapter_test.cpp new file mode 100644 index 0000000..32349ac --- /dev/null +++ b/tests/xllm_service/http_service/usage_proto_adapter_test.cpp @@ -0,0 +1,92 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm-service/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "http_service/usage_proto_adapter.h" + +#include + +#include +#include + +namespace xllm_service { +namespace { + +struct UsageCase { + size_t num_prompt_tokens; + size_t num_generated_tokens; + size_t num_total_tokens; + size_t num_prefix_cache_hit_tokens; +}; + +const std::vector kUsageCases = { + {/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/0, + /*num_total_tokens=*/8, + /*num_prefix_cache_hit_tokens=*/0}, + {/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/10, + /*num_prefix_cache_hit_tokens=*/3}, + {/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/10, + /*num_prefix_cache_hit_tokens=*/8}, +}; + +llm::Usage make_usage(const UsageCase& usage_case) { + llm::Usage usage; + usage.num_prompt_tokens = usage_case.num_prompt_tokens; + usage.num_generated_tokens = usage_case.num_generated_tokens; + usage.num_total_tokens = usage_case.num_total_tokens; + usage.num_prefix_cache_hit_tokens = usage_case.num_prefix_cache_hit_tokens; + return usage; +} + +TEST(UsageProtoAdapterTest, ConvertsPrefixCacheHitsToOpenAIPromptTokenDetails) { + for (const UsageCase& usage_case : kUsageCases) { + xllm::proto::Usage proto_usage = + to_openai_usage_proto(make_usage(usage_case)); + + EXPECT_EQ(proto_usage.prompt_tokens(), usage_case.num_prompt_tokens); + EXPECT_EQ(proto_usage.completion_tokens(), usage_case.num_generated_tokens); + EXPECT_EQ(proto_usage.total_tokens(), usage_case.num_total_tokens); + ASSERT_TRUE(proto_usage.has_prompt_tokens_details()); + EXPECT_TRUE(proto_usage.prompt_tokens_details().has_cached_tokens()); + EXPECT_EQ(proto_usage.prompt_tokens_details().cached_tokens(), + usage_case.num_prefix_cache_hit_tokens); + } +} + +TEST(UsageProtoAdapterTest, + SplitsAnthropicInputTokensIntoUncachedAndCacheReadTokens) { + for (const UsageCase& usage_case : kUsageCases) { + xllm::proto::AnthropicUsage proto_usage = + to_anthropic_usage_proto(make_usage(usage_case)); + + EXPECT_EQ( + proto_usage.input_tokens(), + usage_case.num_prompt_tokens - usage_case.num_prefix_cache_hit_tokens); + EXPECT_EQ(proto_usage.cache_read_input_tokens(), + usage_case.num_prefix_cache_hit_tokens); + EXPECT_EQ(proto_usage.output_tokens(), usage_case.num_generated_tokens); + EXPECT_EQ( + proto_usage.input_tokens() + proto_usage.cache_read_input_tokens(), + usage_case.num_prompt_tokens); + EXPECT_FALSE(proto_usage.has_cache_creation_input_tokens()); + } +} + +} // namespace +} // namespace xllm_service diff --git a/tests/xllm_service/scheduler/CMakeLists.txt b/tests/xllm_service/scheduler/CMakeLists.txt index 7c89d3c..14a1ad5 100644 --- a/tests/xllm_service/scheduler/CMakeLists.txt +++ b/tests/xllm_service/scheduler/CMakeLists.txt @@ -9,13 +9,3 @@ cc_test( :xllm_chat_parse_runtime GTest::gtest_main ) - -cc_test( - NAME - response_handler_usage_test - SRCS - response_handler_usage_test.cpp - DEPS - :scheduler - GTest::gtest_main -) diff --git a/tests/xllm_service/scheduler/response_handler_usage_test.cpp b/tests/xllm_service/scheduler/response_handler_usage_test.cpp deleted file mode 100644 index e86956a..0000000 --- a/tests/xllm_service/scheduler/response_handler_usage_test.cpp +++ /dev/null @@ -1,53 +0,0 @@ -/* Copyright 2025-2026 The xLLM Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - https://github.com/jd-opensource/xllm-service/blob/main/LICENSE - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include - -#include "scheduler/response_handler.h" - -namespace xllm_service { -namespace { - -TEST(ResponseHandlerUsageTest, IncludesCachedTokensInOpenAIUsage) { - llm::Usage usage; - usage.num_prompt_tokens = 8; - usage.num_generated_tokens = 2; - usage.num_total_tokens = 10; - usage.num_cached_tokens = 6; - xllm::proto::Usage proto_usage; - - set_openai_usage(&proto_usage, usage); - - EXPECT_EQ(proto_usage.prompt_tokens(), 8); - EXPECT_EQ(proto_usage.completion_tokens(), 2); - EXPECT_EQ(proto_usage.total_tokens(), 10); - ASSERT_TRUE(proto_usage.has_prompt_tokens_details()); - EXPECT_TRUE(proto_usage.prompt_tokens_details().has_cached_tokens()); - EXPECT_EQ(proto_usage.prompt_tokens_details().cached_tokens(), 6); -} - -TEST(ResponseHandlerUsageTest, IncludesZeroCachedTokensInOpenAIUsage) { - llm::Usage usage; - xllm::proto::Usage proto_usage; - - set_openai_usage(&proto_usage, usage); - - ASSERT_TRUE(proto_usage.has_prompt_tokens_details()); - EXPECT_TRUE(proto_usage.prompt_tokens_details().has_cached_tokens()); - EXPECT_EQ(proto_usage.prompt_tokens_details().cached_tokens(), 0); -} - -} // namespace -} // namespace xllm_service diff --git a/third_party/xllm b/third_party/xllm index adf07f7..8039573 160000 --- a/third_party/xllm +++ b/third_party/xllm @@ -1 +1 @@ -Subproject commit adf07f70cdeaa49ba8a7ed2dd0db1d0f2b718539 +Subproject commit 8039573ec9671e59e87927c883427c6ff975c053 diff --git a/xllm_service/common/xllm/output.h b/xllm_service/common/xllm/output.h index 6458ec9..80ee71d 100644 --- a/xllm_service/common/xllm/output.h +++ b/xllm_service/common/xllm/output.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include #include #include #include @@ -48,7 +49,7 @@ struct Usage { size_t num_total_tokens = 0; // the number of prompt tokens served from prefix cache. - size_t num_cached_tokens = 0; + size_t num_prefix_cache_hit_tokens = 0; }; struct LogProbData { diff --git a/xllm_service/http_service/CMakeLists.txt b/xllm_service/http_service/CMakeLists.txt index acd4aaa..c8afbd1 100644 --- a/xllm_service/http_service/CMakeLists.txt +++ b/xllm_service/http_service/CMakeLists.txt @@ -13,6 +13,18 @@ cc_library( nlohmann_json::nlohmann_json ) +cc_library( + NAME + usage_proto_adapter + HDRS + usage_proto_adapter.h + SRCS + usage_proto_adapter.cpp + DEPS + :common + proto_xllm +) + cc_library( NAME anthropic_adapter @@ -29,6 +41,7 @@ cc_library( DEPS :chat_template :message_projection + :usage_proto_adapter :common glog::glog nlohmann_json::nlohmann_json @@ -50,6 +63,7 @@ cc_library( DEPS :anthropic_adapter :common + :usage_proto_adapter glog::glog nlohmann_json::nlohmann_json proto_xllm diff --git a/xllm_service/http_service/anthropic_adapter.cpp b/xllm_service/http_service/anthropic_adapter.cpp index fcee364..ce28079 100644 --- a/xllm_service/http_service/anthropic_adapter.cpp +++ b/xllm_service/http_service/anthropic_adapter.cpp @@ -18,7 +18,6 @@ limitations under the License. #include #include -#include #include #include #include @@ -31,21 +30,10 @@ limitations under the License. #include "api_service/chat_json_parser.h" #include "chat_template/message_projection.h" #include "common/xllm/uuid.h" +#include "http_service/usage_proto_adapter.h" namespace xllm_service { -void set_anthropic_usage(xllm::proto::AnthropicUsage* proto_usage, - const llm::Usage& usage) { - const size_t num_cached_tokens = - std::min(usage.num_cached_tokens, usage.num_prompt_tokens); - proto_usage->set_input_tokens( - static_cast(usage.num_prompt_tokens - num_cached_tokens)); - proto_usage->set_output_tokens( - static_cast(usage.num_generated_tokens)); - proto_usage->set_cache_read_input_tokens( - static_cast(num_cached_tokens)); -} - namespace { thread_local llm::ShortUUID short_uuid; @@ -389,7 +377,8 @@ void fill_usage(const llm::RequestOutput& request_output, if (!request_output.usage.has_value()) { return; } - set_anthropic_usage(response->mutable_usage(), request_output.usage.value()); + *response->mutable_usage() = + to_anthropic_usage_proto(request_output.usage.value()); } bool normalize_stream_event_json(const xllm::proto::AnthropicStreamEvent& event, diff --git a/xllm_service/http_service/anthropic_adapter.h b/xllm_service/http_service/anthropic_adapter.h index f5f7cf2..68b8676 100644 --- a/xllm_service/http_service/anthropic_adapter.h +++ b/xllm_service/http_service/anthropic_adapter.h @@ -32,9 +32,6 @@ struct AnthropicAdaptResult { std::string error; }; -void set_anthropic_usage(xllm::proto::AnthropicUsage* proto_usage, - const llm::Usage& usage); - std::string new_anthropic_id(); AnthropicAdaptResult parse_anthropic_json( diff --git a/xllm_service/http_service/anthropic_stream_encoder.cpp b/xllm_service/http_service/anthropic_stream_encoder.cpp index 05f7d57..d2ee111 100644 --- a/xllm_service/http_service/anthropic_stream_encoder.cpp +++ b/xllm_service/http_service/anthropic_stream_encoder.cpp @@ -20,6 +20,7 @@ limitations under the License. #include "api_service/anthropic_stream_utils.h" #include "common/xllm/uuid.h" +#include "http_service/usage_proto_adapter.h" namespace xllm_service { namespace { @@ -211,7 +212,7 @@ void AnthropicStreamEncoder::add_message_delta( auto* usage = event.mutable_usage(); if (request_output.usage.has_value()) { - set_anthropic_usage(usage, request_output.usage.value()); + *usage = to_anthropic_usage_proto(request_output.usage.value()); } else { usage->set_input_tokens(0); usage->set_output_tokens(0); diff --git a/xllm_service/http_service/usage_proto_adapter.cpp b/xllm_service/http_service/usage_proto_adapter.cpp new file mode 100644 index 0000000..6731c89 --- /dev/null +++ b/xllm_service/http_service/usage_proto_adapter.cpp @@ -0,0 +1,44 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm-service/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "http_service/usage_proto_adapter.h" + +#include + +namespace xllm_service { + +xllm::proto::Usage to_openai_usage_proto(const llm::Usage& usage) { + xllm::proto::Usage proto_usage; + proto_usage.set_prompt_tokens(static_cast(usage.num_prompt_tokens)); + proto_usage.set_completion_tokens( + static_cast(usage.num_generated_tokens)); + proto_usage.set_total_tokens(static_cast(usage.num_total_tokens)); + proto_usage.mutable_prompt_tokens_details()->set_cached_tokens( + static_cast(usage.num_prefix_cache_hit_tokens)); + return proto_usage; +} + +xllm::proto::AnthropicUsage to_anthropic_usage_proto(const llm::Usage& usage) { + xllm::proto::AnthropicUsage proto_usage; + proto_usage.set_input_tokens(static_cast( + usage.num_prompt_tokens - usage.num_prefix_cache_hit_tokens)); + proto_usage.set_output_tokens( + static_cast(usage.num_generated_tokens)); + proto_usage.set_cache_read_input_tokens( + static_cast(usage.num_prefix_cache_hit_tokens)); + return proto_usage; +} + +} // namespace xllm_service diff --git a/xllm_service/http_service/usage_proto_adapter.h b/xllm_service/http_service/usage_proto_adapter.h new file mode 100644 index 0000000..bffce62 --- /dev/null +++ b/xllm_service/http_service/usage_proto_adapter.h @@ -0,0 +1,28 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm-service/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include "anthropic.pb.h" +#include "chat.pb.h" +#include "common/xllm/output.h" + +namespace xllm_service { + +xllm::proto::Usage to_openai_usage_proto(const llm::Usage& usage); + +xllm::proto::AnthropicUsage to_anthropic_usage_proto(const llm::Usage& usage); + +} // namespace xllm_service diff --git a/xllm_service/proto/xllm_rpc_service.proto b/xllm_service/proto/xllm_rpc_service.proto index c53e31f..185bfcf 100644 --- a/xllm_service/proto/xllm_rpc_service.proto +++ b/xllm_service/proto/xllm_rpc_service.proto @@ -100,7 +100,7 @@ message OutputUsage { // the total number of tokens used in the request (prompt + completion). int32 num_total_tokens = 3; // the number of prompt tokens served from prefix cache. - int32 num_cached_tokens = 4; + int32 num_prefix_cache_hit_tokens = 4; } message LogProbData { diff --git a/xllm_service/rpc_service/CMakeLists.txt b/xllm_service/rpc_service/CMakeLists.txt index 5e4ef79..8afa0b4 100644 --- a/xllm_service/rpc_service/CMakeLists.txt +++ b/xllm_service/rpc_service/CMakeLists.txt @@ -2,6 +2,18 @@ include(cc_binary) include(cc_library) include(cc_test) +cc_library( + NAME + disagg_generation_adapter + HDRS + disagg_generation_adapter.h + SRCS + disagg_generation_adapter.cpp + DEPS + :common + proto::proto_rpc_service +) + cc_library( NAME xllm_rpc_service @@ -11,6 +23,7 @@ cc_library( service.cpp DEPS :common + :disagg_generation_adapter :scheduler absl::random_random absl::strings @@ -28,10 +41,12 @@ cc_binary( SRCS rpc_service_test.cpp DEPS + :disagg_generation_adapter :xllm_rpc_service gflags::gflags glog::glog GTest::gtest_main + proto_xllm ) add_test(NAME XllmRpcServiceTest COMMAND xllm_rpc_service_test) diff --git a/xllm_service/rpc_service/disagg_generation_adapter.cpp b/xllm_service/rpc_service/disagg_generation_adapter.cpp new file mode 100644 index 0000000..805940b --- /dev/null +++ b/xllm_service/rpc_service/disagg_generation_adapter.cpp @@ -0,0 +1,127 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm-service/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "rpc_service/disagg_generation_adapter.h" + +#include +#include +#include +#include + +namespace xllm_service { +namespace { + +RequestOutputConversionResult invalid_usage(std::string message) { + return {llm::Status(llm::StatusCode::INVALID_ARGUMENT, std::move(message)), + std::nullopt}; +} + +RequestOutputConversionResult validate_usage(const proto::OutputUsage& usage) { + if (usage.num_prompt_tokens() < 0 || usage.num_generated_tokens() < 0 || + usage.num_total_tokens() < 0 || usage.num_prefix_cache_hit_tokens() < 0) { + return invalid_usage("token counts must be non-negative"); + } + if (usage.num_prefix_cache_hit_tokens() > usage.num_prompt_tokens()) { + return invalid_usage( + "prefix cache hit tokens must not exceed prompt tokens"); + } + const int64_t expected_total = + static_cast(usage.num_prompt_tokens()) + + static_cast(usage.num_generated_tokens()); + if (static_cast(usage.num_total_tokens()) != expected_total) { + return invalid_usage( + "total tokens must equal prompt tokens plus generated tokens"); + } + return {}; +} + +} // namespace + +RequestOutputConversionResult request_output_from_disagg_generation( + const proto::DisaggStreamGeneration& generation) { + if (generation.has_usage()) { + RequestOutputConversionResult validation = + validate_usage(generation.usage()); + if (!validation.status.ok()) { + return validation; + } + } + + llm::RequestOutput request_output; + request_output.request_id = generation.req_id(); + request_output.service_request_id = generation.service_req_id(); + if (generation.has_gen_status()) { + request_output.status = llm::Status( + static_cast(generation.gen_status().status_code()), + generation.gen_status().status_msg()); + } + if (generation.has_usage()) { + llm::Usage usage; + usage.num_prompt_tokens = + static_cast(generation.usage().num_prompt_tokens()); + usage.num_generated_tokens = + static_cast(generation.usage().num_generated_tokens()); + usage.num_total_tokens = + static_cast(generation.usage().num_total_tokens()); + usage.num_prefix_cache_hit_tokens = + static_cast(generation.usage().num_prefix_cache_hit_tokens()); + request_output.usage = std::move(usage); + } + request_output.finished_on_prefill_instance = + generation.finished_on_prefill_instance(); + request_output.finished = generation.finished(); + request_output.outputs.reserve(generation.outputs_size()); + for (const proto::SequenceOutput& output : generation.outputs()) { + llm::SequenceOutput sequence_output; + sequence_output.index = static_cast(output.index()); + sequence_output.text = output.text(); + sequence_output.token_ids = std::vector(output.token_ids().begin(), + output.token_ids().end()); + if (!output.finish_reason().empty()) { + sequence_output.finish_reason = output.finish_reason(); + } + if (!output.logprobs().empty()) { + std::vector logprobs; + logprobs.reserve(output.logprobs_size()); + for (const proto::LogProb& logprob : output.logprobs()) { + llm::LogProb converted_logprob; + converted_logprob.token = logprob.log_prob_data().token(); + converted_logprob.token_id = logprob.log_prob_data().token_id(); + converted_logprob.logprob = logprob.log_prob_data().logprob(); + converted_logprob.finished_token = + logprob.log_prob_data().finished_token(); + if (!logprob.top_logprobs().empty()) { + std::vector top_logprobs; + top_logprobs.reserve(logprob.top_logprobs_size()); + for (const proto::LogProbData& top_logprob : logprob.top_logprobs()) { + llm::LogProbData converted_top_logprob; + converted_top_logprob.token = top_logprob.token(); + converted_top_logprob.token_id = top_logprob.token_id(); + converted_top_logprob.logprob = top_logprob.logprob(); + converted_top_logprob.finished_token = top_logprob.finished_token(); + top_logprobs.emplace_back(std::move(converted_top_logprob)); + } + converted_logprob.top_logprobs = std::move(top_logprobs); + } + logprobs.emplace_back(std::move(converted_logprob)); + } + sequence_output.logprobs = std::move(logprobs); + } + request_output.outputs.emplace_back(std::move(sequence_output)); + } + return {llm::Status(), std::move(request_output)}; +} + +} // namespace xllm_service diff --git a/xllm_service/rpc_service/disagg_generation_adapter.h b/xllm_service/rpc_service/disagg_generation_adapter.h new file mode 100644 index 0000000..53b4113 --- /dev/null +++ b/xllm_service/rpc_service/disagg_generation_adapter.h @@ -0,0 +1,34 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm-service/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "common/xllm/output.h" +#include "common/xllm/status.h" +#include "xllm_rpc_service.pb.h" + +namespace xllm_service { + +struct RequestOutputConversionResult { + llm::Status status; + std::optional output; +}; + +RequestOutputConversionResult request_output_from_disagg_generation( + const proto::DisaggStreamGeneration& generation); + +} // namespace xllm_service diff --git a/xllm_service/rpc_service/rpc_service_test.cpp b/xllm_service/rpc_service/rpc_service_test.cpp index d8d765c..5f662e6 100644 --- a/xllm_service/rpc_service/rpc_service_test.cpp +++ b/xllm_service/rpc_service/rpc_service_test.cpp @@ -13,79 +13,176 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include #include -#include "rpc_service/service.h" +#include +#include -namespace xllm_service::test { +#include "disagg_pd.pb.h" +#include "rpc_service/disagg_generation_adapter.h" -class XllmRpcServiceTest : public ::testing::Test { - protected: - void SetUp() override { google::InitGoogleLogging("XllmRpcServiceTest"); } +namespace xllm_service { +namespace { - void TearDown() override { google::ShutdownGoogleLogging(); } -}; +proto::DisaggStreamGeneration make_generation(int32_t num_prompt_tokens, + int32_t num_generated_tokens, + int32_t num_total_tokens, + int32_t num_cache_hit_tokens) { + proto::DisaggStreamGeneration generation; + generation.set_req_id("request-123"); + generation.mutable_usage()->set_num_prompt_tokens(num_prompt_tokens); + generation.mutable_usage()->set_num_generated_tokens(num_generated_tokens); + generation.mutable_usage()->set_num_total_tokens(num_total_tokens); + generation.mutable_usage()->set_num_prefix_cache_hit_tokens( + num_cache_hit_tokens); + return generation; +} + +TEST(DisaggGenerationAdapterTest, + PrefixCacheHitTokensRemainWireCompatibleAtFieldFour) { + const google::protobuf::FieldDescriptor* sender_field = + xllm::proto::OutputUsage::descriptor()->FindFieldByName( + "num_prefix_cache_hit_tokens"); + const google::protobuf::FieldDescriptor* receiver_field = + proto::OutputUsage::descriptor()->FindFieldByName( + "num_prefix_cache_hit_tokens"); + + ASSERT_NE(sender_field, nullptr); + ASSERT_NE(receiver_field, nullptr); + EXPECT_EQ(sender_field->number(), 4); + EXPECT_EQ(receiver_field->number(), 4); + EXPECT_EQ(sender_field->type(), receiver_field->type()); + + xllm::proto::OutputUsage sender_usage; + sender_usage.set_num_prompt_tokens(8); + sender_usage.set_num_generated_tokens(2); + sender_usage.set_num_total_tokens(10); + sender_usage.set_num_prefix_cache_hit_tokens(6); + + proto::OutputUsage receiver_usage; + ASSERT_TRUE(receiver_usage.ParseFromString(sender_usage.SerializeAsString())); + EXPECT_EQ(receiver_usage.num_prefix_cache_hit_tokens(), 6); +} + +TEST(DisaggGenerationAdapterTest, RejectsNegativeTokenCounts) { + const proto::DisaggStreamGeneration generation = + make_generation(/*num_prompt_tokens=*/-1, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/1, + /*num_cache_hit_tokens=*/0); -TEST_F(XllmRpcServiceTest, OutputUsageUsesFieldFourForCachedTokens) { - const google::protobuf::FieldDescriptor* field = - proto::OutputUsage::descriptor()->FindFieldByName("num_cached_tokens"); + RequestOutputConversionResult result = + request_output_from_disagg_generation(generation); - ASSERT_NE(field, nullptr); - EXPECT_EQ(field->number(), 4); + EXPECT_FALSE(result.status.ok()); + EXPECT_EQ(result.status.code(), llm::StatusCode::INVALID_ARGUMENT); + EXPECT_NE(result.status.message().find("non-negative"), std::string::npos); + EXPECT_FALSE(result.output.has_value()); } -TEST_F(XllmRpcServiceTest, ConvertsCachedTokensIntoInternalUsage) { - proto::DisaggStreamGeneration generation; - generation.mutable_usage()->set_num_prompt_tokens(8); - generation.mutable_usage()->set_num_generated_tokens(2); - generation.mutable_usage()->set_num_total_tokens(10); - generation.mutable_usage()->set_num_cached_tokens(6); +TEST(DisaggGenerationAdapterTest, RejectsCacheHitsGreaterThanPromptTokens) { + const proto::DisaggStreamGeneration generation = + make_generation(/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/10, + /*num_cache_hit_tokens=*/9); + + RequestOutputConversionResult result = + request_output_from_disagg_generation(generation); + + EXPECT_FALSE(result.status.ok()); + EXPECT_EQ(result.status.code(), llm::StatusCode::INVALID_ARGUMENT); + EXPECT_NE(result.status.message().find("prefix cache hit"), + std::string::npos); + EXPECT_FALSE(result.output.has_value()); +} + +TEST(DisaggGenerationAdapterTest, RejectsInconsistentTotalTokens) { + const proto::DisaggStreamGeneration generation = + make_generation(/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/11, + /*num_cache_hit_tokens=*/6); - llm::RequestOutput output = make_request_output(generation); + RequestOutputConversionResult result = + request_output_from_disagg_generation(generation); + + EXPECT_FALSE(result.status.ok()); + EXPECT_EQ(result.status.code(), llm::StatusCode::INVALID_ARGUMENT); + EXPECT_NE(result.status.message().find("total tokens"), std::string::npos); + EXPECT_FALSE(result.output.has_value()); +} +TEST(DisaggGenerationAdapterTest, ConvertsCompleteValidGeneration) { + proto::DisaggStreamGeneration generation = + make_generation(/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/10, + /*num_cache_hit_tokens=*/6); + generation.set_service_req_id("service-request-456"); + generation.mutable_gen_status()->set_status_code( + static_cast(llm::StatusCode::OK)); + generation.mutable_gen_status()->set_status_msg("complete"); + generation.set_finished(true); + generation.set_finished_on_prefill_instance(true); + + proto::SequenceOutput* sequence = generation.add_outputs(); + sequence->set_index(3); + sequence->set_text("answer"); + sequence->add_token_ids(101); + sequence->add_token_ids(102); + sequence->set_finish_reason("stop"); + proto::LogProb* logprob = sequence->add_logprobs(); + logprob->mutable_log_prob_data()->set_token("answer"); + logprob->mutable_log_prob_data()->set_token_id(102); + logprob->mutable_log_prob_data()->set_logprob(-0.25f); + logprob->mutable_log_prob_data()->set_finished_token(true); + proto::LogProbData* top_logprob = logprob->add_top_logprobs(); + top_logprob->set_token("result"); + top_logprob->set_token_id(103); + top_logprob->set_logprob(-0.5f); + top_logprob->set_finished_token(false); + + RequestOutputConversionResult result = + request_output_from_disagg_generation(generation); + + ASSERT_TRUE(result.status.ok()) << result.status.message(); + ASSERT_TRUE(result.output.has_value()); + const llm::RequestOutput& output = result.output.value(); + EXPECT_EQ(output.request_id, "request-123"); + EXPECT_EQ(output.service_request_id, "service-request-456"); + ASSERT_TRUE(output.status.has_value()); + EXPECT_TRUE(output.status->ok()); + EXPECT_EQ(output.status->message(), "complete"); + EXPECT_TRUE(output.finished); + EXPECT_TRUE(output.finished_on_prefill_instance); ASSERT_TRUE(output.usage.has_value()); - EXPECT_EQ(output.usage->num_cached_tokens, 6u); + EXPECT_EQ(output.usage->num_prompt_tokens, 8u); + EXPECT_EQ(output.usage->num_generated_tokens, 2u); + EXPECT_EQ(output.usage->num_total_tokens, 10u); + EXPECT_EQ(output.usage->num_prefix_cache_hit_tokens, 6u); + + ASSERT_EQ(output.outputs.size(), 1u); + const llm::SequenceOutput& converted_sequence = output.outputs.front(); + EXPECT_EQ(converted_sequence.index, 3u); + EXPECT_EQ(converted_sequence.text, "answer"); + EXPECT_EQ(converted_sequence.token_ids, (std::vector{101, 102})); + ASSERT_TRUE(converted_sequence.finish_reason.has_value()); + EXPECT_EQ(converted_sequence.finish_reason.value(), "stop"); + ASSERT_TRUE(converted_sequence.logprobs.has_value()); + ASSERT_EQ(converted_sequence.logprobs->size(), 1u); + const llm::LogProb& converted_logprob = converted_sequence.logprobs->front(); + EXPECT_EQ(converted_logprob.token, "answer"); + EXPECT_EQ(converted_logprob.token_id, 102); + EXPECT_FLOAT_EQ(converted_logprob.logprob, -0.25f); + EXPECT_TRUE(converted_logprob.finished_token); + ASSERT_TRUE(converted_logprob.top_logprobs.has_value()); + ASSERT_EQ(converted_logprob.top_logprobs->size(), 1u); + EXPECT_EQ(converted_logprob.top_logprobs->front().token, "result"); + EXPECT_EQ(converted_logprob.top_logprobs->front().token_id, 103); + EXPECT_FLOAT_EQ(converted_logprob.top_logprobs->front().logprob, -0.5f); + EXPECT_FALSE(converted_logprob.top_logprobs->front().finished_token); } -// TODO -// TEST_F(XllmRpcServiceTest, RegisterInstance) { -// RpcServiceConfig config; -// HttpServiceConfig http_config; -// ModelConfig model_config; -// auto xllm_service = -// std::make_shared(config, model_config, -// http_config); -// std::string inst_name = "127.0.0.1@nic0"; -// InstanceMetaInfo metainfo(inst_name, "127.0.0.1:7777", -// InstanceType::PREFILL); EXPECT_EQ(ErrorCode::OK, -// xllm_service->register_instance(inst_name, metainfo)); - -// metainfo.type = InstanceType::DECODE; -// EXPECT_EQ(ErrorCode::INSTANCE_EXISTED, -// xllm_service->register_instance(inst_name, metainfo)); -// } - -// TEST_F(XllmRpcServiceTest, UpdateInstanceMetainfo) { -// RpcServiceConfig config; -// HttpServiceConfig http_config; -// ModelConfig model_config; -// auto xllm_service = -// std::make_shared(config, model_config, -// http_config); -// std::string inst_name = "127.0.0.1@nic0"; -// InstanceMetaInfo metainfo(inst_name, "127.0.0.1:7777", -// InstanceType::PREFILL); EXPECT_EQ(ErrorCode::OK, -// xllm_service->register_instance(inst_name, metainfo)); -// metainfo.type = InstanceType::DECODE; -// EXPECT_EQ(ErrorCode::OK, -// xllm_service->update_instance_metainfo(inst_name, metainfo)); - -// std::string inst_name2 = "127.0.0.1@nic2"; -// InstanceMetaInfo metainfo2( -// inst_name2, "127.0.0.1:7778", InstanceType::PREFILL); -// EXPECT_EQ(ErrorCode::INSTANCE_NOT_EXISTED, -// xllm_service->update_instance_metainfo(inst_name2, metainfo)); -// } - -} // namespace xllm_service::test +} // namespace +} // namespace xllm_service diff --git a/xllm_service/rpc_service/service.cpp b/xllm_service/rpc_service/service.cpp index 64e0266..2dfd5b4 100644 --- a/xllm_service/rpc_service/service.cpp +++ b/xllm_service/rpc_service/service.cpp @@ -22,72 +22,11 @@ limitations under the License. #include "common/types.h" #include "common/utils.h" #include "common/xllm/status.h" +#include "rpc_service/disagg_generation_adapter.h" #include "scheduler/scheduler.h" namespace xllm_service { -llm::RequestOutput make_request_output( - const proto::DisaggStreamGeneration& request) { - llm::RequestOutput request_output; - request_output.request_id = request.req_id(); - request_output.service_request_id = request.service_req_id(); - if (request.has_gen_status()) { - request_output.status = llm::Status( - static_cast(request.gen_status().status_code()), - request.gen_status().status_msg()); - } - if (request.has_usage()) { - llm::Usage usage; - usage.num_prompt_tokens = request.usage().num_prompt_tokens(); - usage.num_generated_tokens = request.usage().num_generated_tokens(); - usage.num_total_tokens = request.usage().num_total_tokens(); - usage.num_cached_tokens = request.usage().num_cached_tokens(); - request_output.usage = std::move(usage); - } - request_output.finished_on_prefill_instance = - request.finished_on_prefill_instance(); - request_output.finished = request.finished(); - request_output.outputs.reserve(request.outputs_size()); - for (const auto& output : request.outputs()) { - llm::SequenceOutput sequence_output; - sequence_output.index = output.index(); - sequence_output.text = output.text(); - sequence_output.token_ids = std::vector(output.token_ids().begin(), - output.token_ids().end()); - if (!output.finish_reason().empty()) { - sequence_output.finish_reason = output.finish_reason(); - } - if (!output.logprobs().empty()) { - std::vector logprobs; - logprobs.reserve(output.logprobs_size()); - for (const auto& logprob : output.logprobs()) { - llm::LogProb lp; - lp.token = logprob.log_prob_data().token(); - lp.token_id = logprob.log_prob_data().token_id(); - lp.logprob = logprob.log_prob_data().logprob(); - lp.finished_token = logprob.log_prob_data().finished_token(); - if (!logprob.top_logprobs().empty()) { - std::vector top_logprobs; - top_logprobs.reserve(logprob.top_logprobs_size()); - for (const auto& top_logprob : logprob.top_logprobs()) { - llm::LogProbData lpd; - lpd.token = top_logprob.token(); - lpd.token_id = top_logprob.token_id(); - lpd.logprob = top_logprob.logprob(); - lpd.finished_token = top_logprob.finished_token(); - top_logprobs.emplace_back(std::move(lpd)); - } - lp.top_logprobs = std::move(top_logprobs); - } - logprobs.emplace_back(std::move(lp)); - } - sequence_output.logprobs = std::move(logprobs); - } - request_output.outputs.emplace_back(std::move(sequence_output)); - } - return request_output; -} - XllmRpcServiceImpl::XllmRpcServiceImpl(const Options& options, Scheduler* scheduler) : options_(options), scheduler_(scheduler) {} @@ -207,11 +146,18 @@ void XllmRpcService::Generations(google::protobuf::RpcController* cntl_base, brpc::ClosureGuard done_guard(done); // TODO: use threadpool here - for (const auto& request : req->gens()) { - llm::RequestOutput request_output = make_request_output(request); - - resp->mutable_all_status()->Add()->set_ok( - xllm_rpc_service_impl_->handle_generation(request_output)); + for (const proto::DisaggStreamGeneration& generation : req->gens()) { + RequestOutputConversionResult conversion = + request_output_from_disagg_generation(generation); + proto::Status* status = resp->mutable_all_status()->Add(); + if (!conversion.status.ok()) { + LOG(ERROR) << "Rejecting invalid generation for request " + << generation.req_id() << ": " << conversion.status.message(); + status->set_ok(false); + continue; + } + status->set_ok( + xllm_rpc_service_impl_->handle_generation(conversion.output.value())); } } diff --git a/xllm_service/rpc_service/service.h b/xllm_service/rpc_service/service.h index a2e8c7f..99e6529 100644 --- a/xllm_service/rpc_service/service.h +++ b/xllm_service/rpc_service/service.h @@ -31,9 +31,6 @@ namespace xllm_service { class Scheduler; class InstanceMgr; -llm::RequestOutput make_request_output( - const proto::DisaggStreamGeneration& request); - class XllmRpcServiceImpl final { public: XllmRpcServiceImpl(const Options& options, Scheduler* scheduler); diff --git a/xllm_service/scheduler/CMakeLists.txt b/xllm_service/scheduler/CMakeLists.txt index 1f00afa..c68f6a3 100644 --- a/xllm_service/scheduler/CMakeLists.txt +++ b/xllm_service/scheduler/CMakeLists.txt @@ -55,6 +55,7 @@ cc_library( DEPS :anthropic_adapter :anthropic_stream_encoder + :usage_proto_adapter :chat_template :deepseek_v4_cpp_chat_template :model_type diff --git a/xllm_service/scheduler/response_handler.cpp b/xllm_service/scheduler/response_handler.cpp index 34aaf45..de01687 100644 --- a/xllm_service/scheduler/response_handler.cpp +++ b/xllm_service/scheduler/response_handler.cpp @@ -20,6 +20,7 @@ limitations under the License. #include "common/anthropic_tracer.h" #include "http_service/anthropic_adapter.h" #include "http_service/anthropic_stream_encoder.h" +#include "http_service/usage_proto_adapter.h" #include "scheduler/xllm_chat_parse_bridge.h" #include "xllm/xllm/api_service/stream_output_parser.h" #include "xllm/xllm/api_service/utils.h" @@ -27,18 +28,6 @@ limitations under the License. namespace xllm_service { -void set_openai_usage(xllm::proto::Usage* proto_usage, - const llm::Usage& usage) { - CHECK(proto_usage != nullptr); - proto_usage->set_prompt_tokens(static_cast(usage.num_prompt_tokens)); - proto_usage->set_completion_tokens( - static_cast(usage.num_generated_tokens)); - proto_usage->set_total_tokens(static_cast(usage.num_total_tokens)); - auto* prompt_tokens_details = proto_usage->mutable_prompt_tokens_details(); - prompt_tokens_details->set_cached_tokens( - static_cast(usage.num_cached_tokens)); -} - namespace { size_t find_tool_start(const std::string& text) { @@ -408,7 +397,7 @@ bool ResponseHandler::send_delta_to_client( response.set_id(request_id); response.set_created(created_time); response.set_model(model); - set_openai_usage(response.mutable_usage(), usage); + *response.mutable_usage() = to_openai_usage_proto(usage); if (!call_data->write(response)) { return false; } @@ -484,7 +473,7 @@ bool ResponseHandler::send_delta_to_client( response.set_created(created_time); response.set_model(model); response.mutable_choices(); - set_openai_usage(response.mutable_usage(), usage); + *response.mutable_usage() = to_openai_usage_proto(usage); if (!call_data->write(response)) { return false; } @@ -764,7 +753,7 @@ bool ResponseHandler::send_result_to_client( // add usage statistics if (req_output.usage.has_value()) { const auto& usage = req_output.usage.value(); - set_openai_usage(response.mutable_usage(), usage); + *response.mutable_usage() = to_openai_usage_proto(usage); } return call_data->write_and_finish(response); @@ -807,7 +796,7 @@ bool ResponseHandler::send_result_to_client( // add usage statistics if (req_output.usage.has_value()) { const auto& usage = req_output.usage.value(); - set_openai_usage(response.mutable_usage(), usage); + *response.mutable_usage() = to_openai_usage_proto(usage); } return call_data->write_and_finish(response); diff --git a/xllm_service/scheduler/response_handler.h b/xllm_service/scheduler/response_handler.h index ec4ef15..2ae29db 100644 --- a/xllm_service/scheduler/response_handler.h +++ b/xllm_service/scheduler/response_handler.h @@ -34,8 +34,6 @@ namespace xllm_service { class AnthropicStreamEncoder; -void set_openai_usage(xllm::proto::Usage* proto_usage, const llm::Usage& usage); - struct ChatStreamParseState { std::unordered_set first_message_sent; std::shared_ptr stream_parser; From bb4243261c1a30d0431e8b481a4326cad6c2c017 Mon Sep 17 00:00:00 2001 From: phantomlei3 Date: Fri, 24 Jul 2026 15:14:26 +0800 Subject: [PATCH 3/5] build: align service with xllm main --- third_party/xllm | 2 +- xllm_service/proto/CMakeLists.txt | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/third_party/xllm b/third_party/xllm index 8039573..dfc4054 160000 --- a/third_party/xllm +++ b/third_party/xllm @@ -1 +1 @@ -Subproject commit 8039573ec9671e59e87927c883427c6ff975c053 +Subproject commit dfc40549644817783ae52f06ce49d5d42526c1e7 diff --git a/xllm_service/proto/CMakeLists.txt b/xllm_service/proto/CMakeLists.txt index b9fe77d..bff09d2 100644 --- a/xllm_service/proto/CMakeLists.txt +++ b/xllm_service/proto/CMakeLists.txt @@ -24,6 +24,7 @@ proto_library( ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/common.proto ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/completion.proto ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/sample.proto + ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/text_generation.proto ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/rec.proto ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/disagg_pd.proto ${CMAKE_SOURCE_DIR}/third_party/xllm/xllm/proto/xllm_service.proto From f5e7dff42cf936e78ad40cc6fba6f20bf2f25263 Mon Sep 17 00:00:00 2001 From: phantomlei3 Date: Fri, 24 Jul 2026 15:34:10 +0800 Subject: [PATCH 4/5] refactor: keep upstream cache usage naming --- .../http_service/anthropic_adapter_test.cpp | 2 +- .../anthropic_stream_encoder_test.cpp | 2 +- .../http_service/usage_proto_adapter_test.cpp | 19 +++++++++---------- third_party/xllm | 2 +- xllm_service/common/xllm/output.h | 2 +- .../http_service/usage_proto_adapter.cpp | 8 ++++---- xllm_service/proto/xllm_rpc_service.proto | 2 +- .../rpc_service/disagg_generation_adapter.cpp | 8 ++++---- xllm_service/rpc_service/rpc_service_test.cpp | 14 ++++++-------- 9 files changed, 28 insertions(+), 31 deletions(-) diff --git a/tests/xllm_service/http_service/anthropic_adapter_test.cpp b/tests/xllm_service/http_service/anthropic_adapter_test.cpp index 4d38c5b..839aba2 100644 --- a/tests/xllm_service/http_service/anthropic_adapter_test.cpp +++ b/tests/xllm_service/http_service/anthropic_adapter_test.cpp @@ -807,7 +807,7 @@ TEST(AnthropicAdapterTest, BuildsNonStreamAnthropicJson) { usage.num_prompt_tokens = 3; usage.num_generated_tokens = 4; usage.num_total_tokens = 7; - usage.num_prefix_cache_hit_tokens = 2; + usage.num_cached_tokens = 2; output.usage = usage; xllm::proto::AnthropicMessagesResponse response; diff --git a/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp b/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp index 634fc5b..9f55018 100644 --- a/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp +++ b/tests/xllm_service/http_service/anthropic_stream_encoder_test.cpp @@ -243,7 +243,7 @@ TEST(AnthropicStreamEncoderTest, FinishClosesBlockAndEmitsMessageDeltaStop) { llm::Usage usage; usage.num_prompt_tokens = 3; usage.num_generated_tokens = 5; - usage.num_prefix_cache_hit_tokens = 2; + usage.num_cached_tokens = 2; final.usage = usage; std::vector done; diff --git a/tests/xllm_service/http_service/usage_proto_adapter_test.cpp b/tests/xllm_service/http_service/usage_proto_adapter_test.cpp index 32349ac..b904308 100644 --- a/tests/xllm_service/http_service/usage_proto_adapter_test.cpp +++ b/tests/xllm_service/http_service/usage_proto_adapter_test.cpp @@ -27,22 +27,22 @@ struct UsageCase { size_t num_prompt_tokens; size_t num_generated_tokens; size_t num_total_tokens; - size_t num_prefix_cache_hit_tokens; + size_t num_cached_tokens; }; const std::vector kUsageCases = { {/*num_prompt_tokens=*/8, /*num_generated_tokens=*/0, /*num_total_tokens=*/8, - /*num_prefix_cache_hit_tokens=*/0}, + /*num_cached_tokens=*/0}, {/*num_prompt_tokens=*/8, /*num_generated_tokens=*/2, /*num_total_tokens=*/10, - /*num_prefix_cache_hit_tokens=*/3}, + /*num_cached_tokens=*/3}, {/*num_prompt_tokens=*/8, /*num_generated_tokens=*/2, /*num_total_tokens=*/10, - /*num_prefix_cache_hit_tokens=*/8}, + /*num_cached_tokens=*/8}, }; llm::Usage make_usage(const UsageCase& usage_case) { @@ -50,7 +50,7 @@ llm::Usage make_usage(const UsageCase& usage_case) { usage.num_prompt_tokens = usage_case.num_prompt_tokens; usage.num_generated_tokens = usage_case.num_generated_tokens; usage.num_total_tokens = usage_case.num_total_tokens; - usage.num_prefix_cache_hit_tokens = usage_case.num_prefix_cache_hit_tokens; + usage.num_cached_tokens = usage_case.num_cached_tokens; return usage; } @@ -65,7 +65,7 @@ TEST(UsageProtoAdapterTest, ConvertsPrefixCacheHitsToOpenAIPromptTokenDetails) { ASSERT_TRUE(proto_usage.has_prompt_tokens_details()); EXPECT_TRUE(proto_usage.prompt_tokens_details().has_cached_tokens()); EXPECT_EQ(proto_usage.prompt_tokens_details().cached_tokens(), - usage_case.num_prefix_cache_hit_tokens); + usage_case.num_cached_tokens); } } @@ -75,11 +75,10 @@ TEST(UsageProtoAdapterTest, xllm::proto::AnthropicUsage proto_usage = to_anthropic_usage_proto(make_usage(usage_case)); - EXPECT_EQ( - proto_usage.input_tokens(), - usage_case.num_prompt_tokens - usage_case.num_prefix_cache_hit_tokens); + EXPECT_EQ(proto_usage.input_tokens(), + usage_case.num_prompt_tokens - usage_case.num_cached_tokens); EXPECT_EQ(proto_usage.cache_read_input_tokens(), - usage_case.num_prefix_cache_hit_tokens); + usage_case.num_cached_tokens); EXPECT_EQ(proto_usage.output_tokens(), usage_case.num_generated_tokens); EXPECT_EQ( proto_usage.input_tokens() + proto_usage.cache_read_input_tokens(), diff --git a/third_party/xllm b/third_party/xllm index dfc4054..86cfaa5 160000 --- a/third_party/xllm +++ b/third_party/xllm @@ -1 +1 @@ -Subproject commit dfc40549644817783ae52f06ce49d5d42526c1e7 +Subproject commit 86cfaa572552fb3c655b90c1a74571e4c5c2b667 diff --git a/xllm_service/common/xllm/output.h b/xllm_service/common/xllm/output.h index 80ee71d..908fcb5 100644 --- a/xllm_service/common/xllm/output.h +++ b/xllm_service/common/xllm/output.h @@ -49,7 +49,7 @@ struct Usage { size_t num_total_tokens = 0; // the number of prompt tokens served from prefix cache. - size_t num_prefix_cache_hit_tokens = 0; + size_t num_cached_tokens = 0; }; struct LogProbData { diff --git a/xllm_service/http_service/usage_proto_adapter.cpp b/xllm_service/http_service/usage_proto_adapter.cpp index 6731c89..416a90c 100644 --- a/xllm_service/http_service/usage_proto_adapter.cpp +++ b/xllm_service/http_service/usage_proto_adapter.cpp @@ -26,18 +26,18 @@ xllm::proto::Usage to_openai_usage_proto(const llm::Usage& usage) { static_cast(usage.num_generated_tokens)); proto_usage.set_total_tokens(static_cast(usage.num_total_tokens)); proto_usage.mutable_prompt_tokens_details()->set_cached_tokens( - static_cast(usage.num_prefix_cache_hit_tokens)); + static_cast(usage.num_cached_tokens)); return proto_usage; } xllm::proto::AnthropicUsage to_anthropic_usage_proto(const llm::Usage& usage) { xllm::proto::AnthropicUsage proto_usage; - proto_usage.set_input_tokens(static_cast( - usage.num_prompt_tokens - usage.num_prefix_cache_hit_tokens)); + proto_usage.set_input_tokens( + static_cast(usage.num_prompt_tokens - usage.num_cached_tokens)); proto_usage.set_output_tokens( static_cast(usage.num_generated_tokens)); proto_usage.set_cache_read_input_tokens( - static_cast(usage.num_prefix_cache_hit_tokens)); + static_cast(usage.num_cached_tokens)); return proto_usage; } diff --git a/xllm_service/proto/xllm_rpc_service.proto b/xllm_service/proto/xllm_rpc_service.proto index 185bfcf..c53e31f 100644 --- a/xllm_service/proto/xllm_rpc_service.proto +++ b/xllm_service/proto/xllm_rpc_service.proto @@ -100,7 +100,7 @@ message OutputUsage { // the total number of tokens used in the request (prompt + completion). int32 num_total_tokens = 3; // the number of prompt tokens served from prefix cache. - int32 num_prefix_cache_hit_tokens = 4; + int32 num_cached_tokens = 4; } message LogProbData { diff --git a/xllm_service/rpc_service/disagg_generation_adapter.cpp b/xllm_service/rpc_service/disagg_generation_adapter.cpp index 805940b..9586488 100644 --- a/xllm_service/rpc_service/disagg_generation_adapter.cpp +++ b/xllm_service/rpc_service/disagg_generation_adapter.cpp @@ -30,10 +30,10 @@ RequestOutputConversionResult invalid_usage(std::string message) { RequestOutputConversionResult validate_usage(const proto::OutputUsage& usage) { if (usage.num_prompt_tokens() < 0 || usage.num_generated_tokens() < 0 || - usage.num_total_tokens() < 0 || usage.num_prefix_cache_hit_tokens() < 0) { + usage.num_total_tokens() < 0 || usage.num_cached_tokens() < 0) { return invalid_usage("token counts must be non-negative"); } - if (usage.num_prefix_cache_hit_tokens() > usage.num_prompt_tokens()) { + if (usage.num_cached_tokens() > usage.num_prompt_tokens()) { return invalid_usage( "prefix cache hit tokens must not exceed prompt tokens"); } @@ -75,8 +75,8 @@ RequestOutputConversionResult request_output_from_disagg_generation( static_cast(generation.usage().num_generated_tokens()); usage.num_total_tokens = static_cast(generation.usage().num_total_tokens()); - usage.num_prefix_cache_hit_tokens = - static_cast(generation.usage().num_prefix_cache_hit_tokens()); + usage.num_cached_tokens = + static_cast(generation.usage().num_cached_tokens()); request_output.usage = std::move(usage); } request_output.finished_on_prefill_instance = diff --git a/xllm_service/rpc_service/rpc_service_test.cpp b/xllm_service/rpc_service/rpc_service_test.cpp index 5f662e6..8943776 100644 --- a/xllm_service/rpc_service/rpc_service_test.cpp +++ b/xllm_service/rpc_service/rpc_service_test.cpp @@ -33,8 +33,7 @@ proto::DisaggStreamGeneration make_generation(int32_t num_prompt_tokens, generation.mutable_usage()->set_num_prompt_tokens(num_prompt_tokens); generation.mutable_usage()->set_num_generated_tokens(num_generated_tokens); generation.mutable_usage()->set_num_total_tokens(num_total_tokens); - generation.mutable_usage()->set_num_prefix_cache_hit_tokens( - num_cache_hit_tokens); + generation.mutable_usage()->set_num_cached_tokens(num_cache_hit_tokens); return generation; } @@ -42,10 +41,9 @@ TEST(DisaggGenerationAdapterTest, PrefixCacheHitTokensRemainWireCompatibleAtFieldFour) { const google::protobuf::FieldDescriptor* sender_field = xllm::proto::OutputUsage::descriptor()->FindFieldByName( - "num_prefix_cache_hit_tokens"); + "num_cached_tokens"); const google::protobuf::FieldDescriptor* receiver_field = - proto::OutputUsage::descriptor()->FindFieldByName( - "num_prefix_cache_hit_tokens"); + proto::OutputUsage::descriptor()->FindFieldByName("num_cached_tokens"); ASSERT_NE(sender_field, nullptr); ASSERT_NE(receiver_field, nullptr); @@ -57,11 +55,11 @@ TEST(DisaggGenerationAdapterTest, sender_usage.set_num_prompt_tokens(8); sender_usage.set_num_generated_tokens(2); sender_usage.set_num_total_tokens(10); - sender_usage.set_num_prefix_cache_hit_tokens(6); + sender_usage.set_num_cached_tokens(6); proto::OutputUsage receiver_usage; ASSERT_TRUE(receiver_usage.ParseFromString(sender_usage.SerializeAsString())); - EXPECT_EQ(receiver_usage.num_prefix_cache_hit_tokens(), 6); + EXPECT_EQ(receiver_usage.num_cached_tokens(), 6); } TEST(DisaggGenerationAdapterTest, RejectsNegativeTokenCounts) { @@ -160,7 +158,7 @@ TEST(DisaggGenerationAdapterTest, ConvertsCompleteValidGeneration) { EXPECT_EQ(output.usage->num_prompt_tokens, 8u); EXPECT_EQ(output.usage->num_generated_tokens, 2u); EXPECT_EQ(output.usage->num_total_tokens, 10u); - EXPECT_EQ(output.usage->num_prefix_cache_hit_tokens, 6u); + EXPECT_EQ(output.usage->num_cached_tokens, 6u); ASSERT_EQ(output.outputs.size(), 1u); const llm::SequenceOutput& converted_sequence = output.outputs.front(); From 672f8d65087284d03010d9469beea1e091f83de2 Mon Sep 17 00:00:00 2001 From: phantomlei3 Date: Fri, 24 Jul 2026 16:57:20 +0800 Subject: [PATCH 5/5] test: move rpc service tests under tests --- tests/xllm_service/CMakeLists.txt | 1 + tests/xllm_service/rpc_service/CMakeLists.txt | 15 +++++++++++++++ .../rpc_service/rpc_service_test.cpp | 3 +-- xllm_service/rpc_service/CMakeLists.txt | 16 ---------------- 4 files changed, 17 insertions(+), 18 deletions(-) create mode 100644 tests/xllm_service/rpc_service/CMakeLists.txt rename {xllm_service => tests/xllm_service}/rpc_service/rpc_service_test.cpp (98%) diff --git a/tests/xllm_service/CMakeLists.txt b/tests/xllm_service/CMakeLists.txt index ac87b45..69a63b8 100644 --- a/tests/xllm_service/CMakeLists.txt +++ b/tests/xllm_service/CMakeLists.txt @@ -1,5 +1,6 @@ add_subdirectory(chat_template) add_subdirectory(common) add_subdirectory(http_service) +add_subdirectory(rpc_service) add_subdirectory(scheduler) add_subdirectory(tokenizer) diff --git a/tests/xllm_service/rpc_service/CMakeLists.txt b/tests/xllm_service/rpc_service/CMakeLists.txt new file mode 100644 index 0000000..da8465f --- /dev/null +++ b/tests/xllm_service/rpc_service/CMakeLists.txt @@ -0,0 +1,15 @@ +include(cc_test) + +cc_test( + NAME + xllm_rpc_service_test + SRCS + rpc_service_test.cpp + DEPS + :disagg_generation_adapter + :xllm_rpc_service + gflags::gflags + glog::glog + GTest::gtest_main + proto_xllm +) diff --git a/xllm_service/rpc_service/rpc_service_test.cpp b/tests/xllm_service/rpc_service/rpc_service_test.cpp similarity index 98% rename from xllm_service/rpc_service/rpc_service_test.cpp rename to tests/xllm_service/rpc_service/rpc_service_test.cpp index 8943776..14dcba7 100644 --- a/xllm_service/rpc_service/rpc_service_test.cpp +++ b/tests/xllm_service/rpc_service/rpc_service_test.cpp @@ -37,8 +37,7 @@ proto::DisaggStreamGeneration make_generation(int32_t num_prompt_tokens, return generation; } -TEST(DisaggGenerationAdapterTest, - PrefixCacheHitTokensRemainWireCompatibleAtFieldFour) { +TEST(DisaggGenerationAdapterTest, CachedTokensRemainWireCompatibleAtFieldFour) { const google::protobuf::FieldDescriptor* sender_field = xllm::proto::OutputUsage::descriptor()->FindFieldByName( "num_cached_tokens"); diff --git a/xllm_service/rpc_service/CMakeLists.txt b/xllm_service/rpc_service/CMakeLists.txt index 8afa0b4..203fc60 100644 --- a/xllm_service/rpc_service/CMakeLists.txt +++ b/xllm_service/rpc_service/CMakeLists.txt @@ -1,6 +1,5 @@ include(cc_binary) include(cc_library) -include(cc_test) cc_library( NAME @@ -35,21 +34,6 @@ cc_library( ) target_link_libraries(xllm_rpc_service PRIVATE brpc-static) -cc_binary( - NAME - xllm_rpc_service_test - SRCS - rpc_service_test.cpp - DEPS - :disagg_generation_adapter - :xllm_rpc_service - gflags::gflags - glog::glog - GTest::gtest_main - proto_xllm -) -add_test(NAME XllmRpcServiceTest COMMAND xllm_rpc_service_test) - cc_binary( NAME xllm_rpc_serving