ref(upstream): FULL TREE — Deep-Spark xllm (1470) + ds_vllm csrc/models (703)

Replaces cherry-picked upstream_ref with complete source trees.

xllm/ — Iluvatar official C++ inference engine (15MB, 1470 files)
  Complete: kernels → layers → models → runtime → scheduler → api
  Excluded: .git, binary images, third_party submodule checkouts

ds_vllm/ — Iluvatar official vllm fork (8MB, 703 files)
  Included: csrc/ (ALL CUDA kernels), fused_moe/, qwen3_5 model, _custom_ops
  Excluded: tests, benchmarks, docs, examples (not needed for reference)

Critical call chains now fully traceable:
  MoE: moe_topk_softmax_kernels.cuh → ixformer.h → fused_moe.cpp → layer
  GDN: qwen3_gated_delta_net_base.cpp → qwen3_5_gated_delta_net.cpp
  Attention: ixformer.h → xllm_paged_attention → attention.cpp

upstream_ref/xllm/tests/CMakeLists.txt

@@ -0,0 +1,7 @@
+if(NOT BUILD_TESTING)
+  return()
+endif()
+
+add_subdirectory(api_service)
+add_subdirectory(core)
+add_subdirectory(function_call)

upstream_ref/xllm/tests/api_service/CMakeLists.txt

@@ -0,0 +1,68 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    anthropic_protocol_test
+  SRCS
+    anthropic_protocol_test.cpp
+  DEPS
+    proto::xllm_proto
+    GTest::gtest
+    GTest::gtest_main
+    nlohmann_json::nlohmann_json
+    glog::glog
+)
+
+# Integration test for Anthropic service - requires running server.
+# Tests are prefixed with DISABLED_ in code to skip by default.
+# Run manually: ./anthropic_service_test --gtest_also_run_disabled_tests
+cc_test(
+  NAME
+    anthropic_service_test
+  SRCS
+    anthropic_service_test.cpp
+  DEPS
+    GTest::gtest
+    GTest::gtest_main
+    nlohmann_json::nlohmann_json
+    glog::glog
+)
+target_link_libraries(anthropic_service_test PRIVATE brpc leveldb::leveldb OpenSSL::SSL OpenSSL::Crypto protobuf::libprotobuf)
+add_dependencies(anthropic_service_test brpc-static)
+
+cc_test(
+  NAME
+    api_service_test
+  SRCS
+    chat_json_parser_test.cpp
+  DEPS
+    api_service
+    GTest::gtest_main
+    nlohmann_json::nlohmann_json
+)
+
+cc_test(
+  NAME
+    sample_service_impl_test
+  SRCS
+    sample_service_impl_test.cpp
+  DEPS
+    api_service
+    GTest::gtest_main
+)
+
+# Integration smoke tests for OpenAI-compatible endpoints. These require a
+# running server and are disabled by default in code.
+cc_test(
+  NAME
+    openai_service_test
+  SRCS
+    openai_service_test.cpp
+  DEPS
+    GTest::gtest
+    GTest::gtest_main
+    nlohmann_json::nlohmann_json
+    glog::glog
+)
+target_link_libraries(openai_service_test PRIVATE brpc leveldb::leveldb OpenSSL::SSL OpenSSL::Crypto protobuf::libprotobuf)
+add_dependencies(openai_service_test brpc-static)

upstream_ref/xllm/tests/api_service/anthropic_protocol_test.cpp

@@ -0,0 +1,600 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <glog/logging.h>
+#include <google/protobuf/util/json_util.h>
+#include <gtest/gtest.h>
+
+#include <nlohmann/json.hpp>
+#include <string>
+#include <vector>
+
+#include "anthropic.pb.h"
+
+namespace xllm {
+namespace {
+
+using proto::AnthropicContentBlock;
+using proto::AnthropicContentBlockList;
+using proto::AnthropicMessage;
+using proto::AnthropicMessagesRequest;
+using proto::AnthropicMessagesResponse;
+using proto::AnthropicTool;
+using proto::AnthropicToolChoice;
+using proto::AnthropicUsage;
+
+class AnthropicProtocolTest : public ::testing::Test {
+ protected:
+  void SetUp() override {}
+
+  // Helper to convert nlohmann::json to google::protobuf::Struct
+  static google::protobuf::Struct json_to_struct(const nlohmann::json& j) {
+    google::protobuf::Struct pb_struct;
+    std::string json_str = j.dump();
+    google::protobuf::util::JsonStringToMessage(json_str, &pb_struct);
+    return pb_struct;
+  }
+
+  // Helper to convert google::protobuf::Struct to nlohmann::json
+  static nlohmann::json struct_to_json(
+      const google::protobuf::Struct& pb_struct) {
+    std::string json_str;
+    google::protobuf::util::JsonPrintOptions options;
+    options.preserve_proto_field_names = true;
+    google::protobuf::util::MessageToJsonString(pb_struct, &json_str, options);
+    return nlohmann::json::parse(json_str);
+  }
+};
+
+// Test AnthropicContentBlock with text content
+TEST_F(AnthropicProtocolTest, AnthropicContentBlockText) {
+  AnthropicContentBlock content_block;
+  content_block.set_type("text");
+  content_block.set_text("Hello, world!");
+
+  EXPECT_EQ(content_block.type(), "text");
+  EXPECT_EQ(content_block.text(), "Hello, world!");
+  EXPECT_FALSE(content_block.has_source());
+  EXPECT_FALSE(content_block.has_id());
+  EXPECT_FALSE(content_block.has_name());
+  EXPECT_FALSE(content_block.has_input());
+  EXPECT_EQ(content_block.tool_result_content_case(),
+            AnthropicContentBlock::TOOL_RESULT_CONTENT_NOT_SET);
+  EXPECT_FALSE(content_block.has_is_error());
+}
+
+// Test AnthropicContentBlock with tool use content
+TEST_F(AnthropicProtocolTest, AnthropicContentBlockToolUse) {
+  AnthropicContentBlock content_block;
+  content_block.set_type("tool_use");
+  content_block.set_id("toolu_123");
+  content_block.set_name("get_weather");
+
+  nlohmann::json input_json = {{"location", "Paris"}, {"unit", "celsius"}};
+  *content_block.mutable_input() = json_to_struct(input_json);
+
+  EXPECT_EQ(content_block.type(), "tool_use");
+  EXPECT_EQ(content_block.id(), "toolu_123");
+  EXPECT_EQ(content_block.name(), "get_weather");
+
+  // Verify input
+  auto parsed_input = struct_to_json(content_block.input());
+  EXPECT_EQ(parsed_input["location"], "Paris");
+  EXPECT_EQ(parsed_input["unit"], "celsius");
+}
+
+// Test AnthropicContentBlock with tool result content
+TEST_F(AnthropicProtocolTest, AnthropicContentBlockToolResult) {
+  AnthropicContentBlock content_block;
+  content_block.set_type("tool_result");
+  content_block.set_id("call_123");
+  content_block.set_content_string("The weather in Paris is sunny.");
+  content_block.set_is_error(false);
+
+  EXPECT_EQ(content_block.type(), "tool_result");
+  EXPECT_EQ(content_block.id(), "call_123");
+  EXPECT_EQ(content_block.content_string(), "The weather in Paris is sunny.");
+  EXPECT_FALSE(content_block.is_error());
+}
+
+// Test AnthropicMessage for user role
+TEST_F(AnthropicProtocolTest, AnthropicMessageUser) {
+  AnthropicMessage message;
+  message.set_role("user");
+  message.set_content_string("Hello, assistant!");
+
+  EXPECT_EQ(message.role(), "user");
+  EXPECT_EQ(message.content_string(), "Hello, assistant!");
+}
+
+// Test AnthropicMessage for assistant role with text content
+TEST_F(AnthropicProtocolTest, AnthropicMessageAssistantText) {
+  AnthropicMessage message;
+  message.set_role("assistant");
+
+  // Create content block list
+  auto* content_blocks = message.mutable_content_blocks();
+  auto* text_block = content_blocks->add_blocks();
+  text_block->set_type("text");
+  text_block->set_text("Hello, user!");
+
+  EXPECT_EQ(message.role(), "assistant");
+  EXPECT_TRUE(message.has_content_blocks());
+  EXPECT_EQ(message.content_blocks().blocks_size(), 1);
+  EXPECT_EQ(message.content_blocks().blocks(0).type(), "text");
+  EXPECT_EQ(message.content_blocks().blocks(0).text(), "Hello, user!");
+}
+
+// Test AnthropicTool model
+TEST_F(AnthropicProtocolTest, AnthropicTool) {
+  nlohmann::json tool_schema = {
+      {"type", "object"},
+      {"properties",
+       {{"location", {{"type", "string"}}},
+        {"unit", {{"type", "string"}, {"enum", {"celsius", "fahrenheit"}}}}}},
+      {"required", {"location"}}};
+
+  AnthropicTool tool;
+  tool.set_name("get_weather");
+  tool.set_description("Get the current weather for a location");
+  *tool.mutable_input_schema() = json_to_struct(tool_schema);
+
+  EXPECT_EQ(tool.name(), "get_weather");
+  EXPECT_EQ(tool.description(), "Get the current weather for a location");
+
+  // Verify input_schema
+  auto parsed_schema = struct_to_json(tool.input_schema());
+  EXPECT_EQ(parsed_schema["type"], "object");
+  EXPECT_TRUE(parsed_schema.contains("properties"));
+  EXPECT_TRUE(parsed_schema["properties"].contains("location"));
+}
+
+// Test AnthropicToolChoice with auto type
+TEST_F(AnthropicProtocolTest, AnthropicToolChoiceAuto) {
+  AnthropicToolChoice tool_choice;
+  tool_choice.set_type("auto");
+
+  EXPECT_EQ(tool_choice.type(), "auto");
+  EXPECT_FALSE(tool_choice.has_name());
+}
+
+// Test AnthropicToolChoice with specific tool
+TEST_F(AnthropicProtocolTest, AnthropicToolChoiceSpecific) {
+  AnthropicToolChoice tool_choice;
+  tool_choice.set_type("tool");
+  tool_choice.set_name("get_weather");
+
+  EXPECT_EQ(tool_choice.type(), "tool");
+  EXPECT_EQ(tool_choice.name(), "get_weather");
+}
+
+// Test basic AnthropicMessagesRequest
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestBasic) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(100);
+
+  // Add message
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("What is the weather like today?");
+
+  EXPECT_EQ(request.model(), "my_model");
+  EXPECT_EQ(request.messages_size(), 1);
+  EXPECT_EQ(request.messages(0).role(), "user");
+  EXPECT_EQ(request.max_tokens(), 100);
+  EXPECT_FALSE(request.has_stream());
+  EXPECT_FALSE(request.has_system_string());
+  EXPECT_FALSE(request.has_temperature());
+}
+
+// Test AnthropicMessagesRequest with system prompt
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestWithSystem) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(100);
+  request.set_system_string("You are a helpful weather assistant.");
+
+  // Add message
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("What is the weather like today?");
+
+  EXPECT_EQ(request.system_string(), "You are a helpful weather assistant.");
+}
+
+// Test AnthropicMessagesRequest with tools
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestWithTools) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(100);
+
+  // Add message
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("What is the weather in Paris?");
+
+  // Add tool
+  nlohmann::json tool_schema = {
+      {"type", "object"},
+      {"properties",
+       {{"location", {{"type", "string"}}},
+        {"unit", {{"type", "string"}, {"enum", {"celsius", "fahrenheit"}}}}}},
+      {"required", {"location"}}};
+
+  auto* tool = request.add_tools();
+  tool->set_name("get_weather");
+  tool->set_description("Get the current weather for a location");
+  *tool->mutable_input_schema() = json_to_struct(tool_schema);
+
+  // Set tool choice
+  auto* tool_choice = request.mutable_tool_choice();
+  tool_choice->set_type("auto");
+
+  EXPECT_EQ(request.tools_size(), 1);
+  EXPECT_EQ(request.tools(0).name(), "get_weather");
+  EXPECT_EQ(request.tool_choice().type(), "auto");
+}
+
+// Test AnthropicMessagesRequest with streaming
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestStreaming) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(1000);
+  request.set_stream(true);
+
+  // Add message
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("Tell me a story.");
+
+  EXPECT_TRUE(request.stream());
+}
+
+// Test AnthropicMessagesRequest with all parameters
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestAllParams) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(100);
+  request.set_temperature(0.7f);
+  request.set_top_p(0.9f);
+  request.set_top_k(50);
+
+  // Add message
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("Hello");
+
+  EXPECT_FLOAT_EQ(request.temperature(), 0.7f);
+  EXPECT_FLOAT_EQ(request.top_p(), 0.9f);
+  EXPECT_EQ(request.top_k(), 50);
+}
+
+// Test AnthropicUsage model
+TEST_F(AnthropicProtocolTest, AnthropicUsage) {
+  AnthropicUsage usage;
+  usage.set_input_tokens(50);
+  usage.set_output_tokens(75);
+
+  EXPECT_EQ(usage.input_tokens(), 50);
+  EXPECT_EQ(usage.output_tokens(), 75);
+  EXPECT_FALSE(usage.has_cache_creation_input_tokens());
+  EXPECT_FALSE(usage.has_cache_read_input_tokens());
+}
+
+// Test AnthropicUsage model with cache tokens
+TEST_F(AnthropicProtocolTest, AnthropicUsageWithCache) {
+  AnthropicUsage usage;
+  usage.set_input_tokens(50);
+  usage.set_output_tokens(75);
+  usage.set_cache_creation_input_tokens(25);
+  usage.set_cache_read_input_tokens(15);
+
+  EXPECT_EQ(usage.input_tokens(), 50);
+  EXPECT_EQ(usage.output_tokens(), 75);
+  EXPECT_EQ(usage.cache_creation_input_tokens(), 25);
+  EXPECT_EQ(usage.cache_read_input_tokens(), 15);
+}
+
+// Test AnthropicMessagesResponse model
+TEST_F(AnthropicProtocolTest, AnthropicMessagesResponse) {
+  AnthropicMessagesResponse response;
+  response.set_id("msg_123");
+  response.set_type("message");
+  response.set_role("assistant");
+  response.set_model("my_model");
+  response.set_stop_reason("end_turn");
+
+  // Add content block
+  auto* content_block = response.add_content();
+  content_block->set_type("text");
+  content_block->set_text("Hello, user!");
+
+  // Set usage
+  auto* usage = response.mutable_usage();
+  usage->set_input_tokens(10);
+  usage->set_output_tokens(20);
+
+  EXPECT_EQ(response.id(), "msg_123");
+  EXPECT_EQ(response.type(), "message");
+  EXPECT_EQ(response.role(), "assistant");
+  EXPECT_EQ(response.content_size(), 1);
+  EXPECT_EQ(response.content(0).type(), "text");
+  EXPECT_EQ(response.content(0).text(), "Hello, user!");
+  EXPECT_EQ(response.model(), "my_model");
+  EXPECT_EQ(response.stop_reason(), "end_turn");
+  EXPECT_EQ(response.usage().input_tokens(), 10);
+  EXPECT_EQ(response.usage().output_tokens(), 20);
+}
+
+// Test AnthropicMessagesResponse with specific ID
+TEST_F(AnthropicProtocolTest, AnthropicMessagesResponseAutoId) {
+  AnthropicMessagesResponse response;
+  response.set_id("msg_test123");
+  response.set_type("message");
+  response.set_role("assistant");
+  response.set_model("my_model");
+
+  // Add content block
+  auto* content_block = response.add_content();
+  content_block->set_type("text");
+  content_block->set_text("Hello!");
+
+  // Set usage
+  auto* usage = response.mutable_usage();
+  usage->set_input_tokens(5);
+  usage->set_output_tokens(10);
+
+  EXPECT_EQ(response.id(), "msg_test123");
+  // Check that ID starts with "msg_"
+  EXPECT_EQ(response.id().substr(0, 4), "msg_");
+}
+
+// Test AnthropicTool with default input_schema type
+TEST_F(AnthropicProtocolTest, AnthropicToolInputSchemaDefault) {
+  nlohmann::json tool_schema = {
+      {"properties", {{"param", {{"type", "string"}}}}}};
+
+  AnthropicTool tool;
+  tool.set_name("simple_tool");
+  *tool.mutable_input_schema() = json_to_struct(tool_schema);
+
+  auto parsed_schema = struct_to_json(tool.input_schema());
+  EXPECT_TRUE(parsed_schema.contains("properties"));
+}
+
+// Test message with multiple content blocks
+TEST_F(AnthropicProtocolTest, AnthropicMessageMultipleContentBlocks) {
+  AnthropicMessage message;
+  message.set_role("assistant");
+
+  auto* content_blocks = message.mutable_content_blocks();
+
+  // Add text block
+  auto* text_block = content_blocks->add_blocks();
+  text_block->set_type("text");
+  text_block->set_text("I'll help you check the weather.");
+
+  // Add tool use block
+  auto* tool_block = content_blocks->add_blocks();
+  tool_block->set_type("tool_use");
+  tool_block->set_id("toolu_456");
+  tool_block->set_name("get_weather");
+  nlohmann::json input_json = {{"location", "Tokyo"}};
+  *tool_block->mutable_input() = json_to_struct(input_json);
+
+  EXPECT_EQ(message.content_blocks().blocks_size(), 2);
+  EXPECT_EQ(message.content_blocks().blocks(0).type(), "text");
+  EXPECT_EQ(message.content_blocks().blocks(1).type(), "tool_use");
+  EXPECT_EQ(message.content_blocks().blocks(1).id(), "toolu_456");
+}
+
+// Test AnthropicMessagesRequest with stop_sequences
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestStopSequences) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(100);
+
+  // Add stop sequences
+  request.add_stop_sequences("END");
+  request.add_stop_sequences("STOP");
+
+  // Add message
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("Generate text until END");
+
+  EXPECT_EQ(request.stop_sequences_size(), 2);
+  EXPECT_EQ(request.stop_sequences(0), "END");
+  EXPECT_EQ(request.stop_sequences(1), "STOP");
+}
+
+// Test JSON serialization round-trip
+TEST_F(AnthropicProtocolTest, JsonSerializationRoundTrip) {
+  AnthropicMessagesRequest original;
+  original.set_model("my_model");
+  original.set_max_tokens(100);
+  original.set_temperature(0.7f);
+
+  auto* message = original.add_messages();
+  message->set_role("user");
+  message->set_content_string("Hello!");
+
+  // Serialize to JSON
+  std::string json_str;
+  google::protobuf::util::JsonPrintOptions options;
+  options.preserve_proto_field_names = true;
+  auto status =
+      google::protobuf::util::MessageToJsonString(original, &json_str, options);
+  ASSERT_TRUE(status.ok());
+
+  // Deserialize back
+  AnthropicMessagesRequest parsed;
+  status = google::protobuf::util::JsonStringToMessage(json_str, &parsed);
+  ASSERT_TRUE(status.ok());
+
+  // Verify
+  EXPECT_EQ(parsed.model(), original.model());
+  EXPECT_EQ(parsed.max_tokens(), original.max_tokens());
+  EXPECT_FLOAT_EQ(parsed.temperature(), original.temperature());
+  EXPECT_EQ(parsed.messages_size(), 1);
+  EXPECT_EQ(parsed.messages(0).content_string(), "Hello!");
+}
+
+// Test AnthropicContentBlock with image source
+TEST_F(AnthropicProtocolTest, AnthropicContentBlockImage) {
+  AnthropicContentBlock content_block;
+  content_block.set_type("image");
+
+  nlohmann::json source_json = {
+      {"type", "base64"}, {"media_type", "image/png"}, {"data", "xxxxxx"}};
+  *content_block.mutable_source() = json_to_struct(source_json);
+
+  EXPECT_EQ(content_block.type(), "image");
+  EXPECT_TRUE(content_block.has_source());
+
+  auto parsed_source = struct_to_json(content_block.source());
+  EXPECT_EQ(parsed_source["type"], "base64");
+  EXPECT_EQ(parsed_source["media_type"], "image/png");
+}
+
+// Test tool result with list content
+TEST_F(AnthropicProtocolTest, AnthropicContentBlockToolResultList) {
+  AnthropicContentBlock content_block;
+  content_block.set_type("tool_result");
+  content_block.set_id("call_789");
+
+  // Use content_list for complex content
+  auto* content_list = content_block.mutable_content_list();
+  nlohmann::json item_json = {{"type", "text"}, {"text", "Result item 1"}};
+  *content_list->add_items() = json_to_struct(item_json);
+
+  EXPECT_EQ(content_block.type(), "tool_result");
+  EXPECT_TRUE(content_block.has_content_list());
+  EXPECT_EQ(content_block.content_list().items_size(), 1);
+}
+
+// Test response with tool_use stop_reason
+TEST_F(AnthropicProtocolTest, AnthropicMessagesResponseToolUse) {
+  AnthropicMessagesResponse response;
+  response.set_id("msg_tool_123");
+  response.set_type("message");
+  response.set_role("assistant");
+  response.set_model("my_model");
+  response.set_stop_reason("tool_use");
+
+  // Add text content
+  auto* text_block = response.add_content();
+  text_block->set_type("text");
+  text_block->set_text("Let me check the weather for you.");
+
+  // Add tool_use content
+  auto* tool_block = response.add_content();
+  tool_block->set_type("tool_use");
+  tool_block->set_id("toolu_weather_1");
+  tool_block->set_name("get_weather");
+  nlohmann::json input_json = {{"location", "San Francisco"},
+                               {"unit", "fahrenheit"}};
+  *tool_block->mutable_input() = json_to_struct(input_json);
+
+  // Set usage
+  auto* usage = response.mutable_usage();
+  usage->set_input_tokens(30);
+  usage->set_output_tokens(45);
+
+  EXPECT_EQ(response.stop_reason(), "tool_use");
+  EXPECT_EQ(response.content_size(), 2);
+  EXPECT_EQ(response.content(0).type(), "text");
+  EXPECT_EQ(response.content(1).type(), "tool_use");
+  EXPECT_EQ(response.content(1).name(), "get_weather");
+
+  auto parsed_input = struct_to_json(response.content(1).input());
+  EXPECT_EQ(parsed_input["location"], "San Francisco");
+}
+
+// Test request with system blocks
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestSystemBlocks) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(100);
+
+  // Use system_blocks instead of system_string
+  auto* system_blocks = request.mutable_system_blocks();
+  auto* block = system_blocks->add_blocks();
+  block->set_type("text");
+  block->set_text("You are a helpful assistant specialized in weather.");
+
+  auto* message = request.add_messages();
+  message->set_role("user");
+  message->set_content_string("What's the weather?");
+
+  EXPECT_TRUE(request.has_system_blocks());
+  EXPECT_FALSE(request.has_system_string());
+  EXPECT_EQ(request.system_blocks().blocks_size(), 1);
+  EXPECT_EQ(request.system_blocks().blocks(0).text(),
+            "You are a helpful assistant specialized in weather.");
+}
+
+// Test AnthropicToolChoice with "any" type
+TEST_F(AnthropicProtocolTest, AnthropicToolChoiceAny) {
+  AnthropicToolChoice tool_choice;
+  tool_choice.set_type("any");
+
+  EXPECT_EQ(tool_choice.type(), "any");
+  EXPECT_FALSE(tool_choice.has_name());
+}
+
+// Test complex conversation with multiple turns
+TEST_F(AnthropicProtocolTest, AnthropicMessagesRequestMultiTurn) {
+  AnthropicMessagesRequest request;
+  request.set_model("my_model");
+  request.set_max_tokens(200);
+
+  // User message 1
+  auto* msg1 = request.add_messages();
+  msg1->set_role("user");
+  msg1->set_content_string("What's the weather in Paris?");
+
+  // Assistant message with tool use
+  auto* msg2 = request.add_messages();
+  msg2->set_role("assistant");
+  auto* content_blocks2 = msg2->mutable_content_blocks();
+  auto* tool_use_block = content_blocks2->add_blocks();
+  tool_use_block->set_type("tool_use");
+  tool_use_block->set_id("toolu_123");
+  tool_use_block->set_name("get_weather");
+  nlohmann::json input_json = {{"location", "Paris"}};
+  *tool_use_block->mutable_input() = json_to_struct(input_json);
+
+  // User message with tool result
+  auto* msg3 = request.add_messages();
+  msg3->set_role("user");
+  auto* content_blocks3 = msg3->mutable_content_blocks();
+  auto* tool_result_block = content_blocks3->add_blocks();
+  tool_result_block->set_type("tool_result");
+  tool_result_block->set_id("toolu_123");
+  tool_result_block->set_content_string("The weather in Paris is 22°C, sunny.");
+
+  EXPECT_EQ(request.messages_size(), 3);
+  EXPECT_EQ(request.messages(0).role(), "user");
+  EXPECT_EQ(request.messages(1).role(), "assistant");
+  EXPECT_EQ(request.messages(2).role(), "user");
+  EXPECT_EQ(request.messages(1).content_blocks().blocks(0).type(), "tool_use");
+  EXPECT_EQ(request.messages(2).content_blocks().blocks(0).type(),
+            "tool_result");
+}
+
+}  // namespace
+}  // namespace xllm

upstream_ref/xllm/tests/api_service/chat_json_parser_test.cpp

@@ -0,0 +1,383 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "api_service/chat_json_parser.h"
+
+#include <gtest/gtest.h>
+
+#include <nlohmann/json.hpp>
+
+namespace xllm {
+
+class PreprocessChatJsonTest : public ::testing::Test {
+ protected:
+  void expect_success(const std::string& input,
+                      const ChatJsonParser& parser,
+                      const std::string& expected_output) {
+    auto [status, result] = parser.preprocess(input);
+    ASSERT_TRUE(status.ok()) << "Unexpected error: " << status.message();
+    auto result_json = nlohmann::json::parse(result);
+    auto expected_json = nlohmann::json::parse(expected_output);
+    EXPECT_EQ(result_json, expected_json);
+  }
+
+  void expect_error(const std::string& input,
+                    const ChatJsonParser& parser,
+                    const std::string& expected_error_substring) {
+    auto [status, result] = parser.preprocess(input);
+    ASSERT_FALSE(status.ok()) << "Expected error but got success";
+    EXPECT_NE(status.message().find(expected_error_substring),
+              std::string::npos)
+        << "Error message '" << status.message()
+        << "' does not contain expected substring '" << expected_error_substring
+        << "'";
+  }
+};
+
+// =============================================================================
+// Basic functionality tests
+// =============================================================================
+
+TEST_F(PreprocessChatJsonTest, PassThroughNonArrayContent) {
+  // String content should pass through unchanged
+  std::string input = R"({
+    "messages": [{"role": "user", "content": "Hello"}]
+  })";
+  LlmChatJsonParser llm_parser;
+  VlmChatJsonParser vlm_parser;
+  expect_success(input, llm_parser, input);
+  expect_success(input, vlm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, PassThroughNoMessages) {
+  // JSON without messages field should pass through
+  std::string input = R"({"model": "test"})";
+  LlmChatJsonParser llm_parser;
+  expect_success(input, llm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, CombineTextArrayIntoString) {
+  // Array of text items should be combined into single string for
+  // non-multimodal
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "text", "text": "Hello"},
+        {"type": "text", "text": "World"}
+      ]
+    }]
+  })";
+  std::string expected = R"({
+    "messages": [{"role": "user", "content": "Hello\nWorld"}]
+  })";
+  LlmChatJsonParser llm_parser;
+  VlmChatJsonParser vlm_parser;
+  expect_success(input, llm_parser, expected);
+  // For multimodal, array is preserved (not combined)
+  expect_success(input, vlm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, SingleTextItemCombined) {
+  // Single text item in array should be converted to string for non-multimodal
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [{"type": "text", "text": "Hello"}]
+    }]
+  })";
+  std::string expected = R"({
+    "messages": [{"role": "user", "content": "Hello"}]
+  })";
+  LlmChatJsonParser llm_parser;
+  VlmChatJsonParser vlm_parser;
+  expect_success(input, llm_parser, expected);
+  // For multimodal, array is preserved
+  expect_success(input, vlm_parser, input);
+}
+
+// =============================================================================
+// Multimodal content tests (Issue #801)
+// =============================================================================
+
+TEST_F(PreprocessChatJsonTest, ImageUrlPassesThroughOnMultimodal) {
+  // image_url content should pass through unchanged on multimodal endpoint
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,abc"}},
+        {"type": "text", "text": "What is this?"}
+      ]
+    }]
+  })";
+  VlmChatJsonParser vlm_parser;
+  // Should pass through unchanged for multimodal
+  expect_success(input, vlm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, ImageUrlErrorsOnTextOnly) {
+  // image_url content should error on text-only endpoint with helpful message
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,abc"}},
+        {"type": "text", "text": "What is this?"}
+      ]
+    }]
+  })";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "multimodal backend");
+  expect_error(input, llm_parser, "-backend vlm");
+}
+
+TEST_F(PreprocessChatJsonTest, MultipleMessagesWithMixedContent) {
+  // Multiple messages: some text-only, some with images
+  // On multimodal, all arrays are preserved (no combining)
+  std::string input = R"({
+    "messages": [
+      {
+        "role": "system",
+        "content": [{"type": "text", "text": "You are helpful."}]
+      },
+      {
+        "role": "user",
+        "content": [
+          {"type": "image_url", "image_url": {"url": "data:image/png;base64,xyz"}},
+          {"type": "text", "text": "Describe this image"}
+        ]
+      }
+    ]
+  })";
+  VlmChatJsonParser vlm_parser;
+  // On multimodal: all arrays preserved unchanged
+  expect_success(input, vlm_parser, input);
+}
+
+// =============================================================================
+// Error handling tests
+// =============================================================================
+
+TEST_F(PreprocessChatJsonTest, InvalidJsonReturnsError) {
+  std::string input = "not valid json";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "Invalid JSON");
+}
+
+TEST_F(PreprocessChatJsonTest, NonObjectMessageReturnsError) {
+  std::string input = R"({"messages": ["not an object"]})";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "must be an object");
+}
+
+TEST_F(PreprocessChatJsonTest, NonObjectContentItemReturnsError) {
+  std::string input = R"({
+    "messages": [{"role": "user", "content": ["not an object"]}]
+  })";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "must be an object");
+}
+
+TEST_F(PreprocessChatJsonTest, MissingTextFieldReturnsError) {
+  std::string input = R"({
+    "messages": [{"role": "user", "content": [{"type": "text"}]}]
+  })";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "Missing or invalid 'text' field");
+}
+
+TEST_F(PreprocessChatJsonTest, NonStringTextFieldReturnsError) {
+  std::string input = R"({
+    "messages": [{"role": "user", "content": [{"type": "text", "text": 123}]}]
+  })";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "Missing or invalid 'text' field");
+}
+
+TEST_F(PreprocessChatJsonTest, MalformedTextInMultimodalContent) {
+  // Multimodal mode skips parsing entirely - validation happens downstream
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "image_url", "image_url": {"url": "..."}},
+        {"type": "text"}
+      ]
+    }]
+  })";
+  VlmChatJsonParser vlm_parser;
+  // Should pass through unchanged without validation
+  expect_success(input, vlm_parser, input);
+}
+
+// =============================================================================
+// Edge cases
+// =============================================================================
+
+TEST_F(PreprocessChatJsonTest, EmptyContentArray) {
+  // Empty content array - should result in empty string for non-multimodal
+  std::string input = R"({
+    "messages": [{"role": "user", "content": []}]
+  })";
+  std::string expected = R"({
+    "messages": [{"role": "user", "content": ""}]
+  })";
+  LlmChatJsonParser llm_parser;
+  VlmChatJsonParser vlm_parser;
+  expect_success(input, llm_parser, expected);
+  // For multimodal, empty array is preserved
+  expect_success(input, vlm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, PreservesOtherFields) {
+  // Other fields in the request should be preserved
+  std::string input = R"({
+    "model": "test-model",
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}],
+    "temperature": 0.7,
+    "max_tokens": 100
+  })";
+  std::string expected = R"({
+    "model": "test-model",
+    "messages": [{"role": "user", "content": "Hi"}],
+    "temperature": 0.7,
+    "max_tokens": 100
+  })";
+  LlmChatJsonParser llm_parser;
+  VlmChatJsonParser vlm_parser;
+  expect_success(input, llm_parser, expected);
+  // For multimodal, array is preserved
+  expect_success(input, vlm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, UnknownContentTypeOnMultimodal) {
+  // Unknown content types should pass through on multimodal
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [{"type": "video", "video": {"url": "..."}}]
+    }]
+  })";
+  VlmChatJsonParser vlm_parser;
+  expect_success(input, vlm_parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, UnknownContentTypeErrorsOnTextOnly) {
+  // Unknown content types should error on text-only with helpful message
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [{"type": "video", "video": {"url": "..."}}]
+    }]
+  })";
+  LlmChatJsonParser llm_parser;
+  expect_error(input, llm_parser, "multimodal backend");
+}
+
+// =============================================================================
+// Anthropic parser tests
+// =============================================================================
+
+TEST_F(PreprocessChatJsonTest, AnthropicStringContentRemapped) {
+  std::string input = R"({
+    "messages": [{"role": "user", "content": "Hello"}]
+  })";
+  std::string expected = R"({
+    "messages": [{"role": "user", "content_string": "Hello"}]
+  })";
+  AnthropicChatJsonParser parser;
+  expect_success(input, parser, expected);
+}
+
+TEST_F(PreprocessChatJsonTest, AnthropicArrayContentRemapped) {
+  std::string input = R"({
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "text", "text": "Hello"},
+        {"type": "image", "source": {"data": "abc"}}
+      ]
+    }]
+  })";
+  std::string expected = R"({
+    "messages": [{
+      "role": "user",
+      "content_blocks": {
+        "blocks": [
+          {"type": "text", "text": "Hello"},
+          {"type": "image", "source": {"data": "abc"}}
+        ]
+      }
+    }]
+  })";
+  AnthropicChatJsonParser parser;
+  expect_success(input, parser, expected);
+}
+
+TEST_F(PreprocessChatJsonTest, AnthropicSystemStringRemapped) {
+  std::string input = R"({
+    "system": "You are helpful.",
+    "messages": [{"role": "user", "content": "Hi"}]
+  })";
+  std::string expected = R"({
+    "system_string": "You are helpful.",
+    "messages": [{"role": "user", "content_string": "Hi"}]
+  })";
+  AnthropicChatJsonParser parser;
+  expect_success(input, parser, expected);
+}
+
+TEST_F(PreprocessChatJsonTest, AnthropicSystemArrayRemapped) {
+  std::string input = R"({
+    "system": [{"type": "text", "text": "You are helpful."}],
+    "messages": [{"role": "user", "content": "Hi"}]
+  })";
+  std::string expected = R"({
+    "system_blocks": {"blocks": [{"type": "text", "text": "You are helpful."}]},
+    "messages": [{"role": "user", "content_string": "Hi"}]
+  })";
+  AnthropicChatJsonParser parser;
+  expect_success(input, parser, expected);
+}
+
+TEST_F(PreprocessChatJsonTest, AnthropicNoContentNoSystem) {
+  std::string input = R"({"model": "claude-3"})";
+  AnthropicChatJsonParser parser;
+  expect_success(input, parser, input);
+}
+
+TEST_F(PreprocessChatJsonTest, AnthropicInvalidJsonReturnsError) {
+  std::string input = "not valid json";
+  AnthropicChatJsonParser parser;
+  expect_error(input, parser, "Invalid JSON");
+}
+
+TEST_F(PreprocessChatJsonTest, AnthropicPreservesOtherFields) {
+  std::string input = R"({
+    "model": "claude-3",
+    "max_tokens": 1024,
+    "messages": [{"role": "user", "content": "Hello"}]
+  })";
+  std::string expected = R"({
+    "model": "claude-3",
+    "max_tokens": 1024,
+    "messages": [{"role": "user", "content_string": "Hello"}]
+  })";
+  AnthropicChatJsonParser parser;
+  expect_success(input, parser, expected);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/api_service/openai_service_test.cpp

@@ -0,0 +1,300 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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.
+==============================================================================*/
+
+// Testing steps:
+// 1. start the xllm server first
+// 2. run the test with disabled cases enabled
+// XLLM_TEST_BASE_URL=http://127.0.0.1:9977 XLLM_TEST_MODEL=Qwen3-8B
+//     ./build/lib.linux-aarch64-cpython-311/xllm/openai_service_test
+//     --gtest_also_run_disabled_tests
+
+#include <brpc/channel.h>
+#include <brpc/controller.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <cstdlib>
+#include <map>
+#include <nlohmann/json.hpp>
+#include <string>
+#include <utility>
+
+namespace xllm {
+namespace {
+
+struct TestConfig {
+  std::string base_url;
+  std::string model;
+  std::string api_key;
+
+  static TestConfig get() {
+    TestConfig config;
+    const char* url_env = std::getenv("XLLM_TEST_BASE_URL");
+    config.base_url = url_env ? url_env : "http://127.0.0.1:9977";
+
+    const char* model_env = std::getenv("XLLM_TEST_MODEL");
+    config.model = model_env ? model_env : "my_model";
+
+    const char* key_env = std::getenv("XLLM_TEST_API_KEY");
+    config.api_key = key_env ? key_env : "xllm-test-123456";
+    return config;
+  }
+};
+
+struct HttpResult {
+  bool controller_failed = false;
+  int status_code = 0;
+  std::string content_type;
+  std::string error_text;
+  std::string body;
+  nlohmann::json json = nullptr;
+};
+
+class HttpClient {
+ public:
+  explicit HttpClient(const std::string& base_url) {
+    brpc::ChannelOptions options;
+    options.protocol = brpc::PROTOCOL_HTTP;
+    options.connection_type = brpc::CONNECTION_TYPE_POOLED;
+    options.timeout_ms = 60000;
+    options.max_retry = 3;
+
+    if (channel_.Init(base_url.c_str(), &options) != 0) {
+      LOG(ERROR) << "Failed to init channel for " << base_url;
+    }
+  }
+
+  HttpResult post(const std::string& path,
+                  const nlohmann::json& body,
+                  const std::map<std::string, std::string>& headers) {
+    brpc::Controller cntl;
+    cntl.http_request().uri() = path;
+    cntl.http_request().set_method(brpc::HTTP_METHOD_POST);
+    cntl.http_request().set_content_type("application/json");
+    for (const auto& [key, value] : headers) {
+      cntl.http_request().SetHeader(key, value);
+    }
+    cntl.request_attachment().append(body.dump());
+
+    channel_.CallMethod(nullptr, &cntl, nullptr, nullptr, nullptr);
+
+    HttpResult result;
+    result.controller_failed = cntl.Failed();
+    result.status_code = cntl.http_response().status_code();
+    result.content_type = cntl.http_response().content_type();
+    result.error_text = cntl.ErrorText();
+    result.body = cntl.response_attachment().to_string();
+    if (!result.body.empty()) {
+      try {
+        result.json = nlohmann::json::parse(result.body);
+      } catch (const std::exception&) {
+        result.json = nullptr;
+      }
+    }
+    return result;
+  }
+
+ private:
+  brpc::Channel channel_;
+};
+
+std::string describe_result(const HttpResult& result) {
+  std::string description = "status=" + std::to_string(result.status_code);
+  if (!result.content_type.empty()) {
+    description += ", content_type=" + result.content_type;
+  }
+  if (!result.error_text.empty()) {
+    description += ", error=" + result.error_text;
+  }
+  if (!result.body.empty()) {
+    description += ", body=" + result.body;
+  }
+  return description;
+}
+
+void expect_error_contains(const HttpResult& result,
+                           const std::string& expected_fragment) {
+  const std::string haystack =
+      result.error_text + "\n" + result.body + "\n" + result.json.dump();
+  EXPECT_NE(haystack.find(expected_fragment), std::string::npos)
+      << "Expected fragment '" << expected_fragment
+      << "' in response, got: " << describe_result(result);
+}
+
+void expect_logprobs_shape(const nlohmann::json& choice) {
+  ASSERT_TRUE(choice.contains("logprobs"));
+  ASSERT_TRUE(choice["logprobs"].is_object());
+  ASSERT_TRUE(choice["logprobs"].contains("tokens"));
+  ASSERT_TRUE(choice["logprobs"].contains("token_ids"));
+  ASSERT_TRUE(choice["logprobs"].contains("token_logprobs"));
+  ASSERT_TRUE(choice["logprobs"]["tokens"].is_array());
+  ASSERT_TRUE(choice["logprobs"]["token_ids"].is_array());
+  ASSERT_TRUE(choice["logprobs"]["token_logprobs"].is_array());
+  EXPECT_EQ(choice["logprobs"]["tokens"].size(),
+            choice["logprobs"]["token_ids"].size());
+  EXPECT_EQ(choice["logprobs"]["tokens"].size(),
+            choice["logprobs"]["token_logprobs"].size());
+}
+
+class DISABLED_OpenAIServerFeaturesTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    config_ = TestConfig::get();
+    client_ = std::make_unique<HttpClient>(config_.base_url);
+  }
+
+  std::map<std::string, std::string> get_headers() const {
+    return {{"Authorization", "Bearer " + config_.api_key},
+            {"Content-Type", "application/json"}};
+  }
+
+  nlohmann::json make_sample_request(const std::string& prompt,
+                                     int logprobs = 3) const {
+    return {{"model", config_.model},
+            {"prompt", prompt},
+            {"selector", {{"type", "literal"}, {"value", "<emb_0>"}}},
+            {"logprobs", logprobs},
+            {"request_id", "sample-it"}};
+  }
+
+  void expect_sample_choice(const nlohmann::json& choice,
+                            std::size_t expected_index,
+                            std::size_t max_logprobs) const {
+    EXPECT_EQ(choice["index"], expected_index);
+    ASSERT_TRUE(choice.contains("finish_reason"));
+    ASSERT_TRUE(choice.contains("text"));
+    expect_logprobs_shape(choice);
+    EXPECT_LE(choice["logprobs"]["tokens"].size(), max_logprobs);
+    const std::string finish_reason = choice["finish_reason"];
+    EXPECT_TRUE(finish_reason == "selector_match" ||
+                finish_reason == "empty_logprobs")
+        << "Unexpected finish_reason: " << finish_reason;
+    if (!choice["logprobs"]["tokens"].empty()) {
+      EXPECT_EQ(choice["text"], choice["logprobs"]["tokens"][0]);
+    } else {
+      EXPECT_TRUE(choice["text"].empty());
+      EXPECT_EQ(finish_reason, "empty_logprobs");
+    }
+  }
+
+  TestConfig config_;
+  std::unique_ptr<HttpClient> client_;
+};
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest, SampleSingleMatch) {
+  const HttpResult result =
+      client_->post("/v1/sample",
+                    make_sample_request("Classify <emb_0> in one token."),
+                    get_headers());
+
+  ASSERT_FALSE(result.controller_failed) << describe_result(result);
+  ASSERT_EQ(result.status_code, 200) << describe_result(result);
+  EXPECT_EQ(result.content_type, "application/json") << describe_result(result);
+  ASSERT_TRUE(result.json.is_object()) << describe_result(result);
+  EXPECT_EQ(result.json["id"], "sample-it");
+  EXPECT_EQ(result.json["object"], "sample_completion");
+  EXPECT_EQ(result.json["model"], config_.model);
+  ASSERT_TRUE(result.json.contains("choices"));
+  ASSERT_EQ(result.json["choices"].size(), 1);
+  expect_sample_choice(result.json["choices"][0], 0, 3);
+}
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest, SampleMultipleMatchesStayOrdered) {
+  const HttpResult result = client_->post(
+      "/v1/sample", make_sample_request("A<emb_0>B<emb_0>C"), get_headers());
+
+  ASSERT_FALSE(result.controller_failed) << describe_result(result);
+  ASSERT_EQ(result.status_code, 200) << describe_result(result);
+  ASSERT_TRUE(result.json.is_object()) << describe_result(result);
+  ASSERT_TRUE(result.json.contains("choices"));
+  ASSERT_EQ(result.json["choices"].size(), 2);
+  expect_sample_choice(result.json["choices"][0], 0, 3);
+  expect_sample_choice(result.json["choices"][1], 1, 3);
+}
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest,
+       SampleSelectorMissReturnsEmptyChoices) {
+  const HttpResult result = client_->post(
+      "/v1/sample", make_sample_request("plain text"), get_headers());
+
+  ASSERT_FALSE(result.controller_failed) << describe_result(result);
+  ASSERT_EQ(result.status_code, 200) << describe_result(result);
+  ASSERT_TRUE(result.json.is_object()) << describe_result(result);
+  EXPECT_EQ(result.json["id"], "sample-it");
+  EXPECT_EQ(result.json["object"], "sample_completion");
+  ASSERT_TRUE(result.json.contains("choices"));
+  EXPECT_TRUE(result.json["choices"].empty());
+}
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest,
+       SampleRejectsUnsupportedSelectorType) {
+  nlohmann::json request = make_sample_request("A<emb_0>");
+  request["selector"]["type"] = "regex";
+  const HttpResult result = client_->post("/v1/sample", request, get_headers());
+
+  EXPECT_NE(result.status_code, 200) << describe_result(result);
+  expect_error_contains(result, "literal");
+}
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest, SampleRejectsOutOfRangeLogprobs) {
+  nlohmann::json request = make_sample_request("A<emb_0>");
+  request["logprobs"] = 0;
+  const HttpResult result = client_->post("/v1/sample", request, get_headers());
+
+  EXPECT_NE(result.status_code, 200) << describe_result(result);
+  expect_error_contains(result, "between 1 and 5");
+}
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest, CompletionsRegressionSmoke) {
+  nlohmann::json request = {{"model", config_.model},
+                            {"prompt", "Say hi."},
+                            {"max_tokens", 1},
+                            {"temperature", 0.0}};
+  const HttpResult result =
+      client_->post("/v1/completions", request, get_headers());
+
+  ASSERT_FALSE(result.controller_failed) << describe_result(result);
+  ASSERT_EQ(result.status_code, 200) << describe_result(result);
+  EXPECT_EQ(result.content_type, "application/json") << describe_result(result);
+  ASSERT_TRUE(result.json.is_object()) << describe_result(result);
+  ASSERT_TRUE(result.json.contains("choices"));
+  ASSERT_EQ(result.json["choices"].size(), 1);
+  EXPECT_EQ(result.json["choices"][0]["index"], 0);
+  ASSERT_TRUE(result.json.contains("usage"));
+}
+
+TEST_F(DISABLED_OpenAIServerFeaturesTest, ChatCompletionsRegressionSmoke) {
+  nlohmann::json request = {
+      {"model", config_.model},
+      {"messages", {{{"role", "user"}, {"content", "Say hi."}}}},
+      {"max_tokens", 1},
+      {"temperature", 0.0}};
+  const HttpResult result =
+      client_->post("/v1/chat/completions", request, get_headers());
+
+  ASSERT_FALSE(result.controller_failed) << describe_result(result);
+  ASSERT_EQ(result.status_code, 200) << describe_result(result);
+  EXPECT_EQ(result.content_type, "application/json") << describe_result(result);
+  ASSERT_TRUE(result.json.is_object()) << describe_result(result);
+  ASSERT_TRUE(result.json.contains("choices"));
+  ASSERT_EQ(result.json["choices"].size(), 1);
+  ASSERT_TRUE(result.json["choices"][0].contains("message"));
+  ASSERT_TRUE(result.json["choices"][0]["message"].contains("role"));
+  ASSERT_TRUE(result.json["choices"][0]["message"].contains("content"));
+  ASSERT_TRUE(result.json.contains("usage"));
+}
+
+}  // namespace
+}  // namespace xllm

upstream_ref/xllm/tests/api_service/sample_service_impl_test.cpp

@@ -0,0 +1,335 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "sample_service_impl.h"
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <cstring>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "core/framework/tokenizer/tokenizer.h"
+
+namespace xllm {
+namespace {
+
+class CharTokenizer final : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens = true) const override {
+    if (ids == nullptr) {
+      return false;
+    }
+    ids->clear();
+    if (add_special_tokens) {
+      ids->push_back(kBosTokenId);
+    }
+    size_t pos = 0;
+    while (pos < text.size()) {
+      if (text.size() - pos >= kEmbTokenLen &&
+          std::memcmp(text.data() + pos, kEmbToken, kEmbTokenLen) == 0) {
+        ids->push_back(kEmbTokenId);
+        pos += kEmbTokenLen;
+        continue;
+      }
+      ids->push_back(static_cast<unsigned char>(text[pos]));
+      ++pos;
+    }
+    return true;
+  }
+
+  std::optional<int32_t> token_to_id(
+      const std::string_view& token) const override {
+    if (token == std::string_view(kEmbToken, kEmbTokenLen)) {
+      return kEmbTokenId;
+    }
+    return std::nullopt;
+  }
+
+  std::string id_to_token(int32_t id) const override {
+    if (id == kBosTokenId) {
+      return "<bos>";
+    }
+    if (id == kEmbTokenId) {
+      return std::string(kEmbToken, kEmbTokenLen);
+    }
+    return std::string(1, static_cast<char>(id));
+  }
+
+ private:
+  static constexpr int32_t kBosTokenId = 1;
+  static constexpr int32_t kEmbTokenId = 100000;
+  static constexpr char kEmbToken[] = "<emb_0>";
+  static constexpr size_t kEmbTokenLen = sizeof(kEmbToken) - 1;
+};
+
+class UnstableLiteralTokenizer final : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens = true) const override {
+    if (ids == nullptr) {
+      return false;
+    }
+    ids->clear();
+    if (add_special_tokens) {
+      ids->push_back(1);
+    }
+    for (char ch : text) {
+      ids->push_back(static_cast<unsigned char>(ch));
+    }
+    return true;
+  }
+};
+
+proto::SampleRequest make_valid_request() {
+  proto::SampleRequest request;
+  request.set_model("mtp");
+  request.set_prompt("plain text");
+  auto* selector = request.mutable_selector();
+  selector->set_type("literal");
+  selector->set_value("<emb_0>");
+  return request;
+}
+
+TEST(SampleServiceImplTest, ValidateRequestRejectsUnsupportedSelectorType) {
+  auto request = make_valid_request();
+  request.mutable_selector()->set_type("regex");
+
+  const Status status = sample_service_internal::validate_request(request);
+
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_NE(status.message().find("literal"), std::string::npos);
+}
+
+TEST(SampleServiceImplTest, ValidateRequestRejectsOutOfRangeLogprobs) {
+  auto request = make_valid_request();
+  request.set_logprobs(0);
+
+  Status status = sample_service_internal::validate_request(request);
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_NE(status.message().find("between 1 and 5"), std::string::npos);
+
+  request.set_logprobs(6);
+  status = sample_service_internal::validate_request(request);
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_NE(status.message().find("between 1 and 5"), std::string::npos);
+}
+
+TEST(SampleServiceImplTest, ValidateRequestRejectsMissingRequiredFields) {
+  auto request = make_valid_request();
+  request.clear_model();
+
+  Status status = sample_service_internal::validate_request(request);
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_EQ(status.message(), "model is required");
+
+  request = make_valid_request();
+  request.clear_prompt();
+  status = sample_service_internal::validate_request(request);
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_EQ(status.message(), "prompt is required");
+
+  request = make_valid_request();
+  request.clear_selector();
+  status = sample_service_internal::validate_request(request);
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_EQ(status.message(), "selector is required");
+
+  request = make_valid_request();
+  request.mutable_selector()->clear_value();
+  status = sample_service_internal::validate_request(request);
+  EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
+  EXPECT_EQ(status.message(), "selector.value is required");
+}
+
+TEST(SampleServiceImplTest, ValidateRuntimeConfigRejectsScheduleOverlap) {
+  Status status = sample_service_internal::validate_runtime_config(true);
+  EXPECT_EQ(status.code(), StatusCode::UNAVAILABLE);
+  EXPECT_NE(status.message().find("does not support async scheduling"),
+            std::string::npos);
+
+  status = sample_service_internal::validate_runtime_config(false);
+  EXPECT_TRUE(status.ok());
+}
+
+TEST(SampleServiceImplTest, BuildRequestParamsAndEmptyResponseForSelectorMiss) {
+  CharTokenizer tokenizer;
+  auto request = make_valid_request();
+
+  RequestParams params;
+  ASSERT_TRUE(sample_service_internal::build_request_params(
+      request, tokenizer, &params));
+  EXPECT_TRUE(params.logprobs);
+  EXPECT_EQ(params.top_logprobs,
+            sample_service_internal::kDefaultSampleLogprobs);
+  EXPECT_TRUE(params.is_sample_request);
+  EXPECT_TRUE(params.sample_slots.empty());
+  EXPECT_FALSE(params.request_id.empty());
+
+  proto::SampleResponse response;
+  ASSERT_TRUE(sample_service_internal::build_empty_response(
+      request, tokenizer, params.request_id, &response));
+
+  EXPECT_EQ(response.id(), params.request_id);
+  EXPECT_EQ(response.object(), "sample_completion");
+  EXPECT_EQ(response.model(), request.model());
+  EXPECT_EQ(response.choices_size(), 0);
+  EXPECT_GT(response.created(), 0U);
+  ASSERT_TRUE(response.has_usage());
+  const int32_t expected_prompt_tokens =
+      static_cast<int32_t>(request.prompt().size() + 1);
+  EXPECT_EQ(response.usage().prompt_tokens(), expected_prompt_tokens);
+  EXPECT_EQ(response.usage().completion_tokens(), 0);
+  EXPECT_EQ(response.usage().total_tokens(), expected_prompt_tokens);
+}
+
+TEST(SampleServiceImplTest,
+     BuildRequestParamsKeepsExplicitRequestIdAndMatchedSampleSlots) {
+  CharTokenizer tokenizer;
+  auto request = make_valid_request();
+  request.set_prompt("A<emb_0>B<emb_0>C");
+  request.set_request_id("sample-explicit");
+  request.set_logprobs(4);
+
+  RequestParams params;
+  ASSERT_TRUE(sample_service_internal::build_request_params(
+      request, tokenizer, &params));
+
+  EXPECT_EQ(params.request_id, "sample-explicit");
+  EXPECT_TRUE(params.logprobs);
+  EXPECT_EQ(params.top_logprobs, 4);
+  EXPECT_TRUE(params.is_sample_request);
+  ASSERT_EQ(params.sample_slots.size(), 2);
+  EXPECT_EQ(params.sample_slots[0].request_id, "sample-explicit");
+  EXPECT_EQ(params.sample_slots[0].sample_id, 0);
+  EXPECT_EQ(params.sample_slots[0].token_position, 1);
+  EXPECT_EQ(params.sample_slots[1].request_id, "sample-explicit");
+  EXPECT_EQ(params.sample_slots[1].sample_id, 1);
+  EXPECT_EQ(params.sample_slots[1].token_position, 3);
+}
+
+TEST(SampleServiceImplTest, BuildRequestParamsRejectsUnstableLiteralToken) {
+  UnstableLiteralTokenizer tokenizer;
+  auto request = make_valid_request();
+  request.set_prompt("A<emb_0>B");
+  request.set_request_id("sample-explicit");
+
+  RequestParams params;
+  EXPECT_FALSE(sample_service_internal::build_request_params(
+      request, tokenizer, &params));
+}
+
+TEST(SampleServiceImplTest,
+     BuildResponseSortsBySampleIdAndSerializesTopLogprobs) {
+  RequestOutput req_output;
+  Usage usage;
+  usage.num_prompt_tokens = 8;
+  usage.num_generated_tokens = 2;
+  usage.num_total_tokens = 10;
+  req_output.usage = usage;
+
+  SequenceOutput missing_output;
+  missing_output.index = 1;
+  missing_output.finish_reason = "empty_logprobs";
+
+  SequenceOutput sampled_output;
+  sampled_output.index = 0;
+  sampled_output.text = "stale";
+  LogProb sampled_logprob;
+  sampled_logprob.token = "True";
+  sampled_logprob.token_id = 101;
+  sampled_logprob.logprob = -0.10f;
+  std::vector<LogProbData> top_logprobs;
+  LogProbData top1;
+  top1.token = "True";
+  top1.token_id = 101;
+  top1.logprob = -0.10f;
+  top_logprobs.push_back(top1);
+  LogProbData top2;
+  top2.token = "False";
+  top2.token_id = 102;
+  top2.logprob = -2.30f;
+  top_logprobs.push_back(top2);
+  sampled_logprob.top_logprobs = top_logprobs;
+  sampled_output.logprobs = std::vector<LogProb>{sampled_logprob};
+
+  req_output.outputs = {missing_output, sampled_output};
+
+  proto::SampleResponse response;
+  ASSERT_TRUE(sample_service_internal::build_response(
+      "sample-123", "mtp", 1773369600U, req_output, &response));
+
+  EXPECT_EQ(response.id(), "sample-123");
+  EXPECT_EQ(response.object(), "sample_completion");
+  EXPECT_EQ(response.created(), 1773369600U);
+  EXPECT_EQ(response.model(), "mtp");
+  ASSERT_EQ(response.choices_size(), 2);
+
+  EXPECT_EQ(response.choices(0).index(), 0);
+  EXPECT_EQ(response.choices(0).text(), "True");
+  EXPECT_EQ(response.choices(0).finish_reason(), "selector_match");
+  ASSERT_TRUE(response.choices(0).has_logprobs());
+  EXPECT_EQ(response.choices(0).logprobs().tokens_size(), 2);
+  EXPECT_EQ(response.choices(0).logprobs().tokens(0), "True");
+  EXPECT_EQ(response.choices(0).logprobs().tokens(1), "False");
+  EXPECT_EQ(response.choices(0).logprobs().token_ids(0), 101);
+  EXPECT_EQ(response.choices(0).logprobs().token_ids(1), 102);
+
+  EXPECT_EQ(response.choices(1).index(), 1);
+  EXPECT_TRUE(response.choices(1).text().empty());
+  EXPECT_EQ(response.choices(1).finish_reason(), "empty_logprobs");
+  ASSERT_TRUE(response.choices(1).has_logprobs());
+  EXPECT_EQ(response.choices(1).logprobs().tokens_size(), 0);
+  EXPECT_EQ(response.choices(1).logprobs().token_ids_size(), 0);
+  EXPECT_EQ(response.choices(1).logprobs().token_logprobs_size(), 0);
+
+  ASSERT_TRUE(response.has_usage());
+  EXPECT_EQ(response.usage().prompt_tokens(), 8);
+  EXPECT_EQ(response.usage().completion_tokens(), 2);
+  EXPECT_EQ(response.usage().total_tokens(), 10);
+}
+
+TEST(SampleServiceImplTest, BuildResponseFallsBackToSelectedTokenLogprob) {
+  RequestOutput req_output;
+
+  SequenceOutput sampled_output;
+  sampled_output.index = 0;
+  LogProb sampled_logprob;
+  sampled_logprob.token = "Maybe";
+  sampled_logprob.token_id = 201;
+  sampled_logprob.logprob = -0.25f;
+  sampled_output.logprobs = std::vector<LogProb>{sampled_logprob};
+  req_output.outputs = {sampled_output};
+
+  proto::SampleResponse response;
+  ASSERT_TRUE(sample_service_internal::build_response(
+      "sample-456", "mtp", 1773369601U, req_output, &response));
+
+  ASSERT_EQ(response.choices_size(), 1);
+  EXPECT_EQ(response.choices(0).text(), "Maybe");
+  EXPECT_EQ(response.choices(0).finish_reason(), "selector_match");
+  ASSERT_TRUE(response.choices(0).has_logprobs());
+  EXPECT_EQ(response.choices(0).logprobs().tokens_size(), 1);
+  EXPECT_EQ(response.choices(0).logprobs().tokens(0), "Maybe");
+  EXPECT_EQ(response.choices(0).logprobs().token_ids(0), 201);
+}
+
+}  // namespace
+}  // namespace xllm

upstream_ref/xllm/tests/core/CMakeLists.txt

@@ -0,0 +1,9 @@
+add_subdirectory(common)
+add_subdirectory(distributed_runtime)
+add_subdirectory(framework)
+add_subdirectory(kernels)
+add_subdirectory(layers)
+add_subdirectory(platform)
+add_subdirectory(runtime)
+add_subdirectory(scheduler)
+add_subdirectory(util)

upstream_ref/xllm/tests/core/common/CMakeLists.txt

@@ -0,0 +1,15 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    common_test
+  SRCS
+    rate_limiter_test.cpp
+  DEPS
+    common
+    absl::synchronization
+    absl::time
+    GTest::gtest_main
+    gflags::gflags
+    glog::glog
+)

upstream_ref/xllm/tests/core/common/rate_limiter_test.cpp

@@ -0,0 +1,26 @@
+#include "rate_limiter.h"
+
+#include <gtest/gtest.h>
+
+#include "global_flags.h"
+
+namespace xllm {
+
+TEST(RequestLimiterTest, Basic) {
+  // Set the maximum number of concurrent requests to 1.
+  FLAGS_max_concurrent_requests = 1;
+  RateLimiter rate_limiter;
+  // The current number of concurrent requests is 0, no rate limiting is
+  // applied.
+  EXPECT_EQ(rate_limiter.is_limited(), false);
+  // The current number of concurrent requests is 1, rate limiting is applied.
+  EXPECT_EQ(rate_limiter.is_limited(), true);
+  // Decrease the number of concurrent requests by one, changing the concurrency
+  // from 1 to 0.
+  rate_limiter.decrease_one_request();
+  // The current number of concurrent requests is 0, no rate limiting is
+  // applied.
+  EXPECT_EQ(rate_limiter.is_limited(), false);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/distributed_runtime/CMakeLists.txt

@@ -0,0 +1,2 @@
+# vlm_master_test is a standalone manual test binary and has no registered
+# CTest target yet.

upstream_ref/xllm/tests/core/distributed_runtime/vlm_master_test.cpp

@@ -0,0 +1,172 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 <absl/strings/str_split.h>
+#include <c10/core/Device.h>
+#include <c10/core/ScalarType.h>
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+// #include <pybind11/embed.h>
+
+#include <folly/init/Init.h>
+#include <torch/torch.h>
+
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <string>
+
+#include "distributed_runtime/vlm_master.h"
+#include "framework/request/sequence.h"
+#include "framework/sampling/sampling_params.h"
+#include "runtime/utils.h"
+
+std::vector<char> get_the_bytes(std::string filename) {
+  std::ifstream input(filename, std::ios::binary);
+  std::vector<char> bytes((std::istreambuf_iterator<char>(input)),
+                          (std::istreambuf_iterator<char>()));
+
+  input.close();
+  return bytes;
+}
+
+torch::Tensor load_tensor(std::string filename) {
+  std::vector<char> f = get_the_bytes(filename);
+  torch::IValue x = torch::pickle_load(f);
+  torch::Tensor my_tensor = x.toTensor();
+  return my_tensor;
+}
+
+bool run_inference(const std::string& input_embedding_path,
+                   xllm::VLMMaster& master,
+                   const xllm::RequestParams& sp) {
+  torch::Tensor input_embedding = xllm::load_tensor(input_embedding_path);
+  std::string prompt =
+      "<|im_start|>system\nYou are a helpful "
+      "assistant.<|im_end|>\n<|im_start|>user\n(<image>./</"
+      "image>)\nStructured output text information in the "
+      "graph?<|im_end|>\n<|im_"
+      "start|>assistant\n";
+  bool stream = false;
+  xllm::OutputCallback callback = [](xllm::RequestOutput output) -> bool {
+    if (output.finished) {
+      std::cout << "output.outputs.size(): " << output.outputs.size()
+                << std::endl;
+      if (output.outputs.size() > 0) {
+        const auto& text = output.outputs[0].text;
+        std::cout << "Callback called. text: " << text << std::endl;
+      } else {
+        std::cout << "Callback called. output.outputs.size() = 0" << std::endl;
+      }
+      return true;
+    } else {
+      std::cout << "not finished" << std::endl;
+      return false;
+    }
+  };
+  /*
+  std::future<bool> future = master.handle_request(
+      input_embedding, prompt, sp, callback);
+  bool result = future.get();
+  std::cout << "Final result: " << (result ? "Success" : "Failure")
+            << std::endl;*/
+  master.run_until_complete();
+  return true;
+}
+
+bool run_batch_inference(const std::string& input_embedding_path,
+                         xllm::VLMMaster& master,
+                         const xllm::RequestParams& sp) {
+  std::vector<torch::Tensor> input_embeddings;
+  torch::Tensor input_embedding = xllm::load_tensor(input_embedding_path);
+  xllm::print_tensor(input_embedding, "input_embedding", -1, false, false);
+  input_embeddings.emplace_back(input_embedding);
+  input_embeddings.emplace_back(input_embedding);
+  std::vector<std::string> prompts;
+  std::string prompt =
+      "<|im_start|>system\nYou are a helpful "
+      "assistant.<|im_end|>\n<|im_start|>user\n(<image>./</"
+      "image>)\nStructured output text information in the "
+      "graph?<|im_end|>\n<|im_"
+      "start|>assistant\n";
+  prompts.emplace_back(prompt);
+  prompts.emplace_back(prompt);
+  std::vector<xllm::RequestParams> sps;
+  sps.emplace_back(sp);
+  bool stream = false;
+  xllm::BatchOutputCallback batch_callback =
+      [](size_t index, xllm::RequestOutput output) -> bool {
+    if (output.finished) {
+      std::cout << "output.outputs.size(): " << output.outputs.size()
+                << std::endl;
+      if (output.outputs.size() > 0) {
+        const auto& text = output.outputs[0].text;
+        std::cout << "Callback called. "
+                  << "index: " << index << ", text: " << text << std::endl;
+      } else {
+        std::cout << "Callback called. output.outputs.size() = 0" << std::endl;
+      }
+      return true;
+    } else {
+      std::cout << "not finished" << std::endl;
+      return false;
+    }
+  };
+  master.handle_batch_request(input_embeddings, prompts, sps, batch_callback);
+  /*
+  std::vector<bool> results = futures.get();
+  for (const auto& result : results) {
+    std::cout << "Final result: " << (result ? "Success" : "Failure")
+              << std::endl;
+  }*/
+  master.run_until_complete();
+  return true;
+  // return results[0];
+}
+
+int main(int argc, char** argv) {
+  if (argc < 3) {
+    std::cerr << "Usage: " << argv[0] << " <model_path>" << std::endl;
+    return 1;
+  }
+  FLAGS_minloglevel = 0;
+  folly::Init init(&argc, &argv);
+  std::string model_path = argv[1];
+  std::string input_embedding_path = argv[2];
+  xllm::VLMMaster::Options option;
+  // option.model_path() = "/ktd/llava-1.5-7b-hf";
+  // option.model_path() = "/ktd/MiniCPM-V-2_6";
+  option.model_path() = model_path;
+  option.max_tokens_per_batch() = 2048;
+  // option.enable_prefix_cache() = false;
+
+  std::cout << "begin init xllm::VLMMaster==========" << std::endl;
+  xllm::VLMMaster master(option);
+
+  // torch::Tensor image_tensor;
+  // torch::Tensor image_tensor = torch::empty({});
+  // torch::load(image_tensor, "/ktd/xllm/image_tensor.pt");
+  // torch::Tensor image_tensor = load_tensor("/ktd/xllm/image_tensor.pt");
+  // std::string prompt = "USER: <image>\nWhat are these?\nASSISTANT:";
+
+  xllm::RequestParams sp;
+  sp.max_tokens = 1024;
+  sp.temperature = 0;
+  sp.stop_token_ids = {151645, 151643};
+  // run_inference(input_embedding_path, master);
+  run_batch_inference(input_embedding_path, master);
+
+  return 0;
+}

upstream_ref/xllm/tests/core/framework/CMakeLists.txt

@@ -0,0 +1,27 @@
+include(cc_test)
+
+add_subdirectory(batch)
+add_subdirectory(block)
+add_subdirectory(chat_template)
+add_subdirectory(eplb)
+add_subdirectory(kv_cache)
+add_subdirectory(kv_cache_transfer)
+add_subdirectory(parallel_state)
+add_subdirectory(prefix_cache)
+add_subdirectory(request)
+add_subdirectory(sampling)
+add_subdirectory(state_dict)
+add_subdirectory(tokenizer)
+
+cc_test(
+  NAME
+    hf_model_loader_test
+  SRCS
+    hf_model_loader_test.cpp
+  DEPS
+    :model_loader
+    :xllm_server
+    GTest::gtest_main
+)
+target_link_libraries(hf_model_loader_test PRIVATE
+  "$<LINK_GROUP:RESCAN,xtensor,xllm_server>")

upstream_ref/xllm/tests/core/framework/batch/CMakeLists.txt

@@ -0,0 +1,21 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    batch_test
+  SRCS
+    batch_test.cpp
+  DEPS
+    :batch
+    absl::time
+    GTest::gtest_main
+    $<$<BOOL:${USE_MLU}>:torch_mlu>
+    $<$<BOOL:${USE_NPU}>:torch_npu>
+)
+target_link_libraries(batch_test
+                      PUBLIC
+                      Python::Python
+                      $<$<BOOL:${USE_NPU}>:ascendcl>
+                      $<$<BOOL:${USE_NPU}>:hccl>
+                      $<$<BOOL:${USE_NPU}>:c_sec>
+                      $<$<BOOL:${USE_NPU}>:nnopbase>)

upstream_ref/xllm/tests/core/framework/block/CMakeLists.txt

@@ -0,0 +1,25 @@
+include(cc_test)
+
+if(USE_NPU)
+  cc_test(
+    NAME
+      block_test
+    SRCS
+      block_manager_test.cpp
+      single_block_manager_test.cpp
+    DEPS
+      :xllm_server
+      :block
+      :flags
+      :kv_cache
+      :prefix_cache
+      absl::random_random
+      Boost::serialization
+      GTest::gtest_main
+  )
+
+  target_link_libraries(block_test PRIVATE Folly::folly OpenSSL::SSL OpenSSL::Crypto protobuf::libprotobuf c_sec)
+  target_link_libraries(block_test PRIVATE
+    "$<LINK_GROUP:RESCAN,xtensor,xllm_server>")
+  add_dependencies(block_test brpc-static)
+endif()

upstream_ref/xllm/tests/core/framework/block/block_manager_test.cpp

@@ -0,0 +1,328 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+Copyright 2024 The ScaleLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+
+#include <type_traits>
+#include <utility>
+
+#include "block_manager_impl.h"
+#include "block_manager_pool.h"
+#include "common/global_flags.h"
+#include "framework/request/incremental_decoder.h"
+
+namespace xllm {
+
+namespace {
+
+template <typename T>
+class ScopedValue final {
+ public:
+  ScopedValue(T* target, T value) : target_(target), old_(*target) {
+    *target_ = value;
+  }
+  ~ScopedValue() { *target_ = old_; }
+
+  ScopedValue(const ScopedValue&) = delete;
+  ScopedValue& operator=(const ScopedValue&) = delete;
+
+ private:
+  T* target_;
+  T old_;
+};
+
+template <typename T, typename = void>
+struct HasEnableLinearStateOption : std::false_type {};
+
+template <typename T>
+struct HasEnableLinearStateOption<
+    T,
+    std::void_t<decltype(std::declval<T&>().enable_linear_state(true)),
+                decltype(std::declval<const T&>().enable_linear_state())>>
+    : std::true_type {};
+
+template <typename T, typename = void>
+struct HasSequenceSingleBlockApi : std::false_type {};
+
+template <typename T>
+struct HasSequenceSingleBlockApi<
+    T,
+    std::void_t<decltype(std::declval<const T&>().has_single_block_id()),
+                decltype(std::declval<const T&>().get_single_block_id()),
+                decltype(std::declval<T&>().reset_single_block())>>
+    : std::true_type {};
+
+template <typename OptionsT>
+bool EnableLinearStateOrFail(OptionsT& options) {
+  if constexpr (HasEnableLinearStateOption<OptionsT>::value) {
+    options.enable_linear_state(true);
+    return true;
+  }
+  ADD_FAILURE() << "Task 2 missing APIs: BlockManagerPool::Options "
+                   "enable_linear_state";
+  return false;
+}
+
+template <typename SeqT>
+bool HasSingleBlockIdOrFail(const SeqT& seq) {
+  if constexpr (HasSequenceSingleBlockApi<SeqT>::value) {
+    return seq.has_single_block_id();
+  }
+  ADD_FAILURE() << "Missing APIs: Sequence single-block handle";
+  return false;
+}
+
+template <typename SeqT>
+int32_t GetSingleBlockIdOrFail(const SeqT& seq) {
+  if constexpr (HasSequenceSingleBlockApi<SeqT>::value) {
+    return seq.get_single_block_id();
+  }
+  ADD_FAILURE() << "Missing APIs: Sequence single-block handle";
+  return -1;
+}
+
+Sequence MakeSequence(size_t index, const std::vector<int32_t>& prompt_tokens) {
+  RequestSamplingParam sampling_param;
+  sampling_param.beam_width = 0;
+  sampling_param.is_embeddings = false;
+
+  StoppingChecker stopping_checker;
+
+  SequenceParams params;
+  params.seq_capacity = prompt_tokens.size() + 8;
+  params.echo = false;
+  params.skip_special_tokens = true;
+  params.streaming = false;
+  params.enable_schedule_overlap = false;
+  params.rec_type = RecType::kNone;
+  params.bos_token_id = 0;
+  params.request_id = "block_manager_pool_test";
+  params.sampling_param = &sampling_param;
+  params.stopping_checker = &stopping_checker;
+
+  IncrementalDecoder decoder(
+      /*prompt=*/"prompt",
+      /*num_prompt_tokens=*/prompt_tokens.size(),
+      /*echo=*/params.echo,
+      /*skip_special_tokens=*/params.skip_special_tokens);
+
+  return Sequence(index,
+                  prompt_tokens,
+                  /*input_embedding=*/torch::Tensor(),
+                  /*mm_data=*/MMData(),
+                  decoder,
+                  params);
+}
+
+}  // namespace
+
+TEST(BlockManagerTest, Basic) {
+  const uint32_t n_blocks = 10;
+  const uint32_t block_size = 2;
+  BlockManager::Options options;
+  options.num_blocks(n_blocks).block_size(block_size);
+  BlockManagerImpl manager(options);
+
+  EXPECT_EQ(manager.num_free_blocks(), n_blocks - 1);
+  EXPECT_EQ(manager.block_size(), block_size);
+
+  // Allocate a block
+  {
+    Block block = manager.allocate();
+    EXPECT_EQ(block.id(), 1);
+    EXPECT_EQ(block.size(), block_size);
+    EXPECT_EQ(block.is_shared(), false);
+    EXPECT_EQ(block.ref_count(), 1);
+
+    EXPECT_EQ(manager.num_free_blocks(), n_blocks - 2);
+  }
+  // the block should be freed after the scope
+  EXPECT_EQ(manager.num_free_blocks(), n_blocks - 1);
+
+  // Allocate a list of blocks
+  {
+    std::vector<Block> blocks;
+    for (uint32_t i = 1; i < n_blocks; ++i) {
+      auto block = manager.allocate();
+      EXPECT_EQ(block.id(), i);
+      EXPECT_EQ(block.size(), block_size);
+      EXPECT_EQ(block.is_shared(), false);
+      EXPECT_EQ(block.ref_count(), 1);
+      blocks.push_back(std::move(block));
+    }
+    EXPECT_EQ(manager.num_free_blocks(), 0);
+    for (const auto& block : blocks) {
+      EXPECT_EQ(block.ref_count(), 1);
+      EXPECT_EQ(block.is_shared(), false);
+    }
+
+    // Test CHECK failure
+    EXPECT_DEATH(manager.allocate(), "No more blocks available");
+  }
+
+  // all blocks should be freed after the scope
+  EXPECT_EQ(manager.num_free_blocks(), n_blocks - 1);
+
+  // Test shared blocks
+  {
+    Block block = manager.allocate();
+    EXPECT_EQ(block.ref_count(), 1);
+    EXPECT_EQ(block.is_shared(), false);
+    // test copy constructor
+    {
+      // NOLINTNEXTLINE
+      const Block block2 = block;
+      EXPECT_EQ(block.ref_count(), 2);
+      EXPECT_EQ(block.is_shared(), true);
+      EXPECT_EQ(block2.ref_count(), 2);
+      EXPECT_EQ(block2.is_shared(), true);
+      EXPECT_EQ(block2, block);
+    }
+    EXPECT_EQ(block.ref_count(), 1);
+    EXPECT_EQ(block.is_shared(), false);
+
+    // test assignment operator
+    {
+      Block block4 = manager.allocate();
+      block4 = block;
+      EXPECT_EQ(block.ref_count(), 2);
+      EXPECT_EQ(block.is_shared(), true);
+      EXPECT_EQ(block4.ref_count(), 2);
+      EXPECT_EQ(block4.is_shared(), true);
+      EXPECT_EQ(block4, block);
+
+      Block invalid_block;
+      invalid_block = block;
+      EXPECT_EQ(block.ref_count(), 3);
+      EXPECT_EQ(block.is_shared(), true);
+      EXPECT_EQ(invalid_block.ref_count(), 3);
+      EXPECT_EQ(invalid_block.is_shared(), true);
+      EXPECT_EQ(invalid_block, block);
+    }
+    EXPECT_EQ(block.ref_count(), 1);
+    EXPECT_EQ(block.is_shared(), false);
+
+    // test move constructor
+    {
+      Block block3 = std::move(block);
+      EXPECT_FALSE(block.is_valid());
+
+      EXPECT_EQ(block3.ref_count(), 1);
+      EXPECT_EQ(block3.is_shared(), false);
+      EXPECT_FALSE(block3 == block);
+    }
+    EXPECT_FALSE(block.is_valid());
+  }
+}
+
+TEST(BlockManagerPoolTest, AllocateAssignsSingleBlockWhenEnabled) {
+  ScopedValue<int32_t> max_seqs_guard(&FLAGS_max_seqs_per_batch, 0);
+
+  BlockManagerPool::Options options;
+  options.num_blocks(8).host_num_blocks(0).block_size(1).enable_prefix_cache(
+      false);
+  ASSERT_TRUE(EnableLinearStateOrFail(options));
+  BlockManagerPool pool(options, /*dp_size=*/1);
+
+  Sequence seq = MakeSequence(0, /*prompt_tokens=*/{1, 2, 3});
+  EXPECT_TRUE(pool.allocate(&seq));
+  EXPECT_TRUE(HasSingleBlockIdOrFail(seq));
+  EXPECT_GE(GetSingleBlockIdOrFail(seq), 0);
+}
+
+TEST(BlockManagerPoolTest, DeallocateReleasesSingleBlockId) {
+  ScopedValue<int32_t> max_seqs_guard(&FLAGS_max_seqs_per_batch, 0);
+
+  BlockManagerPool::Options options;
+  options.num_blocks(8).host_num_blocks(0).block_size(1).enable_prefix_cache(
+      false);
+  ASSERT_TRUE(EnableLinearStateOrFail(options));
+  BlockManagerPool pool(options, /*dp_size=*/1);
+
+  Sequence seq1 = MakeSequence(0, /*prompt_tokens=*/{1, 2, 3});
+  ASSERT_TRUE(pool.allocate(&seq1));
+  const int32_t id1 = GetSingleBlockIdOrFail(seq1);
+  pool.deallocate(&seq1);
+  EXPECT_FALSE(HasSingleBlockIdOrFail(seq1));
+
+  Sequence seq2 = MakeSequence(1, /*prompt_tokens=*/{4, 5, 6});
+  ASSERT_TRUE(pool.allocate(&seq2));
+  EXPECT_EQ(GetSingleBlockIdOrFail(seq2), id1);
+}
+
+TEST(BlockManagerPoolTest, TryAllocateKvFailureRollsBackSingleBlock) {
+  // unified scheduler-side single-block pool has 2 ids.
+  ScopedValue<int32_t> max_seqs_guard(&FLAGS_max_seqs_per_batch, 0);
+
+  BlockManagerPool::Options options;
+  options.num_blocks(3).host_num_blocks(0).block_size(1).enable_prefix_cache(
+      false);
+  ASSERT_TRUE(EnableLinearStateOrFail(options));
+  BlockManagerPool pool(options, /*dp_size=*/1);
+
+  // This sequence needs far more KV blocks than available, forcing KV failure
+  // after embedding and linear ids are allocated.
+  std::vector<int32_t> huge_prompt(100, 1);
+  Sequence fail_seq = MakeSequence(0, huge_prompt);
+  EXPECT_FALSE(pool.try_allocate(&fail_seq));
+  EXPECT_FALSE(HasSingleBlockIdOrFail(fail_seq));
+
+  // The unified slot must have been rolled back, leaving enough capacity for
+  // two new sequences to allocate.
+  Sequence seq1 = MakeSequence(1, /*prompt_tokens=*/{1});
+  Sequence seq2 = MakeSequence(2, /*prompt_tokens=*/{2});
+  EXPECT_TRUE(pool.try_allocate(&seq1));
+  EXPECT_TRUE(pool.try_allocate(&seq2));
+  EXPECT_TRUE(HasSingleBlockIdOrFail(seq1));
+  EXPECT_TRUE(HasSingleBlockIdOrFail(seq2));
+}
+
+TEST(BlockManagerPoolTest, AllocateAssignsSingleBlockWhenLinearStateDisabled) {
+  ScopedValue<int32_t> max_seqs_guard(&FLAGS_max_seqs_per_batch, 2);
+
+  BlockManagerPool::Options options;
+  options.num_blocks(8).host_num_blocks(0).block_size(1).enable_prefix_cache(
+      false);
+  BlockManagerPool pool(options, /*dp_size=*/1);
+
+  Sequence seq = MakeSequence(0, /*prompt_tokens=*/{1, 2});
+  EXPECT_TRUE(pool.allocate(&seq));
+  EXPECT_TRUE(HasSingleBlockIdOrFail(seq));
+}
+
+TEST(BlockManagerPoolTest, SequenceCopyDoesNotReuseSingleBlockSlot) {
+  ScopedValue<int32_t> max_seqs_guard(&FLAGS_max_seqs_per_batch, 2);
+
+  BlockManagerPool::Options options;
+  options.num_blocks(8).host_num_blocks(0).block_size(1).enable_prefix_cache(
+      false);
+  ASSERT_TRUE(EnableLinearStateOrFail(options));
+  BlockManagerPool pool(options, /*dp_size=*/1);
+
+  Sequence src = MakeSequence(0, /*prompt_tokens=*/{1, 2, 3});
+  ASSERT_TRUE(pool.allocate(&src));
+  ASSERT_TRUE(HasSingleBlockIdOrFail(src));
+
+  Sequence clone(src);
+  EXPECT_FALSE(HasSingleBlockIdOrFail(clone));
+  EXPECT_EQ(clone.get_single_block_id(), -1);
+
+  ASSERT_TRUE(pool.allocate(&clone));
+  EXPECT_TRUE(HasSingleBlockIdOrFail(clone));
+  EXPECT_NE(GetSingleBlockIdOrFail(clone), GetSingleBlockIdOrFail(src));
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/block/single_block_manager_test.cpp

@@ -0,0 +1,124 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "single_block_manager.h"
+
+#include <gtest/gtest.h>
+
+namespace xllm {
+
+TEST(SingleBlockManagerTest, AllocateAndFreeRoundTrip) {
+  SingleBlockManager manager(3, "single");
+
+  EXPECT_EQ(manager.num_total_blocks(), 3);
+  EXPECT_EQ(manager.num_blocks_in_prefix_cache(), 0);
+  EXPECT_EQ(manager.num_free_blocks(), 3);
+  EXPECT_EQ(manager.num_used_blocks(), 0);
+  EXPECT_DOUBLE_EQ(manager.kv_cache_utilization(), 0.0);
+
+  auto blocks = manager.allocate(2);
+  ASSERT_EQ(blocks.size(), 2);
+  EXPECT_EQ(manager.num_free_blocks(), 1);
+  EXPECT_EQ(manager.num_used_blocks(), 2);
+  EXPECT_DOUBLE_EQ(manager.kv_cache_utilization(), 2.0 / 3.0);
+
+  manager.deallocate(Slice<Block>(blocks));
+  EXPECT_EQ(manager.num_free_blocks(), 1);
+  EXPECT_EQ(manager.num_used_blocks(), 0);
+
+  blocks.clear();
+  EXPECT_EQ(manager.num_free_blocks(), 3);
+  EXPECT_EQ(manager.num_used_blocks(), 0);
+  EXPECT_DOUBLE_EQ(manager.kv_cache_utilization(), 0.0);
+}
+
+TEST(SingleBlockManagerTest, AllocateReturnsEmptyWhenExhausted) {
+  SingleBlockManager manager(2, "single");
+  auto blocks = manager.allocate(2);
+  ASSERT_EQ(blocks.size(), 2);
+  EXPECT_TRUE(manager.allocate(1).empty());
+}
+
+TEST(SingleBlockManagerTest, AllocateSingleDiesWhenExhausted) {
+  SingleBlockManager manager(1, "single");
+  Block block = manager.allocate();
+  EXPECT_EQ(block.id(), 0);
+  EXPECT_DEATH(manager.allocate(), "No more single blocks available");
+}
+
+TEST(SingleBlockManagerTest, PrefixCacheStyleApisAreSafeNoopsWhenDisabled) {
+  SingleBlockManager manager(2, "single");
+
+  const int32_t token_ids_arr[] = {1, 2, 3};
+  const Slice<int32_t> token_ids(token_ids_arr, 3);
+
+  // allocate_shared should be safely substitutable with BlockManager behavior
+  // when prefix cache is disabled.
+  const auto shared = manager.allocate_shared(token_ids);
+  EXPECT_TRUE(shared.empty());
+
+  std::vector<Block> blocks = manager.allocate(1);
+  ASSERT_EQ(blocks.size(), 1u);
+
+  // cache() overloads should be safe no-ops.
+  EXPECT_NO_FATAL_FAILURE(manager.cache(token_ids, blocks));
+  EXPECT_NO_FATAL_FAILURE(manager.cache(blocks));
+
+  // get_merged_kvcache_event should be a safe no-op.
+  KvCacheEvent event;
+  // Avoid inserting a default-constructed XXH3Key, whose bytes would be
+  // uninitialized and could trigger undefined behavior in hashing.
+  uint8_t key_bytes[XXH3_128BITS_HASH_VALUE_LEN] = {};
+  event.stored_cache.emplace(key_bytes);
+  EXPECT_NO_FATAL_FAILURE(manager.get_merged_kvcache_event(&event));
+  EXPECT_EQ(event.stored_cache.size(), 1u);
+}
+
+TEST(SingleBlockManagerTest, UsedBlocksAccountingDoesNotLeakWithAliases) {
+  SingleBlockManager manager(1, "single");
+  EXPECT_EQ(manager.num_used_blocks(), 0u);
+  EXPECT_EQ(manager.num_free_blocks(), 1u);
+
+  {
+    Block block = manager.allocate();
+    EXPECT_EQ(manager.num_used_blocks(), 1u);
+    EXPECT_EQ(manager.num_free_blocks(), 0u);
+
+    // Create an alias to simulate an external holder (e.g. prefix-cache style
+    // reference) that outlives the "sequence-owned" reference.
+    Block alias = block;
+    EXPECT_EQ(block.ref_count(), 2u);
+
+    // Deallocate the sequence-owned reference while an alias still exists.
+    manager.deallocate(Slice<Block>(&block, 1));
+    EXPECT_EQ(manager.num_used_blocks(), 1u);
+
+    // Drop the sequence-owned reference without calling deallocate again.
+    block = Block();
+    EXPECT_EQ(alias.ref_count(), 1u);
+    EXPECT_EQ(manager.num_used_blocks(), 1u);
+    EXPECT_EQ(manager.num_free_blocks(), 0u);
+  }
+
+  // When the last alias is released, used block accounting should converge.
+  EXPECT_EQ(manager.num_used_blocks(), 0u);
+  EXPECT_EQ(manager.num_free_blocks(), 1u);
+}
+
+TEST(SingleBlockManagerTest, ConstructorRejectsZeroBlocks) {
+  EXPECT_DEATH({ SingleBlockManager manager(0, "single"); }, "No blocks");
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/chat_template/CMakeLists.txt

@@ -0,0 +1,34 @@
+include(cc_test)
+
+cc_test (
+  NAME
+    chat_template_test
+  SRCS
+    chat_template_test.cpp
+  DEPS
+    :chat_template
+    :jinja_chat_template
+    :deepseek_v32_cpp_template
+    GTest::gtest_main
+    gflags
+)
+
+cc_test(
+  NAME
+    jinja_chat_template_test
+  SRCS
+    jinja_chat_template_test.cpp
+  DEPS
+    :chat_template
+    GTest::gtest_main
+)
+
+cc_test (
+  NAME
+    deepseek_v32_cpp_template_test
+  SRCS
+    deepseek_v32_cpp_template_test.cpp
+  DEPS
+    :deepseek_v32_cpp_template
+    GTest::gtest_main
+)

upstream_ref/xllm/tests/core/framework/chat_template/chat_template_test.cpp

@@ -0,0 +1,67 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/chat_template/chat_template.h"
+
+#include <gflags/gflags.h>
+#include <gtest/gtest.h>
+
+#include "framework/chat_template/deepseek_v32_cpp_template.h"
+#include "framework/chat_template/jinja_chat_template.h"
+
+DECLARE_bool(use_cpp_chat_template);
+
+namespace xllm {
+namespace {
+
+class ScopedUseCppChatTemplate final {
+ public:
+  explicit ScopedUseCppChatTemplate(bool enabled)
+      : old_value_(FLAGS_use_cpp_chat_template) {
+    FLAGS_use_cpp_chat_template = enabled;
+  }
+
+  ~ScopedUseCppChatTemplate() { FLAGS_use_cpp_chat_template = old_value_; }
+
+ private:
+  bool old_value_;
+};
+
+TEST(ChatTemplateFactory, DeepseekV32FallsBackToJinjaWhenFlagDisabled) {
+  ScopedUseCppChatTemplate scoped_flag(/*enabled=*/false);
+  TokenizerArgs args;
+
+  std::unique_ptr<ChatTemplate> impl =
+      ChatTemplate::create(args, /*model_type=*/"deepseek_v32");
+
+  ASSERT_TRUE(impl != nullptr);
+  EXPECT_NE(dynamic_cast<JinjaChatTemplate*>(impl.get()), nullptr);
+  EXPECT_EQ(dynamic_cast<DeepseekV32CppTemplate*>(impl.get()), nullptr);
+}
+
+TEST(ChatTemplateFactory, NonDeepseekModelUsesJinjaWhenFlagEnabled) {
+  ScopedUseCppChatTemplate scoped_flag(/*enabled=*/true);
+  TokenizerArgs args;
+
+  std::unique_ptr<ChatTemplate> impl =
+      ChatTemplate::create(args, /*model_type=*/"qwen3");
+
+  ASSERT_TRUE(impl != nullptr);
+  EXPECT_NE(dynamic_cast<JinjaChatTemplate*>(impl.get()), nullptr);
+  EXPECT_EQ(dynamic_cast<DeepseekV32CppTemplate*>(impl.get()), nullptr);
+}
+
+}  // namespace
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/chat_template/deepseek_v32_cpp_template_test.cpp

@@ -0,0 +1,164 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/chat_template/deepseek_v32_cpp_template.h"
+
+#include <gtest/gtest.h>
+
+namespace xllm {
+
+TEST(DeepseekV32CppTemplate, BasicUserMessage) {
+  TokenizerArgs args;
+  args.bos_token("<｜begin▁of▁sentence｜>");
+  DeepseekV32CppTemplate encoder(args);
+
+  ChatMessages messages;
+  messages.emplace_back("user", "hello");
+
+  nlohmann::ordered_json kwargs = nlohmann::json::object();
+  auto prompt = encoder.apply(messages, /*json_tools=*/{}, kwargs);
+  ASSERT_TRUE(prompt.has_value());
+
+  EXPECT_NE(prompt->find("<｜begin▁of▁sentence｜>"), std::string::npos);
+  EXPECT_NE(prompt->find("<｜User｜>hello<｜Assistant｜>"), std::string::npos);
+}
+
+TEST(DeepseekV32CppTemplate, DefaultThinkingModeIsChat) {
+  TokenizerArgs args;
+  args.bos_token("<｜begin▁of▁sentence｜>");
+  DeepseekV32CppTemplate encoder(args);
+
+  ChatMessages messages;
+  messages.emplace_back("user", "hello");
+
+  nlohmann::ordered_json kwargs = nlohmann::json::object();
+  auto prompt = encoder.apply(messages, /*json_tools=*/{}, kwargs);
+  ASSERT_TRUE(prompt.has_value());
+
+  EXPECT_NE(prompt->find("</think>"), std::string::npos);
+  EXPECT_EQ(prompt->find("<think>"), std::string::npos);
+}
+
+TEST(DeepseekV32CppTemplate, ThinkingModeEnabledByKwargs) {
+  TokenizerArgs args;
+  args.bos_token("<｜begin▁of▁sentence｜>");
+  DeepseekV32CppTemplate encoder(args);
+
+  ChatMessages messages;
+  messages.emplace_back("user", "hello");
+
+  nlohmann::ordered_json kwargs = nlohmann::json::object();
+  kwargs["thinking"] = true;
+
+  auto prompt = encoder.apply(messages, /*json_tools=*/{}, kwargs);
+  ASSERT_TRUE(prompt.has_value());
+
+  EXPECT_NE(prompt->find("<think>"), std::string::npos);
+}
+
+TEST(DeepseekV32CppTemplate, ToolsInjectionFormat) {
+  TokenizerArgs args;
+  args.bos_token("<｜begin▁of▁sentence｜>");
+  DeepseekV32CppTemplate encoder(args);
+
+  ChatMessages messages;
+  messages.emplace_back("user", "weather in beijing");
+
+  std::vector<JsonTool> tools;
+  JsonTool tool;
+  tool.type = "function";
+  tool.function.name = "get_weather";
+  tool.function.description = "query weather";
+  tool.function.parameters = nlohmann::json{
+      {"type", "object"}, {"properties", {{"city", {{"type", "string"}}}}}};
+  tools.push_back(tool);
+
+  nlohmann::ordered_json kwargs = nlohmann::json::object();
+  auto prompt = encoder.apply(messages, tools, kwargs);
+  ASSERT_TRUE(prompt.has_value());
+
+  // vLLM DSML format
+  EXPECT_NE(prompt->find("## Tools"), std::string::npos);
+  EXPECT_NE(prompt->find("get_weather"), std::string::npos);
+  EXPECT_NE(prompt->find("<functions>"), std::string::npos);
+  EXPECT_NE(prompt->find("</functions>"), std::string::npos);
+  // User message after tools
+  EXPECT_NE(prompt->find("<｜User｜>weather in beijing"
+                         "<｜Assistant｜>"),
+            std::string::npos);
+}
+
+TEST(DeepseekV32CppTemplate, ToolsInjectedAsNewSystemMessage) {
+  TokenizerArgs args;
+  args.bos_token("<｜begin▁of▁sentence｜>");
+  DeepseekV32CppTemplate encoder(args);
+
+  ChatMessages messages;
+  messages.emplace_back("system", "You are helpful.");
+  messages.emplace_back("user", "hi");
+
+  std::vector<JsonTool> tools;
+  JsonTool tool;
+  tool.type = "function";
+  tool.function.name = "search";
+  tool.function.description = "search the web";
+  tool.function.parameters = nlohmann::json{{"type", "object"}};
+  tools.push_back(tool);
+
+  nlohmann::ordered_json kwargs = nlohmann::json::object();
+  auto prompt = encoder.apply(messages, tools, kwargs);
+  ASSERT_TRUE(prompt.has_value());
+
+  // Tools section before system content
+  size_t tools_pos = prompt->find("## Tools");
+  size_t content_pos = prompt->find("You are helpful.");
+  ASSERT_NE(tools_pos, std::string::npos);
+  ASSERT_NE(content_pos, std::string::npos);
+  EXPECT_LT(tools_pos, content_pos);
+}
+
+TEST(DeepseekV32CppTemplate, DropThinkingOnlyWhenLastMessageIsUser) {
+  TokenizerArgs args;
+  args.bos_token("<｜begin▁of▁sentence｜>");
+  DeepseekV32CppTemplate encoder(args);
+
+  // user -> assistant(with reasoning) -> tool
+  ChatMessages messages;
+  messages.emplace_back("user", "weather?");
+
+  Message assistant_msg("assistant", "calling tool");
+  assistant_msg.reasoning_content = "thinking about it";
+  Message::ToolCall tc;
+  tc.id = "1";
+  tc.type = "function";
+  tc.function.name = "get_weather";
+  tc.function.arguments = R"({"city":"beijing"})";
+  assistant_msg.tool_calls = Message::ToolCallVec{tc};
+  messages.push_back(assistant_msg);
+
+  messages.emplace_back("tool", "sunny");
+
+  nlohmann::ordered_json kwargs = nlohmann::json::object();
+  kwargs["thinking"] = true;
+
+  auto prompt = encoder.apply(messages, /*json_tools=*/{}, kwargs);
+  ASSERT_TRUE(prompt.has_value());
+
+  // Last message is "tool", NOT "user",
+  // so thinking should NOT be dropped.
+  EXPECT_NE(prompt->find("thinking about it"), std::string::npos);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/chat_template/jinja_chat_template_test.cpp

@@ -0,0 +1,88 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "jinja_chat_template.h"
+
+#include <gtest/gtest.h>
+
+namespace xllm {
+
+class TestableJinjaChatTemplate : public JinjaChatTemplate {
+ public:
+  TestableJinjaChatTemplate(const TokenizerArgs& args)
+      : JinjaChatTemplate(args) {}
+
+  using JinjaChatTemplate::apply;
+};
+
+TEST(JinjaChatTemplate, OpenChatModel) {
+  // clang-format off
+  const std::string template_str =
+      "<s>"
+      "{% for message in messages %}"
+        "{{ 'GPT4 Correct ' + message['role'] + ': ' + message['content'] + '<|end_of_turn|>'}}"
+      "{% endfor %}"
+      "{% if add_generation_prompt %}{{ 'GPT4 Correct Assistant:' }}{% endif %}";
+
+  nlohmann::ordered_json messages = {
+      {{"role", "system"}, {"content", "you are a helpful assistant."}},
+      {{"role", "user"}, {"content", "hi"}},
+      {{"role", "assistant"}, {"content", "what i can do for you?"}},
+      {{"role", "user"}, {"content", "how are you?"}}};
+  const std::string expected =
+    "<s>"
+    "GPT4 Correct system: you are a helpful assistant.<|end_of_turn|>"
+    "GPT4 Correct user: hi<|end_of_turn|>"
+    "GPT4 Correct assistant: what i can do for you?<|end_of_turn|>"
+    "GPT4 Correct user: how are you?<|end_of_turn|>"
+    "GPT4 Correct Assistant:";
+  // clang-format on
+
+  TokenizerArgs args;
+  args.chat_template(template_str);
+  args.bos_token("");
+  args.eos_token("<|end_of_turn|>");
+  TestableJinjaChatTemplate template_(args);
+  auto result = template_.apply(messages);
+  ASSERT_TRUE(result.has_value());
+
+  EXPECT_EQ(result.value(), expected);
+}
+
+TEST(JinjaChatTemplate, AppliesChatTemplateKwargs) {
+  const std::string template_str =
+      "{% if enable_thinking %}<think>{% endif %}"
+      "{% for message in messages %}"
+      "{{ message['role'] + ': ' + message['content'] }}"
+      "{% endfor %}"
+      "{% if not enable_thinking %}<no_think>{% endif %}";
+
+  nlohmann::ordered_json messages = {
+      {{"role", "user"}, {"content", "describe this image"}}};
+  nlohmann::ordered_json chat_template_kwargs = {{"enable_thinking", false}};
+
+  TokenizerArgs args;
+  args.chat_template(template_str);
+  args.bos_token("");
+  args.eos_token("");
+  TestableJinjaChatTemplate template_(args);
+  const nlohmann::ordered_json tools = nlohmann::json::array();
+  auto result = template_.apply(messages, tools, chat_template_kwargs);
+  ASSERT_TRUE(result.has_value());
+
+  EXPECT_EQ(result.value(), "user: describe this image<no_think>");
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/eplb/CMakeLists.txt

@@ -0,0 +1,13 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    eplb_policy_test
+  SRCS
+    eplb_policy_test.cpp
+  DEPS
+    torch
+    :eplb
+    GTest::gtest_main
+)
+target_link_libraries(eplb_policy_test PRIVATE $<$<BOOL:${USE_NPU}>:c_sec>)

upstream_ref/xllm/tests/core/framework/eplb/eplb_policy_test.cpp

@@ -0,0 +1,44 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "eplb_policy.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "platform/device.h"
+
+namespace xllm {
+
+TEST(EplbPolicyTest, Build) {
+  // use init device to trigger the loading of torch backend for different
+  // devices
+  //  since the allocation of pinnned memory on cpu is still backend-dependent.
+  torch::Device device(Device::type_torch(), 0);
+  std::string rank_table_file;
+  EplbPolicy eplb_policy(5, 4, 1);
+  std::vector<torch::Tensor> tensors;
+  tensors.push_back(torch::arange(0, 16));
+
+  auto expert_load = torch::stack(tensors, 0);
+  expert_load[0] =
+      torch::tensor({100, 100, 100, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 100});
+  auto [rebalance_expert, enable_update_vec] =
+      eplb_policy.rebalance_experts(expert_load);
+  LOG(INFO) << "rebalance_expert:" << rebalance_expert;
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/hf_model_loader_test.cpp

@@ -0,0 +1,234 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "hf_model_loader.h"
+
+#include <gtest/gtest.h>
+
+#include <string>
+#include <type_traits>
+
+#include "core/framework/model/rec_causal_lm.h"
+#include "core/platform/device.h"
+#include "models/model_registry.h"
+
+namespace xllm {
+
+namespace {
+
+using ExpectedRecModelFactory =
+    std::function<std::unique_ptr<RecCausalLM>(const ModelContext& context)>;
+
+static_assert(std::is_same_v<RecModelFactory, ExpectedRecModelFactory>,
+              "RecModelFactory must return std::unique_ptr<RecCausalLM>.");
+static_assert(std::is_base_of_v<CausalLM, RecCausalLM>,
+              "RecCausalLM must derive from CausalLM.");
+
+class DummyRecCausalLM final : public RecCausalLM {
+ public:
+  explicit DummyRecCausalLM(const torch::TensorOptions& options)
+      : options_(options) {}
+
+  ModelOutput forward(const torch::Tensor& tokens,
+                      const torch::Tensor& positions,
+                      std::vector<KVCache>& kv_caches,
+                      const ModelInputParams& parameters) override {
+    UNUSED_PARAMETER(tokens);
+    UNUSED_PARAMETER(positions);
+    UNUSED_PARAMETER(kv_caches);
+    UNUSED_PARAMETER(parameters);
+    return ModelOutput();
+  }
+
+  torch::Tensor logits(const torch::Tensor& hidden_states,
+                       const torch::Tensor& seleted_idxes) override {
+    UNUSED_PARAMETER(hidden_states);
+    UNUSED_PARAMETER(seleted_idxes);
+    return torch::Tensor();
+  }
+
+  void load_model(std::unique_ptr<ModelLoader> loader) override {
+    UNUSED_PARAMETER(loader);
+  }
+
+  torch::Device device() const override { return options_.device(); }
+
+  void prepare_expert_weight(int32_t layer_id,
+                             const std::vector<int32_t>& expert_ids) override {
+    UNUSED_PARAMETER(layer_id);
+    UNUSED_PARAMETER(expert_ids);
+  }
+
+  void update_expert_weight(int32_t layer_id) override {
+    UNUSED_PARAMETER(layer_id);
+  }
+
+  const torch::TensorOptions& options() const override { return options_; }
+
+ private:
+  torch::TensorOptions options_;
+};
+
+}  // namespace
+
+TEST(HFModelLoaderTest, LoadCompressedTensorsFp8StaticConfig) {
+  JsonReader reader;
+  ASSERT_TRUE(reader.parse_text(R"json(
+    {
+      "quantization_config": {
+        "config_groups": {
+          "group_0": {
+            "input_activations": {
+              "dynamic": false,
+              "num_bits": 8,
+              "type": "float"
+            },
+            "weights": {
+              "num_bits": 8,
+              "type": "float"
+            }
+          }
+        },
+        "ignore": [
+          "lm_head",
+          "model.layers.1.mlp.down_proj"
+        ],
+        "quant_method": "compressed-tensors"
+      }
+    }
+  )json"));
+
+  QuantArgs quant_args;
+  if (Device::type_str() == "cuda") {
+    ASSERT_TRUE(load_quant_cfg(reader, quant_args));
+    EXPECT_EQ(quant_args.quant_method(), kQuantMethodFp8);
+    EXPECT_EQ(quant_args.bits(), 8);
+    EXPECT_EQ(quant_args.moe_weight_bits(), 8);
+    EXPECT_FALSE(quant_args.activation_dynamic());
+    ASSERT_EQ(quant_args.ignored_modules().size(), 2);
+    EXPECT_EQ(quant_args.ignored_modules()[0], "lm_head");
+    EXPECT_EQ(quant_args.ignored_modules()[1], "model.layers.1.mlp.down_proj");
+  }
+}
+
+TEST(HFModelLoaderTest, KeepLegacyFp8ConfigUnchanged) {
+  JsonReader reader;
+  ASSERT_TRUE(reader.parse_text(R"json(
+    {
+      "quantization_config": {
+        "activation_scheme": "static",
+        "quant_method": "fp8"
+      }
+    }
+  )json"));
+
+  QuantArgs quant_args;
+  ASSERT_TRUE(load_quant_cfg(reader, quant_args));
+  EXPECT_EQ(quant_args.quant_method(), kQuantMethodFp8);
+  EXPECT_FALSE(quant_args.activation_dynamic());
+}
+
+TEST(HFModelLoaderTest, RegisterRecFactoryAcceptsRecCausalLmReturnType) {
+  const std::string factory_name = "rec_causallm_factory_contract_test";
+  RecModelFactory unregistered_factory =
+      ModelRegistry::get_rec_model_factory(factory_name);
+  EXPECT_FALSE(static_cast<bool>(unregistered_factory));
+
+  ModelRegistry::register_rec_model_factory(
+      factory_name,
+      [](const ModelContext& context) -> std::unique_ptr<RecCausalLM> {
+        UNUSED_PARAMETER(context);
+        return nullptr;
+      });
+
+  RecModelFactory factory = ModelRegistry::get_rec_model_factory(factory_name);
+  EXPECT_TRUE(static_cast<bool>(factory));
+}
+
+TEST(HFModelLoaderTest, RecFactoryCreatesRecCausalLmInstance) {
+  const std::string kFactoryName = "rec_causallm_instance_contract_test";
+  ModelRegistry::register_rec_model_factory(
+      kFactoryName,
+      [](const ModelContext& context) -> std::unique_ptr<RecCausalLM> {
+        UNUSED_PARAMETER(context);
+        const torch::TensorOptions options =
+            torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
+        return std::make_unique<DummyRecCausalLM>(options);
+      });
+
+  RecModelFactory factory = ModelRegistry::get_rec_model_factory(kFactoryName);
+  ASSERT_TRUE(static_cast<bool>(factory));
+
+  ModelContext context;
+  std::unique_ptr<RecCausalLM> rec_model = factory(context);
+  ASSERT_NE(rec_model, nullptr);
+
+  CausalLM* causal_model = dynamic_cast<CausalLM*>(rec_model.get());
+  EXPECT_NE(causal_model, nullptr);
+  EXPECT_EQ(rec_model->device(), torch::Device(torch::kCPU));
+}
+
+#if defined(USE_NPU)
+TEST(HFModelLoaderTest, Qwen35MtpModelArgsFromDenseConfig) {
+  auto loader = ModelRegistry::get_model_args_loader("qwen3_5_mtp");
+  ASSERT_TRUE(loader != nullptr);
+
+  JsonReader reader;
+  ASSERT_TRUE(reader.parse_text(R"json(
+    {
+      "model_type": "qwen3_5",
+      "text_config": {
+        "mtp_num_hidden_layers": 1,
+        "layer_types": ["linear_attention"]
+      }
+    }
+  )json"));
+
+  ModelArgs args;
+  ASSERT_TRUE(loader(reader, &args));
+  EXPECT_EQ(args.model_type(), "qwen3_5_mtp");
+  EXPECT_EQ(args.num_nextn_predict_layers(), 1);
+  EXPECT_EQ(args.n_layers(), 1);
+  ASSERT_EQ(args.layer_types().size(), 1);
+  EXPECT_EQ(args.layer_types()[0], "full_attention");
+}
+
+TEST(HFModelLoaderTest, Qwen35MtpModelArgsFromMoeConfig) {
+  auto loader = ModelRegistry::get_model_args_loader("qwen3_5_moe_mtp");
+  ASSERT_TRUE(loader != nullptr);
+
+  JsonReader reader;
+  ASSERT_TRUE(reader.parse_text(R"json(
+    {
+      "model_type": "qwen3_5_moe",
+      "text_config": {
+        "mtp_num_hidden_layers": 2,
+        "layer_types": ["linear_attention", "linear_attention"]
+      }
+    }
+  )json"));
+
+  ModelArgs args;
+  ASSERT_TRUE(loader(reader, &args));
+  EXPECT_EQ(args.model_type(), "qwen3_5_moe_mtp");
+  EXPECT_EQ(args.num_nextn_predict_layers(), 2);
+  EXPECT_EQ(args.n_layers(), 2);
+  ASSERT_EQ(args.layer_types().size(), 2);
+  EXPECT_EQ(args.layer_types()[0], "full_attention");
+  EXPECT_EQ(args.layer_types()[1], "full_attention");
+}
+#endif
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/kv_cache/CMakeLists.txt

@@ -0,0 +1,18 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    embedding_cache_test
+  SRCS
+    embedding_cache_test.cpp
+  DEPS
+    :kv_cache
+    GTest::gtest_main
+)
+target_link_libraries(embedding_cache_test
+                      PUBLIC
+                      Python::Python
+                      $<$<BOOL:${USE_NPU}>:ascendcl>
+                      $<$<BOOL:${USE_NPU}>:hccl>
+                      $<$<BOOL:${USE_NPU}>:c_sec>
+                      $<$<BOOL:${USE_NPU}>:nnopbase>)

upstream_ref/xllm/tests/core/framework/kv_cache/embedding_cache_test.cpp

@@ -0,0 +1,145 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "embedding_cache.h"
+
+#include <gtest/gtest.h>
+
+#include "platform/device.h"
+
+namespace xllm {
+
+namespace {
+
+bool tensor_equal(const torch::Tensor& lhs, const torch::Tensor& rhs) {
+  return lhs.defined() && rhs.defined() && torch::equal(lhs, rhs);
+}
+
+}  // namespace
+
+TEST(EmbeddingCacheTest, WritePrefillTargetContextAndClear) {
+  torch::Device device(Device::type_torch(), 0);
+  EmbeddingCache cache(/*total_nums=*/4);
+
+  std::vector<int32_t> ids = {3, 2};
+  std::vector<std::string> request_ids = {"req_0", "req_1"};
+  torch::Tensor target_tokens = torch::tensor({31, 41}, torch::kInt);
+  torch::Tensor target_embeddings = torch::tensor({{1.0f, 2.0f}, {3.0f, 4.0f}});
+
+  cache.write_prefill_target_context(
+      ids, request_ids, target_tokens, target_embeddings);
+
+  std::vector<EmbeddingCache::DecodeState> states =
+      cache.read_decode_states(ids, request_ids);
+  ASSERT_EQ(states.size(), ids.size());
+  EXPECT_TRUE(states[0].valid);
+  EXPECT_EQ(states[0].request_id, "req_0");
+  EXPECT_EQ(states[0].token_id, 31);
+  EXPECT_EQ(states[0].position_offset, 0);
+  EXPECT_FALSE(states[0].all_draft_accepted);
+  EXPECT_EQ(states[0].prev_token_id, -1);
+  EXPECT_TRUE(tensor_equal(states[0].embedding, target_embeddings[0]));
+  EXPECT_TRUE(states[1].valid);
+  EXPECT_EQ(states[1].token_id, 41);
+  EXPECT_EQ(states[1].position_offset, 0);
+  EXPECT_TRUE(tensor_equal(states[1].embedding, target_embeddings[1]));
+
+  cache.clear(ids);
+  states = cache.read_decode_states(ids, request_ids);
+  EXPECT_FALSE(states[0].valid);
+  EXPECT_EQ(states[0].token_id, 0);
+  EXPECT_EQ(states[0].position_offset, 0);
+  EXPECT_FALSE(states[0].embedding.defined());
+  EXPECT_FALSE(states[1].valid);
+  EXPECT_EQ(states[1].token_id, 0);
+  EXPECT_EQ(states[1].position_offset, 0);
+  EXPECT_FALSE(states[1].embedding.defined());
+}
+
+TEST(EmbeddingCacheTest, WritePrefillTargetContextSelectsEmbeddings) {
+  EmbeddingCache cache(/*total_nums=*/4);
+
+  std::vector<int32_t> ids = {1, 2};
+  std::vector<std::string> request_ids = {"req_0", "req_1"};
+  torch::Tensor target_tokens = torch::tensor({51, 61}, torch::kInt);
+  torch::Tensor full_embeddings =
+      torch::tensor({{1.0f, 1.1f}, {2.0f, 2.1f}, {3.0f, 3.1f}});
+  torch::Tensor selected_idxes = torch::tensor({2, 0}, torch::kInt);
+
+  cache.write_prefill_target_context(
+      ids, request_ids, target_tokens, full_embeddings, selected_idxes);
+
+  std::vector<EmbeddingCache::DecodeState> states =
+      cache.read_decode_states(ids, request_ids);
+  ASSERT_EQ(states.size(), ids.size());
+  EXPECT_EQ(states[0].token_id, 51);
+  EXPECT_TRUE(tensor_equal(states[0].embedding, full_embeddings[2]));
+  EXPECT_EQ(states[1].token_id, 61);
+  EXPECT_TRUE(tensor_equal(states[1].embedding, full_embeddings[0]));
+}
+
+TEST(EmbeddingCacheTest, WriteValidateTargetContext) {
+  torch::Device device(Device::type_torch(), 0);
+  EmbeddingCache cache(/*total_nums=*/2);
+  std::vector<int32_t> ids = {0, 1};
+  std::vector<std::string> request_ids = {"req_0", "req_1"};
+  torch::Tensor accepted_tokens =
+      torch::tensor({{11, 12, 13}, {21, -1, -1}}, torch::kInt);
+  torch::Tensor accepted_embeddings =
+      torch::tensor({{{1.0f, 1.1f}, {1.2f, 1.3f}, {1.4f, 1.5f}},
+                     {{2.0f, 2.1f}, {2.2f, 2.3f}, {2.4f, 2.5f}}});
+
+  cache.write_target_context(ids,
+                             request_ids,
+                             accepted_tokens,
+                             accepted_embeddings,
+                             /*num_speculative_tokens=*/2);
+
+  std::vector<EmbeddingCache::DecodeState> states =
+      cache.read_decode_states(ids, request_ids);
+  EXPECT_EQ(states[0].token_id, 13);
+  EXPECT_EQ(states[0].position_offset, 2);
+  EXPECT_TRUE(states[0].all_draft_accepted);
+  EXPECT_EQ(states[0].prev_token_id, 12);
+  EXPECT_TRUE(
+      tensor_equal(states[0].prev_embedding, accepted_embeddings[0][1]));
+  EXPECT_TRUE(tensor_equal(states[0].embedding, accepted_embeddings[0][2]));
+
+  EXPECT_EQ(states[1].token_id, 21);
+  EXPECT_EQ(states[1].position_offset, 0);
+  EXPECT_FALSE(states[1].all_draft_accepted);
+  EXPECT_EQ(states[1].prev_token_id, -1);
+  EXPECT_TRUE(tensor_equal(states[1].embedding, accepted_embeddings[1][0]));
+}
+
+TEST(EmbeddingCacheTest, RequestMismatchMaterializesMissingState) {
+  EmbeddingCache cache(/*total_nums=*/2);
+  std::vector<int32_t> ids = {0};
+  std::vector<std::string> request_ids = {"old_req"};
+  torch::Tensor target_tokens = torch::tensor({31}, torch::kInt);
+  torch::Tensor target_embeddings = torch::tensor({{1.0f, 2.0f}});
+
+  cache.write_prefill_target_context(
+      ids, request_ids, target_tokens, target_embeddings);
+
+  std::vector<EmbeddingCache::DecodeState> states =
+      cache.read_decode_states(ids, {"new_req"});
+  ASSERT_EQ(states.size(), ids.size());
+  EXPECT_FALSE(states[0].valid);
+  EXPECT_EQ(states[0].token_id, 0);
+  EXPECT_FALSE(states[0].embedding.defined());
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/kv_cache_transfer/CMakeLists.txt

@@ -0,0 +1,38 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    push_route_test
+  SRCS
+    push_route_test.cpp
+  DEPS
+    :push_route
+    GTest::gtest_main
+)
+
+cc_test(
+  NAME
+    pd_topology_guard_test
+  SRCS
+    pd_topology_guard_test.cpp
+  DEPS
+    :pd_topology_guard
+    GTest::gtest_main
+)
+
+if(USE_NPU OR USE_MLU)
+cc_test(
+  NAME
+    mooncake_transfer_engine_test
+  SRCS
+    mooncake_transfer_engine_test.cpp
+  DEPS
+    :kv_cache_transfer
+    :xllm_server
+    GTest::gtest_main
+)
+
+# Resolve static link order between xtensor and xllm_server for this test target.
+target_link_libraries(mooncake_transfer_engine_test PRIVATE
+  "$<LINK_GROUP:RESCAN,xtensor,xllm_server>")
+endif()

upstream_ref/xllm/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp

@@ -0,0 +1,191 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/kv_cache_transfer/mooncake_transfer_engine.h"
+
+#include <brpc/controller.h>
+#include <gtest/gtest.h>
+
+#include <unordered_map>
+#include <vector>
+
+#include "framework/kv_cache_transfer/kv_cache_transfer.h"
+
+#define private public
+#include "framework/kv_cache_transfer/mooncake_kv_cache_transfer.h"
+#undef private
+
+namespace xllm {
+
+namespace {
+
+TransferKVInfo make_info(int32_t dst_dp_size,
+                         int32_t dst_tp_size,
+                         int32_t dst_dp_rank) {
+  TransferKVInfo info;
+  info.request_id = "req";
+  info.local_blocks_ids = {11, 12};
+  info.remote_blocks_ids = {21, 22};
+  info.dp_rank = dst_dp_rank;
+  info.remote_instance_info.dp_size = dst_dp_size;
+
+  int32_t dst_world_size = dst_dp_size * dst_tp_size;
+  for (int32_t i = 0; i < dst_world_size; ++i) {
+    info.remote_instance_info.cluster_ids.emplace_back(
+        static_cast<uint64_t>(100 + i));
+    info.remote_instance_info.addrs.emplace_back("addr_" + std::to_string(i));
+    info.remote_instance_info.k_cache_ids.emplace_back(200 + i);
+    info.remote_instance_info.v_cache_ids.emplace_back(300 + i);
+  }
+
+  return info;
+}
+
+ParallelArgs make_args(int32_t rank, int32_t world_size, int32_t dp_size) {
+  return ParallelArgs(rank, world_size, dp_size, nullptr);
+}
+
+void expect_same_merge(
+    const std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo>& lhs,
+    const std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo>& rhs) {
+  ASSERT_EQ(lhs.size(), rhs.size());
+  for (const auto& [key, lhs_info] : lhs) {
+    auto it = rhs.find(key);
+    ASSERT_NE(it, rhs.end());
+    const KVCacheTransfer::KVCacheInfo& rhs_info = it->second;
+    EXPECT_EQ(lhs_info.dst_cluster_id, rhs_info.dst_cluster_id);
+    EXPECT_EQ(lhs_info.dst_addr, rhs_info.dst_addr);
+    EXPECT_EQ(lhs_info.dst_k_cache_id, rhs_info.dst_k_cache_id);
+    EXPECT_EQ(lhs_info.dst_v_cache_id, rhs_info.dst_v_cache_id);
+    EXPECT_EQ(lhs_info.src_blocks, rhs_info.src_blocks);
+    EXPECT_EQ(lhs_info.dst_blocks, rhs_info.dst_blocks);
+  }
+}
+
+}  // namespace
+
+TEST(MooncakeTransferEngineServiceTest, OpenSessionRejectsMissingAddr) {
+  MooncakeTransferEngineService service;
+  proto::SessionInfo request;
+  proto::Status response;
+  brpc::Controller cntl;
+
+  service.OpenSession(&cntl, &request, &response, nullptr);
+
+  EXPECT_FALSE(response.ok());
+}
+
+TEST(MooncakeTransferEngineServiceTest, CloseSessionRejectsMissingAddr) {
+  MooncakeTransferEngineService service;
+  proto::SessionInfo request;
+  proto::Status response;
+  brpc::Controller cntl;
+
+  service.CloseSession(&cntl, &request, &response, nullptr);
+
+  EXPECT_FALSE(response.ok());
+}
+
+TEST(MooncakeTransferEngineServiceTest, CloseSessionWithoutHandleReturnsTrue) {
+  MooncakeTransferEngineService service;
+  proto::SessionInfo request;
+  request.set_addr("127.0.0.1:5001");
+  proto::Status response;
+  brpc::Controller cntl;
+
+  service.CloseSession(&cntl, &request, &response, nullptr);
+
+  EXPECT_TRUE(response.ok());
+}
+
+#if defined(USE_MLU)
+TEST(MooncakeKVCacheTransferDefaultTest, OwnerRankMergesSingleDst) {
+  MooncakeKVCacheTransferDefault transfer(
+      0, 0, torch::Device(torch::kCPU), "test");
+  transfer.has_v_cache_ = false;
+
+  const TransferKVInfo info = make_info(1, 3, 0);
+  const ParallelArgs parallel_args = make_args(2, 8, 1);
+  std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo> merged_kv_infos;
+
+  transfer.merge_kv_blocks(merged_kv_infos, {info}, parallel_args);
+
+  ASSERT_EQ(merged_kv_infos.size(), 1U);
+  const KVCacheTransfer::KVCacheInfo& kv_info = merged_kv_infos.begin()->second;
+  EXPECT_EQ(kv_info.dst_cluster_id, 102U);
+  EXPECT_EQ(kv_info.dst_addr, "addr_2");
+  EXPECT_EQ(kv_info.dst_k_cache_id, 202);
+  EXPECT_EQ(kv_info.dst_v_cache_id, 302);
+  EXPECT_EQ(kv_info.src_blocks, info.local_blocks_ids);
+  EXPECT_EQ(kv_info.dst_blocks, info.remote_blocks_ids);
+}
+
+TEST(MooncakeKVCacheTransferDefaultTest, WrappedOwnerRankKeepsMerge) {
+  MooncakeKVCacheTransferDefault transfer(
+      0, 0, torch::Device(torch::kCPU), "test");
+  transfer.has_v_cache_ = false;
+
+  const TransferKVInfo info = make_info(2, 3, 1);
+  const ParallelArgs parallel_args = make_args(5, 8, 1);
+  std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo> merged_kv_infos;
+
+  transfer.merge_kv_blocks(merged_kv_infos, {info}, parallel_args);
+
+  ASSERT_EQ(merged_kv_infos.size(), 1U);
+  const KVCacheTransfer::KVCacheInfo& kv_info = merged_kv_infos.begin()->second;
+  EXPECT_EQ(kv_info.dst_cluster_id, 105U);
+  EXPECT_EQ(kv_info.dst_addr, "addr_5");
+  EXPECT_EQ(kv_info.dst_k_cache_id, 205);
+  EXPECT_EQ(kv_info.dst_v_cache_id, 305);
+  EXPECT_EQ(kv_info.src_blocks, info.local_blocks_ids);
+  EXPECT_EQ(kv_info.dst_blocks, info.remote_blocks_ids);
+}
+
+TEST(MooncakeKVCacheTransferDefaultTest, HasVCacheUsesBaseMerge) {
+  MooncakeKVCacheTransferDefault transfer(
+      0, 0, torch::Device(torch::kCPU), "test");
+  transfer.has_v_cache_ = true;
+
+  const TransferKVInfo info = make_info(2, 3, 1);
+  const ParallelArgs parallel_args = make_args(5, 8, 1);
+  std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo> merged_kv_infos;
+  std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo> base_kv_infos;
+
+  transfer.merge_kv_blocks(merged_kv_infos, {info}, parallel_args);
+  transfer.KVCacheTransfer::merge_kv_blocks(
+      base_kv_infos, {info}, parallel_args);
+
+  expect_same_merge(merged_kv_infos, base_kv_infos);
+}
+
+TEST(MooncakeKVCacheTransferDefaultTest, SmallSrcTpUsesBaseMerge) {
+  MooncakeKVCacheTransferDefault transfer(
+      0, 0, torch::Device(torch::kCPU), "test");
+  transfer.has_v_cache_ = false;
+
+  const TransferKVInfo info = make_info(1, 4, 0);
+  const ParallelArgs parallel_args = make_args(1, 2, 1);
+  std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo> merged_kv_infos;
+  std::unordered_map<std::string, KVCacheTransfer::KVCacheInfo> base_kv_infos;
+
+  transfer.merge_kv_blocks(merged_kv_infos, {info}, parallel_args);
+  transfer.KVCacheTransfer::merge_kv_blocks(
+      base_kv_infos, {info}, parallel_args);
+
+  expect_same_merge(merged_kv_infos, base_kv_infos);
+}
+#endif
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/kv_cache_transfer/pd_topology_guard_test.cpp

@@ -0,0 +1,160 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/kv_cache_transfer/pd_topology_guard.h"
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace xllm {
+
+namespace {
+
+void set_death_style() { GTEST_FLAG_SET(death_test_style, "threadsafe"); }
+
+InstanceInfo make_info(int32_t dp_size,
+                       const std::vector<uint64_t>& cluster_ids) {
+  InstanceInfo info;
+  info.dp_size = dp_size;
+  info.cluster_ids = cluster_ids;
+  return info;
+}
+
+TEST(PdTopologyGuardTest, HomoTopoBypass) {
+  const InstanceInfo local_info = make_info(2, {0, 1, 2, 3});
+  const InstanceInfo remote_info = make_info(2, {0, 1, 2, 3});
+
+  const PdTopo topo = get_pd_topo(local_info);
+  EXPECT_EQ(topo.dp_size, 2);
+  EXPECT_EQ(topo.tp_size, 2);
+
+  const PdTopoResult result =
+      check_pd_topo(local_info, remote_info, "PULL", false);
+  EXPECT_EQ(result.status, PdTopoStatus::ALLOW_HOMO);
+  EXPECT_TRUE(result.reason.empty());
+}
+
+TEST(PdTopologyGuardTest, TryGetPdTopoReturnTopo) {
+  const InstanceInfo info = make_info(2, {0, 1, 2, 3});
+
+  PdTopo topo;
+  std::string reason;
+  EXPECT_TRUE(try_get_pd_topo(info, &topo, &reason));
+  EXPECT_EQ(topo.dp_size, 2);
+  EXPECT_EQ(topo.tp_size, 2);
+  EXPECT_TRUE(reason.empty());
+}
+
+TEST(PdTopologyGuardTest, HeteroTopoNeedMla) {
+  const InstanceInfo local_info = make_info(2, {0, 1, 2, 3});
+  const InstanceInfo remote_info = make_info(1, {0, 1, 2, 3});
+
+  const PdTopoResult result =
+      check_pd_topo(local_info, remote_info, "PUSH", false);
+  EXPECT_EQ(result.status, PdTopoStatus::DENY_HETERO);
+  EXPECT_EQ(result.reason, "hetero pd requires enable_mla=true");
+}
+
+TEST(PdTopologyGuardTest, HeteroTopoNeedPushKv) {
+  const InstanceInfo local_info = make_info(2, {0, 1, 2, 3});
+  const InstanceInfo remote_info = make_info(1, {0, 1, 2, 3});
+
+  const PdTopoResult result =
+      check_pd_topo(local_info, remote_info, "PULL", true);
+  EXPECT_EQ(result.status, PdTopoStatus::DENY_HETERO);
+  EXPECT_EQ(result.reason, "hetero pd requires kv_mode=PUSH");
+}
+
+TEST(PdTopologyGuardTest, HeteroTopoAllowOnPushMla) {
+  const InstanceInfo local_info = make_info(2, {0, 1, 2, 3});
+  const InstanceInfo remote_info = make_info(1, {0, 1, 2, 3});
+
+  const PdTopoResult result =
+      check_pd_topo(local_info, remote_info, "PUSH", true);
+  EXPECT_EQ(result.status, PdTopoStatus::ALLOW_HETERO);
+  EXPECT_TRUE(result.reason.empty());
+}
+
+TEST(PdTopologyGuardTest, CheckPdTopoRejectInvalidLocalTopo) {
+  const InstanceInfo local_info = make_info(0, {0, 1, 2, 3});
+  const InstanceInfo remote_info = make_info(1, {0, 1, 2, 3});
+
+  const PdTopoResult result =
+      check_pd_topo(local_info, remote_info, "PUSH", true);
+  EXPECT_EQ(result.status, PdTopoStatus::INVALID_LOCAL);
+  EXPECT_EQ(result.reason,
+            "invalid local pd topo: dp_size must be greater than 0");
+}
+
+TEST(PdTopologyGuardTest, CheckPdTopoRejectInvalidRemoteTopo) {
+  const InstanceInfo local_info = make_info(1, {0, 1, 2, 3});
+  const InstanceInfo remote_info = make_info(2, {0, 1, 2});
+
+  const PdTopoResult result =
+      check_pd_topo(local_info, remote_info, "PUSH", true);
+  EXPECT_EQ(result.status, PdTopoStatus::INVALID_REMOTE);
+  EXPECT_EQ(result.reason,
+            "invalid remote pd topo: cluster_ids.size() must be divisible by "
+            "dp_size");
+}
+
+TEST(PdTopologyGuardTest, TryGetPdTopoRejectBadClusterSplit) {
+  const InstanceInfo info = make_info(2, {0, 1, 2});
+
+  PdTopo topo;
+  std::string reason;
+  EXPECT_FALSE(try_get_pd_topo(info, &topo, &reason));
+  EXPECT_EQ(reason, "cluster_ids.size() must be divisible by dp_size");
+}
+
+TEST(PdTopologyGuardTest, TryGetPdTopoRejectEmptyClusterIds) {
+  const InstanceInfo info = make_info(2, {});
+
+  PdTopo topo;
+  std::string reason;
+  EXPECT_FALSE(try_get_pd_topo(info, &topo, &reason));
+  EXPECT_EQ(reason, "cluster_ids must not be empty");
+}
+
+TEST(PdTopologyGuardTest, TryGetPdTopoRejectZeroDpSize) {
+  const InstanceInfo info = make_info(0, {0, 1, 2, 3});
+
+  PdTopo topo;
+  std::string reason;
+  EXPECT_FALSE(try_get_pd_topo(info, &topo, &reason));
+  EXPECT_EQ(reason, "dp_size must be greater than 0");
+}
+
+TEST(PdTopologyGuardTest, GetPdTopoRejectBadClusterSplit) {
+  set_death_style();
+  const InstanceInfo info = make_info(2, {0, 1, 2});
+
+  EXPECT_DEATH(get_pd_topo(info),
+               "cluster_ids.size\\(\\) must be divisible by dp_size");
+}
+
+TEST(PdTopologyGuardTest, GetPdTopoRejectEmptyClusterIds) {
+  set_death_style();
+  const InstanceInfo info = make_info(2, {});
+
+  EXPECT_DEATH(get_pd_topo(info), "cluster_ids must not be empty");
+}
+
+}  // namespace
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/kv_cache_transfer/push_route_test.cpp

@@ -0,0 +1,74 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/kv_cache_transfer/push_route.h"
+
+#include <gtest/gtest.h>
+
+#include <vector>
+
+namespace xllm {
+
+TEST(PushRouteTest, SrcTpLessThanDstTpKeepsMulticast) {
+  EXPECT_FALSE(use_push_owner(2, 8));
+
+  const std::vector<int32_t> dst_ranks = get_dst_ranks(1, 2, 8, 3);
+  const std::vector<int32_t> expect_ranks = {25, 27, 29, 31};
+  EXPECT_EQ(dst_ranks, expect_ranks);
+}
+
+TEST(PushRouteTest, InvalidTpSizeNotUseOwnerAndReturnEmpty) {
+  EXPECT_FALSE(use_push_owner(0, 4));
+  EXPECT_FALSE(use_push_owner(4, 0));
+  EXPECT_FALSE(use_push_owner(-1, 2));
+  EXPECT_FALSE(use_push_owner(2, -1));
+
+  const std::vector<int32_t> dst_ranks = get_dst_ranks(0, 0, 4, 0);
+  EXPECT_TRUE(dst_ranks.empty());
+}
+
+TEST(PushRouteTest, SrcTpEqualsDstTpNotUseOwner) {
+  EXPECT_FALSE(use_push_owner(4, 4));
+
+  const std::vector<int32_t> dst_ranks = get_dst_ranks(2, 4, 4, 1);
+  const std::vector<int32_t> expect_ranks = {6};
+  EXPECT_EQ(dst_ranks, expect_ranks);
+}
+
+TEST(PushRouteTest, SrcTpGreaterThanDstTpUsesOwnerRouting) {
+  EXPECT_TRUE(use_push_owner(8, 3));
+
+  const std::vector<int32_t> owner_ranks = get_dst_ranks(2, 8, 3, 2);
+  const std::vector<int32_t> expect_owner_ranks = {8};
+  EXPECT_EQ(owner_ranks, expect_owner_ranks);
+
+  const std::vector<int32_t> wrapped_owner_ranks = get_dst_ranks(5, 8, 3, 2);
+  const std::vector<int32_t> expect_wrapped_owner_ranks = {8};
+  EXPECT_EQ(wrapped_owner_ranks, expect_wrapped_owner_ranks);
+}
+
+TEST(PushRouteTest, HeteroTpTwoToOneKeepsOddDpRoute) {
+  const std::vector<int32_t> odd_dp_ranks = get_dst_ranks(1, 2, 1, 3);
+  const std::vector<int32_t> expect_odd_dp_ranks = {3};
+  EXPECT_EQ(odd_dp_ranks, expect_odd_dp_ranks);
+}
+
+TEST(PushRouteTest, DstDpRankOffsetApplied) {
+  const std::vector<int32_t> dst_ranks = get_dst_ranks(1, 6, 4, 3);
+  const std::vector<int32_t> expect_ranks = {13};
+  EXPECT_EQ(dst_ranks, expect_ranks);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/parallel_state/CMakeLists.txt

@@ -0,0 +1,57 @@
+include(cc_test)
+
+if(USE_NPU)
+  cc_test(
+    NAME
+      mapping_npu_test
+    SRCS
+      mapping_npu_test.cpp
+    DEPS
+      parallel_state
+      absl::synchronization
+      absl::time
+      GTest::gtest_main
+      xllm_atb_layers
+      nnopbase
+      ascendcl
+      atb
+      c_sec
+      spdlog::spdlog
+  )
+
+  cc_test(
+    NAME
+      npu_dp_ep_padding_test
+    SRCS
+      npu_dp_ep_padding_test.cpp
+    DEPS
+      parallel_state
+      torch
+      absl::synchronization
+      absl::time
+      GTest::gtest_main
+  )
+  target_link_libraries(npu_dp_ep_padding_test
+                        PUBLIC Python::Python
+                               ascendcl
+                               hccl
+                               c_sec
+                               nnopbase)
+endif()
+
+if(USE_MLU)
+  # Add test for reduce_scatter
+  # This test must exist individually, because it contains forked processes
+  # which does not allow any device init on main process
+  cc_test(
+    NAME
+      parallel_state_test
+    SRCS
+      parallel_state_test.cpp
+    DEPS
+      :parallel_state
+      GTest::gtest_main
+      torch
+      glog::glog
+  )
+endif()

upstream_ref/xllm/tests/core/framework/parallel_state/mapping_npu_test.cpp

@@ -0,0 +1,51 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "mapping_npu.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+namespace xllm {
+
+MappingNPU::Options get_mapping_options() {
+  MappingNPU::Options options;
+  options.dp_size(2)
+      .tp_size(8)
+      .moe_tp_size(2)
+      .moe_ep_size(8)
+      .pp_size(1)
+      .sp_size(1);
+  return options;
+}
+
+TEST(TestMappingNPU, ToJson) {
+  std::string rank_table_file;
+  MappingNPU::Options options = get_mapping_options();
+  MappingNPU mapping(rank_table_file, 16, 6, options);
+  nlohmann::json data = mapping.to_json();
+  LOG(INFO) << "Mapping INFO:\n" << data.dump(2);
+  nlohmann::json attn_dp = data["attnDp"];
+  int32_t attn_dp_group_id = attn_dp["groupId"];
+  EXPECT_EQ(attn_dp_group_id, 6);
+  nlohmann::json attn_tp = data["attnTp"];
+  int32_t attn_tp_group_id = attn_tp["groupId"];
+  EXPECT_EQ(attn_tp_group_id, 0);
+  nlohmann::json mlp_tp = data["mlpTp"];
+  int32_t mlp_tp_buffer_size = mlp_tp["bufferSize"];
+  EXPECT_EQ(mlp_tp_buffer_size, 128);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/parallel_state/npu_dp_ep_padding_test.cpp

@@ -0,0 +1,63 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/parallel_state/npu_dp_ep_padding.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include "common/global_flags.h"
+#include "framework/parallel_state/mapping_npu.h"
+
+namespace xllm {
+
+MappingNPU::Options get_mapping_options() {
+  MappingNPU::Options options;
+  options.dp_size(2)
+      .tp_size(8)
+      .moe_tp_size(16)
+      .moe_ep_size(1)
+      .pp_size(1)
+      .sp_size(1);
+  return options;
+}
+
+TEST(DpEpPaddingTest, Build) {
+  std::string rank_table_file;
+  MappingNPU::Options options = get_mapping_options();
+  MappingNPU mapping(rank_table_file, 16, 0, options);
+  nlohmann::json data = mapping.to_json();
+  torch::Tensor token_size_per_dp_group = torch::tensor({10, 10});
+  DpEpPadding dp_ep_padding(token_size_per_dp_group,
+                            8,
+                            data,
+                            torch::Device(torch::kCPU),
+                            torch::Dtype(torch::kInt32),
+                            true);
+  DpEpPaddingData dp_ep_padding_data = dp_ep_padding.build();
+  LOG(INFO) << "attn_padding_idx:" << dp_ep_padding_data.attn_padding_idx();
+  LOG(INFO) << "attn_unpadding_idx:" << dp_ep_padding_data.attn_unpadding_idx();
+  LOG(INFO) << "ffn_padding_idx:" << dp_ep_padding_data.ffn_padding_idx();
+  LOG(INFO) << "ffn_unpadding_idx:" << dp_ep_padding_data.ffn_unpadding_idx();
+  LOG(INFO) << "lm_head_skip_padding_token_indices:"
+            << dp_ep_padding_data.lm_head_skip_padding_token_indices();
+  LOG(INFO) << "gather_prenorm_idx:" << dp_ep_padding_data.gather_prenorm_idx();
+  LOG(INFO) << "padding_idx:" << dp_ep_padding_data.padding_idx();
+  LOG(INFO) << "un_padding_idx:" << dp_ep_padding_data.un_padding_idx();
+  LOG(INFO) << "dynamic_ep_idx:" << dp_ep_padding_data.dynamic_ep_idx();
+  LOG(INFO) << "moe_idx:" << dp_ep_padding_data.moe_idx();
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/parallel_state/parallel_state_test.cpp

@@ -0,0 +1,462 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/parallel_state/parallel_state.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <sys/wait.h>
+#include <torch/torch.h>
+#include <unistd.h>
+
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include "platform/device.h"
+
+#if defined(USE_MLU)
+#include "framework/parallel_state/mlu_process_group.h"
+#elif defined(USE_CUDA)
+#include "framework/parallel_state/cuda_process_group.h"
+#endif
+
+namespace xllm {
+namespace parallel_state {
+namespace test {
+
+// Helper function to create ProcessGroup for multi-device testing
+std::unique_ptr<xllm::ProcessGroup> create_test_process_group(
+    int rank,
+    int world_size,
+    int port,
+    const std::string& host,
+    const torch::Device& device) {
+  return xllm::create_process_group(static_cast<int32_t>(rank),
+                                    static_cast<int32_t>(world_size),
+                                    static_cast<int32_t>(world_size),
+                                    static_cast<int32_t>(port),
+                                    false,
+                                    host,
+                                    "reduce_scatter_test_group",
+                                    device);
+}
+
+// Test parameters structure to pass to child processes
+struct TestParams {
+  int32_t rank;
+  int32_t world_size;
+  int32_t port;
+  std::string host;
+  int32_t device_index;
+  int64_t input_size;
+  int64_t hidden_dim;
+  bool test_padding;
+};
+
+struct AllGatherBaseTestParams {
+  int32_t rank;
+  int32_t world_size;
+  int32_t port;
+  std::string host;
+  int32_t device_index;
+  int64_t input_size;
+  int64_t hidden_dim;
+};
+// Child process test function
+int run_reduce_scatter_test_child(const TestParams& params) {
+  try {
+    // Set device
+    xllm::Device xllm_device(params.device_index);
+    xllm_device.set_device();
+    torch::Device device = xllm_device.unwrap();
+
+    // Create ProcessGroup
+    auto process_group = create_test_process_group(
+        params.rank, params.world_size, params.port, params.host, device);
+
+    if (!process_group) {
+      LOG(ERROR) << "Rank " << params.rank << ": Failed to create ProcessGroup";
+      return 1;
+    }
+
+    LOG(INFO) << "Rank " << params.rank
+              << ": ProcessGroup created successfully";
+
+    // Create tensor options
+    auto options = torch::TensorOptions()
+                       .dtype(torch::kFloat32)
+                       .device(device)
+                       .requires_grad(false);
+
+    // Create test input tensor
+    // Each rank creates a tensor with values equal to (rank + 1)
+    // This allows us to verify that reduce_scatter correctly sums the values
+    torch::Tensor input = torch::full({params.input_size, params.hidden_dim},
+                                      static_cast<float>(params.rank + 1),
+                                      options);
+
+    LOG(INFO) << "Rank " << params.rank << ": Input tensor created, shape: ["
+              << params.input_size << ", " << params.hidden_dim << "]";
+
+    // Perform reduce_scatter
+    torch::Tensor output = reduce_scatter(input, process_group.get());
+
+    // Synchronize device
+    xllm_device.synchronize_default_stream();
+
+    LOG(INFO) << "Rank " << params.rank
+              << ": reduce_scatter completed, output shape: [" << output.size(0)
+              << ", " << output.size(1) << "]";
+
+    // Verify output
+    // After reduce_scatter, each rank should receive a chunk of the reduced
+    // tensor
+    // The reduced value should be the sum across all ranks: (1 + 2 + ... +
+    // world_size) = world_size * (world_size + 1) / 2
+    float expected_value =
+        static_cast<float>(params.world_size * (params.world_size + 1) / 2);
+
+    // Calculate expected chunk size
+    int64_t padded_size = params.input_size;
+    if (params.test_padding && params.input_size % params.world_size != 0) {
+      int64_t remainder = params.input_size % params.world_size;
+      padded_size = params.input_size + params.world_size - remainder;
+    }
+    int64_t chunk_size = padded_size / params.world_size;
+    int64_t expected_output_size = chunk_size;
+    if (params.test_padding && params.input_size % params.world_size != 0) {
+      int64_t global_start = params.rank * chunk_size;
+      int64_t global_end = global_start + chunk_size;
+      if (global_start >= params.input_size) {
+        expected_output_size = 0;
+      } else if (global_end > params.input_size) {
+        expected_output_size = params.input_size - global_start;
+      }
+    }
+
+    // Move output to CPU for verification
+    torch::Tensor output_cpu = output.to(torch::kCPU);
+
+    // Verify output size
+    if (expected_output_size == 0) {
+      if (output_cpu.size(0) != 0) {
+        LOG(ERROR) << "Rank " << params.rank
+                   << ": Expected empty output, but got size: "
+                   << output_cpu.size(0);
+        return 1;
+      }
+      LOG(INFO) << "Rank " << params.rank
+                << ": Output is empty as expected (padding case)";
+      return 0;
+    }
+
+    if (output_cpu.size(0) != expected_output_size) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": Output size mismatch. Expected: " << expected_output_size
+                 << ", got: " << output_cpu.size(0);
+      return 1;
+    }
+
+    if (output_cpu.size(1) != params.hidden_dim) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": Output hidden_dim mismatch. Expected: "
+                 << params.hidden_dim << ", got: " << output_cpu.size(1);
+      return 1;
+    }
+
+    // Verify output values
+    auto output_values = output_cpu.accessor<float, 2>();
+    bool all_match = true;
+    for (int64_t i = 0; i < expected_output_size; ++i) {
+      for (int64_t j = 0; j < params.hidden_dim; ++j) {
+        if (std::abs(output_values[i][j] - expected_value) > 1e-5) {
+          all_match = false;
+          LOG(ERROR) << "Rank " << params.rank << ": Mismatch at [" << i << "]["
+                     << j << "]: expected=" << expected_value
+                     << ", got=" << output_values[i][j];
+        }
+      }
+    }
+
+    if (!all_match) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": reduce_scatter test failed - output values don't match";
+      return 1;
+    }
+
+    LOG(INFO) << "Rank " << params.rank << ": reduce_scatter test passed";
+    return 0;
+
+  } catch (const std::exception& e) {
+    LOG(ERROR) << "Rank " << params.rank << ": Exception: " << e.what();
+    return 1;
+  }
+}
+
+int run_allgather_base_test_child(const AllGatherBaseTestParams& params) {
+  try {
+    xllm::Device xllm_device(params.device_index);
+    xllm_device.set_device();
+    torch::Device device = xllm_device.unwrap();
+
+    auto process_group = create_test_process_group(
+        params.rank, params.world_size, params.port, params.host, device);
+    if (!process_group) {
+      LOG(ERROR) << "Rank " << params.rank << ": Failed to create ProcessGroup";
+      return 1;
+    }
+
+    auto options = torch::TensorOptions()
+                       .dtype(torch::kFloat32)
+                       .device(device)
+                       .requires_grad(false);
+    auto cpu_options = torch::TensorOptions().dtype(torch::kFloat32);
+    torch::Tensor input =
+        torch::arange(params.input_size * params.hidden_dim, cpu_options)
+            .reshape({params.input_size, params.hidden_dim})
+            .to(device) +
+        static_cast<float>(params.rank * 1000);
+
+    torch::Tensor stacked = process_group->allgather_base_sync(input);
+    torch::Tensor async_stacked = torch::empty_like(stacked);
+    process_group->allgather_base_async(input, async_stacked)->wait();
+    xllm_device.synchronize_default_stream();
+    if (stacked.dim() != 3 || stacked.size(0) != params.world_size ||
+        stacked.size(1) != params.input_size ||
+        stacked.size(2) != params.hidden_dim) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": allgather_base output shape mismatch";
+      return 1;
+    }
+
+    torch::Tensor stacked_cpu = stacked.to(torch::kCPU);
+    for (int32_t src_rank = 0; src_rank < params.world_size; ++src_rank) {
+      torch::Tensor expected =
+          torch::arange(params.input_size * params.hidden_dim, cpu_options)
+              .reshape({params.input_size, params.hidden_dim}) +
+          static_cast<float>(src_rank * 1000);
+      if (!torch::equal(stacked_cpu[src_rank], expected)) {
+        LOG(ERROR) << "Rank " << params.rank
+                   << ": allgather_base slice mismatch for src_rank="
+                   << src_rank;
+        return 1;
+      }
+    }
+    if (!torch::equal(async_stacked.to(torch::kCPU), stacked_cpu)) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": allgather_base_async result mismatch";
+      return 1;
+    }
+
+    torch::Tensor gathered = gather(input, process_group.get(), 0);
+    xllm_device.synchronize_default_stream();
+    std::vector<torch::Tensor> gathered_parts;
+    gathered_parts.reserve(params.world_size);
+    for (int32_t src_rank = 0; src_rank < params.world_size; ++src_rank) {
+      gathered_parts.push_back(stacked_cpu[src_rank]);
+    }
+    torch::Tensor expected_gather = torch::cat(gathered_parts, 0);
+    if (!torch::equal(gathered.to(torch::kCPU), expected_gather)) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": gather result mismatch after allgather_base";
+      return 1;
+    }
+
+    auto gather_ctx = launch_gather(
+        input,
+        process_group.get(),
+        std::vector<int32_t>(params.world_size, params.input_size));
+    torch::Tensor gathered_async = finish_gather(std::move(gather_ctx));
+    xllm_device.synchronize_default_stream();
+    if (!torch::equal(gathered_async.to(torch::kCPU), expected_gather)) {
+      LOG(ERROR) << "Rank " << params.rank
+                 << ": async gather result mismatch for equal token_num_list";
+      return 1;
+    }
+
+    return 0;
+  } catch (const std::exception& e) {
+    LOG(ERROR) << "Rank " << params.rank << ": Exception: " << e.what();
+    return 1;
+  }
+}
+// Multi-process test fixture
+class ReduceScatterMultiDeviceTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Initialize test parameters
+    world_size_ = 2;
+    port_ = 29501;
+    host_ = "127.0.0.1";
+    input_size_ = 8;
+    hidden_dim_ = 128;
+  }
+
+  void TearDown() override {
+    // Clean up if needed
+  }
+
+  // Run test with multiple processes
+  void RunMultiProcessTest(int64_t input_size, bool test_padding) {
+    std::vector<pid_t> child_pids;
+    std::vector<int> child_statuses(world_size_);
+
+    // Fork child processes
+    for (int32_t rank = 0; rank < world_size_; ++rank) {
+      pid_t pid = fork();
+      if (pid == 0) {
+        // Child process
+        TestParams params;
+        params.rank = rank;
+        params.world_size = world_size_;
+        params.port = port_;
+        params.host = host_;
+        params.device_index = rank % Device::device_count();
+        params.input_size = input_size;
+        params.hidden_dim = hidden_dim_;
+        params.test_padding = test_padding;
+
+        int exit_code = run_reduce_scatter_test_child(params);
+        _exit(exit_code);
+      } else if (pid > 0) {
+        // Parent process
+        child_pids.push_back(pid);
+      } else {
+        // Fork failed
+        LOG(FATAL) << "Failed to fork child process for rank " << rank;
+      }
+    }
+
+    // Wait for all child processes to complete
+    bool all_passed = true;
+    for (size_t i = 0; i < child_pids.size(); ++i) {
+      int status;
+      pid_t waited_pid = waitpid(child_pids[i], &status, 0);
+      if (waited_pid == child_pids[i]) {
+        if (WIFEXITED(status)) {
+          child_statuses[i] = WEXITSTATUS(status);
+          if (child_statuses[i] != 0) {
+            all_passed = false;
+            LOG(ERROR) << "Child process for rank " << i << " exited with code "
+                       << child_statuses[i];
+          }
+        } else {
+          all_passed = false;
+          LOG(ERROR) << "Child process for rank " << i
+                     << " did not exit normally";
+        }
+      } else {
+        all_passed = false;
+        LOG(ERROR) << "Failed to wait for child process for rank " << i;
+      }
+    }
+
+    CHECK(all_passed) << "One or more child processes failed";
+  }
+
+  int32_t world_size_;
+  int32_t port_;
+  std::string host_;
+  int64_t input_size_;
+  int64_t hidden_dim_;
+};
+
+class AllGatherBaseMultiDeviceTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    world_size_ = 2;
+    port_ = 29531;
+    host_ = "127.0.0.1";
+    input_size_ = 4;
+    hidden_dim_ = 6;
+  }
+
+  void RunMultiProcessTest(int64_t input_size, int64_t hidden_dim) {
+    std::vector<pid_t> child_pids;
+    std::vector<int> child_statuses(world_size_);
+
+    for (int32_t rank = 0; rank < world_size_; ++rank) {
+      pid_t pid = fork();
+      if (pid == 0) {
+        AllGatherBaseTestParams params;
+        params.rank = rank;
+        params.world_size = world_size_;
+        params.port = port_;
+        params.host = host_;
+        params.device_index = rank % Device::device_count();
+        params.input_size = input_size;
+        params.hidden_dim = hidden_dim;
+        _exit(run_allgather_base_test_child(params));
+      } else if (pid > 0) {
+        child_pids.push_back(pid);
+      } else {
+        LOG(FATAL) << "Failed to fork child process for rank " << rank;
+      }
+    }
+
+    bool all_passed = true;
+    for (size_t i = 0; i < child_pids.size(); ++i) {
+      int status;
+      pid_t waited_pid = waitpid(child_pids[i], &status, 0);
+      if (waited_pid == child_pids[i]) {
+        if (WIFEXITED(status)) {
+          child_statuses[i] = WEXITSTATUS(status);
+          if (child_statuses[i] != 0) {
+            all_passed = false;
+            LOG(ERROR) << "Child process for rank " << i << " exited with code "
+                       << child_statuses[i];
+          }
+        } else {
+          all_passed = false;
+          LOG(ERROR) << "Child process for rank " << i
+                     << " did not exit normally";
+        }
+      } else {
+        all_passed = false;
+        LOG(ERROR) << "Failed to wait for child process for rank " << i;
+      }
+    }
+
+    CHECK(all_passed) << "One or more child processes failed";
+  }
+
+  int32_t world_size_ = 0;
+  int32_t port_ = 0;
+  std::string host_;
+  int64_t input_size_ = 0;
+  int64_t hidden_dim_ = 0;
+};
+TEST_F(ReduceScatterMultiDeviceTest, BasicTest) {
+  // Test with input size divisible by world_size
+  RunMultiProcessTest(8, false);
+}
+
+TEST_F(ReduceScatterMultiDeviceTest, PaddingTest) {
+  // Test with input size not divisible by world_size (requires padding)
+  RunMultiProcessTest(7, true);
+}
+
+TEST_F(ReduceScatterMultiDeviceTest, LargeInputTest) {
+  // Test with larger input
+  RunMultiProcessTest(32, false);
+}
+
+TEST_F(AllGatherBaseMultiDeviceTest, BasicTest) {
+  RunMultiProcessTest(input_size_, hidden_dim_);
+}
+}  // namespace test
+}  // namespace parallel_state
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/prefix_cache/CMakeLists.txt

@@ -0,0 +1,18 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    prefix_test
+  SRCS
+    prefix_cache_test.cpp
+  DEPS
+    :flags
+    :kv_cache
+    :prefix_cache
+    :block
+    absl::random_random
+    Boost::serialization
+    GTest::gtest_main
+)
+target_link_libraries(prefix_test PRIVATE brpc OpenSSL::SSL OpenSSL::Crypto Folly::folly :xllm_server)
+add_dependencies(prefix_test brpc-static)

upstream_ref/xllm/tests/core/framework/prefix_cache/prefix_cache_test.cpp

@@ -0,0 +1,384 @@
+#include "prefix_cache.h"
+
+#include <absl/random/random.h>
+#include <gtest/gtest.h>
+#include <string.h>
+
+#include <iostream>
+
+#include "framework/block/block_manager_impl.h"
+
+namespace xllm {
+
+void test_basic_operation(BlockManagerImpl* block_manager,
+                          PrefixCache* prefix_cache,
+                          uint32_t block_size) {
+  EXPECT_EQ(prefix_cache->num_blocks(), 0);
+
+  // token_ids number must be greater than  2 * block_size here
+  std::vector<int32_t> token_ids = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  Slice<int32_t> slice_token_ids(token_ids);
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  uint32_t n_blocks = token_ids.size() / block_size;
+  {
+    std::vector<Block> token_blocks = block_manager->allocate(n_blocks);
+
+    prefix_cache->insert(slice_token_ids, token_blocks);
+  }
+
+  EXPECT_EQ(prefix_cache->num_blocks(), n_blocks);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  {
+    auto block_matched =
+        prefix_cache->match(slice_token_ids.slice(block_size, 2 * block_size));
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  EXPECT_EQ(prefix_cache->evict(1), 1);
+
+  EXPECT_EQ(prefix_cache->num_blocks(), n_blocks - 1);
+
+  {
+    auto block_matched =
+        prefix_cache->match(slice_token_ids.slice(0, block_size));
+    EXPECT_EQ(block_matched.size(), 1);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids.slice(block_size));
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+}
+
+TEST(PrefixCacheTest, BasicOperation) {
+  const uint32_t block_size = 4;
+  const uint32_t total_blocks = 5;
+  BlockManager::Options options;
+  options.num_blocks(total_blocks).block_size(block_size);
+  BlockManagerImpl block_manager(options);
+
+  PrefixCache prefix_cache(block_size);
+
+  test_basic_operation(&block_manager, &prefix_cache, block_size);
+}
+
+void test_insert_operation(BlockManagerImpl* block_manager,
+                           PrefixCache* prefix_cache,
+                           uint32_t block_size) {
+  EXPECT_EQ(prefix_cache->num_blocks(), 0);
+
+  // insert two-block firstly
+  // token_ids number must be greater than  2 * block_size here
+  std::vector<int32_t> token_ids = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  Slice<int32_t> slice_token_ids(token_ids);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  uint32_t n_blocks = token_ids.size() / block_size;
+
+  {
+    std::vector<Block> token_blocks = block_manager->allocate(n_blocks);
+
+    prefix_cache->insert(slice_token_ids, token_blocks);
+
+    EXPECT_EQ(prefix_cache->num_blocks(), n_blocks);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  // insert another two-block
+  std::vector<int32_t> token_ids_1 = {9, 10, 11, 12, 13, 14, 15, 16, 17};
+  Slice<int32_t> slice_token_ids_1(token_ids_1);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids_1);
+
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  n_blocks = token_ids_1.size() / block_size;
+
+  {
+    std::vector<Block> token_blocks_1 = block_manager->allocate(n_blocks);
+
+    prefix_cache->insert(slice_token_ids_1, token_blocks_1);
+
+    EXPECT_EQ(prefix_cache->num_blocks(), 2 * n_blocks);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids_1);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  EXPECT_EQ(prefix_cache->evict(1), 1);
+  EXPECT_EQ(prefix_cache->num_blocks(), 2 * n_blocks - 1);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids_1);
+    EXPECT_EQ(block_matched.size(), n_blocks - 1);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  EXPECT_EQ(prefix_cache->evict(1), 1);
+  EXPECT_EQ(prefix_cache->num_blocks(), 2 * n_blocks - 2);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids_1);
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+
+    prefix_cache->insert(slice_token_ids, block_matched);
+  }
+
+  EXPECT_EQ(prefix_cache->num_blocks(), n_blocks);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  prefix_cache->evict(1);
+  EXPECT_EQ(prefix_cache->num_blocks(), n_blocks - 1);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids_1);
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks - 1);
+  }
+}
+
+TEST(PrefixCacheTest, InsertOperation) {
+  const uint32_t block_size = 4;
+  const uint32_t total_blocks = 5;
+  BlockManager::Options options;
+  options.num_blocks(total_blocks).block_size(block_size);
+  BlockManagerImpl block_manager(options);
+
+  PrefixCache prefix_cache(block_size);
+
+  test_insert_operation(&block_manager, &prefix_cache, block_size);
+}
+
+void test_evict_operation(BlockManagerImpl* block_manager,
+                          PrefixCache* prefix_cache,
+                          uint32_t block_size) {
+  EXPECT_EQ(prefix_cache->num_blocks(), 0);
+
+  prefix_cache->evict(1);
+  EXPECT_EQ(prefix_cache->num_blocks(), 0);
+
+  // insert two-block firstly
+  // token_ids number must be greater than  2 * block_size here
+  std::vector<int32_t> token_ids = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  Slice<int32_t> slice_token_ids(token_ids);
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  uint32_t n_blocks = token_ids.size() / block_size;
+
+  {
+    std::vector<Block> token_blocks = block_manager->allocate(n_blocks);
+
+    prefix_cache->insert(slice_token_ids, token_blocks);
+
+    EXPECT_EQ(prefix_cache->num_blocks(), n_blocks);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+
+  EXPECT_EQ(block_manager->num_free_blocks(),
+            block_manager->num_total_blocks() - n_blocks);
+
+  EXPECT_EQ(prefix_cache->evict(n_blocks), n_blocks);
+
+  EXPECT_EQ(block_manager->num_free_blocks(),
+            block_manager->num_total_blocks());
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), 0);
+  }
+
+  {
+    std::vector<Block> token_blocks = block_manager->allocate(n_blocks);
+
+    prefix_cache->insert(slice_token_ids, token_blocks);
+
+    EXPECT_EQ(prefix_cache->num_blocks(), n_blocks);
+  }
+
+  {
+    auto block_matched = prefix_cache->match(slice_token_ids);
+    EXPECT_EQ(block_matched.size(), n_blocks);
+  }
+}
+
+TEST(PrefixCacheTest, EvictOperation) {
+  const uint32_t block_size = 4;
+  const uint32_t total_blocks = 5;
+  BlockManager::Options options;
+  options.num_blocks(total_blocks).block_size(block_size);
+  BlockManagerImpl block_manager(options);
+
+  PrefixCache prefix_cache(block_size);
+
+  test_evict_operation(&block_manager, &prefix_cache, block_size);
+}
+
+TEST(HashUtilTest, XXHash3) {
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_EQ(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {1, 2, 3, 5, 4};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_NE(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {2, 1, 3, 5, 4};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_NE(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {2, 1, 3, 5, 4};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_NE(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {1, 2, 3, 4};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_NE(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {1, 2};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_NE(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+
+  {
+    std::vector<int32_t> tokens_1 = {1, 2, 3, 4, 5};
+    uint8_t hash_value_1[XXH3_128BITS_HASH_VALUE_LEN];
+
+    std::vector<int32_t> tokens_2 = {1, 2, 3, 4, 5, 1, 2, 3, 4, 5};
+    uint8_t hash_value_2[XXH3_128BITS_HASH_VALUE_LEN];
+
+    xxh3_128bits_hash(nullptr, tokens_1, hash_value_1);
+    xxh3_128bits_hash(nullptr, tokens_2, hash_value_2);
+
+    EXPECT_NE(strncmp(reinterpret_cast<const char*>(hash_value_1),
+                      reinterpret_cast<const char*>(hash_value_2),
+                      XXH3_128BITS_HASH_VALUE_LEN),
+              0);
+  }
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/request/CMakeLists.txt

@@ -0,0 +1,12 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    sample_slot_test
+  SRCS
+    sample_slot_test.cpp
+  DEPS
+    :request
+    GTest::gtest
+    GTest::gtest_main
+)

upstream_ref/xllm/tests/core/framework/request/sample_slot_test.cpp

@@ -0,0 +1,425 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "sample_slot.h"
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "common/global_flags.h"
+#include "framework/block/block_manager_impl.h"
+#include "platform/device.h"
+#include "request.h"
+#include "request_state.h"
+
+namespace xllm {
+namespace {
+
+class CharTokenizer final : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens = true) const override {
+    if (ids == nullptr) {
+      return false;
+    }
+    ids->clear();
+    if (add_special_tokens) {
+      ids->push_back(kBosTokenId);
+    }
+    size_t pos = 0;
+    while (pos < text.size()) {
+      if (text.size() - pos >= kEmbTokenLen &&
+          std::memcmp(text.data() + pos, kEmbToken, kEmbTokenLen) == 0) {
+        ids->push_back(kEmbTokenId);
+        pos += kEmbTokenLen;
+        continue;
+      }
+      ids->push_back(static_cast<unsigned char>(text[pos]));
+      ++pos;
+    }
+    return true;
+  }
+
+  std::string decode(const Slice<int32_t>& ids,
+                     bool skip_special_tokens) const override {
+    std::string text;
+    for (const auto token_id : ids) {
+      if (skip_special_tokens && token_id == kBosTokenId) {
+        continue;
+      }
+      if (token_id == kEmbTokenId) {
+        text.append(kEmbToken, kEmbTokenLen);
+        continue;
+      }
+      text.push_back(static_cast<char>(token_id));
+    }
+    return text;
+  }
+
+  std::optional<int32_t> token_to_id(
+      const std::string_view& token) const override {
+    if (token == std::string_view(kEmbToken, kEmbTokenLen)) {
+      return kEmbTokenId;
+    }
+    return std::nullopt;
+  }
+
+  std::string id_to_token(int32_t id) const override {
+    if (id == kBosTokenId) {
+      return "<bos>";
+    }
+    if (id == kEmbTokenId) {
+      return std::string(kEmbToken, kEmbTokenLen);
+    }
+    return std::string(1, static_cast<char>(id));
+  }
+
+  std::unique_ptr<Tokenizer> clone() const override {
+    return std::make_unique<CharTokenizer>();
+  }
+
+ private:
+  static constexpr int32_t kBosTokenId = 1;
+  static constexpr int32_t kEmbTokenId = 100000;
+  static constexpr char kEmbToken[] = "<emb_0>";
+  static constexpr size_t kEmbTokenLen = sizeof(kEmbToken) - 1;
+};
+
+class ScopedBoolFlag final {
+ public:
+  ScopedBoolFlag(bool* flag, bool value) : flag_(flag), old_value_(*flag) {
+    *flag_ = value;
+  }
+
+  ~ScopedBoolFlag() { *flag_ = old_value_; }
+
+ private:
+  bool* flag_;
+  bool old_value_;
+};
+
+TEST(SampleSlotTest, BuildSampleSlotsKeepsMatchOrderAndSampleIds) {
+  CharTokenizer tokenizer;
+  std::vector<SampleSlot> sample_slots;
+
+  ASSERT_TRUE(build_sample_slots(
+      "sample-req", "A<emb_0>B<emb_0>C", "<emb_0>", tokenizer, &sample_slots));
+
+  ASSERT_EQ(sample_slots.size(), 2);
+
+  EXPECT_EQ(sample_slots[0].request_id, "sample-req");
+  EXPECT_EQ(sample_slots[0].sample_id, 0);
+  EXPECT_EQ(sample_slots[0].token_position, 1);
+
+  EXPECT_EQ(sample_slots[1].sample_id, 1);
+  EXPECT_EQ(sample_slots[1].token_position, 3);
+}
+
+class UnstableLiteralTokenizer final : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens = true) const override {
+    if (ids == nullptr) {
+      return false;
+    }
+    ids->clear();
+    if (add_special_tokens) {
+      ids->push_back(1);
+    }
+    for (char ch : text) {
+      ids->push_back(static_cast<unsigned char>(ch));
+    }
+    return true;
+  }
+};
+
+TEST(SampleSlotTest, BuildSampleSlotsRejectsUnstableLiteralToken) {
+  UnstableLiteralTokenizer tokenizer;
+  std::vector<SampleSlot> sample_slots;
+
+  EXPECT_FALSE(build_sample_slots(
+      "sample-req", "A<emb_0>B<emb_0>C", "<emb_0>", tokenizer, &sample_slots));
+  EXPECT_TRUE(sample_slots.empty());
+}
+
+TEST(SampleSlotTest, RequestPropagatesSampleSlotsToSequenceRuntime) {
+  RequestSamplingParam sampling_param;
+  StoppingChecker stopping_checker;
+  RequestState request_state(
+      "abc",
+      std::vector<int32_t>{10, 11, 12},
+      sampling_param,
+      SchedulerParam{},
+      stopping_checker,
+      /*seq_capacity=*/8,
+      /*n=*/1,
+      /*best_of=*/1,
+      /*logprobs=*/false,
+      /*stream=*/false,
+      /*echo=*/false,
+      /*skip_special_tokens=*/true,
+      /*enable_schedule_overlap=*/false,
+      [](const RequestOutput&) { return true; },
+      OutputsFunc{});
+
+  SampleSlot first_slot;
+  first_slot.request_id = "sample-req";
+  first_slot.sample_id = 0;
+  first_slot.token_position = 2;
+
+  SampleSlot second_slot;
+  second_slot.request_id = "sample-req";
+  second_slot.sample_id = 1;
+  second_slot.token_position = 4;
+
+  request_state.sample_slots = {first_slot, second_slot};
+
+  Request request("sample-req", "", "", request_state);
+
+  ASSERT_EQ(request.sequences().size(), 1);
+  const auto& runtime_sample_slots = request.sequences()[0]->sample_slots();
+  ASSERT_EQ(runtime_sample_slots.size(), 2);
+  EXPECT_EQ(runtime_sample_slots[0].sample_id, 0);
+  EXPECT_EQ(runtime_sample_slots[0].token_position, 2);
+  EXPECT_EQ(runtime_sample_slots[1].sample_id, 1);
+}
+
+TEST(SampleSlotTest, RequestOutputSplitsSampleResultsBySampleId) {
+  torch::Device device(Device::type_torch(), 0);
+  BlockManager::Options options;
+  options.num_blocks(4).block_size(4);
+  BlockManagerImpl manager(options);
+
+  CharTokenizer tokenizer;
+  RequestSamplingParam sampling_param;
+  sampling_param.logprobs = true;
+  sampling_param.top_logprobs = 2;
+
+  StoppingChecker stopping_checker;
+  stopping_checker.set_max_generated_tokens(1);
+
+  RequestState request_state(
+      "abc",
+      std::vector<int32_t>{1, 'a', 'b', 'c'},
+      sampling_param,
+      SchedulerParam{},
+      stopping_checker,
+      /*seq_capacity=*/8,
+      /*n=*/1,
+      /*best_of=*/1,
+      /*logprobs=*/true,
+      /*stream=*/false,
+      /*echo=*/false,
+      /*skip_special_tokens=*/true,
+      /*enable_schedule_overlap=*/false,
+      [](const RequestOutput&) { return true; },
+      OutputsFunc{});
+
+  SampleSlot first_slot;
+  first_slot.request_id = "sample-req";
+  first_slot.sample_id = 0;
+  first_slot.token_position = 1;
+
+  SampleSlot second_slot = first_slot;
+  second_slot.sample_id = 1;
+  second_slot.token_position = 2;
+
+  request_state.sample_slots = {first_slot, second_slot};
+
+  Request request("sample-req", "", "", request_state);
+  auto* seq = request.sequences()[0].get();
+  seq->add_kv_blocks(manager.allocate(1));
+  seq->kv_state().set_kv_cache_tokens_num(seq->num_prompt_tokens());
+
+  std::vector<int64_t> top_tokens = {'X', 'Y'};
+  std::vector<float> top_logprobs = {-0.10f, -1.20f};
+  Token first_token('X');
+  first_token.logprob = -0.10f;
+  first_token.top_tokens = top_tokens;
+  first_token.top_logprobs = top_logprobs;
+  seq->append_token(first_token);
+
+  Token missing_logprob_token('Z');
+  seq->append_token(missing_logprob_token);
+
+  RequestOutput output = request.generate_output(tokenizer);
+
+  ASSERT_TRUE(output.status.has_value());
+  EXPECT_TRUE(output.status->ok());
+  ASSERT_TRUE(output.usage.has_value());
+  EXPECT_EQ(output.usage->num_generated_tokens, 2);
+  ASSERT_EQ(output.outputs.size(), 2);
+
+  EXPECT_EQ(output.outputs[0].index, 0U);
+  EXPECT_EQ(output.outputs[0].text, "X");
+  ASSERT_TRUE(output.outputs[0].logprobs.has_value());
+  ASSERT_EQ(output.outputs[0].logprobs->size(), 1);
+  EXPECT_EQ(output.outputs[0].logprobs->front().token, "X");
+  ASSERT_TRUE(output.outputs[0].logprobs->front().top_logprobs.has_value());
+  ASSERT_EQ(output.outputs[0].logprobs->front().top_logprobs->size(), 2);
+  EXPECT_EQ(output.outputs[0].logprobs->front().top_logprobs->at(0).token, "X");
+
+  EXPECT_EQ(output.outputs[1].index, 1U);
+  EXPECT_TRUE(output.outputs[1].text.empty());
+  EXPECT_FALSE(output.outputs[1].logprobs.has_value());
+  ASSERT_TRUE(output.outputs[1].finish_reason.has_value());
+  EXPECT_EQ(output.outputs[1].finish_reason.value(), "empty_logprobs");
+}
+
+TEST(SampleSlotTest, RequestOutputStableSortsOutOfOrderSampleIds) {
+  torch::Device device(Device::type_torch(), 0);
+  BlockManager::Options options;
+  options.num_blocks(4).block_size(4);
+  BlockManagerImpl manager(options);
+
+  CharTokenizer tokenizer;
+  RequestSamplingParam sampling_param;
+  sampling_param.logprobs = true;
+
+  StoppingChecker stopping_checker;
+  stopping_checker.set_max_generated_tokens(1);
+
+  RequestState request_state(
+      "abc",
+      std::vector<int32_t>{1, 'a', 'b', 'c'},
+      sampling_param,
+      SchedulerParam{},
+      stopping_checker,
+      /*seq_capacity=*/8,
+      /*n=*/1,
+      /*best_of=*/1,
+      /*logprobs=*/true,
+      /*stream=*/false,
+      /*echo=*/false,
+      /*skip_special_tokens=*/true,
+      /*enable_schedule_overlap=*/false,
+      [](const RequestOutput&) { return true; },
+      OutputsFunc{});
+
+  SampleSlot slot2;
+  slot2.request_id = "sample-req";
+  slot2.sample_id = 2;
+  slot2.token_position = 1;
+
+  SampleSlot slot0 = slot2;
+  slot0.sample_id = 0;
+  slot0.token_position = 2;
+
+  SampleSlot slot1 = slot2;
+  slot1.sample_id = 1;
+  slot1.token_position = 3;
+
+  request_state.sample_slots = {slot2, slot0, slot1};
+
+  Request request("sample-req", "", "", request_state);
+  auto* seq = request.sequences()[0].get();
+  seq->add_kv_blocks(manager.allocate(1));
+  seq->kv_state().set_kv_cache_tokens_num(seq->num_prompt_tokens());
+
+  Token slot2_token('C');
+  slot2_token.logprob = -0.30f;
+  seq->append_token(slot2_token);
+
+  Token slot0_token('A');
+  slot0_token.logprob = -0.10f;
+  seq->append_token(slot0_token);
+
+  Token slot1_token('B');
+  slot1_token.logprob = -0.20f;
+  seq->append_token(slot1_token);
+
+  RequestOutput output = request.generate_output(tokenizer);
+
+  ASSERT_EQ(output.outputs.size(), 3);
+  EXPECT_EQ(output.outputs[0].index, 0U);
+  EXPECT_EQ(output.outputs[0].text, "A");
+  EXPECT_EQ(output.outputs[1].index, 1U);
+  EXPECT_EQ(output.outputs[1].text, "B");
+  EXPECT_EQ(output.outputs[2].index, 2U);
+  EXPECT_EQ(output.outputs[2].text, "C");
+}
+
+TEST(SampleSlotTest, OneRecOutputCarriesTokenLogprobsWhenEnabled) {
+  ScopedBoolFlag enable_output_sku_logprobs(&FLAGS_enable_output_sku_logprobs,
+                                            true);
+  ScopedBoolFlag enable_convert_tokens_to_item(
+      &FLAGS_enable_convert_tokens_to_item, false);
+
+  CharTokenizer tokenizer;
+  RequestSamplingParam sampling_param;
+  sampling_param.logprobs = true;
+
+  StoppingChecker stopping_checker;
+  stopping_checker.set_max_generated_tokens(3);
+
+  RequestState request_state(
+      /*prompt=*/"",
+      std::vector<int32_t>{11, 12},
+      sampling_param,
+      SchedulerParam{},
+      stopping_checker,
+      /*seq_capacity=*/8,
+      /*n=*/1,
+      /*best_of=*/1,
+      /*logprobs=*/true,
+      /*stream=*/false,
+      /*echo=*/false,
+      /*skip_special_tokens=*/true,
+      /*enable_schedule_overlap=*/false,
+      [](const RequestOutput&) { return true; },
+      OutputsFunc{});
+  request_state.rec_type = RecType::kOneRec;
+
+  Request request("onerec-score",
+                  /*x_request_id=*/"",
+                  /*x_request_time=*/"",
+                  request_state);
+  auto* seq = request.sequences()[0].get();
+
+  Token first_token(101);
+  first_token.logprob = -0.10f;
+  seq->append_token(first_token);
+
+  Token second_token(102);
+  second_token.logprob = -0.20f;
+  seq->append_token(second_token);
+
+  Token third_token(103);
+  third_token.logprob = -0.30f;
+  seq->append_token(third_token);
+
+  RequestOutput output = request.generate_output(tokenizer);
+
+  ASSERT_EQ(output.outputs.size(), 1);
+  ASSERT_EQ(output.outputs[0].token_ids.size(), 3U);
+  ASSERT_EQ(output.outputs[0].token_ids_logprobs.size(), 3U);
+  ASSERT_TRUE(output.outputs[0].token_ids_logprobs[0].has_value());
+  ASSERT_TRUE(output.outputs[0].token_ids_logprobs[1].has_value());
+  ASSERT_TRUE(output.outputs[0].token_ids_logprobs[2].has_value());
+  EXPECT_FLOAT_EQ(output.outputs[0].token_ids_logprobs[0].value(), -0.10f);
+  EXPECT_FLOAT_EQ(output.outputs[0].token_ids_logprobs[1].value(), -0.20f);
+  EXPECT_FLOAT_EQ(output.outputs[0].token_ids_logprobs[2].value(), -0.30f);
+}
+
+}  // namespace
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/sampling/CMakeLists.txt

@@ -0,0 +1,26 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    sampler_test
+  SRCS
+    rejection_sampler_test.cpp
+    "${PROJECT_SOURCE_DIR}/xllm/core/framework/sampling/rejection_sampler.cpp"
+    sampling_params_test.cpp
+  DEPS
+    absl::strings
+    GTest::gtest_main
+    :flags
+    :sampler
+    glog::glog
+    :state_dict
+)
+target_link_libraries(sampler_test PRIVATE brpc OpenSSL::SSL OpenSSL::Crypto leveldb::leveldb protobuf::libprotobuf)
+target_link_libraries(sampler_test
+                      PUBLIC
+                      Python::Python
+                      $<$<BOOL:${USE_NPU}>:ascendcl>
+                      $<$<BOOL:${USE_NPU}>:hccl>
+                      $<$<BOOL:${USE_NPU}>:c_sec>
+                      $<$<BOOL:${USE_NPU}>:nnopbase>)
+add_dependencies(sampler_test brpc-static)

upstream_ref/xllm/tests/core/framework/sampling/rejection_sampler_test.cpp

@@ -0,0 +1,628 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+Copyright 2024 The ScaleLLM Authors. All Rights Reserved.
+
+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/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 "rejection_sampler.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+#include <torch/types.h>
+
+#include <cstdint>
+
+#include "platform/device.h"
+#include "sampler.h"
+
+namespace xllm {
+
+namespace {
+// Helper function to get test device: MLU if available, otherwise CPU
+torch::Device get_test_device() {
+  std::string backend = Device::type_str();
+  if ((backend == "mlu" || backend == "cuda") && Device::device_count() > 0) {
+    return torch::Device(Device::type_torch(), 0);
+  }
+  return torch::Device(torch::kCPU);
+}
+
+// Helper function to get test tensor options with automatic device selection
+torch::TensorOptions get_test_options(
+    torch::ScalarType dtype = torch::kFloat32) {
+  return torch::dtype(dtype).device(get_test_device());
+}
+}  // namespace
+
+TEST(RejectionSamplerTest, Basic) {
+  // test with hand-crafted example
+  const auto options = get_test_options(torch::kFloat32);
+  const auto device = get_test_device();
+
+  // set random seed
+  torch::manual_seed(100);
+
+  const auto draft_token_ids =
+      torch::tensor({{1, 2, 3}}, options.dtype(torch::kInt64));
+
+  // shape: [1, 3, 5]
+  auto draft_probs = torch::tensor({{0.2104, 0.2163, 0.1912, 0.1937, 0.1884},
+                                    {0.2100, 0.1803, 0.2398, 0.2088, 0.1610},
+                                    {0.1838, 0.2079, 0.2270, 0.2451, 0.1362}},
+                                   options)
+                         .reshape({1, 3, 5});
+  // shape: [1, 3, 5]
+  auto target_probs = torch::tensor({{0.1299, 0.2462, 0.1821, 0.1354, 0.3064},
+                                     {0.1159, 0.2839, 0.1603, 0.2451, 0.1949},
+                                     {0.0002, 0.0433, 0.6629, 0.1469, 0.1467}},
+                                    options)
+                          .reshape({1, 3, 5});
+
+  // selected_target_probs:  [0.2462  0.1603  0.1469]
+  // selected_draft_probs:   [0.2163  0.2398  0.2451]
+  // acceptance_probs:       [1.1382  0.6685  0.5993]
+  // uniform_rand:           [0.4785  0.6589  0.9399]
+  // accepted:               [  1        1       0  ]
+  auto uniform_rand = torch::tensor({{0.4785, 0.6589, 0.9399}}, options);
+  auto bonus_token_ids = torch::tensor({{5}}, options.dtype(torch::kInt64));
+
+  auto [output, masked_output] =
+      RejectionSampler::random_sample(draft_token_ids,
+                                      draft_probs,
+                                      target_probs,
+                                      uniform_rand,
+                                      bonus_token_ids,
+                                      true);
+  auto desired_output =
+      torch::tensor({{1, 2, 2, 5}}, options.dtype(torch::kInt64));
+  EXPECT_TRUE(torch::allclose(output, desired_output));
+
+  auto desired_masked_output =
+      torch::tensor({{1, 2, 2, -1}}, options.dtype(torch::kInt64));
+  EXPECT_TRUE(torch::allclose(masked_output, desired_masked_output));
+}
+
+TEST(RejectionSamplerTest, BasicSelectedOnlyDraftProbs) {
+  // test with hand-crafted example using selected-only draft probs
+  const auto options = get_test_options(torch::kFloat32);
+
+  // set random seed
+  torch::manual_seed(100);
+
+  const auto draft_token_ids =
+      torch::tensor({{1, 2, 3}}, options.dtype(torch::kInt64));
+
+  // selected-only draft probs shape: [1, 3]
+  auto draft_probs = torch::tensor({{0.2163, 0.2398, 0.2451}}, options);
+
+  // shape: [1, 3, 5]
+  auto target_probs = torch::tensor({{0.1299, 0.2462, 0.1821, 0.1354, 0.3064},
+                                     {0.1159, 0.2839, 0.1603, 0.2451, 0.1949},
+                                     {0.0002, 0.0433, 0.6629, 0.1469, 0.1467}},
+                                    options)
+                          .reshape({1, 3, 5});
+
+  // selected_target_probs:  [0.2462  0.1603  0.1469]
+  // selected_draft_probs:   [0.2163  0.2398  0.2451]
+  // acceptance_probs:       [1.1382  0.6685  0.5993]
+  // uniform_rand:           [0.4785  0.6589  0.9399]
+  // accepted:               [  1        1       0  ]
+  auto uniform_rand = torch::tensor({{0.4785, 0.6589, 0.9399}}, options);
+  auto bonus_token_ids = torch::tensor({{5}}, options.dtype(torch::kInt64));
+
+  auto [output, masked_output] =
+      RejectionSampler::random_sample(draft_token_ids,
+                                      draft_probs,
+                                      target_probs,
+                                      uniform_rand,
+                                      bonus_token_ids,
+                                      true);
+  auto desired_output =
+      torch::tensor({{1, 2, 2, 5}}, options.dtype(torch::kInt64));
+  EXPECT_TRUE(torch::allclose(output, desired_output));
+
+  auto desired_masked_output =
+      torch::tensor({{1, 2, 2, -1}}, options.dtype(torch::kInt64));
+  EXPECT_TRUE(torch::allclose(masked_output, desired_masked_output));
+}
+
+TEST(RejectionSamplerTest, Mask) {
+  // test accepted mask
+  const auto options = get_test_options(torch::kBool);
+
+  // clang-format off
+  auto accepted = torch::tensor({
+        {0, 1, 0, 1},
+        {1, 0, 1, 1},
+        {1, 1, 0, 1},
+        {1, 1, 1, 1}},
+        options);
+  auto desired_mask = torch::tensor({
+        {1, 0, 0, 0, 0},
+        {1, 1, 0, 0, 0},
+        {1, 1, 1, 0, 0},
+        {1, 1, 1, 1, 1}},
+        options);
+  // clang-format on
+  auto mask = RejectionSampler::build_accepted_mask(accepted);
+  EXPECT_TRUE(torch::allclose(mask, desired_mask));
+}
+
+TEST(RejectionSamplerTest, Greedy) {
+  const auto options = get_test_options(torch::kFloat32);
+  const auto device = get_test_device();
+
+  int64_t batch_size = 2;
+  int64_t n_speculative_tokens = 3;
+  int64_t vocab_size = 4;
+  int64_t n_bonus_tokens = 1;
+
+  const auto draft_token_ids =
+      torch::randint(0,
+                     vocab_size,
+                     {batch_size, n_speculative_tokens},
+                     torch::dtype(torch::kInt64).device(device));
+  auto target_probs =
+      torch::randn({batch_size, n_speculative_tokens, vocab_size}, options)
+          .softmax(/*dim=*/-1, /*dtype=*/torch::kFloat32);
+  const auto bonus_token_ids =
+      torch::randint(0,
+                     vocab_size,
+                     {batch_size, n_bonus_tokens},
+                     torch::dtype(torch::kInt64).device(device));
+
+  auto [output, masked_output] =
+      RejectionSampler::greedy_sample(draft_token_ids,
+                                      target_probs,
+                                      bonus_token_ids,
+                                      /*mask_out_rejected_tokens=*/false);
+  EXPECT_FALSE(masked_output.defined());
+
+  const auto desired_output = target_probs.argmax(/*dim=*/-1);
+
+  // check target tokens
+  EXPECT_TRUE(torch::allclose(
+      output.slice(/*dim=*/-1, /*start=*/0, /*end=*/n_speculative_tokens),
+      desired_output));
+  // check bonus tokens
+  EXPECT_TRUE(torch::allclose(output.slice(/*dim=*/-1,
+                                           /*start=*/n_speculative_tokens),
+                              bonus_token_ids));
+}
+
+TEST(RejectionSamplerTest, LogProbs) {
+  const auto options = get_test_options(torch::kFloat32);
+  const auto device = get_test_device();
+  const auto do_sample = torch::tensor({false, true, false, true}, device);
+  const int64_t max_top_logprobs = 2;
+  RejectionSampler sampler(do_sample,
+                           do_sample.all().item<bool>(),
+                           !do_sample.any().item<bool>(),
+                           /*logprobs=*/true,
+                           max_top_logprobs);
+
+  int64_t batch_size = 4;
+  int64_t n_speculative_tokens = 4;
+  int64_t vocab_size = 8;
+
+  const auto draft_token_ids =
+      torch::randint(0,
+                     vocab_size,
+                     {batch_size, n_speculative_tokens},
+                     torch::dtype(torch::kInt64).device(device));
+  auto draft_probs =
+      torch::randn({batch_size, n_speculative_tokens, vocab_size}, options)
+          .softmax(/*dim=*/-1);
+
+  auto target_logits =
+      torch::randn({batch_size, n_speculative_tokens + 1, vocab_size}, options);
+  const auto bonus_token_ids =
+      torch::randint(0,
+                     vocab_size,
+                     {batch_size, 1},
+                     torch::dtype(torch::kInt64).device(device));
+
+  auto output = sampler.forward(draft_token_ids,
+                                draft_probs,
+                                target_logits,
+                                bonus_token_ids,
+                                /*mask_out_rejected_tokens=*/false);
+
+  const auto logprobs =
+      torch::log_softmax(target_logits, /*dim=*/-1, /*dtype=*/torch::kFloat32);
+  const auto selected_tokens = output.next_tokens;
+
+  const auto selected_logprobs =
+      logprobs.gather(/*dim=*/-1, selected_tokens.unsqueeze(/*dim=*/-1))
+          .squeeze(/*dim=*/-1);
+  EXPECT_TRUE(torch::equal(output.logprobs, selected_logprobs));
+
+  auto [top_k_values, top_k_indices] = logprobs.topk(
+      max_top_logprobs, /*dim=*/-1, /*largest=*/true, /*sorted=*/true);
+  EXPECT_TRUE(torch::equal(output.top_logprobs, top_k_values));
+  EXPECT_TRUE(torch::equal(output.top_tokens, top_k_indices));
+}
+
+TEST(RejectionSamplerTest, ConstructorDoesNotMutateDoSampleShape) {
+  const auto device = get_test_device();
+  auto do_sample = torch::tensor({false, true}, torch::device(device));
+
+  ASSERT_EQ(do_sample.dim(), 1);
+  ASSERT_EQ(do_sample.sizes(), torch::IntArrayRef({2}));
+
+  RejectionSampler sampler(do_sample,
+                           do_sample.all().item<bool>(),
+                           !do_sample.any().item<bool>(),
+                           /*logprobs=*/false,
+                           /*max_top_logprobs=*/0);
+
+  EXPECT_EQ(do_sample.dim(), 1);
+  EXPECT_EQ(do_sample.sizes(), torch::IntArrayRef({2}));
+}
+
+TEST(RejectionSamplerTest,
+     ReusingDoSampleAfterRejectionSamplerKeepsSamplerOutput1D) {
+  const auto options = get_test_options(torch::kFloat32);
+  const auto device = get_test_device();
+  auto do_sample = torch::tensor({false, true}, torch::device(device));
+
+  RejectionSampler rejection_sampler(do_sample,
+                                     do_sample.all().item<bool>(),
+                                     !do_sample.any().item<bool>(),
+                                     /*logprobs=*/false,
+                                     /*max_top_logprobs=*/0);
+  (void)rejection_sampler;
+
+  SamplingParameters params;
+  params.selected_token_idxes =
+      torch::tensor({0, 1}, torch::dtype(torch::kInt64).device(device));
+  params.sample_idxes =
+      torch::tensor({0, 1}, torch::dtype(torch::kInt64).device(device));
+  params.do_sample = do_sample;
+  params.all_random_sample = false;
+  params.all_greedy_sample = false;
+
+  auto logits =
+      torch::tensor({{3.0f, 1.0f, 0.5f}, {0.1f, 0.2f, 4.0f}}, options);
+  auto output = Sampler().forward(logits, params);
+
+  EXPECT_EQ(output.probs.dim(), 2);
+  EXPECT_EQ(output.probs.size(0), 2);
+  EXPECT_EQ(output.next_tokens.dim(), 1);
+  EXPECT_EQ(output.next_tokens.size(0), 2);
+}
+
+TEST(RejectionSamplerTest, Random) {
+  const auto options = get_test_options(torch::kFloat32);
+
+  // set random seed
+  torch::manual_seed(100);
+
+  int64_t vocab_size = 50;
+  int64_t num_samples = 500000;
+
+  auto target_prob = torch::randn({vocab_size}, options).softmax(/*dim=*/-1);
+  auto target_probs =
+      target_prob.reshape({1, 1, -1}).repeat({num_samples, 1, 1});
+
+  auto draft_probs =
+      torch::randn({num_samples, 1, vocab_size}, options).softmax(/*dim=*/-1);
+  auto draft_token_ids = Sampler::random_sample(draft_probs);
+
+  // not used
+  auto bonus_token_ids =
+      torch::ones({num_samples, 1}, options.dtype(torch::kInt64));
+
+  auto uniform_rand = torch::rand(draft_token_ids.sizes(), options);
+  auto [output, masked_output] =
+      RejectionSampler::random_sample(draft_token_ids,
+                                      draft_probs,
+                                      target_probs,
+                                      uniform_rand,
+                                      bonus_token_ids,
+                                      false);
+  EXPECT_FALSE(masked_output.defined());
+
+  // remove bonus token
+  auto token_ids = output
+                       .slice(/*dim=*/-1,
+                              /*start=*/0,
+                              /*end=*/-1)
+                       .flatten();
+
+  // calculate the probability of each sampled token
+  auto bincount = token_ids.bincount(/*weights=*/torch::nullopt,
+                                     /*minlength=*/vocab_size);
+  auto sample_prob = bincount.to(torch::kFloat) / num_samples;
+
+  EXPECT_TRUE(torch::allclose(target_prob,
+                              sample_prob,
+                              /*rtol=*/1e-2,
+                              /*atol=*/1e-2));
+}
+
+TEST(RejectionSamplerTest, RandomSelectedOnlyMatchesDenseWhenAccepted) {
+  const auto options = get_test_options(torch::kFloat32);
+  const auto device = get_test_device();
+
+  // make all tokens accepted by setting target probs > selected draft probs.
+  torch::manual_seed(123);
+  int64_t batch_size = 64;
+  int64_t n_spec = 4;
+  int64_t vocab_size = 16;
+
+  auto target_probs = torch::randn({batch_size, n_spec, vocab_size}, options)
+                          .softmax(/*dim=*/-1, /*dtype=*/torch::kFloat32);
+
+  auto draft_token_ids =
+      torch::randint(0,
+                     vocab_size,
+                     {batch_size, n_spec},
+                     torch::dtype(torch::kInt64).device(device));
+
+  // Build dense draft probs with 50% of selected target probs at draft tokens.
+  auto draft_probs_dense =
+      torch::zeros({batch_size, n_spec, vocab_size}, target_probs.options());
+  auto selected_target_probs =
+      target_probs.gather(/*dim=*/-1, draft_token_ids.unsqueeze(-1))
+          .squeeze(-1);
+  auto selected_draft_probs = selected_target_probs * 0.5;
+  draft_probs_dense.scatter_(/*dim=*/-1,
+                             draft_token_ids.unsqueeze(-1),
+                             selected_draft_probs.unsqueeze(-1));
+
+  // Ensure acceptance by using uniform_rand == 0.
+  auto uniform_rand = torch::zeros({batch_size, n_spec}, options);
+  auto bonus_token_ids =
+      torch::randint(0,
+                     vocab_size,
+                     {batch_size, 1},
+                     torch::dtype(torch::kInt64).device(device));
+
+  auto [dense_output, dense_masked] =
+      RejectionSampler::random_sample(draft_token_ids,
+                                      draft_probs_dense,
+                                      target_probs,
+                                      uniform_rand,
+                                      bonus_token_ids,
+                                      true);
+
+  auto [selected_output, selected_masked] =
+      RejectionSampler::random_sample(draft_token_ids,
+                                      selected_draft_probs,
+                                      target_probs,
+                                      uniform_rand,
+                                      bonus_token_ids,
+                                      true);
+
+  EXPECT_TRUE(torch::equal(dense_output, selected_output));
+  EXPECT_TRUE(torch::equal(dense_masked, selected_masked));
+}
+
+TEST(RejectionSamplerTest, RandomFused) {
+  // Skip test if not running on MLU backend
+  std::string backend = Device::type_str();
+  if (backend != "mlu") {
+    GTEST_SKIP() << "Skipping RandomFused test: fused kernel only available on "
+                    "MLU backend.";
+  }
+  if (Device::device_count() == 0) {
+    GTEST_SKIP() << "Skipping RandomFused test: no MLU devices available";
+  }
+
+  // Prepare random test data
+  torch::ScalarType dtype(torch::kFloat32);
+  torch::Device device(Device::type_torch(), 0);
+  const auto options = torch::dtype(dtype).device(device);
+  torch::manual_seed(100);
+
+  int64_t n_spec = 3;
+  int64_t vocab_size = 50;
+  int64_t num_samples = 1000;
+
+  auto target_prob_base = torch::randn({vocab_size}, options).softmax(-1);
+  auto target_probs =
+      target_prob_base.reshape({1, 1, -1}).repeat({num_samples, n_spec, 1});
+  auto draft_probs =
+      torch::randn({num_samples, n_spec, vocab_size}, options).softmax(-1);
+
+  // Sample draft tokens and bonus tokens
+  auto draft_token_ids = Sampler::random_sample(draft_probs);
+  auto bonus_token_ids = torch::randint(
+      0, vocab_size, {num_samples, 1}, options.dtype(torch::kInt64));
+
+  // Shared random tensor, used for acceptance check
+  auto uniform_rand = torch::rand(draft_token_ids.sizes(), options);
+
+  // Fused kernel output
+  auto [fused_output_unmasked, fused_output_masked] =
+      RejectionSampler::random_sample_fused(draft_token_ids,
+                                            draft_probs,
+                                            target_probs,
+                                            uniform_rand,
+                                            bonus_token_ids,
+                                            /*mask_out_rejected_tokens=*/true);
+
+  // Reference random_sample output
+  auto [ref_output_unmasked, ref_output_masked] =
+      RejectionSampler::random_sample(draft_token_ids,
+                                      draft_probs,
+                                      target_probs,
+                                      uniform_rand,
+                                      bonus_token_ids,
+                                      /*mask_out_rejected_tokens=*/true);
+
+  // Check output shape
+  EXPECT_EQ(fused_output_masked.size(0), num_samples);
+  EXPECT_EQ(fused_output_masked.size(1), n_spec + 1);
+
+  // Mask output should match exactly (same tokens should be accepted/rejected)
+  auto fused_mask = (fused_output_masked != -1);
+  auto ref_mask = (ref_output_masked != -1);
+  EXPECT_TRUE(torch::equal(fused_mask, ref_mask))
+      << "Mismatch in acceptance decision between Fused and Ref "
+         "implementation!";
+
+  // If a draft token is accepted, it must exactly match the input and the
+  // reference output
+  for (int j = 0; j < n_spec; ++j) {
+    auto current_col = fused_output_masked.slice(1, j, j + 1);
+    auto next_col = fused_output_masked.slice(1, j + 1, j + 2);
+    auto is_accepted_draft = (current_col != -1) & (next_col != -1);
+
+    if (is_accepted_draft.any().item<bool>()) {
+      auto fused_drafts = current_col.masked_select(is_accepted_draft);
+      auto input_drafts =
+          draft_token_ids.slice(1, j, j + 1).masked_select(is_accepted_draft);
+
+      EXPECT_TRUE(torch::equal(fused_drafts, input_drafts))
+          << "Fused kernel altered an accepted draft token at index " << j;
+
+      auto ref_drafts =
+          ref_output_masked.slice(1, j, j + 1).masked_select(is_accepted_draft);
+      EXPECT_TRUE(torch::equal(fused_drafts, ref_drafts))
+          << "Mismatch with Reference on accepted draft token at index " << j;
+    }
+  }
+
+  // Bonus token column should match input when all previous tokens accepted
+  auto last_col = fused_output_masked.slice(1, n_spec, n_spec + 1);
+  auto fully_accepted_mask = (last_col != -1);
+
+  if (fully_accepted_mask.any().item<bool>()) {
+    auto valid_bonus_out = last_col.masked_select(fully_accepted_mask);
+    auto valid_bonus_in = bonus_token_ids.masked_select(fully_accepted_mask);
+    EXPECT_TRUE(torch::equal(valid_bonus_out, valid_bonus_in))
+        << "Bonus token mismatch for fully accepted sequences";
+  } else {
+    LOG(INFO)
+        << "No fully accepted sequences in this batch, skipping bonus check.";
+  }
+
+  // After the first -1 in each row, all remaining values must be -1
+  auto cpu_output = fused_output_masked.to(torch::kCPU);
+  auto output_a = cpu_output.accessor<int64_t, 2>();
+
+  for (int i = 0; i < num_samples; ++i) {
+    bool rejected = false;
+    for (int j = 0; j < n_spec + 1; ++j) {
+      if (output_a[i][j] == -1) {
+        rejected = true;
+      } else {
+        if (rejected) {
+          ADD_FAILURE() << "Found valid token after -1 at row " << i << " col "
+                        << j;
+        }
+      }
+    }
+  }
+}
+
+TEST(RejectionSamplerTest, RandomFusedRecoveryDistribution) {
+  // Check MLU device
+  if (Device::type_str() != "mlu" || Device::device_count() == 0) {
+    GTEST_SKIP() << "Skipping test: MLU device required.";
+  }
+
+  torch::manual_seed(42);
+  torch::Device device(Device::type_torch(), 0);
+
+  // Keep vocab_size small for statistical verification
+  int64_t vocab_size = 4;
+  int64_t num_samples = 5000;
+  // Use n_spec > 1 to cover broadcasting/reshaping edge cases
+  int64_t n_spec = 2;
+
+  // Draft model confidently predicts index 0 (prob=1.0) for all tokens
+  auto draft_probs = torch::zeros({num_samples, n_spec, vocab_size},
+                                  torch::dtype(torch::kFloat32).device(device));
+  draft_probs.index_put_({"...", 0}, 1.0f);
+
+  // Define target probability distribution
+  auto target_prob_single =
+      torch::zeros({vocab_size}, torch::dtype(torch::kFloat32).device(device));
+  target_prob_single[0] = 0.1f;
+  target_prob_single[1] = 0.6f;
+  target_prob_single[2] = 0.2f;
+  target_prob_single[3] = 0.1f;
+
+  // Call contiguous() to avoid stride=0 error in kernel
+  auto target_probs = target_prob_single.reshape({1, 1, -1})
+                          .repeat({num_samples, n_spec, 1})
+                          .contiguous();
+
+  // Draft Token IDs set to 0, use kInt32 type
+  auto draft_token_ids = torch::zeros(
+      {num_samples, n_spec}, torch::dtype(torch::kInt32).device(device));
+  auto bonus_token_ids = torch::zeros(
+      {num_samples, 1}, torch::dtype(torch::kInt32).device(device));
+
+  // All drafts at index 0 are rejected (uniform_rand > target_prob /
+  // draft_prob)
+  auto uniform_rand = torch::full({num_samples, n_spec},
+                                  0.5f,
+                                  torch::dtype(torch::kFloat32).device(device));
+  // Run fused kernel and get output
+  auto [output, masked_output] =
+      RejectionSampler::random_sample_fused(draft_token_ids,
+                                            draft_probs,
+                                            target_probs,
+                                            uniform_rand,
+                                            bonus_token_ids,
+                                            /*mask_out_rejected_tokens=*/true);
+
+  // All drafts at index 0 are rejected, n_spec positions have identical
+  // distributions, analyze first column
+  auto recovered_tokens = output.slice(1, 0, 1).flatten().cpu();
+  std::vector<int64_t> counts(vocab_size, 0);
+  auto accessor = recovered_tokens.accessor<int64_t, 1>();
+  for (int i = 0; i < num_samples; ++i) {
+    int64_t token = accessor[i];
+    // Ensure generated token in legal range
+    CHECK_GE(token, 0) << "Generated token index cannot be negative!";
+    CHECK_LT(token, vocab_size)
+        << "Generated token index out of bounds (>= vocab_size)!";
+    counts[token]++;
+  }
+
+  // Token 0 should never appear: probability is clamped to 0
+  EXPECT_EQ(counts[0], 0)
+      << "Token 0 should have zero probability in recovery.";
+
+  // Calculate empirical probabilities
+  double total = static_cast<double>(num_samples);
+  double p1 = counts[1] / total;
+  double p2 = counts[2] / total;
+  double p3 = counts[3] / total;
+
+  LOG(INFO) << "[RandomFusedStats] Empirical Probs: "
+            << "P(1)=" << p1 << ", P(2)=" << p2 << ", P(3)=" << p3;
+
+  // Expected probabilities (normalized)
+  // Raw: [0, 0.6, 0.2, 0.1] => Normalized: [0, 0.666..., 0.222..., 0.111...]
+  double expected_p1 = 0.6 / 0.9;  // ~0.6667
+  double expected_p2 = 0.2 / 0.9;  // ~0.2222
+  double expected_p3 = 0.1 / 0.9;  // ~0.1111
+
+  // Tolerance 0.02 (2%), reasonable for 5000 samples
+  double tolerance = 0.02;
+  EXPECT_NEAR(p1, expected_p1, tolerance)
+      << "Distribution mismatch for Token 1";
+  EXPECT_NEAR(p2, expected_p2, tolerance)
+      << "Distribution mismatch for Token 2";
+  EXPECT_NEAR(p3, expected_p3, tolerance)
+      << "Distribution mismatch for Token 3";
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/sampling/sampling_params_test.cpp

@@ -0,0 +1,118 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+Copyright 2024 The ScaleLLM Authors. All Rights Reserved.
+
+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/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 "sampling_params.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+namespace xllm {
+
+TEST(SamplingParamsTest, NormalConcat) {
+  // construct sampling_parameters_1
+  RequestSamplingParam request_1, request_2;
+  std::vector<int32_t> selected_token_idxes_1{11, 23};
+  std::vector<int32_t> sample_idxes_1{0, 1};
+  std::vector<std::vector<int64_t>> unique_token_ids_vec_1{
+      std::vector<int64_t>{
+          151645, 100022, 104202, 104167, 198, 77091, 872, 220, 151644},
+      std::vector<int64_t>{
+          151645, 100022, 104202, 104167, 198, 77091, 872, 220, 151644}};
+  std::vector<std::vector<int32_t>> unique_token_counts_vec_1{
+      std::vector<int32_t>{1, 1, 1, 1, 3, 1, 1, 1, 2},
+      std::vector<int32_t>{1, 1, 1, 1, 3, 1, 1, 1, 2}};
+  std::vector<int32_t> unique_token_lens_vec_1{9, 9};
+
+  SamplingParameters sampling_parameters_1;
+  sampling_parameters_1.init(
+      std::vector<const RequestSamplingParam*>{&request_1, &request_2},
+      selected_token_idxes_1,
+      sample_idxes_1,
+      unique_token_ids_vec_1,
+      unique_token_counts_vec_1,
+      unique_token_lens_vec_1);
+
+  // construct sampling_parameters_2
+  RequestSamplingParam request_3, request_4;
+  std::vector<int32_t> selected_token_idxes_2{13, 28};
+  std::vector<int32_t> sample_idxes_2{0, 1};
+  std::vector<std::vector<int64_t>> unique_token_ids_vec_2{
+      std::vector<int64_t>{151645,
+                           119414,
+                           100287,
+                           26288,
+                           101239,
+                           198,
+                           77091,
+                           106055,
+                           872,
+                           220,
+                           151644},
+      std::vector<int64_t>{0,
+                           62112,
+                           9370,
+                           107425,
+                           151645,
+                           99489,
+                           106309,
+                           198,
+                           77091,
+                           71618,
+                           872,
+                           220,
+                           151644}};
+  std::vector<std::vector<int32_t>> unique_token_counts_vec_2{
+      std::vector<int32_t>{1, 1, 1, 1, 1, 3, 1, 1, 1, 1, 2},
+      std::vector<int32_t>{0, 1, 1, 1, 1, 1, 1, 3, 1, 1, 1, 1, 2}};
+  std::vector<int32_t> unique_token_lens_vec_2{11, 12};
+
+  SamplingParameters sampling_parameters_2;
+  sampling_parameters_2.init(
+      std::vector<const RequestSamplingParam*>{&request_3, &request_4},
+      selected_token_idxes_2,
+      sample_idxes_2,
+      unique_token_ids_vec_2,
+      unique_token_counts_vec_2,
+      unique_token_lens_vec_2);
+
+  // construct expected output
+  torch::Tensor result_selected_token_idxes = torch::tensor({11, 23, 37, 52});
+  torch::Tensor result_sample_idxes = torch::tensor({0, 1, 2, 3});
+
+  // execute concat
+  sampling_parameters_1.concat(sampling_parameters_2);
+
+  // check results
+  EXPECT_TRUE(torch::equal(sampling_parameters_1.selected_token_idxes,
+                           result_selected_token_idxes));
+  EXPECT_TRUE(
+      torch::equal(sampling_parameters_1.sample_idxes, result_sample_idxes));
+}
+
+TEST(SamplingParamsTest, AbnormalConcat) {
+  // construct both of default sampling_parameters
+  SamplingParameters sampling_parameters_1, sampling_parameters_2;
+
+  // execute concat
+  sampling_parameters_1.concat(sampling_parameters_2);
+
+  // check results
+  EXPECT_FALSE(sampling_parameters_1.selected_token_idxes.defined());
+  EXPECT_FALSE(sampling_parameters_1.sample_idxes.defined());
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/state_dict/CMakeLists.txt

@@ -0,0 +1,13 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    rec_vocab_dict_test
+  SRCS
+    rec_vocab_dict_test.cpp
+  DEPS
+    :flags
+    :state_dict
+    GTest::gtest_main
+    glog::glog
+)

upstream_ref/xllm/tests/core/framework/state_dict/rec_vocab_dict_test.cpp

@@ -0,0 +1,171 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "framework/state_dict/rec_vocab_dict.h"
+
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <filesystem>
+#include <fstream>
+#include <utility>
+#include <vector>
+
+#include "common/global_flags.h"
+
+namespace xllm {
+namespace {
+
+class ScopedConstrainedDecodingFlag final {
+ public:
+  explicit ScopedConstrainedDecodingFlag(bool value)
+      : old_value_(FLAGS_enable_constrained_decoding) {
+    FLAGS_enable_constrained_decoding = value;
+  }
+
+  ~ScopedConstrainedDecodingFlag() {
+    FLAGS_enable_constrained_decoding = old_value_;
+  }
+
+ private:
+  bool old_value_;
+};
+
+void write_vocab_file(
+    const std::filesystem::path& path,
+    const std::vector<std::pair<int64_t, RecTokenTriple>>& records) {
+  std::ofstream ofs(path, std::ios::binary | std::ios::trunc);
+  CHECK(ofs.is_open()) << "Failed to open test vocab file: " << path;
+  for (const auto& record : records) {
+    const int64_t item_id = record.first;
+    const RecTokenTriple& tokens = record.second;
+    ofs.write(reinterpret_cast<const char*>(&item_id), sizeof(item_id));
+    ofs.write(reinterpret_cast<const char*>(tokens.data()),
+              REC_TOKEN_SIZE * sizeof(int32_t));
+  }
+}
+
+std::vector<int32_t> sorted_next_tokens(
+    const std::unordered_set<int32_t>& token_set) {
+  std::vector<int32_t> tokens(token_set.begin(), token_set.end());
+  std::sort(tokens.begin(), tokens.end());
+  return tokens;
+}
+
+std::vector<int32_t> prefix1_values_for_token(const RecConstraintTables& tables,
+                                              int32_t t0) {
+  const int32_t begin = tables.prefix1_offsets[static_cast<size_t>(t0)];
+  const int32_t end = tables.prefix1_offsets[static_cast<size_t>(t0) + 1];
+  return std::vector<int32_t>(tables.prefix1_values.begin() + begin,
+                              tables.prefix1_values.begin() + end);
+}
+
+std::vector<int32_t> prefix2_values_for_tokens(
+    const RecConstraintTables& tables,
+    int32_t t0,
+    int32_t t1) {
+  const int32_t begin = tables.prefix1_offsets[static_cast<size_t>(t0)];
+  const int32_t end = tables.prefix1_offsets[static_cast<size_t>(t0) + 1];
+  for (int32_t index = begin; index < end; ++index) {
+    if (tables.prefix1_values[static_cast<size_t>(index)] != t1) {
+      continue;
+    }
+    const int32_t prefix2_begin =
+        tables.prefix2_value_offsets[static_cast<size_t>(index)];
+    const int32_t prefix2_end =
+        tables.prefix2_value_offsets[static_cast<size_t>(index) + 1];
+    return std::vector<int32_t>(tables.prefix2_values.begin() + prefix2_begin,
+                                tables.prefix2_values.begin() + prefix2_end);
+  }
+  return {};
+}
+
+std::vector<int64_t> prefix1_pair_keys_for_token(
+    const RecConstraintTables& tables,
+    int32_t t0) {
+  const int32_t begin = tables.prefix1_offsets[static_cast<size_t>(t0)];
+  const int32_t end = tables.prefix1_offsets[static_cast<size_t>(t0) + 1];
+  return std::vector<int64_t>(tables.prefix1_pair_keys.begin() + begin,
+                              tables.prefix1_pair_keys.begin() + end);
+}
+
+}  // namespace
+
+TEST(RecVocabDictTest, BuildConstraintTablesMatchesLegacyPrefixMap) {
+  ScopedConstrainedDecodingFlag flag(/*value=*/true);
+  const std::filesystem::path vocab_path =
+      std::filesystem::path(::testing::TempDir()) / "rec_vocab_dict_test.bin";
+
+  write_vocab_file(vocab_path,
+                   {
+                       {100, RecTokenTriple{1, 2, 3}},
+                       {101, RecTokenTriple{1, 2, 4}},
+                       {102, RecTokenTriple{1, 5, 6}},
+                       {103, RecTokenTriple{7, 8, 9}},
+                       {104, RecTokenTriple{7, 8, 10}},
+                       {105, RecTokenTriple{1, 2, 3}},
+                   });
+
+  RecVocabDict vocab_dict;
+  ASSERT_TRUE(vocab_dict.initialize(vocab_path.string()));
+
+  const RecConstraintTables tables =
+      vocab_dict.build_constraint_tables(/*vocab_size=*/16);
+
+  EXPECT_EQ(tables.vocab_size, 16);
+  EXPECT_EQ(tables.first_token_ids, std::vector<int32_t>({1, 7}));
+  EXPECT_EQ(prefix1_values_for_token(tables, /*t0=*/1),
+            std::vector<int32_t>({2, 5}));
+  EXPECT_EQ(prefix1_values_for_token(tables, /*t0=*/7),
+            std::vector<int32_t>({8}));
+  EXPECT_EQ(prefix1_pair_keys_for_token(tables, /*t0=*/1),
+            std::vector<int64_t>({18, 21}));
+  EXPECT_EQ(prefix1_pair_keys_for_token(tables, /*t0=*/7),
+            std::vector<int64_t>({120}));
+  EXPECT_TRUE(std::is_sorted(tables.prefix1_pair_keys.begin(),
+                             tables.prefix1_pair_keys.end()));
+  EXPECT_EQ(tables.prefix1_pair_keys.size(), tables.prefix1_values.size());
+  EXPECT_TRUE(prefix1_values_for_token(tables, /*t0=*/0).empty());
+  EXPECT_EQ(prefix2_values_for_tokens(tables, /*t0=*/1, /*t1=*/2),
+            std::vector<int32_t>({3, 4}));
+  EXPECT_EQ(prefix2_values_for_tokens(tables, /*t0=*/1, /*t1=*/5),
+            std::vector<int32_t>({6}));
+  EXPECT_EQ(prefix2_values_for_tokens(tables, /*t0=*/7, /*t1=*/8),
+            std::vector<int32_t>({9, 10}));
+  EXPECT_EQ(tables.prefix2_value_offsets.size(),
+            tables.prefix1_values.size() + 1);
+  EXPECT_EQ(tables.max_first_degree, 2);
+  EXPECT_EQ(tables.max_prefix1_degree, 2);
+  EXPECT_EQ(tables.max_prefix2_degree, 2);
+
+  std::vector<int32_t> empty_prefix;
+  EXPECT_EQ(sorted_next_tokens(vocab_dict.get_next_tokens_by_prefix_tokens(
+                Slice<int32_t>(empty_prefix))),
+            tables.first_token_ids);
+
+  std::vector<int32_t> prefix1{1};
+  EXPECT_EQ(sorted_next_tokens(vocab_dict.get_next_tokens_by_prefix_tokens(
+                Slice<int32_t>(prefix1))),
+            prefix1_values_for_token(tables, /*t0=*/1));
+
+  std::vector<int32_t> prefix2{1, 2};
+  EXPECT_EQ(sorted_next_tokens(vocab_dict.get_next_tokens_by_prefix_tokens(
+                Slice<int32_t>(prefix2))),
+            prefix2_values_for_tokens(tables, /*t0=*/1, /*t1=*/2));
+
+  std::filesystem::remove(vocab_path);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/framework/tokenizer/CMakeLists.txt

@@ -0,0 +1,12 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    fast_tokenizer_test
+  SRCS
+    fast_tokenizer_test.cpp
+  DEPS
+    :tokenizer
+    glog::glog
+    GTest::gtest_main
+)

upstream_ref/xllm/tests/core/framework/tokenizer/fast_tokenizer_test.cpp

@@ -0,0 +1,423 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+Copyright 2024 The ScaleLLM Authors. All Rights Reserved.
+
+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/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 "framework/tokenizer/fast_tokenizer.h"
+
+#include <gtest/gtest.h>
+
+#include <filesystem>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#include "framework/tokenizer/tokenizer_args.h"
+
+namespace xllm {
+
+std::string CreateTestTokenizerJson(const std::string& filepath,
+                                    const std::string& truncation_json,
+                                    const std::string& padding_json);
+
+// Helper function to create a minimal valid tokenizer.json file for testing
+// This creates a simple BPE tokenizer with a small vocabulary
+std::string CreateTestTokenizerJson(const std::string& filepath) {
+  return CreateTestTokenizerJson(filepath, "null", "null");
+}
+
+std::string CreateTestTokenizerJson(const std::string& filepath,
+                                    const std::string& truncation_json,
+                                    const std::string& padding_json) {
+  // Minimal valid tokenizer.json for testing
+  // This is a simplified BPE tokenizer configuration compatible with
+  // HuggingFace tokenizers
+  const std::string tokenizer_json = R"({
+  "version": "1.0",
+  "truncation": __TRUNCATION__,
+  "padding": __PADDING__,
+  "added_tokens": [
+    {"id": 0, "content": "<|bos|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true},
+    {"id": 1, "content": "<|eos|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true},
+    {"id": 2, "content": "hello", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": false},
+    {"id": 3, "content": "world", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": false},
+    {"id": 4, "content": "test", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": false}
+  ],
+  "normalizer": {
+    "type": "NFC"
+  },
+  "pre_tokenizer": {
+    "type": "Whitespace"
+  },
+  "post_processor": null,
+  "decoder": {
+    "type": "ByteLevel",
+    "add_prefix_space": false,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": null,
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "vocab": {
+      "<|bos|>": 0,
+      "<|eos|>": 1,
+      "hello": 2,
+      "world": 3,
+      "test": 4,
+      "h": 5,
+      "e": 6,
+      "l": 7,
+      "o": 8,
+      "w": 9,
+      "r": 10,
+      "d": 11,
+      "t": 12,
+      "s": 13,
+      " ": 14
+    },
+    "merges": []
+  }
+})";
+  std::string rendered = tokenizer_json;
+  size_t pos = rendered.find("__TRUNCATION__");
+  rendered.replace(pos, std::string("__TRUNCATION__").size(), truncation_json);
+  pos = rendered.find("__PADDING__");
+  rendered.replace(pos, std::string("__PADDING__").size(), padding_json);
+
+  std::ofstream file(filepath);
+  if (!file.is_open()) {
+    return "";
+  }
+  file << rendered;
+  file.close();
+  return filepath;
+}
+
+class FastTokenizerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Create a temporary directory for test files
+    test_dir_ = std::filesystem::temp_directory_path() / "fast_tokenizer_test";
+    std::filesystem::create_directories(test_dir_);
+
+    // Create test tokenizer.json file
+    tokenizer_json_path_ = test_dir_ / "tokenizer.json";
+    CreateTestTokenizerJson(tokenizer_json_path_.string());
+  }
+
+  void TearDown() override {
+    // Clean up test files
+    if (std::filesystem::exists(test_dir_)) {
+      std::filesystem::remove_all(test_dir_);
+    }
+  }
+
+  std::filesystem::path test_dir_;
+  std::filesystem::path tokenizer_json_path_;
+};
+
+// Test that BOS token is added when add_bos_token is true
+TEST_F(FastTokenizerTest, AddBosToken) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = true;
+  args.bos_token() = "<|bos|>";
+  args.add_eos_token() = false;
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("hello world", &ids);
+
+  ASSERT_TRUE(success);
+  ASSERT_FALSE(ids.empty());
+
+  // Check that BOS token (ID 0) is at the beginning
+  EXPECT_EQ(ids[0], 0) << "BOS token should be at the beginning";
+
+  // Verify that the rest of the tokens are present
+  // The tokenizer should encode "hello world" into token IDs
+  EXPECT_GT(ids.size(), 1) << "Should have more than just BOS token";
+}
+
+// Test that EOS token is added when add_eos_token is true
+TEST_F(FastTokenizerTest, AddEosToken) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = false;
+  args.add_eos_token() = true;
+  args.eos_token() = "<|eos|>";
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("hello world", &ids);
+
+  ASSERT_TRUE(success);
+  ASSERT_FALSE(ids.empty());
+
+  // Check that EOS token (ID 1) is at the end
+  EXPECT_EQ(ids.back(), 1) << "EOS token should be at the end";
+
+  // Verify that the rest of the tokens are present
+  EXPECT_GT(ids.size(), 1) << "Should have more than just EOS token";
+}
+
+// Test that both BOS and EOS tokens are added when both flags are true
+TEST_F(FastTokenizerTest, AddBothBosAndEosTokens) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = true;
+  args.bos_token() = "<|bos|>";
+  args.add_eos_token() = true;
+  args.eos_token() = "<|eos|>";
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("hello world", &ids);
+
+  ASSERT_TRUE(success);
+  ASSERT_GE(ids.size(), 2) << "Should have at least BOS and EOS tokens";
+
+  // Check that BOS token (ID 0) is at the beginning
+  EXPECT_EQ(ids[0], 0) << "BOS token should be at the beginning";
+
+  // Check that EOS token (ID 1) is at the end
+  EXPECT_EQ(ids.back(), 1) << "EOS token should be at the end";
+
+  // Verify that there are tokens in between
+  EXPECT_GT(ids.size(), 2) << "Should have tokens between BOS and EOS";
+}
+
+// Test that no special tokens are added when both flags are false
+TEST_F(FastTokenizerTest, NoSpecialTokens) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = false;
+  args.add_eos_token() = false;
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("hello world", &ids);
+
+  ASSERT_TRUE(success);
+  ASSERT_FALSE(ids.empty());
+
+  // Check that BOS token (ID 0) is NOT at the beginning
+  EXPECT_NE(ids[0], 0) << "BOS token should not be present";
+
+  // Check that EOS token (ID 1) is NOT at the end
+  EXPECT_NE(ids.back(), 1) << "EOS token should not be present";
+}
+
+// Test that BOS token is not added when bos_token is empty
+TEST_F(FastTokenizerTest, AddBosTokenWithEmptyToken) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = true;
+  args.bos_token() = "";  // Empty token
+  args.add_eos_token() = false;
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("hello world", &ids);
+
+  ASSERT_TRUE(success);
+  ASSERT_FALSE(ids.empty());
+
+  // BOS token should not be added because bos_token is empty
+  EXPECT_NE(ids[0], 0)
+      << "BOS token should not be added when bos_token is empty";
+}
+
+// Test that EOS token is not added when eos_token is empty
+TEST_F(FastTokenizerTest, AddEosTokenWithEmptyToken) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = false;
+  args.add_eos_token() = true;
+  args.eos_token() = "";  // Empty token
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("hello world", &ids);
+
+  ASSERT_TRUE(success);
+  ASSERT_FALSE(ids.empty());
+
+  // EOS token should not be added because eos_token is empty
+  EXPECT_NE(ids.back(), 1)
+      << "EOS token should not be added when eos_token is empty";
+}
+
+// Test that BOS token is not duplicated when it already exists
+// This simulates the case where the underlying tokenizer already added BOS
+TEST_F(FastTokenizerTest, SkipBosTokenWhenAlreadyPresent) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = true;
+  args.bos_token() = "<|bos|>";
+  args.add_eos_token() = false;
+
+  FastTokenizer tokenizer(args);
+
+  // First encode a text that doesn't start with BOS
+  std::vector<int32_t> ids1;
+  bool success1 = tokenizer.encode("hello world", &ids1);
+  ASSERT_TRUE(success1);
+  ASSERT_FALSE(ids1.empty());
+
+  // Verify BOS token was added
+  EXPECT_EQ(ids1[0], 0) << "BOS token should be added";
+  size_t size_with_bos = ids1.size();
+
+  // Now encode text that starts with BOS token directly
+  // We'll encode the BOS token itself, which should result in BOS being the
+  // first token
+  std::vector<int32_t> ids2;
+  bool success2 = tokenizer.encode("<|bos|> hello world", &ids2);
+  ASSERT_TRUE(success2);
+  ASSERT_FALSE(ids2.empty());
+
+  // The BOS token should be present, but we should not have added it twice
+  // Count how many times BOS token (ID 0) appears at the beginning
+  int bos_count_at_start = 0;
+  for (size_t i = 0; i < ids2.size() && ids2[i] == 0; ++i) {
+    bos_count_at_start++;
+  }
+
+  // Should have at most one BOS token at the beginning
+  EXPECT_LE(bos_count_at_start, 1)
+      << "BOS token should not be duplicated when already present";
+}
+
+// Test that EOS token is not duplicated when it already exists
+TEST_F(FastTokenizerTest, SkipEosTokenWhenAlreadyPresent) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = false;
+  args.add_eos_token() = true;
+  args.eos_token() = "<|eos|>";
+
+  FastTokenizer tokenizer(args);
+
+  // First encode a text that doesn't end with EOS
+  std::vector<int32_t> ids1;
+  bool success1 = tokenizer.encode("hello world", &ids1);
+  ASSERT_TRUE(success1);
+  ASSERT_FALSE(ids1.empty());
+
+  // Verify EOS token was added
+  EXPECT_EQ(ids1.back(), 1) << "EOS token should be added";
+  size_t size_with_eos = ids1.size();
+
+  // Now encode text that ends with EOS token directly
+  std::vector<int32_t> ids2;
+  bool success2 = tokenizer.encode("hello world <|eos|>", &ids2);
+  ASSERT_TRUE(success2);
+  ASSERT_FALSE(ids2.empty());
+
+  // The EOS token should be present, but we should not have added it twice
+  // Count how many times EOS token (ID 1) appears at the end
+  int eos_count_at_end = 0;
+  for (int i = ids2.size() - 1; i >= 0 && ids2[i] == 1; --i) {
+    eos_count_at_end++;
+  }
+
+  // Should have at most one EOS token at the end
+  EXPECT_LE(eos_count_at_end, 1)
+      << "EOS token should not be duplicated when already present";
+}
+
+// Test that both BOS and EOS tokens are not duplicated when both already exist
+TEST_F(FastTokenizerTest, SkipBothBosAndEosTokensWhenAlreadyPresent) {
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = tokenizer_json_path_.string();
+  args.add_bos_token() = true;
+  args.bos_token() = "<|bos|>";
+  args.add_eos_token() = true;
+  args.eos_token() = "<|eos|>";
+
+  FastTokenizer tokenizer(args);
+
+  // Encode text that already contains both BOS and EOS tokens
+  std::vector<int32_t> ids;
+  bool success = tokenizer.encode("<|bos|> hello world <|eos|>", &ids);
+  ASSERT_TRUE(success);
+  ASSERT_FALSE(ids.empty());
+
+  // Count BOS tokens at the beginning
+  int bos_count_at_start = 0;
+  for (size_t i = 0; i < ids.size() && ids[i] == 0; ++i) {
+    bos_count_at_start++;
+  }
+
+  // Count EOS tokens at the end
+  int eos_count_at_end = 0;
+  for (int i = ids.size() - 1; i >= 0 && ids[i] == 1; --i) {
+    eos_count_at_end++;
+  }
+
+  // Should have at most one BOS token at the beginning
+  EXPECT_LE(bos_count_at_start, 1)
+      << "BOS token should not be duplicated when already present";
+
+  // Should have at most one EOS token at the end
+  EXPECT_LE(eos_count_at_end, 1)
+      << "EOS token should not be duplicated when already present";
+}
+
+TEST_F(FastTokenizerTest, IgnoreTokenizerJsonPaddingAndTruncation) {
+  std::filesystem::path padded_tokenizer_path =
+      test_dir_ / "tokenizer_with_padding.json";
+  ASSERT_FALSE(
+      CreateTestTokenizerJson(
+          padded_tokenizer_path.string(),
+          R"({"direction":"Right","max_length":4,"strategy":"LongestFirst","stride":0})",
+          R"({"strategy":{"Fixed":4},"direction":"Right","pad_to_multiple_of":null,"pad_id":0,"pad_type_id":0,"pad_token":"<|bos|>"})")
+          .empty());
+
+  TokenizerArgs args;
+  args.tokenizer_type() = "fast";
+  args.vocab_file() = padded_tokenizer_path.string();
+
+  FastTokenizer tokenizer(args);
+
+  std::vector<int32_t> ids;
+  ASSERT_TRUE(tokenizer.encode("hello", &ids));
+  EXPECT_EQ(ids.size(), 1)
+      << "FastTokenizer should ignore tokenizer.json padding/truncation.";
+  EXPECT_EQ(ids[0], 2);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/kernels/CMakeLists.txt

@@ -0,0 +1,7 @@
+if(USE_CUDA)
+  add_subdirectory(cuda)
+endif()
+
+if(USE_NPU)
+  add_subdirectory(npu)
+endif()

upstream_ref/xllm/tests/core/kernels/cuda/CMakeLists.txt

@@ -0,0 +1,71 @@
+include(cc_test)
+
+add_subdirectory(moe)
+add_subdirectory(xattention)
+
+cc_test(
+  NAME
+    activation_test
+  SRCS
+    activation_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+
+option(XLLM_ENABLE_BLOCK_COPY_TEST
+       "Build and register the expensive CUDA block_copy_test"
+       OFF)
+
+if(XLLM_ENABLE_BLOCK_COPY_TEST)
+  cc_test(
+    NAME
+      block_copy_test
+    SRCS
+      block_copy_test.cpp
+    DEPS
+      :cuda_kernels
+      torch
+      GTest::gtest_main
+      glog::glog
+  )
+endif()
+
+cc_test(
+  NAME
+    fused_qknorm_rope_test
+  SRCS
+    fused_qknorm_rope_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+)
+
+cc_test(
+  NAME
+    cutlass_scaled_mm_test
+  SRCS
+    cutlass_scaled_mm_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+target_link_libraries(cutlass_scaled_mm_test PRIVATE brpc)
+
+cc_test(
+  NAME
+    fp8_quant_test
+  SRCS
+    fp8_quant_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+target_link_libraries(fp8_quant_test PRIVATE brpc)

upstream_ref/xllm/tests/core/kernels/cuda/activation_test.cpp

@@ -0,0 +1,112 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/cuda.h>
+#include <torch/torch.h>
+
+#include <string>
+#include <vector>
+
+#include "cuda_ops_api.h"
+
+namespace xllm::kernel::cuda {
+namespace test {
+
+namespace {
+
+torch::Tensor torch_reference_act_and_mul(const torch::Tensor& input,
+                                          const std::string& act_mode) {
+  CHECK_EQ(input.size(-1) % 2, 0) << "Last dim must be even.";
+  const int64_t d = input.size(-1) / 2;
+  auto x = input.slice(-1, 0, d);
+  auto y = input.slice(-1, d, 2 * d);
+
+  if (act_mode == "silu") {
+    return (x * torch::sigmoid(x)) * y;
+  }
+  if (act_mode == "gelu") {
+    return torch::gelu(x, "none") * y;
+  }
+  if (act_mode == "gelu_tanh") {
+    return torch::gelu(x, "tanh") * y;
+  }
+  LOG(FATAL) << "Unsupported act mode in test: " << act_mode;
+  return torch::Tensor();
+}
+
+std::string to_string(torch::ScalarType dtype) {
+  switch (dtype) {
+    case torch::kFloat32:
+      return "float32";
+    case torch::kFloat16:
+      return "float16";
+    case torch::kBFloat16:
+      return "bfloat16";
+    default:
+      return "unknown";
+  }
+}
+
+class ActAndMulKernelTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    torch::manual_seed(2026);
+    device_ = torch::Device(torch::kCUDA, 0);
+  }
+
+  void run_and_check(torch::ScalarType dtype,
+                     const std::string& act_mode,
+                     int64_t d) const {
+    const auto opts = torch::TensorOptions().device(device_).dtype(dtype);
+    torch::Tensor input = torch::randn({4, 7, 2 * d}, opts) * 0.5;
+    torch::Tensor output = torch::empty({4, 7, d}, opts);
+
+    torch::Tensor reference = torch_reference_act_and_mul(input, act_mode);
+    act_and_mul(output, input, act_mode);
+
+    const double atol = 5e-3;
+    const double rtol = 5e-3;
+    EXPECT_TRUE(torch::allclose(output, reference, rtol, atol))
+        << "Mismatch for act_mode=" << act_mode
+        << ", dtype=" << to_string(dtype) << ", d=" << d;
+  }
+
+  torch::Device device_ = torch::Device(torch::kCPU);
+};
+
+TEST_F(ActAndMulKernelTest, MatchesTorchReference) {
+  const std::vector<torch::ScalarType> dtypes = {torch::kFloat16,
+                                                 torch::kBFloat16};
+  const std::vector<std::string> act_modes = {"silu", "gelu", "gelu_tanh"};
+  // Cover scalar fallback (d<VEC_SIZE), vectorized path, and tail cleanup.
+  const std::vector<int64_t> dims = {3, 64, 129};
+
+  for (auto dtype : dtypes) {
+    for (const auto& act_mode : act_modes) {
+      for (int64_t d : dims) {
+        run_and_check(dtype, act_mode, d);
+      }
+    }
+  }
+}
+
+}  // namespace
+
+}  // namespace test
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/tests/core/kernels/cuda/block_copy_test.cpp

@@ -0,0 +1,688 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <ATen/cuda/CUDAEvent.h>
+#include <c10/cuda/CUDAStream.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/cuda.h>
+#include <torch/torch.h>
+
+#include <functional>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "core/kernels/cuda/cuda_ops_api.h"
+
+namespace xllm::kernel::cuda {
+namespace test {
+namespace {
+
+struct BlockCopyCaseConfig {
+  int64_t num_layers;
+  int64_t num_blocks;
+  int64_t block_size;
+  int64_t num_heads;
+  int64_t head_dim;
+  std::vector<int32_t> src_blocks;
+  std::vector<int32_t> dst_blocks;
+  std::vector<int32_t> cum_sum;
+};
+
+struct PerfBenchmarkCase {
+  std::string name;
+  BlockCopyCaseConfig config;
+  torch::ScalarType dtype;
+  int32_t warmup_iters;
+  int32_t measure_iters;
+};
+
+struct PerfCompareResult {
+  double kernel_ms;
+  double native_ms;
+  double speedup;
+  double logical_copy_gbps;
+  double traffic_gbps;
+};
+
+struct BlockCopyLaunchInputs {
+  torch::Tensor key_ptr_tensor;
+  torch::Tensor value_ptr_tensor;
+  torch::Tensor src_tensor;
+  torch::Tensor dst_tensor;
+  torch::Tensor cum_sum_tensor;
+  int64_t numel_per_block;
+};
+
+struct NativeBlockCopyLaunchInputs {
+  torch::Tensor src_tensor;
+  torch::Tensor dst_tensor;
+};
+
+void apply_reference_block_copy(const std::vector<torch::Tensor>& key_caches,
+                                const std::vector<torch::Tensor>& value_caches,
+                                const std::vector<int32_t>& src_blocks,
+                                const std::vector<int32_t>& dst_blocks,
+                                const std::vector<int32_t>& cum_sum,
+                                std::vector<torch::Tensor>& ref_k_caches,
+                                std::vector<torch::Tensor>& ref_v_caches) {
+  for (size_t layer_idx = 0; layer_idx < key_caches.size(); ++layer_idx) {
+    ref_k_caches[layer_idx] = key_caches[layer_idx].clone();
+    ref_v_caches[layer_idx] = value_caches[layer_idx].clone();
+  }
+
+  for (size_t group_idx = 0; group_idx < src_blocks.size(); ++group_idx) {
+    const int32_t src_block = src_blocks[group_idx];
+    const int32_t dst_begin = group_idx == 0 ? 0 : cum_sum[group_idx - 1];
+    const int32_t dst_end = cum_sum[group_idx];
+    for (int32_t dst_idx = dst_begin; dst_idx < dst_end; ++dst_idx) {
+      const int32_t dst_block = dst_blocks[dst_idx];
+      for (size_t layer_idx = 0; layer_idx < ref_k_caches.size(); ++layer_idx) {
+        ref_k_caches[layer_idx][dst_block].copy_(
+            ref_k_caches[layer_idx][src_block]);
+        ref_v_caches[layer_idx][dst_block].copy_(
+            ref_v_caches[layer_idx][src_block]);
+      }
+    }
+  }
+}
+
+std::vector<int64_t> flatten_src_blocks_for_native(
+    const std::vector<int32_t>& src_blocks,
+    const std::vector<int32_t>& cum_sum) {
+  std::vector<int64_t> flat_src_blocks;
+  flat_src_blocks.reserve(cum_sum.empty() ? 0 : cum_sum.back());
+  for (size_t group_idx = 0; group_idx < src_blocks.size(); ++group_idx) {
+    const int32_t begin = group_idx == 0 ? 0 : cum_sum[group_idx - 1];
+    const int32_t end = cum_sum[group_idx];
+    for (int32_t dst_idx = begin; dst_idx < end; ++dst_idx) {
+      flat_src_blocks.push_back(src_blocks[group_idx]);
+    }
+  }
+  return flat_src_blocks;
+}
+
+void native_block_copy(std::vector<torch::Tensor>& key_caches,
+                       std::vector<torch::Tensor>& value_caches,
+                       const NativeBlockCopyLaunchInputs& launch_inputs) {
+  for (size_t layer_idx = 0; layer_idx < key_caches.size(); ++layer_idx) {
+    auto selected_keys =
+        torch::index_select(key_caches[layer_idx], 0, launch_inputs.src_tensor);
+    auto selected_values = torch::index_select(
+        value_caches[layer_idx], 0, launch_inputs.src_tensor);
+    key_caches[layer_idx].index_copy_(
+        0, launch_inputs.dst_tensor, selected_keys);
+    value_caches[layer_idx].index_copy_(
+        0, launch_inputs.dst_tensor, selected_values);
+  }
+}
+
+BlockCopyLaunchInputs prepare_block_copy_launch_inputs(
+    const std::vector<torch::Tensor>& key_caches,
+    const std::vector<torch::Tensor>& value_caches,
+    const std::vector<int32_t>& src_blocks,
+    const std::vector<int32_t>& dst_blocks,
+    const std::vector<int32_t>& cum_sum,
+    const torch::Device& device) {
+  std::vector<int64_t> key_ptrs;
+  std::vector<int64_t> value_ptrs;
+  key_ptrs.reserve(key_caches.size());
+  value_ptrs.reserve(value_caches.size());
+  for (size_t layer_idx = 0; layer_idx < key_caches.size(); ++layer_idx) {
+    key_ptrs.push_back(
+        reinterpret_cast<int64_t>(key_caches[layer_idx].data_ptr()));
+    value_ptrs.push_back(
+        reinterpret_cast<int64_t>(value_caches[layer_idx].data_ptr()));
+  }
+
+  const auto ptr_opts =
+      torch::TensorOptions().device(device).dtype(torch::kInt64);
+  const auto idx_opts =
+      torch::TensorOptions().device(device).dtype(torch::kInt32);
+  return {
+      .key_ptr_tensor = torch::tensor(key_ptrs, ptr_opts),
+      .value_ptr_tensor = torch::tensor(value_ptrs, ptr_opts),
+      .src_tensor = torch::tensor(src_blocks, idx_opts),
+      .dst_tensor = torch::tensor(dst_blocks, idx_opts),
+      .cum_sum_tensor = torch::tensor(cum_sum, idx_opts),
+      .numel_per_block = key_caches[0][0].numel(),
+  };
+}
+
+NativeBlockCopyLaunchInputs prepare_native_block_copy_launch_inputs(
+    const std::vector<int32_t>& src_blocks,
+    const std::vector<int32_t>& dst_blocks,
+    const std::vector<int32_t>& cum_sum,
+    const torch::Device& device) {
+  auto flat_src_blocks = flatten_src_blocks_for_native(src_blocks, cum_sum);
+  std::vector<int64_t> flat_dst_blocks(dst_blocks.begin(), dst_blocks.end());
+  return {
+      .src_tensor = torch::tensor(
+          flat_src_blocks,
+          torch::TensorOptions().device(device).dtype(torch::kLong)),
+      .dst_tensor = torch::tensor(
+          flat_dst_blocks,
+          torch::TensorOptions().device(device).dtype(torch::kLong)),
+  };
+}
+
+void kernel_block_copy(const BlockCopyLaunchInputs& launch_inputs,
+                       torch::ScalarType dtype) {
+  block_copy(launch_inputs.key_ptr_tensor,
+             launch_inputs.value_ptr_tensor,
+             launch_inputs.src_tensor,
+             launch_inputs.dst_tensor,
+             launch_inputs.cum_sum_tensor,
+             launch_inputs.numel_per_block,
+             dtype);
+}
+
+double measure_cuda_time_ms(const std::function<void()>& fn,
+                            int32_t warmup_iters,
+                            int32_t measure_iters) {
+  for (int32_t iter = 0; iter < warmup_iters; ++iter) {
+    fn();
+  }
+  torch::cuda::synchronize();
+
+  const auto stream = c10::cuda::getCurrentCUDAStream();
+  at::cuda::CUDAEvent start_event(cudaEventDefault);
+  at::cuda::CUDAEvent stop_event(cudaEventDefault);
+  start_event.record(stream);
+  for (int32_t iter = 0; iter < measure_iters; ++iter) {
+    fn();
+  }
+  stop_event.record(stream);
+  stop_event.synchronize();
+
+  const float elapsed_ms = start_event.elapsed_time(stop_event);
+  return static_cast<double>(elapsed_ms) / measure_iters;
+}
+
+std::vector<torch::Tensor> make_random_caches(const BlockCopyCaseConfig& config,
+                                              const torch::Device& device,
+                                              torch::ScalarType dtype) {
+  std::vector<torch::Tensor> caches;
+  caches.reserve(config.num_layers);
+  const auto opts = torch::TensorOptions().device(device).dtype(dtype);
+  for (int64_t layer_idx = 0; layer_idx < config.num_layers; ++layer_idx) {
+    caches.push_back(torch::randn({config.num_blocks,
+                                   config.block_size,
+                                   config.num_heads,
+                                   config.head_dim},
+                                  opts));
+  }
+  return caches;
+}
+
+void expect_caches_allclose(const std::vector<torch::Tensor>& lhs,
+                            const std::vector<torch::Tensor>& rhs,
+                            double rtol,
+                            double atol) {
+  ASSERT_EQ(lhs.size(), rhs.size());
+  for (size_t idx = 0; idx < lhs.size(); ++idx) {
+    EXPECT_TRUE(torch::allclose(lhs[idx], rhs[idx], rtol, atol))
+        << "cache mismatch at layer=" << idx;
+  }
+}
+
+void run_accuracy_compare_case(const BlockCopyCaseConfig& config,
+                               torch::ScalarType dtype,
+                               double rtol,
+                               double atol) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test.";
+  }
+
+  torch::manual_seed(2026);
+  const auto device = torch::Device(torch::kCUDA, 0);
+
+  auto base_key_caches = make_random_caches(config, device, dtype);
+  auto base_value_caches = make_random_caches(config, device, dtype);
+  auto kernel_key_caches = base_key_caches;
+  auto kernel_value_caches = base_value_caches;
+  auto native_key_caches = base_key_caches;
+  auto native_value_caches = base_value_caches;
+
+  auto kernel_launch_inputs =
+      prepare_block_copy_launch_inputs(kernel_key_caches,
+                                       kernel_value_caches,
+                                       config.src_blocks,
+                                       config.dst_blocks,
+                                       config.cum_sum,
+                                       device);
+  auto native_launch_inputs = prepare_native_block_copy_launch_inputs(
+      config.src_blocks, config.dst_blocks, config.cum_sum, device);
+
+  kernel_block_copy(kernel_launch_inputs, dtype);
+  native_block_copy(
+      native_key_caches, native_value_caches, native_launch_inputs);
+  torch::cuda::synchronize();
+
+  expect_caches_allclose(kernel_key_caches, native_key_caches, rtol, atol);
+  expect_caches_allclose(kernel_value_caches, native_value_caches, rtol, atol);
+}
+
+std::string dtype_to_string(torch::ScalarType dtype) {
+  switch (dtype) {
+    case torch::kHalf:
+      return "fp16";
+    case torch::kBFloat16:
+      return "bf16";
+    case torch::kFloat:
+      return "fp32";
+    default:
+      return c10::toString(dtype);
+  }
+}
+
+int64_t get_total_dst_copies(const BlockCopyCaseConfig& config) {
+  return static_cast<int64_t>(config.dst_blocks.size());
+}
+
+double get_logical_copy_bytes_per_iter(const BlockCopyCaseConfig& config,
+                                       torch::ScalarType dtype) {
+  const int64_t numel_per_block =
+      config.block_size * config.num_heads * config.head_dim;
+  const int64_t bytes_per_elem = c10::elementSize(dtype);
+  const int64_t total_dst_copies = get_total_dst_copies(config);
+  const int64_t total_elements =
+      2LL * config.num_layers * total_dst_copies * numel_per_block;
+  return static_cast<double>(total_elements) * bytes_per_elem;
+}
+
+double get_traffic_bytes_per_iter(const BlockCopyCaseConfig& config,
+                                  torch::ScalarType dtype) {
+  return get_logical_copy_bytes_per_iter(config, dtype) * 2.0;
+}
+
+std::string format_perf_case_summary(const PerfBenchmarkCase& benchmark_case,
+                                     const PerfCompareResult& result) {
+  const auto& config = benchmark_case.config;
+  const int64_t total_dst_copies = get_total_dst_copies(config);
+  const double avg_fanout =
+      config.src_blocks.empty()
+          ? 0.0
+          : static_cast<double>(total_dst_copies) /
+                static_cast<double>(config.src_blocks.size());
+  std::ostringstream oss;
+  oss << "block_copy bench [" << benchmark_case.name
+      << "] dtype=" << dtype_to_string(benchmark_case.dtype)
+      << ", layers=" << config.num_layers << ", blocks=" << config.num_blocks
+      << ", block_size=" << config.block_size << ", heads=" << config.num_heads
+      << ", head_dim=" << config.head_dim
+      << ", groups=" << config.src_blocks.size()
+      << ", total_dst=" << total_dst_copies << ", avg_fanout=" << avg_fanout
+      << ", kernel=" << result.kernel_ms << " ms"
+      << ", native=" << result.native_ms << " ms"
+      << ", speedup=" << result.speedup << "x"
+      << ", logical_bw=" << result.logical_copy_gbps << " GB/s"
+      << ", traffic_bw=" << result.traffic_gbps << " GB/s";
+  return oss.str();
+}
+
+PerfCompareResult run_perf_compare_case(const BlockCopyCaseConfig& config,
+                                        torch::ScalarType dtype,
+                                        int32_t warmup_iters,
+                                        int32_t measure_iters) {
+  torch::manual_seed(2026);
+  const auto device = torch::Device(torch::kCUDA, 0);
+
+  auto kernel_key_caches = make_random_caches(config, device, dtype);
+  auto kernel_value_caches = make_random_caches(config, device, dtype);
+  auto native_key_caches = kernel_key_caches;
+  auto native_value_caches = kernel_value_caches;
+
+  auto kernel_launch_inputs =
+      prepare_block_copy_launch_inputs(kernel_key_caches,
+                                       kernel_value_caches,
+                                       config.src_blocks,
+                                       config.dst_blocks,
+                                       config.cum_sum,
+                                       device);
+  auto native_launch_inputs = prepare_native_block_copy_launch_inputs(
+      config.src_blocks, config.dst_blocks, config.cum_sum, device);
+
+  const double kernel_ms = measure_cuda_time_ms(
+      [&]() { kernel_block_copy(kernel_launch_inputs, dtype); },
+      warmup_iters,
+      measure_iters);
+
+  const double native_ms = measure_cuda_time_ms(
+      [&]() {
+        native_block_copy(
+            native_key_caches, native_value_caches, native_launch_inputs);
+      },
+      warmup_iters,
+      measure_iters);
+
+  expect_caches_allclose(kernel_key_caches, native_key_caches, 1e-5, 1e-5);
+  expect_caches_allclose(kernel_value_caches, native_value_caches, 1e-5, 1e-5);
+
+  const double logical_copy_bytes =
+      get_logical_copy_bytes_per_iter(config, dtype);
+  const double traffic_bytes = get_traffic_bytes_per_iter(config, dtype);
+  const double speedup = native_ms / kernel_ms;
+  const double logical_copy_gbps = logical_copy_bytes / (kernel_ms * 1.0e6);
+  const double traffic_gbps = traffic_bytes / (kernel_ms * 1.0e6);
+
+  EXPECT_GT(kernel_ms, 0.0);
+  EXPECT_GT(native_ms, 0.0);
+  return PerfCompareResult{
+      .kernel_ms = kernel_ms,
+      .native_ms = native_ms,
+      .speedup = speedup,
+      .logical_copy_gbps = logical_copy_gbps,
+      .traffic_gbps = traffic_gbps,
+  };
+}
+
+void run_multi_shape_perf_benchmark(
+    const std::vector<PerfBenchmarkCase>& benchmark_cases) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test.";
+  }
+  for (const auto& benchmark_case : benchmark_cases) {
+    SCOPED_TRACE(benchmark_case.name);
+    const auto result = run_perf_compare_case(benchmark_case.config,
+                                              benchmark_case.dtype,
+                                              benchmark_case.warmup_iters,
+                                              benchmark_case.measure_iters);
+    LOG(INFO) << format_perf_case_summary(benchmark_case, result);
+  }
+}
+
+}  // namespace
+
+TEST(BlockCopyTest, KernelMatchesReferenceFp16) {
+  run_accuracy_compare_case(
+      BlockCopyCaseConfig{
+          .num_layers = 3,
+          .num_blocks = 8,
+          .block_size = 4,
+          .num_heads = 2,
+          .head_dim = 8,
+          .src_blocks = {1, 4},
+          .dst_blocks = {2, 3, 6},
+          .cum_sum = {2, 3},
+      },
+      torch::kHalf,
+      1e-5,
+      1e-5);
+}
+
+TEST(BlockCopyTest, KernelMatchesReferenceFp32) {
+  run_accuracy_compare_case(
+      BlockCopyCaseConfig{
+          .num_layers = 4,
+          .num_blocks = 10,
+          .block_size = 8,
+          .num_heads = 3,
+          .head_dim = 16,
+          .src_blocks = {1, 4, 7},
+          .dst_blocks = {2, 3, 5, 8, 9},
+          .cum_sum = {2, 4, 5},
+      },
+      torch::kFloat,
+      1e-6,
+      1e-6);
+}
+
+TEST(BlockCopyTest, KernelMatchesNativeFp16) {
+  run_accuracy_compare_case(
+      BlockCopyCaseConfig{
+          .num_layers = 6,
+          .num_blocks = 32,
+          .block_size = 16,
+          .num_heads = 4,
+          .head_dim = 32,
+          .src_blocks = {1, 4, 9, 12},
+          .dst_blocks = {2, 3, 5, 6, 10, 11, 20},
+          .cum_sum = {2, 4, 6, 7},
+      },
+      torch::kHalf,
+      1e-5,
+      1e-5);
+}
+
+TEST(BlockCopyTest, KernelMatchesNativeFp32) {
+  run_accuracy_compare_case(
+      BlockCopyCaseConfig{
+          .num_layers = 5,
+          .num_blocks = 24,
+          .block_size = 12,
+          .num_heads = 3,
+          .head_dim = 24,
+          .src_blocks = {1, 4, 9},
+          .dst_blocks = {2, 3, 5, 6, 10, 11},
+          .cum_sum = {2, 4, 6},
+      },
+      torch::kFloat,
+      1e-6,
+      1e-6);
+}
+
+TEST(BlockCopyTest, PerfCompareKernelVsNativeMultiShapeFp16) {
+  run_multi_shape_perf_benchmark({
+      PerfBenchmarkCase{
+          .name = "tiny_balanced",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 4,
+                  .num_blocks = 32,
+                  .block_size = 16,
+                  .num_heads = 4,
+                  .head_dim = 32,
+                  .src_blocks = {1, 4, 8, 12},
+                  .dst_blocks = {2, 3, 5, 6, 9, 10, 13, 14},
+                  .cum_sum = {2, 4, 6, 8},
+              },
+          .dtype = torch::kHalf,
+          .warmup_iters = 20,
+          .measure_iters = 150,
+      },
+      PerfBenchmarkCase{
+          .name = "tiny_high_fanout",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 4,
+                  .num_blocks = 48,
+                  .block_size = 16,
+                  .num_heads = 4,
+                  .head_dim = 32,
+                  .src_blocks = {1, 8},
+                  .dst_blocks = {2, 3, 4, 5, 6, 9, 10, 11, 12, 13},
+                  .cum_sum = {5, 10},
+              },
+          .dtype = torch::kHalf,
+          .warmup_iters = 20,
+          .measure_iters = 150,
+      },
+      PerfBenchmarkCase{
+          .name = "medium_balanced",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 8,
+                  .num_blocks = 64,
+                  .block_size = 64,
+                  .num_heads = 8,
+                  .head_dim = 128,
+                  .src_blocks = {1, 4, 8, 12, 16, 20, 24, 28},
+                  .dst_blocks = {2,
+                                 3,
+                                 5,
+                                 6,
+                                 9,
+                                 10,
+                                 13,
+                                 14,
+                                 17,
+                                 18,
+                                 21,
+                                 22,
+                                 25,
+                                 26,
+                                 29,
+                                 30},
+                  .cum_sum = {2, 4, 6, 8, 10, 12, 14, 16},
+              },
+          .dtype = torch::kHalf,
+          .warmup_iters = 20,
+          .measure_iters = 120,
+      },
+      PerfBenchmarkCase{
+          .name = "large_many_layers",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 32,
+                  .num_blocks = 256,
+                  .block_size = 64,
+                  .num_heads = 8,
+                  .head_dim = 128,
+                  .src_blocks = {1, 9, 17, 25, 33, 41, 49, 57},
+                  .dst_blocks = {2,
+                                 3,
+                                 10,
+                                 11,
+                                 18,
+                                 19,
+                                 26,
+                                 27,
+                                 34,
+                                 35,
+                                 42,
+                                 43,
+                                 50,
+                                 51,
+                                 58,
+                                 59},
+                  .cum_sum = {2, 4, 6, 8, 10, 12, 14, 16},
+              },
+          .dtype = torch::kHalf,
+          .warmup_iters = 20,
+          .measure_iters = 80,
+      },
+      PerfBenchmarkCase{
+          .name = "large_high_fanout",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 16,
+                  .num_blocks = 256,
+                  .block_size = 64,
+                  .num_heads = 8,
+                  .head_dim = 128,
+                  .src_blocks = {1, 33, 65, 97},
+                  .dst_blocks = {2,  3,  4,  5,  6,  34, 35, 36,  37,  38,
+                                 66, 67, 68, 69, 70, 98, 99, 100, 101, 102},
+                  .cum_sum = {5, 10, 15, 20},
+              },
+          .dtype = torch::kHalf,
+          .warmup_iters = 20,
+          .measure_iters = 80,
+      },
+      PerfBenchmarkCase{
+          .name = "many_groups_sparse",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 16,
+                  .num_blocks = 256,
+                  .block_size = 32,
+                  .num_heads = 8,
+                  .head_dim = 128,
+                  .src_blocks = {1, 5, 9, 13, 17, 21, 25, 29, 33, 37, 41, 45},
+                  .dst_blocks = {2,
+                                 6,
+                                 10,
+                                 14,
+                                 18,
+                                 22,
+                                 26,
+                                 30,
+                                 34,
+                                 38,
+                                 42,
+                                 46,
+                                 3,
+                                 7,
+                                 11,
+                                 15,
+                                 19,
+                                 23},
+                  .cum_sum = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 18},
+              },
+          .dtype = torch::kHalf,
+          .warmup_iters = 20,
+          .measure_iters = 100,
+      },
+  });
+}
+
+TEST(BlockCopyTest, PerfCompareKernelVsNativeMultiShapeFp32) {
+  run_multi_shape_perf_benchmark({
+      PerfBenchmarkCase{
+          .name = "fp32_medium_balanced",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 8,
+                  .num_blocks = 64,
+                  .block_size = 32,
+                  .num_heads = 8,
+                  .head_dim = 64,
+                  .src_blocks = {1, 4, 8, 12, 16, 20},
+                  .dst_blocks = {2, 3, 5, 6, 9, 10, 13, 14, 17, 18, 21, 22},
+                  .cum_sum = {2, 4, 6, 8, 10, 12},
+              },
+          .dtype = torch::kFloat,
+          .warmup_iters = 20,
+          .measure_iters = 120,
+      },
+      PerfBenchmarkCase{
+          .name = "fp32_large_many_layers",
+          .config =
+              BlockCopyCaseConfig{
+                  .num_layers = 24,
+                  .num_blocks = 192,
+                  .block_size = 64,
+                  .num_heads = 8,
+                  .head_dim = 64,
+                  .src_blocks = {1, 9, 17, 25, 33, 41, 49, 57},
+                  .dst_blocks = {2,
+                                 3,
+                                 10,
+                                 11,
+                                 18,
+                                 19,
+                                 26,
+                                 27,
+                                 34,
+                                 35,
+                                 42,
+                                 43,
+                                 50,
+                                 51,
+                                 58,
+                                 59},
+                  .cum_sum = {2, 4, 6, 8, 10, 12, 14, 16},
+              },
+          .dtype = torch::kFloat,
+          .warmup_iters = 20,
+          .measure_iters = 80,
+      },
+  });
+}
+
+}  // namespace test
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/tests/core/kernels/cuda/cutlass_scaled_mm_test.cpp

@@ -0,0 +1,296 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+ *
+ * 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/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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "cuda_ops_api.h"
+#include "cutlass_extensions/common.hpp"
+
+class CutlassScaledMMTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    device_ = torch::Device(torch::kCUDA);
+
+    // Check if FP8 is supported
+    int compute_capability = xllm::kernel::cuda::get_sm_version_num();
+    if (compute_capability < 89) {
+      GTEST_SKIP() << "FP8 requires compute capability >= 8.9 (Ada Lovelace or "
+                      "Hopper), current: "
+                   << compute_capability;
+    }
+  }
+
+  torch::Device device_ = torch::kCPU;
+};
+
+// Test basic FP8 W8A8 matrix multiplication
+TEST_F(CutlassScaledMMTest, BasicFP8W8A8Test) {
+  const int64_t M = 128;
+  const int64_t N = 256;
+  const int64_t K = 512;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+
+  auto fp16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat16);
+
+  // Create input matrices in FP8 format (row-major for A, column-major for B)
+  // Scale down the input range to avoid FP8 quantization error
+  torch::Tensor a = torch::randn({M, K}, fp32_options) * 0.5f;
+  torch::Tensor b = torch::randn({K, N}, fp32_options) * 0.5f;
+
+  // Convert to FP8
+  torch::Tensor a_fp8 = a.to(torch::kFloat8_e4m3fn);
+  torch::Tensor b_fp8 = b.to(torch::kFloat8_e4m3fn);
+
+  // Transpose B to column-major format
+  b_fp8 = b_fp8.t().contiguous().t();
+
+  // Create scales (per-tensor scaling)
+  torch::Tensor a_scales = torch::ones({1}, fp32_options);
+  torch::Tensor b_scales = torch::ones({1}, fp32_options);
+
+  // Create output tensor (must be Float16 for FP8 GEMM)
+  torch::Tensor c = torch::zeros({M, N}, fp16_options);
+
+  // Call cutlass_scaled_mm
+  ASSERT_NO_THROW({
+    xllm::kernel::cuda::cutlass_scaled_mm(
+        c, a_fp8, b_fp8, a_scales, b_scales, std::nullopt);
+  });
+
+  // Verify output shape
+  EXPECT_EQ(c.size(0), M);
+  EXPECT_EQ(c.size(1), N);
+
+  // Compute reference result using FP32
+  torch::Tensor c_ref = torch::matmul(a, b);
+
+  // Check if results are close (allowing for FP8 quantization error)
+  auto max_diff = (c.to(torch::kFloat32) - c_ref).abs().max().item<float>();
+  LOG(INFO) << "Max difference between FP8 and FP32 result: " << max_diff;
+
+  // FP8 has limited precision (4-bit exponent, 3-bit mantissa), so we use a
+  // loose tolerance
+  EXPECT_LT(max_diff, 2.0f);
+}
+
+// Test FP8 W8A8 with bias
+TEST_F(CutlassScaledMMTest, FP8W8A8WithBiasTest) {
+  const int64_t M = 64;
+  const int64_t N = 128;
+  const int64_t K = 256;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+
+  auto fp16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat16);
+
+  // Create input matrices
+  // Scale down the input range to avoid FP8 quantization error
+  torch::Tensor a = torch::randn({M, K}, fp32_options) * 0.5f;
+  torch::Tensor b = torch::randn({K, N}, fp32_options) * 0.5f;
+  torch::Tensor bias = torch::randn({N}, fp16_options) * 0.5f;
+
+  // Convert to FP8
+  torch::Tensor a_fp8 = a.to(torch::kFloat8_e4m3fn);
+  torch::Tensor b_fp8 = b.to(torch::kFloat8_e4m3fn);
+
+  // Transpose B to column-major format
+  b_fp8 = b_fp8.t().contiguous().t();
+
+  // Create scales
+  torch::Tensor a_scales = torch::ones({1}, fp32_options);
+  torch::Tensor b_scales = torch::ones({1}, fp32_options);
+
+  // Create output tensor (must be Float16 for FP8 GEMM)
+  torch::Tensor c = torch::zeros({M, N}, fp16_options);
+
+  // Call cutlass_scaled_mm with bias
+  ASSERT_NO_THROW({
+    xllm::kernel::cuda::cutlass_scaled_mm(
+        c, a_fp8, b_fp8, a_scales, b_scales, bias);
+  });
+
+  // Verify output shape
+  EXPECT_EQ(c.size(0), M);
+  EXPECT_EQ(c.size(1), N);
+
+  // Compute reference result
+  torch::Tensor c_ref =
+      torch::matmul(a, b) + bias.to(torch::kFloat32).unsqueeze(0);
+
+  // Check if results are close (allowing for FP8 quantization error)
+  auto max_diff = (c.to(torch::kFloat32) - c_ref).abs().max().item<float>();
+  LOG(INFO) << "Max difference with bias: " << max_diff;
+
+  // FP8 has limited precision, so we use a loose tolerance
+  EXPECT_LT(max_diff, 2.0f);
+}
+
+// Test FP8 W8A8 with per-token/per-channel scaling
+TEST_F(CutlassScaledMMTest, FP8W8A8WithScalingTest) {
+  const int64_t M = 64;
+  const int64_t N = 128;
+  const int64_t K = 256;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+
+  auto fp16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat16);
+
+  // Create input matrices
+  torch::Tensor a = torch::randn({M, K}, fp32_options) * 0.5f;
+  torch::Tensor b = torch::randn({K, N}, fp32_options) * 0.5f;
+
+  // Create scales for per-token/per-channel quantization
+  torch::Tensor a_scales = torch::rand({M}, fp32_options) * 0.1f + 0.9f;
+  torch::Tensor b_scales = torch::rand({N}, fp32_options) * 0.1f + 0.9f;
+
+  // Apply scaling and convert to FP8
+  torch::Tensor a_scaled = a / a_scales.unsqueeze(1);
+  torch::Tensor b_scaled = b / b_scales.unsqueeze(0);
+
+  torch::Tensor a_fp8 = a_scaled.to(torch::kFloat8_e4m3fn);
+  torch::Tensor b_fp8 = b_scaled.to(torch::kFloat8_e4m3fn);
+
+  // Transpose B to column-major format
+  b_fp8 = b_fp8.t().contiguous().t();
+
+  // Create output tensor (must be Float16 for FP8 GEMM)
+  torch::Tensor c = torch::zeros({M, N}, fp16_options);
+
+  // Call cutlass_scaled_mm
+  ASSERT_NO_THROW({
+    xllm::kernel::cuda::cutlass_scaled_mm(
+        c, a_fp8, b_fp8, a_scales, b_scales, std::nullopt);
+  });
+
+  // Verify output shape
+  EXPECT_EQ(c.size(0), M);
+  EXPECT_EQ(c.size(1), N);
+
+  // Compute reference result
+  torch::Tensor c_ref = torch::matmul(a, b);
+
+  // Check if results are close
+  auto max_diff = (c.to(torch::kFloat32) - c_ref).abs().max().item<float>();
+  auto mean_diff = (c.to(torch::kFloat32) - c_ref).abs().mean().item<float>();
+  LOG(INFO) << "Max difference with scaling: " << max_diff;
+  LOG(INFO) << "Mean difference with scaling: " << mean_diff;
+
+  EXPECT_LT(max_diff, 2.0f);
+  EXPECT_LT(mean_diff, 0.5f);
+}
+
+// Test different matrix sizes
+TEST_F(CutlassScaledMMTest, DifferentSizesTest) {
+  std::vector<std::tuple<int64_t, int64_t, int64_t>> test_sizes = {
+      {16, 32, 64},     // Small
+      {128, 128, 128},  // Square
+      {256, 512, 384},  // Medium
+      {512, 1024, 768}  // Large
+  };
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+
+  auto fp16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat16);
+
+  for (const auto& [M, N, K] : test_sizes) {
+    LOG(INFO) << "Testing size: M=" << M << ", N=" << N << ", K=" << K;
+
+    // Create input matrices
+    torch::Tensor a = torch::randn({M, K}, fp32_options);
+    torch::Tensor b = torch::randn({K, N}, fp32_options);
+
+    torch::Tensor a_fp8 = a.to(torch::kFloat8_e4m3fn);
+    torch::Tensor b_fp8 = b.to(torch::kFloat8_e4m3fn);
+
+    // Transpose B to column-major format
+    b_fp8 = b_fp8.t().contiguous().t();
+
+    torch::Tensor a_scales = torch::ones({1}, fp32_options);
+    torch::Tensor b_scales = torch::ones({1}, fp32_options);
+    // Create output tensor (must be Float16 for FP8 GEMM)
+    torch::Tensor c = torch::zeros({M, N}, fp16_options);
+
+    // Should not throw for valid sizes
+    ASSERT_NO_THROW({
+      xllm::kernel::cuda::cutlass_scaled_mm(
+          c, a_fp8, b_fp8, a_scales, b_scales, std::nullopt);
+    });
+
+    // Verify output shape
+    EXPECT_EQ(c.size(0), M);
+    EXPECT_EQ(c.size(1), N);
+  }
+}
+
+// Test error handling for invalid inputs
+TEST_F(CutlassScaledMMTest, InvalidInputTest) {
+  const int64_t M = 64;
+  const int64_t N = 128;
+  const int64_t K = 256;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+
+  auto fp16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat16);
+
+  torch::Tensor a_fp8 =
+      torch::randn({M, K}, fp32_options).to(torch::kFloat8_e4m3fn);
+  torch::Tensor b_fp8 =
+      torch::randn({K, N}, fp32_options).to(torch::kFloat8_e4m3fn);
+  b_fp8 = b_fp8.t().contiguous().t();
+
+  torch::Tensor a_scales = torch::ones({1}, fp32_options);
+  torch::Tensor b_scales = torch::ones({1}, fp32_options);
+  // Create output tensor (must be Float16 for FP8 GEMM)
+  torch::Tensor c = torch::zeros({M, N}, fp16_options);
+
+  // Test mismatched dimensions
+  torch::Tensor b_wrong =
+      torch::randn({K + 1, N}, fp32_options).to(torch::kFloat8_e4m3fn);
+  b_wrong = b_wrong.t().contiguous().t();
+
+  EXPECT_THROW(
+      {
+        xllm::kernel::cuda::cutlass_scaled_mm(
+            c, a_fp8, b_wrong, a_scales, b_scales, std::nullopt);
+      },
+      c10::Error);
+
+  // Test invalid bias size
+  torch::Tensor bias_wrong = torch::randn(
+      {N + 1}, torch::TensorOptions().device(device_).dtype(torch::kFloat16));
+
+  EXPECT_THROW(
+      {
+        xllm::kernel::cuda::cutlass_scaled_mm(
+            c, a_fp8, b_fp8, a_scales, b_scales, bias_wrong);
+      },
+      c10::Error);
+}

upstream_ref/xllm/tests/core/kernels/cuda/fp8_quant_test.cpp

@@ -0,0 +1,378 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+ *
+ * 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/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.
+ * ===========================================================================*/
+// clang-format off
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "cuda_ops_api.h"
+#include "cutlass_extensions/common.hpp"
+// clang-format on
+class StaticScaledFP8QuantTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    device_ = torch::Device(torch::kCUDA);
+
+    // Check if FP8 is supported (requires compute capability >= 8.9)
+    int compute_capability = xllm::kernel::cuda::get_sm_version_num();
+    if (compute_capability < 89) {
+      GTEST_SKIP() << "FP8 requires compute capability >= 8.9 (Ada Lovelace or "
+                      "Hopper), current: "
+                   << compute_capability;
+    }
+  }
+
+  torch::Device device_ = torch::kCPU;
+};
+
+// Test basic FP8 quantization with float32 input
+TEST_F(StaticScaledFP8QuantTest, BasicFloat32InputTest) {
+  const int64_t num_tokens = 128;
+  const int64_t hidden_size = 256;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create input tensor with values in a reasonable range for FP8
+  torch::Tensor input = torch::randn({num_tokens, hidden_size}, fp32_options);
+
+  // Create scale tensor (scale should be positive)
+  torch::Tensor scale = torch::tensor({1.0f}, fp32_options);
+
+  // Create output tensor
+  torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+  // Call static_scaled_fp8_quant
+  ASSERT_NO_THROW(
+      { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+  // Verify output shape
+  EXPECT_EQ(out.size(0), num_tokens);
+  EXPECT_EQ(out.size(1), hidden_size);
+
+  // Convert back to float and verify the quantization is reasonable
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+
+  // Check that output values are finite
+  EXPECT_TRUE(out_fp32.isfinite().all().item<bool>());
+
+  // Verify that quantization preserves relative ordering for most values
+  auto input_sign = input.sign();
+  auto out_sign = out_fp32.sign();
+  auto sign_match_ratio =
+      (input_sign == out_sign).to(torch::kFloat32).mean().item<float>();
+  LOG(INFO) << "Sign match ratio: " << sign_match_ratio;
+  EXPECT_GT(sign_match_ratio, 0.9f);
+}
+
+// Test FP8 quantization with float16 input
+TEST_F(StaticScaledFP8QuantTest, Float16InputTest) {
+  const int64_t num_tokens = 64;
+  const int64_t hidden_size = 512;
+
+  auto fp16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat16);
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create float16 input tensor
+  torch::Tensor input = torch::randn({num_tokens, hidden_size}, fp16_options);
+
+  // Create scale tensor
+  torch::Tensor scale = torch::tensor({0.5f}, fp32_options);
+
+  // Create output tensor
+  torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+  // Call static_scaled_fp8_quant
+  ASSERT_NO_THROW(
+      { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+  // Verify output shape
+  EXPECT_EQ(out.size(0), num_tokens);
+  EXPECT_EQ(out.size(1), hidden_size);
+
+  // Verify output is finite
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+  EXPECT_TRUE(out_fp32.isfinite().all().item<bool>());
+}
+
+// Test FP8 quantization with bfloat16 input
+TEST_F(StaticScaledFP8QuantTest, BFloat16InputTest) {
+  const int64_t num_tokens = 32;
+  const int64_t hidden_size = 128;
+
+  auto bf16_options =
+      torch::TensorOptions().device(device_).dtype(torch::kBFloat16);
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create bfloat16 input tensor
+  torch::Tensor input = torch::randn({num_tokens, hidden_size}, bf16_options);
+
+  // Create scale tensor
+  torch::Tensor scale = torch::tensor({2.0f}, fp32_options);
+
+  // Create output tensor
+  torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+  // Call static_scaled_fp8_quant
+  ASSERT_NO_THROW(
+      { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+  // Verify output shape
+  EXPECT_EQ(out.size(0), num_tokens);
+  EXPECT_EQ(out.size(1), hidden_size);
+
+  // Verify output is finite
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+  EXPECT_TRUE(out_fp32.isfinite().all().item<bool>());
+}
+
+// Test FP8 quantization with different scales
+TEST_F(StaticScaledFP8QuantTest, DifferentScalesTest) {
+  const int64_t num_tokens = 64;
+  const int64_t hidden_size = 256;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create input tensor with small values
+  torch::Tensor input =
+      torch::randn({num_tokens, hidden_size}, fp32_options) * 0.1f;
+
+  std::vector<float> test_scales = {0.1f, 0.5f, 1.0f, 2.0f, 10.0f};
+
+  for (float scale_val : test_scales) {
+    LOG(INFO) << "Testing with scale: " << scale_val;
+
+    torch::Tensor scale = torch::tensor({scale_val}, fp32_options);
+    torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+    ASSERT_NO_THROW(
+        { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+    // Verify output shape
+    EXPECT_EQ(out.size(0), num_tokens);
+    EXPECT_EQ(out.size(1), hidden_size);
+
+    // Verify output is finite
+    torch::Tensor out_fp32 = out.to(torch::kFloat32);
+    EXPECT_TRUE(out_fp32.isfinite().all().item<bool>());
+  }
+}
+
+// Test FP8 quantization with various tensor sizes
+TEST_F(StaticScaledFP8QuantTest, DifferentSizesTest) {
+  std::vector<std::pair<int64_t, int64_t>> test_sizes = {
+      {1, 64},      // Single token
+      {16, 128},    // Small batch
+      {64, 256},    // Medium batch
+      {128, 512},   // Large batch
+      {256, 1024},  // Very large batch
+      {512, 4096},  // Large hidden size
+  };
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  for (const auto& [num_tokens, hidden_size] : test_sizes) {
+    LOG(INFO) << "Testing size: num_tokens=" << num_tokens
+              << ", hidden_size=" << hidden_size;
+
+    torch::Tensor input = torch::randn({num_tokens, hidden_size}, fp32_options);
+    torch::Tensor scale = torch::tensor({1.0f}, fp32_options);
+    torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+    ASSERT_NO_THROW(
+        { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+    EXPECT_EQ(out.size(0), num_tokens);
+    EXPECT_EQ(out.size(1), hidden_size);
+  }
+}
+
+// Test FP8 quantization with 3D tensor (batched)
+TEST_F(StaticScaledFP8QuantTest, BatchedTensor3DTest) {
+  const int64_t batch_size = 4;
+  const int64_t seq_len = 32;
+  const int64_t hidden_size = 128;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create 3D input tensor [batch_size, seq_len, hidden_size]
+  torch::Tensor input =
+      torch::randn({batch_size, seq_len, hidden_size}, fp32_options);
+
+  torch::Tensor scale = torch::tensor({1.0f}, fp32_options);
+  torch::Tensor out =
+      torch::empty({batch_size, seq_len, hidden_size}, fp8_options);
+
+  // Call static_scaled_fp8_quant
+  ASSERT_NO_THROW(
+      { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+  // Verify output shape
+  EXPECT_EQ(out.size(0), batch_size);
+  EXPECT_EQ(out.size(1), seq_len);
+  EXPECT_EQ(out.size(2), hidden_size);
+
+  // Verify output is finite
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+  EXPECT_TRUE(out_fp32.isfinite().all().item<bool>());
+}
+
+// Test quantization accuracy with known values
+TEST_F(StaticScaledFP8QuantTest, QuantizationAccuracyTest) {
+  const int64_t num_tokens = 64;
+  const int64_t hidden_size = 128;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create input tensor with values scaled to FP8 range
+  // FP8 e4m3 has max value around 448
+  torch::Tensor input =
+      torch::randn({num_tokens, hidden_size}, fp32_options) * 0.5f;
+
+  torch::Tensor scale = torch::tensor({1.0f}, fp32_options);
+  torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+  xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale);
+
+  // Convert back to float for comparison
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+
+  // Calculate quantization error
+  // Note: FP8 quantization has limited precision (4-bit exponent, 3-bit
+  // mantissa) so we expect some error
+  auto abs_error = (out_fp32 - input).abs();
+  auto max_error = abs_error.max().item<float>();
+  auto mean_error = abs_error.mean().item<float>();
+
+  LOG(INFO) << "Max quantization error: " << max_error;
+  LOG(INFO) << "Mean quantization error: " << mean_error;
+
+  // FP8 has limited precision, but error should be bounded
+  // For values around 0.5, the error should be relatively small
+  EXPECT_LT(mean_error, 0.5f);
+}
+
+// Test with scale that compensates for large input values
+TEST_F(StaticScaledFP8QuantTest, LargeInputWithScaleTest) {
+  const int64_t num_tokens = 32;
+  const int64_t hidden_size = 64;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create input tensor with larger values
+  torch::Tensor input =
+      torch::randn({num_tokens, hidden_size}, fp32_options) * 10.0f;
+
+  // Use a larger scale to bring values into FP8 range
+  // The kernel divides by scale, so larger scale compresses the range
+  float input_max = input.abs().max().item<float>();
+  float scale_val = input_max / 448.0f;  // FP8 e4m3 max is ~448
+  if (scale_val < 1.0f) scale_val = 1.0f;
+
+  torch::Tensor scale = torch::tensor({scale_val}, fp32_options);
+  torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+  ASSERT_NO_THROW(
+      { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+  // Verify output is finite
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+  EXPECT_TRUE(out_fp32.isfinite().all().item<bool>());
+
+  // Verify scaled output is within FP8 representable range
+  float out_max = out_fp32.abs().max().item<float>();
+  LOG(INFO) << "Input max: " << input_max << ", Scale: " << scale_val
+            << ", Output max: " << out_max;
+  EXPECT_LE(out_max, 450.0f);  // Allow some tolerance
+}
+
+// Test that zero values remain zero after quantization
+TEST_F(StaticScaledFP8QuantTest, ZeroValuesTest) {
+  const int64_t num_tokens = 16;
+  const int64_t hidden_size = 32;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create input tensor with zeros
+  torch::Tensor input = torch::zeros({num_tokens, hidden_size}, fp32_options);
+
+  torch::Tensor scale = torch::tensor({1.0f}, fp32_options);
+  torch::Tensor out = torch::empty({num_tokens, hidden_size}, fp8_options);
+
+  xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale);
+
+  // Verify all outputs are zero
+  torch::Tensor out_fp32 = out.to(torch::kFloat32);
+  EXPECT_TRUE((out_fp32 == 0).all().item<bool>());
+}
+
+// Test contiguity requirements
+TEST_F(StaticScaledFP8QuantTest, ContiguityTest) {
+  const int64_t num_tokens = 32;
+  const int64_t hidden_size = 64;
+
+  auto fp32_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto fp8_options =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat8_e4m3fn);
+
+  // Create contiguous input tensor
+  torch::Tensor input =
+      torch::randn({num_tokens, hidden_size}, fp32_options).contiguous();
+
+  // Ensure input is contiguous in the last dimension
+  EXPECT_EQ(input.stride(-1), 1);
+
+  torch::Tensor scale = torch::tensor({1.0f}, fp32_options);
+  torch::Tensor out =
+      torch::empty({num_tokens, hidden_size}, fp8_options).contiguous();
+
+  ASSERT_NO_THROW(
+      { xllm::kernel::cuda::static_scaled_fp8_quant(out, input, scale); });
+
+  // Verify output
+  EXPECT_EQ(out.size(0), num_tokens);
+  EXPECT_EQ(out.size(1), hidden_size);
+}

upstream_ref/xllm/tests/core/kernels/cuda/fused_qknorm_rope_test.cpp

@@ -0,0 +1,233 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/cuda.h>
+#include <torch/torch.h>
+
+#include "core/kernels/cuda/cuda_ops_api.h"
+
+namespace xllm::kernel::cuda {
+namespace test {
+
+namespace {
+
+torch::Tensor apply_reference_rope(const torch::Tensor& input,
+                                   const torch::Tensor& cos,
+                                   const torch::Tensor& sin,
+                                   int64_t rotary_dim,
+                                   bool interleaved) {
+  auto input_float = input.to(torch::kFloat32);
+  auto input_rot = input_float.slice(-1, 0, rotary_dim);
+  auto input_pass = input_float.slice(-1, rotary_dim, input_float.size(-1));
+
+  const int64_t num_tokens = input_float.size(0);
+  const int64_t num_heads = input_float.size(1);
+  const int64_t half_dim = rotary_dim / 2;
+  auto cos_view = cos.view({num_tokens, 1, half_dim});
+  auto sin_view = sin.view({num_tokens, 1, half_dim});
+
+  torch::Tensor rotated;
+  if (interleaved) {
+    auto reshaped = input_rot.view({num_tokens, num_heads, half_dim, 2});
+    auto even = reshaped.select(-1, 0);
+    auto odd = reshaped.select(-1, 1);
+    auto out_even = even * cos_view - odd * sin_view;
+    auto out_odd = even * sin_view + odd * cos_view;
+    rotated = torch::stack({out_even, out_odd}, -1)
+                  .view({num_tokens, num_heads, rotary_dim});
+  } else {
+    auto first = input_rot.slice(-1, 0, half_dim);
+    auto second = input_rot.slice(-1, half_dim, rotary_dim);
+    rotated = torch::cat({first * cos_view - second * sin_view,
+                          second * cos_view + first * sin_view},
+                         -1);
+  }
+
+  if (rotary_dim < input_float.size(-1)) {
+    return torch::cat({rotated, input_pass}, -1).to(input.scalar_type());
+  }
+  return rotated.to(input.scalar_type());
+}
+
+void fused_qk_norm_rope_reference(torch::Tensor& qkv,
+                                  int64_t num_heads_q,
+                                  int64_t num_heads_k,
+                                  int64_t num_heads_v,
+                                  int64_t head_dim,
+                                  double eps,
+                                  const torch::Tensor& q_weight,
+                                  const torch::Tensor& k_weight,
+                                  const torch::Tensor& cos_sin_cache,
+                                  bool interleaved,
+                                  const torch::Tensor& position_ids) {
+  (void)num_heads_v;
+  const int64_t num_tokens = qkv.size(0);
+  const int64_t q_size = num_heads_q * head_dim;
+  const int64_t k_size = num_heads_k * head_dim;
+  const int64_t rotary_dim = cos_sin_cache.size(1);
+  const int64_t half_dim = rotary_dim / 2;
+
+  auto q_slice =
+      qkv.slice(-1, 0, q_size).view({num_tokens, num_heads_q, head_dim});
+  auto k_slice = qkv.slice(-1, q_size, q_size + k_size)
+                     .view({num_tokens, num_heads_k, head_dim});
+
+  auto q_float = q_slice.to(torch::kFloat32);
+  auto k_float = k_slice.to(torch::kFloat32);
+  auto q_weight_float = q_weight.to(torch::kFloat32).view({1, 1, head_dim});
+  auto k_weight_float = k_weight.to(torch::kFloat32).view({1, 1, head_dim});
+
+  auto q_rms = torch::rsqrt((q_float * q_float).mean(-1, true) + eps);
+  auto k_rms = torch::rsqrt((k_float * k_float).mean(-1, true) + eps);
+  auto q = (q_float * q_rms * q_weight_float).to(q_slice.scalar_type());
+  auto k = (k_float * k_rms * k_weight_float).to(k_slice.scalar_type());
+
+  auto selected_cos_sin = cos_sin_cache.index_select(0, position_ids);
+  auto cos = selected_cos_sin.slice(-1, 0, half_dim);
+  auto sin = selected_cos_sin.slice(-1, half_dim, rotary_dim);
+
+  q = apply_reference_rope(q, cos, sin, rotary_dim, interleaved);
+  k = apply_reference_rope(k, cos, sin, rotary_dim, interleaved);
+  q_slice.copy_(q);
+  k_slice.copy_(k);
+}
+
+class FusedQKNormRopeTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    torch::manual_seed(2026);
+    device_ = torch::Device(torch::kCUDA, 0);
+  }
+
+  torch::Device device_ = torch::Device(torch::kCPU);
+};
+
+TEST_F(FusedQKNormRopeTest, MatchesReferenceNeoX) {
+  const int64_t num_tokens = 17;
+  const int64_t num_heads_q = 8;
+  const int64_t num_heads_k = 4;
+  const int64_t num_heads_v = 4;
+  const int64_t head_dim = 128;
+  const int64_t rotary_dim = 128;
+  const int64_t max_position = 512;
+  const double eps = 1e-6;
+
+  auto half_opts = torch::TensorOptions().device(device_).dtype(torch::kHalf);
+  auto float_opts =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto int_opts = torch::TensorOptions().device(device_).dtype(torch::kInt64);
+
+  auto qkv =
+      torch::randn(
+          {num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_dim},
+          half_opts) *
+      0.2;
+  auto q_weight = torch::randn({head_dim}, half_opts);
+  auto k_weight = torch::randn({head_dim}, half_opts);
+  auto cos_sin_cache = torch::randn({max_position, rotary_dim}, float_opts);
+  auto position_ids = torch::randint(0, max_position, {num_tokens}, int_opts);
+
+  auto qkv_ref = qkv.clone();
+  fused_qk_norm_rope_reference(qkv_ref,
+                               num_heads_q,
+                               num_heads_k,
+                               num_heads_v,
+                               head_dim,
+                               eps,
+                               q_weight,
+                               k_weight,
+                               cos_sin_cache,
+                               false,
+                               position_ids);
+
+  auto qkv_out = qkv.clone();
+  fused_qk_norm_rope(qkv_out,
+                     num_heads_q,
+                     num_heads_k,
+                     num_heads_v,
+                     head_dim,
+                     eps,
+                     q_weight,
+                     k_weight,
+                     cos_sin_cache,
+                     false,
+                     position_ids);
+
+  EXPECT_TRUE(torch::allclose(qkv_out, qkv_ref, 2e-3, 2e-3));
+}
+
+TEST_F(FusedQKNormRopeTest, MatchesReferenceInterleaved) {
+  const int64_t num_tokens = 11;
+  const int64_t num_heads_q = 6;
+  const int64_t num_heads_k = 2;
+  const int64_t num_heads_v = 2;
+  const int64_t head_dim = 64;
+  const int64_t rotary_dim = 64;
+  const int64_t max_position = 256;
+  const double eps = 1e-6;
+
+  auto bf16_opts =
+      torch::TensorOptions().device(device_).dtype(torch::kBFloat16);
+  auto float_opts =
+      torch::TensorOptions().device(device_).dtype(torch::kFloat32);
+  auto int_opts = torch::TensorOptions().device(device_).dtype(torch::kInt64);
+
+  auto qkv =
+      torch::randn(
+          {num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_dim},
+          bf16_opts) *
+      0.15;
+  auto q_weight = torch::randn({head_dim}, bf16_opts);
+  auto k_weight = torch::randn({head_dim}, bf16_opts);
+  auto cos_sin_cache = torch::randn({max_position, rotary_dim}, float_opts);
+  auto position_ids = torch::randint(0, max_position, {num_tokens}, int_opts);
+
+  auto qkv_ref = qkv.clone();
+  fused_qk_norm_rope_reference(qkv_ref,
+                               num_heads_q,
+                               num_heads_k,
+                               num_heads_v,
+                               head_dim,
+                               eps,
+                               q_weight,
+                               k_weight,
+                               cos_sin_cache,
+                               true,
+                               position_ids);
+
+  auto qkv_out = qkv.clone();
+  fused_qk_norm_rope(qkv_out,
+                     num_heads_q,
+                     num_heads_k,
+                     num_heads_v,
+                     head_dim,
+                     eps,
+                     q_weight,
+                     k_weight,
+                     cos_sin_cache,
+                     true,
+                     position_ids);
+
+  EXPECT_TRUE(torch::allclose(qkv_out, qkv_ref, 2e-2, 2e-2));
+}
+
+}  // namespace
+
+}  // namespace test
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/tests/core/kernels/cuda/moe/CMakeLists.txt

@@ -0,0 +1,11 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    moe_fused_topk_test
+  SRCS
+    moe_topk_test.cu
+  DEPS
+    :cuda_kernels
+    GTest::gtest_main
+)

upstream_ref/xllm/tests/core/kernels/cuda/moe/moe_topk_test.cu

@@ -0,0 +1,897 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <random>
+#include <vector>
+
+#include "moe_topk.cuh"
+
+namespace xllm::kernel::cuda {
+namespace test {
+
+using namespace reduce_topk;
+
+// ============================================================================
+// Helper: CPU reference top-K (sort descending by value, ascending by index on
+// tie, then take the first K).
+// ============================================================================
+template <typename T>
+void cpuTopK(const std::vector<T>& values,
+             const std::vector<int32_t>& indices,
+             int k,
+             std::vector<T>& outValues,
+             std::vector<int32_t>& outIndices) {
+  struct Pair {
+    T val;
+    int32_t idx;
+  };
+  std::vector<Pair> pairs(values.size());
+  for (size_t i = 0; i < values.size(); ++i) {
+    pairs[i] = {values[i], indices[i]};
+  }
+  std::sort(pairs.begin(), pairs.end(), [](const Pair& a, const Pair& b) {
+    if (a.val != b.val) return a.val > b.val;
+    return a.idx < b.idx;
+  });
+  outValues.resize(k);
+  outIndices.resize(k);
+  for (int i = 0; i < k; ++i) {
+    outValues[i] = pairs[i].val;
+    outIndices[i] = pairs[i].idx;
+  }
+}
+
+// ============================================================================
+// 1. Host-side pack / unpack roundtrip tests
+//    makeCmpVal and unpack are __host__ __device__, so we can call them from
+//    the CPU when compiling with nvcc.
+// ============================================================================
+
+class TopKRedTypePackUnpackTest : public ::testing::Test {};
+
+TEST_F(TopKRedTypePackUnpackTest, FloatRoundtrip) {
+  using RedType = TopKRedType<float>;
+  const float testValues[] = {
+      0.0f, 1.0f, -1.0f, 3.14159f, -100.5f, 1e10f, -1e10f, 0.001f};
+  const int32_t testIndices[] = {0, 1, 31, 100, 1000, 65534, 65535};
+
+  for (float val : testValues) {
+    for (int32_t idx : testIndices) {
+      auto packed = RedType::makeCmpVal(val, idx);
+      float unpackedVal;
+      int32_t unpackedIdx;
+      RedType::unpack(unpackedVal, unpackedIdx, packed);
+      EXPECT_EQ(unpackedVal, val)
+          << "Value mismatch for val=" << val << " idx=" << idx;
+      EXPECT_EQ(unpackedIdx, idx)
+          << "Index mismatch for val=" << val << " idx=" << idx;
+    }
+  }
+}
+
+TEST_F(TopKRedTypePackUnpackTest, IntRoundtrip) {
+  using RedType = TopKRedType<int>;
+  const int testValues[] = {0, 1, -1, 42, -100, INT32_MAX, INT32_MIN};
+  const int32_t testIndices[] = {0, 1, 100, 65534, 65535};
+
+  for (int val : testValues) {
+    for (int32_t idx : testIndices) {
+      auto packed = RedType::makeCmpVal(val, idx);
+      int unpackedVal;
+      int32_t unpackedIdx;
+      RedType::unpack(unpackedVal, unpackedIdx, packed);
+      EXPECT_EQ(unpackedVal, val)
+          << "Value mismatch for val=" << val << " idx=" << idx;
+      EXPECT_EQ(unpackedIdx, idx)
+          << "Index mismatch for val=" << val << " idx=" << idx;
+    }
+  }
+}
+
+// Regression: negative index -1 must NOT survive a pack/unpack roundtrip.
+// kMaxIdx - (-1) = 65536, truncated to 0 in 16 bits, unpacks to 65535.
+TEST_F(TopKRedTypePackUnpackTest, NegativeIndexDoesNotRoundtrip) {
+  using RedType = TopKRedType<float>;
+  const int32_t badIdx = -1;
+  auto packed = RedType::makeCmpVal(1.0f, badIdx);
+  float v;
+  int32_t idx;
+  RedType::unpack(v, idx, packed);
+  EXPECT_NE(idx, badIdx) << "Negative index should NOT roundtrip correctly";
+  EXPECT_EQ(idx, 65535) << "Expected 65535 after failed roundtrip of -1";
+}
+
+// kMaxIdx (65535) as sentinel DOES roundtrip correctly.
+TEST_F(TopKRedTypePackUnpackTest, SentinelKMaxIdxRoundtrip) {
+  using RedType = TopKRedType<float>;
+  const int32_t sentinel = RedType::kMaxIdx;  // 65535
+  const float minVal = -std::numeric_limits<float>::infinity();
+  auto packed = RedType::makeCmpVal(minVal, sentinel);
+  float v;
+  int32_t idx;
+  RedType::unpack(v, idx, packed);
+  EXPECT_EQ(idx, sentinel);
+  // -inf should also survive (TwiddleIn/Out handles it).
+  EXPECT_EQ(v, minVal);
+}
+
+// Zero index roundtrip.
+TEST_F(TopKRedTypePackUnpackTest, ZeroIndexRoundtrip) {
+  using RedType = TopKRedType<float>;
+  auto packed = RedType::makeCmpVal(2.5f, 0);
+  float v;
+  int32_t idx;
+  RedType::unpack(v, idx, packed);
+  EXPECT_EQ(idx, 0);
+  EXPECT_EQ(v, 2.5f);
+}
+
+// ============================================================================
+// 2. Host-side comparison ordering tests
+// ============================================================================
+
+class TopKRedTypeOrderingTest : public ::testing::Test {};
+
+// Larger value → larger packed representation (same index).
+TEST_F(TopKRedTypeOrderingTest, LargerValueHigherPriority) {
+  using RedType = TopKRedType<float>;
+  const int32_t idx = 42;
+  auto p1 = RedType::makeCmpVal(10.0f, idx);
+  auto p2 = RedType::makeCmpVal(5.0f, idx);
+  EXPECT_GT(p1, p2) << "Larger value should produce larger packed value";
+}
+
+// For equal values, smaller index → larger packed representation (higher
+// priority).
+TEST_F(TopKRedTypeOrderingTest, SmallerIndexHigherPriority) {
+  using RedType = TopKRedType<float>;
+  auto p1 = RedType::makeCmpVal(7.0f, 10);
+  auto p2 = RedType::makeCmpVal(7.0f, 20);
+  EXPECT_GT(p1, p2) << "Smaller index should have higher priority";
+}
+
+// Negative value ordering: -1 > -10.
+TEST_F(TopKRedTypeOrderingTest, NegativeValueOrdering) {
+  using RedType = TopKRedType<float>;
+  auto p1 = RedType::makeCmpVal(-1.0f, 0);
+  auto p2 = RedType::makeCmpVal(-10.0f, 0);
+  EXPECT_GT(p1, p2) << "-1 should rank higher than -10";
+}
+
+// Integer ordering.
+TEST_F(TopKRedTypeOrderingTest, IntOrdering) {
+  using RedType = TopKRedType<int>;
+  auto p1 = RedType::makeCmpVal(100, 0);
+  auto p2 = RedType::makeCmpVal(50, 0);
+  EXPECT_GT(p1, p2);
+  auto p3 = RedType::makeCmpVal(-1, 0);
+  auto p4 = RedType::makeCmpVal(-100, 0);
+  EXPECT_GT(p3, p4);
+}
+
+// Sentinel (minValue, kMaxIdx) should be the smallest packed value among
+// any real candidates, ensuring it always loses in a max-reduction.
+TEST_F(TopKRedTypeOrderingTest, SentinelIsSmallest) {
+  using RedType = TopKRedType<float>;
+  auto sentinel = RedType::makeCmpVal(-std::numeric_limits<float>::infinity(),
+                                      RedType::kMaxIdx);
+  auto real = RedType::makeCmpVal(0.0f, 0);
+  EXPECT_GT(real, sentinel);
+}
+
+// Monotone sweep: ascending values with fixed index should produce ascending
+// packed values.
+TEST_F(TopKRedTypeOrderingTest, MonotoneSweepFloat) {
+  using RedType = TopKRedType<float>;
+  typename RedType::TypeCmp prev = RedType::makeCmpVal(-1000.0f, 0);
+  for (float v = -999.0f; v <= 1000.0f; v += 1.0f) {
+    auto cur = RedType::makeCmpVal(v, 0);
+    EXPECT_GT(cur, prev);
+    prev = cur;
+  }
+}
+
+// ============================================================================
+// 3. Device-side kernel wrappers
+// ============================================================================
+
+// Kernel: each of 32 warp lanes holds one (value, index) pair.
+// Performs warp-level top-K reduction; lane 0 writes K results.
+template <int K, typename Type>
+__global__ void testReduceTopKSingleKernel(const Type* __restrict__ values,
+                                           const int32_t* __restrict__ indices,
+                                           Type* __restrict__ outValues,
+                                           int32_t* __restrict__ outIndices,
+                                           Type minValue,
+                                           int actualK) {
+  auto warp = cg::tiled_partition<kWARP_SIZE>(cg::this_thread_block());
+  const int lane = threadIdx.x;
+
+  Type val = values[lane];
+  int32_t idx = indices[lane];
+
+  Type out[K];
+  int32_t outIdx[K];
+  reduceTopK<K>(warp, out, outIdx, val, idx, minValue, actualK);
+
+  if (lane == 0) {
+    for (int i = 0; i < actualK; ++i) {
+      outValues[i] = out[i];
+      outIndices[i] = outIdx[i];
+    }
+  }
+}
+
+// Kernel: each of 32 warp lanes holds N (value, index) pairs.
+// Input layout: values[lane * N + n], indices[lane * N + n].
+template <int K, typename Type, int N>
+__global__ void testReduceTopKMultiKernel(const Type* __restrict__ values,
+                                          const int32_t* __restrict__ indices,
+                                          Type* __restrict__ outValues,
+                                          int32_t* __restrict__ outIndices,
+                                          Type minValue,
+                                          int actualK) {
+  auto warp = cg::tiled_partition<kWARP_SIZE>(cg::this_thread_block());
+  const int lane = threadIdx.x;
+
+  Type val[N];
+  int32_t idx[N];
+  for (int n = 0; n < N; ++n) {
+    val[n] = values[lane * N + n];
+    idx[n] = indices[lane * N + n];
+  }
+
+  Type out[K];
+  int32_t outIdx[K];
+  reduceTopK<K, Type, N>(warp, out, outIdx, val, idx, minValue, actualK);
+
+  if (lane == 0) {
+    for (int i = 0; i < actualK; ++i) {
+      outValues[i] = out[i];
+      outIndices[i] = outIdx[i];
+    }
+  }
+}
+
+// ============================================================================
+// RAII wrapper for device memory
+// ============================================================================
+template <typename T>
+struct DevBuf {
+  T* ptr = nullptr;
+  size_t count = 0;
+
+  explicit DevBuf(size_t n) : count(n) { cudaMalloc(&ptr, n * sizeof(T)); }
+  ~DevBuf() {
+    if (ptr) cudaFree(ptr);
+  }
+  void upload(const T* host) {
+    cudaMemcpy(ptr, host, count * sizeof(T), cudaMemcpyHostToDevice);
+  }
+  void download(T* host) const {
+    cudaMemcpy(host, ptr, count * sizeof(T), cudaMemcpyDeviceToHost);
+  }
+  // non-copyable
+  DevBuf(const DevBuf&) = delete;
+  DevBuf& operator=(const DevBuf&) = delete;
+};
+
+// ============================================================================
+// 4. Device-side integration tests
+// ============================================================================
+
+class ReduceTopKDeviceTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    int deviceCount = 0;
+    cudaGetDeviceCount(&deviceCount);
+    if (deviceCount == 0) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+  }
+};
+
+// ---------- Single-value reduceTopK tests ----------
+
+// K=1: find the global maximum across 32 warp lanes.
+TEST_F(ReduceTopKDeviceTest, SingleValueTopK1) {
+  constexpr int K = 1;
+  constexpr int N = kWARP_SIZE;
+
+  std::vector<float> h_vals(N);
+  std::vector<int32_t> h_idx(N);
+  for (int i = 0; i < N; ++i) {
+    h_vals[i] = static_cast<float>(i * 3 - 40);  // [-40 .. 53]
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(N), d_outV(K);
+  DevBuf<int32_t> d_idx(N), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKSingleKernel<K>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// K=4: find the top-4 across 32 warp lanes with pseudo-random data.
+TEST_F(ReduceTopKDeviceTest, SingleValueTopK4) {
+  constexpr int K = 4;
+  constexpr int N = kWARP_SIZE;
+
+  std::mt19937 rng(42);
+  std::vector<float> h_vals(N);
+  std::vector<int32_t> h_idx(N);
+  for (int i = 0; i < N; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 1000) / 10.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(N), d_outV(K);
+  DevBuf<int32_t> d_idx(N), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKSingleKernel<K>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// All values identical → top-K should return the K smallest indices.
+TEST_F(ReduceTopKDeviceTest, DuplicateValuesPreferSmallerIndex) {
+  constexpr int K = 3;
+  constexpr int N = kWARP_SIZE;
+
+  std::vector<float> h_vals(N, 42.0f);
+  std::vector<int32_t> h_idx(N);
+  for (int i = 0; i < N; ++i) h_idx[i] = i;
+
+  DevBuf<float> d_vals(N), d_outV(K);
+  DevBuf<int32_t> d_idx(N), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKSingleKernel<K>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  int32_t outI[K];
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_EQ(outI[i], i) << "Should prefer smallest indices when values tie";
+  }
+}
+
+// actualK < K: only the first actualK results are meaningful.
+TEST_F(ReduceTopKDeviceTest, SingleValueActualKLessThanK) {
+  constexpr int K = 4;
+  constexpr int N = kWARP_SIZE;
+  const int actualK = 2;
+
+  std::mt19937 rng(99);
+  std::vector<float> h_vals(N);
+  std::vector<int32_t> h_idx(N);
+  for (int i = 0; i < N; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 500) / 5.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, actualK, refV, refI);
+
+  DevBuf<float> d_vals(N), d_outV(K);
+  DevBuf<int32_t> d_idx(N), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKSingleKernel<K>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  actualK);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < actualK; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// Negative values: top-K should still pick the largest (closest to 0).
+TEST_F(ReduceTopKDeviceTest, SingleValueAllNegative) {
+  constexpr int K = 2;
+  constexpr int N = kWARP_SIZE;
+
+  std::vector<float> h_vals(N);
+  std::vector<int32_t> h_idx(N);
+  for (int i = 0; i < N; ++i) {
+    h_vals[i] = -static_cast<float>(i + 1);  // [-1, -2, ..., -32]
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(N), d_outV(K);
+  DevBuf<int32_t> d_idx(N), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKSingleKernel<K>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  // Top-2 should be -1 (idx 0) and -2 (idx 1).
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// ---------- Multi-value reduceTopK tests (N <= 4 path) ----------
+
+// N=2 per thread, K=2 → total 64 candidates.
+TEST_F(ReduceTopKDeviceTest, MultiValueN2K2) {
+  constexpr int K = 2;
+  constexpr int N_PER_THREAD = 2;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(123);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 2000) / 10.0f - 100.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// N=3 per thread, K=2 → total 96 candidates.
+TEST_F(ReduceTopKDeviceTest, MultiValueN3K2) {
+  constexpr int K = 2;
+  constexpr int N_PER_THREAD = 3;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(321);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 3000) / 10.0f - 150.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// N=4 per thread, K=3 → total 128 candidates (boundary of N<=4 path).
+TEST_F(ReduceTopKDeviceTest, MultiValueN4K3) {
+  constexpr int K = 3;
+  constexpr int N_PER_THREAD = 4;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(456);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 5000) / 10.0f - 250.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// ---------- Multi-value reduceTopK tests (N > 4 path) ----------
+// When N > 4 the implementation splits into numLoops rounds of 4 values each,
+// accumulates intermediate results into a buffer, then does a final reduction.
+// This exercises the code path that previously had the -1 sentinel bug.
+
+// N=8 per thread, K=2 → total 256 candidates.
+TEST_F(ReduceTopKDeviceTest, MultiValueN8K2_LargePath) {
+  constexpr int K = 2;
+  constexpr int N_PER_THREAD = 8;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(789);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 10000) / 10.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// N=8 per thread, K=4 → total 256, wider K.
+TEST_F(ReduceTopKDeviceTest, MultiValueN8K4_LargePath) {
+  constexpr int K = 4;
+  constexpr int N_PER_THREAD = 8;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(1024);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 8000) / 10.0f - 400.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// N=12 per thread, K=3 → total 384, numLoops=3.
+TEST_F(ReduceTopKDeviceTest, MultiValueN12K3_LargePath) {
+  constexpr int K = 3;
+  constexpr int N_PER_THREAD = 12;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(2048);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 6000) / 10.0f - 300.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// N=16 per thread, K=2 → total 512, maximum supported N.
+TEST_F(ReduceTopKDeviceTest, MultiValueN16K2_LargePath) {
+  constexpr int K = 2;
+  constexpr int N_PER_THREAD = 16;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::mt19937 rng(4096);
+  std::vector<float> h_vals(TOTAL);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) {
+    h_vals[i] = static_cast<float>(rng() % 20000) / 100.0f - 100.0f;
+    h_idx[i] = i;
+  }
+
+  std::vector<float> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+// ---------- Edge case: N>4 path with all-duplicate values ----------
+// Exercises the sentinel buffer path when no single candidate stands out.
+TEST_F(ReduceTopKDeviceTest, MultiValueN8K3_AllDuplicate) {
+  constexpr int K = 3;
+  constexpr int N_PER_THREAD = 8;
+  constexpr int TOTAL = kWARP_SIZE * N_PER_THREAD;
+
+  std::vector<float> h_vals(TOTAL, 7.7f);
+  std::vector<int32_t> h_idx(TOTAL);
+  for (int i = 0; i < TOTAL; ++i) h_idx[i] = i;
+
+  DevBuf<float> d_vals(TOTAL), d_outV(K);
+  DevBuf<int32_t> d_idx(TOTAL), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKMultiKernel<K, float, N_PER_THREAD>
+      <<<1, 32>>>(d_vals.ptr,
+                  d_idx.ptr,
+                  d_outV.ptr,
+                  d_outI.ptr,
+                  -std::numeric_limits<float>::infinity(),
+                  K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  float outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  // All values are the same, so top-K indices should be the K smallest.
+  for (int i = 0; i < K; ++i) {
+    EXPECT_FLOAT_EQ(outV[i], 7.7f);
+    EXPECT_EQ(outI[i], i)
+        << "Duplicate-value top-K should prefer smallest indices";
+  }
+}
+
+// ---------- Integer type on device ----------
+
+// Single-value reduceTopK with int type, K=2.
+TEST_F(ReduceTopKDeviceTest, IntTypeSingleValueTopK2) {
+  constexpr int K = 2;
+  constexpr int N = kWARP_SIZE;
+
+  std::vector<int> h_vals(N);
+  std::vector<int32_t> h_idx(N);
+  for (int i = 0; i < N; ++i) {
+    h_vals[i] = (i * 7 + 13) % 100 - 50;
+    h_idx[i] = i;
+  }
+
+  std::vector<int> refV;
+  std::vector<int32_t> refI;
+  cpuTopK(h_vals, h_idx, K, refV, refI);
+
+  DevBuf<int> d_vals(N), d_outV(K);
+  DevBuf<int32_t> d_idx(N), d_outI(K);
+  d_vals.upload(h_vals.data());
+  d_idx.upload(h_idx.data());
+
+  testReduceTopKSingleKernel<K>
+      <<<1, 32>>>(d_vals.ptr, d_idx.ptr, d_outV.ptr, d_outI.ptr, INT32_MIN, K);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+
+  int outV[K];
+  int32_t outI[K];
+  d_outV.download(outV);
+  d_outI.download(outI);
+
+  for (int i = 0; i < K; ++i) {
+    EXPECT_EQ(outV[i], refV[i]);
+    EXPECT_EQ(outI[i], refI[i]);
+  }
+}
+
+}  // namespace test
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/tests/core/kernels/cuda/xattention/CMakeLists.txt

@@ -0,0 +1,40 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    decoder_reshape_and_cache_test
+  SRCS
+    decoder_reshape_and_cache_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+target_link_libraries(decoder_reshape_and_cache_test PRIVATE brpc)
+
+cc_test(
+  NAME
+    prefill_reshape_and_cache_test
+  SRCS
+    prefill_reshape_and_cache_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+target_link_libraries(prefill_reshape_and_cache_test PRIVATE brpc)
+
+cc_test(
+  NAME
+    cache_select_test
+  SRCS
+    cache_select_test.cpp
+  DEPS
+    :cuda_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+target_link_libraries(cache_select_test PRIVATE brpc)

upstream_ref/xllm/tests/core/kernels/cuda/xattention/cache_select_test.cpp

@@ -0,0 +1,319 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <limits>
+#include <vector>
+
+#include "core/kernels/cuda/xattention/xattention_ops_api.h"
+
+class CacheSelctTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    device_ = torch::Device(torch::kCUDA);
+    dtype_ = torch::kFloat16;
+  }
+
+  torch::Device device_ = torch::kCPU;
+  torch::ScalarType dtype_ = torch::kFloat16;
+};
+
+void cache_select(
+    const torch::Tensor& beam_index,  // [batch * beam, 1] - out_token_index
+    std::vector<torch::Tensor>&
+        unshared_k_cache,  // per layer: [max_num_request, beam_size,
+                           // max_decode_step, kv_heads, head_dim]
+    std::vector<torch::Tensor>&
+        unshared_v_cache,  // per layer: [max_num_request, beam_size,
+                           // max_decode_step, kv_heads, head_dim]
+    const torch::Tensor& block_table,  // [batch_size, 1]
+    int64_t decode_step,               // current round (step 0, 1, ...)
+    int64_t beam_size,                 // beam width
+    int64_t layer_num) {               // number of layers
+
+  int64_t batch_size = block_table.size(0);
+  int64_t total_beams = beam_index.size(0);
+  int64_t kv_heads = unshared_k_cache[0].size(3);
+  int64_t head_dim = unshared_k_cache[0].size(4);
+  CHECK_EQ(total_beams, batch_size * beam_size) << "beam_index size mismatch";
+
+  if (layer_num > 0) {
+    int64_t max_num_request = unshared_k_cache[0].size(0);
+    int64_t max_decode_step = unshared_k_cache[0].size(2);
+    auto beam_index_reshaped =
+        beam_index.reshape({batch_size, beam_size})
+            .to(torch::kLong);  // [batch_size, beam_size]
+    auto parent_beam = (beam_index_reshaped / beam_size)
+                           .to(torch::kLong);  // [batch_size, beam_size]
+
+    auto _block_table = block_table.select(1, 0);
+
+    auto cache_opts = unshared_k_cache[0].options();
+    auto ori_k_cache = torch::zeros(
+        {layer_num, batch_size, beam_size, max_decode_step, kv_heads, head_dim},
+        cache_opts);
+    auto ori_v_cache = torch::zeros(
+        {layer_num, batch_size, beam_size, max_decode_step, kv_heads, head_dim},
+        cache_opts);
+
+    for (int64_t layer = 0; layer < layer_num; ++layer) {
+      auto k_cache = unshared_k_cache[layer];
+      auto v_cache = unshared_v_cache[layer];
+      ori_k_cache[layer].copy_(k_cache.index_select(0, _block_table));
+      ori_v_cache[layer].copy_(v_cache.index_select(0, _block_table));
+    }
+
+    for (int64_t layer = 0; layer < layer_num; ++layer) {
+      auto k_cache = unshared_k_cache[layer];
+      auto v_cache = unshared_v_cache[layer];
+
+      for (int64_t b = 0; b < batch_size; ++b) {
+        int64_t request_id = block_table[b].item<int64_t>();
+
+        CHECK_GE(request_id, 0) << "Invalid request_id: " << request_id;
+        CHECK_LT(request_id, max_num_request)
+            << "request_id (" << request_id << ") >= max_num_request ("
+            << max_num_request << ")";
+
+        auto parent_beam_batch = parent_beam[b];  // [beam_size]
+
+        for (int64_t new_beam = 0; new_beam < beam_size; ++new_beam) {
+          int64_t old_beam = parent_beam_batch[new_beam].item<int64_t>();
+
+          CHECK_GE(old_beam, 0) << "Invalid old_beam: " << old_beam;
+          CHECK_LT(old_beam, beam_size)
+              << "old_beam (" << old_beam << ") >= beam_size (" << beam_size
+              << ")";
+          if (new_beam == old_beam) {
+            continue;
+          }
+          for (int64_t step = 0; step <= decode_step; ++step) {
+            k_cache[request_id][new_beam][step].copy_(
+                ori_k_cache[layer][request_id][old_beam][step]);
+            v_cache[request_id][new_beam][step].copy_(
+                ori_v_cache[layer][request_id][old_beam][step]);
+          }
+        }
+      }
+    }
+  }
+}
+void cache_select_intelligent(
+    const torch::Tensor& beam_index,  // [batch * beam, 1] - out_token_index
+    std::vector<torch::Tensor>&
+        unshared_k_cache,  // per layer: [max_num_request, beam_size,
+                           // max_decode_step, kv_heads, head_dim]
+    std::vector<torch::Tensor>&
+        unshared_v_cache,  // per layer: [max_num_request, beam_size,
+                           // max_decode_step, kv_heads, head_dim]
+    const torch::Tensor& block_table,  // [batch_size, 1]
+    int64_t decode_step,               // current round (step 0, 1, ...)
+    int64_t beam_size,                 // beam width
+    int64_t layer_num) {
+  int64_t batch_size = block_table.size(0);
+  int64_t total_beams = beam_index.size(0);
+  CHECK_EQ(total_beams, batch_size * beam_size) << "beam_index size mismatch";
+
+  if (layer_num > 0) {
+    int64_t max_num_request = unshared_k_cache[0].size(0);
+    int64_t max_decode_step = unshared_k_cache[0].size(2);
+
+    CHECK_EQ(unshared_k_cache.size(), static_cast<size_t>(layer_num))
+        << "unshared_k_cache size mismatch";
+    CHECK_EQ(unshared_v_cache.size(), static_cast<size_t>(layer_num))
+        << "unshared_v_cache size mismatch";
+    CHECK_LT(decode_step, max_decode_step)
+        << "decode_step must be less than max_decode_step";
+    auto beam_index_reshaped =
+        beam_index.reshape({batch_size, beam_size})
+            .to(torch::kLong);  // [batch_size, beam_size]
+    auto parent_beam = (beam_index_reshaped / beam_size)
+                           .to(torch::kLong);  // [batch_size, beam_size]
+
+    auto block_table_cpu =
+        block_table.select(1, 0).to(torch::kCPU);  // [batch_size]
+    std::vector<int64_t> dirct_beam(batch_size * beam_size, -1);
+    for (int64_t b = 0; b < batch_size; ++b) {
+      for (int64_t new_beam = 0; new_beam < beam_size; ++new_beam) {
+        int64_t old_beam = parent_beam[b][new_beam].item<int64_t>();
+        dirct_beam[b * beam_size + new_beam] = old_beam > new_beam ? 1 : -1;
+      }
+    }
+    for (int64_t layer = 0; layer < layer_num; ++layer) {
+      auto& k_cache = unshared_k_cache[layer];
+      auto& v_cache = unshared_v_cache[layer];
+
+      for (int64_t b = 0; b < batch_size; ++b) {
+        int64_t request_id = block_table_cpu[b].item<int64_t>();
+
+        CHECK_GE(request_id, 0) << "Invalid request_id: " << request_id;
+        CHECK_LT(request_id, max_num_request)
+            << "request_id (" << request_id << ") >= max_num_request ("
+            << max_num_request << ")";
+
+        auto parent_beam_batch = parent_beam[b];  // [beam_size]
+
+        for (int64_t new_beam = 0; new_beam < beam_size; ++new_beam) {
+          int64_t old_beam = parent_beam_batch[new_beam].item<int64_t>();
+
+          CHECK_GE(old_beam, 0) << "Invalid old_beam: " << old_beam;
+          CHECK_LT(old_beam, beam_size)
+              << "old_beam (" << old_beam << ") >= beam_size (" << beam_size
+              << ")";
+
+          if (new_beam == old_beam) {
+            continue;
+          }
+          if (dirct_beam[b * beam_size + new_beam] == 1) {
+            for (int64_t step = 0; step <= decode_step; ++step) {
+              k_cache[request_id][new_beam][step].copy_(
+                  k_cache[request_id][old_beam][step]);
+              v_cache[request_id][new_beam][step].copy_(
+                  v_cache[request_id][old_beam][step]);
+            }
+          }
+        }
+        for (int64_t new_beam = beam_size - 1; new_beam > 0; --new_beam) {
+          int64_t old_beam = parent_beam_batch[new_beam].item<int64_t>();
+          CHECK_GE(old_beam, 0) << "Invalid old_beam: " << old_beam;
+          CHECK_LT(old_beam, beam_size)
+              << "old_beam (" << old_beam << ") >= beam_size (" << beam_size
+              << ")";
+
+          if (new_beam == old_beam) {
+            continue;
+          }
+          if (dirct_beam[b * beam_size + new_beam] == -1) {
+            for (int64_t step = 0; step <= decode_step; ++step) {
+              k_cache[request_id][new_beam][step].copy_(
+                  k_cache[request_id][old_beam][step]);
+              v_cache[request_id][new_beam][step].copy_(
+                  v_cache[request_id][old_beam][step]);
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+TEST_F(CacheSelctTest, CorrectnessTest) {
+  // Small shapes are enough to catch indexing bugs, while keeping the test
+  // fast.
+  const int64_t batch_size = 1;
+  const int64_t beam_size = 2;
+  const int64_t top_k = 2;
+  int32_t current_step = 1;
+  const int64_t kv_heads = 8;
+  const int64_t head_dim = 128;
+  const int64_t layer_num = 2;
+  const int64_t max_num_request = 1;
+  const int64_t max_decode_step = 3;
+  const auto float_opts = torch::TensorOptions().device(device_).dtype(dtype_);
+  const auto int_opts =
+      torch::TensorOptions().device(device_).dtype(torch::kInt32);
+
+  torch::Tensor beam_index = torch::randint(
+      0, beam_size * top_k, {batch_size * beam_size, 1}, int_opts);
+  auto beam_after_sort =
+      beam_index.gather(0, beam_index.argsort(static_cast<int64_t>(0), false));
+  std::vector<torch::Tensor> base_k_cache;
+  std::vector<torch::Tensor> base_v_cache;
+
+  for (int64_t layer = 0; layer < layer_num; ++layer) {
+    base_k_cache.push_back(torch::randn(
+        {max_num_request, beam_size, max_decode_step, kv_heads, head_dim},
+        float_opts));
+    base_v_cache.push_back(torch::randn(
+        {max_num_request, beam_size, max_decode_step, kv_heads, head_dim},
+        float_opts));
+  }
+
+  std::vector<torch::Tensor> k_cache_intelligent;
+  std::vector<torch::Tensor> v_cache_intelligent;
+  std::vector<torch::Tensor> k_cache_normal;
+  std::vector<torch::Tensor> v_cache_normal;
+  k_cache_intelligent.reserve(layer_num);
+  v_cache_intelligent.reserve(layer_num);
+  k_cache_normal.reserve(layer_num);
+  v_cache_normal.reserve(layer_num);
+  for (int64_t layer = 0; layer < layer_num; ++layer) {
+    k_cache_intelligent.push_back(base_k_cache[layer].clone());
+    v_cache_intelligent.push_back(base_v_cache[layer].clone());
+    k_cache_normal.push_back(base_k_cache[layer].clone());
+    v_cache_normal.push_back(base_v_cache[layer].clone());
+  }
+
+  // CPU reference uses request-level block table: [batch_size, 1].
+  torch::Tensor block_table =
+      torch::arange(batch_size, int_opts).view({batch_size, 1});
+
+  // CUDA kernel uses sequence-level block table shape [batch_size * beam_size,
+  // 1].
+  torch::Tensor seq_block_table =
+      torch::arange(batch_size * beam_size, int_opts)
+          .view({batch_size * beam_size, 1});
+
+  //   use cache intelligent
+  cache_select_intelligent(beam_after_sort,
+                           k_cache_intelligent,
+                           v_cache_intelligent,
+                           block_table,
+                           current_step,
+                           beam_size,
+                           layer_num);
+  //   use cache normal
+  cache_select(beam_after_sort,
+               k_cache_normal,
+               v_cache_normal,
+               block_table,
+               current_step,
+               beam_size,
+               layer_num);
+
+  // CUDA kernel version (single launch over layer axis).
+  std::vector<torch::Tensor> k_cache_cuda;
+  std::vector<torch::Tensor> v_cache_cuda;
+  k_cache_cuda.reserve(layer_num);
+  v_cache_cuda.reserve(layer_num);
+  for (int64_t layer = 0; layer < layer_num; ++layer) {
+    k_cache_cuda.push_back(base_k_cache[layer].clone());
+    v_cache_cuda.push_back(base_v_cache[layer].clone());
+  }
+  xllm::kernel::cuda::cache_select(beam_after_sort,
+                                   k_cache_cuda,
+                                   v_cache_cuda,
+                                   seq_block_table,
+                                   current_step,
+                                   beam_size,
+                                   layer_num);
+  for (int64_t layer = 0; layer < layer_num; ++layer) {
+    EXPECT_TRUE(torch::allclose(
+        k_cache_intelligent[layer], k_cache_normal[layer], 1e-5, 1e-5));
+    EXPECT_TRUE(torch::allclose(
+        v_cache_intelligent[layer], v_cache_normal[layer], 1e-5, 1e-5));
+
+    EXPECT_TRUE(torch::allclose(
+        k_cache_intelligent[layer], k_cache_cuda[layer], 1e-5, 1e-5));
+    EXPECT_TRUE(torch::allclose(
+        v_cache_intelligent[layer], v_cache_cuda[layer], 1e-5, 1e-5));
+  }
+}

upstream_ref/xllm/tests/core/kernels/cuda/xattention/decoder_reshape_and_cache_test.cpp

@@ -0,0 +1,150 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "core/kernels/cuda/xattention/xattention_ops_api.h"
+
+namespace xllm::kernel::cuda {
+namespace test {
+
+class DecoderReshapeAndCacheTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    device_ = torch::Device(torch::kCUDA);
+    dtype_ = torch::kFloat16;
+  }
+
+  torch::Device device_ = torch::kCPU;
+  torch::ScalarType dtype_ = torch::kFloat16;
+};
+
+// This test validates `xllm::kernel::cuda::decoder_reshape_and_cache` against
+// a PyTorch reference implementation that follows the current CUDA kernel
+// contract.
+//
+// Kernel input contract:
+// - proj_k/proj_v: [batch_size, beam_size, kv_heads, head_dim], where
+//
+// Cache layout:
+// - unshared_k/v_cache: [max_num_request, beam_size, max_decode_step, kv_heads,
+//   head_dim]
+//
+//
+// The test checks:
+// 1) full-tensor equivalence (CUDA output vs reference output)
+// 2) per-sequence slice correctness at the target decode step
+void torch_reference(const torch::Tensor& proj_k,
+                     const torch::Tensor& proj_v,
+                     torch::Tensor& unshared_k_cache,
+                     torch::Tensor& unshared_v_cache,
+                     const torch::Tensor& step) {
+  const int64_t batch_size = proj_k.size(0);
+  const int64_t beam_size = proj_k.size(1);
+  const int64_t kv_heads = proj_k.size(2);
+  const int64_t head_dim = proj_k.size(3);
+
+  const int64_t max_num_request = unshared_k_cache.size(0);
+  const int64_t max_decode_step = unshared_k_cache.size(2);
+
+  CHECK_EQ(proj_k.dim(), 4) << "proj_k must be 4-dimensional";
+  CHECK_EQ(proj_v.dim(), 4) << "proj_v must be 4-dimensional";
+  CHECK_EQ(proj_v.sizes(), proj_k.sizes())
+      << "proj_v and proj_k must have same shape";
+
+  CHECK_EQ(unshared_k_cache.dim(), 5)
+      << "unshared_k_cache must be 5-dimensional";
+  CHECK_EQ(unshared_v_cache.sizes(), unshared_k_cache.sizes())
+      << "unshared_v_cache and unshared_k_cache must have same shape";
+  CHECK_EQ(unshared_k_cache.size(3), kv_heads)
+      << "unshared_k_cache kv_heads mismatch";
+  CHECK_EQ(unshared_k_cache.size(4), head_dim)
+      << "unshared_k_cache head_dim mismatch";
+  CHECK_LE(batch_size, max_num_request)
+      << "batch_size must be <= max_num_request";
+
+  CHECK_EQ(step.dim(), 1) << "step must be 1-dimensional";
+  CHECK_EQ(step.size(0), 1) << "step must have shape [1]";
+  const int64_t step_value = step[0].item<int64_t>();
+  CHECK_GE(step_value, 0) << "step must be >= 0";
+  CHECK_LT(step_value, max_decode_step)
+      << "step must be less than max_decode_step";
+
+  for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
+    for (int64_t beam_idx = 0; beam_idx < beam_size; ++beam_idx) {
+      unshared_k_cache[batch_idx][beam_idx]
+          .select(0, step_value)
+          .copy_(proj_k[batch_idx][beam_idx]);
+      unshared_v_cache[batch_idx][beam_idx]
+          .select(0, step_value)
+          .copy_(proj_v[batch_idx][beam_idx]);
+    }
+  }
+}
+
+TEST_F(DecoderReshapeAndCacheTest, CorrectnessTest) {
+  const int64_t batch_size = 1;
+  const int64_t beam_size = 2;
+  const int64_t kv_heads = 8;
+  const int64_t head_dim = 128;
+  const int64_t max_num_request = 2;
+  const int64_t max_decode_step = 3;
+
+  torch::Tensor step = torch::tensor({1}, torch::kInt32).to(device_);
+
+  const auto float_opts = torch::TensorOptions().device(device_).dtype(dtype_);
+  const auto int_opts =
+      torch::TensorOptions().device(device_).dtype(torch::kInt32);
+
+  torch::Tensor proj_k =
+      torch::randn({batch_size, beam_size, kv_heads, head_dim}, float_opts);
+  torch::Tensor proj_v =
+      torch::randn({batch_size, beam_size, kv_heads, head_dim}, float_opts);
+
+  torch::Tensor unshared_k_cache = torch::zeros(
+      {max_num_request, beam_size, max_decode_step, kv_heads, head_dim},
+      float_opts);
+  torch::Tensor unshared_v_cache = torch::zeros(
+      {max_num_request, beam_size, max_decode_step, kv_heads, head_dim},
+      float_opts);
+
+  torch::Tensor ref_k_cache = unshared_k_cache.clone();
+  torch::Tensor ref_v_cache = unshared_v_cache.clone();
+
+  decoder_reshape_and_cache(
+      proj_k, proj_v, unshared_k_cache, unshared_v_cache, step);
+
+  torch_reference(proj_k, proj_v, ref_k_cache, ref_v_cache, step);
+
+  EXPECT_TRUE(torch::allclose(unshared_k_cache, ref_k_cache, 1e-5, 1e-5));
+  EXPECT_TRUE(torch::allclose(unshared_v_cache, ref_v_cache, 1e-5, 1e-5));
+
+  const int64_t step_value = step[0].item<int64_t>();
+  for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
+    for (int64_t beam_idx = 0; beam_idx < beam_size; ++beam_idx) {
+      torch::Tensor copied_k =
+          unshared_k_cache[batch_idx][beam_idx].select(0, step_value);
+      torch::Tensor source_k = proj_k[batch_idx][beam_idx];
+      EXPECT_TRUE(torch::allclose(copied_k, source_k, 1e-5, 1e-5));
+    }
+  }
+}
+
+}  // namespace test
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/tests/core/kernels/cuda/xattention/prefill_reshape_and_cache_test.cpp

@@ -0,0 +1,167 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/cuda.h>
+#include <torch/torch.h>
+
+#include "core/kernels/cuda/xattention/xattention_ops_api.h"
+
+namespace xllm::kernel::cuda {
+namespace test {
+namespace {
+
+void apply_reference_prefill_reshape_and_cache(const torch::Tensor& proj_k,
+                                               const torch::Tensor& proj_v,
+                                               torch::Tensor& shared_k_cache,
+                                               torch::Tensor& shared_v_cache) {
+  const int64_t shared_len = proj_k.size(0);
+  shared_k_cache.slice(0, 0, shared_len).copy_(proj_k);
+  shared_v_cache.slice(0, 0, shared_len).copy_(proj_v);
+}
+
+void run_and_check_prefill_reshape_and_cache(const torch::Tensor& proj_k,
+                                             const torch::Tensor& proj_v,
+                                             torch::Tensor& shared_k_cache,
+                                             torch::Tensor& shared_v_cache) {
+  const int64_t shared_len = proj_k.size(0);
+  const int64_t cache_len = shared_k_cache.size(0);
+
+  auto tail_k_before = shared_k_cache.slice(0, shared_len, cache_len).clone();
+  auto tail_v_before = shared_v_cache.slice(0, shared_len, cache_len).clone();
+
+  auto ref_k_cache = shared_k_cache.clone();
+  auto ref_v_cache = shared_v_cache.clone();
+  apply_reference_prefill_reshape_and_cache(
+      proj_k, proj_v, ref_k_cache, ref_v_cache);
+
+  prefill_reshape_and_cache(proj_k, proj_v, shared_k_cache, shared_v_cache);
+  torch::cuda::synchronize();
+
+  EXPECT_TRUE(torch::equal(shared_k_cache, ref_k_cache));
+  EXPECT_TRUE(torch::equal(shared_v_cache, ref_v_cache));
+
+  EXPECT_TRUE(torch::equal(shared_k_cache.slice(0, shared_len, cache_len),
+                           tail_k_before));
+  EXPECT_TRUE(torch::equal(shared_v_cache.slice(0, shared_len, cache_len),
+                           tail_v_before));
+}
+
+class PrefillReshapeAndCacheTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA not available, skipping test.";
+    }
+    torch::manual_seed(2026);
+    device_ = torch::Device(torch::kCUDA, 0);
+  }
+
+  torch::Device device_ = torch::Device(torch::kCPU);
+};
+
+TEST_F(PrefillReshapeAndCacheTest, MatchesReferenceVectorizedPath) {
+  const int64_t shared_len = 7;
+  const int64_t cache_len = 11;
+  const int64_t kv_heads = 8;
+  const int64_t head_dim = 128;
+
+  auto float_opts = torch::TensorOptions().device(device_).dtype(torch::kHalf);
+  auto proj_k = torch::randn({shared_len, kv_heads, head_dim}, float_opts);
+  auto proj_v = torch::randn({shared_len, kv_heads, head_dim}, float_opts);
+
+  auto shared_k_cache =
+      torch::full({cache_len, kv_heads, head_dim}, -7.0f, float_opts);
+  auto shared_v_cache =
+      torch::full({cache_len, kv_heads, head_dim}, 5.0f, float_opts);
+
+  run_and_check_prefill_reshape_and_cache(
+      proj_k, proj_v, shared_k_cache, shared_v_cache);
+}
+
+TEST_F(PrefillReshapeAndCacheTest, MatchesReferenceQkvSliceFp16) {
+  const int64_t shared_len = 9;
+  const int64_t cache_len = 13;
+  const int64_t num_q_heads = 16;
+  const int64_t kv_heads = 8;
+  const int64_t head_dim = 128;
+  const int64_t q_size = num_q_heads * head_dim;
+  const int64_t kv_size = kv_heads * head_dim;
+
+  auto float_opts = torch::TensorOptions().device(device_).dtype(torch::kHalf);
+  auto qkv = torch::randn({shared_len, q_size + 2 * kv_size}, float_opts) * 0.2;
+
+  auto proj_k = qkv.slice(1, q_size, q_size + kv_size)
+                    .view({shared_len, kv_heads, head_dim});
+  auto proj_v = qkv.slice(1, q_size + kv_size, q_size + 2 * kv_size)
+                    .view({shared_len, kv_heads, head_dim});
+
+  ASSERT_FALSE(proj_k.is_contiguous());
+  ASSERT_FALSE(proj_v.is_contiguous());
+  ASSERT_EQ(proj_k.stride(0), q_size + 2 * kv_size);
+  ASSERT_EQ(proj_v.stride(0), q_size + 2 * kv_size);
+  ASSERT_EQ(proj_k.stride(1), head_dim);
+  ASSERT_EQ(proj_v.stride(1), head_dim);
+  ASSERT_EQ(proj_k.stride(2), 1);
+  ASSERT_EQ(proj_v.stride(2), 1);
+
+  auto shared_k_cache =
+      torch::full({cache_len, kv_heads, head_dim}, -3.0f, float_opts);
+  auto shared_v_cache =
+      torch::full({cache_len, kv_heads, head_dim}, 9.0f, float_opts);
+
+  run_and_check_prefill_reshape_and_cache(
+      proj_k, proj_v, shared_k_cache, shared_v_cache);
+}
+
+TEST_F(PrefillReshapeAndCacheTest, MatchesReferenceQkvSliceBf16) {
+  const int64_t shared_len = 6;
+  const int64_t cache_len = 10;
+  const int64_t num_q_heads = 8;
+  const int64_t kv_heads = 4;
+  const int64_t head_dim = 128;
+  const int64_t q_size = num_q_heads * head_dim;
+  const int64_t kv_size = kv_heads * head_dim;
+
+  auto bf16_opts =
+      torch::TensorOptions().device(device_).dtype(torch::kBFloat16);
+  auto qkv = torch::randn({shared_len, q_size + 2 * kv_size}, bf16_opts) * 0.2;
+
+  auto proj_k = qkv.slice(1, q_size, q_size + kv_size)
+                    .view({shared_len, kv_heads, head_dim});
+  auto proj_v = qkv.slice(1, q_size + kv_size, q_size + 2 * kv_size)
+                    .view({shared_len, kv_heads, head_dim});
+
+  ASSERT_FALSE(proj_k.is_contiguous());
+  ASSERT_FALSE(proj_v.is_contiguous());
+  ASSERT_EQ(proj_k.stride(0), q_size + 2 * kv_size);
+  ASSERT_EQ(proj_v.stride(0), q_size + 2 * kv_size);
+  ASSERT_EQ(proj_k.stride(1), head_dim);
+  ASSERT_EQ(proj_v.stride(1), head_dim);
+  ASSERT_EQ(proj_k.stride(2), 1);
+  ASSERT_EQ(proj_v.stride(2), 1);
+
+  auto shared_k_cache =
+      torch::full({cache_len, kv_heads, head_dim}, -3.0f, bf16_opts);
+  auto shared_v_cache =
+      torch::full({cache_len, kv_heads, head_dim}, 9.0f, bf16_opts);
+
+  run_and_check_prefill_reshape_and_cache(
+      proj_k, proj_v, shared_k_cache, shared_v_cache);
+}
+
+}  // namespace
+}  // namespace test
+}  // namespace xllm::kernel::cuda

upstream_ref/xllm/tests/core/kernels/npu/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(tilelang)

upstream_ref/xllm/tests/core/kernels/npu/tilelang/CMakeLists.txt

@@ -0,0 +1,37 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    rope_wrapper_test
+  SRCS
+    rope_wrapper_test.cpp
+  DEPS
+    :tilelang_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+
+cc_test(
+  NAME
+    fused_gdn_gating_wrapper_test
+  SRCS
+    fused_gdn_gating_wrapper_test.cpp
+  DEPS
+    :tilelang_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)
+
+cc_test(
+  NAME
+    split_qkv_rmsnorm_mrope_wrapper_test
+  SRCS
+    split_qkv_rmsnorm_mrope_wrapper_test.cpp
+  DEPS
+    :tilelang_kernels
+    torch
+    GTest::gtest_main
+    glog::glog
+)

upstream_ref/xllm/tests/core/kernels/npu/tilelang/fused_gdn_gating_wrapper_test.cpp

@@ -0,0 +1,220 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/torch.h>
+#include <torch_npu/torch_npu.h>
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "core/kernels/npu/tilelang/tilelang_ops_api.h"
+
+namespace xllm::kernel::npu::tilelang {
+namespace {
+
+class TileLangFusedGdnGatingWrapperTest : public ::testing::Test {
+ protected:
+  static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); }
+
+  static void TearDownTestSuite() { torch_npu::finalize_npu(); }
+};
+
+struct FusedGdnGatingTestCase {
+  std::string name;
+  int64_t num_batches;
+  int64_t num_heads;
+  int64_t seed;
+  float beta = 1.0F;
+  float threshold = 20.0F;
+};
+
+std::pair<torch::Tensor, torch::Tensor> torch_fused_gdn_gating(
+    const torch::Tensor& A_log,
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& dt_bias,
+    float beta = 1.0F,
+    float threshold = 20.0F) {
+  namespace F = torch::nn::functional;
+
+  auto softplus_out =
+      F::softplus(a.to(torch::kFloat32) + dt_bias,
+                  F::SoftplusFuncOptions().beta(beta).threshold(threshold));
+  auto g = -A_log.exp() * softplus_out;
+  auto beta_output = torch::sigmoid(b.to(torch::kFloat32)).to(torch::kBFloat16);
+  return {g.unsqueeze(0), beta_output.unsqueeze(0)};
+}
+
+void run_fused_gdn_gating_case(const FusedGdnGatingTestCase& test_case) {
+  ASSERT_GT(test_case.num_batches, 0);
+
+  const auto device = torch::Device("npu:0");
+  torch::manual_seed(test_case.seed);
+
+  auto fp32_opts = torch::TensorOptions().dtype(torch::kFloat32).device(device);
+  auto bf16_opts =
+      torch::TensorOptions().dtype(torch::kBFloat16).device(device);
+
+  auto A_log = torch::randn({test_case.num_heads}, fp32_opts);
+  auto a =
+      torch::randn({test_case.num_batches, test_case.num_heads}, bf16_opts);
+  auto b =
+      torch::randn({test_case.num_batches, test_case.num_heads}, bf16_opts);
+  auto dt_bias = torch::randn({test_case.num_heads}, fp32_opts);
+
+  auto [g_ref, beta_ref] = torch_fused_gdn_gating(
+      A_log, a, b, dt_bias, test_case.beta, test_case.threshold);
+  auto [g_out, beta_out] = fused_gdn_gating(
+      A_log, a, b, dt_bias, test_case.beta, test_case.threshold);
+
+  auto g_max_diff = (g_out - g_ref).abs().max().item<float>();
+  int64_t g_max_index = (g_out - g_ref).abs().argmax().item<int64_t>();
+  const int64_t g_max_head = g_max_index % test_case.num_heads;
+  const int64_t g_max_row =
+      (g_max_index / test_case.num_heads) % test_case.num_batches;
+  auto beta_max_diff =
+      (beta_out.to(torch::kFloat32) - beta_ref.to(torch::kFloat32))
+          .abs()
+          .max()
+          .item<float>();
+  int64_t beta_max_index =
+      (beta_out.to(torch::kFloat32) - beta_ref.to(torch::kFloat32))
+          .abs()
+          .argmax()
+          .item<int64_t>();
+  const int64_t beta_max_head = beta_max_index % test_case.num_heads;
+  const int64_t beta_max_row =
+      (beta_max_index / test_case.num_heads) % test_case.num_batches;
+
+  EXPECT_TRUE(torch::allclose(g_out, g_ref, 1e-3, 1e-3))
+      << "g mismatch, max_diff=" << g_max_diff << ", row=" << g_max_row
+      << ", head=" << g_max_head;
+  EXPECT_TRUE(torch::allclose(beta_out, beta_ref, 1e-2, 1e-2))
+      << "beta mismatch, max_diff=" << beta_max_diff << ", row=" << beta_max_row
+      << ", head=" << beta_max_head;
+}
+
+TEST_F(TileLangFusedGdnGatingWrapperTest, MatchesTorchReference) {
+  const std::vector<FusedGdnGatingTestCase> cases = {
+      {
+          .name = "tiny_b1_h8",
+          .num_batches = 1,
+          .num_heads = 8,
+          .seed = 101,
+      },
+      {
+          .name = "tiny_b17_h8",
+          .num_batches = 17,
+          .num_heads = 8,
+          .seed = 101,
+      },
+      {
+          .name = "tiny_b1_h16",
+          .num_batches = 1,
+          .num_heads = 16,
+          .seed = 101,
+      },
+      {
+          .name = "qwen35_tp2_image_b275_h24",
+          .num_batches = 275,
+          .num_heads = 24,
+          .seed = 201,
+      },
+      {
+          .name = "qwen35_tp2_video_b3716_h24",
+          .num_batches = 3716,
+          .num_heads = 24,
+          .seed = 202,
+      },
+      {
+          .name = "qwen35_tp2_video_b4096_h24",
+          .num_batches = 4096,
+          .num_heads = 24,
+          .seed = 203,
+      },
+      {
+          .name = "qwen35_tp2_video_b8192_h24",
+          .num_batches = 8192,
+          .num_heads = 24,
+          .seed = 204,
+      },
+      {
+          .name = "small_b17_h32",
+          .num_batches = 17,
+          .num_heads = 32,
+          .seed = 102,
+      },
+      {
+          .name = "medium_b29_h48",
+          .num_batches = 29,
+          .num_heads = 48,
+          .seed = 103,
+      },
+      {
+          .name = "medium_b131_h64",
+          .num_batches = 131,
+          .num_heads = 64,
+          .seed = 104,
+      },
+      {
+          .name = "medium_b257_h128",
+          .num_batches = 257,
+          .num_heads = 128,
+          .seed = 105,
+      },
+      {
+          .name = "large_b4096_h32",
+          .num_batches = 4096,
+          .num_heads = 32,
+          .seed = 106,
+      },
+      {
+          .name = "chunked_b8192_h32",
+          .num_batches = 8192,
+          .num_heads = 32,
+          .seed = 108,
+      },
+      {
+          .name = "custom_beta2_threshold0p5_b33_h64",
+          .num_batches = 33,
+          .num_heads = 64,
+          .seed = 107,
+          .beta = 2.0F,
+          .threshold = 0.5F,
+      },
+      {
+          .name = "large_b65536_h32",
+          .num_batches = 65536,
+          .num_heads = 32,
+          .seed = 109,
+      },
+      {
+          .name = "large_b262144_h32",
+          .num_batches = 262144,
+          .num_heads = 32,
+          .seed = 110,
+      },
+  };
+
+  for (const auto& test_case : cases) {
+    SCOPED_TRACE(test_case.name);
+    run_fused_gdn_gating_case(test_case);
+  }
+}
+
+}  // namespace
+}  // namespace xllm::kernel::npu::tilelang

upstream_ref/xllm/tests/core/kernels/npu/tilelang/rope_wrapper_test.cpp

@@ -0,0 +1,408 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/torch.h>
+#include <torch_npu/csrc/core/npu/NPUStream.h>
+#include <torch_npu/torch_npu.h>
+
+#include <functional>
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include "acl/acl.h"
+#include "tilelang_ops_api.h"
+
+namespace xllm::kernel::npu::tilelang {
+namespace {
+
+class TileLangRopeWrapperTest : public ::testing::Test {
+ protected:
+  static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); }
+
+  static void TearDownTestSuite() { torch_npu::finalize_npu(); }
+};
+
+struct RopeTestCase {
+  std::string name;
+  int64_t num_tokens;
+  int64_t num_heads;
+  int64_t full_head_dim;
+  int64_t rope_dim;
+  int64_t start_dim;
+  int64_t seed;
+};
+
+torch::Tensor torch_rope_ref(const torch::Tensor& x,
+                             const torch::Tensor& sin,
+                             const torch::Tensor& cos) {
+  auto cos_ref = cos;
+  auto sin_ref = sin;
+  if (cos_ref.dim() == 2) {
+    cos_ref = cos_ref.unsqueeze(1);
+    sin_ref = sin_ref.unsqueeze(1);
+  }
+
+  auto x_fp32 = x.to(torch::kFloat32);
+  auto cos_fp32 = cos_ref.to(torch::kFloat32);
+  auto sin_fp32 = sin_ref.to(torch::kFloat32);
+
+  auto x_reshaped =
+      x_fp32.view({x_fp32.size(0), x_fp32.size(1), x_fp32.size(2) / 2, 2});
+  auto x0 = x_reshaped.index({torch::indexing::Slice(),
+                              torch::indexing::Slice(),
+                              torch::indexing::Slice(),
+                              0});
+  auto x1 = x_reshaped.index({torch::indexing::Slice(),
+                              torch::indexing::Slice(),
+                              torch::indexing::Slice(),
+                              1});
+  auto x_rotated = torch::stack({-x1, x0}, /*dim=*/-1).flatten(-2);
+
+  auto out = x_fp32 * cos_fp32 + x_rotated * sin_fp32;
+  return out.to(torch::kBFloat16);
+}
+
+double measure_npu_event_ms(const std::function<void()>& fn,
+                            int32_t device_id,
+                            int warmup_iters = 5,
+                            int measure_iters = 100) {
+  CHECK_GT(measure_iters, 0) << "measure_iters must be > 0";
+  CHECK_GE(warmup_iters, 0) << "warmup_iters must be >= 0";
+
+  const aclrtStream stream = c10_npu::getCurrentNPUStream(device_id).stream();
+  for (int i = 0; i < warmup_iters; ++i) {
+    fn();
+  }
+  CHECK_EQ(aclrtSynchronizeStream(stream), ACL_SUCCESS)
+      << "warmup stream synchronize failed";
+
+  aclrtEvent start_event = nullptr;
+  aclrtEvent end_event = nullptr;
+  CHECK_EQ(aclrtCreateEvent(&start_event), ACL_SUCCESS)
+      << "aclrtCreateEvent(start) failed";
+  CHECK_EQ(aclrtCreateEvent(&end_event), ACL_SUCCESS)
+      << "aclrtCreateEvent(end) failed";
+
+  CHECK_EQ(aclrtRecordEvent(start_event, stream), ACL_SUCCESS)
+      << "aclrtRecordEvent(start) failed";
+  for (int i = 0; i < measure_iters; ++i) {
+    fn();
+  }
+  CHECK_EQ(aclrtRecordEvent(end_event, stream), ACL_SUCCESS)
+      << "aclrtRecordEvent(end) failed";
+  CHECK_EQ(aclrtSynchronizeEvent(end_event), ACL_SUCCESS)
+      << "aclrtSynchronizeEvent(end) failed";
+
+  float elapsed_ms = 0.0F;
+  CHECK_EQ(aclrtEventElapsedTime(&elapsed_ms, start_event, end_event),
+           ACL_SUCCESS)
+      << "aclrtEventElapsedTime failed";
+  CHECK_EQ(aclrtDestroyEvent(start_event), ACL_SUCCESS)
+      << "aclrtDestroyEvent(start) failed";
+  CHECK_EQ(aclrtDestroyEvent(end_event), ACL_SUCCESS)
+      << "aclrtDestroyEvent(end) failed";
+
+  return static_cast<double>(elapsed_ms) / static_cast<double>(measure_iters);
+}
+
+torch::Tensor maybe_narrow(const torch::Tensor& tensor,
+                           int64_t start_dim,
+                           int64_t rope_dim) {
+  if (start_dim == 0 && rope_dim == tensor.size(2)) {
+    return tensor;
+  }
+  return tensor.narrow(/*dim=*/2, /*start=*/start_dim, /*length=*/rope_dim);
+}
+
+void run_apply_rotary_case(const RopeTestCase& test_case) {
+  ASSERT_GT(test_case.num_tokens, 0);
+  ASSERT_GT(test_case.num_heads, 0);
+  ASSERT_GT(test_case.full_head_dim, 0);
+  ASSERT_GT(test_case.rope_dim, 0);
+  ASSERT_GE(test_case.start_dim, 0);
+  ASSERT_LE(test_case.start_dim + test_case.rope_dim, test_case.full_head_dim);
+
+  const auto npu_device = torch::Device("npu:0");
+  const int32_t device_id = npu_device.index();
+  const auto bf16_opts =
+      torch::TensorOptions().dtype(torch::kBFloat16).device(npu_device);
+
+  torch::manual_seed(test_case.seed);
+  auto q_full = torch::randn(
+      {test_case.num_tokens, test_case.num_heads, test_case.full_head_dim},
+      bf16_opts);
+  auto k_full = torch::randn(
+      {test_case.num_tokens, test_case.num_heads, test_case.full_head_dim},
+      bf16_opts);
+  auto sin_cache =
+      torch::randn({test_case.num_tokens, test_case.rope_dim}, bf16_opts);
+  auto cos_cache =
+      torch::randn({test_case.num_tokens, test_case.rope_dim}, bf16_opts);
+
+  auto q_input = maybe_narrow(q_full, test_case.start_dim, test_case.rope_dim);
+  auto k_input = maybe_narrow(k_full, test_case.start_dim, test_case.rope_dim);
+
+  if (test_case.start_dim > 0) {
+    EXPECT_EQ(q_input.storage_offset(), test_case.start_dim);
+    EXPECT_EQ(k_input.storage_offset(), test_case.start_dim);
+    if (test_case.num_tokens * test_case.num_heads > 1) {
+      EXPECT_FALSE(q_input.is_contiguous());
+      EXPECT_FALSE(k_input.is_contiguous());
+    }
+  }
+
+  auto q_ref = torch_rope_ref(q_input, sin_cache, cos_cache);
+  auto k_ref = torch_rope_ref(k_input, sin_cache, cos_cache);
+  auto q_runtime_full = q_full.clone();
+  auto k_runtime_full = k_full.clone();
+  auto q =
+      maybe_narrow(q_runtime_full, test_case.start_dim, test_case.rope_dim);
+  auto k =
+      maybe_narrow(k_runtime_full, test_case.start_dim, test_case.rope_dim);
+  rope_in_place(q, sin_cache, cos_cache);
+  rope_in_place(k, sin_cache, cos_cache);
+
+  auto q_bench_full = q_full.clone();
+  auto k_bench_full = k_full.clone();
+  auto q_bench =
+      maybe_narrow(q_bench_full, test_case.start_dim, test_case.rope_dim);
+  auto k_bench =
+      maybe_narrow(k_bench_full, test_case.start_dim, test_case.rope_dim);
+  const double ref_elapsed_ms = measure_npu_event_ms(
+      [&]() {
+        [[maybe_unused]] auto q_ref_bench =
+            torch_rope_ref(q_input, sin_cache, cos_cache);
+        [[maybe_unused]] auto k_ref_bench =
+            torch_rope_ref(k_input, sin_cache, cos_cache);
+      },
+      device_id);
+  const double tl_elapsed_ms = measure_npu_event_ms(
+      [&]() {
+        rope_in_place(q_bench, sin_cache, cos_cache);
+        rope_in_place(k_bench, sin_cache, cos_cache);
+      },
+      device_id);
+
+  const double speedup =
+      tl_elapsed_ms > 0.0 ? ref_elapsed_ms / tl_elapsed_ms : 0.0;
+  std::cout << "[rope_wrapper_test] case=" << test_case.name
+            << ", ref_ms=" << ref_elapsed_ms
+            << ", tilelang_ms=" << tl_elapsed_ms << ", speedup=" << speedup
+            << "x" << std::endl;
+
+  auto q_max_diff = (q.to(torch::kFloat32) - q_ref.to(torch::kFloat32))
+                        .abs()
+                        .max()
+                        .item<float>();
+  auto k_max_diff = (k.to(torch::kFloat32) - k_ref.to(torch::kFloat32))
+                        .abs()
+                        .max()
+                        .item<float>();
+
+  EXPECT_TRUE(torch::allclose(q, q_ref, /*rtol=*/1e-2, /*atol=*/1e-2))
+      << "q mismatch: tilelang output differs from interleaved rope reference"
+      << ", max_diff=" << q_max_diff;
+  EXPECT_TRUE(torch::allclose(k, k_ref, /*rtol=*/1e-2, /*atol=*/1e-2))
+      << "k mismatch: tilelang output differs from interleaved rope reference"
+      << ", max_diff=" << k_max_diff;
+}
+
+TEST_F(TileLangRopeWrapperTest, ApplyRotaryMatchesNpuReferenceVariant128x128) {
+  const std::vector<RopeTestCase> cases = {
+      {.name = "baseline_16x4_hd128_rd128",
+       .num_tokens = 16,
+       .num_heads = 4,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 20260213},
+      {.name = "large_tokens_2051x2_hd128_rd128",
+       .num_tokens = 2051,
+       .num_heads = 2,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 20260214},
+      {.name = "tiny_1x1_hd128_rd128",
+       .num_tokens = 1,
+       .num_heads = 1,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 101},
+      {.name = "odd_tokens_7x3_hd128_rd128",
+       .num_tokens = 7,
+       .num_heads = 3,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 102},
+      {.name = "token_dim_64x4_hd128_rd128",
+       .num_tokens = 64,
+       .num_heads = 4,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 107},
+      {.name = "chunk_boundary_8x5_hd128_rd128",
+       .num_tokens = 8,
+       .num_heads = 5,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 103},
+      {.name = "cross_chunk_9x5_hd128_rd128",
+       .num_tokens = 9,
+       .num_heads = 5,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 104},
+      {.name = "head_dim_4x64_hd128_rd128",
+       .num_tokens = 4,
+       .num_heads = 64,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 108},
+      {.name = "medium_127x8_hd128_rd128",
+       .num_tokens = 127,
+       .num_heads = 8,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 105},
+      {.name = "large_heads_33x16_hd128_rd128",
+       .num_tokens = 33,
+       .num_heads = 16,
+       .full_head_dim = 128,
+       .rope_dim = 128,
+       .start_dim = 0,
+       .seed = 106},
+  };
+
+  for (const auto& test_case : cases) {
+    SCOPED_TRACE(::testing::Message() << "case=" << test_case.name
+                                      << ", num_tokens=" << test_case.num_tokens
+                                      << ", num_heads=" << test_case.num_heads);
+    run_apply_rotary_case(test_case);
+  }
+}
+
+TEST_F(TileLangRopeWrapperTest, ApplyRotaryMatchesNpuReferenceVariant576x64) {
+  constexpr int64_t kNumHeads = 1;
+  constexpr int64_t kFullHeadDim = 576;
+  constexpr int64_t kStartDim = 512;
+  constexpr int64_t kRopeDim = 64;
+
+  const std::vector<RopeTestCase> cases = {
+      {.name = "1x576_start512_rope64",
+       .num_tokens = 1,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260226},
+      {.name = "8x576_start512_rope64",
+       .num_tokens = 8,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260227},
+      {.name = "47x576_start512_rope64",
+       .num_tokens = 47,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260301},
+      {.name = "48x576_start512_rope64",
+       .num_tokens = 48,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260302},
+      {.name = "49x576_start512_rope64",
+       .num_tokens = 49,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260303},
+      {.name = "95x576_start512_rope64",
+       .num_tokens = 95,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260304},
+      {.name = "96x576_start512_rope64",
+       .num_tokens = 96,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260305},
+      {.name = "97x576_start512_rope64",
+       .num_tokens = 97,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260306},
+      {.name = "128x576_start512_rope64",
+       .num_tokens = 128,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260228},
+      {.name = "512x576_start512_rope64",
+       .num_tokens = 512,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260307},
+      {.name = "1024x576_start512_rope64",
+       .num_tokens = 1024,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260308},
+      {.name = "2048x576_start512_rope64",
+       .num_tokens = 2048,
+       .num_heads = kNumHeads,
+       .full_head_dim = kFullHeadDim,
+       .rope_dim = kRopeDim,
+       .start_dim = kStartDim,
+       .seed = 20260225},
+  };
+
+  for (const auto& test_case : cases) {
+    SCOPED_TRACE(::testing::Message() << "case=" << test_case.name
+                                      << ", num_tokens=" << test_case.num_tokens
+                                      << ", num_heads=" << test_case.num_heads);
+    run_apply_rotary_case(test_case);
+  }
+}
+
+}  // namespace
+}  // namespace xllm::kernel::npu::tilelang

upstream_ref/xllm/tests/core/kernels/npu/tilelang/split_qkv_rmsnorm_mrope_wrapper_test.cpp

@@ -0,0 +1,661 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gtest/gtest.h>
+#include <torch/torch.h>
+#include <torch_npu/csrc/core/npu/NPUStream.h>
+#include <torch_npu/torch_npu.h>
+
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "acl/acl.h"
+#include "core/kernels/npu/tilelang/tilelang_ops_api.h"
+
+namespace xllm::kernel::npu::tilelang {
+namespace {
+
+constexpr int64_t kHeadSize = 256;
+constexpr double kRmsNormEps = 1e-6;
+
+class TileLangSplitQkvRmsnormMRopeWrapperTest : public ::testing::Test {
+ protected:
+  static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); }
+
+  static void TearDownTestSuite() { torch_npu::finalize_npu(); }
+};
+
+struct SplitQkvTestCase {
+  std::string name;
+  int64_t num_tokens;
+  int64_t num_q_heads;
+  int64_t num_kv_heads;
+  int64_t seed;
+  bool is_interleaved;
+};
+
+torch::Tensor rms_norm_ref(const torch::Tensor& x,
+                           const torch::Tensor& weight,
+                           double eps) {
+  torch::Tensor x_fp32 = x.to(torch::kFloat32);
+  torch::Tensor weight_fp32 = weight.to(torch::kFloat32);
+  torch::Tensor reciprocal_std =
+      torch::rsqrt(torch::mean(x_fp32 * x_fp32, /*dim=*/-1, true) + eps);
+  return x_fp32 * reciprocal_std * weight_fp32;
+}
+
+std::pair<torch::Tensor, torch::Tensor> assemble_non_interleaved_mrope_rows_ref(
+    const torch::Tensor& cos_sin,
+    const std::vector<int64_t>& mrope_section) {
+  const int64_t half_rope_dim = cos_sin.size(2) / 2;
+  const int64_t t_len = mrope_section[0];
+  const int64_t h_len = mrope_section[1];
+  const int64_t w_len = mrope_section[2];
+  const int64_t h_end = t_len + h_len;
+  const int64_t w_end = h_end + w_len;
+
+  torch::Tensor cos_axes =
+      cos_sin.slice(/*dim=*/2, /*start=*/0, /*end=*/half_rope_dim)
+          .to(torch::kFloat32);
+  torch::Tensor sin_axes =
+      cos_sin.slice(/*dim=*/2, /*start=*/half_rope_dim, /*end=*/cos_sin.size(2))
+          .to(torch::kFloat32);
+
+  torch::Tensor cos_rows =
+      torch::zeros({cos_sin.size(1), half_rope_dim},
+                   cos_sin.options().dtype(torch::kFloat32));
+  torch::Tensor sin_rows =
+      torch::zeros({cos_sin.size(1), half_rope_dim},
+                   cos_sin.options().dtype(torch::kFloat32));
+
+  if (t_len > 0) {
+    cos_rows.slice(/*dim=*/1, /*start=*/0, /*end=*/t_len)
+        .copy_(cos_axes.select(/*dim=*/0, /*index=*/0)
+                   .slice(/*dim=*/1, /*start=*/0, /*end=*/t_len));
+    sin_rows.slice(/*dim=*/1, /*start=*/0, /*end=*/t_len)
+        .copy_(sin_axes.select(/*dim=*/0, /*index=*/0)
+                   .slice(/*dim=*/1, /*start=*/0, /*end=*/t_len));
+  }
+  if (h_len > 0) {
+    cos_rows.slice(/*dim=*/1, /*start=*/t_len, /*end=*/h_end)
+        .copy_(cos_axes.select(/*dim=*/0, /*index=*/1)
+                   .slice(/*dim=*/1, /*start=*/t_len, /*end=*/h_end));
+    sin_rows.slice(/*dim=*/1, /*start=*/t_len, /*end=*/h_end)
+        .copy_(sin_axes.select(/*dim=*/0, /*index=*/1)
+                   .slice(/*dim=*/1, /*start=*/t_len, /*end=*/h_end));
+  }
+  if (w_len > 0) {
+    cos_rows.slice(/*dim=*/1, /*start=*/h_end, /*end=*/w_end)
+        .copy_(cos_axes.select(/*dim=*/0, /*index=*/2)
+                   .slice(/*dim=*/1, /*start=*/h_end, /*end=*/w_end));
+    sin_rows.slice(/*dim=*/1, /*start=*/h_end, /*end=*/w_end)
+        .copy_(sin_axes.select(/*dim=*/0, /*index=*/2)
+                   .slice(/*dim=*/1, /*start=*/h_end, /*end=*/w_end));
+  }
+
+  return {cos_rows, sin_rows};
+}
+
+std::pair<torch::Tensor, torch::Tensor> assemble_interleaved_mrope_rows_ref(
+    const torch::Tensor& cos_sin,
+    const std::vector<int64_t>& mrope_section) {
+  const int64_t half_rope_dim = cos_sin.size(2) / 2;
+  const int64_t h_len = mrope_section[1];
+  const int64_t w_len = mrope_section[2];
+
+  torch::Tensor cos_axes =
+      cos_sin.slice(/*dim=*/2, /*start=*/0, /*end=*/half_rope_dim)
+          .to(torch::kFloat32);
+  torch::Tensor sin_axes =
+      cos_sin.slice(/*dim=*/2, /*start=*/half_rope_dim, /*end=*/cos_sin.size(2))
+          .to(torch::kFloat32);
+
+  torch::Tensor cos_rows = cos_axes.select(/*dim=*/0, /*index=*/0).clone();
+  torch::Tensor sin_rows = sin_axes.select(/*dim=*/0, /*index=*/0).clone();
+
+  for (int64_t i = 0; i < half_rope_dim; ++i) {
+    int64_t axis_id = 0;
+    if ((i % 3) == 1 && i < h_len * 3) {
+      axis_id = 1;
+    } else if ((i % 3) == 2 && i < w_len * 3) {
+      axis_id = 2;
+    }
+    cos_rows.slice(/*dim=*/1, /*start=*/i, /*end=*/i + 1)
+        .copy_(cos_axes.select(/*dim=*/0, /*index=*/axis_id)
+                   .slice(/*dim=*/1, /*start=*/i, /*end=*/i + 1));
+    sin_rows.slice(/*dim=*/1, /*start=*/i, /*end=*/i + 1)
+        .copy_(sin_axes.select(/*dim=*/0, /*index=*/axis_id)
+                   .slice(/*dim=*/1, /*start=*/i, /*end=*/i + 1));
+  }
+
+  return {cos_rows, sin_rows};
+}
+
+torch::Tensor merge_cos_sin_for_wrapper(const torch::Tensor& cos_sin) {
+  return cos_sin.permute({1, 0, 2}).contiguous().view(
+      {cos_sin.size(1), cos_sin.size(0) * cos_sin.size(2)});
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+torch_split_qkv_rmsnorm_mrope(const torch::Tensor& qkvg,
+                              const torch::Tensor& q_weight,
+                              const torch::Tensor& k_weight,
+                              const torch::Tensor& cos_sin,
+                              int64_t num_q_heads,
+                              int64_t num_kv_heads,
+                              int64_t head_size,
+                              double eps,
+                              const std::vector<int64_t>& mrope_section,
+                              bool is_interleaved) {
+  CHECK_EQ(mrope_section.size(), 3) << "mrope_section must contain [t, h, w]";
+
+  const int64_t q_size = num_q_heads * head_size;
+  const int64_t kv_size = num_kv_heads * head_size;
+  torch::Tensor q = qkvg.slice(/*dim=*/1, /*start=*/0, /*end=*/q_size)
+                        .view({-1, num_q_heads, head_size});
+  torch::Tensor gate = qkvg.slice(/*dim=*/1,
+                                  /*start=*/q_size,
+                                  /*end=*/q_size * 2)
+                           .view({-1, num_q_heads, head_size});
+  torch::Tensor k =
+      qkvg.slice(/*dim=*/1, /*start=*/q_size * 2, /*end=*/q_size * 2 + kv_size)
+          .view({-1, num_kv_heads, head_size});
+  torch::Tensor v = qkvg.slice(/*dim=*/1,
+                               /*start=*/q_size * 2 + kv_size,
+                               /*end=*/q_size * 2 + kv_size * 2)
+                        .view({-1, num_kv_heads, head_size});
+
+  torch::Tensor q_norm = rms_norm_ref(q, q_weight, eps);
+  torch::Tensor k_norm = rms_norm_ref(k, k_weight, eps);
+  const int64_t half_rope_dim = cos_sin.size(2) / 2;
+  auto [cos_row, sin_row] =
+      is_interleaved
+          ? assemble_interleaved_mrope_rows_ref(cos_sin, mrope_section)
+          : assemble_non_interleaved_mrope_rows_ref(cos_sin, mrope_section);
+
+  auto apply_half_mrope = [&](const torch::Tensor& x) {
+    torch::Tensor out = x.clone();
+    torch::Tensor x1 = out.slice(/*dim=*/2, /*start=*/0, /*end=*/half_rope_dim);
+    torch::Tensor x2 = out.slice(/*dim=*/2,
+                                 /*start=*/half_rope_dim,
+                                 /*end=*/half_rope_dim * 2);
+    torch::Tensor cos_expand = cos_row.unsqueeze(/*dim=*/1);
+    torch::Tensor sin_expand = sin_row.unsqueeze(/*dim=*/1);
+    torch::Tensor out_first = x1 * cos_expand - x2 * sin_expand;
+    torch::Tensor out_second = x2 * cos_expand + x1 * sin_expand;
+    out.slice(/*dim=*/2, /*start=*/0, /*end=*/half_rope_dim).copy_(out_first);
+    out.slice(/*dim=*/2,
+              /*start=*/half_rope_dim,
+              /*end=*/half_rope_dim * 2)
+        .copy_(out_second);
+    return out;
+  };
+
+  torch::Tensor q_out = apply_half_mrope(q_norm).to(qkvg.dtype());
+  torch::Tensor k_out = apply_half_mrope(k_norm).to(qkvg.dtype());
+  return {q_out, k_out, v.contiguous(), gate.contiguous()};
+}
+
+float max_abs_diff(const torch::Tensor& lhs, const torch::Tensor& rhs) {
+  return (lhs.to(torch::kFloat32) - rhs.to(torch::kFloat32))
+      .abs()
+      .max()
+      .item<float>();
+}
+
+void run_case(const SplitQkvTestCase& test_case) {
+  ASSERT_GT(test_case.num_tokens, 0);
+  ASSERT_GT(test_case.num_q_heads, 0);
+  ASSERT_GT(test_case.num_kv_heads, 0);
+
+  const torch::Device device("npu:0");
+  const auto opts =
+      torch::TensorOptions().dtype(torch::kBFloat16).device(device);
+  const std::vector<int64_t> mrope_section = {11, 11, 10};
+  const int64_t rope_dim =
+      2 * (mrope_section[0] + mrope_section[1] + mrope_section[2]);
+  const int64_t q_size = test_case.num_q_heads * kHeadSize;
+  const int64_t kv_size = test_case.num_kv_heads * kHeadSize;
+
+  torch::manual_seed(test_case.seed);
+  torch::Tensor qkv =
+      torch::randn({test_case.num_tokens, q_size * 2 + kv_size * 2}, opts);
+  torch::Tensor q_weight = torch::randn({kHeadSize}, opts);
+  torch::Tensor k_weight = torch::randn({kHeadSize}, opts);
+  torch::Tensor phase = torch::randn(
+      {3, test_case.num_tokens, rope_dim / 2},
+      torch::TensorOptions().dtype(torch::kFloat32).device(device));
+  torch::Tensor cos_sin =
+      torch::cat({torch::cos(phase), torch::sin(phase)}, /*dim=*/2)
+          .to(torch::kBFloat16);
+  torch::Tensor cos_sin_merged = merge_cos_sin_for_wrapper(cos_sin);
+  torch::Tensor gather_pattern = build_split_qkv_rmsnorm_mrope_gather_pattern(
+      rope_dim, mrope_section, test_case.is_interleaved, device);
+
+  auto [q_ref, k_ref, v_ref, gate_ref] =
+      torch_split_qkv_rmsnorm_mrope(qkv,
+                                    q_weight,
+                                    k_weight,
+                                    cos_sin,
+                                    test_case.num_q_heads,
+                                    test_case.num_kv_heads,
+                                    kHeadSize,
+                                    kRmsNormEps,
+                                    mrope_section,
+                                    test_case.is_interleaved);
+  auto [q_out, k_out, v_out, gate_out] =
+      split_qkv_rmsnorm_mrope(qkv,
+                              q_weight,
+                              k_weight,
+                              cos_sin_merged,
+                              gather_pattern,
+                              static_cast<float>(kRmsNormEps),
+                              test_case.num_q_heads,
+                              test_case.num_kv_heads,
+                              kHeadSize);
+
+  EXPECT_EQ(q_out.dim(), 3);
+  EXPECT_EQ(k_out.dim(), 3);
+  EXPECT_EQ(v_out.dim(), 3);
+  EXPECT_EQ(gate_out.dim(), 3);
+  EXPECT_EQ(q_out.size(0), test_case.num_tokens);
+  EXPECT_EQ(q_out.size(1), test_case.num_q_heads);
+  EXPECT_EQ(q_out.size(2), kHeadSize);
+  EXPECT_EQ(k_out.size(0), test_case.num_tokens);
+  EXPECT_EQ(k_out.size(1), test_case.num_kv_heads);
+  EXPECT_EQ(k_out.size(2), kHeadSize);
+  EXPECT_EQ(v_out.size(0), test_case.num_tokens);
+  EXPECT_EQ(v_out.size(1), test_case.num_kv_heads);
+  EXPECT_EQ(v_out.size(2), kHeadSize);
+  EXPECT_EQ(gate_out.size(0), test_case.num_tokens);
+  EXPECT_EQ(gate_out.size(1), test_case.num_q_heads);
+  EXPECT_EQ(gate_out.size(2), kHeadSize);
+
+  aclrtStream stream = c10_npu::getCurrentNPUStream(device.index()).stream();
+  aclrtSynchronizeStream(stream);
+
+  EXPECT_TRUE(torch::allclose(q_out, q_ref, /*rtol=*/1e-2, /*atol=*/1e-2))
+      << "q mismatch, max_diff=" << max_abs_diff(q_out, q_ref);
+  EXPECT_TRUE(torch::allclose(k_out, k_ref, /*rtol=*/1e-2, /*atol=*/1e-2))
+      << "k mismatch, max_diff=" << max_abs_diff(k_out, k_ref);
+  EXPECT_TRUE(torch::allclose(v_out, v_ref, /*rtol=*/0, /*atol=*/0))
+      << "v mismatch, max_diff=" << max_abs_diff(v_out, v_ref);
+  EXPECT_TRUE(torch::allclose(gate_out, gate_ref, /*rtol=*/0, /*atol=*/0))
+      << "gate mismatch, max_diff=" << max_abs_diff(gate_out, gate_ref);
+}
+
+TEST_F(TileLangSplitQkvRmsnormMRopeWrapperTest,
+       BuildGatherPatternMatchesReference) {
+  const torch::Device device("npu:0");
+  const std::vector<int64_t> mrope_section = {11, 11, 10};
+  const int64_t rope_dim = 64;
+
+  torch::Tensor non_interleaved =
+      build_split_qkv_rmsnorm_mrope_gather_pattern(
+          rope_dim, mrope_section, /*is_interleaved=*/false, device)
+          .cpu()
+          .view(torch::kInt32);
+  EXPECT_EQ(non_interleaved.dtype(), torch::kInt32);
+  EXPECT_EQ(non_interleaved.size(0), 256);
+  EXPECT_EQ(non_interleaved[0].item<int32_t>(), 0);
+  EXPECT_EQ(non_interleaved[10].item<int32_t>(), 20);
+  EXPECT_EQ(non_interleaved[11].item<int32_t>(), 150);
+  EXPECT_EQ(non_interleaved[21].item<int32_t>(), 170);
+  EXPECT_EQ(non_interleaved[22].item<int32_t>(), 300);
+  EXPECT_EQ(non_interleaved[31].item<int32_t>(), 318);
+  EXPECT_EQ(non_interleaved[32].item<int32_t>(), 64);
+  EXPECT_EQ(non_interleaved[42].item<int32_t>(), 84);
+  EXPECT_EQ(non_interleaved[43].item<int32_t>(), 214);
+  EXPECT_EQ(non_interleaved[53].item<int32_t>(), 234);
+  EXPECT_EQ(non_interleaved[54].item<int32_t>(), 364);
+  EXPECT_EQ(non_interleaved[63].item<int32_t>(), 382);
+  EXPECT_EQ(non_interleaved[64].item<int32_t>(), 0);
+  EXPECT_EQ(non_interleaved[255].item<int32_t>(), 0);
+
+  torch::Tensor interleaved =
+      build_split_qkv_rmsnorm_mrope_gather_pattern(
+          rope_dim, mrope_section, /*is_interleaved=*/true, device)
+          .cpu()
+          .view(torch::kInt32);
+  EXPECT_EQ(interleaved.dtype(), torch::kInt32);
+  EXPECT_EQ(interleaved.size(0), 256);
+  EXPECT_EQ(interleaved[0].item<int32_t>(), 0);
+  EXPECT_EQ(interleaved[1].item<int32_t>(), 130);
+  EXPECT_EQ(interleaved[2].item<int32_t>(), 260);
+  EXPECT_EQ(interleaved[3].item<int32_t>(), 6);
+  EXPECT_EQ(interleaved[4].item<int32_t>(), 136);
+  EXPECT_EQ(interleaved[5].item<int32_t>(), 266);
+  EXPECT_EQ(interleaved[30].item<int32_t>(), 60);
+  EXPECT_EQ(interleaved[31].item<int32_t>(), 190);
+  EXPECT_EQ(interleaved[32].item<int32_t>(), 64);
+  EXPECT_EQ(interleaved[33].item<int32_t>(), 194);
+  EXPECT_EQ(interleaved[34].item<int32_t>(), 324);
+  EXPECT_EQ(interleaved[35].item<int32_t>(), 70);
+  EXPECT_EQ(interleaved[63].item<int32_t>(), 254);
+  EXPECT_EQ(interleaved[64].item<int32_t>(), 0);
+  EXPECT_EQ(interleaved[255].item<int32_t>(), 0);
+}
+
+TEST_F(TileLangSplitQkvRmsnormMRopeWrapperTest, MatchesTorchReference) {
+  const std::vector<SplitQkvTestCase> cases = {
+      {.name = "tiny_t1_q16_kv4",
+       .num_tokens = 1,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 101,
+       .is_interleaved = false},
+      {.name = "medium_t17_q16_kv4",
+       .num_tokens = 17,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 102,
+       .is_interleaved = false},
+      {.name = "chunked_t4097_q16_kv4",
+       .num_tokens = 4097,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 103,
+       .is_interleaved = false},
+      {.name = "tiny_t1_q16_kv2",
+       .num_tokens = 1,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 201,
+       .is_interleaved = false},
+      {.name = "medium_t17_q16_kv2",
+       .num_tokens = 17,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 202,
+       .is_interleaved = false},
+      {.name = "chunked_t4097_q16_kv2",
+       .num_tokens = 4097,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 203,
+       .is_interleaved = false},
+      {.name = "tiny_t1_q8_kv1",
+       .num_tokens = 1,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 301,
+       .is_interleaved = false},
+      {.name = "medium_t17_q8_kv1",
+       .num_tokens = 17,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 302,
+       .is_interleaved = false},
+      {.name = "chunked_t4097_q8_kv1",
+       .num_tokens = 4097,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 303,
+       .is_interleaved = false},
+      {.name = "tiny_t1_q16_kv4_interleaved",
+       .num_tokens = 1,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 401,
+       .is_interleaved = true},
+      {.name = "medium_t17_q16_kv4_interleaved",
+       .num_tokens = 17,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 402,
+       .is_interleaved = true},
+      {.name = "chunked_t4097_q16_kv4_interleaved",
+       .num_tokens = 4097,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 403,
+       .is_interleaved = true},
+      {.name = "tiny_t1_q16_kv2_interleaved",
+       .num_tokens = 1,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 501,
+       .is_interleaved = true},
+      {.name = "medium_t17_q16_kv2_interleaved",
+       .num_tokens = 17,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 502,
+       .is_interleaved = true},
+      {.name = "chunked_t4097_q16_kv2_interleaved",
+       .num_tokens = 4097,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 503,
+       .is_interleaved = true},
+      {.name = "tiny_t1_q8_kv1_interleaved",
+       .num_tokens = 1,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 601,
+       .is_interleaved = true},
+      {.name = "medium_t17_q8_kv1_interleaved",
+       .num_tokens = 17,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 602,
+       .is_interleaved = true},
+      {.name = "chunked_t4097_q8_kv1_interleaved",
+       .num_tokens = 4097,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 603,
+       .is_interleaved = true},
+      // Odd num_tokens
+      {.name = "odd_t3_q16_kv4",
+       .num_tokens = 3,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 701,
+       .is_interleaved = false},
+      {.name = "odd_t7_q16_kv2_interleaved",
+       .num_tokens = 7,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 702,
+       .is_interleaved = true},
+      {.name = "odd_t15_q8_kv1",
+       .num_tokens = 15,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 703,
+       .is_interleaved = false},
+      {.name = "odd_t33_q16_kv4_interleaved",
+       .num_tokens = 33,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 704,
+       .is_interleaved = true},
+      // Boundary: min specialization bucket
+      {.name = "boundary_t2_q16_kv2",
+       .num_tokens = 2,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 801,
+       .is_interleaved = false},
+      // Boundary: vec_core_num (48) and neighbors
+      {.name = "boundary_t48_q8_kv1_interleaved",
+       .num_tokens = 48,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 802,
+       .is_interleaved = true},
+      {.name = "boundary_t49_q16_kv4",
+       .num_tokens = 49,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 803,
+       .is_interleaved = false},
+      // Large: old 4096 boundary
+      {.name = "large_t4096_q16_kv2",
+       .num_tokens = 4096,
+       .num_q_heads = 16,
+       .num_kv_heads = 2,
+       .seed = 901,
+       .is_interleaved = false},
+      // Large: odd, exceeding old chunking limit
+      {.name = "large_t4099_q8_kv1_interleaved",
+       .num_tokens = 4099,
+       .num_q_heads = 8,
+       .num_kv_heads = 1,
+       .seed = 902,
+       .is_interleaved = true},
+      // Large batch
+      {.name = "large_t8192_q16_kv4",
+       .num_tokens = 8192,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 903,
+       .is_interleaved = false},
+      {.name = "large_t8193_q16_kv4_interleaved",
+       .num_tokens = 8193,
+       .num_q_heads = 16,
+       .num_kv_heads = 4,
+       .seed = 904,
+       .is_interleaved = true},
+      // Qwen3.5-0.8B/2B tp=1, Qwen3.5-4B/9B tp=2
+      {.name = "tiny_t1_q8_kv2",
+       .num_tokens = 1,
+       .num_q_heads = 8,
+       .num_kv_heads = 2,
+       .seed = 1001,
+       .is_interleaved = false},
+      {.name = "large_t4097_q8_kv2",
+       .num_tokens = 4097,
+       .num_q_heads = 8,
+       .num_kv_heads = 2,
+       .seed = 1002,
+       .is_interleaved = false},
+      // Qwen3.5-27B tp=1
+      {.name = "tiny_t1_q24_kv4",
+       .num_tokens = 1,
+       .num_q_heads = 24,
+       .num_kv_heads = 4,
+       .seed = 1101,
+       .is_interleaved = false},
+      {.name = "large_t4097_q24_kv4",
+       .num_tokens = 4097,
+       .num_q_heads = 24,
+       .num_kv_heads = 4,
+       .seed = 1102,
+       .is_interleaved = true},
+      // Qwen3.5-27B tp=2
+      {.name = "tiny_t1_q12_kv2",
+       .num_tokens = 1,
+       .num_q_heads = 12,
+       .num_kv_heads = 2,
+       .seed = 1201,
+       .is_interleaved = false},
+      {.name = "large_t4097_q12_kv2",
+       .num_tokens = 4097,
+       .num_q_heads = 12,
+       .num_kv_heads = 2,
+       .seed = 1202,
+       .is_interleaved = true},
+      // Qwen3.5-27B tp=4
+      {.name = "tiny_t1_q6_kv1",
+       .num_tokens = 1,
+       .num_q_heads = 6,
+       .num_kv_heads = 1,
+       .seed = 1301,
+       .is_interleaved = false},
+      {.name = "large_t4097_q6_kv1",
+       .num_tokens = 4097,
+       .num_q_heads = 6,
+       .num_kv_heads = 1,
+       .seed = 1302,
+       .is_interleaved = true},
+      // Qwen3.5-27B tp=8
+      {.name = "tiny_t1_q3_kv1",
+       .num_tokens = 1,
+       .num_q_heads = 3,
+       .num_kv_heads = 1,
+       .seed = 1401,
+       .is_interleaved = false},
+      {.name = "large_t4097_q3_kv1",
+       .num_tokens = 4097,
+       .num_q_heads = 3,
+       .num_kv_heads = 1,
+       .seed = 1402,
+       .is_interleaved = true},
+      // Qwen3.5-35B tp=4, Qwen3.5-4B tp=4
+      {.name = "tiny_t1_q4_kv1",
+       .num_tokens = 1,
+       .num_q_heads = 4,
+       .num_kv_heads = 1,
+       .seed = 1501,
+       .is_interleaved = false},
+      {.name = "large_t4097_q4_kv1",
+       .num_tokens = 4097,
+       .num_q_heads = 4,
+       .num_kv_heads = 1,
+       .seed = 1502,
+       .is_interleaved = true},
+      // Qwen3.5-122B/397B tp=1
+      {.name = "tiny_t1_q32_kv2",
+       .num_tokens = 1,
+       .num_q_heads = 32,
+       .num_kv_heads = 2,
+       .seed = 1601,
+       .is_interleaved = false},
+      {.name = "large_t4097_q32_kv2",
+       .num_tokens = 4097,
+       .num_q_heads = 32,
+       .num_kv_heads = 2,
+       .seed = 1602,
+       .is_interleaved = true},
+      // Qwen3.5-122B/397B tp=2
+      {.name = "tiny_t1_q16_kv1",
+       .num_tokens = 1,
+       .num_q_heads = 16,
+       .num_kv_heads = 1,
+       .seed = 1701,
+       .is_interleaved = false},
+      {.name = "large_t4097_q16_kv1",
+       .num_tokens = 4097,
+       .num_q_heads = 16,
+       .num_kv_heads = 1,
+       .seed = 1702,
+       .is_interleaved = true},
+      // Qwen3.5-0.8B/2B tp=4, various tp=8/16
+      {.name = "tiny_t1_q2_kv1",
+       .num_tokens = 1,
+       .num_q_heads = 2,
+       .num_kv_heads = 1,
+       .seed = 1801,
+       .is_interleaved = false},
+      {.name = "large_t4097_q2_kv1",
+       .num_tokens = 4097,
+       .num_q_heads = 2,
+       .num_kv_heads = 1,
+       .seed = 1802,
+       .is_interleaved = true},
+  };
+
+  for (const SplitQkvTestCase& test_case : cases) {
+    SCOPED_TRACE(test_case.name);
+    run_case(test_case);
+  }
+}
+
+}  // namespace
+}  // namespace xllm::kernel::npu::tilelang

upstream_ref/xllm/tests/core/layers/CMakeLists.txt

@@ -0,0 +1,7 @@
+if(USE_CUDA)
+  add_subdirectory(cuda)
+endif()
+
+if(USE_MLU)
+  add_subdirectory(mlu)
+endif()

upstream_ref/xllm/tests/core/layers/cuda/CMakeLists.txt

@@ -0,0 +1,19 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    xattention_test
+  SRCS
+    xattention_test.cpp
+  DEPS
+    :cuda_layers
+    :common_layers
+    torch
+    GTest::gtest_main
+    glog::glog
+    :block
+)
+target_link_libraries(xattention_test
+                      PRIVATE
+                      brpc
+                      "$<LINK_GROUP:RESCAN,cuda_layers,common_layers>")

upstream_ref/xllm/tests/core/layers/cuda/xattention_test.cpp

@@ -0,0 +1,276 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/cuda/xattention.h"
+
+#include <gtest/gtest.h>
+#include <torch/cuda.h>
+#include <torch/torch.h>
+
+#include <cmath>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "core/common/global_flags.h"
+#include "framework/kv_cache/kv_cache.h"
+#include "layers/cuda/flashinfer_workspace.h"
+#include "layers/cuda/xattention_workspace.h"
+
+namespace xllm::layer::test {
+namespace {
+
+struct DecodeTestInput {
+  AttentionMetadata attn_metadata;
+  torch::Tensor query;
+  torch::Tensor key;
+  torch::Tensor value;
+};
+
+class XAttentionDecodeCompareTest : public ::testing::Test {
+ protected:
+  static constexpr int32_t kBatchSize = 4;
+  static constexpr int32_t kBeamWidth = 128;
+  static constexpr int32_t kNumHeads = 16;
+  static constexpr int32_t kNumKvHeads = 8;
+  static constexpr int32_t kHeadDim = 128;
+  static constexpr int32_t kSharedSeqLen = 300;
+  static constexpr int32_t kMaxDecodeStep = 2;
+  static constexpr int32_t kCurrentStep = 1;
+
+  void SetUp() override {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "CUDA is not available.";
+    }
+
+    device_ = torch::Device(torch::kCUDA, 0);
+    flashinfer::FlashinferWorkspace::get_instance().initialize(device_);
+    xattention::XAttentionWorkspace::get_instance().initialize(device_);
+  }
+
+  DecodeTestInput create_decode_test_input(torch::ScalarType dtype) {
+    const int32_t total_beam = kBatchSize * kBeamWidth;
+    const int32_t unshared_offset = kSharedSeqLen;
+    const int32_t full_kv_len = unshared_offset + total_beam * kMaxDecodeStep;
+    const int32_t kv_len_per_beam = kSharedSeqLen + (kCurrentStep + 1);
+
+    auto float_opts = torch::TensorOptions().dtype(dtype).device(device_);
+    auto int_opts = torch::TensorOptions().dtype(torch::kInt32).device(device_);
+
+    DecodeTestInput input;
+    auto& meta = input.attn_metadata;
+
+    meta.is_prefill = false;
+    meta.is_chunked_prefill = false;
+    meta.is_dummy = false;
+    meta.is_causal = false;
+    meta.max_query_len = 1;
+    meta.max_seq_len = kv_len_per_beam;
+    meta.compute_dtype = dtype == torch::kBFloat16 ? "bfloat16" : "half";
+    meta.enable_cuda_graph = false;
+
+    meta.plan_info = std::make_shared<PlanInfo>();
+    meta.plan_info->layer_id = 0;
+    meta.shared_plan_info = std::make_shared<PlanInfo>();
+    meta.shared_plan_info->layer_id = 0;
+    meta.unshared_plan_info = std::make_shared<PlanInfo>();
+    meta.unshared_plan_info->layer_id = 0;
+
+    meta.full_k_cache =
+        torch::zeros({full_kv_len, kNumKvHeads, kHeadDim}, float_opts);
+    meta.full_v_cache =
+        torch::zeros({full_kv_len, kNumKvHeads, kHeadDim}, float_opts);
+
+    meta.full_k_cache.slice(0, 0, kSharedSeqLen)
+        .copy_(
+            torch::randn({kSharedSeqLen, kNumKvHeads, kHeadDim}, float_opts) *
+            0.001);
+    meta.full_v_cache.slice(0, 0, kSharedSeqLen)
+        .copy_(
+            torch::randn({kSharedSeqLen, kNumKvHeads, kHeadDim}, float_opts) *
+            0.001);
+
+    meta.unshared_k_cache =
+        meta.full_k_cache.slice(0, unshared_offset, full_kv_len)
+            .view({kBatchSize,
+                   kBeamWidth,
+                   kMaxDecodeStep,
+                   kNumKvHeads,
+                   kHeadDim});
+    meta.unshared_v_cache =
+        meta.full_v_cache.slice(0, unshared_offset, full_kv_len)
+            .view({kBatchSize,
+                   kBeamWidth,
+                   kMaxDecodeStep,
+                   kNumKvHeads,
+                   kHeadDim});
+
+    meta.unshared_k_cache.slice(2, 0, 1).copy_(
+        torch::randn({kBatchSize, kBeamWidth, 1, kNumKvHeads, kHeadDim},
+                     float_opts) *
+        0.001);
+    meta.unshared_v_cache.slice(2, 0, 1).copy_(
+        torch::randn({kBatchSize, kBeamWidth, 1, kNumKvHeads, kHeadDim},
+                     float_opts) *
+        0.001);
+
+    meta.block_table =
+        torch::arange(total_beam, int_opts).view({total_beam, 1});
+    meta.step_tensor = torch::tensor({kCurrentStep}, int_opts);
+
+    std::vector<int32_t> paged_kv_indptr(total_beam + 1, 0);
+    std::vector<int32_t> paged_kv_indices;
+    paged_kv_indices.reserve(total_beam * kv_len_per_beam);
+
+    int32_t cursor = 0;
+    for (int32_t beam_id = 0; beam_id < total_beam; ++beam_id) {
+      paged_kv_indptr[beam_id] = cursor;
+      for (int32_t t = 0; t < kSharedSeqLen; ++t) {
+        paged_kv_indices.push_back(t);
+        ++cursor;
+      }
+      for (int32_t s = 0; s <= kCurrentStep; ++s) {
+        paged_kv_indices.push_back(unshared_offset + beam_id * kMaxDecodeStep +
+                                   s);
+        ++cursor;
+      }
+      paged_kv_indptr[beam_id + 1] = cursor;
+    }
+
+    meta.paged_kv_indptr = torch::tensor(paged_kv_indptr, int_opts);
+    meta.paged_kv_indices = torch::tensor(paged_kv_indices, int_opts);
+    meta.paged_kv_last_page_len = torch::ones({total_beam}, int_opts);
+
+    meta.kv_seq_lens = torch::full({total_beam}, kv_len_per_beam, int_opts);
+    meta.q_seq_lens = torch::ones({total_beam}, int_opts);
+    meta.q_cu_seq_lens = torch::arange(0, total_beam + 1, 1, int_opts);
+    std::vector<int32_t> kv_cu_seq_lens(kBatchSize + 1, 0);
+    for (int32_t i = 1; i <= kBatchSize; ++i) {
+      kv_cu_seq_lens[i] = i * kSharedSeqLen;
+    }
+    meta.kv_cu_seq_lens = torch::tensor(kv_cu_seq_lens, int_opts);
+    meta.qo_indptr = torch::arange(0, total_beam + 1, 1, int_opts);
+
+    XAttentionTwoStageDecodeCache two_stage_cache;
+    auto fp32_opts =
+        torch::TensorOptions().dtype(torch::kFloat32).device(device_);
+    two_stage_cache.shared_lse =
+        torch::zeros({total_beam, kNumHeads, 1}, fp32_opts);
+    two_stage_cache.shared_o =
+        torch::zeros({total_beam, kNumHeads, kHeadDim}, float_opts);
+    two_stage_cache.unshared_lse =
+        torch::zeros({total_beam, kNumHeads, 1}, fp32_opts);
+    two_stage_cache.unshared_o =
+        torch::zeros({total_beam, kNumHeads, kHeadDim}, float_opts);
+
+    two_stage_cache.q_cu_seq_lens_shared =
+        torch::arange(0, (kBatchSize + 1) * kBeamWidth, kBeamWidth, int_opts);
+    two_stage_cache.paged_kv_indptr_expanded =
+        torch::arange(total_beam + 1, int_opts);
+    two_stage_cache.paged_kv_indices_expanded =
+        torch::arange(total_beam, int_opts);
+    two_stage_cache.paged_kv_last_page_len_expanded =
+        torch::full({total_beam}, kCurrentStep + 1, int_opts);
+
+    meta.xattention_two_stage_decode_cache = std::move(two_stage_cache);
+
+    input.query =
+        torch::randn({total_beam, kNumHeads * kHeadDim}, float_opts) * 0.001;
+    input.key =
+        torch::randn({total_beam, kNumKvHeads * kHeadDim}, float_opts) * 0.001;
+    input.value =
+        torch::randn({total_beam, kNumKvHeads * kHeadDim}, float_opts) * 0.001;
+
+    return input;
+  }
+
+  torch::Tensor run_decode_once(DecodeTestInput& input, bool enable_two_stage) {
+    FLAGS_enable_xattention_one_stage = !enable_two_stage;
+    FLAGS_max_tokens_per_batch = kSharedSeqLen;
+
+    XAttentionImpl attention(
+        /*num_heads=*/kNumHeads,
+        /*head_size=*/kHeadDim,
+        /*scale=*/1.0f / std::sqrt(static_cast<float>(kHeadDim)),
+        /*num_kv_heads=*/kNumKvHeads,
+        /*sliding_window=*/-1);
+
+    torch::Tensor output = torch::zeros_like(input.query);
+    KVCache dummy_cache;
+
+    auto result = attention.forward(input.attn_metadata,
+                                    input.query,
+                                    input.key,
+                                    input.value,
+                                    output,
+                                    dummy_cache);
+
+    torch::cuda::synchronize();
+    return std::get<0>(result).clone();
+  }
+
+  void compare_single_and_two_stage(torch::ScalarType dtype,
+                                    double atol,
+                                    double rtol) {
+    constexpr int64_t kSeed = 20260303;
+    torch::manual_seed(kSeed);
+    torch::cuda::manual_seed_all(kSeed);
+    auto single_input = create_decode_test_input(dtype);
+    torch::manual_seed(kSeed);
+    torch::cuda::manual_seed_all(kSeed);
+    auto two_stage_input = create_decode_test_input(dtype);
+
+    two_stage_input.query.copy_(single_input.query);
+    two_stage_input.key.copy_(single_input.key);
+    two_stage_input.value.copy_(single_input.value);
+
+    auto single_output =
+        run_decode_once(single_input, /*enable_two_stage=*/false);
+    auto two_stage_output =
+        run_decode_once(two_stage_input, /*enable_two_stage=*/true);
+
+    auto abs_diff =
+        (single_output - two_stage_output).abs().to(torch::kFloat32);
+    const double max_abs_diff = abs_diff.max().item<double>();
+    const double mean_abs_diff = abs_diff.mean().item<double>();
+
+    EXPECT_TRUE(torch::allclose(single_output, two_stage_output, rtol, atol))
+        << "single-stage and two-stage decode outputs mismatch: "
+        << "max_abs_diff=" << max_abs_diff
+        << ", mean_abs_diff=" << mean_abs_diff << ", atol=" << atol
+        << ", rtol=" << rtol;
+
+    EXPECT_LT(max_abs_diff, atol)
+        << "max_abs_diff exceeds threshold: "
+        << "max_abs_diff=" << max_abs_diff << ", threshold=" << atol;
+  }
+
+  torch::Device device_{torch::kCPU};
+};
+
+TEST_F(XAttentionDecodeCompareTest, SingleVsTwoStageFp16) {
+  compare_single_and_two_stage(torch::kFloat16,
+                               /*atol=*/2e-3,
+                               /*rtol=*/2e-3);
+}
+
+TEST_F(XAttentionDecodeCompareTest, SingleVsTwoStageBf16) {
+  compare_single_and_two_stage(torch::kBFloat16,
+                               /*atol=*/2e-2,
+                               /*rtol=*/2e-2);
+}
+
+}  // namespace
+}  // namespace xllm::layer::test

upstream_ref/xllm/tests/core/layers/mlu/CMakeLists.txt

@@ -0,0 +1,106 @@
+include(cc_test)
+
+# Add test for MLU layers
+cc_test(
+  NAME
+    layer_test
+  SRCS
+    dense_mlp_test.cpp
+    tests_utils.cpp
+    deepseek_v32_sp_utils_test.cpp
+    dp_utils_test.cpp
+    indexer_test.cpp
+    mla_test.cpp
+    deepseek_v2_decoder_layer_test.cpp
+    qwen2_attention_test.cpp
+    qwen2_vision_attention_test.cpp
+  DEPS
+    :common_layers
+    :mlu_layers
+    :block
+    :parallel_state
+    :model
+    :model_context
+    :state_dict
+    glog::glog
+    torch
+    GTest::gtest_main
+)
+
+# MoE Gate unit test for MLU (uses layers/mlu/moe_gate)
+cc_test(
+  NAME
+    moe_layer_test
+  SRCS
+    moe_gate_test.cpp
+    fused_moe_test.cpp
+    deepseek_v2_sparse_moe_block_test.cpp
+    tests_utils.cpp
+  DEPS
+    :common_layers
+    :mlu_layers
+    :parallel_state
+    :model
+    :model_context
+    :state_dict
+    glog::glog
+    torch
+    GTest::gtest_main
+)
+
+# Add test for DeepEP
+# This test must exist individually, because it contains forked processes
+#  which does not allow any device init on main process
+cc_test(
+  NAME
+    deep_ep_test
+  SRCS
+    deep_ep_test.cpp
+    tests_utils.cpp
+  DEPS
+    :common_layers
+    :parallel_state
+    :model
+    :model_context
+    :state_dict
+    GTest::gtest_main
+    torch
+    glog::glog
+)
+
+# Add test for FusedMoE All2All path
+# This test must exist individually, because it contains forked processes
+#  which does not allow any device init on main process
+cc_test(
+  NAME
+    fused_moe_all2all_test
+  SRCS
+    fused_moe_all2all_test.cpp
+    tests_utils.cpp
+  DEPS
+    :common_layers
+    :mlu_layers
+    :parallel_state
+    :model
+    :model_context
+    :state_dict
+    GTest::gtest_main
+    torch
+    glog::glog
+)
+
+cc_test(
+  NAME
+    deepseek_v2_attention_multi_device_test
+  SRCS
+    deepseek_v2_attention_multi_device_test.cpp
+    tests_utils.cpp
+  DEPS
+    :block
+    :mlu_layers
+    :parallel_state
+    :state_dict
+    GTest::gtest_main
+    torch
+    glog::glog
+)

upstream_ref/xllm/tests/core/layers/mlu/deep_ep_test.cpp

@@ -0,0 +1,302 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/common/deep_ep.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <sys/wait.h>
+#include <torch/torch.h>
+#include <unistd.h>
+
+#include <cmath>
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include "framework/parallel_state/parallel_args.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+#include "util/tensor_helper.h"
+
+#if defined(USE_MLU)
+#include "framework/parallel_state/mlu_process_group.h"
+#elif defined(USE_CUDA)
+#include "framework/parallel_state/cuda_process_group.h"
+#endif
+
+namespace xllm {
+namespace layer {
+namespace test {
+
+// Special exit code definition for skipping test
+constexpr int32_t EXIT_CODE_SKIP = 77;
+
+// Helper function to create ProcessGroup
+std::unique_ptr<xllm::ProcessGroup> create_test_process_group(
+    int32_t rank,
+    int32_t world_size,
+    int32_t port,
+    const std::string& host,
+    const torch::Device& device) {
+  return xllm::create_process_group(rank,
+                                    world_size,
+                                    world_size,
+                                    port,
+                                    false,
+                                    host,
+                                    "deep_ep_test_group",
+                                    device);
+}
+
+struct TestParams {
+  int32_t rank;
+  int32_t world_size;
+  int32_t port;
+  std::string host;
+  int32_t device_index;
+  int64_t token_size;
+  int64_t max_tokens;
+  int64_t num_experts;
+};
+
+// Child process test function
+int32_t run_deep_ep_test_child(TestParams params) {
+  try {
+    // 1. Check devices
+    int32_t dev_count = xllm::Device::device_count();
+    if (dev_count < params.world_size) {
+      LOG(WARNING) << "Rank " << params.rank
+                   << ": Insufficient devices. Skipping.";
+      return EXIT_CODE_SKIP;
+    }
+    params.device_index = params.rank % dev_count;
+
+    // 2. Set device
+    xllm::Device xllm_device(params.device_index);
+    xllm_device.set_device();
+    torch::Device device = xllm_device.unwrap();
+
+    // 3. Create ProcessGroup
+    auto process_group = create_test_process_group(
+        params.rank, params.world_size, params.port, params.host, device);
+
+    CHECK(process_group) << "Rank " << params.rank
+                         << ": Failed to create ProcessGroup";
+
+    // 4. Create ParallelArgs
+    ParallelArgs parallel_args(
+        params.rank, params.world_size, process_group.get());
+    parallel_args.moe_ep_group_ = process_group.get();
+    parallel_args.ep_size_ = params.world_size;
+
+    auto options = torch::TensorOptions()
+                       .dtype(torch::kFloat32)
+                       .device(device)
+                       .requires_grad(false);
+
+    // 5. Initialize DeepEPImpl
+    DeepEPImpl deep_ep(params.token_size,
+                       params.token_size,
+                       params.max_tokens,
+                       params.num_experts,
+                       parallel_args,
+                       options);
+
+    LOG(INFO) << "Rank " << params.rank << ": DeepEP created successfully";
+
+    const int64_t hidden_dim = params.token_size / sizeof(float);  // 32
+    const int64_t num_tokens_sent = 4;
+
+    // Prepare Layout
+    torch::Tensor token_count_slice = torch::zeros(
+        {params.num_experts},
+        torch::TensorOptions().dtype(torch::kInt32).device(device));
+
+    float base_val = 0.0f;
+
+    if (params.rank == 0) {
+      token_count_slice[0] = 2;
+      token_count_slice[2] = 2;
+      base_val = 10.0f;
+    } else {
+      token_count_slice[1] = 2;
+      token_count_slice[3] = 2;
+      base_val = 20.0f;
+    }
+
+    // Step 1: Prepare Input Data
+    auto buffers = deep_ep.get_buffer();
+
+    torch::Tensor src_data = torch::zeros(
+        {num_tokens_sent, hidden_dim},
+        torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU));
+    for (size_t i = 0; i < num_tokens_sent; ++i) {
+      src_data[i].fill_(base_val + i);
+    }
+
+    torch::Tensor src_data_bytes =
+        src_data.to(device).view({-1}).view(torch::kUInt8);
+    CHECK(src_data_bytes.numel() <= buffers.dispatch_send_token_tensor.numel());
+    buffers.dispatch_send_token_tensor.view({-1})
+        .slice(0, 0, src_data_bytes.numel())
+        .copy_(src_data_bytes);
+
+    // Step 2: Dispatch Step
+    deep_ep.dispatch_step(num_tokens_sent, token_count_slice);
+
+    // Step 3: Process Dispatch Result
+    int64_t num_experts_per_rank = params.num_experts / params.world_size;
+    int64_t max_recv_capacity = params.max_tokens * params.world_size;
+
+    // Allocate Expert Input as Float
+    torch::Tensor expert_input =
+        torch::zeros({max_recv_capacity, hidden_dim}, options);
+
+    // Convert to Int8 View to match DeepEP internal buffer dtype
+    // DeepEP buffer is Int8, so we must pass a Int8 tensor to gather_split
+    torch::Tensor expert_input_byte = view_as_dtype(expert_input, torch::kInt8);
+
+    auto meta = deep_ep.process_dispatch_result(
+        num_experts_per_rank, expert_input_byte, std::nullopt);
+
+    // Validate count
+    torch::Tensor valid_token_num_cpu = meta.token_sum.to(torch::kCPU);
+    int64_t received_count = valid_token_num_cpu.item<int64_t>();
+    CHECK_EQ(received_count, 4)
+        << "Rank " << params.rank << ": Expected 4 received tokens, got "
+        << received_count;
+
+    // Step 4: Simulate Computation (on Float tensor)
+    torch::Tensor expert_output = expert_input + 1.0f;
+
+    // Step 5: Combine Pack
+    // Convert to Int8 View to match DeepEP internal buffer dtype
+    torch::Tensor expert_output_byte =
+        view_as_dtype(expert_output, torch::kInt8);
+
+    torch::Tensor combine_layout =
+        deep_ep.combine_step_pack(expert_output_byte,
+                                  meta.gather_rank_index,
+                                  meta.token_sum,
+                                  params.token_size,
+                                  torch::kInt8);
+
+    // Step 6: Combine Comm
+    torch::Tensor final_output = deep_ep.combine_step_comm(
+        combine_layout, num_tokens_sent, hidden_dim, torch::kFloat32);
+
+    // Step 7: Verification
+    xllm_device.synchronize_default_stream();
+
+    torch::Tensor final_cpu = final_output.to(torch::kCPU);
+    auto final_acc = final_cpu.accessor<float, 2>();
+
+    for (int64_t i = 0; i < num_tokens_sent; ++i) {
+      float expected_val = base_val + i + 1.0f;
+      for (int64_t j = 0; j < hidden_dim; ++j) {
+        float got = final_acc[i][j];
+        CHECK(std::abs(got - expected_val) <= 1e-4)
+            << "Rank " << params.rank << " Verification Failed at [" << i
+            << "][" << j << "]: Expected " << expected_val << ", Got " << got;
+      }
+    }
+
+    LOG(INFO) << "Rank " << params.rank << ": E2E Test Passed!";
+    return 0;
+
+  } catch (const std::exception& e) {
+    LOG(ERROR) << "Rank " << params.rank << ": Exception: " << e.what();
+    return 1;
+  }
+}
+
+// Multi-process test fixture
+class DeepEPMultiDeviceTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    world_size_ = 2;
+    port_ = 29500;
+    host_ = "127.0.0.1";
+    token_size_ = 128;
+    max_tokens_ = 64;
+    num_global_experts_ = 4;
+  }
+
+  void run_test() {
+    std::vector<pid_t> child_pids;
+
+    for (int32_t rank = 0; rank < world_size_; ++rank) {
+      pid_t pid = fork();
+      if (pid == 0) {
+        TestParams params;
+        params.rank = rank;
+        params.world_size = world_size_;
+        params.port = port_;
+        params.host = host_;
+        params.device_index = -1;
+        params.token_size = token_size_;
+        params.max_tokens = max_tokens_;
+        params.num_experts = num_global_experts_;
+
+        int32_t exit_code = run_deep_ep_test_child(params);
+        _exit(exit_code);
+      } else if (pid > 0) {
+        child_pids.push_back(pid);
+      } else {
+        LOG(FATAL) << "Failed to fork rank " << rank;
+      }
+    }
+
+    bool any_failed = false;
+    bool any_skipped = false;
+
+    for (size_t i = 0; i < child_pids.size(); ++i) {
+      int32_t status;
+      waitpid(child_pids[i], &status, 0);
+      if (WIFEXITED(status)) {
+        int32_t exit_code = WEXITSTATUS(status);
+        if (exit_code == EXIT_CODE_SKIP) {
+          any_skipped = true;
+        } else if (exit_code != 0) {
+          any_failed = true;
+          LOG(ERROR) << "Rank " << i << " failed with code " << exit_code;
+        }
+      } else {
+        any_failed = true;
+        LOG(ERROR) << "Rank " << i << " crashed (signal).";
+      }
+    }
+
+    if (any_skipped) {
+      GTEST_SKIP() << "Test skipped due to insufficient devices.";
+    } else {
+      ASSERT_FALSE(any_failed) << "DeepEP End-to-End Test Failed.";
+    }
+  }
+
+  int32_t world_size_;
+  int32_t port_;
+  std::string host_;
+  int64_t token_size_;
+  int64_t max_tokens_;
+  int64_t num_global_experts_;
+};
+
+TEST_F(DeepEPMultiDeviceTest, EndToEndFlow) { run_test(); }
+
+}  // namespace test
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/deepseek_v2_sparse_moe_block_test.cpp

@@ -0,0 +1,959 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/mlu/deepseek_v2_sparse_moe_block.h"
+
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/fused_moe.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+class DeepseekV2SparseMoEBlockTestPeer {
+ public:
+  static FusedMoE moe(DeepseekV2SparseMoEBlockImpl& block) {
+    return block.moe_;
+  }
+
+  static ProcessGroup* routed_pg(DeepseekV2SparseMoEBlockImpl& block) {
+    return block.routed_pg();
+  }
+
+  static void set_enable_deep_ep(DeepseekV2SparseMoEBlockImpl& block,
+                                 bool enable) {
+    block.enable_deep_ep_ = enable;
+  }
+};
+
+class DeepseekV2SparseMoEBlockTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    options_ = torch::TensorOptions()
+                   .dtype(torch::kBFloat16)
+                   .device(Device::type_torch(), 0)
+                   .requires_grad(false);
+    model_args_ = test::create_default_model_args();
+    model_args_.hidden_size() = 256;
+    model_args_.moe_intermediate_size() = 256;
+    model_args_.n_routed_experts() = 4;
+    model_args_.num_experts_per_tok() = 2;
+    model_args_.n_group() = 1;
+    model_args_.topk_group() = 2;
+    model_args_.n_shared_experts() = 1;
+    model_args_.routed_scaling_factor() = 1.0f;
+    model_args_.norm_topk_prob() = true;
+    model_args_.hidden_act() = "silu";
+    model_args_.scoring_func() = "softmax";
+    model_args_.topk_method() = "greedy";
+    set_pg_ctx();
+  }
+
+  void set_pg_ctx() {
+    global_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, /*world_size=*/2);
+    tp_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, /*world_size=*/2);
+    single_rank_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, /*world_size=*/1);
+
+    parallel_args_ =
+        ParallelArgs(/*rank=*/0, /*world_size=*/2, global_pg_.get());
+    parallel_args_.process_group_ = global_pg_.get();
+    parallel_args_.tp_group_ = tp_pg_.get();
+    parallel_args_.single_rank_group_ = single_rank_pg_.get();
+    parallel_args_.sp_group_ = tp_pg_.get();
+    parallel_args_.ep_size_ = 1;
+  }
+
+  void set_tp_ctx(int64_t world_size, int64_t ep_size) {
+    global_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, world_size);
+    tp_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, world_size);
+    single_rank_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, /*world_size=*/1);
+    dp_pg_.reset();
+
+    parallel_args_ =
+        ParallelArgs(/*rank=*/0, world_size, /*dp_size=*/1, global_pg_.get());
+    parallel_args_.process_group_ = global_pg_.get();
+    parallel_args_.tp_group_ = tp_pg_.get();
+    parallel_args_.single_rank_group_ = single_rank_pg_.get();
+    parallel_args_.sp_group_ = tp_pg_.get();
+    parallel_args_.dp_local_process_group_ = nullptr;
+    parallel_args_.ep_size_ = ep_size;
+    parallel_args_.moe_ep_group_ = global_pg_.get();
+    parallel_args_.moe_tp_group_ = tp_pg_.get();
+  }
+
+  void set_tp_dp_ctx(int64_t world_size,
+                     int64_t dp_size,
+                     int64_t tp_size,
+                     int64_t ep_size) {
+    global_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, world_size);
+    dp_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, dp_size);
+    tp_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, tp_size);
+    single_rank_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, /*world_size=*/1);
+
+    parallel_args_ =
+        ParallelArgs(/*rank=*/0, world_size, dp_size, global_pg_.get());
+    parallel_args_.process_group_ = global_pg_.get();
+    parallel_args_.tp_group_ = tp_pg_.get();
+    parallel_args_.single_rank_group_ = single_rank_pg_.get();
+    parallel_args_.sp_group_ = tp_pg_.get();
+    parallel_args_.dp_local_process_group_ = dp_pg_.get();
+    parallel_args_.ep_size_ = ep_size;
+    parallel_args_.moe_ep_group_ = global_pg_.get();
+    parallel_args_.moe_tp_group_ = tp_pg_.get();
+  }
+
+  torch::Tensor mat(int64_t rows, const std::vector<float>& vals) const {
+    return torch::tensor(vals, fp32_opts()).reshape({rows, 2});
+  }
+
+  torch::TensorOptions fp32_opts() const {
+    return torch::TensorOptions()
+        .dtype(torch::kFloat32)
+        .device(options_.device());
+  }
+
+  std::unordered_map<std::string, torch::Tensor> create_fp_weights(
+      int64_t n_shared_experts) const {
+    std::unordered_map<std::string, torch::Tensor> weight_dict;
+    const int64_t num_experts = model_args_.n_routed_experts();
+    const int64_t hidden_size = model_args_.hidden_size();
+    const int64_t intermediate_size = model_args_.moe_intermediate_size();
+
+    for (int64_t expert_id = 0; expert_id < num_experts; ++expert_id) {
+      const std::string expert_prefix =
+          "experts." + std::to_string(expert_id) + ".";
+      const std::string seed_prefix =
+          "deepseek_v2_sparse_moe_block.expert_" + std::to_string(expert_id);
+      weight_dict[expert_prefix + "gate_proj.weight"] =
+          test::seeded_tensor(seed_prefix + ".gate_proj",
+                              {intermediate_size, hidden_size},
+                              torch::kBFloat16,
+                              options_.device());
+      weight_dict[expert_prefix + "up_proj.weight"] =
+          test::seeded_tensor(seed_prefix + ".up_proj",
+                              {intermediate_size, hidden_size},
+                              torch::kBFloat16,
+                              options_.device());
+      weight_dict[expert_prefix + "down_proj.weight"] =
+          test::seeded_tensor(seed_prefix + ".down_proj",
+                              {hidden_size, intermediate_size},
+                              torch::kBFloat16,
+                              options_.device());
+    }
+
+    weight_dict["gate.weight"] =
+        test::seeded_tensor("deepseek_v2_sparse_moe_block.gate",
+                            {num_experts, hidden_size},
+                            torch::kBFloat16,
+                            options_.device());
+
+    if (n_shared_experts > 0) {
+      const int64_t shared_size = intermediate_size * n_shared_experts;
+      weight_dict["shared_experts.gate_proj.weight"] =
+          test::seeded_tensor("deepseek_v2_sparse_moe_block.shared.gate_proj",
+                              {shared_size, hidden_size},
+                              torch::kBFloat16,
+                              options_.device());
+      weight_dict["shared_experts.up_proj.weight"] =
+          test::seeded_tensor("deepseek_v2_sparse_moe_block.shared.up_proj",
+                              {shared_size, hidden_size},
+                              torch::kBFloat16,
+                              options_.device());
+      weight_dict["shared_experts.down_proj.weight"] =
+          test::seeded_tensor("deepseek_v2_sparse_moe_block.shared.down_proj",
+                              {hidden_size, shared_size},
+                              torch::kBFloat16,
+                              options_.device());
+    }
+
+    return weight_dict;
+  }
+
+  DeepseekV2SparseMoEBlock create_block() const {
+    return DeepseekV2SparseMoEBlock(
+        model_args_, quant_args_, parallel_args_, options_);
+  }
+
+  FusedMoE create_raw_moe() const {
+    const FusedMoEArgs moe_args{.is_gated = true,
+                                .enable_result_reduction = false};
+    return FusedMoE(
+        model_args_, moe_args, quant_args_, parallel_args_, options_);
+  }
+
+  torch::Tensor run_comm(torch::Tensor x, ProcessGroup* pg) const {
+    return x + (pg == tp_pg_.get() ? 10.0f : 20.0f);
+  }
+
+  torch::Tensor run_reduce(torch::Tensor x, ProcessGroup* pg) const {
+    return x + (pg == tp_pg_.get() ? 100.0f : 200.0f);
+  }
+
+  void sync_dev() const {
+    xllm::Device(options_.device()).synchronize_default_stream();
+  }
+
+  ModelArgs model_args_;
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+  std::unique_ptr<test::MockProcessGroup> global_pg_;
+  std::unique_ptr<test::MockProcessGroup> dp_pg_;
+  std::unique_ptr<test::MockProcessGroup> tp_pg_;
+  std::unique_ptr<test::MockProcessGroup> single_rank_pg_;
+
+  v32_sp::DeepseekV32SPContext make_sp_ctx(
+      std::vector<int32_t> tokens_per_rank = {2, 2}) const {
+    v32_sp::DeepseekV32SPContext sp_ctx;
+    sp_ctx.rank = 0;
+    sp_ctx.process_group = tp_pg_.get();
+    sp_ctx.comm_plan.tokens_per_rank = std::move(tokens_per_rank);
+    sp_ctx.comm_plan.padded_tokens_per_rank = sp_ctx.comm_plan.tokens_per_rank;
+    sp_ctx.comm_plan.token_num_offset = 0;
+    sp_ctx.comm_plan.ffn_can_rs =
+        v32_sp::can_ffn_rs(sp_ctx.comm_plan.tokens_per_rank);
+    return sp_ctx;
+  }
+};
+
+TEST_F(DeepseekV2SparseMoEBlockTest, PlanExecEnablesAll2AllOnlyForDecode) {
+  set_tp_ctx(/*world_size=*/2, /*ep_size=*/2);
+  auto block = create_block();
+  DeepseekV2SparseMoEBlockTestPeer::set_enable_deep_ep(*block, true);
+
+  ModelInputParams decode_params;
+  decode_params.dp_global_token_nums = {1, 1};
+  decode_params.dp_is_decode = {1, 1};
+  auto decode_cfg = block->plan_exec(decode_params);
+  EXPECT_TRUE(decode_cfg.enable_all2all);
+  EXPECT_FALSE(decode_cfg.need_dp_gather);
+
+  ModelInputParams mixed_params;
+  mixed_params.dp_global_token_nums = {2, 1};
+  mixed_params.dp_is_decode = {0, 1};
+  auto mixed_cfg = block->plan_exec(mixed_params);
+  EXPECT_FALSE(mixed_cfg.enable_all2all);
+  EXPECT_FALSE(mixed_cfg.need_dp_gather);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, PlanExecSetsDpGatherWhenAll2AllOff) {
+  set_tp_dp_ctx(/*world_size=*/4, /*dp_size=*/2, /*tp_size=*/2, /*ep_size=*/4);
+  auto block = create_block();
+
+  ModelInputParams input_params;
+  input_params.dp_global_token_nums = {3, 1};
+  input_params.dp_is_decode = {0, 0};
+
+  auto cfg = block->plan_exec(input_params);
+  EXPECT_FALSE(cfg.enable_all2all);
+  EXPECT_TRUE(cfg.need_dp_gather);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, PrepInDpGatherBuildsLocalSkip) {
+  set_tp_dp_ctx(/*world_size=*/4, /*dp_size=*/2, /*tp_size=*/2, /*ep_size=*/4);
+  auto block = create_block();
+
+  ModelInputParams input_params;
+  input_params.dp_global_token_nums = {3, 1};
+  auto attn_out =
+      mat(/*rows=*/4, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
+  auto residual =
+      mat(/*rows=*/4, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f});
+  auto full_tokens =
+      mat(/*rows=*/4, {11.0f, 22.0f, 33.0f, 44.0f, 55.0f, 66.0f, 77.0f, 88.0f});
+  tp_pg_->set_allgather_outputs(
+      {full_tokens.slice(0, 0, 2), full_tokens.slice(0, 2, 4)});
+
+  auto prep = block->prep_in(attn_out,
+                             residual,
+                             input_params,
+                             block->plan_exec(input_params),
+                             DeepseekV2AttentionImpl::PostAttnLayout::kTpShard);
+
+  EXPECT_TRUE(prep.need_dp_gather);
+  EXPECT_FALSE(prep.need_tp_pad);
+  EXPECT_FALSE(prep.pad_info.active);
+  test::verify_tensor_close(prep.ffn_in, full_tokens.slice(0, 0, 2));
+  test::verify_tensor_close(
+      prep.skip_local,
+      mat(/*rows=*/3, {11.0f, 22.0f, 33.0f, 44.0f, 55.0f, 66.0f}));
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, GatherInDpGatherRebuildsGlobalTokens) {
+  set_tp_dp_ctx(/*world_size=*/4, /*dp_size=*/2, /*tp_size=*/2, /*ep_size=*/4);
+  auto block = create_block();
+
+  ModelInputParams input_params;
+  input_params.dp_global_token_nums = {3, 1};
+  DeepseekV2SparseMoEBlockImpl::PrepOut prep;
+  prep.ffn_in = mat(/*rows=*/2, {11.0f, 22.0f, 33.0f, 44.0f});
+  prep.need_dp_gather = true;
+  auto dp0_tp1 = mat(/*rows=*/2, {55.0f, 66.0f, 0.0f, 0.0f});
+  auto dp1_tp0 = mat(/*rows=*/2, {101.0f, 202.0f, 0.0f, 0.0f});
+  auto dp1_tp1 = torch::zeros_like(dp1_tp0);
+  global_pg_->set_allgather_outputs({prep.ffn_in, dp0_tp1, dp1_tp0, dp1_tp1});
+
+  auto gathered = block->gather_in(prep, input_params);
+
+  test::verify_tensor_close(
+      gathered,
+      mat(/*rows=*/4,
+          {11.0f, 22.0f, 33.0f, 44.0f, 55.0f, 66.0f, 101.0f, 202.0f}));
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, PrepInAll2AllPadsTpShardInput) {
+  set_tp_ctx(/*world_size=*/2, /*ep_size=*/2);
+  auto block = create_block();
+  DeepseekV2SparseMoEBlockTestPeer::set_enable_deep_ep(*block, true);
+
+  ModelInputParams input_params;
+  input_params.dp_global_token_nums = {1, 1};
+  input_params.dp_is_decode = {1, 1};
+  auto attn_out = mat(/*rows=*/3, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  auto residual = mat(/*rows=*/3, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f});
+
+  auto prep = block->prep_in(attn_out,
+                             residual,
+                             input_params,
+                             block->plan_exec(input_params),
+                             DeepseekV2AttentionImpl::PostAttnLayout::kTpShard);
+
+  EXPECT_FALSE(prep.need_dp_gather);
+  EXPECT_TRUE(prep.need_tp_pad);
+  EXPECT_TRUE(prep.pad_info.active);
+  EXPECT_EQ(prep.pad_info.original_tokens, 3);
+  EXPECT_EQ(prep.pad_info.padded_tokens, 4);
+  test::verify_tensor_close(prep.ffn_in,
+                            mat(/*rows=*/2, {11.0f, 22.0f, 33.0f, 44.0f}));
+  test::verify_tensor_close(prep.skip_local, prep.ffn_in);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, PrepInUsesProvidedExecCfg) {
+  set_tp_dp_ctx(/*world_size=*/4, /*dp_size=*/2, /*tp_size=*/2, /*ep_size=*/4);
+  auto block = create_block();
+
+  ModelInputParams input_params;
+  input_params.dp_global_token_nums = {3, 1};
+  input_params.dp_is_decode = {0, 0};
+  auto planned_cfg = block->plan_exec(input_params);
+  EXPECT_FALSE(planned_cfg.enable_all2all);
+  EXPECT_TRUE(planned_cfg.need_dp_gather);
+
+  DeepseekV2SparseMoEBlockImpl::ExecCfg forced_cfg;
+  forced_cfg.enable_all2all = true;
+  forced_cfg.need_dp_gather = false;
+  auto attn_out = mat(/*rows=*/3, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  auto residual = mat(/*rows=*/3, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f});
+
+  auto prep = block->prep_in(attn_out,
+                             residual,
+                             input_params,
+                             forced_cfg,
+                             DeepseekV2AttentionImpl::PostAttnLayout::kTpShard);
+
+  EXPECT_FALSE(prep.need_dp_gather);
+  EXPECT_TRUE(prep.need_tp_pad);
+  EXPECT_TRUE(prep.pad_info.active);
+  test::verify_tensor_close(prep.skip_local, prep.ffn_in);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, MergeOutTpPadGathersAndUnpads) {
+  set_tp_ctx(/*world_size=*/2, /*ep_size=*/2);
+  auto block = create_block();
+
+  ModelInputParams input_params;
+  DeepseekV2SparseMoEBlockImpl::PrepOut prep;
+  prep.need_tp_pad = true;
+  prep.pad_info = {.original_tokens = 3, .padded_tokens = 4, .active = true};
+  auto shard0 = mat(/*rows=*/2, {1.0f, 2.0f, 3.0f, 4.0f});
+  auto shard1 = mat(/*rows=*/2, {5.0f, 6.0f, 0.0f, 0.0f});
+  prep.skip_local = shard0;
+  tp_pg_->set_allgather_outputs({shard0, shard1});
+
+  auto merged = block->merge_out(shard0, prep, input_params);
+
+  test::verify_tensor_close(
+      merged, mat(/*rows=*/3, {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f}));
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, MergeOutDpGatherSlicesLocalTokens) {
+  set_tp_dp_ctx(/*world_size=*/4, /*dp_size=*/2, /*tp_size=*/2, /*ep_size=*/4);
+  auto block = create_block();
+
+  ModelInputParams input_params;
+  input_params.dp_global_token_nums = {3, 1};
+  DeepseekV2SparseMoEBlockImpl::PrepOut prep;
+  prep.skip_local = mat(/*rows=*/3, {11.0f, 22.0f, 33.0f, 44.0f, 55.0f, 66.0f});
+  prep.need_dp_gather = true;
+  auto ffn_out =
+      mat(/*rows=*/4,
+          {101.0f, 102.0f, 103.0f, 104.0f, 105.0f, 106.0f, 107.0f, 108.0f});
+
+  auto merged = block->merge_out(ffn_out, prep, input_params);
+
+  test::verify_tensor_close(
+      merged,
+      mat(/*rows=*/3, {112.0f, 124.0f, 136.0f, 148.0f, 160.0f, 172.0f}));
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardReducePathCombinesSharedAndRouted) {
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto hidden_states = test::seeded_tensor("deepseek_v2_sparse_moe_block.input",
+                                           {4, model_args_.hidden_size()},
+                                           torch::kBFloat16,
+                                           options_.device());
+  auto routed = raw_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  auto shared = raw_moe->forward_shared(hidden_states);
+  auto* shared_pg = raw_moe->shared_pg();
+  ASSERT_TRUE(shared.defined());
+  ASSERT_EQ(shared_pg, single_rank_pg_.get());
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  auto result = block->forward(
+      hidden_states,
+      /*enable_moe_all2all=*/false,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [](ProcessGroup*) { return false; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  sync_dev();
+
+  auto expected = run_reduce(routed, tp_pg_.get()) + shared;
+  EXPECT_FALSE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 0);
+  EXPECT_EQ(reduce_calls, 1);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardReduceOverlapUsesAsyncReduceFns) {
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto hidden_states =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.async_reduce",
+                          {4, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  auto routed = raw_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  auto shared = raw_moe->forward_shared(hidden_states);
+  ASSERT_TRUE(shared.defined());
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  int launch_reduce_calls = 0;
+  int finish_reduce_calls = 0;
+  auto result = block->forward(
+      hidden_states,
+      /*enable_moe_all2all=*/false,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [](ProcessGroup*) { return false; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+          .launch_reduce = std::function<parallel_state::ReduceAsyncCtx(
+              torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++launch_reduce_calls;
+                return parallel_state::ReduceAsyncCtx{
+                    .tensor = run_reduce(std::move(x), pg),
+                };
+              }),
+          .finish_reduce =
+              std::function<torch::Tensor(parallel_state::ReduceAsyncCtx)>(
+                  [&](parallel_state::ReduceAsyncCtx ctx) {
+                    ++finish_reduce_calls;
+                    return std::move(ctx.tensor);
+                  }),
+      });
+
+  sync_dev();
+
+  auto expected = run_reduce(routed, tp_pg_.get()) + shared;
+  EXPECT_FALSE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 0);
+  EXPECT_EQ(reduce_calls, 0);
+  EXPECT_EQ(launch_reduce_calls, 1);
+  EXPECT_EQ(finish_reduce_calls, 1);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardKeepLocalUsesCommPath) {
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto hidden_states = test::seeded_tensor("deepseek_v2_sparse_moe_block.comm",
+                                           {4, model_args_.hidden_size()},
+                                           torch::kBFloat16,
+                                           options_.device());
+  auto routed = raw_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  auto shared = raw_moe->forward_shared(hidden_states);
+  auto* shared_pg = raw_moe->shared_pg();
+  ASSERT_TRUE(shared.defined());
+  ASSERT_EQ(shared_pg, single_rank_pg_.get());
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  auto result = block->forward(
+      hidden_states,
+      /*enable_moe_all2all=*/false,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [](ProcessGroup*) { return true; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  sync_dev();
+
+  auto expected = run_comm(routed, tp_pg_.get()) + run_comm(shared, shared_pg);
+  EXPECT_TRUE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 2);
+  EXPECT_EQ(reduce_calls, 0);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardIgnoresSharedLocalGate) {
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto hidden_states =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.shared_fallback",
+                          {4, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  auto routed = raw_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  auto shared = raw_moe->forward_shared(hidden_states);
+  ASSERT_TRUE(shared.defined());
+
+  auto* routed_pg = DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block);
+  auto* shared_pg = DeepseekV2SparseMoEBlockTestPeer::moe(*block)->shared_pg();
+  ASSERT_NE(routed_pg, shared_pg);
+  ASSERT_EQ(shared_pg, single_rank_pg_.get());
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  auto result = block->forward(
+      hidden_states,
+      /*enable_moe_all2all=*/false,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [&](ProcessGroup* pg) { return pg == routed_pg; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  sync_dev();
+
+  auto expected = run_comm(routed, routed_pg) + run_comm(shared, shared_pg);
+  EXPECT_TRUE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 2);
+  EXPECT_EQ(reduce_calls, 0);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardWithoutSharedIgnoresNullSharedPg) {
+  model_args_.n_shared_experts() = 0;
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/0));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto hidden_states =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.no_shared",
+                          {4, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  auto routed = raw_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  ASSERT_FALSE(raw_moe->forward_shared(hidden_states).defined());
+
+  auto* routed_pg = DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block);
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  auto result = block->forward(
+      hidden_states,
+      /*enable_moe_all2all=*/false,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [&](ProcessGroup* pg) { return pg == routed_pg; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  sync_dev();
+
+  auto expected = run_comm(routed, routed_pg);
+  EXPECT_TRUE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 1);
+  EXPECT_EQ(reduce_calls, 0);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardSPAddsLocalSharedAfterRoutedComm) {
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto local = test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_local",
+                                   {2, model_args_.hidden_size()},
+                                   torch::kBFloat16,
+                                   options_.device());
+  auto remote = test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_remote",
+                                    {2, model_args_.hidden_size()},
+                                    torch::kBFloat16,
+                                    options_.device());
+  tp_pg_->set_allgather_outputs({local, remote});
+  auto sp_ctx = make_sp_ctx();
+
+  auto gathered = torch::cat({local, remote}, 0);
+  auto routed = raw_moe->forward_experts(
+      gathered, /*enable_all2all_communication=*/false);
+  auto shared = raw_moe->forward_shared(local);
+  auto* routed_pg = DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block);
+  const int64_t local_token_num = local.size(0);
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  auto result = block->forward_sp(
+      local,
+      sp_ctx,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [&](ProcessGroup* pg) { return pg == routed_pg; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                if (pg == routed_pg) {
+                  x = x.slice(0, 0, local_token_num);
+                }
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  sync_dev();
+
+  auto expected =
+      run_comm(routed.slice(0, 0, local_token_num), routed_pg) + shared;
+  EXPECT_TRUE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 1);
+  EXPECT_EQ(reduce_calls, 0);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardSPFallsBackWhenLocalKeepOff) {
+  auto block = create_block();
+  auto raw_moe = create_raw_moe();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto local =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_fallback_local",
+                          {2, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  auto remote =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_fallback_remote",
+                          {2, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  tp_pg_->set_allgather_outputs({local, remote});
+  auto sp_ctx = make_sp_ctx();
+
+  auto gathered = torch::cat({local, remote}, 0);
+  auto routed = raw_moe->forward_experts(
+      gathered, /*enable_all2all_communication=*/false);
+  auto shared = raw_moe->forward_shared(gathered);
+  auto* routed_pg = DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block);
+
+  int comm_calls = 0;
+  int reduce_calls = 0;
+  auto result = block->forward_sp(
+      local,
+      sp_ctx,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [&](ProcessGroup* /*pg*/) { return false; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  sync_dev();
+
+  auto expected = run_reduce(routed, routed_pg) + shared;
+  EXPECT_FALSE(result.keep_local_output);
+  EXPECT_EQ(comm_calls, 0);
+  EXPECT_EQ(reduce_calls, 1);
+  test::verify_tensor_close(result.output, expected, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardSPChunkKeepLocalMatchesBase) {
+  auto block = create_block();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+
+  auto local =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_chunk_local",
+                          {2, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  auto remote =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_chunk_remote",
+                          {2, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  tp_pg_->set_allgather_outputs({local, remote});
+  auto sp_ctx = make_sp_ctx();
+
+  const int64_t local_token_num = local.size(0);
+  int base_comm_calls = 0;
+  int base_reduce_calls = 0;
+  auto base = block->forward_sp(
+      local,
+      sp_ctx,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local =
+              std::function<bool(ProcessGroup*)>([&](ProcessGroup* pg) {
+                return pg ==
+                       DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block);
+              }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++base_comm_calls;
+                if (pg == DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block)) {
+                  x = x.slice(0, 0, local_token_num);
+                }
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++base_reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  int chunk_comm_calls = 0;
+  int chunk_reduce_calls = 0;
+  auto chunked = block->forward_sp(
+      local,
+      sp_ctx,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local =
+              std::function<bool(ProcessGroup*)>([&](ProcessGroup* pg) {
+                return pg ==
+                       DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block);
+              }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++chunk_comm_calls;
+                if (pg == DeepseekV2SparseMoEBlockTestPeer::routed_pg(*block)) {
+                  x = x.slice(0, 0, local_token_num);
+                }
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++chunk_reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      },
+      /*chunk_size=*/1);
+
+  sync_dev();
+
+  EXPECT_TRUE(base.keep_local_output);
+  EXPECT_TRUE(chunked.keep_local_output);
+  EXPECT_EQ(base_comm_calls, 1);
+  EXPECT_EQ(base_reduce_calls, 0);
+  EXPECT_EQ(chunk_comm_calls, 1);
+  EXPECT_EQ(chunk_reduce_calls, 0);
+  test::verify_tensor_close(chunked.output, base.output, 1e-3, 1e-4);
+}
+
+TEST_F(DeepseekV2SparseMoEBlockTest, ForwardSPChunkFallbackMatchesBase) {
+  auto block = create_block();
+  StateDict state_dict(create_fp_weights(/*n_shared_experts=*/1));
+  block->load_state_dict(state_dict);
+
+  auto local =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_chunk_fb_local",
+                          {2, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  auto remote =
+      test::seeded_tensor("deepseek_v2_sparse_moe_block.sp_chunk_fb_remote",
+                          {2, model_args_.hidden_size()},
+                          torch::kBFloat16,
+                          options_.device());
+  tp_pg_->set_allgather_outputs({local, remote});
+  auto sp_ctx = make_sp_ctx();
+
+  int base_comm_calls = 0;
+  int base_reduce_calls = 0;
+  auto base = block->forward_sp(
+      local,
+      sp_ctx,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [&](ProcessGroup* /*pg*/) { return false; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++base_comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++base_reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      });
+
+  int chunk_comm_calls = 0;
+  int chunk_reduce_calls = 0;
+  auto chunked = block->forward_sp(
+      local,
+      sp_ctx,
+      DeepseekV2SparseMoEBlockImpl::CommFns{
+          .can_keep_local = std::function<bool(ProcessGroup*)>(
+              [&](ProcessGroup* /*pg*/) { return false; }),
+          .comm = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++chunk_comm_calls;
+                return run_comm(std::move(x), pg);
+              }),
+          .reduce = std::function<torch::Tensor(torch::Tensor, ProcessGroup*)>(
+              [&](torch::Tensor x, ProcessGroup* pg) {
+                ++chunk_reduce_calls;
+                return run_reduce(std::move(x), pg);
+              }),
+      },
+      /*chunk_size=*/1);
+
+  sync_dev();
+
+  EXPECT_FALSE(base.keep_local_output);
+  EXPECT_FALSE(chunked.keep_local_output);
+  EXPECT_EQ(base_comm_calls, 0);
+  EXPECT_EQ(base_reduce_calls, 1);
+  EXPECT_EQ(chunk_comm_calls, 0);
+  EXPECT_EQ(chunk_reduce_calls, 1);
+  test::verify_tensor_close(chunked.output, base.output, 1e-3, 1e-4);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/deepseek_v32_sp_utils_test.cpp

@@ -0,0 +1,971 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <gflags/gflags.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <numeric>
+#include <optional>
+#include <vector>
+
+#include "core/common/global_flags.h"
+#include "framework/batch/batch_forward_type.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "layers/common/attention_metadata.h"
+#include "layers/mlu/deepseek_v32_sp_context.h"
+#include "layers/mlu/tests_utils.h"
+
+namespace xllm::layer::v32_sp {
+namespace {
+
+class ScriptedAllGatherProcessGroup
+    : public xllm::layer::test::MockProcessGroup {
+ public:
+  ScriptedAllGatherProcessGroup(const torch::Device& device,
+                                int64_t rank,
+                                std::vector<torch::Tensor> scripted_outputs)
+      : MockProcessGroup(device,
+                         rank,
+                         static_cast<int64_t>(scripted_outputs.size())),
+        scripted_outputs_(std::move(scripted_outputs)) {}
+
+  void allgather(const torch::Tensor& input,
+                 std::vector<torch::Tensor>& outputs) override {
+    (void)input;
+    CHECK_EQ(outputs.size(), scripted_outputs_.size());
+    for (size_t i = 0; i < scripted_outputs_.size(); ++i) {
+      outputs[i].copy_(scripted_outputs_[i]);
+    }
+  }
+
+  c10::intrusive_ptr<c10d::Work> allgather_async(
+      const torch::Tensor& input,
+      std::vector<torch::Tensor>& outputs) override {
+    allgather(input, outputs);
+    return xllm::layer::test::make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> allgather_base_async(
+      const torch::Tensor& input,
+      torch::Tensor& output) override {
+    (void)input;
+    CHECK_EQ(output.size(0), static_cast<int64_t>(scripted_outputs_.size()));
+    for (size_t i = 0; i < scripted_outputs_.size(); ++i) {
+      output[static_cast<int64_t>(i)].copy_(scripted_outputs_[i]);
+    }
+    return xllm::layer::test::make_completed_work();
+  }
+
+ private:
+  std::vector<torch::Tensor> scripted_outputs_;
+};
+
+class ScopedFlagValue {
+ public:
+  ScopedFlagValue(bool& flag, bool value) : bool_flag_(&flag), old_bool_(flag) {
+    flag = value;
+  }
+
+  ScopedFlagValue(int32_t& flag, int32_t value)
+      : int_flag_(&flag), old_int_(flag) {
+    flag = value;
+  }
+
+  ~ScopedFlagValue() {
+    if (bool_flag_ != nullptr) {
+      *bool_flag_ = old_bool_;
+    }
+    if (int_flag_ != nullptr) {
+      *int_flag_ = old_int_;
+    }
+  }
+
+ private:
+  bool* bool_flag_ = nullptr;
+  int32_t* int_flag_ = nullptr;
+  bool old_bool_ = false;
+  int32_t old_int_ = 0;
+};
+
+AttentionMetadata make_prefill_metadata(
+    const std::vector<int32_t>& q_seq_lens,
+    const std::vector<int32_t>& ctx_seq_lens = {},
+    const std::optional<torch::Tensor>& block_table = std::nullopt,
+    BatchForwardType batch_forward_type = BatchForwardType::PREFILL) {
+  AttentionMetadata attn_metadata;
+  std::vector<int32_t> q_cu_seq_lens = {0};
+  int32_t total = 0;
+  for (int32_t q_seq_len : q_seq_lens) {
+    total += q_seq_len;
+    q_cu_seq_lens.push_back(total);
+  }
+
+  auto int32_options = torch::TensorOptions().dtype(torch::kInt32);
+  attn_metadata.q_cu_seq_lens = torch::tensor(q_cu_seq_lens, int32_options);
+  attn_metadata.kv_cu_seq_lens = torch::tensor(q_cu_seq_lens, int32_options);
+  attn_metadata.kv_seq_lens = torch::tensor(
+      ctx_seq_lens.empty() ? q_seq_lens : ctx_seq_lens, int32_options);
+  if (block_table.has_value()) {
+    attn_metadata.block_table = *block_table;
+  }
+  attn_metadata.is_prefill = batch_forward_type.is_prefill();
+  attn_metadata.is_chunked_prefill = batch_forward_type.is_chunked_prefill();
+  attn_metadata.is_dummy = false;
+  return attn_metadata;
+}
+
+torch::Tensor make_block_table(const std::vector<std::vector<int32_t>>& rows) {
+  CHECK(!rows.empty());
+  const int64_t row_num = static_cast<int64_t>(rows.size());
+  const int64_t col_num = static_cast<int64_t>(rows.front().size());
+  std::vector<int32_t> flat;
+  flat.reserve(row_num * col_num);
+  for (const auto& row : rows) {
+    CHECK_EQ(row.size(), col_num);
+    flat.insert(flat.end(), row.begin(), row.end());
+  }
+  return torch::tensor(flat, torch::TensorOptions().dtype(torch::kInt32))
+      .view({row_num, col_num});
+}
+
+std::vector<int32_t> extract_segment_req_idx(
+    const std::vector<DeepseekV32SPSegment>& segments) {
+  std::vector<int32_t> values;
+  values.reserve(segments.size());
+  for (const auto& segment : segments) {
+    values.push_back(segment.req_idx);
+  }
+  return values;
+}
+
+std::vector<int32_t> extract_segment_q_tokens(
+    const std::vector<DeepseekV32SPSegment>& segments) {
+  std::vector<int32_t> values;
+  values.reserve(segments.size());
+  for (const auto& segment : segments) {
+    values.push_back(segment.q_tokens);
+  }
+  return values;
+}
+
+std::vector<int32_t> extract_segment_suffix_k_lens(
+    const std::vector<DeepseekV32SPSegment>& segments) {
+  std::vector<int32_t> values;
+  values.reserve(segments.size());
+  for (const auto& segment : segments) {
+    values.push_back(segment.suffix_k_len);
+  }
+  return values;
+}
+
+std::vector<int32_t> extract_segment_ctx_lens(
+    const std::vector<DeepseekV32SPSegment>& segments) {
+  std::vector<int32_t> values;
+  values.reserve(segments.size());
+  for (const auto& segment : segments) {
+    values.push_back(segment.ctx_k_len);
+  }
+  return values;
+}
+
+void expect_int32_contiguous(const torch::Tensor& tensor) {
+  EXPECT_EQ(tensor.scalar_type(), torch::kInt32);
+  EXPECT_TRUE(tensor.is_contiguous());
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildZigzagSplitPlanMatchesSingleRequestWithoutPadding) {
+  const auto all_segments = build_all_sp_segments(4, {16}, {16});
+  const auto runtime_artifacts =
+      build_sp_runtime_artifacts(0, 4, all_segments, /*total_tokens=*/16);
+  const auto& rank0 = runtime_artifacts.comm_plan;
+  const auto& gathered_reorder_index =
+      runtime_artifacts.gathered_reorder_index_cpu;
+  const std::vector<int32_t> local_reorder_index(
+      gathered_reorder_index.begin() + rank0.token_num_offset,
+      gathered_reorder_index.begin() + rank0.token_num_offset +
+          rank0.tokens_per_rank[0]);
+  EXPECT_EQ(rank0.tokens_per_rank, (std::vector<int32_t>{4, 4, 4, 4}));
+  EXPECT_EQ(rank0.padded_tokens_per_rank, (std::vector<int32_t>{4, 4, 4, 4}));
+  EXPECT_EQ(local_reorder_index, (std::vector<int32_t>{0, 1, 14, 15}));
+  EXPECT_EQ(gathered_reorder_index,
+            (std::vector<int32_t>{
+                0, 1, 14, 15, 2, 3, 12, 13, 4, 5, 10, 11, 6, 7, 8, 9}));
+
+  const auto segments = build_local_sp_segments(0, all_segments);
+  EXPECT_EQ(extract_segment_req_idx(segments), (std::vector<int32_t>{0, 0}));
+  EXPECT_EQ(extract_segment_q_tokens(segments), (std::vector<int32_t>{2, 2}));
+  EXPECT_EQ(extract_segment_suffix_k_lens(segments),
+            (std::vector<int32_t>{2, 16}));
+  EXPECT_EQ(extract_segment_ctx_lens(segments), (std::vector<int32_t>{2, 16}));
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildContextRejectsMixedBatchForwardType) {
+  ScopedFlagValue enable_prefill_sp(FLAGS_enable_prefill_sp, true);
+  auto attn_metadata = make_prefill_metadata({8}, {8}, std::nullopt);
+  auto tokens = torch::arange(8, torch::TensorOptions().dtype(torch::kInt64));
+  xllm::layer::test::MockProcessGroup process_group(torch::kCPU,
+                                                    /*rank=*/0,
+                                                    /*world_size=*/2);
+
+  auto context = build_deepseek_v32_sp_context(attn_metadata,
+                                               BatchForwardType::MIXED,
+                                               tokens,
+                                               &process_group,
+                                               /*curr_rank=*/0,
+                                               /*world_size=*/2);
+
+  EXPECT_FALSE(context.has_value());
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildZigzagSplitPlanMatchesSingleRequestWithPadding) {
+  const auto all_segments = build_all_sp_segments(4, {10}, {10});
+  const auto runtime_artifacts =
+      build_sp_runtime_artifacts(2, 4, all_segments, /*total_tokens=*/10);
+  const auto& rank2 = runtime_artifacts.comm_plan;
+  const auto& gathered_reorder_index =
+      runtime_artifacts.gathered_reorder_index_cpu;
+  const std::vector<int32_t> local_reorder_index(
+      gathered_reorder_index.begin() + rank2.token_num_offset,
+      gathered_reorder_index.begin() + rank2.token_num_offset +
+          rank2.tokens_per_rank[2]);
+  EXPECT_EQ(rank2.tokens_per_rank, (std::vector<int32_t>{3, 3, 2, 2}));
+  EXPECT_EQ(rank2.padded_tokens_per_rank, (std::vector<int32_t>{3, 3, 3, 3}));
+  EXPECT_EQ(local_reorder_index, (std::vector<int32_t>{4, 7}));
+
+  const auto segments = build_local_sp_segments(2, all_segments);
+  EXPECT_EQ(extract_segment_req_idx(segments), (std::vector<int32_t>{0, 0}));
+  EXPECT_EQ(extract_segment_q_tokens(segments), (std::vector<int32_t>{1, 1}));
+  EXPECT_EQ(extract_segment_suffix_k_lens(segments),
+            (std::vector<int32_t>{5, 8}));
+  EXPECT_EQ(extract_segment_ctx_lens(segments), (std::vector<int32_t>{5, 8}));
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildZigzagSplitPlanTracksContextLens) {
+  AttentionMetadata attn_metadata = make_prefill_metadata({4, 6}, {16, 22});
+
+  EXPECT_EQ(extract_q_seq_lens(attn_metadata), (std::vector<int32_t>{4, 6}));
+  EXPECT_EQ(extract_ctx_seq_lens(attn_metadata),
+            (std::vector<int32_t>{16, 22}));
+
+  const auto all_segments = build_all_sp_segments(4, {4, 6}, {16, 22});
+  const auto segments = build_local_sp_segments(1, all_segments);
+  EXPECT_EQ(extract_segment_q_tokens(segments),
+            (std::vector<int32_t>{1, 0, 1, 1}));
+  EXPECT_EQ(extract_segment_suffix_k_lens(segments),
+            (std::vector<int32_t>{2, 0, 2, 5}));
+  EXPECT_EQ(extract_segment_ctx_lens(segments),
+            (std::vector<int32_t>{14, 12, 18, 21}));
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildDeepseekV32SPContextBuildsMetadataForCurrentRank) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue world_size_flag(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata =
+      make_prefill_metadata({4, 6, 11}, {12, 16, 32});
+  torch::Tensor tokens =
+      torch::arange(0, 21, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+  EXPECT_EQ(context.total_tokens, 21);
+  EXPECT_EQ(context.rank, 1);
+  EXPECT_EQ(context.comm_plan.tokens_per_rank,
+            (std::vector<int32_t>{6, 6, 5, 4}));
+  EXPECT_EQ(context.comm_plan.padded_tokens_per_rank,
+            (std::vector<int32_t>{6, 6, 6, 6}));
+  EXPECT_EQ(context.comm_plan.token_num_offset, 6);
+
+  EXPECT_TRUE(
+      torch::equal(reorder_to_local_shard(tokens, context),
+                   torch::tensor({1, 5, 8, 12, 13, 19},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.gathered_reorder_index,
+                   torch::tensor({0,  4, 9, 10, 11, 20, 1, 5, 8,  12, 13,
+                                  19, 2, 6, 14, 15, 18, 3, 7, 16, 17},
+                                 torch::TensorOptions().dtype(torch::kInt64))));
+
+  EXPECT_TRUE(
+      torch::equal(context.local_attn_metadata.q_cu_seq_lens,
+                   torch::tensor({0, 1, 1, 2, 3, 5, 6},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.local_attn_metadata.kv_seq_lens,
+                   torch::tensor({2, 0, 2, 5, 4, 10},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_EQ(context.local_attn_metadata.max_query_len, 2);
+  EXPECT_EQ(context.local_attn_metadata.max_seq_len, 10);
+  EXPECT_EQ(extract_segment_req_idx(context.local_segments),
+            (std::vector<int32_t>{0, 0, 1, 1, 2, 2}));
+  EXPECT_EQ(extract_segment_q_tokens(context.local_segments),
+            (std::vector<int32_t>{1, 0, 1, 1, 2, 1}));
+  EXPECT_EQ(extract_segment_suffix_k_lens(context.local_segments),
+            (std::vector<int32_t>{2, 0, 2, 5, 4, 10}));
+  EXPECT_EQ(extract_segment_ctx_lens(context.local_segments),
+            (std::vector<int32_t>{10, 8, 12, 15, 25, 31}));
+  EXPECT_EQ(context.seg_q_starts_cpu, (std::vector<int32_t>{0, 1, 1, 2, 3, 5}));
+  EXPECT_EQ(context.req_q_offsets_cpu, (std::vector<int32_t>{0, 4, 10}));
+  EXPECT_EQ(context.req_ctx_offsets_cpu, (std::vector<int32_t>{0, 12, 28}));
+  expect_int32_contiguous(context.seg_q_cu_lens_2col);
+  expect_int32_contiguous(context.seg_suffix_k_cu_lens_2col);
+  expect_int32_contiguous(context.seg_ctx_k_cu_lens_2col);
+  expect_int32_contiguous(context.seg_ctx_lens_1col);
+  EXPECT_TRUE(torch::equal(
+      context.seg_q_cu_lens_2col,
+      torch::tensor({{0, 1}, {0, 0}, {0, 1}, {0, 1}, {0, 2}, {0, 1}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(torch::equal(
+      context.seg_suffix_k_cu_lens_2col,
+      torch::tensor({{0, 2}, {0, 0}, {0, 2}, {0, 5}, {0, 4}, {0, 10}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(torch::equal(
+      context.seg_ctx_k_cu_lens_2col,
+      torch::tensor({{0, 10}, {0, 8}, {0, 12}, {0, 15}, {0, 25}, {0, 31}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.seg_ctx_lens_1col,
+                   torch::tensor({10, 8, 12, 15, 25, 31},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildDeepseekV32SPContextFallsBackWhenSeqShorterThanWorldSize) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes_flag(FLAGS_nnodes, 2);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({3});
+  torch::Tensor tokens =
+      torch::arange(0, 3, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    0,
+                                    4);
+
+  EXPECT_FALSE(maybe_context.has_value());
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildDeepseekV32SPContextAcceptsChunkedPrefillBatch) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue world_size_flag(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata(
+      {4, 6}, {16, 22}, std::nullopt, BatchForwardType::CHUNKED_PREFILL);
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::CHUNKED_PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+  EXPECT_TRUE(
+      torch::equal(context.local_attn_metadata.kv_seq_lens,
+                   torch::tensor({2, 0, 2, 5},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.seg_ctx_lens_1col,
+                   torch::tensor({14, 12, 18, 21},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_EQ(context.req_q_offsets_cpu, (std::vector<int32_t>{0, 4}));
+  EXPECT_EQ(context.req_ctx_offsets_cpu, (std::vector<int32_t>{0, 16}));
+  EXPECT_TRUE(
+      torch::equal(context.seg_suffix_k_cu_lens_2col,
+                   torch::tensor({{0, 2}, {0, 0}, {0, 2}, {0, 5}},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.seg_ctx_k_cu_lens_2col,
+                   torch::tensor({{0, 14}, {0, 12}, {0, 18}, {0, 21}},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildDeepseekV32SPContextRejectsMixedBatch) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue world_size_flag(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({4, 6}, {16, 22});
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::MIXED,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+
+  EXPECT_FALSE(maybe_context.has_value());
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildSPContextTracksSegmentRuntimeView) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue world_size_flag(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata(
+      {4, 6, 11},
+      {12, 16, 32},
+      make_block_table({{10, 11, 12}, {20, 21, 22}, {30, 31, 32}}));
+  torch::Tensor tokens =
+      torch::arange(0, 21, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+
+  EXPECT_EQ(context.seg_q_starts_cpu, (std::vector<int32_t>{0, 1, 1, 2, 3, 5}));
+  EXPECT_EQ(context.req_q_offsets_cpu, (std::vector<int32_t>{0, 4, 10}));
+  EXPECT_EQ(context.req_ctx_offsets_cpu, (std::vector<int32_t>{0, 12, 28}));
+  EXPECT_TRUE(torch::equal(
+      context.seg_q_cu_lens_2col,
+      torch::tensor({{0, 1}, {0, 0}, {0, 1}, {0, 1}, {0, 2}, {0, 1}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(torch::equal(
+      context.seg_suffix_k_cu_lens_2col,
+      torch::tensor({{0, 2}, {0, 0}, {0, 2}, {0, 5}, {0, 4}, {0, 10}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(torch::equal(
+      context.seg_ctx_k_cu_lens_2col,
+      torch::tensor({{0, 10}, {0, 8}, {0, 12}, {0, 15}, {0, 25}, {0, 31}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.seg_ctx_lens_1col,
+                   torch::tensor({10, 8, 12, 15, 25, 31},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildSegmentTensorCacheTracksPartialPrefixHit) {
+  AttentionMetadata attn_metadata = make_prefill_metadata({4, 6}, {16, 22});
+  const auto all_segments = build_all_sp_segments(4, {4, 6}, {16, 22});
+  const auto local_segments = build_local_sp_segments(1, all_segments);
+  const auto local_attn_metadata =
+      build_local_prefill_attention_metadata(attn_metadata, local_segments);
+
+  EXPECT_TRUE(
+      torch::equal(local_attn_metadata.kv_seq_lens,
+                   torch::tensor({2, 0, 2, 5},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_EQ(local_attn_metadata.max_seq_len, 5);
+  EXPECT_TRUE(
+      torch::equal(build_segment_length_matrix(local_segments,
+                                               &DeepseekV32SPSegment::ctx_k_len,
+                                               torch::Device(torch::kCPU)),
+                   torch::tensor({{0, 14}, {0, 12}, {0, 18}, {0, 21}},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(torch::equal(
+      build_segment_ctx_lens_tensor(local_segments, torch::Device(torch::kCPU)),
+      torch::tensor({14, 12, 18, 21},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(torch::equal(
+      build_segment_length_matrix(local_segments,
+                                  &DeepseekV32SPSegment::suffix_k_len,
+                                  torch::Device(torch::kCPU)),
+      torch::tensor({{0, 2}, {0, 0}, {0, 2}, {0, 5}},
+                    torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildSPContextKeepsExactBlockRollbackLens) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue world_size_flag(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata(
+      {4, 8}, {20, 24}, make_block_table({{100, 101, 102}, {200, 201, 202}}));
+  torch::Tensor tokens =
+      torch::arange(0, 12, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    0,
+                                    4);
+
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+  EXPECT_EQ(context.local_attn_metadata.q_cu_seq_lens.size(0), 5);
+  EXPECT_TRUE(
+      torch::equal(context.seg_ctx_lens_1col,
+                   torch::tensor({17, 16, 17, 24},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_EQ(context.req_ctx_offsets_cpu, (std::vector<int32_t>{0, 20}));
+}
+
+TEST(DeepseekV32SPUtilsTest, BuildSPContextTracksMixedHitMiss) {
+  AttentionMetadata attn_metadata = make_prefill_metadata(
+      {8, 8}, {8, 20}, make_block_table({{1, 2, 3}, {4, 5, 6}}));
+  torch::Tensor tokens =
+      torch::arange(0, 16, torch::TensorOptions().dtype(torch::kInt32));
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    2,
+                                    4);
+
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+
+  EXPECT_EQ(extract_segment_q_tokens(context.local_segments),
+            (std::vector<int32_t>{1, 1, 1, 1}));
+  EXPECT_TRUE(
+      torch::equal(context.seg_suffix_k_cu_lens_2col,
+                   torch::tensor({{0, 3}, {0, 6}, {0, 3}, {0, 6}},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(context.seg_ctx_lens_1col,
+                   torch::tensor({3, 6, 15, 18},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildLocalPrefillMetadataKeepsSuffixOnlyLensForChunkedPrefill) {
+  AttentionMetadata attn_metadata = make_prefill_metadata(
+      {4, 6}, {16, 22}, std::nullopt, BatchForwardType::CHUNKED_PREFILL);
+  const auto all_segments = build_all_sp_segments(4, {4, 6}, {16, 22});
+  const auto local_segments = build_local_sp_segments(1, all_segments);
+  const auto local_attn_metadata =
+      build_local_prefill_attention_metadata(attn_metadata, local_segments);
+
+  EXPECT_TRUE(
+      torch::equal(local_attn_metadata.q_cu_seq_lens,
+                   torch::tensor({0, 1, 1, 2, 3},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(local_attn_metadata.kv_cu_seq_lens,
+                   torch::tensor({0, 2, 2, 4, 9},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_TRUE(
+      torch::equal(local_attn_metadata.kv_seq_lens,
+                   torch::tensor({2, 0, 2, 5},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+  EXPECT_EQ(local_attn_metadata.max_query_len, 1);
+  EXPECT_EQ(local_attn_metadata.max_seq_len, 5);
+  EXPECT_FALSE(
+      torch::equal(local_attn_metadata.kv_seq_lens,
+                   torch::tensor({14, 12, 18, 21},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, SliceLocalPackedUsesRankOffset) {
+  DeepseekV32SPContext context;
+  context.rank = 1;
+  context.comm_plan.tokens_per_rank = {2, 3, 1};
+  context.comm_plan.token_num_offset = 2;
+
+  torch::Tensor packed = torch::tensor(
+      {10, 11, 20, 21, 22, 30}, torch::TensorOptions().dtype(torch::kInt32));
+  torch::Tensor local = slice_local_packed(packed, context);
+
+  EXPECT_TRUE(
+      torch::equal(local,
+                   torch::tensor({20, 21, 22},
+                                 torch::TensorOptions().dtype(torch::kInt32))));
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     BuildDeepseekV32SPContextReturnsNulloptWhenDisabled) {
+  ScopedFlagValue enable_sp(FLAGS_enable_prefill_sp, false);
+  ScopedFlagValue world_size_flag(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({8});
+  torch::Tensor tokens =
+      torch::arange(0, 8, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    0,
+                                    4);
+
+  EXPECT_FALSE(maybe_context.has_value());
+}
+
+TEST(DeepseekV32SPUtilsTest, ReorderByIndexSlicesFirstDimension) {
+  torch::Tensor hidden_states = torch::tensor(
+      {{0.0f, 10.0f}, {1.0f, 11.0f}, {2.0f, 12.0f}, {3.0f, 13.0f}},
+      torch::TensorOptions().dtype(torch::kFloat32));
+  torch::Tensor token_index =
+      torch::tensor({3, 1}, torch::TensorOptions().dtype(torch::kInt32));
+
+  torch::Tensor reordered = reorder_by_index(hidden_states, token_index);
+
+  EXPECT_TRUE(torch::equal(
+      reordered,
+      torch::tensor({{3.0f, 13.0f}, {1.0f, 11.0f}},
+                    torch::TensorOptions().dtype(torch::kFloat32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, RestoreGatheredToGlobalOrderWithoutPadding) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({16});
+  torch::Tensor tokens =
+      torch::arange(0, 16, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    0,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+
+  torch::Tensor gathered =
+      context.gathered_reorder_index.to(torch::kFloat32).contiguous();
+  torch::Tensor restored = restore_gathered_to_global_order(gathered, context);
+
+  EXPECT_TRUE(torch::equal(
+      restored,
+      torch::arange(0, 16, torch::TensorOptions().dtype(torch::kFloat32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, AllGatherAcrossRanksRestoresGlobalOrder) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({10});
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  auto context = maybe_context.value();
+
+  const int64_t padded_token_num = context.comm_plan.padded_tokens_per_rank[0];
+  std::vector<torch::Tensor> scripted_outputs;
+  scripted_outputs.reserve(4);
+  auto across_world =
+      context.gathered_reorder_index.to(torch::kCPU).contiguous();
+  const auto* across_world_ptr = across_world.data_ptr<int64_t>();
+  int32_t token_offset = 0;
+  for (int64_t world_rank = 0; world_rank < 4; ++world_rank) {
+    std::vector<float> values(padded_token_num, -1.0f);
+    const int32_t valid_token_num =
+        context.comm_plan.tokens_per_rank[world_rank];
+    for (int32_t i = 0; i < valid_token_num; ++i) {
+      values[i] = static_cast<float>(across_world_ptr[token_offset + i]);
+    }
+    token_offset += valid_token_num;
+    scripted_outputs.push_back(
+        torch::tensor(values, torch::TensorOptions().dtype(torch::kFloat32)));
+  }
+
+  ScriptedAllGatherProcessGroup scripted_group(
+      torch::Device(torch::kCPU), 1, std::move(scripted_outputs));
+  context.process_group = &scripted_group;
+  torch::Tensor local_tensor = torch::tensor(
+      {2.0f, 3.0f, 8.0f}, torch::TensorOptions().dtype(torch::kFloat32));
+
+  torch::Tensor gathered = all_gather_across_ranks(local_tensor, context);
+  torch::Tensor restored = restore_gathered_to_global_order(gathered, context);
+
+  EXPECT_TRUE(torch::equal(
+      restored,
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kFloat32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, GatherSupportsUnevenThreeDimensionalInputs) {
+  std::vector<torch::Tensor> scripted_outputs;
+  scripted_outputs.reserve(4);
+  for (int64_t world_rank = 0; world_rank < 4; ++world_rank) {
+    scripted_outputs.push_back(
+        torch::full({3, 2, 2},
+                    static_cast<float>(world_rank),
+                    torch::TensorOptions().dtype(torch::kFloat32)));
+  }
+
+  ScriptedAllGatherProcessGroup process_group(
+      torch::Device(torch::kCPU), 1, std::move(scripted_outputs));
+  torch::Tensor local_tensor = torch::full(
+      {2, 2, 2}, 1.0f, torch::TensorOptions().dtype(torch::kFloat32));
+
+  torch::Tensor gathered = xllm::parallel_state::gather(
+      local_tensor, &process_group, std::vector<int32_t>{3, 2, 3, 1});
+
+  EXPECT_EQ(gathered.sizes(), (torch::IntArrayRef{9, 2, 2}));
+  EXPECT_TRUE(torch::allclose(gathered[0], torch::zeros({2, 2})));
+  EXPECT_TRUE(torch::allclose(gathered[2], torch::zeros({2, 2})));
+  EXPECT_TRUE(torch::allclose(gathered[3], torch::ones({2, 2})));
+  EXPECT_TRUE(torch::allclose(gathered[4], torch::ones({2, 2})));
+  EXPECT_TRUE(torch::allclose(gathered[5], torch::full({2, 2}, 2.0f)));
+  EXPECT_TRUE(torch::allclose(gathered[7], torch::full({2, 2}, 2.0f)));
+  EXPECT_TRUE(torch::allclose(gathered[8], torch::full({2, 2}, 3.0f)));
+}
+
+TEST(DeepseekV32SPUtilsTest, LaunchAndFinishGatherMatchBlockingGather) {
+  std::vector<torch::Tensor> scripted_outputs;
+  scripted_outputs.reserve(4);
+  for (int64_t world_rank = 0; world_rank < 4; ++world_rank) {
+    scripted_outputs.push_back(
+        torch::full({3, 2},
+                    static_cast<float>(world_rank),
+                    torch::TensorOptions().dtype(torch::kFloat32)));
+  }
+
+  ScriptedAllGatherProcessGroup process_group(
+      torch::Device(torch::kCPU), 1, std::move(scripted_outputs));
+  torch::Tensor local_tensor =
+      torch::full({2, 2}, 1.0f, torch::TensorOptions().dtype(torch::kFloat32));
+
+  auto gather_ctx = xllm::parallel_state::launch_gather(
+      local_tensor, &process_group, std::vector<int32_t>{3, 2, 3, 1});
+  torch::Tensor gathered = xllm::parallel_state::finish_gather(gather_ctx);
+  auto float_options = torch::TensorOptions().dtype(torch::kFloat32);
+
+  EXPECT_EQ(gathered.sizes(), (torch::IntArrayRef{9, 2}));
+  EXPECT_TRUE(torch::allclose(gathered[0], torch::zeros({2}, float_options)));
+  EXPECT_TRUE(torch::allclose(gathered[3], torch::ones({2}, float_options)));
+  EXPECT_TRUE(
+      torch::allclose(gathered[5], torch::full({2}, 2.0f, float_options)));
+  EXPECT_TRUE(
+      torch::allclose(gathered[8], torch::full({2}, 3.0f, float_options)));
+}
+
+TEST(DeepseekV32SPUtilsTest, PadToSPRowsExpandsLocalShardToPaddedRows) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({10});
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  auto context = maybe_context.value();
+
+  torch::Tensor local_tensor = torch::tensor(
+      {2.0f, 3.0f, 8.0f}, torch::TensorOptions().dtype(torch::kFloat32));
+  torch::Tensor padded = pad_to_sp_rows(local_tensor, context);
+
+  EXPECT_EQ(padded.size(0), context.comm_plan.padded_tokens_per_rank[1]);
+  EXPECT_TRUE(
+      torch::allclose(padded.slice(0, 0, local_tensor.size(0)), local_tensor));
+  EXPECT_TRUE(torch::allclose(
+      padded.slice(0, local_tensor.size(0), padded.size(0)),
+      torch::zeros({padded.size(0) - local_tensor.size(0)},
+                   torch::TensorOptions().dtype(torch::kFloat32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, SPContextBuildsGatheredSlotMapping) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({10});
+  attn_metadata.slot_mapping =
+      torch::arange(10, 20, torch::TensorOptions().dtype(torch::kInt32));
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  auto context = maybe_context.value();
+
+  torch::Tensor expected = reorder_by_index(attn_metadata.slot_mapping,
+                                            context.gathered_reorder_index);
+  EXPECT_TRUE(torch::equal(context.gathered_slot_mapping, expected));
+}
+
+TEST(DeepseekV32SPUtilsTest,
+     NonChunkedSegmentSliceRequiresRestoredGlobalDenseOrder) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 2);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({4});
+  torch::Tensor tokens =
+      torch::arange(0, 4, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    0,
+                                    2);
+  ASSERT_TRUE(maybe_context.has_value());
+  const auto& context = maybe_context.value();
+
+  ASSERT_EQ(context.local_segments.size(), 2U);
+  const auto& right_segment = context.local_segments[1];
+  EXPECT_EQ(right_segment.q_tokens, 1);
+  EXPECT_EQ(right_segment.suffix_k_len, 4);
+  EXPECT_EQ(context.req_q_offsets_cpu[right_segment.req_idx], 0);
+
+  torch::Tensor k_gathered = context.gathered_reorder_index.to(
+      torch::TensorOptions().dtype(torch::kFloat32));
+  torch::Tensor k_restored =
+      restore_gathered_to_global_order(k_gathered, context);
+
+  const int32_t req_q_start = context.req_q_offsets_cpu[right_segment.req_idx];
+  torch::Tensor packed_slice =
+      k_gathered.narrow(0, req_q_start, right_segment.suffix_k_len);
+  torch::Tensor restored_slice =
+      k_restored.narrow(0, req_q_start, right_segment.suffix_k_len);
+  torch::Tensor expected =
+      torch::arange(0, 4, torch::TensorOptions().dtype(torch::kFloat32));
+
+  EXPECT_FALSE(torch::equal(packed_slice, expected));
+  EXPECT_TRUE(torch::equal(restored_slice, expected));
+}
+
+TEST(DeepseekV32SPUtilsTest, AsyncGatherReturnsGatheredOrder) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({10});
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  auto context = maybe_context.value();
+
+  const int64_t padded_token_num = context.comm_plan.padded_tokens_per_rank[0];
+  std::vector<torch::Tensor> scripted_outputs;
+  scripted_outputs.reserve(4);
+  auto across_world =
+      context.gathered_reorder_index.to(torch::kCPU).contiguous();
+  const auto* across_world_ptr = across_world.data_ptr<int64_t>();
+  int32_t token_offset = 0;
+  for (int64_t world_rank = 0; world_rank < 4; ++world_rank) {
+    std::vector<float> values(padded_token_num, -1.0f);
+    const int32_t valid_token_num =
+        context.comm_plan.tokens_per_rank[world_rank];
+    for (int32_t i = 0; i < valid_token_num; ++i) {
+      values[i] = static_cast<float>(across_world_ptr[token_offset + i]);
+    }
+    token_offset += valid_token_num;
+    scripted_outputs.push_back(
+        torch::tensor(values, torch::TensorOptions().dtype(torch::kFloat32)));
+  }
+
+  ScriptedAllGatherProcessGroup scripted_group(
+      torch::Device(torch::kCPU), 1, std::move(scripted_outputs));
+  context.process_group = &scripted_group;
+  torch::Tensor local_tensor = torch::tensor(
+      {2.0f, 3.0f, 8.0f}, torch::TensorOptions().dtype(torch::kFloat32));
+
+  auto gather_handle = xllm::parallel_state::launch_gather(
+      local_tensor, &scripted_group, context.comm_plan.tokens_per_rank);
+  torch::Tensor gathered = xllm::parallel_state::finish_gather(gather_handle);
+  torch::Tensor restored = restore_gathered_to_global_order(gathered, context);
+
+  EXPECT_TRUE(torch::equal(gathered,
+                           context.gathered_reorder_index.to(
+                               torch::TensorOptions().dtype(torch::kFloat32))));
+  EXPECT_TRUE(torch::equal(
+      restored,
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kFloat32))));
+}
+
+TEST(DeepseekV32SPUtilsTest, AsyncGatherHelpersMatchBlockingGather) {
+  ScopedFlagValue use_sp(FLAGS_enable_prefill_sp, true);
+  ScopedFlagValue nnodes(FLAGS_nnodes, 4);
+
+  AttentionMetadata attn_metadata = make_prefill_metadata({10});
+  torch::Tensor tokens =
+      torch::arange(0, 10, torch::TensorOptions().dtype(torch::kInt32));
+  auto maybe_context =
+      build_deepseek_v32_sp_context(attn_metadata,
+                                    BatchForwardType::PREFILL,
+                                    tokens,
+                                    reinterpret_cast<ProcessGroup*>(0x1),
+                                    1,
+                                    4);
+  ASSERT_TRUE(maybe_context.has_value());
+  auto context = maybe_context.value();
+
+  const int64_t padded_token_num = context.comm_plan.padded_tokens_per_rank[0];
+  std::vector<torch::Tensor> scripted_outputs;
+  scripted_outputs.reserve(4);
+  for (int64_t world_rank = 0; world_rank < 4; ++world_rank) {
+    scripted_outputs.push_back(
+        torch::full({padded_token_num, 2},
+                    static_cast<float>(world_rank),
+                    torch::TensorOptions().dtype(torch::kFloat32)));
+  }
+
+  ScriptedAllGatherProcessGroup scripted_group(
+      torch::Device(torch::kCPU), 1, std::move(scripted_outputs));
+  context.process_group = &scripted_group;
+  torch::Tensor local_tensor =
+      torch::full({context.comm_plan.tokens_per_rank[context.rank], 2},
+                  1.0f,
+                  torch::TensorOptions().dtype(torch::kFloat32));
+
+  auto gather_handle = launch_all_gather_across_ranks(local_tensor, context);
+  auto gathered_async =
+      finish_all_gather_across_ranks(std::move(gather_handle));
+  auto gathered_sync = all_gather_across_ranks(local_tensor, context);
+
+  EXPECT_TRUE(torch::equal(gathered_async, gathered_sync));
+}
+
+}  // namespace
+}  // namespace xllm::layer::v32_sp

upstream_ref/xllm/tests/core/layers/mlu/dense_mlp_test.cpp

@@ -0,0 +1,441 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/common/dense_mlp.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/common/linear.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+class DenseMLPTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Initialize default model arguments for testing
+    model_args_ = test::create_default_model_args();
+
+    // Initialize w8a8 quantization arguments
+    quant_args_ = test::create_default_quant_args();
+
+    // Initialize tensor options
+    options_ = torch::TensorOptions()
+                   .dtype(torch::kBFloat16)
+                   .device(Device::type_torch(), 0)
+                   .requires_grad(false);
+
+    // Create mock ProcessGroup and initialize ParallelArgs
+    parallel_args_ = test::create_default_parallel_args(mock_process_group_);
+
+    // Note: MLP will be created by individual test cases with their desired
+    // dimensions
+  }
+
+  void TearDown() override {
+    // Clean up if needed
+  }
+
+  std::unordered_map<std::string, torch::Tensor> create_default_test_weights(
+      int64_t hidden_size,
+      int64_t intermediate_size) {
+    // Create test weights for gate_up_proj (gate + up projection)
+    // Shape: [intermediate_size * 2, hidden_size]
+    auto gate_up_weight =
+        torch::full({intermediate_size * 2, hidden_size}, 5.0f, options_);
+
+    // Create test weights for down_proj (down projection)
+    // Shape: [hidden_size, intermediate_size] for RowParallelLinear
+    auto down_weight =
+        torch::full({hidden_size, intermediate_size}, 3.0f, options_);
+
+    // For w8a8 smoothquant, we need to create quantized weights and scales
+    // Create qweight (int8 quantized weights)
+    auto gate_up_qweight = gate_up_weight.to(torch::kInt8);
+    auto down_qweight = down_weight.to(torch::kInt8);
+
+    // Create per_channel_scale (float32 scales for each channel)
+    // For ColumnParallelLinear: per_channel_scale shape is
+    // [out_features_per_partition]
+    auto gate_up_scale = torch::full({intermediate_size},
+                                     0.1f,
+                                     torch::TensorOptions()
+                                         .dtype(torch::kFloat32)
+                                         .device(options_.device()));
+    auto up_scale = torch::full({intermediate_size},
+                                0.1f,
+                                torch::TensorOptions()
+                                    .dtype(torch::kFloat32)
+                                    .device(options_.device()));
+    auto down_scale = torch::full({hidden_size},
+                                  0.1f,
+                                  torch::TensorOptions()
+                                      .dtype(torch::kFloat32)
+                                      .device(options_.device()));
+
+    // Create smooth quantization scales
+    // For ColumnParallelLinear: smooth shape is [in_features] = [hidden_size]
+    auto gate_up_smooth = torch::full({hidden_size},
+                                      0.05f,
+                                      torch::TensorOptions()
+                                          .dtype(torch::kFloat32)
+                                          .device(options_.device()));
+    auto up_smooth = torch::full({hidden_size},
+                                 0.05f,
+                                 torch::TensorOptions()
+                                     .dtype(torch::kFloat32)
+                                     .device(options_.device()));
+    // For RowParallelLinear: smooth shape is [in_features_per_partition] =
+    // [intermediate_size]
+    auto down_smooth = torch::full({intermediate_size},
+                                   0.05f,
+                                   torch::TensorOptions()
+                                       .dtype(torch::kFloat32)
+                                       .device(options_.device()));
+
+    // Create StateDict with w8a8 smoothquant weights
+    std::unordered_map<std::string, torch::Tensor> weight_dict;
+
+    // Gate projection weights (ColumnParallelLinear)
+    // gate_up_qweight shape: [intermediate_size * 2, hidden_size]
+    // gate_proj.qweight should be first intermediate_size rows:
+    // [intermediate_size, hidden_size]
+    weight_dict["gate_proj.qweight"] =
+        gate_up_qweight.slice(0, 0, intermediate_size);
+    weight_dict["gate_proj.per_channel_scale"] = gate_up_scale;
+    weight_dict["gate_proj.smooth"] = gate_up_smooth;
+
+    // Up projection weights (ColumnParallelLinear)
+    // up_proj.qweight should be last intermediate_size rows:
+    // [intermediate_size, hidden_size]
+    weight_dict["up_proj.qweight"] =
+        gate_up_qweight.slice(0, intermediate_size, intermediate_size * 2);
+    weight_dict["up_proj.per_channel_scale"] = up_scale;
+    weight_dict["up_proj.smooth"] = up_smooth;
+
+    // Down projection weights (RowParallelLinear)
+    weight_dict["down_proj.qweight"] = down_qweight;
+    weight_dict["down_proj.per_channel_scale"] = down_scale;
+    weight_dict["down_proj.smooth"] = down_smooth;
+
+    LOG(INFO) << "Test w8a8 smoothquant weights created successfully";
+    LOG(INFO) << "Gate qweight shape: "
+              << weight_dict["gate_proj.qweight"].sizes();
+    LOG(INFO) << "Gate per_channel_scale shape: "
+              << weight_dict["gate_proj.per_channel_scale"].sizes();
+    LOG(INFO) << "Gate smooth shape: "
+              << weight_dict["gate_proj.smooth"].sizes();
+    LOG(INFO) << "Up qweight shape: " << weight_dict["up_proj.qweight"].sizes();
+    LOG(INFO) << "Up per_channel_scale shape: "
+              << weight_dict["up_proj.per_channel_scale"].sizes();
+    LOG(INFO) << "Up smooth shape: " << weight_dict["up_proj.smooth"].sizes();
+    LOG(INFO) << "Down qweight shape: "
+              << weight_dict["down_proj.qweight"].sizes();
+    LOG(INFO) << "Down per_channel_scale shape: "
+              << weight_dict["down_proj.per_channel_scale"].sizes();
+    LOG(INFO) << "Down smooth shape: "
+              << weight_dict["down_proj.smooth"].sizes();
+
+    return weight_dict;
+  }
+
+  // Helper function to create MLP with custom dimensions
+  DenseMLP create_mlp(int64_t hidden_size, int64_t intermediate_size) {
+    // Create MLP with specified dimensions using the new constructor
+    return DenseMLP(DenseMLPImpl(hidden_size,
+                                 intermediate_size,
+                                 /*is_gated=*/true,
+                                 /*has_bias=*/false,
+                                 /*hidden_act=*/"silu",
+                                 /*enable_result_reduction=*/true,
+                                 quant_args_,
+                                 parallel_args_.tp_group_,
+                                 options_));
+  }
+
+  // Helper function to create test weights for the MLP (w8a8 smoothquant
+  // format)
+  std::unordered_map<std::string, torch::Tensor> create_test_weights(
+      int64_t custom_hidden_size = -1,
+      int64_t custom_intermediate_size = -1) {
+    // Use custom sizes if provided, otherwise use model_args_ values
+    int64_t test_hidden_size = (custom_hidden_size > 0)
+                                   ? custom_hidden_size
+                                   : model_args_.hidden_size();
+    int64_t test_intermediate_size = (custom_intermediate_size > 0)
+                                         ? custom_intermediate_size
+                                         : model_args_.intermediate_size();
+
+    return create_default_test_weights(test_hidden_size,
+                                       test_intermediate_size);
+  }
+
+  // Helper function to verify tensor values are close to expected
+  void verify_tensor_close(const torch::Tensor& actual,
+                           const torch::Tensor& expected,
+                           double rtol = 1e-5,
+                           double atol = 1e-8) {
+    test::verify_tensor_close(actual, expected, rtol, atol);
+  }
+
+  // Helper function to create custom input tensor for precision testing
+  torch::Tensor create_custom_input(const std::vector<int64_t>& shape,
+                                    const std::vector<float>& values) {
+    return test::create_custom_input(shape, values, options_);
+  }
+
+  // Helper function to set expected output for precision verification
+  void set_expected_output(const std::vector<float>& expected_values) {
+    expected_output_ = expected_values;
+  }
+
+  // Helper function to verify precision against expected output
+  void verify_precision(const torch::Tensor& actual_output,
+                        double rtol = 1e-3,
+                        double atol = 1e-4) {
+    test::verify_precision(actual_output, expected_output_, rtol, atol);
+  }
+
+  ModelArgs model_args_;
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+
+  // Helper to create a mock ProcessGroup for testing
+  std::unique_ptr<xllm::ProcessGroup> mock_process_group_;
+
+  // Expected output for precision verification
+  std::vector<float> expected_output_;
+};
+
+TEST_F(DenseMLPTest, Bfloat16LoadStateDictTest) {
+  // Test bfloat16 mode (non-quantized)
+  const int64_t batch_size = 8;
+  const int64_t hidden_size = 1024;
+  const int64_t intermediate_size = 2048;
+
+  // Create non-quantized quant_args for bfloat16 mode
+  QuantArgs bfloat16_quant_args;  // Empty means no quantization
+
+  // Create MLP in bfloat16 mode
+  auto mlp = DenseMLP(DenseMLPImpl(hidden_size,
+                                   intermediate_size,
+                                   /*is_gated=*/true,
+                                   /*has_bias=*/false,
+                                   /*hidden_act=*/"silu",
+                                   /*enable_result_reduction=*/true,
+                                   bfloat16_quant_args,
+                                   parallel_args_.tp_group_,
+                                   options_));
+
+  // Create simple test weights for bfloat16 mode
+  std::unordered_map<std::string, torch::Tensor> weight_dict;
+
+  // For bfloat16 mode, we need regular weight tensors
+  auto gate_up_weight =
+      torch::full({intermediate_size * 2, hidden_size}, 0.1f, options_);
+  auto down_weight =
+      torch::full({hidden_size, intermediate_size}, 0.1f, options_);
+
+  weight_dict["gate_proj.weight"] =
+      gate_up_weight.slice(0, 0, intermediate_size);
+  weight_dict["up_proj.weight"] =
+      gate_up_weight.slice(0, intermediate_size, intermediate_size * 2);
+  weight_dict["down_proj.weight"] = down_weight;
+
+  // Load weights
+  StateDict state_dict(weight_dict);
+  mlp->load_state_dict(state_dict);
+
+  // Test forward pass
+  auto hidden_states = torch::ones({batch_size, hidden_size}, options_);
+  auto output = mlp->forward(hidden_states);
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes().size(), 2) << "Output should be 2D tensor";
+  ASSERT_EQ(output.size(0), batch_size) << "Batch size should match";
+  ASSERT_EQ(output.size(1), hidden_size) << "Hidden size should match";
+
+  // Verify output is not all zeros
+  auto output_sum = torch::sum(output).item<float>();
+  ASSERT_NE(output_sum, 0.0f) << "Output should not be all zeros";
+
+  LOG(INFO) << "Bfloat16 mode test passed - output sum: " << output_sum;
+}
+
+TEST_F(DenseMLPTest, SmoothquantLoadStateDictTest) {
+  // Test loading weights into the MLP
+  const int64_t batch_size = 16;
+  const int64_t hidden_size = model_args_.hidden_size();
+  const int64_t intermediate_size = model_args_.intermediate_size();
+
+  // Create MLP with default dimensions
+  auto mlp = create_mlp(hidden_size, intermediate_size);
+
+  // Create test weights and load them (using default model_args_ dimensions)
+  auto weight_dict = create_test_weights();
+
+  // Load weights into the MLP
+  StateDict state_dict(weight_dict);
+  mlp->load_state_dict(state_dict);
+
+  // Test forward pass with loaded weights
+  auto hidden_states = torch::ones({batch_size, hidden_size}, options_);
+
+  LOG(INFO) << "Testing forward pass with loaded weights";
+  auto output = mlp->forward(hidden_states);
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes().size(), 2) << "Output should be 2D tensor";
+  ASSERT_EQ(output.size(0), batch_size) << "Batch size should match";
+  ASSERT_EQ(output.size(1), hidden_size) << "Hidden size should match";
+
+  // Verify output is not all zeros (weights were loaded)
+  auto output_sum = torch::sum(output).item<float>();
+  ASSERT_NE(output_sum, 0.0f)
+      << "Output should not be all zeros after loading weights";
+
+  LOG(INFO) << "State dict loading test passed - output sum: " << output_sum;
+}
+
+#if defined(USE_CUDA)
+TEST_F(DenseMLPTest, Fp8IgnoredDownProjLoadsAsUnquantized) {
+  QuantArgs fp8_quant_args;
+  fp8_quant_args.quant_method() = kQuantMethodFp8;
+  fp8_quant_args.bits() = 8;
+  fp8_quant_args.activation_dynamic() = false;
+  fp8_quant_args.ignored_modules() = {"model.layers.1.mlp.down_proj"};
+
+  const int64_t hidden_size = 16;
+  const int64_t intermediate_size = 32;
+  auto mlp = DenseMLP(DenseMLPImpl(hidden_size,
+                                   intermediate_size,
+                                   /*is_gated=*/true,
+                                   /*has_bias=*/false,
+                                   /*hidden_act=*/"silu",
+                                   /*enable_result_reduction=*/true,
+                                   fp8_quant_args,
+                                   parallel_args_.tp_group_,
+                                   options_,
+                                   "model.layers.1.mlp"));
+
+  std::unordered_map<std::string, torch::Tensor> weight_dict;
+  auto fp8_weight_options = options_.dtype(torch::kFloat8_e4m3fn);
+  auto scale_options = options_.dtype(torch::kFloat32);
+
+  weight_dict["gate_proj.weight"] =
+      torch::zeros({intermediate_size, hidden_size}, fp8_weight_options);
+  weight_dict["gate_proj.weight_scale"] = torch::ones({1}, scale_options);
+  weight_dict["gate_proj.input_scale"] = torch::ones({1}, scale_options);
+
+  weight_dict["up_proj.weight"] =
+      torch::zeros({intermediate_size, hidden_size}, fp8_weight_options);
+  weight_dict["up_proj.weight_scale"] = torch::ones({1}, scale_options);
+  weight_dict["up_proj.input_scale"] = torch::ones({1}, scale_options);
+
+  weight_dict["down_proj.weight"] =
+      torch::zeros({hidden_size, intermediate_size}, options_);
+
+  StateDict state_dict(weight_dict);
+  mlp->load_state_dict(state_dict);
+
+  const auto params = mlp->named_parameters(/*recurse=*/true);
+  EXPECT_TRUE(params.contains("gate_up_proj.weight_scale"));
+  EXPECT_TRUE(params.contains("gate_up_proj.input_scale"));
+  EXPECT_TRUE(params.contains("down_proj.weight"));
+  EXPECT_FALSE(params.contains("down_proj.weight_scale"));
+  EXPECT_FALSE(params.contains("down_proj.input_scale"));
+}
+#endif
+
+TEST_F(DenseMLPTest, SmoothquantPrecisionVerificationTest) {
+  // Test precision verification with custom input and expected output
+  const int64_t batch_size = 16;
+
+  // Use custom hidden and intermediate sizes for more controllable output
+  const int64_t custom_hidden_size = 7168;
+  const int64_t custom_intermediate_size = 9216;
+
+  // Create custom MLP with smaller dimensions
+  auto custom_mlp = create_mlp(custom_hidden_size, custom_intermediate_size);
+
+  // Create test weights and load them with custom dimensions
+  auto weight_dict =
+      create_test_weights(custom_hidden_size, custom_intermediate_size);
+
+  // Load weights into the MLP
+  StateDict state_dict(weight_dict);
+  custom_mlp->load_state_dict(state_dict);
+
+  // Create custom input tensor for precision testing
+  // You can modify these values as needed for your specific test case
+  std::vector<float> input_values;
+
+  // Fill input_values with test data using custom dimensions
+  input_values.reserve(batch_size * custom_hidden_size);
+  // Populate the input_values vector:
+  for (size_t i = 0; i < batch_size; ++i) {
+    float value = 0.5f;
+    for (size_t j = 0; j < custom_hidden_size; ++j) {
+      input_values.push_back(value);
+    }
+  }
+
+  auto hidden_states =
+      create_custom_input({batch_size, custom_hidden_size}, input_values);
+
+  LOG(INFO) << "Testing precision verification with custom input (hidden_size="
+            << custom_hidden_size << ")";
+  auto output = custom_mlp->forward(hidden_states);
+
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes().size(), 2) << "Output should be 2D tensor";
+  ASSERT_EQ(output.size(0), batch_size) << "Batch size should match";
+  ASSERT_EQ(output.size(1), custom_hidden_size) << "Hidden size should match";
+
+  // Set expected output values for precision verification
+  // TODO: Replace these placeholder values with your expected output
+  // The expected values should be calculated based on your specific test case
+  std::vector<float> expected_values;
+
+  // Fill expected_values with placeholder data using custom dimensions
+  expected_values.reserve(batch_size * custom_hidden_size);
+  for (size_t i = 0; i < batch_size; ++i) {
+    for (size_t j = 0; j < custom_hidden_size; ++j) {
+      expected_values.push_back(1105920.0f);  // calculated via vLLM MLU
+    }
+  }
+
+  set_expected_output(expected_values);
+
+  // Note: The precision verification is commented out until you set the
+  // expected values Uncomment the following line after setting the correct
+  verify_precision(output, 1e-3, 1e-4);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/dp_utils_test.cpp

@@ -0,0 +1,343 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/common/dp_utils.h"
+
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <memory>
+
+#include "layers/mlu/tests_utils.h"
+
+namespace xllm {
+namespace layer {
+namespace test {
+
+ParallelArgs make_tp_args(std::unique_ptr<xllm::ProcessGroup>& process_group,
+                          int32_t rank,
+                          int32_t tp_size) {
+  process_group = std::make_unique<test::MockProcessGroup>(
+      torch::Device(torch::kCPU), rank, tp_size);
+  ParallelArgs args(rank, tp_size, process_group.get());
+  args.tp_group_ = process_group.get();
+  return args;
+}
+
+ParallelArgs make_world_args(std::unique_ptr<xllm::ProcessGroup>& process_group,
+                             std::unique_ptr<xllm::ProcessGroup>& tp_group,
+                             int32_t global_rank,
+                             int32_t dp_size,
+                             int32_t tp_size) {
+  process_group = std::make_unique<test::MockProcessGroup>(
+      torch::Device(torch::kCPU), global_rank, dp_size * tp_size);
+  tp_group = std::make_unique<test::MockProcessGroup>(
+      torch::Device(torch::kCPU), global_rank % tp_size, tp_size);
+  ParallelArgs args(
+      global_rank, dp_size * tp_size, dp_size, process_group.get());
+  args.tp_group_ = tp_group.get();
+  return args;
+}
+
+struct GatherShape {
+  int64_t padded_dp_tokens = 0;
+  int64_t shard_tokens = 0;
+};
+
+GatherShape make_gather_shape(const std::vector<int32_t>& dp_tokens,
+                              int32_t tp_size) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, tp_size);
+  int64_t padded_dp_tokens = get_dp_gather_tokens(dp_tokens, args);
+  return {padded_dp_tokens, padded_dp_tokens / tp_size};
+}
+
+std::vector<torch::Tensor> make_global_shards(
+    const std::vector<int32_t>& dp_tokens,
+    int32_t tp_size) {
+  GatherShape shape = make_gather_shape(dp_tokens, tp_size);
+  std::vector<torch::Tensor> shards;
+  for (size_t dp_rank = 0; dp_rank < dp_tokens.size(); ++dp_rank) {
+    std::vector<float> padded_vals(shape.padded_dp_tokens, -1.0f);
+    for (int64_t i = 0; i < dp_tokens[dp_rank]; ++i) {
+      padded_vals[i] = static_cast<float>(dp_rank * 100 + i);
+    }
+    auto padded =
+        torch::tensor(padded_vals).reshape({shape.padded_dp_tokens, 1});
+    auto split = padded.split(shape.shard_tokens, 0);
+    for (int32_t tp_rank = 0; tp_rank < tp_size; ++tp_rank) {
+      shards.push_back(split[tp_rank].clone());
+    }
+  }
+  return shards;
+}
+
+TEST(DpUtilsTest, RsAttnInputNoPadWhenTpOne) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 1);
+
+  torch::Tensor x = torch::tensor({{1.0f, 2.0f}, {3.0f, 4.0f}});
+  torch::Tensor residual = torch::tensor({{10.0f, 20.0f}, {30.0f, 40.0f}});
+
+  auto rs_result =
+      reduce_scatter_attn_input(x, residual, /*target_tokens=*/2, args);
+
+  EXPECT_FALSE(rs_result.second.active);
+  EXPECT_EQ(rs_result.second.original_tokens, 2);
+  EXPECT_EQ(rs_result.second.padded_tokens, 2);
+  test::verify_tensor_close(rs_result.first,
+                            torch::tensor({{11.0f, 22.0f}, {33.0f, 44.0f}}));
+}
+
+TEST(DpUtilsTest, GetRsTokensKeepsTpOneShape) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 1);
+
+  EXPECT_EQ(get_reduce_scatter_tokens(/*num_tokens=*/5, args), 5);
+}
+
+TEST(DpUtilsTest, GetRsTokensAlignsTpShard) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 4);
+
+  EXPECT_EQ(get_reduce_scatter_tokens(/*num_tokens=*/2, args), 4);
+  EXPECT_EQ(get_reduce_scatter_tokens(/*num_tokens=*/5, args), 8);
+  EXPECT_EQ(get_reduce_scatter_tokens(/*num_tokens=*/8, args), 8);
+}
+
+TEST(DpUtilsTest, PadTokensPreservesTrailingDims) {
+  auto x = torch::arange(12, torch::kFloat32).reshape({2, 2, 3});
+
+  auto [padded, info] = pad_tokens(x, /*target_tokens=*/4);
+
+  EXPECT_TRUE(info.active);
+  EXPECT_EQ(info.original_tokens, 2);
+  EXPECT_EQ(info.padded_tokens, 4);
+  EXPECT_EQ(padded.sizes().vec(), std::vector<int64_t>({4, 2, 3}));
+  test::verify_tensor_close(padded.slice(0, 0, 2), x);
+  test::verify_tensor_close(padded.slice(0, 2, 4),
+                            torch::zeros({2, 2, 3}, x.options()));
+}
+
+TEST(DpUtilsTest, RsAttnInputPadsWhenTokensLtTp) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 4);
+
+  torch::Tensor x = torch::tensor({{1.0f, 2.0f}, {3.0f, 4.0f}});
+  torch::Tensor residual = torch::tensor({{10.0f, 20.0f}, {30.0f, 40.0f}});
+
+  auto rs_result =
+      reduce_scatter_attn_input(x, residual, /*target_tokens=*/4, args);
+
+  EXPECT_TRUE(rs_result.second.active);
+  EXPECT_EQ(rs_result.second.original_tokens, 2);
+  EXPECT_EQ(rs_result.second.padded_tokens, 4);
+  test::verify_tensor_close(rs_result.first, torch::tensor({{11.0f, 22.0f}}));
+}
+
+TEST(DpUtilsTest, RsAttnInputPadsToAlignedTokens) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 1, 4);
+
+  torch::Tensor x = torch::tensor(
+      {{1.0f, 1.5f}, {2.0f, 2.5f}, {3.0f, 3.5f}, {4.0f, 4.5f}, {5.0f, 5.5f}});
+  torch::Tensor residual = torch::full({5, 2}, 100.0f);
+
+  auto rs_result =
+      reduce_scatter_attn_input(x, residual, /*target_tokens=*/8, args);
+
+  EXPECT_TRUE(rs_result.second.active);
+  EXPECT_EQ(rs_result.second.original_tokens, 5);
+  EXPECT_EQ(rs_result.second.padded_tokens, 8);
+  test::verify_tensor_close(rs_result.first,
+                            torch::tensor({{3.0f, 3.5f}, {4.0f, 4.5f}}));
+}
+
+TEST(DpUtilsTest, GetDpGatherTokensAlignsMixedDpTokens) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 4);
+
+  EXPECT_EQ(get_dp_gather_tokens({513, 1, 128}, args), 516);
+}
+
+TEST(DpUtilsTest, GetDpGatherTokensKeepsShardTokensPositive) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 4);
+  GatherShape shape = make_gather_shape({0, 1, 33}, /*tp_size=*/4);
+
+  EXPECT_EQ(get_dp_gather_tokens({0, 1, 33}, args), 36);
+  EXPECT_EQ(shape.shard_tokens, 9);
+  EXPECT_GT(shape.shard_tokens, 0);
+}
+
+TEST(DpUtilsTest, GatherGlobalTokensTpOneHandlesEqualDpTokens) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  std::unique_ptr<xllm::ProcessGroup> tp_group;
+  ParallelArgs args =
+      make_world_args(process_group, tp_group, /*global_rank=*/2, 3, 1);
+  const std::vector<int32_t> dp_tokens = {4, 4, 4};
+
+  auto* mock_pg = dynamic_cast<test::MockProcessGroup*>(process_group.get());
+  ASSERT_NE(mock_pg, nullptr);
+  auto shards = make_global_shards(dp_tokens, /*tp_size=*/1);
+  mock_pg->set_allgather_outputs(shards);
+
+  auto output = gather_global_tokens(shards[2], dp_tokens, args);
+
+  EXPECT_EQ(output.size(0), 12);
+  test::verify_tensor_close(output,
+                            torch::tensor({{0.0f},
+                                           {1.0f},
+                                           {2.0f},
+                                           {3.0f},
+                                           {100.0f},
+                                           {101.0f},
+                                           {102.0f},
+                                           {103.0f},
+                                           {200.0f},
+                                           {201.0f},
+                                           {202.0f},
+                                           {203.0f}}));
+}
+
+TEST(DpUtilsTest, GatherGlobalTokensTpOneHandlesUnevenDpTokens) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  std::unique_ptr<xllm::ProcessGroup> tp_group;
+  ParallelArgs args =
+      make_world_args(process_group, tp_group, /*global_rank=*/1, 3, 1);
+  const std::vector<int32_t> dp_tokens = {4, 0, 2};
+
+  auto* mock_pg = dynamic_cast<test::MockProcessGroup*>(process_group.get());
+  ASSERT_NE(mock_pg, nullptr);
+  auto shards = make_global_shards(dp_tokens, /*tp_size=*/1);
+  mock_pg->set_allgather_outputs(shards);
+
+  auto output = gather_global_tokens(shards[1], dp_tokens, args);
+
+  EXPECT_EQ(output.size(0), 6);
+  test::verify_tensor_close(
+      output,
+      torch::tensor({{0.0f}, {1.0f}, {2.0f}, {3.0f}, {200.0f}, {201.0f}}));
+}
+
+TEST(DpUtilsTest, GatherGlobalTokensRestoresDpMajorOrder) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  std::unique_ptr<xllm::ProcessGroup> tp_group;
+  ParallelArgs args =
+      make_world_args(process_group, tp_group, /*global_rank=*/1, 3, 2);
+  const std::vector<int32_t> dp_tokens = {5, 2, 4};
+
+  auto* mock_pg = dynamic_cast<test::MockProcessGroup*>(process_group.get());
+  ASSERT_NE(mock_pg, nullptr);
+  auto shards = make_global_shards(dp_tokens, /*tp_size=*/2);
+  mock_pg->set_allgather_outputs(shards);
+
+  auto output = gather_global_tokens(shards[1], dp_tokens, args);
+
+  test::verify_tensor_close(output,
+                            torch::tensor({{0.0f},
+                                           {1.0f},
+                                           {2.0f},
+                                           {3.0f},
+                                           {4.0f},
+                                           {100.0f},
+                                           {101.0f},
+                                           {200.0f},
+                                           {201.0f},
+                                           {202.0f},
+                                           {203.0f}}));
+}
+
+TEST(DpUtilsTest, GatherGlobalTokensDropsZeroTokenDpRank) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  std::unique_ptr<xllm::ProcessGroup> tp_group;
+  ParallelArgs args =
+      make_world_args(process_group, tp_group, /*global_rank=*/6, 3, 4);
+  const std::vector<int32_t> dp_tokens = {0, 1, 33};
+
+  auto* mock_pg = dynamic_cast<test::MockProcessGroup*>(process_group.get());
+  ASSERT_NE(mock_pg, nullptr);
+  auto shards = make_global_shards(dp_tokens, /*tp_size=*/4);
+  mock_pg->set_allgather_outputs(shards);
+
+  auto output = gather_global_tokens(shards[6], dp_tokens, args);
+
+  EXPECT_EQ(output.size(0), 34);
+  test::verify_tensor_close(
+      output.slice(0, 0, 5),
+      torch::tensor({{100.0f}, {200.0f}, {201.0f}, {202.0f}, {203.0f}}));
+  test::verify_tensor_close(
+      output.slice(0, 30, 34),
+      torch::tensor({{229.0f}, {230.0f}, {231.0f}, {232.0f}}));
+}
+
+TEST(DpUtilsTest, GatherGlobalTokensHandlesEqualDpTokens) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  std::unique_ptr<xllm::ProcessGroup> tp_group;
+  ParallelArgs args =
+      make_world_args(process_group, tp_group, /*global_rank=*/4, 3, 2);
+  const std::vector<int32_t> dp_tokens = {7, 7, 7};
+
+  auto* mock_pg = dynamic_cast<test::MockProcessGroup*>(process_group.get());
+  ASSERT_NE(mock_pg, nullptr);
+  auto shards = make_global_shards(dp_tokens, /*tp_size=*/2);
+  mock_pg->set_allgather_outputs(shards);
+
+  auto output = gather_global_tokens(shards[4], dp_tokens, args);
+
+  EXPECT_EQ(output.size(0), 21);
+  test::verify_tensor_close(output.slice(0, 0, 3),
+                            torch::tensor({{0.0f}, {1.0f}, {2.0f}}));
+  test::verify_tensor_close(output.slice(0, 7, 10),
+                            torch::tensor({{100.0f}, {101.0f}, {102.0f}}));
+  test::verify_tensor_close(output.slice(0, 18, 21),
+                            torch::tensor({{204.0f}, {205.0f}, {206.0f}}));
+}
+
+TEST(DpUtilsTest, NeedDpMoEGatherSkipsAll2AllPath) {
+  std::unique_ptr<xllm::ProcessGroup> process_group;
+  ParallelArgs args = make_tp_args(process_group, 0, 1);
+  args.dp_size_ = 2;
+  args.ep_size_ = 2;
+
+  EXPECT_TRUE(need_dp_moe_gather(args, /*enable_moe_all2all=*/false));
+  EXPECT_FALSE(need_dp_moe_gather(args, /*enable_moe_all2all=*/true));
+}
+
+TEST(DpUtilsTest, UnpadTokensRestoresOriginalLength) {
+  auto x = torch::arange(12, torch::kFloat32).reshape({4, 3});
+  PaddingInfo info;
+  info.original_tokens = 2;
+  info.padded_tokens = 4;
+  info.active = true;
+
+  auto output = unpad_tokens(x, info);
+
+  test::verify_tensor_close(output, x.slice(0, 0, 2));
+}
+
+TEST(DpUtilsTest, AllDpRanksAreDecodeNeedsEveryRankDecode) {
+  ModelInputParams decode_params;
+  decode_params.dp_is_decode = {1, 1, 1};
+  EXPECT_TRUE(all_dp_ranks_are_decode(decode_params));
+
+  ModelInputParams mixed_params;
+  mixed_params.dp_is_decode = {1, 0, 1};
+  EXPECT_FALSE(all_dp_ranks_are_decode(mixed_params));
+}
+
+}  // namespace test
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/fused_moe_all2all_test.cpp

@@ -0,0 +1,696 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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.
+==============================================================================*/
+
+// FusedMoE All2All path unit tests
+// Tests for the DeepEP communication mode with multi-device setup
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <sys/wait.h>
+#include <torch/torch.h>
+#include <unistd.h>
+
+#include <cmath>
+#include <cstring>
+#include <functional>
+#include <memory>
+#include <vector>
+
+#include "common/global_flags.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/fused_moe.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+#include "util/tensor_helper.h"
+
+#if defined(USE_MLU)
+#include "framework/parallel_state/mlu_process_group.h"
+#elif defined(USE_CUDA)
+#include "framework/parallel_state/cuda_process_group.h"
+#endif
+
+namespace xllm {
+namespace layer {
+namespace test {
+
+// Special exit code definition for skipping test
+constexpr int32_t EXIT_CODE_SKIP = 77;
+
+// Helper function to create ProcessGroup
+std::unique_ptr<xllm::ProcessGroup> create_test_process_group(
+    int32_t rank,
+    int32_t world_size,
+    int32_t port,
+    const std::string& host,
+    const torch::Device& device) {
+  return xllm::create_process_group(rank,
+                                    world_size,
+                                    world_size,
+                                    port,
+                                    false,
+                                    host,
+                                    "fused_moe_all2all_test_group",
+                                    device);
+}
+
+struct All2AllTestParams {
+  int32_t rank;
+  int32_t world_size;
+  int32_t port;
+  std::string host;
+  int32_t device_index;
+  // Model parameters
+  int64_t hidden_size;
+  int64_t intermediate_size;
+  int64_t num_experts;
+  int64_t top_k;
+  int64_t batch_size;
+  int64_t seq_len;
+  bool is_smoothquant;
+  int64_t moe_weight_bits = 8;
+  int64_t group_size = 0;
+  // Expected output stats (used when perform_precise_validation is true)
+  double expected_min = 0.0;
+  double expected_max = 0.0;
+  double expected_sum = 0.0;
+  // When true, validate output against expected_min/max/sum (avoids fragile
+  // "any non-zero expected" heuristic; allows validation when all are 0.0).
+  bool perform_precise_validation = false;
+};
+
+// Helper to create model args
+ModelArgs create_model_args(const All2AllTestParams& params) {
+  ModelArgs args;
+  args.n_routed_experts() = static_cast<int32_t>(params.num_experts);
+  args.num_experts_per_tok() = static_cast<int32_t>(params.top_k);
+  args.n_group() = 1;
+  args.topk_group() = static_cast<int32_t>(params.top_k);
+  args.routed_scaling_factor() = 1.0f;
+  args.hidden_size() = params.hidden_size;
+  args.moe_intermediate_size() = static_cast<int32_t>(params.intermediate_size);
+  args.n_shared_experts() = 0;
+  args.norm_topk_prob() = true;
+  args.hidden_act() = "silu";
+  args.scoring_func() = "softmax";
+  args.topk_method() = "greedy";
+  return args;
+}
+
+// Helper to create quant args
+QuantArgs create_quant_args(const All2AllTestParams& params) {
+  QuantArgs args;
+  if (params.is_smoothquant) {
+    args.quant_method() = "smoothquant";
+    args.bits() = 8;
+    args.activation_dynamic() = true;
+    args.moe_weight_bits() = params.moe_weight_bits;
+    args.group_size() = params.group_size;
+  }
+  return args;
+}
+
+// Helper to create test weights for All2All MoE using seeded tensors
+std::unordered_map<std::string, torch::Tensor> create_all2all_test_weights(
+    int64_t num_experts,
+    int64_t hidden_size,
+    int64_t intermediate_size,
+    bool is_smoothquant,
+    int64_t moe_weight_bits,
+    int64_t group_size,
+    int64_t world_size,
+    const torch::Device& device) {
+  std::unordered_map<std::string, torch::Tensor> weight_dict;
+
+  for (size_t expert_id = 0; expert_id < num_experts; ++expert_id) {
+    std::string expert_prefix = "experts." + std::to_string(expert_id) + ".";
+    std::string seed_prefix =
+        "fused_moe_all2all_tests.expert_" + std::to_string(expert_id);
+
+    if (is_smoothquant && moe_weight_bits == 4) {
+      CHECK_GT(group_size, 0);
+      CHECK_EQ(intermediate_size % world_size, 0);
+      const int64_t local_intermediate_size = intermediate_size / world_size;
+      test::append_w4a8_expert_weights(weight_dict,
+                                       expert_prefix,
+                                       seed_prefix,
+                                       hidden_size,
+                                       intermediate_size,
+                                       intermediate_size,
+                                       local_intermediate_size,
+                                       group_size,
+                                       device);
+    } else if (is_smoothquant) {
+      // Create quantized weights using seeded tensors
+      auto gate_weight_fp =
+          test::seeded_tensor(seed_prefix + ".gate_proj",
+                              {intermediate_size, hidden_size},
+                              torch::kBFloat16,
+                              device);
+      auto gate_qweight = gate_weight_fp.to(torch::kInt8);
+      auto gate_scale = test::seeded_tensor(seed_prefix + ".gate_proj.scale",
+                                            {intermediate_size},
+                                            torch::kFloat32,
+                                            device);
+      auto gate_smooth = test::seeded_tensor(seed_prefix + ".gate_proj.smooth",
+                                             {hidden_size},
+                                             torch::kFloat32,
+                                             device);
+
+      auto up_weight_fp = test::seeded_tensor(seed_prefix + ".up_proj",
+                                              {intermediate_size, hidden_size},
+                                              torch::kBFloat16,
+                                              device);
+      auto up_qweight = up_weight_fp.to(torch::kInt8);
+      auto up_scale = test::seeded_tensor(seed_prefix + ".up_proj.scale",
+                                          {intermediate_size},
+                                          torch::kFloat32,
+                                          device);
+
+      auto down_weight_fp =
+          test::seeded_tensor(seed_prefix + ".down_proj",
+                              {hidden_size, intermediate_size},
+                              torch::kBFloat16,
+                              device);
+      auto down_qweight = down_weight_fp.to(torch::kInt8);
+      auto down_scale = test::seeded_tensor(seed_prefix + ".down_proj.scale",
+                                            {hidden_size},
+                                            torch::kFloat32,
+                                            device);
+      auto down_smooth = test::seeded_tensor(seed_prefix + ".down_proj.smooth",
+                                             {intermediate_size},
+                                             torch::kFloat32,
+                                             device);
+
+      weight_dict[expert_prefix + "gate_proj.qweight"] = gate_qweight;
+      weight_dict[expert_prefix + "gate_proj.per_channel_scale"] = gate_scale;
+      weight_dict[expert_prefix + "gate_proj.smooth"] = gate_smooth;
+
+      weight_dict[expert_prefix + "up_proj.qweight"] = up_qweight;
+      weight_dict[expert_prefix + "up_proj.per_channel_scale"] = up_scale;
+      weight_dict[expert_prefix + "up_proj.smooth"] = gate_smooth;
+
+      weight_dict[expert_prefix + "down_proj.qweight"] = down_qweight;
+      weight_dict[expert_prefix + "down_proj.per_channel_scale"] = down_scale;
+      weight_dict[expert_prefix + "down_proj.smooth"] = down_smooth;
+    } else {
+      // Create BF16 weights using seeded tensors
+      auto gate_weight = test::seeded_tensor(seed_prefix + ".gate_proj.weight",
+                                             {intermediate_size, hidden_size},
+                                             torch::kBFloat16,
+                                             device);
+      auto up_weight = test::seeded_tensor(seed_prefix + ".up_proj.weight",
+                                           {intermediate_size, hidden_size},
+                                           torch::kBFloat16,
+                                           device);
+      auto down_weight = test::seeded_tensor(seed_prefix + ".down_proj.weight",
+                                             {hidden_size, intermediate_size},
+                                             torch::kBFloat16,
+                                             device);
+
+      weight_dict[expert_prefix + "gate_proj.weight"] = gate_weight;
+      weight_dict[expert_prefix + "up_proj.weight"] = up_weight;
+      weight_dict[expert_prefix + "down_proj.weight"] = down_weight;
+    }
+  }
+
+  // Gate weights (router)
+  auto gate_weight = test::seeded_tensor("fused_moe_all2all_tests.gate.weight",
+                                         {num_experts, hidden_size},
+                                         torch::kBFloat16,
+                                         device);
+  weight_dict["gate.weight"] = gate_weight;
+
+  return weight_dict;
+}
+
+// Child process test function for basic All2All test
+int32_t run_all2all_basic_test_child(All2AllTestParams params) {
+  try {
+    // 0. Set FLAGS_expert_parallel_degree to enable DeepEP
+    // This is required for All2All path to work
+    FLAGS_expert_parallel_degree = 2;
+
+    // 1. Check devices
+    int32_t dev_count = xllm::Device::device_count();
+    if (dev_count < params.world_size) {
+      LOG(WARNING) << "Rank " << params.rank
+                   << ": Insufficient devices. Skipping.";
+      return EXIT_CODE_SKIP;
+    }
+    params.device_index = params.rank % dev_count;
+
+    // 2. Set device
+    xllm::Device xllm_device(params.device_index);
+    xllm_device.set_device();
+    torch::Device device = xllm_device.unwrap();
+
+    // 3. Create ProcessGroup
+    auto process_group = create_test_process_group(
+        params.rank, params.world_size, params.port, params.host, device);
+
+    CHECK(process_group) << "Rank " << params.rank
+                         << ": Failed to create ProcessGroup";
+
+    // 4. Create ParallelArgs with EP mode
+    ParallelArgs parallel_args(
+        params.rank, params.world_size, process_group.get());
+    parallel_args.moe_ep_group_ = process_group.get();
+    parallel_args.ep_size_ = params.world_size;
+    parallel_args.moe_tp_group_ = process_group.get();
+
+    // 5. Create tensor options
+    // Note: smoothquant mode still uses BF16 for input/output, with int8 for
+    // internal quantized computation. Using float32 would cause CNNL GroupGemm
+    // to fail with "Data type mismatch" error.
+    auto options = torch::TensorOptions()
+                       .dtype(torch::kBFloat16)
+                       .device(device)
+                       .requires_grad(false);
+
+    // 6. Create model args and quant args
+    ModelArgs model_args = create_model_args(params);
+    QuantArgs quant_args = create_quant_args(params);
+
+    // 7. Create FusedMoE
+    FusedMoE fused_moe(FusedMoEImpl(model_args,
+                                    FusedMoEArgs{.is_gated = true},
+                                    quant_args,
+                                    parallel_args,
+                                    options));
+
+    // 8. Create and load weights using seeded tensors
+    auto weight_dict = create_all2all_test_weights(params.num_experts,
+                                                   params.hidden_size,
+                                                   params.intermediate_size,
+                                                   params.is_smoothquant,
+                                                   params.moe_weight_bits,
+                                                   params.group_size,
+                                                   params.world_size,
+                                                   device);
+    StateDict state_dict(weight_dict);
+    fused_moe->load_state_dict(state_dict);
+
+    LOG(INFO) << "Rank " << params.rank << ": FusedMoE created and loaded";
+
+    // 9. Create input tensor using seeded tensor for determinism
+    int64_t num_tokens = params.batch_size * params.seq_len;
+    std::string seed_prefix = params.is_smoothquant
+                                  ? "fused_moe_all2all_tests.smoothquant"
+                                  : "fused_moe_all2all_tests.basic";
+    auto hidden_states = test::seeded_tensor(seed_prefix + ".hidden_states",
+                                             {num_tokens, params.hidden_size},
+                                             torch::kBFloat16,
+                                             device);
+
+    // 10. Run forward with All2All enabled
+    auto output =
+        fused_moe->forward_experts(hidden_states,
+                                   /*enable_all2all_communication=*/true);
+
+    // 11. Verify output
+    xllm_device.synchronize_default_stream();
+
+    CHECK_EQ(output.sizes().size(), 2) << "Output should be 2D tensor";
+    CHECK_EQ(output.size(0), num_tokens) << "Token count should match";
+    CHECK_EQ(output.size(1), params.hidden_size) << "Hidden size should match";
+
+    // Compute and log output stats for setting expected values
+    auto flat_output = output.flatten().to(torch::kFloat32).cpu();
+    double actual_min = torch::min(flat_output).item<double>();
+    double actual_max = torch::max(flat_output).item<double>();
+    double actual_sum = torch::sum(flat_output).item<double>();
+
+    LOG(INFO) << "Rank " << params.rank << ": Output stats - "
+              << "min=" << actual_min << ", max=" << actual_max
+              << ", sum=" << actual_sum;
+
+    // Basic sanity check (will be replaced with precise validation later)
+    CHECK_NE(actual_sum, 0.0) << "Output should not be all zeros";
+
+    if (params.perform_precise_validation) {
+      test::expect_tensor_stats(output,
+                                params.expected_min,
+                                params.expected_max,
+                                params.expected_sum);
+      LOG(INFO) << "Rank " << params.rank
+                << ": All2All test passed with precise validation.";
+    } else {
+      LOG(INFO) << "Rank " << params.rank
+                << ": All2All test passed (basic validation only).";
+    }
+
+    return 0;
+
+  } catch (const std::exception& e) {
+    LOG(ERROR) << "Rank " << params.rank << ": Exception: " << e.what();
+    return 1;
+  }
+}
+
+// Child process test function for SmoothQuant All2All test
+int32_t run_all2all_smoothquant_test_child(All2AllTestParams params) {
+  params.is_smoothquant = true;
+  return run_all2all_basic_test_child(params);
+}
+
+int32_t run_all2all_route_guard_test_child(All2AllTestParams params) {
+  try {
+    FLAGS_expert_parallel_degree = 2;
+
+    int32_t dev_count = xllm::Device::device_count();
+    if (dev_count < params.world_size) {
+      LOG(WARNING) << "Rank " << params.rank
+                   << ": Insufficient devices. Skipping.";
+      return EXIT_CODE_SKIP;
+    }
+    params.device_index = params.rank % dev_count;
+
+    xllm::Device xllm_device(params.device_index);
+    xllm_device.set_device();
+    torch::Device device = xllm_device.unwrap();
+
+    auto process_group = create_test_process_group(
+        params.rank, params.world_size, params.port, params.host, device);
+    CHECK(process_group) << "Rank " << params.rank
+                         << ": Failed to create ProcessGroup";
+
+    ParallelArgs parallel_args(
+        params.rank, params.world_size, process_group.get());
+    parallel_args.moe_ep_group_ = process_group.get();
+    parallel_args.ep_size_ = params.world_size;
+    parallel_args.moe_tp_group_ = process_group.get();
+
+    auto options = torch::TensorOptions()
+                       .dtype(torch::kBFloat16)
+                       .device(device)
+                       .requires_grad(false);
+
+    ModelArgs model_args = create_model_args(params);
+    QuantArgs quant_args = create_quant_args(params);
+    FusedMoE fused_moe(FusedMoEImpl(model_args,
+                                    FusedMoEArgs{.is_gated = true},
+                                    quant_args,
+                                    parallel_args,
+                                    options));
+
+    auto weight_dict = create_all2all_test_weights(params.num_experts,
+                                                   params.hidden_size,
+                                                   params.intermediate_size,
+                                                   params.is_smoothquant,
+                                                   params.moe_weight_bits,
+                                                   params.group_size,
+                                                   params.world_size,
+                                                   device);
+    StateDict state_dict(weight_dict);
+    fused_moe->load_state_dict(state_dict);
+
+    int64_t num_tokens = params.batch_size * params.seq_len;
+    auto hidden_states =
+        test::seeded_tensor("fused_moe_all2all_tests.guard.hidden_states",
+                            {num_tokens, params.hidden_size},
+                            torch::kBFloat16,
+                            device);
+    auto route_info = fused_moe->prep_route(hidden_states);
+    fused_moe->forward_experts(hidden_states,
+                               /*enable_all2all_communication=*/true,
+                               route_info);
+
+    xllm_device.synchronize_default_stream();
+    return 0;
+  } catch (const std::exception& e) {
+    LOG(ERROR) << "Rank " << params.rank << ": Exception: " << e.what();
+    return 1;
+  }
+}
+
+// Multi-process test fixture
+class FusedMoEAll2AllMultiDeviceTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    world_size_ = 2;
+    port_ = 29501;  // Different port from deep_ep_tests
+    host_ = "127.0.0.1";
+    hidden_size_ = 512;
+    intermediate_size_ = 256;
+    num_experts_ = 4;
+    top_k_ = 2;
+    batch_size_ = 2;
+    seq_len_ = 4;
+  }
+
+  // Runs multi-process test with params built from fixture (optionally
+  // is_smoothquant). Use this for basic flow tests.
+  void run_test(int32_t (*test_fn)(All2AllTestParams),
+                bool is_smoothquant = false) {
+    run_test_impl(test_fn, [this, is_smoothquant](int32_t rank) {
+      All2AllTestParams params;
+      params.rank = rank;
+      params.world_size = world_size_;
+      params.port = port_;
+      params.host = host_;
+      params.device_index = -1;
+      params.hidden_size = hidden_size_;
+      params.intermediate_size = intermediate_size_;
+      params.num_experts = num_experts_;
+      params.top_k = top_k_;
+      params.batch_size = batch_size_;
+      params.seq_len = seq_len_;
+      params.is_smoothquant = is_smoothquant;
+      return params;
+    });
+  }
+
+  // Runs multi-process test with custom params from factory (e.g. for precise
+  // validation tests). params_factory(rank) is called per rank; the returned
+  // params should have rank set by the factory or will be overwritten.
+  // Named separately to avoid overload ambiguity with run_test(..., bool).
+  void run_test_with_params(
+      int32_t (*test_fn)(All2AllTestParams),
+      std::function<All2AllTestParams(int32_t rank)> params_factory) {
+    run_test_impl(test_fn, params_factory);
+  }
+
+  void run_fail_test(
+      int32_t (*test_fn)(All2AllTestParams),
+      std::function<All2AllTestParams(int32_t rank)> params_factory) {
+    std::vector<pid_t> child_pids;
+
+    for (int32_t rank = 0; rank < world_size_; ++rank) {
+      pid_t pid = fork();
+      if (pid == 0) {
+        All2AllTestParams params = params_factory(rank);
+        params.rank = rank;
+        params.world_size = world_size_;
+        int32_t exit_code = test_fn(params);
+        _exit(exit_code);
+      } else if (pid > 0) {
+        child_pids.push_back(pid);
+      } else {
+        LOG(FATAL) << "Failed to fork rank " << rank;
+      }
+    }
+
+    bool any_skipped = false;
+    bool any_succeeded = false;
+
+    for (size_t i = 0; i < child_pids.size(); ++i) {
+      int32_t status;
+      waitpid(child_pids[i], &status, 0);
+      if (WIFEXITED(status)) {
+        int32_t exit_code = WEXITSTATUS(status);
+        if (exit_code == EXIT_CODE_SKIP) {
+          any_skipped = true;
+        } else if (exit_code == 0) {
+          any_succeeded = true;
+          LOG(ERROR) << "Rank " << i << " unexpectedly succeeded.";
+        }
+      }
+    }
+
+    if (any_skipped) {
+      GTEST_SKIP() << "Test skipped due to insufficient devices.";
+    } else {
+      ASSERT_FALSE(any_succeeded) << "All2All route guard did not trigger.";
+    }
+  }
+
+ private:
+  void run_test_impl(
+      int32_t (*test_fn)(All2AllTestParams),
+      std::function<All2AllTestParams(int32_t rank)> params_factory) {
+    std::vector<pid_t> child_pids;
+
+    for (int32_t rank = 0; rank < world_size_; ++rank) {
+      pid_t pid = fork();
+      if (pid == 0) {
+        All2AllTestParams params = params_factory(rank);
+        params.rank = rank;
+        params.world_size = world_size_;
+        int32_t exit_code = test_fn(params);
+        _exit(exit_code);
+      } else if (pid > 0) {
+        child_pids.push_back(pid);
+      } else {
+        LOG(FATAL) << "Failed to fork rank " << rank;
+      }
+    }
+
+    bool any_failed = false;
+    bool any_skipped = false;
+
+    for (size_t i = 0; i < child_pids.size(); ++i) {
+      int32_t status;
+      waitpid(child_pids[i], &status, 0);
+      if (WIFEXITED(status)) {
+        int32_t exit_code = WEXITSTATUS(status);
+        if (exit_code == EXIT_CODE_SKIP) {
+          any_skipped = true;
+        } else if (exit_code != 0) {
+          any_failed = true;
+          LOG(ERROR) << "Rank " << i << " failed with code " << exit_code;
+        }
+      } else {
+        any_failed = true;
+        LOG(ERROR) << "Rank " << i << " crashed (signal).";
+      }
+    }
+
+    if (any_skipped) {
+      GTEST_SKIP() << "Test skipped due to insufficient devices.";
+    } else {
+      ASSERT_FALSE(any_failed) << "FusedMoE All2All Test Failed.";
+    }
+  }
+
+ protected:
+  int32_t world_size_;
+  int32_t port_;
+  std::string host_;
+  int64_t hidden_size_;
+  int64_t intermediate_size_;
+  int64_t num_experts_;
+  int64_t top_k_;
+  int64_t batch_size_;
+  int64_t seq_len_;
+};
+
+TEST_F(FusedMoEAll2AllMultiDeviceTest, BasicAll2AllFlow) {
+  run_test(run_all2all_basic_test_child, /*is_smoothquant=*/false);
+}
+
+TEST_F(FusedMoEAll2AllMultiDeviceTest, SmoothQuantAll2AllFlow) {
+  run_test(run_all2all_smoothquant_test_child, /*is_smoothquant=*/true);
+}
+
+TEST_F(FusedMoEAll2AllMultiDeviceTest, BasicAll2AllPreciseValidation) {
+  // Expected values obtained from running with seeded tensors:
+  // min=835584, max=1.26976e+06, sum=4.20999e+09
+  run_test_with_params(run_all2all_basic_test_child, [](int32_t /*rank*/) {
+    All2AllTestParams params;
+    params.rank = 0;
+    params.world_size = 2;
+    params.port = 29503;  // Different port to avoid conflicts
+    params.host = "127.0.0.1";
+    params.device_index = -1;
+    params.hidden_size = 512;
+    params.intermediate_size = 256;
+    params.num_experts = 4;
+    params.top_k = 2;
+    params.batch_size = 2;
+    params.seq_len = 4;
+    params.is_smoothquant = false;
+    params.expected_min = 835584.0;
+    params.expected_max = 1269760.0;
+    params.expected_sum = 4209990000.0;
+    params.perform_precise_validation = true;
+    return params;
+  });
+}
+
+TEST_F(FusedMoEAll2AllMultiDeviceTest, SmoothQuantAll2AllPreciseValidation) {
+  // Expected values obtained from running with seeded tensors:
+  // min=0, max=0.104004, sum=3.88289
+  run_test_with_params(
+      run_all2all_smoothquant_test_child, [](int32_t /*rank*/) {
+        All2AllTestParams params;
+        params.rank = 0;
+        params.world_size = 2;
+        params.port = 29504;  // Different port to avoid conflicts
+        params.host = "127.0.0.1";
+        params.device_index = -1;
+        params.hidden_size = 512;
+        params.intermediate_size = 256;
+        params.num_experts = 4;
+        params.top_k = 2;
+        params.batch_size = 2;
+        params.seq_len = 4;
+        params.is_smoothquant = true;
+        // Note: min=0 is valid for SmoothQuant due to quantization effects
+        params.expected_min = 0.0;
+        params.expected_max = 0.104004;
+        params.expected_sum = 3.88289;
+        params.perform_precise_validation = true;
+        return params;
+      });
+}
+
+TEST_F(FusedMoEAll2AllMultiDeviceTest, W4A8All2AllSmoke) {
+  run_test_with_params(run_all2all_smoothquant_test_child,
+                       [](int32_t /*rank*/) {
+                         All2AllTestParams params;
+                         params.rank = 0;
+                         params.world_size = 2;
+                         params.port = 29505;
+                         params.host = "127.0.0.1";
+                         params.device_index = -1;
+                         params.hidden_size = 512;
+                         params.intermediate_size = 256;
+                         params.num_experts = 4;
+                         params.top_k = 2;
+                         params.batch_size = 2;
+                         params.seq_len = 4;
+                         params.is_smoothquant = true;
+                         params.moe_weight_bits = 4;
+                         params.group_size = 128;
+                         params.perform_precise_validation = false;
+                         return params;
+                       });
+}
+
+TEST_F(FusedMoEAll2AllMultiDeviceTest, ExternalRouteGuard) {
+  run_fail_test(run_all2all_route_guard_test_child, [](int32_t /*rank*/) {
+    All2AllTestParams params;
+    params.rank = 0;
+    params.world_size = 2;
+    params.port = 29506;
+    params.host = "127.0.0.1";
+    params.device_index = -1;
+    params.hidden_size = 512;
+    params.intermediate_size = 256;
+    params.num_experts = 4;
+    params.top_k = 2;
+    params.batch_size = 2;
+    params.seq_len = 4;
+    params.is_smoothquant = false;
+    return params;
+  });
+}
+
+}  // namespace test
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/fused_moe_test.cpp

@@ -0,0 +1,632 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/mlu/fused_moe.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+class FusedMoETest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Initialize default model arguments for testing
+    model_args_ = test::create_default_model_args();
+
+    // Initialize w8a8 quantization arguments
+    quant_args_ = test::create_default_quant_args();
+
+    // Initialize tensor options
+    options_ = torch::TensorOptions()
+                   .dtype(torch::kBFloat16)
+                   .device(Device::type_torch(), 0)
+                   .requires_grad(false);
+
+    // Create mock ProcessGroup and initialize ParallelArgs
+    parallel_args_ = test::create_default_parallel_args(mock_process_group_);
+
+    // Note: FusedMoE will be created by individual test cases with their
+    // desired dimensions
+  }
+
+  void TearDown() override {
+    // Clean up if needed
+  }
+
+  // Helper function to create router logits tensor
+  torch::Tensor create_router_logits(const std::vector<int64_t>& shape,
+                                     const std::vector<float>& values) {
+    return test::create_custom_input(shape, values, options_);
+  }
+
+  std::unordered_map<std::string, torch::Tensor> create_default_test_weights(
+      int64_t num_experts,
+      int64_t hidden_size,
+      int64_t intermediate_size) {
+    // Create test weights for each expert
+    std::unordered_map<std::string, torch::Tensor> weight_dict;
+
+    for (size_t expert_id = 0; expert_id < num_experts; ++expert_id) {
+      std::string expert_prefix = "experts." + std::to_string(expert_id) + ".";
+
+      // Create gate_proj weights (ColumnParallelLinear)
+      // Shape: [intermediate_size, hidden_size]
+      auto gate_weight =
+          torch::full({intermediate_size, hidden_size}, 2.0f, options_);
+      auto gate_qweight = gate_weight.to(torch::kInt8);
+      auto gate_scale = torch::full({intermediate_size},
+                                    0.1f,
+                                    torch::TensorOptions()
+                                        .dtype(torch::kFloat32)
+                                        .device(options_.device()));
+      auto gate_smooth = torch::full({hidden_size},
+                                     0.05f,
+                                     torch::TensorOptions()
+                                         .dtype(torch::kFloat32)
+                                         .device(options_.device()));
+
+      // Create up_proj weights (ColumnParallelLinear)
+      // Shape: [intermediate_size, hidden_size]
+      auto up_weight =
+          torch::full({intermediate_size, hidden_size}, 2.0f, options_);
+      auto up_qweight = up_weight.to(torch::kInt8);
+      auto up_scale = torch::full({intermediate_size},
+                                  0.1f,
+                                  torch::TensorOptions()
+                                      .dtype(torch::kFloat32)
+                                      .device(options_.device()));
+      auto up_smooth = torch::full({hidden_size},
+                                   0.05f,
+                                   torch::TensorOptions()
+                                       .dtype(torch::kFloat32)
+                                       .device(options_.device()));
+
+      // Create down_proj weights (RowParallelLinear)
+      // Shape: [hidden_size, intermediate_size]
+      auto down_weight =
+          torch::full({hidden_size, intermediate_size}, 3.0f, options_);
+      auto down_qweight = down_weight.to(torch::kInt8);
+      auto down_scale = torch::full({hidden_size},
+                                    0.1f,
+                                    torch::TensorOptions()
+                                        .dtype(torch::kFloat32)
+                                        .device(options_.device()));
+      auto down_smooth = torch::full({intermediate_size},
+                                     0.05f,
+                                     torch::TensorOptions()
+                                         .dtype(torch::kFloat32)
+                                         .device(options_.device()));
+
+      // Add weights to dictionary
+      // expert
+      weight_dict[expert_prefix + "gate_proj.qweight"] = gate_qweight;
+      weight_dict[expert_prefix + "gate_proj.per_channel_scale"] = gate_scale;
+      weight_dict[expert_prefix + "gate_proj.smooth"] = gate_smooth;
+
+      weight_dict[expert_prefix + "up_proj.qweight"] = up_qweight;
+      weight_dict[expert_prefix + "up_proj.per_channel_scale"] = up_scale;
+      weight_dict[expert_prefix + "up_proj.smooth"] = up_smooth;
+
+      weight_dict[expert_prefix + "down_proj.qweight"] = down_qweight;
+      weight_dict[expert_prefix + "down_proj.per_channel_scale"] = down_scale;
+      weight_dict[expert_prefix + "down_proj.smooth"] = down_smooth;
+    }
+
+    // gate weight generation
+    auto gate_weight = torch::full({num_experts, hidden_size}, 5.0f, options_);
+    auto e_score_correction_bias = torch::full({num_experts}, 0.1f, options_);
+
+    // Create shared experts weights
+    auto shared_expert_up_weight =
+        torch::full({intermediate_size, hidden_size}, 1.5f, options_);
+    auto shared_expert_up_qweight = shared_expert_up_weight.to(torch::kInt8);
+    auto shared_expert_up_scale = torch::full({intermediate_size},
+                                              0.1f,
+                                              torch::TensorOptions()
+                                                  .dtype(torch::kFloat32)
+                                                  .device(options_.device()));
+    auto shared_expert_up_smooth = torch::full({hidden_size},
+                                               0.05f,
+                                               torch::TensorOptions()
+                                                   .dtype(torch::kFloat32)
+                                                   .device(options_.device()));
+    auto shared_expert_gate_weight =
+        torch::full({intermediate_size, hidden_size}, 1.5f, options_);
+    auto shared_expert_gate_qweight =
+        shared_expert_gate_weight.to(torch::kInt8);
+    auto shared_expert_gate_scale = torch::full({intermediate_size},
+                                                0.1f,
+                                                torch::TensorOptions()
+                                                    .dtype(torch::kFloat32)
+                                                    .device(options_.device()));
+    auto shared_expert_gate_smooth =
+        torch::full({hidden_size},
+                    0.05f,
+                    torch::TensorOptions()
+                        .dtype(torch::kFloat32)
+                        .device(options_.device()));
+    auto shared_expert_down_weight =
+        torch::full({hidden_size, intermediate_size}, 1.3f, options_);
+    auto shared_expert_down_qweight =
+        shared_expert_down_weight.to(torch::kInt8);
+    auto shared_expert_down_scale = torch::full({hidden_size},
+                                                0.1f,
+                                                torch::TensorOptions()
+                                                    .dtype(torch::kFloat32)
+                                                    .device(options_.device()));
+    auto shared_expert_down_smooth =
+        torch::full({intermediate_size},
+                    0.05f,
+                    torch::TensorOptions()
+                        .dtype(torch::kFloat32)
+                        .device(options_.device()));
+
+    // gate
+    weight_dict["gate.weight"] = gate_weight;
+    weight_dict["gate.e_score_correction_bias"] = e_score_correction_bias;
+
+    // shared experts
+    weight_dict["shared_experts.up_proj.qweight"] = shared_expert_up_qweight;
+    weight_dict["shared_experts.up_proj.per_channel_scale"] =
+        shared_expert_up_scale;
+    weight_dict["shared_experts.up_proj.smooth"] = shared_expert_up_smooth;
+    weight_dict["shared_experts.gate_proj.qweight"] =
+        shared_expert_gate_qweight;
+    weight_dict["shared_experts.gate_proj.per_channel_scale"] =
+        shared_expert_gate_scale;
+    weight_dict["shared_experts.gate_proj.smooth"] = shared_expert_gate_smooth;
+    weight_dict["shared_experts.down_proj.qweight"] =
+        shared_expert_down_qweight;
+    weight_dict["shared_experts.down_proj.per_channel_scale"] =
+        shared_expert_down_scale;
+    weight_dict["shared_experts.down_proj.smooth"] = shared_expert_down_smooth;
+
+    LOG(INFO) << "Test w8a8 smoothquant weights created successfully for "
+              << num_experts << " experts";
+    LOG(INFO) << "Hidden size: " << hidden_size
+              << ", Intermediate size: " << intermediate_size;
+
+    return weight_dict;
+  }
+
+  std::unordered_map<std::string, torch::Tensor>
+  create_w4a8_groupwise_test_weights(int64_t num_experts,
+                                     int64_t hidden_size,
+                                     int64_t intermediate_size,
+                                     int64_t group_size) {
+    std::unordered_map<std::string, torch::Tensor> weight_dict;
+
+    for (size_t expert_id = 0; expert_id < num_experts; ++expert_id) {
+      std::string expert_prefix = "experts." + std::to_string(expert_id) + ".";
+      std::string seed_prefix =
+          "fused_moe_tests.expert_" + std::to_string(expert_id);
+      test::append_w4a8_expert_weights(weight_dict,
+                                       expert_prefix,
+                                       seed_prefix,
+                                       hidden_size,
+                                       intermediate_size,
+                                       intermediate_size,
+                                       intermediate_size,
+                                       group_size,
+                                       options_.device());
+    }
+
+    weight_dict["gate.weight"] =
+        torch::full({num_experts, hidden_size}, 5.0f, options_);
+    weight_dict["gate.e_score_correction_bias"] =
+        torch::full({num_experts}, 0.1f, options_);
+
+    return weight_dict;
+  }
+
+  FusedMoE create_fused_moe(int64_t num_experts,
+                            int64_t top_k,
+                            int64_t num_expert_group,
+                            int64_t topk_group,
+                            double route_scale,
+                            int64_t hidden_size,
+                            int64_t intermediate_size,
+                            int64_t n_shared_experts = 1,
+                            bool is_gated = true,
+                            int64_t renormalize = 0,
+                            bool enable_result_reduction = true,
+                            const std::string& hidden_act = "silu",
+                            const std::string& scoring_func = "sigmoid",
+                            const std::string& topk_method = "noaux_tc") {
+    ModelArgs args = model_args_;
+    args.n_routed_experts() = static_cast<int32_t>(num_experts);
+    args.num_experts_per_tok() = static_cast<int32_t>(top_k);
+    args.n_group() = static_cast<int32_t>(num_expert_group);
+    args.topk_group() = static_cast<int32_t>(topk_group);
+    args.routed_scaling_factor() = static_cast<float>(route_scale);
+    args.hidden_size() = hidden_size;
+    args.moe_intermediate_size() = static_cast<int32_t>(intermediate_size);
+    args.n_shared_experts() = static_cast<int32_t>(n_shared_experts);
+    args.norm_topk_prob() = (renormalize != 0);
+    args.hidden_act() = hidden_act;
+    args.scoring_func() = scoring_func;
+    args.topk_method() = topk_method;
+    const FusedMoEArgs moe_args{
+        .is_gated = is_gated,
+        .enable_result_reduction = enable_result_reduction};
+    return FusedMoE(
+        FusedMoEImpl(args, moe_args, quant_args_, parallel_args_, options_));
+  }
+
+  void set_tp_ctx(int64_t world_size) {
+    mock_process_group_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, world_size);
+    single_rank_pg_ = std::make_unique<test::MockProcessGroup>(
+        options_.device(), /*rank=*/0, /*world_size=*/1);
+    parallel_args_ = ParallelArgs(/*rank=*/0,
+                                  world_size,
+                                  /*dp_size=*/1,
+                                  mock_process_group_.get());
+    parallel_args_.process_group_ = mock_process_group_.get();
+    parallel_args_.tp_group_ = mock_process_group_.get();
+    parallel_args_.single_rank_group_ = single_rank_pg_.get();
+  }
+
+  // Helper function to create test weights for the FusedMoE (w8a8 smoothquant
+  // format)
+  std::unordered_map<std::string, torch::Tensor> create_test_weights(
+      int64_t num_experts,
+      int64_t custom_hidden_size = -1,
+      int64_t custom_intermediate_size = -1) {
+    // Use custom sizes if provided, otherwise use model_args_ values
+    int64_t test_hidden_size = (custom_hidden_size > 0)
+                                   ? custom_hidden_size
+                                   : model_args_.hidden_size();
+    int64_t test_intermediate_size = (custom_intermediate_size > 0)
+                                         ? custom_intermediate_size
+                                         : model_args_.intermediate_size();
+
+    return create_default_test_weights(
+        num_experts, test_hidden_size, test_intermediate_size);
+  }
+
+  // Helper function to verify tensor values are close to expected
+  void verify_tensor_close(const torch::Tensor& actual,
+                           const torch::Tensor& expected,
+                           double rtol = 1e-5,
+                           double atol = 1e-8) {
+    test::verify_tensor_close(actual, expected, rtol, atol);
+  }
+
+  // Helper function to create custom input tensor for precision testing
+  torch::Tensor create_custom_input(const std::vector<int64_t>& shape,
+                                    const std::vector<float>& values) {
+    return test::create_custom_input(shape, values, options_);
+  }
+
+  // Helper function to set expected output for precision verification
+  void set_expected_output(const std::vector<float>& expected_values) {
+    expected_output_ = expected_values;
+  }
+
+  // Helper function to verify precision against expected output
+  void verify_precision(const torch::Tensor& actual_output,
+                        double rtol = 1e-3,
+                        double atol = 1e-4) {
+    test::verify_precision(actual_output, expected_output_, rtol, atol);
+  }
+
+  ModelArgs model_args_;
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+
+  // Helper to create a mock ProcessGroup for testing
+  std::unique_ptr<xllm::ProcessGroup> mock_process_group_;
+  std::unique_ptr<xllm::ProcessGroup> single_rank_pg_;
+
+  // Expected output for precision verification
+  std::vector<float> expected_output_;
+};
+
+TEST_F(FusedMoETest, LoadStateDictTest) {
+  // Test loading weights into the FusedMoE
+  const int64_t batch_size = 16;
+  const int64_t seq_len = 32;
+  const int64_t hidden_size = 7168;
+  const int64_t intermediate_size = 2048;
+  const int64_t num_experts = 16;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+  const bool gated = true;
+  const int64_t renormalize = 1;
+  const int64_t n_shared_experts = 1;
+
+  // Create FusedMoE with default dimensions
+  auto fused_moe = create_fused_moe(num_experts,
+                                    top_k,
+                                    num_expert_group,
+                                    topk_group,
+                                    route_scale,
+                                    hidden_size,
+                                    intermediate_size,
+                                    n_shared_experts,
+                                    gated,
+                                    renormalize);
+
+  // Create test weights and load them
+  auto weight_dict =
+      create_test_weights(num_experts, hidden_size, intermediate_size);
+
+  // Load weights into the FusedMoE
+  StateDict state_dict(weight_dict);
+  fused_moe->load_state_dict(state_dict);
+
+  // Create input tensors
+  auto hidden_states = create_custom_input(
+      {batch_size * seq_len, hidden_size},
+      std::vector<float>(batch_size * seq_len * hidden_size, 0.05f));
+
+  auto output =
+      fused_moe->forward_experts(hidden_states,
+                                 /*enable_all2all_communication=*/false);
+
+  CHECK_EQ(output.sizes().size(), 2) << "Output should be 2D tensor";
+  CHECK_EQ(output.size(0), batch_size * seq_len)
+      << "The number of tokens should match";
+  CHECK_EQ(output.size(1), hidden_size) << "The hidden size should match";
+  auto output_sum = torch::sum(output).item<float>();
+  CHECK_NE(output_sum, 0.0f)
+      << "Output should not be all zeros after loading weights";
+
+  LOG(INFO) << "State dict loading test passed - output sum: " << output_sum;
+}
+
+TEST_F(FusedMoETest, PrecisionVerificationTest) {
+  // Test loading weights into the FusedMoE
+  const int64_t batch_size = 16;
+  const int64_t seq_len = 32;
+  const int64_t hidden_size = 7168;
+  const int64_t intermediate_size = 2048;
+  const int64_t num_experts = 16;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+  const bool gated = true;
+  const int64_t renormalize = 1;
+  const int64_t n_shared_experts = 1;
+
+  // Create FusedMoE with default dimensions
+  auto fused_moe = create_fused_moe(num_experts,
+                                    top_k,
+                                    num_expert_group,
+                                    topk_group,
+                                    route_scale,
+                                    hidden_size,
+                                    intermediate_size,
+                                    n_shared_experts,
+                                    gated,
+                                    renormalize);
+
+  // Create test weights and load them
+  auto weight_dict =
+      create_test_weights(num_experts, hidden_size, intermediate_size);
+
+  // Load weights into the FusedMoE
+  StateDict state_dict(weight_dict);
+  fused_moe->load_state_dict(state_dict);
+
+  // Create input tensors
+  auto hidden_states = create_custom_input(
+      {batch_size * seq_len, hidden_size},
+      std::vector<float>(batch_size * seq_len * hidden_size, 0.05f));
+
+  auto output =
+      fused_moe->forward_experts(hidden_states,
+                                 /*enable_all2all_communication=*/false);
+
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  // Verify output shape
+  CHECK_EQ(output.sizes().size(), 2) << "Output should be 2D tensor";
+  CHECK_EQ(output.size(0), batch_size * seq_len)
+      << "Batch size * seq_len should match";
+  CHECK_EQ(output.size(1), hidden_size) << "Hidden size should match";
+
+  // Set expected output values for precision verification
+  // The expected values should be calculated based on your specific test case
+  std::vector<float> expected_values;
+
+  // Fill expected_values with placeholder data using custom dimensions
+  expected_values.reserve(batch_size * seq_len * hidden_size);
+  for (size_t i = 0; i < batch_size; ++i) {
+    for (size_t j = 0; j < seq_len; ++j) {
+      for (size_t k = 0; k < hidden_size; ++k) {
+        expected_values.push_back(992.0f);  // calculated via vLLM MLU
+      }
+    }
+  }
+
+  set_expected_output(expected_values);
+
+  verify_precision(output, 1e-3, 1e-4);
+}
+
+TEST_F(FusedMoETest, W4A8GroupwiseSmokeTest) {
+  const int64_t batch_size = 4;
+  const int64_t seq_len = 8;
+  const int64_t hidden_size = 256;
+  const int64_t intermediate_size = 256;
+  const int64_t group_size = 128;
+  const int64_t num_experts = 8;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+
+  quant_args_.moe_weight_bits() = 4;
+  quant_args_.group_size() = group_size;
+
+  auto fused_moe = create_fused_moe(num_experts,
+                                    top_k,
+                                    num_expert_group,
+                                    topk_group,
+                                    route_scale,
+                                    hidden_size,
+                                    intermediate_size,
+                                    /*n_shared_experts=*/0);
+
+  auto weight_dict = create_w4a8_groupwise_test_weights(
+      num_experts, hidden_size, intermediate_size, group_size);
+  StateDict state_dict(weight_dict);
+  fused_moe->load_state_dict(state_dict);
+
+  auto hidden_states = create_custom_input(
+      {batch_size * seq_len, hidden_size},
+      std::vector<float>(batch_size * seq_len * hidden_size, 0.05f));
+
+  auto output =
+      fused_moe->forward_experts(hidden_states,
+                                 /*enable_all2all_communication=*/false);
+
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  EXPECT_EQ(output.dim(), 2);
+  EXPECT_EQ(output.size(0), batch_size * seq_len);
+  EXPECT_EQ(output.size(1), hidden_size);
+  EXPECT_NE(torch::sum(output).item<float>(), 0.0f);
+}
+
+TEST_F(FusedMoETest, PrepRouteMatchesBaseForward) {
+  const int64_t batch_size = 4;
+  const int64_t seq_len = 8;
+  const int64_t hidden_size = 256;
+  const int64_t intermediate_size = 256;
+  const int64_t num_experts = 8;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+
+  auto fused_moe = create_fused_moe(num_experts,
+                                    top_k,
+                                    num_expert_group,
+                                    topk_group,
+                                    route_scale,
+                                    hidden_size,
+                                    intermediate_size,
+                                    /*n_shared_experts=*/0);
+
+  auto weight_dict =
+      create_test_weights(num_experts, hidden_size, intermediate_size);
+  StateDict state_dict(weight_dict);
+  fused_moe->load_state_dict(state_dict);
+
+  auto hidden_states = create_custom_input(
+      {batch_size * seq_len, hidden_size},
+      std::vector<float>(batch_size * seq_len * hidden_size, 0.05f));
+
+  auto expected =
+      fused_moe->forward_experts(hidden_states,
+                                 /*enable_all2all_communication=*/false);
+  auto route_info = fused_moe->prep_route(hidden_states);
+  auto actual =
+      fused_moe->forward_experts(hidden_states,
+                                 /*enable_all2all_communication=*/false,
+                                 route_info);
+
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  verify_tensor_close(actual, expected, 1e-3, 1e-4);
+}
+
+TEST_F(FusedMoETest, NoReductionModeMatchesExternalReduce) {
+  set_tp_ctx(/*world_size=*/2);
+
+  const int64_t batch_size = 4;
+  const int64_t seq_len = 8;
+  const int64_t hidden_size = 256;
+  const int64_t intermediate_size = 256;
+  const int64_t num_experts = 8;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+
+  auto reduced_moe = create_fused_moe(num_experts,
+                                      top_k,
+                                      num_expert_group,
+                                      topk_group,
+                                      route_scale,
+                                      hidden_size,
+                                      intermediate_size,
+                                      /*n_shared_experts=*/1,
+                                      /*is_gated=*/true,
+                                      /*renormalize=*/0,
+                                      /*enable_result_reduction=*/true);
+  auto raw_moe = create_fused_moe(num_experts,
+                                  top_k,
+                                  num_expert_group,
+                                  topk_group,
+                                  route_scale,
+                                  hidden_size,
+                                  intermediate_size,
+                                  /*n_shared_experts=*/1,
+                                  /*is_gated=*/true,
+                                  /*renormalize=*/0,
+                                  /*enable_result_reduction=*/false);
+
+  auto weight_dict =
+      create_test_weights(num_experts, hidden_size, intermediate_size);
+  StateDict state_dict(weight_dict);
+  reduced_moe->load_state_dict(state_dict);
+  raw_moe->load_state_dict(state_dict);
+
+  auto hidden_states = create_custom_input(
+      {batch_size * seq_len, hidden_size},
+      std::vector<float>(batch_size * seq_len * hidden_size, 0.05f));
+
+  auto expected = reduced_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  auto actual = raw_moe->forward_experts(
+      hidden_states, /*enable_all2all_communication=*/false);
+  actual = parallel_state::reduce(actual, parallel_args_.tp_group_);
+
+  auto shared = raw_moe->forward_shared(hidden_states);
+  ASSERT_TRUE(shared.defined());
+  ASSERT_EQ(raw_moe->shared_pg(), parallel_args_.single_rank_group_);
+  shared = parallel_state::reduce(shared, raw_moe->shared_pg());
+  actual = actual + shared;
+
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  verify_tensor_close(actual, expected, 1e-3, 1e-4);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/indexer_test.cpp

@@ -0,0 +1,459 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/mlu/indexer.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <sstream>
+
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/attention.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+class MockDeepseekScalingRotaryEmbedding
+    : public DeepseekScalingRotaryEmbeddingImpl {
+ public:
+  MockDeepseekScalingRotaryEmbedding(int64_t rotary_dim,
+                                     int64_t max_position_embeddings,
+                                     int64_t rope_theta,
+                                     bool interleaved,
+                                     const torch::TensorOptions& options)
+      : DeepseekScalingRotaryEmbeddingImpl(rotary_dim,
+                                           rotary_dim,
+                                           max_position_embeddings,
+                                           max_position_embeddings,
+                                           rope_theta,
+                                           interleaved,
+                                           /*scaling_factor=*/2.5,
+                                           /*extrapolation_factor=*/1.,
+                                           /*attn_factor=*/40,
+                                           /*beta_fast=*/32,
+                                           /*beta_slow=*/1,
+                                           /*mscale=*/1.,
+                                           /*mscale_all_dim=*/1.,
+                                           options) {
+    mock_rope_ = std::make_shared<RotaryEmbeddingImpl>(
+        rotary_dim, max_position_embeddings, rope_theta, interleaved, options);
+  }
+  void forward(torch::Tensor& q,
+               torch::Tensor& k,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt) {
+    return mock_rope_->forward(
+        q, k, positions, cu_query_lens, max_query_len, is_prompt);
+  }
+
+ private:
+  std::shared_ptr<RotaryEmbeddingImpl> mock_rope_;
+};
+
+class IndexerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    torch::Device torch_device(Device::type_torch(), 0);
+    Device device(torch_device);
+    device.set_seed();
+    options_ = torch::TensorOptions()
+                   .dtype(torch::kBFloat16)
+                   .device(torch_device)
+                   .requires_grad(false);
+    int_option_ = options_.dtype(torch::kInt32);
+
+    parallel_args_ = test::create_default_parallel_args(mock_process_group_);
+    FLAGS_block_size = 1;
+  }
+
+  void TearDown() override {}
+
+  torch::Tensor create_random_tensor(
+      const std::vector<int64_t>& shape,
+      float min_val = -1.0f,
+      float max_val = 1.0f,
+      std::optional<torch::ScalarType> dtype = std::nullopt) {
+    auto opts = dtype.has_value() ? options_.dtype(dtype.value()) : options_;
+    return torch::rand(shape, opts) * (max_val - min_val) + min_val;
+  }
+
+  std::unordered_map<std::string, torch::Tensor> create_random_weights(
+      int64_t dim,
+      int64_t index_n_heads,
+      int64_t index_head_dim,
+      int64_t q_lora_rank) {
+    std::unordered_map<std::string, torch::Tensor> weight_dict;
+    weight_dict["wq_b.weight"] = create_random_tensor(
+        {index_n_heads * index_head_dim, q_lora_rank}, -0.1f, 0.1f);
+    weight_dict["wk.weight"] =
+        create_random_tensor({index_head_dim, dim}, -0.1f, 0.1f);
+    weight_dict["weights_proj.weight"] =
+        create_random_tensor({index_n_heads, dim}, -0.1f, 0.1f);
+    weight_dict["k_norm.weight"] =
+        create_random_tensor({index_head_dim}, -0.5f, 0.5f, torch::kFloat32);
+    weight_dict["k_norm.bias"] =
+        create_random_tensor({index_head_dim}, -0.5f, 0.5f, torch::kFloat32);
+
+    return weight_dict;
+  }
+
+  void populate_attention_metadata(AttentionMetadata& metadata,
+                                   int64_t batch_size,
+                                   int64_t max_query_len,
+                                   int64_t max_seq_len,
+                                   bool is_prefill,
+                                   int64_t max_num_blocks_per_seq) {
+    // q_cu_seq_lens
+    metadata.q_cu_seq_lens = torch::arange(
+        0, (batch_size + 1) * max_query_len, max_query_len, int_option_);
+
+    // kv_cu_seq_lens
+    metadata.kv_cu_seq_lens = torch::arange(
+        0, (batch_size + 1) * max_query_len, max_query_len, int_option_);
+
+    metadata.kv_seq_lens =
+        torch::full({batch_size}, max_query_len, int_option_);
+
+    metadata.block_table =
+        torch::zeros({batch_size, max_num_blocks_per_seq}, int_option_);
+
+    for (int64_t b = 0; b < batch_size; ++b) {
+      auto seq = torch::arange(b * max_query_len + 1,
+                               b * max_query_len + 1 + max_query_len,
+                               int_option_);
+      metadata.block_table[b].index_put_(
+          {torch::indexing::Slice(0, max_query_len)}, seq);
+    }
+
+    // slot_mapping
+    metadata.slot_mapping =
+        torch::arange(1, batch_size * max_query_len + 1, int_option_);
+
+    metadata.max_query_len = max_query_len;
+    metadata.max_seq_len = max_seq_len;
+    metadata.total_kv_len = batch_size * max_query_len;
+    metadata.compute_dtype = "bfloat16";
+    metadata.is_prefill = is_prefill;
+    metadata.is_chunked_prefill = false;
+  }
+
+  void populate_chunked_attention_metadata(AttentionMetadata& metadata,
+                                           int64_t batch_size,
+                                           int64_t history_len,
+                                           int64_t current_len,
+                                           int64_t max_num_blocks_per_seq) {
+    int64_t total_len = history_len + current_len;
+
+    metadata.q_cu_seq_lens = torch::arange(
+        0, (batch_size + 1) * current_len, current_len, int_option_);
+
+    metadata.kv_cu_seq_lens =
+        torch::arange(0, (batch_size + 1) * total_len, total_len, int_option_);
+
+    metadata.kv_seq_lens = torch::full({batch_size}, total_len, int_option_);
+
+    metadata.block_table =
+        torch::zeros({batch_size, max_num_blocks_per_seq}, int_option_);
+
+    for (int64_t b = 0; b < batch_size; ++b) {
+      auto seq =
+          torch::arange(b * total_len, b * total_len + total_len, int_option_);
+      metadata.block_table[b].index_put_({torch::indexing::Slice(0, total_len)},
+                                         seq);
+    }
+
+    metadata.slot_mapping =
+        torch::empty({batch_size * current_len}, int_option_);
+
+    for (int64_t b = 0; b < batch_size; ++b) {
+      auto slots = torch::arange(
+          b * total_len + history_len, b * total_len + total_len, int_option_);
+      metadata.slot_mapping.index_put_(
+          {torch::indexing::Slice(b * current_len, (b + 1) * current_len)},
+          slots);
+    }
+
+    metadata.max_query_len = current_len;
+    metadata.max_seq_len = total_len;
+    metadata.total_kv_len = batch_size * total_len;
+    metadata.compute_dtype = "bfloat16";
+    metadata.is_prefill = true;
+    metadata.is_chunked_prefill = true;
+  }
+
+  struct TestConfig {
+    int64_t dim = 7168;
+    int64_t index_n_heads = 64;
+    int64_t index_head_dim = 128;
+    int64_t qk_rope_head_dim = 64;
+    int64_t index_topk = 2048;
+    int64_t q_lora_rank = 1536;
+    int64_t max_position_embeddings = 8192;
+    int64_t rope_theta = 10000;
+    bool rope_interleaved = true;
+    int64_t head_kv = 1;
+    int64_t block_size = 1;
+    int64_t block_num = 10240;
+  };
+
+  struct TestInputs {
+    torch::Tensor x;
+    torch::Tensor q_norm;
+    torch::Tensor positions;
+    torch::Tensor k_cache;
+    std::unordered_map<std::string, torch::Tensor> weights;
+    AttentionMetadata metadata;
+  };
+
+  TestInputs create_inputs(int64_t batch_size,
+                           int64_t max_query_len,
+                           bool is_prefill,
+                           bool chunked_prefill = false,
+                           int64_t history_len = 0,
+                           bool use_default_rope = false) {
+    test_config_ = TestConfig();
+    if (use_default_rope) {
+      rotary_emb_ = std::make_shared<RotaryEmbeddingImpl>(
+          test_config_.qk_rope_head_dim,
+          test_config_.max_position_embeddings,
+          test_config_.rope_theta,
+          test_config_.rope_interleaved,
+          options_);
+    } else {
+      rotary_emb_ = std::make_shared<MockDeepseekScalingRotaryEmbedding>(
+          test_config_.qk_rope_head_dim,
+          test_config_.max_position_embeddings,
+          test_config_.rope_theta,
+          test_config_.rope_interleaved,
+          options_);
+    }
+
+    TestInputs inputs;
+    int64_t num_tokens = batch_size * max_query_len;
+
+    inputs.x =
+        create_random_tensor({num_tokens, test_config_.dim}, -1.0f, 1.0f);
+    inputs.q_norm = create_random_tensor(
+        {num_tokens, test_config_.q_lora_rank}, -1.0f, 1.0f);
+
+    inputs.positions =
+        torch::randint(0, max_query_len, {num_tokens}, int_option_);
+
+    inputs.k_cache = create_random_tensor({test_config_.block_num,
+                                           test_config_.head_kv,
+                                           test_config_.block_size,
+                                           test_config_.index_head_dim},
+                                          -0.5f,
+                                          0.5f);
+
+    inputs.weights = create_random_weights(test_config_.dim,
+                                           test_config_.index_n_heads,
+                                           test_config_.index_head_dim,
+                                           test_config_.q_lora_rank);
+
+    if (chunked_prefill) {
+      populate_chunked_attention_metadata(inputs.metadata,
+                                          batch_size,
+                                          history_len,
+                                          max_query_len,
+                                          history_len + max_query_len);
+    } else {
+      populate_attention_metadata(inputs.metadata,
+                                  batch_size,
+                                  max_query_len,
+                                  test_config_.max_position_embeddings,
+                                  is_prefill,
+                                  num_tokens);
+    }
+    return inputs;
+  }
+
+  std::tuple<torch::Tensor, torch::Tensor> run_indexer(TestInputs& inputs,
+                                                       bool is_prefill,
+                                                       bool enable_fused_qk) {
+    StateDict state_dict(inputs.weights);
+    QuantArgs quant_args;
+    auto indexer = Indexer(IndexerImpl(test_config_.dim,
+                                       test_config_.index_n_heads,
+                                       test_config_.index_head_dim,
+                                       test_config_.qk_rope_head_dim,
+                                       test_config_.index_topk,
+                                       test_config_.q_lora_rank,
+                                       enable_fused_qk,
+                                       rotary_emb_,
+                                       quant_args,
+                                       parallel_args_,
+                                       options_));
+    indexer->load_state_dict(state_dict);
+    return indexer->forward(inputs.x,
+                            inputs.q_norm,
+                            inputs.positions,
+                            inputs.k_cache,
+                            inputs.metadata,
+                            is_prefill);
+  }
+
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  TestConfig test_config_;
+  torch::TensorOptions options_;
+  torch::TensorOptions int_option_;
+  std::unique_ptr<xllm::ProcessGroup> mock_process_group_;
+  std::shared_ptr<RotaryEmbeddingBase> rotary_emb_;
+};
+
+TEST_F(IndexerTest, PrefillBatch) {
+  LOG(INFO) << "Testing Prefill (Small Batch)";
+  int64_t batch_size = 2;
+  int64_t max_query_len = 4096;
+  const bool is_prefill = true;
+  const bool enable_fused_qk = false;
+  int64_t num_tokens = batch_size * max_query_len;
+  TestInputs inputs = create_inputs(batch_size, max_query_len, is_prefill);
+  auto [new_block_tables, new_context_lens] =
+      run_indexer(inputs, is_prefill, enable_fused_qk);
+
+  EXPECT_EQ(new_block_tables.sizes().size(), 2)
+      << "new_block_tables should be 2D tensor";
+  EXPECT_EQ(new_context_lens.sizes().size(), 1)
+      << "new_context_lens should be 1D tensor";
+  EXPECT_EQ(new_block_tables.size(0), num_tokens) << "Batch size should match";
+  EXPECT_EQ(new_block_tables.size(1), test_config_.index_topk)
+      << "Top-k should match";
+
+  // Verify that the first value in new_block_tables is 1 (calculated via vLLM
+  // MLU)
+  EXPECT_EQ(new_block_tables.index({0, 0}).item<int64_t>(), 1)
+      << "The first value in new_block_tables should be 1";
+}
+
+TEST_F(IndexerTest, ChunkedPrefillBatch) {
+  LOG(INFO) << "Testing Chunked Prefill";
+  const int64_t batch_size = 2;
+  const int64_t history_len = 128;
+  const int64_t current_len = 64;
+  int64_t num_new_tokens = batch_size * current_len;
+  const bool is_prefill = true;
+  const bool is_chunked = true;
+  const bool enable_fused_qk = false;
+  TestInputs inputs = create_inputs(
+      batch_size, current_len, is_prefill, is_chunked, history_len);
+  auto [new_block_tables, new_context_lens] =
+      run_indexer(inputs, is_prefill, enable_fused_qk);
+
+  // Validations
+  // Shape Verification
+  EXPECT_EQ(new_block_tables.dim(), 2);
+  EXPECT_EQ(new_block_tables.size(0), num_new_tokens);  // [batch * current_len]
+  EXPECT_EQ(new_block_tables.size(1), test_config_.index_topk);
+
+  // Value Verification
+  auto top1_indices = new_block_tables.index({torch::indexing::Slice(), 0})
+                          .to(torch::kInt64)
+                          .cpu();
+  auto top1_sum = top1_indices.sum().item<int64_t>();
+  auto top1_max = top1_indices.max().item<int64_t>();
+
+  LOG(INFO) << "[top-1 block index] sum: " << top1_sum << ", max: " << top1_max;
+
+  // The expected value is calculated via vLLM MLU
+  int64_t expected_sum = 12288;
+  int64_t expected_max = 192;
+  EXPECT_EQ(top1_sum, expected_sum)
+      << "top-1 block index sum does not match ground truth";
+  EXPECT_EQ(top1_max, expected_max)
+      << "top-1 block index max does not match ground truth";
+}
+
+TEST_F(IndexerTest, CompareFusedVsNonFusedDecode) {
+  LOG(INFO) << "Testing Decode";
+  TestInputs inputs = create_inputs(128, 1, false);
+
+  auto [base_block_tables, base_context_lens] =
+      run_indexer(inputs, false, false);
+  auto [fused_block_tables, fused_context_lens] =
+      run_indexer(inputs, false, true);
+
+  auto fused_block_tables_slice = fused_block_tables.slice(1, 0, 1);
+  auto base_block_tables_slice = base_block_tables.slice(1, 0, 1);
+  test::verify_tensor_close(fused_context_lens.to(torch::kFloat32),
+                            base_context_lens.to(torch::kFloat32));
+  test::verify_tensor_close(fused_block_tables_slice.to(torch::kFloat32),
+                            base_block_tables_slice.to(torch::kFloat32));
+}
+
+TEST_F(IndexerTest, CompareFusedVsNonFusedMultipleRuns) {
+  LOG(INFO) << "Testing with multiple random seeds";
+
+  Device device(options_.device());
+  for (int i = 0; i < 3; ++i) {
+    LOG(INFO) << "Random seed iteration: " << i;
+    device.set_seed(i * 100);
+    TestInputs inputs = create_inputs(128, 1, false);
+
+    auto [base_block_tables, base_context_lens] =
+        run_indexer(inputs, false, false);
+    auto [fused_block_tables, fused_context_lens] =
+        run_indexer(inputs, false, true);
+
+    auto fused_block_tables_slice = fused_block_tables.slice(1, 0, 1);
+    auto base_block_tables_slice = base_block_tables.slice(1, 0, 1);
+    test::verify_tensor_close(fused_context_lens.to(torch::kFloat32),
+                              base_context_lens.to(torch::kFloat32));
+    test::verify_tensor_close(fused_block_tables_slice.to(torch::kFloat32),
+                              base_block_tables_slice.to(torch::kFloat32));
+  }
+}
+
+TEST_F(IndexerTest, CompareFusedVsNonFusedEdgeCaseSmall) {
+  LOG(INFO) << "Testing Edge Case (Very Small Input)";
+  TestInputs inputs = create_inputs(16, 1, false);
+
+  auto [base_block_tables, base_context_lens] =
+      run_indexer(inputs, false, false);
+  auto [fused_block_tables, fused_context_lens] =
+      run_indexer(inputs, false, true);
+
+  auto fused_block_tables_slice = fused_block_tables.slice(1, 0, 1);
+  auto base_block_tables_slice = base_block_tables.slice(1, 0, 1);
+  test::verify_tensor_close(fused_context_lens.to(torch::kFloat32),
+                            base_context_lens.to(torch::kFloat32));
+  test::verify_tensor_close(fused_block_tables_slice.to(torch::kFloat32),
+                            base_block_tables_slice.to(torch::kFloat32));
+}
+
+TEST_F(IndexerTest, DefaultRopeDecodePath) {
+  LOG(INFO) << "Testing default rope decode path";
+  TestInputs inputs = create_inputs(32, 1, false, false, 0, true);
+
+  auto [block_tables, context_lens] = run_indexer(inputs, false, false);
+
+  EXPECT_EQ(block_tables.dim(), 2);
+  EXPECT_EQ(block_tables.size(0), 32);
+  EXPECT_EQ(block_tables.size(1), test_config_.index_topk);
+  EXPECT_EQ(context_lens.dim(), 1);
+  EXPECT_EQ(context_lens.size(0), 32);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/mla_test.cpp

@@ -0,0 +1,460 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/deepseek_v2_attention.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+namespace {
+
+KVCache create_indexed_kv_cache(torch::Tensor key_cache,
+                                torch::Tensor index_cache) {
+  return KVCache(IndexedKVCacheTensors{
+      KVCacheTensors{key_cache, torch::Tensor()}, index_cache});
+}
+
+}  // namespace
+
+class DeepseekMLATest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    torch::Device torch_device(Device::type_torch(), 0);
+    Device device(torch_device);
+    device.set_seed();
+
+    FLAGS_block_size = 1;
+    // Initialize default model arguments for testing
+    model_args_ = create_mla_model_args();
+
+    // Initialize w8a8 quantization arguments
+    quant_args_ = test::create_default_quant_args();
+
+    // Initialize tensor options
+    options_ = torch::TensorOptions()
+                   .dtype(torch::kBFloat16)
+                   .device(torch_device)
+                   .requires_grad(false);
+
+    // Create mock ProcessGroup and initialize ParallelArgs
+    parallel_args_ = test::create_default_parallel_args(mock_process_group_);
+
+    init_test_weights();
+  }
+
+  ModelArgs create_mla_model_args(
+      const std::string& rope_type = "deepseek_yarn") {
+    ModelArgs model_args;
+    model_args.q_lora_rank() = 1536;
+    model_args.kv_lora_rank() = 512;
+    model_args.qk_nope_head_dim() = 128;
+    model_args.qk_rope_head_dim() = 64;
+    model_args.v_head_dim() = 128;
+    model_args.hidden_size() = 7168;
+    model_args.n_heads() = 128;
+    model_args.max_position_embeddings() = 163840;
+    model_args.rope_theta() = 10000;
+    model_args.rms_norm_eps() = 1e-06;
+
+    // rope_scaling config
+    model_args.rope_scaling_original_max_position_embeddings() = 4096;
+    model_args.rope_scaling_factor() = 40;
+    model_args.rope_extrapolation_factor() = 1.;
+    model_args.rope_scaling_attn_factor() = 1.;
+    model_args.rope_scaling_beta_fast() = 32;
+    model_args.rope_scaling_beta_slow() = 1;
+    model_args.rope_scaling_mscale() = 1.;
+    model_args.rope_scaling_mscale_all_dim() = 1.;
+    model_args.rope_scaling_rope_type() = rope_type;
+
+    // indexer
+    model_args.index_head_dim() = 128;
+    model_args.index_n_heads() = 64;
+    model_args.index_topk() = 2048;
+    model_args.enable_mla() = true;
+
+    return model_args;
+  }
+
+  void init_test_weights() {
+    int64_t q_lora_rank = model_args_.q_lora_rank();
+    int64_t kv_lora_rank = model_args_.kv_lora_rank();
+    int64_t qk_nope_head_dim = model_args_.qk_nope_head_dim();
+    int64_t qk_rope_head_dim = model_args_.qk_rope_head_dim();
+    int64_t index_topk = model_args_.index_topk();
+    int64_t index_n_heads = model_args_.index_n_heads();
+    int64_t index_head_dim = model_args_.index_head_dim();
+    int64_t v_head_dim = model_args_.v_head_dim();
+    int64_t hidden_size = model_args_.hidden_size();
+    int64_t num_heads = model_args_.n_heads();
+    int64_t max_position_embeddings = model_args_.max_position_embeddings();
+    int64_t qk_head_dim = qk_nope_head_dim + qk_rope_head_dim;
+    std::unordered_map<std::string, std::vector<int64_t>> qweight_map = {
+        {"model.layers.0.self_attn.o_proj.qweight",
+         {hidden_size, num_heads * v_head_dim}},
+        {"model.layers.0.self_attn.q_b_proj.qweight",
+         {num_heads * qk_head_dim, q_lora_rank}},
+    };
+    std::unordered_map<std::string, std::vector<int64_t>> scale_map = {
+        {"model.layers.0.self_attn.o_proj.per_channel_scale", {hidden_size}},
+        {"model.layers.0.self_attn.q_b_proj.per_channel_scale",
+         {num_heads * qk_head_dim}},
+        {"model.layers.0.self_attn.o_proj.smooth", {num_heads * v_head_dim}},
+        {"model.layers.0.self_attn.q_b_proj.smooth", {q_lora_rank}},
+    };
+    std::unordered_map<std::string, std::vector<int64_t>> weight_map = {
+        {"model.layers.0.self_attn.indexer.k_norm.bias", {index_head_dim}},
+        {"model.layers.0.self_attn.indexer.k_norm.weight", {index_head_dim}},
+        {"model.layers.0.self_attn.kv_a_layernorm.weight", {kv_lora_rank}},
+        {"model.layers.0.self_attn.q_a_layernorm.weight", {q_lora_rank}},
+        {"model.layers.0.self_attn.indexer.weights_proj.weight",
+         {index_n_heads, hidden_size}},
+        {"model.layers.0.self_attn.indexer.wk.weight",
+         {index_head_dim, hidden_size}},
+        {"model.layers.0.self_attn.indexer.wq_b.weight",
+         {index_n_heads * index_head_dim, q_lora_rank}},
+        {"model.layers.0.self_attn.kv_a_proj_with_mqa.weight",
+         {kv_lora_rank + qk_rope_head_dim, hidden_size}},
+        {"model.layers.0.self_attn.kv_b_proj.weight",
+         {num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank}},
+        {"model.layers.0.self_attn.q_a_proj.weight",
+         {q_lora_rank, hidden_size}},
+    };
+
+    auto float_option = options_.dtype(torch::kFloat32);
+    for (auto& [key, shape] : qweight_map) {
+      auto tensor =
+          test::seeded_tensor(key, shape, torch::kInt8, options_.device());
+      weight_dict_[key] = tensor;
+    }
+    for (auto& [key, shape] : scale_map) {
+      auto tensor =
+          test::seeded_tensor(key, shape, torch::kFloat32, options_.device()) *
+              0.01f +
+          0.03f;
+      weight_dict_[key] = tensor;
+    }
+    for (auto& [key, shape] : weight_map) {
+      auto tensor =
+          test::seeded_tensor(key, shape, torch::kBFloat16, options_.device()) *
+              0.01f +
+          0.02f;
+      weight_dict_[key] = tensor;
+    }
+  }
+
+  torch::Tensor create_seeded_hidden_states(const std::string& key,
+                                            int64_t num_tokens,
+                                            int64_t hidden_size) {
+    return test::seeded_tensor(key,
+                               {num_tokens, hidden_size},
+                               torch::kBFloat16,
+                               options_.device()) *
+           0.02f;
+  }
+
+  torch::Tensor create_seeded_cache(const std::string& key,
+                                    torch::IntArrayRef shape) {
+    return test::seeded_tensor(
+               key, shape, torch::kBFloat16, options_.device()) *
+           0.01f;
+  }
+
+  void populate_attention_metadata(AttentionMetadata& metadata,
+                                   int64_t batch_size,
+                                   int64_t max_query_len,
+                                   int64_t max_seq_len,
+                                   bool is_prefill,
+                                   int64_t max_num_batched_tokens) {
+    // Create q_cu_seq_lens tensor (cu_seq_q_lens)
+    // shape = [batch_size + 1], typically [0, 4, 8, 12, ...] if max_query_len=4
+    auto option_int = options_.dtype(torch::kInt32);
+    metadata.q_cu_seq_lens = torch::arange(
+        0, (batch_size + 1) * max_query_len, max_query_len, option_int);
+
+    // Create kv_cu_seq_lens tensor
+    metadata.kv_cu_seq_lens = torch::zeros({batch_size + 1}, option_int);
+
+    // Create seq_lens tensor
+    // Shape: [batch_size]
+    metadata.kv_seq_lens = torch::full({batch_size}, max_query_len, option_int);
+
+    // Create block_table tensor directly assigned to metadata
+    metadata.block_table =
+        torch::zeros({batch_size, max_num_batched_tokens}, option_int);
+
+    // Fill each batch with consecutive numbers
+    for (int64_t b = 0; b < batch_size; ++b) {
+      int64_t start_val = b * max_query_len + 1;
+      int64_t end_val = start_val + max_query_len;
+      // Generate sequence [start_val, ..., end_val-1]
+      auto seq = torch::arange(start_val, end_val, option_int);
+      metadata.block_table[b].index_put_(
+          {torch::indexing::Slice(0, max_query_len)}, seq);
+    }
+
+    // Create slot_mapping tensor directly assigned to metadata
+    metadata.slot_mapping =
+        torch::arange(1, batch_size * max_query_len + 1, option_int);
+
+    metadata.max_query_len = max_query_len;
+    metadata.max_seq_len = max_seq_len;
+    metadata.total_kv_len = batch_size * max_query_len;
+    metadata.compute_dtype = "half";
+    metadata.is_prefill = is_prefill;
+    metadata.is_chunked_prefill = false;
+    metadata.is_dummy = false;
+  }
+
+  torch::Tensor run_single_test(bool use_fused_mla_qkv,
+                                int64_t batch_size,
+                                int64_t max_query_len,
+                                bool is_prefill,
+                                const torch::Tensor& hidden_states,
+                                const torch::Tensor& positions,
+                                KVCache& kv_cache) {
+    OptimizationConfig optimization_config;
+    optimization_config.enable_fused_mla_kernel = use_fused_mla_qkv;
+    optimization_config.enable_fused_indexer_qk = false;
+    auto deepseek_mla = DeepseekV2Attention(model_args_,
+                                            quant_args_,
+                                            parallel_args_,
+                                            options_,
+                                            optimization_config);
+
+    std::string prefix = "model.layers.0.self_attn.";
+    StateDict state_dict(weight_dict_, prefix);
+    deepseek_mla->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+    // Create metadata object and populate it
+    AttentionMetadata metadata;
+    int64_t num_tokens = batch_size * max_query_len;
+    populate_attention_metadata(metadata,
+                                batch_size,
+                                max_query_len,
+                                model_args_.max_position_embeddings(),
+                                is_prefill,
+                                num_tokens);
+
+    auto output = deepseek_mla(positions, hidden_states, metadata, kv_cache);
+    xllm::Device device(options_.device());
+    device.synchronize_default_stream();
+
+    return output;
+  }
+
+  ModelArgs model_args_;
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+
+  // Helper to create a mock ProcessGroup for testing
+  std::unique_ptr<xllm::ProcessGroup> mock_process_group_;
+
+  std::unordered_map<std::string, torch::Tensor> weight_dict_;
+};
+
+TEST_F(DeepseekMLATest, PrefillTestRandomInput) {
+  int64_t batch_size = 2;
+  int64_t max_query_len = 5;
+  int64_t num_tokens = batch_size * max_query_len;
+  int64_t hidden_size = model_args_.hidden_size();
+
+  auto hidden_states = create_seeded_hidden_states(
+      "mla.prefill.hidden_states", num_tokens, hidden_size);
+  auto positions = torch::arange(max_query_len, options_.dtype(torch::kInt32))
+                       .repeat({batch_size});
+
+  int64_t block_num = 100;
+  auto k_cache = create_seeded_cache(
+      "mla.prefill.k_cache",
+      {block_num,
+       1,
+       1,
+       model_args_.qk_rope_head_dim() + model_args_.kv_lora_rank()});
+  auto index_cache =
+      create_seeded_cache("mla.prefill.index_cache",
+                          {block_num, 1, 1, model_args_.index_head_dim()});
+  KVCache kv_cache =
+      create_indexed_kv_cache(std::move(k_cache), std::move(index_cache));
+
+  auto output = run_single_test(false,
+                                batch_size,
+                                max_query_len,
+                                true,
+                                hidden_states,
+                                positions,
+                                kv_cache);
+
+  std::vector<float> results = {-7.25,
+                                -8.1875,
+                                2.5625,
+                                -4.5625,
+                                -1.2656,
+                                -7.3125,
+                                -6.4375,
+                                7.9375,
+                                -3.8906,
+                                5.1875};
+  auto slice_output =
+      output.flatten().slice(0, 0, 10).to(torch::kFloat32).cpu();
+  auto expected =
+      torch::tensor(results, torch::TensorOptions().dtype(torch::kFloat32));
+  ASSERT_TRUE(torch::allclose(slice_output, expected, 1e-3, 1e-4))
+      << "Prefill expected values mismatch. actual=" << slice_output
+      << ", expected=" << expected;
+}
+
+TEST_F(DeepseekMLATest, DecoderTestRandomInput) {
+  int64_t batch_size = 1;
+  int64_t max_query_len = 1;
+  int64_t num_tokens = batch_size * max_query_len;
+  int64_t hidden_size = model_args_.hidden_size();
+
+  auto hidden_states = create_seeded_hidden_states(
+      "mla.decoder.hidden_states", num_tokens, hidden_size);
+  auto positions = torch::arange(max_query_len, options_.dtype(torch::kInt32))
+                       .repeat({batch_size});
+
+  int64_t block_num = 100;
+  auto k_cache = create_seeded_cache(
+      "mla.decoder.k_cache",
+      {block_num,
+       1,
+       1,
+       model_args_.qk_rope_head_dim() + model_args_.kv_lora_rank()});
+  auto index_cache =
+      create_seeded_cache("mla.decoder.index_cache",
+                          {block_num, 1, 1, model_args_.index_head_dim()});
+  KVCache kv_cache =
+      create_indexed_kv_cache(std::move(k_cache), std::move(index_cache));
+
+  auto output_non_fused = run_single_test(false,
+                                          batch_size,
+                                          max_query_len,
+                                          false,
+                                          hidden_states,
+                                          positions,
+                                          kv_cache);
+  auto output_fused = run_single_test(true,
+                                      batch_size,
+                                      max_query_len,
+                                      false,
+                                      hidden_states,
+                                      positions,
+                                      kv_cache);
+  test::verify_tensor_close(output_fused, output_non_fused);
+}
+
+TEST_F(DeepseekMLATest, VariousBatchSizesTest) {
+  std::vector<int64_t> test_cases = {1, 2, 4, 8};
+  int32_t seq_len = 1;
+  for (const auto& batch_size : test_cases) {
+    LOG(INFO) << "Testing batch_size=" << batch_size << ", seq_len=" << seq_len;
+    int64_t num_tokens = batch_size * seq_len;
+    int64_t hidden_size = model_args_.hidden_size();
+    auto hidden_states = create_seeded_hidden_states(
+        "mla.batch.hidden_states." + std::to_string(batch_size),
+        num_tokens,
+        hidden_size);
+    auto positions = torch::arange(seq_len, options_.dtype(torch::kInt32))
+                         .repeat({batch_size});
+
+    int64_t block_num = 100;
+    auto k_cache = create_seeded_cache(
+        "mla.batch.k_cache." + std::to_string(batch_size),
+        {block_num,
+         1,
+         1,
+         model_args_.qk_rope_head_dim() + model_args_.kv_lora_rank()});
+    auto index_cache = create_seeded_cache(
+        "mla.batch.index_cache." + std::to_string(batch_size),
+        {block_num, 1, 1, model_args_.index_head_dim()});
+    KVCache kv_cache =
+        create_indexed_kv_cache(std::move(k_cache), std::move(index_cache));
+
+    bool is_prefill = false;
+    auto output_fused = run_single_test(true,
+                                        batch_size,
+                                        seq_len,
+                                        is_prefill,
+                                        hidden_states,
+                                        positions,
+                                        kv_cache);
+    auto output_non_fused = run_single_test(false,
+                                            batch_size,
+                                            seq_len,
+                                            is_prefill,
+                                            hidden_states,
+                                            positions,
+                                            kv_cache);
+    test::verify_tensor_close(output_fused, output_non_fused);
+  }
+}
+
+TEST_F(DeepseekMLATest, DefaultRopePrefillTest) {
+  model_args_ = create_mla_model_args("default");
+  init_test_weights();
+
+  int64_t batch_size = 2;
+  int64_t max_query_len = 3;
+  int64_t num_tokens = batch_size * max_query_len;
+  int64_t hidden_size = model_args_.hidden_size();
+
+  auto hidden_states = create_seeded_hidden_states(
+      "mla.default_rope.hidden_states", num_tokens, hidden_size);
+  auto positions = torch::arange(max_query_len, options_.dtype(torch::kInt32))
+                       .repeat({batch_size});
+
+  int64_t block_num = 32;
+  auto k_cache = create_seeded_cache(
+      "mla.default_rope.k_cache",
+      {block_num,
+       1,
+       1,
+       model_args_.qk_rope_head_dim() + model_args_.kv_lora_rank()});
+  auto index_cache =
+      create_seeded_cache("mla.default_rope.index_cache",
+                          {block_num, 1, 1, model_args_.index_head_dim()});
+  KVCache kv_cache =
+      create_indexed_kv_cache(std::move(k_cache), std::move(index_cache));
+
+  auto output = run_single_test(false,
+                                batch_size,
+                                max_query_len,
+                                true,
+                                hidden_states,
+                                positions,
+                                kv_cache);
+
+  EXPECT_EQ(output.dim(), 2);
+  EXPECT_EQ(output.size(0), num_tokens);
+  EXPECT_EQ(output.size(1), hidden_size);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/moe_gate_test.cpp

@@ -0,0 +1,271 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/mlu/moe_gate.h"
+
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+class MoEGateTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    model_args_ = test::create_default_model_args();
+    quant_args_ = QuantArgs();  // use empty quant
+    options_ = torch::TensorOptions()
+                   .dtype(torch::kBFloat16)
+                   .device(Device::type_torch(), 0)
+                   .requires_grad(false);
+    parallel_args_ = test::create_default_parallel_args(mock_process_group_);
+  }
+
+  void TearDown() override {
+    // Clean up if needed
+  }
+
+  // Set MoE gate related fields on model_args_. Call before constructing
+  // MoEGateImpl in each test.
+  void set_moe_gate_params(int64_t num_experts,
+                           int64_t top_k,
+                           int64_t num_expert_group,
+                           int64_t topk_group,
+                           double route_scale,
+                           int64_t hidden_size,
+                           bool renormalize,
+                           const std::string& scoring_func,
+                           const std::string& topk_method) {
+    model_args_.n_routed_experts() = static_cast<int32_t>(num_experts);
+    model_args_.num_experts_per_tok() = static_cast<int32_t>(top_k);
+    model_args_.n_group() = static_cast<int32_t>(num_expert_group);
+    model_args_.topk_group() = static_cast<int32_t>(topk_group);
+    model_args_.routed_scaling_factor() = static_cast<float>(route_scale);
+    model_args_.hidden_size() = hidden_size;
+    model_args_.norm_topk_prob() = renormalize;
+    model_args_.scoring_func() = scoring_func;
+    model_args_.topk_method() = topk_method;
+  }
+
+  // Build state dict with seeded tensors for gate (bfloat16, no quant).
+  std::unordered_map<std::string, torch::Tensor> create_gate_weights_seeded(
+      int64_t num_experts,
+      int64_t hidden_size,
+      const std::string& topk_method) {
+    std::unordered_map<std::string, torch::Tensor> weight_dict;
+    // Gate projection: ReplicatedLinear(hidden_size, num_experts, false, ...)
+    // Weight shape [num_experts, hidden_size] (out_features, in_features).
+    auto gate_weight = test::seeded_tensor("moe_gate_tests.gate_proj.weight",
+                                           {num_experts, hidden_size},
+                                           torch::kBFloat16,
+                                           options_.device());
+    weight_dict["weight"] = gate_weight;
+
+    if (topk_method == "noaux_tc") {
+      auto e_bias =
+          test::seeded_tensor("moe_gate_tests.e_score_correction_bias",
+                              {num_experts},
+                              torch::kBFloat16,
+                              options_.device());
+      weight_dict["e_score_correction_bias"] = e_bias;
+    }
+    return weight_dict;
+  }
+
+  ModelArgs model_args_;
+  // Run forward, assert shapes, and verify min/max/sum against expected values.
+  void run_forward_and_expect(MoEGateImpl* moe_gate,
+                              int64_t num_tokens,
+                              int64_t hidden_size,
+                              int64_t top_k,
+                              const std::string& seed_key_prefix,
+                              double expected_rw_min,
+                              double expected_rw_max,
+                              double expected_rw_sum,
+                              double expected_eid_min,
+                              double expected_eid_max,
+                              double expected_eid_sum) {
+    auto hidden_states = test::seeded_tensor(seed_key_prefix + ".hidden_states",
+                                             {num_tokens, hidden_size},
+                                             torch::kBFloat16,
+                                             options_.device());
+    auto [reduce_weight, expert_id] = moe_gate->forward(hidden_states);
+
+    Device device(options_.device());
+    device.synchronize_default_stream();
+
+    ASSERT_TRUE(reduce_weight.defined()) << "reduce_weight should be defined";
+    ASSERT_TRUE(expert_id.defined()) << "expert_id should be defined";
+    ASSERT_EQ(reduce_weight.dim(), 2);
+    ASSERT_EQ(expert_id.dim(), 2);
+    ASSERT_EQ(reduce_weight.size(0), num_tokens);
+    ASSERT_EQ(reduce_weight.size(1), top_k);
+    ASSERT_EQ(expert_id.size(0), num_tokens);
+    ASSERT_EQ(expert_id.size(1), top_k);
+    double rw_sum =
+        torch::sum(reduce_weight.flatten().to(torch::kFloat64)).item<double>();
+    ASSERT_NE(rw_sum, 0.0) << "reduce_weight sum should not be zero";
+
+    test::expect_tensor_stats(
+        reduce_weight, expected_rw_min, expected_rw_max, expected_rw_sum);
+    test::expect_tensor_stats(
+        expert_id, expected_eid_min, expected_eid_max, expected_eid_sum);
+  }
+
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+  std::unique_ptr<xllm::ProcessGroup> mock_process_group_;
+};
+
+// Sigmoid scoring: typical MoE config (multi-expert, top_k=2, groups).
+TEST_F(MoEGateTest, Sigmoid) {
+  const int64_t batch_size = 16;
+  const int64_t seq_len = 32;
+  const int64_t hidden_size = 7168;
+  const int64_t num_experts = 16;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+  const bool renormalize = true;
+  const std::string scoring_func = "sigmoid";
+  const std::string topk_method = "noaux_tc";
+
+  set_moe_gate_params(num_experts,
+                      top_k,
+                      num_expert_group,
+                      topk_group,
+                      route_scale,
+                      hidden_size,
+                      renormalize,
+                      scoring_func,
+                      topk_method);
+  MoEGateImpl moe_gate(model_args_, quant_args_, options_);
+  auto weight_dict =
+      create_gate_weights_seeded(num_experts, hidden_size, topk_method);
+  StateDict state_dict(weight_dict);
+  moe_gate.load_state_dict(state_dict);
+
+  int64_t num_tokens = batch_size * seq_len;
+  run_forward_and_expect(&moe_gate,
+                         num_tokens,
+                         hidden_size,
+                         top_k,
+                         "moe_gate_tests.sigmoid",
+                         /*reduce_weight*/ 1.25,
+                         1.25,
+                         1280.0,
+                         /*expert_id*/ 6.0,
+                         7.0,
+                         6656.0);
+}
+
+// Softmax scoring: same shape as sigmoid, different scoring path.
+TEST_F(MoEGateTest, Softmax) {
+  const int64_t batch_size = 16;
+  const int64_t seq_len = 32;
+  const int64_t hidden_size = 7168;
+  const int64_t num_experts = 16;
+  const int64_t num_expert_group = 4;
+  const int64_t topk_group = 4;
+  const int64_t top_k = 2;
+  const double route_scale = 2.5;
+  const bool renormalize = true;
+  const std::string scoring_func = "softmax";
+  const std::string topk_method = "";
+
+  set_moe_gate_params(num_experts,
+                      top_k,
+                      num_expert_group,
+                      topk_group,
+                      route_scale,
+                      hidden_size,
+                      renormalize,
+                      scoring_func,
+                      topk_method);
+  MoEGateImpl moe_gate(model_args_, quant_args_, options_);
+  auto weight_dict =
+      create_gate_weights_seeded(num_experts, hidden_size, topk_method);
+  StateDict state_dict(weight_dict);
+  moe_gate.load_state_dict(state_dict);
+
+  int64_t num_tokens = batch_size * seq_len;
+  run_forward_and_expect(&moe_gate,
+                         num_tokens,
+                         hidden_size,
+                         top_k,
+                         "moe_gate_tests.softmax",
+                         /*reduce_weight*/ 3.16604e-14,
+                         2.5,
+                         1280.0,
+                         /*expert_id*/ 0.0,
+                         14.0,
+                         4413.0);
+}
+
+// Sigmoid with topk_group=1
+TEST_F(MoEGateTest, SigmoidTopkGroup1) {
+  const int64_t batch_size = 8;
+  const int64_t seq_len = 16;
+  const int64_t hidden_size = 1024;
+  const int64_t num_experts = 8;
+  const int64_t num_expert_group = 1;
+  const int64_t topk_group = 1;
+  const int64_t top_k = 2;
+  const double route_scale = 1.0;
+  const bool renormalize = true;
+  const std::string scoring_func = "sigmoid";
+  const std::string topk_method = "";
+
+  set_moe_gate_params(num_experts,
+                      top_k,
+                      num_expert_group,
+                      topk_group,
+                      route_scale,
+                      hidden_size,
+                      renormalize,
+                      scoring_func,
+                      topk_method);
+  MoEGateImpl moe_gate(model_args_, quant_args_, options_);
+  auto weight_dict =
+      create_gate_weights_seeded(num_experts, hidden_size, topk_method);
+  StateDict state_dict(weight_dict);
+  moe_gate.load_state_dict(state_dict);
+
+  int64_t num_tokens = batch_size * seq_len;
+  run_forward_and_expect(&moe_gate,
+                         num_tokens,
+                         hidden_size,
+                         top_k,
+                         "moe_gate_tests.sigmoid_topk1",
+                         /*reduce_weight*/ 0.5,
+                         0.5,
+                         128.0,
+                         /*expert_id*/ 0.0,
+                         1.0,
+                         128.0);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/qwen2_attention_test.cpp

@@ -0,0 +1,805 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/common/qwen2_attention.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <cmath>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen2AttentionTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    torch::Device device(Device::type_torch(), 0);
+    Device xllm_device(device);
+    xllm_device.set_seed(42);
+    model_args_.model_type() = "qwen2";
+    model_args_.head_dim() = 128;
+    model_args_.hidden_size() = 1024;
+    model_args_.n_heads() = 16;
+    model_args_.n_kv_heads() = 8;
+    model_args_.max_position_embeddings() = 2048;
+    model_args_.rope_theta() = 1000000.0f;
+    model_args_.rms_norm_eps() = 1e-6f;
+    model_args_.rope_scaling_factor() = 1.0f;
+    model_args_.hidden_act() = "silu";
+
+    options_ = torch::TensorOptions().dtype(torch::kBFloat16).device(device);
+
+    process_group_ = create_process_group(
+        0, 1, 1, 3331, false, "localhost", "tp_group", device);
+    parallel_args_.tp_group_ = process_group_.get();
+
+    int64_t block_num = 100;
+    int64_t n_kv_heads = model_args_.n_kv_heads().value();
+    int64_t head_dim = model_args_.head_dim();
+    int64_t block_size = 16;
+
+    auto k_cache = MakeNoise("qwen2_attention_test.k_cache",
+                             {block_num, n_kv_heads, block_size, head_dim},
+                             0.01f);
+    auto v_cache = MakeNoise("qwen2_attention_test.v_cache",
+                             {block_num, n_kv_heads, block_size, head_dim},
+                             0.01f);
+    kv_cache_ = CreateKvCache(k_cache, v_cache);
+
+    context_ = ModelContext(parallel_args_, model_args_, QuantArgs(), options_);
+    InitTestWeights();
+  }
+
+  void InitTestWeights() {
+    int64_t hidden_size = model_args_.hidden_size();
+    int64_t n_heads = model_args_.n_heads();
+    int64_t n_kv_heads = model_args_.n_kv_heads().value();
+    int64_t head_dim = model_args_.head_dim();
+    int64_t q_size = n_heads * head_dim;
+    int64_t kv_size = n_kv_heads * head_dim;
+
+    const std::string weight_seed_prefix = "qwen2_attention_test.";
+    auto seeded = [this, &weight_seed_prefix](const std::string& name,
+                                              torch::IntArrayRef shape) {
+      return test::seeded_tensor(weight_seed_prefix + name,
+                                 shape,
+                                 torch::typeMetaToScalarType(options_.dtype()),
+                                 options_.device());
+    };
+
+    std::unordered_map<std::string, torch::Tensor> weight_map = {
+        {"q_proj.weight", seeded("q_proj.weight", {q_size, hidden_size})},
+        {"k_proj.weight", seeded("k_proj.weight", {kv_size, hidden_size})},
+        {"v_proj.weight", seeded("v_proj.weight", {kv_size, hidden_size})},
+        {"q_proj.bias", seeded("q_proj.bias", {q_size})},
+        {"k_proj.bias", seeded("k_proj.bias", {kv_size})},
+        {"v_proj.bias", seeded("v_proj.bias", {kv_size})},
+        {"o_proj.weight", seeded("o_proj.weight", {hidden_size, q_size})},
+    };
+
+    for (auto& [name, tensor] : weight_map) {
+      tensor = tensor / torch::sqrt(torch::tensor(tensor.size(0), options_));
+      weight_dict_["model.layers.0.self_attn." + name] = tensor;
+    }
+  }
+
+  torch::Tensor MakeNoise(const std::string& key,
+                          torch::IntArrayRef shape,
+                          float stddev) const {
+    auto noise =
+        test::seeded_tensor(key,
+                            shape,
+                            torch::typeMetaToScalarType(options_.dtype()),
+                            options_.device());
+    return (noise - 0.5f) * (std::sqrt(12.0f) * stddev);
+  }
+
+  int64_t GetBlockNum(int64_t seq_len) const {
+    const int64_t block_size = 16;
+    return (seq_len + block_size - 1) / block_size + 1;
+  }
+
+  torch::Tensor MakeBlockTable(int64_t batch_size, int64_t seq_len) const {
+    auto options_int = options_.dtype(torch::kInt32);
+    const int64_t block_num_per_req = GetBlockNum(seq_len);
+    std::vector<int32_t> block_table_vec;
+    block_table_vec.reserve(batch_size * block_num_per_req);
+    for (int64_t b = 0; b < batch_size; ++b) {
+      for (int64_t i = 0; i < block_num_per_req; ++i) {
+        block_table_vec.push_back(
+            static_cast<int32_t>(b * block_num_per_req + i));
+      }
+    }
+    return torch::tensor(block_table_vec, options_int)
+        .reshape({batch_size, block_num_per_req});
+  }
+
+  torch::Tensor MakeSlotMap(int64_t batch_size,
+                            int64_t token_len,
+                            int64_t kv_seq_len) const {
+    auto options_int = options_.dtype(torch::kInt32);
+    const int64_t block_size = 16;
+    const int64_t block_num_per_req = GetBlockNum(kv_seq_len);
+    const int64_t slot_num_per_req = block_num_per_req * block_size;
+    const int64_t start_pos = kv_seq_len - token_len;
+    std::vector<int32_t> slot_map_vec;
+    slot_map_vec.reserve(batch_size * token_len);
+    for (int64_t b = 0; b < batch_size; ++b) {
+      for (int64_t i = 0; i < token_len; ++i) {
+        slot_map_vec.push_back(
+            static_cast<int32_t>(b * slot_num_per_req + start_pos + i));
+      }
+    }
+    return torch::tensor(slot_map_vec, options_int);
+  }
+
+  KVCache CreateKvCache(torch::Tensor key_cache,
+                        torch::Tensor value_cache) const {
+    return KVCache(KVCacheTensors{
+        key_cache,
+        value_cache,
+    });
+  }
+
+  KVCache CreateQuantizedKvCache(torch::Tensor key_cache,
+                                 torch::Tensor value_cache,
+                                 torch::Tensor key_cache_scale,
+                                 torch::Tensor value_cache_scale) const {
+    return KVCache(QuantizedKVCacheTensors{
+        KVCacheTensors{
+            key_cache,
+            value_cache,
+        },
+        key_cache_scale,
+        value_cache_scale,
+    });
+  }
+
+  AttentionMetadata CreateAttentionMetadata(int64_t batch_size,
+                                            int64_t seq_len,
+                                            bool is_prefill,
+                                            int64_t max_seq_len,
+                                            bool is_chunked_prefill = false) {
+    AttentionMetadata metadata;
+    auto options_int = options_.dtype(torch::kInt32);
+
+    if (is_prefill && !is_chunked_prefill) {
+      // Regular prefill: query and kv have same sequence lengths
+      metadata.q_cu_seq_lens =
+          torch::arange(0, (batch_size + 1) * seq_len, seq_len, options_int);
+      metadata.kv_cu_seq_lens = metadata.q_cu_seq_lens;
+      // Keep paged-cache writes aligned with the deterministic block table.
+      metadata.slot_mapping = MakeSlotMap(batch_size, seq_len, seq_len);
+      metadata.kv_seq_lens = torch::full({batch_size}, seq_len, options_int);
+      metadata.block_table = MakeBlockTable(batch_size, seq_len);
+    } else if (is_chunked_prefill) {
+      // Chunked prefill: query has chunk_len, kv has full seq_len
+      int64_t chunk_len = seq_len;   // current chunk length
+      int64_t kv_seq_len = seq_len;  // accumulated kv length
+      const int64_t num_blocks_per_req = GetBlockNum(kv_seq_len);
+      metadata.q_cu_seq_lens = torch::arange(
+          0, (batch_size + 1) * chunk_len, chunk_len, options_int);
+      metadata.kv_cu_seq_lens = torch::arange(
+          0, (batch_size + 1) * kv_seq_len, kv_seq_len, options_int);
+      metadata.slot_mapping =
+          torch::arange(0, batch_size * chunk_len, options_int);
+      metadata.kv_seq_lens = torch::full({batch_size}, kv_seq_len, options_int);
+      // For chunked prefill with dequant_from_paged_cache, block_table must
+      // correspond to slot_mapping. Each batch uses sequential blocks.
+      // batch 0: blocks [0, num_blocks_per_req-1]
+      // batch 1: blocks [num_blocks_per_req, 2*num_blocks_per_req-1]
+      std::vector<int32_t> block_table_vec;
+      for (int64_t b = 0; b < batch_size; ++b) {
+        for (int64_t i = 0; i < num_blocks_per_req; ++i) {
+          block_table_vec.push_back(b * num_blocks_per_req + i);
+        }
+      }
+      metadata.block_table = torch::tensor(block_table_vec, options_int)
+                                 .reshape({batch_size, num_blocks_per_req});
+    } else {
+      // Decode: query length is 1
+      metadata.q_cu_seq_lens = torch::arange(0, batch_size + 1, 1, options_int);
+      metadata.kv_cu_seq_lens =
+          torch::arange(0, batch_size + 1, seq_len, options_int);
+      // Append the decode token to the last slot of each request.
+      metadata.slot_mapping = MakeSlotMap(batch_size, 1, seq_len);
+      metadata.kv_seq_lens = torch::full({batch_size}, seq_len, options_int);
+      metadata.block_table = MakeBlockTable(batch_size, seq_len);
+    }
+
+    metadata.max_query_len = (is_prefill || is_chunked_prefill) ? seq_len : 1;
+    metadata.max_seq_len = max_seq_len;
+    metadata.total_kv_len = batch_size * seq_len;
+    metadata.compute_dtype = "half";
+    metadata.is_prefill = is_prefill && !is_chunked_prefill;
+    metadata.is_chunked_prefill = is_chunked_prefill;
+    metadata.is_dummy = false;
+
+    return metadata;
+  }
+
+  ModelArgs model_args_;
+  ModelContext context_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+  std::unordered_map<std::string, torch::Tensor> weight_dict_;
+  std::unique_ptr<ProcessGroup> process_group_ = nullptr;
+  KVCache kv_cache_;
+};
+
+TEST_F(Qwen2AttentionTest, PrefillTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+  int64_t batch_size = 2;
+  int64_t seq_len = 128;
+  int64_t hidden_size = model_args_.hidden_size();
+  int64_t num_tokens = batch_size * seq_len;
+
+  auto hidden_states = MakeNoise("qwen2_attention_test.prefill.hidden_states",
+                                 {num_tokens, hidden_size},
+                                 0.02f);
+  auto positions = torch::arange(0, seq_len, options_.dtype(torch::kInt32))
+                       .repeat({batch_size});
+
+  auto metadata = CreateAttentionMetadata(batch_size, seq_len, true, seq_len);
+
+  auto output = qwen2_attention(positions, hidden_states, metadata, kv_cache_);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  CHECK_EQ(output.sizes(), torch::IntArrayRef({num_tokens, hidden_size}));
+  auto test_output = output.flatten().slice(0, 0, 10).unsqueeze(0);
+  std::vector<float> expected_values = {0.6796875f,
+                                        0.67578125f,
+                                        0.6875f,
+                                        0.65625f,
+                                        0.6640625f,
+                                        0.6796875f,
+                                        0.68359375f,
+                                        0.67578125f,
+                                        0.6796875f,
+                                        0.66796875f};
+  test::verify_precision(test_output, expected_values, 1e-5, 1e-6);
+}
+
+TEST_F(Qwen2AttentionTest, DecodeTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  int64_t batch_size = 4;
+  int64_t seq_len = 256;
+  int64_t hidden_size = model_args_.hidden_size();
+  int64_t num_tokens = batch_size;
+
+  auto hidden_states = MakeNoise("qwen2_attention_test.decode.hidden_states",
+                                 {num_tokens, hidden_size},
+                                 0.02f);
+  auto positions =
+      torch::full({num_tokens}, seq_len, options_.dtype(torch::kInt32));
+
+  auto metadata =
+      CreateAttentionMetadata(batch_size, seq_len + 1, false, seq_len + 1);
+
+  auto output = qwen2_attention(positions, hidden_states, metadata, kv_cache_);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  CHECK_EQ(output.sizes(), torch::IntArrayRef({num_tokens, hidden_size}));
+  auto test_output = output.flatten().slice(0, 0, 10).unsqueeze(0);
+  std::vector<float> expected_values = {0.0005264282f,
+                                        0.0008239746f,
+                                        0.0005722046f,
+                                        0.0006027222f,
+                                        0.000831604f,
+                                        0.0004405975f,
+                                        0.001037598f,
+                                        0.001083374f,
+                                        0.000289917f,
+                                        0.0007820129f};
+  test::verify_precision(test_output, expected_values, 1e-5, 1e-6);
+}
+
+TEST_F(Qwen2AttentionTest, MixedSequenceLengthTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  std::vector<int64_t> seq_lens = {32, 64, 128};
+  std::vector<int64_t> cu_seq_lens = {0};
+  int64_t total_tokens = 0;
+  int32_t seq_lens_size = seq_lens.size();
+
+  for (auto len : seq_lens) {
+    total_tokens += len;
+    cu_seq_lens.push_back(total_tokens);
+  }
+
+  int64_t hidden_size = model_args_.hidden_size();
+  auto hidden_states = MakeNoise("qwen2_attention_test.mix.hidden_states",
+                                 {total_tokens, hidden_size},
+                                 0.02f);
+
+  std::vector<int32_t> positions_vec;
+  for (size_t i = 0; i < seq_lens_size; ++i) {
+    for (int64_t j = 0; j < seq_lens[i]; ++j) {
+      positions_vec.push_back(j);
+    }
+  }
+  auto positions = torch::tensor(positions_vec, options_.dtype(torch::kInt32));
+
+  AttentionMetadata metadata;
+  auto options_int = options_.dtype(torch::kInt32);
+
+  metadata.q_cu_seq_lens = torch::tensor(cu_seq_lens, options_int);
+  metadata.kv_cu_seq_lens = metadata.q_cu_seq_lens;
+  metadata.kv_seq_lens = torch::tensor(seq_lens, options_int);
+  metadata.block_table = torch::zeros({seq_lens_size, 1}, options_int);
+  metadata.slot_mapping = torch::arange(0, total_tokens, options_int);
+  metadata.max_query_len = *std::max_element(seq_lens.begin(), seq_lens.end());
+  metadata.max_seq_len = model_args_.max_position_embeddings();
+  metadata.compute_dtype = "half";
+  metadata.is_prefill = true;
+  metadata.is_chunked_prefill = false;
+  metadata.is_dummy = false;
+
+  auto output = qwen2_attention(positions, hidden_states, metadata, kv_cache_);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  CHECK_EQ(output.sizes(), torch::IntArrayRef({total_tokens, hidden_size}));
+  auto test_output = output.flatten().slice(0, 0, 10).unsqueeze(0);
+  std::vector<float> expected_values = {0.07763672f,
+                                        0.08349609f,
+                                        0.08496094f,
+                                        0.08349609f,
+                                        0.07958984f,
+                                        0.08740234f,
+                                        0.09130859f,
+                                        0.08398438f,
+                                        0.08642578f,
+                                        0.07958984f};
+  test::verify_precision(test_output, expected_values, 1e-5, 1e-6);
+}
+
+TEST_F(Qwen2AttentionTest, QuantizedKVCachePrefillTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  // Test parameters
+  int64_t batch_size = 2;
+  int64_t seq_len = 128;
+  int64_t hidden_size = model_args_.hidden_size();
+  int64_t num_tokens = batch_size * seq_len;
+  int64_t block_num = 100;
+  int64_t n_kv_heads = model_args_.n_kv_heads().value();
+  int64_t head_dim = model_args_.head_dim();
+  int64_t block_size = 16;
+
+  // Create INT8 KV cache tensors using seeded tensors for reproducibility
+  auto k_cache =
+      test::seeded_tensor("qwen2_quant_test.k_cache",
+                          {block_num, n_kv_heads, block_size, head_dim},
+                          torch::kInt8,
+                          options_.device());
+  auto v_cache =
+      test::seeded_tensor("qwen2_quant_test.v_cache",
+                          {block_num, n_kv_heads, block_size, head_dim},
+                          torch::kInt8,
+                          options_.device());
+
+  // Create float32 scale tensors
+  auto k_cache_scale = test::seeded_tensor("qwen2_quant_test.k_scale",
+                                           {block_num, n_kv_heads, block_size},
+                                           torch::kFloat32,
+                                           options_.device());
+  auto v_cache_scale = test::seeded_tensor("qwen2_quant_test.v_scale",
+                                           {block_num, n_kv_heads, block_size},
+                                           torch::kFloat32,
+                                           options_.device());
+
+  KVCache quant_kv_cache =
+      CreateQuantizedKvCache(k_cache, v_cache, k_cache_scale, v_cache_scale);
+
+  // Create input tensors using seeded tensors
+  auto hidden_states = test::seeded_tensor("qwen2_quant_test.hidden_states",
+                                           {num_tokens, hidden_size},
+                                           torch::kBFloat16,
+                                           options_.device());
+  auto positions = test::seeded_tensor("qwen2_quant_test.positions",
+                                       {num_tokens},
+                                       torch::kInt32,
+                                       options_.device());
+
+  auto metadata = CreateAttentionMetadata(batch_size, seq_len, true, seq_len);
+
+  // Run forward with quantized KV cache
+  auto output =
+      qwen2_attention(positions, hidden_states, metadata, quant_kv_cache);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes(), torch::IntArrayRef({num_tokens, hidden_size}));
+
+  // Verify precision using expect_tensor_stats
+  // Expected values from Phase 1 print test
+  test::expect_tensor_stats(output,
+                            /*expected_min=*/240,
+                            /*expected_max=*/280,
+                            /*expected_sum=*/67867024);
+}
+
+TEST_F(Qwen2AttentionTest, QuantizedKVCacheDecodeDiagnosticTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  // Test parameters - use minimal batch size for diagnosis
+  int64_t batch_size = 1;
+  int64_t seq_len = 16;  // Start with a small sequence length
+  int64_t hidden_size = model_args_.hidden_size();
+  int64_t num_tokens = batch_size;
+  int64_t block_num = 100;
+  int64_t n_kv_heads = model_args_.n_kv_heads().value();
+  int64_t head_dim = model_args_.head_dim();
+  int64_t block_size = 16;
+
+  auto int8_options =
+      torch::TensorOptions().dtype(torch::kInt8).device(options_.device());
+  auto float_options =
+      torch::TensorOptions().dtype(torch::kFloat32).device(options_.device());
+
+  // 1. Use zeros for INT8 cache to avoid random value issues
+  auto k_cache =
+      torch::zeros({block_num, n_kv_heads, block_size, head_dim}, int8_options);
+  auto v_cache =
+      torch::zeros({block_num, n_kv_heads, block_size, head_dim}, int8_options);
+
+  // 2. Use ones for scale to avoid scale=0 issues
+  auto k_cache_scale =
+      torch::ones({block_num, n_kv_heads, block_size}, float_options);
+  auto v_cache_scale =
+      torch::ones({block_num, n_kv_heads, block_size}, float_options);
+
+  // Verify cache and scale shapes and dtypes
+  ASSERT_EQ(k_cache.sizes(),
+            torch::IntArrayRef({block_num, n_kv_heads, block_size, head_dim}));
+  ASSERT_EQ(k_cache.scalar_type(), torch::kInt8);
+  ASSERT_EQ(k_cache_scale.sizes(),
+            torch::IntArrayRef({block_num, n_kv_heads, block_size}));
+  ASSERT_EQ(k_cache_scale.scalar_type(), torch::kFloat32);
+
+  KVCache quant_kv_cache = CreateQuantizedKvCache(std::move(k_cache),
+                                                  std::move(v_cache),
+                                                  std::move(k_cache_scale),
+                                                  std::move(v_cache_scale));
+
+  // Create input tensors using seeded tensors
+  auto hidden_states = test::seeded_tensor("qwen2_decode_diag.hidden_states",
+                                           {num_tokens, hidden_size},
+                                           torch::kBFloat16,
+                                           options_.device());
+  auto positions =
+      torch::full({num_tokens}, seq_len - 1, options_.dtype(torch::kInt32));
+
+  auto metadata = CreateAttentionMetadata(batch_size, seq_len, false, seq_len);
+
+  // Run forward with quantized KV cache
+  auto output =
+      qwen2_attention(positions, hidden_states, metadata, quant_kv_cache);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes(), torch::IntArrayRef({num_tokens, hidden_size}));
+
+  // Print output stats for debugging
+  torch::Tensor flat = output.flatten().to(torch::kFloat32).cpu();
+  double out_min = torch::min(flat).item<double>();
+  double out_max = torch::max(flat).item<double>();
+  double out_sum = torch::sum(flat).item<double>();
+  LOG(INFO) << "Decode Diagnostic output - min: " << out_min
+            << ", max: " << out_max << ", sum: " << out_sum;
+}
+
+// Phase 4: Decode test with controlled seeded values
+TEST_F(Qwen2AttentionTest, QuantizedKVCacheDecodeTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  // Test parameters - match DecodeTest configuration
+  int64_t batch_size = 4;
+  int64_t seq_len = 256;
+  int64_t hidden_size = model_args_.hidden_size();
+  int64_t num_tokens = batch_size;
+  int64_t block_num = 100;
+  int64_t n_kv_heads = model_args_.n_kv_heads().value();
+  int64_t head_dim = model_args_.head_dim();
+  int64_t block_size = 16;
+
+  // Create INT8 KV cache tensors using seeded tensors for reproducibility
+  auto k_cache =
+      test::seeded_tensor("qwen2_quant_decode.k_cache",
+                          {block_num, n_kv_heads, block_size, head_dim},
+                          torch::kInt8,
+                          options_.device());
+  auto v_cache =
+      test::seeded_tensor("qwen2_quant_decode.v_cache",
+                          {block_num, n_kv_heads, block_size, head_dim},
+                          torch::kInt8,
+                          options_.device());
+
+  // Create scale tensors with controlled range [0.5, 1.5] to avoid extreme
+  // values
+  auto k_cache_scale_raw =
+      test::seeded_tensor("qwen2_quant_decode.k_scale",
+                          {block_num, n_kv_heads, block_size},
+                          torch::kFloat32,
+                          options_.device());
+  auto v_cache_scale_raw =
+      test::seeded_tensor("qwen2_quant_decode.v_scale",
+                          {block_num, n_kv_heads, block_size},
+                          torch::kFloat32,
+                          options_.device());
+  // Scale to [0.5, 1.5] range: 0.5 + raw * 1.0
+  auto k_cache_scale = 0.5f + k_cache_scale_raw;
+  auto v_cache_scale = 0.5f + v_cache_scale_raw;
+
+  KVCache quant_kv_cache = CreateQuantizedKvCache(std::move(k_cache),
+                                                  std::move(v_cache),
+                                                  std::move(k_cache_scale),
+                                                  std::move(v_cache_scale));
+
+  // Create input tensors using seeded tensors
+  auto hidden_states = test::seeded_tensor("qwen2_quant_decode.hidden_states",
+                                           {num_tokens, hidden_size},
+                                           torch::kBFloat16,
+                                           options_.device());
+  auto positions =
+      torch::full({num_tokens}, seq_len, options_.dtype(torch::kInt32));
+
+  auto metadata =
+      CreateAttentionMetadata(batch_size, seq_len + 1, false, seq_len + 1);
+
+  // Run forward with quantized KV cache
+  auto output =
+      qwen2_attention(positions, hidden_states, metadata, quant_kv_cache);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes(), torch::IntArrayRef({num_tokens, hidden_size}));
+
+  // Print output stats for debugging
+  torch::Tensor flat = output.flatten().to(torch::kFloat32).cpu();
+  double out_min = torch::min(flat).item<double>();
+  double out_max = torch::max(flat).item<double>();
+  double out_sum = torch::sum(flat).item<double>();
+  LOG(INFO) << "Quantized Decode output - min: " << out_min
+            << ", max: " << out_max << ", sum: " << out_sum;
+
+  // Verify precision using expect_tensor_stats
+  // Expected values established from successful diagnostic run
+  test::expect_tensor_stats(output,
+                            /*expected_min=*/-282,
+                            /*expected_max=*/67,
+                            /*expected_sum=*/-387352.6875,
+                            /*rtol=*/0.01,
+                            /*atol=*/1.0);
+}
+
+// Chunked prefill + quantized KV cache path uses flash attention and
+// dequant_from_paged_cache; parallel reduction in these kernels can be
+// non-deterministic on MLU, so fixed golden min/max/sum are not stable.
+// We validate shape, finite output, and determinism (two runs with same input
+// yield close results) instead of exact tensor stats.
+TEST_F(Qwen2AttentionTest, QuantizedKVCacheChunkedPrefillTest) {
+  auto qwen2_attention = Qwen2Attention(context_);
+  std::string prefix = "model.layers.0.self_attn.";
+  StateDict state_dict(weight_dict_, prefix);
+  qwen2_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  // Test parameters - first prefill a history chunk, then append a new chunk
+  // through the chunked prefill path.
+  int64_t batch_size = 2;
+  int64_t history_len = 32;
+  int64_t chunk_len = 32;
+  int64_t total_seq_len = history_len + chunk_len;
+  int64_t max_seq_len = total_seq_len;
+  int64_t hidden_size = model_args_.hidden_size();
+  int64_t num_tokens = batch_size * chunk_len;
+  int64_t block_size = 16;
+  // Calculate required blocks: each batch needs (max_seq_len/block_size) blocks
+  int64_t num_blocks_per_req = (max_seq_len + block_size - 1) / block_size + 1;
+  int64_t block_num =
+      batch_size * num_blocks_per_req + 10;  // Extra blocks for safety
+  int64_t n_kv_heads = model_args_.n_kv_heads().value();
+  int64_t head_dim = model_args_.head_dim();
+
+  // Create INT8 KV cache tensors using seeded tensors for reproducibility
+  auto k_cache =
+      test::seeded_tensor("qwen2_quant_chunked.k_cache",
+                          {block_num, n_kv_heads, block_size, head_dim},
+                          torch::kInt8,
+                          options_.device());
+  auto v_cache =
+      test::seeded_tensor("qwen2_quant_chunked.v_cache",
+                          {block_num, n_kv_heads, block_size, head_dim},
+                          torch::kInt8,
+                          options_.device());
+
+  // Create scale tensors with controlled range [0.5, 1.5]
+  auto k_cache_scale_raw =
+      test::seeded_tensor("qwen2_quant_chunked.k_scale",
+                          {block_num, n_kv_heads, block_size},
+                          torch::kFloat32,
+                          options_.device());
+  auto v_cache_scale_raw =
+      test::seeded_tensor("qwen2_quant_chunked.v_scale",
+                          {block_num, n_kv_heads, block_size},
+                          torch::kFloat32,
+                          options_.device());
+  auto k_cache_scale = 0.5f + k_cache_scale_raw;
+  auto v_cache_scale = 0.5f + v_cache_scale_raw;
+
+  KVCache quant_kv_cache = CreateQuantizedKvCache(std::move(k_cache),
+                                                  std::move(v_cache),
+                                                  std::move(k_cache_scale),
+                                                  std::move(v_cache_scale));
+
+  auto options_int = options_.dtype(torch::kInt32);
+  auto make_seq_offsets = [&](int64_t len) {
+    return torch::arange(0, (batch_size + 1) * len, len, options_int);
+  };
+  auto make_positions = [&](int64_t start, int64_t len) {
+    return torch::arange(start, start + len, options_int).repeat({batch_size});
+  };
+  auto make_block_table = [&]() {
+    std::vector<int32_t> block_table_vec;
+    block_table_vec.reserve(batch_size * num_blocks_per_req);
+    for (int64_t b = 0; b < batch_size; ++b) {
+      for (int64_t i = 0; i < num_blocks_per_req; ++i) {
+        block_table_vec.push_back(
+            static_cast<int32_t>(b * num_blocks_per_req + i));
+      }
+    }
+    return torch::tensor(block_table_vec, options_int)
+        .reshape({batch_size, num_blocks_per_req});
+  };
+
+  // Populate the history region first. Each request owns a full max_seq_len
+  // slice in the paged cache, so batch b uses slots
+  // [b * max_seq_len, b * max_seq_len + history_len).
+  auto history_hidden_states =
+      test::seeded_tensor("qwen2_quant_chunked.history_hidden_states",
+                          {batch_size * history_len, hidden_size},
+                          torch::kBFloat16,
+                          options_.device());
+  auto history_positions = make_positions(/*start=*/0, history_len);
+
+  AttentionMetadata history_metadata;
+  history_metadata.q_cu_seq_lens = make_seq_offsets(history_len);
+  history_metadata.kv_cu_seq_lens = history_metadata.q_cu_seq_lens;
+  std::vector<int32_t> history_slot_mapping_vec;
+  history_slot_mapping_vec.reserve(batch_size * history_len);
+  for (int64_t b = 0; b < batch_size; ++b) {
+    for (int64_t i = 0; i < history_len; ++i) {
+      history_slot_mapping_vec.push_back(
+          static_cast<int32_t>(b * max_seq_len + i));
+    }
+  }
+  history_metadata.slot_mapping =
+      torch::tensor(history_slot_mapping_vec, options_int);
+  history_metadata.kv_seq_lens =
+      torch::full({batch_size}, history_len, options_int);
+  history_metadata.block_table =
+      torch::zeros({batch_size, num_blocks_per_req}, options_int);
+  history_metadata.max_query_len = history_len;
+  history_metadata.max_seq_len = max_seq_len;
+  history_metadata.total_kv_len = batch_size * history_len;
+  history_metadata.compute_dtype = "half";
+  history_metadata.is_prefill = true;
+  history_metadata.is_chunked_prefill = false;
+  history_metadata.is_dummy = false;
+
+  auto history_output = qwen2_attention(history_positions,
+                                        history_hidden_states,
+                                        history_metadata,
+                                        quant_kv_cache);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+  ASSERT_EQ(history_output.sizes(),
+            torch::IntArrayRef({batch_size * history_len, hidden_size}));
+
+  auto hidden_states = test::seeded_tensor("qwen2_quant_chunked.hidden_states",
+                                           {num_tokens, hidden_size},
+                                           torch::kBFloat16,
+                                           options_.device());
+  auto positions = make_positions(/*start=*/history_len, chunk_len);
+
+  // Create metadata for the second chunk: reuse the history written above and
+  // append the new chunk at the tail of each request's max_seq_len slice.
+  AttentionMetadata metadata;
+  metadata.q_cu_seq_lens = make_seq_offsets(chunk_len);
+  metadata.kv_cu_seq_lens = make_seq_offsets(total_seq_len);
+  std::vector<int32_t> chunk_slot_mapping_vec;
+  chunk_slot_mapping_vec.reserve(batch_size * chunk_len);
+  for (int64_t b = 0; b < batch_size; ++b) {
+    for (int64_t i = 0; i < chunk_len; ++i) {
+      chunk_slot_mapping_vec.push_back(
+          static_cast<int32_t>(b * max_seq_len + history_len + i));
+    }
+  }
+  metadata.slot_mapping = torch::tensor(chunk_slot_mapping_vec, options_int);
+  metadata.kv_seq_lens = torch::full({batch_size}, total_seq_len, options_int);
+  metadata.block_table = make_block_table();
+  metadata.max_query_len = chunk_len;
+  metadata.max_seq_len = max_seq_len;
+  metadata.total_kv_len = batch_size * total_seq_len;
+  metadata.compute_dtype = "half";
+  metadata.is_prefill = false;
+  metadata.is_chunked_prefill = true;
+  metadata.is_dummy = false;
+
+  // First forward
+  auto output =
+      qwen2_attention(positions, hidden_states, metadata, quant_kv_cache);
+  device.synchronize_default_stream();
+
+  // Verify output shape
+  ASSERT_EQ(output.sizes(), torch::IntArrayRef({num_tokens, hidden_size}));
+
+  // Sanity: no NaN/Inf
+  torch::Tensor flat = output.flatten().to(torch::kFloat32).cpu();
+  ASSERT_TRUE(torch::isfinite(flat).all().item<bool>())
+      << "Output contains NaN or Inf";
+
+  // Second forward with same inputs (quant_kv_cache is overwritten with same
+  // K/V by quant_to_paged_cache). If the path were deterministic, both outputs
+  // would match; we use loose tolerance to allow backend non-determinism.
+  auto output2 =
+      qwen2_attention(positions, hidden_states, metadata, quant_kv_cache);
+  device.synchronize_default_stream();
+  ASSERT_TRUE(torch::allclose(output.flatten().to(torch::kFloat32),
+                              output2.flatten().to(torch::kFloat32),
+                              /*rtol=*/0.05,
+                              /*atol=*/0.05))
+      << "Two forwards with same input should produce close results "
+         "(determinism check)";
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/qwen2_vision_attention_test.cpp

@@ -0,0 +1,150 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/common/qwen2_vision_attention.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/mlu/tests_utils.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen2VisionAttentionTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    torch::Device device(Device::type_torch(), 0);
+    Device xllm_device(device);
+    xllm_device.set_seed(42);
+    model_args_.model_type() = "qwen2_vl";
+    model_args_.mm_hidden_size() = 1280;
+    model_args_.mm_head_dim() = 80;
+    model_args_.mm_num_attention_heads() = 16;
+
+    options_ = torch::TensorOptions().dtype(torch::kBFloat16).device(device);
+
+    process_group_ = create_process_group(
+        0, 1, 1, 3331, false, "localhost", "tp_group", device);
+    parallel_args_.tp_group_ = process_group_.get();
+
+    context_ = ModelContext(parallel_args_, model_args_, QuantArgs(), options_);
+    InitTestWeights();
+  }
+
+  void InitTestWeights() {
+    int32_t mm_hidden_size = model_args_.mm_hidden_size();
+    int32_t mm_num_heads = model_args_.mm_num_attention_heads();
+    int32_t mm_head_dim = model_args_.mm_head_dim();
+    int32_t qkv_size = mm_num_heads * mm_head_dim * 3;
+
+    std::unordered_map<std::string, torch::Tensor> weight_map = {
+        {"qkv.weight", torch::randn({qkv_size, mm_hidden_size}, options_)},
+        {"qkv.bias", torch::randn({qkv_size}, options_)},
+        {"proj.weight",
+         torch::randn({mm_hidden_size, mm_hidden_size}, options_)},
+        {"proj.bias", torch::randn({mm_hidden_size}, options_)},
+    };
+
+    for (auto& [name, tensor] : weight_map) {
+      tensor = tensor / torch::sqrt(torch::tensor(tensor.size(0), options_));
+      weight_dict_["model.visual." + name] = tensor;
+    }
+  }
+
+  ModelArgs model_args_;
+  ModelContext context_;
+  ParallelArgs parallel_args_{0, 1, nullptr};
+  torch::TensorOptions options_;
+  std::unordered_map<std::string, torch::Tensor> weight_dict_;
+  std::unique_ptr<ProcessGroup> process_group_ = nullptr;
+};
+
+TEST_F(Qwen2VisionAttentionTest, ForwardTest) {
+  auto vision_attention = Qwen2VisionAttention(context_);
+
+  std::string prefix = "model.visual.";
+  StateDict state_dict(weight_dict_, prefix);
+  vision_attention->load_state_dict(state_dict.get_dict_with_prefix(prefix));
+
+  int32_t batch_size = 2;
+  int32_t seq_len = 40;
+  int32_t mm_hidden_size = model_args_.mm_hidden_size();
+  int32_t num_tokens = batch_size * seq_len;
+
+  auto hidden_states =
+      test::seeded_tensor("qwen2_vision_attention.hidden_states",
+                          {num_tokens, mm_hidden_size},
+                          torch::kBFloat16,
+                          options_.device());
+
+  // Create cu_seq_len (cumulative sequence lengths)
+  std::vector<int32_t> cu_seq_len_vec = {0, seq_len, num_tokens};
+  auto cu_seq_len =
+      torch::tensor(cu_seq_len_vec, options_.dtype(torch::kInt32));
+
+  // Create rotary embeddings (cos and sin)
+  // Shape: (rope_seqlen, rope_dim)
+  int32_t mm_head_dim = model_args_.mm_head_dim();
+  int32_t rope_dim = mm_head_dim;
+  auto m_cos_pos = test::seeded_tensor("qwen2_vision_attention.m_cos_pos",
+                                       {num_tokens, rope_dim},
+                                       torch::kBFloat16,
+                                       options_.device());
+  auto m_sin_pos = test::seeded_tensor("qwen2_vision_attention.m_sin_pos",
+                                       {num_tokens, rope_dim},
+                                       torch::kBFloat16,
+                                       options_.device());
+
+  // Create ModelInputParams
+  ModelInputParams params;
+  auto output = vision_attention->forward(
+      hidden_states, m_cos_pos, m_sin_pos, cu_seq_len, cu_seq_len_vec, params);
+  xllm::Device device(options_.device());
+  device.synchronize_default_stream();
+
+  CHECK_EQ(output.sizes(), torch::IntArrayRef({num_tokens, mm_hidden_size}));
+
+  int32_t check_count = 10;
+  auto test_output = output.flatten().slice(0, 0, check_count);
+  std::vector<float> expected_values = {0.0703125f,
+                                        0.198242f,
+                                        0.0878906f,
+                                        -0.119141f,
+                                        0.142578f,
+                                        -0.410156f,
+                                        -0.233398f,
+                                        0.328125f,
+                                        -0.298828f,
+                                        0.0712891f};
+  test::verify_precision(test_output.unsqueeze(0), expected_values, 1e-4, 1e-5);
+
+  auto output2 = vision_attention->forward(
+      hidden_states, m_cos_pos, m_sin_pos, cu_seq_len, cu_seq_len_vec, params);
+  device.synchronize_default_stream();
+  ASSERT_TRUE(torch::allclose(output.flatten().to(torch::kFloat32),
+                              output2.flatten().to(torch::kFloat32),
+                              /*rtol=*/1e-4,
+                              /*atol=*/1e-5));
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/tests_utils.cpp

@@ -0,0 +1,348 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "layers/mlu/tests_utils.h"
+
+#include <cmath>
+
+#include "core/platform/device.h"
+
+namespace xllm {
+namespace layer {
+namespace test {
+
+// Supports both 2D and 3D input shapes
+torch::Tensor create_custom_input(const std::vector<int64_t>& shape,
+                                  const std::vector<float>& values,
+                                  const torch::TensorOptions& options) {
+  // Only support 2D or 3D
+  CHECK(shape.size() == 2 || shape.size() == 3) << "Shape must be 2D or 3D";
+  int64_t numel = 1;
+  for (auto d : shape) numel *= d;
+  CHECK_EQ(values.size(), numel) << "Values size must match tensor size";
+
+  // Create tensor from values directly with given shape
+  auto tensor =
+      torch::from_blob(const_cast<float*>(values.data()), shape, torch::kFloat)
+          .clone()
+          .to(options);
+  return tensor;
+}
+
+void verify_tensor_close(const torch::Tensor& actual,
+                         const torch::Tensor& expected,
+                         double rtol,
+                         double atol) {
+  ASSERT_TRUE(actual.sizes() == expected.sizes())
+      << "Tensor shapes don't match: actual=" << actual.sizes()
+      << ", expected=" << expected.sizes();
+
+  auto diff = torch::abs(actual - expected);
+  auto max_diff = torch::max(diff).item<float>();
+  auto mean_diff = torch::mean(diff).item<float>();
+
+  LOG(INFO) << "Max difference: " << max_diff;
+  LOG(INFO) << "Mean difference: " << mean_diff;
+
+  ASSERT_TRUE(torch::allclose(actual, expected, rtol, atol))
+      << "Tensors are not close enough. Max diff: " << max_diff;
+}
+
+void verify_precision(const torch::Tensor& actual_output,
+                      const std::vector<float>& expected_values,
+                      double rtol,
+                      double atol) {
+  ASSERT_FALSE(expected_values.empty())
+      << "Expected output not set. Call SetExpectedOutput() first.";
+
+  // Support both 2D and 3D outputs
+  std::vector<int64_t> output_shape(actual_output.sizes().begin(),
+                                    actual_output.sizes().end());
+  int64_t numel = actual_output.numel();
+
+  ASSERT_TRUE(output_shape.size() == 2 || output_shape.size() == 3)
+      << "Output tensor must be 2D or 3D";
+  ASSERT_EQ(expected_values.size(), numel)
+      << "Expected output size mismatch: expected " << expected_values.size()
+      << ", actual tensor numel " << numel;
+
+  // Create expected tensor from values and shape
+  auto expected_tensor = create_custom_input(
+      output_shape, expected_values, actual_output.options());
+
+  LOG(INFO) << "Verifying precision with rtol=" << rtol << ", atol=" << atol;
+  verify_tensor_close(actual_output, expected_tensor, rtol, atol);
+}
+
+void expect_tensor_stats(const torch::Tensor& t,
+                         double expected_min,
+                         double expected_max,
+                         double expected_sum,
+                         double rtol,
+                         double atol) {
+  EXPECT_TRUE(t.defined() && t.numel() > 0)
+      << "Tensor must be defined and non-empty";
+  torch::Tensor flat = t.flatten().to(torch::kFloat32);
+  if (t.device() != torch::kCPU) {
+    flat = flat.cpu();
+  }
+  double actual_min = torch::min(flat).item<double>();
+  double actual_max = torch::max(flat).item<double>();
+  double actual_sum = torch::sum(flat).item<double>();
+
+  auto within_tol = [rtol, atol](double actual, double expected) {
+    double tol = atol + rtol * std::abs(expected);
+    return std::abs(actual - expected) <= tol;
+  };
+
+  EXPECT_TRUE(within_tol(actual_min, expected_min))
+      << "min mismatch: actual=" << actual_min << ", expected=" << expected_min;
+  EXPECT_TRUE(within_tol(actual_max, expected_max))
+      << "max mismatch: actual=" << actual_max << ", expected=" << expected_max;
+  EXPECT_TRUE(within_tol(actual_sum, expected_sum))
+      << "sum mismatch: actual=" << actual_sum << ", expected=" << expected_sum;
+}
+
+ModelArgs create_default_model_args() {
+  ModelArgs model_args;
+  model_args.hidden_size() = 7168;
+  model_args.intermediate_size() = 18432;
+  model_args.hidden_act() = "silu";
+  return model_args;
+}
+
+QuantArgs create_default_quant_args() {
+  QuantArgs quant_args;
+  quant_args.quant_method() = kQuantMethodSmoothquant;
+  quant_args.bits() = 8;
+  quant_args.activation_dynamic() = true;
+  return quant_args;
+}
+
+ParallelArgs create_default_parallel_args(
+    std::unique_ptr<xllm::ProcessGroup>& mock_process_group) {
+  // Create mock ProcessGroup for MLU testing
+  mock_process_group = std::make_unique<MockProcessGroup>(
+      torch::Device(Device::type_torch(), 0));
+
+  // Initialize ParallelArgs with mock ProcessGroup
+  ParallelArgs parallel_args(0, 1, mock_process_group.get());
+
+  // Set tp_group_ for MLU environment
+  parallel_args.tp_group_ = mock_process_group.get();
+  parallel_args.single_rank_group_ = mock_process_group.get();
+  parallel_args.sp_group_ = mock_process_group.get();
+
+  return parallel_args;
+}
+
+// helper function for seeded_tensor
+// Calculate number of elements from shape
+inline int64_t numel_from_shape(torch::IntArrayRef shape) {
+  return std::accumulate(
+      shape.begin(), shape.end(), (int64_t)1, std::multiplies<int64_t>());
+}
+
+// helper function for seeded_tensor
+// FNV-1a 64-bit hash function
+inline uint64_t fnv1a64(const std::string& s) {
+  uint64_t h = 0xcbf29ce484222325ULL;
+  for (unsigned char c : s) {
+    h ^= c;
+    h *= 0x100000001b3ULL;
+  }
+  return h;
+}
+
+// helper struct for seeded_tensor
+// SplitMix64 pseudo-random number generator
+struct SplitMix64 {
+  uint64_t state;
+  explicit SplitMix64(uint64_t seed) : state(seed) {}
+  inline uint64_t next_u64() {
+    state += 0x9E3779B97F4A7C15ULL;
+    uint64_t z = state;
+    z ^= (z >> 30);
+    z *= 0xBF58476D1CE4E5B9ULL;
+    z ^= (z >> 27);
+    z *= 0x94D049BB133111EBULL;
+    z ^= (z >> 31);
+    return z;
+  }
+};
+
+// Generate tensor consistent with Python version
+torch::Tensor seeded_tensor(const std::string& key,
+                            torch::IntArrayRef shape,
+                            torch::ScalarType dtype,
+                            torch::Device device) {
+  const int64_t N = numel_from_shape(shape);
+  // Generate u64 stream
+  SplitMix64 rng(fnv1a64(key));
+  std::vector<uint64_t> buf;
+  buf.reserve(N);
+  for (int64_t i = 0; i < N; ++i) buf.push_back(rng.next_u64());
+
+  // Map and build CPU tensor according to dtype
+  torch::Tensor out_cpu;
+
+  if (torch::isFloatingType(dtype)) {
+    // Floating point: use high 53 bit -> [0,1)
+    std::vector<double> vals;
+    vals.reserve(N);
+    const double inv_2_53 = 1.0 / static_cast<double>(1ULL << 53);
+    for (uint64_t u : buf)
+      vals.push_back(static_cast<double>(u >> 11) * inv_2_53);
+    out_cpu =
+        torch::from_blob(
+            vals.data(), {N}, torch::TensorOptions().dtype(torch::kDouble))
+            .clone()
+            .to(dtype);
+  } else if (dtype == torch::kBool) {
+    std::vector<uint8_t> vals;
+    vals.reserve(N);
+    for (uint64_t u : buf) vals.push_back(static_cast<uint8_t>(u & 1ULL));
+    out_cpu = torch::from_blob(
+                  vals.data(), {N}, torch::TensorOptions().dtype(torch::kBool))
+                  .clone();
+  } else if (torch::isIntegralType(dtype, /*includeBool=*/false)) {
+    // Integer: min + (u % span), use __int128 to handle (cover int64)
+    auto map_mod_span = [&](auto tag) -> torch::Tensor {
+      using T = decltype(tag);
+      std::vector<T> vals;
+      vals.reserve(N);
+      const __int128 minv =
+          static_cast<__int128>(std::numeric_limits<T>::min());
+      const __int128 maxv =
+          static_cast<__int128>(std::numeric_limits<T>::max());
+      const unsigned __int128 span =
+          static_cast<unsigned __int128>(maxv - minv) + 1U;
+      for (uint64_t u : buf) {
+        T x = static_cast<T>(
+            minv +
+            static_cast<__int128>((static_cast<unsigned __int128>(u) % span)));
+        vals.push_back(x);
+      }
+      return torch::from_blob(
+                 vals.data(), {N}, torch::TensorOptions().dtype(dtype))
+          .clone();
+    };
+
+    // handle aliases in a single point for each type.
+    switch (dtype) {
+      case torch::kUInt8:  // alias for torch::kByte
+        out_cpu = map_mod_span(uint8_t{});
+        break;
+      case torch::kInt8:  // alias for torch::kChar
+        out_cpu = map_mod_span(int8_t{});
+        break;
+      case torch::kInt16:  // alias for torch::kShort
+        out_cpu = map_mod_span(int16_t{});
+        break;
+      case torch::kInt32:  // alias for torch::kInt
+        out_cpu = map_mod_span(int32_t{});
+        break;
+      case torch::kInt64:  // alias for torch::kLong
+        out_cpu = map_mod_span(int64_t{});
+        break;
+      default:
+        LOG(FATAL) << "Unsupported integer dtype: " << dtype;
+    }
+  } else {
+    LOG(FATAL) << "Unsupported dtype for seeded_tensor";
+  }
+
+  // Shape & device
+  out_cpu = out_cpu.view(shape).contiguous();
+  if (device != c10::Device(c10::kCPU)) {
+    return out_cpu.to(device);
+  }
+  return out_cpu;
+}
+
+void append_w4a8_expert_weights(
+    std::unordered_map<std::string, torch::Tensor>& weight_dict,
+    const std::string& expert_prefix,
+    const std::string& seed_prefix,
+    int64_t hidden_size,
+    int64_t gate_up_intermediate_size,
+    int64_t down_qweight_intermediate_size,
+    int64_t down_scale_intermediate_size,
+    int64_t group_size,
+    const torch::Device& device) {
+  CHECK_GT(group_size, 0);
+  CHECK_EQ(hidden_size % 2, 0);
+  CHECK_EQ(gate_up_intermediate_size % 2, 0);
+  CHECK_EQ(down_qweight_intermediate_size % 2, 0);
+  CHECK_EQ(hidden_size % group_size, 0);
+  CHECK_EQ(down_scale_intermediate_size % group_size, 0);
+
+  const int64_t hidden_packed = hidden_size / 2;
+  const int64_t down_packed = down_qweight_intermediate_size / 2;
+  const int64_t hidden_group_cols = hidden_size / group_size;
+  const int64_t down_group_cols = down_scale_intermediate_size / group_size;
+
+  auto gate_qweight = seeded_tensor(seed_prefix + ".gate_proj.qweight.w4",
+                                    {gate_up_intermediate_size, hidden_packed},
+                                    torch::kInt8,
+                                    device);
+  auto gate_scale =
+      seeded_tensor(seed_prefix + ".gate_proj.scale.w4",
+                    {gate_up_intermediate_size, hidden_group_cols},
+                    torch::kFloat32,
+                    device);
+  auto gate_smooth = seeded_tensor(seed_prefix + ".gate_proj.smooth",
+                                   {hidden_size},
+                                   torch::kFloat32,
+                                   device);
+
+  auto up_qweight = seeded_tensor(seed_prefix + ".up_proj.qweight.w4",
+                                  {gate_up_intermediate_size, hidden_packed},
+                                  torch::kInt8,
+                                  device);
+  auto up_scale = seeded_tensor(seed_prefix + ".up_proj.scale.w4",
+                                {gate_up_intermediate_size, hidden_group_cols},
+                                torch::kFloat32,
+                                device);
+
+  auto down_qweight = seeded_tensor(seed_prefix + ".down_proj.qweight.w4",
+                                    {hidden_size, down_packed},
+                                    torch::kInt8,
+                                    device);
+  auto down_scale = seeded_tensor(seed_prefix + ".down_proj.scale.w4",
+                                  {hidden_size, down_group_cols},
+                                  torch::kFloat32,
+                                  device);
+  auto down_smooth = seeded_tensor(seed_prefix + ".down_proj.smooth",
+                                   {gate_up_intermediate_size},
+                                   torch::kFloat32,
+                                   device);
+
+  weight_dict[expert_prefix + "gate_proj.qweight"] = gate_qweight;
+  weight_dict[expert_prefix + "gate_proj.per_channel_scale"] = gate_scale;
+  weight_dict[expert_prefix + "gate_proj.smooth"] = gate_smooth;
+
+  weight_dict[expert_prefix + "up_proj.qweight"] = up_qweight;
+  weight_dict[expert_prefix + "up_proj.per_channel_scale"] = up_scale;
+  weight_dict[expert_prefix + "up_proj.smooth"] = gate_smooth;
+
+  weight_dict[expert_prefix + "down_proj.qweight"] = down_qweight;
+  weight_dict[expert_prefix + "down_proj.per_channel_scale"] = down_scale;
+  weight_dict[expert_prefix + "down_proj.smooth"] = down_smooth;
+}
+
+}  // namespace test
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/layers/mlu/tests_utils.h

@@ -0,0 +1,287 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <numeric>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+
+namespace xllm {
+namespace layer {
+namespace test {
+
+class CompletedWork : public c10d::Work {
+ public:
+  CompletedWork() { finish(); }
+};
+
+inline c10::intrusive_ptr<c10d::Work> make_completed_work() {
+  return c10::make_intrusive<CompletedWork>();
+}
+
+// Mock Backend for testing - minimal implementation for tp=1 tests
+class MockBackend : public c10d::Backend {
+ public:
+  MockBackend(int64_t rank, int64_t world_size)
+      : c10d::Backend(rank, world_size), rank_(rank), world_size_(world_size) {}
+
+  c10::intrusive_ptr<c10d::Work> allreduce(
+      std::vector<torch::Tensor>& tensors,
+      const c10d::AllreduceOptions& opts = c10d::AllreduceOptions()) override {
+    // Mock implementation - return a completed work
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> allgather(
+      std::vector<std::vector<torch::Tensor>>& outputTensors,
+      std::vector<torch::Tensor>& inputTensors,
+      const c10d::AllgatherOptions& opts = c10d::AllgatherOptions()) override {
+    // Mock implementation - return a completed work
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> barrier(
+      const c10d::BarrierOptions& opts = c10d::BarrierOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> broadcast(
+      std::vector<torch::Tensor>& tensors,
+      const c10d::BroadcastOptions& opts = c10d::BroadcastOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> reduce(
+      std::vector<torch::Tensor>& tensors,
+      const c10d::ReduceOptions& opts = c10d::ReduceOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> allgather_coalesced(
+      std::vector<std::vector<torch::Tensor>>& outputTensorLists,
+      std::vector<torch::Tensor>& inputTensors,
+      const c10d::AllgatherOptions& opts = c10d::AllgatherOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> reduce_scatter(
+      std::vector<torch::Tensor>& outputTensors,
+      std::vector<std::vector<torch::Tensor>>& inputTensors,
+      const c10d::ReduceScatterOptions& opts =
+          c10d::ReduceScatterOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> alltoall_base(
+      torch::Tensor& outputTensor,
+      torch::Tensor& inputTensor,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const c10d::AllToAllOptions& opts = c10d::AllToAllOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> alltoall(
+      std::vector<torch::Tensor>& outputTensors,
+      std::vector<torch::Tensor>& inputTensors,
+      const c10d::AllToAllOptions& opts = c10d::AllToAllOptions()) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> send(std::vector<torch::Tensor>& tensors,
+                                      int dstRank,
+                                      int tag) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> recv(std::vector<torch::Tensor>& tensors,
+                                      int srcRank,
+                                      int tag) override {
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> recvAnysource(
+      std::vector<torch::Tensor>& tensors,
+      int tag) override {
+    return make_completed_work();
+  }
+
+  int64_t getRank() const { return rank_; }
+
+  int64_t getSize() const { return world_size_; }
+
+  void shutdown() override {
+    // Mock implementation - do nothing
+  }
+
+ private:
+  int64_t rank_;
+  int64_t world_size_;
+};
+
+// Mock ProcessGroup for testing
+class MockProcessGroup : public xllm::ProcessGroup {
+ public:
+  MockProcessGroup(const torch::Device& device,
+                   int64_t rank = 0,
+                   int64_t world_size = 1)
+      : xllm::ProcessGroup(rank, world_size, device) {
+    // Initialize pg_ with a mock backend for testing
+    pg_ = std::make_unique<MockBackend>(rank, world_size);
+  }
+
+  void allreduce(torch::Tensor& input) override {
+    // Mock implementation - do nothing for testing
+  }
+
+  c10::intrusive_ptr<c10d::Work> allreduce_async(
+      torch::Tensor& input) override {
+    allreduce(input);
+    return make_completed_work();
+  }
+
+  void allgather(const torch::Tensor& input,
+                 std::vector<torch::Tensor>& outputs) override {
+    outputs.resize(this->world_size());
+    if (!allgather_outputs_.empty()) {
+      CHECK_EQ(allgather_outputs_.size(), outputs.size())
+          << "mock allgather outputs size mismatch";
+      for (size_t i = 0; i < outputs.size(); ++i) {
+        outputs[i] = allgather_outputs_[i].clone();
+      }
+      return;
+    }
+
+    // Mock implementation - just copy input to outputs
+    for (size_t i = 0; i < this->world_size(); ++i) {
+      outputs[i] = input.clone();
+    }
+  }
+
+  c10::intrusive_ptr<c10d::Work> allgather_async(
+      const torch::Tensor& input,
+      std::vector<torch::Tensor>& outputs) override {
+    allgather(input, outputs);
+    return make_completed_work();
+  }
+
+  c10::intrusive_ptr<c10d::Work> allgather_base_async(
+      const torch::Tensor& input,
+      torch::Tensor& output) override {
+    CHECK(output.defined()) << "mock allgather_base_async requires output";
+    CHECK_EQ(output.size(0), this->world_size())
+        << "mock allgather_base_async world_size mismatch";
+
+    if (!allgather_outputs_.empty()) {
+      CHECK_EQ(allgather_outputs_.size(), static_cast<size_t>(output.size(0)))
+          << "mock allgather_base outputs size mismatch";
+      for (int64_t i = 0; i < output.size(0); ++i) {
+        output[i].copy_(allgather_outputs_[i]);
+      }
+      return make_completed_work();
+    }
+
+    for (int64_t i = 0; i < output.size(0); ++i) {
+      output[i].copy_(input);
+    }
+    return make_completed_work();
+  }
+
+  void reduce_scatter(const torch::Tensor& input,
+                      torch::Tensor& output) override {
+    int64_t world_size = this->world_size();
+    int64_t chunk_size = input.size(0) / world_size;
+    int64_t start = this->rank() * chunk_size;
+    output.copy_(input.slice(0, start, start + chunk_size));
+  }
+
+  void set_allgather_outputs(std::vector<torch::Tensor> outputs) {
+    allgather_outputs_ = std::move(outputs);
+  }
+
+ private:
+  std::vector<torch::Tensor> allgather_outputs_;
+};
+
+// Helper function to create custom input tensor for precision testing
+torch::Tensor create_custom_input(const std::vector<int64_t>& shape,
+                                  const std::vector<float>& values,
+                                  const torch::TensorOptions& options);
+
+// Helper function to verify tensor values are close to expected
+void verify_tensor_close(const torch::Tensor& actual,
+                         const torch::Tensor& expected,
+                         double rtol = 1e-5,
+                         double atol = 1e-8);
+
+// Helper function to verify precision against expected output
+void verify_precision(const torch::Tensor& actual_output,
+                      const std::vector<float>& expected_values,
+                      double rtol = 1e-3,
+                      double atol = 1e-4);
+
+// Expect tensor's min, max, sum (computed in fp32) to match expected values
+// within tolerance. Uses atol + rtol * |expected| for each of min, max, sum.
+void expect_tensor_stats(const torch::Tensor& t,
+                         double expected_min,
+                         double expected_max,
+                         double expected_sum,
+                         double rtol = 1e-2,
+                         double atol = 1e-5);
+
+// Helper function to create default model arguments for testing
+ModelArgs create_default_model_args();
+
+// Helper function to create default quantization arguments for testing
+QuantArgs create_default_quant_args();
+
+// Helper function to create default parallel arguments for testing
+ParallelArgs create_default_parallel_args(
+    std::unique_ptr<xllm::ProcessGroup>& mock_process_group);
+
+// create a tensor with a seeded random number generator (based on key and
+// shape) It is robust enough to generate the same tensor across any device or
+// os
+torch::Tensor seeded_tensor(const std::string& key,
+                            torch::IntArrayRef shape,
+                            torch::ScalarType dtype = torch::kFloat,
+                            torch::Device device = torch::Device(torch::kCPU));
+
+void append_w4a8_expert_weights(
+    std::unordered_map<std::string, torch::Tensor>& weight_dict,
+    const std::string& expert_prefix,
+    const std::string& seed_prefix,
+    int64_t hidden_size,
+    int64_t gate_up_intermediate_size,
+    int64_t down_qweight_intermediate_size,
+    int64_t down_scale_intermediate_size,
+    int64_t group_size,
+    const torch::Device& device);
+
+}  // namespace test
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm/tests/core/platform/CMakeLists.txt

@@ -0,0 +1,15 @@
+include(cc_test)
+
+if(USE_NPU OR USE_CUDA OR USE_ILU)
+  cc_test(
+    NAME
+      multi_platform_vmm_test
+    SRCS
+      shared_vmm_allocator_test.cpp
+    DEPS
+      :platform
+      :flags
+      GTest::gtest_main
+      glog::glog
+  )
+endif()

upstream_ref/xllm/tests/core/platform/shared_vmm_allocator_test.cpp

@@ -0,0 +1,102 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "core/platform/shared_vmm_allocator.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include "core/platform/device.h"
+
+#if defined(USE_NPU)
+#include <acl/acl.h>
+#include <torch_npu/torch_npu.h>
+#endif
+
+namespace {
+
+class SharedVMMAllocatorTestEnvironment : public ::testing::Environment {
+ public:
+  void SetUp() override {
+    google::InitGoogleLogging("platform_vmm_test");
+    google::SetStderrLogging(google::INFO);
+
+#if defined(USE_NPU)
+    int ret = aclrtSetDevice(0);
+    if (ret != 0) {
+      LOG(ERROR) << "ACL set device id: 0 failed, ret:" << ret;
+    }
+    torch_npu::init_npu("npu:0");
+#endif
+  }
+
+  void TearDown() override {
+#if defined(USE_NPU)
+    torch_npu::finalize_npu();
+    aclrtResetDevice(0);
+    aclFinalize();
+#endif
+    google::ShutdownGoogleLogging();
+  }
+};
+
+::testing::Environment* const test_env =
+    ::testing::AddGlobalTestEnvironment(new SharedVMMAllocatorTestEnvironment);
+
+bool HasDevice() { return xllm::Device::device_count() > 0; }
+
+void InitDevice() {
+  xllm::Device device(0);
+  device.set_device();
+  device.init_device_context();
+}
+
+}  // namespace
+
+TEST(SharedVMMAllocatorTest, BasicAllocateAndSwitch) {
+  if (!HasDevice()) {
+    GTEST_SKIP() << "No accelerator device available";
+  }
+
+  InitDevice();
+  xllm::SharedVMMAllocator allocator;
+  const size_t reserve_size = 64 * 1024 * 1024;
+  allocator.init(/*device_id=*/0, reserve_size);
+
+  EXPECT_TRUE(allocator.is_initialized());
+  EXPECT_GE(allocator.reserved_size(), reserve_size);
+  EXPECT_EQ(allocator.current_offset(), 0u);
+
+  void* first = allocator.allocate(1024 * 1024);
+  EXPECT_NE(first, nullptr);
+  allocator.deallocate(first);
+
+  const size_t offset_after_first = allocator.current_offset();
+  EXPECT_GT(offset_after_first, 0u);
+  EXPECT_GE(allocator.mapped_size(), offset_after_first);
+
+  void* second = allocator.allocate(1024 * 1024);
+  EXPECT_NE(second, nullptr);
+  EXPECT_NE(first, second);
+  EXPECT_GE(allocator.current_offset(), offset_after_first);
+  EXPECT_GE(allocator.high_water_mark(), allocator.current_offset());
+
+  allocator.switch_to_new_virtual_space();
+  EXPECT_EQ(allocator.current_offset(), 0u);
+
+  void* third = allocator.allocate(1024 * 1024);
+  EXPECT_NE(third, nullptr);
+  EXPECT_NE(third, first);
+}

upstream_ref/xllm/tests/core/runtime/CMakeLists.txt

@@ -0,0 +1,82 @@
+include(cc_test)
+
+if(USE_NPU)
+  cc_test(
+    NAME
+      acl_graph_executor_test
+    SRCS
+      acl_graph_executor_test.cpp
+    DEPS
+      :xllm_server
+      :runtime
+      :model_loader
+      :batch
+      :block
+      :model
+      :models
+      :sampler
+      :kv_cache
+      GTest::gtest_main
+      torch_npu
+      atb_customize
+  )
+  target_link_libraries(acl_graph_executor_test
+                        PRIVATE
+                        torch_npu
+                        ascendcl
+                        hccl
+                        c_sec
+                        nnopbase)
+  target_link_options(acl_graph_executor_test PRIVATE
+    "-Wl,--whole-archive"
+    "${CMAKE_BINARY_DIR}/third_party/spdlog/libspdlog.a"
+    "-Wl,--no-whole-archive")
+endif()
+
+cc_test(
+  NAME
+    spec_input_builder_test
+  SRCS
+    spec_input_builder_test.cpp
+    "${PROJECT_SOURCE_DIR}/xllm/core/runtime/spec_input_builder.cpp"
+  DEPS
+    torch
+    glog::glog
+    proto::xllm_proto
+    GTest::gtest_main
+)
+
+if(USE_MLU)
+  cc_test(
+    NAME
+      mlu_graph_executor_test
+    SRCS
+      mlu_graph_executor_test.cpp
+    DEPS
+      :runtime
+      :model_loader
+      :batch
+      :kv_cache
+      :platform
+      GTest::gtest_main
+      torch_mlu
+  )
+endif()
+
+if(USE_CUDA)
+  cc_test(
+    NAME
+      cuda_graph_executor_test
+    SRCS
+      cuda_graph_executor_test.cpp
+    DEPS
+      :runtime
+      :cuda_layers
+      :batch
+      :block
+      :kv_cache
+      :kernels
+      GTest::gtest_main
+      torch
+  )
+endif()

upstream_ref/xllm/tests/core/runtime/acl_graph_executor_test.cpp

@@ -0,0 +1,759 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 <acl/acl.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+#include <torch_npu/torch_npu.h>
+
+#include <cstdlib>
+#include <memory>
+#include <vector>
+
+#include "core/framework/batch/batch.h"
+#include "core/framework/block/block.h"
+#include "core/framework/block/block_manager_impl.h"
+#include "core/framework/kv_cache/kv_cache.h"
+#include "core/framework/model/model_args.h"
+#include "core/framework/model/model_output.h"
+#include "core/framework/model_loader.h"
+#include "core/framework/request/sequence.h"
+#include "core/framework/request/stopping_checker.h"
+#include "core/framework/sampling/sampling_params.h"
+#include "core/layers/npu/npu_lm_head_impl.h"
+#include "core/layers/npu/npu_word_embedding_impl.h"
+#include "core/runtime/acl_graph_executor_impl.h"
+#include "core/runtime/base_executor_impl.h"
+#include "core/runtime/options.h"
+
+// Global test environment for ACL graph executor tests
+class AclGraphExecutorTestEnvironment : public ::testing::Environment {
+ public:
+  void SetUp() override {
+    // Initialize glog
+    google::InitGoogleLogging("acl_graph_executor_test");
+    google::SetStderrLogging(google::INFO);
+
+    // Add any other global initialization here
+    std::cout << "Global test environment setup completed" << std::endl;
+    int ret = aclrtSetDevice(0);
+    if (ret != 0) {
+      LOG(ERROR) << "ACL set device id: 0 failed, ret:" << ret;
+    }
+    torch_npu::init_npu("npu:0");
+  }
+
+  void TearDown() override {
+    // Cleanup if needed
+    google::ShutdownGoogleLogging();
+    torch_npu::finalize_npu();
+    aclrtResetDevice(0);
+    aclFinalize();
+    LOG(INFO) << "AclGraphExecutorTestEnvironment TearDown completed.";
+  }
+};
+
+// Register the global test environment
+::testing::Environment* const test_env =
+    ::testing::AddGlobalTestEnvironment(new AclGraphExecutorTestEnvironment);
+
+namespace xllm {
+
+namespace {
+const KVCache& first_full_attention_cache(
+    const std::vector<KVCache>& kv_caches) {
+  for (const auto& kv_cache : kv_caches) {
+    if (!kv_cache.empty()) {
+      auto k_cache = kv_cache.get_k_cache();
+      if (k_cache.defined() && k_cache.numel() > 0) {
+        return kv_cache;
+      }
+    }
+  }
+  LOG(FATAL) << "No full-attention KV cache found";
+  std::abort();
+}
+}  // namespace
+
+// Initialize glog for testing - use a function to ensure proper initialization
+// order
+void InitializeGlog() {
+  static bool initialized = false;
+  if (!initialized) {
+    google::InitGoogleLogging("acl_graph_executor_test");
+    google::SetStderrLogging(google::INFO);
+    initialized = true;
+  }
+}
+
+// Simple CausalLM implementation for testing ACL graph executor
+// Uses basic operations to verify graph capture and replay functionality
+class SimpleCausalLM : public CausalLM {
+ public:
+  SimpleCausalLM(const ModelArgs& args, const torch::Device& device)
+      : args_(args), device_(device) {
+    // Initialize a simple linear layer for testing
+    linear_ = register_module("linear",
+                              torch::nn::Linear(torch::nn::LinearOptions(
+                                  args.hidden_size(), args.hidden_size())));
+
+    // Initialize token embedding table
+    const int64_t vocab_size = std::max(args.vocab_size(), 1000L);
+    token_embedding_table_ = register_parameter(
+        "token_embedding",
+        torch::randn({vocab_size, args.hidden_size()},
+                     torch::dtype(torch::kFloat32).device(device)));
+
+    // Initialize position embedding table
+    const int64_t max_pos = args.max_position_embeddings();
+    pos_embedding_table_ = register_parameter(
+        "pos_embedding",
+        torch::randn({max_pos, args.hidden_size()},
+                     torch::dtype(torch::kFloat32).device(device)));
+    // Initialize block-related tensors for Rec multi-round computation
+    block_size_ = torch::tensor(4L, torch::dtype(torch::kInt64).device(device));
+    scalar_one_ = torch::tensor(1L, torch::dtype(torch::kInt64).device(device));
+
+    // Initialize scalar tensors for computation
+    // const tensors
+    kv_scale_ =
+        torch::tensor(0.01f, torch::dtype(torch::kFloat32).device(device));
+    q_scale_ =
+        torch::tensor(0.01f, torch::dtype(torch::kFloat32).device(device));
+    cache_scale_ =
+        torch::tensor(0.005f, torch::dtype(torch::kFloat32).device(device));
+    block_scale_ =
+        torch::tensor(0.001f, torch::dtype(torch::kFloat32).device(device));
+
+    // Move to device
+    this->to(device);
+  }
+
+  torch::Tensor forward_impl(const torch::Tensor& tokens,
+                             const torch::Tensor& positions,
+                             std::vector<KVCache>& kv_caches,
+                             const ModelInputParams& params) {
+    // Simple computation: token embedding + position embedding + linear layer
+    // This creates temporary tensors that NPUGraph mempool will manage
+    LOG(INFO) << "SimpleCausalLM forward_impl, tokens: " << tokens.sizes()
+              << ", positions: " << positions.sizes()
+              << ", kv_caches: " << kv_caches.size()
+              << ", params: " << params.num_sequences;
+    const int64_t num_tokens = tokens.size(0);
+    const int64_t hidden_size = args_.hidden_size();
+
+    // Create token embeddings using standard embedding lookup
+    auto token_embeddings = torch::embedding(token_embedding_table_, tokens);
+
+    // Create position embeddings using standard embedding lookup
+    auto position_embeddings =
+        torch::embedding(pos_embedding_table_, positions);
+
+    // Combine embeddings
+    auto combined = token_embeddings + position_embeddings;
+
+    // Apply linear layer
+    auto output = linear_->forward(combined);
+
+    // Add some computation using other params to make it more realistic
+    // if (params.kv_seq_lens.defined()) {
+    //   // Use kv_seq_lens in computation
+    //   auto kv_lens_sum = torch::sum(params.kv_seq_lens);
+    //   output = output + kv_lens_sum * kv_scale_;
+    // }
+
+    // if (params.q_seq_lens.defined()) {
+    //   // Use q_seq_lens in computation
+    //   auto q_lens_sum = torch::sum(params.q_seq_lens);
+    //   output = output + q_lens_sum * q_scale_;
+    // }
+
+    if (params.new_cache_slots.defined()) {
+      // Use new_cache_slots in computation
+      auto cache_slots_sum = torch::sum(params.new_cache_slots);
+      output = output + cache_slots_sum * cache_scale_;
+    }
+
+    if (params.block_tables.defined() && !kv_caches.empty()) {
+      // Use block_tables to do embedding lookup from kv_cache - Rec multi-round
+      // computation Calculate max_seq_len from actual seq_len tensor
+      auto max_seq_len = torch::max(params.kv_seq_lens);
+
+      // Calculate max_block_nums_per_seq
+      auto max_block_nums_per_seq = torch::ceil(max_seq_len / block_size_);
+
+      // Get kv_cache tensor from KVCache
+      const auto& kv_cache_tensor =
+          first_full_attention_cache(kv_caches).get_k_cache();
+
+      // Create col_indices and mask
+      int64_t block_table_len = params.block_tables.size(1);
+      auto col_indices = torch::arange(
+          block_table_len, torch::dtype(torch::kInt64).device(device_));
+      auto mask = col_indices < (max_block_nums_per_seq - scalar_one_);
+
+      // Directly compute embedding
+      auto kv_embeddings =
+          torch::embedding(kv_cache_tensor, params.block_tables);
+
+      // Apply mask and sum
+      auto kv_embeddings_masked = kv_embeddings * mask.view({1, -1, 1});
+      auto kv_embeddings_sum = torch::sum(kv_embeddings_masked);
+      output = output + kv_embeddings_sum * block_scale_;
+    }
+
+    return output;
+  }
+
+  // Adapter method to match CausalLM base class interface
+  ModelOutput forward(const torch::Tensor& tokens,
+                      const torch::Tensor& positions,
+                      std::vector<KVCache>& kv_caches,
+                      const ModelInputParams& parameters) override {
+    auto hidden_states = forward_impl(tokens, positions, kv_caches, parameters);
+    return ModelOutput(hidden_states);
+  }
+
+  const torch::TensorOptions& options() const override {
+    static torch::TensorOptions opts =
+        torch::dtype(torch::kFloat32).device(device_);
+    return opts;
+  }
+
+  const ModelArgs& args() const { return args_; }
+
+  // Implement required virtual functions
+  torch::Tensor logits(const torch::Tensor& hidden_states,
+                       const torch::Tensor& selected_idxes) override {
+    // Simple logits computation
+    const int64_t vocab_size = std::max(args_.vocab_size(), 1000L);
+    return torch::randn({hidden_states.size(0), vocab_size},
+                        torch::dtype(torch::kFloat32).device(device_));
+  }
+
+  void load_model(std::unique_ptr<ModelLoader> loader) override {
+    // Simple implementation for testing
+  }
+
+  torch::Device device() const override { return device_; }
+
+  void prepare_expert_weight(int32_t layer_id,
+                             const std::vector<int32_t>& expert_ids) override {
+    // Simple implementation for testing
+  }
+
+  void update_expert_weight(int32_t layer_id) override {
+    // Simple implementation for testing
+  }
+
+  layer::NpuLmHead get_npu_lm_head() override {
+    // Simple implementation for testing
+    return layer::NpuLmHead(nullptr);
+  }
+
+  void set_npu_lm_head(layer::NpuLmHead& head) override {
+    // Simple implementation for testing
+  }
+
+  layer::NpuWordEmbedding get_npu_word_embedding() override {
+    // Simple implementation for testing
+    return layer::NpuWordEmbedding(nullptr);
+  }
+
+  void set_npu_word_embedding(layer::NpuWordEmbedding& embedding) override {
+    // Simple implementation for testing
+  }
+
+ private:
+  ModelArgs args_;
+  torch::Device device_;
+  torch::nn::Linear linear_{nullptr};
+  torch::Tensor token_embedding_table_;
+  torch::Tensor pos_embedding_table_;
+
+  // Pre-allocated constant scalar tensors for computation
+  torch::Tensor kv_scale_;
+  torch::Tensor q_scale_;
+  torch::Tensor cache_scale_;
+  torch::Tensor block_scale_;
+  torch::Tensor block_size_;
+  torch::Tensor scalar_one_;
+};
+
+class AclGraphExecutorTest : public ::testing::Test {
+ protected:
+  AclGraphExecutorTest() = default;
+
+  void SetUp() override {
+    if (initialized_) {
+      return;
+    }
+    initialized_ = true;
+    sequences_.reserve(100);
+
+    // Set up model args
+    model_args_.model_type("test_model");
+    model_args_.dtype("float32");
+    model_args_.hidden_size(128);
+    model_args_.max_position_embeddings(2048);
+    model_args_.vocab_size(1000);  // Set a reasonable vocab size
+
+    // Set up device
+    device_ = std::make_unique<torch::Device>("npu:0");
+
+    // Set up runtime options
+    options_.num_decoding_tokens(1);
+    options_.block_size(4);
+
+    // Create simple model
+    model_ = std::make_unique<SimpleCausalLM>(model_args_, *device_);
+
+    // Initialize block manager
+    const uint32_t n_blocks = 1000;
+    const uint32_t block_size = 4;
+    BlockManager::Options block_options;
+    block_options.num_blocks(n_blocks).block_size(block_size);
+    block_manager_ = std::make_unique<BlockManagerImpl>(block_options);
+
+    // Initialize sampling and stopping parameters
+    sampling_param_.frequency_penalty = 0.1;
+    stopping_checker_.set_max_generated_tokens(20);
+
+    // Initialize sequence parameters
+    seq_params_.seq_capacity = 100;
+    seq_params_.stopping_checker = &stopping_checker_;
+    seq_params_.sampling_param = &sampling_param_;
+    seq_params_.skip_special_tokens = true;
+    seq_params_.echo = false;
+    seq_params_.logprobs = false;
+    seq_params_.enable_schedule_overlap = false;
+
+    // Initialize input embedding and mm_data
+    input_embedding_ =
+        torch::zeros({1, model_args_.hidden_size()},
+                     torch::dtype(torch::kFloat32).device(*device_));
+    mm_data_ = MMData();  // Default constructor creates empty MMData
+
+    // Initialize KV caches
+    kv_caches_.clear();
+    const int64_t hidden_size = model_args_.hidden_size();
+
+    // Create KV cache with shape [n_blocks, block_size, hidden_size]
+    torch::Tensor kv_cache =
+        torch::randn({n_blocks, block_size * hidden_size},
+                     torch::dtype(torch::kFloat32).device(*device_));
+    kv_caches_.emplace_back(KVCacheTensors{kv_cache, kv_cache});
+  }
+
+  void TearDown() override { return; }
+
+  void reset() {
+    for (auto& sequence : sequences_) {
+      auto blocks = sequence.kv_state().kv_blocks();
+      if (!blocks.empty()) {
+        block_manager_->deallocate(blocks);
+      }
+    }
+  }
+
+  // Helper function to create a simple batch
+  std::unique_ptr<Batch> CreateTestBatch() {
+    sequences_.emplace_back(0,
+                            std::vector<int32_t>{1, 3, 5, 7, 5, 4, 3, 2, 1},
+                            input_embedding_,
+                            mm_data_,
+                            fake_decoder_,
+                            seq_params_);
+    auto& sequence = sequences_.back();
+
+    // Allocate blocks and configure sequence
+    sequence.add_kv_blocks(block_manager_->allocate(3));
+    // Set kv_cache_tokens_num to be >= num_prompt_tokens to move to decode
+    // stage
+    sequence.kv_state().incr_kv_cache_tokens_num(
+        /*size=*/9);  // 9 prompt tokens
+    sequence.append_token(100);
+
+    // Create batch with pointer to sequence (batch doesn't own sequence)
+    auto batch = std::make_unique<Batch>();
+    batch->add(&sequence);
+
+    return batch;
+  }
+  bool initialized_ = false;
+  ModelArgs model_args_;
+  std::unique_ptr<torch::Device> device_;
+  runtime::Options options_;
+  std::unique_ptr<CausalLM> model_;
+
+  // Shared resources for all tests
+  std::unique_ptr<BlockManagerImpl> block_manager_;
+  RequestSamplingParam sampling_param_;
+  StoppingChecker stopping_checker_;
+  SequenceParams seq_params_;
+  torch::Tensor input_embedding_;
+  MMData mm_data_;
+  std::vector<KVCache> kv_caches_;
+
+  // Sequences managed by test class
+  std::vector<Sequence> sequences_;
+
+  // Create a sequence in decode phase
+  IncrementalDecoder fake_decoder_ = IncrementalDecoder("", 1, false, false);
+};
+
+// Test that ACL graph executor produces same results as eager execution
+TEST_F(AclGraphExecutorTest, GraphExecutorVsEagerExecution) {
+  // Create test batch
+  auto batch = CreateTestBatch();
+  ASSERT_FALSE(batch->empty());
+
+  // Prepare forward input
+  auto forward_input = batch->prepare_forward_input(
+      options_.num_decoding_tokens(), 0, model_args_);
+  forward_input = forward_input.to(*device_, torch::kFloat32);
+
+  std::cout << "forward_input.token_ids: " << forward_input.token_ids
+            << std::endl;
+  std::cout << "forward_input.positions: " << forward_input.positions
+            << std::endl;
+  std::cout << "forward_input.input_params.q_seq_lens: "
+            << forward_input.input_params.q_seq_lens << std::endl;
+  std::cout << "forward_input.input_params.kv_seq_lens: "
+            << forward_input.input_params.kv_seq_lens << std::endl;
+  std::cout << "forward_input.input_params.new_cache_slots: "
+            << forward_input.input_params.new_cache_slots << std::endl;
+  std::cout << "forward_input.input_params.block_tables: "
+            << forward_input.input_params.block_tables << std::endl;
+  // Test eager execution (direct model forward)
+  auto eager_model_output = model_->forward({forward_input.token_ids},
+                                            {forward_input.positions},
+                                            kv_caches_,
+                                            {forward_input.input_params});
+  auto eager_output = eager_model_output.hidden_states;
+  // Create ACL graph executor
+  auto graph_executor = std::make_unique<::xllm::npu::AclGraphExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+
+  // Test graph execution with NPUGraph mempool optimization
+  auto graph_model_output = graph_executor->run({forward_input.token_ids},
+                                                {forward_input.positions},
+                                                kv_caches_,
+                                                {forward_input.input_params});
+  auto graph_output = graph_model_output.hidden_states;
+  // Compare outputs - should be identical
+  EXPECT_TRUE(
+      torch::allclose(eager_output, graph_output, /*rtol=*/1e-5, /*atol=*/1e-6))
+      << "Eager output:\n"
+      << eager_output << "\nGraph output:\n"
+      << graph_output;
+}
+
+// Test that graph replay produces consistent results across multiple runs
+TEST_F(AclGraphExecutorTest, GraphReplayConsistency) {
+  // Create test batch
+  auto batch = CreateTestBatch();
+  ASSERT_FALSE(batch->empty());
+
+  // Prepare forward input
+  auto forward_input = batch->prepare_forward_input(
+      options_.num_decoding_tokens(), 0, model_args_);
+  forward_input = forward_input.to(*device_, torch::kFloat32);
+
+  // Create ACL graph executor
+  auto graph_executor = std::make_unique<::xllm::npu::AclGraphExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+
+  // First execution (should create graph with NPUGraph mempool)
+  auto output1 = graph_executor->run({forward_input.token_ids},
+                                     {forward_input.positions},
+                                     kv_caches_,
+                                     {forward_input.input_params});
+
+  // Second execution (should replay graph using mempool-managed tensors)
+  auto output2 = graph_executor->run({forward_input.token_ids},
+                                     {forward_input.positions},
+                                     kv_caches_,
+                                     {forward_input.input_params});
+
+  // Compare outputs - should be identical
+  EXPECT_TRUE(torch::allclose(output1.hidden_states,
+                              output2.hidden_states,
+                              /*rtol=*/1e-5,
+                              /*atol=*/1e-6))
+      << "First output:\n"
+      << output1.hidden_states << "\nSecond output:\n"
+      << output2.hidden_states;
+}
+
+// Test graph creation and execution with different batch sizes
+TEST_F(AclGraphExecutorTest, DifferentBatchSizes) {
+  // Test with different batch sizes to ensure graph creation works
+  const std::vector<uint32_t> batch_sizes = {1, 2, 4};
+
+  for (auto batch_size : batch_sizes) {
+    // Clear sequences from previous iteration to avoid block exhaustion
+    sequences_.clear();
+
+    // Create multiple sequences for larger batch sizes
+    auto batch = std::make_unique<Batch>();
+
+    for (uint32_t i = 0; i < batch_size; ++i) {
+      sequences_.emplace_back(i,
+                              std::vector<int32_t>{static_cast<int32_t>(1 + i),
+                                                   static_cast<int32_t>(3 + i),
+                                                   static_cast<int32_t>(5 + i),
+                                                   static_cast<int32_t>(7 + i)},
+                              input_embedding_,
+                              mm_data_,
+                              fake_decoder_,
+                              seq_params_);
+      auto& sequence = sequences_.back();
+      sequence.add_kv_blocks(block_manager_->allocate(2));
+      std::cout << "batch_size: " << batch_size << " i: " << i
+                << " sequence.kv_state().current_max_tokens_capacity(): "
+                << sequence.kv_state().current_max_tokens_capacity()
+                << std::endl;
+      // Set kv_cache_tokens_num to be >= num_prompt_tokens to move to decode
+      // stage
+      sequence.kv_state().incr_kv_cache_tokens_num(
+          /*size=*/4);  // 4 prompt tokens
+
+      sequence.append_token(100 + i);
+      // Add sequence pointer to batch (batch doesn't own sequence)
+      batch->add(&sequence);
+    }
+
+    // Prepare forward input
+    auto forward_input = batch->prepare_forward_input(
+        options_.num_decoding_tokens(), 0, model_args_);
+    forward_input = forward_input.to(*device_, torch::kFloat32);
+    // Create ACL graph executor
+    auto graph_executor = new ::xllm::npu::AclGraphExecutorImpl(
+        model_.get(), model_args_, *device_, options_);
+
+    // Test graph execution
+    auto output = graph_executor->run({forward_input.token_ids},
+                                      {forward_input.positions},
+                                      kv_caches_,
+                                      {forward_input.input_params});
+
+    // Verify output shape
+    EXPECT_EQ(output.hidden_states.size(0),
+              batch_size * options_.num_decoding_tokens())
+        << "Batch size: " << batch_size;
+    EXPECT_EQ(output.hidden_states.size(1), model_args_.hidden_size())
+        << "Batch size: " << batch_size;
+  }
+}
+
+// Test ACL graph executor against original NPU executor implementation
+TEST_F(AclGraphExecutorTest, AclGraphExecutorVsBaseExecutorImpl) {
+  // Create test batch
+  auto batch = CreateTestBatch();
+  ASSERT_FALSE(batch->empty());
+
+  // Prepare forward input
+  auto forward_input = batch->prepare_forward_input(
+      options_.num_decoding_tokens(), 0, model_args_);
+  forward_input = forward_input.to(*device_, torch::kFloat32);
+  // Test NPU Executor Impl (original implementation)
+  auto npu_executor = std::make_unique<BaseExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+
+  auto npu_model_output = npu_executor->run({forward_input.token_ids},
+                                            {forward_input.positions},
+                                            kv_caches_,
+                                            {forward_input.input_params});
+  auto npu_output = npu_model_output.hidden_states;
+
+  // Test ACL Graph Executor with NPUGraph mempool optimization
+  auto graph_executor = std::make_unique<::xllm::npu::AclGraphExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+
+  auto graph_model_output = graph_executor->run({forward_input.token_ids},
+                                                {forward_input.positions},
+                                                kv_caches_,
+                                                {forward_input.input_params});
+  auto graph_output = graph_model_output.hidden_states;
+
+  // Compare outputs - should be identical
+  EXPECT_TRUE(
+      torch::allclose(npu_output, graph_output, /*rtol=*/1e-5, /*atol=*/1e-6))
+      << "NPU Executor output:\n"
+      << npu_output << "\nACL Graph Executor output:\n"
+      << graph_output;
+
+  // Verify output shapes are the same
+  EXPECT_EQ(npu_output.sizes(), graph_output.sizes())
+      << "Output shape mismatch: NPU=" << npu_output.sizes()
+      << ", Graph=" << graph_output.sizes();
+}
+
+// Test multiple runs to verify consistency across different execution modes
+TEST_F(AclGraphExecutorTest, AclGraphExecutorVsBaseExecutorImplMultipleRuns) {
+  // Create test batch
+  auto batch = CreateTestBatch();
+  ASSERT_FALSE(batch->empty());
+
+  // Prepare forward input
+  auto forward_input = batch->prepare_forward_input(
+      options_.num_decoding_tokens(), 0, model_args_);
+  forward_input = forward_input.to(*device_, torch::kFloat32);
+  // Create both executors
+  auto npu_executor = std::make_unique<BaseExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+  auto graph_executor = std::make_unique<::xllm::npu::AclGraphExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+
+  // Run multiple times and compare results
+  const int num_runs = 3;
+  for (int i = 0; i < num_runs; ++i) {
+    // Direct model forward call (baseline)
+    auto direct_model_output = model_->forward({forward_input.token_ids},
+                                               {forward_input.positions},
+                                               kv_caches_,
+                                               {forward_input.input_params});
+    auto direct_output = direct_model_output.hidden_states;
+
+    // NPU Executor run
+    auto npu_model_output = npu_executor->run({forward_input.token_ids},
+                                              {forward_input.positions},
+                                              kv_caches_,
+                                              {forward_input.input_params});
+    auto npu_output = npu_model_output.hidden_states;
+
+    // ACL Graph Executor run with NPUGraph mempool
+    auto graph_model_output = graph_executor->run({forward_input.token_ids},
+                                                  {forward_input.positions},
+                                                  kv_caches_,
+                                                  {forward_input.input_params});
+    auto graph_output = graph_model_output.hidden_states;
+
+    // Compare direct model output with NPU Executor output
+    EXPECT_TRUE(torch::allclose(
+        direct_output, npu_output, /*rtol=*/1e-5, /*atol=*/1e-6))
+        << "Run " << i << " - Direct model vs NPU Executor mismatch:\n"
+        << "Direct model output:\n"
+        << direct_output << "\nNPU Executor output:\n"
+        << npu_output;
+
+    // Compare direct model output with ACL Graph Executor output
+    EXPECT_TRUE(torch::allclose(
+        direct_output, graph_output, /*rtol=*/1e-5, /*atol=*/1e-6))
+        << "Run " << i << " - Direct model vs ACL Graph Executor mismatch:\n"
+        << "Direct model output:\n"
+        << direct_output << "\nACL Graph Executor output:\n"
+        << graph_output;
+
+    // Compare NPU Executor output with ACL Graph Executor output
+    EXPECT_TRUE(
+        torch::allclose(npu_output, graph_output, /*rtol=*/1e-5, /*atol=*/1e-6))
+        << "Run " << i << " - NPU Executor vs ACL Graph Executor mismatch:\n"
+        << "NPU Executor output:\n"
+        << npu_output << "\nACL Graph Executor output:\n"
+        << graph_output;
+  }
+}
+
+TEST_F(AclGraphExecutorTest, BatchInputCarriesLinearStateIds) {
+  model_args_.layer_types({"linear_attention", "full_attention"});
+  auto batch = CreateTestBatch();
+  ASSERT_FALSE(batch->empty());
+  ASSERT_FALSE(sequences_.empty());
+
+  auto& seq = sequences_.back();
+  auto linear_state_block = block_manager_->allocate(1);
+  ASSERT_EQ(linear_state_block.size(), 1);
+  const int32_t expected_linear_state_id = linear_state_block[0].id();
+  seq.set_single_block(std::move(linear_state_block[0]));
+
+  auto forward_input = batch->prepare_forward_input(
+      options_.num_decoding_tokens(), 0, model_args_);
+  ASSERT_EQ(forward_input.input_params.num_sequences, 1);
+  ASSERT_EQ(forward_input.input_params.linear_state_ids.size(), 1);
+  EXPECT_EQ(forward_input.input_params.linear_state_ids[0],
+            expected_linear_state_id);
+  ASSERT_EQ(forward_input.input_params.embedding_ids.size(), 1);
+  EXPECT_EQ(forward_input.input_params.embedding_ids[0],
+            expected_linear_state_id);
+}
+
+TEST(AclGraphExecutorHybridTest, KvCacheSupportsLinearOnlyLayers) {
+  auto conv_cache = torch::zeros({4, 32, 3}, torch::dtype(torch::kFloat32));
+  auto ssm_cache = torch::zeros({4, 8, 64, 64}, torch::dtype(torch::kFloat32));
+  KVCache linear_only_cache(
+      LinearAttentionKVCacheTensors{conv_cache, ssm_cache});
+
+  EXPECT_FALSE(linear_only_cache.empty());
+  EXPECT_FALSE(linear_only_cache.get_conv_cache().defined() == false);
+  EXPECT_FALSE(linear_only_cache.get_ssm_cache().defined() == false);
+  EXPECT_FALSE(linear_only_cache.get_k_cache().defined());
+  EXPECT_FALSE(linear_only_cache.get_v_cache().defined());
+}
+
+TEST(AclGraphExecutorHybridTest, ModelArgsCountsHybridLayerTypes) {
+  ModelArgs args;
+  args.n_layers(4);
+  args.layer_types(
+      {"linear_attention", "full_attention", "linear_attention", "attention"});
+
+  EXPECT_FALSE(is_full_attention_layer(args, 0));
+  EXPECT_TRUE(is_full_attention_layer(args, 1));
+  EXPECT_FALSE(is_full_attention_layer(args, 2));
+  EXPECT_TRUE(is_full_attention_layer(args, 3));
+  EXPECT_TRUE(has_linear_attention_layers(args));
+}
+
+TEST_F(AclGraphExecutorTest, GraphExecutorUsesFirstFullAttentionKvCache) {
+  auto batch = CreateTestBatch();
+  ASSERT_FALSE(batch->empty());
+
+  auto forward_input = batch->prepare_forward_input(
+      options_.num_decoding_tokens(), 0, model_args_);
+  forward_input = forward_input.to(*device_, torch::kFloat32);
+
+  auto conv_cache =
+      torch::zeros({4, 32, 3}, torch::dtype(torch::kFloat32).device(*device_));
+  auto ssm_cache = torch::zeros({4, 8, 64, 64},
+                                torch::dtype(torch::kFloat32).device(*device_));
+  auto full_k = torch::randn({1000, 4 * model_args_.hidden_size()},
+                             torch::dtype(torch::kFloat32).device(*device_));
+  auto full_v = full_k.clone();
+
+  std::vector<KVCache> hybrid_kv_caches;
+  hybrid_kv_caches.emplace_back(
+      LinearAttentionKVCacheTensors{conv_cache, ssm_cache});
+  hybrid_kv_caches.emplace_back(KVCacheTensors{full_k, full_v});
+
+  auto eager_model_output = model_->forward({forward_input.token_ids},
+                                            {forward_input.positions},
+                                            hybrid_kv_caches,
+                                            {forward_input.input_params});
+  auto graph_executor = std::make_unique<::xllm::npu::AclGraphExecutorImpl>(
+      model_.get(), model_args_, *device_, options_);
+  auto graph_model_output = graph_executor->run({forward_input.token_ids},
+                                                {forward_input.positions},
+                                                hybrid_kv_caches,
+                                                {forward_input.input_params});
+
+  EXPECT_TRUE(torch::allclose(eager_model_output.hidden_states,
+                              graph_model_output.hidden_states,
+                              /*rtol=*/1e-5,
+                              /*atol=*/1e-6));
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/runtime/mlu_graph_executor_test.cpp

@@ -0,0 +1,595 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 <framework/core/device.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <vector>
+
+#include "base_executor_impl.h"
+#include "core/framework/batch/batch.h"
+#include "core/framework/kv_cache/kv_cache.h"
+#include "core/framework/model/model_args.h"
+#include "core/framework/model/model_output.h"
+#include "mlu_graph_executor_impl.h"
+#include "platform/device.h"
+#include "runtime/options.h"
+
+namespace xllm {
+class MockCausalLM : public CausalLM {
+ public:
+  MockCausalLM(const torch::TensorOptions& options) : options_(options) {
+    auto weight = torch::randn({1024, 1024}, options_) * 0.02;
+    weight_ = register_parameter("weight", weight, false);
+  }
+
+  ModelOutput forward(const torch::Tensor& tokens,
+                      const torch::Tensor& positions,
+                      std::vector<KVCache>& kv_caches,
+                      const ModelInputParams& params) override {
+    (void)tokens;
+    (void)positions;
+    (void)kv_caches;
+    ++forward_cnt_;
+    last_tokens_size_ = tokens.size(0);
+    last_dp_token_nums_ = params.dp_global_token_nums;
+    auto hidden_states = params.input_embedding.matmul(weight_);
+    if (return_aux_hidden_states_) {
+      auto aux_hidden_states = hidden_states + 1;
+      return ModelOutput(hidden_states, torch::Tensor(), aux_hidden_states);
+    }
+    return ModelOutput(hidden_states);
+  }
+  torch::Tensor logits(const torch::Tensor& hidden_states,
+                       const torch::Tensor& seleted_idxes) override {
+    (void)seleted_idxes;
+    return hidden_states;
+  }
+  int32_t forward_cnt() const { return forward_cnt_; }
+  int64_t last_tokens_size() const { return last_tokens_size_; }
+  const std::vector<int32_t>& last_dp_token_nums() const {
+    return last_dp_token_nums_;
+  }
+  void return_aux_hidden_states(bool value) {
+    return_aux_hidden_states_ = value;
+  }
+  void load_model(std::unique_ptr<ModelLoader> loader) override {}
+  torch::Device device() const override { return options_.device(); }
+  void prepare_expert_weight(int32_t layer_id,
+                             const std::vector<int32_t>& expert_ids) override {}
+  void update_expert_weight(int32_t layer_id) override {}
+  const torch::TensorOptions& options() const override { return options_; }
+
+ private:
+  torch::Tensor input_;
+  torch::Tensor weight_;
+  torch::TensorOptions options_;
+  bool return_aux_hidden_states_ = false;
+  int32_t forward_cnt_ = 0;
+  int64_t last_tokens_size_ = 0;
+  std::vector<int32_t> last_dp_token_nums_;
+};
+
+class MluGraphExecutorTest : public ::testing::Test {
+ protected:
+  MluGraphExecutorTest() = default;
+
+  void SetUp() override {
+    torch::Device device("mlu:0");
+    tensor_options_ = torch::TensorOptions(torch::kBFloat16).device(device);
+
+    model_args_.model_type("test_model");
+    model_args_.dtype("bfloat16");
+    model_args_.hidden_size(1024);
+    model_args_.max_position_embeddings(2048);
+
+    const uint32_t block_size = 16;
+    options_.num_decoding_tokens(1);
+    options_.block_size(block_size);
+
+    model_ = std::make_unique<MockCausalLM>(tensor_options_);
+    rebuild_impl();
+  }
+
+  ForwardInput prepare_inputs(int32_t batch_size, uint64_t seed) {
+    Device device(tensor_options_.device());
+    device.set_seed(seed);
+    const int64_t max_seq_len = model_args_.max_position_embeddings();
+    const uint32_t block_size = options_.block_size();
+    const int64_t num_blocks_per_req =
+        (max_seq_len + block_size - 1) / block_size + 1;
+    auto int_tensor_options = tensor_options_.dtype(torch::kInt32);
+    auto token_ids = torch::full({batch_size}, 1, int_tensor_options);
+    auto positions = torch::full({batch_size}, 1, int_tensor_options);
+    auto new_cache_slots =
+        torch::randint(0, 10, {batch_size}, int_tensor_options);
+    auto block_table = torch::randint(
+        0, 10, {batch_size, num_blocks_per_req}, int_tensor_options);
+    std::vector<int32_t> q_seq_lens_vec(batch_size + 1, 0);
+    std::vector<int32_t> kv_seq_lens_vec(batch_size + 1, 0);
+    for (int32_t i = 0; i < batch_size; ++i) {
+      q_seq_lens_vec[i + 1] = q_seq_lens_vec[i] + 1;
+      kv_seq_lens_vec[i + 1] = kv_seq_lens_vec[i] + 1;
+    }
+    auto q_seq_lens = torch::tensor(q_seq_lens_vec, int_tensor_options);
+    auto kv_seq_lens = torch::tensor(kv_seq_lens_vec, int_tensor_options);
+    auto input_embedding =
+        torch::randn({batch_size, model_args_.hidden_size()}, tensor_options_) *
+        0.1;
+    ModelInputParams input_params;
+    input_params.batch_forward_type = BatchForwardType::DECODE;
+    input_params.num_sequences = batch_size;
+    input_params.kv_max_seq_len = 1;
+    input_params.q_max_seq_len = 1;
+    input_params.dp_global_token_nums = {1};
+    input_params.dp_is_decode = {1};
+    input_params.new_cache_slots = new_cache_slots;
+    input_params.block_tables = block_table;
+    input_params.q_seq_lens = q_seq_lens;
+    input_params.kv_seq_lens = kv_seq_lens;
+    input_params.q_seq_lens_vec = q_seq_lens_vec;
+    input_params.kv_seq_lens_vec = kv_seq_lens_vec;
+    input_params.input_embedding = input_embedding;
+
+    kv_caches_.resize(batch_size);
+    return {token_ids, positions, input_params};
+  }
+
+  void rebuild_impl() {
+    const torch::Device device("mlu:0");
+    impl_ = std::make_unique<::xllm::mlu::MluGraphExecutorImpl>(
+        model_.get(), model_args_, device, options_);
+    base_impl_ = std::make_unique<BaseExecutorImpl>(
+        model_.get(), model_args_, device, options_);
+  }
+
+  ModelArgs model_args_;
+  torch::TensorOptions tensor_options_;
+  runtime::Options options_;
+  std::unique_ptr<MockCausalLM> model_;
+  std::vector<KVCache> kv_caches_;
+  std::unique_ptr<::xllm::mlu::MluGraphExecutorImpl> impl_;
+  std::unique_ptr<BaseExecutorImpl> base_impl_;
+};
+
+// Test graph creation and execution with different batch sizes
+TEST_F(MluGraphExecutorTest, DifferentBatchSizes) {
+  // Test with different batch sizes to ensure graph creation works
+  const std::vector<uint32_t> batch_sizes = {1, 3, 13, 21, 65};
+  for (auto batch_size : batch_sizes) {
+    auto forward_input = prepare_inputs(batch_size, 1);
+    auto eager_model_output = base_impl_->run({forward_input.token_ids},
+                                              {forward_input.positions},
+                                              kv_caches_,
+                                              {forward_input.input_params});
+    auto eager_output = eager_model_output.hidden_states;
+
+    auto graph_model_output = impl_->run({forward_input.token_ids},
+                                         {forward_input.positions},
+                                         kv_caches_,
+                                         {forward_input.input_params});
+    auto graph_output = graph_model_output.hidden_states;
+
+    auto replay_model_output = impl_->run({forward_input.token_ids},
+                                          {forward_input.positions},
+                                          kv_caches_,
+                                          {forward_input.input_params});
+    auto replay_output = replay_model_output.hidden_states;
+
+    CHECK_EQ(eager_output.sizes(), graph_output.sizes());
+    CHECK_EQ(eager_output.sizes(), replay_output.sizes());
+    // Compare outputs - should be identical
+    torch_mlu::synchronize();
+    EXPECT_TRUE(torch::allclose(eager_output, graph_output, 1e-5, 1e-6));
+    EXPECT_TRUE(torch::allclose(eager_output, replay_output, 1e-5, 1e-6));
+  }
+}
+
+// Test multiple runs to verify consistency across different execution modes
+TEST_F(MluGraphExecutorTest, MluGraphExecutorVsBaseExecutorImplMultipleRuns) {
+  int32_t batch_size = 5;
+  int32_t seed = 42;
+  auto forward_input = prepare_inputs(batch_size, seed);
+  auto eager_model_output = base_impl_->run({forward_input.token_ids},
+                                            {forward_input.positions},
+                                            kv_caches_,
+                                            {forward_input.input_params});
+  auto eager_output = eager_model_output.hidden_states;
+
+  auto graph_model_output = impl_->run({forward_input.token_ids},
+                                       {forward_input.positions},
+                                       kv_caches_,
+                                       {forward_input.input_params});
+  auto graph_output = graph_model_output.hidden_states;
+
+  CHECK_EQ(eager_output.sizes(), graph_output.sizes());
+  // Compare outputs - should be identical
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(eager_output, graph_output, 1e-5, 1e-6));
+
+  // Run multiple times and compare results
+  const int num_runs = 5;
+  auto base_forward_input = prepare_inputs(batch_size + 1, seed);
+  auto replay_forward_input = prepare_inputs(batch_size + 1, seed);
+  EXPECT_TRUE(torch::allclose(base_forward_input.input_params.input_embedding,
+                              replay_forward_input.input_params.input_embedding,
+                              1e-5,
+                              1e-6));
+
+  for (int i = 0; i < num_runs; ++i) {
+    auto base_model_output = base_impl_->run({base_forward_input.token_ids},
+                                             {base_forward_input.positions},
+                                             kv_caches_,
+                                             {base_forward_input.input_params});
+    auto base_output = base_model_output.hidden_states;
+
+    auto replay_model_output = impl_->run({replay_forward_input.token_ids},
+                                          {replay_forward_input.positions},
+                                          kv_caches_,
+                                          {replay_forward_input.input_params});
+    auto replay_output = replay_model_output.hidden_states;
+    base_forward_input.input_params.input_embedding = base_output;
+    replay_forward_input.input_params.input_embedding = replay_output;
+    CHECK_EQ(base_output.sizes(), replay_output.sizes());
+  }
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(base_forward_input.input_params.input_embedding,
+                              replay_forward_input.input_params.input_embedding,
+                              1e-5,
+                              1e-6));
+}
+
+TEST_F(MluGraphExecutorTest, DraftDecodeFallsBackToEager) {
+  options_.is_draft_engine(true);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  const uint64_t seed = 7;
+  auto forward_input = prepare_inputs(batch_size, seed);
+
+  auto eager_model_output = base_impl_->run({forward_input.token_ids},
+                                            {forward_input.positions},
+                                            kv_caches_,
+                                            {forward_input.input_params});
+  auto eager_output = eager_model_output.hidden_states;
+
+  auto first_impl_output = impl_
+                               ->run({forward_input.token_ids},
+                                     {forward_input.positions},
+                                     kv_caches_,
+                                     {forward_input.input_params})
+                               .hidden_states;
+  auto second_impl_output = impl_
+                                ->run({forward_input.token_ids},
+                                      {forward_input.positions},
+                                      kv_caches_,
+                                      {forward_input.input_params})
+                                .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(eager_output, first_impl_output, 1e-5, 1e-6));
+  EXPECT_TRUE(
+      torch::allclose(first_impl_output, second_impl_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), 3);
+}
+
+TEST_F(MluGraphExecutorTest, DraftEagerDoesNotExposeAuxWhenDisabled) {
+  model_->return_aux_hidden_states(true);
+  options_.is_draft_engine(true);
+  options_.enable_graph_aux_hidden_states(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  const uint64_t seed = 17;
+  auto forward_input = prepare_inputs(batch_size, seed);
+
+  ModelOutput output = impl_->run({forward_input.token_ids},
+                                  {forward_input.positions},
+                                  kv_caches_,
+                                  {forward_input.input_params});
+
+  EXPECT_FALSE(output.aux_hidden_states.defined());
+  EXPECT_EQ(model_->forward_cnt(), 1);
+}
+
+TEST_F(MluGraphExecutorTest, TargetDecodeCapturesThenReplays) {
+  options_.is_draft_engine(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  const uint64_t seed = 11;
+  auto forward_input = prepare_inputs(batch_size, seed);
+
+  auto first_impl_output = impl_
+                               ->run({forward_input.token_ids},
+                                     {forward_input.positions},
+                                     kv_caches_,
+                                     {forward_input.input_params})
+                               .hidden_states;
+  auto second_impl_output = impl_
+                                ->run({forward_input.token_ids},
+                                      {forward_input.positions},
+                                      kv_caches_,
+                                      {forward_input.input_params})
+                                .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(
+      torch::allclose(first_impl_output, second_impl_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), 1);
+}
+
+TEST_F(MluGraphExecutorTest, PrefillThenDecodeCapturesAndReplays) {
+  options_.is_draft_engine(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  const uint64_t prefill_seed = 23;
+  auto prefill_input = prepare_inputs(batch_size, prefill_seed);
+  prefill_input.input_params.batch_forward_type = BatchForwardType::PREFILL;
+
+  ModelOutput prefill_output = impl_->run({prefill_input.token_ids},
+                                          {prefill_input.positions},
+                                          kv_caches_,
+                                          {prefill_input.input_params});
+
+  const uint64_t decode_seed = 29;
+  auto decode_input = prepare_inputs(batch_size, decode_seed);
+  auto first_decode_output = impl_
+                                 ->run({decode_input.token_ids},
+                                       {decode_input.positions},
+                                       kv_caches_,
+                                       {decode_input.input_params})
+                                 .hidden_states;
+  auto second_decode_output = impl_
+                                  ->run({decode_input.token_ids},
+                                        {decode_input.positions},
+                                        kv_caches_,
+                                        {decode_input.input_params})
+                                  .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(prefill_output.hidden_states.defined());
+  EXPECT_TRUE(
+      torch::allclose(first_decode_output, second_decode_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), 2);
+}
+
+TEST_F(MluGraphExecutorTest, EqualDpDecodePadsToTpGraphSize) {
+  options_.is_draft_engine(false);
+  options_.world_size(8);
+  options_.dp_size(2);
+  rebuild_impl();
+
+  auto forward_input = prepare_inputs(/*batch_size=*/2, /*seed=*/61);
+  forward_input.input_params.dp_global_token_nums = {2, 2};
+  forward_input.input_params.dp_is_decode = {1, 1};
+
+  auto first_output = impl_
+                          ->run({forward_input.token_ids},
+                                {forward_input.positions},
+                                kv_caches_,
+                                {forward_input.input_params})
+                          .hidden_states;
+  auto second_output = impl_
+                           ->run({forward_input.token_ids},
+                                 {forward_input.positions},
+                                 kv_caches_,
+                                 {forward_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_output, second_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), 1);
+  EXPECT_EQ(model_->last_tokens_size(), 4);
+  EXPECT_EQ(model_->last_dp_token_nums(), std::vector<int32_t>({4, 4}));
+}
+
+TEST_F(MluGraphExecutorTest, UnevenDpDecodePadsToTpGraphSize) {
+  options_.is_draft_engine(false);
+  options_.world_size(8);
+  options_.dp_size(2);
+  rebuild_impl();
+
+  auto forward_input = prepare_inputs(/*batch_size=*/2, /*seed=*/67);
+  forward_input.input_params.dp_global_token_nums = {1, 2};
+  forward_input.input_params.dp_is_decode = {1, 1};
+
+  auto first_output = impl_
+                          ->run({forward_input.token_ids},
+                                {forward_input.positions},
+                                kv_caches_,
+                                {forward_input.input_params})
+                          .hidden_states;
+  auto second_output = impl_
+                           ->run({forward_input.token_ids},
+                                 {forward_input.positions},
+                                 kv_caches_,
+                                 {forward_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_output, second_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), 1);
+  EXPECT_EQ(model_->last_tokens_size(), 4);
+  EXPECT_EQ(model_->last_dp_token_nums(), std::vector<int32_t>({4, 4}));
+}
+
+TEST_F(MluGraphExecutorTest, MtpSeqLensCapacityUsesSpecFactor) {
+  options_.is_draft_engine(false);
+  options_.num_speculative_tokens(1);
+  options_.max_seqs_per_batch(2);
+  rebuild_impl();
+
+  auto forward_input = prepare_inputs(/*batch_size=*/4, /*seed=*/71);
+  forward_input.input_params.dp_global_token_nums = {4, 4};
+  forward_input.input_params.dp_is_decode = {1, 1};
+
+  auto first_output = impl_
+                          ->run({forward_input.token_ids},
+                                {forward_input.positions},
+                                kv_caches_,
+                                {forward_input.input_params})
+                          .hidden_states;
+  auto second_output = impl_
+                           ->run({forward_input.token_ids},
+                                 {forward_input.positions},
+                                 kv_caches_,
+                                 {forward_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_output, second_output, 1e-5, 1e-6));
+}
+
+TEST_F(MluGraphExecutorTest, DpDummyFallsBackToEager) {
+  options_.is_draft_engine(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  const uint64_t seed = 31;
+  auto forward_input = prepare_inputs(batch_size, seed);
+  forward_input.input_params.dp_global_token_nums = {batch_size, 0};
+  forward_input.input_params.dp_is_decode = {1, 0};
+
+  const int32_t start_cnt = model_->forward_cnt();
+  auto first_output = impl_
+                          ->run({forward_input.token_ids},
+                                {forward_input.positions},
+                                kv_caches_,
+                                {forward_input.input_params})
+                          .hidden_states;
+  auto second_output = impl_
+                           ->run({forward_input.token_ids},
+                                 {forward_input.positions},
+                                 kv_caches_,
+                                 {forward_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_output, second_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), start_cnt + 2);
+}
+
+TEST_F(MluGraphExecutorTest, DpUnevenDecodeFallsBackToEager) {
+  options_.is_draft_engine(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  auto forward_input = prepare_inputs(batch_size, 43);
+  forward_input.input_params.dp_global_token_nums = {batch_size,
+                                                     batch_size - 1};
+  forward_input.input_params.dp_is_decode = {1, 1};
+  forward_input.input_params.q_max_seq_len = 2;
+
+  const int32_t start_cnt = model_->forward_cnt();
+  auto first_output = impl_
+                          ->run({forward_input.token_ids},
+                                {forward_input.positions},
+                                kv_caches_,
+                                {forward_input.input_params})
+                          .hidden_states;
+  auto second_output = impl_
+                           ->run({forward_input.token_ids},
+                                 {forward_input.positions},
+                                 kv_caches_,
+                                 {forward_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_output, second_output, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), start_cnt + 2);
+}
+
+TEST_F(MluGraphExecutorTest, DpDummyDoesNotPoisonGraphCache) {
+  options_.is_draft_engine(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  auto dummy_input = prepare_inputs(batch_size, 37);
+  dummy_input.input_params.dp_global_token_nums = {batch_size, 0};
+  dummy_input.input_params.dp_is_decode = {1, 0};
+
+  const int32_t start_cnt = model_->forward_cnt();
+  impl_->run({dummy_input.token_ids},
+             {dummy_input.positions},
+             kv_caches_,
+             {dummy_input.input_params});
+
+  auto decode_input = prepare_inputs(batch_size, 41);
+  decode_input.input_params.dp_global_token_nums = {batch_size, batch_size};
+  decode_input.input_params.dp_is_decode = {1, 1};
+
+  auto first_decode = impl_
+                          ->run({decode_input.token_ids},
+                                {decode_input.positions},
+                                kv_caches_,
+                                {decode_input.input_params})
+                          .hidden_states;
+  auto second_decode = impl_
+                           ->run({decode_input.token_ids},
+                                 {decode_input.positions},
+                                 kv_caches_,
+                                 {decode_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_decode, second_decode, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), start_cnt + 2);
+}
+
+TEST_F(MluGraphExecutorTest, DpUnevenDecodeDoesNotPoisonGraphCache) {
+  options_.is_draft_engine(false);
+  rebuild_impl();
+
+  const int32_t batch_size = 5;
+  auto uneven_input = prepare_inputs(batch_size, 47);
+  uneven_input.input_params.dp_global_token_nums = {batch_size, batch_size - 1};
+  uneven_input.input_params.dp_is_decode = {1, 1};
+  uneven_input.input_params.q_max_seq_len = 2;
+
+  const int32_t start_cnt = model_->forward_cnt();
+  impl_->run({uneven_input.token_ids},
+             {uneven_input.positions},
+             kv_caches_,
+             {uneven_input.input_params});
+
+  auto decode_input = prepare_inputs(batch_size, 53);
+  decode_input.input_params.dp_global_token_nums = {batch_size, batch_size};
+  decode_input.input_params.dp_is_decode = {1, 1};
+
+  auto first_decode = impl_
+                          ->run({decode_input.token_ids},
+                                {decode_input.positions},
+                                kv_caches_,
+                                {decode_input.input_params})
+                          .hidden_states;
+  auto second_decode = impl_
+                           ->run({decode_input.token_ids},
+                                 {decode_input.positions},
+                                 kv_caches_,
+                                 {decode_input.input_params})
+                           .hidden_states;
+
+  torch_mlu::synchronize();
+  EXPECT_TRUE(torch::allclose(first_decode, second_decode, 1e-5, 1e-6));
+  EXPECT_EQ(model_->forward_cnt(), start_cnt + 2);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/runtime/spec_input_builder_test.cpp

@@ -0,0 +1,494 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "runtime/spec_input_builder.h"
+
+#include <gtest/gtest.h>
+
+#include <vector>
+
+#include "framework/model/model_input_params.h"
+
+namespace xllm {
+namespace specBuilder {
+namespace {
+
+Slice<int32_t> to_slice(std::vector<int32_t>& vec) {
+  return {vec.data(), static_cast<size_t>(vec.size())};
+}
+
+std::vector<int32_t> to_layout_seq_lens(const std::vector<int32_t>& lens) {
+#if defined(USE_NPU)
+  return lens;
+#else
+  std::vector<int32_t> out;
+  out.reserve(lens.size() + 1);
+  out.emplace_back(0);
+  int32_t sum = 0;
+  for (int32_t len : lens) {
+    sum += len;
+    out.emplace_back(sum);
+  }
+  return out;
+#endif
+}
+
+std::vector<int32_t> tensor_to_vec_int32(const torch::Tensor& tensor) {
+  torch::Tensor cpu_tensor =
+      tensor.to(torch::kCPU).to(torch::kInt).contiguous();
+  const int32_t* data = cpu_tensor.data_ptr<int32_t>();
+  return {data, data + cpu_tensor.numel()};
+}
+
+TEST(SpecDecodeInputBuilderTest, DraftInputsSingleRowPerSeq) {
+  ModelInputParams params;
+  params.num_sequences = 2;
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
+
+  torch::Tensor positions = torch::tensor({4, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2}, {3, 4, 5}}, torch::kInt);
+  auto view = make_decode_cpu_view(
+      torch::Tensor(), positions, block_tables, to_slice(kv_seq_lens));
+
+  DecodeBuildBuffers buf;
+  for (int32_t seq_id = 0; seq_id < params.num_sequences; ++seq_id) {
+    RowSpec row;
+    row.seq_id = seq_id;
+    row.position_offset = 1;
+    row.append_token = false;
+    append_decode_row(view, row, /*block_size=*/4, buf);
+  }
+
+  EXPECT_TRUE(buf.out_token_ids.empty());
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({5, 9}));
+  EXPECT_EQ(buf.out_new_cache_slots, std::vector<int32_t>({5, 21}));
+  EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({6, 10}));
+}
+
+TEST(SpecDecodeInputBuilderTest, ValidateInputsNonAtbExpansion) {
+  ModelInputParams params;
+  params.num_sequences = 2;
+  const int32_t num_speculative_tokens = 2;
+  const int32_t num_val_tokens = num_speculative_tokens + 1;
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
+
+  torch::Tensor token_ids = torch::tensor({10, 20}, torch::kInt);
+  torch::Tensor positions = torch::tensor({4, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2}, {3, 4, 5}}, torch::kInt);
+  auto view = make_decode_cpu_view(
+      token_ids, positions, block_tables, to_slice(kv_seq_lens));
+
+  DecodeBuildBuffers buf;
+  for (int32_t seq_id = 0; seq_id < params.num_sequences; ++seq_id) {
+    for (int32_t val_idx = 0; val_idx < num_val_tokens; ++val_idx) {
+      RowSpec row;
+      row.seq_id = seq_id;
+      if (val_idx == 0) {
+        row.use_input_token = true;
+      } else {
+        row.token_id = -1 * val_idx;
+      }
+      row.position_offset = 1 + val_idx;
+      row.append_q_len_one = true;
+      row.append_block_table = true;
+      append_decode_row(view, row, /*block_size=*/4, buf);
+    }
+  }
+
+  EXPECT_EQ(buf.out_token_ids, std::vector<int32_t>({10, -1, -2, 20, -1, -2}));
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({5, 6, 7, 9, 10, 11}));
+  EXPECT_EQ(buf.out_new_cache_slots,
+            std::vector<int32_t>({5, 6, 7, 21, 22, 23}));
+  EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({6, 7, 8, 10, 11, 12}));
+  EXPECT_EQ(buf.out_q_seq_lens, to_layout_seq_lens({1, 1, 1, 1, 1, 1}));
+  ASSERT_EQ(buf.out_block_tables.size(), 6);
+}
+
+TEST(SpecDecodeInputBuilderTest, AppendDecodeRowTokenKinds) {
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
+  torch::Tensor token_ids = torch::tensor({10, 20}, torch::kInt);
+  torch::Tensor positions = torch::tensor({4, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2}, {3, 4, 5}}, torch::kInt);
+  auto view = make_decode_cpu_view(
+      token_ids, positions, block_tables, to_slice(kv_seq_lens));
+
+  DecodeBuildBuffers buf;
+  append_decode_row(
+      view,
+      {.seq_id = 1, .use_input_token = true, .position_offset = 0},
+      /*block_size=*/4,
+      buf);
+  append_decode_row(view,
+                    {.seq_id = 0, .token_id = 123, .position_offset = 0},
+                    /*block_size=*/4,
+                    buf);
+  append_decode_row(view,
+                    {.seq_id = 0, .token_id = -2, .position_offset = 0},
+                    /*block_size=*/4,
+                    buf);
+
+  EXPECT_EQ(buf.out_token_ids, std::vector<int32_t>({20, 123, -2}));
+}
+
+TEST(SpecDecodeInputBuilderTest, MakeDecodeCpuViewUsesFlatBlockTableLayout) {
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
+  torch::Tensor token_ids = torch::tensor({10, 20}, torch::kInt);
+  torch::Tensor positions = torch::tensor({4, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2, 0}, {3, 4, 5, 0}}, torch::kInt);
+
+  DecodeCpuView view = make_decode_cpu_view(
+      token_ids, positions, block_tables, to_slice(kv_seq_lens));
+
+  EXPECT_EQ(view.num_sequences, 2);
+  EXPECT_EQ(view.block_table_row_stride, 4);
+  EXPECT_EQ(view.block_tables_data,
+            std::vector<int32_t>({0, 1, 2, 0, 3, 4, 5, 0}));
+
+  ModelInputParams params;
+  params.num_sequences = 2;
+  DecodeBuildBuffers buf;
+  append_decode_row(view,
+                    {.seq_id = 1, .token_id = 99, .position_offset = 2},
+                    /*block_size=*/4,
+                    buf);
+
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({10}));
+  EXPECT_EQ(buf.out_new_cache_slots, std::vector<int32_t>({22}));
+  ASSERT_EQ(buf.out_block_tables.size(), 0);
+}
+
+TEST(SpecDecodeInputBuilderTest, ValidateRowsStartFromCorrectedCurrentView) {
+  ModelInputParams params;
+  params.num_sequences = 2;
+  std::vector<int32_t> token_ids = {31, 41};
+  std::vector<int32_t> positions = {6, 9};
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({7, 10});
+
+  DecodeCpuView view;
+  view.token_ids = token_ids;
+  view.positions = positions;
+  view.kv_seq_lens = kv_seq_lens;
+  view.block_tables_cpu =
+      torch::tensor({{0, 1, 2, 0}, {3, 4, 5, 0}}, torch::kInt);
+  view.block_tables_data = {view.block_tables_cpu.data_ptr<int32_t>(),
+                            static_cast<size_t>(view.block_tables_cpu.numel())};
+  view.num_sequences = static_cast<int32_t>(view.block_tables_cpu.size(0));
+  view.block_table_row_stride =
+      static_cast<int32_t>(view.block_tables_cpu.size(1));
+
+  DecodeBuildBuffers buf;
+  append_decode_row(
+      view,
+      {.seq_id = 0, .token_id = view.token_ids[0], .position_offset = 0},
+      /*block_size=*/4,
+      buf);
+  append_decode_row(view,
+                    {.seq_id = 0, .token_id = -1, .position_offset = 1},
+                    /*block_size=*/4,
+                    buf);
+  append_decode_row(
+      view,
+      {.seq_id = 1, .token_id = view.token_ids[1], .position_offset = 0},
+      /*block_size=*/4,
+      buf);
+
+  EXPECT_EQ(buf.out_token_ids, std::vector<int32_t>({31, -1, 41}));
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({6, 7, 9}));
+  EXPECT_EQ(buf.out_new_cache_slots, std::vector<int32_t>({6, 7, 21}));
+  EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({7, 8, 10}));
+}
+
+TEST(SpecDecodeInputBuilderTest, ValidateInputsAtbChunkedPrefillShape) {
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
+  std::vector<int32_t> atb_kv_seq_lens;
+  std::vector<int32_t> atb_q_seq_lens;
+  int32_t atb_kv_max_seq_len = 0;
+  const int32_t num_val_tokens = 3;
+
+  auto kv_slice = to_slice(kv_seq_lens);
+  for (int32_t seq_id = 0; seq_id < 2; ++seq_id) {
+    int32_t kv_len = calc_kv_len(kv_slice, seq_id, /*offset=*/0);
+    int32_t kv_len_after_validation = kv_len + num_val_tokens;
+    update_kv_seq_lens_and_max(
+        atb_kv_seq_lens, kv_len_after_validation, atb_kv_max_seq_len);
+    append_seq_len_by_layout(atb_q_seq_lens, num_val_tokens);
+  }
+
+  EXPECT_EQ(atb_kv_seq_lens, to_layout_seq_lens({8, 12}));
+  EXPECT_EQ(atb_q_seq_lens, to_layout_seq_lens({3, 3}));
+  EXPECT_EQ(atb_kv_max_seq_len, 12);
+}
+
+TEST(SpecDecodeInputBuilderTest, FirstDecodeInputsFixAndNonFixMix) {
+  ModelInputParams params;
+  params.num_sequences = 2;
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({6, 9});
+
+  torch::Tensor token_ids = torch::tensor({100, 200}, torch::kInt);
+  torch::Tensor positions = torch::tensor({5, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2}, {3, 4, 5}}, torch::kInt);
+  auto view = make_decode_cpu_view(
+      token_ids, positions, block_tables, to_slice(kv_seq_lens));
+
+  DecodeBuildBuffers buf;
+  std::vector<int32_t> select_row_idx(2, 0);
+  auto emit_row =
+      [&](int32_t seq_id, int32_t token_id, int32_t position_offset) {
+        RowSpec row;
+        row.seq_id = seq_id;
+        row.token_id = token_id;
+        row.position_offset = position_offset;
+        row.append_q_len_one = true;
+        row.append_block_table = true;
+        append_decode_row(view, row, /*block_size=*/4, buf);
+      };
+
+  emit_row(/*seq_id=*/0, /*token_id=*/90, /*position_offset=*/-1);
+  emit_row(/*seq_id=*/0, /*token_id=*/100, /*position_offset=*/0);
+  select_row_idx[0] = static_cast<int32_t>(buf.out_token_ids.size()) - 1;
+
+  emit_row(/*seq_id=*/1, /*token_id=*/200, /*position_offset=*/0);
+  select_row_idx[1] = static_cast<int32_t>(buf.out_token_ids.size()) - 1;
+
+  EXPECT_EQ(buf.out_token_ids, std::vector<int32_t>({90, 100, 200}));
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({4, 5, 8}));
+  EXPECT_EQ(buf.out_new_cache_slots, std::vector<int32_t>({4, 5, 20}));
+  EXPECT_EQ(buf.out_q_seq_lens, to_layout_seq_lens({1, 1, 1}));
+  EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({5, 6, 9}));
+  EXPECT_EQ(select_row_idx, std::vector<int32_t>({1, 2}));
+  ASSERT_EQ(buf.out_block_tables.size(), 3);
+}
+
+TEST(SpecDecodeInputBuilderTest, AppendDecodeRowWithInputTokenSource) {
+  ModelInputParams params;
+  params.num_sequences = 2;
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
+
+  torch::Tensor token_ids = torch::tensor({10, 20}, torch::kInt);
+  torch::Tensor positions = torch::tensor({4, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2}, {3, 4, 5}}, torch::kInt);
+  auto view = make_decode_cpu_view(
+      token_ids, positions, block_tables, to_slice(kv_seq_lens));
+
+  DecodeBuildBuffers buf;
+  append_decode_row(view,
+                    {.seq_id = 0,
+                     .use_input_token = true,
+                     .position_offset = 1,
+                     .append_q_len_one = true,
+                     .append_block_table = true},
+                    /*block_size=*/4,
+                    buf);
+  append_decode_row(view,
+                    {.seq_id = 1,
+                     .token_id = -2,
+                     .position_offset = 2,
+                     .append_q_len_one = true,
+                     .append_block_table = true},
+                    /*block_size=*/4,
+                    buf);
+
+  EXPECT_EQ(buf.out_token_ids, std::vector<int32_t>({10, -2}));
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({5, 10}));
+  EXPECT_EQ(buf.out_new_cache_slots, std::vector<int32_t>({5, 22}));
+  EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({6, 11}));
+  EXPECT_EQ(buf.out_q_seq_lens, to_layout_seq_lens({1, 1}));
+  ASSERT_EQ(buf.out_block_tables.size(), 2);
+}
+
+TEST(SpecDecodeInputBuilderTest, ResolveTokenWithPositionOffset) {
+  std::vector<int64_t> last_step_tokens = {11, -1, 13, -1, -1, -1};
+  Slice<int64_t> last_step_slice = {
+      last_step_tokens.data(), static_cast<size_t>(last_step_tokens.size())};
+
+  TokenWithOffset direct =
+      resolve_token_with_position_offset(/*input_token_id=*/20,
+                                         /*seq_id=*/0,
+                                         last_step_slice,
+                                         /*last_step_decode_num=*/3);
+  EXPECT_EQ(direct.token_id, 20);
+  EXPECT_EQ(direct.position_offset, 0);
+
+  TokenWithOffset resolved =
+      resolve_token_with_position_offset(/*input_token_id=*/-1,
+                                         /*seq_id=*/0,
+                                         last_step_slice,
+                                         /*last_step_decode_num=*/3);
+  EXPECT_EQ(resolved.token_id, 13);
+  EXPECT_EQ(resolved.position_offset, 1);
+
+  TokenWithOffset no_accept =
+      resolve_token_with_position_offset(/*input_token_id=*/-2,
+                                         /*seq_id=*/1,
+                                         last_step_slice,
+                                         /*last_step_decode_num=*/3);
+  EXPECT_EQ(no_accept.token_id, 0);
+  EXPECT_EQ(no_accept.position_offset, -1);
+}
+
+TEST(SpecDecodeInputBuilderTest, AppendDecodeRowFromLastStep) {
+  ModelInputParams params;
+  params.num_sequences = 2;
+  std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({6, 9});
+
+  torch::Tensor token_ids = torch::tensor({100, -1}, torch::kInt);
+  torch::Tensor positions = torch::tensor({5, 8}, torch::kInt);
+  torch::Tensor block_tables =
+      torch::tensor({{0, 1, 2}, {3, 4, 5}}, torch::kInt);
+  auto view = make_decode_cpu_view(
+      token_ids, positions, block_tables, to_slice(kv_seq_lens));
+
+  std::vector<int64_t> last_step_tokens = {201, 202};
+  Slice<int64_t> last_step_slice = {
+      last_step_tokens.data(), static_cast<size_t>(last_step_tokens.size())};
+
+  DecodeBuildBuffers buf;
+  append_decode_row_from_last_step(view,
+                                   /*seq_id=*/0,
+                                   /*input_token_id=*/view.token_ids[0],
+                                   last_step_slice,
+                                   /*last_step_decode_num=*/2,
+                                   /*block_size=*/4,
+                                   buf);
+  append_decode_row_from_last_step(view,
+                                   /*seq_id=*/1,
+                                   /*input_token_id=*/view.token_ids[1],
+                                   last_step_slice,
+                                   /*last_step_decode_num=*/2,
+                                   /*block_size=*/4,
+                                   buf);
+
+  EXPECT_EQ(buf.out_token_ids, std::vector<int32_t>({100, 202}));
+  EXPECT_EQ(buf.out_positions, std::vector<int32_t>({5, 9}));
+  EXPECT_EQ(buf.out_new_cache_slots, std::vector<int32_t>({5, 21}));
+  EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({6, 10}));
+}
+
+TEST(SpecDecodeInputBuilderTest, QCuSeqLensConsistency) {
+  ModelInputParams params;
+  params.num_sequences = 3;
+  params.q_seq_lens_vec = to_layout_seq_lens({1, 2, 3});
+
+  torch::Tensor q_cu_seq_lens = build_q_cu_seq_lens_tensor(params);
+  EXPECT_EQ(tensor_to_vec_int32(q_cu_seq_lens),
+            std::vector<int32_t>({1, 3, 6}));
+}
+
+TEST(SpecDecodeInputBuilderTest, CalcSlotIdOutOfRangeDeath) {
+  std::vector<int32_t> block_table = {0};
+  EXPECT_DEATH(calc_slot_id(/*position=*/4,
+                            to_slice(block_table),
+                            /*block_size=*/4),
+               "block table index out of range");
+}
+
+TEST(DraftProbsBuilderTest, CompressForCacheDense) {
+  auto draft_probs =
+      torch::tensor({{0.1f, 0.2f, 0.7f}, {0.6f, 0.1f, 0.3f}}, torch::kFloat32);
+  auto token_ids = torch::tensor({1, 0}, torch::kInt64);
+  auto compressed = draftProbs::compress_for_cache(draft_probs, token_ids);
+  auto expected = torch::tensor({0.2f, 0.6f}, torch::kFloat32);
+  EXPECT_TRUE(torch::allclose(compressed, expected));
+}
+
+TEST(DraftProbsBuilderTest, BuildValidateTensorsSelectedOnly) {
+  std::vector<torch::Tensor> token_steps = {
+      torch::tensor({3, 4}, torch::kInt64),
+      torch::tensor({5, 6}, torch::kInt64)};
+  std::vector<torch::Tensor> probs_steps = {
+      torch::tensor({0.3f, 0.4f}, torch::kFloat32),
+      torch::tensor({0.5f, 0.6f}, torch::kFloat32)};
+
+  auto [draft_token_ids, draft_probs] =
+      draftProbs::build_validate_tensors(token_steps,
+                                         probs_steps,
+                                         /*batch_size=*/2,
+                                         /*vocab_size=*/8,
+                                         /*enable_opt_validate_probs=*/true);
+
+  EXPECT_EQ(draft_token_ids.dim(), 2);
+  EXPECT_EQ(draft_probs.dim(), 2);
+  EXPECT_EQ(draft_token_ids.size(0), 2);
+  EXPECT_EQ(draft_token_ids.size(1), 2);
+  EXPECT_EQ(draft_probs.size(0), 2);
+  EXPECT_EQ(draft_probs.size(1), 2);
+  EXPECT_TRUE(torch::allclose(
+      draft_probs,
+      torch::tensor({{0.3f, 0.5f}, {0.4f, 0.6f}}, torch::kFloat32)));
+}
+
+TEST(DraftProbsBuilderTest, BuildValidateTensorsRecoveredDense) {
+  std::vector<torch::Tensor> token_steps = {
+      torch::tensor({1, 2}, torch::kInt64),
+      torch::tensor({0, 3}, torch::kInt64)};
+  std::vector<torch::Tensor> probs_steps = {
+      torch::tensor({0.2f, 0.7f}, torch::kFloat32),
+      torch::tensor({0.9f, 0.1f}, torch::kFloat32)};
+
+  auto [draft_token_ids, draft_probs] =
+      draftProbs::build_validate_tensors(token_steps,
+                                         probs_steps,
+                                         /*batch_size=*/2,
+                                         /*vocab_size=*/5,
+                                         /*enable_opt_validate_probs=*/false);
+
+  EXPECT_EQ(draft_token_ids.dim(), 2);
+  EXPECT_EQ(draft_probs.dim(), 3);
+  EXPECT_EQ(draft_probs.size(0), 2);
+  EXPECT_EQ(draft_probs.size(1), 2);
+  EXPECT_EQ(draft_probs.size(2), 5);
+
+  auto selected =
+      draft_probs.gather(/*dim=*/-1, draft_token_ids.unsqueeze(-1)).squeeze(-1);
+  auto expected_selected =
+      torch::tensor({{0.2f, 0.9f}, {0.7f, 0.1f}}, torch::kFloat32);
+  EXPECT_TRUE(torch::allclose(selected, expected_selected));
+
+  auto row_sums = draft_probs.sum(/*dim=*/-1);
+  EXPECT_TRUE(torch::allclose(row_sums, expected_selected));
+}
+
+TEST(DraftProbsBuilderTest, BuildValidateTensorsDenseInputFallback) {
+  std::vector<torch::Tensor> token_steps = {
+      torch::tensor({2, 1}, torch::kInt64)};
+  std::vector<torch::Tensor> probs_steps = {
+      torch::tensor({{0.1f, 0.2f, 0.7f}, {0.3f, 0.6f, 0.1f}}, torch::kFloat32)};
+
+  auto [draft_token_ids, draft_probs] =
+      draftProbs::build_validate_tensors(token_steps,
+                                         probs_steps,
+                                         /*batch_size=*/2,
+                                         /*vocab_size=*/3,
+                                         /*enable_opt_validate_probs=*/true);
+
+  EXPECT_EQ(draft_token_ids.dim(), 2);
+  EXPECT_EQ(draft_token_ids.size(0), 2);
+  EXPECT_EQ(draft_token_ids.size(1), 1);
+  EXPECT_EQ(draft_probs.dim(), 2);
+  EXPECT_EQ(draft_probs.size(0), 2);
+  EXPECT_EQ(draft_probs.size(1), 1);
+  EXPECT_TRUE(torch::allclose(
+      draft_probs, torch::tensor({{0.7f}, {0.6f}}, torch::kFloat32)));
+}
+
+}  // namespace
+}  // namespace specBuilder
+}  // namespace xllm

upstream_ref/xllm/tests/core/scheduler/CMakeLists.txt

@@ -0,0 +1,24 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    chunked_prefill_scheduler_test
+  SRCS
+    chunked_prefill_scheduler_test.cpp
+    disagg_pd_chunked_prefill_scheduler_test.cpp
+    continuous_scheduler_test.cpp
+    fixed_steps_scheduler_test.cpp
+  DEPS
+    :scheduler
+    :xllm_server
+    GTest::gtest_main
+    $<$<BOOL:${USE_NPU}>:nnopbase>
+)
+target_link_libraries(chunked_prefill_scheduler_test
+                      PUBLIC
+                      Python::Python
+                      $<$<BOOL:${USE_NPU}>:ascendcl>
+                      $<$<BOOL:${USE_NPU}>:hccl>
+                      $<$<BOOL:${USE_NPU}>:c_sec>)
+target_link_libraries(chunked_prefill_scheduler_test PRIVATE
+  "$<LINK_GROUP:RESCAN,xtensor,xllm_server>")

upstream_ref/xllm/tests/core/scheduler/chunked_prefill_scheduler_test.cpp

@@ -0,0 +1,699 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "chunked_prefill_scheduler.h"
+
+#include <absl/time/clock.h>
+#include <gtest/gtest.h>
+
+#include <cmath>
+#include <optional>
+
+#include "distributed_runtime/engine.h"
+#include "util/utils.h"
+
+namespace xllm {
+
+namespace {
+class FakeTokenizer : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens) const {
+    NOT_IMPLEMENTED();
+  }
+  std::string decode(const Slice<int32_t>& ids,
+                     bool skip_special_tokens) const {
+    NOT_IMPLEMENTED();
+  }
+  std::optional<int32_t> token_to_id(const std::string_view& token) const {
+    NOT_IMPLEMENTED();
+  }
+  std::string id_to_token(int32_t id) const { NOT_IMPLEMENTED(); }
+  size_t vocab_size() const { NOT_IMPLEMENTED(); }
+  std::unique_ptr<Tokenizer> clone() const {
+    return std::make_unique<FakeTokenizer>();
+  }
+};
+
+class FakeEngine : public Engine {
+ public:
+  FakeEngine(int32_t num_blocks, int32_t block_size) {
+    BlockManagerPool::Options opt;
+    opt.num_blocks_ = num_blocks;
+    opt.block_size_ = block_size;
+    opt.enable_prefix_cache_ = false;  // we dont consider prefix cache here
+    fake_tokenizer_ = std::make_unique<FakeTokenizer>();
+    fake_block_manager_ = std::make_unique<BlockManagerPool>(opt, 1);
+  }
+  ForwardOutput step(std::vector<Batch>& batch) { NOT_IMPLEMENTED(); }
+  void update_last_step_result(std::vector<Batch>& batch) { NOT_IMPLEMENTED(); }
+  const Tokenizer* tokenizer() const { return fake_tokenizer_.get(); }
+  BlockManagerPool* block_manager_pool() const {
+    return fake_block_manager_.get();
+  }
+  const ModelArgs& model_args() const { NOT_IMPLEMENTED(); }
+  const TokenizerArgs& tokenizer_args() const { NOT_IMPLEMENTED(); }
+  std::vector<int64_t> get_active_activation_memory() const {
+    NOT_IMPLEMENTED();
+  }
+  bool init() override { return true; }
+
+ private:
+  std::unique_ptr<Tokenizer> fake_tokenizer_;
+  std::unique_ptr<BlockManagerPool> fake_block_manager_;
+};
+
+ContinuousScheduler::Options create_scheduler_options(
+    int32_t max_tokens_per_batch,
+    int32_t max_seqs_per_batch,
+    int32_t num_speculative_tokens,
+    int32_t max_tokens_per_chunk_for_prefill,
+    int32_t dp_size,
+    const std::string& priority_strategy = "fcfs",
+    bool enable_profile_kv_blocks = true,
+    bool enable_latency_aware_schedule = false,
+    int32_t max_global_ttft_ms = std::numeric_limits<int32_t>::max(),
+    int32_t max_global_tpot_ms = std::numeric_limits<int32_t>::max()) {
+  ContinuousScheduler::Options opt;
+  opt.num_speculative_tokens_ = num_speculative_tokens;
+  opt.max_tokens_per_chunk_for_prefill_ = max_tokens_per_chunk_for_prefill;
+  opt.max_tokens_per_batch_ = max_tokens_per_batch;
+  opt.max_seqs_per_batch_ = max_seqs_per_batch;
+  opt.dp_size_ = dp_size;
+  opt.priority_strategy_ = priority_strategy;
+  opt.enable_profile_kv_blocks_ = enable_profile_kv_blocks;
+  opt.enable_latency_aware_schedule_ = enable_latency_aware_schedule;
+  opt.max_global_ttft_ms_ = max_global_ttft_ms;
+  opt.max_global_tpot_ms_ = max_global_tpot_ms;
+  return opt;
+}
+
+std::vector<std::shared_ptr<Request>> generate_request(
+    const std::vector<int32_t>& prompt_lens,
+    const std::vector<int32_t>& max_tokens,
+    std::optional<std::vector<bool>> offlines,
+    std::optional<std::vector<int32_t>> priorities,
+    int32_t max_context_len) {
+  std::vector<std::shared_ptr<Request>> requests;
+  EXPECT_TRUE(prompt_lens.size() == max_tokens.size());
+
+  size_t batch_size = prompt_lens.size();
+  std::vector<bool> offline_vec;
+  std::vector<int32_t> priority_vec;
+  if (offlines.has_value()) {
+    offline_vec = *offlines;
+  } else {
+    offline_vec = std::vector<bool>(batch_size, false);
+  }
+
+  if (priorities.has_value()) {
+    priority_vec = priorities.value();
+  } else {
+    priority_vec = std::vector<int32_t>(
+        batch_size, static_cast<int32_t>(RequestPriority::NORMAL));
+  }
+
+  for (size_t i = 0; i < batch_size; ++i) {
+    std::vector<int32_t> prompt_token_ids;
+    prompt_token_ids.resize(prompt_lens[i]);
+    RequestSamplingParam sampling_param;
+    SchedulerParam scheduler_param;
+    scheduler_param.offline = offline_vec[i];
+    scheduler_param.priority = static_cast<RequestPriority>(priority_vec[i]);
+
+    StoppingChecker stopping_checker;
+    stopping_checker.set_max_generated_tokens(max_tokens[i]);
+    stopping_checker.set_max_context_len(max_context_len);
+    stopping_checker.set_ignore_eos(true);
+
+    RequestState req_state("x",
+                           prompt_token_ids,
+                           sampling_param,
+                           scheduler_param,
+                           stopping_checker,
+                           prompt_lens[i] + 30000,
+                           1,
+                           1,
+                           false,
+                           false,
+                           false,
+                           false,
+                           false,
+                           nullptr,
+                           nullptr);
+
+    auto request =
+        std::make_shared<Request>("1", "1", "1", std::move(req_state), "1");
+    requests.emplace_back(request);
+  }
+
+  return requests;
+}
+
+// dont not consider speculative decoding.
+void update_requests(std::vector<std::shared_ptr<Request>> requests) {
+  for (auto req : requests) {
+    for (auto& seq : req->sequences()) {
+      if (seq->kv_state().kv_cache_tokens_num() == 0) {
+        seq->kv_state().incr_kv_cache_tokens_num(seq->num_prompt_tokens());
+      } else {
+        seq->kv_state().incr_kv_cache_tokens_num(1);
+      }
+      Token token(1);
+      seq->append_token(token);
+    }
+  }
+}
+
+}  // namespace
+
+// TEST-1:
+// Three independent prefill requests, according to the configs,
+// verify how many tokens are processed in one scheduling.
+TEST(ChunkedPrefillSchedulerTest, AddNewRequestBase) {
+  std::vector<int32_t> prompt_len{10, 1024, 2048};
+  std::vector<int32_t> num_blocks{16, 128, 128};
+  std::vector<int32_t> block_size{16, 16, 16};
+  std::vector<int32_t> validate_allowed_max_tokens{10, 1024, 1024};
+  for (size_t idx = 0; idx < prompt_len.size(); ++idx) {
+    ContinuousScheduler::Options opt =
+        create_scheduler_options(10000, 256, 0, 1024, 1);
+    auto engine =
+        std::make_unique<FakeEngine>(num_blocks[idx], block_size[idx]);
+    auto scheduler =
+        std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+    EXPECT_TRUE(scheduler != nullptr);
+
+    // create requests
+    auto requests = generate_request(
+        {prompt_len[idx]}, {10}, std::nullopt, std::nullopt, 10000);
+    for (auto req : requests) {
+      scheduler->add_request(req);
+    }
+
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    const std::vector<uint32_t>& allowed_max_tokens =
+        batch[0].get_allowed_max_tokens();
+    for (size_t i = 0; i < batch[0].size(); ++i) {
+      auto seq = batch[0][i];
+      EXPECT_TRUE(allowed_max_tokens[i] == validate_allowed_max_tokens[idx]);
+    }
+  }
+}
+
+// TEST-2:
+// memory or budget not enough
+TEST(ChunkedPrefillSchedulerTest, ResourceNotEnough) {
+  // case1: max tokens budget not enough
+  {
+    // max token budget: 0
+    ContinuousScheduler::Options opt =
+        create_scheduler_options(1, 256, 0, 1024, 1);
+    auto engine = std::make_unique<FakeEngine>(16, 16);
+    auto scheduler =
+        std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+    EXPECT_TRUE(scheduler != nullptr);
+
+    // request prompt len: 100
+    auto requests =
+        generate_request({1, 100}, {1, 10}, std::nullopt, std::nullopt, 10000);
+    scheduler->add_request(requests[0]);
+    scheduler->add_request(requests[1]);
+
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 1);  // can not schedule the second.
+  }
+
+  // case2: blocks memory not enough
+  {
+    ContinuousScheduler::Options opt =
+        create_scheduler_options(1000, 256, 0, 1024, 1);
+    // free block slot: 1
+    auto engine = std::make_unique<FakeEngine>(2, 8);
+    auto scheduler =
+        std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+    EXPECT_TRUE(scheduler != nullptr);
+
+    // request prompt len: 1000
+    auto requests = generate_request(
+        {1, 1000}, {1, 100}, std::nullopt, std::nullopt, 10000);
+    scheduler->add_request(requests[0]);
+    scheduler->add_request(requests[1]);
+
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 1);  // can not schedule the second.
+  }
+}
+
+// TEST-3:
+// schdule decoding requests + some prefill requests
+TEST(ChunkedPrefillSchedulerTest, NormalSchedule) {
+  // set max free blocks: 512, support 512*32=16384 tokens
+  int block_num = 512;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1);
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  // 1. schedule some new prefill requests
+  auto requests = generate_request(
+      {100, 200, 300}, {10, 20, 30}, std::nullopt, std::nullopt, 30000);
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto total_blocks =
+      util::max(engine->block_manager_pool()->num_free_blocks());
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 3);
+  auto seq_use_blocks =
+      (100 / block_size + 1) + (200 / block_size + 1) + (300 / block_size + 1);
+  EXPECT_TRUE(util::max(engine->block_manager_pool()->num_free_blocks()) ==
+              (total_blocks - seq_use_blocks));
+  update_requests(requests);
+
+  // 2. schedule decoding requets
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 3);
+  update_requests(requests);
+
+  // 3. add new prefill requests
+  auto requests1 =
+      generate_request({400, 500}, {40, 50}, std::nullopt, std::nullopt, 30000);
+  for (auto req : requests1) {
+    scheduler->add_request(req);
+  }
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 5);
+  update_requests(requests);
+  update_requests(requests1);
+
+  // 4. add long new prefill requests,
+  // memory is enough, can handled tokens from
+  // max_tokens_per_chunk_for_prefill(1024) -> remain_max_tokens_budget
+  auto requests2 =
+      generate_request({10000}, {10}, std::nullopt, std::nullopt, 30000);
+  for (auto req : requests2) {
+    scheduler->add_request(req);
+  }
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 6);
+  const std::vector<uint32_t>& allowed_max_tokens =
+      batch[0].get_allowed_max_tokens();
+  // memory is enough, can handled tokens from 1024 -> remain_max_tokens_budget
+  EXPECT_TRUE(allowed_max_tokens[5] == 10000 - 5);
+  update_requests(requests);
+  update_requests(requests1);
+  update_requests(requests2);
+
+  // 5. add long new prefill requests,
+  // memory is not enough, only handled `max_tokens_per_chunk_for_prefill`(1024)
+  // tokens
+  auto requests3 =
+      generate_request({10000}, {10}, std::nullopt, std::nullopt, 30000);
+  for (auto req : requests3) {
+    scheduler->add_request(req);
+  }
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 7);
+  const std::vector<uint32_t>& allowed_max_tokens1 =
+      batch[0].get_allowed_max_tokens();
+  // only can handle max_tokens_per_chunk_for_prefill tokens.
+  EXPECT_TRUE(allowed_max_tokens1[6] == 1024);
+}
+
+// TEST-4:
+// test preempt
+TEST(ChunkedPrefillSchedulerTest, PreemptSchedule) {
+  // set max free blocks: 9, support 9*32=288 tokens
+  // actually only 8 free blocks , because default 1 block is for padding
+  int block_num = 9;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1);
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1. schedule some new prefill requests
+  // request-1 has higher priority than request-2
+  auto requests =
+      generate_request({127, 127}, {10, 10}, std::nullopt, std::nullopt, 30000);
+  running_requests = requests;
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  update_requests(running_requests);
+
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  update_requests(running_requests);
+
+  BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+  int free_blocks_before_preempt =
+      util::max(block_manager_pool->num_free_blocks());
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 1);
+  int free_blocks_after_preempt =
+      util::max(block_manager_pool->num_free_blocks());
+  EXPECT_TRUE(free_blocks_after_preempt > free_blocks_before_preempt);
+  EXPECT_TRUE(scheduler->get_waiting_requests_num() == 1);
+  // append a new block
+  block_manager_pool->allocate(batch[0][0]);
+  // remove preempted request from running_requests
+  running_requests.pop_back();
+  update_requests(running_requests);
+
+  // 2.The preempted request was saved in waiting_priority_queue_,
+  // it just like a new prefill request.
+  // Continue to run, the new prefill request will not preempt
+  // request in running_queue.
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 1);
+}
+
+// TEST-5:
+// test on/offline preempt
+TEST(ChunkedPrefillSchedulerTest, OnDecodePreemptOffDecode) {
+  // set max free blocks: 9, support 9*32=288 tokens
+  // actually only 8 free blocks , because default 1 block is for padding
+  int block_num = 9;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1);
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+  BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+  EXPECT_TRUE(scheduler != nullptr);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1. schedule one online and one prefill prefill requests
+  auto requests = generate_request({127, 127},
+                                   {10, 10},
+                                   std::vector<bool>{true, false},
+                                   std::vector<int32_t>{2, 2},
+                                   30000);
+  running_requests = requests;
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  update_requests(running_requests);
+
+  batch = scheduler->prepare_batch_test();
+
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  update_requests(running_requests);
+
+  int free_blocks_before_preempt =
+      util::max(block_manager_pool->num_free_blocks());
+  // 2. after 2 step, preemption should happen
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 1);
+  int free_blocks_after_preempt =
+      util::max(block_manager_pool->num_free_blocks());
+  EXPECT_TRUE(free_blocks_after_preempt > free_blocks_before_preempt);
+  // check the running request is online request
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 1);
+  EXPECT_TRUE(scheduler->get_running_requests()[0]->offline() == false);
+  EXPECT_TRUE(scheduler->get_waiting_requests_num() == 1);
+}
+
+// TEST-6:
+// test on/offline preempt
+TEST(ChunkedPrefillSchedulerTest, OnPrefillPreemptOffDecode) {
+  // set max free blocks: 9, support 9*32=288 tokens
+  // actually only 8 free blocks , because default 1 block is for padding
+  int block_num = 9;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1);
+  FLAGS_prefill_scheduling_memory_usage_threshold = 2;  // release threshold
+
+  {
+    // 1. two offline decode requests then one online prefill request preempt
+    // them
+    auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+    auto scheduler =
+        std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+    BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+    EXPECT_TRUE(scheduler != nullptr);
+
+    std::vector<std::shared_ptr<Request>> running_requests;
+
+    auto requests = generate_request({100, 100},
+                                     {10, 10},
+                                     std::vector<bool>{true, true},
+                                     std::vector<int32_t>{2, 2},
+                                     30000);
+    running_requests = requests;
+    for (auto req : requests) {
+      scheduler->add_request(req);
+    }
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 0);
+    update_requests(running_requests);
+
+    batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 0);
+    update_requests(running_requests);
+
+    auto new_requests = generate_request({80},
+                                         {10},
+                                         std::vector<bool>{false},
+                                         std::vector<int32_t>{2},
+                                         30000);  // use 3 blocks
+    scheduler->add_request(new_requests[0]);
+    batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+
+    // online prefill request preempt offline decode request
+    EXPECT_TRUE(scheduler->get_running_requests().size() == 2);
+    EXPECT_TRUE(scheduler->get_running_requests()[0]->offline() == false);
+    EXPECT_TRUE(scheduler->get_waiting_requests_num() == 1);
+
+    // offline is evicted
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 1);
+  }
+
+  // 2. another case: longer online prefill request arrives, but can not evict
+  // offline because evicting offline is not enough
+  {
+    auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+    auto scheduler =
+        std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+    BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+    EXPECT_TRUE(scheduler != nullptr);
+
+    std::vector<std::shared_ptr<Request>> running_requests;
+    // 1. schedule one online and one offline
+    auto requests = generate_request({100, 100},
+                                     {10, 10},
+                                     std::vector<bool>{true, false},
+                                     std::vector<int32_t>{2, 2},
+                                     30000);
+    running_requests = requests;
+    for (auto req : requests) {
+      scheduler->add_request(req);
+    }
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 0);
+    update_requests(running_requests);
+
+    auto new_requests = generate_request(
+        {200}, {10}, std::vector<bool>{false}, std::vector<int32_t>{2}, 30000);
+    scheduler->add_request(new_requests[0]);
+    batch = scheduler->prepare_batch_test();
+
+    // 2. online is still waiting
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(scheduler->get_waiting_requests().size() == 1);
+    EXPECT_TRUE(scheduler->get_waiting_requests()[0].get() ==
+                new_requests[0].get());
+  }
+}
+
+// TEST-7:
+// test priority schedule
+TEST(ChunkedPrefillSchedulerTest, PrioritySchedule) {
+  // set max free blocks: 12
+  // actually only 11 free blocks , because default 1 block is for padding
+  int block_num = 12;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1, "priority");
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1: HIGH, 2: NORMAL, 3: LOW
+  auto requests = generate_request({127, 127, 127},
+                                   {10, 10, 10},
+                                   std::vector<bool>{false, false, false},
+                                   std::vector<int32_t>{3, 3, 2},
+                                   30000);
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests()[0]->priority() ==
+              RequestPriority::NORMAL /*NORMAL*/);
+  EXPECT_TRUE(scheduler->get_running_requests()[1]->priority() ==
+              RequestPriority::LOW /*LOW*/);
+
+  running_requests = scheduler->get_running_requests();
+  update_requests(running_requests);
+
+  // new HIGH priority request arrives, its prefill starts
+  auto new_requests = generate_request({32},
+                                       {10},
+                                       std::vector<bool>{false},
+                                       std::vector<int32_t>{1},
+                                       30000);  // use 1 blocks
+  scheduler->add_request(new_requests[0]);
+  batch = scheduler->prepare_batch_test();
+  // check there are 3 running requests owing to decode-maximal
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 3);
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 3);
+  running_requests.push_back(new_requests[0]);
+  update_requests(running_requests);
+  // preemption happens, only HIGH and NORMAL decode requests
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests()[0]->priority() ==
+              RequestPriority::HIGH /*HIGH*/);
+  EXPECT_TRUE(scheduler->get_running_requests()[1]->priority() ==
+              RequestPriority::NORMAL /*NORMAL*/);
+}
+
+// TEST-8:
+// test latency budget
+TEST(ChunkedPrefillSchedulerTest, LatencySchedule) {
+  // set max free blocks: 3
+  // actually only 2 free blocks , because default 1 block is for padding
+  int block_num = 12;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 4;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000,
+                               256,
+                               0,
+                               max_tokens_per_chunk_for_prefill,
+                               1,
+                               "fcfs",
+                               false,
+                               true,
+                               350,
+                               150);
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ChunkedPrefillScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  // mannuly created profile data for y=0.5x^2+10x
+  std::vector<std::pair<int32_t, double>> created_profile_data = {
+      {2, 22}, {4, 48}, {6, 78}, {8, 112}};
+  auto profile_manager = scheduler->get_profile_manager();
+  // fit y=0.5x^2+10x
+  profile_manager->train_prefill_time_predictor(created_profile_data);
+
+  auto requests = generate_request(
+      {10, 10, 10}, {10, 10, 10}, std::nullopt, std::nullopt, 30000);
+  // check if time equation fits well
+  EXPECT_TRUE(
+      static_cast<int32_t>(std::round(profile_manager->predict_step_time(
+          requests[0]->sequences()[0].get(), true, true))) == 150);
+  EXPECT_TRUE(static_cast<int32_t>(std::round(
+                  profile_manager->predict_step_time(2, 0, true, true))) == 22);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1. three requests enter prefill
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_EQ(batch.size(), 1);
+  EXPECT_EQ(batch[0].size(), 3);
+  EXPECT_EQ(scheduler->get_running_requests().size(), 3);
+  running_requests = scheduler->get_running_requests();
+  update_requests(running_requests);
+
+  // 2. three requests enter decode, one new request chunked prefill
+  auto new_requests = generate_request(
+      {10}, {10}, std::nullopt, std::nullopt, 30000);  // use 1 blocks
+  scheduler->add_request(new_requests[0]);
+  batch = scheduler->prepare_batch_test();
+  auto running_sequences_budgets = scheduler->get_running_sequences_budgets();
+
+  EXPECT_EQ(batch.size(), 1);
+  // tpot = 150 > 3 * 10 , all requests enter decode. But not enough for extra
+  // whole prefill, then do chunked
+  EXPECT_EQ(batch[0].size(), 4);
+  EXPECT_EQ(scheduler->get_running_requests().size(), 4);
+  EXPECT_EQ(running_sequences_budgets.size(), 4);
+  EXPECT_EQ(running_sequences_budgets[3], max_tokens_per_chunk_for_prefill);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/scheduler/continuous_scheduler_test.cpp

@@ -0,0 +1,752 @@
+#include "continuous_scheduler.h"
+
+#include <absl/time/clock.h>
+#include <gtest/gtest.h>
+
+#include <cmath>
+#include <limits>
+#include <optional>
+
+#include "chunked_prefill_scheduler.h"
+#include "core/common/global_flags.h"
+#include "distributed_runtime/engine.h"
+#include "prefill_only_scheduler.h"
+#include "scheduler_factory.h"
+#include "util/utils.h"
+
+namespace xllm {
+
+namespace {
+class FakeTokenizer : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens = true) const {
+    NOT_IMPLEMENTED();
+  }
+  std::string decode(const Slice<int32_t>& ids,
+                     bool skip_special_tokens) const {
+    NOT_IMPLEMENTED();
+  }
+  std::optional<int32_t> token_to_id(const std::string_view& token) const {
+    NOT_IMPLEMENTED();
+  }
+  std::string id_to_token(int32_t id) const { NOT_IMPLEMENTED(); }
+  size_t vocab_size() const { NOT_IMPLEMENTED(); }
+  std::unique_ptr<Tokenizer> clone() const {
+    return std::make_unique<FakeTokenizer>();
+  }
+};
+
+class FakeEngine : public Engine {
+ public:
+  FakeEngine(int32_t num_blocks,
+             int32_t block_size,
+             bool enable_prefix_cache = false) {
+    BlockManagerPool::Options opt;
+    opt.num_blocks_ = num_blocks;
+    opt.block_size_ = block_size;
+    opt.enable_prefix_cache_ = enable_prefix_cache;
+    fake_tokenizer_ = std::make_unique<FakeTokenizer>();
+    fake_block_manager_ = std::make_unique<BlockManagerPool>(opt, 1);
+  }
+  ForwardOutput step(std::vector<Batch>& batch) { NOT_IMPLEMENTED(); }
+  void update_last_step_result(std::vector<Batch>& batch) { NOT_IMPLEMENTED(); }
+  const Tokenizer* tokenizer() const { return fake_tokenizer_.get(); }
+  BlockManagerPool* block_manager_pool() const {
+    return fake_block_manager_.get();
+  }
+  const ModelArgs& model_args() const { NOT_IMPLEMENTED(); }
+  const TokenizerArgs& tokenizer_args() const { NOT_IMPLEMENTED(); }
+  std::vector<int64_t> get_active_activation_memory() const {
+    NOT_IMPLEMENTED();
+  }
+  bool init() override { return true; }
+
+ private:
+  std::unique_ptr<Tokenizer> fake_tokenizer_;
+  std::unique_ptr<BlockManagerPool> fake_block_manager_;
+};
+
+class ScopedBoolFlagValue {
+ public:
+  ScopedBoolFlagValue(bool& flag, bool value) : flag_(flag), old_(flag) {
+    flag_ = value;
+  }
+
+  ~ScopedBoolFlagValue() { flag_ = old_; }
+
+ private:
+  bool& flag_;
+  bool old_;
+};
+
+ContinuousScheduler::Options create_scheduler_options(
+    int32_t max_tokens_per_batch,
+    int32_t max_seqs_per_batch,
+    int32_t num_speculative_tokens,
+    int32_t max_tokens_per_chunk_for_prefill,
+    int32_t dp_size,
+    const std::string& priority_strategy = "fcfs",
+    bool enable_profile_kv_blocks = true,
+    bool enable_latency_aware_schedule = false,
+    int32_t max_global_ttft_ms = std::numeric_limits<int32_t>::max(),
+    int32_t max_global_tpot_ms = std::numeric_limits<int32_t>::max()) {
+  ContinuousScheduler::Options opt;
+  opt.num_speculative_tokens_ = num_speculative_tokens;
+  opt.max_tokens_per_chunk_for_prefill_ = max_tokens_per_chunk_for_prefill;
+  opt.max_tokens_per_batch_ = max_tokens_per_batch;
+  opt.max_seqs_per_batch_ = max_seqs_per_batch;
+  opt.dp_size_ = dp_size;
+  opt.priority_strategy_ = priority_strategy;
+  opt.enable_profile_kv_blocks_ = enable_profile_kv_blocks;
+  opt.enable_latency_aware_schedule_ = enable_latency_aware_schedule;
+  opt.max_global_ttft_ms_ = max_global_ttft_ms;
+  opt.max_global_tpot_ms_ = max_global_tpot_ms;
+  return opt;
+}
+
+std::vector<std::shared_ptr<Request>> generate_request(
+    const std::vector<int32_t>& prompt_lens,
+    const std::vector<int32_t>& max_tokens,
+    std::optional<std::vector<bool>> offlines,
+    std::optional<std::vector<int32_t>> priorities,
+    std::optional<std::vector<int32_t>> ns,
+    std::optional<std::vector<int32_t>> beam_widths,
+    int32_t max_context_len) {
+  std::vector<std::shared_ptr<Request>> requests;
+  EXPECT_TRUE(prompt_lens.size() == max_tokens.size());
+
+  size_t batch_size = prompt_lens.size();
+  std::vector<bool> offline_vec;
+  std::vector<int32_t> priority_vec;
+  if (offlines.has_value()) {
+    offline_vec = *offlines;
+  } else {
+    offline_vec = std::vector<bool>(batch_size, false);
+  }
+
+  if (priorities.has_value()) {
+    priority_vec = priorities.value();
+  } else {
+    priority_vec = std::vector<int32_t>(
+        batch_size, static_cast<int32_t>(RequestPriority::NORMAL));
+  }
+
+  std::vector<int32_t> n_vec;
+  std::vector<int32_t> beam_width_vec;
+  if (ns.has_value()) {
+    n_vec = *ns;
+  } else {
+    n_vec = std::vector<int32_t>(batch_size, 1);
+  }
+  if (beam_widths.has_value()) {
+    beam_width_vec = *beam_widths;
+  } else {
+    beam_width_vec = std::vector<int32_t>(batch_size, 0);
+  }
+
+  for (size_t i = 0; i < batch_size; ++i) {
+    std::vector<int32_t> prompt_token_ids;
+    prompt_token_ids.resize(prompt_lens[i]);
+    RequestSamplingParam sampling_param;
+    sampling_param.beam_width = beam_width_vec[i];
+    SchedulerParam scheduler_param;
+    scheduler_param.offline = offline_vec[i];
+    scheduler_param.priority = static_cast<RequestPriority>(priority_vec[i]);
+
+    StoppingChecker stopping_checker;
+    stopping_checker.set_max_generated_tokens(max_tokens[i]);
+    stopping_checker.set_max_context_len(max_context_len);
+    stopping_checker.set_ignore_eos(true);
+
+    RequestState req_state("x",
+                           prompt_token_ids,
+                           sampling_param,
+                           scheduler_param,
+                           stopping_checker,
+                           prompt_lens[i] + 30000,
+                           n_vec[i],
+                           1,
+                           false,
+                           false,
+                           false,
+                           false,
+                           false,
+                           nullptr,
+                           nullptr);
+
+    auto request =
+        std::make_shared<Request>("1", "1", "1", std::move(req_state), "1");
+    requests.emplace_back(request);
+  }
+
+  return requests;
+}
+
+std::shared_ptr<Request> generate_request_with_prompt_tokens(
+    const std::vector<int32_t>& prompt_token_ids,
+    int32_t max_tokens,
+    int32_t max_context_len) {
+  RequestSamplingParam sampling_param;
+  SchedulerParam scheduler_param;
+
+  StoppingChecker stopping_checker;
+  stopping_checker.set_max_generated_tokens(max_tokens);
+  stopping_checker.set_max_context_len(max_context_len);
+  stopping_checker.set_ignore_eos(true);
+
+  RequestState req_state("x",
+                         prompt_token_ids,
+                         sampling_param,
+                         scheduler_param,
+                         stopping_checker,
+                         prompt_token_ids.size() + 30000,
+                         1,
+                         1,
+                         false,
+                         false,
+                         false,
+                         false,
+                         false,
+                         nullptr,
+                         nullptr);
+
+  return std::make_shared<Request>("1", "1", "1", std::move(req_state), "1");
+}
+
+// dont not consider speculative decoding.
+void update_requests(std::vector<std::shared_ptr<Request>> requests) {
+  for (auto req : requests) {
+    for (auto& seq : req->sequences()) {
+      if (seq->kv_state().kv_cache_tokens_num() == 0) {
+        seq->kv_state().incr_kv_cache_tokens_num(seq->num_prompt_tokens());
+      } else {
+        seq->kv_state().incr_kv_cache_tokens_num(1);
+      }
+      Token token(1);
+      seq->append_token(token);
+    }
+  }
+}
+
+void make_request_decode_ready(const std::shared_ptr<Request>& request) {
+  for (auto& seq : request->sequences()) {
+    seq->kv_state().set_kv_cache_tokens_num(seq->num_prompt_tokens());
+    Token token(1);
+    seq->append_token(token);
+  }
+}
+
+void set_chunk_kv(const std::shared_ptr<Request>& request, size_t kv_tokens) {
+  for (auto& seq : request->sequences()) {
+    seq->kv_state().set_kv_cache_tokens_num(kv_tokens);
+  }
+}
+
+}  // namespace
+
+TEST(ContinuousSchedulerFactoryTest,
+     ChunkedPrefillWithoutSPUsesChunkedScheduler) {
+  ScopedBoolFlagValue enable_sp(FLAGS_enable_prefill_sp, false);
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000, 256, 0, 1024, 1);
+  opt.enable_chunked_prefill() = true;
+
+  auto engine = std::make_unique<FakeEngine>(32, 32);
+  auto scheduler = create_continuous_scheduler(engine.get(), opt);
+
+  EXPECT_NE(dynamic_cast<ChunkedPrefillScheduler*>(scheduler.get()), nullptr);
+  EXPECT_EQ(dynamic_cast<PrefillOnlyScheduler*>(scheduler.get()), nullptr);
+}
+
+TEST(ContinuousSchedulerFactoryTest,
+     ChunkedPrefillWithSPUsesPrefillOnlyScheduler) {
+  ScopedBoolFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000, 256, 0, 1024, 1);
+  opt.enable_chunked_prefill() = true;
+
+  auto engine = std::make_unique<FakeEngine>(32, 32);
+  auto scheduler = create_continuous_scheduler(engine.get(), opt);
+
+  EXPECT_NE(dynamic_cast<PrefillOnlyScheduler*>(scheduler.get()), nullptr);
+  EXPECT_EQ(dynamic_cast<ChunkedPrefillScheduler*>(scheduler.get()), nullptr);
+}
+
+TEST(ContinuousSchedulerFactoryTest,
+     ChunkedPrefillWithSPAndSpeculativeUsesPrefillOnlyScheduler) {
+  ScopedBoolFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000, 256, 4, 1024, 1);
+  opt.enable_chunked_prefill() = true;
+
+  auto engine = std::make_unique<FakeEngine>(32, 32);
+  auto scheduler = create_continuous_scheduler(engine.get(), opt);
+
+  EXPECT_NE(dynamic_cast<PrefillOnlyScheduler*>(scheduler.get()), nullptr);
+  EXPECT_EQ(dynamic_cast<ChunkedPrefillScheduler*>(scheduler.get()), nullptr);
+}
+
+TEST(ContinuousSchedulerFactoryTest,
+     ChunkedPrefillWithSPDoesNotBuildMixedBatch) {
+  ScopedBoolFlagValue enable_sp(FLAGS_enable_prefill_sp, true);
+  ContinuousScheduler::Options opt = create_scheduler_options(8, 8, 0, 4, 1);
+  opt.enable_chunked_prefill() = true;
+
+  auto engine = std::make_unique<FakeEngine>(32, 32);
+  auto scheduler = create_continuous_scheduler(engine.get(), opt);
+  auto* prefill_only = dynamic_cast<PrefillOnlyScheduler*>(scheduler.get());
+  ASSERT_NE(prefill_only, nullptr);
+
+  auto requests = generate_request({2, 10},
+                                   {8, 8},
+                                   std::nullopt,
+                                   std::nullopt,
+                                   std::nullopt,
+                                   std::nullopt,
+                                   30000);
+  for (auto& req : requests) {
+    prefill_only->add_request(req);
+  }
+
+  auto batches = prefill_only->prepare_batch_test();
+  ASSERT_EQ(batches.size(), 1);
+  ASSERT_EQ(batches[0].size(), 2);
+  const auto& allowed_max_tokens = batches[0].get_allowed_max_tokens();
+  ASSERT_EQ(allowed_max_tokens.size(), 2);
+
+  make_request_decode_ready(requests[0]);
+  set_chunk_kv(requests[1], allowed_max_tokens[1]);
+
+  batches = prefill_only->prepare_batch_test();
+  ASSERT_EQ(batches.size(), 1);
+  ASSERT_EQ(batches[0].size(), 1);
+
+  const auto forward_input =
+      batches[0].prepare_forward_input(1, 0, ModelArgs());
+  EXPECT_TRUE(
+      forward_input.input_params.batch_forward_type.is_chunked_prefill());
+  EXPECT_FALSE(forward_input.input_params.batch_forward_type.is_mixed());
+  EXPECT_EQ(forward_input.input_params.num_sequences, 1);
+  EXPECT_EQ(batches[0].get_allowed_max_tokens()[0],
+            opt.max_tokens_per_chunk_for_prefill());
+}
+
+TEST(SchedulerFactoryTest, DisaggPDChunkedPrefillKind) {
+  ScopedBoolFlagValue use_mix_scheduler(FLAGS_use_mix_scheduler, false);
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000, 256, 2, 1024, 1);
+  opt.enable_disagg_pd() = true;
+  opt.enable_pd_ooc() = false;
+  opt.enable_chunked_prefill() = true;
+
+  EXPECT_EQ(select_scheduler_kind(opt),
+            SchedulerKind::DISAGG_PD_CHUNKED_PREFILL);
+}
+
+TEST(SchedulerFactoryTest, DisaggPDOOCKeepsPDOOCKind) {
+  ScopedBoolFlagValue use_mix_scheduler(FLAGS_use_mix_scheduler, false);
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000, 256, 0, 1024, 1);
+  opt.enable_disagg_pd() = true;
+  opt.enable_pd_ooc() = true;
+  opt.enable_chunked_prefill() = true;
+
+  EXPECT_EQ(select_scheduler_kind(opt), SchedulerKind::PD_OOC);
+}
+
+// TEST-1:
+// test preempt
+TEST(ContinuousSchedulerTest, OnDecodePreemptOffDecode) {
+  // set max free blocks: 9, support 9*32=288 tokens
+  // actually only 8 free blocks , because default 1 block is for padding
+  int block_num = 9;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1);
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ContinuousScheduler>(engine.get(), opt);
+  BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+  EXPECT_TRUE(scheduler != nullptr);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1. schedule two new online prefill requests
+  auto requests = generate_request({127, 127},
+                                   {10, 10},
+                                   std::vector<bool>{true, false},
+                                   std::vector<int32_t>{2, 2},
+                                   std::nullopt,
+                                   std::nullopt,
+                                   30000);
+  running_requests = requests;
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  update_requests(running_requests);
+
+  batch = scheduler->prepare_batch_test();
+
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  update_requests(running_requests);
+
+  int free_blocks_before_preempt =
+      util::max(block_manager_pool->num_free_blocks());
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 1);
+  int free_blocks_after_preempt =
+      util::max(block_manager_pool->num_free_blocks());
+  EXPECT_TRUE(free_blocks_after_preempt > free_blocks_before_preempt);
+
+  // check the running request is online request
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 1);
+  EXPECT_TRUE(scheduler->get_running_requests()[0]->offline() == false);
+  EXPECT_TRUE(scheduler->get_waiting_requests_num() == 1);
+}
+
+// TEST-2:
+// test preempt
+TEST(ContinuousSchedulerTest, OnPrefillPreemptOffDecode) {
+  // set max free blocks: 9, support 9*32=288 tokens
+  // actually only 8 free blocks , because default 1 block is for padding
+  int block_num = 9;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1);
+  FLAGS_prefill_scheduling_memory_usage_threshold = 2;  // release threshold
+
+  {
+    // 1. two offline decode requests then one online prefill request
+    // preempt them
+    auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+    auto scheduler = std::make_unique<ContinuousScheduler>(engine.get(), opt);
+    BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+    EXPECT_TRUE(scheduler != nullptr);
+
+    std::vector<std::shared_ptr<Request>> running_requests;
+
+    auto requests = generate_request({100, 100},
+                                     {10, 10},
+                                     std::vector<bool>{true, true},
+                                     std::vector<int32_t>{2, 2},
+                                     std::nullopt,
+                                     std::nullopt,
+                                     30000);
+    running_requests = requests;
+    for (auto req : requests) {
+      scheduler->add_request(req);
+    }
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 0);
+    update_requests(running_requests);
+
+    batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 0);
+    update_requests(running_requests);
+
+    auto new_requests = generate_request({80},
+                                         {10},
+                                         std::vector<bool>{false},
+                                         std::vector<int32_t>{2},
+                                         std::nullopt,
+                                         std::nullopt,
+                                         30000);  // use 3 blocks
+    scheduler->add_request(new_requests[0]);
+    batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 1);
+
+    // online prefill request preempt offline decode request
+    EXPECT_TRUE(scheduler->get_running_requests().size() == 1);
+    EXPECT_TRUE(scheduler->get_running_requests()[0]->offline() == false);
+    EXPECT_TRUE(scheduler->get_waiting_requests_num() == 1);
+
+    // offline is evicted
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 1);
+  }
+
+  // 2. another case: longer online prefill request arrives, but can not
+  // evict offline because evicting offline is not enough
+  {
+    auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+    auto scheduler = std::make_unique<ContinuousScheduler>(engine.get(), opt);
+    BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+    EXPECT_TRUE(scheduler != nullptr);
+
+    std::vector<std::shared_ptr<Request>> running_requests;
+    // one online, one offline
+    auto requests = generate_request({100, 100},
+                                     {10, 10},
+                                     std::vector<bool>{true, false},
+                                     std::vector<int32_t>{2, 2},
+                                     std::nullopt,
+                                     std::nullopt,
+                                     30000);
+    running_requests = requests;
+    for (auto req : requests) {
+      scheduler->add_request(req);
+    }
+    auto batch = scheduler->prepare_batch_test();
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(util::max(block_manager_pool->num_free_blocks()) == 0);
+    update_requests(running_requests);
+
+    auto new_requests = generate_request({200},
+                                         {10},
+                                         std::vector<bool>{false},
+                                         std::vector<int32_t>{2},
+                                         std::nullopt,
+                                         std::nullopt,
+                                         30000);
+    scheduler->add_request(new_requests[0]);
+    batch = scheduler->prepare_batch_test();
+    // online is still waiting
+    EXPECT_TRUE(batch.size() == 1);
+    EXPECT_TRUE(batch[0].size() == 2);
+    EXPECT_TRUE(scheduler->get_waiting_requests().size() == 1);
+    EXPECT_TRUE(scheduler->get_waiting_requests()[0].get() ==
+                new_requests[0].get());
+  }
+}
+
+// TEST-3:
+// test priority schedule
+TEST(ContinuousSchedulerTest, PrioritySchedule) {
+  // set max free blocks: 12
+  // actually only 11 free blocks , because default 1 block is for padding
+  int block_num = 12;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt = create_scheduler_options(
+      10000, 256, 0, max_tokens_per_chunk_for_prefill, 1, "priority");
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ContinuousScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1: HIGH, 2: NORMAL, 3: LOW
+  auto requests = generate_request({128, 128, 128},
+                                   {10, 10, 10},
+                                   std::vector<bool>{false, false, false},
+                                   std::vector<int32_t>{3, 3, 2},
+                                   std::nullopt,
+                                   std::nullopt,
+                                   30000);
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests()[0]->priority() ==
+              RequestPriority::NORMAL /*NORMAL*/);
+  EXPECT_TRUE(scheduler->get_running_requests()[1]->priority() ==
+              RequestPriority::LOW /*LOW*/);
+  running_requests = scheduler->get_running_requests();
+  update_requests(running_requests);
+
+  // new HIGH priority request arrives, its prefill starts
+  auto new_requests = generate_request({32},
+                                       {10},
+                                       std::vector<bool>{false},
+                                       std::vector<int32_t>{1},
+                                       std::nullopt,
+                                       std::nullopt,
+                                       30000);  // use 1 blocks
+  scheduler->add_request(new_requests[0]);
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 1);
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 1);
+  update_requests(new_requests);
+
+  // only HIGH and NORMAL requests decode
+  batch = scheduler->prepare_batch_test();
+  EXPECT_TRUE(batch.size() == 1);
+  EXPECT_TRUE(batch[0].size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests().size() == 2);
+  EXPECT_TRUE(scheduler->get_running_requests()[0]->priority() ==
+              RequestPriority::HIGH /*HIGH*/);
+  EXPECT_TRUE(scheduler->get_running_requests()[1]->priority() ==
+              RequestPriority::NORMAL /*NORMAL*/);
+}
+
+// TEST-4:
+// beam strict mode should not partially schedule one request.
+TEST(ContinuousSchedulerTest, BeamStrictNoPartialScheduling) {
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(2, 8, 0, 1024, 1, "fcfs");
+  auto engine = std::make_unique<FakeEngine>(64, 32);
+  auto scheduler = std::make_unique<ContinuousScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  auto req1 = generate_request({64},
+                               {10},
+                               std::nullopt,
+                               std::nullopt,
+                               std::vector<int32_t>{1},
+                               std::vector<int32_t>{0},
+                               30000)[0];
+  make_request_decode_ready(req1);
+
+  auto beam_req = generate_request({64},
+                                   {10},
+                                   std::nullopt,
+                                   std::nullopt,
+                                   std::vector<int32_t>{1},
+                                   std::vector<int32_t>{2},
+                                   30000)[0];
+  // Manually duplicate beam sequence to simulate a decoded beam group.
+  beam_req->sequences().emplace_back(
+      std::make_unique<Sequence>(*beam_req->sequences()[0]));
+  ASSERT_EQ(beam_req->sequences().size(), 2u);
+  make_request_decode_ready(beam_req);
+
+  scheduler->add_request(req1);
+  scheduler->add_request(beam_req);
+  auto batch = scheduler->prepare_batch_test();
+  ASSERT_EQ(batch.size(), 1u);
+  EXPECT_EQ(batch[0].size(), 1u);
+  EXPECT_EQ(batch[0][0], req1->sequences()[0].get());
+  EXPECT_EQ(scheduler->get_running_requests().size(), 1u);
+  EXPECT_EQ(scheduler->get_waiting_requests_num(), 0u);
+  EXPECT_NE(batch[0][0], beam_req->sequences()[0].get());
+  EXPECT_NE(batch[0][0], beam_req->sequences()[1].get());
+}
+
+// TEST-5:
+// test latency budget
+TEST(ContinuousSchedulerTest, LatencySchedule) {
+  // block is enough
+  int block_num = 12;
+  int block_size = 32;
+  int max_tokens_per_chunk_for_prefill = 1024;
+  // set chunked max_tokens budgets 10000 per step
+  ContinuousScheduler::Options opt =
+      create_scheduler_options(10000,
+                               256,
+                               0,
+                               max_tokens_per_chunk_for_prefill,
+                               1,
+                               "fcfs",
+                               false,
+                               true,
+                               350,
+                               25);
+  auto engine = std::make_unique<FakeEngine>(block_num, block_size);
+  auto scheduler = std::make_unique<ContinuousScheduler>(engine.get(), opt);
+  EXPECT_TRUE(scheduler != nullptr);
+
+  // mannuly created profile data for y=0.5x^2+10x
+  std::vector<std::pair<int32_t, double>> created_profile_data = {
+      {2, 22}, {4, 48}, {6, 78}, {8, 112}};
+  auto profile_manager = scheduler->get_profile_manager();
+  // fit y=0.5x^2+10x
+  profile_manager->train_prefill_time_predictor(created_profile_data);
+
+  auto requests = generate_request({10, 10, 10},
+                                   {10, 10, 10},
+                                   std::nullopt,
+                                   std::nullopt,
+                                   std::nullopt,
+                                   std::nullopt,
+                                   30000);
+  // check if time equation fits well
+  EXPECT_TRUE(
+      static_cast<int32_t>(std::round(profile_manager->predict_step_time(
+          requests[0]->sequences()[0].get(), true, true))) == 150);
+  EXPECT_TRUE(static_cast<int32_t>(std::round(
+                  profile_manager->predict_step_time(2, 0, true, true))) == 22);
+
+  std::vector<std::shared_ptr<Request>> running_requests;
+
+  // 1. two requests enter prefill
+  for (auto req : requests) {
+    scheduler->add_request(req);
+  }
+  auto batch = scheduler->prepare_batch_test();
+
+  EXPECT_EQ(batch.size(), 1);
+  // 2*150 < ttft_slo=350 < 3 * 150, only two requests enter prefill
+  EXPECT_EQ(batch[0].size(), 2);
+  EXPECT_EQ(scheduler->get_running_requests().size(), 2);
+  running_requests = scheduler->get_running_requests();
+  update_requests(running_requests);
+
+  // 2. one request enter prefill
+  batch = scheduler->prepare_batch_test();
+  EXPECT_EQ(batch.size(), 1);
+  EXPECT_EQ(batch[0].size(), 1);
+  EXPECT_EQ(scheduler->get_running_requests().size(), 1);
+  running_requests = scheduler->get_running_requests();
+  update_requests(running_requests);
+
+  // 3. two requests start decode
+  // batch = scheduler->prepare_batch_test();
+  // EXPECT_TRUE(batch.size() == 1);
+  // // 2*10 < tpot_slo=25 < 3 * 10, only two requests enter decode
+  // EXPECT_TRUE(batch[0].size() == 2);
+  // EXPECT_TRUE(scheduler->get_running_requests().size() == 2);
+}
+
+TEST(BlockManagerPoolTest, AllocateFailureRollsBackSharedPrefixBlocks) {
+  auto engine = std::make_unique<FakeEngine>(3, 4, true);
+  BlockManagerPool* block_manager_pool = engine->block_manager_pool();
+
+  auto cached_request =
+      generate_request_with_prompt_tokens({1, 2, 3, 4, 5, 6, 7, 8}, 1, 30000);
+  auto failed_request = generate_request_with_prompt_tokens(
+      {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, 1, 30000);
+  auto later_request =
+      generate_request_with_prompt_tokens({20, 21, 22, 23}, 1, 30000);
+
+  auto* cached_sequence = cached_request->sequences()[0].get();
+  ASSERT_TRUE(block_manager_pool->allocate(cached_sequence,
+                                           cached_sequence->num_tokens()));
+  cached_sequence->kv_state().set_kv_cache_tokens_num(
+      cached_sequence->num_tokens());
+  block_manager_pool->deallocate(cached_sequence);
+
+  const size_t free_blocks_before_failure =
+      util::max(block_manager_pool->num_free_blocks());
+  const size_t used_blocks_before_failure =
+      util::min(block_manager_pool->num_used_blocks());
+  EXPECT_EQ(free_blocks_before_failure, 0);
+
+  auto* failed_sequence = failed_request->sequences()[0].get();
+  EXPECT_FALSE(block_manager_pool->allocate(failed_sequence,
+                                            failed_sequence->num_tokens()));
+  EXPECT_EQ(failed_sequence->kv_state().num_kv_blocks(), 0);
+  EXPECT_EQ(failed_sequence->kv_state().shared_kv_blocks_num(), 0);
+  EXPECT_EQ(util::max(block_manager_pool->num_free_blocks()),
+            free_blocks_before_failure);
+  EXPECT_EQ(util::min(block_manager_pool->num_used_blocks()),
+            used_blocks_before_failure);
+
+  auto* later_sequence = later_request->sequences()[0].get();
+  EXPECT_TRUE(block_manager_pool->allocate(later_sequence,
+                                           later_sequence->num_tokens()));
+  EXPECT_EQ(later_sequence->kv_state().num_kv_blocks(), 1);
+
+  (void)engine.release();
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp

@@ -0,0 +1,40 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "scheduler/disagg_pd_chunked_prefill_scheduler.h"
+
+#include <gtest/gtest.h>
+
+namespace xllm {
+
+TEST(DisaggPDChunkedPrefillSchedulerTest, PicksCurrentChunkBudget) {
+  const PDChunkBudget budget = pick_pd_chunk_budget(32, 96, 40, 64);
+  EXPECT_EQ(budget.next_tokens, 40);
+  EXPECT_EQ(budget.max_tokens, 72);
+}
+
+TEST(DisaggPDChunkedPrefillSchedulerTest, LastPromptChunkStopsAtPromptEnd) {
+  const PDChunkBudget budget = pick_pd_chunk_budget(80, 96, 40, 64);
+  EXPECT_EQ(budget.next_tokens, 16);
+  EXPECT_EQ(budget.max_tokens, 96);
+}
+
+TEST(DisaggPDChunkedPrefillSchedulerTest, EmptyBudgetRejectsSchedule) {
+  const PDChunkBudget budget = pick_pd_chunk_budget(32, 96, 40, 0);
+  EXPECT_EQ(budget.next_tokens, 0);
+  EXPECT_EQ(budget.max_tokens, 32);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/scheduler/fixed_steps_scheduler_test.cpp

@@ -0,0 +1,231 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "fixed_steps_scheduler.h"
+
+#include <absl/time/time.h>
+#include <gtest/gtest.h>
+
+#include <algorithm>
+
+#include "common/global_flags.h"
+#include "continuous_scheduler.h"
+#include "distributed_runtime/engine.h"
+#include "framework/request/rec_type.h"
+
+namespace xllm {
+
+namespace {
+
+class FakeTokenizer : public Tokenizer {
+ public:
+  bool encode(const std::string_view& text,
+              std::vector<int32_t>* ids,
+              bool add_special_tokens = true) const override {
+    (void)text;
+    (void)ids;
+    (void)add_special_tokens;
+    return false;
+  }
+  std::string decode(const Slice<int32_t>& ids,
+                     bool skip_special_tokens) const override {
+    (void)ids;
+    (void)skip_special_tokens;
+    return "";
+  }
+  std::optional<int32_t> token_to_id(
+      const std::string_view& token) const override {
+    (void)token;
+    return std::nullopt;
+  }
+  std::string id_to_token(int32_t id) const override {
+    (void)id;
+    return "";
+  }
+  size_t vocab_size() const override { return 0; }
+  std::unique_ptr<Tokenizer> clone() const override {
+    return std::make_unique<FakeTokenizer>();
+  }
+};
+
+class FakeEngine : public Engine {
+ public:
+  FakeEngine(int32_t num_blocks, int32_t block_size) {
+    BlockManagerPool::Options opt;
+    opt.num_blocks_ = num_blocks;
+    opt.block_size_ = block_size;
+    opt.enable_prefix_cache_ = false;
+    fake_tokenizer_ = std::make_unique<FakeTokenizer>();
+    fake_block_manager_ = std::make_unique<BlockManagerPool>(opt, 1);
+  }
+  ForwardOutput step(std::vector<Batch>& batch) override {
+    (void)batch;
+    return ForwardOutput();
+  }
+  void update_last_step_result(std::vector<Batch>& batch) override {
+    (void)batch;
+  }
+  const Tokenizer* tokenizer() const override { return fake_tokenizer_.get(); }
+  BlockManagerPool* block_manager_pool() const override {
+    return fake_block_manager_.get();
+  }
+  const ModelArgs& model_args() const override {
+    static ModelArgs args;
+    return args;
+  }
+  const TokenizerArgs& tokenizer_args() const override {
+    static TokenizerArgs args;
+    return args;
+  }
+  std::vector<int64_t> get_active_activation_memory() const override {
+    return {};
+  }
+  bool init() override { return true; }
+
+ private:
+  std::unique_ptr<Tokenizer> fake_tokenizer_;
+  std::unique_ptr<BlockManagerPool> fake_block_manager_;
+};
+
+ContinuousScheduler::Options CreateOptions(
+    int32_t max_tokens_per_batch = 10000,
+    int32_t max_seqs_per_batch = 256,
+    int32_t dp_size = 1,
+    bool enable_schedule_overlap = false,
+    int32_t rec_worker_max_concurrency = 1) {
+  ContinuousScheduler::Options opt;
+  opt.max_tokens_per_batch_ = max_tokens_per_batch;
+  opt.max_seqs_per_batch_ = max_seqs_per_batch;
+  opt.dp_size_ = dp_size;
+  opt.enable_schedule_overlap_ = enable_schedule_overlap;
+  opt.rec_worker_max_concurrency_ = rec_worker_max_concurrency;
+  opt.max_tokens_per_chunk_for_prefill_ = 1024;
+  opt.num_speculative_tokens_ = 0;
+  return opt;
+}
+
+std::vector<std::shared_ptr<Request>> GenRequests(
+    const std::vector<int32_t>& prompt_lens,
+    const std::vector<int32_t>& max_tokens,
+    RecType rec_type,
+    int32_t max_context_len = 30000) {
+  std::vector<std::shared_ptr<Request>> requests;
+  EXPECT_EQ(prompt_lens.size(), max_tokens.size());
+  for (size_t i = 0; i < prompt_lens.size(); ++i) {
+    std::vector<int32_t> prompt_token_ids(prompt_lens[i], 0);
+    RequestSamplingParam sampling_param;
+    SchedulerParam scheduler_param;
+    scheduler_param.offline = false;
+    scheduler_param.priority = RequestPriority::NORMAL;
+    StoppingChecker stopping_checker;
+    stopping_checker.set_max_generated_tokens(max_tokens[i]);
+    stopping_checker.set_max_context_len(max_context_len);
+    stopping_checker.set_ignore_eos(true);
+    RequestState req_state("x",
+                           prompt_token_ids,
+                           sampling_param,
+                           scheduler_param,
+                           stopping_checker,
+                           static_cast<size_t>(prompt_lens[i]) + 30000,
+                           1,
+                           1,
+                           false,
+                           false,
+                           false,
+                           false,
+                           false,
+                           nullptr,
+                           nullptr);
+    req_state.rec_type = rec_type;
+    auto request =
+        std::make_shared<Request>("1", "1", "1", std::move(req_state), "1");
+    requests.emplace_back(request);
+  }
+  return requests;
+}
+
+}  // namespace
+
+TEST(FixedStepsSchedulerTest, AddRequestSuccess) {
+  auto engine = std::make_unique<FakeEngine>(32, 32);
+  auto opt = CreateOptions();
+  FixedStepsScheduler scheduler(engine.get(), opt);
+  auto requests = GenRequests({64}, {10}, RecType::kOneRec);
+  std::shared_ptr<Request> req = requests[0];
+  EXPECT_TRUE(scheduler.add_request(req));
+}
+
+TEST(FixedStepsSchedulerTest, PrepareBatchEmptyWhenNoRequests) {
+  FLAGS_enable_prefix_cache = false;
+  auto engine = std::make_unique<FakeEngine>(32, 32);
+  auto opt = CreateOptions();
+  FixedStepsScheduler scheduler(engine.get(), opt);
+  ContinuousScheduler* base = &scheduler;
+  std::vector<Batch> batches = base->prepare_batch_test();
+  EXPECT_FALSE(batches.empty());
+  EXPECT_TRUE(batches[0].empty());
+}
+
+TEST(FixedStepsSchedulerTest, PrepareBatchOneRecSchedulesRequest) {
+  FLAGS_enable_prefix_cache = false;
+  FLAGS_prefill_scheduling_memory_usage_threshold = 1.0;
+  auto engine = std::make_unique<FakeEngine>(64, 32);
+  auto opt = CreateOptions(10000, 256);
+  FixedStepsScheduler scheduler(engine.get(), opt);
+  auto requests = GenRequests({64, 64}, {10, 10}, RecType::kOneRec);
+  for (auto& req : requests) {
+    scheduler.add_request(req);
+  }
+  ContinuousScheduler* base = &scheduler;
+  std::vector<Batch> batches = base->prepare_batch_test();
+  EXPECT_FALSE(batches.empty());
+  bool has_non_empty = false;
+  for (const auto& b : batches) {
+    if (!b.empty()) {
+      has_non_empty = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(has_non_empty);
+  EXPECT_EQ(base->get_running_requests().size(), 2u);
+}
+
+TEST(FixedStepsSchedulerTest, PrepareBatchRespectsTokenBudget) {
+  FLAGS_enable_prefix_cache = false;
+  FLAGS_prefill_scheduling_memory_usage_threshold = 1.0;
+  auto engine = std::make_unique<FakeEngine>(64, 32);
+  auto opt = CreateOptions(50, 1);
+  FixedStepsScheduler scheduler(engine.get(), opt);
+  auto requests = GenRequests({40, 40}, {10, 10}, RecType::kOneRec);
+  for (auto& req : requests) {
+    scheduler.add_request(req);
+  }
+  ContinuousScheduler* base = &scheduler;
+  base->prepare_batch_test();
+  EXPECT_LE(base->get_running_requests().size(), 1u);
+}
+
+TEST(FixedStepsSchedulerTest, StepCompletesWithRequest) {
+  FLAGS_enable_prefix_cache = false;
+  FLAGS_prefill_scheduling_memory_usage_threshold = 1.0;
+  auto engine = std::make_unique<FakeEngine>(64, 32);
+  auto opt = CreateOptions(10000, 256);
+  FixedStepsScheduler scheduler(engine.get(), opt);
+  auto requests = GenRequests({32}, {10}, RecType::kOneRec);
+  scheduler.add_request(requests[0]);
+  EXPECT_NO_THROW(scheduler.step(absl::Milliseconds(500)));
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/util/CMakeLists.txt

@@ -0,0 +1,18 @@
+include(cc_test)
+
+cc_test(
+  NAME
+    util_test
+  SRCS
+    blocking_counter_test.cpp
+    suffix_decoding_cache_test.cpp
+    threadpool_test.cpp
+  DEPS
+    util
+    absl::synchronization
+    absl::time
+    GTest::gtest_main
+    gflags::gflags
+)
+target_link_libraries(util_test PRIVATE brpc leveldb::leveldb OpenSSL::SSL OpenSSL::Crypto)
+add_dependencies(util_test brpc-static)

upstream_ref/xllm/tests/core/util/blocking_counter_test.cpp

@@ -0,0 +1,92 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "blocking_counter.h"
+
+#include <gtest/gtest.h>
+
+#include "util/threadpool.h"
+
+namespace xllm {
+
+TEST(BlockingCounterTest, BasicTest) {
+  BlockingCounter counter(1);
+  counter.decrement_count();
+  counter.wait();
+  EXPECT_TRUE(true);
+}
+
+TEST(BlockingCounterTest, TwoThreadTest) {
+  ThreadPool threadpool(1);
+  BlockingCounter counter(2);
+
+  int called = 0;
+  threadpool.schedule([&counter, &called]() {
+    counter.decrement_count();
+    ++called;
+  });
+  counter.decrement_count();
+  ++called;
+  counter.wait();
+  EXPECT_EQ(2, called);
+}
+
+TEST(BlockingCounterTest, MultiThreadTest) {
+  ThreadPool threadpool(4);
+  BlockingCounter counter(5);
+
+  int called = 0;
+  threadpool.schedule([&counter, &called]() {
+    counter.decrement_count();
+    ++called;
+  });
+  threadpool.schedule([&counter, &called]() {
+    counter.decrement_count();
+    ++called;
+  });
+  threadpool.schedule([&counter, &called]() {
+    counter.decrement_count();
+    ++called;
+  });
+  threadpool.schedule([&counter, &called]() {
+    counter.decrement_count();
+    ++called;
+  });
+  counter.decrement_count();
+  ++called;
+
+  counter.wait();
+  EXPECT_EQ(5, called);
+}
+
+TEST(BlockingCounterTest, WaitTimeoutTest) {
+  ThreadPool threadpool(2);
+  BlockingCounter counter(3);
+
+  int called = 0;
+  threadpool.schedule([&counter, &called]() {
+    counter.decrement_count();
+    ++called;
+  });
+
+  counter.decrement_count();
+  ++called;
+
+  const std::chrono::milliseconds timeout(100);
+  counter.wait_for(timeout);
+  EXPECT_EQ(2, called);
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/util/suffix_decoding_cache_test.cpp

@@ -0,0 +1,143 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+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/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 "suffix_decoding_cache.h"
+
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <vector>
+
+namespace xllm {
+
+TEST(SuffixDecodingCacheTest, StartStopAndActiveSet) {
+  SuffixDecodingCache cache(/*max_tree_depth=*/16, /*max_cached_requests=*/4);
+
+  std::vector<int32_t> prompt = {10, 11, 12};
+  cache.start_request("req_a", std::span<const int32_t>(prompt));
+
+  EXPECT_TRUE(cache.has_active_request("req_a"));
+  EXPECT_TRUE(cache.has_cached_request("req_a"));
+
+  auto active = cache.active_requests();
+  EXPECT_EQ(active.size(), 1);
+  EXPECT_EQ(active[0], "req_a");
+
+  cache.stop_request("req_a");
+  EXPECT_FALSE(cache.has_active_request("req_a"));
+  EXPECT_TRUE(cache.has_cached_request("req_a"));
+}
+
+TEST(SuffixDecodingCacheTest, AddResponseAndSpeculate) {
+  SuffixDecodingCache cache(/*max_tree_depth=*/32, /*max_cached_requests=*/8);
+
+  std::vector<int32_t> prompt = {1, 2, 3};
+  cache.start_request("req_a", std::span<const int32_t>(prompt));
+
+  std::vector<int32_t> out = {4, 5, 6, 7};
+  cache.add_active_response("req_a", std::span<const int32_t>(out));
+
+  std::vector<int32_t> ctx = {2, 3, 4};
+  SuffixDecodingDraft draft = cache.speculate("req_a",
+                                              std::span<const int32_t>(ctx),
+                                              /*max_spec_tokens=*/4,
+                                              /*max_spec_factor=*/2.0f,
+                                              /*max_spec_offset=*/0.0f,
+                                              /*min_token_prob=*/0.01f,
+                                              /*use_tree_spec=*/false);
+
+  EXPECT_FALSE(draft.token_ids.empty());
+  EXPECT_EQ(draft.token_ids[0], 5);
+}
+
+TEST(SuffixDecodingCacheTest, EvictionByMaxCachedRequests) {
+  SuffixDecodingCache cache(/*max_tree_depth=*/16, /*max_cached_requests=*/1);
+
+  std::vector<int32_t> p0 = {1, 2};
+  std::vector<int32_t> p1 = {3, 4};
+  cache.start_request("req_a", std::span<const int32_t>(p0));
+  cache.stop_request("req_a");
+
+  cache.start_request("req_b", std::span<const int32_t>(p1));
+
+  auto cached = cache.cached_requests();
+  EXPECT_LE(cached.size(), 1);
+  EXPECT_TRUE(cache.has_cached_request("req_b"));
+}
+
+TEST(SuffixDecodingCacheTest, AddPromptOnlyDoesNotPolluteGlobalCache) {
+  SuffixDecodingCache cache(/*max_tree_depth=*/32, /*max_cached_requests=*/8);
+
+  std::vector<int32_t> p0 = {10, 20, 30};
+  cache.start_request("req_a", std::span<const int32_t>(p0));
+
+  std::vector<int32_t> prompt_tail = {40, 50, 60};
+  cache.add_active_prompt("req_a", std::span<const int32_t>(prompt_tail));
+  cache.stop_request("req_a");
+
+  std::vector<int32_t> p1 = {1, 2, 10, 20, 30, 40};
+  cache.start_request("req_b", std::span<const int32_t>(p1));
+
+  std::vector<int32_t> ctx = {20, 30, 40};
+  auto draft = cache.speculate("req_b",
+                               std::span<const int32_t>(ctx),
+                               /*max_spec_tokens=*/3,
+                               /*max_spec_factor=*/2.0f,
+                               /*max_spec_offset=*/0.0f,
+                               /*min_token_prob=*/0.01f,
+                               /*use_tree_spec=*/false);
+
+  EXPECT_TRUE(draft.token_ids.empty());
+}
+
+TEST(SuffixDecodingCacheTest, GlobalCacheSpeculateAcrossRequests) {
+  SuffixDecodingCache cache(/*max_tree_depth=*/32, /*max_cached_requests=*/8);
+
+  std::vector<int32_t> p0 = {10, 20, 30};
+  cache.start_request("req_a", std::span<const int32_t>(p0));
+  std::vector<int32_t> out0 = {40, 50, 60};
+  cache.add_active_response("req_a", std::span<const int32_t>(out0));
+  cache.stop_request("req_a");
+
+  std::vector<int32_t> p1 = {1, 2, 10, 20, 30, 40};
+  cache.start_request("req_b", std::span<const int32_t>(p1));
+
+  std::vector<int32_t> ctx = {20, 30, 40};
+  auto draft = cache.speculate("req_b",
+                               std::span<const int32_t>(ctx),
+                               /*max_spec_tokens=*/3,
+                               /*max_spec_factor=*/2.0f,
+                               /*max_spec_offset=*/0.0f,
+                               /*min_token_prob=*/0.01f,
+                               /*use_tree_spec=*/false);
+
+  EXPECT_FALSE(draft.token_ids.empty());
+  EXPECT_EQ(draft.token_ids[0], 50);
+}
+
+TEST(SuffixDecodingCacheTest, MaxCachedRequestsZeroDisablesGlobalCache) {
+  SuffixDecodingCache cache(/*max_tree_depth=*/16, /*max_cached_requests=*/0);
+
+  std::vector<int32_t> p0 = {1, 2, 3};
+  cache.start_request("req_a", std::span<const int32_t>(p0));
+
+  EXPECT_TRUE(cache.has_active_request("req_a"));
+  EXPECT_FALSE(cache.has_cached_request("req_a"));
+
+  auto cached = cache.cached_requests();
+  EXPECT_TRUE(cached.empty());
+}
+
+}  // namespace xllm

upstream_ref/xllm/tests/core/util/threadpool_test.cpp

@@ -0,0 +1,172 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+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/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 "threadpool.h"
+
+#include <absl/synchronization/notification.h>
+#include <absl/time/clock.h>
+#include <gtest/gtest.h>
+#include <pthread.h>
+#include <sched.h>
+
+namespace xllm {
+
+TEST(ThreadPoolTest, ScheduleEmptyTask) {
+  ThreadPool threadpool(1);
+  absl::Notification notification;
+  threadpool.schedule(nullptr);
+}
+
+TEST(ThreadPoolTest, ScheduleTask) {
+  ThreadPool threadpool(1);
+  absl::Notification notification;
+  bool called = false;
+  threadpool.schedule([&called, &notification]() {
+    called = true;
+    notification.Notify();
+  });
+  notification.WaitForNotification();
+  EXPECT_TRUE(called);
+}
+
+TEST(ThreadPoolTest, ScheduleMultipleTasks) {
+  ThreadPool threadpool(1);
+  std::vector<std::string> completed_tasks;
+  absl::Notification notification;
+
+  // run frist task
+  threadpool.schedule([&completed_tasks, &notification]() {
+    completed_tasks.emplace_back("first");
+    if (completed_tasks.size() == 2) {
+      absl::SleepFor(absl::Milliseconds(100));
+      notification.Notify();
+    }
+  });
+
+  // run second task
+  threadpool.schedule([&completed_tasks, &notification]() {
+    completed_tasks.emplace_back("second");
+    if (completed_tasks.size() == 2) {
+      notification.Notify();
+    }
+  });
+
+  notification.WaitForNotificationWithTimeout(absl::Milliseconds(200));
+  EXPECT_EQ(completed_tasks.size(), 2);
+  EXPECT_EQ(completed_tasks[0], "first");
+  EXPECT_EQ(completed_tasks[1], "second");
+}
+
+TEST(ThreadPoolTest, MultipleThreads) {
+  ThreadPool threadpool(4);
+  std::atomic_uint32_t counter = 0;
+  absl::Notification notification;
+
+  for (int i = 0; i < 10; ++i) {
+    threadpool.schedule([&counter, &notification]() {
+      absl::SleepFor(absl::Milliseconds(100));
+      counter++;
+      if (counter == 10) {
+        notification.Notify();
+      }
+    });
+  }
+
+  EXPECT_TRUE(
+      notification.WaitForNotificationWithTimeout(absl::Milliseconds(400)));
+  EXPECT_EQ(counter, 10);
+}
+
+TEST(ThreadPoolTest, CpuCoreBindingConstructor) {
+  // Construct with cpu_cores binding — should not crash even if binding fails
+  // (e.g., in containers with restricted affinity).
+  std::vector<int32_t> cpu_cores = {0, 0};  // bind both threads to core 0
+  ThreadPool threadpool(2, cpu_cores);
+  EXPECT_EQ(threadpool.size(), 2);
+
+  std::atomic<int> counter{0};
+  absl::Notification notification;
+  for (int i = 0; i < 2; ++i) {
+    threadpool.schedule([&counter, &notification]() {
+      if (++counter == 2) {
+        notification.Notify();
+      }
+    });
+  }
+  EXPECT_TRUE(
+      notification.WaitForNotificationWithTimeout(absl::Milliseconds(500)));
+  EXPECT_EQ(counter, 2);
+}
+
+TEST(ThreadPoolTest, CpuCoreBindingWithInitFunc) {
+  std::vector<int32_t> cpu_cores = {0};
+  std::atomic<bool> init_called{false};
+  absl::Notification init_done;
+  ThreadPool threadpool(
+      1,
+      [&init_called, &init_done]() {
+        init_called = true;
+        init_done.Notify();
+      },
+      cpu_cores);
+  EXPECT_TRUE(
+      init_done.WaitForNotificationWithTimeout(absl::Milliseconds(500)));
+  EXPECT_TRUE(init_called);
+}
+
+TEST(ThreadPoolTest, CpuCoreBindingMismatchFallback) {
+  // Mismatched cpu_cores size — should fall back to no binding gracefully.
+  std::vector<int32_t> cpu_cores = {0, 1};  // 2 cores but 4 threads
+  ThreadPool threadpool(4, cpu_cores);
+  EXPECT_EQ(threadpool.size(), 4);
+
+  std::atomic<int> counter{0};
+  absl::Notification notification;
+  for (int i = 0; i < 4; ++i) {
+    threadpool.schedule([&counter, &notification]() {
+      if (++counter == 4) {
+        notification.Notify();
+      }
+    });
+  }
+  EXPECT_TRUE(
+      notification.WaitForNotificationWithTimeout(absl::Milliseconds(500)));
+  EXPECT_EQ(counter, 4);
+}
+
+TEST(ThreadPoolTest, CpuCoreBindingVerifyAffinity) {
+  // Verify that after construction the thread is actually bound to the
+  // requested core (if the system allows it).
+  const int32_t target_core = 0;
+  std::vector<int32_t> cpu_cores = {target_core};
+
+  absl::Notification done;
+  std::atomic<bool> affinity_ok{false};
+
+  ThreadPool threadpool(1, cpu_cores);
+  threadpool.schedule([&done, &affinity_ok, target_core]() {
+    cpu_set_t cpu_set;
+    CPU_ZERO(&cpu_set);
+    if (pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpu_set) ==
+        0) {
+      affinity_ok = CPU_ISSET(target_core, &cpu_set);
+    }
+    done.Notify();
+  });
+  EXPECT_TRUE(done.WaitForNotificationWithTimeout(absl::Milliseconds(500)));
+  EXPECT_TRUE(affinity_ok);
+}
+
+}  // namespace xllm