Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ option(USE_MUSA "Enable MUSA support" OFF)
option(USE_DCU "Enable DCU support" OFF)
option(USE_MACA "Enable MACA support" OFF)
option(ENABLE_HA "Enable Mooncake etcd-based high availability support" OFF)
option(USE_XLITE "Enable xlite backend (NPU only)" OFF)
add_compile_definitions(YLT_ENABLE_IBV)
add_definitions(-DYLT_ENABLE_IBV)
set(YLT_ENABLE_IBV ON)
Expand Down Expand Up @@ -477,6 +478,12 @@ if(USE_NPU)
set(CMAKE_VERBOSE_MAKEFILE ON)
add_definitions(-DTORCH_HIGHER_THAN_PTA6)

if (USE_XLITE)
add_definitions(-DUSE_XLITE)
include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/xlite.cmake)
message(STATUS "USE_XLITE is ON")
endif()

# Use vcpkg header files as the first priority search directory,
#-> because the scope of third-party software managed by vcpkg is used throughout the entire xllm.
message(STATUS "VCPKG_INCLUDE_DIR = ${CMAKE_BINARY_DIR}/vcpkg_installed/${VCPKG_TARGET_TRIPLET}/include")
Expand Down
23 changes: 23 additions & 0 deletions cmake/xlite.cmake
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
# xlite link helper: find_package(xlite) when USE_NPU AND USE_XLITE are ON.

function(xllm_link_xlite target)
if(NOT USE_NPU OR NOT USE_XLITE)
return()
endif()

if(NOT TARGET xlite::xlite)
execute_process(
COMMAND ${Python_EXECUTABLE} -c "import xlite; print(xlite.cmake_prefix_path)"
OUTPUT_VARIABLE _XLITE_CMAKE_PREFIX
OUTPUT_STRIP_TRAILING_WHITESPACE
RESULT_VARIABLE _XLITE_IMPORT_RESULT)
if(NOT _XLITE_IMPORT_RESULT EQUAL 0 OR _XLITE_CMAKE_PREFIX STREQUAL "")
message(FATAL_ERROR "USE_XLITE is ON but xlite not found. Install xlite or pass -DUSE_XLITE=OFF.")
endif()

find_package(xlite REQUIRED CONFIG PATHS "${_XLITE_CMAKE_PREFIX}" NO_DEFAULT_PATH)
message(STATUS "xlite::xlite found via find_package (${_XLITE_CMAKE_PREFIX})")
endif()

target_link_libraries(${target} PRIVATE xlite::xlite)
endfunction()
3 changes: 3 additions & 0 deletions xllm/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,9 @@ if(USE_MUSA)
)
endif()

if (USE_NPU AND USE_XLITE)
xllm_link_xlite(xllm)
endif()
# install xllm
install(TARGETS xllm RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})

Expand Down
15 changes: 15 additions & 0 deletions xllm/core/framework/model/causal_lm.h
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,10 @@ namespace layer {
struct AttentionMetadata;
}

namespace xlite {
class XliteModelHolder;
}

struct ModelGraphMetadataState {
virtual ~ModelGraphMetadataState() = default;
};
Expand Down Expand Up @@ -151,6 +155,9 @@ class CausalLM : public torch::nn::Module {
NOT_IMPLEMENTED();
return false;
}

// xlite runtime access, nullptr for non-xlite models.
virtual xlite::XliteModelHolder* get_xlite_holder() { return nullptr; }
#endif

virtual layer::LmHead get_lm_head() {
Expand Down Expand Up @@ -444,6 +451,14 @@ class CausalLMImpl : public CausalLM {
requested_rolling_slots,
model_id);
}

// Forward to inner Model.
xlite::XliteModelHolder* get_xlite_holder() override {
if constexpr (detail::has_get_xlite_holder<Model>::value) {
return model_->get_xlite_holder();
}
return CausalLM::get_xlite_holder();
}
#endif

layer::LmHead get_lm_head() override {
Expand Down
9 changes: 9 additions & 0 deletions xllm/core/framework/model/model_traits.h
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,15 @@ struct has_init_or_refresh_rolling_runtime<
std::declval<int32_t>(),
std::declval<const std::string&>()))>> : std::true_type {};

// SFINAE: xlite-backend models expose get_xlite_holder().
template <typename T, typename = void>
struct has_get_xlite_holder : std::false_type {};

template <typename T>
struct has_get_xlite_holder<
T,
std::void_t<decltype(std::declval<T>()->get_xlite_holder())>>
: std::true_type {};
#endif

template <typename T, typename = void>
Expand Down
84 changes: 84 additions & 0 deletions xllm/core/layers/xlite/xlite_attn_meta_builder.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
/* Copyright 2025-2026 The xLLM Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

https://github.com/jd-opensource/xllm/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.
==============================================================================*/

// Build xlite XModelAttnMeta from xllm ModelInputParams.

#pragma once

#include <xlite/xlite.h>

#include <algorithm>
#include <cstdint>
#include <vector>

#include "core/framework/model/model_input_params.h"
#include "core/layers/xlite/xlite_init_utils.h"

namespace xllm::xlite {

class XliteAttnMetaBuilder {
public:
static void Build(const ModelInputParams& params,
const torch::Tensor& positions,
uint32_t block_size,
XModelAttnMeta& m,
int64_t pad_count = 0) {
// version=0: xlite recomputes position from cachedLens (framework's
// positions tensor is off-by-one on decode).
m.version = 0;
m.lens.clear();
m.cachedLens.clear();
m.blockTables.clear();

const auto& host = params.attention.host;
int n = params.meta.num_sequences;
uint32_t bs = block_size;

// block_tables may be undefined in edge cases (DP empty shard).
const bool has_real_seqs =
n > 0 && host.block_tables.defined() && host.block_tables.dim() >= 2;
if (has_real_seqs) {
auto block_acc = host.block_tables.accessor<int32_t, 2>();
for (int s = 0; s < n; ++s) {
int32_t q_len = host.q_seq_lens[s];
int32_t kv_len = host.kv_seq_lens[s];
m.lens.push_back(static_cast<uint32_t>(q_len));
m.cachedLens.push_back(
static_cast<uint32_t>(std::max(0, kv_len - q_len))); // clamp >= 0
int32_t nblocks =
(kv_len + static_cast<int32_t>(bs) - 1) / static_cast<int32_t>(bs);
std::vector<uint32_t> row(nblocks);
for (int32_t b = 0; b < nblocks; ++b) {
row[b] = static_cast<uint32_t>(block_acc[s][b]);
}
m.blockTables.push_back(std::move(row));
}
}

// DP padding: append dummy seq so sum(lens) aligns across DP groups.
if (pad_count > 0) {
m.lens.push_back(static_cast<uint32_t>(pad_count));
m.cachedLens.push_back(0);
int32_t nblocks =
(static_cast<int32_t>(pad_count) + static_cast<int32_t>(bs) - 1) /
static_cast<int32_t>(bs);
std::vector<uint32_t> row(nblocks, 0);
m.blockTables.push_back(std::move(row));
}
InitXTensor(m.vllmPosition, positions);
}
};

} // namespace xllm::xlite
Loading