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/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 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/http_service/usage_proto_adapter_test.cpp b/tests/xllm_service/http_service/usage_proto_adapter_test.cpp new file mode 100644 index 0000000..b904308 --- /dev/null +++ b/tests/xllm_service/http_service/usage_proto_adapter_test.cpp @@ -0,0 +1,91 @@ +/* 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_cached_tokens; +}; + +const std::vector kUsageCases = { + {/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/0, + /*num_total_tokens=*/8, + /*num_cached_tokens=*/0}, + {/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/10, + /*num_cached_tokens=*/3}, + {/*num_prompt_tokens=*/8, + /*num_generated_tokens=*/2, + /*num_total_tokens=*/10, + /*num_cached_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_cached_tokens = usage_case.num_cached_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_cached_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_cached_tokens); + EXPECT_EQ(proto_usage.cache_read_input_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(), + usage_case.num_prompt_tokens); + EXPECT_FALSE(proto_usage.has_cache_creation_input_tokens()); + } +} + +} // namespace +} // namespace xllm_service 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/tests/xllm_service/rpc_service/rpc_service_test.cpp b/tests/xllm_service/rpc_service/rpc_service_test.cpp new file mode 100644 index 0000000..14dcba7 --- /dev/null +++ b/tests/xllm_service/rpc_service/rpc_service_test.cpp @@ -0,0 +1,185 @@ +/* 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 +#include + +#include "disagg_pd.pb.h" +#include "rpc_service/disagg_generation_adapter.h" + +namespace xllm_service { +namespace { + +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_cached_tokens(num_cache_hit_tokens); + return generation; +} + +TEST(DisaggGenerationAdapterTest, CachedTokensRemainWireCompatibleAtFieldFour) { + const google::protobuf::FieldDescriptor* sender_field = + xllm::proto::OutputUsage::descriptor()->FindFieldByName( + "num_cached_tokens"); + const google::protobuf::FieldDescriptor* receiver_field = + proto::OutputUsage::descriptor()->FindFieldByName("num_cached_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_cached_tokens(6); + + proto::OutputUsage receiver_usage; + ASSERT_TRUE(receiver_usage.ParseFromString(sender_usage.SerializeAsString())); + EXPECT_EQ(receiver_usage.num_cached_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); + + 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("non-negative"), std::string::npos); + EXPECT_FALSE(result.output.has_value()); +} + +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); + + 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_prompt_tokens, 8u); + EXPECT_EQ(output.usage->num_generated_tokens, 2u); + EXPECT_EQ(output.usage->num_total_tokens, 10u); + EXPECT_EQ(output.usage->num_cached_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); +} + +} // namespace +} // namespace xllm_service diff --git a/third_party/xllm b/third_party/xllm index adf07f7..86cfaa5 160000 --- a/third_party/xllm +++ b/third_party/xllm @@ -1 +1 @@ -Subproject commit adf07f70cdeaa49ba8a7ed2dd0db1d0f2b718539 +Subproject commit 86cfaa572552fb3c655b90c1a74571e4c5c2b667 diff --git a/xllm_service/common/xllm/output.h b/xllm_service/common/xllm/output.h index 5c1327a..908fcb5 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 @@ -46,6 +47,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/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 64a1d1f..ce28079 100644 --- a/xllm_service/http_service/anthropic_adapter.cpp +++ b/xllm_service/http_service/anthropic_adapter.cpp @@ -30,8 +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 { + namespace { thread_local llm::ShortUUID short_uuid; @@ -375,11 +377,8 @@ 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)); + *response->mutable_usage() = + to_anthropic_usage_proto(request_output.usage.value()); } bool normalize_stream_event_json(const xllm::proto::AnthropicStreamEvent& event, @@ -410,6 +409,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_stream_encoder.cpp b/xllm_service/http_service/anthropic_stream_encoder.cpp index d39b97f..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,11 +212,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)); + *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..416a90c --- /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_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_cached_tokens)); + proto_usage.set_output_tokens( + static_cast(usage.num_generated_tokens)); + proto_usage.set_cache_read_input_tokens( + static_cast(usage.num_cached_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/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 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/CMakeLists.txt b/xllm_service/rpc_service/CMakeLists.txt index 5e4ef79..203fc60 100644 --- a/xllm_service/rpc_service/CMakeLists.txt +++ b/xllm_service/rpc_service/CMakeLists.txt @@ -1,6 +1,17 @@ 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 @@ -11,6 +22,7 @@ cc_library( service.cpp DEPS :common + :disagg_generation_adapter :scheduler absl::random_random absl::strings @@ -22,19 +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 - :xllm_rpc_service - gflags::gflags - glog::glog - GTest::gtest_main -) -add_test(NAME XllmRpcServiceTest COMMAND xllm_rpc_service_test) - cc_binary( NAME xllm_rpc_serving 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..9586488 --- /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_cached_tokens() < 0) { + return invalid_usage("token counts must be non-negative"); + } + if (usage.num_cached_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_cached_tokens = + static_cast(generation.usage().num_cached_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 deleted file mode 100644 index 4661777..0000000 --- a/xllm_service/rpc_service/rpc_service_test.cpp +++ /dev/null @@ -1,69 +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 - -#include "rpc_service/service.h" - -namespace xllm_service::test { - -class XllmRpcServiceTest : public ::testing::Test { - protected: - void SetUp() override { google::InitGoogleLogging("XllmRpcServiceTest"); } - - void TearDown() override { google::ShutdownGoogleLogging(); } -}; -// 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 diff --git a/xllm_service/rpc_service/service.cpp b/xllm_service/rpc_service/service.cpp index 4f9b842..2dfd5b4 100644 --- a/xllm_service/rpc_service/service.cpp +++ b/xllm_service/rpc_service/service.cpp @@ -22,6 +22,7 @@ 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 { @@ -145,64 +146,18 @@ 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()); + 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; } - 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)); - } - - resp->mutable_all_status()->Add()->set_ok( - xllm_rpc_service_impl_->handle_generation(request_output)); + status->set_ok( + xllm_rpc_service_impl_->handle_generation(conversion.output.value())); } } 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 ba293da..de01687 100644 --- a/xllm_service/scheduler/response_handler.cpp +++ b/xllm_service/scheduler/response_handler.cpp @@ -20,12 +20,14 @@ 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" #include "xllm/xllm/function_call/function_call_parser.h" namespace xllm_service { + namespace { size_t find_tool_start(const std::string& text) { @@ -395,12 +397,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)); + *response.mutable_usage() = to_openai_usage_proto(usage); if (!call_data->write(response)) { return false; } @@ -476,12 +473,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)); + *response.mutable_usage() = to_openai_usage_proto(usage); if (!call_data->write(response)) { return false; } @@ -761,12 +753,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)); + *response.mutable_usage() = to_openai_usage_proto(usage); } return call_data->write_and_finish(response); @@ -809,12 +796,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)); + *response.mutable_usage() = to_openai_usage_proto(usage); } return call_data->write_and_finish(response);