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/xllm/pybind/CMakeLists.txt

@@ -0,0 +1,30 @@
+include(pybind_extension)
+
+pybind_extension(
+  NAME
+    xllm_export
+  COPTS
+    -DPY_MODULE_NAME=xllm_export
+  SRCS
+    bind.cpp
+  DEFINES
+    PYBIND11_DETAILED_ERROR_MESSAGES=1
+  LINKDIRS
+    ${TORCH_INSTALL_PREFIX}/lib
+  DEPS
+    :master
+    :request
+    :util
+    absl::strings
+    brpc
+    gflags::gflags
+    glog::glog
+    Python::Module
+    torch_python
+    torch
+    c10
+)
+target_link_options(xllm_export PRIVATE -Wl,-Bsymbolic)
+target_link_libraries(common PRIVATE leveldb::leveldb OpenSSL::SSL OpenSSL::Crypto protobuf::libprotobuf)
+add_dependencies(common brpc-static)
+

upstream_ref/xllm/xllm/pybind/args.py

@@ -0,0 +1,49 @@
+import argparse
+from argparse import Namespace
+
+class ArgumentParser:
+    def __init__(self) -> None:
+        self.parser = argparse.ArgumentParser()
+        self.parser.add_argument('--model', type=str, help='"Name or path of the huggingface model to use."')
+        self.parser.add_argument('--task', type=str, default="generate", help='The task to use the model for. generate/embed.')
+        self.parser.add_argument('--runner', type=str, choices=['pooling'], default=None, help='Optional runner mode for LLM. Currently supports: pooling.')
+        self.parser.add_argument('--devices', type=str, default='auto', help='Devices to run the model on, e.g. cpu, cuda:0, cuda:0,cuda:1, or auto to use all available gpus.')
+        self.parser.add_argument('--draft_model', type=str, default='', help='draft hf model path to the model file.')
+        self.parser.add_argument('--draft_devices', type=str, default='auto', help='Devices to run the draft model on, e.g. cpu, cuda:0, cuda:0,cuda:1, or auto to use all available gpus.')
+        self.parser.add_argument('--block_size', type=int, default=128, help='Number of slots per kv cache block. Default is 128.')
+        self.parser.add_argument('--max_cache_size', type=int, default=0, help='Max gpu memory size for kv cache. Default is 0, which means cache size is caculated by available memory.')
+        self.parser.add_argument('--max_memory_utilization', type=float, default=0.9, help='The fraction of GPU memory to be used for model inference, including model weights and kv cache.')
+        self.parser.add_argument('--disable_prefix_cache', action='store_true', help='disable the prefix cache for the block manager.')
+        self.parser.add_argument('--max_tokens_per_batch', type=int, default=20000, help='Max number of tokens per batch.')
+        self.parser.add_argument('--max_seqs_per_batch', type=int, default=256, help='Max number of sequences per batch.')
+        self.parser.add_argument('--max_tokens_per_chunk_for_prefill', type=int, default=512, help='Max number of tokens per chunk for request in prefill stage.')
+        self.parser.add_argument('--num_speculative_tokens', type=int, default=0, help='Number of speculative tokens.')
+        self.parser.add_argument('--num_request_handling_threads', type=int, default=4, help='Number of handling threads.')
+        self.parser.add_argument('--communication_backend', type=str, default='hccl', help='npu communication backend.')
+        self.parser.add_argument('--rank_tablefile', type=str, default='', help='atb hccl rank table file')
+        self.parser.add_argument('--expert_parallel_degree', type=int, default=0, help='ep degree')
+        self.parser.add_argument('--disable_chunked_prefill', action='store_true', help='Whether to disable chunked prefill.')
+        self.parser.add_argument('--enable_prefill_sp', action='store_true', help='Enable prefill-only sequence parallel.')
+        self.parser.add_argument('--master_node_addr', type=str, default='', help='The master address for multi-node distributed serving(e.g. 10.18.1.1:9999).')
+        self.parser.add_argument('--instance_role', type=str, default='DEFAULT', help='The role of instance(e.g. DEFAULT, PREFILL, DECODE, MIX).')
+        self.parser.add_argument('--device_ip', type=str, default='', help='The device ip.')
+        self.parser.add_argument('--transfer_listen_port', type=int, default=26000, help='The KVCacheTranfer listen port.')
+        self.parser.add_argument('--nnodes', type=int, default=1, help='The number of multi-nodes.')
+        self.parser.add_argument('--node_rank', type=int, default=0, help='The node rank.')
+        self.parser.add_argument('--dp_size', type=int, default=1, help='Data parallel size for MLA attention.')
+        self.parser.add_argument('--ep_size', type=int, default=1, help='Expert parallel size for MoE model.')
+        self.parser.add_argument('--instance_name', type=str, default='', help='instance name')
+        self.parser.add_argument('--enable_disagg_pd', action='store_true', help='Enable disaggregated prefill and decode execution.')
+        self.parser.add_argument('--enable_pd_ooc', action='store_true', help='Enable online-offline co-location in disaggregated prefill-decoding mode.')
+        self.parser.add_argument('--enable_schedule_overlap', action='store_true', help='Whether to enable schedule overlap.')
+        self.parser.add_argument('--kv_cache_transfer_mode', type=str, default='PUSH', help='The mode of kv cache transfer(e.g. PUSH, PULL).')
+        self.parser.add_argument('--enable_multi_stream_parallel', action='store_true', help='Whether to enable computation communication overlap.')
+        self.parser.add_argument('--disable_ttft_profiling', action='store_true', help='Whether to disable TTFT profiling.')
+        self.parser.add_argument('--enable_forward_interruption', action='store_true', help='Whether to enable forward interruption.')
+        self.parser.add_argument('--enable_shm', action='store_true', help='Use shared memory for inter-process communication in the single-machine multi-GPU scenario.')
+        self.parser.add_argument('--input_shm_size', type=int, default=1024, help='The size of input shared memory in MB.')
+        self.parser.add_argument('--output_shm_size', type=int, default=128, help='The size of output shared memory in MB.')
+        self.parser.add_argument('--kv_cache_dtype', type=str, default='auto', help='KV cache data type. "auto" (default) aligns with model dtype, "int8" enables INT8 quantization (MLU only).')
+
+    def parse_args(self) -> Namespace:
+        return self.parser.parse_args()

upstream_ref/xllm/xllm/pybind/bind.cpp

@@ -0,0 +1,357 @@
+/* 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 <pybind11/functional.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+#include <torch/python.h>
+
+#include "api_service/call.h"
+#include "core/common/options.h"
+#include "core/common/types.h"
+#include "core/distributed_runtime/llm_master.h"
+#include "core/distributed_runtime/vlm_master.h"
+#include "core/framework/request/mm_data.h"
+#include "core/framework/request/request_output.h"
+#include "core/framework/request/request_params.h"
+#include "core/framework/request/sample_slot.h"
+#include "models/model_registry.h"
+
+namespace xllm {
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+PYBIND11_MODULE(xllm_export, m) {
+  // 1. export Options
+  py::class_<Options>(m, "Options")
+      .def(py::init())
+      .def_readwrite("model_path", &Options::model_path_)
+      .def_readwrite("devices", &Options::devices_)
+      .def_readwrite("draft_model_path", &Options::draft_model_path_)
+      .def_readwrite("draft_devices", &Options::draft_devices_)
+      .def_readwrite("backend", &Options::backend_)
+      .def_readwrite("block_size", &Options::block_size_)
+      .def_readwrite("max_cache_size", &Options::max_cache_size_)
+      .def_readwrite("max_memory_utilization",
+                     &Options::max_memory_utilization_)
+      .def_readwrite("enable_prefix_cache", &Options::enable_prefix_cache_)
+      .def_readwrite("max_tokens_per_batch", &Options::max_tokens_per_batch_)
+      .def_readwrite("max_seqs_per_batch", &Options::max_seqs_per_batch_)
+      .def_readwrite("max_tokens_per_chunk_for_prefill",
+                     &Options::max_tokens_per_chunk_for_prefill_)
+      .def_readwrite("num_speculative_tokens",
+                     &Options::num_speculative_tokens_)
+      .def_readwrite("num_request_handling_threads",
+                     &Options::num_request_handling_threads_)
+      .def_readwrite("communication_backend", &Options::communication_backend_)
+      .def_readwrite("rank_tablefile", &Options::rank_tablefile_)
+      .def_readwrite("expert_parallel_degree",
+                     &Options::expert_parallel_degree_)
+      .def_readwrite("task_type", &Options::task_type_)
+      .def_readwrite("enable_chunked_prefill",
+                     &Options::enable_chunked_prefill_)
+      .def_readwrite("enable_prefill_sp", &Options::enable_prefill_sp_)
+      .def_readwrite("master_node_addr", &Options::master_node_addr_)
+      .def_readwrite("nnodes", &Options::nnodes_)
+      .def_readwrite("node_rank", &Options::node_rank_)
+      .def_readwrite("dp_size", &Options::dp_size_)
+      .def_readwrite("ep_size", &Options::ep_size_)
+      .def_readwrite("instance_name", &Options::instance_name_)
+      .def_readwrite("enable_disagg_pd", &Options::enable_disagg_pd_)
+      .def_readwrite("enable_pd_ooc", &Options::enable_pd_ooc_)
+      .def_readwrite("enable_schedule_overlap",
+                     &Options::enable_schedule_overlap_)
+      .def_readwrite("instance_role", &Options::instance_role_)
+      .def_readwrite("kv_cache_transfer_mode",
+                     &Options::kv_cache_transfer_mode_)
+      .def_readwrite("device_ip", &Options::device_ip_)
+      .def_readwrite("transfer_listen_port", &Options::transfer_listen_port_)
+      .def_readwrite("disable_ttft_profiling",
+                     &Options::disable_ttft_profiling_)
+      .def_readwrite("enable_forward_interruption",
+                     &Options::enable_forward_interruption_)
+      .def_readwrite("enable_offline_inference",
+                     &Options::enable_offline_inference_)
+      .def_readwrite("spawn_worker_path", &Options::spawn_worker_path_)
+      .def_readwrite("enable_shm", &Options::enable_shm_)
+      .def_readwrite("input_shm_size", &Options::input_shm_size_)
+      .def_readwrite("output_shm_size", &Options::output_shm_size_)
+      .def_readwrite("is_local", &Options::is_local_)
+      .def_readwrite("kv_cache_dtype", &Options::kv_cache_dtype_);
+
+  // 2. export LLMMaster
+  py::class_<LLMMaster>(m, "LLMMaster")
+      .def(py::init<const Options&>(),
+           py::arg("options"),
+           py::call_guard<py::gil_scoped_release>())
+      .def("handle_request",
+           py::overload_cast<std::string,
+                             std::optional<std::vector<int>>,
+                             RequestParams,
+                             std::optional<Call*>,
+                             OutputCallback>(&LLMMaster::handle_request),
+           py::call_guard<py::gil_scoped_release>())
+      .def("handle_request",
+           py::overload_cast<std::vector<Message>,
+                             std::optional<std::vector<int>>,
+                             RequestParams,
+                             std::optional<Call*>,
+                             OutputCallback>(&LLMMaster::handle_request),
+           py::call_guard<py::gil_scoped_release>())
+      .def("handle_batch_request",
+           py::overload_cast<std::vector<std::string>,
+                             std::vector<RequestParams>,
+                             BatchOutputCallback>(
+               &LLMMaster::handle_batch_request),
+           py::call_guard<py::gil_scoped_release>())
+      .def("handle_batch_request",
+           py::overload_cast<std::vector<std::vector<Message>>,
+                             std::vector<RequestParams>,
+                             BatchOutputCallback>(
+               &LLMMaster::handle_batch_request),
+           py::call_guard<py::gil_scoped_release>())
+      .def("run", &LLMMaster::run, py::call_guard<py::gil_scoped_release>())
+      .def("generate",
+           &LLMMaster::generate,
+           py::call_guard<py::gil_scoped_release>())
+      .def("options",
+           &LLMMaster::options,
+           py::call_guard<py::gil_scoped_release>())
+      .def(
+          "build_sample_slots",
+          [](const LLMMaster& self,
+             const std::string& request_id,
+             const std::string& prompt,
+             const std::string& literal) {
+            std::vector<SampleSlot> sample_slots;
+            const bool ok = xllm::build_sample_slots(
+                request_id, prompt, literal, self.tokenizer(), &sample_slots);
+            return std::make_pair(ok, sample_slots);
+          },
+          py::arg("request_id"),
+          py::arg("prompt"),
+          py::arg("literal"),
+          py::call_guard<py::gil_scoped_release>())
+      .def("get_rate_limiter",
+           &LLMMaster::get_rate_limiter,
+           py::call_guard<py::gil_scoped_release>())
+      .def("__repr__", [](const LLMMaster& self) {
+        return "LLMMaster({})"_s.format(self.options());
+      });
+
+  // 3. export SampleSlot
+  py::class_<SampleSlot>(m, "SampleSlot")
+      .def(py::init())
+      .def_readwrite("request_id", &SampleSlot::request_id)
+      .def_readwrite("sample_id", &SampleSlot::sample_id)
+      .def_readwrite("token_position", &SampleSlot::token_position);
+
+  // 4. export RequestParams
+  py::class_<RequestParams>(m, "RequestParams")
+      .def(py::init())
+      .def(py::init([](py::kwargs kwargs) {
+        RequestParams params;
+        py::object obj = py::cast(params);
+        for (const auto& item : kwargs) {
+          if (!py::isinstance<py::str>(item.first)) {
+            throw py::type_error("Keyword argument name must be a string");
+          }
+          py::setattr(obj, item.first, item.second);
+        }
+        return obj.cast<RequestParams>();
+      }))
+      .def_readwrite("request_id", &RequestParams::request_id)
+      .def_readwrite("service_request_id", &RequestParams::service_request_id)
+      .def_readwrite("x_request_id", &RequestParams::x_request_id)
+      .def_readwrite("x_request_time", &RequestParams::x_request_time)
+      .def_readwrite("max_tokens", &RequestParams::max_tokens)
+      .def_readwrite("n", &RequestParams::n)
+      .def_readwrite("best_of", &RequestParams::best_of)
+      .def_readwrite("echo", &RequestParams::echo)
+      .def_readwrite("frequency_penalty", &RequestParams::frequency_penalty)
+      .def_readwrite("presence_penalty", &RequestParams::presence_penalty)
+      .def_readwrite("repetition_penalty", &RequestParams::repetition_penalty)
+      .def_readwrite("temperature", &RequestParams::temperature)
+      .def_readwrite("top_p", &RequestParams::top_p)
+      .def_readwrite("top_k", &RequestParams::top_k)
+      .def_readwrite("logprobs", &RequestParams::logprobs)
+      .def_readwrite("top_logprobs", &RequestParams::top_logprobs)
+      .def_readwrite("skip_special_tokens", &RequestParams::skip_special_tokens)
+      .def_readwrite("ignore_eos", &RequestParams::ignore_eos)
+      .def_readwrite("is_embeddings", &RequestParams::is_embeddings)
+      .def_readwrite("stop", &RequestParams::stop)
+      .def_readwrite("stop_token_ids", &RequestParams::stop_token_ids)
+      .def_readwrite("beam_width", &RequestParams::beam_width)
+      .def_readwrite("num_return_sequences",
+                     &RequestParams::num_return_sequences)
+      .def_readwrite("add_special_tokens", &RequestParams::add_special_tokens)
+      .def_readwrite("is_sample_request", &RequestParams::is_sample_request)
+      .def_readwrite("sample_slots", &RequestParams::sample_slots);
+
+  // 4. export Usage
+  py::class_<Usage>(m, "Usage")
+      .def(py::init())
+      .def_readwrite("num_prompt_tokens", &Usage::num_prompt_tokens)
+      .def_readwrite("num_generated_tokens", &Usage::num_generated_tokens)
+      .def_readwrite("num_total_tokens", &Usage::num_total_tokens)
+      .def_property_readonly(
+          "prompt_tokens",
+          [](const Usage& self) { return self.num_prompt_tokens; })
+      .def_property_readonly(
+          "completion_tokens",
+          [](const Usage& self) { return self.num_generated_tokens; })
+      .def_property_readonly("total_tokens", [](const Usage& self) {
+        return self.num_total_tokens;
+      });
+  // 5. export RequestOutput
+  py::class_<RequestOutput>(m, "RequestOutput")
+      .def(py::init())
+      .def_readwrite("request_id", &RequestOutput::request_id)
+      .def_readwrite("service_request_id", &RequestOutput::service_request_id)
+      .def_readwrite("prompt", &RequestOutput::prompt)
+      .def_readwrite("status", &RequestOutput::status)
+      .def_readwrite("outputs", &RequestOutput::outputs)
+      .def_readwrite("usage", &RequestOutput::usage)
+      .def_readwrite("finished", &RequestOutput::finished)
+      .def_readwrite("cancelled", &RequestOutput::cancelled);
+
+  // 6. export StatusCode
+  py::enum_<StatusCode>(m, "StatusCode")
+      .value("OK", StatusCode::OK)
+      .value("CANCELLED", StatusCode::CANCELLED)
+      .value("UNKNOWN", StatusCode::UNKNOWN)
+      .value("INVALID_ARGUMENT", StatusCode::INVALID_ARGUMENT)
+      .value("DEADLINE_EXCEEDED", StatusCode::DEADLINE_EXCEEDED)
+      .value("RESOURCE_EXHAUSTED", StatusCode::RESOURCE_EXHAUSTED)
+      .export_values();
+
+  // 7. export Status
+  py::class_<Status>(m, "Status")
+      .def(py::init<StatusCode, const std::string&>(),
+           py::arg("code"),
+           py::arg("message"))
+      .def_property_readonly("code", &Status::code)
+      .def_property_readonly("message", &Status::message)
+      .def_property_readonly("ok", &Status::ok)
+      .def("__repr__", [](const Status& self) {
+        if (self.message().empty()) {
+          return "Status(code={})"_s.format(self.code());
+        }
+        return "Status(code={}, message={!r})"_s.format(self.code(),
+                                                        self.message());
+      });
+
+  // 8. export LogProbData
+  py::class_<LogProbData>(m, "LogProbData")
+      .def(py::init())
+      .def_readwrite("token", &LogProbData::token)
+      .def_readwrite("token_id", &LogProbData::token_id)
+      .def_readwrite("logprob", &LogProbData::logprob)
+      .def_readwrite("finished_token", &LogProbData::finished_token)
+      .def("__repr__", [](const LogProbData& self) {
+        return "LogProbData(token={!r}, token_id={}, logprob={})"_s.format(
+            self.token, self.token_id, self.logprob);
+      });
+
+  // 9. export LogProb
+  py::class_<LogProb, LogProbData>(m, "LogProb")
+      .def(py::init())
+      .def_readwrite("top_logprobs", &LogProb::top_logprobs)
+      .def("__repr__", [](const LogProb& self) {
+        return "LogProb(token={!r}, token_id={}, logprob={})"_s.format(
+            self.token, self.token_id, self.logprob);
+      });
+
+  // 10. export SequenceOutput
+  py::class_<SequenceOutput>(m, "SequenceOutput")
+      .def(py::init())
+      .def_readwrite("index", &SequenceOutput::index)
+      .def_readwrite("text", &SequenceOutput::text)
+      .def_readwrite("embedding", &SequenceOutput::embedding)
+      .def_readwrite("token_ids", &SequenceOutput::token_ids)
+      .def_readwrite("finish_reason", &SequenceOutput::finish_reason)
+      .def_readwrite("logprobs", &SequenceOutput::logprobs)
+      .def_readwrite("embeddings", &SequenceOutput::embeddings)
+      .def("__repr__", [](const SequenceOutput& self) {
+        return "SequenceOutput({}: {!r})"_s.format(self.index, self.text);
+      });
+
+  // 11. export MMType
+  py::enum_<MMType::Value>(m, "MMType")
+      .value("NONE", MMType::Value::NONE)
+      .value("IMAGE", MMType::Value::IMAGE)
+      .value("VIDEO", MMType::Value::VIDEO)
+      .value("AUDIO", MMType::Value::AUDIO)
+      .export_values();
+
+  // 12. export MMData
+  py::class_<MMData>(m, "MMData")
+      .def(py::init<int, const MMDict&>(), py::arg("ty"), py::arg("data"))
+      .def("get",
+           [](const MMData& self, const MMKey& key) -> py::object {
+             auto value = self.get<torch::Tensor>(key);
+             if (value.has_value()) {
+               return py::cast(value.value());
+             }
+             return py::none();
+           })
+      .def("get_list",
+           [](const MMData& self, const MMKey& key) -> py::object {
+             auto value = self.get<std::vector<torch::Tensor>>(key);
+             if (value.has_value()) {
+               return py::cast(value.value());
+             }
+             return py::none();
+           })
+      .def("__repr__", [](const MMData& self) {
+        std::stringstream ss;
+        ss << "MMData(" << self.type() << ": " << self.size() << " items)";
+        return ss.str();
+      });
+
+  // 13. export VLMMaster
+  py::class_<VLMMaster>(m, "VLMMaster")
+      .def(py::init<const Options&>(),
+           py::arg("options"),
+           py::call_guard<py::gil_scoped_release>())
+      .def("handle_batch_request",
+           py::overload_cast<std::vector<std::string>,
+                             std::vector<MMData>,
+                             std::vector<RequestParams>,
+                             BatchOutputCallback>(
+               &VLMMaster::handle_batch_request),
+           py::call_guard<py::gil_scoped_release>())
+      .def("handle_batch_request_with_image_urls",
+           py::overload_cast<std::vector<std::string>,
+                             std::vector<std::vector<std::string>>,
+                             std::vector<RequestParams>,
+                             BatchOutputCallback>(
+               &VLMMaster::handle_batch_request_with_image_urls),
+           py::call_guard<py::gil_scoped_release>())
+      .def("generate",
+           &VLMMaster::generate,
+           py::call_guard<py::gil_scoped_release>())
+      .def("__repr__", [](const VLMMaster& self) {
+        return "VLMMaster({})"_s.format(self.options());
+      });
+
+  // 12. export helpers
+  m.def("get_model_backend",
+        &ModelRegistry::get_model_backend,
+        py::arg("model_type"));
+}
+
+}  // namespace xllm

upstream_ref/xllm/xllm/pybind/embedding.py

@@ -0,0 +1,139 @@
+import os
+import signal
+import sys
+import time
+from . import util
+from typing import Any, List, Optional, Union
+
+from xllm_export import (LLMMaster, Options, RequestOutput,
+                         RequestParams)
+from .errors import ValidationError
+
+class Embedding:
+    def __init__(
+        self,
+        model: str,
+        devices: str = 'auto',
+        block_size: int = 128,
+        max_cache_size: int = 0,
+        max_memory_utilization: float = 0.9,
+        disable_prefix_cache: bool = False,
+        max_tokens_per_batch: int = 20000,
+        max_seqs_per_batch: int = 256,
+        max_tokens_per_chunk_for_prefill: int = 512,
+        num_request_handling_threads: int = 4,
+        communication_backend: str = 'lccl',
+        rank_tablefile: str = '',
+        expert_parallel_degree: int = 0,
+        disable_chunked_prefill: bool = False,
+        enable_prefill_sp: bool = False,
+        instance_role: str = 'DEFAULT',
+        nnodes: int = 1,
+        node_rank: int = 0,
+        dp_size: int = 1,
+        ep_size: int = 1,
+        enable_shm: bool = False,
+        is_local: bool = True,
+        input_shm_size: int = 1024,
+        output_shm_size: int = 128,
+        **kwargs: Any,
+    ) -> None:
+        signal.signal(signal.SIGTERM, lambda s, f: sys.exit(0))
+        signal.signal(signal.SIGINT, lambda s, f: sys.exit(0))
+
+        if not os.path.exists(model):
+            raise ValueError(f"model {model} not exists")
+
+        options = Options()
+        options.model_path = model
+        options.task_type = "embed"
+        options.devices = devices
+        options.draft_model_path = None
+        options.draft_devices = None
+        options.backend = "llm"
+        options.block_size = block_size
+        options.max_cache_size = max_cache_size
+        options.max_memory_utilization = max_memory_utilization
+        if disable_prefix_cache:
+            options.enable_prefix_cache = False
+        else:
+            options.enable_prefix_cache = True
+        options.max_tokens_per_batch = max_tokens_per_batch
+        options.max_seqs_per_batch = max_seqs_per_batch
+        options.max_tokens_per_chunk_for_prefill = max_tokens_per_chunk_for_prefill
+        options.num_request_handling_threads = num_request_handling_threads
+        options.communication_backend = communication_backend
+        options.rank_tablefile = rank_tablefile
+        options.expert_parallel_degree = expert_parallel_degree
+        if disable_chunked_prefill:
+            options.enable_chunked_prefill = False
+        else:
+            options.enable_chunked_prefill = True
+        options.enable_prefill_sp = enable_prefill_sp
+        free_port = util.get_free_port()
+        options.master_node_addr = "127.0.0.1:" + str(free_port)
+        options.nnodes = nnodes
+        options.node_rank = node_rank
+        options.dp_size = dp_size
+        options.ep_size = ep_size
+        options.enable_disagg_pd = False
+        options.enable_schedule_overlap = False
+        options.enable_offline_inference = True
+        options.spawn_worker_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+        options.enable_shm = enable_shm
+        options.is_local = is_local
+        options.input_shm_size = input_shm_size
+        options.output_shm_size = output_shm_size
+        self.master = LLMMaster(options)
+
+    def finish(self) -> None:
+        try:
+            #os.kill(os.getpid(), signal.SIGTERM)
+            #os.kill(os.getpid(), signal.SIGKILL)
+            util.terminate_process(os.getpid())
+        except Exception as e:
+            pass
+
+    def embedding(
+        self,
+        inputs: Union[str, List[str]],
+        request_params: Optional[Union[RequestParams, List[RequestParams]]] = None,
+        wait_for_schedule: bool = True,
+    ) -> List[RequestOutput]:
+        if request_params is None:
+            request_params = RequestParams()
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        if isinstance(request_params, RequestParams):
+            request_params.is_embeddings = True
+            request_params = [request_params]
+        else:
+            for i in range(len(request_params)):
+                request_params[i].is_embeddings = True
+
+        outputs = [None] * len(inputs)
+        def callback(index: int, output: RequestOutput) -> bool:
+            outputs[index] = output
+            return True
+
+        # schedule all requests
+        self.master.handle_batch_request(
+            inputs, request_params, callback
+        )
+
+        # TODO: add wait later
+        if wait_for_schedule:
+            pass
+
+        # generate
+        self.master.generate()
+
+        # wait async output
+        for i in range(len(outputs)):
+            while outputs[i] is None:
+                time.sleep(0.01)
+            if outputs[i].status is not None and not outputs[i].status.ok:
+                raise ValidationError(outputs[i].status.code, outputs[i].status.message)
+            outputs[i].prompt = inputs[i]
+
+        return outputs

upstream_ref/xllm/xllm/pybind/errors.py

@@ -0,0 +1,5 @@
+class ValidationError(Exception):
+    def __init__(self, code: int, message: str) -> None:
+        super().__init__(f"[{code}] {message}")
+        self.code: int = code
+        self.message: str = message

upstream_ref/xllm/xllm/pybind/llm.py

@@ -0,0 +1,485 @@
+import json
+import os
+import signal
+import sys
+import time
+import uuid
+from . import util
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import xllm_export
+from xllm_export import (LLMMaster, VLMMaster, Options, RequestOutput,
+                         RequestParams)
+from .errors import ValidationError
+from .params import (
+    BeamSearchParams,
+    PoolingParams,
+    SamplingParams,
+    to_request_params,
+    to_request_params_list,
+)
+
+def _read_json(path: str) -> Dict[str, object]:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _infer_model_backend(model_path: str) -> str:
+    model_index_path = os.path.join(model_path, "model_index.json")
+    if os.path.exists(model_index_path):
+        data = _read_json(model_index_path)
+        if "_diffusers_version" in data:
+            return "dit"
+
+    config_path = os.path.join(model_path, "config.json")
+    if not os.path.exists(config_path):
+        raise ValueError(
+            "config.json or model_index.json is required for backend detection"
+        )
+    data = _read_json(config_path)
+    model_type = data.get("model_type") or data.get("model_name")
+    if not model_type:
+        raise ValueError("config.json must contain model_type or model_name")
+
+    get_backend = getattr(xllm_export, "get_model_backend", None)
+    if not callable(get_backend):
+        raise ValueError(
+            "xllm_export.get_model_backend is not available. "
+            "Please rebuild xllm_export or explicitly specify backend."
+        )
+    try:
+        backend = get_backend(model_type)
+    except Exception as exc:
+        raise ValueError(f"Failed to resolve backend for model_type: {model_type}") from exc
+    if not backend:
+        raise ValueError(f"Unsupported model_type: {model_type}")
+    return backend
+
+
+class BeamSearchOutput:
+    def __init__(self, output: RequestOutput):
+        self.prompt = output.prompt
+        self.sequences = output.outputs
+        self.status = output.status
+        self.usage = output.usage
+        self.request_output = output
+
+
+class EmbeddingOutputs:
+    def __init__(self, output: RequestOutput):
+        embedding = []
+        if output.outputs and len(output.outputs) > 0:
+            embedding = output.outputs[0].embeddings
+        self.embedding = embedding
+        self.embeddings = embedding
+
+
+class EmbeddingOutput:
+    def __init__(self, output: RequestOutput):
+        self.prompt = output.prompt
+        self.outputs = EmbeddingOutputs(output)
+        self.status = output.status
+        self.usage = output.usage
+        self.request_output = output
+
+
+class LLM:
+    @staticmethod
+    def _is_vllm_style_inputs(prompts: object) -> bool:
+        if isinstance(prompts, dict):
+            return True
+        if isinstance(prompts, list) and prompts and all(isinstance(x, dict) for x in prompts):
+            return True
+        return False
+
+    def __init__(
+        self,
+        model: str,
+        task: str = "generate",
+        runner: Optional[str] = None,
+        devices: str = 'auto',
+        draft_model: Optional[str] = None,
+        draft_devices: Optional[str] = None,
+        block_size: int = 128,
+        max_cache_size: int = 0,
+        max_memory_utilization: float = 0.9,
+        disable_prefix_cache: bool = False,
+        max_tokens_per_batch: int = 20480,
+        max_seqs_per_batch: int = 256,
+        max_tokens_per_chunk_for_prefill: int = -1,
+        num_speculative_tokens: int = 0,
+        num_request_handling_threads: int = 4,
+        communication_backend: str = 'hccl',
+        rank_tablefile: str = '',
+        expert_parallel_degree: int = 0,
+        disable_chunked_prefill: bool = False,
+        enable_prefill_sp: bool = False,
+        instance_role: str = 'DEFAULT',
+        device_ip: str = '',
+        transfer_listen_port: int = 26000,
+        nnodes: int = 1,
+        node_rank: int = 0,
+        dp_size: int = 1,
+        ep_size: int = 1,
+        instance_name: str = '',
+        enable_disagg_pd: bool = False,
+        enable_pd_ooc: bool = False,
+        enable_schedule_overlap: bool = False,
+        kv_cache_transfer_mode: str = 'PUSH',
+        disable_ttft_profiling: bool = False,
+        enable_forward_interruption: bool = False,
+        enable_shm: bool = False,
+        is_local: bool = True,
+        input_shm_size: int = 1024,
+        output_shm_size: int = 128,
+        **kwargs: Any,
+    ) -> None:
+        signal.signal(signal.SIGTERM, lambda s, f: sys.exit(0))
+        signal.signal(signal.SIGINT, lambda s, f: sys.exit(0))
+
+        if runner is not None:
+            if runner != "pooling":
+                raise ValueError(f"unsupported runner: {runner}")
+            task = "embed"
+
+        if not os.path.exists(model):
+            raise ValueError(f"model {model} not exists")
+
+        backend = _infer_model_backend(model)
+        if backend == "dit":
+            raise ValueError("LLM does not support DiT backend models")
+        if backend == "vlm" and task != "generate":
+            raise ValueError("VLM backend only supports generate task in LLM")
+
+        options = Options()
+        options.model_path = model
+        options.task_type = task
+        options.devices = devices
+        options.draft_model_path = draft_model
+        options.draft_devices = draft_devices
+        options.backend = backend
+        options.block_size = block_size
+        options.max_cache_size = max_cache_size
+        options.max_memory_utilization = max_memory_utilization
+        if disable_prefix_cache:
+            options.enable_prefix_cache = False
+        else:
+            options.enable_prefix_cache = True
+        options.max_tokens_per_batch = max_tokens_per_batch
+        options.max_seqs_per_batch = max_seqs_per_batch
+        options.max_tokens_per_chunk_for_prefill = max_tokens_per_chunk_for_prefill
+        options.num_speculative_tokens = num_speculative_tokens
+        options.num_request_handling_threads = num_request_handling_threads
+        options.communication_backend = communication_backend
+        options.rank_tablefile = rank_tablefile
+        options.expert_parallel_degree = expert_parallel_degree
+        if disable_chunked_prefill:
+            options.enable_chunked_prefill = False
+        else:
+            options.enable_chunked_prefill = True
+        options.enable_prefill_sp = enable_prefill_sp
+        free_port = util.get_free_port()
+        options.master_node_addr = "127.0.0.1:" + str(free_port)
+        options.device_ip = device_ip
+        options.transfer_listen_port = transfer_listen_port
+        options.nnodes = nnodes
+        options.node_rank = node_rank
+        options.dp_size = dp_size
+        options.ep_size = ep_size
+        options.instance_name = instance_name
+        options.enable_disagg_pd = enable_disagg_pd
+        options.enable_schedule_overlap = False
+        options.enable_pd_ooc = enable_pd_ooc
+        options.kv_cache_transfer_mode = kv_cache_transfer_mode
+        options.disable_ttft_profiling = disable_ttft_profiling
+        options.enable_forward_interruption = enable_forward_interruption
+        options.enable_offline_inference = True
+        options.spawn_worker_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+        options.enable_shm = enable_shm
+        options.is_local = is_local
+        options.input_shm_size = input_shm_size
+        options.output_shm_size = output_shm_size
+        self._backend = backend
+        if backend == "vlm":
+            self.master = VLMMaster(options)
+        else:
+            self.master = LLMMaster(options)
+
+    def finish(self) -> None:
+        try:
+            #os.kill(os.getpid(), signal.SIGTERM)
+            #os.kill(os.getpid(), signal.SIGKILL)
+            util.terminate_process(os.getpid())
+        except Exception as e:
+            pass
+
+    def generate(
+        self,
+        prompts: Union[
+            str,
+            List[str],
+            Dict[str, object],
+            List[Dict[str, object]],
+        ],
+        sampling_params: Optional[Union[
+            SamplingParams,
+            List[SamplingParams],
+        ]] = None,
+        wait_for_schedule: bool = True,
+        **kwargs: Any,
+    ) -> List[RequestOutput]:
+        request_params = kwargs.pop("request_params", None)
+        if kwargs:
+            unknown = ", ".join(kwargs.keys())
+            raise TypeError(f"Unexpected keyword arguments: {unknown}")
+        if request_params is None:
+            request_params = sampling_params
+        elif sampling_params is not None:
+            raise ValueError("sampling_params and request_params cannot both be set")
+
+        mm_datas = None
+        image_urls = None
+        if self._is_vllm_style_inputs(prompts):
+            from . import mm_utils
+            prompts, mm_datas, image_urls = mm_utils.normalize_vllm_style_inputs(prompts)
+        else:
+            if isinstance(prompts, str):
+                prompts = [prompts]
+            if not isinstance(prompts, list) or not all(isinstance(x, str) for x in prompts):
+                raise TypeError("prompts must be str/list[str] or vLLM-style dicts")
+
+        request_params_list = to_request_params_list(
+            request_params, default_cls=SamplingParams)
+        if len(request_params_list) not in (1, len(prompts)):
+            raise ValueError(
+                "The number of request_params must be 1 or equal to the "
+                "number of prompts."
+            )
+
+        outputs = [None] * len(prompts)
+        def callback(index: int, output: RequestOutput) -> bool:
+            outputs[index] = output
+            return True
+
+        # schedule all requests
+        if self._backend == "vlm":
+            if mm_datas is not None:
+                self.master.handle_batch_request(
+                    prompts, mm_datas, request_params_list, callback
+                )
+            else:
+                if image_urls is None:
+                    image_urls = [[] for _ in prompts]
+                self.master.handle_batch_request_with_image_urls(
+                    prompts, image_urls, request_params_list, callback
+                )
+        else:
+            has_images = image_urls is not None and any(image_urls)
+            if mm_datas is not None or has_images:
+                raise ValueError("multi_modal_data is only supported for VLM models")
+            self.master.handle_batch_request(
+                prompts, request_params_list, callback
+            )
+
+        # TODO: add wait later
+        if wait_for_schedule:
+            pass
+
+        # generate
+        self.master.generate()
+
+        count = len(prompts)
+        idx = 0
+        while idx < count:
+            # wait async output
+            if outputs[idx] is None:
+                continue
+            if outputs[idx].status is not None and not outputs[idx].status.ok:
+                raise ValidationError(outputs[idx].status.code, outputs[idx].status.message)
+            outputs[idx].prompt = prompts[idx]
+            idx += 1
+
+        return outputs
+
+    def beam_search(
+        self,
+        prompts: Union[str, Dict[str, str], List[Union[str, Dict[str, str]]]],
+        params: Optional[Union[RequestParams, BeamSearchParams]] = None,
+        wait_for_schedule: bool = True,
+    ) -> List[BeamSearchOutput]:
+        if isinstance(prompts, (str, dict)):
+            prompts = [prompts]
+
+        parsed_prompts: List[str] = []
+        for prompt in prompts:
+            if isinstance(prompt, str):
+                parsed_prompts.append(prompt)
+                continue
+            if isinstance(prompt, dict):
+                if "prompt" not in prompt:
+                    raise ValueError("beam_search prompt dict must contain key 'prompt'")
+                parsed_prompts.append(prompt["prompt"])
+                continue
+            raise TypeError("prompts must be str or dict with key 'prompt'")
+
+        params = to_request_params(params, default_cls=BeamSearchParams)
+        if params.beam_width <= 0:
+            raise ValueError("beam_width must be greater than 0")
+        else:
+            # Beam search relies on top-k logprob candidates from sampler.
+            # Keep this aligned with vLLM's internal default behavior.
+            params.logprobs = True
+            if params.top_logprobs == 0:
+                # if not set top_logprobs, default to returning 2x candidates for better deduplication
+                params.top_logprobs = 2 * params.beam_width
+
+        outputs = self.generate(parsed_prompts,
+                                request_params=params,
+                                wait_for_schedule=wait_for_schedule)
+        return [BeamSearchOutput(output) for output in outputs]
+
+    def embed(
+        self,
+        prompts: Union[str, List[str]],
+        pooling_params: Optional[Union[
+            RequestParams,
+            PoolingParams,
+            List[Union[RequestParams, PoolingParams]],
+        ]] = None,
+        wait_for_schedule: bool = True,
+    ) -> List[EmbeddingOutput]:
+        request_params_list = to_request_params_list(
+            pooling_params, default_cls=PoolingParams)
+        for params in request_params_list:
+            params.is_embeddings = True
+
+        use_params: Union[RequestParams, List[RequestParams]]
+        if len(request_params_list) == 1:
+            use_params = request_params_list[0]
+        else:
+            use_params = request_params_list
+
+        outputs = self.generate(prompts,
+                                request_params=use_params,
+                                wait_for_schedule=wait_for_schedule)
+        return [EmbeddingOutput(output) for output in outputs]
+
+    @staticmethod
+    def _normalize_selector_values(
+        prompts: Sequence[str],
+        selector: Union[str, dict, Sequence[Union[str, dict]]],
+    ) -> List[str]:
+        def get_literal(value: Union[str, dict]) -> str:
+            if isinstance(value, str):
+                return value
+            if isinstance(value, dict):
+                selector_type = value.get("type", "literal")
+                literal = value.get("value", "")
+                if selector_type != "literal":
+                    raise ValueError("selector.type must be literal")
+                if not isinstance(literal, str) or not literal:
+                    raise ValueError("selector.value is required")
+                return literal
+            raise ValueError("selector must be a string or dict")
+
+        if isinstance(selector, (str, dict)):
+            literal = get_literal(selector)
+            return [literal for _ in prompts]
+
+        selector_values = list(selector)
+        if len(selector_values) != len(prompts):
+            raise ValueError("selector count must match prompts count")
+        return [get_literal(item) for item in selector_values]
+
+    @staticmethod
+    def _build_request_params_list(
+        prompts: Sequence[str],
+        request_params: Optional[Union[RequestParams, Sequence[RequestParams]]],
+    ) -> List[RequestParams]:
+        if request_params is None:
+            return [RequestParams() for _ in prompts]
+        if isinstance(request_params, RequestParams):
+            if len(prompts) != 1:
+                raise ValueError(
+                    "request_params must be a list when prompts has multiple items"
+                )
+            return [request_params]
+
+        params_list = list(request_params)
+        if len(params_list) != len(prompts):
+            raise ValueError("request_params count must match prompts count")
+        return params_list
+
+    def sample(
+        self,
+        prompts: Union[str, List[str]],
+        selector: Union[str, dict, Sequence[Union[str, dict]]],
+        request_params: Optional[Union[RequestParams, Sequence[RequestParams]]] = None,
+        logprobs: int = 5,
+        wait_schedule_done: bool = True,
+    ) -> List[RequestOutput]:
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        if not prompts:
+            return []
+
+        selector_values = self._normalize_selector_values(prompts, selector)
+        params_list = self._build_request_params_list(prompts, request_params)
+        if len(params_list) > 1:
+            # sample() 会原地修改每个 RequestParams（如 request_id/sample_slots）。
+            # 若复用同一个对象，会在并发批处理时互相覆盖。
+            unique_param_objects = {id(p) for p in params_list}
+            if len(unique_param_objects) != len(params_list):
+                raise ValueError(
+                    "request_params contains duplicated RequestParams objects. "
+                    "Please create one RequestParams instance per prompt."
+                )
+
+        for i, prompt in enumerate(prompts):
+            params = params_list[i]
+            if not params.request_id:
+                params.request_id = "sample-" + uuid.uuid4().hex
+
+            params.max_tokens = 1
+            params.n = 1
+            params.best_of = 1
+            params.logprobs = True
+            params.top_logprobs = logprobs
+            params.add_special_tokens = True
+            params.is_sample_request = True
+
+            ok, sample_slots = self.master.build_sample_slots(
+                params.request_id,
+                prompt,
+                selector_values[i],
+            )
+            if not ok:
+                raise ValueError(
+                    "Failed to build sample slots. "
+                    "selector.value must be a stable single special token."
+                )
+            params.sample_slots = sample_slots
+
+        outputs = [None] * len(prompts)
+
+        def callback(index: int, output: RequestOutput) -> bool:
+            outputs[index] = output
+            return True
+
+        self.master.handle_batch_request(prompts, params_list, callback)
+
+        if wait_schedule_done:
+            pass
+
+        self.master.generate()
+
+        for i in range(len(outputs)):
+            while outputs[i] is None:
+                time.sleep(0.01)
+            if outputs[i].status is not None and not outputs[i].status.ok:
+                raise RuntimeError(
+                    f"sample request failed: {outputs[i].status.message}"
+                )
+            outputs[i].prompt = prompts[i]
+
+        return outputs

upstream_ref/xllm/xllm/pybind/mm_utils.py

@@ -0,0 +1,113 @@
+import base64
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+
+from PIL import Image
+from xllm_export import MMData
+
+
+def _bytes_to_data_url(data: bytes) -> str:
+    encoded = base64.b64encode(data).decode("ascii")
+    return f"data:image;base64,{encoded}"
+
+
+def _pil_to_data_url(image: Image.Image) -> str:
+    buf = BytesIO()
+    fmt = image.format or "PNG"
+    image.save(buf, format=fmt)
+    return _bytes_to_data_url(buf.getvalue())
+
+
+def normalize_vllm_style_inputs(
+    prompts: Any,
+) -> Tuple[List[str], Optional[List[MMData]], Optional[List[List[str]]]]:
+    if isinstance(prompts, dict):
+        requests = [prompts]
+        return _parse_vllm_style_requests(requests)
+    if isinstance(prompts, list) and prompts and all(isinstance(x, dict) for x in prompts):
+        return _parse_vllm_style_requests(prompts)
+
+    raise TypeError(
+        "VLM-style inputs must be dict/List[dict] with key 'prompt', e.g. "
+        "{'prompt': '...', 'multi_modal_data': {'image': image}}"
+    )
+
+
+def _parse_vllm_style_requests(
+    requests: List[Dict[str, Any]],
+) -> Tuple[List[str], Optional[List[MMData]], Optional[List[List[str]]]]:
+    prompts: List[str] = []
+    mm_datas: List[MMData] = []
+    image_urls: List[List[str]] = []
+    use_mm_data: Optional[bool] = None
+
+    for req in requests:
+        if "prompt" not in req:
+            raise ValueError("Each request dict must contain key 'prompt'")
+
+        prompt = req["prompt"]
+        if not isinstance(prompt, str):
+            raise TypeError("request['prompt'] must be a string")
+        prompts.append(prompt)
+
+        if "multi_modal_data" not in req:
+            if use_mm_data is True:
+                raise TypeError("Cannot mix MMData and empty multi_modal_data in one batch")
+            use_mm_data = False
+            image_urls.append([])
+            continue
+
+        payload = req["multi_modal_data"]
+        if isinstance(payload, MMData):
+            if use_mm_data is False:
+                raise TypeError("Cannot mix MMData and image inputs in one batch")
+            use_mm_data = True
+            mm_datas.append(payload)
+        else:
+            if use_mm_data is True:
+                raise TypeError("Cannot mix MMData and image inputs in one batch")
+            use_mm_data = False
+            image_urls.append(_to_image_urls(payload))
+
+    if use_mm_data:
+        return prompts, mm_datas, None
+    return prompts, None, image_urls
+
+
+def _to_image_urls(payload: Any) -> List[str]:
+    if not isinstance(payload, dict):
+        raise TypeError("multi_modal_data must be dict or MMData")
+
+    if "image" in payload:
+        images = payload["image"]
+        return _normalize_images(images)
+    if "video" in payload:
+        raise NotImplementedError("video multi_modal_data is not supported yet")
+    if "audio" in payload:
+        raise NotImplementedError("audio multi_modal_data is not supported yet")
+
+    raise ValueError(
+        "Unsupported multi_modal_data format. Expected {'image': ...} or MMData."
+    )
+
+
+def _normalize_images(images: Any) -> List[str]:
+    if isinstance(images, (list, tuple)):
+        if len(images) == 0:
+            raise ValueError("multi_modal_data['image'] cannot be empty")
+        return [_to_image_url(img) for img in images]
+    return [_to_image_url(images)]
+
+
+def _to_image_url(image: Any) -> str:
+    if isinstance(image, str):
+        return image
+    if isinstance(image, Image.Image):
+        return _pil_to_data_url(image.convert("RGB"))
+    if isinstance(image, (bytes, bytearray)):
+        return _bytes_to_data_url(bytes(image))
+
+    raise TypeError(
+        "image must be image path/url string, PIL.Image, bytes, "
+        "or a list of these"
+    )

upstream_ref/xllm/xllm/pybind/multimodal.py

@@ -0,0 +1,71 @@
+from typing import Any, Dict, List, cast
+from functools import lru_cache
+from io import BytesIO
+from PIL import Image
+import torch
+
+@lru_cache(maxsize=1)
+def __cache_image_processor(
+    processor_name: str,
+    *args: Any,
+    trust_remote_code: bool = False,
+    **kwargs: Any,
+):
+    """Load an image processor for the given model name via HuggingFace."""
+    # don't put this import at the top level
+    # it will call torch.cuda.device_count()
+    from transformers import AutoImageProcessor
+    from transformers.image_processing_utils import BaseImageProcessor
+
+    try:
+        processor = AutoImageProcessor.from_pretrained(
+            processor_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs)
+    except ValueError as e:
+        if not trust_remote_code:
+            err_msg = (
+                "Failed to load the image processor. If the image processor is "
+                "a custom processor not yet available in the HuggingFace "
+                "transformers library, consider setting "
+                "`trust_remote_code=True` in LLM or using the "
+                "`--trust-remote-code` flag in the CLI.")
+            raise RuntimeError(err_msg) from e
+        else:
+            raise e
+
+    return cast(BaseImageProcessor, processor)
+
+def try_cat_feature(item):
+    if isinstance(item, torch.Tensor):
+        return item
+
+    assert isinstance(item, list), f"expected list, got {type(item)}"
+    lst = []
+
+    for i in item:
+        res = try_cat_feature(i)
+        if isinstance(res, list):
+            lst.extend(res)
+        elif isinstance(res, torch.Tensor):
+            lst.append(res)
+        else:
+            raise TypeError(f"expected list or torch.Tensor, got {type(res)}")
+
+    if len(lst) == 1:
+        return lst[0]
+    if any(t.shape[1:] != lst[0].shape[1:] for t in lst):
+        return lst
+
+    return torch.cat(lst)
+
+def preprocess(lst: List[str], model: str) -> Dict[str, Any]:
+    
+    images = [Image.open(BytesIO(item)) for item in lst]
+    image_processor = __cache_image_processor(model, trust_remote_code=True)
+
+    data = image_processor.preprocess(images, return_tensors="pt").data
+    return { key: try_cat_feature(val)
+            for key, val in data.items()
+            }

upstream_ref/xllm/xllm/pybind/params.py

@@ -0,0 +1,78 @@
+from typing import Any, List, Optional, Type, Union
+
+from xllm_export import RequestParams
+
+
+class _RequestParamsProxy:
+    def __init__(self, **kwargs: Any) -> None:
+        object.__setattr__(self, "_request_params", RequestParams())
+        for key, value in kwargs.items():
+            self._set_field(key, value)
+
+    def _set_field(self, key: str, value: Any) -> None:
+        if not hasattr(self._request_params, key):
+            raise TypeError(f"Unexpected parameter: {key}")
+        setattr(self._request_params, key, value)
+
+    def __getattr__(self, key: str) -> Any:
+        return getattr(self._request_params, key)
+
+    def __setattr__(self, key: str, value: Any) -> None:
+        if key == "_request_params":
+            object.__setattr__(self, key, value)
+            return
+        self._set_field(key, value)
+
+    def to_request_params(self) -> RequestParams:
+        return self._request_params
+
+
+class SamplingParams(_RequestParamsProxy):
+    pass
+
+
+class BeamSearchParams(SamplingParams):
+    def __init__(self,
+                 beam_width: int = 1,
+                 max_tokens: int = 16,
+                 **kwargs: Any) -> None:
+        super().__init__(beam_width=beam_width, max_tokens=max_tokens, **kwargs)
+
+
+class PoolingParams(_RequestParamsProxy):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.is_embeddings = True
+
+
+ParamLike = Union[RequestParams, _RequestParamsProxy]
+ParamsLike = Optional[Union[ParamLike, List[ParamLike]]]
+
+
+def to_request_params(
+    params: Optional[ParamLike],
+    default_cls: Type[_RequestParamsProxy] = SamplingParams,
+) -> RequestParams:
+    if params is None:
+        return default_cls().to_request_params()
+    if isinstance(params, RequestParams):
+        return params
+    if isinstance(params, _RequestParamsProxy):
+        return params.to_request_params()
+    raise TypeError(
+        "Unsupported params type. Expected RequestParams, SamplingParams, "
+        "BeamSearchParams, or PoolingParams."
+    )
+
+
+def to_request_params_list(
+    params: ParamsLike,
+    default_cls: Type[_RequestParamsProxy] = SamplingParams,
+) -> List[RequestParams]:
+    if params is None:
+        return [default_cls().to_request_params()]
+    if isinstance(params, list):
+        if len(params) == 0:
+            return [default_cls().to_request_params()]
+        return [to_request_params(item, default_cls=default_cls) for item in params]
+    return [to_request_params(params, default_cls=default_cls)]

upstream_ref/xllm/xllm/pybind/util.py

@@ -0,0 +1,35 @@
+import os
+import psutil
+import signal
+import socket
+import sys
+from typing import Union
+
+def terminate_process(pid: int, timeout: Union[int, float] = 30) -> None:
+    try:
+        parent = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        return
+
+    children = parent.children(recursive=True)
+    procs = children + [parent]
+
+    for p in procs:
+        try:
+            p.terminate()
+        except psutil.NoSuchProcess:
+            pass
+
+    gone, alive = psutil.wait_procs(procs, timeout=timeout)
+    for p in alive:
+        try:
+            p.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+def get_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(('0.0.0.0', 0))
+        _, port = s.getsockname()
+    return port
+

upstream_ref/xllm/xllm/pybind/vlm.py

@@ -0,0 +1,180 @@
+import os
+import signal
+import sys
+from . import util
+from typing import List, Optional, Union, Dict, Any
+from xllm_export import (VLMMaster, Options, RequestOutput,
+                         RequestParams, MMData)
+from .errors import ValidationError
+from .params import (
+    SamplingParams,
+    to_request_params_list,
+)
+
+
+class VLM:
+    def __init__(
+        self,
+        model: str,
+        task: str = "generate",
+        devices: str = 'auto',
+        draft_model: Optional[str] = None,
+        draft_devices: Optional[str] = None,
+        block_size: int = 128,
+        max_cache_size: int = 0,
+        max_memory_utilization: float = 0.9,
+        disable_prefix_cache: bool = False,
+        max_tokens_per_batch: int = 50000,
+        max_seqs_per_batch: int = 256,
+        max_tokens_per_chunk_for_prefill: int = 512,
+        num_speculative_tokens: int = 0,
+        num_request_handling_threads: int = 4,
+        communication_backend: str = 'hccl',
+        rank_tablefile: str = '',
+        expert_parallel_degree: int = 0,
+        disable_chunked_prefill: bool = False,
+        enable_prefill_sp: bool = False,
+        instance_role: str = 'DEFAULT',
+        device_ip: str = '',
+        transfer_listen_port: int = 26000,
+        nnodes: int = 1,
+        node_rank: int = 0,
+        dp_size: int = 1,
+        ep_size: int = 1,
+        instance_name: str = '',
+        enable_disagg_pd: bool = False,
+        enable_schedule_overlap: bool = False,
+        kv_cache_transfer_mode: str = 'PUSH',
+        enable_shm: bool = False,
+        is_local: bool = True,
+        input_shm_size: int = 1024,
+        output_shm_size: int = 128,
+        **kwargs: Any,
+    ) -> None:
+        signal.signal(signal.SIGTERM, lambda s, f: sys.exit(0))
+        signal.signal(signal.SIGINT, lambda s, f: sys.exit(0))
+
+        if not os.path.exists(model):
+            raise ValueError(f"model {model} not exists")
+        self.model = model
+
+        options = Options()
+        options.model_path = model
+        options.task_type = task
+        options.devices = devices
+        options.draft_model_path = draft_model
+        options.draft_devices = draft_devices
+        options.backend ="vlm"
+        options.block_size = block_size
+        options.max_cache_size = max_cache_size
+        options.max_memory_utilization = max_memory_utilization
+        if disable_prefix_cache:
+            options.enable_prefix_cache = False
+        else:
+            options.enable_prefix_cache = True
+        options.max_tokens_per_batch = max_tokens_per_batch
+        options.max_seqs_per_batch = max_seqs_per_batch
+        options.max_tokens_per_chunk_for_prefill = max_tokens_per_chunk_for_prefill
+        options.num_speculative_tokens = num_speculative_tokens
+        options.num_request_handling_threads = num_request_handling_threads
+        options.communication_backend = communication_backend
+        options.rank_tablefile = rank_tablefile
+        options.expert_parallel_degree = expert_parallel_degree
+        if disable_chunked_prefill:
+            options.enable_chunked_prefill = False
+        else:
+            options.enable_chunked_prefill = True
+        options.enable_prefill_sp = enable_prefill_sp
+        free_port = util.get_free_port()
+        options.master_node_addr = "127.0.0.1:" + str(free_port)
+        options.device_ip = device_ip
+        options.transfer_listen_port = transfer_listen_port
+        options.nnodes = nnodes
+        options.node_rank = node_rank
+        options.dp_size = dp_size
+        options.ep_size = ep_size
+        options.instance_name = instance_name
+        options.enable_disagg_pd = enable_disagg_pd
+        options.enable_schedule_overlap = False
+        options.kv_cache_transfer_mode = kv_cache_transfer_mode
+        options.enable_offline_inference = True
+        options.spawn_worker_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+        options.enable_shm = enable_shm
+        options.is_local = is_local
+        options.input_shm_size = input_shm_size
+        options.output_shm_size = output_shm_size
+        self.master = VLMMaster(options)
+
+    def finish(self) -> None:
+        try:
+            #os.kill(os.getpid(), signal.SIGTERM)
+            #os.kill(os.getpid(), signal.SIGKILL)
+            util.terminate_process(os.getpid())
+        except Exception as e:
+            pass
+
+    def generate(
+        self,
+        prompts: Union[
+            str,
+            List[str],
+            Dict[str, Any],
+            List[Dict[str, Any]],
+        ],
+        sampling_params: Optional[Union[
+            SamplingParams,
+            List[SamplingParams],
+        ]] = None,
+        wait_for_schedule: bool = True,
+        **kwargs: Any,
+    ) -> List[RequestOutput]:
+        from . import mm_utils
+        prompts, mm_datas, image_urls = mm_utils.normalize_vllm_style_inputs(prompts)
+
+        request_params = kwargs.pop("request_params", None)
+        if kwargs:
+            unknown = ", ".join(kwargs.keys())
+            raise TypeError(f"Unexpected keyword arguments: {unknown}")
+        if request_params is None:
+            request_params = sampling_params
+        elif sampling_params is not None:
+            raise ValueError("sampling_params and request_params cannot both be set")
+
+        request_params_list = to_request_params_list(
+            request_params, default_cls=SamplingParams
+        )
+        if len(request_params_list) not in (1, len(prompts)):
+            raise ValueError(
+                "The number of request_params must be 1 or equal to the number of prompts."
+            )
+
+        outputs = [None] * len(prompts)
+        def callback(index: int, output: RequestOutput) -> bool:
+            outputs[index] = output
+            return True
+        # schedule the batch requests
+        if image_urls is not None:
+            self.master.handle_batch_request_with_image_urls(
+                prompts, image_urls, request_params_list, callback
+            )
+        else:
+            self.master.handle_batch_request(
+                prompts, mm_datas, request_params_list, callback
+            )
+
+        # wait for batch request to be scheduled
+        if wait_for_schedule:
+            pass
+
+        # run until all scheduled requsts complete
+        self.master.generate()
+
+        # throw an exception if there is any error
+        for index, output in enumerate(outputs):
+            if output is None:
+                raise RuntimeError("Request failed, no output received")
+            if output.status is not None and not output.status.ok:
+                raise ValidationError(output.status.code, output.status.message)
+            # carry over the prompt to the output
+            output.prompt = prompts[index]
+        return outputs